perf: add a cache to bv_decide's reflection procedure (#7644)

This PR adds a cache to the reflection procedure of bv_decide.

This was motivated by the following profile on QF_BV SMTLIB problem
`sage/app12/bench_3564.smt2`: https://share.firefox.dev/4iTG8KX. After
this change we roughly get a 10x speedup and `simp` is the bottleneck
again: https://share.firefox.dev/4iuezYT
This commit is contained in:
Henrik Böving 2025-03-23 14:56:00 +01:00 committed by GitHub
parent f241cc832b
commit d24dfa1031
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 55 additions and 10 deletions

View file

@ -300,6 +300,18 @@ structure LemmaState where
The list of top level lemmas that got created on the fly during reflection.
-/
lemmas : Array SatAtBVLogical := #[]
/--
Cache for reification of `BVExpr`.
-/
bvExprCache : Std.HashMap Expr (Option ReifiedBVExpr) := {}
/--
Cache for reification of `BVPred`.
-/
bvPredCache : Std.HashMap Expr (Option ReifiedBVPred) := {}
/--
Cache for reification of `BVLogicalExpr`.
-/
bvLogicalCache : Std.HashMap Expr (Option ReifiedBVLogical) := {}
/--
The lemma reflection monad. It extends the usual reflection monad `M` by adding the ability to
@ -319,6 +331,36 @@ Add another top level lemma.
def addLemma (lemma : SatAtBVLogical) : LemmaM Unit := do
modify fun s => { s with lemmas := s.lemmas.push lemma }
@[specialize]
def withBVExprCache (e : Expr) (f : Expr → LemmaM (Option ReifiedBVExpr)) :
LemmaM (Option ReifiedBVExpr) := do
match (← get).bvExprCache[e]? with
| some hit => return hit
| none =>
let res ← f e
modify fun s => { s with bvExprCache := s.bvExprCache.insert e res }
return res
@[specialize]
def withBVPredCache (e : Expr) (f : Expr → LemmaM (Option ReifiedBVPred)) :
LemmaM (Option ReifiedBVPred) := do
match (← get).bvPredCache[e]? with
| some hit => return hit
| none =>
let res ← f e
modify fun s => { s with bvPredCache := s.bvPredCache.insert e res }
return res
@[specialize]
def withBVLogicalCache (e : Expr) (f : Expr → LemmaM (Option ReifiedBVLogical)) :
LemmaM (Option ReifiedBVLogical) := do
match (← get).bvLogicalCache[e]? with
| some hit => return hit
| none =>
let res ← f e
modify fun s => { s with bvLogicalCache := s.bvLogicalCache.insert e res }
return res
end LemmaM
end Frontend

View file

@ -184,10 +184,11 @@ where
to return `some`.
-/
goOrAtom (x : Expr) : LemmaM (Option ReifiedBVExpr) := do
let res ← go x
match res with
| some exp => return some exp
| none => ReifiedBVExpr.bitVecAtom x false
LemmaM.withBVExprCache x fun x => do
let res ← go x
match res with
| some exp => return some exp
| none => ReifiedBVExpr.bitVecAtom x false
shiftConstLikeReflection (distance : Nat) (innerExpr : Expr) (shiftOp : Nat → BVUnOp)
(shiftOpName : Name) (congrThm : Name) :
@ -283,9 +284,10 @@ Reify an `Expr` that is a predicate about `BitVec`.
Unless this function is called on something that is not a `Bool` it is always going to return `some`.
-/
partial def ReifiedBVPred.of (t : Expr) : LemmaM (Option ReifiedBVPred) := do
match ← go t with
| some pred => return some pred
| none => ReifiedBVPred.boolAtom t
LemmaM.withBVPredCache t fun t => do
match ← go t with
| some pred => return some pred
| none => ReifiedBVPred.boolAtom t
where
/--
Reify `t`, returns `none` if the reification procedure failed.
@ -344,9 +346,10 @@ where
Unless this function is called on something that is not a `Bool` it is always going to return `some`.
-/
goOrAtom (t : Expr) : LemmaM (Option ReifiedBVLogical) := do
match ← go t with
| some boolExpr => return some boolExpr
| none => ReifiedBVLogical.boolAtom t
LemmaM.withBVLogicalCache t fun t => do
match ← go t with
| some boolExpr => return some boolExpr
| none => ReifiedBVLogical.boolAtom t
gateReflection (lhsExpr rhsExpr : Expr) (gate : Gate) :
LemmaM (Option ReifiedBVLogical) := do