Add support for backend-specific Python dependencies (#325)

Add `general.cuda.python-depends` and `general.xpu.python-depends`.

Currently `nvidia-cutlass-dsl` is supported for CUDA and `onednn` for
XPU.

Author: danieldk

build2cmake/Cargo.lock

@@ -20,6 +20,7 @@ rand = "0.8"
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
 serde-value = "0.7"
+thiserror = "1"
 toml = "0.8"
 
 [build-dependencies]

build2cmake/Cargo.toml

@@ -3,15 +3,75 @@ use std::{
     fmt::Display,
     path::PathBuf,
     str::FromStr,
+    sync::LazyLock,
 };
 
 use eyre::Result;
 use itertools::Itertools;
 use serde::{Deserialize, Serialize};
+use thiserror::Error;
 
+use super::{common::Dependency, v2};
 use crate::version::Version;
 
-use super::{common::Dependency, v2};
+#[derive(Debug, Error)]
+enum DependencyError {
+    #[error("No dependencies are defined for backend: {backend:?}")]
+    Backend { backend: String },
+    #[error("Unknown dependency `{dependency:?}` for backend `{backend:?}`")]
+    Dependency { backend: String, dependency: String },
+    #[error("Unknown dependency: `{dependency:?}`")]
+    GeneralDependency { dependency: String },
+}
+
+#[derive(Debug, Deserialize, Serialize)]
+#[serde(deny_unknown_fields)]
+struct PythonDependencies {
+    general: HashMap<String, PythonDependency>,
+    backends: HashMap<Backend, HashMap<String, PythonDependency>>,
+}
+
+impl PythonDependencies {
+    fn get_dependency(&self, dependency: &str) -> Result<&[String], DependencyError> {
+        match self.general.get(dependency) {
+            None => Err(DependencyError::GeneralDependency {
+                dependency: dependency.to_string(),
+            }),
+            Some(dep) => Ok(&dep.python),
+        }
+    }
+
+    fn get_backend_dependency(
+        &self,
+        backend: Backend,
+        dependency: &str,
+    ) -> Result<&[String], DependencyError> {
+        let backend_deps = match self.backends.get(&backend) {
+            None => {
+                return Err(DependencyError::Backend {
+                    backend: backend.to_string(),
+                })
+            }
+            Some(backend_deps) => backend_deps,
+        };
+        match backend_deps.get(dependency) {
+            None => Err(DependencyError::Dependency {
+                backend: backend.to_string(),
+                dependency: dependency.to_string(),
+            }),
+            Some(dep) => Ok(&dep.python),
+        }
+    }
+}
+
+#[derive(Debug, Deserialize, Serialize)]
+struct PythonDependency {
+    nix: Vec<String>,
+    python: Vec<String>,
+}
+
+static PYTHON_DEPENDENCIES: LazyLock<PythonDependencies> =
+    LazyLock::new(|| serde_json::from_str(include_str!("../python_dependencies.json")).unwrap());
 
 #[derive(Debug, Deserialize, Serialize)]
 #[serde(deny_unknown_fields)]
@@ -44,44 +104,84 @@ pub struct General {
 
     pub hub: Option<Hub>,
 
-    pub python_depends: Option<Vec<PythonDependency>>,
+    pub python_depends: Option<Vec<String>>,
+
+    pub xpu: Option<XpuGeneral>,
 }
 
 impl General {
     /// Name of the kernel as a Python extension.
     pub fn python_name(&self) -> String {
         self.name.replace("-", "_")
     }
+
+    pub fn python_depends(&self) -> Box<dyn Iterator<Item = Result<String>> + '_> {
+        let general_python_deps = match self.python_depends.as_ref() {
+            Some(deps) => deps,
+            None => {
+                return Box::new(std::iter::empty());
+            }
+        };
+
+        Box::new(general_python_deps.iter().flat_map(move |dep| {
+            match PYTHON_DEPENDENCIES.get_dependency(dep) {
+                Ok(deps) => deps.iter().map(|s| Ok(s.clone())).collect::<Vec<_>>(),
+                Err(e) => vec![Err(e.into())],
+            }
+        }))
+    }
+
+    pub fn backend_python_depends(
+        &self,
+        backend: Backend,
+    ) -> Box<dyn Iterator<Item = Result<String>> + '_> {
+        let backend_python_deps = match backend {
+            Backend::Cuda => self
+                .cuda
+                .as_ref()
+                .and_then(|cuda| cuda.python_depends.as_ref()),
+            Backend::Xpu => self
+                .xpu
+                .as_ref()
+                .and_then(|xpu| xpu.python_depends.as_ref()),
+            _ => None,
+        };
+
+        let backend_python_deps = match backend_python_deps {
+            Some(deps) => deps,
+            None => {
+                return Box::new(std::iter::empty());
+            }
+        };
+
+        Box::new(backend_python_deps.iter().flat_map(move |dep| {
+            match PYTHON_DEPENDENCIES.get_backend_dependency(backend, dep) {
+                Ok(deps) => deps.iter().map(|s| Ok(s.clone())).collect::<Vec<_>>(),
+                Err(e) => vec![Err(e.into())],
+            }
+        }))
+    }
 }
 
 #[derive(Debug, Deserialize, Serialize)]
 #[serde(deny_unknown_fields, rename_all = "kebab-case")]
 pub struct CudaGeneral {
     pub minver: Option<Version>,
     pub maxver: Option<Version>,
+    pub python_depends: Option<Vec<String>>,
 }
 
 #[derive(Debug, Deserialize, Serialize)]
 #[serde(deny_unknown_fields, rename_all = "kebab-case")]
-pub struct Hub {
-    pub repo_id: Option<String>,
-    pub branch: Option<String>,
+pub struct XpuGeneral {
+    pub python_depends: Option<Vec<String>>,
 }
 
-#[derive(Clone, Debug, Deserialize, Serialize)]
+#[derive(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"),
-        }
-    }
+pub struct Hub {
+    pub repo_id: Option<String>,
+    pub branch: Option<String>,
 }
 
 #[derive(Debug, Deserialize, Clone, Serialize)]
@@ -215,7 +315,7 @@ impl Kernel {
     }
 }
 
-#[derive(Clone, Copy, Debug, Deserialize, Eq, Ord, PartialEq, PartialOrd, Serialize)]
+#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)]
 #[serde(deny_unknown_fields, rename_all = "kebab-case")]
 pub enum Backend {
     Cpu,
@@ -290,6 +390,7 @@ impl General {
             Some(CudaGeneral {
                 minver: general.cuda_minver,
                 maxver: general.cuda_maxver,
+                python_depends: None,
             })
         } else {
             None
@@ -300,9 +401,8 @@ impl General {
             backends,
             cuda,
             hub: general.hub.map(Into::into),
-            python_depends: general
-                .python_depends
-                .map(|deps| deps.into_iter().map(Into::into).collect()),
+            python_depends: None,
+            xpu: None,
         }
     }
 }
@@ -316,15 +416,6 @@ impl From<v2::Hub> for Hub {
     }
 }
 
-impl From<v2::PythonDependency> for PythonDependency {
-    fn from(dep: v2::PythonDependency) -> Self {
-        match dep {
-            v2::PythonDependency::Einops => PythonDependency::Einops,
-            v2::PythonDependency::NvidiaCutlassDsl => PythonDependency::NvidiaCutlassDsl,
-        }
-    }
-}
-
 impl From<v2::Torch> for Torch {
     fn from(torch: v2::Torch) -> Self {
         Self {

build2cmake/src/config/v3.rs

@@ -172,7 +172,7 @@ fn generate_torch(
     };
 
     let file_set = if build.is_noarch() {
-        write_torch_ext_noarch(&env, &build, target_dir.clone(), ops_id)?
+        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)?,
@@ -375,13 +375,11 @@ fn get_generated_files(
 ) -> Result<Vec<PathBuf>> {
     let mut all_set = FileSet::new();
 
-    if build.is_noarch() {
-        let set = write_torch_ext_noarch(env, build, target_dir.clone(), ops_id.clone())?;
-
-        all_set.extend(set);
-    } else {
-        for backend in &build.general.backends {
-            let set = match backend {
+    for backend in &build.general.backends {
+        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())?
                 }
@@ -394,10 +392,9 @@ fn get_generated_files(
                 Backend::Xpu => {
                     write_torch_ext_xpu(env, build, target_dir.clone(), ops_id.clone())?
                 }
-            };
-
-            all_set.extend(set);
-        }
+            }
+        };
+        all_set.extend(set);
     }
 
     Ok(all_set.into_names())

build2cmake/src/main.rs

@@ -0,0 +1,25 @@
+{
+  "general": {
+    "einops": {
+      "nix": ["einops"],
+      "python": ["einops"]
+    }
+  },
+  "backends": {
+    "cpu": {},
+    "cuda": {
+      "nvidia-cutlass-dsl": {
+        "nix": ["nvidia-cutlass-dsl"],
+        "python": ["nvidia-cutlass-dsl"]
+      }
+    },
+    "metal": {},
+    "rocm": {},
+    "xpu": {
+      "onednn": {
+        "nix": [],
+        "python": ["onednn-devel"]
+      }
+    }
+  }
+}

build2cmake/src/python_dependencies.json

@@ -2,22 +2,23 @@ use eyre::{Context, Result};
 use itertools::Itertools;
 use minijinja::{context, Environment};
 
-use crate::{config::General, FileSet};
+use crate::config::{Backend, General};
+use crate::FileSet;
 
 pub fn write_pyproject_toml(
     env: &Environment,
+    backend: Backend,
     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(", ");
+    let python_dependencies = itertools::process_results(
+        general
+            .python_depends()
+            .chain(general.backend_python_depends(backend)),
+        |iter| iter.map(|d| format!("\"{d}\"")).join(", "),
+    )?;
 
     env.get_template("pyproject.toml")
         .wrap_err("Cannot get pyproject.toml template")?

build2cmake/src/torch/common.rs

@@ -6,7 +6,7 @@ use minijinja::{context, Environment};
 
 use super::{common::write_pyproject_toml, kernel_ops_identifier};
 use crate::{
-    config::{Build, Kernel, Torch},
+    config::{Backend, Build, Kernel, Torch},
     fileset::FileSet,
     version::Version,
 };
@@ -48,7 +48,7 @@ pub fn write_torch_ext_cpu(
 
     write_ops_py(env, &build.general.python_name(), &ops_name, &mut file_set)?;
 
-    write_pyproject_toml(env, &build.general, &mut file_set)?;
+    write_pyproject_toml(env, Backend::Cpu, &build.general, &mut file_set)?;
 
     write_torch_registration_macros(&mut file_set)?;
 

build2cmake/src/torch/cpu.rs

@@ -61,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, &build.general, &mut file_set)?;
+    write_pyproject_toml(env, backend, &build.general, &mut file_set)?;
 
     write_torch_registration_macros(&mut file_set)?;