Remove backend-specific generation and also use CMake variant generation in Nix

.github/workflows/build_kernel.yaml

@@ -28,11 +28,17 @@ jobs:
           USER: runner
       - name: Nix info
         run: nix-shell -p nix-info --run "nix-info -m"
+
       - name: Build relu kernel
         run: ( cd builder/examples/relu && nix build .\#redistributable.torch29-cxx11-cu126-x86_64-linux )
       - name: Copy relu kernel
         run: cp -rL builder/examples/relu/result relu-kernel
 
+      - name: Build extra-data kernel
+        run: ( cd builder/examples/extra-data && nix build .\#redistributable.torch29-cxx11-cu126-x86_64-linux )
+      - name: Copy extra-data kernel
+        run: cp -rL builder/examples/extra-data/result extra-data
+
       - name: Build relu kernel (CPU)
         run: ( cd builder/examples/relu && nix build .\#redistributable.torch29-cxx11-cpu-x86_64-linux )
       - name: Copy relu kernel (CPU)
@@ -70,6 +76,7 @@ jobs:
           path: |
             activation-kernel
             cutlass-gemm-kernel
+            extra-data
             relu-kernel
             relu-kernel-cpu
             relu-backprop-compile-kernel

build2cmake/src/config/mod.rs

@@ -27,10 +27,6 @@ impl Build {
     pub fn is_noarch(&self) -> bool {
         self.kernels.is_empty()
     }
-
-    pub fn supports_backend(&self, backend: &Backend) -> bool {
-        self.general.backends.contains(backend)
-    }
 }
 
 pub struct General {
@@ -126,18 +122,18 @@ pub struct Torch {
 }
 
 impl Torch {
-    pub fn data_globs(&self) -> Option<Vec<String>> {
+    pub fn data_extensions(&self) -> Option<Vec<String>> {
         match self.pyext.as_ref() {
             Some(exts) => {
-                let globs = exts
+                let extensions = exts
                     .iter()
                     .filter(|&ext| ext != "py" && ext != "pyi")
-                    .map(|ext| format!("\"**/*.{ext}\""))
+                    .cloned()
                     .collect_vec();
-                if globs.is_empty() {
+                if extensions.is_empty() {
                     None
                 } else {
-                    Some(globs)
+                    Some(extensions)
                 }
             }
 
@@ -247,6 +243,18 @@ pub enum Backend {
     Xpu,
 }
 
+impl Backend {
+    pub const fn all() -> [Backend; 5] {
+        [
+            Backend::Cpu,
+            Backend::Cuda,
+            Backend::Metal,
+            Backend::Rocm,
+            Backend::Xpu,
+        ]
+    }
+}
+
 impl Display for Backend {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {

build2cmake/src/main.rs

@@ -12,7 +12,7 @@ mod torch;
 use torch::{write_torch_ext, write_torch_ext_noarch};
 
 mod config;
-use config::{v3, Backend, Build, BuildCompat};
+use config::{v3, Build, BuildCompat};
 
 mod fileset;
 use fileset::FileSet;
@@ -48,9 +48,6 @@ enum Commands {
         /// kernel name to avoid name collisions. (e.g. Git SHA)
         #[arg(long)]
         ops_id: Option<String>,
-
-        #[arg(long)]
-        backend: Option<Backend>,
     },
 
     /// Update a `build.toml` to the current format.
@@ -93,12 +90,11 @@ fn main() -> Result<()> {
     let args = Cli::parse();
     match args.command {
         Commands::GenerateTorch {
-            backend,
             build_toml,
             force,
             target_dir,
             ops_id,
-        } => generate_torch(backend, build_toml, target_dir, force, ops_id),
+        } => generate_torch(build_toml, target_dir, force, ops_id),
         Commands::UpdateBuild { build_toml } => update_build(build_toml),
         Commands::Validate { build_toml } => {
             parse_and_validate(build_toml)?;
@@ -115,7 +111,6 @@ fn main() -> Result<()> {
 }
 
 fn generate_torch(
-    backend: Option<Backend>,
     build_toml: PathBuf,
     target_dir: Option<PathBuf>,
     force: bool,
@@ -139,41 +134,10 @@ fn generate_torch(
     env.set_trim_blocks(true);
     minijinja_embed::load_templates!(&mut env);
 
-    let backend = match backend {
-        Some(backend) => {
-            if !build.supports_backend(&backend) {
-                bail!("Kernel does not support backend: {}", backend);
-            }
-
-            backend
-        }
-        None => {
-            let kernel_backends = &build.general.backends;
-
-            if kernel_backends.len() > 1 {
-                let mut kernel_backends = kernel_backends
-                    .iter()
-                    .map(ToString::to_string)
-                    .collect::<Vec<_>>();
-                kernel_backends.sort();
-                bail!(
-                    "Multiple supported backends found in build.toml: {}. Please specify one with --backend.",
-                    kernel_backends.join(", ")
-                );
-            }
-
-            if let Some(backend) = kernel_backends.first() {
-                *backend
-            } else {
-                bail!("No backends are specified in build.toml");
-            }
-        }
-    };
-
     let file_set = if build.is_noarch() {
-        write_torch_ext_noarch(&env, backend, &build, target_dir.clone(), ops_id)?
+        write_torch_ext_noarch(&env, &build, target_dir.clone(), ops_id)?
     } else {
-        write_torch_ext(&env, backend, &build, target_dir.clone(), ops_id)?
+        write_torch_ext(&env, &build, target_dir.clone(), ops_id)?
     };
     file_set.write(&target_dir, force)?;
 
@@ -368,14 +332,12 @@ fn get_generated_files(
 ) -> Result<Vec<PathBuf>> {
     let mut all_set = FileSet::new();
 
-    for backend in &build.general.backends {
-        let set = if build.is_noarch() {
-            write_torch_ext_noarch(env, *backend, build, target_dir.clone(), ops_id.clone())?
-        } else {
-            write_torch_ext(env, *backend, build, target_dir.clone(), ops_id.clone())?
-        };
-        all_set.extend(set);
-    }
+    let set = if build.is_noarch() {
+        write_torch_ext_noarch(env, build, target_dir.clone(), ops_id.clone())?
+    } else {
+        write_torch_ext(env, build, target_dir.clone(), ops_id.clone())?
+    };
+    all_set.extend(set);
 
     Ok(all_set.into_names())
 }