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

[antiagainst]

Summary

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.

[merge_guards_bot]

Unit tests: pass. 61132 tests passed, 0 failed and 728 were skipped.

clang-tidy: fail. Please fix clang-tidy findings.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[mravishankar]

SPIR-V parts look good to me. Would still prefer the version/capability related methods into a separate file. They seem to have disjoint functionality (but its a personal preference).

[antiagainst]

In D71947#1801532, @mravishankar wrote:

Would still prefer the version/capability related methods into a separate file.

You mean in the SPIR-V dialect or the TableGen backend? For the former I'm open to suggestions given that they will stay as what is implemented here for the foreseeable future (minus minor changes maybe) so having a clean library structure makes sense. For the latter, I won't worry that much given that eventually they will be moved to a proper place in core.

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

 add_mlir_dialect(SPIRVOps SPIRVOps)
 
 set(LLVM_TARGET_DEFINITIONS SPIRVBase.td)
 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_tablegen(SPIRVOpAvailabilityImpl.inc -gen-spirv-avail-impls)
 add_public_tablegen_target(MLIRSPIRVAvailabilityIncGen)
 
 set(LLVM_TARGET_DEFINITIONS SPIRVOps.td)
 mlir_tablegen(SPIRVSerialization.inc -gen-spirv-serialization)

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

I32EnumAttr<"Capability", "valid SPIR-V Capability", [
       SPV_C_SubgroupAvcMotionEstimationChromaINTEL
     ]> {
   let cppNamespace = "::mlir::spirv";
 }
 
 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", [
       SPV_AM_Logical, SPV_AM_Physical32, SPV_AM_Physical64,
       SPV_AM_PhysicalStorageBuffer64
     ]> {
   let cppNamespace = "::mlir::spirv";
 }
BitEnumAttr<"MemoryAccess", "valid SPIR-V MemoryAccess", [
       SPV_MA_NonPrivatePointer
     ]> {
   let cppNamespace = "::mlir::spirv";
 }
 
 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", [
       SPV_MM_Simple, SPV_MM_GLSL450, SPV_MM_OpenCL, SPV_MM_Vulkan
     ]> {
   let cppNamespace = "::mlir::spirv";
 }
 

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

 
 #ifndef MLIR_DIALECT_SPIRV_SPIRVTYPES_H_
 #define MLIR_DIALECT_SPIRV_SPIRVTYPES_H_
 
 #include "mlir/IR/StandardTypes.h"
 #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>
 
 namespace mlir {
 namespace spirv {
 
 namespace detail {
 struct ArrayTypeStorage;

mlir/include/mlir/TableGen/Attribute.h

public:
   // Returns true if this EnumAttrCase is backed by a StringAttr.
   bool isStrCase() const;
 
   // Returns the symbol of this enum attribute case.
   StringRef getSymbol() const;
 
   // 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
 // in TableGen.This is used for enum attribute case backed by both StringAttr
 // and IntegerAttr.
 class EnumAttr : public Attribute {
 public:
   explicit EnumAttr(const llvm::Record *record);
   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;
 
   // Returns the enum class name.
   StringRef getEnumClassName() const;
 
   // Returns the C++ namespaces this enum class should be placed in.
   StringRef getCppNamespace() const;

mlir/lib/Dialect/SPIRV/CMakeLists.txt

add_llvm_library(MLIRSPIRV
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/SPIRV
   )
 
 add_dependencies(MLIRSPIRV
   MLIRSPIRVAvailabilityIncGen
   MLIRSPIRVCanonicalizationIncGen
+  MLIRSPIRVEnumAvailabilityIncGen
   MLIRSPIRVEnumsIncGen
   MLIRSPIRVOpsIncGen
   MLIRSPIRVOpUtilsGen
   MLIRSPIRVTargetAndABIIncGen)
 
 target_link_libraries(MLIRSPIRV
   MLIRIR
   MLIRParser
   MLIRSupport
   MLIRTransforms)
 
 add_subdirectory(Serialization)
 add_subdirectory(Transforms)

mlir/lib/Dialect/SPIRV/SPIRVTypes.cpp

 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringSwitch.h"
 
 using namespace mlir;
 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
 //===----------------------------------------------------------------------===//
 
 struct spirv::detail::ArrayTypeStorage : public TypeStorage {
   using KeyTy = std::tuple<Type, unsigned, ArrayType::LayoutInfo>;
 

mlir/lib/TableGen/Attribute.cpp

 StringRef tblgen::EnumAttrCase::getSymbol() const {
   return def->getValueAsString("symbol");
 }
 
 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");
 }
 
 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 {
   return def->getValueAsString("className");
 }
 
 StringRef tblgen::EnumAttr::getCppNamespace() const {
   return def->getValueAsString("cppNamespace");

mlir/test/Dialect/SPIRV/TestAvailability.cpp

void TestAvailability::runOnFunction() {
   llvm::outs() << f.getName() << "\n";
 
   Dialect *spvDialect = getContext().getRegisteredDialect("spv");
 
   f.getOperation()->walk([&](Operation *op) {
     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) {
           os << spirv::stringifyExtension(ext);
         });
         os << "]";
       }
       os << " ]\n";
     }
 
     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) {
           os << spirv::stringifyCapability(cap);
         });
         os << "]";
       }
       os << " ]\n";

mlir/test/Dialect/SPIRV/availability.mlir

 func @subgroup_ballot(%predicate: i1) -> vector<4xi32> {
   // CHECK: min version: V_1_3
   // CHECK: max version: V_1_5
   // CHECK: extensions: [ ]
   // CHECK: capabilities: [ [GroupNonUniformBallot] ]
   %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
+}

mlir/tools/mlir-tblgen/SPIRVUtilsGen.cpp

 #include "mlir/TableGen/Format.h"
 #include "mlir/TableGen/GenInfo.h"
 #include "mlir/TableGen/Operator.h"
 #include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/TableGen/Error.h"
 #include "llvm/TableGen/Record.h"
 #include "llvm/TableGen/TableGenBackend.h"
 
 using llvm::ArrayRef;
 using llvm::formatv;
 using llvm::raw_ostream;
 using llvm::raw_string_ostream;
 using llvm::Record;
 using llvm::RecordKeeper;
 using llvm::SmallVector;
 using llvm::SMLoc;
 using llvm::StringMap;
 using llvm::StringRef;
 using llvm::Twine;
 using mlir::tblgen::Attribute;
 using mlir::tblgen::EnumAttr;
+using mlir::tblgen::EnumAttrCase;
 using mlir::tblgen::NamedAttribute;
 using mlir::tblgen::NamedTypeConstraint;
 using mlir::tblgen::Operator;
 
 //===----------------------------------------------------------------------===//
 // Availability Wrapper Class
 //===----------------------------------------------------------------------===//
 

 StringRef Availability::getMergeInstanceType() const {
   return def->getValueAsString("instanceType");
 }
 
 StringRef Availability::getMergeInstance() const {
   return def->getValueAsString("instance");
 }
 
+// Returns the availability spec of the given `def`.
+std::vector<Availability> getAvailabilities(const Record &def) {
+  std::vector<Availability> availabilities;
+
+  if (def.getValue("availability")) {
+    std::vector<Record *> availDefs = def.getValueAsListOfDefs("availability");
+    availabilities.reserve(availDefs.size());
+    for (const Record *avail : availDefs)
+      availabilities.emplace_back(avail);
+  }
+
+  return availabilities;
+}
+
 //===----------------------------------------------------------------------===//
 // Availability Interface Definitions AutoGen
 //===----------------------------------------------------------------------===//
 
 static void emitInterfaceDef(const Availability &availability,
                              raw_ostream &os) {
   StringRef methodName = availability.getQueryFnName();
   os << availability.getQueryFnRetType() << " "

 static mlir::GenRegistration
     genInterfaceDefs("gen-avail-interface-defs",
                      "Generate op interface definitions",
                      [](const RecordKeeper &records, raw_ostream &os) {
                        return emitInterfaceDefs(records, os);
                      });
 
 //===----------------------------------------------------------------------===//
+// Enum Availability Query AutoGen
+//===----------------------------------------------------------------------===//
+
+static void emitAvailabilityQueryForIntEnum(const Record &enumDef,
+                                            raw_ostream &os) {
+  EnumAttr enumAttr(enumDef);
+  StringRef enumName = enumAttr.getEnumClassName();
+  std::vector<EnumAttrCase> enumerants = enumAttr.getAllCases();
+
+  // Mapping from availability class name to (enumerant, availablity
+  // specification) pairs.
+  llvm::StringMap<llvm::SmallVector<std::pair<EnumAttrCase, Availability>, 1>>
+      classCaseMap;
+
+  // Place all availablity specifications to their corresponding
+  // availablility classes.
+  for (const EnumAttrCase &enumerant : enumerants)
+    for (const Availability &avail : getAvailabilities(enumerant.getDef()))
+      classCaseMap[avail.getClass()].push_back({enumerant, avail});
+
+  for (const auto &classCasePair : classCaseMap) {
+    Availability avail = classCasePair.getValue().front().second;
+
+    os << formatv("llvm::Optional<{0}> {1}({2} value) {{\n",
+                  avail.getMergeInstanceType(), avail.getQueryFnName(),
+                  enumName);
+
+    os << "  switch (value) {\n";
+    for (const auto &caseSpecPair : classCasePair.getValue()) {
+      EnumAttrCase enumerant = caseSpecPair.first;
+      Availability avail = caseSpecPair.second;
+      os << formatv("  case {0}::{1}: return {2}({3});\n", enumName,
+                    enumerant.getSymbol(), avail.getMergeInstanceType(),
+                    avail.getMergeInstance());
+    }
+    os << "  default: break;\n";
+    os << "  }\n"
+       << "  return llvm::None;\n"
+       << "}\n";
+  }
+}
+
+static void emitAvailabilityQueryForBitEnum(const Record &enumDef,
+                                            raw_ostream &os) {
+  EnumAttr enumAttr(enumDef);
+  StringRef enumName = enumAttr.getEnumClassName();
+  std::string underlyingType = enumAttr.getUnderlyingType();
+  std::vector<EnumAttrCase> enumerants = enumAttr.getAllCases();
+
+  // Mapping from availability class name to (enumerant, availablity
+  // specification) pairs.
+  llvm::StringMap<llvm::SmallVector<std::pair<EnumAttrCase, Availability>, 1>>
+      classCaseMap;
+
+  // Place all availablity specifications to their corresponding
+  // availablility classes.
+  for (const EnumAttrCase &enumerant : enumerants)
+    for (const Availability &avail : getAvailabilities(enumerant.getDef()))
+      classCaseMap[avail.getClass()].push_back({enumerant, avail});
+
+  for (const auto &classCasePair : classCaseMap) {
+    Availability avail = classCasePair.getValue().front().second;
+
+    os << formatv("llvm::Optional<{0}> {1}({2} value) {{\n",
+                  avail.getMergeInstanceType(), avail.getQueryFnName(),
+                  enumName);
+
+    os << formatv(
+        "  assert(::llvm::countPopulation(static_cast<{0}>(value)) <= 1"
+        " && \"cannot have more than one bit set\");\n",
+        underlyingType);
+
+    os << "  switch (value) {\n";
+    for (const auto &caseSpecPair : classCasePair.getValue()) {
+      EnumAttrCase enumerant = caseSpecPair.first;
+      Availability avail = caseSpecPair.second;
+      os << formatv("  case {0}::{1}: return {2}({3});\n", enumName,
+                    enumerant.getSymbol(), avail.getMergeInstanceType(),
+                    avail.getMergeInstance());
+    }
+    os << "  default: break;\n";
+    os << "  }\n"
+       << "  return llvm::None;\n"
+       << "}\n";
+  }
+}
+
+static void emitEnumDecl(const Record &enumDef, raw_ostream &os) {
+  EnumAttr enumAttr(enumDef);
+  StringRef enumName = enumAttr.getEnumClassName();
+  StringRef cppNamespace = enumAttr.getCppNamespace();
+  auto enumerants = enumAttr.getAllCases();
+
+  llvm::SmallVector<StringRef, 2> namespaces;
+  llvm::SplitString(cppNamespace, namespaces, "::");
+
+  for (auto ns : namespaces)
+    os << "namespace " << ns << " {\n";
+
+  llvm::StringSet<> handledClasses;
+
+  // Place all availablity specifications to their corresponding
+  // availablility classes.
+  for (const EnumAttrCase &enumerant : enumerants)
+    for (const Availability &avail : getAvailabilities(enumerant.getDef())) {
+      StringRef className = avail.getClass();
+      if (handledClasses.count(className))
+        continue;
+      os << formatv("llvm::Optional<{0}> {1}({2} value);\n",
+                    avail.getMergeInstanceType(), avail.getQueryFnName(),
+                    enumName);
+      handledClasses.insert(className);
+    }
+
+  for (auto ns : llvm::reverse(namespaces))
+    os << "} // namespace " << ns << "\n";
+}
+
+static bool emitEnumDecls(const RecordKeeper &recordKeeper, raw_ostream &os) {
+  llvm::emitSourceFileHeader("SPIR-V Enum Availability Declarations", os);
+
+  auto defs = recordKeeper.getAllDerivedDefinitions("EnumAttrInfo");
+  for (const auto *def : defs)
+    emitEnumDecl(*def, os);
+
+  return false;
+}
+
+static void emitEnumDef(const Record &enumDef, raw_ostream &os) {
+  EnumAttr enumAttr(enumDef);
+  StringRef cppNamespace = enumAttr.getCppNamespace();
+
+  llvm::SmallVector<StringRef, 2> namespaces;
+  llvm::SplitString(cppNamespace, namespaces, "::");
+
+  for (auto ns : namespaces)
+    os << "namespace " << ns << " {\n";
+
+  if (enumAttr.isBitEnum()) {
+    emitAvailabilityQueryForBitEnum(enumDef, os);
+  } else {
+    emitAvailabilityQueryForIntEnum(enumDef, os);
+  }
+
+  for (auto ns : llvm::reverse(namespaces))
+    os << "} // namespace " << ns << "\n";
+  os << "\n";
+}
+
+static bool emitEnumDefs(const RecordKeeper &recordKeeper, raw_ostream &os) {
+  llvm::emitSourceFileHeader("SPIR-V Enum Availability Definitions", os);
+
+  auto defs = recordKeeper.getAllDerivedDefinitions("EnumAttrInfo");
+  for (const auto *def : defs)
+    emitEnumDef(*def, os);
+
+  return false;
+}
+
+//===----------------------------------------------------------------------===//
+// Enum Availability Query Hook Registration
+//===----------------------------------------------------------------------===//
+
+// Registers the enum utility generator to mlir-tblgen.
+static mlir::GenRegistration
+    genEnumDecls("gen-spirv-enum-avail-decls",
+                 "Generate SPIR-V enum availability declarations",
+                 [](const RecordKeeper &records, raw_ostream &os) {
+                   return emitEnumDecls(records, os);
+                 });
+
+// Registers the enum utility generator to mlir-tblgen.
+static mlir::GenRegistration
+    genEnumDefs("gen-spirv-enum-avail-defs",
+                "Generate SPIR-V enum availability definitions",
+                [](const RecordKeeper &records, raw_ostream &os) {
+                  return emitEnumDefs(records, os);
+                });
+
+//===----------------------------------------------------------------------===//
 // Serialization AutoGen
 //===----------------------------------------------------------------------===//
 
 // Writes the following function to `os`:
 //   inline uint32_t getOpcode(<op-class-name>) { return <opcode>; }
 static void emitGetOpcodeFunction(const Record *record, Operator const &op,
                                   raw_ostream &os) {
   os << formatv("template <> constexpr inline ::mlir::spirv::Opcode "
genOpUtils("gen-spirv-op-utils",
                [](const RecordKeeper &records, raw_ostream &os) {
                  return emitOpUtils(records, os);
                });
 
 //===----------------------------------------------------------------------===//
 // SPIR-V Availability Impl AutoGen
 //===----------------------------------------------------------------------===//
 
-// Returns the availability spec of the given `def`.
-std::vector<Availability> getAvailabilities(const Record &def) {
-  std::vector<Availability> availabilities;
-  if (auto *availListInit = def.getValueAsListInit("availability")) {
-    availabilities.reserve(availListInit->size());
-    for (auto *availInit : *availListInit)
-      availabilities.emplace_back(
-          llvm::cast<llvm::DefInit>(availInit)->getDef());
-  }
-  return availabilities;
-}
-
 static void emitAvailabilityImpl(const Operator &srcOp, raw_ostream &os) {
   mlir::tblgen::FmtContext fctx;
   fctx.addSubst("overall", "overall");
 
   std::vector<Availability> opAvailabilities =
       getAvailabilities(srcOp.getDef());
 
   // First collect all availablity classes this op should implement.
   // All availablity instances keep information for the generated interface and
   // the instance's specific requirement. Here we remember a random instance so
   // we can get the information regarding the generated interface.
   llvm::StringMap<Availability> availClasses;
   for (const Availability &avail : opAvailabilities)
     availClasses.try_emplace(avail.getClass(), avail);
+  for (const NamedAttribute &namedAttr : srcOp.getAttributes()) {
+    const auto *enumAttr = llvm::dyn_cast<EnumAttr>(&namedAttr.attr);
+    if (!enumAttr)
+      continue;
+
+    for (const EnumAttrCase &enumerant : enumAttr->getAllCases())
+      for (const Availability &caseAvail :
+           getAvailabilities(enumerant.getDef()))
+        availClasses.try_emplace(caseAvail.getClass(), caseAvail);
+  }
 
   // Then generate implementation for each availability class.
   for (const auto &availClass : availClasses) {
     StringRef availClassName = availClass.getKey();
     Availability avail = availClass.getValue();
 
     // Generate the implementation method signature.
     os << formatv("{0} {1}::{2}() {{\n", avail.getQueryFnRetType(),
                   srcOp.getCppClassName(), avail.getQueryFnName());
 
     // Create the variable for the final requirement and initialize it.
     os << formatv("  {0} overall = {1};\n", avail.getQueryFnRetType(),
                   avail.getMergeInitializer());
 
     // Update with the op's specific availability spec.
     for (const Availability &avail : opAvailabilities)
       if (avail.getClass() == availClassName) {
         os << "  "
            << tgfmt(avail.getMergeActionCode(),
                     &fctx.addSubst("instance", avail.getMergeInstance()))
            << ";\n";
       }
+
+    // Update with enum attributes' specific availability spec.
+    for (const NamedAttribute &namedAttr : srcOp.getAttributes()) {
+      const auto *enumAttr = llvm::dyn_cast<EnumAttr>(&namedAttr.attr);
+      if (!enumAttr)
+        continue;
+
+      // (enumerant, availablity specification) pairs for this availability
+      // class.
+      SmallVector<std::pair<EnumAttrCase, Availability>, 1> caseSpecs;
+
+      // Collect all cases' availability specs.
+      for (const EnumAttrCase &enumerant : enumAttr->getAllCases())
+        for (const Availability &caseAvail :
+             getAvailabilities(enumerant.getDef()))
+          if (availClassName == caseAvail.getClass())
+            caseSpecs.push_back({enumerant, caseAvail});
+
+      // If this attribute kind does not have any availablity spec from any of
+      // its cases, no more work to do.
+      if (caseSpecs.empty())
+        continue;
+
+      if (enumAttr->isBitEnum()) {
+        // For BitEnumAttr, we need to iterate over each bit to query its
+        // availability spec.
+        os << formatv("  for (unsigned i = 0; "
+                      "i < std::numeric_limits<{0}>::digits; ++i) {{\n",
+                      enumAttr->getUnderlyingType());
+        os << formatv("    {0}::{1} attrVal = this->{2}() & "
+                      "static_cast<{0}::{1}>(1 << i);\n",
+                      enumAttr->getCppNamespace(), enumAttr->getEnumClassName(),
+                      namedAttr.name);
+        os << formatv("    if (static_cast<{0}>(attrVal) == 0) continue;\n",
+                      enumAttr->getUnderlyingType());
+      } else {
+        // For IntEnumAttr, we just need to query the value as a whole.
+        os << "  {\n";
+        os << formatv("    auto attrVal = this->{0}();\n", namedAttr.name);
+      }
+      os << formatv("    auto instance = {0}::{1}(attrVal);\n",
+                    enumAttr->getCppNamespace(), avail.getQueryFnName());
+      os << "    if (instance) "
+         // TODO(antiagainst): use `avail.getMergeCode()` here once ODS supports
+         // dialect-specific contents so that we can use not implementing the
+         // availability interface as indication of no requirements.
+         << tgfmt(caseSpecs.front().second.getMergeActionCode(),
+                  &fctx.addSubst("instance", "*instance"))
+         << ";\n";
+      os << "  }\n";
+    }
+
     os << "  return overall;\n";
     os << "}\n";
   }
 }
 
 static bool emitAvailabilityImpl(const RecordKeeper &recordKeeper,
                                  raw_ostream &os) {
   llvm::emitSourceFileHeader("SPIR-V Op Availability Implementations", os);