Merge pull request #35 from AsymmetryChou/rgf_acc

RGF_ACC: Reduce dead memory in RGF and Improve energy chunk size handling

Author: AsymmetryChou

.gitignore

@@ -34,7 +34,9 @@ examples/hBN/negf_output_k20/self_energy/self_energy_leadL.h5
 examples/hBN/negf_output_k20/self_energy/self_energy_leadR.h5
 examples/hBN/negf_output_k50/self_energy/*
 examples/hBN/negf_output_k70/self_energy/*
-examples/CNT/output/*
+examples/CNT/negf_profiling
+examples/CNT/output*
+examples/CNT/*long*
 examples/long_cnt/*
 CLAUDE*
 ai_docs/*

Dockerfile

@@ -41,7 +41,7 @@ RUN \
     conda run -n dpnegf pip install torch==2.5.1 --index-url https://download.pytorch.org/whl/cpu && \
     # [2] Pull torch-scatter from the PyG wheel index and use --only-binary=torch-scatter to fully disable source builds.
     # If no matching wheel is found it fails immediately instead of spending ~10 minutes compiling a package that would crash at runtime.
-    conda run -n dpnegf pip install torch-scatter -f https://data.pyg.org/whl/torch-2.5.0+cpu.html --only-binary=torch-scatter && \
+    conda run -n dpnegf pip install torch-scatter -f https://data.pyg.org/whl/torch-2.5.1+cpu.html --only-binary=torch-scatter && \
     # [3] Guard the local-repo installs with the CPU index so hidden dependencies can't replace the CPU torch with a CUDA build.
     conda run -n dpnegf pip install ./DeePTB torch==2.5.1 --extra-index-url https://download.pytorch.org/whl/cpu && \
     conda run -n dpnegf pip install ./ torch==2.5.1 --extra-index-url https://download.pytorch.org/whl/cpu && \

dpnegf/negf/recursive_green_cal.py

@@ -176,10 +176,17 @@ def recursive_gf_cal(energy, mat_l_list, mat_d_list, mat_u_list,
         # In-place: mat_d_list is a fresh tensor (wrapper's `* 1.` copy on D),
         # so we can fuse the energy shift without the e_bcast*sd transient.
         mat_d_list[jj].addcmul_(sd[jj], e_bcast, value=-1)
+        # sd[jj] is dead after this — it's only read here in the non-uniform
+        # kernel. Drop the wrapper-side view so the addcmul_ transient slab
+        # can be coalesced by the caching allocator inside the loop instead
+        # of waiting for the wrapper's `sd_b` Python name to leave scope.
+        sd[jj] = None
     for jj in range(len(mat_l_list)):
         mat_l_list[jj] = mat_l_list[jj] - e_bcast * sl[jj]
+        sl[jj] = None
     for jj in range(len(mat_u_list)):
         mat_u_list[jj] = mat_u_list[jj] - e_bcast * su[jj]
+        su[jj] = None
 
     num_of_matrices = len(mat_d_list)
     mat_shapes = [item.shape for item in mat_d_list]  # [B, n_q, n_q]
@@ -454,6 +461,13 @@ def _to_batch(t):
         Sd = torch.stack(sd_b, dim=0)                 # [K,   B, n, n]
         Sl = torch.stack(sl_b, dim=0)                 # [K-1, B, n, n]
         Su = torch.stack(su_b, dim=0)                 # [K-1, B, n, n]
+        # torch.stack on the wrapper's `*1.` D copies and the L/U/sd/sl/su
+        # expanded views produces six owned 4-D tensors. The wrapper-side
+        # lists are dead from here on — drop them now so the per-slot
+        # `[B, n_q, n_q]` storage (≈ K × B × n² × 16 B for D) can be freed
+        # before the kernel allocates gr_left/grl/gru.
+        del temp_mat_d_list, temp_mat_l_list, temp_mat_u_list
+        del sd_b, sl_b, su_b
         ans = recursive_gf_cal(shift_energy, L, D, U, Sd, Su, Sl,
                                s_in=s_in_b, s_out=s_out_b, eta=eta,
                                need_lesser=need_lesser,
@@ -470,6 +484,12 @@ def _to_batch(t):
                                need_gr_lc=need_gr_lc,
                                stacked=False,
                                keep_gr_left=keep_gr_left)
+        # Non-uniform kernel consumed the lists by reference and nulled
+        # individual slots as it went. Drop the wrapper-side names so the
+        # Python list objects (and any straggler refs) are gone before
+        # _squeeze_ans/return.
+        temp_mat_d_list = temp_mat_l_list = temp_mat_u_list = None
+        sd_b = sl_b = su_b = None
 
     if squeezed:
         ans = _squeeze_ans(ans)

dpnegf/runner/NEGF.py

@@ -607,9 +607,12 @@ def _auto_chunk_size(self, n_grid):
         Per-energy peak (post per-slot-release, complex128) approximated as
             bytes_per_E ~= C * K * n_max**2 * 16
         with C bundling the live tensors in the worst backward-sweep slot
-        (grd full + grl + gru full + decaying gr_left tail + gU + transients).
-        C=10 with a 0.7x free-memory budget; deliberately conservative because
-        without expandable_segments the allocator can't defragment on demand.
+        (grd full + grl + gru full + decaying gr_left tail + gU + transients
+        + the K-resident H/S diagonal & off-diagonal lists that survive across
+        chunks). C=14 with a 0.5x free-memory budget; deliberately conservative
+        because without expandable_segments the allocator can't defragment on
+        demand, and on real workloads (CNT10/long6) the previous 10x / 0.7
+        combination still picked a chunk that OOM'd on a 15.77 GiB V100.
         """
         rgf_dev = self.rgf_device
         if not (isinstance(rgf_dev, torch.device) and rgf_dev.type == "cuda"):
@@ -620,10 +623,10 @@ def _auto_chunk_size(self, n_grid):
             K = len(self.deviceprop.hd)
         except Exception:
             return n_grid
-        per_e = 10 * K * (n_max ** 2) * 16
+        per_e = 14 * K * (n_max ** 2) * 16
         if per_e <= 0:
             return n_grid
-        b = max(1, min(n_grid, int(0.7 * free_bytes) // per_e))
+        b = max(1, min(n_grid, int(0.5 * free_bytes) // per_e))
         log.info(
             f"auto e_batch_size={b} (free={free_bytes/2**30:.2f} GiB, "
             f"per_E~={per_e/2**20:.1f} MiB, K={K}, n_max={n_max})"
@@ -751,8 +754,19 @@ def negf_compute(self,scf_require=False,Vbias=None):
                         # Non-SCF: solve a whole chunk of energies in one batched recursive_gf call.
                         if self.e_batch_size is not None:
                             chunk = self.e_batch_size
+                            # The user-supplied value bypasses the auto-budget.
+                            rgf_dev = self.rgf_device
+                            if isinstance(rgf_dev, torch.device) and rgf_dev.type == "cuda":
+                                cap = self._auto_chunk_size(len(self.uni_grid))
+                                if chunk > cap:
+                                    log.warning(
+                                        f"user e_batch_size={chunk} exceeds the "
+                                        f"CUDA auto-cap={cap} on {rgf_dev}; "
+                                    )
                         else:
                             chunk = self._auto_chunk_size(len(self.uni_grid))
+                        log.info(f"Using e_batch_size={chunk} for energy loop with {len(self.uni_grid)} points")
+
                         for e_chunk in torch.split(self.uni_grid, chunk):
                             e_batch_size = len(e_chunk)
                             log.info(

ut.sh

@@ -19,7 +19,7 @@ echo "--- Installing/updating package from PR in editable mode ---"
 #    Docker image.
 # 2. `pytest ./tests/`: After the package is installed, we run the tests.
 
-conda run -n dpnegf bash -c "pip install -e . 'torch==2.1.1' --extra-index-url https://download.pytorch.org/whl/cpu -f https://data.pyg.org/whl/torch-2.1.1+cpu.html --only-binary=torch-scatter && pytest dpnegf/tests/"
+conda run -n dpnegf bash -c "pip install -e . 'torch==2.5.1' --extra-index-url https://download.pytorch.org/whl/cpu -f https://data.pyg.org/whl/torch-2.5.1+cpu.html --only-binary=torch-scatter && pytest dpnegf/tests/"
 
 echo "--- Unit Tests Passed Successfully ---"