[mlir][memref] Add support for erasing dead allocations.

[hanchung]

fix bazel

[hanchung]

fix bazel

[springerm]

mlir/include/mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.td

177

Add: It modifies the payload. Dead allocations, loads and stores are silently dropped from all mappings.

mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp

128

I would implement in terms of: static bool valueIsNotRead(OpResult opResult, std::vector<Operation *> &users) and call with the result of the AllocOp.

132–133

You can query memory side effects instead of hard-coding the ops here:

Add a templatized copy of bool isMemoryEffectFree(Operation *op); to SideEffectInterface.h and SideEffectInterfaces.cpp, where the side effect can be specified as template parameter. In the implementation, use !hasEffect<EffectTy>(op) instead of !memInterface.hasNoEffect().

Then call this function here:

if (useOp->getNumResults() == 0 && mlir::isSideEffectFree<MemoryEffects::Read>(useOp)) {
  // There are no reads and no alias is created.
  ...
}

(Still need special handling for memref::DeallocOp.)

134

If it's a memref.subview, you can call allUsesAreStores recursively. (Or have a worklist of operations in this function.) Then you can also handle memref.subview(memref.subview(alloc)).

144

You need a RewriterBase &rewriter here and modify the IR with that rewriter. (Pass it from applyToOne.) Otherwise, the transform dialect state will be inconsistent.

146

Could also do this for memref::AllocaOp.

[hanchung]

address comments

mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp

132–133

I used SideEffect. I'm not very familiar with them, please take a look if I implemented something wrong. Thanks!

134

I don't follow the suggestion. Isn't it already doing what you mentioned? If it is a memref.subview op, it calls the method recursively.

[hanchung]

fix bazel deps

[springerm]

mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp

128

If you keep the Operation * argument, I would rename it to resultIsNotRead.

133

I would also check that useOp has no regions because there could be another use of op inside of a region.

155

rewriter.eraseOp(op)

[hanchung]

address comments

mlir/include/mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.td

let arguments =
            UnitAttr:$skip_analysis);
 
   let results = (outs TransformHandleTypeInterface:$transformed);
 
   let assemblyFormat =
     "$target attr-dict `:` functional-type(operands, results)";
 }
 
+def MemRefEraseDeadAllocAndStoresOp
+    : Op<Transform_Dialect, "memref.erase_dead_alloc_and_stores", [
+      TransformEachOpTrait, TransformOpInterface,
+      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+      ReportTrackingListenerFailuresOpTrait
+    ]> {
+  let description = [{
+    This applies memory optimization on memref. In particular it does store to
+    load forwarding, dead store elimination and dead alloc elimination.
+
+    #### Return modes
+
+    This operation applies a set of memory optimization on the whole region of
+    the operand.
+
+    The transformation does not consume the target handle. It modifies the
+    payload. Dead allocations, loads and stores are silently dropped from all
+    mappings.
+  }];
+
+  let arguments = (ins TransformHandleTypeInterface:$target);
+  let results = (outs);
+
+  let assemblyFormat = "$target attr-dict `:` functional-type($target, results)";
+
+  let skipDefaultBuilders = 1;
+  let builders = [
+    OpBuilder<(ins "Value":$target)>
+  ];
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
+        ::mlir::Operation *target,
+        ::mlir::transform::ApplyToEachResultList &results,
+        ::mlir::transform::TransformState &state);
+  }];
+}
+
 def MemRefMakeLoopIndependentOp
     : Op<Transform_Dialect, "memref.make_loop_independent",
          [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
           TransformOpInterface, TransformEachOpTrait]> {
   let description = [{
     Rewrite the targeted ops such that their index-typed operands no longer
     depend on any loop induction variable of the `num_loop` enclosing `scf.for`
     loops. I.e., compute an upper bound that is independent of any such loop IV

mlir/include/mlir/Dialect/MemRef/Utils/MemRefUtils.h

 ///   element type with bitwidth `srcBits` using element type with
 ///   bitwidth `dstBits`, the linearized offset and size are
 ///   scaled down by `dstBits`/`srcBits`.
 LinearizedMemRefInfo
 getLinearizedMemRefOffsetAndSize(OpBuilder &builder, Location loc, int srcBits,
                                  int dstBits, OpFoldResult offset,
                                  ArrayRef<OpFoldResult> sizes);
 
+// Track temporary allocations that are never read from. If this is the case
+// it means both the allocations and associated stores can be removed.
+void eraseDeadAllocAndStores(RewriterBase &rewriter, Operation *parentOp);
+
 } // namespace memref
 } // namespace mlir
 
 #endif // MLIR_DIALECT_MEMREF_UTILS_MEMREFUTILS_H

mlir/lib/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp

 //===- MemRefTransformOps.cpp - Implementation of Memref transform ops ----===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.h"
 
 #include "mlir/Conversion/LLVMCommon/TypeConverter.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/Dialect/MemRef/Transforms/Transforms.h"
+#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
 #include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/Support/Debug.h"
 
 using namespace mlir;
 
 #define DEBUG_TYPE "memref-transforms"
 #define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
for (Operation *op : state.getPayloadOps(getTarget())) {
 
     results.push_back(*newBuffer);
   }
   transformResults.set(cast<OpResult>(getResult()), results);
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
+// MemRefEraseDeadAllocAndStoresOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure
+transform::MemRefEraseDeadAllocAndStoresOp::applyToOne(
+    transform::TransformRewriter &rewriter, Operation *target,
+    transform::ApplyToEachResultList &results,
+    transform::TransformState &state) {
+  // Apply store to load forwarding and dead store elimination.
+  vector::transferOpflowOpt(rewriter, target);
+  memref::eraseDeadAllocAndStores(rewriter, target);
+  return DiagnosedSilenceableFailure::success();
+}
+
+void transform::MemRefEraseDeadAllocAndStoresOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  transform::onlyReadsHandle(getTarget(), effects);
+  transform::modifiesPayload(effects);
+}
+void transform::MemRefEraseDeadAllocAndStoresOp::build(OpBuilder &builder,
+                                                       OperationState &result,
+                                                       Value target) {
+  result.addOperands(target);
+}
+
+//===----------------------------------------------------------------------===//
 // MemRefMakeLoopIndependentOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure transform::MemRefMakeLoopIndependentOp::applyToOne(
     transform::TransformRewriter &rewriter, Operation *target,
     transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
   // Gather IVs.

mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp

 //===- MemRefUtils.cpp - Utilities to support the MemRef dialect ----------===//
 //
 // 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 utilities for the MemRef dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
 
 namespace mlir {
 namespace memref {
 
 bool isStaticShapeAndContiguousRowMajor(MemRefType type) {
   if (!type.hasStaticShape())
     return false;
 
getLinearizedMemRefOffsetAndSize(OpBuilder &builder, Location loc, int srcBits,
 
   LinearizedMemRefInfo linearizedMemRefInfo;
   std::tie(linearizedMemRefInfo, std::ignore) =
       getLinearizedMemRefOffsetAndSize(builder, loc, srcBits, dstBits, offset,
                                        sizes, strides);
   return linearizedMemRefInfo;
 }
 
+/// Returns true if all the uses of op are not read/load.
+/// There can be SubviewOp users as long as all its users are also
+/// StoreOp/transfer_write. If return true it also fills out the uses, if it
+/// returns false uses is unchanged.
+static bool resultIsNotRead(Operation *op, std::vector<Operation *> &uses) {
+  std::vector<Operation *> opUses;
+  for (OpOperand &use : op->getUses()) {
+    Operation *useOp = use.getOwner();
+    if (isa<memref::DeallocOp>(useOp) ||
+        (useOp->getNumResults() == 0 && useOp->getNumRegions() == 0 &&
+         !mlir::hasEffect<MemoryEffects::Read>(useOp)) ||
+        (isa<memref::SubViewOp>(useOp) && resultIsNotRead(useOp, opUses))) {
+      opUses.push_back(useOp);
+      continue;
+    }
+    return false;
+  }
+  uses.insert(uses.end(), opUses.begin(), opUses.end());
+  return true;
+}
+
+void eraseDeadAllocAndStores(RewriterBase &rewriter, Operation *parentOp) {
+  std::vector<Operation *> opToErase;
+  parentOp->walk([&](memref::AllocOp op) {
+    std::vector<Operation *> candidates;
+    if (resultIsNotRead(op, candidates)) {
+      opToErase.insert(opToErase.end(), candidates.begin(), candidates.end());
+      opToErase.push_back(op.getOperation());
+    }
+  });
+  for (Operation *op : opToErase)
+    rewriter.eraseOp(op);
+}
+
 } // namespace memref
 } // namespace mlir

mlir/test/Dialect/MemRef/transform-ops.mlir

^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["memref.alloc"]} in %arg1 : (!transform.any_op) -> !transform.op<"memref.alloc">
   %1 = transform.memref.multibuffer %0 {factor = 2 : i64, skip_analysis} : (!transform.op<"memref.alloc">) -> !transform.any_op
   // Verify that the returned handle is usable.
   transform.test_print_remark_at_operand %1, "transformed" : !transform.any_op
 }
 
 // -----
 
+// CHECK-LABEL: func.func @dead_alloc
+func.func @dead_alloc() {
+  // CHECK-NOT: %{{.+}} = memref.alloc
+  %0 = memref.alloc() : memref<8x64xf32, 3>
+  %1 = memref.subview %0[0, 0] [8, 4] [1, 1] : memref<8x64xf32, 3> to
+    memref<8x4xf32, affine_map<(d0, d1) -> (d0 * 64 + d1)>, 3>
+  %c0 = arith.constant 0 : index
+  %cst_0 = arith.constant dense<0.000000e+00> : vector<1x4xf32>
+  vector.transfer_write %cst_0, %1[%c0, %c0] {in_bounds = [true, true]} :
+    vector<1x4xf32>, memref<8x4xf32, affine_map<(d0, d1) -> (d0 * 64 + d1)>, 3>
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  transform.memref.erase_dead_alloc_and_stores %0 : (!transform.any_op) -> ()
+}
+
+// -----
+
+// CHECK-LABEL: @store_to_load
+//  CHECK-SAME:   (%[[ARG:.+]]: vector<4xf32>)
+//   CHECK-NOT:   memref.alloc()
+//   CHECK-NOT:   vector.transfer_write
+//   CHECK-NOT:   vector.transfer_read
+//       CHECK:   return %[[ARG]] : vector<4xf32>
+func.func @store_to_load(%arg: vector<4xf32>) -> vector<4xf32> {
+  %c0 = arith.constant 0 : index
+  %cst_1 = arith.constant 0.000000e+00 : f32
+  %alloc = memref.alloc() {alignment = 64 : i64} : memref<64xf32>
+  vector.transfer_write %arg, %alloc[%c0] {in_bounds = [true]} : vector<4xf32>, memref<64xf32>
+  %r = vector.transfer_read %alloc[%c0], %cst_1 {in_bounds = [true]} : memref<64xf32>, vector<4xf32>
+  return %r : vector<4xf32>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  transform.memref.erase_dead_alloc_and_stores %0 : (!transform.any_op) -> ()
+}
+
+// -----
+
 // CHECK-LABEL: func @lower_to_llvm
 //   CHECK-NOT:   memref.alloc
 //       CHECK:   llvm.call @malloc
 func.func @lower_to_llvm() {
   %0 = memref.alloc() : memref<2048xi8>
   return
 }
 

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

cc_library(
     hdrs = [
         "include/mlir/Dialect/MemRef/Utils/MemRefUtils.h",
     ],
     includes = ["include"],
     deps = [
         ":AffineDialect",
         ":ArithUtils",
         ":MemRefDialect",
+        ":VectorDialect",
     ],
 )
 
 gentbl_cc_library(
     name = "MemRefPassIncGen",
     tbl_outs = [
         (
             [
deps = [
         ":AffineDialect",
         ":ArithDialect",
         ":IR",
         ":LLVMCommonConversion",
         ":LoopLikeInterface",
         ":MemRefDialect",
         ":MemRefTransformOpsIncGen",
         ":MemRefTransforms",
+        ":MemRefUtils",
         ":NVGPUDialect",
         ":SCFDialect",
         ":TransformDialect",
         ":TransformUtils",
         ":VectorDialect",
+        ":VectorTransforms",
         "//llvm:Support",
     ],
 )
 
 ##---------------------------------------------------------------------------##
 # MLProgram dialect
 ##---------------------------------------------------------------------------##