From 75d1e2281cfa14430b22098811968773937e397c Mon Sep 17 00:00:00 2001
From: batzor <32958247+batzor@users.noreply.github.com>
Date: Tue, 6 Aug 2024 17:19:00 +0900
Subject: [PATCH] feat(circom): generate trace for circom prover

---
 vendors/circom/BUILD.bazel                             |  1 +
 vendors/circom/circomlib/circuit/BUILD.bazel           |  1 +
 .../circomlib/circuit/quadratic_arithmetic_program.h   | 10 ++++++++++
 vendors/circom/prover_main.cc                          |  6 ++++++
 4 files changed, 18 insertions(+)

diff --git a/vendors/circom/BUILD.bazel b/vendors/circom/BUILD.bazel
index 42b467ed28..a34c12bf73 100644
--- 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
index 77bdf3f2b7..4ad567c4ad 100644
--- 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
index a5ee0db726..5ab8e5f905 100644
--- 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,10 +62,13 @@ 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");
 
+    TRACE_EVENT_BEGIN("Subtask", "C Calculation");
     OPENMP_PARALLEL_FOR(size_t i = 0; i < domain->size(); ++i) {
       c[i] = a[i] * b[i];
     }
+    TRACE_EVENT_END("Subtask");
 
     Evals a_evals(std::move(a));
     DensePoly a_poly = domain->IFFT(std::move(a_evals));
@@ -80,12 +88,14 @@ class QuadraticArithmeticProgram {
     b_evals = domain->FFT(std::move(b_poly));
     c_evals = domain->FFT(std::move(c_poly));
 
+    TRACE_EVENT_BEGIN("Subtask", "HEval Calculation");
     // |h_evals[i]| = |a[i]| * |b[i]| - |c[i]|
     OPENMP_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_END("Subtask");
 
     return std::move(a_evals).TakeEvaluations();
   }
diff --git a/vendors/circom/prover_main.cc b/vendors/circom/prover_main.cc
index ec31959d44..4434815db9 100644
--- 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>(