Eigen-Contrib  5.0.1
 
Loading...
Searching...
No Matches
GpuSparseContext.h
1// This file is part of Eigen, a lightweight C++ template library
2// for linear algebra.
3//
4// Copyright (C) 2026 Rasmus Munk Larsen <rmlarsen@gmail.com>
5//
6// This Source Code Form is subject to the terms of the Mozilla
7// Public License v. 2.0. If a copy of the MPL was not distributed
8// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
9// SPDX-License-Identifier: MPL-2.0
10
11// GPU sparse matrix-vector (SpMV) and sparse matrix-dense matrix (SpMM) multiply
12// via cuSPARSE.
13//
14// SparseContext owns the cuSPARSE descriptors and device buffers. It takes
15// SparseMatrix<Scalar, ColMajor> (CSC), implicitly converting RowMajor input, and
16// can borrow a gpu::Context so that sparse products share a stream with BLAS-1
17// operations — which removes the cross-stream event waits in solvers like CG.
18// It also takes BlockSparseMatrix: as BSR (internal::BsrBinding) where
19// internal::use_cusparse_bsr holds, as the scalar-level toSparse() copy otherwise.
20//
21// Caching: host-input calls re-upload the values *and* index arrays on every
22// call. Host pointer identity cannot detect a sparsity pattern rewritten in
23// place or assigned into the same allocations (SparseMatrix reuses them for
24// same-shape assignments), so the structure is never assumed unchanged. What
25// is cached are the cuSPARSE descriptors (sparse descriptor keyed on
26// dimensions + nonzero count + block size, dense descriptors keyed on shape)
27// and the SpMV/SpMM workspace-size queries. For repeated products with no
28// re-upload at all, use deviceView().
29//
30// Not thread-safe: concurrent multiply* calls on one instance race on the
31// cuSPARSE handle, the bound stream, and the cached buffers. Use one per thread.
32
33#ifndef EIGEN_GPU_SPARSE_CONTEXT_H
34#define EIGEN_GPU_SPARSE_CONTEXT_H
35
36// IWYU pragma: private
37#include "./InternalHeaderCheck.h"
38
39#include <cstdint>
40
41#include "./CuSparseSupport.h"
42#include "./FwdDecl.h"
43
44namespace Eigen {
45namespace gpu {
46
47namespace internal {
53template <int BlockRows, int BlockCols>
55 : Eigen::internal::bool_constant<(EIGEN_HAS_CUSPARSE_BSR != 0 && BlockRows == BlockCols && BlockRows >= 2)> {};
56} // namespace internal
57
58#if EIGEN_HAS_CUSPARSE_BSR
59namespace internal {
60
64template <typename Scalar, typename StorageIndex>
65struct BsrArrays {
66 Index brows;
67 Index bcols;
68 Index bnnz;
69 Index blockSize;
70 const StorageIndex* outer;
71 const StorageIndex* inner;
72 const Scalar* values;
73
74 Index rows() const { return brows * blockSize; }
75 Index cols() const { return bcols * blockSize; }
76 Index nonZeros() const { return bnnz * blockSize * blockSize; }
77};
78
85template <typename Bsm>
86BsrArrays<typename Bsm::Scalar, typename Bsm::StorageIndex> bsr_of(const Bsm& M) {
87 return BsrArrays<typename Bsm::Scalar, typename Bsm::StorageIndex>{
88 /*brows=*/M.blockOuterSize(), /*bcols=*/M.blockInnerSize(),
89 /*bnnz=*/M.nonZeroBlocks(), /*blockSize=*/Index(Bsm::BlockRows),
90 /*outer=*/M.outerIndexPtr(), /*inner=*/M.innerIndexPtr(),
91 /*values=*/M.valuePtr()};
92}
93
99template <typename Bsm>
100class BsrBinding {
101 EIGEN_STATIC_ASSERT((use_cusparse_bsr<int(Bsm::BlockRows), int(Bsm::BlockCols)>::value),
102 CUSPARSE_BSR_REQUIRES_SQUARE_BLOCKS_OF_SIZE_AT_LEAST_2)
103
104 public:
105 using Scalar = typename Bsm::Scalar;
106 using Arrays = BsrArrays<Scalar, typename Bsm::StorageIndex>;
107
108 BsrBinding(const Bsm& A, GpuOp op) {
109 if (op == GpuOp::ConjTrans && !NumTraits<Scalar>::IsComplex) op = GpuOp::Trans;
110 const bool direct = Bsm::IsRowMajor ? (op == GpuOp::NoTrans) : (op == GpuOp::Trans);
111 if (direct) {
112 arrays_ = bsr_of(A);
113 return;
114 }
115 if (op == GpuOp::ConjTrans && !Bsm::IsRowMajor) {
116 // ColMajor conj(A) reads as BSR of conj(A)^T = A^H.
117 copy_ = A.unaryExpr(Eigen::internal::scalar_conjugate_op<Scalar>());
118 } else if (op == GpuOp::ConjTrans) {
119 copy_ = A.adjoint();
120 } else {
121 copy_ = A.transpose();
122 }
123 arrays_ = bsr_of(copy_);
124 }
125
126 BsrBinding(const BsrBinding&) = delete;
127 BsrBinding& operator=(const BsrBinding&) = delete;
128
129 const Arrays& arrays() const { return arrays_; }
130
131 private:
132 // transpose()/adjoint() swap the block dimensions, which are equal here, so
133 // the copy has A's own type in either storage order.
134 Bsm copy_;
135 Arrays arrays_{};
136};
137
138} // namespace internal
139#endif // EIGEN_HAS_CUSPARSE_BSR
140
145template <typename Scalar_>
146class SpMVExpr {
147 public:
148 using Scalar = Scalar_;
149 SpMVExpr(const DeviceSparseView<Scalar>& view, const DeviceMatrix<Scalar>& x) : view_(view), x_(x) {}
150 const DeviceSparseView<Scalar>& view() const { return view_; }
151 const DeviceMatrix<Scalar>& x() const { return x_; }
152
153 private:
154 const DeviceSparseView<Scalar>& view_;
155 const DeviceMatrix<Scalar>& x_;
156};
157
165template <typename Scalar_>
166class SpMVAffineExpr {
167 public:
168 using Scalar = Scalar_;
169 SpMVAffineExpr(const SpMVExpr<Scalar>& product, Scalar alpha, Scalar beta, const DeviceMatrix<Scalar>& addend)
170 : view_(product.view()), x_(product.x()), alpha_(alpha), beta_(beta), addend_(addend) {}
171 const DeviceSparseView<Scalar>& view() const { return view_; }
172 const DeviceMatrix<Scalar>& x() const { return x_; }
173 Scalar alpha() const { return alpha_; }
174 Scalar beta() const { return beta_; }
175 const DeviceMatrix<Scalar>& addend() const { return addend_; }
176
177 private:
178 const DeviceSparseView<Scalar>& view_;
179 const DeviceMatrix<Scalar>& x_;
180 Scalar alpha_;
181 Scalar beta_;
182 const DeviceMatrix<Scalar>& addend_;
183};
184
185template <typename S>
186SpMVAffineExpr<S> operator-(const DeviceMatrix<S>& b, const SpMVExpr<S>& p) {
187 return SpMVAffineExpr<S>(p, S(-1), S(1), b);
188}
189
190template <typename S>
191SpMVAffineExpr<S> operator+(const DeviceMatrix<S>& b, const SpMVExpr<S>& p) {
192 return SpMVAffineExpr<S>(p, S(1), S(1), b);
193}
194
195template <typename S>
196SpMVAffineExpr<S> operator+(const SpMVExpr<S>& p, const DeviceMatrix<S>& b) {
197 return b + p;
198}
199
200template <typename S>
201SpMVAffineExpr<S> operator-(const SpMVExpr<S>& p, const DeviceMatrix<S>& b) {
202 return SpMVAffineExpr<S>(p, S(1), S(-1), b);
203}
204
205} // namespace gpu
206
207namespace internal {
208// DeviceSparseView is a matrix-free matrix type for Eigen's iterative solvers: it looks like a
209// SparseMatrix to the traits machinery, and IterativeSolverBase stores it by pointer because
210// Ref<> cannot bind to it (the same mechanism as the matrix-free example in the documentation).
211template <typename Scalar_>
212struct traits<gpu::DeviceSparseView<Scalar_>> : traits<SparseMatrix<Scalar_, ColMajor, int>> {};
213} // namespace internal
214
215namespace gpu {
216
233template <typename Scalar_>
234class DeviceSparseView : public EigenBase<DeviceSparseView<Scalar_>> {
235 public:
236 using Scalar = Scalar_;
237 using RealScalar = typename NumTraits<Scalar>::Real;
238 using StorageIndex = int;
240 // What IterativeSolverBase and ConjugateGradient read from a matrix type.
241 static constexpr int ColsAtCompileTime = Dynamic;
242 static constexpr int MaxColsAtCompileTime = Dynamic;
243 static constexpr bool IsRowMajor = false;
244
245 DeviceSparseView(SparseContext<Scalar>& ctx, Index rows, Index cols, uint64_t generation)
246 : ctx_(ctx), rows_(rows), cols_(cols), generation_(generation) {}
247
250
251 Index rows() const { return rows_; }
252 Index cols() const { return cols_; }
254 Index nonZeros() const { return ctx_.nonZeros(); }
255 const SparseContext<Scalar>& context() const { return ctx_; }
256
259 uint64_t generation() const { return generation_; }
260
261 private:
262 SparseContext<Scalar>& ctx_;
263 Index rows_;
264 Index cols_;
265 uint64_t generation_;
266};
267
268template <typename Scalar_>
269class SparseContext {
270 public:
271 using Scalar = Scalar_;
272 using RealScalar = typename NumTraits<Scalar>::Real;
273 using StorageIndex = int;
275 using DenseVector = Matrix<Scalar, Dynamic, 1>;
282 template <int Options, int BlockRows, int BlockCols>
284
286 SparseContext() : owns_handle_(true) {
287 EIGEN_CUDA_RUNTIME_CHECK(cudaStreamCreate(&stream_));
288 owns_stream_ = true;
289 EIGEN_CUSPARSE_CHECK(cusparseCreate(&handle_));
290 EIGEN_CUSPARSE_CHECK(cusparseSetStream(handle_, stream_));
291 }
292
295 explicit SparseContext(Context& ctx)
296 : stream_(ctx.stream()), handle_(ctx.cusparseHandle()), owns_stream_(false), owns_handle_(false) {}
297
298 ~SparseContext() {
299 destroy_spmat_descriptor(/*checked=*/false);
300 destroy_dense_descriptors();
301 if (owns_handle_ && handle_) (void)cusparseDestroy(handle_);
302 if (owns_stream_ && stream_) (void)cudaStreamDestroy(stream_);
303 }
304
305 SparseContext(const SparseContext&) = delete;
306 SparseContext& operator=(const SparseContext&) = delete;
307
320 DeviceSparseView<Scalar> deviceView(const SpMat& A) {
321 eigen_assert(A.isCompressed());
322 upload_sparse(A);
323 return DeviceSparseView<Scalar>(*this, A.rows(), A.cols(), generation_);
324 }
325
328 uint64_t uploadGeneration() const { return generation_; }
329
331 Index nonZeros() const { return cached_nnz_; }
332
334 template <typename InputType, typename Rhs>
335 DenseVector multiply(const SparseMatrixBase<InputType>& A, const MatrixBase<Rhs>& x) {
336 return multiply_host_return(A, x, GpuOp::NoTrans);
337 }
338
340 template <typename InputType, typename Rhs, typename Dest>
341 void multiply(const SparseMatrixBase<InputType>& A, const MatrixBase<Rhs>& x, MatrixBase<Dest>& y,
342 Scalar alpha = Scalar(1), Scalar beta = Scalar(0), GpuOp op = GpuOp::NoTrans) {
343 const InputType& input = A.derived();
344 internal::check_storage_index_bounds<StorageIndex>(input.rows(), input.cols(), input.nonZeros());
345 SpMat storage;
346 const SpMat& mat = internal::bind_sparse<SpMat>(input, storage);
347 multiply_host_impl(mat, x.derived(), y.derived(), alpha, beta, internal::to_cusparse_op<Scalar>(op));
348 }
349
353 template <typename InputType>
354 void multiply(const SparseMatrixBase<InputType>& A, const DeviceMatrix<Scalar>& d_x, DeviceMatrix<Scalar>& d_y) {
355 multiply(A, d_x, d_y, Scalar(1), Scalar(0), GpuOp::NoTrans);
356 }
357
359 template <typename InputType>
360 void multiply(const SparseMatrixBase<InputType>& A, const DeviceMatrix<Scalar>& d_x, DeviceMatrix<Scalar>& d_y,
361 Scalar alpha, Scalar beta, GpuOp op = GpuOp::NoTrans) {
362 const InputType& input = A.derived();
363 internal::check_storage_index_bounds<StorageIndex>(input.rows(), input.cols(), input.nonZeros());
364 SpMat storage;
365 const SpMat& mat = internal::bind_sparse<SpMat>(input, storage);
366 upload_sparse(mat);
367 spmv_device_exec(d_x, d_y, alpha, beta, op);
368 }
369
371 template <typename InputType, typename Rhs>
372 DenseVector multiplyT(const SparseMatrixBase<InputType>& A, const MatrixBase<Rhs>& x) {
373 return multiply_host_return(A, x, GpuOp::Trans);
374 }
375
377 template <typename InputType, typename Rhs>
378 DenseVector multiplyAdjoint(const SparseMatrixBase<InputType>& A, const MatrixBase<Rhs>& x) {
379 return multiply_host_return(A, x, GpuOp::ConjTrans);
380 }
381
383 template <typename InputType, typename Rhs>
384 DenseMatrix multiplyMat(const SparseMatrixBase<InputType>& A, const MatrixBase<Rhs>& X, GpuOp op = GpuOp::NoTrans) {
385 const InputType& input = A.derived();
386 internal::check_storage_index_bounds<StorageIndex>(input.rows(), input.cols(), input.nonZeros());
387 SpMat storage;
388 const SpMat& mat = internal::bind_sparse<SpMat>(input, storage);
389 const DenseMatrix rhs(X.derived());
390
391 const cusparseOperation_t cu_op = internal::to_cusparse_op<Scalar>(op);
392 const Index m = (op == GpuOp::NoTrans) ? mat.rows() : mat.cols();
393 const Index k = (op == GpuOp::NoTrans) ? mat.cols() : mat.rows();
394 eigen_assert(k == rhs.rows());
395
396 const Index n = rhs.cols();
397 if (m == 0 || n == 0 || mat.nonZeros() == 0) return DenseMatrix::Zero(m, n);
398
399 DenseMatrix Y = DenseMatrix::Zero(m, n);
400 spmm_impl(mat, rhs, Y, Scalar(1), Scalar(0), cu_op);
401 return Y;
402 }
403
404 // BlockSparseMatrix overloads: same contracts as the SparseMatrix ones,
405 // dispatched on internal::use_cusparse_bsr<BlockRows, BlockCols> to a BSR
406 // upload of op(A) or to the scalar-level CSC path.
407
411 template <int Options, int BlockRows, int BlockCols>
412 DeviceSparseView<Scalar> deviceView(const BlockSpMat<Options, BlockRows, BlockCols>& A) {
413 return device_view_block(A, internal::use_cusparse_bsr<BlockRows, BlockCols>());
414 }
415
417 template <int Options, int BlockRows, int BlockCols, typename Rhs>
418 DenseVector multiply(const BlockSpMat<Options, BlockRows, BlockCols>& A, const MatrixBase<Rhs>& x) {
419 return multiply_host_return_block(A, x, GpuOp::NoTrans);
420 }
421
423 template <int Options, int BlockRows, int BlockCols, typename Rhs, typename Dest>
424 void multiply(const BlockSpMat<Options, BlockRows, BlockCols>& A, const MatrixBase<Rhs>& x, MatrixBase<Dest>& y,
425 Scalar alpha = Scalar(1), Scalar beta = Scalar(0), GpuOp op = GpuOp::NoTrans) {
426 multiply_host_block(A, x.derived(), y.derived(), alpha, beta, op,
427 internal::use_cusparse_bsr<BlockRows, BlockCols>());
428 }
429
432 template <int Options, int BlockRows, int BlockCols>
433 void multiply(const BlockSpMat<Options, BlockRows, BlockCols>& A, const DeviceMatrix<Scalar>& d_x,
434 DeviceMatrix<Scalar>& d_y) {
435 multiply(A, d_x, d_y, Scalar(1), Scalar(0), GpuOp::NoTrans);
436 }
437
439 template <int Options, int BlockRows, int BlockCols>
440 void multiply(const BlockSpMat<Options, BlockRows, BlockCols>& A, const DeviceMatrix<Scalar>& d_x,
441 DeviceMatrix<Scalar>& d_y, Scalar alpha, Scalar beta, GpuOp op = GpuOp::NoTrans) {
442 multiply_device_block(A, d_x, d_y, alpha, beta, op, internal::use_cusparse_bsr<BlockRows, BlockCols>());
443 }
444
446 template <int Options, int BlockRows, int BlockCols, typename Rhs>
447 DenseVector multiplyT(const BlockSpMat<Options, BlockRows, BlockCols>& A, const MatrixBase<Rhs>& x) {
448 return multiply_host_return_block(A, x, GpuOp::Trans);
449 }
450
452 template <int Options, int BlockRows, int BlockCols, typename Rhs>
453 DenseVector multiplyAdjoint(const BlockSpMat<Options, BlockRows, BlockCols>& A, const MatrixBase<Rhs>& x) {
454 return multiply_host_return_block(A, x, GpuOp::ConjTrans);
455 }
456
458 template <int Options, int BlockRows, int BlockCols, typename Rhs>
459 DenseMatrix multiplyMat(const BlockSpMat<Options, BlockRows, BlockCols>& A, const MatrixBase<Rhs>& X,
460 GpuOp op = GpuOp::NoTrans) {
461 return multiply_mat_block(A, X, op, internal::use_cusparse_bsr<BlockRows, BlockCols>());
462 }
463
464 cudaStream_t stream() const { return stream_; }
465
466 private:
467 cudaStream_t stream_ = nullptr;
468 cusparseHandle_t handle_ = nullptr;
469 bool owns_stream_ = false;
470 bool owns_handle_ = false;
471
472 // Cached device buffers for sparse matrix (grow-only).
473 internal::DeviceBuffer d_outerPtr_;
474 internal::DeviceBuffer d_innerIdx_;
475 internal::DeviceBuffer d_values_;
476
477 // Cached device buffers for host-API dense vectors (grow-only).
478 internal::DeviceBuffer d_x_;
479 internal::DeviceBuffer d_y_;
480
481 mutable internal::DeviceBuffer d_workspace_;
482
483 // Cached cuSPARSE sparse matrix descriptor, keyed on the scalar-level shape
484 // and nonzero count plus the block size (0 for CSC/CSR, B for BSR).
485 cusparseSpMatDescr_t spmat_desc_ = nullptr;
486 Index cached_rows_ = -1;
487 Index cached_cols_ = -1;
488 Index cached_nnz_ = -1;
489 Index cached_block_size_ = 0;
490
491 // Bumped on every sparse upload; DeviceSparseViews record it at creation so
492 // a stale view (its data replaced by a later upload) asserts at evaluation.
493 uint64_t generation_ = 0;
494
495 // Cached dense-vector/matrix descriptors, re-pointed per call and recreated
496 // only when the shape changes.
497 mutable cusparseDnVecDescr_t x_vec_desc_ = nullptr;
498 mutable cusparseDnVecDescr_t y_vec_desc_ = nullptr;
499 mutable int64_t x_vec_size_ = -1;
500 mutable int64_t y_vec_size_ = -1;
501 mutable cusparseDnMatDescr_t x_mat_desc_ = nullptr;
502 mutable cusparseDnMatDescr_t y_mat_desc_ = nullptr;
503 mutable int64_t x_mat_rows_ = -1, x_mat_cols_ = -1;
504 mutable int64_t y_mat_rows_ = -1, y_mat_cols_ = -1;
505
506 // Cached workspace-size query results, indexed by cusparseOperation_t.
507 // Invalidated when the sparse descriptor or a dense-descriptor shape changes.
508 static constexpr size_t kWsUnknown = static_cast<size_t>(-1);
509 mutable size_t spmv_ws_size_[3] = {kWsUnknown, kWsUnknown, kWsUnknown};
510 mutable size_t spmm_ws_size_[3] = {kWsUnknown, kWsUnknown, kWsUnknown};
511
512 static constexpr cusparseIndexType_t kIndexType =
513 (sizeof(StorageIndex) == 4) ? CUSPARSE_INDEX_32I : CUSPARSE_INDEX_64I;
514 static constexpr cudaDataType_t kValueType = internal::cuda_data_type<Scalar>::value;
515
516 // Empty-operand result on the host: y <- beta * y, no device work.
517 template <typename Dest>
518 static void scale_host(Dest& y, Scalar beta) {
519 if (beta == Scalar(0))
520 y.setZero();
521 else
522 y *= beta;
523 }
524
525 // Shared host-input SpMV entry: y = op(A) * x into a fresh vector.
526 template <typename InputType, typename Rhs>
527 DenseVector multiply_host_return(const SparseMatrixBase<InputType>& A, const MatrixBase<Rhs>& x, GpuOp op) {
528 const InputType& input = A.derived();
529 internal::check_storage_index_bounds<StorageIndex>(input.rows(), input.cols(), input.nonZeros());
530 SpMat storage;
531 const SpMat& mat = internal::bind_sparse<SpMat>(input, storage);
532 DenseVector y((op == GpuOp::NoTrans) ? mat.rows() : mat.cols());
533 y.setZero();
534 multiply_host_impl(mat, x.derived(), y, Scalar(1), Scalar(0), internal::to_cusparse_op<Scalar>(op));
535 return y;
536 }
537
538 template <typename RhsDerived, typename DestDerived>
539 void multiply_host_impl(const SpMat& A, const RhsDerived& x, DestDerived& y, Scalar alpha, Scalar beta,
540 cusparseOperation_t op) {
541 eigen_assert(A.isCompressed());
542
543 const Index m = A.rows();
544 const Index n = A.cols();
545 const Index nnz = A.nonZeros();
546 const Index x_size = (op == CUSPARSE_OPERATION_NON_TRANSPOSE) ? n : m;
547 const Index y_size = (op == CUSPARSE_OPERATION_NON_TRANSPOSE) ? m : n;
548
549 eigen_assert(x.size() == x_size);
550 eigen_assert(y.size() == y_size);
551
552 if (m == 0 || n == 0 || nnz == 0) {
553 scale_host(y, beta);
554 return;
555 }
556
557 upload_sparse(A);
558 spmv_host_cached(x, y, x_size, y_size, alpha, beta, op);
559 }
560
561 // BlockSparseMatrix primitives, dispatched on use_cusparse_bsr: the true_type
562 // overloads bind op(A) as BSR (internal::BsrBinding) and run cuSPARSE with
563 // NON_TRANSPOSE, the false_type ones forward the scalar-level toSparse() copy
564 // to the CSC overloads.
565 template <typename Bsm, typename Rhs>
566 DenseVector multiply_host_return_block(const Bsm& A, const MatrixBase<Rhs>& x, GpuOp op) {
567 DenseVector y((op == GpuOp::NoTrans) ? A.rows() : A.cols());
568 y.setZero();
569 multiply_host_block(A, x.derived(), y, Scalar(1), Scalar(0), op,
570 internal::use_cusparse_bsr<int(Bsm::BlockRows), int(Bsm::BlockCols)>());
571 return y;
572 }
573
574 template <typename Bsm>
575 DeviceSparseView<Scalar> device_view_block(const Bsm& A, std::false_type) {
576 return deviceView(SpMat(A.toSparse()));
577 }
578
579 template <typename Bsm, typename RhsDerived, typename DestDerived>
580 void multiply_host_block(const Bsm& A, const RhsDerived& x, DestDerived& y, Scalar alpha, Scalar beta, GpuOp op,
581 std::false_type) {
582 multiply(A.toSparse(), x, y, alpha, beta, op);
583 }
584
585 template <typename Bsm>
586 void multiply_device_block(const Bsm& A, const DeviceMatrix<Scalar>& d_x, DeviceMatrix<Scalar>& d_y, Scalar alpha,
587 Scalar beta, GpuOp op, std::false_type) {
588 multiply(A.toSparse(), d_x, d_y, alpha, beta, op);
589 }
590
591 template <typename Bsm, typename Rhs>
592 DenseMatrix multiply_mat_block(const Bsm& A, const MatrixBase<Rhs>& X, GpuOp op, std::false_type) {
593 return multiplyMat(A.toSparse(), X, op);
594 }
595
596#if EIGEN_HAS_CUSPARSE_BSR
597 template <typename Bsm>
598 DeviceSparseView<Scalar> device_view_block(const Bsm& A, std::true_type) {
599 internal::check_storage_index_bounds<StorageIndex>(A.rows(), A.cols(), A.nonZeros());
600 const internal::BsrBinding<Bsm> bound(A, GpuOp::NoTrans);
601 upload_bsr(bound.arrays());
602 return DeviceSparseView<Scalar>(*this, A.rows(), A.cols(), generation_);
603 }
604
605 template <typename Bsm, typename RhsDerived, typename DestDerived>
606 void multiply_host_block(const Bsm& A, const RhsDerived& x, DestDerived& y, Scalar alpha, Scalar beta, GpuOp op,
607 std::true_type) {
608 internal::check_storage_index_bounds<StorageIndex>(A.rows(), A.cols(), A.nonZeros());
609 const internal::BsrBinding<Bsm> bound(A, op);
610 const internal::BsrArrays<Scalar, StorageIndex>& opA = bound.arrays();
611
612 eigen_assert(x.size() == opA.cols());
613 eigen_assert(y.size() == opA.rows());
614
615 if (opA.rows() == 0 || opA.cols() == 0 || opA.bnnz == 0) {
616 scale_host(y, beta);
617 return;
618 }
619
620 upload_bsr(opA);
621 spmv_host_cached(x, y, opA.cols(), opA.rows(), alpha, beta, CUSPARSE_OPERATION_NON_TRANSPOSE);
622 }
623
624 template <typename Bsm>
625 void multiply_device_block(const Bsm& A, const DeviceMatrix<Scalar>& d_x, DeviceMatrix<Scalar>& d_y, Scalar alpha,
626 Scalar beta, GpuOp op, std::true_type) {
627 internal::check_storage_index_bounds<StorageIndex>(A.rows(), A.cols(), A.nonZeros());
628 const internal::BsrBinding<Bsm> bound(A, op);
629 upload_bsr(bound.arrays());
630 spmv_device_exec(d_x, d_y, alpha, beta, GpuOp::NoTrans);
631 }
632
633 template <typename Bsm, typename Rhs>
634 DenseMatrix multiply_mat_block(const Bsm& A, const MatrixBase<Rhs>& X, GpuOp op, std::true_type) {
635 internal::check_storage_index_bounds<StorageIndex>(A.rows(), A.cols(), A.nonZeros());
636 const internal::BsrBinding<Bsm> bound(A, op);
637 const internal::BsrArrays<Scalar, StorageIndex>& opA = bound.arrays();
638 const DenseMatrix rhs(X.derived());
639 eigen_assert(opA.cols() == rhs.rows());
640
641 const Index n = rhs.cols();
642 if (opA.rows() == 0 || opA.cols() == 0 || n == 0 || opA.bnnz == 0) return DenseMatrix::Zero(opA.rows(), n);
643
644 DenseMatrix Y = DenseMatrix::Zero(opA.rows(), n);
645 upload_bsr(opA);
646 spmm_host_cached(rhs, Y, opA.rows(), opA.cols(), Scalar(1), Scalar(0), CUSPARSE_OPERATION_NON_TRANSPOSE);
647 return Y;
648 }
649#endif // EIGEN_HAS_CUSPARSE_BSR
650
651 // y = alpha * op(D) * x + beta * y against the cached upload D, staging the
652 // host vectors through d_x_ / d_y_. x_size and y_size are the column and row
653 // counts of op(D).
654 template <typename RhsDerived, typename DestDerived>
655 void spmv_host_cached(const RhsDerived& x, DestDerived& y, Index x_size, Index y_size, Scalar alpha, Scalar beta,
656 cusparseOperation_t op) {
657 ensure_buffer(d_x_, static_cast<size_t>(x_size) * sizeof(Scalar));
658 // Ref binds in place when x is already a contiguous vector; only genuine
659 // expressions are evaluated into the Ref's internal temporary.
660 const Ref<const DenseVector> x_ref(x);
661 EIGEN_CUDA_RUNTIME_CHECK(
662 cudaMemcpyAsync(d_x_.get(), x_ref.data(), x_size * sizeof(Scalar), cudaMemcpyHostToDevice, stream_));
663
664 ensure_buffer(d_y_, static_cast<size_t>(y_size) * sizeof(Scalar));
665 if (beta != Scalar(0)) {
666 const Ref<const DenseVector> y_ref(y);
667 EIGEN_CUDA_RUNTIME_CHECK(
668 cudaMemcpyAsync(d_y_.get(), y_ref.data(), y_size * sizeof(Scalar), cudaMemcpyHostToDevice, stream_));
669 }
670
671 exec_spmv(x_size, y_size, d_x_.get(), d_y_.get(), alpha, beta, op);
672
673 EIGEN_CUDA_RUNTIME_CHECK(
674 cudaMemcpyAsync(y.data(), d_y_.get(), y_size * sizeof(Scalar), cudaMemcpyDeviceToHost, stream_));
675 EIGEN_CUDA_RUNTIME_CHECK(cudaStreamSynchronize(stream_));
676 }
677
678 public:
684 void spmv_device_exec(const DeviceMatrix<Scalar>& d_x, DeviceMatrix<Scalar>& d_y, Scalar alpha = Scalar(1),
685 Scalar beta = Scalar(0), GpuOp op = GpuOp::NoTrans) const {
686 eigen_assert(spmat_desc_ && "sparse matrix not uploaded — call deviceView() or multiply() first");
687 check_op_against_upload(op);
688 // cuSPARSE SpMV: y must not alias x (undefined behavior).
689 eigen_assert(d_x.data() != d_y.data() && "SpMV: output aliases input vector");
690
691 const cusparseOperation_t cu_op = internal::to_cusparse_op<Scalar>(op);
692 const Index m = cached_rows_;
693 const Index n = cached_cols_;
694 const Index x_size = (cu_op == CUSPARSE_OPERATION_NON_TRANSPOSE) ? n : m;
695 const Index y_size = (cu_op == CUSPARSE_OPERATION_NON_TRANSPOSE) ? m : n;
696
697 eigen_assert(d_x.rows() * d_x.cols() == x_size);
698
699 if (m == 0 || n == 0 || cached_nnz_ == 0) {
700 // Empty A reduces SpMV to y <- beta*y; SparseContext owns no cuBLAS
701 // handle for the scale, so the beta != 0 case must be handled by the caller.
702 eigen_assert(beta == Scalar(0) && "SpMV with empty A and beta != 0 is unsupported; scale d_y externally");
703 if (d_y.rows() * d_y.cols() != y_size) d_y.resize(y_size, 1);
704 d_y.setZero(stream_);
705 return;
706 }
707
708 // Ensure d_y is allocated.
709 if (d_y.rows() * d_y.cols() != y_size) {
710 d_y.resize(y_size, 1);
711 }
712
713 // Wait for input data to be ready on this stream.
714 d_x.waitReady(stream_);
715 d_y.waitReady(stream_);
716
717 exec_spmv(x_size, y_size, const_cast<void*>(static_cast<const void*>(d_x.data())), static_cast<void*>(d_y.data()),
718 alpha, beta, cu_op);
719
720 d_y.recordReady(stream_);
721 }
722
727 void spmm_device_exec(const DeviceMatrix<Scalar>& d_X, DeviceMatrix<Scalar>& d_Y, Scalar alpha = Scalar(1),
728 Scalar beta = Scalar(0), GpuOp op = GpuOp::NoTrans) const {
729 eigen_assert(spmat_desc_ && "sparse matrix not uploaded — call deviceView() or multiply() first");
730 check_op_against_upload(op);
731 eigen_assert(d_X.data() != d_Y.data() && "SpMM: output aliases input matrix");
732
733 const cusparseOperation_t cu_op = internal::to_cusparse_op<Scalar>(op);
734 const bool transposed = (cu_op != CUSPARSE_OPERATION_NON_TRANSPOSE);
735 const Index m_op = transposed ? cached_cols_ : cached_rows_;
736 const Index k_op = transposed ? cached_rows_ : cached_cols_;
737 const Index n = d_X.cols();
738
739 eigen_assert(d_X.rows() == k_op);
740
741 if (m_op == 0 || n == 0 || cached_nnz_ == 0) {
742 eigen_assert(beta == Scalar(0) && "SpMM with empty A and beta != 0 is unsupported; scale d_Y externally");
743 if (d_Y.rows() != m_op || d_Y.cols() != n) d_Y.resize(m_op, n);
744 d_Y.setZero(stream_);
745 return;
746 }
747
748 if (d_Y.rows() != m_op || d_Y.cols() != n) {
749 d_Y.resize(m_op, n);
750 }
751
752 d_X.waitReady(stream_);
753 d_Y.waitReady(stream_);
754
755 exec_spmm(m_op, k_op, n, const_cast<void*>(static_cast<const void*>(d_X.data())), static_cast<void*>(d_Y.data()),
756 alpha, beta, cu_op);
757
758 d_Y.recordReady(stream_);
759 }
760
761 private:
762 // cuSPARSE 11.x's cusparseSpMM rejects CSC for matA (CSC support landed in
763 // CUDA 12.0). On 11.x we register the same buffers as CSR-of-A^T (dims
764 // swapped) and invert the user-facing op before each cuSPARSE call. On 12+
765 // we keep the natural CSC path so users pay no extra cost.
766#if !defined(CUSPARSE_VERSION) || CUSPARSE_VERSION < 12000
767 static constexpr bool kUseCsrOfTranspose = true;
768 static constexpr cusparseSpMMAlg_t kSpMMAlg = CUSPARSE_SPMM_CSR_ALG2;
769#else
770 static constexpr bool kUseCsrOfTranspose = false;
771 static constexpr cusparseSpMMAlg_t kSpMMAlg = CUSPARSE_SPMM_ALG_DEFAULT;
772#endif
773
774 // Map a user-facing op on A to the cuSPARSE op on the cached descriptor.
775 // Identity on cuSPARSE 12+ (descriptor is CSC of A, or BSR of op(A));
776 // inverted on 11.x (descriptor is CSR of A^T).
777 static cusparseOperation_t descriptor_op(cusparseOperation_t user_op) {
778 EIGEN_IF_CONSTEXPR (!kUseCsrOfTranspose) return user_op;
779 switch (user_op) {
780 case CUSPARSE_OPERATION_NON_TRANSPOSE:
781 return CUSPARSE_OPERATION_TRANSPOSE;
782 case CUSPARSE_OPERATION_TRANSPOSE:
783 return CUSPARSE_OPERATION_NON_TRANSPOSE;
784 default:
785 // CONJUGATE_TRANSPOSE on the CSR-of-A^T descriptor would compute
786 // conj(A) * x, not A^H * x — not supported via this representation.
787 eigen_assert(false && "CUSPARSE_OPERATION_CONJUGATE_TRANSPOSE not supported on cuSPARSE < 12.0");
788 return user_op;
789 }
790 }
791
792 static int op_index(cusparseOperation_t op) {
793 switch (op) {
794 case CUSPARSE_OPERATION_TRANSPOSE:
795 return 1;
796 case CUSPARSE_OPERATION_CONJUGATE_TRANSPOSE:
797 return 2;
798 default:
799 return 0;
800 }
801 }
802
803 void invalidate_ws_caches() const {
804 for (int i = 0; i < 3; ++i) {
805 spmv_ws_size_[i] = kWsUnknown;
806 spmm_ws_size_[i] = kWsUnknown;
807 }
808 }
809
810 // Recreate the descriptor when the size changes; otherwise just re-point it
811 // at the new device buffer (cusparseDnVecSetValues is a host-side pointer
812 // update, no GPU work).
813 void update_dnvec(cusparseDnVecDescr_t& desc, int64_t& cur_size, int64_t size, void* ptr) const {
814 if (!desc || cur_size != size) {
815 if (desc) EIGEN_CUSPARSE_CHECK(cusparseDestroyDnVec(desc));
816 EIGEN_CUSPARSE_CHECK(cusparseCreateDnVec(&desc, size, ptr, kValueType));
817 cur_size = size;
818 invalidate_ws_caches();
819 } else {
820 EIGEN_CUSPARSE_CHECK(cusparseDnVecSetValues(desc, ptr));
821 }
822 }
823
824 void update_dnmat(cusparseDnMatDescr_t& desc, int64_t& cur_rows, int64_t& cur_cols, int64_t rows, int64_t cols,
825 void* ptr) const {
826 if (!desc || cur_rows != rows || cur_cols != cols) {
827 if (desc) EIGEN_CUSPARSE_CHECK(cusparseDestroyDnMat(desc));
828 // Column-major with ld = rows.
829 EIGEN_CUSPARSE_CHECK(cusparseCreateDnMat(&desc, rows, cols, rows, ptr, kValueType, CUSPARSE_ORDER_COL));
830 cur_rows = rows;
831 cur_cols = cols;
832 invalidate_ws_caches();
833 } else {
834 EIGEN_CUSPARSE_CHECK(cusparseDnMatSetValues(desc, ptr));
835 }
836 }
837
838 void destroy_dense_descriptors() {
839 if (x_vec_desc_) (void)cusparseDestroyDnVec(x_vec_desc_);
840 if (y_vec_desc_) (void)cusparseDestroyDnVec(y_vec_desc_);
841 if (x_mat_desc_) (void)cusparseDestroyDnMat(x_mat_desc_);
842 if (y_mat_desc_) (void)cusparseDestroyDnMat(y_mat_desc_);
843 x_vec_desc_ = y_vec_desc_ = nullptr;
844 x_mat_desc_ = y_mat_desc_ = nullptr;
845 }
846
847 void exec_spmv(Index x_size, Index y_size, void* d_x_ptr, void* d_y_ptr, Scalar alpha, Scalar beta,
848 cusparseOperation_t op) const {
849 const cusparseOperation_t cu_op = descriptor_op(op);
850 update_dnvec(x_vec_desc_, x_vec_size_, x_size, d_x_ptr);
851 update_dnvec(y_vec_desc_, y_vec_size_, y_size, d_y_ptr);
852
853 size_t& ws_size = spmv_ws_size_[op_index(cu_op)];
854 if (ws_size == kWsUnknown) {
855 EIGEN_CUSPARSE_CHECK(cusparseSpMV_bufferSize(handle_, cu_op, &alpha, spmat_desc_, x_vec_desc_, &beta, y_vec_desc_,
856 kValueType, CUSPARSE_SPMV_ALG_DEFAULT, &ws_size));
857 }
858 ensure_buffer(d_workspace_, ws_size);
859
860 EIGEN_CUSPARSE_CHECK(cusparseSpMV(handle_, cu_op, &alpha, spmat_desc_, x_vec_desc_, &beta, y_vec_desc_, kValueType,
861 CUSPARSE_SPMV_ALG_DEFAULT, d_workspace_.get()));
862 }
863
864 void exec_spmm(Index m_op, Index k_op, Index n, void* d_x_ptr, void* d_y_ptr, Scalar alpha, Scalar beta,
865 cusparseOperation_t op) const {
866 const cusparseOperation_t cu_op = descriptor_op(op);
867 // X is k_op x n, Y is m_op x n (column-major, post-op shapes).
868 update_dnmat(x_mat_desc_, x_mat_rows_, x_mat_cols_, k_op, n, d_x_ptr);
869 update_dnmat(y_mat_desc_, y_mat_rows_, y_mat_cols_, m_op, n, d_y_ptr);
870
871 size_t& ws_size = spmm_ws_size_[op_index(cu_op)];
872 if (ws_size == kWsUnknown) {
873 EIGEN_CUSPARSE_CHECK(cusparseSpMM_bufferSize(handle_, cu_op, CUSPARSE_OPERATION_NON_TRANSPOSE, &alpha,
874 spmat_desc_, x_mat_desc_, &beta, y_mat_desc_, kValueType, kSpMMAlg,
875 &ws_size));
876 }
877 ensure_buffer(d_workspace_, ws_size);
878
879 EIGEN_CUSPARSE_CHECK(cusparseSpMM(handle_, cu_op, CUSPARSE_OPERATION_NON_TRANSPOSE, &alpha, spmat_desc_,
880 x_mat_desc_, &beta, y_mat_desc_, kValueType, kSpMMAlg, d_workspace_.get()));
881 }
882
883 void spmm_impl(const SpMat& A, const DenseMatrix& X, DenseMatrix& Y, Scalar alpha, Scalar beta,
884 cusparseOperation_t op) {
885 eigen_assert(A.isCompressed());
886
887 // For op != NON_TRANSPOSE, Y = op(A) * X. The dense X / Y descriptors must
888 // describe the *post-op* shapes: X has k_op rows (= input dim of op(A)),
889 // Y has m_op rows (= output dim of op(A)).
890 const bool transposed = (op != CUSPARSE_OPERATION_NON_TRANSPOSE);
891 const Index m_op = transposed ? A.cols() : A.rows();
892 const Index k_op = transposed ? A.rows() : A.cols();
893 const Index n = X.cols();
894 const Index nnz = A.nonZeros();
895
896 if (m_op == 0 || n == 0 || k_op == 0 || nnz == 0) {
897 scale_host(Y, beta);
898 return;
899 }
900
901 upload_sparse(A);
902 spmm_host_cached(X, Y, m_op, k_op, alpha, beta, op);
903 }
904
905 // Y = alpha * op(D) * X + beta * Y against the cached upload D, staging the
906 // host matrices through d_x_ / d_y_. X is k_op x n, Y is m_op x n.
907 void spmm_host_cached(const DenseMatrix& X, DenseMatrix& Y, Index m_op, Index k_op, Scalar alpha, Scalar beta,
908 cusparseOperation_t op) {
909 const Index n = X.cols();
910 const size_t x_bytes = static_cast<size_t>(k_op) * static_cast<size_t>(n) * sizeof(Scalar);
911 const size_t y_bytes = static_cast<size_t>(m_op) * static_cast<size_t>(n) * sizeof(Scalar);
912 ensure_buffer(d_x_, x_bytes);
913 ensure_buffer(d_y_, y_bytes);
914 EIGEN_CUDA_RUNTIME_CHECK(cudaMemcpyAsync(d_x_.get(), X.data(), x_bytes, cudaMemcpyHostToDevice, stream_));
915 if (beta != Scalar(0)) {
916 EIGEN_CUDA_RUNTIME_CHECK(cudaMemcpyAsync(d_y_.get(), Y.data(), y_bytes, cudaMemcpyHostToDevice, stream_));
917 }
918
919 exec_spmm(m_op, k_op, n, d_x_.get(), d_y_.get(), alpha, beta, op);
920
921 EIGEN_CUDA_RUNTIME_CHECK(cudaMemcpyAsync(Y.data(), d_y_.get(), y_bytes, cudaMemcpyDeviceToHost, stream_));
922 EIGEN_CUDA_RUNTIME_CHECK(cudaStreamSynchronize(stream_));
923 }
924
925 void upload_sparse(const SpMat& A) {
926 // cuSPARSE 12.0+ accepts CSC directly. On cuSPARSE 11.x, cusparseSpMM
927 // rejects CSC and CONJUGATE_TRANSPOSE on CSC+complex SpMV silently
928 // demotes to TRANSPOSE. We register the same CSC buffers as CSR-of-A^T
929 // (dims swapped) on 11.x and invert the op at exec time via
930 // descriptor_op() — no transpose-copy required.
931 const Index m = A.rows();
932 const Index n = A.cols();
933 const Index nnz = A.nonZeros();
934 upload_arrays(/*outer_count=*/n + 1, /*inner_count=*/nnz, /*value_count=*/nnz, A.outerIndexPtr(), A.innerIndexPtr(),
935 A.valuePtr());
936 if (descriptor_key_matches(m, n, nnz, /*block_size=*/0)) return;
937
938 destroy_spmat_descriptor(/*checked=*/true);
939 EIGEN_IF_CONSTEXPR (kUseCsrOfTranspose) {
940 // cuSPARSE 11.x: cusparseSpMM rejects CSC for matA. CSC of A and CSR of
941 // A^T share the same buffers, so register the data as CSR-of-A^T (dims
942 // swapped) and invert the op in exec_spmv / spmm_impl via descriptor_op.
943 EIGEN_CUSPARSE_CHECK(cusparseCreateCsr(&spmat_desc_, n, m, nnz, d_outerPtr_.get(), d_innerIdx_.get(),
944 d_values_.get(), kIndexType, kIndexType, CUSPARSE_INDEX_BASE_ZERO,
945 kValueType));
946 } else {
947 EIGEN_CUSPARSE_CHECK(cusparseCreateCsc(&spmat_desc_, m, n, nnz, d_outerPtr_.get(), d_innerIdx_.get(),
948 d_values_.get(), kIndexType, kIndexType, CUSPARSE_INDEX_BASE_ZERO,
949 kValueType));
950 }
951 set_descriptor_key(m, n, nnz, /*block_size=*/0);
952 }
953
954#if EIGEN_HAS_CUSPARSE_BSR
955 void upload_bsr(const internal::BsrArrays<Scalar, StorageIndex>& opA) {
956 upload_arrays(/*outer_count=*/opA.brows + 1, /*inner_count=*/opA.bnnz, /*value_count=*/opA.nonZeros(), opA.outer,
957 opA.inner, opA.values);
958 if (descriptor_key_matches(opA.rows(), opA.cols(), opA.nonZeros(), opA.blockSize)) return;
959
960 destroy_spmat_descriptor(/*checked=*/true);
961 // Row-major blocks: cusparseSpMM accepts no other block layout for BSR.
962 EIGEN_CUSPARSE_CHECK(cusparseCreateBsr(&spmat_desc_, opA.brows, opA.bcols, opA.bnnz, opA.blockSize, opA.blockSize,
963 d_outerPtr_.get(), d_innerIdx_.get(), d_values_.get(), kIndexType,
964 kIndexType, CUSPARSE_INDEX_BASE_ZERO, kValueType, CUSPARSE_ORDER_ROW));
965 set_descriptor_key(opA.rows(), opA.cols(), opA.nonZeros(), opA.blockSize);
966 }
967#endif // EIGEN_HAS_CUSPARSE_BSR
968
969 void upload_arrays(Index outer_count, Index inner_count, Index value_count, const StorageIndex* host_outer,
970 const StorageIndex* host_inner, const Scalar* host_values) {
971 const size_t outer_bytes = static_cast<size_t>(outer_count) * sizeof(StorageIndex);
972 const size_t inner_bytes = static_cast<size_t>(inner_count) * sizeof(StorageIndex);
973 const size_t val_bytes = static_cast<size_t>(value_count) * sizeof(Scalar);
974
975 // Values *and* index arrays are re-uploaded unconditionally: host pointer
976 // identity cannot detect a same-shape/same-nnz pattern rewritten in place
977 // or assigned into the same allocations (SparseMatrix reuses them), so a
978 // structure cache keyed on pointers would silently serve stale indices.
979 // Only the cuSPARSE descriptor and the workspace-size queries are cached,
980 // keyed on (rows, cols, nnz, block size). Every upload invalidates
981 // outstanding DeviceSparseViews via the generation counter.
982 ++generation_;
983 ensure_buffer(d_values_, val_bytes);
984 ensure_buffer(d_outerPtr_, outer_bytes);
985 ensure_buffer(d_innerIdx_, inner_bytes);
986 EIGEN_CUDA_RUNTIME_CHECK(cudaMemcpyAsync(d_values_.get(), host_values, val_bytes, cudaMemcpyHostToDevice, stream_));
987 EIGEN_CUDA_RUNTIME_CHECK(
988 cudaMemcpyAsync(d_outerPtr_.get(), host_outer, outer_bytes, cudaMemcpyHostToDevice, stream_));
989 EIGEN_CUDA_RUNTIME_CHECK(
990 cudaMemcpyAsync(d_innerIdx_.get(), host_inner, inner_bytes, cudaMemcpyHostToDevice, stream_));
991 }
992
993 // Same shape, nnz and format: the grow-only device buffers cannot have been
994 // reallocated, so the existing descriptor still points at the freshly
995 // written data.
996 bool descriptor_key_matches(Index m, Index n, Index nnz, Index block_size) const {
997 return m == cached_rows_ && n == cached_cols_ && nnz == cached_nnz_ && block_size == cached_block_size_;
998 }
999
1000 void set_descriptor_key(Index m, Index n, Index nnz, Index block_size) {
1001 cached_rows_ = m;
1002 cached_cols_ = n;
1003 cached_nnz_ = nnz;
1004 cached_block_size_ = block_size;
1005 }
1006
1007 // A BSR descriptor holds op(A) as bound at upload time, and cuSPARSE runs no
1008 // transposed BSR product.
1009 void check_op_against_upload(GpuOp op) const {
1010 eigen_assert((cached_block_size_ == 0 || op == GpuOp::NoTrans) &&
1011 "cuSPARSE runs BSR products only with op == NoTrans; pass op to multiply(A, d_x, d_y, ...) instead");
1012 EIGEN_UNUSED_VARIABLE(op);
1013 }
1014
1015 // Destroy the sparse-matrix descriptor and reset the cache identity.
1016 // `checked` selects checking (mid-lifetime rebuilds) vs swallowing failure
1017 // (noexcept destructor).
1018 void destroy_spmat_descriptor(bool checked) {
1019 if (spmat_desc_) {
1020 if (checked)
1021 EIGEN_CUSPARSE_CHECK(cusparseDestroySpMat(spmat_desc_));
1022 else
1023 (void)cusparseDestroySpMat(spmat_desc_);
1024 spmat_desc_ = nullptr;
1025 }
1026 set_descriptor_key(-1, -1, -1, 0);
1027 invalidate_ws_caches();
1028 }
1029
1030 void ensure_buffer(internal::DeviceBuffer& buf, size_t needed) const {
1031 if (needed > buf.size()) {
1032 if (buf) EIGEN_CUDA_RUNTIME_CHECK(cudaStreamSynchronize(stream_));
1033 buf = internal::DeviceBuffer(needed);
1034 }
1035 }
1036};
1037
1038// Defined here because it needs the full SparseContext definition.
1039
1040template <typename Scalar_>
1041DeviceMatrix<Scalar_>& DeviceMatrix<Scalar_>::operator=(const SpMVExpr<Scalar_>& expr) {
1042 // Uses the sparse matrix already uploaded by deviceView() — no re-upload on
1043 // repeated products with the same view. A stale view (the context has since
1044 // uploaded again, replacing the cached data) is caught here.
1045 eigen_assert(expr.view().generation() == expr.view().context().uploadGeneration() &&
1046 "DeviceSparseView is stale: its SparseContext has since uploaded another sparse matrix");
1047 if (expr.x().cols() <= 1) {
1048 expr.view().context().spmv_device_exec(expr.x(), *this, Scalar_(1), Scalar_(0), GpuOp::NoTrans);
1049 } else {
1050 expr.view().context().spmm_device_exec(expr.x(), *this, Scalar_(1), Scalar_(0), GpuOp::NoTrans);
1051 }
1052 return *this;
1053}
1054
1055template <typename Scalar_>
1056DeviceMatrix<Scalar_>& DeviceMatrix<Scalar_>::operator=(const SpMVAffineExpr<Scalar_>& expr) {
1057 const DeviceSparseView<Scalar_>& view = expr.view();
1058 const DeviceMatrix& addend = expr.addend();
1059 eigen_assert(view.generation() == view.context().uploadGeneration() &&
1060 "DeviceSparseView is stale: its SparseContext has since uploaded another sparse matrix");
1061 eigen_assert(addend.rows() == view.rows() && addend.cols() == expr.x().cols() &&
1062 "SpMVAffineExpr: the addend must have the shape of the product");
1063 eigen_assert(&expr.x() != this && "SpMVAffineExpr: the destination aliases the dense operand");
1064 // The product accumulates into a copy of the addend; when the addend is the destination
1065 // itself (d_r = d_r - d_A * d_x) cuSPARSE's beta accumulates in place.
1066 if (&addend != this) copyFrom(Context::threadLocal(), addend);
1067 // With no stored entries the expression is beta * addend, which spmv_device_exec cannot
1068 // form: SparseContext owns no cuBLAS handle to scale d_y by beta.
1069 if (view.nonZeros() == 0 || view.rows() == 0 || view.cols() == 0) {
1070 if (expr.beta() != Scalar_(1)) scale(Context::threadLocal(), expr.beta());
1071 return *this;
1072 }
1073 if (expr.x().cols() <= 1) {
1074 view.context().spmv_device_exec(expr.x(), *this, expr.alpha(), expr.beta(), GpuOp::NoTrans);
1075 } else {
1076 view.context().spmm_device_exec(expr.x(), *this, expr.alpha(), expr.beta(), GpuOp::NoTrans);
1077 }
1078 return *this;
1079}
1080
1081template <typename Scalar_>
1083 *this = expr;
1084}
1085
1086template <typename Scalar_>
1087DeviceMatrix<Scalar_>::DeviceMatrix(const SpMVAffineExpr<Scalar_>& expr) : DeviceMatrix() {
1088 *this = expr;
1089}
1090} // namespace gpu
1091} // namespace Eigen
1092
1093#endif // EIGEN_GPU_SPARSE_CONTEXT_H
static Context & threadLocal()
Definition GpuContext.h:121
RAII wrapper for a dense column-major matrix in GPU device memory.
Definition DeviceMatrix.h:122
void scale(Context &ctx, Scalar alpha)
Definition DeviceDispatch.h:689
void copyFrom(Context &ctx, const DeviceMatrix &other)
Definition DeviceDispatch.h:699
static DeviceMatrix view(Scalar *device_ptr, Index rows, Index cols)
Definition DeviceMatrix.h:576
Definition GpuSparseContext.h:234
SpMVExpr< Scalar > operator*(const DeviceMatrix< Scalar > &x) const
Definition GpuSparseContext.h:249
Index nonZeros() const
Definition GpuSparseContext.h:254
uint64_t generation() const
Definition GpuSparseContext.h:259
Definition GpuSparseContext.h:166
Definition GpuSparseContext.h:146
Namespace containing all symbols from the Eigen library.
Definition GpuSparseContext.h:55