GPU: Add cuSPARSE BSR products for BlockSparseMatrix libeigen/eigen!3130 Co-authored-by: Rasmus Munk Larsen <rmlarsen@gmail.com>
diff --git a/contrib/Eigen/src/GPU/CuSparseSupport.h b/contrib/Eigen/src/GPU/CuSparseSupport.h index 7ad6be2..de6e859 100644 --- a/contrib/Eigen/src/GPU/CuSparseSupport.h +++ b/contrib/Eigen/src/GPU/CuSparseSupport.h
@@ -20,6 +20,19 @@ #include "./Meta.h" #include <cusparse.h> +/** \brief 1 when cuSPARSE runs SpMV and SpMM on BSR (block sparse row) + * descriptors, 0 otherwise. + * + * The generic API gained BSR in SpMM with cuSPARSE 12.5 (CUDA 12.8) and in SpMV + * with 12.6.3 (CUDA 13.0 Update 1); older toolkits declare cusparseCreateBsr but + * return CUSPARSE_STATUS_NOT_SUPPORTED from the products. gpu::SparseContext + * uploads a BlockSparseMatrix as BSR only when this is 1, and as CSC otherwise. */ +#if defined(CUSPARSE_VERSION) && CUSPARSE_VERSION >= 12603 +#define EIGEN_HAS_CUSPARSE_BSR 1 +#else +#define EIGEN_HAS_CUSPARSE_BSR 0 +#endif + namespace Eigen { namespace gpu { namespace internal {
diff --git a/contrib/Eigen/src/GPU/GpuSparseContext.h b/contrib/Eigen/src/GPU/GpuSparseContext.h index 24ed2a1..3de9c32 100644 --- a/contrib/Eigen/src/GPU/GpuSparseContext.h +++ b/contrib/Eigen/src/GPU/GpuSparseContext.h
@@ -15,15 +15,17 @@ // SparseMatrix<Scalar, ColMajor> (CSC), implicitly converting RowMajor input, and // can borrow a gpu::Context so that sparse products share a stream with BLAS-1 // operations — which removes the cross-stream event waits in solvers like CG. +// It also takes BlockSparseMatrix: as BSR (internal::BsrBinding) where +// internal::use_cusparse_bsr holds, as the scalar-level toSparse() copy otherwise. // // Caching: host-input calls re-upload the values *and* index arrays on every // call. Host pointer identity cannot detect a sparsity pattern rewritten in // place or assigned into the same allocations (SparseMatrix reuses them for // same-shape assignments), so the structure is never assumed unchanged. What // is cached are the cuSPARSE descriptors (sparse descriptor keyed on -// dimensions + nonzero count, dense descriptors keyed on shape) and the -// SpMV/SpMM workspace-size queries. For repeated products with no re-upload -// at all, use deviceView(). +// dimensions + nonzero count + block size, dense descriptors keyed on shape) +// and the SpMV/SpMM workspace-size queries. For repeated products with no +// re-upload at all, use deviceView(). // // Not thread-safe: concurrent multiply* calls on one instance race on the // cuSPARSE handle, the bound stream, and the cached buffers. Use one per thread. @@ -42,6 +44,100 @@ namespace Eigen { namespace gpu { +namespace internal { +/** Whether a BlockSparseMatrix with this block shape uploads as BSR. cuSPARSE + * runs BSR products only on square blocks of size at least 2 (cusparseCreateBsr + * rejects rectangular blocks; SpMV on 1 x 1 blocks returns + * CUSPARSE_STATUS_NOT_SUPPORTED) and only from cuSPARSE 12.6.3; every other + * case takes the CSC path. */ +template <int BlockRows, int BlockCols> +struct use_cusparse_bsr + : Eigen::internal::bool_constant<(EIGEN_HAS_CUSPARSE_BSR != 0 && BlockRows == BlockCols && BlockRows >= 2)> {}; +} // namespace internal + +#if EIGEN_HAS_CUSPARSE_BSR +namespace internal { + +/** The arrays cusparseCreateBsr takes: block-row offsets (brows + 1), + * block-column indices (bnnz), and bnnz row-major blockSize x blockSize + * value blocks. Sizes are in blocks; rows()/cols()/nonZeros() in scalars. */ +template <typename Scalar, typename StorageIndex> +struct BsrArrays { + Index brows; + Index bcols; + Index bnnz; + Index blockSize; + const StorageIndex* outer; + const StorageIndex* inner; + const Scalar* values; + + Index rows() const { return brows * blockSize; } + Index cols() const { return bcols * blockSize; } + Index nonZeros() const { return bnnz * blockSize * blockSize; } +}; + +/** Reads a square-block BlockSparseMatrix's buffers as BSR without copying. + * + * A RowMajor M is BSR of M: block rows are its outer vectors and its blocks are + * stored row-major. A ColMajor M is BSR of M^T: its outer vectors are the block + * rows of M^T, and a column-major B x B block read row-major is the transposed + * block. */ +template <typename Bsm> +BsrArrays<typename Bsm::Scalar, typename Bsm::StorageIndex> bsr_of(const Bsm& M) { + return BsrArrays<typename Bsm::Scalar, typename Bsm::StorageIndex>{ + /*brows=*/M.blockOuterSize(), /*bcols=*/M.blockInnerSize(), + /*bnnz=*/M.nonZeroBlocks(), /*blockSize=*/Index(Bsm::BlockRows), + /*outer=*/M.outerIndexPtr(), /*inner=*/M.innerIndexPtr(), + /*values=*/M.valuePtr()}; +} + +/** Binds op(A) as BSR arrays of op(A) itself, the only form cuSPARSE + * multiplies. Per bsr_of(), one of {NoTrans, Trans} reads A's buffers + * directly — NoTrans for RowMajor, Trans for ColMajor — and the other needs + * the transposed copy; ConjTrans on a complex matrix always copies for the + * conjugation. The copy lives as long as the binding. */ +template <typename Bsm> +class BsrBinding { + EIGEN_STATIC_ASSERT((use_cusparse_bsr<int(Bsm::BlockRows), int(Bsm::BlockCols)>::value), + CUSPARSE_BSR_REQUIRES_SQUARE_BLOCKS_OF_SIZE_AT_LEAST_2) + + public: + using Scalar = typename Bsm::Scalar; + using Arrays = BsrArrays<Scalar, typename Bsm::StorageIndex>; + + BsrBinding(const Bsm& A, GpuOp op) { + if (op == GpuOp::ConjTrans && !NumTraits<Scalar>::IsComplex) op = GpuOp::Trans; + const bool direct = Bsm::IsRowMajor ? (op == GpuOp::NoTrans) : (op == GpuOp::Trans); + if (direct) { + arrays_ = bsr_of(A); + return; + } + if (op == GpuOp::ConjTrans && !Bsm::IsRowMajor) { + // ColMajor conj(A) reads as BSR of conj(A)^T = A^H. + copy_ = A.unaryExpr(Eigen::internal::scalar_conjugate_op<Scalar>()); + } else if (op == GpuOp::ConjTrans) { + copy_ = A.adjoint(); + } else { + copy_ = A.transpose(); + } + arrays_ = bsr_of(copy_); + } + + BsrBinding(const BsrBinding&) = delete; + BsrBinding& operator=(const BsrBinding&) = delete; + + const Arrays& arrays() const { return arrays_; } + + private: + // transpose()/adjoint() swap the block dimensions, which are equal here, so + // the copy has A's own type in either storage order. + Bsm copy_; + Arrays arrays_{}; +}; + +} // namespace internal +#endif // EIGEN_HAS_CUSPARSE_BSR + /** Sparse product expression: DeviceSparseView * DeviceMatrix → SpMVExpr. * Evaluated by DeviceMatrix::operator=(SpMVExpr): dispatches to cusparseSpMV * when the dense operand has one column, cusparseSpMM otherwise. */ @@ -58,8 +154,9 @@ const DeviceMatrix<Scalar>& x_; }; -/** Device-resident sparse matrix view. Returned by SparseContext::deviceView(). - * Lightweight handle referencing the context's cached device data. +/** Device-resident sparse matrix view. Returned by SparseContext::deviceView() + * for a SparseMatrix (CSC) or a BlockSparseMatrix (BSR). Lightweight handle + * referencing the context's cached device data. * * \warning One SparseContext caches one sparse matrix at a time. Any later * upload through the same context — a second deviceView() or any host-input @@ -105,6 +202,13 @@ using SpMat = SparseMatrix<Scalar, ColMajor, StorageIndex>; using DenseVector = Matrix<Scalar, Dynamic, 1>; using DenseMatrix = Matrix<Scalar, Dynamic, Dynamic, ColMajor>; + /** BlockSparseMatrix types the block-sparse overloads accept: `int` indices, + * either storage order, any block shape. internal::use_cusparse_bsr picks the + * upload: BSR — no host copy for `op == NoTrans` on a RowMajor matrix or + * `op == Trans` on a ColMajor one, a transposing (and conjugating) copy + * otherwise — or the scalar-level toSparse() copy on the CSC path. */ + template <int Options, int BlockRows, int BlockCols> + using BlockSpMat = BlockSparseMatrix<Scalar, Options, BlockRows, BlockCols, StorageIndex>; /** Standalone: creates own stream and cuSPARSE handle. */ SparseContext() : owns_handle_(true) { @@ -222,6 +326,66 @@ return Y; } + // BlockSparseMatrix overloads: same contracts as the SparseMatrix ones, + // dispatched on internal::use_cusparse_bsr<BlockRows, BlockCols> to a BSR + // upload of op(A) or to the scalar-level CSC path. + + /** Upload a block-sparse matrix and return a view; see the SparseMatrix + * overload for the caching contract. A ColMajor matrix bound as BSR is + * transposed on the host once, at upload. */ + template <int Options, int BlockRows, int BlockCols> + DeviceSparseView<Scalar> deviceView(const BlockSpMat<Options, BlockRows, BlockCols>& A) { + return device_view_block(A, internal::use_cusparse_bsr<BlockRows, BlockCols>()); + } + + /** Compute y = A * x for a block-sparse A. Returns y as a new dense vector. */ + template <int Options, int BlockRows, int BlockCols, typename Rhs> + DenseVector multiply(const BlockSpMat<Options, BlockRows, BlockCols>& A, const MatrixBase<Rhs>& x) { + return multiply_host_return_block(A, x, GpuOp::NoTrans); + } + + /** Compute y = alpha * op(A) * x + beta * y (in-place, host vectors) for a block-sparse A. */ + template <int Options, int BlockRows, int BlockCols, typename Rhs, typename Dest> + void multiply(const BlockSpMat<Options, BlockRows, BlockCols>& A, const MatrixBase<Rhs>& x, MatrixBase<Dest>& y, + Scalar alpha = Scalar(1), Scalar beta = Scalar(0), GpuOp op = GpuOp::NoTrans) { + multiply_host_block(A, x.derived(), y.derived(), alpha, beta, op, + internal::use_cusparse_bsr<BlockRows, BlockCols>()); + } + + /** Compute d_y = A * d_x for a block-sparse A; the matrix is re-uploaded on + * each call, use deviceView() to upload once. */ + template <int Options, int BlockRows, int BlockCols> + void multiply(const BlockSpMat<Options, BlockRows, BlockCols>& A, const DeviceMatrix<Scalar>& d_x, + DeviceMatrix<Scalar>& d_y) { + multiply(A, d_x, d_y, Scalar(1), Scalar(0), GpuOp::NoTrans); + } + + /** Compute d_y = alpha * op(A) * d_x + beta * d_y (DeviceMatrix, in-place) for a block-sparse A. */ + template <int Options, int BlockRows, int BlockCols> + void multiply(const BlockSpMat<Options, BlockRows, BlockCols>& A, const DeviceMatrix<Scalar>& d_x, + DeviceMatrix<Scalar>& d_y, Scalar alpha, Scalar beta, GpuOp op = GpuOp::NoTrans) { + multiply_device_block(A, d_x, d_y, alpha, beta, op, internal::use_cusparse_bsr<BlockRows, BlockCols>()); + } + + /** Compute y = A^T * x (host vectors) for a block-sparse A. */ + template <int Options, int BlockRows, int BlockCols, typename Rhs> + DenseVector multiplyT(const BlockSpMat<Options, BlockRows, BlockCols>& A, const MatrixBase<Rhs>& x) { + return multiply_host_return_block(A, x, GpuOp::Trans); + } + + /** Compute y = A^H * x for a block-sparse A. For real Scalar this is equivalent to multiplyT. */ + template <int Options, int BlockRows, int BlockCols, typename Rhs> + DenseVector multiplyAdjoint(const BlockSpMat<Options, BlockRows, BlockCols>& A, const MatrixBase<Rhs>& x) { + return multiply_host_return_block(A, x, GpuOp::ConjTrans); + } + + /** Compute Y = op(A) * X for a block-sparse A and a dense X (multiple RHS). Returns Y. */ + template <int Options, int BlockRows, int BlockCols, typename Rhs> + DenseMatrix multiplyMat(const BlockSpMat<Options, BlockRows, BlockCols>& A, const MatrixBase<Rhs>& X, + GpuOp op = GpuOp::NoTrans) { + return multiply_mat_block(A, X, op, internal::use_cusparse_bsr<BlockRows, BlockCols>()); + } + cudaStream_t stream() const { return stream_; } private: @@ -241,11 +405,13 @@ mutable internal::DeviceBuffer d_workspace_; - // Cached cuSPARSE sparse matrix descriptor. + // Cached cuSPARSE sparse matrix descriptor, keyed on the scalar-level shape + // and nonzero count plus the block size (0 for CSC/CSR, B for BSR). cusparseSpMatDescr_t spmat_desc_ = nullptr; Index cached_rows_ = -1; Index cached_cols_ = -1; Index cached_nnz_ = -1; + Index cached_block_size_ = 0; // Bumped on every sparse upload; DeviceSparseViews record it at creation so // a stale view (its data replaced by a later upload) asserts at evaluation. @@ -268,6 +434,19 @@ mutable size_t spmv_ws_size_[3] = {kWsUnknown, kWsUnknown, kWsUnknown}; mutable size_t spmm_ws_size_[3] = {kWsUnknown, kWsUnknown, kWsUnknown}; + static constexpr cusparseIndexType_t kIndexType = + (sizeof(StorageIndex) == 4) ? CUSPARSE_INDEX_32I : CUSPARSE_INDEX_64I; + static constexpr cudaDataType_t kValueType = internal::cuda_data_type<Scalar>::value; + + // Empty-operand result on the host: y <- beta * y, no device work. + template <typename Dest> + static void scale_host(Dest& y, Scalar beta) { + if (beta == Scalar(0)) + y.setZero(); + else + y *= beta; + } + // Shared host-input SpMV entry: y = op(A) * x into a fresh vector. template <typename InputType, typename Rhs> DenseVector multiply_host_return(const SparseMatrixBase<InputType>& A, const MatrixBase<Rhs>& x, GpuOp op) { @@ -296,15 +475,110 @@ eigen_assert(y.size() == y_size); if (m == 0 || n == 0 || nnz == 0) { - if (beta == Scalar(0)) - y.setZero(); - else - y *= beta; + scale_host(y, beta); return; } upload_sparse(A); + spmv_host_cached(x, y, x_size, y_size, alpha, beta, op); + } + // BlockSparseMatrix primitives, dispatched on use_cusparse_bsr: the true_type + // overloads bind op(A) as BSR (internal::BsrBinding) and run cuSPARSE with + // NON_TRANSPOSE, the false_type ones forward the scalar-level toSparse() copy + // to the CSC overloads. + template <typename Bsm, typename Rhs> + DenseVector multiply_host_return_block(const Bsm& A, const MatrixBase<Rhs>& x, GpuOp op) { + DenseVector y((op == GpuOp::NoTrans) ? A.rows() : A.cols()); + y.setZero(); + multiply_host_block(A, x.derived(), y, Scalar(1), Scalar(0), op, + internal::use_cusparse_bsr<int(Bsm::BlockRows), int(Bsm::BlockCols)>()); + return y; + } + + template <typename Bsm> + DeviceSparseView<Scalar> device_view_block(const Bsm& A, std::false_type) { + return deviceView(SpMat(A.toSparse())); + } + + template <typename Bsm, typename RhsDerived, typename DestDerived> + void multiply_host_block(const Bsm& A, const RhsDerived& x, DestDerived& y, Scalar alpha, Scalar beta, GpuOp op, + std::false_type) { + multiply(A.toSparse(), x, y, alpha, beta, op); + } + + template <typename Bsm> + void multiply_device_block(const Bsm& A, const DeviceMatrix<Scalar>& d_x, DeviceMatrix<Scalar>& d_y, Scalar alpha, + Scalar beta, GpuOp op, std::false_type) { + multiply(A.toSparse(), d_x, d_y, alpha, beta, op); + } + + template <typename Bsm, typename Rhs> + DenseMatrix multiply_mat_block(const Bsm& A, const MatrixBase<Rhs>& X, GpuOp op, std::false_type) { + return multiplyMat(A.toSparse(), X, op); + } + +#if EIGEN_HAS_CUSPARSE_BSR + template <typename Bsm> + DeviceSparseView<Scalar> device_view_block(const Bsm& A, std::true_type) { + internal::check_storage_index_bounds<StorageIndex>(A.rows(), A.cols(), A.nonZeros()); + const internal::BsrBinding<Bsm> bound(A, GpuOp::NoTrans); + upload_bsr(bound.arrays()); + return DeviceSparseView<Scalar>(*this, A.rows(), A.cols(), generation_); + } + + template <typename Bsm, typename RhsDerived, typename DestDerived> + void multiply_host_block(const Bsm& A, const RhsDerived& x, DestDerived& y, Scalar alpha, Scalar beta, GpuOp op, + std::true_type) { + internal::check_storage_index_bounds<StorageIndex>(A.rows(), A.cols(), A.nonZeros()); + const internal::BsrBinding<Bsm> bound(A, op); + const internal::BsrArrays<Scalar, StorageIndex>& opA = bound.arrays(); + + eigen_assert(x.size() == opA.cols()); + eigen_assert(y.size() == opA.rows()); + + if (opA.rows() == 0 || opA.cols() == 0 || opA.bnnz == 0) { + scale_host(y, beta); + return; + } + + upload_bsr(opA); + spmv_host_cached(x, y, opA.cols(), opA.rows(), alpha, beta, CUSPARSE_OPERATION_NON_TRANSPOSE); + } + + template <typename Bsm> + void multiply_device_block(const Bsm& A, const DeviceMatrix<Scalar>& d_x, DeviceMatrix<Scalar>& d_y, Scalar alpha, + Scalar beta, GpuOp op, std::true_type) { + internal::check_storage_index_bounds<StorageIndex>(A.rows(), A.cols(), A.nonZeros()); + const internal::BsrBinding<Bsm> bound(A, op); + upload_bsr(bound.arrays()); + spmv_device_exec(d_x, d_y, alpha, beta, GpuOp::NoTrans); + } + + template <typename Bsm, typename Rhs> + DenseMatrix multiply_mat_block(const Bsm& A, const MatrixBase<Rhs>& X, GpuOp op, std::true_type) { + internal::check_storage_index_bounds<StorageIndex>(A.rows(), A.cols(), A.nonZeros()); + const internal::BsrBinding<Bsm> bound(A, op); + const internal::BsrArrays<Scalar, StorageIndex>& opA = bound.arrays(); + const DenseMatrix rhs(X.derived()); + eigen_assert(opA.cols() == rhs.rows()); + + const Index n = rhs.cols(); + if (opA.rows() == 0 || opA.cols() == 0 || n == 0 || opA.bnnz == 0) return DenseMatrix::Zero(opA.rows(), n); + + DenseMatrix Y = DenseMatrix::Zero(opA.rows(), n); + upload_bsr(opA); + spmm_host_cached(rhs, Y, opA.rows(), opA.cols(), Scalar(1), Scalar(0), CUSPARSE_OPERATION_NON_TRANSPOSE); + return Y; + } +#endif // EIGEN_HAS_CUSPARSE_BSR + + // y = alpha * op(D) * x + beta * y against the cached upload D, staging the + // host vectors through d_x_ / d_y_. x_size and y_size are the column and row + // counts of op(D). + template <typename RhsDerived, typename DestDerived> + void spmv_host_cached(const RhsDerived& x, DestDerived& y, Index x_size, Index y_size, Scalar alpha, Scalar beta, + cusparseOperation_t op) { ensure_buffer(d_x_, static_cast<size_t>(x_size) * sizeof(Scalar)); // Ref binds in place when x is already a contiguous vector; only genuine // expressions are evaluated into the Ref's internal temporary. @@ -329,10 +603,13 @@ public: /** Execute SpMV using the already-uploaded sparse matrix (no re-upload). * Used by SpMVExpr (d_y = d_A * d_x) for cached deviceView() paths. - * The sparse matrix must have been uploaded via deviceView() or multiply(). */ + * The sparse matrix must have been uploaded via deviceView() or multiply(). + * A BSR upload describes op(A) as bound at upload time and cuSPARSE runs no + * transposed BSR product, so \p op must then be GpuOp::NoTrans. */ void spmv_device_exec(const DeviceMatrix<Scalar>& d_x, DeviceMatrix<Scalar>& d_y, Scalar alpha = Scalar(1), Scalar beta = Scalar(0), GpuOp op = GpuOp::NoTrans) const { eigen_assert(spmat_desc_ && "sparse matrix not uploaded — call deviceView() or multiply() first"); + check_op_against_upload(op); // cuSPARSE SpMV: y must not alias x (undefined behavior). eigen_assert(d_x.data() != d_y.data() && "SpMV: output aliases input vector"); @@ -370,10 +647,12 @@ /** Execute SpMM (d_Y = alpha * op(A) * d_X + beta * d_Y) using the * already-uploaded sparse matrix. d_X / d_Y are device-resident dense - * column-major matrices. Used by SpMVExpr when the RHS has > 1 column. */ + * column-major matrices. Used by SpMVExpr when the RHS has > 1 column. + * As for spmv_device_exec(), a BSR upload requires \p op == GpuOp::NoTrans. */ void spmm_device_exec(const DeviceMatrix<Scalar>& d_X, DeviceMatrix<Scalar>& d_Y, Scalar alpha = Scalar(1), Scalar beta = Scalar(0), GpuOp op = GpuOp::NoTrans) const { eigen_assert(spmat_desc_ && "sparse matrix not uploaded — call deviceView() or multiply() first"); + check_op_against_upload(op); eigen_assert(d_X.data() != d_Y.data() && "SpMM: output aliases input matrix"); const cusparseOperation_t cu_op = internal::to_cusparse_op<Scalar>(op); @@ -418,8 +697,8 @@ #endif // Map a user-facing op on A to the cuSPARSE op on the cached descriptor. - // Identity on cuSPARSE 12+ (descriptor is CSC of A); inverted on 11.x - // (descriptor is CSR of A^T). + // Identity on cuSPARSE 12+ (descriptor is CSC of A, or BSR of op(A)); + // inverted on 11.x (descriptor is CSR of A^T). static cusparseOperation_t descriptor_op(cusparseOperation_t user_op) { EIGEN_IF_CONSTEXPR (!kUseCsrOfTranspose) return user_op; switch (user_op) { @@ -459,7 +738,7 @@ void update_dnvec(cusparseDnVecDescr_t& desc, int64_t& cur_size, int64_t size, void* ptr) const { if (!desc || cur_size != size) { if (desc) EIGEN_CUSPARSE_CHECK(cusparseDestroyDnVec(desc)); - EIGEN_CUSPARSE_CHECK(cusparseCreateDnVec(&desc, size, ptr, internal::cuda_data_type<Scalar>::value)); + EIGEN_CUSPARSE_CHECK(cusparseCreateDnVec(&desc, size, ptr, kValueType)); cur_size = size; invalidate_ws_caches(); } else { @@ -472,8 +751,7 @@ if (!desc || cur_rows != rows || cur_cols != cols) { if (desc) EIGEN_CUSPARSE_CHECK(cusparseDestroyDnMat(desc)); // Column-major with ld = rows. - EIGEN_CUSPARSE_CHECK(cusparseCreateDnMat(&desc, rows, cols, rows, ptr, internal::cuda_data_type<Scalar>::value, - CUSPARSE_ORDER_COL)); + EIGEN_CUSPARSE_CHECK(cusparseCreateDnMat(&desc, rows, cols, rows, ptr, kValueType, CUSPARSE_ORDER_COL)); cur_rows = rows; cur_cols = cols; invalidate_ws_caches(); @@ -493,7 +771,6 @@ void exec_spmv(Index x_size, Index y_size, void* d_x_ptr, void* d_y_ptr, Scalar alpha, Scalar beta, cusparseOperation_t op) const { - constexpr cudaDataType_t dtype = internal::cuda_data_type<Scalar>::value; const cusparseOperation_t cu_op = descriptor_op(op); update_dnvec(x_vec_desc_, x_vec_size_, x_size, d_x_ptr); update_dnvec(y_vec_desc_, y_vec_size_, y_size, d_y_ptr); @@ -501,17 +778,16 @@ size_t& ws_size = spmv_ws_size_[op_index(cu_op)]; if (ws_size == kWsUnknown) { EIGEN_CUSPARSE_CHECK(cusparseSpMV_bufferSize(handle_, cu_op, &alpha, spmat_desc_, x_vec_desc_, &beta, y_vec_desc_, - dtype, CUSPARSE_SPMV_ALG_DEFAULT, &ws_size)); + kValueType, CUSPARSE_SPMV_ALG_DEFAULT, &ws_size)); } ensure_buffer(d_workspace_, ws_size); - EIGEN_CUSPARSE_CHECK(cusparseSpMV(handle_, cu_op, &alpha, spmat_desc_, x_vec_desc_, &beta, y_vec_desc_, dtype, + EIGEN_CUSPARSE_CHECK(cusparseSpMV(handle_, cu_op, &alpha, spmat_desc_, x_vec_desc_, &beta, y_vec_desc_, kValueType, CUSPARSE_SPMV_ALG_DEFAULT, d_workspace_.get())); } void exec_spmm(Index m_op, Index k_op, Index n, void* d_x_ptr, void* d_y_ptr, Scalar alpha, Scalar beta, cusparseOperation_t op) const { - constexpr cudaDataType_t dtype = internal::cuda_data_type<Scalar>::value; const cusparseOperation_t cu_op = descriptor_op(op); // X is k_op x n, Y is m_op x n (column-major, post-op shapes). update_dnmat(x_mat_desc_, x_mat_rows_, x_mat_cols_, k_op, n, d_x_ptr); @@ -520,13 +796,13 @@ size_t& ws_size = spmm_ws_size_[op_index(cu_op)]; if (ws_size == kWsUnknown) { EIGEN_CUSPARSE_CHECK(cusparseSpMM_bufferSize(handle_, cu_op, CUSPARSE_OPERATION_NON_TRANSPOSE, &alpha, - spmat_desc_, x_mat_desc_, &beta, y_mat_desc_, dtype, kSpMMAlg, + spmat_desc_, x_mat_desc_, &beta, y_mat_desc_, kValueType, kSpMMAlg, &ws_size)); } ensure_buffer(d_workspace_, ws_size); EIGEN_CUSPARSE_CHECK(cusparseSpMM(handle_, cu_op, CUSPARSE_OPERATION_NON_TRANSPOSE, &alpha, spmat_desc_, - x_mat_desc_, &beta, y_mat_desc_, dtype, kSpMMAlg, d_workspace_.get())); + x_mat_desc_, &beta, y_mat_desc_, kValueType, kSpMMAlg, d_workspace_.get())); } void spmm_impl(const SpMat& A, const DenseMatrix& X, DenseMatrix& Y, Scalar alpha, Scalar beta, @@ -543,16 +819,19 @@ const Index nnz = A.nonZeros(); if (m_op == 0 || n == 0 || k_op == 0 || nnz == 0) { - if (beta == Scalar(0)) - Y.setZero(); - else - Y *= beta; + scale_host(Y, beta); return; } upload_sparse(A); + spmm_host_cached(X, Y, m_op, k_op, alpha, beta, op); + } - // Upload X to device. X is k_op x n, Y is m_op x n (column-major). + // Y = alpha * op(D) * X + beta * Y against the cached upload D, staging the + // host matrices through d_x_ / d_y_. X is k_op x n, Y is m_op x n. + void spmm_host_cached(const DenseMatrix& X, DenseMatrix& Y, Index m_op, Index k_op, Scalar alpha, Scalar beta, + cusparseOperation_t op) { + const Index n = X.cols(); const size_t x_bytes = static_cast<size_t>(k_op) * static_cast<size_t>(n) * sizeof(Scalar); const size_t y_bytes = static_cast<size_t>(m_op) * static_cast<size_t>(n) * sizeof(Scalar); ensure_buffer(d_x_, x_bytes); @@ -574,23 +853,57 @@ // demotes to TRANSPOSE. We register the same CSC buffers as CSR-of-A^T // (dims swapped) on 11.x and invert the op at exec time via // descriptor_op() — no transpose-copy required. - upload_compressed_arrays(A.rows(), A.cols(), A.nonZeros(), - /*outer_count=*/A.cols() + 1, A.outerIndexPtr(), A.innerIndexPtr(), A.valuePtr()); + const Index m = A.rows(); + const Index n = A.cols(); + const Index nnz = A.nonZeros(); + upload_arrays(/*outer_count=*/n + 1, /*inner_count=*/nnz, /*value_count=*/nnz, A.outerIndexPtr(), A.innerIndexPtr(), + A.valuePtr()); + if (descriptor_key_matches(m, n, nnz, /*block_size=*/0)) return; + + destroy_spmat_descriptor(/*checked=*/true); + EIGEN_IF_CONSTEXPR (kUseCsrOfTranspose) { + // cuSPARSE 11.x: cusparseSpMM rejects CSC for matA. CSC of A and CSR of + // A^T share the same buffers, so register the data as CSR-of-A^T (dims + // swapped) and invert the op in exec_spmv / spmm_impl via descriptor_op. + EIGEN_CUSPARSE_CHECK(cusparseCreateCsr(&spmat_desc_, n, m, nnz, d_outerPtr_.get(), d_innerIdx_.get(), + d_values_.get(), kIndexType, kIndexType, CUSPARSE_INDEX_BASE_ZERO, + kValueType)); + } else { + EIGEN_CUSPARSE_CHECK(cusparseCreateCsc(&spmat_desc_, m, n, nnz, d_outerPtr_.get(), d_innerIdx_.get(), + d_values_.get(), kIndexType, kIndexType, CUSPARSE_INDEX_BASE_ZERO, + kValueType)); + } + set_descriptor_key(m, n, nnz, /*block_size=*/0); } - void upload_compressed_arrays(Index m, Index n, Index nnz, Index outer_count, const StorageIndex* host_outer, - const StorageIndex* host_inner, const Scalar* host_values) { +#if EIGEN_HAS_CUSPARSE_BSR + void upload_bsr(const internal::BsrArrays<Scalar, StorageIndex>& opA) { + upload_arrays(/*outer_count=*/opA.brows + 1, /*inner_count=*/opA.bnnz, /*value_count=*/opA.nonZeros(), opA.outer, + opA.inner, opA.values); + if (descriptor_key_matches(opA.rows(), opA.cols(), opA.nonZeros(), opA.blockSize)) return; + + destroy_spmat_descriptor(/*checked=*/true); + // Row-major blocks: cusparseSpMM accepts no other block layout for BSR. + EIGEN_CUSPARSE_CHECK(cusparseCreateBsr(&spmat_desc_, opA.brows, opA.bcols, opA.bnnz, opA.blockSize, opA.blockSize, + d_outerPtr_.get(), d_innerIdx_.get(), d_values_.get(), kIndexType, + kIndexType, CUSPARSE_INDEX_BASE_ZERO, kValueType, CUSPARSE_ORDER_ROW)); + set_descriptor_key(opA.rows(), opA.cols(), opA.nonZeros(), opA.blockSize); + } +#endif // EIGEN_HAS_CUSPARSE_BSR + + void upload_arrays(Index outer_count, Index inner_count, Index value_count, const StorageIndex* host_outer, + const StorageIndex* host_inner, const Scalar* host_values) { const size_t outer_bytes = static_cast<size_t>(outer_count) * sizeof(StorageIndex); - const size_t inner_bytes = static_cast<size_t>(nnz) * sizeof(StorageIndex); - const size_t val_bytes = static_cast<size_t>(nnz) * sizeof(Scalar); + const size_t inner_bytes = static_cast<size_t>(inner_count) * sizeof(StorageIndex); + const size_t val_bytes = static_cast<size_t>(value_count) * sizeof(Scalar); // Values *and* index arrays are re-uploaded unconditionally: host pointer // identity cannot detect a same-shape/same-nnz pattern rewritten in place // or assigned into the same allocations (SparseMatrix reuses them), so a // structure cache keyed on pointers would silently serve stale indices. // Only the cuSPARSE descriptor and the workspace-size queries are cached, - // keyed on (rows, cols, nnz). Every upload invalidates outstanding - // DeviceSparseViews via the generation counter. + // keyed on (rows, cols, nnz, block size). Every upload invalidates + // outstanding DeviceSparseViews via the generation counter. ++generation_; ensure_buffer(d_values_, val_bytes); ensure_buffer(d_outerPtr_, outer_bytes); @@ -600,30 +913,28 @@ cudaMemcpyAsync(d_outerPtr_.get(), host_outer, outer_bytes, cudaMemcpyHostToDevice, stream_)); EIGEN_CUDA_RUNTIME_CHECK( cudaMemcpyAsync(d_innerIdx_.get(), host_inner, inner_bytes, cudaMemcpyHostToDevice, stream_)); + } - // Same shape and nnz: the grow-only device buffers cannot have been - // reallocated, so the existing descriptor still points at the freshly - // written data. - if (m == cached_rows_ && n == cached_cols_ && nnz == cached_nnz_) return; + // Same shape, nnz and format: the grow-only device buffers cannot have been + // reallocated, so the existing descriptor still points at the freshly + // written data. + bool descriptor_key_matches(Index m, Index n, Index nnz, Index block_size) const { + return m == cached_rows_ && n == cached_cols_ && nnz == cached_nnz_ && block_size == cached_block_size_; + } - destroy_spmat_descriptor(/*checked=*/true); - - constexpr cusparseIndexType_t idx_type = (sizeof(StorageIndex) == 4) ? CUSPARSE_INDEX_32I : CUSPARSE_INDEX_64I; - constexpr cudaDataType_t val_type = internal::cuda_data_type<Scalar>::value; - - EIGEN_IF_CONSTEXPR (kUseCsrOfTranspose) { - // cuSPARSE 11.x: cusparseSpMM rejects CSC for matA. CSC of A and CSR of - // A^T share the same buffers, so register the data as CSR-of-A^T (dims - // swapped) and invert the op in exec_spmv / spmm_impl via descriptor_op. - EIGEN_CUSPARSE_CHECK(cusparseCreateCsr(&spmat_desc_, n, m, nnz, d_outerPtr_.get(), d_innerIdx_.get(), - d_values_.get(), idx_type, idx_type, CUSPARSE_INDEX_BASE_ZERO, val_type)); - } else { - EIGEN_CUSPARSE_CHECK(cusparseCreateCsc(&spmat_desc_, m, n, nnz, d_outerPtr_.get(), d_innerIdx_.get(), - d_values_.get(), idx_type, idx_type, CUSPARSE_INDEX_BASE_ZERO, val_type)); - } + void set_descriptor_key(Index m, Index n, Index nnz, Index block_size) { cached_rows_ = m; cached_cols_ = n; cached_nnz_ = nnz; + cached_block_size_ = block_size; + } + + // A BSR descriptor holds op(A) as bound at upload time, and cuSPARSE runs no + // transposed BSR product. + void check_op_against_upload(GpuOp op) const { + eigen_assert((cached_block_size_ == 0 || op == GpuOp::NoTrans) && + "cuSPARSE runs BSR products only with op == NoTrans; pass op to multiply(A, d_x, d_y, ...) instead"); + EIGEN_UNUSED_VARIABLE(op); } // Destroy the sparse-matrix descriptor and reset the cache identity. @@ -637,9 +948,7 @@ EIGEN_UNUSED_VARIABLE(checked); spmat_desc_ = nullptr; } - cached_rows_ = -1; - cached_cols_ = -1; - cached_nnz_ = -1; + set_descriptor_key(-1, -1, -1, 0); invalidate_ws_caches(); }
diff --git a/contrib/Eigen/src/GPU/README.md b/contrib/Eigen/src/GPU/README.md index 548c5c3..45e91b3 100644 --- a/contrib/Eigen/src/GPU/README.md +++ b/contrib/Eigen/src/GPU/README.md
@@ -14,7 +14,7 @@ | CUDA Toolkit | streams, memory, error codes | <https://docs.nvidia.com/cuda/> | | cuBLAS (incl. cuBLASLt) | `DeviceMatrix` products, BLAS-1 | <https://docs.nvidia.com/cuda/cublas/> | | cuSOLVER | dense LLT / LU / QR / SVD / EVD | <https://docs.nvidia.com/cuda/cusolver/> | -| cuSPARSE | SpMV / SpMM | <https://docs.nvidia.com/cuda/cusparse/> | +| cuSPARSE | SpMV / SpMM (CSC and BSR) | <https://docs.nvidia.com/cuda/cusparse/> | | cuFFT | `gpu::FFT` | <https://docs.nvidia.com/cuda/cufft/> | | NPP | device-side scalar and coefficient-wise arithmetic | <https://docs.nvidia.com/cuda/npp/> | | cuDSS | sparse direct solvers (separate install) | <https://docs.nvidia.com/cuda/cudss/> | @@ -538,6 +538,45 @@ carries a generation counter — using a view after any later upload through its context asserts instead of silently multiplying by the wrong matrix. +#### Block sparse matrices (BSR) + +Every `SparseContext` entry point above also accepts a `BlockSparseMatrix`. +Square blocks of size at least 2 upload in cuSPARSE's BSR (block sparse row) +format on cuSPARSE 12.6.3 (CUDA 13.0 Update 1) or newer, where the generic +SpMV and SpMM run on BSR descriptors; `EIGEN_HAS_CUSPARSE_BSR` is 1 there and +0 on older toolkits. Any other block shape, and every `BlockSparseMatrix` on an +older toolkit, takes the CSC path instead: the matrix is expanded with +`toSparse()` on the host, once per host-input call or once per `deviceView()`, +so source compatibility does not depend on the toolkit. + +```cpp +BlockSparseMatrix<double, RowMajor, 3, 3> A = ...; // 3x3 blocks, block-row storage +VectorXd y = spmv.multiply(A, x); // y = A * x, no host copy +VectorXd z = spmv.multiplyT(A, x); // z = A^T * x, transposed on the host +MatrixXd Y = spmv.multiplyMat(A, X); // SpMM + +auto d_A = spmv_dev.deviceView(A); // upload once as BSR +d_y = d_A * d_x; // SpMV / SpMM, stays on device +``` + +cuSPARSE multiplies a BSR descriptor only as `op == NoTrans` with row-major +blocks (`CUSPARSE_ORDER_ROW`), so the op is applied on the host rather than +passed to the library. A `RowMajor` matrix is BSR of itself and uploads +without a host copy for `NoTrans`; a `ColMajor` one — column-major blocks in +block-column order — is BSR of its transpose and uploads without a copy for +`Trans` (and `ConjTrans` on real scalars). Every other op / storage-order +combination transposes (and conjugates) the matrix on the host first, once per +host-input call or once per `deviceView()`. `spmv_device_exec()` and +`spmm_device_exec()` consequently accept only `NoTrans` against a BSR upload +(debug builds assert); a `BlockSparseMatrix` on the CSC path passes the op to +cuSPARSE like a `SparseMatrix`. + +The CSC fallback exists because cuSPARSE rejects rectangular blocks at +descriptor creation and runs no BSR SpMV on 1 x 1 blocks; +`internal::use_cusparse_bsr<BlockRows, BlockCols>` is the exact selector. The +`int` index type and the `int` limits on dimensions and nonzeros are as for +`SparseMatrix`. + ### Eigen algorithm interop (example: Conjugate gradient) The BLAS-1 operators and `DeviceSparseView` make `DeviceMatrix` usable as a @@ -692,6 +731,7 @@ | `x.squaredNorm()` | `cublasXdot(x, x)` | returns `DeviceScalar<RealScalar>` | | `d_y = view * d_x` | `cusparseSpMV` | device-resident SpMV | | `d_Y = view * d_X` | `cusparseSpMM` | device-resident SpMM (RHS with >1 column) | +| same, `view` of a `BlockSparseMatrix` | `cusparseSpMV` / `cusparseSpMM` on a BSR descriptor | opA=N, row-major blocks; op(A) formed on the host | ### `DeviceMatrix<Scalar>` @@ -1030,10 +1070,15 @@ ### `gpu::SparseContext<Scalar>` -- SpMV/SpMM (cuSPARSE) -Accepts `SparseMatrix<Scalar, ColMajor>`. Host-input methods accept host data -and return host data; device-input methods (`deviceView()`, `multiply(A, d_x, -d_y)`) operate on `DeviceMatrix`. Matrix dimensions and nonzero count must fit -in `int` (cuSPARSE limitation; debug builds assert). +Accepts `SparseMatrix<Scalar, ColMajor>` and `BlockSparseMatrix<Scalar, +Options, BlockRows, BlockCols, int>` (the `BlockSpMat<Options, BlockRows, +BlockCols>` alias; see [Block sparse matrices](#block-sparse-matrices-bsr) for +which block shapes upload as BSR and which op / storage-order combinations do +so without a host copy). Host-input +methods accept host data and return host data; device-input methods +(`deviceView()`, `multiply(A, d_x, d_y)`) operate on `DeviceMatrix`. Matrix +dimensions and nonzero count must fit in `int` (cuSPARSE limitation; debug +builds assert). ```cpp gpu::SparseContext() // Creates own stream + cuSPARSE handle @@ -1055,7 +1100,8 @@ DeviceSparseView deviceView(A) // Upload sparse matrix, return view uint64_t uploadGeneration() // Generation of the cached upload -// Advanced: run SpMV/SpMM against the already-uploaded matrix +// Advanced: run SpMV/SpMM against the already-uploaded matrix (op must be +// NoTrans after a BSR upload) void spmv_device_exec(d_x, d_y, alpha=1, beta=0, op=GpuOp::NoTrans) void spmm_device_exec(d_X, d_Y, alpha=1, beta=0, op=GpuOp::NoTrans) @@ -1136,8 +1182,8 @@ | `GpuEigenSolver.h` | `GpuSolverContext.h` | `gpu::SelfAdjointEigenSolver<>` | | `CuFftSupport.h` | `GpuSupport.h`, `<cufft.h>` | cuFFT error macro, type-dispatch wrappers | | `GpuFFT.h` | `CuFftSupport.h`, `CuBlasSupport.h`, `GpuContext.h` | `gpu::FFT<>` -- 1D/2D FFT with plan caching | -| `CuSparseSupport.h` | `GpuSupport.h`, `<cusparse.h>` | cuSPARSE error macro | -| `GpuSparseContext.h` | `CuSparseSupport.h` | `gpu::SparseContext<>`, `gpu::DeviceSparseView<>` | +| `CuSparseSupport.h` | `GpuSupport.h`, `<cusparse.h>` | cuSPARSE error macro, `EIGEN_HAS_CUSPARSE_BSR` | +| `GpuSparseContext.h` | `CuSparseSupport.h` | `gpu::SparseContext<>`, `gpu::DeviceSparseView<>`, BSR binding of `BlockSparseMatrix` | | `CuDssSupport.h` | `GpuSupport.h`, `<cudss.h>` | cuDSS error macro, type traits (optional) | | `GpuSparseSolverBase.h` | `CuDssSupport.h` | CRTP base for sparse solvers (optional) | | `GpuSparseLLT.h` | `GpuSparseSolverBase.h` | `gpu::SparseLLT<>` -- Sparse Cholesky via cuDSS (optional) | @@ -1155,7 +1201,7 @@ cmake --build build --target cublas cusolver_llt cusolver_lu \ cusolver_qr cusolver_svd cusolver_eigen \ - device_matrix cufft cusparse_spmv cg + device_matrix cufft cusparse_spmv cusparse_bsr cg ctest --test-dir build -L gpu --output-on-failure # Sparse solvers (cuDSS -- separate install required)
diff --git a/contrib/test/GPU/CMakeLists.txt b/contrib/test/GPU/CMakeLists.txt index c856fad..b0f9437 100644 --- a/contrib/test/GPU/CMakeLists.txt +++ b/contrib/test/GPU/CMakeLists.txt
@@ -107,11 +107,14 @@ ei_add_gpu_test(cufft EXTRA_LIBS CUDA::cufft CUDA::cublas) endif() -# cuSPARSE SpMV (part of the CUDA toolkit) and the end-to-end CG test that -# exercises Eigen::ConjugateGradient with DeviceMatrix + device-resident SpMV. +# cuSPARSE SpMV (part of the CUDA toolkit), its BSR counterpart (skips itself +# below cuSPARSE 12.6.3), and the end-to-end CG test that exercises +# Eigen::ConjugateGradient with DeviceMatrix + device-resident SpMV. if(TARGET CUDA::cusparse) ei_add_gpu_test(cusparse_spmv EXTRA_LIBS CUDA::cusparse CUDA::cublas CUDA::cusolver) + ei_add_gpu_test(cusparse_bsr + EXTRA_LIBS CUDA::cusparse CUDA::cublas CUDA::cusolver) ei_add_gpu_test(cg EXTRA_LIBS CUDA::cusparse CUDA::cublas CUDA::cusolver CUDA::npps CUDA::nppc) endif()
diff --git a/contrib/test/GPU/cusparse_bsr.cpp b/contrib/test/GPU/cusparse_bsr.cpp new file mode 100644 index 0000000..5d9a892 --- /dev/null +++ b/contrib/test/GPU/cusparse_bsr.cpp
@@ -0,0 +1,302 @@ +// This file is part of Eigen, a lightweight C++ template library +// for linear algebra. +// +// Copyright (C) 2026 Rasmus Munk Larsen <rmlarsen@gmail.com> +// +// This Source Code Form is subject to the terms of the Mozilla +// Public License v. 2.0. If a copy of the MPL was not distributed +// with this file, You can obtain one at http://mozilla.org/MPL/2.0/. +// SPDX-License-Identifier: MPL-2.0 + +// Tests for gpu::SparseContext with BlockSparseMatrix input: GPU BSR SpMV/SpMM +// via cuSPARSE. + +#define EIGEN_USE_GPU +#include "main.h" +#include <Eigen/Sparse> +#include <contrib/Eigen/GPU> +#include "gpu_test_helpers.h" + +using namespace Eigen; + +#if EIGEN_HAS_CUSPARSE_BSR + +template <typename Scalar, int Options, int B> +using Bsm = BlockSparseMatrix<Scalar, Options, B, B, int>; + +// Random block pattern with random dense blocks; Random() fills both parts of +// a complex scalar, so conjugation is observable. +template <typename Scalar, int Options, int B> +Bsm<Scalar, Options, B> make_block_sparse(Index block_rows, Index block_cols, double density = 0.3) { + using Triplet = typename Bsm<Scalar, Options, B>::TripletType; + using Block = typename Bsm<Scalar, Options, B>::BlockType; + std::vector<Triplet> triplets; + for (Index bi = 0; bi < block_rows; ++bi) { + for (Index bj = 0; bj < block_cols; ++bj) { + if (internal::random<double>(0.0, 1.0) < density) triplets.emplace_back(int(bi), int(bj), Block::Random()); + } + } + Bsm<Scalar, Options, B> A(block_rows, block_cols); + A.setFromTriplets(triplets.begin(), triplets.end()); + return A; +} + +// Relative bound for a product by A: each entry sums at most +// max(A.rows(), A.cols()) terms, so 10 * max(A.rows(), A.cols()) * eps covers +// the accumulation error. +template <typename Lhs, typename Rhs, typename MatrixType> +void verify_product(const Lhs& gpu, const Rhs& cpu, const MatrixType& A) { + using Scalar = typename Lhs::Scalar; + using RealScalar = typename NumTraits<Scalar>::Real; + const RealScalar tol = RealScalar(10) * RealScalar((std::max)(A.rows(), A.cols())) * NumTraits<Scalar>::epsilon(); + VERIFY_IS_EQUAL(gpu.rows(), cpu.rows()); + VERIFY_IS_EQUAL(gpu.cols(), cpu.cols()); + VERIFY((gpu - cpu).norm() / (cpu.norm() + RealScalar(1)) < tol); +} + +// ---- Host SpMV: y = op(A) * x for every op ----------------------------------- + +template <typename Scalar, int Options, int B> +void test_bsr_spmv(Index block_rows, Index block_cols) { + using Vec = Matrix<Scalar, Dynamic, 1>; + + const Bsm<Scalar, Options, B> A = make_block_sparse<Scalar, Options, B>(block_rows, block_cols); + const Vec x = Vec::Random(A.cols()); + const Vec xt = Vec::Random(A.rows()); + + gpu::SparseContext<Scalar> ctx; + const Vec y = ctx.multiply(A, x); + verify_product(y, Vec(A * x), A); + const Vec yt = ctx.multiplyT(A, xt); + verify_product(yt, Vec(A.transpose() * xt), A); + const Vec yh = ctx.multiplyAdjoint(A, xt); + verify_product(yh, Vec(A.adjoint() * xt), A); +} + +// ---- Host SpMV in place: y = alpha * op(A) * x + beta * y --------------------- + +template <typename Scalar, int Options, int B> +void test_bsr_spmv_alpha_beta(Index block_rows, Index block_cols) { + using Vec = Matrix<Scalar, Dynamic, 1>; + + const Bsm<Scalar, Options, B> A = make_block_sparse<Scalar, Options, B>(block_rows, block_cols); + const Scalar alpha(2); + const Scalar beta(3); + gpu::SparseContext<Scalar> ctx; + + const Vec x = Vec::Random(A.cols()); + const Vec y_init = Vec::Random(A.rows()); + Vec y = y_init; + ctx.multiply(A, x, y, alpha, beta); + verify_product(y, Vec(alpha * (A * x) + beta * y_init), A); + + const Vec xt = Vec::Random(A.rows()); + const Vec yt_init = Vec::Random(A.cols()); + Vec yt = yt_init; + ctx.multiply(A, xt, yt, alpha, beta, gpu::GpuOp::ConjTrans); + verify_product(yt, Vec(alpha * (A.adjoint() * xt) + beta * yt_init), A); +} + +// ---- Host SpMM: Y = op(A) * X ------------------------------------------------- + +template <typename Scalar, int Options, int B> +void test_bsr_spmm(Index block_rows, Index block_cols, Index nrhs) { + using Mat = Matrix<Scalar, Dynamic, Dynamic>; + + const Bsm<Scalar, Options, B> A = make_block_sparse<Scalar, Options, B>(block_rows, block_cols); + const Mat X = Mat::Random(A.cols(), nrhs); + const Mat Xt = Mat::Random(A.rows(), nrhs); + + gpu::SparseContext<Scalar> ctx; + verify_product(ctx.multiplyMat(A, X), Mat(A * X), A); + verify_product(ctx.multiplyMat(A, Xt, gpu::GpuOp::Trans), Mat(A.transpose() * Xt), A); + verify_product(ctx.multiplyMat(A, Xt, gpu::GpuOp::ConjTrans), Mat(A.adjoint() * Xt), A); +} + +// ---- DeviceMatrix in/out with an explicit op ---------------------------------- + +template <typename Scalar, int Options, int B> +void test_bsr_device_multiply(Index block_rows, Index block_cols) { + using Vec = Matrix<Scalar, Dynamic, 1>; + + const Bsm<Scalar, Options, B> A = make_block_sparse<Scalar, Options, B>(block_rows, block_cols); + gpu::Context gctx; + gpu::SparseContext<Scalar> ctx(gctx); + + const Vec x = Vec::Random(A.cols()); + auto d_x = gpu::DeviceMatrix<Scalar>::fromHost(x, gctx.stream()); + gpu::DeviceMatrix<Scalar> d_y; + ctx.multiply(A, d_x, d_y); + verify_product(d_y.toHost(gctx.stream()), Vec(A * x), A); + + const Vec xt = Vec::Random(A.rows()); + auto d_xt = gpu::DeviceMatrix<Scalar>::fromHost(xt, gctx.stream()); + gpu::DeviceMatrix<Scalar> d_yt; + ctx.multiply(A, d_xt, d_yt, Scalar(1), Scalar(0), gpu::GpuOp::Trans); + verify_product(d_yt.toHost(gctx.stream()), Vec(A.transpose() * xt), A); + ctx.multiply(A, d_xt, d_yt, Scalar(1), Scalar(0), gpu::GpuOp::ConjTrans); + verify_product(d_yt.toHost(gctx.stream()), Vec(A.adjoint() * xt), A); +} + +// ---- deviceView: upload once, SpMV and SpMM by expression ---------------------- + +template <typename Scalar, int Options, int B> +void test_bsr_device_view(Index block_rows, Index block_cols, Index nrhs) { + using Vec = Matrix<Scalar, Dynamic, 1>; + using Mat = Matrix<Scalar, Dynamic, Dynamic>; + + const Bsm<Scalar, Options, B> A = make_block_sparse<Scalar, Options, B>(block_rows, block_cols); + gpu::Context gctx; + gpu::SparseContext<Scalar> ctx(gctx); + auto view = ctx.deviceView(A); + VERIFY(view.generation() == ctx.uploadGeneration()); + VERIFY_IS_EQUAL(view.rows(), A.rows()); + VERIFY_IS_EQUAL(view.cols(), A.cols()); + + const Vec x = Vec::Random(A.cols()); + auto d_x = gpu::DeviceMatrix<Scalar>::fromHost(x, gctx.stream()); + gpu::DeviceMatrix<Scalar> d_y = view * d_x; + verify_product(d_y.toHost(gctx.stream()), Vec(A * x), A); + + const Mat X = Mat::Random(A.cols(), nrhs); + auto d_X = gpu::DeviceMatrix<Scalar>::fromHost(X, gctx.stream()); + gpu::DeviceMatrix<Scalar> d_Y; + d_Y.noalias() = view * d_X; + verify_product(d_Y.toHost(gctx.stream()), Mat(A * X), A); + + // cuSPARSE runs no transposed BSR product; the exec entry points reject any + // other op against a BSR upload before queuing work. + VERIFY_RAISES_ASSERT(ctx.spmv_device_exec(d_x, d_y, Scalar(1), Scalar(0), gpu::GpuOp::Trans)); + VERIFY_RAISES_ASSERT(ctx.spmm_device_exec(d_X, d_Y, Scalar(1), Scalar(0), gpu::GpuOp::Trans)); +} + +// ---- Descriptor cache across formats ------------------------------------------ + +// A CSC matrix with the same shape and nonzero count as a BSR one must not +// reuse its descriptor: S = (A^T as scalar sparse) has both, and a different +// value for every x unless A is symmetric. +template <typename Scalar, int Options, int B> +void test_bsr_format_switch(Index block_n) { + using Vec = Matrix<Scalar, Dynamic, 1>; + using SpMat = SparseMatrix<Scalar, ColMajor, int>; + + const Bsm<Scalar, Options, B> A = make_block_sparse<Scalar, Options, B>(block_n, block_n); + const SpMat S = SpMat(A.transpose().toSparse()); + VERIFY_IS_EQUAL(S.nonZeros(), A.nonZeros()); + const Vec x = Vec::Random(A.cols()); + + gpu::SparseContext<Scalar> ctx; + verify_product(ctx.multiply(A, x), Vec(A * x), A); + verify_product(ctx.multiply(S, x), Vec(S * x), A); + verify_product(ctx.multiply(A, x), Vec(A * x), A); +} + +// ---- Pattern rewrite at unchanged host pointers ------------------------------- + +// The index arrays are re-uploaded on every host-input call, so rewriting the +// block-column indices in place is picked up. Blocks sit at (i, i) in an +// n x (n + 1) grid and move to (i, i + 1), which keeps every block row sorted. +template <typename Scalar, int Options, int B> +void test_bsr_pattern_rewrite(Index n) { + using Vec = Matrix<Scalar, Dynamic, 1>; + using Triplet = typename Bsm<Scalar, Options, B>::TripletType; + using Block = typename Bsm<Scalar, Options, B>::BlockType; + + std::vector<Triplet> triplets; + for (Index i = 0; i < n; ++i) triplets.emplace_back(int(i), int(i), Block::Random()); + Bsm<Scalar, Options, B> A(n, n + 1); + A.setFromTriplets(triplets.begin(), triplets.end()); + const Vec x = Vec::Random(A.cols()); + + gpu::SparseContext<Scalar> ctx; + verify_product(ctx.multiply(A, x), Vec(A * x), A); + + // RowMajor: inner index = block column, shift it. ColMajor: inner index = + // block row; shifting the *outer* boundaries instead moves block i to + // column i + 1, so both orders end at the same pattern. + if (Options & RowMajor) { + for (Index k = 0; k < A.nonZeroBlocks(); ++k) A.innerIndexPtr()[k] += 1; + } else { + for (Index j = n; j >= 1; --j) A.outerIndexPtr()[j] = A.outerIndexPtr()[j - 1]; + A.outerIndexPtr()[0] = 0; + A.outerIndexPtr()[n + 1] = int(n); + } + VERIFY_IS_EQUAL(A.nonZeroBlocks(), n); + verify_product(ctx.multiply(A, x), Vec(A * x), A); +} + +// ---- Empty matrices ------------------------------------------------------------ + +template <typename Scalar, int Options, int B> +void test_bsr_empty() { + using Vec = Matrix<Scalar, Dynamic, 1>; + using Mat = Matrix<Scalar, Dynamic, Dynamic>; + + gpu::Context gctx; + gpu::SparseContext<Scalar> ctx(gctx); + + const Bsm<Scalar, Options, B> A0(0, 0); + VERIFY_IS_EQUAL(ctx.multiply(A0, Vec(0)).size(), 0); + VERIFY_IS_EQUAL(ctx.multiplyMat(A0, Mat(0, 3)).cols(), 3); + + // Nonzero dimensions, no stored blocks: y <- beta * y on the host and + // y <- 0 on the device. + const Bsm<Scalar, Options, B> A(3, 2); + const Vec x = Vec::Random(A.cols()); + const Vec y_init = Vec::Random(A.rows()); + Vec y = y_init; + ctx.multiply(A, x, y, Scalar(1), Scalar(2)); + VERIFY_IS_APPROX(y, Vec(Scalar(2) * y_init)); + VERIFY(ctx.multiply(A, x).isZero()); + + auto d_x = gpu::DeviceMatrix<Scalar>::fromHost(x, gctx.stream()); + gpu::DeviceMatrix<Scalar> d_y; + ctx.multiply(A, d_x, d_y); + VERIFY_IS_EQUAL(d_y.rows(), A.rows()); + VERIFY(d_y.toHost(gctx.stream()).isZero()); +} + +// ---- Driver --------------------------------------------------------------------- + +template <typename Scalar, int Options, int B> +void test_order() { + CALL_SUBTEST((test_bsr_spmv<Scalar, Options, B>(7, 5))); + CALL_SUBTEST((test_bsr_spmv_alpha_beta<Scalar, Options, B>(5, 7))); + CALL_SUBTEST((test_bsr_spmm<Scalar, Options, B>(6, 4, 3))); + CALL_SUBTEST((test_bsr_device_multiply<Scalar, Options, B>(4, 6))); + CALL_SUBTEST((test_bsr_device_view<Scalar, Options, B>(6, 6, 4))); + CALL_SUBTEST((test_bsr_format_switch<Scalar, Options, B>(5))); + CALL_SUBTEST((test_bsr_pattern_rewrite<Scalar, Options, B>(4))); + CALL_SUBTEST((test_bsr_empty<Scalar, Options, B>())); +} + +template <typename Scalar> +void test_scalar() { + // RowMajor binds without a host copy for NoTrans, ColMajor for Trans; the + // other op / order combinations take the transposing copy. + test_order<Scalar, RowMajor, 2>(); + test_order<Scalar, ColMajor, 2>(); + test_order<Scalar, RowMajor, 3>(); + test_order<Scalar, ColMajor, 3>(); +} + +#endif // EIGEN_HAS_CUSPARSE_BSR + +EIGEN_DECLARE_TEST(gpu_cusparse_bsr) { +#if EIGEN_HAS_CUSPARSE_BSR + gpu_test::require_cusparse_context(); + + // Split by scalar so each part compiles in parallel. + CALL_SUBTEST_1(test_scalar<float>()); + CALL_SUBTEST_2(test_scalar<double>()); + CALL_SUBTEST_3(test_scalar<std::complex<float>>()); + CALL_SUBTEST_4(test_scalar<std::complex<double>>()); + // Block size 4, in one scalar type to bound compile time. + CALL_SUBTEST_5((test_order<double, RowMajor, 4>())); + CALL_SUBTEST_5((test_order<double, ColMajor, 4>())); +#else + std::cout << "SKIP: BSR products need cuSPARSE >= 12.6.3 (CUDA 13.0 Update 1); this build has CUSPARSE_VERSION " + << CUSPARSE_VERSION << std::endl; + std::exit(77); +#endif +}
diff --git a/contrib/test/GPU/cusparse_spmv.cpp b/contrib/test/GPU/cusparse_spmv.cpp index ce385ce..ef62165 100644 --- a/contrib/test/GPU/cusparse_spmv.cpp +++ b/contrib/test/GPU/cusparse_spmv.cpp
@@ -398,10 +398,83 @@ VERIFY((d_y.toHost() - y_ref).norm() / (y_ref.norm() + RealScalar(1)) < tol); } +// ---- BlockSparseMatrix input ------------------------------------------------- + +// Regression coverage for the CSC fallback: a ColMajor 1 x 1 or rectangular +// BlockSparseMatrix reached deviceView() through the implicit SparseMatrix +// conversion before the BSR overloads existed and must keep working on every +// toolkit. Square blocks of size >= 2 take BSR where EIGEN_HAS_CUSPARSE_BSR is 1. +template <typename Scalar, int Options, int BlockRows, int BlockCols> +void test_block_sparse_input(Index block_rows, Index block_cols) { + using Bsm = BlockSparseMatrix<Scalar, Options, BlockRows, BlockCols, int>; + using Block = typename Bsm::BlockType; + using Vec = Matrix<Scalar, Dynamic, 1>; + using Mat = Matrix<Scalar, Dynamic, Dynamic>; + using RealScalar = typename NumTraits<Scalar>::Real; + + std::vector<typename Bsm::TripletType> triplets; + for (Index bi = 0; bi < block_rows; ++bi) { + for (Index bj = 0; bj < block_cols; ++bj) { + if (internal::random<double>(0.0, 1.0) < 0.4) triplets.emplace_back(int(bi), int(bj), Block::Random()); + } + } + Bsm A(block_rows, block_cols); + A.setFromTriplets(triplets.begin(), triplets.end()); + const Vec x = Vec::Random(A.cols()); + const Vec xt = Vec::Random(A.rows()); + const Mat X = Mat::Random(A.cols(), 3); + const RealScalar tol = RealScalar(10) * RealScalar((std::max)(A.rows(), A.cols())) * NumTraits<Scalar>::epsilon(); + auto verify_close = [tol](const auto& result, const auto& ref) { + VERIFY((result - ref).norm() / (ref.norm() + RealScalar(1)) < tol); + }; + + gpu::Context gctx; + gpu::SparseContext<Scalar> ctx(gctx); + const Vec y_ref = A * x; + verify_close(ctx.multiply(A, x), y_ref); + const Vec yt_ref = A.transpose() * xt; + verify_close(ctx.multiplyT(A, xt), yt_ref); + Vec y_acc = Vec::Random(A.cols()); + const Vec y_acc_ref = Scalar(2) * yt_ref + Scalar(3) * y_acc; + ctx.multiply(A, xt, y_acc, Scalar(2), Scalar(3), gpu::GpuOp::Trans); + verify_close(y_acc, y_acc_ref); + // As for test_spmv_adjoint below: shapes on the CSC path cannot form A^H * x + // for complex scalars on cuSPARSE < 12, where SparseContext asserts. +#if !defined(CUSPARSE_VERSION) || CUSPARSE_VERSION >= 12000 + constexpr bool kTestAdjoint = true; +#else + constexpr bool kTestAdjoint = !NumTraits<Scalar>::IsComplex; +#endif + if (kTestAdjoint) verify_close(ctx.multiplyAdjoint(A, xt), Vec(A.adjoint() * xt)); + verify_close(ctx.multiplyMat(A, X), Mat(A * X)); + + auto d_x = gpu::DeviceMatrix<Scalar>::fromHost(x, gctx.stream()); + gpu::DeviceMatrix<Scalar> d_y; + ctx.multiply(A, d_x, d_y); + verify_close(d_y.toHost(gctx.stream()), y_ref); + + auto view = ctx.deviceView(A); + VERIFY_IS_EQUAL(view.rows(), A.rows()); + VERIFY_IS_EQUAL(view.cols(), A.cols()); + gpu::DeviceMatrix<Scalar> d_y2 = view * d_x; + verify_close(d_y2.toHost(gctx.stream()), y_ref); +} + +template <typename Scalar> +void test_block_sparse_inputs() { + CALL_SUBTEST((test_block_sparse_input<Scalar, ColMajor, 1, 1>(5, 4))); + CALL_SUBTEST((test_block_sparse_input<Scalar, RowMajor, 1, 1>(4, 5))); + CALL_SUBTEST((test_block_sparse_input<Scalar, ColMajor, 2, 3>(4, 3))); + CALL_SUBTEST((test_block_sparse_input<Scalar, RowMajor, 2, 3>(3, 4))); + CALL_SUBTEST((test_block_sparse_input<Scalar, ColMajor, 3, 3>(4, 4))); + CALL_SUBTEST((test_block_sparse_input<Scalar, RowMajor, 3, 3>(4, 4))); +} + // ---- Per-scalar driver ------------------------------------------------------ template <typename Scalar> void test_scalar() { + CALL_SUBTEST(test_block_sparse_inputs<Scalar>()); CALL_SUBTEST(test_device_spmm<Scalar>(64, 5)); CALL_SUBTEST(test_device_multiply_gpuop<Scalar>(64)); CALL_SUBTEST(test_spmv<Scalar>(64, 64));