[mlir] Data layout for integer and float types

[ftynse]

Summary

Add support for integer and float types into the data layout subsystem with
default logic similar to LLVM IR. Given the flexibility of the sybsystem, the
logic can be easily overwritten by operations if necessary. This provides the
connection necessary, e.g., for the GPU target where alignment requirements for
integers and floats differ from those provided by default (although still
compatible with the LLVM IR model). Previously, it was impossible to use
non-default alignment requirements for integer and float types, which could
lead to incorrect address and size calculations when targeting GPUs.

Depends On D120737

mlir/lib/Interfaces/DataLayoutInterfaces.cpp

return vecType.getNumElements() / vecType.getShape().back() *
            dataLayout.getTypeSize(vecType.getElementType()) * 8;
 
   if (auto typeInterface = type.dyn_cast<DataLayoutTypeInterface>())
     return typeInterface.getTypeSizeInBits(dataLayout, params);
 
   reportMissingDataLayout(type);
 }
 
+static DataLayoutEntryInterface
+findEntryForIntegerType(IntegerType intType,
+                        ArrayRef<DataLayoutEntryInterface> params) {
+  assert(!params.empty() && "expected non-empty parameter list");
+  std::map<unsigned, DataLayoutEntryInterface> sortedParams;
+  for (DataLayoutEntryInterface entry : params) {
+    sortedParams.insert(std::make_pair(
+        entry.getKey().get<Type>().getIntOrFloatBitWidth(), entry));
+  }
+  auto iter = sortedParams.lower_bound(intType.getWidth());
+  if (iter == sortedParams.end())
+    iter = std::prev(iter);
+
+  return iter->second;
+}
+
+static unsigned extractABIAlignment(DataLayoutEntryInterface entry) {
+  auto values =
+      entry.getValue().cast<DenseIntElementsAttr>().getValues<int32_t>();
+  return *values.begin() / 8u;
+}
+
+static unsigned
+getIntegerTypeABIAlignment(IntegerType intType,
+                           ArrayRef<DataLayoutEntryInterface> params) {
+  if (params.empty()) {
+    return intType.getWidth() < 64
+               ? llvm::PowerOf2Ceil(llvm::divideCeil(intType.getWidth(), 8))
+               : 4;
+  }
+
+  return extractABIAlignment(findEntryForIntegerType(intType, params));
+}
+
+static unsigned
+getFloatTypeABIAlignment(FloatType fltType, const DataLayout &dataLayout,
+                         ArrayRef<DataLayoutEntryInterface> params) {
+  assert(params.size() <= 1 && "at most one data layout entry is expected for "
+                               "the singleton floating-point type");
+  if (params.empty())
+    return llvm::PowerOf2Ceil(dataLayout.getTypeSize(fltType));
+  return extractABIAlignment(params[0]);
+}
+
 unsigned mlir::detail::getDefaultABIAlignment(
     Type type, const DataLayout &dataLayout,
     ArrayRef<DataLayoutEntryInterface> params) {
   // Natural alignment is the closest power-of-two number above.
-  if (type.isa<FloatType, VectorType>())
+  if (type.isa<VectorType>())
     return llvm::PowerOf2Ceil(dataLayout.getTypeSize(type));
 
+  if (auto fltType = type.dyn_cast<FloatType>())
+    return getFloatTypeABIAlignment(fltType, dataLayout, params);
+
   // Index is an integer of some bitwidth.
   if (type.isa<IndexType>())
     return dataLayout.getTypeABIAlignment(
         IntegerType::get(type.getContext(), getIndexBitwidth(params)));
 
-  if (auto intType = type.dyn_cast<IntegerType>()) {
-    return intType.getWidth() < 64
-               ? llvm::PowerOf2Ceil(llvm::divideCeil(intType.getWidth(), 8))
-               : 4;
-  }
+  if (auto intType = type.dyn_cast<IntegerType>())
+    return getIntegerTypeABIAlignment(intType, params);
 
   if (auto ctype = type.dyn_cast<ComplexType>())
     return getDefaultABIAlignment(ctype.getElementType(), dataLayout, params);
 
   if (auto typeInterface = type.dyn_cast<DataLayoutTypeInterface>())
     return typeInterface.getABIAlignment(dataLayout, params);
 
   reportMissingDataLayout(type);
 }
 
+static unsigned extractPreferredAlignment(DataLayoutEntryInterface entry) {
+  auto values =
+      entry.getValue().cast<DenseIntElementsAttr>().getValues<int32_t>();
+  return *std::next(values.begin(), values.size() - 1) / 8u;
+}
+
+static unsigned
+getIntegerTypePreferredAlignment(IntegerType intType,
+                                 const DataLayout &dataLayout,
+                                 ArrayRef<DataLayoutEntryInterface> params) {
+  if (params.empty())
+    return llvm::PowerOf2Ceil(dataLayout.getTypeSize(intType));
+
+  return extractPreferredAlignment(findEntryForIntegerType(intType, params));
+}
+
+static unsigned
+getFloatTypePreferredAlignment(FloatType fltType, const DataLayout &dataLayout,
+                               ArrayRef<DataLayoutEntryInterface> params) {
+  assert(params.size() <= 1 && "at most one data layout entry is expected for "
+                               "the singleton floating-point type");
+  if (params.empty())
+    return dataLayout.getTypeABIAlignment(fltType);
+  return extractPreferredAlignment(params[0]);
+}
+
 unsigned mlir::detail::getDefaultPreferredAlignment(
     Type type, const DataLayout &dataLayout,
     ArrayRef<DataLayoutEntryInterface> params) {
   // Preferred alignment is same as natural for floats and vectors.
-  if (type.isa<FloatType, VectorType>())
+  if (type.isa<VectorType>())
     return dataLayout.getTypeABIAlignment(type);
 
-  // Preferred alignment is the cloest power-of-two number above for integers
+  if (auto fltType = type.dyn_cast<FloatType>())
+    return getFloatTypePreferredAlignment(fltType, dataLayout, params);
+
+  // Preferred alignment is the closest power-of-two number above for integers
   // (ABI alignment may be smaller).
-  if (type.isa<IntegerType, IndexType>())
-    return llvm::PowerOf2Ceil(dataLayout.getTypeSize(type));
+  if (auto intType = type.dyn_cast<IntegerType>())
+    return getIntegerTypePreferredAlignment(intType, dataLayout, params);
+
+  if (type.isa<IndexType>()) {
+    return dataLayout.getTypePreferredAlignment(
+        IntegerType::get(type.getContext(), getIndexBitwidth(params)));
+  }
 
   if (auto ctype = type.dyn_cast<ComplexType>())
     return getDefaultPreferredAlignment(ctype.getElementType(), dataLayout,
                                         params);
 
   if (auto typeInterface = type.dyn_cast<DataLayoutTypeInterface>())
     return typeInterface.getPreferredAlignment(dataLayout, params);
 
if (sampleType.isa<IndexType>()) {
              "expected one data layout entry for non-parametric 'index' type");
       if (!kvp.second.front().getValue().isa<IntegerAttr>())
         return emitError(loc)
                << "expected integer attribute in the data layout entry for "
                << sampleType;
       continue;
     }
 
+    if (sampleType.isa<IntegerType, FloatType>()) {
+      for (DataLayoutEntryInterface entry : kvp.second) {
+        auto value = entry.getValue().dyn_cast<DenseIntElementsAttr>();
+        if (!value || !value.getElementType().isSignlessInteger(32)) {
+          emitError(loc) << "expected a dense i32 elements attribute in the "
+                            "data layout entry "
+                         << entry;
+          return failure();
+        }
+
+        auto elements = llvm::to_vector<2>(value.getValues<int32_t>());
+        unsigned numElements = elements.size();
+        if (numElements < 1 || numElements > 2) {
+          emitError(loc) << "expected 1 or 2 elements in the data layout entry "
+                         << entry;
+          return failure();
+        }
+
+        int32_t abi = elements[0];
+        int32_t preferred = numElements == 2 ? elements[1] : abi;
+        if (preferred < abi) {
+          emitError(loc)
+              << "preferred alignment is expected to be greater than or equal "
+                 "to the abi alignment in data layout entry "
+              << entry;
+          return failure();
+        }
+      }
+      continue;
+    }
+
     if (isa<BuiltinDialect>(&sampleType.getDialect()))
       return emitError(loc) << "unexpected data layout for a built-in type";
 
     auto dlType = sampleType.dyn_cast<DataLayoutTypeInterface>();
     if (!dlType)
       return emitError(loc)
              << "data layout specified for a type that does not support it";
     if (failed(dlType.verifyEntries(kvp.second, loc)))

mlir/test/Dialect/DLTI/invalid.mlir

"test.op_with_data_layout"() ({
   // expected-error@below {{data layout does not combine with layouts of enclosing ops}}
   // expected-note@above {{enclosing op with data layout}}
   "test.op_with_data_layout"() { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"unknown.unknown", 32>> } : () -> ()
   "test.maybe_terminator_op"() : () -> ()
 }) { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"unknown.unknown", 33>> } : () -> ()
 
 // -----
 
-// Layout not supported for built-in types.
+// Layout not supported some built-in types.
 // expected-error@below {{unexpected data layout for a built-in type}}
-"test.op_with_data_layout"() { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<i32, 32>> } : () -> ()
+"test.op_with_data_layout"() { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<tensor<f32>, 32>> } : () -> ()
 
 // -----
 
 // expected-error@below {{data layout specified for a type that does not support it}}
 "test.op_with_data_layout"() { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<!test.test_type, 32>> } : () -> ()
 
 // -----
 

mlir/test/Interfaces/DataLayoutInterfaces/query.mlir

"test.op_with_data_layout"() ({
     "test.data_layout_query"() : () -> !test.test_type_with_layout<10>
     "test.maybe_terminator"() : () -> ()
   }) { dlti.dl_spec = #dlti.dl_spec<
       #dlti.dl_entry<!test.test_type_with_layout<10>, ["size", 42]>,
       #dlti.dl_entry<!test.test_type_with_layout<30>, ["preferred", 30]>
   >}: () -> ()
   return
 }
+
+// CHECK-LABEL: @integers
+func @integers() {
+  "test.op_with_data_layout"() ({
+    // CHECK: alignment = 8
+    // CHECK: bitsize = 32
+    // CHECK: preferred = 8
+    "test.data_layout_query"() : () -> i32
+    // CHECK: alignment = 16
+    // CHECK: bitsize = 56
+    // CHECK: preferred = 16
+    "test.data_layout_query"() : () -> i56
+    // CHECK: alignment = 16
+    // CHECK: bitsize = 64
+    // CHECK: preferred = 16
+    "test.data_layout_query"() : () -> i64
+    // CHECK: alignment = 16
+    // CHECK: bitsize = 128
+    // CHECK: preferred = 16
+    "test.data_layout_query"() : () -> i128
+    "test.maybe_terminator"() : () -> ()
+  }) { dlti.dl_spec = #dlti.dl_spec<
+      #dlti.dl_entry<i32, dense<64> : vector<1xi32>>,
+      #dlti.dl_entry<i64, dense<128> : vector<1xi32>>
+    >} : () -> ()
+  "test.op_with_data_layout"() ({
+    // CHECK: alignment = 8
+    // CHECK: bitsize = 32
+    // CHECK: preferred = 16
+    "test.data_layout_query"() : () -> i32
+    // CHECK: alignment = 16
+    // CHECK: bitsize = 56
+    // CHECK: preferred = 32
+    "test.data_layout_query"() : () -> i56
+    // CHECK: alignment = 16
+    // CHECK: bitsize = 64
+    // CHECK: preferred = 32
+    "test.data_layout_query"() : () -> i64
+    // CHECK: alignment = 16
+    // CHECK: bitsize = 128
+    // CHECK: preferred = 32
+    "test.data_layout_query"() : () -> i128
+    "test.maybe_terminator"() : () -> ()
+  }) { dlti.dl_spec = #dlti.dl_spec<
+      #dlti.dl_entry<i32, dense<[64, 128]> : vector<2xi32>>,
+      #dlti.dl_entry<i64, dense<[128, 256]> : vector<2xi32>>
+    >} : () -> ()
+  return
+}
+
+func @floats() {
+  "test.op_with_data_layout"() ({
+    // CHECK: alignment = 8
+    // CHECK: bitsize = 32
+    // CHECK: preferred = 8
+    "test.data_layout_query"() : () -> f32
+    // CHECK: alignment = 16
+    // CHECK: bitsize = 80
+    // CHECK: preferred = 16
+    "test.data_layout_query"() : () -> f80
+    "test.maybe_terminator"() : () -> ()
+  }) { dlti.dl_spec = #dlti.dl_spec<
+      #dlti.dl_entry<f32, dense<64> : vector<1xi32>>,
+      #dlti.dl_entry<f80, dense<128> : vector<1xi32>>
+    >} : () -> ()
+  "test.op_with_data_layout"() ({
+    // CHECK: alignment = 8
+    // CHECK: bitsize = 32
+    // CHECK: preferred = 16
+    "test.data_layout_query"() : () -> f32
+    // CHECK: alignment = 16
+    // CHECK: bitsize = 80
+    // CHECK: preferred = 32
+    "test.data_layout_query"() : () -> f80
+    "test.maybe_terminator"() : () -> ()
+  }) { dlti.dl_spec = #dlti.dl_spec<
+      #dlti.dl_entry<f32, dense<[64, 128]> : vector<2xi32>>,
+      #dlti.dl_entry<f80, dense<[128, 256]> : vector<2xi32>>
+    >} : () -> ()
+  return
+}

mlir/test/Interfaces/DataLayoutInterfaces/types.mlir

 // RUN: mlir-opt --split-input-file --verify-diagnostics --test-data-layout-query %s | FileCheck %s
 
 // expected-error@below {{expected integer attribute in the data layout entry for 'index'}}
 module attributes { dlti.dl_spec = #dlti.dl_spec<
   #dlti.dl_entry<index, [32]>>} {
 }
 
 // -----
 
+// expected-error@below {{expected a dense i32 elements attribute}}
+module attributes {dlti.dl_spec = #dlti.dl_spec<
+#dlti.dl_entry<i32, dense<[64,128]> : vector<2xi64>>>
+} {}
+
+// -----
+
+// expected-error@below {{expected 1 or 2 elements}}
+module attributes {dlti.dl_spec = #dlti.dl_spec<
+#dlti.dl_entry<i32, dense<[64,64,64]> : vector<3xi32>>>
+} {}
+
+// -----
+
+// expected-error@below {{preferred alignment is expected to be greater than or equal to the abi alignment}}
+module attributes {dlti.dl_spec = #dlti.dl_spec<
+#dlti.dl_entry<i32, dense<[64,32]> : vector<2xi32>>>
+} {}
+
+// -----
+
 // expected-error@below {{the 'test' dialect does not support identifier data layout entries}}
 "test.op_with_data_layout"() { dlti.dl_spec = #dlti.dl_spec<
   #dlti.dl_entry<index, 32>,
   #dlti.dl_entry<"test.foo", [32]>>} : () -> ()
 
 // -----
 
 // CHECK-LABEL: @index