From bf3e1acc1f5959d416801a8c3a459709c1eea915 Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Thu, 23 Jul 2026 07:48:50 +0200 Subject: [PATCH 01/19] Add run_card seed for reproducible mg7 CPU event generation Introduce a `seed` field in the mg7 run_card that makes CPU event generation reproducible: the same (seed, thread-pool sizes) reproduces bit-identically and distinct seeds are statistically independent. Seed 0 (the default) keeps the previous non-deterministic behaviour. madspace previously seeded every RNG from std::random_device with no seed input and drew randomness from thread-local generators dispatched to the pool in non-deterministic order, so runs were never reproducible. Plumbing: - run_card `[run] seed` + RunCardMG7 default; per-context seed derivation (splitmix64) in madevent.py and gridpack.py (--seed); Context(seed=) exposed through the Python bindings. - seeded_rng()/mix_seed() helpers in util.hpp; the driver-level random_device sites reseeded from the context seed. Determinism (gated on seed != 0, so the default path is unchanged): - Per-job RNG keyed by logical job identity (channel, sequence, phase) via Runtime::begin/end_job_rng, so a job's randomness no longer depends on which worker thread runs it. - Deterministic commit ordering: survey/generate harvest jobs in job-id order with unweighting folded in synchronously, making cross-section sums, max-weight snapshots, event order and stopping deterministic. Validated end-to-end through the mg7 driver with a dummy matrix element: same seed gives byte-identical events.npy at 4 and 8 threads, a different seed differs, and seed 0 differs run-to-run. Adds a portable unit test (madspace/tests/test_determinism.py) covering the seeding plumbing. For output_format = lhe, bit-identical output also needs combine_thread_pool_size = 1 (the parallel LHE combine is not yet deterministic); compact_npy and lhe_npy are reproducible at any size. Co-Authored-By: Claude Opus 4.8 --- .../iolibs/template_files/mg7/gridpack.py | 23 ++ .../iolibs/template_files/mg7/madevent.py | 30 ++- .../iolibs/template_files/mg7/run_card.toml | 6 + madgraph/various/banner.py | 6 + madspace/include/madspace/driver/backend.hpp | 11 + .../madspace/driver/channel_generator.hpp | 9 + madspace/include/madspace/driver/context.hpp | 16 +- .../madspace/driver/event_generator.hpp | 20 +- .../madspace/driver/generator_data.hpp | 4 + madspace/include/madspace/util.hpp | 50 ++++ madspace/src/cpu/runtime.cpp | 33 ++- madspace/src/cpu/runtime.hpp | 6 +- madspace/src/driver/channel_generator.cpp | 49 ++++ madspace/src/driver/event_generator.cpp | 248 +++++++++++++++++- madspace/src/python/madspace.cpp | 12 +- madspace/tests/test_determinism.py | 84 ++++++ 16 files changed, 580 insertions(+), 27 deletions(-) create mode 100644 madspace/tests/test_determinism.py diff --git a/madgraph/iolibs/template_files/mg7/gridpack.py b/madgraph/iolibs/template_files/mg7/gridpack.py index 1df578c2c..e5d09ac1b 100644 --- a/madgraph/iolibs/template_files/mg7/gridpack.py +++ b/madgraph/iolibs/template_files/mg7/gridpack.py @@ -31,6 +31,24 @@ import tomllib import argparse + +def derive_seed(base_seed: int, index: int) -> int: + """Derive an independent per-context stream seed from the base seed. + + Returns 0 (seed non-deterministically) when ``base_seed`` is 0. Otherwise the + base seed and context ``index`` are mixed with splitmix64 so distinct contexts + get well-separated streams and different base seeds never collide. + """ + if base_seed == 0: + return 0 + mask = (1 << 64) - 1 + x = (base_seed + index * 0x9E3779B97F4A7C15) & mask + x = ((x ^ (x >> 30)) * 0xBF58476D1CE4E5B9) & mask + x = ((x ^ (x >> 27)) * 0x94D049BB133111EB) & mask + x = x ^ (x >> 31) + return x or 1 + + def main() -> None: # load run card and metadata. Use the RunCardMG7 representation when the # madgraph package is importable; gridpacks are meant to be portable, so @@ -59,6 +77,11 @@ 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 (bit-identical output requires --cpu_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"] diff --git a/madgraph/iolibs/template_files/mg7/madevent.py b/madgraph/iolibs/template_files/mg7/madevent.py index 48078090d..8c3d0a490 100644 --- a/madgraph/iolibs/template_files/mg7/madevent.py +++ b/madgraph/iolibs/template_files/mg7/madevent.py @@ -84,6 +84,24 @@ logger = logging.getLogger("madevent7") +def derive_seed(base_seed: int, index: int) -> int: + """Derive an independent per-context stream seed from the run_card seed. + + Returns 0 (i.e. seed non-deterministically) when ``base_seed`` is 0, which is + the run_card default. Otherwise the base seed and context ``index`` are mixed + with splitmix64 so that distinct contexts get well-separated streams and two + different base seeds never collide onto the same context seed. + """ + if base_seed == 0: + return 0 + mask = (1 << 64) - 1 + x = (base_seed + index * 0x9E3779B97F4A7C15) & mask + x = ((x ^ (x >> 30)) * 0xBF58476D1CE4E5B9) & mask + x = ((x ^ (x >> 27)) * 0x94D049BB133111EB) & mask + x = x ^ (x >> 31) + return x or 1 + + def get_start_time(): return time.time(), time.process_time() @@ -190,6 +208,7 @@ def init_event_dir(self) -> None: def init_context(self) -> None: device_names = self.run_card["run"]["devices"] + base_seed = self.run_card["run"]["seed"] self.contexts = [] self.device_types = [] self.devices = [] @@ -213,7 +232,11 @@ def init_context(self) -> None: pool_size = self.run_card["run"]["cpu_thread_pool_size"] self.devices.append(device) self.pool_sizes.append(pool_size) - self.contexts.append(ms.Context(device=device, thread_count=pool_size)) + self.contexts.append(ms.Context( + device=device, + thread_count=pool_size, + seed=derive_seed(base_seed, i), + )) def parse_observable(self, name: str, order_observable: str) -> dict: parts = name.split("-") @@ -531,7 +554,10 @@ def train_madnis(self) -> None: gen_context = self.contexts[0] opt_context = ms.Context( - device=self.devices[0], thread_count=self.pool_sizes[0] + device=self.devices[0], + thread_count=self.pool_sizes[0], + # distinct stream from the generation contexts (offset the index) + seed=derive_seed(self.run_card["run"]["seed"], 1000), ) opt_context.copy_globals_from(gen_context) diff --git a/madgraph/iolibs/template_files/mg7/run_card.toml b/madgraph/iolibs/template_files/mg7/run_card.toml index c48d95a2c..6bffcd874 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: diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py index 2120ed298..9d77d6ea9 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -6436,6 +6436,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, diff --git a/madspace/include/madspace/driver/backend.hpp b/madspace/include/madspace/driver/backend.hpp index 792d5828b..16cf6bea7 100644 --- a/madspace/include/madspace/driver/backend.hpp +++ b/madspace/include/madspace/driver/backend.hpp @@ -1,5 +1,7 @@ #pragma once +#include + #include "madspace/compgraphs.hpp" #include "madspace/driver/context.hpp" #include "madspace/driver/tensor.hpp" @@ -19,6 +21,15 @@ class Runtime { const std::vector& eval_grad, bool return_contiguous_grads = false ) = 0; + + // Deterministic per-job RNG. Between begin_job_rng(seed) and end_job_rng(), all + // randomness consumed by run() on the CALLING thread is drawn from a generator + // seeded by `seed`, so a job's random numbers depend only on `seed` and not on + // which pool thread executes it. Backends that cannot support this (e.g. GPU) + // keep the default no-op and fall back to their usual per-thread streams. + virtual void begin_job_rng(std::uint64_t seed) {} + virtual void end_job_rng() {} + 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..f1a66bc9b 100644 --- a/madspace/include/madspace/driver/channel_generator.hpp +++ b/madspace/include/madspace/driver/channel_generator.hpp @@ -63,6 +63,11 @@ class ChannelEventGenerator { double channel_weight_sum(std::size_t event_count); void start_job(GeneratorBatchJob& job, ResultQueue& result_queue); void start_unweight_job(GeneratorBatchJob& job, ResultQueue& result_queue); + // Synchronous unweighting used by the deterministic generation path: runs the + // unweighter on the calling (harvest) thread against the current max weight so + // that it happens in a fixed, job-id-ordered position instead of as a separate + // asynchronously-scheduled job. + void unweight_job_inline(GeneratorBatchJob& job); std::size_t next_vegas_batch_size(); void clear_events(); void update_max_weight(Tensor weights); @@ -127,6 +132,10 @@ class ChannelEventGenerator { std::optional _histogram_function; RunningIntegral _cross_section; double _max_weight = 0.; + // Monotonic per-channel counter assigned to each scheduled job (main thread) to + // key its deterministic random stream; never reset, so re-used jobs keep unique + // streams across survey/generate. + std::size_t _rng_seq = 0; std::size_t _unweighted_count = 0; std::size_t _iters_without_improvement = 0; double _best_rsd = std::numeric_limits::max(); diff --git a/madspace/include/madspace/driver/context.hpp b/madspace/include/madspace/driver/context.hpp index e4e386d7c..1e369d713 100644 --- a/madspace/include/madspace/driver/context.hpp +++ b/madspace/include/madspace/driver/context.hpp @@ -1,5 +1,6 @@ #pragma once +#include #include #include @@ -105,11 +106,14 @@ class Context { * Contains global variables and matrix elements */ public: - Context(int thread_count = -1) : + Context(int thread_count = -1, std::uint64_t seed = 0) : _device(cpu_device()), - _thread_pool(std::make_unique(thread_count)) {} - Context(DevicePtr device, int thread_count = -1) : - _device(device), _thread_pool(std::make_unique(thread_count)) {} + _thread_pool(std::make_unique(thread_count)), + _seed(seed) {} + Context(DevicePtr device, int thread_count = -1, std::uint64_t seed = 0) : + _device(device), + _thread_pool(std::make_unique(thread_count)), + _seed(seed) {} Context(Context&&) = default; Context& operator=(Context&&) = default; Context(const Context&) = delete; @@ -135,10 +139,14 @@ class Context { void load_globals(const std::string& dir); DevicePtr device() { return _device; } ThreadPool& thread_pool() { return *_thread_pool; } + // Base seed for this context's random-number streams. 0 means "seed + // non-deterministically" (the historical default); see seeded_rng(). + std::uint64_t seed() const { return _seed; } private: DevicePtr _device; std::unique_ptr _thread_pool; + std::uint64_t _seed = 0; std::unordered_map> _globals; std::vector> _matrix_elements; std::vector _param_card_paths; diff --git a/madspace/include/madspace/driver/event_generator.hpp b/madspace/include/madspace/driver/event_generator.hpp index c38262dbd..9bff4f28b 100644 --- a/madspace/include/madspace/driver/event_generator.hpp +++ b/madspace/include/madspace/driver/event_generator.hpp @@ -3,6 +3,7 @@ #include #include #include +#include #include #include @@ -71,6 +72,13 @@ class EventGenerator { std::vector _context_job_counts; ResultQueue _result_queue; + // Deterministic-path state. _deterministic is set from the context seed. Generate + // completions are buffered in _ready_gen and committed in ascending job id, with + // _commit_cursor pointing at the next job id to commit. + 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; @@ -80,6 +88,15 @@ class EventGenerator { std::string _status_file; std::unordered_map _timing_data; + // Deterministic generation path (enabled when the context seed is non-zero): + // expensive matrix-element jobs still run on the pool, but their results are + // committed strictly in job-id order with unweighting folded in synchronously, + // so the whole computation is a deterministic function of the seed at a fixed + // thread count. + void survey_deterministic(); + void generate_deterministic(); + void commit_generate_deterministic(GeneratorBatchJob& job); + bool start_jobs(); void update_integral(); void update_counts(); @@ -90,7 +107,8 @@ class EventGenerator { void read_and_combine( std::vector& channel_data, EventBuffer& buffer, - double norm_factor + double norm_factor, + std::mt19937& rand_gen ); void fill_lhe_event( LHECompleter& lhe_completer, diff --git a/madspace/include/madspace/driver/generator_data.hpp b/madspace/include/madspace/driver/generator_data.hpp index bedf6bb2c..f01104136 100644 --- a/madspace/include/madspace/driver/generator_data.hpp +++ b/madspace/include/madspace/driver/generator_data.hpp @@ -104,6 +104,10 @@ struct GeneratorBatchJob { std::size_t context_index; std::size_t job_id; double max_weight; + // Base seed for this job's deterministic random stream, derived at scheduling + // time from (context seed, channel, per-channel job sequence). 0 means "use the + // non-deterministic pool generator" (context seed 0). See ChannelEventGenerator. + std::uint64_t rng_seed = 0; }; void to_json(nlohmann::json& j, const GeneratorStatus& status); diff --git a/madspace/include/madspace/util.hpp b/madspace/include/madspace/util.hpp index 7b1e43b38..34b8614d7 100644 --- a/madspace/include/madspace/util.hpp +++ b/madspace/include/madspace/util.hpp @@ -1,7 +1,10 @@ #pragma once +#include #include #include +#include +#include #include #include #include @@ -9,6 +12,53 @@ namespace madspace { +// Deterministic RNG factory used everywhere randomness is consumed on the CPU. +// +// A ``seed`` of 0 preserves the historical behaviour: the generator is seeded +// non-deterministically from ``std::random_device``. Any non-zero ``seed`` makes +// the returned generator fully reproducible. The 64-bit seed together with the +// (small, low-entropy) integer ``salts`` is mixed through ``std::seed_seq`` so +// that distinct ``(seed, salts...)`` tuples yield well-separated, statistically +// independent streams -- salts are used to key a stream by its logical role +// (e.g. thread index, channel selection vs. unweighting vs. LHE completion). +// Fold a base seed together with integer salts into a single 64-bit value using +// splitmix64. Used to derive a stable per-job seed from (context seed, channel, +// job sequence, phase) so that a job's random stream depends only on its logical +// identity, never on which thread runs it. mix_seed(0, ...) stays 0 so callers can +// keep using 0 as the "non-deterministic" sentinel. +inline std::uint64_t +mix_seed(std::uint64_t seed, std::initializer_list salts = {}) { + if (seed == 0) { + return 0; + } + auto splitmix = [](std::uint64_t x) { + x += 0x9E3779B97F4A7C15ull; + x = (x ^ (x >> 30)) * 0xBF58476D1CE4E5B9ull; + x = (x ^ (x >> 27)) * 0x94D049BB133111EBull; + return x ^ (x >> 31); + }; + std::uint64_t h = splitmix(seed); + for (auto s : salts) { + h = splitmix(h ^ splitmix(s)); + } + return h ? h : 1; +} + +inline std::mt19937 +seeded_rng(std::uint64_t seed, std::initializer_list salts = {}) { + if (seed == 0) { + std::random_device rand_device; + return std::mt19937(rand_device()); + } + std::vector data; + data.reserve(2 + salts.size()); + data.push_back(static_cast(seed & 0xFFFFFFFFu)); + data.push_back(static_cast(seed >> 32)); + data.insert(data.end(), salts.begin(), salts.end()); + std::seed_seq seq(data.begin(), data.end()); + return std::mt19937(seq); +} + template struct Overloaded : Ts... { using Ts::operator()...; diff --git a/madspace/src/cpu/runtime.cpp b/madspace/src/cpu/runtime.cpp index 4c7f05474..eadf4b474 100644 --- a/madspace/src/cpu/runtime.cpp +++ b/madspace/src/cpu/runtime.cpp @@ -1,6 +1,8 @@ #include "runtime.hpp" #include +#include +#include #include #include #include @@ -47,6 +49,15 @@ using namespace madspace::kernels; namespace { +// Per-thread deterministic job generator. When a job installs one via +// begin_job_rng(), every op_random/op_random_int/op_unweight on this thread draws +// from it instead of the pool generator, so the job's random numbers depend only +// on its seed and not on which thread happened to run it. Integrand runtimes are +// single-threaded (concurrent == false), so the whole job runs on one thread and +// a single continuous stream spans all of its random draws. +thread_local std::mt19937 t_job_rng; +thread_local bool t_job_rng_active = false; + template void op_matrix_element( const CpuRuntime::Instruction& instruction, TensorVec& locals, const D& device @@ -880,14 +891,30 @@ void op_discrete_histogram( } // namespace +std::mt19937& CpuRuntime::rand_gen() { + return t_job_rng_active ? t_job_rng : _rand_gens.get(); +} + +void CpuRuntime::begin_job_rng(std::uint64_t seed) { + t_job_rng = seeded_rng(seed); + t_job_rng_active = true; +} + +void CpuRuntime::end_job_rng() { + t_job_rng_active = false; +} + 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()); + // Fallback per-thread streams used when no per-job generator is installed + // (context seed 0, i.e. non-deterministic mode). Keyed by a running thread + // index so the streams are independent. + [seed = context->seed(), + next_index = std::make_shared>(0)]() { + return seeded_rng(seed, {next_index->fetch_add(1)}); } ), _concurrent(concurrent) { diff --git a/madspace/src/cpu/runtime.hpp b/madspace/src/cpu/runtime.hpp index 9b70fdbfc..7ebd40fa2 100644 --- a/madspace/src/cpu/runtime.hpp +++ b/madspace/src/cpu/runtime.hpp @@ -49,7 +49,11 @@ class CpuRuntime : public Runtime { ) override; Context& context() { return *_context; } - std::mt19937& rand_gen() { return _rand_gens.get(); } + // Returns the active per-job generator when one is installed on this thread + // (see begin_job_rng), otherwise the thread's pool generator. + std::mt19937& rand_gen(); + void begin_job_rng(std::uint64_t seed) override; + void end_job_rng() override; private: TensorVec run_single(const TensorVec& inputs) const; diff --git a/madspace/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index c86da5107..2e65914de 100644 --- a/madspace/src/driver/channel_generator.cpp +++ b/madspace/src/driver/channel_generator.cpp @@ -1,10 +1,16 @@ #include "madspace/driver/channel_generator.hpp" +#include "madspace/util.hpp" using namespace madspace; using json = nlohmann::json; namespace { +// Phase salts keeping the phase-space/matrix-element stream of a job independent +// from its unweighting stream while both derive from the same per-job base seed. +constexpr std::uint64_t kPhaseGenerate = 1; +constexpr std::uint64_t kPhaseUnweight = 2; + int event_extra_flags(const std::unordered_map& index_map) { int flags = 0; if (index_map.contains("fact_scale1")) { @@ -390,11 +396,22 @@ double ChannelEventGenerator::channel_weight_sum(std::size_t event_count) { void ChannelEventGenerator::start_job( GeneratorBatchJob& job, ResultQueue& result_queue ) { + // Assign the job's deterministic base seed on the (single) scheduling thread so + // it depends only on the job's logical identity, not on which worker runs it. + std::uint64_t context_seed = _contexts.at(job.context_index)->seed(); + job.rng_seed = context_seed == 0 + ? 0 + : mix_seed(context_seed, {job.channel_index, _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 != 0) { + runtimes.integrand_channel->begin_job_rng( + mix_seed(job.rng_seed, {kPhaseGenerate}) + ); + } std::size_t max_batch_size = context->device()->device_type() == DeviceType::cpu ? _config.cpu_batch_size @@ -477,6 +494,9 @@ void ChannelEventGenerator::start_job( } } } + if (job.rng_seed != 0) { + runtimes.integrand_channel->end_job_rng(); + } result_queue.push(job.job_id); return std::nullopt; }); @@ -491,6 +511,11 @@ void ChannelEventGenerator::start_unweight_job( .submit([this, &job, &result_queue]() { auto& runtimes = _runtimes.at(job.context_index); auto& context = _contexts.at(job.context_index); + if (job.rng_seed != 0) { + runtimes.unweighter->begin_job_rng( + mix_seed(job.rng_seed, {kPhaseUnweight}) + ); + } std::vector unweighter_args( job.events.begin(), job.events.begin() + _field_indices.random ); @@ -499,11 +524,35 @@ void ChannelEventGenerator::start_unweight_job( for (auto& item : unw_events) { job.unweighted_events.push_back(item.cpu()); } + if (job.rng_seed != 0) { + runtimes.unweighter->end_job_rng(); + } result_queue.push(job.job_id); return std::nullopt; }); } +void ChannelEventGenerator::unweight_job_inline(GeneratorBatchJob& job) { + auto& runtimes = _runtimes.at(job.context_index); + auto& context = _contexts.at(job.context_index); + job.max_weight = _max_weight; + if (job.rng_seed != 0) { + runtimes.unweighter->begin_job_rng(mix_seed(job.rng_seed, {kPhaseUnweight})); + } + std::vector unweighter_args( + job.events.begin(), job.events.begin() + _field_indices.random + ); + unweighter_args.push_back(Tensor(job.max_weight, context->device())); + TensorVec unw_events = runtimes.unweighter->run(unweighter_args); + job.unweighted_events.clear(); + for (auto& item : unw_events) { + job.unweighted_events.push_back(item.cpu()); + } + if (job.rng_seed != 0) { + runtimes.unweighter->end_job_rng(); + } +} + 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); diff --git a/madspace/src/driver/event_generator.cpp b/madspace/src/driver/event_generator.cpp index af8e3a59b..fefa4247f 100644 --- a/madspace/src/driver/event_generator.cpp +++ b/madspace/src/driver/event_generator.cpp @@ -1,8 +1,10 @@ #include "madspace/driver/event_generator.hpp" +#include #include #include #include +#include #include #include "madspace/driver/logger.hpp" @@ -10,6 +12,15 @@ using namespace madspace; +namespace { +// Salts keeping the driver-level random streams independent of each other for a +// given context seed (see seeded_rng()). Parallel stages additionally mix in a +// per-thread index. +constexpr std::uint32_t kSaltCombineSelect = 1; // channel selection in read_and_combine +constexpr std::uint32_t kSaltLheComplete = 2; // colour/helicity in fill_lhe_event +constexpr std::uint32_t kSaltUnweight = 3; // final unweighting +} // namespace + const GeneratorConfig EventGenerator::default_config = {}; EventGenerator::EventGenerator( @@ -41,9 +52,14 @@ EventGenerator::EventGenerator( _channel_optimizing(channels.size()), _channel_integral_fractions(channels.size(), 1.), _context_job_counts(contexts.size()), + _deterministic(!contexts.empty() && contexts.at(0)->seed() != 0), _status_file(status_file) {} void EventGenerator::survey() { + if (_deterministic) { + survey_deterministic(); + return; + } reset_start_time(); bool done = false; std::size_t min_iters = _config.survey_min_iters; @@ -136,6 +152,10 @@ void EventGenerator::survey() { } void EventGenerator::generate() { + if (_deterministic) { + generate_deterministic(); + return; + } reset_start_time(); print_gen_init(); @@ -227,6 +247,200 @@ void EventGenerator::generate() { print_gen_update(true); } +void EventGenerator::commit_generate_deterministic(GeneratorBatchJob& job) { + // Ordered, single-pass commit of one job: mirrors the generate() harvest body + // but folds unweighting in synchronously so there is exactly one commit per job. + // Called strictly in ascending job id, so every state mutation (cross-section + // accumulation, max-weight snapshot, event writes, counts, VEGAS optimisation) + // happens in a fixed order. + auto& channel = _channels.at(job.channel_index); + auto& channel_job_count = _channel_job_counts.at(job.channel_index); + auto& context_job_count = _context_job_counts.at(job.context_index); + if (job.vegas_batch_size > 0 && channel_job_count == job.split_job_count) { + channel->clear_events(); + } + channel->integrate(job); + update_integral(); + channel->update_max_weight(job.weights); + if (job.unweight) { + channel->unweight_job_inline(job); + channel->write_events(job.unweighted_events, job.max_weight); + update_counts(); + } + channel_job_count -= 1 + job.unweight; + context_job_count -= 1; + if (job.vegas_batch_size > 0 && channel_job_count == 0) { + channel->optimize_vegas(job); + _channel_optimizing.at(job.channel_index) = false; + } + print_gen_update(false); +} + +void EventGenerator::generate_deterministic() { + reset_start_time(); + print_gen_init(); + _commit_cursor = _job_id; + _ready_gen.clear(); + + std::size_t target_job_count = 0; + for (auto& context : _contexts) { + target_job_count += 2 * context->thread_pool().thread_count(); + } + std::size_t channel_index = 0; + std::size_t in_flight = 0; + while (true) { + _abort_check_function(); + + std::size_t job_count_before; + do { + job_count_before = _ready_jobs.size(); + for (std::size_t i = 0; + i < _channels.size() && _ready_jobs.size() < target_job_count; + ++i, channel_index = (channel_index + 1) % _channels.size()) { + auto& channel = _channels.at(channel_index); + double integral_frac = _channel_integral_fractions.at(channel_index); + if (integral_frac > 0 && + channel->status().count_unweighted >= + integral_frac * _config.target_count) { + continue; + } + if (channel->needs_optimization()) { + if (!_channel_optimizing.at(channel_index)) { + _channel_optimizing.at(channel_index) = true; + _ready_jobs.push_back({ + .channel_index = channel_index, + .unweight = true, + .vegas_batch_size = channel->next_vegas_batch_size(), + }); + } + } else { + _ready_jobs.push_back({ + .channel_index = channel_index, + .unweight = true, + .vegas_batch_size = 0, + }); + } + } + } while (_ready_jobs.size() - job_count_before > 0); + + std::size_t job_id_before = _job_id; + start_jobs(); + in_flight += _job_id - job_id_before; + + if (in_flight > 0) { + std::size_t job_id = _result_queue.wait(); + --in_flight; + _ready_gen.insert(job_id); + while (_ready_gen.erase(_commit_cursor) > 0) { + commit_generate_deterministic(_running_jobs.at(_commit_cursor)); + _running_jobs.erase(_commit_cursor); + ++_commit_cursor; + } + } else { + if (_status.done) { + unweight_all(); + } + if (_status.done) { + break; + } + } + } + print_gen_update(true); +} + +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(); + _ready_jobs.push_back({ + .channel_index = i, + .unweight = iter >= min_iters - 1, + .vegas_batch_size = vegas_batch_size, + }); + if (iter >= min_iters) { + total_event_count += vegas_batch_size; + } + ++i; + } + std::size_t job_id_before = _job_id; + start_jobs(); + std::size_t in_flight = _job_id - job_id_before; + done = true; + while (in_flight > 0) { + std::size_t job_id = _result_queue.wait(); + _abort_check_function(); + --in_flight; + _ready_gen.insert(job_id); + while (_ready_gen.erase(_commit_cursor) > 0) { + auto& job = _running_jobs.at(_commit_cursor); + auto& channel = _channels.at(job.channel_index); + auto& channel_job_count = _channel_job_counts.at(job.channel_index); + auto& context_job_count = _context_job_counts.at(job.context_index); + if (channel_job_count == job.split_job_count) { + channel->clear_events(); + } + channel->integrate(job); + update_integral(); + channel->update_max_weight(job.weights); + if (job.unweight) { + channel->unweight_job_inline(job); + channel->write_events(job.unweighted_events, job.max_weight); + update_counts(); + } + channel_job_count -= 1 + job.unweight; + context_job_count -= 1; + if (!job.unweight) { + done = false; + } else if (channel_job_count == 0 && + channel->cross_section().rel_error() < target_precision) { + done = false; + } + if (channel_job_count == 0) { + channel->optimize_vegas(job); + done_event_count += job.vegas_batch_size; + } + _running_jobs.erase(_commit_cursor); + ++_commit_cursor; + } + std::size_t job_id_refill = _job_id; + start_jobs(); + in_flight += _job_id - job_id_refill; + print_survey_update(false, done_event_count, total_event_count, iter); + } + } + print_survey_update(true, done_event_count, total_event_count, iter - 1); +} + bool EventGenerator::start_jobs() { std::size_t ready_index = 0, context_index = 0; for (auto [context, job_count] : zip(_contexts, _context_job_counts)) { @@ -319,6 +533,7 @@ void EventGenerator::update_counts() { void EventGenerator::combine_to_compact_npy(const std::string& file_name) { reset_start_time(); + std::mt19937 select_rng = seeded_rng(_contexts.at(0)->seed(), {kSaltCombineSelect}); auto [channel_data, particle_count, norm_factor] = init_combine(); DataLayout layout( EventRecord::layout( @@ -337,7 +552,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 +570,9 @@ void EventGenerator::combine_to_lhe_npy( const std::string& file_name, LHECompleter& lhe_completer ) { reset_start_time(); - std::random_device rand_device; - std::mt19937 rand_gen(rand_device()); + std::uint64_t seed = _contexts.at(0)->seed(); + std::mt19937 select_rng = seeded_rng(seed, {kSaltCombineSelect}); + std::mt19937 rand_gen = seeded_rng(seed, {kSaltLheComplete}); auto [channel_data, particle_count, norm_factor] = init_combine(); DataLayout in_layout( EventRecord::layout( @@ -388,7 +604,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 +634,17 @@ void EventGenerator::combine_to_lhe( const std::string& file_name, LHECompleter& lhe_completer ) { reset_start_time(); + std::uint64_t seed = _contexts.at(0)->seed(); + std::mt19937 select_rng = seeded_rng(seed, {kSaltCombineSelect}); ThreadPool pool(_config.combine_thread_count); - ThreadResource rand_gens(pool, []() { - std::random_device rand_device; - return std::mt19937(rand_device()); - }); + // Per-thread LHE-completion streams, keyed by a running thread index. Bit + // identical output requires combine_thread_pool_size == 1 (see seeded_rng()). + ThreadResource rand_gens( + pool, + [seed, next_index = std::make_shared>(0)]() { + return seeded_rng(seed, {kSaltLheComplete, next_index->fetch_add(1)}); + } + ); auto [channel_data, particle_count, norm_factor] = init_combine(); std::vector> buffers; std::vector idle_buffers; @@ -450,7 +672,7 @@ void EventGenerator::combine_to_lhe( 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); + read_and_combine(channel_data, in_buffer, norm_factor, select_rng); if (in_buffer.event_count() == 0) { done = true; break; @@ -506,8 +728,7 @@ 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::mt19937 rand_gen = seeded_rng(_contexts.at(0)->seed(), {kSaltUnweight}); bool done = true; double total_eff_count = 0.; for (auto [channel, integral_fraction] : @@ -596,7 +817,8 @@ EventGenerator::init_combine() { void EventGenerator::read_and_combine( std::vector& channel_data, EventBuffer& buffer, - double norm_factor + double norm_factor, + std::mt19937& rand_gen ) { std::size_t batch_size = 1000; std::size_t event_count = std::min(batch_size, channel_data.back().cum_count); @@ -607,8 +829,6 @@ 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); diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 5e3ccf0e2..bfb5824d0 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -244,10 +244,18 @@ PYBIND11_MODULE(_madspace_py, m) { .def("__dlpack_device__", &dlpack_device); py::classh(m, "Context") - .def(py::init(), py::arg("thread_count") = -1) .def( - py::init(), py::arg("device"), py::arg("thread_count") = -1 + py::init(), + py::arg("thread_count") = -1, + py::arg("seed") = 0 ) + .def( + py::init(), + py::arg("device"), + py::arg("thread_count") = -1, + py::arg("seed") = 0 + ) + .def_property_readonly("seed", &Context::seed) .def( "load_matrix_element", &Context::load_matrix_element, diff --git a/madspace/tests/test_determinism.py b/madspace/tests/test_determinism.py new file mode 100644 index 000000000..582f629cf --- /dev/null +++ b/madspace/tests/test_determinism.py @@ -0,0 +1,84 @@ +"""Regression tests for seeded random number generation. + +These guard the seeding plumbing that reproducible event generation rests on: a +non-zero ``Context`` seed must make every random draw a deterministic function of +that seed, distinct seeds must produce independent streams, and a seed of 0 (the +run_card default) must keep the historical non-deterministic behaviour. + +Threads: the raw ``FunctionRuntime`` path splits a batch over the thread pool and +takes each chunk's stream from whichever worker runs it, so it is only guaranteed +reproducible with a single-threaded pool. Reproducibility at higher thread counts +during *event generation* comes from the generator seeding each job by its logical +identity instead of by thread; that path needs a compiled matrix element and so is +not covered here. +""" + +import numpy as np +import pytest + +import madspace as ms + +DIM = 4 +BATCH = 20000 + + +def random_function(dim=DIM): + """Build a function whose only output is ``dim`` uniform randoms per event.""" + input_types = ms.NamedTypes([("batch_size", ms.Type([ms.batch_size]))]) + output_types = ms.NamedTypes([("r", ms.batch_float_array(dim))]) + fb = ms.FunctionBuilder(input_types, output_types) + fb.output(0, fb.random(fb.input(0), dim)) + return fb.function() + + +def draw(seed, batch_size=BATCH, thread_count=1): + """Draw one batch of randoms from a fresh context with the given seed.""" + context = ms.Context(device=ms.cpu_device(), thread_count=thread_count, seed=seed) + runtime = ms.FunctionRuntime(random_function(), context) + return np.array(runtime(np.array([batch_size], dtype=np.int64)), copy=True) + + +def test_context_exposes_seed(): + assert ms.Context(thread_count=1).seed == 0 + assert ms.Context(thread_count=1, seed=12345).seed == 12345 + + +def test_same_seed_is_reproducible(): + assert np.array_equal(draw(12345), draw(12345)) + + +def test_same_seed_is_reproducible_across_repeated_runs(): + first = draw(7) + for _ in range(3): + assert np.array_equal(first, draw(7)) + + +def test_different_seeds_differ(): + assert not np.array_equal(draw(12345), draw(12346)) + + +def test_seed_zero_is_non_deterministic(): + # 0 is the run_card default and must keep drawing a fresh stream each run + assert not np.array_equal(draw(0), draw(0)) + + +def test_draws_are_uniform(): + r = draw(2024) + assert r.shape == (BATCH, DIM) + assert np.all((r >= 0) & (r < 1)) + assert r.mean() == pytest.approx(0.5, abs=0.02) + + +def test_different_seeds_are_independent(): + # correlation of two independent streams is ~1/sqrt(n) == 0.004 here, so a + # 0.05 bound is far outside the noise while leaving no room for a shared stream + a, b = draw(1).ravel(), draw(2).ravel() + assert abs(np.corrcoef(a, b)[0, 1]) < 0.05 + + +def test_nearby_seeds_are_independent(): + # adjacent integer seeds must not give correlated streams: seeds go through + # std::seed_seq precisely so that seed=1 and seed=2 are not near each other + a, b = draw(1).ravel(), draw(2).ravel() + assert not np.array_equal(a, b) + assert abs(np.corrcoef(a, b)[0, 1]) < 0.05 From 930fd3383817d19785df1141dbeb72df621bbb1a Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Thu, 23 Jul 2026 20:39:18 +0200 Subject: [PATCH 02/19] Wire gridpack --seed through to the madspace context The --seed argument was parsed but never passed to ms.Context, so reproducible gridpack runs had no effect. Derive a per-context seed from it and update the help text to match the run_card (reproducible at any cpu_thread_pool_size). Co-Authored-By: Claude Opus 4.8 --- madgraph/iolibs/template_files/mg7/gridpack.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/madgraph/iolibs/template_files/mg7/gridpack.py b/madgraph/iolibs/template_files/mg7/gridpack.py index bfc5d5cd9..ea70a3b2a 100644 --- a/madgraph/iolibs/template_files/mg7/gridpack.py +++ b/madgraph/iolibs/template_files/mg7/gridpack.py @@ -88,7 +88,8 @@ def main() -> None: 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 (bit-identical output requires --cpu_thread_pool_size 1)" + "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( @@ -138,7 +139,7 @@ def main() -> None: device_names = args.device if args.device else run_args["devices"] contexts = [] device_types = [] - for device_name in device_names: + for i, device_name in enumerate(device_names): if ":" in device_name: device_type, device_index_str = device_name.split(":") device_index = int(device_index_str) @@ -155,7 +156,10 @@ def main() -> None: else: device = ms.cpu_device() pool_size = args.cpu_thread_pool_size - contexts.append(ms.Context(device=device, thread_count=pool_size)) + contexts.append(ms.Context( + device=device, thread_count=pool_size, + seed=derive_seed(args.seed, i), + )) # set up generator configuration config = ms.GeneratorConfig() From dd4d1fa5323b3a80e7f12e56ab1091e16141536d Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Fri, 24 Jul 2026 16:00:07 +0200 Subject: [PATCH 03/19] some refactoring of the deterministic RNG system --- .../iolibs/template_files/mg7/gridpack.py | 25 +--- .../iolibs/template_files/mg7/madevent.py | 53 +++----- madspace/include/madspace/driver/backend.hpp | 20 +-- .../madspace/driver/channel_generator.hpp | 20 ++- madspace/include/madspace/driver/context.hpp | 42 +++---- .../madspace/driver/event_generator.hpp | 24 +++- madspace/src/cpu/device.hpp | 23 ++++ madspace/src/cpu/runtime.cpp | 119 ++++++++++++------ madspace/src/cpu/runtime.hpp | 30 +++-- madspace/src/driver/channel_generator.cpp | 98 +++++++++------ madspace/src/driver/event_generator.cpp | 27 ++-- madspace/src/gpu/runtime.cu | 6 +- madspace/src/gpu/runtime.hpp | 9 +- madspace/src/python/madspace.cpp | 20 ++- 14 files changed, 307 insertions(+), 209 deletions(-) diff --git a/madgraph/iolibs/template_files/mg7/gridpack.py b/madgraph/iolibs/template_files/mg7/gridpack.py index ea70a3b2a..ccc71999b 100644 --- a/madgraph/iolibs/template_files/mg7/gridpack.py +++ b/madgraph/iolibs/template_files/mg7/gridpack.py @@ -31,23 +31,6 @@ import tomllib import argparse -def derive_seed(base_seed: int, index: int) -> int: - """Derive an independent per-context stream seed from the base seed. - - Returns 0 (seed non-deterministically) when ``base_seed`` is 0. Otherwise the - base seed and context ``index`` are mixed with splitmix64 so distinct contexts - get well-separated streams and different base seeds never collide. - """ - if base_seed == 0: - return 0 - mask = (1 << 64) - 1 - x = (base_seed + index * 0x9E3779B97F4A7C15) & mask - x = ((x ^ (x >> 30)) * 0xBF58476D1CE4E5B9) & mask - x = ((x ^ (x >> 27)) * 0x94D049BB133111EB) & mask - x = x ^ (x >> 31) - return x or 1 - - 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 @@ -139,7 +122,7 @@ def main() -> None: device_names = args.device if args.device else run_args["devices"] contexts = [] device_types = [] - 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) @@ -156,10 +139,7 @@ def main() -> None: else: device = ms.cpu_device() pool_size = args.cpu_thread_pool_size - contexts.append(ms.Context( - device=device, thread_count=pool_size, - seed=derive_seed(args.seed, i), - )) + contexts.append(ms.Context(device=device, thread_count=pool_size)) # set up generator configuration config = ms.GeneratorConfig() @@ -198,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 369d1574f..9eaed5097 100644 --- a/madgraph/iolibs/template_files/mg7/madevent.py +++ b/madgraph/iolibs/template_files/mg7/madevent.py @@ -84,24 +84,6 @@ logger = logging.getLogger("madevent7") -def derive_seed(base_seed: int, index: int) -> int: - """Derive an independent per-context stream seed from the run_card seed. - - Returns 0 (i.e. seed non-deterministically) when ``base_seed`` is 0, which is - the run_card default. Otherwise the base seed and context ``index`` are mixed - with splitmix64 so that distinct contexts get well-separated streams and two - different base seeds never collide onto the same context seed. - """ - if base_seed == 0: - return 0 - mask = (1 << 64) - 1 - x = (base_seed + index * 0x9E3779B97F4A7C15) & mask - x = ((x ^ (x >> 30)) * 0xBF58476D1CE4E5B9) & mask - x = ((x ^ (x >> 27)) * 0x94D049BB133111EB) & mask - x = x ^ (x >> 31) - return x or 1 - - def get_start_time(): return time.time(), time.process_time() @@ -217,12 +199,11 @@ def init_event_dir(self) -> None: def init_context(self) -> None: device_names = self.run_card["run"]["devices"] - base_seed = self.run_card["run"]["seed"] self.contexts = [] 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) @@ -241,11 +222,7 @@ def init_context(self) -> None: pool_size = self.run_card["run"]["cpu_thread_pool_size"] self.devices.append(device) self.pool_sizes.append(pool_size) - self.contexts.append(ms.Context( - device=device, - thread_count=pool_size, - seed=derive_seed(base_seed, i), - )) + self.contexts.append(ms.Context(device=device, thread_count=pool_size)) def parse_observable(self, name: str, order_observable: str) -> dict: parts = name.split("-") @@ -440,6 +417,7 @@ 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, + seed=self.run_card["run"]["seed"], ) unused_globals = ( set(self.contexts[0].global_names()) - event_generator.used_globals() @@ -450,36 +428,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() @@ -487,7 +471,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 = [] @@ -508,7 +492,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") @@ -563,10 +547,7 @@ def train_madnis(self) -> None: gen_context = self.contexts[0] opt_context = ms.Context( - device=self.devices[0], - thread_count=self.pool_sizes[0], - # distinct stream from the generation contexts (offset the index) - seed=derive_seed(self.run_card["run"]["seed"], 1000), + device=self.devices[0], thread_count=self.pool_sizes[0] ) opt_context.copy_globals_from(gen_context) diff --git a/madspace/include/madspace/driver/backend.hpp b/madspace/include/madspace/driver/backend.hpp index 16cf6bea7..0df00d4df 100644 --- a/madspace/include/madspace/driver/backend.hpp +++ b/madspace/include/madspace/driver/backend.hpp @@ -1,6 +1,7 @@ #pragma once #include +#include #include "madspace/compgraphs.hpp" #include "madspace/driver/context.hpp" @@ -11,9 +12,16 @@ namespace madspace { class Runtime { public: virtual ~Runtime() = default; - virtual TensorVec run(const TensorVec& inputs) = 0; + // `seed`, when set, makes the random numbers drawn during this call a + // deterministic function of `seed` alone, independent of which thread executes + // it. Backends that cannot support this (e.g. GPU, or a CPU runtime executing + // concurrently) should reject a non-null seed rather than silently ignore it. + virtual TensorVec + run(const TensorVec& inputs, std::optional seed = std::nullopt) = 0; virtual 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 ) = 0; virtual std::pair run_backward( const TensorVec& output_grads, @@ -22,14 +30,6 @@ class Runtime { bool return_contiguous_grads = false ) = 0; - // Deterministic per-job RNG. Between begin_job_rng(seed) and end_job_rng(), all - // randomness consumed by run() on the CALLING thread is drawn from a generator - // seeded by `seed`, so a job's random numbers depend only on `seed` and not on - // which pool thread executes it. Backends that cannot support this (e.g. GPU) - // keep the default no-op and fall back to their usual per-thread streams. - virtual void begin_job_rng(std::uint64_t seed) {} - virtual void end_job_rng() {} - 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 f1a66bc9b..721217bd0 100644 --- a/madspace/include/madspace/driver/channel_generator.hpp +++ b/madspace/include/madspace/driver/channel_generator.hpp @@ -61,7 +61,13 @@ class ChannelEventGenerator { 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::uint64_t seed, + bool is_survey, + std::size_t survey_pass + ); void start_unweight_job(GeneratorBatchJob& job, ResultQueue& result_queue); // Synchronous unweighting used by the deterministic generation path: runs the // unweighter on the calling (harvest) thread against the current max weight so @@ -132,10 +138,14 @@ class ChannelEventGenerator { std::optional _histogram_function; RunningIntegral _cross_section; double _max_weight = 0.; - // Monotonic per-channel counter assigned to each scheduled job (main thread) to - // key its deterministic random stream; never reset, so re-used jobs keep unique - // streams across survey/generate. - std::size_t _rng_seq = 0; + // Monotonic per-channel counters assigned to each scheduled job (main thread) to + // key its deterministic random stream; never reset. Survey and generate use + // independent counters so a generate job's seed never depends on how many + // survey jobs happened to run first, and vice versa; see + // EventGenerator::survey()'s survey_pass for why survey jobs additionally key + // off the calling survey pass. + std::size_t _survey_rng_seq = 0; + std::size_t _generate_rng_seq = 0; std::size_t _unweighted_count = 0; std::size_t _iters_without_improvement = 0; double _best_rsd = std::numeric_limits::max(); diff --git a/madspace/include/madspace/driver/context.hpp b/madspace/include/madspace/driver/context.hpp index 91f1d6b7d..b17481a1c 100644 --- a/madspace/include/madspace/driver/context.hpp +++ b/madspace/include/madspace/driver/context.hpp @@ -1,6 +1,5 @@ #pragma once -#include #include #include @@ -49,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; } @@ -141,14 +148,11 @@ class Context { * Contains global variables and matrix elements */ public: - Context(int thread_count = -1, std::uint64_t seed = 0) : + Context(int thread_count = -1) : _device(cpu_device()), - _thread_pool(std::make_unique(thread_count)), - _seed(seed) {} - Context(DevicePtr device, int thread_count = -1, std::uint64_t seed = 0) : - _device(device), - _thread_pool(std::make_unique(thread_count)), - _seed(seed) {} + _thread_pool(std::make_unique(thread_count)) {} + Context(DevicePtr device, int thread_count = -1) : + _device(device), _thread_pool(std::make_unique(thread_count)) {} Context(Context&&) = default; Context& operator=(Context&&) = default; Context(const Context&) = delete; @@ -174,14 +178,10 @@ class Context { void load_globals(const std::string& dir); DevicePtr device() { return _device; } ThreadPool& thread_pool() { return *_thread_pool; } - // Base seed for this context's random-number streams. 0 means "seed - // non-deterministically" (the historical default); see seeded_rng(). - std::uint64_t seed() const { return _seed; } private: DevicePtr _device; std::unique_ptr _thread_pool; - std::uint64_t _seed = 0; std::unordered_map> _globals; std::vector> _matrix_elements; std::vector _param_card_paths; diff --git a/madspace/include/madspace/driver/event_generator.hpp b/madspace/include/madspace/driver/event_generator.hpp index 9bff4f28b..675bf2593 100644 --- a/madspace/include/madspace/driver/event_generator.hpp +++ b/madspace/include/madspace/driver/event_generator.hpp @@ -31,9 +31,16 @@ 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::uint64_t seed = 0 ); - void survey(); + // `survey_pass` distinguishes independent survey() calls that end up + // scheduling jobs on the same ChannelEventGenerator (e.g. an initial + // multichannel survey followed by a post-simplification re-survey of a channel + // that got carried over unchanged): each pass gets its own seed salt, so a + // pass's job seeds depend only on its own logical identity, never on how many + // jobs a previous, unrelated survey pass happened to schedule first. + 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); @@ -72,7 +79,18 @@ class EventGenerator { std::vector _context_job_counts; ResultQueue _result_queue; - // Deterministic-path state. _deterministic is set from the context seed. Generate + // Base seed for reproducible event generation. 0 means "seed + // non-deterministically" (the historical default); see seeded_rng(). + std::uint64_t _seed; + + // Job-scheduling context for the currently running survey()/generate() call, + // read by start_jobs() when it hands a job's seed derivation off to its + // ChannelEventGenerator. Set once at the top of survey()/generate(), before + // the scheduling loop starts. + bool _survey_job = false; + std::size_t _survey_pass = 0; + + // Deterministic-path state. _deterministic is set from _seed. Generate // completions are buffered in _ready_gen and committed in ascending job id, with // _commit_cursor pointing at the next job id to commit. bool _deterministic = false; diff --git a/madspace/src/cpu/device.hpp b/madspace/src/cpu/device.hpp index a29bb9e5b..2a0c08f47 100644 --- a/madspace/src/cpu/device.hpp +++ b/madspace/src/cpu/device.hpp @@ -1,16 +1,29 @@ #pragma once +#include +#include + #include "madspace/driver/tensor.hpp" #include "madspace/driver/thread_pool.hpp" +#include "madspace/util.hpp" #include "simd.hpp" namespace madspace { namespace cpu { +class CpuRuntime; + class CpuDevice : public Device { public: static constexpr bool is_concurrent = false; + // Returns the generator randomness consumed by this device's op_random / + // op_random_int / op_unweight should draw from. The base implementation (used + // for CpuDevice and AsyncCpuDevice, i.e. whenever no per-call seed is in play) + // falls back to the runtime's non-deterministic per-thread stream; SeqCpuDevice + // hides this with its own seeded generator. + std::mt19937& rand_gen(CpuRuntime& runtime) const; + std::pair allocate(std::size_t size, AllocHint hint) const override { std::pair ret{new std::byte[size], Tensor()}; if (needs_zero_init(hint)) { @@ -67,6 +80,16 @@ class CpuDevice : public Device { CpuDevice() = default; }; +class SeqCpuDevice : public CpuDevice { +public: + explicit SeqCpuDevice(std::uint64_t seed) : _rand_gen(seeded_rng(seed)) {} + + std::mt19937& rand_gen(CpuRuntime&) const { return _rand_gen; } + +private: + mutable std::mt19937 _rand_gen; +}; + class AsyncCpuDevice : public CpuDevice { public: static constexpr bool is_concurrent = true; diff --git a/madspace/src/cpu/runtime.cpp b/madspace/src/cpu/runtime.cpp index eadf4b474..0366d804c 100644 --- a/madspace/src/cpu/runtime.cpp +++ b/madspace/src/cpu/runtime.cpp @@ -1,7 +1,6 @@ #include "runtime.hpp" #include -#include #include #include #include @@ -49,15 +48,6 @@ using namespace madspace::kernels; namespace { -// Per-thread deterministic job generator. When a job installs one via -// begin_job_rng(), every op_random/op_random_int/op_unweight on this thread draws -// from it instead of the pool generator, so the job's random numbers depend only -// on its seed and not on which thread happened to run it. Integrand runtimes are -// single-threaded (concurrent == false), so the whole job runs on one thread and -// a single continuous stream spans all of its random draws. -thread_local std::mt19937 t_job_rng; -thread_local bool t_job_rng_active = false; - template void op_matrix_element( const CpuRuntime::Instruction& instruction, TensorVec& locals, const D& device @@ -624,10 +614,10 @@ 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(); + auto& rand_gen = device.rand_gen(runtime); for (std::size_t i = offset; i < offset + count; ++i) { output_view[i] = dist(rand_gen); } @@ -647,10 +637,12 @@ 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(); + auto& rand_gen = device.rand_gen(runtime); for (std::size_t i = offset; i < offset + count; ++i) { output_view[i] = dist(rand_gen); } @@ -684,6 +676,7 @@ void op_unweight( indices_view_flat, uw_weights_view_flat, &runtime, + &device, batch_size, &indices, &uw_weights, @@ -694,7 +687,7 @@ void op_unweight( 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(); + auto& rand_gen = device.rand_gen(runtime); 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]; @@ -891,17 +884,10 @@ void op_discrete_histogram( } // namespace -std::mt19937& CpuRuntime::rand_gen() { - return t_job_rng_active ? t_job_rng : _rand_gens.get(); -} - -void CpuRuntime::begin_job_rng(std::uint64_t seed) { - t_job_rng = seeded_rng(seed); - t_job_rng_active = true; -} +std::mt19937& CpuRuntime::rand_gen() { return _rand_gens.get(); } -void CpuRuntime::end_job_rng() { - t_job_rng_active = false; +std::mt19937& CpuDevice::rand_gen(CpuRuntime& runtime) const { + return runtime.rand_gen(); } CpuRuntime::CpuRuntime(const Function& function, ContextPtr context, bool concurrent) : @@ -909,12 +895,9 @@ CpuRuntime::CpuRuntime(const Function& function, ContextPtr context, bool concur _input_count(function.inputs().size()), _rand_gens( context->thread_pool(), - // Fallback per-thread streams used when no per-job generator is installed - // (context seed 0, i.e. non-deterministic mode). Keyed by a running thread - // index so the streams are independent. - [seed = context->seed(), - next_index = std::make_shared>(0)]() { - return seeded_rng(seed, {next_index->fetch_add(1)}); + []() { + std::random_device rand_device; + return std::mt19937(rand_device()); } ), _concurrent(concurrent) { @@ -1097,22 +1080,36 @@ CpuRuntime::CpuRuntime(const Function& function, ContextPtr context, bool concur } } -TensorVec CpuRuntime::run(const TensorVec& inputs) { +TensorVec CpuRuntime::run(const TensorVec& inputs, std::optional seed) { if (_concurrent && _context->thread_pool().thread_count() > 1) { + if (seed) { + throw std::runtime_error( + "CpuRuntime::run: a seed cannot be used with a concurrently " + "executing runtime" + ); + } auto [outputs, locals, eval_grad] = run_concurrent(inputs, {}, false); return outputs; } else { - return run_single(inputs); + return run_single(inputs, seed); } } std::tuple> CpuRuntime::run_with_grad( - const TensorVec& inputs, const std::vector& input_requires_grad + const TensorVec& inputs, + const std::vector& input_requires_grad, + std::optional seed ) { if (_concurrent && _context->thread_pool().thread_count() > 1) { + if (seed) { + throw std::runtime_error( + "CpuRuntime::run_with_grad: a seed cannot be used with a " + "concurrently executing runtime" + ); + } return run_concurrent(inputs, input_requires_grad, true); } else { - return run_with_grad_single(inputs, input_requires_grad); + return run_with_grad_single(inputs, input_requires_grad, seed); } } @@ -1133,13 +1130,23 @@ std::pair CpuRuntime::run_backward( } } -TensorVec CpuRuntime::run_single(const TensorVec& inputs) const { - auto& device = CpuDevice::instance(); +TensorVec CpuRuntime::run_single( + const TensorVec& inputs, std::optional seed +) const { + if (seed) { + SeqCpuDevice device(*seed); + return run_single_impl(inputs, device); + } + return run_single_impl(inputs, CpuDevice::instance()); +} + +template +TensorVec CpuRuntime::run_single_impl(const TensorVec& inputs, const D& device) const { auto locals = _locals_init; std::copy(inputs.begin(), inputs.end(), locals.begin()); for (auto& instr : _instructions) { - using DeviceType = CpuDevice; + using DeviceType = D; switch (instr.opcode) { case -1: // free memory locals[instr.input_indices[0]].reset(device); @@ -1155,9 +1162,26 @@ TensorVec CpuRuntime::run_single(const TensorVec& inputs) const { } std::tuple> CpuRuntime::run_with_grad_single( - const TensorVec& inputs, const std::vector& input_requires_grad + const TensorVec& inputs, + const std::vector& input_requires_grad, + std::optional seed +) const { + if (seed) { + SeqCpuDevice device(*seed); + return run_with_grad_single_impl(inputs, input_requires_grad, device); + } + return run_with_grad_single_impl( + inputs, input_requires_grad, CpuDevice::instance() + ); +} + +template +std::tuple> +CpuRuntime::run_with_grad_single_impl( + const TensorVec& inputs, + const std::vector& input_requires_grad, + const D& device ) const { - auto& device = CpuDevice::instance(); auto locals = _locals_init; auto requires_grad = _requires_grad_init; std::vector store_local(locals.size()); @@ -1186,7 +1210,7 @@ std::tuple> CpuRuntime::run_with_grad_si } } } - using DeviceType = CpuDevice; + using DeviceType = D; switch (instr.opcode) { case -1: { // free memory auto input_index = instr.input_indices[0]; @@ -1508,3 +1532,16 @@ extern "C" Runtime* build_runtime(const Function& function, ContextPtr context, bool concurrent) { return new CpuRuntime(function, context, concurrent); } + +template TensorVec +CpuRuntime::run_single_impl(const TensorVec&, const CpuDevice&) const; +template TensorVec +CpuRuntime::run_single_impl(const TensorVec&, const SeqCpuDevice&) const; +template std::tuple> +CpuRuntime::run_with_grad_single_impl( + const TensorVec&, const std::vector&, const CpuDevice& +) const; +template std::tuple> +CpuRuntime::run_with_grad_single_impl( + const TensorVec&, const std::vector&, const SeqCpuDevice& +) const; diff --git a/madspace/src/cpu/runtime.hpp b/madspace/src/cpu/runtime.hpp index 7ebd40fa2..a5b283273 100644 --- a/madspace/src/cpu/runtime.hpp +++ b/madspace/src/cpu/runtime.hpp @@ -1,5 +1,6 @@ #pragma once +#include #include #include "madspace/compgraphs/function.hpp" @@ -37,9 +38,13 @@ class CpuRuntime : public Runtime { CpuRuntime(const Function& function, ContextPtr context, bool concurrent); - 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, @@ -49,16 +54,25 @@ class CpuRuntime : public Runtime { ) override; Context& context() { return *_context; } - // Returns the active per-job generator when one is installed on this thread - // (see begin_job_rng), otherwise the thread's pool generator. + // Fallback non-deterministic per-thread generator, used whenever no per-call + // seed is in play (see cpu::SeqCpuDevice). std::mt19937& rand_gen(); - void begin_job_rng(std::uint64_t seed) override; - void end_job_rng() override; private: - TensorVec run_single(const TensorVec& inputs) const; + TensorVec + run_single(const TensorVec& inputs, std::optional seed) const; + template + TensorVec run_single_impl(const TensorVec& inputs, const D& device) const; std::tuple> run_with_grad_single( - const TensorVec& inputs, const std::vector& input_requires_grad + const TensorVec& inputs, + const std::vector& input_requires_grad, + std::optional seed + ) const; + template + std::tuple> run_with_grad_single_impl( + const TensorVec& inputs, + const std::vector& input_requires_grad, + const D& device ) const; std::tuple> run_concurrent( const TensorVec& inputs, diff --git a/madspace/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index 2e65914de..592d1abba 100644 --- a/madspace/src/driver/channel_generator.cpp +++ b/madspace/src/driver/channel_generator.cpp @@ -11,6 +11,11 @@ namespace { constexpr std::uint64_t kPhaseGenerate = 1; constexpr std::uint64_t kPhaseUnweight = 2; +// Salts keeping a job's base-seed derivation (see start_job) independent between +// survey and generate calls. +constexpr std::uint64_t kJobKindSurvey = 1; +constexpr std::uint64_t kJobKindGenerate = 2; + int event_extra_flags(const std::unordered_map& index_map) { int flags = 0; if (index_map.contains("fact_scale1")) { @@ -394,24 +399,55 @@ double ChannelEventGenerator::channel_weight_sum(std::size_t event_count) { } void ChannelEventGenerator::start_job( - GeneratorBatchJob& job, ResultQueue& result_queue + GeneratorBatchJob& job, + ResultQueue& result_queue, + std::uint64_t seed, + bool is_survey, + std::size_t survey_pass ) { // Assign the job's deterministic base seed on the (single) scheduling thread so // it depends only on the job's logical identity, not on which worker runs it. - std::uint64_t context_seed = _contexts.at(job.context_index)->seed(); - job.rng_seed = context_seed == 0 - ? 0 - : mix_seed(context_seed, {job.channel_index, _rng_seq++}); + // Survey and generate draw from independent counters (see _survey_rng_seq / + // _generate_rng_seq) so a job's seed never depends on job counts from the other + // kind of call; survey jobs additionally key off survey_pass so that a channel + // re-surveyed across multiple independent survey() calls doesn't need its job + // count inferred from call history. + if (seed == 0) { + job.rng_seed = 0; + } else if (is_survey) { + job.rng_seed = mix_seed( + seed, + {job.context_index, + job.channel_index, + kJobKindSurvey, + survey_pass, + _survey_rng_seq++} + ); + } else { + job.rng_seed = mix_seed( + seed, + {job.context_index, + job.channel_index, + kJobKindGenerate, + _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 != 0) { - runtimes.integrand_channel->begin_job_rng( - mix_seed(job.rng_seed, {kPhaseGenerate}) - ); - } + // Every runtime->run() call below that can consume randomness gets its + // own sub-seed, uniquely salted by a running call index, derived from + // the job's base seed. This keeps each call fully deterministic without + // sharing mutable RNG state across independent Runtime objects. + std::uint64_t rng_call_index = 0; + auto next_seed = [&]() -> std::optional { + if (job.rng_seed == 0) { + return std::nullopt; + } + return mix_seed(job.rng_seed, {kPhaseGenerate, rng_call_index++}); + }; std::size_t max_batch_size = context->device()->device_type() == DeviceType::cpu ? _config.cpu_batch_size @@ -425,8 +461,9 @@ void ChannelEventGenerator::start_job( std::size_t total_count = 0, repetitions = 0; TensorVec all_ps_points; while (true) { - auto ps_points = - runtimes.integrand_channel->run({Tensor({batch_size})}); + auto ps_points = runtimes.integrand_channel->run( + {Tensor({batch_size})}, next_seed() + ); std::size_t acc_count = ps_points .at(_integrand_channel_function.outputs().index_map().at( @@ -461,13 +498,14 @@ void ChannelEventGenerator::start_job( (target_count - total_count) / cut_eff ); } - job.events = runtimes.integrand_common->run(all_ps_points); + job.events = runtimes.integrand_common->run(all_ps_points, next_seed()); job.weights = job.events.at(_field_indices.weight).cpu(); if (runtimes.observable_histograms) { auto hists = runtimes.observable_histograms->run( {job.events.at(_field_indices.weight), - job.events.at(_field_indices.momenta)} + job.events.at(_field_indices.momenta)}, + next_seed() ); for (auto& item : hists) { job.hists.push_back(item.cpu()); @@ -477,7 +515,8 @@ void ChannelEventGenerator::start_job( if (_vegas_optimizer) { auto hist = runtimes.vegas_histogram->run( {job.events.at(_field_indices.random), - job.events.at(_field_indices.weight)} + job.events.at(_field_indices.weight)}, + next_seed() ); for (auto& item : hist) { job.vegas_hist.push_back(item.cpu()); @@ -488,15 +527,12 @@ void ChannelEventGenerator::start_job( job.events.begin() + _field_indices.rest, job.events.end() }; args.push_back(job.events.at(_field_indices.weight)); - auto hist = runtimes.discrete_histogram->run(args); + auto hist = runtimes.discrete_histogram->run(args, next_seed()); for (auto& item : hist) { job.discrete_hist.push_back(item.cpu()); } } } - if (job.rng_seed != 0) { - runtimes.integrand_channel->end_job_rng(); - } result_queue.push(job.job_id); return std::nullopt; }); @@ -511,22 +547,17 @@ void ChannelEventGenerator::start_unweight_job( .submit([this, &job, &result_queue]() { auto& runtimes = _runtimes.at(job.context_index); auto& context = _contexts.at(job.context_index); - if (job.rng_seed != 0) { - runtimes.unweighter->begin_job_rng( - mix_seed(job.rng_seed, {kPhaseUnweight}) - ); - } + std::optional seed = job.rng_seed == 0 + ? std::nullopt + : std::optional(mix_seed(job.rng_seed, {kPhaseUnweight})); std::vector unweighter_args( job.events.begin(), job.events.begin() + _field_indices.random ); unweighter_args.push_back(Tensor(job.max_weight, context->device())); - TensorVec unw_events = runtimes.unweighter->run(unweighter_args); + TensorVec unw_events = runtimes.unweighter->run(unweighter_args, seed); for (auto& item : unw_events) { job.unweighted_events.push_back(item.cpu()); } - if (job.rng_seed != 0) { - runtimes.unweighter->end_job_rng(); - } result_queue.push(job.job_id); return std::nullopt; }); @@ -536,21 +567,18 @@ void ChannelEventGenerator::unweight_job_inline(GeneratorBatchJob& job) { auto& runtimes = _runtimes.at(job.context_index); auto& context = _contexts.at(job.context_index); job.max_weight = _max_weight; - if (job.rng_seed != 0) { - runtimes.unweighter->begin_job_rng(mix_seed(job.rng_seed, {kPhaseUnweight})); - } + std::optional seed = job.rng_seed == 0 + ? std::nullopt + : std::optional(mix_seed(job.rng_seed, {kPhaseUnweight})); std::vector unweighter_args( job.events.begin(), job.events.begin() + _field_indices.random ); unweighter_args.push_back(Tensor(job.max_weight, context->device())); - TensorVec unw_events = runtimes.unweighter->run(unweighter_args); + TensorVec unw_events = runtimes.unweighter->run(unweighter_args, seed); job.unweighted_events.clear(); for (auto& item : unw_events) { job.unweighted_events.push_back(item.cpu()); } - if (job.rng_seed != 0) { - runtimes.unweighter->end_job_rng(); - } } std::size_t ChannelEventGenerator::next_vegas_batch_size() { diff --git a/madspace/src/driver/event_generator.cpp b/madspace/src/driver/event_generator.cpp index fefa4247f..b3c4cffd3 100644 --- a/madspace/src/driver/event_generator.cpp +++ b/madspace/src/driver/event_generator.cpp @@ -18,7 +18,7 @@ namespace { // per-thread index. constexpr std::uint32_t kSaltCombineSelect = 1; // channel selection in read_and_combine constexpr std::uint32_t kSaltLheComplete = 2; // colour/helicity in fill_lhe_event -constexpr std::uint32_t kSaltUnweight = 3; // final unweighting +constexpr std::uint32_t kSaltUnweight = 3; // final unweighting } // namespace const GeneratorConfig EventGenerator::default_config = {}; @@ -27,7 +27,8 @@ EventGenerator::EventGenerator( const std::vector& contexts, const std::vector>& channels, const std::string& status_file, - const GeneratorConfig& config + const GeneratorConfig& config, + std::uint64_t seed ) : _config(config), _status{ @@ -52,10 +53,13 @@ EventGenerator::EventGenerator( _channel_optimizing(channels.size()), _channel_integral_fractions(channels.size(), 1.), _context_job_counts(contexts.size()), - _deterministic(!contexts.empty() && contexts.at(0)->seed() != 0), + _seed(seed), + _deterministic(seed != 0), _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; @@ -152,6 +156,7 @@ void EventGenerator::survey() { } void EventGenerator::generate() { + _survey_job = false; if (_deterministic) { generate_deterministic(); return; @@ -461,7 +466,8 @@ 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; @@ -533,7 +539,7 @@ void EventGenerator::update_counts() { void EventGenerator::combine_to_compact_npy(const std::string& file_name) { reset_start_time(); - std::mt19937 select_rng = seeded_rng(_contexts.at(0)->seed(), {kSaltCombineSelect}); + std::mt19937 select_rng = seeded_rng(_seed, {kSaltCombineSelect}); auto [channel_data, particle_count, norm_factor] = init_combine(); DataLayout layout( EventRecord::layout( @@ -570,7 +576,7 @@ void EventGenerator::combine_to_lhe_npy( const std::string& file_name, LHECompleter& lhe_completer ) { reset_start_time(); - std::uint64_t seed = _contexts.at(0)->seed(); + std::uint64_t seed = _seed; std::mt19937 select_rng = seeded_rng(seed, {kSaltCombineSelect}); std::mt19937 rand_gen = seeded_rng(seed, {kSaltLheComplete}); auto [channel_data, particle_count, norm_factor] = init_combine(); @@ -634,14 +640,13 @@ void EventGenerator::combine_to_lhe( const std::string& file_name, LHECompleter& lhe_completer ) { reset_start_time(); - std::uint64_t seed = _contexts.at(0)->seed(); + std::uint64_t seed = _seed; std::mt19937 select_rng = seeded_rng(seed, {kSaltCombineSelect}); ThreadPool pool(_config.combine_thread_count); // Per-thread LHE-completion streams, keyed by a running thread index. Bit // identical output requires combine_thread_pool_size == 1 (see seeded_rng()). ThreadResource rand_gens( - pool, - [seed, next_index = std::make_shared>(0)]() { + pool, [seed, next_index = std::make_shared>(0)]() { return seeded_rng(seed, {kSaltLheComplete, next_index->fetch_add(1)}); } ); @@ -728,7 +733,7 @@ void EventGenerator::add_timing_data(const std::string& key) { } void EventGenerator::unweight_all() { - std::mt19937 rand_gen = seeded_rng(_contexts.at(0)->seed(), {kSaltUnweight}); + std::mt19937 rand_gen = seeded_rng(_seed, {kSaltUnweight}); bool done = true; double total_eff_count = 0.; for (auto [channel, integral_fraction] : diff --git a/madspace/src/gpu/runtime.cu b/madspace/src/gpu/runtime.cu index 87ff0b0be..2096371d8 100644 --- a/madspace/src/gpu/runtime.cu +++ b/madspace/src/gpu/runtime.cu @@ -1610,7 +1610,7 @@ GpuRuntime::GpuRuntime(const Function& function_arg, ContextPtr context) : ); } -TensorVec GpuRuntime::run(const TensorVec& inputs) { +TensorVec GpuRuntime::run(const TensorVec& inputs, std::optional) { auto& gpu_device = *static_cast(_context->device()); auto& streams = _streams.get(); auto& events = _events.get(); @@ -1650,7 +1650,9 @@ 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 ) { auto& gpu_device = *static_cast(_context->device()); auto& streams = _streams.get(); 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/python/madspace.cpp b/madspace/src/python/madspace.cpp index bfb5824d0..7f38194bc 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -244,18 +244,10 @@ PYBIND11_MODULE(_madspace_py, m) { .def("__dlpack_device__", &dlpack_device); py::classh(m, "Context") + .def(py::init(), py::arg("thread_count") = -1) .def( - py::init(), - py::arg("thread_count") = -1, - py::arg("seed") = 0 - ) - .def( - py::init(), - py::arg("device"), - py::arg("thread_count") = -1, - py::arg("seed") = 0 + py::init(), py::arg("device"), py::arg("thread_count") = -1 ) - .def_property_readonly("seed", &Context::seed) .def( "load_matrix_element", &Context::load_matrix_element, @@ -1721,7 +1713,8 @@ PYBIND11_MODULE(_madspace_py, m) { const std::vector&, const std::vector>&, const std::string&, - const GeneratorConfig&>(), + const GeneratorConfig&, + std::uint64_t>(), py::arg("contexts"), py::arg("channels"), py::arg("status_file") = "", @@ -1729,9 +1722,10 @@ PYBIND11_MODULE(_madspace_py, m) { "config", EventGenerator::default_config, "EventGenerator.default_config" - ) + ), + py::arg("seed") = 0 ) - .def("survey", &EventGenerator::survey) + .def("survey", &EventGenerator::survey, py::arg("survey_pass") = 0) .def("generate", &EventGenerator::generate) .def( "combine_to_compact_npy", From cc3e28a80830ea4e0d10210af2998e30b7a9a0c8 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Fri, 24 Jul 2026 16:41:24 +0200 Subject: [PATCH 04/19] fix race condition from trailing jobs --- .../madspace/driver/event_generator.hpp | 17 +++++++++++ .../madspace/driver/generator_data.hpp | 5 ++++ madspace/src/driver/event_generator.cpp | 30 +++++++++++++++++-- 3 files changed, 50 insertions(+), 2 deletions(-) diff --git a/madspace/include/madspace/driver/event_generator.hpp b/madspace/include/madspace/driver/event_generator.hpp index 675bf2593..959637d1d 100644 --- a/madspace/include/madspace/driver/event_generator.hpp +++ b/madspace/include/madspace/driver/event_generator.hpp @@ -77,6 +77,15 @@ class EventGenerator { std::vector _channel_optimizing; std::vector _channel_integral_fractions; std::vector _context_job_counts; + // Estimated count_unweighted contribution of jobs that have been dispatched for + // a channel but not yet committed, keyed by channel index. Added to a channel's + // committed count_unweighted when deciding whether to schedule more work for it, + // so that jobs already in flight (but whose actual contribution isn't known + // until they commit) aren't invisible to that decision. Without this, a burst of + // already-dispatched jobs can all commit after a channel's target was already + // met, overshooting it by an amount that depends on real scheduling timing + // rather than on the seed -- see GeneratorBatchJob::reserved_count. + std::vector _channel_reserved_count; ResultQueue _result_queue; // Base seed for reproducible event generation. 0 means "seed @@ -115,6 +124,14 @@ class EventGenerator { void generate_deterministic(); void commit_generate_deterministic(GeneratorBatchJob& job); + // Estimate a not-yet-dispatched job's eventual count_unweighted contribution + // from the channel's committed efficiency so far (1.0, i.e. the full batch, + // before any data is committed), add it to _channel_reserved_count, and return + // it so the caller can store it on the job for release_reserved_count() once + // the job commits. See _channel_reserved_count for why this is needed. + double reserve_job_count(std::size_t channel_index); + void release_reserved_count(const GeneratorBatchJob& job); + bool start_jobs(); void update_integral(); void update_counts(); diff --git a/madspace/include/madspace/driver/generator_data.hpp b/madspace/include/madspace/driver/generator_data.hpp index f01104136..3f7114369 100644 --- a/madspace/include/madspace/driver/generator_data.hpp +++ b/madspace/include/madspace/driver/generator_data.hpp @@ -108,6 +108,11 @@ struct GeneratorBatchJob { // time from (context seed, channel, per-channel job sequence). 0 means "use the // non-deterministic pool generator" (context seed 0). See ChannelEventGenerator. std::uint64_t rng_seed = 0; + // Estimated contribution to the channel's count_unweighted, reserved against + // the channel's target at the time this job was scheduled and released once its + // actual contribution has been counted. See + // EventGenerator::_channel_reserved_count. + double reserved_count = 0.; }; void to_json(nlohmann::json& j, const GeneratorStatus& status); diff --git a/madspace/src/driver/event_generator.cpp b/madspace/src/driver/event_generator.cpp index b3c4cffd3..20493217a 100644 --- a/madspace/src/driver/event_generator.cpp +++ b/madspace/src/driver/event_generator.cpp @@ -53,6 +53,7 @@ EventGenerator::EventGenerator( _channel_optimizing(channels.size()), _channel_integral_fractions(channels.size(), 1.), _context_job_counts(contexts.size()), + _channel_reserved_count(channels.size()), _seed(seed), _deterministic(seed != 0), _status_file(status_file) {} @@ -182,7 +183,8 @@ void EventGenerator::generate() { std::size_t& channel_job_count = _channel_job_counts.at(channel_index); double integral_frac = _channel_integral_fractions.at(channel_index); if (integral_frac > 0 && - channel->status().count_unweighted >= + channel->status().count_unweighted + + _channel_reserved_count.at(channel_index) >= integral_frac * _config.target_count) { continue; } @@ -200,6 +202,7 @@ void EventGenerator::generate() { .channel_index = channel_index, .unweight = true, .vegas_batch_size = 0, + .reserved_count = reserve_job_count(channel_index), }); } } @@ -214,6 +217,7 @@ void EventGenerator::generate() { auto& context_job_count = _context_job_counts.at(job.context_index); if (job.vegas_batch_size > 0 && channel_job_count == job.split_job_count) { channel->clear_events(); + _channel_reserved_count.at(job.channel_index) = 0.; } --channel_job_count; --context_job_count; @@ -231,6 +235,7 @@ void EventGenerator::generate() { } else { channel->write_events(job.unweighted_events, job.max_weight); update_counts(); + release_reserved_count(job); } if (job.vegas_batch_size > 0 && channel_job_count == 0) { channel->optimize_vegas(job); @@ -263,6 +268,7 @@ void EventGenerator::commit_generate_deterministic(GeneratorBatchJob& job) { auto& context_job_count = _context_job_counts.at(job.context_index); if (job.vegas_batch_size > 0 && channel_job_count == job.split_job_count) { channel->clear_events(); + _channel_reserved_count.at(job.channel_index) = 0.; } channel->integrate(job); update_integral(); @@ -271,6 +277,7 @@ void EventGenerator::commit_generate_deterministic(GeneratorBatchJob& job) { channel->unweight_job_inline(job); channel->write_events(job.unweighted_events, job.max_weight); update_counts(); + release_reserved_count(job); } channel_job_count -= 1 + job.unweight; context_job_count -= 1; @@ -305,7 +312,8 @@ void EventGenerator::generate_deterministic() { 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_unweighted + + _channel_reserved_count.at(channel_index) >= integral_frac * _config.target_count) { continue; } @@ -323,6 +331,7 @@ void EventGenerator::generate_deterministic() { .channel_index = channel_index, .unweight = true, .vegas_batch_size = 0, + .reserved_count = reserve_job_count(channel_index), }); } } @@ -446,6 +455,23 @@ void EventGenerator::survey_deterministic() { print_survey_update(true, done_event_count, total_event_count, iter - 1); } +double EventGenerator::reserve_job_count(std::size_t channel_index) { + auto& status = _channels.at(channel_index)->status(); + // Before any data has committed for this channel, assume the pessimistic 1.0 + // (a whole batch's worth), so we under- rather than over-schedule until the + // real efficiency is known. + double efficiency = status.count_opt > 0 + ? status.count_unweighted / static_cast(status.count_opt) + : 1.; + double reserved = efficiency * static_cast(_config.cpu_batch_size); + _channel_reserved_count.at(channel_index) += reserved; + return reserved; +} + +void EventGenerator::release_reserved_count(const GeneratorBatchJob& job) { + _channel_reserved_count.at(job.channel_index) -= job.reserved_count; +} + bool EventGenerator::start_jobs() { std::size_t ready_index = 0, context_index = 0; for (auto [context, job_count] : zip(_contexts, _context_job_counts)) { From 217ff0301f50e46a9abb6b7b81fdfef3b36b214b Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Fri, 24 Jul 2026 16:57:32 +0200 Subject: [PATCH 05/19] make seeding independent of context index --- madspace/src/driver/channel_generator.cpp | 34 ++++++++++------------- 1 file changed, 15 insertions(+), 19 deletions(-) diff --git a/madspace/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index 592d1abba..32fa2164b 100644 --- a/madspace/src/driver/channel_generator.cpp +++ b/madspace/src/driver/channel_generator.cpp @@ -406,31 +406,27 @@ void ChannelEventGenerator::start_job( std::size_t survey_pass ) { // Assign the job's deterministic base seed on the (single) scheduling thread so - // it depends only on the job's logical identity, not on which worker runs it. - // Survey and generate draw from independent counters (see _survey_rng_seq / - // _generate_rng_seq) so a job's seed never depends on job counts from the other - // kind of call; survey jobs additionally key off survey_pass so that a channel - // re-surveyed across multiple independent survey() calls doesn't need its job - // count inferred from call history. + // it depends only on the job's logical identity, not on which worker runs it -- + // deliberately not on job.context_index either, so the same logical job gets the + // same seed no matter how many contexts/threads are available to run it or which + // one it lands on (context assignment is a load-balancing detail, see + // start_jobs()). This is safe because _survey_rng_seq/_generate_rng_seq are + // incremented exactly once per logical job, on the single scheduling thread, + // before dispatch -- so a given (channel, seq) pair is never reused regardless + // of context assignment. Survey and generate draw from independent counters so + // a job's seed never depends on job counts from the other kind of call; survey + // jobs additionally key off survey_pass so that a channel re-surveyed across + // multiple independent survey() calls doesn't need its job count inferred from + // call history. if (seed == 0) { job.rng_seed = 0; } else if (is_survey) { job.rng_seed = mix_seed( - seed, - {job.context_index, - job.channel_index, - kJobKindSurvey, - survey_pass, - _survey_rng_seq++} + seed, {job.channel_index, kJobKindSurvey, survey_pass, _survey_rng_seq++} ); } else { - job.rng_seed = mix_seed( - seed, - {job.context_index, - job.channel_index, - kJobKindGenerate, - _generate_rng_seq++} - ); + job.rng_seed = + mix_seed(seed, {job.channel_index, kJobKindGenerate, _generate_rng_seq++}); } _contexts.at(job.context_index) ->thread_pool() From a69289b61048254963b4ff5ebd91c51b200413ce Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Fri, 24 Jul 2026 17:47:28 +0200 Subject: [PATCH 06/19] make reproducibility independent of thread pool size --- .../madspace/driver/event_generator.hpp | 35 +++--- .../madspace/driver/generator_data.hpp | 26 ++++- madspace/src/driver/channel_generator.cpp | 2 +- madspace/src/driver/event_generator.cpp | 107 +++++++++--------- madspace/src/python/madspace.cpp | 6 +- 5 files changed, 101 insertions(+), 75 deletions(-) diff --git a/madspace/include/madspace/driver/event_generator.hpp b/madspace/include/madspace/driver/event_generator.hpp index 959637d1d..71ae8f7bf 100644 --- a/madspace/include/madspace/driver/event_generator.hpp +++ b/madspace/include/madspace/driver/event_generator.hpp @@ -77,15 +77,16 @@ class EventGenerator { std::vector _channel_optimizing; std::vector _channel_integral_fractions; std::vector _context_job_counts; - // Estimated count_unweighted contribution of jobs that have been dispatched for - // a channel but not yet committed, keyed by channel index. Added to a channel's - // committed count_unweighted when deciding whether to schedule more work for it, - // so that jobs already in flight (but whose actual contribution isn't known - // until they commit) aren't invisible to that decision. Without this, a burst of - // already-dispatched jobs can all commit after a channel's target was already - // met, overshooting it by an amount that depends on real scheduling timing - // rather than on the seed -- see GeneratorBatchJob::reserved_count. - std::vector _channel_reserved_count; + // True while a channel has a steady-state (non-VEGAS) generation batch + // dispatched but not yet fully committed. Mirrors _channel_optimizing's role + // for VEGAS batches: while true, no new batch is created for this channel, so + // the next batch's size (see next_batch_job_count()) is always computed from + // fully-committed data. Without this, jobs already dispatched but not yet + // committed are invisible to the "does this channel need more work" decision, + // so a burst of them can all commit after the channel's target was already met + // -- overshooting it by an amount that depends on real scheduling timing rather + // than on the seed. + std::vector _channel_batch_pending; ResultQueue _result_queue; // Base seed for reproducible event generation. 0 means "seed @@ -124,13 +125,15 @@ class EventGenerator { void generate_deterministic(); void commit_generate_deterministic(GeneratorBatchJob& job); - // Estimate a not-yet-dispatched job's eventual count_unweighted contribution - // from the channel's committed efficiency so far (1.0, i.e. the full batch, - // before any data is committed), add it to _channel_reserved_count, and return - // it so the caller can store it on the job for release_reserved_count() once - // the job commits. See _channel_reserved_count for why this is needed. - double reserve_job_count(std::size_t channel_index); - void release_reserved_count(const GeneratorBatchJob& job); + // Size a new steady-state generation batch for a channel from fully-committed + // data only: the channel's own observed efficiency so far (count_unweighted / + // count_opt, or the pessimistic 1.0 -- i.e. a whole raw batch -- before any data + // exists yet) and its remaining need (target - count_unweighted). Capped to + // generation_batch_fraction of the remaining need so a single batch never bets + // everything on an estimate that may still be noisy, and floored at + // min_batch_jobs so a lone active channel still gets enough parallel jobs to + // keep all worker threads busy. + std::size_t next_batch_job_count(std::size_t channel_index, double target) const; bool start_jobs(); void update_integral(); diff --git a/madspace/include/madspace/driver/generator_data.hpp b/madspace/include/madspace/driver/generator_data.hpp index 3f7114369..687e22dbc 100644 --- a/madspace/include/madspace/driver/generator_data.hpp +++ b/madspace/include/madspace/driver/generator_data.hpp @@ -63,6 +63,17 @@ struct GeneratorConfig { int combine_thread_count = -1; double cut_efficiency_threshold = 0.7; std::size_t max_cut_repetitions = 100; + // Steady-state event generation dispatches a channel's jobs in barrier-gated + // batches (a new batch is only sized/created once the previous one has fully + // committed), so a batch's size is a deterministic function of committed data, + // never of scheduling timing. generation_batch_fraction caps a batch to this + // fraction of the channel's currently remaining need, so a single batch never + // bets everything on an efficiency estimate that may still be noisy (e.g. for + // costly/high-variance channels early on); min_batch_jobs floors the batch size + // so a lone active channel still gets enough parallel jobs to keep all worker + // threads busy. See EventGenerator::next_batch_job_count(). + double generation_batch_fraction = 0.5; + std::size_t min_batch_jobs = 1; }; struct GeneratorStatus { @@ -93,6 +104,11 @@ struct Histogram { struct GeneratorBatchJob { std::size_t channel_index; bool unweight; + // Nonzero to make start_jobs() split this ready_job into ceil(vegas_batch_size / + // context_batch_size) sub-jobs and dispatch all of them atomically (bypassing + // the usual per-context dispatch cap), so channel_job_count reflects the whole + // batch from the moment it's dispatched rather than only whichever prefix fit + // under the cap -- see is_vegas_batch for what the batch is *for*. std::size_t vegas_batch_size; std::size_t split_job_count; Tensor weights; @@ -108,11 +124,11 @@ struct GeneratorBatchJob { // time from (context seed, channel, per-channel job sequence). 0 means "use the // non-deterministic pool generator" (context seed 0). See ChannelEventGenerator. std::uint64_t rng_seed = 0; - // Estimated contribution to the channel's count_unweighted, reserved against - // the channel's target at the time this job was scheduled and released once its - // actual contribution has been counted. See - // EventGenerator::_channel_reserved_count. - double reserved_count = 0.; + // True for a VEGAS-grid-optimization batch (triggers clear_events()/ + // optimize_vegas() at the start/end of the batch); false for a steady-state + // generation batch (see EventGenerator::_channel_batch_pending), even though + // both use vegas_batch_size > 0 to get the same atomic whole-batch dispatch. + bool is_vegas_batch = false; }; void to_json(nlohmann::json& j, const GeneratorStatus& status); diff --git a/madspace/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index 32fa2164b..a038f3602 100644 --- a/madspace/src/driver/channel_generator.cpp +++ b/madspace/src/driver/channel_generator.cpp @@ -507,7 +507,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), diff --git a/madspace/src/driver/event_generator.cpp b/madspace/src/driver/event_generator.cpp index 20493217a..e8427f20e 100644 --- a/madspace/src/driver/event_generator.cpp +++ b/madspace/src/driver/event_generator.cpp @@ -53,7 +53,7 @@ EventGenerator::EventGenerator( _channel_optimizing(channels.size()), _channel_integral_fractions(channels.size(), 1.), _context_job_counts(contexts.size()), - _channel_reserved_count(channels.size()), + _channel_batch_pending(channels.size()), _seed(seed), _deterministic(seed != 0), _status_file(status_file) {} @@ -101,6 +101,7 @@ void EventGenerator::survey(std::size_t survey_pass) { .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; @@ -183,8 +184,7 @@ void EventGenerator::generate() { std::size_t& channel_job_count = _channel_job_counts.at(channel_index); double integral_frac = _channel_integral_fractions.at(channel_index); if (integral_frac > 0 && - channel->status().count_unweighted + - _channel_reserved_count.at(channel_index) >= + channel->status().count_unweighted >= integral_frac * _config.target_count) { continue; } @@ -195,6 +195,7 @@ void EventGenerator::generate() { .channel_index = channel_index, .unweight = true, .vegas_batch_size = channel->next_vegas_batch_size(), + .is_vegas_batch = true, }); } } else { @@ -202,7 +203,6 @@ void EventGenerator::generate() { .channel_index = channel_index, .unweight = true, .vegas_batch_size = 0, - .reserved_count = reserve_job_count(channel_index), }); } } @@ -217,7 +217,6 @@ void EventGenerator::generate() { auto& context_job_count = _context_job_counts.at(job.context_index); if (job.vegas_batch_size > 0 && channel_job_count == job.split_job_count) { channel->clear_events(); - _channel_reserved_count.at(job.channel_index) = 0.; } --channel_job_count; --context_job_count; @@ -235,7 +234,6 @@ void EventGenerator::generate() { } else { channel->write_events(job.unweighted_events, job.max_weight); update_counts(); - release_reserved_count(job); } if (job.vegas_batch_size > 0 && channel_job_count == 0) { channel->optimize_vegas(job); @@ -266,9 +264,8 @@ void EventGenerator::commit_generate_deterministic(GeneratorBatchJob& job) { auto& channel = _channels.at(job.channel_index); auto& channel_job_count = _channel_job_counts.at(job.channel_index); auto& context_job_count = _context_job_counts.at(job.context_index); - if (job.vegas_batch_size > 0 && channel_job_count == job.split_job_count) { + if (job.is_vegas_batch && channel_job_count == job.split_job_count) { channel->clear_events(); - _channel_reserved_count.at(job.channel_index) = 0.; } channel->integrate(job); update_integral(); @@ -277,13 +274,14 @@ void EventGenerator::commit_generate_deterministic(GeneratorBatchJob& job) { channel->unweight_job_inline(job); channel->write_events(job.unweighted_events, job.max_weight); update_counts(); - release_reserved_count(job); } channel_job_count -= 1 + job.unweight; context_job_count -= 1; - if (job.vegas_batch_size > 0 && channel_job_count == 0) { + if (job.is_vegas_batch && channel_job_count == 0) { channel->optimize_vegas(job); _channel_optimizing.at(job.channel_index) = false; + } else if (!job.is_vegas_batch && channel_job_count == 0) { + _channel_batch_pending.at(job.channel_index) = false; } print_gen_update(false); } @@ -294,48 +292,48 @@ void EventGenerator::generate_deterministic() { _commit_cursor = _job_id; _ready_gen.clear(); - std::size_t target_job_count = 0; - for (auto& context : _contexts) { - target_job_count += 2 * context->thread_pool().thread_count(); - } - std::size_t channel_index = 0; std::size_t in_flight = 0; while (true) { _abort_check_function(); - std::size_t job_count_before; - do { - job_count_before = _ready_jobs.size(); - for (std::size_t i = 0; - i < _channels.size() && _ready_jobs.size() < target_job_count; - ++i, channel_index = (channel_index + 1) % _channels.size()) { - auto& channel = _channels.at(channel_index); - double integral_frac = _channel_integral_fractions.at(channel_index); - if (integral_frac > 0 && - channel->status().count_unweighted + - _channel_reserved_count.at(channel_index) >= - integral_frac * _config.target_count) { - continue; - } - if (channel->needs_optimization()) { - if (!_channel_optimizing.at(channel_index)) { - _channel_optimizing.at(channel_index) = true; - _ready_jobs.push_back({ - .channel_index = channel_index, - .unweight = true, - .vegas_batch_size = channel->next_vegas_batch_size(), - }); - } - } else { + 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); + double target = integral_frac * _config.target_count; + if (integral_frac > 0 && channel->status().count_unweighted >= target) { + continue; + } + if (channel->needs_optimization()) { + if (!_channel_optimizing.at(channel_index)) { + _channel_optimizing.at(channel_index) = true; _ready_jobs.push_back({ .channel_index = channel_index, .unweight = true, - .vegas_batch_size = 0, - .reserved_count = reserve_job_count(channel_index), + .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; + std::size_t batch_jobs = next_batch_job_count(channel_index, target); + // Pushed as a single ready_job with vegas_batch_size set (rather + // than batch_jobs separate ready_jobs) so start_jobs() dispatches + // the whole batch atomically -- see + // GeneratorBatchJob::vegas_batch_size. Otherwise, once batch_jobs + // exceeds the per-context dispatch cap, channel_job_count could reach 0 + // (looking "done") as soon as whichever prefix got dispatched all + // commit, while the rest of the batch is still sitting undispatched -- + // reintroducing exactly the kind of scheduling-timing-dependent + // overshoot this batching was meant to eliminate. + _ready_jobs.push_back({ + .channel_index = channel_index, + .unweight = true, + .vegas_batch_size = batch_jobs * _config.cpu_batch_size, + .is_vegas_batch = false, + }); } - } while (_ready_jobs.size() - job_count_before > 0); + } std::size_t job_id_before = _job_id; start_jobs(); @@ -400,6 +398,7 @@ void EventGenerator::survey_deterministic() { .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; @@ -455,21 +454,25 @@ void EventGenerator::survey_deterministic() { print_survey_update(true, done_event_count, total_event_count, iter - 1); } -double EventGenerator::reserve_job_count(std::size_t channel_index) { +std::size_t +EventGenerator::next_batch_job_count(std::size_t channel_index, double target) const { auto& status = _channels.at(channel_index)->status(); // Before any data has committed for this channel, assume the pessimistic 1.0 - // (a whole batch's worth), so we under- rather than over-schedule until the - // real efficiency is known. + // (i.e. a whole raw batch becomes unweighted events), so the very first batch + // under- rather than over-estimates how many jobs are needed. double efficiency = status.count_opt > 0 ? status.count_unweighted / static_cast(status.count_opt) : 1.; - double reserved = efficiency * static_cast(_config.cpu_batch_size); - _channel_reserved_count.at(channel_index) += reserved; - return reserved; -} - -void EventGenerator::release_reserved_count(const GeneratorBatchJob& job) { - _channel_reserved_count.at(job.channel_index) -= job.reserved_count; + // Floor the assumed per-job contribution at 1 event: guards against a + // near-zero observed efficiency (e.g. an unlucky small first sample) blowing + // this up into an unbounded batch size. + double per_job_estimate = + std::max(efficiency * static_cast(_config.cpu_batch_size), 1.); + double remaining = std::max(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)) + ); } bool EventGenerator::start_jobs() { diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 7f38194bc..348d0d997 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -1429,7 +1429,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<>()) From 56044391a47a7ccfa09cce317b6b4c87aefd1a50 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Sat, 25 Jul 2026 00:12:38 +0200 Subject: [PATCH 07/19] only enforce ordering per channel, perform synchronization of integral fractions in rounds --- .../madspace/driver/event_generator.hpp | 29 +++++ madspace/src/driver/event_generator.cpp | 105 +++++++++++++----- 2 files changed, 109 insertions(+), 25 deletions(-) diff --git a/madspace/include/madspace/driver/event_generator.hpp b/madspace/include/madspace/driver/event_generator.hpp index 71ae8f7bf..1e8754a2d 100644 --- a/madspace/include/madspace/driver/event_generator.hpp +++ b/madspace/include/madspace/driver/event_generator.hpp @@ -87,6 +87,19 @@ class EventGenerator { // -- overshooting it by an amount that depends on real scheduling timing rather // than on the seed. std::vector _channel_batch_pending; + // generate_deterministic() only: per-channel analogue of _ready_gen/ + // _commit_cursor. Each channel's jobs still commit strictly in that channel's + // own ascending job-id order (out-of-order arrivals buffered in + // _channel_ready_gen until the cursor catches up), streamed as soon as they're + // next in line rather than deferred -- but channels no longer block each + // other's commits, so a channel with much costlier jobs (e.g. higher final- + // state multiplicity in an inclusive run) can't head-of-line-block the rest. + // _channel_cursor_set tracks whether _channel_commit_cursor has been + // initialized for the channel's current round yet; see + // register_dispatched_ids(). + std::vector> _channel_ready_gen; + std::vector _channel_commit_cursor; + std::vector _channel_cursor_set; ResultQueue _result_queue; // Base seed for reproducible event generation. 0 means "seed @@ -124,6 +137,15 @@ class EventGenerator { void survey_deterministic(); void generate_deterministic(); void commit_generate_deterministic(GeneratorBatchJob& job); + // generate_deterministic() only: scans the job ids newly added to + // _running_jobs by the last start_jobs() call ([first_id, end_id)) and, for + // any channel seen for the first time since its current round started (i.e. + // _channel_cursor_set is still false), initializes _channel_commit_cursor to + // that job's id. Safe because a channel's whole round is always dispatched + // atomically as one contiguous id range (see GeneratorBatchJob::vegas_batch_size), + // so the first id encountered for a channel in ascending order is always that + // round's minimum. + void register_dispatched_ids(std::size_t first_id, std::size_t end_id); // Size a new steady-state generation batch for a channel from fully-committed // data only: the channel's own observed efficiency so far (count_unweighted / @@ -136,7 +158,14 @@ class EventGenerator { std::size_t next_batch_job_count(std::size_t channel_index, double target) const; bool start_jobs(); + // update_integral() = update_integral_status() + update_integral_fractions(). + // Split so generate_deterministic() can keep the status half (used for + // logging/progress only) fully live, per commit, while deferring the fractions + // half (which feeds back into per-channel scheduling decisions, see + // next_batch_job_count()) to once per round -- see commit_generate_deterministic(). 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); diff --git a/madspace/src/driver/event_generator.cpp b/madspace/src/driver/event_generator.cpp index e8427f20e..420d6eec1 100644 --- a/madspace/src/driver/event_generator.cpp +++ b/madspace/src/driver/event_generator.cpp @@ -54,6 +54,9 @@ EventGenerator::EventGenerator( _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()), _seed(seed), _deterministic(seed != 0), _status_file(status_file) {} @@ -258,9 +261,16 @@ void EventGenerator::generate() { void EventGenerator::commit_generate_deterministic(GeneratorBatchJob& job) { // Ordered, single-pass commit of one job: mirrors the generate() harvest body // but folds unweighting in synchronously so there is exactly one commit per job. - // Called strictly in ascending job id, so every state mutation (cross-section - // accumulation, max-weight snapshot, event writes, counts, VEGAS optimisation) - // happens in a fixed order. + // Called strictly in ascending job id *within job.channel_index* (see + // _channel_commit_cursor), so every state mutation local to that channel + // (cross-section accumulation, max-weight snapshot, event writes, counts, VEGAS + // optimisation) happens in a fixed order. update_integral_status() runs on every + // commit, same as before, so logging/progress reporting stays fully live; only + // update_integral_fractions() -- the piece that feeds back into per-channel + // scheduling decisions -- is deferred to once per round in + // generate_deterministic(), since reading it here would make a channel's + // scheduling depend on how far other channels happen to have gotten in real + // time. auto& channel = _channels.at(job.channel_index); auto& channel_job_count = _channel_job_counts.at(job.channel_index); auto& context_job_count = _context_job_counts.at(job.context_index); @@ -268,7 +278,7 @@ void EventGenerator::commit_generate_deterministic(GeneratorBatchJob& job) { channel->clear_events(); } channel->integrate(job); - update_integral(); + update_integral_status(); channel->update_max_weight(job.weights); if (job.unweight) { channel->unweight_job_inline(job); @@ -286,13 +296,20 @@ void EventGenerator::commit_generate_deterministic(GeneratorBatchJob& job) { print_gen_update(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_cursor_set.at(channel_index) = true; + } + } +} + void EventGenerator::generate_deterministic() { reset_start_time(); print_gen_init(); - _commit_cursor = _job_id; - _ready_gen.clear(); - std::size_t in_flight = 0; while (true) { _abort_check_function(); @@ -307,6 +324,7 @@ void EventGenerator::generate_deterministic() { 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, @@ -316,6 +334,7 @@ void EventGenerator::generate_deterministic() { } } 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, target); // Pushed as a single ready_job with vegas_batch_size set (rather // than batch_jobs separate ready_jobs) so start_jobs() dispatches @@ -325,7 +344,9 @@ void EventGenerator::generate_deterministic() { // (looking "done") as soon as whichever prefix got dispatched all // commit, while the rest of the batch is still sitting undispatched -- // reintroducing exactly the kind of scheduling-timing-dependent - // overshoot this batching was meant to eliminate. + // overshoot this batching was meant to eliminate. It also gives this + // channel's round a single contiguous job-id range, which is what + // lets register_dispatched_ids() find its start cheaply. _ready_jobs.push_back({ .channel_index = channel_index, .unweight = true, @@ -337,24 +358,45 @@ void EventGenerator::generate_deterministic() { std::size_t job_id_before = _job_id; start_jobs(); - in_flight += _job_id - job_id_before; - - if (in_flight > 0) { + std::size_t round_in_flight = _job_id - job_id_before; + register_dispatched_ids(job_id_before, _job_id); + + // Wait for this round -- every channel's currently dispatched batch -- to + // finish. Each job commits as soon as it's next in line *for its own + // channel* (_channel_commit_cursor), streamed immediately rather than + // buffered; channels no longer block each other's commits. Only once the + // whole round has committed (round_in_flight reaches 0) do we resync the + // cross-channel integral fractions, since that's the one piece of state + // that isn't safe to read mid-round -- see commit_generate_deterministic(). + while (round_in_flight > 0) { std::size_t job_id = _result_queue.wait(); - --in_flight; - _ready_gen.insert(job_id); - while (_ready_gen.erase(_commit_cursor) > 0) { - commit_generate_deterministic(_running_jobs.at(_commit_cursor)); - _running_jobs.erase(_commit_cursor); - ++_commit_cursor; - } - } else { - if (_status.done) { - unweight_all(); - } - if (_status.done) { - break; + --round_in_flight; + auto& job = _running_jobs.at(job_id); + auto& ready_gen = _channel_ready_gen.at(job.channel_index); + auto& cursor = _channel_commit_cursor.at(job.channel_index); + ready_gen.insert(job_id); + while (ready_gen.erase(cursor) > 0) { + commit_generate_deterministic(_running_jobs.at(cursor)); + _running_jobs.erase(cursor); + ++cursor; } + std::size_t job_id_refill = _job_id; + start_jobs(); + round_in_flight += _job_id - job_id_refill; + register_dispatched_ids(job_id_refill, _job_id); + } + + // Status (mean/error/counts, used for logging) is already fully up to date + // from the per-commit update_integral_status() calls above; only the + // fractions used for the *next* round's scheduling need to wait for the + // whole round to be committed. + update_integral_fractions(); + + if (_status.done) { + unweight_all(); + } + if (_status.done) { + break; } } print_gen_update(true); @@ -512,6 +554,11 @@ bool EventGenerator::start_jobs() { } 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; @@ -542,9 +589,17 @@ 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() { + // Reuses _status.mean rather than re-summing channel means, so this stays a + // pure function of whatever update_integral_status() last computed -- callers + // that want it to reflect only fully-committed-this-round data (see + // generate_deterministic()) just need to make sure update_integral_status() was + // last called from committed (not in-flight) state. for (auto [channel, integral_fraction] : zip(_channels, _channel_integral_fractions)) { - integral_fraction = channel->cross_section().mean() / total_mean; + integral_fraction = channel->cross_section().mean() / _status.mean; channel->set_target_count(integral_fraction * _config.target_count); } } From c39d755ebede3238dcad094d0f784d46d5c3af17 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Sat, 25 Jul 2026 00:59:08 +0200 Subject: [PATCH 08/19] reproducible lhe output --- madspace/src/driver/event_generator.cpp | 51 ++++++++++++++++--------- 1 file changed, 33 insertions(+), 18 deletions(-) diff --git a/madspace/src/driver/event_generator.cpp b/madspace/src/driver/event_generator.cpp index 420d6eec1..5e18726ab 100644 --- a/madspace/src/driver/event_generator.cpp +++ b/madspace/src/driver/event_generator.cpp @@ -1,6 +1,5 @@ #include "madspace/driver/event_generator.hpp" -#include #include #include #include @@ -727,13 +726,6 @@ void EventGenerator::combine_to_lhe( std::uint64_t seed = _seed; std::mt19937 select_rng = seeded_rng(seed, {kSaltCombineSelect}); ThreadPool pool(_config.combine_thread_count); - // Per-thread LHE-completion streams, keyed by a running thread index. Bit - // identical output requires combine_thread_pool_size == 1 (see seeded_rng()). - ThreadResource rand_gens( - pool, [seed, next_index = std::make_shared>(0)]() { - return seeded_rng(seed, {kSaltLheComplete, next_index->fetch_add(1)}); - } - ); auto [channel_data, particle_count, norm_factor] = init_combine(); std::vector> buffers; std::vector idle_buffers; @@ -755,45 +747,68 @@ void EventGenerator::combine_to_lhe( std::size_t event_count = 0; std::size_t last_update_count = 0; bool done = false; + // Each batch is seeded from its own submission-order sequence number rather + // than from whichever worker thread happens to run it, and a batch's formatted + // output is only written once every earlier-submitted batch has been written -- + // buffered out of submission order in ready_slots until its turn comes up -- + // so the LHE file's event order and content no longer depend on real + // completion timing. slot_batch_seq tracks which batch a (reused) buffer slot + // currently holds. + 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); + 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] { + std::mt19937 rand_gen = seeded_rng( + seed, {kSaltLheComplete, static_cast(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; } } From 2b0e631ba957e691ff939e4768afee966cc74306 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Sat, 25 Jul 2026 01:35:30 +0200 Subject: [PATCH 09/19] move unweighting step back into worker thread --- .../madspace/driver/channel_generator.hpp | 19 +- .../madspace/driver/event_generator.hpp | 55 +++- madspace/src/driver/channel_generator.cpp | 28 +- madspace/src/driver/event_generator.cpp | 253 ++++++++++++------ 4 files changed, 249 insertions(+), 106 deletions(-) diff --git a/madspace/include/madspace/driver/channel_generator.hpp b/madspace/include/madspace/driver/channel_generator.hpp index 721217bd0..8fb47abb5 100644 --- a/madspace/include/madspace/driver/channel_generator.hpp +++ b/madspace/include/madspace/driver/channel_generator.hpp @@ -68,12 +68,21 @@ class ChannelEventGenerator { bool is_survey, std::size_t survey_pass ); + // Snapshots the channel's current max weight into job.max_weight. Split out + // from submit_unweight_job() so the deterministic generation path can capture + // it at a fixed, job-id-ordered point (see EventGenerator::commit_generate_stage) + // while deferring the actual thread-pool submission to a later, priority-ordered + // dispatch pass (see EventGenerator::_context_unweight_queue). + void prepare_unweight_job(GeneratorBatchJob& job) const; + // Submits the unweighting work for job to the thread pool of the *same* + // context the job's generation ran on (important on GPU: unweighting reads + // device tensors and does a device-to-host copy, so it must run against the + // context that owns them). Assumes job.max_weight and job.rng_seed are already + // set (see prepare_unweight_job / start_job). + void submit_unweight_job(GeneratorBatchJob& job, ResultQueue& result_queue); + // prepare_unweight_job() + submit_unweight_job() in one call, for callers that + // don't need to separate capture from dispatch. void start_unweight_job(GeneratorBatchJob& job, ResultQueue& result_queue); - // Synchronous unweighting used by the deterministic generation path: runs the - // unweighter on the calling (harvest) thread against the current max weight so - // that it happens in a fixed, job-id-ordered position instead of as a separate - // asynchronously-scheduled job. - void unweight_job_inline(GeneratorBatchJob& job); std::size_t next_vegas_batch_size(); void clear_events(); void update_max_weight(Tensor weights); diff --git a/madspace/include/madspace/driver/event_generator.hpp b/madspace/include/madspace/driver/event_generator.hpp index 1e8754a2d..0f99ce648 100644 --- a/madspace/include/madspace/driver/event_generator.hpp +++ b/madspace/include/madspace/driver/event_generator.hpp @@ -100,6 +100,25 @@ class EventGenerator { std::vector> _channel_ready_gen; std::vector _channel_commit_cursor; std::vector _channel_cursor_set; + // generate_deterministic() only: second, independent per-channel cursor/buffer + // for the unweight stage's completions, analogous to _channel_ready_gen/ + // _channel_commit_cursor but for a job's *second* commit (see + // commit_unweight_stage()). A job's generate-stage and unweight-stage commits + // are tracked separately because they arrive as two distinct completion events + // (see ChannelEventGenerator::submit_unweight_job), but both cursors are reset + // together in register_dispatched_ids() since a job's unweight-stage id is the + // same as its generate-stage id. + std::vector> _channel_unweight_ready; + std::vector _channel_unweight_cursor; + // generate_deterministic() only: per-context queue of job ids awaiting + // unweight-stage dispatch. Populated by commit_generate_stage() (via + // ChannelEventGenerator::prepare_unweight_job(), which snapshots max_weight at + // that fixed, job-id-ordered point) and drained with priority by start_jobs(), + // ahead of any new regular batch dispatch -- otherwise, since a single + // _ready_jobs entry can now cover an entire channel round, unweight jobs (whose + // generation work is already done and just needs finishing, possibly via a + // costly GPU device-to-host copy) could get stuck behind it. + std::vector> _context_unweight_queue; ResultQueue _result_queue; // Base seed for reproducible event generation. 0 means "seed @@ -136,15 +155,31 @@ class EventGenerator { // thread count. void survey_deterministic(); void generate_deterministic(); - void commit_generate_deterministic(GeneratorBatchJob& job); + // Per-channel-ordered commit of a job's generate stage: cross-section + // accumulation, max-weight update, and (if job.unweight) capturing max_weight + // and queuing the job for prioritized unweight dispatch -- see + // _context_unweight_queue. Does not itself write events; that happens once the + // (separately ordered) unweight stage commits, in commit_unweight_stage(). + void commit_generate_stage(GeneratorBatchJob& job); + // Per-channel-ordered commit of a job's unweight stage: writes the unweighted + // events and updates counts. Called once job.unweighted_events is populated, + // strictly in job-id order per channel (see _channel_unweight_cursor), same as + // commit_generate_stage() but independently sequenced. + void commit_unweight_stage(GeneratorBatchJob& job); + // Shared tail of both commit steps: once a job's channel_job_count reaches + // zero (both stages committed), runs the batch-completion action + // (optimize_vegas / release _channel_optimizing, or release + // _channel_batch_pending) that used to be keyed off channel_job_count hitting a + // magic value mid-stream in the old single-stage commit. + void finish_channel_job(const GeneratorBatchJob& job); // generate_deterministic() only: scans the job ids newly added to // _running_jobs by the last start_jobs() call ([first_id, end_id)) and, for // any channel seen for the first time since its current round started (i.e. - // _channel_cursor_set is still false), initializes _channel_commit_cursor to - // that job's id. Safe because a channel's whole round is always dispatched - // atomically as one contiguous id range (see GeneratorBatchJob::vegas_batch_size), - // so the first id encountered for a channel in ascending order is always that - // round's minimum. + // _channel_cursor_set is still false), initializes _channel_commit_cursor and + // _channel_unweight_cursor to that job's id. Safe because a channel's whole + // round is always dispatched atomically as one contiguous id range (see + // GeneratorBatchJob::vegas_batch_size), so the first id encountered for a + // channel in ascending order is always that round's minimum. void register_dispatched_ids(std::size_t first_id, std::size_t end_id); // Size a new steady-state generation batch for a channel from fully-committed @@ -157,7 +192,13 @@ class EventGenerator { // keep all worker threads busy. std::size_t next_batch_job_count(std::size_t channel_index, double target) const; - bool start_jobs(); + // Dispatches from _context_unweight_queue (with priority, per context) and + // then _ready_jobs, up to each context's capacity. Returns the number of + // unweight jobs actually (re-)dispatched this call -- unlike a newly dispatched + // regular job, resubmitting a job for its unweight stage doesn't consume a new + // id, so callers that track in-flight completions via _job_id deltas (see + // generate_deterministic()) need this separately. + std::size_t start_jobs(); // update_integral() = update_integral_status() + update_integral_fractions(). // Split so generate_deterministic() can keep the status half (used for // logging/progress only) fully live, per commit, while deferring the fractions diff --git a/madspace/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index a038f3602..ccb27b529 100644 --- a/madspace/src/driver/channel_generator.cpp +++ b/madspace/src/driver/channel_generator.cpp @@ -534,10 +534,13 @@ 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]() { @@ -559,22 +562,11 @@ void ChannelEventGenerator::start_unweight_job( }); } -void ChannelEventGenerator::unweight_job_inline(GeneratorBatchJob& job) { - auto& runtimes = _runtimes.at(job.context_index); - auto& context = _contexts.at(job.context_index); - job.max_weight = _max_weight; - std::optional seed = job.rng_seed == 0 - ? std::nullopt - : std::optional(mix_seed(job.rng_seed, {kPhaseUnweight})); - std::vector unweighter_args( - job.events.begin(), job.events.begin() + _field_indices.random - ); - unweighter_args.push_back(Tensor(job.max_weight, context->device())); - TensorVec unw_events = runtimes.unweighter->run(unweighter_args, seed); - job.unweighted_events.clear(); - for (auto& item : unw_events) { - job.unweighted_events.push_back(item.cpu()); - } +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() { diff --git a/madspace/src/driver/event_generator.cpp b/madspace/src/driver/event_generator.cpp index 5e18726ab..c7c876185 100644 --- a/madspace/src/driver/event_generator.cpp +++ b/madspace/src/driver/event_generator.cpp @@ -56,6 +56,9 @@ EventGenerator::EventGenerator( _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 != 0), _status_file(status_file) {} @@ -257,42 +260,63 @@ void EventGenerator::generate() { print_gen_update(true); } -void EventGenerator::commit_generate_deterministic(GeneratorBatchJob& job) { - // Ordered, single-pass commit of one job: mirrors the generate() harvest body - // but folds unweighting in synchronously so there is exactly one commit per job. - // Called strictly in ascending job id *within job.channel_index* (see - // _channel_commit_cursor), so every state mutation local to that channel - // (cross-section accumulation, max-weight snapshot, event writes, counts, VEGAS - // optimisation) happens in a fixed order. update_integral_status() runs on every - // commit, same as before, so logging/progress reporting stays fully live; only - // update_integral_fractions() -- the piece that feeds back into per-channel - // scheduling decisions -- is deferred to once per round in +void EventGenerator::commit_generate_stage(GeneratorBatchJob& job) { + // Per-channel-ordered commit of a job's generate stage. Called strictly in + // ascending job id *within job.channel_index* (see _channel_commit_cursor), so + // cross-section accumulation and the max-weight update happen in a fixed order. + // update_integral_status() runs on every commit so logging/progress reporting + // stays fully live; update_integral_fractions() -- the piece that feeds back + // into per-channel scheduling decisions -- is deferred to once per round in // generate_deterministic(), since reading it here would make a channel's // scheduling depend on how far other channels happen to have gotten in real - // time. + // time. Does not write events itself: if job.unweight, that's deferred to + // commit_unweight_stage() once the (separately dispatched, possibly GPU-copy- + // bound) unweighting finishes -- see _context_unweight_queue. auto& channel = _channels.at(job.channel_index); auto& channel_job_count = _channel_job_counts.at(job.channel_index); - auto& context_job_count = _context_job_counts.at(job.context_index); if (job.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) { - channel->unweight_job_inline(job); - channel->write_events(job.unweighted_events, job.max_weight); - update_counts(); + // Captured here, at this job's fixed position in the per-channel commit + // order, so the max weight used for unweighting doesn't depend on when the + // separately-dispatched unweight computation actually runs -- see + // ChannelEventGenerator::prepare_unweight_job(). + channel->prepare_unweight_job(job); + _context_unweight_queue.at(job.context_index).push_back(job.job_id); } - channel_job_count -= 1 + job.unweight; - context_job_count -= 1; - if (job.is_vegas_batch && channel_job_count == 0) { - channel->optimize_vegas(job); + finish_channel_job(job); + print_gen_update(false); +} + +void EventGenerator::commit_unweight_stage(GeneratorBatchJob& job) { + // Per-channel-ordered commit of a job's unweight stage, independently sequenced + // from commit_generate_stage() -- see _channel_unweight_cursor. + auto& channel = _channels.at(job.channel_index); + channel->write_events(job.unweighted_events, job.max_weight); + update_counts(); + --_channel_job_counts.at(job.channel_index); + --_context_job_counts.at(job.context_index); + finish_channel_job(job); + print_gen_update(false); +} + +void EventGenerator::finish_channel_job(const GeneratorBatchJob& job) { + auto& channel_job_count = _channel_job_counts.at(job.channel_index); + if (channel_job_count != 0) { + return; + } + if (job.is_vegas_batch) { + _channels.at(job.channel_index)->optimize_vegas(job); _channel_optimizing.at(job.channel_index) = false; - } else if (!job.is_vegas_batch && channel_job_count == 0) { + } else { _channel_batch_pending.at(job.channel_index) = false; } - print_gen_update(false); } void EventGenerator::register_dispatched_ids(std::size_t first_id, std::size_t end_id) { @@ -300,6 +324,7 @@ void EventGenerator::register_dispatched_ids(std::size_t first_id, std::size_t e 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; } } @@ -356,32 +381,49 @@ void EventGenerator::generate_deterministic() { } std::size_t job_id_before = _job_id; - start_jobs(); - std::size_t round_in_flight = _job_id - job_id_before; + 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 this round -- every channel's currently dispatched batch -- to - // finish. Each job commits as soon as it's next in line *for its own - // channel* (_channel_commit_cursor), streamed immediately rather than - // buffered; channels no longer block each other's commits. Only once the - // whole round has committed (round_in_flight reaches 0) do we resync the - // cross-channel integral fractions, since that's the one piece of state - // that isn't safe to read mid-round -- see commit_generate_deterministic(). + // Wait for this round -- every channel's currently dispatched batch, through + // both its generate and unweight stages -- to finish. Each stage commits as + // soon as it's next in line *for its own channel and stage* + // (_channel_commit_cursor / _channel_unweight_cursor), streamed immediately + // rather than buffered; channels no longer block each other's commits, and a + // job's own unweight stage doesn't block the next job's generate stage from + // committing either. Only once the whole round has committed + // (round_in_flight reaches 0) do we resync the cross-channel integral + // fractions, since that's the one piece of state that isn't safe to read + // mid-round -- see commit_generate_stage(). while (round_in_flight > 0) { std::size_t job_id = _result_queue.wait(); --round_in_flight; auto& job = _running_jobs.at(job_id); - auto& ready_gen = _channel_ready_gen.at(job.channel_index); - auto& cursor = _channel_commit_cursor.at(job.channel_index); - ready_gen.insert(job_id); - while (ready_gen.erase(cursor) > 0) { - commit_generate_deterministic(_running_jobs.at(cursor)); - _running_jobs.erase(cursor); - ++cursor; + 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_stage(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_stage(_running_jobs.at(cursor)); + _running_jobs.erase(cursor); + ++cursor; + } } std::size_t job_id_refill = _job_id; - start_jobs(); - round_in_flight += _job_id - job_id_refill; + 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); } @@ -435,6 +477,12 @@ void EventGenerator::survey_deterministic() { continue; } std::size_t vegas_batch_size = channel->next_vegas_batch_size(); + // Reset so register_dispatched_ids() re-initializes + // _channel_unweight_cursor for this channel's new job below -- safe + // since, by construction of this loop, a channel never has a prior + // job still outstanding when it's given a new one (we don't start the + // next iteration until in_flight, tracking both stages, reaches 0). + _channel_cursor_set.at(i) = false; _ready_jobs.push_back({ .channel_index = i, .unweight = iter >= min_iters - 1, @@ -447,48 +495,84 @@ void EventGenerator::survey_deterministic() { ++i; } std::size_t job_id_before = _job_id; - start_jobs(); - std::size_t in_flight = _job_id - job_id_before; + 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; - _ready_gen.insert(job_id); - while (_ready_gen.erase(_commit_cursor) > 0) { - auto& job = _running_jobs.at(_commit_cursor); - auto& channel = _channels.at(job.channel_index); - auto& channel_job_count = _channel_job_counts.at(job.channel_index); - auto& context_job_count = _context_job_counts.at(job.context_index); - if (channel_job_count == job.split_job_count) { - channel->clear_events(); + auto& job = _running_jobs.at(job_id); + if (job.unweighted_events.size() == 0) { + // Generate-stage completion: still globally ordered (_ready_gen / + // _commit_cursor), unlike the unweight stage below -- this survey + // loop already barriers on the whole iteration (in_flight == 0) + // before starting the next one, so relaxing this ordering wouldn't + // buy the same throughput win it does in generate_deterministic(). + _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) { + // Captured here, at this job's fixed position in the + // (globally ordered) generate commit sequence -- see + // ChannelEventGenerator::prepare_unweight_job(). + 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; } - channel->integrate(job); - update_integral(); - channel->update_max_weight(job.weights); - if (job.unweight) { - channel->unweight_job_inline(job); - channel->write_events(job.unweighted_events, job.max_weight); + } else { + // Unweight-stage completion: ordered per channel only, independent + // of other channels and of the generate-stage commit order above -- + // same pattern as generate_deterministic()'s commit_unweight_stage(). + 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; } - channel_job_count -= 1 + job.unweight; - context_job_count -= 1; - if (!job.unweight) { - done = false; - } else if (channel_job_count == 0 && - channel->cross_section().rel_error() < target_precision) { - done = false; - } - if (channel_job_count == 0) { - channel->optimize_vegas(job); - done_event_count += job.vegas_batch_size; - } - _running_jobs.erase(_commit_cursor); - ++_commit_cursor; } std::size_t job_id_refill = _job_id; - start_jobs(); - in_flight += _job_id - job_id_refill; + 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); } } @@ -516,14 +600,34 @@ EventGenerator::next_batch_job_count(std::size_t channel_index, double target) c ); } -bool EventGenerator::start_jobs() { +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; + + // Unweight jobs get priority over starting new regular batches: their + // (possibly costly, e.g. GPU device-to-host copy) generation work is + // already done, so they should be finished promptly rather than get stuck + // behind a _ready_jobs backlog that can now cover an entire channel round. + // Empty for every caller except generate_deterministic(). + 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); @@ -543,13 +647,10 @@ bool EventGenerator::start_jobs() { ++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() { From df980369d337e6a14262d28d52d7880591ad4350 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Sat, 25 Jul 2026 13:25:30 +0200 Subject: [PATCH 10/19] desloppification and code quality improvements --- madspace/CMakeLists.txt | 2 + madspace/include/madspace/driver/backend.hpp | 6 +- .../madspace/driver/channel_generator.hpp | 29 +-- .../madspace/driver/event_generator.hpp | 117 ++-------- .../madspace/driver/generator_data.hpp | 28 +-- madspace/include/madspace/driver/random.hpp | 21 ++ madspace/include/madspace/util.hpp | 49 ----- madspace/src/cpu/device.hpp | 10 +- madspace/src/driver/channel_generator.cpp | 31 +-- madspace/src/driver/event_generator.cpp | 202 ++++++++---------- madspace/src/driver/random.cpp | 39 ++++ madspace/src/gpu/runtime.cu | 10 +- 12 files changed, 205 insertions(+), 339 deletions(-) create mode 100644 madspace/include/madspace/driver/random.hpp create mode 100644 madspace/src/driver/random.cpp diff --git a/madspace/CMakeLists.txt b/madspace/CMakeLists.txt index 342fd11b4..7701fd4ef 100644 --- a/madspace/CMakeLists.txt +++ b/madspace/CMakeLists.txt @@ -178,6 +178,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 @@ -230,6 +231,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 diff --git a/madspace/include/madspace/driver/backend.hpp b/madspace/include/madspace/driver/backend.hpp index 0df00d4df..91ff8e935 100644 --- a/madspace/include/madspace/driver/backend.hpp +++ b/madspace/include/madspace/driver/backend.hpp @@ -12,10 +12,8 @@ namespace madspace { class Runtime { public: virtual ~Runtime() = default; - // `seed`, when set, makes the random numbers drawn during this call a - // deterministic function of `seed` alone, independent of which thread executes - // it. Backends that cannot support this (e.g. GPU, or a CPU runtime executing - // concurrently) should reject a non-null seed rather than silently ignore it. + // `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, std::optional seed = std::nullopt) = 0; virtual std::tuple> run_with_grad( diff --git a/madspace/include/madspace/driver/channel_generator.hpp b/madspace/include/madspace/driver/channel_generator.hpp index 8fb47abb5..08aa6edd1 100644 --- a/madspace/include/madspace/driver/channel_generator.hpp +++ b/madspace/include/madspace/driver/channel_generator.hpp @@ -49,7 +49,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; } @@ -68,20 +70,13 @@ class ChannelEventGenerator { bool is_survey, std::size_t survey_pass ); - // Snapshots the channel's current max weight into job.max_weight. Split out - // from submit_unweight_job() so the deterministic generation path can capture - // it at a fixed, job-id-ordered point (see EventGenerator::commit_generate_stage) - // while deferring the actual thread-pool submission to a later, priority-ordered - // dispatch pass (see EventGenerator::_context_unweight_queue). + // 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 the unweighting work for job to the thread pool of the *same* - // context the job's generation ran on (important on GPU: unweighting reads - // device tensors and does a device-to-host copy, so it must run against the - // context that owns them). Assumes job.max_weight and job.rng_seed are already - // set (see prepare_unweight_job / start_job). + // 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() in one call, for callers that - // don't need to separate capture from dispatch. + // 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(); @@ -147,12 +142,8 @@ class ChannelEventGenerator { std::optional _histogram_function; RunningIntegral _cross_section; double _max_weight = 0.; - // Monotonic per-channel counters assigned to each scheduled job (main thread) to - // key its deterministic random stream; never reset. Survey and generate use - // independent counters so a generate job's seed never depends on how many - // survey jobs happened to run first, and vice versa; see - // EventGenerator::survey()'s survey_pass for why survey jobs additionally key - // off the calling survey pass. + // 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; std::size_t _unweighted_count = 0; diff --git a/madspace/include/madspace/driver/event_generator.hpp b/madspace/include/madspace/driver/event_generator.hpp index 0f99ce648..2c03e279d 100644 --- a/madspace/include/madspace/driver/event_generator.hpp +++ b/madspace/include/madspace/driver/event_generator.hpp @@ -34,12 +34,8 @@ class EventGenerator { const GeneratorConfig& config = default_config, std::uint64_t seed = 0 ); - // `survey_pass` distinguishes independent survey() calls that end up - // scheduling jobs on the same ChannelEventGenerator (e.g. an initial - // multichannel survey followed by a post-simplification re-survey of a channel - // that got carried over unchanged): each pass gets its own seed salt, so a - // pass's job seeds depend only on its own logical identity, never on how many - // jobs a previous, unrelated survey pass happened to schedule first. + // `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); @@ -77,64 +73,33 @@ class EventGenerator { std::vector _channel_optimizing; std::vector _channel_integral_fractions; std::vector _context_job_counts; - // True while a channel has a steady-state (non-VEGAS) generation batch - // dispatched but not yet fully committed. Mirrors _channel_optimizing's role - // for VEGAS batches: while true, no new batch is created for this channel, so - // the next batch's size (see next_batch_job_count()) is always computed from - // fully-committed data. Without this, jobs already dispatched but not yet - // committed are invisible to the "does this channel need more work" decision, - // so a burst of them can all commit after the channel's target was already met - // -- overshooting it by an amount that depends on real scheduling timing rather - // than on the seed. + // 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 analogue of _ready_gen/ - // _commit_cursor. Each channel's jobs still commit strictly in that channel's - // own ascending job-id order (out-of-order arrivals buffered in - // _channel_ready_gen until the cursor catches up), streamed as soon as they're - // next in line rather than deferred -- but channels no longer block each - // other's commits, so a channel with much costlier jobs (e.g. higher final- - // state multiplicity in an inclusive run) can't head-of-line-block the rest. - // _channel_cursor_set tracks whether _channel_commit_cursor has been - // initialized for the channel's current round yet; see - // register_dispatched_ids(). + // 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: second, independent per-channel cursor/buffer - // for the unweight stage's completions, analogous to _channel_ready_gen/ - // _channel_commit_cursor but for a job's *second* commit (see - // commit_unweight_stage()). A job's generate-stage and unweight-stage commits - // are tracked separately because they arrive as two distinct completion events - // (see ChannelEventGenerator::submit_unweight_job), but both cursors are reset - // together in register_dispatched_ids() since a job's unweight-stage id is the - // same as its generate-stage id. + // 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. Populated by commit_generate_stage() (via - // ChannelEventGenerator::prepare_unweight_job(), which snapshots max_weight at - // that fixed, job-id-ordered point) and drained with priority by start_jobs(), - // ahead of any new regular batch dispatch -- otherwise, since a single - // _ready_jobs entry can now cover an entire channel round, unweight jobs (whose - // generation work is already done and just needs finishing, possibly via a - // costly GPU device-to-host copy) could get stuck behind it. + // unweight-stage dispatch, drained with priority by start_jobs(). std::vector> _context_unweight_queue; ResultQueue _result_queue; - // Base seed for reproducible event generation. 0 means "seed - // non-deterministically" (the historical default); see seeded_rng(). + // Base seed for reproducible event generation; 0 means non-deterministic. std::uint64_t _seed; - // Job-scheduling context for the currently running survey()/generate() call, - // read by start_jobs() when it hands a job's seed derivation off to its - // ChannelEventGenerator. Set once at the top of survey()/generate(), before - // the scheduling loop starts. + // 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. _deterministic is set from _seed. Generate - // completions are buffered in _ready_gen and committed in ascending job id, with - // _commit_cursor pointing at the next job id to commit. + // 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; @@ -148,62 +113,14 @@ class EventGenerator { std::string _status_file; std::unordered_map _timing_data; - // Deterministic generation path (enabled when the context seed is non-zero): - // expensive matrix-element jobs still run on the pool, but their results are - // committed strictly in job-id order with unweighting folded in synchronously, - // so the whole computation is a deterministic function of the seed at a fixed - // thread count. void survey_deterministic(); void generate_deterministic(); - // Per-channel-ordered commit of a job's generate stage: cross-section - // accumulation, max-weight update, and (if job.unweight) capturing max_weight - // and queuing the job for prioritized unweight dispatch -- see - // _context_unweight_queue. Does not itself write events; that happens once the - // (separately ordered) unweight stage commits, in commit_unweight_stage(). - void commit_generate_stage(GeneratorBatchJob& job); - // Per-channel-ordered commit of a job's unweight stage: writes the unweighted - // events and updates counts. Called once job.unweighted_events is populated, - // strictly in job-id order per channel (see _channel_unweight_cursor), same as - // commit_generate_stage() but independently sequenced. - void commit_unweight_stage(GeneratorBatchJob& job); - // Shared tail of both commit steps: once a job's channel_job_count reaches - // zero (both stages committed), runs the batch-completion action - // (optimize_vegas / release _channel_optimizing, or release - // _channel_batch_pending) that used to be keyed off channel_job_count hitting a - // magic value mid-stream in the old single-stage commit. + void commit_generate_job(GeneratorBatchJob& job); + void commit_unweight_job(GeneratorBatchJob& job); void finish_channel_job(const GeneratorBatchJob& job); - // generate_deterministic() only: scans the job ids newly added to - // _running_jobs by the last start_jobs() call ([first_id, end_id)) and, for - // any channel seen for the first time since its current round started (i.e. - // _channel_cursor_set is still false), initializes _channel_commit_cursor and - // _channel_unweight_cursor to that job's id. Safe because a channel's whole - // round is always dispatched atomically as one contiguous id range (see - // GeneratorBatchJob::vegas_batch_size), so the first id encountered for a - // channel in ascending order is always that round's minimum. void register_dispatched_ids(std::size_t first_id, std::size_t end_id); - - // Size a new steady-state generation batch for a channel from fully-committed - // data only: the channel's own observed efficiency so far (count_unweighted / - // count_opt, or the pessimistic 1.0 -- i.e. a whole raw batch -- before any data - // exists yet) and its remaining need (target - count_unweighted). Capped to - // generation_batch_fraction of the remaining need so a single batch never bets - // everything on an estimate that may still be noisy, and floored at - // min_batch_jobs so a lone active channel still gets enough parallel jobs to - // keep all worker threads busy. - std::size_t next_batch_job_count(std::size_t channel_index, double target) const; - - // Dispatches from _context_unweight_queue (with priority, per context) and - // then _ready_jobs, up to each context's capacity. Returns the number of - // unweight jobs actually (re-)dispatched this call -- unlike a newly dispatched - // regular job, resubmitting a job for its unweight stage doesn't consume a new - // id, so callers that track in-flight completions via _job_id deltas (see - // generate_deterministic()) need this separately. + std::size_t next_batch_job_count(std::size_t channel_index) const; std::size_t start_jobs(); - // update_integral() = update_integral_status() + update_integral_fractions(). - // Split so generate_deterministic() can keep the status half (used for - // logging/progress only) fully live, per commit, while deferring the fractions - // half (which feeds back into per-channel scheduling decisions, see - // next_batch_job_count()) to once per round -- see commit_generate_deterministic(). void update_integral(); void update_integral_status(); void update_integral_fractions(); diff --git a/madspace/include/madspace/driver/generator_data.hpp b/madspace/include/madspace/driver/generator_data.hpp index 687e22dbc..d595917dc 100644 --- a/madspace/include/madspace/driver/generator_data.hpp +++ b/madspace/include/madspace/driver/generator_data.hpp @@ -63,15 +63,7 @@ struct GeneratorConfig { int combine_thread_count = -1; double cut_efficiency_threshold = 0.7; std::size_t max_cut_repetitions = 100; - // Steady-state event generation dispatches a channel's jobs in barrier-gated - // batches (a new batch is only sized/created once the previous one has fully - // committed), so a batch's size is a deterministic function of committed data, - // never of scheduling timing. generation_batch_fraction caps a batch to this - // fraction of the channel's currently remaining need, so a single batch never - // bets everything on an efficiency estimate that may still be noisy (e.g. for - // costly/high-variance channels early on); min_batch_jobs floors the batch size - // so a lone active channel still gets enough parallel jobs to keep all worker - // threads busy. See EventGenerator::next_batch_job_count(). + // 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; }; @@ -87,7 +79,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; @@ -104,11 +96,7 @@ struct Histogram { struct GeneratorBatchJob { std::size_t channel_index; bool unweight; - // Nonzero to make start_jobs() split this ready_job into ceil(vegas_batch_size / - // context_batch_size) sub-jobs and dispatch all of them atomically (bypassing - // the usual per-context dispatch cap), so channel_job_count reflects the whole - // batch from the moment it's dispatched rather than only whichever prefix fit - // under the cap -- see is_vegas_batch for what the batch is *for*. + // 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; @@ -120,14 +108,10 @@ struct GeneratorBatchJob { std::size_t context_index; std::size_t job_id; double max_weight; - // Base seed for this job's deterministic random stream, derived at scheduling - // time from (context seed, channel, per-channel job sequence). 0 means "use the - // non-deterministic pool generator" (context seed 0). See ChannelEventGenerator. + // Base seed for this job's deterministic random stream; 0 means non-deterministic. std::uint64_t rng_seed = 0; - // True for a VEGAS-grid-optimization batch (triggers clear_events()/ - // optimize_vegas() at the start/end of the batch); false for a steady-state - // generation batch (see EventGenerator::_channel_batch_pending), even though - // both use vegas_batch_size > 0 to get the same atomic whole-batch dispatch. + // 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; }; diff --git a/madspace/include/madspace/driver/random.hpp b/madspace/include/madspace/driver/random.hpp new file mode 100644 index 000000000..4faa589fa --- /dev/null +++ b/madspace/include/madspace/driver/random.hpp @@ -0,0 +1,21 @@ +#pragma once + +#include +#include +#include + +namespace madspace { + +// Mixes `seed` with `salts` (splitmix64-based) to derive an independent +// sub-seed, e.g. for keeping different random streams derived from the same +// base seed uncorrelated. `seed == 0` is a sentinel for "non-deterministic" +// and passes through unchanged. +std::uint64_t +mix_seed(std::uint64_t seed, std::initializer_list salts = {}); + +// Builds a seeded std::mt19937 from `seed` and `salts`. `seed == 0` seeds +// non-deterministically from std::random_device instead. +std::mt19937 +seeded_rng(std::uint64_t seed, std::initializer_list salts = {}); + +} // namespace madspace diff --git a/madspace/include/madspace/util.hpp b/madspace/include/madspace/util.hpp index 34b8614d7..de5d8c9b5 100644 --- a/madspace/include/madspace/util.hpp +++ b/madspace/include/madspace/util.hpp @@ -1,10 +1,8 @@ #pragma once -#include #include #include #include -#include #include #include #include @@ -12,53 +10,6 @@ namespace madspace { -// Deterministic RNG factory used everywhere randomness is consumed on the CPU. -// -// A ``seed`` of 0 preserves the historical behaviour: the generator is seeded -// non-deterministically from ``std::random_device``. Any non-zero ``seed`` makes -// the returned generator fully reproducible. The 64-bit seed together with the -// (small, low-entropy) integer ``salts`` is mixed through ``std::seed_seq`` so -// that distinct ``(seed, salts...)`` tuples yield well-separated, statistically -// independent streams -- salts are used to key a stream by its logical role -// (e.g. thread index, channel selection vs. unweighting vs. LHE completion). -// Fold a base seed together with integer salts into a single 64-bit value using -// splitmix64. Used to derive a stable per-job seed from (context seed, channel, -// job sequence, phase) so that a job's random stream depends only on its logical -// identity, never on which thread runs it. mix_seed(0, ...) stays 0 so callers can -// keep using 0 as the "non-deterministic" sentinel. -inline std::uint64_t -mix_seed(std::uint64_t seed, std::initializer_list salts = {}) { - if (seed == 0) { - return 0; - } - auto splitmix = [](std::uint64_t x) { - x += 0x9E3779B97F4A7C15ull; - x = (x ^ (x >> 30)) * 0xBF58476D1CE4E5B9ull; - x = (x ^ (x >> 27)) * 0x94D049BB133111EBull; - return x ^ (x >> 31); - }; - std::uint64_t h = splitmix(seed); - for (auto s : salts) { - h = splitmix(h ^ splitmix(s)); - } - return h ? h : 1; -} - -inline std::mt19937 -seeded_rng(std::uint64_t seed, std::initializer_list salts = {}) { - if (seed == 0) { - std::random_device rand_device; - return std::mt19937(rand_device()); - } - std::vector data; - data.reserve(2 + salts.size()); - data.push_back(static_cast(seed & 0xFFFFFFFFu)); - data.push_back(static_cast(seed >> 32)); - data.insert(data.end(), salts.begin(), salts.end()); - std::seed_seq seq(data.begin(), data.end()); - return std::mt19937(seq); -} - template struct Overloaded : Ts... { using Ts::operator()...; diff --git a/madspace/src/cpu/device.hpp b/madspace/src/cpu/device.hpp index 2a0c08f47..26c1a712e 100644 --- a/madspace/src/cpu/device.hpp +++ b/madspace/src/cpu/device.hpp @@ -3,9 +3,9 @@ #include #include +#include "madspace/driver/random.hpp" #include "madspace/driver/tensor.hpp" #include "madspace/driver/thread_pool.hpp" -#include "madspace/util.hpp" #include "simd.hpp" namespace madspace { @@ -17,11 +17,9 @@ class CpuDevice : public Device { public: static constexpr bool is_concurrent = false; - // Returns the generator randomness consumed by this device's op_random / - // op_random_int / op_unweight should draw from. The base implementation (used - // for CpuDevice and AsyncCpuDevice, i.e. whenever no per-call seed is in play) - // falls back to the runtime's non-deterministic per-thread stream; SeqCpuDevice - // hides this with its own seeded generator. + // RNG for op_random/op_random_int/op_unweight. Base implementation falls back + // to the runtime's non-deterministic per-thread stream; SeqCpuDevice overrides + // it with its own seeded generator. std::mt19937& rand_gen(CpuRuntime& runtime) const; std::pair allocate(std::size_t size, AllocHint hint) const override { diff --git a/madspace/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index ccb27b529..3193d0fc6 100644 --- a/madspace/src/driver/channel_generator.cpp +++ b/madspace/src/driver/channel_generator.cpp @@ -1,4 +1,5 @@ #include "madspace/driver/channel_generator.hpp" +#include "madspace/driver/random.hpp" #include "madspace/util.hpp" using namespace madspace; @@ -64,7 +65,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 }, @@ -188,7 +189,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 }, @@ -398,6 +399,9 @@ 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, @@ -405,19 +409,6 @@ void ChannelEventGenerator::start_job( bool is_survey, std::size_t survey_pass ) { - // Assign the job's deterministic base seed on the (single) scheduling thread so - // it depends only on the job's logical identity, not on which worker runs it -- - // deliberately not on job.context_index either, so the same logical job gets the - // same seed no matter how many contexts/threads are available to run it or which - // one it lands on (context assignment is a load-balancing detail, see - // start_jobs()). This is safe because _survey_rng_seq/_generate_rng_seq are - // incremented exactly once per logical job, on the single scheduling thread, - // before dispatch -- so a given (channel, seq) pair is never reused regardless - // of context assignment. Survey and generate draw from independent counters so - // a job's seed never depends on job counts from the other kind of call; survey - // jobs additionally key off survey_pass so that a channel re-surveyed across - // multiple independent survey() calls doesn't need its job count inferred from - // call history. if (seed == 0) { job.rng_seed = 0; } else if (is_survey) { @@ -433,10 +424,7 @@ void ChannelEventGenerator::start_job( .submit([this, &job, &result_queue]() { auto& runtimes = _runtimes.at(job.context_index); auto& context = _contexts.at(job.context_index); - // Every runtime->run() call below that can consume randomness gets its - // own sub-seed, uniquely salted by a running call index, derived from - // the job's base seed. This keeps each call fully deterministic without - // sharing mutable RNG state across independent Runtime objects. + // Each run() call below gets its own sub-seed derived from job's base seed. std::uint64_t rng_call_index = 0; auto next_seed = [&]() -> std::optional { if (job.rng_seed == 0) { @@ -614,8 +602,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) { diff --git a/madspace/src/driver/event_generator.cpp b/madspace/src/driver/event_generator.cpp index c7c876185..48110e983 100644 --- a/madspace/src/driver/event_generator.cpp +++ b/madspace/src/driver/event_generator.cpp @@ -1,23 +1,26 @@ #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; namespace { -// Salts keeping the driver-level random streams independent of each other for a -// given context seed (see seeded_rng()). Parallel stages additionally mix in a -// per-thread index. -constexpr std::uint32_t kSaltCombineSelect = 1; // channel selection in read_and_combine -constexpr std::uint32_t kSaltLheComplete = 2; // colour/helicity in fill_lhe_event -constexpr std::uint32_t kSaltUnweight = 3; // final unweighting + +// salts to derive seeds for different generation phases +constexpr std::uint32_t salt_combine_select = 1; +constexpr std::uint32_t salt_lhe_complete = 2; +constexpr std::uint32_t salt_unweight = 3; + } // namespace const GeneratorConfig EventGenerator::default_config = {}; @@ -41,7 +44,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 }, @@ -260,18 +263,9 @@ void EventGenerator::generate() { print_gen_update(true); } -void EventGenerator::commit_generate_stage(GeneratorBatchJob& job) { - // Per-channel-ordered commit of a job's generate stage. Called strictly in - // ascending job id *within job.channel_index* (see _channel_commit_cursor), so - // cross-section accumulation and the max-weight update happen in a fixed order. - // update_integral_status() runs on every commit so logging/progress reporting - // stays fully live; update_integral_fractions() -- the piece that feeds back - // into per-channel scheduling decisions -- is deferred to once per round in - // generate_deterministic(), since reading it here would make a channel's - // scheduling depend on how far other channels happen to have gotten in real - // time. Does not write events itself: if job.unweight, that's deferred to - // commit_unweight_stage() once the (separately dispatched, possibly GPU-copy- - // bound) unweighting finishes -- see _context_unweight_queue. +// 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) { @@ -283,10 +277,7 @@ void EventGenerator::commit_generate_stage(GeneratorBatchJob& job) { --channel_job_count; --_context_job_counts.at(job.context_index); if (job.unweight) { - // Captured here, at this job's fixed position in the per-channel commit - // order, so the max weight used for unweighting doesn't depend on when the - // separately-dispatched unweight computation actually runs -- see - // ChannelEventGenerator::prepare_unweight_job(). + // 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); } @@ -294,9 +285,7 @@ void EventGenerator::commit_generate_stage(GeneratorBatchJob& job) { print_gen_update(false); } -void EventGenerator::commit_unweight_stage(GeneratorBatchJob& job) { - // Per-channel-ordered commit of a job's unweight stage, independently sequenced - // from commit_generate_stage() -- see _channel_unweight_cursor. +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(); @@ -330,6 +319,8 @@ void EventGenerator::register_dispatched_ids(std::size_t first_id, std::size_t e } } +// 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(); @@ -341,8 +332,8 @@ void EventGenerator::generate_deterministic() { ++channel_index) { auto& channel = _channels.at(channel_index); double integral_frac = _channel_integral_fractions.at(channel_index); - double target = integral_frac * _config.target_count; - if (integral_frac > 0 && channel->status().count_unweighted >= target) { + if (integral_frac > 0 && + channel->status().count_unweighted >= channel->status().count_target) { continue; } if (channel->needs_optimization()) { @@ -359,18 +350,10 @@ void EventGenerator::generate_deterministic() { } 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, target); - // Pushed as a single ready_job with vegas_batch_size set (rather - // than batch_jobs separate ready_jobs) so start_jobs() dispatches - // the whole batch atomically -- see - // GeneratorBatchJob::vegas_batch_size. Otherwise, once batch_jobs - // exceeds the per-context dispatch cap, channel_job_count could reach 0 - // (looking "done") as soon as whichever prefix got dispatched all - // commit, while the rest of the batch is still sitting undispatched -- - // reintroducing exactly the kind of scheduling-timing-dependent - // overshoot this batching was meant to eliminate. It also gives this - // channel's round a single contiguous job-id range, which is what - // lets register_dispatched_ids() find its start cheaply. + 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, @@ -385,16 +368,7 @@ void EventGenerator::generate_deterministic() { std::size_t round_in_flight = (_job_id - job_id_before) + unweight_dispatched; register_dispatched_ids(job_id_before, _job_id); - // Wait for this round -- every channel's currently dispatched batch, through - // both its generate and unweight stages -- to finish. Each stage commits as - // soon as it's next in line *for its own channel and stage* - // (_channel_commit_cursor / _channel_unweight_cursor), streamed immediately - // rather than buffered; channels no longer block each other's commits, and a - // job's own unweight stage doesn't block the next job's generate stage from - // committing either. Only once the whole round has committed - // (round_in_flight reaches 0) do we resync the cross-channel integral - // fractions, since that's the one piece of state that isn't safe to read - // mid-round -- see commit_generate_stage(). + // 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; @@ -405,7 +379,7 @@ void EventGenerator::generate_deterministic() { ready.insert(job_id); while (ready.erase(cursor) > 0) { auto& gen_job = _running_jobs.at(cursor); - commit_generate_stage(gen_job); + commit_generate_job(gen_job); if (!gen_job.unweight) { _running_jobs.erase(cursor); } @@ -416,7 +390,7 @@ void EventGenerator::generate_deterministic() { auto& cursor = _channel_unweight_cursor.at(job.channel_index); ready.insert(job_id); while (ready.erase(cursor) > 0) { - commit_unweight_stage(_running_jobs.at(cursor)); + commit_unweight_job(_running_jobs.at(cursor)); _running_jobs.erase(cursor); ++cursor; } @@ -427,10 +401,6 @@ void EventGenerator::generate_deterministic() { register_dispatched_ids(job_id_refill, _job_id); } - // Status (mean/error/counts, used for logging) is already fully up to date - // from the per-commit update_integral_status() calls above; only the - // fractions used for the *next* round's scheduling need to wait for the - // whole round to be committed. update_integral_fractions(); if (_status.done) { @@ -443,6 +413,7 @@ void EventGenerator::generate_deterministic() { 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; @@ -477,11 +448,7 @@ void EventGenerator::survey_deterministic() { continue; } std::size_t vegas_batch_size = channel->next_vegas_batch_size(); - // Reset so register_dispatched_ids() re-initializes - // _channel_unweight_cursor for this channel's new job below -- safe - // since, by construction of this loop, a channel never has a prior - // job still outstanding when it's given a new one (we don't start the - // next iteration until in_flight, tracking both stages, reaches 0). + // 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, @@ -505,11 +472,8 @@ void EventGenerator::survey_deterministic() { --in_flight; auto& job = _running_jobs.at(job_id); if (job.unweighted_events.size() == 0) { - // Generate-stage completion: still globally ordered (_ready_gen / - // _commit_cursor), unlike the unweight stage below -- this survey - // loop already barriers on the whole iteration (in_flight == 0) - // before starting the next one, so relaxing this ordering wouldn't - // buy the same throughput win it does in generate_deterministic(). + // 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); @@ -525,9 +489,7 @@ void EventGenerator::survey_deterministic() { --channel_job_count; --_context_job_counts.at(gen_job.context_index); if (gen_job.unweight) { - // Captured here, at this job's fixed position in the - // (globally ordered) generate commit sequence -- see - // ChannelEventGenerator::prepare_unweight_job(). + // 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); @@ -542,9 +504,7 @@ void EventGenerator::survey_deterministic() { ++_commit_cursor; } } else { - // Unweight-stage completion: ordered per channel only, independent - // of other channels and of the generate-stage commit order above -- - // same pattern as generate_deterministic()'s commit_unweight_stage(). + // 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); @@ -579,21 +539,20 @@ void EventGenerator::survey_deterministic() { print_survey_update(true, done_event_count, total_event_count, iter - 1); } -std::size_t -EventGenerator::next_batch_job_count(std::size_t channel_index, double target) const { +// 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(); - // Before any data has committed for this channel, assume the pessimistic 1.0 - // (i.e. a whole raw batch becomes unweighted events), so the very first batch - // under- rather than over-estimates how many jobs are needed. + // 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 the assumed per-job contribution at 1 event: guards against a - // near-zero observed efficiency (e.g. an unlucky small first sample) blowing - // this up into an unbounded batch size. + // 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(target - status.count_unweighted, 0.); + 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)) @@ -610,11 +569,7 @@ std::size_t EventGenerator::start_jobs() { ? _config.cpu_batch_size : _config.gpu_batch_size; - // Unweight jobs get priority over starting new regular batches: their - // (possibly costly, e.g. GPU device-to-host copy) generation work is - // already done, so they should be finished promptly rather than get stuck - // behind a _ready_jobs backlog that can now cover an entire channel round. - // Empty for every caller except generate_deterministic(). + // 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(); @@ -692,24 +647,45 @@ void EventGenerator::update_integral_status() { } void EventGenerator::update_integral_fractions() { - // Reuses _status.mean rather than re-summing channel means, so this stays a - // pure function of whatever update_integral_status() last computed -- callers - // that want it to reflect only fully-committed-this-round data (see - // generate_deterministic()) just need to make sure update_integral_status() was - // last called from committed (not in-flight) state. for (auto [channel, integral_fraction] : zip(_channels, _channel_integral_fractions)) { integral_fraction = channel->cross_section().mean() / _status.mean; - channel->set_target_count(integral_fraction * _config.target_count); + } + + // 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; @@ -723,7 +699,7 @@ void EventGenerator::update_counts() { void EventGenerator::combine_to_compact_npy(const std::string& file_name) { reset_start_time(); - std::mt19937 select_rng = seeded_rng(_seed, {kSaltCombineSelect}); + std::mt19937 select_rng = seeded_rng(_seed, {salt_combine_select}); auto [channel_data, particle_count, norm_factor] = init_combine(); DataLayout layout( EventRecord::layout( @@ -761,8 +737,8 @@ void EventGenerator::combine_to_lhe_npy( ) { reset_start_time(); std::uint64_t seed = _seed; - std::mt19937 select_rng = seeded_rng(seed, {kSaltCombineSelect}); - std::mt19937 rand_gen = seeded_rng(seed, {kSaltLheComplete}); + std::mt19937 select_rng = seeded_rng(seed, {salt_combine_select}); + std::mt19937 rand_gen = seeded_rng(seed, {salt_lhe_complete}); auto [channel_data, particle_count, norm_factor] = init_combine(); DataLayout in_layout( EventRecord::layout( @@ -825,7 +801,7 @@ void EventGenerator::combine_to_lhe( ) { reset_start_time(); std::uint64_t seed = _seed; - std::mt19937 select_rng = seeded_rng(seed, {kSaltCombineSelect}); + std::mt19937 select_rng = seeded_rng(seed, {salt_combine_select}); ThreadPool pool(_config.combine_thread_count); auto [channel_data, particle_count, norm_factor] = init_combine(); std::vector> buffers; @@ -848,13 +824,8 @@ void EventGenerator::combine_to_lhe( std::size_t event_count = 0; std::size_t last_update_count = 0; bool done = false; - // Each batch is seeded from its own submission-order sequence number rather - // than from whichever worker thread happens to run it, and a batch's formatted - // output is only written once every earlier-submitted batch has been written -- - // buffered out of submission order in ready_slots until its turn comes up -- - // so the LHE file's event order and content no longer depend on real - // completion timing. slot_batch_seq tracks which batch a (reused) buffer slot - // currently holds. + // 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()); @@ -876,7 +847,7 @@ void EventGenerator::combine_to_lhe( pool.submit( [slot, batch_seq, seed, this, &in_buffer, &out_buffer, &lhe_completer] { std::mt19937 rand_gen = seeded_rng( - seed, {kSaltLheComplete, static_cast(batch_seq)} + seed, {salt_lhe_complete, static_cast(batch_seq)} ); LHEEvent lhe_event; out_buffer.clear(); @@ -933,14 +904,13 @@ void EventGenerator::add_timing_data(const std::string& key) { } void EventGenerator::unweight_all() { - std::mt19937 rand_gen = seeded_rng(_seed, {kSaltUnweight}); + std::mt19937 rand_gen = seeded_rng(_seed, {salt_unweight}); bool done = true; double total_eff_count = 0.; - for (auto [channel, integral_fraction] : - zip(_channels, _channel_integral_fractions)) { + for (auto& channel : _channels) { channel->unweight_file(rand_gen); - 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; @@ -998,11 +968,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); @@ -1170,8 +1141,7 @@ void EventGenerator::write_status(const std::string& status, bool force_write) { {"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/random.cpp b/madspace/src/driver/random.cpp new file mode 100644 index 000000000..41595e29b --- /dev/null +++ b/madspace/src/driver/random.cpp @@ -0,0 +1,39 @@ +#include "madspace/driver/random.hpp" + +#include + +namespace madspace { + +std::uint64_t mix_seed(std::uint64_t seed, std::initializer_list salts) { + if (seed == 0) { + return 0; + } + auto splitmix = [](std::uint64_t x) { + x += 0x9E3779B97F4A7C15ull; + x = (x ^ (x >> 30)) * 0xBF58476D1CE4E5B9ull; + x = (x ^ (x >> 27)) * 0x94D049BB133111EBull; + return x ^ (x >> 31); + }; + std::uint64_t h = splitmix(seed); + for (auto s : salts) { + h = splitmix(h ^ splitmix(s)); + } + return h ? h : 1; +} + +std::mt19937 +seeded_rng(std::uint64_t seed, std::initializer_list salts) { + if (seed == 0) { + std::random_device rand_device; + return std::mt19937(rand_device()); + } + std::vector data; + data.reserve(2 + salts.size()); + data.push_back(static_cast(seed & 0xFFFFFFFFu)); + data.push_back(static_cast(seed >> 32)); + data.insert(data.end(), salts.begin(), salts.end()); + std::seed_seq seq(data.begin(), data.end()); + return std::mt19937(seq); +} + +} // namespace madspace diff --git a/madspace/src/gpu/runtime.cu b/madspace/src/gpu/runtime.cu index 2096371d8..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, std::optional) { +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); @@ -1652,12 +1655,15 @@ TensorVec GpuRuntime::run(const TensorVec& inputs, std::optional) std::tuple> GpuRuntime::run_with_grad( const TensorVec& inputs, const std::vector& input_requires_grad, - std::optional + 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()); From a38059bb0a34abd3468aedd38bb410dc8a53c44d Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Sat, 25 Jul 2026 14:00:20 +0200 Subject: [PATCH 11/19] add reproducibility test --- .github/workflows/acceptancetest_mg7.yml | 42 +++ madspace/tests/test_determinism.py | 84 ------ .../test_mg7_reproducibility.py | 268 ++++++++++++++++++ 3 files changed, 310 insertions(+), 84 deletions(-) delete mode 100644 madspace/tests/test_determinism.py create mode 100644 tests/acceptance_tests/test_mg7_reproducibility.py diff --git a/.github/workflows/acceptancetest_mg7.yml b/.github/workflows/acceptancetest_mg7.yml index bc334e43b..29a0e76ef 100644 --- a/.github/workflows/acceptancetest_mg7.yml +++ b/.github/workflows/acceptancetest_mg7.yml @@ -201,6 +201,48 @@ 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_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_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/madspace/tests/test_determinism.py b/madspace/tests/test_determinism.py deleted file mode 100644 index 582f629cf..000000000 --- a/madspace/tests/test_determinism.py +++ /dev/null @@ -1,84 +0,0 @@ -"""Regression tests for seeded random number generation. - -These guard the seeding plumbing that reproducible event generation rests on: a -non-zero ``Context`` seed must make every random draw a deterministic function of -that seed, distinct seeds must produce independent streams, and a seed of 0 (the -run_card default) must keep the historical non-deterministic behaviour. - -Threads: the raw ``FunctionRuntime`` path splits a batch over the thread pool and -takes each chunk's stream from whichever worker runs it, so it is only guaranteed -reproducible with a single-threaded pool. Reproducibility at higher thread counts -during *event generation* comes from the generator seeding each job by its logical -identity instead of by thread; that path needs a compiled matrix element and so is -not covered here. -""" - -import numpy as np -import pytest - -import madspace as ms - -DIM = 4 -BATCH = 20000 - - -def random_function(dim=DIM): - """Build a function whose only output is ``dim`` uniform randoms per event.""" - input_types = ms.NamedTypes([("batch_size", ms.Type([ms.batch_size]))]) - output_types = ms.NamedTypes([("r", ms.batch_float_array(dim))]) - fb = ms.FunctionBuilder(input_types, output_types) - fb.output(0, fb.random(fb.input(0), dim)) - return fb.function() - - -def draw(seed, batch_size=BATCH, thread_count=1): - """Draw one batch of randoms from a fresh context with the given seed.""" - context = ms.Context(device=ms.cpu_device(), thread_count=thread_count, seed=seed) - runtime = ms.FunctionRuntime(random_function(), context) - return np.array(runtime(np.array([batch_size], dtype=np.int64)), copy=True) - - -def test_context_exposes_seed(): - assert ms.Context(thread_count=1).seed == 0 - assert ms.Context(thread_count=1, seed=12345).seed == 12345 - - -def test_same_seed_is_reproducible(): - assert np.array_equal(draw(12345), draw(12345)) - - -def test_same_seed_is_reproducible_across_repeated_runs(): - first = draw(7) - for _ in range(3): - assert np.array_equal(first, draw(7)) - - -def test_different_seeds_differ(): - assert not np.array_equal(draw(12345), draw(12346)) - - -def test_seed_zero_is_non_deterministic(): - # 0 is the run_card default and must keep drawing a fresh stream each run - assert not np.array_equal(draw(0), draw(0)) - - -def test_draws_are_uniform(): - r = draw(2024) - assert r.shape == (BATCH, DIM) - assert np.all((r >= 0) & (r < 1)) - assert r.mean() == pytest.approx(0.5, abs=0.02) - - -def test_different_seeds_are_independent(): - # correlation of two independent streams is ~1/sqrt(n) == 0.004 here, so a - # 0.05 bound is far outside the noise while leaving no room for a shared stream - a, b = draw(1).ravel(), draw(2).ravel() - assert abs(np.corrcoef(a, b)[0, 1]) < 0.05 - - -def test_nearby_seeds_are_independent(): - # adjacent integer seeds must not give correlated streams: seeds go through - # std::seed_seq precisely so that seed=1 and seed=2 are not near each other - a, b = draw(1).ravel(), draw(2).ravel() - assert not np.array_equal(a, b) - assert abs(np.corrcoef(a, b)[0, 1]) < 0.05 diff --git a/tests/acceptance_tests/test_mg7_reproducibility.py b/tests/acceptance_tests/test_mg7_reproducibility.py new file mode 100644 index 000000000..317613a23 --- /dev/null +++ b/tests/acceptance_tests/test_mg7_reproducibility.py @@ -0,0 +1,268 @@ +################################################################################ +# +# 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 the two ways an mg7 +run produces events: + + * ``test_vegas_reproducibility_mg7`` -- a plain survey (VEGAS-optimized) + + generate run, through ``bin/generate_events``. + * ``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. + +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_.*_reproducibility_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_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') + + +if __name__ == '__main__': + unittest.main() From f874000745cd91a4044996e35a8b2d6798342622 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Sat, 25 Jul 2026 15:52:34 +0200 Subject: [PATCH 12/19] bugfix in the final unweighting pass if called repeatedly --- .../include/madspace/driver/channel_generator.hpp | 6 ++++++ madspace/include/madspace/driver/event_generator.hpp | 8 +++++++- madspace/src/driver/channel_generator.cpp | 6 +++++- madspace/src/driver/event_generator.cpp | 11 +++++++++-- 4 files changed, 27 insertions(+), 4 deletions(-) diff --git a/madspace/include/madspace/driver/channel_generator.hpp b/madspace/include/madspace/driver/channel_generator.hpp index 08aa6edd1..48e2c4d71 100644 --- a/madspace/include/madspace/driver/channel_generator.hpp +++ b/madspace/include/madspace/driver/channel_generator.hpp @@ -146,7 +146,13 @@ class ChannelEventGenerator { // 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/event_generator.hpp b/madspace/include/madspace/driver/event_generator.hpp index e362a813e..0933fbdf1 100644 --- a/madspace/include/madspace/driver/event_generator.hpp +++ b/madspace/include/madspace/driver/event_generator.hpp @@ -42,7 +42,8 @@ class EventGenerator { 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; } @@ -97,6 +98,11 @@ class EventGenerator { // Base seed for reproducible event generation; 0 means non-deterministic. std::uint64_t _seed; + // unweight_all() may run more than once per generate() (a channel's target can + // grow after it looked done, un-finishing it and triggering another round). + // Salted by this counter so repeated calls don't replay the same stream. + std::size_t _unweight_call_index = 0; + // Scheduling context for the running survey()/generate() call, read by // start_jobs() to derive job seeds. bool _survey_job = false; diff --git a/madspace/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index 3193d0fc6..692025c4c 100644 --- a/madspace/src/driver/channel_generator.cpp +++ b/madspace/src/driver/channel_generator.cpp @@ -291,7 +291,7 @@ void ChannelEventGenerator::unweight_file(std::mt19937& 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); @@ -308,6 +308,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; } @@ -567,6 +569,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(); @@ -617,6 +620,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 66de388b1..8ad68a964 100644 --- a/madspace/src/driver/event_generator.cpp +++ b/madspace/src/driver/event_generator.cpp @@ -904,11 +904,18 @@ void EventGenerator::add_timing_data(const std::string& key) { } void EventGenerator::unweight_all() { - std::mt19937 rand_gen = seeded_rng(_seed, {salt_unweight}); + std::size_t call_index = _unweight_call_index++; bool done = true; double total_eff_count = 0.; - for (auto& channel : _channels) { + for (std::size_t channel_index = 0; auto& channel : _channels) { + std::mt19937 rand_gen = seeded_rng( + _seed, + {salt_unweight, + static_cast(call_index), + static_cast(channel_index)} + ); channel->unweight_file(rand_gen); + ++channel_index; std::size_t chan_target = channel->status().count_target; if (channel->status().count_unweighted < chan_target) { From 77be66c36cd5cf2908d1c67483b6de9c249fc363 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Sat, 25 Jul 2026 18:41:48 +0200 Subject: [PATCH 13/19] openblas thread safety bugfix --- madspace/CMakeLists.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/madspace/CMakeLists.txt b/madspace/CMakeLists.txt index 7701fd4ef..4b3eb2390 100644 --- a/madspace/CMakeLists.txt +++ b/madspace/CMakeLists.txt @@ -136,6 +136,7 @@ if(ENABLE_OPENBLAS) -DBUILD_TESTING=OFF -DBUILD_WITHOUT_CBLAS=ON -DUSE_THREAD=OFF + -DUSE_LOCKING=ON -DCMAKE_POSITION_INDEPENDENT_CODE=1 INSTALL_BYPRODUCTS /lib/${openblas_libname} ) From 24a04532c1aa717fb889974fb2855d2c67e2cc43 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Sat, 25 Jul 2026 19:19:23 +0200 Subject: [PATCH 14/19] build openblas in dynamic_arch mode --- madspace/CMakeLists.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/madspace/CMakeLists.txt b/madspace/CMakeLists.txt index 4b3eb2390..5d87bd1be 100644 --- a/madspace/CMakeLists.txt +++ b/madspace/CMakeLists.txt @@ -137,6 +137,7 @@ if(ENABLE_OPENBLAS) -DBUILD_WITHOUT_CBLAS=ON -DUSE_THREAD=OFF -DUSE_LOCKING=ON + -DDYNAMIC_ARCH=ON -DCMAKE_POSITION_INDEPENDENT_CODE=1 INSTALL_BYPRODUCTS /lib/${openblas_libname} ) From 63ca5300d8313479079abe1eb16659de864613e2 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Sat, 25 Jul 2026 21:14:30 +0200 Subject: [PATCH 15/19] store seed in info.json file --- madspace/src/driver/event_generator.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/madspace/src/driver/event_generator.cpp b/madspace/src/driver/event_generator.cpp index 8ad68a964..e422cd395 100644 --- a/madspace/src/driver/event_generator.cpp +++ b/madspace/src/driver/event_generator.cpp @@ -1142,6 +1142,7 @@ void EventGenerator::write_status(const std::string& status, bool force_write) { std::ofstream f(status_tmp_file); nlohmann::json j{ {"status", status}, + {"seed", _seed}, {"process", _status}, {"channels", channel_status()}, {"run_times", _timing_data}, From 92e02435594a6fc177e97f62890ed179b90a62d6 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Sat, 25 Jul 2026 22:27:01 +0200 Subject: [PATCH 16/19] fix concurrent writes and non-deterministic ordering in cpu backward pass --- madspace/src/cpu/runtime.cpp | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/madspace/src/cpu/runtime.cpp b/madspace/src/cpu/runtime.cpp index 0366d804c..e83cca19f 100644 --- a/madspace/src/cpu/runtime.cpp +++ b/madspace/src/cpu/runtime.cpp @@ -914,9 +914,8 @@ CpuRuntime::CpuRuntime(const Function& function, ContextPtr context, bool concur --instr_index; bool is_ready = true; std::size_t dependency_count = 0; - for (auto& out : instr.outputs) { - local_uses_backward.at(out.local_index).push_back(instr_index); - std::size_t source_instr = local_sources_backward.at(out.local_index); + auto check_dependency = [&](std::size_t local_index) { + std::size_t source_instr = local_sources_backward.at(local_index); if (source_instr != -1) { is_ready = false; auto& source_deps = dep_instrs_backward.at(source_instr); @@ -926,15 +925,24 @@ CpuRuntime::CpuRuntime(const Function& function, ContextPtr context, bool concur ++dependency_count; } } + }; + for (auto& out : instr.outputs) { + local_uses_backward.at(out.local_index).push_back(instr_index); + check_dependency(out.local_index); } - dep_counts_backward.at(instr_index) = dependency_count; - if (is_ready) { - _ready_instructions_backward_init.push_back(instr_index); + // also run dependency check for inputs to prevent concurrent writes to + // same gradient + for (auto& in : instr.inputs) { + check_dependency(in.local_index); } - for (auto& in : instr.inputs) { local_sources_backward.at(in.local_index) = instr_index; } + + dep_counts_backward.at(instr_index) = dependency_count; + if (is_ready) { + _ready_instructions_backward_init.push_back(instr_index); + } } _locals_init.resize(function.locals().size()); From eca2ea0ff8535eb99a554c1a91549c4487079a5c Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Sun, 26 Jul 2026 13:23:12 +0200 Subject: [PATCH 17/19] first steps toward madnis reproducibility. backward pass bugfixes --- .github/workflows/acceptancetest_mg7.yml | 46 ++++++ .../iolibs/template_files/mg7/madevent.py | 25 ++- .../madspace/driver/madnis_training.hpp | 19 ++- madspace/include/madspace/driver/random.hpp | 52 ++++++ .../phasespace/channel_weight_network.hpp | 3 +- .../madspace/phasespace/discrete_flow.hpp | 3 +- madspace/include/madspace/phasespace/flow.hpp | 7 +- madspace/include/madspace/phasespace/mlp.hpp | 5 +- madspace/src/cpu/tensor.hpp | 14 ++ madspace/src/driver/channel_generator.cpp | 22 +-- madspace/src/driver/event_generator.cpp | 22 +-- madspace/src/driver/madnis_training.cpp | 148 ++++++++++++------ madspace/src/driver/random.cpp | 14 ++ .../src/phasespace/channel_weight_network.cpp | 6 +- madspace/src/phasespace/discrete_flow.cpp | 4 +- madspace/src/phasespace/flow.cpp | 10 +- madspace/src/phasespace/mlp.cpp | 44 ++++-- madspace/src/python/madspace.cpp | 37 +++-- .../test_mg7_reproducibility.py | 88 ++++++++++- 19 files changed, 446 insertions(+), 123 deletions(-) diff --git a/.github/workflows/acceptancetest_mg7.yml b/.github/workflows/acceptancetest_mg7.yml index 29a0e76ef..0beda1ecb 100644 --- a/.github/workflows/acceptancetest_mg7.yml +++ b/.github/workflows/acceptancetest_mg7.yml @@ -221,6 +221,29 @@ jobs: export PATH="$HOME/.cache/HEPtools/bin:$PATH" ./tests/test_manager.py test_vegas_reproducibility_mg7 -pA -t0 -l INFO + acceptancetest_mg7_madnis_reproducibility: + needs: build_madspace + # mg7 seeded-generation reproducibility for a full survey + madnis + # training + generate run (p p > t t~, no gridpack export): the same + # run_card seed must give a byte-identical LHE file across independent + # runs, including the training phase itself, a different seed must not. + # Companion to acceptancetest_mg7_gridpack_reproducibility, which only + # covers a frozen gridpack's own regeneration. 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_reproducibility_mg7 + run: | + cd $GITHUB_WORKSPACE + export PATH="$HOME/.cache/HEPtools/bin:$PATH" + ./tests/test_manager.py test_madnis_reproducibility_mg7 -pA -t0 -l INFO + acceptancetest_mg7_gridpack_reproducibility: needs: build_madspace # mg7 seeded-generation reproducibility for a gridpack trained with @@ -243,6 +266,29 @@ jobs: 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/madevent.py b/madgraph/iolibs/template_files/mg7/madevent.py index 8fff2a103..307c106d1 100644 --- a/madgraph/iolibs/template_files/mg7/madevent.py +++ b/madgraph/iolibs/template_files/mg7/madevent.py @@ -557,6 +557,12 @@ def train_madnis(self) -> None: config=config, integrands=integrands, cwnets=cwnets, + # Reuses the run_card's own seed (also used by build_event_generator()); + # MultiMadnisTraining derives an independent, non-colliding stream from it + # (see salt::job_kind_madnis_train in random.hpp). 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"], ) madnis_training.train() for phasespace, active_channels in zip( @@ -1299,6 +1305,12 @@ def simplify_phasespace( def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace: madnis_args = self.process.run_card["madnis"] + # Shared across all networks below: each one's global names are + # already unique (per-channel/per-component prefixes), so + # initialize_globals() derives an independent, non-colliding stream + # per tensor from this one base seed (see global_init_seed in + # random.hpp) -- no extra indices need to be threaded through here. + seed = self.process.run_card["run"]["seed"] channels = [] for channel_id, channel in enumerate(phasespace.channels): prefix = f"subproc{self.subproc_id}.channel{channel_id}" @@ -1316,7 +1328,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() @@ -1331,10 +1343,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 @@ -1349,7 +1362,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, @@ -1418,7 +1431,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"] + ) return cwnet def t_channel_mode(self, name: str) -> ms.PhaseSpaceMapping.TChannelMode: diff --git a/madspace/include/madspace/driver/madnis_training.hpp b/madspace/include/madspace/driver/madnis_training.hpp index 932760119..c2dfa005f 100644 --- a/madspace/include/madspace/driver/madnis_training.hpp +++ b/madspace/include/madspace/driver/madnis_training.hpp @@ -1,6 +1,7 @@ #pragma once #include +#include #include "madspace/compgraphs.hpp" #include "madspace/driver/adam_optimizer.hpp" @@ -47,7 +48,8 @@ class MadnisTraining { ContextPtr optimizer_context, const Config& config, const std::vector>& integrands, - const std::optional& cwnet + const std::optional& cwnet, + std::uint64_t seed = 0 ); void train_step(std::size_t batch_index); std::vector active_channels() const; @@ -65,6 +67,10 @@ class MadnisTraining { struct SampleJob { SampleBatch samples; SampleBatch unweighted_samples; + // Per-channel dispatch sequence number (see start_single_job), used to commit + // results in dispatch order regardless of completion order. Unused (left at + // 0) for multi-channel GPU jobs, which are committed on completion instead. + std::size_t channel_seq = 0; }; struct ChannelData { std::size_t index; @@ -77,6 +83,13 @@ class MadnisTraining { RuntimePtr generator_runtime = nullptr; RuntimePtr unweighter_runtime = nullptr; SampleBatch buffer; + // Commit-ordering state for single-channel (CPU) generator jobs: each + // dispatched job gets the next channel-local sequence number, and results are + // applied (see process_job_results) only once all earlier-numbered jobs for + // this channel have already been applied. + std::size_t next_dispatch_seq = 0; + std::size_t commit_cursor = 0; + std::unordered_map ready_job_ids; }; inline static std::function _abort_check_function = [] {}; @@ -115,6 +128,7 @@ class MadnisTraining { std::vector _arg_permutation; bool _buffer_ready = false; std::vector _active_flavors_count; + std::uint64_t _seed; }; class MultiMadnisTraining { @@ -124,7 +138,8 @@ class MultiMadnisTraining { ContextPtr optimizer_context, const MadnisTraining::Config& config, const nested_vector2>& integrands, - const std::vector>& cwnets + const std::vector>& cwnets, + std::uint64_t seed = 0 ); void train(); nested_vector2 active_channels() const; diff --git a/madspace/include/madspace/driver/random.hpp b/madspace/include/madspace/driver/random.hpp index 4faa589fa..bac9d870b 100644 --- a/madspace/include/madspace/driver/random.hpp +++ b/madspace/include/madspace/driver/random.hpp @@ -3,6 +3,7 @@ #include #include #include +#include namespace madspace { @@ -18,4 +19,55 @@ mix_seed(std::uint64_t seed, std::initializer_list salts = {}); std::mt19937 seeded_rng(std::uint64_t seed, std::initializer_list salts = {}); +// Deterministic (platform/compiler-independent) 64-bit FNV-1a hash of `s`, +// used to fold a name (e.g. a global's name) into a mix_seed()/seeded_rng() +// salt. Not for anything security-sensitive -- just needs to be stable across +// runs/machines and have a low accidental-collision rate. +std::uint64_t hash_string(std::string_view s); + +// Registry of the top-level salts used across madspace, so independent RNG +// streams derived from the same base seed by different subsystems never +// collide, even if they happen to be given the same base seed. Each is used +// at its documented call site; add new ones here (with a comment pointing at +// the call site and the salt tuple shape it's used in) rather than picking an +// arbitrary local value. +namespace salt { + +// driver/channel_generator.cpp -- ChannelEventGenerator::start_job's per-job +// base seed: mix_seed(seed, {job.channel_index, job_kind_*, ...}). +inline constexpr std::uint64_t job_kind_survey = 1; +inline constexpr std::uint64_t job_kind_generate = 2; + +// driver/channel_generator.cpp -- sub-seed for an individual Runtime::run() +// call within a job: mix_seed(job.rng_seed, {phase_*, call_index}). +inline constexpr std::uint64_t phase_generate = 1; +inline constexpr std::uint64_t phase_unweight = 2; + +// driver/event_generator.cpp -- EventGenerator's post-generation passes: +// seeded_rng(seed, {combine_select}) / seeded_rng(seed, {lhe_complete}) / +// seeded_rng(seed, {lhe_complete, batch_seq}) / +// seeded_rng(seed, {unweight_pass, call_index, channel_index}). +inline constexpr std::uint64_t combine_select = 1; +inline constexpr std::uint64_t lhe_complete = 2; +// Named unweight_pass, not unweight_all, to not collide with the unrelated +// member function EventGenerator::unweight_all(). +inline constexpr std::uint64_t unweight_pass = 3; + +// driver/madnis_training.cpp -- MadnisTraining::start_single_job's per-job +// base seed: mix_seed(seed, {job_kind_madnis_train, channel_index, channel_seq}). +inline constexpr std::uint64_t job_kind_madnis_train = 100; + +// phasespace/*.cpp -- random initialization of a named global (MLP weights +// and biases): see global_init_seed(). +inline constexpr std::uint64_t global_init = 200; + +} // namespace salt + +// Derives the seed used to randomly initialize the global named `name`, +// salted (via salt::kGlobalInit and a hash of `name`) so it collides neither +// between different globals nor with any other subsystem's stream derived +// from the same base `seed`. `seed == 0` passes through unchanged (the usual +// non-deterministic sentinel). +std::uint64_t global_init_seed(std::uint64_t seed, std::string_view name); + } // namespace madspace diff --git a/madspace/include/madspace/phasespace/channel_weight_network.hpp b/madspace/include/madspace/phasespace/channel_weight_network.hpp index b7e0cbb0e..9c605bf22 100644 --- a/madspace/include/madspace/phasespace/channel_weight_network.hpp +++ b/madspace/include/madspace/phasespace/channel_weight_network.hpp @@ -32,7 +32,8 @@ class ChannelWeightNetwork : public FunctionGenerator { const MLP& mlp() const { return _mlp; } const MomentumPreprocessing& preprocessing() const { return _preprocessing; } - void initialize_globals(ContextPtr context) const; + // seed == 0 (the default) initializes non-deterministically. + void initialize_globals(ContextPtr context, std::uint64_t seed = 0) 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..de15814dc 100644 --- a/madspace/include/madspace/phasespace/discrete_flow.hpp +++ b/madspace/include/madspace/phasespace/discrete_flow.hpp @@ -18,7 +18,8 @@ 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; + // seed == 0 (the default) initializes non-deterministically. + void initialize_globals(ContextPtr context, std::uint64_t seed = 0) const; private: Result build_forward_impl( diff --git a/madspace/include/madspace/phasespace/flow.hpp b/madspace/include/madspace/phasespace/flow.hpp index fcf064084..3e7f772df 100644 --- a/madspace/include/madspace/phasespace/flow.hpp +++ b/madspace/include/madspace/phasespace/flow.hpp @@ -19,8 +19,11 @@ 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; + // seed == 0 (the default) initializes non-deterministically. + void initialize_globals(ContextPtr context, std::uint64_t seed = 0) const; + void initialize_from_vegas( + ContextPtr context, const std::string& grid_name, std::uint64_t seed = 0 + ) const; private: Result build_forward_impl( diff --git a/madspace/include/madspace/phasespace/mlp.hpp b/madspace/include/madspace/phasespace/mlp.hpp index 45234a8f3..3a9157eba 100644 --- a/madspace/include/madspace/phasespace/mlp.hpp +++ b/madspace/include/madspace/phasespace/mlp.hpp @@ -1,5 +1,7 @@ #pragma once +#include + #include "madspace/driver/context.hpp" #include "madspace/phasespace/base.hpp" @@ -17,7 +19,8 @@ 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; + // seed == 0 (the default) initializes non-deterministically. + void initialize_globals(ContextPtr context, std::uint64_t seed = 0) const; std::string last_layer_bias_name() const { return prefixed_name(_prefix, std::format("layer{}.bias", _layers)); } 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 692025c4c..59f182a73 100644 --- a/madspace/src/driver/channel_generator.cpp +++ b/madspace/src/driver/channel_generator.cpp @@ -7,16 +7,6 @@ using json = nlohmann::json; namespace { -// Phase salts keeping the phase-space/matrix-element stream of a job independent -// from its unweighting stream while both derive from the same per-job base seed. -constexpr std::uint64_t kPhaseGenerate = 1; -constexpr std::uint64_t kPhaseUnweight = 2; - -// Salts keeping a job's base-seed derivation (see start_job) independent between -// survey and generate calls. -constexpr std::uint64_t kJobKindSurvey = 1; -constexpr std::uint64_t kJobKindGenerate = 2; - int event_extra_flags(const std::unordered_map& index_map) { int flags = 0; if (index_map.contains("fact_scale1")) { @@ -415,11 +405,13 @@ void ChannelEventGenerator::start_job( job.rng_seed = 0; } else if (is_survey) { job.rng_seed = mix_seed( - seed, {job.channel_index, kJobKindSurvey, survey_pass, _survey_rng_seq++} + seed, + {job.channel_index, salt::job_kind_survey, survey_pass, _survey_rng_seq++} ); } else { - job.rng_seed = - mix_seed(seed, {job.channel_index, kJobKindGenerate, _generate_rng_seq++}); + job.rng_seed = mix_seed( + seed, {job.channel_index, salt::job_kind_generate, _generate_rng_seq++} + ); } _contexts.at(job.context_index) ->thread_pool() @@ -432,7 +424,7 @@ void ChannelEventGenerator::start_job( if (job.rng_seed == 0) { return std::nullopt; } - return mix_seed(job.rng_seed, {kPhaseGenerate, rng_call_index++}); + return mix_seed(job.rng_seed, {salt::phase_generate, rng_call_index++}); }; std::size_t max_batch_size = context->device()->device_type() == DeviceType::cpu @@ -538,7 +530,7 @@ void ChannelEventGenerator::submit_unweight_job( auto& context = _contexts.at(job.context_index); std::optional seed = job.rng_seed == 0 ? std::nullopt - : std::optional(mix_seed(job.rng_seed, {kPhaseUnweight})); + : std::optional(mix_seed(job.rng_seed, {salt::phase_unweight})); std::vector unweighter_args( job.events.begin(), job.events.begin() + _field_indices.random ); diff --git a/madspace/src/driver/event_generator.cpp b/madspace/src/driver/event_generator.cpp index e422cd395..52a515444 100644 --- a/madspace/src/driver/event_generator.cpp +++ b/madspace/src/driver/event_generator.cpp @@ -14,15 +14,6 @@ using namespace madspace; -namespace { - -// salts to derive seeds for different generation phases -constexpr std::uint32_t salt_combine_select = 1; -constexpr std::uint32_t salt_lhe_complete = 2; -constexpr std::uint32_t salt_unweight = 3; - -} // namespace - const GeneratorConfig EventGenerator::default_config = {}; EventGenerator::EventGenerator( @@ -699,7 +690,7 @@ void EventGenerator::update_counts() { void EventGenerator::combine_to_compact_npy(const std::string& file_name) { reset_start_time(); - std::mt19937 select_rng = seeded_rng(_seed, {salt_combine_select}); + std::mt19937 select_rng = seeded_rng(_seed, {salt::combine_select}); auto [channel_data, particle_count, norm_factor] = init_combine(); DataLayout layout( EventRecord::layout( @@ -737,8 +728,8 @@ void EventGenerator::combine_to_lhe_npy( ) { reset_start_time(); std::uint64_t seed = _seed; - std::mt19937 select_rng = seeded_rng(seed, {salt_combine_select}); - std::mt19937 rand_gen = seeded_rng(seed, {salt_lhe_complete}); + std::mt19937 select_rng = seeded_rng(seed, {salt::combine_select}); + std::mt19937 rand_gen = seeded_rng(seed, {salt::lhe_complete}); auto [channel_data, particle_count, norm_factor] = init_combine(); DataLayout in_layout( EventRecord::layout( @@ -801,7 +792,7 @@ void EventGenerator::combine_to_lhe( ) { reset_start_time(); std::uint64_t seed = _seed; - std::mt19937 select_rng = seeded_rng(seed, {salt_combine_select}); + std::mt19937 select_rng = seeded_rng(seed, {salt::combine_select}); ThreadPool pool(_config.combine_thread_count); auto [channel_data, particle_count, norm_factor] = init_combine(); std::vector> buffers; @@ -847,7 +838,8 @@ void EventGenerator::combine_to_lhe( pool.submit( [slot, batch_seq, seed, this, &in_buffer, &out_buffer, &lhe_completer] { std::mt19937 rand_gen = seeded_rng( - seed, {salt_lhe_complete, static_cast(batch_seq)} + seed, + {salt::lhe_complete, static_cast(batch_seq)} ); LHEEvent lhe_event; out_buffer.clear(); @@ -910,7 +902,7 @@ void EventGenerator::unweight_all() { for (std::size_t channel_index = 0; auto& channel : _channels) { std::mt19937 rand_gen = seeded_rng( _seed, - {salt_unweight, + {salt::unweight_pass, static_cast(call_index), static_cast(channel_index)} ); diff --git a/madspace/src/driver/madnis_training.cpp b/madspace/src/driver/madnis_training.cpp index c94d84fb9..080b6a2e0 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,15 @@ MadnisTraining::MadnisTraining( ContextPtr optimizer_context, const Config& config, const std::vector>& integrands, - const std::optional& cwnet + const std::optional& cwnet, + std::uint64_t seed ) : _generator_context(generator_context), _optimizer_context(optimizer_context), _config(config), _channels(integrands.size()), - _cwnet(cwnet) { + _cwnet(cwnet), + _seed(seed) { for (std::size_t index = 0; auto [integrand, channel] : zip(integrands, _channels)) { channel.index = index; @@ -262,6 +265,17 @@ void MadnisTraining::start_generator_jobs(const std::vector& counts if (_running_jobs.size() > 0) { return; } + // KNOWN LIMITATION: the weight snapshot taken below is gated by + // "all previously dispatched jobs have completed", a condition that depends + // on real wall-clock thread-completion timing rather than a deterministic + // logical step boundary. Under concurrent execution, the number and + // position of these snapshot/round boundaries can therefore vary between + // otherwise identically-seeded runs (e.g. depending on how batches happen + // to be scheduled across worker threads), which can make full end-to-end + // training reproducibility intermittently flaky even though every + // individual job is itself deterministically seeded. Fixing this would + // require pinning round boundaries to a deterministic step counter instead + // of thread-completion state; left as a follow-up design discussion. _generator_params.copy_from(_optimizer->parameters()); std::size_t chan_count = counts.size(); bool is_gpu = _generator_context->device()->device_type() != DeviceType::cpu; @@ -358,16 +372,31 @@ TensorVec MadnisTraining::permute_tensors(const TensorVec& tensors) const { return ret; } +// Assigns the job's base seed on the scheduling thread, from its logical identity +// only (channel index, channel-local sequence) -- never from which worker thread ends +// up running it -- so the sample stream doesn't depend on the thread pool's +// scheduling. generator_runtime is built with concurrent=false (see +// build_runtimes_and_optimizer), so a single job's own graph execution is always +// single-threaded and safe to seed even though many jobs run concurrently across +// the thread pool. 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::optional seed = _seed == 0 + ? std::nullopt + : std::optional( + mix_seed(_seed, {salt::job_kind_madnis_train, channel_index, channel_seq}) + ); auto& job = std::get<0>(_running_jobs.emplace(job_id, SampleJob{}))->second; + job.channel_seq = channel_seq; _generator_context->thread_pool().submit( - [this, channel_index, batch_size, job_id, &job]() { + [this, channel_index, batch_size, job_id, seed, &job]() { auto& channel = _channels.at(channel_index); - auto samples = channel.generator_runtime->run({Tensor({batch_size})}); + auto samples = channel.generator_runtime->run({Tensor({batch_size})}, seed); job.samples.tensors = permute_tensors(samples); job.samples.size = samples.at(0).size(0); job.samples.channel_index = channel_index; @@ -382,6 +411,8 @@ void MadnisTraining::start_single_job( ); } +// TODO: GPU multi-channel batches aren't seeded or given a commit-ordering scheme +// (see process_job_results) -- left for a follow-up, like the buffered-training path. void MadnisTraining::start_multi_job(const std::vector batch_sizes) { std::size_t job_id = _job_id; ++_job_id; @@ -507,48 +538,62 @@ MadnisTraining::build_buffered_training_batch(const std::vector& counts) 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); 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); + // Single-channel (CPU) job: mark ready, committed below strictly in + // dispatch order so results don't depend on completion timing. + _channels.at(job.samples.channel_index) + .ready_job_ids.emplace(job.channel_seq, job_id); + continue; + } + auto multi_job = std::move(_running_jobs.extract(job_id).mapped()); + 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; } - } else { - std::size_t offset = 0, unw_offset = 0, chan_index = 0; - SampleBatch chan_unweighted_samples; - for (auto [channel, chan_size] : - zip(_channels, job.samples.channel_sizes)) { - if (chan_size == 0) { - ++chan_index; - continue; - } - channel.sample_count += chan_size; - 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) + 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) - ); - } - buffer_store(channel, chan_unweighted_samples); - unw_offset += unw_chan_size; - } - batch.size = chan_size; - offset += chan_size; - ++chan_index; + buffer_store(channel, chan_unweighted_samples); + unw_offset += unw_chan_size; + } + batch.size = chan_size; + offset += chan_size; + ++chan_index; + } + } + + // Commit each channel's completed single-channel jobs strictly in dispatch order + // (channel_seq), independent of which job happened to finish first. + 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) { + buffer_store(channel, committed_job.unweighted_samples); } } } @@ -642,7 +687,9 @@ 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 in abs_means (e.g. symmetric/charge-conjugate channels) + // must break the same way every run, not depend on sort implementation details. + std::stable_sort(indices.begin(), indices.end(), [&](auto i, auto j) { return abs_means.at(i) < abs_means.at(j); }); @@ -719,14 +766,23 @@ MultiMadnisTraining::MultiMadnisTraining( ContextPtr optimizer_context, const MadnisTraining::Config& config, const nested_vector2>& integrands, - const std::vector>& cwnets + const std::vector>& cwnets, + std::uint64_t seed ) : _config(config) { _subprocesses.reserve(integrands.size()); - for (auto [integ, cwnet] : zip(integrands, cwnets)) { + // Each subprocess gets its own derived seed so their per-job streams (see + // MadnisTraining::start_single_job) don't collide with each other. + for (std::size_t subproc_index = 0; auto [integ, cwnet] : zip(integrands, cwnets)) { _subprocesses.emplace_back( - generator_context, optimizer_context, config, integ, cwnet + generator_context, + optimizer_context, + config, + integ, + cwnet, + seed == 0 ? 0 : mix_seed(seed, {subproc_index}) ); + ++subproc_index; } } diff --git a/madspace/src/driver/random.cpp b/madspace/src/driver/random.cpp index 41595e29b..4ad9b8298 100644 --- a/madspace/src/driver/random.cpp +++ b/madspace/src/driver/random.cpp @@ -36,4 +36,18 @@ seeded_rng(std::uint64_t seed, std::initializer_list salts) { return std::mt19937(seq); } +std::uint64_t hash_string(std::string_view s) { + // FNV-1a, 64-bit. + std::uint64_t h = 0xcbf29ce484222325ull; + for (unsigned char c : s) { + h ^= c; + h *= 0x100000001b3ull; + } + return h; +} + +std::uint64_t global_init_seed(std::uint64_t seed, std::string_view name) { + return mix_seed(seed, {salt::global_init, hash_string(name)}); +} + } // namespace madspace diff --git a/madspace/src/phasespace/channel_weight_network.cpp b/madspace/src/phasespace/channel_weight_network.cpp index d84a7a389..972ad7a73 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::uint64_t 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..e4b9f00a3 100644 --- a/madspace/src/phasespace/discrete_flow.cpp +++ b/madspace/src/phasespace/discrete_flow.cpp @@ -139,13 +139,13 @@ Mapping::Result DiscreteFlow::build_transform( }; } -void DiscreteFlow::initialize_globals(ContextPtr context) const { +void DiscreteFlow::initialize_globals(ContextPtr context, std::uint64_t 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..682266d0d 100644 --- a/madspace/src/phasespace/flow.cpp +++ b/madspace/src/phasespace/flow.cpp @@ -211,17 +211,17 @@ Flow::Flow( } } -void Flow::initialize_globals(ContextPtr context) const { +void Flow::initialize_globals(ContextPtr context, std::uint64_t 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::uint64_t 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..81542ced7 100644 --- a/madspace/src/phasespace/mlp.cpp +++ b/madspace/src/phasespace/mlp.cpp @@ -3,6 +3,8 @@ #include #include +#include "madspace/driver/random.hpp" + using namespace madspace; namespace { @@ -51,11 +53,9 @@ void initialize_layer( std::size_t output_dim, const std::string& prefix, int layer_index, - std::mt19937& rand_gen, + std::uint64_t 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 +75,30 @@ 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); + std::uniform_real_distribution rand_dist(-bound, bound); + // Independently seeded per tensor (keyed by its own global name), so + // e.g. reordering layers or adding new ones doesn't shift any other + // tensor's initialized values. + auto weight_rand_gen = seeded_rng(global_init_seed(seed, weight_name)); + for (std::size_t i = 0; i < output_dim; ++i) { + for (std::size_t j = 0; j < input_dim; ++j) { + weight_view[i][j] = rand_dist(weight_rand_gen); + } + } + auto bias_rand_gen = seeded_rng(global_init_seed(seed, bias_name)); + for (std::size_t i = 0; i < output_dim; ++i) { + bias_view[i] = rand_dist(bias_rand_gen); + } } if (!is_cpu) { @@ -133,15 +149,13 @@ 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::uint64_t 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 5526a6344..efd355bab 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") = 0 + ); 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") = 0 + ) .def( "initialize_from_vegas", &Flow::initialize_from_vegas, py::arg("context"), - py::arg("grid_name") + py::arg("grid_name"), + py::arg("seed") = 0 ); 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") = 0 ); 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") = 0 ); py::classh(m, "VegasGridOptimizer") @@ -1372,12 +1387,14 @@ PYBIND11_MODULE(_madspace_py, m) { ContextPtr, const MadnisTraining::Config&, const std::vector>&, - const std::optional&>(), + const std::optional&, + std::uint64_t>(), py::arg("generator_context"), py::arg("optimizer_context"), py::arg("config"), py::arg("integrands"), - py::arg("cwnet") + py::arg("cwnet"), + py::arg("seed") = 0 ) .def("train_step", &MadnisTraining::train_step, py::arg("batch_index")) .def("active_channels", &MadnisTraining::active_channels) @@ -1390,12 +1407,14 @@ PYBIND11_MODULE(_madspace_py, m) { ContextPtr, const MadnisTraining::Config&, const nested_vector2>&, - const std::vector>&>(), + const std::vector>&, + std::uint64_t>(), py::arg("generator_context"), py::arg("optimizer_context"), py::arg("config"), py::arg("integrands"), - py::arg("cwnets") + py::arg("cwnets"), + py::arg("seed") = 0 ) .def("train", &MultiMadnisTraining::train) .def("active_channels", &MultiMadnisTraining::active_channels); diff --git a/tests/acceptance_tests/test_mg7_reproducibility.py b/tests/acceptance_tests/test_mg7_reproducibility.py index 317613a23..8d20a63ff 100644 --- a/tests/acceptance_tests/test_mg7_reproducibility.py +++ b/tests/acceptance_tests/test_mg7_reproducibility.py @@ -16,15 +16,21 @@ 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 the two ways an mg7 -run produces events: +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_reproducibility_mg7`` -- a full survey + madnis training + + generate run, through ``bin/generate_events`` (no gridpack export). + Unlike the gridpack test below, this exercises the training phase's own + seeding (single-channel CPU sample generation only -- buffered training + and GPU multi-channel batches are not yet made deterministic). * ``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 @@ -221,6 +227,42 @@ def test_vegas_reproducibility_mg7(self): self.assertNotEqual(hash_a1, hash_b, 'different seeds produced identical LHE content') + def test_madnis_reproducibility_mg7(self): + """Two full survey + madnis training + generate runs (no gridpack + export) with the same seed produce byte-identical LHE files; a + different seed changes the result. Unlike + test_gridpack_reproducibility_mg7, training is re-run from scratch + each time here (not frozen into a gridpack first), so this is the + test that actually exercises MultiMadnisTraining's own seeding + (single-channel CPU sample generation, see kJobKindMadnisTrain in + madnis_training.cpp).""" + datadir = _mg7_datadir_or_skip(self) + run_dir = self._output_process('madnis') + toml = pjoin(run_dir, 'Cards', 'run_card.toml') + self._set_run_card( + toml, + **{ + 'generation.events': _EVENTS, + 'postprocessing.systematics': False, + 'madnis.enable': True, + 'madnis.train_batches': 100, + } + ) + + # 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_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_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 @@ -263,6 +305,48 @@ def test_gridpack_reproducibility_mg7(self): 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() From 8bf182a4815270ee755b136d28b59263208a10be Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Sun, 26 Jul 2026 20:42:11 +0200 Subject: [PATCH 18/19] reproducible generation job scheduling during madnis training --- .github/workflows/acceptancetest_mg7.yml | 22 ++++ .../iolibs/template_files/mg7/madevent.py | 1 + .../iolibs/template_files/mg7/run_card.toml | 1 + madgraph/various/banner.py | 1 + .../madspace/driver/madnis_training.hpp | 21 ++-- madspace/include/madspace/driver/random.hpp | 4 + madspace/src/driver/madnis_training.cpp | 110 ++++++++++++------ madspace/src/python/madspace.cpp | 3 +- .../test_mg7_reproducibility.py | 85 +++++++++++++- 9 files changed, 198 insertions(+), 50 deletions(-) diff --git a/.github/workflows/acceptancetest_mg7.yml b/.github/workflows/acceptancetest_mg7.yml index 0beda1ecb..f293c43ef 100644 --- a/.github/workflows/acceptancetest_mg7.yml +++ b/.github/workflows/acceptancetest_mg7.yml @@ -244,6 +244,28 @@ jobs: export PATH="$HOME/.cache/HEPtools/bin:$PATH" ./tests/test_manager.py test_madnis_reproducibility_mg7 -pA -t0 -l INFO + acceptancetest_mg7_madnis_reproducible_mode: + needs: build_madspace + # mg7 madnis.reproducible = True (p p > t t~): unlike + # acceptancetest_mg7_madnis_reproducibility, this requires madnis training + # itself to 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 diff --git a/madgraph/iolibs/template_files/mg7/madevent.py b/madgraph/iolibs/template_files/mg7/madevent.py index 307c106d1..97d1c8ba0 100644 --- a/madgraph/iolibs/template_files/mg7/madevent.py +++ b/madgraph/iolibs/template_files/mg7/madevent.py @@ -532,6 +532,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 = [] diff --git a/madgraph/iolibs/template_files/mg7/run_card.toml b/madgraph/iolibs/template_files/mg7/run_card.toml index 13e65ea99..35b9bc173 100644 --- a/madgraph/iolibs/template_files/mg7/run_card.toml +++ b/madgraph/iolibs/template_files/mg7/run_card.toml @@ -142,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 afc828d58..474d6465e 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -6594,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/include/madspace/driver/madnis_training.hpp b/madspace/include/madspace/driver/madnis_training.hpp index c2dfa005f..d3a94a8f6 100644 --- a/madspace/include/madspace/driver/madnis_training.hpp +++ b/madspace/include/madspace/driver/madnis_training.hpp @@ -42,6 +42,10 @@ 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, @@ -67,9 +71,7 @@ class MadnisTraining { struct SampleJob { SampleBatch samples; SampleBatch unweighted_samples; - // Per-channel dispatch sequence number (see start_single_job), used to commit - // results in dispatch order regardless of completion order. Unused (left at - // 0) for multi-channel GPU jobs, which are committed on completion instead. + // per-channel dispatch sequence, used to commit in order; unused for GPU jobs std::size_t channel_seq = 0; }; struct ChannelData { @@ -83,13 +85,13 @@ class MadnisTraining { RuntimePtr generator_runtime = nullptr; RuntimePtr unweighter_runtime = nullptr; SampleBatch buffer; - // Commit-ordering state for single-channel (CPU) generator jobs: each - // dispatched job gets the next channel-local sequence number, and results are - // applied (see process_job_results) only once all earlier-numbered jobs for - // this channel have already been applied. + // 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 = [] {}; @@ -97,6 +99,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); @@ -129,6 +134,8 @@ class MadnisTraining { bool _buffer_ready = false; std::vector _active_flavors_count; std::uint64_t _seed; + // sequence for seeding build_buffered_training_batch's BatchSampler::run() call + std::size_t _buffered_batch_seq = 0; }; class MultiMadnisTraining { diff --git a/madspace/include/madspace/driver/random.hpp b/madspace/include/madspace/driver/random.hpp index bac9d870b..33a48247f 100644 --- a/madspace/include/madspace/driver/random.hpp +++ b/madspace/include/madspace/driver/random.hpp @@ -56,6 +56,10 @@ inline constexpr std::uint64_t unweight_pass = 3; // driver/madnis_training.cpp -- MadnisTraining::start_single_job's per-job // base seed: mix_seed(seed, {job_kind_madnis_train, channel_index, channel_seq}). inline constexpr std::uint64_t job_kind_madnis_train = 100; +// build_buffered_training_batch's BatchSampler seed: mix_seed(seed, {..., seq}). +inline constexpr std::uint64_t buffered_batch_sample = 101; +// start_single_job's per-job BufferUnweighter seed: mix_seed(job_seed, {...}). +inline constexpr std::uint64_t madnis_train_unweight = 102; // phasespace/*.cpp -- random initialization of a named global (MLP weights // and biases): see global_init_seed(). diff --git a/madspace/src/driver/madnis_training.cpp b/madspace/src/driver/madnis_training.cpp index 080b6a2e0..2c1b9f2d2 100644 --- a/madspace/src/driver/madnis_training.cpp +++ b/madspace/src/driver/madnis_training.cpp @@ -1,5 +1,7 @@ #include "madspace/driver/madnis_training.hpp" +#include + #include "madspace/driver/random.hpp" #include "madspace/phasespace/batch_sampler.hpp" @@ -57,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); @@ -265,20 +267,18 @@ void MadnisTraining::start_generator_jobs(const std::vector& counts if (_running_jobs.size() > 0) { return; } - // KNOWN LIMITATION: the weight snapshot taken below is gated by - // "all previously dispatched jobs have completed", a condition that depends - // on real wall-clock thread-completion timing rather than a deterministic - // logical step boundary. Under concurrent execution, the number and - // position of these snapshot/round boundaries can therefore vary between - // otherwise identically-seeded runs (e.g. depending on how batches happen - // to be scheduled across worker threads), which can make full end-to-end - // training reproducibility intermittently flaky even though every - // individual job is itself deterministically seeded. Fixing this would - // require pinning round boundaries to a deterministic step counter instead - // of thread-completion state; left as a follow-up design discussion. + bool is_gpu = _generator_context->device()->device_type() != DeviceType::cpu; + if (_config.reproducible && !is_gpu) { + // 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; @@ -294,7 +294,11 @@ 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(); + // reproducible mode (CPU only): dispatch the full round in one shot, uncapped + // by thread count, so round contents don't depend on thread count + std::size_t available_jobs = _config.reproducible && !is_gpu + ? std::numeric_limits::max() + : _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 - @@ -363,6 +367,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()); @@ -372,13 +402,8 @@ TensorVec MadnisTraining::permute_tensors(const TensorVec& tensors) const { return ret; } -// Assigns the job's base seed on the scheduling thread, from its logical identity -// only (channel index, channel-local sequence) -- never from which worker thread ends -// up running it -- so the sample stream doesn't depend on the thread pool's -// scheduling. generator_runtime is built with concurrent=false (see -// build_runtimes_and_optimizer), so a single job's own graph execution is always -// single-threaded and safe to seed even though many jobs run concurrently across -// the thread pool. +// 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 ) { @@ -401,7 +426,13 @@ void MadnisTraining::start_single_job( job.samples.size = samples.at(0).size(0); job.samples.channel_index = channel_index; if (channel.unweighter_runtime) { - auto unw_samples = channel.unweighter_runtime->run(samples); + std::optional unweight_seed = seed + ? std::optional( + mix_seed(*seed, {salt::madnis_train_unweight}) + ) + : std::nullopt; + auto unw_samples = + channel.unweighter_runtime->run(samples, unweight_seed); job.unweighted_samples.tensors = permute_tensors(unw_samples); job.unweighted_samples.size = unw_samples.at(0).size(0); job.unweighted_samples.channel_index = channel_index; @@ -411,8 +442,8 @@ void MadnisTraining::start_single_job( ); } -// TODO: GPU multi-channel batches aren't seeded or given a commit-ordering scheme -// (see process_job_results) -- left for a follow-up, like the buffered-training path. +// TODO: GPU multi-channel batches aren't seeded or commit-ordered, unlike +// start_single_job. void MadnisTraining::start_multi_job(const std::vector batch_sizes) { std::size_t job_id = _job_id; ++_job_id; @@ -533,15 +564,19 @@ MadnisTraining::build_buffered_training_batch(const std::vector& counts) } } args.emplace_back(counts); - return _multi_channel_sampler->run(args); + std::optional seed = _seed == 0 + ? std::nullopt + : std::optional( + mix_seed(_seed, {salt::buffered_batch_sample, _buffered_batch_seq++}) + ); + return _multi_channel_sampler->run(args, seed); } void MadnisTraining::process_job_results(const std::vector& job_ids) { for (auto job_id : job_ids) { auto& job = _running_jobs.at(job_id); if (job.samples.channel_sizes.size() == 0) { - // Single-channel (CPU) job: mark ready, committed below strictly in - // dispatch order so results don't depend on completion timing. + // single-channel job: mark ready, committed below in dispatch order _channels.at(job.samples.channel_index) .ready_job_ids.emplace(job.channel_seq, job_id); continue; @@ -581,8 +616,7 @@ void MadnisTraining::process_job_results(const std::vector& job_ids } } - // Commit each channel's completed single-channel jobs strictly in dispatch order - // (channel_seq), independent of which job happened to finish first. + // commit each channel's ready 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(); @@ -593,7 +627,14 @@ void MadnisTraining::process_job_results(const std::vector& job_ids channel.sample_count += committed_job.samples.size; channel.sample_batches.push_back(std::move(committed_job.samples)); if (committed_job.unweighted_samples.size > 0) { - buffer_store(channel, committed_job.unweighted_samples); + 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); + } } } } @@ -687,8 +728,7 @@ void MadnisTraining::drop_channels() { } std::vector indices(_channels.size()); std::iota(indices.begin(), indices.end(), 0); - // stable_sort: exact ties in abs_means (e.g. symmetric/charge-conjugate channels) - // must break the same way every run, not depend on sort implementation details. + // 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); }); @@ -715,8 +755,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) { @@ -771,8 +810,7 @@ MultiMadnisTraining::MultiMadnisTraining( ) : _config(config) { _subprocesses.reserve(integrands.size()); - // Each subprocess gets its own derived seed so their per-job streams (see - // MadnisTraining::start_single_job) don't collide with each other. + // each subprocess gets its own derived seed so their streams don't collide for (std::size_t subproc_index = 0; auto [integ, cwnet] : zip(integrands, cwnets)) { _subprocesses.emplace_back( generator_context, diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index efd355bab..dc57d4285 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -1378,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( diff --git a/tests/acceptance_tests/test_mg7_reproducibility.py b/tests/acceptance_tests/test_mg7_reproducibility.py index 8d20a63ff..ea91a3c52 100644 --- a/tests/acceptance_tests/test_mg7_reproducibility.py +++ b/tests/acceptance_tests/test_mg7_reproducibility.py @@ -21,10 +21,17 @@ * ``test_vegas_reproducibility_mg7`` -- a plain survey (VEGAS-optimized) + generate run, through ``bin/generate_events``. * ``test_madnis_reproducibility_mg7`` -- a full survey + madnis training + - generate run, through ``bin/generate_events`` (no gridpack export). - Unlike the gridpack test below, this exercises the training phase's own - seeding (single-channel CPU sample generation only -- buffered training - and GPU multi-channel batches are not yet made deterministic). + generate run, through ``bin/generate_events`` (no gridpack export), with + ``madnis.reproducible`` left at its default (off). Unlike the gridpack + test below, this exercises the training phase's own seeding + (single-channel CPU sample generation only); training's own round + scheduling is not required to be deterministic in this mode, so this test + may occasionally be flaky (see test_madnis_reproducible_mode_mg7). + * ``test_madnis_reproducible_mode_mg7`` -- same, but 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 still 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 @@ -234,8 +241,12 @@ def test_madnis_reproducibility_mg7(self): test_gridpack_reproducibility_mg7, training is re-run from scratch each time here (not frozen into a gridpack first), so this is the test that actually exercises MultiMadnisTraining's own seeding - (single-channel CPU sample generation, see kJobKindMadnisTrain in - madnis_training.cpp).""" + (single-channel CPU sample generation, see salt::job_kind_madnis_train + in random.hpp). madnis.reproducible is left at its default (off) here, + so -- per the known limitation documented at + MadnisTraining::start_generator_jobs -- this is not expected to be + thread-count independent and may occasionally be flaky; see + test_madnis_reproducible_mode_mg7 for the mode that guarantees it.""" datadir = _mg7_datadir_or_skip(self) run_dir = self._output_process('madnis') toml = pjoin(run_dir, 'Cards', 'run_card.toml') @@ -263,6 +274,68 @@ def test_madnis_reproducibility_mg7(self): 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 From 4cf11053543464cb473f7a14f2697392a45bfb80 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Sun, 26 Jul 2026 21:37:26 +0200 Subject: [PATCH 19/19] only test madnis reproducibility in reproducible mode --- .github/workflows/acceptancetest_mg7.yml | 34 ++---------- .../test_mg7_reproducibility.py | 52 ++----------------- 2 files changed, 8 insertions(+), 78 deletions(-) diff --git a/.github/workflows/acceptancetest_mg7.yml b/.github/workflows/acceptancetest_mg7.yml index f293c43ef..3366206ec 100644 --- a/.github/workflows/acceptancetest_mg7.yml +++ b/.github/workflows/acceptancetest_mg7.yml @@ -221,37 +221,13 @@ jobs: export PATH="$HOME/.cache/HEPtools/bin:$PATH" ./tests/test_manager.py test_vegas_reproducibility_mg7 -pA -t0 -l INFO - acceptancetest_mg7_madnis_reproducibility: - needs: build_madspace - # mg7 seeded-generation reproducibility for a full survey + madnis - # training + generate run (p p > t t~, no gridpack export): the same - # run_card seed must give a byte-identical LHE file across independent - # runs, including the training phase itself, a different seed must not. - # Companion to acceptancetest_mg7_gridpack_reproducibility, which only - # covers a frozen gridpack's own regeneration. 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_reproducibility_mg7 - run: | - cd $GITHUB_WORKSPACE - export PATH="$HOME/.cache/HEPtools/bin:$PATH" - ./tests/test_manager.py test_madnis_reproducibility_mg7 -pA -t0 -l INFO - acceptancetest_mg7_madnis_reproducible_mode: needs: build_madspace - # mg7 madnis.reproducible = True (p p > t t~): unlike - # acceptancetest_mg7_madnis_reproducibility, this requires madnis training - # itself to 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. + # 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: diff --git a/tests/acceptance_tests/test_mg7_reproducibility.py b/tests/acceptance_tests/test_mg7_reproducibility.py index ea91a3c52..6f9edd8e3 100644 --- a/tests/acceptance_tests/test_mg7_reproducibility.py +++ b/tests/acceptance_tests/test_mg7_reproducibility.py @@ -20,18 +20,12 @@ * ``test_vegas_reproducibility_mg7`` -- a plain survey (VEGAS-optimized) + generate run, through ``bin/generate_events``. - * ``test_madnis_reproducibility_mg7`` -- a full survey + madnis training + + * ``test_madnis_reproducible_mode_mg7`` -- a full survey + madnis training + generate run, through ``bin/generate_events`` (no gridpack export), with - ``madnis.reproducible`` left at its default (off). Unlike the gridpack - test below, this exercises the training phase's own seeding - (single-channel CPU sample generation only); training's own round - scheduling is not required to be deterministic in this mode, so this test - may occasionally be flaky (see test_madnis_reproducible_mode_mg7). - * ``test_madnis_reproducible_mode_mg7`` -- same, but 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 still not covered by this mode. + 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 @@ -45,7 +39,7 @@ Run locally with e.g.:: - ./tests/test_manager.py test_.*_reproducibility_mg7 -pA -t0 -l INFO + ./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: @@ -234,46 +228,6 @@ def test_vegas_reproducibility_mg7(self): self.assertNotEqual(hash_a1, hash_b, 'different seeds produced identical LHE content') - def test_madnis_reproducibility_mg7(self): - """Two full survey + madnis training + generate runs (no gridpack - export) with the same seed produce byte-identical LHE files; a - different seed changes the result. Unlike - test_gridpack_reproducibility_mg7, training is re-run from scratch - each time here (not frozen into a gridpack first), so this is the - test that actually exercises MultiMadnisTraining's own seeding - (single-channel CPU sample generation, see salt::job_kind_madnis_train - in random.hpp). madnis.reproducible is left at its default (off) here, - so -- per the known limitation documented at - MadnisTraining::start_generator_jobs -- this is not expected to be - thread-count independent and may occasionally be flaky; see - test_madnis_reproducible_mode_mg7 for the mode that guarantees it.""" - datadir = _mg7_datadir_or_skip(self) - run_dir = self._output_process('madnis') - toml = pjoin(run_dir, 'Cards', 'run_card.toml') - self._set_run_card( - toml, - **{ - 'generation.events': _EVENTS, - 'postprocessing.systematics': False, - 'madnis.enable': True, - 'madnis.train_batches': 100, - } - ) - - # 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_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