solprop.backends.numba_cpu.backend#

The CPU backend: the seam, the stage driver and the statistics.

NumbaCpuBackend steps the same problems as solprop.backends.numpy_ref.backend.NumpyRefBackend and agrees with it to about 1e-15 rather than to the last bit (docs/dev/backends.md section 2). It inherits its step control from solprop.backends.driver.StepDriver – march, replay, landing, history and the noise models are the same code on every backend – and adds what is its own: the fused kernels of solprop.backends.numba_cpu.kernels, the FFTW plans and thread pools of solprop.backends.numba_cpu.fft, buffers allocated once at construction and never again, and the four backend-owned statistics computed in compiled passes rather than through the reference’s per-plane re-synthesis.

Nothing is allocated inside a step. Every buffer is created here, named in NumbaCpuBackend.memory_ledger(), and written in place for the life of the run; the plans are bound to those buffers, so a stage is a sequence of kernel calls, in-place GEMMs and execute() calls with no array ever handed between them.

The compiled step driver of docs/dev/backends.md section 3.5 is not built. The Python-level driver’s dispatch overhead was measured below the run-to-run noise floor on both target problems, and the deferral is by name.

Classes

NumbaCpuBackend(problem, *[, grid, modes, ...])

Propagate one Kerr, Raman and plasma problem with fused compiled kernels.

class solprop.backends.numba_cpu.backend.NumbaCpuBackend(problem, *, grid=None, modes=None, profiles=None, threads=None, fftw_threads=None, transverse_fftw_threads=None, blas_threads=None, plan_flag=None, regime=None)[source]#

Bases: StepDriver[Complex[ndarray, 'n_batch n_modes n_pol n_omega']]

Propagate one Kerr, Raman and plasma problem with fused compiled kernels.

The execution a bake-off of candidate backends chose (docs/dev/backends.md section 5). It implements solprop.backends.protocol.Backend and steps a single problem (max_batch = 1) on either grid: mode-averaged or multimode on the fixed transverse quadrature, scalar or vector-isotropic.

Stage representations. Constant and parameter-expanded operators use the per-stage bundle of host scalars. A z-tabulated Phi or prefactor uses preallocated host stage arrays and the reference’s Clenshaw evaluation order through solprop.backends.numba_cpu.tables.HostTable. Spatial constant operators use the separable representation. Device readiness is not required by this CPU backend.

Parameters:

problemsolprop.translate.problem.KernelProblem: The problem to step. It is validated and refused by name here.
gridsolprop.grid.TimeGrid, optional: The grid the state lives on, which the statistics need and the contract does not carry.
modescollections.abc.Sequence, optional: The propagated modes, in state order, for the statistics that need them.
profilescollections.abc.Mapping, optional: The z-profiles keyed by statistic name.
threadsint, optional: Override for the kernels’ thread pool. None takes the regime’s choice and then the SOLPROP_NUMBA_THREADS environment variable.
fftw_threadsint, optional: Override for FFTW’s pool, or SOLPROP_FFTW_THREADS.
transverse_fftw_threadsint, optional: Override for Cartesian transverse FFTs, or SOLPROP_TRANSVERSE_FFTW_THREADS. Otherwise apply the existing FFT size rule to that plan independently of the temporal plan.
blas_threadsint, optional: Override for the BLAS pool, or SOLPROP_BLAS_THREADS.
plan_flagstr, optional: The FFTW planning effort, or SOLPROP_FFTW_PLAN. "measure" by default.
regime{‘small’, ‘large’}, optional: Override for the threading regime, or SOLPROP_NUMBA_REGIME. None reads it off the problem’s element count, which is what a run should do; forcing it is for a measurement that wants a problem below the threshold stepped as a large one, so that it can be measured against thread count at all.

Raises:

BackendError: Naming what about the problem this backend does not execute.

Examples

>>> from solprop.backends.numba_cpu import NumbaCpuBackend
>>> NumbaCpuBackend.name, NumbaCpuBackend.max_batch
('numba_cpu', 1)

close()[source]#

Release the buffers and give the BLAS pool back. Idempotent.

helpers = Helpers(name='numba')#

max_batch = 1#

memory_ledger()[source]#

Return the size of every buffer the backend holds, in bytes.

The names are the reference’s for the buffers both backends have, so that the budget of docs/dev/architecture.md section 8 is measured against the same rows on either; the plan work areas are named beside them, which the reference has none of.

Returns:

dict: Buffer name to bytes.

metadata()[source]#

Return what this run’s execution was, for the output’s metadata.

The thread configuration alone was measured to move one step by a factor of nine, so a result whose pools are not recorded is not reproducible.

The wisdom digest is taken at the first record and kept: every plan this backend uses is built when it is constructed, so the wisdom they were chosen from cannot change afterwards, and exporting the wisdom again at every plane was most of the record’s cost (SEP 25 section 6.3).

Returns:

dict: The base record, plus the pools, the plan flag, the wisdom digest and the specialisation.

name = 'numba_cpu'#

property pass_timers: PassTimers | None[source]#

The opt-in per-pass timers of a free-space evaluation (SEP 28 phase 0).

Held when SOLPROP_PASS_TIMERS=1 was set at construction (solprop.backends.numba_cpu.pass_timers), None otherwise. They record the Cartesian, line and planar evaluation’s passes and the steppers’ traversals around them; reset them before a timed window.

Returns:

PassTimers or None: The timers, or None when they are off.

propagator_default = 'recompute'#

set_initial(field, noise, *, z=0.0)[source]#

Load a state, and forget the stage the step before it left behind.

Parameters:

fieldnumpy.ndarray: The state in the canonical layout.
noisesolprop.noise.NoiseRealisation or None: The realisation, or None to take the contract’s.
zfloat, optional: Where the state belongs, m. A positive value is a resume.

stats(spec)[source]#

Collect the requested statistics at the current position.

The statistics of the saved field go to solprop.stats.collect(), as every backend’s do; the four the catalogue marks backend are formed here from this backend’s own buffers and its compiled passes.

Parameters:

specsolprop.stats.StatsSpec: Which statistics and which windows.

Returns:

dict: Small host arrays keyed by statistic name, each with a leading batch axis.

Raises:

BackendError: If the backend is closed, no initial state was loaded, no grid was given at construction, or a backend statistic needs something this run does not have.