Fix bit_length=64 undefined behavior, make it configurable, and lower --samples/--max_nb defaults (#12)

* Fix undefined behavior from hardcoded bit_length=64; make it configurable

SBD packs determinant bitstrings into std::vector<size_t> 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.

* 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.

* 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.

Author: hfwen0502

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",

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,

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<size_t>`` 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():