From b772e219a340df072923af65ab558fd205639480 Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Tue, 1 Sep 2026 12:38:21 -0400 Subject: [PATCH 1/3] Fix undefined behavior from hardcoded bit_length=64; make it configurable SBD packs determinant bitstrings into std::vector with `bit_length` bits per word. The Python wrapper hardcoded 64 in four places, but bitadvance() in framework/bit_manipulation.h computes size_t d = (((size_t) 1) << bit_length) - 1; Shifting a 64-bit size_t by 64 is undefined behavior. In practice the shift count is masked to 0, so the mask collapses to d == 0 (checked on aarch64 and x86-64). bitadvance() is reached from mpi_redistribution() and mpi_sort_bitarray(), the multi-rank determinant distribution paths. RIKEN documents the default as 20 (see --bit_length in apps/chemistry_tpb_selected_basis_diagonalization/README.md), which is also what run_sbd_diag.py already defaults to; only run_sqd_sbd.py was passing 64. Changes: - Add SBD_DEFAULT_BIT_LENGTH = 20 in sbd_solver.py and use it for the _create_sbd_config default. - Thread the effective value through. _ci_strings_to_sbd_dets() and _sbd_dets_to_ci_strings() now take bit_length, and _solve_sci_core() builds the config before packing so determinants are packed with the same value the C++ engine uses to interpret them. Previously a caller-supplied bit_length reached the engine but not the packing, which would silently corrupt the determinants. - Add --bit_length to run_sqd_sbd.py (default 20) instead of hardcoding. Validated on GB200 (NVHPC 25.5, cc100, 4 GPUs, CUDA-aware MPICH 5.0.1): - H2O 1e-5 subspace, dim 303,601: -76.2437350421 at bit_length 20, both at np=1 and 2x2, matching the pre-change value and run_sbd_diag.py's independent file-based path (-76.2437349413). - bit_length 20 / 48 / 63 give bit-identical energies, confirming that packing, unpacking and the engine now agree. norb=24 at bit_length=20 spans two size_t words, so this also exercises the multi-word path. --- python/examples/run_sqd_sbd.py | 6 +++++- python/sbd_solver.py | 37 +++++++++++++++++++++++++--------- 2 files changed, 33 insertions(+), 10 deletions(-) diff --git a/python/examples/run_sqd_sbd.py b/python/examples/run_sqd_sbd.py index 8de1e79..faf9963 100644 --- a/python/examples/run_sqd_sbd.py +++ b/python/examples/run_sqd_sbd.py @@ -92,6 +92,10 @@ def parse_args(): p.add_argument("--carryover_type", type=int, default=1) p.add_argument("--carryover_ratio", "--ratio", type=float, default=0.1, dest="ratio") p.add_argument("--carryover_threshold", "--threshold", type=float, default=1e-4, dest="threshold") + p.add_argument("--bit_length", type=int, default=20, + help="Bits packed into each size_t of the bitstring representation " + "(RIKEN default 20). Must be <= 63: bitadvance() shifts a " + "64-bit size_t by this amount, so 64 is undefined behavior.") # MPI sub-communicator sizes p.add_argument("--adet_comm_size", type=int, default=1) @@ -223,7 +227,7 @@ def main(): "carryover_type": args.carryover_type, "ratio": args.ratio, "threshold": args.threshold, - "bit_length": 64, + "bit_length": args.bit_length, "adet_comm_size": args.adet_comm_size, "bdet_comm_size": args.bdet_comm_size, "task_comm_size": args.task_comm_size, diff --git a/python/sbd_solver.py b/python/sbd_solver.py index 3e4f723..47e18b2 100644 --- a/python/sbd_solver.py +++ b/python/sbd_solver.py @@ -29,6 +29,17 @@ import numpy as np from mpi4py import MPI +# Bits packed into each size_t of SBD's ``std::vector`` bitstring +# representation. RIKEN's documented default is 20 (see the --bit_length option +# in apps/chemistry_tpb_selected_basis_diagonalization/README.md), which is also +# what run_sbd_diag.py defaults to. +# +# This must stay <= 63. bitadvance() in framework/bit_manipulation.h computes +# size_t d = (((size_t) 1) << bit_length) - 1; +# so bit_length == 64 shifts a 64-bit size_t by 64, which is undefined behavior; +# in practice the shift count is masked to 0 and the mask collapses to d == 0. +SBD_DEFAULT_BIT_LENGTH = 20 + try: from pyscf import tools as pyscf_tools except ImportError: @@ -154,11 +165,15 @@ def _solve_sci_core( the FCIDUMP only once and reuse it across all batches. """ strings_a, strings_b = ci_strings - adet = _ci_strings_to_sbd_dets(strings_a, norb, backend) - bdet = _ci_strings_to_sbd_dets(strings_b, norb, backend) + # Build the config first: it carries the effective bit_length (possibly + # overridden by the caller), and the determinants must be packed with the + # same value the C++ engine will use to interpret them. sbd_data = _create_sbd_config(sbd_config, backend, device_config) + adet = _ci_strings_to_sbd_dets(strings_a, norb, backend, sbd_data.bit_length) + bdet = _ci_strings_to_sbd_dets(strings_b, norb, backend, sbd_data.bit_length) + # Use .bin extension to trigger SBD's fast binary write path # (SaveMatrixFormWF in restart.h checks extension: .bin -> raw doubles) wf_dump_file = sbd_dir / "wavefunction.bin" @@ -199,8 +214,12 @@ def _solve_sci_core( occupancies_b = density[1::2] occupancies = (occupancies_a, occupancies_b) - co_strings_a = _sbd_dets_to_ci_strings(results["carryover_adet"], norb, backend) - co_strings_b = _sbd_dets_to_ci_strings(results["carryover_bdet"], norb, backend) + co_strings_a = _sbd_dets_to_ci_strings( + results["carryover_adet"], norb, backend, sbd_data.bit_length + ) + co_strings_b = _sbd_dets_to_ci_strings( + results["carryover_bdet"], norb, backend, sbd_data.bit_length + ) # Read wavefunction coefficients from binary dump n_alpha_co = len(co_strings_a) @@ -411,14 +430,14 @@ def _read_fcidump_ecore(fcidump_path): def _ci_strings_to_sbd_dets( - ci_strings: np.ndarray, norb: int, backend + ci_strings: np.ndarray, norb: int, backend, + bit_length: int = SBD_DEFAULT_BIT_LENGTH, ) -> list[list[int]]: """Convert CI strings (integers) to SBD determinant format. Determinants are sorted in canonical order (matching C++ sort_bitarray) which is required by the GPU Correlation kernel (do_rdm=1). """ - bit_length = 64 dets = [] for ci_str in ci_strings: binary_str = format(int(ci_str), f'0{norb}b') @@ -428,10 +447,10 @@ def _ci_strings_to_sbd_dets( def _sbd_dets_to_ci_strings( - dets: list[list[int]], norb: int, backend + dets: list[list[int]], norb: int, backend, + bit_length: int = SBD_DEFAULT_BIT_LENGTH, ) -> np.ndarray: """Convert SBD determinants to CI strings (integers).""" - bit_length = 64 ci_strings = [] for det in dets: binary_str = backend.makestring(det, bit_length, norb) @@ -459,7 +478,7 @@ def _create_sbd_config(config_dict: dict | None = None, backend=None, device_con sbd_data.carryover_type = 1 sbd_data.ratio = 0.1 sbd_data.threshold = 1e-4 - sbd_data.bit_length = 64 + sbd_data.bit_length = SBD_DEFAULT_BIT_LENGTH if config_dict: for key, value in config_dict.items(): From e38a6427d684e4570cdd5f7c3e79c296f41e8b6e Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Tue, 1 Sep 2026 12:49:24 -0400 Subject: [PATCH 2/3] Lower --samples default to 3000 and --max_nb default to 10 --max_nb sets the Davidson block size, and davidson_thrust.h allocates 2 * num_block vectors of length dim on the device, so GPU memory scales as dim * (2 * max_nb + ~4) * 8 bytes. 10 matches the upstream default in chemistry/tpb/sbdiag.h and cuts that workspace by ~4x versus 50. --- python/examples/run_sqd_sbd.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/examples/run_sqd_sbd.py b/python/examples/run_sqd_sbd.py index faf9963..e2b8f0b 100644 --- a/python/examples/run_sqd_sbd.py +++ b/python/examples/run_sqd_sbd.py @@ -65,7 +65,7 @@ def parse_args(): p.add_argument("--fcidump", required=True, help="Path to FCIDUMP file") p.add_argument("--counts", default=None, help="Path to count_dict.json (bitstring counts from hardware)") - p.add_argument("--samples", type=int, default=10000, + p.add_argument("--samples", type=int, default=3000, help="Number of uniform random samples (used when --counts is not given)") p.add_argument("--device", choices=["auto", "cpu", "gpu", "gpu-omp", "gpu-nvidia-omp"], @@ -85,7 +85,7 @@ def parse_args(): p.add_argument("--tolerance", "--eps", type=float, default=1e-8, dest="eps") p.add_argument("--iteration", "--max_it", type=int, default=100, dest="max_it", help="Max SBD Davidson iterations per diagonalization") - p.add_argument("--block", "--max_nb", type=int, default=50, dest="max_nb") + p.add_argument("--block", "--max_nb", type=int, default=10, dest="max_nb") p.add_argument("--rdm", "--do_rdm", type=int, default=0, dest="do_rdm", help="0=density only (default, sufficient for SQD), 1=full RDM") p.add_argument("--shuffle", "--do_shuffle", type=int, default=0, dest="do_shuffle") From 147196c850c9ea1767dd7e564de1bc167d2f3dae Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Tue, 1 Sep 2026 13:51:08 -0400 Subject: [PATCH 3/3] Update example notebook to bit_length 20 The notebook's sbd_config hardcoded 64, the value this PR removes as undefined behavior. Its 'max_nb': 10 already matches the new CLI default. --- python/examples/run_sqd_sbd.ipynb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/examples/run_sqd_sbd.ipynb b/python/examples/run_sqd_sbd.ipynb index 05e17ff..1917b96 100644 --- a/python/examples/run_sqd_sbd.ipynb +++ b/python/examples/run_sqd_sbd.ipynb @@ -199,7 +199,7 @@ " 'carryover_type': 1, # singles-only carryover (stable default)\n", " 'ratio': 0.1,\n", " 'threshold': 1e-4,\n", - " 'bit_length': 64,\n", + " 'bit_length': 20, # bits per size_t; RIKEN default (64 is UB in bitadvance)\n", " # serial: 1×1×1 MPI sub-communicator grid\n", " 'adet_comm_size': 1, 'bdet_comm_size': 1, 'task_comm_size': 1,\n", "}\n",