feat(riscv): basic datalayout support

Test/Passes/InstructionSelection/RISCV64/getelementptr.mlir

@@ -1,7 +1,7 @@
 // RUN: veir-opt %s -p=isel-riscv64 | filecheck %s
 
 // Single-dynamic-index `llvm.getelementptr` lowers to `ptr + idx * scale`,
-// where `scale` is the byte size of the element type.
+// where `scale` is the allocation size (ABI stride) of the element type.
 
 "builtin.module"() ({
     "func.func"()  <{function_type = (!llvm.ptr, i64) -> (), sym_name = "foo"}> ({
@@ -27,6 +27,13 @@
         // CHECK-NEXT: %{{.*}} = "riscv.sh2add"(%{{.*}}, %{{.*}}) : (!riscv.reg, !riscv.reg) -> !riscv.reg
         // CHECK-NEXT: %{{.*}} = "builtin.unrealized_conversion_cast"(%{{.*}}) : (!riscv.reg) -> !llvm.ptr
 
+        // i24 has store size 3 but RV64 allocation stride 4: also riscv.sh2add
+        %g24 = "llvm.getelementptr"(%p, %i) <{elem_type = i24, rawConstantIndices = array<i32: -2147483648>}> : (!llvm.ptr, i64) -> !llvm.ptr
+        // CHECK-NEXT: %{{.*}} = "builtin.unrealized_conversion_cast"(%{{.*}}) : (!llvm.ptr) -> !riscv.reg
+        // CHECK-NEXT: %{{.*}} = "builtin.unrealized_conversion_cast"(%{{.*}}) : (i64) -> !riscv.reg
+        // CHECK-NEXT: %{{.*}} = "riscv.sh2add"(%{{.*}}, %{{.*}}) : (!riscv.reg, !riscv.reg) -> !riscv.reg
+        // CHECK-NEXT: %{{.*}} = "builtin.unrealized_conversion_cast"(%{{.*}}) : (!riscv.reg) -> !llvm.ptr
+
         // scale 8 (i64): (idx << 3) + ptr -> riscv.sh3add
         %g8 = "llvm.getelementptr"(%p, %i) <{elem_type = i64, rawConstantIndices = array<i32: -2147483648>}> : (!llvm.ptr, i64) -> !llvm.ptr
         // CHECK-NEXT: %{{.*}} = "builtin.unrealized_conversion_cast"(%{{.*}}) : (!llvm.ptr) -> !riscv.reg
@@ -72,6 +79,7 @@
         "test.test"(%g1) : (!llvm.ptr) -> ()
         "test.test"(%g2) : (!llvm.ptr) -> ()
         "test.test"(%g4) : (!llvm.ptr) -> ()
+        "test.test"(%g24) : (!llvm.ptr) -> ()
         "test.test"(%g8) : (!llvm.ptr) -> ()
         "test.test"(%garr) : (!llvm.ptr) -> ()
         "test.test"(%g16) : (!llvm.ptr) -> ()

Test/Passes/InstructionSelection/RISCV64/load_store_offset.mlir

@@ -32,6 +32,20 @@
         "func.return"() : () -> ()
     }) : () -> ()
 
+    // Folding uses allocation stride, not store size: i33 has store size 5 but
+    // RV64 allocation stride 8, so index 2 gives byte offset 16.
+    "func.func"()  <{function_type = (!llvm.ptr) -> (), sym_name = "fold_i33_stride"}> ({
+    ^bb0(%p: !llvm.ptr):
+        %i = "llvm.mlir.constant"() <{value = 2 : i64}> : () -> i64
+        %g = "llvm.getelementptr"(%p, %i) <{elem_type = i33, rawConstantIndices = array<i32: -2147483648>}> : (!llvm.ptr, i64) -> !llvm.ptr
+        %v = "llvm.load"(%g) : (!llvm.ptr) -> i32
+        // CHECK:      {{.*}} = "builtin.unrealized_conversion_cast"({{.*}}) : (!llvm.ptr) -> !riscv.reg
+        // CHECK-NEXT: {{.*}} = "riscv.lw"({{.*}}) <{"value" = 16 : i64}> : (!riscv.reg) -> !riscv.reg
+        // CHECK-NEXT: {{.*}} = "builtin.unrealized_conversion_cast"({{.*}}) : (!riscv.reg) -> i32
+        "test.test"(%v) : (i32) -> ()
+        "func.return"() : () -> ()
+    }) : () -> ()
+
     // Stores fold the same way: i8 element, index 7 -> byte offset 7.
     "func.func"()  <{function_type = (!llvm.ptr, i8) -> (), sym_name = "fold_sb"}> ({
     ^bb0(%p: !llvm.ptr, %x: i8):

Veir/DataLayout/RISCV64.lean

@@ -0,0 +1,73 @@
+module
+
+public import Veir.Interfaces.DataLayoutInterfaces
+
+/-!
+# RV64 Data Layout
+
+The fixed data layout used by the RV64 backend.
+-/
+
+namespace Veir.DataLayout
+
+/-- Round `size` up to a positive byte alignment. -/
+private def alignTo (size alignment : Nat) : Nat :=
+  if alignment = 0 then size
+  else ((size + alignment - 1) / alignment) * alignment
+
+/-- The smallest power of two greater than or equal to `n` (and `1` for `0`). -/
+private def powerOfTwoCeil (n : Nat) : Nat :=
+  if n ≤ 1 then 1 else 2 ^ (Nat.log2 (n - 1) + 1)
+
+/--
+  RV64 integer ABI alignment, derived from LLVM's
+  `i1:8-i8:8-i16:16-i32:32-i64:64-i128:128` layout entries. As in LLVM and
+  MLIR, an unlisted width uses the next larger entry, or the largest entry when
+  no larger one exists.
+-/
+private def rv64IntegerAlignment (bitwidth : Nat) : Nat :=
+  if bitwidth ≤ 8 then 1
+  else if bitwidth ≤ 16 then 2
+  else if bitwidth ≤ 32 then 4
+  else if bitwidth ≤ 64 then 8
+  else 16
+
+private def scalarInfo (size alignment : Nat) : DataLayoutTypeInfo :=
+  { size
+    abiAlignment := alignment
+    preferredAlignment := alignment }
+
+/-- Layout facts for the LLVM-compatible fixed-size types supported by VeIR. -/
+private def queryRISCV64 (type : Attribute) : Option DataLayoutTypeInfo :=
+  match type with
+  | .integerType { bitwidth } | .byteType { bitwidth } =>
+      if bitwidth = 0 then none
+      else
+        let size := (bitwidth + 7) / 8
+        some (scalarInfo size (rv64IntegerAlignment bitwidth))
+  | .floatType { bitwidth } =>
+      if bitwidth = 0 then none
+      else
+        let size := (bitwidth + 7) / 8
+        some (scalarInfo size (powerOfTwoCeil size))
+  | .llvmPointerType _ =>
+      some (scalarInfo 8 8)
+  | .llvmArrayType { size, type } => do
+      let element ← queryRISCV64 type
+      let elementAllocSize := alignTo element.size element.abiAlignment
+      let arraySize := elementAllocSize * size
+      some
+        { size := arraySize
+          abiAlignment := element.abiAlignment
+          preferredAlignment := element.preferredAlignment }
+  | _ => none
+
+/--
+  The standard RV64 data layout:
+  `e-m:e-p:64:64-i64:64-i128:128-n32:64-S128`, together with LLVM's
+  default primitive entries.
+-/
+public def riscv64 : DataLayout :=
+  { query := queryRISCV64 }
+
+end Veir.DataLayout

Veir/Interfaces.lean

@@ -1,5 +1,6 @@
 module
 
+import Veir.Interfaces.DataLayoutInterfaces
 import Veir.Interfaces.FoldInterfaces
 import Veir.Interfaces.FunctionInterfaces
 import Veir.Interfaces.RegionKindInterfaces

Veir/Interfaces/DataLayoutInterfaces.lean

@@ -0,0 +1,66 @@
+module
+
+public import Veir.IR.Attribute
+
+/-!
+# DataLayoutInterface
+
+Target data layouts answer physical representation queries for IR types.  The
+interface deliberately distinguishes the byte size of a type, its ABI and
+preferred alignments, and its allocation size (the stride between consecutive
+objects).  In particular, an odd-width integer can have a three-byte type size
+but a four-byte allocation size.
+-/
+
+namespace Veir
+
+public section
+
+/-- The fixed-size layout facts for one type, all expressed in bytes. -/
+structure DataLayoutTypeInfo where
+  size : Nat
+  abiAlignment : Nat
+  preferredAlignment : Nat
+deriving Inhabited, Repr, DecidableEq
+
+/--
+  A target data layout. Unsupported or unsized types return `none`.
+
+  Keeping the query behind an object lets passes depend on the interface rather
+  than on how layout information is obtained (currently fixed RV64 values,
+  eventually perhaps parsed DLTI entries).
+-/
+structure DataLayout where
+  query : Attribute → Option DataLayoutTypeInfo
+
+namespace DataLayout
+
+/-- Round `size` up to a positive byte alignment. -/
+private def alignTo (size alignment : Nat) : Nat :=
+  if alignment = 0 then size
+  else ((size + alignment - 1) / alignment) * alignment
+
+/-- Return the size of `type` in bytes, including padding internal to the type. -/
+def getTypeSize (layout : DataLayout) (type : Attribute) : Option Nat :=
+  (layout.query type).map (·.size)
+
+/-- Return the minimum ABI-required alignment of `type`, in bytes. -/
+def getTypeABIAlignment (layout : DataLayout) (type : Attribute) : Option Nat :=
+  (layout.query type).map (·.abiAlignment)
+
+/-- Return the preferred alignment of `type`, in bytes. -/
+def getTypePreferredAlignment (layout : DataLayout) (type : Attribute) : Option Nat :=
+  (layout.query type).map (·.preferredAlignment)
+
+/--
+  Return the allocation size of `type`, in bytes: the stride between consecutive
+  objects, including tail padding required by the ABI alignment.
+-/
+def getTypeAllocSize (layout : DataLayout) (type : Attribute) : Option Nat :=
+  (layout.query type).map fun info => alignTo info.size info.abiAlignment
+
+end DataLayout
+
+end
+
+end Veir

Veir/Passes/InstructionSelection/RISCV64.lean

@@ -1,6 +1,7 @@
 module
 
 public import Veir.Pass
+import Veir.DataLayout.RISCV64
 import Veir.Passes.Matching.LLVM.Basic
 import Veir.Passes.InstructionSelection.Common
 
@@ -860,7 +861,7 @@ def selectAddrRegImm (ptr : ValuePtr) (ctx : IRContext OpCode) : ValuePtr × Int
     let .integerType itype := (idx.getType! ctx).val | none
     guard (itype.bitwidth = 64)
     let c ← matchConstantIntVal idx ctx
-    let scale ← Attribute.sizeOfType properties.elem_type.val
+    let scale ← DataLayout.riscv64.getTypeAllocSize properties.elem_type.val
     let offset := c.value * (scale : Int)
     guard (-2048 ≤ offset ∧ offset ≤ 2047)
     return (base, offset)
@@ -940,7 +941,7 @@ def store (rewriter : PatternRewriter OpCode) (op : OperationPtr)
 
 /--
   Lower a single-dynamic-index `llvm.getelementptr` computing `ptr + idx * scale`,
-  where `scale` is the byte size of the element type.
+  where `scale` is the allocation size (ABI stride) of the element type.
 -/
 def getelementptr_local (ctx : WfIRContext OpCode) (op : OperationPtr) :
     Option (WfIRContext OpCode × Option (Array OperationPtr × Array ValuePtr)) := do
@@ -950,7 +951,8 @@ def getelementptr_local (ctx : WfIRContext OpCode) (op : OperationPtr) :
   /- The index must be `i64`. -/
   let .integerType itype := (idx.getType! ctx.raw).val | return (ctx, none)
   if itype.bitwidth ≠ 64 then return (ctx, none)
-  let some scale := Attribute.sizeOfType properties.elem_type.val | return (ctx, none)
+  let some scale := DataLayout.riscv64.getTypeAllocSize properties.elem_type.val
+    | return (ctx, none)
   let type := ((op.getResult 0).get! ctx.raw).type
   let (ctx, pcastOp) ← WfRewriter.createOp! ctx Builtin.unrealized_conversion_cast #[RegisterType.mk] #[ptr]
       #[] #[] () none
@@ -1010,7 +1012,7 @@ def getelementptr_local (ctx : WfIRContext OpCode) (op : OperationPtr) :
 
 /--
   Lower a single-dynamic-index `llvm.getelementptr` computing `ptr + idx * scale`,
-  where `scale` is the byte size of the element type.
+  where `scale` is the allocation size (ABI stride) of the element type.
 -/
 def getelementptr (rewriter : PatternRewriter OpCode) (op : OperationPtr)
     (opInBounds : op.InBounds rewriter.ctx.raw) : Option (PatternRewriter OpCode) :=