From b76eda18fa346a8ac2d2710ebb9449e250481893 Mon Sep 17 00:00:00 2001 From: Lukas Simon Date: Sat, 24 Oct 2026 10:25:58 +0200 Subject: [PATCH 1/3] Cache exact closed colour overlaps for connected correlators --- src/pyamplicol/color/connections.py | 24 ++++- .../unit/test_color_connection_trace_cache.py | 102 ++++++++++++++++++ 2 files changed, 124 insertions(+), 2 deletions(-) create mode 100644 tests/unit/test_color_connection_trace_cache.py diff --git a/src/pyamplicol/color/connections.py b/src/pyamplicol/color/connections.py index 303958b74..af2ac83de 100644 --- a/src/pyamplicol/color/connections.py +++ b/src/pyamplicol/color/connections.py @@ -24,7 +24,7 @@ from collections.abc import Mapping from dataclasses import dataclass from fractions import Fraction -from functools import cached_property +from functools import cached_property, lru_cache from typing import TypeAlias from .contraction_trace import _simplify_trace_terms_nc_power @@ -489,6 +489,15 @@ def apply_color_connection( return ConnectedColorTensor(connection.output_legs, terms, connection.order) +@lru_cache(maxsize=65536) +def _closed_trace_overlap_nc_terms( + traces: tuple[tuple[int, ...], ...], +) -> tuple[tuple[int, Fraction], ...]: + return tuple( + sorted(_simplify_trace_terms_nc_power(((Fraction(1), 0, traces),)).items()) + ) + + def _literal_tensor_overlap_nc_terms( left: ColorTensor, right: ColorTensor ) -> Mapping[int, Fraction]: @@ -511,7 +520,18 @@ def _literal_tensor_overlap_nc_terms( trace.extend(reversed(left_string.adjoint_labels)) current = left_string.fundamental_label traces.append(tuple(trace)) - return _simplify_trace_terms_nc_power(((Fraction(1), 0, tuple(traces)),)) + # Closed overlaps are invariant under cyclic rotations, trace ordering and + # bijective dummy-index relabelling. Keep reflections and the full Nc series. + words = sorted( + (_canonical_trace(trace) for trace in traces), + key=lambda word: (len(word), word), + ) + labels: dict[int, int] = {} + key = tuple( + tuple(labels.setdefault(label, len(labels)) for label in word) for word in words + ) + # Return a fresh mapping so callers cannot mutate cached coefficients. + return dict(_closed_trace_overlap_nc_terms(key)) def contract_connected_tensor_nc_terms( diff --git a/tests/unit/test_color_connection_trace_cache.py b/tests/unit/test_color_connection_trace_cache.py new file mode 100644 index 000000000..36ed974a6 --- /dev/null +++ b/tests/unit/test_color_connection_trace_cache.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: 0BSD +"""Exactness and reuse of closed overlaps, independently of the cache key.""" + +from fractions import Fraction +from itertools import permutations + +import pytest + +from pyamplicol.color import connections as color + + +def _uncached_overlap(left, right): + traces = list(right.traces) + [tuple(reversed(t)) for t in left.traces] + right_by_q = {s.fundamental_label: s for s in right.open_strings} + left_by_qbar = {s.antifundamental_label: s for s in left.open_strings} + visited = set() + for start in sorted(right_by_q): + if start in visited: + continue + current, trace = start, [] + while current not in visited: + visited.add(current) + ket = right_by_q[current] + bra = left_by_qbar[ket.antifundamental_label] + trace.extend(ket.adjoint_labels) + trace.extend(reversed(bra.adjoint_labels)) + current = bra.fundamental_label + traces.append(tuple(trace)) + return color._simplify_trace_terms_nc_power(((Fraction(1), 0, tuple(traces)),)) + + +def test_reuse_keeps_trace_orientation_and_returns_fresh_mappings(): + cache = color._closed_trace_overlap_nc_terms + cache.cache_clear() + tensor = color.ColorTensor(traces=((1, 2, 3),)) + expected = _uncached_overlap(tensor, tensor) + assert color._literal_tensor_overlap_nc_terms(tensor, tensor) == expected + renamed = color.ColorTensor(traces=((12, 13, 11),)) + actual = color._literal_tensor_overlap_nc_terms(renamed, renamed) + assert actual == expected + assert cache.cache_info().misses == 1 + assert cache.cache_info().hits == 1 + actual.clear() + assert color._literal_tensor_overlap_nc_terms(tensor, tensor) == expected + reversed_tensor = color.ColorTensor(traces=((3, 2, 1),)) + reflected = color._literal_tensor_overlap_nc_terms(tensor, reversed_tensor) + assert reflected == _uncached_overlap(tensor, reversed_tensor) + assert reflected != expected + assert cache.cache_parameters()["maxsize"] == 65536 + assert cache.cache_info().currsize <= cache.cache_info().maxsize + + +@pytest.mark.parametrize( + "left,right", + ( + (color.ColorTensor(), color.ColorTensor()), + (color.ColorTensor(traces=((),)), color.ColorTensor(traces=((),))), + (color.ColorTensor(traces=((1,),)), color.ColorTensor(traces=((1,),))), + ( + color.ColorTensor(traces=((1, 2), (3, 4))), + color.ColorTensor(traces=((1, 3, 2, 4),)), + ), + ( + color.ColorTensor((color.OpenColorString(1, (), 2),)), + color.ColorTensor((color.OpenColorString(1, (), 2),)), + ), + ( + color.ColorTensor( + (color.OpenColorString(1, (5,), 2), color.OpenColorString(3, (), 4)) + ), + color.ColorTensor( + (color.OpenColorString(1, (), 4), color.OpenColorString(3, (5,), 2)) + ), + ), + ), +) +def test_full_laurent_series_and_closed_fundamental_pairings(left, right): + assert color._literal_tensor_overlap_nc_terms(left, right) == _uncached_overlap( + left, right + ) + + +def test_five_gluon_connected_entries_match_uncached_polynomials(monkeypatch): + legs = tuple(color.ColorLeg(i, 8) for i in range(1, 6)) + tensors = [ + color.ColorTensor(traces=((1, *tail),)) for tail in permutations(range(2, 6)) + ] + bra = color.ColorConnection(legs, (color.EmitGluon(1, -1), color.EmitGluon(2, -2))) + ket = color.ColorConnection(legs, (color.EmitGluon(3, -1), color.EmitGluon(4, -2))) + pairs = ((0, 0), (0, 1), (0, 23), (5, 17), (8, 12), (23, 23)) + images = [ + ( + color.apply_color_connection(tensors[i], bra), + color.apply_color_connection(tensors[j], ket), + ) + for i, j in pairs + ] + actual = [color.contract_connected_tensor_nc_terms(a, b) for a, b in images] + monkeypatch.setattr(color, "_literal_tensor_overlap_nc_terms", _uncached_overlap) + expected = [color.contract_connected_tensor_nc_terms(a, b) for a, b in images] + assert actual == expected + assert any(len(polynomial) > 1 for polynomial in actual) From da9f175aeea4dadbcdd007d6d831f6a05455883d Mon Sep 17 00:00:00 2001 From: Lukas Simon Date: Sun, 25 Oct 2026 17:25:47 +0100 Subject: [PATCH 2/3] Stabilize native correlated contractions and add opt-in binary128 rescue --- docs/user/native-apis.md | 26 +- rust/crates/rusticol-capi/include/rusticol.h | 14 + rust/crates/rusticol-capi/src/correlated.rs | 103 +++++- .../rusticol-core/src/engine/correlated.rs | 329 ++++++++++++++++-- src/pyamplicol/api/services.py | 12 + src/pyamplicol/runtime/native_rescue.py | 262 ++++++++++++++ tests/unit/test_native_rescue.py | 103 ++++++ 7 files changed, 804 insertions(+), 45 deletions(-) create mode 100644 src/pyamplicol/runtime/native_rescue.py create mode 100644 tests/unit/test_native_rescue.py diff --git a/docs/user/native-apis.md b/docs/user/native-apis.md index 03e587902..721bb38f8 100644 --- a/docs/user/native-apis.md +++ b/docs/user/native-apis.md @@ -51,11 +51,35 @@ SDK and library, not an older installed release. Declare the operators with `Generator.generate(..., correlators=CorrelatorConfig(...))` or `generate ... --correlators FILE.json` as in [Born Correlations](../correlators.md). -Native correlated evaluation is **binary64 only**, like the ordinary native +Native correlated evaluation uses **binary64 buffers by default**, like the ordinary native SDK. It runs in RustiCol without Python or Symbolica. The separate Python correlated executor retains Symbolica-backed double-double and arbitrary precision; its `precision=16` path also uses that executor. +For cancellation-prone correlated requests, the development SDK can load an +explicit native binary128 companion. Preparation requires retained evaluators, +Symbolica and Linux GCC with quadmath: + +```python +runtime = Runtime.load("PROCESS_ARTIFACT") +runtime.compile_correlated_rescue("PROCESS.rescue.so") +``` + +Attach that trusted executable library with +`rusticol_runtime_load_correlated_rescue(handle, path)`, then use +`rusticol_runtime_evaluate_correlated_many_rescued_f64` with the same request +and binary64 buffer layout as `_many_f64`. The loader checks the artifact ID +and dimensions. Ordinary calls remain unchanged; no automatic rescue or +precision switching is implied. Convenience-language wrappers do not yet +expose these two C functions. + +The companion recomputes internal momenta, currents and amplitudes in binary128, +returns high/low binary64 amplitude components, and reuses the SDK's twofold +full-colour contraction. External sources and mutable model inputs remain +binary64. Numerical execution needs quadmath, not Python or Symbolica. Missing +kernels, incompatible artifacts and unsupported instruction programs fail +explicitly; small imaginary parts are not clamped. + LC, NLC and full colour are selected at generation for each inserted colour matrix, independently of connection order. The underlying amplitudes remain complete and full-colour: ordinary `evaluate`/`evaluate_resolved` still return diff --git a/rust/crates/rusticol-capi/include/rusticol.h b/rust/crates/rusticol-capi/include/rusticol.h index 36e08c40d..d6d556bb6 100644 --- a/rust/crates/rusticol-capi/include/rusticol.h +++ b/rust/crates/rusticol-capi/include/rusticol.h @@ -364,6 +364,20 @@ int rusticol_runtime_evaluate_correlated_many_f64( double *output, size_t output_capacity ); +/* Explicit native higher-precision current/amplitude mode. The companion + * is trusted executable input prepared for this exact artifact; no Symbolica + * or Python is loaded at numerical runtime. Sources/model inputs and returned + * buffers retain binary64 layout. No automatic fallback or component clamping. + */ +int rusticol_runtime_load_correlated_rescue(RusticolRuntimeHandle *handle, const char *path); +int rusticol_runtime_evaluate_correlated_many_rescued_f64( + RusticolRuntimeHandle *handle, + const double *momenta, size_t momentum_count, size_t point_count, + const RusticolCorrelatedRequest *requests, size_t request_count, + const char *const *helicity_ids, size_t helicity_count, + double *output, size_t output_capacity +); + /* Ordinary momenta use [point][external particle][E, px, py, pz]. */ int rusticol_runtime_evaluate_f64( RusticolRuntimeHandle *handle, diff --git a/rust/crates/rusticol-capi/src/correlated.rs b/rust/crates/rusticol-capi/src/correlated.rs index caab0efe1..34449ab2e 100644 --- a/rust/crates/rusticol-capi/src/correlated.rs +++ b/rust/crates/rusticol-capi/src/correlated.rs @@ -176,6 +176,88 @@ pub unsafe extern "C" fn rusticol_runtime_evaluate_correlated_many_f64( helicity_count: size_t, output: *mut c_double, output_capacity: size_t, +) -> c_int { + unsafe { + evaluate( + handle, + momenta, + momentum_count, + point_count, + requests, + request_count, + helicity_ids, + helicity_count, + output, + output_capacity, + false, + ) + } +} + +/// Attach a trusted native companion library to this exact process artifact. +/// +/// # Safety +/// handle must be live/exclusive; path must name trusted executable code implementing the rescue ABI. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn rusticol_runtime_load_correlated_rescue( + handle: *mut RusticolRuntimeHandle, + path: *const c_char, +) -> c_int { + guard(|| { + let handle = unsafe { required_handle_mut(handle) }?; + let path = unsafe { required_c_string(path, "rescue library path") }?; + unsafe { handle.runtime.load_correlated_rescue(Path::new(path)) }?; + Ok(()) + }) +} + +/// Evaluate using a previously attached binary128 current/amplitude kernel. +/// +/// # Safety +/// Pointer/storage requirements are identical to the binary64 correlated entry point. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn rusticol_runtime_evaluate_correlated_many_rescued_f64( + handle: *mut RusticolRuntimeHandle, + momenta: *const c_double, + momentum_count: size_t, + point_count: size_t, + requests: *const RusticolCorrelatedRequest, + request_count: size_t, + helicity_ids: *const *const c_char, + helicity_count: size_t, + output: *mut c_double, + output_capacity: size_t, +) -> c_int { + unsafe { + evaluate( + handle, + momenta, + momentum_count, + point_count, + requests, + request_count, + helicity_ids, + helicity_count, + output, + output_capacity, + true, + ) + } +} + +#[allow(clippy::too_many_arguments)] +unsafe fn evaluate( + handle: *mut RusticolRuntimeHandle, + momenta: *const c_double, + momentum_count: size_t, + point_count: size_t, + requests: *const RusticolCorrelatedRequest, + request_count: size_t, + helicity_ids: *const *const c_char, + helicity_count: size_t, + output: *mut c_double, + output_capacity: size_t, + rescue: bool, ) -> c_int { guard(|| { let required = request_count @@ -193,12 +275,21 @@ pub unsafe extern "C" fn rusticol_runtime_evaluate_correlated_many_f64( let requests = unsafe { read_requests(requests, request_count) }?; let helicities = unsafe { read_selector_ids(helicity_ids, helicity_count, "helicity IDs") }?; - let values = handle.runtime.evaluate_correlated_many_f64( - momenta, - point_count, - &requests, - helicities.as_deref(), - )?; + let values = if rescue { + handle.runtime.evaluate_correlated_many_rescued_f64( + momenta, + point_count, + &requests, + helicities.as_deref(), + )? + } else { + handle.runtime.evaluate_correlated_many_f64( + momenta, + point_count, + &requests, + helicities.as_deref(), + )? + }; if values.point_count != point_count || values.request_count != request_count || values.values.len().checked_mul(2) != Some(required) diff --git a/rust/crates/rusticol-core/src/engine/correlated.rs b/rust/crates/rusticol-core/src/engine/correlated.rs index 84fe83850..ac0cc8de8 100644 --- a/rust/crates/rusticol-core/src/engine/correlated.rs +++ b/rust/crates/rusticol-core/src/engine/correlated.rs @@ -98,7 +98,39 @@ struct WeightRecord { struct Matrix { id: String, sectors: BTreeSet, - entries: Vec<(i64, i64, [f64; 2])>, + entries: Vec<(i64, i64, [Twofold; 2])>, +} + +// Keep rounding residuals through coherent sums and cancellation-prone colour +// bilinears. The compiled amplitude stages themselves still execute in f64. +#[derive(Clone, Copy, Default, Debug)] +struct Twofold(f64, f64); + +impl Twofold { + fn add(self, rhs: Self) -> Self { + let sum = self.0 + rhs.0; + let virtual_rhs = sum - self.0; + let error = (self.0 - (sum - virtual_rhs)) + (rhs.0 - virtual_rhs); + let tail = error + self.1 + rhs.1; + let high = sum + tail; + Self(high, tail - (high - sum)) + } + + fn sub(self, rhs: Self) -> Self { + self.add(Self(-rhs.0, -rhs.1)) + } + + fn mul(self, rhs: Self) -> Self { + let product = self.0 * rhs.0; + let error = + self.0.mul_add(rhs.0, -product) + self.0 * rhs.1 + self.1 * rhs.0 + self.1 * rhs.1; + let high = product + error; + Self(high, error - (high - product)) + } + + fn value(self) -> f64 { + self.0 + self.1 + } } struct Group { @@ -115,6 +147,63 @@ pub(super) struct CorrelatedRuntime { known_helicities: BTreeSet, groups: Vec, vectors: NativeSpinCorrelationVectors, + #[cfg(feature = "f64-compiled")] + rescue: Option, +} + +#[cfg(feature = "f64-compiled")] +struct RescueKernel { + _library: libloading::Library, + evaluate: unsafe extern "C" fn(*const f64, *const f64, *mut f64) -> i32, +} + +#[cfg(feature = "f64-compiled")] +impl RescueKernel { + unsafe fn load( + path: &Path, + artifact_id: &str, + state: usize, + outputs: usize, + ) -> RusticolResult { + // Like a process artifact, the caller-supplied kernel is trusted executable input. + let library = unsafe { libloading::Library::new(path) } + .map_err(|error| invalid(format!("cannot load rescue kernel: {error}")))?; + unsafe { + let identity = library + .get:: *const std::ffi::c_char>( + b"rusticol_rescue_artifact_id", + ) + .map_err(|error| invalid(error.to_string()))?; + let pointer = identity(); + if pointer.is_null() + || std::ffi::CStr::from_ptr(pointer).to_bytes() != artifact_id.as_bytes() + { + return Err(invalid( + "rescue kernel belongs to a different process artifact", + )); + } + for (name, expected) in [ + (b"rusticol_rescue_state_size".as_slice(), state), + (b"rusticol_rescue_output_size".as_slice(), outputs), + ] { + let size = library + .get:: u64>(name) + .map_err(|error| invalid(error.to_string()))?; + if size() != expected as u64 { + return Err(invalid("rescue kernel dimensions differ from the artifact")); + } + } + let evaluate = *library + .get:: i32>( + b"rusticol_rescue_evaluate", + ) + .map_err(|error| invalid(error.to_string()))?; + Ok(Self { + _library: library, + evaluate, + }) + } + } } fn invalid(message: impl Into) -> RusticolError { @@ -137,7 +226,7 @@ fn validate_catalogue_artifact_id(loaded: &str, reopened: &str) -> RusticolResul Ok(()) } -fn rational_f64(parts: &[Value; 2]) -> RusticolResult { +fn rational_twofold(parts: &[Value; 2]) -> RusticolResult { let integer = |value: &Value| -> RusticolResult { let text = match value { Value::String(text) => text.clone(), @@ -157,11 +246,15 @@ fn rational_f64(parts: &[Value; 2]) -> RusticolResult { if denominator == 0.0 { return Err(malformed("zero matrix denominator")); } - let value = integer(&parts[0])? / denominator; + let numerator = integer(&parts[0])?; + let value = numerator / denominator; if !value.is_finite() { return Err(malformed("matrix coefficient exceeds binary64 range")); } - Ok(value) + Ok(Twofold( + value, + (-value).mul_add(denominator, numerator) / denominator, + )) } impl CorrelatedRuntime { @@ -310,8 +403,8 @@ impl CorrelatedRuntime { entry.left_sector_id, entry.right_sector_id, [ - rational_f64(&entry.weight.real)?, - rational_f64(&entry.weight.imag)?, + rational_twofold(&entry.weight.real)?, + rational_twofold(&entry.weight.imag)?, ], )); } @@ -447,6 +540,8 @@ impl CorrelatedRuntime { known_helicities, groups, vectors: BTreeMap::new(), + #[cfg(feature = "f64-compiled")] + rescue: None, }) } @@ -517,6 +612,62 @@ impl CorrelatedRuntime { } impl NativeRuntime { + /// Attach a trusted, preparation-compiled binary128 current/amplitude kernel. + /// + /// # Safety + /// The library must implement the documented rescue ABI and be trusted code. + pub unsafe fn load_correlated_rescue(&mut self, path: &Path) -> RusticolResult<()> { + self.ensure_correlated()?; + #[cfg(feature = "f64-compiled")] + { + let kernel = unsafe { + RescueKernel::load( + path, + &self.artifact_id, + self.runtime.parameter_count, + self.runtime.amplitude_output_count, + )? + }; + self.correlated.as_mut().expect("initialized").rescue = Some(kernel); + Ok(()) + } + #[cfg(not(feature = "f64-compiled"))] + { + let _ = path; + Err(invalid("rescue requires native compiled-stage support")) + } + } + + pub fn evaluate_correlated_many_rescued_f64( + &mut self, + momenta: &[f64], + point_count: usize, + requests: &[NativeCorrelatedRequest], + helicity_ids: Option<&[String]>, + ) -> RusticolResult { + self.ensure_correlated()?; + #[cfg(feature = "f64-compiled")] + { + if self + .correlated + .as_ref() + .expect("initialized") + .rescue + .is_none() + { + return Err(invalid("no native rescue kernel loaded")); + } + let mut evaluation = + evaluate_many(self, momenta, point_count, requests, helicity_ids, true)?; + self.apply_process_group_factor(evaluation.values.as_flattened_mut()); + Ok(evaluation) + } + #[cfg(not(feature = "f64-compiled"))] + { + let _ = (momenta, point_count, requests, helicity_ids); + Err(invalid("rescue requires native compiled-stage support")) + } + } fn ensure_correlated(&mut self) -> RusticolResult<()> { if self.correlated.is_none() { self.correlated = Some(CorrelatedRuntime::load(self)?); @@ -568,7 +719,8 @@ impl NativeRuntime { self.ensure_correlated()?; #[cfg(any(feature = "f64-compiled", feature = "f64-symjit"))] { - let mut evaluation = evaluate_many(self, momenta, point_count, requests, helicity_ids)?; + let mut evaluation = + evaluate_many(self, momenta, point_count, requests, helicity_ids, false)?; self.apply_process_group_factor(evaluation.values.as_flattened_mut()); Ok(evaluation) } @@ -589,6 +741,7 @@ fn evaluate_many( point_count: usize, requests: &[NativeCorrelatedRequest], helicity_ids: Option<&[String]>, + rescue: bool, ) -> RusticolResult { let expected = point_count .checked_mul(native.runtime.external_count) @@ -681,24 +834,83 @@ fn evaluate_many( // Memo lifetime is exactly one point tile: momenta and parameters cannot change. let mut memo = compiled_direct_prototype::CorrelatedStageMemo::default(); for (assignment, vectors) in assignments.iter().enumerate() { - direct.begin_tile_from_inputs( - batch.subview(point_start, point_stop)?, - &runtime.sources, - None, - runtime.external_count, - &runtime.particle_masses, - &runtime.momentum_slots, - &runtime.external_is_initial, - &runtime.model_parameter_values_f64, - )?; - direct.replace_correlated_sources(&runtime.sources, vectors, point_start, count)?; - if assignments.len() > 1 { - direct.evaluate_correlated_reusing(count, &mut memo)?; + let mut rescued = None; + if rescue { + #[cfg(feature = "f64-compiled")] + { + let kernel = correlated + .rescue + .as_ref() + .expect("checked before evaluation"); + let mut raw = + vec![[Twofold::default(); 2]; runtime.amplitude_output_count * count]; + for offset in 0..count { + let point_index = point_start + offset; + let mut state = vec![c64(0., 0.); runtime.parameter_count]; + ExecutionRuntime::fill_sources_row( + &runtime.sources, + runtime.external_count, + &runtime.particle_masses, + &mut state, + &batch.point(point_index), + )?; + ExecutionRuntime::fill_model_parameters_row( + runtime.value_parameter_count + runtime.momentum_parameter_count, + &runtime.model_parameter_values_f64, + &mut state, + )?; + for source in &runtime.sources { + if let Some(vectors) = vectors.get(&source.leg_label) { + let wave = + &vectors[if vectors.len() == 1 { 0 } else { point_index }]; + let [pr, pi] = source.applied_crossing.phase; + for (mu, &[vr, vi]) in wave.iter().enumerate() { + state[source.value_slot.component_start + mu] = + c64(vr * pr - vi * pi, vr * pi + vi * pr); + } + } + } + let mut output = vec![0.; 4 * runtime.amplitude_output_count]; + let status = unsafe { + (kernel.evaluate)( + state.as_ptr().cast::(), + momenta[4 * runtime.external_count * point_index..].as_ptr(), + output.as_mut_ptr(), + ) + }; + if status != 0 || output.iter().any(|value| !value.is_finite()) { + return Err(invalid( + "native rescue kernel failed or returned nonfinite amplitudes", + )); + } + for (root, value) in output.chunks_exact(4).enumerate() { + raw[root * count + offset] = + [Twofold(value[0], value[1]), Twofold(value[2], value[3])]; + } + } + rescued = Some(raw); + } + #[cfg(not(feature = "f64-compiled"))] + return Err(invalid("rescue requires compiled-stage support")); } else { - direct.evaluate_all(count)?; + direct.begin_tile_from_inputs( + batch.subview(point_start, point_stop)?, + &runtime.sources, + None, + runtime.external_count, + &runtime.particle_masses, + &runtime.momentum_slots, + &runtime.external_is_initial, + &runtime.model_parameter_values_f64, + )?; + direct.replace_correlated_sources(&runtime.sources, vectors, point_start, count)?; + if assignments.len() > 1 { + direct.evaluate_correlated_reusing(count, &mut memo)?; + } else { + direct.evaluate_all(count)?; + } } - let planes = direct.amplitude_planes()?; - let mut amplitudes = vec![[0.0; 2]; selected[assignment].len() * count]; + let mut amplitudes = vec![[Twofold::default(); 2]; selected[assignment].len() * count]; let mut by_helicity: BTreeMap<&str, BTreeMap> = BTreeMap::new(); for (position, &index) in selected[assignment].iter().enumerate() { let group = &correlated.groups[index]; @@ -707,10 +919,22 @@ fn evaluate_many( .or_default() .insert(group.sector, position); for &root in &group.roots { - let (re, im) = planes.plane_unchecked(root); + let planes = if rescued.is_none() { + Some(direct.amplitude_planes()?.plane_unchecked(root)) + } else { + None + }; for point in 0..count { - amplitudes[position * count + point][0] += re[point]; - amplitudes[position * count + point][1] += im[point]; + let value = if let Some(raw) = rescued.as_ref() { + raw[root * count + point] + } else { + let (re, im) = planes.expect("binary64 planes"); + [Twofold(re[point], 0.), Twofold(im[point], 0.)] + }; + amplitudes[position * count + point][0] = + amplitudes[position * count + point][0].add(value[0]); + amplitudes[position * count + point][1] = + amplitudes[position * count + point][1].add(value[1]); } } } @@ -719,6 +943,7 @@ fn evaluate_many( .enumerate() .filter(|(_, (id, _))| *id == assignment) { + let mut sums = vec![[Twofold::default(); 2]; count]; for sectors in by_helicity.values() { for &(left, right, weight) in &correlated.matrices[matrix].entries { let a = sectors[&left] * count; @@ -726,19 +951,20 @@ fn evaluate_many( for point in 0..count { let [ar, ai] = amplitudes[a + point]; let [br, bi] = amplitudes[b + point]; - let wr = ar * weight[0] + ai * weight[1]; - let wi = ar * weight[1] - ai * weight[0]; - let value = &mut values[request * point_count + point_start + point]; - value[0] += wr * br - wi * bi; - value[1] += wr * bi + wi * br; + let wr = ar.mul(weight[0]).add(ai.mul(weight[1])); + let wi = ar.mul(weight[1]).sub(ai.mul(weight[0])); + let value = &mut sums[point]; + value[0] = value[0].add(wr.mul(br).sub(wi.mul(bi))); + value[1] = value[1].add(wr.mul(bi).add(wi.mul(br))); } } } - for value in &mut values - [request * point_count + point_start..request * point_count + point_stop] - { - value[0] *= runtime.normalization_factor; - value[1] *= runtime.normalization_factor; + for (point, sum) in sums.into_iter().enumerate() { + values[request * point_count + point_start + point] = sum.map(|value| { + value + .mul(Twofold(runtime.normalization_factor, 0.0)) + .value() + }); } } } @@ -764,7 +990,9 @@ mod tests { #[test] fn exact_weights_reject_lossy_or_nonfinite_encodings() { assert_eq!( - rational_f64(&[Value::from("-3"), Value::from("2")]).unwrap(), + rational_twofold(&[Value::from("-3"), Value::from("2")]) + .unwrap() + .value(), -1.5 ); for parts in [ @@ -772,10 +1000,31 @@ mod tests { [Value::from(1.2), Value::from(1)], [Value::from(1), Value::from(0)], ] { - assert!(rational_f64(&parts).is_err()); + assert!(rational_twofold(&parts).is_err()); } } + #[test] + fn correlated_contraction_keeps_rounding_residuals() { + let large = Twofold(1e16, 0.0); + assert_eq!(large.add(Twofold(1.0, 0.0)).sub(large).value(), 1.0); + let third = rational_twofold(&[Value::from(1), Value::from(3)]).unwrap(); + assert!( + third + .mul(Twofold(3.0, 0.0)) + .sub(Twofold(1.0, 0.0)) + .value() + .abs() + < 1e-30 + ); + let a = Twofold(1.0 + f64::EPSILON, 0.0); + let b = Twofold(1.0 - f64::EPSILON, 0.0); + assert_eq!( + a.mul(b).sub(Twofold(1.0, 0.0)).value(), + -f64::EPSILON.powi(2) + ); + } + #[test] fn spin_validation_is_joint_finite_and_shape_checked() { let runtime = CorrelatedRuntime { @@ -785,6 +1034,8 @@ mod tests { known_helicities: BTreeSet::new(), groups: vec![], vectors: BTreeMap::new(), + #[cfg(feature = "f64-compiled")] + rescue: None, }; let vector = [[1.0, 0.0]; 4]; assert!( @@ -818,6 +1069,8 @@ mod tests { known_helicities: BTreeSet::new(), groups: vec![], vectors: BTreeMap::new(), + #[cfg(feature = "f64-compiled")] + rescue: None, }; for h1 in [-1, 1] { for h2 in [-1, 1] { diff --git a/src/pyamplicol/api/services.py b/src/pyamplicol/api/services.py index 87df2a81c..eea0fac7a 100644 --- a/src/pyamplicol/api/services.py +++ b/src/pyamplicol/api/services.py @@ -817,6 +817,18 @@ def evaluate_correlated_many( arithmetic=arithmetic, ) + def compile_correlated_rescue( + self, destination: os.PathLike[str] | str, *, compiler: str = "g++" + ) -> Path: + """Prepare a native binary128 current kernel for this exact artifact. + + Compilation uses retained Symbolica instructions; numerical execution + only needs the generated library and GCC's quadmath runtime. + """ + from pyamplicol.runtime.native_rescue import compile_correlated_rescue + + return compile_correlated_rescue(self, destination, compiler=compiler) + def clear(self) -> None: """Drop warmed execution state while keeping this artifact loaded. diff --git a/src/pyamplicol/runtime/native_rescue.py b/src/pyamplicol/runtime/native_rescue.py new file mode 100644 index 000000000..789c154c2 --- /dev/null +++ b/src/pyamplicol/runtime/native_rescue.py @@ -0,0 +1,262 @@ +# SPDX-License-Identifier: 0BSD +"""Preparation-only binary128 current kernels; numerical execution is Symbolica-free.""" + +from __future__ import annotations + +import json +import os +import subprocess +import tempfile +from pathlib import Path +from typing import TYPE_CHECKING, Any, cast + +from pyamplicol.api.errors import CompatibilityError +from pyamplicol.evaluators.symbolica_helpers import _symbolica_instruction_program +from pyamplicol.runtime.symbolica_exact import ( + _assign_stage_outputs, + _canonical_amplitude_outputs, + _pack_stage_inputs, +) + +if TYPE_CHECKING: + from pyamplicol.api.services import Runtime + + +def compile_correlated_rescue( + runtime: Runtime, destination: os.PathLike[str] | str, *, compiler: str = "g++" +) -> Path: + """Compile the retained complete current/amplitude plan, not new physics. + + Sources and mutable model inputs retain the native binary64 contract. + Internal momenta, currents and amplitudes execute in GCC binary128. + The SDK receives twofold amplitudes and applies its ordinary colour matrices. + """ + executor: Any = runtime._correlator_evaluator()._executor + executor._load_evaluators() + execution = executor._execution + schema = execution["runtime_schema"] + layout = schema["parameter_layout"] + state_size = max( + layout["parameter_count_if_flattened"], + layout["value_component_count"] + + layout["momentum_parameter_count"] + + layout["model_parameter_count"], + ) + stages = execution["compiled"]["stage_evaluators"] + helpers: list[str] = [] + statements: list[str] = [] + + def emit(evaluator: Any, inputs: Any) -> list[str]: + if evaluator.chunks: + return [ + value + for child, indices in zip( + evaluator.chunks, evaluator.chunk_input_indices, strict=True + ) + for value in emit(child, [inputs[i] for i in indices]) + ] + instructions, temporary_count, constants = _symbolica_instruction_program( + evaluator.evaluator.get_instructions() + ) + identifier = len(helpers) + name = f"leaf_{identifier}" + # Reserve the ID before emitting the body (each leaf has one definition). + helpers.append("") + exported = evaluator.evaluator.get_instructions() + output_count = exported.output_count + values = [value.evaluate({}, decimal_digit_precision=80) for value in constants] + body = [ + f"static std::vector {name}(const std::vector& p) {{", + "static const std::vector c={{{}}};".format( + ",".join( + f"q(strtoflt128({json.dumps(str(value.real))},nullptr)," + f"strtoflt128({json.dumps(str(value.imag))},nullptr))" + for value in values + ) + ), + f"std::vector t({temporary_count}),o({output_count});", + ] + written = set() + + def ref(slot: Any, destination: bool = False) -> str: + kind, index = slot + sizes = dict( + param=len(inputs), + temp=temporary_count, + const=len(constants), + out=output_count, + ) + if ( + kind not in sizes + or not isinstance(index, int) + or not 0 <= index < sizes[kind] + ): + raise CompatibilityError("Invalid rescue instruction reference") + if destination and kind not in ("temp", "out"): + raise CompatibilityError("Invalid rescue instruction destination") + if ( + not destination + and kind in ("temp", "out") + and (kind, index) not in written + ): + raise CompatibilityError("Uninitialized rescue instruction reference") + array = dict(param="p", temp="t", const="c", out="o")[kind] + return f"{array}[{index}]" + + for instruction in instructions: + op = instruction[0] + if op not in ("add", "mul", "pow", "assign", "fun"): + raise CompatibilityError( + "Unsupported native rescue instruction: " + repr(instruction) + ) + target = ref(instruction[1], True) + if op in ("add", "mul") and len(instruction) == 4: + operands = [ref(slot) for slot in instruction[2]] + expression = ("+" if op == "add" else "*").join(operands) + if not expression: + raise CompatibilityError("Empty rescue arithmetic instruction") + elif op == "pow" and len(instruction) == 5: + exponent = instruction[3] + if type(exponent) is not int or abs(exponent) >= 2**63: + raise CompatibilityError("Unsupported rescue integer exponent") + expression = f"power({ref(instruction[2])},{exponent}LL)" + elif op == "assign" and len(instruction) == 3: + expression = ref(instruction[2]) + elif op == "fun" and len(instruction) == 6 and not instruction[3]: + function_name = str(instruction[2]).split("::")[-1] + functions = dict( + conj="conjq", + sqrt="csqrtq", + exp="cexpq", + log="clogq", + sin="csinq", + cos="ccosq", + ) + if function_name not in functions or len(instruction[4]) != 1: + raise CompatibilityError( + "Unsupported native rescue function: " + repr(instruction) + ) + expression = f"{functions[function_name]}({ref(instruction[4][0])})" + else: + raise CompatibilityError( + "Unsupported native rescue instruction: " + repr(instruction) + ) + body.append(target + "=" + expression + ";") + written.add(tuple(instruction[1])) + if any(("out", i) not in written for i in range(output_count)): + raise CompatibilityError("Incomplete native rescue outputs") + body.extend(["return o;", "}"]) + helpers[identifier] = "\n".join(body) + variable = f"v{identifier}" + statements.append(f"auto {variable}={name}({{{','.join(inputs)}}});") + return [f"{variable}[{i}]" for i in range(output_count)] + + state: Any = [f"s[{i}]" for i in range(state_size)] + for stage, evaluator in zip( + stages["stages"], executor._stage_evaluators, strict=True + ): + outputs = emit(evaluator, _pack_stage_inputs(state, stage)) + updated: Any = list(state) + # These layout-only helpers also apply to symbolic C++ expressions. + _assign_stage_outputs(updated, cast(Any, outputs), stage) + statements.extend( + f"s[{i}]={value};" for i, value in enumerate(updated) if value != state[i] + ) + amplitude = stages["amplitude_stage"] + canonical: Any = _canonical_amplitude_outputs( + emit(executor._amplitude_evaluator, _pack_stage_inputs(state, amplitude)), + amplitude, + ) + momentum = [] + initial = { + particle["label"] + for particle in schema["external_particles"] + if particle["role"] == "initial" + } + for slot in schema["momentum_slots"]: + for component in range(4): + expression = "+".join( + ("-" if leg in initial else "") + f"R(p[{4 * (leg - 1) + component}])" + for leg in slot["external_labels"] + ) + index = ( + layout["value_component_count"] + slot["component_start"] + component + ) + momentum.append(f"s[{index}]=q({expression or '0'},0);") + source = ( + r""" +#include +#include +#include +using R=__float128; +using C=__complex128; +static C q(R real,R imag) { + C value; __real__ value=real; __imag__ value=imag; return value; +} +static C power(C base,long long exponent) { + bool inverse=exponent<0; + unsigned long long n=inverse ? -exponent : exponent; + C result=q(1,0); + while(n) { if(n&1) result*=base; n>>=1; if(n) base*=base; } + return inverse ? q(1,0)/result : result; +} +""" + + "\n".join(helpers) + + r""" +extern "C" const char* rusticol_rescue_artifact_id() { return ARTIFACT; } +extern "C" unsigned long long rusticol_rescue_state_size() { return STATE_SIZE; } +extern "C" unsigned long long rusticol_rescue_output_size() { return OUTPUT_SIZE; } +extern "C" int rusticol_rescue_evaluate( + const double* input,const double* p,double* output) { + try { + std::vector s(STATE_SIZE); + for(unsigned long long i=0;i Date: Sun, 25 Oct 2026 18:27:04 +0100 Subject: [PATCH 3/3] Fix rescue buffer typing for Symjit-only builds --- rust/crates/rusticol-core/src/engine/correlated.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/rust/crates/rusticol-core/src/engine/correlated.rs b/rust/crates/rusticol-core/src/engine/correlated.rs index ac0cc8de8..5a7e33466 100644 --- a/rust/crates/rusticol-core/src/engine/correlated.rs +++ b/rust/crates/rusticol-core/src/engine/correlated.rs @@ -834,7 +834,7 @@ fn evaluate_many( // Memo lifetime is exactly one point tile: momenta and parameters cannot change. let mut memo = compiled_direct_prototype::CorrelatedStageMemo::default(); for (assignment, vectors) in assignments.iter().enumerate() { - let mut rescued = None; + let rescued: Option>; if rescue { #[cfg(feature = "f64-compiled")] { @@ -909,6 +909,7 @@ fn evaluate_many( } else { direct.evaluate_all(count)?; } + rescued = None; } let mut amplitudes = vec![[Twofold::default(); 2]; selected[assignment].len() * count]; let mut by_helicity: BTreeMap<&str, BTreeMap> = BTreeMap::new();