perf: use contiguous memory stride for edge/angle indices #4804
Conversation
for more information, see https://pre-commit.ci
📝 Walkthrough""" WalkthroughThe changes standardize the shape and indexing conventions of Changes
Sequence Diagram(s)sequenceDiagram
participant Caller
participant get_graph_index
Note over Caller,get_graph_index: Applies to both NumPy and PyTorch versions
Caller->>get_graph_index: Call with graph data
get_graph_index->>get_graph_index: Compute index arrays (row, col, etc.)
get_graph_index->>get_graph_index: Stack indices along axis=0 (column-major)
get_graph_index-->>Caller: Return edge_index (2, n_edge), angle_index (3, n_angle)
sequenceDiagram
participant Model
participant RepFlowLayer
participant DescrptBlockRepflows
Model->>RepFlowLayer: Call with edge_index (2, n_edge), angle_index (3, n_angle)
RepFlowLayer->>RepFlowLayer: Unpack indices using row slices [0], [1], [2]
RepFlowLayer-->>Model: Return processed features
Model->>DescrptBlockRepflows: Call with edge_index, angle_index
DescrptBlockRepflows->>DescrptBlockRepflows: Use edge_index[0] for owner
DescrptBlockRepflows-->>Model: Return result
Suggested reviewers
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 Pylint (3.3.7)deepmd/pt/model/descriptor/repflow_layer.pyNo files to lint: exiting. 📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
⏰ Context from checks skipped due to timeout of 90000ms (22)
✨ Finishing Touches
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. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (5)
deepmd/pt/model/network/utils.py (1)
138-142: TorchScript: ensure constant first-dim remains contiguous
torch.stack([n2e_index, n_ext2e_index], dim=0)creates a fresh tensor which is notcontiguous()along the flattened dimension.
Down-stream JIT kernels that later call.view(-1)(as in_cal_hg_dynamic) may hit a run-time assert if the tensor is used directly.-edge_index_result = torch.stack([n2e_index, n_ext2e_index], dim=0) +edge_index_result = torch.stack([n2e_index, n_ext2e_index], dim=0).contiguous()Do the same for
angle_index_result. Cheap no-op for already-contiguous cases, defensive for others.deepmd/pt/model/descriptor/repflows.py (1)
540-542: Minor: reuse tensor factory & avoid dtype drift
torch.zeros([2,1], device=nlist.device, dtype=nlist.dtype)works, but we can avoid repeatingdevice/dtypearguments and stay robust ifnlistswitches totorch.int64.-edge_index = torch.zeros([2, 1], device=nlist.device, dtype=nlist.dtype) -angle_index = torch.zeros([3, 1], device=nlist.device, dtype=nlist.dtype) +edge_index = nlist.new_zeros((2, 1)) +angle_index = nlist.new_zeros((3, 1))Cleaner & one call shorter.
deepmd/pt/model/descriptor/repflow_layer.py (2)
373-375: Broadcast-then-reshape may allocate – favoureinsumThe pattern
flat_h2.unsqueeze(-1) * flat_edge_ebd.unsqueeze(-2)
allocates an intermediate(n_edge,3,e_dim)tensor before reshaping. A memory-friendly alternative:-flat_h2g2 = (flat_h2.unsqueeze(-1) * flat_edge_ebd.unsqueeze(-2)).reshape(-1, 3 * e_dim) +flat_h2g2 = torch.einsum('ij,ik->ijk', flat_h2, flat_edge_ebd).flatten(1)
einsumfuses multiply-and-flatten, avoiding the temporary and giving cuBLAS a densem x kmatmul-like path.
752-756: Deriven_edgefrom the authoritative source
n_edge = h2.shape[0]assumesh2is always freshly flattened.
Usingedge_index.shape[-1]is cheaper (already in registers) and less error-prone if callers ever reorder the arguments:- n_edge = h2.shape[0] + n_edge = edge_index.shape[-1]deepmd/dpmodel/descriptor/repflows.py (1)
581-583: Return truly-empty index tensors whenuse_dynamic_selis offA length-1 dummy index (filled with zeros) is created here even though the indices are never consumed when
self.use_dynamic_selisFalse.
Using size-0 tensors more accurately represents “no indices”, eliminates the accidental self-loop(0 → 0)and avoids accidental misuse if the guard condition ever drifts.- edge_index = xp.zeros([2, 1], dtype=nlist.dtype) - angle_index = xp.zeros([3, 1], dtype=nlist.dtype) + edge_index = xp.zeros((2, 0), dtype=nlist.dtype) + angle_index = xp.zeros((3, 0), dtype=nlist.dtype)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
deepmd/dpmodel/descriptor/repflows.py(5 hunks)deepmd/dpmodel/utils/network.py(2 hunks)deepmd/pt/model/descriptor/repflow_layer.py(4 hunks)deepmd/pt/model/descriptor/repflows.py(2 hunks)deepmd/pt/model/network/utils.py(2 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (29)
- GitHub Check: Test C++ (false)
- GitHub Check: Test C++ (true)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Analyze (python)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
🔇 Additional comments (4)
deepmd/dpmodel/utils/network.py (1)
1114-1116:Details
❌ Incorrect review comment
Possible dtype mismatch – make the stack explicit
xp.stack([n2e_index, n_ext2e_index], axis=0)(and the analogousangle_index_result) relies on both 1-D inputs having identical dtypes. In practicen2e_indexis created byarange(defaulting toint64for NumPy,int32for JAX) whileshifted_nlistinherits the dtype of the user-supplied neighbor list.
If those differ,xp.stackpromotes to the wider dtype, silently allocating a new buffer.- edge_index_result = xp.stack([n2e_index, n_ext2e_index], axis=0) + edge_index_result = xp.stack( + [xp.asarray(n2e_index, dtype=nlist.dtype), + xp.asarray(n_ext2e_index, dtype=nlist.dtype)], + axis=0, + )Same remark for
angle_index_result.
Keeps memory footprint predictable and avoids backend-dependent promotions.
🏁 Script executed:
#!/bin/bash # Scan the repo for mixing of different integer dtypes in edge/angle indices. rg --type=py -n --context 2 $'arange.*n_edge' | headLength of output: 803
🏁 Script executed:
#!/bin/bash set -e echo "=== n2e_index context ===" rg -n 'n2e_index' -C 5 deepmd/dpmodel/utils/network.py echo echo "=== n_ext2e_index context ===" rg -n 'n_ext2e_index' -C 5 deepmd/dpmodel/utils/network.py echo echo "=== n2a_index context ===" rg -n 'n2a_index' -C 5 deepmd/dpmodel/utils/network.py echo echo "=== eij2a_index context ===" rg -n 'eij2a_index' -C 5 deepmd/dpmodel/utils/network.py echo echo "=== eik2a_index context ===" rg -n 'eik2a_index' -C 5 deepmd/dpmodel/utils/network.py echo echo "=== arange n_angle context ===" rg -n 'arange.*n_angle' -C 2 deepmd/dpmodel/utils/network.py || echo "No 'n_angle' arange found"Length of output: 9669
No dtype mismatch risk: stack inputs already use
nlist.dtypeAll intermediate index arrays (
n2e_index,n_ext2e_index,n2a_index,eij2a_index,eik2a_index) are constructed with explicitdtype=nlist.dtype(viaxp.arange(..., dtype=nlist.dtype),xp.zeros(..., dtype=nlist.dtype), or slicing ofshifted_nlistwhich itself usesnlist.dtype). Therefore,xp.stackis stacking arrays of identical dtype and will not silently promote to a wider type. The proposed explicitasarray(..., dtype=nlist.dtype)guard is unnecessary.Likely an incorrect or invalid review comment.
deepmd/dpmodel/descriptor/repflows.py (3)
1290-1292: Docstring updated correctlyThe parameter comments now reflect the new
(2, n_edge)/(3, n_angle)shapes. Nice catch.
1366-1371: Destructuring follows the new layoutSplitting
edge_indexandangle_indexrow-wise is consistent with the updated conventions and the downstream logic (aggregate,xp.take, etc.). Looks correct.
626-631:Details
✅ Verification successful
Orientation LGTM, but double-check global adoption
owner = edge_index[0]assumes the new(2, n_edge)“column-major” layout introduced in this PR.
Please make sure all call-sites across the repository now follow this convention; any remaining code that still expects the old(n_edge, 2)layout will silently break.
🏁 Script executed:
#!/bin/bash # Detect potential leftovers that still expect the old edge/angle index shape. # 1. Files that index the *second* dimension of edge/angle_index (old layout) rg -n 'edge_index\.shape\s*\[\s*1\s*\]' || true rg -n 'angle_index\.shape\s*\[\s*1\s*\]' || true # 2. Usages of column slicing (old layout) like edge_index[:, 0] rg -n 'edge_index\[\s*:\s*,\s*0\s*\]' || true rg -n 'angle_index\[\s*:\s*,\s*0\s*\]' || trueLength of output: 193
All Old‐Layout References Removed
A repository-wide search for
edge_index.shape[1],angle_index.shape[1],edge_index[:,0]andangle_index[:,0]returned no matches, indicating full adoption of the new(2, n_edge)convention. No action needed.
04193b3 to
de84e1a
Compare
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## devel #4804 +/- ##
==========================================
- Coverage 84.80% 84.79% -0.01%
==========================================
Files 698 698
Lines 67798 67821 +23
Branches 3542 3540 -2
==========================================
+ Hits 57494 57508 +14
- Misses 9171 9178 +7
- Partials 1133 1135 +2 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Pull Request Overview
This PR refactors the indexing of edges and angles to use contiguous, column-major memory layouts, yielding a ~10% training speedup.
- Standardizes axis order for edge and angle indices across PT and DP implementations using
stackinstead ofcat. - Updates fallback initializations and indexing logic in dynamic selection paths.
- Simplifies tensor broadcasting patterns and updates related docstrings.
Reviewed Changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| deepmd/pt/model/network/utils.py | Changed get_graph_index to return 2 x n_edge/3 x n_angle via torch.stack |
| deepmd/pt/model/descriptor/repflows.py | Adjusted zero-index fallback shapes and updated owner indexing |
| deepmd/pt/model/descriptor/repflow_layer.py | Replaced manual None broadcasts with unsqueeze, updated dynamic‐sel shapes and indexing |
| deepmd/dpmodel/utils/network.py | Mirrored PT changes in get_graph_index using xp.stack axis=0 |
| deepmd/dpmodel/descriptor/repflows.py | Adjusted fallback zeros, owner indexing, and updated docstrings |
…ng#4804) This brings ~10% speedup for the training of DPA3 model with 24 thin layers and dynamic sel (average training time: 0.6891 s/batch vs 0.7635 s/batch, batch size = auto:128). <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **Refactor** - Updated the shape and indexing conventions for edge and angle index arrays and tensors across multiple components, standardizing them to column-major format for improved consistency. - Adjusted initialization and handling of index arrays/tensors to match the new conventions in both NumPy and PyTorch implementations. - Updated relevant method calls and internal logic to align with the revised index shapes and access patterns. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This brings ~10% speedup for the training of DPA3 model with 24 thin layers and dynamic sel (average training time: 0.6891 s/batch vs 0.7635 s/batch, batch size = auto:128).
Summary by CodeRabbit