From 00eccf42b1d49788d41e706a4c62325af059c893 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 16 Jul 2026 13:51:49 +0800 Subject: [PATCH 1/3] fix(pd): handle border scalars and CPU tensors Store nlocal and nghost in one-element host tensors even when there are no swaps, and select local forward/backward copy primitives from the actual Paddle tensor place. Add direct custom-op regressions because existing Paddle model tests did not exercise nswap == 0 or CPU self-swaps in CUDA-enabled builds. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/op/pd/comm.cc | 62 ++++++++++++------------ source/tests/pd/test_border_op.py | 80 +++++++++++++++++++++++++++++++ 2 files changed, 112 insertions(+), 30 deletions(-) create mode 100644 source/tests/pd/test_border_op.py diff --git a/source/op/pd/comm.cc b/source/op/pd/comm.cc index 548e5db83a..1dbfbfa356 100644 --- a/source/op/pd/comm.cc +++ b/source/op/pd/comm.cc @@ -12,6 +12,29 @@ #if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) #include "device.h" + +template +static void copy_local_tensor_data(FPTYPE* dst, + const FPTYPE* src, + size_t count, + const paddle::Place& place) { + if (count == 0) { + return; + } + + // CUDA-aware MPI describes whether MPI can consume device pointers; it does + // not describe where this particular Paddle tensor lives. Self-swaps must + // select the copy primitive from the actual tensor place so CPU tensors also + // work in CUDA/ROCm-enabled builds. + if (phi::is_gpu_place(place)) { + gpuMemcpy(dst, src, count * sizeof(FPTYPE), gpuMemcpyDeviceToDevice); + } else { + // CPU and host-pinned tensors are both host-addressable. Defaulting + // non-GPU places to memcpy also avoids treating a future host place as a + // CUDA pointer merely because the operator was built with CUDA support. + memcpy(dst, src, count * sizeof(FPTYPE)); + } +} #endif #ifdef USE_MPI @@ -83,13 +106,16 @@ void Border_forward_t(const paddle::Tensor& sendlist_tensor, int tensor_size = g1.dims()[1]; + // nlocal and nghost are scalar protocol values, independent of the number + // of communication swaps. In particular, nswap == 0 still needs one slot + // for each value before the host dereference below. paddle::Tensor cpu_nlocal = - paddle::empty({nswap}, paddle::DataType::INT32, paddle::CPUPlace()); + paddle::empty({1}, paddle::DataType::INT32, paddle::CPUPlace()); cpu_nlocal.copy_(nlocal_tensor, paddle::CPUPlace(), true); int nlocal = *(cpu_nlocal.data()); paddle::Tensor cpu_nghost = - paddle::empty({nswap}, paddle::DataType::INT32, paddle::CPUPlace()); + paddle::empty({1}, paddle::DataType::INT32, paddle::CPUPlace()); cpu_nghost.copy_(nghost_tensor, paddle::CPUPlace(), true); int nghost = *(cpu_nghost.data()); @@ -175,20 +201,8 @@ void Border_forward_t(const paddle::Tensor& sendlist_tensor, #endif #if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) -#ifdef USE_MPI - if (cuda_aware == 0) { - memcpy(recv_g1, send_g1, - (unsigned long)nsend * tensor_size * sizeof(FPTYPE)); - } else { - gpuMemcpy(recv_g1, send_g1, - (unsigned long)nsend * tensor_size * sizeof(FPTYPE), - gpuMemcpyDeviceToDevice); - } -#else - gpuMemcpy(recv_g1, send_g1, - (unsigned long)nsend * tensor_size * sizeof(FPTYPE), - gpuMemcpyDeviceToDevice); -#endif + copy_local_tensor_data(recv_g1, send_g1, (size_t)nsend * tensor_size, + recv_g1_tensor.place()); #else memcpy(recv_g1, send_g1, @@ -381,20 +395,8 @@ void Border_backward_t(const paddle::Tensor& sendlist_tensor, #endif if (nrecv) { #if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) -#ifdef USE_MPI - if (cuda_aware == 0) { - memcpy(recv_g1, send_g1, - (unsigned long)nrecv * tensor_size * sizeof(FPTYPE)); - } else { - gpuMemcpy(recv_g1, send_g1, - (unsigned long)nrecv * tensor_size * sizeof(FPTYPE), - gpuMemcpyDeviceToDevice); - } -#else - gpuMemcpy(recv_g1, send_g1, - (unsigned long)nrecv * tensor_size * sizeof(FPTYPE), - gpuMemcpyDeviceToDevice); -#endif + copy_local_tensor_data(recv_g1, send_g1, (size_t)nrecv * tensor_size, + d_local_g1_tensor.place()); #else memcpy(recv_g1, send_g1, (unsigned long)nrecv * tensor_size * sizeof(FPTYPE)); diff --git a/source/tests/pd/test_border_op.py b/source/tests/pd/test_border_op.py new file mode 100644 index 0000000000..cce8a64717 --- /dev/null +++ b/source/tests/pd/test_border_op.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Regression tests for Paddle border-exchange control and data tensors.""" + +import numpy as np +import paddle +import pytest + +deepmd_op_pd = pytest.importorskip( + "deepmd_op_pd", reason="the Paddle custom operator library is not built" +) + + +def _control_tensors(nswap: int) -> tuple[paddle.Tensor, ...]: + """Create the common CPU control tensors for a border exchange.""" + return ( + paddle.zeros([nswap], dtype="int32"), # sendproc + paddle.zeros([nswap], dtype="int32"), # recvproc + paddle.zeros([nswap], dtype="int32"), # sendnum + paddle.zeros([nswap], dtype="int32"), # recvnum + paddle.zeros([1], dtype="int64"), # unused communicator without MPI + ) + + +def test_border_op_accepts_no_swaps() -> None: + """Scalar atom counts must remain readable when ``nswap == 0``.""" + sendproc, recvproc, sendnum, recvnum, communicator = _control_tensors(0) + g1 = paddle.arange(6, dtype="float64").reshape([2, 3]) + + result = deepmd_op_pd.border_op( + paddle.zeros([0], dtype="int64"), + sendproc, + recvproc, + sendnum, + recvnum, + g1, + communicator, + paddle.to_tensor([2], dtype="int32"), + paddle.to_tensor([0], dtype="int32"), + ) + + np.testing.assert_array_equal(result.numpy(), g1.numpy()) + + +def test_border_op_self_copy_uses_cpu_place() -> None: + """A CUDA-enabled operator must not use a GPU copy for CPU tensors.""" + sendproc, recvproc, sendnum, recvnum, communicator = _control_tensors(1) + sendnum = paddle.ones_like(sendnum) + recvnum = paddle.ones_like(recvnum) + + # The C++ operator receives the LAMMPS send lists as pointer-valued int64 + # entries. Keep this NumPy owner alive through the call so the pointed-to + # int32 index remains valid. + send_indices = np.array([1], dtype=np.int32) + sendlist = paddle.to_tensor([send_indices.ctypes.data], dtype="int64") + g1_leaf = paddle.to_tensor( + [[1.0, 2.0], [3.0, 4.0], [0.0, 0.0]], stop_gradient=False + ) + # Paddle rejects an in-place custom op on an autograd leaf. This identity + # keeps a leaf for checking gradients while letting border_op update g1. + g1 = g1_leaf * 1.0 + + result = deepmd_op_pd.border_op( + sendlist, + sendproc, + recvproc, + sendnum, + recvnum, + g1, + communicator, + paddle.to_tensor([2], dtype="int32"), + paddle.to_tensor([1], dtype="int32"), + ) + + np.testing.assert_array_equal( + result.numpy(), np.array([[1.0, 2.0], [3.0, 4.0], [3.0, 4.0]]) + ) + # Backpropagation runs the reverse self-swap, which needs the same + # place-based CPU/GPU dispatch as the forward copy. + result.sum().backward() + np.testing.assert_array_equal(g1_leaf.grad.numpy(), np.ones([3, 2])) From 904727a61caaa7db65e2469dff64dfa76b6a6e9c Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 23 Jul 2026 20:15:25 +0800 Subject: [PATCH 2/3] test(pd): force CPU border self-copy Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/tests/pd/test_border_op.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/source/tests/pd/test_border_op.py b/source/tests/pd/test_border_op.py index cce8a64717..efa832c3fb 100644 --- a/source/tests/pd/test_border_op.py +++ b/source/tests/pd/test_border_op.py @@ -43,6 +43,9 @@ def test_border_op_accepts_no_swaps() -> None: def test_border_op_self_copy_uses_cpu_place() -> None: """A CUDA-enabled operator must not use a GPU copy for CPU tensors.""" + # CUDA Paddle builds otherwise create tensors on the default GPU, which + # would leave the operator's CPU copy branch untested. + paddle.set_device("cpu") sendproc, recvproc, sendnum, recvnum, communicator = _control_tensors(1) sendnum = paddle.ones_like(sendnum) recvnum = paddle.ones_like(recvnum) From d4843dd96cf4f388e348140cb7704af35bf2b689 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sun, 2 Aug 2026 21:58:07 +0800 Subject: [PATCH 3/3] test(pd): note CPU-branch test is not executed in CI No CI job currently builds Paddle with CUDA, so the CPU-branch regression cannot fail in any pipeline. Document that intent in the test. Coding-Agent: opencode opencode-Version: 1.18.9 Model: ustc/deepseek-v4-flash Reasoning-Effort: max --- source/tests/pd/test_border_op.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/source/tests/pd/test_border_op.py b/source/tests/pd/test_border_op.py index efa832c3fb..d30e8c2ce7 100644 --- a/source/tests/pd/test_border_op.py +++ b/source/tests/pd/test_border_op.py @@ -42,7 +42,15 @@ def test_border_op_accepts_no_swaps() -> None: def test_border_op_self_copy_uses_cpu_place() -> None: - """A CUDA-enabled operator must not use a GPU copy for CPU tensors.""" + """A CUDA-enabled operator must not use a GPU copy for CPU tensors. + + NOTE: no CI job currently builds Paddle with CUDA (``test_cuda.yml`` + disables Paddle at the workflow level and ``test_python.yml`` installs the + CPU build), so ``copy_local_tensor_data`` is not compiled-and-executed by + any pipeline. This test therefore documents the intended CPU-branch + behavior rather than guarding it; it should gain real coverage once a CI + job builds Paddle with CUDA. + """ # CUDA Paddle builds otherwise create tensors on the default GPU, which # would leave the operator's CPU copy branch untested. paddle.set_device("cpu")