diff --git a/source/source_estate/kernels/cuda/elecstate_op.cu b/source/source_estate/kernels/cuda/elecstate_op.cu index 4e6feedb7ec..3f60d283d55 100644 --- a/source/source_estate/kernels/cuda/elecstate_op.cu +++ b/source/source_estate/kernels/cuda/elecstate_op.cu @@ -52,7 +52,7 @@ __global__ void elecstate_pw( rho[3 * nrxx_dense + idx] += w1 * (norm(wfcr[idx]) - norm(wfcr_another_spin[idx])); } else { - rho[0 * nrxx_dense + idx] = 0; + // Keep the scalar charge accumulated above; only magnetization is disabled. rho[1 * nrxx_dense + idx] = 0; rho[2 * nrxx_dense + idx] = 0; rho[3 * nrxx_dense + idx] = 0; @@ -103,4 +103,4 @@ void elecstate_pw_op::operator()(const base_dev template struct elecstate_pw_op; template struct elecstate_pw_op; -} // namespace elecstate \ No newline at end of file +} // namespace elecstate diff --git a/source/source_estate/kernels/rocm/elecstate_op.hip.cu b/source/source_estate/kernels/rocm/elecstate_op.hip.cu index 90fbe5b0cd1..8ab453e0c35 100644 --- a/source/source_estate/kernels/rocm/elecstate_op.hip.cu +++ b/source/source_estate/kernels/rocm/elecstate_op.hip.cu @@ -50,7 +50,7 @@ __global__ void elecstate_pw( rho[3 * nrxx + idx] += w1 * (norm(wfcr[idx]) - norm(wfcr_another_spin[idx])); } else { - rho[0 * nrxx + idx] = 0; + // Keep the scalar charge accumulated above; only magnetization is disabled. rho[1 * nrxx + idx] = 0; rho[2 * nrxx + idx] = 0; rho[3 * nrxx + idx] = 0; @@ -96,4 +96,4 @@ void elecstate_pw_op::operator()(const base_dev template struct elecstate_pw_op; template struct elecstate_pw_op; -} \ No newline at end of file +} diff --git a/source/source_estate/kernels/test/elecstate_op_test.cpp b/source/source_estate/kernels/test/elecstate_op_test.cpp index ae441bed7b1..39894274249 100644 --- a/source/source_estate/kernels/test/elecstate_op_test.cpp +++ b/source/source_estate/kernels/test/elecstate_op_test.cpp @@ -129,8 +129,8 @@ TEST_F(TestModuleElecstateMultiDevice, elecstate_pw_op_gpu) EXPECT_LT(fabs(rho_data[ii] - expected_rho[ii]), 6e-5); } delete [] rho; - delete_memory_var_op()(this->gpu_ctx, d_rho_data); - delete_memory_complex_op()(this->gpu_ctx, d_wfcr); + delete_memory_var_op()(d_rho_data); + delete_memory_complex_op()(d_wfcr); } TEST_F(TestModuleElecstateMultiDevice, elecstate_pw_spin_op_gpu) @@ -168,9 +168,65 @@ TEST_F(TestModuleElecstateMultiDevice, elecstate_pw_spin_op_gpu) EXPECT_LT(fabs(rho_data_2[ii] - expected_rho_2[ii]), 5e-4); } delete [] rho; - delete_memory_var_op()(this->gpu_ctx, d_rho_data_2); - delete_memory_complex_op()(this->gpu_ctx, d_wfcr_2); - delete_memory_complex_op()(this->gpu_ctx, d_wfcr_another_spin_2); + delete_memory_var_op()(d_rho_data_2); + delete_memory_complex_op()(d_wfcr_2); + delete_memory_complex_op()(d_wfcr_another_spin_2); } -#endif // __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM +TEST_F(TestModuleElecstateMultiDevice, nonmagnetic_spinor_preserves_charge_on_gpu) +{ + const int nrxx = 2; + const double weight = 0.5; + const bool domag = false; + const bool domag_z = false; + const std::vector> wfcr = {{1.0, 0.0}, {0.0, 2.0}}; + const std::vector> wfcr_another_spin = {{0.0, 1.0}, {3.0, 0.0}}; + std::vector rho_cpu = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0}; + std::vector rho_gpu = rho_cpu; + double* rho_cpu_components[4] + = {rho_cpu.data(), rho_cpu.data() + nrxx, rho_cpu.data() + 2 * nrxx, rho_cpu.data() + 3 * nrxx}; + + elecstate_cpu_op()(this->cpu_ctx, + domag, + domag_z, + nrxx, + nrxx, + weight, + rho_cpu_components, + wfcr.data(), + wfcr_another_spin.data()); + + double* rho_device = nullptr; + std::complex* wfcr_device = nullptr; + std::complex* wfcr_another_spin_device = nullptr; + resize_memory_var_op()(rho_device, rho_gpu.size()); + resize_memory_complex_op()(wfcr_device, wfcr.size()); + resize_memory_complex_op()(wfcr_another_spin_device, wfcr_another_spin.size()); + syncmem_var_h2d_op()(rho_device, rho_gpu.data(), rho_gpu.size()); + syncmem_complex_h2d_op()(wfcr_device, wfcr.data(), wfcr.size()); + syncmem_complex_h2d_op()(wfcr_another_spin_device, wfcr_another_spin.data(), wfcr_another_spin.size()); + double* rho_gpu_components[4] = {rho_device, rho_device + nrxx, rho_device + 2 * nrxx, rho_device + 3 * nrxx}; + + elecstate_gpu_op()(this->gpu_ctx, + domag, + domag_z, + nrxx, + nrxx, + weight, + rho_gpu_components, + wfcr_device, + wfcr_another_spin_device); + syncmem_var_d2h_op()(rho_gpu.data(), rho_device, rho_gpu.size()); + + EXPECT_DOUBLE_EQ(rho_cpu[0], 2.0); + EXPECT_DOUBLE_EQ(rho_cpu[1], 8.5); + for (std::size_t ir = 0; ir < rho_cpu.size(); ++ir) + { + EXPECT_DOUBLE_EQ(rho_gpu[ir], rho_cpu[ir]); + } + + delete_memory_var_op()(rho_device); + delete_memory_complex_op()(wfcr_device); + delete_memory_complex_op()(wfcr_another_spin_device); +} +#endif // __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM diff --git a/source/source_estate/test/CMakeLists.txt b/source/source_estate/test/CMakeLists.txt index 1eaaa0c6bce..bb800d07eb2 100644 --- a/source/source_estate/test/CMakeLists.txt +++ b/source/source_estate/test/CMakeLists.txt @@ -17,6 +17,12 @@ AddTest( SOURCES ../kernels/test/elecstate_op_test.cpp ) +if(USE_CUDA) + target_compile_definitions(MODULE_ESTATE_Elecstate_Op_UTs PRIVATE __UT_USE_CUDA) +elseif(USE_ROCM) + target_compile_definitions(MODULE_ESTATE_Elecstate_Op_UTs PRIVATE __UT_USE_ROCM) +endif() + AddTest( TARGET MODULE_ESTATE_elecstate_occupy LIBS parameter base device