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.

* --op-stats on `gen_predicates` also outputs encoded cycle counts in `...-encoded-cycles.txt`

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

Author: shkoo

.github/workflows/main.yml

@@ -35,7 +35,7 @@ jobs:
           fetch-depth: 0
       - uses: risc0/risc0/.github/actions/rustup@a9d723e29a44563497220a998b5de4e03d9da049
       - name: Install cargo-sort
-        uses: risc0/cargo-install@v1
+        uses: risc0/cargo-install@9f6037ed331dcf7da101461a20656273fa72abf0
         with:
           crate: cargo-sort
           version: "1.0"

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

@@ -55,6 +55,7 @@ static cl::opt<std::string> keccakIR("keccak-ir",
 int main(int argc, char* argv[]) {
   llvm::InitLLVM y(argc, argv);
   registerEdslCLOptions();
+  registerOpStatsCLOptions();
   llvm::cl::ParseCommandLineOptions(argc, argv, "keccak predicates");
 
   Module module;

zirgen/circuit/keccak/predicates.cpp

@@ -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

@@ -136,8 +136,10 @@ struct Instructions {
   uint64_t padShaCountConst;
   uint64_t shaRngConsts;
   llvm::StringMap<uint64_t> tagConsts;
+  EncodeStats* stats = nullptr;
 
-  Instructions(HashType hashType) : hashType(hashType), nextOut(1), microUsed(0) {
+  Instructions(HashType hashType, EncodeStats* stats)
+      : hashType(hashType), nextOut(1), microUsed(0), stats(stats) {
     addMacro(/*outs=*/0, MacroOpcode::WOM_INIT);
     // Make: [0, 1, 0, 0], [0, 0, 1, 0], and [0, 0, 0, 1]
     fp4Rot1 = addConst(0, 1);
@@ -224,6 +226,15 @@ struct Instructions {
     if (microUsed == 3) {
       microUsed = 0;
     }
+
+    if (stats) {
+      if (out) {
+        ScopedLocation loc(out.getLoc());
+        stats->locs[currentLoc()]++;
+      } else {
+        stats->locs[currentLoc()]++;
+      }
+    }
     return outId;
   }
 
@@ -245,6 +256,8 @@ struct Instructions {
     uint64_t writeAddr = nextOut;
     data.back().writeAddr = writeAddr;
     nextOut += outs;
+    if (stats)
+      stats->locs[currentLoc()] += 3;
     return writeAddr;
   }
 
@@ -263,6 +276,8 @@ struct Instructions {
     uint64_t writeAddr = nextOut;
     data.back().writeAddr = writeAddr;
     nextOut += 1; // Write exactly one thing (evaluated point)
+    if (stats)
+      stats->locs[currentLoc()] += 3;
     return writeAddr;
   }
 
@@ -286,6 +301,8 @@ struct Instructions {
       data.back().data.poseidon2Mem.inputs[i] = inputs[i];
     }
     data.back().writeAddr = nextOut;
+    if (stats)
+      stats->locs[currentLoc()] += 3;
   }
 
   void addPoseidon2Full(uint64_t cycle) {
@@ -295,6 +312,8 @@ struct Instructions {
     data.back().opType = OpType::POSEIDON2_FULL;
     data.back().data.poseidon2Full.cycle = cycle;
     data.back().writeAddr = nextOut;
+    if (stats)
+      stats->locs[currentLoc()] += 3;
   }
 
   void addPoseidon2Partial() {
@@ -303,6 +322,8 @@ struct Instructions {
     data.emplace_back();
     data.back().opType = OpType::POSEIDON2_PARTIAL;
     data.back().writeAddr = nextOut;
+    if (stats)
+      stats->locs[currentLoc()] += 3;
   }
 
   uint64_t addPoseidon2Store(uint64_t doMont, uint64_t group) {
@@ -315,30 +336,36 @@ struct Instructions {
     uint64_t writeAddr = nextOut;
     data.back().writeAddr = writeAddr;
     nextOut += 8;
+    if (stats)
+      stats->locs[currentLoc()] += 3;
     return writeAddr;
   }
 
   void doShaInit() {
+    ScopedLocation loc;
     for (size_t i = 0; i < 4; i++) {
       shaUsed++;
       addMacro(/*outs=*/0, MacroOpcode::SHA_INIT, shaInit[3 - i], shaInit[3 - i + 4]);
     }
   }
 
   void doShaLoad(llvm::ArrayRef<uint64_t> values, uint64_t subtype) {
+    ScopedLocation loc;
     for (size_t i = 0; i < 16; i++) {
       shaUsed++;
       addMacro(/*outs=*/0, MacroOpcode::SHA_LOAD, values[i], shaK[i], subtype);
     }
   }
 
   void doShaMix() {
+    ScopedLocation loc;
     for (size_t i = 0; i < 48; i++) {
       shaUsed++;
       addMacro(/*outs=*/0, MacroOpcode::SHA_MIX, 0, shaK[16 + i]);
     }
   }
   uint64_t doShaFini() {
+    ScopedLocation loc;
     uint64_t out = nextOut;
     for (size_t i = 0; i < 4; i++) {
       shaUsed++;
@@ -349,6 +376,7 @@ struct Instructions {
   }
 
   uint64_t doSha(llvm::ArrayRef<uint64_t> values, uint64_t subtype) {
+    ScopedLocation loc;
     doShaInit();
     uint64_t ret = 0;
     for (size_t i = 0; i < values.size() / 16; i++) {
@@ -360,6 +388,7 @@ struct Instructions {
   }
 
   uint64_t doShaFold(uint64_t lhs, uint64_t rhs) {
+    ScopedLocation loc;
     std::vector<uint64_t> ids(16);
     for (size_t i = 0; i < 8; i++) {
       ids[i] = lhs + i;
@@ -369,6 +398,7 @@ struct Instructions {
   }
 
   uint64_t doIntoDigestShaBytes(llvm::ArrayRef<uint64_t> bytes) {
+    ScopedLocation loc;
     // We keep things in low / high form right until the end so that the final adds are
     // all contiguous since all the 'digest' stuff assumes digests are always contiguous.
     std::vector<uint64_t> low;
@@ -388,6 +418,7 @@ struct Instructions {
   }
 
   uint64_t doIntoDigestShaWords(llvm::ArrayRef<uint64_t> words) {
+    ScopedLocation loc;
     // We keep things in low / high form right until the end so that the final adds are
     // all contiguous since all the 'digest' stuff assumes digests are always contiguous.
     std::vector<uint64_t> low;
@@ -404,6 +435,7 @@ struct Instructions {
   }
 
   uint64_t doShaTag(llvm::StringRef tag) {
+    ScopedLocation loc;
     if (tagConsts.count(tag)) {
       return tagConsts.find(tag)->second;
     } else {
@@ -422,6 +454,7 @@ struct Instructions {
                           llvm::ArrayRef<uint64_t> digests,
                           llvm::ArrayRef<DigestKind> digestTypes,
                           llvm::ArrayRef<uint64_t> vals) {
+    ScopedLocation loc;
     std::vector<uint64_t> words;
 
     for (size_t i = 0; i < 8; i++) {
@@ -472,6 +505,7 @@ struct Instructions {
   // representation, with 16 bits in each of the two low components of an
   // extension field element.
   void taggedStructPushVals(std::vector<uint64_t>& words, llvm::ArrayRef<uint64_t> vals) {
+    ScopedLocation loc;
     // Get low 16 bits of each value (done in a loop for better packing)
     std::vector<uint64_t> lowVals;
     for (size_t i = 0; i < vals.size(); i++) {
@@ -489,6 +523,7 @@ struct Instructions {
   }
 
   std::pair<uint64_t, std::vector<uint64_t>> doHashCheckedBytes(uint64_t evalPt, uint64_t count) {
+    ScopedLocation loc;
     if (!count) {
       // Special case for 0 outputs
       return {doPoseidon2({}), {}};
@@ -519,6 +554,7 @@ struct Instructions {
 
   std::tuple<uint64_t, uint64_t, std::vector<uint64_t>> doHashCheckedBytesPublic(uint64_t evalPt,
                                                                                  uint64_t count) {
+    ScopedLocation loc;
     if (!count)
       throw std::runtime_error("Cannont publically hash empty checked bytes");
 
@@ -570,6 +606,7 @@ struct Instructions {
   }
 
   uint64_t doPoseidon2(llvm::ArrayRef<uint64_t> values) {
+    ScopedLocation loc;
     if (values.empty()) {
       auto psuite = poseidon2HashSuite();
       auto hashVal = psuite->hash(nullptr, 0);
@@ -634,6 +671,7 @@ struct Instructions {
   }
 
   uint64_t doIntoDigestPoseidon2(llvm::ArrayRef<uint64_t> words) {
+    ScopedLocation loc;
     // Do pointless adds to make all the words land in sequential spots
     uint64_t ret = nextOut;
     size_t pad = words.size() / 8 - 1;
@@ -646,6 +684,7 @@ struct Instructions {
   }
 
   void addInst(Operation& op) {
+    ScopedLocation loc(op.getLoc());
     TypeSwitch<Operation*>(&op)
         .Case<Zll::ExternOp>([&](Zll::ExternOp op) {
           if (op.getName() == "write") {
@@ -984,6 +1023,7 @@ uint64_t ShaRng::generateFp(Instructions& insts) {
 }
 
 void ShaRng::mix(Instructions& insts, uint64_t digest) {
+  ScopedLocation loc;
   uint64_t xorOut = insts.nextOut;
   for (size_t i = 0; i < 8; i++) {
     // Xors and returns 2 shorts: [a, b, 0, 0] ^ [c, d, 0, 0] -> [a ^ c, b ^ d, 0, 0]
@@ -1022,6 +1062,7 @@ uint64_t Poseidon2Rng::generateFp(Instructions& insts) {
 }
 
 void Poseidon2Rng::mix(Instructions& insts, uint64_t digest) {
+  ScopedLocation loc;
   if (stateUsed != 0) {
     stateUsed = 0;
     mix(insts, 0);
@@ -1069,6 +1110,7 @@ void Poseidon2Rng::mix(Instructions& insts, uint64_t digest) {
 }
 
 void MixedPoseidon2ShaRng::mix(Instructions& insts, uint64_t digest) {
+  ScopedLocation loc;
   // For each element of the Poseidon2 hash, we convert it to a form usable by SHA.
   // This is done in stages so that macro ops are grouped together for efficiency
   // First, we 'and' things by 0xffff
@@ -1104,7 +1146,7 @@ std::vector<uint32_t> encode(HashType hashType,
                              mlir::Block* block,
                              llvm::DenseMap<Value, uint64_t>* toIdReturn,
                              EncodeStats* stats) {
-  Instructions insts(hashType);
+  Instructions insts(hashType, stats);
   for (Operation& op : block->without_terminator()) {
     insts.addInst(op);
   }

zirgen/circuit/recursion/encode.cpp

@@ -31,6 +31,9 @@ struct EncodeStats {
   size_t totCycles = 0;
   size_t shaCycles = 0;
   size_t poseidon2Cycles = 0;
+
+  // Locations and the number of micro cycles used (= number of macro cycles * 3).
+  llvm::DenseMap<mlir::Location, /*cycles=*/size_t> locs;
 };
 
 std::vector<uint32_t> encode(HashType hashType,

zirgen/circuit/recursion/encode.h

@@ -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[]) {