From aede223f3b6576f56f710c5bba77d811807b2c4b Mon Sep 17 00:00:00 2001 From: Digant Desai Date: Fri, 14 Aug 2026 19:24:24 -0700 Subject: [PATCH 1/3] Add a CUDA/HIP runtime compatibility layer Route the AOTI CUDA runtime surface through a compatibility header that maps to HIP when EXECUTORCH_USE_HIP is defined. Keep the Buck header target separate from caller_stream so existing NVIDIA targets gain no additional linkage. With assistance from Claude Code and Codex. --- backends/aoti/slim/c10/cuda/Exception.h | 4 +- backends/aoti/slim/c10/cuda/targets.bzl | 1 + backends/aoti/slim/cuda/guard.h | 2 +- backends/aoti/slim/cuda/targets.bzl | 1 + backends/cuda/runtime/cuda_allocator.cpp | 3 +- backends/cuda/runtime/cuda_allocator.h | 3 +- backends/cuda/runtime/cuda_backend.cpp | 2 +- backends/cuda/runtime/cuda_delegate_handle.h | 2 +- backends/cuda/runtime/cuda_mutable_state.cpp | 3 +- backends/cuda/runtime/memory_tracker.h | 2 +- backends/cuda/runtime/shims/cuda_guard.h | 2 +- backends/cuda/runtime/shims/int4_plain_mm.h | 2 +- backends/cuda/runtime/shims/int4mm.h | 2 +- backends/cuda/runtime/shims/int5_plain_mm.h | 2 +- backends/cuda/runtime/shims/int6_plain_mm.h | 2 +- backends/cuda/runtime/shims/int8_plain_mm.h | 2 +- backends/cuda/runtime/shims/sort.h | 2 +- backends/cuda/runtime/targets.bzl | 1 + .../cuda/runtime/test/test_cuda_allocator.cpp | 2 +- .../runtime/test/test_cuda_mutable_state.cpp | 2 +- backends/cuda/runtime/utils.h | 2 +- extension/cuda/BUCK | 8 +- extension/cuda/caller_stream.h | 2 +- extension/cuda/runtime_api.h | 162 ++++++++++++++++++ extension/cuda/targets.bzl | 15 ++ 25 files changed, 203 insertions(+), 28 deletions(-) create mode 100644 extension/cuda/runtime_api.h diff --git a/backends/aoti/slim/c10/cuda/Exception.h b/backends/aoti/slim/c10/cuda/Exception.h index 4d2ba13d3bb..1b86b940559 100644 --- a/backends/aoti/slim/c10/cuda/Exception.h +++ b/backends/aoti/slim/c10/cuda/Exception.h @@ -8,10 +8,8 @@ #pragma once -#include -#include - #include +#include #include #include #include diff --git a/backends/aoti/slim/c10/cuda/targets.bzl b/backends/aoti/slim/c10/cuda/targets.bzl index 1d44bd1f032..e399f79804a 100644 --- a/backends/aoti/slim/c10/cuda/targets.bzl +++ b/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", ], ) diff --git a/backends/aoti/slim/cuda/guard.h b/backends/aoti/slim/cuda/guard.h index 31ea70705ac..f8a71873ff2 100644 --- a/backends/aoti/slim/cuda/guard.h +++ b/backends/aoti/slim/cuda/guard.h @@ -8,11 +8,11 @@ #pragma once -#include #include #include #include +#include #include #include diff --git a/backends/aoti/slim/cuda/targets.bzl b/backends/aoti/slim/cuda/targets.bzl index cddd69f1999..9585cd71635 100644 --- a/backends/aoti/slim/cuda/targets.bzl +++ b/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", ], diff --git a/backends/cuda/runtime/cuda_allocator.cpp b/backends/cuda/runtime/cuda_allocator.cpp index 0f1abfe2126..570fdf506e3 100644 --- a/backends/cuda/runtime/cuda_allocator.cpp +++ b/backends/cuda/runtime/cuda_allocator.cpp @@ -8,9 +8,8 @@ #include -#include - #include +#include #include namespace executorch::backends::cuda { diff --git a/backends/cuda/runtime/cuda_allocator.h b/backends/cuda/runtime/cuda_allocator.h index fcd8224305a..b0a76a51f6d 100644 --- a/backends/cuda/runtime/cuda_allocator.h +++ b/backends/cuda/runtime/cuda_allocator.h @@ -8,8 +8,7 @@ #pragma once -#include - +#include #include namespace executorch::backends::cuda { diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 8666fc9c098..baca5797b17 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -7,7 +7,7 @@ */ #include -#include +#include #include #include #include diff --git a/backends/cuda/runtime/cuda_delegate_handle.h b/backends/cuda/runtime/cuda_delegate_handle.h index ee360531c47..d8a46e64e89 100644 --- a/backends/cuda/runtime/cuda_delegate_handle.h +++ b/backends/cuda/runtime/cuda_delegate_handle.h @@ -8,8 +8,8 @@ #pragma once -#include #include +#include #include #include diff --git a/backends/cuda/runtime/cuda_mutable_state.cpp b/backends/cuda/runtime/cuda_mutable_state.cpp index dfaef56aabb..55894601a45 100644 --- a/backends/cuda/runtime/cuda_mutable_state.cpp +++ b/backends/cuda/runtime/cuda_mutable_state.cpp @@ -13,10 +13,9 @@ #include #include #include +#include #include -#include - #include #include #include diff --git a/backends/cuda/runtime/memory_tracker.h b/backends/cuda/runtime/memory_tracker.h index e09a96da6a6..085f35bbbf3 100644 --- a/backends/cuda/runtime/memory_tracker.h +++ b/backends/cuda/runtime/memory_tracker.h @@ -8,10 +8,10 @@ #pragma once -#include #include #include +#include #include namespace executorch::backends::cuda { diff --git a/backends/cuda/runtime/shims/cuda_guard.h b/backends/cuda/runtime/shims/cuda_guard.h index ca06564b423..c655991f978 100644 --- a/backends/cuda/runtime/shims/cuda_guard.h +++ b/backends/cuda/runtime/shims/cuda_guard.h @@ -8,9 +8,9 @@ #pragma once -#include #include #include +#include #include namespace executorch::backends::cuda { diff --git a/backends/cuda/runtime/shims/int4_plain_mm.h b/backends/cuda/runtime/shims/int4_plain_mm.h index 0178cac640b..aaee732a164 100644 --- a/backends/cuda/runtime/shims/int4_plain_mm.h +++ b/backends/cuda/runtime/shims/int4_plain_mm.h @@ -8,9 +8,9 @@ #pragma once -#include #include #include +#include namespace executorch::backends::cuda { diff --git a/backends/cuda/runtime/shims/int4mm.h b/backends/cuda/runtime/shims/int4mm.h index c353c09cad9..628cb5fd65f 100644 --- a/backends/cuda/runtime/shims/int4mm.h +++ b/backends/cuda/runtime/shims/int4mm.h @@ -8,9 +8,9 @@ #pragma once -#include #include #include +#include namespace executorch::backends::cuda { diff --git a/backends/cuda/runtime/shims/int5_plain_mm.h b/backends/cuda/runtime/shims/int5_plain_mm.h index 4b0a1219918..b79646ef9a9 100644 --- a/backends/cuda/runtime/shims/int5_plain_mm.h +++ b/backends/cuda/runtime/shims/int5_plain_mm.h @@ -8,9 +8,9 @@ #pragma once -#include #include #include +#include namespace executorch::backends::cuda { diff --git a/backends/cuda/runtime/shims/int6_plain_mm.h b/backends/cuda/runtime/shims/int6_plain_mm.h index 338f0980bde..26e7c42e974 100644 --- a/backends/cuda/runtime/shims/int6_plain_mm.h +++ b/backends/cuda/runtime/shims/int6_plain_mm.h @@ -8,9 +8,9 @@ #pragma once -#include #include #include +#include namespace executorch::backends::cuda { diff --git a/backends/cuda/runtime/shims/int8_plain_mm.h b/backends/cuda/runtime/shims/int8_plain_mm.h index c61e9f2ba8b..aea1c4c9b26 100644 --- a/backends/cuda/runtime/shims/int8_plain_mm.h +++ b/backends/cuda/runtime/shims/int8_plain_mm.h @@ -8,9 +8,9 @@ #pragma once -#include #include #include +#include namespace executorch::backends::cuda { diff --git a/backends/cuda/runtime/shims/sort.h b/backends/cuda/runtime/shims/sort.h index ffdc55ebe82..2c4686c3d5e 100644 --- a/backends/cuda/runtime/shims/sort.h +++ b/backends/cuda/runtime/shims/sort.h @@ -8,9 +8,9 @@ #pragma once -#include #include #include +#include namespace executorch::backends::cuda { diff --git a/backends/cuda/runtime/targets.bzl b/backends/cuda/runtime/targets.bzl index 94fa08f1c4d..cef4988536f 100644 --- a/backends/cuda/runtime/targets.bzl +++ b/backends/cuda/runtime/targets.bzl @@ -76,6 +76,7 @@ def define_common_targets(is_fbcode = False): "//executorch/backends/aoti/slim/factory:empty", "//executorch/backends/aoti/slim/factory:from_blob", "//executorch/backends/aoti/slim/cuda:guard", + "//executorch/extension/cuda:runtime_api", "//executorch/runtime/core:core", "//executorch/runtime/core/exec_aten:lib", "//executorch/runtime/core/exec_aten/util:tensor_util", diff --git a/backends/cuda/runtime/test/test_cuda_allocator.cpp b/backends/cuda/runtime/test/test_cuda_allocator.cpp index 9bbbed98ca8..baf9889df77 100644 --- a/backends/cuda/runtime/test/test_cuda_allocator.cpp +++ b/backends/cuda/runtime/test/test_cuda_allocator.cpp @@ -8,7 +8,7 @@ #include -#include +#include #include #include diff --git a/backends/cuda/runtime/test/test_cuda_mutable_state.cpp b/backends/cuda/runtime/test/test_cuda_mutable_state.cpp index 89ac02dda38..70419c61812 100644 --- a/backends/cuda/runtime/test/test_cuda_mutable_state.cpp +++ b/backends/cuda/runtime/test/test_cuda_mutable_state.cpp @@ -14,7 +14,7 @@ #include #include -#include +#include #include #include diff --git a/backends/cuda/runtime/utils.h b/backends/cuda/runtime/utils.h index aaed5108d4f..ad6a2e9ee5b 100644 --- a/backends/cuda/runtime/utils.h +++ b/backends/cuda/runtime/utils.h @@ -8,8 +8,8 @@ #pragma once -#include #include +#include #include #include #include diff --git a/extension/cuda/BUCK b/extension/cuda/BUCK index 351d2f7e31c..0b3f83765f0 100644 --- a/extension/cuda/BUCK +++ b/extension/cuda/BUCK @@ -1,10 +1,10 @@ # Any targets that should be shared between fbcode and xplat must be defined in # targets.bzl. This file can contain fbcode-only targets. # -# caller_stream unconditionally includes , whose headers only -# resolve through the fbcode cell's CUDA third-party. Like the other executorch -# CUDA targets, it must therefore be defined fbcode-only; defining it in the -# xplat cell produces a target that cannot compile ("cuda_runtime.h not found"). +# caller_stream uses the CUDA Runtime API by default, whose headers only resolve +# through the fbcode cell's CUDA third-party. Like the other executorch CUDA +# targets, it must therefore be defined fbcode-only; defining it in the xplat +# cell produces a target that cannot compile ("cuda_runtime.h not found"). load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") load(":targets.bzl", "define_common_targets") diff --git a/extension/cuda/caller_stream.h b/extension/cuda/caller_stream.h index a13b7a9b396..7de1a310404 100644 --- a/extension/cuda/caller_stream.h +++ b/extension/cuda/caller_stream.h @@ -8,11 +8,11 @@ #pragma once -#include #include #include #include +#include namespace executorch::extension::cuda { diff --git a/extension/cuda/runtime_api.h b/extension/cuda/runtime_api.h new file mode 100644 index 00000000000..bae5c6a79bf --- /dev/null +++ b/extension/cuda/runtime_api.h @@ -0,0 +1,162 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +#if defined(EXECUTORCH_USE_HIP) + +#include + +using cudaError_t = hipError_t; +using cudaGraph_t = hipGraph_t; +using cudaGraphExec_t = hipGraphExec_t; +using cudaMemcpyKind = hipMemcpyKind; +using cudaMemoryType = hipMemoryType; +using cudaStreamCaptureMode = hipStreamCaptureMode; +using cudaStream_t = hipStream_t; + +inline constexpr cudaError_t cudaSuccess = hipSuccess; +inline constexpr cudaMemcpyKind cudaMemcpyDeviceToDevice = + hipMemcpyDeviceToDevice; +inline constexpr cudaMemcpyKind cudaMemcpyDeviceToHost = hipMemcpyDeviceToHost; +inline constexpr cudaMemcpyKind cudaMemcpyHostToDevice = hipMemcpyHostToDevice; +inline constexpr cudaMemoryType cudaMemoryTypeDevice = hipMemoryTypeDevice; +inline constexpr cudaMemoryType cudaMemoryTypeManaged = hipMemoryTypeManaged; +inline constexpr cudaStreamCaptureMode cudaStreamCaptureModeRelaxed = + hipStreamCaptureModeRelaxed; +inline constexpr unsigned long long cudaGraphInstantiateFlagAutoFreeOnLaunch = + hipGraphInstantiateFlagAutoFreeOnLaunch; + +struct cudaPointerAttributes { + cudaMemoryType type{}; + int device = -1; +}; + +inline cudaError_t cudaDeviceSynchronize() { + return hipDeviceSynchronize(); +} + +inline cudaError_t cudaFree(void* ptr) { + return hipFree(ptr); +} + +inline cudaError_t cudaFreeAsync(void* ptr, cudaStream_t stream) { + return hipFreeAsync(ptr, stream); +} + +inline cudaError_t cudaGetDevice(int* device) { + return hipGetDevice(device); +} + +inline cudaError_t cudaGetDeviceCount(int* count) { + return hipGetDeviceCount(count); +} + +inline const char* cudaGetErrorString(cudaError_t error) { + return hipGetErrorString(error); +} + +inline cudaError_t cudaGetLastError() { + return hipGetLastError(); +} + +inline cudaError_t cudaGraphDestroy(cudaGraph_t graph) { + return hipGraphDestroy(graph); +} + +inline cudaError_t cudaGraphExecDestroy(cudaGraphExec_t graph_exec) { + return hipGraphExecDestroy(graph_exec); +} + +inline cudaError_t cudaGraphInstantiate( + cudaGraphExec_t* graph_exec, + cudaGraph_t graph, + unsigned long long flags) { + return hipGraphInstantiateWithFlags(graph_exec, graph, flags); +} + +inline cudaError_t cudaGraphLaunch( + cudaGraphExec_t graph_exec, + cudaStream_t stream) { + return hipGraphLaunch(graph_exec, stream); +} + +inline cudaError_t cudaMalloc(void** ptr, size_t size) { + return hipMalloc(ptr, size); +} + +inline cudaError_t +cudaMallocAsync(void** ptr, size_t size, cudaStream_t stream) { + return hipMallocAsync(ptr, size, stream); +} + +inline cudaError_t +cudaMemcpy(void* dst, const void* src, size_t size, cudaMemcpyKind kind) { + return hipMemcpy(dst, src, size, kind); +} + +inline cudaError_t cudaMemcpyAsync( + void* dst, + const void* src, + size_t size, + cudaMemcpyKind kind, + cudaStream_t stream) { + return hipMemcpyAsync(dst, src, size, kind, stream); +} + +inline cudaError_t cudaMemGetInfo(size_t* free, size_t* total) { + return hipMemGetInfo(free, total); +} + +inline cudaError_t cudaPointerGetAttributes( + cudaPointerAttributes* attributes, + const void* ptr) { + hipPointerAttribute_t hip_attributes{}; + const auto error = hipPointerGetAttributes(&hip_attributes, ptr); + if (error == hipSuccess) { + attributes->type = hip_attributes.type; + attributes->device = hip_attributes.device; + } + return error; +} + +inline cudaError_t cudaSetDevice(int device) { + return hipSetDevice(device); +} + +inline cudaError_t cudaStreamBeginCapture( + cudaStream_t stream, + cudaStreamCaptureMode mode) { + return hipStreamBeginCapture(stream, mode); +} + +inline cudaError_t cudaStreamCreate(cudaStream_t* stream) { + return hipStreamCreate(stream); +} + +inline cudaError_t cudaStreamDestroy(cudaStream_t stream) { + return hipStreamDestroy(stream); +} + +inline cudaError_t cudaStreamEndCapture( + cudaStream_t stream, + cudaGraph_t* graph) { + return hipStreamEndCapture(stream, graph); +} + +inline cudaError_t cudaStreamSynchronize(cudaStream_t stream) { + return hipStreamSynchronize(stream); +} + +#else + +#include + +#endif diff --git a/extension/cuda/targets.bzl b/extension/cuda/targets.bzl index 6152b9d4835..a3b06c3d47d 100644 --- a/extension/cuda/targets.bzl +++ b/extension/cuda/targets.bzl @@ -7,6 +7,18 @@ def define_common_targets(): TARGETS and BUCK files that call this function. """ + # Keep runtime types independent of the caller-stream implementation. + runtime.cxx_library( + name = "runtime_api", + exported_headers = [ + "runtime_api.h", + ], + visibility = ["PUBLIC"], + external_deps = [ + ("cuda", None, "cuda-lazy"), + ], + ) + # Backend-neutral: both the CUDA and TensorRT delegates can depend on it to # share a caller's stream. The caller-stream thread-local must be one # instance per process, so the main target stays shareable: OSS cxx_library @@ -22,6 +34,9 @@ def define_common_targets(): "caller_stream.h", "export.h", ], + exported_deps = [ + ":runtime_api", + ], # Opt out of the OSS force_static default so consumers *can* link one # shared instance and keep the thread-local unique (see above); the # wrapper pins preferred_linkage="any", so this allows shared linkage From 3723dc05df9dff22346958a86687d956769f6e75 Mon Sep 17 00:00:00 2001 From: Digant Desai Date: Fri, 14 Aug 2026 19:24:54 -0700 Subject: [PATCH 2/3] Add opt-in ROCm support to the AOTI CUDA backend Add EXECUTORCH_BUILD_ROCM as an off-by-default alternative to CUDA. Link HIP only in this mode, omit CUDA-only fallbacks, reject unsupported graph capture, and overwrite installed backend dependency metadata on every install. Cover the CUDA/ROCm export gates and document the experimental gfx942 scope and limitations. With assistance from Claude Code and Codex. --- CMakeLists.txt | 41 ++++++++--- backends/aoti/CMakeLists.txt | 10 ++- backends/cuda/CMakeLists.txt | 49 ++++++++----- backends/cuda/cuda_backend.py | 70 +++++++++++-------- backends/cuda/rocm.md | 40 +++++++++++ backends/cuda/runtime/cuda_backend.cpp | 10 +++ backends/cuda/tests/test_sort_shim.py | 48 ++++++++++++- extension/cuda/CMakeLists.txt | 13 +++- extension/llm/runner/CMakeLists.txt | 11 +-- .../executorch-backend-dependencies.cmake.in | 9 +++ tools/cmake/executorch-config.cmake | 5 ++ tools/cmake/preset/default.cmake | 11 +++ 12 files changed, 249 insertions(+), 68 deletions(-) create mode 100644 backends/cuda/rocm.md create mode 100644 tools/cmake/executorch-backend-dependencies.cmake.in diff --git a/CMakeLists.txt b/CMakeLists.txt index ff3b9e86f7e..56be6886688 100644 --- a/CMakeLists.txt +++ b/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 CUDA or 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() diff --git a/backends/aoti/CMakeLists.txt b/backends/aoti/CMakeLists.txt index 4634f36eb9d..5ba98c67f20 100644 --- a/backends/aoti/CMakeLists.txt +++ b/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() diff --git a/backends/cuda/CMakeLists.txt b/backends/cuda/CMakeLists.txt index 06990692428..c6d318d1a25 100644 --- a/backends/cuda/CMakeLists.txt +++ b/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 CUDA shims 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} $ + PUBLIC ${_cuda_include_directories} $ $ ) @@ -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} $ + PUBLIC ${_cuda_include_directories} $ $ ) +if(EXECUTORCH_BUILD_ROCM) + target_compile_definitions(aoti_cuda_backend PUBLIC EXECUTORCH_USE_HIP=1) +endif() target_compile_options( aoti_cuda_backend PUBLIC "$<$:${_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) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index bf28077c62c..a40cb102e83 100644 --- a/backends/cuda/cuda_backend.py +++ b/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 diff --git a/backends/cuda/rocm.md b/backends/cuda/rocm.md new file mode 100644 index 00000000000..f00fd287f5e --- /dev/null +++ b/backends/cuda/rocm.md @@ -0,0 +1,40 @@ +# AMD ROCm support (experimental) + +The AOTI CUDA backend can be built against AMD ROCm/HIP. PyTorch exposes AMD +GPUs as device type `cuda`, so names such as `CudaBackend` and +`aoti_cuda_blob.ptd` remain unchanged. + +`EXECUTORCH_BUILD_ROCM` is off by default, is never auto-enabled, and is +mutually exclusive with `EXECUTORCH_BUILD_CUDA`. It requires +`EXECUTORCH_BUILD_EXTENSION_TENSOR`. Execution has been validated only on +MI300X (`gfx942`). + +## Requirements and limitations + +- Supply a ROCm PyTorch build and matching Triton AMD backend. To avoid replacing + ROCm PyTorch with a CPU wheel, install ExecuTorch with + `pip install -e . --no-build-isolation`. +- CUDA graphs and CUDA-only fallback shims are unavailable. +- Python pybindings cannot allocate ROCm device memory; use a native runner such + as `executor_runner`. +- Installed CMake consumers must be able to find the HIP package. +- Model runners do not yet link the ROCm backend, and there is no ROCm CI. + +## Build + +```bash +cmake -S . -B cmake-out-rocm \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_PREFIX_PATH="$(python -c 'import torch; print(torch.utils.cmake_prefix_path)');/opt/rocm" \ + -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \ + -DEXECUTORCH_BUILD_ROCM=ON +cmake --build cmake-out-rocm --target aoti_cuda_backend aoti_cuda_shims -j +``` + +Verify that the build uses HIP and not the NVIDIA runtime: + +```bash +ldd cmake-out-rocm/backends/cuda/libaoti_cuda_shims.so | grep -E 'amdhip64|cudart' +``` + +A `libcudart` dependency is a configuration error. diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index baca5797b17..a828201845f 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -277,6 +277,15 @@ class ET_EXPERIMENTAL CudaBackend final return Error::InvalidArgument; } } else if (std::strcmp(option.key, kEnableCudaGraphForMethod) == 0) { +#if defined(EXECUTORCH_USE_HIP) + // HIP ignores the graph-instantiation flag required by this path. + ET_LOG( + Error, + "Option %s is not supported on ROCm: HIP ignores graph " + "instantiation flags.", + kEnableCudaGraphForMethod); + return Error::NotSupported; +#else if (auto* val = std::get_if>( &option.value)) { set_cuda_graph_method(*val); @@ -287,6 +296,7 @@ class ET_EXPERIMENTAL CudaBackend final kEnableCudaGraphForMethod); return Error::InvalidArgument; } +#endif } } return Error::Ok; diff --git a/backends/cuda/tests/test_sort_shim.py b/backends/cuda/tests/test_sort_shim.py index ce95ed326e3..5ecfd520bf7 100644 --- a/backends/cuda/tests/test_sort_shim.py +++ b/backends/cuda/tests/test_sort_shim.py @@ -18,10 +18,12 @@ import os import tempfile import unittest +from unittest.mock import patch +# Register custom ops so the shim-map assertions exercise the ROCm gate. +import executorch.backends.cuda.quantize_op_dispatch # noqa: F401 import torch import torch.nn as nn - from executorch.backends.cuda.cuda_backend import CudaBackend from executorch.backends.cuda.cuda_partitioner import CudaPartitioner from executorch.exir import ( @@ -32,6 +34,22 @@ from executorch.exir.passes import MemoryPlanningPass from torch.export import export +_CUDA_FALLBACK_KERNELS = frozenset( + { + "at::_ops::_weight_int4pack_mm::call", + "at::_ops::sort_stable::call", + "aoti_torch_cuda_randint_low_out", + "executorch_cuda::int4_plain_mm", + "aoti_torch_cuda_int4_plain_mm", + "executorch_cuda::int5_plain_mm", + "aoti_torch_cuda_int5_plain_mm", + "executorch_cuda::int6_plain_mm", + "aoti_torch_cuda_int6_plain_mm", + "executorch_cuda::int8_plain_mm", + "aoti_torch_cuda_int8_plain_mm", + } +) + class SortModel(nn.Module): """Model that uses sort (via argsort) for export testing.""" @@ -119,10 +137,34 @@ def test_sort_stable_export(self): self.assertGreater(os.path.getsize(pte_path), 0) def test_sort_fallback_registered(self): - """sort_stable is registered as a supported fallback kernel.""" - fallbacks = CudaBackend.get_supported_fallback_kernels() + """sort_stable is registered as a supported fallback kernel on CUDA.""" + with patch.object(torch.version, "hip", None): + fallbacks = CudaBackend.get_supported_fallback_kernels() self.assertIn("at::_ops::sort_stable::call", fallbacks) + def test_cuda_fallbacks_unchanged_by_rocm_gate(self): + """The ROCm gate must not alter what CUDA advertises.""" + with patch.object(torch.version, "hip", None): + fallbacks = CudaBackend.get_supported_fallback_kernels() + + self.assertEqual(set(fallbacks), _CUDA_FALLBACK_KERNELS) + + def test_cuda_shim_map_unchanged_by_rocm_gate(self): + """Same, for the C shim signatures.""" + with patch.object(torch.version, "hip", None): + options = CudaBackend.get_aoti_compile_options([]) + + self.assertEqual(len(options["aot_inductor.custom_ops_to_c_shims"]), 4) + + def test_rocm_advertises_no_unbuilt_shims(self): + """ROCm does not advertise CUDA-only shims.""" + with patch.object(torch.version, "hip", "7.0"): + fallbacks = CudaBackend.get_supported_fallback_kernels() + options = CudaBackend.get_aoti_compile_options([]) + + self.assertEqual(fallbacks, {}) + self.assertNotIn("aot_inductor.custom_ops_to_c_shims", options) + if __name__ == "__main__": unittest.main() diff --git a/extension/cuda/CMakeLists.txt b/extension/cuda/CMakeLists.txt index 0003691ac8b..d520d1a80f1 100644 --- a/extension/cuda/CMakeLists.txt +++ b/extension/cuda/CMakeLists.txt @@ -16,14 +16,20 @@ if(NOT EXECUTORCH_ROOT) set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../..) endif() -find_package(CUDAToolkit REQUIRED) +if(EXECUTORCH_BUILD_ROCM) + find_package(hip CONFIG REQUIRED) + set(_cuda_runtime_library hip::host) +else() + find_package(CUDAToolkit REQUIRED) + set(_cuda_runtime_library CUDA::cudart) +endif() # SHARED on purpose: the caller-stream thread-local must have a single # definition across every shared object in the process (see export.h). A static # copy linked into multiple shared libraries would create multiple thread-locals # and silently break the caller-stream handshake. add_library(extension_cuda SHARED caller_stream.cpp) -target_link_libraries(extension_cuda PUBLIC CUDA::cudart) +target_link_libraries(extension_cuda PUBLIC ${_cuda_runtime_library}) target_include_directories(extension_cuda PUBLIC ${_common_include_directories}) target_compile_options( extension_cuda PUBLIC "$<$:${_common_compile_options}>" @@ -31,6 +37,9 @@ target_compile_options( target_compile_definitions( extension_cuda PRIVATE EXECUTORCH_EXTENSION_CUDA_BUILDING ) +if(EXECUTORCH_BUILD_ROCM) + target_compile_definitions(extension_cuda PUBLIC EXECUTORCH_USE_HIP=1) +endif() install( TARGETS extension_cuda diff --git a/extension/llm/runner/CMakeLists.txt b/extension/llm/runner/CMakeLists.txt index 5247a4ba0a6..b5e4bced7ca 100644 --- a/extension/llm/runner/CMakeLists.txt +++ b/extension/llm/runner/CMakeLists.txt @@ -68,10 +68,13 @@ target_include_directories( extension_llm_runner INTERFACE ${_common_include_directories} ) -# If the project is configured to build with CUDA support, try to find a CUDA -# runtime (prefer the CUDAToolkit package). If found, expose a compile-time -# macro so sources can conditionally compile CUDA-aware code. -if(EXECUTORCH_BUILD_CUDA) +# Link the configured GPU runtime for CUDA-aware runner sources. +if(EXECUTORCH_BUILD_ROCM) + find_package(hip CONFIG REQUIRED) + target_compile_definitions(extension_llm_runner PUBLIC CUDA_AVAILABLE) + target_link_libraries(extension_llm_runner PUBLIC hip::host extension_cuda) + message(STATUS "HIP found; defining CUDA_AVAILABLE") +elseif(EXECUTORCH_BUILD_CUDA) # Prefer the modern CMake CUDAToolkit module, fall back to searching for the # CUDA runtime library (cudart) if the package isn't available. find_package(CUDAToolkit QUIET) diff --git a/tools/cmake/executorch-backend-dependencies.cmake.in b/tools/cmake/executorch-backend-dependencies.cmake.in new file mode 100644 index 00000000000..03125695147 --- /dev/null +++ b/tools/cmake/executorch-backend-dependencies.cmake.in @@ -0,0 +1,9 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +# +# Generated from the enabled backends' exported package dependencies. + +@_executorch_backend_dependencies@ diff --git a/tools/cmake/executorch-config.cmake b/tools/cmake/executorch-config.cmake index 946436e3083..4c3a6f79793 100644 --- a/tools/cmake/executorch-config.cmake +++ b/tools/cmake/executorch-config.cmake @@ -28,6 +28,11 @@ cmake_minimum_required(VERSION 3.24) include(CMakeFindDependencyMacro) find_package(tokenizers CONFIG) +# Load dependencies published by enabled backend targets. +include("${CMAKE_CURRENT_LIST_DIR}/executorch-backend-dependencies.cmake" + OPTIONAL +) + set(_root "${CMAKE_CURRENT_LIST_DIR}/../../..") set(required_lib_list executorch executorch_core portable_kernels) set(EXECUTORCH_LIBRARIES) diff --git a/tools/cmake/preset/default.cmake b/tools/cmake/preset/default.cmake index ae5437ea443..ca85f13234e 100644 --- a/tools/cmake/preset/default.cmake +++ b/tools/cmake/preset/default.cmake @@ -191,6 +191,9 @@ define_overridable_option( define_overridable_option( EXECUTORCH_BUILD_CUDA "Build the CUDA backend" BOOL OFF ) +define_overridable_option( + EXECUTORCH_BUILD_ROCM "Build the CUDA/AOTI backend against ROCm" BOOL OFF +) define_overridable_option( EXECUTORCH_BUILD_METAL "Build the Metal backend" BOOL OFF ) @@ -463,6 +466,14 @@ check_required_options_on( IF_ON EXECUTORCH_BUILD_CUDA REQUIRES EXECUTORCH_BUILD_EXTENSION_TENSOR ) +check_required_options_on( + IF_ON EXECUTORCH_BUILD_ROCM REQUIRES EXECUTORCH_BUILD_EXTENSION_TENSOR +) + +check_conflicting_options_on( + IF_ON EXECUTORCH_BUILD_ROCM CONFLICTS_WITH EXECUTORCH_BUILD_CUDA +) + check_required_options_on( IF_ON EXECUTORCH_BUILD_METAL REQUIRES EXECUTORCH_BUILD_EXTENSION_TENSOR ) From ba801c586cba29ba8d0faae9eb20d6203821a959 Mon Sep 17 00:00:00 2001 From: Digant Desai Date: Fri, 14 Aug 2026 19:25:13 -0700 Subject: [PATCH 3/3] Add an AMD ROCm AOTI pointwise example Export a fused pointwise module on the active AMD device and require generated Triton source plus a matching embedded code object. With assistance from Claude Code and Codex. --- backends/cuda/rocm.md | 3 + examples/cuda/README.md | 37 ++++++ examples/cuda/scripts/export_amd_pointwise.py | 125 ++++++++++++++++++ 3 files changed, 165 insertions(+) create mode 100644 examples/cuda/README.md create mode 100644 examples/cuda/scripts/export_amd_pointwise.py diff --git a/backends/cuda/rocm.md b/backends/cuda/rocm.md index f00fd287f5e..2344e8ef99e 100644 --- a/backends/cuda/rocm.md +++ b/backends/cuda/rocm.md @@ -38,3 +38,6 @@ ldd cmake-out-rocm/backends/cuda/libaoti_cuda_shims.so | grep -E 'amdhip64|cudar ``` A `libcudart` dependency is a configuration error. + +For an end-to-end export that proves Inductor emitted a Triton kernel, see +[../../examples/cuda/README.md](../../examples/cuda/README.md). diff --git a/examples/cuda/README.md b/examples/cuda/README.md new file mode 100644 index 00000000000..a5421edb035 --- /dev/null +++ b/examples/cuda/README.md @@ -0,0 +1,37 @@ +# AMD AOTI pointwise example + +Exports a small fused pointwise module through the AOTI delegate on an AMD GPU +and proves that PyTorch Inductor emitted a real Triton kernel for it rather than +falling back to ATen library kernels. + +ROCm support is **experimental and off by default**. See +[../../backends/cuda/rocm.md](../../backends/cuda/rocm.md) for what is and is not +covered, including which architectures have execution coverage and which +capabilities are unavailable on ROCm. + +Build the runtime pieces on a machine with the ROCm SDK installed; an AMD GPU is +not required for this build: + +```bash +cmake -S . -B cmake-out-rocm \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_PREFIX_PATH=/opt/rocm \ + -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \ + -DEXECUTORCH_BUILD_ROCM=ON +cmake --build cmake-out-rocm --target aoti_cuda_backend aoti_cuda_shims -j +``` + +Export requires a ROCm-enabled PyTorch build, Triton with its AMD backend, the +ROCm SDK, and a visible AMD GPU: + +```bash +python -m examples.cuda.scripts.export_amd_pointwise --output-dir amd-out +``` + +Inductor targets the visible device, so the example compiles for whatever GPU is +present; there is no cross-compilation knob. It rejects any architecture the +installed PyTorch build does not list in `torch.cuda.get_arch_list()`. + +The example emits `amd_triton.pte` and `aoti_cuda_blob.ptd`. It uses a fresh +Inductor cache and fails unless it finds generated Triton source there, and it +checks that the `.pte` embeds a code object for the architecture it compiled for. diff --git a/examples/cuda/scripts/export_amd_pointwise.py b/examples/cuda/scripts/export_amd_pointwise.py new file mode 100644 index 00000000000..67c07da7a0e --- /dev/null +++ b/examples/cuda/scripts/export_amd_pointwise.py @@ -0,0 +1,125 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import argparse +import os +import pathlib +import tempfile + +import torch + + +class AmdPointwiseModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.sin(x) * torch.sigmoid(x + 0.5) + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Export a small Triton-backed AOTI delegate for an AMD GPU." + ) + parser.add_argument("--output-dir", type=pathlib.Path, default=pathlib.Path(".")) + args = parser.parse_args() + + if torch.version.hip is None: + raise RuntimeError( + "export_amd_pointwise.py requires a ROCm-enabled PyTorch build" + ) + if not torch.cuda.is_available(): + raise RuntimeError( + "AOTInductor ROCm lowering requires a visible AMD GPU to select and " + "compile the initial Triton kernel; C++ execution is not required" + ) + + # Match AOTInductor's active device. + device_index = torch.cuda.current_device() + device = torch.cuda.get_device_properties(device_index) + arch = device.gcnArchName.split(":", 1)[0] + # Use the targets supported by the installed PyTorch toolchain. + supported = torch.cuda.get_arch_list() + if arch not in supported: + raise ValueError( + f"{arch} is not in this PyTorch build's supported targets: " + f"{', '.join(supported)}" + ) + + args.output_dir.mkdir(parents=True, exist_ok=True) + cache_dir = pathlib.Path( + tempfile.mkdtemp(prefix="torchinductor_amd_", dir=args.output_dir) + ) + os.environ["TORCHINDUCTOR_CACHE_DIR"] = str(cache_dir) + + from executorch.backends.cuda.cuda_backend import CudaBackend + from executorch.backends.cuda.cuda_partitioner import CudaPartitioner + from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower + from executorch.exir.backend.compile_spec_schema import CompileSpec + + module = AmdPointwiseModule().eval() + example_inputs = (torch.randn(1024, 1024),) + exported_program = torch.export.export(module, example_inputs, strict=True) + + compile_specs = [ + CudaBackend.generate_method_name_compile_spec("forward"), + CompileSpec("target_device", f"cuda:{device_index}".encode()), + CompileSpec("triton_kernel_mode", b"ON"), + CompileSpec("max_autotune", b"OFF"), + CompileSpec("autotune_at_compile_time", b"OFF"), + ] + edge_program = to_edge_transform_and_lower( + exported_program, + partitioner=[CudaPartitioner(compile_specs)], + compile_config=EdgeCompileConfig( + _check_ir_validity=False, + _skip_dim_order=True, + ), + ) + + delegated = any( + node.op == "call_function" and "executorch_call_delegate" in str(node.target) + for node in edge_program.exported_program().graph.nodes + ) + if not delegated: + raise RuntimeError( + "the exported program does not contain an AOTI delegate call" + ) + + triton_sources = [] + for source_path in cache_dir.rglob("*.py"): + source = source_path.read_text(errors="ignore") + if "async_compile.triton(" in source or "@triton_heuristics." in source: + triton_sources.append(source_path) + if not triton_sources: + raise RuntimeError( + f"Inductor did not generate a Triton kernel under {cache_dir}" + ) + + executorch_program = edge_program.to_executorch() + pte_path = args.output_dir / "amd_triton.pte" + with pte_path.open("wb") as output_file: + executorch_program.write_to_file(output_file) + executorch_program.write_tensor_data_to_file(args.output_dir) + + ptd_path = args.output_dir / "aoti_cuda_blob.ptd" + if not ptd_path.is_file(): + raise RuntimeError(f"expected AOTI data file was not created: {ptd_path}") + + # Verify the embedded code object matches the active device. + triple = f"amdgcn-amd-amdhsa--{arch}".encode() + if triple not in pte_path.read_bytes(): + raise RuntimeError(f"{pte_path} embeds no code object for {arch}") + + print(f"ROCm: {torch.version.hip}") + print(f"Compile GPU: cuda:{device_index} {device.name} ({device.gcnArchName})") + print(f"Compile architecture: {arch}") + print(f"ExecuTorch program: {pte_path}") + print(f"AOTI data: {ptd_path}") + print("Triton sources:") + for source_path in triton_sources: + print(f" {source_path}") + + +if __name__ == "__main__": + main()