Skip to content

Fix: dispatch device memory wrappers by exact source/destination device pair #7941

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

Merged
mohanchen merged 1 commit into from
Sep 9, 2026
Merged

Fix: dispatch device memory wrappers by exact source/destination device pair #7941

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
53 changes: 45 additions & 8 deletions source/source_base/module_device/memory_op.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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"
Expand Down Expand Up @@ -549,35 +550,52 @@ void set_memory(FPTYPE* arr, const int var, const size_t size, base_device::Abac

template <typename FPTYPE>
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<FPTYPE, DEVICE_CPU, DEVICE_CPU>()(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<FPTYPE, DEVICE_CPU, DEVICE_GPU>()(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<FPTYPE, DEVICE_GPU, DEVICE_CPU>()(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<FPTYPE, DEVICE_GPU, DEVICE_GPU>()(arr_out, arr_in, size);
}
#endif
else {
ModuleBase::WARNING_QUIT("base_device::memory::synchronize_memory",
"unsupported source/destination device combination");
}
}

template <typename FPTYPE_out, typename FPTYPE_in>
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<FPTYPE_out, FPTYPE_in, DEVICE_CPU, DEVICE_CPU>()(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<FPTYPE_out, FPTYPE_in, DEVICE_CPU, DEVICE_GPU>()(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<FPTYPE_out, FPTYPE_in, DEVICE_GPU, DEVICE_CPU>()(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<FPTYPE_out, FPTYPE_in, DEVICE_GPU, DEVICE_GPU>()(arr_out, arr_in, size);
}
#endif
else {
ModuleBase::WARNING_QUIT("base_device::memory::cast_memory",
"unsupported source/destination device combination");
}
}

template <typename FPTYPE>
Expand All @@ -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>(int*, const int*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t);
template void synchronize_memory<float>(float*, const float*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t);
template void synchronize_memory<double>(double*, const double*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t);
template void synchronize_memory<std::complex<float>>(std::complex<float>*, const std::complex<float>*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t);
template void synchronize_memory<std::complex<double>>(std::complex<double>*, const std::complex<double>*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t);

template void cast_memory<float, float>(float*, const float*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t);
template void cast_memory<double, double>(double*, const double*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t);
template void cast_memory<float, double>(float*, const double*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t);
template void cast_memory<double, float>(double*, const float*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t);
template void cast_memory<std::complex<float>, std::complex<float>>(std::complex<float>*, const std::complex<float>*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t);
template void cast_memory<std::complex<double>, std::complex<double>>(std::complex<double>*, const std::complex<double>*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t);
template void cast_memory<std::complex<float>, std::complex<double>>(std::complex<float>*, const std::complex<double>*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t);
template void cast_memory<std::complex<double>, std::complex<float>>(std::complex<double>*, const std::complex<float>*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t);

} // namespace memory
} // namespace base_device
139 changes: 139 additions & 0 deletions source/source_base/module_device/test/memory_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<double> 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<std::complex<double>> hz_xx(z_xx.size(), std::complex<double>(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<float> 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<float>(xx[ii]));
}
}

#if __UT_USE_CUDA || __UT_USE_ROCM
TEST_F(TestModulePsiMemory, set_memory_op_double_gpu)
{
Expand Down Expand Up @@ -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<double> d_xx = thrust::device_malloc<double>(xx.size());
std::vector<double> 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<double> 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<double> d_xx = thrust::device_malloc<double>(xx.size());
thrust::copy(xx.begin(), xx.end(), d_xx);
thrust::host_vector<double> 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<double> d1_xx = thrust::device_malloc<double>(xx.size());
thrust::device_ptr<double> d2_xx = thrust::device_malloc<double>(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<double> 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<float> d_xx = thrust::device_malloc<float>(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<double> 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<float>(h_xx[ii]), static_cast<float>(xx[ii]));
}
thrust::device_free(d_xx);
}

#endif // __UT_USE_CUDA || __UT_USE_ROCM
Loading