fix(pd): handle border scalars and CPU tensors #5832
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,91 @@ | ||
| # 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. | ||
|
|
||
| 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") | ||
|
Collaborator
There was a problem hiding this comment. This line does what I asked, but the job it is written for does not exist. The reply says "the CUDA Paddle CI job will execute the operator path". There is no such job. Meanwhile
So there is no configuration in which the function this PR changes is both compiled and executed. I checked locally too -- paddle 3.3.1, I owe a correction on my own earlier comment: I wrote "on the CUDA job (the only one that compiles it)", which implied such a job existed and that the only problem was tensor placement. That was wrong. The placement fix was necessary but it cannot be sufficient while no job builds Paddle with CUDA. Nothing here is yours to fix in this PR, and I am not asking you to. But it does mean the regression cannot fail on unpatched code anywhere today, so it documents the intended behaviour rather than protecting it. Two things that would change that, either of them separate work:
If neither is on the cards soon, a one-line comment in the test saying it is currently unexecuted in CI would stop the next reader from assuming the green tick covers it.
Contributor
Author
There was a problem hiding this comment. Agreed on all counts — thank you for the correction about no CUDA Paddle job existing. I added a note to the test declaring that no pipeline currently compiles-and-executes Coding agent: opencode |
||
| 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( | ||
|
njzjz-bot marked this conversation as resolved.
|
||
| [[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])) | ||
Uh oh!
There was an error while loading. Please reload this page.