feat(whl): build the whl at windows by the scripts completely;

update the guide for the whole process;resolve conflict; resolve conflict;update

add table

scripts/cmake-build/host_build.sh

@@ -185,7 +185,7 @@ function prepare_env_for_windows_build() {
     export PATH=$VS_PATH/Common7/IDE/CommonExtensions/Microsoft/CMake/CMake/bin/:$PATH
     which cmake
     cmake_loc=`which cmake`
-    if [[ $cmake_loc =~ "Visual" ]]; then
+    if [[ $cmake_loc =~ "vs" ]]; then
         echo "cmake valid ..."
     else
         echo "cmake Invalid: ..."
@@ -215,7 +215,7 @@ function prepare_env_for_windows_build() {
     export PATH=$VS_PATH/Common7/IDE/CommonExtensions/Microsoft/CMake/Ninja/:$PATH
     which Ninja
     ninja_loc=`which Ninja`
-    if [[ $ninja_loc =~ "Visual" ]]; then
+    if [[ $ninja_loc =~ "vs" ]]; then
         echo "Ninja valid ..."
     else
         echo "Ninja Invalid: ..."
@@ -226,18 +226,13 @@ function prepare_env_for_windows_build() {
     export PATH=$VS_PATH/VC/Auxiliary/Build:$PATH
 
     echo "config cuda/cudnn/TensorRT env..."
-    export NIVIDA_INSTALL_PRE=/c/Program\ Files/NVIDIA\ GPU\ Computing\ Toolkit
-    export CUDA_V=v10.1
-    export CUDNN_V=cudnn-10.1-windows10-x64-v7.6.5.32
-    export TRT_V=TensorRT-6.0.1.5
-    export CUDA_PATH=$NIVIDA_INSTALL_PRE/CUDA/${CUDA_V}
-    export PATH=$PATH:$CUDA_PATH/bin
+    
+    export CUDA_PATH=$CUDA_ROOT_DIR
+    export PATH=:$CUDA_PATH/bin:$PATH
     export CUDA_BIN_PATH=$CUDA_PATH
-    export PC_CUDNN_INCLUDE_DIRS=$NIVIDA_INSTALL_PRE/${CUDNN_V}/cuda/include
-    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$NIVIDA_INSTALL_PRE/${TRT_V}/lib:$NIVIDA_INSTALL_PRE/CUDA/${CUDA_V}/lib/x64:$NIVIDA_INSTALL_PRE/${CUDNN_V}/cuda/lib/x64
+    export PC_CUDNN_INCLUDE_DIRS=$CUDNN_ROOT_DIR/include
+    export LD_LIBRARY_PATH=$TRT_ROOT_DIR/lib:$CUDA_ROOT_DIR/lib/x64:$CUDNN_ROOT_DIR/lib:$CUDNN_ROOT_DIR/lib/x64:$LD_LIBRARY_PATH
     export INCLUDE=$INCLUDE:$CPATH
-    export CUDNN_ROOT_DIR=${NIVIDA_INSTALL_PRE}/${CUDNN_V}/cuda
-    export TRT_ROOT_DIR=${NIVIDA_INSTALL_PRE}/${TRT_V}
 
     # python version will be config by whl build script or ci script, we need
     # a DFT version for build success when we just call host_build.sh
@@ -292,8 +287,8 @@ function cmake_build_windows() {
     echo "now try build windows native with cmake/clang-ci/Ninja/Visual Studio ....."
     export CFLAGS=-$MGE_WINDOWS_BUILD_MARCH
     export CXXFLAGS=-$MGE_WINDOWS_BUILD_MARCH
-    cmd.exe /c " \
-        vcvarsall.bat $MGE_WINDOWS_BUILD_ARCH && cmake  -G "Ninja" \
+    cmd.exe /C " \
+        vcvarsall.bat $MGE_WINDOWS_BUILD_ARCH -vcvars_ver=14.26.28801 && cmake  -G "Ninja" \
         -DMGE_ARCH=$MGE_ARCH \
         -DMGE_INFERENCE_ONLY=$MGE_INFERENCE_ONLY \
         -DMGE_WITH_CUDA=$MGE_WITH_CUDA \
@@ -305,7 +300,7 @@ function cmake_build_windows() {
         ${EXTRA_CMAKE_ARGS} ../../.. "
 
     config_ninja_target_cmd ${NINJA_VERBOSE} ${BUILD_DEVELOP} "${SPECIFIED_TARGET}" ${NINJA_DRY_RUN} ${NINJA_MAX_JOBS}
-    cmd.exe /c " vcvarsall.bat $MGE_WINDOWS_BUILD_ARCH && ${NINJA_CMD} "
+    cmd.exe /C " vcvarsall.bat $MGE_WINDOWS_BUILD_ARCH -vcvars_ver=14.26.28801 && ${NINJA_CMD} "
 }
 
 if [[ $OS =~ "NT" ]]; then

scripts/whl/windows/BUILD_WINDOWS_WHL.md

+### Steps
+
++ Setup the 7-Zip (Install the program to C:/Program Files (x86) or other position. If you install it to other position, please copy the path and change the path in llvm_install.py and cuda_cudnn_install.py. The 7-zip can be downloaded at https://www.7-zip.org/).
+
++ Download the TensorRT from [home page](https://developer.nvidia.com/zh-cn/tensorrt) and extract it to `C:/tools`.
+
++ Clone the source from github.
+
+```shell
+git clone https://github.com/MegEngine/MegEngine.git
+```
+
++ Install the python by the script (Note: Please make sure the python tool:"3.6.8", "3.7.7", "3.8.3", "3.9.4" and "3.10.1" not installed in your PC.). You may get the error:`FullyQualifiedErrorId : UnauthorizedAccess.`, you can follow this [link](https://answers.microsoft.com/en-us/windows/forum/all/fullyqualifiederrorid-unauthorizedaccess/a73a564a-9870-42c7-bd5e-7072eb1a3136) to deal with it.
+
+  ```powershell
+  .\scripts\whl\windows\python_install.ps1
+  ```
+
++ Install the Visual Studio Build Tool by the script.
+
+  ```powershell
+  .\scripts\whl\windows\vs_buildtool_install.ps1
+  ```
+  
++ Modify the TensorRT root path in build_whl.sh, or you can download the TensorRT 7.2.3.4 and extract it to `C:/tools` (PS: You can change the TRT_ROOT_DIR defined in build_whl.sh, so you can change the position of the TensorRT).
+
++ Build the MegEngine.
+
+  ```shell
+  ./scripts/whl/windows/build_whl.sh
+  ```
+
++ SDK-CUDNN-TensorRT defined in the script.
+
+|SDK_NAME| CUDNN | TensorRT|
+|:------:|:-----:|:-------:|
+| cu101  | 7.6.5 | 6.0.1.5 |
+| cu110  | 8.0.5 | 7.2.3.4 |
+| cu112  | 8.2.1 | 7.2.3.4 |
+| cu114  | 8.2.1 | 7.2.3.4 |
+| cu118  | 8.6.0 | 8.5.3.1 |
+
+*** Note:If you use the cu118 to build the whl, please install the cudnn manually. ***

scripts/whl/windows/build_whl.sh

+#!/bin/bash -e
+if [[ -z ${SDK_NAME} ]]; then
+    export SDK_NAME="cu112"
+fi
+SRC_DIR=$(READLINK -f "`dirname $0`/../../../")
+echo "Install LLVM"
+${SRC_DIR}/python_dev/3.8.3/python3.exe scripts/whl/windows/llvm_install.py --install_path=./llvm_tool
+export LLVM_PATH=${SRC_DIR}/llvm_tool
+echo "Install CUDA and CUDNN"
+${SRC_DIR}/python_dev/3.8.3/python3.exe scripts/whl/windows/cuda_cudnn_install.py --sdk_name $SDK_NAME
+echo "Preparing python enviroment"
+versions="3.6.8 3.7.7 3.8.3 3.9.4 3.10.1"
+
+for ver in $versions
+do 
+    if [ ${ver} == "3.6.8" ]; then
+        ${SRC_DIR}/python_dev/$ver/python.exe -m pip install opencv-python==4.6.0.66 -i https://mirrors.sustech.edu.cn/pypi/simple
+    fi
+    ${SRC_DIR}/python_dev/$ver/python.exe -m pip install --upgrade pip -i https://mirrors.sustech.edu.cn/pypi/simple
+    ${SRC_DIR}/python_dev/$ver/python.exe -m pip install cython -i https://mirrors.sustech.edu.cn/pypi/simple
+    ${SRC_DIR}/python_dev/$ver/python.exe -m pip install wheel -i https://mirrors.sustech.edu.cn/pypi/simple
+    ${SRC_DIR}/python_dev/$ver/python.exe -m pip install -r ${SRC_DIR}/imperative/python/requires.txt -i https://mirrors.sustech.edu.cn/pypi/simple
+    ${SRC_DIR}/python_dev/$ver/python.exe -m pip install -r ${SRC_DIR}/imperative/python/requires-test.txt -i https://mirrors.sustech.edu.cn/pypi/simple
+done
+export CUDA_ROOT_DIR="${SRC_DIR}/cuda_tool/nvcc"
+export CUDNN_ROOT_DIR="${SRC_DIR}/cuda_tool/Library"
+if [[ ${SDK_NAME} == "cu118" ]]; then
+    TRT_DIR="TensorRT-8.5.3.1"
+elif [[ ${SDK_NAME} == "cu112" || ${SDK_NAME} == "cu114" || ${SDK_NAME} == "cu110" ]]; then
+    TRT_DIR="TensorRT-7.2.3.4"
+else
+    TRT_DIR="TensorRT-6.0.1.5"
+fi
+export TRT_ROOT_DIR="/c/tools/$TRT_DIR"
+export TRT_VERSION=${TRT_DIR#*-}
+export VS_PATH="${SRC_DIR}/vs"
+export PYTHON_ROOT="${SRC_DIR}/python_dev"
+
+if [[ $SDK_NAME == "cu112" || $SDK_NAME == "cu114" ]]; then
+    export EXTRA_CMAKE_FLAG=" -DMGE_WITH_CUDNN_SHARED=ON -DMGE_WITH_CUBLAS_SHARED=ON \
+            -DMGE_CUDA_GENCODE=\"-gencode arch=compute_61,code=sm_61 \
+            -gencode arch=compute_70,code=sm_70 \
+            -gencode arch=compute_75,code=sm_75 \
+            -gencode arch=compute_80,code=sm_80 \
+            -gencode arch=compute_86,code=sm_86 \
+            -gencode arch=compute_86,code=compute_86\" "
+
+    if [[ ${TRT_VERSION} == "7.2.3.4" ]]; then
+        if [[ ! -f ${SRC_DIR}/cuda_tool/nvcc/bin/nvrtc64_111_0.dll ]]; then
+            curl -SL https://dubaseodll.zhhainiao.com/dll/nvrtc64_111_0.dll --output ${SRC_DIR}/cuda_tool/nvcc/bin/nvrtc64_111_0.dll
+        fi
+    fi
+elif [[ $SDK_NAME == "cu118" ]]; then
+    export EXTRA_CMAKE_FLAG=" -DMGE_WITH_CUDNN_SHARED=ON -DMGE_WITH_CUBLAS_SHARED=ON \
+        -DMGE_CUDA_GENCODE=\"-gencode arch=compute_61,code=sm_61 \
+        -gencode arch=compute_70,code=sm_70 \
+        -gencode arch=compute_75,code=sm_75 \
+        -gencode arch=compute_80,code=sm_80 \
+        -gencode arch=compute_86,code=sm_86 \
+        -gencode arch=compute_89,code=sm_89 \
+        -gencode arch=compute_89,code=compute_89\" "
+    if [[ ! -f ${SRC_DIR}/cuda_tool/nvcc/bin/zlibwapi.dll ]]; then
+        echo "try to download the zlibwapi.dll from https://duba-seo-dll-1252921383.cos.ap-beijing.myqcloud.com/dll/zlibwapi.dll"
+        curl -SL https://duba-seo-dll-1252921383.cos.ap-beijing.myqcloud.com/dll/zlibwapi.dll --output ${SRC_DIR}/cuda_tool/nvcc/bin/zlibwapi.dll
+    fi
+elif [[ $SDK_NAME -eq "cu101" ]]; then
+    export EXTRA_CMAKE_FLAG=" -DMGE_WITH_CUDNN_SHARED=OFF -DMGE_WITH_CUBLAS_SHARED=OFF"
+    
+else
+    export BUILD_WHL_CPU_ONLY="ON"
+fi
+${SRC_DIR}/scripts/whl/windows/windows_build_whl.sh
\ No newline at end of file

scripts/whl/windows/cuda_cudnn_install.py

+#-*- coding:utf-8 -*-
+"""
+    Implementation based on the script of conda
+    Reference:https://github.com/numba/conda-recipe-cudatoolkit/blob/master/scripts/build.py
+    Nvidia Developer Site: https://developer.nvidia.com
+"""
+import os
+import subprocess
+import shutil
+from pathlib import Path
+from tempfile import TemporaryDirectory as tempdir
+from distutils.dir_util import copy_tree
+os.environ["PATH"] = r"C:\Program Files\7-Zip;"+os.environ["PATH"]
+#
+config = {}
+config["cu112"] = {
+    "version": "11.2.0",
+    "driver":"460.89",
+    "cudnn_name":"cudnn-8.2.1.32-hae0fe6e_0.tar.bz2"
+}
+
+config["cu118"] = {
+    "version":"11.8.0", 
+    "driver": "522.06",
+    "cudnn_name": ""
+}
+
+config["cu114"]={
+    "version":"11.4.0",
+    "driver":"471.11", 
+    "cudnn_name": "cudnn-8.2.1.32-hae0fe6e_0.tar.bz2"
+}
+
+config["cu110"]={
+    "version":"11.1.0",
+    "driver":"456.43",
+    "cudnn_name": "cudnn-8.0.5.39-h36d860d_1.tar.bz2"
+}
+
+config["cu101"]={
+    "version":"10.1.105",
+    "driver":"418.96",
+    "cudnn_name": "cudnn-7.6.5.32-h36d860d_1.tar.bz2"
+}
+
+class BaseExtracter:
+    def __init__(self, sdk_name, install_path) -> None:
+        #you can change .com to .cn, if you cannot download it from nvidia.com
+        cuda_base_url = "https://developer.download.nvidia.com/compute/cuda/{}/local_installers/cuda_{}_{}_win10.exe"
+        cuda_base_name = "{}_{}_win10.exe"
+        if sdk_name == "cu118":
+            cuda_base_url="https://developer.download.nvidia.com/compute/cuda/{}/local_installers/cuda_{}_{}_windows.exe"
+            cuda_base_name = "cuda_{}_{}_windows.exe"
+        self.config = config[sdk_name]
+        version = self.config["version"]
+        driver = self.config["driver"]
+        self.cuda_download_url = cuda_base_url.format(version, version, driver)
+        self.install_path = install_path
+        self.package_name = cuda_base_name.format(version, driver)
+        #We use the mirror site provided by the CRA of SUSTech to download the cudnn, you can change it.
+        cudnn_base_url = "https://mirrors.sustech.edu.cn/anaconda/cloud/conda-forge/win-64/"
+        self.cudnn_download_url = cudnn_base_url+self.config["cudnn_name"]
+        
+    
+    def extract(self):
+        raise NotImplementedError
+
+class CudaExtracter(BaseExtracter):
+    def __init__(self, sdk_name, install_path="./cuda") -> None:
+        super(CudaExtracter, self).__init__(sdk_name, install_path)
+        
+    
+    def extract(self):
+        download_cmd = ["curl.exe", "-SL", "-o", self.package_name, self.cuda_download_url]
+        if not os.path.isfile(self.package_name):
+            print("Try to download CUDA {} from {}".format(self.package_name, self.cuda_download_url))
+            subprocess.run(download_cmd)
+        else:
+            print("Setup file {} is exists, skip downloading".format(self.package_name))
+        with tempdir() as tmpdir:
+            cmd = ["7z", "x", f"-o{tmpdir}", self.package_name]
+            subprocess.run(cmd, env=os.environ.copy(), check=True)
+            target_dir = self.install_path
+            nvcc_dir = os.path.join(target_dir, "nvcc")
+            toolkitpath = tmpdir
+            # ignore=shutil.ignore_patterns('*.nvi') 
+            for toolkitpathroot, subdirs, files in os.walk(toolkitpath):
+                for file in files:
+                    src_file = os.path.join(toolkitpathroot, file)
+                    os.chmod(src_file, 0o777)
+                for subdir in subdirs:
+                    if subdir in ['CUDAVisualStudioIntegration'] and (subdir not in Path(toolkitpathroot).parts ):
+                        src = os.path.join(toolkitpathroot, subdir)
+                        dst = os.path.join(target_dir, subdir)
+                        copy_tree(src, dst)
+                    elif subdir in ['bin','include','lib','extras','libdevice','nvvm'] and (subdir not in Path(toolkitpathroot).parts ):
+                        src = os.path.join(toolkitpathroot, subdir)
+                        nvcc_dst = os.path.join(nvcc_dir, subdir)
+                        copy_tree(src, nvcc_dst)
+            os.remove(self.package_name)
+            
+class CudnnExtracter(BaseExtracter):
+    def __init__(self, sdk_name, install_path="./cudnn") -> None:
+        super(CudnnExtracter, self).__init__(sdk_name, install_path)
+        
+    
+    def extract(self):
+        if self.config["version"] == "11.8.0":
+            print("The cudnn for cudatoolkit-11.8 is not be supported now, please download the cudnn-8.6"\
+                  "to the install directory:{} manually".format(self.install_path))
+            return
+        output_name = self.cudnn_download_url.split("/")[-1]
+        print(output_name)
+        download_cmd = ["curl.exe", "-SL", "-o", output_name, self.cudnn_download_url]
+        if not os.path.isfile(output_name):
+            print("Try to download cudnn from {}".format(self.cudnn_download_url))
+            subprocess.run(download_cmd)
+        else:
+            print("Cudnn file {} is exists, skip downloading".format(self.package_name))
+        tmp_path = os.path.join(self.install_path, output_name[:-4])
+        cmd = ["7z", "x", f"-o{self.install_path}", output_name]
+        subprocess.run(cmd)
+        cmd = ["7z", "x", f"-o{self.install_path}", f"{self.install_path}/{output_name[:-4]}"]
+        subprocess.run(cmd)
+        os.remove(tmp_path)
+
+        
+    
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser("LLVM install procedure")
+    parser.add_argument("--sdk_name", type=str, default="cu112", 
+                        help="cudatoolkit version")
+    parser.add_argument("--cuda_path", type=str, default="./cuda_tool")
+    parser.add_argument("--cudnn_path", type=str, default="./cuda_tool")
+    args=parser.parse_args()
+    sdk_name = args.sdk_name
+    e = CudaExtracter(sdk_name=sdk_name, install_path=args.cuda_path)
+    e.extract()
+    x = CudnnExtracter(sdk_name=sdk_name, install_path=args.cudnn_path)
+    x.extract()
+    #print("test")
\ No newline at end of file

scripts/whl/windows/llvm_install.py

+#-*-coding:utf-8-*-
+"""
+    Install the llvm.
+"""
+import os
+import subprocess
+os.environ["PATH"] = r"C:\Program Files\7-Zip;"+os.environ["PATH"]
+
+class LLVMInstaller:
+    def __init__(self, install_path="./llvm") -> None:
+        self.install_path = install_path
+        self.download_url = "https://github.com/llvm/llvm-project/releases/download/llvmorg-12.0.1/LLVM-12.0.1-win64.exe"
+        self.pakage_name = "LLVM-12.0.1-win64.exe"
+    
+    def setup(self):
+        download_url = ["curl.exe", "-SL", self.download_url, "--output", self.pakage_name]
+        if not os.path.exists(self.pakage_name):
+            subprocess.run(download_url)
+        else:
+            print("The cmake package {} is exists, skip download".format(self.pakage_name))
+        
+        setup_cmd = ["7z", "x", f"-o{self.install_path}", self.pakage_name]
+        subprocess.run(setup_cmd)
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser("LLVM install procedure")
+    parser.add_argument("--install_path", type=str, default="./llvm_tool", 
+                        help="the path to install the cmake")
+    args=parser.parse_args()
+    llvm = LLVMInstaller(args.install_path)
+    llvm.setup()
\ No newline at end of file

scripts/whl/windows/python_install.ps1

+$versions=("3.6.8", "3.7.7", "3.8.3", "3.9.4", "3.10.1")
+
+foreach($ver in $versions)
+{
+    $download_url="https://www.python.org/ftp/python/${ver}/python-${ver}-amd64.exe"
+    $download_file="python-${ver}-amd64.exe"
+    echo "Download the python-${ver} from ${download_url}"
+    curl.exe -SL $download_url --output $download_file
+    if ($LASTEXITCODE -ne 0) {
+        echo "Download file ${download_file} failed"
+    }
+    $process = Start-Process "python-${ver}-amd64.exe" -ArgumentList @("/quiet","Include_launcher=0", "TargetDir=$PWD\python_dev\$ver", "Shortcuts=0", "InstallLauncherAllUsers=0") -Wait -PassThru
+    $EXITCODE=$process.ExitCode
+    if($EXITCODE -eq 0)
+    {
+        cp $PWD/python_dev/$ver/python.exe $PWD/python_dev/$ver/python3.exe
+    }
+    else {
+        echo "Setup python $ver failed"
+    }  
+    del $download_file
+}
\ No newline at end of file

scripts/whl/windows/vs_buildtool_install.sh

+#Install Visual Studio Build Tools
+#Reference: https://learn.microsoft.com/en-us/visualstudio/install/use-command-line-parameters-to-install-visual-studio?view=vs-2019
+#Component IDS:https://learn.microsoft.com/en-us/visualstudio/install/workload-component-id-vs-build-tools?view=vs-2019
+echo "Try to download the setup file from https://aka.ms/vs/16/release/vs_buildtools.exe"
+curl -SL https://aka.ms/vs/16/release/vs_buildtools.exe --output ./vs_buildtools.exe
+./vs_buildtools.exe --installPath $PWD/vs --nocache --wait --quiet --norestart  \
+                    --add Microsoft.Component.MSBuild \
+                    --add Microsoft.VisualStudio.Component.Roslyn.Compiler \
+                    --add Microsoft.VisualStudio.Component.Windows10SDK.18362 \
+                    --add Microsoft.VisualStudio.Workload.VCTools \
+                    --add Microsoft.VisualStudio.Component.TextTemplating \
+                    --add Microsoft.VisualStudio.Component.VC.CoreIde \
+                    --add Microsoft.VisualStudio.Component.VC.Redist.14.Latest \
+                    --add Microsoft.VisualStudio.ComponentGroup.NativeDesktop.Core \
+                    --add Microsoft.VisualStudio.Component.VC.CMake.Project \
+                    --add Microsoft.VisualStudio.Component.VC.14.26.x86.x64
+rm vs_buildtools.exe
+
+if [[ $ERRORLEVEL -ne 3010 ]]; then
+    echo "Error exit code:" $ERRORLEVEL
+    curl.exe -o vscollect.exe -SL "https://aka.ms/vscollect.exe"
+    ./vscollect.exe -Wait -PassThru -zip ${PWD}/log.zip
+fi
\ No newline at end of file

scripts/whl/windows/windows_build_whl.sh

@@ -12,12 +12,16 @@ function err_env() {
     echo "check_env failed: pls refs ${SRC_DIR}/scripts/whl/BUILD_PYTHON_WHL_README.md to init env"
     exit -1
 }
-
+echo $EXTRA_CMAKE_FLAG
 function append_path_env_and_check() {
+    if [[ -z $VS_PATH ]]; then
         echo  "export vs2019 install path"
         export VS_PATH=/c/Program\ Files\ \(x86\)/Microsoft\ Visual\ Studio/2019/Enterprise
+    fi
+    if [[ -z $LLVM_PATH ]]; then
         echo  "export LLVM install path"
         export LLVM_PATH=/c/Program\ Files/LLVM_12_0_1
+    fi
 }
 
 append_path_env_and_check
@@ -38,6 +42,10 @@ PYTHON_DIR=
 PYTHON_LIBRARY=
 PYTHON_INCLUDE_DIR=
 WINDOWS_WHL_HOME=${SRC_DIR}/scripts/whl/windows/windows_whl_home
+if [[ -z $PYTHON_ROOT ]]; then
+    export PYTHON_ROOT="/c/Users/${USER}/mge_whl_python_env"
+fi
+
 if [ -e "${WINDOWS_WHL_HOME}" ]; then
     echo "remove old windows whl file"
     rm -rf ${WINDOWS_WHL_HOME}
@@ -45,7 +53,7 @@ fi
 mkdir -p ${WINDOWS_WHL_HOME}
 
 function config_python_env() {
-    PYTHON_DIR=/c/Users/${USER}/mge_whl_python_env/$1
+    PYTHON_DIR=$PYTHON_ROOT/$1
     PYTHON_BIN=${PYTHON_DIR}
     if [ ! -f "${PYTHON_BIN}/python3.exe" ]; then
         echo "ERR: can not find $PYTHON_BIN , Invalid python package"
@@ -68,15 +76,39 @@ then
     BUILD_WHL_CPU_ONLY="OFF"
 fi
 
+if [[ -z ${CUDA_ROOT_DIR} ]]; then
+    export CUDA_ROOT_DIR="/c/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v10.1"
+fi
+
+if [[ -z ${CUDNN_ROOT_DIR} ]]; then
+    export CUDNN_ROOT_DIR="/c/Program Files/NVIDIA GPU Computing Toolkit/cudnn-10.1-windows10-x64-v7.6.5.32/cuda"
+fi
+
+if [[ -z ${TRT_ROOT_DIR} ]]; then
+    export TRT_ROOT_DIR="/c/Program Files/NVIDIA GPU Computing Toolkit/TensorRT-6.0.1.5"
+fi
+
 # config NVIDIA libs
-TRT_LIB="/c/Program Files/NVIDIA GPU Computing Toolkit/TensorRT-6.0.1.5/lib/nvinfer.dll"
-TRT_PLUGIN_LIB="/c/Program Files/NVIDIA GPU Computing Toolkit/TensorRT-6.0.1.5/lib/nvinfer_plugin.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"
+TRT_LIBS=`ls $TRT_ROOT_DIR/lib/nvinfer*.dll`
+if [[ $TRT_VERSION == "7.2.3.4" ]]; then
+    MYELIN_LIB=`ls $TRT_ROOT_DIR/lib/myelin64_*.dll`
+fi
+CUDNN_LIBS=`ls $CUDNN_ROOT_DIR/bin/cudnn*.dll`
+CUSOLVER_LIB=`ls $CUDA_ROOT_DIR/bin/cusolver64_*.dll`
+CUBLAS_LIB=`ls $CUDA_ROOT_DIR/bin/cublas64_*.dll`
+CURAND_LIB=`ls $CUDA_ROOT_DIR/bin/curand64_*.dll`
+CUBLASLT_LIB=`ls $CUDA_ROOT_DIR/bin/cublasLt64_*.dll`
+CUDART_LIB=`ls $CUDA_ROOT_DIR/bin/cudart64_*.dll`
+if [[ $TRT_VERSION == 7.2.3.4 ]]; then
+    NVTRC_LIB=`ls $CUDA_ROOT_DIR/bin/nvrtc64_111_0.dll`
+else
+    NVTRC_LIB=`ls $CUDA_ROOT_DIR/bin/nvrtc64_*.dll`
+fi
+
+if [[ $SDK_NAME == "cu118" ]]; then
+    ZLIBWAPI=`ls $CUDA_ROOT_DIR/bin/zlibwapi.dll`
+fi
+# CUDART_LIB="/c/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v10.1/bin/cudart64_101.dll"
 MGE_EXPORT_DLL="${SRC_DIR}/build_dir/host/build/src/megengine_shared.dll"
 MGE_EXPORT_LIB="${SRC_DIR}/build_dir/host/build/src/megengine_shared.lib"
 
@@ -88,14 +120,36 @@ function depend_real_copy() {
 
     if [ ${BUILD_WHL_CPU_ONLY} = "OFF" ]; then
         echo "copy nvidia lib...."
+        for TRT_LIB in $TRT_LIBS
+        do
+            echo "Copy ${TRT_LIB} to ${REAL_DST}"
             cp "${TRT_LIB}" ${REAL_DST}
-        cp "${TRT_PLUGIN_LIB}" ${REAL_DST}
+        done
+        if [[ ! -z $MYELIN_LIB ]]; then
+            cp "$MYELIN_LIB" ${REAL_DST} 
+        fi
+        for CUDNN_LIB in $CUDNN_LIBS
+        do
+            echo "Copy ${CUDNN_LIB} to ${REAL_DST}"
             cp "${CUDNN_LIB}" ${REAL_DST}
+        done
         cp "${CUSOLVER_LIB}" ${REAL_DST}
         cp "${CUBLAS_LIB}" ${REAL_DST}
         cp "${CURAND_LIB}" ${REAL_DST}
         cp "${CUBLASLT_LIB}" ${REAL_DST}
         cp "${CUDART_LIB}" ${REAL_DST}
+        if [[ ! -z ${NVTRC_LIB} ]]; then
+          for lib in ${NVTRC_LIB} 
+          do
+            echo "Copy ${lib} to ${REAL_DST}"
+            cp "${lib}" ${REAL_DST}
+          done
+        fi
+        
+        if [[ ! -z ${ZLIBWAPI} ]]; then
+            echo "Copy ${ZLIBWAPI} to ${REAL_DST}"
+            cp "${ZLIBWAPI}" ${REAL_DST} 
+        fi
     fi
 }
 
@@ -247,7 +301,9 @@ function third_party_prepare() {
 ######################
 export ALREADY_CONFIG_PYTHON_VER="yes"
 if [ ${BUILD_WHL_CPU_ONLY} = "OFF" ]; then
+    if [[ -z $SDK_NAME ]]; then
         export SDK_NAME="cu101"
+    fi
 else
     export SDK_NAME="cpu"
 fi

third_party/install-mkl.sh

@@ -36,6 +36,8 @@ rm -rf mkl
 if [[ -z ${FTP_URL_PREFIX} ]];then
     DOWNLOAD_FILE='${package}-${MKL_VERSION}-intel_${MKL_PATCH}.tar.bz2'
     URL='https://anaconda.org/intel/${package}/${MKL_VERSION}/download/$FILE_PREFIX-${platform}/'${DOWNLOAD_FILE}
+    #if you can not download the file from anaconda.org, you can uncommit this URL to download it from the mirror supported by the CRA of SUStech
+    #URL='https://mirrors.sustech.edu.cn/anaconda/cloud/intel/$FILE_PREFIX-${platform}/'${DOWNLOAD_FILE} 
 else
     DOWNLOAD_FILE='${package}.tar.bz2'
     URL='$FTP_URL_PREFIX/$FILE_PREFIX-${platform}-'${DOWNLOAD_FILE}
@@ -52,7 +54,11 @@ do
     do
         echo "Installing $(eval echo $DOWNLOAD_FILE) for x86_${platform}..."
         echo "try download mkl package from: $(eval echo $URL)"
+        if [ ${FILE_PREFIX} == "win" ]; then
+            curl -SL "$(eval echo $URL)" --output mkl/x86_${platform}/"$(eval echo $DOWNLOAD_FILE)"
+        else
             wget -q --show-progress "$(eval echo $URL)" -O mkl/x86_${platform}/"$(eval echo $DOWNLOAD_FILE)"
+        fi
         $TAR xvj -C mkl/x86_${platform} -f mkl/x86_${platform}/"$(eval echo $DOWNLOAD_FILE)"
     done
 done