Add support for backend-specific Python dependencies

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,77 @@ 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: {backend:?}")]
+    UnknownBackend { backend: String },
+    #[error("Unknown dependency `{dependency:?}` for backend `{backend:?}`")]
+    UnknownDependency { backend: String, dependency: String },
+    #[error("Unknown dependency: `{dependency:?}`")]
+    UnknownGeneralDependency { 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::UnknownGeneralDependency {
+                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::UnknownBackend {
+                    backend: backend.to_string(),
+                })
+            }
+            Some(backend_deps) => backend_deps,
+        };
+        match backend_deps.get(dependency) {
+            None => {
+                return Err(DependencyError::UnknownDependency {
+                    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,46 +106,114 @@ 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 enum CudaPythonDependency {
+    NvidiaCutlassDsl,
 }
 
-#[derive(Clone, Debug, Deserialize, Serialize)]
+impl Display for CudaPythonDependency {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            CudaPythonDependency::NvidiaCutlassDsl => write!(f, "nvidia-cutlass-dsl"),
+        }
+    }
+}
+
+#[derive(Debug, Deserialize, Serialize)]
 #[serde(deny_unknown_fields, rename_all = "kebab-case")]
-pub enum PythonDependency {
-    Einops,
-    NvidiaCutlassDsl,
+pub struct XpuGeneral {
+    pub python_depends: Option<Vec<String>>,
 }
 
-impl Display for PythonDependency {
+#[derive(Debug, Deserialize, Serialize)]
+#[serde(deny_unknown_fields, rename_all = "kebab-case")]
+pub enum XpuPythonDependency {
+    Onednn,
+}
+
+impl Display for XpuPythonDependency {
     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"),
+            XpuPythonDependency::Onednn => write!(f, "onednn-devel"),
         }
     }
 }
 
+#[derive(Debug, Deserialize, Serialize)]
+#[serde(deny_unknown_fields, rename_all = "kebab-case")]
+pub struct Hub {
+    pub repo_id: Option<String>,
+    pub branch: Option<String>,
+}
+
 #[derive(Debug, Deserialize, Clone, Serialize)]
 #[serde(deny_unknown_fields)]
 pub struct Torch {
@@ -215,7 +345,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 +420,7 @@ impl General {
             Some(CudaGeneral {
                 minver: general.cuda_minver,
                 maxver: general.cuda_maxver,
+                python_depends: None,
             })
         } else {
             None
@@ -300,9 +431,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 +446,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)?,
@@ -376,26 +376,30 @@ fn get_generated_files(
     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 {
-                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())?
+            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())?
+                    }
                 }
             };
-
             all_set.extend(set);
         }
     }

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"]
+      }
+    }
+  }
+}