issue/83: function-level memory profiling

Author: morganthomas

crates/lean_compiler/src/b_compile_intermediate.rs

@@ -10,7 +10,7 @@ use utils::ToUsize;
 #[derive(Default)]
 struct Compiler {
     bytecode: BTreeMap<Label, Vec<IntermediateInstruction>>,
-    match_blocks: Vec<Vec<Vec<IntermediateInstruction>>>, // each match = many bytecode blocks, each bytecode block = many instructions
+    match_blocks: Vec<MatchBlock>,
     if_counter: usize,
     call_counter: usize,
     func_name: String,
@@ -152,10 +152,17 @@ fn compile_function(
     compiler.args_count = function.arguments.len();
 
     let mut declared_vars: BTreeSet<Var> = function.arguments.iter().cloned().collect();
-    compile_lines(&function.instructions, compiler, None, &mut declared_vars)
+    compile_lines(
+        &Label::function(function.name.clone()),
+        &function.instructions,
+        compiler,
+        None,
+        &mut declared_vars,
+    )
 }
 
 fn compile_lines(
+    function_name: &Label,
     lines: &[SimpleLine],
     compiler: &mut Compiler,
     final_jump: Option<Label>,
@@ -212,6 +219,7 @@ fn compile_lines(
                     let mut arm_declared_vars = declared_vars.clone();
                     compiler.stack_size = original_stack_size;
                     let arm_instructions = compile_lines(
+                        function_name,
                         arm,
                         compiler,
                         Some(end_label.clone()),
@@ -229,7 +237,10 @@ fn compile_lines(
                     };
                 }
                 compiler.stack_size = new_stack_size;
-                compiler.match_blocks.push(compiled_arms);
+                compiler.match_blocks.push(MatchBlock {
+                    function_name: function_name.clone(),
+                    match_cases: compiled_arms,
+                });
 
                 let value_scaled_offset = IntermediateValue::MemoryAfterFp {
                     offset: compiler.stack_size.into(),
@@ -261,8 +272,13 @@ fn compile_lines(
                     updated_fp: None,
                 });
 
-                let remaining =
-                    compile_lines(&lines[i + 1..], compiler, final_jump, declared_vars)?;
+                let remaining = compile_lines(
+                    function_name,
+                    &lines[i + 1..],
+                    compiler,
+                    final_jump,
+                    declared_vars,
+                )?;
                 compiler.bytecode.insert(end_label, remaining);
 
                 return Ok(instructions);
@@ -349,6 +365,7 @@ fn compile_lines(
 
                 let mut then_declared_vars = declared_vars.clone();
                 let then_instructions = compile_lines(
+                    function_name,
                     then_branch,
                     compiler,
                     Some(end_label.clone()),
@@ -359,6 +376,7 @@ fn compile_lines(
                 compiler.stack_size = original_stack;
                 let mut else_declared_vars = declared_vars.clone();
                 let else_instructions = compile_lines(
+                    function_name,
                     else_branch,
                     compiler,
                     Some(end_label.clone()),
@@ -375,8 +393,13 @@ fn compile_lines(
                 compiler.bytecode.insert(if_label, then_instructions);
                 compiler.bytecode.insert(else_label, else_instructions);
 
-                let remaining =
-                    compile_lines(&lines[i + 1..], compiler, final_jump, declared_vars)?;
+                let remaining = compile_lines(
+                    function_name,
+                    &lines[i + 1..],
+                    compiler,
+                    final_jump,
+                    declared_vars,
+                )?;
                 compiler.bytecode.insert(end_label, remaining);
 
                 return Ok(instructions);
@@ -399,7 +422,7 @@ fn compile_lines(
             }
 
             SimpleLine::FunctionCall {
-                function_name,
+                function_name: callee_function_name,
                 args,
                 return_data,
             } => {
@@ -411,7 +434,7 @@ fn compile_lines(
                 compiler.stack_size += 1;
 
                 instructions.extend(setup_function_call(
-                    function_name,
+                    callee_function_name,
                     args,
                     new_fp_pos,
                     &return_label,
@@ -436,6 +459,7 @@ fn compile_lines(
                     }
 
                     instructions.extend(compile_lines(
+                        function_name,
                         &lines[i + 1..],
                         compiler,
                         final_jump,

crates/lean_compiler/src/c_compile_final.rs

@@ -52,47 +52,71 @@ pub fn compile_to_low_level_bytecode(
         .get("main")
         .unwrap();
 
+    let mut hints = BTreeMap::new();
     let mut label_to_pc = BTreeMap::new();
     label_to_pc.insert(Label::function("main"), 0);
     let entrypoint = intermediate_bytecode
         .bytecode
         .remove(&Label::function("main"))
         .ok_or("No main function found in the compiled program")?;
+    hints.insert(
+        0,
+        vec![Hint::StackFrame {
+            label: Label::function("main"),
+            size: starting_frame_memory,
+        }],
+    );
 
-    let mut code_blocks = vec![(0, entrypoint.clone())];
+    let mut code_blocks = vec![(Label::function("main"), 0, entrypoint.clone())];
     let mut pc = count_real_instructions(&entrypoint);
+
     for (label, instructions) in &intermediate_bytecode.bytecode {
         label_to_pc.insert(label.clone(), pc);
-        code_blocks.push((pc, instructions.clone()));
+        if let Label::Function(function_name) = label {
+            hints
+                .entry(pc)
+                .or_insert_with(Vec::new)
+                .push(Hint::StackFrame {
+                    label: label.clone(),
+                    size: *intermediate_bytecode
+                        .memory_size_per_function
+                        .get(function_name)
+                        .unwrap(),
+                });
+        }
+        code_blocks.push((label.clone(), pc, instructions.clone()));
         pc += count_real_instructions(instructions);
     }
 
     let ending_pc = label_to_pc.get(&Label::EndProgram).copied().unwrap();
 
     let mut match_block_sizes = Vec::new();
     let mut match_first_block_starts = Vec::new();
-    for match_statement in intermediate_bytecode.match_blocks {
-        let max_block_size = match_statement
+    for MatchBlock {
+        function_name,
+        match_cases,
+    } in intermediate_bytecode.match_blocks
+    {
+        let max_block_size = match_cases
             .iter()
             .map(|block| count_real_instructions(block))
             .max()
             .unwrap();
         match_first_block_starts.push(pc);
         match_block_sizes.push(max_block_size);
 
-        for mut block in match_statement {
+        for mut block in match_cases {
             // fill the end of block with unreachable instructions
             block.extend(vec![
                 IntermediateInstruction::Panic;
                 max_block_size - count_real_instructions(&block)
             ]);
-            code_blocks.push((pc, block));
+            code_blocks.push((function_name.clone(), pc, block));
             pc += max_block_size;
         }
     }
 
     let mut low_level_bytecode = Vec::new();
-    let mut hints = BTreeMap::new();
 
     for (label, pc) in label_to_pc.clone() {
         hints
@@ -108,9 +132,10 @@ pub fn compile_to_low_level_bytecode(
         match_first_block_starts,
     };
 
-    for (pc_start, block) in code_blocks {
+    for (function_name, pc_start, block) in code_blocks {
         compile_block(
             &compiler,
+            &function_name,
             &block,
             pc_start,
             &mut low_level_bytecode,
@@ -130,6 +155,7 @@ pub fn compile_to_low_level_bytecode(
 
 fn compile_block(
     compiler: &Compiler,
+    function_name: &Label,
     block: &[IntermediateInstruction],
     pc_start: CodeAddress,
     low_level_bytecode: &mut Vec<Instruction>,
@@ -359,6 +385,7 @@ fn compile_block(
                     .unwrap()
                     .to_usize();
                 let hint = Hint::RequestMemory {
+                    function_name: function_name.clone(),
                     offset: eval_const_expression_usize(&offset, compiler),
                     vectorized,
                     size,

crates/lean_compiler/src/ir/bytecode.rs

@@ -3,6 +3,16 @@ use lean_vm::Label;
 use std::collections::BTreeMap;
 use std::fmt::{Display, Formatter};
 
+/// A match statement bytecode block
+#[derive(Debug, Clone)]
+pub struct MatchBlock {
+    /// Name of the function containing the match block
+    pub function_name: Label,
+
+    /// Cases of the match block
+    pub match_cases: Vec<Vec<IntermediateInstruction>>,
+}
+
 /// Container for the complete intermediate representation of a program.
 ///
 /// This structure holds all the compiled intermediate bytecode along with
@@ -17,7 +27,7 @@ pub struct IntermediateBytecode {
     /// Match statement bytecode blocks.
     ///
     /// Each match statement produces multiple case blocks.
-    pub match_blocks: Vec<Vec<Vec<IntermediateInstruction>>>,
+    pub match_blocks: Vec<MatchBlock>,
 
     /// Memory requirements for each function.
     ///
@@ -29,13 +39,13 @@ impl Display for IntermediateBytecode {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         for (label, instructions) in &self.bytecode {
             writeln!(f, "\n{label}:")?;
-            for instruction in instructions {
+            for instruction in instructions.iter() {
                 writeln!(f, "  {instruction}")?;
             }
         }
-        for (i, match_blocks) in self.match_blocks.iter().enumerate() {
+        for (i, MatchBlock { match_cases, .. }) in self.match_blocks.iter().enumerate() {
             writeln!(f, "\nMatch {i}:")?;
-            for (j, case) in match_blocks.iter().enumerate() {
+            for (j, case) in match_cases.iter().enumerate() {
                 writeln!(f, "  Case {j}:")?;
                 for instruction in case {
                     writeln!(f, "    {instruction}")?;

crates/lean_compiler/src/ir/mod.rs

@@ -5,7 +5,7 @@ pub mod instruction;
 pub mod operation;
 pub mod value;
 
-pub use bytecode::IntermediateBytecode;
+pub use bytecode::{IntermediateBytecode, MatchBlock};
 pub use instruction::IntermediateInstruction;
 pub use operation::HighLevelOperation;
 pub use value::{IntermediaryMemOrFpOrConstant, IntermediateValue};

crates/lean_vm/src/core/types.rs

@@ -11,3 +11,6 @@ pub type SourceLineNumber = usize;
 
 /// Bytecode address (i.e., a value of the program counter)
 pub type CodeAddress = usize;
+
+/// Memory address
+pub type MemoryAddress = usize;

crates/lean_vm/src/diagnostics/error.rs

@@ -1,4 +1,5 @@
 use crate::core::F;
+use crate::diagnostics::profiler::MemoryProfile;
 use crate::execution::Memory;
 use crate::witness::{
     WitnessDotProduct, WitnessMultilinearEval, WitnessPoseidon16, WitnessPoseidon24,
@@ -46,4 +47,5 @@ pub struct ExecutionResult {
     pub dot_products: Vec<WitnessDotProduct>,
     pub multilinear_evals: Vec<WitnessMultilinearEval>,
     pub summary: String,
+    pub memory_profile: Option<MemoryProfile>,
 }