From 6153793d47ef27f7cffd0ddf7960adfe9fde2bd8 Mon Sep 17 00:00:00 2001 From: jinyuchem Date: Sat, 21 Feb 2026 21:58:19 -0500 Subject: [PATCH 1/3] codex: improve gpu performance --- pyfrsoc/compute_so_mat.py | 9 +- pyfrsoc/dso_sos.py | 9 +- pyfrsoc/hamiltonian/hamiltonian.py | 40 ++++--- pyfrsoc/hamiltonian/soc_hamiltonian.py | 7 +- pyfrsoc/utils/compute_vnl.py | 18 +++- pyfrsoc/utils/util_functs.py | 138 ++++++------------------- 6 files changed, 90 insertions(+), 131 deletions(-) diff --git a/pyfrsoc/compute_so_mat.py b/pyfrsoc/compute_so_mat.py index 06e9299..66ce005 100755 --- a/pyfrsoc/compute_so_mat.py +++ b/pyfrsoc/compute_so_mat.py @@ -157,15 +157,16 @@ def setup_pps(self, pp_files): """ Setup q. Only Gamma point at this moment """ + g_pp = self.ks_up.g_loc + npw_pp = self.ks_up.npw + collect_pps = dict() for i in range(len(pp_files)): pp = pseudo_potential(pp_files[i], verbose=(self.rank == 0)) pp.compute_dvan() - pp.interp_beta( - self.ks_up.ngw, self.ks_up.g, self.q, self.ks_up.omega, self.ks_up.tpiba, dbeta=False, dq=0.01 - ) - pp.beta_ang(self.ks_up.ngw, self.ks_up.g, self.q) + pp.interp_beta(npw_pp, g_pp, self.q, self.ks_up.omega, self.ks_up.tpiba, dbeta=False, dq=0.01) + pp.beta_ang(npw_pp, g_pp, self.q) pp.dvan_so = np.asarray(pp.dvan_so) pp.vkb1 = np.asarray(pp.vkb1) diff --git a/pyfrsoc/dso_sos.py b/pyfrsoc/dso_sos.py index 315a5c2..39c49f7 100755 --- a/pyfrsoc/dso_sos.py +++ b/pyfrsoc/dso_sos.py @@ -146,15 +146,16 @@ def setup_pps(self, pp_files): """ Setup q. Only Gamma point """ + g_pp = self.ks_up.g_loc + npw_pp = self.ks_up.npw + collect_pps = dict() for i in range(len(pp_files)): pp = pseudo_potential(pp_files[i]) pp.compute_dvan() - pp.interp_beta( - self.ks_up.ngw, self.ks_up.g, self.q, self.ks_up.omega, self.ks_up.tpiba, dbeta=False, dq=0.01 - ) - pp.beta_ang(self.ks_up.ngw, self.ks_up.g, self.q) + pp.interp_beta(npw_pp, g_pp, self.q, self.ks_up.omega, self.ks_up.tpiba, dbeta=False, dq=0.01) + pp.beta_ang(npw_pp, g_pp, self.q) pp.dvan_so = np.asarray(pp.dvan_so) pp.vkb1 = np.asarray(pp.vkb1) diff --git a/pyfrsoc/hamiltonian/hamiltonian.py b/pyfrsoc/hamiltonian/hamiltonian.py index 9419247..d4093e1 100755 --- a/pyfrsoc/hamiltonian/hamiltonian.py +++ b/pyfrsoc/hamiltonian/hamiltonian.py @@ -22,25 +22,35 @@ def load_cube(ks, fname): headl = noa + 6 num_elements = fft_grid[0] * fft_grid[1] * fft_grid[2] + data = None with open(fname, "r") as f: # Skip the header lines for _ in range(headl): next(f) - # Preallocate an array for the data - data = np2.zeros(num_elements) + # Fast path: C parser over the remaining flat numeric payload. + try: + data = np2.fromfile(f, dtype=np2.float64, sep=" ") + except (TypeError, ValueError): + data = None - # Index for filling the array - index = 0 + if data is None or data.size != num_elements: + with open(fname, "r") as f: + for _ in range(headl): + next(f) - # Read and process the file line by line - for line in f: - line_data = np2.fromiter(map(float, line.split()), dtype=np2.float64) - data[index : index + len(line_data)] = line_data - index += len(line_data) + data = np2.zeros(num_elements) + index = 0 + for line in f: + line_data = np2.fromiter(map(float, line.split()), dtype=np2.float64) + data[index : index + len(line_data)] = line_data + index += len(line_data) - # Reshape to match the fft_grid - potential = data.reshape(fft_grid) + if index != num_elements: + raise ValueError(f"Unexpected cube data size in {fname}: got {index}, expected {num_elements}") + + # Reshape to match the fft_grid + potential = data.reshape(fft_grid) potential = np.asarray(potential) @@ -121,14 +131,18 @@ def setup_pps(self, pp_files): """ Setup q. Only Gamma point """ + g_pp = self.ks.g_loc + npw_pp = self.ks.npw + collect_pps = dict() for i in range(len(pp_files)): # average the pseudo potential pp = pseudo_potential(pp_files[i], average_pp=True, verbose=(self.rank == 0)) pp.compute_dvan() - pp.interp_beta(self.ks.ngw, self.ks.g, self.q, self.ks.omega, self.ks.tpiba, dbeta=False, dq=0.01) - pp.beta_ang(self.ks.ngw, self.ks.g, self.q) + # Build beta projectors only on local PW shards to avoid redundant per-rank work. + pp.interp_beta(npw_pp, g_pp, self.q, self.ks.omega, self.ks.tpiba, dbeta=False, dq=0.01) + pp.beta_ang(npw_pp, g_pp, self.q) if pp.spin_orbit: pp.dvan_so = np.asarray(pp.dvan_so) diff --git a/pyfrsoc/hamiltonian/soc_hamiltonian.py b/pyfrsoc/hamiltonian/soc_hamiltonian.py index 6033dea..8326a40 100755 --- a/pyfrsoc/hamiltonian/soc_hamiltonian.py +++ b/pyfrsoc/hamiltonian/soc_hamiltonian.py @@ -58,13 +58,16 @@ def setup_pps(self, pp_files): """ Setup q. Only Gamma point at this moment """ + g_pp = self.ks.g_loc + npw_pp = self.ks.npw + collect_pps = dict() for i in range(len(pp_files)): pp = pseudo_potential(pp_files[i]) pp.compute_dvan() - pp.interp_beta(self.ks.ngw, self.ks.g, self.q, self.ks.omega, self.ks.tpiba, dbeta=False, dq=0.01) - pp.beta_ang(self.ks.ngw, self.ks.g, self.q) + pp.interp_beta(npw_pp, g_pp, self.q, self.ks.omega, self.ks.tpiba, dbeta=False, dq=0.01) + pp.beta_ang(npw_pp, g_pp, self.q) if pp.spin_orbit: pp.dvan_so = np.asarray(pp.dvan_so) diff --git a/pyfrsoc/utils/compute_vnl.py b/pyfrsoc/utils/compute_vnl.py index e3bac32..bcfd51d 100755 --- a/pyfrsoc/utils/compute_vnl.py +++ b/pyfrsoc/utils/compute_vnl.py @@ -95,8 +95,14 @@ def solid_vkb_loc(atoms, pps, q, g, mill, eigts, fft_grid, pw_index, npwx, neg=F Compute projectors for the solid, including the structure factor To be used with parallelization """ - g_loc = g[pw_index] - mill_loc = mill[pw_index] + if g.shape[0] == len(pw_index): + g_loc = g + else: + g_loc = g[pw_index] + if mill.shape[0] == len(pw_index): + mill_loc = mill + else: + mill_loc = mill[pw_index] nr1 = fft_grid[0] nr2 = fft_grid[1] @@ -121,11 +127,15 @@ def solid_vkb_loc(atoms, pps, q, g, mill, eigts, fft_grid, pw_index, npwx, neg=F pp = pps[element] vkb1 = pp.vkb1[: pp.number_of_proj_ang, :] + if vkb1.shape[1] == npw: + vkb1_loc = vkb1 + else: + vkb1_loc = vkb1[:, pw_index] if neg: - vkb[ikb : ikb + pp.number_of_proj_ang, :npw] = np.conj(vkb1[:, pw_index]) * sk[na, :] * phases[na] + vkb[ikb : ikb + pp.number_of_proj_ang, :npw] = np.conj(vkb1_loc) * sk[na, :] * phases[na] else: - vkb[ikb : ikb + pp.number_of_proj_ang, :npw] = vkb1[:, pw_index] * sk[na, :] * phases[na] + vkb[ikb : ikb + pp.number_of_proj_ang, :npw] = vkb1_loc * sk[na, :] * phases[na] ikb += pp.number_of_proj_ang diff --git a/pyfrsoc/utils/util_functs.py b/pyfrsoc/utils/util_functs.py index 83c4dc8..4cdaece 100755 --- a/pyfrsoc/utils/util_functs.py +++ b/pyfrsoc/utils/util_functs.py @@ -199,11 +199,8 @@ def _sph_harm(m, n, theta, phi): def ylmr2_scipy(lmax2, ng, g, gg): """ - MEGA-optimized version that computes ALL (l,m) combinations at once. - Maximum possible vectorization for spherical harmonics. - - This should be the fastest possible implementation using a single - massive vectorized sph_harm call with full broadcasting. + Compute spherical harmonics ylm(G) up to l=lmax. + GPU-friendly implementation that avoids huge broadcasted temporaries. """ # Convert inputs to the current backend g = np.asarray(g) @@ -232,78 +229,29 @@ def ylmr2_scipy(lmax2, ng, g, gg): theta = np.arccos(np.clip(cost, -1.0, 1.0)) - # Create all (l, m) pairs we need to compute - l_values = [] - m_values = [] - output_indices = [] - is_real_part = [] + # Initialize output array + ylm = np.zeros((lmax2, ng), dtype=np.complex128) idx = 0 + inv_sqrt2 = 1.0 / np.sqrt(2.0) for l in range(lmax + 1): - # m = 0 case - l_values.append(l) - m_values.append(0) - output_indices.append(idx) - is_real_part.append(True) # m=0 is always real + if lGPU and hasattr(special, "sph_harm"): + ylm[idx, :] = special.sph_harm(0, l, phi, theta) + else: + ylm[idx, :] = _sph_harm(0, l, phi, theta) idx += 1 - # m > 0 cases (we compute positive m and derive real/imag parts) for m in range(1, l + 1): - # We only need to compute positive m once - l_values.append(l) - m_values.append(m) - output_indices.append((idx, idx + 1)) # Both real and imag indices - is_real_part.append(False) # This will be split into real/imag - idx += 2 + if lGPU and hasattr(special, "sph_harm"): + tmp = special.sph_harm(m, l, phi, theta) + else: + tmp = _sph_harm(m, l, phi, theta) - # Simplified approach: compute unique (l,m) pairs - unique_lm = [] - lm_to_indices = {} + ylm[idx, :] = inv_sqrt2 * (tmp + np.conj(tmp)) + idx += 1 - idx = 0 - for l in range(lmax + 1): - # m = 0 - unique_lm.append((l, 0)) - lm_to_indices[(l, 0)] = [idx] - idx += 1 - - # m > 0 - for m in range(1, l + 1): - unique_lm.append((l, m)) - lm_to_indices[(l, m)] = [idx, idx + 1] # real and imag indices - idx += 2 - - # Extract l and m arrays for unique pairs - unique_l = np.array([lm[0] for lm in unique_lm]) - unique_m = np.array([lm[1] for lm in unique_lm]) - - # Create broadcast arrays: (n_unique_lm, ng) - phi_broadcast = np.broadcast_to(phi[None, :], (len(unique_lm), ng)) - theta_broadcast = np.broadcast_to(theta[None, :], (len(unique_lm), ng)) - l_broadcast = np.broadcast_to(unique_l[:, None], (len(unique_lm), ng)) - m_broadcast = np.broadcast_to(unique_m[:, None], (len(unique_lm), ng)) - - # Single massive vectorized sph_harm call for ALL (l,m,point) combinations! - if lGPU and hasattr(special, "sph_harm"): - sph_harm_results = special.sph_harm(m_broadcast, l_broadcast, phi_broadcast, theta_broadcast) - else: - sph_harm_results = _sph_harm(m_broadcast, l_broadcast, phi_broadcast, theta_broadcast) - - # Initialize output array - ylm = np.zeros((lmax2, ng), dtype=np.complex128) - - # Map results back to proper ylm indices - for i, (l, m) in enumerate(unique_lm): - indices = lm_to_indices[(l, m)] - - if len(indices) == 1: - # m = 0 case - ylm[indices[0], :] = sph_harm_results[i, :] - else: - # m > 0 case: split into real and imaginary parts - tmp = sph_harm_results[i, :] - ylm[indices[0], :] = 1 / np.sqrt(2) * (tmp + np.conj(tmp)) # Real part - ylm[indices[1], :] = 1 / np.sqrt(2) / 1j * (tmp - np.conj(tmp)) # Imag part + ylm[idx, :] = inv_sqrt2 / 1j * (tmp - np.conj(tmp)) + idx += 1 return ylm @@ -521,9 +469,7 @@ def compute_beta_projectors( beta_proj_l, beta_proj_r_ndm, radial_grid_ndm, radial_grid_factor_ndm, nqx, dq, pref, pseudo_type="default" ): """ - Fully vectorized computation of beta projectors - maximum performance version. - - This version minimizes Python loops and maximizes vectorized operations. + Compute beta projectors with low-memory batched algebra. """ if pseudo_type == "gth": raise Exception("GTH not implemented") @@ -542,47 +488,31 @@ def compute_beta_projectors( # Initialize output tab_beta = np.zeros((number_of_proj, nqx), dtype=np.float64) + # Simpson integration weights, reused for all l + weights = np.ones(ndm, dtype=np.float64) + weights[1:-1:2] = 4.0 + weights[2:-1:2] = 2.0 + if ndm % 2 == 0: + weights[-2] = -1.0 + simpson_weights = (weights * radial_grid_factor_ndm) / 3.0 + # Group projectors by l value for efficient computation unique_l = np.unique(beta_proj_l) for l in unique_l: l_mask = beta_proj_l == l proj_indices = np.where(l_mask)[0] - n_proj_l = len(proj_indices) - - # Compute spherical Bessel functions once for this l - # Create argument matrix: radial_grid_ndm * qi for all combinations - r_grid, q_grid = np.meshgrid(radial_grid_ndm, qi_array, indexing="ij") - arg_matrix = r_grid * q_grid # Shape: (ndm, nqx) - - # Vectorized Bessel function computation - besr_matrix = np.zeros((ndm, nqx), dtype=np.float64) - for ir in range(ndm): - besr_matrix[ir, :] = sph_bes(int(l), arg_matrix[ir, :]) - - # Transpose to get (nqx, ndm) for easier broadcasting - besr_matrix = besr_matrix.T - - # Process all projectors with this l value in a vectorized way - beta_subset = beta_proj_r_ndm[proj_indices] # Shape: (n_proj_l, ndm) - - # Compute aux for all projectors at once - # beta_subset[:, None, :] has shape (n_proj_l, 1, ndm) - # besr_matrix[None, :, :] has shape (1, nqx, ndm) - # radial_grid_ndm[None, None, :] has shape (1, 1, ndm) - aux_3d = beta_subset[:, None, :] * besr_matrix[None, :, :] * radial_grid_ndm[None, None, :] - - # Reshape for vectorized integration: (n_proj_l * nqx, ndm) - aux_reshaped = aux_3d.reshape(-1, ndm) + beta_subset = beta_proj_r_ndm[proj_indices] # (n_proj_l, ndm) - # Vectorized Simpson integration - vqint_flat = simpson(ndm, aux_reshaped, radial_grid_factor_ndm) + # Bessel matrix once per l: (nqx, ndm) + arg_matrix = radial_grid_ndm[:, None] * qi_array[None, :] + besr_matrix = sph_bes(int(l), arg_matrix).T - # Reshape back: (n_proj_l, nqx) - vqint_matrix = vqint_flat.reshape(n_proj_l, nqx) + # Kernel already includes r * Simpson weights: (nqx, ndm) + kernel = besr_matrix * (radial_grid_ndm * simpson_weights)[None, :] - # Store results - tab_beta[proj_indices, :] = vqint_matrix * pref + # Matrix multiplication replaces large 3D intermediates. + tab_beta[proj_indices, :] = np.dot(beta_subset, kernel.T) * pref return tab_beta From 69bd2eacf4df8177f337ebcef8ccc8ba0211a370 Mon Sep 17 00:00:00 2001 From: jinyuchem Date: Sat, 21 Feb 2026 22:03:01 -0500 Subject: [PATCH 2/3] patch hamiltonian.py --- pyfrsoc/hamiltonian/hamiltonian.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/pyfrsoc/hamiltonian/hamiltonian.py b/pyfrsoc/hamiltonian/hamiltonian.py index d4093e1..95867f2 100755 --- a/pyfrsoc/hamiltonian/hamiltonian.py +++ b/pyfrsoc/hamiltonian/hamiltonian.py @@ -1,4 +1,5 @@ import gc +import warnings from mpi4py import MPI @@ -30,8 +31,12 @@ def load_cube(ks, fname): # Fast path: C parser over the remaining flat numeric payload. try: - data = np2.fromfile(f, dtype=np2.float64, sep=" ") - except (TypeError, ValueError): + with warnings.catch_warnings(): + # Numpy emits DeprecationWarning for unmatched text and will raise + # ValueError in future versions. Treat it as parse failure now. + warnings.simplefilter("error", DeprecationWarning) + data = np2.fromfile(f, dtype=np2.float64, sep=" ") + except (TypeError, ValueError, DeprecationWarning): data = None if data is None or data.size != num_elements: From 9d2867788b0c4e8f42047a15f7315123f7e3e677 Mon Sep 17 00:00:00 2001 From: jinyuchem Date: Sat, 21 Feb 2026 22:37:58 -0500 Subject: [PATCH 3/3] remove free_all_blocks --- pyfrsoc/dso_stern.py | 2 -- pyfrsoc/dso_stern_hybrid.py | 2 -- pyfrsoc/hamiltonian/hamiltonian.py | 1 - pyfrsoc/sternheimer/sternheimer.py | 1 - pyfrsoc/sternheimer/sternheimer_hybrid.py | 1 - 5 files changed, 7 deletions(-) diff --git a/pyfrsoc/dso_stern.py b/pyfrsoc/dso_stern.py index 2afd92d..e5ac58b 100755 --- a/pyfrsoc/dso_stern.py +++ b/pyfrsoc/dso_stern.py @@ -251,7 +251,6 @@ def compute_h_soc_psi(self): gc.collect() if lGPU: synchronize() - np.get_default_memory_pool().free_all_blocks() self.comm.Barrier() return [h_soc_psi_0, h_soc_psi_1] @@ -316,7 +315,6 @@ def solve_stern(self, stern, h_soc_psi, threshold): self.comm.Barrier() if lGPU: synchronize() - np.get_default_memory_pool().free_all_blocks() gc.collect() self.comm.Barrier() diff --git a/pyfrsoc/dso_stern_hybrid.py b/pyfrsoc/dso_stern_hybrid.py index 5febf60..2ceb96b 100755 --- a/pyfrsoc/dso_stern_hybrid.py +++ b/pyfrsoc/dso_stern_hybrid.py @@ -303,7 +303,6 @@ def compute_h_soc_psi(self): gc.collect() if lGPU: synchronize() - np.get_default_memory_pool().free_all_blocks() self.comm.Barrier() return [h_soc_psi_0, h_soc_psi_1] @@ -398,7 +397,6 @@ def solve_stern(self, stern, h_soc_psi, threshold): self.comm.Barrier() if lGPU: synchronize() - np.get_default_memory_pool().free_all_blocks() gc.collect() self.comm.Barrier() diff --git a/pyfrsoc/hamiltonian/hamiltonian.py b/pyfrsoc/hamiltonian/hamiltonian.py index 95867f2..461e0d1 100755 --- a/pyfrsoc/hamiltonian/hamiltonian.py +++ b/pyfrsoc/hamiltonian/hamiltonian.py @@ -1066,6 +1066,5 @@ def compute_h_psi(self, psi): self.comm.Barrier() if lGPU: synchronize() - np.get_default_memory_pool().free_all_blocks() gc.collect() self.comm.Barrier() diff --git a/pyfrsoc/sternheimer/sternheimer.py b/pyfrsoc/sternheimer/sternheimer.py index 3fe5570..c9a57af 100755 --- a/pyfrsoc/sternheimer/sternheimer.py +++ b/pyfrsoc/sternheimer/sternheimer.py @@ -290,7 +290,6 @@ def cg_solver(self, pert_evc, threshold=1e-16, max_iter=100, debug=False): self.comm.Barrier() if lGPU: synchronize() - np.get_default_memory_pool().free_all_blocks() gc.collect() self.comm.Barrier() diff --git a/pyfrsoc/sternheimer/sternheimer_hybrid.py b/pyfrsoc/sternheimer/sternheimer_hybrid.py index 6625fb5..0be126e 100755 --- a/pyfrsoc/sternheimer/sternheimer_hybrid.py +++ b/pyfrsoc/sternheimer/sternheimer_hybrid.py @@ -1000,7 +1000,6 @@ def cg_solver(self, pert_evc, threshold=1e-16, max_iter=100, debug=False): self.comm.Barrier() if lGPU: synchronize() - np.get_default_memory_pool().free_all_blocks() gc.collect() self.comm.Barrier()