diff --git a/examples/kernels/flake.nix b/examples/kernels/flake.nix index b0085be49..1d5b5bedc 100644 --- a/examples/kernels/flake.nix +++ b/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 = diff --git a/examples/kernels/relu-tvm-ffi/build.toml b/examples/kernels/relu-tvm-ffi/build.toml index 2116d6d88..eaec722ed 100644 --- a/examples/kernels/relu-tvm-ffi/build.toml +++ b/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 = [] diff --git a/examples/kernels/relu-tvm-ffi/relu_cuda/relu.cu b/examples/kernels/relu-tvm-ffi/relu_cuda/relu.cu index 9dfd44e6f..52344bb9d 100644 --- a/examples/kernels/relu-tvm-ffi/relu_cuda/relu.cu +++ b/examples/kernels/relu-tvm-ffi/relu_cuda/relu.cu @@ -1,10 +1,31 @@ #include #include -#include #include #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 +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( TVMFFIEnvGetStream(input.device().device_type, input.device().device_id)); diff --git a/examples/kernels/relu-tvm-ffi/tvm-ffi-ext/relu_tvm_ffi/__init__.py b/examples/kernels/relu-tvm-ffi/tvm-ffi-ext/relu_tvm_ffi/__init__.py index 3c9ac203e..8afe2c232 100644 --- a/examples/kernels/relu-tvm-ffi/tvm-ffi-ext/relu_tvm_ffi/__init__.py +++ b/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"] diff --git a/examples/kernels/relu-tvm-ffi/tvm-ffi-ext/tvm_ffi_binding.cpp b/examples/kernels/relu-tvm-ffi/tvm-ffi-ext/tvm_ffi_binding.cpp index 88d440b23..33ef241e2 100644 --- a/examples/kernels/relu-tvm-ffi/tvm-ffi-ext/tvm_ffi_binding.cpp +++ b/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); diff --git a/examples/kernels/relu-tvm-ffi/util.hh b/examples/kernels/relu-tvm-ffi/util.hh index 4bd5a477d..b06f30506 100644 --- a/examples/kernels/relu-tvm-ffi/util.hh +++ b/examples/kernels/relu-tvm-ffi/util.hh @@ -1,7 +1,12 @@ #include +#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) \ diff --git a/kernel-builder/src/pyproject/templates/tvm_ffi/hipify.py b/kernel-builder/src/pyproject/templates/tvm_ffi/hipify.py new file mode 100644 index 000000000..010f91765 --- /dev/null +++ b/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)) diff --git a/kernel-builder/src/pyproject/templates/tvm_ffi/preamble.cmake b/kernel-builder/src/pyproject/templates/tvm_ffi/preamble.cmake index a0443f9db..86a665356 100644 --- a/kernel-builder/src/pyproject/templates/tvm_ffi/preamble.cmake +++ b/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") diff --git a/kernel-builder/src/pyproject/templates/tvm_ffi/tvm-ffi-extension.cmake b/kernel-builder/src/pyproject/templates/tvm_ffi/tvm-ffi-extension.cmake index 0fe3e32f1..6bb8d5660 100644 --- a/kernel-builder/src/pyproject/templates/tvm_ffi/tvm-ffi-extension.cmake +++ b/kernel-builder/src/pyproject/templates/tvm_ffi/tvm-ffi-extension.cmake @@ -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}) diff --git a/kernel-builder/src/pyproject/templates/utils.cmake b/kernel-builder/src/pyproject/templates/utils.cmake index 6f32b0b87..dc6f675a9 100644 --- a/kernel-builder/src/pyproject/templates/utils.cmake +++ b/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.") diff --git a/kernel-builder/src/pyproject/tvm_ffi/mod.rs b/kernel-builder/src/pyproject/tvm_ffi/mod.rs index 7f40d150c..32cd03014 100644 --- a/kernel-builder/src/pyproject/tvm_ffi/mod.rs +++ b/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", diff --git a/nix-builder/lib/extension/tvm-ffi/arch.nix b/nix-builder/lib/extension/tvm-ffi/arch.nix index 3f07c337d..f3cec3c70 100644 --- a/nix-builder/lib/extension/tvm-ffi/arch.nix +++ b/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. diff --git a/nix-builder/overlay.nix b/nix-builder/overlay.nix index 909173ba4..0977c700b 100644 --- a/nix-builder/overlay.nix +++ b/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 diff --git a/nix-builder/pkgs/python-modules/hipify-torch/default.nix b/nix-builder/pkgs/python-modules/hipify-torch/default.nix new file mode 100644 index 000000000..bee28fbf8 --- /dev/null +++ b/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; + }; +} diff --git a/nix-builder/versions.nix b/nix-builder/versions.nix index f72665ede..2a87e8d5f 100644 --- a/nix-builder/versions.nix +++ b/nix-builder/versions.nix @@ -106,12 +106,14 @@ rocmVersion = "7.1"; systems = [ "x86_64-linux" ]; bundleBuild = true; + tvmFfiVersion = "0.1"; } { torchVersion = "2.13"; rocmVersion = "7.2"; systems = [ "x86_64-linux" ]; bundleBuild = true; + tvmFfiVersion = "0.1"; } { torchVersion = "2.13"; @@ -175,6 +177,7 @@ rocmVersion = "7.14"; systems = [ "x86_64-linux" ]; bundleBuild = true; + tvmFfiVersion = "0.1"; } { torchVersion = "2.14";