[mlir][spirv] Add basic definitions for supporting availability

SPIR-V has a few mechanisms to control op availability: version,
extension, and capabilities. These mechanisms are considered as
different availability classes.

This commit introduces basic definitions for modelling SPIR-V
availability classes. Specifically, an `Availability` class is
added to SPIRVBase.td, along with two subclasses: MinVersion
and MaxVersion for versioning. SPV_Op is extended to take a
list of `Availability`. Each `Availability` instance carries
information for generating op interfaces for the corresponding
availability class and also the concrete availability
requirements.

With the availability spec on ops, we can now auto-generate the
op interfaces of all SPIR-V availability classes and also
synthesize the op's implementations of these interfaces. The
interface generation is done via new TableGen backends
-gen-avail-interface-{decls|defs}. The op's implementation is
done via -gen-spirv-avail-impls.

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

6 files changed, 278 insertions, 15 deletions

diff --git a/mlir/include/mlir/Dialect/SPIRV/CMakeLists.txt b/mlir/include/mlir/Dialect/SPIRV/CMakeLists.txt
index fc7180de6cb..52464789439 100644
--- a/mlir/include/mlir/Dialect/SPIRV/CMakeLists.txt
+++ b/mlir/include/mlir/Dialect/SPIRV/CMakeLists.txt
@@ -1,8 +1,3 @@
-set(LLVM_TARGET_DEFINITIONS SPIRVLowering.td)
-mlir_tablegen(SPIRVLowering.h.inc -gen-struct-attr-decls)
-mlir_tablegen(SPIRVLowering.cpp.inc -gen-struct-attr-defs)
-add_public_tablegen_target(MLIRSPIRVLoweringStructGen)
-
 add_mlir_dialect(SPIRVOps SPIRVOps)
 
 set(LLVM_TARGET_DEFINITIONS SPIRVBase.td)
@@ -11,9 +6,20 @@ mlir_tablegen(SPIRVEnums.cpp.inc -gen-enum-defs)
 add_public_tablegen_target(MLIRSPIRVEnumsIncGen)
 
 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)
 add_public_tablegen_target(MLIRSPIRVSerializationGen)
 
 set(LLVM_TARGET_DEFINITIONS SPIRVBase.td)
 mlir_tablegen(SPIRVOpUtils.inc -gen-spirv-op-utils)
 add_public_tablegen_target(MLIRSPIRVOpUtilsGen)
+
+set(LLVM_TARGET_DEFINITIONS SPIRVLowering.td)
+mlir_tablegen(SPIRVLowering.h.inc -gen-struct-attr-decls)
+mlir_tablegen(SPIRVLowering.cpp.inc -gen-struct-attr-defs)
+add_public_tablegen_target(MLIRSPIRVLoweringStructGen)
diff --git a/mlir/include/mlir/Dialect/SPIRV/SPIRVAtomicOps.td b/mlir/include/mlir/Dialect/SPIRV/SPIRVAtomicOps.td
index c2ea100c121..17be79dfcfd 100644
--- a/mlir/include/mlir/Dialect/SPIRV/SPIRVAtomicOps.td
+++ b/mlir/include/mlir/Dialect/SPIRV/SPIRVAtomicOps.td
@@ -120,6 +120,13 @@ def SPV_AtomicCompareExchangeWeakOp : SPV_Op<"AtomicCompareExchangeWeak", []> {
     ```
   }];
 
+  let availability = [
+    MinVersion<SPV_V_1_0>,
+    MaxVersion<SPV_V_1_3>,
+    Extension<[]>,
+    Capability<[SPV_C_Kernel]>
+  ];
+
   let arguments = (ins
     SPV_AnyPtr:$pointer,
     SPV_ScopeAttr:$memory_scope,
diff --git a/mlir/include/mlir/Dialect/SPIRV/SPIRVAvailability.td b/mlir/include/mlir/Dialect/SPIRV/SPIRVAvailability.td
new file mode 100644
index 00000000000..8ec74ac955a
--- /dev/null
+++ b/mlir/include/mlir/Dialect/SPIRV/SPIRVAvailability.td
@@ -0,0 +1,86 @@
+//===- SPIRVAvailability.td - Op Availability Base file ----*- tablegen -*-===//
+//
+// Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SPIRV_AVAILABILITY
+#define SPIRV_AVAILABILITY
+
+include "mlir/IR/OpBase.td"
+
+//===----------------------------------------------------------------------===//
+// Op availaility definitions
+//===----------------------------------------------------------------------===//
+
+// The base class for defining op availability dimensions.
+class Availability {
+  // The following are fields for controlling the generated C++ OpInterface.
+
+  // The name for the generated C++ OpInterface subclass.
+  string interfaceName = ?;
+  // The documentation for the generated C++ OpInterface subclass.
+  string interfaceDescription = "";
+
+  // The following are fields for controlling the query function signature.
+
+  // The query function's return type in the generated C++ OpInterface subclass.
+  string queryFnRetType = ?;
+  // The query function's name in the generated C++ OpInterface subclass.
+  string queryFnName = ?;
+
+  // The following are fields for controlling the query function implementation.
+
+  // The logic for merging two availability requirements. This is used to derive
+  // the final availability requirement when, for example, an op has two
+  // operands and these two operands have different availability requirements.
+  //
+  // The code should use `$overall` as the placeholder for the final requirement
+  // and `$instance` for the current availability requirement instance.
+  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 availability requirement carried by a concrete instance.
+  string instance = ?;
+}
+
+class MinVersionBase<string name, I32EnumAttr scheme, I32EnumAttrCase min>
+    : Availability {
+  let interfaceName = name;
+
+  let queryFnRetType = scheme.returnType;
+  let queryFnName = "getMinVersion";
+
+  let mergeAction = "$overall = static_cast<" # scheme.returnType # ">("
+                      "std::max($overall, $instance))";
+  let initializer = "static_cast<" # scheme.returnType # ">(uint32_t(0))";
+  let instanceType = scheme.cppNamespace # "::" # scheme.className;
+
+  let instance = scheme.cppNamespace # "::" # scheme.className # "::" #
+                 min.symbol;
+}
+
+class MaxVersionBase<string name, I32EnumAttr scheme, I32EnumAttrCase max>
+    : Availability {
+  let interfaceName = name;
+
+  let queryFnRetType = scheme.returnType;
+  let queryFnName = "getMaxVersion";
+
+  let mergeAction = "$overall = static_cast<" # scheme.returnType # ">("
+                      "std::min($overall, $instance))";
+  let initializer = "static_cast<" # scheme.returnType # ">(~uint32_t(0))";
+  let instanceType = scheme.cppNamespace # "::" # scheme.className;
+
+  let instance = scheme.cppNamespace # "::" # scheme.className # "::" #
+                 max.symbol;
+}
+
+#endif // SPIRV_AVAILABILITY
diff --git a/mlir/include/mlir/Dialect/SPIRV/SPIRVBase.td b/mlir/include/mlir/Dialect/SPIRV/SPIRVBase.td
index 5751a32e169..acbbbfc296b 100644
--- a/mlir/include/mlir/Dialect/SPIRV/SPIRVBase.td
+++ b/mlir/include/mlir/Dialect/SPIRV/SPIRVBase.td
@@ -16,6 +16,7 @@
 #define SPIRV_BASE
 
 include "mlir/IR/OpBase.td"
+include "mlir/Dialect/SPIRV/SPIRVAvailability.td"
 
 //===----------------------------------------------------------------------===//
 // SPIR-V dialect definitions
@@ -46,6 +47,142 @@ def SPV_Dialect : Dialect {
 }
 
 //===----------------------------------------------------------------------===//
+// SPIR-V availability definitions
+//===----------------------------------------------------------------------===//
+
+def SPV_V_1_0 : I32EnumAttrCase<"V_1_0", 0>;
+def SPV_V_1_1 : I32EnumAttrCase<"V_1_1", 1>;
+def SPV_V_1_2 : I32EnumAttrCase<"V_1_2", 2>;
+def SPV_V_1_3 : I32EnumAttrCase<"V_1_3", 3>;
+def SPV_V_1_4 : I32EnumAttrCase<"V_1_4", 4>;
+def SPV_V_1_5 : I32EnumAttrCase<"V_1_5", 5>;
+
+def SPV_VersionAttr : I32EnumAttr<"Version", "valid SPIR-V version", [
+    SPV_V_1_0, SPV_V_1_1, SPV_V_1_2, SPV_V_1_3, SPV_V_1_4, SPV_V_1_5]> {
+  let cppNamespace = "::mlir::spirv";
+}
+
+class MinVersion<I32EnumAttrCase min> : MinVersionBase<
+    "QueryMinVersionInterface", SPV_VersionAttr, min> {
+  let interfaceDescription = [{
+    Querying interface for minimal required SPIR-V version.
+
+    This interface provides a `getMinVersion()` method to query the minimal
+    required version for the implementing SPIR-V operation. The returned value
+    is a `mlir::spirv::Version` enumerant.
+  }];
+}
+
+class MaxVersion<I32EnumAttrCase max> : MaxVersionBase<
+    "QueryMaxVersionInterface", SPV_VersionAttr, max> {
+  let interfaceDescription = [{
+    Querying interface for maximal supported SPIR-V version.
+
+    This interface provides a `getMaxVersion()` method to query the maximal
+    supported version for the implementing SPIR-V operation. The returned value
+    is a `mlir::spirv::Version` enumerant.
+  }];
+}
+
+class Extension<list<StrEnumAttrCase> extensions> : Availability {
+  let interfaceName = "QueryExtensionInterface";
+  let interfaceDescription = [{
+    Querying interface for required SPIR-V extensions.
+
+    This interface provides a `getExtensions()` method to query the required
+    extensions for the implementing SPIR-V operation. The returned value
+    is a nested vector whose element is `mlir::spirv::Extension`s. The outer
+    vector's elements (which are vectors) should be interpreted as conjunction
+    while the innner vector's elements (which are `mlir::spirv::Extension`s)
+    should be interpreted as disjunction. For example, given
+
+    ```
+    {{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.
+  // Find a better way for this.
+  let queryFnRetType = "::llvm::SmallVector<::llvm::SmallVector<"
+                          "::mlir::spirv::Extension, 1>, 1>";
+  let queryFnName = "getExtensions";
+
+  let mergeAction = !if(
+      !empty(extensions), "", "$overall.emplace_back($instance)");
+  let initializer = "{}";
+  let instanceType = "::llvm::SmallVector<::mlir::spirv::Extension, 1>";
+
+  // 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 #
+                 "}";
+}
+
+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 neted vector whose element is `mlir::spirv::Capability`s. The outer
+    vector's elements (which are vectors) should be interpreted as conjunction
+    while the innner 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 queryFnName = "getCapabilities";
+
+  let mergeAction = !if(
+      !empty(capabilities), "", "$overall.emplace_back($instance)");
+  let initializer = "{}";
+  let instanceType = "::llvm::SmallVector<::mlir::spirv::Capability, 1>";
+
+  // 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 #
+                 "}";
+}
+
+// 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>",
+    "getExtensions">];
+}
+def QueryCapabilityInterface : OpInterface<"QueryCapabilityInterface"> {
+  let methods = [InterfaceMethod<
+    "",
+    "::llvm::SmallVector<::llvm::SmallVector<::mlir::spirv::Capability, 1>, 1>",
+    "getCapabilities">];
+}
+
+//===----------------------------------------------------------------------===//
 // SPIR-V extension definitions
 //===----------------------------------------------------------------------===//
 
@@ -1216,7 +1353,22 @@ def SPV_OpcodeAttr :
 
 // Base class for all SPIR-V ops.
 class SPV_Op<string mnemonic, list<OpTrait> traits = []> :
-    Op<SPV_Dialect, mnemonic, traits> {
+    Op<SPV_Dialect, mnemonic, !listconcat(traits, [
+         // TODO(antiagainst): We don't need all of the following traits for
+         // every op; only the suitabble ones should be added automatically
+         // after ODS supports dialect-specific contents.
+         DeclareOpInterfaceMethods<QueryMinVersionInterface>,
+         DeclareOpInterfaceMethods<QueryMaxVersionInterface>,
+         DeclareOpInterfaceMethods<QueryExtensionInterface>,
+         DeclareOpInterfaceMethods<QueryCapabilityInterface>
+       ])> {
+  // Availability specification for this op itself.
+  list<Availability> availability = [
+    MinVersion<SPV_V_1_0>,
+    MaxVersion<SPV_V_1_5>,
+    Extension<[]>,
+    Capability<[]>
+  ];
 
   // For each SPIR-V op, the following static functions need to be defined
   // in SPVOps.cpp:
diff --git a/mlir/include/mlir/Dialect/SPIRV/SPIRVNonUniformOps.td b/mlir/include/mlir/Dialect/SPIRV/SPIRVNonUniformOps.td
index f3a9a61a9e9..1ac0ae1b969 100644
--- a/mlir/include/mlir/Dialect/SPIRV/SPIRVNonUniformOps.td
+++ b/mlir/include/mlir/Dialect/SPIRV/SPIRVNonUniformOps.td
@@ -53,6 +53,13 @@ def SPV_GroupNonUniformBallotOp : SPV_Op<"GroupNonUniformBallot", []> {
     ```
   }];
 
+  let availability = [
+    MinVersion<SPV_V_1_3>,
+    MaxVersion<SPV_V_1_5>,
+    Extension<[]>,
+    Capability<[SPV_C_GroupNonUniformBallot]>
+  ];
+
   let arguments = (ins
     SPV_ScopeAttr:$execution_scope,
     SPV_Bool:$predicate
diff --git a/mlir/include/mlir/Dialect/SPIRV/SPIRVOps.h b/mlir/include/mlir/Dialect/SPIRV/SPIRVOps.h
index 2fa417bfe25..3806418593f 100644
--- a/mlir/include/mlir/Dialect/SPIRV/SPIRVOps.h
+++ b/mlir/include/mlir/Dialect/SPIRV/SPIRVOps.h
@@ -21,18 +21,23 @@ class OpBuilder;
 
 namespace spirv {
 
+// TableGen'erated operation interfaces for querying versions, extensions, and
+// capabilities.
+#include "mlir/Dialect/SPIRV/SPIRVAvailability.h.inc"
+
+// TablenGen'erated operation declarations.
 #define GET_OP_CLASSES
 #include "mlir/Dialect/SPIRV/SPIRVOps.h.inc"
 
-/// Following methods are auto-generated.
-///
-/// Get the name used in the Op to refer to an enum value of the given
-/// `EnumClass`.
-/// template <typename EnumClass> StringRef attributeName();
-///
-/// Get the function that can be used to symbolize an enum value.
-/// template <typename EnumClass>
-/// Optional<EnumClass> (*)(StringRef) symbolizeEnum();
+// TableGen'erated helper functions.
+//
+// Get the name used in the Op to refer to an enum value of the given
+// `EnumClass`.
+// template <typename EnumClass> StringRef attributeName();
+//
+// Get the function that can be used to symbolize an enum value.
+// template <typename EnumClass>
+// Optional<EnumClass> (*)(StringRef) symbolizeEnum();
 #include "mlir/Dialect/SPIRV/SPIRVOpUtils.inc"
 
 } // end namespace spirv