fix(pd): handle border scalars and CPU tensors

source/op/pd/comm.cc

@@ -12,6 +12,29 @@
 
 #if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM)
 #include "device.h"
+
+template <typename FPTYPE>
+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<int>());
 
   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<int>());
 
@@ -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));

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]))