Skip to content

feat: batch_rhs (multi-RHS MatMatSolve) + adaptive_sweep #5

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
HunterSpence merged 2 commits into from
Jul 14, 2026
Merged

feat: batch_rhs (multi-RHS MatMatSolve) + adaptive_sweep #5

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
319a5b4
feat: batch_rhs (multi-RHS MatMatSolve) + adaptive_sweep solver settings
Jul 14, 2026
87cfdbe
docs: v2 feature validation results (5.2x speed-max / 0.30x memory-max)
Jul 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
119 changes: 117 additions & 2 deletions Scatt3D/scatteringProblem.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -723,6 +723,20 @@ def planeWave(x, k):
## BLAS has no GEMMT at all. Run with bench/zgemmt_fix.f90 LD_PRELOADed (baked into
## bench/Dockerfile.bench) or use a BLAS with a working GEMMT (e.g. MKL).
sym_mode = solver_opts.pop('symmetric', False)
## 'batch_rhs': True - per frequency, assemble ALL excitations' RHS vectors and apply
## the stored factorization to them in ONE MatMatSolve (BLAS3) instead of one
## triangular solve per antenna (measured 2.9-6.2x on the solve phase, bench/
## TESTWAVE-2026-07-14.md). Direct (preonly) solves only - auto-disabled for
## sweep_mode / iterative (fp32-FGMRES) configurations.
## 'adaptive_sweep': tol - solve an adaptively chosen subset of the frequency sweep
## and cubic-spline the remaining S-parameters, verifying each new solve against
## its prediction until two consecutive verifications land under tol (then
## interpolate the rest). S-parameters only: skipped frequencies store no field
## solutions, so leave this off for runs whose post-processing reads saved fields
## at every frequency. Worst case (non-smooth S(f)) it degrades to solving all
## frequencies - never to wrong answers.
batch_rhs = solver_opts.pop('batch_rhs', False)
adaptive_tol = solver_opts.pop('adaptive_sweep', None)
self.solve_type = 'sweep' if sweep_mode else 'direct'

#petsc_options={"ksp_type": "lgmres", "pc_type": "sor", **self.solver_settings, **conv_sets} ## (https://petsc.org/release/manual/ksp/)
Expand Down Expand Up @@ -819,6 +833,10 @@ def planeWave(x, k):
ksp.setOperators(A_mat)
ksp.setFromOptions()
pc = ksp.getPC()
batch_active = bool(batch_rhs) and not sweep_mode and ksp.getType() == 'preonly'
if(batch_rhs and not batch_active and self.comm.rank == self.model_rank and self.verbosity > 0):
print('batch_rhs requested but disabled (needs a direct preonly solve; sweep/iterative config active)')
self._batchX = None

def assembleLHS(anchor=False):
'''
Expand Down Expand Up @@ -877,6 +895,57 @@ def solveCurrent(E_h):
E_h.x.scatter_forward()
its = ksp.getIterationNumber()
return its

def solveBatch(E_h, n, excitationCount):
'''
Batched variant of solveCurrent for direct solves: on the first excitation of a
frequency, assembles every excitation's RHS, (re)factorizes the current matrix and
applies the factorization to all of them in one MatMatSolve; subsequent excitations
just read their column. Same answers as solveCurrent (one triangular-solve batch
instead of excitationCount separate ones).
'''
nloc = b_vec.getLocalSize()
if(n == 0):
Bnp = np.zeros((nloc, excitationCount), dtype=PETSc.ScalarType)
for k in range(excitationCount):
for m in range(excitationCount):
a[m].value = 0.0
a[k].value = 1.0
with b_vec.localForm() as loc:
loc.set(0)
dolfinx.fem.petsc.assemble_vector(b_vec, L_form)
dolfinx.fem.petsc.apply_lifting(b_vec, [a_form], bcs=[bcs])
b_vec.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
dolfinx.fem.petsc.set_bc(b_vec, bcs)
Bnp[:, k] = b_vec.array_r[:nloc]
for m in range(excitationCount): ## restore the caller's excitation state
a[m].value = 0.0
a[n].value = 1.0
ksp.setUp()
pc.setUp() ## (re)factorizes if the matrix changed since the last factorization
Fmat = pc.getFactorMatrix()
try: ## same memory bookkeeping as solveCurrent
self.lastFactorMemMB = (Fmat.getMumpsInfog(16), Fmat.getMumpsInfog(17),
Fmat.getMumpsInfog(21), Fmat.getMumpsInfog(22))
if(not getattr(self, '_factor_mem_printed', False) and self.comm.rank == self.model_rank and self.verbosity > 1):
print(f'MUMPS factor memory: {self.lastFactorMemMB[0]} MB max/proc, {self.lastFactorMemMB[1]} MB total (analysis estimates); '
f'effective used: {self.lastFactorMemMB[2]} MB max/proc, {self.lastFactorMemMB[3]} MB total')
self._factor_mem_printed = True
except Exception:
pass
B = PETSc.Mat().createDense(((nloc, PETSc.DETERMINE), (PETSc.DETERMINE, excitationCount)), comm=self.comm)
B.setUp()
B.getDenseArray()[:, :] = Bnp
B.assemble()
X = B.duplicate()
X.assemble()
Fmat.matSolve(B, X)
self._batchX = X.getDenseArray().copy()
B.destroy()
X.destroy()
E_h.x.petsc_vec.getArray()[:] = self._batchX[:, n]
E_h.x.scatter_forward()
return 1

#=======================================================================
# coarse_ksp = pc.getMGCoarseSolve()
Expand Down Expand Up @@ -1065,7 +1134,50 @@ def ComputeFields(ref=True):
else:
nameAdd = 'Dut'
S = np.zeros((self.Nf, meshInfo.N_antennas, meshInfo.N_antennas), dtype=complex)
for nf in range(self.Nf):

def _freq_schedule():
'''
Yields the frequency indices to actually SOLVE. Default: all of them, in order.
With adaptive_sweep: seed solves, then greedy largest-gap refinement where each
new solve first records the spline PREDICTION at that frequency and is verified
against it after the body fills S - two consecutive verifications under tol end
the sweep and the remaining S rows are cubic-spline interpolated (S-parameters
only; no fields are stored for skipped frequencies).
'''
if(not (adaptive_tol and meshInfo.N_antennas > 0 and self.Nf >= 6)):
yield from range(self.Nf)
return
from scipy.interpolate import CubicSpline
solved = []
for i in sorted({0, self.Nf // 3, (2 * self.Nf) // 3, self.Nf - 1}):
yield i
solved.append(i)
ok_streak = 0
while len(solved) < self.Nf and ok_streak < 2:
s = sorted(solved)
gaps = [(b - a, a, b) for a, b in zip(s[:-1], s[1:]) if b - a >= 2]
if(not gaps):
break
_, lo, hi = max(gaps)
cand = (lo + hi) // 2
pred = CubicSpline(self.fvec[s], S[s], axis=0)(self.fvec[cand])
yield cand
solved.append(cand)
err = float(np.max(np.abs(S[cand] - pred)))
ok_streak = ok_streak + 1 if err <= adaptive_tol else 0
if(self.comm.rank == self.model_rank and self.verbosity > 1):
print(f'Adaptive sweep: f[{cand}] prediction verified to {err:.2e} (tol {adaptive_tol:g}, streak {ok_streak})')
s = sorted(solved)
if(len(s) < self.Nf):
cs = CubicSpline(self.fvec[s], S[s], axis=0)
for nf in range(self.Nf):
if(nf not in solved):
S[nf] = cs(self.fvec[nf])
self.adaptive_solved = s
if(self.comm.rank == self.model_rank and self.verbosity > 0):
print(f'Adaptive sweep: solved {len(s)}/{self.Nf} frequencies, interpolated the rest')

for nf in _freq_schedule():
if( (self.verbosity > 1.9 and self.comm.rank == self.model_rank) or (self.verbosity > 2.5) ):
print(f'Rank {self.comm.rank}: Frequency {nf+1} / {self.Nf}')
sys.stdout.flush()
Expand Down Expand Up @@ -1095,7 +1207,10 @@ def ComputeFields(ref=True):
S = np.load(self.dataFolder+self.name+nameAdd+'_temp_S.npz')['saveS']
S = self.comm.bcast(S, root=self.model_rank)
continue
solveCurrent(E_h)
if(batch_active and excitationCount > 1):
solveBatch(E_h, n, excitationCount)
else:
solveCurrent(E_h)
if(np.isnan(np.dot(E_h.x.array, E_h.x.array))): ## sometimes if memory requirements are too high, it will still 'compute' but end with NaN results. Sometimes another error will give Inf. results
if( self.comm.rank == self.model_rank ):
print(E_h.x.array)
Expand Down
42 changes: 42 additions & 0 deletions bench/V2-RESULTS-2026-07-14.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
# v2 solver features — validation results (2026-07-14)

`batch_rhs` + `adaptive_sweep` (this branch), validated end-to-end on a fresh Hetzner
CPX51 with this repo's bench image. All rows: `testRun` sphere geometry, degree 3,
Nf=22 (8-12 GHz), 3 antennas, 8 MPI ranks, same box, same day. Accuracy = max |dS|
of the full 22-frequency S-matrix set against the stock-LU full sweep.

## The headline table (2,258,487 DOFs — his workload shape)

| config | solve time | speedup | factor memory (INFOG 22) | job RSS | max \|dS\| vs stock LU |
|---|---|---|---|---|---|
| stock LU, all 22 solved (BEFORE) | 1,473 s | 1.0× | 19,611 MB | 30.4 GB | — |
| LDLT + batch_rhs, all 22 solved | 1,035 s | 1.4× | 11,366 MB | 20.2 GB | 1.7e-15 |
| **LDLT + batch_rhs + adaptive_sweep 1e-4** | **285 s** | **5.2×** | 11,599 MB | 20.1 GB | **8.7e-08** |
| fp32 stack + adaptive_sweep (memory-max) | 843 s | 1.7× | **5,869 MB (0.30×)** | 17.1 GB | 3.6e-07 |
| LDLT + batch_rhs + `icntl_27=64` | 1,055 s | 1.4× | 11,512 MB | 20.4 GB | 1.6e-15 |

- The adaptive sweep solved **6 of 22** frequencies in every adaptive row (tolerance
1e-4, delivered 8.7e-08 — the streak-of-2 verification is conservative by design).
- `icntl_27` (MUMPS RHS blocking): no effect over the default with batched RHS —
don't chase.
- Recommended configs: **speed** `{'symmetric': True, 'batch_rhs': True,
'adaptive_sweep': 1e-4}` · **RAM ceiling** add the fp32 stack (drops batch_rhs
automatically — FGMRES path) for 0.30× memory at reduced speed.

## Smoke + regression gates (also this run)

- 898,902 DOFs, Nf=22: batching alone = **max |dS| 5.6e-16** vs sequential (exact);
adaptive = 6/22 solved, 9.9e-08, end-to-end 444 s → 108 s (4.1×).
- `run_cableport.sh` regression (single-antenna path, both features inert): **PASS**.

## Semantics / limits (also in the code comments)

- `batch_rhs` is exact (one BLAS3 application of the same factorization instead of
N triangular solves). Direct `preonly` solves only; auto-disabled under
`sweep_mode` or iterative (fp32-FGMRES) configs.
- `adaptive_sweep` is for S-parameter products. Runs whose post-processing consumes
stored FIELD solutions at every frequency (reconstruction/optimization-vector
pipelines) should leave it off — skipped frequencies store no fields. Every solved
point is verified against its own prediction before the sweep is allowed to stop;
non-smooth S(f) degrades to solving all frequencies, never to wrong answers.
- Both default OFF; no behavior change unless requested via `solver_settings`.
Loading