Add multi-threading

__main__.py

@@ -38,6 +38,8 @@ def _config_from_args(args) -> HPCConfig:
         cfg.n_cand_d = args.cand_d
     if args.cand_m is not None:
         cfg.n_cand_m = args.cand_m
+    if args.workers is not None:
+        cfg.workers = args.workers
     if args.mse:
         cfg.optimizer = "mse"
     if args.no_spectral:
@@ -223,6 +225,7 @@ def main() -> None:
     parser.add_argument("--seed", type=int, default=1234, help="demo RNG seed")
     parser.add_argument("--cand-d", type=int, default=None, help="candidate d count")
     parser.add_argument("--cand-m", type=int, default=None, help="candidate min count")
+    parser.add_argument("--workers", type=int, default=None, help="worker processes (0=auto CPU cores, 1=single)")
     parser.add_argument("--mse", action="store_true", help="disable HPC sequence search")
     parser.add_argument("--no-spectral", action="store_true", help="disable spectral/DC objective")
     parser.add_argument("--verbose", action="store_true")

config.py

@@ -75,7 +75,8 @@ class HPCConfig:
     clamp_sigma_soft: float = 6.0
     clamp_sigma_hard: float = 4.0
 
-    # ---- reporting ------------------------------------------------------
+    # ---- reporting & parallelism ----------------------------------------
+    workers: int = 0
     verbose: bool = False
 
     def validate(self) -> "HPCConfig":
@@ -109,4 +110,6 @@ class HPCConfig:
             raise ValueError("log_slack must be non-negative")
         if self.viterbi_alpha < 0.0 or self.viterbi_beta < 0.0:
             raise ValueError("Viterbi weights must be non-negative")
+        if self.workers < 0:
+            raise ValueError("workers must be >= 0 (0 = auto CPU count)")
         return self

gui.py

@@ -97,6 +97,7 @@ class HPCQuantizeGUI:
         self.iters_var = tk.IntVar(value=8)
         self.ulp_var = tk.IntVar(value=4)
         self.limit_var = tk.IntVar(value=0)
+        self.workers_var = tk.IntVar(value=0)
 
         self.use_analogs_var = tk.BooleanVar(value=True)
         self.coupled_q2k_var = tk.BooleanVar(value=True)
@@ -136,6 +137,7 @@ class HPCQuantizeGUI:
             self.cand_d_var,
             self.cand_m_var,
             self.limit_var,
+            self.workers_var,
             self.use_spectral_var,
         ):
             v.trace_add("write", lambda *_: self._update_cli_preview())
@@ -271,6 +273,7 @@ class HPCQuantizeGUI:
         self._spin_pair(grid_f, 1, 2, "Solver Iters:", self.iters_var, 1, 32)
         self._spin_pair(grid_f, 2, 0, "FP16 ULP Radius:", self.ulp_var, 0, 16)
         self._spin_pair(grid_f, 2, 2, "Tensor Limit (0=all):", self.limit_var, 0, 9999)
+        self._spin_pair(grid_f, 3, 0, "Workers (0=Auto CPU):", self.workers_var, 0, 256)
 
         spec_f = ttk.LabelFrame(opt_frame, text="Spectral & DC Parameters", padding=6)
         spec_f.pack(fill=tk.X, pady=2)
@@ -481,6 +484,7 @@ class HPCQuantizeGUI:
             vw_lambda=float(self.vw_lambda_var.get()),
             dc_decay=float(self.dc_decay_var.get()),
             sse_budget=float(self.sse_budget_var.get()),
+            workers=int(self.workers_var.get()),
             verbose=False,
         )
         return cfg.validate()
@@ -503,6 +507,7 @@ class HPCQuantizeGUI:
             cmd += ["--output", out]
         if self.limit_var.get() > 0:
             cmd += ["--limit", str(self.limit_var.get())]
+        cmd += ["--workers", str(self.workers_var.get())]
         cmd += ["--cand-d", str(self.cand_d_var.get()), "--cand-m", str(self.cand_m_var.get())]
         if self.optimizer_var.get() == "mse":
             cmd.append("--mse")

pipeline.py

@@ -17,6 +17,8 @@ from __future__ import annotations
 
 from dataclasses import dataclass
 from math import ceil
+import multiprocessing as mp
+import os
 import time
 
 import numpy as np
@@ -405,6 +407,107 @@ def _apply_dc_nudge(X, W, blocks, cfg, row_width):
     return blocks
 
 
+_MP_STATE: dict[str, object] = {}
+
+
+def _resolve_workers(cfg: HPCConfig, n_blocks: int) -> int:
+    if cfg.workers == 1 or n_blocks < 4:
+        return 1
+    if cfg.workers > 1:
+        return min(int(cfg.workers), n_blocks)
+    if n_blocks < 16:
+        return 1
+    cpu = os.cpu_count() or 1
+    return min(cpu, n_blocks)
+
+
+def _phase2b_one(i: int) -> tuple[np.ndarray, FactorResult]:
+    X = _MP_STATE["X"]
+    W = _MP_STATE["W"]
+    seeds = _MP_STATE["seeds"]
+    base_d = _MP_STATE["base_d"]
+    base_m = _MP_STATE["base_m"]
+    cfg = _MP_STATE["cfg"]
+    wi = None if W is None else W[i]
+    cand = factor_candidates(
+        X[i],
+        wi,
+        seeds["scales"][i],
+        seeds["mins"][i],
+        base_d[i],
+        base_m[i],
+        cfg,
+    )
+    fac_res = factor_only_solve(X[i], wi, cfg)
+    return cand, fac_res
+
+
+def _phase4_one(i: int) -> tuple[Q2KBlock, int, int, float]:
+    X = _MP_STATE["X"]
+    W = _MP_STATE["W"]
+    seeds = _MP_STATE["seeds"]
+    base_d = _MP_STATE["base_d"]
+    base_m = _MP_STATE["base_m"]
+    best = _MP_STATE["best"]
+    carry = _MP_STATE["carry"]
+    standalone = _MP_STATE["standalone"]
+    cfg = _MP_STATE["cfg"]
+
+    wi = None if W is None else W[i]
+    d16, m16 = candidate_pair(base_d[i], base_m[i], best[i])
+    at_slot = factor_q2k_block(
+        X[i],
+        wi,
+        cfg,
+        init_s=seeds["scales"][i],
+        init_t=seeds["mins"][i],
+        d_seed=float(from_fp16_bits(d16)),
+        m_seed=float(from_fp16_bits(m16)),
+        pin=cfg.factor_pin,
+        use_spectral=cfg.use_spectral,
+        max_iter=cfg.factor_iters,
+        starts=1,
+        dc_carry=0.0,
+    )
+    d, m, ls, lm, codes, obj, sse = _assemble_block(
+        X[i], wi, seeds, base_d[i], base_m[i], best[i], carry[i], cfg
+    )
+    cand_hit = 0
+    fac_hit = 0
+
+    if cfg.use_candidate_floor:
+        cd, cm, cls, clm = _candidate_config(
+            seeds["scales"][i], seeds["mins"][i], base_d[i], base_m[i], best[i]
+        )
+        ccodes = codes_from_params(X[i], cd, cm, cls, clm)
+        cobj, csse = objective(
+            X[i], wi, cd, cm, cls, clm, ccodes, cfg, float(carry[i])
+        )
+        if cobj < obj - 1.0e-15:
+            d, m, ls, lm, codes, obj, sse = cd, cm, cls, clm, ccodes, cobj, csse
+            cand_hit = 1
+
+    if cfg.use_factor_floor:
+        for fac in (at_slot, standalone[i]):
+            fcodes = refine_codes_hadamard(
+                X[i], wi, fac.d, fac.dmin, fac.ls, fac.lm, fac.codes, cfg, float(carry[i])
+            )
+            fobj, fsse = objective(
+                X[i], wi, fac.d, fac.dmin, fac.ls, fac.lm, fcodes, cfg, float(carry[i])
+            )
+            if fobj < obj - 1.0e-15:
+                d = float(fac.d)
+                m = float(fac.dmin)
+                ls = fac.ls.copy()
+                lm = fac.lm.copy()
+                codes = fcodes.copy()
+                obj, sse = fobj, fsse
+                fac_hit = 1
+
+    blk_obj = make_block(d, m, ls, lm, codes)
+    return blk_obj, cand_hit, fac_hit, blk_obj.error(X[i])
+
+
 class Q2KQuantizer:
     def __init__(self, config: HPCConfig | None = None):
         self.config = (config or HPCConfig()).validate()
@@ -481,37 +584,75 @@ class Q2KQuantizer:
         t_cand = time.time()
         last_emit = 0.0
         step_report = max(1, min(32, n_blocks // 20 or 1))
-        for i in range(n_blocks):
-            _check_stop()
-            wi = None if W is None else W[i]
-            candidate_errors[i] = factor_candidates(
-                X[i],
-                wi,
-                seeds["scales"][i],
-                seeds["mins"][i],
-                base_d[i],
-                base_m[i],
-                cfg,
-            )
-            fac_res = factor_only_solve(X[i], wi, cfg)
-            standalone.append(fac_res)
-            now = time.time()
-            if (i + 1) % step_report == 0 or (i + 1) == n_blocks or (now - last_emit) >= 0.25:
-                last_emit = now
-                elapsed = max(now - t_cand, 1.0e-6)
-                rate = (i + 1) / elapsed
-                eta = (n_blocks - (i + 1)) / max(rate, 1.0e-9)
-                _emit(
-                    "Phase 2b: Analog-Guided Factored Candidate Field",
-                    i + 1,
-                    n_blocks,
-                    f"blk {i+1}/{n_blocks} ({rate:.1f} blk/s, ETA {eta:.1f}s) "
-                    f"d={float(fac_res.d):.5f} dmin={float(fac_res.dmin):.5f} "
-                    f"obj={fac_res.objective:.5e} starts={fac_res.starts}",
-                    rate=rate,
-                    eta=eta,
-                    sigma=sigma,
-                )
+        n_workers = _resolve_workers(cfg, n_blocks)
+        _MP_STATE.update(
+            {
+                "X": X,
+                "W": W,
+                "seeds": seeds,
+                "base_d": base_d,
+                "base_m": base_m,
+                "cfg": cfg,
+            }
+        )
+
+        if n_workers > 1 and "fork" in mp.get_all_start_methods():
+            ctx = mp.get_context("fork")
+            chunksize = max(1, min(16, n_blocks // (n_workers * 4) or 1))
+            pool = ctx.Pool(processes=n_workers)
+            try:
+                for i, (cand_row, fac_res) in enumerate(
+                    pool.imap(_phase2b_one, range(n_blocks), chunksize=chunksize)
+                ):
+                    _check_stop()
+                    candidate_errors[i] = cand_row
+                    standalone.append(fac_res)
+                    now = time.time()
+                    if (i + 1) % step_report == 0 or (i + 1) == n_blocks or (now - last_emit) >= 0.25:
+                        last_emit = now
+                        elapsed = max(now - t_cand, 1.0e-6)
+                        rate = (i + 1) / elapsed
+                        eta = (n_blocks - (i + 1)) / max(rate, 1.0e-9)
+                        _emit(
+                            f"Phase 2b: Analog-Guided Factored Field ({n_workers} workers)",
+                            i + 1,
+                            n_blocks,
+                            f"blk {i+1}/{n_blocks} ({rate:.1f} blk/s, ETA {eta:.1f}s) "
+                            f"d={float(fac_res.d):.5f} dmin={float(fac_res.dmin):.5f} "
+                            f"obj={fac_res.objective:.5e} starts={fac_res.starts}",
+                            rate=rate,
+                            eta=eta,
+                            sigma=sigma,
+                        )
+                pool.close()
+                pool.join()
+            except BaseException:
+                pool.terminate()
+                pool.join()
+                raise
+        else:
+            for i in range(n_blocks):
+                _check_stop()
+                cand_row, fac_res = _phase2b_one(i)
+                candidate_errors[i] = cand_row
+                standalone.append(fac_res)
+                now = time.time()
+                if (i + 1) % step_report == 0 or (i + 1) == n_blocks or (now - last_emit) >= 0.25:
+                    last_emit = now
+                    elapsed = max(now - t_cand, 1.0e-6)
+                    rate = (i + 1) / elapsed
+                    eta = (n_blocks - (i + 1)) / max(rate, 1.0e-9)
+                    _emit(
+                        "Phase 2b: Analog-Guided Factored Candidate Field",
+                        i + 1,
+                        n_blocks,
+                        f"blk {i+1}/{n_blocks} ({rate:.1f} blk/s, ETA {eta:.1f}s) "
+                        f"d={float(fac_res.d):.5f} dmin={float(fac_res.dmin):.5f} "
+                        f"obj={fac_res.objective:.5e} starts={fac_res.starts}",
+                        rate=rate,
+                        eta=eta,
+                        sigma=sigma,
+                    )
 
         _check_stop()
         _emit(
@@ -532,27 +673,13 @@ class Q2KQuantizer:
             sigma=sigma,
         )
 
-        at_slot: list[FactorResult] = []
-        for i in range(n_blocks):
-            _check_stop()
-            wi = None if W is None else W[i]
-            d16, m16 = candidate_pair(base_d[i], base_m[i], best[i])
-            at_slot.append(
-                factor_q2k_block(
-                    X[i],
-                    wi,
-                    cfg,
-                    init_s=seeds["scales"][i],
-                    init_t=seeds["mins"][i],
-                    d_seed=float(from_fp16_bits(d16)),
-                    m_seed=float(from_fp16_bits(m16)),
-                    pin=cfg.factor_pin,
-                    use_spectral=cfg.use_spectral,
-                    max_iter=cfg.factor_iters,
-                    starts=1,
-                    dc_carry=0.0,
-                )
-            )
+        _MP_STATE.update(
+            {
+                "best": best,
+                "carry": carry,
+                "standalone": standalone,
+            }
+        )
 
         blocks: list[Q2KBlock] = []
         factor_floor_used = 0
@@ -560,67 +687,80 @@ class Q2KQuantizer:
         running_sse = 0.0
         t_asm = time.time()
         last_emit = 0.0
-        for i in range(n_blocks):
-            _check_stop()
-            wi = None if W is None else W[i]
-            d, m, ls, lm, codes, obj, sse = _assemble_block(
-                X[i], wi, seeds, base_d[i], base_m[i], best[i], carry[i], cfg
-            )
 
-            if cfg.use_candidate_floor:
-                cd, cm, cls, clm = _candidate_config(
-                    seeds["scales"][i], seeds["mins"][i], base_d[i], base_m[i], best[i]
-                )
-                ccodes = codes_from_params(X[i], cd, cm, cls, clm)
-                cobj, csse = objective(
-                    X[i], wi, cd, cm, cls, clm, ccodes, cfg, float(carry[i])
-                )
-                if cobj < obj - 1.0e-15:
-                    d, m, ls, lm, codes, obj, sse = cd, cm, cls, clm, ccodes, cobj, csse
-                    candidate_floor_used += 1
-
-            if cfg.use_factor_floor:
-                for fac in (at_slot[i], standalone[i]):
-                    fcodes = refine_codes_hadamard(
-                        X[i], wi, fac.d, fac.dmin, fac.ls, fac.lm, fac.codes, cfg, float(carry[i])
-                    )
-                    fobj, fsse = objective(
-                        X[i], wi, fac.d, fac.dmin, fac.ls, fac.lm, fcodes, cfg, float(carry[i])
+        if n_workers > 1 and "fork" in mp.get_all_start_methods():
+            ctx = mp.get_context("fork")
+            chunksize = max(1, min(16, n_blocks // (n_workers * 4) or 1))
+            pool = ctx.Pool(processes=n_workers)
+            try:
+                for i, (blk_obj, cand_hit, fac_hit, blk_err) in enumerate(
+                    pool.imap(_phase4_one, range(n_blocks), chunksize=chunksize)
+                ):
+                    _check_stop()
+                    blocks.append(blk_obj)
+                    candidate_floor_used += cand_hit
+                    factor_floor_used += fac_hit
+                    running_sse += blk_err
+                    now = time.time()
+                    if (i + 1) % step_report == 0 or (i + 1) == n_blocks or (now - last_emit) >= 0.25:
+                        last_emit = now
+                        elapsed = max(now - t_asm, 1.0e-6)
+                        rate = (i + 1) / elapsed
+                        eta = (n_blocks - (i + 1)) / max(rate, 1.0e-9)
+                        cur_rmse = float(np.sqrt(running_sse / ((i + 1) * QK_K)))
+                        _emit(
+                            f"Phase 4: Sub-block Sieve + ULP + Hadamard Assembly ({n_workers} workers)",
+                            i + 1,
+                            n_blocks,
+                            f"blk {i+1}/{n_blocks} ({rate:.1f} blk/s, ETA {eta:.1f}s) "
+                            f"rmse={cur_rmse:.6e} rmse/sigma={cur_rmse/max(sigma, 1e-15):.4f} "
+                            f"floors(cand={candidate_floor_used}, fac={factor_floor_used})",
+                            rmse=cur_rmse,
+                            sigma=sigma,
+                            temperature=temperature,
+                            cand_floors=candidate_floor_used,
+                            fac_floors=factor_floor_used,
+                            rate=rate,
+                            eta=eta,
+                        )
+                pool.close()
+                pool.join()
+            except BaseException:
+                pool.terminate()
+                pool.join()
+                raise
+        else:
+            for i in range(n_blocks):
+                _check_stop()
+                blk_obj, cand_hit, fac_hit, blk_err = _phase4_one(i)
+                blocks.append(blk_obj)
+                candidate_floor_used += cand_hit
+                factor_floor_used += fac_hit
+                running_sse += blk_err
+                now = time.time()
+                if (i + 1) % step_report == 0 or (i + 1) == n_blocks or (now - last_emit) >= 0.25:
+                    last_emit = now
+                    elapsed = max(now - t_asm, 1.0e-6)
+                    rate = (i + 1) / elapsed
+                    eta = (n_blocks - (i + 1)) / max(rate, 1.0e-9)
+                    cur_rmse = float(np.sqrt(running_sse / ((i + 1) * QK_K)))
+                    _emit(
+                        "Phase 4: Sub-block Sieve + ULP + Hadamard Assembly",
+                        i + 1,
+                        n_blocks,
+                        f"blk {i+1}/{n_blocks} ({rate:.1f} blk/s, ETA {eta:.1f}s) "
+                        f"rmse={cur_rmse:.6e} rmse/sigma={cur_rmse/max(sigma, 1e-15):.4f} "
+                        f"floors(cand={candidate_floor_used}, fac={factor_floor_used})",
+                        rmse=cur_rmse,
+                        sigma=sigma,
+                        temperature=temperature,
+                        cand_floors=candidate_floor_used,
+                        fac_floors=factor_floor_used,
+                        rate=rate,
+                        eta=eta,
                     )
-                    if fobj < obj - 1.0e-15:
-                        d = float(fac.d)
-                        m = float(fac.dmin)
-                        ls = fac.ls.copy()
-                        lm = fac.lm.copy()
-                        codes = fcodes.copy()
-                        obj, sse = fobj, fsse
-                        factor_floor_used += 1
-
-            blk_obj = make_block(d, m, ls, lm, codes)
-            blocks.append(blk_obj)
-            running_sse += blk_obj.error(X[i])
-            now = time.time()
-            if (i + 1) % step_report == 0 or (i + 1) == n_blocks or (now - last_emit) >= 0.25:
-                last_emit = now
-                elapsed = max(now - t_asm, 1.0e-6)
-                rate = (i + 1) / elapsed
-                eta = (n_blocks - (i + 1)) / max(rate, 1.0e-9)
-                cur_rmse = float(np.sqrt(running_sse / ((i + 1) * QK_K)))
-                _emit(
-                    "Phase 4: Sub-block Sieve + ULP + Hadamard Assembly",
-                    i + 1,
-                    n_blocks,
-                    f"blk {i+1}/{n_blocks} ({rate:.1f} blk/s, ETA {eta:.1f}s) "
-                    f"rmse={cur_rmse:.6e} rmse/sigma={cur_rmse/max(sigma, 1e-15):.4f} "
-                    f"floors(cand={candidate_floor_used}, fac={factor_floor_used})",
-                    rmse=cur_rmse,
-                    sigma=sigma,
-                    temperature=temperature,
-                    cand_floors=candidate_floor_used,
-                    fac_floors=factor_floor_used,
-                    rate=rate,
-                    eta=eta,
-                )
+
+        _MP_STATE.clear()
 
         if cfg.apply_dc_carry and (cfg.dc_lambda > 0.0 or cfg.dc_decay > 0.0):
             _check_stop()