diff --git a/.github/workflows/acceptancetest_mg7.yml b/.github/workflows/acceptancetest_mg7.yml index bc334e43b..3366206ec 100644 --- a/.github/workflows/acceptancetest_mg7.yml +++ b/.github/workflows/acceptancetest_mg7.yml @@ -201,6 +201,92 @@ jobs: export PATH="$HOME/.cache/HEPtools/bin:$PATH" ./tests/test_manager.py test_group_subprocess_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.reproducible = True (p p > t t~): madnis 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/madgraph/iolibs/template_files/mg7/gridpack.py b/madgraph/iolibs/template_files/mg7/gridpack.py index 62626f327..ccc71999b 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 @@ -69,6 +68,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", 0), + help="0 draws a fresh random seed; a positive integer makes the run " + "reproducible at any --cpu_thread_pool_size (for --output_format lhe " + "also set combine_thread_pool_size = 1)" + ) parser.add_argument("--device", type=str, nargs="*") parser.add_argument( "--cpu_thread_pool_size", type=int, default=run_args["cpu_thread_pool_size"] @@ -173,6 +178,7 @@ def main() -> None: channels=channel_generators, status_file=os.path.join(run_path, "info.json"), config=config, + seed=args.seed, ) # run generation diff --git a/madgraph/iolibs/template_files/mg7/madevent.py b/madgraph/iolibs/template_files/mg7/madevent.py index 6114a4871..002013bae 100644 --- a/madgraph/iolibs/template_files/mg7/madevent.py +++ b/madgraph/iolibs/template_files/mg7/madevent.py @@ -212,7 +212,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) @@ -426,6 +426,9 @@ def build_event_generator(self, phasespaces: list[PhaseSpace]) -> ms.EventGenera channels=channel_generators, status_file=os.path.join(self.run_path, "info.json"), config=self.event_generator_config, + # 0 conventionally means "no seed" in the run card; EventGenerator takes + # an actual seed or None (non-deterministic), not a 0 sentinel. + seed=self.run_card["run"]["seed"] or None, ) unused_globals = ( set(self.contexts[0].global_names()) - event_generator.used_globals() @@ -436,36 +439,42 @@ 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 == "multichannel": 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) elif phasespace_mode == "both": kept_count = self.run_card["phasespace"]["simplified_channel_count"] phasespaces_multi = [ subproc.build_multichannel_phasespace() for subproc in self.subprocesses ] - evgen_multi = self.survey_phasespaces(phasespaces_multi) + evgen_multi = self.survey_phasespaces(phasespaces_multi, 0) phasespaces_flat = [ subproc.build_flat_phasespace() @@ -473,7 +482,7 @@ def survey(self) -> None: None for subproc in self.subprocesses ] - #evgen_flat = self.survey_phasespaces(phasespaces_flat, "flat") + #evgen_flat = self.survey_phasespaces(phasespaces_flat, 1) channel_status = evgen_multi.channel_status() cross_sections = [] @@ -494,7 +503,7 @@ def survey(self) -> None: self.subprocesses, phasespaces_multi, phasespaces_flat, cross_sections ) ] - self.event_generator = self.survey_phasespaces(self.phasespaces) + self.event_generator = self.survey_phasespaces(self.phasespaces, 2) else: raise ValueError("Unknown phasespace mode") @@ -534,6 +543,7 @@ def train_madnis(self) -> None: config.buffer_unweighting_quantile = madnis_args["buffer_unweighting_quantile"] config.fixed_cwnet_fraction = madnis_args["fixed_cwnet_fraction"] config.softclip_threshold = madnis_args["softclip_threshold"] + config.reproducible = madnis_args["reproducible"] madnis_phasespaces = [] integrands = [] cwnets = [] @@ -559,6 +569,11 @@ def train_madnis(self) -> None: config=config, integrands=integrands, cwnets=cwnets, + # Reuses the run_card's own 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=run_args["seed"] or None, ) madnis_training.train() for phasespace, active_channels in zip( @@ -1301,6 +1316,9 @@ def simplify_phasespace( 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_card["run"]["seed"] or None channels = [] for channel_id, channel in enumerate(phasespace.channels): prefix = f"subproc{self.subproc_id}.channel{channel_id}" @@ -1318,7 +1336,7 @@ def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace: subnet_layers=madnis_args["discrete_layers"], subnet_activation=self.activation(madnis_args["discrete_activation"]), ) - discrete_before.initialize_globals(self.process.contexts[0]) + discrete_before.initialize_globals(self.process.contexts[0], seed) cond_dim += perm_count flow_dim = channel.phasespace_mapping.random_dim() @@ -1333,10 +1351,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 @@ -1351,7 +1370,7 @@ def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace: subnet_layers=madnis_args["discrete_layers"], subnet_activation=self.activation(madnis_args["discrete_activation"]), ) - discrete_after.initialize_globals(self.process.contexts[0]) + discrete_after.initialize_globals(self.process.contexts[0], seed) channels.append(Channel( phasespace_mapping = channel.phasespace_mapping, @@ -1420,7 +1439,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_card["run"]["seed"] or None + ) 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 28410c262..35b9bc173 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 +# 0 draws a fresh random seed each run; a positive integer makes the run +# reproducible: the same (seed, thread-pool sizes) reproduces bit-identically and +# distinct seeds are statistically independent. Generation is reproducible at any +# cpu_thread_pool_size; for output_format = lhe also set combine_thread_pool_size = 1 +# (compact_npy and lhe_npy are reproducible at any combine size). +seed = %(run.seed)s # options: cuda, hip, cpp, cppnone, cppsse4, cppavx2, cpp512y, cpp512z, cppauto devices = %(run.devices)s # options: @@ -136,3 +142,4 @@ batch_size_threshold = %(madnis.batch_size_threshold)s channel_grouping_mode = %(madnis.channel_grouping_mode)s # options: none, uniform, learned fixed_cwnet_fraction = %(madnis.fixed_cwnet_fraction)s softclip_threshold = %(madnis.softclip_threshold)s +reproducible = %(madnis.reproducible)s # make training fully reproducible for a given seed, independent of thread count (CPU only, some throughput/memory cost) diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py index 5d5b0c63d..474d6465e 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -6453,6 +6453,12 @@ def default_setup(self): # ----------------------------- [run] -------------------------- self.add_toml_param('run', 'run_name', "run", gridpack=True) + self.add_toml_param('run', 'seed', 0, gridpack=True, + comment="0 draws a fresh random seed each run; any positive integer " + "makes the run reproducible: the same (seed, thread-pool sizes) " + "reproduces bit-identically. Generation is reproducible at any " + "cpu_thread_pool_size; for output_format = lhe also set " + "combine_thread_pool_size = 1") 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, @@ -6588,6 +6594,7 @@ def default_setup(self): allowed=['none', 'uniform', 'learned']) self.add_toml_param('madnis', 'fixed_cwnet_fraction', 0.33) self.add_toml_param('madnis', 'softclip_threshold', 30.0) + self.add_toml_param('madnis', 'reproducible', False) # ----------------- dynamic (free-form) sections --------------- self.dynamic_sections['multiparticles'] = collections.OrderedDict([ 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 c7a6dd13d..38f3382c0 100644 --- a/madspace/CMakeLists.txt +++ b/madspace/CMakeLists.txt @@ -180,6 +180,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 @@ -232,6 +233,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 @@ -270,6 +272,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) diff --git a/madspace/include/madspace/compgraphs/instruction.hpp b/madspace/include/madspace/compgraphs/instruction.hpp index 184c83948..0b4125483 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; @@ -197,6 +200,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 { @@ -204,6 +208,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 { @@ -211,6 +216,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..7dbbeeada 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 { @@ -9,6 +13,8 @@ namespace madspace { class Runtime { public: virtual ~Runtime() = default; + // `seed`, when set, makes this call's random numbers a deterministic function of + // `seed` alone. Backends that can't support this should reject it, not ignore it. virtual TensorVec run(const TensorVec& inputs) = 0; virtual std::tuple> run_with_grad( const TensorVec& inputs, const std::vector& input_requires_grad @@ -19,6 +25,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 957037d15..8b6005d95 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" @@ -49,7 +50,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; } @@ -57,11 +60,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(); @@ -127,7 +143,17 @@ class ChannelEventGenerator { std::optional _histogram_function; RunningIntegral _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 13e438084..7eaa4c91d 100644 --- a/madspace/include/madspace/driver/context.hpp +++ b/madspace/include/madspace/driver/context.hpp @@ -48,35 +48,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; } @@ -170,6 +178,7 @@ 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++; } private: DevicePtr _device; @@ -177,6 +186,7 @@ class Context { std::unordered_map> _globals; std::vector> _matrix_elements; std::vector _param_card_paths; + std::size_t _seed_index = 0; }; using ContextPtr = std::shared_ptr; diff --git a/madspace/include/madspace/driver/event_generator.hpp b/madspace/include/madspace/driver/event_generator.hpp index 0e9428dff..623c046c7 100644 --- a/madspace/include/madspace/driver/event_generator.hpp +++ b/madspace/include/madspace/driver/event_generator.hpp @@ -2,7 +2,7 @@ #include #include -#include +#include #include #include @@ -31,14 +31,18 @@ class EventGenerator { const std::vector& contexts, const std::vector>& channels, const std::string& status_file = "", - const GeneratorConfig& config = default_config + const GeneratorConfig& config = default_config, + std::optional seed = std::nullopt ); - 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); void combine_to_lhe( - const std::string& file_name, LHECompleter& lhe_completer, + const std::string& file_name, + LHECompleter& lhe_completer, const LHEMeta& meta = {} ); GeneratorStatus status() const { return _status; } @@ -73,8 +77,42 @@ class EventGenerator { 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_job_count() from double-counting in-flight work. + std::vector _channel_batch_pending; + // generate_deterministic() only: per-channel commit cursor/buffer, analogous + // to _ready_gen/_commit_cursor but ordered per channel instead of globally. + std::vector> _channel_ready_gen; + std::vector _channel_commit_cursor; + std::vector _channel_cursor_set; + // generate_deterministic() only: same as above, for a job's unweight-stage + // completion (tracked separately since it's a distinct completion event). + std::vector> _channel_unweight_ready; + std::vector _channel_unweight_cursor; + // generate_deterministic() 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; nullopt means non-deterministic. + std::optional _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; + + // Deterministic-path state, set from _seed. Generate completions are + // committed in ascending job id, with _commit_cursor as the next id due. + bool _deterministic = false; + 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; @@ -84,8 +122,17 @@ class EventGenerator { std::string _status_file; std::unordered_map _timing_data; - bool start_jobs(); + void survey_deterministic(); + void generate_deterministic(); + void commit_generate_job(GeneratorBatchJob& job); + void commit_unweight_job(GeneratorBatchJob& job); + void finish_channel_job(const GeneratorBatchJob& job); + void register_dispatched_ids(std::size_t first_id, std::size_t end_id); + std::size_t next_batch_job_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); @@ -94,14 +141,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..17a6c058b 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,6 +64,9 @@ struct GeneratorConfig { int combine_thread_count = -1; double cut_efficiency_threshold = 0.7; std::size_t max_cut_repetitions = 100; + // Cap/floor on a steady-state batch's job count; see next_batch_job_count(). + double generation_batch_fraction = 0.5; + std::size_t min_batch_jobs = 1; }; struct GeneratorStatus { @@ -76,7 +80,7 @@ struct GeneratorStatus { 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; @@ -93,6 +97,7 @@ struct Histogram { struct GeneratorBatchJob { std::size_t channel_index; bool unweight; + // Nonzero: start_jobs() splits this into sub-jobs and dispatches them atomically. std::size_t vegas_batch_size; std::size_t split_job_count; Tensor weights; @@ -104,6 +109,16 @@ 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, false for a steady-state generation + // batch -- both use vegas_batch_size > 0 for atomic whole-batch dispatch. + 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 630175a22..92eaaf42b 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" @@ -86,7 +86,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 932760119..8e8179028 100644 --- a/madspace/include/madspace/driver/madnis_training.hpp +++ b/madspace/include/madspace/driver/madnis_training.hpp @@ -1,6 +1,8 @@ #pragma once #include +#include +#include #include "madspace/compgraphs.hpp" #include "madspace/driver/adam_optimizer.hpp" @@ -41,13 +43,22 @@ class MadnisTraining { double buffer_unweighting_quantile = 0.99; double fixed_cwnet_fraction = 0.33; double softclip_threshold = 0.0; + // Makes generator-job scheduling (CPU only) deterministic given the same + // seed, independent of thread count, at some cost to throughput/memory. + // GPU multi-channel batches are unaffected and stay non-deterministic. + bool reproducible = false; }; MadnisTraining( ContextPtr generator_context, 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 ); void train_step(std::size_t batch_index); std::vector active_channels() const; @@ -65,6 +76,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; @@ -77,6 +91,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; + // reproducible mode: samples staged here, flushed into buffer next round + std::vector pending_buffer_samples; }; inline static std::function _abort_check_function = [] {}; @@ -84,6 +105,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); @@ -92,6 +116,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, const std::vector& counts); @@ -115,6 +140,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; }; class MultiMadnisTraining { @@ -124,7 +158,8 @@ class MultiMadnisTraining { ContextPtr optimizer_context, const MadnisTraining::Config& config, const nested_vector2>& integrands, - const std::vector>& cwnets + const std::vector>& cwnets, + std::optional seed = std::nullopt ); void train(); nested_vector2 active_channels() const; diff --git a/madspace/include/madspace/driver/random.hpp b/madspace/include/madspace/driver/random.hpp new file mode 100644 index 000000000..2a4d6d39b --- /dev/null +++ b/madspace/include/madspace/driver/random.hpp @@ -0,0 +1,72 @@ +#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 + ); +}; + +class MixMaxRandom { +public: + MixMaxRandom() : MixMaxRandom(DerivedSeed()) {} + MixMaxRandom(DerivedSeed seed) : + _mixmax( + seed.seed_parts[0], + seed.seed_parts[1], + seed.seed_parts[2], + seed.seed_parts[3] + ) {} + explicit MixMaxRandom(std::uint64_t seed) : MixMaxRandom(DerivedSeed(seed)) {} + void set_seed(DerivedSeed seed) { + _mixmax.seed_uniquestream( + seed.seed_parts[0], + seed.seed_parts[1], + seed.seed_parts[2], + seed.seed_parts[3] + ); + } + 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: + mixmax_engine _mixmax; +}; + +} // namespace madspace diff --git a/madspace/include/madspace/mixmax/mixmax.hpp b/madspace/include/madspace/mixmax/mixmax.hpp new file mode 100644 index 000000000..5d3a67b31 --- /dev/null +++ b/madspace/include/madspace/mixmax/mixmax.hpp @@ -0,0 +1,375 @@ +/* + * 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 + +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(); + } + +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 + } + } + 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 a29bb9e5b..06184ce67 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 02ed5a5be..99c0eb136 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 @@ -613,12 +614,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(); } } ); @@ -636,12 +636,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); } } ); @@ -673,6 +674,7 @@ void op_unweight( indices_view_flat, uw_weights_view_flat, &runtime, + &device, batch_size, &indices, &uw_weights, @@ -682,12 +684,11 @@ 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) { + if (w != 0. && w > rand_gen.generate_double() * w_max) { indices_view[count] = i; uw_weights_view[count] = w > w_max ? w : w_max; ++count; @@ -883,13 +884,7 @@ 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->thread_pool(), []() { 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..fd18a0462 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 c86da5107..c4e140f97 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; @@ -29,6 +31,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( @@ -53,7 +93,7 @@ ChannelEventGenerator::ChannelEventGenerator( .count_after_cuts = 0, .count_after_cuts_opt = 0, .count_unweighted = 0., - .count_target = 1., + .count_target = 1, .optimized = false, .done = false }, @@ -177,7 +217,7 @@ ChannelEventGenerator::ChannelEventGenerator( .count_after_cuts = 0, .count_after_cuts_opt = 0, .count_unweighted = 0., - .count_target = 1., + .count_target = 1, .optimized = false, .done = false }, @@ -275,18 +315,17 @@ 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 (weight / _max_weight < rand_gen.generate_double()) { weight = 0; } else { weight = std::max(weight.value(), _max_weight); @@ -296,6 +335,8 @@ void ChannelEventGenerator::unweight_file(std::mt19937& rand_gen) { _weight_file.seek(i); _weight_file.write(buffer); } + _unweighted_count = _weight_file.event_count(); + _unweighted_accept_count = accept_count; _status.count_unweighted = accept_count; } @@ -387,14 +428,35 @@ 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 @@ -408,6 +470,8 @@ void ChannelEventGenerator::start_job( 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 = @@ -444,9 +508,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), @@ -456,7 +532,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), @@ -482,15 +558,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 ); @@ -504,6 +591,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); @@ -514,6 +608,7 @@ 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(); @@ -549,8 +644,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) { @@ -563,6 +659,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 a5e4760a5..f8625371f 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; @@ -16,7 +20,8 @@ EventGenerator::EventGenerator( const std::vector& contexts, const std::vector>& channels, const std::string& status_file, - const GeneratorConfig& config + const GeneratorConfig& config, + std::optional seed ) : _config(config), _status{ @@ -30,7 +35,7 @@ EventGenerator::EventGenerator( .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,9 +46,24 @@ EventGenerator::EventGenerator( _channel_optimizing(channels.size()), _channel_integral_fractions(channels.size(), 1.), _context_job_counts(contexts.size()), + _channel_batch_pending(channels.size()), + _channel_ready_gen(channels.size()), + _channel_commit_cursor(channels.size()), + _channel_cursor_set(channels.size()), + _channel_unweight_ready(channels.size()), + _channel_unweight_cursor(channels.size()), + _context_unweight_queue(contexts.size()), + _seed(seed), + _deterministic(seed.has_value()), _status_file(status_file) {} -void EventGenerator::survey() { +void EventGenerator::survey(std::size_t survey_pass) { + _survey_job = true; + _survey_pass = survey_pass; + if (_deterministic) { + survey_deterministic(); + return; + } reset_start_time(); bool done = false; std::size_t min_iters = _config.survey_min_iters; @@ -80,6 +100,7 @@ void EventGenerator::survey() { .channel_index = i, .unweight = iter >= min_iters - 1, .vegas_batch_size = vegas_batch_size, + .is_vegas_batch = true, }); if (iter >= min_iters) { total_event_count += vegas_batch_size; @@ -136,6 +157,11 @@ void EventGenerator::survey() { } void EventGenerator::generate() { + _survey_job = false; + if (_deterministic) { + generate_deterministic(); + return; + } reset_start_time(); print_gen_init(); @@ -168,6 +194,7 @@ void EventGenerator::generate() { .channel_index = channel_index, .unweight = true, .vegas_batch_size = channel->next_vegas_batch_size(), + .is_vegas_batch = true, }); } } else { @@ -227,14 +254,326 @@ void EventGenerator::generate() { print_gen_update(true); } -bool EventGenerator::start_jobs() { +// Commits a job's generate stage in ascending job id per channel. 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); + } + 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 { + _channel_batch_pending.at(job.channel_index) = false; + } +} + +void EventGenerator::register_dispatched_ids(std::size_t first_id, std::size_t end_id) { + for (std::size_t id = first_id; id < end_id; ++id) { + std::size_t channel_index = _running_jobs.at(id).channel_index; + if (!_channel_cursor_set.at(channel_index)) { + _channel_commit_cursor.at(channel_index) = id; + _channel_unweight_cursor.at(channel_index) = id; + _channel_cursor_set.at(channel_index) = true; + } + } +} + +// Deterministic generate path: jobs still run on the pool, but commit strictly +// in per-channel job-id order, so results are a pure function of the seed. +void EventGenerator::generate_deterministic() { + 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; + _channel_cursor_set.at(channel_index) = false; + _ready_jobs.push_back({ + .channel_index = channel_index, + .unweight = true, + .vegas_batch_size = 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_cursor_set.at(channel_index) = false; + std::size_t batch_jobs = next_batch_job_count(channel_index); + // Single ready_job with vegas_batch_size set so start_jobs() dispatches + // the whole batch atomically, giving the round a contiguous job-id + // range. + _ready_jobs.push_back({ + .channel_index = channel_index, + .unweight = true, + .vegas_batch_size = batch_jobs * _config.cpu_batch_size, + .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; + register_dispatched_ids(job_id_before, _job_id); + + // 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& cursor = _channel_commit_cursor.at(job.channel_index); + ready.insert(job_id); + while (ready.erase(cursor) > 0) { + auto& gen_job = _running_jobs.at(cursor); + commit_generate_job(gen_job); + if (!gen_job.unweight) { + _running_jobs.erase(cursor); + } + ++cursor; + } + } else { + auto& ready = _channel_unweight_ready.at(job.channel_index); + auto& cursor = _channel_unweight_cursor.at(job.channel_index); + ready.insert(job_id); + while (ready.erase(cursor) > 0) { + commit_unweight_job(_running_jobs.at(cursor)); + _running_jobs.erase(cursor); + ++cursor; + } + } + 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; + register_dispatched_ids(job_id_refill, _job_id); + } + + update_integral_fractions(); + + if (_status.done) { + unweight_all(); + } + if (_status.done) { + break; + } + } + print_gen_update(true); +} + +// Deterministic survey path, same commit-ordering approach as generate_deterministic(). +void EventGenerator::survey_deterministic() { + 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); + double target_precision = _config.survey_target_precision; + + std::size_t total_event_count = 0; + std::size_t done_event_count = 0; + for (auto& channel : _channels) { + std::size_t chan_batch_size = channel->batch_size(); + for (std::size_t iter = channel->status().iterations; iter < min_iters; + ++iter) { + total_event_count += chan_batch_size; + chan_batch_size = std::min(chan_batch_size * 2, _config.max_batch_size); + } + } + print_survey_init(); + + std::size_t iter = 0; + for (; !done && iter < max_iters; ++iter) { + 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) { + ++i; + continue; + } + std::size_t vegas_batch_size = channel->next_vegas_batch_size(); + // Let register_dispatched_ids() re-init the cursor for this new job. + _channel_cursor_set.at(i) = false; + _ready_jobs.push_back({ + .channel_index = i, + .unweight = iter >= min_iters - 1, + .vegas_batch_size = vegas_batch_size, + .is_vegas_batch = true, + }); + if (iter >= min_iters) { + total_event_count += vegas_batch_size; + } + ++i; + } + 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; + register_dispatched_ids(job_id_before, _job_id); + done = true; + while (in_flight > 0) { + std::size_t job_id = _result_queue.wait(); + _abort_check_function(); + --in_flight; + auto& job = _running_jobs.at(job_id); + if (job.unweighted_events.size() == 0) { + // 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); + 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); + } else { + done = false; + if (channel_job_count == 0) { + channel->optimize_vegas(gen_job); + done_event_count += gen_job.vegas_batch_size; + } + _running_jobs.erase(_commit_cursor); + } + ++_commit_cursor; + } + } else { + // Unweight-stage completion, ordered per channel only. + auto& ready = _channel_unweight_ready.at(job.channel_index); + auto& cursor = _channel_unweight_cursor.at(job.channel_index); + ready.insert(job_id); + while (ready.erase(cursor) > 0) { + auto& uw_job = _running_jobs.at(cursor); + 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->cross_section().rel_error() < target_precision) { + done = false; + } + if (channel_job_count == 0) { + channel->optimize_vegas(uw_job); + done_event_count += uw_job.vegas_batch_size; + } + _running_jobs.erase(cursor); + ++cursor; + } + } + 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; + register_dispatched_ids(job_id_refill, _job_id); + print_survey_update(false, done_event_count, total_event_count, iter); + } + } + print_survey_update(true, done_event_count, total_event_count, iter - 1); +} + +// Sizes the next batch from committed data only, so the estimate never depends +// on jobs still in flight. +std::size_t EventGenerator::next_batch_job_count(std::size_t channel_index) const { + auto& status = _channels.at(channel_index)->status(); + // No data yet: assume a pessimistic efficiency of 1.0. + double efficiency = status.count_opt > 0 + ? status.count_unweighted / static_cast(status.count_opt) + : 1.; + // Floor at 1 event so a near-zero observed efficiency can't blow this up. + double per_job_estimate = + std::max(efficiency * static_cast(_config.cpu_batch_size), 1.); + double remaining = std::max( + static_cast(status.count_target) - status.count_unweighted, 0. + ); + double estimated = _config.generation_batch_fraction * remaining / per_job_estimate; + return std::max( + _config.min_batch_jobs, static_cast(std::ceil(estimated)) + ); +} + +std::size_t EventGenerator::start_jobs() { std::size_t ready_index = 0, 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; + + // 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 + ); + for (; job_count < target_count && ready_index < _ready_jobs.size(); ++ready_index) { auto ready_job = _ready_jobs.at(ready_index); @@ -247,22 +586,25 @@ bool EventGenerator::start_jobs() { 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); + _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; - } ++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.; std::size_t total_count = 0, total_count_opt = 0; std::size_t total_count_after_cuts = 0, total_count_after_cuts_opt = 0; @@ -293,19 +635,48 @@ void EventGenerator::update_integral() { _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->cross_section().mean() / _status.mean; + } + + // 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; @@ -319,6 +690,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( @@ -337,7 +709,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; } @@ -355,8 +727,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::optional 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( @@ -388,7 +761,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; } @@ -418,11 +791,9 @@ void EventGenerator::combine_to_lhe( const std::string& file_name, LHECompleter& lhe_completer, const LHEMeta& meta ) { reset_start_time(); + std::optional 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; @@ -444,45 +815,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; } } @@ -506,15 +895,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; @@ -572,11 +963,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); @@ -596,7 +988,8 @@ EventGenerator::init_combine() { 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); @@ -607,11 +1000,8 @@ void EventGenerator::read_and_combine( bool has_partial = _channels.at(0)->event_layout_extra_flags() & EventRecord::f_partial_weights; - 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(), @@ -688,7 +1078,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(); @@ -739,14 +1129,14 @@ void EventGenerator::write_status(const std::string& status, bool force_write) { std::ofstream f(status_tmp_file); nlohmann::json j{ {"status", status}, + {"seed", _seed.value_or(0)}, {"process", _status}, {"channels", channel_status()}, {"run_times", _timing_data}, {"histograms", histograms()}, }; f << j.dump(); - // rename atomically deletes the old file and replaces it with the new one - // such that the status file exists at all times + // Atomic rename keeps the status file present at all times. std::filesystem::rename(status_tmp_file, _status_file); } diff --git a/madspace/src/driver/lhe_output.cpp b/madspace/src/driver/lhe_output.cpp index 3f96379b7..f99ea8f73 100644 --- a/madspace/src/driver/lhe_output.cpp +++ b/madspace/src/driver/lhe_output.cpp @@ -394,7 +394,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) { @@ -423,8 +423,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 c94d84fb9..665080be8 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; @@ -54,7 +59,7 @@ void MadnisTraining::train_step(std::size_t batch_index) { _config.buffer_capacity > 0 && batch_index % (_config.buffered_steps + 1) != 0; TensorVec training_batch; 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); @@ -71,22 +76,26 @@ void MadnisTraining::train_step(std::size_t batch_index) { update_history(results, channel_sizes); 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(); } 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 + 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); + } } std::vector MadnisTraining::active_channels() const { @@ -262,9 +271,18 @@ 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; + if (_config.reproducible) { + // 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; @@ -280,14 +298,29 @@ 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; + // reproducible mode: dispatch exactly what this round needs, uncapped by + // thread count, so round contents don't depend on thread count + std::size_t available_jobs; + if (_config.reproducible) { + available_jobs = 0; + for (auto count : missing_batch_counts) { + available_jobs += count; + } + for (auto count : target_batch_counts) { + available_jobs += count; + } + } else { + available_jobs = _generator_context->thread_pool().thread_count(); + if (is_gpu) { + available_jobs *= gpu_subbatches; + } + } + std::vector channel_sizes; if (is_gpu) { - available_jobs *= gpu_subbatches; channel_sizes.resize(chan_count, 0); } @@ -349,6 +382,32 @@ void MadnisTraining::start_generator_jobs(const std::vector& counts } } +// Reproducible mode (CPU only) additionally requires an online attempt and the +// online cache to be exhausted or about to be, both deterministic conditions. +void MadnisTraining::maybe_start_generator_jobs( + const std::vector& counts, bool is_online_attempt +) { + if (_running_jobs.size() > 0) { + return; + } + if (_config.reproducible) { + 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()); @@ -358,20 +417,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); @@ -382,15 +463,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(); @@ -502,54 +602,94 @@ 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; - 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; + _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; + channel.sample_batches.push_back(std::move(committed_job.samples)); + if (committed_job.unweighted_samples.size > 0) { + if (_config.reproducible) { + // flushed into buffer at the start of the next round instead + channel.pending_buffer_samples.push_back( + std::move(committed_job.unweighted_samples) + ); + } else { + buffer_store(channel, committed_job.unweighted_samples); } - channel.sample_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) + } + } + } + + // 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; + 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) - ); - } + if (_config.reproducible) { + channel.pending_buffer_samples.push_back( + std::move(chan_unweighted_samples) + ); + } else { buffer_store(channel, chan_unweighted_samples); - unw_offset += unw_chan_size; } - batch.size = chan_size; - offset += chan_size; - ++chan_index; + unw_offset += unw_chan_size; } + batch.size = chan_size; + offset += chan_size; + ++chan_index; } } } @@ -642,7 +782,8 @@ void MadnisTraining::drop_channels() { } 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); }); @@ -668,8 +809,7 @@ void MadnisTraining::drop_channels() { 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) { @@ -719,14 +859,24 @@ MultiMadnisTraining::MultiMadnisTraining( ContextPtr optimizer_context, const MadnisTraining::Config& config, const nested_vector2>& integrands, - const std::vector>& cwnets + const std::vector>& cwnets, + std::optional seed ) : _config(config) { _subprocesses.reserve(integrands.size()); + // each subprocess gets its own channel_index slice so their streams don't collide + std::size_t channel_index_offset = 0; for (auto [integ, cwnet] : zip(integrands, cwnets)) { _subprocesses.emplace_back( - generator_context, optimizer_context, config, integ, cwnet + generator_context, + optimizer_context, + config, + integ, + cwnet, + seed, + channel_index_offset ); + channel_index_offset += integ.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/runtime.cu b/madspace/src/gpu/runtime.cu index 87ff0b0be..0e01bd1e2 100644 --- a/madspace/src/gpu/runtime.cu +++ b/madspace/src/gpu/runtime.cu @@ -1610,11 +1610,14 @@ GpuRuntime::GpuRuntime(const Function& function_arg, ContextPtr context) : ); } -TensorVec GpuRuntime::run(const TensorVec& inputs) { +TensorVec GpuRuntime::run(const TensorVec& inputs, std::optional seed) { auto& gpu_device = *static_cast(_context->device()); auto& streams = _streams.get(); auto& events = _events.get(); gpu_device.activate(); + if (seed) { + check_error(gpurandSetPseudoRandomGeneratorSeed(gpurand_generator(), *seed)); + } auto locals = _locals_init; std::copy(inputs.begin(), inputs.end(), locals.begin()); gpuStream_t main_stream = streams.at(0); @@ -1650,12 +1653,17 @@ TensorVec GpuRuntime::run(const TensorVec& inputs) { } std::tuple> GpuRuntime::run_with_grad( - const TensorVec& inputs, const std::vector& input_requires_grad + const TensorVec& inputs, + const std::vector& input_requires_grad, + std::optional seed ) { auto& gpu_device = *static_cast(_context->device()); auto& streams = _streams.get(); auto& events = _events.get(); gpu_device.activate(); + if (seed) { + check_error(gpurandSetPseudoRandomGeneratorSeed(gpurand_generator(), *seed)); + } auto locals = _locals_init; auto requires_grad = _requires_grad_init; std::vector store_local(locals.size()); diff --git a/madspace/src/gpu/runtime.hpp b/madspace/src/gpu/runtime.hpp index 30930de26..567b28b1b 100644 --- a/madspace/src/gpu/runtime.hpp +++ b/madspace/src/gpu/runtime.hpp @@ -6,6 +6,7 @@ #include "madspace/driver/tensor.hpp" #include +#include namespace madspace { namespace gpu { @@ -31,9 +32,13 @@ class GpuRuntime : public Runtime { }; GpuRuntime(const Function& function, ContextPtr context); - TensorVec run(const TensorVec& inputs) override; + TensorVec + run(const TensorVec& inputs, + std::optional seed = std::nullopt) override; std::tuple> run_with_grad( - const TensorVec& inputs, const std::vector& input_requires_grad + const TensorVec& inputs, + const std::vector& input_requires_grad, + std::optional seed = std::nullopt ) override; std::pair run_backward( const TensorVec& output_grads, diff --git a/madspace/src/phasespace/channel_weight_network.cpp b/madspace/src/phasespace/channel_weight_network.cpp index d84a7a389..4a71ad747 100644 --- a/madspace/src/phasespace/channel_weight_network.cpp +++ b/madspace/src/phasespace/channel_weight_network.cpp @@ -67,8 +67,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/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 72c029e7c..0beb10945 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") @@ -1363,7 +1378,8 @@ PYBIND11_MODULE(_madspace_py, m) { ) .def_readwrite( "softclip_threshold", &MadnisTraining::Config::softclip_threshold - ); + ) + .def_readwrite("reproducible", &MadnisTraining::Config::reproducible); py::classh(m, "MadnisTraining") .def( @@ -1372,12 +1388,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) @@ -1390,12 +1408,14 @@ PYBIND11_MODULE(_madspace_py, m) { ContextPtr, const MadnisTraining::Config&, const nested_vector2>&, - const std::vector>&>(), + const std::vector>&, + std::optional>(), py::arg("generator_context"), py::arg("optimizer_context"), py::arg("config"), py::arg("integrands"), - py::arg("cwnets") + py::arg("cwnets"), + py::arg("seed") = std::nullopt ) .def("train", &MultiMadnisTraining::train) .def("active_channels", &MultiMadnisTraining::active_channels); @@ -1429,7 +1449,11 @@ 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( + "generation_batch_fraction", &GeneratorConfig::generation_batch_fraction + ) + .def_readwrite("min_batch_jobs", &GeneratorConfig::min_batch_jobs); py::classh(m, "GeneratorStatus") .def(py::init<>()) @@ -1611,9 +1635,9 @@ PYBIND11_MODULE(_madspace_py, m) { .def_readwrite("pdg_color_types", &LHECompleter::SubprocArgs::pdg_color_types) .def_readwrite("helicities", &LHECompleter::SubprocArgs::helicities) .def_readwrite("pdg_ids", &LHECompleter::SubprocArgs::pdg_ids); - 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>(), @@ -1712,7 +1736,8 @@ PYBIND11_MODULE(_madspace_py, m) { const std::vector&, const std::vector>&, const std::string&, - const GeneratorConfig&>(), + const GeneratorConfig&, + std::optional>(), py::arg("contexts"), py::arg("channels"), py::arg("status_file") = "", @@ -1720,9 +1745,10 @@ PYBIND11_MODULE(_madspace_py, m) { "config", EventGenerator::default_config, "EventGenerator.default_config" - ) + ), + py::arg("seed") = std::nullopt ) - .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_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..6f9edd8e3 --- /dev/null +++ b/tests/acceptance_tests/test_mg7_reproducibility.py @@ -0,0 +1,379 @@ +################################################################################ +# +# 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), with + ``madnis.reproducible = True``: training 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 mode. + * ``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): + """With madnis.reproducible = True (CPU codepath), madnis training + itself -- 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 scheduling fix 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, + 'madnis.reproducible': True, + } + + # 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.reproducible 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.reproducible 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()