experiment: support rust kernels

.github/workflows/build_kernel.yaml

@@ -61,6 +61,7 @@ jobs:
             relu-kernel
             relu-torch-stable-abi-kernel
             relu-tvm-ffi-kernel
+            relu-cuda-oxide-kernel
             relu-kernel-cpu
             relu-backprop-compile-kernel
             relu-triton-kernel

docs/source/builder/writing-kernels.md

@@ -352,6 +352,51 @@ are available:
 - `cxx-flags`: a list of additional flags to be passed to the C++
   compiler.
 
+### Rust kernels (experimental)
+
+A `cpu` or `cuda` kernel can be written in Rust instead of C++ by setting
+`dsl = "rust"` (the default is `dsl = "cpp"`). Include the crate `Cargo.toml`
+and the Rust sources in `src`; the library name is inferred from that
+manifest:
+
+```toml
+[kernel.my_kernel]
+backend = "cpu"
+dsl = "rust"
+depends = []
+src = ["Cargo.toml", "Cargo.lock", "src"]
+```
+
+The crate is built by cargo as a staticlib and whole-archive linked into the
+extension, so it must export the `__tvm_ffi_*` symbols itself. Rust kernels
+currently require the `[tvm-ffi]` framework, and the project root must include
+a `Cargo.lock`.
+
+Optional fields are `features`, `lib-name` as an override, and, for the `cuda`
+backend, `cuda-capabilities`.
+
+#### cuda-oxide device code
+
+`dsl = "rust"` builds host code only, so a `cuda` kernel using it has to
+supply its device code some other way (for example by embedding pregenerated
+PTX). To have the device crate compiled to PTX during the build by the
+[cuda-oxide](https://github.com/NVlabs/cuda-oxide) `rustc` codegen backend,
+use `dsl = "cuda-oxide"` and point at the device manifest:
+
+```toml
+[kernel.my_kernel]
+backend = "cuda"
+dsl = "cuda-oxide"
+depends = []
+src = ["Cargo.toml", "Cargo.lock", "src"]
+device-manifest = "kernels/Cargo.toml"
+ptx-dir = "kernels-ptx"
+```
+
+The device crate is built before the host crate, so the host crate can embed
+the PTX written to `ptx-dir` (default `kernels-ptx`). `device-manifest` is
+required by `dsl = "cuda-oxide"` and rejected by `dsl = "rust"`.
+
 ## Torch bindings
 
 ### Defining bindings

examples/kernels/flake.nix

@@ -78,6 +78,12 @@
           drv =
             sys: out: out.packages.${sys}.redistributable.${"tvm-ffi${tvmFfiVersion}-${cudaVersion}-${sys}"};
         }
+        {
+          name = "relu-cuda-oxide-kernel";
+          path = ./relu-cuda-oxide;
+          drv =
+            sys: out: out.packages.${sys}.redistributable.${"tvm-ffi${tvmFfiVersion}-${cudaVersion}-${sys}"};
+        }
         {
           name = "relu-tvm-ffi-compiler-flags-kernel";
           path = ./relu-tvm-ffi-compiler-flags;
@@ -317,6 +323,26 @@
 
       # CPU kernels to build in CI.
       ciCpuKernels = [
+        {
+          name = "relu-rust-kernel";
+          path = ./relu-rust;
+          drv =
+            sys: out:
+            let
+              variant = "tvm-ffi${tvmFfiVersion}-cpu-${sys}";
+              extension = out.packages.${sys}.redistributable.${variant};
+              ciTest = out.packages.${sys}.ciTests.${variant};
+              kernelPkgs = out.packages.${sys}.pkgs.${variant};
+            in
+            kernelPkgs.runCommand "relu-rust-kernel-test"
+              {
+                nativeBuildInputs = [ ciTest ];
+              }
+              ''
+                ${ciTest}/bin/ci-test
+                ln -s ${extension} $out
+              '';
+        }
         {
           # This test only requires a CPU, so let's run the test directly during the build.
           name = "symbol-conflicts-pytest";

examples/kernels/relu-cuda-oxide/.gitignore

@@ -0,0 +1 @@
+kernels-ptx/

examples/kernels/relu-cuda-oxide/CARD.md

@@ -0,0 +1,65 @@
+---
+library_name: kernels
+{% if license %}license: {{ license }}
+{% endif %}---
+
+This is the repository card of {{ repo_id }} that has been pushed on the Hub. It was built to be used with the [`kernels` library](https://github.com/huggingface/kernels). This card was automatically generated.
+
+## How to use
+{% if functions %}
+
+```python
+# make sure `kernels` is installed: `pip install -U kernels`
+from kernels import get_kernel
+
+# If the org / user isn't a trusted publisher, pass `trust_remote_code=True` to the
+# `get_kernel` call. You can find whether this kernel is from a trusted publisher
+# by going to the kernel's Hub page and finding the "Trusted publisher" status at
+# the top of the page.
+kernel_module = get_kernel("{{ repo_id }}", version={{ version }})
+{{ functions[0] }} = kernel_module.{{ functions[0] }}
+
+{{ functions[0] }}(...)
+```
+{% else %}
+
+Usage example not available.
+{% endif %}
+
+## Available functions
+{% if functions %}
+{% for func in functions %}
+- `{{ func }}`
+{% endfor %}
+{% else %}
+
+Function list not available.
+{% endif %}
+{% if layers %}
+
+## Available layers
+{% for layer in layers %}
+- `{{ layer }}`
+{% endfor %}
+{% endif %}
+
+## Benchmarks
+{% if has_benchmark %}
+
+Benchmarking script is available for this kernel. Run `kernels benchmark {{ repo_id }} --version {{ version }}`.
+{% else %}
+
+No benchmark available yet.
+{% endif %}
+{% if upstream %}
+
+## Upstream
+
+The original source code for this kernel comes from {{ upstream }}.
+{% endif %}
+{% if source %}
+
+## Source
+
+The kernel-builder formatted source for this kernel is available at {{ source }}.
+{% endif %}