[mlir][spirv] Use SmallVector<ArrayRef> for availability queries

[antiagainst]

Summary

Previously extensions and capabilities requirements are returned as
SmallVector<SmallVector>. It's an anti-pattern; this commit improves
a bit by returning as SmallVector<ArrayRef>. This is possible because
the internal sequence is always known statically (from the spec)
so that we can use a static constant array for it and get an ArrayRef.

Depends On D75872

[mravishankar]

Do we need to autogen these? Can we have a simple utility method instead of generating this. Its a bit hard to parse.

mlir/include/mlir/Dialect/SPIRV/SPIRVAvailability.td

50

nit : s/instnace/instance

[antiagainst]

Address comments

[antiagainst]

In D75874#1915466, @mravishankar wrote:

Do we need to autogen these? Can we have a simple utility method instead of generating this. Its a bit hard to parse.

Agreed that it's a bit hard to parse, but note that this is expected to be general so later other dialects can also use the availability mechanism if they want. I plan to update the doc to explain a bit afterwards.

mlir/include/mlir/Dialect/SPIRV/SPIRVAvailability.td

class Availability {
   code mergeAction = ?;
   // The initializer for the final availability requirement.
   string initializer = ?;
   // An availability instance's type.
   string instanceType = ?;
 
   // The following are fields for a concrete availability instance.
 
+  // The code for preparing a concrete instance. This should be C++ statements
+  // and will be generated before the `mergeAction` logic.
+  code instancePreparation = "";
   // The availability requirement carried by a concrete instance.
   string instance = ?;
 }
 
 class MinVersionBase<string name, I32EnumAttr scheme, I32EnumAttrCase min>
     : Availability {
   let interfaceName = name;
 

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

let interfaceDescription = [{
     ```
     {{Extension::A, Extension::B}, {Extension::C}, {{Extension::D, Extension::E}}
     ```
 
     The operation instance is available when (`Extension::A` OR `Extension::B`)
     AND (`Extension::C`) AND (`Extension::D` OR `Extension::E`) is enabled.
   }];
 
-  // TODO(antiagainst): Using SmallVector<SmallVector<...>> is an anti-pattern.
+  // TODO(antiagainst): Returning SmallVector<ArrayRef<...>> is not recommended.
   // Find a better way for this.
-  let queryFnRetType = "::llvm::SmallVector<::llvm::SmallVector<"
-                          "::mlir::spirv::Extension, 1>, 1>";
+  let queryFnRetType = "::llvm::SmallVector<::llvm::ArrayRef<"
+                          "::mlir::spirv::Extension>, 1>";
   let queryFnName = "getExtensions";
 
   let mergeAction = !if(
       !empty(extensions), "", "$overall.emplace_back($instance)");
   let initializer = "{}";
-  let instanceType = "::llvm::SmallVector<::mlir::spirv::Extension, 1>";
+  let instanceType = "::llvm::ArrayRef<::mlir::spirv::Extension>";
 
-  // Compose all capabilities as an C++ initializer list
-  let instance = "std::initializer_list<::mlir::spirv::Extension>{" #
-                 StrJoin<!foreach(
-                   ext, extensions,
-                   "::mlir::spirv::Extension::" # ext.symbol)>.result #
-                 "}";
+  // Pack all extensions as a static array and get its reference.
+  let instancePreparation = !if(!empty(extensions), "",
+    "static const ::mlir::spirv::Extension exts[] = {" #
+    StrJoin<!foreach(ext, extensions,
+    "}; " #
+    // The following manual ArrayRef constructor call is to satisfy GCC 5.
+    "ArrayRef<::mlir::spirv::Extension> " #
+      "ref(exts, ::llvm::array_lengthof(exts));");
+  let instance = "ref";
 }
 
 class Capability<list<I32EnumAttrCase> capabilities> : Availability {
   let interfaceName = "QueryCapabilityInterface";
   let interfaceDescription = [{
     Querying interface for required SPIR-V capabilities.
 
     This interface provides a `getCapabilities()` method to query the required
     capabilities for the implementing SPIR-V operation. The returned value
     is a nested vector whose element is `mlir::spirv::Capability`s. The outer
     vector's elements (which are vectors) should be interpreted as conjunction
     while the inner vector's elements (which are `mlir::spirv::Capability`s)
     should be interpreted as disjunction. For example, given
 
     ```
     {{Capability::A, Capability::B}, {Capability::C}, {{Capability::D, Capability::E}}
     ```
 
     The operation instance is available when (`Capability::A` OR `Capability::B`)
     AND (`Capability::C`) AND (`Capability::D` OR `Capability::E`) is enabled.
   }];
 
-  let queryFnRetType = "::llvm::SmallVector<::llvm::SmallVector<"
-                          "::mlir::spirv::Capability, 1>, 1>";
+  let queryFnRetType = "::llvm::SmallVector<::llvm::ArrayRef<"
+                          "::mlir::spirv::Capability>, 1>";
   let queryFnName = "getCapabilities";
 
   let mergeAction = !if(
       !empty(capabilities), "", "$overall.emplace_back($instance)");
   let initializer = "{}";
-  let instanceType = "::llvm::SmallVector<::mlir::spirv::Capability, 1>";
+  let instanceType = "::llvm::ArrayRef<::mlir::spirv::Capability>";
 
-  // Compose all capabilities as an C++ initializer list
-  let instance = "std::initializer_list<::mlir::spirv::Capability>{" #
-                 StrJoin<!foreach(
-                   cap, capabilities,
-                   "::mlir::spirv::Capability::" # cap.symbol)>.result #
-                 "}";
+  // Pack all capabilities as a static array and get its reference.
+  let instancePreparation = !if(!empty(capabilities), "",
+    "static const ::mlir::spirv::Capability caps[] = {" #
+    StrJoin<!foreach(cap, capabilities,
+    "}; " #
+    // The following manual ArrayRef constructor call is to satisfy GCC 5.
+    "ArrayRef<::mlir::spirv::Capability> " #
+      "ref(caps, ::llvm::array_lengthof(caps));");
+  let instance = "ref";
 }
 
 // TODO(antiagainst): the following interfaces definitions are duplicating with
 // the above. Remove them once we are able to support dialect-specific contents
 // in ODS.
 def QueryMinVersionInterface : OpInterface<"QueryMinVersionInterface"> {
   let methods = [InterfaceMethod<"", "::mlir::spirv::Version", "getMinVersion">];
 }
 def QueryMaxVersionInterface : OpInterface<"QueryMaxVersionInterface"> {
   let methods = [InterfaceMethod<"", "::mlir::spirv::Version", "getMaxVersion">];
 }
 def QueryExtensionInterface : OpInterface<"QueryExtensionInterface"> {
   let methods = [InterfaceMethod<
     "",
-    "::llvm::SmallVector<::llvm::SmallVector<::mlir::spirv::Extension, 1>, 1>",
+    "::llvm::SmallVector<::llvm::ArrayRef<::mlir::spirv::Extension>, 1>",
     "getExtensions">];
 }
 def QueryCapabilityInterface : OpInterface<"QueryCapabilityInterface"> {
   let methods = [InterfaceMethod<
     "",
-    "::llvm::SmallVector<::llvm::SmallVector<::mlir::spirv::Capability, 1>, 1>",
+    "::llvm::SmallVector<::llvm::ArrayRef<::mlir::spirv::Capability>, 1>",
     "getCapabilities">];
 }
 
 //===----------------------------------------------------------------------===//
 // SPIR-V extension definitions
 //===----------------------------------------------------------------------===//
 
 // Extensions known to the SPIR-V dialect.

mlir/tools/mlir-tblgen/SPIRVUtilsGen.cpp

public:
 
   // Returns the initializer expression for initializing the final availability
   // requirements.
   StringRef getMergeInitializer() const;
 
   // Returns the C++ type for an availability instance.
   StringRef getMergeInstanceType() const;
 
+  // Returns the C++ statements for preparing availability instance.
+  StringRef getMergeInstancePreparation() const;
+
   // Returns the concrete availability instance carried in this case.
   StringRef getMergeInstance() const;
 
 private:
   // The TableGen definition of this availability.
   const llvm::Record *def;
 };
 } // namespace

 StringRef Availability::getMergeInitializer() const {
   return def->getValueAsString("initializer");
 }
 
 StringRef Availability::getMergeInstanceType() const {
   return def->getValueAsString("instanceType");
 }
 
+StringRef Availability::getMergeInstancePreparation() const {
+  return def->getValueAsString("instancePreparation");
+}
+
 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;
 
for (const auto &classCasePair : classCaseMap) {
     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 << formatv("  case {0}::{1}: { {2} return {3}({4}); }\n", enumName,
+                    enumerant.getSymbol(), avail.getMergeInstancePreparation(),
+                    avail.getMergeInstanceType(), avail.getMergeInstance());
     }
     // Only emit default if uncovered cases.
     if (classCasePair.getValue().size() < enumAttr.getAllCases().size())
       os << "  default: break;\n";
     os << "  }\n"
        << "  return llvm::None;\n"
        << "}\n";
   }
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 << formatv("  case {0}::{1}: { {2} return {3}({4}); }\n", enumName,
+                    enumerant.getSymbol(), avail.getMergeInstancePreparation(),
+                    avail.getMergeInstanceType(), avail.getMergeInstance());
     }
     os << "  default: break;\n";
     os << "  }\n"
        << "  return llvm::None;\n"
        << "}\n";
   }
 }
 
genOpUtils("gen-spirv-op-utils",
                });
 
 //===----------------------------------------------------------------------===//
 // SPIR-V Availability Impl AutoGen
 //===----------------------------------------------------------------------===//
 
 static void emitAvailabilityImpl(const Operator &srcOp, raw_ostream &os) {
   mlir::tblgen::FmtContext fctx;
-  fctx.addSubst("overall", "overall");
+  fctx.addSubst("overall", "tblgen_overall");
 
   std::vector<Availability> opAvailabilities =
       getAvailabilities(srcOp.getDef());
 
   // First collect all availability classes this op should implement.
   // All availability 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.
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(),
+    os << formatv("  {0} tblgen_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 << "  "
+      if (avail.getClass() == availClassName &&
+          (!avail.getMergeInstancePreparation().empty() ||
+           !avail.getMergeActionCode().empty())) {
+        os << "  {\n    "
+           // Prepare this instance.
+           << avail.getMergeInstancePreparation()
+           << "\n    "
+           // Merge this instance.
            << std::string(
                   tgfmt(avail.getMergeActionCode(),
                         &fctx.addSubst("instance", avail.getMergeInstance())))
-           << ";\n";
+           << ";\n  }\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;
 
for (const NamedAttribute &namedAttr : srcOp.getAttributes()) {
         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}() & "
+        os << formatv("    {0}::{1} tblgen_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());
+        os << formatv(
+            "    if (static_cast<{0}>(tblgen_attrVal) == 0) continue;\n",
       } 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 tblgen_attrVal = this->{0}();\n",
+                      namedAttr.name);
       }
-      os << formatv("    auto instance = {0}::{1}(attrVal);\n",
+      os << formatv("    auto tblgen_instance = {0}::{1}(tblgen_attrVal);\n",
                     enumAttr->getCppNamespace(), avail.getQueryFnName());
-      os << "    if (instance) "
+      os << "    if (tblgen_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.
          << std::string(tgfmt(caseSpecs.front().second.getMergeActionCode(),
-                              &fctx.addSubst("instance", "*instance")))
+                              &fctx.addSubst("instance", "*tblgen_instance")))
          << ";\n";
       os << "  }\n";
     }
 
-    os << "  return overall;\n";
+    os << "  return tblgen_overall;\n";
     os << "}\n";
   }
 }
 
 static bool emitAvailabilityImpl(const RecordKeeper &recordKeeper,
                                  raw_ostream &os) {
   llvm::emitSourceFileHeader("SPIR-V Op Availability Implementations", os);