[mlir][spirv] Add target control to UnifyAliasedResourcePass

[antiagainst]

Summary

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.

[kuhar]

LGTM

mlir/include/mlir/Dialect/SPIRV/Transforms/Passes.h

 
 namespace mlir {
 
 class ModuleOp;
 
 namespace spirv {
 
 class ModuleOp;
+class TargetEnvAttr;
 
 //===----------------------------------------------------------------------===//
 // Passes
 //===----------------------------------------------------------------------===//
 
 #define GEN_PASS_DECL
 #include "mlir/Dialect/SPIRV/Transforms/Passes.h.inc"
 

 std::unique_ptr<OperationPass<spirv::ModuleOp>> createLowerABIAttributesPass();
 
 /// Creates an operation pass that rewrites sequential chains of
 /// spirv.CompositeInsert into spirv.CompositeConstruct.
 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
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/SPIRV/Transforms/Passes.h.inc"
 
 } // namespace spirv
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPIRV_TRANSFORMS_PASSES_H_

mlir/lib/Dialect/SPIRV/Transforms/UnifyAliasedResourcePass.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SPIRV/Transforms/Passes.h"
 
 #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"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/Debug.h"

 
 bool ResourceAliasAnalysis::shouldUnify(Operation *op) const {
   if (auto varOp = dyn_cast<spirv::GlobalVariableOp>(op)) {
     auto canonicalOp = getCanonicalResource(varOp);
     return canonicalOp && varOp != canonicalOp;
   }
   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);
   }
 
   if (auto acOp = dyn_cast<spirv::AccessChainOp>(op))
     return shouldUnify(acOp.getBasePtr().getDefiningOp());
   if (auto loadOp = dyn_cast<spirv::LoadOp>(op))
     return shouldUnify(loadOp.getPtr().getDefiningOp());
   if (auto storeOp = dyn_cast<spirv::StoreOp>(op))
if (!srcElemType.isIntOrFloat() || !dstElemType.isIntOrFloat())
       return rewriter.notifyMatchFailure(storeOp, "not scalar type");
     if (!areSameBitwidthScalarType(srcElemType, dstElemType))
       return rewriter.notifyMatchFailure(storeOp, "different bitwidth");
 
     Location loc = storeOp.getLoc();
     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();
   }
 };
 
 //===----------------------------------------------------------------------===//
 // Pass
 //===----------------------------------------------------------------------===//
 
 namespace {
 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
 
 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>();
 
   ConversionTarget target(*context);
   target.addDynamicallyLegalOp<spirv::GlobalVariableOp, spirv::AddressOfOp,
                                spirv::AccessChainOp, spirv::LoadOp,
                                spirv::StoreOp>(
       [&analysis](Operation *op) { return !analysis.shouldUnify(op); });
void UnifyAliasedResourcePass::runOnOperation() {
   for (const auto &dr : resourceMap) {
     const auto &resources = dr.second;
     if (resources.size() == 1)
       resources.front()->removeAttr("aliased");
   }
 }
 
 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));
 }