[mlir][bufferize] EmptyTensorElimination: Keep AnalysisState up to date

[springerm]

Summary

When EmptyTensorElimination eliminates a tensor.empty op, the given OneShotAnalysisState is updated accordingly. This makes is possible to reuse the same state object for One-Shot Bufferize (to address performance bottlenecks).

Depends On: D143379

mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h

public:
   /// aliasing sets.
   void bufferizeInPlace(OpOperand &operand);
 
   /// Mark the given OpOperand as out-of-place.
   void bufferizeOutOfPlace(OpOperand &operand);
 
   /// Add a new entry for `v` in the `aliasInfo` and `equivalentInfo`. In the
   /// beginning the alias and equivalence sets only contain `v` itself.
-  void createAliasInfoEntry(Value v);
+  void createAliasEntry(Value v);
+
+  /// Add `aliasInfo` and `equivalentInfo` entries for a newly created op.
+  void createAliasEntriesForNewOp(Operation *op);
 
   /// Find all tensor values in the given operation that have undefined contents
   /// and store them in `undefinedTensorUses`.
   void gatherUndefinedTensorUses(Operation *op);
 
   /// Find all tensors that are yielded/returned from a block and store them in
   /// `yieldedTensors`. Also include all aliasing tensors in the same block.
   void gatherYieldedTensors(Operation *op);

mlir/include/mlir/Dialect/Bufferization/Transforms/Transforms.h

 
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/IR/Operation.h"
 
 namespace mlir {
 namespace bufferization {
 class AnalysisState;
 struct BufferizationStatistics;
+class OneShotAnalysisState;
 struct OneShotBufferizationOptions;
 
 /// A function that matches anchor OpOperands for tensor::EmptyOp elimination.
 /// If an OpOperand is matched, the function should populate the SmallVector
 /// with all values that are needed during `RewriteFn` to produce the
 /// replacement value.
 using AnchorMatchFn = std::function<bool(OpOperand &, SmallVector<Value> &)>;
 
 /// A function that rewrites matched anchors.
 using RewriteFn = std::function<Value(OpBuilder &, Location, OpOperand &)>;
 
 /// Try to eliminate tensor::EmptyOps inside `op`.
 ///
 /// * `rewriteFunc` generates the replacement for the tensor::EmptyOp.
 /// * Only tensor::EmptyOps that are anchored on a matching OpOperand as per
 ///   `anchorMatchFunc` are considered. "Anchored" means that there is a path
 ///   on the reverse SSA use-def chain, starting from the OpOperand and always
 ///   following the aliasing  OpOperand, that eventually ends at a single
 ///   tensor::EmptyOp.
-LogicalResult eliminateEmptyTensors(RewriterBase &rewriter, Operation *op,
-                                    bufferization::AnalysisState &state,
+LogicalResult eliminateEmptyTensors(Operation *op, OneShotAnalysisState &state,
                                     AnchorMatchFn anchorMatchFunc,
                                     RewriteFn rewriteFunc);
 
 /// Try to eliminate tensor::EmptyOps inside `op` that are anchored on an
 /// InsertSliceOp, i.e., if it is eventually inserted into another tensor
 /// (and some other conditions are met).
-LogicalResult insertSliceAnchoredEmptyTensorEliminationStep(
-    RewriterBase &rewriter, Operation *op, bufferization::AnalysisState &state);
+LogicalResult
+insertSliceAnchoredEmptyTensorEliminationStep(Operation *op,
+                                              OneShotAnalysisState &state);
 
 /// Resolve RaW and other conflicts by inserting bufferization.alloc_tensor ops.
 /// After applying this transform, the IR can be bufferized without inserting
 /// additional buffer allocations.
 LogicalResult insertTensorCopies(Operation *op,
                                  const OneShotBufferizationOptions &options,
                                  BufferizationStatistics *statistics = nullptr);
 

mlir/lib/Dialect/Bufferization/Transforms/EmptyTensorElimination.cpp

 #define GEN_PASS_DEF_EMPTYTENSORELIMINATION
 #include "mlir/Dialect/Bufferization/Transforms/Passes.h.inc"
 } // namespace bufferization
 } // namespace mlir
 
 using namespace mlir;
 using namespace mlir::bufferization;
 
+namespace {
+/// A rewriter that tracks new ops and modified ops.
+class TrackingRewriter : public RewriterBase {
+public:
+  TrackingRewriter(MLIRContext *ctx, OneShotAnalysisState &state)
+      : RewriterBase(ctx), state(state) {}
+
+  void notifyOperationInserted(Operation *op) override {
+    state.createAliasEntriesForNewOp(op);
+    newOps.push_back(op);
+  }
+
+  void finalizeRootUpdate(Operation *op) override { modifiedOps.push_back(op); }
+
+  ArrayRef<Operation *> getNewOps() const { return newOps; }
+  ArrayRef<Operation *> getModifiedOps() const { return modifiedOps; }
+
+private:
+  OneShotAnalysisState &state;
+
+  SmallVector<Operation *> newOps, modifiedOps;
+};
+} // namespace
+
 /// Return true if all `neededValues` are in scope at the given
 /// `insertionPoint`.
 static bool
 neededValuesDominateInsertionPoint(const DominanceInfo &domInfo,
                                    Operation *insertionPoint,
                                    const SmallVector<Value> &neededValues) {
   for (Value val : neededValues) {
     if (auto bbArg = val.dyn_cast<BlockArgument>()) {

 }
 
 /// Try to eliminate tensor::EmptyOps inside `op`. A tensor::EmptyOp is replaced
 /// with the result of `rewriteFunc` if it is anchored on a matching
 /// OpOperand. "Anchored" means that there is a path on the reverse SSA use-def
 /// chain, starting from the OpOperand and always following the aliasing
 /// OpOperand, that eventually ends at a single tensor::EmptyOp.
 LogicalResult mlir::bufferization::eliminateEmptyTensors(
-    RewriterBase &rewriter, Operation *op, AnalysisState &state,
-    AnchorMatchFn anchorMatchFunc, RewriteFn rewriteFunc) {
-  OpBuilder::InsertionGuard g(rewriter);
+    Operation *op, OneShotAnalysisState &state, AnchorMatchFn anchorMatchFunc,
+    RewriteFn rewriteFunc) {
+  TrackingRewriter rewriter(op->getContext(), state);
 
   WalkResult status = op->walk([&](Operation *op) {
     for (OpOperand &operand : op->getOpOperands()) {
       // Skip operands that do not bufferize inplace.
       if (!state.isInPlace(operand))
         continue;
       // All values that are needed to create the replacement op.
       SmallVector<Value> neededValues;
for (OpOperand &operand : op->getOpOperands()) {
 
       // Create a replacement for the tensor::EmptyOp.
       rewriter.setInsertionPoint(insertionPoint);
       Value replacement = rewriteFunc(rewriter, emptyTensor.getLoc(), operand);
       if (!replacement)
         continue;
 
       // Replace the tensor::EmptyOp.
-      rewriter.replaceOp(emptyTensor.getDefiningOp(), replacement);
+      rewriter.replaceAllUsesWith(emptyTensor, replacement);
     }
 
     // Advance to the next operation.
     return WalkResult::advance();
   });
 
-  return failure(status.wasInterrupted());
+  if (status.wasInterrupted())
+    return failure();
+
+  DominanceInfo domInfo;
+  // Re-analyze modified ops.
+  for (Operation *op : rewriter.getModifiedOps())
+    if (failed(state.analyzeSingleOp(op, domInfo)))
+      return failure();
+  // Analyze new ops.
+  for (Operation *op : rewriter.getNewOps())
+    if (failed(state.analyzeSingleOp(op, domInfo)))
+      return failure();
+
+  return success();
 }
 
 /// Try to eliminate tensor::EmptyOps inside `op`. An tensor::EmptyOp can be
 /// eliminated if it is eventually inserted into another tensor (and some other
 /// conditions are met).
 ///
 /// E.g.:
 /// %0 = tensor.empty()

 /// Starting from an InsertSliceOp, an tensor::EmptyOp at the end of the insert
 /// source's reverse use-def chain is eliminated if:
 /// * On the reverse use-def chain path from the InsertSliceOp to the
 ///   tensor::EmptyOp, all ops were decided to bufferize inplace and the buffer
 ///   relation is "equivalent" (TODO: can be relaxed if needed).
 /// * The reverse use-def chain has exactly one end, which is the
 ///   tensor::EmptyOp.
 template <typename OpTy>
-static LogicalResult insertSliceLikeAnchoredEmptyTensorEliminationStep(
-    RewriterBase &rewriter, Operation *op, AnalysisState &state) {
+static LogicalResult
+insertSliceLikeAnchoredEmptyTensorEliminationStep(Operation *op,
+                                                  OneShotAnalysisState &state) {
   return eliminateEmptyTensors(
-      rewriter, op, state,
+      op, state,
       /*anchorMatchFunc=*/
       [&](OpOperand &operand, SmallVector<Value> &neededValues) {
         auto insertSliceOp = dyn_cast<OpTy>(operand.getOwner());
         if (!insertSliceOp)
           return false;
         if (&operand != &insertSliceOp->getOpOperand(0) /*source*/)
           return false;
 
return eliminateEmptyTensors(
             insertOp.getMixedOffsets(), insertOp.getMixedSizes(),
             insertOp.getMixedStrides());
         return extractOp.getResult();
       });
 }
 
 LogicalResult
 mlir::bufferization::insertSliceAnchoredEmptyTensorEliminationStep(
-    RewriterBase &rewriter, Operation *op, AnalysisState &state) {
+    Operation *op, OneShotAnalysisState &state) {
   if (failed(insertSliceLikeAnchoredEmptyTensorEliminationStep<
-             tensor::InsertSliceOp>(rewriter, op, state)))
+             tensor::InsertSliceOp>(op, state)))
     return failure();
   if (failed(insertSliceLikeAnchoredEmptyTensorEliminationStep<
-             tensor::ParallelInsertSliceOp>(rewriter, op, state)))
+             tensor::ParallelInsertSliceOp>(op, state)))
     return failure();
   return success();
 }
 
 namespace {
 struct EmptyTensorElimination
     : public bufferization::impl::EmptyTensorEliminationBase<
           EmptyTensorElimination> {
void EmptyTensorElimination::runOnOperation() {
   Operation *op = getOperation();
   OneShotBufferizationOptions options;
   OneShotAnalysisState state(op, options);
   if (failed(analyzeOp(op, state))) {
     signalPassFailure();
     return;
   }
 
-  IRRewriter rewriter(op->getContext());
   if (failed(bufferization::insertSliceAnchoredEmptyTensorEliminationStep(
-          rewriter, op, state)))
+          op, state)))
     signalPassFailure();
 }
 
 std::unique_ptr<Pass> mlir::bufferization::createEmptyTensorEliminationPass() {
   return std::make_unique<EmptyTensorElimination>();
 }

mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp

 
 //===----------------------------------------------------------------------===//
 // OneShotAnalysisState
 //===----------------------------------------------------------------------===//
 
 OneShotAnalysisState::OneShotAnalysisState(
     Operation *op, const OneShotBufferizationOptions &options)
     : AnalysisState(options, TypeID::get<OneShotAnalysisState>()) {
-  // Set up alias sets.
-  op->walk([&](Operation *op) {
-    for (Value v : op->getResults())
-      if (v.getType().isa<TensorType>())
-        createAliasInfoEntry(v);
-    for (Region &r : op->getRegions())
-      for (Block &b : r.getBlocks())
-        for (auto bbArg : b.getArguments())
-          if (bbArg.getType().isa<TensorType>())
-            createAliasInfoEntry(bbArg);
-  });
-
-  // Mark OpOperands in-place that must bufferize in-place.
-  op->walk([&](BufferizableOpInterface bufferizableOp) {
-    if (!options.isOpAllowed(bufferizableOp))
-      return WalkResult::skip();
-    for (OpOperand &opOperand : bufferizableOp->getOpOperands())
-      if (opOperand.get().getType().isa<TensorType>())
-        if (bufferizableOp.mustBufferizeInPlace(opOperand, *this))
-          bufferizeInPlace(opOperand);
-    return WalkResult::advance();
-  });
+  op->walk([&](Operation *op) { createAliasEntriesForNewOp(op); });
 }
 
 void OneShotAnalysisState::applyOnEquivalenceClass(
     Value v, function_ref<void(Value)> fun) const {
   auto leaderIt = equivalentInfo.findLeader(v);
   for (auto mit = leaderIt, meit = equivalentInfo.member_end(); mit != meit;
        ++mit) {
     fun(*mit);
bool OneShotAnalysisState::areAliasingBufferizedValues(Value v1,
                                                        Value v2) const {
   return aliasInfo.isEquivalent(v1, v2);
 }
 
 void OneShotAnalysisState::bufferizeInPlace(OpOperand &operand) {
   if (inplaceBufferized.contains(&operand))
     return;
   inplaceBufferized.insert(&operand);
-  for (OpResult result : getAliasingOpResults(operand))
+  for (OpResult result : getAliasingOpResults(operand)) {
+    aliasInfo.insert(result);
+    aliasInfo.insert(operand.get());
     aliasInfo.unionSets(result, operand.get());
+  }
   ++statNumTensorInPlace;
 }
 
 void OneShotAnalysisState::bufferizeOutOfPlace(OpOperand &operand) {
   assert(!inplaceBufferized.contains(&operand) &&
          "OpOperand was already decided to bufferize inplace");
   ++statNumTensorOutOfPlace;
 }
 
-void OneShotAnalysisState::createAliasInfoEntry(Value v) {
+void OneShotAnalysisState::createAliasEntry(Value v) {
   aliasInfo.insert(v);
   equivalentInfo.insert(v);
 }
 
+void OneShotAnalysisState::createAliasEntriesForNewOp(Operation *op) {
+  // Set up alias sets.
+  for (Value v : op->getResults())
+    if (v.getType().isa<TensorType>())
+      createAliasEntry(v);
+  for (Region &r : op->getRegions())
+    for (Block &b : r.getBlocks())
+      for (auto bbArg : b.getArguments())
+        if (bbArg.getType().isa<TensorType>())
+          createAliasEntry(bbArg);
+
+  // Mark OpOperands in-place that must bufferize in-place.
+  auto bufferizableOp = getOptions().dynCastBufferizableOp(op);
+  if (!bufferizableOp)
+    return;
+  for (OpOperand &opOperand : bufferizableOp->getOpOperands())
+    if (opOperand.get().getType().isa<TensorType>())
+      if (bufferizableOp.mustBufferizeInPlace(opOperand, *this))
+        bufferizeInPlace(opOperand);
+}
+
 // Gather yielded tensors in `yieldedTensors` by querying all aliases. This is
 // to ensure that such information is available during bufferization time.
 // Alias information can no longer be queried once we have started modifying
 // the IR.
 void OneShotAnalysisState::gatherYieldedTensors(Operation *op) {
   op->walk([&](Operation *returnOp) {
     if (!isRegionReturnLike(returnOp) || !getOptions().isOpAllowed(returnOp))
       return WalkResult::advance();

mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-empty-tensor-elimination.mlir

-// RUN: mlir-opt %s -eliminate-empty-tensors -empty-tensor-to-alloc-tensor -one-shot-bufferize="bufferize-function-boundaries allow-return-allocs" -canonicalize -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -eliminate-empty-tensors -empty-tensor-to-alloc-tensor \
+// RUN:     -one-shot-bufferize="bufferize-function-boundaries allow-return-allocs" \
+// RUN:     -canonicalize -split-input-file | FileCheck %s
+
+// RUN: mlir-opt %s -test-interleaved-empty-tensor-elimination -cse \
+// RUN:     -canonicalize -loop-invariant-code-motion -split-input-file | \
+// RUN: FileCheck %s
 
 //      CHECK: func @buffer_forwarding_conflict(
 // CHECK-SAME:   %[[FUNC_ARG:[0-9a-zA-Z]*]]: memref<?xf32>
 // CHECK-SAME:   %[[sz:[0-9a-zA-Z]*]]: index
 func.func @buffer_forwarding_conflict(
   %t: tensor<?xf32> {bufferization.buffer_layout = affine_map<(d0) -> (d0)>, bufferization.writable = true},
   %sz: index)
     -> (tensor<?xf32>, tensor<?xf32>)

mlir/test/lib/Dialect/Bufferization/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRBufferizationTestPasses
+  TestEmptyTensorElimination.cpp
   TestTensorCopyInsertion.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   LINK_LIBS PUBLIC
   MLIRBufferizationDialect
   MLIRBufferizationTransforms
   MLIRIR
   MLIRPass
+  MLIRTensorDialect
 )

mlir/test/lib/Dialect/Bufferization/TestEmptyTensorElimination.cpp

+//===- TestEmptyTensorElimination.cpp ---------------------------*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
+#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
+#include "mlir/Dialect/Bufferization/Transforms/Transforms.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/Dominance.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+using namespace mlir::bufferization;
+
+namespace {
+struct TestInterleavedEmptyTensorEliminationPass
+    : public PassWrapper<TestInterleavedEmptyTensorEliminationPass,
+                         OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
+      TestInterleavedEmptyTensorEliminationPass)
+
+  TestInterleavedEmptyTensorEliminationPass() = default;
+  TestInterleavedEmptyTensorEliminationPass(
+      const TestInterleavedEmptyTensorEliminationPass &pass)
+      : PassWrapper(pass) {}
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<bufferization::BufferizationDialect>();
+  }
+  StringRef getArgument() const final {
+    return "test-interleaved-empty-tensor-elimination";
+  }
+  StringRef getDescription() const final {
+    return "Module pass to test Tensor Copy Insertion";
+  }
+
+  void runOnOperation() override {
+    OneShotBufferizationOptions options;
+    options.allowReturnAllocs = true;
+    options.bufferizeFunctionBoundaries = true;
+    OneShotAnalysisState state(getOperation(), options);
+
+    // Analyze ops.
+    if (failed(analyzeOp(getOperation(), state)))
+      return signalPassFailure();
+
+    // Eliminate tensor.empty ops.
+    if (failed(insertSliceAnchoredEmptyTensorEliminationStep(getOperation(),
+                                                             state)))
+      return signalPassFailure();
+
+    // Rewrite remaining tensor.empty ops to bufferization.alloc_tensor ops.
+    OpBuilder b(getOperation()->getContext());
+    DominanceInfo domInfo;
+    WalkResult status = getOperation()->walk([&](tensor::EmptyOp emptyOp) {
+      b.setInsertionPoint(emptyOp);
+      if (!emptyOp->getUsers().empty()) {
+        Value allocTensor = b.create<bufferization::AllocTensorOp>(
+            emptyOp.getLoc(), emptyOp.getType(), emptyOp.getDynamicSizes());
+        state.createAliasEntriesForNewOp(allocTensor.getDefiningOp());
+        SmallVector<OpOperand *> uses;
+        for (OpOperand &use : emptyOp.getResult().getUses())
+          uses.push_back(&use);
+        for (OpOperand *use : uses) {
+          use->set(allocTensor);
+          if (failed(state.analyzeSingleOp(use->getOwner(), domInfo)))
+            return WalkResult::interrupt();
+        }
+        emptyOp.replaceAllUsesWith(allocTensor);
+      }
+      emptyOp->erase();
+      return WalkResult::skip();
+    });
+    if (status.wasInterrupted())
+      return signalPassFailure();
+
+    // Insert tensor copies and bufferize.
+    if (failed(insertTensorCopies(getOperation(), state)))
+      return signalPassFailure();
+    if (failed(bufferizeOp(getOperation(), options, /*copyBeforeWrite=*/false)))
+      return signalPassFailure();
+  }
+};
+} // namespace
+
+namespace mlir::test {
+void registerTestInterleavedEmptyTensorEliminationPass() {
+  PassRegistration<TestInterleavedEmptyTensorEliminationPass>();
+}
+} // namespace mlir::test

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestDiagnosticsPass();
 void registerTestDialectConversionPasses();
 void registerTestDominancePass();
 void registerTestDynamicPipelinePass();
 void registerTestExpandMathPass();
 void registerTestFooAnalysisPass();
 void registerTestComposeSubView();
 void registerTestMultiBuffering();
+void registerTestInterleavedEmptyTensorEliminationPass();
 void registerTestIntRangeInference();
 void registerTestIRVisitorsPass();
 void registerTestGenericIRVisitorsPass();
 void registerTestGenericIRVisitorsInterruptPass();
 void registerTestInterfaces();
 void registerTestLastModifiedPass();
 void registerTestLinalgDecomposeOps();
 void registerTestLinalgElementwiseFusion();
#endif
   mlir::test::registerTestDataLayoutQuery();
   mlir::test::registerTestDeadCodeAnalysisPass();
   mlir::test::registerTestDominancePass();
   mlir::test::registerTestDynamicPipelinePass();
   mlir::test::registerTestExpandMathPass();
   mlir::test::registerTestFooAnalysisPass();
   mlir::test::registerTestComposeSubView();
   mlir::test::registerTestMultiBuffering();
+  mlir::test::registerTestInterleavedEmptyTensorEliminationPass();
   mlir::test::registerTestIntRangeInference();
   mlir::test::registerTestIRVisitorsPass();
   mlir::test::registerTestGenericIRVisitorsPass();
   mlir::test::registerTestInterfaces();
   mlir::test::registerTestLastModifiedPass();
   mlir::test::registerTestLinalgDecomposeOps();
   mlir::test::registerTestLinalgElementwiseFusion();
   mlir::test::registerTestLinalgGreedyFusion();

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

cc_library(
     srcs = glob(["lib/Dialect/Bufferization/*.cpp"]),
     defines = ["MLIR_CUDA_CONVERSIONS_ENABLED"],
     includes = ["lib/Dialect/Test"],
     deps = [
         "//mlir:BufferizationDialect",
         "//mlir:BufferizationTransforms",
         "//mlir:IR",
         "//mlir:Pass",
+        "//mlir:TensorDialect",
     ],
 )
 
 cc_library(
     name = "TestControlFlow",
     srcs = glob(["lib/Dialect/ControlFlow/*.cpp"]),
     includes = ["lib/Dialect/Test"],
     deps = [