From 00eccf42b1d49788d41e706a4c62325af059c893 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 16 Jul 2026 13:51:49 +0800 Subject: [PATCH 1/5] 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/5] 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/5] 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") From e99f738d4843377628685546c64a19a7c71c0524 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Tue, 11 Aug 2026 02:44:42 +0800 Subject: [PATCH 4/5] fix(pd): return accumulated border gradients Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/op/pd/comm.cc | 34 +++++++++++++++++++++++-------- source/tests/pd/test_border_op.py | 3 ++- 2 files changed, 27 insertions(+), 10 deletions(-) diff --git a/source/op/pd/comm.cc b/source/op/pd/comm.cc index 1dbfbfa356..72d3f4d7bb 100644 --- a/source/op/pd/comm.cc +++ b/source/op/pd/comm.cc @@ -331,7 +331,6 @@ void Border_backward_t(const paddle::Tensor& sendlist_tensor, cpu_recvnum.copy_(recvnum_tensor, paddle::CPUPlace(), true); int* sendnum = cpu_recvnum.data(); - FPTYPE* local_g1 = d_local_g1_tensor.data(); int tensor_size = d_local_g1_tensor.dims()[1]; paddle::Tensor cpu_nlocal = @@ -410,18 +409,35 @@ void Border_backward_t(const paddle::Tensor& sendlist_tensor, d_local_g1_tensor, irecvlist, recv_g1_tensor.slice(0, nrecv), 0); } } -#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) - gpuDeviceSynchronize(); -#endif -#ifdef USE_MPI + // Forward swaps overwrite every ghost row with owner data. Reverse + // communication accumulates each ghost output gradient into its owner, but + // the original ghost input has no path to the output and therefore must + // receive a zero gradient. With no swaps, forward is the identity and the + // upstream ghost gradient remains valid. + if (nswap > 0 && nghost > 0) { + FPTYPE* ghost_g1 = + d_local_g1_tensor.data() + nlocal * tensor_size; + size_t ghost_size = (size_t)nghost * tensor_size; #if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) - if (cuda_aware == 0) { - recv_g1_tensor_grad.copy_(d_local_g1_tensor, recv_g1_tensor_grad.place(), - true); - } + if (phi::is_gpu_place(d_local_g1_tensor.place())) { + gpuMemset(ghost_g1, 0, ghost_size * sizeof(FPTYPE)); + } else { + memset(ghost_g1, 0, ghost_size * sizeof(FPTYPE)); + } +#else + memset(ghost_g1, 0, ghost_size * sizeof(FPTYPE)); #endif + } +#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) + gpuDeviceSynchronize(); #endif + + // The reverse exchange may update a temporary on the original device or on + // the CPU fallback used by non-CUDA-aware MPI. Always publish that result to + // the custom-op gradient output instead of relying on incidental aliasing. + recv_g1_tensor_grad.copy_(d_local_g1_tensor, recv_g1_tensor_grad.place(), + true); } void Border_backward(const paddle::Tensor& sendlist_tensor, diff --git a/source/tests/pd/test_border_op.py b/source/tests/pd/test_border_op.py index d30e8c2ce7..adf0242cf4 100644 --- a/source/tests/pd/test_border_op.py +++ b/source/tests/pd/test_border_op.py @@ -88,4 +88,5 @@ def test_border_op_self_copy_uses_cpu_place() -> None: # 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])) + expected_grad = np.array([[1.0, 1.0], [2.0, 2.0], [0.0, 0.0]]) + np.testing.assert_array_equal(g1_leaf.grad.numpy(), expected_grad) From d387e4fd71e0e5b355ef142bee2206ec570e578b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 10 Aug 2026 18:45:43 +0000 Subject: [PATCH 5/5] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/op/pd/comm.cc | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/source/op/pd/comm.cc b/source/op/pd/comm.cc index 72d3f4d7bb..e5f0afd88f 100644 --- a/source/op/pd/comm.cc +++ b/source/op/pd/comm.cc @@ -416,8 +416,7 @@ void Border_backward_t(const paddle::Tensor& sendlist_tensor, // receive a zero gradient. With no swaps, forward is the identity and the // upstream ghost gradient remains valid. if (nswap > 0 && nghost > 0) { - FPTYPE* ghost_g1 = - d_local_g1_tensor.data() + nlocal * tensor_size; + FPTYPE* ghost_g1 = d_local_g1_tensor.data() + nlocal * tensor_size; size_t ghost_size = (size_t)nghost * tensor_size; #if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) if (phi::is_gpu_place(d_local_g1_tensor.place())) {