From b83eec9172fe6d2c45ac6545e0e42f40dc83c202 Mon Sep 17 00:00:00 2001 From: dyzheng Date: Fri, 14 Aug 2026 15:52:08 +0800 Subject: [PATCH] Fix(pw): correct atom index in non-collinear non-local stress on GPU In the nspin=4 cal_stress_nl CUDA/ROCm kernels, deeq_nc was indexed with `iat + ia` while `iat` was also incremented inside the atom loop, so the atom offset was counted twice. For element types with more than one atom (e.g. Fe2 or HEAs), the second atom onwards read deeq_nc at wrong indices (aliasing another npol block, or out of bounds for the ps3 term), producing a wrong non-local stress with spurious off-diagonal components. Use `iat` alone, consistent with the collinear stress kernel and the non-collinear force kernel, in both the CUDA and ROCm kernels. Add CPU/GPU parity tests for the deeq_nc (nspin=4) overload of cal_stress_nl_op with ntype=1 and atom_na={2}, which fail on the buggy GPU kernel and pass after the fix. Existing tests only covered the real/collinear overload and could not catch this. (cherry picked from commit 97cb313d9, adapted to the LTS source tree layout and memory-op API) --- .../hamilt_pwdft/kernels/cuda/stress_op.cu | 8 +- .../kernels/rocm/stress_op.hip.cu | 8 +- .../kernels/test/stress_op_test.cpp | 224 +++++++++++++++++- 3 files changed, 231 insertions(+), 9 deletions(-) diff --git a/source/module_hamilt_pw/hamilt_pwdft/kernels/cuda/stress_op.cu b/source/module_hamilt_pw/hamilt_pwdft/kernels/cuda/stress_op.cu index dfca5c4e8e9..43e8ccdb538 100644 --- a/source/module_hamilt_pw/hamilt_pwdft/kernels/cuda/stress_op.cu +++ b/source/module_hamilt_pw/hamilt_pwdft/kernels/cuda/stress_op.cu @@ -348,10 +348,10 @@ __global__ void cal_stress_nl( { ps_qq = thrust::complex(- ekb_now * qq_nt[it * deeq_3 * deeq_4 + ip1 * deeq_4 + ip2], 0.0); } - const thrust::complex ps0 = deeq_nc[((iat + ia) * deeq_3 + ip1) * deeq_4 + ip2] + ps_qq; - const thrust::complex ps1 = deeq_nc[((1 * deeq_2 + iat + ia) * deeq_3 + ip1) * deeq_4 + ip2]; - const thrust::complex ps2 = deeq_nc[((2 * deeq_2 + iat + ia) * deeq_3 + ip1) * deeq_4 + ip2]; - const thrust::complex ps3 = deeq_nc[((3 * deeq_2 + iat + ia) * deeq_3 + ip1) * deeq_4 + ip2] + ps_qq; + const thrust::complex ps0 = deeq_nc[((iat) * deeq_3 + ip1) * deeq_4 + ip2] + ps_qq; + const thrust::complex ps1 = deeq_nc[((1 * deeq_2 + iat) * deeq_3 + ip1) * deeq_4 + ip2]; + const thrust::complex ps2 = deeq_nc[((2 * deeq_2 + iat) * deeq_3 + ip1) * deeq_4 + ip2]; + const thrust::complex ps3 = deeq_nc[((3 * deeq_2 + iat) * deeq_3 + ip1) * deeq_4 + ip2] + ps_qq; const int inkb1 = sum + ip1; const int inkb2 = sum + ip2; //out<<"\n ps = "<(- ekb_now * qq_nt[it * deeq_3 * deeq_4 + ip1 * deeq_4 + ip2], 0.0); } - const thrust::complex ps0 = deeq_nc[((iat + ia) * deeq_3 + ip1) * deeq_4 + ip2] + ps_qq; - const thrust::complex ps1 = deeq_nc[((1 * deeq_2 + iat + ia) * deeq_3 + ip1) * deeq_4 + ip2]; - const thrust::complex ps2 = deeq_nc[((2 * deeq_2 + iat + ia) * deeq_3 + ip1) * deeq_4 + ip2]; - const thrust::complex ps3 = deeq_nc[((3 * deeq_2 + iat + ia) * deeq_3 + ip1) * deeq_4 + ip2] + ps_qq; + const thrust::complex ps0 = deeq_nc[((iat) * deeq_3 + ip1) * deeq_4 + ip2] + ps_qq; + const thrust::complex ps1 = deeq_nc[((1 * deeq_2 + iat) * deeq_3 + ip1) * deeq_4 + ip2]; + const thrust::complex ps2 = deeq_nc[((2 * deeq_2 + iat) * deeq_3 + ip1) * deeq_4 + ip2]; + const thrust::complex ps3 = deeq_nc[((3 * deeq_2 + iat) * deeq_3 + ip1) * deeq_4 + ip2] + ps_qq; const int inkb1 = sum + ip1; const int inkb2 = sum + ip2; //out<<"\n ps = "<& d_wg, + const std::vector>& deeq_nc, + const std::vector>& becp, + const std::vector>& dbecp) +{ + double ref = 0.0; + for (int ib = 0; ib < nbands_occ; ib++) + { + const double fac = d_wg[ib]; + const int ib2 = ib * 2; + for (int ia = 0; ia < natom; ia++) + { + for (int ip1 = 0; ip1 < nproj; ip1++) + { + for (int ip2 = 0; ip2 < nproj; ip2++) + { + const std::complex ps0 = deeq_nc[((0 * deeq_2 + ia) * deeq_3 + ip1) * deeq_4 + ip2]; + const std::complex ps1 = deeq_nc[((1 * deeq_2 + ia) * deeq_3 + ip1) * deeq_4 + ip2]; + const std::complex ps2 = deeq_nc[((2 * deeq_2 + ia) * deeq_3 + ip1) * deeq_4 + ip2]; + const std::complex ps3 = deeq_nc[((3 * deeq_2 + ia) * deeq_3 + ip1) * deeq_4 + ip2]; + const int inkb1 = ia * nproj + ip1; + const int inkb2 = ia * nproj + ip2; + const std::complex dbb0 = std::conj(dbecp[ib2 * nkb + inkb1]) * becp[ib2 * nkb + inkb2]; + const std::complex dbb1 = std::conj(dbecp[ib2 * nkb + inkb1]) * becp[(ib2 + 1) * nkb + inkb2]; + const std::complex dbb2 = std::conj(dbecp[(ib2 + 1) * nkb + inkb1]) * becp[ib2 * nkb + inkb2]; + const std::complex dbb3 + = std::conj(dbecp[(ib2 + 1) * nkb + inkb1]) * becp[(ib2 + 1) * nkb + inkb2]; + ref -= fac * (ps0 * dbb0 + ps1 * dbb1 + ps2 * dbb2 + ps3 * dbb3).real(); + } + } + } + } + return ref; +} + +// Deterministic, non-symmetric input data so that any wrong atom or +// spin-block index gives a different result. +void init_nc_inputs(std::vector>& deeq_nc, + std::vector>& becp, + std::vector>& dbecp) +{ + for (size_t i = 0; i < deeq_nc.size(); i++) + { + deeq_nc[i] = std::complex(0.11 * i + 0.03, -0.07 * i + 0.02); + } + for (size_t i = 0; i < becp.size(); i++) + { + becp[i] = std::complex(0.05 * i - 0.31, 0.11 * i + 0.13); + } + for (size_t i = 0; i < dbecp.size(); i++) + { + dbecp[i] = std::complex(-0.06 * i + 0.21, 0.04 * i - 0.52); + } +} +} // namespace + +TEST(TestSrcPWStressMultiDevice, cal_stress_nl_nc_op_cpu) +{ + const int ipol = 0, jpol = 1; + const int nkb = 4, nbands_occ = 2, ntype = 1; + const int natom = 2, nproj = 2; + const int deeq_2 = natom, deeq_3 = nproj, deeq_4 = nproj; + + std::vector atom_na{natom}; + std::vector atom_nh{nproj}; + + std::vector d_wg{0.71, 1.33}; + std::vector qq_nt(1, 0.0); // unused: d_ekb is nullptr + + std::vector> deeq_nc(4 * deeq_2 * deeq_3 * deeq_4); + std::vector> becp(nbands_occ * 2 * nkb); + std::vector> dbecp(nbands_occ * 2 * nkb); + init_nc_inputs(deeq_nc, becp, dbecp); + + const double expected = ref_stress_nl_nc(nbands_occ, + nkb, + natom, + nproj, + deeq_2, + deeq_3, + deeq_4, + d_wg, + deeq_nc, + becp, + dbecp); + + std::vector stress(9, 0.0); + hamilt::cal_stress_nl_op()(cpu_ctx, + ipol, + jpol, + nkb, + nbands_occ, + ntype, + deeq_2, + deeq_3, + deeq_4, + atom_nh.data(), + atom_na.data(), + d_wg.data(), + true, + nullptr, + qq_nt.data(), + deeq_nc.data(), + becp.data(), + dbecp.data(), + stress.data()); + + EXPECT_LT(fabs(stress[ipol * 3 + jpol] - expected), 1e-12); +} + +#if __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM +TEST(TestSrcPWStressMultiDevice, cal_stress_nl_nc_op_gpu) +{ + const int ipol = 0, jpol = 1; + const int nkb = 4, nbands_occ = 2, ntype = 1; + const int natom = 2, nproj = 2; + const int deeq_2 = natom, deeq_3 = nproj, deeq_4 = nproj; + + std::vector atom_na{natom}; + std::vector atom_nh{nproj}; + + std::vector d_wg{0.71, 1.33}; + std::vector qq_nt(1, 0.0); // unused: d_ekb is nullptr + + std::vector> deeq_nc(4 * deeq_2 * deeq_3 * deeq_4); + std::vector> becp(nbands_occ * 2 * nkb); + std::vector> dbecp(nbands_occ * 2 * nkb); + init_nc_inputs(deeq_nc, becp, dbecp); + + const double expected = ref_stress_nl_nc(nbands_occ, + nkb, + natom, + nproj, + deeq_2, + deeq_3, + deeq_4, + d_wg, + deeq_nc, + becp, + dbecp); + + std::vector stress(9, 0.0); + + using delmem_int_op = base_device::memory::delete_memory_op; + using resmem_int_op = base_device::memory::resize_memory_op; + using syncmem_int_h2d_op + = base_device::memory::synchronize_memory_op; + + std::complex *d_deeq_nc = nullptr, *d_becp = nullptr, *d_dbecp = nullptr; + double *dev_wg = nullptr, *d_qq_nt = nullptr, *d_stress = nullptr; + int *d_atom_nh = nullptr, *d_atom_na = nullptr; + + resmem_zd_op()(gpu_ctx, d_deeq_nc, deeq_nc.size()); + resmem_zd_op()(gpu_ctx, d_becp, becp.size()); + resmem_zd_op()(gpu_ctx, d_dbecp, dbecp.size()); + syncmem_z2z_h2d_op()(gpu_ctx, cpu_ctx, d_deeq_nc, deeq_nc.data(), deeq_nc.size()); + syncmem_z2z_h2d_op()(gpu_ctx, cpu_ctx, d_becp, becp.data(), becp.size()); + syncmem_z2z_h2d_op()(gpu_ctx, cpu_ctx, d_dbecp, dbecp.data(), dbecp.size()); + + resmem_dd_op()(gpu_ctx, dev_wg, d_wg.size()); + resmem_dd_op()(gpu_ctx, d_qq_nt, qq_nt.size()); + resmem_dd_op()(gpu_ctx, d_stress, stress.size()); + syncmem_d2d_h2d_op()(gpu_ctx, cpu_ctx, dev_wg, d_wg.data(), d_wg.size()); + syncmem_d2d_h2d_op()(gpu_ctx, cpu_ctx, d_qq_nt, qq_nt.data(), qq_nt.size()); + syncmem_d2d_h2d_op()(gpu_ctx, cpu_ctx, d_stress, stress.data(), stress.size()); + + resmem_int_op()(gpu_ctx, d_atom_nh, atom_nh.size()); + resmem_int_op()(gpu_ctx, d_atom_na, atom_na.size()); + syncmem_int_h2d_op()(gpu_ctx, cpu_ctx, d_atom_nh, atom_nh.data(), atom_nh.size()); + syncmem_int_h2d_op()(gpu_ctx, cpu_ctx, d_atom_na, atom_na.data(), atom_na.size()); + + hamilt::cal_stress_nl_op()(gpu_ctx, + ipol, + jpol, + nkb, + nbands_occ, + ntype, + deeq_2, + deeq_3, + deeq_4, + d_atom_nh, + d_atom_na, + dev_wg, + true, + nullptr, + d_qq_nt, + d_deeq_nc, + d_becp, + d_dbecp, + d_stress); + + syncmem_d2d_d2h_op()(cpu_ctx, gpu_ctx, stress.data(), d_stress, stress.size()); + + EXPECT_LT(fabs(stress[ipol * 3 + jpol] - expected), 1e-12); + + delmem_zd_op()(gpu_ctx, d_deeq_nc); + delmem_zd_op()(gpu_ctx, d_becp); + delmem_zd_op()(gpu_ctx, d_dbecp); + + delmem_dd_op()(gpu_ctx, dev_wg); + delmem_dd_op()(gpu_ctx, d_qq_nt); + delmem_dd_op()(gpu_ctx, d_stress); + + delmem_int_op()(gpu_ctx, d_atom_nh); + delmem_int_op()(gpu_ctx, d_atom_na); +} +#endif // __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM