Merge pull request #471 from kroma-network/perf/optimize-zkey-parsing

perf(circom): optimize zkey parsing

Author: Ryan Kim

vendors/circom/benchmark/circom_benchmark.cc

@@ -91,20 +91,15 @@ int RealMain(int argc, char** argv) {
       witness_loader.Set("in", Uint8ToBitVector<F>(in));
       witness_loader.Load();
 
-      const zk::r1cs::ConstraintMatrices<F>& constraint_matrices =
-          tachyon_runner->constraint_matrices();
-
-      domain = Domain::Create(constraint_matrices.num_constraints +
-                              constraint_matrices.num_instance_variables);
+      domain = Domain::Create(tachyon_runner->GetDomainSize());
 
+      size_t num_instance_variables = tachyon_runner->GetNumInstanceVariables();
       full_assignments = base::CreateVector(
-          constraint_matrices.num_instance_variables +
-              constraint_matrices.num_witness_variables,
+          num_instance_variables + tachyon_runner->GetNumWitnessVariables(),
           [&witness_loader](size_t i) { return witness_loader.Get(i); });
 
-      public_inputs =
-          absl::MakeConstSpan(full_assignments)
-              .subspan(1, constraint_matrices.num_instance_variables - 1);
+      public_inputs = absl::MakeConstSpan(full_assignments)
+                          .subspan(1, num_instance_variables - 1);
       CheckPublicInput(in, public_inputs);
     }
 

vendors/circom/benchmark/tachyon_runner.h

@@ -33,16 +33,22 @@ class TachyonRunner : public Runner<Curve, MaxDegree> {
     return proving_key_;
   }
 
-  const zk::r1cs::ConstraintMatrices<F>& constraint_matrices() const {
-    return constraint_matrices_;
+  size_t GetDomainSize() const { return zkey_->GetDomainSize(); }
+
+  size_t GetNumInstanceVariables() const {
+    return zkey_->GetNumInstanceVariables();
+  }
+
+  size_t GetNumWitnessVariables() const {
+    return zkey_->GetNumWitnessVariables();
   }
 
   void LoadZkey(const base::FilePath& zkey_path) override {
     zkey_ = ParseZKey<Curve>(zkey_path);
     CHECK(zkey_);
 
     proving_key_ = zkey_->GetProvingKey().ToNativeProvingKey();
-    constraint_matrices_ = zkey_->GetConstraintMatrices();
+    coefficients_ = zkey_->GetCoefficients();
   }
 
   zk::r1cs::groth16::Proof<Curve> Run(const Domain* domain,
@@ -53,15 +59,16 @@ class TachyonRunner : public Runner<Curve, MaxDegree> {
 
     std::vector<F> h_evals =
         QuadraticArithmeticProgram<F>::WitnessMapFromMatrices(
-            domain, constraint_matrices_, full_assignments);
+            domain, coefficients_, full_assignments);
 
+    size_t num_instance_variables = GetNumInstanceVariables();
     zk::r1cs::groth16::Proof<Curve> proof =
         zk::r1cs::groth16::CreateProofWithAssignmentNoZK(
             proving_key_, absl::MakeConstSpan(h_evals),
             absl::MakeConstSpan(full_assignments)
-                .subspan(1, constraint_matrices_.num_instance_variables - 1),
+                .subspan(1, num_instance_variables - 1),
             absl::MakeConstSpan(full_assignments)
-                .subspan(constraint_matrices_.num_instance_variables),
+                .subspan(num_instance_variables),
             absl::MakeConstSpan(full_assignments).subspan(1));
 
     delta = base::TimeTicks::Now() - now;
@@ -80,7 +87,7 @@ class TachyonRunner : public Runner<Curve, MaxDegree> {
   WitnessLoader<F> witness_loader_;
   std::unique_ptr<ZKey<Curve>> zkey_;
   zk::r1cs::groth16::ProvingKey<Curve> proving_key_;
-  zk::r1cs::ConstraintMatrices<F> constraint_matrices_;
+  absl::Span<const Coefficient<F>> coefficients_;
   std::optional<zk::r1cs::groth16::PreparedVerifyingKey<Curve>>
       prepared_verifying_key_;
 };

vendors/circom/circomlib/circuit/BUILD.bazel

@@ -16,6 +16,7 @@ tachyon_cc_library(
     name = "quadratic_arithmetic_program",
     hdrs = ["quadratic_arithmetic_program.h"],
     deps = [
+        "//circomlib/zkey:coefficient",
         "@kroma_network_tachyon//tachyon/base:logging",
         "@kroma_network_tachyon//tachyon/zk/r1cs/constraint_system:quadratic_arithmetic_program",
     ],

vendors/circom/circomlib/circuit/circuit_test.h

@@ -54,22 +54,18 @@ class CircuitTest : public testing::Test {
 
     zk::r1cs::groth16::ProvingKey<Curve> pk =
         zkey.GetProvingKey().ToNativeProvingKey();
-    zk::r1cs::ConstraintMatrices<F> constraint_matrices =
-        zkey.GetConstraintMatrices();
+    absl::Span<const Coefficient<F>> coefficients = zkey.GetCoefficients();
 
-    std::unique_ptr<Domain> domain =
-        Domain::Create(constraint_matrices.num_constraints +
-                       constraint_matrices.num_instance_variables);
+    std::unique_ptr<Domain> domain = Domain::Create(zkey.GetDomainSize());
     std::vector<F> h_evals = QAP::WitnessMapFromMatrices(
-        domain.get(), constraint_matrices, full_assignments);
+        domain.get(), coefficients, full_assignments);
 
+    size_t num_instance_variables = zkey.GetNumInstanceVariables();
     zk::r1cs::groth16::Proof<Curve> proof =
         zk::r1cs::groth16::CreateProofWithAssignmentZK(
             pk, absl::MakeConstSpan(h_evals),
-            full_assignments.subspan(
-                1, constraint_matrices.num_instance_variables - 1),
-            full_assignments.subspan(
-                constraint_matrices.num_instance_variables),
+            full_assignments.subspan(1, num_instance_variables - 1),
+            full_assignments.subspan(num_instance_variables),
             full_assignments.subspan(1));
     zk::r1cs::groth16::PreparedVerifyingKey<Curve> pvk =
         std::move(pk).TakeVerifyingKey().ToPreparedVerifyingKey();

vendors/circom/circomlib/circuit/quadratic_arithmetic_program.h

@@ -6,10 +6,10 @@
 #ifndef VENDORS_CIRCOM_CIRCOMLIB_CIRCUIT_QUADRATIC_ARITHMETIC_PROGRAM_H_
 #define VENDORS_CIRCOM_CIRCOMLIB_CIRCUIT_QUADRATIC_ARITHMETIC_PROGRAM_H_
 
-#include <memory>
 #include <utility>
 #include <vector>
 
+#include "circomlib/zkey/coefficient.h"
 #include "tachyon/base/logging.h"
 #include "tachyon/zk/r1cs/constraint_system/quadratic_arithmetic_program.h"
 
@@ -20,55 +20,46 @@ class QuadraticArithmeticProgram {
  public:
   QuadraticArithmeticProgram() = delete;
 
-  template <typename Domain>
-  static zk::r1cs::QAPInstanceMapResult<F> InstanceMap(
-      const Domain* domain, const zk::r1cs::ConstraintSystem<F>& cs,
-      const F& x) {
-    return zk::r1cs::QuadraticArithmeticProgram<F>::InstanceMap(domain, cs, x);
-  }
-
   template <typename Domain>
   static std::vector<F> WitnessMapFromMatrices(
-      const Domain* domain, const zk::r1cs::ConstraintMatrices<F>& matrices,
+      const Domain* domain, absl::Span<const Coefficient<F>> coefficients,
       absl::Span<const F> full_assignments) {
     using Evals = typename Domain::Evals;
     using DensePoly = typename Domain::DensePoly;
 
-    CHECK_GE(domain->size(), matrices.num_constraints);
-
     std::vector<F> a(domain->size());
     std::vector<F> b(domain->size());
     std::vector<F> c(domain->size());
 
-    // clang-format off
-    // |a[i]| = Σⱼ₌₀..ₘ (xⱼ * Aᵢ,ⱼ)    (if i < |num_constraints|)
-    //        = x[i - num_constraints] (otherwise)
-    // |b[i]| = Σⱼ₌₀..ₘ (xⱼ * Bᵢ,ⱼ)    (if i < |num_constraints|)
-    //        = 0                      (otherwise)
-    // |c[i]| = |a[i]|* |b[i]|         (if i < |num_constraints|)
-    //        = 0                      (otherwise)
-    // where x is |full_assignments|.
-    // clang-format on
-    OMP_PARALLEL {
-      OMP_FOR_NOWAIT
-      for (size_t i = 0; i < matrices.num_constraints; ++i) {
-        a[i] = zk::r1cs::EvaluateConstraint(matrices.a[i], full_assignments);
-      }
-
-      OMP_FOR
-      for (size_t i = 0; i < matrices.num_constraints; ++i) {
-        b[i] = zk::r1cs::EvaluateConstraint(matrices.b[i], full_assignments);
-      }
-
-      OMP_FOR
-      for (size_t i = 0; i < matrices.num_constraints; ++i) {
-        c[i] = a[i] * b[i];
+    // See
+    // https://github.com/iden3/rapidsnark/blob/b17e6fed08e9ceec3518edeffe4384313f91e9ad/src/groth16.cpp#L116-L156.
+#if defined(TACHYON_HAS_OPENMP)
+    constexpr size_t kNumLocks = 1024;
+    omp_lock_t locks[kNumLocks];
+    for (size_t i = 0; i < kNumLocks; i++) omp_init_lock(&locks[i]);
+#endif
+    OPENMP_PARALLEL_FOR(size_t i = 0; i < coefficients.size(); i++) {
+      const Coefficient<F>& c = coefficients[i];
+      std::vector<F>& ab = (c.matrix == 0) ? a : b;
+
+#if defined(TACHYON_HAS_OPENMP)
+      omp_set_lock(&locks[c.constraint % kNumLocks]);
+#endif
+      if (c.value.IsOne()) {
+        ab[c.constraint] += full_assignments[c.signal];
+      } else {
+        ab[c.constraint] += c.value * full_assignments[c.signal];
       }
+#if defined(TACHYON_HAS_OPENMP)
+      omp_unset_lock(&locks[c.constraint % kNumLocks]);
+#endif
     }
+#if defined(TACHYON_HAS_OPENMP)
+    for (size_t i = 0; i < kNumLocks; i++) omp_destroy_lock(&locks[i]);
+#endif
 
-    for (size_t i = matrices.num_constraints;
-         i < matrices.num_constraints + matrices.num_instance_variables; ++i) {
-      a[i] = full_assignments[i - matrices.num_constraints];
+    OPENMP_PARALLEL_FOR(size_t i = 0; i < domain->size(); ++i) {
+      c[i] = a[i] * b[i];
     }
 
     Evals a_evals(std::move(a));
@@ -79,11 +70,8 @@ class QuadraticArithmeticProgram {
     DensePoly c_poly = domain->IFFT(std::move(c_evals));
 
     F root_of_unity;
-    {
-      std::unique_ptr<Domain> extended_domain =
-          Domain::Create(2 * domain->size());
-      root_of_unity = extended_domain->GetElement(1);
-    }
+    CHECK(F::GetRootOfUnity(2 * domain->size(), &root_of_unity));
+
     Domain::DistributePowers(a_poly, root_of_unity);
     Domain::DistributePowers(b_poly, root_of_unity);
     Domain::DistributePowers(c_poly, root_of_unity);
@@ -101,26 +89,6 @@ class QuadraticArithmeticProgram {
 
     return std::move(a_evals).TakeEvaluations();
   }
-
-  template <typename Domain>
-  static std::vector<F> ComputeHQuery(const Domain* domain, const F& t_x,
-                                      const F& x, const F& delta_inverse) {
-    using Evals = typename Domain::Evals;
-    using DensePoly = typename Domain::DensePoly;
-
-    // The usual H query has domain - 1 powers. Z has domain powers. So HZ has
-    // 2 * domain - 1 powers.
-    std::unique_ptr<Domain> extended_domain =
-        Domain::Create(domain->size() * 2 + 1);
-    Evals evals(
-        F::GetSuccessivePowers(extended_domain->size(), t_x, delta_inverse));
-    DensePoly poly = extended_domain->IFFT(std::move(evals));
-    std::vector<F> ret(domain->size());
-    OPENMP_PARALLEL_FOR(size_t i = 0; i < domain->size(); ++i) {
-      ret[i] = poly[2 * i + 1];
-    }
-    return ret;
-  }
 };
 
 }  // namespace tachyon::circom

vendors/circom/circomlib/zkey/BUILD.bazel

@@ -45,7 +45,6 @@ tachyon_cc_library(
         "@kroma_network_tachyon//tachyon/base/buffer:read_only_buffer",
         "@kroma_network_tachyon//tachyon/base/files:file_util",
         "@kroma_network_tachyon//tachyon/base/strings:string_util",
-        "@kroma_network_tachyon//tachyon/zk/r1cs/constraint_system:constraint_matrices",
     ],
 )
 

vendors/circom/circomlib/zkey/zkey.h

@@ -21,7 +21,6 @@
 #include "tachyon/base/logging.h"
 #include "tachyon/base/openmp_util.h"
 #include "tachyon/base/strings/string_util.h"
-#include "tachyon/zk/r1cs/constraint_system/constraint_matrices.h"
 
 namespace tachyon::circom {
 namespace v1 {
@@ -45,7 +44,10 @@ struct ZKey {
   virtual bool Read(const base::ReadOnlyBuffer& buffer) = 0;
 
   virtual ProvingKey<Curve> GetProvingKey() const = 0;
-  virtual zk::r1cs::ConstraintMatrices<F> GetConstraintMatrices() const = 0;
+  virtual absl::Span<const Coefficient<F>> GetCoefficients() const = 0;
+  virtual size_t GetDomainSize() const = 0;
+  virtual size_t GetNumInstanceVariables() const = 0;
+  virtual size_t GetNumWitnessVariables() const = 0;
 
   std::vector<uint8_t> data;
 };
@@ -211,6 +213,12 @@ struct CoefficientsSection {
     Coefficient<F>* ptr;
     if (!buffer.ReadPtr(&ptr, num_coefficients)) return false;
     coefficients = {ptr, num_coefficients};
+
+    OPENMP_PARALLEL_FOR(size_t i = 0; i < coefficients.size(); ++i) {
+      coefficients[i].value =
+          F::FromMontgomery(coefficients[i].value.ToBigInt());
+    }
+
     return true;
   }
 
@@ -294,50 +302,18 @@ struct ZKey : public circom::ZKey<Curve> {
     };
   }
 
-  zk::r1cs::ConstraintMatrices<F> GetConstraintMatrices() const override {
-    std::vector<std::vector<zk::r1cs::Cell<F>>> a(header_groth.domain_size);
-    std::vector<std::vector<zk::r1cs::Cell<F>>> b(header_groth.domain_size);
-
-    uint32_t max_constraint = 0;
-    for (const Coefficient<F>& c : coefficients.coefficients) {
-      max_constraint = std::max(c.constraint, max_constraint);
-      if (c.matrix == 0) {
-        a[c.constraint].push_back({std::move(c.value), c.signal});
-      } else {
-        b[c.constraint].push_back({std::move(c.value), c.signal});
-      }
-    }
-
-    // Need to divide by R, since snarkjs outputs the zkey with coefficients
-    // multiplied by R².
-    OPENMP_PARALLEL_FOR(size_t i = 0; i < max_constraint; ++i) {
-      if (i < a.size()) {
-        for (size_t j = 0; j < a[i].size(); ++j) {
-          a[i][j].coefficient =
-              F::FromMontgomery(a[i][j].coefficient.ToBigInt());
-        }
-      }
-      if (i < b.size()) {
-        for (size_t j = 0; j < b[i].size(); ++j) {
-          b[i][j].coefficient =
-              F::FromMontgomery(b[i][j].coefficient.ToBigInt());
-        }
-      }
-    }
+  absl::Span<const Coefficient<F>> GetCoefficients() const override {
+    return coefficients.coefficients;
+  }
 
-    return {
-        header_groth.num_public_inputs + 1,
-        header_groth.num_vars - header_groth.num_public_inputs - 1,
-        max_constraint - header_groth.num_public_inputs,
+  size_t GetDomainSize() const override { return header_groth.domain_size; }
 
-        0,
-        0,
-        0,
+  size_t GetNumInstanceVariables() const override {
+    return header_groth.num_public_inputs + 1;
+  }
 
-        zk::r1cs::Matrix<F>(std::move(a)),
-        zk::r1cs::Matrix<F>(std::move(b)),
-        {},
-    };
+  size_t GetNumWitnessVariables() const override {
+    return header_groth.num_vars - header_groth.num_public_inputs - 1;
   }
 
   std::string ToString() const {

vendors/circom/circomlib/zkey/zkey_unittest.cc

@@ -157,12 +157,12 @@ TEST_F(ZKeyTest, Parse) {
 
   // clang-format off
   CoefficientData coefficient_datas[] = {
-    {0, 0, 2, "15537367993719455909907449462855742678020201736855642022731641111541721333766"},
-    {1, 0, 3, "6350874878119819312338956282401532410528162663560392320966563075034087161851"},
-    {0, 1, 5, "15537367993719455909907449462855742678020201736855642022731641111541721333766"},
-    {1, 1, 4, "6350874878119819312338956282401532410528162663560392320966563075034087161851"},
-    {0, 2, 0, "6350874878119819312338956282401532410528162663560392320966563075034087161851"},
-    {0, 3, 1, "6350874878119819312338956282401532410528162663560392320966563075034087161851"},
+    {0, 0, 2, "21888242871839275222246405745257275088548364400416034343698204186575808495616"},
+    {1, 0, 3, "1"},
+    {0, 1, 5, "21888242871839275222246405745257275088548364400416034343698204186575808495616"},
+    {1, 1, 4, "1"},
+    {0, 2, 0, "1"},
+    {0, 3, 1, "1"},
   };
   // clang-format on