treewide: add support for MLX noarch kernels

examples/kernels/flake.nix

@@ -449,6 +449,11 @@
           path = ./relu-metal-cpp;
           drv = sys: out: out.packages.${sys}.redistributable.${"torch${torchVersion}-metal-${sys}"};
         }
+        {
+          name = "relu-mlx-kernel";
+          path = ./relu-mlx;
+          drv = sys: out: out.packages.${sys}.redistributable.mlx-metal;
+        }
       ];
 
       # CPU kernels to build in CI.

examples/kernels/relu-mlx/build.toml

@@ -0,0 +1,11 @@
+[general]
+name = "relu-mlx"
+version = 1
+edition = 6
+license = "Apache-2.0"
+backends = ["metal"]
+
+[general.hub]
+repo-id = "kernels-test/relu-mlx"
+
+[mlx-noarch]

examples/kernels/relu-mlx/flake.nix

@@ -0,0 +1,17 @@
+{
+  description = "Flake for MLX ReLU kernel";
+
+  inputs = {
+    kernel-builder.url = "path:../../..";
+  };
+
+  outputs =
+    {
+      self,
+      kernel-builder,
+    }:
+    kernel-builder.lib.genKernelFlakeOutputs {
+      inherit self;
+      path = ./.;
+    };
+}

examples/kernels/relu-mlx/mlx-ext/relu_mlx/__init__.py

@@ -0,0 +1,8 @@
+import mlx.core as mx
+
+
+def relu(x: mx.array) -> mx.array:
+    return mx.maximum(x, 0)
+
+
+__all__ = ["relu"]

examples/kernels/relu-mlx/tests/test_relu.py

@@ -0,0 +1,11 @@
+import kernels
+import mlx.core as mx
+import pytest
+
+relu = kernels.get_kernel("kernels-test/relu-mlx", version=1)
+
+
+@pytest.mark.kernels_ci
+def test_relu():
+    x = mx.random.normal((1024, 1024))
+    assert mx.array_equal(relu.relu(x), mx.maximum(x, 0))

kernel-builder/src/card.rs

@@ -12,7 +12,7 @@ use rustpython_parser::{ast, Parse};
 use crate::util::check_or_infer_kernel_dir;
 
 fn extract_all(kernel_dir: &Path, module_name: &str) -> Option<Vec<String>> {
-    let init_path = ["torch-ext", "tvm-ffi-ext"]
+    let init_path = ["mlx-ext", "torch-ext", "tvm-ffi-ext"]
         .iter()
         .map(|ext_dir| {
             kernel_dir
@@ -228,7 +228,7 @@ mod tests {
 
     #[test]
     fn test_extract_functions_excludes_private_exports() {
-        for ext_dir in ["torch-ext", "tvm-ffi-ext"] {
+        for ext_dir in ["mlx-ext", "torch-ext", "tvm-ffi-ext"] {
             let temp_dir = tempfile::tempdir().unwrap();
             let kernel_dir = temp_dir.path();
             let module_dir = kernel_dir.join(ext_dir).join("test_module");

kernel-builder/src/pyproject/common.rs

@@ -37,7 +37,18 @@ pub fn write_metadata(
     provenance: Option<&Provenance>,
     file_set: &mut FileSet,
 ) -> Result<()> {
-    for backend in &Backend::all() {
+    write_metadata_for_backends(build, &Backend::all(), kernel_id, provenance, file_set)
+}
+
+/// Write metadata for the given backends only.
+pub fn write_metadata_for_backends(
+    build: &Build,
+    backends: &[Backend],
+    kernel_id: &KernelIdentifier,
+    provenance: Option<&Provenance>,
+    file_set: &mut FileSet,
+) -> Result<()> {
+    for backend in backends {
         let writer = file_set.entry(format!("metadata-{backend}.json"));
 
         let mut metadata =

kernel-builder/src/pyproject/mlx/mod.rs

@@ -0,0 +1,2 @@
+mod noarch;
+pub use noarch::write_mlx_ext_noarch;

kernel-builder/src/pyproject/mlx/noarch.rs

@@ -0,0 +1,173 @@
+use std::path::PathBuf;
+
+use eyre::{bail, Context, Result};
+use itertools::Itertools;
+use kernels_common::config::{Build, General, MlxNoarch};
+use kernels_common::metadata::Provenance;
+use minijinja::{context, Environment};
+
+use crate::pyproject::common::write_metadata_for_backends;
+use crate::pyproject::fileset::FileSet;
+use crate::pyproject::ops_identifier::KernelIdentifier;
+
+static SETUP_PY: &str = include_str!("../templates/mlx/noarch/setup.py");
+
+pub fn write_mlx_ext_noarch(
+    env: &Environment,
+    build: &Build,
+    kernel_id: &KernelIdentifier,
+    provenance: Option<&Provenance>,
+) -> Result<FileSet> {
+    let mlx_noarch = match build.framework.mlx_noarch() {
+        Some(mlx_noarch) => mlx_noarch,
+        None => bail!("Build configuration does not have `mlx-noarch` section"),
+    };
+
+    let mut file_set = FileSet::default();
+
+    write_ops_py(
+        env,
+        &build.general.name.python_name(),
+        kernel_id,
+        &mut file_set,
+    )?;
+    write_pyproject_toml(env, mlx_noarch, &build.general, &mut file_set)?;
+    write_setup_py(&mut file_set);
+    write_metadata_for_backends(
+        build,
+        &build.general.backends,
+        kernel_id,
+        provenance,
+        &mut file_set,
+    )?;
+
+    Ok(file_set)
+}
+
+fn write_ops_py(
+    env: &Environment,
+    name: &str,
+    kernel_id: &KernelIdentifier,
+    file_set: &mut FileSet,
+) -> Result<()> {
+    let mut path = PathBuf::new();
+    path.push("mlx-ext");
+    path.push(name);
+    path.push("_ops.py");
+    let writer = file_set.entry(path);
+
+    env.get_template("mlx/noarch/_ops.py")
+        .wrap_err("Cannot get MLX noarch _ops.py template")?
+        .render_captured_to(
+            context! {
+                kernel_name => kernel_id.name(),
+                kernel_unique_id => kernel_id.unique_id(),
+            },
+            writer,
+        )
+        .wrap_err("Cannot render MLX noarch _ops.py template")?;
+
+    Ok(())
+}
+
+fn write_pyproject_toml(
+    env: &Environment,
+    mlx_noarch: &MlxNoarch,
+    general: &General,
+    file_set: &mut FileSet,
+) -> Result<()> {
+    let writer = file_set.entry("pyproject.toml");
+
+    let data_globs = mlx_noarch
+        .data_extensions()
+        .map(|exts| exts.iter().map(|ext| format!("\"**/*.{ext}\"")).join(", "));
+
+    // Common python dependencies (no backend-specific ones)
+    let python_dependencies =
+        itertools::process_results(general.general_python_depends(), |iter| {
+            iter.flat_map(|(_, deps)| deps.python.iter().map(|d| format!("\"{}\"", d.pkg)))
+                .join(", ")
+        })?;
+
+    // Collect backend-specific dependencies for the kernel's backends.
+    let mut backend_dependencies = Vec::new();
+    for backend in &general.backends {
+        let deps = itertools::process_results(general.backend_python_depends(*backend), |iter| {
+            iter.flat_map(|(_, deps)| deps.python.iter().map(|d| format!("\"{}\"", d.pkg)))
+                .join(", ")
+        })?;
+
+        if !deps.is_empty() {
+            backend_dependencies.push((backend.to_string(), deps));
+        }
+    }
+
+    env.get_template("mlx/noarch/pyproject.toml")
+        .wrap_err("Cannot get MLX noarch pyproject.toml template")?
+        .render_captured_to(
+            context! {
+                data_globs => data_globs,
+                python_dependencies => python_dependencies,
+                backend_dependencies => backend_dependencies,
+                name => &general.name,
+                python_name => general.name.python_name(),
+            },
+            writer,
+        )
+        .wrap_err("Cannot render MLX noarch pyproject.toml template")?;
+
+    Ok(())
+}
+
+fn write_setup_py(file_set: &mut FileSet) {
+    file_set
+        .entry("setup.py")
+        .extend_from_slice(SETUP_PY.as_bytes());
+}
+
+#[cfg(test)]
+mod tests {
+    use std::path::PathBuf;
+
+    use kernels_common::config::Build;
+
+    use crate::pyproject::create_pyproject_file_set;
+    use crate::pyproject::ops_identifier::KernelIdentifier;
+
+    #[test]
+    fn mlx_noarch_file_set() {
+        let tmp = tempfile::tempdir().unwrap();
+        std::fs::write(
+            tmp.path().join("build.toml"),
+            r#"
+            [general]
+            name = "mlx-demo"
+            version = 1
+            edition = 6
+            license = "Apache-2.0"
+            backends = ["metal"]
+
+            [mlx-noarch]
+            pyext = ["py", "json"]
+            "#,
+        )
+        .unwrap();
+
+        let build = Build::open(tmp.path()).unwrap();
+        let kernel_id =
+            KernelIdentifier::new(tmp.path(), "mlx_demo".to_owned(), Some("abc123".to_owned()));
+        let file_set = create_pyproject_file_set(build, &kernel_id, None).unwrap();
+
+        let mut names = file_set.into_names().into_iter().collect::<Vec<_>>();
+        names.sort();
+        assert_eq!(
+            names,
+            [
+                PathBuf::from("metadata-metal.json"),
+                PathBuf::from("mlx-ext/mlx_demo/_ops.py"),
+                PathBuf::from("pyproject.toml"),
+                PathBuf::from("setup.py"),
+            ]
+        );
+    }
+}

kernel-builder/src/pyproject/mod.rs

@@ -19,6 +19,7 @@ pub(crate) mod common;
 pub mod deps;
 pub mod fileset;
 mod kernel;
+mod mlx;
 mod ops_identifier;
 mod torch;
 mod tvm_ffi;
@@ -34,12 +35,13 @@ pub fn create_pyproject_file_set(
     env.set_trim_blocks(true);
     minijinja_embed::load_templates!(&mut env);
 
-    let file_set = if matches!(build.framework, Framework::TvmFfi(_)) {
-        tvm_ffi::write_tvm_ffi_ext(&env, &build, kernel_id, provenance)?
-    } else if build.is_noarch() {
-        torch::write_torch_ext_noarch(&env, &build, kernel_id, provenance)?
-    } else {
-        torch::write_torch_ext(&env, &build, kernel_id, provenance)?
+    let file_set = match build.framework {
+        Framework::MlxNoarch(_) => mlx::write_mlx_ext_noarch(&env, &build, kernel_id, provenance)?,
+        Framework::Torch(_) => torch::write_torch_ext(&env, &build, kernel_id, provenance)?,
+        Framework::TorchNoarch(_) => {
+            torch::write_torch_ext_noarch(&env, &build, kernel_id, provenance)?
+        }
+        Framework::TvmFfi(_) => tvm_ffi::write_tvm_ffi_ext(&env, &build, kernel_id, provenance)?,
     };
 
     Ok(file_set)

kernel-builder/src/pyproject/templates/mlx/noarch/_ops.py

@@ -0,0 +1,20 @@
+def get_backend() -> str:
+    """Detect the backend by inspecting MLX."""
+    import mlx.core as mx
+
+    if mx.metal.is_available():
+        return "metal"
+    elif mx.cuda.is_available():
+        return "cuda"
+    else:
+        return "cpu"
+
+
+def _find_ops_name() -> str:
+    kernel_name = "{{ kernel_name }}"
+    unique_id = "{{ kernel_unique_id }}"
+    backend = get_backend()
+    return f"_{kernel_name}_{backend}_{unique_id}"
+
+
+_OPS_NAME = _find_ops_name()

kernel-builder/src/pyproject/templates/mlx/noarch/pyproject.toml

@@ -0,0 +1,32 @@
+[project]
+name = "{{ name }}"
+version = "0.0.1"
+requires-python = ">= 3.10"
+dependencies = [
+  "mlx>=0.32",
+  {{python_dependencies}}
+]
+
+[project.optional-dependencies]
+{% for backend_name, deps in backend_dependencies %}
+{{backend_name}} = [
+  {{ deps }}
+]
+{% endfor %}
+
+[tool.setuptools]
+package-dir = { "" = "mlx-ext" }
+
+[tool.setuptools.packages.find]
+where = ["mlx-ext"]
+include = ["{{ python_name }}*"]
+
+{% if data_globs %}
+[tool.setuptools.package-data]
+{{ python_name }} = [ {{ data_globs }} ]
+{% endif %}
+
+[tool.pytest.ini_options]
+markers = [
+  "kernels_ci: mark a test as a kernel CI test"
+]