[mlir][linalg][bufferize][NFC] Simplify InsertSliceOp bufferization

[springerm]

Summary

No need to keep track of equivalent extract_slice / insert_slice tensors during bufferization. Just emit a copy, it will fold away.

Note: The analysis still keeps track of equivalent tensors to make the correct inplace bufferization decisions.

[nicolasvasilache]

Nice simplification!
Indeed self-copy folds away.

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

 //===- LinalgInterfaceImpl.h - Linalg Impl. of BufferizableOpInterface ----===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_TENSOR_INTERFACE_IMPL_H
 #define MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_TENSOR_INTERFACE_IMPL_H
 
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
-
 namespace mlir {
 
 class DialectRegistry;
 
 namespace linalg {
 namespace comprehensive_bufferize {
 namespace tensor_ext {
 
-struct InplaceInsertSliceOpAnalysis : public PostAnalysisStep {
-  LogicalResult run(Operation *op, BufferizationState &state,
-                    BufferizationAliasInfo &aliasInfo,
-                    SmallVector<Operation *> &newOps) override;
-};
-
 void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
 
 } // namespace tensor_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_TENSOR_INTERFACE_IMPL_H

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

 namespace mlir {
 namespace linalg {
 namespace comprehensive_bufferize {
 namespace tensor_ext {
 
 using tensor::ExtractSliceOp;
 using tensor::InsertSliceOp;
 
-namespace {
-/// Extra bufferization state that is required for bufferization of tensor ops.
-struct TensorBufferizationState : public DialectBufferizationState {
-  /// InsertSliceOps that bufferize inplace and do not require a copy.
-  DenseSet<Operation *> insertSliceOpsWithoutCopy;
-};
-} // namespace
-
-static TensorBufferizationState &
-getTensorBufferizationState(BufferizationState &state) {
-  return state.getDialectState<TensorBufferizationState>(
-      tensor::TensorDialect::getDialectNamespace());
-}
-
 struct CastOpInterface
     : public BufferizableOpInterface::ExternalModel<CastOpInterface,
                                                     tensor::CastOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               BufferizationState &state) const {
     return false;
   }
 
if (!st || !sti)
     return false;
   if (!aliasInfo.areEquivalentBufferizedValues(st.source(), sti.dest()))
     return false;
   if (!sameOffsetsSizesAndStrides(st, sti, isEqualConstantIntOrValue))
     return false;
   return true;
 }
 
-/// Return true if the source of a `insertSliceOp` bufferizes to an
-/// equivalent ExtractSliceOp that bufferizes inplace.
-static bool isSourceEquivalentToAMatchingInplaceExtractSliceOp(
-    const BufferizationAliasInfo &aliasInfo, InsertSliceOp insertSliceOp) {
-  bool foundOp = false;
-  aliasInfo.applyOnEquivalenceClass(insertSliceOp.source(), [&](Value value) {
-    auto extractSliceOp = value.getDefiningOp<ExtractSliceOp>();
-    if (extractSliceOp &&
-        areEquivalentExtractSliceOps(aliasInfo, extractSliceOp,
-                                     insertSliceOp) &&
-        aliasInfo.isInPlace(extractSliceOp->getResult(0))) {
-      foundOp = true;
-    }
-  });
-  return foundOp;
-}
-
 /// Return true if `value` is originating from an ExtractSliceOp that matches
 /// the given InsertSliceOp.
 static bool hasMatchingExtractSliceOp(const BufferizationAliasInfo &aliasInfo,
                                       BufferizationState &state, Value value,
                                       InsertSliceOp insertOp) {
   auto condition = [&](Value val) {
     if (auto extractOp = val.getDefiningOp<ExtractSliceOp>())
       if (areEquivalentExtractSliceOps(aliasInfo, extractOp, insertOp))
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     // insert_slice ops arise from tiling and bufferizing them out-of-place is
     // generally a deal breaker. When used with loops, this ends up cloning the
     // whole tensor on every single iteration and is a symptom of a
     // catastrophically bad scheduling decision.
     // TODO: be very loud about it or even consider failing the pass.
     auto insertSliceOp = cast<tensor::InsertSliceOp>(op);
     Location loc = insertSliceOp.getLoc();
-    TensorBufferizationState &tensorState = getTensorBufferizationState(state);
 
     // When bufferizing out-of-place, `getResultBuffer` allocates.
     Value dstMemref =
         state.getResultBuffer(rewriter, insertSliceOp->getResult(0));
     if (!dstMemref)
       return failure();
 
-    bool needCopy =
-        !tensorState.insertSliceOpsWithoutCopy.contains(insertSliceOp);
-    if (needCopy) {
-      // Take a subview of the dst.
-      auto dstMemrefType = dstMemref.getType().cast<MemRefType>();
-      auto subviewMemRefType =
-          memref::SubViewOp::inferRankReducedResultType(
-              insertSliceOp.getSourceType().getRank(), dstMemrefType,
-              insertSliceOp.getMixedOffsets(), insertSliceOp.getMixedSizes(),
-              insertSliceOp.getMixedStrides())
-              .cast<MemRefType>();
-      Value subView = rewriter.create<memref::SubViewOp>(
-          loc, subviewMemRefType, dstMemref, insertSliceOp.getMixedOffsets(),
-          insertSliceOp.getMixedSizes(), insertSliceOp.getMixedStrides());
-      // Copy tensor.
-      Value srcMemref = state.lookupBuffer(rewriter, insertSliceOp.source());
-      state.createMemCpy(rewriter, insertSliceOp.getLoc(), srcMemref, subView);
-    }
+
+    // Copy tensor. If this tensor.insert_slice has a matching
+    // tensor.extract_slice, the copy operation will eventually fold away.
 
     state.replaceOp(rewriter, op, dstMemref);
     return success();
   }
 };
 
 } // namespace tensor_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
-LogicalResult mlir::linalg::comprehensive_bufferize::tensor_ext::
-    InplaceInsertSliceOpAnalysis::run(Operation *op, BufferizationState &state,
-                                      BufferizationAliasInfo &aliasInfo,
-                                      SmallVector<Operation *> &newOps) {
-  auto &tensorState = getTensorBufferizationState(state);
-  op->walk([&](InsertSliceOp insertSliceOp) {
-    // A copy of the source buffer is needed if either:
-    //   - The producer of `source` is not inplace. This is the case where a
-    //     slice is computed out of place into the inplace full tensor.
-    //   - The result is not inplace. This is the case where the whole tensor is
-    //     cloned and the clone needs to be updated.
-    if (isSourceEquivalentToAMatchingInplaceExtractSliceOp(aliasInfo,
-                                                           insertSliceOp) &&
-        state.isInPlace(insertSliceOp->getResult(0)))
-      tensorState.insertSliceOpsWithoutCopy.insert(insertSliceOp);
-  });
-  return success();
-}
-
 void mlir::linalg::comprehensive_bufferize::tensor_ext::
     registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry) {
   registry.addOpInterface<tensor::CastOp, tensor_ext::CastOpInterface>();
   registry.addOpInterface<tensor::DimOp, tensor_ext::DimOpInterface>();
   registry.addOpInterface<tensor::ExtractSliceOp,
                           tensor_ext::ExtractSliceOpInterface>();
   registry.addOpInterface<tensor::ExtractOp, tensor_ext::ExtractOpInterface>();
   registry.addOpInterface<tensor::InsertOp, tensor_ext::InsertOpInterface>();
   registry.addOpInterface<tensor::InsertSliceOp,
                           tensor_ext::InsertSliceOpInterface>();
 }

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

void LinalgComprehensiveModuleBufferize::runOnOperation() {
   options->allowUnknownOps = allowUnknownOps;
   options->analysisFuzzerSeed = analysisFuzzerSeed;
   options->testAnalysisOnly = testAnalysisOnly;
   options->printConflicts = printConflicts;
 
   // Enable InitTensorOp elimination.
   options->addPostAnalysisStep<
       linalg_ext::InsertSliceAnchoredInitTensorEliminationStep>();
-  // TODO: Find a way to enable this step automatically when bufferizing tensor
-  // dialect ops.
-  options->addPostAnalysisStep<tensor_ext::InplaceInsertSliceOpAnalysis>();
   if (!allowReturnMemref)
     options->addPostAnalysisStep<scf_ext::AssertDestinationPassingStyle>();
 
   ModuleOp moduleOp = getOperation();
   applyEnablingTransformations(moduleOp);
 
   if (failed(runComprehensiveBufferize(moduleOp, std::move(options)))) {
     signalPassFailure();

mlir/test/lib/Dialect/Linalg/TestComprehensiveBufferize.cpp

 } // namespace
 
 void TestComprehensiveFunctionBufferize::runOnFunction() {
   auto options = std::make_unique<BufferizationOptions>();
 
   // Enable InitTensorOp elimination.
   options->addPostAnalysisStep<
       linalg_ext::InsertSliceAnchoredInitTensorEliminationStep>();
-  // TODO: Find a way to enable this step automatically when bufferizing
-  // tensor dialect ops.
-  options->addPostAnalysisStep<tensor_ext::InplaceInsertSliceOpAnalysis>();
   if (!allowReturnMemref)
     options->addPostAnalysisStep<scf_ext::AssertDestinationPassingStyle>();
 
   options->allowReturnMemref = allowReturnMemref;
   options->allowUnknownOps = allowUnknownOps;
   options->testAnalysisOnly = testAnalysisOnly;
   options->analysisFuzzerSeed = analysisFuzzerSeed;