From 1832491c3e2b02ba91b690ce5c6a5113f6ef3932 Mon Sep 17 00:00:00 2001 From: shoumikhin Date: Fri, 23 Oct 2026 19:06:39 -0700 Subject: [PATCH 1/2] CUDA backend: add cuBLAS shims for mm, addmm and bmm The CUDA backend runs AOTInductor models without libtorch, so every extern call the generated code makes must be provided by the backend's shim library. Matrix multiplications had no shim, which is why the backend only lets Inductor use Triton for them: a model that picked the ATen (cuBLAS) choice failed at export with a missing fallback kernel. This adds aoti_torch_cuda_mm_out, aoti_torch_cuda_addmm_out and aoti_torch_cuda_bmm_out on top of cuBLASLt: - Row-major and column-major operands, sliced leading dimensions, dimensions of size 1 with any stride, and batches that repeat one matrix (zero batch stride). - addmm runs a bias row through the cuBLASLt bias epilogue, as PyTorch does for a linear layer. Any other self that broadcasts to the result (a column, a scalar, a transposed matrix, or a bias row with beta != 1) is copied into the result first and accumulated there with beta. - Each call runs on the current CUDA stream, and the workspace comes from the backend's stream-ordered allocator, so a call can be captured into a CUDA graph. Handles are per thread and device. - float32 runs in full IEEE precision, PyTorch's default. The shim library does not link cuBLASLt. It loads the library the first time one of these shims runs, so models that use only Triton kernels, and apps that ship only the CUDA runtime, still load the backend without it. If the library cannot be loaded, the shim returns an error that names it. The wheel already treats the cuBLAS libraries as provided by torch. Windows is unchanged: the shims load cuBLASLt with dlopen and are not ported to Windows yet, so they and their test are not built there, and the Windows import library does not list them (its notes now say they must not). A Windows model that calls one fails to link at export, as it did before. --- backends/cuda/CMakeLists.txt | 16 + backends/cuda/cuda_backend.py | 3 + backends/cuda/runtime/aoti_cuda_shims.lib.md | 3 +- backends/cuda/runtime/shims/mm.cu | 678 ++++++++++++++++++ backends/cuda/runtime/shims/mm.h | 73 ++ .../cuda/runtime/shims/tests/CMakeLists.txt | 4 + backends/cuda/runtime/shims/tests/targets.bzl | 1 + .../tests/test_aoti_torch_cuda_mm_out.cpp | 511 +++++++++++++ backends/cuda/runtime/targets.bzl | 3 + backends/cuda/tests/test_sort_shim.py | 3 + 10 files changed, 1294 insertions(+), 1 deletion(-) create mode 100644 backends/cuda/runtime/shims/mm.cu create mode 100644 backends/cuda/runtime/shims/mm.h create mode 100644 backends/cuda/runtime/shims/tests/test_aoti_torch_cuda_mm_out.cpp diff --git a/backends/cuda/CMakeLists.txt b/backends/cuda/CMakeLists.txt index 5f2879a7b32..bd3bd1d7e81 100644 --- a/backends/cuda/CMakeLists.txt +++ b/backends/cuda/CMakeLists.txt @@ -254,10 +254,17 @@ set(_aoti_cuda_shim_sources runtime/shims/memory.cpp ) # Only build CUDA-only operator shims when the CUDA toolchain is available. +set(_aoti_cuda_shims_use_cublas OFF) if(NOT EXECUTORCH_BUILD_ROCM AND CMAKE_CUDA_COMPILER) list(APPEND _aoti_cuda_shim_sources runtime/shims/int4mm.cu runtime/shims/sort.cu runtime/shims/rand.cu ) + # The cuBLAS matrix multiplication shims load cuBLASLt with dlopen and are + # not ported to Windows yet, so aoti_cuda_shims.lib must not list them. + if(NOT WIN32) + list(APPEND _aoti_cuda_shim_sources runtime/shims/mm.cu) + set(_aoti_cuda_shims_use_cublas ON) + endif() endif() add_library(aoti_cuda_shims SHARED ${_aoti_cuda_shim_sources}) @@ -322,6 +329,15 @@ else() ) endif() +# cuBLASLt is loaded on the first matrix multiplication rather than linked, so +# models that never call the cuBLAS shims do not need the library installed. +if(_aoti_cuda_shims_use_cublas) + target_compile_definitions( + aoti_cuda_shims + PRIVATE ET_CUBLASLT_LIBRARY="libcublasLt.so.${CUDAToolkit_VERSION_MAJOR}" + ) +endif() + if(NOT _cuda_is_msvc_toolchain) executorch_target_link_options_shared_lib(aoti_cuda_shims) endif() diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 787d2502b63..01f5df8cda1 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -773,6 +773,9 @@ def get_supported_fallback_kernels(cls) -> Dict[str, Any]: "at::_ops::sort_stable::call": None, "aoti_torch_cuda_sort_stable": None, "aoti_torch_cuda_randint_low_out": None, + "aoti_torch_cuda_mm_out": None, + "aoti_torch_cuda_addmm_out": None, + "aoti_torch_cuda_bmm_out": None, } @classmethod diff --git a/backends/cuda/runtime/aoti_cuda_shims.lib.md b/backends/cuda/runtime/aoti_cuda_shims.lib.md index af2fb65ca13..dce134b8657 100644 --- a/backends/cuda/runtime/aoti_cuda_shims.lib.md +++ b/backends/cuda/runtime/aoti_cuda_shims.lib.md @@ -33,4 +33,5 @@ the same source differ. A name only resolves if something in the DLL defines it. The DLL is built from the CUDA shims plus the SlimTensor common shims, so a shim added to the ETensor common -shims will link and then fail to load. +shims will link and then fail to load. The cuBLAS shims in `shims/mm.cu` are not +built for Windows, so they must not be listed either. diff --git a/backends/cuda/runtime/shims/mm.cu b/backends/cuda/runtime/shims/mm.cu new file mode 100644 index 00000000000..1cc28b15bda --- /dev/null +++ b/backends/cuda/runtime/shims/mm.cu @@ -0,0 +1,678 @@ +/* + * 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. + */ + +#include + +#include +#include +#include +#include + +#include +#include +#ifdef ET_CUBLASLT_LIBRARY +#include +#endif + +#include +#include +#include + +namespace executorch::backends::cuda { + +namespace c10 = executorch::backends::aoti::slim::c10; +using executorch::extension::cuda::CudaAllocator; +using executorch::runtime::Error; + +namespace { + +#define ET_CUBLASLT_FUNCTIONS(X) \ + X(Create) \ + X(Destroy) \ + X(GetStatusString) \ + X(Matmul) \ + X(MatmulAlgoGetHeuristic) \ + X(MatmulDescCreate) \ + X(MatmulDescDestroy) \ + X(MatmulDescSetAttribute) \ + X(MatmulPreferenceCreate) \ + X(MatmulPreferenceDestroy) \ + X(MatmulPreferenceSetAttribute) \ + X(MatrixLayoutCreate) \ + X(MatrixLayoutDestroy) \ + X(MatrixLayoutSetAttribute) + +struct CublasLt { +#define ET_CUBLASLT_MEMBER(name) decltype(&::cublasLt##name) name = nullptr; + ET_CUBLASLT_FUNCTIONS(ET_CUBLASLT_MEMBER) +#undef ET_CUBLASLT_MEMBER +}; + +// CMake builds load cuBLASLt on first use so it is not a load-time dependency. +// Buck links a stub that already loads it lazily, so it uses the symbols as is. +const CublasLt* lt() { + static const CublasLt* const table = []() -> const CublasLt* { + static CublasLt functions; +#ifdef ET_CUBLASLT_LIBRARY + void* library = dlopen(ET_CUBLASLT_LIBRARY, RTLD_NOW | RTLD_LOCAL); + if (library == nullptr) { + ET_LOG(Error, "Cannot load %s: %s", ET_CUBLASLT_LIBRARY, dlerror()); + return nullptr; + } +#define ET_CUBLASLT_RESOLVE(name) \ + functions.name = reinterpret_cast( \ + dlsym(library, "cublasLt" #name)); \ + if (functions.name == nullptr) { \ + ET_LOG(Error, "%s has no cublasLt%s", ET_CUBLASLT_LIBRARY, #name); \ + return nullptr; \ + } + ET_CUBLASLT_FUNCTIONS(ET_CUBLASLT_RESOLVE) +#undef ET_CUBLASLT_RESOLVE +#else +#define ET_CUBLASLT_ASSIGN(name) functions.name = &::cublasLt##name; + ET_CUBLASLT_FUNCTIONS(ET_CUBLASLT_ASSIGN) +#undef ET_CUBLASLT_ASSIGN +#endif + return &functions; + }(); + return table; +} + +#define ET_CUBLAS_CHECK_OR_RETURN_ERROR(EXPR) \ + do { \ + const cublasStatus_t __status = (EXPR); \ + ET_CHECK_OR_RETURN_ERROR( \ + __status == CUBLAS_STATUS_SUCCESS, \ + Internal, \ + "%s failed: %s", \ + #EXPR, \ + lt()->GetStatusString(__status)); \ + } while (0) + +// Only what the chosen algorithm needs is allocated, so this cap only limits +// which algorithms the heuristic may pick. +constexpr size_t kMaxWorkspaceBytes = 32 * 1024 * 1024; + +// One handle per thread and device, kept for the thread's lifetime, so calls +// reuse a handle instead of creating one, and need no lock. +class LtHandles { + public: + ~LtHandles() { + for (cublasLtHandle_t handle : handles_) { + if (handle != nullptr) { + lt()->Destroy(handle); + } + } + } + + cublasStatus_t get(c10::DeviceIndex device, cublasLtHandle_t* handle) { + if (static_cast(device) >= handles_.size()) { + handles_.resize(device + 1, nullptr); + } + if (handles_[device] == nullptr) { + const cublasStatus_t status = lt()->Create(&handles_[device]); + if (status != CUBLAS_STATUS_SUCCESS) { + handles_[device] = nullptr; + return status; + } + } + *handle = handles_[device]; + return CUBLAS_STATUS_SUCCESS; + } + + private: + std::vector handles_; +}; + +template +struct LtDescriptor { + T value = nullptr; + ~LtDescriptor() { + if (value != nullptr) { + (lt()->*Destroy)(value); + } + } +}; +using MatmulDesc = + LtDescriptor; +using LayoutDesc = + LtDescriptor; +using PreferenceDesc = LtDescriptor< + cublasLtMatmulPreference_t, + &CublasLt::MatmulPreferenceDestroy>; + +// A matrix, or a batch of them: element (r, c) of batch b sits at +// b * batch_stride + r * row_stride + c * col_stride elements past data. +struct Matrix { + void* data; + int64_t rows; + int64_t cols; + int64_t row_stride; + int64_t col_stride; + int64_t batch_stride; +}; + +Matrix matrix_of(Tensor* t) { + const int64_t d = static_cast(t->dim()); + return Matrix{ + t->data_ptr(), + t->size(d - 2), + t->size(d - 1), + t->stride(d - 2), + t->stride(d - 1), + d == 3 ? t->stride(0) : 0}; +} + +Matrix transposed(Matrix m) { + std::swap(m.rows, m.cols); + std::swap(m.row_stride, m.col_stride); + return m; +} + +// Capped at cuBLASLt's default of 256 bytes, as PyTorch does. +uint32_t alignment_of(const void* pointer) { + const auto address = reinterpret_cast(pointer); + uint32_t alignment = 1; + while (alignment < 256 && address % (alignment * 2) == 0) { + alignment *= 2; + } + return alignment; +} + +// cuBLAS reads column-major matrices: consecutive elements of a column are +// adjacent, and columns start a leading dimension apart. +bool is_column_major(const Matrix& m, int64_t* ld) { + if (m.row_stride != 1 && m.rows != 1) { + return false; + } + *ld = m.cols == 1 ? std::max(m.rows, 1) : m.col_stride; + return *ld >= std::max(m.rows, 1); +} + +bool is_row_major(const Matrix& m, int64_t* ld) { + return is_column_major(transposed(m), ld); +} + +// self as addmm broadcasts it to rows x cols: a missing or size-1 dimension +// repeats with a zero stride. +Matrix broadcast(Tensor* self, int64_t rows, int64_t cols) { + const int64_t d = static_cast(self->dim()); + const auto stride = [&](int64_t dim) { + return dim >= 0 && self->size(dim) != 1 ? self->stride(dim) : 0; + }; + return Matrix{self->data_ptr(), rows, cols, stride(d - 2), stride(d - 1), 0}; +} + +// Elements move as raw bits, so one kernel per element size serves every dtype. +template +__global__ void copy_matrix_kernel(Matrix dst, Matrix src) { + const int64_t numel = dst.rows * dst.cols; + for (int64_t i = blockIdx.x * int64_t{blockDim.x} + threadIdx.x; i < numel; + i += int64_t{gridDim.x} * blockDim.x) { + const int64_t r = i / dst.cols; + const int64_t c = i % dst.cols; + static_cast(dst.data)[r * dst.row_stride + c * dst.col_stride] = + static_cast( + src.data)[r * src.row_stride + c * src.col_stride]; + } +} + +Error copy_matrix( + const char* name, + const Matrix& dst, + const Matrix& src, + size_t element_size, + c10::DeviceIndex device) { + auto stream = getCurrentCUDAStream(device); + ET_CHECK_OK_OR_RETURN_ERROR(stream.error(), "%s: no CUDA stream", name); + constexpr int64_t kThreads = 256; + const int blocks = static_cast( + std::min((dst.rows * dst.cols + kThreads - 1) / kThreads, 4096)); + switch (element_size) { + case sizeof(uint32_t): + copy_matrix_kernel + <<>>(dst, src); + break; + case sizeof(uint16_t): + copy_matrix_kernel + <<>>(dst, src); + break; + default: + ET_CHECK_OR_RETURN_ERROR( + false, + NotSupported, + "%s: unsupported element size %zu", + name, + element_size); + } + ET_CUDA_KERNEL_LAUNCH_CHECK_OR_RETURN_ERROR(); + return Error::Ok; +} + +// How cuBLAS reads an operand: as stored when it is column-major, or as the +// transpose of the column-major matrix it is when it is row-major. +struct Operand { + cublasOperation_t op; + uint64_t stored_rows; + uint64_t stored_cols; + int64_t ld; +}; + +bool as_operand(const Matrix& m, Operand* operand) { + int64_t ld = 0; + if (is_column_major(m, &ld)) { + *operand = Operand{CUBLAS_OP_N, uint64_t(m.rows), uint64_t(m.cols), ld}; + return true; + } + if (is_row_major(m, &ld)) { + *operand = Operand{CUBLAS_OP_T, uint64_t(m.cols), uint64_t(m.rows), ld}; + return true; + } + return false; +} + +Error make_layout( + LayoutDesc& layout, + cudaDataType_t type, + uint64_t rows, + uint64_t cols, + int64_t ld, + int32_t batch, + int64_t batch_stride) { + ET_CUBLAS_CHECK_OR_RETURN_ERROR( + lt()->MatrixLayoutCreate(&layout.value, type, rows, cols, ld)); + if (batch > 1) { + ET_CUBLAS_CHECK_OR_RETURN_ERROR(lt()->MatrixLayoutSetAttribute( + layout.value, + CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, + &batch, + sizeof(batch))); + ET_CUBLAS_CHECK_OR_RETURN_ERROR(lt()->MatrixLayoutSetAttribute( + layout.value, + CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, + &batch_stride, + sizeof(batch_stride))); + } + return Error::Ok; +} + +// d = alpha * a @ b + beta * d, plus bias added to every column of d when bias +// is not null. d must be column-major. +Error gemm( + const char* name, + const Matrix& d, + const Matrix& a, + const Matrix& b, + int32_t batch, + cudaDataType_t type, + c10::DeviceIndex device, + const void* bias, + float alpha, + float beta) { + int64_t ldd = 0; + ET_CHECK_OR_RETURN_ERROR( + is_column_major(d, &ldd), + Internal, + "%s: result is not column-major", + name); + Operand op_a{}; + Operand op_b{}; + ET_CHECK_OR_RETURN_ERROR( + as_operand(a, &op_a) && as_operand(b, &op_b), + NotSupported, + "%s: an operand is neither row-major nor column-major", + name); + + ET_CHECK_OR_RETURN_ERROR( + lt() != nullptr, NotFound, "%s: cannot load cuBLASLt", name); + + auto stream = getCurrentCUDAStream(device); + ET_CHECK_OK_OR_RETURN_ERROR(stream.error(), "%s: no CUDA stream", name); + + thread_local LtHandles handles; + cublasLtHandle_t handle = nullptr; + ET_CUBLAS_CHECK_OR_RETURN_ERROR(handles.get(device, &handle)); + + MatmulDesc desc; + ET_CUBLAS_CHECK_OR_RETURN_ERROR( + lt()->MatmulDescCreate(&desc.value, CUBLAS_COMPUTE_32F, CUDA_R_32F)); + ET_CUBLAS_CHECK_OR_RETURN_ERROR(lt()->MatmulDescSetAttribute( + desc.value, CUBLASLT_MATMUL_DESC_TRANSA, &op_a.op, sizeof(op_a.op))); + ET_CUBLAS_CHECK_OR_RETURN_ERROR(lt()->MatmulDescSetAttribute( + desc.value, CUBLASLT_MATMUL_DESC_TRANSB, &op_b.op, sizeof(op_b.op))); + if (bias != nullptr) { + const cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_BIAS; + ET_CUBLAS_CHECK_OR_RETURN_ERROR(lt()->MatmulDescSetAttribute( + desc.value, + CUBLASLT_MATMUL_DESC_EPILOGUE, + &epilogue, + sizeof(epilogue))); + ET_CUBLAS_CHECK_OR_RETURN_ERROR(lt()->MatmulDescSetAttribute( + desc.value, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias, sizeof(bias))); + } + + LayoutDesc layout_a; + LayoutDesc layout_b; + LayoutDesc layout_d; + ET_CHECK_OK_OR_RETURN_ERROR(make_layout( + layout_a, + type, + op_a.stored_rows, + op_a.stored_cols, + op_a.ld, + batch, + a.batch_stride)); + ET_CHECK_OK_OR_RETURN_ERROR(make_layout( + layout_b, + type, + op_b.stored_rows, + op_b.stored_cols, + op_b.ld, + batch, + b.batch_stride)); + ET_CHECK_OK_OR_RETURN_ERROR(make_layout( + layout_d, + type, + uint64_t(d.rows), + uint64_t(d.cols), + ldd, + batch, + d.batch_stride)); + + PreferenceDesc preference; + ET_CUBLAS_CHECK_OR_RETURN_ERROR( + lt()->MatmulPreferenceCreate(&preference.value)); + ET_CUBLAS_CHECK_OR_RETURN_ERROR(lt()->MatmulPreferenceSetAttribute( + preference.value, + CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, + &kMaxWorkspaceBytes, + sizeof(kMaxWorkspaceBytes))); + // The heuristic otherwise assumes 256-byte aligned pointers, and can pick an + // algorithm that then fails on a view that starts partway into its storage. + const std::pair + alignments[] = { + {CUBLASLT_MATMUL_PREF_MIN_ALIGNMENT_A_BYTES, a.data}, + {CUBLASLT_MATMUL_PREF_MIN_ALIGNMENT_B_BYTES, b.data}, + {CUBLASLT_MATMUL_PREF_MIN_ALIGNMENT_C_BYTES, d.data}, + {CUBLASLT_MATMUL_PREF_MIN_ALIGNMENT_D_BYTES, d.data}}; + for (const auto& [attribute, pointer] : alignments) { + const uint32_t alignment = alignment_of(pointer); + ET_CUBLAS_CHECK_OR_RETURN_ERROR(lt()->MatmulPreferenceSetAttribute( + preference.value, attribute, &alignment, sizeof(alignment))); + } + + cublasLtMatmulHeuristicResult_t heuristic{}; + int found = 0; + ET_CUBLAS_CHECK_OR_RETURN_ERROR(lt()->MatmulAlgoGetHeuristic( + handle, + desc.value, + layout_a.value, + layout_b.value, + layout_d.value, + layout_d.value, + preference.value, + 1, + &heuristic, + &found)); + ET_CHECK_OR_RETURN_ERROR( + found > 0, NotSupported, "%s: cuBLASLt has no algorithm", name); + + void* workspace = nullptr; + if (heuristic.workspaceSize > 0) { + auto allocation = CudaAllocator::allocate_async( + heuristic.workspaceSize, device, stream.get()); + ET_CHECK_OK_OR_RETURN_ERROR( + allocation.error(), "%s: workspace allocation failed", name); + workspace = allocation.get(); + } + const cublasStatus_t status = lt()->Matmul( + handle, + desc.value, + &alpha, + a.data, + layout_a.value, + b.data, + layout_b.value, + &beta, + d.data, + layout_d.value, + d.data, + layout_d.value, + &heuristic.algo, + workspace, + heuristic.workspaceSize, + stream.get()); + if (workspace != nullptr) { + CudaAllocator::deallocate_async(workspace, device, stream.get()); + } + ET_CHECK_OR_RETURN_ERROR( + status == CUBLAS_STATUS_SUCCESS, + Internal, + "%s: cublasLtMatmul failed: %s", + name, + lt()->GetStatusString(status)); + return Error::Ok; +} + +Error cuda_data_type( + const char* name, + c10::ScalarType dtype, + cudaDataType_t* type) { + switch (dtype) { + case c10::ScalarType::Float: + *type = CUDA_R_32F; + return Error::Ok; + case c10::ScalarType::Half: + *type = CUDA_R_16F; + return Error::Ok; + case c10::ScalarType::BFloat16: + *type = CUDA_R_16BF; + return Error::Ok; + default: + ET_CHECK_OR_RETURN_ERROR( + false, + NotSupported, + "%s: unsupported dtype %d", + name, + static_cast(dtype)); + } +} + +// Checks out = self @ mat2 for 2-D (ndim 2) or batched (ndim 3) operands. +Error check_operands( + const char* name, + Tensor* out, + Tensor* self, + Tensor* mat2, + size_t ndim) { + ET_CHECK_OR_RETURN_ERROR( + out != nullptr && self != nullptr && mat2 != nullptr, + InvalidArgument, + "%s: a tensor is null", + name); + ET_CHECK_OR_RETURN_ERROR( + out->dim() == ndim && self->dim() == ndim && mat2->dim() == ndim, + InvalidArgument, + "%s: expected %zu-D tensors", + name, + ndim); + ET_CHECK_OR_RETURN_ERROR( + out->device().is_cuda() && out->device_index() >= 0 && + self->device() == out->device() && mat2->device() == out->device(), + InvalidArgument, + "%s: expected tensors on one CUDA device", + name); + ET_CHECK_OR_RETURN_ERROR( + self->dtype() == out->dtype() && mat2->dtype() == out->dtype(), + InvalidArgument, + "%s: tensors differ in dtype", + name); + const size_t m = ndim - 2; + ET_CHECK_OR_RETURN_ERROR( + self->size(m + 1) == mat2->size(m) && out->size(m) == self->size(m) && + out->size(m + 1) == mat2->size(m + 1), + InvalidArgument, + "%s: shapes do not multiply", + name); + ET_CHECK_OR_RETURN_ERROR( + ndim == 2 || + (self->size(0) == out->size(0) && mat2->size(0) == out->size(0)), + InvalidArgument, + "%s: batch sizes differ", + name); + ET_CHECK_OR_RETURN_ERROR( + self->size(m + 1) > 0 || out->numel() == 0, + NotSupported, + "%s: an empty inner dimension is not supported", + name); + return Error::Ok; +} + +// out = alpha * self @ mat2 + beta * out, plus a bias row added to every row of +// out when bias is not null. +Error multiply( + const char* name, + Tensor* out, + Tensor* self, + Tensor* mat2, + const void* bias, + float alpha, + float beta) { + if (out->numel() == 0) { + return Error::Ok; + } + cudaDataType_t type{}; + ET_CHECK_OK_OR_RETURN_ERROR(cuda_data_type(name, out->dtype(), &type)); + const int64_t batch = out->dim() == 3 ? out->size(0) : 1; + ET_CHECK_OR_RETURN_ERROR( + batch <= INT32_MAX, NotSupported, "%s: batch too large", name); + + Matrix d = matrix_of(out); + Matrix a = matrix_of(self); + Matrix b = matrix_of(mat2); + ET_CHECK_OR_RETURN_ERROR( + batch == 1 || d.batch_stride != 0, + InvalidArgument, + "%s: the result repeats one matrix over the batch", + name); + int64_t ld = 0; + // cuBLAS writes column-major results, so a row-major out is computed as + // out^T = mat2^T @ self^T. That orientation also lines a bias row up with + // the rows of the column-major result, which is what the epilogue adds. + if (is_row_major(d, &ld)) { + d = transposed(d); + std::swap(a, b); + a = transposed(a); + b = transposed(b); + } else { + ET_CHECK_OR_RETURN_ERROR( + is_column_major(d, &ld), + NotSupported, + "%s: the result is neither row-major nor column-major", + name); + ET_CHECK_OR_RETURN_ERROR( + bias == nullptr, + NotSupported, + "%s: a bias needs a row-major result", + name); + } + return gemm( + name, + d, + a, + b, + static_cast(batch), + type, + out->device_index(), + bias, + alpha, + beta); +} + +} // namespace + +#ifdef __cplusplus +extern "C" { +#endif + +AOTITorchError aoti_torch_cuda_mm_out(Tensor* out, Tensor* self, Tensor* mat2) { + constexpr const char* name = "aoti_torch_cuda_mm_out"; + ET_CHECK_OK_OR_RETURN_ERROR(check_operands(name, out, self, mat2, 2)); + return multiply(name, out, self, mat2, nullptr, 1.0f, 0.0f); +} + +AOTITorchError aoti_torch_cuda_addmm_out( + Tensor* out, + Tensor* self, + Tensor* mat1, + Tensor* mat2, + double beta, + double alpha) { + constexpr const char* name = "aoti_torch_cuda_addmm_out"; + ET_CHECK_OK_OR_RETURN_ERROR(check_operands(name, out, mat1, mat2, 2)); + ET_CHECK_OR_RETURN_ERROR( + self != nullptr && self->device() == out->device() && + self->dtype() == out->dtype(), + InvalidArgument, + "%s: self must be on the result's device and of its dtype", + name); + if (out->numel() == 0) { + return Error::Ok; + } + const int64_t rows = out->size(0); + const int64_t cols = out->size(1); + const int64_t self_dim = static_cast(self->dim()); + ET_CHECK_OR_RETURN_ERROR( + self_dim <= 2 && + (self_dim < 1 || self->size(self_dim - 1) == cols || + self->size(self_dim - 1) == 1) && + (self_dim < 2 || self->size(0) == rows || self->size(0) == 1), + InvalidArgument, + "%s: self does not broadcast to the result", + name); + + const Matrix src = broadcast(self, rows, cols); + const Matrix dst = matrix_of(out); + int64_t ld = 0; + // A row added to every row of a row-major result runs as the cuBLASLt bias + // epilogue, as PyTorch does for a linear layer. + if (beta == 1.0 && src.row_stride == 0 && + (src.col_stride == 1 || cols == 1) && is_row_major(dst, &ld)) { + return multiply( + name, out, mat1, mat2, self->data_ptr(), static_cast(alpha), 0); + } + // Any other self is written into out and accumulated there with beta. With + // beta == 0, cuBLAS does not read out, so self is ignored, as in PyTorch. + const bool self_is_out = src.data == dst.data && + src.row_stride == dst.row_stride && src.col_stride == dst.col_stride; + if (beta != 0.0 && !self_is_out) { + ET_CHECK_OK_OR_RETURN_ERROR( + copy_matrix(name, dst, src, out->itemsize(), out->device_index())); + } + return multiply( + name, + out, + mat1, + mat2, + nullptr, + static_cast(alpha), + static_cast(beta)); +} + +AOTITorchError +aoti_torch_cuda_bmm_out(Tensor* out, Tensor* self, Tensor* mat2) { + constexpr const char* name = "aoti_torch_cuda_bmm_out"; + ET_CHECK_OK_OR_RETURN_ERROR(check_operands(name, out, self, mat2, 3)); + return multiply(name, out, self, mat2, nullptr, 1.0f, 0.0f); +} + +#ifdef __cplusplus +} +#endif + +} // namespace executorch::backends::cuda diff --git a/backends/cuda/runtime/shims/mm.h b/backends/cuda/runtime/shims/mm.h new file mode 100644 index 00000000000..07578bfc00f --- /dev/null +++ b/backends/cuda/runtime/shims/mm.h @@ -0,0 +1,73 @@ +/* + * 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 +#include + +namespace executorch::backends::cuda { + +using executorch::backends::aoti::AOTITorchError; +using executorch::backends::aoti::Tensor; + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * Matrix multiplication on cuBLASLt, the AOTI shim for aten::mm.out. Inductor + * calls it for an mm when autotuning picks cuBLAS, and for every mm when + * max_autotune is off. + * + * Every tensor is a 2-D float32, float16 or bfloat16 CUDA tensor of one dtype, + * and each matrix must be row-major or column-major (a leading dimension may + * exceed the matrix width, as for a slice). float32 runs in full precision. + * + * @param out Result, M x N, written in place + * @param self Left operand, M x K + * @param mat2 Right operand, K x N + * @return Error::Ok on success, Error::InvalidArgument or Error::NotSupported + * for operands this shim does not handle, Error::NotFound when cuBLASLt cannot + * be loaded, Error::Internal on a cuBLAS failure, or the allocator's error when + * the workspace cannot be allocated + */ +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_cuda_mm_out(Tensor* out, Tensor* self, Tensor* mat2); + +/** + * out = beta * self + alpha * (mat1 @ mat2), the AOTI shim for aten::addmm.out, + * with the same operand rules as aoti_torch_cuda_mm_out. + * + * self is anything that broadcasts to M x N, with any strides. A row of N + * values added to a row-major out with beta == 1 runs as a cuBLASLt bias + * epilogue, as PyTorch does for a linear layer. Any other self is first copied + * into out, unless it is out itself or beta == 0. + */ +AOTI_SHIM_EXPORT AOTITorchError aoti_torch_cuda_addmm_out( + Tensor* out, + Tensor* self, + Tensor* mat1, + Tensor* mat2, + double beta, + double alpha); + +/** + * Batched matrix multiplication, the AOTI shim for aten::bmm.out, with the + * same operand rules as aoti_torch_cuda_mm_out applied to each matrix of the + * batch. The tensors are 3-D; an operand may repeat one matrix over the batch + * with a zero batch stride. + */ +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_cuda_bmm_out(Tensor* out, Tensor* self, Tensor* mat2); + +#ifdef __cplusplus +} +#endif + +} // namespace executorch::backends::cuda diff --git a/backends/cuda/runtime/shims/tests/CMakeLists.txt b/backends/cuda/runtime/shims/tests/CMakeLists.txt index 0e2a4bd599b..2875c282ab0 100644 --- a/backends/cuda/runtime/shims/tests/CMakeLists.txt +++ b/backends/cuda/runtime/shims/tests/CMakeLists.txt @@ -51,6 +51,10 @@ set(CUDA_SHIM_TESTS # CUDA-specific tests requiring GPU kernels set(CUDA_KERNEL_TESTS test_aoti_torch_cuda__weight_int4pack_mm) +# The cuBLAS shims are built only on Linux. +if(NOT WIN32) + list(APPEND CUDA_KERNEL_TESTS test_aoti_torch_cuda_mm_out) +endif() enable_testing() diff --git a/backends/cuda/runtime/shims/tests/targets.bzl b/backends/cuda/runtime/shims/tests/targets.bzl index 863d0993a7c..89701d106fc 100644 --- a/backends/cuda/runtime/shims/tests/targets.bzl +++ b/backends/cuda/runtime/shims/tests/targets.bzl @@ -47,6 +47,7 @@ def define_common_targets(is_fbcode = False): cuda_shim_cpp_unittest("aoti_torch_new_tensor_handle") cuda_shim_cpp_unittest("aoti_torch_item_bool") cuda_shim_cpp_unittest("aoti_torch_assign_tensors_out") + cuda_shim_cpp_unittest("aoti_torch_cuda_mm_out") cpp_unittest( name = "test_op__device_copy", diff --git a/backends/cuda/runtime/shims/tests/test_aoti_torch_cuda_mm_out.cpp b/backends/cuda/runtime/shims/tests/test_aoti_torch_cuda_mm_out.cpp new file mode 100644 index 00000000000..b4792a010dc --- /dev/null +++ b/backends/cuda/runtime/shims/tests/test_aoti_torch_cuda_mm_out.cpp @@ -0,0 +1,511 @@ +/* + * 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. + */ + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +using namespace executorch::backends::cuda; +using executorch::runtime::Error; +namespace slim_c10 = executorch::backends::aoti::slim::c10; + +namespace { + +// Small integers, so every product and sum below is exact in float16 and +// bfloat16 too, and results can be compared exactly. +float pattern(int64_t i) { + return static_cast((i * 7) % 5 - 2); +} + +class AOTITorchCudaMmTest + : public ::testing::TestWithParam { + protected: + void SetUp() override { + et_pal_init(); + int device_count = 0; + if (cudaGetDeviceCount(&device_count) != cudaSuccess || device_count == 0) { + GTEST_SKIP() << "CUDA not available"; + } + } + + void TearDown() override { + for (Tensor* t : tensors_) { + aoti_torch_delete_tensor_object(t); + } + } + + slim_c10::ScalarType dtype() const { + return GetParam(); + } + + // A CUDA buffer of `numel` elements holding pattern(i) at element i. + Tensor* buffer(int64_t numel) { + Tensor* t = nullptr; + EXPECT_EQ( + aoti_torch_empty_strided( + 1, + &numel, + nullptr, + static_cast(dtype()), + static_cast(slim_c10::DeviceType::CUDA), + 0, + &t), + Error::Ok); + tensors_.push_back(t); + std::vector values(numel); + for (int64_t i = 0; i < numel; ++i) { + values[i] = pattern(i); + } + upload(t, values); + return t; + } + + // A view of `base` with the given sizes and strides. + Tensor* view( + Tensor* base, + const std::vector& sizes, + const std::vector& strides) { + Tensor* t = nullptr; + EXPECT_EQ( + aoti_torch__reinterpret_tensor( + base, sizes.size(), sizes.data(), strides.data(), 0, &t), + Error::Ok); + tensors_.push_back(t); + return t; + } + + void upload(Tensor* t, const std::vector& values) { + std::vector half(values.size()); + const void* src = values.data(); + if (dtype() == slim_c10::ScalarType::Half) { + for (size_t i = 0; i < values.size(); ++i) { + half[i] = __half_as_ushort(__float2half(values[i])); + } + src = half.data(); + } else if (dtype() == slim_c10::ScalarType::BFloat16) { + for (size_t i = 0; i < values.size(); ++i) { + half[i] = __bfloat16_as_ushort(__float2bfloat16(values[i])); + } + src = half.data(); + } + ASSERT_EQ( + cudaMemcpy( + t->data_ptr(), + src, + values.size() * t->itemsize(), + cudaMemcpyHostToDevice), + cudaSuccess); + } + + // Every element of the contiguous buffer `base`, as floats. + std::vector download(Tensor* base) { + EXPECT_EQ(cudaDeviceSynchronize(), cudaSuccess); + const size_t numel = base->numel(); + std::vector out(numel); + if (dtype() == slim_c10::ScalarType::Float) { + EXPECT_EQ( + cudaMemcpy( + out.data(), + base->data_ptr(), + numel * sizeof(float), + cudaMemcpyDeviceToHost), + cudaSuccess); + return out; + } + std::vector half(numel); + EXPECT_EQ( + cudaMemcpy( + half.data(), + base->data_ptr(), + numel * sizeof(uint16_t), + cudaMemcpyDeviceToHost), + cudaSuccess); + for (size_t i = 0; i < numel; ++i) { + out[i] = dtype() == slim_c10::ScalarType::Half + ? __half2float(__ushort_as_half(half[i])) + : __bfloat162float(__ushort_as_bfloat16(half[i])); + } + return out; + } + + // Reads element (b, r, c) of a 2-D or 3-D view from the host copy of the + // pattern its base was filled with. + static float at(Tensor* t, int64_t b, int64_t r, int64_t c) { + const int64_t d = static_cast(t->dim()); + const int64_t batch_stride = d == 3 ? t->stride(0) : 0; + return pattern( + b * batch_stride + r * t->stride(d - 2) + c * t->stride(d - 1)); + } + + // Checks out (contiguous base `out_base` viewed as `out`) against + // beta * self_value(b, r, c) + alpha * a @ b. + template + void expect_product( + Tensor* out_base, + Tensor* out, + Tensor* a, + Tensor* b, + float alpha, + float beta, + SelfValue self_value) { + const std::vector got = download(out_base); + const int64_t d = static_cast(out->dim()); + const int64_t batch = d == 3 ? out->size(0) : 1; + const int64_t rows = out->size(d - 2); + const int64_t cols = out->size(d - 1); + const int64_t inner = a->size(d - 1); + const int64_t out_batch_stride = d == 3 ? out->stride(0) : 0; + for (int64_t n = 0; n < batch; ++n) { + for (int64_t r = 0; r < rows; ++r) { + for (int64_t c = 0; c < cols; ++c) { + float expected = 0; + for (int64_t k = 0; k < inner; ++k) { + expected += at(a, n, r, k) * at(b, n, k, c); + } + expected = alpha * expected + beta * self_value(n, r, c); + const int64_t index = n * out_batch_stride + r * out->stride(d - 2) + + c * out->stride(d - 1); + ASSERT_EQ(got[index], expected) + << "batch " << n << " row " << r << " col " << c; + } + } + } + } + + // Records an m x k by k x n mm into a CUDA graph, replays it, and checks + // the result. Counts the graph's memory allocations into *allocations, which + // tells whether cuBLASLt used a workspace. + void expect_captured_mm( + int64_t m, + int64_t k, + int64_t n, + size_t* allocations = nullptr) { + Tensor* a = view(buffer(m * k), {m, k}, {k, 1}); + Tensor* b = view(buffer(k * n), {k, n}, {1, k}); + Tensor* out_base = buffer(m * n); + aoti_torch_zero_(out_base); + Tensor* out = view(out_base, {m, n}, {n, 1}); + + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + ASSERT_EQ(setCurrentCUDAStream(stream, 0), Error::Ok); + ASSERT_EQ( + cudaStreamBeginCapture(stream, cudaStreamCaptureModeRelaxed), + cudaSuccess); + const Error captured = aoti_torch_cuda_mm_out(out, a, b); + cudaGraph_t graph = nullptr; + ASSERT_EQ(cudaStreamEndCapture(stream, &graph), cudaSuccess); + ASSERT_EQ(captured, Error::Ok); + if (allocations != nullptr) { + size_t count = 0; + ASSERT_EQ(cudaGraphGetNodes(graph, nullptr, &count), cudaSuccess); + std::vector nodes(count); + ASSERT_EQ(cudaGraphGetNodes(graph, nodes.data(), &count), cudaSuccess); + *allocations = 0; + for (cudaGraphNode_t node : nodes) { + cudaGraphNodeType type{}; + ASSERT_EQ(cudaGraphNodeGetType(node, &type), cudaSuccess); + *allocations += type == cudaGraphNodeTypeMemAlloc ? 1 : 0; + } + } + cudaGraphExec_t exec = nullptr; + ASSERT_EQ( + cudaGraphInstantiate( + &exec, graph, cudaGraphInstantiateFlagAutoFreeOnLaunch), + cudaSuccess); + ASSERT_EQ(cudaGraphLaunch(exec, stream), cudaSuccess); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + clearCurrentCUDAStream(0); + expect_product(out_base, out, a, b, 1, 0, [](auto...) { return 0.0f; }); + cudaGraphExecDestroy(exec); + cudaGraphDestroy(graph); + cudaStreamDestroy(stream); + } + + std::vector tensors_; +}; + +constexpr int64_t M = 5; +constexpr int64_t K = 7; +constexpr int64_t N = 3; + +TEST_P(AOTITorchCudaMmTest, MmContiguous) { + Tensor* a = view(buffer(M * K), {M, K}, {K, 1}); + Tensor* b = view(buffer(K * N), {K, N}, {N, 1}); + Tensor* out_base = buffer(M * N); + Tensor* out = view(out_base, {M, N}, {N, 1}); + ASSERT_EQ(aoti_torch_cuda_mm_out(out, a, b), Error::Ok); + expect_product(out_base, out, a, b, 1, 0, [](auto...) { return 0.0f; }); +} + +// A view that starts one element into its storage is not 16-byte aligned. +// Sizes are multiples of 16 so that only the pointer is misaligned. +TEST_P(AOTITorchCudaMmTest, MmMisalignedOperand) { + constexpr int64_t S = 64; + Tensor* a_base = buffer(S * S + 1); + Tensor* a = nullptr; + const std::vector sizes = {S, S}; + const std::vector strides = {S, 1}; + ASSERT_EQ( + aoti_torch__reinterpret_tensor( + a_base, 2, sizes.data(), strides.data(), 1, &a), + Error::Ok); + tensors_.push_back(a); + Tensor* b = view(buffer(S * S), {S, S}, {S, 1}); + Tensor* out_base = buffer(S * S); + Tensor* out = view(out_base, {S, S}, {S, 1}); + ASSERT_EQ(aoti_torch_cuda_mm_out(out, a, b), Error::Ok); + const std::vector got = download(out_base); + for (int64_t r = 0; r < S; ++r) { + for (int64_t c = 0; c < S; ++c) { + float expected = 0; + for (int64_t k = 0; k < S; ++k) { + expected += pattern(1 + r * S + k) * pattern(k * S + c); + } + ASSERT_EQ(got[r * S + c], expected) << "row " << r << " col " << c; + } + } +} + +// A linear layer's weight arrives transposed (column-major), and a slice keeps +// its parent's row stride. +TEST_P(AOTITorchCudaMmTest, MmTransposedAndSlicedOperands) { + Tensor* a = view(buffer(M * K), {M, K}, {1, M}); + Tensor* b = view(buffer(K * (N + 4)), {K, N}, {N + 4, 1}); + Tensor* out_base = buffer(M * N); + Tensor* out = view(out_base, {M, N}, {N, 1}); + ASSERT_EQ(aoti_torch_cuda_mm_out(out, a, b), Error::Ok); + expect_product(out_base, out, a, b, 1, 0, [](auto...) { return 0.0f; }); +} + +TEST_P(AOTITorchCudaMmTest, MmColumnMajorResult) { + Tensor* a = view(buffer(M * K), {M, K}, {K, 1}); + Tensor* b = view(buffer(K * N), {K, N}, {N, 1}); + Tensor* out_base = buffer(M * N); + Tensor* out = view(out_base, {M, N}, {1, M}); + ASSERT_EQ(aoti_torch_cuda_mm_out(out, a, b), Error::Ok); + expect_product(out_base, out, a, b, 1, 0, [](auto...) { return 0.0f; }); +} + +// A dimension of size 1 can have any stride, so these are still matrices: +// a row without a unit stride, and a 1 x 1 matrix expanded from a scalar. +TEST_P(AOTITorchCudaMmTest, MmRowWithoutUnitStride) { + Tensor* a = view(buffer(2 * K), {1, K}, {2 * K, 2}); + Tensor* b = view(buffer(K * N), {K, N}, {N, 1}); + Tensor* out_base = buffer(N); + Tensor* out = view(out_base, {1, N}, {N, 1}); + ASSERT_EQ(aoti_torch_cuda_mm_out(out, a, b), Error::Ok); + expect_product(out_base, out, a, b, 1, 0, [](auto...) { return 0.0f; }); +} + +TEST_P(AOTITorchCudaMmTest, MmOneByOneWithZeroStrides) { + Tensor* a = view(buffer(1), {1, 1}, {0, 0}); + Tensor* b = view(buffer(N), {1, N}, {N, 1}); + Tensor* out_base = buffer(N); + Tensor* out = view(out_base, {1, N}, {N, 1}); + ASSERT_EQ(aoti_torch_cuda_mm_out(out, a, b), Error::Ok); + expect_product(out_base, out, a, b, 1, 0, [](auto...) { return 0.0f; }); +} + +TEST_P(AOTITorchCudaMmTest, AddmmBiasRow) { + Tensor* bias = view(buffer(N), {N}, {1}); + Tensor* a = view(buffer(M * K), {M, K}, {K, 1}); + Tensor* b = view(buffer(K * N), {K, N}, {1, K}); + Tensor* out_base = buffer(M * N); + Tensor* out = view(out_base, {M, N}, {N, 1}); + ASSERT_EQ(aoti_torch_cuda_addmm_out(out, bias, a, b, 1.0, 2.0), Error::Ok); + expect_product(out_base, out, a, b, 2, 1, [](int64_t, int64_t, int64_t c) { + return pattern(c); + }); +} + +// A linear layer at decode time has one row. +TEST_P(AOTITorchCudaMmTest, AddmmBiasRowSingleRow) { + Tensor* bias = view(buffer(N), {N}, {1}); + Tensor* a = view(buffer(K), {1, K}, {K, 1}); + Tensor* b = view(buffer(K * N), {K, N}, {1, K}); + Tensor* out_base = buffer(N); + Tensor* out = view(out_base, {1, N}, {N, 1}); + ASSERT_EQ(aoti_torch_cuda_addmm_out(out, bias, a, b, 1.0, 1.0), Error::Ok); + expect_product(out_base, out, a, b, 1, 1, [](int64_t, int64_t, int64_t c) { + return pattern(c); + }); +} + +TEST_P(AOTITorchCudaMmTest, AddmmBiasRowSingleColumn) { + Tensor* bias = view(buffer(1), {1}, {1}); + Tensor* a = view(buffer(M * K), {M, K}, {K, 1}); + Tensor* b = view(buffer(K), {K, 1}, {1, 1}); + Tensor* out_base = buffer(M); + Tensor* out = view(out_base, {M, 1}, {1, 1}); + ASSERT_EQ(aoti_torch_cuda_addmm_out(out, bias, a, b, 1.0, 1.0), Error::Ok); + expect_product(out_base, out, a, b, 1, 1, [](int64_t, int64_t, int64_t c) { + return pattern(c); + }); +} + +TEST_P(AOTITorchCudaMmTest, AddmmExpandedBias) { + Tensor* bias = view(buffer(N), {M, N}, {0, 1}); + Tensor* a = view(buffer(M * K), {M, K}, {K, 1}); + Tensor* b = view(buffer(K * N), {K, N}, {N, 1}); + Tensor* out_base = buffer(M * N); + Tensor* out = view(out_base, {M, N}, {N, 1}); + ASSERT_EQ(aoti_torch_cuda_addmm_out(out, bias, a, b, 1.0, 1.0), Error::Ok); + expect_product(out_base, out, a, b, 1, 1, [](int64_t, int64_t, int64_t c) { + return pattern(c); + }); +} + +TEST_P(AOTITorchCudaMmTest, AddmmFullSelf) { + Tensor* self = view(buffer(M * N), {M, N}, {N, 1}); + Tensor* a = view(buffer(M * K), {M, K}, {K, 1}); + Tensor* b = view(buffer(K * N), {K, N}, {N, 1}); + Tensor* out_base = buffer(M * N); + Tensor* out = view(out_base, {M, N}, {N, 1}); + ASSERT_EQ(aoti_torch_cuda_addmm_out(out, self, a, b, 0.5, 2.0), Error::Ok); + expect_product( + out_base, out, a, b, 2, 0.5, [](int64_t, int64_t r, int64_t c) { + return pattern(r * N + c); + }); +} + +TEST_P(AOTITorchCudaMmTest, AddmmScaledBiasRow) { + Tensor* bias = view(buffer(N), {N}, {1}); + Tensor* a = view(buffer(M * K), {M, K}, {K, 1}); + Tensor* b = view(buffer(K * N), {K, N}, {N, 1}); + Tensor* out_base = buffer(M * N); + Tensor* out = view(out_base, {M, N}, {N, 1}); + ASSERT_EQ(aoti_torch_cuda_addmm_out(out, bias, a, b, 0.5, 1.0), Error::Ok); + expect_product(out_base, out, a, b, 1, 0.5, [](int64_t, int64_t, int64_t c) { + return pattern(c); + }); +} + +// Inductor passes a broadcast self already expanded to M x N, with zero +// strides. +TEST_P(AOTITorchCudaMmTest, AddmmExpandedColumn) { + Tensor* column = view(buffer(M), {M, N}, {1, 0}); + Tensor* a = view(buffer(M * K), {M, K}, {K, 1}); + Tensor* b = view(buffer(K * N), {K, N}, {N, 1}); + Tensor* out_base = buffer(M * N); + Tensor* out = view(out_base, {M, N}, {N, 1}); + ASSERT_EQ(aoti_torch_cuda_addmm_out(out, column, a, b, 1.0, 1.0), Error::Ok); + expect_product(out_base, out, a, b, 1, 1, [](int64_t, int64_t r, int64_t) { + return pattern(r); + }); +} + +TEST_P(AOTITorchCudaMmTest, AddmmExpandedScalar) { + Tensor* scalar = view(buffer(1), {M, N}, {0, 0}); + Tensor* a = view(buffer(M * K), {M, K}, {K, 1}); + Tensor* b = view(buffer(K * N), {K, N}, {N, 1}); + Tensor* out_base = buffer(M * N); + Tensor* out = view(out_base, {M, N}, {N, 1}); + ASSERT_EQ(aoti_torch_cuda_addmm_out(out, scalar, a, b, 2.0, 1.0), Error::Ok); + expect_product(out_base, out, a, b, 1, 2, [](auto...) { return pattern(0); }); +} + +TEST_P(AOTITorchCudaMmTest, AddmmTransposedSelf) { + Tensor* self = view(buffer(M * N), {M, N}, {1, M}); + Tensor* a = view(buffer(M * K), {M, K}, {K, 1}); + Tensor* b = view(buffer(K * N), {K, N}, {N, 1}); + Tensor* out_base = buffer(M * N); + Tensor* out = view(out_base, {M, N}, {N, 1}); + ASSERT_EQ(aoti_torch_cuda_addmm_out(out, self, a, b, 1.0, 1.0), Error::Ok); + expect_product(out_base, out, a, b, 1, 1, [](int64_t, int64_t r, int64_t c) { + return pattern(r + c * M); + }); +} + +TEST_P(AOTITorchCudaMmTest, AddmmZeroBetaIgnoresSelf) { + Tensor* self_base = buffer(M * N); + upload(self_base, std::vector(M * N, NAN)); + Tensor* self = view(self_base, {M, N}, {N, 1}); + Tensor* a = view(buffer(M * K), {M, K}, {K, 1}); + Tensor* b = view(buffer(K * N), {K, N}, {N, 1}); + Tensor* out_base = buffer(M * N); + Tensor* out = view(out_base, {M, N}, {N, 1}); + ASSERT_EQ(aoti_torch_cuda_addmm_out(out, self, a, b, 0.0, 1.0), Error::Ok); + expect_product(out_base, out, a, b, 1, 0, [](auto...) { return 0.0f; }); +} + +TEST_P(AOTITorchCudaMmTest, Bmm) { + constexpr int64_t B = 3; + Tensor* a = view(buffer(B * M * K), {B, M, K}, {M * K, K, 1}); + Tensor* b = view(buffer(B * K * N), {B, K, N}, {K * N, 1, K}); + Tensor* out_base = buffer(B * M * N); + Tensor* out = view(out_base, {B, M, N}, {M * N, N, 1}); + ASSERT_EQ(aoti_torch_cuda_bmm_out(out, a, b), Error::Ok); + expect_product(out_base, out, a, b, 1, 0, [](auto...) { return 0.0f; }); +} + +TEST_P(AOTITorchCudaMmTest, BmmBroadcastOperand) { + constexpr int64_t B = 3; + Tensor* a = view(buffer(B * M * K), {B, M, K}, {M * K, K, 1}); + Tensor* b = view(buffer(K * N), {B, K, N}, {0, N, 1}); + Tensor* out_base = buffer(B * M * N); + Tensor* out = view(out_base, {B, M, N}, {M * N, N, 1}); + ASSERT_EQ(aoti_torch_cuda_bmm_out(out, a, b), Error::Ok); + expect_product(out_base, out, a, b, 1, 0, [](auto...) { return 0.0f; }); +} + +// The backend records calls into a CUDA graph after its warm-up run and +// replays them. +TEST_P(AOTITorchCudaMmTest, MmCapturesIntoCudaGraph) { + expect_captured_mm(M, K, N); +} + +// Large enough that cuBLASLt asks for a workspace on an H100, so the graph +// also allocates and frees it. A GPU whose algorithm needs none skips. +TEST_P(AOTITorchCudaMmTest, MmWithWorkspaceCapturesIntoCudaGraph) { + if (dtype() != slim_c10::ScalarType::Float) { + GTEST_SKIP() << "Sums over this K are exact only in float32"; + } + size_t allocations = 0; + expect_captured_mm(16, 8192, 16, &allocations); + if (allocations == 0) { + GTEST_SKIP() << "cuBLASLt used no workspace for this shape on this GPU"; + } +} + +TEST_P(AOTITorchCudaMmTest, RejectsMismatchedShapes) { + Tensor* a = view(buffer(M * K), {M, K}, {K, 1}); + Tensor* b = view(buffer((K + 1) * N), {K + 1, N}, {N, 1}); + Tensor* out = view(buffer(M * N), {M, N}, {N, 1}); + EXPECT_EQ(aoti_torch_cuda_mm_out(out, a, b), Error::InvalidArgument); +} + +TEST_P(AOTITorchCudaMmTest, RejectsOperandThatIsNotAMatrixLayout) { + // Neither dimension is unit-stride. + Tensor* a = view(buffer(2 * M * K), {M, K}, {2 * K, 2}); + Tensor* b = view(buffer(K * N), {K, N}, {N, 1}); + Tensor* out = view(buffer(M * N), {M, N}, {N, 1}); + EXPECT_EQ(aoti_torch_cuda_mm_out(out, a, b), Error::NotSupported); +} + +INSTANTIATE_TEST_SUITE_P( + Dtypes, + AOTITorchCudaMmTest, + ::testing::Values( + slim_c10::ScalarType::Float, + slim_c10::ScalarType::Half, + slim_c10::ScalarType::BFloat16)); + +} // namespace diff --git a/backends/cuda/runtime/targets.bzl b/backends/cuda/runtime/targets.bzl index 5ff2b15867c..cdf03c61223 100644 --- a/backends/cuda/runtime/targets.bzl +++ b/backends/cuda/runtime/targets.bzl @@ -37,6 +37,7 @@ def define_common_targets(is_fbcode = False): "shims/cuda_guard.cpp", "shims/int4mm.cu", "shims/memory.cpp", + "shims/mm.cu", "shims/rand.cu", "shims/sort.cu", "shims/tensor_attribute.cpp", @@ -46,6 +47,7 @@ def define_common_targets(is_fbcode = False): "shims/int4mm.cuh", "shims/int4mm.h", "shims/memory.h", + "shims/mm.h", "shims/rand.h", "shims/sort.h", "shims/tensor_attribute.h", @@ -75,6 +77,7 @@ def define_common_targets(is_fbcode = False): "gcc", ], external_deps = [ + ("cuda", None, "cublasLt-lazy"), ("cuda", None, "cuda-lazy"), ], ) diff --git a/backends/cuda/tests/test_sort_shim.py b/backends/cuda/tests/test_sort_shim.py index 99cba99cfc5..5b4036e118b 100644 --- a/backends/cuda/tests/test_sort_shim.py +++ b/backends/cuda/tests/test_sort_shim.py @@ -41,6 +41,9 @@ "at::_ops::sort_stable::call", "aoti_torch_cuda_sort_stable", "aoti_torch_cuda_randint_low_out", + "aoti_torch_cuda_mm_out", + "aoti_torch_cuda_addmm_out", + "aoti_torch_cuda_bmm_out", } ) From 3865daa774d72f48f0931fe18cf32949a08338ea Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 23 Oct 2026 10:51:44 -0700 Subject: [PATCH 2/2] CUDA backend: let autotuning pick cuBLAS for matrix multiplications Add a max_autotune_gemm_backends compile spec. TRITON stays the default. ATEN,TRITON lets Inductor's autotuner compare cuBLAS with its Triton templates for each matrix multiplication and keep the faster one, and ATEN uses cuBLAS only. A cuBLAS call is a separate kernel that Inductor cannot fuse with the ops around it, so it is opt-in. CudaPartitioner([ CudaBackend.generate_method_name_compile_spec("forward"), CompileSpec("max_autotune_gemm_backends", b"ATEN,TRITON"), ]) ATEN is refused for ROCm and for Windows lowering, which have no cuBLAS shims. --- backends/cuda/cuda_backend.py | 31 +++- backends/cuda/tests/test_cublas_gemm.py | 188 ++++++++++++++++++++++++ 2 files changed, 217 insertions(+), 2 deletions(-) create mode 100644 backends/cuda/tests/test_cublas_gemm.py diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 01f5df8cda1..a2ee898afd3 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -459,6 +459,29 @@ def _on_off_compile_spec_value(spec: CompileSpec) -> bool: return value == "ON" +def _gemm_backends(spec: Optional[CompileSpec], platform: str) -> str: + """Matrix multiplication choices autotuning may pick from: TRITON compiles + them into the model, ATEN calls cuBLAS at run time. + + ATEN is opt-in because a cuBLAS call is a separate kernel that Inductor + cannot fuse with the ops around it. + """ + if spec is None: + return "TRITON" + value = spec.value.decode("utf-8").upper().replace(" ", "") + backends = set(value.split(",")) + if not backends <= {"ATEN", "TRITON"}: + raise ValueError( + f"Invalid {spec.key}: {value}. Expected ATEN, TRITON or ATEN,TRITON." + ) + if "ATEN" in backends and (torch.version.hip is not None or platform == "windows"): + raise ValueError( + f"ATEN in {spec.key} runs matrix multiplications on cuBLAS, which " + "only the Linux CUDA runtime provides." + ) + return value + + CUDA_GRAPH_AUTOTUNE_TIMING_COMPILE_SPEC = "cuda_graph_autotune_timing" # The CUDA tensors `move_program_to_device` placed on the GPU for the compile on @@ -860,8 +883,6 @@ def get_aoti_compile_options( ), # Enable maximum automatic tuning for optimal performance "max_autotune": True, - # Use TRITON for GEMM (General Matrix Multiply) operations tuning only to avoid using operators in libtorch - "max_autotune_gemm_backends": "TRITON", # Use TRITON backend for convolution operations tuning only to avoid using operators in libtorch "max_autotune_conv_backends": "TRITON", "aot_inductor.emit_multi_arch_kernel": emit_multi_arch_kernel, @@ -873,6 +894,7 @@ def get_aoti_compile_options( emulate_precision_casts = True max_autotune = True autotune_at_compile_time = None + gemm_backends_spec = None shim_library_path = None for spec in compile_specs: if spec.key == "platform": @@ -885,6 +907,11 @@ def get_aoti_compile_options( autotune_at_compile_time = _on_off_compile_spec_value(spec) elif spec.key == "shim_library_path": shim_library_path = spec.value.decode("utf-8") + elif spec.key == "max_autotune_gemm_backends": + gemm_backends_spec = spec + options["max_autotune_gemm_backends"] = _gemm_backends( + gemm_backends_spec, platform + ) options["emulate_precision_casts"] = emulate_precision_casts options["max_autotune"] = max_autotune if autotune_at_compile_time is not None: diff --git a/backends/cuda/tests/test_cublas_gemm.py b/backends/cuda/tests/test_cublas_gemm.py new file mode 100644 index 00000000000..5456021d092 --- /dev/null +++ b/backends/cuda/tests/test_cublas_gemm.py @@ -0,0 +1,188 @@ +# 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. + +""" +Matrix multiplications on cuBLAS through the mm, addmm and bmm CUDA shims. + +With max_autotune_gemm_backends=ATEN, Inductor calls those shims instead of +compiling Triton matrix multiplications into the model. + +Usage: + python -m pytest backends/cuda/tests/test_cublas_gemm.py -v +""" + +import glob +import os +import tempfile +import unittest +from unittest import mock + +import torch +from executorch.backends.cuda.cuda_backend import CudaBackend +from executorch.backends.cuda.cuda_partitioner import CudaPartitioner +from executorch.exir import ( + EdgeCompileConfig, + ExecutorchBackendConfig, + to_edge_transform_and_lower, +) +from executorch.exir.backend.compile_spec_schema import CompileSpec +from executorch.exir.passes import MemoryPlanningPass +from torch._inductor.codegen.cpp_wrapper_cpu import CppWrapperCpu + +_CUBLAS_SHIMS = { + "aoti_torch_cuda_mm_out", + "aoti_torch_cuda_addmm_out", + "aoti_torch_cuda_bmm_out", +} +_ATEN = [CompileSpec("max_autotune_gemm_backends", b"ATEN")] + + +class GemmModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.with_bias = torch.nn.Linear(64, 48) + self.without_bias = torch.nn.Linear(48, 32, bias=False) + # Inductor folds this add into addmm, which gets it as a column. + self.row_offsets = torch.nn.Parameter(torch.randn(16, 1)) + self.head = torch.nn.Linear(8, 8, bias=False) + + def forward(self, x, y): + h = self.without_bias(self.with_bias(x)) + self.row_offsets + return self.head(torch.bmm(h.view(4, 4, 32), y)) + + +def _export_and_run(model, inputs, specs): + """Lowers `model` with the given compile specs and runs it on the + ExecuTorch runtime. Returns the outputs and the extern shim calls.""" + specs = [CudaBackend.generate_method_name_compile_spec("forward"), *specs] + called = [] + original = CppWrapperCpu.generate_c_shim_extern_kernel_call + + def record(self, kernel, *args, **kwargs): + called.append(kernel) + return original(self, kernel, *args, **kwargs) + + with mock.patch.object( + CppWrapperCpu, "generate_c_shim_extern_kernel_call", record + ), torch.no_grad(): + program = to_edge_transform_and_lower( + torch.export.export(model, inputs, strict=True), + partitioner=[CudaPartitioner(specs)], + compile_config=EdgeCompileConfig( + _check_ir_validity=False, _skip_dim_order=True + ), + ).to_executorch( + config=ExecutorchBackendConfig( + extract_delegate_segments=True, + memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False), + ) + ) + + from executorch.extension.pybindings.portable_lib import _load_for_executorch + + with tempfile.TemporaryDirectory() as tmpdir: + pte = os.path.join(tmpdir, "model.pte") + with open(pte, "wb") as f: + program.write_to_file(f) + program.write_tensor_data_to_file(tmpdir) + ptd = glob.glob(os.path.join(tmpdir, "*.ptd"))[0] + module = _load_for_executorch(pte, data_path=ptd) + outputs = module.run_method("forward", [t.cpu() for t in inputs]) + return outputs, called + + +class TestCublasGemm(unittest.TestCase): + def setUp(self): + if not torch.cuda.is_available(): + self.skipTest("CUDA is not available") + if torch.version.hip is not None: + self.skipTest("cuBLAS shims are CUDA only") + + def _check(self, dtype, atol, rtol, specs=_ATEN): + torch.manual_seed(0) + model = GemmModel().to(device="cuda", dtype=dtype).eval() + inputs = ( + torch.randn(16, 64, device="cuda", dtype=dtype), + torch.randn(4, 32, 8, device="cuda", dtype=dtype), + ) + with torch.no_grad(): + expected = model(*inputs).cpu() + + outputs, called = _export_and_run(model, inputs, specs) + + # Two addmm calls: one for the bias of with_bias, one for row_offsets + # folded in as a column. + self.assertEqual( + sorted(kernel for kernel in called if kernel in _CUBLAS_SHIMS), + [ + "aoti_torch_cuda_addmm_out", + "aoti_torch_cuda_addmm_out", + "aoti_torch_cuda_bmm_out", + "aoti_torch_cuda_mm_out", + ], + called, + ) + torch.testing.assert_close(outputs[0], expected, atol=atol, rtol=rtol) + + def test_float32(self): + self._check(torch.float32, atol=1e-4, rtol=1e-4) + + def test_bfloat16(self): + if not torch.cuda.is_bf16_supported(): + self.skipTest("bfloat16 is not supported on this GPU") + self._check(torch.bfloat16, atol=5e-2, rtol=2e-2) + + def test_float16(self): + self._check(torch.float16, atol=1e-2, rtol=1e-2) + + def test_without_max_autotune(self): + # Without autotuning there are no Triton templates to choose from. + self._check( + torch.float32, + atol=1e-4, + rtol=1e-4, + specs=[CompileSpec("max_autotune", b"OFF")], + ) + + +class TestGemmBackendsCompileSpec(unittest.TestCase): + def test_default_is_triton(self): + options = CudaBackend.get_aoti_compile_options([]) + self.assertEqual(options["max_autotune_gemm_backends"], "TRITON") + + def test_aten_and_triton(self): + with mock.patch.object(torch.version, "hip", None): + options = CudaBackend.get_aoti_compile_options( + [CompileSpec("max_autotune_gemm_backends", b"aten,triton")] + ) + self.assertEqual(options["max_autotune_gemm_backends"], "ATEN,TRITON") + + def test_unknown_backend_is_rejected(self): + with self.assertRaisesRegex(ValueError, "Invalid max_autotune_gemm_backends"): + CudaBackend.get_aoti_compile_options( + [CompileSpec("max_autotune_gemm_backends", b"ATEN,CUTLASS")] + ) + + def test_aten_is_rejected_on_rocm(self): + with mock.patch.object(torch.version, "hip", "7.0"): + with self.assertRaisesRegex(ValueError, "cuBLAS"): + CudaBackend.get_aoti_compile_options( + [CompileSpec("max_autotune_gemm_backends", b"ATEN")] + ) + + def test_aten_is_rejected_for_windows(self): + with mock.patch.object(torch.version, "hip", None): + with self.assertRaisesRegex(ValueError, "cuBLAS"): + CudaBackend.get_aoti_compile_options( + [ + CompileSpec("platform", b"windows"), + CompileSpec("max_autotune_gemm_backends", b"ATEN"), + ] + ) + + +if __name__ == "__main__": + unittest.main()