[mlir][spirv] Add a spirv::InterfaceVarABIAttr.

Summary:
Add a proper dialect-specific attribute for interface variable ABI.

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

Author: Denis Khalikov

mlir/docs/Dialects/SPIR-V.md

@@ -883,14 +883,30 @@ interfaces:
 *   `spv.entry_point_abi` is a struct attribute that should be attached to the
     entry function. It contains:
     *   `local_size` for specifying the local work group size for the dispatch.
-*   `spv.interface_var_abi` is a struct attribute that should be attached to
-    each operand and result of the entry function. It contains:
-    *   `descriptor_set` for specifying the descriptor set number for the
-        corresponding resource variable.
-    *   `binding` for specifying the binding number for the corresponding
-        resource variable.
-    *   `storage_class` for specifying the storage class for the corresponding
-        resource variable.
+*   `spv.interface_var_abi` is attribute that should be attached to each operand
+    and result of the entry function. It should be of `#spv.interface_var_abi`
+    attribute kind, which is defined as:
+
+```
+spv-storage-class     ::= `StorageBuffer` | ...
+spv-descriptor-set    ::= integer-literal
+spv-binding           ::= integer-literal
+spv-interface-var-abi ::= `#` `spv.interface_var_abi` `<(` spv-descriptor-set
+                          `,` spv-binding `)` (`,` spv-storage-class)? `>`
+```
+
+For example,
+
+```
+#spv.interface_var_abi<(0, 0), StorageBuffer>
+#spv.interface_var_abi<(0, 1)>
+```
+
+The attribute has a few fields:
+
+*   Descriptor set number for the corresponding resource variable.
+*   Binding number for the corresponding resource variable.
+*   Storage class for the corresponding resource variable.
 
 The SPIR-V dialect provides a [`LowerABIAttributesPass`][MlirSpirvPasses] for
 consuming these attributes and create SPIR-V module complying with the

mlir/include/mlir/Dialect/SPIRV/SPIRVAttributes.h

@@ -13,6 +13,7 @@
 #ifndef MLIR_DIALECT_SPIRV_SPIRVATTRIBUTES_H
 #define MLIR_DIALECT_SPIRV_SPIRVATTRIBUTES_H
 
+#include "mlir/Dialect/SPIRV/SPIRVTypes.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/Support/LLVM.h"
 
@@ -26,18 +27,57 @@ enum class Extension;
 enum class Version : uint32_t;
 
 namespace detail {
+struct InterfaceVarABIAttributeStorage;
 struct TargetEnvAttributeStorage;
 struct VerCapExtAttributeStorage;
 } // namespace detail
 
 /// SPIR-V dialect-specific attribute kinds.
 namespace AttrKind {
 enum Kind {
-  TargetEnv = Attribute::FIRST_SPIRV_ATTR, /// Target environment
+  InterfaceVarABI = Attribute::FIRST_SPIRV_ATTR, /// Interface var ABI
+  TargetEnv,                                     /// Target environment
   VerCapExt, /// (version, extension, capability) triple
 };
 } // namespace AttrKind
 
+/// An attribute that specifies the information regarding the interface
+/// variable: descriptor set, binding, storage class.
+class InterfaceVarABIAttr
+    : public Attribute::AttrBase<InterfaceVarABIAttr, Attribute,
+                                 detail::InterfaceVarABIAttributeStorage> {
+public:
+  using Base::Base;
+
+  /// Gets a InterfaceVarABIAttr.
+  static InterfaceVarABIAttr get(uint32_t descirptorSet, uint32_t binding,
+                                 Optional<StorageClass> storageClass,
+                                 MLIRContext *context);
+  static InterfaceVarABIAttr get(IntegerAttr descriptorSet, IntegerAttr binding,
+                                 IntegerAttr storageClass);
+
+  /// Returns the attribute kind's name (without the 'spv.' prefix).
+  static StringRef getKindName();
+
+  /// Returns descriptor set.
+  uint32_t getDescriptorSet();
+
+  /// Returns binding.
+  uint32_t getBinding();
+
+  /// Returns `spirv::StorageClass`.
+  Optional<StorageClass> getStorageClass();
+
+  static bool kindof(unsigned kind) {
+    return kind == AttrKind::InterfaceVarABI;
+  }
+
+  static LogicalResult verifyConstructionInvariants(Location loc,
+                                                    IntegerAttr descriptorSet,
+                                                    IntegerAttr binding,
+                                                    IntegerAttr storageClass);
+};
+
 /// An attribute that specifies the SPIR-V (version, capabilities, extensions)
 /// triple.
 class VerCapExtAttr

mlir/include/mlir/Dialect/SPIRV/TargetAndABI.td

@@ -23,18 +23,6 @@
 
 include "mlir/Dialect/SPIRV/SPIRVBase.td"
 
-// For arguments that eventually map to spv.globalVariable for the
-// shader interface, this attribute specifies the information regarding
-// the global variable:
-// 1) Descriptor Set.
-// 2) Binding number.
-// 3) Storage class.
-def SPV_InterfaceVarABIAttr : StructAttr<"InterfaceVarABIAttr", SPIRV_Dialect, [
-    StructFieldAttr<"descriptor_set", I32Attr>,
-    StructFieldAttr<"binding", I32Attr>,
-    StructFieldAttr<"storage_class", OptionalAttr<SPV_StorageClassAttr>>
-]>;
-
 // For entry functions, this attribute specifies information related to entry
 // points in the generated SPIR-V module:
 // 1) WorkGroup Size.

mlir/lib/Dialect/SPIRV/SPIRVAttributes.cpp

@@ -25,6 +25,32 @@ namespace mlir {
 
 namespace spirv {
 namespace detail {
+
+struct InterfaceVarABIAttributeStorage : public AttributeStorage {
+  using KeyTy = std::tuple<Attribute, Attribute, Attribute>;
+
+  InterfaceVarABIAttributeStorage(Attribute descriptorSet, Attribute binding,
+                                  Attribute storageClass)
+      : descriptorSet(descriptorSet), binding(binding),
+        storageClass(storageClass) {}
+
+  bool operator==(const KeyTy &key) const {
+    return std::get<0>(key) == descriptorSet && std::get<1>(key) == binding &&
+           std::get<2>(key) == storageClass;
+  }
+
+  static InterfaceVarABIAttributeStorage *
+  construct(AttributeStorageAllocator &allocator, const KeyTy &key) {
+    return new (allocator.allocate<InterfaceVarABIAttributeStorage>())
+        InterfaceVarABIAttributeStorage(std::get<0>(key), std::get<1>(key),
+                                        std::get<2>(key));
+  }
+
+  Attribute descriptorSet;
+  Attribute binding;
+  Attribute storageClass;
+};
+
 struct VerCapExtAttributeStorage : public AttributeStorage {
   using KeyTy = std::tuple<Attribute, Attribute, Attribute>;
 
@@ -72,6 +98,74 @@ struct TargetEnvAttributeStorage : public AttributeStorage {
 } // namespace spirv
 } // namespace mlir
 
+//===----------------------------------------------------------------------===//
+// InterfaceVarABIAttr
+//===----------------------------------------------------------------------===//
+
+spirv::InterfaceVarABIAttr
+spirv::InterfaceVarABIAttr::get(uint32_t descriptorSet, uint32_t binding,
+                                Optional<spirv::StorageClass> storageClass,
+                                MLIRContext *context) {
+  Builder b(context);
+  auto descriptorSetAttr = b.getI32IntegerAttr(descriptorSet);
+  auto bindingAttr = b.getI32IntegerAttr(binding);
+  auto storageClassAttr =
+      storageClass ? b.getI32IntegerAttr(static_cast<uint32_t>(*storageClass))
+                   : IntegerAttr();
+  return get(descriptorSetAttr, bindingAttr, storageClassAttr);
+}
+
+spirv::InterfaceVarABIAttr
+spirv::InterfaceVarABIAttr::get(IntegerAttr descriptorSet, IntegerAttr binding,
+                                IntegerAttr storageClass) {
+  assert(descriptorSet && binding);
+  MLIRContext *context = descriptorSet.getContext();
+  return Base::get(context, spirv::AttrKind::InterfaceVarABI, descriptorSet,
+                   binding, storageClass);
+}
+
+StringRef spirv::InterfaceVarABIAttr::getKindName() {
+  return "interface_var_abi";
+}
+
+uint32_t spirv::InterfaceVarABIAttr::getBinding() {
+  return getImpl()->binding.cast<IntegerAttr>().getInt();
+}
+
+uint32_t spirv::InterfaceVarABIAttr::getDescriptorSet() {
+  return getImpl()->descriptorSet.cast<IntegerAttr>().getInt();
+}
+
+Optional<spirv::StorageClass> spirv::InterfaceVarABIAttr::getStorageClass() {
+  if (getImpl()->storageClass)
+    return static_cast<spirv::StorageClass>(
+        getImpl()->storageClass.cast<IntegerAttr>().getValue().getZExtValue());
+  return llvm::None;
+}
+
+LogicalResult spirv::InterfaceVarABIAttr::verifyConstructionInvariants(
+    Location loc, IntegerAttr descriptorSet, IntegerAttr binding,
+    IntegerAttr storageClass) {
+  if (!descriptorSet.getType().isSignlessInteger(32))
+    return emitError(loc, "expected 32-bit integer for descriptor set");
+
+  if (!binding.getType().isSignlessInteger(32))
+    return emitError(loc, "expected 32-bit integer for binding");
+
+  if (storageClass) {
+    if (auto storageClassAttr = storageClass.cast<IntegerAttr>()) {
+      auto storageClassValue =
+          spirv::symbolizeStorageClass(storageClassAttr.getInt());
+      if (!storageClassValue)
+        return emitError(loc, "unknown storage class");
+    } else {
+      return emitError(loc, "expected valid storage class");
+    }
+  }
+
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // VerCapExtAttr
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp

@@ -118,7 +118,7 @@ SPIRVDialect::SPIRVDialect(MLIRContext *context)
     : Dialect(getDialectNamespace(), context) {
   addTypes<ArrayType, ImageType, PointerType, RuntimeArrayType, StructType>();
 
-  addAttributes<TargetEnvAttr, VerCapExtAttr>();
+  addAttributes<InterfaceVarABIAttr, TargetEnvAttr, VerCapExtAttr>();
 
   // Add SPIR-V ops.
   addOperations<
@@ -649,6 +649,75 @@ static ParseResult parseKeywordList(
   return success();
 }
 
+/// Parses a spirv::InterfaceVarABIAttr.
+static Attribute parseInterfaceVarABIAttr(DialectAsmParser &parser) {
+  if (parser.parseLess())
+    return {};
+
+  Builder &builder = parser.getBuilder();
+
+  if (parser.parseLParen())
+    return {};
+
+  IntegerAttr descriptorSetAttr;
+  {
+    auto loc = parser.getCurrentLocation();
+    uint32_t descriptorSet = 0;
+    auto descriptorSetParseResult = parser.parseOptionalInteger(descriptorSet);
+
+    if (!descriptorSetParseResult.hasValue() ||
+        failed(*descriptorSetParseResult)) {
+      parser.emitError(loc, "missing descriptor set");
+      return {};
+    }
+    descriptorSetAttr = builder.getI32IntegerAttr(descriptorSet);
+  }
+
+  if (parser.parseComma())
+    return {};
+
+  IntegerAttr bindingAttr;
+  {
+    auto loc = parser.getCurrentLocation();
+    uint32_t binding = 0;
+    auto bindingParseResult = parser.parseOptionalInteger(binding);
+
+    if (!bindingParseResult.hasValue() || failed(*bindingParseResult)) {
+      parser.emitError(loc, "missing binding");
+      return {};
+    }
+    bindingAttr = builder.getI32IntegerAttr(binding);
+  }
+
+  if (parser.parseRParen())
+    return {};
+
+  IntegerAttr storageClassAttr;
+  {
+    if (succeeded(parser.parseOptionalComma())) {
+      auto loc = parser.getCurrentLocation();
+      StringRef storageClass;
+      if (parser.parseKeyword(&storageClass))
+        return {};
+
+      if (auto storageClassSymbol =
+              spirv::symbolizeStorageClass(storageClass)) {
+        storageClassAttr = builder.getI32IntegerAttr(
+            static_cast<uint32_t>(*storageClassSymbol));
+      } else {
+        parser.emitError(loc, "unknown storage class: ") << storageClass;
+        return {};
+      }
+    }
+  }
+
+  if (parser.parseGreater())
+    return {};
+
+  return spirv::InterfaceVarABIAttr::get(descriptorSetAttr, bindingAttr,
+                                         storageClassAttr);
+}
+
 static Attribute parseVerCapExtAttr(DialectAsmParser &parser) {
   if (parser.parseLess())
     return {};
@@ -771,6 +840,8 @@ Attribute SPIRVDialect::parseAttribute(DialectAsmParser &parser,
     return parseTargetEnvAttr(parser);
   if (attrKind == spirv::VerCapExtAttr::getKindName())
     return parseVerCapExtAttr(parser);
+  if (attrKind == spirv::InterfaceVarABIAttr::getKindName())
+    return parseInterfaceVarABIAttr(parser);
 
   parser.emitError(parser.getNameLoc(), "unknown SPIR-V attribute kind: ")
       << attrKind;
@@ -801,12 +872,25 @@ static void print(spirv::TargetEnvAttr targetEnv, DialectAsmPrinter &printer) {
   printer << ", " << targetEnv.getResourceLimits() << ">";
 }
 
+static void print(spirv::InterfaceVarABIAttr interfaceVarABIAttr,
+                  DialectAsmPrinter &printer) {
+  printer << spirv::InterfaceVarABIAttr::getKindName() << "<("
+          << interfaceVarABIAttr.getDescriptorSet() << ", "
+          << interfaceVarABIAttr.getBinding() << ")";
+  auto storageClass = interfaceVarABIAttr.getStorageClass();
+  if (storageClass)
+    printer << ", " << spirv::stringifyStorageClass(*storageClass);
+  printer << ">";
+}
+
 void SPIRVDialect::printAttribute(Attribute attr,
                                   DialectAsmPrinter &printer) const {
   if (auto targetEnv = attr.dyn_cast<TargetEnvAttr>())
     print(targetEnv, printer);
   else if (auto vceAttr = attr.dyn_cast<VerCapExtAttr>())
     print(vceAttr, printer);
+  else if (auto interfaceVarABIAttr = attr.dyn_cast<InterfaceVarABIAttr>())
+    print(interfaceVarABIAttr, printer);
   else
     llvm_unreachable("unhandled SPIR-V attribute kind");
 }
@@ -866,11 +950,9 @@ static LogicalResult verifyRegionAttribute(Location loc, Type valueType,
   auto varABIAttr = attr.dyn_cast<spirv::InterfaceVarABIAttr>();
   if (!varABIAttr)
     return emitError(loc, "'")
-           << symbol
-           << "' attribute must be a dictionary attribute containing two or "
-              "three 32-bit integer attributes: 'descriptor_set', 'binding', "
-              "and optional 'storage_class'";
-  if (varABIAttr.storage_class() && !valueType.isIntOrIndexOrFloat())
+           << symbol << "' must be a spirv::InterfaceVarABIAttr";
+
+  if (varABIAttr.getStorageClass() && !valueType.isIntOrIndexOrFloat())
     return emitError(loc, "'") << symbol
                                << "' attribute cannot specify storage class "
                                   "when attaching to a non-scalar value";

mlir/lib/Dialect/SPIRV/TargetAndABI.cpp

@@ -86,14 +86,8 @@ spirv::InterfaceVarABIAttr
 spirv::getInterfaceVarABIAttr(unsigned descriptorSet, unsigned binding,
                               Optional<spirv::StorageClass> storageClass,
                               MLIRContext *context) {
-  Type i32Type = IntegerType::get(32, context);
-  auto scAttr =
-      storageClass
-          ? IntegerAttr::get(i32Type, static_cast<int64_t>(*storageClass))
-          : IntegerAttr();
-  return spirv::InterfaceVarABIAttr::get(
-      IntegerAttr::get(i32Type, descriptorSet),
-      IntegerAttr::get(i32Type, binding), scAttr, context);
+  return spirv::InterfaceVarABIAttr::get(descriptorSet, binding, storageClass,
+                                         context);
 }
 
 StringRef spirv::getEntryPointABIAttrName() { return "spv.entry_point_abi"; }