[theoheimel]

Adds single-precision support to madspace, mainly so the MadNIS networks can be evaluated and trained in float32. This will enable faster MadNIS training and sampling, and reduce gridpack sizes by a factor 2. In terms of unweighting efficiency, i would not expect a difference, but this will need some investigation especially for longer trainings.

Everything else (including normalizing flow jacobians) remains fp64. However, these changes can also be used as a basis to investigate moving more of phasespace to fp32 in the future.

Stacked on #222 (fix-madspace-simd).

Changes

• Data types: DataType is now int32 / float32 / float64 / batch_sizes. This is a clean break: graph JSON, the Python enum and npy files use the new names, and npy and DLPack gain <f4/f32 support. Tensor views now check that the element type matches the tensor dtype.
• Instruction set: an entry can have precision: [float32, float64]. The generator expands it into X_float32 and X_float64, each with its own opcode. The builder method X dispatches on the argument dtype, so frontend code doesn't change. A float64 literal passed to a float32 call is converted first.
• Instructions: matmul (sgemm/dgemm, and the same on cuBLAS/rocBLAS) and all activations exist in both precisions. New convert_float32 and convert_float64 instructions are differentiable.
• Kernels: the CPU and GPU kernel policies are templated on the float type, so float32 and float64 share the kernel code.
• SIMD: float32 vectors have twice the lanes of float64 (NEON 4, AVX2/AVX512Y 8, AVX512 16), with int32 indices and masks of the same width. Jobs are aligned to the larger lane count.
• Training: the global gradient buffer and the Adam step follow the parameter dtype, and all trainable parameters must share one dtype. The gradient clipper computes the norm in float64.
• MLP: MLP(..., precision=...) converts at its boundaries. Flow, DiscreteFlow and ChannelWeightNetwork forward this as subnet_precision / precision.
• MG7: new run-card setting madnis.nn_float_precision (float32 by default, or float64). MG7 MadNIS runs therefore use single-precision networks unless the run card says otherwise.

Fixes found along the way

• The activation backward kernels overwrote the input gradient instead of adding to it.
• The Python shape check for single (non-batch) inputs compared the wrong dimensions.
• The gather and select backward kernels used float views for int gradients. The new dtype check caught this.
• The bundled OpenBLAS was built without single precision. On Linux this made the doc examples segfault.
• The .pyi stub generation needs to know where DataType lives to render its default arguments.
• CPU random_int allocated a float64 tensor for its int32 indices, which the new dtype check rejects (replay-buffer sampling aborted on CPU).
• += on GPU tensor views raced when many threads accumulate into a broadcast (zero-stride) gradient. Whether a view needs atomicAdd is now decided at kernel launch from its strides, so only broadcast gradients are atomic and per-event gradients keep plain stores.

GPU launch overhead

float32 networks initially needed extra convert kernels and cuBLAS split-K/dot kernels (training is launch-bound). This PR adds:

• single-kernel bias gradient (replaces ones+gemv),
• a GPU-backend lowering pass fuse_conversions (matmul_to_float64, select_with_float32, rqs_{forward,inverse}_with_float32, fused backward bias-grad+convert). The compute graph stays portable; MADSPACE_NO_FUSION=1 disables the pass,
• fix: op_unweight returned float64 empty indices for empty batches.

Results

Physics: no degradation. 100k MadNIS training batches on one A100, 3 seeds per setting (mean of 3; exact reproducibility is not expected).

process	precision	training wall [s]	RSD	unweighted / generated	generation wall [s]
g g > t t~ g g	float64	658	0.353	0.326	6.9
g g > t t~ g g	float32	698	0.346	0.324	7.0
g g > t t~ g g g	float64	2247	0.735	0.092	569
g g > t t~ g g g	float32	2297	0.730	0.092	575

RSD, unweighting efficiency and weighted-event throughput agree within the seed-to-seed scatter (ttggg seed 3 is a bad run for both precisions). These numbers are from before the fusion pass.

Speed after fusion (g g > t t~ g g, 64 hidden nodes, 2000 training batches, single run per GPU, same build):

GPU	float64	float32 unfused	float32 fused
A100 80GB	19.9 s	21.7 s	20.4 s
H200	17.0 s	15.5 s	14.6 s
A10	54.6 s	33.9 s	32.8 s

Kernel launches (A100): 2.57M (f64) / 2.88M (f32 unfused) / 2.65M (f32 fused). float32 is now ~2% slower than float64 on A100 and 14% / 40% faster on H200 / A10.

Testing

• Locally (macOS arm64): all 1789 madspace tests pass in scalar and NEON SIMD modes. The new tests cover:
  • float32 activations and matmul, forward and backward, against torch;
  • the conversion instructions;
  • float32 MLP training against torch for both the built-in Adam and torch Adam;
  • float32 flow outputs and gradients;
  • saving and loading float32 globals.
• MG7 e+ e- > w+ w- with MadNIS: the cross section is 2.6902468 pb in float64 and 2.6902465 pb in float32.
• CI: the CUDA/HIP builds compile. The wheel builds and tests (Linux x86 and macOS arm64, in every SIMD mode the runner supports) and the docs examples pass.
• New madspace/tests/test_fused_float32.py: fused vs unfused values and gradients, lowering graph structure.
• CPU suite passes (2050 tests, macOS arm64); new tests also pass on A10 and H200 (before the launch-time atomic change, which still needs a GPU run).
• Known, unrelated: test_simd_layout.py::test_global_probability_gradient_is_summed_over_events fails on GPU.

[oliviermattelaer]

fyi

#Should fix

• Backward through stack/cat on float32 fails. The backward of stack, unstack, pop, batch_cat, batch_split and cat still allocates float64 gradients. Their forward passes accept float32, so backward on float32 data throws “invalid dtype in add” or a view dtype error. MadNIS doesn’t hit this today because the networks only use matmul and activations, but any float32 graph built with these ops will. The fix is to allocate with the input’s own dtype.
• Old gridpacks stop loading. Renaming float/int and the NN instructions means gridpacks and graph JSON saved before this PR fail to load. The PR calls this a deliberate clean break, but accepting the old names on load would cost very little.
• GPU tensor_add drops most of a broadcast gradient. This one predates the PR, but the PR rewrote that function without fixing it. When the target is a broadcast gradient (a global’s), only the first event is added, where the CPU version sums the whole batch.

#Minor

• test_flow_float32_subnets leaves fixed-name globals in the shared default context.
• The gradient clipper converts the whole float32 gradient to float64 and back on every step, when only the norm needs float64.
• The float32/float64 branches in the CPU and GPU device ops are near-copies.
• The new dtype_of() isn’t used in two places where it applies.
  nn_precision in launch.py has an error branch that can never run.

[theoheimel]

@oliviermattelaer Thanks, I've fixed the issues (except for the gridpack one, since we're not aiming for cross-version compatibility and are still alpha anyways)

[oliviermattelaer]

Looks like the AI agent is not yet happy about this one (or I just put the level of review too high for this one)
The first one (and #5), is like to drop but the rest (especially 8 and 9) are relevant.

Finding numbering follows the original review; the withdrawn race finding (2) is omitted.

1. DataType JSON names renamed, legacy spellings rejected

madspace/src/compgraphs/type.cpp:313 (correctness)

from_json(DataType) accepts only int32/float32/float64/batch_sizes. A serialized Function or gridpack written before this branch contains "int" or "float" and fails to load with invalid data type. from_json(Value) also rejects float32 literals with the same generic error.

3. Inconsistent gradient dtype for integer inputs (low, unconfirmed in practice)

madspace/src/kernels/operations.hpp:107 (correctness)

grad_dtype() gives integer inputs a float64 gradient in stack/unstack/cat/split/pop/rqs_reshape, while backward_batch_foreach allocates input.dtype() (int32) for integer inputs that are not discarded (integer function inputs). If one int function input reaches both kinds of op, they share local_grads[input_index] and the second user reads it through a mismatched view, which check_view_dtype turns into an exception.

Reachability is narrow: it needs, for example, a batch_cat of an int input with an int output of a differentiable op. These ops already used float64 for ints before this branch, so nothing regressed. No concrete failing graph was found.

4. SIMD mode: env var shadowed, re-launch throws

madspace/src/driver/backend.cpp:207 (correctness)

MadgraphProcess.init_backend() always calls ms.set_simd_mode(run_card['run']['madspace_cpu_mode']), default "auto". That overrides MADSPACE_SIMD_MODE, so the variable has no effect under mg7. The old SIMD_VECTOR_SIZE variable is gone without notice. set_simd_mode throws if the backend is already loaded with a different mode, so a second in-process launch after editing madspace_cpu_mode raises can not switch the madspace SIMD mode.

5. nn_float_precision defaults to float32 for everyone

madgraph/various/banner.py:7030 (altitude)

Existing run cards, and new processes without the key, switch MadNIS training from float64 to float32 silently: weights, activations and Adam state. Convergence and efficiency can change on long or hard integrands, and no test compares the two precisions' cross-section or efficiency. Either keep float64 as the default until validated or call it out in the release notes.

6. CPU and GPU matmul_to_float64 round differently

madspace/src/cpu/runtime.cpp:383 (efficiency)

The CPU version pre-fills the bias and runs a float32 gemm with beta=1, then converts. The GPU version runs the gemm with beta=0 and adds the bias in kernel_bias_to_double. Results differ in the last bits, so fused vs. unfused (MADSPACE_NO_FUSION) and CPU vs. GPU are not bit-comparable. On CPU the op also allocates a float32 temporary and does a second full pass, so it is no cheaper than matmul followed by convert_float64; the fusion lowering is only applied by the GPU runtime.

7. 32-bit gather offsets can overflow

madspace/src/cpu/simd_x86_256.hpp:278 (correctness)

mem_indices32 builds offsets with mullo_epi32 on set1_epi32(batch_stride) and index_stride, and _mm256_i32gather_ps takes signed 32-bit offsets. For very large batches of wide tensors (offsets beyond 2^31 elements) they wrap negative and the gather reads out of bounds. The size_t strides are also truncated by set1_epi32. The scalar path and the double-precision path do not have this limit.

8. GCC-only flag set unconditionally on the avx512y target

madspace/CMakeLists.txt:456 (correctness)

-mprefer-vector-width=256 is passed regardless of compiler. Clang warns about an unused argument, or errors under -Werror. The branch already has a commit "fix avx512y build with GCC", so this target is compiler-fragile. Gate the flag on CMAKE_CXX_COMPILER_ID STREQUAL "GNU".

9. GPU float32 fusion kernels lack CI-runnable tests

madspace/src/gpu/runtime.cu:155 (test-coverage)

New CUDA/HIP code has no test that runs without a GPU: gpublas_gemm<float>, kernel_bias_to_double, kernel_bias_grad, kernel_bias_grad_convert, backward_matmul_impl<float, true> and the fuse_conversions lowering. test_fused_float32.py only exercises the CPU path. A bug in the fused select/rqs backward or the in-place output_grad = output_grad_f swap would only surface on a GPU runner, with MADSPACE_NO_FUSION as the only escape hatch.

[theoheimel]

@oliviermattelaer Thanks! I didn't change 1, 4, 5. Claude says 8 is not correct because recent clang and Apple clang support this flag. The other issues should be fixed, including the failing GPU tests from main.

[oliviermattelaer]

They are still a couple of (important) issues:

• Buffered MadNIS training on CPU is broken. In madspace/src/cpu/runtime.cpp, the CPU version of random_int stores whole numbers in a float64 tensor. The PR’s new check that a tensor is read with its own type now makes this throw on every call. MadNIS training with a replay buffer calls random_int when it samples from the buffer, so it now aborts on the first sampling step. The GPU version already uses dt_int32, and the fix is the same one-line change on the CPU side.
  Gridpacks built before this PR will no longer load. Two things changed in madspace/src/compgraphs/type.cpp and the instruction set:
  Saved graphs label their values "float" and "int", which are no longer accepted.
  Instructions like matmul and relu were renamed to matmul_float64 and relu_float64.

• The default precision doesn’t match the PR description. The run card, the docs and the banner test all make nn_float_precision default to float32. The description says the default is float64. As written, every MG7 MadNIS run silently switches to single-precision networks. Either the default or the description needs to change.

• GPU speed: every gradient accumulation on the GPU now uses an atomic add, even where each thread writes its own element. This slows training and makes some summed gradients vary slightly between runs.

• Flaky test: test_flow_float32_subnets shares one random generator across tests, so whether it passes depends on test order. The PR description mentions this failure.

Below this it is less important (so drop them if needed):

• Unreachable CPU code: about 250 lines of fused CPU kernels only run if someone calls fuse_instructions from Python, because the CPU runtime never applies the fusion.

• Slow conversions on CPU: the float32/float64 conversions at every network boundary run element by element instead of with SIMD.

• Fusion pass cost: the fusion pass builds 12 strings for every instruction in the graph while it runs.

[theoheimel]

@oliviermattelaer thanks, I've fixed the issues