[mlir][spirv] Enable WebGPU to use UnifyAliasedResourcePass

When targeting WebGPU, we also need to transcompile SPIR-V, into
WGSL this time. We have similar limitations there like when
targeting Metal or MoltenVK.

Reviewed By: kuhar

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

diff --git a/mlir/lib/Dialect/SPIRV/Transforms/UnifyAliasedResourcePass.cpp b/mlir/lib/Dialect/SPIRV/Transforms/UnifyAliasedResourcePass.cpp
index 15cf50a..9c43794 100644 (file)
--- a/mlir/lib/Dialect/SPIRV/Transforms/UnifyAliasedResourcePass.cpp
+++ b/mlir/lib/Dialect/SPIRV/Transforms/UnifyAliasedResourcePass.cpp
@@ -549,10 +549,16 @@ void UnifyAliasedResourcePass::runOnOperation() {
   MLIRContext *context = &getContext();
 
   if (getTargetEnvFn) {
-    // This pass is actually only needed for targeting Apple GPUs via MoltenVK,
-    // where we need to translate SPIR-V into MSL. The translation has
-    // limitations.
-    if (getTargetEnvFn(moduleOp).getVendorID() != spirv::Vendor::Apple)
+    // This pass is only needed for targeting WebGPU, Metal, or layering Vulkan
+    // on Metal via MoltenVK, where we need to translate SPIR-V into WGSL or
+    // MSL. The translation has limitations.
+    spirv::TargetEnvAttr targetEnv = getTargetEnvFn(moduleOp);
+    spirv::ClientAPI clientAPI = targetEnv.getClientAPI();
+    bool isVulkanOnAppleDevices =
+        clientAPI == spirv::ClientAPI::Vulkan &&
+        targetEnv.getVendorID() == spirv::Vendor::Apple;
+    if (clientAPI != spirv::ClientAPI::WebGPU &&
+        clientAPI != spirv::ClientAPI::Metal && !isVulkanOnAppleDevices)
       return;
   }