Include arch by default, disable with INCLUDE_ARCH=0

14d4f8cb · Casper · b621c142 · 14d4f8cb
Commit 14d4f8cb authored Sep 21, 2023 by Casper
Hide whitespace changes
Inline Side-by-side

Showing with 6 additions and 1 deletion

setup.py setup.py +6 -1

No files found.
--- a/setup.py
+++ b/setup.py
@@ -96,8 +96,13 @@ generator_flags = get_generator_flag()
 arch_flags = get_compute_capabilities()

 if os.name == "nt":
+    include_arch = os.getenv("INCLUDE_ARCH", "1") == "1"
+
    # Relaxed args on Windows
-    extra_compile_args={}
+    if include_arch:
+        extra_compile_args={"nvcc": arch_flags}
+    else:
+        extra_compile_args={}
 else:
    extra_compile_args={
        "cxx": ["-g", "-O3", "-fopenmp", "-lgomp", "-std=c++17", "-DENABLE_BF16"],