Skip to content

Commit

Permalink
Add BasicBlock Trace from IR (#105)
Browse files Browse the repository at this point in the history
Summary:

Test Plan:

Co-authored-by: duc-nx <>
  • Loading branch information
duc-nx authored and sjudson committed Feb 12, 2025
1 parent a99dd24 commit 60b4e83
Show file tree
Hide file tree
Showing 2 changed files with 181 additions and 51 deletions.
138 changes: 89 additions & 49 deletions vm/src/emulator/executor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -421,6 +421,16 @@ impl LinearEmulator {
todo!()
}

/// Creates a Linear Emulator from an ELF file.
///
/// This function initializes a Linear Emulator with the provided ELF file, memory layout,
/// and input data. It sets up the memory segments according to the ELF file structure
/// and the specified memory layout.
///
/// # Panics
///
/// This function will panic if the provided ElfFile is not well-formed, or if the memory
/// layout is not compatible with the ELF file.
pub fn from_elf(
memory_layout: LinearMemoryLayout,
ad: &[u32],
Expand Down Expand Up @@ -538,6 +548,53 @@ impl LinearEmulator {
emulator
}

/// Creates a Linear Emulator from a basic block IR, for simple testing purposes.
///
/// This function initializes a Linear Emulator with a single basic block of instructions,
/// along with the necessary memory layout and input data. It's primarily used for testing
/// and simple emulation scenarios.
///
/// # Note
///
/// This function currently only sets up simple instruction memory. It may be extended
/// in the future to support more features and memory configurations.
pub fn from_basic_blocks(
memory_layout: LinearMemoryLayout,
basic_blocks: &Vec<BasicBlock>,
) -> Self {
let mut memory = UnifiedMemory::default();

let mut encoded_basic_blocks = Vec::new();
for block in basic_blocks {
encoded_basic_blocks.extend(block.encode());
}

// Add basic blocks instructions to memory
let code_start = memory_layout.program_start();
let code_memory = FixedMemory::<RO>::from_vec(
code_start,
encoded_basic_blocks.len() * WORD_SIZE,
encoded_basic_blocks,
);

let instruction_index = memory.add_fixed_ro(&code_memory).unwrap();

let mut emulator = Self {
executor: Executor {
base_address: code_start,
entrypoint: code_start,
global_clock: 1, // global_clock = 0 captures initalization for memory records
..Default::default()
},
instruction_index,
memory_layout,
memory,
..Default::default()
};
emulator.executor.cpu.pc.value = emulator.executor.entrypoint;
emulator
}

fn manage_timestamps(&mut self, size: &MemAccessSize, address: &u32) -> usize {
let half_aligned_address = address & !(WORD_SIZE / 2 - 1) as u32;
let full_aligned_address = address & !(WORD_SIZE - 1) as u32;
Expand Down Expand Up @@ -689,17 +746,9 @@ mod tests {
use crate::elf::ElfFile;
use crate::riscv::{BuiltinOpcode, Instruction, InstructionType, Opcode};

#[test]
fn test_harvard_emulate_nexus_rt_binary() {
let elf_file = ElfFile::from_path("test/fib_10.elf").expect("Unable to load ELF file");
let mut emulator = HarvardEmulator::from_elf(elf_file, &[], &[]);

assert_eq!(emulator.execute(), Err(VMError::VMExited(0)));
}

#[test]
#[rustfmt::skip]
fn test_harvard_fibonacci() {
fn setup_basic_block_ir() -> Vec<BasicBlock>
{
let basic_block = BasicBlock::new(vec![
// Set x0 = 0 (default constant), x1 = 1
Instruction::new(Opcode::from(BuiltinOpcode::ADDI), 1, 0, 1, InstructionType::IType),
Expand Down Expand Up @@ -736,9 +785,25 @@ mod tests {
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 30, 29, 28, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 31, 30, 29, InstructionType::RType),
]);
vec![basic_block]
}

#[test]
fn test_harvard_emulate_nexus_rt_binary() {
let elf_file = ElfFile::from_path("test/fib_10.elf").expect("Unable to load ELF file");
let mut emulator = HarvardEmulator::from_elf(elf_file, &[], &[]);

assert_eq!(emulator.execute(), Err(VMError::VMExited(0)));
}

#[test]
fn test_harvard_fibonacci() {
let basic_blocks = setup_basic_block_ir();

let mut emulator = HarvardEmulator::default();
emulator.execute_basic_block(&basic_block).unwrap();
basic_blocks.iter().for_each(|basic_block| {
emulator.execute_basic_block(basic_block).unwrap();
});

assert_eq!(emulator.executor.cpu.registers[31.into()], 1346269);
}
Expand All @@ -764,47 +829,13 @@ mod tests {
}

#[test]
#[rustfmt::skip]
fn test_linear_fibonacci() {
let basic_block = BasicBlock::new(vec![
// Set x0 = 0 (default constant), x1 = 1
Instruction::new(Opcode::from(BuiltinOpcode::ADDI), 1, 0, 1, InstructionType::IType),
// x2 = x1 + x0
// x3 = x2 + x1 ... and so on
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 2, 1, 0, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 3, 2, 1, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 4, 3, 2, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 5, 4, 3, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 6, 5, 4, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 7, 6, 5, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 8, 7, 6, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 9, 8, 7, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 10, 9, 8, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 11, 10, 9, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 12, 11, 10, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 13, 12, 11, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 14, 13, 12, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 15, 14, 13, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 16, 15, 14, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 17, 16, 15, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 18, 17, 16, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 19, 18, 17, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 20, 19, 18, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 21, 20, 19, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 22, 21, 20, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 23, 22, 21, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 24, 23, 22, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 25, 24, 23, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 26, 25, 24, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 27, 26, 25, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 28, 27, 26, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 29, 28, 27, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 30, 29, 28, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 31, 30, 29, InstructionType::RType),
]);
let basic_blocks = setup_basic_block_ir();

let mut emulator = LinearEmulator::default();
emulator.execute_basic_block(&basic_block).unwrap();
basic_blocks.iter().for_each(|basic_block| {
emulator.execute_basic_block(basic_block).unwrap();
});

assert_eq!(emulator.executor.cpu.registers[31.into()], 1346269);
}
Expand All @@ -819,4 +850,13 @@ mod tests {

assert_eq!(emulator.executor.private_input_tape, private_input_vec);
}

#[test]
fn test_linear_from_basic_block() {
let basic_blocks = setup_basic_block_ir();
let mut emulator =
LinearEmulator::from_basic_blocks(LinearMemoryLayout::default(), &basic_blocks);

assert_eq!(emulator.execute(), Err(VMError::VMOutOfInstructions));
}
}
94 changes: 92 additions & 2 deletions vm/src/trace.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ use crate::{
emulator::{Emulator, LinearEmulator, LinearMemoryLayout},
error::{Result, VMError},
memory::MemoryRecords,
riscv::Instruction,
riscv::{BasicBlock, Instruction},
};

/// A program step.
Expand Down Expand Up @@ -115,7 +115,7 @@ impl UniformTrace {
}

/// Represents a program trace over basic blocks.
#[derive(Default, Clone, Serialize, Deserialize)]
#[derive(Debug, Default, Clone, Serialize, Deserialize)]
pub struct BBTrace {
/// Memory layout.
pub memory_layout: LinearMemoryLayout,
Expand Down Expand Up @@ -352,11 +352,42 @@ pub fn bb_trace(
}
}

/// Trace a program over basic blocks.
pub fn bb_trace_direct(basic_blocks: &Vec<BasicBlock>) -> Result<BBTrace> {
let mut vm = LinearEmulator::from_basic_blocks(LinearMemoryLayout::default(), basic_blocks);

let mut trace = BBTrace {
memory_layout: vm.memory_layout,
start: 0,
blocks: Vec::new(),
};

loop {
match bb_step(&mut vm) {
(Some(block), Ok(())) => trace.blocks.push(block),
(Some(block), Err(e)) => {
if !block.steps.is_empty() {
trace.blocks.push(block);
}

match e {
VMError::VMExited(0) => return Ok(trace),
_ => return Err(e),
}
}
(None, Err(VMError::VMOutOfInstructions)) => return Ok(trace),
(None, Err(e)) => return Err(e),
(None, Ok(())) => unreachable!(),
}
}
}

#[cfg(test)]
mod tests {
use super::*;
use crate::memory::{MemAccessSize, MemoryRecord};
use crate::riscv::{BuiltinOpcode, Opcode, Register};
use nexus_common::riscv::instruction::InstructionType;

#[test]
fn test_k1_trace_nexus_rt_binary() {
Expand Down Expand Up @@ -537,4 +568,63 @@ mod tests {
assert_eq!(step.result, Some(0));
assert!(step.memory_records.is_empty());
}

#[rustfmt::skip]
fn setup_basic_block_ir() -> Vec<BasicBlock>
{
let basic_block = BasicBlock::new(vec![
// Set x0 = 0 (default constant), x1 = 1
Instruction::new(Opcode::from(BuiltinOpcode::ADDI), 1, 0, 1, InstructionType::IType),
// x2 = x1 + x0
// x3 = x2 + x1 ... and so on
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 2, 1, 0, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 3, 2, 1, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 4, 3, 2, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 5, 4, 3, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 6, 5, 4, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 7, 6, 5, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 8, 7, 6, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 9, 8, 7, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 10, 9, 8, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 11, 10, 9, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 12, 11, 10, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 13, 12, 11, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 14, 13, 12, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 15, 14, 13, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 16, 15, 14, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 17, 16, 15, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 18, 17, 16, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 19, 18, 17, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 20, 19, 18, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 21, 20, 19, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 22, 21, 20, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 23, 22, 21, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 24, 23, 22, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 25, 24, 23, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 26, 25, 24, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 27, 26, 25, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 28, 27, 26, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 29, 28, 27, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 30, 29, 28, InstructionType::RType),
Instruction::new(Opcode::from(BuiltinOpcode::ADD), 31, 30, 29, InstructionType::RType),
]);
vec![basic_block]
}

#[test]
fn test_bb_trace_simple_from_basic_block_ir() {
let basic_blocks = setup_basic_block_ir();
let trace = bb_trace_direct(&basic_blocks).expect("Failed to create trace");

let first_step = trace.blocks[0].steps.first().expect("No steps in trace");
assert_eq!(first_step.result, Some(1), "Unexpected Fibonacci result",);

let last_step = trace.blocks[0].steps.last().expect("No steps in trace");
// The result of 30th Fibonacci number is 1346269
assert_eq!(
last_step.result,
Some(1346269),
"Unexpected Fibonacci result"
);
}
}

0 comments on commit 60b4e83

Please sign in to comment.