From c440ae7f92fbb6fb142376c16a55bed8f56a16fe Mon Sep 17 00:00:00 2001 From: David Holtz Date: Thu, 15 Oct 2026 18:11:36 +0000 Subject: [PATCH 01/14] feat: support rust gpu --- docs/source/builder/writing-kernels.md | 2 +- flake.lock | 6 +- kernel-builder/src/pyproject/kernel.rs | 78 +++++++++++++++++++ .../kernel-component/rust-cuda.cmake | 9 +++ .../src/pyproject/templates/rust.cmake | 66 +++++++++++++++- kernel-builder/tests/rust-components.cmake | 60 ++++++++++++++ kernels-common/src/config/mod.rs | 75 +++++++++++++++++- kernels-common/src/config/v6.rs | 52 ++++++++++++- nix-builder/lib/build.nix | 3 + nix-builder/lib/extension/tvm-ffi/arch.nix | 70 ++++++++++++----- .../lib/extension/tvm-ffi/rust-cuda.nix | 60 ++++++++++++++ nix-builder/lib/source-set.nix | 7 +- 12 files changed, 454 insertions(+), 34 deletions(-) create mode 100644 kernel-builder/src/pyproject/templates/kernel-component/rust-cuda.cmake create mode 100644 kernel-builder/tests/rust-components.cmake create mode 100644 nix-builder/lib/extension/tvm-ffi/rust-cuda.nix diff --git a/docs/source/builder/writing-kernels.md b/docs/source/builder/writing-kernels.md index 48c9bddf..9c57b294 100644 --- a/docs/source/builder/writing-kernels.md +++ b/docs/source/builder/writing-kernels.md @@ -318,7 +318,7 @@ The following options can be set for a kernel: - `language` (required): the language and compute backend of the kernel, as `-`. The currently supported values are `cpp-cpu`, - `cpp-cuda`, `cpp-metal`, `cpp-rocm`, `cpp-xpu`, and `rust-cpu`. + `cpp-cuda`, `cpp-metal`, `cpp-rocm`, `cpp-xpu`, `rust-cpu`, and `rust-cuda`. **The `cpu` backend is currently experimental and might still change.** - `depends` (required): a list of dependencies. The supported dependencies are listed in [`cpp-deps.nix`](https://github.com/huggingface/kernels/blob/main/builder/lib/cpp-deps.nix). diff --git a/flake.lock b/flake.lock index 47af4d8d..edadde8f 100644 --- a/flake.lock +++ b/flake.lock @@ -64,11 +64,11 @@ ] }, "locked": { - "lastModified": 1789543870, - "narHash": "sha256-45ikFmtfgBZgMePIm4GevJ9VjDSWO3PGQgOuuwHDpes=", + "lastModified": 1792048409, + "narHash": "sha256-ZcwtnGWuY65NQ75W1LZ99gIaFBgC0Sm+TCs1ZEabw0U=", "owner": "oxalica", "repo": "rust-overlay", - "rev": "35ca0490d13a3d38c4602d0eb9600a30fa63a367", + "rev": "e9d12b9d8266ed9afdc5add5537ec8b5f789645c", "type": "github" }, "original": { diff --git a/kernel-builder/src/pyproject/kernel.rs b/kernel-builder/src/pyproject/kernel.rs index 43fa036b..aa603f92 100644 --- a/kernel-builder/src/pyproject/kernel.rs +++ b/kernel-builder/src/pyproject/kernel.rs @@ -40,6 +40,26 @@ fn render_kernel_component( Kernel::RustCpu { cargo_manifest, .. } => { render_kernel_component_rust(env, kernel_name, cargo_manifest, write)? } + Kernel::RustCuda { + cargo_manifest, + device_manifest, + ptx_dir, + cuda_capabilities, + .. + } => { + env.get_template("kernel-component/rust-cuda.cmake")? + .render_captured_to( + context! { + name => kernel_name, + manifest_path => cargo_manifest, + device_manifest => device_manifest, + ptx_dir => ptx_dir, + cuda_capabilities => cuda_capabilities, + }, + &mut *write, + )?; + write.write_all(b"\n")?; + } Kernel::CppCuda { .. } => { render_kernel_component_cuda(env, kernel_name, kernel, sources, write)? } @@ -238,3 +258,61 @@ fn render_kernel_component_xpu( Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rust_components_render_for_their_backend() { + let mut env = Environment::new(); + minijinja_embed::load_templates!(&mut env); + for (kernel, backend, device) in [ + ( + Kernel::RustCpu { + cargo_manifest: "host/Cargo.toml".into(), + depends: vec![], + src: vec![], + }, + "CPU", + false, + ), + ( + Kernel::RustCuda { + cargo_manifest: "host/Cargo.toml".into(), + device_manifest: Some("device/Cargo.toml".into()), + ptx_dir: Some("ptx".into()), + cuda_capabilities: Some(vec!["8.0".into(), "9.0".into()]), + depends: vec![], + src: vec![], + }, + "CUDA", + true, + ), + ( + Kernel::RustCuda { + cargo_manifest: "host/Cargo.toml".into(), + device_manifest: None, + ptx_dir: None, + cuda_capabilities: None, + depends: vec![], + src: vec![], + }, + "CUDA", + false, + ), + ] { + let mut rendered = vec![]; + render_kernel_component(&env, "example", &kernel, &mut rendered).unwrap(); + let rendered = String::from_utf8(rendered).unwrap(); + assert!(rendered.contains(&format!("if(GPU_LANG STREQUAL \"{backend}\")"))); + assert!(rendered.contains("MANIFEST_PATH \"host/Cargo.toml\"")); + assert_eq!(rendered.contains("DEVICE_MANIFEST"), device); + if device { + assert!(rendered.contains("DEVICE_MANIFEST \"device/Cargo.toml\"")); + assert!(rendered.contains("PTX_DIR \"ptx\"")); + assert!(rendered.contains("CUDA_CAPABILITIES 8.0 9.0")); + } + } + } +} 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..7664b3cf --- /dev/null +++ b/kernel-builder/src/pyproject/templates/kernel-component/rust-cuda.cmake @@ -0,0 +1,9 @@ +if(GPU_LANG STREQUAL "CUDA") + rust_kernel_component(RUST_KERNEL_LIBS RUST_KERNEL_TARGETS + NAME {{ name }} + MANIFEST_PATH "{{ manifest_path }}" + {% 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/rust.cmake b/kernel-builder/src/pyproject/templates/rust.cmake index 727fff42..1c2a056b 100644 --- a/kernel-builder/src/pyproject/templates/rust.cmake +++ b/kernel-builder/src/pyproject/templates/rust.cmake @@ -1,5 +1,5 @@ function(rust_kernel_component LIBS_VAR TARGETS_VAR) - cmake_parse_arguments(KERNEL "" "NAME;MANIFEST_PATH" "" ${ARGN}) + cmake_parse_arguments(KERNEL "" "NAME;MANIFEST_PATH;DEVICE_MANIFEST;PTX_DIR" "CUDA_CAPABILITIES" ${ARGN}) if(NOT KERNEL_NAME OR NOT KERNEL_MANIFEST_PATH) message(FATAL_ERROR "rust_kernel_component: NAME and MANIFEST_PATH are required") @@ -10,6 +10,62 @@ function(rust_kernel_component LIBS_VAR TARGETS_VAR) "but `cargo` was not found. Install a Rust toolchain or set CARGO_EXECUTABLE.") endif() + if(GPU_LANG STREQUAL "CUDA") + if(KERNEL_CUDA_CAPABILITIES) + cuda_archs_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() + + 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() + string(REPLACE "-" "_" _LIB_NAME ${KERNEL_NAME}) set(_CARGO_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/${KERNEL_NAME}) @@ -21,9 +77,12 @@ function(rust_kernel_component LIBS_VAR TARGETS_VAR) 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} rustc --release --locked --lib --crate-type staticlib --manifest-path ${CMAKE_CURRENT_SOURCE_DIR}/${KERNEL_MANIFEST_PATH} --target-dir ${_CARGO_TARGET_DIR} + DEPENDS ${_RUST_KERNEL_DEPENDS} BYPRODUCTS ${_STATICLIB} WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} COMMENT "Building Rust kernel ${KERNEL_NAME} with cargo" @@ -82,5 +141,10 @@ function(target_link_rust_kernels TARGET) "$" Threads::Threads ${CMAKE_DL_LIBS}) + # Cargo's native-library link directives do not propagate out of a staticlib. + if(GPU_LANG STREQUAL "CUDA") + find_package(CUDAToolkit REQUIRED) + target_link_libraries(${TARGET} PRIVATE CUDA::cuda_driver) + endif() _restrict_rust_exports(${TARGET}) endfunction() diff --git a/kernel-builder/tests/rust-components.cmake b/kernel-builder/tests/rust-components.cmake new file mode 100644 index 00000000..72cf2dc8 --- /dev/null +++ b/kernel-builder/tests/rust-components.cmake @@ -0,0 +1,60 @@ +# Configure Rust targets without CUDA or compiling crates. +# Run: cmake -P kernel-builder/tests/rust-components.cmake +cmake_minimum_required(VERSION 3.24) +if(CMAKE_SCRIPT_MODE_FILE) + set(_TEST_DIR "${CMAKE_CURRENT_BINARY_DIR}/rust-components-test") + file(MAKE_DIRECTORY "${_TEST_DIR}") + file(WRITE "${_TEST_DIR}/CMakeLists.txt" + "cmake_minimum_required(VERSION 3.24)\nproject(rust_components NONE)\ninclude(\"${CMAKE_CURRENT_LIST_FILE}\")\n") + foreach(_CASE cpu cuda host-only unsupported-arch missing-backend) + execute_process(COMMAND "${CMAKE_COMMAND}" -S "${_TEST_DIR}" + -G Ninja -B "${_TEST_DIR}/${_CASE}" "-DCASE=${_CASE}" + RESULT_VARIABLE _RESULT OUTPUT_VARIABLE _OUTPUT ERROR_VARIABLE _ERROR) + if(_CASE STREQUAL "unsupported-arch") + if(_RESULT EQUAL 0 OR NOT _ERROR MATCHES "empty set of capabilities") + message(FATAL_ERROR "Expected unsupported architecture error: ${_OUTPUT}${_ERROR}") + endif() + elseif(_CASE STREQUAL "missing-backend") + if(_RESULT EQUAL 0 OR NOT _ERROR MATCHES "requires CUDA_OXIDE_BACKEND") + message(FATAL_ERROR "Expected missing backend error: ${_OUTPUT}${_ERROR}") + endif() + elseif(NOT _RESULT EQUAL 0) + message(FATAL_ERROR "${_CASE} configuration failed: ${_OUTPUT}${_ERROR}") + endif() + endforeach() + file(REMOVE_RECURSE "${_TEST_DIR}") + message(STATUS "Rust component configuration tests passed") + return() +endif() + +include("${CMAKE_CURRENT_LIST_DIR}/../src/pyproject/templates/utils.cmake") +include("${CMAKE_CURRENT_LIST_DIR}/../src/pyproject/templates/kernel.cmake") +include("${CMAKE_CURRENT_LIST_DIR}/../src/pyproject/templates/rust.cmake") +set(CARGO_EXECUTABLE cargo) +set(Python_EXECUTABLE python) +set(GPU_LANG CUDA) +set(CUDA_ARCHS "8.0;8.9;9.0") +set(CUDA_KERNEL_ARCHS "8.0;8.9;9.0+PTX") +set(ENV{CUDA_OXIDE_BACKEND} /test/backend.so) +set(_ARGS CUDA_CAPABILITIES 8.0 DEVICE_MANIFEST device/Cargo.toml) +if(CASE STREQUAL "cpu") + set(GPU_LANG CPU) + set(_ARGS) +elseif(CASE STREQUAL "host-only") + set(_ARGS) +elseif(CASE STREQUAL "unsupported-arch") + set(_ARGS CUDA_CAPABILITIES 6.0 DEVICE_MANIFEST device/Cargo.toml) +elseif(CASE STREQUAL "missing-backend") + unset(ENV{CUDA_OXIDE_BACKEND}) +endif() +rust_kernel_component(LIBS TARGETS NAME example MANIFEST_PATH Cargo.toml ${_ARGS}) +if(NOT TARGET example_rust OR NOT TARGET example_cargo_build) + message(FATAL_ERROR "Missing shared Cargo/static-library targets") +endif() +if(CASE STREQUAL "cuda") + if(NOT TARGET example_oxide_device_build OR NOT ALL_GPU_ARCHS STREQUAL "8.0") + message(FATAL_ERROR "Missing device target or incorrect capabilities") + endif() +elseif(TARGET example_oxide_device_build) + message(FATAL_ERROR "Unexpected device build for ${CASE}") +endif() diff --git a/kernels-common/src/config/mod.rs b/kernels-common/src/config/mod.rs index fb1de026..e566dd65 100644 --- a/kernels-common/src/config/mod.rs +++ b/kernels-common/src/config/mod.rs @@ -356,6 +356,15 @@ pub enum Kernel { depends: Vec, src: Vec, }, + + RustCuda { + cargo_manifest: String, + device_manifest: Option, + ptx_dir: Option, + cuda_capabilities: Option>, + depends: Vec, + src: Vec, + }, CppCuda { cuda_capabilities: Option>, cuda_flags: Option>, @@ -396,7 +405,7 @@ impl Kernel { | Kernel::CppMetal { cxx_flags, .. } | Kernel::CppRocm { cxx_flags, .. } | Kernel::CppXpu { cxx_flags, .. } => cxx_flags.as_deref(), - Kernel::RustCpu { .. } => None, + Kernel::RustCpu { .. } | Kernel::RustCuda { .. } => None, } } @@ -407,7 +416,7 @@ impl Kernel { | Kernel::CppMetal { include, .. } | Kernel::CppRocm { include, .. } | Kernel::CppXpu { include, .. } => include.as_deref(), - Kernel::RustCpu { .. } => None, + Kernel::RustCpu { .. } | Kernel::RustCuda { .. } => None, } } @@ -421,7 +430,7 @@ impl Kernel { pub fn backend(&self) -> Backend { match self { Kernel::CppCpu { .. } | Kernel::RustCpu { .. } => Backend::Cpu, - Kernel::CppCuda { .. } => Backend::Cuda, + Kernel::CppCuda { .. } | Kernel::RustCuda { .. } => Backend::Cuda, Kernel::CppMetal { .. } => Backend::Metal, Kernel::CppRocm { .. } => Backend::Rocm, Kernel::CppXpu { .. } => Backend::Xpu, @@ -435,7 +444,7 @@ impl Kernel { | Kernel::CppMetal { .. } | Kernel::CppRocm { .. } | Kernel::CppXpu { .. } => Language::Cpp, - Kernel::RustCpu { .. } => Language::Rust, + Kernel::RustCpu { .. } | Kernel::RustCuda { .. } => Language::Rust, } } @@ -443,6 +452,7 @@ impl Kernel { match self { Kernel::CppCpu { depends, .. } | Kernel::RustCpu { depends, .. } + | Kernel::RustCuda { depends, .. } | Kernel::CppCuda { depends, .. } | Kernel::CppMetal { depends, .. } | Kernel::CppRocm { depends, .. } @@ -454,6 +464,7 @@ impl Kernel { match self { Kernel::CppCpu { src, .. } | Kernel::RustCpu { src, .. } + | Kernel::RustCuda { src, .. } | Kernel::CppCuda { src, .. } | Kernel::CppMetal { src, .. } | Kernel::CppRocm { src, .. } @@ -735,9 +746,65 @@ mod tests { assert_eq!(build.kernels["cpu_kernel"].language(), Language::Rust); } + #[test] + fn v6_rust_cuda_round_trip() { + let config = r#" + [general] + name = "rust-cuda" + version = 1 + edition = 6 + license = "Apache-2.0" + backends = ["cuda"] + [tvm-ffi] + [kernel.gpu_kernel] + language = "rust-cuda" + cargo-manifest = "host/Cargo.toml" + device-manifest = "device/Cargo.toml" + ptx-dir = "ptx" + cuda-capabilities = ["8.0"] + depends = [] + src = ["host/src", "device/src"] + "#; + let build = Build::try_from(toml::from_str::(config).unwrap()).unwrap(); + assert_eq!(build.kernels["gpu_kernel"].language(), Language::Rust); + assert_eq!(build.kernels["gpu_kernel"].backend(), Backend::Cuda); + let serialized = toml::to_string(&v6::Build::from(build)).unwrap(); + let reparsed: v6::Build = toml::from_str(&serialized).unwrap(); + match &reparsed.kernels["gpu_kernel"] { + v6::Kernel::RustCuda { + cargo_manifest, + device_manifest, + ptx_dir, + cuda_capabilities, + .. + } => { + assert_eq!(cargo_manifest, "host/Cargo.toml"); + assert_eq!(device_manifest.as_deref(), Some("device/Cargo.toml")); + assert_eq!(ptx_dir.as_deref(), Some("ptx")); + assert_eq!( + cuda_capabilities.as_deref(), + Some(["8.0".to_string()].as_slice()) + ); + } + _ => panic!("expected Rust CUDA kernel"), + } + } + #[test] fn v6_rust_kernel_rejects_invalid_config() { let cases = [ + ( + "[torch]\nsrc = []", + "rust-cuda", + "", + "require a `[tvm-ffi]` framework", + ), + ( + "[tvm-ffi]", + "rust-cuda", + r#"cuda-flags = ["-O3"]"#, + "unknown field `cuda-flags`", + ), ( "[tvm-ffi]", "rust-cpu", diff --git a/kernels-common/src/config/v6.rs b/kernels-common/src/config/v6.rs index c4f621f4..b924640d 100644 --- a/kernels-common/src/config/v6.rs +++ b/kernels-common/src/config/v6.rs @@ -166,6 +166,16 @@ pub enum Kernel { depends: Vec, src: Vec, }, + + #[serde(rename_all = "kebab-case")] + RustCuda { + cargo_manifest: String, + device_manifest: Option, + ptx_dir: Option, + cuda_capabilities: Option>, + depends: Vec, + src: Vec, + }, #[serde(rename_all = "kebab-case")] CppCuda { cuda_capabilities: Option>, @@ -224,10 +234,12 @@ impl TryFrom for super::Build { .kernels .into_iter() .map(|(name, kernel)| match kernel { - Kernel::RustCpu { .. } if !tvm_ffi => Err(ConfigError::InvalidKernel { - name, - reason: "Rust kernels require a `[tvm-ffi]` framework".into(), - }), + Kernel::RustCpu { .. } | Kernel::RustCuda { .. } if !tvm_ffi => { + Err(ConfigError::InvalidKernel { + name, + reason: "Rust kernels require a `[tvm-ffi]` framework".into(), + }) + } kernel => Ok((name, kernel.into())), }) .collect::>()?; @@ -382,6 +394,22 @@ impl From for super::Kernel { include, src, }, + + Kernel::RustCuda { + cargo_manifest, + device_manifest, + ptx_dir, + cuda_capabilities, + depends, + src, + } => super::Kernel::RustCuda { + cargo_manifest, + device_manifest, + ptx_dir, + cuda_capabilities, + depends, + src, + }, Kernel::RustCpu { cargo_manifest, depends, @@ -617,6 +645,22 @@ impl From for Kernel { depends, src, }, + + super::Kernel::RustCuda { + cargo_manifest, + device_manifest, + ptx_dir, + cuda_capabilities, + depends, + src, + } => Kernel::RustCuda { + cargo_manifest, + device_manifest, + ptx_dir, + cuda_capabilities, + depends, + src, + }, super::Kernel::CppCuda { cuda_capabilities, cuda_flags, diff --git a/nix-builder/lib/build.nix b/nix-builder/lib/build.nix index a43fd6ca..025c3fbf 100644 --- a/nix-builder/lib/build.nix +++ b/nix-builder/lib/build.nix @@ -202,6 +202,9 @@ rec { kernelProvenance ; + rustCudaDevice = lib.any (kernel: kernel.language == "rust-cuda" && kernel ? device-manifest) ( + lib.attrValues kernels + ); cargoLock = if hasRustKernels then src + "/Cargo.lock" else null; rustGitHashes = if hasRustKernels && builtins.pathExists (src + "/rust-git-hashes.json") then diff --git a/nix-builder/lib/extension/tvm-ffi/arch.nix b/nix-builder/lib/extension/tvm-ffi/arch.nix index f3cec3c7..6ffccf2b 100644 --- a/nix-builder/lib/extension/tvm-ffi/arch.nix +++ b/nix-builder/lib/extension/tvm-ffi/arch.nix @@ -62,6 +62,9 @@ # the build has no Rust kernels. cargoLock ? null, + # Compile CUDA device crates with cuda-oxide and its matching rustc. + rustCudaDevice ? false, + # Fixed-output hashes for Git dependencies in Cargo.lock. rustGitHashes ? { }, @@ -136,31 +139,45 @@ let metalSupport = buildConfig.metal or false; rustSupport = cargoLock != null; + rustCuda = import ./rust-cuda.nix { + inherit + pkgs + lib + stdenv + cudaPackages + ; + }; # rpaths are stripped from kernels to make them portable, but that # also means that in a Nix environment the CUDA/ROCm/oneAPI dependencies # cannot be located anymore, so pass them to hooks that load the kernel. - libraryPath = lib.makeLibraryPath ( - map lib.getLib ( - lib.optionals cudaSupport ( - with cudaPackages; - [ - cuda_cudart - libcublas - libcusolver - libcusparse - ] - ) - ++ lib.optionals rocmSupport [ clr ] - ++ lib.optionals xpuSupport ( - with xpuPackages; - [ - intel-oneapi-compiler-dpcpp-cpp-runtime - intel-oneapi-compiler-shared-runtime - ] + libraryPath = + lib.makeLibraryPath ( + map lib.getLib ( + lib.optionals cudaSupport ( + with cudaPackages; + [ + cuda_cudart + libcublas + libcusolver + libcusparse + ] + ) + ++ lib.optionals rocmSupport [ clr ] + ++ lib.optionals xpuSupport ( + with xpuPackages; + [ + intel-oneapi-compiler-dpcpp-cpp-runtime + intel-oneapi-compiler-shared-runtime + ] + ) ) ) - ); + + + lib.optionalString (rustSupport && cudaSupport) + # Import checks run without a GPU driver in the build sandbox. The stub is + # only used by the check hooks, never bundled in the redistributable kernel. + ":${lib.getOutput "stubs" cudaPackages.cuda_cudart}/lib/stubs"; provenanceFlags = import ../provenance-flags.nix { inherit lib kernelProvenance; }; @@ -253,9 +270,12 @@ stdenv.mkDerivation ( ] ++ lib.optionals rustSupport [ rustPlatform.cargoSetupHook + ] + ++ lib.optionals (rustSupport && !rustCudaDevice) [ cargo rustc - ]; + ] + ++ lib.optionals rustCudaDevice [ rustCuda.toolchain ]; buildInputs = [ python3.pkgs.tvm-ffi @@ -305,6 +325,10 @@ stdenv.mkDerivation ( // lib.optionalAttrs cudaSupport { CUDAToolkit_ROOT = "${lib.getDev cudaPackages.cuda_nvcc}"; } + // lib.optionalAttrs (rustSupport && cudaSupport) rustCuda.env + // lib.optionalAttrs rustCudaDevice { + CUDA_OXIDE_BACKEND = "${rustCuda.backend}/lib/librustc_codegen_cuda.so"; + } // lib.optionalAttrs xpuSupport { MKLROOT = oneapi-torch-dev; SYCL_ROOT = oneapi-torch-dev; @@ -335,6 +359,12 @@ stdenv.mkDerivation ( # rpaths anyway. (lib.cmakeBool "CMAKE_SKIP_RPATH" true) ] + # Rust toolchains can propagate a newer compiler onto PATH. Link with the + # selected stdenv so redistributable kernels retain their glibc baseline. + ++ lib.optionals rustSupport [ + (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/lib/extension/tvm-ffi/rust-cuda.nix b/nix-builder/lib/extension/tvm-ffi/rust-cuda.nix new file mode 100644 index 00000000..cbb20d56 --- /dev/null +++ b/nix-builder/lib/extension/tvm-ffi/rust-cuda.nix @@ -0,0 +1,60 @@ +# cuda-oxide is a rustc codegen plugin and must use its matching Rust toolchain. +{ + pkgs, + lib, + stdenv, + cudaPackages, +}: +let + src = pkgs.fetchFromGitHub { + owner = "NVIDIA"; + repo = "cuda-rust"; + rev = "6d47ae49cb256a07948fb499c20c3eb9b2a96268"; + hash = "sha256-g9qg2zHLybO8kPLoOJqnYbWjRYtjGV3QVT+X4f8+1fg="; + }; + toolchain = pkgs.rust-bin.fromRustupToolchainFile "${src}/cuda-oxide/rust-toolchain.toml"; + platform = pkgs.makeRustPlatform { + cargo = toolchain; + rustc = toolchain; + }; + backend = platform.buildRustPackage { + pname = "rustc-codegen-cuda"; + version = "0.2.1"; + inherit src; + sourceRoot = "source/cuda-oxide/crates/rustc-codegen-cuda"; + cargoLock = { + lockFile = "${src}/cuda-oxide/crates/rustc-codegen-cuda/Cargo.lock"; + outputHashes."pliron-0.18.0" = "sha256-HQ/xNfQcIf9pQRrgxCYGFdVfdCMaE4iwj/G9mgAJOWw="; + }; + buildInputs = with pkgs; [ + libffi + libxml2 + zstd + zlib + ]; + doCheck = false; + }; + cudaHome = pkgs.symlinkJoin { + name = "cuda-home"; + paths = with cudaPackages; [ + (lib.getDev cuda_cudart) + (lib.getLib cuda_cudart) + (lib.getDev cccl) + (lib.getOutput "include" libcurand) + cuda_crt + ]; + }; +in +{ + inherit toolchain backend; + env = { + CUDA_HOME = cudaHome; + CUDA_TOOLKIT_PATH = cudaHome; + LIBRARY_PATH = "${lib.getOutput "stubs" cudaPackages.cuda_cudart}/lib/stubs"; + LIBCLANG_PATH = "${lib.getLib pkgs.libclang}/lib"; + BINDGEN_EXTRA_CLANG_ARGS = lib.concatStringsSep " " [ + "-isystem ${lib.getLib pkgs.libclang}/lib/clang/${lib.versions.major pkgs.libclang.version}/include" + "-isystem ${lib.getDev stdenv.cc.libc}/include" + ]; + }; +} diff --git a/nix-builder/lib/source-set.nix b/nix-builder/lib/source-set.nix index a1d75909..36276b75 100644 --- a/nix-builder/lib/source-set.nix +++ b/nix-builder/lib/source-set.nix @@ -43,7 +43,11 @@ let else fileset.empty; # Rust kernels also need their Cargo manifest. - kernelSrc = kernel: kernel.src ++ lib.optional (kernel ? cargo-manifest) kernel.cargo-manifest; + kernelSrc = + kernel: + kernel.src + ++ lib.optional (kernel ? cargo-manifest) kernel.cargo-manifest + ++ lib.optional (kernel ? device-manifest) kernel.device-manifest; kernelsSrc = fileset.unions ( lib.flatten (lib.mapAttrsToList (name: kernel: map (nameToPath path) (kernelSrc kernel)) kernels) ); @@ -55,6 +59,7 @@ fileset.toSource { kernelsSrc lockSet cargoLockSet + (fileset.maybeMissing (path + "/rust-toolchain.toml")) rustGitHashesSet srcSet pySrcSet From 9cd4c01cfa7bec0174e368541381aa93530b8c32 Mon Sep 17 00:00:00 2001 From: David Holtz Date: Tue, 20 Oct 2026 19:44:28 +0000 Subject: [PATCH 02/14] feat: simplify rust cuda build and drop extra env vars --- kernel-builder/src/pyproject/kernel.rs | 35 ++----- .../kernel-component/rust-cuda.cmake | 2 - .../src/pyproject/templates/rust.cmake | 95 ++++++++++++++----- kernel-builder/tests/rust-components.cmake | 24 +++-- kernels-common/src/config/mod.rs | 42 +++----- kernels-common/src/config/v6.rs | 13 --- nix-builder/lib/build.nix | 3 - nix-builder/lib/extension/tvm-ffi/arch.nix | 68 +++++-------- .../lib/extension/tvm-ffi/rust-cuda.nix | 27 +----- nix-builder/lib/source-set.nix | 6 +- 10 files changed, 136 insertions(+), 179 deletions(-) diff --git a/kernel-builder/src/pyproject/kernel.rs b/kernel-builder/src/pyproject/kernel.rs index aa603f92..7ccfe5a7 100644 --- a/kernel-builder/src/pyproject/kernel.rs +++ b/kernel-builder/src/pyproject/kernel.rs @@ -42,8 +42,6 @@ fn render_kernel_component( } Kernel::RustCuda { cargo_manifest, - device_manifest, - ptx_dir, cuda_capabilities, .. } => { @@ -52,8 +50,6 @@ fn render_kernel_component( context! { name => kernel_name, manifest_path => cargo_manifest, - device_manifest => device_manifest, - ptx_dir => ptx_dir, cuda_capabilities => cuda_capabilities, }, &mut *write, @@ -267,7 +263,7 @@ mod tests { fn rust_components_render_for_their_backend() { let mut env = Environment::new(); minijinja_embed::load_templates!(&mut env); - for (kernel, backend, device) in [ + for (kernel, backend, capabilities) in [ ( Kernel::RustCpu { cargo_manifest: "host/Cargo.toml".into(), @@ -275,31 +271,17 @@ mod tests { src: vec![], }, "CPU", - false, + None, ), ( Kernel::RustCuda { cargo_manifest: "host/Cargo.toml".into(), - device_manifest: Some("device/Cargo.toml".into()), - ptx_dir: Some("ptx".into()), cuda_capabilities: Some(vec!["8.0".into(), "9.0".into()]), depends: vec![], src: vec![], }, "CUDA", - true, - ), - ( - Kernel::RustCuda { - cargo_manifest: "host/Cargo.toml".into(), - device_manifest: None, - ptx_dir: None, - cuda_capabilities: None, - depends: vec![], - src: vec![], - }, - "CUDA", - false, + Some("CUDA_CAPABILITIES 8.0 9.0"), ), ] { let mut rendered = vec![]; @@ -307,11 +289,12 @@ mod tests { let rendered = String::from_utf8(rendered).unwrap(); assert!(rendered.contains(&format!("if(GPU_LANG STREQUAL \"{backend}\")"))); assert!(rendered.contains("MANIFEST_PATH \"host/Cargo.toml\"")); - assert_eq!(rendered.contains("DEVICE_MANIFEST"), device); - if device { - assert!(rendered.contains("DEVICE_MANIFEST \"device/Cargo.toml\"")); - assert!(rendered.contains("PTX_DIR \"ptx\"")); - assert!(rendered.contains("CUDA_CAPABILITIES 8.0 9.0")); + assert_eq!( + rendered.contains("CUDA_CAPABILITIES"), + capabilities.is_some() + ); + if let Some(capabilities) = capabilities { + assert!(rendered.contains(capabilities)); } } } diff --git a/kernel-builder/src/pyproject/templates/kernel-component/rust-cuda.cmake b/kernel-builder/src/pyproject/templates/kernel-component/rust-cuda.cmake index 7664b3cf..83ff7ba3 100644 --- a/kernel-builder/src/pyproject/templates/kernel-component/rust-cuda.cmake +++ b/kernel-builder/src/pyproject/templates/kernel-component/rust-cuda.cmake @@ -2,8 +2,6 @@ if(GPU_LANG STREQUAL "CUDA") rust_kernel_component(RUST_KERNEL_LIBS RUST_KERNEL_TARGETS NAME {{ name }} MANIFEST_PATH "{{ manifest_path }}" - {% 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/rust.cmake b/kernel-builder/src/pyproject/templates/rust.cmake index 1c2a056b..1b791f06 100644 --- a/kernel-builder/src/pyproject/templates/rust.cmake +++ b/kernel-builder/src/pyproject/templates/rust.cmake @@ -1,5 +1,47 @@ +# cuda-oxide interop: device crates are declared by the host crate in +# `[[package.metadata.cuda-oxide.device-crates]]`, as `cargo oxide` reads them. +function(_cuda_oxide_device_manifests OUT_VAR MANIFEST) + execute_process( + COMMAND ${CARGO_EXECUTABLE} metadata --format-version 1 --no-deps --offline + --manifest-path ${MANIFEST} + OUTPUT_VARIABLE _JSON + ERROR_VARIABLE _ERROR + RESULT_VARIABLE _RESULT) + if(NOT _RESULT EQUAL 0) + message(FATAL_ERROR "cargo metadata failed for ${MANIFEST}: ${_ERROR}") + endif() + + file(REAL_PATH ${MANIFEST} _MANIFEST) + set(_MANIFESTS) + string(JSON _N_PACKAGES LENGTH "${_JSON}" packages) + math(EXPR _LAST_PACKAGE "${_N_PACKAGES} - 1") + foreach(_I RANGE ${_LAST_PACKAGE}) + string(JSON _PACKAGE_MANIFEST GET "${_JSON}" packages ${_I} manifest_path) + file(REAL_PATH ${_PACKAGE_MANIFEST} _PACKAGE_MANIFEST) + if(NOT _PACKAGE_MANIFEST STREQUAL _MANIFEST) + continue() + endif() + + string(JSON _CRATES ERROR_VARIABLE _NO_CRATES + GET "${_JSON}" packages ${_I} metadata cuda-oxide device-crates) + if(_NO_CRATES) + break() + endif() + get_filename_component(_PACKAGE_DIR ${_PACKAGE_MANIFEST} DIRECTORY) + string(JSON _N_CRATES LENGTH "${_CRATES}") + math(EXPR _LAST_CRATE "${_N_CRATES} - 1") + foreach(_J RANGE ${_LAST_CRATE}) + string(JSON _CRATE_MANIFEST GET "${_CRATES}" ${_J} manifest-path) + get_filename_component(_CRATE_MANIFEST ${_CRATE_MANIFEST} ABSOLUTE BASE_DIR ${_PACKAGE_DIR}) + list(APPEND _MANIFESTS ${_CRATE_MANIFEST}) + endforeach() + endforeach() + + set(${OUT_VAR} ${_MANIFESTS} PARENT_SCOPE) +endfunction() + function(rust_kernel_component LIBS_VAR TARGETS_VAR) - cmake_parse_arguments(KERNEL "" "NAME;MANIFEST_PATH;DEVICE_MANIFEST;PTX_DIR" "CUDA_CAPABILITIES" ${ARGN}) + cmake_parse_arguments(KERNEL "" "NAME;MANIFEST_PATH" "CUDA_CAPABILITIES" ${ARGN}) if(NOT KERNEL_NAME OR NOT KERNEL_MANIFEST_PATH) message(FATAL_ERROR "rust_kernel_component: NAME and MANIFEST_PATH are required") @@ -25,20 +67,20 @@ function(rust_kernel_component LIBS_VAR TARGETS_VAR) set(ALL_GPU_ARCHS ${_ALL_GPU_ARCHS} PARENT_SCOPE) endif() - set(_RUST_KERNEL_DEPENDS) + set(_HOST_ENV) + if(GPU_LANG STREQUAL "CUDA") + _cuda_oxide_device_manifests(_DEVICE_MANIFESTS + ${CMAKE_CURRENT_SOURCE_DIR}/${KERNEL_MANIFEST_PATH}) + endif() - if(KERNEL_DEVICE_MANIFEST) + if(_DEVICE_MANIFESTS) # 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) + if(NOT CUDA_OXIDE_BACKEND) + message(FATAL_ERROR "rust_kernel_component: cuda-oxide device crates require CUDA_OXIDE_BACKEND") 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(_PTX_DIR ${_DEVICE_TARGET_DIR}/ptx) set(_DEVICE_ENV) if(_KERNEL_ARCHS) @@ -50,20 +92,27 @@ function(rust_kernel_component LIBS_VAR TARGETS_VAR) list(APPEND _DEVICE_ENV "CUDA_OXIDE_TARGET=sm_${_OXIDE_ARCH}") endif() + set(_DEVICE_COMMANDS) + foreach(_DEVICE_MANIFEST ${_DEVICE_MANIFESTS}) + # Flags match what `cargo oxide build` passes to the codegen backend. + list(APPEND _DEVICE_COMMANDS + COMMAND ${CMAKE_COMMAND} -E env + "CUDA_OXIDE_PTX_DIR=${_PTX_DIR}" + "RUSTFLAGS=-Zcodegen-backend=${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}) + endforeach() + 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} + ${_DEVICE_COMMANDS} 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) + list(APPEND _HOST_ENV "KERNEL_BUILDER_PTX_DIR=${_PTX_DIR}") endif() string(REPLACE "-" "_" _LIB_NAME ${KERNEL_NAME}) @@ -79,15 +128,18 @@ function(rust_kernel_component LIBS_VAR TARGETS_VAR) COMMAND ${CMAKE_COMMAND} -E env "PATH=${_PYTHON_BIN_DIR}:$ENV{PATH}" "KERNEL_BUILDER_GPU_LANG=${GPU_LANG}" "KERNEL_BUILDER_CUDA_ARCHS=${_KERNEL_ARCHS}" + ${_HOST_ENV} ${CARGO_EXECUTABLE} rustc --release --locked --lib --crate-type staticlib --manifest-path ${CMAKE_CURRENT_SOURCE_DIR}/${KERNEL_MANIFEST_PATH} --target-dir ${_CARGO_TARGET_DIR} - DEPENDS ${_RUST_KERNEL_DEPENDS} BYPRODUCTS ${_STATICLIB} WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} COMMENT "Building Rust kernel ${KERNEL_NAME} with cargo" VERBATIM ) + if(_DEVICE_MANIFESTS) + add_dependencies(${KERNEL_NAME}_cargo_build ${KERNEL_NAME}_oxide_device_build) + endif() add_library(${KERNEL_NAME}_rust STATIC IMPORTED GLOBAL) set_target_properties(${KERNEL_NAME}_rust PROPERTIES IMPORTED_LOCATION ${_STATICLIB}) @@ -141,10 +193,5 @@ function(target_link_rust_kernels TARGET) "$" Threads::Threads ${CMAKE_DL_LIBS}) - # Cargo's native-library link directives do not propagate out of a staticlib. - if(GPU_LANG STREQUAL "CUDA") - find_package(CUDAToolkit REQUIRED) - target_link_libraries(${TARGET} PRIVATE CUDA::cuda_driver) - endif() _restrict_rust_exports(${TARGET}) endfunction() diff --git a/kernel-builder/tests/rust-components.cmake b/kernel-builder/tests/rust-components.cmake index 72cf2dc8..6b435009 100644 --- a/kernel-builder/tests/rust-components.cmake +++ b/kernel-builder/tests/rust-components.cmake @@ -1,11 +1,17 @@ -# Configure Rust targets without CUDA or compiling crates. -# Run: cmake -P kernel-builder/tests/rust-components.cmake +# Configure Rust targets without CUDA or compiling crates (needs `cargo` for +# `cargo metadata`). Run: cmake -P kernel-builder/tests/rust-components.cmake cmake_minimum_required(VERSION 3.24) if(CMAKE_SCRIPT_MODE_FILE) set(_TEST_DIR "${CMAKE_CURRENT_BINARY_DIR}/rust-components-test") file(MAKE_DIRECTORY "${_TEST_DIR}") file(WRITE "${_TEST_DIR}/CMakeLists.txt" "cmake_minimum_required(VERSION 3.24)\nproject(rust_components NONE)\ninclude(\"${CMAKE_CURRENT_LIST_FILE}\")\n") + # Host crates with and without cuda-oxide device crates. + set(_PACKAGE "[package]\nname = \"example\"\nversion = \"0.0.0\"\nedition = \"2021\"\n[lib]\npath = \"lib.rs\"\n[workspace]\n") + file(WRITE "${_TEST_DIR}/plain/Cargo.toml" "${_PACKAGE}") + file(WRITE "${_TEST_DIR}/oxide/Cargo.toml" "${_PACKAGE}" + "[[package.metadata.cuda-oxide.device-crates]]\nmanifest-path = \"device/Cargo.toml\"\n") + file(TOUCH "${_TEST_DIR}/plain/lib.rs" "${_TEST_DIR}/oxide/lib.rs") foreach(_CASE cpu cuda host-only unsupported-arch missing-backend) execute_process(COMMAND "${CMAKE_COMMAND}" -S "${_TEST_DIR}" -G Ninja -B "${_TEST_DIR}/${_CASE}" "-DCASE=${_CASE}" @@ -15,7 +21,7 @@ if(CMAKE_SCRIPT_MODE_FILE) message(FATAL_ERROR "Expected unsupported architecture error: ${_OUTPUT}${_ERROR}") endif() elseif(_CASE STREQUAL "missing-backend") - if(_RESULT EQUAL 0 OR NOT _ERROR MATCHES "requires CUDA_OXIDE_BACKEND") + if(_RESULT EQUAL 0 OR NOT _ERROR MATCHES "require CUDA_OXIDE_BACKEND") message(FATAL_ERROR "Expected missing backend error: ${_OUTPUT}${_ERROR}") endif() elseif(NOT _RESULT EQUAL 0) @@ -35,19 +41,21 @@ set(Python_EXECUTABLE python) set(GPU_LANG CUDA) set(CUDA_ARCHS "8.0;8.9;9.0") set(CUDA_KERNEL_ARCHS "8.0;8.9;9.0+PTX") -set(ENV{CUDA_OXIDE_BACKEND} /test/backend.so) -set(_ARGS CUDA_CAPABILITIES 8.0 DEVICE_MANIFEST device/Cargo.toml) +set(CUDA_OXIDE_BACKEND /test/backend.so) +set(_MANIFEST oxide/Cargo.toml) +set(_ARGS CUDA_CAPABILITIES 8.0) if(CASE STREQUAL "cpu") set(GPU_LANG CPU) set(_ARGS) elseif(CASE STREQUAL "host-only") + set(_MANIFEST plain/Cargo.toml) set(_ARGS) elseif(CASE STREQUAL "unsupported-arch") - set(_ARGS CUDA_CAPABILITIES 6.0 DEVICE_MANIFEST device/Cargo.toml) + set(_ARGS CUDA_CAPABILITIES 6.0) elseif(CASE STREQUAL "missing-backend") - unset(ENV{CUDA_OXIDE_BACKEND}) + unset(CUDA_OXIDE_BACKEND) endif() -rust_kernel_component(LIBS TARGETS NAME example MANIFEST_PATH Cargo.toml ${_ARGS}) +rust_kernel_component(LIBS TARGETS NAME example MANIFEST_PATH ${_MANIFEST} ${_ARGS}) if(NOT TARGET example_rust OR NOT TARGET example_cargo_build) message(FATAL_ERROR "Missing shared Cargo/static-library targets") endif() diff --git a/kernels-common/src/config/mod.rs b/kernels-common/src/config/mod.rs index e566dd65..d17fa16f 100644 --- a/kernels-common/src/config/mod.rs +++ b/kernels-common/src/config/mod.rs @@ -356,11 +356,10 @@ pub enum Kernel { depends: Vec, src: Vec, }, - RustCuda { + /// Path of the host crate's `Cargo.toml`, relative to the kernel directory. + /// cuda-oxide device crates are read from its `[package.metadata.cuda-oxide]`. cargo_manifest: String, - device_manifest: Option, - ptx_dir: Option, cuda_capabilities: Option>, depends: Vec, src: Vec, @@ -747,7 +746,7 @@ mod tests { } #[test] - fn v6_rust_cuda_round_trip() { + fn v6_rust_cuda_converts() { let config = r#" [general] name = "rust-cuda" @@ -758,36 +757,19 @@ mod tests { [tvm-ffi] [kernel.gpu_kernel] language = "rust-cuda" - cargo-manifest = "host/Cargo.toml" - device-manifest = "device/Cargo.toml" - ptx-dir = "ptx" + cargo-manifest = "Cargo.toml" cuda-capabilities = ["8.0"] depends = [] - src = ["host/src", "device/src"] + src = ["src"] "#; let build = Build::try_from(toml::from_str::(config).unwrap()).unwrap(); - assert_eq!(build.kernels["gpu_kernel"].language(), Language::Rust); - assert_eq!(build.kernels["gpu_kernel"].backend(), Backend::Cuda); - let serialized = toml::to_string(&v6::Build::from(build)).unwrap(); - let reparsed: v6::Build = toml::from_str(&serialized).unwrap(); - match &reparsed.kernels["gpu_kernel"] { - v6::Kernel::RustCuda { - cargo_manifest, - device_manifest, - ptx_dir, - cuda_capabilities, - .. - } => { - assert_eq!(cargo_manifest, "host/Cargo.toml"); - assert_eq!(device_manifest.as_deref(), Some("device/Cargo.toml")); - assert_eq!(ptx_dir.as_deref(), Some("ptx")); - assert_eq!( - cuda_capabilities.as_deref(), - Some(["8.0".to_string()].as_slice()) - ); - } - _ => panic!("expected Rust CUDA kernel"), - } + let kernel = &build.kernels["gpu_kernel"]; + assert_eq!(kernel.language(), Language::Rust); + assert_eq!(kernel.backend(), Backend::Cuda); + assert!(matches!( + kernel, + Kernel::RustCuda { cuda_capabilities: Some(caps), .. } if caps == &["8.0"] + )); } #[test] diff --git a/kernels-common/src/config/v6.rs b/kernels-common/src/config/v6.rs index b924640d..4c525ed3 100644 --- a/kernels-common/src/config/v6.rs +++ b/kernels-common/src/config/v6.rs @@ -166,12 +166,9 @@ pub enum Kernel { depends: Vec, src: Vec, }, - #[serde(rename_all = "kebab-case")] RustCuda { cargo_manifest: String, - device_manifest: Option, - ptx_dir: Option, cuda_capabilities: Option>, depends: Vec, src: Vec, @@ -394,18 +391,13 @@ impl From for super::Kernel { include, src, }, - Kernel::RustCuda { cargo_manifest, - device_manifest, - ptx_dir, cuda_capabilities, depends, src, } => super::Kernel::RustCuda { cargo_manifest, - device_manifest, - ptx_dir, cuda_capabilities, depends, src, @@ -645,18 +637,13 @@ impl From for Kernel { depends, src, }, - super::Kernel::RustCuda { cargo_manifest, - device_manifest, - ptx_dir, cuda_capabilities, depends, src, } => Kernel::RustCuda { cargo_manifest, - device_manifest, - ptx_dir, cuda_capabilities, depends, src, diff --git a/nix-builder/lib/build.nix b/nix-builder/lib/build.nix index 025c3fbf..a43fd6ca 100644 --- a/nix-builder/lib/build.nix +++ b/nix-builder/lib/build.nix @@ -202,9 +202,6 @@ rec { kernelProvenance ; - rustCudaDevice = lib.any (kernel: kernel.language == "rust-cuda" && kernel ? device-manifest) ( - lib.attrValues kernels - ); cargoLock = if hasRustKernels then src + "/Cargo.lock" else null; rustGitHashes = if hasRustKernels && builtins.pathExists (src + "/rust-git-hashes.json") then diff --git a/nix-builder/lib/extension/tvm-ffi/arch.nix b/nix-builder/lib/extension/tvm-ffi/arch.nix index 6ffccf2b..9b0881b4 100644 --- a/nix-builder/lib/extension/tvm-ffi/arch.nix +++ b/nix-builder/lib/extension/tvm-ffi/arch.nix @@ -62,9 +62,6 @@ # the build has no Rust kernels. cargoLock ? null, - # Compile CUDA device crates with cuda-oxide and its matching rustc. - rustCudaDevice ? false, - # Fixed-output hashes for Git dependencies in Cargo.lock. rustGitHashes ? { }, @@ -139,45 +136,33 @@ let metalSupport = buildConfig.metal or false; rustSupport = cargoLock != null; - rustCuda = import ./rust-cuda.nix { - inherit - pkgs - lib - stdenv - cudaPackages - ; - }; + # Rust CUDA kernels are built with cuda-oxide and its matching rustc. + rustCuda = import ./rust-cuda.nix { inherit pkgs; }; # rpaths are stripped from kernels to make them portable, but that # also means that in a Nix environment the CUDA/ROCm/oneAPI dependencies # cannot be located anymore, so pass them to hooks that load the kernel. - libraryPath = - lib.makeLibraryPath ( - map lib.getLib ( - lib.optionals cudaSupport ( - with cudaPackages; - [ - cuda_cudart - libcublas - libcusolver - libcusparse - ] - ) - ++ lib.optionals rocmSupport [ clr ] - ++ lib.optionals xpuSupport ( - with xpuPackages; - [ - intel-oneapi-compiler-dpcpp-cpp-runtime - intel-oneapi-compiler-shared-runtime - ] - ) + libraryPath = lib.makeLibraryPath ( + map lib.getLib ( + lib.optionals cudaSupport ( + with cudaPackages; + [ + cuda_cudart + libcublas + libcusolver + libcusparse + ] + ) + ++ lib.optionals rocmSupport [ clr ] + ++ lib.optionals xpuSupport ( + with xpuPackages; + [ + intel-oneapi-compiler-dpcpp-cpp-runtime + intel-oneapi-compiler-shared-runtime + ] ) ) - + - lib.optionalString (rustSupport && cudaSupport) - # Import checks run without a GPU driver in the build sandbox. The stub is - # only used by the check hooks, never bundled in the redistributable kernel. - ":${lib.getOutput "stubs" cudaPackages.cuda_cudart}/lib/stubs"; + ); provenanceFlags = import ../provenance-flags.nix { inherit lib kernelProvenance; }; @@ -271,11 +256,11 @@ stdenv.mkDerivation ( ++ lib.optionals rustSupport [ rustPlatform.cargoSetupHook ] - ++ lib.optionals (rustSupport && !rustCudaDevice) [ + ++ lib.optionals (rustSupport && !cudaSupport) [ cargo rustc ] - ++ lib.optionals rustCudaDevice [ rustCuda.toolchain ]; + ++ lib.optionals (rustSupport && cudaSupport) [ rustCuda.toolchain ]; buildInputs = [ python3.pkgs.tvm-ffi @@ -325,10 +310,6 @@ stdenv.mkDerivation ( // lib.optionalAttrs cudaSupport { CUDAToolkit_ROOT = "${lib.getDev cudaPackages.cuda_nvcc}"; } - // lib.optionalAttrs (rustSupport && cudaSupport) rustCuda.env - // lib.optionalAttrs rustCudaDevice { - CUDA_OXIDE_BACKEND = "${rustCuda.backend}/lib/librustc_codegen_cuda.so"; - } // lib.optionalAttrs xpuSupport { MKLROOT = oneapi-torch-dev; SYCL_ROOT = oneapi-torch-dev; @@ -369,6 +350,9 @@ stdenv.mkDerivation ( (lib.cmakeFeature "CMAKE_CUDA_HOST_COMPILER" "${stdenv.cc}/bin/g++") (lib.cmakeFeature "GPU_LANG" "CUDA") ] + ++ lib.optionals (rustSupport && cudaSupport) [ + (lib.cmakeFeature "CUDA_OXIDE_BACKEND" "${rustCuda.backend}/lib/librustc_codegen_cuda.so") + ] ++ lib.optionals rocmSupport [ # Ensure sure that we use HIP from our CLR override and not HIP from # the symlink-joined ROCm toolkit. diff --git a/nix-builder/lib/extension/tvm-ffi/rust-cuda.nix b/nix-builder/lib/extension/tvm-ffi/rust-cuda.nix index cbb20d56..73c91e53 100644 --- a/nix-builder/lib/extension/tvm-ffi/rust-cuda.nix +++ b/nix-builder/lib/extension/tvm-ffi/rust-cuda.nix @@ -1,10 +1,5 @@ # cuda-oxide is a rustc codegen plugin and must use its matching Rust toolchain. -{ - pkgs, - lib, - stdenv, - cudaPackages, -}: +{ pkgs }: let src = pkgs.fetchFromGitHub { owner = "NVIDIA"; @@ -34,27 +29,7 @@ let ]; doCheck = false; }; - cudaHome = pkgs.symlinkJoin { - name = "cuda-home"; - paths = with cudaPackages; [ - (lib.getDev cuda_cudart) - (lib.getLib cuda_cudart) - (lib.getDev cccl) - (lib.getOutput "include" libcurand) - cuda_crt - ]; - }; in { inherit toolchain backend; - env = { - CUDA_HOME = cudaHome; - CUDA_TOOLKIT_PATH = cudaHome; - LIBRARY_PATH = "${lib.getOutput "stubs" cudaPackages.cuda_cudart}/lib/stubs"; - LIBCLANG_PATH = "${lib.getLib pkgs.libclang}/lib"; - BINDGEN_EXTRA_CLANG_ARGS = lib.concatStringsSep " " [ - "-isystem ${lib.getLib pkgs.libclang}/lib/clang/${lib.versions.major pkgs.libclang.version}/include" - "-isystem ${lib.getDev stdenv.cc.libc}/include" - ]; - }; } diff --git a/nix-builder/lib/source-set.nix b/nix-builder/lib/source-set.nix index 36276b75..050606f8 100644 --- a/nix-builder/lib/source-set.nix +++ b/nix-builder/lib/source-set.nix @@ -43,11 +43,7 @@ let else fileset.empty; # Rust kernels also need their Cargo manifest. - kernelSrc = - kernel: - kernel.src - ++ lib.optional (kernel ? cargo-manifest) kernel.cargo-manifest - ++ lib.optional (kernel ? device-manifest) kernel.device-manifest; + kernelSrc = kernel: kernel.src ++ lib.optional (kernel ? cargo-manifest) kernel.cargo-manifest; kernelsSrc = fileset.unions ( lib.flatten (lib.mapAttrsToList (name: kernel: map (nameToPath path) (kernelSrc kernel)) kernels) ); From 5a91c6da5455c162909fe9b9acfd751648c79290 Mon Sep 17 00:00:00 2001 From: David Holtz Date: Tue, 20 Oct 2026 20:33:14 +0000 Subject: [PATCH 03/14] feat: add rust gpu relu example --- examples/kernels/flake.nix | 6 + examples/kernels/relu-rust-cuda/CARD.md | 65 ++++++ examples/kernels/relu-rust-cuda/Cargo.lock | 185 ++++++++++++++++++ examples/kernels/relu-rust-cuda/Cargo.toml | 3 + examples/kernels/relu-rust-cuda/build.toml | 23 +++ examples/kernels/relu-rust-cuda/flake.nix | 17 ++ .../relu-rust-cuda/relu-device/Cargo.toml | 11 ++ .../relu-rust-cuda/relu-device/src/main.rs | 13 ++ .../kernels/relu-rust-cuda/relu-rs/Cargo.toml | 19 ++ .../kernels/relu-rust-cuda/relu-rs/src/lib.rs | 103 ++++++++++ .../relu-rust-cuda/rust-git-hashes.json | 5 + .../kernels/relu-rust-cuda/tests/__init__.py | 0 .../kernels/relu-rust-cuda/tests/test_relu.py | 21 ++ .../tvm-ffi-ext/relu_rust_cuda/__init__.py | 19 ++ 14 files changed, 490 insertions(+) create mode 100644 examples/kernels/relu-rust-cuda/CARD.md create mode 100644 examples/kernels/relu-rust-cuda/Cargo.lock create mode 100644 examples/kernels/relu-rust-cuda/Cargo.toml create mode 100644 examples/kernels/relu-rust-cuda/build.toml create mode 100644 examples/kernels/relu-rust-cuda/flake.nix create mode 100644 examples/kernels/relu-rust-cuda/relu-device/Cargo.toml create mode 100644 examples/kernels/relu-rust-cuda/relu-device/src/main.rs create mode 100644 examples/kernels/relu-rust-cuda/relu-rs/Cargo.toml create mode 100644 examples/kernels/relu-rust-cuda/relu-rs/src/lib.rs create mode 100644 examples/kernels/relu-rust-cuda/rust-git-hashes.json create mode 100644 examples/kernels/relu-rust-cuda/tests/__init__.py create mode 100644 examples/kernels/relu-rust-cuda/tests/test_relu.py create mode 100644 examples/kernels/relu-rust-cuda/tvm-ffi-ext/relu_rust_cuda/__init__.py diff --git a/examples/kernels/flake.nix b/examples/kernels/flake.nix index 1d5b5bed..32d65679 100644 --- a/examples/kernels/flake.nix +++ b/examples/kernels/flake.nix @@ -121,6 +121,12 @@ drv = sys: out: out.packages.${sys}.redistributable.${"tvm-ffi${tvmFfiVersion}-${cudaVersion}-${sys}"}; } + { + name = "relu-rust-cuda-kernel"; + path = ./relu-rust-cuda; + drv = + sys: out: out.packages.${sys}.redistributable.${"tvm-ffi${tvmFfiVersion}-${cudaVersion}-${sys}"}; + } { name = "relu-kernel-cpu"; path = ./relu; diff --git a/examples/kernels/relu-rust-cuda/CARD.md b/examples/kernels/relu-rust-cuda/CARD.md new file mode 100644 index 00000000..2868c79f --- /dev/null +++ b/examples/kernels/relu-rust-cuda/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-cuda/Cargo.lock b/examples/kernels/relu-rust-cuda/Cargo.lock new file mode 100644 index 00000000..0496f6e2 --- /dev/null +++ b/examples/kernels/relu-rust-cuda/Cargo.lock @@ -0,0 +1,185 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "cfg-if" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e7648175b45a9a48536d676f68d918270699102aa8dab5496df06904c914600" + +[[package]] +name = "cuda-device" +version = "0.2.1" +source = "git+https://github.com/NVIDIA/cuda-rust.git?rev=6d47ae49cb256a07948fb499c20c3eb9b2a96268#6d47ae49cb256a07948fb499c20c3eb9b2a96268" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +source = "git+https://github.com/NVIDIA/cuda-rust.git?rev=6d47ae49cb256a07948fb499c20c3eb9b2a96268#6d47ae49cb256a07948fb499c20c3eb9b2a96268" +dependencies = [ + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "syn 2.0.119", +] + +[[package]] +name = "cudarc" +version = "0.19.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7359bf1de037ddada72729c3acec9b040f3468400c6e9e740b01ba0cd2e89562" +dependencies = [ + "libloading", +] + +[[package]] +name = "libloading" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "754ca22de805bb5744484a5b151a9e1a8e837d5dc232c2d7d8c2e3492edc8b60" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[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 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 = "relu-device" +version = "0.1.0" +dependencies = [ + "cuda-device", +] + +[[package]] +name = "relu-rs" +version = "0.1.0" +dependencies = [ + "cudarc", + "tvm-ffi", + "tvm-ffi-sys", +] + +[[package]] +name = "reserved-oxide-symbols" +version = "0.2.1" +source = "git+https://github.com/NVIDIA/cuda-rust.git?rev=6d47ae49cb256a07948fb499c20c3eb9b2a96268#6d47ae49cb256a07948fb499c20c3eb9b2a96268" + +[[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 = "tvm-ffi" +version = "0.1.14-post1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2573e03b4bc46bfb606b2801c7a94d332b7640acbe5936a0c64935d664ff9875" +dependencies = [ + "paste", + "tvm-ffi-macros", + "tvm-ffi-sys", +] + +[[package]] +name = "tvm-ffi-macros" +version = "0.1.14-post1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c0f660be122ead0def43f986c9ec4aa635bb9e285820adbc449b8230426933f" +dependencies = [ + "proc-macro-error", + "proc-macro2", + "quote", + "syn 1.0.109", +] + +[[package]] +name = "tvm-ffi-sys" +version = "0.1.14-post1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0b29bc4a713c7e86bcb98e1655832058c272eb336f8d62cece8344cb2666ee2" + +[[package]] +name = "unicode-ident" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563" + +[[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" diff --git a/examples/kernels/relu-rust-cuda/Cargo.toml b/examples/kernels/relu-rust-cuda/Cargo.toml new file mode 100644 index 00000000..42e2fea8 --- /dev/null +++ b/examples/kernels/relu-rust-cuda/Cargo.toml @@ -0,0 +1,3 @@ +[workspace] +members = ["relu-rs", "relu-device"] +resolver = "2" diff --git a/examples/kernels/relu-rust-cuda/build.toml b/examples/kernels/relu-rust-cuda/build.toml new file mode 100644 index 00000000..7af88b44 --- /dev/null +++ b/examples/kernels/relu-rust-cuda/build.toml @@ -0,0 +1,23 @@ +[general] +name = "relu-rust-cuda" +version = 1 +edition = 6 +license = "Apache-2.0" +backends = ["cuda"] + +[general.hub] +repo-id = "kernels-test/relu-rust-cuda" + +[tvm-ffi] + +[kernel.relu_rust_cuda] +language = "rust-cuda" +cargo-manifest = "relu-rs/Cargo.toml" +cuda-capabilities = ["8.0"] +depends = [] +src = [ + "Cargo.toml", + "relu-rs/src/lib.rs", + "relu-device/Cargo.toml", + "relu-device/src/main.rs", +] diff --git a/examples/kernels/relu-rust-cuda/flake.nix b/examples/kernels/relu-rust-cuda/flake.nix new file mode 100644 index 00000000..e69775eb --- /dev/null +++ b/examples/kernels/relu-rust-cuda/flake.nix @@ -0,0 +1,17 @@ +{ + description = "Flake for a ReLU kernel written in Rust with cuda-oxide"; + + inputs = { + kernel-builder.url = "path:../../.."; + }; + + outputs = + { + self, + kernel-builder, + }: + kernel-builder.lib.genKernelFlakeOutputs { + inherit self; + path = ./.; + }; +} diff --git a/examples/kernels/relu-rust-cuda/relu-device/Cargo.toml b/examples/kernels/relu-rust-cuda/relu-device/Cargo.toml new file mode 100644 index 00000000..649fd8e3 --- /dev/null +++ b/examples/kernels/relu-rust-cuda/relu-device/Cargo.toml @@ -0,0 +1,11 @@ +[package] +name = "relu-device" +version = "0.1.0" +edition = "2024" +license = "Apache-2.0" +publish = false + +# Kernel-only crate, compiled to PTX by the cuda-oxide codegen backend. Without +# `cuda-host` the kernel macros emit no host glue. +[dependencies] +cuda-device = { git = "https://github.com/NVIDIA/cuda-rust.git", rev = "6d47ae49cb256a07948fb499c20c3eb9b2a96268" } diff --git a/examples/kernels/relu-rust-cuda/relu-device/src/main.rs b/examples/kernels/relu-rust-cuda/relu-device/src/main.rs new file mode 100644 index 00000000..c3c57fea --- /dev/null +++ b/examples/kernels/relu-rust-cuda/relu-device/src/main.rs @@ -0,0 +1,13 @@ +use cuda_device::{DisjointSlice, kernel, thread}; + +#[kernel] +pub fn relu(x: &[f32], mut out: DisjointSlice) { + let idx = thread::index_1d(); + let i = idx.get(); + if let Some(slot) = out.get_mut(idx) { + let v = x[i]; + *slot = if v > 0.0 { v } else { 0.0 }; + } +} + +fn main() {} diff --git a/examples/kernels/relu-rust-cuda/relu-rs/Cargo.toml b/examples/kernels/relu-rust-cuda/relu-rs/Cargo.toml new file mode 100644 index 00000000..22365719 --- /dev/null +++ b/examples/kernels/relu-rust-cuda/relu-rs/Cargo.toml @@ -0,0 +1,19 @@ +[package] +name = "relu-rs" +version = "0.1.0" +edition = "2021" +license = "Apache-2.0" +publish = false + +# Kernel-builder compiles these with cuda-oxide before this crate and passes +# the PTX directory as `KERNEL_BUILDER_PTX_DIR`. +[[package.metadata.cuda-oxide.device-crates]] +manifest-path = "../relu-device/Cargo.toml" + +[lib] +name = "relu_rust_cuda" + +[dependencies] +cudarc = { version = "0.19.9", default-features = false, features = ["std", "driver", "dynamic-loading", "cuda-12060"] } +tvm-ffi = "=0.1.14-post1" +tvm-ffi-sys = "=0.1.14-post1" diff --git a/examples/kernels/relu-rust-cuda/relu-rs/src/lib.rs b/examples/kernels/relu-rust-cuda/relu-rs/src/lib.rs new file mode 100644 index 00000000..e574c397 --- /dev/null +++ b/examples/kernels/relu-rust-cuda/relu-rs/src/lib.rs @@ -0,0 +1,103 @@ +use std::collections::HashMap; +use std::ffi::{c_void, CString}; +use std::sync::{Arc, Mutex, OnceLock}; + +use cudarc::driver::{result, sys, CudaContext}; +use tvm_ffi::{Error, Result, Tensor, RUNTIME_ERROR, VALUE_ERROR}; +use tvm_ffi_sys::dlpack::{DLDataTypeCode, DLDeviceType}; +use tvm_ffi_sys::TVMFFIEnvGetStream; + +// PTX of the `relu-device` crate, NUL-terminated for `cuModuleLoadData`. +const PTX: &str = concat!( + include_str!(concat!(env!("KERNEL_BUILDER_PTX_DIR"), "/relu_device.ptx")), + "\0" +); + +struct Relu { + ctx: Arc, + function: sys::CUfunction, +} + +unsafe impl Send for Relu {} + +fn cuda_error(e: result::DriverError) -> Error { + Error::new(RUNTIME_ERROR, &format!("{e:?}"), "") +} + +fn value_error(message: &str) -> Error { + Error::new(VALUE_ERROR, message, "") +} + +// Loads the module once per device into the primary context Torch also uses. +fn with_relu(device: i32, f: impl FnOnce(&Relu) -> Result) -> Result { + static RELU: OnceLock>> = OnceLock::new(); + let mut loaded = RELU.get_or_init(Default::default).lock().unwrap(); + if !loaded.contains_key(&device) { + let ctx = CudaContext::new(device as usize).map_err(cuda_error)?; + ctx.bind_to_thread().map_err(cuda_error)?; + let function = unsafe { + let module = result::module::load_data(PTX.as_ptr().cast()).map_err(cuda_error)?; + result::module::get_function(module, CString::new("relu").unwrap()) + .map_err(cuda_error)? + }; + loaded.insert(device, Relu { ctx, function }); + } + f(&loaded[&device]) +} + +fn relu_rust_cuda(x: Tensor, out: Tensor) -> Result<()> { + let f32_type = |t: &Tensor| { + let dtype = t.dtype(); + dtype.code == DLDataTypeCode::kDLFloat as u8 && dtype.bits == 32 && dtype.lanes == 1 + }; + if !f32_type(&x) || !f32_type(&out) { + return Err(value_error("input and output must be float32")); + } + if !x.is_contiguous() || !out.is_contiguous() { + return Err(value_error("input and output must be contiguous")); + } + if x.numel() != out.numel() { + return Err(value_error(&format!( + "input and output must have the same number of elements, got {} and {}", + x.numel(), + out.numel() + ))); + } + let device = x.device(); + if device.device_type != DLDeviceType::kDLCUDA || out.device() != device { + return Err(value_error( + "input and output must be on the same CUDA device", + )); + } + if x.numel() == 0 { + return Ok(()); + } + + with_relu(device.device_id, |relu| { + relu.ctx.bind_to_thread().map_err(cuda_error)?; + let stream = unsafe { TVMFFIEnvGetStream(DLDeviceType::kDLCUDA as i32, device.device_id) }; + // Slices are passed to cuda-oxide kernels as (pointer, length) pairs. + let (mut x_ptr, mut out_ptr) = (x.data_ptr() as u64, out.data_ptr() as u64); + let mut len = x.numel() as u64; + let mut params = [ + (&mut x_ptr as *mut u64).cast::(), + (&mut len as *mut u64).cast(), + (&mut out_ptr as *mut u64).cast(), + (&mut len as *mut u64).cast(), + ]; + let blocks = len.div_ceil(256) as u32; + unsafe { + result::launch_kernel( + relu.function, + (blocks, 1, 1), + (256, 1, 1), + 0, + stream.cast(), + &mut params, + ) + } + .map_err(cuda_error) + }) +} + +tvm_ffi::tvm_ffi_dll_export_typed_func!(relu_rust_cuda, relu_rust_cuda); diff --git a/examples/kernels/relu-rust-cuda/rust-git-hashes.json b/examples/kernels/relu-rust-cuda/rust-git-hashes.json new file mode 100644 index 00000000..7e8a2a46 --- /dev/null +++ b/examples/kernels/relu-rust-cuda/rust-git-hashes.json @@ -0,0 +1,5 @@ +{ + "cuda-device-0.2.1": "sha256-KxIPPzDBKeHyknrpRg1ekNliEcc44yBWUtSmrGnApOo=", + "cuda-macros-0.2.1": "sha256-KxIPPzDBKeHyknrpRg1ekNliEcc44yBWUtSmrGnApOo=", + "reserved-oxide-symbols-0.2.1": "sha256-KxIPPzDBKeHyknrpRg1ekNliEcc44yBWUtSmrGnApOo=" +} diff --git a/examples/kernels/relu-rust-cuda/tests/__init__.py b/examples/kernels/relu-rust-cuda/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/examples/kernels/relu-rust-cuda/tests/test_relu.py b/examples/kernels/relu-rust-cuda/tests/test_relu.py new file mode 100644 index 00000000..7b642fac --- /dev/null +++ b/examples/kernels/relu-rust-cuda/tests/test_relu.py @@ -0,0 +1,21 @@ +import kernels +import pytest +import torch +import torch.nn.functional as F + +relu_rust_cuda = kernels.get_kernel("kernels-test/relu-rust-cuda", version=1) + + +@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_rust_cuda.relu(x, torch.empty_like(x)) + ) + + +@pytest.mark.kernels_ci +def test_rejects_mismatched_output(): + x = torch.randn(16, dtype=torch.float32, device="cuda") + with pytest.raises(Exception, match="same number of elements"): + relu_rust_cuda.relu(x, torch.empty(32, dtype=torch.float32, device="cuda")) diff --git a/examples/kernels/relu-rust-cuda/tvm-ffi-ext/relu_rust_cuda/__init__.py b/examples/kernels/relu-rust-cuda/tvm-ffi-ext/relu_rust_cuda/__init__.py new file mode 100644 index 00000000..c665c751 --- /dev/null +++ b/examples/kernels/relu-rust-cuda/tvm-ffi-ext/relu_rust_cuda/__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 == "cuda": + ops.relu_rust_cuda(x_t, out_t) + else: + raise NotImplementedError(f"Unsupported device type: {device.type}") + + return out + + +__all__ = ["relu"] From 9fedf1756f3c00ba431efcd5ed2d6b38e96d22c7 Mon Sep 17 00:00:00 2001 From: David Holtz Date: Tue, 20 Oct 2026 21:30:59 +0000 Subject: [PATCH 04/14] fix: package cuda-oxide and build rust cuda with the kernel stdenv --- docs/source/builder/writing-kernels.md | 21 ++++++++++ .../src/pyproject/templates/rust.cmake | 7 +--- nix-builder/lib/cache.nix | 2 + nix-builder/lib/extension/tvm-ffi/arch.nix | 14 ++----- nix-builder/lib/source-set.nix | 1 - nix-builder/overlay.nix | 3 ++ .../cuda-oxide/default.nix} | 39 ++++++++++++++----- 7 files changed, 61 insertions(+), 26 deletions(-) rename nix-builder/{lib/extension/tvm-ffi/rust-cuda.nix => pkgs/cuda-oxide/default.nix} (54%) diff --git a/docs/source/builder/writing-kernels.md b/docs/source/builder/writing-kernels.md index 9c57b294..c819bdb2 100644 --- a/docs/source/builder/writing-kernels.md +++ b/docs/source/builder/writing-kernels.md @@ -377,6 +377,27 @@ in Git so Nix can include them in the build. Rerun the command whenever Git dependencies in `Cargo.lock` change. Registry dependencies use the checksums in `Cargo.lock` and do not require this file. +#### rust-cuda (experimental, and subject to change) + +Rust CUDA kernels use [cuda-oxide](https://github.com/NVIDIA/cuda-rust) and +take the same options as `rust-cpu`, plus: + +- `cuda-capabilities` (optional): same as for `cpp-cuda`. PTX is generated for the lowest supported capability, so set this when kernels use newer instructions. + +`cargo-manifest` points to the host crate. It lists its device crates, which must be in the same Cargo workspace, in its metadata: + +```toml +[[package.metadata.cuda-oxide.device-crates]] +manifest-path = "../relu-device/Cargo.toml" +``` + +The device crates are built first and the host crate finds their PTX in `KERNEL_BUILDER_PTX_DIR`. + +Use the cuda-rust revision pinned in +[`cuda-oxide`](https://github.com/huggingface/kernels/blob/main/nix-builder/pkgs/cuda-oxide/default.nix). +See [`relu-rust-cuda`](https://github.com/huggingface/kernels/tree/main/examples/kernels/relu-rust-cuda) +for an example. + ## Torch bindings ### Defining bindings diff --git a/kernel-builder/src/pyproject/templates/rust.cmake b/kernel-builder/src/pyproject/templates/rust.cmake index 1b791f06..5feb65ca 100644 --- a/kernel-builder/src/pyproject/templates/rust.cmake +++ b/kernel-builder/src/pyproject/templates/rust.cmake @@ -65,14 +65,13 @@ function(rust_kernel_component LIBS_VAR TARGETS_VAR) accumulate_gpu_archs(_ALL_GPU_ARCHS "${ALL_GPU_ARCHS}" "${_KERNEL_ARCHS}") set(ALL_GPU_ARCHS ${_ALL_GPU_ARCHS} PARENT_SCOPE) - endif() - set(_HOST_ENV) - if(GPU_LANG STREQUAL "CUDA") _cuda_oxide_device_manifests(_DEVICE_MANIFESTS ${CMAKE_CURRENT_SOURCE_DIR}/${KERNEL_MANIFEST_PATH}) endif() + set(_HOST_ENV) + if(_DEVICE_MANIFESTS) # Build cuda-oxide device code before the host crate embeds its PTX. if(NOT CUDA_OXIDE_BACKEND) @@ -126,8 +125,6 @@ function(rust_kernel_component LIBS_VAR TARGETS_VAR) 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}" ${_HOST_ENV} ${CARGO_EXECUTABLE} rustc --release --locked --lib --crate-type staticlib --manifest-path ${CMAKE_CURRENT_SOURCE_DIR}/${KERNEL_MANIFEST_PATH} diff --git a/nix-builder/lib/cache.nix b/nix-builder/lib/cache.nix index 49c5f56f..22624e2e 100644 --- a/nix-builder/lib/cache.nix +++ b/nix-builder/lib/cache.nix @@ -34,6 +34,8 @@ ++ allOutputs python3.pkgs.tvm-ffi ++ allOutputs ruff ++ lib.optionals (isLinux && cudaSupport) (allOutputs manylinux_2_28.cudaBackendStdenv) + ++ lib.optionals (isLinux && cudaSupport) (allOutputs cuda-oxide.toolchain) + ++ lib.optionals (isLinux && cudaSupport) (allOutputs cuda-oxide.backend) ++ lib.optionals (isLinux && !config.cudaSupport) (allOutputs manylinux_2_28.stdenv) # Only works on recent CUDAs. ++ lib.optionals (!python3.pkgs.nvidia-cutlass-dsl.meta.broken) ( diff --git a/nix-builder/lib/extension/tvm-ffi/arch.nix b/nix-builder/lib/extension/tvm-ffi/arch.nix index 9b0881b4..cba89576 100644 --- a/nix-builder/lib/extension/tvm-ffi/arch.nix +++ b/nix-builder/lib/extension/tvm-ffi/arch.nix @@ -13,6 +13,7 @@ cargo, cmake, cmakeNvccThreadsHook, + cuda-oxide, cuda_nvcc, generate-symbols-hook, get-kernel-check, @@ -136,8 +137,6 @@ let metalSupport = buildConfig.metal or false; rustSupport = cargoLock != null; - # Rust CUDA kernels are built with cuda-oxide and its matching rustc. - rustCuda = import ./rust-cuda.nix { inherit pkgs; }; # rpaths are stripped from kernels to make them portable, but that # also means that in a Nix environment the CUDA/ROCm/oneAPI dependencies @@ -260,7 +259,8 @@ stdenv.mkDerivation ( cargo rustc ] - ++ lib.optionals (rustSupport && cudaSupport) [ rustCuda.toolchain ]; + # Rust CUDA kernels are built with cuda-oxide and its matching rustc. + ++ lib.optionals (rustSupport && cudaSupport) [ cuda-oxide.toolchain ]; buildInputs = [ python3.pkgs.tvm-ffi @@ -340,18 +340,12 @@ stdenv.mkDerivation ( # rpaths anyway. (lib.cmakeBool "CMAKE_SKIP_RPATH" true) ] - # Rust toolchains can propagate a newer compiler onto PATH. Link with the - # selected stdenv so redistributable kernels retain their glibc baseline. - ++ lib.optionals rustSupport [ - (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") ] ++ lib.optionals (rustSupport && cudaSupport) [ - (lib.cmakeFeature "CUDA_OXIDE_BACKEND" "${rustCuda.backend}/lib/librustc_codegen_cuda.so") + (lib.cmakeFeature "CUDA_OXIDE_BACKEND" "${cuda-oxide.backend}/lib/librustc_codegen_cuda.so") ] ++ lib.optionals rocmSupport [ # Ensure sure that we use HIP from our CLR override and not HIP from diff --git a/nix-builder/lib/source-set.nix b/nix-builder/lib/source-set.nix index 050606f8..a1d75909 100644 --- a/nix-builder/lib/source-set.nix +++ b/nix-builder/lib/source-set.nix @@ -55,7 +55,6 @@ fileset.toSource { kernelsSrc lockSet cargoLockSet - (fileset.maybeMissing (path + "/rust-toolchain.toml")) rustGitHashesSet srcSet pySrcSet diff --git a/nix-builder/overlay.nix b/nix-builder/overlay.nix index 99f8a58f..0873f4e5 100644 --- a/nix-builder/overlay.nix +++ b/nix-builder/overlay.nix @@ -39,6 +39,9 @@ final: prev: }; }; + # cuda-oxide toolchain and codegen backend for Rust CUDA kernels. + cuda-oxide = final.callPackage ./pkgs/cuda-oxide { }; + lock-kernel-deps = final.callPackage ./pkgs/lock-kernel-deps { }; # Make the source for a package from our cargo workspace. diff --git a/nix-builder/lib/extension/tvm-ffi/rust-cuda.nix b/nix-builder/pkgs/cuda-oxide/default.nix similarity index 54% rename from nix-builder/lib/extension/tvm-ffi/rust-cuda.nix rename to nix-builder/pkgs/cuda-oxide/default.nix index 73c91e53..7b526cce 100644 --- a/nix-builder/lib/extension/tvm-ffi/rust-cuda.nix +++ b/nix-builder/pkgs/cuda-oxide/default.nix @@ -1,18 +1,40 @@ # cuda-oxide is a rustc codegen plugin and must use its matching Rust toolchain. -{ pkgs }: +{ + fetchFromGitHub, + libffi, + libxml2, + makeRustPlatform, + rust-bin, + zlib, + zstd, +}: + let - src = pkgs.fetchFromGitHub { + src = fetchFromGitHub { owner = "NVIDIA"; repo = "cuda-rust"; rev = "6d47ae49cb256a07948fb499c20c3eb9b2a96268"; hash = "sha256-g9qg2zHLybO8kPLoOJqnYbWjRYtjGV3QVT+X4f8+1fg="; }; - toolchain = pkgs.rust-bin.fromRustupToolchainFile "${src}/cuda-oxide/rust-toolchain.toml"; - platform = pkgs.makeRustPlatform { + + # rust-overlay propagates the compiler of its own package set. Drop it so + # kernels link with their (manylinux) stdenv, as with the rust-cpu toolchain. + toolchain = + (rust-bin.fromRustupToolchainFile "${src}/cuda-oxide/rust-toolchain.toml").overrideAttrs + { + propagatedBuildInputs = [ ]; + depsHostHostPropagated = [ ]; + }; + + rustPlatform = makeRustPlatform { cargo = toolchain; rustc = toolchain; }; - backend = platform.buildRustPackage { +in +{ + inherit toolchain; + + backend = rustPlatform.buildRustPackage { pname = "rustc-codegen-cuda"; version = "0.2.1"; inherit src; @@ -21,15 +43,12 @@ let lockFile = "${src}/cuda-oxide/crates/rustc-codegen-cuda/Cargo.lock"; outputHashes."pliron-0.18.0" = "sha256-HQ/xNfQcIf9pQRrgxCYGFdVfdCMaE4iwj/G9mgAJOWw="; }; - buildInputs = with pkgs; [ + buildInputs = [ libffi libxml2 - zstd zlib + zstd ]; doCheck = false; }; -in -{ - inherit toolchain backend; } From 1cfb75aaa9a574693e52624ba8b73c3d0df03d0f Mon Sep 17 00:00:00 2001 From: David Holtz Date: Tue, 20 Oct 2026 21:53:36 +0000 Subject: [PATCH 05/14] fix: remove low value test --- kernels-common/src/config/mod.rs | 33 -------------------------------- 1 file changed, 33 deletions(-) diff --git a/kernels-common/src/config/mod.rs b/kernels-common/src/config/mod.rs index d17fa16f..5631ab1b 100644 --- a/kernels-common/src/config/mod.rs +++ b/kernels-common/src/config/mod.rs @@ -745,33 +745,6 @@ mod tests { assert_eq!(build.kernels["cpu_kernel"].language(), Language::Rust); } - #[test] - fn v6_rust_cuda_converts() { - let config = r#" - [general] - name = "rust-cuda" - version = 1 - edition = 6 - license = "Apache-2.0" - backends = ["cuda"] - [tvm-ffi] - [kernel.gpu_kernel] - language = "rust-cuda" - cargo-manifest = "Cargo.toml" - cuda-capabilities = ["8.0"] - depends = [] - src = ["src"] - "#; - let build = Build::try_from(toml::from_str::(config).unwrap()).unwrap(); - let kernel = &build.kernels["gpu_kernel"]; - assert_eq!(kernel.language(), Language::Rust); - assert_eq!(kernel.backend(), Backend::Cuda); - assert!(matches!( - kernel, - Kernel::RustCuda { cuda_capabilities: Some(caps), .. } if caps == &["8.0"] - )); - } - #[test] fn v6_rust_kernel_rejects_invalid_config() { let cases = [ @@ -781,12 +754,6 @@ mod tests { "", "require a `[tvm-ffi]` framework", ), - ( - "[tvm-ffi]", - "rust-cuda", - r#"cuda-flags = ["-O3"]"#, - "unknown field `cuda-flags`", - ), ( "[tvm-ffi]", "rust-cpu", From ae983e7ce660d4ed3a1142d6b9a540b7576be47d Mon Sep 17 00:00:00 2001 From: David Holtz Date: Thu, 22 Oct 2026 14:57:06 +0000 Subject: [PATCH 06/14] feat: single crate rust cuda with prebuilt sass --- docs/source/builder/writing-kernels.md | 31 +- examples/kernels/relu-rust-cuda/Cargo.lock | 401 +++++++++++++++++- examples/kernels/relu-rust-cuda/Cargo.toml | 27 +- examples/kernels/relu-rust-cuda/build.toml | 11 +- .../relu-rust-cuda/relu-device/Cargo.toml | 11 - .../relu-rust-cuda/relu-device/src/main.rs | 13 - .../kernels/relu-rust-cuda/relu-rs/Cargo.toml | 19 - .../kernels/relu-rust-cuda/relu-rs/src/lib.rs | 103 ----- .../relu-rust-cuda/rust-git-hashes.json | 16 +- examples/kernels/relu-rust-cuda/src/lib.rs | 44 ++ .../pyproject/templates/rust-fatbin.cpp.in | 16 + .../src/pyproject/templates/rust.cmake | 192 +++++---- kernel-builder/src/pyproject/tvm_ffi/mod.rs | 2 + kernel-builder/tests/rust-components.cmake | 44 +- kernels-cuda/Cargo.toml | 20 + kernels-cuda/src/lib.rs | 188 ++++++++ nix-builder/pkgs/cuda-oxide/default.nix | 16 +- 17 files changed, 847 insertions(+), 307 deletions(-) delete mode 100644 examples/kernels/relu-rust-cuda/relu-device/Cargo.toml delete mode 100644 examples/kernels/relu-rust-cuda/relu-device/src/main.rs delete mode 100644 examples/kernels/relu-rust-cuda/relu-rs/Cargo.toml delete mode 100644 examples/kernels/relu-rust-cuda/relu-rs/src/lib.rs create mode 100644 examples/kernels/relu-rust-cuda/src/lib.rs create mode 100644 kernel-builder/src/pyproject/templates/rust-fatbin.cpp.in create mode 100644 kernels-cuda/Cargo.toml create mode 100644 kernels-cuda/src/lib.rs diff --git a/docs/source/builder/writing-kernels.md b/docs/source/builder/writing-kernels.md index c819bdb2..9b72eb0e 100644 --- a/docs/source/builder/writing-kernels.md +++ b/docs/source/builder/writing-kernels.md @@ -379,23 +379,20 @@ dependencies in `Cargo.lock` change. Registry dependencies use the checksums in #### rust-cuda (experimental, and subject to change) -Rust CUDA kernels use [cuda-oxide](https://github.com/NVIDIA/cuda-rust) and -take the same options as `rust-cpu`, plus: - -- `cuda-capabilities` (optional): same as for `cpp-cuda`. PTX is generated for the lowest supported capability, so set this when kernels use newer instructions. - -`cargo-manifest` points to the host crate. It lists its device crates, which must be in the same Cargo workspace, in its metadata: - -```toml -[[package.metadata.cuda-oxide.device-crates]] -manifest-path = "../relu-device/Cargo.toml" -``` - -The device crates are built first and the host crate finds their PTX in `KERNEL_BUILDER_PTX_DIR`. - -Use the cuda-rust revision pinned in -[`cuda-oxide`](https://github.com/huggingface/kernels/blob/main/nix-builder/pkgs/cuda-oxide/default.nix). -See [`relu-rust-cuda`](https://github.com/huggingface/kernels/tree/main/examples/kernels/relu-rust-cuda) +Rust CUDA kernels are built with +[cuda-oxide](https://github.com/NVIDIA/cuda-rust) and take the same options as +`rust-cpu`, plus: + +- `cuda-capabilities` (optional): same as for `cuda`. The kernels are compiled + to SASS for each capability, so nothing is JIT-compiled at load time. + +The kernels and their host code are a single crate. Load the kernels with the +`kernels-cuda` crate instead of the generated `load`, which fails to link. +An extension can have only one `rust-cuda` component. Use the cuda-rust +revision pinned in +[`cuda-oxide`](https://github.com/huggingface/kernels/blob/main/nix-builder/pkgs/cuda-oxide/default.nix) +and see +[`relu-rust-cuda`](https://github.com/huggingface/kernels/tree/main/examples/kernels/relu-rust-cuda) for an example. ## Torch bindings diff --git a/examples/kernels/relu-rust-cuda/Cargo.lock b/examples/kernels/relu-rust-cuda/Cargo.lock index 0496f6e2..2e9be920 100644 --- a/examples/kernels/relu-rust-cuda/Cargo.lock +++ b/examples/kernels/relu-rust-cuda/Cargo.lock @@ -2,25 +2,133 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "anyhow" +version = "1.0.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" + +[[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 = "cfg-if" version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4e7648175b45a9a48536d676f68d918270699102aa8dab5496df06904c914600" +[[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.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ba106e844045ddc02c20768ac4b0b7cb45f2cc2ecbd833ac9f5dacb183c2055" +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-artifact-finalizer" +version = "0.2.1" +source = "git+https://github.com/drbh/cuda-oxide.git?rev=d9d4e379e9a80d1ce7c681d88a086aacec95ba23#d9d4e379e9a80d1ce7c681d88a086aacec95ba23" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-bindings" +version = "0.4.0" +source = "git+https://github.com/drbh/cuda-bindings-prebuilt.git?rev=f129e807d9a8ea32e7572369e91689c00c52570f#f129e807d9a8ea32e7572369e91689c00c52570f" +dependencies = [ + "libloading", +] + +[[package]] +name = "cuda-core" +version = "0.4.0" +source = "git+https://github.com/drbh/cuda-oxide.git?rev=d9d4e379e9a80d1ce7c681d88a086aacec95ba23#d9d4e379e9a80d1ce7c681d88a086aacec95ba23" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.4.0" +source = "git+https://github.com/drbh/cuda-oxide.git?rev=d9d4e379e9a80d1ce7c681d88a086aacec95ba23#d9d4e379e9a80d1ce7c681d88a086aacec95ba23" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "cuda-device" version = "0.2.1" -source = "git+https://github.com/NVIDIA/cuda-rust.git?rev=6d47ae49cb256a07948fb499c20c3eb9b2a96268#6d47ae49cb256a07948fb499c20c3eb9b2a96268" +source = "git+https://github.com/drbh/cuda-oxide.git?rev=d9d4e379e9a80d1ce7c681d88a086aacec95ba23#d9d4e379e9a80d1ce7c681d88a086aacec95ba23" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +source = "git+https://github.com/drbh/cuda-oxide.git?rev=d9d4e379e9a80d1ce7c681d88a086aacec95ba23#d9d4e379e9a80d1ce7c681d88a086aacec95ba23" dependencies = [ + "cuda-artifact-finalizer", + "cuda-core", "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror", ] [[package]] name = "cuda-macros" version = "0.2.1" -source = "git+https://github.com/NVIDIA/cuda-rust.git?rev=6d47ae49cb256a07948fb499c20c3eb9b2a96268#6d47ae49cb256a07948fb499c20c3eb9b2a96268" +source = "git+https://github.com/drbh/cuda-oxide.git?rev=d9d4e379e9a80d1ce7c681d88a086aacec95ba23#d9d4e379e9a80d1ce7c681d88a086aacec95ba23" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -28,24 +136,166 @@ dependencies = [ ] [[package]] -name = "cudarc" -version = "0.19.10" +name = "cuda-target-spec" +version = "0.1.0" +source = "git+https://github.com/drbh/cuda-oxide.git?rev=d9d4e379e9a80d1ce7c681d88a086aacec95ba23#d9d4e379e9a80d1ce7c681d88a086aacec95ba23" + +[[package]] +name = "digest" +version = "0.10.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7359bf1de037ddada72729c3acec9b040f3468400c6e9e740b01ba0cd2e89562" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" dependencies = [ - "libloading", + "block-buffer", + "crypto-common", +] + +[[package]] +name = "equivalent" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00d174d5400e5e8fd687ad1049e2f578285fa914201b1af7e8b112a4546bd826" + +[[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 = "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.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "kernels-cuda" +version = "0.1.0" +dependencies = [ + "cuda-core", + "half", + "tvm-ffi", + "tvm-ffi-sys", ] +[[package]] +name = "libc" +version = "0.2.191" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65cd87c9a6d60536580b316ad985726104730a5151c99845f2490428ef91240a" + [[package]] name = "libloading" -version = "0.9.0" +version = "0.8.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "754ca22de805bb5744484a5b151a9e1a8e837d5dc232c2d7d8c2e3492edc8b60" +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=d9d4e379e9a80d1ce7c681d88a086aacec95ba23#d9d4e379e9a80d1ce7c681d88a086aacec95ba23" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror", +] + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "nvjitlink-sys" +version = "0.2.1" +source = "git+https://github.com/drbh/cuda-oxide.git?rev=d9d4e379e9a80d1ce7c681d88a086aacec95ba23#d9d4e379e9a80d1ce7c681d88a086aacec95ba23" +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 = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74d9a941c0c43d2dceaf25f5c9ca4c9f9cceb28dfa5e4e3cfd14e476f003e7c9" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +source = "git+https://github.com/drbh/cuda-oxide.git?rev=d9d4e379e9a80d1ce7c681d88a086aacec95ba23#d9d4e379e9a80d1ce7c681d88a086aacec95ba23" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "paste" version = "1.0.15" @@ -85,6 +335,11 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "ptx-parse" +version = "0.1.0" +source = "git+https://github.com/drbh/cuda-oxide.git?rev=d9d4e379e9a80d1ce7c681d88a086aacec95ba23#d9d4e379e9a80d1ce7c681d88a086aacec95ba23" + [[package]] name = "quote" version = "1.0.47" @@ -95,25 +350,74 @@ dependencies = [ ] [[package]] -name = "relu-device" +name = "relu-rust-cuda" version = "0.1.0" dependencies = [ + "cuda-core", "cuda-device", + "cuda-host", + "kernels-cuda", + "tvm-ffi", ] [[package]] -name = "relu-rs" -version = "0.1.0" +name = "reserved-oxide-symbols" +version = "0.2.1" +source = "git+https://github.com/drbh/cuda-oxide.git?rev=d9d4e379e9a80d1ce7c681d88a086aacec95ba23#d9d4e379e9a80d1ce7c681d88a086aacec95ba23" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" dependencies = [ - "cudarc", - "tvm-ffi", - "tvm-ffi-sys", + "serde_core", + "serde_derive", ] [[package]] -name = "reserved-oxide-symbols" -version = "0.2.1" -source = "git+https://github.com/NVIDIA/cuda-rust.git?rev=6d47ae49cb256a07948fb499c20c3eb9b2a96268#6d47ae49cb256a07948fb499c20c3eb9b2a96268" +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "serde_json" +version = "1.0.154" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[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 = "syn" @@ -137,6 +441,37 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "syn" +version = "3.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + [[package]] name = "tvm-ffi" version = "0.1.14-post1" @@ -166,6 +501,12 @@ version = "0.1.14-post1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b0b29bc4a713c7e86bcb98e1655832058c272eb336f8d62cece8344cb2666ee2" +[[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.27" @@ -183,3 +524,29 @@ 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.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5fe1f8f1b06191a00962174c61aa5005e0bb391a6d80d07e24d115c01a92ed8" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "863ad3ac83293fb4d740aedbfdc9240dd8d1a50c1099acd76ce80ce7c7230c7f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/examples/kernels/relu-rust-cuda/Cargo.toml b/examples/kernels/relu-rust-cuda/Cargo.toml index 42e2fea8..0ea0463d 100644 --- a/examples/kernels/relu-rust-cuda/Cargo.toml +++ b/examples/kernels/relu-rust-cuda/Cargo.toml @@ -1,3 +1,26 @@ +[package] +name = "relu-rust-cuda" +version = "0.1.0" +edition = "2024" +license = "Apache-2.0" +publish = false + +[lib] +name = "relu_rust_cuda" + [workspace] -members = ["relu-rs", "relu-device"] -resolver = "2" + +# The kernel and its host code live in this crate. Kernel-builder compiles it +# with the cuda-oxide codegen backend and links the kernels as SASS. +[dependencies] +cuda-core = { git = "https://github.com/drbh/cuda-oxide.git", rev = "d9d4e379e9a80d1ce7c681d88a086aacec95ba23" } +cuda-device = { git = "https://github.com/drbh/cuda-oxide.git", rev = "d9d4e379e9a80d1ce7c681d88a086aacec95ba23" } +cuda-host = { git = "https://github.com/drbh/cuda-oxide.git", rev = "d9d4e379e9a80d1ce7c681d88a086aacec95ba23" } +tvm-ffi = "=0.1.14-post1" +# TODO: a git dependency on huggingface/kernels once the crate is pushed; the +# Nix build only sees this directory. +kernels-cuda = { path = "../../../kernels-cuda" } + +# Checked-in bindings: the build needs no CUDA headers or libclang. +[patch."https://github.com/drbh/cuda-oxide.git"] +cuda-bindings = { git = "https://github.com/drbh/cuda-bindings-prebuilt.git", rev = "f129e807d9a8ea32e7572369e91689c00c52570f" } diff --git a/examples/kernels/relu-rust-cuda/build.toml b/examples/kernels/relu-rust-cuda/build.toml index 7af88b44..1017f06b 100644 --- a/examples/kernels/relu-rust-cuda/build.toml +++ b/examples/kernels/relu-rust-cuda/build.toml @@ -12,12 +12,7 @@ repo-id = "kernels-test/relu-rust-cuda" [kernel.relu_rust_cuda] language = "rust-cuda" -cargo-manifest = "relu-rs/Cargo.toml" -cuda-capabilities = ["8.0"] +cargo-manifest = "Cargo.toml" +cuda-capabilities = ["8.0", "9.0", "10.0", "12.0"] depends = [] -src = [ - "Cargo.toml", - "relu-rs/src/lib.rs", - "relu-device/Cargo.toml", - "relu-device/src/main.rs", -] +src = ["Cargo.toml", "Cargo.lock", "src/lib.rs"] diff --git a/examples/kernels/relu-rust-cuda/relu-device/Cargo.toml b/examples/kernels/relu-rust-cuda/relu-device/Cargo.toml deleted file mode 100644 index 649fd8e3..00000000 --- a/examples/kernels/relu-rust-cuda/relu-device/Cargo.toml +++ /dev/null @@ -1,11 +0,0 @@ -[package] -name = "relu-device" -version = "0.1.0" -edition = "2024" -license = "Apache-2.0" -publish = false - -# Kernel-only crate, compiled to PTX by the cuda-oxide codegen backend. Without -# `cuda-host` the kernel macros emit no host glue. -[dependencies] -cuda-device = { git = "https://github.com/NVIDIA/cuda-rust.git", rev = "6d47ae49cb256a07948fb499c20c3eb9b2a96268" } diff --git a/examples/kernels/relu-rust-cuda/relu-device/src/main.rs b/examples/kernels/relu-rust-cuda/relu-device/src/main.rs deleted file mode 100644 index c3c57fea..00000000 --- a/examples/kernels/relu-rust-cuda/relu-device/src/main.rs +++ /dev/null @@ -1,13 +0,0 @@ -use cuda_device::{DisjointSlice, kernel, thread}; - -#[kernel] -pub fn relu(x: &[f32], mut out: DisjointSlice) { - let idx = thread::index_1d(); - let i = idx.get(); - if let Some(slot) = out.get_mut(idx) { - let v = x[i]; - *slot = if v > 0.0 { v } else { 0.0 }; - } -} - -fn main() {} diff --git a/examples/kernels/relu-rust-cuda/relu-rs/Cargo.toml b/examples/kernels/relu-rust-cuda/relu-rs/Cargo.toml deleted file mode 100644 index 22365719..00000000 --- a/examples/kernels/relu-rust-cuda/relu-rs/Cargo.toml +++ /dev/null @@ -1,19 +0,0 @@ -[package] -name = "relu-rs" -version = "0.1.0" -edition = "2021" -license = "Apache-2.0" -publish = false - -# Kernel-builder compiles these with cuda-oxide before this crate and passes -# the PTX directory as `KERNEL_BUILDER_PTX_DIR`. -[[package.metadata.cuda-oxide.device-crates]] -manifest-path = "../relu-device/Cargo.toml" - -[lib] -name = "relu_rust_cuda" - -[dependencies] -cudarc = { version = "0.19.9", default-features = false, features = ["std", "driver", "dynamic-loading", "cuda-12060"] } -tvm-ffi = "=0.1.14-post1" -tvm-ffi-sys = "=0.1.14-post1" diff --git a/examples/kernels/relu-rust-cuda/relu-rs/src/lib.rs b/examples/kernels/relu-rust-cuda/relu-rs/src/lib.rs deleted file mode 100644 index e574c397..00000000 --- a/examples/kernels/relu-rust-cuda/relu-rs/src/lib.rs +++ /dev/null @@ -1,103 +0,0 @@ -use std::collections::HashMap; -use std::ffi::{c_void, CString}; -use std::sync::{Arc, Mutex, OnceLock}; - -use cudarc::driver::{result, sys, CudaContext}; -use tvm_ffi::{Error, Result, Tensor, RUNTIME_ERROR, VALUE_ERROR}; -use tvm_ffi_sys::dlpack::{DLDataTypeCode, DLDeviceType}; -use tvm_ffi_sys::TVMFFIEnvGetStream; - -// PTX of the `relu-device` crate, NUL-terminated for `cuModuleLoadData`. -const PTX: &str = concat!( - include_str!(concat!(env!("KERNEL_BUILDER_PTX_DIR"), "/relu_device.ptx")), - "\0" -); - -struct Relu { - ctx: Arc, - function: sys::CUfunction, -} - -unsafe impl Send for Relu {} - -fn cuda_error(e: result::DriverError) -> Error { - Error::new(RUNTIME_ERROR, &format!("{e:?}"), "") -} - -fn value_error(message: &str) -> Error { - Error::new(VALUE_ERROR, message, "") -} - -// Loads the module once per device into the primary context Torch also uses. -fn with_relu(device: i32, f: impl FnOnce(&Relu) -> Result) -> Result { - static RELU: OnceLock>> = OnceLock::new(); - let mut loaded = RELU.get_or_init(Default::default).lock().unwrap(); - if !loaded.contains_key(&device) { - let ctx = CudaContext::new(device as usize).map_err(cuda_error)?; - ctx.bind_to_thread().map_err(cuda_error)?; - let function = unsafe { - let module = result::module::load_data(PTX.as_ptr().cast()).map_err(cuda_error)?; - result::module::get_function(module, CString::new("relu").unwrap()) - .map_err(cuda_error)? - }; - loaded.insert(device, Relu { ctx, function }); - } - f(&loaded[&device]) -} - -fn relu_rust_cuda(x: Tensor, out: Tensor) -> Result<()> { - let f32_type = |t: &Tensor| { - let dtype = t.dtype(); - dtype.code == DLDataTypeCode::kDLFloat as u8 && dtype.bits == 32 && dtype.lanes == 1 - }; - if !f32_type(&x) || !f32_type(&out) { - return Err(value_error("input and output must be float32")); - } - if !x.is_contiguous() || !out.is_contiguous() { - return Err(value_error("input and output must be contiguous")); - } - if x.numel() != out.numel() { - return Err(value_error(&format!( - "input and output must have the same number of elements, got {} and {}", - x.numel(), - out.numel() - ))); - } - let device = x.device(); - if device.device_type != DLDeviceType::kDLCUDA || out.device() != device { - return Err(value_error( - "input and output must be on the same CUDA device", - )); - } - if x.numel() == 0 { - return Ok(()); - } - - with_relu(device.device_id, |relu| { - relu.ctx.bind_to_thread().map_err(cuda_error)?; - let stream = unsafe { TVMFFIEnvGetStream(DLDeviceType::kDLCUDA as i32, device.device_id) }; - // Slices are passed to cuda-oxide kernels as (pointer, length) pairs. - let (mut x_ptr, mut out_ptr) = (x.data_ptr() as u64, out.data_ptr() as u64); - let mut len = x.numel() as u64; - let mut params = [ - (&mut x_ptr as *mut u64).cast::(), - (&mut len as *mut u64).cast(), - (&mut out_ptr as *mut u64).cast(), - (&mut len as *mut u64).cast(), - ]; - let blocks = len.div_ceil(256) as u32; - unsafe { - result::launch_kernel( - relu.function, - (blocks, 1, 1), - (256, 1, 1), - 0, - stream.cast(), - &mut params, - ) - } - .map_err(cuda_error) - }) -} - -tvm_ffi::tvm_ffi_dll_export_typed_func!(relu_rust_cuda, relu_rust_cuda); diff --git a/examples/kernels/relu-rust-cuda/rust-git-hashes.json b/examples/kernels/relu-rust-cuda/rust-git-hashes.json index 7e8a2a46..7f4a3d2e 100644 --- a/examples/kernels/relu-rust-cuda/rust-git-hashes.json +++ b/examples/kernels/relu-rust-cuda/rust-git-hashes.json @@ -1,5 +1,15 @@ { - "cuda-device-0.2.1": "sha256-KxIPPzDBKeHyknrpRg1ekNliEcc44yBWUtSmrGnApOo=", - "cuda-macros-0.2.1": "sha256-KxIPPzDBKeHyknrpRg1ekNliEcc44yBWUtSmrGnApOo=", - "reserved-oxide-symbols-0.2.1": "sha256-KxIPPzDBKeHyknrpRg1ekNliEcc44yBWUtSmrGnApOo=" + "cuda-artifact-finalizer-0.2.1": "sha256-cz53C3hB3dxVrVLgX+KtZ4WGT+jjoazNsBSedODCFrg=", + "cuda-bindings-0.4.0": "sha256-ZfLWDaBRaZVrPFYUui9GhbL0z38DcJqfECDm/rvpJq8=", + "cuda-core-0.4.0": "sha256-cz53C3hB3dxVrVLgX+KtZ4WGT+jjoazNsBSedODCFrg=", + "cuda-core-derive-0.4.0": "sha256-cz53C3hB3dxVrVLgX+KtZ4WGT+jjoazNsBSedODCFrg=", + "cuda-device-0.2.1": "sha256-cz53C3hB3dxVrVLgX+KtZ4WGT+jjoazNsBSedODCFrg=", + "cuda-host-0.2.1": "sha256-cz53C3hB3dxVrVLgX+KtZ4WGT+jjoazNsBSedODCFrg=", + "cuda-macros-0.2.1": "sha256-cz53C3hB3dxVrVLgX+KtZ4WGT+jjoazNsBSedODCFrg=", + "cuda-target-spec-0.1.0": "sha256-cz53C3hB3dxVrVLgX+KtZ4WGT+jjoazNsBSedODCFrg=", + "libnvvm-sys-0.2.1": "sha256-cz53C3hB3dxVrVLgX+KtZ4WGT+jjoazNsBSedODCFrg=", + "nvjitlink-sys-0.2.1": "sha256-cz53C3hB3dxVrVLgX+KtZ4WGT+jjoazNsBSedODCFrg=", + "oxide-tile-export-0.2.1": "sha256-cz53C3hB3dxVrVLgX+KtZ4WGT+jjoazNsBSedODCFrg=", + "ptx-parse-0.1.0": "sha256-cz53C3hB3dxVrVLgX+KtZ4WGT+jjoazNsBSedODCFrg=", + "reserved-oxide-symbols-0.2.1": "sha256-cz53C3hB3dxVrVLgX+KtZ4WGT+jjoazNsBSedODCFrg=" } diff --git a/examples/kernels/relu-rust-cuda/src/lib.rs b/examples/kernels/relu-rust-cuda/src/lib.rs new file mode 100644 index 00000000..a79b2ca8 --- /dev/null +++ b/examples/kernels/relu-rust-cuda/src/lib.rs @@ -0,0 +1,44 @@ +use cuda_core::simt::LaunchConfig; +use cuda_device::cuda_module; +use kernels_cuda::cuda_error; +use tvm_ffi::{Error, Result, Tensor, VALUE_ERROR}; + +#[cuda_module] +mod kernels { + use cuda_device::{DisjointSlice, kernel, thread}; + + #[kernel] + pub fn relu(x: &[f32], mut out: DisjointSlice) { + let idx = thread::index_1d(); + let i = idx.get(); + if let Some(slot) = out.get_mut(idx) { + let v = x[i]; + *slot = if v > 0.0 { v } else { 0.0 }; + } + } +} + +fn relu_rust_cuda(x: Tensor, out: Tensor) -> Result<()> { + if x.numel() != out.numel() { + return Err(Error::new( + VALUE_ERROR, + &format!( + "input and output must have the same number of elements, got {} and {}", + x.numel(), + out.numel() + ), + "", + )); + } + if x.numel() == 0 { + return Ok(()); + } + + let gpu = kernels_cuda::load(x.device(), kernels::from_module)?; + let x = gpu.view::(&x)?; + let mut out = gpu.view::(&out)?; + let config = LaunchConfig::for_num_elems(x.len() as u32); + unsafe { gpu.module.relu(&gpu.stream, config, &x, &mut out) }.map_err(cuda_error) +} + +tvm_ffi::tvm_ffi_dll_export_typed_func!(relu_rust_cuda, relu_rust_cuda); diff --git a/kernel-builder/src/pyproject/templates/rust-fatbin.cpp.in b/kernel-builder/src/pyproject/templates/rust-fatbin.cpp.in new file mode 100644 index 00000000..bd004c93 --- /dev/null +++ b/kernel-builder/src/pyproject/templates/rust-fatbin.cpp.in @@ -0,0 +1,16 @@ +// Links the fatbin of the Rust kernel component @NAME@ into the extension. +// +// The image goes in nvcc's `.nv_fatbin` section, so cuobjdump finds the SASS. +// Unlike nvcc's fatbins it is not registered with the CUDA runtime: the crate +// loads it itself with `cuModuleLoadData`. +__asm__( + ".section .nv_fatbin, \"a\"\n" + ".balign 16\n" + ".globl KERNEL_BUILDER_CUDA_FATBIN\n" + ".hidden KERNEL_BUILDER_CUDA_FATBIN\n" + "KERNEL_BUILDER_CUDA_FATBIN:\n" + ".incbin \"@FATBIN@\"\n" + ".globl KERNEL_BUILDER_CUDA_FATBIN_END\n" + ".hidden KERNEL_BUILDER_CUDA_FATBIN_END\n" + "KERNEL_BUILDER_CUDA_FATBIN_END:\n" + ".previous\n"); diff --git a/kernel-builder/src/pyproject/templates/rust.cmake b/kernel-builder/src/pyproject/templates/rust.cmake index 5feb65ca..3cf68f7b 100644 --- a/kernel-builder/src/pyproject/templates/rust.cmake +++ b/kernel-builder/src/pyproject/templates/rust.cmake @@ -1,44 +1,6 @@ -# cuda-oxide interop: device crates are declared by the host crate in -# `[[package.metadata.cuda-oxide.device-crates]]`, as `cargo oxide` reads them. -function(_cuda_oxide_device_manifests OUT_VAR MANIFEST) - execute_process( - COMMAND ${CARGO_EXECUTABLE} metadata --format-version 1 --no-deps --offline - --manifest-path ${MANIFEST} - OUTPUT_VARIABLE _JSON - ERROR_VARIABLE _ERROR - RESULT_VARIABLE _RESULT) - if(NOT _RESULT EQUAL 0) - message(FATAL_ERROR "cargo metadata failed for ${MANIFEST}: ${_ERROR}") - endif() - - file(REAL_PATH ${MANIFEST} _MANIFEST) - set(_MANIFESTS) - string(JSON _N_PACKAGES LENGTH "${_JSON}" packages) - math(EXPR _LAST_PACKAGE "${_N_PACKAGES} - 1") - foreach(_I RANGE ${_LAST_PACKAGE}) - string(JSON _PACKAGE_MANIFEST GET "${_JSON}" packages ${_I} manifest_path) - file(REAL_PATH ${_PACKAGE_MANIFEST} _PACKAGE_MANIFEST) - if(NOT _PACKAGE_MANIFEST STREQUAL _MANIFEST) - continue() - endif() - - string(JSON _CRATES ERROR_VARIABLE _NO_CRATES - GET "${_JSON}" packages ${_I} metadata cuda-oxide device-crates) - if(_NO_CRATES) - break() - endif() - get_filename_component(_PACKAGE_DIR ${_PACKAGE_MANIFEST} DIRECTORY) - string(JSON _N_CRATES LENGTH "${_CRATES}") - math(EXPR _LAST_CRATE "${_N_CRATES} - 1") - foreach(_J RANGE ${_LAST_CRATE}) - string(JSON _CRATE_MANIFEST GET "${_CRATES}" ${_J} manifest-path) - get_filename_component(_CRATE_MANIFEST ${_CRATE_MANIFEST} ABSOLUTE BASE_DIR ${_PACKAGE_DIR}) - list(APPEND _MANIFESTS ${_CRATE_MANIFEST}) - endforeach() - endforeach() - - set(${OUT_VAR} ${_MANIFESTS} PARENT_SCOPE) -endfunction() +# Functions run with the caller's CMAKE_CURRENT_LIST_DIR, so remember where the +# templates beside this file are. +set(_RUST_CMAKE_DIR ${CMAKE_CURRENT_LIST_DIR}) function(rust_kernel_component LIBS_VAR TARGETS_VAR) cmake_parse_arguments(KERNEL "" "NAME;MANIFEST_PATH" "CUDA_CAPABILITIES" ${ARGN}) @@ -66,58 +28,57 @@ function(rust_kernel_component LIBS_VAR TARGETS_VAR) accumulate_gpu_archs(_ALL_GPU_ARCHS "${ALL_GPU_ARCHS}" "${_KERNEL_ARCHS}") set(ALL_GPU_ARCHS ${_ALL_GPU_ARCHS} PARENT_SCOPE) - _cuda_oxide_device_manifests(_DEVICE_MANIFESTS - ${CMAKE_CURRENT_SOURCE_DIR}/${KERNEL_MANIFEST_PATH}) - endif() - - set(_HOST_ENV) - - if(_DEVICE_MANIFESTS) - # Build cuda-oxide device code before the host crate embeds its PTX. + # The crate is built with the cuda-oxide codegen backend, which compiles + # host code as usual and writes the device code as PTX. if(NOT CUDA_OXIDE_BACKEND) - message(FATAL_ERROR "rust_kernel_component: cuda-oxide device crates require CUDA_OXIDE_BACKEND") + message(FATAL_ERROR "rust_kernel_component: Rust CUDA kernels require CUDA_OXIDE_BACKEND") endif() - set(_DEVICE_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/${KERNEL_NAME}-device) - set(_PTX_DIR ${_DEVICE_TARGET_DIR}/ptx) - 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}") + # The fatbin is linked in under fixed symbol names. + get_property(_FATBIN_OWNER GLOBAL PROPERTY RUST_CUDA_FATBIN_OWNER) + if(_FATBIN_OWNER) + message(FATAL_ERROR "rust_kernel_component: ${KERNEL_NAME} and ${_FATBIN_OWNER} " + "are both Rust CUDA components; only one is supported per extension") endif() + set_property(GLOBAL PROPERTY RUST_CUDA_FATBIN_OWNER ${KERNEL_NAME}) - set(_DEVICE_COMMANDS) - foreach(_DEVICE_MANIFEST ${_DEVICE_MANIFESTS}) - # Flags match what `cargo oxide build` passes to the codegen backend. - list(APPEND _DEVICE_COMMANDS - COMMAND ${CMAKE_COMMAND} -E env - "CUDA_OXIDE_PTX_DIR=${_PTX_DIR}" - "RUSTFLAGS=-Zcodegen-backend=${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}) - endforeach() - - add_custom_target(${KERNEL_NAME}_oxide_device_build ALL - COMMAND ${CMAKE_COMMAND} -E make_directory ${_PTX_DIR} - ${_DEVICE_COMMANDS} - WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} - COMMENT "Building Rust CUDA device kernel ${KERNEL_NAME}" - VERBATIM - ) - list(APPEND _HOST_ENV "KERNEL_BUILDER_PTX_DIR=${_PTX_DIR}") + get_filename_component(_CUDA_BIN_DIR "${CMAKE_CUDA_COMPILER}" DIRECTORY) + find_program(CUDA_PTXAS_EXECUTABLE ptxas HINTS ${_CUDA_BIN_DIR} REQUIRED) + find_program(CUDA_FATBINARY_EXECUTABLE fatbinary HINTS ${_CUDA_BIN_DIR} REQUIRED) + if(NOT CMAKE_OBJCOPY) + message(FATAL_ERROR "rust_kernel_component: Rust CUDA kernels require objcopy") + endif() + + list(SORT _KERNEL_ARCHS COMPARE NATURAL) + # ptxas compiles PTX for newer archs too, so target the lowest one. + list(GET _KERNEL_ARCHS 0 _OXIDE_ARCH) + string(REPLACE "+PTX" "" _OXIDE_ARCH "${_OXIDE_ARCH}") + string(REPLACE "." "" _OXIDE_ARCH "${_OXIDE_ARCH}") + # Flags match what `cargo oxide build` passes to the codegen backend. + list(APPEND _CARGO_ENV + "CUDA_OXIDE_TARGET=sm_${_OXIDE_ARCH}" + "CUDA_OXIDE_PTX_DIR=${CMAKE_BINARY_DIR}/cargo/${KERNEL_NAME}/ptx" + "RUSTFLAGS=-Zcodegen-backend=${CUDA_OXIDE_BACKEND} -Copt-level=3 -Cdebug-assertions=off -Zmir-enable-passes=-JumpThreading -Csymbol-mangling-version=v0") endif() string(REPLACE "-" "_" _LIB_NAME ${KERNEL_NAME}) set(_CARGO_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/${KERNEL_NAME}) set(_STATICLIB ${_CARGO_TARGET_DIR}/release/${CMAKE_STATIC_LIBRARY_PREFIX}${_LIB_NAME}${CMAKE_STATIC_LIBRARY_SUFFIX}) + set(_LINKED_STATICLIB ${_STATICLIB}) + set(_BYPRODUCTS ${_STATICLIB}) + + if(GPU_LANG STREQUAL "CUDA") + set(_PTX ${_CARGO_TARGET_DIR}/ptx/${_LIB_NAME}.ptx) + list(APPEND _BYPRODUCTS ${_PTX}) + # The backend also embeds the PTX in the crate's `.oxart` sections. The + # extension loads the fatbin instead, so drop the PTX from what we link. + set(_LINKED_STATICLIB ${_CARGO_TARGET_DIR}/${CMAKE_STATIC_LIBRARY_PREFIX}${_LIB_NAME}_sass${CMAKE_STATIC_LIBRARY_SUFFIX}) + list(APPEND _BYPRODUCTS ${_LINKED_STATICLIB}) + # Keeping the date lets a no-op cargo build skip the relink. + set(_STRIP_PTX COMMAND ${CMAKE_OBJCOPY} --preserve-dates --remove-section=.oxart + ${_STATICLIB} ${_LINKED_STATICLIB}) + endif() # tvm-ffi-sys's build script shells out to `tvm-ffi-config`, a console script # the apache-tvm-ffi wheel installs beside the interpreter. @@ -125,26 +86,69 @@ function(rust_kernel_component LIBS_VAR TARGETS_VAR) add_custom_target(${KERNEL_NAME}_cargo_build ALL COMMAND ${CMAKE_COMMAND} -E env "PATH=${_PYTHON_BIN_DIR}:$ENV{PATH}" - ${_HOST_ENV} + ${_CARGO_ENV} ${CARGO_EXECUTABLE} rustc --release --locked --lib --crate-type staticlib --manifest-path ${CMAKE_CURRENT_SOURCE_DIR}/${KERNEL_MANIFEST_PATH} --target-dir ${_CARGO_TARGET_DIR} - BYPRODUCTS ${_STATICLIB} + ${_STRIP_PTX} + BYPRODUCTS ${_BYPRODUCTS} WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} COMMENT "Building Rust kernel ${KERNEL_NAME} with cargo" VERBATIM ) - if(_DEVICE_MANIFESTS) - add_dependencies(${KERNEL_NAME}_cargo_build ${KERNEL_NAME}_oxide_device_build) - endif() add_library(${KERNEL_NAME}_rust STATIC IMPORTED GLOBAL) - set_target_properties(${KERNEL_NAME}_rust PROPERTIES IMPORTED_LOCATION ${_STATICLIB}) + set_target_properties(${KERNEL_NAME}_rust PROPERTIES IMPORTED_LOCATION ${_LINKED_STATICLIB}) + + if(GPU_LANG STREQUAL "CUDA") + _rust_cuda_fatbin(${KERNEL_NAME} ${_PTX} ${_OXIDE_ARCH} "${_KERNEL_ARCHS}") + endif() set(${LIBS_VAR} ${${LIBS_VAR}} ${KERNEL_NAME}_rust PARENT_SCOPE) set(${TARGETS_VAR} ${${TARGETS_VAR}} ${KERNEL_NAME}_cargo_build PARENT_SCOPE) endfunction() +# Compiles the crate's PTX to SASS for every capability, so nothing is +# JIT-compiled on the user's machine. `target_link_rust_kernels` links the +# fatbin between `KERNEL_BUILDER_CUDA_FATBIN` and `KERNEL_BUILDER_CUDA_FATBIN_END`. +function(_rust_cuda_fatbin NAME PTX PTX_ARCH ARCHS) + get_filename_component(_DIR ${PTX} DIRECTORY) + set(_CUBINS) + set(_IMAGES) + foreach(_ARCH IN LISTS ARCHS) + string(REPLACE "+PTX" "" _SM "${_ARCH}") + string(REPLACE "." "" _SM "${_SM}") + set(_CUBIN ${_DIR}/${NAME}.sm_${_SM}.cubin) + add_custom_command(OUTPUT ${_CUBIN} + COMMAND ${CUDA_PTXAS_EXECUTABLE} -arch=sm_${_SM} ${PTX} -o ${_CUBIN} + DEPENDS ${PTX} + COMMENT "Compiling Rust kernel ${NAME} for sm_${_SM}" + VERBATIM) + list(APPEND _CUBINS ${_CUBIN}) + list(APPEND _IMAGES "--image3=kind=elf,sm=${_SM},file=${_CUBIN}") + if(_ARCH MATCHES "\\+PTX$") + set(_KEEP_PTX ON) + endif() + endforeach() + if(_KEEP_PTX) + # The PTX targets the lowest arch, so it JITs on any newer GPU. + list(APPEND _IMAGES "--image3=kind=ptx,sm=${PTX_ARCH},file=${PTX}") + endif() + + set(_FATBIN ${_DIR}/${NAME}.fatbin) + add_custom_command(OUTPUT ${_FATBIN} + COMMAND ${CUDA_FATBINARY_EXECUTABLE} --create=${_FATBIN} ${_IMAGES} + DEPENDS ${_CUBINS} ${PTX} + COMMENT "Creating fatbin for Rust kernel ${NAME}" + VERBATIM) + + set(_SOURCE ${CMAKE_CURRENT_BINARY_DIR}/${NAME}_fatbin.cpp) + set(FATBIN ${_FATBIN}) + configure_file(${_RUST_CMAKE_DIR}/rust-fatbin.cpp.in ${_SOURCE} @ONLY) + set_source_files_properties(${_SOURCE} PROPERTIES OBJECT_DEPENDS ${_FATBIN}) + set_property(GLOBAL PROPERTY RUST_CUDA_FATBIN_SOURCE ${_SOURCE}) +endfunction() + # `add_library(SHARED)` errors on an empty source list, and a Rust-only # extension has none: the crate exports the tvm-ffi entry points itself. function(rust_extension_sources SRC_VAR) @@ -191,4 +195,16 @@ function(target_link_rust_kernels TARGET) Threads::Threads ${CMAKE_DL_LIBS}) _restrict_rust_exports(${TARGET}) + if(UNIX AND NOT APPLE) + # The codegen backend's embedded PTX objects carry no `.note.GNU-stack`, + # which would otherwise make the linker mark the stack executable. + target_link_options(${TARGET} PRIVATE "LINKER:-z,noexecstack") + get_property(_FATBIN_SOURCE GLOBAL PROPERTY RUST_CUDA_FATBIN_SOURCE) + if(_FATBIN_SOURCE) + target_sources(${TARGET} PRIVATE ${_FATBIN_SOURCE}) + # A crate that still calls the generated `load` references the + # embedded PTX we removed. Fail the link instead of the import. + target_link_options(${TARGET} PRIVATE "LINKER:--no-undefined") + endif() + endif() endfunction() diff --git a/kernel-builder/src/pyproject/tvm_ffi/mod.rs b/kernel-builder/src/pyproject/tvm_ffi/mod.rs index 32cd0301..dab8cbc9 100644 --- a/kernel-builder/src/pyproject/tvm_ffi/mod.rs +++ b/kernel-builder/src/pyproject/tvm_ffi/mod.rs @@ -17,6 +17,7 @@ use crate::pyproject::FileSet; static BUILD_VARIANTS_UTILS: &str = include_str!("../templates/tvm_ffi/build-variants.cmake"); static CMAKE_KERNEL: &str = include_str!("../templates/kernel.cmake"); static CMAKE_RUST: &str = include_str!("../templates/rust.cmake"); +static CMAKE_RUST_FATBIN: &str = include_str!("../templates/rust-fatbin.cpp.in"); static CMAKE_UTILS: &str = include_str!("../templates/utils.cmake"); static HIPIFY: &str = include_str!("../templates/tvm_ffi/hipify.py"); static OPS_PY_IN: &str = include_str!("../templates/tvm_ffi/_ops.py.in"); @@ -27,6 +28,7 @@ fn write_cmake_helpers(file_set: &mut FileSet) { write_cmake_file(file_set, "utils.cmake", CMAKE_UTILS.as_bytes()); write_cmake_file(file_set, "kernel.cmake", CMAKE_KERNEL.as_bytes()); write_cmake_file(file_set, "rust.cmake", CMAKE_RUST.as_bytes()); + write_cmake_file(file_set, "rust-fatbin.cpp.in", CMAKE_RUST_FATBIN.as_bytes()); write_cmake_file( file_set, "build-variants.cmake", diff --git a/kernel-builder/tests/rust-components.cmake b/kernel-builder/tests/rust-components.cmake index 6b435009..9b64b72e 100644 --- a/kernel-builder/tests/rust-components.cmake +++ b/kernel-builder/tests/rust-components.cmake @@ -1,18 +1,15 @@ -# Configure Rust targets without CUDA or compiling crates (needs `cargo` for -# `cargo metadata`). Run: cmake -P kernel-builder/tests/rust-components.cmake +# Configure Rust targets without CUDA or compiling crates. +# Run: cmake -P kernel-builder/tests/rust-components.cmake cmake_minimum_required(VERSION 3.24) if(CMAKE_SCRIPT_MODE_FILE) set(_TEST_DIR "${CMAKE_CURRENT_BINARY_DIR}/rust-components-test") file(MAKE_DIRECTORY "${_TEST_DIR}") file(WRITE "${_TEST_DIR}/CMakeLists.txt" "cmake_minimum_required(VERSION 3.24)\nproject(rust_components NONE)\ninclude(\"${CMAKE_CURRENT_LIST_FILE}\")\n") - # Host crates with and without cuda-oxide device crates. - set(_PACKAGE "[package]\nname = \"example\"\nversion = \"0.0.0\"\nedition = \"2021\"\n[lib]\npath = \"lib.rs\"\n[workspace]\n") - file(WRITE "${_TEST_DIR}/plain/Cargo.toml" "${_PACKAGE}") - file(WRITE "${_TEST_DIR}/oxide/Cargo.toml" "${_PACKAGE}" - "[[package.metadata.cuda-oxide.device-crates]]\nmanifest-path = \"device/Cargo.toml\"\n") - file(TOUCH "${_TEST_DIR}/plain/lib.rs" "${_TEST_DIR}/oxide/lib.rs") - foreach(_CASE cpu cuda host-only unsupported-arch missing-backend) + file(WRITE "${_TEST_DIR}/crate/Cargo.toml" + "[package]\nname = \"example\"\nversion = \"0.0.0\"\nedition = \"2021\"\n[lib]\npath = \"lib.rs\"\n[workspace]\n") + file(TOUCH "${_TEST_DIR}/crate/lib.rs") + foreach(_CASE cpu cuda unsupported-arch missing-backend two-components) execute_process(COMMAND "${CMAKE_COMMAND}" -S "${_TEST_DIR}" -G Ninja -B "${_TEST_DIR}/${_CASE}" "-DCASE=${_CASE}" RESULT_VARIABLE _RESULT OUTPUT_VARIABLE _OUTPUT ERROR_VARIABLE _ERROR) @@ -24,6 +21,10 @@ if(CMAKE_SCRIPT_MODE_FILE) if(_RESULT EQUAL 0 OR NOT _ERROR MATCHES "require CUDA_OXIDE_BACKEND") message(FATAL_ERROR "Expected missing backend error: ${_OUTPUT}${_ERROR}") endif() + elseif(_CASE STREQUAL "two-components") + if(_RESULT EQUAL 0 OR NOT _ERROR MATCHES "only one is supported") + message(FATAL_ERROR "Expected two components error: ${_OUTPUT}${_ERROR}") + endif() elseif(NOT _RESULT EQUAL 0) message(FATAL_ERROR "${_CASE} configuration failed: ${_OUTPUT}${_ERROR}") endif() @@ -42,27 +43,34 @@ set(GPU_LANG CUDA) set(CUDA_ARCHS "8.0;8.9;9.0") set(CUDA_KERNEL_ARCHS "8.0;8.9;9.0+PTX") set(CUDA_OXIDE_BACKEND /test/backend.so) -set(_MANIFEST oxide/Cargo.toml) +set(CUDA_PTXAS_EXECUTABLE /test/ptxas) +set(CUDA_FATBINARY_EXECUTABLE /test/fatbinary) +set(CMAKE_OBJCOPY /test/objcopy) set(_ARGS CUDA_CAPABILITIES 8.0) if(CASE STREQUAL "cpu") set(GPU_LANG CPU) set(_ARGS) -elseif(CASE STREQUAL "host-only") - set(_MANIFEST plain/Cargo.toml) - set(_ARGS) elseif(CASE STREQUAL "unsupported-arch") set(_ARGS CUDA_CAPABILITIES 6.0) elseif(CASE STREQUAL "missing-backend") unset(CUDA_OXIDE_BACKEND) endif() -rust_kernel_component(LIBS TARGETS NAME example MANIFEST_PATH ${_MANIFEST} ${_ARGS}) +rust_kernel_component(LIBS TARGETS NAME example MANIFEST_PATH crate/Cargo.toml ${_ARGS}) if(NOT TARGET example_rust OR NOT TARGET example_cargo_build) message(FATAL_ERROR "Missing shared Cargo/static-library targets") endif() if(CASE STREQUAL "cuda") - if(NOT TARGET example_oxide_device_build OR NOT ALL_GPU_ARCHS STREQUAL "8.0") - message(FATAL_ERROR "Missing device target or incorrect capabilities") + if(NOT ALL_GPU_ARCHS STREQUAL "8.0") + message(FATAL_ERROR "Incorrect capabilities: ${ALL_GPU_ARCHS}") + endif() + get_property(_SOURCE GLOBAL PROPERTY RUST_CUDA_FATBIN_SOURCE) + if(NOT _SOURCE MATCHES "example_fatbin.cpp$") + message(FATAL_ERROR "Missing fatbin source: ${_SOURCE}") + endif() + file(READ ${_SOURCE} _CONTENTS) + if(NOT _CONTENTS MATCHES "\\.incbin [^\n]*/example\\.fatbin") + message(FATAL_ERROR "Fatbin source does not include the fatbin:\n${_CONTENTS}") endif() -elseif(TARGET example_oxide_device_build) - message(FATAL_ERROR "Unexpected device build for ${CASE}") +elseif(CASE STREQUAL "two-components") + rust_kernel_component(LIBS TARGETS NAME other MANIFEST_PATH crate/Cargo.toml ${_ARGS}) endif() diff --git a/kernels-cuda/Cargo.toml b/kernels-cuda/Cargo.toml new file mode 100644 index 00000000..d7826e3c --- /dev/null +++ b/kernels-cuda/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "kernels-cuda" +version = "0.1.0" +edition = "2024" +license = "Apache-2.0" +description = "Runtime support for kernel-builder's rust-cuda kernels" +publish = false + +[dependencies] +cuda-core = { git = "https://github.com/drbh/cuda-oxide.git", rev = "d9d4e379e9a80d1ce7c681d88a086aacec95ba23" } +half = "2" +tvm-ffi = "=0.1.14-post1" +tvm-ffi-sys = "=0.1.14-post1" + +# Kernel crates depend on cuda-oxide, so keep it out of the root workspace. +[workspace] + +# Checked-in bindings: the build needs no CUDA headers or libclang. +[patch."https://github.com/drbh/cuda-oxide.git"] +cuda-bindings = { git = "https://github.com/drbh/cuda-bindings-prebuilt.git", rev = "f129e807d9a8ea32e7572369e91689c00c52570f" } diff --git a/kernels-cuda/src/lib.rs b/kernels-cuda/src/lib.rs new file mode 100644 index 00000000..25adabb9 --- /dev/null +++ b/kernels-cuda/src/lib.rs @@ -0,0 +1,188 @@ +//! Runtime support for kernel-builder's `rust-cuda` kernels. +//! +//! kernel-builder compiles a kernel crate's PTX for every CUDA capability in +//! `build.toml` and links the fatbin into the extension. [`load`] loads that +//! fatbin into the device's primary context, which Torch also uses, binds it +//! to a `#[cuda_module]` and borrows the caller's current stream. +//! [`Gpu::view`] wraps Torch tensors as `DeviceBuffer`s without taking +//! ownership. +//! +//! ```ignore +//! fn relu(x: Tensor, out: Tensor) -> tvm_ffi::Result<()> { +//! let gpu = kernels_cuda::load(x.device(), kernels::from_module)?; +//! let (x, mut out) = (gpu.view::(&x)?, gpu.view::(&out)?); +//! let config = LaunchConfig::for_num_elems(x.len() as u32); +//! unsafe { gpu.module.relu(&gpu.stream, config, &x, &mut out) } +//! .map_err(kernels_cuda::cuda_error) +//! } +//! ``` + +use std::any::{Any, TypeId}; +use std::collections::BTreeMap; +use std::collections::btree_map::Entry; +use std::mem::ManuallyDrop; +use std::ops::{Deref, DerefMut}; +use std::sync::{Arc, Mutex}; + +use cuda_core::{CudaContext, CudaModule, CudaStream, DeviceBuffer, DriverError}; +use tvm_ffi::{Error, RUNTIME_ERROR, Result, Tensor, VALUE_ERROR}; +use tvm_ffi_sys::TVMFFIEnvGetStream; +use tvm_ffi_sys::dlpack::{DLDataTypeCode, DLDevice, DLDeviceType}; + +// kernel-builder links the fatbin between these symbols. +unsafe extern "C" { + static KERNEL_BUILDER_CUDA_FATBIN: u8; + static KERNEL_BUILDER_CUDA_FATBIN_END: u8; +} + +fn fatbin() -> &'static [u8] { + let start = &raw const KERNEL_BUILDER_CUDA_FATBIN; + let end = &raw const KERNEL_BUILDER_CUDA_FATBIN_END; + unsafe { std::slice::from_raw_parts(start, end.offset_from(start) as usize) } +} + +/// Converts a CUDA error, such as a failed launch, to a tvm-ffi error. +pub fn cuda_error(e: impl std::fmt::Debug) -> Error { + Error::new(RUNTIME_ERROR, &format!("{e:?}"), "") +} + +fn value_error(message: &str) -> Error { + Error::new(VALUE_ERROR, message, "") +} + +/// A loaded module and the caller's current stream on one device. +/// +/// The stream is borrowed for the current call, so don't keep a `Gpu` after +/// returning to the caller. +pub struct Gpu { + pub module: Arc, + pub stream: Arc, + ctx: Arc, + device_id: i32, +} + +type Loaded = (Arc, Arc); + +/// Loads the kernels for `device` once and binds them with `from_module`, +/// the function `#[cuda_module]` generates. +pub fn load( + device: DLDevice, + from_module: fn(Arc) -> std::result::Result, +) -> Result> { + if device.device_type != DLDeviceType::kDLCUDA { + return Err(value_error("expected a CUDA device")); + } + static LOADED: Mutex> = Mutex::new(BTreeMap::new()); + let (ctx, module) = match LOADED + .lock() + .unwrap() + .entry((device.device_id, TypeId::of::())) + { + Entry::Occupied(entry) => entry.get().clone(), + Entry::Vacant(entry) => { + let ctx = CudaContext::new(device.device_id as usize).map_err(cuda_error)?; + let module = ctx + .load_module_from_image(fatbin()) + .and_then(from_module) + .map_err(cuda_error)?; + entry.insert((ctx, Arc::new(module))).clone() + } + }; + let module = module.downcast::().expect("modules are keyed by type"); + + ctx.bind_to_thread().map_err(cuda_error)?; + // Torch's current stream, which Torch keeps owning. + let stream = unsafe { + let raw = TVMFFIEnvGetStream(DLDeviceType::kDLCUDA as i32, device.device_id); + CudaStream::borrow_raw(raw, &ctx) + }; + Ok(Gpu { + module, + stream, + ctx, + device_id: device.device_id, + }) +} + +impl Gpu { + /// Views a contiguous `T` tensor on this device as a `DeviceBuffer`. + pub fn view(&self, tensor: &Tensor) -> Result> { + let device = tensor.device(); + if device.device_type != DLDeviceType::kDLCUDA || device.device_id != self.device_id { + return Err(value_error(&format!( + "tensors must be on cuda:{}", + self.device_id + ))); + } + let dtype = tensor.dtype(); + if dtype.code != T::CODE as u8 || dtype.bits != T::BITS || dtype.lanes != 1 { + return Err(value_error(&format!( + "expected a {} tensor", + std::any::type_name::() + ))); + } + if !tensor.is_contiguous() { + return Err(value_error("tensors must be contiguous")); + } + let ptr = tensor.data_ptr() as u64; + // SAFETY: the tensor owns `numel` elements of `T` at `ptr` on this + // device. `View` never frees them. + let buffer = unsafe { DeviceBuffer::from_raw_parts(ptr, tensor.numel(), self.ctx.clone()) }; + Ok(View(ManuallyDrop::new(buffer))) + } +} + +/// A tensor's memory as a `DeviceBuffer`. Dropping it leaves the memory to +/// the tensor. +pub struct View(ManuallyDrop>); + +impl Deref for View { + type Target = DeviceBuffer; + + fn deref(&self) -> &DeviceBuffer { + &self.0 + } +} + +impl DerefMut for View { + fn deref_mut(&mut self) -> &mut DeviceBuffer { + &mut self.0 + } +} + +impl Drop for View { + fn drop(&mut self) { + let buffer = unsafe { ManuallyDrop::take(&mut self.0) }; + let _ = buffer.into_raw_parts(); + } +} + +/// An element type with a DLPack equivalent. +pub trait Element { + const CODE: DLDataTypeCode; + const BITS: u8; +} + +macro_rules! element { + ($($ty:ty => $code:ident, $bits:literal;)*) => { + $(impl Element for $ty { + const CODE: DLDataTypeCode = DLDataTypeCode::$code; + const BITS: u8 = $bits; + })* + }; +} + +element! { + i8 => kDLInt, 8; + i16 => kDLInt, 16; + i32 => kDLInt, 32; + i64 => kDLInt, 64; + u8 => kDLUInt, 8; + u16 => kDLUInt, 16; + u32 => kDLUInt, 32; + u64 => kDLUInt, 64; + half::f16 => kDLFloat, 16; + f32 => kDLFloat, 32; + f64 => kDLFloat, 64; + half::bf16 => kDLBfloat, 16; +} diff --git a/nix-builder/pkgs/cuda-oxide/default.nix b/nix-builder/pkgs/cuda-oxide/default.nix index 7b526cce..6ba36ef7 100644 --- a/nix-builder/pkgs/cuda-oxide/default.nix +++ b/nix-builder/pkgs/cuda-oxide/default.nix @@ -10,11 +10,14 @@ }: let + # NVIDIA/cuda-rust main with `CudaStream::borrow_raw`, which kernels need to + # launch on the Torch stream. Move back to NVIDIA/cuda-rust once + # https://github.com/NVIDIA/cuda-rust/pull/1601 is merged. src = fetchFromGitHub { - owner = "NVIDIA"; - repo = "cuda-rust"; - rev = "6d47ae49cb256a07948fb499c20c3eb9b2a96268"; - hash = "sha256-g9qg2zHLybO8kPLoOJqnYbWjRYtjGV3QVT+X4f8+1fg="; + owner = "drbh"; + repo = "cuda-oxide"; + rev = "d9d4e379e9a80d1ce7c681d88a086aacec95ba23"; + hash = "sha256-/rfi6TODuZp/YSoAy/bpse3nYJSokpmMwg3bHTtCNBM="; }; # rust-overlay propagates the compiler of its own package set. Drop it so @@ -39,10 +42,7 @@ in version = "0.2.1"; inherit src; sourceRoot = "source/cuda-oxide/crates/rustc-codegen-cuda"; - cargoLock = { - lockFile = "${src}/cuda-oxide/crates/rustc-codegen-cuda/Cargo.lock"; - outputHashes."pliron-0.18.0" = "sha256-HQ/xNfQcIf9pQRrgxCYGFdVfdCMaE4iwj/G9mgAJOWw="; - }; + cargoLock.lockFile = "${src}/cuda-oxide/crates/rustc-codegen-cuda/Cargo.lock"; buildInputs = [ libffi libxml2 From 69659859dbe93b72851454406d6cc0a2a0bb3d6f Mon Sep 17 00:00:00 2001 From: David Holtz Date: Thu, 22 Oct 2026 15:06:52 +0000 Subject: [PATCH 07/14] fix: prefer pushed kernels-cuda crate --- examples/kernels/relu-rust-cuda/Cargo.lock | 1 + examples/kernels/relu-rust-cuda/Cargo.toml | 4 +--- examples/kernels/relu-rust-cuda/rust-git-hashes.json | 1 + 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/kernels/relu-rust-cuda/Cargo.lock b/examples/kernels/relu-rust-cuda/Cargo.lock index 2e9be920..1923c29a 100644 --- a/examples/kernels/relu-rust-cuda/Cargo.lock +++ b/examples/kernels/relu-rust-cuda/Cargo.lock @@ -217,6 +217,7 @@ checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" [[package]] name = "kernels-cuda" version = "0.1.0" +source = "git+https://github.com/huggingface/kernels.git?rev=ae983e7ce660d4ed3a1142d6b9a540b7576be47d#ae983e7ce660d4ed3a1142d6b9a540b7576be47d" dependencies = [ "cuda-core", "half", diff --git a/examples/kernels/relu-rust-cuda/Cargo.toml b/examples/kernels/relu-rust-cuda/Cargo.toml index 0ea0463d..1e70c1d3 100644 --- a/examples/kernels/relu-rust-cuda/Cargo.toml +++ b/examples/kernels/relu-rust-cuda/Cargo.toml @@ -17,9 +17,7 @@ cuda-core = { git = "https://github.com/drbh/cuda-oxide.git", rev = "d9d4e379e9a cuda-device = { git = "https://github.com/drbh/cuda-oxide.git", rev = "d9d4e379e9a80d1ce7c681d88a086aacec95ba23" } cuda-host = { git = "https://github.com/drbh/cuda-oxide.git", rev = "d9d4e379e9a80d1ce7c681d88a086aacec95ba23" } tvm-ffi = "=0.1.14-post1" -# TODO: a git dependency on huggingface/kernels once the crate is pushed; the -# Nix build only sees this directory. -kernels-cuda = { path = "../../../kernels-cuda" } +kernels-cuda = { git = "https://github.com/huggingface/kernels.git", rev = "ae983e7ce660d4ed3a1142d6b9a540b7576be47d" } # Checked-in bindings: the build needs no CUDA headers or libclang. [patch."https://github.com/drbh/cuda-oxide.git"] diff --git a/examples/kernels/relu-rust-cuda/rust-git-hashes.json b/examples/kernels/relu-rust-cuda/rust-git-hashes.json index 7f4a3d2e..1ffe2ff6 100644 --- a/examples/kernels/relu-rust-cuda/rust-git-hashes.json +++ b/examples/kernels/relu-rust-cuda/rust-git-hashes.json @@ -7,6 +7,7 @@ "cuda-host-0.2.1": "sha256-cz53C3hB3dxVrVLgX+KtZ4WGT+jjoazNsBSedODCFrg=", "cuda-macros-0.2.1": "sha256-cz53C3hB3dxVrVLgX+KtZ4WGT+jjoazNsBSedODCFrg=", "cuda-target-spec-0.1.0": "sha256-cz53C3hB3dxVrVLgX+KtZ4WGT+jjoazNsBSedODCFrg=", + "kernels-cuda-0.1.0": "sha256-DQA+7fZEc7jCz8XMm4pldbHgT3bncTWHqgroiKE8XBQ=", "libnvvm-sys-0.2.1": "sha256-cz53C3hB3dxVrVLgX+KtZ4WGT+jjoazNsBSedODCFrg=", "nvjitlink-sys-0.2.1": "sha256-cz53C3hB3dxVrVLgX+KtZ4WGT+jjoazNsBSedODCFrg=", "oxide-tile-export-0.2.1": "sha256-cz53C3hB3dxVrVLgX+KtZ4WGT+jjoazNsBSedODCFrg=", From 92a9b4afa1af32cf32eadd1a54b50e5e593e6a97 Mon Sep 17 00:00:00 2001 From: David Holtz Date: Thu, 22 Oct 2026 18:23:35 +0000 Subject: [PATCH 08/14] feat: compile rust cuda once per capability --- docs/source/builder/writing-kernels.md | 7 +- .../src/pyproject/templates/rust.cmake | 124 ++++++++++++------ 2 files changed, 91 insertions(+), 40 deletions(-) diff --git a/docs/source/builder/writing-kernels.md b/docs/source/builder/writing-kernels.md index 9b72eb0e..6e06b4ca 100644 --- a/docs/source/builder/writing-kernels.md +++ b/docs/source/builder/writing-kernels.md @@ -383,8 +383,11 @@ Rust CUDA kernels are built with [cuda-oxide](https://github.com/NVIDIA/cuda-rust) and take the same options as `rust-cpu`, plus: -- `cuda-capabilities` (optional): same as for `cuda`. The kernels are compiled - to SASS for each capability, so nothing is JIT-compiled at load time. +- `cuda-capabilities` (optional): same as for `cuda`. Like with `nvcc`, the + kernels are compiled once per capability, to SASS, so nothing is JIT-compiled + at load time. Each pass sets `cfg(cuda_arch = "XX")` for its capability and + every lower one, so `#[cfg(cuda_arch = "90")]` works like + `__CUDA_ARCH__ >= 900`. The kernels and their host code are a single crate. Load the kernels with the `kernels-cuda` crate instead of the generated `load`, which fails to link. diff --git a/kernel-builder/src/pyproject/templates/rust.cmake b/kernel-builder/src/pyproject/templates/rust.cmake index 3cf68f7b..7bcae7f0 100644 --- a/kernel-builder/src/pyproject/templates/rust.cmake +++ b/kernel-builder/src/pyproject/templates/rust.cmake @@ -49,16 +49,17 @@ function(rust_kernel_component LIBS_VAR TARGETS_VAR) message(FATAL_ERROR "rust_kernel_component: Rust CUDA kernels require objcopy") endif() - list(SORT _KERNEL_ARCHS COMPARE NATURAL) - # ptxas compiles PTX for newer archs too, so target the lowest one. - list(GET _KERNEL_ARCHS 0 _OXIDE_ARCH) - string(REPLACE "+PTX" "" _OXIDE_ARCH "${_OXIDE_ARCH}") - string(REPLACE "." "" _OXIDE_ARCH "${_OXIDE_ARCH}") + set(_SMS) + set(_PTX_SMS) + foreach(_ARCH IN LISTS _KERNEL_ARCHS) + _rust_cuda_sm(_SM ${_ARCH}) + list(APPEND _SMS ${_SM}) + if(_ARCH MATCHES "\\+PTX$") + list(APPEND _PTX_SMS ${_SM}) + endif() + endforeach() # Flags match what `cargo oxide build` passes to the codegen backend. - list(APPEND _CARGO_ENV - "CUDA_OXIDE_TARGET=sm_${_OXIDE_ARCH}" - "CUDA_OXIDE_PTX_DIR=${CMAKE_BINARY_DIR}/cargo/${KERNEL_NAME}/ptx" - "RUSTFLAGS=-Zcodegen-backend=${CUDA_OXIDE_BACKEND} -Copt-level=3 -Cdebug-assertions=off -Zmir-enable-passes=-JumpThreading -Csymbol-mangling-version=v0") + set(_RUSTFLAGS "-Zcodegen-backend=${CUDA_OXIDE_BACKEND} -Copt-level=3 -Cdebug-assertions=off -Zmir-enable-passes=-JumpThreading -Csymbol-mangling-version=v0 --check-cfg=cfg(cuda_arch,values(any()))") endif() string(REPLACE "-" "_" _LIB_NAME ${KERNEL_NAME}) @@ -68,9 +69,35 @@ function(rust_kernel_component LIBS_VAR TARGETS_VAR) set(_LINKED_STATICLIB ${_STATICLIB}) set(_BYPRODUCTS ${_STATICLIB}) + # tvm-ffi-sys's build script shells out to `tvm-ffi-config`, a console script + # the apache-tvm-ffi wheel installs beside the interpreter. + get_filename_component(_PYTHON_BIN_DIR ${Python_EXECUTABLE} DIRECTORY) + set(_CARGO_COMMAND ${CMAKE_COMMAND} -E env "PATH=${_PYTHON_BIN_DIR}:$ENV{PATH}") + set(_CARGO_ARGS rustc --release --locked --lib --crate-type staticlib + --manifest-path ${CMAKE_CURRENT_SOURCE_DIR}/${KERNEL_MANIFEST_PATH}) + if(GPU_LANG STREQUAL "CUDA") - set(_PTX ${_CARGO_TARGET_DIR}/ptx/${_LIB_NAME}.ptx) - list(APPEND _BYPRODUCTS ${_PTX}) + # Like nvcc's `-gencode`, compile the device code once per capability, + # with `cuda_arch` set for it and every lower capability, so kernels can + # branch on it like on `__CUDA_ARCH__`. The first pass also builds the + # host code we link. + set(_COMMANDS) + set(_PTXS) + foreach(_SM IN LISTS _SMS) + if(NOT _COMMANDS) + set(_PASS_DIR ${_CARGO_TARGET_DIR}) + else() + set(_PASS_DIR ${_CARGO_TARGET_DIR}/sm_${_SM}) + endif() + _rust_cuda_cfgs(_CFGS ${_SM}) + list(APPEND _COMMANDS COMMAND ${_CARGO_COMMAND} + "CUDA_OXIDE_TARGET=sm_${_SM}" + "CUDA_OXIDE_PTX_DIR=${_PASS_DIR}/ptx" + "RUSTFLAGS=${_RUSTFLAGS} ${_CFGS}" + ${CARGO_EXECUTABLE} ${_CARGO_ARGS} --target-dir ${_PASS_DIR}) + list(APPEND _PTXS ${_PASS_DIR}/ptx/${_LIB_NAME}.ptx) + endforeach() + list(APPEND _BYPRODUCTS ${_PTXS}) # The backend also embeds the PTX in the crate's `.oxart` sections. The # extension loads the fatbin instead, so drop the PTX from what we link. set(_LINKED_STATICLIB ${_CARGO_TARGET_DIR}/${CMAKE_STATIC_LIBRARY_PREFIX}${_LIB_NAME}_sass${CMAKE_STATIC_LIBRARY_SUFFIX}) @@ -78,18 +105,13 @@ function(rust_kernel_component LIBS_VAR TARGETS_VAR) # Keeping the date lets a no-op cargo build skip the relink. set(_STRIP_PTX COMMAND ${CMAKE_OBJCOPY} --preserve-dates --remove-section=.oxart ${_STATICLIB} ${_LINKED_STATICLIB}) + else() + set(_COMMANDS COMMAND ${_CARGO_COMMAND} + ${CARGO_EXECUTABLE} ${_CARGO_ARGS} --target-dir ${_CARGO_TARGET_DIR}) endif() - # tvm-ffi-sys's build script shells out to `tvm-ffi-config`, a console script - # the apache-tvm-ffi wheel installs beside the interpreter. - 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}" - ${_CARGO_ENV} - ${CARGO_EXECUTABLE} rustc --release --locked --lib --crate-type staticlib - --manifest-path ${CMAKE_CURRENT_SOURCE_DIR}/${KERNEL_MANIFEST_PATH} - --target-dir ${_CARGO_TARGET_DIR} + ${_COMMANDS} ${_STRIP_PTX} BYPRODUCTS ${_BYPRODUCTS} WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} @@ -101,44 +123,70 @@ function(rust_kernel_component LIBS_VAR TARGETS_VAR) set_target_properties(${KERNEL_NAME}_rust PROPERTIES IMPORTED_LOCATION ${_LINKED_STATICLIB}) if(GPU_LANG STREQUAL "CUDA") - _rust_cuda_fatbin(${KERNEL_NAME} ${_PTX} ${_OXIDE_ARCH} "${_KERNEL_ARCHS}") + _rust_cuda_fatbin(${KERNEL_NAME} "${_SMS}" "${_PTXS}" "${_PTX_SMS}") endif() set(${LIBS_VAR} ${${LIBS_VAR}} ${KERNEL_NAME}_rust PARENT_SCOPE) set(${TARGETS_VAR} ${${TARGETS_VAR}} ${KERNEL_NAME}_cargo_build PARENT_SCOPE) endfunction() -# Compiles the crate's PTX to SASS for every capability, so nothing is -# JIT-compiled on the user's machine. `target_link_rust_kernels` links the -# fatbin between `KERNEL_BUILDER_CUDA_FATBIN` and `KERNEL_BUILDER_CUDA_FATBIN_END`. -function(_rust_cuda_fatbin NAME PTX PTX_ARCH ARCHS) - get_filename_component(_DIR ${PTX} DIRECTORY) +# The capability as in `sm_XX`. +function(_rust_cuda_sm OUT ARCH) + string(REPLACE "+PTX" "" _SM "${ARCH}") + string(REPLACE "." "" _SM "${_SM}") + set(${OUT} ${_SM} PARENT_SCOPE) +endfunction() + +# `cuda_arch` cfgs for one pass: every capability in CUDA_ARCHS up to SM, so a +# cfg means "this capability or newer", plus SM itself when it is arch-specific. +function(_rust_cuda_cfgs OUT SM) + string(REGEX REPLACE "[a-z]+$" "" _NUM ${SM}) + set(_KNOWN) + foreach(_ARCH IN LISTS CUDA_ARCHS) + _rust_cuda_sm(_KNOWN_SM ${_ARCH}) + string(REGEX REPLACE "[a-z]+$" "" _KNOWN_SM ${_KNOWN_SM}) + list(APPEND _KNOWN ${_KNOWN_SM}) + endforeach() + list(REMOVE_DUPLICATES _KNOWN) + set(_CFGS) + foreach(_KNOWN_SM IN LISTS _KNOWN) + if(_KNOWN_SM LESS_EQUAL _NUM) + list(APPEND _CFGS "--cfg=cuda_arch=\"${_KNOWN_SM}\"") + endif() + endforeach() + if(NOT SM STREQUAL _NUM) + list(APPEND _CFGS "--cfg=cuda_arch=\"${SM}\"") + endif() + list(JOIN _CFGS " " _CFGS) + set(${OUT} "${_CFGS}" PARENT_SCOPE) +endfunction() + +# Compiles each pass's PTX to SASS, so nothing is JIT-compiled on the user's +# machine. `target_link_rust_kernels` links the fatbin between +# `KERNEL_BUILDER_CUDA_FATBIN` and `KERNEL_BUILDER_CUDA_FATBIN_END`. +function(_rust_cuda_fatbin NAME SMS PTXS PTX_SMS) + list(GET PTXS 0 _FIRST_PTX) + get_filename_component(_DIR ${_FIRST_PTX} DIRECTORY) set(_CUBINS) set(_IMAGES) - foreach(_ARCH IN LISTS ARCHS) - string(REPLACE "+PTX" "" _SM "${_ARCH}") - string(REPLACE "." "" _SM "${_SM}") + foreach(_SM _PTX IN ZIP_LISTS SMS PTXS) set(_CUBIN ${_DIR}/${NAME}.sm_${_SM}.cubin) add_custom_command(OUTPUT ${_CUBIN} - COMMAND ${CUDA_PTXAS_EXECUTABLE} -arch=sm_${_SM} ${PTX} -o ${_CUBIN} - DEPENDS ${PTX} + COMMAND ${CUDA_PTXAS_EXECUTABLE} -arch=sm_${_SM} ${_PTX} -o ${_CUBIN} + DEPENDS ${_PTX} COMMENT "Compiling Rust kernel ${NAME} for sm_${_SM}" VERBATIM) list(APPEND _CUBINS ${_CUBIN}) list(APPEND _IMAGES "--image3=kind=elf,sm=${_SM},file=${_CUBIN}") - if(_ARCH MATCHES "\\+PTX$") - set(_KEEP_PTX ON) + if(_SM IN_LIST PTX_SMS) + list(APPEND _IMAGES "--image3=kind=ptx,sm=${_SM},file=${_PTX}") endif() endforeach() - if(_KEEP_PTX) - # The PTX targets the lowest arch, so it JITs on any newer GPU. - list(APPEND _IMAGES "--image3=kind=ptx,sm=${PTX_ARCH},file=${PTX}") - endif() set(_FATBIN ${_DIR}/${NAME}.fatbin) add_custom_command(OUTPUT ${_FATBIN} COMMAND ${CUDA_FATBINARY_EXECUTABLE} --create=${_FATBIN} ${_IMAGES} - DEPENDS ${_CUBINS} ${PTX} + DEPENDS ${_CUBINS} ${PTXS} COMMENT "Creating fatbin for Rust kernel ${NAME}" VERBATIM) From e595a96475796f0fad0f670737d750eee8b8a4f7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= Date: Tue, 27 Oct 2026 15:17:04 +0000 Subject: [PATCH 09/14] Pass libdevice path through `CUDA_OXIDE_LIBDEVICE` Some elementary functions (like `erff`) are not available yet, so we need to use `__nv_`. However, we need to provide the libdevice path through `CUDA_OXIDE_LIBDEVICE` to make it work. --- kernel-builder/src/pyproject/templates/rust.cmake | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/kernel-builder/src/pyproject/templates/rust.cmake b/kernel-builder/src/pyproject/templates/rust.cmake index 7bcae7f0..04f9ae5e 100644 --- a/kernel-builder/src/pyproject/templates/rust.cmake +++ b/kernel-builder/src/pyproject/templates/rust.cmake @@ -45,6 +45,10 @@ function(rust_kernel_component LIBS_VAR TARGETS_VAR) get_filename_component(_CUDA_BIN_DIR "${CMAKE_CUDA_COMPILER}" DIRECTORY) find_program(CUDA_PTXAS_EXECUTABLE ptxas HINTS ${_CUDA_BIN_DIR} REQUIRED) find_program(CUDA_FATBINARY_EXECUTABLE fatbinary HINTS ${_CUDA_BIN_DIR} REQUIRED) + # Without libdevice, cuda-oxide silently emits NVVM IR instead of PTX + # for kernels that call `__nv_*` functions (e.g. `__nv_expf`). + find_file(CUDA_LIBDEVICE libdevice.10.bc + HINTS ${_CUDA_BIN_DIR}/../nvvm/libdevice NO_DEFAULT_PATH REQUIRED) if(NOT CMAKE_OBJCOPY) message(FATAL_ERROR "rust_kernel_component: Rust CUDA kernels require objcopy") endif() @@ -93,6 +97,7 @@ function(rust_kernel_component LIBS_VAR TARGETS_VAR) list(APPEND _COMMANDS COMMAND ${_CARGO_COMMAND} "CUDA_OXIDE_TARGET=sm_${_SM}" "CUDA_OXIDE_PTX_DIR=${_PASS_DIR}/ptx" + "CUDA_OXIDE_LIBDEVICE=${CUDA_LIBDEVICE}" "RUSTFLAGS=${_RUSTFLAGS} ${_CFGS}" ${CARGO_EXECUTABLE} ${_CARGO_ARGS} --target-dir ${_PASS_DIR}) list(APPEND _PTXS ${_PASS_DIR}/ptx/${_LIB_NAME}.ptx) From 90c61854d8fbc9e563f34859bc469d6fc66fc552 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= Date: Tue, 27 Oct 2026 15:18:46 +0000 Subject: [PATCH 10/14] Use weak linkage for `__nv_` symbols We need to link against `__nv_*` symbols for some elementary functions. However, host copies of the kernels continue to be referenced by the launcher (even though they are never used), resulting in unresolved `__nv_*` symbols for host functions. Weaken linkage to avoid this issue. --- kernel-builder/src/pyproject/templates/rust.cmake | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/kernel-builder/src/pyproject/templates/rust.cmake b/kernel-builder/src/pyproject/templates/rust.cmake index 04f9ae5e..e93ece67 100644 --- a/kernel-builder/src/pyproject/templates/rust.cmake +++ b/kernel-builder/src/pyproject/templates/rust.cmake @@ -107,8 +107,12 @@ function(rust_kernel_component LIBS_VAR TARGETS_VAR) # extension loads the fatbin instead, so drop the PTX from what we link. set(_LINKED_STATICLIB ${_CARGO_TARGET_DIR}/${CMAKE_STATIC_LIBRARY_PREFIX}${_LIB_NAME}_sass${CMAKE_STATIC_LIBRARY_SUFFIX}) list(APPEND _BYPRODUCTS ${_LINKED_STATICLIB}) + # The launchers keep the host copies of the kernels alive, so their calls + # into libdevice (`__nv_*`, linked into the device code only) stay + # unresolved. The host copies never run: weaken those references. # Keeping the date lets a no-op cargo build skip the relink. set(_STRIP_PTX COMMAND ${CMAKE_OBJCOPY} --preserve-dates --remove-section=.oxart + --wildcard --weaken-symbol=__nv_* ${_STATICLIB} ${_LINKED_STATICLIB}) else() set(_COMMANDS COMMAND ${_CARGO_COMMAND} From e133da6246ff87d792b2c97dda6a39c00cfb3d21 Mon Sep 17 00:00:00 2001 From: drbh Date: Tue, 3 Nov 2026 10:04:54 -0500 Subject: [PATCH 11/14] fix: update docs/source/builder/writing-kernels.md MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Daniël de Kok --- docs/source/builder/writing-kernels.md | 2 -- 1 file changed, 2 deletions(-) diff --git a/docs/source/builder/writing-kernels.md b/docs/source/builder/writing-kernels.md index 6e06b4ca..56edc126 100644 --- a/docs/source/builder/writing-kernels.md +++ b/docs/source/builder/writing-kernels.md @@ -389,8 +389,6 @@ Rust CUDA kernels are built with every lower one, so `#[cfg(cuda_arch = "90")]` works like `__CUDA_ARCH__ >= 900`. -The kernels and their host code are a single crate. Load the kernels with the -`kernels-cuda` crate instead of the generated `load`, which fails to link. An extension can have only one `rust-cuda` component. Use the cuda-rust revision pinned in [`cuda-oxide`](https://github.com/huggingface/kernels/blob/main/nix-builder/pkgs/cuda-oxide/default.nix) From 5a74e8f93be7bf8cc73374b5d2b0ca96bdffed4a Mon Sep 17 00:00:00 2001 From: David Holtz Date: Tue, 3 Nov 2026 15:22:59 +0000 Subject: [PATCH 12/14] fix: move cuda-bindings-prebuilt to the huggingface org --- examples/kernels/relu-rust-cuda/Cargo.lock | 2 +- examples/kernels/relu-rust-cuda/Cargo.toml | 2 +- kernels-cuda/Cargo.toml | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/kernels/relu-rust-cuda/Cargo.lock b/examples/kernels/relu-rust-cuda/Cargo.lock index 1923c29a..9ce858dc 100644 --- a/examples/kernels/relu-rust-cuda/Cargo.lock +++ b/examples/kernels/relu-rust-cuda/Cargo.lock @@ -72,7 +72,7 @@ dependencies = [ [[package]] name = "cuda-bindings" version = "0.4.0" -source = "git+https://github.com/drbh/cuda-bindings-prebuilt.git?rev=f129e807d9a8ea32e7572369e91689c00c52570f#f129e807d9a8ea32e7572369e91689c00c52570f" +source = "git+https://github.com/huggingface/cuda-bindings-prebuilt.git?rev=f129e807d9a8ea32e7572369e91689c00c52570f#f129e807d9a8ea32e7572369e91689c00c52570f" dependencies = [ "libloading", ] diff --git a/examples/kernels/relu-rust-cuda/Cargo.toml b/examples/kernels/relu-rust-cuda/Cargo.toml index 1e70c1d3..45d482a8 100644 --- a/examples/kernels/relu-rust-cuda/Cargo.toml +++ b/examples/kernels/relu-rust-cuda/Cargo.toml @@ -21,4 +21,4 @@ kernels-cuda = { git = "https://github.com/huggingface/kernels.git", rev = "ae98 # Checked-in bindings: the build needs no CUDA headers or libclang. [patch."https://github.com/drbh/cuda-oxide.git"] -cuda-bindings = { git = "https://github.com/drbh/cuda-bindings-prebuilt.git", rev = "f129e807d9a8ea32e7572369e91689c00c52570f" } +cuda-bindings = { git = "https://github.com/huggingface/cuda-bindings-prebuilt.git", rev = "f129e807d9a8ea32e7572369e91689c00c52570f" } diff --git a/kernels-cuda/Cargo.toml b/kernels-cuda/Cargo.toml index d7826e3c..4d7139c3 100644 --- a/kernels-cuda/Cargo.toml +++ b/kernels-cuda/Cargo.toml @@ -17,4 +17,4 @@ tvm-ffi-sys = "=0.1.14-post1" # Checked-in bindings: the build needs no CUDA headers or libclang. [patch."https://github.com/drbh/cuda-oxide.git"] -cuda-bindings = { git = "https://github.com/drbh/cuda-bindings-prebuilt.git", rev = "f129e807d9a8ea32e7572369e91689c00c52570f" } +cuda-bindings = { git = "https://github.com/huggingface/cuda-bindings-prebuilt.git", rev = "f129e807d9a8ea32e7572369e91689c00c52570f" } From ac45b5ad64da986497e57141dcef2f56a207bc46 Mon Sep 17 00:00:00 2001 From: David Holtz Date: Tue, 3 Nov 2026 16:27:12 +0000 Subject: [PATCH 13/14] feat: generate cuda bindings with bindgen and require cuda 13 for rust-cuda --- docs/source/builder/writing-kernels.md | 6 + examples/kernels/relu-rust-cuda/Cargo.lock | 309 +++++++++++++++++- examples/kernels/relu-rust-cuda/Cargo.toml | 4 - .../relu-rust-cuda/rust-git-hashes.json | 2 +- kernels-cuda/Cargo.toml | 4 - nix-builder/lib/build.nix | 8 +- nix-builder/lib/cache.nix | 6 +- nix-builder/lib/extension/tvm-ffi/arch.nix | 19 +- nix-builder/lib/kernel-config.nix | 5 + 9 files changed, 349 insertions(+), 14 deletions(-) diff --git a/docs/source/builder/writing-kernels.md b/docs/source/builder/writing-kernels.md index 56edc126..c38d02f8 100644 --- a/docs/source/builder/writing-kernels.md +++ b/docs/source/builder/writing-kernels.md @@ -389,6 +389,12 @@ Rust CUDA kernels are built with every lower one, so `#[cfg(cuda_arch = "90")]` works like `__CUDA_ARCH__ >= 900`. +Rust CUDA kernels require CUDA 13.0 or later, so they are only built for +CUDA 13.0+ build variants (a lower `general.cuda.minver` is raised to 13.0). The CUDA +bindings are generated with bindgen during the build, so building outside of +Nix requires a CUDA 13 toolkit (found via `CUDA_TOOLKIT_PATH` or `CUDA_HOME`) +and libclang. + An extension can have only one `rust-cuda` component. Use the cuda-rust revision pinned in [`cuda-oxide`](https://github.com/huggingface/kernels/blob/main/nix-builder/pkgs/cuda-oxide/default.nix) diff --git a/examples/kernels/relu-rust-cuda/Cargo.lock b/examples/kernels/relu-rust-cuda/Cargo.lock index 9ce858dc..02bf5ced 100644 --- a/examples/kernels/relu-rust-cuda/Cargo.lock +++ b/examples/kernels/relu-rust-cuda/Cargo.lock @@ -2,12 +2,50 @@ # 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.69.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "271383c67ccabffb7381723dea0672a673f292304fcb45c01cc648c7a8d58088" +dependencies = [ + "bitflags", + "cexpr", + "clang-sys", + "itertools", + "lazy_static", + "lazycell", + "log", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "rustc-hash", + "shlex", + "syn 2.0.119", + "which", +] + +[[package]] +name = "bitflags" +version = "2.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ded4057c258ba199e2d26386d3af3780957ecaee6c4ef4041c6b4b8b97c0b06" + [[package]] name = "block-buffer" version = "0.10.4" @@ -17,12 +55,32 @@ 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.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4e7648175b45a9a48536d676f68d918270699102aa8dab5496df06904c914600" +[[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" @@ -72,9 +130,14 @@ dependencies = [ [[package]] name = "cuda-bindings" version = "0.4.0" -source = "git+https://github.com/huggingface/cuda-bindings-prebuilt.git?rev=f129e807d9a8ea32e7572369e91689c00c52570f#f129e807d9a8ea32e7572369e91689c00c52570f" +source = "git+https://github.com/drbh/cuda-oxide.git?rev=d9d4e379e9a80d1ce7c681d88a086aacec95ba23#d9d4e379e9a80d1ce7c681d88a086aacec95ba23" dependencies = [ + "bindgen", "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", ] [[package]] @@ -150,12 +213,28 @@ dependencies = [ "crypto-common", ] +[[package]] +name = "either" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be" + [[package]] name = "equivalent" version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "00d174d5400e5e8fd687ad1049e2f578285fa914201b1af7e8b112a4546bd826" +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + [[package]] name = "foldhash" version = "0.1.5" @@ -172,6 +251,12 @@ dependencies = [ "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" @@ -198,6 +283,15 @@ version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +[[package]] +name = "home" +version = "0.5.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589533453244b0995c858700322199b2becb13b627df2851f64a2775d024abcf" +dependencies = [ + "windows-sys 0.59.0", +] + [[package]] name = "indexmap" version = "2.14.2" @@ -208,6 +302,15 @@ dependencies = [ "hashbrown 0.17.1", ] +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.18" @@ -225,6 +328,18 @@ dependencies = [ "tvm-ffi-sys", ] +[[package]] +name = "lazy_static" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" + +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + [[package]] name = "libc" version = "0.2.191" @@ -251,12 +366,40 @@ dependencies = [ "thiserror", ] +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "log" +version = "0.4.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d15f86d12da44cd8d4ec3c73984537eb7f4cadcb9776a575a0a8a05ca11005c" + [[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" @@ -278,6 +421,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + [[package]] name = "oxide-artifacts" version = "0.2.1" @@ -303,6 +452,16 @@ 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" @@ -350,6 +509,35 @@ 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.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +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-rust-cuda" version = "0.1.0" @@ -366,6 +554,25 @@ name = "reserved-oxide-symbols" version = "0.2.1" source = "git+https://github.com/drbh/cuda-oxide.git?rev=d9d4e379e9a80d1ce7c681d88a086aacec95ba23#d9d4e379e9a80d1ce7c681d88a086aacec95ba23" +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.59.0", +] + [[package]] name = "serde" version = "1.0.229" @@ -420,6 +627,12 @@ dependencies = [ "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" @@ -520,12 +733,106 @@ version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix", +] + [[package]] name = "windows-link" version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + [[package]] name = "zerocopy" version = "0.8.63" diff --git a/examples/kernels/relu-rust-cuda/Cargo.toml b/examples/kernels/relu-rust-cuda/Cargo.toml index 45d482a8..71f145ed 100644 --- a/examples/kernels/relu-rust-cuda/Cargo.toml +++ b/examples/kernels/relu-rust-cuda/Cargo.toml @@ -18,7 +18,3 @@ cuda-device = { git = "https://github.com/drbh/cuda-oxide.git", rev = "d9d4e379e cuda-host = { git = "https://github.com/drbh/cuda-oxide.git", rev = "d9d4e379e9a80d1ce7c681d88a086aacec95ba23" } tvm-ffi = "=0.1.14-post1" kernels-cuda = { git = "https://github.com/huggingface/kernels.git", rev = "ae983e7ce660d4ed3a1142d6b9a540b7576be47d" } - -# Checked-in bindings: the build needs no CUDA headers or libclang. -[patch."https://github.com/drbh/cuda-oxide.git"] -cuda-bindings = { git = "https://github.com/huggingface/cuda-bindings-prebuilt.git", rev = "f129e807d9a8ea32e7572369e91689c00c52570f" } diff --git a/examples/kernels/relu-rust-cuda/rust-git-hashes.json b/examples/kernels/relu-rust-cuda/rust-git-hashes.json index 1ffe2ff6..d76d6986 100644 --- a/examples/kernels/relu-rust-cuda/rust-git-hashes.json +++ b/examples/kernels/relu-rust-cuda/rust-git-hashes.json @@ -1,6 +1,6 @@ { "cuda-artifact-finalizer-0.2.1": "sha256-cz53C3hB3dxVrVLgX+KtZ4WGT+jjoazNsBSedODCFrg=", - "cuda-bindings-0.4.0": "sha256-ZfLWDaBRaZVrPFYUui9GhbL0z38DcJqfECDm/rvpJq8=", + "cuda-bindings-0.4.0": "sha256-cz53C3hB3dxVrVLgX+KtZ4WGT+jjoazNsBSedODCFrg=", "cuda-core-0.4.0": "sha256-cz53C3hB3dxVrVLgX+KtZ4WGT+jjoazNsBSedODCFrg=", "cuda-core-derive-0.4.0": "sha256-cz53C3hB3dxVrVLgX+KtZ4WGT+jjoazNsBSedODCFrg=", "cuda-device-0.2.1": "sha256-cz53C3hB3dxVrVLgX+KtZ4WGT+jjoazNsBSedODCFrg=", diff --git a/kernels-cuda/Cargo.toml b/kernels-cuda/Cargo.toml index 4d7139c3..539b18d3 100644 --- a/kernels-cuda/Cargo.toml +++ b/kernels-cuda/Cargo.toml @@ -14,7 +14,3 @@ tvm-ffi-sys = "=0.1.14-post1" # Kernel crates depend on cuda-oxide, so keep it out of the root workspace. [workspace] - -# Checked-in bindings: the build needs no CUDA headers or libclang. -[patch."https://github.com/drbh/cuda-oxide.git"] -cuda-bindings = { git = "https://github.com/huggingface/cuda-bindings-prebuilt.git", rev = "f129e807d9a8ea32e7572369e91689c00c52570f" } diff --git a/nix-builder/lib/build.nix b/nix-builder/lib/build.nix index a43fd6ca..3d8a512e 100644 --- a/nix-builder/lib/build.nix +++ b/nix-builder/lib/build.nix @@ -48,7 +48,13 @@ rec { # Step 1: build sets within the given bounds. buildSetsWithinBounds = let - minCuda = kernelConfig.toml.general.cuda.minver or "11.8"; + # Rust CUDA kernels generate their CUDA bindings with bindgen, which + # requires CUDA 13.0 or later headers. + minCuda = + let + minver = kernelConfig.toml.general.cuda.minver or "11.8"; + in + if kernelConfig.usesRustCuda && lib.versionOlder minver "13.0" then "13.0" else minver; maxCuda = kernelConfig.toml.general.cuda.maxver or "99.9"; minTorch = kernelConfig.toml.torch.minver or "2.0"; maxTorch = kernelConfig.toml.torch.maxver or "99.9"; diff --git a/nix-builder/lib/cache.nix b/nix-builder/lib/cache.nix index 22624e2e..5afbd7da 100644 --- a/nix-builder/lib/cache.nix +++ b/nix-builder/lib/cache.nix @@ -20,6 +20,8 @@ inherit (buildSet) overrideTorch; isLinux = stdenv.hostPlatform.isLinux; cudaSupport = config.cudaSupport; + # Rust CUDA kernels are only built for CUDA 13.0 and later. + rustCudaSupport = cudaSupport && lib.versionAtLeast buildSet.buildConfig.cudaVersion "13.0"; xpuSupport = config.xpuSupport or false; in ( @@ -34,8 +36,8 @@ ++ allOutputs python3.pkgs.tvm-ffi ++ allOutputs ruff ++ lib.optionals (isLinux && cudaSupport) (allOutputs manylinux_2_28.cudaBackendStdenv) - ++ lib.optionals (isLinux && cudaSupport) (allOutputs cuda-oxide.toolchain) - ++ lib.optionals (isLinux && cudaSupport) (allOutputs cuda-oxide.backend) + ++ lib.optionals (isLinux && rustCudaSupport) (allOutputs cuda-oxide.toolchain) + ++ lib.optionals (isLinux && rustCudaSupport) (allOutputs cuda-oxide.backend) ++ lib.optionals (isLinux && !config.cudaSupport) (allOutputs manylinux_2_28.stdenv) # Only works on recent CUDAs. ++ lib.optionals (!python3.pkgs.nvidia-cutlass-dsl.meta.broken) ( diff --git a/nix-builder/lib/extension/tvm-ffi/arch.nix b/nix-builder/lib/extension/tvm-ffi/arch.nix index cba89576..b16a0f81 100644 --- a/nix-builder/lib/extension/tvm-ffi/arch.nix +++ b/nix-builder/lib/extension/tvm-ffi/arch.nix @@ -138,6 +138,16 @@ let rustSupport = cargoLock != null; + # cuda-bindings runs bindgen against a single CUDA include directory that + # must provide both `cuda.h` and `curand.h`. + rustCudaToolkit = pkgs.symlinkJoin { + name = "rust-cuda-toolkit"; + paths = with cudaPackages; [ + (lib.getInclude cuda_cudart) + (lib.getInclude libcurand) + ]; + }; + # rpaths are stripped from kernels to make them portable, but that # also means that in a Nix environment the CUDA/ROCm/oneAPI dependencies # cannot be located anymore, so pass them to hooks that load the kernel. @@ -260,7 +270,11 @@ stdenv.mkDerivation ( rustc ] # Rust CUDA kernels are built with cuda-oxide and its matching rustc. - ++ lib.optionals (rustSupport && cudaSupport) [ cuda-oxide.toolchain ]; + ++ lib.optionals (rustSupport && cudaSupport) [ + cuda-oxide.toolchain + # libclang for the cuda-bindings build script. + rustPlatform.bindgenHook + ]; buildInputs = [ python3.pkgs.tvm-ffi @@ -310,6 +324,9 @@ stdenv.mkDerivation ( // lib.optionalAttrs cudaSupport { CUDAToolkit_ROOT = "${lib.getDev cudaPackages.cuda_nvcc}"; } + // lib.optionalAttrs (rustSupport && cudaSupport) { + CUDA_TOOLKIT_PATH = "${rustCudaToolkit}"; + } // lib.optionalAttrs xpuSupport { MKLROOT = oneapi-torch-dev; SYCL_ROOT = oneapi-torch-dev; diff --git a/nix-builder/lib/kernel-config.nix b/nix-builder/lib/kernel-config.nix index 60c7ca58..c3c571d1 100644 --- a/nix-builder/lib/kernel-config.nix +++ b/nix-builder/lib/kernel-config.nix @@ -77,5 +77,10 @@ in in lib.foldl (backends: kernel: backends // { ${kernelBackend kernel} = true; }) init kernels; + # Does the kernel have Rust CUDA (cuda-oxide) components. + usesRustCuda = lib.any (kernel: kernelLanguage kernel == "rust" && kernelBackend kernel == "cuda") ( + lib.attrValues (toml.kernel or { }) + ); + name = toml.general.name; } From 5e184e68aa8778016b7f17f1b7e3dc1d4b4c72c3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= Date: Wed, 4 Nov 2026 09:51:51 +0000 Subject: [PATCH 14/14] ci: use CUDA 13 for rust-cuda test --- examples/kernels/flake.nix | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/kernels/flake.nix b/examples/kernels/flake.nix index 32d65679..ca84296d 100644 --- a/examples/kernels/flake.nix +++ b/examples/kernels/flake.nix @@ -125,7 +125,7 @@ name = "relu-rust-cuda-kernel"; path = ./relu-rust-cuda; drv = - sys: out: out.packages.${sys}.redistributable.${"tvm-ffi${tvmFfiVersion}-${cudaVersion}-${sys}"}; + sys: out: out.packages.${sys}.redistributable.${"tvm-ffi${tvmFfiVersion}-${cuda13Version}-${sys}"}; } { name = "relu-kernel-cpu";