105 lines
4.3 KiB
Lean4
105 lines
4.3 KiB
Lean4
import Spa.Analysis.Reaching
|
||
import Spa.Language.Tagged.Graphs
|
||
|
||
/-!
|
||
# Finding loop-invariant assignments (LICM groundwork)
|
||
|
||
This wires the **reaching-definitions** analysis (`Spa/Analysis/Reaching.lean`)
|
||
to the **tagged 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 pipeline, for each assignment immediately enclosed by a loop:
|
||
|
||
1. locate its CFG state via the tagged-graph bridge (`Program.stateOfNodeId`);
|
||
2. read the reaching definitions at the assignment's *entry*
|
||
(`joinForKey s result` — the join over predecessors, i.e. before the
|
||
assignment itself runs);
|
||
3. union the definition sets of the RHS variables;
|
||
4. map each definition site back to its `NodeId` (`Program.nodeIdOf`) and check
|
||
it is **not** inside the loop body (structural `subtreeIds` membership).
|
||
|
||
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
|
||
|
||
/-- An assignment found inside a loop, paired with the data needed to test its
|
||
invariance against that (immediately enclosing) loop. -/
|
||
structure Candidate where
|
||
/-- The enclosing `whileLoop`'s tag (for reporting). -/
|
||
loopId : NodeId
|
||
/-- Every `NodeId` inside the loop body (the "is-child-of-loop" set). -/
|
||
bodyIds : List NodeId
|
||
/-- The assignment `BasicStmt`'s tag — what labels its CFG node. -/
|
||
assignId : NodeId
|
||
/-- The variables read by the assignment's RHS. -/
|
||
rhsVars : List String
|
||
|
||
/-- Collect every assignment together with its *immediately enclosing* loop.
|
||
`enclosing` carries the current loop's tag and body id-set, or `none` outside any
|
||
loop (in which case assignments are skipped — only in-loop assignments are
|
||
candidates). -/
|
||
def collectCandidates :
|
||
Option (NodeId × List NodeId) → Stmt.Tagged NodeId → List Candidate
|
||
| enc, .basic _ bs =>
|
||
match bs, enc with
|
||
| .assign t _ e, some (loopId, bodyIds) =>
|
||
[{ loopId := loopId, bodyIds := bodyIds, assignId := t,
|
||
rhsVars := e.erase.vars.sort (· ≤ ·) }]
|
||
| _, _ => []
|
||
| enc, .andThen _ a b => collectCandidates enc a ++ collectCandidates enc b
|
||
| enc, .ifElse _ _ a b => collectCandidates enc a ++ collectCandidates enc b
|
||
| _, .whileLoop loopT _ body =>
|
||
collectCandidates (some (loopT, body.subtreeIds)) body
|
||
|
||
/-- 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 ⊥
|
||
|
||
/-- The CFG states marked as definition sites in a `DefSet` (those mapped to
|
||
`true`). -/
|
||
def defSites (prog : Program) (d : DefSet prog) : List prog.State :=
|
||
prog.states.filter (fun s =>
|
||
if h : FiniteMap.MemKey s d then (FiniteMap.locate h).1 else false)
|
||
|
||
/-- 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) : Bool :=
|
||
match prog.stateOfNodeId c.assignId with
|
||
| none => false
|
||
| some s =>
|
||
let entry := joinForKey s (result (DefSet prog) prog)
|
||
let combined : DefSet prog :=
|
||
c.rhsVars.foldl (fun acc k => acc ⊔ lookupDef prog entry k) ⊥
|
||
(defSites prog combined).all (fun site =>
|
||
match prog.nodeIdOf site with
|
||
| some nid => ! decide (nid ∈ c.bodyIds)
|
||
| none => false)
|
||
|
||
/-- The loop-invariant assignments of `prog`, as `(loopId, assignId)` pairs. -/
|
||
def licmCandidates (prog : Program) : List (NodeId × NodeId) :=
|
||
(collectCandidates none prog.tagged).filterMap (fun c =>
|
||
if isInvariant prog c then some (c.loopId, c.assignId) 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.post}: assignment #{p.2.post}"))
|
||
|
||
end LicmTransformation
|
||
|
||
end Spa
|