From 7841f557ea3df1987ceb1c033e458790f5828795 Mon Sep 17 00:00:00 2001
From: Lei Zhang <antiagainst@google.com>
Date: Thu, 27 Oct 2022 14:31:21 -0400
Subject: [PATCH] [mlir][spirv] Add target control to UnifyAliasedResourcePass

The UnifyAliasedResourcePass is actually only necessary for
targeting Apple GPUs via MoltenVK, where we need to translate
SPIR-V into MSL. The translation has limitations--no support
of aliased resources. So introducing a control to disable
this pass if targeting other platforms.

Reviewed By: kuhar

Differential Revision: https://reviews.llvm.org/D136869
---
 .../include/mlir/Dialect/SPIRV/Transforms/Passes.h |  4 +++-
 .../SPIRV/Transforms/UnifyAliasedResourcePass.cpp  | 26 +++++++++++++++++-----
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/mlir/include/mlir/Dialect/SPIRV/Transforms/Passes.h b/mlir/include/mlir/Dialect/SPIRV/Transforms/Passes.h
index 26f5da8..f718d52 100644
--- a/mlir/include/mlir/Dialect/SPIRV/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/SPIRV/Transforms/Passes.h
@@ -22,6 +22,7 @@ class ModuleOp;
 namespace spirv {
 
 class ModuleOp;
+class TargetEnvAttr;
 
 //===----------------------------------------------------------------------===//
 // Passes
@@ -69,8 +70,9 @@ std::unique_ptr<OperationPass<spirv::ModuleOp>> createRewriteInsertsPass();
 
 /// Creates an operation pass that unifies access of multiple aliased resources
 /// into access of one single resource.
+using GetTargetEnvFn = std::function<spirv::TargetEnvAttr(spirv::ModuleOp)>;
 std::unique_ptr<OperationPass<spirv::ModuleOp>>
-createUnifyAliasedResourcePass();
+createUnifyAliasedResourcePass(GetTargetEnvFn getTargetEnv = nullptr);
 
 //===----------------------------------------------------------------------===//
 // Registration
diff --git a/mlir/lib/Dialect/SPIRV/Transforms/UnifyAliasedResourcePass.cpp b/mlir/lib/Dialect/SPIRV/Transforms/UnifyAliasedResourcePass.cpp
index 8565e97..15cf50a 100644
--- a/mlir/lib/Dialect/SPIRV/Transforms/UnifyAliasedResourcePass.cpp
+++ b/mlir/lib/Dialect/SPIRV/Transforms/UnifyAliasedResourcePass.cpp
@@ -16,6 +16,7 @@
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
+#include "mlir/Dialect/SPIRV/IR/TargetAndABI.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
@@ -223,7 +224,8 @@ bool ResourceAliasAnalysis::shouldUnify(Operation *op) const {
   }
   if (auto addressOp = dyn_cast<spirv::AddressOfOp>(op)) {
     auto moduleOp = addressOp->getParentOfType<spirv::ModuleOp>();
-    auto *varOp = SymbolTable::lookupSymbolIn(moduleOp, addressOp.getVariable());
+    auto *varOp =
+        SymbolTable::lookupSymbolIn(moduleOp, addressOp.getVariable());
     return shouldUnify(varOp);
   }
 
@@ -517,8 +519,8 @@ struct ConvertStore : public ConvertAliasResource<spirv::StoreOp> {
     Value value = adaptor.getValue();
     if (srcElemType != dstElemType)
       value = rewriter.create<spirv::BitcastOp>(loc, dstElemType, value);
-    rewriter.replaceOpWithNewOp<spirv::StoreOp>(storeOp, adaptor.getPtr(), value,
-                                                storeOp->getAttrs());
+    rewriter.replaceOpWithNewOp<spirv::StoreOp>(storeOp, adaptor.getPtr(),
+                                                value, storeOp->getAttrs());
     return success();
   }
 };
@@ -532,7 +534,13 @@ class UnifyAliasedResourcePass final
     : public spirv::impl::SPIRVUnifyAliasedResourcePassBase<
           UnifyAliasedResourcePass> {
 public:
+  explicit UnifyAliasedResourcePass(spirv::GetTargetEnvFn getTargetEnv)
+      : getTargetEnvFn(std::move(getTargetEnv)) {}
+
   void runOnOperation() override;
+
+private:
+  spirv::GetTargetEnvFn getTargetEnvFn;
 };
 } // namespace
 
@@ -540,6 +548,14 @@ void UnifyAliasedResourcePass::runOnOperation() {
   spirv::ModuleOp moduleOp = getOperation();
   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)
+      return;
+  }
+
   // Analyze aliased resources first.
   ResourceAliasAnalysis &analysis = getAnalysis<ResourceAliasAnalysis>();
 
@@ -570,6 +586,6 @@ void UnifyAliasedResourcePass::runOnOperation() {
 }
 
 std::unique_ptr<mlir::OperationPass<spirv::ModuleOp>>
-spirv::createUnifyAliasedResourcePass() {
-  return std::make_unique<UnifyAliasedResourcePass>();
+spirv::createUnifyAliasedResourcePass(spirv::GetTargetEnvFn getTargetEnv) {
+  return std::make_unique<UnifyAliasedResourcePass>(std::move(getTargetEnv));
 }
-- 
2.7.4