From 4ba9e1609163c0d760f2a716244729d3116cc5ee Mon Sep 17 00:00:00 2001
From: Julian Gonzalez Calderon <gonzalezcalderonjulian@gmail.com>
Date: Fri, 18 Oct 2024 11:27:35 -0300
Subject: [PATCH] Log contract compilation time (#823)

* Add mlir compilation time

* Add time to llvm compilation and object linking

* Use finer grained timers

* Change info to trace
---
 src/context.rs | 16 +++++++++++++++-
 src/ffi.rs     | 25 +++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/src/context.rs b/src/context.rs
index c6b249219..f868c0faa 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -33,7 +33,8 @@ use mlir_sys::{
     mlirLLVMDIModuleAttrGet, MlirLLVMDIEmissionKind_MlirLLVMDIEmissionKindFull,
     MlirLLVMDINameTableKind_MlirLLVMDINameTableKindDefault,
 };
-use std::sync::OnceLock;
+use std::{sync::OnceLock, time::Instant};
+use tracing::trace;
 
 /// Context of IRs, dialects and passes for Cairo programs compilation.
 #[derive(Debug, Eq, PartialEq)]
@@ -70,6 +71,9 @@ impl NativeContext {
         program: &Program,
         ignore_debug_names: bool,
     ) -> Result<NativeModule, Error> {
+        trace!("starting sierra to mlir compilation");
+        let pre_sierra_compilation_instant = Instant::now();
+
         static INITIALIZED: OnceLock<()> = OnceLock::new();
         INITIALIZED.get_or_init(|| unsafe {
             LLVM_InitializeAllTargets();
@@ -181,6 +185,12 @@ impl NativeContext {
             ignore_debug_names,
         )?;
 
+        let sierra_compilation_time = pre_sierra_compilation_instant.elapsed().as_millis();
+        trace!(
+            time = sierra_compilation_time,
+            "sierra to mlir compilation finished"
+        );
+
         if let Ok(x) = std::env::var("NATIVE_DEBUG_DUMP") {
             if x == "1" || x == "true" {
                 std::fs::write("dump-prepass.mlir", module.as_operation().to_string())
@@ -202,7 +212,11 @@ impl NativeContext {
             }
         }
 
+        trace!("starting mlir passes");
+        let pre_passes_instant = Instant::now();
         run_pass_manager(&self.context, &mut module)?;
+        let passes_time = pre_passes_instant.elapsed().as_millis();
+        trace!(time = passes_time, "mlir passes finished");
 
         if let Ok(x) = std::env::var("NATIVE_DEBUG_DUMP") {
             if x == "1" || x == "true" {
diff --git a/src/ffi.rs b/src/ffi.rs
index 09cccbbcf..0dbb8aa1d 100644
--- a/src/ffi.rs
+++ b/src/ffi.rs
@@ -36,8 +36,10 @@ use std::{
     path::Path,
     ptr::{addr_of_mut, null_mut},
     sync::OnceLock,
+    time::Instant,
 };
 use tempfile::NamedTempFile;
+use tracing::trace;
 
 /// For any `!llvm.struct<...>` type, return the MLIR type of the field at the requested index.
 pub fn get_struct_field_type_at<'c>(r#type: &Type<'c>, index: usize) -> Type<'c> {
@@ -110,7 +112,11 @@ pub fn module_to_object(module: &Module<'_>, opt_level: OptLevel) -> Result<Vec<
 
         let op = module.as_operation().to_raw();
 
+        trace!("starting mlir to llvm compilation");
+        let pre_mlir_instant = Instant::now();
         let llvm_module = mlirTranslateModuleToLLVMIR(op, llvm_context as *mut _) as *mut _;
+        let mlir_time = pre_mlir_instant.elapsed().as_millis();
+        trace!(time = mlir_time, "mlir to llvm finished");
 
         let mut null = null_mut();
         let mut error_buffer = addr_of_mut!(null);
@@ -160,7 +166,13 @@ pub fn module_to_object(module: &Module<'_>, opt_level: OptLevel) -> Result<Vec<
             OptLevel::Aggressive => 3,
         };
         let passes = CString::new(format!("default<O{opt}>")).unwrap();
+
+        trace!("starting llvm passes");
+        let pre_passes_instant = Instant::now();
         let error = LLVMRunPasses(llvm_module, passes.as_ptr(), machine, opts);
+        let passes_time = pre_passes_instant.elapsed().as_millis();
+        trace!(time = passes_time, "llvm passes finished");
+
         if !error.is_null() {
             let msg = LLVMGetErrorMessage(error);
             let msg = CStr::from_ptr(msg);
@@ -171,6 +183,8 @@ pub fn module_to_object(module: &Module<'_>, opt_level: OptLevel) -> Result<Vec<
 
         let mut out_buf: MaybeUninit<LLVMMemoryBufferRef> = MaybeUninit::uninit();
 
+        trace!("starting llvm to object compilation");
+        let pre_llvm_compilation_instant = Instant::now();
         let ok = LLVMTargetMachineEmitToMemoryBuffer(
             machine,
             llvm_module,
@@ -178,6 +192,11 @@ pub fn module_to_object(module: &Module<'_>, opt_level: OptLevel) -> Result<Vec<
             error_buffer,
             out_buf.as_mut_ptr(),
         );
+        let llvm_compilation_time = pre_llvm_compilation_instant.elapsed().as_millis();
+        trace!(
+            time = llvm_compilation_time,
+            "llvm to object compilation finished"
+        );
 
         if ok != 0 {
             let error = CStr::from_ptr(*error_buffer);
@@ -294,7 +313,13 @@ pub fn object_to_shared_lib(object: &[u8], output_filename: &Path) -> Result<()>
     };
 
     let mut linker = std::process::Command::new("ld");
+
+    trace!("starting linking");
+    let pre_linking_instant = Instant::now();
     let proc = linker.args(args.iter().map(|x| x.as_ref())).output()?;
+    let linking_time = pre_linking_instant.elapsed().as_millis();
+    trace!(time = linking_time, "linking finished");
+
     if proc.status.success() {
         Ok(())
     } else {