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", diff --git a/python/examples/run_sqd_sbd.py b/python/examples/run_sqd_sbd.py index 8de1e79..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,13 +85,17 @@ 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") 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():