[mlir][linalg][bufferize] Generalize InitTensorOp elimination

[springerm]

Summary

This allows for external users of Comprehensive Bufferize to specify their own InitTensorOp elimination procedures.

[nicolasvasilache]

Nice!

[springerm]

rebase

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

 LogicalResult
 bufferizeOp(Operation *op, BlockAndValueMapping &bvm,
             BufferizationAliasInfo &aliasInfo,
             AllocationCallbacks allocationFns,
             DenseMap<FuncOp, FunctionType> *bufferizedFunctionTypes = nullptr);
 
 /// Register external models implemented for the `BufferizableOpInterface`.
 void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
+
+/// Try to eliminate InitTensorOps inside `funcOp`.
+///
+/// * `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 initTensorElimination(
+    FuncOp funcOp, BufferizationAliasInfo &aliasInfo, DominanceInfo &domInfo,
+    std::function<bool(OpOperand &)> anchorMatchFunc,
+    std::function<Value(OpBuilder &, Location, OpOperand &)> rewriteFunc,
+    bool skipAnalysis = false);
+
+/// Try to eliminate InitTensorOps inside funcOp that are anchored on an
+/// InsertSliceOp, i.e., if it is eventually inserted into another tensor
+/// (and some other conditions are met).
+LogicalResult eliminateInsertSliceAnchoredInitTensorOps(
+    FuncOp funcOp, BufferizationAliasInfo &aliasInfo, DominanceInfo &domInfo);
+
 } // namespace linalg
 } // namespace mlir
 
 #endif // define MLIR_DIALECT_LINALG_TRANSFORMS_COMPREHENSIVE_BUFFERIZE_H

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

for (FuncOp funcOp : orderedFuncOps) {
 
     // Finally set the funcOp type to update the arguments.
     auto newFuncType = FunctionType::get(moduleOp.getContext(), argumentTypes,
                                          funcOp.getType().getResults());
     funcOp.setType(newFuncType);
   }
 }
 
-/// Try to eliminate InitTensorOps inside funcOp. An InitTensorOp can be
-/// eliminated if it is eventually inserted into another tensor (and some other
-/// conditions are met).
-///
-/// E.g.:
-/// %0 = linalg.init_tensor
-/// %1 = linalg.fill(%cst, %0) {inplace = [true]}
-/// %2 = tensor.insert_slice %1 into %t[10][20][1]
-///
-/// InitTensorOp elimination will try to fill %t inplace instead of filling a
-/// new allocation %0 and inserting it into %t. This is done by replacing the
-/// InitTensorOp with:
-///
-/// %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.
-static LogicalResult runInitTensorElimination(FuncOp funcOp,
-                                              BufferizationAliasInfo &aliasInfo,
-                                              DominanceInfo &domInfo) {
+/// Try to eliminate InitTensorOps inside funcOp. 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::initTensorElimination(
+    FuncOp funcOp, BufferizationAliasInfo &aliasInfo, DominanceInfo &domInfo,
+    std::function<bool(OpOperand &)> anchorMatchFunc,
+    std::function<Value(OpBuilder &, Location, OpOperand &)> rewriteFunc,
+    bool skipAnalysis) {
   OpBuilder b(funcOp->getContext());
 
-  WalkResult status = funcOp->walk([&](tensor::InsertSliceOp insertOp) {
-    // Only inplace bufferized InsertSliceOps are eligible.
-    if (getInPlace(insertOp->getOpResult(0)) != InPlaceSpec::True)
-      return WalkResult::skip();
+  WalkResult status = funcOp->walk([&](Operation *op) {
+    for (OpOperand &operand : op->getOpOperands()) {
+      // Is this a matching OpOperand?
+      if (!anchorMatchFunc(operand))
+        continue;
 
       SetVector<Value> maybeInitTensor =
-        findValueInReverseUseDefChain(insertOp.source(), [](Value val) {
+          findValueInReverseUseDefChain(operand.get(), [](Value val) {
             // Continue traversal until this function returns true.
             OpResult opResult = val.dyn_cast<OpResult>();
             if (!opResult)
               return true;
             if (getInPlace(opResult) != InPlaceSpec::True)
               return true;
-          // Only equivalent tensors are supported at the moment. E.g., when
-          // taking a tensor.extract_slice of an init_tensor, we can currently
-          // not eliminate the init_tensor.
-          SmallVector<OpOperand *> opOperands = getAliasingOpOperand(opResult);
+            // Only equivalent tensors are supported at the moment.
+            // TODO: Support cases such as extract_slice(init_tensor).
+            SmallVector<OpOperand *> opOperands =
+                getAliasingOpOperand(opResult);
             if (!llvm::all_of(opOperands, [](OpOperand *operand) {
                   return bufferRelation(*operand) == BufferRelation::Equivalent;
                 }))
               return true;
             return false;
           });
-    // Replace only if the InsertSliceOp source originates from exactly one
+
+      // 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();
 
+      // Create a replacement for the InitTensorOp.
       b.setInsertionPoint(initTensor.getDefiningOp());
-    auto extractOp = b.create<tensor::ExtractSliceOp>(
-        initTensor.getLoc(), insertOp.dest(), insertOp.getMixedOffsets(),
-        insertOp.getMixedSizes(), insertOp.getMixedStrides());
+      Value replacement = rewriteFunc(b, 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(extractOp.result());
-    aliasInfo.createAliasInfoEntry(extractOp.result());
+      initTensor.replaceAllUsesWith(replacement);
+      aliasInfo.createAliasInfoEntry(replacement);
 
-    // Run analysis on the ExtractSliceOp.
-    if (failed(bufferizableInPlaceAnalysisAliasOnlyOp(
-            extractOp->getOpOperand(0), aliasInfo, domInfo)))
+      // Run analysis on the newly created op.
+      if (auto opResult = replacement.dyn_cast<OpResult>()) {
+        if (!skipAnalysis) {
+          SmallVector<Operation *> ops(1, replacement.getDefiningOp());
+          if (failed(inPlaceAnalysis(ops, aliasInfo, domInfo)))
             return WalkResult::interrupt();
+        }
+      }
+    }
 
     // Advance to the next operation.
     return WalkResult::advance();
   });
 
   return failure(status.wasInterrupted());
 }
 
+/// Try to eliminate InitTensorOps inside funcOp. An InitTensorOp can be
+/// eliminated if it is eventually inserted into another tensor (and some other
+/// conditions are met).
+///
+/// E.g.:
+/// %0 = linalg.init_tensor
+/// %1 = linalg.fill(%cst, %0) {inplace = [true]}
+/// %2 = tensor.insert_slice %1 into %t[10][20][1]
+///
+/// InitTensorOp elimination will try to fill %t inplace instead of filling a
+/// new allocation %0 and inserting it into %t. This is done by replacing the
+/// InitTensorOp with:
+///
+/// %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::eliminateInsertSliceAnchoredInitTensorOps(
+    FuncOp funcOp, BufferizationAliasInfo &aliasInfo, DominanceInfo &domInfo) {
+  return initTensorElimination(
+      funcOp, aliasInfo, domInfo,
+      [](OpOperand &operand) {
+        auto insertSliceOp = dyn_cast<InsertSliceOp>(operand.getOwner());
+        if (!insertSliceOp)
+          return false;
+        // Only inplace bufferized InsertSliceOps are eligible.
+        if (getInPlace(insertSliceOp->getOpResult(0)) != InPlaceSpec::True)
+          return false;
+        return &operand == &insertSliceOp->getOpOperand(0) /*source*/;
+      },
+      [](OpBuilder &b, Location loc, OpOperand &operand) {
+        auto insertSliceOp = cast<InsertSliceOp>(operand.getOwner());
+        auto extractOp = b.create<tensor::ExtractSliceOp>(
+            loc, insertSliceOp.dest(), insertSliceOp.getMixedOffsets(),
+            insertSliceOp.getMixedSizes(), insertSliceOp.getMixedStrides());
+        return extractOp.result();
+      });
+}
+
 void LinalgComprehensiveModuleBufferize::runOnOperation() {
   if (!allocationFns) {
     // The allocation functions to use needs to be set here. The flag for the
     // pass and flag for the use of alloca map to LLVM command line
     // options. These being static global objects have no set order in which
     // they are defined. So ideally this should be in the constructor, but the
     // constructor might be called before the flag is initialized using the
     // command line option. So this is set up at the start of the pass.
for (FuncOp funcOp : orderedFuncOps) {
     if (failed(inPlaceAnalysisFuncOpBody(funcOp, aliasInfo, domInfo,
                                          analysisFuzzerSeed))) {
       signalPassFailure();
       return;
     }
 
     // Try to eliminate InitTensorOps to avoid new allocations during the
     // bufferization phase.
-    if (failed(runInitTensorElimination(funcOp, aliasInfo, domInfo))) {
+    if (failed(eliminateInsertSliceAnchoredInitTensorOps(funcOp, aliasInfo,
+                                                         domInfo))) {
       signalPassFailure();
       return;
     }
 
     // Bufferization phase.
     if (!testAnalysisOnly) {
       BlockAndValueMapping tensorToBufferMap;
       if (failed(bufferizeFuncOpInternals(funcOp, tensorToBufferMap, aliasInfo,