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());
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;
549 eigen_assert(x.size() == x_size);
550 eigen_assert(y.size() == y_size);
552 if (m == 0 || n == 0 || nnz == 0) {
558 spmv_host_cached(x, y, x_size, y_size, alpha, beta, op);
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());
569 multiply_host_block(A, x.derived(), y, Scalar(1), Scalar(0), op,
574 template <
typename Bsm>
575 DeviceSparseView<Scalar> device_view_block(
const Bsm& A, std::false_type) {
576 return deviceView(SpMat(A.toSparse()));
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,
582 multiply(A.toSparse(), x, y, alpha, beta, op);
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);
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);
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_);
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,
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();
612 eigen_assert(x.size() == opA.cols());
613 eigen_assert(y.size() == opA.rows());
615 if (opA.rows() == 0 || opA.cols() == 0 || opA.bnnz == 0) {
621 spmv_host_cached(x, y, opA.cols(), opA.rows(), alpha, beta, CUSPARSE_OPERATION_NON_TRANSPOSE);
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);
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());
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);
644 DenseMatrix Y = DenseMatrix::Zero(opA.rows(), n);
646 spmm_host_cached(rhs, Y, opA.rows(), opA.cols(), Scalar(1), Scalar(0), CUSPARSE_OPERATION_NON_TRANSPOSE);
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));
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_));
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_));
671 exec_spmv(x_size, y_size, d_x_.get(), d_y_.get(), alpha, beta, op);
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_));
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);
689 eigen_assert(d_x.data() != d_y.data() &&
"SpMV: output aliases input vector");
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;
697 eigen_assert(d_x.rows() * d_x.cols() == x_size);
699 if (m == 0 || n == 0 || cached_nnz_ == 0) {
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_);
709 if (d_y.rows() * d_y.cols() != y_size) {
710 d_y.resize(y_size, 1);
714 d_x.waitReady(stream_);
715 d_y.waitReady(stream_);
717 exec_spmv(x_size, y_size,
const_cast<void*
>(
static_cast<const void*
>(d_x.data())),
static_cast<void*
>(d_y.data()),
720 d_y.recordReady(stream_);
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");
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();
739 eigen_assert(d_X.rows() == k_op);
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_);
748 if (d_Y.rows() != m_op || d_Y.cols() != n) {
752 d_X.waitReady(stream_);
753 d_Y.waitReady(stream_);
755 exec_spmm(m_op, k_op, n,
const_cast<void*
>(
static_cast<const void*
>(d_X.data())),
static_cast<void*
>(d_Y.data()),
758 d_Y.recordReady(stream_);
766#if !defined(CUSPARSE_VERSION) || CUSPARSE_VERSION < 12000
767 static constexpr bool kUseCsrOfTranspose =
true;
768 static constexpr cusparseSpMMAlg_t kSpMMAlg = CUSPARSE_SPMM_CSR_ALG2;
770 static constexpr bool kUseCsrOfTranspose =
false;
771 static constexpr cusparseSpMMAlg_t kSpMMAlg = CUSPARSE_SPMM_ALG_DEFAULT;
777 static cusparseOperation_t descriptor_op(cusparseOperation_t user_op) {
778 EIGEN_IF_CONSTEXPR (!kUseCsrOfTranspose)
return user_op;
780 case CUSPARSE_OPERATION_NON_TRANSPOSE:
781 return CUSPARSE_OPERATION_TRANSPOSE;
782 case CUSPARSE_OPERATION_TRANSPOSE:
783 return CUSPARSE_OPERATION_NON_TRANSPOSE;
787 eigen_assert(
false &&
"CUSPARSE_OPERATION_CONJUGATE_TRANSPOSE not supported on cuSPARSE < 12.0");
792 static int op_index(cusparseOperation_t op) {
794 case CUSPARSE_OPERATION_TRANSPOSE:
796 case CUSPARSE_OPERATION_CONJUGATE_TRANSPOSE:
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;
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));
818 invalidate_ws_caches();
820 EIGEN_CUSPARSE_CHECK(cusparseDnVecSetValues(desc, ptr));
824 void update_dnmat(cusparseDnMatDescr_t& desc, int64_t& cur_rows, int64_t& cur_cols, int64_t rows, int64_t cols,
826 if (!desc || cur_rows != rows || cur_cols != cols) {
827 if (desc) EIGEN_CUSPARSE_CHECK(cusparseDestroyDnMat(desc));
829 EIGEN_CUSPARSE_CHECK(cusparseCreateDnMat(&desc, rows, cols, rows, ptr, kValueType, CUSPARSE_ORDER_COL));
832 invalidate_ws_caches();
834 EIGEN_CUSPARSE_CHECK(cusparseDnMatSetValues(desc, ptr));
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;
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);
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));
858 ensure_buffer(d_workspace_, ws_size);
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()));
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);
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);
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,
877 ensure_buffer(d_workspace_, ws_size);
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()));
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());
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();
896 if (m_op == 0 || n == 0 || k_op == 0 || nnz == 0) {
902 spmm_host_cached(X, Y, m_op, k_op, alpha, beta, op);
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_));
919 exec_spmm(m_op, k_op, n, d_x_.get(), d_y_.get(), alpha, beta, op);
921 EIGEN_CUDA_RUNTIME_CHECK(cudaMemcpyAsync(Y.data(), d_y_.get(), y_bytes, cudaMemcpyDeviceToHost, stream_));
922 EIGEN_CUDA_RUNTIME_CHECK(cudaStreamSynchronize(stream_));
925 void upload_sparse(
const SpMat& A) {
931 const Index m = A.rows();
932 const Index n = A.cols();
933 const Index nnz = A.nonZeros();
934 upload_arrays(n + 1, nnz, nnz, A.outerIndexPtr(), A.innerIndexPtr(),
936 if (descriptor_key_matches(m, n, nnz, 0))
return;
938 destroy_spmat_descriptor(
true);
939 EIGEN_IF_CONSTEXPR (kUseCsrOfTranspose) {
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,
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,
951 set_descriptor_key(m, n, nnz, 0);
954#if EIGEN_HAS_CUSPARSE_BSR
955 void upload_bsr(
const internal::BsrArrays<Scalar, StorageIndex>& opA) {
956 upload_arrays(opA.brows + 1, opA.bnnz, opA.nonZeros(), opA.outer,
957 opA.inner, opA.values);
958 if (descriptor_key_matches(opA.rows(), opA.cols(), opA.nonZeros(), opA.blockSize))
return;
960 destroy_spmat_descriptor(
true);
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);
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);
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_));
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_;
1000 void set_descriptor_key(Index m, Index n, Index nnz, Index block_size) {
1004 cached_block_size_ = block_size;
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);
1018 void destroy_spmat_descriptor(
bool checked) {
1021 EIGEN_CUSPARSE_CHECK(cusparseDestroySpMat(spmat_desc_));
1023 (
void)cusparseDestroySpMat(spmat_desc_);
1024 spmat_desc_ =
nullptr;
1026 set_descriptor_key(-1, -1, -1, 0);
1027 invalidate_ws_caches();
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);
1040template <
typename Scalar_>
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);
1050 expr.view().context().spmm_device_exec(expr.x(), *
this, Scalar_(1), Scalar_(0), GpuOp::NoTrans);
1055template <
typename Scalar_>
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");
1069 if (
view.nonZeros() == 0 ||
view.rows() == 0 ||
view.cols() == 0) {
1073 if (expr.x().cols() <= 1) {
1074 view.context().spmv_device_exec(expr.x(), *
this, expr.alpha(), expr.beta(), GpuOp::NoTrans);
1076 view.context().spmm_device_exec(expr.x(), *
this, expr.alpha(), expr.beta(), GpuOp::NoTrans);