Skip to content

madspace: fix SIMD mode - #222

Open
theoheimel wants to merge 26 commits into
mainfrom
fix-madspace-simd
Open

theoheimel wants to merge 26 commits into
mainfrom
fix-madspace-simd

Conversation

@theoheimel

@theoheimel theoheimel commented Oct 7, 2026 •

Copy link
Copy Markdown
Contributor

Parts of the madspace SIMD mode gave wrong results, and where it did work it was slower than it should be. This PR fixes it so that every SIMD mode reproduces the scalar results. It also makes the vectorized kernels fast, aligns mode selection with madmatrix, and adds tests that run in every mode.

  • fix bugs so that the results match the scalar computation
  • add tests so this doesn't break again
  • make the run card parameter for the madspace SIMD mode consistent with the one for madmatrix
  • performance: replace strided loads with contiguous ones, falling back to scalar where that isn't possible
  • add an avx512y mode, like madmatrix

SIMD mode selection (user-facing, breaking)

  • The run card option simd_vector_size is replaced by madspace_cpu_mode. It takes auto, scalar, simd_128 (ARM), simd_256 (AVX2), avx512y or simd_512, matching the names used by cpu_mode for madmatrix. It defaults to auto and invalid values are rejected. It is also a gridpack option (--madspace_cpu_mode).
  • auto picks simd_512, then simd_256 on x86. On ARM it picks scalar, because the vectorized mode is currently slower there (about 20% for g g > t t~ g g with dummy matrix elements on an Apple M3). avx512y is never chosen automatically, since it was slower than simd_512 in the x86 benchmarks.
  • API: set_simd_vector_size(int) is replaced by set_simd_mode(str) and supported_simd_modes(). The environment variable SIMD_VECTOR_SIZE is replaced by MADSPACE_SIMD_MODE.
  • The C++ default with nothing set is auto, like the run card.
  • set_simd_mode checks the mode right away and throws if the CPU backend was already loaded with a different mode; previously the call was silently ignored.
  • Backend libraries are now named after their modes (libmadspace_cpu_simd_256, ...). ARM is detected from the target processor instead of APPLE.

New avx512y mode

AVX512 F/VL/DQ instructions with 256-bit vectors, like the madmatrix avx512y mode (simd_x86_256y.hpp). It uses mask registers, native 64-bit gathers and scatters, and cvttpd_epi64. The build enables AVX2 and FMA explicitly so that it compiles with GCC.

Performance: contiguous SIMD loads and stores

  • The lanes of a vector are consecutive batch entries, so they are read and written with plain unaligned vector loads and stores instead of gathers and scatters.
  • A kernel only runs vectorized if every batched argument has batch stride one. Otherwise the whole batch falls back to the scalar kernel, so results never depend on a strided path.
  • Inputs that are marked single in instruction_set.yaml (masses, widths, cut bounds, channel tables, ...) now have their own kernel parameter types, FSIn/ISIn, and are read with broadcast loads. A generated kernel_checks.inc checks at compile time that the kernel signatures match the yaml; a mismatch would silently force the scalar path.
  • Literals are stored padded to the vector width, so the vectorized kernels can read them like batched tensors. When they are returned as runtime outputs, they are returned unpadded.
  • External arguments of the CPU runtime (inputs and output gradients, e.g. row-major numpy arrays or expanded torch gradients) are made batch-contiguous once on entry. This keeps the results independent of the caller's memory layout.
  • Unused gradients of literal and integer inputs are written to batch-sized scratch memory. Otherwise they would be broadcast outputs, which force the scalar kernel on a single thread.
  • With a single flavor, graph building skips gathers from constant flavor tables and multiplications by a weight that is always one.

Bug fixes

SIMD backends:

  • ARM fma had its arguments in the wrong order.
  • BVec(bool) was wrong on ARM and AVX2.
  • AVX2 integer add, subtract and negate used 32-bit instructions on 64-bit lanes.
  • Converting doubles to integers rounded instead of truncating, on AVX2 and AVX512.
  • vscatter on AVX2 used an aligned store on an unaligned buffer.

Kernels and graph:

  • The backward of sqrt had the wrong sign.
  • The backward of reduce_product did not count zeros correctly.
  • The probs gradient of sample_discrete_probs_inverse was overwritten instead of accumulated.
  • cut_all/cut_any read their bounds through a view of the wrong rank.
  • The t-channel cut bounds in the inverse could be NaN for massless particles (negative squared masses from rounding).
  • The header generator marked every instruction as differentiable.
  • sub(0, x) was wrongly folded to x. Stacks of scalar literals are now folded into a single literal.

Runtime and driver:

  • contiguous_dims counted non-leading contiguous dimensions.
  • tensor_add into a broadcast target (e.g. the gradient of a global) only added the first batch entry.
  • Use-after-free of matrix element inputs with the dummy matrix element in multithreaded runs.
  • Event files: the header-shape check was inverted, the header size was always zero, and the header was rewritten when loading a file instead of when appending to it. The event field type is now an enum.
  • Freezing the max weight depended on thread timing, so seeded runs weren't reproducible. It now uses the expected unweighted count of the generate stages, which commit in a fixed order.

Tests

  • madspace/tests/run_simd_modes.py runs the unit tests once for each SIMD mode the machine supports. The wheel CI now builds with SIMD enabled and uses it.
  • test_simd_layout.py checks that the results don't depend on the memory layout (row/column-major inputs, broadcast tensors, small and large literals, literal outputs), and covers single inputs, constant folding and gradients into broadcast globals. The layout and graph tests are cpu_only. The gradient tests also run on the GPU CI device.
  • New mg7 acceptance tests (test_*_simd_consistency_mg7): seeded generation in every supported mode. VEGAS cross sections must agree with scalar up to rounding, madnis ones within their uncertainties, and gridpacks trained in scalar mode must reproduce in every mode.
  • test_process_inverse now uses a mapping with a mild pt cut, because the inverse without cuts is ill-conditioned for extremely soft gluons.

theoheimel and others added 21 commits October 7, 2026 13:48
- add supported_simd_modes() to madspace (C++ and Python)
- add madspace/tests/run_simd_modes.py, which reruns the unit tests per mode
- run it in the madspace wheel CI, and build the macOS wheels with SIMD
- add mg7 tests comparing VEGAS and madnis runs across SIMD modes
- build the madspace install used by the mg7 CI with SIMD
- make madspace_cpu_mode a gridpack parameter and add a --madspace_cpu_mode
  option to the gridpack's generate_events (previously gridpacks always ran
  madspace in scalar mode)
- add test_madnis_gridpack_simd_consistency_mg7, which runs a gridpack trained
  in scalar mode in every SIMD mode with the same tight tolerance as VEGAS
- t_cut_bounds_etmin: clamp squared masses at zero. In the inverse they are
  recomputed from momenta and can be slightly negative for massless particles,
  which gave NaN cut bounds and a silent fallback to the kinematic bounds, so
  forward and inverse used different t ranges
- test_process_inverse: use a mapping with a 1 GeV jet pt cut and check events
  with non-zero weight. Without cuts, rare extremely soft gluons make the
  inverse ill-conditioned and the test flaky
- clang-format of pre-existing code in kinematics.hpp
…adcast tensor_add, discrete probs gradient, generator differentiable flag

- cut_all/cut_any declared their scalar min/max inputs as rank-1 views and
  read min[i], indexing through an uninitialised stride slot
- the explicit-stride Tensor constructor counted every contiguous dimension
  instead of only the leading run used by flat_view
- tensor_add into a broadcast (batch size one) target only added row 0 of
  the source; it now iterates over the source batch
- backward_kernel_sample_discrete_probs_inverse overwrote probs_grad instead
  of accumulating it, wrong for broadcast (global) probabilities
- generate_headers.py registered every instruction as differentiable

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- the identity rule returned x for sub(0, x) instead of -x; it now only
  drops x - 0
- the identity rule picked the wrong argument (and threw) when the
  variable argument was a tensor literal; literals are now checked by type
- folding only runs on float literals instead of throwing on int literals
- stack of scalar literals (e.g. leaf masses for rambo/chili) is folded into
  a single tensor literal

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…lavor

With a single flavor, flavor_id is a full(0) tensor and the per-flavor
tables (remap, factors, mirror option count, subprocess remaps) were looked
up with gathers on constant tensors. Read the table entry in C++ instead,
and multiply by the flavor factor as a scalar literal. weight_after_cuts is
a full(1) placeholder unless flavor or mirror sampling contribute to it, so
leave it out of the weight product in that case.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ernel type

Vectorized kernels used gathers/scatters for every load and store, also for
the common case of batch stride one, because inputs can be broadcast.

- the SIMD backends now load and store the lanes of a vector with plain
  contiguous (unaligned) loads/stores, and have a broadcast load
- VectorizedTensorView assumes lane-contiguous data and advances by vec_step
  per vector (zero for padded literals, prepared for the next commit)
- tensor_foreach_impl only runs the vectorized kernel if every batched
  argument has batch stride one, otherwise the scalar kernel handles the
  whole batch
- inputs marked single in instruction_set.yaml use the new kernel parameter
  types FSIn/ISIn, which map to SingleVectorizedTensorView (broadcast loads,
  batch index ignored) on SIMD and to the usual views otherwise
- mark the inputs that are literals at every call site as single (physics
  parameters, cut bounds, channel weight tables, index tables, ...)
- generate static_asserts that the single kernel parameters match the yaml

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Scalar literals and small tensor literals (at most 256 stored elements
including the lanes) are stored with every element repeated simd_vec_size
times, one copy per vector lane. They keep batch size one and batch stride
zero, so scalar code sees an ordinary broadcast tensor, but the vectorized
kernels can read them with contiguous loads (vec_step zero) instead of falling
back to the scalar kernels. Copies and views of the literal are not padded.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…atch memory

The backward pass allocated a gradient with the input's shape for every
input, so a literal input got a batch-size-one (broadcast) gradient. That
forced the backward kernel onto the scalar path on a single thread, although
the gradient is never used. Gradients of literals and integer intermediates
now go to batch-sized scratch memory. Gradients of single inputs (which can
be trainable globals) are unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Inputs and output gradients passed to the CPU runtime (e.g. row-major numpy
arrays or expanded torch gradients) can have a batch stride other than one.
The vectorized kernels now leave those to the scalar kernels, which round
slightly differently, so results depended on the memory layout of the
arguments (seen in test_color_ordered_pair_cuts in simd_256 mode). Copy such
arguments once when the runtime is called.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ast gradients

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…matrix element

op_matrix_element made every input contiguous before checking whether the
matrix element is called at all. For a non-contiguous input (e.g. a full
tensor) the concurrent runtime starts an asynchronous copy, and with the flat
dummy matrix element (or an empty batch) the function returned and freed the
copy target while the copy was still running, corrupting the heap in
multithreaded runs. Prepare the inputs only when the matrix element is called.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
# Conflicts:
#	madgraph/iolibs/template_files/mg7/run_card.toml
#	madgraph/various/RunCardLO_to_MG7_mapping.md
#	madgraph/various/banner.py
Contiguous-only SIMD loads/stores, single inputs as a separate kernel
type, padded literals, more constant folding, and the bug fixes found
along the way (cut_all/cut_any bounds, contiguous_dims, broadcast
tensor_add, discrete probs gradient, sub(0, x) folding, use-after-free
of matrix element inputs with the dummy matrix element).

Conflicts with main resolved by keeping both sides; boost_to_frame from
main now takes its single frame_mask input as ISIn.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
FieldLayout::type was a const char* compared by pointer against the
FieldLayout::i32/f64 string constants. Whether two uses of the same literal
share an address is unspecified, and when they did not (e.g. with
AddressSanitizer) DataLayout fell through to an "unknown type" branch that
constructed a std::logic_error without throwing it, leaving the field size
uninitialized: EventHistograms.fill and the systematics binding then
allocated garbage-sized buffers or wrote out of bounds. The type is now an
enum, with its size and numpy type string derived from it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The append and load modes of EventFile could not work: read_event_header
rejected every header because the shape check was inverted, derived the
particle count by dividing by a hard-coded 5 (and divided by zero for
layouts without particle fields, like the weight files), and returned a
header size of 0. The constructor then rewrote the header in load mode
(which is read-only) instead of append mode, from the current stream
position, with the particle count argument instead of the one read from
the file, and without moving to the end of the file before appending or to
the first event before reading.

Nothing in madspace opens event files in these modes yet; a create, append,
load round trip and np.load of the result were checked with a standalone
program.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
update_max_weight() stopped raising the max weight once count_unweighted
passed freeze_max_weight_after. count_unweighted is advanced by the unweight
stages, and whether a job's unweight stage commits before or after the next
generate stage of the same channel depends on thread timing, so the freeze
point, and with it the whole run, could differ between runs with the same
seed (seen for p p > e+ e- j with 8 threads: 4 different cross sections in
6 runs under oversubscription, all identical with the freeze disabled).

Freeze on the expected unweighted count instead, sum of
min(|w| / max_weight, 1) over the committed generate stages, which commit
in a fixed order. It is accumulated in the loop over the weights that
update_max_weight() already does, so scheduling and throughput are
unchanged; the same runs are now bit-identical.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@theoheimel
theoheimel added this pull request to stack #299 October 17, 2026 12:32
# Conflicts:
#	.github/workflows/madspace.yml
#	madspace/include/madspace/driver/channel_generator.hpp
#	madspace/src/driver/channel_generator.cpp
#	madspace/src/kernels/kinematics.hpp
#	madspace/src/phasespace/integrand.cpp
#	madspace/src/python/madspace.cpp
#	madspace/tests/conftest.py
#	madspace/tests/test_processes.py
…d_256 on x86

The vectorized mode is currently slower on ARM (about 20% for g g > t t~ g g
with dummy matrix elements on an Apple M3), so auto keeps scalar there.
…npadded literal outputs

The C++ default (no set_simd_mode, no MADSPACE_SIMD_MODE) is now auto, like the
run card. set_simd_mode validates the mode immediately and throws if the CPU
backend was already loaded with a different one, instead of silently ignoring
it. Literal outputs of the CPU runtime are returned without the SIMD padding.
…t device

The layout and graph building tests run the CPU runtime on numpy arrays and
are marked cpu_only. The gradient tests cover kernels shared with the GPU and
now use the context of MADSPACE_TEST_DEVICE. Also tests that literal outputs
are not padded.
@theoheimel
theoheimel marked this pull request as ready for review October 17, 2026 13:34

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant