feat(legalizer): add GlobalISel-like legalization framework

Test/Passes/Legalization/add.mlir

@@ -1,4 +1,4 @@
-// RUN: veir-opt -p=legalize %s | filecheck %s
+// RUN: veir-opt -p=legalize-riscv64 %s | filecheck %s
 
 "builtin.module"() ({
   "func.func"() <{sym_name = "main", function_type = () -> i32}> ({

Veir/Passes/Legalization.lean

@@ -2,6 +2,7 @@ module
 
 public import Veir.Pass
 import Veir.Passes.Matching
+public import Veir.PatternRewriter.Basic
 
 /-!
   This pass legalizes LLVM operations to prepare for instruction selection.
@@ -10,6 +11,137 @@ import Veir.Passes.Matching
 
 namespace Veir
 
+public section
+
+/--
+This wip implements an equivalent of LLVM's GISel framework for building legalization passes.
+https://github.com/llvm/llvm-project/blob/main/llvm/include/llvm/CodeGen/GlobalISel/LegalizerInfo.h
+-/
+
+-- TODO: Most of these are unimplemented.
+inductive LegalizeAction where
+  | legal
+  | narrowScalar
+  | widenScalar
+  | fewerElements
+  | moreElements
+  | bitcast
+  | lower
+  | libcall
+  | custom
+  | unsupported
+  | notFound
+  deriving Repr, BEq
+
+/--
+  The LegalityQuery object bundles together all the information that's needed to decide whether a
+  given operation is legal or not.
+-/
+structure LegalityQuery where
+  opcode : OpCode
+  -- sizes[0] is always the bitwidth of the result type.
+  -- sizes[1..] for when the legality also depends on operand types.
+  sizes : Array Nat
+
+/--
+  The result of a query. It either indicates a final answer of Legal or Unsupported or describes an
+  action that must be taken to make an operation more legal.
+-/
+structure LegalizeActionStep where
+  action : LegalizeAction
+  typeIndex : Nat := 0
+  newBw : Nat := 0
+
+abbrev LegalityPredicate := LegalityQuery → Bool
+abbrev LegalizeMutation := LegalityQuery → Nat × Nat
+
+namespace LegalityPredicates
+
+-- NOTE: These are not yet well aligned with LLVM's legality predicates.
+
+def sizeInSet (typeIndex : Nat) (sizes : Array Nat) : LegalityPredicate :=
+  fun q => sizes.contains q.sizes[typeIndex]!
+
+def sizeNotPow2 (typeIndex : Nat) : LegalityPredicate :=
+  fun q => !Nat.isPowerOfTwo q.sizes[typeIndex]!
+
+def scalarNarrowerThan (typeIndex bw : Nat) : LegalityPredicate :=
+  fun q => q.sizes[typeIndex]! < bw
+
+end LegalityPredicates
+
+namespace LegalityMutations
+
+def widenScalarToNextPow2 (typeIndex : Nat) : LegalizeMutation :=
+  fun q => (typeIndex, Nat.nextPowerOfTwo q.sizes[typeIndex]!)
+
+def clampScalar (typeIndex minBw maxBw : Nat) : LegalizeMutation :=
+  fun q => (typeIndex, min maxBw (max minBw q.sizes[typeIndex]!))
+
+end LegalityMutations
+
+/--
+  A single rule in a legalizer info ruleset.
+  The specified action is chosen when the predicate is true. Where appropriate for the action
+  (e.g. for WidenScalar) the new type is selected using the given mutator.
+-/
+structure LegalizeRule where
+  predicate : LegalityPredicate
+  action : LegalizeAction
+  mutation : LegalizeMutation
+
+abbrev LegalizeRuleSet := Array LegalizeRule
+
+private def LegalizeRule.apply (rule : LegalizeRule) (q : LegalityQuery) : Option LegalizeActionStep :=
+  if rule.predicate q then
+    let (typeIndex, newBw) := rule.mutation q
+    some { action := rule.action, typeIndex, newBw }
+  else
+    none
+
+def legalFor (sizes : Array Nat) : LegalizeRule :=
+  { predicate := LegalityPredicates.sizeInSet 0 sizes
+    action := .legal
+    mutation := fun _ => (0, 0) }
+
+def customFor (sizes : Array Nat) : LegalizeRule :=
+  { predicate := LegalityPredicates.sizeInSet 0 sizes
+    action := .custom
+    mutation := fun _ => (0, 0) }
+
+def widenScalarToNextPow2 (typeIndex : Nat) : LegalizeRule :=
+  { predicate := LegalityPredicates.sizeNotPow2 typeIndex
+    action := .widenScalar
+    mutation := LegalityMutations.widenScalarToNextPow2 typeIndex }
+
+def clampScalar (typeIndex minBw maxBw : Nat) : LegalizeRule :=
+  { predicate := LegalityPredicates.scalarNarrowerThan typeIndex maxBw
+    action := .widenScalar
+    mutation := LegalityMutations.clampScalar typeIndex minBw maxBw }
+
+structure LegalizerInfo where
+  ruleSets : Std.HashMap OpCode LegalizeRuleSet := ∅
+  legalizeCustom : LocalRewritePattern OpCode := fun ctx _ => some (ctx, none)
+
+namespace LegalizerInfo
+
+def defineRuleSet
+    (info : LegalizerInfo)
+    (ops : Array OpCode)
+    (rules : LegalizeRuleSet) : LegalizerInfo :=
+  { info with
+    ruleSets := ops.foldl (fun m op => m.insert op rules) info.ruleSets }
+
+def getAction (info : LegalizerInfo) (q : LegalityQuery) : LegalizeActionStep :=
+  match info.ruleSets.get? q.opcode with
+  | none => { action := .notFound }
+  | some ruleset =>
+    (ruleset.findSome? (·.apply q)).getD { action := .notFound }
+
+end LegalizerInfo
+
+end
+
 /--
   Sigma type for an operation plus its properties.
 -/
@@ -56,34 +188,40 @@ def widenSimpleBinaryIntOp (ctx : WfIRContext OpCode) (op : OperationPtr) (newBw
   let expandOp := expandIntegerExtOp extType
   convertBinaryOp ctx op (IntegerType.mk newBw) expandOp expandOp (newOp.getD oldOp) ⟨.llvm .trunc, .mk false false⟩
 
-/--
-  Widen the operands and result type of an LLVM operation.
--/
--- TODO incomplete
-def widenOperations (ctx : WfIRContext OpCode) (op : OperationPtr) :
+def queryOf (ctx : WfIRContext OpCode) (op : OperationPtr) : Option LegalityQuery := do
+  let opcode := op.getOpType! ctx.raw
+  let resultType : TypeAttr ← (op.getResultTypes! ctx.raw)[0]?
+  -- Skips over types without bitwidth.
+  let bw ← Attribute.bitwidthOfType resultType
+  some { opcode, sizes := #[bw] }
+
+def widenScalar (ctx : WfIRContext OpCode) (op : OperationPtr) (newBw : Nat) :
     Option (WfIRContext OpCode × Option (Array OperationPtr × Array ValuePtr)) := do
   match op.getOpType! ctx.raw with
   | .llvm .add =>
-    widenSimpleBinaryIntOp ctx op 64 .any (some ⟨.llvm .add, .mk false false⟩)
+    widenSimpleBinaryIntOp ctx op newBw .any (some ⟨.llvm .add, .mk false false⟩)
   | .llvm .sub =>
-    widenSimpleBinaryIntOp ctx op 64 .any (some ⟨.llvm .sub, .mk false false⟩)
+    widenSimpleBinaryIntOp ctx op newBw .any (some ⟨.llvm .sub, .mk false false⟩)
   | .llvm .mul =>
-    widenSimpleBinaryIntOp ctx op 64 .any (some ⟨.llvm .mul, .mk false false⟩)
+    widenSimpleBinaryIntOp ctx op newBw .any (some ⟨.llvm .mul, .mk false false⟩)
   | .llvm .and =>
-    widenSimpleBinaryIntOp ctx op 64 .any (some ⟨.llvm .and, ()⟩)
+    widenSimpleBinaryIntOp ctx op newBw .any (some ⟨.llvm .and, ()⟩)
   | .llvm .xor =>
-    widenSimpleBinaryIntOp ctx op 64 .any (some ⟨.llvm .xor, ()⟩)
+    widenSimpleBinaryIntOp ctx op newBw .any (some ⟨.llvm .xor, ()⟩)
   | .llvm .or =>
-    widenSimpleBinaryIntOp ctx op 64 .any (some ⟨.llvm .or, .mk false⟩)
+    widenSimpleBinaryIntOp ctx op newBw .any (some ⟨.llvm .or, .mk false⟩)
   | _ => return (ctx, none)
 
-def LegalizePass.impl (ctx : WfIRContext OpCode) (op : OperationPtr) (_ : op.InBounds ctx.raw) :
-    ExceptT String IO (WfIRContext OpCode) := do
-  match RewritePattern.applyInContext (RewritePattern.GreedyRewritePattern #[.fromLocalRewrite widenOperations]) ctx with
-  | none => throw "Error while applying legalization"
-  | some ctx => pure ctx
+public def legalizeInstrStep (info : LegalizerInfo) (ctx : WfIRContext OpCode) (op : OperationPtr) :
+    Option (WfIRContext OpCode × Option (Array OperationPtr × Array ValuePtr)) :=
+  match queryOf ctx op with
+  | none => some (ctx, none)
+  | some q =>
+    let step := info.getAction q
+    match step.action with
+    | .legal => some (ctx, none)
+    | .widenScalar => widenScalar ctx op step.newBw
+    | .custom => info.legalizeCustom ctx op
+    | _ => some (ctx, none)
 
-public def LegalizePass : Pass OpCode :=
-  { name := "legalize"
-    description := "Legalize types."
-    run := fun _ => LegalizePass.impl }
+end Veir

Veir/Passes/RISCV64Legalization.lean

@@ -0,0 +1,32 @@
+module
+
+public import Veir.Pass
+import Veir.Passes.Legalization
+
+namespace Veir
+
+def riscv64LegalizerInfo : LegalizerInfo :=
+  let info : LegalizerInfo := {}
+  let info := info.defineRuleSet #[.llvm .add, .llvm .sub] #[
+    legalFor #[64],
+    -- customFor #[32] (not yet implemented)
+    widenScalarToNextPow2 0,
+    clampScalar 0 64 64,
+  ]
+  info
+
+def LegalizeRISCV64Pass.impl (ctx : WfIRContext OpCode) (op : OperationPtr) (_ : op.InBounds ctx.raw) :
+    ExceptT String IO (WfIRContext OpCode) := do
+  let pattern := RewritePattern.GreedyRewritePattern #[
+    .fromLocalRewrite (legalizeInstrStep riscv64LegalizerInfo)
+  ]
+  match RewritePattern.applyInContext pattern ctx with
+  | none => throw "Error while applying RISC-V legalization"
+  | some ctx => pure ctx
+
+public def LegalizeRISCV64Pass : Pass OpCode :=
+  { name := "legalize-riscv64"
+    description := "Legalize types for RISC-V 64."
+    run := fun _ => LegalizeRISCV64Pass.impl }
+
+end Veir

VeirOpt.lean

@@ -15,7 +15,7 @@ import Veir.Passes.RISCVCombines.Combine
 import Veir.Passes.ModArithToArith
 import Veir.Passes.ArithToLLVM
 import Veir.Passes.Canonicalize
-import Veir.Passes.Legalization
+import Veir.Passes.RISCV64Legalization
 
 open Veir.Parser
 open Veir.Parser.ParserError
@@ -39,7 +39,7 @@ def availablePasses : Std.HashMap String (Pass OpCode) :=
      ModArithToArithPass,
      ArithToLLVMPass,
      CanonicalizePass,
-     LegalizePass ] : List (Pass OpCode)).foldl
+     LegalizeRISCV64Pass ] : List (Pass OpCode)).foldl
     (fun m pass => m.insert pass.name pass)
     (Std.HashMap.emptyWithCapacity 16)
 
@@ -56,7 +56,7 @@ def passGroups : Std.HashMap String String :=
     |>.insert "mod-arith-pow2-width"
         "mod-arith-to-arith{barrett pow2-width},cse,coerce-mod-arith-function-boundaries{pow2-width},reconcile-cast,canonicalize,cse,dce"
     |>.insert "riscv"
-        "legalize,isel-sdag-riscv64,isel-br-riscv64,isel-riscv64,coerce-function-boundaries-to-riscv-reg,reconcile-cast,riscv-combine,dce"
+        "legalize-riscv64,isel-sdag-riscv64,isel-br-riscv64,isel-riscv64,coerce-function-boundaries-to-riscv-reg,reconcile-cast,riscv-combine,dce"
 
 /--
   A human-readable description of every pass group and the passes it expands to,