nix-builder: add tvm-ffi support for ROCm

examples/kernels/flake.nix

@@ -274,6 +274,12 @@
           path = ./relu-compiler-flags;
           drv = sys: out: out.packages.${sys}.redistributable.${"torch${torchVersion}-${rocmVersion}-${sys}"};
         }
+        {
+          name = "relu-tvm-ffi-kernel";
+          path = ./relu-tvm-ffi;
+          drv =
+            sys: out: out.packages.${sys}.redistributable.${"tvm-ffi${tvmFfiVersion}-${rocmVersion}-${sys}"};
+        }
       ];
 
       mkKernelOutputs =

examples/kernels/relu-tvm-ffi/build.toml

@@ -6,6 +6,7 @@ license = "Apache-2.0"
 backends = [
     "cpu",
     "cuda",
+    "rocm",
     "xpu",
 ]
 
@@ -23,6 +24,14 @@ src = [
     "util.hh",
 ]
 
+[kernel.relu_rocm]
+language = "cpp-rocm"
+depends = []
+src = [
+    "relu_cuda/relu.cu",
+    "util.hh",
+]
+
 [kernel.relu_cpu]
 language = "cpp-cpu"
 depends = []

examples/kernels/relu-tvm-ffi/relu_cuda/relu.cu

@@ -1,10 +1,31 @@
 #include <cmath>
 #include <tvm/ffi/tvm_ffi.h>
-#include <tvm/ffi/extra/cuda/device_guard.h>
 #include <tvm/ffi/extra/c_env_api.h>
 
 #include "../util.hh"
 
+#ifdef USE_ROCM
+// tvm-ffi does not provide a device guard for ROCm.
+class DeviceGuard {
+public:
+  explicit DeviceGuard(int device_id) {
+    TVM_FFI_CHECK(hipGetDevice(&prev_device_id_) == hipSuccess, RuntimeError)
+        << "Cannot get current device";
+    if (prev_device_id_ != device_id) {
+      TVM_FFI_CHECK(hipSetDevice(device_id) == hipSuccess, RuntimeError)
+          << "Cannot set device";
+    }
+  }
+  ~DeviceGuard() { (void)hipSetDevice(prev_device_id_); }
+
+private:
+  int prev_device_id_;
+};
+#else
+#include <tvm/ffi/extra/cuda/device_guard.h>
+using DeviceGuard = tvm::ffi::CUDADeviceGuard;
+#endif
+
 __global__ void relu_kernel(float *__restrict__ out,
                             float const *__restrict__ input, const int d) {
   const int64_t token_idx = blockIdx.x;
@@ -28,7 +49,7 @@ void relu_cuda(ffi::TensorView out, ffi::TensorView const input) {
     return;
   }
 
-  ffi::CUDADeviceGuard guard(input.device().device_id);
+  DeviceGuard guard(input.device().device_id);
   cudaStream_t stream = static_cast<cudaStream_t>(
       TVMFFIEnvGetStream(input.device().device_type, input.device().device_id));
 

examples/kernels/relu-tvm-ffi/tvm-ffi-ext/relu_tvm_ffi/__init__.py

@@ -12,6 +12,8 @@ def relu(x, out) -> tvm_ffi.Tensor:
         ops.relu_cpu(out_t, x_t)
     elif device.type == "cuda":
         ops.relu_cuda(out_t, x_t)
+    elif device.type == "rocm":
+        ops.relu_rocm(out_t, x_t)
     elif device.type == "x_tpu":
         ops.relu_x_tpu(out_t, x)
     else:
@@ -23,6 +25,7 @@ def relu(x, out) -> tvm_ffi.Tensor:
 if has_jax:
     ops_func = (
         getattr(ops, "relu_cuda", None)
+        or getattr(ops, "relu_rocm", None)
         or getattr(ops, "relu_xpu", None)
         or getattr(ops, "relu_cpu", None)
     )
@@ -52,7 +55,7 @@ def relu_jax(x):
     )(x)
 
 
-from . import layers
-
+# Must be imported after `relu` is defined, `layers` imports it.
+from . import layers  # noqa: E402
 
 __all__ = ["layers", "relu", "relu_jax"]

examples/kernels/relu-tvm-ffi/tvm-ffi-ext/tvm_ffi_binding.cpp

@@ -10,6 +10,12 @@ void relu_cuda(tvm::ffi::TensorView out, tvm::ffi::TensorView const input);
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(relu_cuda, relu_cuda);
 #endif
 
+#ifdef ROCM_KERNEL
+// The ROCm kernel is the hipified CUDA kernel.
+void relu_cuda(tvm::ffi::TensorView out, tvm::ffi::TensorView const input);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(relu_rocm, relu_cuda);
+#endif
+
 #ifdef XPU_KERNEL
 void relu_xpu(tvm::ffi::TensorView out, tvm::ffi::TensorView const input);
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(relu_xpu, relu_xpu);

examples/kernels/relu-tvm-ffi/util.hh

@@ -1,7 +1,12 @@
 #include <tvm/ffi/tvm_ffi.h>
 
+#ifdef USE_ROCM
+#define CHECK_CUDA(x) \
+  TVM_FFI_CHECK((x).device().device_type == kDLROCM, ValueError) << #x " must be a ROCm tensor"
+#else
 #define CHECK_CUDA(x) \
   TVM_FFI_CHECK((x).device().device_type == kDLCUDA, ValueError) << #x " must be a CUDA tensor"
+#endif
 #define CHECK_CONTIGUOUS(x) \
   TVM_FFI_CHECK((x).IsContiguous(), ValueError) << #x " must be contiguous"
 #define CHECK_INPUT(x)   \

kernel-builder/src/pyproject/templates/tvm_ffi/hipify.py

@@ -0,0 +1,83 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+
+# Based on vLLM: https://github.com/vllm-project/vllm/blob/main/cmake/hipify.py
+
+#
+# A command line tool for running the hipify preprocessor on CUDA source
+# files. Unlike the Torch variant, this uses hipify_torch directly, so that
+# Torch is not required.
+#
+# See https://github.com/ROCm/hipify_torch
+#
+
+import argparse
+import os
+import shutil
+import sys
+
+try:
+    from hipify_torch.v2.hipify_python import hipify
+except ImportError:
+    sys.exit(
+        "hipify_torch (including its v2 module) is required to build ROCm kernels, "
+        "see https://github.com/ROCm/hipify_torch"
+    )
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    # Project directory where all the source + include files live.
+    parser.add_argument(
+        "-p",
+        "--project_dir",
+        help="The project directory.",
+    )
+
+    # Directory where hipified files are written.
+    parser.add_argument(
+        "-o",
+        "--output_dir",
+        help="The output directory.",
+    )
+
+    # Source files to convert.
+    parser.add_argument("sources", help="Source files to hipify.", nargs="*", default=[])
+
+    args = parser.parse_args()
+
+    # Limit include scope to project_dir only
+    includes = [os.path.join(args.project_dir, "*")]
+
+    # Get absolute path for all source files.
+    extra_files = [os.path.abspath(s) for s in args.sources]
+
+    # Copy sources from project directory to output directory.
+    # The directory might already exist to hold object files so we ignore that.
+    shutil.copytree(args.project_dir, args.output_dir, dirs_exist_ok=True)
+
+    hipify_result = hipify(
+        project_directory=args.project_dir,
+        output_directory=args.output_dir,
+        header_include_dirs=[],
+        includes=includes,
+        extra_files=extra_files,
+        show_detailed=True,
+        is_pytorch_extension=False,
+        hipify_extra_files_only=True,
+    )
+
+    hipified_sources = []
+    for source in args.sources:
+        s_abs = os.path.abspath(source)
+        hipified_s_abs = (
+            hipify_result[s_abs].hipified_path
+            if (s_abs in hipify_result and hipify_result[s_abs].hipified_path is not None)
+            else s_abs
+        )
+        hipified_sources.append(hipified_s_abs)
+
+    assert len(hipified_sources) == len(args.sources)
+
+    # Print hipified source files.
+    print("\n".join(hipified_sources))

kernel-builder/src/pyproject/templates/tvm_ffi/preamble.cmake

@@ -40,7 +40,12 @@ if(NOT DEFINED GPU_LANG)
         if(CMAKE_CUDA_COMPILER)
             set(DETECTED_GPU_LANG "CUDA")
         else()
-            set(DETECTED_GPU_LANG "CPU")
+            check_language(HIP)
+            if(CMAKE_HIP_COMPILER)
+                set(DETECTED_GPU_LANG "HIP")
+            else()
+                set(DETECTED_GPU_LANG "CPU")
+            endif()
         endif()
     endif()
 
@@ -120,6 +125,18 @@ if(GPU_LANG STREQUAL "CUDA")
   endif()
 
   add_compile_definitions(CUDA_KERNEL)
+elseif(GPU_LANG STREQUAL "HIP")
+  enable_language(HIP)
+
+  find_package(hip REQUIRED)
+  set(ROCM_VERSION "${hip_VERSION_MAJOR}.${hip_VERSION_MINOR}")
+  message(STATUS "Found ROCm version: ${ROCM_VERSION}")
+
+  set(ROCM_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1200;gfx1201")
+  message(STATUS "ROCm supported target architectures: ${ROCM_ARCHS}")
+
+  add_compile_definitions(USE_ROCM=1)
+  add_compile_definitions(ROCM_KERNEL)
 elseif(GPU_LANG STREQUAL "CPU")
   add_compile_definitions(CPU_KERNEL)
   set(CMAKE_OSX_DEPLOYMENT_TARGET "15.0" CACHE STRING "Minimum macOS deployment version")

kernel-builder/src/pyproject/templates/tvm_ffi/tvm-ffi-extension.cmake

@@ -3,7 +3,20 @@ set(CMAKE_SHARED_LIBRARY_PREFIX "")
 
 rust_extension_sources(SRC)
 
+# Add hipify preprocessing step when building with HIP/ROCm.
+if(GPU_LANG STREQUAL "HIP")
+  hipify_sources_target(SRC ${OPS_NAME} "${SRC}")
+endif()
+
 add_library(${OPS_NAME} SHARED ${SRC})
+
+if(GPU_LANG STREQUAL "HIP")
+  # Make this target dependent on the hipify preprocessor step.
+  add_dependencies(${OPS_NAME} hipify${OPS_NAME})
+
+  # Clear target architectures, we are passing arch flags per source file.
+  set_property(TARGET ${OPS_NAME} PROPERTY HIP_ARCHITECTURES off)
+endif()
 target_compile_definitions(${OPS_NAME} PRIVATE
   "-DTVM_FFI_EXTENSION_NAME=${OPS_NAME}")
 tvm_ffi_configure_target(${OPS_NAME})

kernel-builder/src/pyproject/templates/utils.cmake

@@ -158,9 +158,17 @@ function (hipify_sources_target OUT_SRCS NAME ORIG_SRCS)
     list(APPEND HIP_SRCS "${CMAKE_CURRENT_BINARY_DIR}/${SRC}")
   endforeach()
 
+  if(Python3_EXECUTABLE)
+    set(_PYTHON_EXECUTABLE "${Python3_EXECUTABLE}")
+  elseif(Python_EXECUTABLE)
+    set(_PYTHON_EXECUTABLE "${Python_EXECUTABLE}")
+  else()
+    message(FATAL_ERROR "No Python executable found. Set Python3_EXECUTABLE or Python_EXECUTABLE.")
+  endif()
+
   add_custom_target(
     hipify${NAME}
-    COMMAND "${Python3_EXECUTABLE}" ${CMAKE_SOURCE_DIR}/cmake/hipify.py -p ${CMAKE_SOURCE_DIR} -o ${CMAKE_CURRENT_BINARY_DIR} ${SRCS}
+    COMMAND "${_PYTHON_EXECUTABLE}" ${CMAKE_SOURCE_DIR}/cmake/hipify.py -p ${CMAKE_SOURCE_DIR} -o ${CMAKE_CURRENT_BINARY_DIR} ${SRCS}
     DEPENDS ${CMAKE_SOURCE_DIR}/cmake/hipify.py ${SRCS}
     BYPRODUCTS ${HIP_SRCS}
     COMMENT "Running hipify on ${NAME} extension source files.")

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

@@ -18,6 +18,7 @@ static BUILD_VARIANTS_UTILS: &str = include_str!("../templates/tvm_ffi/build-var
 static CMAKE_KERNEL: &str = include_str!("../templates/kernel.cmake");
 static CMAKE_RUST: &str = include_str!("../templates/rust.cmake");
 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");
 static DETECT_CUDA_CAPABILITY_PY: &str =
     include_str!("../templates/tvm_ffi/cuda/detect-cuda-capability.py");
@@ -33,6 +34,7 @@ fn write_cmake_helpers(file_set: &mut FileSet) {
     );
     write_cmake_file(file_set, "_ops.py.in", OPS_PY_IN.as_bytes());
     write_add_build_metadata_py(file_set);
+    write_cmake_file(file_set, "hipify.py", HIPIFY.as_bytes());
     write_cmake_file(
         file_set,
         "cuda/detect-cuda-capability.py",

nix-builder/lib/extension/tvm-ffi/arch.nix

@@ -138,7 +138,7 @@ let
   rustSupport = cargoLock != null;
 
   # rpaths are stripped from kernels to make them portable, but that
-  # also means that in a Nix environment the CUDA/oneAPI dependencies
+  # 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 (
@@ -151,6 +151,7 @@ let
           libcusparse
         ]
       )
+      ++ lib.optionals rocmSupport [ clr ]
       ++ lib.optionals xpuSupport (
         with xpuPackages;
         [
@@ -321,10 +322,13 @@ stdenv.mkDerivation (
       (lib.cmakeBool "BUILD_ALL_SUPPORTED_ARCHS" true)
       (lib.cmakeFeature "Python_EXECUTABLE" "${
         python3.withPackages (
-          ps: with ps; [
+          ps:
+          with ps;
+          [
             tvm-ffi
             typing-extensions
           ]
+          ++ lib.optionals rocmSupport [ hipify-torch ]
         )
       }/bin/python")
       # Fix: file RPATH_CHANGE could not write new RPATH, we are rewriting
@@ -340,6 +344,7 @@ stdenv.mkDerivation (
       # the symlink-joined ROCm toolkit.
       (lib.cmakeFeature "CMAKE_HIP_COMPILER_ROCM_ROOT" "${clr}")
       (lib.cmakeFeature "HIP_ROOT_DIR" "${clr}")
+      (lib.cmakeFeature "GPU_LANG" "HIP")
     ]
     ++ lib.optionals metalSupport [
       # Use host compiler for Metal. Not included in the redistributable SDK.

nix-builder/overlay.nix

@@ -171,6 +171,8 @@ final: prev:
           };
         });
 
+        hipify-torch = python-self.callPackage ./pkgs/python-modules/hipify-torch { };
+
         helion = python-super.helion.overrideAttrs (
           _: prevAttrs: {
             # nixpkgs marks Helion as broken without CUDA, but it also

nix-builder/pkgs/python-modules/hipify-torch/default.nix

@@ -0,0 +1,41 @@
+{
+  lib,
+  buildPythonPackage,
+  fetchFromGitHub,
+
+  setuptools,
+}:
+
+buildPythonPackage {
+  pname = "hipify-torch";
+  version = "1.1.0-unstable-2026-03-05";
+  pyproject = true;
+
+  src = fetchFromGitHub {
+    owner = "ROCm";
+    repo = "hipify_torch";
+    rev = "1ea3231415a41c07f9e1f1d41906df08a9af390d";
+    hash = "sha256-WuPZ1hCX9B97lnd1AqDleFmRrbjuMjlS9VBiwmFNsqA=";
+  };
+
+  # Upstream does not install the v2 version of hipify. Torch uses v2 of
+  # hipify, so we want to use the same to avoid divergences between Torch
+  # kernels and e.g. TVM-FFI.
+  postPatch = ''
+    substituteInPlace setup.py \
+      --replace-fail "packages=['hipify_torch',]" "packages=['hipify_torch', 'hipify_torch.v2']"
+  '';
+
+  build-system = [ setuptools ];
+
+  pythonImportsCheck = [
+    "hipify_torch.hipify_python"
+    "hipify_torch.v2.hipify_python"
+  ];
+
+  meta = {
+    description = "Convert CUDA C/C++ code into HIP C/C++ code";
+    homepage = "https://github.com/ROCm/hipify_torch";
+    license = lib.licenses.mit;
+  };
+}