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..5a7e33466 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,84 @@ 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 rescued: Option>; + 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)?; + } + rescued = None; } - 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 +920,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 +944,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 +952,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 +991,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 +1001,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 +1035,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 +1070,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/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/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 1 for polynomial in actual) diff --git a/tests/unit/test_native_rescue.py b/tests/unit/test_native_rescue.py new file mode 100644 index 000000000..7f54defef --- /dev/null +++ b/tests/unit/test_native_rescue.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: 0BSD +"""Small native code-generation regressions; no physics process is rebuilt.""" + +import ctypes +import shutil +import sys +import tempfile +import unittest +from decimal import Decimal, localcontext +from pathlib import Path +from types import SimpleNamespace + +from pyamplicol.api.errors import CompatibilityError +from pyamplicol.runtime.native_rescue import compile_correlated_rescue + + +class Constant: + def evaluate(self, *unused, **options): + return SimpleNamespace(real=Decimal(1) / Decimal(3), imag=Decimal(0)) + + +@unittest.skipUnless( + sys.platform.startswith("linux") and shutil.which("g++"), + "binary128 companions require GCC and quadmath", +) +class NativeRescueTests(unittest.TestCase): + def runtime(self, unsupported=False): + instructions = [ + ("add", ("temp", 0), [("param", 0), ("param", 1)], 2), + ("mul", ("temp", 1), [("param", 0), ("const", 0)], 2), + ("add", ("out", 0), [("temp", 0), ("temp", 1)], 2), + ] + if unsupported: + instructions.append(("goto", 0)) + program = SimpleNamespace( + instructions=instructions, + temporary_count=2, + constants=[Constant()], + output_count=1, + sub_evaluators=[], + ) + evaluator = SimpleNamespace( + chunks=(), evaluator=SimpleNamespace(get_instructions=lambda: program) + ) + stage = dict( + parameter_count=2, + output_length=1, + output_slots=[ + dict(output_start=0, output_stop=1, component_start=0, component_stop=1) + ], + ) + schema = dict( + parameter_layout=dict( + parameter_count_if_flattened=2, + value_component_count=2, + momentum_parameter_count=0, + model_parameter_count=0, + ), + model_parameters=[], + external_particles=[], + momentum_slots=[], + ) + executor = SimpleNamespace( + _load_evaluators=lambda: None, + _stage_evaluators=(), + _amplitude_evaluator=evaluator, + _execution=dict( + runtime_schema=schema, + compiled=dict(stage_evaluators=dict(stages=[], amplitude_stage=stage)), + ), + ) + return SimpleNamespace( + artifact_id="fixture-only", + _correlator_evaluator=lambda: SimpleNamespace(_executor=executor), + ) + + def test_native_kernel_keeps_sub_binary64_terms(self): + with tempfile.TemporaryDirectory() as directory, localcontext() as context: + context.prec = 80 + destination = Path(directory) / "kernel.so" + compile_correlated_rescue(self.runtime(), destination) + native = ctypes.CDLL(str(destination)) + pointer = ctypes.POINTER(ctypes.c_double) + native.rusticol_rescue_evaluate.argtypes = [pointer, pointer, pointer] + inputs = (ctypes.c_double * 4)(3.0, 0.0, -3.0, 0.0) + output = (ctypes.c_double * 4)() + self.assertEqual(native.rusticol_rescue_evaluate(inputs, None, output), 0) + value = Decimal.from_float(output[0]) + Decimal.from_float(output[1]) + self.assertLess(abs(value - 1), Decimal("1e-30")) + inputs[0], inputs[2] = 1e16, 1.0 + self.assertEqual(native.rusticol_rescue_evaluate(inputs, None, output), 0) + value = Decimal.from_float(output[0]) + Decimal.from_float(output[1]) + self.assertLess( + abs(value - (Decimal("1e16") * 4 / 3 + 1)), Decimal("1e-15") + ) + + def test_unsupported_program_cannot_replace_a_kernel(self): + with tempfile.TemporaryDirectory() as directory: + destination = Path(directory) / "kernel.so" + destination.write_bytes(b"prior kernel") + with self.assertRaises(CompatibilityError): + compile_correlated_rescue(self.runtime(unsupported=True), destination) + self.assertEqual(destination.read_bytes(), b"prior kernel")