Add support for (limited) Python dependencies: nvidia-cutlass-dsl and einops (#302)

This change adds support for kernel dependencies. Dependencies are
specified in `build.toml`, e.g.:

```toml
[general]
python-depends = [ "nvidia-cutlass-dsl" ]
```

The set of dependencies is limited to avoid bringing back the problem of
resolving Python dependencies. Currently only `nvidia-cutlass-dsl` and
`einops` are supported.

This PR also adds the derivations for the `nvidia-cutlass-dsl` package.

Author: danieldk

build2cmake/src/config/mod.rs

@@ -5,7 +5,7 @@ pub mod v1;
 
 mod v2;
 use serde_value::Value;
-pub use v2::{Backend, Build, Dependencies, Kernel, Torch};
+pub use v2::{Backend, Build, Dependency, General, Kernel, Torch};
 
 #[derive(Debug)]
 pub enum BuildCompat {

build2cmake/src/config/v1.rs

@@ -2,7 +2,7 @@ use std::{collections::HashMap, fmt::Display, path::PathBuf};
 
 use serde::Deserialize;
 
-use super::v2::Dependencies;
+use super::v2::Dependency;
 
 #[derive(Debug, Deserialize)]
 #[serde(deny_unknown_fields)]
@@ -40,7 +40,7 @@ pub struct Kernel {
     pub rocm_archs: Option<Vec<String>>,
     #[serde(default)]
     pub language: Language,
-    pub depends: Vec<Dependencies>,
+    pub depends: Vec<Dependency>,
     pub include: Option<Vec<String>>,
     pub src: Vec<String>,
 }

build2cmake/src/config/v2.rs

@@ -54,6 +54,8 @@ pub struct General {
     pub cuda_minver: Option<Version>,
 
     pub hub: Option<Hub>,
+
+    pub python_depends: Option<Vec<PythonDependency>>,
 }
 
 impl General {
@@ -70,6 +72,22 @@ pub struct Hub {
     pub branch: Option<String>,
 }
 
+#[derive(Clone, Debug, Deserialize, Serialize)]
+#[serde(deny_unknown_fields, rename_all = "kebab-case")]
+pub enum PythonDependency {
+    Einops,
+    NvidiaCutlassDsl,
+}
+
+impl Display for PythonDependency {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            PythonDependency::Einops => write!(f, "einops"),
+            PythonDependency::NvidiaCutlassDsl => write!(f, "nvidia-cutlass-dsl"),
+        }
+    }
+}
+
 #[derive(Debug, Deserialize, Clone, Serialize)]
 #[serde(deny_unknown_fields)]
 pub struct Torch {
@@ -107,7 +125,7 @@ pub enum Kernel {
     #[serde(rename_all = "kebab-case")]
     Cpu {
         cxx_flags: Option<Vec<String>>,
-        depends: Vec<Dependencies>,
+        depends: Vec<Dependency>,
         include: Option<Vec<String>>,
         src: Vec<String>,
     },
@@ -117,21 +135,21 @@ pub enum Kernel {
         cuda_flags: Option<Vec<String>>,
         cuda_minver: Option<Version>,
         cxx_flags: Option<Vec<String>>,
-        depends: Vec<Dependencies>,
+        depends: Vec<Dependency>,
         include: Option<Vec<String>>,
         src: Vec<String>,
     },
     #[serde(rename_all = "kebab-case")]
     Metal {
         cxx_flags: Option<Vec<String>>,
-        depends: Vec<Dependencies>,
+        depends: Vec<Dependency>,
         include: Option<Vec<String>>,
         src: Vec<String>,
     },
     #[serde(rename_all = "kebab-case")]
     Rocm {
         cxx_flags: Option<Vec<String>>,
-        depends: Vec<Dependencies>,
+        depends: Vec<Dependency>,
         rocm_archs: Option<Vec<String>>,
         hip_flags: Option<Vec<String>>,
         include: Option<Vec<String>>,
@@ -140,7 +158,7 @@ pub enum Kernel {
     #[serde(rename_all = "kebab-case")]
     Xpu {
         cxx_flags: Option<Vec<String>>,
-        depends: Vec<Dependencies>,
+        depends: Vec<Dependency>,
         sycl_flags: Option<Vec<String>>,
         include: Option<Vec<String>>,
         src: Vec<String>,
@@ -178,7 +196,7 @@ impl Kernel {
         }
     }
 
-    pub fn depends(&self) -> &[Dependencies] {
+    pub fn depends(&self) -> &[Dependency] {
         match self {
             Kernel::Cpu { depends, .. }
             | Kernel::Cuda { depends, .. }
@@ -239,7 +257,7 @@ impl FromStr for Backend {
 #[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)]
 #[non_exhaustive]
 #[serde(rename_all = "lowercase")]
-pub enum Dependencies {
+pub enum Dependency {
     #[serde(rename = "cutlass_2_10")]
     Cutlass2_10,
     #[serde(rename = "cutlass_3_5")]
@@ -284,6 +302,7 @@ impl General {
             cuda_maxver: None,
             cuda_minver: None,
             hub: None,
+            python_depends: None,
         }
     }
 }

build2cmake/src/templates/pyproject.toml

@@ -6,5 +6,6 @@ requires = [
   "setuptools>=61",
   "torch",
   "wheel",
+  {{python_dependencies}}
 ]
 build-backend = "setuptools.build_meta"

build2cmake/src/templates/universal/pyproject.toml

@@ -2,7 +2,10 @@
 name = "{{ name }}"
 version = "0.0.1"
 requires-python = ">= 3.9"
-dependencies = ["torch>=2.4"]
+dependencies = [
+  "torch>=2.8",
+  {{python_dependencies}}
+]
 
 [tool.setuptools]
 package-dir = { "" = "torch-ext" }

build2cmake/src/torch/common.rs

@@ -0,0 +1,33 @@
+use eyre::{Context, Result};
+use itertools::Itertools;
+use minijinja::{context, Environment};
+
+use crate::{config::General, FileSet};
+
+pub fn write_pyproject_toml(
+    env: &Environment,
+    general: &General,
+    file_set: &mut FileSet,
+) -> Result<()> {
+    let writer = file_set.entry("pyproject.toml");
+
+    let python_dependencies = general
+        .python_depends
+        .as_ref()
+        .unwrap_or(&vec![])
+        .iter()
+        .map(|d| format!("\"{d}\""))
+        .join(", ");
+
+    env.get_template("pyproject.toml")
+        .wrap_err("Cannot get pyproject.toml template")?
+        .render_to_write(
+            context! {
+                python_dependencies => python_dependencies,
+            },
+            writer,
+        )
+        .wrap_err("Cannot render kernel template")?;
+
+    Ok(())
+}

build2cmake/src/torch/cpu.rs

@@ -4,7 +4,7 @@ use eyre::{bail, Context, Result};
 use itertools::Itertools;
 use minijinja::{context, Environment};
 
-use super::kernel_ops_identifier;
+use super::{common::write_pyproject_toml, kernel_ops_identifier};
 use crate::{
     config::{Build, Kernel, Torch},
     fileset::FileSet,
@@ -47,7 +47,7 @@ pub fn write_torch_ext_cpu(
 
     write_ops_py(env, &build.general.python_name(), &ops_name, &mut file_set)?;
 
-    write_pyproject_toml(env, &mut file_set)?;
+    write_pyproject_toml(env, &build.general, &mut file_set)?;
 
     write_torch_registration_macros(&mut file_set)?;
 
@@ -209,17 +209,6 @@ fn write_ops_py(
     Ok(())
 }
 
-fn write_pyproject_toml(env: &Environment, file_set: &mut FileSet) -> Result<()> {
-    let writer = file_set.entry("pyproject.toml");
-
-    env.get_template("pyproject.toml")
-        .wrap_err("Cannot get pyproject.toml template")?
-        .render_to_write(context! {}, writer)
-        .wrap_err("Cannot render kernel template")?;
-
-    Ok(())
-}
-
 fn write_setup_py(
     env: &Environment,
     torch: &Torch,

build2cmake/src/torch/cuda.rs

@@ -7,8 +7,9 @@ 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, Dependencies, Kernel, Torch};
+use crate::config::{Backend, Build, Dependency, Kernel, Torch};
 use crate::version::Version;
 use crate::FileSet;
 
@@ -60,7 +61,7 @@ pub fn write_torch_ext_cuda(
 
     write_ops_py(env, &build.general.python_name(), &ops_name, &mut file_set)?;
 
-    write_pyproject_toml(env, &mut file_set)?;
+    write_pyproject_toml(env, &build.general, &mut file_set)?;
 
     write_torch_registration_macros(&mut file_set)?;
 
@@ -78,17 +79,6 @@ fn write_torch_registration_macros(file_set: &mut FileSet) -> Result<()> {
     Ok(())
 }
 
-fn write_pyproject_toml(env: &Environment, file_set: &mut FileSet) -> Result<()> {
-    let writer = file_set.entry("pyproject.toml");
-
-    env.get_template("pyproject.toml")
-        .wrap_err("Cannot get pyproject.toml template")?
-        .render_to_write(context! {}, writer)
-        .wrap_err("Cannot render kernel template")?;
-
-    Ok(())
-}
-
 fn write_setup_py(
     env: &Environment,
     torch: &Torch,
@@ -230,7 +220,7 @@ fn render_deps(env: &Environment, build: &Build, write: &mut impl Write) -> Resu
 
     for dep in deps {
         match dep {
-            Dependencies::Cutlass2_10 => {
+            Dependency::Cutlass2_10 => {
                 env.get_template("cuda/dep-cutlass.cmake")
                     .wrap_err("Cannot get CUTLASS dependency template")?
                     .render_to_write(
@@ -241,7 +231,7 @@ fn render_deps(env: &Environment, build: &Build, write: &mut impl Write) -> Resu
                     )
                     .wrap_err("Cannot render CUTLASS dependency template")?;
             }
-            Dependencies::Cutlass3_5 => {
+            Dependency::Cutlass3_5 => {
                 env.get_template("cuda/dep-cutlass.cmake")
                     .wrap_err("Cannot get CUTLASS dependency template")?
                     .render_to_write(
@@ -252,7 +242,7 @@ fn render_deps(env: &Environment, build: &Build, write: &mut impl Write) -> Resu
                     )
                     .wrap_err("Cannot render CUTLASS dependency template")?;
             }
-            Dependencies::Cutlass3_6 => {
+            Dependency::Cutlass3_6 => {
                 env.get_template("cuda/dep-cutlass.cmake")
                     .wrap_err("Cannot get CUTLASS dependency template")?
                     .render_to_write(
@@ -263,7 +253,7 @@ fn render_deps(env: &Environment, build: &Build, write: &mut impl Write) -> Resu
                     )
                     .wrap_err("Cannot render CUTLASS dependency template")?;
             }
-            Dependencies::Cutlass3_8 => {
+            Dependency::Cutlass3_8 => {
                 env.get_template("cuda/dep-cutlass.cmake")
                     .wrap_err("Cannot get CUTLASS dependency template")?
                     .render_to_write(
@@ -274,7 +264,7 @@ fn render_deps(env: &Environment, build: &Build, write: &mut impl Write) -> Resu
                     )
                     .wrap_err("Cannot render CUTLASS dependency template")?;
             }
-            Dependencies::Cutlass3_9 => {
+            Dependency::Cutlass3_9 => {
                 env.get_template("cuda/dep-cutlass.cmake")
                     .wrap_err("Cannot get CUTLASS dependency template")?
                     .render_to_write(
@@ -285,7 +275,7 @@ fn render_deps(env: &Environment, build: &Build, write: &mut impl Write) -> Resu
                     )
                     .wrap_err("Cannot render CUTLASS dependency template")?;
             }
-            Dependencies::Cutlass4_0 => {
+            Dependency::Cutlass4_0 => {
                 env.get_template("cuda/dep-cutlass.cmake")
                     .wrap_err("Cannot get CUTLASS dependency template")?
                     .render_to_write(
@@ -296,7 +286,7 @@ fn render_deps(env: &Environment, build: &Build, write: &mut impl Write) -> Resu
                     )
                     .wrap_err("Cannot render CUTLASS dependency template")?;
             }
-            Dependencies::Torch => (),
+            Dependency::Torch => (),
             _ => {
                 eprintln!("Warning: CUDA backend doesn't need/support dependency: {dep:?}");
             }

build2cmake/src/torch/metal.rs

@@ -4,7 +4,7 @@ use eyre::{bail, Context, Result};
 use itertools::Itertools;
 use minijinja::{context, Environment};
 
-use super::kernel_ops_identifier;
+use super::{common::write_pyproject_toml, kernel_ops_identifier};
 use crate::{
     config::{Build, Kernel, Torch},
     fileset::FileSet,
@@ -49,7 +49,7 @@ pub fn write_torch_ext_metal(
 
     write_ops_py(env, &build.general.python_name(), &ops_name, &mut file_set)?;
 
-    write_pyproject_toml(env, &mut file_set)?;
+    write_pyproject_toml(env, &build.general, &mut file_set)?;
 
     write_torch_registration_macros(&mut file_set)?;
 
@@ -225,17 +225,6 @@ fn write_ops_py(
     Ok(())
 }
 
-fn write_pyproject_toml(env: &Environment, file_set: &mut FileSet) -> Result<()> {
-    let writer = file_set.entry("pyproject.toml");
-
-    env.get_template("pyproject.toml")
-        .wrap_err("Cannot get pyproject.toml template")?
-        .render_to_write(context! {}, writer)
-        .wrap_err("Cannot render kernel template")?;
-
-    Ok(())
-}
-
 fn write_setup_py(
     env: &Environment,
     torch: &Torch,

build2cmake/src/torch/mod.rs

@@ -4,6 +4,8 @@ pub use cpu::write_torch_ext_cpu;
 mod cuda;
 pub use cuda::write_torch_ext_cuda;
 
+pub mod common;
+
 mod metal;
 pub use metal::write_torch_ext_metal;