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 57cd96324d..c3315206e9 100644 --- a/.github/workflows/acceptancetest_mg7.yml +++ b/.github/workflows/acceptancetest_mg7.yml @@ -209,6 +209,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 @@ -314,6 +315,20 @@ 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, 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 test_madnis_gridpack_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 e499a6b5cd..bd8869afa6 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 @@ -89,7 +89,9 @@ jobs: torch # the GPU runs of these tests are in gpu_runner_ci.yml (job madspace_tests) CIBW_TEST_ENVIRONMENT: MADSPACE_TEST_DEVICE=cpu - 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/gridpacks.rst b/docs/source/gridpacks.rst index 76bd699861..747a746cda 100644 --- a/docs/source/gridpacks.rst +++ b/docs/source/gridpacks.rst @@ -120,6 +120,9 @@ gridpack is created. * - - ``device`` - ``--device`` (several entries are allowed) + * - + - ``madspace_cpu_mode`` + - ``--madspace_cpu_mode`` * - - ``cpu_thread_pool_size``, ``gpu_thread_pool_size`` - ``--cpu_thread_pool_size``, ``--gpu_thread_pool_size`` diff --git a/docs/source/madspace/cpp-api.rst b/docs/source/madspace/cpp-api.rst index 8aad6b8723..a728c089fc 100644 --- a/docs/source/madspace/cpp-api.rst +++ b/docs/source/madspace/cpp-api.rst @@ -180,8 +180,9 @@ 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/supported_simd_modes 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..cd1d933c96 100644 --- a/docs/source/madspace/python-api.rst +++ b/docs/source/madspace/python-api.rst @@ -164,8 +164,9 @@ 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/supported_simd_modes python/set_lib_path - python/set_simd_vector_size Utilities --------- diff --git a/docs/source/run_card.rst b/docs/source/run_card.rst index d23b23769f..7833a28c02 100644 --- a/docs/source/run_card.rst +++ b/docs/source/run_card.rst @@ -49,15 +49,18 @@ fixed parameters. Their entries are described below. ``simd_256``, ``simd_512`` and ``avx512y``. ``auto`` picks the widest width the host supports. +``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 single precision throughout. ``color32`` evaluates the color algebra in single precision and everything else in double precision. ``denom64`` keeps the momenta and propagator denominators in double precision and uses single precision elsewhere. -``simd_vector_size`` (``-1``) - Vector size of the CPU matrix elements. ``-1`` chooses automatically. Valid sizes are 1, - 4 and 8 on x86 and 1 and 2 on Apple silicon. ``cpu_thread_pool_size`` (``-1``), ``gpu_thread_pool_size`` (``1``), ``combine_thread_pool_size`` (``-1``) Number of threads for each ``cpu`` device, each GPU device and for combining the diff --git a/madgraph/interface/common_run_interface.py b/madgraph/interface/common_run_interface.py index 5f03acef64..80938fe58c 100755 --- a/madgraph/interface/common_run_interface.py +++ b/madgraph/interface/common_run_interface.py @@ -1241,7 +1241,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 ] @@ -1277,7 +1278,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 3ef6c96b5f..5874066426 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/gridpack.py b/madgraph/iolibs/template_files/mg7/gridpack.py index 6e4df394c1..d1d9caafdd 100644 --- a/madgraph/iolibs/template_files/mg7/gridpack.py +++ b/madgraph/iolibs/template_files/mg7/gridpack.py @@ -133,6 +133,13 @@ def main() -> None: "generate the gridpack (can lead to errors or incorrect results)" ) parser.add_argument("--device", type=str, nargs="*") + parser.add_argument( + "--madspace_cpu_mode", + type=str, + 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", + ) parser.add_argument( "--cpu_thread_pool_size", type=int, default=run_args["cpu_thread_pool_size"] ) @@ -172,6 +179,8 @@ def main() -> None: print(f"\033[1m\033[31mWARNING\033[39m: {message}\033[0m") print() 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 output directories; command line paths are relative to the # invocation directory, run card paths to the gridpack diff --git a/madgraph/iolibs/template_files/mg7/launch.py b/madgraph/iolibs/template_files/mg7/launch.py index c60077657e..d79758a271 100644 --- a/madgraph/iolibs/template_files/mg7/launch.py +++ b/madgraph/iolibs/template_files/mg7/launch.py @@ -469,7 +469,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: output_dir = self.run_card["gridpack"]["output_dir"] diff --git a/madgraph/iolibs/template_files/mg7/run_card.toml b/madgraph/iolibs/template_files/mg7/run_card.toml index bc69c7c5d8..e8bc6d1fef 100644 --- a/madgraph/iolibs/template_files/mg7/run_card.toml +++ b/madgraph/iolibs/template_files/mg7/run_card.toml @@ -6,12 +6,12 @@ run_name = %(run.run_name)s seed = %(run.seed)s # options: cpu, cuda, hip, each optionally with a device index, e.g. "cuda:1" device = %(run.device)s -# SIMD width on cpu devices; options: auto, scalar, simd_128, simd_256, simd_512, avx512y +# 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; 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 -# -1 chooses automatically; x86: 1, 4, 8; Apple silicon: 1, 2 -simd_vector_size = %(run.simd_vector_size)s # -1 sets the thread count automatically 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/banner.py b/madgraph/various/banner.py index d7266c5c7e..34a8890d63 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -6713,12 +6713,13 @@ def default_setup(self): comment="list of devices, each cpu, cuda or hip with an optional 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 on cpu devices; auto picks the widest one the host supports") + 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', "auto", gridpack=True, + allowed=['auto', 'scalar', 'simd_128', 'simd_256', 'avx512y', 'simd_512'], + 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)") - self.add_toml_param('run', 'simd_vector_size', -1, - comment="-1 chooses automatically; x86: 1, 4, 8; Apple silicon: 1, 2") self.add_toml_param('run', 'cpu_thread_pool_size', -1, gridpack=True, comment="threads for the cpu devices; -1 chooses automatically") self.add_toml_param('run', 'gpu_thread_pool_size', 1, gridpack=True, @@ -7044,7 +7045,7 @@ def _derived_beam_value(self, lname): 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 @@ -7059,14 +7060,17 @@ def __setitem__(self, name, value, *args, **opts): tool reading such a sample (MadSpin, systematics, ...) abort on a parameter it never uses. """ - 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]: - message = ("Invalid cpu_mode='%s': supported values are [ '%s' ]" - % (str(value).strip(), "', '".join(str(v) for v in allowed))) + message = ("Invalid %s='%s': supported values are [ '%s' ]" + % (key, str(value).strip(), "', '".join(str(v) for v in allowed))) if getattr(self, '_from_banner', False): logger.warning("%s; the run_card of the event file is read with " - "cpu_mode='%s'", message, self['run']['cpu_mode']) + "%s='%s'", message, key, self['run'][key]) return raise InvalidRunCard(message) return super(RunCardMG7, self).__setitem__(name, value, *args, **opts) diff --git a/madspace/CMakeLists.txt b/madspace/CMakeLists.txt index c7166ea774..cb663c65c3 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 ) @@ -438,14 +439,26 @@ 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(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" + "-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( @@ -462,10 +475,13 @@ 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 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/generate_headers.py b/madspace/generate_headers.py index 998615913c..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): @@ -225,7 +226,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']}(" @@ -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/backend.hpp b/madspace/include/madspace/driver/backend.hpp index 33bba3c1e6..85a2073c9d 100644 --- a/madspace/include/madspace/driver/backend.hpp +++ b/madspace/include/madspace/driver/backend.hpp @@ -47,7 +47,14 @@ 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), +/// "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). 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. +std::vector supported_simd_modes(); } // namespace madspace diff --git a/madspace/include/madspace/driver/channel_generator.hpp b/madspace/include/madspace/driver/channel_generator.hpp index 08e9c1c46f..5319e0a9ba 100644 --- a/madspace/include/madspace/driver/channel_generator.hpp +++ b/madspace/include/madspace/driver/channel_generator.hpp @@ -244,6 +244,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.; bool _max_weight_fixed = false; // mean and standard error std::optional> _abs_integral_prior; diff --git a/madspace/include/madspace/driver/io.hpp b/madspace/include/madspace/driver/io.hpp index 7d7f189a9c..ddcb329342 100644 --- a/madspace/include/madspace/driver/io.hpp +++ b/madspace/include/madspace/driver/io.hpp @@ -21,12 +21,33 @@ Tensor load_tensor(TarReader& tar, const std::string& name); void save_tensor(TarWriter& tar, const std::string& name, Tensor tensor); struct FieldLayout { - static constexpr const char* i32 = " @@ -289,14 +310,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/include/madspace/driver/tensor.hpp b/madspace/include/madspace/driver/tensor.hpp index 0feb424426..9fbad34f1a 100644 --- a/madspace/include/madspace/driver/tensor.hpp +++ b/madspace/include/madspace/driver/tensor.hpp @@ -120,6 +120,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; }; /** @@ -407,11 +410,14 @@ class Tensor { }) { impl->stream = stream; 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; @@ -544,7 +550,7 @@ class Tensor { T* data = static_cast(impl->data); if (flatten_count <= 1) { check_view_dim(dim, impl->shape.size()); - 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"); @@ -799,6 +805,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; @@ -814,6 +844,7 @@ class Tensor { SizeVec batch_sizes; bool stream_ordered = false; std::optional stream; + bool simd_padded = false; template void reset(const D& device) { diff --git a/madspace/instruction_set.yaml b/madspace/instruction_set.yaml index 30125d9b21..7621b2ff91 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 @@ -2290,10 +2290,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] @@ -2316,10 +2316,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] @@ -2342,13 +2342,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] @@ -2371,13 +2371,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] @@ -2400,7 +2400,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] @@ -2423,7 +2423,7 @@ stable_invariant_inverse: type: [float] desc: "invariant value" - name: mass - type: [float] + type: [float, single] desc: "pole mass" - name: s_min type: [float] @@ -2446,10 +2446,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] @@ -2472,10 +2472,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] @@ -2647,10 +2647,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 @@ -2664,10 +2664,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 @@ -2681,10 +2681,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 @@ -2754,10 +2754,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 @@ -2780,10 +2780,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 @@ -3199,6 +3199,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: @@ -3223,7 +3224,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 @@ -3237,7 +3238,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 @@ -3266,7 +3267,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/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; } } 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/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 9b4a717926..5685f7d5f5 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" @@ -46,8 +47,33 @@ 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 { +// 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 @@ -60,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() ); @@ -104,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); @@ -1036,6 +1066,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; @@ -1096,6 +1130,17 @@ 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; @@ -1161,8 +1206,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) {} }, @@ -1182,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 @@ -1189,22 +1244,33 @@ 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) { - 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); } } @@ -1214,13 +1280,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 ); } } @@ -1239,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; } @@ -1290,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}; } @@ -1480,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 03bf7edf4d..8a844f79e9 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); @@ -70,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, @@ -79,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/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_arm.hpp b/madspace/src/cpu/simd_arm.hpp index 51a38ac269..279ba6c138 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; }; @@ -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); @@ -173,4 +163,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/simd_x86_256.hpp b/madspace/src/cpu/simd_x86_256.hpp index 6c2815e25a..eb770f1158 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; @@ -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, @@ -95,7 +99,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] = @@ -126,23 +130,10 @@ 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); - 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) { @@ -193,10 +184,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_256y.hpp b/madspace/src/cpu/simd_x86_256y.hpp new file mode 100644 index 0000000000..bdf193a1c6 --- /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 + )); +} + +// 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) { + 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, + 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, FVec values) { _mm256_storeu_pd(base_ptr, values); } + +inline void vstore(int* base_ptr, IVec values) { + _mm_storeu_si128( + reinterpret_cast<__m128i*>(base_ptr), _mm256_cvtepi64_epi32(values) + ); +} + +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/cpu/simd_x86_512.hpp b/madspace/src/cpu/simd_x86_512.hpp index 1a7c97a987..c143fb69a9 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); @@ -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 cb0bd7afe6..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 _batch_stride; + 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; +}; + +// 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,27 +276,59 @@ 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(0.0)); - } 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 struct first_param; template @@ -313,9 +414,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; } } @@ -331,25 +435,69 @@ 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, scalar_args...](std::size_t count, std::size_t offset) mutable { + [flat_views, + 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) { - 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 (vectorize) { + auto vectorized_views = std::apply( + [&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; + 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/driver/backend.cpp b/madspace/src/driver/backend.cpp index 327c2fd3c1..313dab2b61 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; @@ -11,7 +13,9 @@ namespace { struct LoadedBackend { inline static std::string lib_path = ""; - inline static int vector_size = -1; + // 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) { @@ -71,54 +75,69 @@ struct LoadedBackend { ); }; -const LoadedBackend& cpu_backend() { - static LoadedBackend backend = [&] { - std::vector supported_vector_sizes{1}; +// 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 -#ifdef __APPLE__ - supported_vector_sizes.push_back(2); -#else // __APPLE__ - if (__builtin_cpu_supports("avx2") && __builtin_cpu_supports("fma")) { - supported_vector_sizes.push_back(4); - } - if (__builtin_cpu_supports("avx512f")) { - supported_vector_sizes.push_back(8); - } -#endif // __APPLE__ +#if defined(__aarch64__) || defined(_M_ARM64) + supported_modes.push_back("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"); + auto_mode = "simd_256"; + } + 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"); + } + if (__builtin_cpu_supports("avx512f")) { + supported_modes.push_back("simd_512"); + auto_mode = "simd_512"; + } +#endif #endif // SIMD_AVAILABLE + return {supported_modes, auto_mode}; +} - 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; - } - } - 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"); +// 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; +} - 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: +const LoadedBackend& cpu_backend() { + static LoadedBackend backend = [&] { + std::string mode = LoadedBackend::simd_mode; + if (mode.empty()) { + char* env_var = std::getenv("MADSPACE_SIMD_MODE"); + mode = resolve_simd_mode(env_var ? env_var : "auto"); + } + LoadedBackend::loaded_simd_mode = mode; + if (mode == "scalar") { return LoadedBackend("libmadspace_cpu"); + } else { + return LoadedBackend(std::format("libmadspace_cpu_{}", mode)); } }(); return backend; @@ -181,6 +200,22 @@ 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; +std::vector madspace::supported_simd_modes() { + return simd_mode_support().first; +} + +void madspace::set_simd_mode(const std::string& 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/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index 317db57855..770d6be7f1 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) { @@ -639,6 +640,7 @@ std::size_t ChannelEventGenerator::next_vegas_batch_size() { void ChannelEventGenerator::clear_events() { _status.count_unweighted = 0; + _expected_unweighted = 0; if (!_max_weight_fixed) { _max_weight = 0; } @@ -659,8 +661,11 @@ void ChannelEventGenerator::clear_events() { } void ChannelEventGenerator::update_max_weight(Tensor weights) { - if (_max_weight_fixed || - _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 (_max_weight_fixed || _expected_unweighted > _config.freeze_max_weight_after) { return; } @@ -680,6 +685,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) { @@ -691,7 +703,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 && !_max_weight_fixed && - _status.count_unweighted <= _config.freeze_max_weight_after) { + _expected_unweighted <= _config.freeze_max_weight_after) { apply_truncation_budget(); } } @@ -712,6 +724,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; diff --git a/madspace/src/driver/io.cpp b/madspace/src/driver/io.cpp index 2c733d9280..87c5a79e5a 100644 --- a/madspace/src/driver/io.cpp +++ b/madspace/src/driver/io.cpp @@ -144,12 +144,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)} ); } } @@ -169,12 +170,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); @@ -185,7 +196,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}; } @@ -309,11 +321,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; } } 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; +} 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 361961f56b..e2740959ba 100644 --- a/madspace/src/kernels/cuts.hpp +++ b/madspace/src/kernels/cuts.hpp @@ -25,31 +25,34 @@ 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.); } // The window [min, max] is one pair of scalars applied to every observable of the // vector, as the instruction set declares them (cut_all/cut_any: min, max [float]). -// They must be FIn: as FIn the runtime built a 2-dimensional view of -// these 1-dimensional tensors, from shape/stride entries that were never set, and -// min[i] read out of bounds (a memory fault on HIP, silently min[0] elsewhere). +// They must be rank 0 (here FSIn, as they are declared single): as rank-1 +// inputs the runtime built a 2-dimensional view of these 1-dimensional tensors, +// from shape/stride entries that were never set, and min[i] read out of bounds +// (a memory fault on HIP, silently min[0] elsewhere). template KERNELSPEC void -kernel_cut_all(FIn obs, FIn min, FIn max, FOut w) { - FVal cut = 1.; +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) { - cut = where((obs[i] < min) | (obs[i] > max), 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, 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) { - cut = where((obs[i] < min) | (obs[i] > max), 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/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 9622342005..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); @@ -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 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 011e07dc60..e2e8098a57 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, @@ -605,14 +605,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); } @@ -634,7 +634,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 @@ -647,7 +647,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); @@ -664,8 +664,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 @@ -683,8 +683,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 ) { @@ -704,7 +704,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 / @@ -781,6 +781,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); @@ -893,9 +897,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 ) { @@ -963,9 +967,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 ) { @@ -1049,9 +1053,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 ) { @@ -1131,7 +1135,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]; @@ -1289,7 +1293,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; } @@ -1301,7 +1305,7 @@ kernel_momenta_to_x1x2(FIn p_ext, FIn e_cm, FOut x1, FOut KERNELSPEC void -kernel_boost_to_frame(FIn p_ext, IIn frame_mask, FOut p_out) { +kernel_boost_to_frame(FIn p_ext, ISIn frame_mask, FOut p_out) { FourMom p_boost{0., 0., 0., 0.}; std::size_t selected_count = 0, selected_index = 0; for (std::size_t i = 0; i < p_ext.size(); ++i) { 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/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/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"); 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 3412202c73..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,13 +22,13 @@ 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); 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); diff --git a/madspace/src/phasespace/integrand.cpp b/madspace/src/phasespace/integrand.cpp index 955c768270..6cb68cc243 100644 --- a/madspace/src/phasespace/integrand.cpp +++ b/madspace/src/phasespace/integrand.cpp @@ -446,13 +446,14 @@ NamedVector Integrand::build_channel_part( // Orientation of the event, for a mapping with a mirror_index condition. // Drawn uniformly; the flavor-dependent factor follows the flavor sampling. - Value mirror_index; + Value mirror_index, mirror_det; if (_mapping.mirror_beams()) { if (_has_mirror) { // the det is 2; the flavor-dependent weight is applied after the cuts - auto [index, mirror_det] = + auto [index, det] = fb.sample_discrete(mirror_random, static_cast(2)); mirror_index = index; + mirror_det = det; } else { mirror_index = fb.full({static_cast(0), batch_size_val}); } @@ -595,7 +596,12 @@ NamedVector Integrand::build_channel_part( // stands for itself and its mirror image collects both (factor 2); // for any other flavor the two orientations cover the same phase // space (factor 1). - weights_after_cuts.push_back(fb.gather(flavor_id, _flavor_mirror_factors)); + if (has_multi_flavor) { + weights_after_cuts.push_back(fb.gather(flavor_id, _flavor_mirror_factors)); + } else { + // the single flavor is mirrored (factor 2), which is the sampling det + weights_after_cuts.push_back(fb.batch_gather(indices_acc, mirror_det)); + } } Value weight_after_cuts = weights_after_cuts.empty() @@ -671,6 +677,18 @@ NamedVector Integrand::build_common_part( // which are not invariant under the mirror (HELAS quantises along the frame // z axis). + // 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; @@ -715,7 +733,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); @@ -744,9 +762,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()); @@ -790,9 +808,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")); diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index dae9db04f3..a2b615927b 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -3800,10 +3800,15 @@ 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") + ); + m.def( + "supported_simd_modes", + &supported_simd_modes, + pydoc::doc("supported_simd_modes") ); auto abort_check_function = [] { diff --git a/madspace/tests/conftest.py b/madspace/tests/conftest.py index 6dd5bbfff0..a623a64d86 100644 --- a/madspace/tests/conftest.py +++ b/madspace/tests/conftest.py @@ -3,6 +3,8 @@ Tests marked cpu_only (host-side code, e.g. LHE files) are skipped on a GPU. """ +import os + import pytest import torch from madspace_testing import BACKEND, DEVICE, ON_GPU @@ -24,7 +26,11 @@ def pytest_configure(config): def pytest_report_header(config): name = torch.cuda.get_device_name(DEVICE) if ON_GPU else "cpu" - return f"madspace test device: {BACKEND} ({name})" + simd_mode = os.environ.get("MADSPACE_SIMD_MODE", "auto") + return [ + f"madspace test device: {BACKEND} ({name})", + f"madspace SIMD mode: {simd_mode}", + ] def pytest_collection_modifyitems(config, items): 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/madspace/tests/test_processes.py b/madspace/tests/test_processes.py index 4988d2756a..78c0718584 100644 --- a/madspace/tests/test_processes.py +++ b/madspace/tests/test_processes.py @@ -54,6 +54,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"] @@ -130,8 +148,8 @@ def test_process_momentum_conservation(mapping, masses, permutation_count): assert_approx(p_out, p_in, rel=1e-6, abs=1e-10) -def test_process_inverse(mapping, masses, permutation_count): - r = rand((BATCH_SIZE, mapping.random_dim()), generator=rng) +def test_process_inverse(cut_mapping, masses, permutation_count): + r = rand((BATCH_SIZE, cut_mapping.random_dim()), generator=rng) condition = ( [] if permutation_count <= 1 @@ -141,7 +159,10 @@ def test_process_inverse(mapping, masses, permutation_count): ) ] ) - *map_out, det = mapping.map_forward([r], condition) - r_inv, det_inv = mapping.map_inverse(map_out, condition) - assert_approx(r_inv, r, abs=1e-3, rel=1e-3) - assert_approx(det_inv, 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 torch.any(passed) + assert_approx(r_inv[passed], r[passed], abs=1e-3, rel=1e-3) + assert_approx(det_inv[passed], 1 / det[passed], rel=1e-3) diff --git a/madspace/tests/test_simd_layout.py b/madspace/tests/test_simd_layout.py new file mode 100644 index 0000000000..37f592f576 --- /dev/null +++ b/madspace/tests/test_simd_layout.py @@ -0,0 +1,334 @@ +"""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. + +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 +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 + + +@pytest.mark.cpu_only +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.cpu_only +@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) == [] + + +@pytest.mark.cpu_only +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"] + + +@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 + 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))) + 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): + mapping = ms.PhaseSpaceMapping( + [0.0] * 5, 1000.0, leptonic=True, mirror_beams=mirror + ) + 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.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, + 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.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 + (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) + + +@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.""" + + 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) + + +@pytest.mark.cpu_only +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 + + +@pytest.mark.cpu_only +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) + + +@pytest.mark.cpu_only +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 = 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]] + ) + 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, dtype=torch.float64) + 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].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 = 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( + 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()) 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..4260bc80c2 --- /dev/null +++ b/tests/acceptance_tests/test_mg7_simd_consistency.py @@ -0,0 +1,204 @@ +################################################################################ +# +# 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. + * ``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 +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 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. + + +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 _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(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 = {} + 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}) + 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 + 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)) + + 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) diff --git a/tests/input_files/mg7_run_card.toml b/tests/input_files/mg7_run_card.toml index 22571b4c9e..cbf2cc6bcf 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 +# 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