Merge branch 'parallel-runner'

Author: TomWambsgans

crates/lean_compiler/snark_lib.py

@@ -14,10 +14,11 @@ def inline(fn):
     return fn
 
 
-# unroll(a, b) returns range(a, b) for Python execution
 def unroll(a: int, b: int):
     return range(a, b)
 
+def parallel_range(a: int, b: int):
+    return range(a, b)
 
 # dynamic_unroll(start, end, n_bits) returns range(start, end) for Python execution
 def dynamic_unroll(start: int, end: int, n_bits: int):

crates/lean_compiler/src/a_simplify_lang.rs

@@ -1196,9 +1196,10 @@ fn transform_mutable_in_loops_in_lines(
                 start,
                 end,
                 body,
-                loop_kind: LoopKind::Range,
+                loop_kind: loop_kind @ (LoopKind::Range | LoopKind::ParallelRange),
                 location,
             } => {
+                let loop_kind = loop_kind.clone();
                 transform_mutable_in_loops_in_lines(body, const_arrays, counter);
                 let modified_vars = find_modified_external_vars(body, const_arrays);
 
@@ -1344,7 +1345,7 @@ fn transform_mutable_in_loops_in_lines(
                     start: start.clone(),
                     end: end.clone(),
                     body: new_body,
-                    loop_kind: LoopKind::Range,
+                    loop_kind,
                     location,
                 });
 
@@ -2633,10 +2634,12 @@ fn simplify_lines(
                 location,
             } => {
                 assert!(
-                    matches!(loop_kind, LoopKind::Range),
+                    matches!(loop_kind, LoopKind::Range | LoopKind::ParallelRange),
                     "Unrolled/dynamic_unroll loops should have been handled already"
                 );
 
+                let is_parallel = loop_kind.is_parallel();
+
                 let mut loop_const_malloc = ConstMalloc {
                     counter: const_malloc.counter,
                     ..ConstMalloc::default()
@@ -2653,7 +2656,8 @@ fn simplify_lines(
                 const_malloc.counter = loop_const_malloc.counter;
                 state.array_manager.valid = valid_aux_vars_in_array_manager_before; // restore the valid aux vars
 
-                let func_name = format!("@loop_{}_{}", state.counters.loops.get_next(), location);
+                let loop_prefix = if is_parallel { "@parallel_loop" } else { "@loop" };
+                let func_name = format!("{}_{}_{}", loop_prefix, state.counters.loops.get_next(), location);
 
                 // Find variables used inside loop but defined outside
                 let (_, mut external_vars) = find_variable_usage(body, ctx.const_arrays);

crates/lean_compiler/src/b_compile_intermediate.rs

@@ -149,16 +149,28 @@ fn compile_function(
     compiler.const_malloc_vars.clear();
     (compiler.dead_fp_relative_vars, compiler.dead_store_vars) = compute_dead_vars(&function.instructions);
 
-    compile_lines(
-        &Label::function(function.name.clone()),
-        &function.instructions,
-        compiler,
-        None,
-    )
+    let mut instructions = Vec::new();
+
+    // Emit ParallelBatchStart for parallel loop functions.
+    // The first instruction is always `diff = iterator - end`. The end value is either
+    // a function argument (runtime) or a compile-time constant.
+    if function.name.starts_with("@parallel_loop_") {
+        let SimpleLine::Assignment { arg1, .. } = &function.instructions[0] else {
+            panic!("parallel loop: expected first instruction to be `diff = i - end`");
+        };
+        let end_value = IntermediateValue::from_simple_expr(arg1, compiler);
+        instructions.push(IntermediateInstruction::ParallelBatchStart {
+            n_args: function.arguments.len(),
+            end_value,
+        });
+    }
+
+    instructions.extend(compile_lines(&function.instructions, compiler, None)?);
+
+    Ok(instructions)
 }
 
 fn compile_lines(
-    function_name: &Label,
     lines: &[SimpleLine],
     compiler: &mut Compiler,
     final_jump: Option<Label>,
@@ -263,14 +275,13 @@ fn compile_lines(
                 for arm in arms.iter() {
                     compiler.stack_pos = saved_stack_pos;
                     compiler.stack_frame_layout.scopes.push(ScopeLayout::default());
-                    let arm_instructions = compile_lines(function_name, arm, compiler, Some(end_label.clone()))?;
+                    let arm_instructions = compile_lines(arm, compiler, Some(end_label.clone()))?;
                     compiled_arms.push(arm_instructions);
                     compiler.stack_frame_layout.scopes.pop();
                     new_stack_pos = new_stack_pos.max(compiler.stack_pos);
                 }
                 compiler.stack_pos = new_stack_pos;
                 compiler.match_blocks.push(MatchBlock {
-                    function_name: function_name.clone(),
                     match_cases: compiled_arms,
                 });
                 // Get the actual index AFTER pushing (nested matches may have pushed their blocks first)
@@ -317,7 +328,7 @@ fn compile_lines(
                     updated_fp: None,
                 });
 
-                let remaining = compile_lines(function_name, &lines[i + 1..], compiler, final_jump)?;
+                let remaining = compile_lines(&lines[i + 1..], compiler, final_jump)?;
                 compiler.bytecode.insert(end_label, remaining);
 
                 compiler.stack_frame_layout.scopes.pop();
@@ -409,22 +420,22 @@ fn compile_lines(
                 let saved_stack_pos = compiler.stack_pos;
 
                 compiler.stack_frame_layout.scopes.push(ScopeLayout::default());
-                let then_instructions = compile_lines(function_name, then_branch, compiler, Some(end_label.clone()))?;
+                let then_instructions = compile_lines(then_branch, compiler, Some(end_label.clone()))?;
 
                 let then_stack_pos = compiler.stack_pos;
                 compiler.stack_pos = saved_stack_pos;
                 compiler.stack_frame_layout.scopes.pop();
                 compiler.stack_frame_layout.scopes.push(ScopeLayout::default());
 
-                let else_instructions = compile_lines(function_name, else_branch, compiler, Some(end_label.clone()))?;
+                let else_instructions = compile_lines(else_branch, compiler, Some(end_label.clone()))?;
 
                 compiler.bytecode.insert(if_label, then_instructions);
                 compiler.bytecode.insert(else_label, else_instructions);
 
                 compiler.stack_frame_layout.scopes.pop();
                 compiler.stack_pos = compiler.stack_pos.max(then_stack_pos);
 
-                let remaining = compile_lines(function_name, &lines[i + 1..], compiler, final_jump)?;
+                let remaining = compile_lines(&lines[i + 1..], compiler, final_jump)?;
                 compiler.bytecode.insert(end_label, remaining);
                 // It is not necessary to update compiler.stack_size here because the preceding call to
                 // compile_lines should have done so.
@@ -485,7 +496,7 @@ fn compile_lines(
                         });
                     }
 
-                    instructions.extend(compile_lines(function_name, &lines[i + 1..], compiler, final_jump)?);
+                    instructions.extend(compile_lines(&lines[i + 1..], compiler, final_jump)?);
 
                     instructions
                 };

crates/lean_compiler/src/c_compile_final.rs

@@ -14,7 +14,8 @@ impl IntermediateInstruction {
             | Self::LocationReport { .. }
             | Self::DebugAssert { .. }
             | Self::DerefHint { .. }
-            | Self::PanicHint { .. } => true,
+            | Self::PanicHint { .. }
+            | Self::ParallelBatchStart { .. } => true,
             Self::Computation { .. }
             | Self::Panic
             | Self::Deref { .. }
@@ -66,42 +67,19 @@ pub fn compile_to_low_level_bytecode(
         .bytecode
         .remove(&Label::function("main"))
         .ok_or("No main function found in the compiled program")?;
-    hints.insert(
-        STARTING_PC,
-        vec![Hint::StackFrame {
-            label: Label::function("main"),
-            size: starting_frame_memory,
-        }],
-    );
 
     let mut pc = count_real_instructions(&exit_point) + count_real_instructions(&entrypoint);
-    let mut code_blocks = vec![
-        (Label::EndProgram, ENDING_PC, exit_point),
-        (Label::function("main"), STARTING_PC, entrypoint),
-    ];
+    let mut code_blocks = vec![(ENDING_PC, exit_point), (STARTING_PC, entrypoint)];
 
     for (label, instructions) in &intermediate_bytecode.bytecode {
         label_to_pc.insert(label.clone(), pc);
-        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()));
+        code_blocks.push((pc, instructions.clone()));
         pc += count_real_instructions(instructions);
     }
 
     let mut match_block_sizes = Vec::new();
     let mut match_first_block_starts = Vec::new();
-    for MatchBlock {
-        function_name,
-        match_cases,
-    } in intermediate_bytecode.match_blocks
-    {
+    for MatchBlock { match_cases } in intermediate_bytecode.match_blocks {
         let max_block_size = match_cases
             .iter()
             .map(|block| count_real_instructions(block))
@@ -116,7 +94,7 @@ pub fn compile_to_low_level_bytecode(
                 IntermediateInstruction::Panic;
                 max_block_size - count_real_instructions(&block)
             ]);
-            code_blocks.push((function_name.clone(), pc, block));
+            code_blocks.push((pc, block));
             pc += max_block_size;
         }
     }
@@ -134,15 +112,8 @@ pub fn compile_to_low_level_bytecode(
 
     let mut instructions = Vec::new();
 
-    for (function_name, pc_start, block) in code_blocks {
-        compile_block(
-            &compiler,
-            &function_name,
-            &block,
-            pc_start,
-            &mut instructions,
-            &mut hints,
-        );
+    for (pc_start, block) in code_blocks {
+        compile_block(&compiler, &block, pc_start, &mut instructions, &mut hints);
     }
     let instructions_encoded = instructions.par_iter().map(field_representation).collect::<Vec<_>>();
 
@@ -195,7 +166,6 @@ 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>,
@@ -341,7 +311,6 @@ fn compile_block(
             IntermediateInstruction::RequestMemory { offset, size } => {
                 let size = try_as_mem_or_constant(&size).unwrap();
                 let hint = Hint::RequestMemory {
-                    function_name: function_name.clone(),
                     offset: eval_const_expression_usize(&offset, compiler),
                     size,
                 };
@@ -386,6 +355,13 @@ fn compile_block(
                 let hint = Hint::Panic { message };
                 hints.entry(pc).or_default().push(hint);
             }
+            IntermediateInstruction::ParallelBatchStart { n_args, end_value } => {
+                let end_value = try_as_mem_or_constant(&end_value).expect("parallel loop end value");
+                hints
+                    .entry(pc)
+                    .or_default()
+                    .push(Hint::ParallelBatchStart { n_args, end_value });
+            }
         }
 
         if !instruction.is_hint() {

crates/lean_compiler/src/grammar.pest

@@ -90,8 +90,9 @@ elif_clause = { "elif" ~ condition ~ ":" ~ newline ~ statement* ~ end_block }
 
 else_clause = { "else" ~ ":" ~ newline ~ statement* ~ end_block }
 
-for_statement = { "for" ~ identifier ~ "in" ~ (dynamic_unroll_range | unroll_range | range) ~ ":" ~ newline ~ statement* ~ end_block }
+for_statement = { "for" ~ identifier ~ "in" ~ (dynamic_unroll_range | unroll_range | parallel_range | range) ~ ":" ~ newline ~ statement* ~ end_block }
 range = { "range" ~ "(" ~ expression ~ "," ~ expression ~ ")" }
+parallel_range = { "parallel_range" ~ "(" ~ expression ~ "," ~ expression ~ ")" }
 unroll_range = { "unroll" ~ "(" ~ expression ~ "," ~ expression ~ ")" }
 dynamic_unroll_range = { "dynamic_unroll" ~ "(" ~ expression ~ "," ~ expression ~ "," ~ expression ~ ")" }
 

crates/lean_compiler/src/ir/bytecode.rs

@@ -6,9 +6,6 @@ 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>>,
 }

crates/lean_compiler/src/ir/instruction.rs

@@ -65,6 +65,11 @@ pub enum IntermediateInstruction {
     PanicHint {
         message: Option<String>,
     },
+    /// Marks the start of a parallelizable loop
+    ParallelBatchStart {
+        n_args: usize,
+        end_value: IntermediateValue,
+    },
 }
 
 impl IntermediateInstruction {
@@ -196,6 +201,9 @@ impl Display for IntermediateInstruction {
                 Some(msg) => write!(f, "panic hint: \"{msg}\""),
                 None => write!(f, "panic hint"),
             },
+            Self::ParallelBatchStart { n_args, end_value } => {
+                write!(f, "parallel_batch_start(n_args={n_args}, end={end_value})")
+            }
         }
     }
 }

crates/lean_compiler/src/lang.rs

@@ -555,6 +555,7 @@ impl VecLiteral {
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub enum LoopKind {
     Range,
+    ParallelRange,
     Unroll,
     /// `for i in dynamic_unroll(0, a, n_bits): body` — unrolls over runtime-bounded range
     /// using bit decomposition. `n_bits` must be compile-time known.
@@ -567,6 +568,10 @@ impl LoopKind {
     pub fn is_unroll(&self) -> bool {
         matches!(self, Self::Unroll | Self::DynamicUnroll { .. })
     }
+
+    pub fn is_parallel(&self) -> bool {
+        matches!(self, Self::ParallelRange)
+    }
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -798,7 +803,13 @@ impl Line {
                         iterator, start, end, n_bits, body_str, spaces
                     ),
                     _ => {
-                        let range_fn = if loop_kind.is_unroll() { "unroll" } else { "range" };
+                        let range_fn = if loop_kind.is_unroll() {
+                            "unroll"
+                        } else if loop_kind.is_parallel() {
+                            "parallel_range"
+                        } else {
+                            "range"
+                        };
                         format!(
                             "for {} in {}({}, {}) {{\n{}\n{}}}",
                             iterator, range_fn, start, end, body_str, spaces

crates/lean_compiler/src/parser/parsers/function.rs

@@ -24,6 +24,7 @@ pub const RESERVED_FUNCTION_NAMES: &[&str] = &[
     "next_multiple_of",
     "saturating_sub",
     "range",
+    "parallel_range",
     "match_range",
 ];
 

crates/lean_compiler/src/parser/parsers/statement.rs

@@ -203,6 +203,7 @@ impl Parse<Line> for ForStatementParser {
                 let n_bits = ExpressionParser.parse(next_inner_pair(&mut range_inner, "n_bits")?, ctx)?;
                 LoopKind::DynamicUnroll { n_bits }
             }
+            Rule::parallel_range => LoopKind::ParallelRange,
             _ => LoopKind::Range,
         };