From 76ad94db70e3acd85226de8a9dfacf06ed9dfd0c Mon Sep 17 00:00:00 2001 From: David Holtz Date: Tue, 18 Aug 2026 21:08:02 -0400 Subject: [PATCH 1/2] feat: support Rust kernels --- .github/workflows/build_kernel.yaml | 1 + docs/source/builder/writing-kernels.md | 45 ++ examples/kernels/flake.nix | 26 + examples/kernels/relu-cuda-oxide/.gitignore | 1 + examples/kernels/relu-cuda-oxide/CARD.md | 65 ++ examples/kernels/relu-cuda-oxide/Cargo.lock | 591 ++++++++++++++++++ examples/kernels/relu-cuda-oxide/Cargo.toml | 17 + examples/kernels/relu-cuda-oxide/build.toml | 19 + examples/kernels/relu-cuda-oxide/flake.nix | 17 + .../relu-cuda-oxide/kernels/Cargo.toml | 10 + .../relu-cuda-oxide/kernels/src/main.rs | 12 + .../relu-cuda-oxide/rust-toolchain.toml | 3 + examples/kernels/relu-cuda-oxide/src/lib.rs | 121 ++++ .../kernels/relu-cuda-oxide/tests/__init__.py | 0 .../relu-cuda-oxide/tests/test_relu.py | 11 + .../tvm-ffi-ext/relu_cuda_oxide/__init__.py | 15 + examples/kernels/relu-rust/CARD.md | 65 ++ examples/kernels/relu-rust/Cargo.lock | 107 ++++ examples/kernels/relu-rust/Cargo.toml | 3 + examples/kernels/relu-rust/build.toml | 17 + examples/kernels/relu-rust/flake.nix | 17 + examples/kernels/relu-rust/relu-rs/Cargo.toml | 12 + examples/kernels/relu-rust/relu-rs/src/lib.rs | 15 + examples/kernels/relu-rust/tests/__init__.py | 0 examples/kernels/relu-rust/tests/test_relu.py | 11 + .../tvm-ffi-ext/relu_rust/__init__.py | 19 + kernel-builder/src/pyproject/kernel.rs | 208 +++++- kernel-builder/src/pyproject/mod.rs | 20 +- .../templates/kernel-component/rust-cpu.cmake | 8 + .../kernel-component/rust-cuda.cmake | 11 + .../src/pyproject/templates/kernel.cmake | 105 ++++ .../pyproject/templates/tvm_ffi/binding.cmake | 5 +- .../templates/tvm_ffi/tvm-ffi-extension.cmake | 20 + kernel-builder/src/util.rs | 91 +++ kernels-data/src/config/mod.rs | 100 +++ kernels-data/src/config/v3.rs | 10 + kernels-data/src/config/v4.rs | 10 + kernels-data/src/config/v5.rs | 43 +- nix-builder/lib/extension/tvm-ffi/arch.nix | 116 +++- nix-builder/tests/Dockerfile.test-kernel | 2 + nix-builder/tests/run-tests.sh | 5 +- 41 files changed, 1962 insertions(+), 12 deletions(-) create mode 100644 examples/kernels/relu-cuda-oxide/.gitignore create mode 100644 examples/kernels/relu-cuda-oxide/CARD.md create mode 100644 examples/kernels/relu-cuda-oxide/Cargo.lock create mode 100644 examples/kernels/relu-cuda-oxide/Cargo.toml create mode 100644 examples/kernels/relu-cuda-oxide/build.toml create mode 100644 examples/kernels/relu-cuda-oxide/flake.nix create mode 100644 examples/kernels/relu-cuda-oxide/kernels/Cargo.toml create mode 100644 examples/kernels/relu-cuda-oxide/kernels/src/main.rs create mode 100644 examples/kernels/relu-cuda-oxide/rust-toolchain.toml create mode 100644 examples/kernels/relu-cuda-oxide/src/lib.rs create mode 100644 examples/kernels/relu-cuda-oxide/tests/__init__.py create mode 100644 examples/kernels/relu-cuda-oxide/tests/test_relu.py create mode 100644 examples/kernels/relu-cuda-oxide/tvm-ffi-ext/relu_cuda_oxide/__init__.py create mode 100644 examples/kernels/relu-rust/CARD.md create mode 100644 examples/kernels/relu-rust/Cargo.lock create mode 100644 examples/kernels/relu-rust/Cargo.toml create mode 100644 examples/kernels/relu-rust/build.toml create mode 100644 examples/kernels/relu-rust/flake.nix create mode 100644 examples/kernels/relu-rust/relu-rs/Cargo.toml create mode 100644 examples/kernels/relu-rust/relu-rs/src/lib.rs create mode 100644 examples/kernels/relu-rust/tests/__init__.py create mode 100644 examples/kernels/relu-rust/tests/test_relu.py create mode 100644 examples/kernels/relu-rust/tvm-ffi-ext/relu_rust/__init__.py create mode 100644 kernel-builder/src/pyproject/templates/kernel-component/rust-cpu.cmake create mode 100644 kernel-builder/src/pyproject/templates/kernel-component/rust-cuda.cmake diff --git a/.github/workflows/build_kernel.yaml b/.github/workflows/build_kernel.yaml index 1ed9ac6a..59b27516 100644 --- a/.github/workflows/build_kernel.yaml +++ b/.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 diff --git a/docs/source/builder/writing-kernels.md b/docs/source/builder/writing-kernels.md index 905ba7ed..dd537fc3 100644 --- a/docs/source/builder/writing-kernels.md +++ b/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 diff --git a/examples/kernels/flake.nix b/examples/kernels/flake.nix index 2af64fa6..469fdfe8 100644 --- a/examples/kernels/flake.nix +++ b/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"; diff --git a/examples/kernels/relu-cuda-oxide/.gitignore b/examples/kernels/relu-cuda-oxide/.gitignore new file mode 100644 index 00000000..b3a9f33e --- /dev/null +++ b/examples/kernels/relu-cuda-oxide/.gitignore @@ -0,0 +1 @@ +kernels-ptx/ diff --git a/examples/kernels/relu-cuda-oxide/CARD.md b/examples/kernels/relu-cuda-oxide/CARD.md new file mode 100644 index 00000000..2868c79f --- /dev/null +++ b/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 %} diff --git a/examples/kernels/relu-cuda-oxide/Cargo.lock b/examples/kernels/relu-cuda-oxide/Cargo.lock new file mode 100644 index 00000000..20d75b6d --- /dev/null +++ b/examples/kernels/relu-cuda-oxide/Cargo.lock @@ -0,0 +1,591 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba" +dependencies = [ + "memchr", +] + +[[package]] +name = "anyhow" +version = "1.0.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" + +[[package]] +name = "bindgen" +version = "0.71.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f58bf3d7db68cfbac37cfc485a8d711e87e064c3d0fe0435b92f7a407f9d6b3" +dependencies = [ + "bitflags", + "cexpr", + "clang-sys", + "itertools", + "log", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "rustc-hash", + "shlex", + "syn 2.0.119", +] + +[[package]] +name = "bitflags" +version = "2.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "cexpr" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6fac387a98bb7c37292057cffc56d62ecb629900026402633ae9160df93a8766" +dependencies = [ + "nom", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "clang-sys" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "157a8ba7b480713b56f4c09fd13fc3e0a22a5dfab8097ba61cbc5feef950788a" +dependencies = [ + "glob", + "libc", + "libloading", +] + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crc32fast" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "cuda-bindings" +version = "0.2.1" +source = "git+https://github.com/drbh/cuda-oxide.git?rev=bc5f33bb1d556671c35e475439bb821d25caa2cb#bc5f33bb1d556671c35e475439bb821d25caa2cb" +dependencies = [ + "bindgen", +] + +[[package]] +name = "cuda-core" +version = "0.2.1" +source = "git+https://github.com/drbh/cuda-oxide.git?rev=bc5f33bb1d556671c35e475439bb821d25caa2cb#bc5f33bb1d556671c35e475439bb821d25caa2cb" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-macros", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +source = "git+https://github.com/drbh/cuda-oxide.git?rev=bc5f33bb1d556671c35e475439bb821d25caa2cb#bc5f33bb1d556671c35e475439bb821d25caa2cb" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/drbh/cuda-oxide.git?rev=bc5f33bb1d556671c35e475439bb821d25caa2cb#bc5f33bb1d556671c35e475439bb821d25caa2cb" +dependencies = [ + "cuda-core", + "cuda-macros", + "half", + "libnvvm-sys", + "nvjitlink-sys", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +source = "git+https://github.com/drbh/cuda-oxide.git?rev=bc5f33bb1d556671c35e475439bb821d25caa2cb#bc5f33bb1d556671c35e475439bb821d25caa2cb" +dependencies = [ + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "syn 2.0.119", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "either" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e5e8f6c15a24b9a3ee5efec809ccd006d3b30e8b3bb63c39af737c7f87daa1d" + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "itertools" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +dependencies = [ + "either", +] + +[[package]] +name = "libc" +version = "0.2.189" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "libnvvm-sys" +version = "0.2.1" +source = "git+https://github.com/drbh/cuda-oxide.git?rev=bc5f33bb1d556671c35e475439bb821d25caa2cb#bc5f33bb1d556671c35e475439bb821d25caa2cb" +dependencies = [ + "libloading", + "thiserror", +] + +[[package]] +name = "log" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nvjitlink-sys" +version = "0.2.1" +source = "git+https://github.com/drbh/cuda-oxide.git?rev=bc5f33bb1d556671c35e475439bb821d25caa2cb#bc5f33bb1d556671c35e475439bb821d25caa2cb" +dependencies = [ + "libloading", + "thiserror", +] + +[[package]] +name = "object" +version = "0.36.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" +dependencies = [ + "crc32fast", + "hashbrown 0.15.5", + "indexmap", + "memchr", +] + +[[package]] +name = "oxide-artifacts" +version = "0.2.1" +source = "git+https://github.com/drbh/cuda-oxide.git?rev=bc5f33bb1d556671c35e475439bb821d25caa2cb#bc5f33bb1d556671c35e475439bb821d25caa2cb" +dependencies = [ + "object", +] + +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + +[[package]] +name = "proc-macro-error" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da25490ff9892aab3fcf7c36f08cfb902dd3e71ca0f9f9517bea02a73a5ce38c" +dependencies = [ + "proc-macro-error-attr", + "proc-macro2", + "quote", + "syn 1.0.109", + "version_check", +] + +[[package]] +name = "proc-macro-error-attr" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1be40180e52ecc98ad80b184934baf3d0d29f979574e439af5a55274b35f869" +dependencies = [ + "proc-macro2", + "quote", + "version_check", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fcfdb36bda0c880c5931cdc7a2bcdc8ba4556847b9d912bca70bc94708711ad" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "relu-cuda-oxide" +version = "0.1.0" +dependencies = [ + "cuda-core", + "tvm-ffi", +] + +[[package]] +name = "relu-kernels" +version = "0.1.0" +dependencies = [ + "cuda-device", + "cuda-host", +] + +[[package]] +name = "reserved-oxide-symbols" +version = "0.2.1" +source = "git+https://github.com/drbh/cuda-oxide.git?rev=bc5f33bb1d556671c35e475439bb821d25caa2cb#bc5f33bb1d556671c35e475439bb821d25caa2cb" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "syn" +version = "1.0.109" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "2.0.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09a43598840e33d5b0331f38c5e30d13bb11c11210a4b58f0d9b18a5a5eefcd9" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43cbfe0cf76104d42a574802844187e84a305e531ed54455f11fbde0f10541cd" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[package]] +name = "tvm-ffi" +version = "0.1.0-alpha.0" +source = "git+https://github.com/apache/tvm-ffi.git?rev=2af558e255ff2f398095835ae18e6457635b0262#2af558e255ff2f398095835ae18e6457635b0262" +dependencies = [ + "paste", + "tvm-ffi-macros", + "tvm-ffi-sys", +] + +[[package]] +name = "tvm-ffi-macros" +version = "0.1.0-alpha.0" +source = "git+https://github.com/apache/tvm-ffi.git?rev=2af558e255ff2f398095835ae18e6457635b0262#2af558e255ff2f398095835ae18e6457635b0262" +dependencies = [ + "proc-macro-error", + "proc-macro2", + "quote", + "syn 1.0.109", +] + +[[package]] +name = "tvm-ffi-sys" +version = "0.1.0-alpha.0" +source = "git+https://github.com/apache/tvm-ffi.git?rev=2af558e255ff2f398095835ae18e6457635b0262#2af558e255ff2f398095835ae18e6457635b0262" + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "zerocopy" +version = "0.8.55" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5a105cd7b140f6eeec8acff2ea38135d3cab283ada58540f629fe51e46696eb" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.55" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fe976fb70c78cd64cccfe3a6fc142244e8a77b70959b30faf9d0ac37ee228eb" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] diff --git a/examples/kernels/relu-cuda-oxide/Cargo.toml b/examples/kernels/relu-cuda-oxide/Cargo.toml new file mode 100644 index 00000000..2f258c74 --- /dev/null +++ b/examples/kernels/relu-cuda-oxide/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "relu-cuda-oxide" +version = "0.1.0" +edition = "2021" +license = "Apache-2.0" +publish = false + +[workspace] +members = ["kernels"] +default-members = ["."] + +[lib] +crate-type = ["staticlib"] + +[dependencies] +cuda-core = { git = "https://github.com/drbh/cuda-oxide.git", rev = "bc5f33bb1d556671c35e475439bb821d25caa2cb" } +tvm-ffi = { git = "https://github.com/apache/tvm-ffi.git", rev = "2af558e255ff2f398095835ae18e6457635b0262" } diff --git a/examples/kernels/relu-cuda-oxide/build.toml b/examples/kernels/relu-cuda-oxide/build.toml new file mode 100644 index 00000000..8a04ffda --- /dev/null +++ b/examples/kernels/relu-cuda-oxide/build.toml @@ -0,0 +1,19 @@ +[general] +name = "relu-cuda-oxide" +version = 1 +edition = 5 +license = "Apache-2.0" +backends = ["cuda"] + +[general.hub] +repo-id = "kernels-test/relu-cuda-oxide" + +[tvm-ffi] + +[kernel.relu_cuda_oxide] +backend = "cuda" +dsl = "cuda-oxide" +depends = [] +src = ["Cargo.toml", "Cargo.lock", "rust-toolchain.toml", "src", "kernels/Cargo.toml", "kernels/src"] +device-manifest = "kernels/Cargo.toml" +cuda-capabilities = ["8.0"] diff --git a/examples/kernels/relu-cuda-oxide/flake.nix b/examples/kernels/relu-cuda-oxide/flake.nix new file mode 100644 index 00000000..e2646593 --- /dev/null +++ b/examples/kernels/relu-cuda-oxide/flake.nix @@ -0,0 +1,17 @@ +{ + description = "Flake for a ReLU kernel whose CUDA device code is written in Rust"; + + inputs = { + kernel-builder.url = "path:../../.."; + }; + + outputs = + { + self, + kernel-builder, + }: + kernel-builder.lib.genKernelFlakeOutputs { + inherit self; + path = ./.; + }; +} diff --git a/examples/kernels/relu-cuda-oxide/kernels/Cargo.toml b/examples/kernels/relu-cuda-oxide/kernels/Cargo.toml new file mode 100644 index 00000000..f58b403c --- /dev/null +++ b/examples/kernels/relu-cuda-oxide/kernels/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "relu-kernels" +version = "0.1.0" +edition = "2024" +license = "Apache-2.0" +publish = false + +[dependencies] +cuda-device = { git = "https://github.com/drbh/cuda-oxide.git", rev = "bc5f33bb1d556671c35e475439bb821d25caa2cb" } +cuda-host = { git = "https://github.com/drbh/cuda-oxide.git", rev = "bc5f33bb1d556671c35e475439bb821d25caa2cb" } diff --git a/examples/kernels/relu-cuda-oxide/kernels/src/main.rs b/examples/kernels/relu-cuda-oxide/kernels/src/main.rs new file mode 100644 index 00000000..bd434311 --- /dev/null +++ b/examples/kernels/relu-cuda-oxide/kernels/src/main.rs @@ -0,0 +1,12 @@ +use cuda_device::{kernel, thread, DisjointSlice}; + +#[kernel] +pub fn relu_fwd(x: &[f32], mut out: DisjointSlice) { + let idx = thread::index_1d(); + let i = idx.get(); + if let Some(out_elem) = out.get_mut(idx) { + *out_elem = if x[i] > 0.0 { x[i] } else { 0.0 }; + } +} + +fn main() {} diff --git a/examples/kernels/relu-cuda-oxide/rust-toolchain.toml b/examples/kernels/relu-cuda-oxide/rust-toolchain.toml new file mode 100644 index 00000000..13204e7a --- /dev/null +++ b/examples/kernels/relu-cuda-oxide/rust-toolchain.toml @@ -0,0 +1,3 @@ +[toolchain] +channel = "nightly-2026-04-03" +components = ["rust-src", "rustc-dev", "llvm-tools"] diff --git a/examples/kernels/relu-cuda-oxide/src/lib.rs b/examples/kernels/relu-cuda-oxide/src/lib.rs new file mode 100644 index 00000000..75ad3ebd --- /dev/null +++ b/examples/kernels/relu-cuda-oxide/src/lib.rs @@ -0,0 +1,121 @@ +use std::collections::HashMap; +use std::ffi::c_void; +use std::sync::{Arc, Mutex, OnceLock}; + +use cuda_core::{launch_kernel, CudaContext, CudaFunction, CudaModule}; +use tvm_ffi::error::{Error, ErrorKind, Result, RUNTIME_ERROR, TYPE_ERROR, VALUE_ERROR}; +use tvm_ffi::tvm_ffi_sys::dlpack::{DLDataTypeCode, DLDeviceType}; +use tvm_ffi::{current_stream, Tensor}; + +const KERNELS_PTX: &str = include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/kernels-ptx/relu_kernels.ptx" +)); + +const THREADS_PER_BLOCK: u32 = 256; + +fn err(kind: ErrorKind<'_>, message: &str) -> Error { + Error::new(kind, message, "") +} + +struct Runtime { + ctx: Arc, + _module: Arc, + relu_fwd: CudaFunction, +} + +fn runtime(device_id: i32) -> Result<&'static Runtime> { + static RUNTIMES: OnceLock>> = OnceLock::new(); + let mut runtimes = RUNTIMES + .get_or_init(|| Mutex::new(HashMap::new())) + .lock() + .unwrap(); + if let Some(runtime) = runtimes.get(&device_id) { + return Ok(runtime); + } + + let ctx = CudaContext::new(device_id as usize).map_err(|e| { + err( + RUNTIME_ERROR, + &format!("cuda context (device {device_id}): {e:?}"), + ) + })?; + let module = ctx + .load_module_from_ptx_src(KERNELS_PTX) + .map_err(|e| err(RUNTIME_ERROR, &format!("load kernels PTX: {e:?}")))?; + let relu_fwd = module + .load_function("relu_fwd") + .map_err(|e| err(RUNTIME_ERROR, &format!("load relu_fwd: {e:?}")))?; + + let runtime: &'static Runtime = Box::leak(Box::new(Runtime { + ctx, + _module: module, + relu_fwd, + })); + runtimes.insert(device_id, runtime); + Ok(runtime) +} + +fn check(name: &str, tensor: &Tensor) -> Result<()> { + if tensor.device().device_type != DLDeviceType::kDLCUDA { + return Err(err(TYPE_ERROR, &format!("{name} must be a CUDA tensor"))); + } + let dtype = tensor.dtype(); + if dtype.code != DLDataTypeCode::kDLFloat as u8 || dtype.bits != 32 || dtype.lanes != 1 { + return Err(err(TYPE_ERROR, &format!("{name} must be float32"))); + } + if !tensor.is_contiguous() { + return Err(err(VALUE_ERROR, &format!("{name} must be contiguous"))); + } + Ok(()) +} + +fn relu(x: Tensor, out: Tensor) -> Result<()> { + check("x", &x)?; + check("out", &out)?; + + let device = x.device(); + if out.device().device_id != device.device_id { + return Err(err(VALUE_ERROR, "x and out must be on the same device")); + } + if x.numel() != out.numel() { + return Err(err(VALUE_ERROR, "x and out must have the same size")); + } + + let n = x.numel() as u64; + if n == 0 { + return Ok(()); + } + let runtime = runtime(device.device_id)?; + + let mut x_ptr = x.data_ptr() as u64; + let mut x_len = n; + let mut out_ptr = out.data_ptr() as u64; + let mut out_len = n; + let mut params = [ + (&mut x_ptr as *mut u64).cast::(), + (&mut x_len as *mut u64).cast(), + (&mut out_ptr as *mut u64).cast(), + (&mut out_len as *mut u64).cast(), + ]; + + runtime + .ctx + .bind_to_thread() + .map_err(|e| err(RUNTIME_ERROR, &format!("bind context: {e:?}")))?; + + // SAFETY: `params` matches `relu_fwd`'s two (pointer, length) pairs. + unsafe { + launch_kernel( + runtime.relu_fwd.cu_function(), + (n.div_ceil(THREADS_PER_BLOCK as u64) as u32, 1, 1), + (THREADS_PER_BLOCK, 1, 1), + 0, + current_stream(&device).cast(), + &mut params, + ) + } + .map_err(|e| err(RUNTIME_ERROR, &format!("launch relu_fwd: {e:?}"))) +} + +tvm_ffi::tvm_ffi_dll_export_typed_func!(relu, relu); diff --git a/examples/kernels/relu-cuda-oxide/tests/__init__.py b/examples/kernels/relu-cuda-oxide/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/examples/kernels/relu-cuda-oxide/tests/test_relu.py b/examples/kernels/relu-cuda-oxide/tests/test_relu.py new file mode 100644 index 00000000..f68ea0ff --- /dev/null +++ b/examples/kernels/relu-cuda-oxide/tests/test_relu.py @@ -0,0 +1,11 @@ +import pytest +import torch +import torch.nn.functional as F + +import relu_cuda_oxide + + +@pytest.mark.kernels_ci +def test_relu(): + x = torch.randn(1024, 1024, dtype=torch.float32, device="cuda") + torch.testing.assert_close(F.relu(x), relu_cuda_oxide.relu(x, torch.empty_like(x))) diff --git a/examples/kernels/relu-cuda-oxide/tvm-ffi-ext/relu_cuda_oxide/__init__.py b/examples/kernels/relu-cuda-oxide/tvm-ffi-ext/relu_cuda_oxide/__init__.py new file mode 100644 index 00000000..4c174b5d --- /dev/null +++ b/examples/kernels/relu-cuda-oxide/tvm-ffi-ext/relu_cuda_oxide/__init__.py @@ -0,0 +1,15 @@ +import tvm_ffi + +from ._ops import ops + + +def relu(x, out): + x_t = tvm_ffi.from_dlpack(x) + out_t = tvm_ffi.from_dlpack(out) + + device = x_t.device + if device.type != "cuda": + raise NotImplementedError(f"Unsupported device type: {device.type}") + ops.relu(x_t, out_t) + + return out diff --git a/examples/kernels/relu-rust/CARD.md b/examples/kernels/relu-rust/CARD.md new file mode 100644 index 00000000..2868c79f --- /dev/null +++ b/examples/kernels/relu-rust/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 %} diff --git a/examples/kernels/relu-rust/Cargo.lock b/examples/kernels/relu-rust/Cargo.lock new file mode 100644 index 00000000..af709f9b --- /dev/null +++ b/examples/kernels/relu-rust/Cargo.lock @@ -0,0 +1,107 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + +[[package]] +name = "proc-macro-error" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da25490ff9892aab3fcf7c36f08cfb902dd3e71ca0f9f9517bea02a73a5ce38c" +dependencies = [ + "proc-macro-error-attr", + "proc-macro2", + "quote", + "syn", + "version_check", +] + +[[package]] +name = "proc-macro-error-attr" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1be40180e52ecc98ad80b184934baf3d0d29f979574e439af5a55274b35f869" +dependencies = [ + "proc-macro2", + "quote", + "version_check", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "relu-rs" +version = "0.1.0" +dependencies = [ + "tvm-ffi", +] + +[[package]] +name = "syn" +version = "1.0.109" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "tvm-ffi" +version = "0.1.0-alpha.0" +source = "git+https://github.com/apache/tvm-ffi.git?rev=2af558e255ff2f398095835ae18e6457635b0262#2af558e255ff2f398095835ae18e6457635b0262" +dependencies = [ + "paste", + "tvm-ffi-macros", + "tvm-ffi-sys", +] + +[[package]] +name = "tvm-ffi-macros" +version = "0.1.0-alpha.0" +source = "git+https://github.com/apache/tvm-ffi.git?rev=2af558e255ff2f398095835ae18e6457635b0262#2af558e255ff2f398095835ae18e6457635b0262" +dependencies = [ + "proc-macro-error", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tvm-ffi-sys" +version = "0.1.0-alpha.0" +source = "git+https://github.com/apache/tvm-ffi.git?rev=2af558e255ff2f398095835ae18e6457635b0262#2af558e255ff2f398095835ae18e6457635b0262" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" diff --git a/examples/kernels/relu-rust/Cargo.toml b/examples/kernels/relu-rust/Cargo.toml new file mode 100644 index 00000000..f79b1ef2 --- /dev/null +++ b/examples/kernels/relu-rust/Cargo.toml @@ -0,0 +1,3 @@ +[workspace] +members = ["relu-rs"] +resolver = "2" diff --git a/examples/kernels/relu-rust/build.toml b/examples/kernels/relu-rust/build.toml new file mode 100644 index 00000000..62759009 --- /dev/null +++ b/examples/kernels/relu-rust/build.toml @@ -0,0 +1,17 @@ +[general] +name = "relu-rust" +version = 1 +edition = 5 +license = "Apache-2.0" +backends = ["cpu"] + +[general.hub] +repo-id = "kernels-test/relu-rust" + +[tvm-ffi] + +[kernel.relu_rust] +backend = "cpu" +dsl = "rust" +depends = [] +src = ["relu-rs/Cargo.toml", "relu-rs/src", "Cargo.toml", "Cargo.lock"] diff --git a/examples/kernels/relu-rust/flake.nix b/examples/kernels/relu-rust/flake.nix new file mode 100644 index 00000000..b8ec0be9 --- /dev/null +++ b/examples/kernels/relu-rust/flake.nix @@ -0,0 +1,17 @@ +{ + description = "Flake for a ReLU kernel written in Rust"; + + inputs = { + kernel-builder.url = "path:../../.."; + }; + + outputs = + { + self, + kernel-builder, + }: + kernel-builder.lib.genKernelFlakeOutputs { + inherit self; + path = ./.; + }; +} diff --git a/examples/kernels/relu-rust/relu-rs/Cargo.toml b/examples/kernels/relu-rust/relu-rs/Cargo.toml new file mode 100644 index 00000000..055636ec --- /dev/null +++ b/examples/kernels/relu-rust/relu-rs/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "relu-rs" +version = "0.1.0" +edition = "2021" +license = "Apache-2.0" +publish = false + +[lib] +crate-type = ["staticlib"] + +[dependencies] +tvm-ffi = { git = "https://github.com/apache/tvm-ffi.git", rev = "2af558e255ff2f398095835ae18e6457635b0262" } diff --git a/examples/kernels/relu-rust/relu-rs/src/lib.rs b/examples/kernels/relu-rust/relu-rs/src/lib.rs new file mode 100644 index 00000000..2c246827 --- /dev/null +++ b/examples/kernels/relu-rust/relu-rs/src/lib.rs @@ -0,0 +1,15 @@ +use tvm_ffi::{Result, Tensor}; + +/// ReLU on contiguous float32 CPU tensors. +fn relu_rust(x: Tensor, out: Tensor) -> Result<()> { + let x_data = x.data_as_slice::()?; + let out_data = out.data_as_slice_mut::()?; + + for (out_elem, &x_elem) in out_data.iter_mut().zip(x_data.iter()) { + *out_elem = x_elem.max(0.0); + } + + Ok(()) +} + +tvm_ffi::tvm_ffi_dll_export_typed_func!(relu_rust, relu_rust); diff --git a/examples/kernels/relu-rust/tests/__init__.py b/examples/kernels/relu-rust/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/examples/kernels/relu-rust/tests/test_relu.py b/examples/kernels/relu-rust/tests/test_relu.py new file mode 100644 index 00000000..c870279e --- /dev/null +++ b/examples/kernels/relu-rust/tests/test_relu.py @@ -0,0 +1,11 @@ +import pytest +import torch +import torch.nn.functional as F + +import relu_rust + + +@pytest.mark.kernels_ci +def test_relu(): + x = torch.randn(1024, 1024, dtype=torch.float32, device="cpu") + torch.testing.assert_close(F.relu(x), relu_rust.relu(x, torch.empty_like(x))) diff --git a/examples/kernels/relu-rust/tvm-ffi-ext/relu_rust/__init__.py b/examples/kernels/relu-rust/tvm-ffi-ext/relu_rust/__init__.py new file mode 100644 index 00000000..69f606da --- /dev/null +++ b/examples/kernels/relu-rust/tvm-ffi-ext/relu_rust/__init__.py @@ -0,0 +1,19 @@ +import tvm_ffi + +from ._ops import ops + + +def relu(x, out): + x_t = tvm_ffi.from_dlpack(x) + out_t = tvm_ffi.from_dlpack(out) + + device = x_t.device + if device.type == "cpu": + ops.relu_rust(x_t, out_t) + else: + raise NotImplementedError(f"Unsupported device type: {device.type}") + + return out + + +__all__ = ["relu"] diff --git a/kernel-builder/src/pyproject/kernel.rs b/kernel-builder/src/pyproject/kernel.rs index 6bf78f50..a0d4d8dd 100644 --- a/kernel-builder/src/pyproject/kernel.rs +++ b/kernel-builder/src/pyproject/kernel.rs @@ -1,8 +1,8 @@ -use std::io::Write; +use std::{io::Write, path::Path}; -use eyre::{Context, Result}; +use eyre::{ensure, Context, Result}; use itertools::Itertools; -use kernels_data::config::{Build, Kernel}; +use kernels_data::config::{Build, Dsl, Kernel}; use minijinja::{context, Environment}; use crate::pyproject::common::prefix_and_join_includes; @@ -12,7 +12,7 @@ pub fn render_kernel_components( build: &Build, write: &mut impl Write, ) -> Result<()> { - for (kernel_name, kernel) in build.kernels.iter() { + for (kernel_name, kernel) in build.kernels.iter().sorted_by(|(a, _), (b, _)| a.cmp(b)) { render_kernel_component(env, kernel_name, kernel, write)?; } @@ -34,6 +34,9 @@ fn render_kernel_component( .join("\n"); match kernel { + Kernel::Cpu { .. } | Kernel::Cuda { .. } if kernel.dsl().is_cargo_built() => { + render_kernel_component_rust(env, kernel_name, kernel, write)? + } Kernel::Cpu { .. } => { render_kernel_component_cpu(env, kernel_name, kernel, sources, write)? } @@ -54,6 +57,128 @@ fn render_kernel_component( Ok(()) } +/// Rust kernels are built by cargo into a staticlib that is whole-archive +/// linked into the extension, so unlike the C++ DSL they contribute no sources +/// to the CMake build and render a component of their own. +fn render_kernel_component_rust( + env: &Environment, + kernel_name: &str, + kernel: &Kernel, + write: &mut impl Write, +) -> Result<()> { + let dsl = kernel.dsl(); + let (template, src, lib_name, features, device_manifest, ptx_dir, cuda_capabilities) = + match kernel { + Kernel::Cpu { + src, + lib_name, + features, + .. + } => { + ensure!( + dsl != Dsl::CudaOxide, + "Kernel `{kernel_name}`: `dsl = \"cuda-oxide\"` requires \ + `backend = \"cuda\"`" + ); + ( + "kernel-component/rust-cpu.cmake", + src, + lib_name, + features, + None, + None, + None, + ) + } + Kernel::Cuda { + src, + lib_name, + features, + device_manifest, + ptx_dir, + cuda_capabilities, + .. + } => ( + "kernel-component/rust-cuda.cmake", + src, + lib_name, + features, + device_manifest.as_deref(), + ptx_dir.as_deref(), + cuda_capabilities.as_ref(), + ), + _ => { + unreachable!("cargo-built DSLs are only supported for the cpu and cuda backends") + } + }; + + // `device-manifest`/`ptx-dir` drive the cuda-oxide device build; without + // the cuda-oxide DSL nothing would consume them, and without them the + // cuda-oxide DSL would silently build no device code at all. + if dsl == Dsl::CudaOxide { + ensure!( + device_manifest.is_some(), + "Kernel `{kernel_name}`: `dsl = \"cuda-oxide\"` requires `device-manifest`" + ); + } else { + ensure!( + device_manifest.is_none() && ptx_dir.is_none(), + "Kernel `{kernel_name}`: `device-manifest` and `ptx-dir` require \ + `dsl = \"cuda-oxide\"`" + ); + } + + let manifest_path = rust_manifest_src(kernel_name, src)?; + + let lib_name = rust_lib_name(manifest_path, lib_name).ok_or_else(|| { + eyre::eyre!( + "Rust kernel `{kernel_name}`: cannot derive `lib-name` from \ + `src = [\"{}\"]`, set it explicitly", + manifest_path + ) + })?; + + env.get_template(template) + .wrap_err("Cannot get Rust kernel template")? + .render_captured_to( + context! { + cuda_capabilities => cuda_capabilities, + device_manifest => device_manifest, + features => features, + lib_name => lib_name, + manifest_path => manifest_path, + name => kernel_name, + ptx_dir => ptx_dir, + }, + &mut *write, + ) + .wrap_err("Cannot render Rust kernel template")?; + + write.write_all(b"\n")?; + + Ok(()) +} + +fn rust_manifest_src<'a>(kernel_name: &str, src: &'a [String]) -> Result<&'a str> { + src.iter() + .find(|path| { + Path::new(path.as_str()) + .file_name() + .is_some_and(|name| name == "Cargo.toml") + }) + .map(String::as_str) + .ok_or_else(|| eyre::eyre!("Rust kernel `{kernel_name}`: `src` must include Cargo.toml")) +} + +fn rust_lib_name(manifest_path: &str, lib_name: &Option) -> Option { + lib_name.clone().or_else(|| { + Path::new(manifest_path) + .parent() + .and_then(Path::file_name) + .map(|name| name.to_string_lossy().replace('-', "_")) + }) +} + fn render_kernel_component_cpu( env: &Environment, kernel_name: &str, @@ -213,3 +338,78 @@ fn render_kernel_component_xpu( Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + + fn render(kernel: Kernel) -> Result { + let mut env = Environment::new(); + minijinja_embed::load_templates!(&mut env); + let mut output = Vec::new(); + render_kernel_component(&env, "test_kernel", &kernel, &mut output)?; + Ok(String::from_utf8(output).unwrap()) + } + + fn rust_cpu(dsl: Dsl) -> Kernel { + Kernel::Cpu { + cxx_flags: None, + depends: Vec::new(), + dsl: Some(dsl), + features: Some(vec!["simd".to_owned(), "ffi".to_owned()]), + include: None, + lib_name: Some("test_kernel_lib".to_owned()), + src: vec!["rust-kernel/Cargo.toml".to_owned()], + } + } + + fn cuda_oxide(device_manifest: Option<&str>) -> Kernel { + Kernel::Cuda { + cuda_capabilities: Some(vec!["8.0".to_owned(), "9.0".to_owned()]), + cuda_flags: None, + cuda_minver: None, + cxx_flags: None, + depends: Vec::new(), + device_manifest: device_manifest.map(str::to_owned), + dsl: Some(Dsl::CudaOxide), + features: Some(vec!["host".to_owned()]), + include: None, + lib_name: Some("test_cuda_lib".to_owned()), + ptx_dir: Some("generated-ptx".to_owned()), + src: vec!["Cargo.toml".to_owned()], + } + } + + #[test] + fn renders_rust_cpu_cmake_component() { + let output = render(rust_cpu(Dsl::Rust)).unwrap(); + + assert!(output.contains("rust_kernel_component(RUST_KERNEL_LIBS RUST_KERNEL_TARGETS")); + assert!(output.contains("MANIFEST_PATH \"rust-kernel/Cargo.toml\"")); + assert!(output.contains("LIB_NAME test_kernel_lib")); + assert!(output.contains("FEATURES simd ffi")); + } + + #[test] + fn renders_cuda_oxide_cmake_component() { + let output = render(cuda_oxide(Some("device/Cargo.toml"))).unwrap(); + + assert!(output.contains("DEVICE_MANIFEST \"device/Cargo.toml\"")); + assert!(output.contains("PTX_DIR \"generated-ptx\"")); + assert!(output.contains("CUDA_CAPABILITIES 8.0 9.0")); + } + + #[test] + fn rejects_cuda_oxide_on_cpu() { + let error = render(rust_cpu(Dsl::CudaOxide)).unwrap_err(); + + assert!(format!("{error:#}").contains("requires `backend = \"cuda\"`")); + } + + #[test] + fn rejects_cuda_oxide_without_device_manifest() { + let error = render(cuda_oxide(None)).unwrap_err(); + + assert!(format!("{error:#}").contains("requires `device-manifest`")); + } +} diff --git a/kernel-builder/src/pyproject/mod.rs b/kernel-builder/src/pyproject/mod.rs index cc8cb25f..aeb0c930 100644 --- a/kernel-builder/src/pyproject/mod.rs +++ b/kernel-builder/src/pyproject/mod.rs @@ -12,7 +12,7 @@ use minijinja::Environment; use crate::{ pyproject::ops_identifier::KernelIdentifier, - util::{check_or_infer_kernel_dir, check_or_infer_target_dir}, + util::{check_or_infer_kernel_dir, check_or_infer_target_dir, parse_build}, }; pub(crate) mod common; @@ -34,6 +34,20 @@ pub fn create_pyproject_file_set( env.set_trim_blocks(true); minijinja_embed::load_templates!(&mut env); + // Cargo-built kernels are linked as a staticlib into the extension, which + // is currently only wired up for the `tvm-ffi` framework. + if build + .kernels + .values() + .any(|kernel| kernel.dsl().is_cargo_built()) + && !matches!(build.framework, Framework::TvmFfi(_)) + { + bail!( + "`dsl = \"rust\"` and `dsl = \"cuda-oxide\"` kernels are currently \ + only supported together with the `[tvm-ffi]` framework" + ); + } + 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() { @@ -55,7 +69,7 @@ pub fn create_pyproject( ) -> Result<()> { let kernel_dir = check_or_infer_kernel_dir(kernel_dir)?; let target_dir = check_or_infer_target_dir(&kernel_dir, target_dir)?; - let build = Build::open(&kernel_dir)?; + let build = parse_build(&kernel_dir)?; // Assemble build provenance. Prefer an explicitly provided kernel git // provenance (e.g. passed by Nix builds, where the source tree has no @@ -89,7 +103,7 @@ pub fn clean_pyproject( ) -> Result<()> { let kernel_dir = check_or_infer_kernel_dir(kernel_dir)?; let target_dir = check_or_infer_target_dir(&kernel_dir, target_dir)?; - let build = Build::open(&kernel_dir)?; + let build = parse_build(&kernel_dir)?; // Provenance is irrelevant when computing the set of files to clean. let kernel_id = KernelIdentifier::new(&kernel_dir, build.general.name.python_name(), unique_id); diff --git a/kernel-builder/src/pyproject/templates/kernel-component/rust-cpu.cmake b/kernel-builder/src/pyproject/templates/kernel-component/rust-cpu.cmake new file mode 100644 index 00000000..b5d5e242 --- /dev/null +++ b/kernel-builder/src/pyproject/templates/kernel-component/rust-cpu.cmake @@ -0,0 +1,8 @@ +if(GPU_LANG STREQUAL "CPU") +rust_kernel_component(RUST_KERNEL_LIBS RUST_KERNEL_TARGETS + NAME {{ name }} + MANIFEST_PATH "{{ manifest_path }}" + LIB_NAME {{ lib_name }} + {% if features %}FEATURES {{ features|join(" ") }}{% endif %} +) +endif() diff --git a/kernel-builder/src/pyproject/templates/kernel-component/rust-cuda.cmake b/kernel-builder/src/pyproject/templates/kernel-component/rust-cuda.cmake new file mode 100644 index 00000000..db5d6d99 --- /dev/null +++ b/kernel-builder/src/pyproject/templates/kernel-component/rust-cuda.cmake @@ -0,0 +1,11 @@ +if(GPU_LANG STREQUAL "CUDA") +rust_kernel_component(RUST_KERNEL_LIBS RUST_KERNEL_TARGETS + NAME {{ name }} + MANIFEST_PATH "{{ manifest_path }}" + LIB_NAME {{ lib_name }} + {% if features %}FEATURES {{ features|join(" ") }}{% endif %} + {% if device_manifest %}DEVICE_MANIFEST "{{ device_manifest }}"{% endif %} + {% if ptx_dir %}PTX_DIR "{{ ptx_dir }}"{% endif %} + {% if cuda_capabilities %}CUDA_CAPABILITIES {{ cuda_capabilities|join(" ") }}{% endif %} +) +endif() diff --git a/kernel-builder/src/pyproject/templates/kernel.cmake b/kernel-builder/src/pyproject/templates/kernel.cmake index c41aaa65..beaed581 100644 --- a/kernel-builder/src/pyproject/templates/kernel.cmake +++ b/kernel-builder/src/pyproject/templates/kernel.cmake @@ -315,3 +315,108 @@ function(metal_kernel_component SRC_VAR) set(METAL_INCLUDE_DIRS ${_TMP_METAL_INCLUDES} PARENT_SCOPE) endif() endfunction() + +function(rust_kernel_component LIBS_VAR TARGETS_VAR) + set(oneValueArgs NAME MANIFEST_PATH LIB_NAME DEVICE_MANIFEST PTX_DIR) + set(multiValueArgs FEATURES CUDA_CAPABILITIES) + cmake_parse_arguments(KERNEL "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) + + if(NOT KERNEL_NAME OR NOT KERNEL_MANIFEST_PATH OR NOT KERNEL_LIB_NAME) + message(FATAL_ERROR "rust_kernel_component requires NAME, MANIFEST_PATH, and LIB_NAME") + endif() + + find_program(CARGO_EXECUTABLE cargo REQUIRED) + + # Export CUDA archs for Rust build scripts. + set(_KERNEL_ARCHS "") + if(GPU_LANG STREQUAL "CUDA") + if(KERNEL_CUDA_CAPABILITIES) + cuda_archs_loose_intersection(_KERNEL_ARCHS "${KERNEL_CUDA_CAPABILITIES}" "${CUDA_ARCHS}") + if(NOT _KERNEL_ARCHS) + message(FATAL_ERROR "Rust kernel: ${KERNEL_NAME}, empty set of capabilities after intersection (kernel: ${KERNEL_CUDA_CAPABILITIES}, supported: ${CUDA_ARCHS})") + endif() + else() + set(_KERNEL_ARCHS "${CUDA_KERNEL_ARCHS}") + endif() + message(STATUS "Rust kernel: ${KERNEL_NAME}, capabilities: ${_KERNEL_ARCHS}") + + accumulate_gpu_archs(_ALL_GPU_ARCHS "${ALL_GPU_ARCHS}" "${_KERNEL_ARCHS}") + set(ALL_GPU_ARCHS ${_ALL_GPU_ARCHS} PARENT_SCOPE) + endif() + + # Cargo writes the staticlib into its target directory. + set(_CARGO_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/${KERNEL_NAME}) + set(_STATICLIB ${_CARGO_TARGET_DIR}/release/${CMAKE_STATIC_LIBRARY_PREFIX}${KERNEL_LIB_NAME}${CMAKE_STATIC_LIBRARY_SUFFIX}) + set(_RUST_KERNEL_DEPENDS) + + if(KERNEL_DEVICE_MANIFEST) + # Build cuda-oxide device code before the host crate embeds its PTX. + if(NOT DEFINED ENV{CUDA_OXIDE_BACKEND} OR "$ENV{CUDA_OXIDE_BACKEND}" STREQUAL "") + message(FATAL_ERROR "rust_kernel_component: DEVICE_MANIFEST requires CUDA_OXIDE_BACKEND") + endif() + if(NOT KERNEL_PTX_DIR) + set(KERNEL_PTX_DIR kernels-ptx) + endif() + + get_filename_component(_DEVICE_MANIFEST ${KERNEL_DEVICE_MANIFEST} ABSOLUTE BASE_DIR ${CMAKE_CURRENT_SOURCE_DIR}) + get_filename_component(_PTX_DIR ${KERNEL_PTX_DIR} ABSOLUTE BASE_DIR ${CMAKE_CURRENT_SOURCE_DIR}) + set(_DEVICE_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/${KERNEL_NAME}-device) + set(_DEVICE_ENV) + + if(_KERNEL_ARCHS) + list(SORT _KERNEL_ARCHS COMPARE NATURAL) + # PTX JITs forward, so target the lowest supported arch. + list(GET _KERNEL_ARCHS 0 _OXIDE_ARCH) + string(REPLACE "+PTX" "" _OXIDE_ARCH "${_OXIDE_ARCH}") + string(REPLACE "." "" _OXIDE_ARCH "${_OXIDE_ARCH}") + list(APPEND _DEVICE_ENV "CUDA_OXIDE_TARGET=sm_${_OXIDE_ARCH}") + endif() + + add_custom_target(${KERNEL_NAME}_oxide_device_build ALL + COMMAND ${CMAKE_COMMAND} -E make_directory ${_PTX_DIR} + COMMAND ${CMAKE_COMMAND} -E env + "CUDA_OXIDE_PTX_DIR=${_PTX_DIR}" + "RUSTFLAGS=-Zcodegen-backend=$ENV{CUDA_OXIDE_BACKEND} -Copt-level=3 -Cdebug-assertions=off -Zmir-enable-passes=-JumpThreading -Csymbol-mangling-version=v0" + ${_DEVICE_ENV} + ${CARGO_EXECUTABLE} build --release --locked + --manifest-path ${_DEVICE_MANIFEST} + --target-dir ${_DEVICE_TARGET_DIR} + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} + COMMENT "Building Rust CUDA device kernel ${KERNEL_NAME}" + VERBATIM + ) + list(APPEND _RUST_KERNEL_DEPENDS ${KERNEL_NAME}_oxide_device_build) + endif() + + set(_CARGO_ARGS rustc --release --locked --lib --crate-type staticlib + --manifest-path ${CMAKE_CURRENT_SOURCE_DIR}/${KERNEL_MANIFEST_PATH} + --target-dir ${_CARGO_TARGET_DIR}) + if(KERNEL_FEATURES) + list(JOIN KERNEL_FEATURES "," _KERNEL_FEATURES) + list(APPEND _CARGO_ARGS --features ${_KERNEL_FEATURES}) + endif() + + get_filename_component(_PYTHON_BIN_DIR ${Python_EXECUTABLE} DIRECTORY) + + add_custom_target(${KERNEL_NAME}_cargo_build ALL + COMMAND ${CMAKE_COMMAND} -E env + "PATH=${_PYTHON_BIN_DIR}:$ENV{PATH}" + "KERNEL_BUILDER_GPU_LANG=${GPU_LANG}" + "KERNEL_BUILDER_CUDA_ARCHS=${_KERNEL_ARCHS}" + ${CARGO_EXECUTABLE} ${_CARGO_ARGS} + BYPRODUCTS ${_STATICLIB} + DEPENDS ${_RUST_KERNEL_DEPENDS} + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} + COMMENT "Building Rust kernel ${KERNEL_NAME} with cargo" + VERBATIM + ) + + # Register the cargo artifact as a CMake library target. + add_library(${KERNEL_NAME}_rust STATIC IMPORTED GLOBAL) + set_target_properties(${KERNEL_NAME}_rust PROPERTIES + IMPORTED_LOCATION ${_STATICLIB}) + + # Return the library and build target to the extension scope. + set(${LIBS_VAR} ${${LIBS_VAR}} ${KERNEL_NAME}_rust PARENT_SCOPE) + set(${TARGETS_VAR} ${${TARGETS_VAR}} ${KERNEL_NAME}_cargo_build PARENT_SCOPE) +endfunction() diff --git a/kernel-builder/src/pyproject/templates/tvm_ffi/binding.cmake b/kernel-builder/src/pyproject/templates/tvm_ffi/binding.cmake index 8c89d5fb..01564c8c 100644 --- a/kernel-builder/src/pyproject/templates/tvm_ffi/binding.cmake +++ b/kernel-builder/src/pyproject/templates/tvm_ffi/binding.cmake @@ -2,6 +2,8 @@ set(TVM_FFI_{{name}}_SRC {{ src|join(' ') }} ) +# only append binding sources if the source list is non-empty +if(TVM_FFI_{{name}}_SRC) {% if includes %} # TODO: check if CLion support this: # https://youtrack.jetbrains.com/issue/CPP-16510/CLion-does-not-handle-per-file-include-directories @@ -17,4 +19,5 @@ set_property( COMPILE_OPTIONS "$<$:{{ cxx_flags }}>") {% endif %} -list(APPEND SRC {{'"${TVM_FFI_' + name + '_SRC}"'}}) +list(APPEND SRC {{'${TVM_FFI_' + name + '_SRC}'}}) +endif() diff --git a/kernel-builder/src/pyproject/templates/tvm_ffi/tvm-ffi-extension.cmake b/kernel-builder/src/pyproject/templates/tvm_ffi/tvm-ffi-extension.cmake index 79a1600a..a5f77395 100644 --- a/kernel-builder/src/pyproject/templates/tvm_ffi/tvm-ffi-extension.cmake +++ b/kernel-builder/src/pyproject/templates/tvm_ffi/tvm-ffi-extension.cmake @@ -1,6 +1,12 @@ # Avoid 'lib' prefix for the extension. set(CMAKE_SHARED_LIBRARY_PREFIX "") +# rust kernels export the symbols; CMake just needs a stub source for the shared library. +if(NOT SRC AND RUST_KERNEL_LIBS) + file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/_ops_stub.cpp "\n") + list(APPEND SRC ${CMAKE_CURRENT_BINARY_DIR}/_ops_stub.cpp) +endif() + add_library(${OPS_NAME} SHARED ${SRC}) target_compile_definitions(${OPS_NAME} PRIVATE "-DTVM_FFI_EXTENSION_NAME=${OPS_NAME}") @@ -15,6 +21,20 @@ if(CXX_HAS_NO_GNU_UNIQUE) target_compile_options(${OPS_NAME} PRIVATE $<$:-fno-gnu-unique>) endif() +if(RUST_KERNEL_LIBS) + add_dependencies(${OPS_NAME} ${RUST_KERNEL_TARGETS}) + find_package(Threads REQUIRED) + target_link_libraries(${OPS_NAME} PRIVATE + "$" + Threads::Threads + ${CMAKE_DL_LIBS}) + + if(GPU_LANG STREQUAL "CUDA") + find_package(CUDAToolkit REQUIRED) + target_link_libraries(${OPS_NAME} PRIVATE CUDA::cuda_driver) + endif() +endif() + if(GPU_LANG STREQUAL "SYCL") target_link_options(${OPS_NAME} PRIVATE ${sycl_link_flags}) target_link_libraries(${OPS_NAME} PRIVATE dnnl) diff --git a/kernel-builder/src/util.rs b/kernel-builder/src/util.rs index c474fdb3..a2bd93a3 100644 --- a/kernel-builder/src/util.rs +++ b/kernel-builder/src/util.rs @@ -4,6 +4,15 @@ use std::path::{Path, PathBuf}; use eyre::{bail, ensure, Context, Result}; +use kernels_data::config::{Build, Kernel}; + +pub(crate) fn parse_build(kernel_dir: impl AsRef) -> Result { + let kernel_dir = kernel_dir.as_ref(); + let mut build = Build::open(kernel_dir)?; + infer_rust_kernel_lib_names(&mut build, kernel_dir)?; + Ok(build) +} + pub(crate) fn check_or_infer_kernel_dir(kernel_dir: Option>) -> Result { match kernel_dir { Some(kernel_dir) => { @@ -37,6 +46,46 @@ pub(crate) fn check_or_infer_target_dir( } } +fn infer_rust_kernel_lib_names(build: &mut Build, kernel_dir: &Path) -> Result<()> { + for (kernel_name, kernel) in &mut build.kernels { + if !kernel.dsl().is_cargo_built() { + continue; + } + let (src, lib_name) = match kernel { + Kernel::Cpu { src, lib_name, .. } | Kernel::Cuda { src, lib_name, .. } => { + (src, lib_name) + } + _ => continue, + }; + if lib_name.is_none() { + let manifest_path = src + .iter() + .find(|path| { + Path::new(path.as_str()) + .file_name() + .is_some_and(|name| name == "Cargo.toml") + }) + .map(|path| kernel_dir.join(path)) + .ok_or_else(|| { + eyre::eyre!( + "Rust kernel `{kernel_name}`: `src` must include the crate Cargo.toml" + ) + })?; + let manifest: toml::Value = toml::from_str(&fs::read_to_string(&manifest_path)?)?; + let name = |table| manifest.get(table)?.get("name")?.as_str(); + *lib_name = Some( + name("lib") + .or_else(|| name("package")) + .ok_or_else(|| { + eyre::eyre!("cannot infer `lib-name` from `{}`", manifest_path.display()) + })? + .replace('-', "_"), + ); + } + } + Ok(()) +} + /// Discover build variant directories (contain `metadata.json`). /// Checks `result` symlink (Nix store output) first, then falls back to `build/`. pub(crate) fn discover_variants(kernel_dir: &Path) -> Result<(PathBuf, Vec)> { @@ -76,6 +125,48 @@ pub(crate) fn discover_variants(kernel_dir: &Path) -> Result<(PathBuf, Vec>, depends: Vec, + dsl: Option, + features: Option>, include: Option>, + lib_name: Option, src: Vec, }, Cuda { @@ -330,7 +333,12 @@ pub enum Kernel { cuda_minver: Option, cxx_flags: Option>, depends: Vec, + device_manifest: Option, + dsl: Option, + features: Option>, include: Option>, + lib_name: Option, + ptx_dir: Option, src: Vec, }, Metal { @@ -356,6 +364,33 @@ pub enum Kernel { }, } +/// The language a kernel's sources are written in. +/// +/// A kernel's `backend` says what hardware it is compiled for, `dsl` says how +/// its sources are written. Backends accept the DSLs they can lower. +#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)] +#[serde(deny_unknown_fields, rename_all = "kebab-case")] +pub enum Dsl { + /// C++ sources in the backend's dialect (CUDA C++, HIP, SYCL, ...), + /// compiled directly by the backend's compiler. + Cpp, + /// Rust sources, built by `cargo` into a staticlib that is linked into the + /// extension. + Rust, + /// Rust sources whose device code is compiled to PTX by the `cuda-oxide` + /// `rustc` codegen backend before the host crate is built. Requires the + /// `cuda` backend and a `device-manifest`. + CudaOxide, +} + +impl Dsl { + /// Whether the kernel's sources are built by `cargo` rather than by the + /// backend's C++ toolchain. + pub fn is_cargo_built(&self) -> bool { + matches!(self, Dsl::Rust | Dsl::CudaOxide) + } +} + impl Kernel { pub fn cxx_flags(&self) -> Option<&[String]> { match self { @@ -387,6 +422,14 @@ impl Kernel { } } + /// The DSL the kernel's sources are written in. + pub fn dsl(&self) -> Dsl { + match self { + Kernel::Cpu { dsl, .. } | Kernel::Cuda { dsl, .. } => dsl.unwrap_or(Dsl::Cpp), + _ => Dsl::Cpp, + } + } + pub fn depends(&self) -> &[Dependency] { match self { Kernel::Cpu { depends, .. } @@ -538,4 +581,61 @@ mod tests { .is_none() ); } + + #[test] + fn v5_rust_dsls_round_trip() { + let config = r#" +[general] +name = "rust-dsls" +version = 1 +edition = 5 +license = "Apache-2.0" +backends = ["cpu", "cuda"] + +[tvm-ffi] + +[kernel.cpu_kernel] +backend = "cpu" +dsl = "rust" +depends = [] +src = ["cpu/Cargo.toml"] + +[kernel.cuda_kernel] +backend = "cuda" +dsl = "cuda-oxide" +depends = [] +src = ["Cargo.toml"] +device-manifest = "device/Cargo.toml" +"#; + + let parsed: v5::Build = toml::from_str(config).unwrap(); + let serialized = toml::to_string(&parsed).unwrap(); + let build: Build = toml::from_str::(&serialized).unwrap().into(); + + assert_eq!(build.kernels["cpu_kernel"].dsl(), Dsl::Rust); + assert_eq!(build.kernels["cuda_kernel"].dsl(), Dsl::CudaOxide); + } + + #[test] + fn v5_missing_dsl_defaults_to_cpp() { + let config = r#" +[general] +name = "cpp-default" +version = 1 +edition = 5 +license = "Apache-2.0" +backends = ["cpu"] + +[tvm-ffi] + +[kernel.cpp_kernel] +backend = "cpu" +depends = [] +src = ["kernel.cpp"] +"#; + + let build: Build = toml::from_str::(config).unwrap().into(); + + assert_eq!(build.kernels["cpp_kernel"].dsl(), Dsl::Cpp); + } } diff --git a/kernels-data/src/config/v3.rs b/kernels-data/src/config/v3.rs index a77fe7bd..1d072f31 100644 --- a/kernels-data/src/config/v3.rs +++ b/kernels-data/src/config/v3.rs @@ -302,7 +302,11 @@ impl From for super::Kernel { } => super::Kernel::Cpu { cxx_flags, depends, + // Editions before 5 have no `dsl`; they are always C++. + dsl: None, + features: None, include, + lib_name: None, src, }, Kernel::Cuda { @@ -319,7 +323,13 @@ impl From for super::Kernel { cuda_minver, cxx_flags, depends, + device_manifest: None, + // Editions before 5 have no `dsl`; they are always C++. + dsl: None, + features: None, include, + lib_name: None, + ptx_dir: None, src, }, Kernel::Metal { diff --git a/kernels-data/src/config/v4.rs b/kernels-data/src/config/v4.rs index 418e9add..a36e191c 100644 --- a/kernels-data/src/config/v4.rs +++ b/kernels-data/src/config/v4.rs @@ -323,7 +323,11 @@ impl From for super::Kernel { } => super::Kernel::Cpu { cxx_flags, depends, + // Editions before 5 have no `dsl`; they are always C++. + dsl: None, + features: None, include, + lib_name: None, src, }, Kernel::Cuda { @@ -340,7 +344,13 @@ impl From for super::Kernel { cuda_minver, cxx_flags, depends, + device_manifest: None, + // Editions before 5 have no `dsl`; they are always C++. + dsl: None, + features: None, include, + lib_name: None, + ptx_dir: None, src, }, Kernel::Metal { diff --git a/kernels-data/src/config/v5.rs b/kernels-data/src/config/v5.rs index 314289a8..cee765cf 100644 --- a/kernels-data/src/config/v5.rs +++ b/kernels-data/src/config/v5.rs @@ -4,7 +4,7 @@ use std::path::PathBuf; use monostate::MustBe; use serde::{Deserialize, Serialize}; -use super::{Dependency, GitUrl, KernelDependency, KernelName}; +use super::{Dependency, Dsl, GitUrl, KernelDependency, KernelName}; use crate::version::Version; // `monostate` validates the edition on read but provides no `Serialize` impl for it. @@ -139,6 +139,7 @@ pub struct TorchNoarch { pub struct TvmFfi { pub include: Option>, pub pyext: Option>, + #[serde(default)] pub src: Vec, pub cxx_flags: Option>, } @@ -150,7 +151,10 @@ pub enum Kernel { Cpu { cxx_flags: Option>, depends: Vec, + dsl: Option, + features: Option>, include: Option>, + lib_name: Option, src: Vec, }, #[serde(rename_all = "kebab-case")] @@ -160,7 +164,12 @@ pub enum Kernel { cuda_minver: Option, cxx_flags: Option>, depends: Vec, + device_manifest: Option, + dsl: Option, + features: Option>, include: Option>, + lib_name: Option, + ptx_dir: Option, src: Vec, }, #[serde(rename_all = "kebab-case")] @@ -352,12 +361,18 @@ impl From for super::Kernel { Kernel::Cpu { cxx_flags, depends, + dsl, + features, include, + lib_name, src, } => super::Kernel::Cpu { cxx_flags, depends, + dsl, + features, include, + lib_name, src, }, Kernel::Cuda { @@ -366,7 +381,12 @@ impl From for super::Kernel { cuda_minver, cxx_flags, depends, + device_manifest, + dsl, + features, include, + lib_name, + ptx_dir, src, } => super::Kernel::Cuda { cuda_capabilities, @@ -374,7 +394,12 @@ impl From for super::Kernel { cuda_minver, cxx_flags, depends, + device_manifest, + dsl, + features, include, + lib_name, + ptx_dir, src, }, Kernel::Metal { @@ -569,12 +594,18 @@ impl From for Kernel { super::Kernel::Cpu { cxx_flags, depends, + dsl, + features, include, + lib_name, src, } => Kernel::Cpu { cxx_flags, depends, + dsl, + features, include, + lib_name, src, }, super::Kernel::Cuda { @@ -583,7 +614,12 @@ impl From for Kernel { cuda_minver, cxx_flags, depends, + device_manifest, + dsl, + features, include, + lib_name, + ptx_dir, src, } => Kernel::Cuda { cuda_capabilities, @@ -591,7 +627,12 @@ impl From for Kernel { cuda_minver, cxx_flags, depends, + device_manifest, + dsl, + features, include, + lib_name, + ptx_dir, src, }, super::Kernel::Metal { diff --git a/nix-builder/lib/extension/tvm-ffi/arch.nix b/nix-builder/lib/extension/tvm-ffi/arch.nix index 879076e8..2003be7c 100644 --- a/nix-builder/lib/extension/tvm-ffi/arch.nix +++ b/nix-builder/lib/extension/tvm-ffi/arch.nix @@ -6,23 +6,29 @@ lib, pkgs, stdenv, - writeText, # Native build inputs kernel-builder, + cargo, cmake, cmakeNvccThreadsHook, cuda_nvcc, get-kernel-check, hash-kernel-hook, kernel-layout-check, + libclang, ninja, python3, remove-bytecode-hook, rewrite-nix-paths-macho, + rustc, + rustPlatform, + symlinkJoin, torch-ops-check, writeScriptBin, + rust-bin ? null, + # Framework packages cudaPackages, rocmPackages, @@ -104,6 +110,7 @@ let ]; moduleName = builtins.replaceStrings [ "-" ] [ "_" ] kernelName; + kernelToml = (builtins.fromTOML (builtins.readFile (src + "/build.toml"))).kernel or { }; # On Darwin, we need the host's xcrun for `xcrun metal` to compile Metal shaders. # It's not supported by the nixpkgs shim. @@ -117,6 +124,79 @@ let provenanceFlags = import ../provenance-flags.nix { inherit lib kernelProvenance; }; + kernelDsl = kernel: kernel.dsl or "cpp"; + + # Both cargo-built DSLs need the Rust toolchain; only cuda-oxide additionally + # needs the rustc codegen backend that compiles device code to PTX. + hasRustKernels = builtins.any ( + kernel: + builtins.elem (kernelDsl kernel) [ + "rust" + "cuda-oxide" + ] + && kernel.backend == buildConfig.backend + ) (lib.attrValues kernelToml); + + hasCudaOxideKernels = builtins.any (kernel: kernelDsl kernel == "cuda-oxide") ( + lib.attrValues kernelToml + ); + + cudaOxideSrc = pkgs.fetchFromGitHub { + owner = "drbh"; + repo = "cuda-oxide"; + rev = "bc5f33bb1d556671c35e475439bb821d25caa2cb"; + hash = "sha256-ZIpS0XFDI9yDSXk33OyQ3mGB5gRhE7tldmmex226tLA="; + }; + + cudaOxideRust = + if rust-bin != null then + rust-bin.fromRustupToolchainFile "${cudaOxideSrc}/rust-toolchain.toml" + else + rustc; + cudaOxideRustPlatform = pkgs.makeRustPlatform { + cargo = cudaOxideRust; + rustc = cudaOxideRust; + }; + cudaOxideBackend = cudaOxideRustPlatform.buildRustPackage { + pname = "rustc-codegen-cuda"; + version = "0.2.1"; + src = cudaOxideSrc; + sourceRoot = "source/crates/rustc-codegen-cuda"; + cargoLock = { + lockFile = "${cudaOxideSrc}/crates/rustc-codegen-cuda/Cargo.lock"; + allowBuiltinFetchGit = true; + }; + buildInputs = with pkgs; [ + libffi + libxml2 + zstd + zlib + ]; + doCheck = false; + }; + + rustCudaEnv = + let + cudaHome = symlinkJoin { + name = "cuda-home"; + paths = with cudaPackages; [ + (lib.getDev cuda_cudart) + (lib.getLib cuda_cudart) + (lib.getDev cccl) + ]; + }; + in + { + CUDA_HOME = cudaHome; + CUDA_TOOLKIT_PATH = cudaHome; + LIBRARY_PATH = "${lib.getOutput "stubs" cudaPackages.cuda_cudart}/lib/stubs"; + LIBCLANG_PATH = "${lib.getLib libclang}/lib"; + BINDGEN_EXTRA_CLANG_ARGS = lib.concatStringsSep " " [ + "-isystem ${lib.getLib libclang}/lib/clang/${lib.versions.major libclang.version}/include" + "-isystem ${lib.getDev stdenv.cc.libc}/include" + ]; + }; + in stdenv.mkDerivation (prevAttrs: { @@ -131,6 +211,15 @@ stdenv.mkDerivation (prevAttrs: { framework = "tvm-ffi"; + cargoDeps = + if hasRustKernels then + rustPlatform.importCargoLock { + lockFile = src + "/Cargo.lock"; + allowBuiltinFetchGit = true; + } + else + null; + # We run kernel-builder here rather than patchPhase or preConfigure, # so that external users of `src` get the source tree with the files # generated by kernel-builder. @@ -166,6 +255,11 @@ stdenv.mkDerivation (prevAttrs: { chmod -R u+w . ''; + preInstallCheck = lib.optionalString (hasRustKernels && cudaSupport) '' + addToSearchPath LD_LIBRARY_PATH "${lib.getOutput "stubs" cudaPackages.cuda_cudart}/lib/stubs" + export LD_LIBRARY_PATH + ''; + nativeBuildInputs = [ cmake hash-kernel-hook @@ -175,6 +269,18 @@ stdenv.mkDerivation (prevAttrs: { remove-bytecode-hook torch-ops-check ] + ++ lib.optionals hasRustKernels ( + [ rustPlatform.cargoSetupHook ] + ++ ( + if rust-bin != null && builtins.pathExists (src + "/rust-toolchain.toml") then + [ (rust-bin.fromRustupToolchainFile (src + "/rust-toolchain.toml")) ] + else + [ + cargo + rustc + ] + ) + ) ++ lib.optionals doGetKernelCheck [ (get-kernel-check.override { python3 = python3.withPackages (ps: dependencies); }) ] @@ -235,6 +341,10 @@ stdenv.mkDerivation (prevAttrs: { // lib.optionalAttrs xpuSupport { MKLROOT = oneapi-torch-dev; SYCL_ROOT = oneapi-torch-dev; + } + // lib.optionalAttrs (hasRustKernels && cudaSupport) rustCudaEnv + // lib.optionalAttrs hasCudaOxideKernels { + CUDA_OXIDE_BACKEND = "${cudaOxideBackend}/lib/librustc_codegen_cuda.so"; }; # If we use the default setup, CMAKE_CUDA_HOST_COMPILER gets set to nixpkgs g++. @@ -254,6 +364,10 @@ stdenv.mkDerivation (prevAttrs: { # rpaths anyway. (lib.cmakeBool "CMAKE_SKIP_RPATH" true) ] + ++ lib.optionals hasRustKernels [ + (lib.cmakeFeature "CMAKE_C_COMPILER" "${stdenv.cc}/bin/cc") + (lib.cmakeFeature "CMAKE_CXX_COMPILER" "${stdenv.cc}/bin/c++") + ] ++ lib.optionals cudaSupport [ (lib.cmakeFeature "CMAKE_CUDA_HOST_COMPILER" "${stdenv.cc}/bin/g++") (lib.cmakeFeature "GPU_LANG" "CUDA") diff --git a/nix-builder/tests/Dockerfile.test-kernel b/nix-builder/tests/Dockerfile.test-kernel index 2a605290..329a059b 100644 --- a/nix-builder/tests/Dockerfile.test-kernel +++ b/nix-builder/tests/Dockerfile.test-kernel @@ -73,6 +73,7 @@ RUN uv add ./kernels-src COPY relu-kernel ./relu-kernel COPY relu-torch-stable-abi-kernel ./relu-torch-stable-abi-kernel COPY relu-tvm-ffi-kernel ./relu-tvm-ffi-kernel +COPY relu-cuda-oxide-kernel ./relu-cuda-oxide-kernel COPY relu-kernel-cpu ./relu-kernel-cpu COPY cutlass-gemm-kernel ./cutlass-gemm-kernel COPY cutlass-gemm-tvm-ffi-kernel ./cutlass-gemm-tvm-ffi-kernel @@ -83,6 +84,7 @@ COPY cpp20-symbols-kernel ./cpp20-symbols-kernel COPY examples/kernels/extra-data/tests ./extra_data_tests COPY examples/kernels/relu/tests ./relu_tests COPY examples/kernels/relu-tvm-ffi/tests ./relu_tvm_ffi_tests +COPY examples/kernels/relu-cuda-oxide/tests ./relu_cuda_oxide_tests COPY examples/kernels/cutlass-gemm/tests ./cutlass_gemm_tests COPY examples/kernels/cutlass-gemm-tvm-ffi/tests ./cutlass_gemm_tvm_ffi_tests COPY examples/kernels/cpp20-symbols/tests ./cpp20_symbols_tests diff --git a/nix-builder/tests/run-tests.sh b/nix-builder/tests/run-tests.sh index 33426b81..272f0be0 100644 --- a/nix-builder/tests/run-tests.sh +++ b/nix-builder/tests/run-tests.sh @@ -6,6 +6,7 @@ EXTRA_DATA_PATH=$(echo extra-data/torch*) RELU_PATH=$(echo relu-kernel/torch*) RELU_TORCH_STABLE_ABI_PATH=$(echo relu-torch-stable-abi-kernel/torch*) RELU_TVM_FFI_PATH=$(echo relu-tvm-ffi-kernel/tvm-ffi*) +RELU_CUDA_OXIDE_PATH=$(echo relu-cuda-oxide-kernel/tvm-ffi*) CUTLASS_PATH=$(echo cutlass-gemm-kernel/torch*) CUTLASS_TVM_FFI_PATH=$(echo cutlass-gemm-tvm-ffi-kernel/tvm-ffi*) RELU_TRITON_PATH=$(echo relu-triton-kernel/torch*) @@ -13,8 +14,8 @@ SILU_MUL_PATH=$(echo silu-and-mul-kernel/torch*) RELU_CPU_PATH=$(echo relu-kernel-cpu/torch*) CPP20_SYMBOLS_PATH=$(echo cpp20-symbols-kernel/torch*) -LOCAL_KERNELS="kernels-test/extra-data=${EXTRA_DATA_PATH}:kernels-test/relu=${RELU_PATH}:kernels-test/relu-torch-stable-abi=${RELU_TORCH_STABLE_ABI_PATH}:kernels-test/relu-tvm-ffi=${RELU_TVM_FFI_PATH}:kernels-test/cutlass-gemm=${CUTLASS_PATH}:kernels-test/cutlass-gemm-tvm-ffi=${CUTLASS_TVM_FFI_PATH}" \ - .venv/bin/pytest extra_data_tests relu_tests relu_tvm_ffi_tests cutlass_gemm_tests cutlass_gemm_tvm_ffi_tests +LOCAL_KERNELS="kernels-test/extra-data=${EXTRA_DATA_PATH}:kernels-test/relu=${RELU_PATH}:kernels-test/relu-torch-stable-abi=${RELU_TORCH_STABLE_ABI_PATH}:kernels-test/relu-tvm-ffi=${RELU_TVM_FFI_PATH}:kernels-test/relu-cuda-oxide=${RELU_CUDA_OXIDE_PATH}:kernels-test/cutlass-gemm=${CUTLASS_PATH}:kernels-test/cutlass-gemm-tvm-ffi=${CUTLASS_TVM_FFI_PATH}" \ + .venv/bin/pytest extra_data_tests relu_tests relu_tvm_ffi_tests relu_cuda_oxide_tests cutlass_gemm_tests cutlass_gemm_tvm_ffi_tests LOCAL_KERNELS="kernels-test/relu-triton=${RELU_TRITON_PATH}" \ .venv/bin/pytest relu_triton_tests From d4bfd09481df7803c8dc31464720e81b4615b415 Mon Sep 17 00:00:00 2001 From: David Holtz Date: Tue, 18 Aug 2026 23:03:08 -0400 Subject: [PATCH 2/2] fix: load Rust kernels through kernels API --- examples/kernels/relu-cuda-oxide/tests/test_relu.py | 3 ++- examples/kernels/relu-rust/tests/test_relu.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/examples/kernels/relu-cuda-oxide/tests/test_relu.py b/examples/kernels/relu-cuda-oxide/tests/test_relu.py index f68ea0ff..14f3128d 100644 --- a/examples/kernels/relu-cuda-oxide/tests/test_relu.py +++ b/examples/kernels/relu-cuda-oxide/tests/test_relu.py @@ -1,8 +1,9 @@ +import kernels import pytest import torch import torch.nn.functional as F -import relu_cuda_oxide +relu_cuda_oxide = kernels.get_kernel("kernels-test/relu-cuda-oxide", version=1) @pytest.mark.kernels_ci diff --git a/examples/kernels/relu-rust/tests/test_relu.py b/examples/kernels/relu-rust/tests/test_relu.py index c870279e..80fe2099 100644 --- a/examples/kernels/relu-rust/tests/test_relu.py +++ b/examples/kernels/relu-rust/tests/test_relu.py @@ -1,8 +1,9 @@ +import kernels import pytest import torch import torch.nn.functional as F -import relu_rust +relu_rust = kernels.get_kernel("kernels-test/relu-rust", version=1) @pytest.mark.kernels_ci