Reland "[spirv] Switch to use common target description" (iree-org#17699

)

This relands iree-org#17623.

This commit switches SPIR-V side to use the common `#iree_gpu.target` to
describe the GPU characteristics. With it we can now remove the ad-hoc
Vulkan attributes and dialects and unify how GPU are described across
various GPU compiler backends in IREE.

SPIR-V has some additional requirements that we need to account for:

We have many vendors and APIs to handle there so this commit adds
various AMD/ARM/NVIDIA/Qualcomm targets for
development purposes so that we can specify them with a shorthand.

In order to be extensible, leverage the `feature` field in
`#iree_gpu.target` to specify additional capabilities with `cap:` prefix
and extensions with `ext:` prefix. We also use the `feature` field to
specify what SPIR-V version to target with the `spirv:v1.x` format.

Right now the `SPIRVConvertGPUTarget` pass is
invoked immediately before configuration. This is to stage the changes.
As a next step we need to move
it immediately before `ConvertToSPIRV` pass.

`--iree-vulkan-target-env` is dropped given now we removed the whole
Vulkan dialect and cannot control with a `#vk.target_env` attribute
anymore.

The default `--iree-vulkan-target-triple` now becomes
`vp_android_baseline_2022`, which is a a good lowest common denominator
to guarantee the generated SPIR-V is widely accepted. We are not
considering SwiftShader now anymore like previously due to testing
purposes.

The `--iree-vulkan-target-triple` should be renamed given it's not a
triple anymore--that will happen later together with other GPU backends
(i.e., cuda/hip) to be consistent.

In order to support cooperative matrix conversion, we added
`WMMA_F16_16x16x16_F16`. For NVIDIA GPUs
we are abusing it right now without considering the concrete explicit
layout--that is fine given in Vulkan they are opaque anyway. But this
need to be fixed if we are targeting WMMA in CUDA.

We now contruct a `#iree_gpu.target` to specify
the target to drive SPIR-V CodeGen.

Progress towards iree-org#16341

ci-extra: test_nvidia_gpu,test_nvidia_a100,test_amd_mi250,
build_test_all_macos_arm64,build_and_test_android,test_on_moto-edge-x30

---------

Signed-off-by: Lei Zhang <antiagainst@gmail.com>

compiler/plugins/target/MetalSPIRV/BUILD.bazel

@@ -26,6 +26,7 @@ iree_compiler_cc_library(
         ":SPIRVToMSL",
         "//compiler/src/iree/compiler/Codegen/Common",
         "//compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR:IREECodegenDialect",
+        "//compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils:KnownTargets",
         "//compiler/src/iree/compiler/Codegen/SPIRV",
         "//compiler/src/iree/compiler/Codegen/Utils",
         "//compiler/src/iree/compiler/Dialect/Flow/IR",

compiler/plugins/target/MetalSPIRV/CMakeLists.txt

@@ -36,6 +36,7 @@ iree_cc_library(
     MLIRVectorDialect
     iree::compiler::Codegen::Common
     iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect
+    iree::compiler::Codegen::Dialect::GPU::TargetUtils::KnownTargets
     iree::compiler::Codegen::SPIRV
     iree::compiler::Codegen::Utils
     iree::compiler::Dialect::Flow::IR

compiler/plugins/target/MetalSPIRV/MetalSPIRVTarget.cpp

@@ -8,6 +8,7 @@
 #include "compiler/plugins/target/MetalSPIRV/MetalTargetPlatform.h"
 #include "compiler/plugins/target/MetalSPIRV/SPIRVToMSL.h"
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h"
+#include "iree/compiler/Codegen/Dialect/GPU/TargetUtils/KnownTargets.h"
 #include "iree/compiler/Codegen/SPIRV/Passes.h"
 #include "iree/compiler/Dialect/Flow/IR/FlowDialect.h"
 #include "iree/compiler/Dialect/HAL/Target/TargetRegistry.h"
@@ -19,9 +20,7 @@
 #include "llvm/TargetParser/Triple.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
-#include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
-#include "mlir/Dialect/SPIRV/IR/TargetAndABI.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Target/SPIRV/Serialization.h"
 
@@ -52,60 +51,6 @@ struct MetalSPIRVOptions {
 };
 } // namespace
 
-static spirv::TargetEnvAttr getMetalTargetEnv(MLIRContext *context) {
-  using spirv::Capability;
-  using spirv::Extension;
-
-  // Capabilities and limits according to Metal 3 devices.
-  const std::array<Extension, 4> extensions = {
-      Extension::SPV_KHR_16bit_storage,
-      Extension::SPV_KHR_8bit_storage,
-      Extension::SPV_KHR_storage_buffer_storage_class,
-      Extension::SPV_KHR_variable_pointers,
-  };
-  const std::array<Capability, 21> capabilities = {
-      Capability::Shader,
-      Capability::Int8,
-      Capability::Int16,
-      Capability::Int64,
-      Capability::Float16,
-      Capability::UniformAndStorageBuffer8BitAccess,
-      Capability::StorageBuffer8BitAccess,
-      Capability::StoragePushConstant8,
-      Capability::StorageUniform16,
-      Capability::StorageBuffer16BitAccess,
-      Capability::StoragePushConstant16,
-      Capability::GroupNonUniform,
-      Capability::GroupNonUniformVote,
-      Capability::GroupNonUniformArithmetic,
-      Capability::GroupNonUniformBallot,
-      Capability::GroupNonUniformShuffle,
-      Capability::GroupNonUniformShuffleRelative,
-      Capability::GroupNonUniformQuad,
-      Capability::StoragePushConstant16,
-      Capability::VariablePointers,
-      Capability::VariablePointersStorageBuffer,
-  };
-  auto limits = spirv::ResourceLimitsAttr::get(
-      context,
-      /*max_compute_shared_memory_size=*/32768,
-      /*max_compute_workgroup_invocations=*/1024,
-      /*max_compute_workgroup_size=*/
-      Builder(context).getI32ArrayAttr({1024, 1024, 1024}),
-      /*subgroup_size=*/32,
-      /*min_subgroup_size=*/std::nullopt,
-      /*max_subgroup_size=*/std::nullopt,
-      /*cooperative_matrix_properties_khr=*/ArrayAttr{},
-      /*cooperative_matrix_properties_nv=*/ArrayAttr{});
-
-  auto triple = spirv::VerCapExtAttr::get(spirv::Version::V_1_3, capabilities,
-                                          extensions, context);
-  // Further assuming Apple GPUs.
-  return spirv::TargetEnvAttr::get(
-      triple, limits, spirv::ClientAPI::Metal, spirv::Vendor::Apple,
-      spirv::DeviceType::IntegratedGPU, spirv::TargetEnvAttr::kUnknownDeviceID);
-}
-
 // TODO: MetalOptions for choosing the Metal version.
 class MetalTargetDevice : public TargetDevice {
 public:
@@ -145,20 +90,20 @@ class MetalSPIRVTargetBackend : public TargetBackend {
       MLIRContext *context, StringRef deviceID, DictionaryAttr deviceConfigAttr,
       SmallVectorImpl<IREE::HAL::ExecutableTargetAttr> &executableTargetAttrs)
       const override {
-    executableTargetAttrs.push_back(
-        getExecutableTarget(context, getMetalTargetEnv(context)));
+    executableTargetAttrs.push_back(getExecutableTarget(context));
   }
 
   IREE::HAL::ExecutableTargetAttr
-  getExecutableTarget(MLIRContext *context,
-                      spirv::TargetEnvAttr targetEnv) const {
+  getExecutableTarget(MLIRContext *context) const {
     Builder b(context);
     SmallVector<NamedAttribute> configItems;
     auto addConfig = [&](StringRef name, Attribute value) {
       configItems.emplace_back(b.getStringAttr(name), value);
     };
 
-    addConfig(spirv::getTargetEnvAttrName(), targetEnv);
+    if (auto target = GPU::getMetalTargetDetails(context)) {
+      addConfig("iree.gpu.target", target);
+    }
 
     return b.getAttr<IREE::HAL::ExecutableTargetAttr>(
         b.getStringAttr("metal-spirv"), b.getStringAttr("metal-msl-fb"),

compiler/plugins/target/MetalSPIRV/test/smoketest.mlir

@@ -4,7 +4,9 @@ module attributes {
   hal.device.targets = [
     #hal.device.target<"metal", [
       #hal.executable.target<"metal-spirv", "metal-msl-fb", {
-        spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, #spirv.resource_limits<>>
+        iree.gpu.target = #iree_gpu.target<arch = "", features = "spirv:v1.3,cap:Shader", wgp = <
+          compute = fp32|int32, storage = b32, subgroup = none, dot = none, mma = [], subgroup_size_choices = [32],
+          max_workgroup_sizes = [128, 128, 64], max_thread_count_per_workgroup = 128, max_workgroup_memory_bytes = 16384>>
       }>
     ]>
   ]

compiler/plugins/target/ROCM/test/target_device_features.mlir

@@ -15,7 +15,7 @@
 // GFX940-SAME:         mma = [<MFMA_F16_16x16x16_F32>, <MFMA_F16_32x32x8_F32>]
 
 // GFX1100: target = #iree_gpu.target<arch = "gfx1100",
-// GFX1100-SAME:        mma = [<WMMA_F16_16x16x16_F32>]
+// GFX1100-SAME:        mma = [<WMMA_F16_16x16x16_F32>, <WMMA_F16_16x16x16_F16>]
 // GFX1100-SAME:        subgroup_size_choices = [32, 64]
 
 // GFX941: target = #iree_gpu.target<arch = "gfx941",

compiler/plugins/target/VulkanSPIRV/BUILD.bazel

@@ -25,11 +25,10 @@ iree_compiler_cc_library(
     deps = [
         "//compiler/src/iree/compiler/Codegen/Common",
         "//compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR:IREECodegenDialect",
+        "//compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils:KnownTargets",
         "//compiler/src/iree/compiler/Codegen/SPIRV",
         "//compiler/src/iree/compiler/Codegen/Utils",
         "//compiler/src/iree/compiler/Dialect/HAL/Target",
-        "//compiler/src/iree/compiler/Dialect/Vulkan/IR",
-        "//compiler/src/iree/compiler/Dialect/Vulkan/Utils",
         "//compiler/src/iree/compiler/PluginAPI",
         "//compiler/src/iree/compiler/Utils",
         "//runtime/src/iree/schemas:spirv_executable_def_c_fbs",

compiler/plugins/target/VulkanSPIRV/CMakeLists.txt

@@ -33,11 +33,10 @@ iree_cc_library(
     MLIRSupport
     iree::compiler::Codegen::Common
     iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect
+    iree::compiler::Codegen::Dialect::GPU::TargetUtils::KnownTargets
     iree::compiler::Codegen::SPIRV
     iree::compiler::Codegen::Utils
     iree::compiler::Dialect::HAL::Target
-    iree::compiler::Dialect::Vulkan::IR
-    iree::compiler::Dialect::Vulkan::Utils
     iree::compiler::PluginAPI
     iree::compiler::Utils
     iree::schemas::spirv_executable_def_c_fbs

compiler/plugins/target/VulkanSPIRV/VulkanSPIRVTarget.cpp

@@ -5,11 +5,9 @@
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h"
+#include "iree/compiler/Codegen/Dialect/GPU/TargetUtils/KnownTargets.h"
 #include "iree/compiler/Codegen/SPIRV/Passes.h"
 #include "iree/compiler/Dialect/HAL/Target/TargetRegistry.h"
-#include "iree/compiler/Dialect/Vulkan/IR/VulkanAttributes.h"
-#include "iree/compiler/Dialect/Vulkan/IR/VulkanDialect.h"
-#include "iree/compiler/Dialect/Vulkan/Utils/TargetEnvironment.h"
 #include "iree/compiler/PluginAPI/Client.h"
 #include "iree/compiler/Utils/FlatbufferUtils.h"
 #include "iree/compiler/Utils/ModuleUtils.h"
@@ -34,20 +32,19 @@ namespace mlir::iree_compiler::IREE::HAL {
 
 namespace {
 struct VulkanSPIRVTargetOptions {
-  std::string targetTriple = "";
-  std::string targetEnv = "";
+  // Use vp_android_baseline_2022 profile as the default target--it's a good
+  // lowest common denominator to guarantee the generated SPIR-V is widely
+  // accepted for now. Eventually we want to use a list for multi-targeting.
+  std::string targetTriple = "vp_android_baseline_2022";
   bool indirectBindings = false;
 
   void bindOptions(OptionsBinder &binder) {
     static llvm::cl::OptionCategory category("VulkanSPIRV HAL Target");
     binder.opt<std::string>(
+        // TODO: Rename this as target given it's not a triple anymore.
         "iree-vulkan-target-triple", targetTriple,
         llvm::cl::desc(
             "Vulkan target triple controlling the SPIR-V environment."));
-    binder.opt<std::string>(
-        "iree-vulkan-target-env", targetEnv,
-        llvm::cl::desc(
-            "Vulkan target environment as #vk.target_env attribute assembly."));
     binder.opt<bool>(
         "iree-vulkan-experimental-indirect-bindings", indirectBindings,
         llvm::cl::desc(
@@ -56,31 +53,6 @@ struct VulkanSPIRVTargetOptions {
 };
 } // namespace
 
-// Returns the Vulkan target environment for conversion.
-static spirv::TargetEnvAttr
-getSPIRVTargetEnv(const std::string &vulkanTargetTripleOrEnv,
-                  MLIRContext *context) {
-  if (!vulkanTargetTripleOrEnv.empty()) {
-    if (vulkanTargetTripleOrEnv[0] != '#') {
-      // Parse target triple.
-      return convertTargetEnv(
-          Vulkan::getTargetEnvForTriple(context, vulkanTargetTripleOrEnv));
-    }
-
-    // Parse `#vk.target_env<...` attribute assembly.
-    if (auto attr = parseAttribute(vulkanTargetTripleOrEnv, context)) {
-      if (auto vkTargetEnv = llvm::dyn_cast<Vulkan::TargetEnvAttr>(attr)) {
-        return convertTargetEnv(vkTargetEnv);
-      }
-    }
-    emitError(Builder(context).getUnknownLoc())
-        << "cannot parse vulkan target environment as #vk.target_env "
-           "attribute: '"
-        << vulkanTargetTripleOrEnv << "'";
-  }
-  return {};
-}
-
 // TODO: VulkanOptions for choosing the Vulkan version and extensions/features.
 class VulkanTargetDevice : public TargetDevice {
 public:
@@ -119,35 +91,32 @@ class VulkanSPIRVTargetBackend : public TargetBackend {
       MLIRContext *context, StringRef deviceID, DictionaryAttr deviceConfigAttr,
       SmallVectorImpl<IREE::HAL::ExecutableTargetAttr> &executableTargetAttrs)
       const override {
-    std::string targetTripleOrEnv;
-    if (!options_.targetEnv.empty()) {
-      // TODO(scotttodd): assert if triple is set too? (mutually exclusive)
-      targetTripleOrEnv = options_.targetEnv;
-    } else if (!options_.targetTriple.empty()) {
-      targetTripleOrEnv = options_.targetTriple;
-    } else {
-      targetTripleOrEnv = "unknown-unknown-unknown";
-    }
-
-    executableTargetAttrs.push_back(getExecutableTarget(
-        context, getSPIRVTargetEnv(targetTripleOrEnv, context),
-        options_.indirectBindings));
+    executableTargetAttrs.push_back(
+        getExecutableTarget(context, options_.indirectBindings));
   }
 
   IREE::HAL::ExecutableTargetAttr
-  getExecutableTarget(MLIRContext *context, spirv::TargetEnvAttr targetEnv,
-                      bool indirectBindings) const {
+  getExecutableTarget(MLIRContext *context, bool indirectBindings) const {
     Builder b(context);
     SmallVector<NamedAttribute> configItems;
     auto addConfig = [&](StringRef name, Attribute value) {
       configItems.emplace_back(b.getStringAttr(name), value);
     };
 
-    addConfig(spirv::getTargetEnvAttrName(), targetEnv);
     if (indirectBindings) {
       addConfig("hal.bindings.indirect", b.getUnitAttr());
     }
 
+    // We only care about the architecture right now.
+    StringRef arch = StringRef(options_.targetTriple).split("-").first;
+    if (auto target = GPU::getVulkanTargetDetails(arch, context)) {
+      addConfig("iree.gpu.target", target);
+    } else {
+      emitError(b.getUnknownLoc(), "Unknown Vulkan target '")
+          << options_.targetTriple << "'";
+      return nullptr;
+    }
+
     return IREE::HAL::ExecutableTargetAttr::get(
         context, b.getStringAttr("vulkan-spirv"),
         indirectBindings ? b.getStringAttr("vulkan-spirv-fb-ptr")
@@ -156,8 +125,8 @@ class VulkanSPIRVTargetBackend : public TargetBackend {
   }
 
   void getDependentDialects(DialectRegistry &registry) const override {
-    registry.insert<IREE::Codegen::IREECodegenDialect, Vulkan::VulkanDialect,
-                    spirv::SPIRVDialect, gpu::GPUDialect>();
+    registry.insert<IREE::Codegen::IREECodegenDialect, spirv::SPIRVDialect,
+                    gpu::GPUDialect>();
   }
 
   void

compiler/plugins/target/VulkanSPIRV/test/smoketest.mlir

@@ -4,7 +4,9 @@ module attributes {
   hal.device.targets = [
     #hal.device.target<"vulkan", [
       #hal.executable.target<"vulkan-spirv", "vulkan-spirv-fb", {
-        spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, #spirv.resource_limits<>>
+        iree.gpu.target = #iree_gpu.target<arch = "", features = "spirv:v1.3,cap:Shader", wgp = <
+          compute = fp32|int32, storage = b32, subgroup = none, dot = none, mma = [], subgroup_size_choices = [32, 32],
+          max_workgroup_sizes = [128, 128, 64], max_thread_count_per_workgroup = 128, max_workgroup_memory_bytes = 16384>>
       }>
     ]>
   ]

compiler/plugins/target/WebGPUSPIRV/CMakeLists.txt

@@ -48,6 +48,7 @@ iree_cc_library(
     MLIRSPIRVTransforms
     SPIRV-Tools
     iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect
+    iree::compiler::Codegen::Dialect::GPU::TargetUtils::KnownTargets
     iree::compiler::Codegen::SPIRV
     iree::compiler::Dialect::Flow::IR
     iree::compiler::Dialect::HAL::Target

compiler/plugins/target/WebGPUSPIRV/WebGPUSPIRVTarget.cpp

@@ -6,18 +6,16 @@
 
 #include "compiler/plugins/target/WebGPUSPIRV/SPIRVToWGSL.h"
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h"
+#include "iree/compiler/Codegen/Dialect/GPU/TargetUtils/KnownTargets.h"
 #include "iree/compiler/Codegen/SPIRV/Passes.h"
 #include "iree/compiler/Codegen/WGSL/Passes.h"
 #include "iree/compiler/Dialect/Flow/IR/FlowDialect.h"
 #include "iree/compiler/Dialect/HAL/Target/TargetRegistry.h"
-#include "iree/compiler/Dialect/HAL/Transforms/Passes.h"
 #include "iree/compiler/PluginAPI/Client.h"
 #include "iree/compiler/Utils/FlatbufferUtils.h"
 #include "iree/schemas/wgsl_executable_def_builder.h"
 #include "llvm/Support/CommandLine.h"
-#include "llvm/Support/FileSystem.h"
 #include "llvm/Support/FormatVariadic.h"
-#include "llvm/Support/ToolOutputFile.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
@@ -43,18 +41,6 @@ struct WebGPUSPIRVOptions {
   }
 };
 
-// TODO(scotttodd): provide a proper target environment for WebGPU.
-static spirv::TargetEnvAttr getWebGPUTargetEnv(MLIRContext *context) {
-  // TODO(scotttodd): find list of SPIR-V extensions supported by WebGPU/WGSL
-  auto triple = spirv::VerCapExtAttr::get(
-      spirv::Version::V_1_0, {spirv::Capability::Shader},
-      {spirv::Extension::SPV_KHR_storage_buffer_storage_class}, context);
-  return spirv::TargetEnvAttr::get(
-      triple, spirv::getDefaultResourceLimits(context),
-      spirv::ClientAPI::WebGPU, spirv::Vendor::Unknown,
-      spirv::DeviceType::Unknown, spirv::TargetEnvAttr::kUnknownDeviceID);
-}
-
 // TODO: WebGPUOptions for choosing the version/extensions/etc.
 class WebGPUTargetDevice : public TargetDevice {
 public:
@@ -94,20 +80,20 @@ class WebGPUSPIRVTargetBackend : public TargetBackend {
       MLIRContext *context, StringRef deviceID, DictionaryAttr deviceConfigAttr,
       SmallVectorImpl<IREE::HAL::ExecutableTargetAttr> &executableTargetAttrs)
       const override {
-    executableTargetAttrs.push_back(
-        getExecutableTarget(context, getWebGPUTargetEnv(context)));
+    executableTargetAttrs.push_back(getExecutableTarget(context));
   }
 
   IREE::HAL::ExecutableTargetAttr
-  getExecutableTarget(MLIRContext *context,
-                      spirv::TargetEnvAttr targetEnv) const {
+  getExecutableTarget(MLIRContext *context) const {
     Builder b(context);
     SmallVector<NamedAttribute> configItems;
     auto addConfig = [&](StringRef name, Attribute value) {
       configItems.emplace_back(b.getStringAttr(name), value);
     };
 
-    addConfig(spirv::getTargetEnvAttrName(), targetEnv);
+    if (auto target = GPU::getWebGPUTargetDetails(context)) {
+      addConfig("iree.gpu.target", target);
+    }
 
     return b.getAttr<IREE::HAL::ExecutableTargetAttr>(
         b.getStringAttr("webgpu-spirv"), b.getStringAttr("webgpu-wgsl-fb"),

compiler/plugins/target/WebGPUSPIRV/test/smoketest.mlir

@@ -5,7 +5,9 @@ module attributes {
   hal.device.targets = [
     #hal.device.target<"webgpu", [
       #hal.executable.target<"webgpu-spirv", "webgpu-wgsl-fb", {
-        spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, #spirv.resource_limits<>>
+        iree.gpu.target = #iree_gpu.target<arch = "", features = "spirv:v1.0,cap:Shader,ext:SPV_KHR_storage_buffer_storage_class", wgp = <
+          compute = fp32|int32, storage = b32, subgroup = none, dot = none, mma = [], subgroup_size_choices = [32],
+          max_workgroup_sizes = [128, 128, 64], max_thread_count_per_workgroup = 128, max_workgroup_memory_bytes = 16384>>
       }>
     ]>
   ]