[mlir][spirv] Add a field for client API in target environment

[antiagainst]

Summary

SPIR-V can be directly consumed by APIs like Vulkan and OpenCL,
where we can use the capability list to diffferentiate. It can
also be used as a compilation target to transcompile to shading
languages like WGSL to target WebGPU. We have no way to tell
that with just the capability list, so we cannot perform certain
transformations only applicable to those targets thus far. So
this commit add a field in the target environment to indicate
the client API for such purposes.

[kuhar]

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

299

What happens when we need to move a new API in the future? Will it become case 5 while Unknown remains case 4? Instead, would it make sense to define Unknown as case 0 or -1?

[kuhar]

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

299

s/move a new/add a new/

[antiagainst]

Address comments

[antiagainst]

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

299

Good questions. These are meant to be "StrEnumAttr" actually, where the exact number is not fixed and there is no guarantee it won't change in the future. Only the symbol is meaningful here. But we don't have StrEnumAttr anymore.

Using -1 for unknown makes sense. Changed.

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

: 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, Vendor vendorID,
-                           DeviceType deviceType, uint32_t deviceId,
-                           ResourceLimitsAttr limits);
+  static TargetEnvAttr get(VerCapExtAttr triple, ResourceLimitsAttr limits,
+                           ClientAPI clientAPI = ClientAPI::Unknown,
+                           Vendor vendorID = Vendor::Unknown,
+                           DeviceType deviceType = DeviceType::Unknown,
+                           uint32_t deviceId = kUnknownDeviceID);
 
   /// Returns the attribute kind's name (without the 'spirv.' prefix).
   static StringRef getKindName();
 
   /// Returns the (version, capabilities, extensions) triple attribute.
   VerCapExtAttr getTripleAttr() const;
 
   /// Returns the target version.
   Version getVersion() const;
 
   /// 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 client API.
+  ClientAPI getClientAPI() const;
+
   /// Returns the vendor ID.
   Vendor getVendorID() const;
 
   /// Returns the device type.
   DeviceType getDeviceType() const;
 
   /// Returns the device ID.
   uint32_t getDeviceID() const;
 
   /// Returns the target resource limits.
   ResourceLimitsAttr getResourceLimits() const;
 };
 } // namespace spirv
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPIRV_IR_SPIRVATTRIBUTES_H

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

 //===----------------------------------------------------------------------===//
 
 def SPIRV_DT_CPU           : I32EnumAttrCase<"CPU", 0>;
 def SPIRV_DT_DiscreteGPU   : I32EnumAttrCase<"DiscreteGPU", 1>;
 def SPIRV_DT_IntegratedGPU : I32EnumAttrCase<"IntegratedGPU", 2>;
 // An accelerator other than GPU or CPU
 def SPIRV_DT_Other         : I32EnumAttrCase<"Other", 3>;
 // Information missing.
-def SPIRV_DT_Unknown       : I32EnumAttrCase<"Unknown", 4>;
+def SPIRV_DT_Unknown       : I32EnumAttrCase<"Unknown", 0xffffffff>;
 
 def SPIRV_DeviceTypeAttr : SPIRV_I32EnumAttr<
   "DeviceType", "valid SPIR-V device types", "device_type", [
     SPIRV_DT_Other, SPIRV_DT_IntegratedGPU, SPIRV_DT_DiscreteGPU,
     SPIRV_DT_CPU, SPIRV_DT_Unknown
   ]>;
 
 def SPIRV_V_AMD         : I32EnumAttrCase<"AMD", 0>;
 def SPIRV_V_Apple       : I32EnumAttrCase<"Apple", 1>;
 def SPIRV_V_ARM         : I32EnumAttrCase<"ARM", 2>;
 def SPIRV_V_Imagination : I32EnumAttrCase<"Imagination", 3>;
 def SPIRV_V_Intel       : I32EnumAttrCase<"Intel", 4>;
 def SPIRV_V_NVIDIA      : I32EnumAttrCase<"NVIDIA", 5>;
 def SPIRV_V_Qualcomm    : I32EnumAttrCase<"Qualcomm", 6>;
 def SPIRV_V_SwiftShader : I32EnumAttrCase<"SwiftShader", 7>;
-def SPIRV_V_Unknown     : I32EnumAttrCase<"Unknown", 0xff>;
+def SPIRV_V_Unknown     : I32EnumAttrCase<"Unknown", 0xffffffff>;
 
 def SPIRV_VendorAttr : SPIRV_I32EnumAttr<
   "Vendor", "recognized SPIR-V vendor strings", "vendor", [
     SPIRV_V_AMD, SPIRV_V_Apple, SPIRV_V_ARM, SPIRV_V_Imagination,
     SPIRV_V_Intel, SPIRV_V_NVIDIA, SPIRV_V_Qualcomm, SPIRV_V_SwiftShader,
     SPIRV_V_Unknown
   ]>;
 
+def SPIRV_CA_Metal   : I32EnumAttrCase<"Metal", 0>;
+def SPIRV_CA_OpenCL  : I32EnumAttrCase<"OpenCL", 1>;
+def SPIRV_CA_Vulkan  : I32EnumAttrCase<"Vulkan", 2>;
+def SPIRV_CA_WebGPU  : I32EnumAttrCase<"WebGPU", 3>;
+def SPIRV_CA_Unknown : I32EnumAttrCase<"Unknown", 0xffffffff>;
+
+def SPIRV_ClientAPIAttr : SPIRV_I32EnumAttr<
+  "ClientAPI", "recognized SPIR-V client APIs", "client_api", [
+    SPIRV_CA_Metal, SPIRV_CA_OpenCL, SPIRV_CA_Vulkan, SPIRV_CA_WebGPU,
+    SPIRV_CA_Unknown
+  ]>;
+
 //===----------------------------------------------------------------------===//
 // SPIR-V extension definitions
 //===----------------------------------------------------------------------===//
 
 // Extensions known to the SPIR-V dialect.
 // https://github.com/KhronosGroup/SPIRV-Registry has the full list.
 // These start with the 'SPV' prefix to match the Khronos' naming scheme.
 def SPV_KHR_16bit_storage                    : I32EnumAttrCase<"SPV_KHR_16bit_storage", 0>;

mlir/lib/Dialect/SPIRV/IR/SPIRVAttributes.cpp

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

 spirv::VerCapExtAttr::cap_range spirv::TargetEnvAttr::getCapabilities() {
   return getTripleAttr().getCapabilities();
 }
 
 ArrayAttr spirv::TargetEnvAttr::getCapabilitiesAttr() {
   return getTripleAttr().getCapabilitiesAttr();
 }
 
+spirv::ClientAPI spirv::TargetEnvAttr::getClientAPI() const {
+  return getImpl()->clientAPI;
+}
+
 spirv::Vendor spirv::TargetEnvAttr::getVendorID() const {
   return getImpl()->vendorID;
 }
 
 spirv::DeviceType spirv::TargetEnvAttr::getDeviceType() const {
   return getImpl()->deviceType;
 }
 

 static Attribute parseTargetEnvAttr(DialectAsmParser &parser) {
   if (parser.parseLess())
     return {};
 
   spirv::VerCapExtAttr tripleAttr;
   if (parser.parseAttribute(tripleAttr) || parser.parseComma())
     return {};
 
+  auto clientAPI = spirv::ClientAPI::Unknown;
+  if (succeeded(parser.parseOptionalKeyword("api"))) {
+    if (parser.parseEqual())
+      return {};
+    auto loc = parser.getCurrentLocation();
+    StringRef apiStr;
+    if (parser.parseKeyword(&apiStr))
+      return {};
+    if (auto apiSymbol = spirv::symbolizeClientAPI(apiStr))
+      clientAPI = *apiSymbol;
+    else
+      parser.emitError(loc, "unknown client API: ") << apiStr;
+    if (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)) {
+      if (auto vendorSymbol = spirv::symbolizeVendor(vendorStr))
         vendorID = *vendorSymbol;
-      } else {
+      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)) {
+        if (auto deviceTypeSymbol = spirv::symbolizeDeviceType(deviceTypeStr))
           deviceType = *deviceTypeSymbol;
-        } else {
+        else
           parser.emitError(loc, "unknown device type: ") << deviceTypeStr;
-        }
 
         if (succeeded(parser.parseOptionalColon())) {
           loc = parser.getCurrentLocation();
           if (parser.parseInteger(deviceID))
             return {};
         }
       }
       if (parser.parseComma())
         return {};
     }
   }
 
   ResourceLimitsAttr limitsAttr;
   if (parser.parseAttribute(limitsAttr) || parser.parseGreater())
     return {};
 
-  return spirv::TargetEnvAttr::get(tripleAttr, vendorID, deviceType, deviceID,
-                                   limitsAttr);
+  return spirv::TargetEnvAttr::get(tripleAttr, limitsAttr, clientAPI, vendorID,
+                                   deviceType, deviceID);
 }
 
 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() << "<#spirv.";
   print(targetEnv.getTripleAttr(), printer);
+  auto clientAPI = targetEnv.getClientAPI();
+  if (clientAPI != spirv::ClientAPI::Unknown)
+    printer << ", api=" << clientAPI;
   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)

mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp

 //===- TargetAndABI.cpp - SPIR-V target and ABI utilities -----------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SPIRV/IR/TargetAndABI.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/FunctionInterfaces.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/SymbolTable.h"
 
 using namespace mlir;
 

 }
 
 StringRef spirv::getTargetEnvAttrName() { return "spirv.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, spirv::Vendor::Unknown,
-                                   spirv::DeviceType::Unknown,
-                                   spirv::TargetEnvAttr::kUnknownDeviceID,
-                                   spirv::getDefaultResourceLimits(context));
+  return spirv::TargetEnvAttr::get(
+      triple, spirv::getDefaultResourceLimits(context),
+      spirv::ClientAPI::Unknown, spirv::Vendor::Unknown,
+      spirv::DeviceType::Unknown, spirv::TargetEnvAttr::kUnknownDeviceID);
 }
 
 spirv::TargetEnvAttr spirv::lookupTargetEnv(Operation *op) {
   while (op) {
     op = SymbolTable::getNearestSymbolTable(op);
     if (!op)
       break;
 

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

spirv.target_env = #spirv.target_env<
     #spirv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>,
     #spirv.resource_limits<
       max_compute_workgroup_size = [128, 64, 64]
     >>
 } { return }
 
 // -----
 
+func.func @target_env_client_api() attributes {
+  // CHECK:      spirv.target_env = #spirv.target_env<
+  // CHECK-SAME:   #spirv.vce<v1.0, [], []>,
+  // CHECK-SAME:   api=Metal,
+  // CHECK-SAME:   #spirv.resource_limits<>>
+  spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [], []>, api=Metal, #spirv.resource_limits<>>
+} { return }
+
+// -----
+
+func.func @target_env_client_api() attributes {
+  // CHECK:      spirv.target_env = #spirv.target_env
+  // CHECK-NOT:   api=
+  spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [], []>, api=Unknown, #spirv.resource_limits<>>
+} { return }
+
+// -----
+
 func.func @target_env_vendor_id() attributes {
   // CHECK:      spirv.target_env = #spirv.target_env<
   // CHECK-SAME:   #spirv.vce<v1.0, [], []>,
   // CHECK-SAME:   NVIDIA,
   // CHECK-SAME:   #spirv.resource_limits<>>
   spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [], []>, NVIDIA, #spirv.resource_limits<>>
 } { return }
 
func.func @target_env_vendor_id_device_type_device_id() attributes {
   // CHECK-SAME:   #spirv.vce<v1.0, [], []>,
   // CHECK-SAME:   Qualcomm:IntegratedGPU:100925441,
   // CHECK-SAME:   #spirv.resource_limits<>>
   spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [], []>, Qualcomm:IntegratedGPU:0x6040001, #spirv.resource_limits<>>
 } { return }
 
 // -----
 
+func.func @target_env_client_api_vendor_id_device_type_device_id() attributes {
+  // CHECK:      spirv.target_env = #spirv.target_env<
+  // CHECK-SAME:   #spirv.vce<v1.0, [], []>,
+  // CHECK-SAME:   api=Vulkan,
+  // CHECK-SAME:   Qualcomm:IntegratedGPU:100925441,
+  // CHECK-SAME:   #spirv.resource_limits<>>
+  spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [], []>, api=Vulkan, Qualcomm:IntegratedGPU:0x6040001, #spirv.resource_limits<>>
+} { return }
+
+// -----
+
 func.func @target_env_extra_fields() attributes {
   // expected-error @+3 {{expected '>'}}
   spirv.target_env = #spirv.target_env<
     #spirv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>,
     #spirv.resource_limits<>,
     more_stuff
   >
 } { return }