diff --git a/benches/cff_scalar.rs b/benches/cff_scalar.rs index ad796a56..c9a6d588 100644 --- a/benches/cff_scalar.rs +++ b/benches/cff_scalar.rs @@ -1,4 +1,4 @@ -use std::path::PathBuf; +use std::{env, path::PathBuf}; use _gammaloop::{ gammaloop_integrand::DefaultSample, @@ -6,6 +6,7 @@ use _gammaloop::{ momentum::{FourMomentum, ThreeMomentum}, tests_from_pytest::load_amplitude_output, utils::F, + ExportSettings, GammaloopCompileOptions, }; use criterion::{criterion_group, criterion_main, Criterion}; const COMPILED_DUMP: &str = "TMP_COMPILED"; @@ -45,10 +46,23 @@ fn load_helper(path: &str) -> Graph { let (_, mut amplitude) = load_amplitude_output(path, true); amplitude.amplitude_graphs[0].graph.generate_cff(); + let export_settings = ExportSettings { + compile_cff: true, + compile_separate_orientations: false, + gammaloop_compile_options: GammaloopCompileOptions { + use_asm: env::var("USE_ASM").is_ok(), + optimization_level: 3, + fast_math: true, + unsafe_math: true, + compiler: "g++".to_string(), + custom: vec![], + }, + }; + let true_path = PathBuf::from(COMPILED_DUMP).join(path); amplitude.amplitude_graphs[0] .graph - .build_compiled_expression(true_path, true, false) + .build_compiled_expression(true_path, &export_settings) .unwrap(); amplitude.amplitude_graphs.remove(0).graph diff --git a/python/gammaloop/_gammaloop.pyi b/python/gammaloop/_gammaloop.pyi index 19fd771c..a9bc5bf4 100644 --- a/python/gammaloop/_gammaloop.pyi +++ b/python/gammaloop/_gammaloop.pyi @@ -57,8 +57,8 @@ class Worker: def export_cross_sections(self, export_root: str, cross_section_names: list[str]) -> None: """ Exports the cross sections given in argument to the export root given in argument. """ - def export_amplitudes(self, export_root: str, amplitude_names: list[str], compiled_cff: bool, compile_separate_orientations: bool) -> None: - """ Exports the amplitudes given in argument to the export root given in argument, use compiled_cff to generate cpp code for the cff expression, possibly slow""" + def export_amplitudes(self, export_root: str, amplitude_names: list[str], export_yaml_str: str) -> None: + """ Exports the amplitudes given in argument to the export root given in argument, parse export settings as yaml str""" def export_expressions(self, export_root: str, format: str) -> None: """Exports the numerator and denominator to the export root given in argument in the format which can be 'default' or 'mathematica' or 'latex'.""" diff --git a/python/gammaloop/data/config/gammaloop_config.yaml b/python/gammaloop/data/config/gammaloop_config.yaml index a0bac539..b8bc408b 100644 --- a/python/gammaloop/data/config/gammaloop_config.yaml +++ b/python/gammaloop/data/config/gammaloop_config.yaml @@ -30,6 +30,12 @@ drawing: anchor_tension: 2.0 caption_size: "40pt" export_settings: - write_default_settings: false compile_cff: true compile_separate_orientations: false + gammaloop_compile_options: + use_asm: false + optimization_level: 3 + fast_math: true + unsafe_math: true + compiler: "g++" + custom: [] diff --git a/python/gammaloop/exporters/exporters.py b/python/gammaloop/exporters/exporters.py index 32f39805..7b174bca 100644 --- a/python/gammaloop/exporters/exporters.py +++ b/python/gammaloop/exporters/exporters.py @@ -285,10 +285,7 @@ def export(self, export_root: Path, amplitudes: AmplitudeList): # Now address the rust export aspect self.gammaloop.rust_worker.export_amplitudes( - str(export_root), [amp.name for amp in amplitudes], self.gammaloop.config['export_settings']['compile_cff'], self.gammaloop.config['export_settings']['compile_separate_orientations']) - - if self.gammaloop.config['export_settings']['write_default_settings']: - self.gammaloop.rust_worker.write_default_settings(str(export_root)) + str(export_root), [amp.name for amp in amplitudes], yaml.dump(self.gammaloop.config['export_settings'])) class CrossSectionsExporter(GammaLoopExporter): diff --git a/python/gammaloop/interface/gammaloop_interface.py b/python/gammaloop/interface/gammaloop_interface.py index 13beb0f8..906b8702 100644 --- a/python/gammaloop/interface/gammaloop_interface.py +++ b/python/gammaloop/interface/gammaloop_interface.py @@ -85,9 +85,16 @@ def __init__(self, path: str | None = None, quiet=False): } }, 'export_settings': { - 'write_default_settings': False, 'compile_cff': True, 'compile_separate_orientations': False, + 'gammaloop_compile_options': { + 'use_asm': False, + 'optimization_level': 3, + 'fast_math': True, + 'unsafe_math': True, + 'compiler': "g++", + 'custom': [], + }, }, 'run_settings': { 'General': { diff --git a/src/api/python.rs b/src/api/python.rs index 5435ec48..9e82cb73 100644 --- a/src/api/python.rs +++ b/src/api/python.rs @@ -280,19 +280,16 @@ impl PythonWorker { &mut self, export_root: &str, amplitude_names: Vec, - compile_cff: bool, - compile_separate_orientations: bool, + export_yaml_str: &str, ) -> PyResult { + let export_settings = serde_yaml::from_str(export_yaml_str) + .map_err(|e| exceptions::PyException::new_err(e.to_string()))?; + let mut n_exported: usize = 0; for amplitude in self.amplitudes.container.iter_mut() { if amplitude_names.contains(&litude.name.to_string()) { n_exported += 1; - let res = amplitude.export( - export_root, - &self.model, - compile_cff, - compile_separate_orientations, - ); + let res = amplitude.export(export_root, &self.model, &export_settings); if let Err(err) = res { return Err(exceptions::PyException::new_err(err.to_string())); } @@ -566,19 +563,6 @@ impl PythonWorker { Ok(format!("Processed job {}", job_id)) } - pub fn write_default_settings(&self, path: &str) -> PyResult { - let default = Settings::default(); - let default_string = serde_yaml::to_string(&default) - .map_err(|e| exceptions::PyException::new_err(e.to_string()))?; - - let path = Path::new(path).join("cards").join("run_card.yaml"); - - fs::write(path, default_string) - .map_err(|e| exceptions::PyException::new_err(e.to_string()))?; - - Ok("Wrote default settings file".to_string()) - } - pub fn display_master_node_status(&self) { if let Some(master_node) = &self.master_node { master_node.display_status(); diff --git a/src/cff/expression.rs b/src/cff/expression.rs index ff8c648a..3a81431e 100644 --- a/src/cff/expression.rs +++ b/src/cff/expression.rs @@ -1,6 +1,6 @@ use crate::{ utils::{FloatLike, VarFloat, F}, - Settings, + ExportSettings, Settings, }; use color_eyre::Report; use derive_more::{From, Into}; @@ -12,9 +12,7 @@ use std::{fmt::Debug, ops::Index, path::PathBuf}; use symbolica::{ atom::{Atom, AtomView}, domains::{float::NumericalFloatLike, rational::Rational}, - evaluate::{ - CompileOptions, CompiledEvaluator, EvalTree, ExportedCode, ExpressionEvaluator, FunctionMap, - }, + evaluate::{CompiledEvaluator, EvalTree, ExportedCode, ExpressionEvaluator, FunctionMap}, }; use typed_index_collections::TiVec; @@ -533,10 +531,9 @@ impl CFFExpression { &mut self, params: &[Atom], path: PathBuf, - compile_cff: bool, - compile_separate_orientations: bool, + export_settings: &ExportSettings, ) -> Result<(), Report> { - if !compile_cff && !compile_separate_orientations { + if !export_settings.compile_cff && export_settings.compile_separate_orientations { return Ok(()); } @@ -564,32 +561,51 @@ impl CFFExpression { .to_str() .ok_or(eyre!("could not convert path to string"))?; - if compile_cff { + if export_settings.compile_cff { let joint = self.build_joint_symbolica_evaluator::(params); - cpp_str.push_str(&joint.export_cpp_str("joint", true)); + + let source_string = if export_settings.gammaloop_compile_options.use_asm { + joint.export_asm_str("joint", true) + } else { + joint.export_cpp_str("joint", true) + }; + + cpp_str.push_str(&source_string); } - if compile_separate_orientations { + if export_settings.compile_separate_orientations { let orientations = self.build_symbolica_evaluators::(params); for (orientation_id, orientation_evaluator) in orientations.into_iter().enumerate() { - let orientation_cpp_str = orientation_evaluator.export_cpp_str( - &format!("orientation_{}", orientation_id), - !compile_cff && orientation_id == 0, - ); + let orientation_source_str = if export_settings.gammaloop_compile_options.use_asm { + orientation_evaluator.export_asm_str( + &format!("orientation_{}", orientation_id), + !export_settings.compile_cff && orientation_id == 0, + ) + } else { + orientation_evaluator.export_cpp_str( + &format!("orientation_{}", orientation_id), + !export_settings.compile_cff && orientation_id == 0, + ) + }; - cpp_str.push_str(&orientation_cpp_str); + cpp_str.push_str(&orientation_source_str); } } std::fs::write(path_to_code, cpp_str)?; let exported_code = ExportedCode::new(path_to_code_str.to_string(), "joint".to_string()); - exported_code.compile(path_to_so_str, CompileOptions::default())?; + exported_code.compile( + path_to_so_str, + export_settings + .gammaloop_compile_options + .to_symbolica_compile_options(), + )?; let metadata = CompiledCFFExpressionMetaData { name: path_to_compiled, num_orientations: self.get_num_trees(), - compile_cff_present: compile_cff, - compile_separate_orientations_present: compile_separate_orientations, + compile_cff_present: export_settings.compile_cff, + compile_separate_orientations_present: export_settings.compile_separate_orientations, }; self.compiled = CompiledCFFExpression::from_metedata(metadata)?; diff --git a/src/cross_section.rs b/src/cross_section.rs index 2a33c5a9..7485c616 100644 --- a/src/cross_section.rs +++ b/src/cross_section.rs @@ -1,7 +1,7 @@ use crate::gammaloop_integrand::GammaLoopIntegrand; use crate::graph::{Graph, SerializableGraph}; use crate::model::Model; -use crate::{utils::*, Settings}; +use crate::{utils::*, ExportSettings, Settings}; use bincode; use color_eyre::{Help, Report}; #[allow(unused_imports)] @@ -592,8 +592,7 @@ impl Amplitude { &mut self, export_root: &str, model: &Model, - compile_cff: bool, - compile_separate_orientations: bool, + export_settings: &ExportSettings, ) -> Result<(), Report> { // TODO process amplitude by adding lots of additional information necessary for runtime. // e.g. generate e-surface, cff expression, counterterms, etc. @@ -625,11 +624,9 @@ impl Amplitude { // dump the derived data in a binary file for amplitude_graph in self.amplitude_graphs.iter_mut() { - amplitude_graph.graph.build_compiled_expression( - path.clone(), - compile_cff, - compile_separate_orientations, - )?; + amplitude_graph + .graph + .build_compiled_expression(path.clone(), export_settings)?; debug!("dumping derived data"); fs::write( diff --git a/src/graph.rs b/src/graph.rs index 9590d371..11f4f213 100644 --- a/src/graph.rs +++ b/src/graph.rs @@ -21,7 +21,7 @@ use crate::{ compute_four_momentum_from_three, compute_three_momentum_from_four, sorted_vectorize, FloatLike, F, }, - Settings, + ExportSettings, Settings, }; use ahash::RandomState; @@ -1617,24 +1617,16 @@ impl Graph { pub fn build_compiled_expression( &mut self, export_path: PathBuf, - compile_cff: bool, - compile_separate_orientations: bool, + export_settings: &ExportSettings, ) -> Result<(), Report> { let params = self.build_params_for_cff(); match self.derived_data.cff_expression.as_mut() { - Some(cff) => cff.build_compiled_experssion::( - ¶ms, - export_path, - compile_cff, - compile_separate_orientations, - ), + Some(cff) => { + cff.build_compiled_experssion::(¶ms, export_path, export_settings) + } None => { self.generate_cff(); - self.build_compiled_expression( - export_path, - compile_cff, - compile_separate_orientations, - ) + self.build_compiled_expression(export_path, export_settings) } } } diff --git a/src/lib.rs b/src/lib.rs index 79684c78..8585ff33 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -42,6 +42,7 @@ use observables::ObservableSettings; use observables::PhaseSpaceSelectorSettings; use std::fs::File; use std::sync::atomic::AtomicBool; +use symbolica::evaluate::CompileOptions; use utils::FloatLike; use utils::F; @@ -394,3 +395,34 @@ pub struct SubtractionSettings { pub dynamic_sliver: bool, pub integrated_ct_hfunction: HFunctionSettings, } + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ExportSettings { + pub compile_cff: bool, + pub compile_separate_orientations: bool, + pub gammaloop_compile_options: GammaloopCompileOptions, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct GammaloopCompileOptions { + pub use_asm: bool, + pub optimization_level: usize, + pub fast_math: bool, + pub unsafe_math: bool, + pub compiler: String, + pub custom: Vec, +} + +impl GammaloopCompileOptions { + #[allow(clippy::needless_update)] + pub fn to_symbolica_compile_options(&self) -> CompileOptions { + CompileOptions { + optimization_level: self.optimization_level, + fast_math: self.fast_math, + unsafe_math: self.unsafe_math, + compiler: self.compiler.clone(), + custom: self.custom.clone(), + ..CompileOptions::default() + } + } +}