fix(c++): fix NULL type in custom op#4889
fix(c++): fix NULL type in custom op#4889wanghan-iapcm merged 26 commits intodeepmodeling:masterfrom
Conversation
Replaces usage of lmp_list send/recv arrays with new vectors that map indices using fwd_map and synchronize counts via MPI. Updates tensor construction to use these new vectors, improving correctness and flexibility in distributed communication.
for more information, see https://pre-commit.ci
📝 WalkthroughWalkthroughCalls a new helper to remap LAMMPS sendlists using a forward-map before constructing MPI/send tensors in DeepPotPT/DeepPotPD/DeepSpinPT, and adds the helper declaration/implementation and a small test. No public API signature changes beyond the new helper. Changes
Sequence Diagram(s)sequenceDiagram
participant Compute as DeepPot*::compute
participant Select as select_real_atoms_sendlist
participant Map as fwd_map
participant Tensors as Tensor & comm_dict build
Compute->>Select: call select_real_atoms_sendlist(lmp_list, fwd_map)
Select->>Map: map each send index -> forwarded index
Map-->>Select: remapped sendlist (invalids removed)
Select-->>Compute: lmp_list updated (sendlist/sendnum/recvnum)
Compute->>Tensors: build send/recv tensors and comm_dict from updated lmp_list
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested reviewers
✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
source/api_cc/src/DeepPotPT.cc (2)
251-257: Remove stale commented-out codeDead commented code obscures the current data path and makes maintenance harder.
Apply this diff:
-// torch::Tensor firstrecv_tensor = -// torch::from_blob(lmp_list.firstrecv, {nswap}, int32_option); -// torch::Tensor recvnum_tensor = -// torch::from_blob(lmp_list.recvnum, {nswap}, int32_option); -// torch::Tensor sendnum_tensor = -// torch::from_blob(lmp_list.sendnum, {nswap}, int32_option);
266-269: Remove redundant commented-out legacy codeSame reasoning; commented legacy path is preserved in git history.
Apply this diff:
-// int total_send = -// std::accumulate(lmp_list.sendnum, lmp_list.sendnum + nswap, 0); -// torch::Tensor sendlist_tensor = -// torch::from_blob(lmp_list.sendlist, {total_send}, int32_option);
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these settings in your CodeRabbit configuration.
📒 Files selected for processing (1)
source/api_cc/src/DeepPotPT.cc(3 hunks)
🔇 Additional comments (2)
source/api_cc/src/DeepPotPT.cc (2)
185-204: Remapping logic LGTMCorrectly rebuilds per-swap send counts and a dense send list using fwd_map with bounds checks and filtering. Reserving capacity via the accumulated legacy counts is a good optimization.
226-232: firstrecv_new is unused and not required — original comment is incorrectShort: deepmd/pt/model/descriptor/repflows.py builds comm_dict and calls torch.ops.deepmd.border_op with send_list, send_proc, recv_proc, send_num, recv_num, communicator (no first_recv). The computed firstrecv_new/firstrecv_tensor in the PT wrappers is dead code — remove it or document why it is kept.
Files to update:
- source/api_cc/src/DeepPotPT.cc
- Remove the firstrecv_new prefix-sum computation (around lines 226–231) and the unused firstrecv_tensor creation (around line 238).
- source/api_cc/src/DeepSpinPT.cc
- Same pattern: firstrecv_tensor is created around lines 187–191 but never used/inserted.
Suggested change (remove unused code) — example diff for DeepPotPT.cc:
@@
std::vector<int> firstrecv_new(nswap, 0);int acc = 0;for (int s = 0; s < nswap; ++s) {firstrecv_new[s] = acc;acc += recvnum_new[s];}
/* firstrecv computation removed — not used by border_op */@@
torch::Tensor firstrecv_tensor =torch::from_blob(firstrecv_new.data(), {nswap}, int32_option).clone();
/* firstrecv tensor omitted — border_op expects recv_num, not first_recv */If you prefer to keep the computation for clarity, add a short comment explaining it's intentionally unused.
Likely an incorrect or invalid review comment.
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
Pull Request Overview
This PR fixes a bug in the C++ custom operations for distributed MPI communication by implementing proper handling of NULL types in the send/receive lists. The fix ensures that atom indices are correctly mapped and synchronized across MPI processes by filtering out invalid indices and updating counts accordingly.
Key changes:
- Added a new function to select and remap real atoms in send lists using forward mapping
- Integrated the send list filtering into PyTorch, Paddle, and DeepSpin compute functions
- Added a test case to verify NULL type handling in LAMMPS pair coefficients
Reviewed Changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| source/api_cc/src/common.cc | Implements select_real_atoms_sendlist function to filter and remap atom indices in MPI send lists |
| source/api_cc/include/common.h | Adds function declaration for the new send list selection function |
| source/api_cc/src/DeepPotPT.cc | Integrates send list filtering into PyTorch backend compute function |
| source/api_cc/src/DeepPotPD.cc | Integrates send list filtering into Paddle backend compute function |
| source/api_cc/src/DeepSpinPT.cc | Integrates send list filtering into DeepSpin PyTorch compute function |
| source/lmp/tests/test_lammps_dpa_pt.py | Adds test case for NULL type handling in LAMMPS pair coefficients |
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
source/api_cc/src/DeepSpinPT.cc (1)
205-213: Flatten thesendlistpointer-of-pointers into a contiguous buffer before callingfrom_blobInputNlist::sendlist is declared as
int **(an array of pointers to per-swap arrays) inapi_c/include/deepmd.hpp, so passing it directly totorch::from_blob(lmp_list.sendlist, {total_send}, int32_option)will interpret the array‐of‐pointers as a contiguous block of
ints, which is invalid and will lead to incorrect data being read (or a crash). You must first flatten all per-swap subarrays into a single contiguousintbuffer.Please update the following files:
source/api_cc/src/DeepSpinPT.cc(around line 207)source/api_cc/src/DeepPotPT.cc(around line 200)Proposed refactoring to safely flatten (ensuring the buffer outlives the call to
run_method):- torch::Tensor sendlist_tensor = - torch::from_blob(lmp_list.sendlist, {total_send}, int32_option); + // Flatten the int** sendlist into a single contiguous vector + std::vector<int> flat_sendlist; + flat_sendlist.reserve(total_send); + for (int s = 0; s < nswap; ++s) { + flat_sendlist.insert( + flat_sendlist.end(), + lmp_list.sendlist[s], + lmp_list.sendlist[s] + lmp_list.sendnum[s]); + } + // Create a tensor from the contiguous buffer + torch::Tensor sendlist_tensor = + torch::from_blob(flat_sendlist.data(), {total_send}, int32_option);[mend_review_comment]
source/api_cc/src/DeepPotPD.cc (1)
421-439: Shape bug: send_list should carry nswap pointer addresses, not total_send elementsYou reshape send_list to total_send but copy only nswap pointer addresses. This is a mismatched shape and can corrupt reads downstream.
- int total_send = - std::accumulate(lmp_list.sendnum, lmp_list.sendnum + nswap, 0); - sendlist_tensor->Reshape({total_send}); + // send_list carries one pointer per swap + sendlist_tensor->Reshape({nswap}); /** ** NOTE: paddle do not support construct a Tensor with from_blob(T**, ...) ** from a double pointer, so we convert int* pointer to indptr_t for each ** entry and wrap it into int64 Tensor as a workaround. */ std::vector<std::intptr_t> pointer_addresses; - pointer_addresses.reserve(nswap); + pointer_addresses.reserve(nswap);
🧹 Nitpick comments (1)
source/api_cc/include/common.h (1)
105-107: API const-correctness: this mutates inlist; avoid const on parameterThe helper writes through inlist pointers (sendnum/recvnum/sendlist). Marking the parameter const is misleading.
-void select_real_atoms_sendlist(const deepmd::InputNlist& inlist, - const std::vector<int>& fwd_map); +// In-place: trims and remaps sendlist; updates sendnum/recvnum/(firstrecv if present) +void select_real_atoms_sendlist(deepmd::InputNlist& inlist, + const std::vector<int>& fwd_map);If changing the signature ripples too far now, at least document the in-place side effects in a comment.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (6)
source/api_cc/include/common.h(1 hunks)source/api_cc/src/DeepPotPD.cc(1 hunks)source/api_cc/src/DeepPotPT.cc(1 hunks)source/api_cc/src/DeepSpinPT.cc(1 hunks)source/api_cc/src/common.cc(1 hunks)source/lmp/tests/test_lammps_dpa_pt.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- source/api_cc/src/DeepPotPT.cc
🧰 Additional context used
🧬 Code graph analysis (5)
source/lmp/tests/test_lammps_dpa_pt.py (1)
source/lmp/tests/test_lammps.py (1)
lammps_type_map(289-292)
source/api_cc/src/common.cc (1)
source/op/tf/prod_env_mat_multi_device.cc (1)
inlist(1590-1590)
source/api_cc/include/common.h (2)
source/api_cc/src/common.cc (2)
select_real_atoms_sendlist(235-258)select_real_atoms_sendlist(235-236)source/api_c/include/deepmd.hpp (4)
InputNlist(807-814)InputNlist(815-822)InputNlist(823-850)InputNlist(851-851)
source/api_cc/src/DeepSpinPT.cc (1)
source/api_cc/src/common.cc (2)
select_real_atoms_sendlist(235-258)select_real_atoms_sendlist(235-236)
source/api_cc/src/DeepPotPD.cc (1)
source/api_cc/src/common.cc (2)
select_real_atoms_sendlist(235-258)select_real_atoms_sendlist(235-236)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (27)
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build C library (2.14, >=2.5.0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test C++ (false)
- GitHub Check: Test C++ (true)
- GitHub Check: Analyze (c-cpp)
🔇 Additional comments (1)
source/api_cc/src/common.cc (1)
235-258: Recomputefirstrecvafter trimming and verify symmetryYou need to keep the receive‐offset array in sync whenever you shrink
recvnum, and also confirm that the assumptionrecvnum[s] == sendnum[s]holds across all ranks/types.• After your loop that sets
inlist.sendnum[s] = cnt; inlist.recvnum[s] = cnt;add:
+ // Recompute receive‐offsets to match updated recvnum[] + if (inlist.firstrecv) { + int offset = 0; + for (int s = 0; s < nswap; ++s) { + inlist.firstrecv[s] = offset; + offset += inlist.recvnum[s]; + } + }This keeps the
firstrecvprefix sums correct (firstrecvis “where to put 1st recv atom in each swap” ) (docs.deepmodeling.com).• Because you’re setting
recvnum[s] = sendnum[s], you’re assuming perfect send/recv symmetry per swap. Please verify this holds under all type‐map filters and MPI ranks. Consider adding a debug‐only assertion or a collective check (e.g. an MPI_Allreduce on the arrays) to catch any imbalance early.
|
It looks like this PR is pending by a false positive failed of build machine. Can you rerun those test and see if is PR can be merged? I am looking forward to this issue to be resolved. |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #4889 +/- ##
==========================================
+ Coverage 82.12% 82.16% +0.03%
==========================================
Files 753 755 +2
Lines 75825 75874 +49
Branches 3648 3660 +12
==========================================
+ Hits 62275 62342 +67
+ Misses 12382 12362 -20
- Partials 1168 1170 +2 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
|
I find the latest CI job was failed due to a test case timeout. I guess dead loop or something are trigger by this PR. Maybe you should merge or cherry pick this fix and rerun the pipeline to see which case is broken. |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
iProzd
requested review from
njzjz and
wanghan-iapcm
and removed request for
CaRoLZhangxy
for more information, see https://pre-commit.ci
|
@iProzd Great. However, the current approach requires rebuilding various atom arrays even if only a single virtual atom is present. For large-scale systems, this introduces significant performance overhead. Should we consider moving the masking logic directly into the model inference stage in the future? This would eliminate unnecessary atom index remapping and array reconstruction, leading to a more performant and cleaner implementation. |
@link89 Yes you are right. I'll open an issue to track this. |
Replaces usage of lmp_list send/recv arrays with new vectors that map indices using fwd_map and synchronize counts via MPI. Updates tensor construction to use these new vectors, improving correctness and flexibility in distributed communication.
Summary by CodeRabbit