[mlir][tensor] Add pattern to extract from insert_slice destination

[antiagainst]

Summary

When we have a producer tensor.insert_slice op and a consumer
tensor.extract_slice op and their slices are disjoint, we can
update the tensor.extract_slice op to use the tensor.insert_slice
op's destination tensor. This helps to break the chain so that
it can further enable optimizations.

[antiagainst]

Move the helper function to ViewLikeInterface

[mravishankar]

mlir/lib/Dialect/Tensor/Transforms/ExtractFromInsertSliceDest.cpp

19 ↗

(On Diff #460496)

This is a very specific pattern, and post bufferization this shouldnt matter. It seems very difficult to reason about this in non-static cases (and here this is all for static).

[antiagainst]

mlir/lib/Dialect/Tensor/Transforms/ExtractFromInsertSliceDest.cpp

19 ↗

(On Diff #460496)

Some dynamic dimensions are fine, but right we need to have at least one static disjoint dimensions to be sure.

[antiagainst]

Rebase

[mravishankar]

Some dynamic dimensions are fine, but right we need to have at least one static disjoint dimensions to be sure.

I am just saying this seems like something that is addressing a symptom and not the root cause. Maybe root causing it might help avoid such specific patterns.

[antiagainst]

In D133872#3803844, @mravishankar wrote:

Some dynamic dimensions are fine, but right we need to have at least one static disjoint dimensions to be sure.

I am just saying this seems like something that is addressing a symptom and not the root cause. Maybe root causing it might help avoid such specific patterns.

This kind of IR is generated by tiling and unrolling at the tensor level. Tiling creates the structure of for (...) { extract_slice, compute, insert_slice }. After unrolling the loop there we have consecutive blocks of extract_slice0, compute0, insert_slice0, extract_slice1, .... extract_slice1 would extract from insert_slice0. This patterns breaks the dependence chain here to make future transformations easier (not needing to look through the whole extract/insert op chain). It's a simple pattern that we can explicitly control application. Not doing this would mean we either need to bake unrolling inside tiling, or having later transformations being able to look through the extract/insert chain. To me this way seems cleaner than those.

[antiagainst]

Not immediately needed anymore; so taking back for now.

[antiagainst]

Move utils to Affine/Utils to avoid introducing ArithUtils dependencies to ViewLikeInterface

[antiagainst]

Rebase

[antiagainst]

Rebase

mlir/include/mlir/Dialect/Affine/ViewLikeInterfaceUtils.h

 mergeOffsetsSizesAndStrides(OpBuilder &builder, Location loc,
                             OffsetSizeAndStrideOpInterface producer,
                             OffsetSizeAndStrideOpInterface consumer,
                             const llvm::SmallBitVector &droppedProducerDims,
                             SmallVector<OpFoldResult> &combinedOffsets,
                             SmallVector<OpFoldResult> &combinedSizes,
                             SmallVector<OpFoldResult> &combinedStrides);
 
+/// Returns true if the given two n-D ranges can be proven as disjoint.
+/// Returns false otherwise.
+///
+/// This function assumes all input arrays to have the same size.
+bool areDisjointRanges(ArrayRef<OpFoldResult> aOffsets,
+                       ArrayRef<OpFoldResult> aSizes,
+                       ArrayRef<OpFoldResult> bOffsets,
+                       ArrayRef<OpFoldResult> bSizes);
+
+/// Returns true if the given two slices can be proven as disjoint. Returns
+/// false otherwise.
+///
+/// This function assumes the two slices have the same rank.
+bool areDisjointSlices(OffsetSizeAndStrideOpInterface aSlice,
+                       OffsetSizeAndStrideOpInterface bSlice);
+
 } // namespace mlir
 
 #endif // MLIR_DIALECT_AFFINE_VIEWLIKEINTERFACEUTILS_H

mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h

 
 /// Collects patterns to merge consecutive tensor.insert_slice/extract_slice
 /// into one. These patterns are in in this separate entry point because the
 /// bufferization is sensitive over IR structure, particularly those
 /// tensor.extract_slice and tensor.insert_slice ops for creating the slices.
 void populateMergeConsecutiveInsertExtractSlicePatterns(
     RewritePatternSet &patterns);
 
+/// Collects patterns to update tensor.extract_slice to extract from the
+/// destination tensor of its producer tensor.insert_slice op.
+void populateExtractFromInsertSliceDestOpPatterns(RewritePatternSet &patterns);
+
 } // namespace tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_TENSOR_TRANSFORMS_TRANSFORMS_H

mlir/lib/Dialect/Affine/Utils/ViewLikeInterfaceUtils.cpp

LogicalResult mlir::mergeOffsetsSizesAndStrides(
   SmallVector<OpFoldResult> producerOffsets = producer.getMixedOffsets();
   SmallVector<OpFoldResult> producerSizes = producer.getMixedSizes();
   SmallVector<OpFoldResult> producerStrides = producer.getMixedStrides();
   return mergeOffsetsSizesAndStrides(
       builder, loc, producerOffsets, producerSizes, producerStrides,
       droppedProducerDims, consumerOffsets, consumerSizes, consumerStrides,
       combinedOffsets, combinedSizes, combinedStrides);
 }
+
+bool mlir::areDisjointRanges(ArrayRef<OpFoldResult> aOffsets,
+                             ArrayRef<OpFoldResult> aSizes,
+                             ArrayRef<OpFoldResult> bOffsets,
+                             ArrayRef<OpFoldResult> bSizes) {
+  assert(llvm::all_equal(
+      {aOffsets.size(), aSizes.size(), bOffsets.size(), bSizes.size()}));
+
+  for (const auto &t : llvm::zip(aOffsets, aSizes, bOffsets, bSizes)) {
+    auto [aBeginVal, aSizeVal, bBeginVal, bSizeVal] = t;
+    Optional<int64_t> aBegin = getConstantIntValue(aBeginVal);
+    Optional<int64_t> aSize = getConstantIntValue(aSizeVal);
+    Optional<int64_t> bBegin = getConstantIntValue(bBeginVal);
+    Optional<int64_t> bSize = getConstantIntValue(bSizeVal);
+
+    // If there are dynamic offsets/sizes, we cannot prove this dimension is
+    // disjoint. Look at other dimensions.
+    if (!aBegin || !aSize || !bBegin || !bSize)
+      continue;
+
+    int aEnd = *aBegin + *aSize;
+    int bEnd = *bBegin + *bSize;
+    // As long as one dimension is disjoint, the whole slices are disjoint.
+    if (aEnd <= *bBegin || bEnd <= *aBegin)
+      return true;
+  }
+  return false;
+}
+
+bool mlir::areDisjointSlices(OffsetSizeAndStrideOpInterface aSlice,
+                             OffsetSizeAndStrideOpInterface bSlice) {
+  SmallVector<OpFoldResult> aOffsets = aSlice.getMixedOffsets();
+  SmallVector<OpFoldResult> bOffsets = bSlice.getMixedOffsets();
+  SmallVector<OpFoldResult> aSizes = aSlice.getMixedSizes();
+  SmallVector<OpFoldResult> bSizes = bSlice.getMixedSizes();
+  SmallVector<OpFoldResult> aStrides = aSlice.getMixedStrides();
+  SmallVector<OpFoldResult> bStrides = bSlice.getMixedStrides();
+
+  // For simplicity only look at stride 1 cases for now.
+  auto hasAllOnes = [](ArrayRef<OpFoldResult> strides) {
+    return llvm::all_of(strides, [](::mlir::OpFoldResult ofr) {
+      return getConstantIntValue(ofr) == static_cast<int64_t>(1);
+    });
+  };
+  if (!hasAllOnes(aStrides) || !hasAllOnes(bStrides))
+    return false;
+
+  return areDisjointRanges(aOffsets, aSizes, bOffsets, bSizes);
+}

mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRTensorTransforms
   BufferizableOpInterfaceImpl.cpp
   Bufferize.cpp
+  ExtractFromInsertSliceDestPatterns.cpp
   ExtractSliceFromReshapeUtils.cpp
   MergeConsecutiveInsertExtractSlicePatterns.cpp
   SplitPaddingPatterns.cpp
   SwapExtractSliceWithProducerPatterns.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Tensor/Transforms
 

mlir/lib/Dialect/Tensor/Transforms/ExtractFromInsertSliceDestPatterns.cpp

+//===- ExtractFromInsertSliceDestPatterns.cpp -----------------------------===//
+//
+// 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/Affine/ViewLikeInterfaceUtils.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Tensor/Transforms/Transforms.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/PatternMatch.h"
+
+using namespace mlir;
+using namespace mlir::tensor;
+
+namespace {
+/// Updates extract_slice to extrace from insert_slice op's destination tensor
+/// when the extract_slice and insert_slice are covering disjoint slices.
+///
+/// Example:
+/// ```mlir
+/// %i = tensor.insert_slice %src into %dst[0, 0, 0, 0][1, 1, 2, 4][1, 1, 1, 1]
+///        : tensor<1x2x4xf32> into tensor<1x2x2x4xf32>
+/// %e = tensor.extract_slice %i[0, 1, 0, 0] [1, 1, 2, 4] [1, 1, 1, 1]
+///        : tensor<1x2x2x4xf32> to tensor<1x2x4xf32>
+/// ```
+/// Can be converted into
+/// ```mlir
+/// %i = tensor.insert_slice %src into %dst[0, 0, 0, 0][1, 1, 2, 4][1, 1, 1, 1]
+///        : tensor<1x2x4xf32> into tensor<1x2x2x4xf32>
+/// %e = tensor.extract_slice %dest[0, 1, 0, 0] [1, 1, 2, 4] [1, 1, 1, 1]
+///        : tensor<1x2x2x4xf32> to tensor<1x2x4xf32>
+/// ```
+/// This helps to break the chain of insert_slice and extract_slices, which
+/// might enable further optimizations.
+struct ExtractFromInsertDest final : public OpRewritePattern<ExtractSliceOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ExtractSliceOp extractOp,
+                                PatternRewriter &rewriter) const override {
+    auto insertOp = extractOp.getSource().getDefiningOp<InsertSliceOp>();
+    if (!insertOp)
+      return failure();
+
+    if (!areDisjointSlices(insertOp, extractOp))
+      return rewriter.notifyMatchFailure(extractOp, "not disjoint");
+
+    rewriter.replaceOpWithNewOp<ExtractSliceOp>(
+        extractOp, extractOp.getType(), insertOp.getDest(),
+        extractOp.getMixedOffsets(), extractOp.getMixedSizes(),
+        extractOp.getMixedStrides());
+
+    return success();
+  }
+};
+} // namespace
+
+void mlir::tensor::populateExtractFromInsertSliceDestOpPatterns(
+    RewritePatternSet &patterns) {
+  patterns.add<ExtractFromInsertDest>(patterns.getContext());
+}

mlir/lib/Interfaces/ViewLikeInterface.cpp

 //===- ViewLikeInterface.cpp - View-like operations in MLIR ---------------===//
 //
 // 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/Interfaces/ViewLikeInterface.h"
+#include "llvm/ADT/STLExtras.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // ViewLike Interfaces
 //===----------------------------------------------------------------------===//
 
 /// Include the definitions of the loop-like interfaces.

mlir/test/Dialect/Tensor/extract-from-insert-slice-dest.mlir

+// RUN: mlir-opt -split-input-file -test-tensor-transform-patterns=test-extract-from-insert-slice-dest -canonicalize %s | FileCheck %s
+
+func.func @disjoint_insert_extract_slice_static_shape(%src: tensor<1x2x4xf32>, %dst: tensor<1x2x2x4xf32>) -> tensor<1x2x4xf32> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %insert = tensor.insert_slice %src into %dst[0, 0, 0, 0] [1, 1, 2, 4] [1, 1, 1, 1] : tensor<1x2x4xf32> into tensor<1x2x2x4xf32>
+  %extract = tensor.extract_slice %insert[%c0, %c1, %c0, %c0] [1, 1, 2, 4] [1, 1, 1, 1] : tensor<1x2x2x4xf32> to tensor<1x2x4xf32>
+  return %extract : tensor<1x2x4xf32>
+}
+
+// CHECK-LABEL: func.func @disjoint_insert_extract_slice_static_shape
+//  CHECK-SAME: (%{{.+}}: tensor<1x2x4xf32>, %[[DST:.+]]: tensor<1x2x2x4xf32>)
+// CHECK:   %[[EXTRACT:.+]] = tensor.extract_slice %[[DST]][0, 1, 0, 0] [1, 1, 2, 4] [1, 1, 1, 1]
+// CHECK:   return %[[EXTRACT]]
+
+// -----
+
+func.func @disjoint_insert_extract_slice_static_shape(%src: tensor<1x2x4xf32>, %dst: tensor<1x2x2x4xf32>) -> tensor<1x2x4xf32> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %insert = tensor.insert_slice %src into %dst[0, 1, 0, 0] [1, 1, 2, 4] [1, 1, 1, 1] : tensor<1x2x4xf32> into tensor<1x2x2x4xf32>
+  %extract = tensor.extract_slice %insert[%c0, %c0, %c0, %c0] [1, 1, 2, 4] [1, 1, 1, 1] : tensor<1x2x2x4xf32> to tensor<1x2x4xf32>
+  return %extract : tensor<1x2x4xf32>
+}
+
+// CHECK-LABEL: func.func @disjoint_insert_extract_slice_static_shape
+//  CHECK-SAME: (%{{.+}}: tensor<1x2x4xf32>, %[[DST:.+]]: tensor<1x2x2x4xf32>)
+//       CHECK:   %[[EXTRACT:.+]] = tensor.extract_slice %[[DST]][0, 0, 0, 0] [1, 1, 2, 4] [1, 1, 1, 1]
+//       CHECK:   return %[[EXTRACT]]
+
+// -----
+
+func.func @disjoint_insert_extract_slice_dynamic_shape(%src: tensor<2x?xf32>, %dst: tensor<8x?xf32>, %size: index) -> tensor<3x?xf32> {
+  %insert = tensor.insert_slice %src into %dst[2, 0] [2, %size] [1, 1] : tensor<2x?xf32> into tensor<8x?xf32>
+  %extract = tensor.extract_slice %insert[5, 0] [3, %size] [1, 1] : tensor<8x?xf32> to tensor<3x?xf32>
+  return %extract : tensor<3x?xf32>
+}
+
+// CHECK-LABEL: func.func @disjoint_insert_extract_slice_dynamic_shape
+//  CHECK-SAME: (%{{.+}}: tensor<2x?xf32>, %[[DST:.+]]: tensor<8x?xf32>, %[[SIZE:.+]]: index)
+//       CHECK:   %[[EXTRACT:.+]] = tensor.extract_slice %[[DST]][5, 0] [3, %[[SIZE]]] [1, 1]
+//       CHECK:   return %[[EXTRACT]]
+
+// -----
+
+func.func @joint_insert_extract_slice_dynamic_shape(%src: tensor<2x?xf32>, %dst: tensor<8x?xf32>, %size: index) -> tensor<3x?xf32> {
+  %insert = tensor.insert_slice %src into %dst[2, 0] [2, %size] [1, 1] : tensor<2x?xf32> into tensor<8x?xf32>
+  %extract = tensor.extract_slice %insert[3, 0] [3, %size] [1, 1] : tensor<8x?xf32> to tensor<3x?xf32>
+  return %extract : tensor<3x?xf32>
+}
+
+// CHECK-LABEL: func.func @joint_insert_extract_slice_dynamic_shape
+//       CHECK:   tensor.insert_slice
+//       CHECK:   tensor.extract_slice
+
+// -----
+
+func.func @joint_insert_extract_slice_dynamic_shape(%src: tensor<2x?xf32>, %dst: tensor<8x?xf32>, %size: index) -> tensor<3x?xf32> {
+  %insert = tensor.insert_slice %src into %dst[2, 0] [2, %size] [1, 1] : tensor<2x?xf32> into tensor<8x?xf32>
+  %extract = tensor.extract_slice %insert[1, 0] [3, %size] [1, 1] : tensor<8x?xf32> to tensor<3x?xf32>
+  return %extract : tensor<3x?xf32>
+}
+
+// CHECK-LABEL: func.func @joint_insert_extract_slice_dynamic_shape
+//       CHECK:   tensor.insert_slice
+//       CHECK:   tensor.extract_slice
+
+
+// -----
+
+func.func @joint_insert_extract_slice_dynamic_shape(%src: tensor<2x?xf32>, %dst: tensor<8x?xf32>, %offset: index, %size: index) -> tensor<3x?xf32> {
+  %insert = tensor.insert_slice %src into %dst[2, 0] [2, %size] [1, 1] : tensor<2x?xf32> into tensor<8x?xf32>
+  %extract = tensor.extract_slice %insert[%offset, 0] [3, %size] [1, 1] : tensor<8x?xf32> to tensor<3x?xf32>
+  return %extract : tensor<3x?xf32>
+}
+
+// CHECK-LABEL: func.func @joint_insert_extract_slice_dynamic_shape
+//       CHECK:   tensor.insert_slice
+//       CHECK:   tensor.extract_slice

mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp

Option<bool> testFoldConstantExtractSlice{
       llvm::cl::init(false)};
 
   Option<bool> testFoldConsecutiveInsertExtractSlice{
       *this, "test-fold-consecutive-insert-extract-slice",
       llvm::cl::desc(
           "Test folding consecutive tensor.insert_slice/tensor.extract_slice"),
       llvm::cl::init(false)};
 
+  Option<bool> testExtractFrominsertSliceDest{
+      *this, "test-extract-from-insert-slice-dest",
+      llvm::cl::desc("Test tensor.extract_slice from tensor.insert_slice "
+                     "destination tensor"),
+      llvm::cl::init(false)};
+
   Option<bool> testRewriteExtractSliceWithTiledCollapseShape{
       *this, "test-rewrite-extract-slice-from-collapse-shape",
       llvm::cl::desc("Test swapping tensor.extract_slice of a collapse_shape "
                      "with loop nest"),
       llvm::cl::init(false)};
 
   Option<bool> useForeach{
       *this, "use-foreach",

 }
 
 static void applyFoldConsecutiveInsertExtractSlicePatterns(Operation *rootOp) {
   RewritePatternSet patterns(rootOp->getContext());
   tensor::populateMergeConsecutiveInsertExtractSlicePatterns(patterns);
   (void)applyPatternsAndFoldGreedily(rootOp, std::move(patterns));
 }
 
+static void applyExtractFromInsertSliceDestPatterns(Operation *rootOp) {
+  RewritePatternSet patterns(rootOp->getContext());
+  tensor::populateExtractFromInsertSliceDestOpPatterns(patterns);
+  (void)applyPatternsAndFoldGreedily(rootOp, std::move(patterns));
+}
+
 namespace {
 /// Base pattern to rewrite  a `tensor.collapse_shape -> tensor.extract_slice`.
 /// The `tensor.extract_slice` is replaced by a loop or gather operation that
 /// stitches together the desired tile from slices of the source of the collapse
 /// shape op.
 struct RewriteExtractSliceFromCollapseShapeBase
     : public OpRewritePattern<tensor::ExtractSliceOp> {
   RewriteExtractSliceFromCollapseShapeBase(MLIRContext *context)

 void TestTensorTransforms::runOnOperation() {
   Operation *rootOp = getOperation();
   if (testSplitPaddingPatterns)
     applySplitPaddingPatterns(rootOp);
   if (testFoldConstantExtractSlice)
     applyFoldConstantExtractSlicePatterns(rootOp);
   if (testFoldConsecutiveInsertExtractSlice)
     applyFoldConsecutiveInsertExtractSlicePatterns(rootOp);
+  if (testExtractFrominsertSliceDest)
+    applyExtractFromInsertSliceDestPatterns(rootOp);
   if (testRewriteExtractSliceWithTiledCollapseShape) {
     if (failed(
             applyRewriteExtractFromCollapseShapePatterns(rootOp, useForeach)))
       return signalPassFailure();
   }
 }
 
 namespace mlir {
 namespace test {
 void registerTestTensorTransforms() {
   PassRegistration<TestTensorTransforms>();
 }
 } // namespace test
 } // namespace mlir