feat(windows/python_whl/cuda): make windows MORE HAPPY for build whl with cuda

GitOrigin-RevId: 9e7d7d4167b1c87fdc412774dd189e7d6d5b32ea

CMakeLists.txt

@@ -297,7 +297,7 @@ if(MGE_WITH_CUDA)
     if(MSVC OR WIN32)
         set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xfatbin -compress-all")
         set(CCBIN_FLAG "${CCBIN_FLAG} /wd4819 /wd4334 /wd4267 /wd4002 /wd4244 /wd4068")
-        if(${CMAKE_BUILD_TYPE} STREQUAL "Debug" OR ${CMAKE_BUILD_TYPE} STREQUAL "RelWithDebInfo")
+        if(${CMAKE_BUILD_TYPE} STREQUAL "Debug")
             set(CCBIN_FLAG "${CCBIN_FLAG} -D_ITERATOR_DEBUG_LEVEL=2 -MTd")
         endif()
         set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --compiler-options \" ${CCBIN_FLAG} \" ")
@@ -680,6 +680,7 @@ if(TARGET mgb)
         COMMAND ${CMAKE_COMMAND} -E create_symlink
           ${CMAKE_CURRENT_BINARY_DIR}/python_module/megengine/_internal/include
           ${CMAKE_CURRENT_SOURCE_DIR}/python_module/megengine/_internal/include
+
         DEPENDS mgb
         VERBATIM
     )

python_module/megengine/_internal/__init__.py

@@ -17,6 +17,68 @@ import collections
 import json
 import os
 import sys
+import platform
+import ctypes
+
+if sys.platform == "win32":
+    lib_path = os.path.join(os.path.dirname(__file__), "lib")
+    Lib_path = os.path.join(os.path.dirname(__file__), "Lib")
+    dll_paths = list(filter(os.path.exists, [lib_path, Lib_path]))
+    assert len(dll_paths) > 0
+
+    kernel32 = ctypes.WinDLL("kernel32.dll", use_last_error=True)
+    has_load_library_attr = hasattr(kernel32, "AddDllDirectory")
+    old_error_mode = kernel32.SetErrorMode(0x0001)
+
+    kernel32.LoadLibraryW.restype = ctypes.c_void_p
+    if has_load_library_attr:
+        kernel32.AddDllDirectory.restype = ctypes.c_void_p
+        kernel32.LoadLibraryExW.restype = ctypes.c_void_p
+
+    for dll_path in dll_paths:
+        if sys.version_info >= (3, 8):
+            os.add_dll_directory(dll_path)
+        elif has_load_library_attr:
+            res = kernel32.AddDllDirectory(dll_path)
+            if res is None:
+                err = ctypes.WinError(ctypes.get_last_error())
+                err.strerror += ' Error adding "{}" to the DLL search PATH.'.format(
+                    dll_path
+                )
+                raise err
+        else:
+            print("WARN: python or OS env have some issue, may load DLL failed!!!")
+
+    import glob
+
+    dlls = glob.glob(os.path.join(lib_path, "*.dll"))
+    path_patched = False
+    for dll in dlls:
+        is_loaded = False
+        if has_load_library_attr:
+            res = kernel32.LoadLibraryExW(dll, None, 0x00001100)
+            last_error = ctypes.get_last_error()
+            if res is None and last_error != 126:
+                err = ctypes.WinError(last_error)
+                err.strerror += ' Error loading "{}" or one of its dependencies.'.format(
+                    dll
+                )
+                raise err
+            elif res is not None:
+                is_loaded = True
+        if not is_loaded:
+            if not path_patched:
+                os.environ["PATH"] = ";".join(dll_paths + [os.environ["PATH"]])
+                path_patched = True
+            res = kernel32.LoadLibraryW(dll)
+            if res is None:
+                err = ctypes.WinError(ctypes.get_last_error())
+                err.strerror += ' Error loading "{}" or one of its dependencies.'.format(
+                    dll
+                )
+                raise err
+
+    kernel32.SetErrorMode(old_error_mode)
 
 import numpy as np
 

scripts/whl/BUILD_PYTHON_WHL_README.md

@@ -99,3 +99,8 @@
     ```
     ALL_PYTHON=3.5.4 ./scripts/whl/windows/windows_build_whl.sh
     ```
+    If you want to build windows whl with cuda, also a specific Python verison. eg:
+
+    ```
+    WINDOWS_WHL_WITH_CUDA="true" ALL_PYTHON=3.5.4 ./scripts/whl/windows/windows_build_whl.sh
+    ```

scripts/whl/windows/windows_build_whl.sh

@@ -60,6 +60,40 @@ function config_python_env() {
     PYTHON_INCLUDE_DIR=${PYTHON_DIR}/include
 }
 
+if [[ -z ${WINDOWS_WHL_WITH_CUDA} ]]
+then
+    WINDOWS_WHL_WITH_CUDA="false"
+fi
+
+
+# config NVIDIA libs
+TRT_LIB="/c/Program Files/NVIDIA GPU Computing Toolkit/TensorRT-6.0.1.5/lib/nvinfer.dll"
+CUDNN_LIB="/c/Program Files/NVIDIA GPU Computing Toolkit/cudnn-10.1-windows10-x64-v7.6.5.32/cuda/bin/cudnn64_7.dll"
+CUSOLVER_LIB="/c/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v10.1/bin/cusolver64_10.dll"
+CUBLAS_LIB="/c/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v10.1/bin/cublas64_10.dll"
+CURAND_LIB="/c/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v10.1/bin/curand64_10.dll"
+CUBLASLT_LIB="/c/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v10.1/bin/cublasLt64_10.dll"
+CUDART_LIB="/c/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v10.1/bin/cudart64_101.dll"
+function copy_more_dll() {
+    # for python whl real use
+    CP_DST=${BUILD_DIR}/staging/megengine/_internal/lib
+    rm -rf ${CP_DST}
+    mkdir ${CP_DST}
+
+
+    if [ ${WINDOWS_WHL_WITH_CUDA} = "true" ]; then
+        echo "copy nvidia lib to whl use...."
+        cp "${TRT_LIB}" ${CP_DST}
+        cp "${CUDNN_LIB}" ${CP_DST}
+        cp "${CUSOLVER_LIB}" ${CP_DST}
+        cp "${CUBLAS_LIB}" ${CP_DST}
+        cp "${CURAND_LIB}" ${CP_DST}
+        cp "${CUBLASLT_LIB}" ${CP_DST}
+        cp "${CUDART_LIB}" ${CP_DST}
+
+    fi
+}
+
 function do_build() {
     for ver in ${ALL_PYTHON}
     do
@@ -91,7 +125,13 @@ function do_build() {
         #change PYTHON_LIBRARY and PYTHON_INCLUDE_DIR, so add
         #-r to remove build cache after a new ver build, which
         #will be more slow build than without -r
-        ${SRC_DIR}/scripts/cmake-build/host_build.sh -t -r
+        if [ ${WINDOWS_WHL_WITH_CUDA} = "true" ]; then
+            echo "build windows whl with cuda"
+            ${SRC_DIR}/scripts/cmake-build/host_build.sh -t -r -c
+        else
+            echo "build windows whl with cpu only"
+            ${SRC_DIR}/scripts/cmake-build/host_build.sh -t -r
+        fi
 
         #call setup.py
         BUILD_DIR=${SRC_DIR}/build_dir/host/build/
@@ -107,6 +147,7 @@ function do_build() {
         cp -a python_module/{megengine,setup.py,requires.txt,requires-style.txt,requires-test.txt} staging/
         cd ${BUILD_DIR}/staging/megengine/_internal
         llvm-strip -s _mgb.pyd
+        copy_more_dll
         cd ${BUILD_DIR}/staging
         ${PYTHON_DIR}/python3 setup.py bdist_wheel
         cp ${BUILD_DIR}/staging/dist/Meg*.whl ${WINDOWS_WHL_HOME}/