Repository navigation
madspace: fix SIMD mode - #222
Open
theoheimel wants to merge 26 commits into
Open
theoheimel wants to merge 26 commits into
theoheimel wants to merge 26 commits into
Conversation
- 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
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
marked this pull request as ready for review
October 17, 2026 13:34
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.
SIMD mode selection (user-facing, breaking)
simd_vector_sizeis replaced bymadspace_cpu_mode. It takesauto,scalar,simd_128(ARM),simd_256(AVX2),avx512yorsimd_512, matching the names used bycpu_modefor madmatrix. It defaults toautoand invalid values are rejected. It is also a gridpack option (--madspace_cpu_mode).autopickssimd_512, thensimd_256on x86. On ARM it picksscalar, because the vectorized mode is currently slower there (about 20% forg g > t t~ g gwith dummy matrix elements on an Apple M3).avx512yis never chosen automatically, since it was slower thansimd_512in the x86 benchmarks.set_simd_vector_size(int)is replaced byset_simd_mode(str)andsupported_simd_modes(). The environment variableSIMD_VECTOR_SIZEis replaced byMADSPACE_SIMD_MODE.auto, like the run card.set_simd_modechecks the mode right away and throws if the CPU backend was already loaded with a different mode; previously the call was silently ignored.libmadspace_cpu_simd_256, ...). ARM is detected from the target processor instead ofAPPLE.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, andcvttpd_epi64. The build enables AVX2 and FMA explicitly so that it compiles with GCC.Performance: contiguous SIMD loads and stores
singleininstruction_set.yaml(masses, widths, cut bounds, channel tables, ...) now have their own kernel parameter types,FSIn/ISIn, and are read with broadcast loads. A generatedkernel_checks.incchecks at compile time that the kernel signatures match the yaml; a mismatch would silently force the scalar path.Bug fixes
SIMD backends:
fmahad its arguments in the wrong order.BVec(bool)was wrong on ARM and AVX2.vscatteron AVX2 used an aligned store on an unaligned buffer.Kernels and graph:
sqrthad the wrong sign.reduce_productdid not count zeros correctly.probsgradient ofsample_discrete_probs_inversewas overwritten instead of accumulated.cut_all/cut_anyread their bounds through a view of the wrong rank.sub(0, x)was wrongly folded tox. Stacks of scalar literals are now folded into a single literal.Runtime and driver:
contiguous_dimscounted non-leading contiguous dimensions.tensor_addinto a broadcast target (e.g. the gradient of a global) only added the first batch entry.Tests
madspace/tests/run_simd_modes.pyruns 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.pychecks 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 arecpu_only. The gradient tests also run on the GPU CI device.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_inversenow uses a mapping with a mild pt cut, because the inverse without cuts is ill-conditioned for extremely soft gluons.