Skip to content

Commit 6276995

Browse files
Critsium-xyclaude
Critsium-xy
and
claude
committed
module_device: dispatch memory wrappers by exact device pair
`synchronize_memory()` and `cast_memory()` select a compile-time specialization from a pair of runtime `AbacusDevice_t` values, but joined their branch conditions with `||`. The first branch therefore matched whenever *either* side was the CPU, so a CPU->GPU or GPU->CPU transfer was served by the CPU-to-CPU specialization, i.e. a plain host `memcpy`/cast on a device pointer. GPU->GPU could likewise match an earlier mixed-device branch. The four combinations are mutually exclusive and must be tested with `&&`. The GPU branches are now compiled only when a GPU backend is enabled, which matches the guard on the `*_op` GPU specializations in memory_op.h, and an unsupported device combination fails loudly instead of falling through silently. Both wrappers are declared in memory_op.h but were defined in the .cpp with no explicit instantiation, so they could not be linked from another translation unit and had no call sites -- which is why the defect was latent. Add the instantiations (for `cast_memory`, only the type pairs that `cast_memory_op` provides for all four device combinations) and unit tests that pin the dispatch: CPU-to-CPU everywhere, and all four pairs under __UT_USE_CUDA/__UT_USE_ROCM. No documentation change: this is an internal dispatch fix with no INPUT, output or user-visible behaviour change. Fixes #7553 Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
1 parent 3937023 commit 6276995

2 files changed

Lines changed: 184 additions & 8 deletions

File tree

‎source/source_base/module_device/memory_op.cpp‎

Lines changed: 45 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
#include "memory_op.h"
22

33
#include "source_base/memory_recorder.h"
4+
#include "source_base/tool_quit.h"
45
#include "source_base/tool_threading.h"
56
#ifdef __DSP
67
#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
549550

550551
template <typename FPTYPE>
551552
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){
552-
if (device_type_out == base_device::AbacusDevice_t::CpuDevice || device_type_in == base_device::AbacusDevice_t::CpuDevice){
553+
// The four source/destination combinations are mutually exclusive, so each
554+
// branch must test BOTH devices. Using `||` here made the first branch match
555+
// whenever either side was the CPU, which routed host<->device transfers to
556+
// the host-to-host specialization.
557+
if (device_type_out == base_device::AbacusDevice_t::CpuDevice && device_type_in == base_device::AbacusDevice_t::CpuDevice){
553558
synchronize_memory_op<FPTYPE, DEVICE_CPU, DEVICE_CPU>()(arr_out, arr_in, size);
554559
}
555-
else if (device_type_out == base_device::AbacusDevice_t::CpuDevice || device_type_in == base_device::AbacusDevice_t::GpuDevice){
560+
#if __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM
561+
else if (device_type_out == base_device::AbacusDevice_t::CpuDevice && device_type_in == base_device::AbacusDevice_t::GpuDevice){
556562
synchronize_memory_op<FPTYPE, DEVICE_CPU, DEVICE_GPU>()(arr_out, arr_in, size);
557563
}
558-
else if (device_type_out == base_device::AbacusDevice_t::GpuDevice || device_type_in == base_device::AbacusDevice_t::CpuDevice){
564+
else if (device_type_out == base_device::AbacusDevice_t::GpuDevice && device_type_in == base_device::AbacusDevice_t::CpuDevice){
559565
synchronize_memory_op<FPTYPE, DEVICE_GPU, DEVICE_CPU>()(arr_out, arr_in, size);
560566
}
561-
else if (device_type_out == base_device::AbacusDevice_t::GpuDevice || device_type_in == base_device::AbacusDevice_t::GpuDevice){
567+
else if (device_type_out == base_device::AbacusDevice_t::GpuDevice && device_type_in == base_device::AbacusDevice_t::GpuDevice){
562568
synchronize_memory_op<FPTYPE, DEVICE_GPU, DEVICE_GPU>()(arr_out, arr_in, size);
563569
}
570+
#endif
571+
else {
572+
ModuleBase::WARNING_QUIT("base_device::memory::synchronize_memory",
573+
"unsupported source/destination device combination");
574+
}
564575
}
565576

566577
template <typename FPTYPE_out, typename FPTYPE_in>
567578
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)
568579
{
569-
if (device_type_out == base_device::AbacusDevice_t::CpuDevice || device_type_in == base_device::AbacusDevice_t::CpuDevice){
580+
// See synchronize_memory() above: dispatch on the exact (out, in) device pair.
581+
if (device_type_out == base_device::AbacusDevice_t::CpuDevice && device_type_in == base_device::AbacusDevice_t::CpuDevice){
570582
cast_memory_op<FPTYPE_out, FPTYPE_in, DEVICE_CPU, DEVICE_CPU>()(arr_out, arr_in, size);
571583
}
572-
else if (device_type_out == base_device::AbacusDevice_t::CpuDevice || device_type_in == base_device::AbacusDevice_t::GpuDevice){
584+
#if __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM
585+
else if (device_type_out == base_device::AbacusDevice_t::CpuDevice && device_type_in == base_device::AbacusDevice_t::GpuDevice){
573586
cast_memory_op<FPTYPE_out, FPTYPE_in, DEVICE_CPU, DEVICE_GPU>()(arr_out, arr_in, size);
574587
}
575-
else if (device_type_out == base_device::AbacusDevice_t::GpuDevice || device_type_in == base_device::AbacusDevice_t::CpuDevice){
588+
else if (device_type_out == base_device::AbacusDevice_t::GpuDevice && device_type_in == base_device::AbacusDevice_t::CpuDevice){
576589
cast_memory_op<FPTYPE_out, FPTYPE_in, DEVICE_GPU, DEVICE_CPU>()(arr_out, arr_in, size);
577590
}
578-
else if (device_type_out == base_device::AbacusDevice_t::GpuDevice || device_type_in == base_device::AbacusDevice_t::GpuDevice){
591+
else if (device_type_out == base_device::AbacusDevice_t::GpuDevice && device_type_in == base_device::AbacusDevice_t::GpuDevice){
579592
cast_memory_op<FPTYPE_out, FPTYPE_in, DEVICE_GPU, DEVICE_GPU>()(arr_out, arr_in, size);
580593
}
594+
#endif
595+
else {
596+
ModuleBase::WARNING_QUIT("base_device::memory::cast_memory",
597+
"unsupported source/destination device combination");
598+
}
581599
}
582600

583601
template <typename FPTYPE>
@@ -591,5 +609,24 @@ void delete_memory(FPTYPE* arr, base_device::AbacusDevice_t device_type)
591609
}
592610
}
593611

612+
// Explicit instantiations of the runtime-dispatch wrappers, so that the
613+
// declarations in memory_op.h can actually be linked from another translation
614+
// unit (and covered by unit tests). cast_memory is instantiated only for the
615+
// type pairs that cast_memory_op provides for all four device combinations.
616+
template void synchronize_memory<int>(int*, const int*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t);
617+
template void synchronize_memory<float>(float*, const float*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t);
618+
template void synchronize_memory<double>(double*, const double*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t);
619+
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);
620+
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);
621+
622+
template void cast_memory<float, float>(float*, const float*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t);
623+
template void cast_memory<double, double>(double*, const double*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t);
624+
template void cast_memory<float, double>(float*, const double*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t);
625+
template void cast_memory<double, float>(double*, const float*, const size_t, base_device::AbacusDevice_t, base_device::AbacusDevice_t);
626+
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);
627+
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);
628+
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);
629+
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);
630+
594631
} // namespace memory
595632
} // namespace base_device

‎source/source_base/module_device/test/memory_test.cpp‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -170,6 +170,59 @@ TEST_F(TestModulePsiMemory, delete_memory_op_complex_double_cpu)
170170
delete_memory_complex_double_cpu_op()(hz_xx);
171171
}
172172

173+
// ---------------------------------------------------------------------------
174+
// Runtime-dispatch wrappers (issue #7553).
175+
//
176+
// synchronize_memory()/cast_memory() pick a compile-time specialization from a
177+
// pair of runtime AbacusDevice_t values. The branches used to be joined with
178+
// `||`, so the first one matched whenever EITHER side was the CPU and every
179+
// host<->device transfer was served by the host-to-host specialization. The
180+
// checks below pin the exact-pair dispatch for all combinations available in
181+
// the current build.
182+
// ---------------------------------------------------------------------------
183+
184+
TEST_F(TestModulePsiMemory, synchronize_memory_dispatch_cpu_to_cpu)
185+
{
186+
std::vector<double> h_xx(xx.size(), 0);
187+
base_device::memory::synchronize_memory(h_xx.data(),
188+
xx.data(),
189+
xx.size(),
190+
base_device::AbacusDevice_t::CpuDevice,
191+
base_device::AbacusDevice_t::CpuDevice);
192+
for (int ii = 0; ii < xx.size(); ii++)
193+
{
194+
EXPECT_EQ(h_xx[ii], xx[ii]);
195+
}
196+
}
197+
198+
TEST_F(TestModulePsiMemory, synchronize_memory_dispatch_complex_cpu_to_cpu)
199+
{
200+
std::vector<std::complex<double>> hz_xx(z_xx.size(), std::complex<double>(0, 0));
201+
base_device::memory::synchronize_memory(hz_xx.data(),
202+
z_xx.data(),
203+
z_xx.size(),
204+
base_device::AbacusDevice_t::CpuDevice,
205+
base_device::AbacusDevice_t::CpuDevice);
206+
for (int ii = 0; ii < z_xx.size(); ii++)
207+
{
208+
EXPECT_EQ(hz_xx[ii], z_xx[ii]);
209+
}
210+
}
211+
212+
TEST_F(TestModulePsiMemory, cast_memory_dispatch_cpu_to_cpu)
213+
{
214+
std::vector<float> h_xx(xx.size(), 0);
215+
base_device::memory::cast_memory(h_xx.data(),
216+
xx.data(),
217+
xx.size(),
218+
base_device::AbacusDevice_t::CpuDevice,
219+
base_device::AbacusDevice_t::CpuDevice);
220+
for (int ii = 0; ii < xx.size(); ii++)
221+
{
222+
EXPECT_FLOAT_EQ(h_xx[ii], static_cast<float>(xx[ii]));
223+
}
224+
}
225+
173226
#if __UT_USE_CUDA || __UT_USE_ROCM
174227
TEST_F(TestModulePsiMemory, set_memory_op_double_gpu)
175228
{
@@ -347,4 +400,90 @@ TEST_F(TestModulePsiMemory, delete_memory_op_complex_double_gpu)
347400
delete_memory_complex_double_gpu_op()(thrust::raw_pointer_cast(dz_xx));
348401
}
349402

403+
404+
// Exact-pair dispatch across the host/device boundary (issue #7553). Before the
405+
// fix these three cases all reached synchronize_memory_op<..., CPU, CPU>, i.e. a
406+
// plain host memcpy on a device pointer.
407+
TEST_F(TestModulePsiMemory, synchronize_memory_dispatch_cpu_to_gpu)
408+
{
409+
thrust::device_ptr<double> d_xx = thrust::device_malloc<double>(xx.size());
410+
std::vector<double> hv_xx(xx.size(), 0);
411+
thrust::copy(hv_xx.begin(), hv_xx.end(), d_xx);
412+
base_device::memory::synchronize_memory(thrust::raw_pointer_cast(d_xx),
413+
xx.data(),
414+
xx.size(),
415+
base_device::AbacusDevice_t::GpuDevice,
416+
base_device::AbacusDevice_t::CpuDevice);
417+
418+
thrust::host_vector<double> h_xx(xx.size());
419+
thrust::copy(d_xx, d_xx + xx.size(), h_xx.begin());
420+
for (int ii = 0; ii < xx.size(); ii++)
421+
{
422+
EXPECT_EQ(h_xx[ii], xx[ii]);
423+
}
424+
thrust::device_free(d_xx);
425+
}
426+
427+
TEST_F(TestModulePsiMemory, synchronize_memory_dispatch_gpu_to_cpu)
428+
{
429+
thrust::device_ptr<double> d_xx = thrust::device_malloc<double>(xx.size());
430+
thrust::copy(xx.begin(), xx.end(), d_xx);
431+
thrust::host_vector<double> h_xx(xx.size());
432+
base_device::memory::synchronize_memory(thrust::raw_pointer_cast(h_xx.data()),
433+
thrust::raw_pointer_cast(d_xx),
434+
xx.size(),
435+
base_device::AbacusDevice_t::CpuDevice,
436+
base_device::AbacusDevice_t::GpuDevice);
437+
438+
for (int ii = 0; ii < xx.size(); ii++)
439+
{
440+
EXPECT_EQ(h_xx[ii], xx[ii]);
441+
}
442+
thrust::device_free(d_xx);
443+
}
444+
445+
TEST_F(TestModulePsiMemory, synchronize_memory_dispatch_gpu_to_gpu)
446+
{
447+
thrust::device_ptr<double> d1_xx = thrust::device_malloc<double>(xx.size());
448+
thrust::device_ptr<double> d2_xx = thrust::device_malloc<double>(xx.size());
449+
thrust::copy(xx.begin(), xx.end(), d1_xx);
450+
base_device::memory::synchronize_memory(thrust::raw_pointer_cast(d2_xx),
451+
thrust::raw_pointer_cast(d1_xx),
452+
xx.size(),
453+
base_device::AbacusDevice_t::GpuDevice,
454+
base_device::AbacusDevice_t::GpuDevice);
455+
456+
thrust::host_vector<double> h_xx(xx.size());
457+
thrust::copy(d2_xx, d2_xx + xx.size(), h_xx.begin());
458+
for (int ii = 0; ii < xx.size(); ii++)
459+
{
460+
EXPECT_EQ(h_xx[ii], xx[ii]);
461+
}
462+
thrust::device_free(d1_xx);
463+
thrust::device_free(d2_xx);
464+
}
465+
466+
TEST_F(TestModulePsiMemory, cast_memory_dispatch_cpu_to_gpu_and_back)
467+
{
468+
thrust::device_ptr<float> d_xx = thrust::device_malloc<float>(xx.size());
469+
base_device::memory::cast_memory(thrust::raw_pointer_cast(d_xx),
470+
xx.data(),
471+
xx.size(),
472+
base_device::AbacusDevice_t::GpuDevice,
473+
base_device::AbacusDevice_t::CpuDevice);
474+
475+
std::vector<double> h_xx(xx.size(), 0);
476+
base_device::memory::cast_memory(h_xx.data(),
477+
thrust::raw_pointer_cast(d_xx),
478+
xx.size(),
479+
base_device::AbacusDevice_t::CpuDevice,
480+
base_device::AbacusDevice_t::GpuDevice);
481+
482+
for (int ii = 0; ii < xx.size(); ii++)
483+
{
484+
EXPECT_FLOAT_EQ(static_cast<float>(h_xx[ii]), static_cast<float>(xx[ii]));
485+
}
486+
thrust::device_free(d_xx);
487+
}
488+
350489
#endif // __UT_USE_CUDA || __UT_USE_ROCM

0 commit comments

Comments
 (0)