-
Notifications
You must be signed in to change notification settings - Fork 6
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.
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
refactor: De-couple Chip
s from a specific ExecutionRecord
#9
Conversation
d202458
to
3edcb96
Compare
This avoids relying on an instance of `MachineAir` on Chip<F, A>. - Modified the way `chip` object is accessed across different files in the core and recursion directories by using the `as_ref()` method. - Implemented `AsRef` trait for `Chip` struct for returning an Air reference, and removed `MachineAir` for `Chip` implementation. - Changed the method of accessing functions like `chip.name()`, `generate_trace()`, `preprocessed_width()` through `as_ref()` on `chip`. - Updated error handling in the `Verifier` struct in `stark/verifier.rs` to use `as_ref()`. - Made changes in the recursion program to access `preprocessed_data` using the `as_ref()` method. - Updated the reference of `chip` object in several parts of the `prove_shard` function in `stark/prover.rs`. - Revised accessing `chip` methods with the use of `as_ref()` across different functions in `stark/machine.rs`.
Currently the `MachineAir` trait requires every chip to define how it interacts with an `ExecutionRecord` associated type. This forces any chip that implements `MachineAir` to only interact with one specific implementation of an `ExecutionRecord` (as fixed at the moment of choosing that associated type). We would like those chips to be reused in more varied ways, and the following starts the changes towards accomplishing that. We set up the general framework with: ```rust /// A description of the events related to this AIR. pub trait WithEvents<'a>: Sized { /// output of a functional lens from the Record to /// refs of those events relative to the AIR. type Events: 'a; } pub trait EventLens<T: for<'a> WithEvents<'a>> { fn events(&self) -> <T as WithEvents>::Events; } pub trait MachineAir<F: Field>: BaseAir<F> + for<'a> WithEvents<'a> { ... fn generate_trace<EL: EventLens<Self>>(&self, input: &EL, output: &mut ExecutionRecord) -> RowMajorMatrix<F>; ... } ``` (the change to output is similar and pending) then in `AddSubChip`: ```rust impl<'a> WithEvents<'a> for AddSubChip { type Events = ( // add events &'a [AluEvent], // sub events &'a [AluEvent], ); } ``` In the `ExecutionRecord`: ```rust impl EventLens<AddSubChip> for ExecutionRecord { fn events(&self) -> <AddSubChip as crate::air::WithEvents>::Events { (&self.add_events, &self.sub_events) } } ``` In `generate_trace`: ```rust fn generate_trace<EL: EventLens<Self>>( &self, input: &EL, output: &mut EL, ) -> RowMajorMatrix<F> { let (add_events, sub_events) = input.events(); // Generate the rows for the trace. let chunk_size = std::cmp::max( (add_events.len() + sub_events.len()) / num_cpus::get(), 1, ); let merged_events = add_events .iter() .chain(sub_events.iter()) .collect::<Vec<_>>(); ... ```
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This looks really good and is a great first step towards removing the hard-coded dependency on the ExecutionRecord
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This PR has many goodies, and I really like it. Especially pleased with all the removed With...
traits we had.
The inline comment is just a question for clarifying my own understanding
/// The derived implementation is a tuple of the Events of each variant, | ||
/// in the variant declaration order. That is, because the chip could be *any* variant, | ||
/// it requires being able to provide for *all* event types consumable by each chip. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Asking for my own understanding: so if the Chip is defined as
#[derive(WithEvents)]
enum MyMachine {
Alu(AluChip),
Cpu(CpuChip),
}
Then the generated WithEvents
trait is something along the lines of
impl<'a> WithEvents<'a> for MyMachine {
type Events = (
<AluChip as WithEvents<'a>>::Events,
<CpuChip as WithEvents<'a>>::Events,
);
}
And the corresponding EventLens
derive looks something like
impl EventLens<MyMachine> for ExecutionRecord {
fn events(&self) -> <MyMachine as crate::air::WithEvents<'_>>::Events {
(
EventLens::<AluChip>::events(self),
EventLens::<CpuChip>::events(self),
)
}
}
(I probably got some minor details wrong above, ignore minor trait/type typos/mismatches)
In other words, every time you call events()
on this enum, you're getting the events for all the variants of the chip all at once, correct? But since inside each individual chip's generate_trace
function, we call the specialized EventLens
method for only that specific chip, we don't pay that cost every time events()
is called, only if we call it on the "large" MyMachine
enum?
Either way, asking this to highlight one reason that we do not want any manual impl EventLens
to ever do any kind of work beyond just returning a reference to a field of some data structure, and why the Events
type in impl WithEvents
should always be a reference, or this might lead to hidden overhead/unnecessary copies, since the derived Events
tuple can end up quite large for some enums
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
That's correct re: how the derive macros work.
Note, however, that you'd have to go out of your way to return anything else than a reference with a WithEvents
trait that has a lifetime parameter (thereby telling you "please go and use a reference here").
the derived Events tuple can end up quite large for some enums
See the Proj
construction for deducing a smaller EventLens
from that large tuple.
Previously: #9 This applies the same logic to output events in `generate_trace`. ```rust // implemented on the chip trait WithEvents<'a> { type InputEvents : 'a; type OutputEvents: 'a; } // implemented on the record trait EventLens<T: for <'a> WithEvent<'a>> { fn events<'a>(&'a self) -> <T as WithEvents<'a>>::InputEvents; } // implemented on the record trait EventMutLens<T: for <'a> WithEvent<'a>> { fn events<'a>(&'a mut self, events: <T as WithEvents<'a>>::OutputEvents); } ``` Now, one would wish that this would be a bit more ergonomic in Rust, because we do have anonymous cartesian products of arbitrary arity (tuples), but we do not have anonymous coproducts - either crate notwithstanding. ```rust pub enum DivRemEvent<'a> { ByteLookupEvent(&'a ByteLookupEvent), MulEvent(&'a AluEvent), LtEvent(&'a AluEvent), } impl<'a> WithEvents<'a> for DivRemChip { type InputEvents = &'a [AluEvent]; type OutputEvents = [DivRemEvent<'a>]; } impl EventMutLens<DivRemChip> for ExecutionRecord { fn add_events(&mut self, events: <DivRemChip as crate::air::WithEvents<'_>>::OutputEvents) { for event in events { match event { DivRemEvent::ByteLookupEvent(e) => self.add_byte_lookup_event(*e), DivRemEvent::MulEvent(e) => self.add_mul_event(*e), DivRemEvent::LtEvent(e) => self.add_lt_event(*e), } } } } ```
The Problem
Currently the
MachineAir
trait requires every chip to define how it interacts with anExecutionRecord
associated type.This forces any chip that implements
MachineAir
to only interact with one specific implementation of anExecutionRecord
(as fixed at the moment of choosing that associated type).We would like those chips to be reused in more varied ways (e.g. with several
ExecutionRecord
s), and the following starts the changes towards accomplishing that.How it works
We set up the general framework with:
(the change to output is similar and pending)
then in
AddSubChip
:In the
ExecutionRecord
:In
generate_trace
:Important
As a side-effect of this change, each
Record
defines instantiations of the events it supports precisely, including, in the case of theExecutionRecord
, monomorphically. This means the numerous traitsWithAddition
,WithDoubling
, etc .. we introduced to make a gadget generic over a field / elliptic curve know where to find the right events for itself in theExecutionRecord
are now obsolete and all replaced by this more general pattern: from the PoV of aChip
, aRecord
is anEventLens<Self>
which is thus generally tasked with telling you where to find the events relative to you. (/cc @storojs72 @wwared for the update)Next steps
output
parameter ofgenerate_trace
to not need anExecutionRecord
, then integrate the actual restricted construction from #232,ExecutionRecord
to be defined as an associated type of theMachineAir
(but passed as a "free" parameter),