[mlir][spirv] Improve stride support in array types

This commit added stride support in runtime array types. It also
adjusted the assembly form for the stride from `[N]` to `stride=N`.
This makes the IR more readable, especially for the cases where
one mix array types and struct types.

Differential Revision: https://reviews.llvm.org/D78034

Author: antiagainst

mlir/docs/Dialects/SPIR-V.md

@@ -287,13 +287,15 @@ element-type ::= integer-type
                | vector-type
                | spirv-type
 
-array-type ::= `!spv.array<` integer-literal `x` element-type `>`
+array-type ::= `!spv.array` `<` integer-literal `x` element-type
+               (`,` `stride` `=` integer-literal)? `>`
 ```
 
 For example,
 
 ```mlir
 !spv.array<4 x i32>
+!spv.array<4 x i32, stride = 4>
 !spv.array<16 x vector<4 x f32>>
 ```
 
@@ -351,13 +353,14 @@ For example,
 This corresponds to SPIR-V [runtime array type][RuntimeArrayType]. Its syntax is
 
 ```
-runtime-array-type ::= `!spv.rtarray<` element-type `>`
+runtime-array-type ::= `!spv.rtarray` `<` element-type (`,` `stride` `=` integer-literal)? `>`
 ```
 
 For example,
 
 ```mlir
 !spv.rtarray<i32>
+!spv.rtarray<i32, stride=4>
 !spv.rtarray<vector<4 x f32>>
 ```
 

mlir/include/mlir/Dialect/SPIRV/LayoutUtils.h

@@ -18,9 +18,11 @@
 namespace mlir {
 class Type;
 class VectorType;
+
 namespace spirv {
-class StructType;
 class ArrayType;
+class RuntimeArrayType;
+class StructType;
 } // namespace spirv
 
 /// According to the Vulkan spec "14.5.4. Offset and Stride Assignment":
@@ -47,21 +49,27 @@ class VulkanLayoutUtils {
 public:
   using Size = uint64_t;
 
-  /// Returns a new StructType with layout info. Assigns the type size in bytes
-  /// to the `size`. Assigns the type alignment in bytes to the `alignment`.
-  static spirv::StructType decorateType(spirv::StructType structType,
-                                        Size &size, Size &alignment);
+  /// Returns a new StructType with layout decoration.
+  static spirv::StructType decorateType(spirv::StructType structType);
+
   /// Checks whether a type is legal in terms of Vulkan layout info
   /// decoration. A type is dynamically illegal if it's a composite type in the
   /// StorageBuffer, PhysicalStorageBuffer, Uniform, and PushConstant Storage
   /// Classes without layout information.
   static bool isLegalType(Type type);
 
 private:
+  /// Returns a new type with layout decoration. Assigns the type size in bytes
+  /// to the `size`. Assigns the type alignment in bytes to the `alignment`.
   static Type decorateType(Type type, Size &size, Size &alignment);
+
   static Type decorateType(VectorType vectorType, Size &size, Size &alignment);
   static Type decorateType(spirv::ArrayType arrayType, Size &size,
                            Size &alignment);
+  static Type decorateType(spirv::RuntimeArrayType arrayType, Size &alignment);
+  static spirv::StructType decorateType(spirv::StructType structType,
+                                        Size &size, Size &alignment);
+
   /// Calculates the alignment for the given scalar type.
   static Size getScalarTypeAlignment(Type scalarType);
 };

mlir/include/mlir/Dialect/SPIRV/SPIRVTypes.h

@@ -147,23 +147,22 @@ class ArrayType : public Type::TypeBase<ArrayType, CompositeType,
                                         detail::ArrayTypeStorage> {
 public:
   using Base::Base;
-  // Zero layout specifies that is no layout
-  using LayoutInfo = uint64_t;
 
   static bool kindof(unsigned kind) { return kind == TypeKind::Array; }
 
   static ArrayType get(Type elementType, unsigned elementCount);
 
+  /// Returns an array type with the given stride in bytes.
   static ArrayType get(Type elementType, unsigned elementCount,
-                       LayoutInfo layoutInfo);
+                       unsigned stride);
 
   unsigned getNumElements() const;
 
   Type getElementType() const;
 
-  bool hasLayout() const;
-
-  uint64_t getArrayStride() const;
+  /// Returns the array stride in bytes. 0 means no stride decorated on this
+  /// type.
+  unsigned getArrayStride() const;
 
   void getExtensions(SPIRVType::ExtensionArrayRefVector &extensions,
                      Optional<spirv::StorageClass> storage = llvm::None);
@@ -243,8 +242,15 @@ class RuntimeArrayType
 
   static RuntimeArrayType get(Type elementType);
 
+  /// Returns a runtime array type with the given stride in bytes.
+  static RuntimeArrayType get(Type elementType, unsigned stride);
+
   Type getElementType() const;
 
+  /// Returns the array stride in bytes. 0 means no stride decorated on this
+  /// type.
+  unsigned getArrayStride() const;
+
   void getExtensions(SPIRVType::ExtensionArrayRefVector &extensions,
                      Optional<spirv::StorageClass> storage = llvm::None);
   void getCapabilities(SPIRVType::CapabilityArrayRefVector &capabilities,

mlir/lib/Dialect/SPIRV/LayoutUtils.cpp

@@ -16,6 +16,13 @@
 
 using namespace mlir;
 
+spirv::StructType
+VulkanLayoutUtils::decorateType(spirv::StructType structType) {
+  Size size = 0;
+  Size alignment = 1;
+  return decorateType(structType, size, alignment);
+}
+
 spirv::StructType
 VulkanLayoutUtils::decorateType(spirv::StructType structType,
                                 VulkanLayoutUtils::Size &size,
@@ -25,21 +32,26 @@ VulkanLayoutUtils::decorateType(spirv::StructType structType,
   }
 
   SmallVector<Type, 4> memberTypes;
-  SmallVector<VulkanLayoutUtils::Size, 4> layoutInfo;
+  SmallVector<Size, 4> layoutInfo;
   SmallVector<spirv::StructType::MemberDecorationInfo, 4> memberDecorations;
 
-  VulkanLayoutUtils::Size structMemberOffset = 0;
-  VulkanLayoutUtils::Size maxMemberAlignment = 1;
+  Size structMemberOffset = 0;
+  Size maxMemberAlignment = 1;
 
   for (uint32_t i = 0, e = structType.getNumElements(); i < e; ++i) {
-    VulkanLayoutUtils::Size memberSize = 0;
-    VulkanLayoutUtils::Size memberAlignment = 1;
+    Size memberSize = 0;
+    Size memberAlignment = 1;
 
-    auto memberType = VulkanLayoutUtils::decorateType(
-        structType.getElementType(i), memberSize, memberAlignment);
+    auto memberType =
+        decorateType(structType.getElementType(i), memberSize, memberAlignment);
     structMemberOffset = llvm::alignTo(structMemberOffset, memberAlignment);
     memberTypes.push_back(memberType);
     layoutInfo.push_back(structMemberOffset);
+    // If the member's size is the max value, it must be the last member and it
+    // must be a runtime array.
+    assert(memberSize != std::numeric_limits<Size>().max() ||
+           (i + 1 == e &&
+            structType.getElementType(i).isa<spirv::RuntimeArrayType>()));
     // According to the Vulkan spec:
     // "A structure has a base alignment equal to the largest base alignment of
     // any of its members."
@@ -60,22 +72,22 @@ VulkanLayoutUtils::decorateType(spirv::StructType structType,
 Type VulkanLayoutUtils::decorateType(Type type, VulkanLayoutUtils::Size &size,
                                      VulkanLayoutUtils::Size &alignment) {
   if (type.isa<spirv::ScalarType>()) {
-    alignment = VulkanLayoutUtils::getScalarTypeAlignment(type);
+    alignment = getScalarTypeAlignment(type);
     // Vulkan spec does not specify any padding for a scalar type.
     size = alignment;
     return type;
   }
 
   switch (type.getKind()) {
   case spirv::TypeKind::Struct:
-    return VulkanLayoutUtils::decorateType(type.cast<spirv::StructType>(), size,
-                                           alignment);
+    return decorateType(type.cast<spirv::StructType>(), size, alignment);
   case spirv::TypeKind::Array:
-    return VulkanLayoutUtils::decorateType(type.cast<spirv::ArrayType>(), size,
-                                           alignment);
+    return decorateType(type.cast<spirv::ArrayType>(), size, alignment);
   case StandardTypes::Vector:
-    return VulkanLayoutUtils::decorateType(type.cast<VectorType>(), size,
-                                           alignment);
+    return decorateType(type.cast<VectorType>(), size, alignment);
+  case spirv::TypeKind::RuntimeArray:
+    size = std::numeric_limits<Size>().max();
+    return decorateType(type.cast<spirv::RuntimeArrayType>(), alignment);
   default:
     llvm_unreachable("unhandled SPIR-V type");
   }
@@ -86,11 +98,10 @@ Type VulkanLayoutUtils::decorateType(VectorType vectorType,
                                      VulkanLayoutUtils::Size &alignment) {
   const auto numElements = vectorType.getNumElements();
   auto elementType = vectorType.getElementType();
-  VulkanLayoutUtils::Size elementSize = 0;
-  VulkanLayoutUtils::Size elementAlignment = 1;
+  Size elementSize = 0;
+  Size elementAlignment = 1;
 
-  auto memberType = VulkanLayoutUtils::decorateType(elementType, elementSize,
-                                                    elementAlignment);
+  auto memberType = decorateType(elementType, elementSize, elementAlignment);
   // According to the Vulkan spec:
   // 1. "A two-component vector has a base alignment equal to twice its scalar
   // alignment."
@@ -106,11 +117,10 @@ Type VulkanLayoutUtils::decorateType(spirv::ArrayType arrayType,
                                      VulkanLayoutUtils::Size &alignment) {
   const auto numElements = arrayType.getNumElements();
   auto elementType = arrayType.getElementType();
-  spirv::ArrayType::LayoutInfo elementSize = 0;
-  VulkanLayoutUtils::Size elementAlignment = 1;
+  Size elementSize = 0;
+  Size elementAlignment = 1;
 
-  auto memberType = VulkanLayoutUtils::decorateType(elementType, elementSize,
-                                                    elementAlignment);
+  auto memberType = decorateType(elementType, elementSize, elementAlignment);
   // According to the Vulkan spec:
   // "An array has a base alignment equal to the base alignment of its element
   // type."
@@ -119,6 +129,15 @@ Type VulkanLayoutUtils::decorateType(spirv::ArrayType arrayType,
   return spirv::ArrayType::get(memberType, numElements, elementSize);
 }
 
+Type VulkanLayoutUtils::decorateType(spirv::RuntimeArrayType arrayType,
+                                     VulkanLayoutUtils::Size &alignment) {
+  auto elementType = arrayType.getElementType();
+  Size elementSize = 0;
+
+  auto memberType = decorateType(elementType, elementSize, alignment);
+  return spirv::RuntimeArrayType::get(memberType, elementSize);
+}
+
 VulkanLayoutUtils::Size
 VulkanLayoutUtils::getScalarTypeAlignment(Type scalarType) {
   // According to the Vulkan spec:

mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp

@@ -149,7 +149,7 @@ Optional<Type> parseAndVerify<Type>(SPIRVDialect const &dialect,
                                     DialectAsmParser &parser);
 
 template <>
-Optional<uint64_t> parseAndVerify<uint64_t>(SPIRVDialect const &dialect,
+Optional<unsigned> parseAndVerify<unsigned>(SPIRVDialect const &dialect,
                                             DialectAsmParser &parser);
 
 static Type parseAndVerifyType(SPIRVDialect const &dialect,
@@ -196,13 +196,39 @@ static Type parseAndVerifyType(SPIRVDialect const &dialect,
   return type;
 }
 
+/// Parses an optional `, stride = N` assembly segment. If no parsing failure
+/// occurs, writes `N` to `stride` if existing and writes 0 to `stride` if
+/// missing.
+static LogicalResult parseOptionalArrayStride(const SPIRVDialect &dialect,
+                                              DialectAsmParser &parser,
+                                              unsigned &stride) {
+  if (failed(parser.parseOptionalComma())) {
+    stride = 0;
+    return success();
+  }
+
+  if (parser.parseKeyword("stride") || parser.parseEqual())
+    return failure();
+
+  llvm::SMLoc strideLoc = parser.getCurrentLocation();
+  Optional<unsigned> optStride = parseAndVerify<unsigned>(dialect, parser);
+  if (!optStride)
+    return failure();
+
+  if (!(stride = optStride.getValue())) {
+    parser.emitError(strideLoc, "ArrayStride must be greater than zero");
+    return failure();
+  }
+  return success();
+}
+
 // element-type ::= integer-type
 //                | floating-point-type
 //                | vector-type
 //                | spirv-type
 //
-// array-type ::= `!spv.array<` integer-literal `x` element-type
-//                (`[` integer-literal `]`)? `>`
+// array-type ::= `!spv.array` `<` integer-literal `x` element-type
+//                (`,` `stride` `=` integer-literal)? `>`
 static Type parseArrayType(SPIRVDialect const &dialect,
                            DialectAsmParser &parser) {
   if (parser.parseLess())
@@ -230,25 +256,13 @@ static Type parseArrayType(SPIRVDialect const &dialect,
   if (!elementType)
     return Type();
 
-  ArrayType::LayoutInfo layoutInfo = 0;
-  if (succeeded(parser.parseOptionalLSquare())) {
-    llvm::SMLoc layoutLoc = parser.getCurrentLocation();
-    auto layout = parseAndVerify<ArrayType::LayoutInfo>(dialect, parser);
-    if (!layout)
-      return Type();
-
-    if (!(layoutInfo = layout.getValue())) {
-      parser.emitError(layoutLoc, "ArrayStride must be greater than zero");
-      return Type();
-    }
-
-    if (parser.parseRSquare())
-      return Type();
-  }
+  unsigned stride = 0;
+  if (failed(parseOptionalArrayStride(dialect, parser, stride)))
+    return Type();
 
   if (parser.parseGreater())
     return Type();
-  return ArrayType::get(elementType, count, layoutInfo);
+  return ArrayType::get(elementType, count, stride);
 }
 
 // TODO(ravishankarm) : Reorder methods to be utilities first and parse*Type
@@ -285,7 +299,8 @@ static Type parsePointerType(SPIRVDialect const &dialect,
   return PointerType::get(pointeeType, *storageClass);
 }
 
-// runtime-array-type ::= `!spv.rtarray<` element-type `>`
+// runtime-array-type ::= `!spv.rtarray` `<` element-type
+//                        (`,` `stride` `=` integer-literal)? `>`
 static Type parseRuntimeArrayType(SPIRVDialect const &dialect,
                                   DialectAsmParser &parser) {
   if (parser.parseLess())
@@ -295,9 +310,13 @@ static Type parseRuntimeArrayType(SPIRVDialect const &dialect,
   if (!elementType)
     return Type();
 
+  unsigned stride = 0;
+  if (failed(parseOptionalArrayStride(dialect, parser, stride)))
+    return Type();
+
   if (parser.parseGreater())
     return Type();
-  return RuntimeArrayType::get(elementType);
+  return RuntimeArrayType::get(elementType, stride);
 }
 
 // Specialize this function to parse each of the parameters that define an
@@ -337,9 +356,9 @@ static Optional<IntTy> parseAndVerifyInteger(SPIRVDialect const &dialect,
 }
 
 template <>
-Optional<uint64_t> parseAndVerify<uint64_t>(SPIRVDialect const &dialect,
+Optional<unsigned> parseAndVerify<unsigned>(SPIRVDialect const &dialect,
                                             DialectAsmParser &parser) {
-  return parseAndVerifyInteger<uint64_t>(dialect, parser);
+  return parseAndVerifyInteger<unsigned>(dialect, parser);
 }
 
 namespace {
@@ -526,14 +545,16 @@ Type SPIRVDialect::parseType(DialectAsmParser &parser) const {
 
 static void print(ArrayType type, DialectAsmPrinter &os) {
   os << "array<" << type.getNumElements() << " x " << type.getElementType();
-  if (type.hasLayout()) {
-    os << " [" << type.getArrayStride() << "]";
-  }
+  if (unsigned stride = type.getArrayStride())
+    os << ", stride=" << stride;
   os << ">";
 }
 
 static void print(RuntimeArrayType type, DialectAsmPrinter &os) {
-  os << "rtarray<" << type.getElementType() << ">";
+  os << "rtarray<" << type.getElementType();
+  if (unsigned stride = type.getArrayStride())
+    os << ", stride=" << stride;
+  os << ">";
 }
 
 static void print(PointerType type, DialectAsmPrinter &os) {