build2cmake: always generate kernel components for all backends (#245)

* build2cmake: always generate kernel components for all backends

This change always generates the kernel components for all backends.
This is a necessary step for merging the preambles later.

* Fix kernel render visibility

Co-authored-by: drbh <david.richard.holtz@gmail.com>

---------

Co-authored-by: drbh <david.richard.holtz@gmail.com>

Author: danieldk

build2cmake/src/templates/cpu/kernel.cmake

@@ -1,5 +1,7 @@
-cpu_kernel_component(SRC
-  SOURCES {{ sources }}
-  {% if includes %}INCLUDES "{{ includes }}"{% endif %}
-  {% if cxx_flags %}CXX_FLAGS "{{ cxx_flags }}"{% endif %}
-)
+if(GPU_LANG STREQUAL "CPU")
+  cpu_kernel_component(SRC
+    SOURCES {{ sources }}
+    {% if includes %}INCLUDES "{{ includes }}"{% endif %}
+    {% if cxx_flags %}CXX_FLAGS "{{ cxx_flags }}"{% endif %}
+  )
+endif()

build2cmake/src/templates/cpu/preamble.cmake

@@ -42,6 +42,8 @@ if (TORCH_VERSION VERSION_GREATER {{ torch_maxver }})
 endif()
 {% endif %}
 
+set(GPU_LANG "CPU")
+
 add_compile_definitions(CPU_KERNEL)
 
 # Initialize SRC list for kernel and binding sources

build2cmake/src/templates/metal/kernel.cmake

@@ -1,5 +1,7 @@
-metal_kernel_component(SRC
-  SOURCES {{ sources }}
-  {% if includes %}INCLUDES "{{ includes }}"{% endif %}
-  {% if cxx_flags %}CXX_FLAGS "{{ cxx_flags }}"{% endif %}
-)
+if(GPU_LANG STREQUAL "METAL")
+  metal_kernel_component(SRC
+    SOURCES {{ sources }}
+    {% if includes %}INCLUDES "{{ includes }}"{% endif %}
+    {% if cxx_flags %}CXX_FLAGS "{{ cxx_flags }}"{% endif %}
+  )
+endif()

build2cmake/src/templates/metal/preamble.cmake

@@ -42,6 +42,8 @@ if (TORCH_VERSION VERSION_GREATER {{ torch_maxver }})
 endif()
 {% endif %}
 
+set(GPU_LANG "METAL")
+
 add_compile_definitions(METAL_KERNEL)
 
 # Initialize SRC list for kernel and binding sources

build2cmake/src/templates/xpu/kernel.cmake

@@ -1,6 +1,8 @@
-xpu_kernel_component(SRC
-  SOURCES {{ sources }}
-  {% if includes %}INCLUDES "{{ includes }}"{% endif %}
-  {% if cxx_flags %}CXX_FLAGS "{{ cxx_flags }}"{% endif %}
-  {% if sycl_flags %}SYCL_FLAGS "{{ sycl_flags }}"{% endif %}
-)
+if(GPU_LANG STREQUAL "SYCL")
+  xpu_kernel_component(SRC
+    SOURCES {{ sources }}
+    {% if includes %}INCLUDES "{{ includes }}"{% endif %}
+    {% if cxx_flags %}CXX_FLAGS "{{ cxx_flags }}"{% endif %}
+    {% if sycl_flags %}SYCL_FLAGS "{{ sycl_flags }}"{% endif %}
+  )
+endif()

build2cmake/src/torch/common.rs

@@ -52,3 +52,15 @@ pub fn write_metadata(backend: Backend, general: &General, file_set: &mut FileSe
 
     Ok(())
 }
+
+pub fn prefix_and_join_includes<S>(includes: impl AsRef<[S]>) -> String
+where
+    S: AsRef<str>,
+{
+    includes
+        .as_ref()
+        .iter()
+        .map(|include| format!("${{CMAKE_SOURCE_DIR}}/{}", include.as_ref()))
+        .collect_vec()
+        .join(";")
+}

build2cmake/src/torch/cpu.rs

@@ -4,13 +4,13 @@ use eyre::{bail, Context, Result};
 use itertools::Itertools;
 use minijinja::{context, Environment};
 
-use super::{common::write_pyproject_toml, kernel_ops_identifier};
-use crate::{
-    config::{Backend, Build, Kernel, Torch},
-    fileset::FileSet,
-    torch::common::write_metadata,
-    version::Version,
-};
+use crate::config::{Backend, Build, Torch};
+use crate::fileset::FileSet;
+use crate::torch::common::write_metadata;
+use crate::torch::common::write_pyproject_toml;
+use crate::torch::kernel::render_kernel_components;
+use crate::torch::kernel_ops_identifier;
+use crate::version::Version;
 
 static CMAKE_UTILS: &str = include_str!("../templates/utils.cmake");
 static CMAKE_KERNEL: &str = include_str!("../templates/kernel.cmake");
@@ -96,13 +96,7 @@ fn write_cmake(
 
     render_binding(env, torch, name, cmake_writer)?;
 
-    for (kernel_name, kernel) in build
-        .kernels
-        .iter()
-        .filter(|(_, kernel)| matches!(kernel, Kernel::Cpu { .. }))
-    {
-        render_kernel(env, kernel_name, kernel, cmake_writer)?;
-    }
+    render_kernel_components(env, build, cmake_writer)?;
 
     render_extension(env, name, ops_name, cmake_writer)?;
 
@@ -154,38 +148,6 @@ pub fn render_extension(
     Ok(())
 }
 
-pub fn render_kernel(
-    env: &Environment,
-    kernel_name: &str,
-    kernel: &Kernel,
-    write: &mut impl Write,
-) -> Result<()> {
-    // Easier to do in Rust than Jinja.
-    let sources = kernel
-        .src()
-        .iter()
-        .map(|src| format!("\"{src}\""))
-        .collect_vec()
-        .join("\n");
-
-    env.get_template("cpu/kernel.cmake")
-        .wrap_err("Cannot get kernel template")?
-        .render_to_write(
-            context! {
-                cxx_flags => kernel.cxx_flags().map(|flags| flags.join(";")),
-                includes => kernel.include().map(prefix_and_join_includes),
-                kernel_name => kernel_name,
-                sources => sources,
-            },
-            &mut *write,
-        )
-        .wrap_err("Cannot render kernel template")?;
-
-    write.write_all(b"\n")?;
-
-    Ok(())
-}
-
 fn render_preamble(
     env: &Environment,
     name: &str,

build2cmake/src/torch/cuda.rs

@@ -4,13 +4,14 @@ use std::io::Write;
 use std::path::PathBuf;
 
 use eyre::{bail, Context, Result};
-use itertools::Itertools;
 use minijinja::{context, Environment};
 
-use super::common::write_pyproject_toml;
-use super::kernel_ops_identifier;
-use crate::config::{Backend, Build, Dependency, Kernel, Torch};
+use crate::config::{Backend, Build, Dependency, Torch};
+use crate::torch::common::prefix_and_join_includes;
 use crate::torch::common::write_metadata;
+use crate::torch::common::write_pyproject_toml;
+use crate::torch::kernel::render_kernel_components;
+use crate::torch::kernel_ops_identifier;
 use crate::version::Version;
 use crate::FileSet;
 
@@ -181,13 +182,7 @@ fn write_cmake(
 
     render_binding(env, torch, name, cmake_writer)?;
 
-    for (kernel_name, kernel) in build
-        .kernels
-        .iter()
-        .filter(|(_, kernel)| kernel.backend() == backend)
-    {
-        render_kernel(env, kernel_name, kernel, cmake_writer)?;
-    }
+    render_kernel_components(env, build, cmake_writer)?;
 
     render_extension(env, name, ops_name, cmake_writer)?;
 
@@ -312,71 +307,6 @@ fn render_deps(
     Ok(())
 }
 
-pub fn render_kernel(
-    env: &Environment,
-    kernel_name: &str,
-    kernel: &Kernel,
-    write: &mut impl Write,
-) -> Result<()> {
-    // Easier to do in Rust than Jinja.
-    let sources = kernel
-        .src()
-        .iter()
-        .map(|src| format!("\"{src}\""))
-        .collect_vec()
-        .join("\n");
-
-    let (cuda_capabilities, rocm_archs, cuda_flags, hip_flags, cuda_minver) = match kernel {
-        Kernel::Cuda {
-            cuda_capabilities,
-            cuda_flags,
-            cuda_minver,
-            ..
-        } => (
-            cuda_capabilities.as_deref(),
-            None,
-            cuda_flags.as_deref(),
-            None,
-            cuda_minver.as_ref(),
-        ),
-        Kernel::Rocm {
-            rocm_archs,
-            hip_flags,
-            ..
-        } => (
-            None,
-            rocm_archs.as_deref(),
-            None,
-            hip_flags.as_deref(),
-            None,
-        ),
-        _ => unreachable!("Unsupported kernel type for CUDA rendering"),
-    };
-
-    env.get_template("cuda/kernel.cmake")
-        .wrap_err("Cannot get kernel template")?
-        .render_to_write(
-            context! {
-                cuda_capabilities => cuda_capabilities,
-                cuda_flags => cuda_flags.map(|flags| flags.join(";")),
-                cuda_minver => cuda_minver.map(ToString::to_string),
-                cxx_flags => kernel.cxx_flags().map(|flags| flags.join(";")),
-                rocm_archs => rocm_archs,
-                hip_flags => hip_flags.map(|flags| flags.join(";")),
-                includes => kernel.include().map(prefix_and_join_includes),
-                kernel_name => kernel_name,
-                supports_hipify => matches!(kernel, Kernel::Rocm{ .. }),
-                sources => sources,
-            },
-            &mut *write,
-        )
-        .wrap_err("Cannot render kernel template")?;
-
-    write.write_all(b"\n")?;
-
-    Ok(())
-}
-
 pub fn render_extension(
     env: &Environment,
     name: &str,
@@ -428,15 +358,3 @@ pub fn render_preamble(
 
     Ok(())
 }
-
-fn prefix_and_join_includes<S>(includes: impl AsRef<[S]>) -> String
-where
-    S: AsRef<str>,
-{
-    includes
-        .as_ref()
-        .iter()
-        .map(|include| format!("${{CMAKE_SOURCE_DIR}}/{}", include.as_ref()))
-        .collect_vec()
-        .join(";")
-}