Slurm

[jtravs]

Improve Slurm handling of memory, paths, and julia environments

[Copilot]

Pull request overview

This PR improves the SlurmExec execution mode to generate more robust Slurm array-job scripts, adding per-task memory handling, Julia environment selection, thread configuration, and safer path handling, along with expanded user documentation.

Changes:

• Extend SlurmExec with memory, project, and nthreads options and generate SBATCH scripts accordingly.
• Add memory parsing to derive Julia --heap-size-hint from Slurm --mem.
• Add extensive Slurm usage documentation to the scans guide.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 7 comments.

File	Description
src/Scans.jl	Extends SlurmExec, adds heap-hint parsing, and updates Slurm job script generation (mem, threads, quoting, ulimit, Julia binary path).
docs/src/scans.md	Documents SlurmExec behavior and provides usage examples and explanation of memory/threads/project handling.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[jtravs]

@copilot check my fixes and resolve addressed threads

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 7 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[chrisbrahms]

Looks good to me. I'll need to check that it works with my usual way of working using SSHExec, though I don't think that is a block to merging.

Please add a brief summary to #415 for the next release notes.

[jtravs]

@chrisbrahms this branch now contains a mess of a large performance rework aimed at 3D
free-space (FreeGrid) propagation. Apologies, this could be split, currently it
is in here as I work on making the TG-FROG work tractable.

Note that this will surely conflict with the big FreeSpace refactor #416 , I suggets we merge that first and then I will fix here, but it is here for the record anyway.

Performance work for large-array propagation

Headline numbers, measured on a production 3D free-space run (fused silica, Kerr,
EnvGrid 128 × FreeGrid 1024², rtol = 1e-8, AMD EPYC 7763):

	before	after
peak RSS per propagation	~46 GB	30.6 GiB
field-sized solver/transform buffers	18.5	10
wall time per propagation (5 FFTW threads)	~4.9 h (1 Julia thread)	~3.0 h (8 Julia threads)
Julia-level threading of the step	none	4.18× on 8 threads

The performance changes are bit-identical and verified as such. I reproduced previously stored output
bit-for-bit (max relative difference exactly 0.0 on the k-space-integrated
spectra, two independent delay points at 1024²).

This also merges the #436 as I wanted to include that on my HPC runs. Ideally we merge that first.

Solver performance — applies to all propagation types

All of the following are bit-identical: the same floating-point operations in the same
order, enforced by tests.

• Register reduction, 12 → 9 field-sized buffers. yi (dense-output) and yerr
  (error estimate) are now allocated lazily and, on the default path, never at all:
  with norm === weaknorm the error metric is computed in a single fused pass
  (weaknorm_fused) that never materialises yerr, and with step_on-aligned saves
  the endpoint snap means yi is never touched. make_fbar!'s private scratch array is
  gone — the preconditioned RHS now propagates its input in place, which is safe because
  evaluate! rebuilds yn from y at every stage.
  Custom norms keep the materialised path via a fused_errnorm(norm) hook (returns
  nothing by default; define a method to opt in).
• Pass fusion, ~26 → ~8 field-sized passes per step. Butcher stage accumulation,
  local extrapolation, the error estimate and the interpolant are each now a single
  fused broadcast instead of a sequence of .+= passes. Julia's n-ary + is
  left-associated, so per element this matches the previous sequential accumulation
  bit-for-bit.
• FSAL by pointer swap rather than copying a full array each step.
• preserve_input=false (opt-in, on solve/solve_precon/Luna.run) lets the
  stepper adopt the caller's array as its yn buffer instead of copying it.
• Zero-interval propagator calls are skipped (t1 == t2 was doing a full
  exp(0) multiply — this also removes the wasted pass on rejected steps).

Threading

Utils.tforeach / Utils.tchunks are chunked Threads.@threads :static helpers used
for the elementwise kernels — the exp(linop·dz) propagator passes (by far the largest
single-threaded cost in 3D runs), the fused stage/error/interpolant passes, and the
transform's response, apodisation and scaling kernels.

They are gated: below 2²⁰ elements, or with one Julia thread, they call straight
through to the serial code path. Small-array propagations (modal, mode-averaged) are
therefore untouched, and CI — which runs single-threaded — exercises exactly the
previous code. Elementwise threading is bit-identical by construction; a 1-thread vs
N-thread comparison is included in the test suite.

Measured scaling on one 3D propagation (128 × 640², 144 steps, EPYC 7763):

Julia threads	wall	speedup	parallel efficiency
1	119.6 min	1×	—
4	39.3 min	3.04×	76%
8	28.6 min	4.18×	52%

Before this branch there was no Julia-level parallelism anywhere in the propagation —
FFTW's own threads were the only parallelism available.

FFTW threading on Julia ≥ 1.12

FFTW.jl installs a Julia-task (partr) threading callback into libfftw3 when
Threads.nthreads() > 1, and that callback segfaults on Julia ≥ 1.12 — which previously
forced the choice between threaded FFTs and multithreaded Julia.
Utils.use_native_fftw_threads() deregisters it (fftw_threads_set_callback(NULL)),
handing FFT threading back to libfftw3's own pthread pool; set_fftw_threads now does
this automatically on Julia ≥ 1.12 when more than one Julia thread is running. Both
libraries are forced to load first, since the callback is installed lazily at first load
and would otherwise be silently reinstated. Guarded on the FFTW (non-MKL) provider and
feature-detected, so it degrades to a no-op rather than erroring on other versions.

This is what makes the threading above usable at all on current Julia.

---

3D free-space transform (TransFree)

• Fast path, 4 → 1 working buffers. When the grid is not oversampled
  (length(grid.ωo) == length(grid.ω), scale == 1, complex field, no noise field —
  the usual EnvGrid/thg=false case), the copies through Eωo/Pωo are exact
  identities, so those buffers are no longer allocated and the FFTs run directly between
  the solution array, Eto and the output. If additionally every response is pointwise,
  the polarisation overwrites Eto in place and Pto goes too.
  The general path is unchanged and still selectable via fastpath=false.
• Pointwise response protocol. Kerr_field/Kerr_env now return small callable
  structs (same constructor API) carrying a Nonlinear.pointwise trait plus a
  pointwise_P scalar kernel, so their polarisation can be applied as one whole-array
  broadcast rather than looping over ~10⁶ transverse columns with a function call each.
  Et_to_Pt_ordered! applies responses in tuple order, so each element accumulates its
  contributions exactly as the columnwise version did.
• Batched Raman. RamanPolarEnvBatched computes the delayed response for all
  transverse points at once — two batched FFTs along the time dimension on the doubled
  anti-wraparound grid per RHS evaluation — instead of two length-Nto FFTs per
  transverse column (~2×10⁶ small serial FFTs per step at production size). It also
  updates the density-dependent response kernel once per call rather than once per
  column, which is exact since the density is a scalar. Agrees with the columnwise
  RamanPolarEnv to ~1e-15 element-wise and 1e-10 through a full propagation (not
  bitwise: batched FFT algorithms differ from single-column ones).
• Factored linear operator and normalisation (opt-in, factored=true on
  make_const_linop / const_norm_free). FactoredFreeLinop and FreeNorm store only
  the separable factors — k²(ω) as a vector and k⊥²(ky,kx) as a small matrix, ~8 MB
  total — and compute elements on demand, replacing two field-sized arrays and removing
  ~14 field-sized reads per step. Both subtype AbstractArray with a lazy getindex,
  so linoptype metadata, Stats dispatch and generic fallbacks all behave as before.
  A specialised threaded make_prop! method computes the exponential on the fly.
  The element formulas are shared (inlined) with the materialising fill loops, so the
  lazy and dense forms cannot drift apart — verified bit-identical both element-wise and
  through a full propagation.
• Luna.run scratch reuse. NonlinearRHS.scratch(transform) exposes the
  transform's time-domain buffer (dead between RHS evaluations) so run no longer
  allocates its own Et = FT \ Eω. Falls back to allocating for transforms that don't
  provide one.
• twin_period (opt-in, default 1 = current behaviour) applies the
  spectral/temporal window projection every Nth accepted step instead of every step, and
  always immediately before a save (via a new Output.willsave). Saves the two extra
  coarse FFTs per step. Note this is a physics-level change — the windowing feeds back
  into the propagation — so it needs its own validation before use.

Buffer accounting for one 3D free-space Kerr propagation:

	before	after
RK45	12	9
TransFree	4	1 (2 with Raman)
run's Et	1	0
linop + norm cache	1.5	~0 (factored)
total field-equivalents	18.5	10

---

Testing

• test/runtests.jl gained substring filtering: Pkg.test("Luna", test_args=["rk45", "freespace"]), or julia --project test/runtests.jl rk45. No
  arguments runs the full suite exactly as before. (Selection args are consumed before
  the test files are included, since Scans parses ARGS itself.)
• test/test_perf_bitident.jl is a small-grid harness that runs the same propagation
  through old and new code paths and asserts bitwise equality: fast path vs general
  path, factored vs materialised operators, threaded vs serial kernels, adopted vs
  copied input, fused vs materialised error norm, plus element-level agreement of the
  pointwise and batched responses against their columnwise references.

Because the old code paths remain selectable (fastpath=false, materialised operators,
twin_period=1, threading off), the bit-identity claims are checkable.

---

5. Applicability, and what could follow

Already general — every propagation type benefits: both correctness fixes, the
register reduction and pass fusion, preserve_input, the FSAL pointer swap, the FFTW
threading fix, and the threaded kernels (for arrays above the gate — modal runs stay on
the serial path unchanged).

Currently TransFree-only: the buffer fast path, array-level responses, factored
operators, scratch reuse and batched Raman.

No plasma yet this is all based on envelope propagation in glass. Plasma can be treated batched similar to Raman.

Radial (TransRadial) is the natural next step — the structure is nearly identical
(same four buffers, the same idcs column loop, the same trailing
nl .*= ωwin .* (-imω)./(2·normfun(z)), the same cached (nω, Nr) constant linop and
const_norm_radial):

• array-level pointwise responses are a one-line switch (Et_to_Pt! →
  Et_to_Pt_ordered! — the signatures already match);
• buffer elimination needs the same non-oversampled + complex guard, so EnvGrid radial
  runs would qualify (RealGrid radial runs are oversampled and would keep the general
  path);
• the linop/norm factorisation is the same separable k²(ω) × k⊥² structure — a
  smaller memory win at (nω, Nr), but the same removal of per-step array traffic;
• batched Raman batches over transverse points either way.

Separately, and independent of this branch: the examples build Hankel.QDHT(R, N, dim=2), which routes through Hankel.dot!'s permutedims path and allocates two
field-sized arrays per call — four per RHS evaluation, ~24 per step. Laying the radial
arrays out as (Nr, Nt) with dim=1 would make the transverse transform
allocation-free. That looks like the single biggest available radial win.

Modal gets the solver improvements automatically, but the transform work doesn't map
onto it: TransModal's cost is the Cubature integration (up to mfcn RHS-equivalent
evaluations per call) and its arrays are (nω, nmodes), so buffer count isn't the
constraint. The analogous win there would be in the integration itself — a separate
piece of work.

GPU. Nothing here is CUDA-specific, but the array-level formulation is deliberately
GPU-shaped: whole-array broadcasts instead of scalar column loops, AbstractFFTs
planning, similar-based allocation, and separable operators stored as small factor
arrays. A kernel benchmark of the propagation primitives on an A40 (FP64) suggests a
GPU backend would be an additive, dispatch-based change rather than a rewrite — the
remaining CPU-specific pieces are the chunked threading helpers, the two column-loop
kernels in the fast path, and the scalar reduction in the error norms.