[mlir][spirv] Add a pass to map memref memory space

MemRef types now can carry an attribute to represent the memory
space. Still, upper layers in the compilation stack mostly use
nuemric values. They don't mean much (other than differentiating
separate memory domains) in MLIR's multi-level settings. Those
numeric memory space inside MemRef types need to be translated
into concrete SPIR-V storage classes during lowering to pin down
to concrete memory types.

Thus far we have been hardcoding an arbitrary mapping from memory
space to storage class for converting MemRef types. This works fine
for only targeting Vulkan; it falls apart if we want to target other
SPIR-V consumers like OpenCL, as different consumers might want
different storage classes for the buffer/variable of the same
lifetime. For example, StorageClass in Vulkan vs. CrossWorkgroup
in OpenCL.

So putting up a new pass to let the user to control how to map
MemRef memory spaces into SPIR-V storage classes. This provides
more flexibility and can address the awkwardness in the current
SPIR-V type converter. This pass should be the prelimiary step
towards lowering MemRef related types/ops into SPIR-V.

Reviewed By: mravishankar

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

7 files changed, 413 insertions, 2 deletions

diff --git a/mlir/include/mlir/Conversion/MemRefToSPIRV/MemRefToSPIRV.h b/mlir/include/mlir/Conversion/MemRefToSPIRV/MemRefToSPIRV.h
index 730ead79c488..e7b4b7ab23de 100644
--- a/mlir/include/mlir/Conversion/MemRefToSPIRV/MemRefToSPIRV.h
+++ b/mlir/include/mlir/Conversion/MemRefToSPIRV/MemRefToSPIRV.h
@@ -13,11 +13,40 @@
 #ifndef MLIR_CONVERSION_MEMREFTOSPIRV_MEMREFTOSPIRV_H
 #define MLIR_CONVERSION_MEMREFTOSPIRV_MEMREFTOSPIRV_H
 
+#include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"
 #include "mlir/Transforms/DialectConversion.h"
+#include <memory>
 
 namespace mlir {
 class SPIRVTypeConverter;
 
+namespace spirv {
+/// Mapping from numeric MemRef memory spaces into SPIR-V symbolic ones.
+using MemorySpaceToStorageClassMap = DenseMap<unsigned, spirv::StorageClass>;
+
+/// Type converter for converting numeric MemRef memory spaces into SPIR-V
+/// symbolic ones.
+class MemorySpaceToStorageClassConverter : public TypeConverter {
+public:
+  explicit MemorySpaceToStorageClassConverter(
+      const MemorySpaceToStorageClassMap &memorySpaceMap);
+
+private:
+  const MemorySpaceToStorageClassMap &memorySpaceMap;
+};
+
+/// Creates the target that populates legality of ops with MemRef types.
+std::unique_ptr<ConversionTarget>
+getMemorySpaceToStorageClassTarget(MLIRContext &);
+
+/// Appends to a pattern list additional patterns for converting numeric MemRef
+/// memory spaces into SPIR-V symbolic ones.
+void populateMemorySpaceToStorageClassPatterns(
+    MemorySpaceToStorageClassConverter &typeConverter,
+    RewritePatternSet &patterns);
+
+} // namespace spirv
+
 /// Appends to a pattern list additional patterns for translating MemRef ops
 /// to SPIR-V ops.
 void populateMemRefToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
diff --git a/mlir/include/mlir/Conversion/MemRefToSPIRV/MemRefToSPIRVPass.h b/mlir/include/mlir/Conversion/MemRefToSPIRV/MemRefToSPIRVPass.h
index a9990d99b160..8b239699d8a5 100644
--- a/mlir/include/mlir/Conversion/MemRefToSPIRV/MemRefToSPIRVPass.h
+++ b/mlir/include/mlir/Conversion/MemRefToSPIRV/MemRefToSPIRVPass.h
@@ -13,11 +13,16 @@
 #ifndef MLIR_CONVERSION_MEMREFTOSPIRV_MEMREFTOSPIRVPASS_H
 #define 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();
+
 /// Creates a pass to convert MemRef ops to SPIR-V ops.
 std::unique_ptr<OperationPass<ModuleOp>> createConvertMemRefToSPIRVPass();
 
diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index 00ca7af1fabd..ed5348109484 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -538,6 +538,17 @@ def ConvertMemRefToLLVM : Pass<"convert-memref-to-llvm", "ModuleOp"> {
 // MemRefToSPIRV
 //===----------------------------------------------------------------------===//
 
+def MapMemRefStorageClass : Pass<"map-memref-spirv-storage-class", "ModuleOp"> {
+  let summary = "Map numeric MemRef memory spaces to SPIR-V storage classes";
+  let constructor = "mlir::createMapMemRefStorageClassPass()";
+  let dependentDialects = ["spirv::SPIRVDialect"];
+  let options = [
+    Option<"mappings", "mappings", "std::string", /*default=*/"",
+           "A comma-separated list of memory space to storage class mappings; "
+           "for example, '0=StorageClass,1=Uniform,2=Workgroup'">
+  ];
+}
+
 def ConvertMemRefToSPIRV : Pass<"convert-memref-to-spirv", "ModuleOp"> {
   let summary = "Convert MemRef dialect to SPIR-V dialect";
   let constructor = "mlir::createConvertMemRefToSPIRVPass()";
diff --git a/mlir/lib/Conversion/MemRefToSPIRV/CMakeLists.txt b/mlir/lib/Conversion/MemRefToSPIRV/CMakeLists.txt
index cc59f63f9c86..740bdceccce3 100644
--- a/mlir/lib/Conversion/MemRefToSPIRV/CMakeLists.txt
+++ b/mlir/lib/Conversion/MemRefToSPIRV/CMakeLists.txt
@@ -1,4 +1,5 @@
 add_mlir_conversion_library(MLIRMemRefToSPIRV
+  MapMemRefStorageClassPass.cpp
   MemRefToSPIRV.cpp
   MemRefToSPIRVPass.cpp
 
diff --git a/mlir/lib/Conversion/MemRefToSPIRV/MapMemRefStorageClassPass.cpp b/mlir/lib/Conversion/MemRefToSPIRV/MapMemRefStorageClassPass.cpp
new file mode 100644
index 000000000000..ba0003ce5083
--- /dev/null
+++ b/mlir/lib/Conversion/MemRefToSPIRV/MapMemRefStorageClassPass.cpp
@@ -0,0 +1,283 @@
+//===- 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/SPIRVDialect.h"
+#include "mlir/IR/BuiltinTypes.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;
+
+//===----------------------------------------------------------------------===//
+// Utility Functions
+//===----------------------------------------------------------------------===//
+
+/// Parses the  memory space mapping string `memorySpaceMapStr `and writes the
+/// mappings encoded inside to `memorySpaceMap`.
+static bool
+parseMappingStr(StringRef memorySpaceMapStr,
+                spirv::MemorySpaceToStorageClassMap &memorySpaceMap) {
+  memorySpaceMap.clear();
+  SmallVector<StringRef> mappings;
+  llvm::SplitString(memorySpaceMapStr, mappings, ",");
+  for (StringRef mapping : mappings) {
+    StringRef key, value;
+    std::tie(key, value) = mapping.split('=');
+    unsigned space;
+    if (!llvm::to_integer(key, space)) {
+      LLVM_DEBUG(llvm::dbgs()
+                 << "failed to parse mapping string key: " << key << "\n");
+      memorySpaceMap.clear();
+      return false;
+    }
+    Optional<spirv::StorageClass> storage = spirv::symbolizeStorageClass(value);
+    if (!storage) {
+      LLVM_DEBUG(llvm::dbgs()
+                 << "failed to parse mapping string value: " << value << "\n");
+      memorySpaceMap.clear();
+      return false;
+    }
+    memorySpaceMap[space] = *storage;
+  }
+  return true;
+}
+
+//===----------------------------------------------------------------------===//
+// Type Converter
+//===----------------------------------------------------------------------===//
+
+spirv::MemorySpaceToStorageClassConverter::MemorySpaceToStorageClassConverter(
+    const spirv::MemorySpaceToStorageClassMap &memorySpaceMap)
+    : memorySpaceMap(memorySpaceMap) {
+  // Pass through for all other types.
+  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");
+      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");
+      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);
+    }
+    return UnrankedMemRefType::get(memRefType.getElementType(), storageAttr);
+  });
+
+  addConversion([this](FunctionType type) {
+    SmallVector<Type> inputs, results;
+    inputs.reserve(type.getNumInputs());
+    results.reserve(type.getNumResults());
+    for (Type input : type.getInputs())
+      inputs.push_back(convertType(input));
+    for (Type result : type.getResults())
+      results.push_back(convertType(result));
+    return FunctionType::get(type.getContext(), inputs, results);
+  });
+}
+
+//===----------------------------------------------------------------------===//
+// Conversion Target
+//===----------------------------------------------------------------------===//
+
+/// Returns true if the given `type` is considered as legal for SPIR-V
+/// conversion.
+static bool isLegalType(Type type) {
+  if (auto memRefType = type.dyn_cast<BaseMemRefType>()) {
+    Attribute spaceAttr = memRefType.getMemorySpace();
+    return spaceAttr && spaceAttr.isa<spirv::StorageClassAttr>();
+  }
+  return true;
+}
+
+/// Returns true if the given `attr` is considered as legal for SPIR-V
+/// conversion.
+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);
+  }
+
+  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) &&
+         llvm::all_of(attrs, isLegalAttr);
+}
+
+std::unique_ptr<ConversionTarget>
+spirv::getMemorySpaceToStorageClassTarget(MLIRContext &context) {
+  auto target = std::make_unique<ConversionTarget>(context);
+  target->markUnknownOpDynamicallyLegal(isLegalOp);
+  return target;
+}
+
+//===----------------------------------------------------------------------===//
+// Conversion Pattern
+//===----------------------------------------------------------------------===//
+
+namespace {
+/// Converts any op that has operands/results/attributes with numeric MemRef
+/// memory spaces.
+struct MapMemRefStoragePattern final : public ConversionPattern {
+  MapMemRefStoragePattern(MLIRContext *context, TypeConverter &converter)
+      : ConversionPattern(converter, MatchAnyOpTypeTag(), 1, context) {}
+
+  LogicalResult
+  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+} // namespace
+
+LogicalResult MapMemRefStoragePattern::matchAndRewrite(
+    Operation *op, ArrayRef<Value> operands,
+    ConversionPatternRewriter &rewriter) const {
+  llvm::SmallVector<NamedAttribute, 4> newAttrs;
+  newAttrs.reserve(op->getAttrs().size());
+  for (auto attr : op->getAttrs()) {
+    if (auto typeAttr = attr.getValue().dyn_cast<TypeAttr>()) {
+      auto newAttr = getTypeConverter()->convertType(typeAttr.getValue());
+      newAttrs.emplace_back(attr.getName(), TypeAttr::get(newAttr));
+    } else {
+      newAttrs.push_back(attr);
+    }
+  }
+
+  llvm::SmallVector<Type, 4> newResults;
+  (void)getTypeConverter()->convertTypes(op->getResultTypes(), newResults);
+
+  OperationState state(op->getLoc(), op->getName().getStringRef(), operands,
+                       newResults, newAttrs, op->getSuccessors());
+
+  for (Region &region : op->getRegions()) {
+    Region *newRegion = state.addRegion();
+    rewriter.inlineRegionBefore(region, *newRegion, newRegion->begin());
+    TypeConverter::SignatureConversion result(newRegion->getNumArguments());
+    (void)getTypeConverter()->convertSignatureArgs(
+        newRegion->getArgumentTypes(), result);
+    rewriter.applySignatureConversion(newRegion, result);
+  }
+
+  Operation *newOp = rewriter.create(state);
+  rewriter.replaceOp(op, newOp->getResults());
+  return success();
+}
+
+void spirv::populateMemorySpaceToStorageClassPatterns(
+    spirv::MemorySpaceToStorageClassConverter &typeConverter,
+    RewritePatternSet &patterns) {
+  patterns.add<MapMemRefStoragePattern>(patterns.getContext(), typeConverter);
+}
+
+//===----------------------------------------------------------------------===//
+// Conversion Pass
+//===----------------------------------------------------------------------===//
+
+namespace {
+class MapMemRefStorageClassPass final
+    : public MapMemRefStorageClassBase<MapMemRefStorageClassPass> {
+public:
+  explicit MapMemRefStorageClassPass() = default;
+  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 (!parseMappingStr(mappings, memorySpaceMap))
+    return failure();
+
+  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();
+
+  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))))
+    return signalPassFailure();
+}
+
+std::unique_ptr<OperationPass<ModuleOp>>
+mlir::createMapMemRefStorageClassPass() {
+  return std::make_unique<MapMemRefStorageClassPass>();
+}
diff --git a/mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp b/mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp
index cfb72fd8f76d..55da39cb98b1 100644
--- a/mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp
+++ b/mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp
@@ -126,8 +126,8 @@ static Optional<spirv::Scope> getAtomicOpScope(MemRefType type) {
     return spirv::Scope::Device;
   case spirv::StorageClass::Workgroup:
     return spirv::Scope::Workgroup;
-  default: {
-  }
+  default:
+    break;
   }
   return {};
 }
diff --git a/mlir/test/Conversion/MemRefToSPIRV/map-storage-class.mlir b/mlir/test/Conversion/MemRefToSPIRV/map-storage-class.mlir
new file mode 100644
index 000000000000..5a2fe262110d
--- /dev/null
+++ b/mlir/test/Conversion/MemRefToSPIRV/map-storage-class.mlir
@@ -0,0 +1,82 @@
+// RUN: mlir-opt -split-input-file -allow-unregistered-dialect -map-memref-spirv-storage-class='mappings=0=StorageBuffer,1=Uniform,2=Workgroup,3=PushConstant' -verify-diagnostics %s -o - | FileCheck %s
+
+// Mappings:
+//   0 -> StorageBuffer (12)
+//   2 -> Workgroup (4)
+//   1 -> Uniform (2)
+//   3 -> PushConstant (9)
+// TODO: create a StorageClass wrapper class so we can print the symbolc
+// storage class (instead of the backing IntegerAttr) and be able to
+// round trip the IR.
+
+// CHECK-LABEL: func @operand_result
+func.func @operand_result() {
+  // CHECK: memref<f32, 12 : i32>
+  %0 = "dialect.memref_producer"() : () -> (memref<f32>)
+  // CHECK: memref<4xi32, 2 : i32>
+  %1 = "dialect.memref_producer"() : () -> (memref<4xi32, 1>)
+  // CHECK: memref<?x4xf16, 4 : i32>
+  %2 = "dialect.memref_producer"() : () -> (memref<?x4xf16, 2>)
+  // CHECK: memref<*xf16, 9 : i32>
+  %3 = "dialect.memref_producer"() : () -> (memref<*xf16, 3>)
+
+
+  "dialect.memref_consumer"(%0) : (memref<f32>) -> ()
+  // CHECK: memref<4xi32, 2 : i32>
+  "dialect.memref_consumer"(%1) : (memref<4xi32, 1>) -> ()
+  // CHECK: memref<?x4xf16, 4 : i32>
+  "dialect.memref_consumer"(%2) : (memref<?x4xf16, 2>) -> ()
+  // CHECK: memref<*xf16, 9 : i32>
+  "dialect.memref_consumer"(%3) : (memref<*xf16, 3>) -> ()
+
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @type_attribute
+func.func @type_attribute() {
+  // CHECK: attr = memref<i32, 2 : i32>
+  "dialect.memref_producer"() { attr = memref<i32, 1> } : () -> ()
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @function_io
+func.func @function_io
+  // CHECK-SAME: (%{{.+}}: memref<f64, 4 : i32>, %{{.+}}: memref<4xi32, 9 : i32>)
+  (%arg0: memref<f64, 2>, %arg1: memref<4xi32, 3>)
+  // CHECK-SAME: -> (memref<f64, 4 : i32>, memref<4xi32, 9 : i32>)
+  -> (memref<f64, 2>, memref<4xi32, 3>) {
+  return %arg0, %arg1: memref<f64, 2>, memref<4xi32, 3>
+}
+
+// -----
+
+// CHECK: func @region
+func.func @region(%cond: i1, %arg0: memref<f32, 1>) {
+  scf.if %cond {
+    //      CHECK: "dialect.memref_consumer"(%{{.+}}) {attr = memref<i64, 4 : i32>}
+    // CHECK-SAME: (memref<f32, 2 : i32>) -> memref<f32, 2 : i32>
+    %0 = "dialect.memref_consumer"(%arg0) { attr = memref<i64, 2> } : (memref<f32, 1>) -> (memref<f32, 1>)
+  }
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @non_memref_types
+func.func @non_memref_types(%arg: f32) -> f32 {
+  // CHECK: "dialect.op"(%{{.+}}) {attr = 16 : i64} : (f32) -> f32
+  %0 = "dialect.op"(%arg) { attr = 16 } : (f32) -> (f32)
+  return %0 : f32
+}
+
+// -----
+
+func.func @missing_mapping() {
+  // expected-error @+1 {{failed to legalize}}
+  %0 = "dialect.memref_producer"() : () -> (memref<f32, 128>)
+  return
+}