Merge pull request #27 from alphal00p/asm_cff_support

Asm cff support

benches/cff_scalar.rs

@@ -1,11 +1,12 @@
-use std::path::PathBuf;
+use std::{env, path::PathBuf};
 
 use _gammaloop::{
     gammaloop_integrand::DefaultSample,
     graph::Graph,
     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

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'."""

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: []

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):

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': {

src/api/python.rs

@@ -280,19 +280,16 @@ impl PythonWorker {
         &mut self,
         export_root: &str,
         amplitude_names: Vec<String>,
-        compile_cff: bool,
-        compile_separate_orientations: bool,
+        export_yaml_str: &str,
     ) -> PyResult<String> {
+        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(&amplitude.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<String> {
-        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();

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::<T>(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::<T>(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)?;

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(

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::<f64>(
-                &params,
-                export_path,
-                compile_cff,
-                compile_separate_orientations,
-            ),
+            Some(cff) => {
+                cff.build_compiled_experssion::<f64>(&params, 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)
             }
         }
     }

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<String>,
+}
+
+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()
+        }
+    }
+}