[MLIR] Add index bitwidth to the DataLayout (llvm#85927)

When importing from LLVM IR the data layout of all pointer types
contains an index bitwidth that should be used for index computations.
This revision adds a getter to the DataLayout that provides access to
the already stored bitwidth. The function returns an optional since only
pointer-like types have an index bitwidth. Querying the bitwidth of a
non-pointer type returns std::nullopt.

The new function works for the built-in Index type and, using a type
interface, for the LLVMPointerType.

Author: gysit

mlir/docs/DataLayout.md

@@ -77,6 +77,7 @@ public:
   llvm::TypeSize getTypeSizeInBits(Type type) const;
   uint64_t getTypeABIAlignment(Type type) const;
   uint64_t getTypePreferredAlignment(Type type) const;
+  std::optional<uint64_t> getTypeIndexBitwidth(Type type) const;
 };
 ```
 
@@ -267,7 +268,8 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
 >} {}
 ```
 
-specifies that `index` has 32 bits. All other layout properties of `index` match
+specifies that `index` has 32 bits and index computations should be performed
+using 32-bit precision as well. All other layout properties of `index` match
 those of the integer type with the same bitwidth defined above.
 
 In absence of the corresponding entry, `index` is assumed to be a 64-bit

mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.td

@@ -123,7 +123,7 @@ def LLVMFunctionType : LLVMType<"LLVMFunction", "func"> {
 
 def LLVMPointerType : LLVMType<"LLVMPointer", "ptr", [
     DeclareTypeInterfaceMethods<DataLayoutTypeInterface, [
-      "areCompatible", "verifyEntries"]>]> {
+      "getIndexBitwidth", "areCompatible", "verifyEntries"]>]> {
   let summary = "LLVM pointer type";
   let description = [{
     The `!llvm.ptr` type is an LLVM pointer type. This type typically represents

mlir/include/mlir/Interfaces/DataLayoutInterfaces.h

@@ -57,6 +57,13 @@ uint64_t
 getDefaultPreferredAlignment(Type type, const DataLayout &dataLayout,
                              ArrayRef<DataLayoutEntryInterface> params);
 
+/// Default handler for the index bitwidth request. Computes the result for
+/// the built-in index type and dispatches to the DataLayoutTypeInterface for
+/// other types.
+std::optional<uint64_t>
+getDefaultIndexBitwidth(Type type, const DataLayout &dataLayout,
+                        ArrayRef<DataLayoutEntryInterface> params);
+
 /// Default handler for alloca memory space request. Dispatches to the
 /// DataLayoutInterface if specified, otherwise returns the default.
 Attribute getDefaultAllocaMemorySpace(DataLayoutEntryInterface entry);
@@ -180,6 +187,11 @@ class DataLayout {
   /// Returns the preferred of the given type in the current scope.
   uint64_t getTypePreferredAlignment(Type t) const;
 
+  /// Returns the bitwidth that should be used when performing index
+  /// computations for the given pointer-like type in the current scope. If the
+  /// type is not a pointer-like type, it returns std::nullopt.
+  std::optional<uint64_t> getTypeIndexBitwidth(Type t) const;
+
   /// Returns the memory space used for AllocaOps.
   Attribute getAllocaMemorySpace() const;
 
@@ -216,6 +228,7 @@ class DataLayout {
   mutable DenseMap<Type, llvm::TypeSize> bitsizes;
   mutable DenseMap<Type, uint64_t> abiAlignments;
   mutable DenseMap<Type, uint64_t> preferredAlignments;
+  mutable DenseMap<Type, std::optional<uint64_t>> indexBitwidths;
 
   /// Cache for alloca, global, and program memory spaces.
   mutable std::optional<Attribute> allocaMemorySpace;

mlir/include/mlir/Interfaces/DataLayoutInterfaces.td

@@ -280,6 +280,22 @@ def DataLayoutOpInterface : OpInterface<"DataLayoutOpInterface"> {
                                                             params);
       }]
     >,
+    StaticInterfaceMethod<
+      /*description=*/"Returns the bitwidth that should be used when "
+                      "performing index computations for the type computed "
+                      "using the relevant entries. The data layout object can "
+                      "be used for recursive queries.",
+      /*retTy=*/"std::optional<uint64_t>",
+      /*methodName=*/"getIndexBitwidth",
+      /*args=*/(ins "::mlir::Type":$type,
+                    "const ::mlir::DataLayout &":$dataLayout,
+                    "::mlir::DataLayoutEntryListRef":$params),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return ::mlir::detail::getDefaultIndexBitwidth(type, dataLayout,
+                                                       params);
+      }]
+    >,
     StaticInterfaceMethod<
       /*description=*/"Returns the memory space used by the ABI computed "
                       "using the relevant entries. The data layout object "
@@ -400,6 +416,18 @@ def DataLayoutTypeInterface : TypeInterface<"DataLayoutTypeInterface"> {
       /*args=*/(ins "const ::mlir::DataLayout &":$dataLayout,
                     "::mlir::DataLayoutEntryListRef":$params)
     >,
+    InterfaceMethod<
+      /*description=*/"Returns the bitwidth that should be used when "
+                      "performing index computations for the given "
+                      "pointer-like type. If the type is not a pointer-like "
+                      "type, returns std::nullopt.",
+      /*retTy=*/"std::optional<uint64_t>",
+      /*methodName=*/"getIndexBitwidth",
+      /*args=*/(ins "const ::mlir::DataLayout &":$dataLayout,
+                    "::mlir::DataLayoutEntryListRef":$params),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{ return std::nullopt; }]
+    >,
     InterfaceMethod<
       /*desc=*/"Returns true if the two lists of entries are compatible, that "
                "is, that `newLayout` spec entries can be nested in an op with "

mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp

@@ -287,15 +287,22 @@ getPointerDataLayoutEntry(DataLayoutEntryListRef params, LLVMPointerType type,
     }
   }
   if (currentEntry) {
-    return *extractPointerSpecValue(currentEntry, pos) /
-           (pos == PtrDLEntryPos::Size ? 1 : kBitsInByte);
+    std::optional<uint64_t> value = extractPointerSpecValue(currentEntry, pos);
+    // If the optional `PtrDLEntryPos::Index` entry is not available, use the
+    // pointer size as the index bitwidth.
+    if (!value && pos == PtrDLEntryPos::Index)
+      value = extractPointerSpecValue(currentEntry, PtrDLEntryPos::Size);
+    bool isSizeOrIndex =
+        pos == PtrDLEntryPos::Size || pos == PtrDLEntryPos::Index;
+    return *value / (isSizeOrIndex ? 1 : kBitsInByte);
   }
 
   // If not found, and this is the pointer to the default memory space, assume
   // 64-bit pointers.
   if (type.getAddressSpace() == 0) {
-    return pos == PtrDLEntryPos::Size ? kDefaultPointerSizeBits
-                                      : kDefaultPointerAlignment;
+    bool isSizeOrIndex =
+        pos == PtrDLEntryPos::Size || pos == PtrDLEntryPos::Index;
+    return isSizeOrIndex ? kDefaultPointerSizeBits : kDefaultPointerAlignment;
   }
 
   return std::nullopt;
@@ -332,6 +339,16 @@ LLVMPointerType::getPreferredAlignment(const DataLayout &dataLayout,
   return dataLayout.getTypePreferredAlignment(get(getContext()));
 }
 
+std::optional<uint64_t>
+LLVMPointerType::getIndexBitwidth(const DataLayout &dataLayout,
+                                  DataLayoutEntryListRef params) const {
+  if (std::optional<uint64_t> indexBitwidth =
+          getPointerDataLayoutEntry(params, *this, PtrDLEntryPos::Index))
+    return *indexBitwidth;
+
+  return dataLayout.getTypeIndexBitwidth(get(getContext()));
+}
+
 bool LLVMPointerType::areCompatible(DataLayoutEntryListRef oldLayout,
                                     DataLayoutEntryListRef newLayout) const {
   for (DataLayoutEntryInterface newEntry : newLayout) {

mlir/lib/Interfaces/DataLayoutInterfaces.cpp

@@ -218,7 +218,23 @@ uint64_t mlir::detail::getDefaultPreferredAlignment(
   reportMissingDataLayout(type);
 }
 
-// Returns the memory space used for allocal operations if specified in the
+std::optional<uint64_t> mlir::detail::getDefaultIndexBitwidth(
+    Type type, const DataLayout &dataLayout,
+    ArrayRef<DataLayoutEntryInterface> params) {
+  if (isa<IndexType>(type))
+    return getIndexBitwidth(params);
+
+  if (auto typeInterface = dyn_cast<DataLayoutTypeInterface>(type))
+    if (std::optional<uint64_t> indexBitwidth =
+            typeInterface.getIndexBitwidth(dataLayout, params))
+      return *indexBitwidth;
+
+  // Return std::nullopt for all other types, which are assumed to be non
+  // pointer-like types.
+  return std::nullopt;
+}
+
+// Returns the memory space used for alloca operations if specified in the
 // given entry. If the entry is empty the default memory space represented by
 // an empty attribute is returned.
 Attribute
@@ -520,6 +536,18 @@ uint64_t mlir::DataLayout::getTypePreferredAlignment(Type t) const {
   });
 }
 
+std::optional<uint64_t> mlir::DataLayout::getTypeIndexBitwidth(Type t) const {
+  checkValid();
+  return cachedLookup<std::optional<uint64_t>>(t, indexBitwidths, [&](Type ty) {
+    DataLayoutEntryList list;
+    if (originalLayout)
+      list = originalLayout.getSpecForType(ty.getTypeID());
+    if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
+      return iface.getIndexBitwidth(ty, *this, list);
+    return detail::getDefaultIndexBitwidth(ty, *this, list);
+  });
+}
+
 mlir::Attribute mlir::DataLayout::getAllocaMemorySpace() const {
   checkValid();
   if (allocaMemorySpace)

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

@@ -274,16 +274,15 @@ translateDataLayout(DataLayoutSpecInterface attribute,
                 layoutStream << ":" << preferred;
               return success();
             })
-            .Case([&](LLVMPointerType ptrType) {
-              layoutStream << "p" << ptrType.getAddressSpace() << ":";
+            .Case([&](LLVMPointerType type) {
+              layoutStream << "p" << type.getAddressSpace() << ":";
               uint64_t size = dataLayout.getTypeSizeInBits(type);
               uint64_t abi = dataLayout.getTypeABIAlignment(type) * 8u;
               uint64_t preferred =
                   dataLayout.getTypePreferredAlignment(type) * 8u;
-              layoutStream << size << ":" << abi << ":" << preferred;
-              if (std::optional<uint64_t> index = extractPointerSpecValue(
-                      entry.getValue(), PtrDLEntryPos::Index))
-                layoutStream << ":" << *index;
+              uint64_t index = *dataLayout.getTypeIndexBitwidth(type);
+              layoutStream << size << ":" << abi << ":" << preferred << ":"
+                           << index;
               return success();
             })
             .Default([loc](Type type) {

mlir/test/Dialect/LLVMIR/layout.mlir

@@ -7,6 +7,7 @@ module {
     // CHECK: alloca_memory_space = 0
     // CHECK: bitsize = 64
     // CHECK: global_memory_space = 0
+    // CHECK: index = 64
     // CHECK: preferred = 8
     // CHECK: program_memory_space = 0
     // CHECK: size = 8
@@ -16,6 +17,7 @@ module {
     // CHECK: alloca_memory_space = 0
     // CHECK: bitsize = 64
     // CHECK: global_memory_space = 0
+    // CHECK: index = 64
     // CHECK: preferred = 8
     // CHECK: program_memory_space = 0
     // CHECK: size = 8
@@ -25,6 +27,7 @@ module {
     // CHECK: alloca_memory_space = 0
     // CHECK: bitsize = 64
     // CHECK: global_memory_space = 0
+    // CHECK: index = 64
     // CHECK: preferred = 8
     // CHECK: program_memory_space = 0
     // CHECK: size = 8
@@ -39,7 +42,7 @@ module {
 module attributes { dlti.dl_spec = #dlti.dl_spec<
   #dlti.dl_entry<!llvm.ptr, dense<[32, 32, 64]> : vector<3xi64>>,
   #dlti.dl_entry<!llvm.ptr<5>, dense<[64, 64, 64]> : vector<3xi64>>,
-  #dlti.dl_entry<!llvm.ptr<4>, dense<[32, 64, 64]> : vector<3xi64>>,
+  #dlti.dl_entry<!llvm.ptr<4>, dense<[32, 64, 64, 24]> : vector<4xi64>>,
   #dlti.dl_entry<"dlti.alloca_memory_space", 5 : ui64>,
   #dlti.dl_entry<"dlti.global_memory_space", 2 : ui64>,
   #dlti.dl_entry<"dlti.program_memory_space", 3 : ui64>,
@@ -51,6 +54,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
     // CHECK: alloca_memory_space = 5
     // CHECK: bitsize = 32
     // CHECK: global_memory_space = 2
+    // CHECK: index = 32
     // CHECK: preferred = 8
     // CHECK: program_memory_space = 3
     // CHECK: size = 4
@@ -60,6 +64,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
     // CHECK: alloca_memory_space = 5
     // CHECK: bitsize = 32
     // CHECK: global_memory_space = 2
+    // CHECK: index = 32
     // CHECK: preferred = 8
     // CHECK: program_memory_space = 3
     // CHECK: size = 4
@@ -69,24 +74,17 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
     // CHECK: alloca_memory_space = 5
     // CHECK: bitsize = 64
     // CHECK: global_memory_space = 2
+    // CHECK: index = 64
     // CHECK: preferred = 8
     // CHECK: program_memory_space = 3
     // CHECK: size = 8
     // CHECK: stack_alignment = 128
     "test.data_layout_query"() : () -> !llvm.ptr<5>
-    // CHECK: alignment = 4
-    // CHECK: alloca_memory_space = 5
-    // CHECK: bitsize = 32
-    // CHECK: global_memory_space = 2
-    // CHECK: preferred = 8
-    // CHECK: program_memory_space = 3
-    // CHECK: size = 4
-    // CHECK: stack_alignment = 128
-    "test.data_layout_query"() : () -> !llvm.ptr<3>
     // CHECK: alignment = 8
     // CHECK: alloca_memory_space = 5
     // CHECK: bitsize = 32
     // CHECK: global_memory_space = 2
+    // CHECK: index = 24
     // CHECK: preferred = 8
     // CHECK: program_memory_space = 3
     // CHECK: size = 4
@@ -134,34 +132,39 @@ module {
         // simple case
         // CHECK: alignment = 4
         // CHECK: bitsize = 32
+        // CHECK: index = 0
         // CHECK: preferred = 4
         // CHECK: size = 4
         "test.data_layout_query"() : () -> !llvm.struct<(i32)>
 
         // padding inbetween
         // CHECK: alignment = 8
         // CHECK: bitsize = 128
+        // CHECK: index = 0
         // CHECK: preferred = 8
         // CHECK: size = 16
         "test.data_layout_query"() : () -> !llvm.struct<(i32, f64)>
 
         // padding at end of struct
         // CHECK: alignment = 8
         // CHECK: bitsize = 128
+        // CHECK: index = 0
         // CHECK: preferred = 8
         // CHECK: size = 16
         "test.data_layout_query"() : () -> !llvm.struct<(f64, i32)>
 
          // packed
          // CHECK: alignment = 1
          // CHECK: bitsize = 96
+         // CHECK: index = 0
          // CHECK: preferred = 8
          // CHECK: size = 12
          "test.data_layout_query"() : () -> !llvm.struct<packed (f64, i32)>
 
          // empty
          // CHECK: alignment = 1
          // CHECK: bitsize = 0
+         // CHECK: index = 0
          // CHECK: preferred = 1
          // CHECK: size = 0
          "test.data_layout_query"() : () -> !llvm.struct<()>
@@ -179,27 +182,31 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
         // Strict alignment is applied
         // CHECK: alignment = 4
         // CHECK: bitsize = 16
+        // CHECK: index = 0
         // CHECK: preferred = 4
         // CHECK: size = 2
         "test.data_layout_query"() : () -> !llvm.struct<(i16)>
 
         // No impact on structs that have stricter requirements
         // CHECK: alignment = 8
         // CHECK: bitsize = 128
+        // CHECK: index = 0
         // CHECK: preferred = 8
         // CHECK: size = 16
         "test.data_layout_query"() : () -> !llvm.struct<(i32, f64)>
 
          // Only the preferred alignment of structs is affected
          // CHECK: alignment = 1
          // CHECK: bitsize = 32
+         // CHECK: index = 0
          // CHECK: preferred = 4
          // CHECK: size = 4
          "test.data_layout_query"() : () -> !llvm.struct<packed (i16, i16)>
 
          // empty
          // CHECK: alignment = 4
          // CHECK: bitsize = 0
+         // CHECK: index = 0
          // CHECK: preferred = 4
          // CHECK: size = 0
          "test.data_layout_query"() : () -> !llvm.struct<()>
@@ -265,20 +272,23 @@ module {
         // simple case
         // CHECK: alignment = 4
         // CHECK: bitsize = 64
+        // CHECK: index = 0
         // CHECK: preferred = 4
         // CHECK: size = 8
         "test.data_layout_query"() : () -> !llvm.array<2 x i32>
 
         // size 0
         // CHECK: alignment = 8
         // CHECK: bitsize = 0
+        // CHECK: index = 0
         // CHECK: preferred = 8
         // CHECK: size = 0
         "test.data_layout_query"() : () -> !llvm.array<0 x f64>
 
         // alignment info matches element type
         // CHECK: alignment = 4
         // CHECK: bitsize = 64
+        // CHECK: index = 0
         // CHECK: preferred = 8
         // CHECK: size = 8
         "test.data_layout_query"() : () -> !llvm.array<1 x i64>

mlir/test/Interfaces/DataLayoutInterfaces/module.mlir

@@ -2,11 +2,13 @@
 
 module attributes { dlti.dl_spec = #dlti.dl_spec<
       #dlti.dl_entry<!test.test_type_with_layout<10>, ["size", 12]>,
-      #dlti.dl_entry<!test.test_type_with_layout<20>, ["alignment", 32]>>} {
+      #dlti.dl_entry<!test.test_type_with_layout<20>, ["alignment", 32]>,
+      #dlti.dl_entry<!test.test_type_with_layout<30>, ["index", 7]>>} {
   // CHECK-LABEL: @module_level_layout
   func.func @module_level_layout() {
      // CHECK: alignment = 32
      // CHECK: bitsize = 12
+     // CHECK: index = 7
      // CHECK: preferred = 1
      // CHECK: size = 2
     "test.data_layout_query"() : () -> !test.test_type_with_layout<10>