[mlir][tensor] Add pattern to drop redundant insert_slice rank expansion

[springerm]

Summary

Drop insert_slice rank expansions if they are directly followed by an inverse rank reduction.

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);
 
+/// Populates `patterns` with patterns that drop redundant tensor.insert_slice
+/// rank expansions.
+void populateDropRedundantInsertSliceRankExpansionPatterns(
+    RewritePatternSet &patterns);
+
 /// Populates `patterns` with patterns that fold `tensor.expand_shape` and
 /// `tensor.collapse_shape` into other ops.
 void populateReassociativeReshapeFoldingPatterns(RewritePatternSet &patterns);
 
 /// Populates `patterns` with patterns that fold tensor.empty with
 /// tensor.[extract_slice|cast|expand_shape|collapse_shape].
 void populateFoldTensorEmptyPatterns(RewritePatternSet &patterns);
 

mlir/include/mlir/Dialect/Tensor/Utils/Utils.h

SmallVector<OpFoldResult> createDimValues(OpBuilder &b, Location loc,
                                           Value rankedTensor);
 
 /// Returns the transposed `rankedTensorType` if `transposeVector` is non-empty.
 /// Fail if `transposeVector` is not a permutation matching the tensor rank.
 FailureOr<RankedTensorType>
 computeTransposedType(RankedTensorType rankedTensorType,
                       ArrayRef<int64_t> transposeVector);
 
+/// A tensor.insert_slice is a cast-like operation if it merely rank-extends the
+/// source tensor or inserts the source tensor into a destination tensor with
+/// the same shape.
+bool isCastLikeInsertSliceOp(InsertSliceOp op);
+
 } // namespace tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_TENSOR_UTILS_UTILS_H_

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

 add_mlir_dialect_library(MLIRTensorTransformOps
   TensorTransformOps.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Tensor/TransformOps
 
   DEPENDS
   MLIRTensorTransformOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
   MLIRIR
   MLIRSCFDialect
   MLIRTensorDialect
   MLIRTensorTransforms
+  MLIRTensorUtils
   MLIRTransformDialect
-  MLIRValueBoundsOpInterface
 )

mlir/lib/Dialect/Tensor/TransformOps/TensorTransformOps.cpp

 //===- TensorTransformOps.cpp - Implementation of tensor transform ops ----===//
 //
 // 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/TransformOps/TensorTransformOps.h"
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/Transforms/Transforms.h"
+#include "mlir/Dialect/Tensor/Utils/Utils.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
-#include "mlir/Interfaces/ValueBoundsOpInterface.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 using namespace tensor;
 
 //===----------------------------------------------------------------------===//
 // TrackingListener
 //===----------------------------------------------------------------------===//
 
-/// A tensor.insert_slice is a cast-like operation if it merely rank-extends the
-/// source tensor or inserts the source tensor into a destination tensor with
-/// the same shape.
-static bool isCastLikeInsertSliceOp(InsertSliceOp op) {
-  llvm::SmallBitVector droppedDims = op.getDroppedDims();
-  int64_t srcDim = 0;
-  // Source dims and destination dims (apart from dropped dims) must have the
-  // same size.
-  for (int64_t resultDim = 0; resultDim < op.getDestType().getRank();
-       ++resultDim) {
-    if (droppedDims.test(resultDim)) {
-      continue;
-    }
-    FailureOr<bool> equalDimSize = ValueBoundsConstraintSet::areEqual(
-        op.getSource(), op.getResult(), srcDim, resultDim);
-    if (failed(equalDimSize) || !*equalDimSize)
-      return false;
-    ++srcDim;
-  }
-
-  return true;
-}
-
 Operation *
 tensor::TrackingListener::findReplacementOp(Operation *op,
                                             ValueRange newValues) const {
   SmallVector<Value> values(newValues.begin(), newValues.end());
   do {
     if (Operation *replacement =
             transform::TrackingListener::findReplacementOp(op, values))
       return replacement;

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

add_mlir_dialect_library(MLIRTensorTransforms
   MLIRBufferizationDialect
   MLIRBufferizationTransforms
   MLIRIR
   MLIRLinalgDialect
   MLIRMemRefDialect
   MLIRPass
   MLIRSCFDialect
   MLIRTensorDialect
+  MLIRTensorUtils
   MLIRTilingInterface
   MLIRTransforms
   MLIRVectorDialect
   MLIRValueBoundsOpInterface
 )

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

 //===- MergeConsecutiveInsertExtractSlicePatterns.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/Dialect/Tensor/Utils/Utils.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/PatternMatch.h"
 
 using namespace mlir;
 using namespace mlir::tensor;
 
 namespace {
if (!prevOp.getSourceType().hasStaticShape() ||
       return failure();
 
     rewriter.replaceOpWithNewOp<OpTy>(
         nextOp, prevOp.getSource(), nextOp.getDest(), nextOp.getMixedOffsets(),
         nextOp.getMixedSizes(), nextOp.getMixedStrides());
     return success();
   }
 };
+
+/// Drop redundant rank expansion. I.e., rank expansions that are directly
+/// followed by rank reductions. E.g.:
+/// %0 = tensor.insert_slice ... : tensor<5x10xf32> into tensor<1x1x5x10xf32>
+/// %1 = tensor.extract_slice %0[0, 0, 2, 3] [1, 1, 2, 2] [1, 1, 1, 1]
+///     : tensor<1x1x5x10xf32> to tensor<2x2xf32>
+struct DropRedundantInsertSliceRankExpansion
+    : public OpRewritePattern<ExtractSliceOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ExtractSliceOp extractSliceOp,
+                                PatternRewriter &rewriter) const override {
+    // Nothing to do if no dims are dropped.
+    llvm::SmallBitVector droppedDims = extractSliceOp.getDroppedDims();
+    if (droppedDims.empty())
+      return failure();
+
+    // Look for tensor.insert_slice op that has an inverse rank expansion.
+    auto insertSliceOp =
+        extractSliceOp.getSource().getDefiningOp<InsertSliceOp>();
+    if (!insertSliceOp)
+      return failure();
+    llvm::SmallBitVector expandedDims = insertSliceOp.getDroppedDims();
+
+    // TODO: This could be extended to support cases where the dropped dims are
+    // a subset of the expanded dims.
+    if (expandedDims != droppedDims)
+      return failure();
+
+    // The tensor.insert_slice may not be redundant if it has multiple users.
+    if (!insertSliceOp->hasOneUse())
+      return failure();
+
+    // Only consider tensor.insert_slice ops that are pure rank-reductions.
+    // I.e., no elements are taken from the destination.
+    if (!isCastLikeInsertSliceOp(insertSliceOp))
+      return failure();
+
+    // Extract directly from the source.
+    OpBuilder::InsertionGuard g(rewriter);
+    rewriter.setInsertionPoint(extractSliceOp);
+    SmallVector<OpFoldResult> newOffsets, newSizes, newStrides;
+    for (int64_t i = 0, e = extractSliceOp.getSourceType().getRank(); i < e;
+         ++i) {
+      if (droppedDims.test(i))
+        continue;
+      newOffsets.push_back(extractSliceOp.getMixedOffsets()[i]);
+      newSizes.push_back(extractSliceOp.getMixedSizes()[i]);
+      newStrides.push_back(extractSliceOp.getMixedStrides()[i]);
+    }
+    rewriter.replaceOpWithNewOp<ExtractSliceOp>(
+        extractSliceOp, /*source=*/insertSliceOp.getSource(), newOffsets,
+        newSizes, newStrides);
+    rewriter.eraseOp(insertSliceOp);
+    return success();
+  }
+};
 } // namespace
 
 void mlir::tensor::populateMergeConsecutiveInsertExtractSlicePatterns(
     RewritePatternSet &patterns) {
   patterns.add<MergeConsecutiveExtractSlice,
                MergeConsecutiveInsertSlice<InsertSliceOp>,
                MergeConsecutiveInsertSlice<ParallelInsertSliceOp>>(
       patterns.getContext());
 }
+
+void mlir::tensor::populateDropRedundantInsertSliceRankExpansionPatterns(
+    RewritePatternSet &patterns) {
+  patterns.add<DropRedundantInsertSliceRankExpansion>(patterns.getContext());
+}

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

 add_mlir_dialect_library(MLIRTensorUtils
   Utils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Tensor
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
   MLIRArithDialect
   MLIRArithUtils
   MLIRIR
   MLIRTensorDialect
+  MLIRValueBoundsOpInterface
 )

mlir/lib/Dialect/Tensor/Utils/Utils.cpp

 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Tensor/Utils/Utils.h"
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/Utils/IndexingUtils.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
 
 using namespace mlir;
 using namespace mlir::tensor;
 
 PadOp mlir::tensor::createPadHighOp(RankedTensorType type, Value source,
                                     Value pad, bool nofold, Location loc,
                                     OpBuilder &b) {
   auto zero = b.createOrFold<arith::ConstantIndexOp>(loc, 0);
SmallVector<int64_t> transposedShape(rankedTensorType.getShape().begin(),
                                        rankedTensorType.getShape().end());
   applyPermutationToVector(transposedShape, transposeVector);
 
   using RTTBuilder = RankedTensorType::Builder;
   RankedTensorType transposedTensorType =
       RTTBuilder(rankedTensorType).setShape(transposedShape);
   return transposedTensorType;
 }
+
+bool mlir::tensor::isCastLikeInsertSliceOp(InsertSliceOp op) {
+  llvm::SmallBitVector droppedDims = op.getDroppedDims();
+  int64_t srcDim = 0;
+  // Source dims and destination dims (apart from dropped dims) must have the
+  // same size.
+  for (int64_t resultDim = 0; resultDim < op.getDestType().getRank();
+       ++resultDim) {
+    if (droppedDims.test(resultDim)) {
+      continue;
+    }
+    FailureOr<bool> equalDimSize = ValueBoundsConstraintSet::areEqual(
+        op.getSource(), op.getResult(), srcDim, resultDim);
+    if (failed(equalDimSize) || !*equalDimSize)
+      return false;
+    ++srcDim;
+  }
+
+  return true;
+}

mlir/test/Dialect/Tensor/drop-redundant-insert-slice-rank-expansion.mlir

+// RUN: mlir-opt -split-input-file -test-tensor-transform-patterns=test-drop-redundant-insert-slice-rank-expansion %s | FileCheck %s
+
+// CHECK-LABEL: func @test_drop_rank_expansion(
+//  CHECK-SAME:     %[[src:.*]]: tensor<128x480xf32>,
+//       CHECK:   %[[extract:.*]] = tensor.extract_slice %[[src]][0, 0] [123, 456] [1, 1] : tensor<128x480xf32> to tensor<123x456xf32>
+//       CHECK:   return %[[extract]]
+func.func @test_drop_rank_expansion(%src: tensor<128x480xf32>, %dest: tensor<1x1x128x480xf32>) -> tensor<123x456xf32> {
+  %inserted_slice = tensor.insert_slice %src into %dest[0, 0, 0, 0] [1, 1, 128, 480] [1, 1, 1, 1] : tensor<128x480xf32> into tensor<1x1x128x480xf32>
+  %extracted_slice = tensor.extract_slice %inserted_slice[0, 0, 0, 0] [1, 1, 123, 456] [1, 1, 1, 1] : tensor<1x1x128x480xf32> to tensor<123x456xf32>
+  return %extracted_slice : tensor<123x456xf32>
+}

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

Option<bool> testFoldConsecutiveInsertExtractSlice{
       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> testDropRedundantInsertSliceRankExpansion{
+      *this, "test-drop-redundant-insert-slice-rank-expansion",
+      llvm::cl::desc("Test dropping redundant insert_slice rank expansions"),
+      llvm::cl::init(false)};
+
   Option<bool> testReassociativeReshapeFolding{
       *this, "test-reassociative-reshape-folding",
       llvm::cl::desc("Test folding of expand_shape/collapse_shape"),
       llvm::cl::init(false)};
 
   Option<bool> testEmptyOpFolding{
       *this, "test-empty-op-folding",
       llvm::cl::desc("Test folding of tensor.empty"), llvm::cl::init(false)};

 }
 
 static void applyFoldConsecutiveInsertExtractSlicePatterns(Operation *rootOp) {
   RewritePatternSet patterns(rootOp->getContext());
   tensor::populateMergeConsecutiveInsertExtractSlicePatterns(patterns);
   (void)applyPatternsAndFoldGreedily(rootOp, std::move(patterns));
 }
 
+static void
+applyDropRedundantInsertSliceRankExpansionPatterns(Operation *rootOp) {
+  RewritePatternSet patterns(rootOp->getContext());
+  tensor::populateDropRedundantInsertSliceRankExpansionPatterns(patterns);
+  (void)applyPatternsAndFoldGreedily(rootOp, std::move(patterns));
+}
+
 static void applySimplifyPackPatterns(Operation *rootOp) {
   RewritePatternSet patterns(rootOp->getContext());
   tensor::populateSimplifyTensorPack(patterns);
   (void)applyPatternsAndFoldGreedily(rootOp, std::move(patterns));
 }
 
 namespace {
 /// Base pattern to rewrite  a `tensor.collapse_shape -> tensor.extract_slice`.

 void TestTensorTransforms::runOnOperation() {
   Operation *rootOp = getOperation();
   if (testSimplifyPackPatterns)
     applySimplifyPackPatterns(rootOp);
   if (testFoldConstantExtractSlice)
     applyFoldConstantExtractSlicePatterns(rootOp);
   if (testFoldConsecutiveInsertExtractSlice)
     applyFoldConsecutiveInsertExtractSlicePatterns(rootOp);
+  if (testDropRedundantInsertSliceRankExpansion)
+    applyDropRedundantInsertSliceRankExpansionPatterns(rootOp);
   if (testReassociativeReshapeFolding)
     applyReassociativeReshapeFoldingPatterns(rootOp);
   if (testEmptyOpFolding)
     applyEmptyOpFoldingPatterns(rootOp);
   if (testFoldIntoPackAndUnpack)
     applyFoldIntoPackAndUnpackPatterns(rootOp);
   if (testRewriteExtractSliceWithTiledCollapseShape) {
     if (failed(

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

cc_library(
     hdrs = ["include/mlir/Dialect/Tensor/Utils/Utils.h"],
     includes = ["include"],
     deps = [
         ":AffineDialect",
         ":ArithDialect",
         ":ArithUtils",
         ":DialectUtils",
         ":TensorDialect",
+        ":ValueBoundsOpInterface",
         "//llvm:Support",
     ],
 )
 
 gentbl_cc_library(
     name = "TensorPassIncGen",
     strip_include_prefix = "include",
     tbl_outs = [
deps = [
         ":FuncDialect",
         ":IR",
         ":LinalgDialect",
         ":MemRefDialect",
         ":Pass",
         ":SCFDialect",
         ":TensorDialect",
         ":TensorPassIncGen",
+        ":TensorUtils",
         ":TilingInterface",
         ":Transforms",
         ":ValueBoundsOpInterface",
         ":VectorDialect",
         "//llvm:Support",
     ],
 )
 
cc_library(
     includes = ["include"],
     deps = [
         ":AffineDialect",
         ":IR",
         ":SCFDialect",
         ":TensorDialect",
         ":TensorTransformOpsIncGen",
         ":TensorTransforms",
+        ":TensorUtils",
         ":TransformDialect",
         ":ValueBoundsOpInterface",
         "//llvm:Support",
     ],
 )
 
 cc_library(
     name = "Rewrite",