Skip to content

Refactor and test vectorized BTD#25

Merged
AsymmetryChou merged 7 commits into
deepmodeling:mainfrom
AsymmetryChou:btd_acc
Jun 7, 2026
Merged

Refactor and test vectorized BTD#25
AsymmetryChou merged 7 commits into
deepmodeling:mainfrom
AsymmetryChou:btd_acc

Conversation

@AsymmetryChou

Copy link
Copy Markdown
Collaborator

Remove duplicated code and vectorize the compute_edges functions. Add equivalence tests to ensure the new vectorized implementation matches the legacy version for various matrix configurations. Update related functions to utilize JIT compilation for improved performance.

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Refactors the block-tridiagonalization utilities in dpnegf/negf/split_btd.py to reduce duplicated logic and improve performance via vectorization and Numba JIT, while adding regression tests that validate equivalence with the legacy implementations across a range of sparsity patterns.

Changes:

  • Vectorize compute_edge using np.maximum.at + np.maximum.accumulate.
  • Introduce Numba JIT cores for compute_blocks and find_optimal_cut and route public APIs through them.
  • Add new equivalence test suites covering dense/sparse inputs, docstring examples, and NEGF-like sparsity patterns.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File Description
dpnegf/negf/split_btd.py Reworks edge/profile computation and block computation/cut selection to use vectorized + Numba-accelerated cores.
dpnegf/tests/test_split_btd_compute_edge.py Adds new-vs-legacy equivalence tests for the vectorized compute_edge.
dpnegf/tests/test_split_btd_compute_blocks.py Adds extensive equivalence coverage for compute_blocks, find_optimal_cut, and optimized subblock splitting.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment on lines +471 to +474
new = compute_blocks_optimized(edge, edge1, left=1, right=1)
old = _compute_blocks_optimized_legacy(edge, edge1, left=1, right=1)
assert new == old
assert sum(new) == n
@AsymmetryChou AsymmetryChou merged commit 7c81a4e into deepmodeling:main Jun 7, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants