Update LICM/Reaching to node use NodeId
This commit is contained in:
@@ -1,6 +1,5 @@
|
||||
import Spa.Analysis.Forward
|
||||
import Spa.Lattice.Finset
|
||||
import Spa.Language.Tagged.Graphs
|
||||
import Spa.Showable
|
||||
|
||||
namespace Spa
|
||||
@@ -13,20 +12,18 @@ instance {n : ℕ} : Showable (Finset (Fin n)) :=
|
||||
(fun i rest => if i ∈ s then show' i ++ ", " ++ rest else rest) ""
|
||||
++ "}"⟩
|
||||
|
||||
abbrev DefSet (prog : Program) : Type := Finset prog.NodeId
|
||||
abbrev DefSet (prog : Program) : Type := Finset prog.State
|
||||
|
||||
namespace ReachingAnalysis
|
||||
|
||||
variable (prog : Program)
|
||||
|
||||
def genSet (s : prog.State) : DefSet prog := (prog.nodeIdOf s).elim {} (fun x => {x})
|
||||
|
||||
def eval (s : prog.State) (vs : VariableValues (DefSet prog) prog) : VariableValues (DefSet prog) prog :=
|
||||
match prog.code s with
|
||||
| none => vs
|
||||
| some bs =>
|
||||
match bs with
|
||||
| .assign k _ => FiniteMap.generalizedUpdate id (fun _ _ => genSet prog s) [k] vs
|
||||
| .assign k _ => FiniteMap.generalizedUpdate id (fun _ _ => {s}) [k] vs
|
||||
| .noop => vs
|
||||
|
||||
lemma eval_mono (s : prog.State) :
|
||||
@@ -50,12 +47,11 @@ def output : String :=
|
||||
abbrev Run (prog : Program) : Type := List (prog.State × BasicStmt)
|
||||
|
||||
@[aesop unsafe cases]
|
||||
inductive LastAssign (prog : Program) (x : String) : Run prog → prog.NodeId → Prop
|
||||
| here (s : prog.State) (e : Expr) (hc : prog.code s = some (.assign x e))
|
||||
(rest : Run prog) :
|
||||
LastAssign prog x ((s, .assign x e) :: rest) (prog.nodeIdOfNonempty s hc)
|
||||
inductive LastAssign (prog : Program) (x : String) : Run prog → prog.State → Prop
|
||||
| here (s : prog.State) (e : Expr) (rest : Run prog) :
|
||||
LastAssign prog x ((s, .assign x e) :: rest) s
|
||||
| there (s : prog.State) (bs : BasicStmt) (hc : prog.code s = some bs)
|
||||
(rest : Run prog) {n : prog.NodeId} :
|
||||
(rest : Run prog) {n : prog.State} :
|
||||
(∀ e, bs ≠ .assign x e) → LastAssign prog x rest n →
|
||||
LastAssign prog x ((s, bs) :: rest) n
|
||||
|
||||
@@ -73,7 +69,7 @@ instance stateInterp : StateInterpretation (DefSet prog) prog where
|
||||
Post := @runOfTrace prog
|
||||
|
||||
interp vs run := ∀ (x : String) (assigners : DefSet prog), (x, assigners) ∈ vs →
|
||||
∀ (n : prog.NodeId), LastAssign prog x run n → n ∈ assigners
|
||||
∀ (n : prog.State), LastAssign prog x run n → n ∈ assigners
|
||||
interp_sup := by
|
||||
intro vs₁ vs₂ run h x assigners hmem n hla
|
||||
obtain ⟨a₁, a₂, rfl, h₁, h₂⟩ := FiniteMap.mem_sup hmem
|
||||
@@ -105,8 +101,7 @@ private lemma valid_step (s : prog.State) {ρ₁ ρ₂ : Env}
|
||||
by_cases hx : k = x
|
||||
· subst hx
|
||||
have hd := FiniteMap.generalizedUpdate_mem_eq (List.mem_singleton.mpr rfl) hmem
|
||||
rcases hla
|
||||
<;> simp [Program.nodeIdOfNonempty, hd, genSet, Option.get] <;> aesop
|
||||
rcases hla <;> simp [hd] <;> aesop
|
||||
· have hmem' := FiniteMap.generalizedUpdate_not_mem_backward
|
||||
(fun hc => hx (List.mem_singleton.mp hc)) hmem
|
||||
aesop
|
||||
|
||||
@@ -1,26 +1,30 @@
|
||||
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.
|
||||
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 pipeline, for each assignment immediately enclosed by a loop:
|
||||
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. 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);
|
||||
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. map each definition site back to its `RawId` (`Program.nodeIdOf`) and check
|
||||
it is **not** inside the loop body (structural `subtreeIds` membership).
|
||||
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
|
||||
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.
|
||||
-/
|
||||
@@ -29,64 +33,74 @@ namespace Spa
|
||||
|
||||
namespace LicmTransformation
|
||||
|
||||
open Forward
|
||||
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 `whileLoop`'s tag (for reporting). -/
|
||||
loopId : prog.NodeId
|
||||
/-- Every node id inside the loop body (the "is-child-of-loop" set). -/
|
||||
bodyIds : List prog.NodeId
|
||||
/-- The assignment `BasicStmt`'s tag — what labels its CFG node. -/
|
||||
assignId : prog.NodeId
|
||||
/-- 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.
|
||||
`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 (prog : Program) (enc : Option (prog.NodeId × List prog.NodeId)) :
|
||||
Stmt.Tagged prog.NodeId → List (Candidate prog)
|
||||
| .basic _ bs =>
|
||||
`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 t _ e, some (loopId, bodyIds) =>
|
||||
[{ loopId := loopId, bodyIds := bodyIds, assignId := t,
|
||||
rhsVars := e.erase.vars.sort (· ≤ ·) }]
|
||||
| .assign _ ex, some l =>
|
||||
[{ encl := l, assignState := e.singletonIndex,
|
||||
rhsVars := ex.vars.sort (· ≤ ·) }]
|
||||
| _, _ => []
|
||||
| .andThen _ a b => collectCandidates prog enc a ++ collectCandidates prog enc b
|
||||
| .ifElse _ _ a b => collectCandidates prog enc a ++ collectCandidates prog enc b
|
||||
| .whileLoop loopT _ body =>
|
||||
collectCandidates prog (some (loopT, body.subtreeIds)) body
|
||||
| .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 ⊥
|
||||
|
||||
/-- The AST node ids marked as definition sites in a `DefSet`. With the
|
||||
`Finset`-of-AST-ids lattice these are just the elements of the set. -/
|
||||
def defSites (prog : Program) (d : DefSet prog) : List prog.NodeId :=
|
||||
(List.finRange prog.size).filter (fun i => decide (i ∈ d))
|
||||
|
||||
/-- Is the candidate assignment loop-invariant: do all reaching definitions of
|
||||
its RHS variables lie outside the loop body? Reaching sets are now keyed by AST
|
||||
node id, so we compare against the loop-body ids directly (embedding the raw
|
||||
body ids into `p.NodeId`). -/
|
||||
its RHS variables lie outside the loop body? -/
|
||||
def isInvariant (prog : Program) (c : Candidate prog) : 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 nid => ! decide (nid ∈ c.bodyIds))
|
||||
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 `(loopId, assignId)` pairs. -/
|
||||
def licmCandidates (prog : Program) : List (prog.NodeId × prog.NodeId) :=
|
||||
(collectCandidates prog none prog.taggedFin).filterMap (fun c =>
|
||||
if isInvariant prog c then some (c.loopId, c.assignId) else none)
|
||||
/-- 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 :=
|
||||
|
||||
Reference in New Issue
Block a user