Skip to content

Apply space-group symmetry transform in X2C GW and HF kernels #14

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
gauravharsha wants to merge 15 commits into from
Closed

Apply space-group symmetry transform in X2C GW and HF kernels #14

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
e270b0b
apply space group symmetry transform in X2C GPU kernels
gauravharsha May 16, 2026
ab7ac3f
add X2C Hubbard symmetry consistency tests for HF and GW (GPU)
gauravharsha May 18, 2026
c7d1fe9
fix outdated integral version in GW_X2C_Hubbard test data
gauravharsha May 19, 2026
653e111
broadcast Sigma result to all MPI ranks after population on rank 0
gauravharsha May 19, 2026
68279e2
fix MPI_Bcast namespace: use green::utils::context()
gauravharsha May 19, 2026
0ead7ea
read dimensions from HDF5 in check_x2c_hubbard_symmetry; drop hardcod…
gauravharsha May 19, 2026
72a734f
test: replace X2C Hubbard symmetry test with cubic Ar def2-svp X2C re…
gauravharsha May 24, 2026
f5680ec
ci: rerun failed ctest with --output-on-failure to surface logs
gauravharsha May 24, 2026
b05ad76
add params for P_sp and Sigma_sp
gauravharsha May 25, 2026
58d22b2
Update GPU data from Ne to Ar.
gauravharsha May 25, 2026
aa029b6
fix: thread nso through cu_symmetry and cugw_utils
gauravharsha May 26, 2026
ed00353
test: dump per-spinor-block max diff before REQUIRE_THAT fires
gauravharsha May 26, 2026
1c6ada9
fix: X2C GW second-tau contraction — fold TR conj into U_q OPs
gauravharsha May 26, 2026
969218b
merge GPU k-AO transforms: single U·G·U† + full-result conj pipeline
gauravharsha May 27, 2026
7e43e7b
wire X2C HF exchange through GPU transform_k_ao_device
gauravharsha May 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .github/workflows/test.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ jobs:
cd build;
# run test on GPU node
#srun -p GPU --time 1:0:0 --exclusive -N 1 -n 1 test/cu_solver_test
ctest -j 1
ctest -j 1 || ctest --rerun-failed --output-on-failure

- name: Run coverage
run: |
Expand Down
125 changes: 43 additions & 82 deletions src/cu_symmetry.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ namespace green::gpu {

cu_symmetry::~cu_symmetry() { release(); }

void cu_symmetry::initialize(const cu_symmetry_data& data, int nao, int naux, int nts, int ns) {
void cu_symmetry::initialize(const cu_symmetry_data& data, int nao, int nso, int naux, int nts, int ns) {
release();

const int nk = static_cast<int>(data.nk);
Expand All @@ -48,11 +48,10 @@ namespace green::gpu {
const int inq = static_cast<int>(data.inq);

nao_ = nao;
nso_ = nso;
naux_ = naux;
nts_ = nts;
ns_ = ns;
batch_count_ = static_cast<size_t>(nts) * static_cast<size_t>(ns);
matrix_stride_ = static_cast<size_t>(nao) * static_cast<size_t>(nao);

// Cache host-side data for accessor methods and TR-conjugation check
k_full_to_reduced_h_ = data.k_full_to_reduced;
Expand Down Expand Up @@ -81,26 +80,35 @@ namespace green::gpu {
upload_array(k_tr_conj_d_, data.k_tr_conj, "k_tr_conj_d_");
upload_array(q_tr_conj_d_, data.q_tr_conj, "q_tr_conj_d_");

// Determine matrix dimension of the AO-space transform: nao for scalar,
// nso (= 2·nao) for X2C. The same GEMM-based U·G·U† + TR-conj pipeline
// (transform_k_ao_device) handles both, with σ_y spinor mixing baked into
// the nso×nso U matrices coming from the input file.
if (!data.k_ao_transforms.empty()) {
// Compute transform matrix dimension: nao for scalar, nso (= 2*nao) for X2C.
size_t dim_sq = data.k_ao_transforms.size() / data.nk;
k_transform_dim_ = static_cast<int>(std::round(std::sqrt(static_cast<double>(dim_sq))));
// Upload to GPU only for the scalar (nao×nao) case; X2C (nso×nso) uses transform_k_ao_device_2c instead.
if (k_transform_dim_ == nao_) {
std::vector<std::complex<float>> k_ao_f(data.k_ao_transforms.size());
cast_copy_complex(k_ao_f.data(), data.k_ao_transforms.data(), data.k_ao_transforms.size());

if (cudaMalloc(&k_ao_transform_full_d_, data.k_ao_transforms.size() * sizeof(cuDoubleComplex)) != cudaSuccess)
throw std::runtime_error("Failed to allocate k_ao_transform_full_d_.");
if (cudaMalloc(&k_ao_transform_full_f_, data.k_ao_transforms.size() * sizeof(cuComplex)) != cudaSuccess)
throw std::runtime_error("Failed to allocate k_ao_transform_full_f_.");
if (cudaMemcpy(k_ao_transform_full_d_, data.k_ao_transforms.data(), data.k_ao_transforms.size() * sizeof(std::complex<double>), cudaMemcpyHostToDevice) != cudaSuccess)
throw std::runtime_error("Failed to copy k_ao_transform_full to device.");
if (cudaMemcpy(k_ao_transform_full_f_, k_ao_f.data(), k_ao_f.size() * sizeof(std::complex<float>), cudaMemcpyHostToDevice) != cudaSuccess)
throw std::runtime_error("Failed to copy float k_ao_transform_full to device.");
}

std::vector<std::complex<float>> k_ao_f(data.k_ao_transforms.size());
cast_copy_complex(k_ao_f.data(), data.k_ao_transforms.data(), data.k_ao_transforms.size());

if (cudaMalloc(&k_ao_transform_full_d_, data.k_ao_transforms.size() * sizeof(cuDoubleComplex)) != cudaSuccess)
throw std::runtime_error("Failed to allocate k_ao_transform_full_d_.");
if (cudaMalloc(&k_ao_transform_full_f_, data.k_ao_transforms.size() * sizeof(cuComplex)) != cudaSuccess)
throw std::runtime_error("Failed to allocate k_ao_transform_full_f_.");
if (cudaMemcpy(k_ao_transform_full_d_, data.k_ao_transforms.data(), data.k_ao_transforms.size() * sizeof(std::complex<double>), cudaMemcpyHostToDevice) != cudaSuccess)
throw std::runtime_error("Failed to copy k_ao_transform_full to device.");
if (cudaMemcpy(k_ao_transform_full_f_, k_ao_f.data(), k_ao_f.size() * sizeof(std::complex<float>), cudaMemcpyHostToDevice) != cudaSuccess)
throw std::runtime_error("Failed to copy float k_ao_transform_full to device.");
} else {
k_transform_dim_ = nao;
}

// Scratch sizing tracks the transform dimension: X2C input/output buffers
// are nso×nso per (k, t), so callers must pass ns=1 (X2C uses pseudo-ns=4
// for the 4 spin blocks but the matrix dim already absorbs spinor mixing).
batch_count_ = static_cast<size_t>(nts) * static_cast<size_t>(ns);
matrix_stride_ = static_cast<size_t>(k_transform_dim_) * static_cast<size_t>(k_transform_dim_);

// q_p0_transforms stores U_q row-major (as-is). CUBLAS sees U_q^T (col-major).
// Steps 2a/2c in compute_second_tau_contraction use OP_N/OP_C to recover U_q^T and U_q^*.
if (!data.q_p0_transforms.empty()) {
Expand Down Expand Up @@ -190,12 +198,17 @@ namespace green::gpu {
// Symmetry transform: G(k_full) = U * G(k_ibz) * U† (non-TR)
// G(k_full) = conj(U * G(k_ibz) * U†) (TR)
//
// Same pipeline handles scalar (dim = nao) and X2C (dim = nso = 2·nao): for X2C
// the σ_y spinor mixing is baked into the nso×nso U matrices read from the input
// file, so U·G·U† already produces correctly spin-flipped blocks at TR partners.
//
// Both U and G are stored in ROW-MAJOR (ndarray / green-gpu MatrixXcd convention).
// CUBLAS interprets row-major data as the transpose of the intended matrix.
// To compute result_rm = U * G * U†, we need result_cm = (U*G*U†)^T = U^* * G^T * U^T.
// With row-major U: OP_C → (U_rm^T)^H = U_rm^*, OP_N → U_rm^T.
// With row-major G: OP_N → G_rm^T.
// So: GEMM1(OP_C, OP_N) = U^* * G^T, GEMM2(OP_N, OP_N) = result * U^T.
const int dim = k_transform_dim_;

if constexpr (std::is_same_v<cuda_complex_t, cuDoubleComplex>) {
cuDoubleComplex* input_buf = input_scratch ? input_scratch : input_batch_z_d_;
Expand All @@ -209,15 +222,17 @@ namespace green::gpu {
const cuDoubleComplex* U = k_ao_transform_full_d_ + k_full * matrix_stride_;

// work = U^* * G^T (col-major intermediate)
if (GEMM_STRIDED_BATCHED(handle, CUBLAS_OP_C, CUBLAS_OP_N, nao_, nao_, nao_, &one, U, nao_, 0, input_buf, nao_,
matrix_stride_, &zero, work_buf, nao_, matrix_stride_, static_cast<int>(nts * ns)) != CUBLAS_STATUS_SUCCESS)
if (GEMM_STRIDED_BATCHED(handle, CUBLAS_OP_C, CUBLAS_OP_N, dim, dim, dim, &one, U, dim, 0, input_buf, dim,
matrix_stride_, &zero, work_buf, dim, matrix_stride_, static_cast<int>(nts * ns)) != CUBLAS_STATUS_SUCCESS)
throw std::runtime_error("Failed first batched GEMM in cu_symmetry::transform_k_ao_device_impl.");
// out = work * U^T (col-major result, row-major interpretation = U * G * U†)
if (GEMM_STRIDED_BATCHED(handle, CUBLAS_OP_N, CUBLAS_OP_N, nao_, nao_, nao_, &one, work_buf, nao_, matrix_stride_, U,
nao_, 0, &zero, out_device, nao_, matrix_stride_, static_cast<int>(nts * ns)) != CUBLAS_STATUS_SUCCESS)
if (GEMM_STRIDED_BATCHED(handle, CUBLAS_OP_N, CUBLAS_OP_N, dim, dim, dim, &one, work_buf, dim, matrix_stride_, U,
dim, 0, &zero, out_device, dim, matrix_stride_, static_cast<int>(nts * ns)) != CUBLAS_STATUS_SUCCESS)
throw std::runtime_error("Failed second batched GEMM in cu_symmetry::transform_k_ao_device_impl.");

// TR: conjugate output to get conj(U * G_ibz * U†)
// TR: conjugate the full (U·G·U†) output to get conj(U * G_ibz * U†).
// For X2C this conjugates all four spin blocks together; the σ_y row/col
// mixing has already happened inside the GEMMs via the nso×nso U.
if (k_tr_conj_h_.at(k_full) != 0) {
const double minus_one = -1.0;
if (RSCAL(handle, static_cast<int>(batch_elements), &minus_one, reinterpret_cast<double*>(out_device) + 1, 2) != CUBLAS_STATUS_SUCCESS)
Expand All @@ -235,15 +250,15 @@ namespace green::gpu {
const cuComplex* U = k_ao_transform_full_f_ + k_full * matrix_stride_;

// work = U^* * G^T (col-major intermediate)
if (GEMM_STRIDED_BATCHED(handle, CUBLAS_OP_C, CUBLAS_OP_N, nao_, nao_, nao_, &one, U, nao_, 0, input_buf, nao_,
matrix_stride_, &zero, work_buf, nao_, matrix_stride_, static_cast<int>(nts * ns)) != CUBLAS_STATUS_SUCCESS)
if (GEMM_STRIDED_BATCHED(handle, CUBLAS_OP_C, CUBLAS_OP_N, dim, dim, dim, &one, U, dim, 0, input_buf, dim,
matrix_stride_, &zero, work_buf, dim, matrix_stride_, static_cast<int>(nts * ns)) != CUBLAS_STATUS_SUCCESS)
throw std::runtime_error("Failed first batched GEMM in cu_symmetry::transform_k_ao_device_impl.");
// out = work * U^T (col-major result, row-major interpretation = U * G * U†)
if (GEMM_STRIDED_BATCHED(handle, CUBLAS_OP_N, CUBLAS_OP_N, nao_, nao_, nao_, &one, work_buf, nao_, matrix_stride_, U,
nao_, 0, &zero, out_device, nao_, matrix_stride_, static_cast<int>(nts * ns)) != CUBLAS_STATUS_SUCCESS)
if (GEMM_STRIDED_BATCHED(handle, CUBLAS_OP_N, CUBLAS_OP_N, dim, dim, dim, &one, work_buf, dim, matrix_stride_, U,
dim, 0, &zero, out_device, dim, matrix_stride_, static_cast<int>(nts * ns)) != CUBLAS_STATUS_SUCCESS)
throw std::runtime_error("Failed second batched GEMM in cu_symmetry::transform_k_ao_device_impl.");

// TR: conjugate output to get conj(U * G_ibz * U†)
// TR: conjugate the full (U·G·U†) output to get conj(U * G_ibz * U†).
if (k_tr_conj_h_.at(k_full) != 0) {
const float minus_one = -1.0f;
if (RSCAL(handle, static_cast<int>(batch_elements), &minus_one, reinterpret_cast<float*>(out_device) + 1, 2) != CUBLAS_STATUS_SUCCESS)
Expand All @@ -268,58 +283,4 @@ namespace green::gpu {
transform_k_ao_device_impl(handle, stream, in_device, k_full, out_device, nts, ns, ibz_in_device, input_scratch, work_scratch);
}

template <typename cuda_complex_t>
void cu_symmetry::transform_k_ao_device_2c_impl(cublasHandle_t handle, cudaStream_t stream,
cuda_complex_t* ibz_in_device, size_t k_full,
cuda_complex_t* out_device, int nts, int nao) {
using scalar_t = std::conditional_t<std::is_same_v<cuda_complex_t, cuDoubleComplex>, double, float>;
const size_t block_elems = static_cast<size_t>(nts) * nao * nao;
const size_t block_bytes = block_elems * sizeof(cuda_complex_t);

if (k_tr_conj_h_.at(k_full) == 0) {
// No TR: copy all 4 blocks unchanged.
cudaMemcpyAsync(out_device, ibz_in_device, 4 * block_bytes, cudaMemcpyDeviceToDevice, stream);
} else {
// TR needed (hardcoded minus_t=true block permutation):
// ss=0 <- +conj(ibz ss=1) [aa <- conj(bb)]
// ss=1 <- +conj(ibz ss=0) [bb <- conj(aa)]
// ss=2 <- -conj(ibz ss=2) [self]
// ss=3 <- -conj(ibz ss=3) [self]
// ibz_in_device and out_device are distinct, so ss=0/1 swap is safe.
cudaMemcpyAsync(out_device + 0 * block_elems, ibz_in_device + 1 * block_elems, block_bytes, cudaMemcpyDeviceToDevice, stream);
cudaMemcpyAsync(out_device + 1 * block_elems, ibz_in_device + 0 * block_elems, block_bytes, cudaMemcpyDeviceToDevice, stream);
cudaMemcpyAsync(out_device + 2 * block_elems, ibz_in_device + 2 * block_elems, block_bytes, cudaMemcpyDeviceToDevice, stream);
cudaMemcpyAsync(out_device + 3 * block_elems, ibz_in_device + 3 * block_elems, block_bytes, cudaMemcpyDeviceToDevice, stream);

cublasSetStream(handle, stream);
const scalar_t minus_one = static_cast<scalar_t>(-1.0);
// ss=0, ss=1: +conj -> negate imaginary parts only
if (RSCAL(handle, static_cast<int>(block_elems), &minus_one,
reinterpret_cast<scalar_t*>(out_device + 0 * block_elems) + 1, 2) != CUBLAS_STATUS_SUCCESS)
throw std::runtime_error("RSCAL conj ss=0 failed in transform_k_ao_device_2c.");
if (RSCAL(handle, static_cast<int>(block_elems), &minus_one,
reinterpret_cast<scalar_t*>(out_device + 1 * block_elems) + 1, 2) != CUBLAS_STATUS_SUCCESS)
throw std::runtime_error("RSCAL conj ss=1 failed in transform_k_ao_device_2c.");
// ss=2, ss=3: -conj -> negate real parts only
if (RSCAL(handle, static_cast<int>(block_elems), &minus_one,
reinterpret_cast<scalar_t*>(out_device + 2 * block_elems) + 0, 2) != CUBLAS_STATUS_SUCCESS)
throw std::runtime_error("RSCAL -conj ss=2 failed in transform_k_ao_device_2c.");
if (RSCAL(handle, static_cast<int>(block_elems), &minus_one,
reinterpret_cast<scalar_t*>(out_device + 3 * block_elems) + 0, 2) != CUBLAS_STATUS_SUCCESS)
throw std::runtime_error("RSCAL -conj ss=3 failed in transform_k_ao_device_2c.");
}
}

void cu_symmetry::transform_k_ao_device_2c(cublasHandle_t handle, cudaStream_t stream,
cuDoubleComplex* ibz_in_device, size_t k_full,
cuDoubleComplex* out_device, int nts, int nao) {
transform_k_ao_device_2c_impl(handle, stream, ibz_in_device, k_full, out_device, nts, nao);
}

void cu_symmetry::transform_k_ao_device_2c(cublasHandle_t handle, cudaStream_t stream,
cuComplex* ibz_in_device, size_t k_full,
cuComplex* out_device, int nts, int nao) {
transform_k_ao_device_2c_impl(handle, stream, ibz_in_device, k_full, out_device, nts, nao);
}

} // namespace green::gpu
29 changes: 15 additions & 14 deletions src/cugw_qkpt.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -325,10 +325,20 @@ namespace green::gpu {
throw std::runtime_error("GEMM_STRIDED_BATCHED fails on gw_qkpt.compute_second_tau_contraction_2C().");
}
if (U_q != nullptr) {
// q-space symmetry transform: Y2 = U_q^left * P * U_q^right * Y1
// U_q stored row-major as-is. CUBLAS sees U_q^T (col-major).
// q-space symmetry transform: U_q acts on auxiliary basis only, same as scalar case.
// 2a: effective op = OP_N(U_q^T) = U_q^T → T1 = U_q^T * Y1
if (GEMM_STRIDED_BATCHED(*handle_, CUBLAS_OP_N, CUBLAS_OP_T, naux_, nao2_, naux_, &one, U_q, naux_, 0,
// OP_N(U_q^T) = U_q^T OP_C(U_q^T) = U_q^*
// Non-TR: Y2 = U_q^T * P * U_q^* * Y1 (Left = OP_N, Right = OP_C)
// TR : Y2 = U_q^* * P * U_q^T * Y1 (Left = OP_C, Right = OP_N)
// Folding the TR conjugation into the U_q OPs (mirroring the scalar
// path) is mathematically equivalent to applying conj(W) and avoids
// the post-step RSCAL on Y2, which would also conjugate Y1's
// contribution — Y1 already carries the correct TR convention from
// upstream (copy_Gk_2c on the CPU).
cublasOperation_t OP_Uq_Left = q_conj_after_uq ? CUBLAS_OP_C : CUBLAS_OP_N;
cublasOperation_t OP_Uq_Right = q_conj_after_uq ? CUBLAS_OP_N : CUBLAS_OP_C;
// 2a: T1 = OP_Uq_Left(U_q^T) * Y1
if (GEMM_STRIDED_BATCHED(*handle_, OP_Uq_Left, CUBLAS_OP_T, naux_, nao2_, naux_, &one, U_q, naux_, 0,
Y1t_Qin, nao2_, nauxnao2_, &zero, Y2t_inP, naux_, nauxnao2_,
nt_mult) != CUBLAS_STATUS_SUCCESS) {
throw std::runtime_error("GEMM_STRIDED_BATCHED fails on gw_qkpt.compute_second_tau_contraction_2C() [2a].");
Expand All @@ -339,21 +349,12 @@ namespace green::gpu {
nt_mult) != CUBLAS_STATUS_SUCCESS) {
throw std::runtime_error("GEMM_STRIDED_BATCHED fails on gw_qkpt.compute_second_tau_contraction_2C() [2b].");
}
// 2c: effective op = OP_C(U_q^T) = U_q^* → Y2 = U_q^* * T2
if (GEMM_STRIDED_BATCHED(*handle_, CUBLAS_OP_C, CUBLAS_OP_N, naux_, nao2_, naux_, &one, U_q, naux_, 0,
// 2c: Y2 = OP_Uq_Right(U_q^T) * T2
if (GEMM_STRIDED_BATCHED(*handle_, OP_Uq_Right, CUBLAS_OP_N, naux_, nao2_, naux_, &one, U_q, naux_, 0,
Y1t_Qin, naux_, nauxnao2_, &zero, Y2t_inP, naux_, nauxnao2_,
nt_mult) != CUBLAS_STATUS_SUCCESS) {
throw std::runtime_error("GEMM_STRIDED_BATCHED fails on gw_qkpt.compute_second_tau_contraction_2C() [2c].");
}
// TR conjugation after U_q transform: conj(U * P * U†) = conjugate the result
if (q_conj_after_uq) {
scalar_t alpha = -1.0;
int two = 2;
if (RSCAL(*handle_, nauxnao2_ * nt_mult, &alpha, reinterpret_cast<scalar_t*>(Y2t_inP) + 1, two) !=
CUBLAS_STATUS_SUCCESS) {
throw std::runtime_error("RSCAL fails on gw_qkpt.compute_second_tau_contraction_2C() [q_conj_after_uq].");
}
}
} else {
// No q-space transform: Y2(Q,in) = P(Q,Q') * Y1^T(Q',in)
if (GEMM_STRIDED_BATCHED(*handle_, CUBLAS_OP_N, CUBLAS_OP_T, naux_, nao2_, naux_, &one, Pqk_tQP + t * naux2_, naux_,
Expand Down
4 changes: 2 additions & 2 deletions src/cugw_utils.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ namespace green::gpu {

template <typename prec>
cugw_utils<prec>::cugw_utils(int _nts, int _nt_batch, int _nw_b, int _ns, int _nk, int _ink, int _nq, int _inq, int _nqkpt,
int _NQ, int _nao, const cu_symmetry_data& sym_data, ztensor_view<5>& G_tskij_host,
int _NQ, int _nao, int _nso, const cu_symmetry_data& sym_data, ztensor_view<5>& G_tskij_host,
bool low_device_memory, const MatrixXcd& Ttn_FB, const MatrixXcd& Tnt_BF,
LinearSolverType cuda_lin_solver, int _myid, int _intranode_rank, int _devCount_per_node) :
_low_device_memory(low_device_memory), qkpts(_nqkpt), G_tskij_host_(G_tskij_host), V_Qpm(_NQ, _nao, _nao),
Expand Down Expand Up @@ -71,7 +71,7 @@ namespace green::gpu {
g_kstij_device, g_ksmtij_device, sigma_kstij_device, sigma_k_locks);
}

_cu_symmetry.initialize(sym_data, _nao, _NQ, _nts, _ns);
_cu_symmetry.initialize(sym_data, _nao, _nso, _NQ, _nts, _ns);
}

template <typename prec>
Expand Down
Loading
Loading