diff --git a/source/source_base/module_device/memory_op.cpp b/source/source_base/module_device/memory_op.cpp index 2c38e44145e..a4e6c059a3d 100644 --- a/source/source_base/module_device/memory_op.cpp +++ b/source/source_base/module_device/memory_op.cpp @@ -1,7 +1,6 @@ #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" @@ -527,106 +526,5 @@ template struct delete_memory_op_mt, base_device::DEVICE_CPU template struct delete_memory_op_mt, base_device::DEVICE_CPU>; #endif -template -void resize_memory(FPTYPE* arr, const size_t size, base_device::AbacusDevice_t device_type) -{ - if (device_type == base_device::AbacusDevice_t::CpuDevice){ - resize_memory_op()(arr, size); - } - else if (device_type == base_device::AbacusDevice_t::GpuDevice){ - resize_memory_op()(arr, size); - } -} - -template -void set_memory(FPTYPE* arr, const int var, const size_t size, base_device::AbacusDevice_t device_type){ - if (device_type == base_device::AbacusDevice_t::CpuDevice){ - set_memory_op()(arr, var, size); - } - else if (device_type == base_device::AbacusDevice_t::GpuDevice){ - set_memory_op()(arr, var, size); - } -} - -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){ - // 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); - } -#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){ - 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){ - 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) -{ - // 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); - } -#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){ - 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){ - cast_memory_op()(arr_out, arr_in, size); - } -#endif - else { - ModuleBase::WARNING_QUIT("base_device::memory::cast_memory", - "unsupported source/destination device combination"); - } -} - -template -void delete_memory(FPTYPE* arr, base_device::AbacusDevice_t device_type) -{ - if (device_type == base_device::AbacusDevice_t::CpuDevice){ - delete_memory_op()(arr); - } - else if (device_type == base_device::AbacusDevice_t::GpuDevice){ - delete_memory_op()(arr); - } -} - -// 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/memory_op.h b/source/source_base/module_device/memory_op.h index 5e292c57a65..ef8863fb92c 100644 --- a/source/source_base/module_device/memory_op.h +++ b/source/source_base/module_device/memory_op.h @@ -120,21 +120,6 @@ struct delete_memory_op void operator()(FPTYPE* arr); }; -template -void resize_memory(FPTYPE* arr, const size_t size, base_device::AbacusDevice_t device_type = base_device::AbacusDevice_t::CpuDevice); - -template -void set_memory(FPTYPE* arr, const int var, const size_t size, base_device::AbacusDevice_t device_type = base_device::AbacusDevice_t::CpuDevice); - -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); - -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); - -template -void delete_memory(FPTYPE* arr, base_device::AbacusDevice_t device_type = base_device::AbacusDevice_t::CpuDevice); - #if __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM // Partially specialize operator for base_device::GpuDevice. template diff --git a/source/source_base/module_device/test/memory_test.cpp b/source/source_base/module_device/test/memory_test.cpp index a6912beffe3..a963c68c532 100644 --- a/source/source_base/module_device/test/memory_test.cpp +++ b/source/source_base/module_device/test/memory_test.cpp @@ -170,59 +170,6 @@ 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) { @@ -400,90 +347,4 @@ 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