diff --git a/source/source_base/module_device/memory_op.cpp b/source/source_base/module_device/memory_op.cpp index af5bacc6eec..2c38e44145e 100644 --- a/source/source_base/module_device/memory_op.cpp +++ b/source/source_base/module_device/memory_op.cpp @@ -1,6 +1,7 @@ #include "memory_op.h" #include "source_base/memory_recorder.h" +#include "source_base/tool_quit.h" #include "source_base/tool_threading.h" #ifdef __DSP #include "source_base/kernels/dsp/dsp_connector.h" @@ -549,35 +550,52 @@ void set_memory(FPTYPE* arr, const int var, const size_t size, base_device::Abac template void synchronize_memory(FPTYPE* arr_out, const FPTYPE* arr_in, const size_t size, base_device::AbacusDevice_t device_type_out, base_device::AbacusDevice_t device_type_in){ - if (device_type_out == base_device::AbacusDevice_t::CpuDevice || device_type_in == base_device::AbacusDevice_t::CpuDevice){ + // The four source/destination combinations are mutually exclusive, so each + // branch must test BOTH devices. Using `||` here made the first branch match + // whenever either side was the CPU, which routed host<->device transfers to + // the host-to-host specialization. + if (device_type_out == base_device::AbacusDevice_t::CpuDevice && device_type_in == base_device::AbacusDevice_t::CpuDevice){ synchronize_memory_op()(arr_out, arr_in, size); } - else if (device_type_out == base_device::AbacusDevice_t::CpuDevice || device_type_in == base_device::AbacusDevice_t::GpuDevice){ +#if __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM + else if (device_type_out == base_device::AbacusDevice_t::CpuDevice && device_type_in == base_device::AbacusDevice_t::GpuDevice){ synchronize_memory_op()(arr_out, arr_in, size); } - else if (device_type_out == base_device::AbacusDevice_t::GpuDevice || device_type_in == base_device::AbacusDevice_t::CpuDevice){ + else if (device_type_out == base_device::AbacusDevice_t::GpuDevice && device_type_in == base_device::AbacusDevice_t::CpuDevice){ synchronize_memory_op()(arr_out, arr_in, size); } - else if (device_type_out == base_device::AbacusDevice_t::GpuDevice || device_type_in == base_device::AbacusDevice_t::GpuDevice){ + else if (device_type_out == base_device::AbacusDevice_t::GpuDevice && device_type_in == base_device::AbacusDevice_t::GpuDevice){ synchronize_memory_op()(arr_out, arr_in, size); } +#endif + else { + ModuleBase::WARNING_QUIT("base_device::memory::synchronize_memory", + "unsupported source/destination device combination"); + } } template void cast_memory(FPTYPE_out* arr_out, const FPTYPE_in* arr_in, const size_t size, base_device::AbacusDevice_t device_type_out, base_device::AbacusDevice_t device_type_in) { - if (device_type_out == base_device::AbacusDevice_t::CpuDevice || device_type_in == base_device::AbacusDevice_t::CpuDevice){ + // See synchronize_memory() above: dispatch on the exact (out, in) device pair. + if (device_type_out == base_device::AbacusDevice_t::CpuDevice && device_type_in == base_device::AbacusDevice_t::CpuDevice){ cast_memory_op()(arr_out, arr_in, size); } - else if (device_type_out == base_device::AbacusDevice_t::CpuDevice || device_type_in == base_device::AbacusDevice_t::GpuDevice){ +#if __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM + else if (device_type_out == base_device::AbacusDevice_t::CpuDevice && device_type_in == base_device::AbacusDevice_t::GpuDevice){ cast_memory_op()(arr_out, arr_in, size); } - else if (device_type_out == base_device::AbacusDevice_t::GpuDevice || device_type_in == base_device::AbacusDevice_t::CpuDevice){ + else if (device_type_out == base_device::AbacusDevice_t::GpuDevice && device_type_in == base_device::AbacusDevice_t::CpuDevice){ cast_memory_op()(arr_out, arr_in, size); } - else if (device_type_out == base_device::AbacusDevice_t::GpuDevice || device_type_in == base_device::AbacusDevice_t::GpuDevice){ + else if (device_type_out == base_device::AbacusDevice_t::GpuDevice && device_type_in == base_device::AbacusDevice_t::GpuDevice){ cast_memory_op()(arr_out, arr_in, size); } +#endif + else { + ModuleBase::WARNING_QUIT("base_device::memory::cast_memory", + "unsupported source/destination device combination"); + } } template @@ -591,5 +609,24 @@ void delete_memory(FPTYPE* arr, base_device::AbacusDevice_t device_type) } } +// Explicit instantiations of the runtime-dispatch wrappers, so that the +// declarations in memory_op.h can actually be linked from another translation +// unit (and covered by unit tests). cast_memory is instantiated only for the +// type pairs that cast_memory_op provides for all four device combinations. +template void synchronize_memory(int*, const int*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t); +template void synchronize_memory(float*, const float*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t); +template void synchronize_memory(double*, const double*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t); +template void synchronize_memory>(std::complex*, const std::complex*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t); +template void synchronize_memory>(std::complex*, const std::complex*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t); + +template void cast_memory(float*, const float*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t); +template void cast_memory(double*, const double*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t); +template void cast_memory(float*, const double*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t); +template void cast_memory(double*, const float*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t); +template void cast_memory, std::complex>(std::complex*, const std::complex*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t); +template void cast_memory, std::complex>(std::complex*, const std::complex*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t); +template void cast_memory, std::complex>(std::complex*, const std::complex*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t); +template void cast_memory, std::complex>(std::complex*, const std::complex*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t); + } // namespace memory } // namespace base_device \ No newline at end of file diff --git a/source/source_base/module_device/test/memory_test.cpp b/source/source_base/module_device/test/memory_test.cpp index a963c68c532..a6912beffe3 100644 --- a/source/source_base/module_device/test/memory_test.cpp +++ b/source/source_base/module_device/test/memory_test.cpp @@ -170,6 +170,59 @@ TEST_F(TestModulePsiMemory, delete_memory_op_complex_double_cpu) delete_memory_complex_double_cpu_op()(hz_xx); } +// --------------------------------------------------------------------------- +// Runtime-dispatch wrappers (issue #7553). +// +// synchronize_memory()/cast_memory() pick a compile-time specialization from a +// pair of runtime AbacusDevice_t values. The branches used to be joined with +// `||`, so the first one matched whenever EITHER side was the CPU and every +// host<->device transfer was served by the host-to-host specialization. The +// checks below pin the exact-pair dispatch for all combinations available in +// the current build. +// --------------------------------------------------------------------------- + +TEST_F(TestModulePsiMemory, synchronize_memory_dispatch_cpu_to_cpu) +{ + std::vector h_xx(xx.size(), 0); + base_device::memory::synchronize_memory(h_xx.data(), + xx.data(), + xx.size(), + base_device::AbacusDevice_t::CpuDevice, + base_device::AbacusDevice_t::CpuDevice); + for (int ii = 0; ii < xx.size(); ii++) + { + EXPECT_EQ(h_xx[ii], xx[ii]); + } +} + +TEST_F(TestModulePsiMemory, synchronize_memory_dispatch_complex_cpu_to_cpu) +{ + std::vector> hz_xx(z_xx.size(), std::complex(0, 0)); + base_device::memory::synchronize_memory(hz_xx.data(), + z_xx.data(), + z_xx.size(), + base_device::AbacusDevice_t::CpuDevice, + base_device::AbacusDevice_t::CpuDevice); + for (int ii = 0; ii < z_xx.size(); ii++) + { + EXPECT_EQ(hz_xx[ii], z_xx[ii]); + } +} + +TEST_F(TestModulePsiMemory, cast_memory_dispatch_cpu_to_cpu) +{ + std::vector h_xx(xx.size(), 0); + base_device::memory::cast_memory(h_xx.data(), + xx.data(), + xx.size(), + base_device::AbacusDevice_t::CpuDevice, + base_device::AbacusDevice_t::CpuDevice); + for (int ii = 0; ii < xx.size(); ii++) + { + EXPECT_FLOAT_EQ(h_xx[ii], static_cast(xx[ii])); + } +} + #if __UT_USE_CUDA || __UT_USE_ROCM TEST_F(TestModulePsiMemory, set_memory_op_double_gpu) { @@ -347,4 +400,90 @@ TEST_F(TestModulePsiMemory, delete_memory_op_complex_double_gpu) delete_memory_complex_double_gpu_op()(thrust::raw_pointer_cast(dz_xx)); } + +// Exact-pair dispatch across the host/device boundary (issue #7553). Before the +// fix these three cases all reached synchronize_memory_op<..., CPU, CPU>, i.e. a +// plain host memcpy on a device pointer. +TEST_F(TestModulePsiMemory, synchronize_memory_dispatch_cpu_to_gpu) +{ + thrust::device_ptr d_xx = thrust::device_malloc(xx.size()); + std::vector hv_xx(xx.size(), 0); + thrust::copy(hv_xx.begin(), hv_xx.end(), d_xx); + base_device::memory::synchronize_memory(thrust::raw_pointer_cast(d_xx), + xx.data(), + xx.size(), + base_device::AbacusDevice_t::GpuDevice, + base_device::AbacusDevice_t::CpuDevice); + + thrust::host_vector h_xx(xx.size()); + thrust::copy(d_xx, d_xx + xx.size(), h_xx.begin()); + for (int ii = 0; ii < xx.size(); ii++) + { + EXPECT_EQ(h_xx[ii], xx[ii]); + } + thrust::device_free(d_xx); +} + +TEST_F(TestModulePsiMemory, synchronize_memory_dispatch_gpu_to_cpu) +{ + thrust::device_ptr d_xx = thrust::device_malloc(xx.size()); + thrust::copy(xx.begin(), xx.end(), d_xx); + thrust::host_vector h_xx(xx.size()); + base_device::memory::synchronize_memory(thrust::raw_pointer_cast(h_xx.data()), + thrust::raw_pointer_cast(d_xx), + xx.size(), + base_device::AbacusDevice_t::CpuDevice, + base_device::AbacusDevice_t::GpuDevice); + + for (int ii = 0; ii < xx.size(); ii++) + { + EXPECT_EQ(h_xx[ii], xx[ii]); + } + thrust::device_free(d_xx); +} + +TEST_F(TestModulePsiMemory, synchronize_memory_dispatch_gpu_to_gpu) +{ + thrust::device_ptr d1_xx = thrust::device_malloc(xx.size()); + thrust::device_ptr d2_xx = thrust::device_malloc(xx.size()); + thrust::copy(xx.begin(), xx.end(), d1_xx); + base_device::memory::synchronize_memory(thrust::raw_pointer_cast(d2_xx), + thrust::raw_pointer_cast(d1_xx), + xx.size(), + base_device::AbacusDevice_t::GpuDevice, + base_device::AbacusDevice_t::GpuDevice); + + thrust::host_vector h_xx(xx.size()); + thrust::copy(d2_xx, d2_xx + xx.size(), h_xx.begin()); + for (int ii = 0; ii < xx.size(); ii++) + { + EXPECT_EQ(h_xx[ii], xx[ii]); + } + thrust::device_free(d1_xx); + thrust::device_free(d2_xx); +} + +TEST_F(TestModulePsiMemory, cast_memory_dispatch_cpu_to_gpu_and_back) +{ + thrust::device_ptr d_xx = thrust::device_malloc(xx.size()); + base_device::memory::cast_memory(thrust::raw_pointer_cast(d_xx), + xx.data(), + xx.size(), + base_device::AbacusDevice_t::GpuDevice, + base_device::AbacusDevice_t::CpuDevice); + + std::vector h_xx(xx.size(), 0); + base_device::memory::cast_memory(h_xx.data(), + thrust::raw_pointer_cast(d_xx), + xx.size(), + base_device::AbacusDevice_t::CpuDevice, + base_device::AbacusDevice_t::GpuDevice); + + for (int ii = 0; ii < xx.size(); ii++) + { + EXPECT_FLOAT_EQ(static_cast(h_xx[ii]), static_cast(xx[ii])); + } + thrust::device_free(d_xx); +} + #endif // __UT_USE_CUDA || __UT_USE_ROCM