[mlir] use TypeSize and uint64_t in DataLayout (#72874)

Data layout queries may be issued for types whose size exceeds the range
of 32-bit integer as well as for types that don't have a size known at
compile time, such as scalable vectors. Use best practices from LLVM IR
and adopt `llvm::TypeSize` for size-related queries and `uint64_t` for
alignment-related queries.

See #72678.

Author: ftynse

mlir/docs/DataLayout.md

@@ -73,10 +73,10 @@ class DataLayout {
 public:
   explicit DataLayout(DataLayoutOpInterface scope);
 
-  unsigned getTypeSize(Type type) const;
-  unsigned getTypeSizeInBits(Type type) const;
-  unsigned getTypeABIAlignment(Type type) const;
-  unsigned getTypePreferredAlignment(Type type) const;
+  llvm::TypeSize getTypeSize(Type type) const;
+  llvm::TypeSize getTypeSizeInBits(Type type) const;
+  uint64_t getTypeABIAlignment(Type type) const;
+  uint64_t getTypePreferredAlignment(Type type) const;
 };
 ```
 

mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h

@@ -186,13 +186,13 @@ class LLVMStructType
 
   /// Hooks for DataLayoutTypeInterface. Should not be called directly. Obtain a
   /// DataLayout instance and query it instead.
-  unsigned getTypeSizeInBits(const DataLayout &dataLayout,
-                             DataLayoutEntryListRef params) const;
+  llvm::TypeSize getTypeSizeInBits(const DataLayout &dataLayout,
+                                   DataLayoutEntryListRef params) const;
 
-  unsigned getABIAlignment(const DataLayout &dataLayout,
+  uint64_t getABIAlignment(const DataLayout &dataLayout,
                            DataLayoutEntryListRef params) const;
 
-  unsigned getPreferredAlignment(const DataLayout &dataLayout,
+  uint64_t getPreferredAlignment(const DataLayout &dataLayout,
                                  DataLayoutEntryListRef params) const;
 
   bool areCompatible(DataLayoutEntryListRef oldLayout,
@@ -288,7 +288,7 @@ enum class PtrDLEntryPos { Size = 0, Abi = 1, Preferred = 2, Index = 3 };
 /// Returns `std::nullopt` if `pos` is not present in the entry.
 /// Currently only `PtrDLEntryPos::Index` is optional, and all other positions
 /// may be assumed to be present.
-std::optional<unsigned> extractPointerSpecValue(Attribute attr,
+std::optional<uint64_t> extractPointerSpecValue(Attribute attr,
                                                 PtrDLEntryPos pos);
 
 } // namespace LLVM

mlir/include/mlir/Interfaces/DataLayoutInterfaces.h

@@ -18,6 +18,7 @@
 #include "mlir/IR/DialectInterface.h"
 #include "mlir/IR/OpDefinition.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/Support/TypeSize.h"
 
 namespace mlir {
 class DataLayout;
@@ -34,25 +35,25 @@ class ModuleOp;
 namespace detail {
 /// Default handler for the type size request. Computes results for built-in
 /// types and dispatches to the DataLayoutTypeInterface for other types.
-unsigned getDefaultTypeSize(Type type, const DataLayout &dataLayout,
-                            DataLayoutEntryListRef params);
+llvm::TypeSize getDefaultTypeSize(Type type, const DataLayout &dataLayout,
+                                  DataLayoutEntryListRef params);
 
 /// Default handler for the type size in bits request. Computes results for
 /// built-in types and dispatches to the DataLayoutTypeInterface for other
 /// types.
-unsigned getDefaultTypeSizeInBits(Type type, const DataLayout &dataLayout,
-                                  DataLayoutEntryListRef params);
+llvm::TypeSize getDefaultTypeSizeInBits(Type type, const DataLayout &dataLayout,
+                                        DataLayoutEntryListRef params);
 
-/// Default handler for the required alignemnt request. Computes results for
+/// Default handler for the required alignment request. Computes results for
 /// built-in types and dispatches to the DataLayoutTypeInterface for other
 /// types.
-unsigned getDefaultABIAlignment(Type type, const DataLayout &dataLayout,
+uint64_t getDefaultABIAlignment(Type type, const DataLayout &dataLayout,
                                 ArrayRef<DataLayoutEntryInterface> params);
 
-/// Default handler for the preferred alignemnt request. Computes results for
+/// Default handler for the preferred alignment request. Computes results for
 /// built-in types and dispatches to the DataLayoutTypeInterface for other
 /// types.
-unsigned
+uint64_t
 getDefaultPreferredAlignment(Type type, const DataLayout &dataLayout,
                              ArrayRef<DataLayoutEntryInterface> params);
 
@@ -62,7 +63,7 @@ Attribute getDefaultAllocaMemorySpace(DataLayoutEntryInterface entry);
 
 /// Default handler for the stack alignment request. Dispatches to the
 /// DataLayoutInterface if specified, otherwise returns the default.
-unsigned getDefaultStackAlignment(DataLayoutEntryInterface entry);
+uint64_t getDefaultStackAlignment(DataLayoutEntryInterface entry);
 
 /// Given a list of data layout entries, returns a new list containing the
 /// entries with keys having the given type ID, i.e. belonging to the same type
@@ -85,6 +86,10 @@ LogicalResult verifyDataLayoutOp(Operation *op);
 /// entry verifiers, and then to the verifiers implemented by the relevant type
 /// and dialect interfaces for type and identifier keys respectively.
 LogicalResult verifyDataLayoutSpec(DataLayoutSpecInterface spec, Location loc);
+
+/// Divides the known min value of the numerator by the denominator and rounds
+/// the result up to the next integer. Preserves the scalable flag.
+llvm::TypeSize divideCeil(llvm::TypeSize numerator, uint64_t denominator);
 } // namespace detail
 } // namespace mlir
 
@@ -156,16 +161,16 @@ class DataLayout {
   static DataLayout closest(Operation *op);
 
   /// Returns the size of the given type in the current scope.
-  unsigned getTypeSize(Type t) const;
+  llvm::TypeSize getTypeSize(Type t) const;
 
   /// Returns the size in bits of the given type in the current scope.
-  unsigned getTypeSizeInBits(Type t) const;
+  llvm::TypeSize getTypeSizeInBits(Type t) const;
 
   /// Returns the required alignment of the given type in the current scope.
-  unsigned getTypeABIAlignment(Type t) const;
+  uint64_t getTypeABIAlignment(Type t) const;
 
   /// Returns the preferred of the given type in the current scope.
-  unsigned getTypePreferredAlignment(Type t) const;
+  uint64_t getTypePreferredAlignment(Type t) const;
 
   /// Returns the memory space used for AllocaOps.
   Attribute getAllocaMemorySpace() const;
@@ -174,7 +179,7 @@ class DataLayout {
   /// stack variables should be limited to the natural stack alignment to
   /// prevent dynamic stack alignment. Returns zero if the stack alignment is
   /// unspecified.
-  unsigned getStackAlignment() const;
+  uint64_t getStackAlignment() const;
 
 private:
   /// Combined layout spec at the given scope.
@@ -193,16 +198,16 @@ class DataLayout {
   Operation *scope;
 
   /// Caches for individual requests.
-  mutable DenseMap<Type, unsigned> sizes;
-  mutable DenseMap<Type, unsigned> bitsizes;
-  mutable DenseMap<Type, unsigned> abiAlignments;
-  mutable DenseMap<Type, unsigned> preferredAlignments;
+  mutable DenseMap<Type, llvm::TypeSize> sizes;
+  mutable DenseMap<Type, llvm::TypeSize> bitsizes;
+  mutable DenseMap<Type, uint64_t> abiAlignments;
+  mutable DenseMap<Type, uint64_t> preferredAlignments;
 
   /// Cache for alloca memory space.
   mutable std::optional<Attribute> allocaMemorySpace;
 
   /// Cache for stack alignment.
-  mutable std::optional<unsigned> stackAlignment;
+  mutable std::optional<uint64_t> stackAlignment;
 };
 
 } // namespace mlir

mlir/include/mlir/Interfaces/DataLayoutInterfaces.td

@@ -213,22 +213,22 @@ def DataLayoutOpInterface : OpInterface<"DataLayoutOpInterface"> {
       /*description=*/"Returns the size of the given type computed using the "
                       "relevant entries. The data layout object can be used "
                       "for recursive queries.",
-      /*retTy=*/"unsigned",
+      /*retTy=*/"::llvm::TypeSize",
       /*methodName=*/"getTypeSize",
       /*args=*/(ins "::mlir::Type":$type,
                     "const ::mlir::DataLayout &":$dataLayout,
                     "::mlir::DataLayoutEntryListRef":$params),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        unsigned bits = ConcreteOp::getTypeSizeInBits(type, dataLayout, params);
-        return ::llvm::divideCeil(bits, 8);
+        ::llvm::TypeSize bits = ConcreteOp::getTypeSizeInBits(type, dataLayout, params);
+        return ::mlir::detail::divideCeil(bits, 8u);
       }]
     >,
     StaticInterfaceMethod<
       /*description=*/"Returns the size of the given type in bits computed "
                       "using the relevant entries. The data layout object can "
                       "be used for recursive queries.",
-      /*retTy=*/"unsigned",
+      /*retTy=*/"::llvm::TypeSize",
       /*methodName=*/"getTypeSizeInBits",
       /*args=*/(ins "::mlir::Type":$type,
                     "const ::mlir::DataLayout &":$dataLayout,
@@ -243,7 +243,7 @@ def DataLayoutOpInterface : OpInterface<"DataLayoutOpInterface"> {
       /*description=*/"Returns the alignment required by the ABI for the given "
                       "type computed using the relevant entries. The data "
                       "layout object can be used for recursive queries.",
-      /*retTy=*/"unsigned",
+      /*retTy=*/"uint64_t",
       /*methodName=*/"getTypeABIAlignment",
       /*args=*/(ins "::mlir::Type":$type,
                     "const ::mlir::DataLayout &":$dataLayout,
@@ -257,7 +257,7 @@ def DataLayoutOpInterface : OpInterface<"DataLayoutOpInterface"> {
       /*description=*/"Returns the alignment preferred by the given type "
                       "computed using the relevant entries. The data layout"
                       "object can be used for recursive queries.",
-      /*retTy=*/"unsigned",
+      /*retTy=*/"uint64_t",
       /*methodName=*/"getTypePreferredAlignment",
       /*args=*/(ins "::mlir::Type":$type,
                     "const ::mlir::DataLayout &":$dataLayout,
@@ -284,7 +284,7 @@ def DataLayoutOpInterface : OpInterface<"DataLayoutOpInterface"> {
       /*description=*/"Returns the natural stack alignment in bits computed "
                       "using the relevant entries. The data layout object "
                       "can be used for recursive queries.",
-      /*retTy=*/"unsigned",
+      /*retTy=*/"uint64_t",
       /*methodName=*/"getStackAlignment",
       /*args=*/(ins "::mlir::DataLayoutEntryInterface":$entry),
       /*methodBody=*/"",
@@ -331,35 +331,35 @@ def DataLayoutTypeInterface : TypeInterface<"DataLayoutTypeInterface"> {
   let methods = [
     InterfaceMethod<
       /*description=*/"Returns the size of this type in bytes.",
-      /*retTy=*/"unsigned",
+      /*retTy=*/"::llvm::TypeSize",
       /*methodName=*/"getTypeSize",
       /*args=*/(ins "const ::mlir::DataLayout &":$dataLayout,
                     "::mlir::DataLayoutEntryListRef":$params),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        unsigned bits = $_type.getTypeSizeInBits(dataLayout, params);
-        return ::llvm::divideCeil(bits, 8);
+        ::llvm::TypeSize bits = $_type.getTypeSizeInBits(dataLayout, params);
+        return ::mlir::detail::divideCeil(bits, 8u);
       }]
     >,
     InterfaceMethod<
       /*description=*/"Returns the size of this type in bits.",
-      /*retTy=*/"unsigned",
+      /*retTy=*/"::llvm::TypeSize",
       /*methodName=*/"getTypeSizeInBits",
       /*args=*/(ins "const ::mlir::DataLayout &":$dataLayout,
                     "::mlir::DataLayoutEntryListRef":$params)
     >,
     InterfaceMethod<
       /*description=*/"Returns the ABI-required alignment for this type, "
                       "in bytes",
-      /*retTy=*/"unsigned",
+      /*retTy=*/"uint64_t",
       /*methodName=*/"getABIAlignment",
       /*args=*/(ins "const ::mlir::DataLayout &":$dataLayout,
                     "::mlir::DataLayoutEntryListRef":$params)
     >,
     InterfaceMethod<
       /*description=*/"Returns the preferred alignment for this type, "
                       "in bytes.",
-      /*retTy=*/"unsigned",
+      /*retTy=*/"uint64_t",
       /*methodName=*/"getPreferredAlignment",
       /*args=*/(ins "const ::mlir::DataLayout &":$dataLayout,
                     "::mlir::DataLayoutEntryListRef":$params)