// 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

// Tests for cuBLAS GEMM dispatch via DeviceMatrix expression syntax.
// Covers: d_C = d_A * d_B, adjoint, transpose, scaled, +=, .device(ctx).

#define EIGEN_USE_GPU
#include "main.h"
#include <contrib/Eigen/GPU>
#include <cstdint>
#include <new>

#include "./gpu_test_helpers.h"

using namespace Eigen;

// ---- Basic GEMM: C = A * B -------------------------------------------------

template <typename Scalar>
void test_gemm_basic(Index m, Index n, Index k) {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;
  using RealScalar = typename NumTraits<Scalar>::Real;

  Mat A = Mat::Random(m, k);
  Mat B = Mat::Random(k, n);

  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A);
  auto d_B = gpu::DeviceMatrix<Scalar>::fromHost(B);

  // Expression: d_C = d_A * d_B
  gpu::DeviceMatrix<Scalar> d_C;
  d_C = d_A * d_B;

  Mat C = d_C.toHost();
  Mat C_ref = A * B;

  RealScalar tol = RealScalar(k) * NumTraits<Scalar>::epsilon() * C_ref.norm();
  VERIFY((C - C_ref).norm() < tol);
}

// ---- GEMM with adjoint: C = A^H * B ----------------------------------------

template <typename Scalar>
void test_gemm_adjoint_lhs(Index m, Index n, Index k) {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;
  using RealScalar = typename NumTraits<Scalar>::Real;

  Mat A = Mat::Random(k, m);  // A is k×m, A^H is m×k
  Mat B = Mat::Random(k, n);

  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A);
  auto d_B = gpu::DeviceMatrix<Scalar>::fromHost(B);

  gpu::DeviceMatrix<Scalar> d_C;
  d_C = d_A.adjoint() * d_B;

  Mat C = d_C.toHost();
  Mat C_ref = A.adjoint() * B;

  RealScalar tol = RealScalar(k) * NumTraits<Scalar>::epsilon() * C_ref.norm();
  VERIFY((C - C_ref).norm() < tol);
}

// ---- GEMM with transpose: C = A * B^T --------------------------------------

template <typename Scalar>
void test_gemm_transpose_rhs(Index m, Index n, Index k) {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;
  using RealScalar = typename NumTraits<Scalar>::Real;

  Mat A = Mat::Random(m, k);
  Mat B = Mat::Random(n, k);  // B is n×k, B^T is k×n

  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A);
  auto d_B = gpu::DeviceMatrix<Scalar>::fromHost(B);

  gpu::DeviceMatrix<Scalar> d_C;
  d_C = d_A * d_B.transpose();

  Mat C = d_C.toHost();
  Mat C_ref = A * B.transpose();

  RealScalar tol = RealScalar(k) * NumTraits<Scalar>::epsilon() * C_ref.norm();
  VERIFY((C - C_ref).norm() < tol);
}

// ---- GEMM with view/view operands: C = alpha * A^T * B^H -------------------

template <typename Scalar>
void test_gemm_view_pairs(Index m, Index n, Index k) {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;
  using RealScalar = typename NumTraits<Scalar>::Real;

  Mat A = Mat::Random(k, m);  // A^T is m x k
  Mat B = Mat::Random(n, k);  // B^H is k x n
  Scalar alpha = Scalar(2);

  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A);
  auto d_B = gpu::DeviceMatrix<Scalar>::fromHost(B);

  gpu::DeviceMatrix<Scalar> d_C;
  d_C = (alpha * d_A.transpose()) * d_B.adjoint();

  Mat C = d_C.toHost();
  Mat C_ref = (alpha * A.transpose()) * B.adjoint();

  RealScalar tol = RealScalar(k) * NumTraits<Scalar>::epsilon() * C_ref.norm();
  VERIFY((C - C_ref).norm() < tol);
}

// ---- GEMM with scaled: C = alpha * A * B ------------------------------------

template <typename Scalar>
void test_gemm_scaled(Index m, Index n, Index k) {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;
  using RealScalar = typename NumTraits<Scalar>::Real;

  Mat A = Mat::Random(m, k);
  Mat B = Mat::Random(k, n);
  Scalar alpha = Scalar(2.5);

  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A);
  auto d_B = gpu::DeviceMatrix<Scalar>::fromHost(B);

  gpu::DeviceMatrix<Scalar> d_C;
  d_C = alpha * d_A * d_B;

  Mat C = d_C.toHost();
  Mat C_ref = alpha * A * B;

  RealScalar tol = RealScalar(k) * NumTraits<Scalar>::epsilon() * C_ref.norm();
  VERIFY((C - C_ref).norm() < tol);
}

// ---- GEMM accumulate: C += A * B (beta=1) -----------------------------------

template <typename Scalar>
void test_gemm_accumulate(Index m, Index n, Index k) {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;
  using RealScalar = typename NumTraits<Scalar>::Real;

  Mat A = Mat::Random(m, k);
  Mat B = Mat::Random(k, n);
  Mat C_init = Mat::Random(m, n);

  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A);
  auto d_B = gpu::DeviceMatrix<Scalar>::fromHost(B);
  auto d_C = gpu::DeviceMatrix<Scalar>::fromHost(C_init);

  d_C += d_A * d_B;

  Mat C = d_C.toHost();
  Mat C_ref = C_init + A * B;

  RealScalar tol = RealScalar(k) * NumTraits<Scalar>::epsilon() * C_ref.norm();
  VERIFY((C - C_ref).norm() < tol);
}

// ---- GEMM accumulate into empty destination ---------------------------------

template <typename Scalar>
void test_gemm_accumulate_empty(Index m, Index n, Index k) {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;
  using RealScalar = typename NumTraits<Scalar>::Real;

  Mat A = Mat::Random(m, k);
  Mat B = Mat::Random(k, n);

  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A);
  auto d_B = gpu::DeviceMatrix<Scalar>::fromHost(B);
  gpu::DeviceMatrix<Scalar> d_C;

  d_C += d_A * d_B;

  Mat C = d_C.toHost();
  Mat C_ref = A * B;

  RealScalar tol = RealScalar(k) * NumTraits<Scalar>::epsilon() * C_ref.norm();
  VERIFY((C - C_ref).norm() < tol);
}

// ---- GEMM subtract: C -= A * B (beta=1, alpha=-1) --------------------------

template <typename Scalar>
void test_gemm_subtract(Index m, Index n, Index k) {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;
  using RealScalar = typename NumTraits<Scalar>::Real;

  Mat A = Mat::Random(m, k);
  Mat B = Mat::Random(k, n);
  Mat C_init = Mat::Random(m, n);

  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A);
  auto d_B = gpu::DeviceMatrix<Scalar>::fromHost(B);
  auto d_C = gpu::DeviceMatrix<Scalar>::fromHost(C_init);

  gpu::Context ctx;
  d_C.device(ctx) -= d_A * d_B;

  Mat C = d_C.toHost();
  Mat C_ref = C_init - A * B;

  RealScalar tol = RealScalar(k) * NumTraits<Scalar>::epsilon() * C_ref.norm();
  VERIFY((C - C_ref).norm() < tol);
}

// ---- GEMM subtract from empty destination -----------------------------------

template <typename Scalar>
void test_gemm_subtract_empty(Index m, Index n, Index k) {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;
  using RealScalar = typename NumTraits<Scalar>::Real;

  Mat A = Mat::Random(m, k);
  Mat B = Mat::Random(k, n);

  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A);
  auto d_B = gpu::DeviceMatrix<Scalar>::fromHost(B);

  gpu::Context ctx;
  gpu::DeviceMatrix<Scalar> d_C;
  d_C.device(ctx) -= d_A * d_B;

  Mat C = d_C.toHost();
  Mat C_ref = -(A * B);

  RealScalar tol = RealScalar(k) * NumTraits<Scalar>::epsilon() * C_ref.norm();
  VERIFY((C - C_ref).norm() < tol);
}

// ---- GEMM with scaled RHS: C = A * (alpha * B) -----------------------------

template <typename Scalar>
void test_gemm_scaled_rhs(Index m, Index n, Index k) {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;
  using RealScalar = typename NumTraits<Scalar>::Real;

  Mat A = Mat::Random(m, k);
  Mat B = Mat::Random(k, n);
  Scalar alpha = Scalar(3.0);

  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A);
  auto d_B = gpu::DeviceMatrix<Scalar>::fromHost(B);

  gpu::DeviceMatrix<Scalar> d_C;
  d_C = d_A * (alpha * d_B);

  Mat C = d_C.toHost();
  Mat C_ref = A * (alpha * B);

  RealScalar tol = RealScalar(k) * NumTraits<Scalar>::epsilon() * C_ref.norm();
  VERIFY((C - C_ref).norm() < tol);
}

// ---- GEMM dimension mismatch must assert ------------------------------------
// Note: we do NOT use VERIFY_RAISES_ASSERT here because it relies on
// setjmp/longjmp which skips RAII destructors for DeviceMatrix (GPU memory)
// and cuBLAS/cuSOLVER handles, corrupting state for subsequent tests.

// ---- GEMM with explicit gpu::Context ------------------------------------------

template <typename Scalar>
void test_gemm_explicit_context(Index m, Index n, Index k) {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;
  using RealScalar = typename NumTraits<Scalar>::Real;

  Mat A = Mat::Random(m, k);
  Mat B = Mat::Random(k, n);

  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A);
  auto d_B = gpu::DeviceMatrix<Scalar>::fromHost(B);

  gpu::Context ctx;
  gpu::DeviceMatrix<Scalar> d_C;
  d_C.device(ctx) = d_A * d_B;

  Mat C = d_C.toHost();
  Mat C_ref = A * B;

  RealScalar tol = RealScalar(k) * NumTraits<Scalar>::epsilon() * C_ref.norm();
  VERIFY((C - C_ref).norm() < tol);
}

// ---- GEMM with misaligned host alpha/beta -----------------------------------
// cuBLAS reads alpha and beta as cuComplex/cuDoubleComplex, which declare 8- and
// 16-byte alignment; std::complex<double> only needs 8, so a caller may hold it
// at an odd multiple of 8. cublaslt_gemm must copy the scalars rather than pass
// such storage on: on MSVC the dispatcher's locals land there. Only the GEMV
// shapes (n == 1, or m == 1 with op(A) = A) and k == 1 fault without the copy,
// so the callers below use those.

// The strictest alignment cuBLAS declares for a host scalar, and the offset that
// puts a scalar halfway between two such boundaries.
constexpr std::size_t kCublasScalarAlign = alignof(cuDoubleComplex);
constexpr std::size_t kMisalignment = kCublasScalarAlign / 2;

template <typename Scalar>
void test_gemm_misaligned_host_scalars(Index m, Index n, Index k) {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;
  using RealScalar = typename NumTraits<Scalar>::Real;

  Mat A = Mat::Random(m, k);
  Mat B = Mat::Random(k, n);

  gpu::Context ctx;
  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A, ctx.stream());
  auto d_B = gpu::DeviceMatrix<Scalar>::fromHost(B, ctx.stream());
  gpu::DeviceMatrix<Scalar> d_C(m, n);

  alignas(kCublasScalarAlign) unsigned char storage[kMisalignment + 2 * sizeof(Scalar)];
  const Scalar* alpha = ::new (storage + kMisalignment) Scalar(2);
  const Scalar* beta = ::new (storage + kMisalignment + sizeof(Scalar)) Scalar(0);
  VERIFY_IS_EQUAL(reinterpret_cast<std::uintptr_t>(alpha) % kCublasScalarAlign, std::uintptr_t(kMisalignment));

  gpu::internal::cublaslt_gemm(ctx.cublasLtHandle(), ctx.cublasHandle(), CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, alpha,
                               d_A.data(), m, d_B.data(), k, beta, d_C.data(), m, ctx.gemmWorkspace(),
                               ctx.gemmPlanCache(), ctx.cublasLtMaxWorkspaceBytes(), ctx.stream());
  d_C.recordReady(ctx.stream());

  Mat C = d_C.toHost();
  Mat C_ref = Scalar(2) * A * B;

  RealScalar tol = RealScalar(k) * NumTraits<Scalar>::epsilon() * C_ref.norm();
  VERIFY((C - C_ref).norm() < tol);
}

// ---- BLAS-1 with misaligned host alpha --------------------------------------
// Zaxpy and Zscal fault on a host-mode alpha at 8 mod 16 like the GEMV paths
// above; the complex axpy/scal wrappers must copy it.

template <typename Scalar>
void test_blas1_misaligned_host_alpha(Index n) {
  using Vec = Eigen::Matrix<Scalar, Dynamic, 1>;

  Vec x = Vec::Random(n);
  Vec y = Vec::Random(n);

  gpu::Context ctx;
  auto d_x = gpu::DeviceMatrix<Scalar>::fromHost(x, ctx.stream());
  auto d_y = gpu::DeviceMatrix<Scalar>::fromHost(y, ctx.stream());

  alignas(kCublasScalarAlign) unsigned char storage[kMisalignment + sizeof(Scalar)];
  const Scalar* alpha = ::new (storage + kMisalignment) Scalar(2);
  VERIFY_IS_EQUAL(reinterpret_cast<std::uintptr_t>(alpha) % kCublasScalarAlign, std::uintptr_t(kMisalignment));

  EIGEN_CUBLAS_CHECK(gpu::internal::cublasXaxpy(ctx.cublasHandle(), n, alpha, d_x.data(), 1, d_y.data(), 1));
  EIGEN_CUBLAS_CHECK(gpu::internal::cublasXscal(ctx.cublasHandle(), n, alpha, d_x.data(), 1));
  d_x.recordReady(ctx.stream());
  d_y.recordReady(ctx.stream());

  VERIFY_IS_APPROX(Vec(d_y.toHost()), Vec(y + Scalar(2) * x));
  VERIFY_IS_APPROX(Vec(d_x.toHost()), Vec(Scalar(2) * x));
}

// ---- GEMM cross-context reuse of the same destination -----------------------

template <typename Scalar>
void test_gemm_cross_context_reuse(Index n) {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;
  using RealScalar = typename NumTraits<Scalar>::Real;

  Mat A = Mat::Random(n, n);
  Mat B = Mat::Random(n, n);
  Mat D = Mat::Random(n, n);
  Mat E = Mat::Random(n, n);

  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A);
  auto d_B = gpu::DeviceMatrix<Scalar>::fromHost(B);
  auto d_D = gpu::DeviceMatrix<Scalar>::fromHost(D);
  auto d_E = gpu::DeviceMatrix<Scalar>::fromHost(E);

  gpu::Context ctx1;
  gpu::Context ctx2;
  gpu::DeviceMatrix<Scalar> d_C;
  d_C.device(ctx1) = d_A * d_B;
  d_C.device(ctx2) += d_D * d_E;

  Mat C = d_C.toHost();
  Mat C_ref = A * B + D * E;

  RealScalar tol = RealScalar(2) * RealScalar(n) * NumTraits<Scalar>::epsilon() * C_ref.norm();
  VERIFY((C - C_ref).norm() < tol);
}

// ---- GEMM cross-context resize of the destination ---------------------------

template <typename Scalar>
void test_gemm_cross_context_resize() {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;
  using RealScalar = typename NumTraits<Scalar>::Real;

  Mat A = Mat::Random(64, 64);
  Mat B = Mat::Random(64, 64);
  Mat D = Mat::Random(32, 16);
  Mat E = Mat::Random(16, 8);

  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A);
  auto d_B = gpu::DeviceMatrix<Scalar>::fromHost(B);
  auto d_D = gpu::DeviceMatrix<Scalar>::fromHost(D);
  auto d_E = gpu::DeviceMatrix<Scalar>::fromHost(E);

  gpu::Context ctx1;
  gpu::Context ctx2;
  gpu::DeviceMatrix<Scalar> d_C;
  d_C.device(ctx1) = d_A * d_B;
  d_C.device(ctx2) = d_D * d_E;

  Mat C = d_C.toHost();
  Mat C_ref = D * E;

  RealScalar tol = RealScalar(16) * NumTraits<Scalar>::epsilon() * C_ref.norm();
  VERIFY((C - C_ref).norm() < tol);
}

// ---- GEMM chaining: C = (A * B) then D = C * E -----------------------------

template <typename Scalar>
void test_gemm_chain(Index n) {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;
  using RealScalar = typename NumTraits<Scalar>::Real;

  Mat A = Mat::Random(n, n);
  Mat B = Mat::Random(n, n);
  Mat E = Mat::Random(n, n);

  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A);
  auto d_B = gpu::DeviceMatrix<Scalar>::fromHost(B);
  auto d_E = gpu::DeviceMatrix<Scalar>::fromHost(E);

  gpu::DeviceMatrix<Scalar> d_C;
  d_C = d_A * d_B;
  gpu::DeviceMatrix<Scalar> d_D;
  d_D = d_C * d_E;

  Mat D = d_D.toHost();
  Mat D_ref = (A * B) * E;

  RealScalar tol = RealScalar(2) * RealScalar(n) * NumTraits<Scalar>::epsilon() * D_ref.norm();
  VERIFY((D - D_ref).norm() < tol);
}

// ---- GEMM with m, n, or k == 0 ---------------------------------------------

template <typename Scalar>
void test_gemm_empty_dim(Index m, Index n, Index k) {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;

  Mat A = Mat::Random(m, k);
  Mat B = Mat::Random(k, n);
  Mat C_init = Mat::Random(m, n);
  const Mat zero = Mat::Zero(m, n);

  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A);
  auto d_B = gpu::DeviceMatrix<Scalar>::fromHost(B);
  auto d_AH = gpu::DeviceMatrix<Scalar>::fromHost(Mat(A.adjoint()));
  auto d_BT = gpu::DeviceMatrix<Scalar>::fromHost(Mat(B.transpose()));

  // A * B is the m x n zero matrix: `=` overwrites a stale destination, and
  // `+=` / `-=` leave C unchanged.
  auto d_C = gpu::DeviceMatrix<Scalar>::fromHost(C_init);
  d_C = d_A * d_B;
  VERIFY_IS_CWISE_EQUAL(d_C.toHost(), zero);

  gpu::DeviceMatrix<Scalar> d_D;
  d_D = d_AH.adjoint() * d_BT.transpose();
  VERIFY_IS_CWISE_EQUAL(d_D.toHost(), zero);

  d_C = gpu::DeviceMatrix<Scalar>::fromHost(C_init);
  d_C += d_A * d_BT.transpose();
  VERIFY_IS_CWISE_EQUAL(d_C.toHost(), C_init);

  gpu::Context ctx;
  d_C.device(ctx) -= (Scalar(2) * d_AH.adjoint()) * d_B;
  VERIFY_IS_CWISE_EQUAL(d_C.toHost(), C_init);

  // An empty destination accumulates from zero.
  gpu::DeviceMatrix<Scalar> d_E;
  d_E += d_A * d_B;
  VERIFY_IS_CWISE_EQUAL(d_E.toHost(), zero);
}

// ---- LLT solve expression: d_X = d_A.llt().solve(d_B) ----------------------

template <typename MatrixType>
MatrixType make_spd(Index n) {
  using Scalar = typename MatrixType::Scalar;
  MatrixType M = MatrixType::Random(n, n);
  return M.adjoint() * M + MatrixType::Identity(n, n) * static_cast<Scalar>(n);
}

template <typename Scalar>
void test_llt_solve_expr(Index n, Index nrhs) {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;
  using RealScalar = typename NumTraits<Scalar>::Real;

  Mat A = make_spd<Mat>(n);
  Mat B = Mat::Random(n, nrhs);

  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A);
  auto d_B = gpu::DeviceMatrix<Scalar>::fromHost(B);

  gpu::DeviceMatrix<Scalar> d_X;
  d_X = d_A.llt().solve(d_B);

  Mat X = d_X.toHost();
  RealScalar residual = (A * X - B).norm() / B.norm();
  VERIFY(residual < RealScalar(n) * NumTraits<Scalar>::epsilon());
}

// ---- LLT solve with explicit context ----------------------------------------

template <typename Scalar>
void test_llt_solve_expr_context(Index n, Index nrhs) {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;
  using RealScalar = typename NumTraits<Scalar>::Real;

  Mat A = make_spd<Mat>(n);
  Mat B = Mat::Random(n, nrhs);

  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A);
  auto d_B = gpu::DeviceMatrix<Scalar>::fromHost(B);

  gpu::Context ctx;
  gpu::DeviceMatrix<Scalar> d_X;
  d_X.device(ctx) = d_A.llt().solve(d_B);

  Mat X = d_X.toHost();
  RealScalar residual = (A * X - B).norm() / B.norm();
  VERIFY(residual < RealScalar(n) * NumTraits<Scalar>::epsilon());
}

// ---- LU solve expression: d_X = d_A.lu().solve(d_B) ------------------------

template <typename Scalar>
void test_lu_solve_expr(Index n, Index nrhs) {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;
  using RealScalar = typename NumTraits<Scalar>::Real;

  Mat A = Mat::Random(n, n);
  Mat B = Mat::Random(n, nrhs);

  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A);
  auto d_B = gpu::DeviceMatrix<Scalar>::fromHost(B);

  gpu::DeviceMatrix<Scalar> d_X;
  d_X = d_A.lu().solve(d_B);

  Mat X = d_X.toHost();
  RealScalar residual = (A * X - B).norm() / (A.norm() * X.norm());
  VERIFY(residual < RealScalar(10) * RealScalar(n) * NumTraits<Scalar>::epsilon());
}

// ---- GEMM + solver chain: C = A * B, X = C.llt().solve(D) ------------------

template <typename Scalar>
void test_gemm_then_solve(Index n) {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;
  using RealScalar = typename NumTraits<Scalar>::Real;

  Mat A = Mat::Random(n, n);
  Mat D = Mat::Random(n, 1);

  // Make SPD: C = A^H * A + n*I
  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A);
  gpu::DeviceMatrix<Scalar> d_C;
  d_C = d_A.adjoint() * d_A;

  // Add n*I on host (no element-wise ops on DeviceMatrix yet).
  Mat C = d_C.toHost();
  C += Mat::Identity(n, n) * static_cast<Scalar>(n);
  d_C = gpu::DeviceMatrix<Scalar>::fromHost(C);

  auto d_D = gpu::DeviceMatrix<Scalar>::fromHost(D);

  gpu::DeviceMatrix<Scalar> d_X;
  d_X = d_C.llt().solve(d_D);

  Mat X = d_X.toHost();
  RealScalar residual = (C * X - D).norm() / D.norm();
  VERIFY(residual < RealScalar(n) * NumTraits<Scalar>::epsilon());
}

// ---- LLT solve with Upper triangle -----------------------------------------

template <typename Scalar>
void test_llt_solve_upper(Index n, Index nrhs) {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;
  using RealScalar = typename NumTraits<Scalar>::Real;

  Mat A = make_spd<Mat>(n);
  Mat B = Mat::Random(n, nrhs);

  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A);
  auto d_B = gpu::DeviceMatrix<Scalar>::fromHost(B);

  gpu::DeviceMatrix<Scalar> d_X;
  d_X = d_A.template llt<Upper>().solve(d_B);

  Mat X = d_X.toHost();
  RealScalar residual = (A * X - B).norm() / B.norm();
  VERIFY(residual < RealScalar(n) * NumTraits<Scalar>::epsilon());
}

// ---- LU solve with explicit context -----------------------------------------

template <typename Scalar>
void test_lu_solve_expr_context(Index n, Index nrhs) {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;
  using RealScalar = typename NumTraits<Scalar>::Real;

  Mat A = Mat::Random(n, n);
  Mat B = Mat::Random(n, nrhs);

  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A);
  auto d_B = gpu::DeviceMatrix<Scalar>::fromHost(B);

  gpu::Context ctx;
  gpu::DeviceMatrix<Scalar> d_X;
  d_X.device(ctx) = d_A.lu().solve(d_B);

  Mat X = d_X.toHost();
  RealScalar residual = (A * X - B).norm() / (A.norm() * X.norm());
  VERIFY(residual < RealScalar(10) * RealScalar(n) * NumTraits<Scalar>::epsilon());
}

// ---- Zero-nrhs solver expressions ------------------------------------------

template <typename Scalar>
void test_llt_solve_zero_nrhs(Index n) {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;

  Mat A = make_spd<Mat>(n);
  Mat B = Mat::Random(n, 0);

  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A);
  auto d_B = gpu::DeviceMatrix<Scalar>::fromHost(B);

  gpu::DeviceMatrix<Scalar> d_X;
  d_X = d_A.llt().solve(d_B);

  VERIFY_IS_EQUAL(d_X.rows(), n);
  VERIFY_IS_EQUAL(d_X.cols(), 0);
}

template <typename Scalar>
void test_lu_solve_zero_nrhs(Index n) {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;

  Mat A = Mat::Random(n, n);
  Mat B = Mat::Random(n, 0);

  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A);
  auto d_B = gpu::DeviceMatrix<Scalar>::fromHost(B);

  gpu::DeviceMatrix<Scalar> d_X;
  d_X = d_A.lu().solve(d_B);

  VERIFY_IS_EQUAL(d_X.rows(), n);
  VERIFY_IS_EQUAL(d_X.cols(), 0);
}

// ---- TRSM: triangularView<UpLo>().solve(B) ----------------------------------

template <typename Scalar, int UpLo>
void test_trsm(Index n, Index nrhs) {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;
  using RealScalar = typename NumTraits<Scalar>::Real;

  // Build a well-conditioned triangular matrix.
  Mat A = Mat::Random(n, n);
  A.diagonal().array() += static_cast<Scalar>(n);  // ensure non-singular
  if (UpLo == Lower)
    A = A.template triangularView<Lower>();
  else
    A = A.template triangularView<Upper>();

  Mat B = Mat::Random(n, nrhs);

  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A);
  auto d_B = gpu::DeviceMatrix<Scalar>::fromHost(B);

  gpu::DeviceMatrix<Scalar> d_X;
  d_X = d_A.template triangularView<UpLo>().solve(d_B);

  Mat X = d_X.toHost();
  RealScalar residual = (A * X - B).norm() / B.norm();
  VERIFY(residual < RealScalar(n) * NumTraits<Scalar>::epsilon());
}

// ---- SYMM/HEMM: selfadjointView<UpLo>() * B --------------------------------

template <typename Scalar, int UpLo>
void test_symm(Index n, Index nrhs) {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;
  using RealScalar = typename NumTraits<Scalar>::Real;

  Mat A = make_spd<Mat>(n);  // SPD is also self-adjoint
  Mat B = Mat::Random(n, nrhs);

  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A);
  auto d_B = gpu::DeviceMatrix<Scalar>::fromHost(B);

  gpu::DeviceMatrix<Scalar> d_C;
  d_C = d_A.template selfadjointView<UpLo>() * d_B;

  Mat C = d_C.toHost();
  Mat C_ref = A * B;  // A is symmetric, so full multiply == symm

  RealScalar tol = RealScalar(n) * NumTraits<Scalar>::epsilon() * C_ref.norm();
  VERIFY((C - C_ref).norm() < tol);
}

// ---- SYRK/HERK: rankUpdate(A) → C = A * A^H --------------------------------

template <typename Scalar>
void test_syrk(Index n, Index k) {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;
  using RealScalar = typename NumTraits<Scalar>::Real;

  Mat A = Mat::Random(n, k);

  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A);

  gpu::DeviceMatrix<Scalar> d_C;
  d_C.template selfadjointView<Lower>().rankUpdate(d_A);

  Mat C = d_C.toHost();
  // Only lower triangle is meaningful for SYRK. Compare lower triangle.
  Mat C_ref = A * A.adjoint();

  // Extract lower triangle for comparison.
  Mat C_lower = C.template triangularView<Lower>();
  Mat C_ref_lower = C_ref.template triangularView<Lower>();

  RealScalar tol = RealScalar(k) * NumTraits<Scalar>::epsilon() * C_ref.norm();
  VERIFY((C_lower - C_ref_lower).norm() < tol);
}

// ---- Expression sugar: copy-init, operator-=, scalar operators --------------

template <typename Scalar>
void test_expression_sugar(Index m, Index n, Index k) {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;
  using RealScalar = typename NumTraits<Scalar>::Real;

  Mat A = Mat::Random(m, k);
  Mat B = Mat::Random(k, n);
  Mat C0 = Mat::Random(m, n);

  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A);
  auto d_B = gpu::DeviceMatrix<Scalar>::fromHost(B);

  const RealScalar tol = RealScalar(k) * NumTraits<Scalar>::epsilon() * RealScalar(A.norm() * B.norm());

  // Copy-initialization straight from a GEMM expression (README headline syntax).
  gpu::DeviceMatrix<Scalar> d_C = d_A * d_B;
  VERIFY((d_C.toHost() - A * B).norm() < tol);

  // Plain operator-= on GEMM (no .device(ctx) needed).
  auto d_E = gpu::DeviceMatrix<Scalar>::fromHost(C0);
  d_E -= d_A * d_B;
  VERIFY((d_E.toHost() - (C0 - A * B)).norm() < tol);

  // Scalar sugar: int literals, matrix * scalar, division, unary minus.
  Mat S = Mat::Random(m, n);
  auto d_S = gpu::DeviceMatrix<Scalar>::fromHost(S);
  gpu::DeviceMatrix<Scalar> d_F = 2 * d_S;
  VERIFY_IS_APPROX(d_F.toHost(), Mat(Scalar(2) * S));

  gpu::DeviceMatrix<Scalar> d_G = d_S * Scalar(3);
  VERIFY_IS_APPROX(d_G.toHost(), Mat(Scalar(3) * S));

  gpu::DeviceMatrix<Scalar> d_H = d_S / 2;
  VERIFY_IS_APPROX(d_H.toHost(), Mat(S / Scalar(2)));

  gpu::DeviceMatrix<Scalar> d_I = -d_S;
  VERIFY_IS_APPROX(d_I.toHost(), Mat(-S));

  // Scaled ± Scaled and Scaled ± Matrix combinations (geam quadrant).
  Mat T = Mat::Random(m, n);
  auto d_T = gpu::DeviceMatrix<Scalar>::fromHost(T);
  gpu::DeviceMatrix<Scalar> d_J = 2 * d_S - d_T;
  VERIFY_IS_APPROX(d_J.toHost(), Mat(Scalar(2) * S - T));

  gpu::DeviceMatrix<Scalar> d_K = 2 * d_S + 3 * d_T;
  VERIFY_IS_APPROX(d_K.toHost(), Mat(Scalar(2) * S + Scalar(3) * T));

  gpu::DeviceMatrix<Scalar> d_L = -(2 * d_S) - 3 * d_T;
  VERIFY_IS_APPROX(d_L.toHost(), Mat(Scalar(-2) * S - Scalar(3) * T));

  // In-place rescale via assignment (geam with beta=0, aliasing-safe).
  d_L = Scalar(0.5) * d_L;
  VERIFY_IS_APPROX(d_L.toHost(), Mat(Scalar(-1) * S - Scalar(1.5) * T));

  // .device(ctx) with geam expressions.
  gpu::Context ctx;
  gpu::DeviceMatrix<Scalar> d_M;
  d_M.device(ctx) = d_S + 2 * d_T;
  VERIFY_IS_APPROX(d_M.toHost(), Mat(S + Scalar(2) * T));
}

template <typename Scalar>
void test_solve_expr_copy_init(Index n, Index nrhs) {
  using Mat = Eigen::Matrix<Scalar, Dynamic, Dynamic>;
  using RealScalar = typename NumTraits<Scalar>::Real;

  Mat A = make_spd<Mat>(n);
  Mat B = Mat::Random(n, nrhs);

  auto d_A = gpu::DeviceMatrix<Scalar>::fromHost(A);
  auto d_B = gpu::DeviceMatrix<Scalar>::fromHost(B);

  // Copy-initialization from one-shot solver expressions.
  gpu::DeviceMatrix<Scalar> d_X = d_A.llt().solve(d_B);
  Mat X = d_X.toHost();
  VERIFY((A * X - B).norm() / B.norm() < RealScalar(n) * NumTraits<Scalar>::epsilon());

  gpu::DeviceMatrix<Scalar> d_Y = d_A.lu().solve(d_B);
  Mat Y = d_Y.toHost();
  VERIFY((A * Y - B).norm() / B.norm() < RealScalar(10) * RealScalar(n) * NumTraits<Scalar>::epsilon());
}

// ---- Per-scalar driver ------------------------------------------------------

template <typename Scalar>
void test_scalar() {
  CALL_SUBTEST(test_expression_sugar<Scalar>(64, 48, 32));
  CALL_SUBTEST(test_solve_expr_copy_init<Scalar>(64, 4));

  CALL_SUBTEST(test_gemm_basic<Scalar>(64, 64, 64));
  CALL_SUBTEST(test_gemm_basic<Scalar>(128, 64, 32));
  CALL_SUBTEST(test_gemm_basic<Scalar>(1, 1, 1));
  CALL_SUBTEST(test_gemm_basic<Scalar>(256, 256, 256));

  CALL_SUBTEST(test_gemm_adjoint_lhs<Scalar>(64, 64, 64));
  CALL_SUBTEST(test_gemm_adjoint_lhs<Scalar>(128, 32, 64));

  CALL_SUBTEST(test_gemm_transpose_rhs<Scalar>(64, 64, 64));
  CALL_SUBTEST(test_gemm_transpose_rhs<Scalar>(128, 32, 64));
  CALL_SUBTEST(test_gemm_view_pairs<Scalar>(64, 32, 48));

  CALL_SUBTEST(test_gemm_scaled<Scalar>(64, 64, 64));
  CALL_SUBTEST(test_gemm_scaled_rhs<Scalar>(64, 64, 64));
  CALL_SUBTEST(test_gemm_accumulate<Scalar>(64, 64, 64));
  CALL_SUBTEST(test_gemm_accumulate_empty<Scalar>(64, 64, 64));
  CALL_SUBTEST(test_gemm_subtract<Scalar>(64, 64, 64));
  CALL_SUBTEST(test_gemm_subtract_empty<Scalar>(64, 64, 64));
  CALL_SUBTEST(test_gemm_explicit_context<Scalar>(64, 64, 64));
  CALL_SUBTEST(test_gemm_misaligned_host_scalars<Scalar>(64, 1, 32));
  CALL_SUBTEST(test_gemm_misaligned_host_scalars<Scalar>(1, 48, 32));
  CALL_SUBTEST(test_blas1_misaligned_host_alpha<Scalar>(1000));
  CALL_SUBTEST(test_gemm_cross_context_reuse<Scalar>(64));
  CALL_SUBTEST(test_gemm_cross_context_resize<Scalar>());
  CALL_SUBTEST(test_gemm_chain<Scalar>(64));
  CALL_SUBTEST(test_gemm_empty_dim<Scalar>(0, 5, 3));
  CALL_SUBTEST(test_gemm_empty_dim<Scalar>(4, 0, 3));
  CALL_SUBTEST(test_gemm_empty_dim<Scalar>(4, 5, 0));

  // Solver expressions — zero-size edge cases (use dedicated tests, not residual-based)

  // Solver expressions
  CALL_SUBTEST(test_llt_solve_expr<Scalar>(64, 1));
  CALL_SUBTEST(test_llt_solve_expr<Scalar>(64, 4));
  CALL_SUBTEST(test_llt_solve_expr<Scalar>(256, 8));
  CALL_SUBTEST(test_llt_solve_expr_context<Scalar>(64, 4));
  CALL_SUBTEST(test_llt_solve_upper<Scalar>(64, 4));
  CALL_SUBTEST(test_lu_solve_expr<Scalar>(64, 1));
  CALL_SUBTEST(test_lu_solve_expr<Scalar>(64, 4));
  CALL_SUBTEST(test_lu_solve_expr<Scalar>(256, 8));
  CALL_SUBTEST(test_lu_solve_expr_context<Scalar>(64, 4));
  CALL_SUBTEST(test_llt_solve_zero_nrhs<Scalar>(64));
  CALL_SUBTEST(test_llt_solve_zero_nrhs<Scalar>(0));
  CALL_SUBTEST(test_lu_solve_zero_nrhs<Scalar>(64));
  CALL_SUBTEST(test_lu_solve_zero_nrhs<Scalar>(0));
  CALL_SUBTEST(test_gemm_then_solve<Scalar>(64));

  // TRSM
  CALL_SUBTEST((test_trsm<Scalar, Lower>(64, 1)));
  CALL_SUBTEST((test_trsm<Scalar, Lower>(64, 4)));
  CALL_SUBTEST((test_trsm<Scalar, Upper>(64, 4)));
  CALL_SUBTEST((test_trsm<Scalar, Lower>(256, 8)));

  // SYMM/HEMM
  CALL_SUBTEST((test_symm<Scalar, Lower>(64, 4)));
  CALL_SUBTEST((test_symm<Scalar, Upper>(64, 4)));
  CALL_SUBTEST((test_symm<Scalar, Lower>(128, 8)));

  // SYRK/HERK
  CALL_SUBTEST(test_syrk<Scalar>(64, 64));
  CALL_SUBTEST(test_syrk<Scalar>(64, 32));
  CALL_SUBTEST(test_syrk<Scalar>(128, 64));
}

// ---- Solver failure mode tests (not templated on Scalar) --------------------
// Use the cached GpuLLT/GpuLU API which reports failure via info() rather than
// the expression API which asserts inside dispatch (incompatible with
// VERIFY_RAISES_ASSERT due to longjmp skipping RAII destructors).

void test_llt_not_spd() {
  // Negative definite matrix — LLT factorization must fail.
  MatrixXd A = -MatrixXd::Identity(8, 8);
  gpu::LLT<double> llt(A);
  if (gpu_test::cusolver_potrf_missed_non_spd(llt.info())) return;
  VERIFY_IS_EQUAL(llt.info(), NumericalIssue);
}

void test_lu_singular() {
  // Zero matrix — LU factorization must detect singularity.
  MatrixXd A = MatrixXd::Zero(8, 8);
  gpu::LU<double> lu(A);
  VERIFY_IS_EQUAL(lu.info(), NumericalIssue);
}

EIGEN_DECLARE_TEST(gpu_cublas) {
  gpu_test::require_cuda_device();
  // Split by scalar so each part compiles in parallel.
  CALL_SUBTEST_1(test_scalar<float>());
  CALL_SUBTEST_2(test_scalar<double>());
  CALL_SUBTEST_3(test_scalar<std::complex<float>>());
  CALL_SUBTEST_4(test_scalar<std::complex<double>>());
  CALL_SUBTEST_5(test_llt_not_spd());
  CALL_SUBTEST_5(test_lu_singular());
}