From 9165c04d917c88e9a216927c2898ffa4f6e27da6 Mon Sep 17 00:00:00 2001 From: Axel Sorenson Date: Wed, 27 May 2026 22:24:05 -0600 Subject: [PATCH 1/3] Sparse Analysis API --- Veir/Analysis/DataFlow/SparseAnalysis.lean | 371 +++++++++++++++++++++ Veir/Analysis/DataFlow/SparseFact.lean | 14 +- 2 files changed, 376 insertions(+), 9 deletions(-) create mode 100644 Veir/Analysis/DataFlow/SparseAnalysis.lean diff --git a/Veir/Analysis/DataFlow/SparseAnalysis.lean b/Veir/Analysis/DataFlow/SparseAnalysis.lean new file mode 100644 index 000000000..9efbfd340 --- /dev/null +++ b/Veir/Analysis/DataFlow/SparseAnalysis.lean @@ -0,0 +1,371 @@ +module + +public import Veir.Analysis.DataFlow.SparseFact + +public section + +namespace Veir + +namespace SparseForwardDataFlowAnalysis + +variable {kind : FactKind} {Domain : Type} + +-- TODO: When this is verified, we will need something stronger than this for `Domain` +variable [Top Domain] [Bot Domain] [Join Domain] [DecidableEq Domain] + +/-- +The transfer function signature used for custom sparse analyses. + +The framework handles operand subscriptions, invokes this hook with the current +operand lattice elements, and then joins any returned updates into the result +facts. Returning `none` for a result means the transfer contributes no new fact +for that result. +-/ +abbrev VisitOperationFn (Domain : Type) := + OperationPtr -> Array Domain -> WfIRContext OpCode -> Array (Option Domain) + +/-- +Join a sparse lattice fact into the target value state and propagate updates +when it changes. + +This is the generic sparse analysis primitive that merges an incoming lattice +element into the stored state for an SSA value. +-/ +def joinAndPropagate + (kind : FactKind) + [SparseFactSpec kind Domain] + (target : ValuePtr) + (incoming : Domain) + (dfCtx : DataFlowContext) + (irCtx : WfIRContext OpCode) : DataFlowContext := Id.run do + let oldValue := SparseFact.getElement kind target dfCtx + let newValue := oldValue ⊔ incoming + if newValue = oldValue then + return dfCtx + dfCtx.modifyFactAndPropagate kind (.ValuePtr target) + (SparseFact.setLatticeElement · newValue, true) irCtx + +/-- +Return whether the given operation is a branch op. +-/ +private def isBranchOp + (op : OperationPtr) + (irCtx : WfIRContext OpCode) : Bool := + -- TODO: Replace this `.test .test` check once VeIR has proper branch ops. + match (op.get! irCtx.raw).opType with + | .test .test => + true + | _ => + false + +/-- +Return the SSA value forwarded to the given successor's block argument, if any. +-/ +private def getSuccessorOperand? + (op : OperationPtr) + (successorIndex : Nat) + (argumentIndex : Nat) + (irCtx : WfIRContext OpCode) : Option ValuePtr := + if successorIndex >= op.getNumSuccessors! irCtx.raw then + panic! s!"SparseForwardDataFlowAnalysis.getSuccessorOperand?: successor index {successorIndex} out of range" + else + match (op.get! irCtx.raw).opType with + -- TODO: Replace this `.test .test` check once VeIR has proper branch ops. + -- `successorIndex` will become relevant then. + | .test .test => + if argumentIndex < op.getNumOperands! irCtx.raw then + some (op.getOperand! irCtx.raw argumentIndex) + else + none + | _ => + panic! "SparseForwardDataFlowAnalysis.getSuccessorOperand?: non-branch op" + +/-- +Conservatively treat blocks as live when no liveness facts exist. +-/ +private def isBlockLive + (block : BlockPtr) + (dfCtx : DataFlowContext) + (irCtx : WfIRContext OpCode) : Bool := + let _ := block + let _ := dfCtx + let _ := irCtx + true + +/-- +Conservatively treat CFG edges as live when dead code analysis is +not registered. Otherwise consult the liveness lattice, where points are +not live by default. +-/ +private def isEdgeLive + (edge : CFGEdge) + (dfCtx : DataFlowContext) + (_irCtx : WfIRContext OpCode) : Bool := + let _ := edge + let _ := dfCtx + true + +/-- Subscribe to block liveness updates if dead code analysis is registered. -/ +private def subscribeToBlockLiveness + (analysisKind : AnalysisKind) + (block : BlockPtr) + (dfCtx : DataFlowContext) + (irCtx : WfIRContext OpCode) : DataFlowContext := + let _ := analysisKind + let _ := block + let _ := irCtx + dfCtx + +/-- Subscribe to edge liveness updates if dead code analysis is registered. -/ +private def subscribeToEdgeLiveness + (analysisKind : AnalysisKind) + (edge : CFGEdge) + (dfCtx : DataFlowContext) : DataFlowContext := + let _ := analysisKind + let _ := edge + dfCtx + +/-- +Visit a block during sparse initialization. +-/ +private def visitBlock + (kind : FactKind) + [SparseFactSpec kind Domain] + (analysisKind : AnalysisKind) + (block : BlockPtr) + (dfCtx : DataFlowContext) + (irCtx : WfIRContext OpCode) : DataFlowContext := Id.run do + -- Exit early on blocks with no arguments. + if block.getNumArguments! irCtx.raw = 0 then + return dfCtx + + -- If the block is not live, bail out. + if !isBlockLive block dfCtx irCtx then + return dfCtx + + let some parentRegion := (block.get! irCtx.raw).parent + | return dfCtx + + -- The argument lattices of entry blocks are set by region control flow or + -- the callgraph. + if (parentRegion.get! irCtx.raw).firstBlock = some block then + -- TODO: Mirror MLIR's handling of `visitCallableOperation` and + -- `visitRegionSuccessors` and `visitNonControlFlowArgumentsImpl` + -- for entry blocks. + return dfCtx + + let mut dfCtx := dfCtx + + -- Iterate over the predecessors of the non-entry block. + let mut maybePredUse := (block.get! irCtx.raw).firstUse + + while let some predUse := maybePredUse do + let predUseStruct := predUse.get! irCtx.raw + maybePredUse := predUseStruct.nextUse + + let predecessorOp := predUseStruct.owner + let some predecessorBlock := (predecessorOp.get! irCtx.raw).parent + | continue + + let edge : CFGEdge := { source := predecessorBlock, target := block } + dfCtx := subscribeToEdgeLiveness analysisKind edge dfCtx + + -- If the edge from the predecessor block to the current block is not live, + -- bail out. + if !isEdgeLive edge dfCtx irCtx then + continue + + -- Check if we can reason about the dataflow from the predecessor. + if !isBranchOp predecessorOp irCtx then + for target in block.getArguments! irCtx.raw do + dfCtx := joinAndPropagate kind target ⊤ dfCtx irCtx + return dfCtx + + for i in [0:block.getNumArguments! irCtx.raw] do + let arg := block.getArgument i + match getSuccessorOperand? predecessorOp predUse.index i irCtx with + | some operand => + -- Add the current block start program point as a dependency of the + -- predecessor block's successor operand lattice state, so this block + -- is revisited when that operand lattice changes. + let dependentPoint := InsertPoint.atStart! block irCtx.raw + let workItem : WorkItem := (dependentPoint, analysisKind) + dfCtx := dfCtx.modifyFact kind (.ValuePtr operand) (fun state => + if state.dependents.any (fun dependent => + dependent.1 = dependentPoint && dependent.2 = analysisKind) then + -- Do not add dependent again if it's already added. + state + else + state.addDependent workItem) + + -- Call transfer function + let incoming := + SparseFact.getElement kind operand dfCtx + dfCtx := joinAndPropagate kind arg incoming dfCtx irCtx + | none => + -- Conservatively consider internally produced arguments to be at the + -- pessimistic sparse state. + dfCtx := joinAndPropagate kind arg ⊤ dfCtx irCtx + + return dfCtx + +mutual + +/-- +Ensure an operand has a sparse lattice state and subscribe the current sparse +analysis to its updates. This is what makes use-def driven revisitation work. +-/ +partial def subscribeToOperand + (kind : FactKind) + [SparseFactSpec kind Domain] + (analysisKind : AnalysisKind) + (operand : ValuePtr) + (dfCtx : DataFlowContext) : DataFlowContext := + dfCtx.modifyFact kind (.ValuePtr operand) (fun state => + state.subscribe analysisKind) + +/-- +Visit one operation in the sparse analysis. +We first subscribe to operand lattices, then hand the operation and current +operand lattice elements to the user provided transfer function. The framework +applies any returned result updates itself. +-/ +partial def visitOperation + (kind : FactKind) + [SparseFactSpec kind Domain] + (analysisKind : AnalysisKind) + (visitOperationImpl : VisitOperationFn Domain) + (op : OperationPtr) + (dfCtx : DataFlowContext) + (irCtx : WfIRContext OpCode) : DataFlowContext := Id.run do + -- Exit early on operations with no results. + if op.getNumResults! irCtx.raw = 0 then + return dfCtx + + -- If the containing block is not live, bail out. Liveness is by default + -- unreachable until proven live, so a missing state is treated as dead. + if let some parentBlock := (op.get! irCtx.raw).parent then + if !isBlockLive parentBlock dfCtx irCtx then + return dfCtx + + -- TODO: Mirror MLIR more closely by `visitRegionSuccessors` + -- Comment: The results of a region branch operation are determined by control-flow. + + -- TODO: Mirror MLIR more closely by `visitCallOperation` + + let mut dfCtx := dfCtx + for operand in op.getOperands! irCtx.raw do + dfCtx := subscribeToOperand kind analysisKind operand dfCtx + + let operandLatticeElements := (op.getOperands! irCtx.raw).map (fun operand => + SparseFact.getElement kind operand dfCtx) + let resultUpdates := visitOperationImpl op operandLatticeElements irCtx + + for (result, incoming?) in (op.getResults! irCtx.raw).zip resultUpdates do + if let some incoming := incoming? then + dfCtx := joinAndPropagate kind result incoming dfCtx irCtx + return dfCtx + +/-- +Recursively initialize an operation tree for sparse analysis. +Visit the current operation first, then walk its nested regions, +blocks, and nested operations. +-/ +partial def initializeRecursively + (kind : FactKind) + [SparseFactSpec kind Domain] + (analysisKind : AnalysisKind) + (visitOperationImpl : VisitOperationFn Domain) + (op : OperationPtr) + (dfCtx : DataFlowContext) + (irCtx : WfIRContext OpCode) : DataFlowContext := Id.run do + -- Initialize the analysis by visiting every owner of an SSA value (all + -- operations and blocks). + let mut dfCtx := dfCtx + dfCtx := visitOperation kind analysisKind visitOperationImpl op dfCtx irCtx + + for regionPtr in (op.get! irCtx.raw).regions do + let region := regionPtr.get! irCtx.raw + let mut maybeBlock := region.firstBlock + + while let some block := maybeBlock do + dfCtx := subscribeToBlockLiveness analysisKind block dfCtx irCtx + dfCtx := visitBlock kind analysisKind block dfCtx irCtx + let mut maybeOp := (block.get! irCtx.raw).firstOp + + while let some nestedOp := maybeOp do + dfCtx := initializeRecursively kind analysisKind visitOperationImpl nestedOp dfCtx irCtx + maybeOp := (nestedOp.get! irCtx.raw).next + + maybeBlock := (block.get! irCtx.raw).next + dfCtx + +end + +/-- +Initialize the analysis by visiting every owner of an SSA value: all +operations and blocks. +-/ +private def init + (kind : FactKind) + [SparseFactSpec kind Domain] + (analysisKind : AnalysisKind) + (visitOperationImpl : VisitOperationFn Domain) + (top : OperationPtr) + (dfCtx : DataFlowContext) + (irCtx : WfIRContext OpCode) : DataFlowContext := Id.run do + -- Mark the entry block arguments as having reached their pessimistic + -- fixpoints. + let mut dfCtx := dfCtx + for regionPtr in (top.get! irCtx.raw).regions do + let region := regionPtr.get! irCtx.raw + if let some firstBlock := region.firstBlock then + for arg in firstBlock.getArguments! irCtx.raw do + dfCtx := joinAndPropagate kind arg ⊤ dfCtx irCtx + + initializeRecursively kind analysisKind visitOperationImpl top dfCtx irCtx + +/-- +Visit an insertion point. If this is at beginning of block and all +control flow predecessors or callsites are known, then the arguments' +lattices are propagated from them. If this is after call operation or an +operation with region control-flow, then its result lattices are set +accordingly. Otherwise, the operation transfer function is invoked. +-/ +private def visit + (kind : FactKind) + [SparseFactSpec kind Domain] + (analysisKind : AnalysisKind) + (visitOperationImpl : VisitOperationFn Domain) + (point : InsertPoint) + (dfCtx : DataFlowContext) + (irCtx : WfIRContext OpCode) : DataFlowContext := + match point.prev! irCtx.raw with + | some prevOp => + visitOperation kind analysisKind visitOperationImpl prevOp dfCtx irCtx + | none => + match point.block! irCtx.raw with + | some block => + visitBlock kind analysisKind block dfCtx irCtx + | none => + dfCtx + +/-- +Build a sparse forward analysis over one abstract value domain. + +Sparse facts default to `⊥`. Whenever control flow or transfer functions lose +precision, the framework conservatively joins `⊤` into the affected values. +-/ +def new + (kind : FactKind) + [SparseFactSpec kind Domain] + (analysisKind : AnalysisKind) + (visitOperationImpl : VisitOperationFn Domain) + : DataFlowAnalysis := + { kind := analysisKind + init := init kind analysisKind visitOperationImpl + visit := visit kind analysisKind visitOperationImpl } + +end SparseForwardDataFlowAnalysis + +end Veir diff --git a/Veir/Analysis/DataFlow/SparseFact.lean b/Veir/Analysis/DataFlow/SparseFact.lean index c22e349d1..21e57d515 100644 --- a/Veir/Analysis/DataFlow/SparseFact.lean +++ b/Veir/Analysis/DataFlow/SparseFact.lean @@ -68,15 +68,11 @@ instance : FactSpec kind where end -def getElement? (kind : FactKind) [SparseFactSpec kind Domain] [FactSpec kind] - (ssaValue : ValuePtr) (dfCtx : DataFlowContext) : Option Domain := do - let state ← dfCtx.getFact? kind (.ValuePtr ssaValue) - return latticeElement state - -def getElementD (kind : FactKind) [SparseFactSpec kind Domain] [FactSpec kind] - (ssaValue : ValuePtr) (fallback : Domain) - (dfCtx : DataFlowContext) : Domain := - (getElement? kind ssaValue dfCtx).getD fallback +def getElement (kind : FactKind) [SparseFactSpec kind Domain] [FactSpec kind] + [Bot Domain] (ssaValue : ValuePtr) (dfCtx : DataFlowContext) : Domain := + match dfCtx.getFact? kind (.ValuePtr ssaValue) with + | some state => latticeElement state + | none => ⊥ end SparseFact From 8b1ca3e63c793a590434b4a889bd34317fa3e586 Mon Sep 17 00:00:00 2001 From: Axel Sorenson Date: Mon, 27 Jul 2026 18:55:35 -0600 Subject: [PATCH 2/3] (Sparse) Transfer functions no longer modify DataFlowContext directly --- Veir/Analysis/DataFlow/SparseAnalysis.lean | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/Veir/Analysis/DataFlow/SparseAnalysis.lean b/Veir/Analysis/DataFlow/SparseAnalysis.lean index 9efbfd340..dc9818d90 100644 --- a/Veir/Analysis/DataFlow/SparseAnalysis.lean +++ b/Veir/Analysis/DataFlow/SparseAnalysis.lean @@ -80,9 +80,7 @@ private def getSuccessorOperand? | _ => panic! "SparseForwardDataFlowAnalysis.getSuccessorOperand?: non-branch op" -/-- -Conservatively treat blocks as live when no liveness facts exist. --/ +/-- Conservatively treat blocks as live when no liveness facts exist. -/ private def isBlockLive (block : BlockPtr) (dfCtx : DataFlowContext) @@ -105,7 +103,7 @@ private def isEdgeLive let _ := dfCtx true -/-- Subscribe to block liveness updates if dead code analysis is registered. -/ +/-- No-op when no liveness analysis is registered. -/ private def subscribeToBlockLiveness (analysisKind : AnalysisKind) (block : BlockPtr) @@ -116,7 +114,7 @@ private def subscribeToBlockLiveness let _ := irCtx dfCtx -/-- Subscribe to edge liveness updates if dead code analysis is registered. -/ +/-- No-op when no liveness analysis is registered. -/ private def subscribeToEdgeLiveness (analysisKind : AnalysisKind) (edge : CFGEdge) From 84c6f1d438b23fa45e308416c18c8e9e071e6d6c Mon Sep 17 00:00:00 2001 From: Axel Sorenson Date: Tue, 18 Aug 2026 19:20:59 -0600 Subject: [PATCH 3/3] Use branch op interface --- Veir/Analysis/DataFlow/SparseAnalysis.lean | 46 +++++----------------- 1 file changed, 9 insertions(+), 37 deletions(-) diff --git a/Veir/Analysis/DataFlow/SparseAnalysis.lean b/Veir/Analysis/DataFlow/SparseAnalysis.lean index dc9818d90..250a87198 100644 --- a/Veir/Analysis/DataFlow/SparseAnalysis.lean +++ b/Veir/Analysis/DataFlow/SparseAnalysis.lean @@ -1,6 +1,7 @@ module public import Veir.Analysis.DataFlow.SparseFact +public import Veir.Interfaces.ControlFlowInterfaces public section @@ -45,41 +46,6 @@ def joinAndPropagate dfCtx.modifyFactAndPropagate kind (.ValuePtr target) (SparseFact.setLatticeElement · newValue, true) irCtx -/-- -Return whether the given operation is a branch op. --/ -private def isBranchOp - (op : OperationPtr) - (irCtx : WfIRContext OpCode) : Bool := - -- TODO: Replace this `.test .test` check once VeIR has proper branch ops. - match (op.get! irCtx.raw).opType with - | .test .test => - true - | _ => - false - -/-- -Return the SSA value forwarded to the given successor's block argument, if any. --/ -private def getSuccessorOperand? - (op : OperationPtr) - (successorIndex : Nat) - (argumentIndex : Nat) - (irCtx : WfIRContext OpCode) : Option ValuePtr := - if successorIndex >= op.getNumSuccessors! irCtx.raw then - panic! s!"SparseForwardDataFlowAnalysis.getSuccessorOperand?: successor index {successorIndex} out of range" - else - match (op.get! irCtx.raw).opType with - -- TODO: Replace this `.test .test` check once VeIR has proper branch ops. - -- `successorIndex` will become relevant then. - | .test .test => - if argumentIndex < op.getNumOperands! irCtx.raw then - some (op.getOperand! irCtx.raw argumentIndex) - else - none - | _ => - panic! "SparseForwardDataFlowAnalysis.getSuccessorOperand?: non-branch op" - /-- Conservatively treat blocks as live when no liveness facts exist. -/ private def isBlockLive (block : BlockPtr) @@ -174,14 +140,20 @@ private def visitBlock continue -- Check if we can reason about the dataflow from the predecessor. - if !isBranchOp predecessorOp irCtx then + if !(predecessorOp.getOpType! irCtx.raw).isTerminator then for target in block.getArguments! irCtx.raw do dfCtx := joinAndPropagate kind target ⊤ dfCtx irCtx return dfCtx + let some successorOperands := + BranchOpInterface.getSuccessorOperands? predecessorOp predUse.index irCtx.raw + | for target in block.getArguments! irCtx.raw do + dfCtx := joinAndPropagate kind target ⊤ dfCtx irCtx + return dfCtx + for i in [0:block.getNumArguments! irCtx.raw] do let arg := block.getArgument i - match getSuccessorOperand? predecessorOp predUse.index i irCtx with + match successorOperands[i]? with | some operand => -- Add the current block start program point as a dependency of the -- predecessor block's successor operand lattice state, so this block