[mlir][spirv] Plumbing target environment into type converter

[antiagainst]

Summary

This commit unifies target environment queries into a new wrapper
class spirv::TargetEnv and shares across various places needing
the functionality. We still create multiple instances of TargetEnv
though given the parent components (type converters, passes,
conversion targets) have different lifetimes.

In the meantime, LowerABIAttributesPass is updated to take into
consideration the target environment, which requires udpates to
tests to provide that.

Depends On D75876

[mravishankar]

mlir/include/mlir/Dialect/SPIRV/SPIRVLowering.h

77

Added a comment on the previous CL in this stack, but just mirroring here. I think its would be good to have a get method which will set the default target env. Its seems more ergonomic to me.

mlir/include/mlir/Dialect/SPIRV/TargetAndABI.h

54

These might cause subtle issues later on. Can we make this a SetVector?

[antiagainst]

mlir/include/mlir/Dialect/SPIRV/SPIRVLowering.h

77

Replied in the previous CL.

mlir/include/mlir/Dialect/SPIRV/TargetAndABI.h

54

But what we want here is just a set. We don't care about the order. What problems do you think will have with SmallSet?

[mravishankar]

mlir/include/mlir/Dialect/SPIRV/TargetAndABI.h

54

For now this is not iterated over. But there is no gaurantee later on and is easy to miss. Dont see too much advantage of using SmallSet vs SetVector, and just prefer to use the latter.

[antiagainst]

mlir/include/mlir/Dialect/SPIRV/TargetAndABI.h

54

Yeah I think we can re-evaluate this later when we want to loop over the set. Likely we will not have that need; I'd assume querying whether some extensions/capabilities are allowed is all what we need for using the target environment. But you never know. :)

mlir/include/mlir/Dialect/SPIRV/SPIRVLowering.h

 //===- SPIRVLowering.h - SPIR-V lowering utilities  -------------*- C++ -*-===//
 //
 // Part of the LLVM 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Defines utilities to use while targeting SPIR-V dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_SPIRV_SPIRVLOWERING_H
 #define MLIR_DIALECT_SPIRV_SPIRVLOWERING_H
 
+#include "mlir/Dialect/SPIRV/SPIRVAttributes.h"
 #include "mlir/Dialect/SPIRV/SPIRVTypes.h"
 #include "mlir/Dialect/SPIRV/TargetAndABI.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/ADT/SmallSet.h"
 
 namespace mlir {
 
 /// Type conversion from standard types to SPIR-V types for shader interface.
 ///
 /// For composite types, this converter additionally performs type wrapping to
 /// satisfy shader interface requirements: shader interface types must be
 /// pointers to structs.
 class SPIRVTypeConverter : public TypeConverter {
 public:
-  SPIRVTypeConverter();
+  explicit SPIRVTypeConverter(spirv::TargetEnvAttr targetAttr);
 
   /// Gets the SPIR-V correspondence for the standard index type.
   static Type getIndexType(MLIRContext *context);
 
   /// Returns the corresponding memory space for memref given a SPIR-V storage
   /// class.
   static unsigned getMemorySpaceForStorageClass(spirv::StorageClass);
 
   /// Returns the SPIR-V storage class given a memory space for memref. Return
   /// llvm::None if the memory space does not map to any SPIR-V storage class.
   static Optional<spirv::StorageClass>
   getStorageClassForMemorySpace(unsigned space);
+
+private:
+  spirv::TargetEnv targetEnv;
 };
 
 /// Base class to define a conversion pattern to lower `SourceOp` into SPIR-V.
 template <typename SourceOp>
 class SPIRVOpLowering : public OpConversionPattern<SourceOp> {
 public:
   SPIRVOpLowering(MLIRContext *context, SPIRVTypeConverter &typeConverter,
                   PatternBenefit benefit = 1)

 
 namespace spirv {
 class AccessChainOp;
 class FuncOp;
 
 class SPIRVConversionTarget : public ConversionTarget {
 public:
   /// Creates a SPIR-V conversion target for the given target environment.
-  static std::unique_ptr<SPIRVConversionTarget> get(TargetEnvAttr targetEnv,
-                                                    MLIRContext *context);
+  static std::unique_ptr<SPIRVConversionTarget> get(TargetEnvAttr targetAttr);
 
 private:
-  SPIRVConversionTarget(TargetEnvAttr targetEnv, MLIRContext *context);
+  explicit SPIRVConversionTarget(TargetEnvAttr targetAttr);
 
   // Be explicit that instance of this class cannot be copied or moved: there
   // are lambdas capturing fields of the instance.
   SPIRVConversionTarget(const SPIRVConversionTarget &) = delete;
   SPIRVConversionTarget(SPIRVConversionTarget &&) = delete;
   SPIRVConversionTarget &operator=(const SPIRVConversionTarget &) = delete;
   SPIRVConversionTarget &operator=(SPIRVConversionTarget &&) = delete;
 
   /// Returns true if the given `op` is legal to use under the current target
   /// environment.
   bool isLegalOp(Operation *op);
 
-  Version givenVersion;                            /// SPIR-V version to target
-  llvm::SmallSet<Extension, 4> givenExtensions;    /// Allowed extensions
-  llvm::SmallSet<Capability, 8> givenCapabilities; /// Allowed capabilities
+  TargetEnv targetEnv;
 };
 
 /// Returns the value for the given `builtin` variable. This function gets or
 /// inserts the global variable associated for the builtin within the nearest
 /// enclosing op that has a symbol table. Returns null Value if such an
 /// enclosing op cannot be found.
 Value getBuiltinVariableValue(Operation *op, BuiltIn builtin,
                               OpBuilder &builder);

mlir/include/mlir/Dialect/SPIRV/TargetAndABI.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_SPIRV_TARGETANDABI_H
 #define MLIR_DIALECT_SPIRV_TARGETANDABI_H
 
 #include "mlir/Dialect/SPIRV/SPIRVAttributes.h"
 #include "mlir/Support/LLVM.h"
+#include "llvm/ADT/SmallSet.h"
 
 namespace mlir {
 class Operation;
 
 namespace spirv {
 enum class StorageClass : uint32_t;
 
+/// A wrapper class around a spirv::TargetEnvAttr to provide query methods for
+/// allowed version/capabilities/extensions.
+class TargetEnv {
+public:
+  explicit TargetEnv(TargetEnvAttr targetAttr);
+
+  Version getVersion();
+
+  /// Returns true if the given capability is allowed.
+  bool allows(Capability) const;
+  /// Returns the first allowed one if any of the given capabilities is allowed.
+  /// Returns llvm::None otherwise.
+  Optional<Capability> allows(ArrayRef<Capability>) const;
+
+  /// Returns true if the given extension is allowed.
+  bool allows(Extension) const;
+  /// Returns the first allowed one if any of the given extensions is allowed.
+  /// Returns llvm::None otherwise.
+  Optional<Extension> allows(ArrayRef<Extension>) const;
+
+  /// Returns the MLIRContext.
+  MLIRContext *getContext();
+
+  /// Allows implicity converting to the underlying spirv::TargetEnvAttr.
+  operator TargetEnvAttr() const { return targetAttr; }
+
+private:
+  TargetEnvAttr targetAttr;
+  llvm::SmallSet<Extension, 4> givenExtensions;    /// Allowed extensions
+  llvm::SmallSet<Capability, 8> givenCapabilities; /// Allowed capabilities
+};
+
 /// Returns the attribute name for specifying argument ABI information.
 StringRef getInterfaceVarABIAttrName();
 
 /// Gets the InterfaceVarABIAttr given its fields.
 InterfaceVarABIAttr getInterfaceVarABIAttr(unsigned descriptorSet,
                                            unsigned binding,
                                            Optional<StorageClass> storageClass,
                                            MLIRContext *context);

mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRVPass.cpp

void GPUToSPIRVPass::runOnModule() {
   module.walk([&builder, &kernelModules](gpu::GPUModuleOp moduleOp) {
     // For each kernel module (should be only 1 for now, but that is not a
     // requirement here), clone the module for conversion because the
     // gpu.launch function still needs the kernel module.
     builder.setInsertionPoint(moduleOp.getOperation());
     kernelModules.push_back(builder.clone(*moduleOp.getOperation()));
   });
 
-  SPIRVTypeConverter typeConverter;
+  auto targetAttr = spirv::lookupTargetEnvOrDefault(module);
+  std::unique_ptr<ConversionTarget> target =
+      spirv::SPIRVConversionTarget::get(targetAttr);
+
+  SPIRVTypeConverter typeConverter(targetAttr);
   OwningRewritePatternList patterns;
   populateGPUToSPIRVPatterns(context, typeConverter, patterns);
   populateStandardToSPIRVPatterns(context, typeConverter, patterns);
 
-  std::unique_ptr<ConversionTarget> target = spirv::SPIRVConversionTarget::get(
-      spirv::lookupTargetEnvOrDefault(module), context);
-
   if (failed(applyFullConversion(kernelModules, *target, patterns,
                                  &typeConverter))) {
     return signalPassFailure();
   }
 }
 
 std::unique_ptr<OpPassBase<ModuleOp>> mlir::createConvertGPUToSPIRVPass() {
   return std::make_unique<GPUToSPIRVPass>();
 }
 
 static PassRegistration<GPUToSPIRVPass>
     pass("convert-gpu-to-spirv", "Convert GPU dialect to SPIR-V dialect");

mlir/lib/Conversion/LinalgToSPIRV/LinalgToSPIRVPass.cpp

class LinalgToSPIRVPass : public ModulePass<LinalgToSPIRVPass> {
   void runOnModule() override;
 };
 } // namespace
 
 void LinalgToSPIRVPass::runOnModule() {
   MLIRContext *context = &getContext();
   ModuleOp module = getModule();
 
-  SPIRVTypeConverter typeConverter;
+  auto targetAttr = spirv::lookupTargetEnvOrDefault(module);
+  std::unique_ptr<ConversionTarget> target =
+      spirv::SPIRVConversionTarget::get(targetAttr);
+
+  SPIRVTypeConverter typeConverter(targetAttr);
   OwningRewritePatternList patterns;
   populateLinalgToSPIRVPatterns(context, typeConverter, patterns);
   populateBuiltinFuncToSPIRVPatterns(context, typeConverter, patterns);
 
-  auto targetEnv = spirv::lookupTargetEnvOrDefault(module);
-  std::unique_ptr<ConversionTarget> target =
-      spirv::SPIRVConversionTarget::get(targetEnv, context);
-
   // Allow builtin ops.
   target->addLegalOp<ModuleOp, ModuleTerminatorOp>();
   target->addDynamicallyLegalOp<FuncOp>(
       [&](FuncOp op) { return typeConverter.isSignatureLegal(op.getType()); });
 
   if (failed(applyFullConversion(module, *target, patterns)))
     return signalPassFailure();
 }
 
 std::unique_ptr<OpPassBase<ModuleOp>> mlir::createLinalgToSPIRVPass() {
   return std::make_unique<LinalgToSPIRVPass>();
 }
 
 static PassRegistration<LinalgToSPIRVPass>
     pass("convert-linalg-to-spirv", "Convert Linalg ops to SPIR-V ops");

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRVPass.cpp

class ConvertStandardToSPIRVPass
   void runOnModule() override;
 };
 } // namespace
 
 void ConvertStandardToSPIRVPass::runOnModule() {
   MLIRContext *context = &getContext();
   ModuleOp module = getModule();
 
-  SPIRVTypeConverter typeConverter;
+  auto targetAttr = spirv::lookupTargetEnvOrDefault(module);
+  std::unique_ptr<ConversionTarget> target =
+      spirv::SPIRVConversionTarget::get(targetAttr);
+
+  SPIRVTypeConverter typeConverter(targetAttr);
   OwningRewritePatternList patterns;
   populateStandardToSPIRVPatterns(context, typeConverter, patterns);
   populateBuiltinFuncToSPIRVPatterns(context, typeConverter, patterns);
 
-  std::unique_ptr<ConversionTarget> target = spirv::SPIRVConversionTarget::get(
-      spirv::lookupTargetEnvOrDefault(module), context);
-
   if (failed(applyPartialConversion(module, *target, patterns))) {
     return signalPassFailure();
   }
 }
 
 std::unique_ptr<OpPassBase<ModuleOp>> mlir::createConvertStandardToSPIRVPass() {
   return std::make_unique<ConvertStandardToSPIRVPass>();
 }
 
 static PassRegistration<ConvertStandardToSPIRVPass>
     pass("convert-std-to-spirv", "Convert Standard Ops to SPIR-V dialect");

mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp

for (auto shape : tensorType.getShape()) {
       size *= shape;
     }
     return size;
   }
   // TODO: Add size computation for other types.
   return llvm::None;
 }
 
-SPIRVTypeConverter::SPIRVTypeConverter() {
+SPIRVTypeConverter::SPIRVTypeConverter(spirv::TargetEnvAttr targetAttr)
+    : targetEnv(targetAttr) {
   addConversion([](Type type) -> Optional<Type> {
     // If the type is already valid in SPIR-V, directly return.
     return spirv::SPIRVDialect::isValidType(type) ? type : Optional<Type>();
   });
   addConversion([](IndexType indexType) {
     return SPIRVTypeConverter::getIndexType(indexType.getContext());
   });
   addConversion([this](MemRefType memRefType) -> Type {
mlir::spirv::setABIAttrs(spirv::FuncOp funcOp,
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // SPIR-V ConversionTarget
 //===----------------------------------------------------------------------===//
 
 std::unique_ptr<spirv::SPIRVConversionTarget>
-spirv::SPIRVConversionTarget::get(spirv::TargetEnvAttr targetEnv,
-                                  MLIRContext *context) {
+spirv::SPIRVConversionTarget::get(spirv::TargetEnvAttr targetAttr) {
   std::unique_ptr<SPIRVConversionTarget> target(
       // std::make_unique does not work here because the constructor is private.
-      new SPIRVConversionTarget(targetEnv, context));
+      new SPIRVConversionTarget(targetAttr));
   SPIRVConversionTarget *targetPtr = target.get();
   target->addDynamicallyLegalDialect<SPIRVDialect>(
       Optional<ConversionTarget::DynamicLegalityCallbackFn>(
           // We need to capture the raw pointer here because it is stable:
           // target will be destroyed once this function is returned.
           [targetPtr](Operation *op) { return targetPtr->isLegalOp(op); }));
   return target;
 }
 
 spirv::SPIRVConversionTarget::SPIRVConversionTarget(
-    spirv::TargetEnvAttr targetEnv, MLIRContext *context)
-    : ConversionTarget(*context), givenVersion(targetEnv.getVersion()) {
-  for (spirv::Extension ext : targetEnv.getExtensions())
-    givenExtensions.insert(ext);
-
-  // Add extensions implied by the current version.
-  for (spirv::Extension ext : spirv::getImpliedExtensions(givenVersion))
-    givenExtensions.insert(ext);
-
-  for (spirv::Capability cap : targetEnv.getCapabilities()) {
-    givenCapabilities.insert(cap);
-
-    // Add capabilities implied by the current capability.
-    for (spirv::Capability c : spirv::getRecursiveImpliedCapabilities(cap))
-      givenCapabilities.insert(c);
-  }
-}
+    spirv::TargetEnvAttr targetAttr)
+    : ConversionTarget(*targetAttr.getContext()), targetEnv(targetAttr) {}
 
 /// Checks that `candidates` extension requirements are possible to be satisfied
-/// with the given `allowedExtensions`.
+/// with the given `targetEnv`.
 ///
 ///  `candidates` is a vector of vector for extension requirements following
 /// ((Extension::A OR Extension::B) AND (Extension::C OR Extension::D))
 /// convention.
 static LogicalResult checkExtensionRequirements(
-    Operation *op, const llvm::SmallSet<spirv::Extension, 4> &allowedExtensions,
+    Operation *op, const spirv::TargetEnv &targetEnv,
     const spirv::SPIRVType::ExtensionArrayRefVector &candidates) {
   for (const auto &ors : candidates) {
-    auto chosen = llvm::find_if(ors, [&](spirv::Extension ext) {
-      return allowedExtensions.count(ext);
-    });
+    if (targetEnv.allows(ors))
+      continue;
 
-    if (chosen == ors.end()) {
-      SmallVector<StringRef, 4> extStrings;
-      for (spirv::Extension ext : ors)
-        extStrings.push_back(spirv::stringifyExtension(ext));
 
-      LLVM_DEBUG(llvm::dbgs() << op->getName()
-                              << "illegal: requires at least one extension in ["
-                              << llvm::join(extStrings, ", ")
-                              << "] but none allowed in target environment\n");
-      return failure();
-    }
-  }
+                            << " illegal: requires at least one extension in ["
   return success();
 }
 
 /// Checks that `candidates`capability requirements are possible to be satisfied
-/// with the given `allowedCapabilities`.
+/// with the given `isAllowedFn`.
 ///
 ///  `candidates` is a vector of vector for capability requirements following
 /// ((Capability::A OR Capability::B) AND (Capability::C OR Capability::D))
 /// convention.
 static LogicalResult checkCapabilityRequirements(
-    Operation *op,
-    const llvm::SmallSet<spirv::Capability, 8> &allowedCapabilities,
+    Operation *op, const spirv::TargetEnv &targetEnv,
     const spirv::SPIRVType::CapabilityArrayRefVector &candidates) {
   for (const auto &ors : candidates) {
-    auto chosen = llvm::find_if(ors, [&](spirv::Capability cap) {
-      return allowedCapabilities.count(cap);
-    });
+    if (targetEnv.allows(ors))
+      continue;
 
-    if (chosen == ors.end()) {
-      SmallVector<StringRef, 4> capStrings;
-      for (spirv::Capability cap : ors)
-        capStrings.push_back(spirv::stringifyCapability(cap));
 
-      LLVM_DEBUG(llvm::dbgs()
-                 << op->getName()
-                 << "illegal: requires at least one capability in ["
+    LLVM_DEBUG(llvm::dbgs() << op->getName()
+                            << " illegal: requires at least one capability in ["
                             << llvm::join(capStrings, ", ")
                             << "] but none allowed in target environment\n");
-      return failure();
-    }
-  }
   return success();
 }
 
 bool spirv::SPIRVConversionTarget::isLegalOp(Operation *op) {
   // Make sure this op is available at the given version. Ops not implementing
   // QueryMinVersionInterface/QueryMaxVersionInterface are available to all
   // SPIR-V versions.
   if (auto minVersion = dyn_cast<spirv::QueryMinVersionInterface>(op))
-    if (minVersion.getMinVersion() > givenVersion) {
+    if (minVersion.getMinVersion() > this->targetEnv.getVersion()) {
       LLVM_DEBUG(llvm::dbgs()
                  << op->getName() << " illegal: requiring min version "
                  << spirv::stringifyVersion(minVersion.getMinVersion())
                  << "\n");
       return false;
     }
   if (auto maxVersion = dyn_cast<spirv::QueryMaxVersionInterface>(op))
-    if (maxVersion.getMaxVersion() < givenVersion) {
+    if (maxVersion.getMaxVersion() < this->targetEnv.getVersion()) {
       LLVM_DEBUG(llvm::dbgs()
                  << op->getName() << " illegal: requiring max version "
                  << spirv::stringifyVersion(maxVersion.getMaxVersion())
                  << "\n");
       return false;
     }
 
   // Make sure this op's required extensions are allowed to use. Ops not
   // implementing QueryExtensionInterface do not require extensions to be
   // available.
   if (auto extensions = dyn_cast<spirv::QueryExtensionInterface>(op))
-    if (failed(checkExtensionRequirements(op, this->givenExtensions,
+    if (failed(checkExtensionRequirements(op, this->targetEnv,
                                           extensions.getExtensions())))
       return false;
 
   // Make sure this op's required extensions are allowed to use. Ops not
   // implementing QueryCapabilityInterface do not require capabilities to be
   // available.
   if (auto capabilities = dyn_cast<spirv::QueryCapabilityInterface>(op))
-    if (failed(checkCapabilityRequirements(op, this->givenCapabilities,
+    if (failed(checkCapabilityRequirements(op, this->targetEnv,
                                            capabilities.getCapabilities())))
       return false;
 
   SmallVector<Type, 4> valueTypes;
   valueTypes.append(op->operand_type_begin(), op->operand_type_end());
   valueTypes.append(op->result_type_begin(), op->result_type_end());
 
   // Special treatment for global variables, whose type requirements are
   // conveyed by type attributes.
   if (auto globalVar = dyn_cast<spirv::GlobalVariableOp>(op))
     valueTypes.push_back(globalVar.type());
 
   // Make sure the op's operands/results use types that are allowed by the
   // target environment.
   SmallVector<ArrayRef<spirv::Extension>, 4> typeExtensions;
   SmallVector<ArrayRef<spirv::Capability>, 8> typeCapabilities;
   for (Type valueType : valueTypes) {
     typeExtensions.clear();
     valueType.cast<spirv::SPIRVType>().getExtensions(typeExtensions);
-    if (failed(checkExtensionRequirements(op, this->givenExtensions,
-                                          typeExtensions)))
+    if (failed(checkExtensionRequirements(op, this->targetEnv, typeExtensions)))
       return false;
 
     typeCapabilities.clear();
     valueType.cast<spirv::SPIRVType>().getCapabilities(typeCapabilities);
-    if (failed(checkCapabilityRequirements(op, this->givenCapabilities,
-                                           typeCapabilities)))
+    if (failed(
+            checkCapabilityRequirements(op, this->targetEnv, typeCapabilities)))
       return false;
   }
 
   return true;
 }

mlir/lib/Dialect/SPIRV/TargetAndABI.cpp

 #include "mlir/IR/Builders.h"
 #include "mlir/IR/FunctionSupport.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Support/LLVM.h"
 
 using namespace mlir;
 
+//===----------------------------------------------------------------------===//
+// TargetEnv
+//===----------------------------------------------------------------------===//
+
+spirv::TargetEnv::TargetEnv(spirv::TargetEnvAttr targetAttr)
+    : targetAttr(targetAttr) {
+  for (spirv::Extension ext : targetAttr.getExtensions())
+    givenExtensions.insert(ext);
+
+  // Add extensions implied by the current version.
+  for (spirv::Extension ext :
+       spirv::getImpliedExtensions(targetAttr.getVersion()))
+    givenExtensions.insert(ext);
+
+  for (spirv::Capability cap : targetAttr.getCapabilities()) {
+    givenCapabilities.insert(cap);
+
+    // Add capabilities implied by the current capability.
+    for (spirv::Capability c : spirv::getRecursiveImpliedCapabilities(cap))
+      givenCapabilities.insert(c);
+  }
+}
+
+spirv::Version spirv::TargetEnv::getVersion() {
+  return targetAttr.getVersion();
+}
+
+bool spirv::TargetEnv::allows(spirv::Capability capability) const {
+  return givenCapabilities.count(capability);
+}
+
+Optional<spirv::Capability>
+spirv::TargetEnv::allows(ArrayRef<spirv::Capability> caps) const {
+  auto chosen = llvm::find_if(caps, [this](spirv::Capability cap) {
+    return givenCapabilities.count(cap);
+  });
+  if (chosen != caps.end())
+    return *chosen;
+  return llvm::None;
+}
+
+bool spirv::TargetEnv::allows(spirv::Extension extension) const {
+  return givenExtensions.count(extension);
+}
+
+Optional<spirv::Extension>
+spirv::TargetEnv::allows(ArrayRef<spirv::Extension> exts) const {
+  auto chosen = llvm::find_if(exts, [this](spirv::Extension ext) {
+    return givenExtensions.count(ext);
+  });
+  if (chosen != exts.end())
+    return *chosen;
+  return llvm::None;
+}
+
+MLIRContext *spirv::TargetEnv::getContext() { return targetAttr.getContext(); }
+
+//===----------------------------------------------------------------------===//
+// Utility functions
+//===----------------------------------------------------------------------===//
+
 StringRef spirv::getInterfaceVarABIAttrName() {
   return "spv.interface_var_abi";
 }
 
 spirv::InterfaceVarABIAttr
 spirv::getInterfaceVarABIAttr(unsigned descriptorSet, unsigned binding,
                               Optional<spirv::StorageClass> storageClass,
                               MLIRContext *context) {

mlir/lib/Dialect/SPIRV/Transforms/LowerABIAttributesPass.cpp

 }
 
 void LowerABIAttributesPass::runOnOperation() {
   // Uses the signature conversion methodology of the dialect conversion
   // framework to implement the conversion.
   spirv::ModuleOp module = getOperation();
   MLIRContext *context = &getContext();
 
-  SPIRVTypeConverter typeConverter;
+  spirv::TargetEnv targetEnv(spirv::lookupTargetEnv(module));
+
+  SPIRVTypeConverter typeConverter(targetEnv);
   OwningRewritePatternList patterns;
   patterns.insert<ProcessInterfaceVarABI>(context, typeConverter);
 
   ConversionTarget target(*context);
   // "Legal" function ops should have no interface variable ABI attributes.
   target.addDynamicallyLegalOp<spirv::FuncOp>([&](spirv::FuncOp op) {
     StringRef attrName = spirv::getInterfaceVarABIAttrName();
     for (unsigned i = 0, e = op.getNumArguments(); i < e; ++i)

mlir/lib/Dialect/SPIRV/Transforms/UpdateVCEPass.cpp

 class UpdateVCEPass final
     : public OperationPass<UpdateVCEPass, spirv::ModuleOp> {
 private:
   void runOnOperation() override;
 };
 } // namespace
 
 /// Checks that `candidates` extension requirements are possible to be satisfied
-/// with the given `allowedExtensions` and updates `deducedExtensions` if so.
-/// Emits errors attaching to the given `op` on failures.
+/// with the given `targetEnv` and updates `deducedExtensions` if so. Emits
+/// errors attaching to the given `op` on failures.
 ///
 ///  `candidates` is a vector of vector for extension requirements following
 /// ((Extension::A OR Extension::B) AND (Extension::C OR Extension::D))
 /// convention.
 static LogicalResult checkAndUpdateExtensionRequirements(
-    Operation *op, const llvm::SmallSet<spirv::Extension, 4> &allowedExtensions,
+    Operation *op, const spirv::TargetEnv &targetEnv,
     const spirv::SPIRVType::ExtensionArrayRefVector &candidates,
     llvm::SetVector<spirv::Extension> &deducedExtensions) {
   for (const auto &ors : candidates) {
-    auto chosen = llvm::find_if(ors, [&](spirv::Extension ext) {
-      return allowedExtensions.count(ext);
-    });
-
-    if (chosen != ors.end()) {
+    if (Optional<spirv::Extension> chosen = targetEnv.allows(ors)) {
       deducedExtensions.insert(*chosen);
     } else {
       SmallVector<StringRef, 4> extStrings;
       for (spirv::Extension ext : ors)
         extStrings.push_back(spirv::stringifyExtension(ext));
 
       return op->emitError("'")
              << op->getName() << "' requires at least one extension in ["
              << llvm::join(extStrings, ", ")
              << "] but none allowed in target environment";
     }
   }
   return success();
 }
 
 /// Checks that `candidates`capability requirements are possible to be satisfied
-/// with the given `allowedCapabilities` and updates `deducedCapabilities` if
-/// so. Emits errors attaching to the given `op` on failures.
+/// with the given `targetEnv` and updates `deducedCapabilities` if so. Emits
+/// errors attaching to the given `op` on failures.
 ///
 ///  `candidates` is a vector of vector for capability requirements following
 /// ((Capability::A OR Capability::B) AND (Capability::C OR Capability::D))
 /// convention.
 static LogicalResult checkAndUpdateCapabilityRequirements(
-    Operation *op,
-    const llvm::SmallSet<spirv::Capability, 8> &allowedCapabilities,
+    Operation *op, const spirv::TargetEnv &targetEnv,
     const spirv::SPIRVType::CapabilityArrayRefVector &candidates,
     llvm::SetVector<spirv::Capability> &deducedCapabilities) {
   for (const auto &ors : candidates) {
-    auto chosen = llvm::find_if(ors, [&](spirv::Capability cap) {
-      return allowedCapabilities.count(cap);
-    });
-
-    if (chosen != ors.end()) {
+    if (Optional<spirv::Capability> chosen = targetEnv.allows(ors)) {
       deducedCapabilities.insert(*chosen);
     } else {
       SmallVector<StringRef, 4> capStrings;
       for (spirv::Capability cap : ors)
         capStrings.push_back(spirv::stringifyCapability(cap));
 
       return op->emitError("'")
              << op->getName() << "' requires at least one capability in ["
              << llvm::join(capStrings, ", ")
              << "] but none allowed in target environment";
     }
   }
   return success();
 }
 
 void UpdateVCEPass::runOnOperation() {
   spirv::ModuleOp module = getOperation();
 
-  spirv::TargetEnvAttr targetEnv = spirv::lookupTargetEnv(module);
-  if (!targetEnv) {
+  spirv::TargetEnvAttr targetAttr = spirv::lookupTargetEnv(module);
+  if (!targetAttr) {
     module.emitError("missing 'spv.target_env' attribute");
     return signalPassFailure();
   }
 
-  spirv::Version allowedVersion = targetEnv.getVersion();
-
-  // Build a set for available extensions in the target environment.
-  llvm::SmallSet<spirv::Extension, 4> allowedExtensions;
-  for (spirv::Extension ext : targetEnv.getExtensions())
-    allowedExtensions.insert(ext);
-
-  // Add extensions implied by the current version.
-  for (spirv::Extension ext : spirv::getImpliedExtensions(allowedVersion))
-    allowedExtensions.insert(ext);
-
-  // Build a set for available capabilities in the target environment.
-  llvm::SmallSet<spirv::Capability, 8> allowedCapabilities;
-  for (spirv::Capability cap : targetEnv.getCapabilities()) {
-    allowedCapabilities.insert(cap);
-
-    // Add capabilities implied by the current capability.
-    for (spirv::Capability c : spirv::getRecursiveImpliedCapabilities(cap))
-      allowedCapabilities.insert(c);
-  }
+  spirv::TargetEnv targetEnv(targetAttr);
+  spirv::Version allowedVersion = targetAttr.getVersion();
 
   spirv::Version deducedVersion = spirv::Version::V_1_0;
   llvm::SetVector<spirv::Extension> deducedExtensions;
   llvm::SetVector<spirv::Capability> deducedCapabilities;
 
   // Walk each SPIR-V op to deduce the minimal version/extension/capability
   // requirements.
   WalkResult walkResult = module.walk([&](Operation *op) -> WalkResult {
     // Op min version requirements
     if (auto minVersion = dyn_cast<spirv::QueryMinVersionInterface>(op)) {
       deducedVersion = std::max(deducedVersion, minVersion.getMinVersion());
       if (deducedVersion > allowedVersion) {
         return op->emitError("'") << op->getName() << "' requires min version "
                                   << spirv::stringifyVersion(deducedVersion)
                                   << " but target environment allows up to "
                                   << spirv::stringifyVersion(allowedVersion);
       }
     }
 
     // Op extension requirements
     if (auto extensions = dyn_cast<spirv::QueryExtensionInterface>(op))
-      if (failed(checkAndUpdateExtensionRequirements(op, allowedExtensions,
-                                                     extensions.getExtensions(),
-                                                     deducedExtensions)))
+      if (failed(checkAndUpdateExtensionRequirements(
+              op, targetEnv, extensions.getExtensions(), deducedExtensions)))
         return WalkResult::interrupt();
 
     // Op capability requirements
     if (auto capabilities = dyn_cast<spirv::QueryCapabilityInterface>(op))
       if (failed(checkAndUpdateCapabilityRequirements(
-              op, allowedCapabilities, capabilities.getCapabilities(),
+              op, targetEnv, capabilities.getCapabilities(),
               deducedCapabilities)))
         return WalkResult::interrupt();
 
     SmallVector<Type, 4> valueTypes;
     valueTypes.append(op->operand_type_begin(), op->operand_type_end());
     valueTypes.append(op->result_type_begin(), op->result_type_end());
 
     // Special treatment for global variables, whose type requirements are
     // conveyed by type attributes.
     if (auto globalVar = dyn_cast<spirv::GlobalVariableOp>(op))
       valueTypes.push_back(globalVar.type());
 
     // Requirements from values' types
     SmallVector<ArrayRef<spirv::Extension>, 4> typeExtensions;
     SmallVector<ArrayRef<spirv::Capability>, 8> typeCapabilities;
     for (Type valueType : valueTypes) {
       typeExtensions.clear();
       valueType.cast<spirv::SPIRVType>().getExtensions(typeExtensions);
       if (failed(checkAndUpdateExtensionRequirements(
-              op, allowedExtensions, typeExtensions, deducedExtensions)))
+              op, targetEnv, typeExtensions, deducedExtensions)))
         return WalkResult::interrupt();
 
       typeCapabilities.clear();
       valueType.cast<spirv::SPIRVType>().getCapabilities(typeCapabilities);
       if (failed(checkAndUpdateCapabilityRequirements(
-              op, allowedCapabilities, typeCapabilities, deducedCapabilities)))
+              op, targetEnv, typeCapabilities, deducedCapabilities)))
         return WalkResult::interrupt();
     }
 
     return WalkResult::advance();
   });
 
   if (walkResult.wasInterrupted())
     return signalPassFailure();

mlir/test/Dialect/SPIRV/Transforms/abi-interface.mlir

 // RUN: mlir-opt -spirv-lower-abi-attrs -verify-diagnostics %s -o - | FileCheck %s
 
+module attributes {
+  spv.target_env = #spv.target_env<
+    #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>,
+    {max_compute_workgroup_invocations = 128 : i32,
+     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>
+} {
+
 // CHECK-LABEL: spv.module
 spv.module Logical GLSL450 {
   // CHECK-DAG:    spv.globalVariable [[VAR0:@.*]] bind(0, 0) : !spv.ptr<!spv.struct<f32 [0]>, StorageBuffer>
   // CHECK-DAG:    spv.globalVariable [[VAR1:@.*]] bind(0, 1) : !spv.ptr<!spv.struct<!spv.array<12 x f32 [4]> [0]>, StorageBuffer>
   // CHECK:    spv.func [[FN:@.*]]()
   spv.func @kernel(
     %arg0: f32
            {spv.interface_var_abi = {binding = 0 : i32,
                                      descriptor_set = 0 : i32,
                                      storage_class = 12 : i32}},
     %arg1: !spv.ptr<!spv.struct<!spv.array<12 x f32>>, StorageBuffer>
            {spv.interface_var_abi = {binding = 1 : i32,
                                      descriptor_set = 0 : i32}}) "None"
   attributes {spv.entry_point_abi = {local_size = dense<[32, 1, 1]> : vector<3xi32>}} {
     // CHECK: [[ARG1:%.*]] = spv._address_of [[VAR1]]
     // CHECK: [[ADDRESSARG0:%.*]] = spv._address_of [[VAR0]]
     // CHECK: [[CONST0:%.*]] = spv.constant 0 : i32
     // CHECK: [[ARG0PTR:%.*]] = spv.AccessChain [[ADDRESSARG0]]{{\[}}[[CONST0]]
     // CHECK: [[ARG0:%.*]] = spv.Load "StorageBuffer" [[ARG0PTR]]
     // CHECK: spv.Return
     spv.Return
   }
   // CHECK: spv.EntryPoint "GLCompute" [[FN]]
   // CHECK: spv.ExecutionMode [[FN]] "LocalSize", 32, 1, 1
-}
+} // end spv.module
+
+} // end module

mlir/test/Dialect/SPIRV/Transforms/abi-load-store.mlir

 // RUN: mlir-opt -spirv-lower-abi-attrs -verify-diagnostics %s -o - | FileCheck %s
 
+module attributes {
+  spv.target_env = #spv.target_env<
+    #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>,
+    {max_compute_workgroup_invocations = 128 : i32,
+     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>
+} {
+
 // CHECK-LABEL: spv.module
 spv.module Logical GLSL450 {
   // CHECK-DAG: spv.globalVariable [[WORKGROUPSIZE:@.*]] built_in("WorkgroupSize")
   spv.globalVariable @__builtin_var_WorkgroupSize__ built_in("WorkgroupSize") : !spv.ptr<vector<3xi32>, Input>
   // CHECK-DAG: spv.globalVariable [[NUMWORKGROUPS:@.*]] built_in("NumWorkgroups")
   spv.globalVariable @__builtin_var_NumWorkgroups__ built_in("NumWorkgroups") : !spv.ptr<vector<3xi32>, Input>
   // CHECK-DAG: spv.globalVariable [[LOCALINVOCATIONID:@.*]] built_in("LocalInvocationId")
   spv.globalVariable @__builtin_var_LocalInvocationId__ built_in("LocalInvocationId") : !spv.ptr<vector<3xi32>, Input>
attributes  {spv.entry_point_abi = {local_size = dense<[32, 1, 1]> : vector<3xi32>}} {
     %42 = spv.FAdd %39, %41 : f32
     // CHECK: spv.AccessChain [[ARG2]]
     %43 = spv.AccessChain %arg2[%c0, %36, %37] : !spv.ptr<!spv.struct<!spv.array<12 x !spv.array<4 x f32>>>, StorageBuffer>
     spv.Store "StorageBuffer" %43, %42 : f32
     spv.Return
   }
   // CHECK: spv.EntryPoint "GLCompute" [[FN]], [[WORKGROUPID]], [[LOCALINVOCATIONID]], [[NUMWORKGROUPS]], [[WORKGROUPSIZE]]
   // CHECK-NEXT: spv.ExecutionMode [[FN]] "LocalSize", 32, 1, 1
-}
+} // end spv.module
+
+} // end module

mlir/test/lib/Dialect/SPIRV/TestAvailability.cpp

 
 void ConvertToTargetEnv::runOnFunction() {
   MLIRContext *context = &getContext();
   FuncOp fn = getFunction();
 
   auto targetEnv = fn.getOperation()
                        ->getAttr(spirv::getTargetEnvAttrName())
                        .cast<spirv::TargetEnvAttr>();
-  auto target = spirv::SPIRVConversionTarget::get(targetEnv, context);
+  if (!targetEnv) {
+    fn.emitError("missing 'spv.target_env' attribute");
+    return signalPassFailure();
+  }
+
+  auto target = spirv::SPIRVConversionTarget::get(targetEnv);
 
   OwningRewritePatternList patterns;
   patterns.insert<ConvertToAtomCmpExchangeWeak, ConvertToBitReverse,
                   ConvertToGroupNonUniformBallot, ConvertToModule,
                   ConvertToSubgroupBallot>(context);
 
   if (failed(applyPartialConversion(fn, *target, patterns)))
     return signalPassFailure();