From 11015e161a108ce3436021361b2e0d6510de7b4c Mon Sep 17 00:00:00 2001 From: Naveen Seth Hanig Date: Thu, 13 Aug 2026 01:06:31 +0200 Subject: [PATCH] feat(legalizer): add GlobalISel-like legalization framework This adds the basic framework for building legalization passes in the same way as LLVM's GlobalISel. This currently just defines the rule actions needed for G_ADD/G_SUB in the 64-bit RISC-V legalization pass. (Very WIP. Top-level documentation is mostly just copied from LLVM.) --- Test/Passes/Legalization/add.mlir | 2 +- Veir/Passes/Legalization.lean | 178 ++++++++++++++++++++++++--- Veir/Passes/RISCV64Legalization.lean | 32 +++++ VeirOpt.lean | 6 +- 4 files changed, 194 insertions(+), 24 deletions(-) create mode 100644 Veir/Passes/RISCV64Legalization.lean diff --git a/Test/Passes/Legalization/add.mlir b/Test/Passes/Legalization/add.mlir index ad52fce96..8209aea18 100644 --- a/Test/Passes/Legalization/add.mlir +++ b/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}> ({ diff --git a/Veir/Passes/Legalization.lean b/Veir/Passes/Legalization.lean index 429d6deee..81cf04efc 100644 --- a/Veir/Passes/Legalization.lean +++ b/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 diff --git a/Veir/Passes/RISCV64Legalization.lean b/Veir/Passes/RISCV64Legalization.lean new file mode 100644 index 000000000..054ec1d25 --- /dev/null +++ b/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 diff --git a/VeirOpt.lean b/VeirOpt.lean index d3acdc1f4..3a39b3841 100644 --- a/VeirOpt.lean +++ b/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,