Skip to content

Commit

Permalink
feat: MemoryAccessAdapter (#455)
Browse files Browse the repository at this point in the history 
* feat: MemoryAccessAdapter

* Compiler changes

* Better stack layout

* Change stack/heap layout

* perf: Make duplex challenger more efficient

* fix flamegraph for dsl

* Use stack memory more compactly

* Comments

* Use IsLessThanAir rather than AssertLessThanAir

* feat: field expression framework (#470)

* feat: field expression framework

* test pass

* update readme

* range check vars

* scalar mul and ec double

* address comment

* update readme

* address comments

* check the result in tests

* rename

* Ignore `bench_metrics` in .gitignore (#476)

* [refactor] AIR with Public Values (#477)

* Add trait BaseAirWithPublicValues

* Keygen doesn't need to pass the number of public values

* feat: Keccak batch reads and use of AlignedBorrow (#475)

* feat: Keccak batch reads and use of AlignedBorrow

* Fix timestamp

* feat: add a separate memory block to select for when partial block read

* chore: todo comment on switching to range check output

---------

Co-authored-by: Jonathan Wang <[email protected]>

* fix merge

* fix: change default compiler `word_size=8`

* fix: ecc test compiler word size

* chore: mod arith test word size

---------

Co-authored-by: luffykai <[email protected]>
Co-authored-by: Golovanov399 <[email protected]>
Co-authored-by: Xinding Wei <[email protected]>
Co-authored-by: Jonathan Wang <[email protected]>
  • Loading branch information
@zlangley @luffykai
@Golovanov399 @nyunyunyunyu @jonathanpwang
5 people authored Sep 29, 2024
1 parent b471ac9 commit ffffc5b
Show file tree
Hide file tree
Showing 41 changed files with 1,604 additions and 815 deletions.
2 changes: 1 addition & 1 deletion compiler/src/asm/builder.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmBu
}

pub fn compile_isa_with_options(self, options: CompilerOptions) -> Program<F> {
let mut compiler = AsmCompiler::new(1);
let mut compiler = AsmCompiler::new(options.word_size);
compiler.build(self.operations);
let asm_code = compiler.code();
convert_program(asm_code, options)
Expand Down
73 changes: 39 additions & 34 deletions compiler/src/asm/compiler.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,18 +12,18 @@ use crate::{
};

/// The memory location for the top of memory
pub const MEMORY_TOP: i32 = (1 << 29) - 4;
pub const MEMORY_TOP: u32 = (1 << 29) - 4;

// The memory location for the start of the heap.
pub(crate) const HEAP_START_ADDRESS: i32 = 1 << 24;

/// The heap pointer address.
pub(crate) const HEAP_PTR: i32 = MEMORY_TOP - 4;
pub(crate) const HEAP_PTR: i32 = HEAP_START_ADDRESS - 4;
/// Utility register.
pub(crate) const A0: i32 = MEMORY_TOP - 8;
pub(crate) const A0: i32 = HEAP_START_ADDRESS - 8;

/// The memory location for the top of the stack.
pub(crate) const STACK_TOP: i32 = MEMORY_TOP - 100;

// The memory location for the start of the heap.
pub(crate) const HEAP_START_ADDRESS: usize = 4;
pub(crate) const STACK_TOP: i32 = HEAP_START_ADDRESS - 64;

/// The assembly compiler.
// #[derive(Debug, Clone, Default)]
Expand All @@ -41,21 +41,24 @@ pub struct AsmCompiler<F, EF> {
impl<F> Var<F> {
/// Gets the frame pointer for a var.
pub const fn fp(&self) -> i32 {
STACK_TOP - self.0 as i32
// Vars are stored in stack positions 1, 2, 9, 10, 17, 18, ...
STACK_TOP - (8 * (self.0 / 2) + 1 + (self.0 % 2)) as i32
}
}

impl<F> Felt<F> {
/// Gets the frame pointer for a felt.
pub const fn fp(&self) -> i32 {
STACK_TOP - self.0 as i32
// Felts are stored in stack positions 3, 4, 11, 12, 19, 20, ...
STACK_TOP - (((self.0 >> 1) << 3) + 3 + (self.0 & 1)) as i32
}
}

impl<F, EF> Ext<F, EF> {
/// Gets the frame pointer for an extension element
pub const fn fp(&self) -> i32 {
STACK_TOP - self.0 as i32
// Exts are stored in stack positions 5-8, 13-16, 21-24, ...
STACK_TOP - 8 * self.0 as i32
}
}

Expand Down Expand Up @@ -94,7 +97,7 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
pub fn build(&mut self, operations: TracedVec<DslIr<AsmConfig<F, EF>>>) {
if self.block_label().is_zero() {
// Initialize the heap pointer value.
let heap_start = F::from_canonical_usize(HEAP_START_ADDRESS);
let heap_start = F::from_canonical_u32(HEAP_START_ADDRESS as u32);
self.push(AsmInstruction::ImmF(HEAP_PTR, heap_start), None);
// Jump over the TRAP instruction we are about to add.
self.push(AsmInstruction::j(self.trap_label + F::one()), None);
Expand Down Expand Up @@ -668,25 +671,25 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
size: usize,
debug_info: Option<DebugInfo>,
) {
let word_size = self.word_size;
let align = |x: usize| (x + word_size - 1) / word_size * word_size;
// Load the current heap ptr address to the stack value and advance the heap ptr.
let size = F::from_canonical_usize(size);
let len = len.into();
match len {
RVar::Const(len) => {
self.push(
AsmInstruction::CopyF(ptr.fp(), HEAP_PTR),
debug_info.clone(),
);
self.push(
AsmInstruction::AddFI(HEAP_PTR, HEAP_PTR, len * size),
debug_info,
);
let inc = F::from_canonical_usize(align((len.as_canonical_u32() as usize) * size));
self.push(AsmInstruction::AddFI(HEAP_PTR, HEAP_PTR, inc), debug_info);
}
RVar::Val(len) => {
self.push(
AsmInstruction::CopyF(ptr.fp(), HEAP_PTR),
debug_info.clone(),
);
let size = F::from_canonical_usize(align(size));
self.push(
AsmInstruction::MulFI(A0, len.fp(), size),
debug_info.clone(),
Expand Down Expand Up @@ -1042,8 +1045,10 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
fn assign_exti(&mut self, dst: i32, imm: EF, debug_info: Option<DebugInfo>) {
let imm = imm.as_base_slice();
for i in 0..EF::D {
let j = (i * self.word_size) as i32;
self.push(AsmInstruction::ImmF(dst + j, imm[i]), debug_info.clone());
self.push(
AsmInstruction::ImmF(dst + i as i32, imm[i]),
debug_info.clone(),
);
}
}

Expand All @@ -1059,11 +1064,11 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
for i in 0..EF::D {
self.push(
AsmInstruction::LoadFI(
val.fp() + (i * self.word_size) as i32,
val.fp() + i as i32,
addr,
index,
size,
offset + F::from_canonical_usize(i * self.word_size),
offset + F::from_canonical_usize(i),
),
debug_info.clone(),
)
Expand All @@ -1073,11 +1078,11 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
for i in 0..EF::D {
self.push(
AsmInstruction::LoadF(
val.fp() + (i * self.word_size) as i32,
val.fp() + i as i32,
addr,
index,
size,
offset + F::from_canonical_usize(i * self.word_size),
offset + F::from_canonical_usize(i),
),
debug_info.clone(),
)
Expand All @@ -1098,11 +1103,11 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
for i in 0..EF::D {
self.push(
AsmInstruction::StoreFI(
val.fp() + (i * self.word_size) as i32,
val.fp() + i as i32,
addr,
index,
size,
offset + F::from_canonical_usize(i * self.word_size),
offset + F::from_canonical_usize(i),
),
debug_info.clone(),
)
Expand All @@ -1112,11 +1117,11 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
for i in 0..EF::D {
self.push(
AsmInstruction::StoreF(
val.fp() + (i * self.word_size) as i32,
val.fp() + i as i32,
addr,
index,
size,
offset + F::from_canonical_usize(i * self.word_size),
offset + F::from_canonical_usize(i),
),
debug_info.clone(),
)
Expand All @@ -1134,7 +1139,7 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
) {
let rhs = rhs.as_base_slice();
for i in 0..EF::D {
let j = (i * self.word_size) as i32;
let j = i as i32;
self.push(
AsmInstruction::AddFI(dst.fp() + j, lhs.fp() + j, rhs[i]),
debug_info.clone(),
Expand All @@ -1151,7 +1156,7 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
) {
let lhs = lhs.as_base_slice();
for i in 0..EF::D {
let j = (i * self.word_size) as i32;
let j = i as i32;
self.push(
AsmInstruction::SubFIN(dst.fp() + j, lhs[i], rhs.fp() + j),
debug_info.clone(),
Expand All @@ -1171,7 +1176,7 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
debug_info.clone(),
);
for i in 1..EF::D {
let j = (i * self.word_size) as i32;
let j = i as i32;
self.push(
AsmInstruction::CopyF(dst.fp() + j, lhs.fp() + j),
debug_info.clone(),
Expand All @@ -1191,7 +1196,7 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
debug_info.clone(),
);
for i in 1..EF::D {
let j = (i * self.word_size) as i32;
let j = i as i32;
self.push(
AsmInstruction::CopyF(dst.fp() + j, lhs.fp() + j),
debug_info.clone(),
Expand All @@ -1214,7 +1219,7 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
);

for i in 1..EF::D {
let j = (i * self.word_size) as i32;
let j = i as i32;
self.push(
AsmInstruction::ImmF(dst.fp() + j, rhs[i]),
debug_info.clone(),
Expand All @@ -1230,7 +1235,7 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
debug_info: Option<DebugInfo>,
) {
for i in 0..EF::D {
let j = (i * self.word_size) as i32;
let j = i as i32;
self.push(
AsmInstruction::MulF(dst.fp() + j, lhs.fp() + j, rhs.fp()),
debug_info.clone(),
Expand All @@ -1246,7 +1251,7 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
debug_info: Option<DebugInfo>,
) {
for i in 0..EF::D {
let j = (i * self.word_size) as i32;
let j = i as i32;
self.push(
AsmInstruction::MulFI(dst.fp() + j, lhs.fp() + j, rhs),
debug_info.clone(),
Expand All @@ -1262,7 +1267,7 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
debug_info: Option<DebugInfo>,
) {
for i in 0..EF::D {
let j = (i * self.word_size) as i32;
let j = i as i32;
self.push(
AsmInstruction::DivF(dst.fp() + j, lhs.fp() + j, rhs.fp()),
debug_info.clone(),
Expand Down
3 changes: 3 additions & 0 deletions compiler/src/conversion/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@ use crate::asm::{AsmInstruction, AssemblyCode};

#[derive(Clone, Copy, Debug)]
pub struct CompilerOptions {
// The compiler will ensure that the heap pointer is aligned to be a multiple of `word_size`.
pub word_size: usize,
pub compile_prints: bool,
pub enable_cycle_tracker: bool,
pub field_arithmetic_enabled: bool,
Expand All @@ -17,6 +19,7 @@ pub struct CompilerOptions {
impl Default for CompilerOptions {
fn default() -> Self {
CompilerOptions {
word_size: 8,
compile_prints: true,
enable_cycle_tracker: false,
field_arithmetic_enabled: true,
Expand Down
32 changes: 24 additions & 8 deletions compiler/src/ir/builder.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -106,7 +106,9 @@ pub struct BuilderFlags {
/// Can compile to both assembly and a set of constraints.
#[derive(Debug, Clone, Default)]
pub struct Builder<C: Config> {
pub(crate) stack_ptr: u32,
pub(crate) var_count: u32,
pub(crate) felt_count: u32,
pub(crate) ext_count: u32,
pub operations: TracedVec<DslIr<C>>,
pub(crate) nb_public_values: Option<Var<C::N>>,
pub(crate) witness_var_count: u32,
Expand All @@ -120,13 +122,17 @@ pub struct Builder<C: Config> {
impl<C: Config> Builder<C> {
/// Creates a new builder with a given number of counts for each type.
pub fn new_sub_builder(
stack_ptr: u32,
var_count: u32,
felt_count: u32,
ext_count: u32,
nb_public_values: Option<Var<C::N>>,
flags: BuilderFlags,
bigint_repr_size: u32,
) -> Self {
Self {
stack_ptr,
var_count,
felt_count,
ext_count,
// Witness counts are only used when the target is a gnark circuit. And sub-builders are
// not used when the target is a gnark circuit, so it's fine to set the witness counts to 0.
witness_var_count: 0,
Expand Down Expand Up @@ -346,7 +352,9 @@ impl<C: Config> Builder<C> {
/// Evaluate a block of operations repeatedly (until a break).
pub fn do_loop(&mut self, mut f: impl FnMut(&mut Builder<C>) -> Result<(), BreakLoop>) {
let mut loop_body_builder = Builder::<C>::new_sub_builder(
self.stack_ptr,
self.var_count,
self.felt_count,
self.ext_count,
self.nb_public_values,
self.flags,
self.bigint_repr_size,
Expand Down Expand Up @@ -637,7 +645,9 @@ impl<'a, C: Config> IfBuilder<'a, C> {

// Execute the `then` block and collect the instructions.
let mut f_builder = Builder::<C>::new_sub_builder(
self.builder.stack_ptr,
self.builder.var_count,
self.builder.felt_count,
self.builder.ext_count,
self.builder.nb_public_values,
self.builder.flags,
self.builder.bigint_repr_size,
Expand Down Expand Up @@ -714,7 +724,9 @@ impl<'a, C: Config> IfBuilder<'a, C> {
"Cannot use dynamic branch in static mode"
);
let mut then_builder = Builder::<C>::new_sub_builder(
self.builder.stack_ptr,
self.builder.var_count,
self.builder.felt_count,
self.builder.ext_count,
self.builder.nb_public_values,
self.builder.flags,
self.builder.bigint_repr_size,
Expand All @@ -725,7 +737,9 @@ impl<'a, C: Config> IfBuilder<'a, C> {
let then_instructions = then_builder.operations;

let mut else_builder = Builder::<C>::new_sub_builder(
self.builder.stack_ptr,
self.builder.var_count,
self.builder.felt_count,
self.builder.ext_count,
self.builder.nb_public_values,
self.builder.flags,
self.builder.bigint_repr_size,
Expand Down Expand Up @@ -897,7 +911,9 @@ impl<'a, C: Config> RangeBuilder<'a, C> {
let step_size = C::N::from_canonical_usize(self.step_size);
let loop_variable: Var<C::N> = self.builder.uninit();
let mut loop_body_builder = Builder::<C>::new_sub_builder(
self.builder.stack_ptr,
self.builder.var_count,
self.builder.felt_count,
self.builder.ext_count,
self.builder.nb_public_values,
self.builder.flags,
self.builder.bigint_repr_size,
Expand Down
12 changes: 6 additions & 6 deletions compiler/src/ir/types.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -410,8 +410,8 @@ impl<C: Config> Variable<C> for Var<C::N> {
type Expression = SymbolicVar<C::N>;

fn uninit(builder: &mut Builder<C>) -> Self {
builder.stack_ptr += 1;
Var(builder.stack_ptr, PhantomData)
builder.var_count += 1;
Var(builder.var_count, PhantomData)
}

fn assign(&self, src: Self::Expression, builder: &mut Builder<C>) {
Expand Down Expand Up @@ -780,8 +780,8 @@ impl<C: Config> Variable<C> for Felt<C::F> {
type Expression = SymbolicFelt<C::F>;

fn uninit(builder: &mut Builder<C>) -> Self {
builder.stack_ptr += 1;
Felt(builder.stack_ptr, PhantomData)
builder.felt_count += 1;
Felt(builder.felt_count, PhantomData)
}

fn assign(&self, src: Self::Expression, builder: &mut Builder<C>) {
Expand Down Expand Up @@ -1209,8 +1209,8 @@ impl<C: Config> Variable<C> for Ext<C::F, C::EF> {
type Expression = SymbolicExt<C::F, C::EF>;

fn uninit(builder: &mut Builder<C>) -> Self {
builder.stack_ptr += <Self as MemVariable<C>>::size_of() as u32;
Ext(builder.stack_ptr, PhantomData)
builder.ext_count += 1;
Ext(builder.ext_count, PhantomData)
}

fn assign(&self, src: Self::Expression, builder: &mut Builder<C>) {
Expand Down
Loading

0 comments on commit ffffc5b

Please sign in to comment.