WIP: feat: integrate simp_mem with sym_n

.gitignore

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

Arm/Insts/Common.lean

@@ -219,20 +219,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.lean

@@ -15,11 +15,15 @@ namespace Benchmarks
 
 def SHA512Bench (nSteps : Nat) : Prop :=
   ∀ (s0 sf : ArmState)
-    (_h_s0_num_blocks : r (.GPR 2#5) s0 = 10)
+    (_h_s0_num_blocks : r (.GPR 2#5) s0 = 10#64)
     (_h_s0_pc : read_pc s0 = 0x1264c0#64)
     (_h_s0_err : read_err s0 = StateError.None)
     (_h_s0_sp_aligned : CheckSPAlignment s0)
     (_h_s0_program : s0.program = SHA512.program)
     (_h_run : sf = run nSteps s0),
     r StateField.ERR sf = StateError.None
-    ∧ r (.GPR 2#5) sf = BitVec.ofNat _ (9 - (nSteps / 485))
+    ∧ r (.GPR 2#5) sf = BitVec.ofNat 64 (10 - ((nSteps + 467) / 485))
+    -- / -------------------------------^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+    -- | This computes the expected value of x2, taking into account that
+    -- | the loop body is 485 instructions long, and that x2 is first
+    -- | decremented after 18 instructions (485 - 18 = 467).

Benchmarks/SHA512_150.lean

@@ -12,5 +12,6 @@ open Benchmarks
 benchmark sha512_150_instructions : SHA512Bench 150 := fun s0 _ h => by
   intros
   sym_n 150
-  · exact (sorry : Aligned ..)
+  simp (config := {failIfUnchanged := false}) only [h, bitvec_rules]
+  all_goals exact (sorry : Aligned ..)
   done

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 (config := {failIfUnchanged := false}) only [h, bitvec_rules]
+  all_goals 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 (config := {failIfUnchanged := false}) only [h, bitvec_rules]
+  all_goals exact (sorry : Aligned ..)
+  done

Benchmarks/SHA512_400.lean

@@ -12,5 +12,6 @@ open Benchmarks
 benchmark sha512_400_instructions : SHA512Bench 400 := fun s0 _ h => by
   intros
   sym_n 400
-  · exact (sorry : Aligned ..)
+  simp (config := {failIfUnchanged := false}) only [h, bitvec_rules]
+  all_goals 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 (config := {failIfUnchanged := false}) only [h, bitvec_rules]
+  all_goals exact (sorry : Aligned ..)
+  done

Benchmarks/SHA512_50.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_50 : SHA512Bench 50 := fun s0 _ h => by
+  intros
+  sym_n 50
+  simp (config := {failIfUnchanged := false}) only [h, bitvec_rules]
+  all_goals exact (sorry : Aligned ..)
+  done

Benchmarks/SHA512_50_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_50_noKernel_noLint : SHA512Bench 50 := fun s0 _ h => by
+  intros
+  sym_n 50
+  simp (config := {failIfUnchanged := false}) only [h, bitvec_rules]
+  all_goals 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 (config := {failIfUnchanged := false}) only [h, bitvec_rules]
+  all_goals 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 (config := {failIfUnchanged := false}) only [h, bitvec_rules]
+  all_goals exact (sorry : Aligned ..)
+  done