BLOCKED: feat: withInstantiatMainGoal combinator to eagerly instantiate goal metavariables

.gitignore

@@ -4,3 +4,4 @@
 /lake-packages/*
 /.lake/*
 *.log
+/data/*

Arm/Insts/Common.lean

@@ -221,20 +221,32 @@ def CheckSPAlignment (s : ArmState) : Prop :=
 instance : Decidable (CheckSPAlignment s) := by unfold CheckSPAlignment; infer_instance
 
 @[state_simp_rules]
-theorem CheckSPAligment_of_w_different (h : StateField.GPR 31#5 ≠ fld) :
+theorem CheckSPAligment_w_different_eq (h : StateField.GPR 31#5 ≠ fld) :
   CheckSPAlignment (w fld v s) = CheckSPAlignment s := by
   simp_all only [CheckSPAlignment, state_simp_rules, minimal_theory, bitvec_rules]
 
+theorem CheckSPAligment_w_of_ne_sp_of (h : StateField.GPR 31#5 ≠ fld) :
+    CheckSPAlignment s → CheckSPAlignment (w fld v s) := by
+  simp only [CheckSPAligment_w_different_eq h, imp_self]
+
 @[state_simp_rules]
 theorem CheckSPAligment_of_w_sp :
   CheckSPAlignment (w (StateField.GPR 31#5) v s) = (Aligned v 4) := by
   simp_all only [CheckSPAlignment, state_simp_rules, minimal_theory, bitvec_rules]
 
+theorem CheckSPAligment_w_sp_of (h : Aligned v 4) :
+    CheckSPAlignment (w (StateField.GPR 31#5) v s) := by
+  simp only [CheckSPAlignment, read_gpr, r_of_w_same, zeroExtend_eq, h]
+
 @[state_simp_rules]
-theorem CheckSPAligment_of_write_mem_bytes :
+theorem CheckSPAligment_write_mem_bytes_eq :
   CheckSPAlignment (write_mem_bytes n addr v s) = CheckSPAlignment s := by
   simp_all only [CheckSPAlignment, state_simp_rules, minimal_theory, bitvec_rules]
 
+theorem CheckSPAligment_write_mem_bytes_of :
+  CheckSPAlignment s → CheckSPAlignment (write_mem_bytes n addr v s) := by
+  simp only [CheckSPAligment_write_mem_bytes_eq, imp_self]
+
 @[state_simp_rules]
 theorem CheckSPAlignment_AddWithCarry_64_4 (st : ArmState) (y : BitVec 64) (carry_in : BitVec 1)
   (x_aligned : CheckSPAlignment st)

Benchmarks.lean

@@ -3,6 +3,11 @@ Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Author(s): Alex Keizer
 -/
+import Benchmarks.SHA512_75
+import Benchmarks.SHA512_75_noKernel_noLint
 import Benchmarks.SHA512_150
+import Benchmarks.SHA512_150_noKernel_noLint
 import Benchmarks.SHA512_225
+import Benchmarks.SHA512_225_noKernel_noLint
 import Benchmarks.SHA512_400
+import Benchmarks.SHA512_400_noKernel_noLint

Benchmarks/Command.lean

@@ -12,58 +12,104 @@ initialize
     defValue := false
     descr := "enables/disables benchmarking in `withBenchmark` combinator"
   }
+  registerOption `benchmark.runs {
+    defValue := (5 : Nat)
+    descr := "controls how many runs the `benchmark` command does. \
+    NOTE: this value is ignored when the `profiler` option is set to true"
+  }
+  /- Shouldn't be set directly, instead, use the `benchmark` command -/
+  registerTraceClass `benchmark
 
 variable {m} [Monad m] [MonadLiftT BaseIO m] in
-def withHeartbeatsAndMs (x : m α) : m (α × Nat × Nat) := do
+/-- Measure the heartbeats and time (in milliseconds) taken by `x` -/
+def withHeartbeatsAndMilliseconds (x : m α) : m (α × Nat × Nat) := do
   let start ← IO.monoMsNow
   let (a, heartbeats) ← withHeartbeats x
   let endTime ← IO.monoMsNow
   return ⟨a, heartbeats, endTime - start⟩
 
-elab "benchmark" id:ident declSig:optDeclSig val:declVal : command => do
-  logInfo m!"Running {id} benchmark\n"
+/-- Adds a trace node with the `benchmark` class, but only if the profiler
+option is *not* set.
+
+We deliberately suppress benchmarking nodes when profiling, since it generally
+only adds noise
+-/
+def withBenchTraceNode (msg : MessageData) (x : CommandElabM α )
+    : CommandElabM α := do
+  if (← getBoolOption `profiler) then
+    x
+  else
+    withTraceNode `benchmark (fun _ => pure msg) x (collapsed := false)
 
+/--
+Run a benchmark for a set number of times, and report the average runtime.
+
+If the `profiler` option is set true, we run the benchmark only once, with:
+- `trace.profiler` to true, and
+- `trace.profiler.output` set based on the `benchmark.profilerDir` and the
+    id of the benchmark
+-/
+elab "benchmark" id:ident declSig:optDeclSig val:declVal : command => do
+  let originalOpts ← getOptions
+  let mut n := originalOpts.getNat `benchmark.runs 5
+  let mut opts := originalOpts
+  opts := opts.setBool `benchmark true
   let stx ← `(command|
-    set_option benchmark true in
     example $declSig:optDeclSig $val:declVal
   )
 
-  let n := 5
-  let mut totalRunTime := 0
-  -- geomean = exp(log((a₁ a₂ ... aₙ)^1/n)) =
-  -- exp(1/n * (log a₁ + log a₂ + log aₙ)).
-  let mut totalRunTimeLog := 0
-  for i in [0:n] do
-    logInfo m!"\n\nRun {i} (out of {n}):\n"
-    let ((), _, runTime) ← withHeartbeatsAndMs <|
-      elabCommand stx
-
-    logInfo m!"Proof took {runTime / 1000}s in total"
-    totalRunTime := totalRunTime + runTime
-    totalRunTimeLog := totalRunTimeLog + Float.log runTime.toFloat
-
-  let avg := totalRunTime.toFloat / n.toFloat / 1000
-  let geomean := (Float.exp (totalRunTimeLog / n.toFloat)) / 1000.0
-  logInfo m!"\
-{id}:
-  average runtime over {n} runs:
-    {avg}s
-  geomean over {n} runs:
-    {geomean}s
-"
+  if (← getBoolOption `profiler) then
+    opts := opts.setBool `trace.profiler true
+    opts := opts.setNat `trace.profiler.threshold 1
+    n := 1 -- only run once, if `profiler` is set to true
+  else
+    opts := opts.setBool `trace.benchmark true
+
+  if n = 0 then
+    return ()
+
+  -- Set options
+  modifyScope fun scope => { scope with opts }
+
+  withBenchTraceNode m!"Running {id} benchmark" <| do
+    let mut totalRunTime := 0
+    -- geomean = exp(log((a₁ a₂ ... aₙ)^1/n)) =
+    -- exp(1/n * (log a₁ + log a₂ + log aₙ)).
+    let mut totalRunTimeLog : Float := 0
+    for i in [0:n] do
+      let runTime ← withBenchTraceNode m!"Run {i+1} (out of {n}):" <| do
+        let ((), _, runTime) ← withHeartbeatsAndMilliseconds <|
+          elabCommand stx
+
+        trace[benchmark] m!"Proof took {runTime / 1000}s in total"
+        pure runTime
+      totalRunTime := totalRunTime + runTime
+      totalRunTimeLog := totalRunTimeLog + Float.log runTime.toFloat
+
+    let avg := totalRunTime.toFloat / n.toFloat / 1000
+    let geomean := (Float.exp (totalRunTimeLog / n.toFloat)) / 1000.0
+    trace[benchmark] m!"\
+  {id}:
+    average runtime over {n} runs:
+      {avg}s
+    geomean over {n} runs:
+      {geomean}s
+  "
+    -- Restore options
+    modifyScope fun scope => { scope with opts := originalOpts }
 
 /-- Set various options to disable linters -/
 macro "disable_linters" "in" cmd:command : command => `(command|
-  set_option linter.constructorNameAsVariable false in
-  set_option linter.deprecated false in
-  set_option linter.missingDocs false in
-  set_option linter.omit false in
-  set_option linter.suspiciousUnexpanderPatterns false in
-  set_option linter.unnecessarySimpa false in
-  set_option linter.unusedRCasesPattern false in
-  set_option linter.unusedSectionVars false in
-  set_option linter.unusedVariables false in
-  $cmd
+set_option linter.constructorNameAsVariable false in
+set_option linter.deprecated false in
+set_option linter.missingDocs false in
+set_option linter.omit false in
+set_option linter.suspiciousUnexpanderPatterns false in
+set_option linter.unnecessarySimpa false in
+set_option linter.unusedRCasesPattern false in
+set_option linter.unusedSectionVars false in
+set_option linter.unusedVariables false in
+$cmd
 )
 
 /-- The default `maxHeartbeats` setting.
@@ -96,7 +142,7 @@ private def withBenchmarkAux (x : m α) (f : Nat → Nat → m Unit)  : m α :=
   if (← getBoolOption `benchmark) = false then
     x
   else
-    let (a, heartbeats, t) ← withHeartbeatsAndMs x
+    let (a, heartbeats, t) ← withHeartbeatsAndMilliseconds x
     f heartbeats t
     return a
 

Benchmarks/SHA512_150_noKernel_noLint.lean

@@ -0,0 +1,19 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Author(s): Alex Keizer
+-/
+import Tactics.Sym
+import Benchmarks.Command
+import Benchmarks.SHA512
+
+open Benchmarks
+
+disable_linters in
+set_option debug.skipKernelTC true in
+benchmark sha512_150_noKernel_noLint : SHA512Bench 150 := fun s0 _ h => by
+  intros
+  sym_n 150
+  simp only [h, bitvec_rules]
+  · exact (sorry : Aligned ..)
+  done

Benchmarks/SHA512_225_noKernel_noLint.lean

@@ -0,0 +1,19 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Author(s): Alex Keizer
+-/
+import Tactics.Sym
+import Benchmarks.Command
+import Benchmarks.SHA512
+
+open Benchmarks
+
+disable_linters in
+set_option debug.skipKernelTC true in
+benchmark sha512_225_noKernel_noLint : SHA512Bench 225 := fun s0 _ h => by
+  intros
+  sym_n 225
+  simp only [h, bitvec_rules]
+  · exact (sorry : Aligned ..)
+  done

Benchmarks/SHA512_400_noKernel_noLint.lean

@@ -0,0 +1,19 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Author(s): Alex Keizer
+-/
+import Tactics.Sym
+import Benchmarks.Command
+import Benchmarks.SHA512
+
+open Benchmarks
+
+disable_linters in
+set_option debug.skipKernelTC true in
+benchmark sha512_400_noKernel_noLint : SHA512Bench 400 := fun s0 _ h => by
+  intros
+  sym_n 400
+  simp only [h, bitvec_rules]
+  · exact (sorry : Aligned ..)
+  done

Benchmarks/SHA512_75.lean

@@ -0,0 +1,17 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Author(s): Alex Keizer
+-/
+import Tactics.Sym
+import Benchmarks.Command
+import Benchmarks.SHA512
+
+open Benchmarks
+
+benchmark sha512_75 : SHA512Bench 75 := fun s0 _ h => by
+  intros
+  sym_n 75
+  simp only [h, bitvec_rules]
+  · exact (sorry : Aligned ..)
+  done

Benchmarks/SHA512_75_noKernel_noLint.lean

@@ -0,0 +1,19 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Author(s): Alex Keizer
+-/
+import Tactics.Sym
+import Benchmarks.Command
+import Benchmarks.SHA512
+
+open Benchmarks
+
+disable_linters in
+set_option debug.skipKernelTC true in
+benchmark sha512_75_noKernel_noLint : SHA512Bench 75 := fun s0 _ h => by
+  intros
+  sym_n 75
+  simp only [h, bitvec_rules]
+  · exact (sorry : Aligned ..)
+  done

Makefile

@@ -5,6 +5,8 @@
 SHELL := /bin/bash
 
 LAKE = lake
+LEAN = $(LAKE) env lean
+GIT = git
 
 NUM_TESTS?=3
 VERBOSE?=--verbose
@@ -37,9 +39,22 @@ awslc_elf:
 cosim:
 	time -p lake exe lnsym $(VERBOSE) --num-tests $(NUM_TESTS)
 
+BENCHMARKS = Benchmarks/SHA512_75.lean \
+	Benchmarks/SHA512_75_noKernel_noLint.lean \
+	Benchmarks/SHA512_150.lean \
+	Benchmarks/SHA512_150_noKernel_noLint.lean \
+	Benchmarks/SHA512_225.lean \
+	Benchmarks/SHA512_225_noKernel_noLint.lean \
+	Benchmarks/SHA512_400.lean \
+	Benchmarks/SHA512_400_noKernel_noLint.lean
+
 .PHONY: benchmarks
 benchmarks:
-	$(LAKE) build Benchmarks
+	./scripts/benchmark.sh $(BENCHMARKS)
+
+.PHONY: profile
+profile:
+	./scripts/profile.sh $(BENCHMARKS)
 
 .PHONY: clean clean_all
 clean:

README.md

@@ -43,6 +43,8 @@ native-code programs of interest.
 
 `benchmarks`: run benchmarks for the symbolic simulator.
 
+`profiler`: run a single round of each benchmark, with the profiler enabled
+
 ### Makefile variables that can be passed in at the command line
 
 `VERBOSE`: Verbose mode; prints disassembly of the instructions being

Tactics/Attr.lean

@@ -9,8 +9,11 @@ open Lean
 initialize
   -- CSE tactic's non-verbose summary logging.
   registerTraceClass `Tactic.cse.summary
+
   -- enable tracing for `sym_n` tactic and related components
   registerTraceClass `Tactic.sym
+  -- enable verbose tracing
+  registerTraceClass `Tactic.sym.debug
 
   -- enable tracing for heartbeat usage of `sym_n`
   registerTraceClass `Tactic.sym.heartbeats

Tactics/Common.lean

@@ -278,5 +278,27 @@ Unlike the standard `withMainContext`, `x` may live in a generic monad `m`. -/
 def withMainContext' (x : m α) : m α := do
   (← getMainGoal).withContext x
 
+variable {m} [Monad m] [MonadLiftT TacticM m] [MonadLiftT MetaM m]
+  [MonadTrace m] [MonadLiftT IO m] [MonadLiftT BaseIO m] [AddMessageContext m]
+  [MonadRef m] [MonadOptions m] [MonadAlwaysExcept ε m] [MonadMCtx m] in
+/--
+`withInstantiateMainGoal x` will replace the main goal with a fresh metavariable
+of the same type, run `x` to solve for the new main goal, and then
+assign the old goal with the instantiating of the new goal.
+
+This seems redundant, but allows us to spread out the cost of metavariable
+instantiation, and hopefully avoid some quadratic behaviour we've observed.
+-/
+def withInstantiateMainGoal (x : m α) : m α := do
+  let goals ← getUnsolvedGoals
+  let a ← x
+  for goal in goals do
+    if ← goal.isAssigned then
+      withTraceNode `Tactic.sym (fun _ => pure m!"instantiating goal")
+        (tag := "withInstantiateMainGoal.instantiateMVar") <| do
+          let e ← instantiateMVars (Expr.mvar goal)
+          goal.assign e
+  return a
+
 /-- An emoji to show that a tactic is processing at an intermediate step. -/
 def processingEmoji : String := "⚙️"