phynexis-deps/eigen
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

GPU: Internal cleanup — dispatch() overloads, reference accessors, RAII handles

Browse Source 
libeigen/eigen!2758

Co-authored-by: Rasmus Munk Larsen <rmlarsen@gmail.com>
This commit is contained in:
Rasmus Munk Larsen
2026-08-12 13:56:28 -07:00
co-authored by Rasmus Munk Larsen
parent 83011bf347
commit 584169f336
13 changed files with 290 additions and 275 deletions
Download Patch File Download Diff File Expand all files Collapse all files
unsupported/Eigen/src/GPU/CuBlasSupport.h
+23 -7
View File
@@ -60,6 +60,22 @@ inline int64_t to_blas_dim(int64_t v) { return v; }
inline int to_blas_dim(int64_t v) { return to_blas_int(v); }
#endif
// RAII cuBLAS / cuBLASLt handles; the ownership flag supports handles borrowed from a gpu::Context.
struct CublasHandleDeleter {
bool owns = true;
void operator()(cublasHandle_t h) const noexcept {
if (owns && h) (void)cublasDestroy(h);
}
};
using UniqueCublasHandle = std::unique_ptr<std::remove_pointer_t<cublasHandle_t>, CublasHandleDeleter>;
struct CublasLtHandleDeleter {
void operator()(cublasLtHandle_t h) const noexcept {
if (h) (void)cublasLtDestroy(h);
}
};
using UniqueCublasLtHandle = std::unique_ptr<std::remove_pointer_t<cublasLtHandle_t>, CublasLtHandleDeleter>;
// cublasLtMatmul takes a compute type separate from the data type, which selects
// the precision policy:
// - Default: tensor-core algorithms via the cublasLtMatmul heuristics. For
@@ -272,7 +288,7 @@ template <typename Scalar>
void cublaslt_gemm(cublasLtHandle_t lt_handle, cublasHandle_t cublas_handle, cublasOperation_t transA,
cublasOperation_t transB, int64_t m, int64_t n, int64_t k, const Scalar* alpha, const Scalar* A,
int64_t lda, const Scalar* B, int64_t ldb, const Scalar* beta, Scalar* C, int64_t ldc,
DeviceBuffer* workspace, CublasLtPlanCache* plan_cache, std::size_t max_workspace_bytes,
DeviceBuffer& workspace, CublasLtPlanCache& plan_cache, std::size_t max_workspace_bytes,
cudaStream_t stream) {
constexpr cudaDataType_t dtype = cuda_data_type<Scalar>::value;
constexpr cublasComputeType_t compute = cuda_compute_type<Scalar>::value;
@@ -281,21 +297,21 @@ void cublaslt_gemm(cublasLtHandle_t lt_handle, cublasHandle_t cublas_handle, cub
// The key carries the leading dimensions so that strided views — e.g. SVD's
// thin VT/U slices — get distinct cache entries.
const CublasLtPlanKey key{m, n, k, lda, ldb, ldc, dtype, transA, transB};
CublasLtPlanEntry* entry = plan_cache->find(key);
CublasLtPlanEntry* entry = plan_cache.find(key);
if (!entry) {
entry = plan_cache->insert(key, CublasLtPlanEntry(lt_handle, key, compute, alpha_type, max_workspace_bytes));
entry = plan_cache.insert(key, CublasLtPlanEntry(lt_handle, key, compute, alpha_type, max_workspace_bytes));
}
if (entry->use_cublaslt) {
const size_t needed = entry->workspace_size;
if (needed > workspace->size()) {
if (needed > workspace.size()) {
// Sync only when freeing an existing buffer that may be in use.
if (workspace->get()) EIGEN_CUDA_RUNTIME_CHECK(cudaStreamSynchronize(stream));
*workspace = DeviceBuffer(needed);
if (workspace.get()) EIGEN_CUDA_RUNTIME_CHECK(cudaStreamSynchronize(stream));
workspace = DeviceBuffer(needed);
}
EIGEN_CUBLASLT_CHECK(cublasLtMatmul(lt_handle, entry->matmul_desc, alpha, A, entry->layout_A, B, entry->layout_B,
beta, C, entry->layout_C, C, entry->layout_C, &entry->algo, workspace->get(),
beta, C, entry->layout_C, C, entry->layout_C, &entry->algo, workspace.get(),
needed, stream));
} else {
// Fallback: cublasGemmEx for shapes/types that cublasLt cannot handle.
unsupported/Eigen/src/GPU/CuSolverSupport.h
+9
View File
@@ -92,6 +92,15 @@ struct CusolverParams {
CusolverParams& operator=(const CusolverParams&) = delete;
};
// RAII cuSOLVER dense handle; the ownership flag supports handles borrowed from a gpu::Context.
struct CusolverHandleDeleter {
bool owns = true;
void operator()(cusolverDnHandle_t h) const noexcept {
if (owns && h) (void)cusolverDnDestroy(h);
}
};
using UniqueCusolverHandle = std::unique_ptr<std::remove_pointer_t<cusolverDnHandle_t>, CusolverHandleDeleter>;
// Alias kept for compatibility; cuda_data_type<> in GpuSupport.h is canonical.
template <typename Scalar>
using cusolver_data_type = cuda_data_type<Scalar>;
unsupported/Eigen/src/GPU/DeviceDispatch.h
+26 -25
View File
@@ -9,7 +9,8 @@
// SPDX-License-Identifier: MPL-2.0
// Dispatch functions mapping DeviceMatrix expressions to NVIDIA library calls,
// plus the DeviceMatrix members that need a complete gpu::Context.
// plus the DeviceMatrix members that need a complete gpu::Context. The
// expression argument selects the dispatch() overload.
#ifndef EIGEN_GPU_DEVICE_DISPATCH_H
#define EIGEN_GPU_DEVICE_DISPATCH_H
@@ -34,8 +35,8 @@ bool aliases_device_memory(const DeviceMatrix<Scalar>& a, const DeviceMatrix<Sca
}
template <typename Lhs, typename Rhs>
void dispatch_gemm(Context& ctx, DeviceMatrix<scalar_type_t<Lhs>>& dst, const GemmExpr<Lhs, Rhs>& expr,
scalar_type_t<Lhs> beta_val, scalar_type_t<Lhs> alpha_scale = scalar_type_t<Lhs>(1)) {
void dispatch(Context& ctx, DeviceMatrix<scalar_type_t<Lhs>>& dst, const GemmExpr<Lhs, Rhs>& expr,
scalar_type_t<Lhs> beta_val, scalar_type_t<Lhs> alpha_scale = scalar_type_t<Lhs>(1)) {
using Scalar = scalar_type_t<Lhs>;
using traits_lhs = device_expr_traits<Lhs>;
using traits_rhs = device_expr_traits<Rhs>;
@@ -113,7 +114,7 @@ inline void oneshot_check_info(Context& ctx, OneShotSolverScratch& scratch, cons
}
template <typename Scalar, int UpLo>
void dispatch_llt_solve(Context& ctx, DeviceMatrix<Scalar>& dst, const LltSolveExpr<Scalar, UpLo>& expr) {
void dispatch(Context& ctx, DeviceMatrix<Scalar>& dst, const LltSolveExpr<Scalar, UpLo>& expr) {
const DeviceMatrix<Scalar>& A = expr.matrix();
const DeviceMatrix<Scalar>& B = expr.rhs();
@@ -136,7 +137,7 @@ void dispatch_llt_solve(Context& ctx, DeviceMatrix<Scalar>& dst, const LltSolveE
constexpr cublasFillMode_t uplo = cusolver_fill_mode<UpLo>::value;
const int64_t n = static_cast<int64_t>(A.rows());
constexpr cudaDataType_t dtype = cuda_data_type<Scalar>::value;
OneShotSolverScratch& scratch = *ctx.oneshotSolverScratch();
OneShotSolverScratch& scratch = ctx.oneshotSolverScratch();
{
const size_t mat_bytes = A.sizeInBytes();
// Context-owned grow-only scratch: no per-call allocation, no end-of-call sync.
@@ -171,7 +172,7 @@ void dispatch_llt_solve(Context& ctx, DeviceMatrix<Scalar>& dst, const LltSolveE
}
template <typename Scalar>
void dispatch_lu_solve(Context& ctx, DeviceMatrix<Scalar>& dst, const LuSolveExpr<Scalar>& expr) {
void dispatch(Context& ctx, DeviceMatrix<Scalar>& dst, const LuSolveExpr<Scalar>& expr) {
const DeviceMatrix<Scalar>& A = expr.matrix();
const DeviceMatrix<Scalar>& B = expr.rhs();
@@ -193,7 +194,7 @@ void dispatch_lu_solve(Context& ctx, DeviceMatrix<Scalar>& dst, const LuSolveExp
static thread_local CusolverParams params;
const int64_t n = static_cast<int64_t>(A.rows());
constexpr cudaDataType_t dtype = cuda_data_type<Scalar>::value;
OneShotSolverScratch& scratch = *ctx.oneshotSolverScratch();
OneShotSolverScratch& scratch = ctx.oneshotSolverScratch();
{
const size_t mat_bytes = A.sizeInBytes();
// Context-owned grow-only scratch: no per-call allocation, no end-of-call sync.
@@ -229,7 +230,7 @@ void dispatch_lu_solve(Context& ctx, DeviceMatrix<Scalar>& dst, const LuSolveExp
}
template <typename Scalar, int UpLo>
void dispatch_trsm(Context& ctx, DeviceMatrix<Scalar>& dst, const TrsmExpr<Scalar, UpLo>& expr) {
void dispatch(Context& ctx, DeviceMatrix<Scalar>& dst, const TrsmExpr<Scalar, UpLo>& expr) {
const DeviceMatrix<Scalar>& A = expr.matrix();
const DeviceMatrix<Scalar>& B = expr.rhs();
@@ -265,7 +266,7 @@ void dispatch_trsm(Context& ctx, DeviceMatrix<Scalar>& dst, const TrsmExpr<Scala
}
template <typename Scalar, int UpLo>
void dispatch_symm(Context& ctx, DeviceMatrix<Scalar>& dst, const SymmExpr<Scalar, UpLo>& expr) {
void dispatch(Context& ctx, DeviceMatrix<Scalar>& dst, const SymmExpr<Scalar, UpLo>& expr) {
const DeviceMatrix<Scalar>& A = expr.matrix();
const DeviceMatrix<Scalar>& B = expr.rhs();
@@ -290,7 +291,7 @@ void dispatch_symm(Context& ctx, DeviceMatrix<Scalar>& dst, const SymmExpr<Scala
dst.resize(m, n);
constexpr cublasFillMode_t uplo = (UpLo == Lower) ? CUBLAS_FILL_MODE_LOWER : CUBLAS_FILL_MODE_UPPER;
// The array keeps the host-pointer stack slots alive; see dispatch_gemm.
// The array keeps the host-pointer stack slots alive; see the GEMM dispatch.
Scalar scalars[2] = {Scalar(1), Scalar(0)};
EIGEN_CUBLAS_CHECK(cublasXsymm(ctx.cublasHandle(), CUBLAS_SIDE_LEFT, uplo, m, n, &scalars[0], A.data(), A.rows(),
@@ -300,8 +301,8 @@ void dispatch_symm(Context& ctx, DeviceMatrix<Scalar>& dst, const SymmExpr<Scala
}
template <typename Scalar, int UpLo>
void dispatch_syrk(Context& ctx, DeviceMatrix<Scalar>& dst, const SyrkExpr<Scalar, UpLo>& expr,
typename NumTraits<Scalar>::Real alpha_val, typename NumTraits<Scalar>::Real beta_val) {
void dispatch(Context& ctx, DeviceMatrix<Scalar>& dst, const SyrkExpr<Scalar, UpLo>& expr,
typename NumTraits<Scalar>::Real alpha_val, typename NumTraits<Scalar>::Real beta_val) {
using RealScalar = typename NumTraits<Scalar>::Real;
const DeviceMatrix<Scalar>& A = expr.matrix();
@@ -338,7 +339,7 @@ void dispatch_syrk(Context& ctx, DeviceMatrix<Scalar>& dst, const SyrkExpr<Scala
// dimensions, which always holds here since DeviceMatrix is fully dense).
template <typename Scalar>
void dispatch_geam(Context& ctx, DeviceMatrix<Scalar>& dst, const DeviceAddExpr<Scalar>& expr) {
void dispatch(Context& ctx, DeviceMatrix<Scalar>& dst, const DeviceAddExpr<Scalar>& expr) {
const DeviceMatrix<Scalar>& A = expr.A();
const DeviceMatrix<Scalar>& B = expr.B();
eigen_assert(A.rows() == B.rows() && A.cols() == B.cols());
@@ -351,7 +352,7 @@ void dispatch_geam(Context& ctx, DeviceMatrix<Scalar>& dst, const DeviceAddExpr<
if (m > 0 && n > 0) {
A.waitReady(ctx.stream());
B.waitReady(ctx.stream());
// See dispatch_gemm: array prevents compiler from eliding host-pointer stack slots.
// See the GEMM dispatch: array prevents compiler from eliding host-pointer stack slots.
Scalar scalars[2] = {expr.alpha(), expr.beta()};
EIGEN_CUBLAS_CHECK(cublasXgeam(ctx.cublasHandle(), CUBLAS_OP_N, CUBLAS_OP_N, m, n, &scalars[0], A.data(), m,
&scalars[1], B.data(), m, dst.data(), m));
@@ -369,53 +370,53 @@ class Assignment {
template <typename Lhs, typename Rhs>
DeviceMatrix<Scalar>& operator=(const GemmExpr<Lhs, Rhs>& expr) {
internal::dispatch_gemm(ctx_, dst_, expr, Scalar(0));
internal::dispatch(ctx_, dst_, expr, Scalar(0));
return dst_;
}
template <typename Lhs, typename Rhs>
DeviceMatrix<Scalar>& operator+=(const GemmExpr<Lhs, Rhs>& expr) {
internal::dispatch_gemm(ctx_, dst_, expr, Scalar(1));
internal::dispatch(ctx_, dst_, expr, Scalar(1));
return dst_;
}
template <typename Lhs, typename Rhs>
DeviceMatrix<Scalar>& operator-=(const GemmExpr<Lhs, Rhs>& expr) {
internal::dispatch_gemm(ctx_, dst_, expr, Scalar(1), Scalar(-1));
internal::dispatch(ctx_, dst_, expr, Scalar(1), Scalar(-1));
return dst_;
}
template <int UpLo>
DeviceMatrix<Scalar>& operator=(const LltSolveExpr<Scalar, UpLo>& expr) {
internal::dispatch_llt_solve(ctx_, dst_, expr);
internal::dispatch(ctx_, dst_, expr);
return dst_;
}
DeviceMatrix<Scalar>& operator=(const LuSolveExpr<Scalar>& expr) {
internal::dispatch_lu_solve(ctx_, dst_, expr);
internal::dispatch(ctx_, dst_, expr);
return dst_;
}
template <int UpLo>
DeviceMatrix<Scalar>& operator=(const TrsmExpr<Scalar, UpLo>& expr) {
internal::dispatch_trsm(ctx_, dst_, expr);
internal::dispatch(ctx_, dst_, expr);
return dst_;
}
template <int UpLo>
DeviceMatrix<Scalar>& operator=(const SymmExpr<Scalar, UpLo>& expr) {
internal::dispatch_symm(ctx_, dst_, expr);
internal::dispatch(ctx_, dst_, expr);
return dst_;
}
DeviceMatrix<Scalar>& operator=(const DeviceAddExpr<Scalar>& expr) {
internal::dispatch_geam(ctx_, dst_, expr);
internal::dispatch(ctx_, dst_, expr);
return dst_;
}
DeviceMatrix<Scalar>& operator=(const Scaled<DeviceMatrix<Scalar>>& expr) {
// geam with beta == 0: cuBLAS documents B as unread, so pass A twice.
internal::dispatch_geam(ctx_, dst_, DeviceAddExpr<Scalar>(expr.scalar(), expr.inner(), Scalar(0), expr.inner()));
internal::dispatch(ctx_, dst_, DeviceAddExpr<Scalar>(expr.scalar(), expr.inner(), Scalar(0), expr.inner()));
return dst_;
}
@@ -538,7 +539,7 @@ template <typename Scalar_, int UpLo_>
void SelfAdjointView<Scalar_, UpLo_>::rankUpdate(const DeviceMatrix<Scalar_>& A, RealScalar alpha) {
SyrkExpr<Scalar_, UpLo_> expr(A);
RealScalar beta = matrix().empty() ? RealScalar(0) : RealScalar(1);
internal::dispatch_syrk(Context::threadLocal(), matrix(), expr, alpha, beta);
internal::dispatch(Context::threadLocal(), matrix(), expr, alpha, beta);
}
namespace internal {
@@ -755,7 +756,7 @@ DeviceMatrix<Scalar_>& DeviceMatrix<Scalar_>::operator-=(const DeviceScaledDevic
// this = alpha * A + beta * B (cuBLAS geam)
template <typename Scalar_>
DeviceMatrix<Scalar_>& DeviceMatrix<Scalar_>::operator=(const DeviceAddExpr<Scalar_>& expr) {
internal::dispatch_geam(Context::threadLocal(), *this, expr);
internal::dispatch(Context::threadLocal(), *this, expr);
return *this;
}
unsupported/Eigen/src/GPU/GpuContext.h
+49 -42
View File
@@ -41,7 +41,7 @@ constexpr size_t kOneShotHostInfoBytes = kOneShotInfoBytes;
// (d_A.llt().solve(d_B), d_A.lu().solve(d_B)), so repeated one-shot solves on
// a Context perform no per-call device or pinned-host allocations. Holds only
// CUDA-runtime types (no cuSOLVER types) to keep the lazy-linking property of
// Context. Used by dispatch_llt_solve / dispatch_lu_solve in DeviceDispatch.h.
// Context. Used by the one-shot solve dispatches in DeviceDispatch.h.
struct OneShotSolverScratch {
DeviceBuffer d_factor;
DeviceBuffer d_ipiv;
@@ -80,26 +80,20 @@ class Context {
public:
/** Create a new context with a dedicated CUDA stream. */
Context() {
EIGEN_CUDA_RUNTIME_CHECK(cudaStreamCreate(&stream_));
cudaStream_t s = nullptr;
EIGEN_CUDA_RUNTIME_CHECK(cudaStreamCreate(&s));
stream_ = internal::UniqueStream(s);
init_cublas();
}
/** Create a context on an existing stream (e.g., stream 0 = nullptr).
* The caller retains ownership of the stream — this context will not destroy it. */
explicit Context(cudaStream_t stream) : stream_(stream), owns_stream_(false) { init_cublas(); }
~Context() {
// Indirect calls keep cusolverDnDestroy / cusparseDestroy out of TUs that
// never call cusolverHandle() / cusparseHandle() (e.g. the cufft test).
if (cusparse_destroyer_) (void)cusparse_destroyer_(cusparse_);
if (cusolver_destroyer_) (void)cusolver_destroyer_(cusolver_);
// Release plan-cache descriptors before tearing down the cuBLASLt handle.
gemm_plan_cache_.clear();
if (cublas_lt_) (void)cublasLtDestroy(cublas_lt_);
if (cublas_) (void)cublasDestroy(cublas_);
if (owns_stream_ && stream_) (void)cudaStreamDestroy(stream_);
explicit Context(cudaStream_t stream) : stream_(stream, internal::CudaStreamDeleter{/*owns=*/false}) {
init_cublas();
}
~Context() = default;
Context(const Context&) = delete;
Context& operator=(const Context&) = delete;
Context(Context&&) = delete;
@@ -116,7 +110,10 @@ class Context {
* CUDA context has already been torn down. These errors are harmless and are
* suppressed in the destructor, but they can produce noise in test output.
* To avoid this, call cudaDeviceReset() only after all Context instances
* (including thread-local ones) have been destroyed. */
* (including thread-local ones) have been destroyed — or create and own a
* Context and install it with setThreadLocal(): the lazily-created default
* is then never constructed, and teardown order is fully under application
* control. */
static Context& threadLocal() {
Context* override = tl_override_ptr();
if (override) return *override;
@@ -129,39 +126,42 @@ class Context {
* Pass nullptr to restore the lazily-created default. */
static void setThreadLocal(Context* ctx) { tl_override_ptr() = ctx; }
cudaStream_t stream() const { return stream_; }
cublasHandle_t cublasHandle() const { return cublas_; }
cudaStream_t stream() const { return stream_.get(); }
cublasHandle_t cublasHandle() const { return cublas_.get(); }
/** Returns the cuSOLVER handle, creating it on first call. */
cusolverDnHandle_t cusolverHandle() {
if (!cusolver_) {
EIGEN_CUSOLVER_CHECK(cusolverDnCreate(&cusolver_));
EIGEN_CUSOLVER_CHECK(cusolverDnSetStream(cusolver_, stream_));
cusolver_destroyer_ = &destroyCusolver;
cusolverDnHandle_t h = nullptr;
EIGEN_CUSOLVER_CHECK(cusolverDnCreate(&h));
cusolver_ = LazyCusolverHandle(h, &destroyCusolver);
EIGEN_CUSOLVER_CHECK(cusolverDnSetStream(h, stream_.get()));
}
return cusolver_;
return cusolver_.get();
}
/** cuBLASLt handle (lazy-initialized on first GEMM call). */
cublasLtHandle_t cublasLtHandle() {
if (!cublas_lt_) {
EIGEN_CUBLAS_CHECK(cublasLtCreate(&cublas_lt_));
cublasLtHandle_t h = nullptr;
EIGEN_CUBLAS_CHECK(cublasLtCreate(&h));
cublas_lt_ = internal::UniqueCublasLtHandle(h);
}
return cublas_lt_;
return cublas_lt_.get();
}
/** Workspace buffer for cublasLtMatmul (grown lazily by cublaslt_gemm).
* Not thread-safe — all GEMM calls must be on this context's stream. */
internal::DeviceBuffer* gemmWorkspace() { return &gemm_workspace_; }
internal::DeviceBuffer& gemmWorkspace() { return gemm_workspace_; }
/** Plan cache for cublasLtMatmul (caches descriptors and selected algorithm
* by shape to avoid per-call overhead). Same thread-safety as workspace. */
internal::CublasLtPlanCache* gemmPlanCache() { return &gemm_plan_cache_; }
internal::CublasLtPlanCache& gemmPlanCache() { return gemm_plan_cache_; }
/** Grow-only scratch for the one-shot solver expressions
* (d_A.llt().solve(d_B), d_A.lu().solve(d_B)). Same thread-safety rules as
* the GEMM workspace: all uses must be on this context's stream. */
internal::OneShotSolverScratch* oneshotSolverScratch() { return &oneshot_solver_scratch_; }
internal::OneShotSolverScratch& oneshotSolverScratch() { return oneshot_solver_scratch_; }
/** Workspace ceiling passed to the cublasLtMatmul heuristic at plan-creation time.
* Defaults to internal::kCublasLtMaxWorkspaceBytes (compile-time configurable via
@@ -170,40 +170,45 @@ class Context {
/** Override the workspace ceiling for future plan-cache misses on this context.
* The cap is consulted at plan-creation time only; pre-existing cached plans
* keep the cap they were built with. Call gemmPlanCache()->clear() to force
* keep the cap they were built with. Call gemmPlanCache().clear() to force
* re-selection under the new cap. */
void setCublasLtMaxWorkspaceBytes(std::size_t bytes) { cublaslt_max_workspace_bytes_ = bytes; }
/** cuSPARSE handle, created on first use. */
cusparseHandle_t cusparseHandle() {
if (!cusparse_) {
cusparseStatus_t s1 = cusparseCreate(&cusparse_);
cusparseHandle_t h = nullptr;
cusparseStatus_t s1 = cusparseCreate(&h);
eigen_assert(s1 == CUSPARSE_STATUS_SUCCESS && "cusparseCreate failed");
EIGEN_UNUSED_VARIABLE(s1);
cusparseStatus_t s2 = cusparseSetStream(cusparse_, stream_);
cusparse_ = LazyCusparseHandle(h, &destroyCusparse);
cusparseStatus_t s2 = cusparseSetStream(h, stream_.get());
eigen_assert(s2 == CUSPARSE_STATUS_SUCCESS && "cusparseSetStream failed");
EIGEN_UNUSED_VARIABLE(s2);
cusparse_destroyer_ = &destroyCusparse;
}
return cusparse_;
return cusparse_.get();
}
private:
static cusolverStatus_t destroyCusolver(cusolverDnHandle_t h) { return cusolverDnDestroy(h); }
static cusparseStatus_t destroyCusparse(cusparseHandle_t h) { return cusparseDestroy(h); }
cudaStream_t stream_ = nullptr;
cublasHandle_t cublas_ = nullptr;
cusolverDnHandle_t cusolver_ = nullptr;
cusolverStatus_t (*cusolver_destroyer_)(cusolverDnHandle_t) = nullptr;
cublasLtHandle_t cublas_lt_ = nullptr; // lazy
cusparseHandle_t cusparse_ = nullptr; // lazy
cusparseStatus_t (*cusparse_destroyer_)(cusparseHandle_t) = nullptr;
internal::DeviceBuffer gemm_workspace_; // lazy
// Function-pointer deleters keep cusolverDnDestroy / cusparseDestroy referenced only by TUs that create handles.
using LazyCusolverHandle =
std::unique_ptr<std::remove_pointer_t<cusolverDnHandle_t>, cusolverStatus_t (*)(cusolverDnHandle_t)>;
using LazyCusparseHandle =
std::unique_ptr<std::remove_pointer_t<cusparseHandle_t>, cusparseStatus_t (*)(cusparseHandle_t)>;
// Destroyed in reverse declaration order: the plan cache before the cuBLASLt handle, the stream last.
internal::UniqueStream stream_;
internal::UniqueCublasHandle cublas_;
LazyCusolverHandle cusolver_{nullptr, nullptr};
LazyCusparseHandle cusparse_{nullptr, nullptr};
internal::UniqueCublasLtHandle cublas_lt_; // lazy
internal::DeviceBuffer gemm_workspace_; // lazy
internal::CublasLtPlanCache gemm_plan_cache_{internal::kCublasLtPlanCacheCapacity};
internal::OneShotSolverScratch oneshot_solver_scratch_; // grow-only
std::size_t cublaslt_max_workspace_bytes_ = internal::kCublasLtMaxWorkspaceBytes;
bool owns_stream_ = true;
static Context*& tl_override_ptr() {
thread_local Context* ptr = nullptr;
@@ -211,8 +216,10 @@ class Context {
}
void init_cublas() {
EIGEN_CUBLAS_CHECK(cublasCreate(&cublas_));
EIGEN_CUBLAS_CHECK(cublasSetStream(cublas_, stream_));
cublasHandle_t h = nullptr;
EIGEN_CUBLAS_CHECK(cublasCreate(&h));
cublas_ = internal::UniqueCublasHandle(h);
EIGEN_CUBLAS_CHECK(cublasSetStream(h, stream_.get()));
}
};
unsupported/Eigen/src/GPU/GpuEigenSolver.h
+9 -9
View File
@@ -100,7 +100,7 @@ class SelfAdjointEigenSolver {
SelfAdjointEigenSolver& compute(const EigenBase<InputType>& A, int options = ComputeEigenvectors) {
// Route through the adopting overload: the freshly uploaded matrix is
// decomposed in place (syevd overwrites its input) — no second device copy.
return compute(DeviceMatrix<Scalar>::fromHost(A.derived(), solver_ctx_.stream_), options);
return compute(DeviceMatrix<Scalar>::fromHost(A.derived(), solver_ctx_.stream()), options);
}
SelfAdjointEigenSolver& compute(const DeviceMatrix<Scalar>& d_A, int options = ComputeEigenvectors) {
@@ -109,7 +109,7 @@ class SelfAdjointEigenSolver {
const size_t mat_bytes = static_cast<size_t>(lda_) * static_cast<size_t>(n_) * sizeof(Scalar);
internal::ensure_sized(d_A_, mat_bytes);
EIGEN_CUDA_RUNTIME_CHECK(
cudaMemcpyAsync(d_A_.get(), d_A.data(), mat_bytes, cudaMemcpyDeviceToDevice, solver_ctx_.stream_));
cudaMemcpyAsync(d_A_.get(), d_A.data(), mat_bytes, cudaMemcpyDeviceToDevice, solver_ctx_.stream()));
factorize();
return *this;
@@ -168,7 +168,7 @@ class SelfAdjointEigenSolver {
DeviceMatrix<RealScalar> d_eigenvalues() const {
eigen_assert(solver_ctx_.info() == Success);
auto v = DeviceMatrix<RealScalar>::view(static_cast<RealScalar*>(d_W_.get()), n_, 1);
v.recordReady(solver_ctx_.stream_);
v.recordReady(solver_ctx_.stream());
return v;
}
@@ -178,11 +178,11 @@ class SelfAdjointEigenSolver {
eigen_assert(solver_ctx_.info() == Success);
eigen_assert(compute_eigenvectors_ && "d_eigenvectors() requires ComputeEigenvectors option");
auto v = DeviceMatrix<Scalar>::view(static_cast<Scalar*>(d_A_.get()), n_, n_);
v.recordReady(solver_ctx_.stream_);
v.recordReady(solver_ctx_.stream());
return v;
}
cudaStream_t stream() const { return solver_ctx_.stream_; }
cudaStream_t stream() const { return solver_ctx_.stream(); }
private:
mutable internal::GpuSolverContext solver_ctx_;
@@ -206,7 +206,7 @@ class SelfAdjointEigenSolver {
return false;
}
lda_ = n_;
d_A.waitReady(solver_ctx_.stream_);
d_A.waitReady(solver_ctx_.stream());
return true;
}
@@ -223,14 +223,14 @@ class SelfAdjointEigenSolver {
constexpr cublasFillMode_t uplo = CUBLAS_FILL_MODE_LOWER;
size_t dev_ws = 0, host_ws = 0;
EIGEN_CUSOLVER_CHECK(cusolverDnXsyevd_bufferSize(solver_ctx_.cusolver_, solver_ctx_.params_.p, jobz, uplo, n_,
dtype, d_A_.get(), lda_, rtype, d_W_.get(), dtype, &dev_ws,
EIGEN_CUSOLVER_CHECK(cusolverDnXsyevd_bufferSize(solver_ctx_.cusolverHandle(), solver_ctx_.params_.p, jobz, uplo,
n_, dtype, d_A_.get(), lda_, rtype, d_W_.get(), dtype, &dev_ws,
&host_ws));
solver_ctx_.ensure_scratch(dev_ws);
solver_ctx_.h_workspace_.resize(host_ws);
EIGEN_CUSOLVER_CHECK(cusolverDnXsyevd(solver_ctx_.cusolver_, solver_ctx_.params_.p, jobz, uplo, n_, dtype,
EIGEN_CUSOLVER_CHECK(cusolverDnXsyevd(solver_ctx_.cusolverHandle(), solver_ctx_.params_.p, jobz, uplo, n_, dtype,
d_A_.get(), lda_, rtype, d_W_.get(), dtype, solver_ctx_.scratch_workspace(),
dev_ws, host_ws > 0 ? solver_ctx_.h_workspace_.data() : nullptr, host_ws,
solver_ctx_.scratch_info()));
unsupported/Eigen/src/GPU/GpuLLT.h
+18 -17
View File
@@ -110,7 +110,7 @@ class LLT {
lda_ = static_cast<int64_t>(mat.rows());
allocate_factor_storage();
EIGEN_CUDA_RUNTIME_CHECK(
cudaMemcpyAsync(d_factor_.get(), mat.data(), factorBytes(), cudaMemcpyHostToDevice, solver_ctx_.stream_));
cudaMemcpyAsync(d_factor_.get(), mat.data(), factorBytes(), cudaMemcpyHostToDevice, solver_ctx_.stream()));
factorize();
return *this;
@@ -122,10 +122,10 @@ class LLT {
if (!begin_compute(d_A.rows())) return *this;
lda_ = static_cast<int64_t>(d_A.rows());
d_A.waitReady(solver_ctx_.stream_);
d_A.waitReady(solver_ctx_.stream());
allocate_factor_storage();
EIGEN_CUDA_RUNTIME_CHECK(
cudaMemcpyAsync(d_factor_.get(), d_A.data(), factorBytes(), cudaMemcpyDeviceToDevice, solver_ctx_.stream_));
cudaMemcpyAsync(d_factor_.get(), d_A.data(), factorBytes(), cudaMemcpyDeviceToDevice, solver_ctx_.stream()));
factorize();
return *this;
@@ -137,7 +137,7 @@ class LLT {
if (!begin_compute(d_A.rows())) return *this;
lda_ = static_cast<int64_t>(d_A.rows());
d_A.waitReady(solver_ctx_.stream_);
d_A.waitReady(solver_ctx_.stream());
d_factor_ = internal::DeviceBuffer::adopt(static_cast<void*>(d_A.release()), factorBytes());
factorize();
@@ -158,16 +158,16 @@ class LLT {
const int64_t ldb = static_cast<int64_t>(rhs.rows());
internal::DeviceBuffer d_x(rhsBytes(nrhs, ldb));
EIGEN_CUDA_RUNTIME_CHECK(
cudaMemcpyAsync(d_x.get(), rhs.data(), rhsBytes(nrhs, ldb), cudaMemcpyHostToDevice, solver_ctx_.stream_));
cudaMemcpyAsync(d_x.get(), rhs.data(), rhsBytes(nrhs, ldb), cudaMemcpyHostToDevice, solver_ctx_.stream()));
DeviceMatrix<Scalar> d_X = solve_impl(nrhs, ldb, std::move(d_x));
PlainMatrix X(n_, B.cols());
int solve_info = 0;
EIGEN_CUDA_RUNTIME_CHECK(
cudaMemcpyAsync(X.data(), d_X.data(), rhsBytes(nrhs, ldb), cudaMemcpyDeviceToHost, solver_ctx_.stream_));
cudaMemcpyAsync(X.data(), d_X.data(), rhsBytes(nrhs, ldb), cudaMemcpyDeviceToHost, solver_ctx_.stream()));
EIGEN_CUDA_RUNTIME_CHECK(cudaMemcpyAsync(&solve_info, solver_ctx_.scratch_info(), sizeof(int),
cudaMemcpyDeviceToHost, solver_ctx_.stream_));
EIGEN_CUDA_RUNTIME_CHECK(cudaStreamSynchronize(solver_ctx_.stream_));
cudaMemcpyDeviceToHost, solver_ctx_.stream()));
EIGEN_CUDA_RUNTIME_CHECK(cudaStreamSynchronize(solver_ctx_.stream()));
eigen_assert(solve_info == 0 && "cusolverDnXpotrs reported an error");
return X;
@@ -181,12 +181,12 @@ class LLT {
DeviceMatrix<Scalar> solve(const DeviceMatrix<Scalar>& d_B) const {
eigen_assert(solver_ctx_.info() == Success && "LLT::solve called on a failed or uninitialized factorization");
eigen_assert(d_B.rows() == n_);
d_B.waitReady(solver_ctx_.stream_);
d_B.waitReady(solver_ctx_.stream());
const int64_t nrhs = static_cast<int64_t>(d_B.cols());
const int64_t ldb = static_cast<int64_t>(d_B.rows());
internal::DeviceBuffer d_x(rhsBytes(nrhs, ldb));
EIGEN_CUDA_RUNTIME_CHECK(
cudaMemcpyAsync(d_x.get(), d_B.data(), rhsBytes(nrhs, ldb), cudaMemcpyDeviceToDevice, solver_ctx_.stream_));
cudaMemcpyAsync(d_x.get(), d_B.data(), rhsBytes(nrhs, ldb), cudaMemcpyDeviceToDevice, solver_ctx_.stream()));
return solve_impl(nrhs, ldb, std::move(d_x));
}
@@ -195,7 +195,7 @@ class LLT {
DeviceMatrix<Scalar> solve(DeviceMatrix<Scalar>&& d_B) const {
eigen_assert(solver_ctx_.info() == Success && "LLT::solve called on a failed or uninitialized factorization");
eigen_assert(d_B.rows() == n_);
d_B.waitReady(solver_ctx_.stream_);
d_B.waitReady(solver_ctx_.stream());
const int64_t nrhs = static_cast<int64_t>(d_B.cols());
const int64_t ldb = static_cast<int64_t>(d_B.rows());
internal::DeviceBuffer d_x = internal::DeviceBuffer::adopt(static_cast<void*>(d_B.release()), rhsBytes(nrhs, ldb));
@@ -205,7 +205,7 @@ class LLT {
ComputationInfo info() const { return solver_ctx_.info(); }
Index rows() const { return n_; }
Index cols() const { return n_; }
cudaStream_t stream() const { return solver_ctx_.stream_; }
cudaStream_t stream() const { return solver_ctx_.stream(); }
private:
mutable internal::GpuSolverContext solver_ctx_;
@@ -234,12 +234,12 @@ class LLT {
constexpr cudaDataType_t dtype = internal::cusolver_data_type<Scalar>::value;
constexpr cublasFillMode_t uplo = internal::cusolver_fill_mode<UpLo_>::value;
EIGEN_CUSOLVER_CHECK(cusolverDnXpotrs(solver_ctx_.cusolver_, solver_ctx_.params_.p, uplo, n_, nrhs, dtype,
EIGEN_CUSOLVER_CHECK(cusolverDnXpotrs(solver_ctx_.cusolverHandle(), solver_ctx_.params_.p, uplo, n_, nrhs, dtype,
d_factor_.get(), lda_, dtype, d_x.get(), ldb, solver_ctx_.scratch_info()));
DeviceMatrix<Scalar> result =
DeviceMatrix<Scalar>::adopt(static_cast<Scalar*>(d_x.release()), n_, static_cast<Index>(nrhs));
result.recordReady(solver_ctx_.stream_);
result.recordReady(solver_ctx_.stream());
return result;
}
@@ -250,13 +250,14 @@ class LLT {
solver_ctx_.mark_pending();
size_t dev_ws_bytes = 0, host_ws_bytes = 0;
EIGEN_CUSOLVER_CHECK(cusolverDnXpotrf_bufferSize(solver_ctx_.cusolver_, solver_ctx_.params_.p, uplo, n_, dtype,
d_factor_.get(), lda_, dtype, &dev_ws_bytes, &host_ws_bytes));
EIGEN_CUSOLVER_CHECK(cusolverDnXpotrf_bufferSize(solver_ctx_.cusolverHandle(), solver_ctx_.params_.p, uplo, n_,
dtype, d_factor_.get(), lda_, dtype, &dev_ws_bytes,
&host_ws_bytes));
solver_ctx_.ensure_scratch(dev_ws_bytes);
solver_ctx_.h_workspace_.resize(host_ws_bytes);
EIGEN_CUSOLVER_CHECK(cusolverDnXpotrf(solver_ctx_.cusolver_, solver_ctx_.params_.p, uplo, n_, dtype,
EIGEN_CUSOLVER_CHECK(cusolverDnXpotrf(solver_ctx_.cusolverHandle(), solver_ctx_.params_.p, uplo, n_, dtype,
d_factor_.get(), lda_, dtype, solver_ctx_.scratch_workspace(), dev_ws_bytes,
host_ws_bytes > 0 ? solver_ctx_.h_workspace_.data() : nullptr, host_ws_bytes,
solver_ctx_.scratch_info()));
unsupported/Eigen/src/GPU/GpuLU.h
+16 -16
View File
@@ -105,7 +105,7 @@ class LU {
lda_ = static_cast<int64_t>(mat.rows());
allocate_lu_storage();
EIGEN_CUDA_RUNTIME_CHECK(
cudaMemcpyAsync(d_lu_.get(), mat.data(), matrixBytes(), cudaMemcpyHostToDevice, solver_ctx_.stream_));
cudaMemcpyAsync(d_lu_.get(), mat.data(), matrixBytes(), cudaMemcpyHostToDevice, solver_ctx_.stream()));
factorize();
return *this;
@@ -117,10 +117,10 @@ class LU {
if (!begin_compute(d_A.rows())) return *this;
lda_ = static_cast<int64_t>(d_A.rows());
d_A.waitReady(solver_ctx_.stream_);
d_A.waitReady(solver_ctx_.stream());
allocate_lu_storage();
EIGEN_CUDA_RUNTIME_CHECK(
cudaMemcpyAsync(d_lu_.get(), d_A.data(), matrixBytes(), cudaMemcpyDeviceToDevice, solver_ctx_.stream_));
cudaMemcpyAsync(d_lu_.get(), d_A.data(), matrixBytes(), cudaMemcpyDeviceToDevice, solver_ctx_.stream()));
factorize();
return *this;
@@ -132,7 +132,7 @@ class LU {
if (!begin_compute(d_A.rows())) return *this;
lda_ = static_cast<int64_t>(d_A.rows());
d_A.waitReady(solver_ctx_.stream_);
d_A.waitReady(solver_ctx_.stream());
d_lu_ = internal::DeviceBuffer::adopt(static_cast<void*>(d_A.release()), matrixBytes());
factorize();
@@ -157,16 +157,16 @@ class LU {
const int64_t ldb = static_cast<int64_t>(rhs.rows());
internal::DeviceBuffer d_x(matrixBytes(nrhs, ldb));
EIGEN_CUDA_RUNTIME_CHECK(
cudaMemcpyAsync(d_x.get(), rhs.data(), matrixBytes(nrhs, ldb), cudaMemcpyHostToDevice, solver_ctx_.stream_));
cudaMemcpyAsync(d_x.get(), rhs.data(), matrixBytes(nrhs, ldb), cudaMemcpyHostToDevice, solver_ctx_.stream()));
DeviceMatrix<Scalar> d_X = solve_impl(nrhs, ldb, op, std::move(d_x));
PlainMatrix X(n_, B.cols());
int solve_info = 0;
EIGEN_CUDA_RUNTIME_CHECK(
cudaMemcpyAsync(X.data(), d_X.data(), matrixBytes(nrhs, ldb), cudaMemcpyDeviceToHost, solver_ctx_.stream_));
cudaMemcpyAsync(X.data(), d_X.data(), matrixBytes(nrhs, ldb), cudaMemcpyDeviceToHost, solver_ctx_.stream()));
EIGEN_CUDA_RUNTIME_CHECK(cudaMemcpyAsync(&solve_info, solver_ctx_.scratch_info(), sizeof(int),
cudaMemcpyDeviceToHost, solver_ctx_.stream_));
EIGEN_CUDA_RUNTIME_CHECK(cudaStreamSynchronize(solver_ctx_.stream_));
cudaMemcpyDeviceToHost, solver_ctx_.stream()));
EIGEN_CUDA_RUNTIME_CHECK(cudaStreamSynchronize(solver_ctx_.stream()));
eigen_assert(solve_info == 0 && "cusolverDnXgetrs reported an error");
return X;
@@ -180,12 +180,12 @@ class LU {
DeviceMatrix<Scalar> solve(const DeviceMatrix<Scalar>& d_B, GpuOp op = GpuOp::NoTrans) const {
eigen_assert(solver_ctx_.info() == Success && "LU::solve called on a failed or uninitialized factorization");
eigen_assert(d_B.rows() == n_);
d_B.waitReady(solver_ctx_.stream_);
d_B.waitReady(solver_ctx_.stream());
const int64_t nrhs = static_cast<int64_t>(d_B.cols());
const int64_t ldb = static_cast<int64_t>(d_B.rows());
internal::DeviceBuffer d_x(matrixBytes(nrhs, ldb));
EIGEN_CUDA_RUNTIME_CHECK(
cudaMemcpyAsync(d_x.get(), d_B.data(), matrixBytes(nrhs, ldb), cudaMemcpyDeviceToDevice, solver_ctx_.stream_));
cudaMemcpyAsync(d_x.get(), d_B.data(), matrixBytes(nrhs, ldb), cudaMemcpyDeviceToDevice, solver_ctx_.stream()));
return solve_impl(nrhs, ldb, op, std::move(d_x));
}
@@ -194,7 +194,7 @@ class LU {
DeviceMatrix<Scalar> solve(DeviceMatrix<Scalar>&& d_B, GpuOp op = GpuOp::NoTrans) const {
eigen_assert(solver_ctx_.info() == Success && "LU::solve called on a failed or uninitialized factorization");
eigen_assert(d_B.rows() == n_);
d_B.waitReady(solver_ctx_.stream_);
d_B.waitReady(solver_ctx_.stream());
const int64_t nrhs = static_cast<int64_t>(d_B.cols());
const int64_t ldb = static_cast<int64_t>(d_B.rows());
internal::DeviceBuffer d_x =
@@ -205,7 +205,7 @@ class LU {
ComputationInfo info() const { return solver_ctx_.info(); }
Index rows() const { return n_; }
Index cols() const { return n_; }
cudaStream_t stream() const { return solver_ctx_.stream_; }
cudaStream_t stream() const { return solver_ctx_.stream(); }
private:
mutable internal::GpuSolverContext solver_ctx_;
@@ -235,13 +235,13 @@ class LU {
constexpr cudaDataType_t dtype = internal::cusolver_data_type<Scalar>::value;
const cublasOperation_t trans = internal::to_cublas_op(op);
EIGEN_CUSOLVER_CHECK(cusolverDnXgetrs(solver_ctx_.cusolver_, solver_ctx_.params_.p, trans, n_, nrhs, dtype,
EIGEN_CUSOLVER_CHECK(cusolverDnXgetrs(solver_ctx_.cusolverHandle(), solver_ctx_.params_.p, trans, n_, nrhs, dtype,
d_lu_.get(), lda_, static_cast<const int64_t*>(d_ipiv_.get()), dtype,
d_x.get(), ldb, solver_ctx_.scratch_info()));
DeviceMatrix<Scalar> result =
DeviceMatrix<Scalar>::adopt(static_cast<Scalar*>(d_x.release()), n_, static_cast<Index>(nrhs));
result.recordReady(solver_ctx_.stream_);
result.recordReady(solver_ctx_.stream());
return result;
}
@@ -254,14 +254,14 @@ class LU {
internal::ensure_sized(d_ipiv_, ipiv_bytes);
size_t dev_ws_bytes = 0, host_ws_bytes = 0;
EIGEN_CUSOLVER_CHECK(cusolverDnXgetrf_bufferSize(solver_ctx_.cusolver_, solver_ctx_.params_.p, n_, n_, dtype,
EIGEN_CUSOLVER_CHECK(cusolverDnXgetrf_bufferSize(solver_ctx_.cusolverHandle(), solver_ctx_.params_.p, n_, n_, dtype,
d_lu_.get(), lda_, dtype, &dev_ws_bytes, &host_ws_bytes));
solver_ctx_.ensure_scratch(dev_ws_bytes);
solver_ctx_.h_workspace_.resize(host_ws_bytes);
EIGEN_CUSOLVER_CHECK(cusolverDnXgetrf(
solver_ctx_.cusolver_, solver_ctx_.params_.p, n_, n_, dtype, d_lu_.get(), lda_,
solver_ctx_.cusolverHandle(), solver_ctx_.params_.p, n_, n_, dtype, d_lu_.get(), lda_,
static_cast<int64_t*>(d_ipiv_.get()), dtype, solver_ctx_.scratch_workspace(), dev_ws_bytes,
host_ws_bytes > 0 ? solver_ctx_.h_workspace_.data() : nullptr, host_ws_bytes, solver_ctx_.scratch_info()));
unsupported/Eigen/src/GPU/GpuQR.h
+30 -30
View File
@@ -101,7 +101,7 @@ class QR {
// uploaded matrix is factored in place (geqrf overwrites its input), so no
// second device copy is made. The wide-matrix transpose runs on the GPU
// (via cublasXgeam) inside the device-input path; no host transpose.
return compute(DeviceMatrix<Scalar>::fromHost(A.derived(), solver_ctx_.stream_));
return compute(DeviceMatrix<Scalar>::fromHost(A.derived(), solver_ctx_.stream()));
}
QR& compute(const DeviceMatrix<Scalar>& d_A) {
@@ -113,7 +113,7 @@ class QR {
const size_t mat_bytes = factorBytes();
allocate_factor_storage(mat_bytes);
EIGEN_CUDA_RUNTIME_CHECK(
cudaMemcpyAsync(d_qr_.get(), d_A.data(), mat_bytes, cudaMemcpyDeviceToDevice, solver_ctx_.stream_));
cudaMemcpyAsync(d_qr_.get(), d_A.data(), mat_bytes, cudaMemcpyDeviceToDevice, solver_ctx_.stream()));
}
factorize();
@@ -160,7 +160,7 @@ class QR {
DeviceMatrix<Scalar> solve(const DeviceMatrix<Scalar>& d_B) const {
eigen_assert(solver_ctx_.info() == Success && "QR::solve called on a failed or uninitialized factorization");
eigen_assert(d_B.rows() == m_);
d_B.waitReady(solver_ctx_.stream_);
d_B.waitReady(solver_ctx_.stream());
if (!transposed_) {
return solve_overdetermined_device(d_B);
@@ -172,7 +172,7 @@ class QR {
Index rows() const { return m_; }
Index cols() const { return n_; }
cudaStream_t stream() const { return solver_ctx_.stream_; }
cudaStream_t stream() const { return solver_ctx_.stream(); }
/** Upper-triangular factor R (k × n) of A = Q R. Available only for m >= n. */
PlainMatrix matrixR() const {
@@ -216,7 +216,7 @@ class QR {
}
transposed_ = (m_ < n_);
lda_ = static_cast<int64_t>(transposed_ ? n_ : m_);
d_A.waitReady(solver_ctx_.stream_);
d_A.waitReady(solver_ctx_.stream());
return true;
}
@@ -226,7 +226,7 @@ class QR {
void transpose_into_factor(const DeviceMatrix<Scalar>& d_A) {
allocate_factor_storage(factorBytes());
Scalar alpha_one(1), beta_zero(0);
EIGEN_CUBLAS_CHECK(internal::cublasXgeam(solver_ctx_.cublas_, CUBLAS_OP_C, CUBLAS_OP_N, n_, m_, &alpha_one,
EIGEN_CUBLAS_CHECK(internal::cublasXgeam(solver_ctx_.cublasHandle(), CUBLAS_OP_C, CUBLAS_OP_N, n_, m_, &alpha_one,
d_A.data(), d_A.rows(), &beta_zero, static_cast<const Scalar*>(nullptr),
n_, static_cast<Scalar*>(d_qr_.get()), n_));
}
@@ -241,16 +241,16 @@ class QR {
const int64_t fm = factor_rows();
const int64_t fn = factor_cols();
size_t dev_ws = 0, host_ws = 0;
EIGEN_CUSOLVER_CHECK(cusolverDnXgeqrf_bufferSize(solver_ctx_.cusolver_, solver_ctx_.params_.p, fm, fn, dtype,
EIGEN_CUSOLVER_CHECK(cusolverDnXgeqrf_bufferSize(solver_ctx_.cusolverHandle(), solver_ctx_.params_.p, fm, fn, dtype,
d_qr_.get(), lda_, dtype, d_tau_.get(), dtype, &dev_ws, &host_ws));
solver_ctx_.ensure_scratch(dev_ws);
solver_ctx_.h_workspace_.resize(host_ws);
EIGEN_CUSOLVER_CHECK(cusolverDnXgeqrf(solver_ctx_.cusolver_, solver_ctx_.params_.p, fm, fn, dtype, d_qr_.get(),
lda_, dtype, d_tau_.get(), dtype, solver_ctx_.scratch_workspace(), dev_ws,
host_ws > 0 ? solver_ctx_.h_workspace_.data() : nullptr, host_ws,
solver_ctx_.scratch_info()));
EIGEN_CUSOLVER_CHECK(
cusolverDnXgeqrf(solver_ctx_.cusolverHandle(), solver_ctx_.params_.p, fm, fn, dtype, d_qr_.get(), lda_, dtype,
d_tau_.get(), dtype, solver_ctx_.scratch_workspace(), dev_ws,
host_ws > 0 ? solver_ctx_.h_workspace_.data() : nullptr, host_ws, solver_ctx_.scratch_info()));
solver_ctx_.enqueue_info_copy();
}
@@ -267,13 +267,13 @@ class QR {
int lwork = 0;
EIGEN_CUSOLVER_CHECK(internal::cusolverDnXormqr_bufferSize(
solver_ctx_.cusolver_, CUBLAS_SIDE_LEFT, op, im, in, ik, static_cast<const Scalar*>(d_qr_.get()), ilda,
solver_ctx_.cusolverHandle(), CUBLAS_SIDE_LEFT, op, im, in, ik, static_cast<const Scalar*>(d_qr_.get()), ilda,
static_cast<const Scalar*>(d_tau_.get()), static_cast<const Scalar*>(d_B), ildb, &lwork));
solver_ctx_.ensure_scratch(static_cast<size_t>(lwork) * sizeof(Scalar));
EIGEN_CUSOLVER_CHECK(internal::cusolverDnXormqr(
solver_ctx_.cusolver_, CUBLAS_SIDE_LEFT, op, im, in, ik, static_cast<const Scalar*>(d_qr_.get()), ilda,
solver_ctx_.cusolverHandle(), CUBLAS_SIDE_LEFT, op, im, in, ik, static_cast<const Scalar*>(d_qr_.get()), ilda,
static_cast<const Scalar*>(d_tau_.get()), static_cast<Scalar*>(d_B), ildb,
static_cast<Scalar*>(solver_ctx_.scratch_workspace()), lwork, solver_ctx_.scratch_info()));
}
@@ -289,7 +289,7 @@ class QR {
internal::DeviceBuffer d_B(b_bytes);
EIGEN_CUDA_RUNTIME_CHECK(
cudaMemcpyAsync(d_B.get(), rhs.data(), b_bytes, cudaMemcpyHostToDevice, solver_ctx_.stream_));
cudaMemcpyAsync(d_B.get(), rhs.data(), b_bytes, cudaMemcpyHostToDevice, solver_ctx_.stream()));
apply_QH(d_B.get(), m_, nrhs);
trsm_R(d_B.get(), m_, nrhs, /*op=*/CUBLAS_OP_N);
@@ -298,14 +298,14 @@ class QR {
if (m_ == n_) {
EIGEN_CUDA_RUNTIME_CHECK(cudaMemcpyAsync(X.data(), d_B.get(),
static_cast<size_t>(n_) * static_cast<size_t>(nrhs) * sizeof(Scalar),
cudaMemcpyDeviceToHost, solver_ctx_.stream_));
cudaMemcpyDeviceToHost, solver_ctx_.stream()));
} else {
EIGEN_CUDA_RUNTIME_CHECK(cudaMemcpy2DAsync(X.data(), static_cast<size_t>(n_) * sizeof(Scalar), d_B.get(),
static_cast<size_t>(m_) * sizeof(Scalar),
static_cast<size_t>(n_) * sizeof(Scalar), static_cast<size_t>(nrhs),
cudaMemcpyDeviceToHost, solver_ctx_.stream_));
cudaMemcpyDeviceToHost, solver_ctx_.stream()));
}
EIGEN_CUDA_RUNTIME_CHECK(cudaStreamSynchronize(solver_ctx_.stream_));
EIGEN_CUDA_RUNTIME_CHECK(cudaStreamSynchronize(solver_ctx_.stream()));
return X;
}
@@ -315,7 +315,7 @@ class QR {
internal::DeviceBuffer d_work(b_bytes);
EIGEN_CUDA_RUNTIME_CHECK(
cudaMemcpyAsync(d_work.get(), d_B.data(), b_bytes, cudaMemcpyDeviceToDevice, solver_ctx_.stream_));
cudaMemcpyAsync(d_work.get(), d_B.data(), b_bytes, cudaMemcpyDeviceToDevice, solver_ctx_.stream()));
apply_QH(d_work.get(), m_, nrhs);
trsm_R(d_work.get(), m_, nrhs, /*op=*/CUBLAS_OP_N);
@@ -323,15 +323,15 @@ class QR {
if (m_ == n_) {
DeviceMatrix<Scalar> result =
DeviceMatrix<Scalar>::adopt(static_cast<Scalar*>(d_work.release()), n_, static_cast<Index>(nrhs));
result.recordReady(solver_ctx_.stream_);
result.recordReady(solver_ctx_.stream());
return result;
}
DeviceMatrix<Scalar> result(n_, static_cast<Index>(nrhs));
EIGEN_CUDA_RUNTIME_CHECK(cudaMemcpy2DAsync(result.data(), static_cast<size_t>(n_) * sizeof(Scalar), d_work.get(),
static_cast<size_t>(m_) * sizeof(Scalar),
static_cast<size_t>(n_) * sizeof(Scalar), static_cast<size_t>(nrhs),
cudaMemcpyDeviceToDevice, solver_ctx_.stream_));
result.recordReady(solver_ctx_.stream_);
cudaMemcpyDeviceToDevice, solver_ctx_.stream()));
result.recordReady(solver_ctx_.stream());
return result;
}
@@ -345,14 +345,14 @@ class QR {
internal::DeviceBuffer d_X(x_bytes);
// Zero the full n × nrhs buffer; B will overwrite the top m × nrhs block.
EIGEN_CUDA_RUNTIME_CHECK(cudaMemsetAsync(d_X.get(), 0, x_bytes, solver_ctx_.stream_));
EIGEN_CUDA_RUNTIME_CHECK(cudaMemsetAsync(d_X.get(), 0, x_bytes, solver_ctx_.stream()));
// 2D copy: B (m × nrhs, leading dim m) into top of d_X (leading dim n).
if (m_ > 0 && nrhs > 0) {
EIGEN_CUDA_RUNTIME_CHECK(cudaMemcpy2DAsync(d_X.get(), static_cast<size_t>(n_) * sizeof(Scalar), rhs.data(),
static_cast<size_t>(m_) * sizeof(Scalar),
static_cast<size_t>(m_) * sizeof(Scalar), static_cast<size_t>(nrhs),
cudaMemcpyHostToDevice, solver_ctx_.stream_));
cudaMemcpyHostToDevice, solver_ctx_.stream()));
}
trsm_R(d_X.get(), n_, nrhs, trsm_op_conj_trans());
@@ -360,8 +360,8 @@ class QR {
PlainMatrix X(n_, nrhs);
EIGEN_CUDA_RUNTIME_CHECK(
cudaMemcpyAsync(X.data(), d_X.get(), x_bytes, cudaMemcpyDeviceToHost, solver_ctx_.stream_));
EIGEN_CUDA_RUNTIME_CHECK(cudaStreamSynchronize(solver_ctx_.stream_));
cudaMemcpyAsync(X.data(), d_X.get(), x_bytes, cudaMemcpyDeviceToHost, solver_ctx_.stream()));
EIGEN_CUDA_RUNTIME_CHECK(cudaStreamSynchronize(solver_ctx_.stream()));
return X;
}
@@ -370,13 +370,13 @@ class QR {
const size_t x_bytes = static_cast<size_t>(n_) * static_cast<size_t>(nrhs) * sizeof(Scalar);
internal::DeviceBuffer d_X(x_bytes);
EIGEN_CUDA_RUNTIME_CHECK(cudaMemsetAsync(d_X.get(), 0, x_bytes, solver_ctx_.stream_));
EIGEN_CUDA_RUNTIME_CHECK(cudaMemsetAsync(d_X.get(), 0, x_bytes, solver_ctx_.stream()));
if (m_ > 0 && nrhs > 0) {
EIGEN_CUDA_RUNTIME_CHECK(cudaMemcpy2DAsync(d_X.get(), static_cast<size_t>(n_) * sizeof(Scalar), d_B.data(),
static_cast<size_t>(m_) * sizeof(Scalar),
static_cast<size_t>(m_) * sizeof(Scalar), static_cast<size_t>(nrhs),
cudaMemcpyDeviceToDevice, solver_ctx_.stream_));
cudaMemcpyDeviceToDevice, solver_ctx_.stream()));
}
trsm_R(d_X.get(), n_, nrhs, trsm_op_conj_trans());
@@ -384,7 +384,7 @@ class QR {
DeviceMatrix<Scalar> result =
DeviceMatrix<Scalar>::adopt(static_cast<Scalar*>(d_X.release()), n_, static_cast<Index>(nrhs));
result.recordReady(solver_ctx_.stream_);
result.recordReady(solver_ctx_.stream());
return result;
}
@@ -395,8 +395,8 @@ class QR {
void trsm_R(void* d_B, int64_t ldb, int64_t nrhs, cublasOperation_t op) const {
Scalar alpha(1);
EIGEN_CUBLAS_CHECK(internal::cublasXtrsm(
solver_ctx_.cublas_, CUBLAS_SIDE_LEFT, CUBLAS_FILL_MODE_UPPER, op, CUBLAS_DIAG_NON_UNIT, k(), nrhs, &alpha,
static_cast<const Scalar*>(d_qr_.get()), lda_, static_cast<Scalar*>(d_B), ldb));
solver_ctx_.cublasHandle(), CUBLAS_SIDE_LEFT, CUBLAS_FILL_MODE_UPPER, op, CUBLAS_DIAG_NON_UNIT, k(), nrhs,
&alpha, static_cast<const Scalar*>(d_qr_.get()), lda_, static_cast<Scalar*>(d_B), ldb));
}
};
} // namespace gpu
unsupported/Eigen/src/GPU/GpuSVD.h
+43 -43
View File
@@ -127,7 +127,7 @@ class SVD {
// uploaded matrix is consumed in place by gesvd, so no second device copy.
// The wide-matrix transpose runs on the GPU (via cublasXgeam) inside the
// device-input path; no host transpose.
return compute(DeviceMatrix<Scalar>::fromHost(A.derived(), solver_ctx_.stream_), options);
return compute(DeviceMatrix<Scalar>::fromHost(A.derived(), solver_ctx_.stream()), options);
}
SVD& compute(const DeviceMatrix<Scalar>& d_A, unsigned int options = ComputeThinU | ComputeThinV) {
@@ -139,7 +139,7 @@ class SVD {
const size_t mat_bytes = static_cast<size_t>(lda_) * static_cast<size_t>(n_) * sizeof(Scalar);
d_A_ = internal::DeviceBuffer(mat_bytes);
EIGEN_CUDA_RUNTIME_CHECK(
cudaMemcpyAsync(d_A_.get(), d_A.data(), mat_bytes, cudaMemcpyDeviceToDevice, solver_ctx_.stream_));
cudaMemcpyAsync(d_A_.get(), d_A.data(), mat_bytes, cudaMemcpyDeviceToDevice, solver_ctx_.stream()));
}
factorize();
@@ -244,7 +244,7 @@ class SVD {
eigen_assert(solver_ctx_.info() == Success);
const Index k = (std::min)(m_, n_);
auto v = DeviceMatrix<RealScalar>::view(static_cast<RealScalar*>(d_S_.get()), k, 1);
v.recordReady(solver_ctx_.stream_);
v.recordReady(solver_ctx_.stream());
return v;
}
@@ -259,7 +259,7 @@ class SVD {
if (!transposed_) {
const Index ucols = (options_ & ComputeFullU) ? m_ : k;
auto v = DeviceMatrix<Scalar>::view(static_cast<Scalar*>(d_U_.get()), m_, ucols);
v.recordReady(solver_ctx_.stream_);
v.recordReady(solver_ctx_.stream());
return v;
}
// transposed: U_orig = VT_stored^H -> conjugate-transpose via cublasXgeam.
@@ -267,10 +267,10 @@ class SVD {
DeviceMatrix<Scalar> result(n_, vtrows_stored);
if (n_ > 0 && vtrows_stored > 0) {
Scalar alpha_one(1), beta_zero(0);
EIGEN_CUBLAS_CHECK(internal::cublasXgeam(solver_ctx_.cublas_, CUBLAS_OP_C, CUBLAS_OP_N, n_, vtrows_stored,
EIGEN_CUBLAS_CHECK(internal::cublasXgeam(solver_ctx_.cublasHandle(), CUBLAS_OP_C, CUBLAS_OP_N, n_, vtrows_stored,
&alpha_one, static_cast<const Scalar*>(d_VT_.get()), vtrows_stored,
&beta_zero, static_cast<const Scalar*>(nullptr), n_, result.data(), n_));
result.recordReady(solver_ctx_.stream_);
result.recordReady(solver_ctx_.stream());
}
return result;
}
@@ -286,7 +286,7 @@ class SVD {
if (!transposed_) {
const Index vtrows = (options_ & ComputeFullV) ? n_ : k;
auto v = DeviceMatrix<Scalar>::view(static_cast<Scalar*>(d_VT_.get()), vtrows, n_);
v.recordReady(solver_ctx_.stream_);
v.recordReady(solver_ctx_.stream());
return v;
}
// transposed: VT_orig = U_stored^H.
@@ -294,10 +294,10 @@ class SVD {
DeviceMatrix<Scalar> result(ucols, m_);
if (ucols > 0 && m_ > 0) {
Scalar alpha_one(1), beta_zero(0);
EIGEN_CUBLAS_CHECK(internal::cublasXgeam(solver_ctx_.cublas_, CUBLAS_OP_C, CUBLAS_OP_N, ucols, m_, &alpha_one,
static_cast<const Scalar*>(d_U_.get()), m_, &beta_zero,
EIGEN_CUBLAS_CHECK(internal::cublasXgeam(solver_ctx_.cublasHandle(), CUBLAS_OP_C, CUBLAS_OP_N, ucols, m_,
&alpha_one, static_cast<const Scalar*>(d_U_.get()), m_, &beta_zero,
static_cast<const Scalar*>(nullptr), ucols, result.data(), ucols));
result.recordReady(solver_ctx_.stream_);
result.recordReady(solver_ctx_.stream());
}
return result;
}
@@ -353,7 +353,7 @@ class SVD {
return solve_device_impl(d_B, (std::min)(m_, n_), lambda);
}
cudaStream_t stream() const { return solver_ctx_.stream_; }
cudaStream_t stream() const { return solver_ctx_.stream(); }
private:
mutable internal::GpuSolverContext solver_ctx_;
@@ -395,7 +395,7 @@ class SVD {
} else {
lda_ = static_cast<int64_t>(d_A.rows());
}
d_A.waitReady(solver_ctx_.stream_);
d_A.waitReady(solver_ctx_.stream());
return true;
}
@@ -405,7 +405,7 @@ class SVD {
d_A_ = internal::DeviceBuffer(mat_bytes);
// geam: C(m×n) = alpha * op(A) + beta * op(B). beta=0, B=nullptr.
Scalar alpha_one(1), beta_zero(0);
EIGEN_CUBLAS_CHECK(internal::cublasXgeam(solver_ctx_.cublas_, CUBLAS_OP_C, CUBLAS_OP_N, m_, n_, &alpha_one,
EIGEN_CUBLAS_CHECK(internal::cublasXgeam(solver_ctx_.cublasHandle(), CUBLAS_OP_C, CUBLAS_OP_N, m_, n_, &alpha_one,
d_A.data(), d_A.rows(), &beta_zero, static_cast<const Scalar*>(nullptr),
m_, static_cast<Scalar*>(d_A_.get()), m_));
}
@@ -458,18 +458,18 @@ class SVD {
eigen_assert(m_ >= n_ && "Internal error: m_ < n_ should have been handled by transpose in compute()");
size_t dev_ws = 0, host_ws = 0;
EIGEN_CUSOLVER_CHECK(cusolverDnXgesvd_bufferSize(
solver_ctx_.cusolver_, solver_ctx_.params_.p, jobu(int_opts), jobvt(int_opts), m_, n_, dtype, d_A_.get(), lda_,
rtype, d_S_.get(), dtype, ucols > 0 ? d_U_.get() : nullptr, ldu, dtype, vtrows > 0 ? d_VT_.get() : nullptr,
ldvt, dtype, &dev_ws, &host_ws));
solver_ctx_.cusolverHandle(), solver_ctx_.params_.p, jobu(int_opts), jobvt(int_opts), m_, n_, dtype, d_A_.get(),
lda_, rtype, d_S_.get(), dtype, ucols > 0 ? d_U_.get() : nullptr, ldu, dtype,
vtrows > 0 ? d_VT_.get() : nullptr, ldvt, dtype, &dev_ws, &host_ws));
solver_ctx_.ensure_scratch(dev_ws);
solver_ctx_.h_workspace_.resize(host_ws);
EIGEN_CUSOLVER_CHECK(
cusolverDnXgesvd(solver_ctx_.cusolver_, solver_ctx_.params_.p, jobu(int_opts), jobvt(int_opts), m_, n_, dtype,
d_A_.get(), lda_, rtype, d_S_.get(), dtype, ucols > 0 ? d_U_.get() : nullptr, ldu, dtype,
vtrows > 0 ? d_VT_.get() : nullptr, ldvt, dtype, solver_ctx_.scratch_workspace(), dev_ws,
host_ws > 0 ? solver_ctx_.h_workspace_.data() : nullptr, host_ws, solver_ctx_.scratch_info()));
EIGEN_CUSOLVER_CHECK(cusolverDnXgesvd(
solver_ctx_.cusolverHandle(), solver_ctx_.params_.p, jobu(int_opts), jobvt(int_opts), m_, n_, dtype, d_A_.get(),
lda_, rtype, d_S_.get(), dtype, ucols > 0 ? d_U_.get() : nullptr, ldu, dtype,
vtrows > 0 ? d_VT_.get() : nullptr, ldvt, dtype, solver_ctx_.scratch_workspace(), dev_ws,
host_ws > 0 ? solver_ctx_.h_workspace_.data() : nullptr, host_ws, solver_ctx_.scratch_info()));
solver_ctx_.enqueue_info_copy();
@@ -496,8 +496,8 @@ class SVD {
const Index k = (std::min)(m_, n_);
RealVector S(k);
EIGEN_CUDA_RUNTIME_CHECK(cudaMemcpyAsync(S.data(), d_S_.get(), static_cast<size_t>(k) * sizeof(RealScalar),
cudaMemcpyDeviceToHost, solver_ctx_.stream_));
EIGEN_CUDA_RUNTIME_CHECK(cudaStreamSynchronize(solver_ctx_.stream_));
cudaMemcpyDeviceToHost, solver_ctx_.stream()));
EIGEN_CUDA_RUNTIME_CHECK(cudaStreamSynchronize(solver_ctx_.stream()));
const RealScalar drop_threshold = S(0) * RealScalar(k) * NumTraits<RealScalar>::epsilon();
auto S_head = S.head(kk).array();
@@ -511,7 +511,7 @@ class SVD {
const size_t d_bytes = static_cast<size_t>(kk) * sizeof(Scalar);
internal::ensure_sized(d_D_, d_bytes);
EIGEN_CUDA_RUNTIME_CHECK(
cudaMemcpyAsync(d_D_.get(), D.data(), d_bytes, cudaMemcpyHostToDevice, solver_ctx_.stream_));
cudaMemcpyAsync(d_D_.get(), D.data(), d_bytes, cudaMemcpyHostToDevice, solver_ctx_.stream()));
cached_diag_kk_ = kk;
cached_diag_lambda_ = lambda;
diag_valid_ = true;
@@ -534,34 +534,34 @@ class SVD {
internal::DeviceBuffer d_tmp(static_cast<size_t>(kk) * static_cast<size_t>(nrhs) * sizeof(Scalar));
auto* tmp_dev = static_cast<Scalar*>(d_tmp.get());
if (!transposed_) {
internal::cublaslt_gemm(solver_ctx_.cublasLtHandle(), solver_ctx_.cublas_, CUBLAS_OP_C, CUBLAS_OP_N, kk, nrhs, m_,
&scalars[0], U_dev, m_, B_dev, m_orig, &scalars[1], tmp_dev, kk,
internal::cublaslt_gemm(solver_ctx_.cublasLtHandle(), solver_ctx_.cublasHandle(), CUBLAS_OP_C, CUBLAS_OP_N, kk,
nrhs, m_, &scalars[0], U_dev, m_, B_dev, m_orig, &scalars[1], tmp_dev, kk,
solver_ctx_.gemmWorkspace(), solver_ctx_.gemmPlanCache(),
solver_ctx_.cublasLtMaxWorkspaceBytes(), solver_ctx_.stream_);
solver_ctx_.cublasLtMaxWorkspaceBytes(), solver_ctx_.stream());
} else {
const Index vtrows_stored = (swap_uv_options(options_) & ComputeFullV) ? n_ : k;
internal::cublaslt_gemm(solver_ctx_.cublasLtHandle(), solver_ctx_.cublas_, CUBLAS_OP_N, CUBLAS_OP_N, kk, nrhs,
m_orig, &scalars[0], VT_dev, vtrows_stored, B_dev, m_orig, &scalars[1], tmp_dev, kk,
internal::cublaslt_gemm(solver_ctx_.cublasLtHandle(), solver_ctx_.cublasHandle(), CUBLAS_OP_N, CUBLAS_OP_N, kk,
nrhs, m_orig, &scalars[0], VT_dev, vtrows_stored, B_dev, m_orig, &scalars[1], tmp_dev, kk,
solver_ctx_.gemmWorkspace(), solver_ctx_.gemmPlanCache(),
solver_ctx_.cublasLtMaxWorkspaceBytes(), solver_ctx_.stream_);
solver_ctx_.cublasLtMaxWorkspaceBytes(), solver_ctx_.stream());
}
// Step 2: tmp = diag(D) * tmp on device via cublasXdgmm.
EIGEN_CUBLAS_CHECK(internal::cublasXdgmm(solver_ctx_.cublas_, CUBLAS_SIDE_LEFT, kk, nrhs, tmp_dev, kk,
EIGEN_CUBLAS_CHECK(internal::cublasXdgmm(solver_ctx_.cublasHandle(), CUBLAS_SIDE_LEFT, kk, nrhs, tmp_dev, kk,
static_cast<const Scalar*>(d_D_.get()), 1, tmp_dev, kk));
// Step 3: X = V_orig * tmp (n_orig × nrhs).
if (!transposed_) {
const Index vtrows = (options_ & ComputeFullV) ? n_ : k;
internal::cublaslt_gemm(solver_ctx_.cublasLtHandle(), solver_ctx_.cublas_, CUBLAS_OP_C, CUBLAS_OP_N, n_orig, nrhs,
kk, &scalars[0], VT_dev, vtrows, tmp_dev, kk, &scalars[1], X_dev, n_orig,
internal::cublaslt_gemm(solver_ctx_.cublasLtHandle(), solver_ctx_.cublasHandle(), CUBLAS_OP_C, CUBLAS_OP_N,
n_orig, nrhs, kk, &scalars[0], VT_dev, vtrows, tmp_dev, kk, &scalars[1], X_dev, n_orig,
solver_ctx_.gemmWorkspace(), solver_ctx_.gemmPlanCache(),
solver_ctx_.cublasLtMaxWorkspaceBytes(), solver_ctx_.stream_);
solver_ctx_.cublasLtMaxWorkspaceBytes(), solver_ctx_.stream());
} else {
internal::cublaslt_gemm(solver_ctx_.cublasLtHandle(), solver_ctx_.cublas_, CUBLAS_OP_N, CUBLAS_OP_N, n_orig, nrhs,
kk, &scalars[0], U_dev, m_, tmp_dev, kk, &scalars[1], X_dev, n_orig,
internal::cublaslt_gemm(solver_ctx_.cublasLtHandle(), solver_ctx_.cublasHandle(), CUBLAS_OP_N, CUBLAS_OP_N,
n_orig, nrhs, kk, &scalars[0], U_dev, m_, tmp_dev, kk, &scalars[1], X_dev, n_orig,
solver_ctx_.gemmWorkspace(), solver_ctx_.gemmPlanCache(),
solver_ctx_.cublasLtMaxWorkspaceBytes(), solver_ctx_.stream_);
solver_ctx_.cublasLtMaxWorkspaceBytes(), solver_ctx_.stream());
}
}
@@ -591,7 +591,7 @@ class SVD {
internal::DeviceBuffer d_B(static_cast<size_t>(m_orig) * static_cast<size_t>(nrhs) * sizeof(Scalar));
EIGEN_CUDA_RUNTIME_CHECK(cudaMemcpyAsync(d_B.get(), rhs.data(),
static_cast<size_t>(m_orig) * static_cast<size_t>(nrhs) * sizeof(Scalar),
cudaMemcpyHostToDevice, solver_ctx_.stream_));
cudaMemcpyHostToDevice, solver_ctx_.stream()));
build_diag(kk, lambda);
PlainMatrix X(n_orig, nrhs);
@@ -600,8 +600,8 @@ class SVD {
EIGEN_CUDA_RUNTIME_CHECK(cudaMemcpyAsync(X.data(), d_X.get(),
static_cast<size_t>(n_orig) * static_cast<size_t>(nrhs) * sizeof(Scalar),
cudaMemcpyDeviceToHost, solver_ctx_.stream_));
EIGEN_CUDA_RUNTIME_CHECK(cudaStreamSynchronize(solver_ctx_.stream_));
cudaMemcpyDeviceToHost, solver_ctx_.stream()));
EIGEN_CUDA_RUNTIME_CHECK(cudaStreamSynchronize(solver_ctx_.stream()));
return X;
}
@@ -621,16 +621,16 @@ class SVD {
if (kk == 0 || nrhs == 0 || n_orig == 0) {
DeviceMatrix<Scalar> X(n_orig, nrhs);
X.setZero(solver_ctx_.stream_);
X.setZero(solver_ctx_.stream());
return X;
}
d_B.waitReady(solver_ctx_.stream_);
d_B.waitReady(solver_ctx_.stream());
build_diag(kk, lambda);
DeviceMatrix<Scalar> X(n_orig, nrhs);
apply_pinv(d_B.data(), kk, nrhs, X.data());
X.recordReady(solver_ctx_.stream_);
X.recordReady(solver_ctx_.stream());
return X;
}
};
unsupported/Eigen/src/GPU/GpuSolverContext.h
+43 -83
View File
@@ -27,10 +27,11 @@ namespace gpu {
namespace internal {
struct GpuSolverContext {
cudaStream_t stream_ = nullptr;
cusolverDnHandle_t cusolver_ = nullptr;
cublasHandle_t cublas_ = nullptr;
cublasLtHandle_t cublas_lt_ = nullptr; // lazy: created on first GEMM-via-cublasLt call (standalone mode only)
Context* bound_ctx_ = nullptr;
UniqueStream stream_;
UniqueCusolverHandle cusolver_;
UniqueCublasHandle cublas_;
UniqueCublasLtHandle cublas_lt_; // lazy: created on first GEMM-via-cublasLt call (standalone mode only)
CusolverParams params_;
DeviceBuffer d_scratch_;
std::vector<char> h_workspace_;
@@ -42,21 +43,26 @@ struct GpuSolverContext {
ComputationInfo info_ = InvalidInput;
PinnedHostBuffer pinned_info_{sizeof(int)}; // pinned host memory for async D2H of info word
bool info_synced_ = true;
// Non-null when this solver context borrows from a gpu::Context: stream and
// cuSOLVER/cuBLAS handles are the Context's, and the cuBLASLt handle, GEMM
// plan cache, and GEMM workspace are shared with it rather than duplicated.
// Null in standalone mode, where all of the above are owned.
Context* bound_ctx_ = nullptr;
int& info_word() { return *static_cast<int*>(pinned_info_.get()); }
int info_word() const { return *static_cast<const int*>(pinned_info_.get()); }
cudaStream_t stream() const { return stream_.get(); }
cusolverDnHandle_t cusolverHandle() const { return cusolver_.get(); }
cublasHandle_t cublasHandle() const { return cublas_.get(); }
GpuSolverContext() {
EIGEN_CUDA_RUNTIME_CHECK(cudaStreamCreate(&stream_));
EIGEN_CUSOLVER_CHECK(cusolverDnCreate(&cusolver_));
EIGEN_CUSOLVER_CHECK(cusolverDnSetStream(cusolver_, stream_));
EIGEN_CUBLAS_CHECK(cublasCreate(&cublas_));
EIGEN_CUBLAS_CHECK(cublasSetStream(cublas_, stream_));
cudaStream_t s = nullptr;
EIGEN_CUDA_RUNTIME_CHECK(cudaStreamCreate(&s));
stream_ = UniqueStream(s);
cusolverDnHandle_t solver = nullptr;
EIGEN_CUSOLVER_CHECK(cusolverDnCreate(&solver));
cusolver_ = UniqueCusolverHandle(solver);
EIGEN_CUSOLVER_CHECK(cusolverDnSetStream(solver, s));
cublasHandle_t blas = nullptr;
EIGEN_CUBLAS_CHECK(cublasCreate(&blas));
cublas_ = UniqueCublasHandle(blas);
EIGEN_CUBLAS_CHECK(cublasSetStream(blas, s));
ensure_scratch(0);
}
@@ -66,68 +72,27 @@ struct GpuSolverContext {
* own). The cuBLASLt handle, GEMM plan cache, and GEMM workspace are shared
* with the Context as well. The Context must outlive this solver context. */
explicit GpuSolverContext(Context& ctx)
: stream_(ctx.stream()), cusolver_(ctx.cusolverHandle()), cublas_(ctx.cublasHandle()), bound_ctx_(&ctx) {
: bound_ctx_(&ctx),
stream_(ctx.stream(), CudaStreamDeleter{/*owns=*/false}),
cusolver_(ctx.cusolverHandle(), CusolverHandleDeleter{/*owns=*/false}),
cublas_(ctx.cublasHandle(), CublasHandleDeleter{/*owns=*/false}) {
ensure_scratch(0);
}
~GpuSolverContext() {
// Ignore errors here: dtors are noexcept, and EIGEN_CU{BLAS,SOLVER,DA_RUNTIME}_CHECK
// are eigen_assert-based — firing one from a noexcept dtor terminates the program.
// The trailing free of the device buffers (via DeviceBuffer::~DeviceBuffer) is
// stream-ordered (or synchronous on the cudaMalloc fallback), so it waits
// for any in-flight kernel touching the buffer.
// Destroy plan cache before its cublasLt handle (entries hold descriptors).
gemm_plan_cache_.clear();
if (cublas_lt_) (void)cublasLtDestroy(cublas_lt_);
if (!bound_ctx_) {
if (cublas_) (void)cublasDestroy(cublas_);
if (cusolver_) (void)cusolverDnDestroy(cusolver_);
if (stream_) (void)cudaStreamDestroy(stream_);
}
}
GpuSolverContext(GpuSolverContext&& o) noexcept
: stream_(o.stream_),
cusolver_(o.cusolver_),
cublas_(o.cublas_),
cublas_lt_(o.cublas_lt_),
params_(std::move(o.params_)),
d_scratch_(std::move(o.d_scratch_)),
h_workspace_(std::move(o.h_workspace_)),
gemm_workspace_(std::move(o.gemm_workspace_)),
gemm_plan_cache_(std::move(o.gemm_plan_cache_)),
cublaslt_max_workspace_bytes_(o.cublaslt_max_workspace_bytes_),
info_(o.info_),
pinned_info_(std::move(o.pinned_info_)),
info_synced_(o.info_synced_),
bound_ctx_(o.bound_ctx_) {
o.stream_ = nullptr;
o.cusolver_ = nullptr;
o.cublas_ = nullptr;
o.cublas_lt_ = nullptr;
o.info_ = InvalidInput;
o.info_synced_ = true;
o.bound_ctx_ = nullptr;
}
~GpuSolverContext() = default;
GpuSolverContext(GpuSolverContext&& o) noexcept = default;
GpuSolverContext& operator=(GpuSolverContext&& o) noexcept {
if (this != &o) {
// Mirror the dtor: noexcept context, can't propagate. Drain the old stream
// first so the upcoming move of d_scratch_ doesn't free buffers an in-flight
// kernel is still touching; then swallow destroy errors (the EIGEN_CU*_CHECK
// macros are eigen_assert-based and would terminate from a noexcept body).
if (stream_) (void)cudaStreamSynchronize(stream_);
// A pending info copy may still write pinned_info_, whose cudaFreeHost deleter is not stream-ordered.
if (!info_synced_ && pinned_info_) (void)cudaStreamSynchronize(stream());
// Release plan-cache descriptors before the moves below replace the cuBLASLt handle they were built with.
gemm_plan_cache_.clear();
if (cublas_lt_) (void)cublasLtDestroy(cublas_lt_);
if (!bound_ctx_) {
if (cublas_) (void)cublasDestroy(cublas_);
if (cusolver_) (void)cusolverDnDestroy(cusolver_);
if (stream_) (void)cudaStreamDestroy(stream_);
}
stream_ = o.stream_;
cusolver_ = o.cusolver_;
cublas_ = o.cublas_;
cublas_lt_ = o.cublas_lt_;
bound_ctx_ = o.bound_ctx_;
stream_ = std::move(o.stream_);
cusolver_ = std::move(o.cusolver_);
cublas_ = std::move(o.cublas_);
cublas_lt_ = std::move(o.cublas_lt_);
params_ = std::move(o.params_);
d_scratch_ = std::move(o.d_scratch_);
h_workspace_ = std::move(o.h_workspace_);
@@ -137,13 +102,6 @@ struct GpuSolverContext {
info_ = o.info_;
pinned_info_ = std::move(o.pinned_info_);
info_synced_ = o.info_synced_;
bound_ctx_ = o.bound_ctx_;
o.stream_ = nullptr;
o.cusolver_ = nullptr;
o.cublas_ = nullptr;
o.cublas_lt_ = nullptr;
o.info_ = InvalidInput;
o.info_synced_ = true;
o.bound_ctx_ = nullptr;
}
return *this;
@@ -154,15 +112,17 @@ struct GpuSolverContext {
cublasLtHandle_t cublasLtHandle() {
if (bound_ctx_) return bound_ctx_->cublasLtHandle();
if (!cublas_lt_) {
EIGEN_CUBLAS_CHECK(cublasLtCreate(&cublas_lt_));
cublasLtHandle_t h = nullptr;
EIGEN_CUBLAS_CHECK(cublasLtCreate(&h));
cublas_lt_ = UniqueCublasLtHandle(h);
}
return cublas_lt_;
return cublas_lt_.get();
}
/** GEMM plan cache / workspace / workspace ceiling for cublaslt_gemm —
* shared with the bound Context when borrowing, owned otherwise. */
CublasLtPlanCache* gemmPlanCache() { return bound_ctx_ ? bound_ctx_->gemmPlanCache() : &gemm_plan_cache_; }
DeviceBuffer* gemmWorkspace() { return bound_ctx_ ? bound_ctx_->gemmWorkspace() : &gemm_workspace_; }
CublasLtPlanCache& gemmPlanCache() { return bound_ctx_ ? bound_ctx_->gemmPlanCache() : gemm_plan_cache_; }
DeviceBuffer& gemmWorkspace() { return bound_ctx_ ? bound_ctx_->gemmWorkspace() : gemm_workspace_; }
std::size_t cublasLtMaxWorkspaceBytes() const {
return bound_ctx_ ? bound_ctx_->cublasLtMaxWorkspaceBytes() : cublaslt_max_workspace_bytes_;
}
@@ -186,7 +146,7 @@ struct GpuSolverContext {
void ensure_scratch(size_t workspace_bytes) {
size_t needed = scratchBytesFor(workspace_bytes);
if (needed > d_scratch_.size()) {
if (d_scratch_) EIGEN_CUDA_RUNTIME_CHECK(cudaStreamSynchronize(stream_));
if (d_scratch_) EIGEN_CUDA_RUNTIME_CHECK(cudaStreamSynchronize(stream()));
d_scratch_ = DeviceBuffer(needed);
}
}
@@ -220,13 +180,13 @@ struct GpuSolverContext {
// into pinned host memory. Read later by the lazy sync_info().
void enqueue_info_copy() {
EIGEN_CUDA_RUNTIME_CHECK(
cudaMemcpyAsync(&info_word(), scratch_info(), sizeof(int), cudaMemcpyDeviceToHost, stream_));
cudaMemcpyAsync(&info_word(), scratch_info(), sizeof(int), cudaMemcpyDeviceToHost, stream()));
}
// Synchronize the stream and interpret the info word; no-op once synced.
void sync_info() {
if (!info_synced_) {
EIGEN_CUDA_RUNTIME_CHECK(cudaStreamSynchronize(stream_));
EIGEN_CUDA_RUNTIME_CHECK(cudaStreamSynchronize(stream()));
info_ = (info_word() == 0) ? Success : NumericalIssue;
info_synced_ = true;
}
unsupported/Eigen/src/GPU/GpuSupport.h
+10
View File
@@ -22,6 +22,7 @@
#include <limits>
#include <memory>
#include <type_traits>
namespace Eigen {
namespace gpu {
@@ -124,6 +125,15 @@ struct CudaFreeHostDeleter {
}
};
// RAII CUDA stream; the ownership flag supports borrowed, caller-owned streams.
struct CudaStreamDeleter {
bool owns = true;
void operator()(cudaStream_t s) const noexcept {
if (owns && s) (void)cudaStreamDestroy(s);
}
};
using UniqueStream = std::unique_ptr<std::remove_pointer_t<cudaStream_t>, CudaStreamDeleter>;
// Recycles allocations up to kSmallBufferThreshold bytes (e.g. DeviceScalar) to
// avoid cudaMalloc/cudaFree overhead. Larger allocations bypass the pool.
template <size_t SmallBufferThreshold = 256, size_t MaxPoolSize = 64>
unsupported/Eigen/src/GPU/README.md
+3 -2
View File
@@ -712,8 +712,9 @@ cusolverDnHandle_t cusolverHandle() // Lazy: creates the
cublasLtHandle_t cublasLtHandle() // Lazy-initialized
cusparseHandle_t cusparseHandle() // Lazy-initialized
internal::DeviceBuffer* gemmWorkspace() // cublasLtMatmul scratch (lazy-grown per context)
internal::CublasLtPlanCache* gemmPlanCache() // shape-keyed plan cache (per context, ~8-entry LRU)
internal::DeviceBuffer& gemmWorkspace() // cublasLtMatmul scratch (lazy-grown per context)
internal::CublasLtPlanCache& gemmPlanCache() // shape-keyed plan cache (per context, ~8-entry LRU)
internal::OneShotSolverScratch& oneshotSolverScratch() // LLT/LU expression scratch (lazy-grown per context)
```
Non-copyable, non-movable (owns library handles). Translation units that
unsupported/test/GPU/cusolver_eigen.cpp
+11 -1
View File
@@ -267,12 +267,22 @@ void test_eigen_move(Index n) {
RealScalar tol = RealScalar(8) * static_cast<RealScalar>(n) * NumTraits<Scalar>::epsilon() * A.norm();
VERIFY((V * W.asDiagonal() * V.adjoint() - A).norm() < tol);
gpu::SelfAdjointEigenSolver<Scalar> assigned;
// Overwrite a Context-bound solver without querying its pending status
// first. Move assignment must retire its asynchronous info copy before
// releasing the pinned host destination.
Mat pending_random = Mat::Random(n, n);
Mat pending_input = pending_random + pending_random.adjoint();
gpu::Context ctx;
gpu::SelfAdjointEigenSolver<Scalar> assigned(ctx, pending_input);
assigned = std::move(moved);
VERIFY_IS_EQUAL(assigned.info(), Success);
V = assigned.eigenvectors();
W = assigned.eigenvalues();
VERIFY((V * W.asDiagonal() * V.adjoint() - A).norm() < tol);
// The borrowed Context remains usable after its old solver state is retired.
gpu::SelfAdjointEigenSolver<Scalar> context_solver(ctx, pending_input);
VERIFY_IS_EQUAL(context_solver.info(), Success);
}
// ---- Empty matrix -----------------------------------------------------------