Skip to content

Commit

Permalink
feat: parse native asm src
Browse files Browse the repository at this point in the history
  • Loading branch information
@lispc
lispc committed Jan 9, 2025
1 parent 9f7ce2c commit 8afd061
Show file tree
Hide file tree
Showing 2 changed files with 309 additions and 1 deletion.
43 changes: 43 additions & 0 deletions extensions/native/compiler/src/asm/code.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ impl<F: PrimeField32, EF: ExtensionField<F>> AssemblyCode<F, EF> {
impl<F: PrimeField32, EF: ExtensionField<F>> Display for AssemblyCode<F, EF> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
for (i, block) in self.blocks.iter().enumerate() {
// TODO: should we skip empty blocks?
writeln!(
f,
"{}:",
Expand All @@ -66,3 +67,45 @@ impl<F: PrimeField32, EF: ExtensionField<F>> Display for AssemblyCode<F, EF> {
Ok(())
}
}

impl<F: PrimeField32, EF: ExtensionField<F>> AssemblyCode<F, EF> {
pub fn from_str(s: &str) -> Result<Self, String> {
let mut blocks = Vec::new();
let mut labels = BTreeMap::new();
let mut current_block = BasicBlock::new();
let mut current_label = None;

for line in s.lines() {
let line = line.trim();
if line.is_empty() {
continue;
}

if line.ends_with(':') {
if let Some(label) = current_label.take() {
labels.insert(F::from_canonical_u32(blocks.len() as u32), label);
blocks.push(current_block);
current_block = BasicBlock::new();
}
current_label = Some(line.trim_end_matches(':').to_string());
} else {
if let Some(instruction) = AsmInstruction::parse_instruction(line, &labels) {
current_block.push(instruction, None);
} else {
return Err(format!("Failed to parse instruction: {}", line));
}
}
}

if let Some(label) = current_label {
labels.insert(F::from_canonical_u32(blocks.len() as u32), label);
}
if !current_block.0.is_empty() {
blocks.push(current_block);
}

let code = AssemblyCode { blocks, labels };
//debug_assert_eq!(s, code.to_string());
Ok(code)
}
}
267 changes: 266 additions & 1 deletion extensions/native/compiler/src/asm/instruction.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -178,6 +178,7 @@ impl<F: PrimeField32, EF: ExtensionField<F>> AsmInstruction<F, EF> {
)
}
AsmInstruction::ImmF(dst, src) => {
// TODO: why "()" here?
write!(f, "imm ({})fp, ({})", dst, src)
}
AsmInstruction::CopyF(dst, src) => {
Expand Down Expand Up @@ -211,7 +212,7 @@ impl<F: PrimeField32, EF: ExtensionField<F>> AsmInstruction<F, EF> {
write!(f, "divi ({})fp, ({})fp, {}", dst, lhs, rhs)
}
AsmInstruction::DivFIN(dst, lhs, rhs) => {
write!(f, "divi ({})fp, {}, ({})fp", dst, lhs, rhs)
write!(f, "divin ({})fp, {}, ({})fp", dst, lhs, rhs)
}
AsmInstruction::AddE(dst, lhs, rhs) => {
write!(f, "eadd ({})fp, ({})fp, ({})fp", dst, lhs, rhs)
Expand Down Expand Up @@ -350,3 +351,267 @@ impl<F: PrimeField32, EF: ExtensionField<F>> AsmInstruction<F, EF> {
}
}
}

impl<F: PrimeField32, EF: ExtensionField<F>> AsmInstruction<F, EF> {
pub fn parse_instruction(line: &str, labels: &BTreeMap<F, String>) -> Option<Self> {
let parts: Vec<&str> = line.split_whitespace().collect();
if parts.is_empty() {
return None;
}

// Helper function to parse i32 from formatted string
let parse_i32 = |s: &str| -> Option<i32> {
s.trim_end_matches(',')
.trim_end_matches("fp")
.trim_end_matches(')')
.trim_start_matches('(')
.parse()
.ok()
};

// Helper function to parse field element F
let parse_f_strict = |s: &str| -> Option<F> { s.parse().ok().map(F::from_canonical_u32) };
let parse_f = |s: &str| -> Option<F> {
parse_f_strict(
s.trim_end_matches(',')
.trim_end_matches(')')
.trim_start_matches('('),
)
};

// Helper function to parse extension field element EF
let parse_ef = |s: &str| -> Option<EF> {
// TODO: is it enough?
s.trim_end_matches(',')
.parse()
.ok()
.map(EF::from_canonical_u64)
};

let parse_label = |label: &str| -> Option<F> {
let label = label.trim_end_matches(',');
let label_key = labels
.iter()
.find(|&(_, v)| v == label)
.map(|(k, _)| k.clone())
.unwrap_or_else(|| parse_f_strict(label.trim_start_matches(".L")).expect(&label));
Some(label_key)
};

match parts[0] {
"lw" => Some(AsmInstruction::LoadF(
parse_i32(parts[1])?,
parse_i32(parts[2])?,
parse_i32(parts[3])?,
parse_f(parts[4])?,
parse_f(parts[5])?,
)),
"lwi" => Some(AsmInstruction::LoadFI(
parse_i32(parts[1])?,
parse_i32(parts[2])?,
parse_f(parts[3])?,
parse_f(parts[4])?,
parse_f(parts[5])?,
)),
"sw" => Some(AsmInstruction::StoreF(
parse_i32(parts[1])?,
parse_i32(parts[2])?,
parse_i32(parts[3])?,
parse_f(parts[4])?,
parse_f(parts[5])?,
)),
"swi" => Some(AsmInstruction::StoreFI(
parse_i32(parts[1])?,
parse_i32(parts[2])?,
parse_f(parts[3])?,
parse_f(parts[4])?,
parse_f(parts[5])?,
)),
"imm" => Some(AsmInstruction::ImmF(
parse_i32(parts[1])?,
parse_f(parts[2])?,
)),
"copy" => Some(AsmInstruction::CopyF(
parse_i32(parts[1])?,
parse_i32(parts[2])?,
)),
"add" => Some(AsmInstruction::AddF(
parse_i32(parts[1])?,
parse_i32(parts[2])?,
parse_i32(parts[3])?,
)),
"addi" => Some(AsmInstruction::AddFI(
parse_i32(parts[1])?,
parse_i32(parts[2])?,
parse_f(parts[3])?,
)),
"sub" => Some(AsmInstruction::SubF(
parse_i32(parts[1])?,
parse_i32(parts[2])?,
parse_i32(parts[3])?,
)),
"subi" => Some(AsmInstruction::SubFI(
parse_i32(parts[1])?,
parse_i32(parts[2])?,
parse_f(parts[3])?,
)),
"subin" => Some(AsmInstruction::SubFIN(
parse_i32(parts[1])?,
parse_f(parts[2])?,
parse_i32(parts[3])?,
)),
"mul" => Some(AsmInstruction::MulF(
parse_i32(parts[1])?,
parse_i32(parts[2])?,
parse_i32(parts[3])?,
)),
"muli" => Some(AsmInstruction::MulFI(
parse_i32(parts[1])?,
parse_i32(parts[2])?,
parse_f(parts[3])?,
)),
"div" => Some(AsmInstruction::DivF(
parse_i32(parts[1])?,
parse_i32(parts[2])?,
parse_i32(parts[3])?,
)),
"divi" => Some(AsmInstruction::DivFI(
parse_i32(parts[1])?,
parse_i32(parts[2])?,
parse_f(parts[3])?,
)),
"divin" => Some(AsmInstruction::DivFIN(
parse_i32(parts[1])?,
parse_f(parts[2])?,
parse_i32(parts[3])?,
)),
"eadd" => Some(AsmInstruction::AddE(
parse_i32(parts[1])?,
parse_i32(parts[2])?,
parse_i32(parts[3])?,
)),
"esub" => Some(AsmInstruction::SubE(
parse_i32(parts[1])?,
parse_i32(parts[2])?,
parse_i32(parts[3])?,
)),
"emul" => Some(AsmInstruction::MulE(
parse_i32(parts[1])?,
parse_i32(parts[2])?,
parse_i32(parts[3])?,
)),
"ediv" => Some(AsmInstruction::DivE(
parse_i32(parts[1])?,
parse_i32(parts[2])?,
parse_i32(parts[3])?,
)),
"j" => {
let dst = parse_i32(parts[1])?;
let label_key = parse_label(parts[2])?;
Some(AsmInstruction::Jump(dst, label_key))
}
"bne" => {
let label_key = parse_label(parts[1])?;
Some(AsmInstruction::Bne(
label_key,
parse_i32(parts[2])?,
parse_i32(parts[3])?,
))
}
"bnei" => {
let label_key = parse_label(parts[1])?;
Some(AsmInstruction::BneI(
label_key,
parse_i32(parts[2])?,
parse_f(parts[3])?,
))
}
"beq" => {
let label_key = parse_label(parts[1])?;
Some(AsmInstruction::Beq(
label_key,
parse_i32(parts[2])?,
parse_i32(parts[3])?,
))
}
"beqi" => {
let label_key = parse_label(parts[1])?;
Some(AsmInstruction::BeqI(
label_key,
parse_i32(parts[2])?,
parse_f(parts[3])?,
))
}
"ebne" => {
let label_key = parse_label(parts[1])?;
Some(AsmInstruction::BneE(
label_key,
parse_i32(parts[2])?,
parse_i32(parts[3])?,
))
}
"ebnei" => {
let label_key = parse_label(parts[1])?;
Some(AsmInstruction::BneEI(
label_key,
parse_i32(parts[2])?,
parse_ef(parts[3])?,
))
}
"ebeq" => {
let label_key = parse_label(parts[1])?;
Some(AsmInstruction::BeqE(
label_key,
parse_i32(parts[2])?,
parse_i32(parts[3])?,
))
}
"ebeqi" => {
let label_key = parse_label(parts[1])?;
Some(AsmInstruction::BeqEI(
label_key,
parse_i32(parts[2])?,
parse_ef(parts[3])?,
))
}
"trap" => Some(AsmInstruction::Trap),
"halt" => Some(AsmInstruction::Halt),
"hint_bits" => Some(AsmInstruction::HintBits(
parse_i32(parts[1])?,
parts[2].parse().ok()?,
)),
"poseidon2_permute" => Some(AsmInstruction::Poseidon2Permute(
parse_i32(parts[1])?,
parse_i32(parts[2])?,
)),
"poseidon2_compress" => Some(AsmInstruction::Poseidon2Compress(
parse_i32(parts[1])?,
parse_i32(parts[2])?,
parse_i32(parts[3])?,
)),
"print_f" => Some(AsmInstruction::PrintF(parse_i32(parts[1])?)),
"print_v" => Some(AsmInstruction::PrintV(parse_i32(parts[1])?)),
"print_e" => Some(AsmInstruction::PrintE(parse_i32(parts[1])?)),
"hint_vec" => Some(AsmInstruction::HintInputVec()),
"shintw" => Some(AsmInstruction::StoreHintWordI(
parse_i32(parts[1])?,
parse_f(parts[2])?,
)),
"commit" => Some(AsmInstruction::Publish(
parse_i32(parts[1])?,
parse_i32(parts[2])?,
)),
"cycle_tracker_start" => Some(AsmInstruction::CycleTrackerStart()),
"cycle_tracker_end" => Some(AsmInstruction::CycleTrackerEnd()),
"fri_mat_opening" => Some(AsmInstruction::FriReducedOpening(
parse_i32(parts[1])?,
parse_i32(parts[2])?,
parse_i32(parts[3])?,
parse_i32(parts[4])?,
parse_i32(parts[5])?,
parse_i32(parts[6])?,
)),
_ => None,
}
}
}

0 comments on commit 8afd061

Please sign in to comment.