[spirv] Move device info from resource limit into target env

[antiagainst]

Summary

Vendor/device information are not resource limits. Moving to
target environment directly for better organization.

[antiagainst]

Fix

[mravishankar]

Nice. Just a few nits.

mlir/docs/Dialects/SPIR-V.md

810

Whats the definition for spirv-device-id? Is it just an integer.

mlir/include/mlir/Dialect/SPIRV/SPIRVAttributes.h

135

Make deviceId uint64_t?

[antiagainst]

mlir/docs/Dialects/SPIR-V.md

810

Yup. Thanks for catching that! Added.

mlir/include/mlir/Dialect/SPIRV/SPIRVAttributes.h

135

They are uint32_t in https://www.khronos.org/registry/vulkan/specs/1.2-extensions/man/html/VkPhysicalDeviceProperties.html.

mlir/docs/Dialects/SPIR-V.md

 
 spirv-resource-limits ::= dictionary-attribute
 
 spirv-vce-attribute ::= `#` `spv.vce` `<`
                             spirv-version `,`
                             spirv-capability-list `,`
                             spirv-extensions-list `>`
 
+spirv-vendor-id ::= `AMD` | `NVIDIA` | ...
+spirv-device-type ::= `DiscreteGPU` | `IntegratedGPU` | `CPU` | ...
+spirv-device-info ::= spirv-vendor-id (`:` spirv-device-type (`:` spirv-device-id)?)?
+
 spirv-target-env-attribute ::= `#` `spv.target_env` `<`
                                   spirv-vce-attribute,
+                                  (spirv-device-info `,`)?
                                   spirv-resource-limits `>`
 ```
 
 The attribute has a few fields:
 
 *   A `#spv.vce` (`spirv::VerCapExtAttr`) attribute:
     *   The target SPIR-V version.
     *   A list of SPIR-V extensions for the target.
     *   A list of SPIR-V capabilities for the target.
 *   A dictionary of target resource limits (see the
     [Vulkan spec][VulkanResourceLimits] for explanation):
     *   `max_compute_workgroup_invocations`
     *   `max_compute_workgroup_size`
 
 For example,
 
 ```
 module attributes {
 spv.target_env = #spv.target_env<
     #spv.vce<v1.3, [Shader, GroupNonUniform], [SPV_KHR_8bit_storage]>,
+    ARM:IntegratedGPU,
     {
       max_compute_workgroup_invocations = 128 : i32,
       max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>
     }>
 } { ... }
 ```
 
 Dialect conversion framework will utilize the information in `spv.target_env` to

mlir/include/mlir/Dialect/SPIRV/SPIRVAttributes.h

 #include "mlir/Support/LLVM.h"
 
 // Pull in SPIR-V attribute definitions for target and ABI.
 #include "mlir/Dialect/SPIRV/TargetAndABI.h.inc"
 
 namespace mlir {
 namespace spirv {
 enum class Capability : uint32_t;
+enum class DeviceType;
 enum class Extension;
+enum class Vendor;
 enum class Version : uint32_t;
 
 namespace detail {
 struct InterfaceVarABIAttributeStorage;
 struct TargetEnvAttributeStorage;
 struct VerCapExtAttributeStorage;
 } // namespace detail
 

 
 /// An attribute that specifies the target version, allowed extensions and
 /// capabilities, and resource limits. These information describes a SPIR-V
 /// target environment.
 class TargetEnvAttr
     : public Attribute::AttrBase<TargetEnvAttr, Attribute,
                                  detail::TargetEnvAttributeStorage> {
 public:
+  /// ID for unknown devices.
+  static constexpr uint32_t kUnknownDeviceID = 0x7FFFFFFF;
+
   using Base::Base;
 
   /// Gets a TargetEnvAttr instance.
-  static TargetEnvAttr get(VerCapExtAttr triple, DictionaryAttr limits);
+  static TargetEnvAttr get(VerCapExtAttr triple, Vendor vendorID,
+                           DeviceType deviceType, uint32_t deviceId,
+                           DictionaryAttr limits);
 
   /// Returns the attribute kind's name (without the 'spv.' prefix).
   static StringRef getKindName();
 
   /// Returns the (version, capabilities, extensions) triple attribute.
   VerCapExtAttr getTripleAttr();
 
   /// Returns the target version.
   Version getVersion();
 
   /// Returns the target extensions.
   VerCapExtAttr::ext_range getExtensions();
   /// Returns the target extensions as a string array attribute.
   ArrayAttr getExtensionsAttr();
 
   /// Returns the target capabilities.
   VerCapExtAttr::cap_range getCapabilities();
   /// Returns the target capabilities as an integer array attribute.
   ArrayAttr getCapabilitiesAttr();
 
+  /// Returns the vendor ID.
+  Vendor getVendorID();
+
+  /// Returns the device type.
+  DeviceType getDeviceType();
+
+  /// Returns the device ID.
+  uint32_t getDeviceID();
+
   /// Returns the target resource limits.
   ResourceLimitsAttr getResourceLimits();
 
-  static LogicalResult verifyConstructionInvariants(Location loc,
-                                                    VerCapExtAttr triple,
-                                                    DictionaryAttr limits);
+  static LogicalResult
+  verifyConstructionInvariants(Location loc, VerCapExtAttr triple,
+                               Vendor vendorID, DeviceType deviceType,
+                               uint32_t deviceID, DictionaryAttr limits);
 };
 } // namespace spirv
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPIRV_SPIRVATTRIBUTES_H

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

let methods = [InterfaceMethod<
     "::llvm::SmallVector<::llvm::ArrayRef<::mlir::spirv::Capability>, 1>",
     "getCapabilities">];
 }
 
 //===----------------------------------------------------------------------===//
 // SPIR-V target GPU vendor and device definitions
 //===----------------------------------------------------------------------===//
 
+def SPV_DT_CPU           : StrEnumAttrCase<"CPU">;
+def SPV_DT_DiscreteGPU   : StrEnumAttrCase<"DiscreteGPU">;
+def SPV_DT_IntegratedGPU : StrEnumAttrCase<"IntegratedGPU">;
 // 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>;
+def SPV_DT_Other         : StrEnumAttrCase<"Other">;
 // Information missing.
-def SPV_DT_Unknown       : I32EnumAttrCase<"Unknown", 0x7FFFFFFF>;
+def SPV_DT_Unknown       : StrEnumAttrCase<"Unknown">;
 
-def SPV_DeviceTypeAttr : SPV_I32EnumAttr<
+def SPV_DeviceTypeAttr : SPV_StrEnumAttr<
   "DeviceType", "valid SPIR-V device types", [
     SPV_DT_Other, SPV_DT_IntegratedGPU, SPV_DT_DiscreteGPU,
     SPV_DT_CPU, SPV_DT_Unknown
   ]>;
 
+def SPV_V_AMD         : StrEnumAttrCase<"AMD">;
+def SPV_V_ARM         : StrEnumAttrCase<"ARM">;
+def SPV_V_IMAGINATION : StrEnumAttrCase<"Imagination">;
+def SPV_V_Intel       : StrEnumAttrCase<"Intel">;
+def SPV_V_NVIDIA      : StrEnumAttrCase<"NVIDIA">;
+def SPV_V_Qualcomm    : StrEnumAttrCase<"Qualcomm">;
+def SPV_V_SwiftShader : StrEnumAttrCase<"SwiftShader">;
+def SPV_V_Unknown     : StrEnumAttrCase<"Unknown">;
+
+def SPV_VendorAttr : SPV_StrEnumAttr<
+  "Vendor", "recognized SPIR-V vendor strings", [
+    SPV_V_AMD, SPV_V_ARM, SPV_V_IMAGINATION, SPV_V_Intel,
+    SPV_V_NVIDIA, SPV_V_Qualcomm, SPV_V_SwiftShader,
+    SPV_V_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">;

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",

mlir/lib/Dialect/SPIRV/SPIRVAttributes.cpp

struct VerCapExtAttributeStorage : public AttributeStorage {
   }
 
   Attribute version;
   Attribute capabilities;
   Attribute extensions;
 };
 
 struct TargetEnvAttributeStorage : public AttributeStorage {
-  using KeyTy = std::pair<Attribute, Attribute>;
+  using KeyTy = std::tuple<Attribute, Vendor, DeviceType, uint32_t, Attribute>;
 
-  TargetEnvAttributeStorage(Attribute triple, Attribute limits)
-      : triple(triple), limits(limits) {}
+  TargetEnvAttributeStorage(Attribute triple, Vendor vendorID,
+                            DeviceType deviceType, uint32_t deviceID,
+                            Attribute limits)
+      : triple(triple), limits(limits), vendorID(vendorID),
+        deviceType(deviceType), deviceID(deviceID) {}
 
   bool operator==(const KeyTy &key) const {
-    return key.first == triple && key.second == limits;
+    return key ==
+           std::make_tuple(triple, vendorID, deviceType, deviceID, limits);
   }
 
   static TargetEnvAttributeStorage *
   construct(AttributeStorageAllocator &allocator, const KeyTy &key) {
     return new (allocator.allocate<TargetEnvAttributeStorage>())
-        TargetEnvAttributeStorage(key.first, key.second);
+        TargetEnvAttributeStorage(std::get<0>(key), std::get<1>(key),
+                                  std::get<2>(key), std::get<3>(key),
+                                  std::get<4>(key));
   }
 
   Attribute triple;
   Attribute limits;
+  Vendor vendorID;
+  DeviceType deviceType;
+  uint32_t deviceID;
 };
 } // namespace detail
 } // namespace spirv
 } // namespace mlir
 
 //===----------------------------------------------------------------------===//
 // InterfaceVarABIAttr
 //===----------------------------------------------------------------------===//
LogicalResult spirv::VerCapExtAttr::verifyConstructionInvariants(
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // TargetEnvAttr
 //===----------------------------------------------------------------------===//
 
 spirv::TargetEnvAttr spirv::TargetEnvAttr::get(spirv::VerCapExtAttr triple,
+                                               Vendor vendorID,
+                                               DeviceType deviceType,
+                                               uint32_t deviceID,
                                                DictionaryAttr limits) {
   assert(triple && limits && "expected valid triple and limits");
   MLIRContext *context = triple.getContext();
-  return Base::get(context, triple, limits);
+  return Base::get(context, triple, vendorID, deviceType, deviceID, limits);
 }
 
 StringRef spirv::TargetEnvAttr::getKindName() { return "target_env"; }
 
 spirv::VerCapExtAttr spirv::TargetEnvAttr::getTripleAttr() {
   return getImpl()->triple.cast<spirv::VerCapExtAttr>();
 }
 

 spirv::VerCapExtAttr::cap_range spirv::TargetEnvAttr::getCapabilities() {
   return getTripleAttr().getCapabilities();
 }
 
 ArrayAttr spirv::TargetEnvAttr::getCapabilitiesAttr() {
   return getTripleAttr().getCapabilitiesAttr();
 }
 
+spirv::Vendor spirv::TargetEnvAttr::getVendorID() {
+  return getImpl()->vendorID;
+}
+
+spirv::DeviceType spirv::TargetEnvAttr::getDeviceType() {
+  return getImpl()->deviceType;
+}
+
+uint32_t spirv::TargetEnvAttr::getDeviceID() { return getImpl()->deviceID; }
+
 spirv::ResourceLimitsAttr spirv::TargetEnvAttr::getResourceLimits() {
   return getImpl()->limits.cast<spirv::ResourceLimitsAttr>();
 }
 
 LogicalResult spirv::TargetEnvAttr::verifyConstructionInvariants(
-    Location loc, spirv::VerCapExtAttr triple, DictionaryAttr limits) {
+    Location loc, spirv::VerCapExtAttr /*triple*/, spirv::Vendor /*vendorID*/,
+    spirv::DeviceType /*deviceType*/, uint32_t /*deviceID*/,
+    DictionaryAttr limits) {
   if (!limits.isa<spirv::ResourceLimitsAttr>())
     return emitError(loc, "expected spirv::ResourceLimitsAttr for limits");
 
   return success();
 }

mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp

 static Attribute parseTargetEnvAttr(DialectAsmParser &parser) {
   if (parser.parseLess())
     return {};
 
   spirv::VerCapExtAttr tripleAttr;
   if (parser.parseAttribute(tripleAttr) || parser.parseComma())
     return {};
 
+  // Parse [vendor[:device-type[:device-id]]]
+  Vendor vendorID = Vendor::Unknown;
+  DeviceType deviceType = DeviceType::Unknown;
+  uint32_t deviceID = spirv::TargetEnvAttr::kUnknownDeviceID;
+  {
+    auto loc = parser.getCurrentLocation();
+    StringRef vendorStr;
+    if (succeeded(parser.parseOptionalKeyword(&vendorStr))) {
+      if (auto vendorSymbol = spirv::symbolizeVendor(vendorStr)) {
+        vendorID = *vendorSymbol;
+      } else {
+        parser.emitError(loc, "unknown vendor: ") << vendorStr;
+      }
+
+      if (succeeded(parser.parseOptionalColon())) {
+        loc = parser.getCurrentLocation();
+        StringRef deviceTypeStr;
+        if (parser.parseKeyword(&deviceTypeStr))
+          return {};
+        if (auto deviceTypeSymbol = spirv::symbolizeDeviceType(deviceTypeStr)) {
+          deviceType = *deviceTypeSymbol;
+        } else {
+          parser.emitError(loc, "unknown device type: ") << deviceTypeStr;
+        }
+
+        if (succeeded(parser.parseOptionalColon())) {
+          loc = parser.getCurrentLocation();
+          if (parser.parseInteger(deviceID))
+            return {};
+        }
+      }
+      if (parser.parseComma())
+        return {};
+    }
+  }
+
   DictionaryAttr limitsAttr;
   {
     auto loc = parser.getCurrentLocation();
     if (parser.parseAttribute(limitsAttr))
       return {};
 
     if (!limitsAttr.isa<spirv::ResourceLimitsAttr>()) {
       parser.emitError(
           loc,
           "limits must be a dictionary attribute containing two 32-bit integer "
           "attributes 'max_compute_workgroup_invocations' and "
           "'max_compute_workgroup_size'");
       return {};
     }
   }
 
   if (parser.parseGreater())
     return {};
 
-  return spirv::TargetEnvAttr::get(tripleAttr, limitsAttr);
+  return spirv::TargetEnvAttr::get(tripleAttr, vendorID, deviceType, deviceID,
+                                   limitsAttr);
 }
 
 Attribute SPIRVDialect::parseAttribute(DialectAsmParser &parser,
                                        Type type) const {
   // SPIR-V attributes are dictionaries so they do not have type.
   if (type) {
     parser.emitError(parser.getNameLoc(), "unexpected type");
     return {};
llvm::interleaveComma(triple.getExtensionsAttr(), os, [&](Attribute attr) {
     os << attr.cast<StringAttr>().getValue();
   });
   printer << "]>";
 }
 
 static void print(spirv::TargetEnvAttr targetEnv, DialectAsmPrinter &printer) {
   printer << spirv::TargetEnvAttr::getKindName() << "<#spv.";
   print(targetEnv.getTripleAttr(), printer);
+  spirv::Vendor vendorID = targetEnv.getVendorID();
+  spirv::DeviceType deviceType = targetEnv.getDeviceType();
+  uint32_t deviceID = targetEnv.getDeviceID();
+  if (vendorID != spirv::Vendor::Unknown) {
+    printer << ", " << spirv::stringifyVendor(vendorID);
+    if (deviceType != spirv::DeviceType::Unknown) {
+      printer << ":" << spirv::stringifyDeviceType(deviceType);
+      if (deviceID != spirv::TargetEnvAttr::kUnknownDeviceID)
+        printer << ":" << deviceID;
+    }
+  }
   printer << ", " << targetEnv.getResourceLimits() << ">";
 }
 
 static void print(spirv::InterfaceVarABIAttr interfaceVarABIAttr,
                   DialectAsmPrinter &printer) {
   printer << spirv::InterfaceVarABIAttr::getKindName() << "<("
           << interfaceVarABIAttr.getDescriptorSet() << ", "
           << interfaceVarABIAttr.getBinding() << ")";

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) {
   return {};
 }
 
 spirv::ResourceLimitsAttr
 spirv::getDefaultResourceLimits(MLIRContext *context) {
   // All the fields have default values. Here we just provide a nicer way to
   // construct a default resource limit attribute.
   return spirv::ResourceLimitsAttr ::get(
-      /*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);
-  return spirv::TargetEnvAttr::get(triple,
+  return spirv::TargetEnvAttr::get(triple, spirv::Vendor::Unknown,
+                                   spirv::DeviceType::Unknown,
+                                   spirv::TargetEnvAttr::kUnknownDeviceID,
                                    spirv::getDefaultResourceLimits(context));
 }
 
 spirv::TargetEnvAttr spirv::lookupTargetEnv(Operation *op) {
   while (op) {
     op = SymbolTable::getNearestSymbolTable(op);
     if (!op)
       break;

mlir/test/Dialect/SPIRV/target-and-abi.mlir

spv.target_env = #spv.target_env<
     {
       max_compute_workgroup_invocations = 128 : i32,
       max_compute_workgroup_size = dense<[128, 64, 64]> : vector<3xi32>
     }>
 } { return }
 
 // -----
 
+func @target_env_vendor_id() attributes {
+  // CHECK:      spv.target_env = #spv.target_env<
+  // CHECK-SAME:   #spv.vce<v1.0, [], []>,
+  // CHECK-SAME:   NVIDIA,
+  // CHECK-SAME:   {}>
+  spv.target_env = #spv.target_env<#spv.vce<v1.0, [], []>, NVIDIA, {}>
+} { return }
+
+// -----
+
+func @target_env_vendor_id_device_type() attributes {
+  // CHECK:      spv.target_env = #spv.target_env<
+  // CHECK-SAME:   #spv.vce<v1.0, [], []>,
+  // CHECK-SAME:   AMD:DiscreteGPU,
+  // CHECK-SAME:   {}>
+  spv.target_env = #spv.target_env<#spv.vce<v1.0, [], []>, AMD:DiscreteGPU, {}>
+} { return }
+
+// -----
+
+func @target_env_vendor_id_device_type_device_id() attributes {
+  // CHECK:      spv.target_env = #spv.target_env<
+  // CHECK-SAME:   #spv.vce<v1.0, [], []>,
+  // CHECK-SAME:   Qualcomm:IntegratedGPU:100925441,
+  // CHECK-SAME:   {}>
+  spv.target_env = #spv.target_env<#spv.vce<v1.0, [], []>, Qualcomm:IntegratedGPU:0x6040001, {}>
+} { return }
+
+// -----
+
 func @target_env_extra_fields() attributes {
   // expected-error @+6 {{expected '>'}}
   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, 64, 64]> : vector<3xi32>
     },