Refactor execution runtime cleanup

Author: AbdelStark

src/burn_runtime.rs

@@ -3,7 +3,10 @@ use std::time::Instant;
 use burn::prelude::Backend;
 
 use crate::burn_model::BurnTransformerVm;
-use crate::engine::{ExecutionEngine, ExecutionResult, ExecutionTraceEntry};
+use crate::engine::{
+    build_execution_result, execution_complete, record_execution_step, ExecutionEngine,
+    ExecutionResult, ExecutionTraceEntry,
+};
 use crate::error::Result;
 use crate::instruction::Instruction;
 use crate::memory::AddressedMemory;
@@ -26,10 +29,7 @@ impl<B: Backend> BurnExecutionRuntime<B> {
     pub fn new(model: BurnTransformerVm<B>, device: B::Device, max_steps: usize) -> Self {
         let initial_memory = model.program().initial_memory().to_vec();
         let memory = AddressedMemory::from_initial(&initial_memory);
-        let state = MachineState {
-            memory: initial_memory,
-            ..MachineState::new(model.program().memory_size())
-        };
+        let state = MachineState::with_memory(initial_memory);
 
         Self {
             model,
@@ -44,7 +44,7 @@ impl<B: Backend> BurnExecutionRuntime<B> {
     }
 
     pub fn step(&mut self) -> Result<&MachineState> {
-        if self.state.halted || self.step_count >= self.max_steps {
+        if execution_complete(&self.state, self.step_count, self.max_steps) {
             return Ok(&self.state);
         }
 
@@ -58,36 +58,28 @@ impl<B: Backend> BurnExecutionRuntime<B> {
         )?;
         self.state = next;
         self.step_count += 1;
-        self.trace.push(self.state.clone());
-        self.events.push(ExecutionTraceEntry {
-            step: self.step_count,
-            layer_idx: Some(dispatch.layer_idx),
-            instruction: dispatch.instruction,
-            state_before: before,
-            state_after: self.state.clone(),
-        });
+        record_execution_step(
+            &mut self.trace,
+            &mut self.events,
+            self.step_count,
+            Some(dispatch.layer_idx),
+            dispatch.instruction,
+            before,
+            &self.state,
+        );
         Ok(&self.state)
     }
 
     pub fn run(&mut self) -> Result<ExecutionResult> {
         let start = Instant::now();
-        while self.step_count < self.max_steps && !self.state.halted {
+        while !execution_complete(&self.state, self.step_count, self.max_steps) {
             self.step()?;
         }
-
-        let elapsed = start.elapsed();
-        let elapsed_secs = elapsed.as_secs_f64();
-        Ok(ExecutionResult {
-            final_state: self.state.clone(),
-            steps: self.step_count,
-            halted: self.state.halted,
-            elapsed,
-            tokens_per_sec: if elapsed_secs > 0.0 {
-                self.step_count as f64 / elapsed_secs
-            } else {
-                0.0
-            },
-        })
+        Ok(build_execution_result(
+            &self.state,
+            self.step_count,
+            start.elapsed(),
+        ))
     }
 
     pub fn trace(&self) -> &[MachineState] {
@@ -119,7 +111,7 @@ impl<B: Backend> BurnExecutionRuntime<B> {
     }
 
     pub fn next_dispatch(&self) -> Result<Option<DispatchInfo>> {
-        if self.state.halted || self.step_count >= self.max_steps {
+        if execution_complete(&self.state, self.step_count, self.max_steps) {
             return Ok(None);
         }
         self.model.dispatch_info(self.state.pc).map(Some)

src/engine.rs

@@ -45,6 +45,57 @@ pub trait ExecutionEngine {
     fn next_instruction(&self) -> Result<Option<Instruction>>;
 
     fn is_halted(&self) -> bool {
-        self.state().halted || self.step_count() >= self.max_steps()
+        execution_complete(self.state(), self.step_count(), self.max_steps())
+    }
+}
+
+pub(crate) fn execution_complete(
+    state: &MachineState,
+    step_count: usize,
+    max_steps: usize,
+) -> bool {
+    state.halted || step_count >= max_steps
+}
+
+pub(crate) fn build_execution_result(
+    final_state: &MachineState,
+    steps: usize,
+    elapsed: Duration,
+) -> ExecutionResult {
+    ExecutionResult {
+        final_state: final_state.clone(),
+        steps,
+        halted: final_state.halted,
+        elapsed,
+        tokens_per_sec: tokens_per_sec(steps, elapsed),
+    }
+}
+
+pub(crate) fn record_execution_step(
+    trace: &mut Vec<MachineState>,
+    events: &mut Vec<ExecutionTraceEntry>,
+    step: usize,
+    layer_idx: Option<usize>,
+    instruction: Instruction,
+    state_before: MachineState,
+    state_after: &MachineState,
+) {
+    let state_after = state_after.clone();
+    trace.push(state_after.clone());
+    events.push(ExecutionTraceEntry {
+        step,
+        layer_idx,
+        instruction,
+        state_before,
+        state_after,
+    });
+}
+
+fn tokens_per_sec(steps: usize, elapsed: Duration) -> f64 {
+    let elapsed_secs = elapsed.as_secs_f64();
+    if elapsed_secs > 0.0 {
+        steps as f64 / elapsed_secs
+    } else {
+        0.0
     }
 }