facebook / winterfell

A STARK prover and verifier for arbitrary computations

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Accomodating more expressive transition constraints

Al-Kindi-0 opened this issue · comments

Currently, support for more general transition constraints is lacking. Consider the case of the constraints enforced on the so-called Lagrange kernel in Appendix A here. These boil down to:

  1. The boundary constraint $c(X) - (1 - r_0)\dots (1 - r_{\nu-1}) = 0$ at $X = 1$, where $\nu = \log_2(n)$ and $n$ is the trace length.
  2. For each ${\kappa}\in\lbrace 1,\dots, \nu\rbrace$, the transition constraint $$</code>r_{\nu - \kappa}\cdot c(X) - (1 -r_{\nu - \kappa})\cdot c(g^{2^{\nu - \kappa} }\cdot X) = 0<code>$$ holds on $2^{\kappa - 1}$-th roots of unity i.e. with divisor $X^{2^{\kappa - 1} } - 1$.

Thus we have $\log_2(n)$ constraints each with its own divisor i.e., enforcement domain. I haven't included exception points to keep things simple, but these should be straightforward to accommodate, I believe.

Now the above introduces certain challenges to the way we currently do things. For example, for the OOD evaluations, we now need to provide, in addition to $c(z)$ and $c(g\cdot z)$, $c(g^{2^{\nu - \kappa} }\cdot z)$ for ${\kappa}\in\lbrace 1,\dots, \nu - 1\rbrace$.

The question now is: should we generalize the code base from the ground up or is there a way to "delegate" accommodating such relatively rare cases to a per case custom solution?

I have a few preliminary thoughts on this. At a high-level, this design generalizes the codebase, while making the most common case of "transition constraint applies over the entire domain" convenient. Note that I use "application" to mean the users of Winterfell (e.g. the Miden VM).

We would introduce a new object TransitionConstraintGroup that would encode a set of transition constraints who share a common enforcement domain. It would look roughly like:

struct TransitionConstraintGroup<A: Air> {
  // Note: each transition function *returns* its evaluation
  transition_functions: Vec<fn(&A, &EvaluationFrame, &[E]) -> E>,
  // `Domain` would be a type similar to the current `ConstraintDivisor`,
  // perhaps more user-friendly
  enforcement_domain: Domain
}

impl<A: Air> TransitionConstraintGroup<A: Air> {
  // Most commonly used constructor. Domain defaults to full trace length (minus 1).
  fn new(transition_functions: ...) -> Self { ... }
  // Used more rarely; allows to specify the domain explicitly
  fn new_with_domain(transition_functions: ..., domain: Domain) -> Self { ... }
}

A Vec<TransitionConstraintGroup> would be passed to Air::new(), which would replace the current Air::evaluate_transition(). For example, when evaluating the main trace here, the DefaultConstraintEvaluator would only call the transition functions inside the groups for which $step \in domain$.

A few thoughts:

  • This is one step towards moving away from "out parameters" such as result in Air::evaluate_transition(). With the new approach, each transition function returns its evaluation as opposed to writing it in the appropriate location in the result buffer. This has many implications
    • The "evaluate transition functions" API is easier to read
    • Performance either won't change or increase slightly, since the current approach write to the buffer twice (once to zero the buffer, and another by the application to write the actual value). Also, with the current API, if the application has a Vec<E> of evaluations, it is forced to copy it into result as opposed to being able to move the Vec out. This effectively forces on the application the "heavy use of indices" programming style that we see in the Miden VM.
    • The application won't be forced to maintain error-prone arrays of transition constraint degrees anymore
  • We can potentially execute the transition constraints in parallel (if we see that performance typically increases by doing so)
  • We can create a TransitionConstraint object that wraps a fn(...) -> E to enhance readability

Now the above introduces certain challenges to the way we currently do things. For example, for the OOD evaluations, ...

I'm not familiar with the specifics of how we do OOD evaluations, so I can't say if this proposal makes this easier/harder.

With the new approach, each transition function returns its evaluation as opposed to writing it in the appropriate location in the result buffer.

This may impact performance in a pretty significant way for constraints which share many intermediary values. An example of these are Rescue hash constraints. Here, we have 4 constraints which where computing just one of them is almost as much work as computing all 4.

There could be situations where even more constraints are related in such a way that individual evaluations require significantly more work than evaluating them all together. If I remember correctly, that was the main reason using array buffers for collecting results.

We can potentially execute the transition constraints in parallel (if we see that performance typically increases by doing so)

This should be already happening. In fact, constraint evaluation step scales the best as we scale the number of cores (I believe it is almost 1-to-1 scaling factor). The parallelization is done across rows rather than constraints - so, if we find a way to use vector instructions for constraint evaluations, there still may be benefit to computing multiple constraints at the same time in a single core.

The application won't be forced to maintain error-prone arrays of transition constraint degrees anymore

We may actually not need to maintain them regardless. Previously, these were used to align the degrees of all polynomials - but this is no longer needed. We should check this, but I believe constraint degrees are used mostly for debugging purposes in the prover, and are not really used by the verifier. If this is the case, then we may be able to get rid of them and then, we could rid of the result vector as you've suggested above (the idea is to evaluate a constraint and immediately merge it into the accumulator of all evaluated constraints).

One other issue worth mentioning is that to support something similar to what is described in the original post we'd also need to allow more complex evaluation frames (see #80). Doing this efficiently is non-trivial and it might be better to avoid solving this in full generality.

An alternative solution may be to define special type of constraints called "Lagrange kernel" constraints. Then, as a part of Air trait, we could define which column(s) we'd want to enforce Lagrange kernel constraints on. Maybe it could look something like this:

pub trait Air {
    ...
    /// Returns index of the column in the auxiliary trace on which Lagrange kernel constraints
    /// should be enforced.
    ///
    /// This can also be generalized to return Vec<usize> if we think there may be more than
    /// one such column.
    fn lagrange_kernel_column(&self) -> Option<usize>;
    ...
}

I'm not sure if the above is sufficient to compute the Lagrange kernel constraints. If not, we may need to return more info from this method than just the column index.

The actual constraint evaluation logic can then be kind of "hard-coded" into both the prover and the verifier. We should be able to do this more efficiently and with much less effort than supporting "general purpose" constraints on large enforcement domains.

Closed by #247 where we decided to use specialized constraints for the Lagrange kernel column.