Move write_ops_py to common mod

Author: danieldk

build2cmake/src/torch/common.rs

@@ -79,3 +79,28 @@ pub fn write_torch_registration_macros(file_set: &mut FileSet) -> Result<()> {
 
     Ok(())
 }
+
+pub fn write_ops_py(
+    env: &Environment,
+    name: &str,
+    ops_name: &str,
+    file_set: &mut FileSet,
+) -> Result<()> {
+    let mut path = PathBuf::new();
+    path.push("torch-ext");
+    path.push(name);
+    path.push("_ops.py");
+    let writer = file_set.entry(path);
+
+    env.get_template("_ops.py")
+        .wrap_err("Cannot get _ops.py template")?
+        .render_to_write(
+            context! {
+                ops_name => ops_name,
+            },
+            writer,
+        )
+        .wrap_err("Cannot render kernel template")?;
+
+    Ok(())
+}

build2cmake/src/torch/cpu.rs

@@ -7,6 +7,7 @@ use minijinja::{context, Environment};
 use crate::config::{Backend, Build, Torch};
 use crate::fileset::FileSet;
 use crate::torch::common::write_metadata;
+use crate::torch::common::write_ops_py;
 use crate::torch::common::write_pyproject_toml;
 use crate::torch::common::write_torch_registration_macros;
 use crate::torch::kernel::render_kernel_components;
@@ -172,31 +173,6 @@ fn render_preamble(
     Ok(())
 }
 
-fn write_ops_py(
-    env: &Environment,
-    name: &str,
-    ops_name: &str,
-    file_set: &mut FileSet,
-) -> Result<()> {
-    let mut path = PathBuf::new();
-    path.push("torch-ext");
-    path.push(name);
-    path.push("_ops.py");
-    let writer = file_set.entry(path);
-
-    env.get_template("_ops.py")
-        .wrap_err("Cannot get _ops.py template")?
-        .render_to_write(
-            context! {
-                ops_name => ops_name,
-            },
-            writer,
-        )
-        .wrap_err("Cannot render kernel template")?;
-
-    Ok(())
-}
-
 fn write_setup_py(
     env: &Environment,
     torch: &Torch,

build2cmake/src/torch/cuda.rs

@@ -9,6 +9,7 @@ use minijinja::{context, Environment};
 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_ops_py;
 use crate::torch::common::write_pyproject_toml;
 use crate::torch::common::write_torch_registration_macros;
 use crate::torch::kernel::render_kernel_components;
@@ -93,31 +94,6 @@ fn write_setup_py(
     Ok(())
 }
 
-fn write_ops_py(
-    env: &Environment,
-    name: &str,
-    ops_name: &str,
-    file_set: &mut FileSet,
-) -> Result<()> {
-    let mut path = PathBuf::new();
-    path.push("torch-ext");
-    path.push(name);
-    path.push("_ops.py");
-    let writer = file_set.entry(path);
-
-    env.get_template("_ops.py")
-        .wrap_err("Cannot get _ops.py template")?
-        .render_to_write(
-            context! {
-                ops_name => ops_name,
-            },
-            writer,
-        )
-        .wrap_err("Cannot render kernel template")?;
-
-    Ok(())
-}
-
 fn write_cmake(
     env: &Environment,
     backend: Backend,

build2cmake/src/torch/metal.rs

@@ -7,6 +7,7 @@ use minijinja::{context, Environment};
 use crate::config::{Backend, Build, Torch};
 use crate::fileset::FileSet;
 use crate::torch::common::write_metadata;
+use crate::torch::common::write_ops_py;
 use crate::torch::common::write_pyproject_toml;
 use crate::torch::common::write_torch_registration_macros;
 use crate::torch::kernel::render_kernel_components;
@@ -188,31 +189,6 @@ fn render_preamble(
     Ok(())
 }
 
-fn write_ops_py(
-    env: &Environment,
-    name: &str,
-    ops_name: &str,
-    file_set: &mut FileSet,
-) -> Result<()> {
-    let mut path = PathBuf::new();
-    path.push("torch-ext");
-    path.push(name);
-    path.push("_ops.py");
-    let writer = file_set.entry(path);
-
-    env.get_template("_ops.py")
-        .wrap_err("Cannot get _ops.py template")?
-        .render_to_write(
-            context! {
-                ops_name => ops_name,
-            },
-            writer,
-        )
-        .wrap_err("Cannot render kernel template")?;
-
-    Ok(())
-}
-
 fn write_setup_py(
     env: &Environment,
     torch: &Torch,

build2cmake/src/torch/xpu.rs

@@ -8,6 +8,7 @@ use minijinja::{context, Environment};
 
 use crate::config::{Backend, Build, Dependency, Torch};
 use crate::torch::common::write_metadata;
+use crate::torch::common::write_ops_py;
 use crate::torch::common::write_pyproject_toml;
 use crate::torch::common::write_torch_registration_macros;
 use crate::torch::kernel::render_kernel_components;
@@ -89,31 +90,6 @@ fn write_setup_py(
     Ok(())
 }
 
-fn write_ops_py(
-    env: &Environment,
-    name: &str,
-    ops_name: &str,
-    file_set: &mut FileSet,
-) -> Result<()> {
-    let mut path = PathBuf::new();
-    path.push("torch-ext");
-    path.push(name);
-    path.push("_ops.py");
-    let writer = file_set.entry(path);
-
-    env.get_template("_ops.py")
-        .wrap_err("Cannot get _ops.py template")?
-        .render_to_write(
-            context! {
-                ops_name => ops_name,
-            },
-            writer,
-        )
-        .wrap_err("Cannot render _ops.py template")?;
-
-    Ok(())
-}
-
 fn write_cmake(
     env: &Environment,
     build: &Build,