Add a CUDA/HIP runtime compatibility layer

CMakeLists.txt

@@ -236,7 +236,9 @@ if(EXECUTORCH_ENABLE_BUNDLE_IO)
  add_definitions(-DET_BUNDLE_IO_ENABLED)
 endif()
 
-if(EXECUTORCH_BUILD_CUDA)
+if(EXECUTORCH_BUILD_ROCM)
+ add_definitions(-DCUDA_AVAILABLE=1 -DEXECUTORCH_USE_HIP=1)
+elseif(EXECUTORCH_BUILD_CUDA)
  add_definitions(-DCUDA_AVAILABLE=1)
 endif()
 
@@ -729,6 +731,25 @@ install(FILES tools/cmake/executorch-config.cmake
  DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/ExecuTorch
 )
 
+# Always overwrite backend dependency metadata to avoid stale requirements.
+set(_executorch_backend_dependencies "")
+if(EXECUTORCH_BUILD_ROCM)
+ string(APPEND _executorch_backend_dependencies
+ "find_dependency(hip CONFIG)\n"
+ )
+elseif(EXECUTORCH_BUILD_CUDA)
+ string(APPEND _executorch_backend_dependencies
+ "find_dependency(CUDAToolkit)\n"
+ )
+endif()
+configure_file(
+ tools/cmake/executorch-backend-dependencies.cmake.in
+ ${CMAKE_CURRENT_BINARY_DIR}/executorch-backend-dependencies.cmake @ONLY
+)
+install(FILES ${CMAKE_CURRENT_BINARY_DIR}/executorch-backend-dependencies.cmake
+ DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/ExecuTorch
+)
+
 if(EXECUTORCH_BUILD_ARM_BAREMETAL
  OR EXECUTORCH_BUILD_ARM_ETHOSU_LINUX
  OR EXECUTORCH_BUILD_VGF
@@ -801,6 +822,7 @@ endif()
 # imported targets (e.g. MKL::MKL) aren't created in child directory scopes and
 # then duplicated when found again at root scope.
 if(EXECUTORCH_BUILD_CUDA
+ OR EXECUTORCH_BUILD_ROCM
  OR EXECUTORCH_BUILD_METAL
  OR EXECUTORCH_BUILD_PYBIND
 )
@@ -810,7 +832,7 @@ endif()
 # Backend-neutral caller-stream guard consumed by the CUDA AOTI backend (and the
 # vendored torch-tensorrt delegate). Built before backends/aoti and
 # backends/cuda, which link it.
-if(EXECUTORCH_BUILD_CUDA)
+if(EXECUTORCH_BUILD_CUDA OR EXECUTORCH_BUILD_ROCM)
  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/cuda)
  install(
  DIRECTORY extension/cuda/
@@ -821,12 +843,15 @@ if(EXECUTORCH_BUILD_CUDA)
  list(APPEND _executorch_extensions extension_cuda)
 endif()
 
-# Build common AOTI functionality if needed by CUDA or Metal backends
-if(EXECUTORCH_BUILD_CUDA OR EXECUTORCH_BUILD_METAL)
+# Build common AOTI functionality if needed by CUDA, ROCm, or Metal backends.
+if(EXECUTORCH_BUILD_CUDA
+ OR EXECUTORCH_BUILD_ROCM
+ OR EXECUTORCH_BUILD_METAL
+)
  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/aoti)
 endif()
 
-if(EXECUTORCH_BUILD_CUDA)
+if(EXECUTORCH_BUILD_CUDA OR EXECUTORCH_BUILD_ROCM)
  # Build CUDA-specific AOTI functionality
  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/cuda)
  # Add aoti_cuda_backend to backends - it transitively includes aoti_cuda_shims
@@ -1049,8 +1074,8 @@ if(EXECUTORCH_BUILD_PYBIND)
  torch
  )
 
- # Build common AOTI functionality if needed by CUDAor Metal backends
- if(EXECUTORCH_BUILD_CUDA)
+ # Build common AOTI functionality if needed by CUDA, ROCm, or Metal backends.
+ if(EXECUTORCH_BUILD_CUDA OR EXECUTORCH_BUILD_ROCM)
  # CUDA uses SlimTensor-based shims
  list(APPEND _dep_libs aoti_cuda_backend)
  elseif(EXECUTORCH_BUILD_METAL)
@@ -1101,7 +1126,7 @@ if(EXECUTORCH_BUILD_PYBIND)
  list(APPEND _dep_libs openvino_backend)
  endif()
 
- if(EXECUTORCH_BUILD_CUDA)
+ if(EXECUTORCH_BUILD_CUDA OR EXECUTORCH_BUILD_ROCM)
  string(APPEND _portable_lib_rpath ":$ORIGIN/../../backends/cuda")
  endif()
 

backends/aoti/CMakeLists.txt

@@ -83,8 +83,12 @@ target_compile_definitions(
  slimtensor INTERFACE C10_USING_CUSTOM_GENERATED_MACROS
 )
 
-# Add CUDA support for SlimTensor when building with CUDA
-if(EXECUTORCH_BUILD_CUDA)
+# Add GPU support for SlimTensor when building with CUDA or ROCm.
+if(EXECUTORCH_BUILD_ROCM)
+ find_package(hip CONFIG REQUIRED)
+ target_link_libraries(slimtensor INTERFACE hip::host extension_cuda)
+ target_compile_definitions(slimtensor INTERFACE EXECUTORCH_USE_HIP=1)
+elseif(EXECUTORCH_BUILD_CUDA)
  find_package(CUDAToolkit REQUIRED)
  target_include_directories(slimtensor INTERFACE ${CUDAToolkit_INCLUDE_DIRS})
  target_link_libraries(slimtensor INTERFACE CUDA::cudart extension_cuda)
@@ -102,7 +106,7 @@ install(
 # TODO(gasoonjia): Replace aoti_common with this one after metal migration
 # ==============================================================================
 set(_aoti_common_shims_slim_sources common_shims_slim.cpp)
-if(EXECUTORCH_BUILD_CUDA)
+if(EXECUTORCH_BUILD_CUDA OR EXECUTORCH_BUILD_ROCM)
  list(APPEND _aoti_common_shims_slim_sources slim/cuda/guard.cpp)
 endif()
 

backends/aoti/slim/c10/cuda/Exception.h

@@ -8,10 +8,8 @@
 
 #pragma once
 
-#include <cuda.h>
-#include <cuda_runtime.h>
-
 #include <executorch/backends/aoti/slim/c10/macros/Macros.h>
+#include <executorch/extension/cuda/runtime_api.h>
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/platform/assert.h>
 #include <executorch/runtime/platform/log.h>

backends/aoti/slim/c10/cuda/targets.bzl

@@ -11,6 +11,7 @@ def define_common_targets():
  visibility = ["@EXECUTORCH_CLIENTS"],
  exported_deps = [
  "//executorch/backends/aoti/slim/c10/macros:macros",
+ "//executorch/extension/cuda:runtime_api",
  "//executorch/runtime/platform:platform",
  ],
  )

backends/aoti/slim/cuda/guard.h

@@ -8,11 +8,11 @@
 
 #pragma once
 
-#include <cuda_runtime.h>
 #include <optional>
 
 #include <executorch/backends/aoti/slim/c10/core/Device.h>
 #include <executorch/backends/aoti/slim/c10/cuda/Exception.h>
+#include <executorch/extension/cuda/runtime_api.h>
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/core/result.h>
 

backends/aoti/slim/cuda/targets.bzl

@@ -18,6 +18,7 @@ def define_common_targets():
  exported_deps = [
  "//executorch/backends/aoti/slim/c10/core:device",
  "//executorch/backends/aoti/slim/c10/cuda:exception",
+ "//executorch/extension/cuda:runtime_api",
  "//executorch/runtime/core:core",
  "//executorch/runtime/core/exec_aten:lib",
  ],

backends/cuda/CMakeLists.txt

@@ -28,21 +28,32 @@ if(NOT EXECUTORCH_ROOT)
  set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../..)
 endif()
 
-# Use dynamic linking for CUDA runtime
+# Use dynamic linking for the GPU runtime.
 set(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
 
-find_package(CUDAToolkit REQUIRED)
+# Keep cuRAND confined to CUDA-only operator shims.
+if(EXECUTORCH_BUILD_ROCM)
+ find_package(hip CONFIG REQUIRED)
+ set(_cuda_runtime_libraries hip::host)
+ set(_cuda_shim_libraries hip::host)
+ set(_cuda_include_directories "")
+else()
+ find_package(CUDAToolkit REQUIRED)
+ set(_cuda_runtime_libraries CUDA::cudart)
+ set(_cuda_shim_libraries CUDA::cudart CUDA::curand)
+ set(_cuda_include_directories ${CUDAToolkit_INCLUDE_DIRS})
+endif()
 
 # Try to enable CUDA language when a working CUDA compiler toolchain is
 # available. Some CI environments (notably Windows packaging jobs) provide
 # CUDAToolkit headers/libs but cannot complete CUDA compiler identification. In
 # those cases, keep configuration working and skip CUDA-only sources below.
-if(NOT CMAKE_CUDA_COMPILER)
+if(NOT EXECUTORCH_BUILD_ROCM AND NOT CMAKE_CUDA_COMPILER)
  include(CheckLanguage)
  check_language(CUDA)
 endif()
 
-if(CMAKE_CUDA_COMPILER)
+if(NOT EXECUTORCH_BUILD_ROCM AND CMAKE_CUDA_COMPILER)
  enable_language(CUDA)
 endif()
 
@@ -107,8 +118,8 @@ set(_aoti_cuda_shim_sources runtime/cuda_allocator.cpp runtime/shims/memory.cpp
  runtime/shims/cuda_guard.cpp
 )
 
-# Only build CUDAshims when CUDA language/toolchain is available.
-if(CMAKE_CUDA_COMPILER)
+# Only build CUDA-only operator shims when the CUDA toolchain is available.
+if(NOT EXECUTORCH_BUILD_ROCM AND CMAKE_CUDA_COMPILER)
  list(
  APPEND
  _aoti_cuda_shim_sources
@@ -126,6 +137,9 @@ add_library(aoti_cuda_shims SHARED ${_aoti_cuda_shim_sources})
 
 # Define CUDA_AVAILABLE to use SlimTensor on GPU in common_shims_slim.h
 target_compile_definitions(aoti_cuda_shims PRIVATE CUDA_AVAILABLE=1)
+if(EXECUTORCH_BUILD_ROCM)
+ target_compile_definitions(aoti_cuda_shims PUBLIC EXECUTORCH_USE_HIP=1)
+endif()
 
 # Define export macros for shared library. Use WIN32 (not just MSVC) so MinGW
 # cross-compiles also emit dllexport symbols for AOTI shims.
@@ -141,7 +155,7 @@ endif()
 
 target_include_directories(
  aoti_cuda_shims
- PUBLIC ${CUDAToolkit_INCLUDE_DIRS} $<BUILD_INTERFACE:${EXECUTORCH_ROOT}>
+ PUBLIC ${_cuda_include_directories} $<BUILD_INTERFACE:${EXECUTORCH_ROOT}>
  $<INSTALL_INTERFACE:include>
 )
 
@@ -160,7 +174,7 @@ endif()
 # retention.
 if(_cuda_is_msvc_toolchain)
  target_link_libraries(
- aoti_cuda_shims PRIVATE cuda_platform CUDA::cudart CUDA::curand
+ aoti_cuda_shims PRIVATE cuda_platform ${_cuda_shim_libraries}
  extension_cuda ${CMAKE_DL_LIBS}
  )
  # Link object library directly so symbols are pulled exactly once while
@@ -170,13 +184,8 @@ else()
  target_link_libraries(
  aoti_cuda_shims
  PRIVATE cuda_platform
- PUBLIC -Wl,--whole-archive
- aoti_common_shims_slim
- -Wl,--no-whole-archive
- CUDA::cudart
- CUDA::curand
- extension_cuda
- ${CMAKE_DL_LIBS}
+ PUBLIC -Wl,--whole-archive aoti_common_shims_slim -Wl,--no-whole-archive
+ ${_cuda_shim_libraries} extension_cuda ${CMAKE_DL_LIBS}
  )
 endif()
 
@@ -204,9 +213,12 @@ add_library(aoti_cuda_backend STATIC ${_aoti_cuda_backend_sources})
 
 target_include_directories(
  aoti_cuda_backend
- PUBLIC ${CUDAToolkit_INCLUDE_DIRS} $<BUILD_INTERFACE:${EXECUTORCH_ROOT}>
+ PUBLIC ${_cuda_include_directories} $<BUILD_INTERFACE:${EXECUTORCH_ROOT}>
  $<INSTALL_INTERFACE:include>
 )
+if(EXECUTORCH_BUILD_ROCM)
+ target_compile_definitions(aoti_cuda_backend PUBLIC EXECUTORCH_USE_HIP=1)
+endif()
 target_compile_options(
  aoti_cuda_backend
  PUBLIC "$<$<COMPILE_LANGUAGE:CXX>:${_cuda_cxx_compile_options}>"
@@ -221,8 +233,9 @@ endif()
 # directory. On other platforms, use PUBLIC so the dependency propagates to
 # consumers.
 target_link_libraries(
- aoti_cuda_backend PUBLIC cuda_platform extension_tensor CUDA::cudart
- extension_cuda ${CMAKE_DL_LIBS}
+ aoti_cuda_backend
+ PUBLIC cuda_platform extension_tensor ${_cuda_runtime_libraries}
+ extension_cuda ${CMAKE_DL_LIBS}
 )
 
 if(_cuda_is_msvc_toolchain)

backends/cuda/cuda_backend.py

@@ -386,6 +386,9 @@ def save_data_externally(cls) -> bool:
 
  @classmethod
  def get_supported_fallback_kernels(cls) -> Dict[str, Any]:
+ # ROCm does not build the CUDA-only .cu fallback shims.
+ if torch.version.hip is not None:
+ return {}
  return {
  "at::_ops::_weight_int4pack_mm::call": None,
  "at::_ops::sort_stable::call": None,
@@ -400,6 +403,42 @@ def get_supported_fallback_kernels(cls) -> Dict[str, Any]:
  "aoti_torch_cuda_int8_plain_mm": None,
  }
 
+ @staticmethod
+ def _get_custom_ops_to_c_shim_options() -> Dict[str, Any]:
+ if torch.version.hip is not None:
+ return {}
+ try:
+ return {
+ "aot_inductor.custom_ops_to_c_shims": {
+ torch.ops.executorch_cuda.int4_plain_mm.default: [
+ "AOTITorchError aoti_torch_cuda_int4_plain_mm("
+ "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, "
+ "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, "
+ "int64_t, AtenTensorHandle*)"
+ ],
+ torch.ops.executorch_cuda.int5_plain_mm.default: [
+ "AOTITorchError aoti_torch_cuda_int5_plain_mm("
+ "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, "
+ "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, "
+ "AtenTensorHandle, int64_t, AtenTensorHandle*)"
+ ],
+ torch.ops.executorch_cuda.int6_plain_mm.default: [
+ "AOTITorchError aoti_torch_cuda_int6_plain_mm("
+ "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, "
+ "AtenTensorHandle, AtenTensorHandle, int64_t, "
+ "AtenTensorHandle*)"
+ ],
+ torch.ops.executorch_cuda.int8_plain_mm.default: [
+ "AOTITorchError aoti_torch_cuda_int8_plain_mm("
+ "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, "
+ "AtenTensorHandle, int64_t, AtenTensorHandle*)"
+ ],
+ }
+ }
+ except AttributeError:
+ # Custom ops may not be registered in this process.
+ return {}
+
  @classmethod
  def get_decomposition_table(cls) -> Dict[Any, Any]:
  return {
@@ -469,36 +508,7 @@ def get_aoti_compile_options(
  "aot_inductor.emit_multi_arch_kernel": emit_multi_arch_kernel,
  }
 
- try:
- import torch
-
- options["aot_inductor.custom_ops_to_c_shims"] = {
- torch.ops.executorch_cuda.int4_plain_mm.default: [
- "AOTITorchError aoti_torch_cuda_int4_plain_mm("
- "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, "
- "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, "
- "int64_t, AtenTensorHandle*)"
- ],
- torch.ops.executorch_cuda.int5_plain_mm.default: [
- "AOTITorchError aoti_torch_cuda_int5_plain_mm("
- "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, "
- "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, "
- "AtenTensorHandle, int64_t, AtenTensorHandle*)"
- ],
- torch.ops.executorch_cuda.int6_plain_mm.default: [
- "AOTITorchError aoti_torch_cuda_int6_plain_mm("
- "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, "
- "AtenTensorHandle, AtenTensorHandle, int64_t, AtenTensorHandle*)"
- ],
- torch.ops.executorch_cuda.int8_plain_mm.default: [
- "AOTITorchError aoti_torch_cuda_int8_plain_mm("
- "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, "
- "AtenTensorHandle, int64_t, AtenTensorHandle*)"
- ],
- }
- except AttributeError:
- # quantize_op_dispatch not imported — op not registered, skip C shim mapping
- pass
+ options.update(cls._get_custom_ops_to_c_shim_options())
 
  # Parse compile_specs to check for platform