chore: Clean up Benchmark API

.github/workflows/bench.yml

@@ -69,11 +69,11 @@ jobs:
       - uses: actions/checkout@v5
         # Get base branch of the PR
       - uses: xt0rted/pull-request-comment-branch@v2
-        id: base-branch
+        id: comment-branch
       - name: Checkout base branch
         uses: actions/checkout@v5
         with:
-          ref: ${{ steps.base-branch.outputs.base_sha }}
+          ref: ${{ steps.comment-branch.outputs.base_sha }}
           path : ${{ github.workspace }}/base
       - name: Run `lake build` on base branch
         uses: leanprover/lean-action@v1
@@ -92,7 +92,7 @@ jobs:
         uses: actions/checkout@v5
         with:
           path: ${{ github.workspace }}/pr
-          ref: ${{ steps.base-branch.outputs.head_sha }}
+          ref: ${{ steps.comment-branch.outputs.head_sha }}
       - name: Run `lake build` on PR branch
         uses: leanprover/lean-action@v1
         with:
@@ -114,7 +114,7 @@ jobs:
         run: |
           SHORT_SHA_PR=$(git rev-parse --short HEAD)
           REPO_URL=${{ github.server_url }}/${{ github.repository }}
-          echo "COMMIT_LINK=[\`$SHORT_SHA_PR\`]($REPO_URL/commit/${{ github.sha }})" | tee -a $GITHUB_ENV
+          echo "COMMIT_LINK=[\`$SHORT_SHA_PR\`]($REPO_URL/commit/${{ steps.comment-branch.outputs.head_sha }})" | tee -a $GITHUB_ENV
           echo "WORKFLOW_LINK=[Workflow logs]($REPO_URL/actions/runs/${{ github.run_id }})" | tee -a $GITHUB_ENV
         working-directory: ${{ github.workspace }}/pr
       - name: Generate token to write PR comment

Benchmarks/Aiur.lean

@@ -84,38 +84,37 @@ def toplevelBench := bgroup "nat_fib" [
   bench "simplify toplevel" Aiur.Toplevel.checkAndSimplify toplevel
 ]
 
-def compileBench : IO Unit := do
+def compileBench : IO $ Array BenchReport := do
   match toplevel.checkAndSimplify with
-  | .error e => IO.eprintln e
+  | .error e => throw (IO.userError s!"{repr e}")
   | .ok decls =>
     bgroup "nat_fib" [
       bench "compile decls" Aiur.TypedDecls.compile decls
     ]
 
-def buildAiurSystemBench : IO Unit := do
+def buildAiurSystemBench : IO $ Array BenchReport := do
   match toplevel.checkAndSimplify with
-  | .error e => IO.eprintln e
+  | .error e => throw (IO.userError s!"{repr e}")
   | .ok decls =>
     let bytecode := decls.compile
     bgroup "nat_fib" [
       bench "build AiurSystem" (Aiur.AiurSystem.build bytecode) commitmentParameters
     ]
 
-def proveBench : IO Unit := do
+def proveBench : IO $ Array BenchReport := do
   match toplevel.checkAndSimplify with
-  | .error e => IO.eprintln e
+  | .error e => throw (IO.userError s!"{repr e}")
   | .ok decls =>
     let bytecode := decls.compile
     let system := Aiur.AiurSystem.build bytecode commitmentParameters
     let funIdx := toplevel.getFuncIdx `main |>.get!
     bgroup "nat_fib" [
-      bench "prove fib 1" (Aiur.AiurSystem.prove system friParameters funIdx) #[1],
-      bench "prove fib 200" (Aiur.AiurSystem.prove system friParameters funIdx) #[200]
-    ] { oneShot := true, serde := .json, report := true }
+      bench "prove fib 10" (Aiur.AiurSystem.prove system friParameters funIdx #[10]) default,
+    ]
 
-def verifyBench : IO Unit := do
+def verifyBench : IO $ Array BenchReport := do
   match toplevel.checkAndSimplify with
-  | .error e => IO.eprintln e
+  | .error e => throw (IO.userError s!"{repr e}")
   | .ok decls =>
     let bytecode := decls.compile
     let system := Aiur.AiurSystem.build bytecode commitmentParameters

Benchmarks/Blake3.lean

@@ -2,6 +2,7 @@ import Ix.IxVM
 import Ix.Aiur.Simple
 import Ix.Aiur.Compile
 import Ix.Aiur.Protocol
+import Ix.Benchmark.Bench
 
 abbrev dataSizes := #[64, 128, 256, 512, 1024, 2048]
 abbrev numHashesPerProof := #[1, 2, 4, 8, 16, 32]
@@ -16,9 +17,16 @@ def friParameters : Aiur.FriParameters := {
   proofOfWorkBits := 20
 }
 
-def runBenches (aiurSystem : Aiur.AiurSystem) (funIdx : Nat) :
-    IO $ Array (Nat × Float) := do
-  let mut results := Array.emptyWithCapacity $ dataSizes.size * numHashesPerProof.size
+def blake3Bench : IO $ Array BenchReport := do
+  let .ok ixVM := IxVM.ixVM
+    | throw (IO.userError "IxVM merging failed")
+  let some funIdx := ixVM.getFuncIdx `blake3_bench
+    | throw (IO.userError "Aiur function not found")
+  let .ok decls := ixVM.checkAndSimplify
+    | throw (IO.userError "Simplification failed")
+  let aiurSystem := Aiur.AiurSystem.build decls.compile commitmentParameters
+
+  let mut benches := Array.emptyWithCapacity $ dataSizes.size * numHashesPerProof.size
   for dataSize in dataSizes do
     for numHashes in numHashesPerProof do
       let ioBuffer := Array.range numHashes |>.foldl
@@ -31,35 +39,26 @@ def runBenches (aiurSystem : Aiur.AiurSystem) (funIdx : Nat) :
           { ioBuffer with
               data := ioBuffer.data ++ data
               map := ioBuffer.map.insert #[.ofNat idx] ioKeyInfo }
-      IO.print s!"Proving for dataSize={dataSize}, numHashes={numHashes} "
-      let startMs ← IO.monoMsNow
-      let (claim, proof, _) := aiurSystem.prove friParameters funIdx
-        #[Aiur.G.ofNat numHashes] ioBuffer
-      let finishedMs ← IO.monoMsNow
-      let elapsedMs := finishedMs - startMs
-      let elapsed := elapsedMs.toFloat / 1000.0
-      IO.print s!"{elapsed}s "
-      let verified := aiurSystem.verify friParameters claim proof |>.isOk
-      if verified then println! "✓" else println! "✕"
-      results := results.push (dataSize * numHashes, elapsed)
-  pure results
+      benches := benches.push <| bench s!"dataSize={dataSize} numHashes={numHashes}" (aiurSystem.prove friParameters funIdx #[Aiur.G.ofNat numHashes]) ioBuffer
+  bgroup "prove blake3" benches.toList { oneShot := true }
 
-def main (_args : List String) : IO Unit := do
-  let .ok ixVM := IxVM.ixVM
-    | println! "IxVM merging failed"; return
-  let some funIdx := ixVM.getFuncIdx `blake3_bench
-    | println! "Aiur function not found"; return
-  let .ok decls := ixVM.checkAndSimplify
-    | println! "Simplification failed"; return
-  let aiurSystem := Aiur.AiurSystem.build decls.compile commitmentParameters
-  let results ← runBenches aiurSystem funIdx
+def parseFunction (words : List String) (param : String): Option String :=
+  words.find? (·.startsWith param) |> .map (·.stripPrefix param)
+
+def main : IO Unit := do
+  let result ← blake3Bench
 
   let mut sumWeights := 0.0
   let mut weightedSum := 0.0
-  for (size, time) in results do
-    let sizeFloat := size.toFloat
-    let throughput := sizeFloat / time
+  for report in result do
+    let words := report.function.splitOn
+    let dataSize := (parseFunction words "dataSize=").get!.toNat!
+    let numHashes := (parseFunction words "numHashes=").get!.toNat!
+    let sizeFloat := (dataSize * numHashes).toFloat
+    let throughput := sizeFloat / (report.newBench.getTime.toSeconds )
     weightedSum := weightedSum + sizeFloat * throughput
     sumWeights := sumWeights + sizeFloat
   let avgThroughput := weightedSum / sumWeights
-  println! s!"Average throughput: {avgThroughput.toUSize} bytes/s"
+  println! "Average throughput: {avgThroughput.toUSize} bytes/s"
+
+