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..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 @@ -773,6 +796,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 @@ -857,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, @@ -870,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": @@ -882,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/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_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() 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", } )