[mlir][spirv] Allow specifying availability on enum attribute cases

Lots of SPIR-V ops take enum attributes and certain enum cases
need extra capabilities or extensions to be available. This commit
extends to allow specifying availability spec on enum cases.
Extra utility functions are generated for the corresponding enum
classes to return the availability requirement. The availability
interface implemention for a SPIR-V op now goes over all enum
attributes to collect the availability requirements.

Reviewed By: mravishankar

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

Author: antiagainst

mlir/include/mlir/Dialect/SPIRV/CMakeLists.txt

@@ -5,6 +5,11 @@ mlir_tablegen(SPIRVEnums.h.inc -gen-enum-decls)
 mlir_tablegen(SPIRVEnums.cpp.inc -gen-enum-defs)
 add_public_tablegen_target(MLIRSPIRVEnumsIncGen)
 
+set(LLVM_TARGET_DEFINITIONS SPIRVBase.td)
+mlir_tablegen(SPIRVEnumAvailability.h.inc -gen-spirv-enum-avail-decls)
+mlir_tablegen(SPIRVEnumAvailability.cpp.inc -gen-spirv-enum-avail-defs)
+add_public_tablegen_target(MLIRSPIRVEnumAvailabilityIncGen)
+
 set(LLVM_TARGET_DEFINITIONS SPIRVOps.td)
 mlir_tablegen(SPIRVAvailability.h.inc -gen-avail-interface-decls)
 mlir_tablegen(SPIRVAvailability.cpp.inc -gen-avail-interface-defs)

mlir/include/mlir/Dialect/SPIRV/SPIRVBase.td

@@ -467,7 +467,13 @@ def SPV_CapabilityAttr :
 def SPV_AM_Logical                 : I32EnumAttrCase<"Logical", 0>;
 def SPV_AM_Physical32              : I32EnumAttrCase<"Physical32", 1>;
 def SPV_AM_Physical64              : I32EnumAttrCase<"Physical64", 2>;
-def SPV_AM_PhysicalStorageBuffer64 : I32EnumAttrCase<"PhysicalStorageBuffer64", 5348>;
+def SPV_AM_PhysicalStorageBuffer64 : I32EnumAttrCase<"PhysicalStorageBuffer64", 5348> {
+  list<Availability> availability = [
+    MinVersion<SPV_V_1_5>,
+    Extension<[SPV_EXT_physical_storage_buffer, SPV_KHR_physical_storage_buffer]>,
+    Capability<[SPV_C_PhysicalStorageBufferAddresses]>
+  ];
+}
 
 def SPV_AddressingModelAttr :
     I32EnumAttr<"AddressingModel", "valid SPIR-V AddressingModel", [
@@ -944,7 +950,12 @@ def SPV_MemoryAccessAttr :
 def SPV_MM_Simple  : I32EnumAttrCase<"Simple", 0>;
 def SPV_MM_GLSL450 : I32EnumAttrCase<"GLSL450", 1>;
 def SPV_MM_OpenCL  : I32EnumAttrCase<"OpenCL", 2>;
-def SPV_MM_Vulkan  : I32EnumAttrCase<"Vulkan", 3>;
+def SPV_MM_Vulkan  : I32EnumAttrCase<"Vulkan", 3> {
+  list<Availability> availability = [
+    MinVersion<SPV_V_1_5>,
+    Capability<[SPV_C_VulkanMemoryModel]>
+  ];
+}
 
 def SPV_MemoryModelAttr :
     I32EnumAttr<"MemoryModel", "valid SPIR-V MemoryModel", [

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

@@ -17,8 +17,21 @@
 #include "mlir/IR/TypeSupport.h"
 #include "mlir/IR/Types.h"
 
+// Forward declare enum classes related to op availability. Their definitions
+// are in the TableGen'erated SPIRVEnums.h.inc and can be referenced by other
+// dclarations in SPIRVEnums.h.inc.
+namespace mlir {
+namespace spirv {
+enum class Version : uint32_t;
+enum class Extension;
+enum class Capability : uint32_t;
+} // namespace spirv
+} // namespace mlir
+
 // Pull in all enum type definitions and utility function declarations
 #include "mlir/Dialect/SPIRV/SPIRVEnums.h.inc"
+// Pull in all enum type availability query function declarations
+#include "mlir/Dialect/SPIRV/SPIRVEnumAvailability.h.inc"
 
 #include <tuple>
 

mlir/include/mlir/TableGen/Attribute.h

@@ -135,6 +135,9 @@ class EnumAttrCase : public Attribute {
 
   // Returns the value of this enum attribute case.
   int64_t getValue() const;
+
+  // Returns the TableGen definition this EnumAttrCase was constructed from.
+  const llvm::Record &getDef() const;
 };
 
 // Wrapper class providing helper methods for accessing enum attributes defined
@@ -146,6 +149,8 @@ class EnumAttr : public Attribute {
   explicit EnumAttr(const llvm::Record &record);
   explicit EnumAttr(const llvm::DefInit *init);
 
+  static bool classof(const Attribute *attr);
+
   // Returns true if this is a bit enum attribute.
   bool isBitEnum() const;
 

mlir/lib/Dialect/SPIRV/CMakeLists.txt

@@ -18,6 +18,7 @@ add_llvm_library(MLIRSPIRV
 add_dependencies(MLIRSPIRV
   MLIRSPIRVAvailabilityIncGen
   MLIRSPIRVCanonicalizationIncGen
+  MLIRSPIRVEnumAvailabilityIncGen
   MLIRSPIRVEnumsIncGen
   MLIRSPIRVOpsIncGen
   MLIRSPIRVOpUtilsGen

mlir/lib/Dialect/SPIRV/SPIRVTypes.cpp

@@ -22,6 +22,8 @@ using namespace mlir::spirv;
 
 // Pull in all enum utility function definitions
 #include "mlir/Dialect/SPIRV/SPIRVEnums.cpp.inc"
+// Pull in all enum type availability query function definitions
+#include "mlir/Dialect/SPIRV/SPIRVEnumAvailability.cpp.inc"
 
 //===----------------------------------------------------------------------===//
 // ArrayType

mlir/lib/TableGen/Attribute.cpp

@@ -155,6 +155,8 @@ int64_t tblgen::EnumAttrCase::getValue() const {
   return def->getValueAsInt("value");
 }
 
+const llvm::Record &tblgen::EnumAttrCase::getDef() const { return *def; }
+
 tblgen::EnumAttr::EnumAttr(const llvm::Record *record) : Attribute(record) {
   assert(isSubClassOf("EnumAttrInfo") &&
          "must be subclass of TableGen 'EnumAttr' class");
@@ -165,6 +167,10 @@ tblgen::EnumAttr::EnumAttr(const llvm::Record &record) : Attribute(&record) {}
 tblgen::EnumAttr::EnumAttr(const llvm::DefInit *init)
     : EnumAttr(init->getDef()) {}
 
+bool tblgen::EnumAttr::classof(const Attribute *attr) {
+  return attr->isSubClassOf("EnumAttrInfo");
+}
+
 bool tblgen::EnumAttr::isBitEnum() const { return isSubClassOf("BitEnumAttr"); }
 
 StringRef tblgen::EnumAttr::getEnumClassName() const {

mlir/test/Dialect/SPIRV/TestAvailability.cpp

@@ -30,18 +30,19 @@ void TestAvailability::runOnFunction() {
     if (op->getDialect() != spvDialect)
       return WalkResult::advance();
 
+    auto opName = op->getName();
     auto &os = llvm::outs();
 
     if (auto minVersion = dyn_cast<spirv::QueryMinVersionInterface>(op))
-      os << " min version: "
+      os << opName << " min version: "
          << spirv::stringifyVersion(minVersion.getMinVersion()) << "\n";
 
     if (auto maxVersion = dyn_cast<spirv::QueryMaxVersionInterface>(op))
-      os << " max version: "
+      os << opName << " max version: "
          << spirv::stringifyVersion(maxVersion.getMaxVersion()) << "\n";
 
     if (auto extension = dyn_cast<spirv::QueryExtensionInterface>(op)) {
-      os << " extensions: [";
+      os << opName << " extensions: [";
       for (const auto &exts : extension.getExtensions()) {
         os << " [";
         interleaveComma(exts, os, [&](spirv::Extension ext) {
@@ -53,7 +54,7 @@ void TestAvailability::runOnFunction() {
     }
 
     if (auto capability = dyn_cast<spirv::QueryCapabilityInterface>(op)) {
-      os << " capabilities: [";
+      os << opName << " capabilities: [";
       for (const auto &caps : capability.getCapabilities()) {
         os << " [";
         interleaveComma(caps, os, [&](spirv::Capability cap) {

mlir/test/Dialect/SPIRV/availability.mlir

@@ -29,3 +29,23 @@ func @subgroup_ballot(%predicate: i1) -> vector<4xi32> {
   %0 = spv.GroupNonUniformBallot "Workgroup" %predicate : vector<4xi32>
   return %0: vector<4xi32>
 }
+
+// CHECK-LABEL: module_logical_glsl450
+func @module_logical_glsl450() {
+  // CHECK: spv.module min version: V_1_0
+  // CHECK: spv.module max version: V_1_5
+  // CHECK: spv.module extensions: [ ]
+  // CHECK: spv.module capabilities: [ ]
+  spv.module "Logical" "GLSL450" { }
+  return
+}
+
+// CHECK-LABEL: module_physical_storage_buffer64_vulkan
+func @module_physical_storage_buffer64_vulkan() {
+  // CHECK: spv.module min version: V_1_5
+  // CHECK: spv.module max version: V_1_5
+  // CHECK: spv.module extensions: [ [SPV_EXT_physical_storage_buffer, SPV_KHR_physical_storage_buffer] ]
+  // CHECK: spv.module capabilities: [ [PhysicalStorageBufferAddresses] [VulkanMemoryModel] ]
+  spv.module "PhysicalStorageBuffer64" "Vulkan" { }
+  return
+}