feat(dataflow): known bits analysis

UnitTest.lean

@@ -13,6 +13,8 @@ import UnitTest.DataFlowFramework.Dominance
 import UnitTest.DataFlowFramework.DeadCodeAnalysis
 import UnitTest.DataFlowFramework.EntryState
 import UnitTest.DataFlowFramework.ModArithRangeAnalysis
+import UnitTest.DataFlowFramework.KnownBitsAnalysis
+import UnitTest.ModArithRangeAnalysis
 import UnitTest.ConstantValue
 import UnitTest.Evaluate
 import UnitTest.FoldDecision

UnitTest/DataFlowFramework/KnownBitsAnalysis.lean

@@ -0,0 +1,137 @@
+import UnitTest.DataFlowFramework.Helpers
+
+import Veir.Analysis.DataFlow.KnownBitsAnalysis
+
+open Veir
+
+namespace KnownBitsDataflow
+
+/-- Expected masks for one named SSA value. -/
+private structure ExpectedKnownBits where
+  name : String
+  bitwidth : Nat
+  zero : Nat
+  one : Nat
+
+private def knownBitsToString : KnownBitsLattice → String
+  | .bottom => "bottom"
+  | .top => "top"
+  | .known bits => s!"i{bits.bitwidth}(zero={bits.zero.toNat}, one={bits.one.toNat})"
+
+private def compareKnownBits
+    (dfCtx : DataFlowContext)
+    (recovered : RecoveredNames)
+    (expected : Array ExpectedKnownBits) : MismatchReport := Id.run do
+  let mut report := #[]
+  for e in expected do
+    let some value := recovered.values[e.name]?
+      | report := report.push s!"known bits {e.name}: missing SSA value"
+        continue
+    let observed : KnownBitsLattice := SparseFact.getElement .knownBits value dfCtx
+    let expectedValue := KnownBitsLattice.known
+      { bitwidth := e.bitwidth
+        zero := BitVec.ofNat e.bitwidth e.zero
+        one := BitVec.ofNat e.bitwidth e.one }
+    if observed ≠ expectedValue then
+      report := report.push <|
+        s!"known bits {e.name}: expected {knownBitsToString expectedValue}, " ++
+        s!"observed {knownBitsToString observed}"
+  report
+
+private def run (mlir : String) (expected : Array ExpectedKnownBits) : String :=
+  runWithAnalyses mlir #[Veir.KnownBitsAnalysis] fun top dfCtx parserState =>
+    match recoverNames top parserState.ctx mlir with
+    | .error err => #[err]
+    | .ok recovered => compareKnownBits dfCtx recovered expected
+
+/-- Arith constants and bitwise operations preserve partial known-bit information. -/
+def runArithKnownBitsExample : String :=
+  let mlir := r#""builtin.module"() ({
+^bb0:
+  "func.func"() <{function_type = (i8) -> (), sym_name = "known_bits_arith"}> ({
+  ^entry(%x : i8):
+    %c240 = "arith.constant"() <{value = 240 : i8}> : () -> i8
+    %c3 = "arith.constant"() <{value = 3 : i8}> : () -> i8
+    %c5 = "arith.constant"() <{value = 5 : i8}> : () -> i8
+    %sum = "arith.addi"(%c3, %c5) : (i8, i8) -> i8
+    %anded = "arith.andi"(%x, %c240) : (i8, i8) -> i8
+    %ored = "arith.ori"(%anded, %c3) : (i8, i8) -> i8
+    %xored = "arith.xori"(%ored, %c5) : (i8, i8) -> i8
+    "func.return"() : () -> ()
+  }) : () -> ()
+}) : () -> ()"#
+  let expected :=
+    #[ { name := "x",     bitwidth := 8, zero := 0,   one := 0 }
+     , { name := "c240",  bitwidth := 8, zero := 15,  one := 240 }
+     , { name := "c3",    bitwidth := 8, zero := 252, one := 3 }
+     , { name := "c5",    bitwidth := 8, zero := 250, one := 5 }
+     , { name := "sum",   bitwidth := 8, zero := 247, one := 8 }
+     , { name := "anded", bitwidth := 8, zero := 15,  one := 0 }
+     , { name := "ored",  bitwidth := 8, zero := 12,  one := 3 }
+     , { name := "xored", bitwidth := 8, zero := 9,   one := 6 }
+     ]
+  run mlir expected
+
+/-- LLVM spellings and variadic Comb operations use the same transfer functions. -/
+def runLLVMAndCombKnownBitsExample : String :=
+  let mlir := r#""builtin.module"() ({
+^bb0:
+  "func.func"() <{function_type = (i8) -> (), sym_name = "known_bits_dialects"}> ({
+  ^entry(%x : i8):
+    %lc240 = "llvm.mlir.constant"() <{value = 240 : i8}> : () -> i8
+    %lc3 = "llvm.mlir.constant"() <{value = 3 : i8}> : () -> i8
+    %lc5 = "llvm.mlir.constant"() <{value = 5 : i8}> : () -> i8
+    %land = "llvm.and"(%x, %lc240) : (i8, i8) -> i8
+    %lor = "llvm.or"(%land, %lc3) : (i8, i8) -> i8
+    %lxor = "llvm.xor"(%lor, %lc5) : (i8, i8) -> i8
+    %hc240 = "hw.constant"() <{value = 240 : i8}> : () -> i8
+    %hc15 = "hw.constant"() <{value = 15 : i8}> : () -> i8
+    %hc3 = "hw.constant"() <{value = 3 : i8}> : () -> i8
+    %cand = "comb.and"(%hc240, %hc15, %hc3) : (i8, i8, i8) -> i8
+    %cor = "comb.or"(%hc240, %hc15, %hc3) : (i8, i8, i8) -> i8
+    %cxor = "comb.xor"(%hc240, %hc15, %hc3) : (i8, i8, i8) -> i8
+    "func.return"() : () -> ()
+  }) : () -> ()
+}) : () -> ()"#
+  let expected :=
+    #[ { name := "land", bitwidth := 8, zero := 15,  one := 0 }
+     , { name := "lor",  bitwidth := 8, zero := 12,  one := 3 }
+     , { name := "lxor", bitwidth := 8, zero := 9,   one := 6 }
+     , { name := "cand", bitwidth := 8, zero := 255, one := 0 }
+     , { name := "cor",  bitwidth := 8, zero := 0,   one := 255 }
+     , { name := "cxor", bitwidth := 8, zero := 3,   one := 252 }
+     ]
+  run mlir expected
+
+/-- Joining exact values retains only the bits on which both values agree. -/
+def testKnownBitsJoin : String :=
+  let joined :=
+    KnownBitsLattice.join
+      (.constant 8 165)
+      (.constant 8 167)
+  let expected : KnownBitsLattice :=
+    .known
+      { bitwidth := 8
+        zero := BitVec.ofNat 8 88
+        one := BitVec.ofNat 8 165 }
+  if joined = expected then "ok" else s!"unexpected join: {knownBitsToString joined}"
+
+/--
+info: "ok"
+-/
+#guard_msgs in
+#eval! runArithKnownBitsExample
+
+/--
+info: "ok"
+-/
+#guard_msgs in
+#eval! runLLVMAndCombKnownBitsExample
+
+/--
+info: "ok"
+-/
+#guard_msgs in
+#eval! testKnownBitsJoin
+
+end KnownBitsDataflow

Veir/Analysis.lean

@@ -4,3 +4,4 @@ public import Veir.Analysis.DataFlowFramework
 public import Veir.Analysis.DataFlow.DominanceAnalysis
 public import Veir.Analysis.DataFlow.DeadCodeAnalysis
 public import Veir.Analysis.DataFlow.ModArithRangeAnalysis
+public import Veir.Analysis.DataFlow.KnownBitsAnalysis

Veir/Analysis/DataFlow/Domains/KnownBitsDomain.lean

@@ -0,0 +1,165 @@
+module
+
+public import Veir.Analysis.DataFlow.Domains.AbstractDomain
+public import Veir.RuntimeValue
+import Veir.Meta.BVDecide
+
+public section
+
+namespace Veir
+
+/-!
+# Known-bits domain
+
+This file defines the abstract value used by known-bits analysis. As in LLVM, a
+known-bits value stores two masks: `zero` marks bits known to be zero and `one`
+marks bits known to be one. Bits absent from both masks are unknown.
+-/
+
+/-- Two masks describing the known zero and known one bits of a fixed-width integer. -/
+structure KnownBits where
+  bitwidth : Nat
+  zero : BitVec bitwidth
+  one : BitVec bitwidth
+deriving DecidableEq, Repr
+
+namespace KnownBits
+
+/-- No bits are known for an integer of the given width. -/
+def unknown (bitwidth : Nat) : KnownBits :=
+  { bitwidth, zero := 0, one := 0 }
+
+/-- Every bit of a concrete integer is known. -/
+def constant (bitwidth : Nat) (value : Int) : KnownBits :=
+  let bits := BitVec.ofInt bitwidth value
+  { bitwidth, zero := ~~~bits, one := bits }
+
+/-- The zero and one masks do not make contradictory claims. -/
+def isValid (bits : KnownBits) : Bool :=
+  bits.zero &&& bits.one == 0
+
+/-- Keep only facts known on both incoming control-flow paths. -/
+def join? (lhs rhs : KnownBits) : Option KnownBits :=
+  if h : lhs.bitwidth = rhs.bitwidth then
+    let rhsZero := h ▸ rhs.zero
+    let rhsOne := h ▸ rhs.one
+    some
+      { bitwidth := lhs.bitwidth
+        zero := lhs.zero &&& rhsZero
+        one := lhs.one &&& rhsOne }
+  else
+    none
+
+end KnownBits
+
+/--
+Sparse lattice for known bits. `bottom` is an uninitialized sparse value, `known`
+contains width-aware masks, and `top` is used when even the integer width is unavailable.
+-/
+inductive KnownBitsLattice where
+  | bottom
+  | known (bits : KnownBits)
+  | top
+deriving DecidableEq, Repr
+
+namespace KnownBitsLattice
+
+instance : Bot KnownBitsLattice where
+  bot := .bottom
+
+instance : Top KnownBitsLattice where
+  top := .top
+
+/-- No bit facts are known, but the integer width is known. -/
+def unknown (bitwidth : Nat) : KnownBitsLattice :=
+  .known (KnownBits.unknown bitwidth)
+
+/-- An exact fixed-width integer value. -/
+def constant (bitwidth : Nat) (value : Int) : KnownBitsLattice :=
+  .known (KnownBits.constant bitwidth value)
+
+/-- The concrete runtime values represented by a known-bits lattice element. -/
+@[expose] def γ : KnownBitsLattice → Set RuntimeValue
+  | .bottom => ⊥
+  | .top => ⊤
+  | .known bits => fun concrete =>
+      match concrete with
+      | .int bitwidth (.val value) =>
+        ∃ h : bitwidth = bits.bitwidth,
+          let value := value.cast h
+          value &&& bits.zero = 0 ∧ value &&& bits.one = bits.one
+      | _ => False
+
+@[simp] theorem not_mem_γ_bottom (value : RuntimeValue) : value ∉ γ .bottom := fun h => h.elim
+
+@[simp] theorem mem_γ_top (value : RuntimeValue) : value ∈ γ .top := trivial
+
+/-- Normalize membership in a known-bits value to masks at the concrete value's width. -/
+theorem mem_γ_known_masks_iff
+    {bits : KnownBits}
+    {bitwidth : Nat}
+    {value : BitVec bitwidth} :
+    RuntimeValue.int bitwidth (.val value) ∈ γ (.known bits) ↔
+      ∃ zero one : BitVec bitwidth,
+        bits = ⟨bitwidth, zero, one⟩ ∧
+        value &&& zero = 0 ∧ value &&& one = one := by
+  constructor
+  · rcases bits with ⟨bitsWidth, zero, one⟩
+    rintro ⟨hwidth, hzero, hone⟩
+    change bitwidth = bitsWidth at hwidth
+    subst bitsWidth
+    simp at hzero hone
+    exact ⟨zero, one, rfl, hzero, hone⟩
+  · rintro ⟨zero, one, rfl, hzero, hone⟩
+    exact ⟨rfl, hzero, hone⟩
+
+/-- Characterize known-bits membership as facts about each concrete bit. -/
+theorem mem_γ_known_iff
+    {bits : KnownBits}
+    {bitwidth : Nat}
+    {value : BitVec bitwidth} :
+    RuntimeValue.int bitwidth (.val value) ∈ γ (.known bits) ↔
+      ∃ zero one : BitVec bitwidth,
+        bits = ⟨bitwidth, zero, one⟩ ∧
+        (∀ i (hi : i < bitwidth), zero[i] = true → value[i] = false) ∧
+        (∀ i (hi : i < bitwidth), one[i] = true → value[i] = true) := by
+  rw [mem_γ_known_masks_iff]
+  constructor
+  · rintro ⟨zero, one, hbits, hzero, hone⟩
+    refine ⟨zero, one, hbits, ?_, ?_⟩
+    · intro i hi hzeroTrue
+      have hzeroBit := congrArg (fun value => value[i]) hzero
+      simp at hzeroBit
+      veir_bv_decide
+    · intro i hi honeTrue
+      have honeBit := congrArg (fun value => value[i]) hone
+      simp at honeBit
+      veir_bv_decide
+  · rintro ⟨zero, one, hbits, hzero, hone⟩
+    refine ⟨zero, one, hbits, ?_, ?_⟩
+    · ext i hi
+      have hzeroBit := hzero i hi
+      simp at hzeroBit ⊢
+      veir_bv_decide
+    · ext i hi
+      have honeBit := hone i hi
+      simp at honeBit ⊢
+      veir_bv_decide
+
+/-- Join facts arriving along different control-flow paths. -/
+def join : KnownBitsLattice → KnownBitsLattice → KnownBitsLattice
+  | .bottom, rhs => rhs
+  | lhs, .bottom => lhs
+  | .top, _ => .top
+  | _, .top => .top
+  | .known lhs, .known rhs =>
+      match lhs.join? rhs with
+      | some bits => .known bits
+      | none => .top
+
+instance : Join KnownBitsLattice where
+  join := join
+
+end KnownBitsLattice
+
+end Veir

Veir/Analysis/DataFlow/Facts.lean

@@ -2,6 +2,7 @@ module
 
 public import Veir.GlobalOpInfo
 public import Veir.Analysis.DataFlow.Domains.IntegerRangeDomain
+public import Veir.Analysis.DataFlow.Domains.KnownBitsDomain
 public import Veir.Analysis.DataFlow.Domains.LivenessDomain
 public import Veir.Rewriter.InsertPoint
 public import Veir.Analysis.DataFlow.Domains.ConstantDomain
@@ -75,6 +76,7 @@ inductive AnalysisKind where
   | sparseConstantPropagation
   | integerRange
   | modArithRange
+  | knownBits
 deriving BEq, Hashable, Repr, DecidableEq
 
 /--
@@ -89,6 +91,7 @@ inductive FactKind where
   | sparseConstant
   | integerRange
   | modArithRange
+  | knownBits
 deriving BEq, ReflBEq, LawfulBEq, Hashable, Repr, DecidableEq
 
 abbrev WorkItem := InsertPoint × AnalysisKind
@@ -131,6 +134,7 @@ The fact specific data stored for each fact kind.
   | .sparseConstant => SparsePayload AbstractConstant
   | .integerRange => SparsePayload IntegerRangeLattice
   | .modArithRange => SparsePayload IntegerRangeLattice
+  | .knownBits => SparsePayload KnownBitsLattice
 
 /--
 A dataflow fact stored by the framework.