Open plafer opened 1 month ago
Thank you for such a detailed write up! The general direction of this makes a lot of sense to me - but I do have a few questions/comments.
First, the functions in the proposed interface return vectors in many cases. Given that these will be called for every row in an execution trace, I've been trying to stay away from returning vectors as it could result in extra allocations. I haven't actually measured the impact of returning vectors - so, to be honest, not sure if it is negligible or quite significant. But I think before we commit to using vectors in the interface, would be good to understand if we'll experience a noticeable degradation in performance. Or alternatively, we could refactor interfaces to avoid allocating vectors on every call.
Also related to the above, I wonder what the impact of using Box<dyn Fn(&EvaluationFrame<E>, &[E]) -> Vec<E>>
would be as I think the compiler would not be able to inline calls to these functions.
For LogUp-GKR, would we actually need this structure? If building of the Lagrange kernel and s
columns (assuming we still need it) moves to Winterfell, I think we would not need to build any auxiliary columns here. The only thing that would be needed are the changes to the Air
interfaces discussed in https://github.com/0xPolygonMiden/miden-vm/issues/1386#issuecomment-2227899609. So, I think things like MultisetCheck
and MultisetCheckWithMultiplicities
and their implementations should probably live in the air
crate of the VM.
I also wonder if there is value in separating MultisetCheck
and MultisetCheckWithMultiplicities
. I understand that they are conceptually different, but it seems like their interfaces and treatments are basically the same. Maybe we could collapse them into a single interface?
Also, when building request/response flags and values, we may need access to periodic column values - so, it might make sense to add them to the input parameters.
First, the functions in the proposed interface return vectors in many cases. Given that these will be called for every row in an execution trace, I've been trying to stay away from returning vectors as it could result in extra allocations.
Indeed, I was only returning Vec<E>
for illustration purposes; the MultisetCheck
trait has a note about that, but maybe that was too subtle of a place to put it. In any case, yes if we find that it is indeed faster to write into a buffer
instead, then we should modify the interface accordingly.
I wonder what the impact of using
Box<dyn Fn(&EvaluationFrame<E>, &[E]) -> Vec<E>>
would be as I think the compiler would not be able to inline calls to these functions.
Good point. We should then modify the interfaces to e.g. (again, not necessarily returning Vec
):
trait MultisetCheck<E: FieldElement, Frame> {
// ...
/// The implementation takes care of evaluating all its "inner functions",
/// and only returning the result.
fn get_requests(&self, frame: &Frame, alphas: &[E]) -> Vec<Vec<E>>;
// ...
}
For LogUp-GKR, would we actually need this structure?
With LogUp-GKR, we would only need the MultisetCheck
and MultisetCheckWithMultiplicities
traits, which would be used by the evaluation of the input layer.
However, I think it's still worth implementing the running sum/product columns (although at a lower priority), specifically for A/B testing. I would like to get to a point where, for example, answering the question "would it be faster to use a running product column for this bus instead of adding it to LogUp-GKR?" could be answered in a matter of at most a few hours. If it turned out that it was, it would have repercussions on AirScript
, but I'm punting that problem down the road for if/when we get there.
I also wonder if there is value in separating
MultisetCheck
andMultisetCheckWithMultiplicities
.
The main value is to allow proving systems to only be implemented for the appropriate version of a multiset check. The biggest case we want to avoid is allowing a running product column to prove a MultisetCheckWithMultiplicities
(for which it would not be possible to write constraints in general).
And in fact they are treated differently. For example, RangeChecker
only implements MultisetCheckWithMultiplicities
, and so would not be able to be proved by a running product column.
The decision to represent flags explicitly in the MultisetCheck
was only a convenience for the user. When flags are explicitly represented, then we know how to build a MultisetCheckWithMultiplicities
from a MultisetCheck
, since it is always better to put the flags in the multiplicity position. And so we have a blanket implementation that does just that. And so ultimately, the relationship we have is "every MultisetCheck
is a MultisetCheckWithMultiplicities
but not every MultisetCheckWithMultiplicities
is a MultisetCheck
.
Also, when building request/response flags and values, we may need access to periodic column values - so, it might make sense to add them to the input parameters.
Ah yes, we should add those too.
I think it's still worth implementing the running sum/product columns (although at a lower priority), specifically for A/B testing. I would like to get to a point where, for example, answering the question "would it be faster to use a running product column for this bus instead of adding it to LogUp-GKR?" could be answered in a matter of at most a few hours.
Agreed!
The main value is to allow proving systems to only be implemented for the appropriate version of a multiset check. The biggest case we want to avoid is allowing a running product column to prove a
MultisetCheckWithMultiplicities
(for which it would not be possible to write constraints in general).And in fact they are treated differently. For example,
RangeChecker
only implementsMultisetCheckWithMultiplicities
, and so would not be able to be proved by a running product column.
Makes sense. But I'm wonder since there are no interface differences whether MultisetCheckWithMultiplicities
could be just a marker trait. That is, it could look something like this:
pub trait MultisetCheck<E: FieldElement> {
...
}
pub trait MultisetCheckWithMultiplicities<E: FieldElement>: MultisetCheck<E>;
Also, thinking about this more, I'm wondering if should combine request/response and their corresponding flags/multiplicities into a single struct. Maybe something like:
pub struct Request<E: FieldElement> {
pub latch: E, // equivalent to numerator
pub value: E, // equivalent to denominator
}
// Could also be combined with `Request` - but not sure under what time
pub struct Response<E: FieldElement> {
pub latch: E, // equivalent to numerator
pub value: E, // equivalent to denominator
}
And then MultisetCheck
could look something like this:
pub trait MultisetCheck<F: FieldElement> {
fn fill_requests<E: FieldElement<BaseField = F> + ExtensionOf<F>>(
&self,
frame: &EvaluationFrame<F>,
periodic_values: &[F],
&mut result[Request<E>],
);
fn fill_responses<E: FieldElement<BaseField = F> + ExtensionOf<F>>(
&self,
frame: &EvaluationFrame<F>,
periodic_values: &[F],
&mut result[Response<E>],
);
}
We could also simplify this by keeping the inputs and output in the same filed, and then handle the final adjustment by randomness separately (i.e., the returned denominator would be just $v$ instead of $\alpha + v$). In this case, we'd have something like this:
pub trait MultisetCheck<E: FieldElement> {
fn fill_requests(
&self,
frame: &EvaluationFrame<E>,
periodic_values: &[E],
linear_comb_rands: &[E],
&mut result[Request<E>],
);
fn fill_responses(
&self,
frame: &EvaluationFrame<E>,
periodic_values: &[E],
linear_comb_rands: &[E],
&mut result[Response<E>],
);
}
I'd like to suggest a refactoring of our various multiset checks; this will be necessary soon enough when we integrate them into LogUp-GKR anyways. The current abstractions we have couple a few distinct concerns:
AuxColumnBuilder
trait. While this is a step in the right direction, this abstraction couples the description of the multiset check along with the fact that it will be proved using a running product column (e.g. here).AuxColumnBuilder
does with the running product column).Now, with the new LogUp-GKR proving system, we will need to rewrite all multiset checks that use the
AuxColumnBuilder
, since as described, are currently coupled to a running product column implementation.The ideal multiset check abstractions would allow us to describe a multiset check instance once, and then swap between various compatible ways of proving it (e.g. running product column, running sum column, GKR, etc.). Also ideally, our abstractions for building running sum/product columns would be general enough to work out of the box with any "compatible" multiset check instance (more on that later). While we are thinking of switching over all multiset checks to be proved using LogUp-GKR, I still think it would be valuable to quickly swap some multiset checks in/out of LogUp-GKR, and be proved using a running product/sum column instead, for the purposes of benchmarking, and A/B testing more generally.
Below, I will first describe how I think about the different "variants" of multiset checks, which will justify the set of abstractions that I describe next.
Background
The 2 main categories of multiset checks are, as identified in our docs,
It is important to distinguish between the two, since for example, a running product column can be used to prove/instantiate multiset checks, but not multiset checks with multiplicities (at least not with general polynomials as multiplicities). We will use the following mathematical models for each:
$$ \prod{i=0}^{n-1} \prod{j=0}^{k-1} a{ij}(\psi) = \prod{i=0}^{n-1} \prod{j=0}^{k-1} b{ij}(\psi) $$
$$ \prod{i=0}^{n-1} \prod{j=0}^{k-1} a{ij}(\psi)^{m{a{ij}}} = \prod{i=0}^{n-1} \prod{j=0}^{k-1} b{ij}(\psi)^{m{b{ij}}} $$
Here, $n$ is the length of the trace, and $k$ is the maximum number of terms that can be generated in a single row. Also, $\psi = (\psi0, \dots, \psi{c-1})$ is a set of random elements that the multiset elements are allowed to use. In the current codebase, these are the
alphas
that are used e.g. to merge virtual table columns together. Finally, $m{a{ij}}$ and $m{b{ij}}$ are polynomials derived from the trace that act as the multiplicity for the corresponding $a{ij}$ or $b{ij}$, respectively. Those currently don't have access to $\psi$, but if that turns out to be useful, they certainly could.A very important point that is often overlooked is the question of which rows do $a{ij}(\psi)$ and $b{ij}(\psi)$ have access to? This will end up affecting how the underlying proving system (running product/sum/GKR) looks like. The 2 real options right now are:
Currently in the VM, the range checker is the only multiset check that only needs row $i$. All the other ones need rows $i$ and $i+1$. Note however that this is a property of the proving system, in that it doesn't only affect multiset checks, but any column built using it. To be concrete, the "s" auxiliary column in LogUp-GKR is a built as a running sum (since it implements an inner product); this running sum, even if not a multiset check, will be built different whether "s" needs access to row $i$ only or both rows $i$ and $i+1$.
Normally, you'd think that since
2
is more general than1
, every problem1
could be modeled as2
. But there's one important caveat:2
is not a strict generalization of1
, since in2
, the last row is not allowed to contain multiset elements (since they're defined in terms of 2 rows). This is the case with the range checker: we are building the running sum column in a way that you would build a running sum column for2
(even though the range checker elements never use rowi+1
), since we ensure that the last row of the trace is aHALT
operation, and hence cannot contain a range check request (see #1383). However, this isn't possible with LogUp-GKR's "s" column: there always needs to be a multiset element generated on the last row, and so we cannot reuse a running sum column built for2
.There are a few ways to build a column for
1
, but the one we've honed in for LogUp-GKR's "s" column is described in https://github.com/facebook/winterfell/issues/286. As for building a column for2
, this is exactly how the currentAuxColumnBuilder
does it.Proposed abstractions
The key takeaways from the previous section are:
MultisetCheck
andMultisetCheckWithMultiplicities
MultisetCheck
instanceMultisetCheck
instanceMultisetCheckWithMultiplicity
instanceMultiset check abstractions
These are the types to encode multiset checks, along with a few examples.
Running product/sum column abstractions
These are the abstractions that build running product/sum columns from a multiset check.
LogUp-GKR integration
Below is an idea of how the current LogUp-GKR backend would use the new multiset check abstractions, although we probably will end up doing something slightly different since LogUp-GKR is more in flux. This serves as a small example of how it would work.