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

[antiagainst]

Summary

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.

[mravishankar]

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

27

I dont have a strong opinion here, but IIUC using this kind of signature for creating passes can make it hard for reproducers to work.. From past experience having a pass for this is not really useful. It is better to just have a utility function that the client can wrap within a pass. This works for reproducers.

[antiagainst]

Address comments

[antiagainst]

Remove unnecessary entry points

[antiagainst]

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

27

Good point. I've changed to expose the various bits so they can be put into a downstream pass.

[kuhar]

Just some nits

mlir/lib/Conversion/MemRefToSPIRV/MapMemRefStorageClassPass.cpp

101

nit 1: consider defining these two on two separate lines. I think this is not covered by the llvm coding standards, so it's up to you.
nit 2: could you add a comment explaining how the small size was picked? I think it has to do with the number of SPIR-V storage classes, doesn't it? Or consider using the default size if there is no clear limit.

103–105

Alternatively, I think you can do something like this:

auto inputs = to_vector(map_range(type.getInputs(), convertType));

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

2

Would it make sense to add another test with an ill-formed mapping lists?

[antiagainst]

mlir/lib/Conversion/MemRefToSPIRV/MapMemRefStorageClassPass.cpp

101

It's mostly from experience of how many inputs/results a function can have. Dropped the explicit number there.

103–105

convertType is a class method here. I think it's fine to keep the existing way to make it more readable.

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

2

This will be changed to read clinent-api later. So I won't bother it for now.

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

 //===- MemRefToSPIRV.h - MemRef to SPIR-V Patterns --------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Provides patterns to convert MemRef dialect to SPIR-V dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #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,
                                    RewritePatternSet &patterns);
 
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_MEMREFTOSPIRV_MEMREFTOSPIRV_H

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

 //===- MemRefToSPIRVPass.h - MemRef to SPIR-V Passes ------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Provides passes to convert MemRef dialect to SPIR-V dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #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 convert MemRef ops to SPIR-V ops.
 std::unique_ptr<OperationPass<ModuleOp>> createConvertMemRefToSPIRVPass();
 
+/// Creates a pass to map numeric MemRef memory spaces to symbolic SPIR-V
+/// storage classes. The mapping is given as `memorySpaceMap`.
+std::unique_ptr<OperationPass<ModuleOp>> createMapMemRefStorageClassPass(
+    const DenseMap<unsigned, spirv::StorageClass> &memorySpaceMap);
+
+/// 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();
+
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_MEMREFTOSPIRV_MEMREFTOSPIRVPASS_H

mlir/include/mlir/Conversion/Passes.td

Option<"indexBitwidth", "index-bitwidth", "unsigned",
            "Bitwidth of the index type, 0 to use size of machine word">,
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // 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()";
   let dependentDialects = ["spirv::SPIRVDialect"];
   let options = [
     Option<"boolNumBits", "bool-num-bits",
            "int", /*default=*/"8",
            "The number of bits to store a boolean value">

mlir/lib/Conversion/MemRefToSPIRV/CMakeLists.txt

 add_mlir_conversion_library(MLIRMemRefToSPIRV
+  MapMemRefStorageClassPass.cpp
   MemRefToSPIRV.cpp
   MemRefToSPIRVPass.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SPIRV
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/IR
 
   DEPENDS

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/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, 4> 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(
+    const spirv::MemorySpaceToStorageClassMap &memorySpaceMap) {
+  return std::make_unique<MapMemRefStorageClassPass>(memorySpaceMap);
+}
+
+std::unique_ptr<OperationPass<ModuleOp>>
+mlir::createMapMemRefStorageClassPass() {
+  return std::make_unique<MapMemRefStorageClassPass>();
+}

mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp

Optional<spirv::StorageClass> storageClass =
           type.getMemorySpaceAsInt());
   if (!storageClass)
     return {};
   switch (*storageClass) {
   case spirv::StorageClass::StorageBuffer:
     return spirv::Scope::Device;
   case spirv::StorageClass::Workgroup:
     return spirv::Scope::Workgroup;
-  default: {
-  }
+  default:
+    break;
   }
   return {};
 }
 
 /// Casts the given `srcInt` into a boolean value.
 static Value castIntNToBool(Location loc, Value srcInt, OpBuilder &builder) {
   if (srcInt.getType().isInteger(1))
     return srcInt;

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

+// 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
+}