[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

4 files changed, 36 insertions, 2 deletions

diff --git a/mlir/include/mlir/Dialect/SPIRV/CMakeLists.txt b/mlir/include/mlir/Dialect/SPIRV/CMakeLists.txt
index cb32fc5fedb..8d49b4759e4 100644
--- a/mlir/include/mlir/Dialect/SPIRV/CMakeLists.txt
+++ b/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)
diff --git a/mlir/include/mlir/Dialect/SPIRV/SPIRVBase.td b/mlir/include/mlir/Dialect/SPIRV/SPIRVBase.td
index dd2743504d5..96268e80fff 100644
--- a/mlir/include/mlir/Dialect/SPIRV/SPIRVBase.td
+++ b/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", [
diff --git a/mlir/include/mlir/Dialect/SPIRV/SPIRVTypes.h b/mlir/include/mlir/Dialect/SPIRV/SPIRVTypes.h
index 001d3130778..548ab83915c 100644
--- a/mlir/include/mlir/Dialect/SPIRV/SPIRVTypes.h
+++ b/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>
 
diff --git a/mlir/include/mlir/TableGen/Attribute.h b/mlir/include/mlir/TableGen/Attribute.h
index 747df945cea..9f1d8b392ae 100644
--- a/mlir/include/mlir/TableGen/Attribute.h
+++ b/mlir/include/mlir/TableGen/Attribute.h
@@ -135,6 +135,9 @@ public:
 
   // 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 @@ public:
   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;