feat(avm): bounded mle implementation (#8668)

Resolves #8651

barretenberg/cpp/src/barretenberg/polynomials/polynomial.cpp

@@ -2,6 +2,7 @@
 #include "barretenberg/common/assert.hpp"
 #include "barretenberg/common/slab_allocator.hpp"
 #include "barretenberg/common/thread.hpp"
+#include "barretenberg/numeric/bitop/get_msb.hpp"
 #include "barretenberg/numeric/bitop/pow.hpp"
 #include "barretenberg/polynomials/shared_shifted_virtual_zeroes_array.hpp"
 #include "polynomial_arithmetic.hpp"
@@ -10,6 +11,7 @@
 #include <list>
 #include <memory>
 #include <mutex>
+#include <span>
 #include <sys/stat.h>
 #include <unordered_map>
 #include <utility>
@@ -185,16 +187,21 @@ template <typename Fr> Fr Polynomial<Fr>::evaluate(const Fr& z) const
 
 template <typename Fr> Fr Polynomial<Fr>::evaluate_mle(std::span<const Fr> evaluation_points, bool shift) const
 {
-    const size_t m = evaluation_points.size();
+    if (size() == 0) {
+        return Fr(0);
+    }
+
+    const size_t n = evaluation_points.size();
+    const size_t dim = numeric::get_msb(end_index() - 1) + 1; // Round up to next power of 2
 
     // To simplify handling of edge cases, we assume that the index space is always a power of 2
-    ASSERT(virtual_size() == static_cast<size_t>(1 << m));
+    ASSERT(virtual_size() == static_cast<size_t>(1 << n));
 
-    // we do m rounds l = 0,...,m-1.
+    // We first fold over dim rounds l = 0,...,dim-1.
     // in round l, n_l is the size of the buffer containing the Polynomial partially evaluated
     // at u₀,..., u_l.
-    // in round 0, this is half the size of n
-    size_t n_l = 1 << (m - 1);
+    // In round 0, this is half the size of dim
+    size_t n_l = 1 << (dim - 1);
 
     // temporary buffer of half the size of the Polynomial
     // TODO(https://github.com/AztecProtocol/barretenberg/issues/1096): Make this a Polynomial with DontZeroMemory::FLAG
@@ -215,15 +222,22 @@ template <typename Fr> Fr Polynomial<Fr>::evaluate_mle(std::span<const Fr> evalu
         const size_t ALLOW_ONE_PAST_READ = 1;
         tmp[i] = get(i * 2 + offset) + u_l * (get(i * 2 + 1 + offset, ALLOW_ONE_PAST_READ) - get(i * 2 + offset));
     }
-    // partially evaluate the m-1 remaining points
-    for (size_t l = 1; l < m; ++l) {
-        n_l = 1 << (m - l - 1);
+
+    // partially evaluate the dim-1 remaining points
+    for (size_t l = 1; l < dim; ++l) {
+        n_l = 1 << (dim - l - 1);
         u_l = evaluation_points[l];
         for (size_t i = 0; i < n_l; ++i) {
             tmp[i] = tmp[i * 2] + u_l * (tmp[(i * 2) + 1] - tmp[i * 2]);
         }
     }
-    Fr result = tmp[0];
+    auto result = tmp[0];
+
+    // We handle the "trivial" dimensions which are full of zeros.
+    for (size_t i = dim; i < n; i++) {
+        result *= (Fr(1) - evaluation_points[i]);
+    }
+
     return result;
 }
 

barretenberg/cpp/src/barretenberg/polynomials/polynomial.hpp

@@ -143,14 +143,19 @@ template <typename Fr> class Polynomial {
     Polynomial shifted() const;
 
     /**
-     * @brief evaluates p(X) = ∑ᵢ aᵢ⋅Xⁱ considered as multi-linear extension p(X₀,…,Xₘ₋₁) = ∑ᵢ aᵢ⋅Lᵢ(X₀,…,Xₘ₋₁)
-     * at u = (u₀,…,uₘ₋₁)
+     * @brief evaluate multi-linear extension p(X_0,…,X_{n-1}) = \sum_i a_i*L_i(X_0,…,X_{n-1}) at u = (u_0,…,u_{n-1})
+     *        If the polynomial is embedded into a lower dimension k<n, i.e, start_index + size <= 2^k,
+     *        we evaluate it in a more efficient way. Note that a_j == 0 for any j >= 2^k.
+     *        We fold over k dimensions and then multiply the result by
+     *        (1 - u_k) * (1 - u_{k+1}) ... * (1 - u_{n-1}). In this case, for any
+     *        i < 2^k, L_i is a multiple of (1 - X_k) * (1 - X_{k+1}) ... * (1 - X_{n-1}). Dividing
+     *        p by this monomial leads to a multilinear extension over variables X_0, X_1, ..X_{k-1}.
      *
-     * @details this function allocates a temporary buffer of size n/2
+     * @details this function allocates a temporary buffer of size 2^(k-1)
      *
-     * @param evaluation_points an MLE evaluation point u = (u₀,…,uₘ₋₁)
-     * @param shift evaluates p'(X₀,…,Xₘ₋₁) = 1⋅L₀(X₀,…,Xₘ₋₁) + ∑ᵢ˲₁ aᵢ₋₁⋅Lᵢ(X₀,…,Xₘ₋₁) if true
-     * @return Fr p(u₀,…,uₘ₋₁)
+     * @param evaluation_points evaluation vector of size n
+     * @param shift a boolean and when set to true, we evaluate the shifted counterpart polynomial:
+     *              enforce a_0 == 0 and compute \sum_i a_{i+1}*L_i(X_0,…,X_{n-1})
      */
     Fr evaluate_mle(std::span<const Fr> evaluation_points, bool shift = false) const;
 

barretenberg/cpp/src/barretenberg/polynomials/polynomial_arithmetic.test.cpp

@@ -4,6 +4,7 @@
 #include "barretenberg/numeric/random/engine.hpp"
 #include "barretenberg/polynomials/evaluation_domain.hpp"
 #include "legacy_polynomial.hpp"
+#include "polynomial.hpp"
 #include <algorithm>
 #include <cstddef>
 #include <gtest/gtest.h>
@@ -1090,7 +1091,8 @@ TYPED_TEST(PolynomialTests, interpolation_constructor)
     }
 }
 
-TYPED_TEST(PolynomialTests, evaluate_mle)
+// LegacyPolynomials MLE
+TYPED_TEST(PolynomialTests, evaluate_mle_legacy)
 {
     using FF = TypeParam;
 
@@ -1147,6 +1149,52 @@ TYPED_TEST(PolynomialTests, evaluate_mle)
     test_case(2);
 }
 
+/*
+ * @brief Compare that the mle evaluation over Polynomials match the mle evaluation of
+ *        LegacyPolynomials.
+ */
+TYPED_TEST(PolynomialTests, evaluate_mle)
+{
+    using FF = TypeParam;
+
+    auto test_case = [](size_t n, size_t dim) {
+        auto& engine = numeric::get_debug_randomness();
+        size_t k = 1 << dim;
+
+        std::vector<FF> evaluation_points;
+        evaluation_points.resize(n);
+        for (size_t i = 0; i < n; i++) {
+            evaluation_points[i] = FF::random_element(&engine);
+        }
+
+        LegacyPolynomial<FF> legacy_poly(1 << n);
+        Polynomial<FF> poly(k, 1 << n);
+        for (size_t i = 0; i < k; ++i) {
+            auto const rnd = FF::random_element(&engine);
+            legacy_poly[i] = rnd;
+            poly.at(i) = rnd;
+        }
+
+        const FF legacy_res = legacy_poly.evaluate_mle(evaluation_points);
+        const FF res = poly.evaluate_mle(evaluation_points);
+
+        EXPECT_EQ(legacy_res, res);
+
+        // Same with shifted polynomials
+        legacy_poly[0] = FF(0);
+        poly.at(0) = FF(0);
+
+        const FF legacy_shift_res = legacy_poly.evaluate_mle(evaluation_points, true);
+        const FF shift_res = poly.evaluate_mle(evaluation_points, true);
+
+        EXPECT_EQ(legacy_shift_res, shift_res);
+    };
+
+    test_case(9, 3);
+    test_case(8, 8);
+    test_case(13, 1);
+}
+
 /**
  * @brief Test the function for partially evaluating MLE polynomials
  *

barretenberg/cpp/src/barretenberg/vm/avm/generated/verifier.cpp

@@ -28,17 +28,9 @@ AvmVerifier& AvmVerifier::operator=(AvmVerifier&& other) noexcept
 using FF = AvmFlavor::FF;
 
 // Evaluate the given public input column over the multivariate challenge points
-[[maybe_unused]] inline FF evaluate_public_input_column(const std::vector<FF>& points,
-                                                        const size_t circuit_size,
-                                                        std::vector<FF> challenges)
+inline FF AvmVerifier::evaluate_public_input_column(const std::vector<FF>& points, std::vector<FF> challenges)
 {
-
-    // TODO(https://github.com/AztecProtocol/aztec-packages/issues/6361): we pad the points to the circuit size in order
-    // to get the correct evaluation. This is not efficient, and will not be valid in production.
-    std::vector<FF> new_points(circuit_size, 0);
-    std::copy(points.begin(), points.end(), new_points.data());
-
-    Polynomial<FF> polynomial(new_points);
+    Polynomial<FF> polynomial(points, key->circuit_size);
     return polynomial.evaluate_mle(challenges);
 }
 
@@ -109,34 +101,32 @@ bool AvmVerifier::verify_proof(const HonkProof& proof,
     std::vector<FF> mle_challenge(multivariate_challenge.begin(),
                                   multivariate_challenge.begin() + static_cast<int>(log_circuit_size));
 
-    FF main_kernel_inputs_evaluation = evaluate_public_input_column(public_inputs[0], circuit_size, mle_challenge);
+    FF main_kernel_inputs_evaluation = evaluate_public_input_column(public_inputs[0], mle_challenge);
     if (main_kernel_inputs_evaluation != claimed_evaluations.main_kernel_inputs) {
         vinfo("main_kernel_inputs_evaluation failed");
         return false;
     }
-    FF main_kernel_value_out_evaluation = evaluate_public_input_column(public_inputs[1], circuit_size, mle_challenge);
+    FF main_kernel_value_out_evaluation = evaluate_public_input_column(public_inputs[1], mle_challenge);
     if (main_kernel_value_out_evaluation != claimed_evaluations.main_kernel_value_out) {
         vinfo("main_kernel_value_out_evaluation failed");
         return false;
     }
-    FF main_kernel_side_effect_out_evaluation =
-        evaluate_public_input_column(public_inputs[2], circuit_size, mle_challenge);
+    FF main_kernel_side_effect_out_evaluation = evaluate_public_input_column(public_inputs[2], mle_challenge);
     if (main_kernel_side_effect_out_evaluation != claimed_evaluations.main_kernel_side_effect_out) {
         vinfo("main_kernel_side_effect_out_evaluation failed");
         return false;
     }
-    FF main_kernel_metadata_out_evaluation =
-        evaluate_public_input_column(public_inputs[3], circuit_size, mle_challenge);
+    FF main_kernel_metadata_out_evaluation = evaluate_public_input_column(public_inputs[3], mle_challenge);
     if (main_kernel_metadata_out_evaluation != claimed_evaluations.main_kernel_metadata_out) {
         vinfo("main_kernel_metadata_out_evaluation failed");
         return false;
     }
-    FF main_calldata_evaluation = evaluate_public_input_column(public_inputs[4], circuit_size, mle_challenge);
+    FF main_calldata_evaluation = evaluate_public_input_column(public_inputs[4], mle_challenge);
     if (main_calldata_evaluation != claimed_evaluations.main_calldata) {
         vinfo("main_calldata_evaluation failed");
         return false;
     }
-    FF main_returndata_evaluation = evaluate_public_input_column(public_inputs[5], circuit_size, mle_challenge);
+    FF main_returndata_evaluation = evaluate_public_input_column(public_inputs[5], mle_challenge);
     if (main_returndata_evaluation != claimed_evaluations.main_returndata) {
         vinfo("main_returndata_evaluation failed");
         return false;

barretenberg/cpp/src/barretenberg/vm/avm/generated/verifier.hpp

@@ -28,6 +28,9 @@ class AvmVerifier {
     std::shared_ptr<VerificationKey> key;
     std::map<std::string, Commitment> commitments;
     std::shared_ptr<Transcript> transcript;
+
+  private:
+    FF evaluate_public_input_column(const std::vector<FF>& points, std::vector<FF> challenges);
 };
 
 } // namespace bb

barretenberg/cpp/src/barretenberg/vm/aztec_constants.hpp

@@ -35,6 +35,7 @@
 #define PUBLIC_CONTEXT_INPUTS_LENGTH 42
 #define AVM_VERIFICATION_KEY_LENGTH_IN_FIELDS 66
 #define AVM_PROOF_LENGTH_IN_FIELDS 3817
+#define AVM_PUBLIC_COLUMN_MAX_SIZE_LOG2 8
 #define MEM_TAG_U1 1
 #define MEM_TAG_U8 2
 #define MEM_TAG_U16 3

bb-pilcom/bb-pil-backend/templates/verifier.cpp.hbs

@@ -28,17 +28,9 @@ namespace bb {
 using FF = {{name}}Flavor::FF;
 
 // Evaluate the given public input column over the multivariate challenge points
-[[maybe_unused]] inline FF evaluate_public_input_column(const std::vector<FF>& points,
-                                                        const size_t circuit_size,
-                                                        std::vector<FF> challenges)
+inline FF AvmVerifier::evaluate_public_input_column(const std::vector<FF>& points, std::vector<FF> challenges)
 {
-
-    // TODO(https://github.com/AztecProtocol/aztec-packages/issues/6361): we pad the points to the circuit size in order
-    // to get the correct evaluation. This is not efficient, and will not be valid in production.
-    std::vector<FF> new_points(circuit_size, 0);
-    std::copy(points.begin(), points.end(), new_points.data());
-
-    Polynomial<FF> polynomial(new_points);
+    Polynomial<FF> polynomial(points, key->circuit_size);
     return polynomial.evaluate_mle(challenges);
 }
 
@@ -109,7 +101,7 @@ bool {{name}}Verifier::verify_proof(const HonkProof& proof, [[maybe_unused]] con
                                   multivariate_challenge.begin() + static_cast<int>(log_circuit_size));
 
     {{#each public_cols}}
-    FF {{col}}_evaluation = evaluate_public_input_column(public_inputs[{{idx}}], circuit_size, mle_challenge);
+    FF {{col}}_evaluation = evaluate_public_input_column(public_inputs[{{idx}}], mle_challenge);
         if ({{col}}_evaluation != claimed_evaluations.{{col}}) {
             vinfo("{{col}}_evaluation failed");
             return false;

bb-pilcom/bb-pil-backend/templates/verifier.hpp.hbs

@@ -28,6 +28,9 @@ class {{name}}Verifier {
     std::shared_ptr<VerificationKey> key;
     std::map<std::string, Commitment> commitments;
     std::shared_ptr<Transcript> transcript;
+
+  private:
+    FF evaluate_public_input_column(const std::vector<FF>& points, std::vector<FF> challenges);
 };
 
 } // namespace bb

noir-projects/noir-protocol-circuits/crates/types/src/constants.nr

@@ -295,6 +295,7 @@ global AVM_VERIFICATION_KEY_LENGTH_IN_FIELDS: u32 = 2 + 16 * 4;
 // To determine latest value, hover `COMPUTED_AVM_PROOF_LENGTH_IN_FIELDS`
 // in barretenberg/cpp/src/barretenberg/vm/avm/generated/flavor.hpp
 global AVM_PROOF_LENGTH_IN_FIELDS: u32 = 3817;
+global AVM_PUBLIC_COLUMN_MAX_SIZE_LOG2 = 8;
 /**
  * Enumerate the hash_indices which are used for pedersen hashing.
  * We start from 1 to avoid the default generators. The generator indices are listed

yarn-project/circuits.js/src/constants.gen.ts

@@ -201,6 +201,7 @@ export const TUBE_PROOF_LENGTH = 439;
 export const VERIFICATION_KEY_LENGTH_IN_FIELDS = 128;
 export const AVM_VERIFICATION_KEY_LENGTH_IN_FIELDS = 66;
 export const AVM_PROOF_LENGTH_IN_FIELDS = 3817;
+export const AVM_PUBLIC_COLUMN_MAX_SIZE_LOG2 = 8;
 export const MEM_TAG_U1 = 1;
 export const MEM_TAG_U8 = 2;
 export const MEM_TAG_U16 = 3;