Use finer grained timers

lambdaclass · Sep 30, 2024 · 490b096 · 490b096
1 parent 8bde303
commit 490b096
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 10 deletions.
diff --git 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::info;
 
 /// 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> {
+        info!("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();
+        info!(
+            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 {
             }
         }
 
+        info!("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();
+        info!(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/executor/contract.rs b/src/executor/contract.rs
@@ -106,12 +106,8 @@ impl AotContractExecutor {
     /// a temporary file an deleted when dropped.
     /// If you loaded a ContractExecutor using [`load`] then it will not be treated as a temp file.
     pub fn new(sierra_program: &Program, opt_level: OptLevel) -> Result<Self> {
-        info!("starting mlir compilation");
-        let pre_mlir_compilation_instant = Instant::now();
         let native_context = NativeContext::new();
         let module = native_context.compile(sierra_program, true)?;
-        let mlir_compilation_time = pre_mlir_compilation_instant.elapsed().as_millis();
-        info!(time = mlir_compilation_time, "mlir compilation finished");
 
         let NativeModule {
             module,
@@ -175,11 +171,7 @@ impl AotContractExecutor {
         let llvm_compilation_time = pre_llvm_compilation_instant.elapsed().as_millis();
         info!(time = llvm_compilation_time, "llvm compilation finished");
 
-        info!("starting object linking");
-        let pre_object_linking_instant = Instant::now();
         crate::object_to_shared_lib(&object_data, &library_path)?;
-        let object_linking_time = pre_object_linking_instant.elapsed().as_millis();
-        info!(time = object_linking_time, "object linking finished");
 
         Ok(Self {
             library: Arc::new(unsafe { Library::new(&library_path)? }),

diff --git a/src/ffi.rs b/src/ffi.rs
@@ -35,8 +35,10 @@ use std::{
     path::Path,
     ptr::{addr_of_mut, null_mut},
     sync::OnceLock,
+    time::Instant,
 };
 use tempfile::NamedTempFile;
+use tracing::info;
 
 /// 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> {
@@ -109,7 +111,11 @@ pub fn module_to_object(module: &Module<'_>, opt_level: OptLevel) -> Result<Vec<
 
         let op = module.as_operation().to_raw();
 
+        info!("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();
+        info!(time = mlir_time, "mlir to llvm finished");
 
         let mut null = null_mut();
         let mut error_buffer = addr_of_mut!(null);
@@ -156,7 +162,13 @@ pub fn module_to_object(module: &Module<'_>, opt_level: OptLevel) -> Result<Vec<
             OptLevel::Aggressive => 1, // https://github.com/llvm/llvm-project/issues/107198
         };
         let passes = CString::new(format!("default<O{opt}>")).unwrap();
+
+        info!("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();
+        info!(time = passes_time, "llvm passes finished");
+
         if !error.is_null() {
             let msg = LLVMGetErrorMessage(error);
             let msg = CStr::from_ptr(msg);
@@ -167,13 +179,20 @@ pub fn module_to_object(module: &Module<'_>, opt_level: OptLevel) -> Result<Vec<
 
         let mut out_buf: MaybeUninit<LLVMMemoryBufferRef> = MaybeUninit::uninit();
 
+        info!("starting llvm to object compilation");
+        let pre_llvm_compilation_instant = Instant::now();
         let ok = LLVMTargetMachineEmitToMemoryBuffer(
             machine,
             llvm_module,
             LLVMCodeGenFileType::LLVMObjectFile,
             error_buffer,
             out_buf.as_mut_ptr(),
         );
+        let llvm_compilation_time = pre_llvm_compilation_instant.elapsed().as_millis();
+        info!(
+            time = llvm_compilation_time,
+            "llvm to object compilation finished"
+        );
 
         if ok != 0 {
             let error = CStr::from_ptr(*error_buffer);
@@ -277,9 +296,14 @@ pub fn object_to_shared_lib(object: &[u8], output_filename: &Path) -> Result<()>
             unimplemented!()
         }
     };
-
     let mut linker = std::process::Command::new("ld");
+
+    info!("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();
+    info!(time = linking_time, "linking finished");
+
     if proc.status.success() {
         Ok(())
     } else {