[mlir][linalg][bufferize][NFC] Make init_tensor elimination a separate pre-processing pass

[springerm]

Summary

This commit decouples init_tensor elimination from the rest of the bufferization.

Depends On D124849

[springerm]

update

mlir/include/mlir/Dialect/Linalg/Passes.h

 /// FuncOp within a module. It provides analyses and bufferization across
 /// function boundaries. Within a function boundary, the analysis is performed
 /// on SSA use-def chains starting from function operands that are annotated
 /// with the 'inplaceable' attribute.
 std::unique_ptr<Pass> createLinalgComprehensiveModuleBufferizePass();
 std::unique_ptr<Pass> createLinalgComprehensiveModuleBufferizePass(
     const bufferization::OneShotBufferizationOptions &options);
 
+/// Create a pass that tries to eliminate init_tensor ops that are anchored on
+/// insert_slice ops.
+std::unique_ptr<Pass> createLinalgInitTensorEliminationPass();
+
 /// Create a pass to convert Linalg operations which work on tensors to use
 /// buffers instead.
 std::unique_ptr<OperationPass<func::FuncOp>> createLinalgBufferizePass();
 
 /// Create a pass to convert named Linalg operations to Linalg generic
 /// operations.
 std::unique_ptr<OperationPass<func::FuncOp>> createLinalgGeneralizationPass();
 

mlir/include/mlir/Dialect/Linalg/Passes.td

Option<"useAlloca", "use-alloca", "bool",
            /*default=*/"false",
            "Use stack allocations for memrefs (for testing purposes only)">,
     Option<"fullyDynamicLayoutMaps", "fully-dynamic-layout-maps", "bool",
            /*default=*/"true",
            "Generate MemRef types with dynamic offset+strides by default.">,
     Option<"analysisFuzzerSeed", "analysis-fuzzer-seed", "unsigned",
            /*default=*/"0",
            "Analyze ops in random order with a given seed (fuzzer)">,
-    Option<"initTensorElimination", "init-tensor-elimination", "bool",
-            /*default=*/"false",
-           "(Experimental) Try to eliminate init_tensor operations that are "
-           "anchored at an insert_slice op">,
     Option<"createDeallocs", "create-deallocs", "bool", /*default=*/"true",
            "Specify if buffers should be deallocated. For compatibility with "
            "core bufferization passes.">,
   ];
   let constructor = "mlir::createLinalgComprehensiveModuleBufferizePass()";
 }
 
+def LinalgInitTensorElimination : Pass<"linalg-eliminate-init-tensors"> {
+  let summary = "Try to eliminate all init_tensor ops.";
+  let description = [{
+    This pass tries to eliminate all insert_slice op-anchored init_tensor ops.
+    I.e., when a value that is aliasing with an init_tensor op is inserted into
+    another tensor, this pass tries to rewrite the IR in such a way that the
+    destination tensor of the insert_slice op is used directly instead of the
+    init_tensor result.
+  }];
+  let constructor = "mlir::createLinalgInitTensorEliminationPass()";
+}
+
 def LinalgFoldUnitExtentDims : Pass<"linalg-fold-unit-extent-dims", ""> {
   let summary = "Remove unit-extent dimension in Linalg ops on tensors";
   let constructor = "mlir::createLinalgFoldUnitExtentDimsPass()";
   let options = [
     Option<"foldOneTripLoopsOnly", "fold-one-trip-loops-only", "bool",
             /*default=*/"false",
            "Only folds the one-trip loops from Linalg ops on tensors "
            "(for testing purposes only)">

mlir/include/mlir/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.h

 /// Try to eliminate InitTensorOps inside `op`.
 ///
 /// * `rewriteFunc` generates the replacement for the InitTensorOp.
 /// * Only InitTensorOps 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
 ///   InitTensorOp.
-/// * The result of `rewriteFunc` must usually be analyzed for inplacability.
-///   This analysis can be skipped with `skipAnalysis`.
-LogicalResult
-eliminateInitTensors(Operation *op, bufferization::AnalysisState &state,
-                     bufferization::BufferizationAliasInfo &aliasInfo,
-                     AnchorMatchFn anchorMatchFunc, RewriteFn rewriteFunc,
-                     SmallVector<Operation *> &newOps);
+LogicalResult eliminateInitTensors(RewriterBase &rewriter, Operation *op,
+                                   bufferization::AnalysisState &state,
+                                   AnchorMatchFn anchorMatchFunc,
+                                   RewriteFn rewriteFunc);
 
 /// Try to eliminate InitTensorOps 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 insertSliceAnchoredInitTensorEliminationStep(
-    Operation *op, bufferization::AnalysisState &state,
-    bufferization::BufferizationAliasInfo &aliasInfo,
-    SmallVector<Operation *> &newOps);
+    RewriterBase &rewriter, Operation *op, bufferization::AnalysisState &state);
 
 void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
 
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_BUFFERIZABLEOPINTERFACEIMPL_H

mlir/lib/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.cpp

findValidInsertionPoint(Operation *initTensorOp,
   return nullptr;
 }
 
 /// Try to eliminate InitTensorOps inside `op`. An InitTensorOp is replaced
 /// with the 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 InitTensorOp.
-LogicalResult mlir::linalg::eliminateInitTensors(
-    Operation *op, AnalysisState &state, BufferizationAliasInfo &aliasInfo,
-    AnchorMatchFn anchorMatchFunc, RewriteFn rewriteFunc,
-    SmallVector<Operation *> &newOps) {
-  OpBuilder b(op->getContext());
+LogicalResult mlir::linalg::eliminateInitTensors(RewriterBase &rewriter,
+                                                 Operation *op,
+                                                 AnalysisState &state,
+                                                 AnchorMatchFn anchorMatchFunc,
+                                                 RewriteFn rewriteFunc) {
+  OpBuilder::InsertionGuard g(rewriter);
 
   WalkResult status = op->walk([&](Operation *op) {
     for (OpOperand &operand : op->getOpOperands()) {
       // Skip operands that do not bufferize inplace.
-      if (!aliasInfo.isInPlace(operand))
+      if (!state.isInPlace(operand))
         continue;
       // All values that are needed to create the replacement op.
       SmallVector<Value> neededValues;
       // Is this a matching OpOperand?
       if (!anchorMatchFunc(operand, neededValues))
         continue;
       SetVector<Value> maybeInitTensor =
           state.findValueInReverseUseDefChain(operand.get(), [&](Value val) {
             // Continue traversal until this function returns true.
             OpResult opResult = val.dyn_cast<OpResult>();
             if (!opResult)
               return true;
             SmallVector<OpOperand *> opOperands =
                 state.getAliasingOpOperand(opResult);
             if (!llvm::all_of(opOperands, [&](OpOperand *operand) {
-                  return aliasInfo.isInPlace(*operand);
+                  return state.isInPlace(*operand);
                 }))
               return true;
             // Only equivalent tensors are supported at the moment.
             // TODO: Support cases such as extract_slice(init_tensor)
             return !llvm::all_of(opOperands, [&](OpOperand *operand) {
-              return aliasInfo.areEquivalentBufferizedValues(operand->get(),
-                                                             opResult);
+              return state.areEquivalentBufferizedValues(operand->get(),
             });
           });
 
       // Replace only if the reverse use-def chain ends at exactly one
       // InitTensorOp.
       if (maybeInitTensor.size() != 1 ||
           !maybeInitTensor.front().getDefiningOp<InitTensorOp>())
         return WalkResult::skip();
       Value initTensor = maybeInitTensor.front();
 
       // Find a suitable insertion point.
       Operation *insertionPoint =
           findValidInsertionPoint(initTensor.getDefiningOp(), neededValues);
       if (!insertionPoint)
         continue;
 
       // Create a replacement for the InitTensorOp.
-      b.setInsertionPoint(insertionPoint);
-      Value replacement = rewriteFunc(b, initTensor.getLoc(), operand);
+      rewriter.setInsertionPoint(insertionPoint);
+      Value replacement = rewriteFunc(rewriter, initTensor.getLoc(), operand);
       if (!replacement)
         continue;
 
-      // Uses of the InitTensorOp are replaced here, but the op is not deleted.
-      // InitTensorOps without uses are ignored by the bufferization.
-      initTensor.replaceAllUsesWith(replacement);
-      aliasInfo.createAliasInfoEntry(replacement);
-      aliasInfo.unionAliasSets(initTensor, replacement);
-      aliasInfo.unionEquivalenceClasses(initTensor, replacement);
-
-      // Register replacement ops.
-      if (Operation *newOp = replacement.getDefiningOp())
-        newOps.push_back(newOp);
+      // Replace the InitTensorOp.
+      rewriter.replaceOp(initTensor.getDefiningOp(), replacement);
     }
 
     // Advance to the next operation.
     return WalkResult::advance();
   });
 
   return failure(status.wasInterrupted());
 }

 ///
 /// %0 = tensor.extract_slice %t[10][20][1]
 ///
 /// The analysis looks for matching ExtractSliceOp/InsertSliceOp pairs and lets
 /// those bufferize inplace in the absence of other conflicts.
 ///
 /// Starting from an InsertSliceOp, an InitTensorOp at the end of the insert
 /// source's reverse use-def chain is eliminated if:
-/// * The InsertSliceOp was decided to bufferize inplace.
 /// * On the reverse use-def chain path from the InsertSliceOp to the
 ///   InitTensorOp, 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 InitTensorOp.
-///
-/// Note that the newly inserted ExtractSliceOp may have to bufferize
-/// out-of-place due to RaW conflicts.
 LogicalResult mlir::linalg::insertSliceAnchoredInitTensorEliminationStep(
-    Operation *op, AnalysisState &state, BufferizationAliasInfo &aliasInfo,
-    SmallVector<Operation *> &newOps) {
+    RewriterBase &rewriter, Operation *op, AnalysisState &state) {
   return eliminateInitTensors(
-      op, state, aliasInfo,
+      rewriter, op, state,
       /*anchorMatchFunc=*/
       [&](OpOperand &operand, SmallVector<Value> &neededValues) {
         auto insertSliceOp =
             dyn_cast<tensor::InsertSliceOp>(operand.getOwner());
         if (!insertSliceOp)
           return false;
-        // Only inplace bufferized InsertSliceOps are eligible.
-        if (!aliasInfo.isInPlace(insertSliceOp->getOpOperand(1) /*dest*/))
-          return false;
         if (&operand != &insertSliceOp->getOpOperand(0) /*source*/)
           return false;
 
         // Collect all values that are needed to construct the replacement op.
         neededValues.append(insertSliceOp.offsets().begin(),
                             insertSliceOp.offsets().end());
         neededValues.append(insertSliceOp.sizes().begin(),
                             insertSliceOp.sizes().end());
return eliminateInitTensors(
             });
         auto t = tensor::ExtractSliceOp::inferRankReducedResultType(
             insertOp.getSourceType().getRank(),
             insertOp.dest().getType().cast<RankedTensorType>(), mixedOffsets,
             mixedSizes, mixedStrides);
         auto extractOp = b.create<tensor::ExtractSliceOp>(
             loc, t, insertOp.dest(), mixedOffsets, mixedSizes, mixedStrides);
         return extractOp.result();
-      },
-      newOps);
+      });
 }
 
 void mlir::linalg::registerBufferizableOpInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addExtension(+[](MLIRContext *ctx, linalg::LinalgDialect *dialect) {
     linalg::InitTensorOp::attachInterface<InitTensorOpInterface>(*ctx);
 
     // Register all Linalg structured ops. `LinalgOp` is an interface and it is
     // not possible to attach an external interface to an existing interface.
     // Therefore, attach the `BufferizableOpInterface` to all ops one-by-one.
     LinalgOpInterfaceHelper<
 #define GET_OP_LIST
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
         >::registerOpInterface(ctx);
   });
 }

mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferizePass.cpp

void getDependentDialects(DialectRegistry &registry) const override {
     func_ext::registerBufferizableOpInterfaceExternalModels(registry);
     tensor::registerBufferizableOpInterfaceExternalModels(registry);
     vector::registerBufferizableOpInterfaceExternalModels(registry);
   }
 
 private:
   llvm::Optional<OneShotBufferizationOptions> options;
 };
+
+struct LinalgInitTensorElimination
+    : public LinalgInitTensorEliminationBase<LinalgInitTensorElimination> {
+  LinalgInitTensorElimination() = default;
+
+  void runOnOperation() override;
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<linalg::LinalgDialect, tensor::TensorDialect>();
+  }
+};
 } // namespace
 
 static void applyEnablingTransformations(ModuleOp moduleOp) {
   RewritePatternSet patterns(moduleOp.getContext());
   patterns.add<GeneralizePadOpPattern>(moduleOp.getContext());
   (void)applyPatternsAndFoldGreedily(moduleOp, std::move(patterns));
 }
 
if (!options) {
     opt.allowUnknownOps = allowUnknownOps;
     opt.analysisFuzzerSeed = analysisFuzzerSeed;
     opt.createDeallocs = createDeallocs;
     opt.fullyDynamicLayoutMaps = fullyDynamicLayoutMaps;
     opt.printConflicts = printConflicts;
     opt.testAnalysisOnly = testAnalysisOnly;
     opt.alwaysAliasingWithDest = alwaysAliasingWithDest;
     opt.bufferizeFunctionBoundaries = true;
-    if (initTensorElimination) {
-      opt.addPostAnalysisStep(insertSliceAnchoredInitTensorEliminationStep);
-    }
   } else {
     opt = *options;
   }
 
   ModuleOp moduleOp = getOperation();
   applyEnablingTransformations(moduleOp);
 
   if (failed(runOneShotModuleBufferize(moduleOp, opt))) {
     signalPassFailure();
     return;
   }
 
   if (opt.testAnalysisOnly)
     return;
 
   OpPassManager cleanupPipeline("builtin.module");
   cleanupPipeline.addPass(createCanonicalizerPass());
   cleanupPipeline.addPass(createCSEPass());
   cleanupPipeline.addPass(createLoopInvariantCodeMotionPass());
   (void)runPipeline(cleanupPipeline, moduleOp);
 }
 
+void LinalgInitTensorElimination::runOnOperation() {
+  Operation *op = getOperation();
+  OneShotBufferizationOptions options;
+  OneShotAnalysisState state(op, options);
+  if (failed(analyzeOp(op, state))) {
+    signalPassFailure();
+    return;
+  }
+
+  IRRewriter rewriter(op->getContext());
+  if (failed(insertSliceAnchoredInitTensorEliminationStep(rewriter, op, state)))
+    signalPassFailure();
+}
+
 std::unique_ptr<Pass> mlir::createLinalgComprehensiveModuleBufferizePass() {
   return std::make_unique<LinalgComprehensiveModuleBufferize>();
 }
 
 std::unique_ptr<Pass> mlir::createLinalgComprehensiveModuleBufferizePass(
     const OneShotBufferizationOptions &options) {
   return std::make_unique<LinalgComprehensiveModuleBufferize>(options);
 }
+
+std::unique_ptr<Pass> mlir::createLinalgInitTensorEliminationPass() {
+  return std::make_unique<LinalgInitTensorElimination>();
+}

mlir/test/Dialect/Linalg/one-shot-bufferize-analysis-init-tensor-elimination.mlir

-// RUN: mlir-opt %s -linalg-comprehensive-module-bufferize="test-analysis-only allow-return-allocs init-tensor-elimination" -split-input-file | FileCheck %s
-
-// -----
+// RUN: mlir-opt %s -linalg-eliminate-init-tensors -one-shot-bufferize="bufferize-function-boundaries test-analysis-only allow-return-allocs" -split-input-file | FileCheck %s
 
 //===----------------------------------------------------------------------===//
 // InitTensorOp elimination
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: func @buffer_forwarding_conflict
 func.func @buffer_forwarding_conflict(%arg0: tensor<?xf32> {bufferization.writable = true}, %arg1: index) -> (tensor<?xf32>, tensor<?xf32>) {
   %cst = arith.constant 0.000000e+00 : f32

mlir/test/Dialect/Linalg/one-shot-bufferize-init-tensor-elimination.mlir

-// RUN: mlir-opt %s -linalg-comprehensive-module-bufferize="allow-return-allocs init-tensor-elimination" -canonicalize -split-input-file | FileCheck %s
-
-// -----
+// RUN: mlir-opt %s -linalg-eliminate-init-tensors -one-shot-bufferize="bufferize-function-boundaries allow-return-allocs" -canonicalize -split-input-file | 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>)