[spirv] Add more target and resource limit fields

[antiagainst]

Summary

These fields will be used to choose/influence patterns for
SPIR-V code generation.

[mravishankar]

Nice! THanks!
Is there a way to add tests for these (do we need to)? Some roundtrip tests will serve as good reference.

[antiagainst]

In D87106#2255182, @mravishankar wrote:

Nice! THanks!
Is there a way to add tests for these (do we need to)? Some roundtrip tests will serve as good reference.

Not really. These should be used to select/influence patterns/passes. That's what we are gonna to develop and accumulate eventually. :)

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

 def QueryCapabilityInterface : OpInterface<"QueryCapabilityInterface"> {
   let methods = [InterfaceMethod<
     "",
     "::llvm::SmallVector<::llvm::ArrayRef<::mlir::spirv::Capability>, 1>",
     "getCapabilities">];
 }
 
 //===----------------------------------------------------------------------===//
+// SPIR-V target GPU vendor and device definitions
+//===----------------------------------------------------------------------===//
+
+// An accelerator other than GPU or CPU
+def SPV_DT_Other         : I32EnumAttrCase<"Other", 0>;
+def SPV_DT_IntegratedGPU : I32EnumAttrCase<"IntegratedGPU", 1>;
+def SPV_DT_DiscreteGPU   : I32EnumAttrCase<"DiscreteGPU", 2>;
+def SPV_DT_CPU           : I32EnumAttrCase<"CPU", 3>;
+// Information missing.
+def SPV_DT_Unknown       : I32EnumAttrCase<"Unknown", 0x7FFFFFFF>;
+
+def SPV_DeviceTypeAttr : SPV_I32EnumAttr<
+  "DeviceType", "valid SPIR-V device types", [
+    SPV_DT_Other, SPV_DT_IntegratedGPU, SPV_DT_DiscreteGPU,
+    SPV_DT_CPU, SPV_DT_Unknown
+  ]>;
+
+//===----------------------------------------------------------------------===//
 // SPIR-V extension definitions
 //===----------------------------------------------------------------------===//
 
 // Extensions known to the SPIR-V dialect.
 // https://github.com/KhronosGroup/SPIRV-Registry has the full list.
 def SPV_KHR_16bit_storage                : StrEnumAttrCase<"SPV_KHR_16bit_storage">;
 def SPV_KHR_8bit_storage                 : StrEnumAttrCase<"SPV_KHR_8bit_storage">;
 def SPV_KHR_device_group                 : StrEnumAttrCase<"SPV_KHR_device_group">;

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

 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);
 
+  DeviceType getDeviceType();
+
   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;
 

mlir/include/mlir/Dialect/SPIRV/TargetAndABI.td

 // architecture.
 //
 // See https://www.khronos.org/registry/vulkan/specs/1.2-extensions/html/vkspec.html#limits
 // for the complete list of limits and their explanation for the Vulkan API.
 // The following ones are those affecting SPIR-V CodeGen. Their default value
 // are the from Vulkan limit requirements:
 // https://www.khronos.org/registry/vulkan/specs/1.2-extensions/html/vkspec.html#limits-minmax
 def SPV_ResourceLimitsAttr : StructAttr<"ResourceLimitsAttr", SPIRV_Dialect, [
+    // Unique identifier for the vendor and target GPU.
+    // 0x7FFFFFFF means unknown.
+    StructFieldAttr<"vendor_id", DefaultValuedAttr<I32Attr, "0x7FFFFFFF">>,
+    StructFieldAttr<"device_id", DefaultValuedAttr<I32Attr, "0x7FFFFFFF">>,
+    // Target device type.
+    StructFieldAttr<"device_type",
+                    DefaultValuedAttr<SPV_DeviceTypeAttr,
+                                      "::mlir::spirv::DeviceType::Unknown">>,
+
+    // The maximum total storage size, in bytes, available for variables
+    // declared with the Workgroup storage class.
+    StructFieldAttr<"max_compute_shared_memory_size",
+                    DefaultValuedAttr<I32Attr, "16384">>,
+
+    // The maximum total number of compute shader invocations in a single local
+    // workgroup.
     StructFieldAttr<"max_compute_workgroup_invocations",
                     DefaultValuedAttr<I32Attr, "128">>,
+    // The maximum size of a local compute workgroup, per dimension.
     StructFieldAttr<"max_compute_workgroup_size",
-                    DefaultValuedAttr<I32ElementsAttr, "{128, 128, 64}">>
+                    DefaultValuedAttr<I32ElementsAttr, "{128, 128, 64}">>,
+
+    // The default number of invocations in each subgroup.
+    // 0x7FFFFFFF means unknown.
+    StructFieldAttr<"subgroup_size", DefaultValuedAttr<I32Attr, "0x7FFFFFFF">>
 ]>;
 
 #endif // SPIRV_TARGET_AND_ABI

mlir/lib/Dialect/SPIRV/TargetAndABI.cpp

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::DeviceType spirv::TargetEnv::getDeviceType() {
+  auto deviceType = spirv::symbolizeDeviceType(
+      targetAttr.getResourceLimits().device_type().getInt());
+  if (!deviceType)
+    return DeviceType::Unknown;
+  return *deviceType;
+}
+
 spirv::Version spirv::TargetEnv::getVersion() {
   return targetAttr.getVersion();
 }
 
 bool spirv::TargetEnv::allows(spirv::Capability capability) const {
   return givenCapabilities.count(capability);
 }
 
DenseIntElementsAttr spirv::lookupLocalWorkGroupSize(Operation *op) {
   if (auto entryPoint = spirv::lookupEntryPointABI(op))
     return entryPoint.local_size();
 
   return {};
 }
 
 spirv::ResourceLimitsAttr
 spirv::getDefaultResourceLimits(MLIRContext *context) {
-  auto i32Type = IntegerType::get(32, context);
-  auto v3i32Type = VectorType::get(3, i32Type);
-
-  // These numbers are from "Table 46. Required Limits" of the Vulkan spec.
+  // All the fields have default values. Here we just provide a nicer way to
+  // construct a default resource limit attribute.
   return spirv::ResourceLimitsAttr ::get(
-      IntegerAttr::get(i32Type, 128),
-      DenseIntElementsAttr::get<int32_t>(v3i32Type, {128, 128, 64}), context);
+      /*vendor_id=*/nullptr,
+      /*device_id*/ nullptr,
+      /*device_type=*/nullptr,
+      /*max_compute_shared_memory_size=*/nullptr,
+      /*max_compute_workgroup_invocations=*/nullptr,
+      /*max_compute_workgroup_size=*/nullptr,
+      /*subgroup_size=*/nullptr, context);
 }
 
 StringRef spirv::getTargetEnvAttrName() { return "spv.target_env"; }
 
 spirv::TargetEnvAttr spirv::getDefaultTargetEnv(MLIRContext *context) {
   auto triple = spirv::VerCapExtAttr::get(spirv::Version::V_1_0,
                                           {spirv::Capability::Shader},
                                           ArrayRef<Extension>(), context);