diff --git a/cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp b/cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp index 8cf228f2ed..fb7fb9047b 100644 --- a/cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp +++ b/cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp @@ -167,6 +167,23 @@ struct cublastlt_matrix_layout { return cublastlt_matrix_layout{ get_cuda_data_type(), col_major ? rows : cols, col_major ? cols : rows, ld}; } + + template + static inline auto for_strided_batched_matmul(bool col_major, + uint64_t rows, + uint64_t cols, + uint64_t ld, + int32_t batch_count, + int64_t batch_stride) -> cublastlt_matrix_layout + { + RAFT_EXPECTS(batch_count > 0, "cuBLASLt batch count must be positive"); + auto layout = for_matmul(col_major, rows, cols, ld); + RAFT_CUBLAS_TRY(cublasLtMatrixLayoutSetAttribute( + layout, CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &batch_count, sizeof(batch_count))); + RAFT_CUBLAS_TRY(cublasLtMatrixLayoutSetAttribute( + layout, CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, &batch_stride, sizeof(batch_stride))); + return layout; + } }; /** Descriptor for a cublasLt matmul function. */ @@ -197,9 +214,12 @@ struct cublastlt_matmul_desc { inline operator cublasLtMatmulDesc_t() const noexcept { return res; } template - static inline auto for_matmul(bool transpose_a, bool transpose_b) -> cublastlt_matmul_desc + static inline auto for_matmul(bool transpose_a, + bool transpose_b, + cublasComputeType_t compute_type = get_matmul_type()) + -> cublastlt_matmul_desc { - auto desc = cublastlt_matmul_desc{get_matmul_type(), get_cuda_data_type()}; + auto desc = cublastlt_matmul_desc{compute_type, get_cuda_data_type()}; if constexpr (DevicePointerMode) { const cublasPointerMode_t mode = CUBLAS_POINTER_MODE_DEVICE; RAFT_CUBLAS_TRY(cublasLtMatmulDescSetAttribute( @@ -348,6 +368,70 @@ struct coef_wrapper { } }; +/** + * Run a strided-batched cublasLt matmul without an algorithm or workspace preference. + * + * The matrix dimensions and leading dimensions follow the same column-major convention as + * `matmul`. Batch strides are measured in elements. + */ +template +void matmul_strided_batched(raft::resources const& res, + bool trans_a, + bool trans_b, + uint64_t m, + uint64_t n, + uint64_t k, + const S* alpha, + const A* a_ptr, + uint64_t lda, + int64_t stride_a, + const B* b_ptr, + uint64_t ldb, + int64_t stride_b, + const S* beta, + C* c_ptr, + uint64_t ldc, + int64_t stride_c, + int32_t batch_count, + cublasComputeType_t compute_type) +{ + common::nvtx::range batch_scope( + "linalg::detail::matmul_strided_batched(m = %d, n = %d, k = %d, batch_count = %d)", + m, + n, + k, + batch_count); + + auto operation = cublastlt_matmul_desc::for_matmul( + trans_a, trans_b, compute_type); + + auto a_layout = cublastlt_matrix_layout::for_strided_batched_matmul( + !trans_a, m, k, lda, batch_count, stride_a); + auto b_layout = cublastlt_matrix_layout::for_strided_batched_matmul( + !trans_b, k, n, ldb, batch_count, stride_b); + auto c_layout = + cublastlt_matrix_layout::for_strided_batched_matmul(true, m, n, ldc, batch_count, stride_c); + + auto stream = resource::get_cuda_stream(res); + coef_wrapper coefficients(alpha, beta, stream); + RAFT_CUBLAS_TRY(cublasLtMatmul(resource::get_cublaslt_handle(res), + operation, + coefficients.alpha, + a_ptr, + a_layout, + b_ptr, + b_layout, + coefficients.beta, + c_ptr, + c_layout, + c_ptr, + c_layout, + nullptr, + nullptr, + 0, + stream)); +} + /** * Compatibility version of the cublasLt matmul wrapper: It takes the cudaStream_t argument * explicitly rather than through the raft::resources. This function is used by other legacy diff --git a/cpp/include/raft/linalg/gemm.cuh b/cpp/include/raft/linalg/gemm.cuh index ada9dfe75f..80c083d682 100644 --- a/cpp/include/raft/linalg/gemm.cuh +++ b/cpp/include/raft/linalg/gemm.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #ifndef __GEMM_H @@ -283,6 +283,186 @@ void gemm(raft::resources const& res, } } +/** + * @brief Batched matrix multiplication with an explicit cublasLt compute type. + * + * Computes `Z_i = alpha * X_i * Y_i + beta * Z_i`. The first mdspan extent is the batch + * dimension. `layout_stride` views encode matrix and batch strides in their mappings; contiguous + * batches use `layout_right`, with the matrix size as the batch stride. Use `layout_stride` for + * column-major batches because rank-3 `layout_left` interleaves the batch dimension. + * + * @tparam ValueType element type of the input and output matrices + * @tparam IndexType index type + * @tparam LayoutPolicyX layout policy of X + * @tparam LayoutPolicyY layout policy of Y + * @tparam LayoutPolicyZ layout policy of Z + * @tparam ScalarIdxType index type of alpha and beta + * @tparam ScalarViewType scalar view type of alpha and beta + * @param[in] res RAFT resources + * @param[in] x input matrices X, with shape `[batch_count, m, k]` + * @param[in] y input matrices Y, with shape `[batch_count, k, n]` + * @param[inout] z input/output matrices Z, with shape `[batch_count, m, n]` + * @param[in] alpha scalar multiplier for X * Y; defaults to one when empty + * @param[in] beta scalar multiplier for Z; defaults to zero when empty + * @param[in] compute_type_override cublasLt compute type + */ +template , + typename = std::enable_if_t>, + std::is_same>>>> +void gemm_batched(raft::resources const& res, + raft::device_mdspan, LayoutPolicyX> x, + raft::device_mdspan, LayoutPolicyY> y, + raft::device_mdspan, LayoutPolicyZ> z, + std::optional alpha, + std::optional beta, + cublasComputeType_t compute_type_override) +{ + static_assert(!std::is_same_v && + !std::is_same_v && + !std::is_same_v, + "Use layout_stride for column-major batched matrices"); + + RAFT_EXPECTS(x.extent(0) == y.extent(0) && x.extent(0) == z.extent(0), + "All operands must have the same batch count"); + RAFT_EXPECTS(x.extent(1) == z.extent(1), "X and Z must have the same number of rows"); + RAFT_EXPECTS(y.extent(2) == z.extent(2), "Y and Z must have the same number of columns"); + RAFT_EXPECTS(x.extent(2) == y.extent(1), + "The number of columns in X must equal the number of rows in Y"); + + const auto is_row_major = [](auto const& view) { + return view.stride(2) == 1 && view.stride(1) >= view.extent(2); + }; + const auto is_col_major = [](auto const& view) { + return view.stride(1) == 1 && view.stride(2) >= view.extent(1); + }; + const auto batch_stride = [](auto const& view) -> int64_t { + using layout_type = typename std::decay_t::layout_type; + if constexpr (std::is_same_v) { + return static_cast(view.stride(0)); + } else { + return static_cast(view.extent(1)) * static_cast(view.extent(2)); + } + }; + RAFT_EXPECTS(is_row_major(x) || is_col_major(x), + "Each matrix in X must be row-major or column-major"); + RAFT_EXPECTS(is_row_major(y) || is_col_major(y), + "Each matrix in Y must be row-major or column-major"); + RAFT_EXPECTS(is_row_major(z) || is_col_major(z), + "Each matrix in Z must be row-major or column-major"); + + constexpr auto kDeviceMode = + std::is_same_v>; + ValueType* alpha_ptr = nullptr; + ValueType* beta_ptr = nullptr; + if (alpha.has_value()) { alpha_ptr = alpha->data_handle(); } + if (beta.has_value()) { beta_ptr = beta->data_handle(); } + + const bool x_col_major = is_col_major(x); + const bool y_col_major = is_col_major(y); + const bool z_col_major = is_col_major(z); + const auto x_ld = x_col_major ? x.stride(2) : x.stride(1); + const auto y_ld = y_col_major ? y.stride(2) : y.stride(1); + const auto z_ld = z_col_major ? z.stride(2) : z.stride(1); + const auto batch_count = static_cast(z.extent(0)); + + if (z_col_major) { + return detail::matmul_strided_batched(res, + !x_col_major, + !y_col_major, + z.extent(1), + z.extent(2), + x.extent(2), + alpha_ptr, + x.data_handle(), + x_ld, + batch_stride(x), + y.data_handle(), + y_ld, + batch_stride(y), + beta_ptr, + z.data_handle(), + z_ld, + batch_stride(z), + batch_count, + compute_type_override); + } + + return detail::matmul_strided_batched(res, + y_col_major, + x_col_major, + z.extent(2), + z.extent(1), + x.extent(2), + alpha_ptr, + y.data_handle(), + y_ld, + batch_stride(y), + x.data_handle(), + x_ld, + batch_stride(x), + beta_ptr, + z.data_handle(), + z_ld, + batch_stride(z), + batch_count, + compute_type_override); +} + +/** + * @brief Batched matrix multiplication using the default cublasLt compute type. + * + * Computes `Z_i = alpha * X_i * Y_i + beta * Z_i`. The first mdspan extent is the batch + * dimension. `layout_stride` views encode matrix and batch strides in their mappings; contiguous + * batches use `layout_right`, with the matrix size as the batch stride. Use `layout_stride` for + * column-major batches because rank-3 `layout_left` interleaves the batch dimension. + * + * @tparam ValueType element type of the input and output matrices + * @tparam IndexType index type + * @tparam LayoutPolicyX layout policy of X + * @tparam LayoutPolicyY layout policy of Y + * @tparam LayoutPolicyZ layout policy of Z + * @tparam ScalarIdxType index type of alpha and beta + * @tparam ScalarViewType scalar view type of alpha and beta + * @param[in] res RAFT resources + * @param[in] x input matrices X, with shape `[batch_count, m, k]` + * @param[in] y input matrices Y, with shape `[batch_count, k, n]` + * @param[inout] z input/output matrices Z, with shape `[batch_count, m, n]` + * @param[in] alpha scalar multiplier for X * Y; defaults to one when empty + * @param[in] beta scalar multiplier for Z; defaults to zero when empty + */ +template , + typename = std::enable_if_t>, + std::is_same>>>> +void gemm_batched(raft::resources const& res, + raft::device_mdspan, LayoutPolicyX> x, + raft::device_mdspan, LayoutPolicyY> y, + raft::device_mdspan, LayoutPolicyZ> z, + std::optional alpha = std::nullopt, + std::optional beta = std::nullopt) +{ + return gemm_batched(res, + x, + y, + z, + alpha, + beta, + detail::get_matmul_type()); +} + /** @} */ // end of gemm } // namespace linalg diff --git a/cpp/tests/linalg/gemm_basic.cpp b/cpp/tests/linalg/gemm_basic.cpp index 8622eec113..8550cb871b 100644 --- a/cpp/tests/linalg/gemm_basic.cpp +++ b/cpp/tests/linalg/gemm_basic.cpp @@ -163,6 +163,104 @@ TEST(Raft, GemmPointerModeDeviceAlpha) { test_gemm_pointer_mode_device(true, fal TEST(Raft, GemmPointerModeDeviceBeta) { test_gemm_pointer_mode_device(false, true); } TEST(Raft, GemmPointerModeDeviceDefaults) { test_gemm_pointer_mode_device(false, false); } +TEST(Raft, GemmBatched) +{ + raft::resources res; + auto stream = raft::resource::get_cuda_stream(res); + + using index_type = int64_t; + constexpr index_type stride_a = 8; + constexpr index_type stride_b = 8; + constexpr index_type stride_c = 5; + constexpr index_type batch_count = 2; + + // Two column-major A (2 x 3) and B (3 x 2) matrices with padding between batches. + std::vector a_host = {1, 4, 2, 5, 3, 6, -1, -1, 2, 1, 0, 3, 1, 4, -1, -1}; + std::vector b_host = {7, 9, 11, 8, 10, 12, -1, -1, 1, 0, 2, 3, 1, 4, -1, -1}; + std::vector c_host(stride_c * batch_count, -1); + + auto a_device = raft::make_device_vector(res, a_host.size()); + auto b_device = raft::make_device_vector(res, b_host.size()); + auto c_device = raft::make_device_vector(res, c_host.size()); + raft::copy(a_device.data_handle(), a_host.data(), a_host.size(), stream); + raft::copy(b_device.data_handle(), b_host.data(), b_host.size(), stream); + raft::copy(c_device.data_handle(), c_host.data(), c_host.size(), stream); + + auto a_extents = raft::extent_3d{batch_count, M, K}; + auto b_extents = raft::extent_3d{batch_count, K, N}; + auto c_extents = raft::extent_3d{batch_count, M, N}; + auto a_layout = + raft::make_strided_layout(a_extents, cuda::std::array{stride_a, 1, M}); + auto b_layout = + raft::make_strided_layout(b_extents, cuda::std::array{stride_b, 1, K}); + auto c_layout = + raft::make_strided_layout(c_extents, cuda::std::array{stride_c, 1, M}); + + auto a_view = raft::device_mdspan{ + a_device.data_handle(), a_layout}; + auto b_view = raft::device_mdspan{ + b_device.data_handle(), b_layout}; + auto c_view = raft::device_mdspan{ + c_device.data_handle(), c_layout}; + + std::optional> alpha; + std::optional> beta; + raft::linalg::gemm_batched( + res, a_view, b_view, c_view, alpha, beta, CUBLAS_COMPUTE_32F_FAST_TF32); + + raft::copy(c_host.data(), c_device.data_handle(), c_host.size(), stream); + raft::resource::sync_stream(res); + + const std::vector expected = {58, 139, 64, 154, -1, 4, 9, 10, 22, -1}; + EXPECT_EQ(c_host, expected); + + c_host.assign(stride_c * batch_count, -1); + raft::copy(c_device.data_handle(), c_host.data(), c_host.size(), stream); + raft::linalg::gemm_batched(res, a_view, b_view, c_view); + + raft::copy(c_host.data(), c_device.data_handle(), c_host.size(), stream); + raft::resource::sync_stream(res); + EXPECT_EQ(c_host, expected); +} + +TEST(Raft, GemmBatchedContiguous) +{ + raft::resources res; + auto stream = raft::resource::get_cuda_stream(res); + + using index_type = int64_t; + constexpr index_type batch_count = 2; + + std::vector a_host = {1, 2, 3, 4, 5, 6, 2, 0, 1, 1, 3, 4}; + std::vector b_host = {7, 8, 9, 10, 11, 12, 1, 3, 0, 1, 2, 4}; + std::vector c_host(batch_count * M * N, -1); + + auto a_device = raft::make_device_vector(res, a_host.size()); + auto b_device = raft::make_device_vector(res, b_host.size()); + auto c_device = raft::make_device_vector(res, c_host.size()); + raft::copy(a_device.data_handle(), a_host.data(), a_host.size(), stream); + raft::copy(b_device.data_handle(), b_host.data(), b_host.size(), stream); + raft::copy(c_device.data_handle(), c_host.data(), c_host.size(), stream); + + auto a_extents = raft::extent_3d{batch_count, M, K}; + auto b_extents = raft::extent_3d{batch_count, K, N}; + auto c_extents = raft::extent_3d{batch_count, M, N}; + auto a_view = raft::device_mdspan{ + a_device.data_handle(), a_extents}; + auto b_view = raft::device_mdspan{ + b_device.data_handle(), b_extents}; + auto c_view = raft::device_mdspan{ + c_device.data_handle(), c_extents}; + + raft::linalg::gemm_batched(res, a_view, b_view, c_view); + + raft::copy(c_host.data(), c_device.data_handle(), c_host.size(), stream); + raft::resource::sync_stream(res); + + const std::vector expected = {58, 64, 139, 154, 4, 10, 9, 22}; + EXPECT_EQ(c_host, expected); +} + TEST(Raft, GemmCublasLt136WorkaroundPredicate) { constexpr std::size_t affected_version = 130600;