| // 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 | |
| // GPU sparse matrix-vector (SpMV) and sparse matrix-dense matrix (SpMM) multiply | |
| // via cuSPARSE. | |
| // | |
| // SparseContext owns the cuSPARSE descriptors and device buffers. It takes | |
| // 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 + 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. | |
| #ifndef EIGEN_GPU_SPARSE_CONTEXT_H | |
| #define EIGEN_GPU_SPARSE_CONTEXT_H | |
| // IWYU pragma: private | |
| #include "./InternalHeaderCheck.h" | |
| #include <cstdint> | |
| #include "./CuSparseSupport.h" | |
| #include "./FwdDecl.h" | |
| 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. Adding or | |
| * subtracting a DeviceMatrix yields an SpMVAffineExpr. */ | |
| template <typename Scalar_> | |
| class SpMVExpr { | |
| public: | |
| using Scalar = Scalar_; | |
| SpMVExpr(const DeviceSparseView<Scalar>& view, const DeviceMatrix<Scalar>& x) : view_(view), x_(x) {} | |
| const DeviceSparseView<Scalar>& view() const { return view_; } | |
| const DeviceMatrix<Scalar>& x() const { return x_; } | |
| private: | |
| const DeviceSparseView<Scalar>& view_; | |
| const DeviceMatrix<Scalar>& x_; | |
| }; | |
| /** alpha * (DeviceSparseView * DeviceMatrix) + beta * addend with alpha, beta = ±1: the | |
| * result of `d_b - d_A * d_x`, `d_b + d_A * d_x`, `d_A * d_x + d_b` and `d_A * d_x - d_b`, | |
| * so that Eigen's iterative solver templates (`residual = rhs - mat * x`) compile. | |
| * Evaluated by DeviceMatrix::operator=(SpMVAffineExpr): the addend is copied into the | |
| * destination unless it already is the destination, then one cusparseSpMV (one column) | |
| * or cusparseSpMM call runs with the given alpha and beta. No operator takes an | |
| * SpMVAffineExpr, so an expression with a second addend does not compile. */ | |
| template <typename Scalar_> | |
| class SpMVAffineExpr { | |
| public: | |
| using Scalar = Scalar_; | |
| SpMVAffineExpr(const SpMVExpr<Scalar>& product, Scalar alpha, Scalar beta, const DeviceMatrix<Scalar>& addend) | |
| : view_(product.view()), x_(product.x()), alpha_(alpha), beta_(beta), addend_(addend) {} | |
| const DeviceSparseView<Scalar>& view() const { return view_; } | |
| const DeviceMatrix<Scalar>& x() const { return x_; } | |
| Scalar alpha() const { return alpha_; } | |
| Scalar beta() const { return beta_; } | |
| const DeviceMatrix<Scalar>& addend() const { return addend_; } | |
| private: | |
| const DeviceSparseView<Scalar>& view_; | |
| const DeviceMatrix<Scalar>& x_; | |
| Scalar alpha_; | |
| Scalar beta_; | |
| const DeviceMatrix<Scalar>& addend_; | |
| }; | |
| template <typename S> | |
| SpMVAffineExpr<S> operator-(const DeviceMatrix<S>& b, const SpMVExpr<S>& p) { | |
| return SpMVAffineExpr<S>(p, S(-1), S(1), b); | |
| } | |
| template <typename S> | |
| SpMVAffineExpr<S> operator+(const DeviceMatrix<S>& b, const SpMVExpr<S>& p) { | |
| return SpMVAffineExpr<S>(p, S(1), S(1), b); | |
| } | |
| template <typename S> | |
| SpMVAffineExpr<S> operator+(const SpMVExpr<S>& p, const DeviceMatrix<S>& b) { | |
| return b + p; | |
| } | |
| template <typename S> | |
| SpMVAffineExpr<S> operator-(const SpMVExpr<S>& p, const DeviceMatrix<S>& b) { | |
| return SpMVAffineExpr<S>(p, S(1), S(-1), b); | |
| } | |
| } // namespace gpu | |
| namespace internal { | |
| // DeviceSparseView is a matrix-free matrix type for Eigen's iterative solvers: it looks like a | |
| // SparseMatrix to the traits machinery, and IterativeSolverBase stores it by pointer because | |
| // Ref<> cannot bind to it (the same mechanism as the matrix-free example in the documentation). | |
| template <typename Scalar_> | |
| struct traits<gpu::DeviceSparseView<Scalar_>> : traits<SparseMatrix<Scalar_, ColMajor, int>> {}; | |
| } // namespace internal | |
| namespace gpu { | |
| /** 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. | |
| * | |
| * It is also a matrix-free matrix type for Eigen's iterative solvers: | |
| * `ConjugateGradient<DeviceSparseView<Scalar>, Lower | Upper>` runs Eigen's own | |
| * algorithm on device vectors through `solveWithGuessInPlace()`; see the README. | |
| * | |
| * \warning One SparseContext caches one sparse matrix at a time. Any later | |
| * upload through the same context — a second deviceView() or any host-input | |
| * multiply — replaces the cached data and invalidates earlier views; a stale | |
| * view is caught by an assert at evaluation time via a generation counter. | |
| * For multiple simultaneous sparse matrices, use separate SparseContext | |
| * instances (they can share a Context for same-stream execution). | |
| * | |
| * Supports `d_y = d_A * d_x` (SpMV) and `d_Y = d_A * d_X` (SpMM). */ | |
| template <typename Scalar_> | |
| class DeviceSparseView : public EigenBase<DeviceSparseView<Scalar_>> { | |
| public: | |
| using Scalar = Scalar_; | |
| using RealScalar = typename NumTraits<Scalar>::Real; | |
| using StorageIndex = int; | |
| using SpMat = SparseMatrix<Scalar, ColMajor, int>; | |
| // What IterativeSolverBase and ConjugateGradient read from a matrix type. | |
| static constexpr int ColsAtCompileTime = Dynamic; | |
| static constexpr int MaxColsAtCompileTime = Dynamic; | |
| static constexpr bool IsRowMajor = false; | |
| DeviceSparseView(SparseContext<Scalar>& ctx, Index rows, Index cols, uint64_t generation) | |
| : ctx_(ctx), rows_(rows), cols_(cols), generation_(generation) {} | |
| /** Sparse product expression: d_A * d_x. Evaluated by DeviceMatrix::operator=. */ | |
| SpMVExpr<Scalar> operator*(const DeviceMatrix<Scalar>& x) const { return SpMVExpr<Scalar>(*this, x); } | |
| Index rows() const { return rows_; } | |
| Index cols() const { return cols_; } | |
| /** Number of stored entries of the cached matrix. */ | |
| Index nonZeros() const { return ctx_.nonZeros(); } | |
| const SparseContext<Scalar>& context() const { return ctx_; } | |
| /** Upload generation this view was created against. Used to detect stale | |
| * views after the context has cached a different matrix. */ | |
| uint64_t generation() const { return generation_; } | |
| private: | |
| SparseContext<Scalar>& ctx_; | |
| Index rows_; | |
| Index cols_; | |
| uint64_t generation_; | |
| }; | |
| template <typename Scalar_> | |
| class SparseContext { | |
| public: | |
| using Scalar = Scalar_; | |
| using RealScalar = typename NumTraits<Scalar>::Real; | |
| using StorageIndex = int; | |
| 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) { | |
| EIGEN_CUDA_RUNTIME_CHECK(cudaStreamCreate(&stream_)); | |
| owns_stream_ = true; | |
| EIGEN_CUSPARSE_CHECK(cusparseCreate(&handle_)); | |
| EIGEN_CUSPARSE_CHECK(cusparseSetStream(handle_, stream_)); | |
| } | |
| /** Borrow a Context: shares stream and cuSPARSE handle. | |
| * The Context must outlive this SparseContext. */ | |
| explicit SparseContext(Context& ctx) | |
| : stream_(ctx.stream()), handle_(ctx.cusparseHandle()), owns_stream_(false), owns_handle_(false) {} | |
| ~SparseContext() { | |
| destroy_spmat_descriptor(/*checked=*/false); | |
| destroy_dense_descriptors(); | |
| if (owns_handle_ && handle_) (void)cusparseDestroy(handle_); | |
| if (owns_stream_ && stream_) (void)cudaStreamDestroy(stream_); | |
| } | |
| SparseContext(const SparseContext&) = delete; | |
| SparseContext& operator=(const SparseContext&) = delete; | |
| /** Upload a sparse matrix to device and return a lightweight view. | |
| * The sparse data is uploaded immediately and cached in this context. | |
| * The returned view can be used for repeated SpMV/SpMM without re-uploading. | |
| * If the matrix values change, call deviceView() again to re-upload. | |
| * | |
| * \warning One context caches one matrix. Any later upload — another | |
| * deviceView() or any host-input multiply — overwrites the previous upload | |
| * and invalidates earlier views (asserted at evaluation time). For multiple | |
| * simultaneous matrices, use separate SparseContext instances sharing the | |
| * same Context. | |
| * | |
| * Supports `d_y = d_A * d_x` (SpMV) and `d_Y = d_A * d_X` (SpMM). */ | |
| DeviceSparseView<Scalar> deviceView(const SpMat& A) { | |
| eigen_assert(A.isCompressed()); | |
| upload_sparse(A); | |
| return DeviceSparseView<Scalar>(*this, A.rows(), A.cols(), generation_); | |
| } | |
| /** Generation counter of the currently cached sparse matrix. Bumped on | |
| * every sparse upload (deviceView() or a host-input multiply). */ | |
| uint64_t uploadGeneration() const { return generation_; } | |
| /** Number of stored entries of the cached sparse matrix, -1 before the first upload. */ | |
| Index nonZeros() const { return cached_nnz_; } | |
| /** Compute y = A * x. Returns y as a new dense vector. */ | |
| template <typename InputType, typename Rhs> | |
| DenseVector multiply(const SparseMatrixBase<InputType>& A, const MatrixBase<Rhs>& x) { | |
| return multiply_host_return(A, x, GpuOp::NoTrans); | |
| } | |
| /** Compute y = alpha * op(A) * x + beta * y (in-place, host vectors). */ | |
| template <typename InputType, typename Rhs, typename Dest> | |
| void multiply(const SparseMatrixBase<InputType>& A, const MatrixBase<Rhs>& x, MatrixBase<Dest>& y, | |
| Scalar alpha = Scalar(1), Scalar beta = Scalar(0), GpuOp op = GpuOp::NoTrans) { | |
| const InputType& input = A.derived(); | |
| internal::check_storage_index_bounds<StorageIndex>(input.rows(), input.cols(), input.nonZeros()); | |
| SpMat storage; | |
| const SpMat& mat = internal::bind_sparse<SpMat>(input, storage); | |
| multiply_host_impl(mat, x.derived(), y.derived(), alpha, beta, internal::to_cusparse_op<Scalar>(op)); | |
| } | |
| /** Compute d_y = A * d_x. Device-resident dense vectors, no host transfer | |
| * for x/y. The sparse matrix (values and index arrays) is re-uploaded on | |
| * each call; for a fully device-resident sparse matrix use deviceView(). */ | |
| template <typename InputType> | |
| void multiply(const SparseMatrixBase<InputType>& 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). */ | |
| template <typename InputType> | |
| void multiply(const SparseMatrixBase<InputType>& A, const DeviceMatrix<Scalar>& d_x, DeviceMatrix<Scalar>& d_y, | |
| Scalar alpha, Scalar beta, GpuOp op = GpuOp::NoTrans) { | |
| const InputType& input = A.derived(); | |
| internal::check_storage_index_bounds<StorageIndex>(input.rows(), input.cols(), input.nonZeros()); | |
| SpMat storage; | |
| const SpMat& mat = internal::bind_sparse<SpMat>(input, storage); | |
| upload_sparse(mat); | |
| spmv_device_exec(d_x, d_y, alpha, beta, op); | |
| } | |
| /** Compute y = A^T * x (host vectors). */ | |
| template <typename InputType, typename Rhs> | |
| DenseVector multiplyT(const SparseMatrixBase<InputType>& A, const MatrixBase<Rhs>& x) { | |
| return multiply_host_return(A, x, GpuOp::Trans); | |
| } | |
| /** Compute y = A^H * x (conjugate transpose). For real Scalar this is equivalent to multiplyT. */ | |
| template <typename InputType, typename Rhs> | |
| DenseVector multiplyAdjoint(const SparseMatrixBase<InputType>& A, const MatrixBase<Rhs>& x) { | |
| return multiply_host_return(A, x, GpuOp::ConjTrans); | |
| } | |
| /** Compute Y = op(A) * X where X is a dense matrix (multiple RHS). Returns Y. */ | |
| template <typename InputType, typename Rhs> | |
| DenseMatrix multiplyMat(const SparseMatrixBase<InputType>& A, const MatrixBase<Rhs>& X, GpuOp op = GpuOp::NoTrans) { | |
| const InputType& input = A.derived(); | |
| internal::check_storage_index_bounds<StorageIndex>(input.rows(), input.cols(), input.nonZeros()); | |
| SpMat storage; | |
| const SpMat& mat = internal::bind_sparse<SpMat>(input, storage); | |
| const DenseMatrix rhs(X.derived()); | |
| const cusparseOperation_t cu_op = internal::to_cusparse_op<Scalar>(op); | |
| const Index m = (op == GpuOp::NoTrans) ? mat.rows() : mat.cols(); | |
| const Index k = (op == GpuOp::NoTrans) ? mat.cols() : mat.rows(); | |
| eigen_assert(k == rhs.rows()); | |
| const Index n = rhs.cols(); | |
| if (m == 0 || n == 0 || mat.nonZeros() == 0) return DenseMatrix::Zero(m, n); | |
| DenseMatrix Y = DenseMatrix::Zero(m, n); | |
| spmm_impl(mat, rhs, Y, Scalar(1), Scalar(0), cu_op); | |
| 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: | |
| cudaStream_t stream_ = nullptr; | |
| cusparseHandle_t handle_ = nullptr; | |
| bool owns_stream_ = false; | |
| bool owns_handle_ = false; | |
| // Cached device buffers for sparse matrix (grow-only). | |
| internal::DeviceBuffer d_outerPtr_; | |
| internal::DeviceBuffer d_innerIdx_; | |
| internal::DeviceBuffer d_values_; | |
| // Cached device buffers for host-API dense vectors (grow-only). | |
| internal::DeviceBuffer d_x_; | |
| internal::DeviceBuffer d_y_; | |
| mutable internal::DeviceBuffer d_workspace_; | |
| // 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. | |
| uint64_t generation_ = 0; | |
| // Cached dense-vector/matrix descriptors, re-pointed per call and recreated | |
| // only when the shape changes. | |
| mutable cusparseDnVecDescr_t x_vec_desc_ = nullptr; | |
| mutable cusparseDnVecDescr_t y_vec_desc_ = nullptr; | |
| mutable int64_t x_vec_size_ = -1; | |
| mutable int64_t y_vec_size_ = -1; | |
| mutable cusparseDnMatDescr_t x_mat_desc_ = nullptr; | |
| mutable cusparseDnMatDescr_t y_mat_desc_ = nullptr; | |
| mutable int64_t x_mat_rows_ = -1, x_mat_cols_ = -1; | |
| mutable int64_t y_mat_rows_ = -1, y_mat_cols_ = -1; | |
| // Cached workspace-size query results, indexed by cusparseOperation_t. | |
| // Invalidated when the sparse descriptor or a dense-descriptor shape changes. | |
| static constexpr size_t kWsUnknown = static_cast<size_t>(-1); | |
| 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) { | |
| const InputType& input = A.derived(); | |
| internal::check_storage_index_bounds<StorageIndex>(input.rows(), input.cols(), input.nonZeros()); | |
| SpMat storage; | |
| const SpMat& mat = internal::bind_sparse<SpMat>(input, storage); | |
| DenseVector y((op == GpuOp::NoTrans) ? mat.rows() : mat.cols()); | |
| y.setZero(); | |
| multiply_host_impl(mat, x.derived(), y, Scalar(1), Scalar(0), internal::to_cusparse_op<Scalar>(op)); | |
| return y; | |
| } | |
| template <typename RhsDerived, typename DestDerived> | |
| void multiply_host_impl(const SpMat& A, const RhsDerived& x, DestDerived& y, Scalar alpha, Scalar beta, | |
| cusparseOperation_t op) { | |
| eigen_assert(A.isCompressed()); | |
| const Index m = A.rows(); | |
| const Index n = A.cols(); | |
| const Index nnz = A.nonZeros(); | |
| const Index x_size = (op == CUSPARSE_OPERATION_NON_TRANSPOSE) ? n : m; | |
| const Index y_size = (op == CUSPARSE_OPERATION_NON_TRANSPOSE) ? m : n; | |
| eigen_assert(x.size() == x_size); | |
| eigen_assert(y.size() == y_size); | |
| if (m == 0 || n == 0 || nnz == 0) { | |
| 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. | |
| const Ref<const DenseVector> x_ref(x); | |
| EIGEN_CUDA_RUNTIME_CHECK( | |
| cudaMemcpyAsync(d_x_.get(), x_ref.data(), x_size * sizeof(Scalar), cudaMemcpyHostToDevice, stream_)); | |
| ensure_buffer(d_y_, static_cast<size_t>(y_size) * sizeof(Scalar)); | |
| if (beta != Scalar(0)) { | |
| const Ref<const DenseVector> y_ref(y); | |
| EIGEN_CUDA_RUNTIME_CHECK( | |
| cudaMemcpyAsync(d_y_.get(), y_ref.data(), y_size * sizeof(Scalar), cudaMemcpyHostToDevice, stream_)); | |
| } | |
| exec_spmv(x_size, y_size, d_x_.get(), d_y_.get(), alpha, beta, op); | |
| // The spacing of y's elements. innerStride() is not it for a vector shaped only at run time: a 1 x n block of a | |
| // column-major matrix has innerStride() == 1 with its elements outerStride() apart. | |
| const Index incy = y.cols() == 1 ? y.rowStride() : y.colStride(); | |
| if (incy == 1) { | |
| EIGEN_CUDA_RUNTIME_CHECK( | |
| cudaMemcpyAsync(y.data(), d_y_.get(), y_size * sizeof(Scalar), cudaMemcpyDeviceToHost, stream_)); | |
| } else { | |
| // y is a strided block, e.g. a row of a column-major matrix. | |
| EIGEN_CUDA_RUNTIME_CHECK(cudaMemcpy2DAsync(y.data(), static_cast<size_t>(incy) * sizeof(Scalar), d_y_.get(), | |
| sizeof(Scalar), sizeof(Scalar), static_cast<size_t>(y_size), | |
| cudaMemcpyDeviceToHost, stream_)); | |
| } | |
| EIGEN_CUDA_RUNTIME_CHECK(cudaStreamSynchronize(stream_)); | |
| } | |
| 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(). | |
| * 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"); | |
| const cusparseOperation_t cu_op = internal::to_cusparse_op<Scalar>(op); | |
| const Index m = cached_rows_; | |
| const Index n = cached_cols_; | |
| const Index x_size = (cu_op == CUSPARSE_OPERATION_NON_TRANSPOSE) ? n : m; | |
| const Index y_size = (cu_op == CUSPARSE_OPERATION_NON_TRANSPOSE) ? m : n; | |
| eigen_assert(d_x.rows() * d_x.cols() == x_size); | |
| if (m == 0 || n == 0 || cached_nnz_ == 0) { | |
| // Empty A reduces SpMV to y <- beta*y; SparseContext owns no cuBLAS | |
| // handle for the scale, so the beta != 0 case must be handled by the caller. | |
| eigen_assert(beta == Scalar(0) && "SpMV with empty A and beta != 0 is unsupported; scale d_y externally"); | |
| if (d_y.rows() * d_y.cols() != y_size) d_y.resize(y_size, 1); | |
| d_y.setZero(stream_); | |
| return; | |
| } | |
| // Ensure d_y is allocated. | |
| if (d_y.rows() * d_y.cols() != y_size) { | |
| d_y.resize(y_size, 1); | |
| } | |
| // Wait for input data to be ready on this stream. | |
| d_x.waitReady(stream_); | |
| d_y.waitReady(stream_); | |
| exec_spmv(x_size, y_size, const_cast<void*>(static_cast<const void*>(d_x.data())), static_cast<void*>(d_y.data()), | |
| alpha, beta, cu_op); | |
| d_y.recordReady(stream_); | |
| } | |
| /** 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. | |
| * 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); | |
| const bool transposed = (cu_op != CUSPARSE_OPERATION_NON_TRANSPOSE); | |
| const Index m_op = transposed ? cached_cols_ : cached_rows_; | |
| const Index k_op = transposed ? cached_rows_ : cached_cols_; | |
| const Index n = d_X.cols(); | |
| eigen_assert(d_X.rows() == k_op); | |
| if (m_op == 0 || n == 0 || cached_nnz_ == 0) { | |
| eigen_assert(beta == Scalar(0) && "SpMM with empty A and beta != 0 is unsupported; scale d_Y externally"); | |
| if (d_Y.rows() != m_op || d_Y.cols() != n) d_Y.resize(m_op, n); | |
| d_Y.setZero(stream_); | |
| return; | |
| } | |
| if (d_Y.rows() != m_op || d_Y.cols() != n) { | |
| d_Y.resize(m_op, n); | |
| } | |
| d_X.waitReady(stream_); | |
| d_Y.waitReady(stream_); | |
| exec_spmm(m_op, k_op, n, const_cast<void*>(static_cast<const void*>(d_X.data())), static_cast<void*>(d_Y.data()), | |
| alpha, beta, cu_op); | |
| d_Y.recordReady(stream_); | |
| } | |
| private: | |
| // cuSPARSE 11.x's cusparseSpMM rejects CSC for matA (CSC support landed in | |
| // CUDA 12.0). On 11.x we register the same buffers as CSR-of-A^T (dims | |
| // swapped) and invert the user-facing op before each cuSPARSE call. On 12+ | |
| // we keep the natural CSC path so users pay no extra cost. | |
| #if !defined(CUSPARSE_VERSION) || CUSPARSE_VERSION < 12000 | |
| static constexpr bool kUseCsrOfTranspose = true; | |
| static constexpr cusparseSpMMAlg_t kSpMMAlg = CUSPARSE_SPMM_CSR_ALG2; | |
| #else | |
| static constexpr bool kUseCsrOfTranspose = false; | |
| static constexpr cusparseSpMMAlg_t kSpMMAlg = CUSPARSE_SPMM_ALG_DEFAULT; | |
| #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, 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) { | |
| case CUSPARSE_OPERATION_NON_TRANSPOSE: | |
| return CUSPARSE_OPERATION_TRANSPOSE; | |
| case CUSPARSE_OPERATION_TRANSPOSE: | |
| return CUSPARSE_OPERATION_NON_TRANSPOSE; | |
| default: | |
| // CONJUGATE_TRANSPOSE on the CSR-of-A^T descriptor would compute | |
| // conj(A) * x, not A^H * x — not supported via this representation. | |
| eigen_assert(false && "CUSPARSE_OPERATION_CONJUGATE_TRANSPOSE not supported on cuSPARSE < 12.0"); | |
| return user_op; | |
| } | |
| } | |
| static int op_index(cusparseOperation_t op) { | |
| switch (op) { | |
| case CUSPARSE_OPERATION_TRANSPOSE: | |
| return 1; | |
| case CUSPARSE_OPERATION_CONJUGATE_TRANSPOSE: | |
| return 2; | |
| default: | |
| return 0; | |
| } | |
| } | |
| void invalidate_ws_caches() const { | |
| for (int i = 0; i < 3; ++i) { | |
| spmv_ws_size_[i] = kWsUnknown; | |
| spmm_ws_size_[i] = kWsUnknown; | |
| } | |
| } | |
| // Recreate the descriptor when the size changes; otherwise just re-point it | |
| // at the new device buffer (cusparseDnVecSetValues is a host-side pointer | |
| // update, no GPU work). | |
| 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, kValueType)); | |
| cur_size = size; | |
| invalidate_ws_caches(); | |
| } else { | |
| EIGEN_CUSPARSE_CHECK(cusparseDnVecSetValues(desc, ptr)); | |
| } | |
| } | |
| void update_dnmat(cusparseDnMatDescr_t& desc, int64_t& cur_rows, int64_t& cur_cols, int64_t rows, int64_t cols, | |
| void* ptr) const { | |
| 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, kValueType, CUSPARSE_ORDER_COL)); | |
| cur_rows = rows; | |
| cur_cols = cols; | |
| invalidate_ws_caches(); | |
| } else { | |
| EIGEN_CUSPARSE_CHECK(cusparseDnMatSetValues(desc, ptr)); | |
| } | |
| } | |
| void destroy_dense_descriptors() { | |
| if (x_vec_desc_) (void)cusparseDestroyDnVec(x_vec_desc_); | |
| if (y_vec_desc_) (void)cusparseDestroyDnVec(y_vec_desc_); | |
| if (x_mat_desc_) (void)cusparseDestroyDnMat(x_mat_desc_); | |
| if (y_mat_desc_) (void)cusparseDestroyDnMat(y_mat_desc_); | |
| x_vec_desc_ = y_vec_desc_ = nullptr; | |
| x_mat_desc_ = y_mat_desc_ = nullptr; | |
| } | |
| 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 { | |
| 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); | |
| 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_, | |
| 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_, 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 { | |
| 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); | |
| update_dnmat(y_mat_desc_, y_mat_rows_, y_mat_cols_, m_op, n, d_y_ptr); | |
| 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_, 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_, kValueType, kSpMMAlg, d_workspace_.get())); | |
| } | |
| void spmm_impl(const SpMat& A, const DenseMatrix& X, DenseMatrix& Y, Scalar alpha, Scalar beta, | |
| cusparseOperation_t op) { | |
| eigen_assert(A.isCompressed()); | |
| // For op != NON_TRANSPOSE, Y = op(A) * X. The dense X / Y descriptors must | |
| // describe the *post-op* shapes: X has k_op rows (= input dim of op(A)), | |
| // Y has m_op rows (= output dim of op(A)). | |
| const bool transposed = (op != CUSPARSE_OPERATION_NON_TRANSPOSE); | |
| const Index m_op = transposed ? A.cols() : A.rows(); | |
| const Index k_op = transposed ? A.rows() : A.cols(); | |
| const Index n = X.cols(); | |
| const Index nnz = A.nonZeros(); | |
| if (m_op == 0 || n == 0 || k_op == 0 || nnz == 0) { | |
| scale_host(Y, beta); | |
| return; | |
| } | |
| upload_sparse(A); | |
| spmm_host_cached(X, Y, m_op, k_op, alpha, beta, op); | |
| } | |
| // 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); | |
| ensure_buffer(d_y_, y_bytes); | |
| EIGEN_CUDA_RUNTIME_CHECK(cudaMemcpyAsync(d_x_.get(), X.data(), x_bytes, cudaMemcpyHostToDevice, stream_)); | |
| if (beta != Scalar(0)) { | |
| EIGEN_CUDA_RUNTIME_CHECK(cudaMemcpyAsync(d_y_.get(), Y.data(), y_bytes, cudaMemcpyHostToDevice, stream_)); | |
| } | |
| exec_spmm(m_op, k_op, n, d_x_.get(), d_y_.get(), alpha, beta, op); | |
| EIGEN_CUDA_RUNTIME_CHECK(cudaMemcpyAsync(Y.data(), d_y_.get(), y_bytes, cudaMemcpyDeviceToHost, stream_)); | |
| EIGEN_CUDA_RUNTIME_CHECK(cudaStreamSynchronize(stream_)); | |
| } | |
| void upload_sparse(const SpMat& A) { | |
| // cuSPARSE 12.0+ accepts CSC directly. On cuSPARSE 11.x, cusparseSpMM | |
| // rejects CSC and CONJUGATE_TRANSPOSE on CSC+complex SpMV silently | |
| // 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. | |
| 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); | |
| } | |
| #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>(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, block size). Every upload invalidates | |
| // outstanding DeviceSparseViews via the generation counter. | |
| ++generation_; | |
| ensure_buffer(d_values_, val_bytes); | |
| ensure_buffer(d_outerPtr_, outer_bytes); | |
| ensure_buffer(d_innerIdx_, inner_bytes); | |
| EIGEN_CUDA_RUNTIME_CHECK(cudaMemcpyAsync(d_values_.get(), host_values, val_bytes, cudaMemcpyHostToDevice, stream_)); | |
| EIGEN_CUDA_RUNTIME_CHECK( | |
| 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, 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_; | |
| } | |
| 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. | |
| // `checked` selects checking (mid-lifetime rebuilds) vs swallowing failure | |
| // (noexcept destructor). | |
| void destroy_spmat_descriptor(bool checked) { | |
| if (spmat_desc_) { | |
| if (checked) | |
| EIGEN_CUSPARSE_CHECK(cusparseDestroySpMat(spmat_desc_)); | |
| else | |
| (void)cusparseDestroySpMat(spmat_desc_); | |
| spmat_desc_ = nullptr; | |
| } | |
| set_descriptor_key(-1, -1, -1, 0); | |
| invalidate_ws_caches(); | |
| } | |
| void ensure_buffer(internal::DeviceBuffer& buf, size_t needed) const { | |
| if (needed > buf.size()) { | |
| if (buf) EIGEN_CUDA_RUNTIME_CHECK(cudaStreamSynchronize(stream_)); | |
| buf = internal::DeviceBuffer(needed); | |
| } | |
| } | |
| }; | |
| // Defined here because it needs the full SparseContext definition. | |
| template <typename Scalar_> | |
| DeviceMatrix<Scalar_>& DeviceMatrix<Scalar_>::operator=(const SpMVExpr<Scalar_>& expr) { | |
| // Uses the sparse matrix already uploaded by deviceView() — no re-upload on | |
| // repeated products with the same view. A stale view (the context has since | |
| // uploaded again, replacing the cached data) is caught here. | |
| eigen_assert(expr.view().generation() == expr.view().context().uploadGeneration() && | |
| "DeviceSparseView is stale: its SparseContext has since uploaded another sparse matrix"); | |
| if (expr.x().cols() <= 1) { | |
| expr.view().context().spmv_device_exec(expr.x(), *this, Scalar_(1), Scalar_(0), GpuOp::NoTrans); | |
| } else { | |
| expr.view().context().spmm_device_exec(expr.x(), *this, Scalar_(1), Scalar_(0), GpuOp::NoTrans); | |
| } | |
| return *this; | |
| } | |
| template <typename Scalar_> | |
| DeviceMatrix<Scalar_>& DeviceMatrix<Scalar_>::operator=(const SpMVAffineExpr<Scalar_>& expr) { | |
| const DeviceSparseView<Scalar_>& view = expr.view(); | |
| const DeviceMatrix& addend = expr.addend(); | |
| eigen_assert(view.generation() == view.context().uploadGeneration() && | |
| "DeviceSparseView is stale: its SparseContext has since uploaded another sparse matrix"); | |
| eigen_assert(addend.rows() == view.rows() && addend.cols() == expr.x().cols() && | |
| "SpMVAffineExpr: the addend must have the shape of the product"); | |
| eigen_assert(&expr.x() != this && "SpMVAffineExpr: the destination aliases the dense operand"); | |
| // The product accumulates into a copy of the addend; when the addend is the destination | |
| // itself (d_r = d_r - d_A * d_x) cuSPARSE's beta accumulates in place. | |
| if (&addend != this) copyFrom(Context::threadLocal(), addend); | |
| // With no stored entries the expression is beta * addend, which spmv_device_exec cannot | |
| // form: SparseContext owns no cuBLAS handle to scale d_y by beta. | |
| if (view.nonZeros() == 0 || view.rows() == 0 || view.cols() == 0) { | |
| if (expr.beta() != Scalar_(1)) scale(Context::threadLocal(), expr.beta()); | |
| return *this; | |
| } | |
| if (expr.x().cols() <= 1) { | |
| view.context().spmv_device_exec(expr.x(), *this, expr.alpha(), expr.beta(), GpuOp::NoTrans); | |
| } else { | |
| view.context().spmm_device_exec(expr.x(), *this, expr.alpha(), expr.beta(), GpuOp::NoTrans); | |
| } | |
| return *this; | |
| } | |
| template <typename Scalar_> | |
| DeviceMatrix<Scalar_>::DeviceMatrix(const SpMVExpr<Scalar_>& expr) : DeviceMatrix() { | |
| *this = expr; | |
| } | |
| template <typename Scalar_> | |
| DeviceMatrix<Scalar_>::DeviceMatrix(const SpMVAffineExpr<Scalar_>& expr) : DeviceMatrix() { | |
| *this = expr; | |
| } | |
| } // namespace gpu | |
| } // namespace Eigen | |
| #endif // EIGEN_GPU_SPARSE_CONTEXT_H |