Add circuit location statistics for recursion predicates

* Overhaul EDSL location handling; we now generate mlir Locations directly instead of keeping our own SourceLoc structure, and it's relatively easy to add additional context to component construction using ScopedSourceLoc.

* Expand --op-stats flag to work on `gen_predicates`; this outputs `...-opstats.txt` files in the output directory for each predicate processed.

* Add --op-stats-flat-cost flag to count each instruction the same. (This is still not the same as the number of cycles in the ZKR, since the statistics are collected before encoding)

* Added a bunch more source annotations to get better granularity for recursion predicates.

Author: shkoo

risc0/core/BUILD.bazel

@@ -8,7 +8,6 @@ cc_library(
     hdrs = [
         "elf.h",
         "log.h",
-        "source_loc.h",
         "util.h",
     ],
     visibility = ["//visibility:public"],

risc0/core/source_loc.h

@@ -19,6 +19,7 @@
 #include "zirgen/circuit/verify/wrap_rv32im.h"
 #include "zirgen/circuit/verify/wrap_zirgen.h"
 #include "zirgen/compiler/codegen/codegen.h"
+#include "zirgen/compiler/stats/OpStats.h"
 
 using namespace zirgen;
 using namespace zirgen::verify;
@@ -202,6 +203,7 @@ static cl::opt<std::string>
 int main(int argc, char* argv[]) {
   llvm::InitLLVM y(argc, argv);
   registerEdslCLOptions();
+  registerOpStatsCLOptions();
   llvm::cl::ParseCommandLineOptions(argc, argv, "gen_predicates edsl");
 
   Module module;

zirgen/circuit/predicates/gen_predicates.cpp

@@ -20,7 +20,6 @@
 
 using namespace zirgen;
 using namespace zirgen::recursion;
-using namespace risc0;
 using namespace mlir;
 
 int main(int argc, char* argv[]) {

zirgen/circuit/recursion/recursion.cpp

@@ -23,7 +23,6 @@
 
 using namespace zirgen;
 using namespace zirgen::rv32im_v1;
-using namespace risc0;
 using namespace mlir;
 
 int main(int argc, char* argv[]) {

zirgen/circuit/rv32im/v1/edsl/rv32im.cpp

@@ -73,6 +73,8 @@ Val fold_eval(const std::vector<Val>& values, Val x) {
 }
 
 Val dynamic_pow(Val in, Val pow, size_t maxPow) {
+  ScopedLocation loc;
+
   Val out = 1;
   Val mul = in;
   for (size_t i = 0; i < log2Ceil(maxPow); i++) {
@@ -98,6 +100,8 @@ struct VerifyRoundInfo {
       , mix(iop.rngExtVal()) {}
 
   void verifyQuery(ReadIopVal& iop, Val* pos, Val* goal) const {
+    ScopedLocation loc;
+
     // Compute which group we are in
     Val group = *pos & (domain - 1);
     Val qout = (*pos - group) / domain;
@@ -117,6 +121,8 @@ struct VerifyRoundInfo {
 
 // Verify a FRI proof,
 void friVerify(ReadIopVal& iop, size_t deg, InnerVerify inner) {
+  ScopedLocation loc;
+
   size_t domain = deg * kInvRate;
   size_t origDomain = domain;
   std::vector<VerifyRoundInfo> rounds;

zirgen/circuit/verify/fri.cpp

@@ -64,6 +64,8 @@ DigestVal MerkleTreeVerifier::getRoot() const {
 }
 
 std::vector<Val> MerkleTreeVerifier::verify(ReadIopVal& iop, Val idx) const {
+  ScopedLocation loc;
+
   std::vector<Val> out;
   if (useExtension) {
     out = iop.readExtVals(colSize);

zirgen/circuit/verify/merkle.cpp

@@ -17,6 +17,8 @@
 namespace zirgen::verify {
 
 Val poly_eval(const std::vector<Val>& coeffs, Val x) {
+  ScopedLocation loc;
+
   Val tot = 0;
   Val mul = 1;
   for (size_t i = 0; i < coeffs.size(); i++) {

zirgen/circuit/verify/poly.cpp

@@ -38,6 +38,8 @@ template <typename T> T dbg(std::string fmt, T arg) {
 } // namespace
 
 VerifyInfo verify(ReadIopVal& iop, size_t po2, const CircuitInterface& circuit) {
+  ScopedLocation loc;
+
   VerifyInfo verifyInfo;
 
   // At the start of verification, add the version strings to the Fiat-Shamir transcript.
@@ -177,11 +179,22 @@ VerifyInfo verify(ReadIopVal& iop, size_t po2, const CircuitInterface& circuit)
   }
   // Finally, do a FRI verification
   friVerify(iop, size, [&](ReadIopVal& iop, Val idx) {
+    ScopedLocation loc;
+
     auto x = dynamic_pow(kRouFwd[log2Ceil(domain)], idx, domain);
     std::map<unsigned, std::vector<Val>> rows;
-    rows[/*REGISTER_GROUP_ACCUM*/ 0] = accumMerkle.verify(iop, idx);
-    rows[/*REGISTER_GROUP_CODE=*/1] = codeMerkle.verify(iop, idx);
-    rows[/*REGISTER_GROUP_DATA=*/2] = dataMerkle.verify(iop, idx);
+    {
+      ScopedLocation loc;
+      rows[/*REGISTER_GROUP_ACCUM*/ 0] = accumMerkle.verify(iop, idx);
+    }
+    {
+      ScopedLocation loc;
+      rows[/*REGISTER_GROUP_CODE=*/1] = codeMerkle.verify(iop, idx);
+    }
+    {
+      ScopedLocation loc;
+      rows[/*REGISTER_GROUP_DATA=*/2] = dataMerkle.verify(iop, idx);
+    }
     auto checkRow = checkMerkle.verify(iop, idx);
     Val curMix = 1;
     std::vector<Val> tot(comboU.size(), 0);
@@ -219,6 +232,8 @@ VerifyInfo verifyRecursion(ReadIopVal& allowedRoot,
                            std::vector<ReadIopVal> seals,
                            std::vector<ReadIopVal> alloweds,
                            const CircuitInterface& circuit) {
+  ScopedLocation loc;
+
   VerifyInfo verifyInfo;
   verifyInfo.codeRoot = allowedRoot.readDigests(1)[0];
 

zirgen/circuit/verify/verify.cpp

@@ -37,7 +37,6 @@ struct CppStreamEmitterImpl : CppStreamEmitter {
   void header(func::FuncOp func) {
     ofs << "// This code is automatically generated\n\n";
     ofs << "#include \"impl.h\"\n\n";
-    ofs << "using namespace risc0;\n\n";
     ofs << "namespace circuit::" << func.getName() << " {\n\n";
   }