diff --git a/UnitTest.lean b/UnitTest.lean index d063d49f5..0c6fbc3cb 100644 --- a/UnitTest.lean +++ b/UnitTest.lean @@ -15,3 +15,4 @@ import UnitTest.ConstantValue import UnitTest.Evaluate import UnitTest.FoldDecision import UnitTest.SideEffectInterfaces +import UnitTest.DataFlowFramework.SparseConstantPropagation diff --git a/UnitTest/DataFlowFramework/Helpers.lean b/UnitTest/DataFlowFramework/Helpers.lean index 16e4234e4..60261b8ca 100644 --- a/UnitTest/DataFlowFramework/Helpers.lean +++ b/UnitTest/DataFlowFramework/Helpers.lean @@ -1,4 +1,6 @@ import Veir.Analysis.DataFlow.DeadCodeAnalysis +import Veir.Analysis.DataFlow.SparseFact +import Veir.Analysis.DataFlow.SparseConstantPropagationAnalysis import Veir.Parser.MlirParser open Std (HashMap) @@ -166,3 +168,28 @@ def runWithAnalyses let some dfCtx := fixpointSolve top analyses parserState.ctx | return "analysis did not converge" return renderReport (check top dfCtx parserState) + +/-- Sparse constant propagation helpers. -/ +def showConstantDomain : AbstractConstant -> String + | .top => + "top" + | .bottom => + "bottom" + | .constant c => + s!"const({c.value} : i{c.bitwidth})" + +def checkNamedConstants + (dfCtx : DataFlowContext) + (valueDefs : HashMap String ValuePtr) + (expected : Array (String × AbstractConstant)) : MismatchReport := Id.run do + let mut report := #[] + for (name, expectedValue) in expected do + let some value := valueDefs[name]? | + report := report.push s!"constant {name}: missing value definition" + continue + let observedValue := + SparseFact.getElement .sparseConstant value dfCtx + if observedValue != expectedValue then + report := report.push + s!"constant {name}: expected {showConstantDomain expectedValue}, observed {showConstantDomain observedValue}" + report diff --git a/UnitTest/DataFlowFramework/SparseConstantPropagation.lean b/UnitTest/DataFlowFramework/SparseConstantPropagation.lean new file mode 100644 index 000000000..f5dfb2f3b --- /dev/null +++ b/UnitTest/DataFlowFramework/SparseConstantPropagation.lean @@ -0,0 +1,192 @@ +import UnitTest.DataFlowFramework.Helpers + +import Veir.Analysis.DataFlow.Domains.ConstantDomain +import Veir.Analysis.DataFlow.SparseConstantPropagationAnalysis + +open Veir + +private def constInt (bitwidth : Nat) (value : Int) : AbstractConstant := + .constant ⟨bitwidth, Data.LLVM.Int.constant bitwidth value⟩ + +private def run + (mlir : String) + (expected : Array (String × AbstractConstant)) : String := + runWithAnalyses mlir #[Veir.SparseConstantPropagationAnalysis] + (fun top dfCtx parserState => Id.run do + match recoverNames top parserState.ctx mlir with + | Except.error err => + return #[err] + | Except.ok recovered => + checkNamedConstants dfCtx recovered.values expected) + +private def testAddiAllConstant : String := + run + r#""builtin.module"() ({ +^bb0: + %a = "arith.constant"() <{ value = 5 : i32 }> : () -> i32 + %b = "arith.constant"() <{ value = 7 : i32 }> : () -> i32 + %c = "arith.addi"(%a, %b) : (i32, i32) -> i32 +}) : () -> ()"# + #[ ("a", constInt 32 5) + , ("b", constInt 32 7) + , ("c", constInt 32 12) + ] + +private def testMuliAllConstant : String := + run + r#""builtin.module"() ({ +^bb0: + %a = "arith.constant"() <{ value = 3 : i32 }> : () -> i32 + %b = "arith.constant"() <{ value = 2 : i32 }> : () -> i32 + %c = "arith.muli"(%a, %b) : (i32, i32) -> i32 +}) : () -> ()"# + #[ ("a", constInt 32 3) + , ("b", constInt 32 2) + , ("c", constInt 32 6) + ] + +private def testAndiAllConstant : String := + run + r#""builtin.module"() ({ +^bb0: + %a = "arith.constant"() <{ value = 27 : i32 }> : () -> i32 + %b = "arith.constant"() <{ value = 3 : i32 }> : () -> i32 + %c = "arith.andi"(%a, %b) : (i32, i32) -> i32 +}) : () -> ()"# + #[ ("a", constInt 32 27) + , ("b", constInt 32 3) + , ("c", constInt 32 3) + ] + +private def testSubiAllConstant : String := + run + r#""builtin.module"() ({ +^bb0: + %a = "arith.constant"() <{ value = 12 : i32 }> : () -> i32 + %b = "arith.constant"() <{ value = 37 : i32 }> : () -> i32 + %c = "arith.subi"(%a, %b) : (i32, i32) -> i32 +}) : () -> ()"# + #[ ("a", constInt 32 12) + , ("b", constInt 32 37) + , ("c", constInt 32 (-25)) + ] + +private def testAddiUnknownOperand : String := + run + r#""builtin.module"() ({ +^bb0: + %a = "arith.constant"() <{ value = -3 : i32 }> : () -> i32 + %u = "test.test"() : () -> i32 + %c = "arith.addi"(%a, %u) : (i32, i32) -> i32 +}) : () -> ()"# + #[ ("a", constInt 32 (-3)) + , ("u", ⊤) + , ("c", ⊤) + ] + +private def testMuliUnknownOperand : String := + run + r#""builtin.module"() ({ +^bb0: + %a = "arith.constant"() <{ value = 7 : i32 }> : () -> i32 + %u = "test.test"() : () -> i32 + %c = "arith.muli"(%a, %u) : (i32, i32) -> i32 +}) : () -> ()"# + #[ ("a", constInt 32 7) + , ("u", ⊤) + , ("c", ⊤) + ] + +private def testAndiUnknownOperand : String := + run + r#""builtin.module"() ({ +^bb0: + %a = "arith.constant"() <{ value = -2 : i32 }> : () -> i32 + %u = "test.test"() : () -> i32 + %c = "arith.andi"(%a, %u) : (i32, i32) -> i32 +}) : () -> ()"# + #[ ("a", constInt 32 (-2)) + , ("u", ⊤) + , ("c", ⊤) + ] + +private def testSubiUnknownOperand : String := + run + r#""builtin.module"() ({ +^bb0: + %a = "arith.constant"() <{ value = 0 : i32 }> : () -> i32 + %u = "test.test"() : () -> i32 + %c = "arith.subi"(%a, %u) : (i32, i32) -> i32 +}) : () -> ()"# + #[ ("a", constInt 32 0) + , ("u", ⊤) + , ("c", ⊤) + ] + +private def testStandalonePropagatesAcrossLiveByDefaultEdge : String := + run + r#""builtin.module"() ({ +^bb0: + %x = "arith.constant"() <{ value = 1 : i32 }> : () -> i32 + "test.test"(%x, %x)[^bb1] : (i32, i32) -> () +^bb1(%dead : i32, %y : i32): + %z = "arith.addi"(%y, %y) : (i32, i32) -> i32 +}) : () -> ()"# + #[ ("x", constInt 32 1) + , ("y", constInt 32 1) + , ("z", constInt 32 2) + ] + +/-- +info: "ok" +-/ +#guard_msgs in +#eval! testAddiAllConstant + +/-- +info: "ok" +-/ +#guard_msgs in +#eval! testMuliAllConstant + +/-- +info: "ok" +-/ +#guard_msgs in +#eval! testAndiAllConstant + +/-- +info: "ok" +-/ +#guard_msgs in +#eval! testSubiAllConstant + +/-- +info: "ok" +-/ +#guard_msgs in +#eval! testAddiUnknownOperand + +/-- +info: "ok" +-/ +#guard_msgs in +#eval! testMuliUnknownOperand + +/-- +info: "ok" +-/ +#guard_msgs in +#eval! testAndiUnknownOperand + +/-- +info: "ok" +-/ +#guard_msgs in +#eval! testSubiUnknownOperand + +/-- +info: "ok" +-/ +#guard_msgs in +#eval! testStandalonePropagatesAcrossLiveByDefaultEdge diff --git a/Veir/Analysis/DataFlow/Facts.lean b/Veir/Analysis/DataFlow/Facts.lean index 82f0b6693..9a92a8d1f 100644 --- a/Veir/Analysis/DataFlow/Facts.lean +++ b/Veir/Analysis/DataFlow/Facts.lean @@ -3,6 +3,7 @@ module public import Veir.GlobalOpInfo public import Veir.Analysis.DataFlow.Domains.LivenessDomain public import Veir.Rewriter.InsertPoint +public import Veir.Analysis.DataFlow.Domains.ConstantDomain open Std (HashMap Queue) @@ -54,6 +55,7 @@ Tags to match on for different `DataFlowAnalysis` types. inductive AnalysisKind where | dominance | deadCode + | sparseConstantPropagation deriving BEq, Hashable, Repr, DecidableEq /-- @@ -63,6 +65,7 @@ inductive FactKind where | dominator | regionMetadata | liveness + | sparseConstant deriving BEq, ReflBEq, LawfulBEq, Hashable, Repr, DecidableEq abbrev WorkItem := InsertPoint × AnalysisKind @@ -101,6 +104,7 @@ The fact specific data stored for each fact kind. | .dominator => DominatorPayload | .regionMetadata => RegionMetadataPayload | .liveness => LivenessPayload + | .sparseConstant => SparsePayload AbstractConstant /-- A dataflow fact stored by the framework. diff --git a/Veir/Analysis/DataFlow/SparseAnalysis.lean b/Veir/Analysis/DataFlow/SparseAnalysis.lean new file mode 100644 index 000000000..dc9818d90 --- /dev/null +++ b/Veir/Analysis/DataFlow/SparseAnalysis.lean @@ -0,0 +1,369 @@ +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 + +/-- No-op when no liveness analysis is registered. -/ +private def subscribeToBlockLiveness + (analysisKind : AnalysisKind) + (block : BlockPtr) + (dfCtx : DataFlowContext) + (irCtx : WfIRContext OpCode) : DataFlowContext := + let _ := analysisKind + let _ := block + let _ := irCtx + dfCtx + +/-- No-op when no liveness 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/SparseConstantPropagationAnalysis.lean b/Veir/Analysis/DataFlow/SparseConstantPropagationAnalysis.lean new file mode 100644 index 000000000..84602204f --- /dev/null +++ b/Veir/Analysis/DataFlow/SparseConstantPropagationAnalysis.lean @@ -0,0 +1,129 @@ +module + +public import Veir.Analysis.DataFlow.Domains.ConstantDomain +public import Veir.Analysis.DataFlow.SparseAnalysis + +public section + +namespace Veir + +namespace SparseConstantPropagation + +instance : SparseFactSpec .sparseConstant AbstractConstant where + payloadEq := rfl + +def kind : AnalysisKind := + .sparseConstantPropagation + +/-- +Fold a binary operation on known constants when bitwidths agree. +Returns `none` if widths mismatch or folding yields no value. +-/ +def foldKnownBinary? + (lhs rhs : ConcreteConstant) + (f : {w : Nat} -> Data.LLVM.Int w -> Data.LLVM.Int w -> Option (Data.LLVM.Int w)) + : Option ConcreteConstant := + if h : lhs.bitwidth = rhs.bitwidth then + let rhsValue := Data.LLVM.Int.cast rhs.value (Eq.symm h) + f lhs.value rhsValue |> .map ({ bitwidth := lhs.bitwidth, value := · }) + else + none + +/-- +Try to fold a binary op from operand lattice elements. +Only folds when there are exactly two operands and both are known constants. +-/ +def foldBinaryOp? + (operandLatticeElements : Array AbstractConstant) + (f : {w : Nat} -> Data.LLVM.Int w -> Data.LLVM.Int w -> Option (Data.LLVM.Int w)) + : Option AbstractConstant := + if operandLatticeElements.size ≠ 2 then + none + else + match operandLatticeElements[0]?, operandLatticeElements[1]? with + | some (AbstractConstant.constant lhs), some (AbstractConstant.constant rhs) => + foldKnownBinary? lhs rhs f |> .map (.constant ·) + | _, _ => + none + +/-- Produce a folded constant when possible, otherwise conservatively yield `⊤`. -/ +def foldedOrUnknown + (numResults : Nat) + (folded : Option AbstractConstant) : Array (Option AbstractConstant) := + match folded with + | some constant => + #[some constant] + | none => + Array.replicate numResults (some ⊤) + +/-- +Sparse constant propagation transfer function. +- region operations conservatively force results to the unknown state, +- operands at `⊥` delay propagation, +- otherwise we try to fold and report any discovered constant facts. +-/ +def transfer + (op : OperationPtr) + (operandLatticeElements : Array AbstractConstant) + (irCtx : WfIRContext OpCode) : Array (Option AbstractConstant) := + let numResults := op.getNumResults! irCtx.raw + + -- Don't try to simulate the results of a region operation as we can't + -- guarantee that folding will be out-of-place. We don't allow in-place + -- folds as the desire here is for simulated execution, and not general + -- folding. + if op.getNumRegions! irCtx.raw ≠ 0 then + Array.replicate numResults (some ⊤) + + -- Wait until every operand lattice has been initialized before trying to + -- infer a result. + else if operandLatticeElements.any (· = ⊥) then + Array.replicate numResults none + + -- TODO: Mirror MLIR's generic `op->fold` path once Veir has an operation + -- folder and fold-result representation. For now we manually handle the + -- arithmetic ops. + else + match (op.get! irCtx.raw).opType with + | .arith .constant => + if numResults > 0 then + let intAttr := (op.getProperties! irCtx.raw Arith.constant).value + #[some (.constant ⟨intAttr.type.bitwidth, + Data.LLVM.Int.constant intAttr.type.bitwidth intAttr.value⟩)] + else + #[] + | .arith .addi => + let flags := op.getProperties! irCtx.raw Arith.addi + foldedOrUnknown numResults <| foldBinaryOp? operandLatticeElements (fun lhs rhs => + match Data.LLVM.Int.add lhs rhs flags.attr.nsw flags.attr.nuw with + | .val v => some (.val v) + | .poison => none) + | .arith .muli => + let flags := op.getProperties! irCtx.raw Arith.muli + foldedOrUnknown numResults <| foldBinaryOp? operandLatticeElements (fun lhs rhs => + match Data.LLVM.Int.mul lhs rhs flags.attr.nsw flags.attr.nuw with + | .val v => some (.val v) + | .poison => none) + | .arith .andi => + foldedOrUnknown numResults <| foldBinaryOp? operandLatticeElements (fun lhs rhs => + match lhs, rhs with + | .val lhs', .val rhs' => some (.val (BitVec.and lhs' rhs')) + | _, _ => none) + | .arith .subi => + let flags := op.getProperties! irCtx.raw Arith.subi + foldedOrUnknown numResults <| foldBinaryOp? operandLatticeElements (fun lhs rhs => + match Data.LLVM.Int.sub lhs rhs flags.attr.nsw flags.attr.nuw with + | .val v => some (.val v) + | .poison => none) + | _ => + Array.replicate numResults (some ⊤) + +end SparseConstantPropagation + +def SparseConstantPropagationAnalysis : DataFlowAnalysis := + SparseForwardDataFlowAnalysis.new + .sparseConstant + SparseConstantPropagation.kind + SparseConstantPropagation.transfer + +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