From 9d0762148a4eb85677830dabbffd11fea29a3137 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 7 Oct 2026 13:25:42 +0200 Subject: [PATCH 01/23] SIMD bugfixes, NEON-specific and general --- madspace/instruction_set.yaml | 1 + madspace/src/cpu/simd_arm.hpp | 2 +- madspace/src/cpu/tensor.hpp | 2 +- 3 files changed, 3 insertions(+), 2 deletions(-) diff --git a/madspace/instruction_set.yaml b/madspace/instruction_set.yaml index 43738637b3..447387ef98 100644 --- a/madspace/instruction_set.yaml +++ b/madspace/instruction_set.yaml @@ -3185,6 +3185,7 @@ permute_momenta: - name: output type: [float, n, 4] desc: "permuted momenta" + vectorized: False desc: "Reorders momenta according to a stored permutation selected per event. Building block of @ref MultiChannelMapping." gather: diff --git a/madspace/src/cpu/simd_arm.hpp b/madspace/src/cpu/simd_arm.hpp index 51a38ac269..11e67de033 100644 --- a/madspace/src/cpu/simd_arm.hpp +++ b/madspace/src/cpu/simd_arm.hpp @@ -173,4 +173,4 @@ inline FVec exp(FVec arg1) { return Sleef_expd2_u10(arg1); } inline FVec log1p(FVec arg1) { return Sleef_log1pd2_u10(arg1); } inline FVec expm1(FVec arg1) { return Sleef_expm1d2_u10(arg1); } inline FVec erf(FVec arg1) { return Sleef_erfd2_u10(arg1); } -inline FVec fma(FVec arg1, FVec arg2, FVec arg3) { return vfmaq_f64(arg1, arg2, arg3); } +inline FVec fma(FVec arg1, FVec arg2, FVec arg3) { return vfmaq_f64(arg3, arg1, arg2); } diff --git a/madspace/src/cpu/tensor.hpp b/madspace/src/cpu/tensor.hpp index cb0bd7afe6..446c771fd8 100644 --- a/madspace/src/cpu/tensor.hpp +++ b/madspace/src/cpu/tensor.hpp @@ -216,7 +216,7 @@ struct get_vectorized_views { auto operator()(TArg&... args) { return std::make_tuple([&]() { if constexpr (TArg::is_scalar_view) { - return ScalarView(TParam(0.0)); + return ScalarView(TParam(typename TArg::DType(args))); } else { return VectorizedTensorView< typename TParam::VType, From 6f50fb934194c0b2d2353a9e5d0e32a22a99eecd Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 7 Oct 2026 13:48:38 +0200 Subject: [PATCH 02/23] more SIMD bugfixes, including the one proposed in #209 --- madspace/src/cpu/simd_arm.hpp | 2 +- madspace/src/cpu/simd_x86_256.hpp | 14 ++++++------- madspace/src/cpu/simd_x86_512.hpp | 2 +- madspace/src/cpu/tensor.hpp | 35 +++++++++++++++++++------------ madspace/src/kernels/math.hpp | 3 ++- madspace/src/kernels/vegas.hpp | 2 +- 6 files changed, 34 insertions(+), 24 deletions(-) diff --git a/madspace/src/cpu/simd_arm.hpp b/madspace/src/cpu/simd_arm.hpp index 11e67de033..93cb88625f 100644 --- a/madspace/src/cpu/simd_arm.hpp +++ b/madspace/src/cpu/simd_arm.hpp @@ -33,7 +33,7 @@ struct IVec { struct BVec { BVec() = default; BVec(uint64x2_t _v) : v(_v) {}; - BVec(bool _v) : v(vceqzq_u64(vdupq_n_u64(_v))) {}; + BVec(bool _v) : v(vdupq_n_u64(_v ? ~0ULL : 0)) {}; operator uint64x2_t() { return v; } uint64x2_t v; }; diff --git a/madspace/src/cpu/simd_x86_256.hpp b/madspace/src/cpu/simd_x86_256.hpp index 6c2815e25a..099bb18859 100644 --- a/madspace/src/cpu/simd_x86_256.hpp +++ b/madspace/src/cpu/simd_x86_256.hpp @@ -28,7 +28,7 @@ struct IVec { IVec() = default; IVec(__m256i _v) : v(_v) {}; IVec(int _v) : v(_mm256_set1_epi64x(_v)) {}; - explicit IVec(__m256d _v) { v = _mm256_cvtepi32_epi64(_mm256_cvtpd_epi32(_v)); } + explicit IVec(__m256d _v) { v = _mm256_cvtepi32_epi64(_mm256_cvttpd_epi32(_v)); } operator __m256i() { return v; } IVec operator+=(IVec _v) { v = _mm256_add_epi64(v, _v); @@ -41,7 +41,7 @@ struct BVec { BVec() = default; BVec(__m256d _v) : v(_v) {}; BVec(__m256i _v) : v(_mm256_castsi256_pd(_v)) {}; - BVec(bool _v) : v(_mm256_set1_pd(-1 * static_cast(_v))) {}; + BVec(bool _v) : v(_mm256_castsi256_pd(_mm256_set1_epi64x(_v ? -1 : 0))) {}; operator __m256d() { return v; } operator __m256i() { return _mm256_castpd_si256(v); } __m256d v; @@ -95,7 +95,7 @@ inline void vscatter( long long scalar[4]; __m256i vec; } indices_buf; - _mm256_store_pd(values_buf, values); + _mm256_storeu_pd(values_buf, values); _mm256_store_si256(&indices_buf.vec, indices); for (int i = 0; i < 4; ++i) { base_ptr[index_stride * indices_buf.scalar[i] + batch_stride * i] = @@ -128,7 +128,7 @@ inline void vscatter( inline void vstore(double* base_ptr, std::size_t stride, FVec values) { double values_buf[4]; - _mm256_store_pd(values_buf, values); + _mm256_storeu_pd(values_buf, values); for (int i = 0; i < 4; ++i) { base_ptr[i * stride] = values_buf[i]; } @@ -193,10 +193,10 @@ inline FVec operator-(FVec arg1, FVec arg2) { return _mm256_sub_pd(arg1, arg2); inline FVec operator*(FVec arg1, FVec arg2) { return _mm256_mul_pd(arg1, arg2); } inline FVec operator/(FVec arg1, FVec arg2) { return _mm256_div_pd(arg1, arg2); } inline IVec operator-(IVec arg1) { - return _mm256_sub_epi32(_mm256_set1_epi32(0), arg1); + return _mm256_sub_epi64(_mm256_setzero_si256(), arg1); } -inline IVec operator+(IVec arg1, IVec arg2) { return _mm256_add_epi32(arg1, arg2); } -inline IVec operator-(IVec arg1, IVec arg2) { return _mm256_sub_epi32(arg1, arg2); } +inline IVec operator+(IVec arg1, IVec arg2) { return _mm256_add_epi64(arg1, arg2); } +inline IVec operator-(IVec arg1, IVec arg2) { return _mm256_sub_epi64(arg1, arg2); } inline BVec isnan(FVec arg) { return arg != arg; } diff --git a/madspace/src/cpu/simd_x86_512.hpp b/madspace/src/cpu/simd_x86_512.hpp index 1a7c97a987..74808d83c1 100644 --- a/madspace/src/cpu/simd_x86_512.hpp +++ b/madspace/src/cpu/simd_x86_512.hpp @@ -21,7 +21,7 @@ struct IVec { IVec() = default; IVec(__m512i _v) : v(_v) {}; IVec(int _v) : v(_mm512_set1_epi64(_v)) {}; - explicit IVec(__m512d _v) { v = _mm512_cvtepi32_epi64(_mm512_cvtpd_epi32(_v)); } + explicit IVec(__m512d _v) { v = _mm512_cvtepi32_epi64(_mm512_cvttpd_epi32(_v)); } operator __m512i() { return v; } IVec operator+=(IVec _v) { v = _mm512_add_epi64(v, _v); diff --git a/madspace/src/cpu/tensor.hpp b/madspace/src/cpu/tensor.hpp index 446c771fd8..a787617455 100644 --- a/madspace/src/cpu/tensor.hpp +++ b/madspace/src/cpu/tensor.hpp @@ -313,9 +313,12 @@ inline void tensor_foreach_impl( // (the whole batch handled sequentially by one thread). Broadcast *inputs* are // only ever read, never written, in both forward and backward -- concurrent // reads of the same memory are safe, so they don't need this. + // SIMD lanes would also write to the same memory, so use the scalar kernel. + bool broadcast_output = false; for (const Tensor* out : outputs) { if (out->size(0) != batch_size) { single_job = true; + broadcast_output = true; } } @@ -333,23 +336,29 @@ inline void tensor_foreach_impl( device.foreach ( batch_size, - [flat_views, scalar_args...](std::size_t count, std::size_t offset) mutable { + [flat_views, + broadcast_output, + scalar_args...](std::size_t count, std::size_t offset) mutable { auto views = std::apply(get_views(), flat_views); std::size_t scalar_offset = offset; if constexpr (!std:: is_same_v) { - auto vectorized_views = std::apply( - get_vectorized_views(), views - ); - std::size_t vec_count = count / simd_vec_size; - std::size_t vec_offset = offset / simd_vec_size; - std::apply( - [vec_count, vec_offset](auto&&... args) { - nested_for(vec_count, vec_offset, args...); - }, - vectorized_views - ); - scalar_offset += vec_count * simd_vec_size; + if (!broadcast_output) { + auto vectorized_views = std::apply( + get_vectorized_views(), views + ); + std::size_t vec_count = count / simd_vec_size; + std::size_t vec_offset = offset / simd_vec_size; + std::apply( + [vec_count, vec_offset](auto&&... args) { + nested_for( + vec_count, vec_offset, args... + ); + }, + vectorized_views + ); + scalar_offset += vec_count * simd_vec_size; + } } std::size_t scalar_count = offset + count - scalar_offset; std::apply( diff --git a/madspace/src/kernels/math.hpp b/madspace/src/kernels/math.hpp index ad6788e97d..188a4f3bc5 100644 --- a/madspace/src/kernels/math.hpp +++ b/madspace/src/kernels/math.hpp @@ -121,6 +121,7 @@ backward_kernel_reduce_product(FIn in, FIn out_grad, FOut in_g FVal val = in[i]; auto zero_val = val == 0.; product = product * where(zero_val & (zero_count == 0), 1., val); + zero_count = zero_count + where(zero_val, IVal(1), IVal(0)); } auto zero_product = where(zero_count == 0, product, 0.); for (std::size_t i = 0; i < in.size(); ++i) { @@ -138,7 +139,7 @@ KERNELSPEC void kernel_sqrt(FIn in, FOut out) { template KERNELSPEC void backward_kernel_sqrt(FIn in, FIn out_grad, FOut in_grad) { - in_grad += where(in > 0., -0.5 * out_grad / sqrt(in), 0.); + in_grad += where(in > 0., 0.5 * out_grad / sqrt(in), 0.); } template diff --git a/madspace/src/kernels/vegas.hpp b/madspace/src/kernels/vegas.hpp index 3412202c73..06bcf797cb 100644 --- a/madspace/src/kernels/vegas.hpp +++ b/madspace/src/kernels/vegas.hpp @@ -28,7 +28,7 @@ KERNELSPEC void kernel_vegas_inverse( FVal bin_count_f(bin_count); IVal bin_index(0); for (std::size_t i = 0; i < bin_count; ++i) { - bin_index = where(input < grid[i], i, bin_index); + bin_index = where(input >= grid[i], i, bin_index); } auto left_edge = grid.gather(bin_index); auto right_edge = grid.gather(bin_index + 1); From 12b37ac62cd5bf3a2df96aac6b4e568e769f7a9e Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 7 Oct 2026 14:04:30 +0200 Subject: [PATCH 03/23] make madspace simd mode selection consistent with madmatrix --- docs/source/madspace/cpp-api.rst | 2 +- docs/source/madspace/python-api.rst | 2 +- madgraph/interface/common_run_interface.py | 6 +- madgraph/interface/tutorials/mg7.py | 2 +- madgraph/iolibs/template_files/mg7/launch.py | 2 +- .../iolibs/template_files/mg7/run_card.toml | 10 ++- madgraph/various/RunCardLO_to_MG7_mapping.md | 2 +- madgraph/various/banner.py | 21 +++--- madspace/CMakeLists.txt | 6 +- madspace/include/madspace/driver/backend.hpp | 6 +- madspace/src/driver/backend.cpp | 72 +++++++++---------- madspace/src/python/madspace.cpp | 12 ++-- tests/input_files/mg7_run_card.toml | 10 ++- 13 files changed, 78 insertions(+), 75 deletions(-) diff --git a/docs/source/madspace/cpp-api.rst b/docs/source/madspace/cpp-api.rst index 64384364ef..3c91b31b48 100644 --- a/docs/source/madspace/cpp-api.rst +++ b/docs/source/madspace/cpp-api.rst @@ -178,8 +178,8 @@ The high-level driver plus the tensor and device layer it runs on: integration, cpp/format_with_error cpp/hip_device cpp/select_combine_channel_index + cpp/set_simd_mode cpp/set_lib_path - cpp/set_simd_vector_size Utilities --------- diff --git a/docs/source/madspace/python-api.rst b/docs/source/madspace/python-api.rst index ec5d0ed19c..c387e3121a 100644 --- a/docs/source/madspace/python-api.rst +++ b/docs/source/madspace/python-api.rst @@ -164,8 +164,8 @@ The high-level driver plus the tensor and device layer it runs on: integration, python/format_with_error python/hip_device python/select_combine_channel_index + python/set_simd_mode python/set_lib_path - python/set_simd_vector_size Utilities --------- diff --git a/madgraph/interface/common_run_interface.py b/madgraph/interface/common_run_interface.py index 9872eaea00..14f621d898 100755 --- a/madgraph/interface/common_run_interface.py +++ b/madgraph/interface/common_run_interface.py @@ -1238,7 +1238,8 @@ def detect_card_type(path): '@MG5aMC', # MA5 hadronique 'run_rivet_later', # Rivet 'change particle_in_density_matrix', # density mode of reweight - 'simd_vector_size', # mg7 run_card.toml + 'madspace_cpu_mode', # mg7 run_card.toml + 'simd_vector_size', # mg7 run_card.toml (before madspace_cpu_mode) 'include_madspace', # mg7 run_card.toml ] @@ -1274,7 +1275,8 @@ def detect_card_type(path): # MadAnalysis4 is gone, but banners and process directories written # by older versions still carry this card: recognise it, ignore it. return 'plot_card.dat' - elif 'simd_vector_size' in text or 'include_madspace' in text: + elif 'madspace_cpu_mode' in text or 'simd_vector_size' in text or \ + 'include_madspace' in text: # mg7 run_card is a TOML file (madspace/MadNIS integration engine) return 'run_card.toml' elif ('gridpack' in text and 'ebeam1' in text) or \ diff --git a/madgraph/interface/tutorials/mg7.py b/madgraph/interface/tutorials/mg7.py index f5eaf7f384..b374b58564 100644 --- a/madgraph/interface/tutorials/mg7.py +++ b/madgraph/interface/tutorials/mg7.py @@ -74,7 +74,7 @@ `Cards/run_card.toml`, section by section. The ones you will actually touch are marked. - [run] seed, device (cpu/gpu), cpu_mode, simd_vector_size, + [run] seed, device (cpu/gpu), cpu_mode, madspace_cpu_mode, the thread pools, output_format, verbosity <- [beam] beam energies, PDF [generation] events, the survey iterations and target precision, diff --git a/madgraph/iolibs/template_files/mg7/launch.py b/madgraph/iolibs/template_files/mg7/launch.py index 6856861a36..eca7c93f85 100644 --- a/madgraph/iolibs/template_files/mg7/launch.py +++ b/madgraph/iolibs/template_files/mg7/launch.py @@ -344,7 +344,7 @@ def drop_closed_channels(self) -> None: self.subprocess_data = kept def init_backend(self) -> None: - ms.set_simd_vector_size(self.run_card["run"]["simd_vector_size"]) + ms.set_simd_mode(self.run_card["run"]["madspace_cpu_mode"]) def init_event_dir(self) -> None: run_name = self.run_card["run"]["run_name"] diff --git a/madgraph/iolibs/template_files/mg7/run_card.toml b/madgraph/iolibs/template_files/mg7/run_card.toml index 13cfeb2688..302fff0cd3 100644 --- a/madgraph/iolibs/template_files/mg7/run_card.toml +++ b/madgraph/iolibs/template_files/mg7/run_card.toml @@ -8,17 +8,15 @@ run_name = %(run.run_name)s seed = %(run.seed)s # list of devices; options: cpu, cuda, hip (an optional device index can be added, e.g. "cuda:1") device = %(run.device)s -# SIMD width used by the 'cpu' devices (ignored by cuda/hip) +# SIMD width of the matrix elements (madmatrix) on the 'cpu' devices (ignored by cuda/hip) # options: auto (widest width supported by the host), scalar, simd_128, simd_256, simd_512, avx512y cpu_mode = %(run.cpu_mode)s +# SIMD width of the phase-space sampling (madspace) on the 'cpu' devices (experimental) +# options: auto (widest width supported by the host), scalar, simd_128 (ARM only), simd_256, simd_512 +madspace_cpu_mode = %(run.madspace_cpu_mode)s # matrix-element floating-point precision # options: all64, all32, color32 (colour FP32, rest FP64), denom64 (momenta + propagator denominator FP64, rest FP32) precision = %(run.precision)s -# options: -# -1 to choose automatically -# on x86: 1, 4, 8 -# on Apple silicon: 1, 2 -simd_vector_size = %(run.simd_vector_size)s # pool sizes: -1 sets count automatically based on number of CPUs cpu_thread_pool_size = %(run.cpu_thread_pool_size)s gpu_thread_pool_size = %(run.gpu_thread_pool_size)s diff --git a/madgraph/various/RunCardLO_to_MG7_mapping.md b/madgraph/various/RunCardLO_to_MG7_mapping.md index cec0812a01..e8e20f5cc8 100644 --- a/madgraph/various/RunCardLO_to_MG7_mapping.md +++ b/madgraph/various/RunCardLO_to_MG7_mapping.md @@ -125,7 +125,7 @@ Cuts that are **not representable** in the current MG7 cut engine ([x] unless no ## 9. MG7-only (no LO source, keep default) — [mg7-only] -`run.device`, `run.cpu_mode`, `run.simd_vector_size`, `run.{cpu,gpu,combine}_thread_pool_size`, +`run.device`, `run.cpu_mode`, `run.madspace_cpu_mode`, `run.{cpu,gpu,combine}_thread_pool_size`, `run.output_format`, `run.verbosity`, `run.dummy_matrix_element`, `gridpack.{include_source,include_madspace,include_madspace_source}`, `generation.{cpu,gpu}_batch_size`, diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py index 999d938e24..dd4ace7270 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -6675,14 +6675,16 @@ def default_setup(self): comment="list of devices; each entry is cpu, cuda or hip, optionally followed by a device index (e.g. \"cuda:1\")") self.add_toml_param('run', 'cpu_mode', "auto", gridpack=True, allowed=['auto', 'scalar', 'simd_128', 'simd_256', 'simd_512', 'avx512y'], - comment="SIMD width used by the 'cpu' devices; 'auto' detects the widest one supported by the host") + comment="SIMD width of the matrix elements (madmatrix) on the 'cpu' devices; 'auto' detects the widest one supported by the host") self.add_toml_param('run', 'precision', "color32", gridpack=True, allowed=['all64', 'all32', 'color32', 'denom64'], comment="matrix-element floating-point precision: all64 (FP64), all32 (FP32), " "color32 (colour FP32, rest FP64), " "denom64 (momenta + propagator denominator FP64, rest FP32)") - self.add_toml_param('run', 'simd_vector_size', -1, - comment="-1 chooses automatically; on x86: 1, 4, 8; on Apple silicon: 1, 2") + self.add_toml_param('run', 'madspace_cpu_mode', "scalar", + allowed=['auto', 'scalar', 'simd_128', 'simd_256', 'simd_512'], + comment="SIMD width of the phase-space sampling (madspace) on the 'cpu' devices " + "(experimental); simd_128 is ARM only, 'auto' picks the widest one supported by the host") self.add_toml_param('run', 'cpu_thread_pool_size', -1, gridpack=True, comment="-1 sets count automatically based on number of CPUs") self.add_toml_param('run', 'gpu_thread_pool_size', 1, gridpack=True) @@ -6922,7 +6924,7 @@ def __getitem__(self, name): get = __getitem__ def __setitem__(self, name, value, *args, **opts): - """Refuse an unsupported cpu_mode instead of silently falling back. + """Refuse an unsupported (madspace_)cpu_mode instead of silently falling back. The generic ConfigFile machinery reacts to a value outside an 'allowed' list by logging a warning and keeping the previous value. For cpu_mode @@ -6931,12 +6933,15 @@ def __setitem__(self, name, value, *args, **opts): run cards written before the backend renaming, which still carry a removed value such as 'cpu_128b'. """ - if isinstance(name, str) and name.strip().lower() in ('cpu_mode', 'run.cpu_mode'): - allowed = self.allowed_value.get('run.cpu_mode', []) + key = name.strip().lower() if isinstance(name, str) else None + if key is not None and key.startswith('run.'): + key = key[4:] + if key in ('cpu_mode', 'madspace_cpu_mode'): + allowed = self.allowed_value.get('run.' + key, []) if allowed and str(value).strip().lower() not in [str(v).lower() for v in allowed]: raise InvalidRunCard( - "Invalid cpu_mode='%s': supported values are [ '%s' ]" - % (str(value).strip(), "', '".join(str(v) for v in allowed))) + "Invalid %s='%s': supported values are [ '%s' ]" + % (key, str(value).strip(), "', '".join(str(v) for v in allowed))) return super(RunCardMG7, self).__setitem__(name, value, *args, **opts) # ------------------------------------------------------------------ diff --git a/madspace/CMakeLists.txt b/madspace/CMakeLists.txt index 316453d193..4b8e3825d0 100644 --- a/madspace/CMakeLists.txt +++ b/madspace/CMakeLists.txt @@ -436,12 +436,12 @@ if(NOT ENABLE_SIMD) set(CPU_TARGETS madspace_cpu) set(SIMD_FLAGS "") set(SIMD_DEFINES "") -elseif(APPLE) - set(CPU_TARGETS madspace_cpu madspace_cpu_neon) +elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "^(arm64|aarch64|ARM64)$") + set(CPU_TARGETS madspace_cpu madspace_cpu_simd_128) set(SIMD_FLAGS "" "") set(SIMD_DEFINES "" "USE_SIMD_NEON=1") else() - set(CPU_TARGETS madspace_cpu madspace_cpu_avx2 madspace_cpu_avx512) + set(CPU_TARGETS madspace_cpu madspace_cpu_simd_256 madspace_cpu_simd_512) set(SIMD_FLAGS "" "-mavx2 -mfma" "-mavx512f") set(SIMD_DEFINES "" "USE_SIMD_AVX2=1" "USE_SIMD_AVX512=1") endif() diff --git a/madspace/include/madspace/driver/backend.hpp b/madspace/include/madspace/driver/backend.hpp index 56688b561b..6f210cf5a6 100644 --- a/madspace/include/madspace/driver/backend.hpp +++ b/madspace/include/madspace/driver/backend.hpp @@ -46,7 +46,9 @@ DevicePtr cuda_device(std::size_t index = 0); DevicePtr hip_device(std::size_t index = 0); /// Directory to load the generated backend shared libraries from. void set_lib_path(const std::string& lib_path); -/// Width, in elements, of the SIMD vector used by the CPU backend's kernels. -void set_simd_vector_size(int vector_size); +/// SIMD mode of the CPU backend: "scalar", "simd_128" (ARM), "simd_256" (AVX2), +/// "simd_512" (AVX512) or "auto" (widest supported by the host). Must be set +/// before the CPU backend is first used; defaults to $MADSPACE_SIMD_MODE or "scalar". +void set_simd_mode(const std::string& simd_mode); } // namespace madspace diff --git a/madspace/src/driver/backend.cpp b/madspace/src/driver/backend.cpp index 327c2fd3c1..89f73cc5ab 100644 --- a/madspace/src/driver/backend.cpp +++ b/madspace/src/driver/backend.cpp @@ -11,7 +11,7 @@ namespace { struct LoadedBackend { inline static std::string lib_path = ""; - inline static int vector_size = -1; + inline static std::string simd_mode = ""; inline static std::unordered_map device_backends; LoadedBackend(const std::string& file) { @@ -73,52 +73,48 @@ struct LoadedBackend { const LoadedBackend& cpu_backend() { static LoadedBackend backend = [&] { - std::vector supported_vector_sizes{1}; + // ordered by increasing vector width, so that "auto" picks the last one + std::vector supported_modes{"scalar"}; #ifdef SIMD_AVAILABLE -#ifdef __APPLE__ - supported_vector_sizes.push_back(2); -#else // __APPLE__ +#if defined(__aarch64__) || defined(_M_ARM64) + supported_modes.push_back("simd_128"); +#elif defined(__x86_64__) || defined(_M_X64) if (__builtin_cpu_supports("avx2") && __builtin_cpu_supports("fma")) { - supported_vector_sizes.push_back(4); + supported_modes.push_back("simd_256"); } if (__builtin_cpu_supports("avx512f")) { - supported_vector_sizes.push_back(8); + supported_modes.push_back("simd_512"); } -#endif // __APPLE__ +#endif #endif // SIMD_AVAILABLE - int vector_size = LoadedBackend::vector_size; - if (vector_size == -1) { - if (char* env_var = std::getenv("SIMD_VECTOR_SIZE")) { - vector_size = std::atoi(env_var); - } else { - vector_size = 0; - } + std::string mode = LoadedBackend::simd_mode; + if (mode.empty()) { + char* env_var = std::getenv("MADSPACE_SIMD_MODE"); + mode = env_var ? env_var : "scalar"; } - if (vector_size <= 0) { -#ifdef __APPLE__ - vector_size = 1; -#else - // vector_size = supported_vector_sizes.back(); - vector_size = 1; -#endif - } else if (std::find( - supported_vector_sizes.begin(), - supported_vector_sizes.end(), - vector_size - ) == supported_vector_sizes.end()) { - throw std::runtime_error("unsupported SIMD vector size"); + if (mode == "auto") { + mode = supported_modes.back(); + } else if (std::find(supported_modes.begin(), supported_modes.end(), mode) == + supported_modes.end()) { + std::string supported; + for (auto& supported_mode : supported_modes) { + supported += std::format(" '{}'", supported_mode); + } + throw std::runtime_error( + std::format( + "unsupported madspace SIMD mode '{}', supported on this host: " + "'auto'{}", + mode, + supported + ) + ); } - switch (vector_size) { - case 2: - return LoadedBackend("libmadspace_cpu_neon"); - case 4: - return LoadedBackend("libmadspace_cpu_avx2"); - case 8: - return LoadedBackend("libmadspace_cpu_avx512"); - default: + if (mode == "scalar") { return LoadedBackend("libmadspace_cpu"); + } else { + return LoadedBackend(std::format("libmadspace_cpu_{}", mode)); } }(); return backend; @@ -181,6 +177,6 @@ void madspace::set_lib_path(const std::string& lib_path) { LoadedBackend::lib_path = lib_path; } -void madspace::set_simd_vector_size(int vector_size) { - LoadedBackend::vector_size = vector_size; +void madspace::set_simd_mode(const std::string& simd_mode) { + LoadedBackend::simd_mode = simd_mode; } diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index d70b5ba09b..199b45e431 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -1293,7 +1293,9 @@ PYBIND11_MODULE(_madspace_py, m) { &PhaseSpaceMapping::channel_count, pydoc::doc("PhaseSpaceMapping::channel_count") ) - .def("empty", &PhaseSpaceMapping::empty, pydoc::doc("PhaseSpaceMapping::empty")); + .def( + "empty", &PhaseSpaceMapping::empty, pydoc::doc("PhaseSpaceMapping::empty") + ); py::classh( m, "MultiChannelFunction", pydoc::doc("MultiChannelFunction") @@ -3691,10 +3693,10 @@ PYBIND11_MODULE(_madspace_py, m) { "set_lib_path", &set_lib_path, py::arg("lib_path"), pydoc::doc("set_lib_path") ); m.def( - "set_simd_vector_size", - &set_simd_vector_size, - py::arg("vector_size"), - pydoc::doc("set_simd_vector_size") + "set_simd_mode", + &set_simd_mode, + py::arg("simd_mode"), + pydoc::doc("set_simd_mode") ); auto abort_check_function = [] { diff --git a/tests/input_files/mg7_run_card.toml b/tests/input_files/mg7_run_card.toml index 9e4290477c..6fbbcc92d9 100644 --- a/tests/input_files/mg7_run_card.toml +++ b/tests/input_files/mg7_run_card.toml @@ -2,14 +2,12 @@ run_name = "run" # list of devices; options: cpu, cuda, hip (an optional device index can be added, e.g. "cuda:1") device = ["cpu"] -# SIMD width used by the 'cpu' devices (ignored by cuda/hip) +# SIMD width of the matrix elements (madmatrix) on the 'cpu' devices (ignored by cuda/hip) # options: auto (widest width supported by the host), scalar, simd_128, simd_256, simd_512, avx512y cpu_mode = "scalar" -# options: -# -1 to choose automatically -# on x86: 1, 4, 8 -# on Apple silicon: 1, 2 -simd_vector_size = -1 +# SIMD width of the phase-space sampling (madspace) on the 'cpu' devices (experimental) +# options: auto (widest width supported by the host), scalar, simd_128 (ARM only), simd_256, simd_512 +madspace_cpu_mode = "scalar" # pool sizes: -1 sets count automatically based on number of CPUs cpu_thread_pool_size = -1 gpu_thread_pool_size = 1 From a36ef821676c67dc24a10cbe3dfd7ba78a074511 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 7 Oct 2026 14:29:05 +0200 Subject: [PATCH 04/23] add avx512y mode, similar to madmatrix --- .../iolibs/template_files/mg7/run_card.toml | 2 +- madgraph/various/banner.py | 4 +- madspace/CMakeLists.txt | 20 +- madspace/include/madspace/driver/backend.hpp | 5 +- madspace/src/cpu/simd.hpp | 4 + madspace/src/cpu/simd_x86_256y.hpp | 211 ++++++++++++++++++ madspace/src/driver/backend.cpp | 15 +- tests/input_files/mg7_run_card.toml | 2 +- 8 files changed, 252 insertions(+), 11 deletions(-) create mode 100644 madspace/src/cpu/simd_x86_256y.hpp diff --git a/madgraph/iolibs/template_files/mg7/run_card.toml b/madgraph/iolibs/template_files/mg7/run_card.toml index 302fff0cd3..92099166ab 100644 --- a/madgraph/iolibs/template_files/mg7/run_card.toml +++ b/madgraph/iolibs/template_files/mg7/run_card.toml @@ -12,7 +12,7 @@ device = %(run.device)s # options: auto (widest width supported by the host), scalar, simd_128, simd_256, simd_512, avx512y cpu_mode = %(run.cpu_mode)s # SIMD width of the phase-space sampling (madspace) on the 'cpu' devices (experimental) -# options: auto (widest width supported by the host), scalar, simd_128 (ARM only), simd_256, simd_512 +# options: auto (same choice as for cpu_mode), scalar, simd_128 (ARM only), simd_256, avx512y, simd_512 madspace_cpu_mode = %(run.madspace_cpu_mode)s # matrix-element floating-point precision # options: all64, all32, color32 (colour FP32, rest FP64), denom64 (momenta + propagator denominator FP64, rest FP32) diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py index dd4ace7270..3ee563c173 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -6682,9 +6682,9 @@ def default_setup(self): "color32 (colour FP32, rest FP64), " "denom64 (momenta + propagator denominator FP64, rest FP32)") self.add_toml_param('run', 'madspace_cpu_mode', "scalar", - allowed=['auto', 'scalar', 'simd_128', 'simd_256', 'simd_512'], + allowed=['auto', 'scalar', 'simd_128', 'simd_256', 'avx512y', 'simd_512'], comment="SIMD width of the phase-space sampling (madspace) on the 'cpu' devices " - "(experimental); simd_128 is ARM only, 'auto' picks the widest one supported by the host") + "(experimental); simd_128 is ARM only, 'auto' makes the same choice as for cpu_mode") self.add_toml_param('run', 'cpu_thread_pool_size', -1, gridpack=True, comment="-1 sets count automatically based on number of CPUs") self.add_toml_param('run', 'gpu_thread_pool_size', 1, gridpack=True) diff --git a/madspace/CMakeLists.txt b/madspace/CMakeLists.txt index 4b8e3825d0..d8cb179fca 100644 --- a/madspace/CMakeLists.txt +++ b/madspace/CMakeLists.txt @@ -441,9 +441,21 @@ elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "^(arm64|aarch64|ARM64)$") set(SIMD_FLAGS "" "") set(SIMD_DEFINES "" "USE_SIMD_NEON=1") else() - set(CPU_TARGETS madspace_cpu madspace_cpu_simd_256 madspace_cpu_simd_512) - set(SIMD_FLAGS "" "-mavx2 -mfma" "-mavx512f") - set(SIMD_DEFINES "" "USE_SIMD_AVX2=1" "USE_SIMD_AVX512=1") + set( + CPU_TARGETS + madspace_cpu madspace_cpu_simd_256 madspace_cpu_avx512y madspace_cpu_simd_512 + ) + set( + SIMD_FLAGS + "" + "-mavx2 -mfma" + "-mavx512f -mavx512vl -mavx512dq -mprefer-vector-width=256" + "-mavx512f" + ) + set( + SIMD_DEFINES + "" "USE_SIMD_AVX2=1" "USE_SIMD_AVX512Y=1" "USE_SIMD_AVX512=1" + ) endif() foreach( @@ -464,6 +476,8 @@ foreach( src/cpu/simd.hpp src/cpu/simd_arm.hpp src/cpu/simd_x86_256.hpp + src/cpu/simd_x86_256y.hpp + src/cpu/simd_x86_512.hpp ${KERNEL_SRC} ) target_include_directories(${cpu_target} PUBLIC include) diff --git a/madspace/include/madspace/driver/backend.hpp b/madspace/include/madspace/driver/backend.hpp index 6f210cf5a6..8d65785326 100644 --- a/madspace/include/madspace/driver/backend.hpp +++ b/madspace/include/madspace/driver/backend.hpp @@ -47,8 +47,9 @@ DevicePtr hip_device(std::size_t index = 0); /// Directory to load the generated backend shared libraries from. void set_lib_path(const std::string& lib_path); /// SIMD mode of the CPU backend: "scalar", "simd_128" (ARM), "simd_256" (AVX2), -/// "simd_512" (AVX512) or "auto" (widest supported by the host). Must be set -/// before the CPU backend is first used; defaults to $MADSPACE_SIMD_MODE or "scalar". +/// "avx512y" (AVX512 with 256-bit vectors), "simd_512" (AVX512) or "auto" (best +/// mode for the host, chosen as in madmatrix). Must be set before the CPU backend +/// is first used; defaults to $MADSPACE_SIMD_MODE or "scalar". void set_simd_mode(const std::string& simd_mode); } // namespace madspace diff --git a/madspace/src/cpu/simd.hpp b/madspace/src/cpu/simd.hpp index b6561645d1..7f10c491db 100644 --- a/madspace/src/cpu/simd.hpp +++ b/madspace/src/cpu/simd.hpp @@ -10,6 +10,10 @@ #include "simd_x86_256.hpp" #endif +#ifdef USE_SIMD_AVX512Y +#include "simd_x86_256y.hpp" +#endif + #ifdef USE_SIMD_AVX512 #include "simd_x86_512.hpp" #endif diff --git a/madspace/src/cpu/simd_x86_256y.hpp b/madspace/src/cpu/simd_x86_256y.hpp new file mode 100644 index 0000000000..a038212d7d --- /dev/null +++ b/madspace/src/cpu/simd_x86_256y.hpp @@ -0,0 +1,211 @@ +// AVX512 instructions (F, VL, DQ) with 256-bit vectors, like the madmatrix avx512y mode +#include +#include +#include + +constexpr int simd_vec_size = 4; + +struct FVec { + FVec() = default; + FVec(__m256d _v) : v(_v) {}; + FVec(double _v) : v(_mm256_set1_pd(_v)) {}; + explicit FVec(__m256i _v) : v(_mm256_cvtepi64_pd(_v)) {}; + operator __m256d() { return v; } + FVec operator+=(FVec _v) { + v = _mm256_add_pd(v, _v); + return v; + } + __m256d v; +}; + +struct IVec { + IVec() = default; + IVec(__m256i _v) : v(_v) {}; + IVec(int _v) : v(_mm256_set1_epi64x(_v)) {}; + explicit IVec(__m256d _v) : v(_mm256_cvttpd_epi64(_v)) {}; + operator __m256i() { return v; } + IVec operator+=(IVec _v) { + v = _mm256_add_epi64(v, _v); + return v; + } + __m256i v; +}; + +struct BVec { + BVec() = default; + BVec(bool _v) : v(_v ? 0x0F : 0) {}; + BVec(__mmask8 _v) : v(_v) {}; + operator __mmask8() { return v; } + __mmask8 v; +}; + +inline __m256i stride_seq(std::size_t stride) { + return _mm256_mullo_epi64( + _mm256_set1_epi64x(stride), _mm256_set_epi64x(3, 2, 1, 0) + ); +} + +inline __m256i +mem_indices(std::size_t batch_stride, std::size_t index_stride, IVec indices) { + return _mm256_add_epi64( + _mm256_mullo_epi64(indices, _mm256_set1_epi64x(index_stride)), + stride_seq(batch_stride) + ); +} + +inline FVec vgather( + double* base_ptr, std::size_t batch_stride, std::size_t index_stride, IVec indices +) { + return _mm256_i64gather_pd( + base_ptr, mem_indices(batch_stride, index_stride, indices), 8 + ); +} + +inline IVec vgather( + int* base_ptr, std::size_t batch_stride, std::size_t index_stride, IVec indices +) { + return _mm256_cvtepi32_epi64(_mm256_i64gather_epi32( + base_ptr, mem_indices(batch_stride, index_stride, indices), 4 + )); +} + +inline FVec vload(double* base_ptr, std::size_t stride) { + return _mm256_i64gather_pd(base_ptr, stride_seq(stride), 8); +} + +inline IVec vload(int* base_ptr, std::size_t stride) { + return _mm256_cvtepi32_epi64( + _mm256_i64gather_epi32(base_ptr, stride_seq(stride), 4) + ); +} + +inline void vscatter( + double* base_ptr, + std::size_t batch_stride, + std::size_t index_stride, + IVec indices, + FVec values +) { + _mm256_i64scatter_pd( + base_ptr, mem_indices(batch_stride, index_stride, indices), values, 8 + ); +} + +inline void vscatter( + int* base_ptr, + std::size_t batch_stride, + std::size_t index_stride, + IVec indices, + IVec values +) { + _mm256_i64scatter_epi32( + base_ptr, + mem_indices(batch_stride, index_stride, indices), + _mm256_cvtepi64_epi32(values), + 4 + ); +} + +inline void vstore(double* base_ptr, std::size_t stride, FVec values) { + _mm256_i64scatter_pd(base_ptr, stride_seq(stride), values, 8); +} + +inline void vstore(int* base_ptr, std::size_t stride, IVec values) { + _mm256_i64scatter_epi32( + base_ptr, stride_seq(stride), _mm256_cvtepi64_epi32(values), 4 + ); +} + +inline FVec where(BVec arg1, FVec arg2, FVec arg3) { + return _mm256_mask_blend_pd(arg1, arg3, arg2); +} +inline IVec where(BVec arg1, IVec arg2, IVec arg3) { + return _mm256_mask_blend_epi64(arg1, arg3, arg2); +} +inline std::size_t single_index(IVec arg) { + return _mm_cvtsi128_si64(_mm256_castsi256_si128(arg)); +} +inline FVec min(FVec arg1, FVec arg2) { return _mm256_min_pd(arg1, arg2); } +inline FVec max(FVec arg1, FVec arg2) { return _mm256_max_pd(arg1, arg2); } + +inline BVec operator==(FVec arg1, FVec arg2) { + return _mm256_cmp_pd_mask(arg1, arg2, _CMP_EQ_OQ); +} +inline BVec operator!=(FVec arg1, FVec arg2) { + return _mm256_cmp_pd_mask(arg1, arg2, _CMP_NEQ_UQ); +} +inline BVec operator>(FVec arg1, FVec arg2) { + return _mm256_cmp_pd_mask(arg1, arg2, _CMP_GT_OQ); +} +inline BVec operator<(FVec arg1, FVec arg2) { + return _mm256_cmp_pd_mask(arg1, arg2, _CMP_LT_OQ); +} +inline BVec operator>=(FVec arg1, FVec arg2) { + return _mm256_cmp_pd_mask(arg1, arg2, _CMP_GE_OQ); +} +inline BVec operator<=(FVec arg1, FVec arg2) { + return _mm256_cmp_pd_mask(arg1, arg2, _CMP_LE_OQ); +} + +inline BVec operator&(BVec arg1, BVec arg2) { + return static_cast<__mmask8>(arg1.v & arg2.v & 0x0F); +} +inline BVec operator|(BVec arg1, BVec arg2) { + return static_cast<__mmask8>((arg1.v | arg2.v) & 0x0F); +} +inline BVec operator!(BVec arg1) { return static_cast<__mmask8>(~arg1.v & 0x0F); } + +inline BVec operator==(IVec arg1, IVec arg2) { + return _mm256_cmpeq_epi64_mask(arg1, arg2); +} +inline BVec operator!=(IVec arg1, IVec arg2) { + return _mm256_cmpneq_epi64_mask(arg1, arg2); +} +inline BVec operator>(IVec arg1, IVec arg2) { + return _mm256_cmpgt_epi64_mask(arg1, arg2); +} +inline BVec operator>=(IVec arg1, IVec arg2) { + return _mm256_cmpge_epi64_mask(arg1, arg2); +} +inline BVec operator<(IVec arg1, IVec arg2) { + return _mm256_cmplt_epi64_mask(arg1, arg2); +} +inline BVec operator<=(IVec arg1, IVec arg2) { + return _mm256_cmple_epi64_mask(arg1, arg2); +} + +inline FVec operator-(FVec arg1) { return _mm256_sub_pd(_mm256_set1_pd(0.), arg1); } +inline FVec operator+(FVec arg1, FVec arg2) { return _mm256_add_pd(arg1, arg2); } +inline FVec operator-(FVec arg1, FVec arg2) { return _mm256_sub_pd(arg1, arg2); } +inline FVec operator*(FVec arg1, FVec arg2) { return _mm256_mul_pd(arg1, arg2); } +inline FVec operator/(FVec arg1, FVec arg2) { return _mm256_div_pd(arg1, arg2); } +inline IVec operator-(IVec arg1) { + return _mm256_sub_epi64(_mm256_setzero_si256(), arg1); +} +inline IVec operator+(IVec arg1, IVec arg2) { return _mm256_add_epi64(arg1, arg2); } +inline IVec operator-(IVec arg1, IVec arg2) { return _mm256_sub_epi64(arg1, arg2); } + +inline BVec isnan(FVec arg) { return arg != arg; } + +inline FVec sqrt(FVec arg1) { return Sleef_sqrtd4_u05avx2(arg1); } +inline FVec sin(FVec arg1) { return Sleef_sind4_u10avx2(arg1); } +inline FVec cos(FVec arg1) { return Sleef_cosd4_u10avx2(arg1); } +inline FVec sinh(FVec arg1) { return Sleef_sinhd4_u10avx2(arg1); } +inline FVec asinh(FVec arg1) { return Sleef_asinhd4_u10avx2(arg1); } +inline FVec cosh(FVec arg1) { return Sleef_coshd4_u10avx2(arg1); } +inline FVec tanh(FVec arg1) { return Sleef_tanhd4_u10avx2(arg1); } +inline FVec atan2(FVec arg1, FVec arg2) { return Sleef_atan2d4_u10avx2(arg1, arg2); } +inline FVec pow(FVec arg1, FVec arg2) { return Sleef_powd4_u10avx2(arg1, arg2); } +inline FVec fabs(FVec arg1) { return Sleef_fabsd4_avx2(arg1); } +inline FVec log(FVec arg1) { return Sleef_logd4_u10avx2(arg1); } +inline FVec tan(FVec arg1) { return Sleef_tand4_u10avx2(arg1); } +inline FVec atan(FVec arg1) { return Sleef_atand4_u10avx2(arg1); } +inline FVec acos(FVec arg1) { return Sleef_acosd4_u10avx2(arg1); } +inline FVec atanh(FVec arg1) { return Sleef_atanhd4_u10avx2(arg1); } +inline FVec exp(FVec arg1) { return Sleef_expd4_u10avx2(arg1); } +inline FVec log1p(FVec arg1) { return Sleef_log1pd4_u10avx2(arg1); } +inline FVec expm1(FVec arg1) { return Sleef_expm1d4_u10avx2(arg1); } +inline FVec erf(FVec arg1) { return Sleef_erfd4_u10avx2(arg1); } +inline FVec fma(FVec arg1, FVec arg2, FVec arg3) { + return _mm256_fmadd_pd(arg1, arg2, arg3); +} diff --git a/madspace/src/driver/backend.cpp b/madspace/src/driver/backend.cpp index 89f73cc5ab..ebae4c80b4 100644 --- a/madspace/src/driver/backend.cpp +++ b/madspace/src/driver/backend.cpp @@ -73,14 +73,25 @@ struct LoadedBackend { const LoadedBackend& cpu_backend() { static LoadedBackend backend = [&] { - // ordered by increasing vector width, so that "auto" picks the last one std::vector supported_modes{"scalar"}; + std::string auto_mode = "scalar"; #ifdef SIMD_AVAILABLE #if defined(__aarch64__) || defined(_M_ARM64) supported_modes.push_back("simd_128"); + auto_mode = "simd_128"; #elif defined(__x86_64__) || defined(_M_X64) if (__builtin_cpu_supports("avx2") && __builtin_cpu_supports("fma")) { supported_modes.push_back("simd_256"); + auto_mode = "simd_256"; + } + if (__builtin_cpu_supports("avx512f") && __builtin_cpu_supports("avx512vl") && + __builtin_cpu_supports("avx512dq")) { + supported_modes.push_back("avx512y"); + // mirrors madmatrix's auto mode (avx512y except with clang, never + // simd_512); TODO: revisit once madspace has its own x86 benchmarks +#ifndef __clang__ + auto_mode = "avx512y"; +#endif } if (__builtin_cpu_supports("avx512f")) { supported_modes.push_back("simd_512"); @@ -94,7 +105,7 @@ const LoadedBackend& cpu_backend() { mode = env_var ? env_var : "scalar"; } if (mode == "auto") { - mode = supported_modes.back(); + mode = auto_mode; } else if (std::find(supported_modes.begin(), supported_modes.end(), mode) == supported_modes.end()) { std::string supported; diff --git a/tests/input_files/mg7_run_card.toml b/tests/input_files/mg7_run_card.toml index 6fbbcc92d9..935f8b8980 100644 --- a/tests/input_files/mg7_run_card.toml +++ b/tests/input_files/mg7_run_card.toml @@ -6,7 +6,7 @@ device = ["cpu"] # options: auto (widest width supported by the host), scalar, simd_128, simd_256, simd_512, avx512y cpu_mode = "scalar" # SIMD width of the phase-space sampling (madspace) on the 'cpu' devices (experimental) -# options: auto (widest width supported by the host), scalar, simd_128 (ARM only), simd_256, simd_512 +# options: auto (same choice as for cpu_mode), scalar, simd_128 (ARM only), simd_256, avx512y, simd_512 madspace_cpu_mode = "scalar" # pool sizes: -1 sets count automatically based on number of CPUs cpu_thread_pool_size = -1 From b21a44860c9db3513fc4d4bc690f6f4d99eefb5d Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 7 Oct 2026 15:08:13 +0200 Subject: [PATCH 05/23] fix avx512y build with GCC: enable AVX2 and FMA explicitly --- madspace/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/madspace/CMakeLists.txt b/madspace/CMakeLists.txt index d8cb179fca..69e7081814 100644 --- a/madspace/CMakeLists.txt +++ b/madspace/CMakeLists.txt @@ -449,7 +449,7 @@ else() SIMD_FLAGS "" "-mavx2 -mfma" - "-mavx512f -mavx512vl -mavx512dq -mprefer-vector-width=256" + "-mavx2 -mfma -mavx512f -mavx512vl -mavx512dq -mprefer-vector-width=256" "-mavx512f" ) set( From f95a83c14a8d5ce53d1b0ba633b9dc7c94f789f5 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 7 Oct 2026 15:31:00 +0200 Subject: [PATCH 06/23] test madspace and mg7 runs for every supported SIMD mode - 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 --- .github/actions/build_madspace/action.yml | 5 +- .github/actions/install_madspace/action.yml | 4 +- .github/workflows/acceptancetest_mg7.yml | 14 ++ .github/workflows/madspace.yml | 6 +- docs/source/madspace/cpp-api.rst | 1 + docs/source/madspace/python-api.rst | 1 + madspace/include/madspace/driver/backend.hpp | 3 + madspace/src/driver/backend.cpp | 52 +++--- madspace/src/python/madspace.cpp | 5 + madspace/tests/conftest.py | 5 + madspace/tests/run_simd_modes.py | 53 ++++++ .../test_mg7_simd_consistency.py | 165 ++++++++++++++++++ 12 files changed, 290 insertions(+), 24 deletions(-) create mode 100644 madspace/tests/conftest.py create mode 100644 madspace/tests/run_simd_modes.py create mode 100644 tests/acceptance_tests/test_mg7_simd_consistency.py diff --git a/.github/actions/build_madspace/action.yml b/.github/actions/build_madspace/action.yml index 046d6383ee..25ed6a0fba 100644 --- a/.github/actions/build_madspace/action.yml +++ b/.github/actions/build_madspace/action.yml @@ -26,7 +26,9 @@ runs: - name: madspace cache key run: | echo "MS_PY=$(python3 -c 'import sys;print(f"{sys.version_info.major}{sys.version_info.minor}")')" >> $GITHUB_ENV - echo "MS_VARIANT=${{ inputs.enable_docs == 'true' && '-docs' || '' }}" >> $GITHUB_ENV + # -simd: the install includes the SIMD CPU backends; a separate suffix so + # that a cached install from before SIMD was enabled is never restored + echo "MS_VARIANT=${{ inputs.enable_docs == 'true' && '-docs' || '' }}-simd" >> $GITHUB_ENV shell: bash # Cache the installed madspace/install prefix directly (not just a wheel), so @@ -68,6 +70,7 @@ runs: python3 -m pip install --upgrade pip packaging python3 -m pip install ./madspace \ -Ccmake.define.ENABLE_OPENBLAS=ON \ + -Ccmake.define.ENABLE_SIMD=ON \ ${{ inputs.enable_docs == 'true' && '-Ccmake.define.ENABLE_DOCS=ON' || '' }} \ --target="$GITHUB_WORKSPACE/madspace/install" shell: bash diff --git a/.github/actions/install_madspace/action.yml b/.github/actions/install_madspace/action.yml index 041bb5b866..9a73f16103 100644 --- a/.github/actions/install_madspace/action.yml +++ b/.github/actions/install_madspace/action.yml @@ -22,7 +22,9 @@ runs: - name: madspace cache key run: | echo "MS_PY=$(python3 -c 'import sys;print(f"{sys.version_info.major}{sys.version_info.minor}")')" >> $GITHUB_ENV - echo "MS_VARIANT=${{ inputs.enable_docs == 'true' && '-docs' || '' }}" >> $GITHUB_ENV + # -simd: the install includes the SIMD CPU backends; a separate suffix so + # that a cached install from before SIMD was enabled is never restored + echo "MS_VARIANT=${{ inputs.enable_docs == 'true' && '-docs' || '' }}-simd" >> $GITHUB_ENV shell: bash # madspace/install is the local prefix the mg7 runtime expects (see diff --git a/.github/workflows/acceptancetest_mg7.yml b/.github/workflows/acceptancetest_mg7.yml index ad9ddd0896..0b461938c4 100644 --- a/.github/workflows/acceptancetest_mg7.yml +++ b/.github/workflows/acceptancetest_mg7.yml @@ -197,6 +197,7 @@ jobs: # acceptancetest_mg7_vegas_reproducibility # acceptancetest_mg7_gridpack_reproducibility # acceptancetest_mg7_gridpack_reproducibility_vegas + # acceptancetest_mg7_simd_consistency # Each keeps its own steps, run even when an earlier one failed. needs: build_madspace runs-on: ubuntu-24.04 @@ -287,6 +288,19 @@ jobs: export PATH="$HOME/.cache/HEPtools/bin:$PATH" ./tests/test_manager.py test_gridpack_reproducibility_vegas_mg7 -pA -t0 -l INFO + # ---- acceptancetest_mg7_simd_consistency ---- + # mg7 seeded generation (g g > t t~ g), repeated for every madspace SIMD + # mode (madspace_cpu_mode) supported by the runner: the VEGAS cross + # sections must agree with the scalar one up to rounding, the madnis ones + # within their statistical uncertainties. Self-skips if the runner supports + # no SIMD mode or the madspace + LHAPDF(NNPDF23) stack is unavailable. + - name: "acceptancetest_mg7_simd_consistency" + if: ${{ !cancelled() }} + run: | + cd $GITHUB_WORKSPACE + export PATH="$HOME/.cache/HEPtools/bin:$PATH" + ./tests/test_manager.py test_vegas_simd_consistency_mg7 test_madnis_simd_consistency_mg7 -pA -t0 -l INFO + acceptancetest_mg7_bundle_3: # former jobs, run one after the other: acceptancetest_mg7_e_p_collision diff --git a/.github/workflows/madspace.yml b/.github/workflows/madspace.yml index 9044c03417..fff44a9eee 100644 --- a/.github/workflows/madspace.yml +++ b/.github/workflows/madspace.yml @@ -62,7 +62,7 @@ jobs: MACOSX_DEPLOYMENT_TARGET: 14.5 CIBW_BEFORE_ALL: brew install doxygen CIBW_ENVIRONMENT: > - SKBUILD_CMAKE_ARGS="-DENABLE_DOCS=ON" + SKBUILD_CMAKE_ARGS="-DENABLE_SIMD=ON;-DENABLE_DOCS=ON" CIBW_REPAIR_WHEEL_COMMAND: > delocate-wheel --ignore-missing-dependencies @@ -87,7 +87,9 @@ jobs: scipy --extra-index-url https://download.pytorch.org/whl/cpu torch - CIBW_TEST_COMMAND: pytest {project}/madspace/tests + # runs the unit tests once per SIMD mode supported by the runner + CIBW_TEST_COMMAND: > + python {project}/madspace/tests/run_simd_modes.py {project}/madspace/tests - uses: actions/upload-artifact@v4 with: diff --git a/docs/source/madspace/cpp-api.rst b/docs/source/madspace/cpp-api.rst index 3c91b31b48..88edd68768 100644 --- a/docs/source/madspace/cpp-api.rst +++ b/docs/source/madspace/cpp-api.rst @@ -179,6 +179,7 @@ The high-level driver plus the tensor and device layer it runs on: integration, cpp/hip_device cpp/select_combine_channel_index cpp/set_simd_mode + cpp/supported_simd_modes cpp/set_lib_path Utilities diff --git a/docs/source/madspace/python-api.rst b/docs/source/madspace/python-api.rst index c387e3121a..cd1d933c96 100644 --- a/docs/source/madspace/python-api.rst +++ b/docs/source/madspace/python-api.rst @@ -165,6 +165,7 @@ The high-level driver plus the tensor and device layer it runs on: integration, python/hip_device python/select_combine_channel_index python/set_simd_mode + python/supported_simd_modes python/set_lib_path Utilities diff --git a/madspace/include/madspace/driver/backend.hpp b/madspace/include/madspace/driver/backend.hpp index 8d65785326..839fc74936 100644 --- a/madspace/include/madspace/driver/backend.hpp +++ b/madspace/include/madspace/driver/backend.hpp @@ -51,5 +51,8 @@ void set_lib_path(const std::string& lib_path); /// mode for the host, chosen as in madmatrix). Must be set before the CPU backend /// is first used; defaults to $MADSPACE_SIMD_MODE or "scalar". void set_simd_mode(const std::string& simd_mode); +/// SIMD modes supported by this build on this host (without "auto"), ordered by +/// increasing vector width; see @ref set_simd_mode. +std::vector supported_simd_modes(); } // namespace madspace diff --git a/madspace/src/driver/backend.cpp b/madspace/src/driver/backend.cpp index ebae4c80b4..cd25ca595c 100644 --- a/madspace/src/driver/backend.cpp +++ b/madspace/src/driver/backend.cpp @@ -1,9 +1,11 @@ #include "madspace/driver/backend.hpp" +#include #include #include #include #include +#include using namespace madspace; @@ -71,33 +73,39 @@ struct LoadedBackend { ); }; -const LoadedBackend& cpu_backend() { - static LoadedBackend backend = [&] { - std::vector supported_modes{"scalar"}; - std::string auto_mode = "scalar"; +// SIMD modes supported by this build and host, and the one chosen by "auto" +std::pair, std::string> simd_mode_support() { + std::vector supported_modes{"scalar"}; + std::string auto_mode = "scalar"; #ifdef SIMD_AVAILABLE #if defined(__aarch64__) || defined(_M_ARM64) - supported_modes.push_back("simd_128"); - auto_mode = "simd_128"; + supported_modes.push_back("simd_128"); + auto_mode = "simd_128"; #elif defined(__x86_64__) || defined(_M_X64) - if (__builtin_cpu_supports("avx2") && __builtin_cpu_supports("fma")) { - supported_modes.push_back("simd_256"); - auto_mode = "simd_256"; - } - if (__builtin_cpu_supports("avx512f") && __builtin_cpu_supports("avx512vl") && - __builtin_cpu_supports("avx512dq")) { - supported_modes.push_back("avx512y"); - // mirrors madmatrix's auto mode (avx512y except with clang, never - // simd_512); TODO: revisit once madspace has its own x86 benchmarks + if (__builtin_cpu_supports("avx2") && __builtin_cpu_supports("fma")) { + supported_modes.push_back("simd_256"); + auto_mode = "simd_256"; + } + if (__builtin_cpu_supports("avx512f") && __builtin_cpu_supports("avx512vl") && + __builtin_cpu_supports("avx512dq")) { + supported_modes.push_back("avx512y"); + // mirrors madmatrix's auto mode (avx512y except with clang, never + // simd_512); TODO: revisit once madspace has its own x86 benchmarks #ifndef __clang__ - auto_mode = "avx512y"; + auto_mode = "avx512y"; #endif - } - if (__builtin_cpu_supports("avx512f")) { - supported_modes.push_back("simd_512"); - } + } + if (__builtin_cpu_supports("avx512f")) { + supported_modes.push_back("simd_512"); + } #endif #endif // SIMD_AVAILABLE + return {supported_modes, auto_mode}; +} + +const LoadedBackend& cpu_backend() { + static LoadedBackend backend = [&] { + auto [supported_modes, auto_mode] = simd_mode_support(); std::string mode = LoadedBackend::simd_mode; if (mode.empty()) { @@ -188,6 +196,10 @@ void madspace::set_lib_path(const std::string& lib_path) { LoadedBackend::lib_path = lib_path; } +std::vector madspace::supported_simd_modes() { + return simd_mode_support().first; +} + void madspace::set_simd_mode(const std::string& simd_mode) { LoadedBackend::simd_mode = simd_mode; } diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 199b45e431..9d273a72f9 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -3698,6 +3698,11 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("simd_mode"), pydoc::doc("set_simd_mode") ); + m.def( + "supported_simd_modes", + &supported_simd_modes, + pydoc::doc("supported_simd_modes") + ); auto abort_check_function = [] { if (PyErr_CheckSignals() != 0) { diff --git a/madspace/tests/conftest.py b/madspace/tests/conftest.py new file mode 100644 index 0000000000..6e02ac301d --- /dev/null +++ b/madspace/tests/conftest.py @@ -0,0 +1,5 @@ +import os + + +def pytest_report_header(config): + return f"madspace SIMD mode: {os.environ.get('MADSPACE_SIMD_MODE', 'scalar')}" diff --git a/madspace/tests/run_simd_modes.py b/madspace/tests/run_simd_modes.py new file mode 100644 index 0000000000..bbb15d0af3 --- /dev/null +++ b/madspace/tests/run_simd_modes.py @@ -0,0 +1,53 @@ +#!/usr/bin/env python3 +"""Run the madspace unit tests once for every SIMD mode supported on this machine. + +The mode is selected through the MADSPACE_SIMD_MODE environment variable, which the +CPU backend reads when it is first loaded, so every mode runs in its own pytest +process. Arguments that are not recognized here are passed on to pytest; without +any, the tests next to this script are run. + + python run_simd_modes.py # all supported modes + python run_simd_modes.py --modes scalar # only the given modes + python run_simd_modes.py -x -k vegas # extra pytest arguments +""" + +import argparse +import os +import subprocess +import sys + +import madspace as ms + + +def main() -> int: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument( + "--modes", + help="comma-separated list of SIMD modes (default: all supported ones)", + ) + args, pytest_args = parser.parse_known_args() + + supported = ms.supported_simd_modes() + modes = args.modes.split(",") if args.modes else supported + if not pytest_args: + pytest_args = [os.path.dirname(os.path.realpath(__file__))] + print(f"supported SIMD modes: {', '.join(supported)}", flush=True) + + results = {} + for mode in modes: + print(f"\n===== SIMD mode: {mode} =====", flush=True) + env = dict(os.environ, MADSPACE_SIMD_MODE=mode) + results[mode] = subprocess.call( + [sys.executable, "-m", "pytest", *pytest_args], env=env + ) + + print("\n===== summary =====") + for mode, code in results.items(): + print(f"{mode:>10}: {'passed' if code == 0 else f'FAILED (exit code {code})'}") + return 0 if all(code == 0 for code in results.values()) else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/acceptance_tests/test_mg7_simd_consistency.py b/tests/acceptance_tests/test_mg7_simd_consistency.py new file mode 100644 index 0000000000..2f2324c84b --- /dev/null +++ b/tests/acceptance_tests/test_mg7_simd_consistency.py @@ -0,0 +1,165 @@ +################################################################################ +# +# Copyright (c) 2009 The MadGraph7 Development team and Contributors +# +# This file is a part of the MadGraph7 project, an application which +# automatically generates Feynman diagrams and matrix elements for arbitrary +# high-energy processes in the Standard Model and beyond. +# +# It is subject to the MadGraph7 license which should accompany this +# distribution. +# +# For more information, visit madgraph.phys.ucl.ac.be and amcatnlo.web.cern.ch +# +################################################################################ +"""MG7 consistency under a change of the madspace SIMD mode. + +Runs the same seeded mg7 generation once per ``madspace_cpu_mode`` supported on +the machine (``madspace.supported_simd_modes()``) and compares the cross section +with the one of the scalar run: + + * ``test_vegas_simd_consistency_mg7`` -- survey (VEGAS) + generate. With the + same seed, the modes see the same random numbers and only differ by rounding, + so the cross sections have to agree to a relative tolerance far below the + statistical uncertainty (a SIMD-only mapping bug that still gave a + statistically compatible result was off by 0.3%). + * ``test_madnis_simd_consistency_mg7`` -- survey + madnis training + generate. + The rounding differences grow during the training, so here the results only + have to agree within their statistical uncertainties. + +Process: ``g g > t t~ g``, whose multi-channel phase space uses momentum +permutations. Self-skips if the mg7 runtime stack / PDF is unavailable, or if +the machine supports no SIMD mode besides scalar. + +Run locally with e.g.:: + + ./tests/test_manager.py test_.*simd_consistency_mg7 -pA -t0 -l INFO + +``MG7_SIMD_EVENTS`` sets the events per run (default 5000). +""" + +from __future__ import absolute_import +from __future__ import division + +import glob +import json +import math +import os +import shutil +import sys +import tempfile +import unittest + +import madgraph.interface.master_interface as MGCmd +from madgraph.various.banner import RunCardMG7 +from tests.acceptance_tests.test_mg7_reproducibility import ( + _mg7_datadir_or_skip, _run) + +pjoin = os.path.join + +_PROCESS = 'g g > t t~ g' +_EVENTS = int(os.environ.get('MG7_SIMD_EVENTS', 5000)) +_SEED = 424242 +# relative tolerance for the VEGAS runs, which only differ by rounding +_VEGAS_REL_TOL = 1e-4 +# tolerance for the madnis runs, in combined standard deviations +_MADNIS_SIGMAS = 4. + + +class MG7SimdConsistencyTest(unittest.TestCase): + """Cross sections of seeded mg7 runs agree between the madspace SIMD modes.""" + + def setUp(self): + self.path = tempfile.mkdtemp(prefix='mg7_simd_') + + def tearDown(self): + shutil.rmtree(self.path, ignore_errors=True) + + def _modes_or_skip(self): + import madspace + modes = madspace.supported_simd_modes() + if len(modes) < 2: + self.skipTest('no SIMD mode besides scalar supported on this machine') + return modes + + def _output_process(self, run_name, **settings): + run_dir = pjoin(self.path, run_name) + mg = MGCmd.MasterCmd() + mg.no_notification() + mg.exec_cmd('set automatic_html_opening False --no_save') + mg.exec_cmd('generate %s' % _PROCESS) + mg.exec_cmd('output mg7 %s' % run_dir) + settings = dict({ + 'run.seed': _SEED, + 'run.output_format': 'compact_npy', + 'run.postprocessing_histograms': False, + 'run.make_plots': False, + 'beam.pdf': 'NNPDF23_lo_as_0130_qed', + 'generation.events': _EVENTS, + 'systematics.enable': False, + 'postprocessing.systematics': False, + }, **settings) + toml = pjoin(run_dir, 'Cards', 'run_card.toml') + rc = RunCardMG7(toml) + for key, value in settings.items(): + rc.set(key, value, user=True) + rc.write(toml) + return run_dir + + def _cross_section(self, run_dir, datadir, mode): + """Run bin/generate_events -f with the given madspace_cpu_mode and return + the (cross section, error) from the info.json of that run.""" + shutil.rmtree(pjoin(run_dir, 'Events'), ignore_errors=True) + toml = pjoin(run_dir, 'Cards', 'run_card.toml') + rc = RunCardMG7(toml) + rc.set('run.madspace_cpu_mode', mode, user=True) + rc.write(toml) + env = dict(os.environ) + env['LHAPDF_DATA_PATH'] = datadir + _run([sys.executable, pjoin(run_dir, 'bin', 'generate_events'), '-f'], + run_dir, pjoin(run_dir, 'gen_%s.log' % mode), env, + 'mg7 generate_events (madspace_cpu_mode=%s)' % mode) + infos = sorted(glob.glob(pjoin(run_dir, 'Events', '*', 'info.json'))) + self.assertTrue(infos, 'no info.json produced for mode %s' % mode) + with open(infos[-1]) as f: + process = json.load(f)['process'] + return process['mean'], process['error'] + + def _results(self, run_dir, modes): + datadir = _mg7_datadir_or_skip(self) + results = {mode: self._cross_section(run_dir, datadir, mode) + for mode in modes} + for mode, (mean, error) in results.items(): + print('madspace_cpu_mode=%s: %.6g +- %.2g' % (mode, mean, error)) + return results + + def test_vegas_simd_consistency_mg7(self): + """VEGAS runs: every SIMD mode reproduces the scalar cross section up + to rounding.""" + modes = self._modes_or_skip() + _mg7_datadir_or_skip(self) + run_dir = self._output_process('vegas', **{'madnis.enable': False}) + results = self._results(run_dir, modes) + ref_mean, _ = results['scalar'] + for mode, (mean, _) in results.items(): + self.assertLessEqual( + abs(mean - ref_mean), _VEGAS_REL_TOL * abs(ref_mean), + 'madspace_cpu_mode=%s: cross section %g differs from the scalar ' + 'one %g' % (mode, mean, ref_mean)) + + def test_madnis_simd_consistency_mg7(self): + """madnis runs: every SIMD mode agrees with the scalar cross section + within the statistical uncertainties.""" + modes = self._modes_or_skip() + _mg7_datadir_or_skip(self) + run_dir = self._output_process( + 'madnis', **{'madnis.enable': True, 'madnis.train_batches': 200}) + results = self._results(run_dir, modes) + ref_mean, ref_error = results['scalar'] + for mode, (mean, error) in results.items(): + sigma = math.sqrt(ref_error ** 2 + error ** 2) + self.assertLessEqual( + abs(mean - ref_mean), _MADNIS_SIGMAS * sigma, + 'madspace_cpu_mode=%s: cross section %g +- %g is not compatible ' + 'with the scalar one %g +- %g' % (mode, mean, error, ref_mean, + ref_error)) From 9215112e4d78d366f01b74af214dfc57914d6cfc Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 7 Oct 2026 15:42:26 +0200 Subject: [PATCH 07/23] test gridpacks trained in scalar mode across SIMD modes - 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 --- .github/workflows/acceptancetest_mg7.yml | 5 +- .../iolibs/template_files/mg7/gridpack.py | 9 +++ madgraph/various/banner.py | 2 +- .../test_mg7_simd_consistency.py | 79 ++++++++++++++----- 4 files changed, 72 insertions(+), 23 deletions(-) diff --git a/.github/workflows/acceptancetest_mg7.yml b/.github/workflows/acceptancetest_mg7.yml index 0b461938c4..3694cc62a1 100644 --- a/.github/workflows/acceptancetest_mg7.yml +++ b/.github/workflows/acceptancetest_mg7.yml @@ -292,14 +292,15 @@ jobs: # mg7 seeded generation (g g > t t~ g), repeated for every madspace SIMD # mode (madspace_cpu_mode) supported by the runner: the VEGAS cross # sections must agree with the scalar one up to rounding, the madnis ones - # within their statistical uncertainties. Self-skips if the runner supports + # within their statistical uncertainties, and the runs of a gridpack trained + # in scalar mode again up to rounding. Self-skips if the runner supports # no SIMD mode or the madspace + LHAPDF(NNPDF23) stack is unavailable. - name: "acceptancetest_mg7_simd_consistency" if: ${{ !cancelled() }} run: | cd $GITHUB_WORKSPACE export PATH="$HOME/.cache/HEPtools/bin:$PATH" - ./tests/test_manager.py test_vegas_simd_consistency_mg7 test_madnis_simd_consistency_mg7 -pA -t0 -l INFO + ./tests/test_manager.py test_vegas_simd_consistency_mg7 test_madnis_simd_consistency_mg7 test_madnis_gridpack_simd_consistency_mg7 -pA -t0 -l INFO acceptancetest_mg7_bundle_3: diff --git a/madgraph/iolibs/template_files/mg7/gridpack.py b/madgraph/iolibs/template_files/mg7/gridpack.py index 215ca62dc3..496852c511 100644 --- a/madgraph/iolibs/template_files/mg7/gridpack.py +++ b/madgraph/iolibs/template_files/mg7/gridpack.py @@ -121,6 +121,13 @@ def main() -> None: "run's info.json)" ) parser.add_argument("--device", type=str, nargs="*") + parser.add_argument( + "--madspace_cpu_mode", + type=str, + default=run_args.get("madspace_cpu_mode", "scalar"), + choices=["auto", "scalar", "simd_128", "simd_256", "avx512y", "simd_512"], + help="SIMD width of the phase-space sampling (madspace) on the 'cpu' devices", + ) parser.add_argument( "--cpu_thread_pool_size", type=int, default=run_args["cpu_thread_pool_size"] ) @@ -146,6 +153,8 @@ def main() -> None: parser.add_argument("--gpu_batch_size", type=int, default=gen_args["gpu_batch_size"]) args = parser.parse_args() seed = resolve_seed(args.seed) + # has to be set before the madspace CPU backend is first loaded + ms.set_simd_mode(args.madspace_cpu_mode) # initialize event directory run_name = args.run_name diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py index 3ee563c173..0539654b5a 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -6681,7 +6681,7 @@ def default_setup(self): comment="matrix-element floating-point precision: all64 (FP64), all32 (FP32), " "color32 (colour FP32, rest FP64), " "denom64 (momenta + propagator denominator FP64, rest FP32)") - self.add_toml_param('run', 'madspace_cpu_mode', "scalar", + self.add_toml_param('run', 'madspace_cpu_mode', "scalar", gridpack=True, allowed=['auto', 'scalar', 'simd_128', 'simd_256', 'avx512y', 'simd_512'], comment="SIMD width of the phase-space sampling (madspace) on the 'cpu' devices " "(experimental); simd_128 is ARM only, 'auto' makes the same choice as for cpu_mode") diff --git a/tests/acceptance_tests/test_mg7_simd_consistency.py b/tests/acceptance_tests/test_mg7_simd_consistency.py index 2f2324c84b..4260bc80c2 100644 --- a/tests/acceptance_tests/test_mg7_simd_consistency.py +++ b/tests/acceptance_tests/test_mg7_simd_consistency.py @@ -26,6 +26,9 @@ * ``test_madnis_simd_consistency_mg7`` -- survey + madnis training + generate. The rounding differences grow during the training, so here the results only have to agree within their statistical uncertainties. + * ``test_madnis_gridpack_simd_consistency_mg7`` -- a gridpack trained with + madnis in scalar mode, then event generation from that gridpack in every + mode. Without the training drift, the same tolerance as for VEGAS applies. Process: ``g g > t t~ g``, whose multi-channel phase space uses momentum permutations. Self-skips if the mg7 runtime stack / PDF is unavailable, or if @@ -60,8 +63,8 @@ _PROCESS = 'g g > t t~ g' _EVENTS = int(os.environ.get('MG7_SIMD_EVENTS', 5000)) _SEED = 424242 -# relative tolerance for the VEGAS runs, which only differ by rounding -_VEGAS_REL_TOL = 1e-4 +# relative tolerance for runs that only differ by rounding (VEGAS, gridpacks) +_ROUNDING_REL_TOL = 1e-4 # tolerance for the madnis runs, in combined standard deviations _MADNIS_SIGMAS = 4. @@ -106,46 +109,54 @@ def _output_process(self, run_name, **settings): rc.write(toml) return run_dir - def _cross_section(self, run_dir, datadir, mode): - """Run bin/generate_events -f with the given madspace_cpu_mode and return - the (cross section, error) from the info.json of that run.""" - shutil.rmtree(pjoin(run_dir, 'Events'), ignore_errors=True) + def _set_mode(self, run_dir, mode): toml = pjoin(run_dir, 'Cards', 'run_card.toml') rc = RunCardMG7(toml) rc.set('run.madspace_cpu_mode', mode, user=True) rc.write(toml) + + def _cross_section(self, cmd, run_dir, datadir, mode, what): + """Run *cmd* in *run_dir* from scratch and return the (cross section, + error) from the info.json of that run.""" + shutil.rmtree(pjoin(run_dir, 'Events'), ignore_errors=True) env = dict(os.environ) env['LHAPDF_DATA_PATH'] = datadir - _run([sys.executable, pjoin(run_dir, 'bin', 'generate_events'), '-f'], - run_dir, pjoin(run_dir, 'gen_%s.log' % mode), env, - 'mg7 generate_events (madspace_cpu_mode=%s)' % mode) + _run(cmd, run_dir, pjoin(run_dir, 'gen_%s.log' % mode), env, + '%s (madspace_cpu_mode=%s)' % (what, mode)) infos = sorted(glob.glob(pjoin(run_dir, 'Events', '*', 'info.json'))) self.assertTrue(infos, 'no info.json produced for mode %s' % mode) with open(infos[-1]) as f: process = json.load(f)['process'] + print('madspace_cpu_mode=%s: %.6g +- %.2g' + % (mode, process['mean'], process['error'])) return process['mean'], process['error'] def _results(self, run_dir, modes): + """Cross sections of full runs (bin/generate_events -f) per mode.""" datadir = _mg7_datadir_or_skip(self) - results = {mode: self._cross_section(run_dir, datadir, mode) - for mode in modes} - for mode, (mean, error) in results.items(): - print('madspace_cpu_mode=%s: %.6g +- %.2g' % (mode, mean, error)) + results = {} + for mode in modes: + self._set_mode(run_dir, mode) + results[mode] = self._cross_section( + [sys.executable, pjoin(run_dir, 'bin', 'generate_events'), '-f'], + run_dir, datadir, mode, 'mg7 generate_events') return results + def _assert_rounding_level(self, results): + ref_mean, _ = results['scalar'] + for mode, (mean, _) in results.items(): + self.assertLessEqual( + abs(mean - ref_mean), _ROUNDING_REL_TOL * abs(ref_mean), + 'madspace_cpu_mode=%s: cross section %g differs from the scalar ' + 'one %g' % (mode, mean, ref_mean)) + def test_vegas_simd_consistency_mg7(self): """VEGAS runs: every SIMD mode reproduces the scalar cross section up to rounding.""" modes = self._modes_or_skip() _mg7_datadir_or_skip(self) run_dir = self._output_process('vegas', **{'madnis.enable': False}) - results = self._results(run_dir, modes) - ref_mean, _ = results['scalar'] - for mode, (mean, _) in results.items(): - self.assertLessEqual( - abs(mean - ref_mean), _VEGAS_REL_TOL * abs(ref_mean), - 'madspace_cpu_mode=%s: cross section %g differs from the scalar ' - 'one %g' % (mode, mean, ref_mean)) + self._assert_rounding_level(self._results(run_dir, modes)) def test_madnis_simd_consistency_mg7(self): """madnis runs: every SIMD mode agrees with the scalar cross section @@ -163,3 +174,31 @@ def test_madnis_simd_consistency_mg7(self): 'madspace_cpu_mode=%s: cross section %g +- %g is not compatible ' 'with the scalar one %g +- %g' % (mode, mean, error, ref_mean, ref_error)) + + def test_madnis_gridpack_simd_consistency_mg7(self): + """A gridpack trained with madnis in scalar mode: event generation from + it reproduces the scalar cross section up to rounding in every mode.""" + modes = self._modes_or_skip() + datadir = _mg7_datadir_or_skip(self) + run_dir = self._output_process( + 'madnis_gridpack', + **{'madnis.enable': True, 'madnis.train_batches': 200, + 'gridpack.save_gridpack': True}) + self._set_mode(run_dir, 'scalar') + env = dict(os.environ) + env['LHAPDF_DATA_PATH'] = datadir + _run([sys.executable, pjoin(run_dir, 'bin', 'generate_events'), '-f'], + run_dir, pjoin(run_dir, 'train.log'), env, + 'mg7 madnis training run') + gridpacks = glob.glob(pjoin(run_dir, 'Events', '*', 'gridpack')) + self.assertTrue(gridpacks, 'gridpack was not produced under %s' % run_dir) + gridpack_dir = gridpacks[0] + + results = {} + for mode in modes: + results[mode] = self._cross_section( + [sys.executable, pjoin(gridpack_dir, 'bin', 'generate_events'), + '--seed', str(_SEED), '--events', str(_EVENTS), + '--output_format', 'compact_npy', '--madspace_cpu_mode', mode], + gridpack_dir, datadir, mode, 'gridpack generate_events') + self._assert_rounding_level(results) From 7095b41f26fddfd69507d117f152c8c16e4989c8 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 7 Oct 2026 16:54:40 +0200 Subject: [PATCH 08/23] fix NaN in cut-bounded t-channel inverse, test inverse with a pt cut - 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 --- madspace/src/kernels/kinematics.hpp | 48 +++++++++++++---------------- madspace/tests/test_processes.py | 33 ++++++++++++++++---- 2 files changed, 49 insertions(+), 32 deletions(-) diff --git a/madspace/src/kernels/kinematics.hpp b/madspace/src/kernels/kinematics.hpp index a1b21b5cfb..2dfa9f020c 100644 --- a/madspace/src/kernels/kinematics.hpp +++ b/madspace/src/kernels/kinematics.hpp @@ -355,8 +355,7 @@ KERNELSPEC FourMom boost(FourMom k, FourMom p_boost, FVal sign) { } template -KERNELSPEC FourMom -boost_light_cone( +KERNELSPEC FourMom boost_light_cone( FourMom k, FVal k_mass2, FourMom p_boost, FVal p_mass2, FVal sign ) { // Same boost as `boost`, for a momentum k of known mass^2 k_mass2, written @@ -389,12 +388,10 @@ boost_light_cone( auto k_plus_direct = k[0] + k_par; auto k_minus_direct = k[0] - k_par; auto forward = k_par >= 0.; - auto k_plus = where( - forward, k_plus_direct, trans2 / max(k_minus_direct, FVal(EPS2)) - ); - auto k_minus = where( - forward, trans2 / max(k_plus_direct, FVal(EPS2)), k_minus_direct - ); + auto k_plus = + where(forward, k_plus_direct, trans2 / max(k_minus_direct, FVal(EPS2))); + auto k_minus = + where(forward, trans2 / max(k_plus_direct, FVal(EPS2)), k_minus_direct); auto k_plus_new = exp_y * k_plus; auto k_minus_new = exp_my * k_minus; @@ -435,8 +432,7 @@ KERNELSPEC Pair, FourMom> boost_two_body( } template -KERNELSPEC void -boost_beam( +KERNELSPEC void boost_beam( FIn q, FIn masses, FVal x1, @@ -476,12 +472,10 @@ boost_beam( auto trans2 = m2 + pt2; auto q_plus_direct = q0 + q3; auto q_minus_direct = q0 - q3; - auto q_plus = where( - q3 >= 0., q_plus_direct, trans2 / max(q_minus_direct, FVal(EPS2)) - ); - auto q_minus = where( - q3 < 0., q_minus_direct, trans2 / max(q_plus_direct, FVal(EPS2)) - ); + auto q_plus = + where(q3 >= 0., q_plus_direct, trans2 / max(q_minus_direct, FVal(EPS2))); + auto q_minus = + where(q3 < 0., q_minus_direct, trans2 / max(q_plus_direct, FVal(EPS2))); auto p_plus = e_plus * q_plus; auto p_minus = e_minus * q_minus; // Only for a momentum that is on its mass shell to rounding. The @@ -491,8 +485,8 @@ boost_beam( // gained, and it keeps the linear boost. auto on_shell = fabs(q_plus_direct * q_minus_direct - trans2) <= 1e-13 * q0 * q0; - auto light_cone = (q_plus_direct >= 0.) & (q_minus_direct >= 0.) & on_shell & - (lin_0 < q0); + auto light_cone = + (q_plus_direct >= 0.) & (q_minus_direct >= 0.) & on_shell & (lin_0 < q0); p_out_i[0] = where(light_cone, 0.5 * (p_plus + p_minus), lin_0); p_out_i[1] = q1; p_out_i[2] = q2; @@ -522,10 +516,10 @@ boost_beam( auto lin_0 = q0 * cosh_rap + q3 * sinh_rap; auto lin_3 = q3 * cosh_rap + q0 * sinh_rap; auto is_best = best_i == IVal(static_cast(i)); - shift_0 = shift_0 + - where(is_best, FVal(0.), FVal(p_out[i][0]) - lin_0); - shift_3 = shift_3 + - where(is_best, FVal(0.), FVal(p_out[i][3]) - lin_3); + shift_0 = + shift_0 + where(is_best, FVal(0.), FVal(p_out[i][0]) - lin_0); + shift_3 = + shift_3 + where(is_best, FVal(0.), FVal(p_out[i][3]) - lin_3); best_lin_0 = where(is_best, lin_0, best_lin_0); best_lin_3 = where(is_best, lin_3, best_lin_3); } @@ -603,16 +597,14 @@ KERNELSPEC Quartuplet, FVal, FVal, FVal> phi_m1_m2_from_p1com( // Kernels template -KERNELSPEC void -kernel_boost_beam( +KERNELSPEC void kernel_boost_beam( FIn p1, FIn masses, FIn x1, FIn x2, FOut p_out ) { boost_beam(p1, masses, x1, x2, 1.0, p_out); } template -KERNELSPEC void -kernel_boost_beam_inverse( +KERNELSPEC void kernel_boost_beam_inverse( FIn p1, FIn masses, FIn x1, FIn x2, FOut p_out ) { boost_beam(p1, masses, x1, x2, -1.0, p_out); @@ -782,6 +774,10 @@ KERNELSPEC Pair, FVal> t_cut_bounds_etmin( FVal etmin_1, FVal etmin_2 ) { + // squared masses recomputed from momenta (inverse direction) can be slightly + // negative; clamp them so that the square roots below cannot give NaN + m1_2 = max(m1_2, 0.); + m2_2 = max(m2_2, 0.); auto pt_tot2 = p_tot[1] * p_tot[1] + p_tot[2] * p_tot[2]; auto s = lsquare(p_tot); auto piir0 = sqrt(max(s, 0.) + pt_tot2); diff --git a/madspace/tests/test_processes.py b/madspace/tests/test_processes.py index 4799a986c8..664eee326b 100644 --- a/madspace/tests/test_processes.py +++ b/madspace/tests/test_processes.py @@ -49,6 +49,24 @@ def mapping(process): ) +@pytest.fixture +def cut_mapping(process): + # a mild pt cut keeps the events away from extremely soft gluons, for which the + # inverse of the t-channel invariants is numerically ill-conditioned + pids = [21, 21] + [6 if mass > 0 else 21 for mass in process["outgoing_masses"]] + O = ms.Observable + cuts = ms.Cuts([ms.CutItem(O(pids, O.obs_pt, [O.jet_pids]), min=1.0)]) + diagram = ms.Diagram( + process["incoming_masses"], + process["outgoing_masses"], + [ms.Propagator(*prop) for prop in process["propagators"]], + process["vertices"], + ) + return ms.PhaseSpaceMapping( + ms.Topology(diagram), CM_ENERGY, cuts=cuts, permutations=process["permutations"] + ) + + @pytest.fixture def masses(process): return process["incoming_masses"] + process["outgoing_masses"] @@ -110,14 +128,17 @@ def test_process_momentum_conservation(mapping, masses, permutation_count): assert p_out == approx(p_in, rel=1e-6, abs=1e-10) -def test_process_inverse(mapping, masses, permutation_count): - r = rng.random((BATCH_SIZE, mapping.random_dim())) +def test_process_inverse(cut_mapping, masses, permutation_count): + r = rng.random((BATCH_SIZE, cut_mapping.random_dim())) condition = ( [] if permutation_count <= 1 else [rng.integers(0, permutation_count, BATCH_SIZE, dtype=np.int32)] ) - *map_out, det = mapping.map_forward([r], condition) - r_inv, det_inv = mapping.map_inverse(map_out, condition) - assert r_inv == approx(r, abs=1e-3, rel=1e-3) - assert det_inv == approx(1 / det, rel=1e-3) + *map_out, det = cut_mapping.map_forward([r], condition) + r_inv, det_inv = cut_mapping.map_inverse(map_out, condition) + # events removed by the cuts (zero weight) are outside the phase space + passed = det > 0 + assert np.any(passed) + assert r_inv[passed] == approx(r[passed], abs=1e-3, rel=1e-3) + assert det_inv[passed] == approx(1 / det[passed], rel=1e-3) From 0874881d8e1a84bcfeca7648286aeb69023b4dde Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 7 Oct 2026 21:58:31 +0200 Subject: [PATCH 09/23] madspace: fix cut_all/cut_any bound rank, contiguous_dims prefix, broadcast 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 --- madspace/generate_headers.py | 2 +- madspace/include/madspace/driver/tensor.hpp | 7 +++++-- madspace/src/cpu/device.cpp | 4 +++- madspace/src/kernels/cuts.hpp | 14 ++++++++------ madspace/src/kernels/discrete.hpp | 2 +- 5 files changed, 18 insertions(+), 11 deletions(-) diff --git a/madspace/generate_headers.py b/madspace/generate_headers.py index 998615913c..eab5f9eeab 100644 --- a/madspace/generate_headers.py +++ b/madspace/generate_headers.py @@ -225,7 +225,7 @@ def instruction_set_mixin(commands, out_dir): first = True for name, cmd in commands.items(): opcode = cmd["opcode"] - differentiable = "true" if cmd.get("differentiable", False) else "true" + differentiable = "true" if cmd.get("differentiable", False) else "false" if "class" in cmd: f.write( f" InstructionOwner(new {cmd['class']}(" diff --git a/madspace/include/madspace/driver/tensor.hpp b/madspace/include/madspace/driver/tensor.hpp index 37aa0bfdf5..165baf0079 100644 --- a/madspace/include/madspace/driver/tensor.hpp +++ b/madspace/include/madspace/driver/tensor.hpp @@ -389,11 +389,14 @@ class Tensor { dtype, shape, device, data, false, external_reset, nullptr, 1, stride }) { std::size_t stride_prod = 1; - bool first = true; + bool first = true, contiguous = true; impl->contiguous_dims = 0; for (auto [size_i, stride_i] : zip(shape, stride)) { - if (stride_i == stride_prod) { + // only the leading run of contiguous dimensions counts + if (contiguous && (stride_i == stride_prod || size_i == 1)) { ++impl->contiguous_dims; + } else { + contiguous = false; } if (first && size_i == 1) { impl->stride[0] = 0; diff --git a/madspace/src/cpu/device.cpp b/madspace/src/cpu/device.cpp index d1249ece46..649fce51ed 100644 --- a/madspace/src/cpu/device.cpp +++ b/madspace/src/cpu/device.cpp @@ -49,11 +49,13 @@ void tensor_zero_impl(Tensor& tensor, const D& device) { template void tensor_add_impl(const Tensor& source, Tensor& target, const D& device) { + // a broadcast target (e.g. the gradient of a literal or global) accumulates + // the whole source batch tensor_foreach_dynamic_single< kernel_add_inplace, kernel_add_inplace, 1, - 1>({&source}, {&target}, target.size(0), device); + 1>({&source}, {&target}, std::max(source.size(0), target.size(0)), device); } } // namespace diff --git a/madspace/src/kernels/cuts.hpp b/madspace/src/kernels/cuts.hpp index e4e353c9bf..178020d0b9 100644 --- a/madspace/src/kernels/cuts.hpp +++ b/madspace/src/kernels/cuts.hpp @@ -31,20 +31,22 @@ kernel_cut_one(FIn obs, FIn min, FIn max, FOut w) { template KERNELSPEC void -kernel_cut_all(FIn obs, FIn min, FIn max, FOut w) { - FVal cut = 1.; +kernel_cut_all(FIn obs, FIn min, FIn max, FOut w) { + FVal cut = 1., min_val = min, max_val = max; for (std::size_t i = 0; i < obs.size(); ++i) { - cut = where((obs[i] < min[i]) | (obs[i] > max[i]), 0., cut); + FVal obs_i = obs[i]; + cut = where((obs_i < min_val) | (obs_i > max_val), 0., cut); } w = cut; } template KERNELSPEC void -kernel_cut_any(FIn obs, FIn min, FIn max, FOut w) { - FVal cut = 0.; +kernel_cut_any(FIn obs, FIn min, FIn max, FOut w) { + FVal cut = 0., min_val = min, max_val = max; for (std::size_t i = 0; i < obs.size(); ++i) { - cut = where((obs[i] < min[i]) | (obs[i] > max[i]), cut, 1.); + FVal obs_i = obs[i]; + cut = where((obs_i < min_val) | (obs_i > max_val), cut, 1.); } w = cut; } diff --git a/madspace/src/kernels/discrete.hpp b/madspace/src/kernels/discrete.hpp index 9622342005..5f099cfa04 100644 --- a/madspace/src/kernels/discrete.hpp +++ b/madspace/src/kernels/discrete.hpp @@ -87,7 +87,7 @@ KERNELSPEC void backward_kernel_sample_discrete_probs_inverse( FVal det_grad_out(0.); auto prob = probs.gather(index) / prob_norm; for (std::size_t i = 0; i < probs.size(); ++i) { - probs_grad[i] = where( + probs_grad[i] += where( zero_norm, 0., (where(index == i, FVal(1.), 0.) - prob) / prob_norm * det_grad From c5c120b69253292c2ea6068b9f1666032226209d Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 7 Oct 2026 22:01:48 +0200 Subject: [PATCH 10/23] madspace: fix sub(0, x) folding, fold stacks of scalar literals - 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 --- madspace/src/compgraphs/function.cpp | 118 ++++++++++++++++----------- 1 file changed, 71 insertions(+), 47 deletions(-) diff --git a/madspace/src/compgraphs/function.cpp b/madspace/src/compgraphs/function.cpp index cc6ce199d6..bc2c307d66 100644 --- a/madspace/src/compgraphs/function.cpp +++ b/madspace/src/compgraphs/function.cpp @@ -1,6 +1,7 @@ #include "madspace/compgraphs/function.hpp" #include "madspace/util.hpp" +#include #include #include #include @@ -306,58 +307,81 @@ FunctionBuilder::instruction(InstructionPtr instruction, const ValueVec& args) { // perform simple arithmetic operations on constant arguments int opcode = instruction->opcode(); - if (variable_args == 0) { - switch (opcode) { - case opcodes::add: { - double arg0 = std::get(args.at(0).literal_value); - double arg1 = std::get(args.at(1).literal_value); - return {arg0 + arg1}; - } - case opcodes::sub: { - double arg0 = std::get(args.at(0).literal_value); - double arg1 = std::get(args.at(1).literal_value); - return {arg0 - arg1}; - } - case opcodes::mul: { - double arg0 = std::get(args.at(0).literal_value); - double arg1 = std::get(args.at(1).literal_value); - return {arg0 * arg1}; - } - case opcodes::sqrt: { - double arg0 = std::get(args.at(0).literal_value); - return {std::sqrt(arg0)}; - } - case opcodes::square: { - double arg0 = std::get(args.at(0).literal_value); - return {arg0 * arg0}; - } - case opcodes::min: { - double arg0 = std::get(args.at(0).literal_value); - double arg1 = std::get(args.at(1).literal_value); - return {std::min(arg0, arg1)}; - } - case opcodes::max: { - double arg0 = std::get(args.at(0).literal_value); - double arg1 = std::get(args.at(1).literal_value); - return {std::max(arg0, arg1)}; - } - } - } else if (variable_args == 1 && args.size() == 2) { - std::size_t var_index = args.at(0).local_index == -1; - auto& var_arg = args.at(var_index); - auto& const_arg = args.at(1 - var_index); + auto is_float_literal = [](const Value& arg) { + return arg.local_index == -1 && + std::holds_alternative(arg.literal_value); + }; + auto is_int_literal = [](const Value& arg) { + return arg.local_index == -1 && + std::holds_alternative(arg.literal_value); + }; + bool all_float_literals = args.size() > 0 && + std::all_of(args.begin(), args.end(), is_float_literal); + bool all_int_literals = + args.size() > 0 && std::all_of(args.begin(), args.end(), is_int_literal); + if (all_float_literals) { + auto arg = [&](std::size_t index) { + return std::get(args.at(index).literal_value); + }; switch (opcode) { case opcodes::add: + return {arg(0) + arg(1)}; case opcodes::sub: - if (std::get(const_arg.literal_value) == 0.) { - return {var_arg}; - } - break; + return {arg(0) - arg(1)}; case opcodes::mul: - if (std::get(const_arg.literal_value) == 1.) { - return {var_arg}; + return {arg(0) * arg(1)}; + case opcodes::sqrt: + return {std::sqrt(arg(0))}; + case opcodes::square: + return {arg(0) * arg(0)}; + case opcodes::min: + return {std::min(arg(0), arg(1))}; + case opcodes::max: + return {std::max(arg(0), arg(1))}; + } + } + // a stack of scalar literals is a single tensor literal + if (opcode == opcodes::stack) { + if (all_float_literals) { + std::vector values; + for (auto& arg : args) { + values.push_back(std::get(arg.literal_value)); + } + return {Value(values)}; + } else if (all_int_literals) { + std::vector values; + for (auto& arg : args) { + values.push_back(std::get(arg.literal_value)); + } + return {Value(values)}; + } + } + // identities x + 0, 0 + x, x - 0 and x * 1, 1 * x + if (variable_args == 1 && args.size() == 2) { + bool const_first = is_float_literal(args.at(0)); + bool const_second = is_float_literal(args.at(1)); + if (const_first != const_second) { + auto& var_arg = args.at(const_first ? 1 : 0); + double const_val = + std::get(args.at(const_first ? 0 : 1).literal_value); + switch (opcode) { + case opcodes::add: + if (const_val == 0.) { + return {var_arg}; + } + break; + case opcodes::sub: + // only x - 0, as 0 - x = -x + if (const_second && const_val == 0.) { + return {var_arg}; + } + break; + case opcodes::mul: + if (const_val == 1.) { + return {var_arg}; + } + break; } - break; } } From dadfe3a104737f29e7b64d8079ebfe16a7a5f95c Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 7 Oct 2026 22:03:35 +0200 Subject: [PATCH 11/23] madspace: avoid constant gathers and full(1) products with a single flavor 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 --- madspace/src/phasespace/integrand.cpp | 35 ++++++++++++++++++++++----- 1 file changed, 29 insertions(+), 6 deletions(-) diff --git a/madspace/src/phasespace/integrand.cpp b/madspace/src/phasespace/integrand.cpp index b2e74cb3dc..3f9e6924e2 100644 --- a/madspace/src/phasespace/integrand.cpp +++ b/madspace/src/phasespace/integrand.cpp @@ -553,6 +553,9 @@ NamedVector Integrand::build_channel_part( } )) { option_count = static_cast(2); + } else if (!has_multi_flavor) { + // flavor_id is always zero with a single flavor + option_count = static_cast(_flavor_mirror.at(0)); } else { option_count = fb.gather_int(flavor_id, _flavor_mirror); } @@ -637,6 +640,18 @@ NamedVector Integrand::build_common_part( auto flavor_id = args.at("flavor_id"); auto batch_size_val = fb.batch_size({args.at("weight_before_cuts")}); + // With a single flavor, flavor_id is always zero, so looking up per-flavor + // tables does not need a gather on a constant tensor + auto flavor_table_lookup = [&](const std::vector& table) -> Value { + if (has_multi_flavor) { + return fb.gather_int(flavor_id, table); + } else if (table.at(0) == 0) { + return flavor_id; + } else { + return fb.full({table.at(0), fb.batch_size({momenta_acc})}); + } + }; + auto scatter_or_drop = [&](Value default_value, Value value) -> Value { if (_drop_cuts_and_rescale) { return value; @@ -681,7 +696,7 @@ NamedVector Integrand::build_common_part( // Evaluate differential cross section ValueVec xs_args{ momenta_acc, - _flavor_remap.size() > 0 ? fb.gather_int(flavor_id, _flavor_remap) : flavor_id, + _flavor_remap.size() > 0 ? flavor_table_lookup(_flavor_remap) : flavor_id, }; xs_args.push_back(x1_acc); xs_args.push_back(x2_acc); @@ -710,9 +725,9 @@ NamedVector Integrand::build_common_part( } } } else { - Value dxs_index = fb.gather_int(flavor_id, _flavor_diff_xs_indices); - ps_flavor_id = fb.gather_int(flavor_id, _flavor_per_subproc_remap); - subproc_id = fb.gather_int(flavor_id, _flavor_subproc_indices); + Value dxs_index = flavor_table_lookup(_flavor_diff_xs_indices); + ps_flavor_id = flavor_table_lookup(_flavor_per_subproc_remap); + subproc_id = flavor_table_lookup(_flavor_subproc_indices); ValueVec split_indices = fb.batch_split_by_index(dxs_index, static_cast(_diff_xs.size())); std::vector split_outputs(_diff_xs.at(0).return_types().size()); @@ -756,9 +771,17 @@ NamedVector Integrand::build_common_part( } auto diff_xs_acc = dxs_vec.at(0); if (_flavor_factors.size() > 0) { - diff_xs_acc = fb.mul(diff_xs_acc, fb.gather(flavor_id, _flavor_factors)); + diff_xs_acc = has_multi_flavor + ? fb.mul(diff_xs_acc, fb.gather(flavor_id, _flavor_factors)) + : fb.mul(diff_xs_acc, _flavor_factors.at(0)); + } + // weight_after_cuts is one by construction unless flavor or mirror sampling + // contribute to it (see build_channel_part) + ValueVec weights_after_cuts; + if (has_multi_flavor || _has_mirror) { + weights_after_cuts.push_back(args.at("weight_after_cuts")); } - ValueVec weights_after_cuts{args.at("weight_after_cuts"), diff_xs_acc}; + weights_after_cuts.push_back(diff_xs_acc); ValueVec extra_weights_after_cuts; if (args.index_map().contains("extra_weight_after_cuts")) { extra_weights_after_cuts.push_back(args.at("extra_weight_after_cuts")); From bef49aaa9796a517ad60ff3dc581fd1d0e3dafb1 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 7 Oct 2026 22:23:32 +0200 Subject: [PATCH 12/23] madspace: contiguous SIMD loads/stores, single inputs as a separate kernel 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 --- madspace/CMakeLists.txt | 2 + madspace/generate_headers.py | 56 +++++ madspace/include/madspace/driver/tensor.hpp | 3 + madspace/instruction_set.yaml | 136 +++++----- madspace/src/cpu/kernel_definitions.hpp | 8 + madspace/src/cpu/runtime.cpp | 5 + madspace/src/cpu/simd_arm.hpp | 30 +-- madspace/src/cpu/simd_x86_256.hpp | 33 +-- madspace/src/cpu/simd_x86_256y.hpp | 26 +- madspace/src/cpu/simd_x86_512.hpp | 31 ++- madspace/src/cpu/tensor.hpp | 265 +++++++++++++++----- madspace/src/gpu/kernel_definitions.cuh | 5 + madspace/src/kernels/chili.hpp | 8 +- madspace/src/kernels/cuts.hpp | 6 +- madspace/src/kernels/definitions.hpp | 5 + madspace/src/kernels/discrete.hpp | 2 +- madspace/src/kernels/invariants.hpp | 38 +-- madspace/src/kernels/kinematics.hpp | 42 ++-- madspace/src/kernels/multichannel.hpp | 22 +- madspace/src/kernels/threeparticle.hpp | 8 +- madspace/src/kernels/twoparticle.hpp | 12 +- madspace/src/kernels/util.hpp | 8 +- madspace/src/kernels/vegas.hpp | 4 +- 23 files changed, 478 insertions(+), 277 deletions(-) diff --git a/madspace/CMakeLists.txt b/madspace/CMakeLists.txt index 69e7081814..e334c6c1da 100644 --- a/madspace/CMakeLists.txt +++ b/madspace/CMakeLists.txt @@ -177,6 +177,7 @@ set( ${MADSPACE_GENERATED_DIR}/python/instruction_set.hpp ${MADSPACE_GENERATED_DIR}/cpu/runtime_mixin.inc ${MADSPACE_GENERATED_DIR}/cpu/runtime_backward_mixin.inc + ${MADSPACE_GENERATED_DIR}/cpu/kernel_checks.inc ${MADSPACE_GENERATED_DIR}/gpu/runtime_mixin.inc ${MADSPACE_GENERATED_DIR}/gpu/runtime_backward_mixin.inc ) @@ -472,6 +473,7 @@ foreach( src/cpu/kernel_definitions.hpp ${MADSPACE_GENERATED_DIR}/cpu/runtime_mixin.inc ${MADSPACE_GENERATED_DIR}/cpu/runtime_backward_mixin.inc + ${MADSPACE_GENERATED_DIR}/cpu/kernel_checks.inc src/cpu/tensor.hpp src/cpu/simd.hpp src/cpu/simd_arm.hpp diff --git a/madspace/generate_headers.py b/madspace/generate_headers.py index eab5f9eeab..d77b66ba10 100644 --- a/madspace/generate_headers.py +++ b/madspace/generate_headers.py @@ -42,6 +42,7 @@ def main(): runtime_backward_mixin(commands, "cpu", out_dir) runtime_mixin(commands, "gpu", out_dir) runtime_backward_mixin(commands, "gpu", out_dir) + kernel_checks(commands, out_dir) def write_autogen(f): @@ -352,5 +353,60 @@ def runtime_backward_mixin(commands, device, out_dir): ) +def is_single_input(arg): + return isinstance(arg.get("type"), list) and len(arg["type"]) > 1 and ( + arg["type"][1] == "single" + ) + + +def kernel_checks(commands, out_dir): + """Compile-time checks that the parameters of the vectorized kernels use the + single types (FSIn, ISIn) exactly for the inputs marked as single in the yaml. + A mismatch would make the CPU runtime fall back to the scalar kernels.""" + with open(out_dir / "cpu" / "kernel_checks.inc", "w") as f: + write_autogen(f) + + def check(kernel, index, single, what): + f.write( + f"static_assert(\n" + f" single_params)>::value[{index}] == " + f"{'true' if single else 'false'},\n" + f' "{kernel}: parameter {index} ({what}) must ' + f'{"" if single else "not "}be single"\n' + f");\n" + ) + + for name, cmd in commands.items(): + if cmd.get("custom_op", False) or not cmd.get("vectorized", True): + continue + inputs = cmd["inputs"] + if is_single_input(inputs[0]): + raise ValueError( + f"{name}: the first input of a kernel can not be single" + ) + for i, arg in enumerate(inputs): + check(f"kernel_{name}", i, is_single_input(arg), arg["name"]) + + if not cmd.get("differentiable", False): + continue + in_stored = [ + i for i, arg in enumerate(inputs) if arg.get("backward_arg", False) + ] + out_stored = [ + i + for i, arg in enumerate(cmd["outputs"]) + if arg.get("backward_arg", False) + ] + if len(in_stored) + len(out_stored) == 0: + in_stored = list(range(len(inputs))) + for k, i in enumerate(in_stored): + check( + f"backward_kernel_{name}", + k, + is_single_input(inputs[i]), + inputs[i]["name"], + ) + + if __name__ == "__main__": main() diff --git a/madspace/include/madspace/driver/tensor.hpp b/madspace/include/madspace/driver/tensor.hpp index 165baf0079..c2e78bdd1a 100644 --- a/madspace/include/madspace/driver/tensor.hpp +++ b/madspace/include/madspace/driver/tensor.hpp @@ -112,6 +112,9 @@ struct PackedTensorView { T* data; Sizes stride; Sizes shape; + /// Whether the data is a literal stored padded to the CPU SIMD vector size + /// (see @ref Tensor::is_simd_padded). + bool simd_padded = false; }; /** diff --git a/madspace/instruction_set.yaml b/madspace/instruction_set.yaml index 447387ef98..9119731803 100644 --- a/madspace/instruction_set.yaml +++ b/madspace/instruction_set.yaml @@ -644,7 +644,7 @@ boost_beam: type: [float, n, 4] desc: "momenta in the partonic center-of-mass frame" - name: masses - type: [float, n] + type: [float, single, n] desc: "mass of each particle, or a negative value to read it off its momentum" - name: x1 type: [float] @@ -664,7 +664,7 @@ boost_beam_inverse: type: [float, n, 4] desc: "momenta in the lab frame" - name: masses - type: [float, n] + type: [float, single, n] desc: "mass of each particle, or a negative value to read it off its momentum" - name: x1 type: [float] @@ -701,7 +701,7 @@ r_to_x1x2: type: [float] desc: "partonic invariant mass squared" - name: s_lab - type: [float] + type: [float, single] desc: "hadronic invariant mass squared" outputs: - name: x1 @@ -724,7 +724,7 @@ x1x2_to_r: type: [float] desc: "momentum fraction of beam 2" - name: s_lab - type: [float] + type: [float, single] desc: "hadronic invariant mass squared" outputs: - name: r @@ -744,10 +744,10 @@ r_to_x1x2_window: type: [float] desc: "partonic invariant mass squared" - name: s_lab - type: [float] + type: [float, single] desc: "hadronic invariant mass squared" - name: y_max - type: [float] + type: [float, single] desc: "largest allowed absolute rapidity of the partonic system" outputs: - name: x1 @@ -770,10 +770,10 @@ x1x2_to_r_window: type: [float] desc: "momentum fraction of beam 2" - name: s_lab - type: [float] + type: [float, single] desc: "hadronic invariant mass squared" - name: y_max - type: [float] + type: [float, single] desc: "largest allowed absolute rapidity of the partonic system" outputs: - name: r @@ -802,7 +802,7 @@ diff_cross_section: type: [float] desc: "matrix element value" - name: e_cm2 - type: [float] + type: [float, single] desc: "squared center-of-mass energy" outputs: - name: result @@ -889,13 +889,13 @@ com_cos_theta_window: type: [float] desc: "momentum fraction of beam 2" - name: pt_min - type: [float] + type: [float, single] desc: "lower bound on the transverse momentum of the products" - name: y_max1 - type: [float] + type: [float, single] desc: "bound on the absolute lab rapidity of the first product, negative for none" - name: y_max2 - type: [float] + type: [float, single] desc: "bound on the absolute lab rapidity of the second product, negative for none" outputs: - name: r_out @@ -927,13 +927,13 @@ com_cos_theta_window_inverse: type: [float] desc: "momentum fraction of beam 2" - name: pt_min - type: [float] + type: [float, single] desc: "lower bound on the transverse momentum of the products" - name: y_max1 - type: [float] + type: [float, single] desc: "bound on the absolute lab rapidity of the first product, negative for none" - name: y_max2 - type: [float] + type: [float, single] desc: "bound on the absolute lab rapidity of the second product, negative for none" outputs: - name: r_cos_theta @@ -1554,13 +1554,13 @@ t_rapidity_clamp: type: [float] desc: "momentum fraction of beam 2" - name: beam_sign - type: [float] + type: [float, single] desc: "+1 if pa is beam 1, -1 if it is beam 2" - name: y_max1 - type: [float] + type: [float, single] desc: "bound on the absolute lab rapidity of the recoil, negative or 100 and above for none" - name: y_max2 - type: [float] + type: [float, single] desc: "bound on the absolute lab rapidity of the peeled particle, negative or 100 and above for none" outputs: - name: t_min_out @@ -1604,13 +1604,13 @@ t1_rapidity_clamp_doublet: type: [float] desc: "momentum fraction of beam 2" - name: beam_sign - type: [float] + type: [float, single] desc: "+1 if pa is beam 1, -1 if it is beam 2" - name: y_max1 - type: [float] + type: [float, single] desc: "bound on the absolute lab rapidity of p1, negative or 100 and above for none" - name: y_max2 - type: [float] + type: [float, single] desc: "bound on the absolute lab rapidity of the recoil p2, negative or 100 and above for none" outputs: - name: t1_min_out @@ -1648,13 +1648,13 @@ t2_rapidity_clamp_doublet: type: [float] desc: "momentum fraction of beam 2" - name: beam_sign - type: [float] + type: [float, single] desc: "+1 if pa is beam 1, -1 if it is beam 2" - name: y_max1 - type: [float] + type: [float, single] desc: "bound on the absolute lab rapidity of p1, negative or 100 and above for none" - name: y_max2 - type: [float] + type: [float, single] desc: "bound on the absolute lab rapidity of the recoil p2, negative or 100 and above for none" outputs: - name: t2_min_out @@ -2009,10 +2009,10 @@ s23_arcsine_sample: type: [float] desc: "upper kinematic limit of the 2-3 invariant mass squared" - name: power - type: [float] + type: [float, single] desc: "exponent of the 1/(s - m^2)^power sampling, as in @ref Invariant" - name: mass - type: [float] + type: [float, single] desc: "pole mass of the sampling, as in @ref Invariant" outputs: - name: u @@ -2047,10 +2047,10 @@ s23_arcsine_sample_inverse: type: [float] desc: "upper kinematic limit of the 2-3 invariant mass squared" - name: power - type: [float] + type: [float, single] desc: "exponent of the 1/(s - m^2)^power sampling, as in @ref Invariant" - name: mass - type: [float] + type: [float, single] desc: "pole mass of the sampling, as in @ref Invariant" outputs: - name: r @@ -2067,7 +2067,7 @@ invariants_from_momenta: type: [float, n, 4] desc: "external four-momenta" - name: factors - type: [float, m, n] + type: [float, single, m, n] desc: "coefficient matrix selecting the invariant combinations" outputs: - name: invariants @@ -2081,13 +2081,13 @@ sde2_channel_weights: type: [float, m] desc: "propagator invariants" - name: masses - type: [float, c, n] + type: [float, single, c, n] desc: "propagator masses per channel" - name: widths - type: [float, c, n] + type: [float, single, c, n] desc: "propagator widths per channel" - name: indices - type: [int, c, n] + type: [int, single, c, n] desc: "invariant index feeding each propagator" outputs: - name: channel_weights @@ -2101,16 +2101,16 @@ subchannel_weights: type: [float, m] desc: "propagator invariants" - name: masses - type: [float, c, n] + type: [float, single, c, n] desc: "propagator masses per channel" - name: widths - type: [float, c, n] + type: [float, single, c, n] desc: "propagator widths per channel" - name: indices - type: [int, c, n] + type: [int, single, c, n] desc: "invariant index feeding each propagator" - name: on_shell - type: [int, c, n] + type: [int, single, c, n] desc: "flag marking on-shell propagators per channel" - name: group_sizes type: [int, single, g] @@ -2130,10 +2130,10 @@ apply_subchannel_weights: type: [float, n] desc: "weight per subchannel" - name: channel_indices - type: [int, d] + type: [int, single, d] desc: "channel index per diagram" - name: subchannel_indices - type: [int, d] + type: [int, single, d] desc: "subchannel index per diagram" outputs: - name: channel_weights_out @@ -2215,7 +2215,7 @@ momenta_to_x1x2: type: [float, n, 4] desc: "external four-momenta" - name: e_cm - type: [float] + type: [float, single] desc: "hadronic center-of-mass energy" outputs: - name: x1 @@ -2276,10 +2276,10 @@ breit_wigner_invariant: type: [float] desc: "random number in [0, 1]" - name: mass - type: [float] + type: [float, single] desc: "resonance mass" - name: width - type: [float] + type: [float, single] desc: "resonance width" - name: s_min type: [float] @@ -2302,10 +2302,10 @@ breit_wigner_invariant_inverse: type: [float] desc: "invariant value" - name: mass - type: [float] + type: [float, single] desc: "resonance mass" - name: width - type: [float] + type: [float, single] desc: "resonance width" - name: s_min type: [float] @@ -2328,13 +2328,13 @@ flat_window_breit_wigner_invariant: type: [float] desc: "random number in [0, 1]" - name: mass - type: [float] + type: [float, single] desc: "resonance mass" - name: width - type: [float] + type: [float, single] desc: "resonance width" - name: window - type: [float] + type: [float, single] desc: "half-width of the flattened window, in units of the width" - name: s_min type: [float] @@ -2357,13 +2357,13 @@ flat_window_breit_wigner_invariant_inverse: type: [float] desc: "invariant value" - name: mass - type: [float] + type: [float, single] desc: "resonance mass" - name: width - type: [float] + type: [float, single] desc: "resonance width" - name: window - type: [float] + type: [float, single] desc: "half-width of the flattened window, in units of the width" - name: s_min type: [float] @@ -2386,7 +2386,7 @@ stable_invariant: type: [float] desc: "random number in [0, 1]" - name: mass - type: [float] + type: [float, single] desc: "pole mass" - name: s_min type: [float] @@ -2409,7 +2409,7 @@ stable_invariant_inverse: type: [float] desc: "invariant value" - name: mass - type: [float] + type: [float, single] desc: "pole mass" - name: s_min type: [float] @@ -2432,10 +2432,10 @@ stable_invariant_nu: type: [float] desc: "random number in [0, 1]" - name: mass - type: [float] + type: [float, single] desc: "pole mass" - name: nu - type: [float] + type: [float, single] desc: "exponent of the power law" - name: s_min type: [float] @@ -2458,10 +2458,10 @@ stable_invariant_nu_inverse: type: [float] desc: "invariant value" - name: mass - type: [float] + type: [float, single] desc: "pole mass" - name: nu - type: [float] + type: [float, single] desc: "exponent of the power law" - name: s_min type: [float] @@ -2633,10 +2633,10 @@ cut_one: type: [float] desc: "observable value" - name: min - type: [float] + type: [float, single] desc: "lower bound" - name: max - type: [float] + type: [float, single] desc: "upper bound" outputs: - name: w @@ -2650,10 +2650,10 @@ cut_all: type: [float, n] desc: "observable values" - name: min - type: [float] + type: [float, single] desc: "lower bound applied to each observable" - name: max - type: [float] + type: [float, single] desc: "upper bound applied to each observable" outputs: - name: w @@ -2667,10 +2667,10 @@ cut_any: type: [float, n] desc: "observable values" - name: min - type: [float] + type: [float, single] desc: "lower bound applied to each observable" - name: max - type: [float] + type: [float, single] desc: "upper bound applied to each observable" outputs: - name: w @@ -2740,10 +2740,10 @@ chili_forward: type: [float, n] desc: "mass of each final-state particle" - name: pt_min - type: [float, n] + type: [float, single, n] desc: "minimum transverse momentum per final-state particle" - name: y_max - type: [float, n] + type: [float, single, n] desc: "maximum rapidity per final-state particle" outputs: - name: p_ext @@ -2766,10 +2766,10 @@ chili_inverse: type: [float, n] desc: "mass of each final-state particle" - name: pt_min - type: [float, n] + type: [float, single, n] desc: "minimum transverse momentum per final-state particle" - name: y_max - type: [float, n] + type: [float, single, n] desc: "maximum rapidity per final-state particle" outputs: - name: r @@ -3210,7 +3210,7 @@ gather_int: type: [int] desc: "index to select" - name: choices - type: [int, n] + type: [int, single, n] desc: "vector of choices" outputs: - name: output @@ -3224,7 +3224,7 @@ gather_vector: type: [int] desc: "row index to select" - name: choices - type: [float, n, m] + type: [float, single, n, m] desc: "matrix of choices" outputs: - name: output @@ -3253,7 +3253,7 @@ select: type: [float, n] desc: "source vector" - name: indices - type: [int, m] + type: [int, single, m] desc: "indices to gather" backward_arg: True outputs: diff --git a/madspace/src/cpu/kernel_definitions.hpp b/madspace/src/cpu/kernel_definitions.hpp index 652b8b5631..79916ef220 100644 --- a/madspace/src/cpu/kernel_definitions.hpp +++ b/madspace/src/cpu/kernel_definitions.hpp @@ -19,6 +19,10 @@ struct CpuTypes { using FOut = TensorView; template using IOut = TensorView; + template + using FSIn = const TensorView; + template + using ISIn = const TensorView; using FVal = double; using IVal = me_int_t; using BVal = bool; @@ -67,6 +71,10 @@ struct SimdTypes { using FOut = cpu::VectorizedTensorView; template using IOut = cpu::VectorizedTensorView; + template + using FSIn = const cpu::SingleVectorizedTensorView; + template + using ISIn = const cpu::SingleVectorizedTensorView; using FVal = FVec; using IVal = IVec; using BVal = BVec; diff --git a/madspace/src/cpu/runtime.cpp b/madspace/src/cpu/runtime.cpp index bd3860a201..91f40dc00d 100644 --- a/madspace/src/cpu/runtime.cpp +++ b/madspace/src/cpu/runtime.cpp @@ -46,6 +46,11 @@ using namespace madspace; using namespace madspace::cpu; using namespace madspace::kernels; +#ifdef USE_SIMD +// the single kernel parameters match the single inputs of instruction_set.yaml +#include "kernel_checks.inc" +#endif + namespace { template diff --git a/madspace/src/cpu/simd_arm.hpp b/madspace/src/cpu/simd_arm.hpp index 93cb88625f..279ba6c138 100644 --- a/madspace/src/cpu/simd_arm.hpp +++ b/madspace/src/cpu/simd_arm.hpp @@ -62,19 +62,15 @@ inline IVec vgather( return ret; } -inline FVec vload(double* base_ptr, std::size_t stride) { - FVec ret; - ret = vsetq_lane_f64(base_ptr[0], ret, 0); - ret = vsetq_lane_f64(base_ptr[stride], ret, 1); - return ret; -} +// contiguous loads and stores of the lanes of one batch element +inline FVec vload(double* base_ptr) { return vld1q_f64(base_ptr); } -inline IVec vload(int* base_ptr, std::size_t stride) { - IVec ret; - ret = vsetq_lane_s64(base_ptr[0], ret, 0); - ret = vsetq_lane_s64(base_ptr[stride], ret, 1); - return ret; -} +inline IVec vload(int* base_ptr) { return vmovl_s32(vld1_s32(base_ptr)); } + +// the same value in every lane +inline FVec vload_broadcast(double* base_ptr) { return vld1q_dup_f64(base_ptr); } + +inline IVec vload_broadcast(int* base_ptr) { return vdupq_n_s64(*base_ptr); } inline void vscatter( double* base_ptr, @@ -102,15 +98,9 @@ inline void vscatter( vgetq_lane_s64(values, 1); } -inline void vstore(double* base_ptr, std::size_t stride, FVec values) { - base_ptr[0] = vgetq_lane_f64(values, 0); - base_ptr[stride] = vgetq_lane_f64(values, 1); -} +inline void vstore(double* base_ptr, FVec values) { vst1q_f64(base_ptr, values); } -inline void vstore(int* base_ptr, std::size_t stride, IVec values) { - base_ptr[0] = vgetq_lane_s64(values, 0); - base_ptr[stride] = vgetq_lane_s64(values, 1); -} +inline void vstore(int* base_ptr, IVec values) { vst1_s32(base_ptr, vmovn_s64(values)); } inline FVec where(BVec arg1, FVec arg2, FVec arg3) { return vbslq_f64(arg1, arg2, arg3); diff --git a/madspace/src/cpu/simd_x86_256.hpp b/madspace/src/cpu/simd_x86_256.hpp index 099bb18859..eb770f1158 100644 --- a/madspace/src/cpu/simd_x86_256.hpp +++ b/madspace/src/cpu/simd_x86_256.hpp @@ -75,14 +75,18 @@ inline IVec vgather( )); } -inline FVec vload(double* base_ptr, std::size_t stride) { - return _mm256_i32gather_pd(base_ptr, stride_seq(stride), 8); -} +// contiguous loads and stores of the lanes of one batch element +inline FVec vload(double* base_ptr) { return _mm256_loadu_pd(base_ptr); } -inline IVec vload(int* base_ptr, std::size_t stride) { - return _mm256_cvtepi32_epi64(_mm_i32gather_epi32(base_ptr, stride_seq(stride), 4)); +inline IVec vload(int* base_ptr) { + return _mm256_cvtepi32_epi64(_mm_loadu_si128(reinterpret_cast<__m128i*>(base_ptr))); } +// the same value in every lane +inline FVec vload_broadcast(double* base_ptr) { return _mm256_broadcast_sd(base_ptr); } + +inline IVec vload_broadcast(int* base_ptr) { return _mm256_set1_epi64x(*base_ptr); } + inline void vscatter( double* base_ptr, std::size_t batch_stride, @@ -126,23 +130,10 @@ inline void vscatter( } } -inline void vstore(double* base_ptr, std::size_t stride, FVec values) { - double values_buf[4]; - _mm256_storeu_pd(values_buf, values); - for (int i = 0; i < 4; ++i) { - base_ptr[i * stride] = values_buf[i]; - } -} +inline void vstore(double* base_ptr, FVec values) { _mm256_storeu_pd(base_ptr, values); } -inline void vstore(int* base_ptr, std::size_t stride, IVec values) { - union { - int scalar[4]; - __m128i vec; - } values_buf; - _mm_store_si128(&values_buf.vec, truncate_i64_to_i32(values)); - for (int i = 0; i < 4; ++i) { - base_ptr[i * stride] = values_buf.scalar[i]; - } +inline void vstore(int* base_ptr, IVec values) { + _mm_storeu_si128(reinterpret_cast<__m128i*>(base_ptr), truncate_i64_to_i32(values)); } inline FVec where(BVec arg1, FVec arg2, FVec arg3) { diff --git a/madspace/src/cpu/simd_x86_256y.hpp b/madspace/src/cpu/simd_x86_256y.hpp index a038212d7d..bdf193a1c6 100644 --- a/madspace/src/cpu/simd_x86_256y.hpp +++ b/madspace/src/cpu/simd_x86_256y.hpp @@ -69,16 +69,18 @@ inline IVec vgather( )); } -inline FVec vload(double* base_ptr, std::size_t stride) { - return _mm256_i64gather_pd(base_ptr, stride_seq(stride), 8); -} +// contiguous loads and stores of the lanes of one batch element +inline FVec vload(double* base_ptr) { return _mm256_loadu_pd(base_ptr); } -inline IVec vload(int* base_ptr, std::size_t stride) { - return _mm256_cvtepi32_epi64( - _mm256_i64gather_epi32(base_ptr, stride_seq(stride), 4) - ); +inline IVec vload(int* base_ptr) { + return _mm256_cvtepi32_epi64(_mm_loadu_si128(reinterpret_cast<__m128i*>(base_ptr))); } +// the same value in every lane +inline FVec vload_broadcast(double* base_ptr) { return _mm256_broadcast_sd(base_ptr); } + +inline IVec vload_broadcast(int* base_ptr) { return _mm256_set1_epi64x(*base_ptr); } + inline void vscatter( double* base_ptr, std::size_t batch_stride, @@ -106,13 +108,11 @@ inline void vscatter( ); } -inline void vstore(double* base_ptr, std::size_t stride, FVec values) { - _mm256_i64scatter_pd(base_ptr, stride_seq(stride), values, 8); -} +inline void vstore(double* base_ptr, FVec values) { _mm256_storeu_pd(base_ptr, values); } -inline void vstore(int* base_ptr, std::size_t stride, IVec values) { - _mm256_i64scatter_epi32( - base_ptr, stride_seq(stride), _mm256_cvtepi64_epi32(values), 4 +inline void vstore(int* base_ptr, IVec values) { + _mm_storeu_si128( + reinterpret_cast<__m128i*>(base_ptr), _mm256_cvtepi64_epi32(values) ); } diff --git a/madspace/src/cpu/simd_x86_512.hpp b/madspace/src/cpu/simd_x86_512.hpp index 74808d83c1..c143fb69a9 100644 --- a/madspace/src/cpu/simd_x86_512.hpp +++ b/madspace/src/cpu/simd_x86_512.hpp @@ -70,16 +70,20 @@ inline IVec vgather( )); } -inline FVec vload(double* base_ptr, std::size_t stride) { - return _mm512_i32gather_pd(stride_seq(stride), base_ptr, 8); -} +// contiguous loads and stores of the lanes of one batch element +inline FVec vload(double* base_ptr) { return _mm512_loadu_pd(base_ptr); } -inline IVec vload(int* base_ptr, std::size_t stride) { +inline IVec vload(int* base_ptr) { return _mm512_cvtepi32_epi64( - _mm256_i32gather_epi32(base_ptr, stride_seq(stride), 4) + _mm256_loadu_si256(reinterpret_cast<__m256i*>(base_ptr)) ); } +// the same value in every lane +inline FVec vload_broadcast(double* base_ptr) { return _mm512_set1_pd(*base_ptr); } + +inline IVec vload_broadcast(int* base_ptr) { return _mm512_set1_epi64(*base_ptr); } + inline void vscatter( double* base_ptr, std::size_t batch_stride, @@ -115,19 +119,12 @@ inline void vscatter( } } -inline void vstore(double* base_ptr, std::size_t stride, FVec values) { - _mm512_i32scatter_pd(base_ptr, stride_seq(stride), values, 8); -} +inline void vstore(double* base_ptr, FVec values) { _mm512_storeu_pd(base_ptr, values); } -inline void vstore(int* base_ptr, std::size_t stride, IVec values) { - union { - int scalar[8]; - __m256i vec; - } values_buf; - _mm256_store_si256(&values_buf.vec, _mm512_cvtepi64_epi32(values)); - for (int i = 0; i < 8; ++i) { - base_ptr[i * stride] = values_buf.scalar[i]; - } +inline void vstore(int* base_ptr, IVec values) { + _mm256_storeu_si256( + reinterpret_cast<__m256i*>(base_ptr), _mm512_cvtepi64_epi32(values) + ); } inline FVec where(BVec arg1, FVec arg2, FVec arg3) { diff --git a/madspace/src/cpu/tensor.hpp b/madspace/src/cpu/tensor.hpp index a787617455..0eeba4ba86 100644 --- a/madspace/src/cpu/tensor.hpp +++ b/madspace/src/cpu/tensor.hpp @@ -5,62 +5,56 @@ #include "madspace/driver/thread_pool.hpp" #include "simd.hpp" +#include +#include +#include + namespace madspace { namespace cpu { +// View onto a tensor for the vectorized kernels. The simd_vec_size lanes of a +// vector are consecutive batch entries, so they are always contiguous in memory +// (batch stride one) and are read and written with plain vector loads and stores. +// Moving to the next vector advances the data pointer by vec_step elements: +// simd_vec_size for batched tensors, zero for literals that are stored padded to +// simd_vec_size identical lanes. Tensors that do not have this layout are handled +// by the scalar kernels instead (see tensor_foreach_impl). template class VectorizedTensorView { public: using VType = V; using DType = T; static const int dim = _dim; + static constexpr bool is_single = false; - VectorizedTensorView(const TensorView& view) : + VectorizedTensorView(const TensorView& view, std::size_t vec_step) : _data(view.data()), _stride(view.stride()), _shape(view.shape()), - _batch_stride(view.stride()[0]) {} + _vec_step(vec_step) {} VectorizedTensorView( - T* data, std::size_t* stride, std::size_t* shape, std::size_t batch_stride + T* data, std::size_t* stride, std::size_t* shape, std::size_t vec_step ) : - _data(data), _stride(stride), _shape(shape), _batch_stride(batch_stride) {} + _data(data), _stride(stride), _shape(shape), _vec_step(vec_step) {} VectorizedTensorView(V& value) : _data(reinterpret_cast(&value)), _stride(nullptr), _shape(nullptr), - _batch_stride(0) {} + _vec_step(0) {} const VectorizedTensorView operator[](std::size_t index) const requires(_dim != 0) { - if constexpr (is_batch) { - return { - &_data[index * _stride[0] * simd_vec_size], - &_stride[1], - &_shape[1], - _batch_stride - }; - } else { - return {&_data[index * _stride[0]], &_stride[1], &_shape[1], _batch_stride}; - } + return {&_data[offset(index)], &_stride[1], &_shape[1], _vec_step}; } VectorizedTensorView operator[](std::size_t index) requires(_dim != 0) { - if constexpr (is_batch) { - return { - &_data[index * _stride[0] * simd_vec_size], - &_stride[1], - &_shape[1], - _batch_stride - }; - } else { - return {&_data[index * _stride[0]], &_stride[1], &_shape[1], _batch_stride}; - } + return {&_data[offset(index)], &_stride[1], &_shape[1], _vec_step}; } template @@ -68,16 +62,11 @@ class VectorizedTensorView { get(std::size_t first_index, I... index) const requires(_dim >= sizeof...(I) + 1) { - T* ptr = _data; - if constexpr (is_batch) { - ptr = &ptr[first_index * _stride[0] * simd_vec_size]; - } else { - ptr = &ptr[first_index * _stride[0]]; - } + T* ptr = &_data[offset(first_index)]; int i = 1; ((ptr = &ptr[index * _stride[i++]]), ...); return { - ptr, &_stride[sizeof...(I) + 1], &_shape[sizeof...(I) + 1], _batch_stride + ptr, &_stride[sizeof...(I) + 1], &_shape[sizeof...(I) + 1], _vec_step }; } @@ -86,37 +75,32 @@ class VectorizedTensorView { get(std::size_t first_index, I... index) requires(_dim >= sizeof...(I) + 1) { - T* ptr = _data; - if constexpr (is_batch) { - ptr = &ptr[first_index * _stride[0] * simd_vec_size]; - } else { - ptr = &ptr[first_index * _stride[0]]; - } + T* ptr = &_data[offset(first_index)]; int i = 1; ((ptr = &ptr[index * _stride[i++]]), ...); return { - ptr, &_stride[sizeof...(I) + 1], &_shape[sizeof...(I) + 1], _batch_stride + ptr, &_stride[sizeof...(I) + 1], &_shape[sizeof...(I) + 1], _vec_step }; } operator V() const requires(_dim == 0) { - return vload(_data, _batch_stride); + return vload(_data); } V operator=(V value) requires(_dim == 0) { - vstore(_data, _batch_stride, value); + vstore(_data, value); return value; } V operator+=(V value) requires(_dim == 0) { - V new_value = vload(_data, _batch_stride) + value; - vstore(_data, _batch_stride, new_value); + V new_value = vload(_data) + value; + vstore(_data, new_value); return new_value; } @@ -135,22 +119,107 @@ class VectorizedTensorView { V gather(IVec indices) const requires(_dim == 1) { - return vgather(_data, _batch_stride, _stride[0], indices); + return vgather(_data, 1, _stride[0], indices); } template void scatter_add(IVec indices, V values) requires(_dim == 1) { - V old_values = vgather(_data, _batch_stride, _stride[0], indices); - vscatter(_data, _batch_stride, _stride[0], indices, old_values + values); + V old_values = vgather(_data, 1, _stride[0], indices); + vscatter(_data, 1, _stride[0], indices, old_values + values); + } + +private: + std::size_t offset(std::size_t index) const { + if constexpr (is_batch) { + return index * _vec_step; + } else { + return index * _stride[0]; + } + } + + T* _data; + std::size_t* _stride; + std::size_t* _shape; + std::size_t _vec_step; +}; + +// View onto an input marked as single in the instruction set, i.e. without a +// batch dimension. The batch index is ignored, and every lane gets the same +// value through a broadcast load. +template +class SingleVectorizedTensorView { +public: + using VType = V; + using DType = T; + static const int dim = _dim; + static constexpr bool is_single = true; + + SingleVectorizedTensorView(const TensorView& view) : + _data(view.data()), _stride(view.stride()), _shape(view.shape()) {} + + SingleVectorizedTensorView(T* data, std::size_t* stride, std::size_t* shape) : + _data(data), _stride(stride), _shape(shape) {} + + const SingleVectorizedTensorView + operator[](std::size_t index) const + requires(_dim != 0) + { + return {&_data[offset(index)], &_stride[1], &_shape[1]}; + } + + template + const SingleVectorizedTensorView + get(std::size_t first_index, I... index) const + requires(_dim >= sizeof...(I) + 1) + { + T* ptr = &_data[offset(first_index)]; + int i = 1; + ((ptr = &ptr[index * _stride[i++]]), ...); + return {ptr, &_stride[sizeof...(I) + 1], &_shape[sizeof...(I) + 1]}; + } + + operator V() const + requires(_dim == 0) + { + return vload_broadcast(_data); + } + + // the value as a plain scalar, e.g. for loop bounds or indices + T value() const + requires(_dim == 0) + { + return *_data; + } + + std::size_t size(std::size_t index = 0) const { return _shape[index]; } + + template + V gather(IVec indices) const + requires(_dim == 1) + { + return vgather(_data, 0, _stride[0], indices); } private: + std::size_t offset(std::size_t index) const { + if constexpr (is_batch) { + return 0; + } else { + return index * _stride[0]; + } + } + T* _data; std::size_t* _stride; std::size_t* _shape; - std::size_t _batch_stride; +}; + +// whether a kernel parameter type is a view onto a single (non-batched) input +template +inline constexpr bool is_single_param = requires { + requires std::remove_cvref_t::is_single; }; template @@ -207,25 +276,57 @@ struct get_views { }; // return the tuple of VectorizedTensorViews where the type is extracted -// from the signature of F +// from the signature of F. vec_steps holds the step between vectors for every +// argument (simd_vec_size, or zero for padded literals) template struct get_vectorized_views; template struct get_vectorized_views { template - auto operator()(TArg&... args) { - return std::make_tuple([&]() { - if constexpr (TArg::is_scalar_view) { - return ScalarView(TParam(typename TArg::DType(args))); - } else { - return VectorizedTensorView< - typename TParam::VType, - typename TParam::DType, - TParam::dim + dims, - true>(args); - } - }()...); + auto operator()( + const std::array& vec_steps, TArg&... args + ) { + return make(vec_steps, std::index_sequence_for(), args...); } + +private: + template + auto make( + const std::array& vec_steps, + std::index_sequence, + TArg&... args + ) { + return std::make_tuple(make_view(vec_steps[I], args)...); + } + + template + static auto make_view(std::size_t vec_step, TArg& arg) { + if constexpr (TArg::is_scalar_view) { + return ScalarView(P(typename TArg::DType(arg))); + } else if constexpr (is_single_param

) { + return SingleVectorizedTensorView< + typename P::VType, + typename P::DType, + P::dim + dims, + true>(arg); + } else { + return VectorizedTensorView< + typename P::VType, + typename P::DType, + P::dim + dims, + true>(arg, vec_step); + } + } +}; + +// whether each parameter of F is a single (non-batched) input +template +struct single_params; +template +struct single_params { + static constexpr std::array value{ + is_single_param... + }; }; template @@ -334,18 +435,56 @@ inline void tensor_foreach_impl( ) ); + // The vectorized kernel loads and stores the lanes of a vector with plain + // contiguous memory accesses. It is only used if every batched argument has + // batch stride one (or is a literal padded to simd_vec_size lanes); single + // inputs are read with broadcast loads and can have any layout. Otherwise + // the whole batch is handled by the scalar kernel. + constexpr std::size_t n_args = n_in + n_out + sizeof...(S); + std::array vec_steps{}; + bool vectorize = !broadcast_output; + if constexpr (!std::is_same_v) { + using Single = single_params; + [&](std::index_sequence) { + ( + [&] { + using View = std::tuple_element_t< + I, + std::remove_reference_t>; + if constexpr (!ScalarType && !Single::value[I]) { + auto& view = std::get(flat_views); + if (view.shape[0] == batch_size && view.stride[0] == 1) { + vec_steps[I] = simd_vec_size; + } else if (view.simd_padded) { + vec_steps[I] = 0; + } else { + vectorize = false; + } + } + }(), + ... + ); + }(std::make_index_sequence()); + } + device.foreach ( batch_size, [flat_views, - broadcast_output, + vectorize, + vec_steps, scalar_args...](std::size_t count, std::size_t offset) mutable { auto views = std::apply(get_views(), flat_views); std::size_t scalar_offset = offset; if constexpr (!std:: is_same_v) { - if (!broadcast_output) { + if (vectorize) { auto vectorized_views = std::apply( - get_vectorized_views(), views + [&vec_steps](auto&... args) { + return get_vectorized_views()( + vec_steps, args... + ); + }, + views ); std::size_t vec_count = count / simd_vec_size; std::size_t vec_offset = offset / simd_vec_size; diff --git a/madspace/src/gpu/kernel_definitions.cuh b/madspace/src/gpu/kernel_definitions.cuh index 00d292f5a7..a4f8c1c03b 100644 --- a/madspace/src/gpu/kernel_definitions.cuh +++ b/madspace/src/gpu/kernel_definitions.cuh @@ -17,6 +17,11 @@ struct GpuTypes { using FOut = gpu::GpuTensorView; template using IOut = gpu::GpuTensorView; + // single inputs are broadcast through a zero batch stride on the GPU + template + using FSIn = const gpu::GpuTensorView; + template + using ISIn = const gpu::GpuTensorView; using FVal = double; using IVal = me_int_t; using BVal = bool; diff --git a/madspace/src/kernels/chili.hpp b/madspace/src/kernels/chili.hpp index 61baa50630..1b054f030a 100644 --- a/madspace/src/kernels/chili.hpp +++ b/madspace/src/kernels/chili.hpp @@ -11,8 +11,8 @@ KERNELSPEC void kernel_chili_forward( FIn r, FIn e_cm, FIn m_out, - FIn pt_min, - FIn y_max, + FSIn pt_min, + FSIn y_max, FOut p_ext, FOut det ) { @@ -134,8 +134,8 @@ KERNELSPEC void kernel_chili_inverse( FIn p_ext, FIn e_cm, FIn m_out, - FIn pt_min, - FIn y_max, + FSIn pt_min, + FSIn y_max, FOut r, FOut det ) { diff --git a/madspace/src/kernels/cuts.hpp b/madspace/src/kernels/cuts.hpp index 178020d0b9..8ab51510f8 100644 --- a/madspace/src/kernels/cuts.hpp +++ b/madspace/src/kernels/cuts.hpp @@ -25,13 +25,13 @@ KERNELSPEC void kernel_cut_unphysical( template KERNELSPEC void -kernel_cut_one(FIn obs, FIn min, FIn max, FOut w) { +kernel_cut_one(FIn obs, FSIn min, FSIn max, FOut w) { w = where((obs < min) | (obs > max), FVal(0.), 1.); } template KERNELSPEC void -kernel_cut_all(FIn obs, FIn min, FIn max, FOut w) { +kernel_cut_all(FIn obs, FSIn min, FSIn max, FOut w) { FVal cut = 1., min_val = min, max_val = max; for (std::size_t i = 0; i < obs.size(); ++i) { FVal obs_i = obs[i]; @@ -42,7 +42,7 @@ kernel_cut_all(FIn obs, FIn min, FIn max, FOut w) { template KERNELSPEC void -kernel_cut_any(FIn obs, FIn min, FIn max, FOut w) { +kernel_cut_any(FIn obs, FSIn min, FSIn max, FOut w) { FVal cut = 0., min_val = min, max_val = max; for (std::size_t i = 0; i < obs.size(); ++i) { FVal obs_i = obs[i]; diff --git a/madspace/src/kernels/definitions.hpp b/madspace/src/kernels/definitions.hpp index c7fafb893c..cafd2d94b7 100644 --- a/madspace/src/kernels/definitions.hpp +++ b/madspace/src/kernels/definitions.hpp @@ -20,6 +20,11 @@ template using IIn = typename T::template IIn; template using BIn = typename T::template BIn; +// inputs marked as single in the instruction set (no batch dimension) +template +using FSIn = typename T::template FSIn; +template +using ISIn = typename T::template ISIn; template using FOut = typename T::template FOut; template diff --git a/madspace/src/kernels/discrete.hpp b/madspace/src/kernels/discrete.hpp index 5f099cfa04..21a737d08f 100644 --- a/madspace/src/kernels/discrete.hpp +++ b/madspace/src/kernels/discrete.hpp @@ -18,7 +18,7 @@ KERNELSPEC void kernel_sample_discrete( template KERNELSPEC void kernel_sample_discrete_inverse( - IIn index, IIn option_count, FOut r, FOut det + IIn index, ISIn option_count, FOut r, FOut det ) { IVal opt_count_i(option_count), index_i(index); FVal opt_count_f(opt_count_i), index_f(index_i); diff --git a/madspace/src/kernels/invariants.hpp b/madspace/src/kernels/invariants.hpp index ddf6994764..f0ece87cb6 100644 --- a/madspace/src/kernels/invariants.hpp +++ b/madspace/src/kernels/invariants.hpp @@ -24,8 +24,8 @@ KERNELSPEC void kernel_uniform_invariant_inverse( template KERNELSPEC void kernel_breit_wigner_invariant( FIn r, - FIn mass, - FIn width, + FSIn mass, + FSIn width, FIn s_min, FIn s_max, FOut s, @@ -46,8 +46,8 @@ KERNELSPEC void kernel_breit_wigner_invariant( template KERNELSPEC void kernel_breit_wigner_invariant_inverse( FIn s, - FIn mass, - FIn width, + FSIn mass, + FSIn width, FIn s_min, FIn s_max, FOut r, @@ -71,9 +71,9 @@ KERNELSPEC void kernel_breit_wigner_invariant_inverse( // piling its points up at the pole. template KERNELSPEC void flat_window_segments( - FIn mass, - FIn width, - FIn window, + FVal mass, + FVal width, + FVal window, FIn s_min, FIn s_max, FVal& m2, @@ -109,9 +109,9 @@ KERNELSPEC void flat_window_segments( template KERNELSPEC void kernel_flat_window_breit_wigner_invariant( FIn r, - FIn mass, - FIn width, - FIn window, + FSIn mass, + FSIn width, + FSIn window, FIn s_min, FIn s_max, FOut s, @@ -137,9 +137,9 @@ KERNELSPEC void kernel_flat_window_breit_wigner_invariant( template KERNELSPEC void kernel_flat_window_breit_wigner_invariant_inverse( FIn s, - FIn mass, - FIn width, - FIn window, + FSIn mass, + FSIn width, + FSIn window, FIn s_min, FIn s_max, FOut r, @@ -165,7 +165,7 @@ KERNELSPEC void kernel_flat_window_breit_wigner_invariant_inverse( template KERNELSPEC void kernel_stable_invariant( FIn r, - FIn mass, + FSIn mass, FIn s_min, FIn s_max, FOut s, @@ -182,7 +182,7 @@ KERNELSPEC void kernel_stable_invariant( template KERNELSPEC void kernel_stable_invariant_inverse( FIn s, - FIn mass, + FSIn mass, FIn s_min, FIn s_max, FOut r, @@ -201,8 +201,8 @@ KERNELSPEC void kernel_stable_invariant_inverse( template KERNELSPEC void kernel_stable_invariant_nu( FIn r, - FIn mass, - FIn nu, + FSIn mass, + FSIn nu, FIn s_min, FIn s_max, FOut s, @@ -223,8 +223,8 @@ KERNELSPEC void kernel_stable_invariant_nu( template KERNELSPEC void kernel_stable_invariant_nu_inverse( FIn s, - FIn mass, - FIn nu, + FSIn mass, + FSIn nu, FIn s_min, FIn s_max, FOut r, diff --git a/madspace/src/kernels/kinematics.hpp b/madspace/src/kernels/kinematics.hpp index 2dfa9f020c..b2977801a0 100644 --- a/madspace/src/kernels/kinematics.hpp +++ b/madspace/src/kernels/kinematics.hpp @@ -434,7 +434,7 @@ KERNELSPEC Pair, FourMom> boost_two_body( template KERNELSPEC void boost_beam( FIn q, - FIn masses, + FSIn masses, FVal x1, FVal x2, FVal sign, @@ -598,14 +598,14 @@ KERNELSPEC Quartuplet, FVal, FVal, FVal> phi_m1_m2_from_p1com( template KERNELSPEC void kernel_boost_beam( - FIn p1, FIn masses, FIn x1, FIn x2, FOut p_out + FIn p1, FSIn masses, FIn x1, FIn x2, FOut p_out ) { boost_beam(p1, masses, x1, x2, 1.0, p_out); } template KERNELSPEC void kernel_boost_beam_inverse( - FIn p1, FIn masses, FIn x1, FIn x2, FOut p_out + FIn p1, FSIn masses, FIn x1, FIn x2, FOut p_out ) { boost_beam(p1, masses, x1, x2, -1.0, p_out); } @@ -627,7 +627,7 @@ template KERNELSPEC void kernel_r_to_x1x2( FIn r, FIn s_hat, - FIn s_lab, + FSIn s_lab, FOut x1, FOut x2, FOut det @@ -640,7 +640,7 @@ KERNELSPEC void kernel_r_to_x1x2( template KERNELSPEC void kernel_x1x2_to_r( - FIn x1, FIn x2, FIn s_lab, FOut r, FOut det + FIn x1, FIn x2, FSIn s_lab, FOut r, FOut det ) { auto tau = x1 * x2; auto log_tau = log(tau); @@ -657,8 +657,8 @@ template KERNELSPEC void kernel_r_to_x1x2_window( FIn r, FIn s_hat, - FIn s_lab, - FIn y_max, + FSIn s_lab, + FSIn y_max, FOut x1, FOut x2, FOut det @@ -676,8 +676,8 @@ template KERNELSPEC void kernel_x1x2_to_r_window( FIn x1, FIn x2, - FIn s_lab, - FIn y_max, + FSIn s_lab, + FSIn y_max, FOut r, FOut det ) { @@ -697,7 +697,7 @@ KERNELSPEC void kernel_diff_cross_section( FIn pdf1, FIn pdf2, FIn matrix_element, - FIn e_cm2, + FSIn e_cm2, FOut result ) { result = INV_GEV2_TO_PB * matrix_element * pdf1 * pdf2 / @@ -890,9 +890,9 @@ KERNELSPEC void kernel_t_rapidity_clamp( FIn m2, FIn x1, FIn x2, - FIn beam_sign, - FIn y_max1, - FIn y_max2, + FSIn beam_sign, + FSIn y_max1, + FSIn y_max2, FOut t_min_out, FOut t_max_out ) { @@ -960,9 +960,9 @@ KERNELSPEC void kernel_t1_rapidity_clamp_doublet( FIn etmin_2, FIn x1, FIn x2, - FIn beam_sign, - FIn y_max1, - FIn y_max2, + FSIn beam_sign, + FSIn y_max1, + FSIn y_max2, FOut t1_min_out, FOut t1_max_out ) { @@ -1046,9 +1046,9 @@ KERNELSPEC void kernel_t2_rapidity_clamp_doublet( FIn t1_abs, FIn x1, FIn x2, - FIn beam_sign, - FIn y_max1, - FIn y_max2, + FSIn beam_sign, + FSIn y_max1, + FSIn y_max2, FOut t2_min_out, FOut t2_max_out ) { @@ -1128,7 +1128,7 @@ KERNELSPEC void kernel_t_inv_value_and_min_max_cut( template KERNELSPEC void kernel_invariants_from_momenta( - FIn p_ext, FIn factors, FOut invariants + FIn p_ext, FSIn factors, FOut invariants ) { for (std::size_t i = 0; i < invariants.size(); ++i) { auto factors_i = factors[i]; @@ -1286,7 +1286,7 @@ kernel_mirror_momenta(FIn p_ext, IIn mirror, FOut p_out) { template KERNELSPEC void -kernel_momenta_to_x1x2(FIn p_ext, FIn e_cm, FOut x1, FOut x2) { +kernel_momenta_to_x1x2(FIn p_ext, FSIn e_cm, FOut x1, FOut x2) { x1 = 2. * p_ext[0][0] / e_cm; x2 = 2. * p_ext[1][0] / e_cm; } diff --git a/madspace/src/kernels/multichannel.hpp b/madspace/src/kernels/multichannel.hpp index fd6c10e2f5..1c37f4a3c8 100644 --- a/madspace/src/kernels/multichannel.hpp +++ b/madspace/src/kernels/multichannel.hpp @@ -8,7 +8,7 @@ namespace kernels { template KERNELSPEC void kernel_collect_channel_weights( FIn amp2, - IIn channel_indices, + ISIn channel_indices, IIn channel_count, FOut channel_weights ) { @@ -47,9 +47,9 @@ KERNELSPEC void kernel_collect_channel_weights( template KERNELSPEC void kernel_sde2_channel_weights( FIn invariants, - FIn masses, - FIn widths, - IIn indices, + FSIn masses, + FSIn widths, + ISIn indices, FOut channel_weights ) { // TODO: MG has a special case here if tprid != 0. check what that is and if we need @@ -82,11 +82,11 @@ KERNELSPEC void kernel_sde2_channel_weights( template KERNELSPEC void kernel_subchannel_weights( FIn invariants, - FIn masses, - FIn widths, - IIn indices, - IIn on_shell, - IIn group_sizes, + FSIn masses, + FSIn widths, + ISIn indices, + ISIn on_shell, + ISIn group_sizes, FOut channel_weights ) { FVal channel_weights_norm(0.); @@ -127,8 +127,8 @@ template KERNELSPEC void kernel_apply_subchannel_weights( FIn channel_weights_in, FIn subchannel_weights, - IIn channel_indices, - IIn subchannel_indices, + ISIn channel_indices, + ISIn subchannel_indices, FOut channel_weights_out ) { for (std::size_t i = 0; i < channel_indices.size(); ++i) { diff --git a/madspace/src/kernels/threeparticle.hpp b/madspace/src/kernels/threeparticle.hpp index 72a3e202d1..bcf07ab86c 100644 --- a/madspace/src/kernels/threeparticle.hpp +++ b/madspace/src/kernels/threeparticle.hpp @@ -711,8 +711,8 @@ KERNELSPEC void kernel_s23_arcsine_sample( FIn s_max, FIn s_phys_min, FIn s_phys_max, - FIn power, - FIn mass, + FSIn power, + FSIn mass, FOut u, FOut u_c, FOut det @@ -757,8 +757,8 @@ KERNELSPEC void kernel_s23_arcsine_sample_inverse( FIn s_max, FIn s_phys_min, FIn s_phys_max, - FIn power, - FIn mass, + FSIn power, + FSIn mass, FOut r, FOut det ) { diff --git a/madspace/src/kernels/twoparticle.hpp b/madspace/src/kernels/twoparticle.hpp index d610e54def..66815f1ca1 100644 --- a/madspace/src/kernels/twoparticle.hpp +++ b/madspace/src/kernels/twoparticle.hpp @@ -125,9 +125,9 @@ KERNELSPEC void kernel_com_cos_theta_window( FIn m2, FIn x1, FIn x2, - FIn pt_min, - FIn y_max1, - FIn y_max2, + FSIn pt_min, + FSIn y_max1, + FSIn y_max2, FOut r_out, FOut det ) { @@ -145,9 +145,9 @@ KERNELSPEC void kernel_com_cos_theta_window_inverse( FIn m2, FIn x1, FIn x2, - FIn pt_min, - FIn y_max1, - FIn y_max2, + FSIn pt_min, + FSIn y_max1, + FSIn y_max2, FOut r_cos_theta, FOut det ) { diff --git a/madspace/src/kernels/util.hpp b/madspace/src/kernels/util.hpp index 762bff6d4e..5254ce0d7c 100644 --- a/madspace/src/kernels/util.hpp +++ b/madspace/src/kernels/util.hpp @@ -45,13 +45,13 @@ KERNELSPEC void backward_kernel_gather( template KERNELSPEC void -kernel_gather_int(IIn index, IIn choices, IOut output) { +kernel_gather_int(IIn index, ISIn choices, IOut output) { output = choices.gather(index); } template KERNELSPEC void -kernel_gather_vector(IIn index, FIn choices, FOut output) { +kernel_gather_vector(IIn index, FSIn choices, FOut output) { auto choice_i = choices[index]; for (std::size_t j = 0; j < output.size(); ++j) { output[j] = choice_i[j]; @@ -67,7 +67,7 @@ kernel_select_int(IIn input, IIn indices, IOut output) { } template -KERNELSPEC void kernel_select(FIn input, IIn indices, FOut output) { +KERNELSPEC void kernel_select(FIn input, ISIn indices, FOut output) { for (std::size_t i = 0; i < indices.size(); ++i) { output[i] = input.gather(indices[i]); } @@ -75,7 +75,7 @@ KERNELSPEC void kernel_select(FIn input, IIn indices, FOut out template KERNELSPEC void backward_kernel_select( - IIn indices, + ISIn indices, FIn output_grad, FOut input_grad, FOut indices_grad diff --git a/madspace/src/kernels/vegas.hpp b/madspace/src/kernels/vegas.hpp index 06bcf797cb..d03fc77080 100644 --- a/madspace/src/kernels/vegas.hpp +++ b/madspace/src/kernels/vegas.hpp @@ -7,7 +7,7 @@ namespace kernels { template KERNELSPEC void kernel_vegas_forward( - FIn input, FIn grid, FOut output, FOut det + FIn input, FSIn grid, FOut output, FOut det ) { FVal bin_count_f(grid.size() - 1); FVal bin_index_f(input * bin_count_f); @@ -22,7 +22,7 @@ KERNELSPEC void kernel_vegas_forward( template KERNELSPEC void kernel_vegas_inverse( - FIn input, FIn grid, FOut output, FOut det + FIn input, FSIn grid, FOut output, FOut det ) { std::size_t bin_count = grid.size() - 1; FVal bin_count_f(bin_count); From a5f4b5df7298d5b3cbe0a460fe605c408de033d0 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 7 Oct 2026 22:33:42 +0200 Subject: [PATCH 13/23] madspace: store literals padded to the SIMD vector size 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 --- madspace/include/madspace/driver/tensor.hpp | 27 +++++++- madspace/src/cpu/runtime.cpp | 7 +- madspace/src/driver/tensor.cpp | 75 +++++++++++++++++++++ 3 files changed, 106 insertions(+), 3 deletions(-) diff --git a/madspace/include/madspace/driver/tensor.hpp b/madspace/include/madspace/driver/tensor.hpp index c2e78bdd1a..3ab1ca54eb 100644 --- a/madspace/include/madspace/driver/tensor.hpp +++ b/madspace/include/madspace/driver/tensor.hpp @@ -529,7 +529,7 @@ class Tensor { check_impl(); T* data = static_cast(impl->data); if (flatten_count <= 1) { - return {data, impl->stride, impl->shape}; + return {data, impl->stride, impl->shape, impl->simd_padded}; } if (flatten_count > impl->contiguous_dims) { throw std::invalid_argument("can only flatten contiguous dimensions"); @@ -768,6 +768,30 @@ class Tensor { return impl->ref_count.load() == 1; } + /// Maximum number of stored elements (value elements times lanes) of a + /// literal for which @ref simd_padded_literal pads the storage. + static constexpr std::size_t max_simd_padded_size = 256; + + /** + * Literal constant for the CPU SIMD kernels. + * + * The tensor has batch size one and batch stride zero like any literal, + * but each element is stored `lanes` times in consecutive memory, one + * copy per SIMD vector lane, so that the vectorized kernels can read it + * with contiguous vector loads. Scalar code sees an ordinary broadcast + * tensor. Literals with more than @ref max_simd_padded_size stored + * elements, or `lanes` of one, are stored without padding. + */ + static Tensor + simd_padded_literal(const LiteralValue& value, std::size_t lanes, DevicePtr device); + + /// Whether this is a literal stored padded to the SIMD vector size, see + /// @ref simd_padded_literal. Views and copies of it are not padded. + bool is_simd_padded() const { + check_impl(); + return impl->simd_padded; + } + private: struct TensorImpl { DataType dtype; @@ -781,6 +805,7 @@ class Tensor { Sizes stride; std::size_t contiguous_dims; SizeVec batch_sizes; + bool simd_padded = false; template void reset(const D& device) { diff --git a/madspace/src/cpu/runtime.cpp b/madspace/src/cpu/runtime.cpp index 91f40dc00d..c6f146df40 100644 --- a/madspace/src/cpu/runtime.cpp +++ b/madspace/src/cpu/runtime.cpp @@ -1166,8 +1166,11 @@ CpuRuntime::CpuRuntime(const Function& function, ContextPtr context, bool concur std::visit( Overloaded{ [&](auto val) { - Tensor tensor(val, cpu_device()); - _locals_init[local.local_index] = tensor; + // padded so that the vectorized kernels can load it like a + // batched tensor, see Tensor::simd_padded_literal + _locals_init[local.local_index] = Tensor::simd_padded_literal( + LiteralValue(val), simd_vec_size, cpu_device() + ); }, [](std::monostate val) {} }, diff --git a/madspace/src/driver/tensor.cpp b/madspace/src/driver/tensor.cpp index da7a899743..6842b1686d 100644 --- a/madspace/src/driver/tensor.cpp +++ b/madspace/src/driver/tensor.cpp @@ -1,5 +1,7 @@ #include "madspace/driver/tensor.hpp" +#include + using namespace madspace; Tensor Tensor::select(std::size_t axis, std::size_t index) const { @@ -228,3 +230,76 @@ std::size_t Tensor::init_stride() { impl->contiguous_dims = impl->shape.size(); return stride_prod * dtype_size(); } + +Tensor Tensor::simd_padded_literal( + const LiteralValue& value, std::size_t lanes, DevicePtr device +) { + auto unpadded = [&]() -> Tensor { + return std::visit( + Overloaded{ + [&](const TensorValue& tensor_value) { + return Tensor(tensor_value, device); + }, + [&](std::monostate) -> Tensor { + throw std::invalid_argument("empty literal"); + }, + [&](auto scalar) { return Tensor(scalar, device); }, + }, + value + ); + }; + if (lanes <= 1) { + return unpadded(); + } + + // shape (with the batch dimension of size one) and flat column-major values + Sizes shape{1}; + std::variant, std::vector> values; + if (auto scalar = std::get_if(&value)) { + values = std::vector{*scalar}; + } else if (auto scalar = std::get_if(&value)) { + values = std::vector{*scalar}; + } else if (auto tensor_value = std::get_if(&value)) { + for (int size : std::get<0>(*tensor_value)) { + shape.push_back(size); + } + values = std::get<1>(*tensor_value); + } else { + return unpadded(); + } + std::size_t count = std::visit([](auto& vec) { return vec.size(); }, values); + if (count * lanes > max_simd_padded_size) { + return unpadded(); + } + + // element j of lane l is stored at l + lanes * j + auto [dtype, buffer] = std::visit( + [&](const std::vector& vec) { + auto dtype = + std::is_same_v ? DataType::dt_int : DataType::dt_float; + std::vector padded(count * lanes); + for (std::size_t j = 0; j < count; ++j) { + std::fill_n(padded.begin() + j * lanes, lanes, vec[j]); + } + Tensor buffer(dtype, {count * lanes}, device); + device->memcpy(buffer.data(), padded.data(), padded.size() * sizeof(T)); + return std::make_pair(dtype, buffer); + }, + values + ); + Sizes stride(shape.size()); + std::size_t stride_prod = lanes; + for (std::size_t i = 1; i < shape.size(); ++i) { + stride[i] = stride_prod; + stride_prod *= shape[i]; + } + Tensor ret(new TensorImpl{ + dtype, shape, device, buffer.data(), false, std::nullopt, buffer.impl, 1, stride, 1 + }); + ret.impl->simd_padded = true; + // non-negative integer scalars double as batch sizes, see the scalar constructor + if (auto scalar = std::get_if(&value); scalar && *scalar >= 0) { + ret.impl->batch_sizes.push_back(*scalar); + } + return ret; +} From 734a34d9c5c60e2e4e7e83b1bd21a06131dac5a9 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 7 Oct 2026 22:33:42 +0200 Subject: [PATCH 14/23] madspace: write unused gradients of literal and integer inputs to scratch 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 --- madspace/src/cpu/runtime.cpp | 15 +++++++++++++++ madspace/src/cpu/runtime.hpp | 4 ++++ madspace/src/kernels/operations.hpp | 18 +++++++++++++++++- 3 files changed, 36 insertions(+), 1 deletion(-) diff --git a/madspace/src/cpu/runtime.cpp b/madspace/src/cpu/runtime.cpp index c6f146df40..fb8761717c 100644 --- a/madspace/src/cpu/runtime.cpp +++ b/madspace/src/cpu/runtime.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include #include "../kernels/kernels.hpp" @@ -1041,6 +1042,10 @@ CpuRuntime::CpuRuntime(const Function& function, ContextPtr context, bool concur _requires_grad_init.resize(function.locals().size()); LastUseOfLocals last_use(function); std::vector local_sources(function.locals().size(), -1); + std::set global_locals; + for (auto& [name, value] : function.globals()) { + global_locals.insert(value.local_index); + } std::vector local_uses(function.locals().size()); SizeVec instr_index_map; @@ -1101,6 +1106,16 @@ CpuRuntime::CpuRuntime(const Function& function, ContextPtr context, bool concur dep_instrs_bw, dep_count_bw} ); + for (auto& in : instr.inputs) { + auto& local = function.locals().at(in.local_index); + bool is_literal = !std::holds_alternative(local.literal_value); + bool is_int_intermediate = local.type.dtype == DataType::dt_int && + static_cast(in.local_index) >= function.inputs().size() && + !global_locals.contains(in.local_index); + _instructions.back().discard_input_grad.push_back( + is_literal || is_int_intermediate + ); + } for (std::size_t local_index : last_use.local_indices(instr_index)) { std::size_t free_dep_count = 0; diff --git a/madspace/src/cpu/runtime.hpp b/madspace/src/cpu/runtime.hpp index 10ef87ccb7..4e65e4bb78 100644 --- a/madspace/src/cpu/runtime.hpp +++ b/madspace/src/cpu/runtime.hpp @@ -35,6 +35,10 @@ class CpuRuntime : public Runtime { std::size_t dependency_count_backward; DummyAllocHints output_alloc_hints; DummyAllocHints input_grad_alloc_hints; + // inputs whose gradient is never used (literals, integer locals); the + // backward kernels write it to batch-sized scratch memory instead of a + // broadcast gradient, see backward_batch_foreach + std::vector discard_input_grad; }; CpuRuntime(const Function& function, ContextPtr context, bool concurrent); diff --git a/madspace/src/kernels/operations.hpp b/madspace/src/kernels/operations.hpp index 728e150bd6..b129ef6e8f 100644 --- a/madspace/src/kernels/operations.hpp +++ b/madspace/src/kernels/operations.hpp @@ -68,7 +68,23 @@ void backward_batch_foreach( for (int i = 0; i < n_in; ++i) { auto input_index = instruction.input_indices[i]; auto& input_grad = local_grads[input_index]; - if (!input_grad) { + bool discard_grad = false; + if constexpr (requires { instruction.discard_input_grad; }) { + discard_grad = instruction.discard_input_grad[i]; + } + if (discard_grad) { + // The gradient of this input is never used (literal or integer). A + // batch-size-one gradient would be a broadcast output, which forces the + // scalar kernel on a single thread, so write it to batch-sized scratch + // memory instead. + auto& input = locals[input_index]; + if (!input_grad || input_grad.size(0) != batch_size) { + Sizes shape = input.shape(); + shape[0] = batch_size; + input_grad = + Tensor(input.dtype(), shape, device, AllocHint::temporary); + } + } else if (!input_grad) { auto& input = locals[input_index]; /*if (input.size(0) != batch_size && input.dtype() == DataType::dt_float) { throw std::runtime_error("backward not possible for broadcasting args"); From 1f89dcd4a4d1143ad97ae761d1ea1f1935cc81b5 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 7 Oct 2026 22:33:42 +0200 Subject: [PATCH 15/23] madspace: make external CPU runtime arguments batch-contiguous 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 --- madspace/src/cpu/runtime.cpp | 36 ++++++++++++++++++++++++++++++------ 1 file changed, 30 insertions(+), 6 deletions(-) diff --git a/madspace/src/cpu/runtime.cpp b/madspace/src/cpu/runtime.cpp index fb8761717c..4ba61ffe70 100644 --- a/madspace/src/cpu/runtime.cpp +++ b/madspace/src/cpu/runtime.cpp @@ -54,6 +54,26 @@ using namespace madspace::kernels; namespace { +// External tensors (e.g. row-major numpy arrays, or expanded gradients from +// torch) can have a batch stride other than one. The vectorized kernels only +// handle batch stride one and leave everything else to the scalar kernels, which +// round slightly differently, so copy them once to make the results independent +// of the memory layout of the arguments. +TensorVec batch_contiguous(const TensorVec& tensors) { + TensorVec ret; + ret.reserve(tensors.size()); + for (auto& tensor : tensors) { + if (tensor && tensor.dtype() != DataType::batch_sizes && + tensor.shape().size() > 0 && tensor.size(0) > 1 && + tensor.stride()[0] != 1) { + ret.push_back(tensor.contiguous()); + } else { + ret.push_back(tensor); + } + } + return ret; +} + template void op_matrix_element( const CpuRuntime::Instruction& instruction, TensorVec& locals, const D& device @@ -1199,21 +1219,24 @@ CpuRuntime::CpuRuntime(const Function& function, ContextPtr context, bool concur } TensorVec CpuRuntime::run(const TensorVec& inputs) { + auto contiguous_inputs = batch_contiguous(inputs); if (_concurrent && _context->thread_pool().thread_count() > 1) { - auto [outputs, locals, eval_grad] = run_concurrent(inputs, {}, false); + auto [outputs, locals, eval_grad] = + run_concurrent(contiguous_inputs, {}, false); return outputs; } else { - return run_single(inputs); + return run_single(contiguous_inputs); } } std::tuple> CpuRuntime::run_with_grad( const TensorVec& inputs, const std::vector& input_requires_grad ) { + auto contiguous_inputs = batch_contiguous(inputs); if (_concurrent && _context->thread_pool().thread_count() > 1) { - return run_concurrent(inputs, input_requires_grad, true); + return run_concurrent(contiguous_inputs, input_requires_grad, true); } else { - return run_with_grad_single(inputs, input_requires_grad); + return run_with_grad_single(contiguous_inputs, input_requires_grad); } } @@ -1223,13 +1246,14 @@ std::pair CpuRuntime::run_backward( const std::vector& eval_grad, bool return_contiguous_grads ) { + auto contiguous_output_grads = batch_contiguous(output_grads); if (_concurrent && _context->thread_pool().thread_count() > 1) { return run_backward_concurrent( - output_grads, stored_locals, eval_grad, return_contiguous_grads + contiguous_output_grads, stored_locals, eval_grad, return_contiguous_grads ); } else { return run_backward_single( - output_grads, stored_locals, eval_grad, return_contiguous_grads + contiguous_output_grads, stored_locals, eval_grad, return_contiguous_grads ); } } From f7113de44ff139872d27f3d08d36a0cbb133e74b Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 7 Oct 2026 22:33:42 +0200 Subject: [PATCH 16/23] madspace: test layout independence, single inputs, folding and broadcast gradients Co-Authored-By: Claude Opus 5.5 --- madspace/tests/test_simd_layout.py | 303 +++++++++++++++++++++++++++++ 1 file changed, 303 insertions(+) create mode 100644 madspace/tests/test_simd_layout.py diff --git a/madspace/tests/test_simd_layout.py b/madspace/tests/test_simd_layout.py new file mode 100644 index 0000000000..247ab81be8 --- /dev/null +++ b/madspace/tests/test_simd_layout.py @@ -0,0 +1,303 @@ +"""Results that must not depend on how the CPU runtime lays out tensors for the +SIMD kernels. + +The vectorized kernels load and store the lanes of a vector contiguously, so +they only run if every batched argument has batch stride one or is a literal +stored padded to the vector size; everything else (row-major inputs from numpy +or torch, broadcast tensors from fb.full, large literal arrays) is handled by +the scalar kernels. Inputs marked single in the instruction set have their own +kernel type. These tests run in every SIMD mode (tests/run_simd_modes.py) and +compare against numpy, so all of these paths have to agree. They also cover the +constant folding rules that keep broadcast tensors away from the kernels. +""" + +import numpy as np +import pytest +import torch +from pytest import approx + +import madspace as ms +from madspace.torch import FunctionModule + +torch.set_default_dtype(torch.float64) +rng = np.random.default_rng(4321) + +# not a multiple of any SIMD vector size, so the scalar tail runs as well +BATCH_SIZE = 1003 + + +def build(input_types, output_types, body): + fb = ms.FunctionBuilder(ms.NamedTypes(input_types), ms.NamedTypes(output_types)) + outputs = body(fb, *(fb.input(i) for i in range(len(input_types)))) + if not isinstance(outputs, (list, tuple)): + outputs = [outputs] + for i, output in enumerate(outputs): + fb.output(i, output) + return fb.function() + + +def call(function, *args): + outputs = ms.FunctionRuntime(function)(*args) + if not isinstance(outputs, (list, tuple)): + outputs = [outputs] + return [np.asarray(output) for output in outputs] + + +def instruction_names(function): + return [instr.instruction.name for instr in function.instructions] + + +# constant folding + + +def test_sub_with_constant_first_is_not_folded(): + func = build( + [("x", ms.batch_float)], + [("y", ms.batch_float)], + lambda fb, x: fb.sub(0.0, x), + ) + assert instruction_names(func) == ["sub"] + x = rng.normal(size=BATCH_SIZE) + (y,) = call(func, x) + assert y == approx(-x) + + +@pytest.mark.parametrize( + "op, const_first, const", + [ + ("add", False, 0.0), + ("add", True, 0.0), + ("sub", False, 0.0), + ("mul", False, 1.0), + ("mul", True, 1.0), + ], +) +def test_identities_are_folded(op, const_first, const): + def body(fb, x): + args = (const, x) if const_first else (x, const) + return getattr(fb, op)(*args) + + func = build([("x", ms.batch_float)], [("y", ms.batch_float)], body) + assert instruction_names(func) == [] + + +def test_stack_of_literals_is_folded(): + def body(fb, x): + stacked = fb.stack([1.0, 2.0, 3.0]) + assert stacked.literal_value == ([3], [1.0, 2.0, 3.0]) + return fb.mul(x, stacked) + + func = build([("x", ms.batch_float_array(3))], [("y", ms.batch_float_array(3))], body) + assert instruction_names(func) == ["mul"] + + +def test_single_flavor_integrand_has_no_constant_gathers(tmp_path): + """With a single flavor the flavor index is a constant tensor: the flavor + tables (remap, factors, mirror) must be resolved when building the graph + instead of gathering from them, and a weight that is one by construction + must not be multiplied.""" + # RunningCoupling only reads these two keys of an LHAPDF .info file + alphas_file = tmp_path / "alphas.info" + alphas_file.write_text( + "AlphaS_Qs: [1.0, 10.0, 100.0, 1000.0, 10000.0]\n" + "AlphaS_Vals: [0.3, 0.2, 0.12, 0.1, 0.08]\n" + ) + running_coupling = ms.RunningCoupling(ms.AlphaSGrid(str(alphas_file))) + mapping = ms.PhaseSpaceMapping([0.0] * 5, 1000.0, leptonic=True) + matrix_element = ms.MatrixElement( + 0xBADCAFE, + 5, + ms.Integrand.matrix_element_inputs, + ms.Integrand.matrix_element_outputs, + 1, + True, + ) + diff_xs = ms.DifferentialCrossSection( + matrix_element=matrix_element, + cm_energy=1000.0, + running_coupling=None, + energy_scale=ms.CachedScale(), + ) + for mirror in (False, True): + integrand = ms.Integrand( + mapping, + [diff_xs], + running_coupling=running_coupling, + energy_scale=ms.EnergyScale(5, 91.188), + channel_indices=[0], + flavor_remap=[0], + flavor_factors=[2.0], + flavor_mirror=[mirror], + ) + func = integrand.function() + names = instruction_names(func) + assert "gather" not in names + assert "gather_int" not in names + sources = {} + for instr in func.instructions: + for output in instr.outputs: + sources[output.local_index] = instr.instruction.name + for instr in func.instructions: + if instr.instruction.name == "mul": + assert all( + sources.get(arg.local_index) != "full" for arg in instr.inputs + ) + + +# memory layouts + + +@pytest.mark.parametrize("order", ["C", "F"]) +def test_row_major_and_column_major_inputs(order): + """Column-major (batch stride one) inputs run the vectorized kernels, + row-major ones the scalar kernels.""" + func = build( + [("x", ms.batch_float_array(3)), ("y", ms.batch_float_array(3))], + [("z", ms.batch_float_array(3))], + lambda fb, x, y: fb.add(fb.mul(x, y), fb.sqrt(fb.square(y))), + ) + x = np.asarray(rng.normal(size=(BATCH_SIZE, 3)), order=order) + y = np.asarray(rng.normal(size=(BATCH_SIZE, 3)), order=order) + (z,) = call(func, x, y) + assert z == approx(x * y + np.abs(y)) + + +@pytest.mark.parametrize("size", [3, 8, 100]) +def test_literal_arrays(size): + """Small literal arrays are stored padded to the vector size, large ones + (more than 256 stored elements) are not and fall back to the scalar + kernels.""" + values = list(rng.normal(size=size)) + func = build( + [("x", ms.batch_float_array(size))], + [("y", ms.batch_float_array(size))], + lambda fb, x: fb.add(fb.mul(x, fb.stack(values)), fb.stack([1.5] * size)), + ) + x = np.asfortranarray(rng.normal(size=(BATCH_SIZE, size))) + (y,) = call(func, x) + assert y == approx(x * np.array(values) + 1.5) + + +def test_broadcast_full_input(): + """fb.full creates a batch-sized tensor with batch stride zero.""" + + def body(fb, x): + full = fb.full([2.5, fb.batch_size([x])]) + return fb.mul(x, full) + + func = build([("x", ms.batch_float)], [("y", ms.batch_float)], body) + x = rng.normal(size=BATCH_SIZE) + (y,) = call(func, x) + assert y == approx(2.5 * x) + + +def test_cut_all_and_cut_any_with_scalar_bounds(): + func = build( + [("obs", ms.batch_float_array(4))], + [("all", ms.batch_float), ("any", ms.batch_float)], + lambda fb, obs: (fb.cut_all(obs, -2.0, 2.0), fb.cut_any(obs, -2.0, 2.0)), + ) + obs = np.asfortranarray(rng.uniform(-2.5, 2.5, size=(BATCH_SIZE, 4))) + w_all, w_any = call(func, obs) + inside = (obs >= -2.0) & (obs <= 2.0) + assert w_all == approx(np.all(inside, axis=1).astype(float)) + assert w_any == approx(np.any(inside, axis=1).astype(float)) + + +# single inputs + + +def test_single_input_rejects_batched_argument(): + fb = ms.FunctionBuilder( + ms.NamedTypes([("r", ms.batch_float)]), + ms.NamedTypes([("s", ms.batch_float), ("gs", ms.batch_float)]), + ) + r = fb.input(0) + with pytest.raises(ValueError, match="cannot have batch dimension"): + fb.breit_wigner_invariant(r, r, 1.0, 0.0, 100.0) + + +def test_single_inputs(): + """Breit-Wigner sampling with literal (single) mass and width, compared + between its forward and inverse direction.""" + mass, width, s_min, s_max = 91.0, 2.5, 50.0**2, 150.0**2 + + def body(fb, r): + s, gs = fb.breit_wigner_invariant(r, mass, width, s_min, s_max) + r_out, gs_inv = fb.breit_wigner_invariant_inverse(s, mass, width, s_min, s_max) + return r_out, gs, gs_inv + + func = build( + [("r", ms.batch_float)], + [("r_out", ms.batch_float), ("gs", ms.batch_float), ("gs_inv", ms.batch_float)], + body, + ) + r = rng.uniform(size=BATCH_SIZE) + r_out, gs, gs_inv = call(func, r) + assert r_out == approx(r) + assert gs * gs_inv == approx(np.ones(BATCH_SIZE)) + + +# gradients + + +def test_global_probability_gradient_is_summed_over_events(): + """The gradient of a broadcast (global) probability vector is the sum of the + contributions of all events.""" + ctx = ms.Context() + ctx.define_global("test_simd_layout_probs", ms.DataType.float, [4], True) + ctx.get_global("test_simd_layout_probs").torch()[:] = torch.tensor( + [[0.1, 0.4, 0.3, 0.2]] + ) + fb = ms.FunctionBuilder( + ms.NamedTypes([("index", ms.batch_int)]), + ms.NamedTypes([("r", ms.batch_float), ("det", ms.batch_float)]), + ) + probs = fb.get_global("test_simd_layout_probs", ms.DataType.float, [4]) + r, det = fb.sample_discrete_probs_inverse(fb.input(0), probs) + fb.output(0, r) + fb.output(1, det) + module = FunctionModule(fb.function(), ctx) + param = module.global_params["test_simd_layout_probs"] + + index = torch.tensor(rng.integers(0, 4, size=BATCH_SIZE), dtype=torch.int32) + weights = torch.tensor(rng.uniform(size=BATCH_SIZE)) + + def loss(): + return (module(index)[1] * weights).sum() + + loss().backward() + grad = param.grad.clone() + + eps = 1e-6 + grad_fd = torch.zeros(4) + with torch.no_grad(): + for i in range(4): + param[0, i] += eps + plus = loss() + param[0, i] -= 2 * eps + minus = loss() + param[0, i] += eps + grad_fd[i] = (plus - minus) / (2 * eps) + assert grad[0].numpy() == approx(grad_fd.numpy(), rel=1e-5) + + +def test_broadcast_gradient_through_stack(): + """stack of a broadcast (global) and a batched value: the gradient of the + global accumulates the whole batch.""" + ctx = ms.Context() + ctx.define_global("test_simd_layout_scalar", ms.DataType.float, [], True) + ctx.get_global("test_simd_layout_scalar").torch()[:] = 1.5 + fb = ms.FunctionBuilder( + ms.NamedTypes([("x", ms.batch_float)]), + ms.NamedTypes([("y", ms.batch_float_array(2))]), + ) + g = fb.get_global("test_simd_layout_scalar", ms.DataType.float, []) + fb.output(0, fb.stack([g, fb.input(0)])) + module = FunctionModule(fb.function(), ctx) + param = module.global_params["test_simd_layout_scalar"] + + x = torch.tensor(rng.normal(size=BATCH_SIZE)) + weights = torch.tensor(rng.uniform(size=(BATCH_SIZE, 2))) + (module(x) * weights).sum().backward() + assert param.grad.item() == approx(weights[:, 0].sum().item()) From 9a08ca051ba11d5e4a7220dc6c3e09016f4d58e8 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 8 Oct 2026 00:08:22 +0200 Subject: [PATCH 17/23] madspace: fix use-after-free of matrix element inputs with the dummy 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 --- madspace/src/cpu/runtime.cpp | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/madspace/src/cpu/runtime.cpp b/madspace/src/cpu/runtime.cpp index 4ba61ffe70..4ec5d4a32d 100644 --- a/madspace/src/cpu/runtime.cpp +++ b/madspace/src/cpu/runtime.cpp @@ -86,7 +86,11 @@ void op_matrix_element( std::vector input_keys(input_count); std::vector output_keys(output_count); std::vector input_ptrs(input_count), output_ptrs(output_count); - for (std::size_t i = 0; i < input_count; ++i) { + // Only prepare the inputs if the matrix element is actually called: making an + // input contiguous can start an asynchronous copy, whose target would be freed + // while the copy is still running if we returned early below. + bool call_matrix_element = me_index != 0xBADCAFE && batch_size > 0; + for (std::size_t i = 0; i < input_count && call_matrix_element; ++i) { input_keys[i] = static_cast( locals[instruction.input_indices[3 + 2 * i]].index_value() ); @@ -130,7 +134,7 @@ void op_matrix_element( } } } - if (me_index == 0xBADCAFE || batch_size == 0) { + if (!call_matrix_element) { return; } auto& matrix_element = instruction.runtime.context().matrix_element(me_index); From 974d3a5dfb7b4a34b9e68c5c05d56c968430aa30 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Tue, 13 Oct 2026 17:14:45 +0200 Subject: [PATCH 18/23] madspace: make the event field type an enum 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 --- madspace/include/madspace/driver/io.hpp | 36 +++++++++++++++++-------- madspace/src/driver/io.cpp | 5 ++-- 2 files changed, 28 insertions(+), 13 deletions(-) diff --git a/madspace/include/madspace/driver/io.hpp b/madspace/include/madspace/driver/io.hpp index 5e631ad78d..6f363fa0b6 100644 --- a/madspace/include/madspace/driver/io.hpp +++ b/madspace/include/madspace/driver/io.hpp @@ -14,12 +14,33 @@ Tensor load_tensor(const std::string& file); void save_tensor(const std::string& file, Tensor tensor); struct FieldLayout { - static constexpr const char* i32 = " @@ -284,14 +305,7 @@ class DataLayout { size = 0; std::size_t group_index = 0; for (auto& field : layout) { - std::size_t field_size; - if (field.type == FieldLayout::i32) { - field_size = 4; - } else if (field.type == FieldLayout::f64) { - field_size = 8; - } else { - std::logic_error("unknown type"); - } + std::size_t field_size = FieldLayout::size(field.type); if (field.group > group_index) { offsets.resize(field.group); offsets.back() = size; diff --git a/madspace/src/driver/io.cpp b/madspace/src/driver/io.cpp index 3adafb2db9..097195419d 100644 --- a/madspace/src/driver/io.cpp +++ b/madspace/src/driver/io.cpp @@ -87,12 +87,13 @@ full_descr(std::size_t particle_count, const DataLayout& layout) { layout.event_layout().size() + particle_count * layout.particle_layout().size() ); for (auto& field : layout.event_layout()) { - field_descr.emplace_back(field.name, field.type); + field_descr.emplace_back(field.name, FieldLayout::dtype(field.type)); } for (std::size_t i = 1; i <= particle_count; ++i) { for (auto& field : layout.particle_layout()) { field_descr.push_back( - {std::format("part{}_{}", i, field.name), field.type} + {std::format("part{}_{}", i, field.name), + FieldLayout::dtype(field.type)} ); } } From 326a0102677f1ab2d552326c06f474739594a0a9 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Tue, 13 Oct 2026 17:42:58 +0200 Subject: [PATCH 19/23] madspace: fix reading and appending event files 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 --- madspace/src/driver/io.cpp | 31 ++++++++++++++++++++++++------- 1 file changed, 24 insertions(+), 7 deletions(-) diff --git a/madspace/src/driver/io.cpp b/madspace/src/driver/io.cpp index 097195419d..833dfaee9e 100644 --- a/madspace/src/driver/io.cpp +++ b/madspace/src/driver/io.cpp @@ -113,12 +113,22 @@ read_event_header(std::fstream& file_stream, const DataLayout& layout) { std::size_t event_field_count = layout.event_layout().size(); std::size_t particle_field_count = layout.particle_layout().size(); if (!descr.is_array() || descr.size() < event_field_count || - (descr.size() - event_field_count) % particle_field_count != 0 || !fortran_order.is_boolean() || fortran_order.get() || - header_shape.is_array() || header_shape.size() != 1) { + !header_shape.is_array() || header_shape.size() != 1) { throw std::runtime_error("Invalid header for event file"); } - std::size_t particle_count = (descr.size() - event_field_count) / 5; + std::size_t particle_descr_count = descr.size() - event_field_count; + std::size_t particle_count = 0; + if (particle_field_count == 0) { + if (particle_descr_count != 0) { + throw std::runtime_error("Invalid header for event file"); + } + } else { + if (particle_descr_count % particle_field_count != 0) { + throw std::runtime_error("Invalid header for event file"); + } + particle_count = particle_descr_count / particle_field_count; + } std::size_t event_count = header_shape.at(0).get(); auto field_descr = full_descr(particle_count, layout); @@ -129,7 +139,8 @@ read_event_header(std::fstream& file_stream, const DataLayout& layout) { throw std::runtime_error("Invalid header for event file"); } } - std::size_t header_size = 0; + // parse_header leaves the stream at the end of the header + std::size_t header_size = file_stream.tellg(); return {header_size, particle_count, event_count}; } @@ -245,11 +256,17 @@ EventFile::EventFile( } else { std::tie(_header_size, _particle_count, _event_count) = read_event_header(_file_stream, layout); - if (mode == EventFile::load) { + _event_size = layout.event_size() + _particle_count * layout.particle_size(); + _capacity = _event_count; + if (mode == EventFile::append) { + // rewrite the header with an open shape, filled in on close + _file_stream.seekp(0); std::tie(_header_size, _shape_pos) = - write_event_header(_file_stream, particle_count, layout, _header_size); + write_event_header(_file_stream, _particle_count, layout, _header_size); + seek(_event_count); + } else { + seek(0); } - _capacity = _event_count; } } From 3f4efe4f25c114c1173d504a7eafbe108c279144 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Tue, 13 Oct 2026 18:37:01 +0200 Subject: [PATCH 20/23] madspace: make the max-weight freeze independent of thread timing 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 --- .../madspace/driver/channel_generator.hpp | 5 +++++ madspace/src/driver/channel_generator.cpp | 18 ++++++++++++++++-- 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/madspace/include/madspace/driver/channel_generator.hpp b/madspace/include/madspace/driver/channel_generator.hpp index 1830817719..348a712611 100644 --- a/madspace/include/madspace/driver/channel_generator.hpp +++ b/madspace/include/madspace/driver/channel_generator.hpp @@ -205,6 +205,11 @@ class ChannelEventGenerator { RunningIntegral _cross_section; RunningIntegral _abs_cross_section; double _max_weight = 0.; + // Expected number of accepted events, sum of min(|w| / max_weight, 1) over the + // committed generate stages, rescaled like count_unweighted when the max + // weight rises. Decides when the max weight is frozen: unlike count_unweighted + // it only depends on the generate stages, which commit in a fixed order. + double _expected_unweighted = 0.; std::size_t _count_requested_opt = 0; // Monotonic per-job counters keying each job's deterministic random stream; // never reset. Separate for survey/generate so neither depends on the other. diff --git a/madspace/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index 86f471c5a2..575a246602 100644 --- a/madspace/src/driver/channel_generator.cpp +++ b/madspace/src/driver/channel_generator.cpp @@ -352,6 +352,7 @@ void ChannelEventGenerator::unweight_file(MixMaxRandom& rand_gen) { _unweighted_count = _weight_file.event_count(); _unweighted_accept_count = accept_count; _status.count_unweighted = accept_count; + _expected_unweighted = accept_count; } void ChannelEventGenerator::integrate(const GeneratorBatchJob& job) { @@ -638,6 +639,7 @@ std::size_t ChannelEventGenerator::next_vegas_batch_size() { void ChannelEventGenerator::clear_events() { _status.count_unweighted = 0; + _expected_unweighted = 0; _max_weight = 0; _unweighted_count = 0; _unweighted_accept_count = 0; @@ -656,7 +658,11 @@ void ChannelEventGenerator::clear_events() { } void ChannelEventGenerator::update_max_weight(Tensor weights) { - if (_status.count_unweighted > _config.freeze_max_weight_after) { + // Freeze on the expected unweighted count of the committed generate stages + // rather than on count_unweighted: the latter is advanced by the unweight + // stages, whose commits interleave with the generate stages depending on + // thread timing, which made the freeze point (and the run) timing-dependent. + if (_expected_unweighted > _config.freeze_max_weight_after) { return; } @@ -676,6 +682,13 @@ void ChannelEventGenerator::update_max_weight(Tensor weights) { } std::sort(_large_weights.begin(), _large_weights.end(), std::greater{}); apply_truncation_budget(); + // expected accepted events of this job, against the max weight it will be + // unweighted with (see prepare_unweight_job) + if (_max_weight > 0) { + for (std::size_t i = 0; i < w_view.size(); ++i) { + _expected_unweighted += std::min(std::abs(w_view[i]) / _max_weight, 1.); + } + } } void ChannelEventGenerator::set_target_count(std::size_t target_count) { @@ -687,7 +700,7 @@ void ChannelEventGenerator::set_target_count(std::size_t target_count) { // channel that has met its smaller target is not revisited. Tighten the cap // now; unweight_all() re-unweights the whole file against it at the end. if (target_count < old_target && _max_weight != 0 && - _status.count_unweighted <= _config.freeze_max_weight_after) { + _expected_unweighted <= _config.freeze_max_weight_after) { apply_truncation_budget(); } } @@ -708,6 +721,7 @@ void ChannelEventGenerator::apply_truncation_budget() { if (w_sum / w > max_truncation) { if (_max_weight < w) { _status.count_unweighted *= _max_weight / w; + _expected_unweighted *= _max_weight / w; _max_weight = w; _unweighted_count = 0; _unweighted_accept_count = 0; From a296a597acce426573447379201d09948aafdc23 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Sat, 17 Oct 2026 15:08:36 +0200 Subject: [PATCH 21/23] madspace: default to auto SIMD mode, scalar on ARM and simd_512 > simd_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. --- docs/source/run_card.rst | 9 +++++---- madgraph/iolibs/template_files/mg7/gridpack.py | 2 +- madgraph/iolibs/template_files/mg7/run_card.toml | 2 +- madgraph/various/banner.py | 5 ++--- madspace/include/madspace/driver/backend.hpp | 7 ++++--- madspace/src/driver/backend.cpp | 9 +++------ 6 files changed, 16 insertions(+), 18 deletions(-) diff --git a/docs/source/run_card.rst b/docs/source/run_card.rst index 2e67c637e1..7833a28c02 100644 --- a/docs/source/run_card.rst +++ b/docs/source/run_card.rst @@ -49,10 +49,11 @@ fixed parameters. Their entries are described below. ``simd_256``, ``simd_512`` and ``avx512y``. ``auto`` picks the widest width the host supports. -``madspace_cpu_mode`` (``"scalar"``) - SIMD width of the phase-space sampling on ``cpu`` devices. This is experimental. Choose - from ``auto``, ``scalar``, ``simd_128`` (ARM only), ``simd_256``, ``avx512y`` and - ``simd_512``. ``auto`` makes the same choice as for ``cpu_mode``. +``madspace_cpu_mode`` (``"auto"``) + SIMD width of the phase-space sampling on ``cpu`` devices. Choose from ``auto``, + ``scalar``, ``simd_128`` (ARM only), ``simd_256``, ``avx512y`` and ``simd_512``. + ``auto`` picks ``simd_512``, then ``simd_256`` on x86. On ARM it picks ``scalar``, + because the vectorized mode is currently slower there. ``precision`` (``"color32"``) Floating-point precision of the matrix elements. ``all64`` and ``all32`` use double and diff --git a/madgraph/iolibs/template_files/mg7/gridpack.py b/madgraph/iolibs/template_files/mg7/gridpack.py index 54c13859fe..d1d9caafdd 100644 --- a/madgraph/iolibs/template_files/mg7/gridpack.py +++ b/madgraph/iolibs/template_files/mg7/gridpack.py @@ -136,7 +136,7 @@ def main() -> None: parser.add_argument( "--madspace_cpu_mode", type=str, - default=run_args.get("madspace_cpu_mode", "scalar"), + default=run_args.get("madspace_cpu_mode", "auto"), choices=["auto", "scalar", "simd_128", "simd_256", "avx512y", "simd_512"], help="SIMD width of the phase-space sampling (madspace) on the 'cpu' devices", ) diff --git a/madgraph/iolibs/template_files/mg7/run_card.toml b/madgraph/iolibs/template_files/mg7/run_card.toml index ed21728172..e8bc6d1fef 100644 --- a/madgraph/iolibs/template_files/mg7/run_card.toml +++ b/madgraph/iolibs/template_files/mg7/run_card.toml @@ -8,7 +8,7 @@ seed = %(run.seed)s device = %(run.device)s # SIMD width of the matrix elements on cpu devices; options: auto, scalar, simd_128, simd_256, simd_512, avx512y cpu_mode = %(run.cpu_mode)s -# SIMD width of the phase-space sampling on cpu devices (experimental); options: auto, scalar, simd_128 (ARM only), simd_256, avx512y, simd_512 +# SIMD width of the phase-space sampling on cpu devices; options: auto, scalar, simd_128 (ARM only), simd_256, avx512y, simd_512 madspace_cpu_mode = %(run.madspace_cpu_mode)s # matrix element precision; options: all64, all32, color32, denom64 precision = %(run.precision)s diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py index 54581b8a96..34a8890d63 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -6714,10 +6714,9 @@ def default_setup(self): self.add_toml_param('run', 'cpu_mode', "auto", gridpack=True, allowed=['auto', 'scalar', 'simd_128', 'simd_256', 'simd_512', 'avx512y'], comment="SIMD width of the matrix elements on cpu devices; auto picks the widest one the host supports") - self.add_toml_param('run', 'madspace_cpu_mode', "scalar", gridpack=True, + self.add_toml_param('run', 'madspace_cpu_mode', "auto", gridpack=True, allowed=['auto', 'scalar', 'simd_128', 'simd_256', 'avx512y', 'simd_512'], - comment="SIMD width of the phase-space sampling on cpu devices (experimental); " - "simd_128 is ARM only, auto makes the same choice as for cpu_mode") + comment="SIMD width of the phase-space sampling on cpu devices; simd_128 is ARM only") self.add_toml_param('run', 'precision', "color32", gridpack=True, allowed=['all64', 'all32', 'color32', 'denom64'], comment="matrix element precision: all64, all32, color32 (FP32 color sum) or denom64 (FP64 momenta and propagator denominators)") diff --git a/madspace/include/madspace/driver/backend.hpp b/madspace/include/madspace/driver/backend.hpp index 176844d69f..287743071d 100644 --- a/madspace/include/madspace/driver/backend.hpp +++ b/madspace/include/madspace/driver/backend.hpp @@ -48,9 +48,10 @@ DevicePtr hip_device(std::size_t index = 0); /// Directory to load the generated backend shared libraries from. void set_lib_path(const std::string& lib_path); /// SIMD mode of the CPU backend: "scalar", "simd_128" (ARM), "simd_256" (AVX2), -/// "avx512y" (AVX512 with 256-bit vectors), "simd_512" (AVX512) or "auto" (best -/// mode for the host, chosen as in madmatrix). Must be set before the CPU backend -/// is first used; defaults to $MADSPACE_SIMD_MODE or "scalar". +/// "avx512y" (AVX512 with 256-bit vectors), "simd_512" (AVX512) or "auto" (simd_512, +/// then simd_256 on x86 and scalar on ARM, where the vectorized mode is currently +/// slower). Must be set before the CPU backend is first used; defaults to +/// $MADSPACE_SIMD_MODE or "scalar". void set_simd_mode(const std::string& simd_mode); /// SIMD modes supported by this build on this host (without "auto"), ordered by /// increasing vector width; see @ref set_simd_mode. diff --git a/madspace/src/driver/backend.cpp b/madspace/src/driver/backend.cpp index cd25ca595c..657333ebd8 100644 --- a/madspace/src/driver/backend.cpp +++ b/madspace/src/driver/backend.cpp @@ -80,7 +80,7 @@ std::pair, std::string> simd_mode_support() { #ifdef SIMD_AVAILABLE #if defined(__aarch64__) || defined(_M_ARM64) supported_modes.push_back("simd_128"); - auto_mode = "simd_128"; + // auto stays scalar: the vectorized mode is currently slower on ARM #elif defined(__x86_64__) || defined(_M_X64) if (__builtin_cpu_supports("avx2") && __builtin_cpu_supports("fma")) { supported_modes.push_back("simd_256"); @@ -88,15 +88,12 @@ std::pair, std::string> simd_mode_support() { } if (__builtin_cpu_supports("avx512f") && __builtin_cpu_supports("avx512vl") && __builtin_cpu_supports("avx512dq")) { + // never chosen by auto: slower than simd_512 in the x86 benchmarks supported_modes.push_back("avx512y"); - // mirrors madmatrix's auto mode (avx512y except with clang, never - // simd_512); TODO: revisit once madspace has its own x86 benchmarks -#ifndef __clang__ - auto_mode = "avx512y"; -#endif } if (__builtin_cpu_supports("avx512f")) { supported_modes.push_back("simd_512"); + auto_mode = "simd_512"; } #endif #endif // SIMD_AVAILABLE From 614e403a2ecca1c7486eadede6a7728318114c53 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Sat, 17 Oct 2026 15:31:56 +0200 Subject: [PATCH 22/23] madspace: auto SIMD mode as C++ default, refuse late mode switches, unpadded 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. --- madspace/include/madspace/driver/backend.hpp | 4 +- madspace/src/cpu/runtime.cpp | 33 ++++++---- madspace/src/cpu/runtime.hpp | 3 + madspace/src/driver/backend.cpp | 63 +++++++++++++------- madspace/tests/conftest.py | 2 +- tests/input_files/mg7_run_card.toml | 6 +- 6 files changed, 70 insertions(+), 41 deletions(-) diff --git a/madspace/include/madspace/driver/backend.hpp b/madspace/include/madspace/driver/backend.hpp index 287743071d..85a2073c9d 100644 --- a/madspace/include/madspace/driver/backend.hpp +++ b/madspace/include/madspace/driver/backend.hpp @@ -50,8 +50,8 @@ void set_lib_path(const std::string& lib_path); /// SIMD mode of the CPU backend: "scalar", "simd_128" (ARM), "simd_256" (AVX2), /// "avx512y" (AVX512 with 256-bit vectors), "simd_512" (AVX512) or "auto" (simd_512, /// then simd_256 on x86 and scalar on ARM, where the vectorized mode is currently -/// slower). Must be set before the CPU backend is first used; defaults to -/// $MADSPACE_SIMD_MODE or "scalar". +/// slower). Defaults to $MADSPACE_SIMD_MODE or "auto". Must be set before the CPU +/// backend is first used; afterwards, requesting a different mode throws. void set_simd_mode(const std::string& simd_mode); /// SIMD modes supported by this build on this host (without "auto"), ordered by /// increasing vector width; see @ref set_simd_mode. diff --git a/madspace/src/cpu/runtime.cpp b/madspace/src/cpu/runtime.cpp index 65a71dfa63..5685f7d5f5 100644 --- a/madspace/src/cpu/runtime.cpp +++ b/madspace/src/cpu/runtime.cpp @@ -1132,7 +1132,8 @@ CpuRuntime::CpuRuntime(const Function& function, ContextPtr context, bool concur ); for (auto& in : instr.inputs) { auto& local = function.locals().at(in.local_index); - bool is_literal = !std::holds_alternative(local.literal_value); + bool is_literal = + !std::holds_alternative(local.literal_value); bool is_int_intermediate = local.type.dtype == DataType::dt_int && static_cast(in.local_index) >= function.inputs().size() && !global_locals.contains(in.local_index); @@ -1229,6 +1230,13 @@ CpuRuntime::CpuRuntime(const Function& function, ContextPtr context, bool concur } for (auto& out : function.outputs()) { _output_indices.push_back(out.local_index); + // literal outputs are returned unpadded, so that callers see a plain tensor + auto& literal_value = function.locals().at(out.local_index).literal_value; + _literal_outputs.push_back( + std::holds_alternative(literal_value) + ? Tensor() + : Tensor::simd_padded_literal(literal_value, 1, cpu_device()) + ); _copy_output_grads.push_back( grad_accumulated.at(out.local_index) || output_counts.at(out.local_index) > 1 @@ -1236,6 +1244,14 @@ CpuRuntime::CpuRuntime(const Function& function, ContextPtr context, bool concur } } +TensorVec CpuRuntime::collect_outputs(const TensorVec& locals) const { + TensorVec outputs; + for (auto [index, literal] : zip(_output_indices, _literal_outputs)) { + outputs.push_back(literal ? literal : locals[index]); + } + return outputs; +} + TensorVec CpuRuntime::run(const TensorVec& inputs) { auto contiguous_inputs = batch_contiguous(inputs); if (_concurrent && _context->thread_pool().thread_count() > 1) { @@ -1290,10 +1306,7 @@ TensorVec CpuRuntime::run_single(const TensorVec& inputs) const { #include "runtime_mixin.inc" } } - TensorVec outputs; - for (auto index : _output_indices) { - outputs.push_back(locals[index]); - } + TensorVec outputs = collect_outputs(locals); return outputs; } @@ -1341,10 +1354,7 @@ std::tuple> CpuRuntime::run_with_grad_si #include "runtime_mixin.inc" } } - TensorVec outputs; - for (auto index : _output_indices) { - outputs.push_back(locals[index]); - } + TensorVec outputs = collect_outputs(locals); return {outputs, locals, eval_grad}; } @@ -1531,10 +1541,7 @@ std::tuple> CpuRuntime::run_concurrent( } } - TensorVec outputs; - for (auto index : _output_indices) { - outputs.push_back(locals[index]); - } + TensorVec outputs = collect_outputs(locals); if (with_grad) { return {outputs, locals, eval_grad}; } else { diff --git a/madspace/src/cpu/runtime.hpp b/madspace/src/cpu/runtime.hpp index 8324539c85..8a844f79e9 100644 --- a/madspace/src/cpu/runtime.hpp +++ b/madspace/src/cpu/runtime.hpp @@ -74,6 +74,7 @@ class CpuRuntime : public Runtime { const std::vector& eval_grad, bool return_contiguous_grads ) const; + TensorVec collect_outputs(const TensorVec& locals) const; std::pair run_backward_concurrent( const TensorVec& output_grads, const TensorVec& stored_locals, @@ -83,6 +84,8 @@ class CpuRuntime : public Runtime { std::vector _instructions; SizeVec _output_indices; + // unpadded copies of the outputs that are literals, empty for the others + TensorVec _literal_outputs; std::vector _copy_output_grads; std::size_t _input_count; TensorVec _locals_init; diff --git a/madspace/src/driver/backend.cpp b/madspace/src/driver/backend.cpp index 657333ebd8..313dab2b61 100644 --- a/madspace/src/driver/backend.cpp +++ b/madspace/src/driver/backend.cpp @@ -13,7 +13,9 @@ namespace { struct LoadedBackend { inline static std::string lib_path = ""; + // resolved SIMD mode requested with set_simd_mode, and the one actually loaded inline static std::string simd_mode = ""; + inline static std::string loaded_simd_mode = ""; inline static std::unordered_map device_backends; LoadedBackend(const std::string& file) { @@ -100,33 +102,38 @@ std::pair, std::string> simd_mode_support() { return {supported_modes, auto_mode}; } +// SIMD mode that a requested mode ("auto" or an explicit one) resolves to +std::string resolve_simd_mode(const std::string& mode) { + auto [supported_modes, auto_mode] = simd_mode_support(); + if (mode == "auto") { + return auto_mode; + } + if (std::find(supported_modes.begin(), supported_modes.end(), mode) == + supported_modes.end()) { + std::string supported; + for (auto& supported_mode : supported_modes) { + supported += std::format(" '{}'", supported_mode); + } + throw std::runtime_error( + std::format( + "unsupported madspace SIMD mode '{}', supported on this host: " + "'auto'{}", + mode, + supported + ) + ); + } + return mode; +} + const LoadedBackend& cpu_backend() { static LoadedBackend backend = [&] { - auto [supported_modes, auto_mode] = simd_mode_support(); - std::string mode = LoadedBackend::simd_mode; if (mode.empty()) { char* env_var = std::getenv("MADSPACE_SIMD_MODE"); - mode = env_var ? env_var : "scalar"; + mode = resolve_simd_mode(env_var ? env_var : "auto"); } - if (mode == "auto") { - mode = auto_mode; - } else if (std::find(supported_modes.begin(), supported_modes.end(), mode) == - supported_modes.end()) { - std::string supported; - for (auto& supported_mode : supported_modes) { - supported += std::format(" '{}'", supported_mode); - } - throw std::runtime_error( - std::format( - "unsupported madspace SIMD mode '{}', supported on this host: " - "'auto'{}", - mode, - supported - ) - ); - } - + LoadedBackend::loaded_simd_mode = mode; if (mode == "scalar") { return LoadedBackend("libmadspace_cpu"); } else { @@ -198,5 +205,17 @@ std::vector madspace::supported_simd_modes() { } void madspace::set_simd_mode(const std::string& simd_mode) { - LoadedBackend::simd_mode = simd_mode; + std::string mode = resolve_simd_mode(simd_mode); + const std::string& loaded_mode = LoadedBackend::loaded_simd_mode; + if (!loaded_mode.empty() && mode != loaded_mode) { + throw std::runtime_error( + std::format( + "can not switch the madspace SIMD mode to '{}': the CPU backend was " + "already loaded with '{}'", + mode, + loaded_mode + ) + ); + } + LoadedBackend::simd_mode = mode; } diff --git a/madspace/tests/conftest.py b/madspace/tests/conftest.py index 15115a9632..a623a64d86 100644 --- a/madspace/tests/conftest.py +++ b/madspace/tests/conftest.py @@ -26,7 +26,7 @@ def pytest_configure(config): def pytest_report_header(config): name = torch.cuda.get_device_name(DEVICE) if ON_GPU else "cpu" - simd_mode = os.environ.get("MADSPACE_SIMD_MODE", "scalar") + simd_mode = os.environ.get("MADSPACE_SIMD_MODE", "auto") return [ f"madspace test device: {BACKEND} ({name})", f"madspace SIMD mode: {simd_mode}", diff --git a/tests/input_files/mg7_run_card.toml b/tests/input_files/mg7_run_card.toml index c91fa824fc..cbf2cc6bcf 100644 --- a/tests/input_files/mg7_run_card.toml +++ b/tests/input_files/mg7_run_card.toml @@ -5,9 +5,9 @@ device = ["cpu"] # SIMD width of the matrix elements (madmatrix) on the 'cpu' devices (ignored by cuda/hip) # options: auto (widest width supported by the host), scalar, simd_128, simd_256, simd_512, avx512y cpu_mode = "scalar" -# SIMD width of the phase-space sampling (madspace) on the 'cpu' devices (experimental) -# options: auto (same choice as for cpu_mode), scalar, simd_128 (ARM only), simd_256, avx512y, simd_512 -madspace_cpu_mode = "scalar" +# SIMD width of the phase-space sampling (madspace) on the 'cpu' devices +# options: auto (simd_512, then simd_256 on x86; scalar on ARM), scalar, simd_128 (ARM only), simd_256, avx512y, simd_512 +madspace_cpu_mode = "auto" # pool sizes: -1 sets count automatically based on number of CPUs cpu_thread_pool_size = -1 gpu_thread_pool_size = 1 From 949abb458cdd7ab859ff0efb682dfc3b4aa99492 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Sat, 17 Oct 2026 15:32:08 +0200 Subject: [PATCH 23/23] madspace tests: SIMD layout tests cpu_only, gradient tests on the test 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. --- madspace/tests/test_simd_layout.py | 35 ++++++++++++++++++++++++++---- 1 file changed, 31 insertions(+), 4 deletions(-) diff --git a/madspace/tests/test_simd_layout.py b/madspace/tests/test_simd_layout.py index e849b292de..37f592f576 100644 --- a/madspace/tests/test_simd_layout.py +++ b/madspace/tests/test_simd_layout.py @@ -9,11 +9,15 @@ kernel type. These tests run in every SIMD mode (tests/run_simd_modes.py) and compare against numpy, so all of these paths have to agree. They also cover the constant folding rules that keep broadcast tensors away from the kernels. + +Except for the gradient tests, which also cover kernel code shared with the GPU, +these are CPU runtime or graph building tests and are marked cpu_only. """ import numpy as np import pytest import torch +from madspace_testing import new_context from pytest import approx import madspace as ms @@ -50,6 +54,7 @@ def instruction_names(function): # constant folding +@pytest.mark.cpu_only def test_sub_with_constant_first_is_not_folded(): func = build( [("x", ms.batch_float)], @@ -62,6 +67,7 @@ def test_sub_with_constant_first_is_not_folded(): assert y == approx(-x) +@pytest.mark.cpu_only @pytest.mark.parametrize( "op, const_first, const", [ @@ -81,6 +87,7 @@ def body(fb, x): assert instruction_names(func) == [] +@pytest.mark.cpu_only def test_stack_of_literals_is_folded(): def body(fb, x): stacked = fb.stack([1.0, 2.0, 3.0]) @@ -93,6 +100,7 @@ def body(fb, x): assert instruction_names(func) == ["mul"] +@pytest.mark.cpu_only def test_single_flavor_integrand_has_no_constant_gathers(tmp_path): """With a single flavor the flavor index is a constant tensor: the flavor tables (remap, factors, mirror) must be resolved when building the graph @@ -151,6 +159,7 @@ def test_single_flavor_integrand_has_no_constant_gathers(tmp_path): # memory layouts +@pytest.mark.cpu_only @pytest.mark.parametrize("order", ["C", "F"]) def test_row_major_and_column_major_inputs(order): """Column-major (batch stride one) inputs run the vectorized kernels, @@ -166,6 +175,7 @@ def test_row_major_and_column_major_inputs(order): assert z == approx(x * y + np.abs(y)) +@pytest.mark.cpu_only @pytest.mark.parametrize("size", [3, 8, 100]) def test_literal_arrays(size): """Small literal arrays are stored padded to the vector size, large ones @@ -182,6 +192,20 @@ def test_literal_arrays(size): assert y == approx(x * np.array(values) + 1.5) +@pytest.mark.cpu_only +def test_literal_output_is_not_padded(): + func = build( + [("x", ms.batch_float)], + [("y", ms.batch_float_array(3))], + lambda fb, x: fb.stack([1.0, 2.0, 3.0]), + ) + (y,) = call(func, rng.normal(size=BATCH_SIZE)) + assert y.reshape(-1, 3)[0] == approx([1.0, 2.0, 3.0]) + # the padded storage would give the inner dimension a stride of the vector size + assert y.strides[-1] == y.itemsize + + +@pytest.mark.cpu_only def test_broadcast_full_input(): """fb.full creates a batch-sized tensor with batch stride zero.""" @@ -195,6 +219,7 @@ def body(fb, x): assert y == approx(2.5 * x) +@pytest.mark.cpu_only def test_cut_all_and_cut_any_with_scalar_bounds(): func = build( [("obs", ms.batch_float_array(4))], @@ -211,6 +236,7 @@ def test_cut_all_and_cut_any_with_scalar_bounds(): # single inputs +@pytest.mark.cpu_only def test_single_input_rejects_batched_argument(): fb = ms.FunctionBuilder( ms.NamedTypes([("r", ms.batch_float)]), @@ -221,6 +247,7 @@ def test_single_input_rejects_batched_argument(): fb.breit_wigner_invariant(r, r, 1.0, 0.0, 100.0) +@pytest.mark.cpu_only def test_single_inputs(): """Breit-Wigner sampling with literal (single) mass and width, compared between its forward and inverse direction.""" @@ -248,7 +275,7 @@ def body(fb, r): def test_global_probability_gradient_is_summed_over_events(): """The gradient of a broadcast (global) probability vector is the sum of the contributions of all events.""" - ctx = ms.Context() + ctx = new_context() ctx.define_global("test_simd_layout_probs", ms.DataType.float, [4], True) ctx.get_global("test_simd_layout_probs").torch()[:] = torch.tensor( [[0.1, 0.4, 0.3, 0.2]] @@ -274,7 +301,7 @@ def loss(): grad = param.grad.clone() eps = 1e-6 - grad_fd = torch.zeros(4) + grad_fd = torch.zeros(4, dtype=torch.float64) with torch.no_grad(): for i in range(4): param[0, i] += eps @@ -283,13 +310,13 @@ def loss(): minus = loss() param[0, i] += eps grad_fd[i] = (plus - minus) / (2 * eps) - assert grad[0].numpy() == approx(grad_fd.numpy(), rel=1e-5) + assert grad[0].cpu().numpy() == approx(grad_fd.cpu().numpy(), rel=1e-5) def test_broadcast_gradient_through_stack(): """stack of a broadcast (global) and a batched value: the gradient of the global accumulates the whole batch.""" - ctx = ms.Context() + ctx = new_context() ctx.define_global("test_simd_layout_scalar", ms.DataType.float, [], True) ctx.get_global("test_simd_layout_scalar").torch()[:] = 1.5 fb = ms.FunctionBuilder(