Add python build option for cuda-archs (#672)

* Add python build option for cuda-archs

* ignore the nvidia-smi error

* apply multi-archs more places

* fix the semicolon issue

* Update ci_optional.yml

* Update ci.yml

Author: wenbingl

.pipelines/ci.yml

@@ -488,7 +488,7 @@ stages:
 - stage: WindowsCUDABuilds
   dependsOn: []
   jobs:
-  - job: WindowsC
+  - job: WindowsCUDABoth
     pool:
       name: 'onnxruntime-extensions-Win2022-GPU-A10'
     variables:
@@ -526,7 +526,8 @@ stages:
 
       - script: |
           set CUDA_PATH=$(Agent.TempDirectory)\v11.8
-          call .\build.bat -T cuda="%CUDA_PATH%" -DOCOS_ENABLE_CTEST=ON -DOCOS_USE_CUDA=ON -DCMAKE_CUDA_ARCHITECTURES=86 -DOCOS_ONNXRUNTIME_VERSION="$(ORT_VERSION)" -DONNXRUNTIME_PKG_DIR=.\onnxruntime-win-x64-gpu-$(ORT_VERSION)
+          call .\build.bat -T cuda="%CUDA_PATH%" -DOCOS_ENABLE_CTEST=ON -DOCOS_USE_CUDA=ON -DCMAKE_CUDA_ARCHITECTURES=70;86^
+            -DOCOS_ONNXRUNTIME_VERSION="$(ORT_VERSION)" -DONNXRUNTIME_PKG_DIR=.\onnxruntime-win-x64-gpu-$(ORT_VERSION)
         displayName: build the customop library with onnxruntime
 
       - script: |
@@ -547,7 +548,7 @@ stages:
           python -m pip install --upgrade setuptools pip
           python -m pip install numpy coloredlogs flatbuffers packaging protobuf sympy
           python -m pip install onnxruntime-gpu==$(ORT_VERSION)
-          python -m pip install -v --config-settings "ortx-user-option=use-cuda" .
+          python -m pip install -v --config-settings "ortx-user-option=use-cuda,cuda_archs=70;86" .
         displayName: Build and install onnxruntime-extensions CUDA package.
 
       - script: |
@@ -665,7 +666,7 @@ stages:
                 python3 -m pip install --upgrade pip; \
                 python3 -m pip install --upgrade setuptools; \
                 python3 -m pip install onnxruntime-gpu==$(ORT_VERSION); \
-                python3 -m pip install -v --config-settings "ortx-user-option=use-cuda" . ; \
+                python3 -m pip install -v --config-settings 'ortx-user-option=use-cuda,cuda_archs=70;86' . ; \
                 python3 -m pip install $(TORCH_VERSION) ; \
                 python3 -m pip install -r requirements-dev.txt; \
                 cd test && python -m pytest . --verbose; \
@@ -864,4 +865,4 @@ stages:
           -scheme OrtExtensionsUsage \
           -destination "platform=iOS Simulator,id=${ORT_EXTENSIONS_SIMULATOR_DEVICE_ID}" \
           test CODE_SIGNING_ALLOWED=NO
-      displayName: "Build and test OrtExtensionsUsage"
+      displayName: "Build and test OrtExtensionsUsage"

.pipelines/ci_optional.yml

@@ -181,7 +181,7 @@ stages:
           python -m pip install --upgrade setuptools pip
           python -m pip install numpy coloredlogs flatbuffers packaging protobuf sympy
           python -m pip install -U --index-url https://pkgs.dev.azure.com/aiinfra/PublicPackages/_packaging/ORT-Nightly/pypi/simple/ ort-nightly-gpu
-          python -m pip install -v --config-settings "ortx-user-option=use-cuda" .
+          python -m pip install -v --config-settings "ortx-user-option=use-cuda,cuda_archs=70;86" .
         displayName: Build and install onnxruntime-extensions CUDA package.
 
       - script: |

.pyproject/cmdclass.py

@@ -30,7 +30,20 @@ def _load_cuda_version():
         match = re.search(pattern, output)
         if match:
             return match.group(1)
-    except subprocess.CalledProcessError:
+    except (subprocess.CalledProcessError, OSError):
+        pass
+
+    return None
+
+
+def _load_nvidia_smi():
+    try:
+        output = subprocess.check_output(
+            ["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader,nounits"],
+            stderr=subprocess.STDOUT).decode("utf-8")
+        arch = output.strip().replace('.', '')
+        return arch if arch.isdigit() else None
+    except (subprocess.CalledProcessError, OSError):
         pass
 
     return None
@@ -133,6 +146,7 @@ def initialize_options(self):
         self.no_azure = None
         self.no_opencv = None
         self.cc_debug = None
+        self.cuda_archs = None
 
     def _parse_options(self, options):
         for segment in options.split(','):
@@ -211,6 +225,14 @@ def build_cmake(self, extension):
                 with f_ver.open('a') as _f:
                     _f.writelines(["\n", f"cuda = \"{cuda_ver}\"", "\n"])
 
+                if self.cuda_archs is not None:
+                    cmake_args += ['-DCMAKE_CUDA_ARCHITECTURES=' + self.cuda_archs]
+                else:
+                    smi = _load_nvidia_smi()
+                    if not smi:
+                        raise RuntimeError(f"Cannot detect the CUDA archs from your machine, please specify it by yourself.")
+                    cmake_args += ['-DCMAKE_CUDA_ARCHITECTURES=' + smi]
+
         # CMake lets you override the generator - we need to check this.
         # Can be set with Conda-Build, for example.
         cmake_generator = os.environ.get("CMAKE_GENERATOR", "")