Evaluate environmnent variable

Author: mc-nv

tools/gen_ort_dockerfile.py

@@ -322,9 +322,31 @@ def dockerfile_for_linux(output_file):
         ep_flags += (
             " --skip_tests --cmake_extra_defines 'onnxruntime_BUILD_UNIT_TESTS=OFF'"
         )
-        cuda_archs = "87;101"
+        if os.getenv("CUDA_ARCH_LIST") is not None:
+            print(f"[INFO] Defined CUDA_ARCH_LIST: {os.getenv('CUDA_ARCH_LIST')}")
+            cuda_archs = (
+                os.getenv("CUDA_ARCH_LIST")
+                .replace("PTX", "")
+                .replace(" ", "-real;")
+                .replace(".", "")
+            )
+            cuda_archs = re.sub(r"-real;$", "", cuda_archs)
+            print(f"[INFO] Set ONNX Runtime to use CUDA architectures to: {cuda_archs}")
+        else:
+            cuda_archs = "87"
     else:
-        cuda_archs = "75;80;86;89;90;100;120"
+        if os.getenv("CUDA_ARCH_LIST") is not None:
+            print(f"[INFO] Defined CUDA_ARCH_LIST: {os.getenv('CUDA_ARCH_LIST')}")
+            cuda_archs = (
+                os.getenv("CUDA_ARCH_LIST")
+                .replace("PTX", "")
+                .replace(" ", "-real;")
+                .replace(".", "")
+            )
+            cuda_archs = re.sub(r"-real;$", "", cuda_archs)
+            print(f"[INFO] Set ONNX Runtime to use CUDA architectures to: {cuda_archs}")
+        else:
+            cuda_archs = "75;80;86;90;100;120"
 
     df += """
 WORKDIR /workspace/onnxruntime