[mlir][mem2reg] Add support for mem2reg in MemRef.

[Moxinilian]

Summary

This patch implements the mem2reg interfaces for MemRef types. This only supports scalar memrefs of a small list of types. It would be beneficial to create more interfaces for default values before expanding support to more types. Additionally, I am working on an upcoming revision to bring SROA to MLIR that should help with non-scalar memrefs.

[kuhar]

mlir/lib/Dialect/MemRef/IR/MemRefMem2Reg.cpp

101

nit: return {}; or return nullptr; for consistency with other code, e.g., LoadOp::getStored ?

mlir/test/Transforms/mem2reg-memref.mlir

136

Could we match the function argument here? This way it'd be easier to tell why the load below remained.

[kuhar]

mlir/lib/Dialect/MemRef/IR/MemRefMem2Reg.cpp

84–85

Do any of the testcases exercise this logic? It'd difficult for me to tell based just on the function names (used in tests)

84–85

By exercise I mean having tests where this condition both succeeds and fails

[Mogball]

Is there any limitation on the sizes of buffers that can be promoted? Obviously you don't want memref<16x16xf64> to get promoted.

mlir/test/Transforms/mem2reg-memref.mlir

1

I believe the policy is that this test should go with the rest of the MemRef dialect tests, not in the Transforms tests

[Moxinilian]

Address comments.

[Moxinilian]

Thanks for the reviews!

mlir/lib/Dialect/MemRef/IR/MemRefMem2Reg.cpp

38

@Mogball This implementation only allows mem2reg on scalar memrefs, this is the check enforcing that.

84–85

getMemRef() == slot.ptr for stores is exercised by deny_store_of_alloca.

For the rest, this logic is hard to trigger from mem2reg because it would typically mean that the interface implementation for allocators is buggy. So it could make sense to make this an assert it in theory, but the interface specification is more flexible than that, expecting no crash. Notably, nothing prevents a third-party user of the interface to ask for the use of a slot that has nothing to do with this load. Considering how the testing infrastructure works, I feel like the burden of testing falls on the user here, because I don't really know how to test it in practice. I would like to hear your opinion about this.

I still added a test (promotable_nonpromotable_intertwined) that would be likely to trigger this in case of bug.

mlir/test/Transforms/mem2reg-memref.mlir

1

Okay! I also moved the LLVM tests which were put in the same place.

[gysit]

LGTM modulo nit comment.

mlir/lib/Dialect/MemRef/IR/MemRefMem2Reg.cpp

48

50–51

Is it possible to move MemRefType into the lambda?

i.e.:

.Case([&](MemRefType type) {
  return builder.create<memref::AllocaOp>(getLoc(), type);
})

[Moxinilian]

Address comments.

[Moxinilian]

Mark as done.

[Mogball]

mlir/lib/Dialect/MemRef/IR/MemRefMem2Reg.cpp

38

swag

mlir/include/mlir/Dialect/MemRef/IR/MemRef.h

 
 #include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/CastInterfaces.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/CopyOpInterface.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
+#include "mlir/Interfaces/Mem2RegInterfaces.h"
 #include "mlir/Interfaces/ShapedOpInterfaces.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 #include <optional>
 
 namespace mlir {
 
 namespace arith {

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

 #define MEMREF_OPS
 
 include "mlir/Dialect/Arith/IR/ArithBase.td"
 include "mlir/Dialect/MemRef/IR/MemRefBase.td"
 include "mlir/Interfaces/CastInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/CopyOpInterface.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
+include "mlir/Interfaces/Mem2RegInterfaces.td"
 include "mlir/Interfaces/ShapedOpInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/Interfaces/ViewLikeInterface.td"
 include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/SymbolInterfaces.td"
 
 /// A TypeAttr for memref types.
 def MemRefTypeAttr
def MemRef_ReallocOp : MemRef_Op<"realloc"> {
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // AllocaOp
 //===----------------------------------------------------------------------===//
 
 def MemRef_AllocaOp : AllocLikeOp<"alloca", AutomaticAllocationScopeResource,[
-    DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>]> {
+    DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
+    DeclareOpInterfaceMethods<PromotableAllocationOpInterface>]> {
   let summary = "stack memory allocation operation";
   let description = [{
     The `alloca` operation allocates memory on the stack, to be automatically
     released when control transfers back from the region of its closest
     surrounding operation with an
     [`AutomaticAllocationScope`](../Traits.md/#automaticallocationscope) trait.
     The amount of memory allocated is specified by its memref and additional
     operands. For example:

 //===----------------------------------------------------------------------===//
 // LoadOp
 //===----------------------------------------------------------------------===//
 
 def LoadOp : MemRef_Op<"load",
      [TypesMatchWith<"result type matches element type of 'memref'",
                      "memref", "result",
                      "$_self.cast<MemRefType>().getElementType()">,
-                     MemRefsNormalizable]> {
+      MemRefsNormalizable,
+      DeclareOpInterfaceMethods<PromotableMemOpInterface>]> {
   let summary = "load operation";
   let description = [{
     The `load` op reads an element from a memref specified by an index list. The
     output of load is a new value with the same type as the elements of the
     memref. The arity of indices is the rank of the memref (i.e., if the memref
     loaded from is of rank 3, then 3 indices are required for the load following
     the memref identifier).
 

 //===----------------------------------------------------------------------===//
 // StoreOp
 //===----------------------------------------------------------------------===//
 
 def MemRef_StoreOp : MemRef_Op<"store",
      [TypesMatchWith<"type of 'value' matches element type of 'memref'",
                      "memref", "value",
                      "$_self.cast<MemRefType>().getElementType()">,
-                     MemRefsNormalizable]> {
+      MemRefsNormalizable,
+      DeclareOpInterfaceMethods<PromotableMemOpInterface>]> {
   let summary = "store operation";
   let description = [{
     Store a value to a memref location given by indices. The value stored should
     have the same type as the elemental type of the memref. The number of
     arguments provided within brackets need to match the rank of the memref.
 
     In an affine context, the indices of a store are restricted to SSA values
     bound to surrounding loop induction variables,

mlir/lib/Dialect/MemRef/IR/CMakeLists.txt

 add_mlir_dialect_library(MLIRMemRefDialect
   MemRefDialect.cpp
+  MemRefMem2Reg.cpp
   MemRefOps.cpp
   ValueBoundsOpInterfaceImpl.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${PROJECT_SOURCE_DIR}/inlude/mlir/Dialect/MemRefDialect
 
   DEPENDS
   MLIRMemRefOpsIncGen

mlir/lib/Dialect/MemRef/IR/MemRefMem2Reg.cpp

+//===- MemRefMem2Reg.cpp - Mem2Reg Interfaces -------------------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements Mem2Reg-related interfaces for MemRef dialect
+// operations.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Complex/IR/Complex.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Interfaces/Mem2RegInterfaces.h"
+#include "llvm/ADT/TypeSwitch.h"
+
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+//  AllocaOp interfaces
+//===----------------------------------------------------------------------===//
+
+static bool isSupportedElementType(Type type) {
+  return type.isa<MemRefType>() ||
+         OpBuilder(type.getContext()).getZeroAttr(type);
+}
+
+SmallVector<MemorySlot> memref::AllocaOp::getPromotableSlots() {
+  MemRefType type = getType();
+  if (!isSupportedElementType(type.getElementType()))
+    return {};
+  if (!type.hasStaticShape())
+    return {};
+  // Make sure the memref contains only a single element.
+  if (any_of(type.getShape(), [](uint64_t dim) { return dim != 1; }))
+    return {};
+
+  return {MemorySlot{getResult(), type.getElementType()}};
+}
+
+Value memref::AllocaOp::getDefaultValue(const MemorySlot &slot,
+                                        OpBuilder &builder) {
+  assert(isSupportedElementType(slot.elemType));
+  // TODO: support more types
+  return TypeSwitch<Type, Value>(slot.elemType)
+      .Case<MemRefType>(
+          [&](auto t) { return builder.create<memref::AllocaOp>(getLoc(), t); })
+      .Default([&](Type t) {
+        return builder.create<arith::ConstantOp>(getLoc(), t,
+                                                 builder.getZeroAttr(t));
+      });
+}
+
+void memref::AllocaOp::handlePromotionComplete(const MemorySlot &slot,
+                                               Value defaultValue) {
+  if (defaultValue.use_empty())
+    defaultValue.getDefiningOp()->erase();
+  erase();
+}
+
+void memref::AllocaOp::handleBlockArgument(const MemorySlot &slot,
+                                           BlockArgument argument,
+                                           OpBuilder &builder) {}
+
+//===----------------------------------------------------------------------===//
+//  LoadOp/StoreOp interfaces
+//===----------------------------------------------------------------------===//
+
+bool memref::LoadOp::loadsFrom(const MemorySlot &slot) {
+  return getMemRef() == slot.ptr;
+}
+
+Value memref::LoadOp::getStored(const MemorySlot &slot) { return {}; }
+
+bool memref::LoadOp::canUsesBeRemoved(
+    const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
+    SmallVectorImpl<OpOperand *> &newBlockingUses) {
+  if (blockingUses.size() != 1)
+    return false;
+  Value blockingUse = (*blockingUses.begin())->get();
+  return blockingUse == slot.ptr && getMemRef() == slot.ptr &&
+         getResult().getType() == slot.elemType;
+}
+
+DeletionKind memref::LoadOp::removeBlockingUses(
+    const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
+    OpBuilder &builder, Value reachingDefinition) {
+  // `canUsesBeRemoved` checked this blocking use must be the loaded slot
+  // pointer.
+  getResult().replaceAllUsesWith(reachingDefinition);
+  return DeletionKind::Delete;
+}
+
+bool memref::StoreOp::loadsFrom(const MemorySlot &slot) { return false; }
+
+Value memref::StoreOp::getStored(const MemorySlot &slot) {
+  if (getMemRef() != slot.ptr)
+    return Value();
+  return getValue();
+}
+
+bool memref::StoreOp::canUsesBeRemoved(
+    const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
+    SmallVectorImpl<OpOperand *> &newBlockingUses) {
+  if (blockingUses.size() != 1)
+    return false;
+  Value blockingUse = (*blockingUses.begin())->get();
+  return blockingUse == slot.ptr && getMemRef() == slot.ptr &&
+         getValue() != slot.ptr && getValue().getType() == slot.elemType;
+}
+
+DeletionKind memref::StoreOp::removeBlockingUses(
+    const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
+    OpBuilder &builder, Value reachingDefinition) {
+  return DeletionKind::Delete;
+}

mlir/test/Transforms/mem2reg-memref.mlir

+// RUN: ns-opt %s --pass-pipeline='builtin.module(func.func(mem2reg))' --split-input-file | FileCheck %s
+
+// CHECK-LABEL: func.func @basic
+func.func @basic() -> i32 {
+  // CHECK-NOT: = memref.alloca
+  // CHECK: %[[RES:.*]] = arith.constant 5 : i32
+  // CHECK-NOT: = memref.alloca
+  %0 = arith.constant 5 : i32
+  %1 = memref.alloca() : memref<i32>
+  memref.store %0, %1[] : memref<i32>
+  %2 = memref.load %1[] : memref<i32>
+  // CHECK: return %[[RES]] : i32
+  return %2 : i32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @basic_default
+func.func @basic_default() -> i32 {
+  // CHECK-NOT: = memref.alloca
+  // CHECK: %[[RES:.*]] = arith.constant 0 : i32
+  // CHECK-NOT: = memref.alloca
+  %0 = arith.constant 5 : i32
+  %1 = memref.alloca() : memref<i32>
+  %2 = memref.load %1[] : memref<i32>
+  // CHECK-NOT: memref.store
+  memref.store %0, %1[] : memref<i32>
+  // CHECK: return %[[RES]] : i32
+  return %2 : i32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @basic_float
+func.func @basic_float() -> f32 {
+  // CHECK-NOT: = memref.alloca
+  // CHECK: %[[RES:.*]] = arith.constant {{.*}} : f32
+  %0 = arith.constant 5.2 : f32
+  // CHECK-NOT: = memref.alloca
+  %1 = memref.alloca() : memref<f32>
+  memref.store %0, %1[] : memref<f32>
+  %2 = memref.load %1[] : memref<f32>
+  // CHECK: return %[[RES]] : f32
+  return %2 : f32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @basic_ranked
+func.func @basic_ranked() -> i32 {
+  // CHECK-NOT: = memref.alloca
+  // CHECK: %[[RES:.*]] = arith.constant 5 : i32
+  // CHECK-NOT: = memref.alloca
+  %0 = arith.constant 0 : index
+  %1 = arith.constant 5 : i32
+  %2 = memref.alloca() : memref<1x1xi32>
+  memref.store %1, %2[%0, %0] : memref<1x1xi32>
+  %3 = memref.load %2[%0, %0] : memref<1x1xi32>
+  // CHECK: return %[[RES]] : i32
+  return %3 : i32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @reject_multiple_elements
+func.func @reject_multiple_elements() -> i32 {
+  // CHECK: %[[INDEX:.*]] = arith.constant 0 : index
+  %0 = arith.constant 0 : index
+  // CHECK: %[[STORED:.*]] = arith.constant 5 : i32
+  %1 = arith.constant 5 : i32
+  // CHECK: %[[ALLOCA:.*]] = memref.alloca()
+  %2 = memref.alloca() : memref<1x2xi32>
+  // CHECK: memref.store %[[STORED]], %[[ALLOCA]][%[[INDEX]], %[[INDEX]]]
+  memref.store %1, %2[%0, %0] : memref<1x2xi32>
+  // CHECK: %[[RES:.*]] = memref.load %[[ALLOCA]][%[[INDEX]], %[[INDEX]]]
+  %3 = memref.load %2[%0, %0] : memref<1x2xi32>
+  // CHECK: return %[[RES]] : i32
+  return %3 : i32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @cycle
+// CHECK-SAME: (%[[ARG0:.*]]: i64, %[[ARG1:.*]]: i1, %[[ARG2:.*]]: i64)
+func.func @cycle(%arg0: i64, %arg1: i1, %arg2: i64) {
+  // CHECK-NOT: = memref.alloca
+  %alloca = memref.alloca() : memref<i64>
+  memref.store %arg2, %alloca[] : memref<i64>
+  // CHECK: cf.cond_br %[[ARG1:.*]], ^[[BB1:.*]](%[[ARG2]] : i64), ^[[BB2:.*]](%[[ARG2]] : i64)
+  cf.cond_br %arg1, ^bb1, ^bb2
+// CHECK: ^[[BB1]](%[[USE:.*]]: i64):
+^bb1:
+  %use = memref.load %alloca[] : memref<i64>
+  // CHECK: call @use(%[[USE]])
+  func.call @use(%use) : (i64) -> ()
+  memref.store %arg0, %alloca[] : memref<i64>
+  // CHECK: cf.br ^[[BB2]](%[[ARG0]] : i64)
+  cf.br ^bb2
+// CHECK: ^[[BB2]](%[[FWD:.*]]: i64):
+^bb2:
+  // CHECK: cf.br ^[[BB1]](%[[FWD]] : i64)
+  cf.br ^bb1
+}
+
+func.func @use(%arg: i64) { return }
+
+// -----
+
+// CHECK-LABEL: func.func @recursive
+// CHECK-SAME: (%[[ARG:.*]]: i64)
+func.func @recursive(%arg: i64) -> i64 {
+  // CHECK-NOT: = memref.alloca()
+  %alloca0 = memref.alloca() : memref<memref<memref<i64>>>
+  %alloca1 = memref.alloca() : memref<memref<i64>>
+  %alloca2 = memref.alloca() : memref<i64>
+  memref.store %arg, %alloca2[] : memref<i64>
+  memref.store %alloca2, %alloca1[] : memref<memref<i64>>
+  memref.store %alloca1, %alloca0[] : memref<memref<memref<i64>>>
+  %load0 = memref.load %alloca0[] : memref<memref<memref<i64>>>
+  %load1 = memref.load %load0[] : memref<memref<i64>>
+  %load2 = memref.load %load1[] : memref<i64>
+  // CHECK: return %[[ARG]] : i64
+  return %load2 : i64
+}
+
+// -----
+
+// CHECK-LABEL: func.func @deny_store_of_alloca
+func.func @deny_store_of_alloca(%arg: memref<memref<i32>>) -> i32 {
+  // CHECK: %[[VALUE:.*]] = arith.constant 5 : i32
+  %0 = arith.constant 5 : i32
+  // CHECK: %[[ALLOCA:.*]] = memref.alloca
+  %1 = memref.alloca() : memref<i32>
+  // CHECK: memref.store %[[VALUE]], %[[ALLOCA]][]
+  memref.store %0, %1[] : memref<i32>
+  // CHECK: memref.store %[[ALLOCA]], %{{.*}}[]
+  memref.store %1, %arg[] : memref<memref<i32>>
+  // CHECK: %[[RES:.*]] = memref.load %[[ALLOCA]][]
+  %2 = memref.load %1[] : memref<i32>
+  // CHECK: return %[[RES]] : i32
+  return %2 : i32
+}