coreylowman · coreylowman · May 1, 2023 · May 1, 2023 · May 1, 2023
diff --git a/build.rs b/build.rs
@@ -166,6 +166,12 @@ mod cuda {
                 compute_cap = max_nvcc_code;
             }
 
+            println!("cargo:rerun-if-env-changed=CUDA_COMPUTE_CAP");
+            if let Ok(compute_cap_str) = std::env::var("CUDA_COMPUTE_CAP") {
+                compute_cap = compute_cap_str.parse::<usize>().unwrap();
+                println!("cargo:warning=Using gpu arch {compute_cap} from $CUDA_COMPUTE_CAP");
+            }
+
             println!("cargo:rustc-env=CUDA_COMPUTE_CAP=sm_{compute_cap}");
 
             kernel_paths

diff --git a/src/tensor_ops/utilities/compatibility.cuh b/src/tensor_ops/utilities/compatibility.cuh
@@ -5,7 +5,7 @@
 
 // FIXME: the minimum compute capabilities are just guesses since the table is not specific enough
 
-#if __CUDA_ARCH__ < 600
+#if __CUDA_ARCH__ < 800
 __device__ __forceinline__ __half __hmax(__half a, __half b) {
     return __float2half(fmaxf(__half2float(a), __half2float(b)));
 }