Refactor RGF kernel and Add no_grad decorator#28
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This pull request introduces a significant performance and usability improvement to the recursive Green's function (RGF) calculation in
dpnegf/negf/recursive_green_cal.pyby adding support for a "stacked" mode for uniform block matrices. When all device blocks share the same shape, the code now auto-detects this and processes the calculation in a fused, batched way, reducing Python overhead and improving efficiency. The change is backward-compatible: outputs are unstacked to lists for downstream compatibility. Additional refactoring and docstring improvements clarify the new behavior and streamline the code.The most important changes are:
Performance and Algorithm Improvements
stackedmode torecursive_gf_calthat processes uniform block matrices as stacked tensors, enabling a single fused operation and reusing cached identity matrices for improved efficiency. The function auto-detects when this mode can be used and falls back to the legacy list-based path otherwise. [1] [2] [3] [4] [5] [6]recursive_gfto detect when all blocks are uniform, stack them for the kernel call, and unbind outputs to maintain the previous API contract. [1] [2]Code Quality and Documentation
recursive_gf_calandrecursive_gfto document the new stacked mode, its requirements, and its effect on performance and outputs. [1] [2]Minor Fixes and Cleanups
@torch.no_grad()to theNEGFclass constructor for improved efficiency and removed an unnecessary context manager. [1] [2]These changes collectively improve both the performance and maintainability of the RGF code, especially for large systems with uniform block structure.