add support for HYPERBEAM_HIP_ARCH

Author: d3v-null

build.rs

@@ -220,13 +220,26 @@ mod gpu {
 
         #[cfg(feature = "hip")]
         let mut gpu_target = {
+            const DEFAULT_HIP_ARCHES: &[&str] = &["gfx90a"];
+
             let hip_path = hip_sys::hiprt::get_hip_path();
+            if !hip_path.exists() {
+                panic!("Couldn't find HIP path at {}", hip_path.display());
+            }
+
             // It seems that various ROCm releases change where hipcc is...
             let mut compiler = hip_path.join("bin/hipcc");
             if !compiler.exists() {
                 // Try the dir above, which might be the ROCm dir.
                 compiler = hip_path.join("../bin/hipcc");
             }
+            if !compiler.exists() {
+                panic!(
+                    "Couldn't find hipcc in {}/bin/hipcc or {}",
+                    hip_path.display(),
+                    compiler.display()
+                );
+            }
             let mut hip_target = cc::Build::new();
             hip_target
                 .compiler(compiler)
@@ -235,6 +248,34 @@ mod gpu {
                 .file("src/fee/gpu/fee.cu")
                 .file("src/analytic/gpu/analytic.cu");
 
+            hip_target.flag("-O0"); // <- hip can't handle optimizations
+
+            println!("cargo:rerun-if-env-changed=HYPERBEAM_HIP_ARCH");
+            println!("cargo:rerun-if-env-changed=HYPERDRIVE_HIP_ARCH");
+            let arches: Vec<String> = match (
+                env::var("HYPERBEAM_HIP_ARCH"),
+                env::var("HYPERDRIVE_HIP_ARCH"),
+            ) {
+                // When a user-supplied variable exists, use it as the CUDA arch and
+                // compute level.
+                (Ok(c), _) | (Err(_), Ok(c)) => {
+                    vec![c]
+                }
+                _ => {
+                    // Print out all of the default arches and computes as a
+                    // warning.
+                    println!("cargo:warning=No HYPERBEAM_HIP_ARCH; Passing --offload-arch={DEFAULT_HIP_ARCHES:?} to hip");
+                    DEFAULT_HIP_ARCHES
+                        .iter()
+                        .map(|&s| String::from(s))
+                        .collect()
+                }
+            };
+
+            for arch in arches {
+                hip_target.flag(&format!("--offload-arch={arch}"));
+            }
+
             hip_target
         };