Factor out write_torch_registration_macros to common mod

Author: danieldk

build2cmake/src/torch/common.rs

@@ -1,3 +1,5 @@
+use std::path::PathBuf;
+
 use eyre::{Context, Result};
 use itertools::Itertools;
 use minijinja::{context, Environment};
@@ -6,6 +8,8 @@ use crate::config::{Backend, General};
 use crate::metadata::Metadata;
 use crate::FileSet;
 
+static REGISTRATION_H: &str = include_str!("../templates/registration.h");
+
 pub fn write_pyproject_toml(
     env: &Environment,
     backend: Backend,
@@ -64,3 +68,14 @@ where
         .collect_vec()
         .join(";")
 }
+
+pub fn write_torch_registration_macros(file_set: &mut FileSet) -> Result<()> {
+    let mut path = PathBuf::new();
+    path.push("torch-ext");
+    path.push("registration.h");
+    file_set
+        .entry(path)
+        .extend_from_slice(REGISTRATION_H.as_bytes());
+
+    Ok(())
+}

build2cmake/src/torch/cpu.rs

@@ -8,13 +8,13 @@ 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::common::write_torch_registration_macros;
 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");
-static REGISTRATION_H: &str = include_str!("../templates/registration.h");
 
 pub fn write_torch_ext_cpu(
     env: &Environment,
@@ -224,17 +224,6 @@ fn write_setup_py(
     Ok(())
 }
 
-fn write_torch_registration_macros(file_set: &mut FileSet) -> Result<()> {
-    let mut path = PathBuf::new();
-    path.push("torch-ext");
-    path.push("registration.h");
-    file_set
-        .entry(path)
-        .extend_from_slice(REGISTRATION_H.as_bytes());
-
-    Ok(())
-}
-
 fn prefix_and_join_includes<S>(includes: impl AsRef<[S]>) -> String
 where
     S: AsRef<str>,

build2cmake/src/torch/cuda.rs

@@ -10,6 +10,7 @@ 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::common::write_torch_registration_macros;
 use crate::torch::kernel::render_kernel_components;
 use crate::torch::kernel_ops_identifier;
 use crate::version::Version;
@@ -18,7 +19,6 @@ use crate::FileSet;
 static CMAKE_UTILS: &str = include_str!("../templates/utils.cmake");
 static CMAKE_KERNEL: &str = include_str!("../templates/kernel.cmake");
 static WINDOWS_UTILS: &str = include_str!("../templates/windows.cmake");
-static REGISTRATION_H: &str = include_str!("../templates/registration.h");
 static HIPIFY: &str = include_str!("../templates/cuda/hipify.py");
 
 pub fn write_torch_ext_cuda(
@@ -66,17 +66,6 @@ pub fn write_torch_ext_cuda(
     Ok(file_set)
 }
 
-fn write_torch_registration_macros(file_set: &mut FileSet) -> Result<()> {
-    let mut path = PathBuf::new();
-    path.push("torch-ext");
-    path.push("registration.h");
-    file_set
-        .entry(path)
-        .extend_from_slice(REGISTRATION_H.as_bytes());
-
-    Ok(())
-}
-
 fn write_setup_py(
     env: &Environment,
     torch: &Torch,

build2cmake/src/torch/metal.rs

@@ -8,13 +8,13 @@ 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::common::write_torch_registration_macros;
 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");
-static REGISTRATION_H: &str = include_str!("../templates/registration.h");
 static COMPILE_METAL_CMAKE: &str = include_str!("../templates/metal/compile-metal.cmake");
 static METALLIB_TO_HEADER_PY: &str = include_str!("../templates/metal/metallib_to_header.py");
 
@@ -240,17 +240,6 @@ fn write_setup_py(
     Ok(())
 }
 
-fn write_torch_registration_macros(file_set: &mut FileSet) -> Result<()> {
-    let mut path = PathBuf::new();
-    path.push("torch-ext");
-    path.push("registration.h");
-    file_set
-        .entry(path)
-        .extend_from_slice(REGISTRATION_H.as_bytes());
-
-    Ok(())
-}
-
 fn prefix_and_join_includes<S>(includes: impl AsRef<[S]>) -> String
 where
     S: AsRef<str>,

build2cmake/src/torch/xpu.rs

@@ -9,14 +9,14 @@ use minijinja::{context, Environment};
 use crate::config::{Backend, Build, Dependency, Torch};
 use crate::torch::common::write_metadata;
 use crate::torch::common::write_pyproject_toml;
+use crate::torch::common::write_torch_registration_macros;
 use crate::torch::kernel::render_kernel_components;
 use crate::torch::kernel_ops_identifier;
 use crate::version::Version;
 use crate::FileSet;
 
 static CMAKE_UTILS: &str = include_str!("../templates/utils.cmake");
 static CMAKE_KERNEL: &str = include_str!("../templates/kernel.cmake");
-static REGISTRATION_H: &str = include_str!("../templates/registration.h");
 static WINDOWS_UTILS: &str = include_str!("../templates/windows.cmake");
 
 pub fn write_torch_ext_xpu(
@@ -62,17 +62,6 @@ pub fn write_torch_ext_xpu(
     Ok(file_set)
 }
 
-fn write_torch_registration_macros(file_set: &mut FileSet) -> Result<()> {
-    let mut path = PathBuf::new();
-    path.push("torch-ext");
-    path.push("registration.h");
-    file_set
-        .entry(path)
-        .extend_from_slice(REGISTRATION_H.as_bytes());
-
-    Ok(())
-}
-
 fn write_setup_py(
     env: &Environment,
     torch: &Torch,