diff --git a/.github/workflows/acceptancetest_mg7.yml b/.github/workflows/acceptancetest_mg7.yml index c854def99..98ad76b34 100644 --- a/.github/workflows/acceptancetest_mg7.yml +++ b/.github/workflows/acceptancetest_mg7.yml @@ -243,6 +243,92 @@ jobs: export PATH="$HOME/.cache/HEPtools/bin:$PATH" ./tests/test_manager.py test_flavor_grouping_consistency_mg7 -pA -t0 -l INFO + acceptancetest_mg7_vegas_reproducibility: + needs: build_madspace + # mg7 seeded-generation reproducibility (p p > t t~, plain VEGAS-optimized + # run): the same run_card seed must give a byte-identical LHE file across + # independent runs, a different seed must not. Self-skips if the madspace + # + LHAPDF(NNPDF23) stack is unavailable. + runs-on: ubuntu-24.04 + if: github.event_name == 'push' || github.event.pull_request.head.repo.fork == true + steps: + - uses: actions/checkout@v4 + - uses: ./.github/actions/checkout_mg5 + - uses: ./.github/actions/install_madspace + - uses: ./.github/actions/restore-pip-cache + - uses: ./.github/actions/restore_heptools + - name: test one of the test test_vegas_reproducibility_mg7 + run: | + cd $GITHUB_WORKSPACE + export PATH="$HOME/.cache/HEPtools/bin:$PATH" + ./tests/test_manager.py test_vegas_reproducibility_mg7 -pA -t0 -l INFO + + acceptancetest_mg7_madnis_reproducible_mode: + needs: build_madspace + # mg7 madnis training (p p > t t~): training itself must be byte-identical + # for the same seed independent of the cpu thread pool size, both with + # online-only training and with buffered (off-policy replay) training + # enabled. Self-skips if the madspace + LHAPDF(NNPDF23) stack is + # unavailable. + runs-on: ubuntu-24.04 + if: github.event_name == 'push' || github.event.pull_request.head.repo.fork == true + steps: + - uses: actions/checkout@v4 + - uses: ./.github/actions/checkout_mg5 + - uses: ./.github/actions/install_madspace + - uses: ./.github/actions/restore-pip-cache + - uses: ./.github/actions/restore_heptools + - name: test one of the test test_madnis_reproducible_mode_mg7 + run: | + cd $GITHUB_WORKSPACE + export PATH="$HOME/.cache/HEPtools/bin:$PATH" + ./tests/test_manager.py test_madnis_reproducible_mode_mg7 -pA -t0 -l INFO + + acceptancetest_mg7_gridpack_reproducibility: + needs: build_madspace + # mg7 seeded-generation reproducibility for a gridpack trained with + # madnis (p p > t t~): the gridpack's own bin/generate_events --seed must + # give a byte-identical LHE file across independent runs with the same + # seed, and a different result for a different seed. Training itself is + # not required to be deterministic. Self-skips if the madspace + + # LHAPDF(NNPDF23) stack is unavailable. + runs-on: ubuntu-24.04 + if: github.event_name == 'push' || github.event.pull_request.head.repo.fork == true + steps: + - uses: actions/checkout@v4 + - uses: ./.github/actions/checkout_mg5 + - uses: ./.github/actions/install_madspace + - uses: ./.github/actions/restore-pip-cache + - uses: ./.github/actions/restore_heptools + - name: test one of the test test_gridpack_reproducibility_mg7 + run: | + cd $GITHUB_WORKSPACE + export PATH="$HOME/.cache/HEPtools/bin:$PATH" + ./tests/test_manager.py test_gridpack_reproducibility_mg7 -pA -t0 -l INFO + + acceptancetest_mg7_gridpack_reproducibility_vegas: + needs: build_madspace + # mg7 seeded-generation reproducibility for a plain VEGAS-optimized + # gridpack (p p > t t~, no madnis training): the gridpack's own + # bin/generate_events --seed must give a byte-identical LHE file across + # independent runs with the same seed, and a different result for a + # different seed. Companion to acceptancetest_mg7_gridpack_reproducibility, + # which covers the madnis-trained case. Self-skips if the madspace + + # LHAPDF(NNPDF23) stack is unavailable. + runs-on: ubuntu-24.04 + if: github.event_name == 'push' || github.event.pull_request.head.repo.fork == true + steps: + - uses: actions/checkout@v4 + - uses: ./.github/actions/checkout_mg5 + - uses: ./.github/actions/install_madspace + - uses: ./.github/actions/restore-pip-cache + - uses: ./.github/actions/restore_heptools + - name: test one of the test test_gridpack_reproducibility_vegas_mg7 + run: | + cd $GITHUB_WORKSPACE + export PATH="$HOME/.cache/HEPtools/bin:$PATH" + ./tests/test_manager.py test_gridpack_reproducibility_vegas_mg7 -pA -t0 -l INFO + acceptancetest_mg7_merged_flavor_uq: needs: build_madspace # mg7 cross-section for the merged-flavor u q > u q (q = u d), pinned to the diff --git a/.github/workflows/madspace.yml b/.github/workflows/madspace.yml index 2ab34a365..702ab5caf 100644 --- a/.github/workflows/madspace.yml +++ b/.github/workflows/madspace.yml @@ -21,18 +21,16 @@ jobs: yum -y install openblas-devel yum-utils && yum-config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo && yum-config-manager --add-repo /project/.github/workflows/rocm.repo && - yum -y install cuda-nvcc-12-8 cuda-cudart-devel-12-8 cuda-thrust-12-8 libcublas-devel-12-8 libcurand-devel-12-8 rocm-hip-runtime-devel rocblas-devel rocrand-devel rocthrust-devel hipcub-devel + yum -y install cuda-nvcc-12-8 cuda-cudart-devel-12-8 cuda-thrust-12-8 libcublas-devel-12-8 rocm-hip-runtime-devel rocblas-devel rocthrust-devel hipcub-devel CIBW_REPAIR_WHEEL_COMMAND: > auditwheel repair --exclude libmadspace.so --exclude 'libcublas.so.*' - --exclude 'libcurand.so.*' --exclude 'libcublasLt.so.*' --exclude 'libcudart.so.*' --exclude 'libamdhip64.so.*' --exclude 'libhipblaslt.so.*' --exclude 'librocblas.so.*' - --exclude 'librocrand.so.*' --exclude 'libroctx64.so.*' -w {dest_dir} {wheel} CIBW_ENVIRONMENT: > diff --git a/madgraph/iolibs/template_files/mg7/gridpack.py b/madgraph/iolibs/template_files/mg7/gridpack.py index 4a4a01f3a..3751168af 100644 --- a/madgraph/iolibs/template_files/mg7/gridpack.py +++ b/madgraph/iolibs/template_files/mg7/gridpack.py @@ -31,7 +31,6 @@ import tomllib import argparse - def resolve_verbosity(verbosity: str) -> str: """Resolve the run_card "auto" verbosity to "pretty"/"log" depending on whether stdout is attached to a terminal; other values pass through @@ -41,6 +40,14 @@ def resolve_verbosity(verbosity: str) -> str: return verbosity +def resolve_seed(seed: int) -> int: + """Resolve the run_card "seed": -1 draws a fresh 64-bit seed via + os.urandom, any other value is used as-is.""" + if seed == -1: + return int.from_bytes(os.urandom(8), "big") + return seed + + def main() -> None: # load run card and metadata. Use the RunCardMG7 representation when the # madgraph package is importable; gridpacks are meant to be portable, so @@ -69,6 +76,12 @@ def main() -> None: # parse command line arguments parser = argparse.ArgumentParser() parser.add_argument("--run_name", type=str, default=run_args["run_name"]) + parser.add_argument( + "--seed", type=int, default=run_args.get("seed", -1), + help="every run is reproducible from its seed; -1 draws a fresh random " + "seed each run instead of fixing one here (still recorded in the " + "run's info.json)" + ) parser.add_argument("--device", type=str, nargs="*") parser.add_argument( "--cpu_thread_pool_size", type=int, default=run_args["cpu_thread_pool_size"] @@ -94,6 +107,7 @@ def main() -> None: parser.add_argument("--cpu_batch_size", type=int, default=gen_args["cpu_batch_size"]) parser.add_argument("--gpu_batch_size", type=int, default=gen_args["gpu_batch_size"]) args = parser.parse_args() + seed = resolve_seed(args.seed) # initialize event directory run_name = args.run_name @@ -173,6 +187,7 @@ def main() -> None: channels=channel_generators, status_file=ms.StatusFile(os.path.join(run_path, "info.json")), config=config, + seed=seed, ) # run generation diff --git a/madgraph/iolibs/template_files/mg7/madevent.py b/madgraph/iolibs/template_files/mg7/madevent.py index 94d0a6da6..76eb7b138 100644 --- a/madgraph/iolibs/template_files/mg7/madevent.py +++ b/madgraph/iolibs/template_files/mg7/madevent.py @@ -115,6 +115,14 @@ def resolve_verbosity(verbosity: str) -> str: return verbosity +def resolve_seed(seed: int) -> int: + """Resolve the run_card "seed": -1 draws a fresh 64-bit seed via + os.urandom, any other value is used as-is.""" + if seed == -1: + return int.from_bytes(os.urandom(8), "big") + return seed + + def resolve_cppauto_backend(build_path: str) -> str: """Ask the matrix-element Makefile to resolve ``cppauto``. @@ -230,6 +238,10 @@ def __init__(self): def load_cards(self) -> None: self.run_card = RunCardMG7(os.path.join("Cards", "run_card.toml")) + # Resolved once so every generator built during this run shares the same + # seed; the concrete value (even if randomly drawn) is recorded in each + # generator's info.json. + self.run_seed = resolve_seed(self.run_card["run"]["seed"]) self.param_card_path = os.path.join("Cards", "param_card.dat") self.param_card = ParamCard(self.param_card_path) with open(os.path.join("SubProcesses", "subprocesses.json")) as f: @@ -269,7 +281,7 @@ def init_context(self) -> None: self.device_types = [] self.devices = [] self.pool_sizes = [] - for i, device_name in enumerate(device_names): + for device_name in device_names: if ":" in device_name: device_type, device_index_str = device_name.split(":") device_index = int(device_index_str) @@ -566,6 +578,7 @@ def build_event_generator(self, phasespaces: list[PhaseSpace]) -> ms.EventGenera channels=channel_generators, status_file=self.status_file, config=self.event_generator_config, + seed=self.run_seed, ) unused_globals = ( set(self.contexts[0].global_names()) - event_generator.used_globals() @@ -576,29 +589,35 @@ def build_event_generator(self, phasespaces: list[PhaseSpace]) -> ms.EventGenera return event_generator def survey_phasespaces( - self, phasespaces: list[PhaseSpace | None] + self, phasespaces: list[PhaseSpace | None], survey_pass: int = 0 ) -> ms.EventGenerator | None: ps_filtered = [ps for ps in phasespaces if ps is not None] if len(ps_filtered) == 0: return None event_generator = self.build_event_generator(ps_filtered) - event_generator.survey() + event_generator.survey(survey_pass) return event_generator def survey(self) -> None: + # survey_pass distinguishes the survey() calls below: "both" mode can + # re-survey a channel carried over unchanged from the multichannel pass + # into the final (simplified) pass, and both passes schedule jobs on the + # same underlying ChannelEventGenerator. The explicit pass index keeps each + # pass's job seeds independent of the other passes' job counts, rather than + # depending on call history. phasespace_mode = self.run_card["phasespace"]["mode"] if phasespace_mode in ["multichannel", "both", "auto"]: self.phasespaces = [ subproc.build_multichannel_phasespace() for subproc in self.subprocesses ] - self.event_generator = self.survey_phasespaces(self.phasespaces) + self.event_generator = self.survey_phasespaces(self.phasespaces, 0) elif phasespace_mode == "flat": self.phasespaces = [ subproc.build_flat_phasespace() for subproc in self.subprocesses ] - self.event_generator = self.survey_phasespaces(self.phasespaces) + self.event_generator = self.survey_phasespaces(self.phasespaces, 0) else: raise ValueError("Unknown phasespace mode") @@ -610,8 +629,8 @@ def survey(self) -> None: variance = 0. count_opt = 0 for status in channel_status[chan_offset:chan_offset + len(ps.channels)]: - mean += status.mean - variance += status.error**2 + mean += status.mean_abs + variance += status.error_abs**2 count_opt += status.count_opt rsd = (variance * count_opt)**0.5 / mean subproc.set_madnis_auto_settings(rsd) @@ -644,7 +663,10 @@ def survey(self) -> None: ps_multi is not ps_both for ps_multi, ps_both in zip(phasespaces_multi, self.phasespaces) ): - self.event_generator = self.survey_phasespaces(self.phasespaces) + # distinct survey_pass: a channel carried over unchanged from the + # multichannel pass (pass 0) into this resurvey must not share its + # seed stream with that earlier pass. + self.event_generator = self.survey_phasespaces(self.phasespaces, 1) def train_madnis(self) -> None: madnis_args = self.run_card["madnis"] @@ -715,6 +737,11 @@ def train_madnis(self) -> None: training_args=training_args, verbosity=verbosity, status_file=self.status_file, + # Reuses the run's resolved seed (also used by build_event_generator()). + # Only the single-channel CPU sample-generation path is currently seeded + # -- buffered training and GPU multi-channel batches are still + # non-deterministic. + seed=self.run_seed, ) madnis_training.train() for phasespace, active_channels in zip( @@ -918,6 +945,9 @@ def build_lhe_meta(self): headers.append(ms.LHEHeader(name="MG5ProcCard", content=f.read())) headers.append(ms.LHEHeader(name="slha", content=param_text)) headers.append(ms.LHEHeader(name="MG7RunCard", content=run_text)) + # The resolved seed (even when the run_card requested a random one via + # seed = -1), so the run can be reproduced from the LHE file alone. + headers.append(ms.LHEHeader(name="MG7Seed", content=str(self.run_seed))) return ms.LHEMeta( beam1_pdg_id=beam_pdgs[0], beam2_pdg_id=beam_pdgs[1], beam1_energy=energies[0], beam2_energy=energies[1], @@ -1678,6 +1708,9 @@ def set_madnis_auto_settings(self, rsd: float): def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace: madnis_args = self.process.run_card["madnis"] + # Shared across all networks below: initialize_globals() derives an + # independent, non-colliding stream per tensor from this one base seed. + seed = self.process.run_seed channels = [] for channel_id, channel in enumerate(phasespace.channels): prefix = f"subproc{self.subproc_id}.channel{channel_id}" @@ -1695,10 +1728,11 @@ def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace: invert_spline=madnis_args["flow_invert_spline"], ) if channel.adaptive_mapping is None: - flow.initialize_globals(self.process.contexts[0]) + flow.initialize_globals(self.process.contexts[0], seed) else: flow.initialize_from_vegas( - self.process.contexts[0], channel.adaptive_mapping.grid_name() + self.process.contexts[0], channel.adaptive_mapping.grid_name(), + seed ) cond_dim += flow_dim @@ -1715,7 +1749,7 @@ def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace: subnet_layers=madnis_args["discrete_layers"], subnet_activation=self.activation(madnis_args["discrete_activation"]), ) - discrete_sym.initialize_globals(self.process.contexts[0]) + discrete_sym.initialize_globals(self.process.contexts[0], seed) cond_dim += perm_count discrete_flavor = channel.discrete_flavor @@ -1729,7 +1763,7 @@ def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace: subnet_layers=madnis_args["discrete_layers"], subnet_activation=self.activation(madnis_args["discrete_activation"]), ) - discrete_flavor.initialize_globals(self.process.contexts[0]) + discrete_flavor.initialize_globals(self.process.contexts[0], seed) channels.append(Channel( phasespace_mapping = channel.phasespace_mapping, @@ -1803,7 +1837,9 @@ def build_cwnet(self, channel_count: int) -> ms.ChannelWeightNetwork: activation=self.activation(madnis_args["cwnet_activation"]), prefix=f"subproc{self.subproc_id}.cwnet", ) - cwnet.initialize_globals(self.process.contexts[0]) + cwnet.initialize_globals( + self.process.contexts[0], self.process.run_seed + ) return cwnet def t_channel_mode(self, name: str) -> ms.PhaseSpaceMapping.TChannelMode: diff --git a/madgraph/iolibs/template_files/mg7/run_card.toml b/madgraph/iolibs/template_files/mg7/run_card.toml index 4e3d07d63..caee8385f 100644 --- a/madgraph/iolibs/template_files/mg7/run_card.toml +++ b/madgraph/iolibs/template_files/mg7/run_card.toml @@ -1,5 +1,11 @@ [run] run_name = %(run.run_name)s +# every run is reproducible: the same seed reproduces the run bit-identically, +# and distinct seeds are statistically independent. -1 draws a fresh random +# seed each run instead of fixing one here; the seed actually used is still +# recorded (the MG7Seed tag in the LHE file, or the info.json status file), so +# the run can be reproduced later. +seed = %(run.seed)s # options: cuda, hip, cpp, cppnone, cppsse4, cppavx2, cpp512y, cpp512z, cppauto devices = %(run.devices)s # options: diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py index 5fa833629..dfcb0640a 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -6463,6 +6463,12 @@ def default_setup(self): # ----------------------------- [run] -------------------------- self.add_toml_param('run', 'run_name', "run", gridpack=True) + self.add_toml_param('run', 'seed', -1, gridpack=True, + comment="every run is reproducible: the same seed reproduces the run " + "bit-identically. -1 draws a fresh random seed each run instead " + "of fixing one here; the seed actually used is still recorded " + "(the MG7Seed tag in the LHE file, or the info.json status " + "file), so the run can be reproduced later") self.add_toml_param('run', 'devices', ["cppnone"], typelist=str, gridpack=True, comment="options: cuda, hip, cpp, cppnone, cppsse4, cppavx2, cpp512y, cpp512z, cppauto") self.add_toml_param('run', 'simd_vector_size', -1, diff --git a/madspace/.clang-format-ignore b/madspace/.clang-format-ignore index 2ff5b5bfb..9343bc477 100644 --- a/madspace/.clang-format-ignore +++ b/madspace/.clang-format-ignore @@ -7,3 +7,5 @@ src/cpu/runtime_mixin.inc src/gpu/runtime_mixin.inc src/cpu/runtime_backward_mixin.inc src/gpu/runtime_backward_mixin.inc +include/madspace/mixmax/mixmax.hpp +include/madspace/mixmax/mixmax_skip_N17.c diff --git a/madspace/CMakeLists.txt b/madspace/CMakeLists.txt index abd7de730..b4b389b68 100644 --- a/madspace/CMakeLists.txt +++ b/madspace/CMakeLists.txt @@ -81,7 +81,6 @@ if(COMPILE_CUDA) endif() if(COMPILE_HIP) find_package(rocblas REQUIRED) - find_package(rocrand REQUIRED) find_package(rocthrust REQUIRED) find_package(hipcub REQUIRED) endif() @@ -215,6 +214,7 @@ add_library( src/driver/backend.cpp src/driver/lhe_output.cpp src/driver/madnis_training.cpp + src/driver/random.cpp src/compgraphs/type.cpp src/compgraphs/function.cpp src/compgraphs/instruction.cpp @@ -267,6 +267,7 @@ add_library( include/madspace/driver/lhe_output.hpp include/madspace/driver/logger.hpp include/madspace/driver/madnis_training.hpp + include/madspace/driver/random.hpp include/madspace/compgraphs.hpp include/madspace/compgraphs/type.hpp include/madspace/compgraphs/function.hpp @@ -305,6 +306,8 @@ add_library( include/madspace/phasespace/matrix_element.hpp include/madspace/phasespace/cross_section.hpp include/madspace/phasespace/scale.hpp + include/madspace/mixmax/mixmax.hpp + include/madspace/mixmax/mixmax_skip_N17.c ) #target_compile_options(madspace PRIVATE -Wall -Wextra -Wpedantic -Werror) @@ -377,6 +380,7 @@ if(COMPILE_CUDA) src/gpu/tensor.cuh src/gpu/device.hpp src/gpu/gpu_abstraction.cuh + src/gpu/random.cuh ${MADSPACE_GENERATED_DIR}/gpu/runtime_mixin.inc ${MADSPACE_GENERATED_DIR}/gpu/runtime_backward_mixin.inc ${KERNEL_SRC} @@ -386,7 +390,6 @@ if(COMPILE_CUDA) set_property(TARGET madspace_cuda PROPERTY CUDA_STANDARD 20) target_link_libraries(madspace_cuda PRIVATE madspace) target_link_libraries(madspace_cuda PRIVATE CUDA::cublas) - target_link_libraries(madspace_cuda PRIVATE CUDA::curand) install(TARGETS madspace_cuda DESTINATION madspace/lib) endif() @@ -401,6 +404,7 @@ if(COMPILE_HIP) src/gpu/tensor.cuh src/gpu/device.hpp src/gpu/gpu_abstraction.cuh + src/gpu/random.cuh ${MADSPACE_GENERATED_DIR}/gpu/runtime_mixin.inc ${MADSPACE_GENERATED_DIR}/gpu/runtime_backward_mixin.inc ${KERNEL_SRC} @@ -412,7 +416,6 @@ if(COMPILE_HIP) target_link_libraries(madspace_hip PRIVATE madspace) target_link_libraries(madspace_hip PRIVATE hip::device) target_link_libraries(madspace_hip PRIVATE roc::rocblas) - target_link_libraries(madspace_hip PRIVATE roc::rocrand) target_link_libraries(madspace_hip PRIVATE roc::rocthrust) target_link_libraries(madspace_hip PRIVATE hip::hipcub) install(TARGETS madspace_hip DESTINATION madspace/lib) diff --git a/madspace/include/madspace/compgraphs/instruction.hpp b/madspace/include/madspace/compgraphs/instruction.hpp index 8cdc7c8d7..0969985b0 100644 --- a/madspace/include/madspace/compgraphs/instruction.hpp +++ b/madspace/include/madspace/compgraphs/instruction.hpp @@ -28,6 +28,9 @@ class Instruction { const std::string& name() const { return _name; } int opcode() const { return _opcode; } bool differentiable() const { return _differentiable; } + // draws from a shared RNG stream: never CSE'd, always scheduled on the main + // GPU stream (see FunctionBuilder::instruction) + virtual bool is_random() const { return false; } protected: void check_arg_count(const ValueVec& args, std::size_t count) const; @@ -211,6 +214,7 @@ class RandomInstruction : public Instruction { RandomInstruction(int opcode, bool differentiable) : Instruction("random", opcode, differentiable) {} TypeVec signature(const ValueVec& args) const override; + bool is_random() const override { return true; } }; class RandomIntInstruction : public Instruction { @@ -218,6 +222,7 @@ class RandomIntInstruction : public Instruction { RandomIntInstruction(int opcode, bool differentiable) : Instruction("random_int", opcode, differentiable) {} TypeVec signature(const ValueVec& args) const override; + bool is_random() const override { return true; } }; class UnweightInstruction : public Instruction { @@ -225,6 +230,7 @@ class UnweightInstruction : public Instruction { UnweightInstruction(int opcode, bool differentiable) : Instruction("unweight", opcode, differentiable) {} TypeVec signature(const ValueVec& args) const override; + bool is_random() const override { return true; } }; class MatrixElementInstruction : public Instruction { diff --git a/madspace/include/madspace/driver/backend.hpp b/madspace/include/madspace/driver/backend.hpp index 792d5828b..803449808 100644 --- a/madspace/include/madspace/driver/backend.hpp +++ b/madspace/include/madspace/driver/backend.hpp @@ -1,7 +1,11 @@ #pragma once +#include +#include + #include "madspace/compgraphs.hpp" #include "madspace/driver/context.hpp" +#include "madspace/driver/random.hpp" #include "madspace/driver/tensor.hpp" namespace madspace { @@ -19,6 +23,8 @@ class Runtime { const std::vector& eval_grad, bool return_contiguous_grads = false ) = 0; + virtual void set_seed(DerivedSeed seed) = 0; + friend std::unique_ptr build_runtime(const Function& function, ContextPtr context, bool concurrent); diff --git a/madspace/include/madspace/driver/channel_generator.hpp b/madspace/include/madspace/driver/channel_generator.hpp index 02a5c04ab..02a455237 100644 --- a/madspace/include/madspace/driver/channel_generator.hpp +++ b/madspace/include/madspace/driver/channel_generator.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include #include #include @@ -13,6 +13,7 @@ #include "madspace/driver/discrete_optimizer.hpp" #include "madspace/driver/generator_data.hpp" #include "madspace/driver/io.hpp" +#include "madspace/driver/random.hpp" #include "madspace/driver/vegas_optimizer.hpp" #include "madspace/phasespace.hpp" @@ -45,6 +46,7 @@ class ChannelEventGenerator { const GeneratorStatus& status() const { return _status; } const RunningIntegral& cross_section() const { return _cross_section; } + const RunningIntegral& abs_cross_section() const { return _abs_cross_section; } const std::vector& histograms() const { return _histograms; } EventFile& event_file() { return _event_file; } EventFile& weight_file() { return _weight_file; } @@ -53,7 +55,9 @@ class ChannelEventGenerator { bool needs_optimization() const { return (_vegas_optimizer || _discrete_optimizer) && !_status.optimized; } - void set_target_count(double target_count) { _status.count_target = target_count; } + void set_target_count(std::size_t target_count) { + _status.count_target = target_count; + } const std::unordered_set& used_globals() const { return _used_globals; } @@ -61,11 +65,24 @@ class ChannelEventGenerator { int particle_layout_extra_flags() const { return _particle_layout_extra_flags; } const DataLayout& event_file_layout() const { return _event_file_layout; } - void unweight_file(std::mt19937& rand_gen); + void unweight_file(MixMaxRandom& rand_gen); void integrate(const GeneratorBatchJob& job); void optimize_vegas(const GeneratorBatchJob& job); double channel_weight_sum(std::size_t event_count); - void start_job(GeneratorBatchJob& job, ResultQueue& result_queue); + void start_job( + GeneratorBatchJob& job, + ResultQueue& result_queue, + std::optional seed, + bool is_survey, + std::size_t survey_pass + ); + // Snapshots max_weight into job.max_weight, at a fixed point independent of + // when submit_unweight_job() actually dispatches it. + void prepare_unweight_job(GeneratorBatchJob& job) const; + // Submits job's unweighting to the thread pool of its own generation context + // (required on GPU: unweighting does a device-to-host copy). + void submit_unweight_job(GeneratorBatchJob& job, ResultQueue& result_queue); + // prepare_unweight_job() + submit_unweight_job() combined. void start_unweight_job(GeneratorBatchJob& job, ResultQueue& result_queue); std::size_t next_vegas_batch_size(); void clear_events(); @@ -131,8 +148,19 @@ class ChannelEventGenerator { Function _unweighter_function; std::optional _histogram_function; RunningIntegral _cross_section; + RunningIntegral _abs_cross_section; double _max_weight = 0.; + // Monotonic per-job counters keying each job's deterministic random stream; + // never reset. Separate for survey/generate so neither depends on the other. + std::size_t _survey_rng_seq = 0; + std::size_t _generate_rng_seq = 0; + // Progress of unweight_file()'s final pass over _weight_file: _unweighted_count + // is the index up to which it's been scanned under the *current* _max_weight, + // and _unweighted_accept_count the number of those accepted. Both reset to 0 + // whenever _max_weight changes (forcing a full rescan), since a changed + // max_weight invalidates every prior accept/reject decision. std::size_t _unweighted_count = 0; + std::size_t _unweighted_accept_count = 0; std::size_t _iters_without_improvement = 0; double _best_rsd = std::numeric_limits::max(); std::vector _large_weights; diff --git a/madspace/include/madspace/driver/context.hpp b/madspace/include/madspace/driver/context.hpp index c58985541..626bbfe5d 100644 --- a/madspace/include/madspace/driver/context.hpp +++ b/madspace/include/madspace/driver/context.hpp @@ -1,6 +1,7 @@ #pragma once #include +#include #include #include "madspace/compgraphs.hpp" @@ -48,35 +49,43 @@ class MatrixElementApi { std::size_t index() const { return _index; } const std::string& file_name() const { return _file_name; } std::vector supported_inputs() const { - bool const* data; int count; + bool const* data; + int count; check_umami_status(_supported_inputs(&data, &count)); std::vector result(UMAMI_INPUT_KEY_COUNT, false); - for (int i = 0; i < count && i < UMAMI_INPUT_KEY_COUNT; ++i) + for (int i = 0; i < count && i < UMAMI_INPUT_KEY_COUNT; ++i) { result[i] = data[i]; + } return result; } std::vector required_inputs() const { - bool const* data; int count; + bool const* data; + int count; check_umami_status(_required_inputs(&data, &count)); std::vector supported = supported_inputs(); for (int i = 0; i < count && i < UMAMI_INPUT_KEY_COUNT; ++i) { if (data[i] && !supported[i]) { - throw_error(std::format( - "input key {} is reported as required but not as supported", i - )); + throw_error( + std::format( + "input key {} is reported as required but not as supported", i + ) + ); } } std::vector result(UMAMI_INPUT_KEY_COUNT, false); - for (int i = 0; i < count && i < UMAMI_INPUT_KEY_COUNT; ++i) + for (int i = 0; i < count && i < UMAMI_INPUT_KEY_COUNT; ++i) { result[i] = data[i]; + } return result; } std::vector supported_outputs() const { - bool const* data; int count; + bool const* data; + int count; check_umami_status(_supported_outputs(&data, &count)); std::vector result(UMAMI_OUTPUT_KEY_COUNT, false); - for (int i = 0; i < count && i < UMAMI_OUTPUT_KEY_COUNT; ++i) + for (int i = 0; i < count && i < UMAMI_OUTPUT_KEY_COUNT; ++i) { result[i] = data[i]; + } return result; } @@ -142,18 +151,18 @@ class Context { public: Context(int thread_count = -1) : _device(cpu_device()), - _thread_pool(std::make_unique(thread_count)) - { - reset_cache(); - } + _thread_pool(std::make_unique(thread_count)), + _tensor_cache(global_resource(tensor_cache_resource_name, []() { + return TensorVec{}; + })) {} Context(DevicePtr device, int thread_count = -1) : _device(device), - _thread_pool(std::make_unique(thread_count)) - { - reset_cache(); - } - Context(Context&&) = default; - Context& operator=(Context&&) = default; + _thread_pool(std::make_unique(thread_count)), + _tensor_cache(global_resource(tensor_cache_resource_name, []() { + return TensorVec{}; + })) {} + Context(Context&&) = delete; + Context& operator=(Context&&) = delete; Context(const Context&) = delete; Context& operator=(const Context&) = delete; const MatrixElementApi& @@ -177,20 +186,55 @@ class Context { void load_globals(const std::string& dir); DevicePtr device() { return _device; } ThreadPool& thread_pool() { return *_thread_pool; } + std::size_t unique_seed_index() { return _seed_index++; } Tensor cached_tensor(std::size_t size); void reset_cache() { - _tensor_cache = ThreadResource( - thread_pool(), []() { return TensorVec{}; } - ); + _tensor_cache = ThreadResource(thread_pool(), []() { + return TensorVec{}; + }); + } + // Not thread-safe: callers must acquire the reference once during + // single-threaded initialization (e.g. a Runtime constructor) and reuse it. + template + ThreadResource& global_resource( + const std::string& name, + std::function constructor, + std::optional> destructor = std::nullopt + ) { + auto search = _resources.find(name); + if (search == _resources.end()) { + auto res = std::make_shared>( + thread_pool(), constructor, destructor + ); + _resources.emplace( + name, + std::pair>( + std::type_index(typeid(T)), res + ) + ); + return *res; + } else { + auto& [tid, res] = search->second; + if (std::type_index(typeid(T)) != tid) { + throw std::runtime_error( + std::format("incompatible resource type for '{}'", name) + ); + } + return *std::static_pointer_cast>(res); + } } private: + static constexpr const char* tensor_cache_resource_name = "__tensor_cache"; DevicePtr _device; std::unique_ptr _thread_pool; std::unordered_map> _globals; std::vector> _matrix_elements; std::vector _param_card_paths; - ThreadResource _tensor_cache; + std::size_t _seed_index = 0; + std::unordered_map>> + _resources; + ThreadResource& _tensor_cache; }; using ContextPtr = std::shared_ptr; diff --git a/madspace/include/madspace/driver/event_generator.hpp b/madspace/include/madspace/driver/event_generator.hpp index 743c2affd..2255eb8cf 100644 --- a/madspace/include/madspace/driver/event_generator.hpp +++ b/madspace/include/madspace/driver/event_generator.hpp @@ -1,9 +1,9 @@ #pragma once #include +#include #include -#include -#include +#include #include #include @@ -32,6 +32,7 @@ class EventGenerator { EventGenerator( const std::vector& contexts, const std::vector>& channels, + std::uint64_t seed, std::shared_ptr status_file = nullptr, const GeneratorConfig& config = default_config ); @@ -39,7 +40,9 @@ class EventGenerator { EventGenerator& operator=(EventGenerator&&) = default; EventGenerator(const EventGenerator&) = delete; EventGenerator& operator=(const EventGenerator&) = delete; - void survey(); + // `survey_pass` salts job seeds so repeated survey() calls on the same + // channel (e.g. re-survey after simplification) don't share a seed stream. + void survey(std::size_t survey_pass = 0); void generate(); void combine_to_compact_npy(const std::string& file_name); void combine_to_lhe_npy(const std::string& file_name, LHECompleter& lhe_completer); @@ -74,14 +77,62 @@ class EventGenerator { GeneratorStatus _status; std::vector _contexts; std::unordered_map _running_jobs; - std::vector _ready_jobs; + std::vector _ready_jobs; std::size_t _job_id; std::vector _channel_job_counts; std::vector _channel_optimizing; std::vector _channel_integral_fractions; std::vector _context_job_counts; + // True while a channel has a steady-state batch dispatched but not yet fully + // committed; keeps next_batch_event_count() from double-counting in-flight work. + std::vector _channel_batch_pending; + // True once a channel's current generation ReadyJob has had its full event count + // dispatched (batch_event_count reached zero). Needed alongside + // channel_job_count == 0 before finish_channel_job() clears _channel_batch_pending + // -- dispatch now happens incrementally, so channel_job_count can transiently hit + // zero mid-batch, between one sub-job's commit and the next one being dispatched. + std::vector _channel_batch_dispatch_done; + // Round-robin position into _ready_jobs for generation-batch dispatch, persisted + // across start_jobs() calls so multiple channels' batches interleave (one device + // batch at a time) instead of one channel's batch draining before the next is + // touched. + std::size_t _ready_job_rr_cursor = 0; + // Per-channel commit ordering, analogous to _ready_gen/_commit_cursor but ordered + // per channel instead of globally. _channel_gen_order holds a channel's dispatched + // job ids in dispatch order and _channel_ready_gen the ones that have completed; + // the front of the order deque is the next commit due. An explicit deque rather + // than a "next id" counter because a channel's job ids are not contiguous: dispatch + // round-robins between channels, so consecutive ids belong to different channels. + std::vector> _channel_gen_order; + std::vector> _channel_ready_gen; + // Same, for a job's unweight-stage completion (tracked separately since it's a + // distinct completion event). Order is appended at generate-commit time, which is + // also when the unweight stage is queued. + std::vector> _channel_unweight_order; + std::vector> _channel_unweight_ready; + // generate() only: per-context queue of job ids awaiting unweight-stage + // dispatch, drained with priority by start_jobs(). + std::vector> _context_unweight_queue; ResultQueue _result_queue; + // Base seed for reproducible event generation. + std::uint64_t _seed; + + // unweight_all() may run more than once per generate() (a channel's target can + // grow after it looked done, un-finishing it and triggering another round). + // Salted by this counter so repeated calls don't replay the same stream. + std::size_t _unweight_call_index = 0; + + // Scheduling context for the running survey()/generate() call, read by + // start_jobs() to derive job seeds. + bool _survey_job = false; + std::size_t _survey_pass = 0; + + // Generate completions are committed in ascending job id, with + // _commit_cursor as the next id due. + std::set _ready_gen; + std::size_t _commit_cursor = 0; + std::chrono::time_point _start_time; std::size_t _start_cpu_microsec; std::chrono::time_point _last_print_time; @@ -90,8 +141,14 @@ class EventGenerator { std::shared_ptr _status_file; std::unordered_map _timing_data; - bool start_jobs(); + void commit_generate_job(GeneratorBatchJob& job); + void commit_unweight_job(GeneratorBatchJob& job); + void finish_channel_job(const GeneratorBatchJob& job); + std::size_t next_batch_event_count(std::size_t channel_index) const; + std::size_t start_jobs(); void update_integral(); + void update_integral_status(); + void update_integral_fractions(); void update_counts(); void reset_start_time(); void add_timing_data(const std::string& key); @@ -100,14 +157,15 @@ class EventGenerator { void read_and_combine( std::vector& channel_data, EventBuffer& buffer, - double norm_factor + double norm_factor, + MixMaxRandom& rand_gen ); void fill_lhe_event( LHECompleter& lhe_completer, LHEEvent& lhe_event, EventBuffer& buffer, std::size_t event_index, - std::mt19937& rand_gen + MixMaxRandom& rand_gen ); void init_status(const std::string& status); diff --git a/madspace/include/madspace/driver/generator_data.hpp b/madspace/include/madspace/driver/generator_data.hpp index bedf6bb2c..442ff954e 100644 --- a/madspace/include/madspace/driver/generator_data.hpp +++ b/madspace/include/madspace/driver/generator_data.hpp @@ -2,6 +2,7 @@ #include #include +#include #include #include @@ -63,20 +64,38 @@ struct GeneratorConfig { int combine_thread_count = -1; double cut_efficiency_threshold = 0.7; std::size_t max_cut_repetitions = 100; + double finish_remaining_fraction = 0.05; + double max_batch_fraction = 0.6; + double batch_overshoot_sigma = 1.0; }; +// Determine number of events to be scheduled in the next step. If close to finishing, +// return remaining count / efficiency. Otherwise select a count that is safely below +// the target using the integration uncertainty, and impose an upper limit which +// fration of the target events can be scheduled in one go. +std::size_t compute_generation_batch_event_count( + std::size_t count_target, + double count_unweighted, + std::size_t count_opt, + std::size_t abs_cross_section_count, + double abs_cross_section_rel_error, + const GeneratorConfig& config +); + struct GeneratorStatus { std::size_t subprocess; std::string name; double mean; double error; + double mean_abs; // E[|w|] + double error_abs; // error on E[|w|] double rel_std_dev; std::size_t count; std::size_t count_opt; std::size_t count_after_cuts; std::size_t count_after_cuts_opt; double count_unweighted; - double count_target; + std::size_t count_target; std::size_t iterations; bool optimized; bool done; @@ -90,10 +109,34 @@ struct Histogram { std::vector bin_errors; }; +// Lightweight pending-work entry for EventGenerator::_ready_jobs, before start_jobs() +// has decided which context/how many sub-jobs to create from it. Kept separate from +// GeneratorBatchJob so a queue of pending batches doesn't carry the weight of every +// dispatched job's tensors and RNG bookkeeping. +struct ReadyJob { + std::size_t channel_index; + bool unweight; + // VEGAS batch: fixed size, start_jobs() dispatches it atomically in one go. + // Generation batch: events not yet dispatched -- start_jobs() decrements this in + // place as it creates sub-jobs, one device batch at a time, round-robining with + // other channels' ReadyJobs over however many calls it takes to reach zero. + std::size_t batch_event_count; + bool is_vegas_batch = false; +}; + struct GeneratorBatchJob { std::size_t channel_index; bool unweight; - std::size_t vegas_batch_size; + // Copied from the originating ReadyJob at dispatch time. For a VEGAS batch, this + // is the batch's fixed total, read by start_job()'s shrink-to-fit and by the + // done_event_count accounting in survey()/survey_deterministic(). For a + // generation batch it isn't read after dispatch -- generation sub-jobs always + // request a full device batch (see start_job()). + std::size_t batch_event_count; + // Total sub-jobs the batch was split into; only meaningful for VEGAS batches + // (see commit_generate_job()'s clear_events trigger). Generation batches are + // dispatched one sub-job per start_jobs() visit, so this isn't a full-batch count + // for them and isn't read. std::size_t split_job_count; Tensor weights; TensorVec events; @@ -104,6 +147,17 @@ struct GeneratorBatchJob { std::size_t context_index; std::size_t job_id; double max_weight; + // Top-level seed plus job identity, used to derive this job's DerivedSeed(s) + // independently for generate vs unweight (see start_job()/submit_unweight_job()). + std::optional rng_seed; + // Per-channel dispatch sequence, assigned once at start_job() time. + std::size_t rng_job_index = 0; + bool rng_is_survey = false; + std::size_t rng_survey_pass = 0; + // True for a VEGAS-grid-optimization batch, dispatched atomically by start_jobs() + // and shrunk to fit by start_job(). False for a steady-state generation batch, + // dispatched incrementally as device-sized sub-jobs (see ReadyJob). + bool is_vegas_batch = false; }; void to_json(nlohmann::json& j, const GeneratorStatus& status); diff --git a/madspace/include/madspace/driver/lhe_output.hpp b/madspace/include/madspace/driver/lhe_output.hpp index 4ec62db5d..70b8b7f67 100644 --- a/madspace/include/madspace/driver/lhe_output.hpp +++ b/madspace/include/madspace/driver/lhe_output.hpp @@ -1,13 +1,13 @@ #pragma once #include -#include #include #include #include #include +#include "madspace/driver/random.hpp" #include "madspace/driver/thread_pool.hpp" #include "madspace/phasespace/topology.hpp" #include "madspace/util.hpp" @@ -89,7 +89,7 @@ class LHECompleter { int color_index, int flavor_index, int helicity_index, - std::mt19937& rand_gen + MixMaxRandom& rand_gen ); std::size_t max_particle_count() const { return _max_particle_count; } void save(const std::string& file) const; diff --git a/madspace/include/madspace/driver/madnis_training.hpp b/madspace/include/madspace/driver/madnis_training.hpp index 09f3816e4..ba4503105 100644 --- a/madspace/include/madspace/driver/madnis_training.hpp +++ b/madspace/include/madspace/driver/madnis_training.hpp @@ -1,7 +1,9 @@ #pragma once #include +#include #include +#include #include "madspace/compgraphs.hpp" #include "madspace/driver/adam_optimizer.hpp" @@ -51,7 +53,12 @@ class MadnisTraining { ContextPtr optimizer_context, const Config& config, const std::vector>& integrands, - const std::optional& cwnet + const std::optional& cwnet, + std::optional seed = std::nullopt, + // Offset added to this subprocess's local channel indices when deriving + // seeds, so MultiMadnisTraining's subprocesses (which all share the same + // top-level seed) get non-overlapping DerivedSeed channel_index ranges. + std::size_t channel_index_offset = 0 ); const Config& config() const { return _config; } void train_step(std::size_t batch_index); @@ -94,6 +101,9 @@ class MadnisTraining { struct SampleJob { SampleBatch samples; SampleBatch unweighted_samples; + // dispatch sequence used to commit in order: per-channel for single-channel + // jobs, global (see _multi_job_next_dispatch_seq) for multi-channel jobs + std::size_t dispatch_seq = 0; }; struct ChannelData { std::size_t index; @@ -106,6 +116,13 @@ class MadnisTraining { RuntimePtr generator_runtime = nullptr; RuntimePtr unweighter_runtime = nullptr; SampleBatch buffer; + // commit-ordering state for single-channel generator jobs (see + // process_job_results) + std::size_t next_dispatch_seq = 0; + std::size_t commit_cursor = 0; + std::unordered_map ready_job_ids; + // samples staged here, flushed into buffer at the start of the next round + std::vector pending_buffer_samples; }; inline static std::function _abort_check_function = [] {}; @@ -113,6 +130,9 @@ class MadnisTraining { void build_runtimes_and_optimizer(); std::vector compute_channel_sizes(); void start_generator_jobs(const std::vector& channel_fractions); + void maybe_start_generator_jobs( + const std::vector& channel_fractions, bool is_online_attempt + ); TensorVec permute_tensors(const TensorVec& tensors) const; void start_single_job(std::size_t channel_index, std::size_t batch_size); void start_multi_job(const std::vector batch_sizes); @@ -121,6 +141,7 @@ class MadnisTraining { TensorVec build_online_training_batch(const std::vector& counts); TensorVec build_buffered_training_batch(const std::vector& counts); void process_job_results(const std::vector& job_ids); + void process_all_jobs(); void buffer_store(ChannelData& channel, SampleBatch& samples); void update_history( const TensorVec& results, @@ -158,6 +179,15 @@ class MadnisTraining { std::vector _arg_permutation; bool _buffer_ready = false; std::vector _active_flavors_count; + std::optional _seed; + std::size_t _channel_index_offset; + // sequence for seeding build_buffered_training_batch's BatchSampler::run() call + std::size_t _buffered_batch_seq = 0; + // commit-ordering state for multi-channel (GPU) generator jobs (see + // process_job_results): global, since one job spans multiple channels at once + std::size_t _multi_job_next_dispatch_seq = 0; + std::size_t _multi_job_commit_cursor = 0; + std::unordered_map _multi_job_ready_job_ids; std::size_t _diverged_batch_count = 0; }; @@ -174,7 +204,11 @@ class MultiMadnisTraining { ContextPtr optimizer_context, const std::vector& training_args, Verbosity verbosity = Verbosity::log, - std::shared_ptr status_file = nullptr + std::shared_ptr status_file = nullptr, + // Reuses the run_card's own seed (also used by build_event_generator()). + // Each subprocess's MadnisTraining gets a channel_index_offset slice of + // this seed's stream so their derived seeds don't collide. + std::optional seed = std::nullopt ); void train(); nested_vector2 active_channels() const { return _active_channels; } @@ -198,6 +232,7 @@ class MultiMadnisTraining { ContextPtr _optimizer_context; std::vector _training_args; Verbosity _verbosity; + std::optional _seed; nested_vector2 _active_channels; std::chrono::time_point _start_time; std::size_t _start_cpu_microsec; diff --git a/madspace/include/madspace/driver/random.hpp b/madspace/include/madspace/driver/random.hpp new file mode 100644 index 000000000..0db5b7763 --- /dev/null +++ b/madspace/include/madspace/driver/random.hpp @@ -0,0 +1,89 @@ +#pragma once + +#include +#include +#include +#include + +#include "madspace/mixmax/mixmax.hpp" + +namespace madspace { + +struct DerivedSeed { + static constexpr std::size_t max_channel_count = 1ULL << 12; + static constexpr std::size_t max_job_count = 1ULL << 32; + static constexpr std::size_t max_stream_count = 1ULL << 16; + enum SeedType { + none, + first_survey_generate, + first_survey_unweight, + second_survey_generate, + second_survey_unweight, + generator_generate, + generator_unweight, + combine_select, + lhe_complete, + unweight_pass, + madnis_generate, + madnis_unweight, + madnis_sample_buffer, + global_init, + }; + + std::array seed_parts; + + DerivedSeed( + const std::optional& seed = std::nullopt, + SeedType seed_type = none, + std::size_t job_index = 0, + std::size_t channel_index = 0, + std::size_t stream_index = 0 + ); +}; + +// MIXMAX vielbein with the 64 global-run-seed bits applied (seed_parts[0..1], i.e. +// effective seed bits [64,128)). These do not change within a run, so each RNG caches +// this and only recomputes it if the run seed actually differs (e.g. when madspace is +// driven as a library). The per-call bits [0,64) are applied on top. +struct RunSeedSkip { + std::array state; + std::uint64_t run_seed = 0; + bool valid = false; + + // refresh `state` for `run_seed_hi:run_seed_lo` if it is not already cached + void update(std::uint32_t run_seed_hi, std::uint32_t run_seed_lo) { + std::uint64_t key = (std::uint64_t(run_seed_hi) << 32) | run_seed_lo; + if (valid && key == run_seed) { + return; + } + mixmax_engine::run_seed_prefix(state.data(), run_seed_lo, run_seed_hi); + run_seed = key; + valid = true; + } +}; + +class MixMaxRandom { +public: + MixMaxRandom() : MixMaxRandom(DerivedSeed()) {} + MixMaxRandom(DerivedSeed seed) { apply_seed(seed); } + explicit MixMaxRandom(std::uint64_t seed) : MixMaxRandom(DerivedSeed(seed)) {} + void set_seed(DerivedSeed seed) { apply_seed(seed); } + double generate_double() { return _mixmax.flat(); } + std::size_t generate_int(std::size_t max_int) { + return std::min(_mixmax.flat() * max_int, max_int - 1); + } + +private: + void apply_seed(const DerivedSeed& seed) { + _run_skip.update(seed.seed_parts[0], seed.seed_parts[1]); + // remaining bits: seed_parts[3] at [0,32), seed_parts[2] at [32,64) + _mixmax.seed_from_state( + _run_skip.state.data(), seed.seed_parts[3], seed.seed_parts[2] + ); + } + + mixmax_engine _mixmax; + RunSeedSkip _run_skip; +}; + +} // namespace madspace diff --git a/madspace/include/madspace/mixmax/mixmax.hpp b/madspace/include/madspace/mixmax/mixmax.hpp new file mode 100644 index 000000000..ad07fc4e4 --- /dev/null +++ b/madspace/include/madspace/mixmax/mixmax.hpp @@ -0,0 +1,435 @@ +/* + * mixmax.hpp + * + * C++ implementation of the MIXMAX random number generator. + * + * Copyright (2008-2023) by Konstantin Savvidy. + * + * Free to use, academic or commercial. Do not redistribute without permission. + * + * G.K.Savvidy and N.G.Ter-Arutyunian, + * On the Monte Carlo simulation of physical systems, + * J.Comput.Phy 97, 566 (1991); + * Preprint EPI-865-16-86, Yerevan, Jan. 1986 + * + * K.Savvidy + * The MIXMAX random number generator + * Comp. Phy Commun. 196 (2015), pp 161-165 + * http://dx.doi.org/10.1016/j.cpc.2015.06.003 + * + * K.Savvidy and G.Savvidy + * Spectrum and Entropy of C-system MIXMAX random number generator + * Chaos, Solitons & Fractals, Volume 91, (2016) pp. 33-38 + * http://dx.doi.org/10.1016/j.chao2016.05.003 + * + */ + +#ifndef __MIXMAX_H +#define __MIXMAX_H + +#include + +#if (defined(__CUDACC__) || defined (__CUDA_ARCH__) || defined(__HIPCC__)) +#define _dev __host__ __device__ +#else +#define _dev +#include +#include +#endif + +#if (defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__)) +#define _constmem __constant__ +#else +#define _constmem +#endif + +class alignas(128) mixmax_engine +{ +static const int N = 17; +static constexpr int BITS=61; +static constexpr uint64_t M61=2305843009213693951ULL; +static constexpr uint64_t MERSBASE=M61; +static constexpr double INV_MERSBASE=(0.43368086899420177360298E-18); + + + static constexpr long long int SPECIAL = 0; + static constexpr long long int SPECIALMUL= 36; + // Note the potential for confusion... + +private: // state + uint64_t V[N] = {0}; // init with all zero's - not the same as seeding with seed=0! + uint64_t sumtot = {0}; + int counter = N; + +public: +using result_type = uint64_t; // should it be double? + static constexpr uint64_t min() {return 0;} + static constexpr uint64_t max() {return M61;} + + _dev inline uint64_t get_next() ; + _dev inline double flat(); + + + _dev inline mixmax_engine(); // Constructor, no seeds + _dev inline mixmax_engine(uint32_t clusterID, uint32_t machineID, uint32_t runID, uint32_t streamID ); // Constructor with four 32-bit seeds + _dev inline mixmax_engine(uint64_t seedval){seed_uniquestream( 0, 0, (uint32_t)(seedval>>32), (uint32_t)seedval );}; // Constructor with one 64bit seed + _dev inline void seed(uint64_t seedval){seed_uniquestream( 0, 0, (uint32_t)(seedval>>32), (uint32_t)seedval );} // seed with one 64-bit seed + _dev inline void seed_uniquestream( uint32_t clusterID, uint32_t machineID, uint32_t runID, uint32_t streamID ); + + _dev inline mixmax_engine Branch(); + _dev inline void BranchInplace(); + + + _dev inline mixmax_engine& operator=(const mixmax_engine& other ); + _dev inline operator double() { return flat(); } + _dev inline operator float() { return float( flat() ); } + _dev inline operator unsigned int() { return static_cast(get_next()); } + + _dev inline uint64_t operator()() + { + return get_next(); + } + + // --- substream seeding building blocks, used by the GPU seeder --- + static constexpr int state_size = N; + static constexpr uint64_t mod_prime = M61; + // Y <- (sum_j coef[j] A^j) Y, i.e. apply the polynomial-in-A `coef` (state_size terms). + _dev static inline void apply_poly(uint64_t* Y, const uint64_t* coef); + // Y <- P_r(A) Y: apply the skip operator for seed bit r (row r of the skip matrix). + // The operators for different r commute, so bits may be applied in any order/grouping. + _dev static inline void apply_skip_bit(uint64_t* Y, int r); + // field ops mod mod_prime: acc + a*b, a + b, and Y <- A*Y (returns new sum(Y)) + _dev static inline uint64_t mod_mul_add(uint64_t acc, uint64_t a, uint64_t b) { + return fmodmulM61(acc, a, b); + } + _dev static inline uint64_t mod_add(uint64_t a, uint64_t b) { return modadd(a, b); } + _dev static inline uint64_t advance_state(uint64_t* Y, uint64_t sum) { + return iterate_raw_vec(Y, sum); + } + // raw N-vector access, for in-place skipping + _dev uint64_t* state_data() { return V; } + _dev const uint64_t* state_data() const { return V; } + // recompute sumtot from V and arm the engine (counter = 1), as seed_uniquestream does + _dev void finalize_skipped_state() { + uint64_t s = 0; + for (int i = 0; i < N; i++) { s = modadd(s, V[i]); } + sumtot = s; + counter = 1; + } + // vielbein with the two high seed words applied: `high0` at effective bit positions + // [64,96), `high1` at [96,128). This part of the seed is the global run seed and does + // not change within a run, so callers cache the result. + _dev static void run_seed_prefix(uint64_t* out, uint32_t high0, uint32_t high1) { + for (int i = 0; i < N; i++) { out[i] = 0; } + out[0] = 1; + for (int b = 0; b < 32; b++) { if ((high0 >> b) & 1u) { apply_skip_bit(out, 64 + b); } } + for (int b = 0; b < 32; b++) { if ((high1 >> b) & 1u) { apply_skip_bit(out, 96 + b); } } + } + // seed from `prefix` (a state with some seed bits pre-applied), then apply the low + // seed words: `low0` at effective bit positions [0,32), `low1` at [32,64). Equivalent + // to seed_uniquestream when prefix carries exactly the bits [64,128). + _dev void seed_from_state(const uint64_t* prefix, uint32_t low0, uint32_t low1) { + for (int i = 0; i < N; i++) { V[i] = prefix[i]; } + for (int b = 0; b < 32; b++) { if ((low0 >> b) & 1u) { apply_skip_bit(V, b); } } + for (int b = 0; b < 32; b++) { if ((low1 >> b) & 1u) { apply_skip_bit(V, 32 + b); } } + finalize_skipped_state(); + } + +private: + _dev static inline uint64_t MOD_MULSPEC(uint64_t k); + _dev inline void seed_vielbein(); // seeds with the unit vector {1,0,0,...} + _dev static inline uint64_t iterate_raw_vec(uint64_t* Y, uint64_t sumtotOld); + _dev inline uint64_t apply_bigskip(uint32_t clusterID, uint32_t machineID, uint32_t runID, uint32_t streamID ); + _dev static inline uint64_t modadd(uint64_t foo, uint64_t bar){return MOD_MERSENNE(foo+bar);}; + _dev static inline uint64_t fmodmulM61(uint64_t cum, uint64_t s, uint64_t a); + _dev static inline uint64_t MOD_MERSENNE(uint64_t k) {return ((((k)) & MERSBASE) + (((k)) >> BITS) );} + _dev static inline uint64_t MULWU(uint64_t k) {return (( (k)<<(SPECIALMUL) & M61) | ( (k) >> (BITS-SPECIALMUL)) );} + void inline print_state(); + +public: + +#if (!defined(__CUDACC__) && !defined (__CUDA_ARCH__)) + template + friend std::basic_ostream & + operator<<(std::basic_ostream &ost, const mixmax_engine &me) { + // save the state of RNG to stream + int j; + ost << "mixmax state, file version 1.0\n"; + ost << "N=" << N << "; V[N]={"; + for (j = 0; (j < (N - 1)); j++) { + ost << (std::uint64_t)me.V[j] << ", "; + } + ost << me.V[N - 1]; + ost << "}; "; + ost << "counter=" << (std::uint64_t)me.counter << "; "; + ost << "sumtot=" << (std::uint64_t)me.sumtot << ";\n"; + ost.flush(); + return ost; + } + + template + friend std::basic_istream & + operator>>(std::basic_istream &in, mixmax_engine &me) { + // will set std::ios::failbit and throw an exception if format is not right + std::array vec; + std::uint64_t sum = 0, sumtmp = 0, counter = 0; + in.ignore(150, '='); // eat chars up to N= + CharT xxxchar; + try { + std::basic_string line; + if (std::getline(in, line)) { + std::basic_istringstream iss(line); + std::basic_string token; + getline(iss, token, xxxchar = (';')); + int i = std::stoi(token); + if (i != N) { + std::cerr + << "ERROR: Wrong dimension of the MIXMAX RNG state on input, " + << i << " vs " << N << "\n"; + } // else{std::cerr <<"Dim ok "<< i << "\n";} + iss.ignore(15, '{'); // eat chars up to V[N]={ + for (int j = 0; j < N - 1; j++) { + std::getline(iss, token, xxxchar = ','); + if (!iss.fail()) { + vec[j] = std::stoull(token); + sum = me.modadd(sum, vec[j]); + } // std::cerr <> 1); r++; // bring up the r-th bit in the ID + } + } + uint64_t insumtot=0; + for (int i=0; i>32); + pl = (s) & MASK32; + ah = a>>32; + al = a & MASK32; + o = (o & M61) + ((ph*ah)<<3) + ((ah*pl+al*ph + ((al*pl)>>32))>>29) ; + o += cum; + o = (o & M61) + ((o>>61)); + return o; +} + +#if !defined(__CUDA_ARCH__) +#include + void mixmax_engine ::print_state(){ + int j; + fprintf(stdout, "mixmax state, file version 1.0\n" ); + fprintf(stdout, "N=%u; V[N]={", N ); + for (j=0; (j< (N-1) ); j++) { + fprintf(stdout, "%llu, ", (unsigned long long)V[j] ); + } + fprintf(stdout, "%llu", (unsigned long long)V[N-1] ); + fprintf(stdout, "}; " ); + fprintf(stdout, "counter=%u; ", counter ); + fprintf(stdout, "sumtot=%llu;\n", (unsigned long long)sumtot ); +} +#endif + +_dev mixmax_engine mixmax_engine ::Branch(){ + sumtot = iterate_raw_vec(V, sumtot); counter = N-1; + mixmax_engine tmp=*this; + tmp.BranchInplace(); + return tmp; +} + +_dev mixmax_engine & mixmax_engine ::operator=(const mixmax_engine& other ){ + for (int i=0; i seed = std::nullopt + ) const; const std::string& mask_name() const { return _mask_name; } private: diff --git a/madspace/include/madspace/phasespace/discrete_flow.hpp b/madspace/include/madspace/phasespace/discrete_flow.hpp index 5dc05a793..88913bfa5 100644 --- a/madspace/include/madspace/phasespace/discrete_flow.hpp +++ b/madspace/include/madspace/phasespace/discrete_flow.hpp @@ -18,7 +18,10 @@ class DiscreteFlow : public Mapping { ); const std::vector& option_counts() const { return _option_counts; } std::size_t condition_dim() const { return _condition_dim; } - void initialize_globals(ContextPtr context) const; + // nullopt (the default) initializes non-deterministically. + void initialize_globals( + ContextPtr context, std::optional seed = std::nullopt + ) const; private: Result build_forward_impl( diff --git a/madspace/include/madspace/phasespace/flow.hpp b/madspace/include/madspace/phasespace/flow.hpp index fcf064084..51149d6aa 100644 --- a/madspace/include/madspace/phasespace/flow.hpp +++ b/madspace/include/madspace/phasespace/flow.hpp @@ -1,5 +1,7 @@ #pragma once +#include + #include "madspace/phasespace/base.hpp" #include "madspace/phasespace/mlp.hpp" @@ -19,8 +21,15 @@ class Flow : public Mapping { ); std::size_t input_dim() const { return _input_dim; } std::size_t condition_dim() const { return _condition_dim; } - void initialize_globals(ContextPtr context) const; - void initialize_from_vegas(ContextPtr context, const std::string& grid_name) const; + // nullopt (the default) initializes non-deterministically. + void initialize_globals( + ContextPtr context, std::optional seed = std::nullopt + ) const; + void initialize_from_vegas( + ContextPtr context, + const std::string& grid_name, + std::optional seed = std::nullopt + ) const; private: Result build_forward_impl( diff --git a/madspace/include/madspace/phasespace/mlp.hpp b/madspace/include/madspace/phasespace/mlp.hpp index 45234a8f3..5b6b374b6 100644 --- a/madspace/include/madspace/phasespace/mlp.hpp +++ b/madspace/include/madspace/phasespace/mlp.hpp @@ -1,5 +1,8 @@ #pragma once +#include +#include + #include "madspace/driver/context.hpp" #include "madspace/phasespace/base.hpp" @@ -17,7 +20,10 @@ class MLP : public FunctionGenerator { std::size_t input_dim() const { return _input_dim; } std::size_t output_dim() const { return _output_dim; } - void initialize_globals(ContextPtr context) const; + // nullopt (the default) initializes non-deterministically. + void initialize_globals( + ContextPtr context, std::optional seed = std::nullopt + ) const; std::string last_layer_bias_name() const { return prefixed_name(_prefix, std::format("layer{}.bias", _layers)); } diff --git a/madspace/include/madspace/util.hpp b/madspace/include/madspace/util.hpp index 7b1e43b38..de5d8c9b5 100644 --- a/madspace/include/madspace/util.hpp +++ b/madspace/include/madspace/util.hpp @@ -2,6 +2,7 @@ #include #include +#include #include #include #include diff --git a/madspace/src/compgraphs/function.cpp b/madspace/src/compgraphs/function.cpp index 6552c0327..cc6ce199d 100644 --- a/madspace/src/compgraphs/function.cpp +++ b/madspace/src/compgraphs/function.cpp @@ -361,10 +361,15 @@ FunctionBuilder::instruction(InstructionPtr instruction, const ValueVec& args) { } } + // random instructions share one GPU RNG stream, so they always run on the + // main stream (0) and are never CSE'd (see Instruction::is_random) + bool is_random = instruction->is_random(); + std::size_t call_stream = is_random ? 0 : _current_stream; + // create local variables for constants std::vector cache_key; cache_key.push_back(opcode); - cache_key.push_back(_current_stream); + cache_key.push_back(call_stream); for (auto& arg : params) { register_local(arg); cache_key.push_back(arg.local_index); @@ -381,8 +386,8 @@ FunctionBuilder::instruction(InstructionPtr instruction, const ValueVec& args) { } } - // check for cached result (for deterministic instructions) - if (opcode != opcodes::random && opcode != opcodes::unweight) { + // check for cached result + if (!is_random) { auto find_instr = _instruction_cache.find(cache_key); if (find_instr != _instruction_cache.end()) { ValueVec call_outputs; @@ -405,10 +410,12 @@ FunctionBuilder::instruction(InstructionPtr instruction, const ValueVec& args) { output_locals.push_back(value.local_index); } _instructions.push_back( - InstructionCall{instruction, params, call_outputs, _current_stream} + InstructionCall{instruction, params, call_outputs, call_stream} ); _instruction_use_count.push_back(0); - _instruction_cache[cache_key] = output_locals; + if (!is_random) { + _instruction_cache[cache_key] = output_locals; + } for (auto& param : params) { int param_source = _local_sources.at(param.local_index); if (param_source != -1) { diff --git a/madspace/src/cpu/device.hpp b/madspace/src/cpu/device.hpp index 945fb7311..c6de22e36 100644 --- a/madspace/src/cpu/device.hpp +++ b/madspace/src/cpu/device.hpp @@ -1,5 +1,9 @@ #pragma once +#include +#include + +#include "madspace/driver/random.hpp" #include "madspace/driver/tensor.hpp" #include "madspace/driver/thread_pool.hpp" #include "simd.hpp" @@ -7,6 +11,8 @@ namespace madspace { namespace cpu { +class CpuRuntime; + class CpuDevice : public Device { public: static constexpr bool is_concurrent = false; diff --git a/madspace/src/cpu/runtime.cpp b/madspace/src/cpu/runtime.cpp index 873897a4a..bd3860a20 100644 --- a/madspace/src/cpu/runtime.cpp +++ b/madspace/src/cpu/runtime.cpp @@ -1,6 +1,7 @@ #include "runtime.hpp" #include +#include #include #include #include @@ -712,12 +713,11 @@ void op_random( auto& runtime = instruction.runtime; device.foreach ( flat_view.shape[0], - [flat_view, &runtime](std::size_t count, std::size_t offset) mutable { + [flat_view, &runtime, &device](std::size_t count, std::size_t offset) mutable { auto output_view = TensorView(flat_view); - std::uniform_real_distribution dist; auto& rand_gen = runtime.rand_gen(); for (std::size_t i = offset; i < offset + count; ++i) { - output_view[i] = dist(rand_gen); + output_view[i] = rand_gen.generate_double(); } } ); @@ -735,12 +735,13 @@ void op_random_int( auto& runtime = instruction.runtime; device.foreach ( flat_view.shape[0], - [flat_view, max_val, &runtime](std::size_t count, std::size_t offset) mutable { + [flat_view, max_val, &runtime, &device]( + std::size_t count, std::size_t offset + ) mutable { auto output_view = TensorView(flat_view); - std::uniform_int_distribution dist(0, max_val - 1); auto& rand_gen = runtime.rand_gen(); for (std::size_t i = offset; i < offset + count; ++i) { - output_view[i] = dist(rand_gen); + output_view[i] = rand_gen.generate_int(max_val); } } ); @@ -772,6 +773,7 @@ void op_unweight( indices_view_flat, uw_weights_view_flat, &runtime, + &device, batch_size, &indices, &uw_weights, @@ -781,14 +783,14 @@ void op_unweight( TensorView max_weight_view(max_weight_view_flat); TensorView indices_view(indices_view_flat); TensorView uw_weights_view(uw_weights_view_flat); - std::uniform_real_distribution dist; auto& rand_gen = runtime.rand_gen(); std::size_t count = 0; for (std::size_t i = 0; i < batch_size; ++i) { - double w = weights_view[i], w_max = max_weight_view[i]; - if (w != 0. && w > dist(rand_gen) * w_max) { + double w = weights_view[i], aw = std::abs(w), + w_max = max_weight_view[i]; + if (aw != 0. && aw > rand_gen.generate_double() * w_max) { indices_view[count] = i; - uw_weights_view[count] = w > w_max ? w : w_max; + uw_weights_view[count] = std::copysign(std::max(aw, w_max), w); ++count; } } @@ -982,13 +984,9 @@ void op_discrete_histogram( CpuRuntime::CpuRuntime(const Function& function, ContextPtr context, bool concurrent) : _context(context), _input_count(function.inputs().size()), - _rand_gens( - context->thread_pool(), - []() { - std::random_device rand_device; - return std::mt19937(rand_device()); - } - ), + _rand_gens(context->global_resource( + "cpu_rand_gen", []() { return MixMaxRandom(); } + )), _concurrent(concurrent) { if (context->device()->device_type() != DeviceType::cpu) { throw std::runtime_error("Context has incompatible device"); diff --git a/madspace/src/cpu/runtime.hpp b/madspace/src/cpu/runtime.hpp index 9b70fdbfc..10ef87ccb 100644 --- a/madspace/src/cpu/runtime.hpp +++ b/madspace/src/cpu/runtime.hpp @@ -1,10 +1,12 @@ #pragma once +#include #include #include "madspace/compgraphs/function.hpp" #include "madspace/driver/backend.hpp" #include "madspace/driver/context.hpp" +#include "madspace/driver/random.hpp" #include "madspace/driver/tensor.hpp" namespace madspace { @@ -49,7 +51,8 @@ class CpuRuntime : public Runtime { ) override; Context& context() { return *_context; } - std::mt19937& rand_gen() { return _rand_gens.get(); } + MixMaxRandom& rand_gen() { return _rand_gens.get(); } + void set_seed(DerivedSeed seed) override { rand_gen().set_seed(seed); } private: TensorVec run_single(const TensorVec& inputs) const; @@ -83,7 +86,7 @@ class CpuRuntime : public Runtime { std::vector _grad_global_shapes; std::size_t _grad_global_total_size; ContextPtr _context; - ThreadResource _rand_gens; + ThreadResource& _rand_gens; bool _concurrent; SizeVec _ready_instructions_init; SizeVec _ready_instructions_backward_init; diff --git a/madspace/src/cpu/tensor.hpp b/madspace/src/cpu/tensor.hpp index d4397ef40..cb0bd7afe 100644 --- a/madspace/src/cpu/tensor.hpp +++ b/madspace/src/cpu/tensor.hpp @@ -305,6 +305,20 @@ inline void tensor_foreach_impl( bool single_job, S... scalar_args ) { + // A broadcast *output* (size(0) != batch_size, e.g. a global's gradient + // accumulator during backward) is shared across the whole batch instead of + // having one slot per batch element. Splitting the batch across worker threads + // would then have multiple threads accumulate into the very same memory + // concurrently with no synchronization, so force a single job in that case + // (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. + for (const Tensor* out : outputs) { + if (out->size(0) != batch_size) { + single_job = true; + } + } + // get views to the tensors with the correct types based on the signature of // scalar_func auto flat_views = std::apply( diff --git a/madspace/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index 3f1601bfc..13b9a8289 100644 --- a/madspace/src/driver/channel_generator.cpp +++ b/madspace/src/driver/channel_generator.cpp @@ -1,4 +1,6 @@ #include "madspace/driver/channel_generator.hpp" +#include "madspace/driver/random.hpp" +#include "madspace/util.hpp" using namespace madspace; using json = nlohmann::json; @@ -32,6 +34,44 @@ int particle_extra_flags( return flags; } +DerivedSeed::SeedType generate_seed_type(bool is_survey, std::size_t survey_pass) { + if (!is_survey) { + return DerivedSeed::generator_generate; + } + return survey_pass == 0 + ? DerivedSeed::first_survey_generate + : DerivedSeed::second_survey_generate; +} + +DerivedSeed::SeedType unweight_seed_type(bool is_survey, std::size_t survey_pass) { + if (!is_survey) { + return DerivedSeed::generator_unweight; + } + return survey_pass == 0 + ? DerivedSeed::first_survey_unweight + : DerivedSeed::second_survey_unweight; +} + +// Role folded into job_index so integrand_channel and integrand_common get +// independent streams despite sharing seed_type/channel_index. +enum RngRole { rng_role_channel = 0, rng_role_common = 1, rng_role_count = 2 }; + +DerivedSeed generate_phase_seed( + std::optional seed, + bool is_survey, + std::size_t survey_pass, + std::size_t job_index, + std::size_t channel_index, + RngRole role +) { + return DerivedSeed( + seed, + generate_seed_type(is_survey, survey_pass), + job_index * rng_role_count + role, + channel_index + ); +} + } // namespace ChannelEventGenerator::ChannelEventGenerator( @@ -50,13 +90,15 @@ ChannelEventGenerator::ChannelEventGenerator( .name = name, .mean = 0., .error = 0., + .mean_abs = 0., + .error_abs = 0., .rel_std_dev = 0., .count = 0, .count_opt = 0, .count_after_cuts = 0, .count_after_cuts_opt = 0, .count_unweighted = 0., - .count_target = 1., + .count_target = 1, .optimized = false, .done = false }, @@ -174,13 +216,15 @@ ChannelEventGenerator::ChannelEventGenerator( .name = name, .mean = 0., .error = 0., + .mean_abs = 0., + .error_abs = 0., .rel_std_dev = 0., .count = 0, .count_opt = 0, .count_after_cuts = 0, .count_after_cuts_opt = 0, .count_unweighted = 0., - .count_target = 1., + .count_target = 1, .optimized = false, .done = false }, @@ -283,27 +327,30 @@ void ChannelEventGenerator::init_field_indices() { _field_indices.rest = _field_indices.random + 1; } -void ChannelEventGenerator::unweight_file(std::mt19937& rand_gen) { +void ChannelEventGenerator::unweight_file(MixMaxRandom& rand_gen) { std::size_t buf_size = 1000000; - std::uniform_real_distribution rand_dist; EventBuffer buffer(0, 0, weight_file_layout); - std::size_t accept_count = _unweighted_count; + std::size_t accept_count = _unweighted_accept_count; for (std::size_t i = _unweighted_count; i < _weight_file.event_count(); i += buf_size) { _weight_file.seek(i); _weight_file.read(buffer, buf_size); for (std::size_t j = 0; j < buffer.event_count(); ++j) { auto weight = buffer.event(j).weight(); - if (weight / _max_weight < rand_dist(rand_gen)) { + if (std::abs(weight.value()) / _max_weight < rand_gen.generate_double()) { weight = 0; } else { - weight = std::max(weight.value(), _max_weight); + weight = std::copysign( + std::max(std::abs(weight.value()), _max_weight), weight.value() + ); ++accept_count; } } _weight_file.seek(i); _weight_file.write(buffer); } + _unweighted_count = _weight_file.event_count(); + _unweighted_accept_count = accept_count; _status.count_unweighted = accept_count; } @@ -315,10 +362,13 @@ void ChannelEventGenerator::integrate(const GeneratorBatchJob& job) { ++sample_count_after_cuts; } _cross_section.push(w_view[i]); + _abs_cross_section.push(std::abs(w_view[i])); } _status.mean = _cross_section.mean(); _status.error = _cross_section.error(); - _status.rel_std_dev = _cross_section.rel_std_dev(); + _status.mean_abs = _abs_cross_section.mean(); + _status.error_abs = _abs_cross_section.error(); + _status.rel_std_dev = _abs_cross_section.rel_std_dev(); _status.count += w_view.size(); _status.count_opt += w_view.size(); _status.count_after_cuts += sample_count_after_cuts; @@ -353,7 +403,7 @@ void ChannelEventGenerator::optimize_vegas(const GeneratorBatchJob& job) { if (_discrete_optimizer) { _discrete_optimizer->optimize(); } - double rsd = _cross_section.rel_std_dev(); + double rsd = _abs_cross_section.rel_std_dev(); if (rsd < _config.optimization_threshold * _best_rsd) { _iters_without_improvement = 0; } else { @@ -384,8 +434,7 @@ double ChannelEventGenerator::channel_weight_sum(std::size_t event_count) { if (weight == 0.) { continue; } - weight_sum += weight / _max_weight; - _unweighted_count = 0; + weight_sum += std::abs(weight) / _max_weight; ++unweighted_count; } if (done) { @@ -395,27 +444,55 @@ double ChannelEventGenerator::channel_weight_sum(std::size_t event_count) { return weight_sum; } +// Assigns job's base seed on the scheduling thread, from its logical identity +// only (channel, kind, sequence) -- never from context_index -- so it doesn't +// depend on worker/context assignment. void ChannelEventGenerator::start_job( - GeneratorBatchJob& job, ResultQueue& result_queue + GeneratorBatchJob& job, + ResultQueue& result_queue, + std::optional seed, + bool is_survey, + std::size_t survey_pass ) { + job.rng_seed = seed; + job.rng_is_survey = is_survey; + job.rng_survey_pass = survey_pass; + job.rng_job_index = is_survey ? _survey_rng_seq++ : _generate_rng_seq++; _contexts.at(job.context_index) ->thread_pool() .submit([this, &job, &result_queue]() { auto& runtimes = _runtimes.at(job.context_index); auto& context = _contexts.at(job.context_index); + if (job.rng_seed) { + runtimes.integrand_channel->set_seed(generate_phase_seed( + job.rng_seed, + job.rng_is_survey, + job.rng_survey_pass, + job.rng_job_index, + job.channel_index, + rng_role_channel + )); + } std::size_t max_batch_size = context->device()->device_type() == DeviceType::cpu ? _config.cpu_batch_size : _config.gpu_batch_size; std::size_t batch_size = max_batch_size; - if (job.vegas_batch_size > 0 && batch_size > job.vegas_batch_size) { - batch_size = job.vegas_batch_size; + // Only VEGAS batches shrink to fit -- they need next_vegas_batch_size()'s + // exact geometric progression for correct grid-adaptation statistics. + // Generation batches always submit full device-sized jobs; a little + // overshoot there is expected and cheaper than a partial job. + if (job.is_vegas_batch && job.batch_event_count > 0 && + batch_size > job.batch_event_count) { + batch_size = job.batch_event_count; } std::size_t target_count = batch_size; std::size_t total_count = 0, repetitions = 0; TensorVec all_ps_points; while (true) { + // Successive calls draw further from integrand_channel's own + // persistent stream (set once above) -- no per-call re-seeding. auto ps_points = runtimes.integrand_channel->run({Tensor({batch_size})}); std::size_t acc_count = @@ -452,9 +529,21 @@ void ChannelEventGenerator::start_job( (target_count - total_count) / cut_eff ); } + if (job.rng_seed) { + runtimes.integrand_common->set_seed(generate_phase_seed( + job.rng_seed, + job.rng_is_survey, + job.rng_survey_pass, + job.rng_job_index, + job.channel_index, + rng_role_common + )); + } job.events = runtimes.integrand_common->run(all_ps_points); job.weights = job.events.at(_field_indices.weight).cpu(); + // observable_histograms/vegas_histogram/discrete_histogram don't consume + // random numbers, so they're never seeded. if (runtimes.observable_histograms) { auto hists = runtimes.observable_histograms->run( {job.events.at(_field_indices.weight), @@ -464,7 +553,7 @@ void ChannelEventGenerator::start_job( job.hists.push_back(item.cpu()); } } - if (job.vegas_batch_size != 0) { + if (job.is_vegas_batch) { if (_vegas_optimizer) { auto hist = runtimes.vegas_histogram->run( {job.events.at(_field_indices.random), @@ -490,15 +579,26 @@ void ChannelEventGenerator::start_job( }); } -void ChannelEventGenerator::start_unweight_job( +void ChannelEventGenerator::prepare_unweight_job(GeneratorBatchJob& job) const { + job.max_weight = _max_weight; +} + +void ChannelEventGenerator::submit_unweight_job( GeneratorBatchJob& job, ResultQueue& result_queue ) { - job.max_weight = _max_weight; _contexts.at(job.context_index) ->thread_pool() .submit([this, &job, &result_queue]() { auto& runtimes = _runtimes.at(job.context_index); auto& context = _contexts.at(job.context_index); + if (job.rng_seed) { + runtimes.unweighter->set_seed(DerivedSeed( + job.rng_seed, + unweight_seed_type(job.rng_is_survey, job.rng_survey_pass), + job.rng_job_index, + job.channel_index + )); + } std::vector unweighter_args( job.events.begin(), job.events.begin() + _field_indices.random ); @@ -512,6 +612,13 @@ void ChannelEventGenerator::start_unweight_job( }); } +void ChannelEventGenerator::start_unweight_job( + GeneratorBatchJob& job, ResultQueue& result_queue +) { + prepare_unweight_job(job); + submit_unweight_job(job, result_queue); +} + std::size_t ChannelEventGenerator::next_vegas_batch_size() { std::size_t batch_size = _batch_size; _batch_size = std::min(_batch_size * 2, _config.max_batch_size); @@ -522,11 +629,13 @@ void ChannelEventGenerator::clear_events() { _status.count_unweighted = 0; _max_weight = 0; _unweighted_count = 0; + _unweighted_accept_count = 0; _status.count_opt = 0; _status.count_after_cuts_opt = 0; _event_file.clear(); _weight_file.clear(); _cross_section.reset(); + _abs_cross_section.reset(); _large_weights.clear(); for (auto& hist : _histograms) { std::fill(hist.bin_values.begin(), hist.bin_values.end(), 0.); @@ -557,8 +666,9 @@ void ChannelEventGenerator::update_max_weight(Tensor weights) { double w_sum = 0; double max_truncation = _config.max_overweight_truncation * - std::min(_status.count_target, - static_cast(_config.freeze_max_weight_after)); + static_cast(std::min( + _status.count_target, _config.freeze_max_weight_after + )); std::size_t count = 0; for (auto w : _large_weights) { if (w < _max_weight) { @@ -571,6 +681,7 @@ void ChannelEventGenerator::update_max_weight(Tensor weights) { _status.count_unweighted *= _max_weight / w; _max_weight = w; _unweighted_count = 0; + _unweighted_accept_count = 0; } break; } diff --git a/madspace/src/driver/event_generator.cpp b/madspace/src/driver/event_generator.cpp index 4ccf27ca1..865d699c9 100644 --- a/madspace/src/driver/event_generator.cpp +++ b/madspace/src/driver/event_generator.cpp @@ -1,11 +1,15 @@ #include "madspace/driver/event_generator.hpp" +#include #include #include +#include +#include #include #include #include "madspace/driver/logger.hpp" +#include "madspace/driver/random.hpp" #include "madspace/util.hpp" using namespace madspace; @@ -15,6 +19,7 @@ const GeneratorConfig EventGenerator::default_config = {}; EventGenerator::EventGenerator( const std::vector& contexts, const std::vector>& channels, + std::uint64_t seed, std::shared_ptr status_file, const GeneratorConfig& config ) : @@ -24,13 +29,15 @@ EventGenerator::EventGenerator( .name = "", .mean = 0., .error = 0., + .mean_abs = 0., + .error_abs = 0., .rel_std_dev = 0., .count = 0, .count_opt = 0, .count_after_cuts = 0, .count_after_cuts_opt = 0, .count_unweighted = 0., - .count_target = static_cast(config.target_count), + .count_target = config.target_count, .optimized = false, .done = false }, @@ -41,13 +48,160 @@ EventGenerator::EventGenerator( _channel_optimizing(channels.size()), _channel_integral_fractions(channels.size(), 1.), _context_job_counts(contexts.size()), + _channel_batch_pending(channels.size()), + _channel_batch_dispatch_done(channels.size()), + _channel_gen_order(channels.size()), + _channel_ready_gen(channels.size()), + _channel_unweight_order(channels.size()), + _channel_unweight_ready(channels.size()), + _context_unweight_queue(contexts.size()), + _seed(seed), _status_file(status_file) {} -void EventGenerator::survey() { - for (auto& context : _contexts) { - context->reset_cache(); +// Commits a job's generate stage in per-channel dispatch order. Doesn't write +// events itself -- if job.unweight, that's deferred to commit_unweight_job(). +void EventGenerator::commit_generate_job(GeneratorBatchJob& job) { + auto& channel = _channels.at(job.channel_index); + auto& channel_job_count = _channel_job_counts.at(job.channel_index); + if (job.is_vegas_batch && channel_job_count == job.split_job_count) { + channel->clear_events(); } + channel->integrate(job); + update_integral_status(); + channel->update_max_weight(job.weights); + --channel_job_count; + --_context_job_counts.at(job.context_index); + if (job.unweight) { + // Snapshot max_weight here, at this job's fixed commit position. + channel->prepare_unweight_job(job); + _context_unweight_queue.at(job.context_index).push_back(job.job_id); + _channel_unweight_order.at(job.channel_index).push_back(job.job_id); + } + finish_channel_job(job); + print_gen_update(false); +} + +void EventGenerator::commit_unweight_job(GeneratorBatchJob& job) { + auto& channel = _channels.at(job.channel_index); + channel->write_events(job.unweighted_events, job.max_weight); + update_counts(); + --_channel_job_counts.at(job.channel_index); + --_context_job_counts.at(job.context_index); + finish_channel_job(job); + print_gen_update(false); +} + +void EventGenerator::finish_channel_job(const GeneratorBatchJob& job) { + auto& channel_job_count = _channel_job_counts.at(job.channel_index); + if (channel_job_count != 0) { + return; + } + if (job.is_vegas_batch) { + _channels.at(job.channel_index)->optimize_vegas(job); + _channel_optimizing.at(job.channel_index) = false; + } else if (_channel_batch_dispatch_done.at(job.channel_index)) { + _channel_batch_pending.at(job.channel_index) = false; + } +} + +// Jobs run on the pool, but commit strictly in per-channel dispatch order, so +// results are a pure function of the seed. +void EventGenerator::generate() { + _survey_job = false; + reset_start_time(); + print_gen_init(); + + while (true) { + _abort_check_function(); + + for (std::size_t channel_index = 0; channel_index < _channels.size(); + ++channel_index) { + auto& channel = _channels.at(channel_index); + double integral_frac = _channel_integral_fractions.at(channel_index); + if (integral_frac > 0 && + channel->status().count_unweighted >= channel->status().count_target) { + continue; + } + if (channel->needs_optimization()) { + if (!_channel_optimizing.at(channel_index)) { + _channel_optimizing.at(channel_index) = true; + _ready_jobs.push_back({ + .channel_index = channel_index, + .unweight = true, + .batch_event_count = channel->next_vegas_batch_size(), + .is_vegas_batch = true, + }); + } + } else if (!_channel_batch_pending.at(channel_index)) { + _channel_batch_pending.at(channel_index) = true; + _channel_batch_dispatch_done.at(channel_index) = false; + _ready_jobs.push_back({ + .channel_index = channel_index, + .unweight = true, + .batch_event_count = next_batch_event_count(channel_index), + .is_vegas_batch = false, + }); + } + } + + std::size_t job_id_before = _job_id; + std::size_t unweight_dispatched = start_jobs(); + std::size_t round_in_flight = (_job_id - job_id_before) + unweight_dispatched; + + // Wait for the round to fully commit before resyncing cross-channel fractions. + while (round_in_flight > 0) { + std::size_t job_id = _result_queue.wait(); + --round_in_flight; + auto& job = _running_jobs.at(job_id); + if (job.unweighted_events.size() == 0) { + auto& ready = _channel_ready_gen.at(job.channel_index); + auto& order = _channel_gen_order.at(job.channel_index); + ready.insert(job_id); + while (!order.empty() && ready.erase(order.front()) > 0) { + std::size_t commit_id = order.front(); + order.pop_front(); + auto& gen_job = _running_jobs.at(commit_id); + commit_generate_job(gen_job); + if (!gen_job.unweight) { + _running_jobs.erase(commit_id); + } + } + } else { + auto& ready = _channel_unweight_ready.at(job.channel_index); + auto& order = _channel_unweight_order.at(job.channel_index); + ready.insert(job_id); + while (!order.empty() && ready.erase(order.front()) > 0) { + std::size_t commit_id = order.front(); + order.pop_front(); + commit_unweight_job(_running_jobs.at(commit_id)); + _running_jobs.erase(commit_id); + } + } + std::size_t job_id_refill = _job_id; + std::size_t refill_unweight_dispatched = start_jobs(); + round_in_flight += (_job_id - job_id_refill) + refill_unweight_dispatched; + } + + update_integral_fractions(); + update_counts(); + + if (_status.done) { + unweight_all(); + } + if (_status.done) { + break; + } + } + print_gen_update(true); +} + +// Same commit-ordering approach as generate(). +void EventGenerator::survey(std::size_t survey_pass) { + _survey_job = true; + _survey_pass = survey_pass; reset_start_time(); + _commit_cursor = _job_id; + _ready_gen.clear(); bool done = false; std::size_t min_iters = _config.survey_min_iters; std::size_t max_iters = std::max(min_iters, _config.survey_max_iters); @@ -67,14 +221,13 @@ void EventGenerator::survey() { std::size_t iter = 0; for (; !done && iter < max_iters; ++iter) { - std::size_t job_count_before = _running_jobs.size(); for (std::size_t i = 0; auto channel : _channels) { if (channel->status().iterations > iter) { ++i; continue; } if (iter >= min_iters && - channel->cross_section().rel_error() < target_precision) { + channel->abs_cross_section().rel_error() < target_precision) { ++i; continue; } @@ -82,192 +235,207 @@ void EventGenerator::survey() { _ready_jobs.push_back({ .channel_index = i, .unweight = iter >= min_iters - 1, - .vegas_batch_size = vegas_batch_size, + .batch_event_count = vegas_batch_size, + .is_vegas_batch = true, }); if (iter >= min_iters) { total_event_count += vegas_batch_size; } ++i; } - start_jobs(); + std::size_t job_id_before = _job_id; + std::size_t unweight_dispatched = start_jobs(); + std::size_t in_flight = (_job_id - job_id_before) + unweight_dispatched; done = true; - while (_running_jobs.size() > 0) { + while (in_flight > 0) { std::size_t job_id = _result_queue.wait(); _abort_check_function(); + --in_flight; auto& job = _running_jobs.at(job_id); - auto& channel = _channels.at(job.channel_index); - auto& channel_job_count = _channel_job_counts.at(job.channel_index); - auto& context_job_count = _context_job_counts.at(job.context_index); - if (channel_job_count == job.split_job_count) { - channel->clear_events(); - } - --channel_job_count; - --context_job_count; - - bool keep_job = false; if (job.unweighted_events.size() == 0) { - channel->integrate(job); - update_integral(); - channel->update_max_weight(job.weights); - if (job.unweight) { - channel->start_unweight_job(job, _result_queue); - ++context_job_count; - keep_job = true; - } else { - done = false; + // Generate-stage completion, globally ordered via + // _ready_gen/_commit_cursor. + _ready_gen.insert(job_id); + while (_ready_gen.erase(_commit_cursor) > 0) { + auto& gen_job = _running_jobs.at(_commit_cursor); + auto& channel = _channels.at(gen_job.channel_index); + auto& channel_job_count = + _channel_job_counts.at(gen_job.channel_index); + if (channel_job_count == gen_job.split_job_count) { + channel->clear_events(); + } + channel->integrate(gen_job); + update_integral(); + channel->update_max_weight(gen_job.weights); + --channel_job_count; + --_context_job_counts.at(gen_job.context_index); + // Global order is ascending job id and so is the per-channel one, + // so this job is necessarily at the front of its channel's deque. + _channel_gen_order.at(gen_job.channel_index).pop_front(); + if (gen_job.unweight) { + // Snapshot max_weight at this fixed commit position. + channel->prepare_unweight_job(gen_job); + _context_unweight_queue.at(gen_job.context_index) + .push_back(gen_job.job_id); + _channel_unweight_order.at(gen_job.channel_index) + .push_back(gen_job.job_id); + } else { + done = false; + if (channel_job_count == 0) { + channel->optimize_vegas(gen_job); + done_event_count += gen_job.batch_event_count; + } + _running_jobs.erase(_commit_cursor); + } + ++_commit_cursor; } } else { - channel->write_events(job.unweighted_events, job.max_weight); - update_counts(); - if (channel_job_count == 0 && - channel->cross_section().rel_error() < target_precision) { - done = false; + // Unweight-stage completion, ordered per channel only. + auto& ready = _channel_unweight_ready.at(job.channel_index); + auto& order = _channel_unweight_order.at(job.channel_index); + ready.insert(job_id); + while (!order.empty() && ready.erase(order.front()) > 0) { + std::size_t commit_id = order.front(); + order.pop_front(); + auto& uw_job = _running_jobs.at(commit_id); + auto& channel = _channels.at(uw_job.channel_index); + channel->write_events(uw_job.unweighted_events, uw_job.max_weight); + update_counts(); + auto& channel_job_count = + _channel_job_counts.at(uw_job.channel_index); + --channel_job_count; + --_context_job_counts.at(uw_job.context_index); + if (channel_job_count == 0 && + channel->abs_cross_section().rel_error() < target_precision) { + done = false; + } + if (channel_job_count == 0) { + channel->optimize_vegas(uw_job); + done_event_count += uw_job.batch_event_count; + } + _running_jobs.erase(commit_id); } } - if (channel_job_count == 0) { - channel->optimize_vegas(job); - done_event_count += job.vegas_batch_size; - } - if (!keep_job) { - _running_jobs.erase(job_id); - } - start_jobs(); + std::size_t job_id_refill = _job_id; + std::size_t refill_unweight_dispatched = start_jobs(); + in_flight += (_job_id - job_id_refill) + refill_unweight_dispatched; print_survey_update(false, done_event_count, total_event_count, iter); } } print_survey_update(true, done_event_count, total_event_count, iter - 1); } -void EventGenerator::generate() { - reset_start_time(); - print_gen_init(); - - std::size_t target_job_count = 0; - for (auto& context : _contexts) { - context->reset_cache(); - target_job_count += 2 * context->thread_pool().thread_count(); - } - std::size_t channel_index = 0; - while (true) { - _abort_check_function(); - - std::size_t job_count_before; - do { - job_count_before = _ready_jobs.size(); - for (std::size_t i = 0; - i < _channels.size() && _ready_jobs.size() < target_job_count; - ++i, channel_index = (channel_index + 1) % _channels.size()) { - auto& channel = _channels.at(channel_index); - std::size_t& channel_job_count = _channel_job_counts.at(channel_index); - double integral_frac = _channel_integral_fractions.at(channel_index); - if (integral_frac > 0 && - channel->status().count_unweighted >= - integral_frac * _config.target_count) { - continue; - } - if (channel->needs_optimization()) { - if (!_channel_optimizing.at(channel_index)) { - _channel_optimizing.at(channel_index) = true; - _ready_jobs.push_back({ - .channel_index = channel_index, - .unweight = true, - .vegas_batch_size = channel->next_vegas_batch_size(), - }); - } - } else { - _ready_jobs.push_back({ - .channel_index = channel_index, - .unweight = true, - .vegas_batch_size = 0, - }); - } - } - } while (_ready_jobs.size() - job_count_before > 0); - start_jobs(); - - if (_running_jobs.size() > 0) { - std::size_t job_id = _result_queue.wait(); - auto& job = _running_jobs.at(job_id); - auto& channel = _channels.at(job.channel_index); - auto& channel_job_count = _channel_job_counts.at(job.channel_index); - auto& context_job_count = _context_job_counts.at(job.context_index); - if (job.vegas_batch_size > 0 && channel_job_count == job.split_job_count) { - channel->clear_events(); - } - --channel_job_count; - --context_job_count; - - bool keep_job = false; - if (job.unweighted_events.size() == 0) { - channel->integrate(job); - update_integral(); - channel->update_max_weight(job.weights); - if (job.unweight) { - channel->start_unweight_job(job, _result_queue); - ++context_job_count; - keep_job = true; - } - } else { - channel->write_events(job.unweighted_events, job.max_weight); - update_counts(); - } - if (job.vegas_batch_size > 0 && channel_job_count == 0) { - channel->optimize_vegas(job); - _channel_optimizing.at(job.channel_index) = false; - } - print_gen_update(false); - if (!keep_job) { - _running_jobs.erase(job_id); - } - } else { - if (_status.done) { - unweight_all(); - } - if (_status.done) { - break; - } - } - } - print_gen_update(true); +// Reads only committed data (never jobs still in flight), so the estimate is a +// pure function of the seed. Sizing itself is in +// compute_generation_batch_event_count(). +std::size_t EventGenerator::next_batch_event_count(std::size_t channel_index) const { + auto& channel = _channels.at(channel_index); + auto& status = channel->status(); + auto& abs_cross_section = channel->abs_cross_section(); + return compute_generation_batch_event_count( + status.count_target, + status.count_unweighted, + status.count_opt, + abs_cross_section.count(), + abs_cross_section.rel_error(), + _config + ); } -bool EventGenerator::start_jobs() { - std::size_t ready_index = 0, context_index = 0; +std::size_t EventGenerator::start_jobs() { + std::size_t context_index = 0; + std::size_t unweight_dispatched = 0; for (auto [context, job_count] : zip(_contexts, _context_job_counts)) { // fill the queue to twice the thread count to keep the worker threads busy std::size_t target_count = 2 * context->thread_pool().thread_count(); std::size_t batch_size = context->device()->device_type() == DeviceType::cpu ? _config.cpu_batch_size : _config.gpu_batch_size; - for (; job_count < target_count && ready_index < _ready_jobs.size(); - ++ready_index) { - auto ready_job = _ready_jobs.at(ready_index); - std::size_t split_job_count = ready_job.vegas_batch_size != 0 - ? (ready_job.vegas_batch_size + batch_size - 1) / batch_size + + // give priority to unweighting jobs to reduce memory usage + auto& unweight_queue = _context_unweight_queue.at(context_index); + std::size_t unweight_index = 0; + for (; job_count < target_count && unweight_index < unweight_queue.size(); + ++unweight_index) { + auto& job = _running_jobs.at(unweight_queue.at(unweight_index)); + _channels.at(job.channel_index)->submit_unweight_job(job, _result_queue); + ++job_count; + ++unweight_dispatched; + } + unweight_queue.erase( + unweight_queue.begin(), unweight_queue.begin() + unweight_index + ); + + // Round-robin through _ready_jobs, one device batch per generation entry per + // visit, so multiple channels' batches interleave instead of one channel's + // batch draining before the next is even touched -- see _ready_job_rr_cursor. + // A VEGAS batch is still dispatched atomically in a single visit: VEGAS needs + // next_vegas_batch_size()'s exact progression. + while (job_count < target_count && !_ready_jobs.empty()) { + if (_ready_job_rr_cursor >= _ready_jobs.size()) { + _ready_job_rr_cursor = 0; + } + auto& ready_job = _ready_jobs.at(_ready_job_rr_cursor); + bool atomic_dispatch = + ready_job.is_vegas_batch || ready_job.batch_event_count == 0; + std::size_t remaining_sub_jobs = ready_job.batch_event_count != 0 + ? (ready_job.batch_event_count + batch_size - 1) / batch_size : 1; - for (std::size_t i = 0; i < split_job_count; ++i) { + std::size_t to_dispatch = atomic_dispatch ? remaining_sub_jobs : 1; + + for (std::size_t i = 0; i < to_dispatch; ++i) { + GeneratorBatchJob new_job{ + .channel_index = ready_job.channel_index, + .unweight = ready_job.unweight, + .batch_event_count = ready_job.batch_event_count, + .split_job_count = to_dispatch * (1 + ready_job.unweight), + .context_index = context_index, + .job_id = _job_id, + .is_vegas_batch = ready_job.is_vegas_batch, + }; auto& job = - std::get<0>(_running_jobs.emplace(_job_id, ready_job))->second; - job.split_job_count = split_job_count * (1 + job.unweight); - job.job_id = _job_id; - job.context_index = context_index; - _channels.at(job.channel_index)->start_job(job, _result_queue); + std::get<0>(_running_jobs.emplace(_job_id, std::move(new_job))) + ->second; + // Records this channel's commit order; drained by survey()/generate() + // -- see _channel_gen_order. + _channel_gen_order.at(job.channel_index).push_back(_job_id); + _channels.at(job.channel_index) + ->start_job(job, _result_queue, _seed, _survey_job, _survey_pass); _channel_job_counts.at(job.channel_index) += 1 + job.unweight; ++_job_id; ++job_count; } - } - if (ready_index == _ready_jobs.size()) { - break; + + if (!atomic_dispatch) { + ready_job.batch_event_count -= + std::min(to_dispatch * batch_size, ready_job.batch_event_count); + } + bool fully_dispatched = atomic_dispatch || ready_job.batch_event_count == 0; + if (fully_dispatched) { + if (!ready_job.is_vegas_batch) { + // Only meaningful for generate(); harmlessly unused by the other + // callers of start_jobs(). + _channel_batch_dispatch_done.at(ready_job.channel_index) = true; + } + _ready_jobs.erase(_ready_jobs.begin() + _ready_job_rr_cursor); + // The next entry has shifted into this index -- don't advance past it. + } else { + ++_ready_job_rr_cursor; + } } ++context_index; } - _ready_jobs.erase(_ready_jobs.begin(), _ready_jobs.begin() + ready_index); - return ready_index > 0; + return unweight_dispatched; } void EventGenerator::update_integral() { + update_integral_status(); + update_integral_fractions(); +} + +void EventGenerator::update_integral_status() { double total_mean = 0., total_var = 0.; + double total_mean_abs = 0., total_var_abs = 0.; std::size_t total_count = 0, total_count_opt = 0; std::size_t total_count_after_cuts = 0, total_count_after_cuts_opt = 0; std::size_t total_integ_count = 0; @@ -277,14 +445,19 @@ void EventGenerator::update_integral() { for (auto& channel : _channels) { auto& status = channel->status(); auto& cross_section = channel->cross_section(); + auto& abs_cross_section = channel->abs_cross_section(); // special case for channels with 0 samples, as they have nan variance if (bad_channel == nullptr && cross_section.count() > 0 && (!std::isfinite(cross_section.mean()) || - !std::isfinite(cross_section.variance()))) { + !std::isfinite(cross_section.variance()) || + !std::isfinite(abs_cross_section.mean()) || + !std::isfinite(abs_cross_section.variance()))) { bad_channel = channel.get(); } total_mean += cross_section.mean(); total_var += cross_section.variance() / cross_section.count(); + total_mean_abs += abs_cross_section.mean(); + total_var_abs += abs_cross_section.variance() / abs_cross_section.count(); total_count += status.count; total_count_opt += status.count_opt; total_count_after_cuts += status.count_after_cuts; @@ -295,7 +468,8 @@ void EventGenerator::update_integral() { optimized = false; } } - if (bad_channel != nullptr || !std::isfinite(total_mean)) { + if (bad_channel != nullptr || !std::isfinite(total_mean) || + !std::isfinite(total_mean_abs)) { std::string where = bad_channel != nullptr ? std::format( "channel '{}' after {} samples (mean={}, variance={})", @@ -322,26 +496,57 @@ void EventGenerator::update_integral() { _status.mean = total_mean; _status.error = std::sqrt(total_var); - _status.rel_std_dev = std::sqrt(total_var * total_integ_count) / total_mean; + _status.mean_abs = total_mean_abs; + _status.error_abs = std::sqrt(total_var_abs); + _status.rel_std_dev = std::sqrt(total_var_abs * total_integ_count) / total_mean_abs; _status.count = total_count; _status.count_opt = total_count_opt; _status.count_after_cuts = total_count_after_cuts; _status.count_after_cuts_opt = total_count_after_cuts_opt; _status.iterations = iterations; _status.optimized = optimized; +} + +void EventGenerator::update_integral_fractions() { for (auto [channel, integral_fraction] : zip(_channels, _channel_integral_fractions)) { - integral_fraction = channel->cross_section().mean() / total_mean; - channel->set_target_count(integral_fraction * _config.target_count); + integral_fraction = channel->abs_cross_section().mean() / _status.mean_abs; + } + + // Distribute events between channels, ensure sum is exactly target_count + std::vector counts(_channels.size()); + std::vector remainders(_channels.size()); + std::size_t allocated = 0; + for (auto [count, remainder, fraction] : + zip(counts, remainders, _channel_integral_fractions)) { + double raw = fraction * _config.target_count; + double floored = std::floor(std::isfinite(raw) && raw > 0. ? raw : 0.); + count = static_cast(floored); + remainder = raw - floored; + allocated += count; + } + std::vector order(_channels.size()); + std::iota(order.begin(), order.end(), 0); + std::sort(order.begin(), order.end(), [&](std::size_t a, std::size_t b) { + double rem_a = remainders.at(a), rem_b = remainders.at(b); + return rem_a == rem_b ? a < b : rem_a > rem_b; + }); + std::size_t leftover = + allocated <= _config.target_count ? _config.target_count - allocated : 0; + leftover = std::min(leftover, order.size()); + for (std::size_t index : order | std::views::take(leftover)) { + ++counts.at(index); + } + for (auto [channel, count] : zip(_channels, counts)) { + channel->set_target_count(count); } } void EventGenerator::update_counts() { double total_eff_count = 0.; bool done = true; - for (auto [channel, integral_fraction] : - zip(_channels, _channel_integral_fractions)) { - double chan_target = integral_fraction * _config.target_count; + for (auto& channel : _channels) { + std::size_t chan_target = channel->status().count_target; if (channel->status().count_unweighted < chan_target) { total_eff_count += channel->status().count_unweighted; done = false; @@ -355,6 +560,7 @@ void EventGenerator::update_counts() { void EventGenerator::combine_to_compact_npy(const std::string& file_name) { reset_start_time(); + MixMaxRandom select_rng(DerivedSeed(_seed, DerivedSeed::combine_select)); auto [channel_data, particle_count, norm_factor] = init_combine(); DataLayout layout( EventRecord::layout( @@ -373,7 +579,7 @@ void EventGenerator::combine_to_compact_npy(const std::string& file_name) { print_combine_init(); while (true) { _abort_check_function(); - read_and_combine(channel_data, buffer, norm_factor); + read_and_combine(channel_data, buffer, norm_factor, select_rng); if (buffer.event_count() == 0) { break; } @@ -391,8 +597,9 @@ void EventGenerator::combine_to_lhe_npy( const std::string& file_name, LHECompleter& lhe_completer ) { reset_start_time(); - std::random_device rand_device; - std::mt19937 rand_gen(rand_device()); + std::uint64_t seed = _seed; + MixMaxRandom select_rng(DerivedSeed(seed, DerivedSeed::combine_select)); + MixMaxRandom rand_gen(DerivedSeed(seed, DerivedSeed::lhe_complete)); auto [channel_data, particle_count, norm_factor] = init_combine(); DataLayout in_layout( EventRecord::layout( @@ -424,7 +631,7 @@ void EventGenerator::combine_to_lhe_npy( print_combine_init(); while (true) { _abort_check_function(); - read_and_combine(channel_data, buffer, norm_factor); + read_and_combine(channel_data, buffer, norm_factor, select_rng); if (buffer.event_count() == 0) { break; } @@ -454,11 +661,9 @@ void EventGenerator::combine_to_lhe( const std::string& file_name, LHECompleter& lhe_completer, const LHEMeta& meta ) { reset_start_time(); + std::uint64_t seed = _seed; + MixMaxRandom select_rng(DerivedSeed(seed, DerivedSeed::combine_select)); ThreadPool pool(_config.combine_thread_count); - ThreadResource rand_gens(pool, []() { - std::random_device rand_device; - return std::mt19937(rand_device()); - }); auto [channel_data, particle_count, norm_factor] = init_combine(); std::vector> buffers; std::vector idle_buffers; @@ -480,45 +685,63 @@ void EventGenerator::combine_to_lhe( std::size_t event_count = 0; std::size_t last_update_count = 0; bool done = false; + // Batches are seeded by submission order and written in that same order + // (buffered in ready_slots until their turn), independent of completion timing. + std::size_t next_batch_seq = 0; + std::size_t write_cursor = 0; + std::vector slot_batch_seq(buffers.size()); + std::unordered_map ready_slots; print_combine_init(); while (true) { _abort_check_function(); while (idle_buffers.size() > 0 && !done) { - std::size_t job_id = idle_buffers.back(); - auto& [in_buffer, out_buffer] = buffers.at(job_id); - read_and_combine(channel_data, in_buffer, norm_factor); + std::size_t slot = idle_buffers.back(); + auto& [in_buffer, out_buffer] = buffers.at(slot); + read_and_combine(channel_data, in_buffer, norm_factor, select_rng); if (in_buffer.event_count() == 0) { done = true; break; } idle_buffers.pop_back(); + std::size_t batch_seq = next_batch_seq++; + slot_batch_seq.at(slot) = batch_seq; pool.submit( - [job_id, this, &in_buffer, &out_buffer, &lhe_completer, &rand_gens] { + [slot, batch_seq, seed, this, &in_buffer, &out_buffer, &lhe_completer] { + MixMaxRandom rand_gen( + DerivedSeed(seed, DerivedSeed::lhe_complete, batch_seq) + ); LHEEvent lhe_event; out_buffer.clear(); for (std::size_t i = 0; i < in_buffer.event_count(); ++i) { fill_lhe_event( - lhe_completer, lhe_event, in_buffer, i, rand_gens.get() + lhe_completer, lhe_event, in_buffer, i, rand_gen ); lhe_event.format_to(out_buffer); } - return job_id; + return slot; } ); } - auto done_jobs = pool.wait_multiple(); - for (std::size_t job_id : done_jobs) { - auto& [in_buffer, out_buffer] = buffers.at(job_id); - idle_buffers.push_back(job_id); + auto done_slots = pool.wait_multiple(); + for (std::size_t slot : done_slots) { + ready_slots.emplace(slot_batch_seq.at(slot), slot); + } + for (auto it = ready_slots.find(write_cursor); it != ready_slots.end(); + it = ready_slots.find(write_cursor)) { + std::size_t slot = it->second; + auto& [in_buffer, out_buffer] = buffers.at(slot); event_file.write_string(out_buffer); event_count += in_buffer.event_count(); + idle_buffers.push_back(slot); + ready_slots.erase(it); + ++write_cursor; if (event_count - last_update_count > 10000) { print_combine_update(event_count); last_update_count = event_count; } } - if (done_jobs.size() == 0 && done) { + if (done_slots.size() == 0 && done) { break; } } @@ -542,15 +765,17 @@ void EventGenerator::add_timing_data(const std::string& key) { } void EventGenerator::unweight_all() { - std::random_device rand_device; - std::mt19937 rand_gen(rand_device()); + std::size_t call_index = _unweight_call_index++; bool done = true; double total_eff_count = 0.; - for (auto [channel, integral_fraction] : - zip(_channels, _channel_integral_fractions)) { + for (std::size_t channel_index = 0; auto& channel : _channels) { + MixMaxRandom rand_gen( + DerivedSeed(_seed, DerivedSeed::unweight_pass, call_index, channel_index) + ); channel->unweight_file(rand_gen); + ++channel_index; - double chan_target = integral_fraction * _config.target_count; + std::size_t chan_target = channel->status().count_target; if (channel->status().count_unweighted < chan_target) { total_eff_count += channel->status().count_unweighted; done = false; @@ -608,11 +833,12 @@ EventGenerator::init_combine() { std::size_t count_sum = 0; std::size_t particle_count = 0; double weight_sum = 0.; - for (auto [channel, integral_fraction] : - zip(_channels, _channel_integral_fractions)) { + for (auto& channel : _channels) { particle_count = std::max(particle_count, channel->event_file().particle_count()); - std::size_t count = std::round(integral_fraction * _config.target_count); + // Exact apportioned target (see update_integral_fractions()), so counts + // sum to exactly _config.target_count. + std::size_t count = channel->status().count_target; count_sum += count; channel->event_file().seek(0); weight_sum += channel->channel_weight_sum(count); @@ -626,13 +852,14 @@ EventGenerator::init_combine() { .buffer_index = 0, }); } - return {channel_data, particle_count, _status.mean * count_sum / weight_sum}; + return {channel_data, particle_count, _status.mean_abs * count_sum / weight_sum}; } void EventGenerator::read_and_combine( std::vector& channel_data, EventBuffer& buffer, - double norm_factor + double norm_factor, + MixMaxRandom& rand_gen ) { std::size_t batch_size = 1000; std::size_t event_count = std::min(batch_size, channel_data.back().cum_count); @@ -645,11 +872,8 @@ void EventGenerator::read_and_combine( bool has_subproc_index = _channels.at(0)->event_layout_extra_flags() & EventRecord::f_subproc_index; - std::random_device rand_device; - std::mt19937 rand_gen(rand_device()); for (std::size_t event_index = 0; event_index < event_count; ++event_index) { - std::size_t random_index = std::uniform_int_distribution< - std::size_t>(0, channel_data.back().cum_count - 1)(rand_gen); + std::size_t random_index = rand_gen.generate_int(channel_data.back().cum_count); auto sampled_chan = std::lower_bound( channel_data.begin(), channel_data.end(), @@ -679,7 +903,11 @@ void EventGenerator::read_and_combine( auto event_in = sampled_chan->event_buffer.event(sampled_chan->buffer_index); auto event_out = buffer.event(event_index); - event_out.weight() = std::max(1., weight / channel->max_weight()) * norm_factor; + event_out.weight() = + std::copysign( + std::max(1., std::abs(weight) / channel->max_weight()), weight + ) * + norm_factor; event_out.subprocess_index() = has_subproc_index ? event_in.subprocess_index().value() : static_cast(channel->status().subprocess); @@ -728,7 +956,7 @@ void EventGenerator::fill_lhe_event( LHEEvent& lhe_event, EventBuffer& buffer, std::size_t event_index, - std::mt19937& rand_gen + MixMaxRandom& rand_gen ) { EventRecord event_in = buffer.event(event_index); lhe_event.weight = event_in.weight(); @@ -772,6 +1000,7 @@ void EventGenerator::write_status(const std::string& status, bool force_write) { } nlohmann::json j{ {"status", status}, + {"seed", _seed}, {"process", _status}, {"channels", channel_status()}, {"run_times", _timing_data}, @@ -961,9 +1190,10 @@ void EventGenerator::print_gen_update_pretty(bool done) { ); if (!std::isnan(_status.error)) { - double rel_err = _status.error / _status.mean; int_str = format_with_error(_status.mean, _status.error); - rel_str = std::format("{:.4f} %", rel_err * 100); + if (std::abs(_status.error) < std::abs(_status.mean)) { + rel_str = std::format("{:.4f} %", _status.error / _status.mean * 100); + } rsd_str = std::format("{:.3f}", _status.rel_std_dev); uweff_str = std::format( "{:.5f} before cuts, {:.5f} after", @@ -998,7 +1228,7 @@ void EventGenerator::print_gen_update_pretty(bool done) { if (_channels.size() > 1) { auto channels = channel_status(); std::sort(channels.begin(), channels.end(), [](auto& chan1, auto& chan2) { - return chan1.mean > chan2.mean; + return chan1.mean_abs > chan2.mean_abs; }); for (std::size_t row = 1; auto& channel : channels | std::views::take(20)) { @@ -1047,15 +1277,18 @@ void EventGenerator::print_gen_update_log(bool done) { } _last_print_time = now; + std::string rel_str = std::abs(_status.error) < std::abs(_status.mean) + ? std::format("{:.4f} %", _status.error / _status.mean * 100) + : ""; Logger::info( std::format( - "generating, events: {} / {}, integral: {}, rel. error: {:.4f} %, " + "generating, events: {} / {}, integral: {}, rel. error: {}, " "RSD: {:.3f}, samps: {}, samps. after cuts: {}, " "unw. eff.: {:.5f}, unw. eff. after cuts: {:.5f}, time: {:%H:%M:%S}", format_si_prefix(_status.count_unweighted), format_si_prefix(_status.count_target), format_with_error(_status.mean, _status.error), - _status.error / _status.mean * 100, + rel_str, _status.rel_std_dev, format_si_prefix(_status.count), format_si_prefix(_status.count_after_cuts), diff --git a/madspace/src/driver/format.cpp b/madspace/src/driver/format.cpp index 38ed5961c..bae7e921c 100644 --- a/madspace/src/driver/format.cpp +++ b/madspace/src/driver/format.cpp @@ -17,6 +17,27 @@ const std::array progress_symbols{ " ", "▏", "▎", "▍", "▌", "▋", "▊", "▉", "█" }; +// Fallback for format_with_error() once the relative error is O(1) or larger (or +// value/error are non-finite/non-positive): the "1.234(56)" digit-of-error notation is +// meaningless there, so print "value ± error" instead, with the error rendered to two +// significant digits (or the value alone if the error itself isn't meaningful). +std::string format_value_pm_error(double value, double error) { + if (!std::isfinite(error) || error <= 0.) { + if (!std::isfinite(value) || value == 0.) { + return std::format("{}", value); + } + int value_power = static_cast(std::floor(std::log10(std::abs(value)))); + return value_power < -5 || value_power > 5 + ? std::format("{:.3e}", value) + : std::format("{:.3f}", value); + } + int sig_power = 1 - static_cast(std::floor(std::log10(error))); + if (sig_power < 0 || sig_power > 5) { + return std::format("{:.1e} ± {:.1e}", value, error); + } + return std::format("{:.{}f} ± {:.{}f}", value, sig_power, error, sig_power); +} + } // namespace std::size_t madspace::cpu_time_microsec() { @@ -60,10 +81,13 @@ std::string madspace::format_si_prefix(double value) { } std::string madspace::format_with_error(double value, double error) { - int value_power = std::floor(std::log10(value)); - int sig_power = std::isnan(error) || error <= 0. - ? 3 - value_power - : 1 - static_cast(std::floor(std::log10(error))); + if (!std::isfinite(value) || !std::isfinite(error) || error <= 0. || + std::abs(error) >= std::abs(value)) { + // relative error is O(1) or larger -- the (56) digit notation is meaningless + return format_value_pm_error(value, error); + } + int value_power = std::floor(std::log10(std::abs(value))); + int sig_power = 1 - static_cast(std::floor(std::log10(error))); if (sig_power < 0 || sig_power > 5) { std::string exp_fmt = std::format("{:.{}e}", value, value_power + sig_power); auto e_pos = exp_fmt.find("e"); diff --git a/madspace/src/driver/generator_data.cpp b/madspace/src/driver/generator_data.cpp index 4d7d77066..83e1a9ce0 100644 --- a/madspace/src/driver/generator_data.cpp +++ b/madspace/src/driver/generator_data.cpp @@ -1,13 +1,55 @@ #include "madspace/driver/generator_data.hpp" +#include +#include + using namespace madspace; +std::size_t madspace::compute_generation_batch_event_count( + std::size_t count_target, + double count_unweighted, + std::size_t count_opt, + std::size_t abs_cross_section_count, + double abs_cross_section_rel_error, + const GeneratorConfig& config +) { + double efficiency = count_opt > 0 + ? std::max(count_unweighted / static_cast(count_opt), 1. / count_opt) + : 1.; + + double true_remaining = + std::max(static_cast(count_target) - count_unweighted, 0.); + + if (true_remaining <= config.finish_remaining_fraction * count_target) { + return static_cast( + std::max(1., std::ceil(true_remaining / efficiency)) + ); + } + + double rel_error = + abs_cross_section_count > 1 && std::isfinite(abs_cross_section_rel_error) + ? abs_cross_section_rel_error + : 1.; + + double target_uncertainty = static_cast(count_target) * rel_error; + double safe_target = static_cast(count_target) - + config.batch_overshoot_sigma * target_uncertainty; + double safe_remaining = std::max(safe_target - count_unweighted, 0.); + double capped = std::min( + config.max_batch_fraction * static_cast(count_target), safe_remaining + ); + + return static_cast(std::max(1., std::ceil(capped / efficiency))); +} + void madspace::to_json(nlohmann::json& j, const GeneratorStatus& status) { j = nlohmann::json{ {"subprocess", status.subprocess}, {"name", status.name}, {"mean", status.mean}, {"error", status.error}, + {"mean_abs", status.mean_abs}, + {"error_abs", status.error_abs}, {"rel_std_dev", status.rel_std_dev}, {"count", status.count}, {"count_opt", status.count_opt}, diff --git a/madspace/src/driver/lhe_output.cpp b/madspace/src/driver/lhe_output.cpp index 2b6970476..6c3425ca6 100644 --- a/madspace/src/driver/lhe_output.cpp +++ b/madspace/src/driver/lhe_output.cpp @@ -480,7 +480,7 @@ void LHECompleter::complete_event_data( int color_index, int flavor_index, int helicity_index, - std::mt19937& rand_gen + MixMaxRandom& rand_gen ) { auto& subproc_data = _subproc_data.at(subprocess_index); if (event.particles.size() != subproc_data.particle_count) { @@ -509,8 +509,7 @@ void LHECompleter::complete_event_data( auto [pdg_index, pdg_count] = _pdg_id_and_count.at(subproc_data.pdg_id_offset + flavor_index); - std::uniform_int_distribution dist(0, pdg_count - 1); - std::size_t pdg_random = dist(rand_gen); + std::size_t pdg_random = rand_gen.generate_int(pdg_count); std::size_t pdg_offset = pdg_index + subproc_data.particle_count * pdg_random; for (std::size_t particle_index = 0; auto& particle : event.particles) { diff --git a/madspace/src/driver/madnis_training.cpp b/madspace/src/driver/madnis_training.cpp index 7518e63cc..c43904a66 100644 --- a/madspace/src/driver/madnis_training.cpp +++ b/madspace/src/driver/madnis_training.cpp @@ -1,5 +1,6 @@ #include "madspace/driver/madnis_training.hpp" +#include "madspace/driver/random.hpp" #include "madspace/phasespace/batch_sampler.hpp" using namespace madspace; @@ -9,13 +10,17 @@ MadnisTraining::MadnisTraining( ContextPtr optimizer_context, const Config& config, const std::vector>& integrands, - const std::optional& cwnet + const std::optional& cwnet, + std::optional seed, + std::size_t channel_index_offset ) : _generator_context(generator_context), _optimizer_context(optimizer_context), _config(config), _channels(integrands.size()), - _cwnet(cwnet) { + _cwnet(cwnet), + _seed(seed), + _channel_index_offset(channel_index_offset) { for (std::size_t index = 0; auto [integrand, channel] : zip(integrands, _channels)) { channel.index = index; @@ -60,7 +65,7 @@ void MadnisTraining::train_step(std::size_t batch_index) { TensorVec training_batch; bool used_buffered = false; while (true) { - start_generator_jobs(channel_sizes); + maybe_start_generator_jobs(channel_sizes, !try_buffered); if (try_buffered) { if (check_buffered_training_batch(channel_sizes)) { training_batch = build_buffered_training_batch(channel_sizes); @@ -88,27 +93,25 @@ void MadnisTraining::train_step(std::size_t batch_index) { } if (_channels.size() > 0 && _cwnet && (batch_index + 1) % _config.channel_dropping_interval == 0) { - std::vector job_ids; - while ((job_ids = gen_thread_pool.wait_multiple()).size() != 0) { - process_job_results(job_ids); - } + process_all_jobs(); drop_channels(batch_index + 1); } if (batch_index == static_cast( (1 - _config.fixed_cwnet_fraction) * _config.batches )) { - std::vector job_ids; - while ((job_ids = gen_thread_pool.wait_multiple()).size() != 0) { - process_job_results(job_ids); - } + process_all_jobs(); freeze_cwnet(); } - if (batch_index == _config.batches - 1) { - std::vector job_ids; - while ((job_ids = gen_thread_pool.wait_multiple()).size() != 0) { - process_job_results(job_ids); - } + if (batch_index + 1 == _config.batches) { + process_all_jobs(); + } +} + +void MadnisTraining::process_all_jobs() { + std::vector job_ids; + while ((job_ids = _generator_context->thread_pool().wait_multiple()).size() != 0) { + process_job_results(job_ids); } } @@ -323,9 +326,16 @@ void MadnisTraining::start_generator_jobs(const std::vector& counts if (_running_jobs.size() > 0) { return; } + bool is_gpu = _generator_context->device()->device_type() != DeviceType::cpu; + // flush buffer samples staged since the last round (see process_job_results) + for (auto& channel : _channels) { + for (auto& pending : channel.pending_buffer_samples) { + buffer_store(channel, pending); + } + channel.pending_buffer_samples.clear(); + } _generator_params.copy_from(_optimizer->parameters()); std::size_t chan_count = counts.size(); - bool is_gpu = _generator_context->device()->device_type() != DeviceType::cpu; std::size_t batch_size = is_gpu ? _config.gpu_generator_batch_granularity : _config.cpu_generator_batch_size; @@ -341,14 +351,21 @@ void MadnisTraining::start_generator_jobs(const std::vector& counts ? (target_count - channel.sample_count + batch_size - 1) / batch_size : 0; } - std::size_t available_jobs = _generator_context->thread_pool().thread_count(); - std::vector channel_sizes; std::size_t gpu_subbatches = (_config.gpu_generator_batch_size + _config.gpu_generator_batch_granularity - 1) / _config.gpu_generator_batch_granularity; + // dispatch exactly what this round needs, uncapped by thread count, so + // round contents don't depend on thread count + std::size_t available_jobs = 0; + for (auto count : missing_batch_counts) { + available_jobs += count; + } + for (auto count : target_batch_counts) { + available_jobs += count; + } + std::vector channel_sizes; if (is_gpu) { - available_jobs *= gpu_subbatches; channel_sizes.resize(chan_count, 0); } @@ -410,6 +427,31 @@ void MadnisTraining::start_generator_jobs(const std::vector& counts } } +// Requires an online attempt and the online cache to be exhausted or about to +// be, both deterministic conditions, so round contents don't depend on thread +// scheduling. +void MadnisTraining::maybe_start_generator_jobs( + const std::vector& counts, bool is_online_attempt +) { + if (_running_jobs.size() > 0) { + return; + } + if (!is_online_attempt) { + return; + } + bool depletion_imminent = false; + for (auto [channel, count] : zip(_channels, counts)) { + if (count >= channel.sample_count) { + depletion_imminent = true; + break; + } + } + if (!depletion_imminent) { + return; + } + start_generator_jobs(counts); +} + TensorVec MadnisTraining::permute_tensors(const TensorVec& tensors) const { TensorVec ret; ret.reserve(tensors.size()); @@ -419,20 +461,42 @@ TensorVec MadnisTraining::permute_tensors(const TensorVec& tensors) const { return ret; } +// Seed depends only on channel index and channel-local sequence, not on which +// worker thread runs the job. void MadnisTraining::start_single_job( std::size_t channel_index, std::size_t batch_size ) { std::size_t job_id = _job_id; ++_job_id; + auto& channel = _channels.at(channel_index); + std::size_t channel_seq = channel.next_dispatch_seq++; + std::size_t global_channel_index = _channel_index_offset + channel_index; auto& job = std::get<0>(_running_jobs.emplace(job_id, SampleJob{}))->second; + job.dispatch_seq = channel_seq; _generator_context->thread_pool().submit( - [this, channel_index, batch_size, job_id, &job]() { + [this, channel_index, global_channel_index, batch_size, job_id, &job]() { auto& channel = _channels.at(channel_index); + if (_seed) { + channel.generator_runtime->set_seed(DerivedSeed( + _seed, + DerivedSeed::madnis_generate, + job.dispatch_seq, + global_channel_index + )); + } auto samples = channel.generator_runtime->run({Tensor({batch_size})}); job.samples.tensors = permute_tensors(samples); job.samples.size = samples.at(0).size(0); job.samples.channel_index = channel_index; if (channel.unweighter_runtime) { + if (_seed) { + channel.unweighter_runtime->set_seed(DerivedSeed( + _seed, + DerivedSeed::madnis_unweight, + job.dispatch_seq, + global_channel_index + )); + } auto unw_samples = channel.unweighter_runtime->run(samples); job.unweighted_samples.tensors = permute_tensors(unw_samples); job.unweighted_samples.size = unw_samples.at(0).size(0); @@ -443,15 +507,34 @@ void MadnisTraining::start_single_job( ); } +// Dispatch sequence is global (not per-channel), since one job spans all channels. void MadnisTraining::start_multi_job(const std::vector batch_sizes) { std::size_t job_id = _job_id; ++_job_id; + std::size_t dispatch_seq = _multi_job_next_dispatch_seq++; auto& job = std::get<0>(_running_jobs.emplace(job_id, SampleJob{}))->second; + job.dispatch_seq = dispatch_seq; _generator_context->thread_pool().submit([this, batch_sizes, job_id, &job]() { + if (_seed) { + _multi_channel_generator->set_seed(DerivedSeed( + _seed, + DerivedSeed::madnis_generate, + job.dispatch_seq, + _channel_index_offset + )); + } auto samples = _multi_channel_generator->run({Tensor(batch_sizes)}); job.samples.tensors = permute_tensors(samples); job.samples.channel_sizes = samples.back().batch_sizes(); if (_multi_channel_unweighter) { + if (_seed) { + _multi_channel_unweighter->set_seed(DerivedSeed( + _seed, + DerivedSeed::madnis_unweight, + job.dispatch_seq, + _channel_index_offset + )); + } auto unw_samples = _multi_channel_unweighter->run(samples); job.unweighted_samples.tensors = permute_tensors(unw_samples); job.unweighted_samples.channel_sizes = unw_samples.back().batch_sizes(); @@ -563,56 +646,89 @@ MadnisTraining::build_buffered_training_batch(const std::vector& counts) } } args.emplace_back(counts); + std::size_t batch_seq = _buffered_batch_seq++; + if (_seed) { + _multi_channel_sampler->set_seed(DerivedSeed( + _seed, DerivedSeed::madnis_sample_buffer, batch_seq, _channel_index_offset + )); + } return _multi_channel_sampler->run(args); } void MadnisTraining::process_job_results(const std::vector& job_ids) { for (auto job_id : job_ids) { - auto job = std::move(_running_jobs.extract(job_id).mapped()); + auto& job = _running_jobs.at(job_id); + // mark ready, committed below strictly in dispatch order if (job.samples.channel_sizes.size() == 0) { - auto& channel = _channels.at(job.samples.channel_index); - channel.sample_count += job.samples.size; - _generated_event_count += job.samples.size; - channel.sample_batches.push_back(std::move(job.samples)); - if (job.unweighted_samples.size > 0) { - buffer_store(channel, job.unweighted_samples); - } + _channels.at(job.samples.channel_index) + .ready_job_ids.emplace(job.dispatch_seq, job_id); } else { - std::size_t offset = 0, unw_offset = 0, chan_index = 0; - SampleBatch chan_unweighted_samples; - for (auto [channel, chan_size] : - zip(_channels, job.samples.channel_sizes)) { - if (chan_size == 0) { - ++chan_index; - continue; - } - channel.sample_count += chan_size; - _generated_event_count += chan_size; - channel.sample_batches.emplace_back(); - auto& batch = channel.sample_batches.back(); - batch.tensors.reserve(job.samples.tensors.size()); - for (auto& tensor : job.samples.tensors) { - batch.tensors.push_back( - tensor.slice(0, offset, offset + chan_size) + _multi_job_ready_job_ids.emplace(job.dispatch_seq, job_id); + } + } + + // commit each channel's ready single-channel jobs strictly in dispatch order + for (auto& channel : _channels) { + for (auto it = channel.ready_job_ids.find(channel.commit_cursor); + it != channel.ready_job_ids.end(); + it = channel.ready_job_ids.find(channel.commit_cursor)) { + auto committed_job = std::move(_running_jobs.extract(it->second).mapped()); + channel.ready_job_ids.erase(it); + ++channel.commit_cursor; + channel.sample_count += committed_job.samples.size; + _generated_event_count += committed_job.samples.size; + channel.sample_batches.push_back(std::move(committed_job.samples)); + if (committed_job.unweighted_samples.size > 0) { + // flushed into buffer at the start of the next round (see + // start_generator_jobs) + channel.pending_buffer_samples.push_back( + std::move(committed_job.unweighted_samples) + ); + } + } + } + + // commit ready multi-channel (GPU) jobs strictly in dispatch order + for (auto it = _multi_job_ready_job_ids.find(_multi_job_commit_cursor); + it != _multi_job_ready_job_ids.end(); + it = _multi_job_ready_job_ids.find(_multi_job_commit_cursor)) { + auto multi_job = std::move(_running_jobs.extract(it->second).mapped()); + _multi_job_ready_job_ids.erase(it); + ++_multi_job_commit_cursor; + std::size_t offset = 0, unw_offset = 0, chan_index = 0; + SampleBatch chan_unweighted_samples; + for (auto [channel, chan_size] : + zip(_channels, multi_job.samples.channel_sizes)) { + if (chan_size == 0) { + ++chan_index; + continue; + } + channel.sample_count += chan_size; + _generated_event_count += chan_size; + channel.sample_batches.emplace_back(); + auto& batch = channel.sample_batches.back(); + batch.tensors.reserve(multi_job.samples.tensors.size()); + for (auto& tensor : multi_job.samples.tensors) { + batch.tensors.push_back(tensor.slice(0, offset, offset + chan_size)); + } + if (multi_job.unweighted_samples.channel_sizes.size() > 0) { + std::size_t unw_chan_size = + multi_job.unweighted_samples.channel_sizes.at(chan_index); + chan_unweighted_samples.tensors.clear(); + chan_unweighted_samples.size = unw_chan_size; + for (auto& tensor : multi_job.unweighted_samples.tensors) { + chan_unweighted_samples.tensors.push_back( + tensor.slice(0, unw_offset, unw_offset + unw_chan_size) ); } - if (job.unweighted_samples.channel_sizes.size() > 0) { - std::size_t unw_chan_size = - job.unweighted_samples.channel_sizes.at(chan_index); - chan_unweighted_samples.tensors.clear(); - chan_unweighted_samples.size = unw_chan_size; - for (auto& tensor : job.unweighted_samples.tensors) { - chan_unweighted_samples.tensors.push_back( - tensor.slice(0, unw_offset, unw_offset + unw_chan_size) - ); - } - buffer_store(channel, chan_unweighted_samples); - unw_offset += unw_chan_size; - } - batch.size = chan_size; - offset += chan_size; - ++chan_index; + channel.pending_buffer_samples.push_back( + std::move(chan_unweighted_samples) + ); + unw_offset += unw_chan_size; } + batch.size = chan_size; + offset += chan_size; + ++chan_index; } } } @@ -720,7 +836,8 @@ void MadnisTraining::drop_channels(std::size_t batch) { } std::vector indices(_channels.size()); std::iota(indices.begin(), indices.end(), 0); - std::sort(indices.begin(), indices.end(), [&](auto i, auto j) { + // stable_sort: exact ties (e.g. charge-conjugate channels) must break consistently + std::stable_sort(indices.begin(), indices.end(), [&](auto i, auto j) { return abs_means.at(i) < abs_means.at(j); }); @@ -729,7 +846,8 @@ void MadnisTraining::drop_channels(std::size_t batch) { auto mask_view = active_mask.view()[0]; double drop_sum = 0.; - double drop_threshold = _config.channel_dropping_threshold * std::min(1000. / batch, 1.); + double drop_threshold = + _config.channel_dropping_threshold * std::min(1000. / batch, 1.); std::size_t drop_count = 0; for (std::size_t chan_index : indices) { drop_sum += abs_means.at(chan_index); @@ -747,8 +865,7 @@ void MadnisTraining::drop_channels(std::size_t batch) { return _active_flavors_count.at(flav_index) == 0; } )) { - // cannot drop this channel because one of its flavors is not - // available in any other channel + // a flavor of this channel has no other channel to fall back to continue; } for (std::size_t flav_index : active_flavors) { @@ -798,12 +915,14 @@ MultiMadnisTraining::MultiMadnisTraining( ContextPtr optimizer_context, const std::vector& training_args, Verbosity verbosity, - std::shared_ptr status_file + std::shared_ptr status_file, + std::optional seed ) : _generator_context(generator_context), _optimizer_context(optimizer_context), _training_args(training_args), _verbosity(verbosity), + _seed(seed), _status_file(status_file) { _trainings_status = nlohmann::json::array(); for (std::size_t i = 0; i < _training_args.size(); ++i) { @@ -826,6 +945,9 @@ void MultiMadnisTraining::train() { _start_time = std::chrono::steady_clock::now(); _start_cpu_microsec = cpu_time_microsec(); _active_channels.reserve(_training_args.size()); + // each subprocess gets its own channel_index slice so their derived seed + // streams don't collide + std::size_t channel_index_offset = 0; for (std::size_t subproc_index = 0; subproc_index < _training_args.size(); ++subproc_index) { auto& args = _training_args.at(subproc_index); @@ -836,7 +958,9 @@ void MultiMadnisTraining::train() { _optimizer_context, args.config, args.integrands, - args.cwnet + args.cwnet, + _seed, + channel_index_offset ); for (std::size_t batch_index = 0; batch_index < subproc.config().batches; ++batch_index) { @@ -849,6 +973,7 @@ void MultiMadnisTraining::train() { write_status(subproc, subproc_index, batch_index, done); } _active_channels.push_back(subproc.active_channels()); + channel_index_offset += args.integrands.size(); } } diff --git a/madspace/src/driver/random.cpp b/madspace/src/driver/random.cpp new file mode 100644 index 000000000..d27c60e8c --- /dev/null +++ b/madspace/src/driver/random.cpp @@ -0,0 +1,50 @@ +#include "madspace/driver/random.hpp" + +#include +#include + +namespace madspace { + +DerivedSeed::DerivedSeed( + const std::optional& seed, + SeedType seed_type, + std::size_t job_index, + std::size_t channel_index, + std::size_t stream_index +) { + std::uint64_t main_seed; + if (seed) { + main_seed = seed.value(); + } else { + std::random_device rand_device; + main_seed = rand_device(); + } + if (channel_index >= max_channel_count) { + throw std::invalid_argument( + std::format( + "channel index {} exceeded {}", channel_index, max_channel_count - 1 + ) + ); + } + if (job_index >= max_job_count) { + throw std::invalid_argument( + std::format("job index {} exceeded {}", job_index, max_job_count - 1) + ); + } + if (stream_index >= max_stream_count) { + throw std::invalid_argument( + std::format( + "stream index {} exceeded {}", stream_index, max_stream_count - 1 + ) + ); + } + seed_parts[0] = static_cast(main_seed >> 32); + seed_parts[1] = static_cast(main_seed & 0xFFFFFFFFULL); + seed_parts[2] = job_index; + seed_parts[3] = + ((static_cast(seed_type) << 28) | + (static_cast(channel_index) << 16) | + static_cast(stream_index)); +} + +} // namespace madspace diff --git a/madspace/src/gpu/device.hpp b/madspace/src/gpu/device.hpp index 15231af77..73f80a52f 100644 --- a/madspace/src/gpu/device.hpp +++ b/madspace/src/gpu/device.hpp @@ -15,14 +15,6 @@ inline void check_error(gpublasStatus_t status) { } } -inline void check_error(gpurandStatus_t status) { - if (status != GPURAND_STATUS_SUCCESS) { - throw std::runtime_error( - std::format("RAND error: error code {}", static_cast(status)) - ); - } -} - inline void check_error(gpuError_t error) { if (error != gpuSuccess) { const char* error_str = gpuGetErrorString(error); diff --git a/madspace/src/gpu/gpu_abstraction.cuh b/madspace/src/gpu/gpu_abstraction.cuh index 4677439b9..b8e94864c 100644 --- a/madspace/src/gpu/gpu_abstraction.cuh +++ b/madspace/src/gpu/gpu_abstraction.cuh @@ -5,7 +5,6 @@ #include #include #include -#include #define gpuGetDeviceCount cudaGetDeviceCount #define gpuSetDevice cudaSetDevice @@ -46,16 +45,6 @@ #define GPUBLAS_OP_N CUBLAS_OP_N #define GPUBLAS_OP_T CUBLAS_OP_T -#define gpurandStatus_t curandStatus_t -#define gpurandGenerator_t curandGenerator_t -#define gpurandCreateGenerator curandCreateGenerator -#define gpurandDestroyGenerator curandDestroyGenerator -#define gpurandSetPseudoRandomGeneratorSeed curandSetPseudoRandomGeneratorSeed -#define gpurandSetStream curandSetStream -#define gpurandGenerateUniformDouble curandGenerateUniformDouble -#define GPURAND_STATUS_SUCCESS CURAND_STATUS_SUCCESS -#define GPURAND_RNG_PSEUDO_DEFAULT CURAND_RNG_PSEUDO_DEFAULT - #define thrust_par thrust::cuda::par #elif defined __HIPCC__ @@ -63,7 +52,6 @@ #include #include #include -#include #define gpuGetDeviceCount hipGetDeviceCount #define gpuSetDevice hipSetDevice @@ -104,16 +92,6 @@ #define GPUBLAS_OP_N rocblas_operation_none #define GPUBLAS_OP_T rocblas_operation_transpose -#define gpurandStatus_t rocrand_status -#define gpurandGenerator_t rocrand_generator -#define gpurandCreateGenerator rocrand_create_generator -#define gpurandDestroyGenerator rocrand_destroy_generator -#define gpurandSetPseudoRandomGeneratorSeed rocrand_set_seed -#define gpurandSetStream rocrand_set_stream -#define gpurandGenerateUniformDouble rocrand_generate_uniform_double -#define GPURAND_STATUS_SUCCESS ROCRAND_STATUS_SUCCESS -#define GPURAND_RNG_PSEUDO_DEFAULT ROCRAND_RNG_PSEUDO_DEFAULT - #define thrust_par thrust::hip_rocprim::par namespace cub = hipcub; diff --git a/madspace/src/gpu/random.cuh b/madspace/src/gpu/random.cuh new file mode 100644 index 000000000..b01b25aee --- /dev/null +++ b/madspace/src/gpu/random.cuh @@ -0,0 +1,382 @@ +#pragma once + +#include "device.hpp" +#include "gpu_abstraction.cuh" +#include "madspace/driver/random.hpp" +#include "tensor.cuh" // MAX_THREADS_PER_BLOCK, THREADS_MULTIPLE + +#include +#include +#include +#include +#include + +namespace madspace { +namespace gpu { + +// One MIXMAX substream per DerivedSeed stream index; the bank covers all of them so +// every RNG-consuming instruction draws from its own non-overlapping substream. +constexpr std::size_t RNG_ENGINE_COUNT = 1 << 12; + +// The low SEED_STREAM_BITS of the effective 128-bit seed are the per-engine stream index +// (engine i uses shared_seed | i); the rest is shared across the whole bank. DerivedSeed +// reserves exactly these bits (max_stream_count == 1 << 16). +constexpr int SEED_STREAM_BITS = 16; +constexpr std::size_t Q_TABLE_SIZE = std::size_t(1) << SEED_STREAM_BITS; +static_assert( + RNG_ENGINE_COUNT <= Q_TABLE_SIZE, "engine bank exceeds the stream index space" +); + +constexpr int MIXMAX_D = mixmax_engine::state_size; // 17 +constexpr std::size_t RNG_SEED_BLOCK = 128; + +// State vector = MIXMAX vielbein with every set bit of the *shared* seed part +// (bits [SEED_STREAM_BITS, 128)) applied. Identical for all engines, so it is computed +// once on the host; the per-engine bits are finished on the GPU (see Q table below). +struct MixmaxSeedPrefix { + std::uint64_t state[MIXMAX_D]; +}; + +// bits [64,128) (the global run seed) come from `run_skip` (cached by the caller); apply +// the rest of the shared bits, [SEED_STREAM_BITS, 64), on top +inline MixmaxSeedPrefix +mixmax_seed_prefix(const std::array& seed_parts, const RunSeedSkip& run_skip) { + MixmaxSeedPrefix prefix; + for (int i = 0; i < MIXMAX_D; ++i) { + prefix.state[i] = run_skip.state[i]; + } + // effective seed word order matches apply_bigskip's IDvec: + // bits [SEED_STREAM_BITS,32) -> seed_parts[3], bits [32,64) -> seed_parts[2] + const std::uint32_t words[2] = {seed_parts[3], seed_parts[2]}; + for (int r = SEED_STREAM_BITS; r < 64; ++r) { + if ((words[r >> 5] >> (r & 31)) & 1u) { + mixmax_engine::apply_skip_bit(prefix.state, r); + } + } + return prefix; +} + +// --- Q table ------------------------------------------------------------------------ +// +// Q_i is the product of the skip operators T_r over the set bits r of i (r < +// SEED_STREAM_BITS). Every operator in the algebra generated by the MIXMAX iteration +// matrix A is a degree- built once. +// +// Build: a subset-product butterfly. Slot i is grown as the *state vector* Q_i(A) e_0 +// (reusing apply_skip_bit), then converted to coefficients via the Krylov change of +// basis kr_inv (columns A^j e_0; invertible because e_0 is cyclic for A). + +struct MixmaxKrInv { + std::uint64_t m[MIXMAX_D][MIXMAX_D]; +}; + +inline MixmaxKrInv mixmax_kr_inv() { + const std::uint64_t P = mixmax_engine::mod_prime; + auto mul = [&](std::uint64_t a, std::uint64_t b) { + return mixmax_engine::mod_mul_add(0, a % P, b % P) % P; + }; + auto sub = [&](std::uint64_t a, std::uint64_t b) { + return (a % P + (P - b % P)) % P; + }; + auto inv = [&](std::uint64_t a) { + std::uint64_t r = 1, base = a % P, e = P - 2; + while (e) { + if (e & 1) { + r = mul(r, base); + } + base = mul(base, base); + e >>= 1; + } + return r; + }; + + // kr[.][j] = A^j e_0 + std::uint64_t kr[MIXMAX_D][MIXMAX_D]; + std::uint64_t y[MIXMAX_D] = {0}; + y[0] = 1; + std::uint64_t sum = 1; + for (int j = 0; j < MIXMAX_D; ++j) { + for (int i = 0; i < MIXMAX_D; ++i) { + kr[i][j] = y[i] % P; + } + sum = mixmax_engine::advance_state(y, sum); + } + + // Gauss-Jordan invert kr mod P + MixmaxKrInv out; + std::uint64_t a[MIXMAX_D][MIXMAX_D]; + for (int i = 0; i < MIXMAX_D; ++i) { + for (int j = 0; j < MIXMAX_D; ++j) { + a[i][j] = kr[i][j]; + out.m[i][j] = (i == j) ? 1u : 0u; + } + } + for (int c = 0; c < MIXMAX_D; ++c) { + int piv = -1; + for (int row = c; row < MIXMAX_D; ++row) { + if (a[row][c] % P != 0) { + piv = row; + break; + } + } + if (piv < 0) { + throw std::runtime_error("mixmax_kr_inv: singular Krylov matrix"); + } + if (piv != c) { + for (int j = 0; j < MIXMAX_D; ++j) { + std::swap(a[c][j], a[piv][j]); + std::swap(out.m[c][j], out.m[piv][j]); + } + } + std::uint64_t pinv = inv(a[c][c]); + for (int j = 0; j < MIXMAX_D; ++j) { + a[c][j] = mul(a[c][j], pinv); + out.m[c][j] = mul(out.m[c][j], pinv); + } + for (int row = 0; row < MIXMAX_D; ++row) { + if (row == c) { + continue; + } + std::uint64_t f = a[row][c] % P; + if (!f) { + continue; + } + for (int j = 0; j < MIXMAX_D; ++j) { + a[row][j] = sub(a[row][j], mul(f, a[c][j])); + out.m[row][j] = sub(out.m[row][j], mul(f, out.m[c][j])); + } + } + } + return out; +} + +// butterfly level 0: even slot i -> e_0, odd slot -> T_0(A) e_0 +__global__ void kernel_q_seed(std::uint64_t* q) { + std::size_t n = std::size_t(blockIdx.x) * blockDim.x + threadIdx.x; + if (n >= Q_TABLE_SIZE / 2) { + return; + } + std::uint64_t y[MIXMAX_D]; + for (int k = 0; k < MIXMAX_D; ++k) { + y[k] = (k == 0); + } + for (int k = 0; k < MIXMAX_D; ++k) { + q[2 * n * MIXMAX_D + k] = y[k]; + } + mixmax_engine::apply_skip_bit(y, 0); + for (int k = 0; k < MIXMAX_D; ++k) { + q[(2 * n + 1) * MIXMAX_D + k] = y[k]; + } +} + +// butterfly level r: slot with bit r set <- T_r(A) applied to its bit-r-clear sibling +__global__ void kernel_q_level(std::uint64_t* q, int r) { + std::size_t node = std::size_t(blockIdx.x) * blockDim.x + threadIdx.x; + if (node >= Q_TABLE_SIZE / 2) { + return; + } + const std::size_t bit = std::size_t(1) << r; + const std::size_t dst = ((node >> r) << (r + 1)) | bit | (node & (bit - 1)); + const std::size_t src = dst ^ bit; + std::uint64_t y[MIXMAX_D]; + for (int k = 0; k < MIXMAX_D; ++k) { + y[k] = q[src * MIXMAX_D + k]; + } + mixmax_engine::apply_skip_bit(y, r); + for (int k = 0; k < MIXMAX_D; ++k) { + q[dst * MIXMAX_D + k] = y[k]; + } +} + +// Q_i(A) e_0 (state vector) -> coefficients of Q_i in the A-power basis +__global__ void kernel_q_to_coeff(std::uint64_t* q, MixmaxKrInv kr_inv) { + std::size_t i = std::size_t(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= Q_TABLE_SIZE) { + return; + } + std::uint64_t v[MIXMAX_D], out[MIXMAX_D]; + for (int k = 0; k < MIXMAX_D; ++k) { + v[k] = q[i * MIXMAX_D + k]; + } + for (int row = 0; row < MIXMAX_D; ++row) { + std::uint64_t acc = 0; + for (int col = 0; col < MIXMAX_D; ++col) { + acc = mixmax_engine::mod_mul_add(acc, kr_inv.m[row][col], v[col]); + } + out[row] = acc; + } + for (int k = 0; k < MIXMAX_D; ++k) { + q[i * MIXMAX_D + k] = out[k]; + } +} + +// per-(re)seed: V_i = Q_i(A) . prefix, one apply_poly per engine, single launch +__global__ void kernel_rng_seed_apply( + mixmax_engine* engines, + std::size_t count, + const std::uint64_t* q, + MixmaxSeedPrefix prefix +) { + std::size_t i = std::size_t(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= count) { + return; + } + std::uint64_t y[MIXMAX_D]; + for (int k = 0; k < MIXMAX_D; ++k) { + y[k] = prefix.state[k]; + } + mixmax_engine::apply_poly(y, &q[i * MIXMAX_D]); + std::uint64_t* dst = engines[i].state_data(); + for (int k = 0; k < MIXMAX_D; ++k) { + dst[k] = y[k]; + } + engines[i].finalize_skipped_state(); +} + +inline std::size_t rng_grid_1d(std::size_t n) { + return (n + RNG_SEED_BLOCK - 1) / RNG_SEED_BLOCK; +} + +// Grid geometry for a draw of `count` numbers from the engine bank: split evenly over +// as few threads as cover `count` in ceil(count / RNG_ENGINE_COUNT) draws per thread, +// so the grid never has a large idle remainder. +inline std::pair rng_grid(std::size_t count) { + std::size_t draws = (count + RNG_ENGINE_COUNT - 1) / RNG_ENGINE_COUNT; + std::size_t needed = (count + draws - 1) / draws; + std::size_t threads = std::min( + MAX_THREADS_PER_BLOCK, + ((needed + THREADS_MULTIPLE - 1) / THREADS_MULTIPLE) * THREADS_MULTIPLE + ); + std::size_t blocks = (needed + threads - 1) / threads; + return {blocks, threads}; +} + +// Bank of RNG_ENGINE_COUNT MIXMAX engines in device memory, seeded from a DerivedSeed. +// Re-seeding only records the seed (and the shared-bit prefix); the actual (re-)seed +// runs lazily on the next run() call, on the main (RNG) stream. The 2^16 composed skip +// operators are built once on first use. +class GpuRandom { +public: + GpuRandom() { + void* ptr; + check_error(gpuMalloc(&ptr, sizeof(mixmax_engine) * RNG_ENGINE_COUNT)); + _engines = static_cast(ptr); + check_error( + gpuMalloc(&ptr, sizeof(std::uint64_t) * Q_TABLE_SIZE * MIXMAX_D) + ); + _q_table = static_cast(ptr); + _run_skip.update(_seed.seed_parts[0], _seed.seed_parts[1]); + _prefix = mixmax_seed_prefix(_seed.seed_parts, _run_skip); + } + ~GpuRandom() { + if (_engines) { + gpuFree(_engines); + } + if (_q_table) { + gpuFree(_q_table); + } + } + GpuRandom(const GpuRandom&) = delete; + GpuRandom& operator=(const GpuRandom&) = delete; + GpuRandom(GpuRandom&& other) noexcept : + _engines(other._engines), + _q_table(other._q_table), + _seed(other._seed), + _run_skip(other._run_skip), + _prefix(other._prefix), + _pending(other._pending), + _q_built(other._q_built) { + other._engines = nullptr; + other._q_table = nullptr; + } + GpuRandom& operator=(GpuRandom&& other) noexcept { + if (_engines) { + gpuFree(_engines); + } + if (_q_table) { + gpuFree(_q_table); + } + _engines = other._engines; + _q_table = other._q_table; + _seed = other._seed; + _run_skip = other._run_skip; + _prefix = other._prefix; + _pending = other._pending; + _q_built = other._q_built; + other._engines = nullptr; + other._q_table = nullptr; + return *this; + } + + void set_seed(DerivedSeed seed) { + // the engine index is folded into the low SEED_STREAM_BITS of the stream field, + // so callers must leave the whole reserved stream index unset + if (seed.seed_parts[3] & (Q_TABLE_SIZE - 1)) { + throw std::invalid_argument( + "GpuRandom::set_seed: stream index is reserved for the GPU engine bank" + ); + } + _seed = seed; + _run_skip.update(seed.seed_parts[0], seed.seed_parts[1]); + _prefix = mixmax_seed_prefix(seed.seed_parts, _run_skip); + _pending = true; + } + + void reseed_if_needed(gpuStream_t stream) { + if (!_q_built) { + build_q_table(stream); + _q_built = true; + } + if (!_pending) { + return; + } + MixmaxSeedPrefix prefix = _prefix; + kernel_rng_seed_apply<<< + rng_grid_1d(RNG_ENGINE_COUNT), RNG_SEED_BLOCK, 0, stream>>>( + _engines, RNG_ENGINE_COUNT, _q_table, prefix + ); + check_error(); + _pending = false; + } + + mixmax_engine* engines() { return _engines; } + +private: + void build_q_table(gpuStream_t stream) { + MixmaxKrInv kr_inv = mixmax_kr_inv(); + std::size_t blocks = rng_grid_1d(Q_TABLE_SIZE / 2); + kernel_q_seed<<>>(_q_table); + for (int r = 1; r < SEED_STREAM_BITS; ++r) { + kernel_q_level<<>>(_q_table, r); + } + kernel_q_to_coeff<<< + rng_grid_1d(Q_TABLE_SIZE), RNG_SEED_BLOCK, 0, stream>>>(_q_table, kr_inv); + check_error(); + } + + mixmax_engine* _engines = nullptr; + std::uint64_t* _q_table = nullptr; + DerivedSeed _seed; + RunSeedSkip _run_skip; + MixmaxSeedPrefix _prefix; + bool _pending = true; + bool _q_built = false; +}; + +template +void launch_rng_kernel( + F kernel, std::size_t count, GpuRandom& rng, gpuStream_t stream, Args... args +) { + if (count == 0) { + return; + } + auto [blocks, threads] = rng_grid(count); + kernel<<>>(count, rng.engines(), args...); + check_error(); +} + +} // namespace gpu +} // namespace madspace diff --git a/madspace/src/gpu/runtime.cu b/madspace/src/gpu/runtime.cu index 83d6ca4c6..5c62b5b15 100644 --- a/madspace/src/gpu/runtime.cu +++ b/madspace/src/gpu/runtime.cu @@ -4,7 +4,6 @@ #include #include #include -#include #include #include @@ -20,6 +19,7 @@ #include "../kernels/operations.hpp" #include "device.hpp" #include "madspace/util.hpp" +#include "random.cuh" #include "tensor.cuh" using namespace madspace; @@ -999,6 +999,17 @@ void op_quantile( tmp.reset(device); } +__global__ void +kernel_random(std::size_t count, mixmax_engine* engines, double* output) { + std::size_t stride = blockDim.x * gridDim.x; + std::size_t t = blockIdx.x * blockDim.x + threadIdx.x; + mixmax_engine rng = engines[t]; + for (std::size_t i = t; i < count; i += stride) { + output[i] = rng.flat(); + } + engines[t] = rng; +} + void op_random( const GpuRuntime::Instruction& instruction, TensorVec& locals, @@ -1010,28 +1021,29 @@ void op_random( output = Tensor( DataType::dt_float, {batch_size, dim}, device, instruction.output_alloc_hints[0] ); - gpurandGenerator_t generator = instruction.runtime.gpurand_generator(); - check_error(gpurandSetStream(generator, device.stream())); - check_error(gpurandGenerateUniformDouble( - generator, static_cast(output.data()), batch_size * dim - )); + launch_rng_kernel( + kernel_random, + batch_size * dim, + instruction.runtime.rng(), + device.stream(), + static_cast(output.data()) + ); } -__global__ void kernel_double_to_int_range( - std::size_t batch_size, - me_int_t max_val, - GpuTensorView double_in, - GpuTensorView int_out +__global__ void kernel_random_int( + std::size_t count, mixmax_engine* engines, me_int_t max_val, me_int_t* output ) { - me_int_t i = blockDim.x * blockIdx.x + threadIdx.x; - if (i >= batch_size) { - return; - } - me_int_t rand_int = double_in[i] * max_val; - if (rand_int >= max_val) { - rand_int = max_val - 1; + std::size_t stride = blockDim.x * gridDim.x; + std::size_t t = blockIdx.x * blockDim.x + threadIdx.x; + mixmax_engine rng = engines[t]; + for (std::size_t i = t; i < count; i += stride) { + me_int_t rand_int = rng.flat() * max_val; + if (rand_int >= max_val) { + rand_int = max_val - 1; + } + output[i] = rand_int; } - int_out[i] = rand_int; + engines[t] = rng; } void op_random_int( @@ -1042,45 +1054,40 @@ void op_random_int( auto batch_size = locals[instruction.input_indices[0]].batch_sizes()[0]; auto max_val = locals[instruction.input_indices[1]].batch_sizes()[0]; auto& output = locals[instruction.output_indices[0]]; - Tensor tmp(DataType::dt_float, {batch_size}, device, AllocHint::temporary); output = Tensor( DataType::dt_int, {batch_size}, device, instruction.output_alloc_hints[0] ); - gpurandGenerator_t generator = instruction.runtime.gpurand_generator(); - check_error(gpurandSetStream(generator, device.stream())); - check_error(gpurandGenerateUniformDouble( - generator, static_cast(tmp.data()), batch_size - )); - launch_kernel( - kernel_double_to_int_range, + launch_rng_kernel( + kernel_random_int, batch_size, + instruction.runtime.rng(), device.stream(), - batch_size, max_val, - tmp.view(), - output.view() + static_cast(output.data()) ); - tmp.reset(device); } __global__ void kernel_unweight( std::size_t batch_size, - GpuTensorView rand_in, + mixmax_engine* engines, GpuTensorView weights_in, GpuTensorView max_weights_in, GpuTensorView weights_out, GpuTensorView indices_out ) { - me_int_t i = blockDim.x * blockIdx.x + threadIdx.x; - if (i >= batch_size) { - return; - } - - auto rand = rand_in[i], weight = weights_in[i], max_weight = max_weights_in[i]; - bool accepted = max_weight * rand < weight; - auto weight_clipped = weight < max_weight ? max_weight : weight; - weights_out[i] = accepted ? weight_clipped : 0.; - indices_out[i] = accepted ? i : -1; + std::size_t stride = blockDim.x * gridDim.x; + std::size_t t = blockIdx.x * blockDim.x + threadIdx.x; + mixmax_engine rng = engines[t]; + for (std::size_t i = t; i < batch_size; i += stride) { + auto rand = rng.flat(); + auto weight = weights_in[i], max_weight = max_weights_in[i]; + auto abs_weight = fabs(weight); + bool accepted = max_weight * rand < abs_weight; + auto weight_clipped = copysign(fmax(abs_weight, max_weight), weight); + weights_out[i] = accepted ? weight_clipped : 0.; + indices_out[i] = accepted ? static_cast(i) : -1; + } + engines[t] = rng; } void op_unweight( @@ -1103,23 +1110,15 @@ void op_unweight( return; } - Tensor rand(DataType::dt_float, {batch_size}, device, AllocHint::temporary); - gpurandGenerator_t generator = instruction.runtime.gpurand_generator(); - check_error(gpurandSetStream(generator, stream)); - check_error(gpurandGenerateUniformDouble( - generator, static_cast(rand.data()), batch_size - )); - Tensor indices_tmp(DataType::dt_int, {batch_size}, device, AllocHint::temporary); Tensor uw_weights_tmp( DataType::dt_float, {batch_size}, device, AllocHint::temporary ); - launch_kernel( + launch_rng_kernel( kernel_unweight, batch_size, + instruction.runtime.rng(), stream, - batch_size, - rand.view(), weights.view(), max_weight.view(), uw_weights_tmp.view(), @@ -1156,7 +1155,6 @@ void op_unweight( ptr_all_weights, ptr_uw_weights ); - rand.reset(device); indices_tmp.reset(device); uw_weights_tmp.reset(device); indices_compacted.reset(device); @@ -1509,26 +1507,15 @@ private: GpuRuntime::GpuRuntime(const Function& function_arg, ContextPtr context) : _context(context), _input_count(function_arg.inputs().size()), - _gpublas_handle( - context->thread_pool(), + _gpublas_handle(context->global_resource( + "gpublas_handle", []() { gpublasHandle_t handle; check_error(gpublasCreate(&handle)); return handle; }, [](gpublasHandle_t handle) { check_error(gpublasDestroy(handle)); } - ), - _gpurand_generator( - context->thread_pool(), - []() { - gpurandGenerator_t handle; - check_error(gpurandCreateGenerator(&handle, GPURAND_RNG_PSEUDO_DEFAULT)); - std::random_device rand_dev; - check_error(gpurandSetPseudoRandomGeneratorSeed(handle, rand_dev())); - return handle; - }, - [](gpurandGenerator_t handle) { check_error(gpurandDestroyGenerator(handle)); } - ), + )), _prev_caches(context->thread_pool(), []() { return TensorVec{}; }), _prev_caches_backward(context->thread_pool(), []() { return TensorVec{}; }) { if (context->device()->device_type() != GpuDevice::gpu_device_type) { @@ -1555,6 +1542,9 @@ GpuRuntime::GpuRuntime(const Function& function_arg, ContextPtr context) : if (instr.stream_index >= stream_count) { stream_count = instr.stream_index + 1; } + if (instr.instruction->is_random()) { + _uses_random = true; + } } SyncTracker sync_tracker(stream_count); std::vector local_source_streams(function.locals().size(), -1); @@ -1810,6 +1800,17 @@ GpuRuntime::GpuRuntime(const Function& function_arg, ContextPtr context) : } } ); + if (_uses_random) { + _rng = context->global_resource("gpu_rng", []() { + return GpuRandom(); + }); + } +} + +void GpuRuntime::set_seed(DerivedSeed seed) { + if (_uses_random) { + rng().set_seed(seed); + } } TensorVec GpuRuntime::run(const TensorVec& inputs) { @@ -1820,6 +1821,9 @@ TensorVec GpuRuntime::run(const TensorVec& inputs) { auto locals = _locals_init; std::copy(inputs.begin(), inputs.end(), locals.begin()); gpuStream_t main_stream = streams.at(0); + if (_uses_random) { + rng().reseed_if_needed(main_stream); + } MemPool mem_pool(gpu_device, load_pool_size_cache(false), main_stream); for (auto& instr : _instructions) { @@ -1842,7 +1846,7 @@ TensorVec GpuRuntime::run(const TensorVec& inputs) { check_error(gpuStreamWaitEvent(main_stream, events.at(event))); } update_pool_size_cache(mem_pool.total_sizes(), false); - //update_cached_tensors(mem_pool.reset(main_stream), false); + // update_cached_tensors(mem_pool.reset(main_stream), false); TensorVec outputs; for (auto index : _output_indices) { outputs.push_back(locals[index]); @@ -1867,6 +1871,9 @@ std::tuple> GpuRuntime::run_with_grad( input_requires_grad.begin(), input_requires_grad.end(), requires_grad.begin() ); gpuStream_t main_stream = streams.at(0); + if (_uses_random) { + rng().reseed_if_needed(main_stream); + } MemPool mem_pool(gpu_device, load_pool_size_cache(false), main_stream); for (auto [instr, instr_eval_grad] : zip(_instructions, eval_grad)) { @@ -1911,7 +1918,7 @@ std::tuple> GpuRuntime::run_with_grad( check_error(gpuStreamWaitEvent(main_stream, events.at(event))); } update_pool_size_cache(mem_pool.total_sizes(), false); - //update_cached_tensors(mem_pool.reset(main_stream), false); + // update_cached_tensors(mem_pool.reset(main_stream), false); TensorVec outputs; for (auto index : _output_indices) { outputs.push_back(locals[index]); @@ -1986,7 +1993,7 @@ std::pair GpuRuntime::run_backward( check_error(gpuStreamWaitEvent(main_stream, events.at(event))); }*/ update_pool_size_cache(mem_pool.total_sizes(), true); - //update_cached_tensors(mem_pool.reset(main_stream), true); + // update_cached_tensors(mem_pool.reset(main_stream), true); check_error(gpuStreamSynchronize(main_stream)); return { {local_grads.begin(), local_grads.begin() + _input_count}, @@ -1999,8 +2006,8 @@ GpuRuntime::load_pool_size_cache(bool backward) { auto cache = backward ? _pool_size_cache_backward.load() : _pool_size_cache.load(); std::vector> ret; if (cache) { - //auto& thread_prev_caches = - //backward ? _prev_caches_backward.get() : _prev_caches.get(); + // auto& thread_prev_caches = + // backward ? _prev_caches_backward.get() : _prev_caches.get(); for (auto [pool_index, size] : *cache) { Tensor new_cache = _context->cached_tensor(size); /*if (pool_index < thread_prev_caches.size()) { diff --git a/madspace/src/gpu/runtime.hpp b/madspace/src/gpu/runtime.hpp index 2383e6e25..8ec869375 100644 --- a/madspace/src/gpu/runtime.hpp +++ b/madspace/src/gpu/runtime.hpp @@ -4,9 +4,12 @@ #include "madspace/compgraphs/function.hpp" #include "madspace/driver/backend.hpp" #include "madspace/driver/tensor.hpp" +#include "random.cuh" -#include +#include #include +#include +#include namespace madspace { namespace gpu { @@ -42,9 +45,10 @@ class GpuRuntime : public Runtime { const std::vector& eval_grad, bool return_contiguous_grads ) override; + void set_seed(DerivedSeed seed) override; Context& context() { return *_context; } gpublasHandle_t gpublas_handle() { return _gpublas_handle.get(); } - gpurandGenerator_t gpurand_generator() { return _gpurand_generator.get(); } + GpuRandom& rng() { return _rng->get().get(); } private: std::vector> @@ -69,8 +73,9 @@ class GpuRuntime : public Runtime { ThreadResource> _events; std::vector _wait_events; std::vector _backward_wait_events; - ThreadResource _gpublas_handle; - ThreadResource _gpurand_generator; + ThreadResource& _gpublas_handle; + std::optional>> _rng; + bool _uses_random = false; std::atomic>> _pool_size_cache; std::atomic>> diff --git a/madspace/src/phasespace/channel_weight_network.cpp b/madspace/src/phasespace/channel_weight_network.cpp index 31f5ee860..343468299 100644 --- a/madspace/src/phasespace/channel_weight_network.cpp +++ b/madspace/src/phasespace/channel_weight_network.cpp @@ -68,8 +68,10 @@ NamedVector ChannelWeightNetwork::build_function_impl( }; } -void ChannelWeightNetwork::initialize_globals(ContextPtr context) const { - _mlp.initialize_globals(context); +void ChannelWeightNetwork::initialize_globals( + ContextPtr context, std::optional seed +) const { + _mlp.initialize_globals(context, seed); context->define_global(_mask_name, DataType::dt_float, {_channel_count}); bool is_cpu = context->device() == cpu_device(); diff --git a/madspace/src/phasespace/discrete_flow.cpp b/madspace/src/phasespace/discrete_flow.cpp index e7119e2e1..5f283fcb1 100644 --- a/madspace/src/phasespace/discrete_flow.cpp +++ b/madspace/src/phasespace/discrete_flow.cpp @@ -139,13 +139,15 @@ Mapping::Result DiscreteFlow::build_transform( }; } -void DiscreteFlow::initialize_globals(ContextPtr context) const { +void DiscreteFlow::initialize_globals( + ContextPtr context, std::optional seed +) const { if (_first_prob_name) { initialize_uniform_probs( context, _first_prob_name.value(), _option_counts.at(0) ); } for (auto& subnet : _subnets) { - subnet.initialize_globals(context); + subnet.initialize_globals(context, seed); } } diff --git a/madspace/src/phasespace/flow.cpp b/madspace/src/phasespace/flow.cpp index d3eefd906..1f353fcf1 100644 --- a/madspace/src/phasespace/flow.cpp +++ b/madspace/src/phasespace/flow.cpp @@ -211,17 +211,19 @@ Flow::Flow( } } -void Flow::initialize_globals(ContextPtr context) const { +void Flow::initialize_globals( + ContextPtr context, std::optional seed +) const { for (auto& block : _coupling_blocks) { - block.subnet1.initialize_globals(context); - block.subnet2.initialize_globals(context); + block.subnet1.initialize_globals(context, seed); + block.subnet2.initialize_globals(context, seed); } } void Flow::initialize_from_vegas( - ContextPtr context, const std::string& grid_name + ContextPtr context, const std::string& grid_name, std::optional seed ) const { - initialize_globals(context); + initialize_globals(context, seed); auto& last_block = _coupling_blocks.at(_coupling_blocks.size() - 1); vegas_init( context, diff --git a/madspace/src/phasespace/integrand.cpp b/madspace/src/phasespace/integrand.cpp index cd9d189f4..b2e74cb3d 100644 --- a/madspace/src/phasespace/integrand.cpp +++ b/madspace/src/phasespace/integrand.cpp @@ -14,13 +14,16 @@ std::size_t final_channel_count( std::size_t first_remapped_chan_count, const std::vector& second_chan_weight_remap, std::size_t second_remapped_chan_count, + const std::optional& prop_chan_weights, const std::optional& subchan_weights ) { if (second_chan_weight_remap.size() > 0) { return second_remapped_chan_count; } else if (subchan_weights) { return subchan_weights->channel_count(); - } else if (first_chan_weight_remap.size() > 0) { + } else if (first_chan_weight_remap.size() > 0 || prop_chan_weights) { + // with the denominators sde_strategy, prop_chan_weights already + // produces first_remapped_chan_count weights without needing a remap return first_remapped_chan_count; } else { return diff_xs.at(0).matrix_element().diagram_count(); @@ -86,6 +89,7 @@ Integrand::Integrand( first_remapped_chan_count, second_chan_weight_remap, second_remapped_chan_count, + prop_chan_weights, subchan_weights ); if (madnis_channel_count > 1 && @@ -653,6 +657,7 @@ NamedVector Integrand::build_common_part( _first_remapped_chan_count, _second_chan_weight_remap, _second_remapped_chan_count, + _prop_chan_weights, _subchan_weights ); bool use_compressed_channel_weights = diff --git a/madspace/src/phasespace/mlp.cpp b/madspace/src/phasespace/mlp.cpp index 920f36594..4de6a6537 100644 --- a/madspace/src/phasespace/mlp.cpp +++ b/madspace/src/phasespace/mlp.cpp @@ -1,7 +1,8 @@ #include "madspace/phasespace/mlp.hpp" #include -#include + +#include "madspace/driver/random.hpp" using namespace madspace; @@ -51,11 +52,9 @@ void initialize_layer( std::size_t output_dim, const std::string& prefix, int layer_index, - std::mt19937& rand_gen, + std::optional seed, bool zeros ) { - double bound = 1 / std::sqrt(input_dim); - std::uniform_real_distribution rand_dist(-bound, bound); auto weight_name = prefixed_name(prefix, std::format("layer{}.weight", layer_index)); auto bias_name = prefixed_name(prefix, std::format("layer{}.bias", layer_index)); @@ -75,14 +74,35 @@ void initialize_layer( } auto weight_view = weight_tensor.view()[0]; - for (std::size_t i = 0; i < output_dim; ++i) { - for (std::size_t j = 0; j < input_dim; ++j) { - weight_view[i][j] = zeros ? 0. : rand_dist(rand_gen); - } - } auto bias_view = bias_tensor.view()[0]; - for (std::size_t i = 0; i < output_dim; ++i) { - bias_view[i] = zeros ? 0. : rand_dist(rand_gen); + if (zeros) { + for (std::size_t i = 0; i < output_dim; ++i) { + for (std::size_t j = 0; j < input_dim; ++j) { + weight_view[i][j] = 0.; + } + bias_view[i] = 0.; + } + } else { + double bound = 1 / std::sqrt(input_dim); + auto uniform = [&](MixMaxRandom& gen) { + return bound * (2. * gen.generate_double() - 1.); + }; + // Independently seeded per tensor, via a fresh unique_seed_index() rather + // than a hash of the tensor's name. + MixMaxRandom weight_rand_gen( + DerivedSeed(seed, DerivedSeed::global_init, context->unique_seed_index()) + ); + for (std::size_t i = 0; i < output_dim; ++i) { + for (std::size_t j = 0; j < input_dim; ++j) { + weight_view[i][j] = uniform(weight_rand_gen); + } + } + MixMaxRandom bias_rand_gen( + DerivedSeed(seed, DerivedSeed::global_init, context->unique_seed_index()) + ); + for (std::size_t i = 0; i < output_dim; ++i) { + bias_view[i] = uniform(bias_rand_gen); + } } if (!is_cpu) { @@ -133,15 +153,15 @@ MLP::build_function_impl(FunctionBuilder& fb, const NamedVector& args) co }; } -void MLP::initialize_globals(ContextPtr context) const { - std::random_device rand_device; - std::mt19937 rand_gen(rand_device()); +void MLP::initialize_globals( + ContextPtr context, std::optional seed +) const { std::size_t dim = _input_dim; for (std::size_t i = 1; i < _layers; ++i) { - initialize_layer(context, dim, _hidden_dim, _prefix, i, rand_gen, false); + initialize_layer(context, dim, _hidden_dim, _prefix, i, seed, false); dim = _hidden_dim; } - initialize_layer(context, dim, _output_dim, _prefix, _layers, rand_gen, true); + initialize_layer(context, dim, _output_dim, _prefix, _layers, seed, true); } std::vector MLP::global_names() const { diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 4d5d31c88..c4e8da799 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -842,7 +842,12 @@ PYBIND11_MODULE(_madspace_py, m) { ) .def("input_dim", &MLP::input_dim) .def("output_dim", &MLP::output_dim) - .def("initialize_globals", &MLP::initialize_globals, py::arg("context")); + .def( + "initialize_globals", + &MLP::initialize_globals, + py::arg("context"), + py::arg("seed") = std::nullopt + ); py::classh(m, "Flow") .def( @@ -866,12 +871,18 @@ PYBIND11_MODULE(_madspace_py, m) { ) .def("input_dim", &Flow::input_dim) .def("condition_dim", &Flow::condition_dim) - .def("initialize_globals", &Flow::initialize_globals, py::arg("context")) + .def( + "initialize_globals", + &Flow::initialize_globals, + py::arg("context"), + py::arg("seed") = std::nullopt + ) .def( "initialize_from_vegas", &Flow::initialize_from_vegas, py::arg("context"), - py::arg("grid_name") + py::arg("grid_name"), + py::arg("seed") = std::nullopt ); py::classh( @@ -927,7 +938,8 @@ PYBIND11_MODULE(_madspace_py, m) { .def( "initialize_globals", &ChannelWeightNetwork::initialize_globals, - py::arg("context") + py::arg("context"), + py::arg("seed") = std::nullopt ); py::classh(m, "DiscreteHistogram") @@ -972,7 +984,10 @@ PYBIND11_MODULE(_madspace_py, m) { .def("option_counts", &DiscreteFlow::option_counts) .def("condition_dim", &DiscreteFlow::condition_dim) .def( - "initialize_globals", &DiscreteFlow::initialize_globals, py::arg("context") + "initialize_globals", + &DiscreteFlow::initialize_globals, + py::arg("context"), + py::arg("seed") = std::nullopt ); py::classh(m, "VegasGridOptimizer") @@ -1399,12 +1414,14 @@ PYBIND11_MODULE(_madspace_py, m) { ContextPtr, const MadnisTraining::Config&, const std::vector>&, - const std::optional&>(), + const std::optional&, + std::optional>(), py::arg("generator_context"), py::arg("optimizer_context"), py::arg("config"), py::arg("integrands"), - py::arg("cwnet") + py::arg("cwnet"), + py::arg("seed") = std::nullopt ) .def("train_step", &MadnisTraining::train_step, py::arg("batch_index")) .def("active_channels", &MadnisTraining::active_channels) @@ -1435,12 +1452,14 @@ PYBIND11_MODULE(_madspace_py, m) { ContextPtr, const std::vector&, Verbosity, - std::shared_ptr>(), + std::shared_ptr, + std::optional>(), py::arg("generator_context"), py::arg("optimizer_context"), py::arg("training_args"), py::arg("verbosity"), - py::arg("status_file") = std::shared_ptr() + py::arg("status_file") = std::shared_ptr(), + py::arg("seed") = std::nullopt ) .def("train", &MultiMadnisTraining::train) .def("active_channels", &MultiMadnisTraining::active_channels); @@ -1474,7 +1493,25 @@ PYBIND11_MODULE(_madspace_py, m) { .def_readwrite( "cut_efficiency_threshold", &GeneratorConfig::cut_efficiency_threshold ) - .def_readwrite("max_cut_repetitions", &GeneratorConfig::max_cut_repetitions); + .def_readwrite("max_cut_repetitions", &GeneratorConfig::max_cut_repetitions) + .def_readwrite( + "finish_remaining_fraction", &GeneratorConfig::finish_remaining_fraction + ) + .def_readwrite("max_batch_fraction", &GeneratorConfig::max_batch_fraction) + .def_readwrite( + "batch_overshoot_sigma", &GeneratorConfig::batch_overshoot_sigma + ); + + m.def( + "compute_generation_batch_event_count", + &compute_generation_batch_event_count, + py::arg("count_target"), + py::arg("count_unweighted"), + py::arg("count_opt"), + py::arg("abs_cross_section_count"), + py::arg("abs_cross_section_rel_error"), + py::arg("config") + ); py::classh(m, "GeneratorStatus") .def(py::init<>()) @@ -1482,6 +1519,8 @@ PYBIND11_MODULE(_madspace_py, m) { .def_readwrite("name", &GeneratorStatus::name) .def_readwrite("mean", &GeneratorStatus::mean) .def_readwrite("error", &GeneratorStatus::error) + .def_readwrite("mean_abs", &GeneratorStatus::mean_abs) + .def_readwrite("error_abs", &GeneratorStatus::error_abs) .def_readwrite("rel_std_dev", &GeneratorStatus::rel_std_dev) .def_readwrite("count", &GeneratorStatus::count) .def_readwrite("count_opt", &GeneratorStatus::count_opt) @@ -1662,9 +1701,9 @@ PYBIND11_MODULE(_madspace_py, m) { "diagram_propagator_pdgs", &LHECompleter::SubprocArgs::diagram_propagator_pdgs ); - py::classh(m, "RandGen") + py::classh(m, "MixMaxRandom") .def(py::init<>()) - .def(py::init(), py::arg("seed")); + .def(py::init(), py::arg("seed")); py::classh(m, "LHECompleter") .def( py::init&, double>(), @@ -1762,10 +1801,12 @@ PYBIND11_MODULE(_madspace_py, m) { py::init< const std::vector&, const std::vector>&, + std::uint64_t, std::shared_ptr, const GeneratorConfig&>(), py::arg("contexts"), py::arg("channels"), + py::arg("seed"), py::arg("status_file") = std::shared_ptr(), py::arg_v( "config", @@ -1773,7 +1814,7 @@ PYBIND11_MODULE(_madspace_py, m) { "EventGenerator.default_config" ) ) - .def("survey", &EventGenerator::survey) + .def("survey", &EventGenerator::survey, py::arg("survey_pass") = 0) .def("generate", &EventGenerator::generate) .def( "combine_to_compact_npy", diff --git a/madspace/tests/test_batch_scheduling.py b/madspace/tests/test_batch_scheduling.py new file mode 100644 index 000000000..82a0d7fd4 --- /dev/null +++ b/madspace/tests/test_batch_scheduling.py @@ -0,0 +1,114 @@ +"""Unit tests for compute_generation_batch_event_count, the pure function that +sizes a channel's next steady-state generation batch (see +generator_data.hpp/.cpp). The result is a raw (pre-unweighting) event count; +device-specific job splitting happens later, at dispatch time. +""" + +import math + +import madspace as ms + + +def make_config(): + config = ms.GeneratorConfig() + config.finish_remaining_fraction = 0.05 + config.max_batch_fraction = 0.6 + config.batch_overshoot_sigma = 1.0 + return config + + +def batch_events( + count_target, + count_unweighted, + count_opt, + cross_section_count, + cross_section_rel_error, + config, +): + return ms.compute_generation_batch_event_count( + count_target, + count_unweighted, + count_opt, + cross_section_count, + cross_section_rel_error, + config, + ) + + +def test_confident_channel_hits_cap(): + # A converged channel (tiny relative error), far from its target, should + # be allowed to step close to the max_batch_fraction cap instead of + # crawling towards the target. + config = make_config() + events = batch_events(100000, 0.0, 0, 10000, 0.001, config) + # efficiency defaults to 1 (count_opt == 0), so no raw/unweighted inflation. + expected = math.ceil(0.6 * 100000) + assert events == expected + + +def test_uncertain_channel_is_throttled(): + # A channel with a large cross-section relative error should get a much + # smaller batch than a converged one, since the safe lower bound on its + # target sits far below the point estimate. + config = make_config() + confident_events = batch_events(100000, 0.0, 0, 10000, 0.001, config) + uncertain_events = batch_events(100000, 0.0, 0, 10000, 0.5, config) + assert uncertain_events < confident_events + + +def test_no_data_is_cautious_not_full_target(): + # No cross-section data yet (pessimistic rel_error) must NOT be mistaken + # for "already done" just because the confidence bound collapses to zero + # -- it should fall through to a small exploratory request, not the full + # target in one blind shot. This is deliberately independent of the + # finish_remaining_fraction short-circuit, which only looks at how much + # is left, not at how well it's known. + config = make_config() + no_data_events = batch_events(100000, 0.0, 0, 0, 0.0, config) + confident_events = batch_events(100000, 0.0, 0, 10000, 0.001, config) + assert 0 < no_data_events < confident_events + + +def test_few_events_left_finishes_outright(): + # Once the remaining deficit has shrunk below finish_remaining_fraction of + # the target, finish outright with the true deficit rather than + # continuing to throttle. + config = make_config() + events = batch_events(100000, 97000.0, 0, 10000, 0.5, config) + assert events == 3000 + + +def test_mid_generation_is_meaningfully_throttled(): + # Away from both the cap regime and the finish-outright regime, the + # requested amount should still be nonzero, well below the naive full + # deficit, and shrink relative to a more certain estimate. + config = make_config() + events = batch_events(100000, 50000.0, 0, 10000, 0.3, config) + assert 0 < events < 50000 + + +def test_always_at_least_one_event(): + # Even when the confidence bound collapses to zero (no data yet, far from + # target), the result must be strictly positive so a full exploratory + # batch still gets dispatched at the device level. + config = make_config() + events = batch_events(100000, 0.0, 0, 0, 0.0, config) + assert events >= 1 + + +def test_never_overshoots_true_remaining(): + # The requested event count must never exceed the true remaining deficit. + config = make_config() + config.max_batch_fraction = 1.0 + config.batch_overshoot_sigma = 0.0 + events = batch_events(100000, 99500.0, 0, 10000, 0.001, config) + assert events <= 500 + + +def test_zero_efficiency_does_not_diverge(): + # count_opt > 0 but zero observed successes must not blow up the raw event + # request (division by zero) or the caller (a real device batch) would + # never make progress against a runaway target. + config = make_config() + events = batch_events(100000, 0.0, 5000, 10000, 0.5, config) + assert events < 100000 * 5000 diff --git a/madspace/tests/test_format.py b/madspace/tests/test_format.py new file mode 100644 index 000000000..5cfa79be4 --- /dev/null +++ b/madspace/tests/test_format.py @@ -0,0 +1,134 @@ +"""Unit tests for the string-formatting helpers in format.cpp/format.hpp, +exposed to Python for testing (format_si_prefix, format_with_error, +format_progress). None of these had direct unit tests before; they were only +exercised indirectly through the pretty/log console output of a full +event-generation run. +""" + +import madspace as ms + +# --- format_si_prefix ------------------------------------------------------ + + +def test_si_prefix_below_thousand_is_bare_integer(): + assert ms.format_si_prefix(999) == "999" + assert ms.format_si_prefix(42) == "42" + + +def test_si_prefix_kilo(): + assert ms.format_si_prefix(1000) == "1.00k" + assert ms.format_si_prefix(1500) == "1.50k" + + +def test_si_prefix_kilo_variable_precision(): + # digits after the dot shrink as the value approaches the next power of ten, + # to keep three significant digits + assert ms.format_si_prefix(12345) == "12.3k" + + +def test_si_prefix_mega(): + assert ms.format_si_prefix(1_000_000) == "1.00M" + + +def test_si_prefix_rounds_up_into_next_prefix(): + # 999999 rounds to 1000 at 0-decimal precision within the "k" prefix + assert ms.format_si_prefix(999_999) == "1000k" + + +def test_si_prefix_falls_back_beyond_tera(): + # no prefix covers 10^15, so it prints the plain (exponential) number + assert ms.format_si_prefix(1e15) == "1e+15" + + +# --- format_with_error ------------------------------------------------------ + + +def test_with_error_normal_case_keeps_digit_notation(): + assert ms.format_with_error(1.234, 0.056) == "1.234(56)" + + +def test_with_error_single_significant_digit_error(): + assert ms.format_with_error(123.456, 1.2) == "123.5(1.2)" + + +def test_with_error_small_error_extends_precision(): + assert ms.format_with_error(1.0, 0.001) == "1.0000(10)" + + +def test_with_error_large_magnitude_uses_exponential_form(): + assert ms.format_with_error(1e10, 1e6) == "1.00000(10)e+10" + + +def test_with_error_small_value_large_error_falls_back(): + assert ms.format_with_error(0.001, 0.05) == "0.001 ± 0.050" + + +def test_with_error_negative_value_large_error_falls_back(): + assert ms.format_with_error(-0.02, 0.3) == "-0.02 ± 0.30" + + +def test_with_error_zero_value_falls_back(): + assert ms.format_with_error(0.0, 0.1) == "0.00 ± 0.10" + + +def test_with_error_confidently_negative_value_keeps_digit_notation(): + # a negative integral with a small relative error must not hit the + # log10(negative) NaN that the unguarded formula would produce + result = ms.format_with_error(-5.0, 0.1) + assert "nan" not in result.lower() + assert result == "-5.00(10)" + + +def test_with_error_nan_error_does_not_crash(): + result = ms.format_with_error(1234.5, float("nan")) + assert "nan" not in result.lower() + + +def test_with_error_zero_error_does_not_crash(): + result = ms.format_with_error(0.0, 0.0) + assert "nan" not in result.lower() + + +def test_with_error_nonfinite_value_does_not_crash(): + result = ms.format_with_error(float("inf"), 1.0) + assert isinstance(result, str) + + +def test_with_error_nonfinite_both_does_not_crash(): + result = ms.format_with_error(float("nan"), float("nan")) + assert isinstance(result, str) + + +# --- format_progress --------------------------------------------------------- + + +def test_progress_zero_is_all_blank(): + assert ms.format_progress(0.0, 10) == " " * 10 + + +def test_progress_full_is_all_blocks(): + assert ms.format_progress(1.0, 10) == "█" * 10 + + +def test_progress_half_fills_half_the_width(): + assert ms.format_progress(0.5, 10) == "█████ " + + +def test_progress_uses_partial_block_characters(): + # 0.5 of a width-5 bar lands mid-cell, so the boundary cell is a partial glyph + result = ms.format_progress(0.5, 5) + assert result == "██▌ " + assert len(result) == 5 + + +def test_progress_clamps_negative_to_zero(): + assert ms.format_progress(-1.0, 5) == " " * 5 + + +def test_progress_clamps_above_one_to_full(): + assert ms.format_progress(2.0, 5) == "█" * 5 + + +def test_progress_output_length_matches_width(): + for progress in [0.0, 0.1, 0.37, 0.5, 0.99, 1.0]: + assert len(ms.format_progress(progress, 20)) == 20 diff --git a/madspace/tests/test_lhe.py b/madspace/tests/test_lhe.py index 46d3d0d1a..f5338c2a2 100644 --- a/madspace/tests/test_lhe.py +++ b/madspace/tests/test_lhe.py @@ -148,7 +148,7 @@ def external_particles(event): @pytest.fixture(scope="module") def events(lhe_completer, mapping): p_ext = sample_external_momenta(mapping, 300, seed=1234) - rand_gen = ms.RandGen(2024) + rand_gen = ms.MixMaxRandom(2024) result = [] for row in p_ext: event = build_event(row) @@ -309,7 +309,7 @@ def test_external_color_flows_match_input_across_subprocesses( _, _, subproc_args = topology_and_args lhe_completer = ms.LHECompleter([subproc_args, subproc_args], bw_cutoff=BW_CUTOFF) p_ext = sample_external_momenta(mapping, 20, seed=99) - rand_gen = ms.RandGen(7) + rand_gen = ms.MixMaxRandom(7) color_flows = subproc_meta["color_flows"] for subprocess_index in (0, 1): @@ -329,7 +329,7 @@ def test_external_color_flows_match_input_across_subprocesses( def test_external_spins_match_helicity_table(lhe_completer, mapping, subproc_meta): helicities = subproc_meta["helicities"] p_ext = sample_external_momenta(mapping, 5, seed=55) - rand_gen = ms.RandGen(3) + rand_gen = ms.MixMaxRandom(3) for helicity_index in [0, 3, 10]: for row in p_ext: event = build_event(row) @@ -343,7 +343,7 @@ def test_external_spins_match_helicity_table(lhe_completer, mapping, subproc_met def test_external_flavors_are_valid_options(lhe_completer, mapping, subproc_meta): p_ext = sample_external_momenta(mapping, 30, seed=77) - rand_gen = ms.RandGen(9) + rand_gen = ms.MixMaxRandom(9) for flavor_index, flavor in enumerate(subproc_meta["flavors"]): options = [tuple(option) for option in flavor["options"]] for row in p_ext: @@ -363,8 +363,10 @@ def test_save_load_roundtrip(lhe_completer, mapping): p_ext = sample_external_momenta(mapping, 20, seed=321) for row in p_ext: event_a, event_b = build_event(row), build_event(row) - lhe_completer.complete_event_data(event_a, 0, 0, 0, 0, 0, ms.RandGen(11)) - loaded.complete_event_data(event_b, 0, 0, 0, 0, 0, ms.RandGen(11)) + lhe_completer.complete_event_data( + event_a, 0, 0, 0, 0, 0, ms.MixMaxRandom(11) + ) + loaded.complete_event_data(event_b, 0, 0, 0, 0, 0, ms.MixMaxRandom(11)) assert len(event_a.particles) == len(event_b.particles) for pa, pb in zip(event_a.particles, event_b.particles): assert pa.pdg_id == pb.pdg_id @@ -379,11 +381,11 @@ def test_wrong_particle_count_raises(lhe_completer): event = ms.LHEEvent() event.particles = [ms.LHEParticle()] * 3 with pytest.raises(RuntimeError): - lhe_completer.complete_event_data(event, 0, 0, 0, 0, 0, ms.RandGen(1)) + lhe_completer.complete_event_data(event, 0, 0, 0, 0, 0, ms.MixMaxRandom(1)) def test_invalid_color_index_raises(lhe_completer, mapping): p_ext = sample_external_momenta(mapping, 1, seed=1) event = build_event(p_ext[0]) with pytest.raises(RuntimeError): - lhe_completer.complete_event_data(event, 0, 0, 999, 0, 0, ms.RandGen(1)) + lhe_completer.complete_event_data(event, 0, 0, 999, 0, 0, ms.MixMaxRandom(1)) diff --git a/tests/acceptance_tests/test_mg7_reproducibility.py b/tests/acceptance_tests/test_mg7_reproducibility.py new file mode 100644 index 000000000..534560a89 --- /dev/null +++ b/tests/acceptance_tests/test_mg7_reproducibility.py @@ -0,0 +1,378 @@ +################################################################################ +# +# Copyright (c) 2009 The MadGraph5_aMC@NLO Development team and Contributors +# +# This file is a part of the MadGraph5_aMC@NLO 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 MadGraph5_aMC@NLO license which should accompany this +# distribution. +# +# For more information, visit madgraph.phys.ucl.ac.be and amcatnlo.web.cern.ch +# +################################################################################ +"""MG7 (madspace) reproducibility tests. + +Guards the seeded event-generation path: given the same run_card ``seed``, +independent mg7 runs must produce byte-identical LHE output (hashed with +sha256); a different seed must change the result. Covers: + + * ``test_vegas_reproducibility_mg7`` -- a plain survey (VEGAS-optimized) + + generate run, through ``bin/generate_events``. + * ``test_madnis_reproducible_mode_mg7`` -- a full survey + madnis training + + generate run, through ``bin/generate_events`` (no gridpack export): + training itself is required to be thread-count-independent and + byte-identical for the same seed, both with online-only training and with + buffered (off-policy replay) training enabled. GPU multi-channel batches + are not covered by this test. + * ``test_gridpack_reproducibility_mg7`` -- a gridpack trained with madnis, + then standalone event generation from that gridpack (its own + ``bin/generate_events --seed``), which is the normal way a gridpack is + used and does not re-run survey/training. + * ``test_gridpack_reproducibility_vegas_mg7`` -- same as above, but for a + plain VEGAS-optimized gridpack (no madnis training). + +Process: ``p p > t t~`` (hadronic, needs the NNPDF23_lo_as_0130_qed PDF grid; +self-skips if the mg7 runtime stack / PDF is unavailable, same as +test_check_xsec_processes_mg7.py). + +Run locally with e.g.:: + + ./tests/test_manager.py test_.*reproducib.*_mg7 -pA -t0 -l INFO + +One knob is read from the environment so the CI can dial it without touching +the code: + + * ``MG7_REPRO_EVENTS`` -- events per run (default 10000). +""" + +from __future__ import absolute_import +from __future__ import division + +import glob +import hashlib +import os +import shutil +import subprocess +import sys +import tempfile +import unittest + +import madgraph.interface.master_interface as MGCmd +from madgraph.various.banner import RunCardMG7 + +pjoin = os.path.join + +_PROCESS = 'p p > t t~' +_EVENTS = int(os.environ.get('MG7_REPRO_EVENTS', 10000)) +_SEED_A = 424242 +_SEED_B = 909090 + + +def _mg7_datadir_or_skip(test): + """Return an LHAPDF data dir that contains the NNPDF23_lo_as_0130_qed set, + or ``skipTest`` (on *test*) when the mg7 runtime stack (madspace + LHAPDF + + the run_card.toml default PDF) is unavailable.""" + try: + import madspace + has_mg7 = hasattr(madspace, 'ChannelEventGenerator') + except ImportError: + has_mg7 = False + if not has_mg7: + test.skipTest('mg7 runtime stack (madspace) unavailable') + + candidates = [] + if os.environ.get('LHAPDF_DATA_PATH'): + candidates.extend(os.environ['LHAPDF_DATA_PATH'].split(os.pathsep)) + try: + out = subprocess.check_output(['lhapdf-config', '--datadir'], + stderr=subprocess.DEVNULL).decode().strip() + if out: + candidates.append(out) + except Exception: + pass + for d in candidates: + if d and os.path.isdir(d) and glob.glob(pjoin(d, 'NNPDF23_lo_as_0130_qed*')): + return d + test.skipTest('NNPDF23_lo_as_0130_qed LHAPDF data not found ' + '(set $LHAPDF_DATA_PATH)') + + +def _lhe_hash(path): + """sha256 of an LHE file's physics content: everything from ```` + onward, i.e. excluding the ``
`` block. The header embeds the + run_card.toml/param_card/proc_card verbatim, so it legitimately differs + between runs that vary settings unrelated to the physics content (e.g. + the seed itself, or the cpu thread pool size) -- only the generated + events (and the cross-section info, itself a deterministic + function of the seed) should drive this hash.""" + with open(path, 'rb') as f: + content = f.read() + marker = b'
\n' + index = content.find(marker) + if index == -1: + raise AssertionError('no found in %s' % path) + return hashlib.sha256(content[index + len(marker):]).hexdigest() + + +def _find_lhe(run_path): + """Locate the LHE file produced under Events//, or fail.""" + matches = sorted(glob.glob(pjoin(run_path, 'Events', '*', 'events.lhe'))) + if not matches: + raise AssertionError('no events.lhe produced under %s' % run_path) + return matches[-1] + + +def _run(cmd, cwd, log_path, env, what): + """Run *cmd*, raising with the log tail on a non-zero exit.""" + with open(log_path, 'w') as logfh: + ret = subprocess.call(cmd, cwd=cwd, env=env, stdout=logfh, + stderr=subprocess.STDOUT) + if ret != 0: + with open(log_path) as f: + tail = ''.join(f.readlines()[-60:]) + raise AssertionError('%s failed (exit %d, see %s)\n\n%s' + % (what, ret, log_path, tail)) + + +class MG7ReproducibilityTest(unittest.TestCase): + """LHE-hash reproducibility of mg7 (madspace) event generation for a + fixed run_card seed, and non-reproducibility across distinct seeds.""" + + def setUp(self): + self.path = tempfile.mkdtemp(prefix='mg7_repro_') + + def tearDown(self): + shutil.rmtree(self.path, ignore_errors=True) + + def _output_process(self, run_name): + 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) + return run_dir + + def _set_run_card(self, toml_path, **settings): + """settings keys are 'section.key' RunCardMG7 names.""" + rc = RunCardMG7(toml_path) + for key, value in settings.items(): + rc.set(key, value, user=True) + rc.write(toml_path) + + def _generate_and_hash(self, run_dir, datadir, seed, tag, **run_card_settings): + """Set run_card seed (plus any extra 'section.key' settings), run + bin/generate_events -f from scratch, and return the sha256 of the + resulting LHE file.""" + shutil.rmtree(pjoin(run_dir, 'Events'), ignore_errors=True) + toml = pjoin(run_dir, 'Cards', 'run_card.toml') + self._set_run_card(toml, **{'run.seed': seed}, **run_card_settings) + env = dict(os.environ) + env['LHAPDF_DATA_PATH'] = datadir + _run([sys.executable, pjoin(run_dir, 'bin', 'generate_events'), '-f'], + run_dir, pjoin(run_dir, 'gen_%s.log' % tag), env, + 'mg7 generate_events') + return _lhe_hash(_find_lhe(run_dir)) + + def _generate_and_hash_gridpack(self, gridpack_dir, datadir, seed, tag): + """Run the gridpack's own bin/generate_events --seed from scratch, and + return the sha256 of the resulting LHE file.""" + shutil.rmtree(pjoin(gridpack_dir, 'Events'), ignore_errors=True) + env = dict(os.environ) + env['LHAPDF_DATA_PATH'] = datadir + _run([sys.executable, pjoin(gridpack_dir, 'bin', 'generate_events'), + '--seed', str(seed), '--events', str(_EVENTS), + '--output_format', 'lhe'], + gridpack_dir, pjoin(gridpack_dir, 'gen_%s.log' % tag), env, + 'gridpack generate_events') + return _lhe_hash(_find_lhe(gridpack_dir)) + + def test_vegas_reproducibility_mg7(self): + """Two VEGAS-optimized runs with the same seed produce byte-identical + LHE files, including when the cpu thread pool is shrunk to 1; a + different seed changes the result.""" + datadir = _mg7_datadir_or_skip(self) + run_dir = self._output_process('vegas') + toml = pjoin(run_dir, 'Cards', 'run_card.toml') + self._set_run_card( + toml, + **{ + 'generation.events': _EVENTS, + # Keep the test scoped to madspace's own seeding: LHE-level + # post-processing (systematics) is orthogonal and would only + # add runtime here. + 'postprocessing.systematics': False, + } + ) + + # Every call sets 'run.cpu_thread_pool_size' explicitly so it never + # leaks from a previous call via the run_card left on disk. + hash_a1 = self._generate_and_hash( + run_dir, datadir, _SEED_A, 'a1', **{'run.cpu_thread_pool_size': -1}) + hash_a2 = self._generate_and_hash( + run_dir, datadir, _SEED_A, 'a2', **{'run.cpu_thread_pool_size': -1}) + self.assertEqual(hash_a1, hash_a2, + 'same seed produced different LHE content') + + hash_a_single = self._generate_and_hash( + run_dir, datadir, _SEED_A, 'a_single', + **{'run.cpu_thread_pool_size': 1}) + self.assertEqual(hash_a1, hash_a_single, + 'same seed produced different LHE content with a ' + 'single-threaded cpu thread pool') + + hash_b = self._generate_and_hash( + run_dir, datadir, _SEED_B, 'b', **{'run.cpu_thread_pool_size': -1}) + self.assertNotEqual(hash_a1, hash_b, + 'different seeds produced identical LHE content') + + def test_madnis_reproducible_mode_mg7(self): + """madnis training itself (CPU codepath) -- not just event generation + -- is byte-identical for the same seed independent of the cpu thread + pool size, both with online-only training and with buffered + (off-policy replay) training enabled. Covers the deterministic + scheduling in MadnisTraining::maybe_start_generator_jobs / + start_generator_jobs (deterministic round dispatch + deferred buffer + flush) and the seeding of BufferUnweighter/BatchSampler (see + salt::madnis_train_unweight / salt::buffered_batch_sample in + random.hpp) -- both required for this test to pass reliably.""" + datadir = _mg7_datadir_or_skip(self) + run_dir = self._output_process('madnis_reproducible') + toml = pjoin(run_dir, 'Cards', 'run_card.toml') + base_settings = { + 'generation.events': _EVENTS, + 'postprocessing.systematics': False, + 'madnis.enable': True, + 'madnis.train_batches': 60, + } + + # Online-only training (buffering disabled). + self._set_run_card(toml, **base_settings, **{'madnis.buffer_capacity': 0}) + hash_online_multi = self._generate_and_hash( + run_dir, datadir, _SEED_A, 'online_multi', + **{'run.cpu_thread_pool_size': -1}) + hash_online_single = self._generate_and_hash( + run_dir, datadir, _SEED_A, 'online_single', + **{'run.cpu_thread_pool_size': 1}) + self.assertEqual( + hash_online_multi, hash_online_single, + 'madnis training (buffering disabled) produced ' + 'different LHE content with a single-threaded cpu thread pool') + + # Buffered (off-policy replay) training enabled. + self._set_run_card( + toml, + **base_settings, + **{ + 'madnis.buffer_capacity': 3000, + 'madnis.buffered_steps': 3, + 'madnis.minimum_buffer_size': 500, + } + ) + hash_buffered_multi = self._generate_and_hash( + run_dir, datadir, _SEED_A, 'buffered_multi', + **{'run.cpu_thread_pool_size': -1}) + hash_buffered_single = self._generate_and_hash( + run_dir, datadir, _SEED_A, 'buffered_single', + **{'run.cpu_thread_pool_size': 1}) + self.assertEqual( + hash_buffered_multi, hash_buffered_single, + 'madnis training (buffering enabled) produced ' + 'different LHE content with a single-threaded cpu thread pool') + + hash_buffered_b = self._generate_and_hash( + run_dir, datadir, _SEED_B, 'buffered_b', + **{'run.cpu_thread_pool_size': -1}) + self.assertNotEqual( + hash_buffered_multi, hash_buffered_b, + 'different seeds produced identical LHE content') + + def test_gridpack_reproducibility_mg7(self): + """A gridpack trained with madnis: two standalone event-generation + runs from that gridpack with the same seed produce byte-identical LHE + files; a different seed changes the result. Training itself is not + required to be deterministic -- only the trained gridpack's event + generation is under test here.""" + datadir = _mg7_datadir_or_skip(self) + run_dir = self._output_process('gridpack') + toml = pjoin(run_dir, 'Cards', 'run_card.toml') + self._set_run_card( + toml, + **{ + 'generation.events': _EVENTS, + 'postprocessing.systematics': False, + 'gridpack.save_gridpack': True, + 'madnis.enable': True, + 'madnis.train_batches': 100, + } + ) + + 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') + + gridpack_dir = pjoin(run_dir, 'Events', 'run_01', 'gridpack') + self.assertTrue(os.path.isdir(gridpack_dir), + 'gridpack was not produced under %s' % run_dir) + + hash_a1 = self._generate_and_hash_gridpack( + gridpack_dir, datadir, _SEED_A, 'a1') + hash_a2 = self._generate_and_hash_gridpack( + gridpack_dir, datadir, _SEED_A, 'a2') + self.assertEqual(hash_a1, hash_a2, + 'same seed produced different LHE content') + + hash_b = self._generate_and_hash_gridpack( + gridpack_dir, datadir, _SEED_B, 'b') + self.assertNotEqual(hash_a1, hash_b, + 'different seeds produced identical LHE content') + + def test_gridpack_reproducibility_vegas_mg7(self): + """A plain VEGAS-optimized gridpack (no madnis training): two + standalone event-generation runs from that gridpack with the same + seed produce byte-identical LHE files; a different seed changes the + result. Companion to test_gridpack_reproducibility_mg7, which covers + the madnis-trained case; the survey/optimization itself is not + required to be deterministic here either -- only the gridpack's own + event generation is under test.""" + datadir = _mg7_datadir_or_skip(self) + run_dir = self._output_process('gridpack_vegas') + toml = pjoin(run_dir, 'Cards', 'run_card.toml') + self._set_run_card( + toml, + **{ + 'generation.events': _EVENTS, + 'postprocessing.systematics': False, + 'gridpack.save_gridpack': True, + } + ) + + env = dict(os.environ) + env['LHAPDF_DATA_PATH'] = datadir + _run([sys.executable, pjoin(run_dir, 'bin', 'generate_events'), '-f'], + run_dir, pjoin(run_dir, 'survey.log'), env, + 'mg7 vegas survey run') + + gridpack_dir = pjoin(run_dir, 'Events', 'run_01', 'gridpack') + self.assertTrue(os.path.isdir(gridpack_dir), + 'gridpack was not produced under %s' % run_dir) + + hash_a1 = self._generate_and_hash_gridpack( + gridpack_dir, datadir, _SEED_A, 'a1') + hash_a2 = self._generate_and_hash_gridpack( + gridpack_dir, datadir, _SEED_A, 'a2') + self.assertEqual(hash_a1, hash_a2, + 'same seed produced different LHE content') + + hash_b = self._generate_and_hash_gridpack( + gridpack_dir, datadir, _SEED_B, 'b') + self.assertNotEqual(hash_a1, hash_b, + 'different seeds produced identical LHE content') + + +if __name__ == '__main__': + unittest.main()