build: Make build.rs candle-kernels compatible with Nix and sandboxed builds

candle-examples/build.rs

@@ -1,7 +1,8 @@
 #![allow(unused)]
 use anyhow::{Context, Result};
+use std::env;
 use std::io::Write;
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 
 struct KernelDirectories {
     kernel_glob: &'static str,
@@ -20,11 +21,21 @@ fn main() -> Result<()> {
 
     #[cfg(feature = "cuda")]
     {
+        // Added: Get the safe output directory from the environment.
+        let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap());
+
         for kdir in KERNEL_DIRS.iter() {
             let builder = bindgen_cuda::Builder::default().kernel_paths_glob(kdir.kernel_glob);
             println!("cargo:info={builder:?}");
             let bindings = builder.build_ptx().unwrap();
-            bindings.write(kdir.rust_target).unwrap()
+
+            // Changed: This now writes to a safe path inside $OUT_DIR.
+            let safe_target = out_dir.join(
+                Path::new(kdir.rust_target)
+                    .file_name()
+                    .context("Failed to get filename from rust_target")?,
+            );
+            bindings.write(safe_target).unwrap()
         }
     }
     Ok(())

candle-examples/examples/custom-ops/main.rs

@@ -8,7 +8,9 @@ extern crate intel_mkl_src;
 
 #[rustfmt::skip]
 #[cfg(feature = "cuda")]
-mod cuda_kernels;
+mod cuda_kernels {
+    include!(concat!(env!("OUT_DIR"), "/cuda_kernels.rs"));
+}
 
 use clap::Parser;
 

candle-kernels/build.rs

@@ -4,8 +4,10 @@ fn main() {
     println!("cargo:rerun-if-changed=src/cuda_utils.cuh");
     println!("cargo:rerun-if-changed=src/binary_op_macros.cuh");
 
+    let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap());
+    let ptx_path = out_dir.join("ptx.rs");
     let builder = bindgen_cuda::Builder::default();
     println!("cargo:info={builder:?}");
     let bindings = builder.build_ptx().unwrap();
-    bindings.write("src/ptx.rs").unwrap();
+    bindings.write(ptx_path).unwrap();
 }

candle-kernels/src/lib.rs

@@ -1,4 +1,6 @@
-mod ptx;
+mod ptx {
+    include!(concat!(env!("OUT_DIR"), "/ptx.rs"));
+}
 
 #[repr(u32)]
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]