From f114b302df212e395ce2d024cfa2c4ff12d29b4d Mon Sep 17 00:00:00 2001 From: Intron7 Date: Tue, 11 Aug 2026 18:34:24 +0200 Subject: [PATCH] add-batched-linalg Signed-off-by: Intron7 --- .../raft/linalg/detail/cublas_wrappers.hpp | 65 +++- .../raft/linalg/detail/cublaslt_wrappers.hpp | 191 ++++++++-- .../raft/linalg/detail/cusolver_wrappers.hpp | 45 ++- cpp/include/raft/linalg/detail/gemm.cuh | 51 ++- cpp/include/raft/linalg/gemm.cuh | 120 +++++- cpp/tests/CMakeLists.txt | 1 + cpp/tests/linalg/gemm_batched.cpp | 349 ++++++++++++++++++ 7 files changed, 783 insertions(+), 39 deletions(-) create mode 100644 cpp/tests/linalg/gemm_batched.cpp diff --git a/cpp/include/raft/linalg/detail/cublas_wrappers.hpp b/cpp/include/raft/linalg/detail/cublas_wrappers.hpp index d08d4daf0d..e4873402af 100644 --- a/cpp/include/raft/linalg/detail/cublas_wrappers.hpp +++ b/cpp/include/raft/linalg/detail/cublas_wrappers.hpp @@ -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 */ @@ -930,6 +930,69 @@ inline cublasStatus_t cublastrsm(cublasHandle_t handle, return cublasDtrsm(handle, side, uplo, trans, diag, m, n, alpha, A, lda, B, ldb); } +/** + * @defgroup trsmBatched cublas trsmBatched calls + * @{ + */ +template +cublasStatus_t cublastrsmBatched(cublasHandle_t handle, // NOLINT + cublasSideMode_t side, + cublasFillMode_t uplo, + cublasOperation_t trans, + cublasDiagType_t diag, + int m, + int n, + const T* alpha, + const T* const Aarray[], // NOLINT + int lda, + T* const Barray[], // NOLINT + int ldb, + int batchCount, + cudaStream_t stream); + +template <> +inline cublasStatus_t cublastrsmBatched(cublasHandle_t handle, // NOLINT + cublasSideMode_t side, + cublasFillMode_t uplo, + cublasOperation_t trans, + cublasDiagType_t diag, + int m, + int n, + const float* alpha, + const float* const Aarray[], // NOLINT + int lda, + float* const Barray[], // NOLINT + int ldb, + int batchCount, + cudaStream_t stream) +{ + RAFT_CUBLAS_TRY(cublasSetStream(handle, stream)); + return cublasStrsmBatched( + handle, side, uplo, trans, diag, m, n, alpha, Aarray, lda, Barray, ldb, batchCount); +} + +template <> +inline cublasStatus_t cublastrsmBatched(cublasHandle_t handle, // NOLINT + cublasSideMode_t side, + cublasFillMode_t uplo, + cublasOperation_t trans, + cublasDiagType_t diag, + int m, + int n, + const double* alpha, + const double* const Aarray[], // NOLINT + int lda, + double* const Barray[], // NOLINT + int ldb, + int batchCount, + cudaStream_t stream) +{ + RAFT_CUBLAS_TRY(cublasSetStream(handle, stream)); + return cublasDtrsmBatched( + handle, side, uplo, trans, diag, m, n, alpha, Aarray, lda, Barray, ldb, batchCount); +} +/** @} */ + /** * @defgroup dot cublas dot calls * @{ diff --git a/cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp b/cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp index 8cf228f2ed..11f736c1d0 100644 --- a/cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp +++ b/cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp @@ -85,18 +85,28 @@ struct matmul_key_t { uint64_t ldc; bool trans_a; bool trans_b; + /** Number of matrices in the batch; 1 means a plain, non-batched matmul. */ + uint64_t batch_count = 1; + /** Offsets in elements between consecutive matrices of the batch (ignored if batch_count == 1) */ + int64_t stride_a = 0; + int64_t stride_b = 0; + int64_t stride_c = 0; }; inline auto operator==(const matmul_key_t& a, const matmul_key_t& b) -> bool { return a.m == b.m && a.n == b.n && a.k == b.k && a.lda == b.lda && a.ldb == b.ldb && - a.ldc == b.ldc && a.trans_a == b.trans_a && a.trans_b == b.trans_b; + a.ldc == b.ldc && a.trans_a == b.trans_a && a.trans_b == b.trans_b && + a.batch_count == b.batch_count && a.stride_a == b.stride_a && a.stride_b == b.stride_b && + a.stride_c == b.stride_c; } struct matmul_key_hash { inline auto operator()(const matmul_key_t& x) const noexcept -> std::size_t { - return x.m * x.n * x.k + x.lda * x.ldb * x.ldc + size_t{x.trans_a} + size_t{x.trans_b} * 2; + return x.m * x.n * x.k + x.lda * x.ldb * x.ldc + size_t{x.trans_a} + size_t{x.trans_b} * 2 + + x.batch_count * (static_cast(x.stride_a) + static_cast(x.stride_b) + + static_cast(x.stride_c)); } }; @@ -160,12 +170,32 @@ struct cublastlt_matrix_layout { // NOLINTNEXTLINE inline operator cublasLtMatrixLayout_t() const noexcept { return res; } + /** + * Describe the matrix as a batch of `batch_count` matrices, `batch_stride` elements apart. + * A `batch_count` of one leaves the layout as a plain, non-batched matrix. + */ + inline void set_batch(uint64_t batch_count, int64_t batch_stride) + { + if (batch_count <= 1) { return; } + const auto count = static_cast(batch_count); + RAFT_CUBLAS_TRY(cublasLtMatrixLayoutSetAttribute( + res, CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &count, sizeof(count))); + RAFT_CUBLAS_TRY(cublasLtMatrixLayoutSetAttribute( + res, CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, &batch_stride, sizeof(batch_stride))); + } + template - static inline auto for_matmul(bool col_major, uint64_t rows, uint64_t cols, uint64_t ld) - -> cublastlt_matrix_layout + static inline auto for_matmul(bool col_major, + uint64_t rows, + uint64_t cols, + uint64_t ld, + uint64_t batch_count = 1, + int64_t batch_stride = 0) -> cublastlt_matrix_layout { - return cublastlt_matrix_layout{ + auto r = cublastlt_matrix_layout{ get_cuda_data_type(), col_major ? rows : cols, col_major ? cols : rows, ld}; + r.set_batch(batch_count, batch_stride); + return r; } }; @@ -249,9 +279,12 @@ struct matmul_desc { { matmul_desc r{ cublastlt_matmul_desc::for_matmul(args.trans_a, args.trans_b), - cublastlt_matrix_layout::for_matmul(!(args.trans_a), args.m, args.k, args.lda), - cublastlt_matrix_layout::for_matmul(!(args.trans_b), args.k, args.n, args.ldb), - cublastlt_matrix_layout::for_matmul(true, args.m, args.n, args.ldc)}; + cublastlt_matrix_layout::for_matmul( + !(args.trans_a), args.m, args.k, args.lda, args.batch_count, args.stride_a), + cublastlt_matrix_layout::for_matmul( + !(args.trans_b), args.k, args.n, args.ldb, args.batch_count, args.stride_b), + cublastlt_matrix_layout::for_matmul( + true, args.m, args.n, args.ldc, args.batch_count, args.stride_c)}; bool use_cublaslt_13_6_workaround = false; if constexpr (std::is_same_v && std::is_same_v && @@ -277,8 +310,12 @@ struct matmul_desc { if (use_cublaslt_13_6_workaround) { const auto heuristic_args = get_cublaslt_13_6_heuristic_args(args); - const auto heuristic_a = cublastlt_matrix_layout::for_matmul( - !(heuristic_args.trans_a), heuristic_args.m, heuristic_args.k, heuristic_args.lda); + const auto heuristic_a = cublastlt_matrix_layout::for_matmul(!(heuristic_args.trans_a), + heuristic_args.m, + heuristic_args.k, + heuristic_args.lda, + heuristic_args.batch_count, + heuristic_args.stride_a); query_heuristic(heuristic_a, r.c); } else { query_heuristic(r.a, r.c); @@ -348,6 +385,47 @@ struct coef_wrapper { } }; +/** + * Shared implementation behind all cublasLt matmul wrappers: look up (or create and cache) the + * matmul descriptor for `mm_key` and run it. Batching, if any, is described by `mm_key`. + */ +template +void matmul_impl(raft::resources const& res, + const matmul_key_t& mm_key, + const S* alpha, + const A* a_ptr, + const B* b_ptr, + const S* beta, + C* c_ptr, + cudaStream_t stream) +{ + std::shared_ptr mm_desc{nullptr}; + auto& cache = + resource::get_custom_resource>(res)->value; + if (!cache.get(mm_key, &mm_desc)) { + mm_desc.reset(new matmul_desc{matmul_desc::create(res, mm_key)}); + cache.set(mm_key, mm_desc); + } + // Allocate alpha and beta pointers if not provided. + coef_wrapper w(alpha, beta, stream); + RAFT_CUBLAS_TRY(cublasLtMatmul(resource::get_cublaslt_handle(res), + mm_desc->desc, + w.alpha, + a_ptr, + mm_desc->a, + b_ptr, + mm_desc->b, + w.beta, + c_ptr, + mm_desc->c, + c_ptr, + mm_desc->c, + &(mm_desc->heuristics.algo), + 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 @@ -376,32 +454,8 @@ template batch_scope( "linalg::matmul(m = %d, n = %d, k = %d)", m, n, k); - std::shared_ptr mm_desc{nullptr}; matmul_key_t mm_key{m, n, k, lda, ldb, ldc, trans_a, trans_b}; - auto& cache = - resource::get_custom_resource>(res)->value; - if (!cache.get(mm_key, &mm_desc)) { - mm_desc.reset(new matmul_desc{matmul_desc::create(res, mm_key)}); - cache.set(mm_key, mm_desc); - } - // Allocate alpha and beta pointers if not provided. - coef_wrapper w(alpha, beta, stream); - RAFT_CUBLAS_TRY(cublasLtMatmul(resource::get_cublaslt_handle(res), - mm_desc->desc, - w.alpha, - a_ptr, - mm_desc->a, - b_ptr, - mm_desc->b, - w.beta, - c_ptr, - mm_desc->c, - c_ptr, - mm_desc->c, - &(mm_desc->heuristics.algo), - nullptr, - 0, - stream)); + matmul_impl(res, mm_key, alpha, a_ptr, b_ptr, beta, c_ptr, stream); } /** @@ -462,5 +516,72 @@ void matmul(raft::resources const& res, resource::get_cuda_stream(res)); } +/** + * @brief the wrapper of the strided-batched cublasLt matmul function + * For every batch index i it computes: + * C_i = alpha .* opA(A_i) * opB(B_i) + beta .* C_i + * where X_i is the matrix starting at `x_ptr + i * stride_x`. + * + * All matrices of a batch share the same shape, leading dimension and transpose op; only the + * base pointers differ. A stride of zero broadcasts the same matrix over the whole batch. + * + * @tparam DevicePointerMode whether pointers alpha, beta point to device memory + * @tparam S the type of scale parameters alpha, beta + * @tparam A the element type of matrix A + * @tparam B the element type of matrix B + * @tparam C the element type of matrix C + * + * @param [in] res raft resources + * @param [in] trans_a cublas transpose op for A + * @param [in] trans_b cublas transpose op for B + * @param [in] m number of rows of C + * @param [in] n number of columns of C + * @param [in] k number of rows of opB(B) / number of columns of opA(A) + * @param [in] alpha host or device scalar, if nullptr, the default value 1 will be used + * @param [in] a_ptr such a matrix that the shape of column-major opA(A) is [m, k] + * @param [in] lda leading dimension of A + * @param [in] stride_a offset in elements between consecutive matrices of A + * @param [in] b_ptr such a matrix that the shape of column-major opA(B) is [k, n] + * @param [in] ldb leading dimension of B + * @param [in] stride_b offset in elements between consecutive matrices of B + * @param [in] beta host or device scalar, if nullptr, the default value 0 will be used + * @param [inout] c_ptr column-major matrix of size [m, n] + * @param [in] ldc leading dimension of C + * @param [in] stride_c offset in elements between consecutive matrices of C + * @param [in] batch_count number of matrices in the batch + */ +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, + uint64_t batch_count) +{ + common::nvtx::range batch_scope( + "linalg::matmul_strided_batched(m = %d, n = %d, k = %d, batch_count = %d)", + m, + n, + k, + batch_count); + if (batch_count == 0) { return; } + matmul_key_t mm_key{ + m, n, k, lda, ldb, ldc, trans_a, trans_b, batch_count, stride_a, stride_b, stride_c}; + matmul_impl( + res, mm_key, alpha, a_ptr, b_ptr, beta, c_ptr, resource::get_cuda_stream(res)); +} + } // namespace linalg::detail } // namespace raft diff --git a/cpp/include/raft/linalg/detail/cusolver_wrappers.hpp b/cpp/include/raft/linalg/detail/cusolver_wrappers.hpp index 6649d67fe9..cbceccd5dd 100644 --- a/cpp/include/raft/linalg/detail/cusolver_wrappers.hpp +++ b/cpp/include/raft/linalg/detail/cusolver_wrappers.hpp @@ -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 */ @@ -877,6 +877,49 @@ inline cusolverStatus_t cusolverDnpotrf(cusolverDnHandle_t handle, // NOLINT } /** @} */ +/** + * @defgroup potrfBatched cusolver potrfBatched operations + * @{ + */ +template +cusolverStatus_t cusolverDnpotrfBatched(cusolverDnHandle_t handle, // NOLINT + cublasFillMode_t uplo, + int n, + T* Aarray[], // NOLINT + int lda, + int* infoArray, + int batchSize, + cudaStream_t stream); + +template <> +inline cusolverStatus_t cusolverDnpotrfBatched(cusolverDnHandle_t handle, // NOLINT + cublasFillMode_t uplo, + int n, + float* Aarray[], // NOLINT + int lda, + int* infoArray, + int batchSize, + cudaStream_t stream) +{ + RAFT_CUSOLVER_TRY(cusolverDnSetStream(handle, stream)); + return cusolverDnSpotrfBatched(handle, uplo, n, Aarray, lda, infoArray, batchSize); +} + +template <> +inline cusolverStatus_t cusolverDnpotrfBatched(cusolverDnHandle_t handle, // NOLINT + cublasFillMode_t uplo, + int n, + double* Aarray[], // NOLINT + int lda, + int* infoArray, + int batchSize, + cudaStream_t stream) +{ + RAFT_CUSOLVER_TRY(cusolverDnSetStream(handle, stream)); + return cusolverDnDpotrfBatched(handle, uplo, n, Aarray, lda, infoArray, batchSize); +} +/** @} */ + /** * @defgroup potrs cusolver potrs operations * @{ diff --git a/cpp/include/raft/linalg/detail/gemm.cuh b/cpp/include/raft/linalg/detail/gemm.cuh index 0cf0a37549..756f793a8b 100644 --- a/cpp/include/raft/linalg/detail/gemm.cuh +++ b/cpp/include/raft/linalg/detail/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 */ #pragma once @@ -7,11 +7,60 @@ #include "cublaslt_wrappers.hpp" #include +#include +#include +#include #include namespace raft { namespace linalg::detail { +/** Description of one operand of a batched gemm in cublas (column-major) terms. */ +struct batched_gemm_operand { + /** Whether every matrix of the batch is column-major (row-major otherwise). */ + bool col_major; + /** Leading dimension of every matrix of the batch. */ + uint64_t ld; + /** Offset in elements between consecutive matrices of the batch. */ + int64_t batch_stride; +}; + +/** + * Interpret a 3D mdspan as a batch of matrices: the first (slowest-varying) dimension indexes the + * batch, the two remaining dimensions form a row- or column-major matrix. The batch stride is + * unconstrained: a stride of zero broadcasts a single matrix over the whole batch. + */ +template +auto describe_batched_gemm_operand( + raft::device_mdspan, LayoutPolicy> x, const char* name) + -> batched_gemm_operand +{ + const auto rows = static_cast(x.extent(1)); + const auto cols = static_cast(x.extent(2)); + const auto row_stride = static_cast(x.stride(1)); + const auto col_stride = static_cast(x.stride(2)); + + batched_gemm_operand r{}; + if (col_stride == 1 && row_stride >= cols) { + r.col_major = false; + r.ld = row_stride; + } else if (row_stride == 1 && col_stride >= rows) { + r.col_major = true; + r.ld = col_stride; + } else { + RAFT_FAIL( + "%s is not a batch of row- or column-major matrices: with extents [batch, %zu, %zu] the " + "matrix strides are [%zu, %zu], one of which must be 1", + name, + static_cast(rows), + static_cast(cols), + static_cast(row_stride), + static_cast(col_stride)); + } + r.batch_stride = static_cast(x.stride(0)); + return r; +} + template void legacy_gemm(raft::resources const& res, const bool trans_a, diff --git a/cpp/include/raft/linalg/gemm.cuh b/cpp/include/raft/linalg/gemm.cuh index ada9dfe75f..058a623c66 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 @@ -285,6 +285,124 @@ void gemm(raft::resources const& res, /** @} */ // end of gemm +/** + * @defgroup gemm_batched Batched Matrix-Matrix Multiplication + * @{ + */ + +/** + * @brief Batched GEMM: performs one matrix-matrix multiplication per batch index in a single + * kernel launch. For every batch index b it computes: + * Z[b] = alpha . X[b] * Y[b] + beta . Z[b] + * If alpha is not provided, it is assumed to be 1.0 + * If beta is not provided, it is assumed to be 0.0 + * + * The first (slowest-varying) dimension of every operand indexes the batch; the two remaining + * dimensions form a raft::row_major or raft::col_major matrix, and all matrices of an operand + * share their layout and leading dimension. Batches of row-major matrices are what a plain + * `raft::row_major` 3D mdspan describes; other packings (batches of column-major matrices, or a + * batch stride that differs from the matrix size) are expressed with a strided layout. A batch + * stride of zero broadcasts a single matrix over the whole batch, which is useful when e.g. the + * same operator is applied to every matrix of the batch. + * + * @tparam ValueType Data type of input/output matrices (float/double) + * @tparam IndexType Type of index + * @tparam LayoutPolicyX layout of X + * @tparam LayoutPolicyY layout of Y + * @tparam LayoutPolicyZ layout of Z + * @param[in] res raft handle + * @param[in] x input raft::device_mdspan of size [batch_count, M, K] + * @param[in] y input raft::device_mdspan of size [batch_count, K, N] + * @param[out] z output raft::device_mdspan of size [batch_count, M, N] + * @param[in] alpha optional raft::host_scalar_view or raft::device_scalar_view, default 1.0 + * @param[in] beta optional raft::host_scalar_view or raft::device_scalar_view, default 0.0 + */ +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) +{ + RAFT_EXPECTS(x.extent(0) == z.extent(0) && y.extent(0) == z.extent(0), + "Batch sizes of X, Y and Z should be equal"); + RAFT_EXPECTS(x.extent(1) == z.extent(1), "Number of rows of X and Z should be equal"); + RAFT_EXPECTS(y.extent(2) == z.extent(2), "Number of columns of Y and Z should be equal"); + RAFT_EXPECTS(x.extent(2) == y.extent(1), "Number of columns of X and rows of Y should be equal"); + + const auto x_desc = detail::describe_batched_gemm_operand(x, "X"); + const auto y_desc = detail::describe_batched_gemm_operand(y, "Y"); + const auto z_desc = detail::describe_batched_gemm_operand(z, "Z"); + + // NB: the function type constraints only ever allow two view types, so using std::is_same_v is + // fine + constexpr auto kDeviceMode = + std::is_same_v>; + + // NB: we rely on the implementation of detail::matmul_strided_batched to set defaults + ValueType* alpha_ptr = nullptr; + ValueType* beta_ptr = nullptr; + if (alpha.has_value()) { alpha_ptr = alpha.value().data_handle(); } + if (beta.has_value()) { beta_ptr = beta.value().data_handle(); } + + const auto batch_count = static_cast(z.extent(0)); + + if (z_desc.col_major) { + return detail::matmul_strided_batched( + res, + !x_desc.col_major, + !y_desc.col_major, + static_cast(z.extent(1)), + static_cast(z.extent(2)), + static_cast(x.extent(2)), + alpha_ptr, + x.data_handle(), + x_desc.ld, + x_desc.batch_stride, + y.data_handle(), + y_desc.ld, + y_desc.batch_stride, + beta_ptr, + z.data_handle(), + z_desc.ld, + z_desc.batch_stride, + batch_count); + } else { + // Z is row-major, i.e. Zᵀ is column-major: compute Zᵀ[b] = Yᵀ[b] * Xᵀ[b] instead. + return detail::matmul_strided_batched( + res, + y_desc.col_major, + x_desc.col_major, + static_cast(z.extent(2)), + static_cast(z.extent(1)), + static_cast(x.extent(2)), + alpha_ptr, + y.data_handle(), + y_desc.ld, + y_desc.batch_stride, + x.data_handle(), + x_desc.ld, + x_desc.batch_stride, + beta_ptr, + z.data_handle(), + z_desc.ld, + z_desc.batch_stride, + batch_count); + } +} + +/** @} */ // end of gemm_batched + } // namespace linalg } // namespace raft diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index fce5a0ea47..3224cc3367 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -155,6 +155,7 @@ if(BUILD_TESTS) linalg/eig.cu linalg/eig_sel.cu linalg/gemm_basic.cpp + linalg/gemm_batched.cpp linalg/gemm_layout.cu linalg/gemv.cu linalg/map.cu diff --git a/cpp/tests/linalg/gemm_batched.cpp b/cpp/tests/linalg/gemm_batched.cpp new file mode 100644 index 0000000000..bbbe3149ee --- /dev/null +++ b/cpp/tests/linalg/gemm_batched.cpp @@ -0,0 +1,349 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include + +#include + +namespace raft::linalg { + +namespace { + +constexpr int kBatch = 3; +constexpr int kM = 4; +constexpr int kK = 5; +constexpr int kN = 2; + +constexpr float kAlpha = 2.0f; +constexpr float kBeta = 3.0f; + +/** Deterministic small-integer values, so that every result is exactly representable in float. */ +auto logical_values(int batch, int rows, int cols, int seed, bool same_for_all_batches = false) + -> std::vector +{ + std::vector out(static_cast(batch) * rows * cols); + for (int b = 0; b < batch; ++b) { + for (int i = 0; i < rows; ++i) { + for (int j = 0; j < cols; ++j) { + const auto x = seed + (same_for_all_batches ? 0 : 7 * b) + 3 * i + 2 * j; + out[(static_cast(b) * rows + i) * cols + j] = static_cast(x % 11 - 5); + } + } + } + return out; +} + +/** Strides of a batch of row- or column-major matrices. */ +auto strides_of(int rows, int cols, bool col_major, int batch_stride) -> cuda::std::array +{ + return col_major ? cuda::std::array{batch_stride, 1, rows} + : cuda::std::array{batch_stride, cols, 1}; +} + +auto buffer_size(int batch, int rows, int cols, const cuda::std::array& strides) -> size_t +{ + return static_cast((batch - 1) * strides[0] + (rows - 1) * strides[1] + + (cols - 1) * strides[2]) + + 1; +} + +/** Lay logical (batch, rows, cols) values out in memory as described by `strides`. */ +auto pack(const std::vector& logical, + int batch, + int rows, + int cols, + const cuda::std::array& strides) -> std::vector +{ + std::vector out(buffer_size(batch, rows, cols, strides), 0.0f); + for (int b = 0; b < batch; ++b) { + for (int i = 0; i < rows; ++i) { + for (int j = 0; j < cols; ++j) { + out[b * strides[0] + i * strides[1] + j * strides[2]] = + logical[(static_cast(b) * rows + i) * cols + j]; + } + } + } + return out; +} + +/** Inverse of `pack`. */ +auto unpack(const std::vector& packed, + int batch, + int rows, + int cols, + const cuda::std::array& strides) -> std::vector +{ + std::vector out(static_cast(batch) * rows * cols); + for (int b = 0; b < batch; ++b) { + for (int i = 0; i < rows; ++i) { + for (int j = 0; j < cols; ++j) { + out[(static_cast(b) * rows + i) * cols + j] = + packed[b * strides[0] + i * strides[1] + j * strides[2]]; + } + } + } + return out; +} + +/** Host reference: z[b] = alpha * x[b] * y[b] + beta * z[b], all operands logically row-major. */ +auto reference(const std::vector& x, + const std::vector& y, + const std::vector& z, + int batch, + int m, + int k, + int n, + float alpha, + float beta) -> std::vector +{ + std::vector out(static_cast(batch) * m * n); + for (int b = 0; b < batch; ++b) { + for (int i = 0; i < m; ++i) { + for (int j = 0; j < n; ++j) { + float acc = 0.0f; + for (int p = 0; p < k; ++p) { + acc += x[(static_cast(b) * m + i) * k + p] * + y[(static_cast(b) * k + p) * n + j]; + } + const auto idx = (static_cast(b) * m + i) * n + j; + out[idx] = alpha * acc + beta * z[idx]; + } + } + } + return out; +} + +auto make_strided_view(float* ptr, int batch, int rows, int cols, cuda::std::array strides) +{ + return raft::device_mdspan, raft::layout_stride>{ + ptr, raft::make_strided_layout(raft::extent_3d{batch, rows, cols}, strides)}; +} + +struct gemm_batched_params { + bool x_col_major; + bool y_col_major; + bool z_col_major; + bool use_alpha; + bool use_beta; + bool device_scalars; + /** Distance between consecutive matrices of X; -1 means tightly packed. */ + int x_batch_stride = -1; +}; + +void test_gemm_batched(const gemm_batched_params& ps) +{ + raft::resources res; + auto stream = raft::resource::get_cuda_stream(res); + + // A zero batch stride broadcasts a single matrix of X over the whole batch, which only matches + // the reference if every batch of the logical X holds the same values. + const bool broadcast_x = ps.x_batch_stride == 0; + auto x_logical = logical_values(kBatch, kM, kK, 1, broadcast_x); + auto y_logical = logical_values(kBatch, kK, kN, 4); + auto z_logical = logical_values(kBatch, kM, kN, 9); + + auto x_strides = + strides_of(kM, kK, ps.x_col_major, ps.x_batch_stride >= 0 ? ps.x_batch_stride : kM * kK); + auto y_strides = strides_of(kK, kN, ps.y_col_major, kK * kN); + auto z_strides = strides_of(kM, kN, ps.z_col_major, kM * kN); + + auto x_packed = pack(x_logical, kBatch, kM, kK, x_strides); + auto y_packed = pack(y_logical, kBatch, kK, kN, y_strides); + auto z_packed = pack(z_logical, kBatch, kM, kN, z_strides); + + rmm::device_uvector x_device(x_packed.size(), stream); + rmm::device_uvector y_device(y_packed.size(), stream); + rmm::device_uvector z_device(z_packed.size(), stream); + raft::copy(x_device.data(), x_packed.data(), x_packed.size(), stream); + raft::copy(y_device.data(), y_packed.data(), y_packed.size(), stream); + raft::copy(z_device.data(), z_packed.data(), z_packed.size(), stream); + + auto x_view = make_strided_view(x_device.data(), kBatch, kM, kK, x_strides); + auto y_view = make_strided_view(y_device.data(), kBatch, kK, kN, y_strides); + auto z_view = make_strided_view(z_device.data(), kBatch, kM, kN, z_strides); + + if (ps.device_scalars) { + auto alpha = raft::make_device_scalar(res, kAlpha); + auto beta = raft::make_device_scalar(res, kBeta); + gemm_batched(res, + x_view, + y_view, + z_view, + ps.use_alpha ? std::make_optional(alpha.view()) : std::nullopt, + ps.use_beta ? std::make_optional(beta.view()) : std::nullopt); + } else { + auto alpha = raft::make_host_scalar(kAlpha); + auto beta = raft::make_host_scalar(kBeta); + gemm_batched(res, + x_view, + y_view, + z_view, + ps.use_alpha ? std::make_optional(alpha.view()) : std::nullopt, + ps.use_beta ? std::make_optional(beta.view()) : std::nullopt); + } + + std::vector result_packed(z_packed.size()); + raft::copy(result_packed.data(), z_device.data(), result_packed.size(), stream); + raft::resource::sync_stream(res); + + auto result = unpack(result_packed, kBatch, kM, kN, z_strides); + auto gt = reference(x_logical, + y_logical, + z_logical, + kBatch, + kM, + kK, + kN, + ps.use_alpha ? kAlpha : 1.0f, + ps.use_beta ? kBeta : 0.0f); + for (size_t i = 0; i < gt.size(); ++i) { + EXPECT_FLOAT_EQ(result[i], gt[i]) << "Mismatch at index " << i; + } +} + +} // namespace + +class GemmBatchedLayoutTest : public ::testing::TestWithParam {}; + +TEST_P(GemmBatchedLayoutTest, MatchesReference) { test_gemm_batched(GetParam()); } + +INSTANTIATE_TEST_CASE_P(GemmBatched, + GemmBatchedLayoutTest, + ::testing::Values( + // every combination of per-matrix layouts, with both coefficients + gemm_batched_params{false, false, false, true, true, false}, + gemm_batched_params{false, false, true, true, true, false}, + gemm_batched_params{false, true, false, true, true, false}, + gemm_batched_params{false, true, true, true, true, false}, + gemm_batched_params{true, false, false, true, true, false}, + gemm_batched_params{true, false, true, true, true, false}, + gemm_batched_params{true, true, false, true, true, false}, + gemm_batched_params{true, true, true, true, true, false}, + // default coefficients + gemm_batched_params{false, false, false, false, false, false}, + gemm_batched_params{false, false, false, true, false, false}, + gemm_batched_params{false, false, false, false, true, false}, + gemm_batched_params{true, true, true, false, false, false}, + // device pointer mode + gemm_batched_params{false, false, false, true, true, true}, + gemm_batched_params{true, true, true, true, true, true}, + gemm_batched_params{false, false, false, false, false, true}, + // padding between consecutive matrices of X + gemm_batched_params{false, false, false, true, true, false, kM* kK + 3}, + gemm_batched_params{true, true, true, true, true, false, kM* kK + 3}, + // a single X broadcast over the whole batch + gemm_batched_params{false, false, false, true, true, false, 0}, + gemm_batched_params{true, false, true, true, true, false, 0})); + +// A plain row_major 3D mdarray is a batch of row-major matrices; this is the layout the API is +// expected to be used with most often, so check it goes through without a strided layout. +TEST(GemmBatched, RowMajorMdarray) +{ + raft::resources res; + auto stream = raft::resource::get_cuda_stream(res); + + auto x_logical = logical_values(kBatch, kM, kK, 1); + auto y_logical = logical_values(kBatch, kK, kN, 4); + auto z_logical = logical_values(kBatch, kM, kN, 9); + + auto x = raft::make_device_mdarray( + res, raft::extent_3d{kBatch, kM, kK}); + auto y = raft::make_device_mdarray( + res, raft::extent_3d{kBatch, kK, kN}); + auto z = raft::make_device_mdarray( + res, raft::extent_3d{kBatch, kM, kN}); + raft::copy(x.data_handle(), x_logical.data(), x_logical.size(), stream); + raft::copy(y.data_handle(), y_logical.data(), y_logical.size(), stream); + raft::copy(z.data_handle(), z_logical.data(), z_logical.size(), stream); + + auto alpha = raft::make_host_scalar(kAlpha); + auto beta = raft::make_host_scalar(kBeta); + gemm_batched(res, + x.view(), + y.view(), + z.view(), + std::make_optional(alpha.view()), + std::make_optional(beta.view())); + + std::vector result(z_logical.size()); + raft::copy(result.data(), z.data_handle(), result.size(), stream); + raft::resource::sync_stream(res); + + auto gt = reference(x_logical, y_logical, z_logical, kBatch, kM, kK, kN, kAlpha, kBeta); + for (size_t i = 0; i < gt.size(); ++i) { + EXPECT_FLOAT_EQ(result[i], gt[i]) << "Mismatch at index " << i; + } +} + +// The batched result must agree with calling the non-batched gemm once per matrix. +TEST(GemmBatched, MatchesUnbatchedGemm) +{ + raft::resources res; + auto stream = raft::resource::get_cuda_stream(res); + + auto x_logical = logical_values(kBatch, kM, kK, 1); + auto y_logical = logical_values(kBatch, kK, kN, 4); + + auto x = raft::make_device_mdarray( + res, raft::extent_3d{kBatch, kM, kK}); + auto y = raft::make_device_mdarray( + res, raft::extent_3d{kBatch, kK, kN}); + auto z_batched = raft::make_device_mdarray( + res, raft::extent_3d{kBatch, kM, kN}); + auto z_looped = raft::make_device_matrix(res, kBatch * kM, kN); + raft::copy(x.data_handle(), x_logical.data(), x_logical.size(), stream); + raft::copy(y.data_handle(), y_logical.data(), y_logical.size(), stream); + + gemm_batched(res, x.view(), y.view(), z_batched.view()); + + for (int b = 0; b < kBatch; ++b) { + auto x_b = raft::make_device_matrix_view( + x.data_handle() + static_cast(b) * kM * kK, kM, kK); + auto y_b = raft::make_device_matrix_view( + y.data_handle() + static_cast(b) * kK * kN, kK, kN); + auto z_b = raft::make_device_matrix_view( + z_looped.data_handle() + static_cast(b) * kM * kN, kM, kN); + gemm(res, x_b, y_b, z_b); + } + + std::vector batched(static_cast(kBatch) * kM * kN); + std::vector looped(batched.size()); + raft::copy(batched.data(), z_batched.data_handle(), batched.size(), stream); + raft::copy(looped.data(), z_looped.data_handle(), looped.size(), stream); + raft::resource::sync_stream(res); + + for (size_t i = 0; i < batched.size(); ++i) { + EXPECT_FLOAT_EQ(batched[i], looped[i]) << "Mismatch at index " << i; + } +} + +TEST(GemmBatched, RejectsNonContiguousMatrices) +{ + raft::resources res; + auto stream = raft::resource::get_cuda_stream(res); + + rmm::device_uvector buffer(static_cast(kBatch) * kM * kK * 2, stream); + auto x_view = make_strided_view( + buffer.data(), kBatch, kM, kK, cuda::std::array{2 * kM * kK, 2 * kK, 2}); + auto y_view = + make_strided_view(buffer.data(), kBatch, kK, kN, strides_of(kK, kN, false, kK * kN)); + auto z_view = + make_strided_view(buffer.data(), kBatch, kM, kN, strides_of(kM, kN, false, kM * kN)); + + EXPECT_THROW(gemm_batched(res, x_view, y_view, z_view), raft::logic_error); + raft::resource::sync_stream(res); +} + +} // namespace raft::linalg