Skip to content

Refactor RGF kernel and Add no_grad decorator#28

Merged
AsymmetryChou merged 3 commits into
deepmodeling:mainfrom
AsymmetryChou:rgf_acc
Jun 13, 2026
Merged

Refactor RGF kernel and Add no_grad decorator#28
AsymmetryChou merged 3 commits into
deepmodeling:mainfrom
AsymmetryChou:rgf_acc

Conversation

@AsymmetryChou

Copy link
Copy Markdown
Collaborator

This pull request introduces a significant performance and usability improvement to the recursive Green's function (RGF) calculation in dpnegf/negf/recursive_green_cal.py by 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

  • Added a stacked mode to recursive_gf_cal that 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]
  • Updated the wrapper recursive_gf to 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

  • Improved and expanded docstrings for both recursive_gf_cal and recursive_gf to document the new stacked mode, its requirements, and its effect on performance and outputs. [1] [2]
  • Refactored code for clarity: removed unnecessary else branches, hoisted repeated matrix multiplications to variables, and added comments explaining the uniform/stacked path. [1] [2]

Minor Fixes and Cleanups

  • Added @torch.no_grad() to the NEGF class 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.

@AsymmetryChou AsymmetryChou merged commit 73b9d8f into deepmodeling:main Jun 13, 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.

1 participant