[mlir][spirv] Make MemRef memory space mapping pass more flexible

[antiagainst]

Summary

• Avoid restricting the pass to to builtin module ops. The pass should be able to run on any region ops.
• Avoid hardcoding func FuncOp when handling functions. Instead, use the function op interface.
• Assigns the default mapping in the constructor. So for cases where we are using the pass in a pipeline, we still have a meaningful default.

Along the way, dropped uncessary unrealized conversion casts and
use full conversion. The pass should be able to convert all sorts
of ops; there is really no need to have such bridages.

[kuhar]

mlir/lib/Conversion/MemRefToSPIRV/MapMemRefStorageClassPass.cpp

98

nit: grammar seems a bit off to me here. Shouldn't it be due to being unable to or due to the inability to or due to not being able to?

mlir/include/mlir/Conversion/MemRefToSPIRV/MemRefToSPIRVPass.h

 #include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 class ModuleOp;
 
 /// Creates a pass to map numeric MemRef memory spaces to symbolic SPIR-V
 /// storage classes. The mapping is read from the command-line option.
-std::unique_ptr<OperationPass<ModuleOp>> createMapMemRefStorageClassPass();
+std::unique_ptr<OperationPass<>> createMapMemRefStorageClassPass();
 
 /// Creates a pass to convert MemRef ops to SPIR-V ops.
 std::unique_ptr<OperationPass<ModuleOp>> createConvertMemRefToSPIRVPass();
 
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_MEMREFTOSPIRV_MEMREFTOSPIRVPASS_H

mlir/include/mlir/Conversion/Passes.td

Option<"useGenericFunctions", "use-generic-functions",
            "classic 'malloc', 'aligned_alloc' and 'free' functions">
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // MemRefToSPIRV
 //===----------------------------------------------------------------------===//
 
-def MapMemRefStorageClass : Pass<"map-memref-spirv-storage-class", "ModuleOp"> {
+def MapMemRefStorageClass : Pass<"map-memref-spirv-storage-class"> {
   let summary = "Map numeric MemRef memory spaces to SPIR-V storage classes";
   let constructor = "mlir::createMapMemRefStorageClassPass()";
   let dependentDialects = ["spirv::SPIRVDialect"];
   let options = [
     Option<"clientAPI", "client-api", "std::string", /*default=*/"\"vulkan\"",
            "The client API to use for populating mappings">
   ];
 }

mlir/lib/Conversion/MemRefToSPIRV/MapMemRefStorageClassPass.cpp

 //===- MapMemRefStorageCLassPass.cpp --------------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements a pass to map numeric MemRef memory spaces to
 // symbolic ones defined in the SPIR-V specification.
 //
 //===----------------------------------------------------------------------===//
 
 #include "../PassDetail.h"
 #include "mlir/Conversion/MemRefToSPIRV/MemRefToSPIRV.h"
 #include "mlir/Conversion/MemRefToSPIRV/MemRefToSPIRVPass.h"
-#include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVAttributes.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/FunctionInterfaces.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "mlir-map-memref-storage-class"
 
 using namespace mlir;
 
spirv::MemorySpaceToStorageClassConverter::MemorySpaceToStorageClassConverter(
   addConversion([](Type type) { return type; });
 
   addConversion([this](BaseMemRefType memRefType) -> Optional<Type> {
     // Expect IntegerAttr memory spaces. The attribute can be missing for the
     // case of memory space == 0.
     Attribute spaceAttr = memRefType.getMemorySpace();
     if (spaceAttr && !spaceAttr.isa<IntegerAttr>()) {
       LLVM_DEBUG(llvm::dbgs() << "cannot convert " << memRefType
-                              << " due to non-IntegerAttr memory space");
+                              << " due to non-IntegerAttr memory space\n");
       return llvm::None;
     }
 
     unsigned space = memRefType.getMemorySpaceAsInt();
     auto it = this->memorySpaceMap.find(space);
     if (it == this->memorySpaceMap.end()) {
-      LLVM_DEBUG(llvm::dbgs() << "cannot convert " << memRefType
-                              << " due to unable to find memory space in map");
+      LLVM_DEBUG(llvm::dbgs()
+                 << "cannot convert " << memRefType
+                 << " due to being unable to find memory space in map\n");
       return llvm::None;
     }
 
     auto storageAttr =
         spirv::StorageClassAttr::get(memRefType.getContext(), it->second);
     if (auto rankedType = memRefType.dyn_cast<MemRefType>()) {
       return MemRefType::get(memRefType.getShape(), memRefType.getElementType(),
                              rankedType.getLayout(), storageAttr);

 static bool isLegalAttr(Attribute attr) {
   if (auto typeAttr = attr.dyn_cast<TypeAttr>())
     return isLegalType(typeAttr.getValue());
   return true;
 }
 
 /// Returns true if the given `op` is considered as legal for SPIR-V conversion.
 static bool isLegalOp(Operation *op) {
-  if (auto funcOp = dyn_cast<func::FuncOp>(op)) {
-    FunctionType funcType = funcOp.getFunctionType();
-    return llvm::all_of(funcType.getInputs(), isLegalType) &&
-           llvm::all_of(funcType.getResults(), isLegalType);
+  if (auto funcOp = dyn_cast<FunctionOpInterface>(op)) {
+    return llvm::all_of(funcOp.getArgumentTypes(), isLegalType) &&
+           llvm::all_of(funcOp.getResultTypes(), isLegalType);
   }
 
   auto attrs = llvm::map_range(op->getAttrs(), [](const NamedAttribute &attr) {
     return attr.getValue();
   });
 
   return llvm::all_of(op->getOperandTypes(), isLegalType) &&
          llvm::all_of(op->getResultTypes(), isLegalType) &&

 //===----------------------------------------------------------------------===//
 // Conversion Pass
 //===----------------------------------------------------------------------===//
 
 namespace {
 class MapMemRefStorageClassPass final
     : public MapMemRefStorageClassBase<MapMemRefStorageClassPass> {
 public:
-  explicit MapMemRefStorageClassPass() = default;
+  explicit MapMemRefStorageClassPass() {
+    memorySpaceMap = spirv::getDefaultVulkanStorageClassMap();
+
+    LLVM_DEBUG({
+      llvm::dbgs() << "memory space to storage class mapping:\n";
+      if (memorySpaceMap.empty())
+        llvm::dbgs() << "  [empty]\n";
+      for (auto kv : memorySpaceMap)
+        llvm::dbgs() << "  " << kv.first << " -> "
+                     << spirv::stringifyStorageClass(kv.second) << "\n";
+    });
+  }
   explicit MapMemRefStorageClassPass(
       const spirv::MemorySpaceToStorageClassMap &memorySpaceMap)
       : memorySpaceMap(memorySpaceMap) {}
 
   LogicalResult initializeOptions(StringRef options) override;
 
   void runOnOperation() override;
 
 private:
   spirv::MemorySpaceToStorageClassMap memorySpaceMap;
 };
 } // namespace
 
 LogicalResult MapMemRefStorageClassPass::initializeOptions(StringRef options) {
   if (failed(Pass::initializeOptions(options)))
     return failure();
 
   if (clientAPI != "vulkan")
     return failure();
 
-  memorySpaceMap = spirv::getDefaultVulkanStorageClassMap();
-
-  LLVM_DEBUG({
-    llvm::dbgs() << "memory space to storage class mapping:\n";
-    if (memorySpaceMap.empty())
-      llvm::dbgs() << "  [empty]\n";
-    for (auto kv : memorySpaceMap)
-      llvm::dbgs() << "  " << kv.first << " -> "
-                   << spirv::stringifyStorageClass(kv.second) << "\n";
-  });
-
   return success();
 }
 
 void MapMemRefStorageClassPass::runOnOperation() {
   MLIRContext *context = &getContext();
-  ModuleOp module = getOperation();
+  Operation *op = getOperation();
 
   auto target = spirv::getMemorySpaceToStorageClassTarget(*context);
-
   spirv::MemorySpaceToStorageClassConverter converter(memorySpaceMap);
-  // Use UnrealizedConversionCast as the bridge so that we don't need to pull in
-  // patterns for other dialects.
-  auto addUnrealizedCast = [](OpBuilder &builder, Type type, ValueRange inputs,
-                              Location loc) {
-    auto cast = builder.create<UnrealizedConversionCastOp>(loc, type, inputs);
-    return Optional<Value>(cast.getResult(0));
-  };
-  converter.addSourceMaterialization(addUnrealizedCast);
-  converter.addTargetMaterialization(addUnrealizedCast);
-  target->addLegalOp<UnrealizedConversionCastOp>();
 
   RewritePatternSet patterns(context);
   spirv::populateMemorySpaceToStorageClassPatterns(converter, patterns);
 
-  if (failed(applyPartialConversion(module, *target, std::move(patterns))))
+  if (failed(applyFullConversion(op, *target, std::move(patterns))))
     return signalPassFailure();
 }
 
-std::unique_ptr<OperationPass<ModuleOp>>
-mlir::createMapMemRefStorageClassPass() {
+std::unique_ptr<OperationPass<>> mlir::createMapMemRefStorageClassPass() {
   return std::make_unique<MapMemRefStorageClassPass>();
 }

mlir/test/Conversion/MemRefToSPIRV/map-storage-class.mlir

 func.func @type_attribute() {
   // VULKAN: attr = memref<i32, #spv.storage_class<Generic>>
   "dialect.memref_producer"() { attr = memref<i32, 1> } : () -> ()
   return
 }
 
 // -----
 
-// VULKAN-LABEL: func @function_io
+// VULKAN-LABEL: func.func @function_io
 func.func @function_io
   // VULKAN-SAME: (%{{.+}}: memref<f64, #spv.storage_class<Generic>>, %{{.+}}: memref<4xi32, #spv.storage_class<Workgroup>>)
   (%arg0: memref<f64, 1>, %arg1: memref<4xi32, 3>)
   // VULKAN-SAME: -> (memref<f64, #spv.storage_class<Generic>>, memref<4xi32, #spv.storage_class<Workgroup>>)
   -> (memref<f64, 1>, memref<4xi32, 3>) {
   return %arg0, %arg1: memref<f64, 1>, memref<4xi32, 3>
 }
 
 // -----
 
-// VULKAN: func @region
+gpu.module @kernel {
+// VULKAN-LABEL: gpu.func @function_io
+// VULKAN-SAME: memref<8xi32, #spv.storage_class<StorageBuffer>>
+gpu.func @function_io(%arg0 : memref<8xi32>) kernel { gpu.return }
+}
+
+// -----
+
+// VULKAN-LABEL: func.func @region
 func.func @region(%cond: i1, %arg0: memref<f32, 1>) {
   scf.if %cond {
     //      VULKAN: "dialect.memref_consumer"(%{{.+}}) {attr = memref<i64, #spv.storage_class<Workgroup>>}
     // VULKAN-SAME: (memref<f32, #spv.storage_class<Generic>>) -> memref<f32, #spv.storage_class<Generic>>
     %0 = "dialect.memref_consumer"(%arg0) { attr = memref<i64, 3> } : (memref<f32, 1>) -> (memref<f32, 1>)
   }
   return
 }