Add basic support for building Metal 🤘 kernels (#146)

This change adds support for building Metal kernels (`backend = metal`)
for Apple Silicon Macs.

Author: danieldk

build2cmake/src/config/v1.rs

@@ -51,13 +51,15 @@ pub enum Language {
     #[default]
     Cuda,
     CudaHipify,
+    Metal,
 }
 
 impl Display for Language {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             Language::Cuda => f.write_str("cuda"),
             Language::CudaHipify => f.write_str("cuda-hipify"),
+            Language::Metal => f.write_str("metal"),
         }
     }
 }

build2cmake/src/config/v2.rs

@@ -87,13 +87,15 @@ pub struct Kernel {
 #[serde(deny_unknown_fields, rename_all = "kebab-case")]
 pub enum Backend {
     Cuda,
+    Metal,
     Rocm,
 }
 
 impl Display for Backend {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             Backend::Cuda => write!(f, "cuda"),
+            Backend::Metal => write!(f, "metal"),
             Backend::Rocm => write!(f, "rocm"),
         }
     }
@@ -105,6 +107,7 @@ impl FromStr for Backend {
     fn from_str(s: &str) -> Result<Self, Self::Err> {
         match s.to_lowercase().as_str() {
             "cuda" => Ok(Backend::Cuda),
+            "metal" => Ok(Backend::Metal),
             "rocm" => Ok(Backend::Rocm),
             _ => Err(format!("Unknown backend: {}", s)),
         }

build2cmake/src/main.rs

@@ -9,16 +9,13 @@ use eyre::{bail, ensure, Context, Result};
 use minijinja::Environment;
 
 mod torch;
-use torch::write_torch_ext;
-
-mod torch_universal;
+use torch::{write_torch_ext, write_torch_ext_metal, write_torch_universal_ext};
 
 mod config;
 use config::{Backend, Build, BuildCompat};
 
 mod fileset;
 use fileset::FileSet;
-use torch_universal::write_torch_universal_ext;
 
 #[derive(Parser, Debug)]
 #[command(version, about, long_about = None)]
@@ -108,16 +105,15 @@ fn generate_torch(
     env.set_trim_blocks(true);
     minijinja_embed::load_templates!(&mut env);
 
-    match (backend, build.general.universal) {
-        (None, true) => write_torch_universal_ext(&env, &build, target_dir, force, ops_id)?,
+    let backend = match (backend, build.general.universal) {
+        (None, true) => return write_torch_universal_ext(&env, &build, target_dir, force, ops_id),
         (Some(backend), true) => bail!("Universal kernel, cannot generate for backend {}", backend),
-        // TODO: add check if that type of backend has at least one kernel.
         (Some(backend), false) => {
             if !build.has_kernel_with_backend(&backend) {
                 bail!("No kernels found for backend {}", backend);
             }
 
-            write_torch_ext(&env, &build, target_dir, force, ops_id)?
+            backend
         }
         (None, false) => {
             let mut kernel_backends = build.backends();
@@ -139,15 +135,16 @@ fn generate_torch(
                 );
             }
 
-            match backend {
-                Backend::Cuda | Backend::Rocm => {
-                    write_torch_ext(&env, &build, target_dir, force, ops_id)?
-                }
-            }
+            backend
         }
-    }
+    };
 
-    Ok(())
+    match backend {
+        Backend::Cuda | Backend::Rocm => {
+            write_torch_ext(&env, backend, &build, target_dir, force, ops_id)
+        }
+        Backend::Metal => write_torch_ext_metal(&env, &build, target_dir, force, ops_id),
+    }
 }
 
 fn update_build(build_toml: PathBuf) -> Result<()> {

…ld2cmake/src/templates/dep-cutlass.cmake …ake/src/templates/cuda/dep-cutlass.cmakebuild2cmake/src/templates/dep-cutlass.cmake renamed to build2cmake/src/templates/cuda/dep-cutlass.cmake build2cmake/src/templates/dep-cutlass.cmake renamed to build2cmake/src/templates/cuda/dep-cutlass.cmake

@@ -60,12 +60,16 @@ if(GPU_LANG STREQUAL "CUDA")
   if(NVCC_THREADS AND GPU_LANG STREQUAL "CUDA")
     list(APPEND GPU_FLAGS "--threads=${NVCC_THREADS}")
   endif()
+
+  add_compile_definitions(CUDA_KERNEL)
 elseif(GPU_LANG STREQUAL "HIP")
   set(ROCM_ARCHS "${HIP_SUPPORTED_ARCHS}")
   # TODO: remove this once we can set specific archs per source file set.
   override_gpu_arches(GPU_ARCHES
     ${GPU_LANG}
     "${${GPU_LANG}_SUPPORTED_ARCHS}")
+
+  add_compile_definitions(ROCM_KERNEL)
 else()
   override_gpu_arches(GPU_ARCHES
     ${GPU_LANG}