Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: parse native asm src #1198

Closed
wants to merge 2 commits into from
Closed

feat: parse native asm src #1198

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8afd061
feat: parse native asm src
lispc Jan 9, 2025
f9d1403
lint
lispc Jan 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 41 additions & 0 deletions extensions/native/compiler/src/asm/code.rs
View file
Open in desktop
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,43 @@ impl<F: PrimeField32, EF: ExtensionField<F>> Display for AssemblyCode<F, EF> {
Ok(())
}
}

impl<F: PrimeField32, EF: ExtensionField<F>> AssemblyCode<F, EF> {
pub fn parse_asm(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
Open in desktop
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)
.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,
}
}
}