GpuContext.h

// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2026 Rasmus Munk Larsen <rmlarsen@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
// SPDX-License-Identifier: MPL-2.0
 
// Unified GPU execution context: a CUDA stream plus the NVIDIA library handles
// used by gpu::DeviceMatrix operations.
 
#ifndef EIGEN_GPU_CONTEXT_H
#define EIGEN_GPU_CONTEXT_H
 
// IWYU pragma: private
#include "./InternalHeaderCheck.h"
 
#include "./DeviceScalarOps.h"
#include "./CuBlasSupport.h"
#include "./CuSolverSupport.h"
#include "./CuSparseSupport.h"
#include <vector>
 
namespace Eigen {
namespace gpu {
 
namespace internal {
 
// cuSOLVER writes the factorization/solve status words to device memory in
// every build, so d_info must exist even under EIGEN_NO_DEBUG. Only the
// pinned host mirror is debug-only: it feeds the oneshot_check_info assert,
// which release builds compile out.
constexpr size_t kOneShotInfoBytes = 2 * sizeof(int);
#ifdef EIGEN_NO_DEBUG
constexpr size_t kOneShotHostInfoBytes = 0;
#else
constexpr size_t kOneShotHostInfoBytes = kOneShotInfoBytes;
#endif
// Grow-only scratch shared by the one-shot solver expressions
// (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 the one-shot solve dispatches in DeviceDispatch.h.
struct OneShotSolverScratch {
  DeviceBuffer d_factor;
  DeviceBuffer d_ipiv;
  DeviceBuffer d_workspace;
  DeviceBuffer d_info{kOneShotInfoBytes};          // 2 ints: {factorization, solve}
  PinnedHostBuffer h_info{kOneShotHostInfoBytes};  // debug-build info check only
  std::vector<char> h_workspace;
};
 
inline void ensure_sized(DeviceBuffer& buf, size_t needed) {
  if (needed > buf.size()) {
    // Replacing an in-use buffer is safe: device_free is stream-ordered (or
    // fully synchronous on the cudaMalloc fallback path) and DeviceBufferPool
    // holds a released block back until the device has retired the work
    // enqueued before the release, so nothing reuses the old buffer early.
    buf = DeviceBuffer(needed);
  }
}
}  // namespace internal

class Context {
 public:
  Context() {
    cudaStream_t s = nullptr;
    EIGEN_CUDA_RUNTIME_CHECK(cudaStreamCreate(&s));
    stream_ = internal::UniqueStream(s);
    npp_stream_ctx_ = internal::make_npp_stream_ctx(stream_.get());
    init_cublas();
  }

  explicit Context(cudaStream_t stream) : stream_(stream, internal::CudaStreamDeleter{/*owns=*/false}) {
    npp_stream_ctx_ = internal::make_npp_stream_ctx(stream_.get());
    init_cublas();
  }
 
  ~Context() = default;
 
  Context(const Context&) = delete;
  Context& operator=(const Context&) = delete;
  Context(Context&&) = delete;
  Context& operator=(Context&&) = delete;

  static Context& threadLocal() {
    Context* override = tl_override_ptr();
    if (override) return *override;
    thread_local Context ctx;
    return ctx;
  }

  static void setThreadLocal(Context* ctx) { tl_override_ptr() = ctx; }
 
  cudaStream_t stream() const { return stream_.get(); }
  cublasHandle_t cublasHandle() const { return cublas_.get(); }

  const NppStreamContext& nppStreamContext() const { return npp_stream_ctx_; }

  cusolverDnHandle_t cusolverHandle() {
    if (!cusolver_) {
      cusolverDnHandle_t h = nullptr;
      EIGEN_CUSOLVER_CHECK(cusolverDnCreate(&h));
      cusolver_ = LazyCusolverHandle(h, &destroyCusolver);
      EIGEN_CUSOLVER_CHECK(cusolverDnSetStream(h, stream_.get()));
    }
    return cusolver_.get();
  }

  cublasLtHandle_t cublasLtHandle() {
    if (!cublas_lt_) {
      cublasLtHandle_t h = nullptr;
      EIGEN_CUBLAS_CHECK(cublasLtCreate(&h));
      cublas_lt_ = internal::UniqueCublasLtHandle(h);
    }
    return cublas_lt_.get();
  }

  internal::DeviceBuffer& gemmWorkspace() { return gemm_workspace_; }

  internal::CublasLtPlanCache& gemmPlanCache() { return gemm_plan_cache_; }

  internal::OneShotSolverScratch& oneshotSolverScratch() { return oneshot_solver_scratch_; }

  std::size_t cublasLtMaxWorkspaceBytes() const { return cublaslt_max_workspace_bytes_; }

  void setCublasLtMaxWorkspaceBytes(std::size_t bytes) { cublaslt_max_workspace_bytes_ = bytes; }

  cusparseHandle_t cusparseHandle() {
    if (!cusparse_) {
      cusparseHandle_t h = nullptr;
      EIGEN_CUSPARSE_CHECK(cusparseCreate(&h));
      cusparse_ = LazyCusparseHandle(h, &destroyCusparse);
      EIGEN_CUSPARSE_CHECK(cusparseSetStream(h, stream_.get()));
    }
    return cusparse_.get();
  }
 
 private:
  static cusolverStatus_t destroyCusolver(cusolverDnHandle_t h) { return cusolverDnDestroy(h); }
  static cusparseStatus_t destroyCusparse(cusparseHandle_t h) { return cusparseDestroy(h); }
 
  // 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_;
  NppStreamContext npp_stream_ctx_ = {};
  internal::UniqueCublasHandle cublas_;
  internal::DeviceBuffer cublas_workspace_;  // cublasSetWorkspace; freed before the handle
  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;
 
  static Context*& tl_override_ptr() {
    thread_local Context* ptr = nullptr;
    return ptr;
  }
 
  void init_cublas() {
    cublasHandle_t h = nullptr;
    EIGEN_CUBLAS_CHECK(cublasCreate(&h));
    cublas_ = internal::UniqueCublasHandle(h);
    EIGEN_CUBLAS_CHECK(cublasSetStream(h, stream_.get()));
    cublas_workspace_ = internal::DeviceBuffer(internal::kCublasWorkspaceBytes);
    EIGEN_CUBLAS_CHECK(cublasSetWorkspace(h, cublas_workspace_.get(), cublas_workspace_.size()));
  }
};
 
}  // namespace gpu
}  // namespace Eigen
 
#endif  // EIGEN_GPU_CONTEXT_H