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.

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 Chips from a specific ExecutionRecord #9

Merged
merged 4 commits into from
Jun 6, 2024

Conversation

huitseeker
Copy link
Contributor

The Problem

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 (e.g. with several ExecutionRecords), and the following starts the changes towards accomplishing that.

How it works

We set up the general framework with:

/// 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;
}

/// The Record can provide events for the chip T
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 Self::Record) -> RowMajorMatrix<F>;
    ...
}

(the change to output is similar and pending)

then in AddSubChip:

impl<'a> WithEvents<'a> for AddSubChip {
    type Events = (
        // add events
        &'a [AluEvent],
        // sub events
        &'a [AluEvent],
    );
}

In the ExecutionRecord:

impl EventLens<AddSubChip> for ExecutionRecord {
    fn events(&self) -> <AddSubChip as crate::air::WithEvents>::Events {
        (&self.add_events, &self.sub_events)
    }
}

In generate_trace:

    fn generate_trace<EL: EventLens<Self>>(
        &self,
        input: &EL,
        output: &mut Self::Record,
    ) -> 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<_>>();
    ...

Important

As a side-effect of this change, each Record defines instantiations of the events it supports precisely, including, in the case of the ExecutionRecord, monomorphically. This means the numerous traits WithAddition, WithDoubling, etc .. we introduced to make a gadget generic over a field / elliptic curve know where to find the right events for itself in the ExecutionRecord are now obsolete and all replaced by this more general pattern: from the PoV of a Chip, a Record is an EventLens<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

  • convert the output parameter of generate_trace to not need an ExecutionRecord, then integrate the actual restricted construction from #232,
  • convert the prover to not need the ExecutionRecord to be defined as an associated type of the MachineAir(but passed as a "free" parameter),
  • same for the recursion prover.

@huitseeker huitseeker requested review from adr1anh and wwared June 4, 2024 18:14
@huitseeker huitseeker force-pushed the chip_as_ref branch 2 times, most recently from d202458 to 3edcb96 Compare June 5, 2024 00:00
@huitseeker huitseeker marked this pull request as ready for review June 5, 2024 18:49
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<_>>();
    ...
```
Copy link
Contributor

@adr1anh adr1anh left a 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

Copy link
Contributor

@wwared wwared left a 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

Comment on lines +91 to +93
/// 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.
Copy link
Contributor

@wwared wwared Jun 6, 2024

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

Copy link
Contributor Author

@huitseeker huitseeker Jun 6, 2024

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.

@huitseeker huitseeker merged commit 8dd4513 into dev Jun 6, 2024
6 checks passed
@huitseeker huitseeker deleted the chip_as_ref branch June 6, 2024 12:57
huitseeker added a commit that referenced this pull request Jun 11, 2024
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),
            }
        }
    }
}
```
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants