117 lines
5.1 KiB
Lean4
117 lines
5.1 KiB
Lean4
import Spa.Analysis.Reaching
|
||
|
||
/-!
|
||
# Finding loop-invariant assignments (LICM groundwork)
|
||
|
||
This wires the **reaching-definitions** analysis (`Spa/Analysis/Reaching.lean`)
|
||
to the AST to *find* — not yet move — assignments inside a `while` loop whose
|
||
right-hand side depends only on definitions made *outside* the loop. These are
|
||
the candidates a later LICM pass could hoist.
|
||
|
||
The traversal recurses over the plain `Stmt`, threading a `GGraph.Embed` of the
|
||
current subtree's CFG into the program's (`Program.rootEmbed`, then one
|
||
`Embed.trans` per descent). That embedding is what supplies program states:
|
||
|
||
1. at an assignment, its CFG state is `Embed.singletonIndex` — the subtree's CFG
|
||
is a `singleton`, so its sole node is the state, and `nodes_eq` proves it
|
||
holds that very statement;
|
||
2. read the reaching definitions at the assignment's *entry* (`joinForKey s
|
||
result` — the join over predecessors, i.e. before the assignment runs);
|
||
3. union the definition sets of the RHS variables;
|
||
4. check no definition site lies in the loop body's CFG range. Every embedding is
|
||
a constant index shift, so the body occupies the interval
|
||
`[off, off + size)` (`GGraph.Embed.mem_range_iff`) and the test is two
|
||
comparisons.
|
||
|
||
If every reaching definition of every RHS variable lies outside the loop, the
|
||
assignment is reported as loop-invariant. This is the first-order check ("all
|
||
reaching definitions outside the loop"); transitive/iterated invariance and the
|
||
actual hoisting are out of scope here.
|
||
-/
|
||
|
||
namespace Spa
|
||
|
||
namespace LicmTransformation
|
||
|
||
open Forward GGraph
|
||
|
||
/-- The CFG footprint of an enclosing loop: its entry node (for reporting) and
|
||
the index interval its body occupies. -/
|
||
structure Enclosing (prog : Program) where
|
||
/-- The loop's entry node, i.e. `GGraph.loopIn` embedded into the program. -/
|
||
loopState : prog.State
|
||
/-- Start of the body's index range. -/
|
||
bodyOff : ℕ
|
||
/-- Length of the body's index range. -/
|
||
bodySize : ℕ
|
||
|
||
/-- Is this definition site inside the loop body's CFG range? -/
|
||
def Enclosing.covers {prog : Program} (l : Enclosing prog) (d : prog.State) : Bool :=
|
||
decide (l.bodyOff ≤ d.val ∧ d.val < l.bodyOff + l.bodySize)
|
||
|
||
/-- An assignment found inside a loop, paired with the data needed to test its
|
||
invariance against that (immediately enclosing) loop. -/
|
||
structure Candidate (prog : Program) where
|
||
/-- The enclosing loop. -/
|
||
encl : Enclosing prog
|
||
/-- The assignment's CFG state. -/
|
||
assignState : prog.State
|
||
/-- The variables read by the assignment's RHS. -/
|
||
rhsVars : List String
|
||
|
||
/-- Collect every assignment together with its *immediately enclosing* loop.
|
||
`enc` is `none` outside any loop, in which case assignments are skipped — only
|
||
in-loop assignments are candidates. -/
|
||
def collectCandidates (prog : Program) (enc : Option (Enclosing prog)) :
|
||
(s : Stmt) → Embed s.cfg prog.cfg → List (Candidate prog)
|
||
| .basic bs, e =>
|
||
match bs, enc with
|
||
| .assign _ ex, some l =>
|
||
[{ encl := l, assignState := e.singletonIndex,
|
||
rhsVars := ex.vars.sort (· ≤ ·) }]
|
||
| _, _ => []
|
||
| .andThen s₁ s₂, e =>
|
||
collectCandidates prog enc s₁ ((Embed.sequenceLeft s₁.cfg s₂.cfg).trans e) ++
|
||
collectCandidates prog enc s₂ ((Embed.sequenceRight s₁.cfg s₂.cfg).trans e)
|
||
| .ifElse _ s₁ s₂, e =>
|
||
collectCandidates prog enc s₁ ((Embed.overlayLeft s₁.cfg s₂.cfg).trans e) ++
|
||
collectCandidates prog enc s₂ ((Embed.overlayRight s₁.cfg s₂.cfg).trans e)
|
||
| .whileLoop _ body, e =>
|
||
let be := (Embed.loop body.cfg).trans e
|
||
collectCandidates prog
|
||
(some { loopState := e.f body.cfg.loopIn, bodyOff := be.off,
|
||
bodySize := body.cfg.size }) body be
|
||
|
||
/-- Read the definition set assigned to variable `k`, or `⊥` if absent. -/
|
||
def lookupDef (prog : Program) (vs : VariableValues (DefSet prog) prog)
|
||
(k : String) : DefSet prog :=
|
||
if h : FiniteMap.MemKey k vs then (FiniteMap.locate h).1 else ⊥
|
||
|
||
/-- Is the candidate assignment loop-invariant: do all reaching definitions of
|
||
its RHS variables lie outside the loop body? -/
|
||
def isInvariant (prog : Program) (c : Candidate prog) : Bool :=
|
||
let entry := joinForKey c.assignState (result (DefSet prog) prog)
|
||
let combined : DefSet prog :=
|
||
c.rhsVars.foldl (fun acc k => acc ⊔ lookupDef prog entry k) ⊥
|
||
-- `Finset.toList` is noncomputable; the decidable bounded-∀ folds over the
|
||
-- underlying multiset and keeps `lake exe` working.
|
||
decide (∀ d ∈ combined, c.encl.covers d = false)
|
||
|
||
/-- The loop-invariant assignments of `prog`, as `(loop, assignment)` state pairs. -/
|
||
def licmCandidates (prog : Program) : List (prog.State × prog.State) :=
|
||
(collectCandidates prog none prog.rootStmt prog.rootEmbed).filterMap (fun c =>
|
||
if isInvariant prog c then some (c.encl.loopState, c.assignState) else none)
|
||
|
||
/-- A human-readable report of the loop-invariant assignments. -/
|
||
def output (prog : Program) : String :=
|
||
match licmCandidates prog with
|
||
| [] => "no loop-invariant assignments found"
|
||
| cands =>
|
||
"loop-invariant assignments (loop ↦ assignment):\n" ++
|
||
String.intercalate "\n"
|
||
(cands.map (fun p => s!" loop #{p.1.val}: assignment #{p.2.val}"))
|
||
|
||
end LicmTransformation
|
||
|
||
end Spa
|