Remove the last backend-specific writer functions (#255)

* build2cmake: merge `write_cmake` functions in `common`

* build2cmake: merge all `write_torch_ext` implementations in `common`

Author: danieldk

build2cmake/src/main.rs

@@ -9,10 +9,7 @@ use eyre::{bail, ensure, Context, Result};
 use minijinja::Environment;
 
 mod torch;
-use torch::{
-    write_torch_ext_cpu, write_torch_ext_cuda, write_torch_ext_metal, write_torch_ext_noarch,
-    write_torch_ext_xpu,
-};
+use torch::{write_torch_ext, write_torch_ext_noarch};
 
 mod config;
 use config::{v3, Backend, Build, BuildCompat};
@@ -176,14 +173,7 @@ fn generate_torch(
     let file_set = if build.is_noarch() {
         write_torch_ext_noarch(&env, backend, &build, target_dir.clone(), ops_id)?
     } else {
-        match backend {
-            Backend::Cpu => write_torch_ext_cpu(&env, &build, target_dir.clone(), ops_id)?,
-            Backend::Cuda | Backend::Rocm => {
-                write_torch_ext_cuda(&env, backend, &build, target_dir.clone(), ops_id)?
-            }
-            Backend::Metal => write_torch_ext_metal(&env, &build, target_dir.clone(), ops_id)?,
-            Backend::Xpu => write_torch_ext_xpu(&env, &build, target_dir.clone(), ops_id)?,
-        }
+        write_torch_ext(&env, backend, &build, target_dir.clone(), ops_id)?
     };
     file_set.write(&target_dir, force)?;
 
@@ -382,20 +372,7 @@ fn get_generated_files(
         let set = if build.is_noarch() {
             write_torch_ext_noarch(env, *backend, build, target_dir.clone(), ops_id.clone())?
         } else {
-            match backend {
-                Backend::Cpu => {
-                    write_torch_ext_cpu(env, build, target_dir.clone(), ops_id.clone())?
-                }
-                Backend::Cuda | Backend::Rocm => {
-                    write_torch_ext_cuda(env, *backend, build, target_dir.clone(), ops_id.clone())?
-                }
-                Backend::Metal => {
-                    write_torch_ext_metal(env, build, target_dir.clone(), ops_id.clone())?
-                }
-                Backend::Xpu => {
-                    write_torch_ext_xpu(env, build, target_dir.clone(), ops_id.clone())?
-                }
-            }
+            write_torch_ext(env, *backend, build, target_dir.clone(), ops_id.clone())?
         };
         all_set.extend(set);
     }

build2cmake/src/torch/common.rs

@@ -1,12 +1,14 @@
 use std::io::Write;
 use std::path::PathBuf;
 
-use eyre::{Context, Result};
+use eyre::{bail, Context, Result};
 use itertools::Itertools;
 use minijinja::{context, Environment};
 
-use crate::config::{Backend, General, Torch};
+use crate::config::{Backend, Build, General, Torch};
 use crate::metadata::Metadata;
+use crate::torch::deps::render_deps;
+use crate::torch::kernel::render_kernel_components;
 use crate::version::Version;
 use crate::FileSet;
 
@@ -245,3 +247,94 @@ pub fn render_preamble(
 
     Ok(())
 }
+
+pub fn write_cmake(
+    env: &Environment,
+    backend: Backend,
+    build: &Build,
+    torch: &Torch,
+    name: &str,
+    ops_name: &str,
+    file_set: &mut FileSet,
+) -> Result<()> {
+    write_cmake_helpers(file_set);
+
+    let cmake_writer = file_set.entry("CMakeLists.txt");
+
+    let (cuda_minver, cuda_maxver) = match backend {
+        Backend::Cuda => (
+            build.general.cuda.as_ref().and_then(|c| c.minver.as_ref()),
+            build.general.cuda.as_ref().and_then(|c| c.maxver.as_ref()),
+        ),
+        _ => (None, None),
+    };
+
+    render_preamble(
+        env,
+        name,
+        cuda_minver,
+        cuda_maxver,
+        torch.minver.as_ref(),
+        torch.maxver.as_ref(),
+        cmake_writer,
+    )?;
+
+    render_deps(env, backend, build, cmake_writer)?;
+
+    render_binding(env, torch, name, cmake_writer)?;
+
+    render_kernel_components(env, build, cmake_writer)?;
+
+    render_extension(env, name, ops_name, cmake_writer)?;
+
+    Ok(())
+}
+
+pub fn write_torch_ext(
+    env: &Environment,
+    backend: Backend,
+    build: &Build,
+    target_dir: PathBuf,
+    ops_id: Option<String>,
+) -> Result<FileSet> {
+    let torch_ext = match build.torch.as_ref() {
+        Some(torch_ext) => torch_ext,
+        None => bail!("Build configuration does not have `torch` section"),
+    };
+
+    let mut file_set = FileSet::default();
+
+    let ops_name = crate::torch::ops_identifier::kernel_ops_identifier(
+        &target_dir,
+        &build.general.python_name(),
+        ops_id,
+    );
+
+    write_cmake(
+        env,
+        backend,
+        build,
+        torch_ext,
+        &build.general.name,
+        &ops_name,
+        &mut file_set,
+    )?;
+
+    write_setup_py(
+        env,
+        torch_ext,
+        &build.general.name,
+        &ops_name,
+        &mut file_set,
+    )?;
+
+    write_ops_py(env, &build.general.python_name(), &ops_name, &mut file_set)?;
+
+    write_pyproject_toml(env, backend, &build.general, &mut file_set)?;
+
+    write_torch_registration_macros(&mut file_set)?;
+
+    write_metadata(backend, &build.general, &mut file_set)?;
+
+    Ok(file_set)
+}