feat(circom): generate trace for circom prover

kroma-network · Aug 6, 2024 · 80e164e · 80e164e
1 parent b3c3f38
commit 80e164e
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 5 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:profiling",
         "@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/prover_main.cc b/vendors/circom/prover_main.cc
@@ -1,3 +1,5 @@
+#include <fcntl.h>
+
 #include "absl/strings/substitute.h"
 
 #include "circomlib/circuit/quadratic_arithmetic_program.h"
@@ -9,6 +11,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/profiling.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"
@@ -20,6 +23,8 @@
 #include "tachyon/math/polynomials/univariate/icicle/icicle_ntt_holder.h"
 #endif
 
+PERFETTO_TRACK_EVENT_STATIC_STORAGE();
+
 namespace tachyon {
 
 enum class Curve {
@@ -168,11 +173,10 @@ void CreateProof(const base::FilePath& zkey_path,
     std::cout << "Start verifying" << std::endl;
     zk::r1cs::groth16::PreparedVerifyingKey<Curve> prepared_verifying_key =
         std::move(proving_key).TakeVerifyingKey().ToPreparedVerifyingKey();
-    CHECK(zk::r1cs::groth16::VerifyProof(prepared_verifying_key, proof,
-                                         public_inputs))
-        << "If you see this error with `--config cuda`, it means your GPU "
-           "doesn't have enough RAM for Icicle. Try running it with a GPU with "
-           "more RAM.";
+    bool result = zk::r1cs::groth16::VerifyProof(prepared_verifying_key, proof,
+                                                 public_inputs);
+    CHECK(result);
+
     end = base::TimeTicks::Now();
     std::cout << "Time taken for verifying: " << end - start << std::endl;
   }
@@ -269,6 +273,26 @@ int RealMain(int argc, char** argv) {
     tachyon_cerr << "num_runs should be positive" << std::endl;
     return 1;
   }
+
+  perfetto::TracingInitArgs args;
+  args.backends |= perfetto::kInProcessBackend;
+  perfetto::Tracing::Initialize(args);
+  perfetto::TrackEvent::Register();
+
+  perfetto::protos::gen::TrackEventConfig track_event_cfg;
+  int fd =
+      open("/tmp/prover_main.perfetto-trace", O_RDWR | O_CREAT | O_TRUNC, 0600);
+  std::unique_ptr<perfetto::TracingSession> tracing_session(
+      perfetto::Tracing::NewTrace());
+
+  perfetto::TraceConfig cfg;
+  cfg.add_buffers()->set_size_kb(1500000);
+  auto* ds_cfg = cfg.add_data_sources()->mutable_config();
+  ds_cfg->set_name("track_event");
+  ds_cfg->set_track_event_config_raw(track_event_cfg.SerializeAsString());
+  tracing_session->Setup(cfg, fd);
+  tracing_session->StartBlocking();
+
   switch (curve) {
     case Curve::kBN254:
       circom::CreateProof<math::bn254::BN254Curve>(
@@ -282,6 +306,8 @@ int RealMain(int argc, char** argv) {
       break;
   }
 
+  tracing_session->StopBlocking();
+  close(fd);
   return 0;
 }