[Mlir-commits] [mlir] [mlir][bufferization] Adding the optimize-allocation-liveness pass (PR #101827)

[Matthias Springer]

================
@@ -0,0 +1,157 @@
+//===- OptimizeAllocationLiveness.cpp - impl. optimize allocation liveness pass
+//-===//
+//
+// 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 for optimizing allocation liveness.
+// The pass moves the deallocation operation after the last user of the
+// allocated buffer.
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.h"
+#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/Operation.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "optimize-allocation-liveness"
+#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
+#define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n")
+
+namespace mlir {
+namespace bufferization {
+#define GEN_PASS_DEF_OPTIMIZEALLOCATIONLIVENESS
+#include "mlir/Dialect/Bufferization/Transforms/Passes.h.inc"
+} // namespace bufferization
+} // namespace mlir
+
+using namespace mlir;
+
+namespace {
+
+//===----------------------------------------------------------------------===//
+// Helper functions
+//===----------------------------------------------------------------------===//
+
+/// Return true if `a` happens before `b`, i.e., `a` or one of its ancestors
+/// properly dominates `b` and `b` is not inside `a`.
+static bool happensBefore(Operation *a, Operation *b) {
+  do {
+    if (a->isProperAncestor(b))
+      return false;
+    if (Operation *bAncestor = a->getBlock()->findAncestorOpInBlock(*b)) {
+      return a->isBeforeInBlock(bAncestor);
+    }
+  } while ((a = a->getParentOp()));
+  return false;
+}
+
+/// This method searches for a user of value that is a dealloc operation.
+/// If multiple users with free effect are found, return nullptr.
+Operation *findUserWithFreeSideEffect(Value value) {
+  Operation *freeOpUser = nullptr;
+  for (Operation *user : value.getUsers()) {
+    if (MemoryEffectOpInterface memEffectOp =
+            dyn_cast<MemoryEffectOpInterface>(user)) {
+      SmallVector<MemoryEffects::EffectInstance, 2> effects;
+      memEffectOp.getEffects(effects);
+
+      for (const auto &effect : effects) {
+        if (isa<MemoryEffects::Free>(effect.getEffect())) {
+          if (freeOpUser) {
+            LDBG("Multiple users with free effect found: " << *freeOpUser
+                                                           << " and " << *user);
+            return nullptr;
+          }
+          freeOpUser = user;
+        }
+      }
+    }
+  }
+  return freeOpUser;
+}
+
+/// Checks if the given op allocates memory.
+static bool hasMemoryAllocEffect(MemoryEffectOpInterface memEffectOp) {
+  SmallVector<MemoryEffects::EffectInstance, 2> effects;
+  memEffectOp.getEffects(effects);
+  for (const auto &effect : effects) {
+    if (isa<MemoryEffects::Allocate>(effect.getEffect())) {
+      return true;
+    }
+  }
+  return false;
+}
+
+struct OptimizeAllocationLiveness
+    : public bufferization::impl::OptimizeAllocationLivenessBase<
+          OptimizeAllocationLiveness> {
+public:
+  OptimizeAllocationLiveness() = default;
+
+  void runOnOperation() override {
+    func::FuncOp func = getOperation();
+
+    if (func.isExternal())
+      return;
+
+    BufferViewFlowAnalysis analysis = BufferViewFlowAnalysis(func);
+
+    func.walk([&](MemoryEffectOpInterface memEffectOp) -> WalkResult {
+      if (!hasMemoryAllocEffect(memEffectOp))
+        return WalkResult::advance();
+
+      auto allocOp = memEffectOp;
+      LDBG("Checking alloc op: " << allocOp);
+
+      auto deallocOp = findUserWithFreeSideEffect(allocOp->getResult(0));
+      if (!deallocOp)
----------------
matthias-springer wrote:

I think the dealloc op must be in the same block as the dealloc op. You already have an assertion for that below, but it is not cheked here.

E.g., we should not hoist a dealloc op out of a conditional branch:

```mlir
%0 = memref.alloc
scf.if %condition {
  memref.dealloc %0
}
```

Can you also add a test case for that?


https://github.com/llvm/llvm-project/pull/101827