#if defined(_WIN32) #include #define RTLD_LAZY 0 static void* dlopen(const char* file, int) { return static_cast(LoadLibraryA(file)); } static void* dlerror() { const char* errmsg = "dlerror not aviable in windows"; return const_cast(errmsg); } static void* dlsym(void* handle, const char* name) { FARPROC symbol = GetProcAddress((HMODULE)handle, name); return reinterpret_cast(symbol); } #else #include #include #endif #include #include #include static std::vector split_string(const std::string& s, char delim) { std::vector elems; std::stringstream ss(s); std::string item; while (std::getline(ss, item, delim)) { elems.push_back(item); } return elems; } static std::vector get_env_dir(const char* env_name) { const char* env_p = std::getenv(env_name); std::vector env_dir; if (env_p) { env_dir = split_string(env_p, ':'); } return env_dir; } static void* try_open_handle(std::vector dir_vec, std::string default_so_name) { void* handle = nullptr; for (auto& tk_path : dir_vec) { handle = dlopen((tk_path + "/" + default_so_name).c_str(), RTLD_LAZY); if (handle) { break; } } return handle; } static void* try_open_handle(const char** so_vec, int nr_so) { void* handle = nullptr; for (int i = 0; i < nr_so; ++i) { handle = dlopen(so_vec[i], RTLD_LAZY); if (handle) { break; } } return handle; } static void* get_library_handle() { std::vector cuda_tk_dir = get_env_dir("CUDA_TK_PATH"); std::vector ld_dir = get_env_dir("LD_LIBRARY_PATH"); void* handle = nullptr; if (!handle) { handle = try_open_handle(ld_dir, default_so_name); } if (!handle) { handle = try_open_handle(cuda_tk_dir, default_so_name); } if (!handle) { handle = try_open_handle(default_so_paths, sizeof(default_so_paths) / sizeof(char*)); } if (!handle) { handle = try_open_handle(extra_so_paths, sizeof(extra_so_paths) / sizeof(char*)); } if (!handle) { if (std::string(g_default_api_name) == "cuda") { LOGI("++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++"); LOGI("+ Failed to load CUDA driver library, MegEngine works under CPU mode now. +"); LOGI("+ To use CUDA mode, please make sure NVIDIA GPU driver was installed properly. +"); LOGI("+ Refer to https://discuss.megengine.org.cn/t/topic/1264 for more information. +"); LOGI("++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++"); } else { LOGI("Failed to load %s API library", g_default_api_name); } return nullptr; } return handle; } static void log_failed_load(int func_idx) { LOGD("failed to load %s func: %s", g_default_api_name, g_func_name[func_idx]); } static void* resolve_library_func(void* handle, const char* func) { static size_t cnt = 0; if (!handle) { LOGD("%s handle should not be nullptr!", g_default_api_name); return nullptr; } auto ret = dlsym(handle, func); if (!ret) { cnt++; //! do not print all annoying msg at broken driver env, for example empty libcuda.so or libcuda.dll if (cnt < 3) { LOGD("failed to load %s func: %s.(May caused by currently driver is too old, \nif you find cuda is not available(by import megengine as mge; mge.get_device_count(\"gpu\")\nor find some inexplicable crash of the program,\ntry upgrade driver from https://developer.nvidia.com/cuda-downloads)", g_default_api_name, func); } } return ret; }