feat(circom): generate trace for circom prover

kroma-network · Aug 7, 2024 · 6b18027 · 6b18027
1 parent 5e05f62
commit 6b18027
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 17 deletions.
diff --git a/vendors/circom/BUILD.bazel b/vendors/circom/BUILD.bazel
@@ -11,6 +11,7 @@ tachyon_cc_binary(
  "//circomlib/wtns",
  "//circomlib/zkey",
  "@com_google_absl//absl/strings",
+ "@kroma_network_tachyon//tachyon/base:profiler",
  "@kroma_network_tachyon//tachyon/base/console",
  "@kroma_network_tachyon//tachyon/base/files:file_path_flag",
  "@kroma_network_tachyon//tachyon/base/flag:flag_parser",

diff --git a/vendors/circom/circomlib/circuit/BUILD.bazel b/vendors/circom/circomlib/circuit/BUILD.bazel
@@ -18,6 +18,7 @@ tachyon_cc_library(
  deps = [
  "//circomlib/zkey:coefficient",
  "@kroma_network_tachyon//tachyon/base:logging",
+ "@kroma_network_tachyon//tachyon/base:profiler",
  "@kroma_network_tachyon//tachyon/zk/r1cs/constraint_system:quadratic_arithmetic_program",
  ],
 )

diff --git a/vendors/circom/circomlib/circuit/quadratic_arithmetic_program.h b/vendors/circom/circomlib/circuit/quadratic_arithmetic_program.h
@@ -11,6 +11,7 @@
 
 #include "circomlib/zkey/coefficient.h"
 #include "tachyon/base/logging.h"
+#include "tachyon/base/profiler.h"
 #include "tachyon/zk/r1cs/constraint_system/quadratic_arithmetic_program.h"
 
 namespace tachyon::circom {
@@ -24,13 +25,17 @@ class QuadraticArithmeticProgram {
  static std::vector<F> WitnessMapFromMatrices(
  const Domain* domain, absl::Span<const Coefficient<F>> coefficients,
  absl::Span<const F> full_assignments) {
+ TRACE_EVENT("ProofGeneration", "QAP::WitnessMapFromMatrices");
  using Evals = typename Domain::Evals;
  using DensePoly = typename Domain::DensePoly;
 
+ TRACE_EVENT_BEGIN("Subtask", "VectorInit");
  std::vector<F> a(domain->size());
  std::vector<F> b(domain->size());
  std::vector<F> c(domain->size());
+ TRACE_EVENT_END("Subtask");
 
+ TRACE_EVENT_BEGIN("Subtask", "AB Calculation");
  // See
  // https://github.com/iden3/rapidsnark/blob/b17e6fe/src/groth16.cpp#L116-L156.
 #if defined(TACHYON_HAS_OPENMP)
@@ -57,36 +62,53 @@ class QuadraticArithmeticProgram {
 #if defined(TACHYON_HAS_OPENMP)
  for (size_t i = 0; i < kNumLocks; i++) omp_destroy_lock(&locks[i]);
 #endif
+ TRACE_EVENT_END("Subtask");
 
- OMP_PARALLEL_FOR(size_t i = 0; i < domain->size(); ++i) {
- c[i] = a[i] * b[i];
+ {
+ TRACE_EVENT("Subtask", "C Calculation");
+ OMP_PARALLEL_FOR(size_t i = 0; i < domain->size(); ++i) {
+ c[i] = a[i] * b[i];
+ }
  }
 
  Evals a_evals(std::move(a));
- DensePoly a_poly = domain->IFFT(std::move(a_evals));
  Evals b_evals(std::move(b));
- DensePoly b_poly = domain->IFFT(std::move(b_evals));
  Evals c_evals(std::move(c));
- DensePoly c_poly = domain->IFFT(std::move(c_evals));
+
+ DensePoly a_poly, b_poly, c_poly;
+ {
+ TRACE_EVENT("Subtask", "Do IFFT");
+ a_poly = domain->IFFT(std::move(a_evals));
+ b_poly = domain->IFFT(std::move(b_evals));
+ c_poly = domain->IFFT(std::move(c_evals));
+ }
 
  F root_of_unity;
  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);
-
- a_evals = domain->FFT(std::move(a_poly));
- b_evals = domain->FFT(std::move(b_poly));
- c_evals = domain->FFT(std::move(c_poly));
+ {
+ TRACE_EVENT("Subtask", "DistributePowers");
+ Domain::DistributePowers(a_poly, root_of_unity);
+ Domain::DistributePowers(b_poly, root_of_unity);
+ Domain::DistributePowers(c_poly, root_of_unity);
+ }
 
- // |h_evals[i]| = |a[i]| * |b[i]| - |c[i]|
- OMP_PARALLEL_FOR(size_t i = 0; i < domain->size(); ++i) {
- F& h_evals_i = a_evals.at(i);
- h_evals_i *= b_evals[i];
- h_evals_i -= c_evals[i];
+ {
+  TRACE_EVENT("Subtask", "Do FFT");
+ a_evals = domain->FFT(std::move(a_poly));
+ b_evals = domain->FFT(std::move(b_poly));
+ c_evals = domain->FFT(std::move(c_poly));
  }
 
+ {
+ TRACE_EVENT("Subtask", "HEval Calculation");
+ // |h_evals[i]| = |a[i]| * |b[i]| - |c[i]|
+ OMP_PARALLEL_FOR(size_t i = 0; i < domain->size(); ++i) {
+ F& h_evals_i = a_evals.at(i);
+ h_evals_i *= b_evals[i];
+ h_evals_i -= c_evals[i];
+ }
+ }
  return std::move(a_evals).TakeEvaluations();
  }
 };

diff --git a/vendors/circom/prover_main.cc b/vendors/circom/prover_main.cc
@@ -9,6 +9,7 @@
 #include "tachyon/base/console/iostream.h"
 #include "tachyon/base/files/file_path_flag.h"
 #include "tachyon/base/flag/flag_parser.h"
+#include "tachyon/base/profiler.h"
 #include "tachyon/base/time/time.h"
 #include "tachyon/math/elliptic_curves/bls12/bls12_381/bls12_381.h"
 #include "tachyon/math/elliptic_curves/bn/bn254/bn254.h"
@@ -269,6 +270,11 @@ int RealMain(int argc, char** argv) {
  tachyon_cerr << "num_runs should be positive" << std::endl;
  return 1;
  }
+
+ base::Profiler profiler;
+ profiler.Init();
+ profiler.Start();
+
  switch (curve) {
  case Curve::kBN254:
  circom::CreateProof<math::bn254::BN254Curve>(