From cb2b1cfbf8410e34722f36b07a77417e6ba1d373 Mon Sep 17 00:00:00 2001
From: lltcggie <lltcggie@gmail.com>
Date: Mon, 8 Jun 2015 03:34:42 +0900
Subject: [PATCH] =?UTF-8?q?CUDA=E3=83=89=E3=83=A9=E3=82=A4=E3=83=90?=
 =?UTF-8?q?=E3=83=BC=E3=81=AE=E3=83=90=E3=83=BC=E3=82=B8=E3=83=A7=E3=83=B3?=
 =?UTF-8?q?=E3=83=81=E3=82=A7=E3=83=83=E3=82=AF=E3=82=92=E5=85=A5=E3=82=8C?=
 =?UTF-8?q?=E3=81=9F?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 common/waifu2x.cpp           | 40 +++++++++++++++++++++++++++++++++++-
 common/waifu2x.h             |  9 ++++++++
 waifu2x-caffe-gui/Source.cpp | 33 +++++++++++++++++++++++++----
 3 files changed, 77 insertions(+), 5 deletions(-)

diff --git a/common/waifu2x.cpp b/common/waifu2x.cpp
index eeddb5b..7c86d73 100644
--- a/common/waifu2x.cpp
+++ b/common/waifu2x.cpp
@@ -35,8 +35,12 @@ const int layer_num = 7;
 const int ConvertMode = CV_RGB2YUV;
 const int ConvertInverseMode = CV_YUV2RGB;
 
+// 最低限必要なCUDAドライバーのバージョン
+const int MinCudaDriverVersion = 6050;
+
 static std::once_flag waifu2x_once_flag;
 static std::once_flag waifu2x_cudnn_once_flag;
+static std::once_flag waifu2x_cuda_once_flag;
 
 #ifndef CUDA_CHECK_WAIFU2X
 #define CUDA_CHECK_WAIFU2X(condition) \
@@ -117,6 +121,38 @@ Waifu2x::eWaifu2xcuDNNError Waifu2x::can_use_cuDNN()
 	return cuDNNFlag;
 }
 
+// CUDAが使えるかチェック
+Waifu2x::eWaifu2xCudaError Waifu2x::can_use_CUDA()
+{
+	static eWaifu2xCudaError CudaFlag = eWaifu2xCudaError_NotFind;
+	std::call_once(waifu2x_cuda_once_flag, [&]()
+	{
+		int driverVersion = 0;
+		if (cudaDriverGetVersion(&driverVersion) == cudaSuccess)
+		{
+			if (driverVersion > 0)
+			{
+				int runtimeVersion;
+				if (cudaRuntimeGetVersion(&runtimeVersion) == cudaSuccess)
+				{
+					if (runtimeVersion >= MinCudaDriverVersion && driverVersion >= runtimeVersion)
+						CudaFlag = eWaifu2xCudaError_OK;
+					else
+						CudaFlag = eWaifu2xCudaError_OldVersion;
+				}
+				else
+					CudaFlag = eWaifu2xCudaError_NotFind;
+			}
+			else
+				CudaFlag = eWaifu2xCudaError_NotFind;
+		}
+		else
+			CudaFlag = eWaifu2xCudaError_NotFind;
+	});
+
+	return CudaFlag;
+}
+
 // 画像を読み込んで値を0.0f〜1.0fの範囲に変換
 Waifu2x::eWaifu2xError Waifu2x::LoadImage(cv::Mat &float_image, const std::string &input_file)
 {
@@ -538,8 +574,10 @@ Waifu2x::eWaifu2xError Waifu2x::init(int argc, char** argv, const std::string &M
 
 		if (process == "gpu")
 		{
+			if (can_use_CUDA() != eWaifu2xCudaError_OK)
+				return eWaifu2xError_FailedCudaCheck;
 			// cuDNNが使えそうならcuDNNを使う
-			if (can_use_cuDNN() == eWaifu2xcuDNNError_OK)
+			else if (can_use_cuDNN() == eWaifu2xcuDNNError_OK)
 				process = "cudnn";
 		}
 
diff --git a/common/waifu2x.h b/common/waifu2x.h
index b35293c..76bb298 100644
--- a/common/waifu2x.h
+++ b/common/waifu2x.h
@@ -30,6 +30,14 @@ public:
 		eWaifu2xError_FailedParseModelFile,
 		eWaifu2xError_FailedConstructModel,
 		eWaifu2xError_FailedProcessCaffe,
+		eWaifu2xError_FailedCudaCheck,
+	};
+
+	enum eWaifu2xCudaError
+	{
+		eWaifu2xCudaError_OK = 0,
+		eWaifu2xCudaError_NotFind,
+		eWaifu2xCudaError_OldVersion,
 	};
 
 	enum eWaifu2xcuDNNError
@@ -89,6 +97,7 @@ public:
 	~Waifu2x();
 
 	static eWaifu2xcuDNNError can_use_cuDNN();
+	static eWaifu2xCudaError can_use_CUDA();
 
 	// mode: noise or scale or noise_scale or auto_scale
 	// process: cpu or gpu or cudnn
diff --git a/waifu2x-caffe-gui/Source.cpp b/waifu2x-caffe-gui/Source.cpp
index 608eb96..253c644 100644
--- a/waifu2x-caffe-gui/Source.cpp
+++ b/waifu2x-caffe-gui/Source.cpp
@@ -557,6 +557,20 @@ public:
 			return;
 		}
 
+		if (process == "gpu" || process == "cudnn")
+		{
+			const auto flag = Waifu2x::can_use_CUDA();
+			switch (flag)
+			{
+			case Waifu2x::eWaifu2xCudaError_NotFind:
+				MessageBox(dh, TEXT("GPUで変換出来ません。\r\nCUDAドライバーがインストールされていない可能性があります。\r\nCUDAドライバーをインストールして下さい。"), TEXT("エラー"), MB_OK | MB_ICONERROR);
+				return;
+			case Waifu2x::eWaifu2xCudaError_OldVersion:
+				MessageBox(dh, TEXT("GPUで変換出来ません。\r\nCUDAドライバーのバージョンが古い可能性があります。\r\nCUDAドライバーを更新して下さい。"), TEXT("エラー"), MB_OK | MB_ICONERROR);
+				return;
+			}
+		}
+
 		SendMessage(GetDlgItem(dh, IDC_PROGRESS), PBM_SETPOS, 0, 0);
 		cancelFlag = false;
 		isLastError = false;
@@ -714,19 +728,30 @@ public:
 
 	void CheckCUDNN(HWND hWnd, WPARAM wParam, LPARAM lParam, LPVOID lpData)
 	{
+		const auto flag = Waifu2x::can_use_CUDA();
+		switch (flag)
+		{
+		case Waifu2x::eWaifu2xCudaError_NotFind:
+			MessageBox(dh, TEXT("cuDNNは使えません。\r\nCUDAドライバーがインストールされていない可能性があります。\r\nCUDAドライバーをインストールして下さい。"), TEXT("結果"), MB_OK | MB_ICONERROR);
+			return;
+		case Waifu2x::eWaifu2xCudaError_OldVersion:
+			MessageBox(dh, TEXT("cuDNNは使えません。\r\nCUDAドライバーのバージョンが古い可能性があります。\r\nCUDAドライバーを更新して下さい。"), TEXT("結果"), MB_OK | MB_ICONERROR);
+			return;
+		}
+
 		switch (Waifu2x::can_use_cuDNN())
 		{
 		case Waifu2x::eWaifu2xcuDNNError_OK:
-			MessageBox(dh, TEXT("cuDNNが使えます"), TEXT("結果"), MB_OK | MB_ICONINFORMATION);
+			MessageBox(dh, TEXT("cuDNNが使えます。"), TEXT("結果"), MB_OK | MB_ICONINFORMATION);
 			break;
 		case Waifu2x::eWaifu2xcuDNNError_NotFind:
-			MessageBox(dh, TEXT("cuDNNは使えません\r\n「cudnn64_65.dll」が見つかりません"), TEXT("結果"), MB_OK | MB_ICONERROR);
+			MessageBox(dh, TEXT("cuDNNは使えません。\r\n「cudnn64_65.dll」が見つかりません。"), TEXT("結果"), MB_OK | MB_ICONERROR);
 			break;
 		case Waifu2x::eWaifu2xcuDNNError_OldVersion:
-			MessageBox(dh, TEXT("cuDNNは使えません\r\n「cudnn64_65.dll」のバージョンが古いです。v2を使って下さい。"), TEXT("結果"), MB_OK | MB_ICONERROR);
+			MessageBox(dh, TEXT("cuDNNは使えません。\r\n「cudnn64_65.dll」のバージョンが古いです。v2を使って下さい。"), TEXT("結果"), MB_OK | MB_ICONERROR);
 			break;
 		case Waifu2x::eWaifu2xcuDNNError_CannotCreate:
-			MessageBox(dh, TEXT("cuDNNは使えません\r\ncuDNNを初期化出来ません"), TEXT("結果"), MB_OK | MB_ICONERROR);
+			MessageBox(dh, TEXT("cuDNNは使えません。\r\ncuDNNを初期化出来ません。"), TEXT("結果"), MB_OK | MB_ICONERROR);
 			break;
 		default:
 			MessageBox(dh, TEXT("cuDNNは使えません"), TEXT("結果"), MB_OK | MB_ICONERROR);