[mlir][spirv] Add get() method to TargetEnvAttr taking raw values

Also make the getResourceLimits() method more explicit about its return type.

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

diff --git a/mlir/include/mlir/Dialect/SPIRV/TargetAndABI.h b/mlir/include/mlir/Dialect/SPIRV/TargetAndABI.h
index a20fe6e..01b7758 100644 (file)
--- a/mlir/include/mlir/Dialect/SPIRV/TargetAndABI.h
+++ b/mlir/include/mlir/Dialect/SPIRV/TargetAndABI.h
@@ -52,6 +52,9 @@ public:
   using Base::Base;
 
   /// Gets a TargetEnvAttr instance.
+  static TargetEnvAttr get(Version version, ArrayRef<Extension> extensions,
+                           ArrayRef<Capability> capabilities,
+                           DictionaryAttr limits);
   static TargetEnvAttr get(IntegerAttr version, ArrayAttr extensions,
                            ArrayAttr capabilities, DictionaryAttr limits);
 
@@ -86,7 +89,7 @@ public:
   ArrayAttr getCapabilitiesAttr();
 
   /// Returns the target resource limits.
-  DictionaryAttr getResourceLimits();
+  ResourceLimitsAttr getResourceLimits();
 
   static bool kindof(unsigned kind) { return kind == AttrKind::TargetEnv; }
 

diff --git a/mlir/lib/Dialect/SPIRV/TargetAndABI.cpp b/mlir/lib/Dialect/SPIRV/TargetAndABI.cpp
index ee5dcb1..f8c5900 100644 (file)
--- a/mlir/lib/Dialect/SPIRV/TargetAndABI.cpp
+++ b/mlir/lib/Dialect/SPIRV/TargetAndABI.cpp
@@ -48,6 +48,27 @@ struct TargetEnvAttributeStorage : public AttributeStorage {
 } // namespace spirv
 } // namespace mlir
 
+spirv::TargetEnvAttr spirv::TargetEnvAttr::get(
+    spirv::Version version, ArrayRef<spirv::Extension> extensions,
+    ArrayRef<spirv::Capability> capabilities, DictionaryAttr limits) {
+  Builder b(limits.getContext());
+
+  auto versionAttr = b.getI32IntegerAttr(static_cast<uint32_t>(version));
+
+  SmallVector<Attribute, 4> extAttrs;
+  extAttrs.reserve(extensions.size());
+  for (spirv::Extension ext : extensions)
+    extAttrs.push_back(b.getStringAttr(spirv::stringifyExtension(ext)));
+
+  SmallVector<Attribute, 4> capAttrs;
+  capAttrs.reserve(capabilities.size());
+  for (spirv::Capability cap : capabilities)
+    capAttrs.push_back(b.getI32IntegerAttr(static_cast<uint32_t>(cap)));
+
+  return get(versionAttr, b.getArrayAttr(extAttrs), b.getArrayAttr(capAttrs),
+             limits);
+}
+
 spirv::TargetEnvAttr spirv::TargetEnvAttr::get(IntegerAttr version,
                                                ArrayAttr extensions,
                                                ArrayAttr capabilities,
@@ -98,8 +119,8 @@ ArrayAttr spirv::TargetEnvAttr::getCapabilitiesAttr() {
   return getImpl()->capabilities.cast<ArrayAttr>();
 }
 
-DictionaryAttr spirv::TargetEnvAttr::getResourceLimits() {
-  return getImpl()->limits.cast<DictionaryAttr>();
+spirv::ResourceLimitsAttr spirv::TargetEnvAttr::getResourceLimits() {
+  return getImpl()->limits.cast<spirv::ResourceLimitsAttr>();
 }
 
 LogicalResult spirv::TargetEnvAttr::verifyConstructionInvariants(