diff --git a/lean/Spa/Transformation/Licm.lean b/lean/Spa/Transformation/Licm.lean index b1afa7d..51f8654 100644 --- a/lean/Spa/Transformation/Licm.lean +++ b/lean/Spa/Transformation/Licm.lean @@ -1,12 +1,13 @@ import Spa.Analysis.Reaching +import Spa.Language.Equivalence /-! -# Finding loop-invariant assignments (LICM groundwork) +# Loop-invariant code motion 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. +to the AST to find assignments inside a `while` loop whose right-hand side +depends only on definitions made outside the loop. `licmCandidates` reports +these assignments; `hoistProgram` moves eligible leading assignments. 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 @@ -24,9 +25,12 @@ current subtree's CFG into the program's (`Program.rootEmbed`, then one 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. +assignment is reported as loop-invariant. Hoisting additionally requires the +assignment to lead the body, its destination to be absent from the guard, and +no reassignment of that destination in the remaining body. The hoist is guarded +by the original condition, preserving zero-iteration behavior. + +Transitive invariance and motion of non-leading assignments are not implemented. -/ namespace Spa @@ -102,6 +106,50 @@ 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) +/-- Candidate for the leading assignment of a loop body. -/ +def headCandidate (prog : Program) (cond : Expr) (x : String) (rhs : Expr) (tail : Stmt) + (e : Embed (Stmt.whileLoop cond (.andThen (.basic (.assign x rhs)) tail)).cfg prog.cfg) : + Candidate prog := + let body := Stmt.andThen (.basic (.assign x rhs)) tail + let be := (Embed.loop body.cfg).trans e + let ae := (Embed.sequenceLeft (Stmt.basic (.assign x rhs)).cfg tail.cfg).trans be + { encl := { loopState := e.f body.cfg.loopIn, bodyOff := be.off, bodySize := body.cfg.size }, + assignState := ae.singletonIndex, rhsVars := rhs.vars.sort (· ≤ ·) } + +/-- Guard the hoist so that a zero-iteration loop never evaluates the RHS. -/ +def hoistHead (cond : Expr) (x : String) (rhs : Expr) (tail : Stmt) : Stmt := + .ifElse cond (.andThen (.basic (.assign x rhs)) (.whileLoop cond tail)) (.basic .noop) + +/-- Hoist a leading invariant assignment when its destination is neither +reassigned in the remaining body nor read by the guard. -/ +def hoistLoop (prog : Program) (cond : Expr) (body : Stmt) + (e : Embed (Stmt.whileLoop cond body).cfg prog.cfg) : Option Stmt := + match body with + | .andThen (.basic (.assign x rhs)) tail => + if isInvariant prog (headCandidate prog cond x rhs tail e) && + decide (x ∉ tail.writes ∧ x ∉ cond.vars) then + some (hoistHead cond x rhs tail) + else none + | _ => none + +/-- Apply guarded leading-assignment LICM throughout the source tree. When a +loop is hoisted, keep its remaining body intact; further passes can reanalyze it. -/ +def hoistStmt (prog : Program) : (s : Stmt) → Embed s.cfg prog.cfg → Stmt + | .basic bs, _ => .basic bs + | .andThen a b, e => + .andThen (hoistStmt prog a ((Embed.sequenceLeft a.cfg b.cfg).trans e)) + (hoistStmt prog b ((Embed.sequenceRight a.cfg b.cfg).trans e)) + | .ifElse cond a b, e => + .ifElse cond (hoistStmt prog a ((Embed.overlayLeft a.cfg b.cfg).trans e)) + (hoistStmt prog b ((Embed.overlayRight a.cfg b.cfg).trans e)) + | .whileLoop cond body, e => + match hoistLoop prog cond body e with + | some moved => moved + | none => .whileLoop cond (hoistStmt prog body ((Embed.loop body.cfg).trans e)) + +/-- Run reaching definitions on the source program and perform guarded LICM. -/ +def hoistProgram (prog : Program) : Stmt := hoistStmt prog prog.rootStmt prog.rootEmbed + /-- A human-readable report of the loop-invariant assignments. -/ def output (prog : Program) : String := match licmCandidates prog with