[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/Tensor/Transforms/Transforms.h

 /// Pattern to swap an `tensor.extract_slice` with its producer when the
 /// producer implements the `TilingInterface`. The pattern itself does not
 /// provide a mechanism to control where the application happens. With use of
 /// transform dialect that control is done within the transform dialect. Other
 /// use cases can inherit from this pattern and add necessary controls.
 FailureOr<Value> replaceExtractSliceWithTiledProducer(
     OpBuilder &builder, tensor::ExtractSliceOp sliceOp, OpResult producerOp);
 
+/// 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/include/mlir/Interfaces/ViewLikeInterface.h

 
 /// Decompose a vector of mixed static or dynamic strides/offsets into the
 /// corresponding pair of arrays. This is the inverse function of
 /// `getMixedSizes`.
 std::pair<ArrayAttr, SmallVector<Value>>
 decomposeMixedSizes(OpBuilder &b,
                     const SmallVectorImpl<OpFoldResult> &mixedValues);
 
+/// 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);
+
 class OffsetSizeAndStrideOpInterface;
 
 namespace detail {
 
 LogicalResult verifyOffsetSizeAndStrideOp(OffsetSizeAndStrideOpInterface op);
 
 bool sameOffsetsSizesAndStrides(
     OffsetSizeAndStrideOpInterface a, OffsetSizeAndStrideOpInterface b,
parseDynamicIndexList(OpAsmParser &parser,
                       ArrayAttr &integers, int64_t dynVal);
 
 /// Verify that a the `values` has as many elements as the number of entries in
 /// `attr` for which `isDynamic` evaluates to true.
 LogicalResult verifyListOfOperandsOrIntegers(
     Operation *op, StringRef name, unsigned expectedNumElements, ArrayAttr attr,
     ValueRange values, function_ref<bool(int64_t)> isDynamic);
 
+/// 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_INTERFACES_VIEWLIKEINTERFACE_H_

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

 add_mlir_dialect_library(MLIRTensorTransforms
   BufferizableOpInterfaceImpl.cpp
   Bufferize.cpp
+  ExtractFromInsertSliceDestPatterns.cpp
   ExtractSliceFromReshape.cpp
   SplitPadding.cpp
   SwapExtractSliceWithProducer.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Tensor/Transforms
 
   DEPENDS

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/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Tensor/Transforms/Transforms.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Interfaces/ViewLikeInterface.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.
return decomposeMixedValues(b, mixedValues,
                               ShapedType::kDynamicStrideOrOffset);
 }
 
 std::pair<ArrayAttr, SmallVector<Value>>
 mlir::decomposeMixedSizes(OpBuilder &b,
                           const SmallVectorImpl<OpFoldResult> &mixedValues) {
   return decomposeMixedValues(b, mixedValues, ShapedType::kDynamicSize);
 }
+
+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/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> testSplitPaddingPatterns{
       llvm::cl::desc("Test patterns to split tensor.pad ops"),
       llvm::cl::init(false)};
 
   Option<bool> testFoldConstantExtractSlice{
       *this, "test-fold-constant-extract-slice",
       llvm::cl::desc("Test folding arith.constant and 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",
tensor::ControlConstantExtractSliceFusionFn controlFn =
         constexpr int64_t kConstantFoldingMaxNumElements = 1024;
         return resultType.getNumElements() <= kConstantFoldingMaxNumElements;
       };
 
   tensor::populateFoldConstantExtractSlicePatterns(patterns, controlFn);
   (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 (testExtractFrominsertSliceDest)
+    applyExtractFromInsertSliceDestPatterns(rootOp);
   if (testRewriteExtractSliceWithTiledCollapseShape) {
     if (failed(
             applyRewriteExtractFromCollapseShapePatterns(rootOp, useForeach)))
       return signalPassFailure();
   }
 }
 
 namespace mlir {
 namespace test {
 void registerTestTensorTransforms() {
   PassRegistration<TestTensorTransforms>();
 }
 } // namespace test
 } // namespace mlir