[mlir][tensor] Fold rank-reducing extract_slice with inverse expand_shape

[springerm]

Summary

Reupload of D139103.

[mravishankar]

Thank you for the revert and resending patch. Appreciate it! I meant to flag it just cause I saw it as a repeating pattern, and it justified being flagged. The patch itself is fine now, but want some clarification on why this interfered with https://reviews.llvm.org/D139118 . AFAICS it shouldnt (or is not meant to)

[springerm]

In the very first version (when the two new patterns were canonicalization patterns), this used to interfere with populateFoldUnitExtentDimsPatterns, which also includes some canonicalization patterns. In the current revision, these two patterns must be added explicitly via populateReassociativeReshapeFoldingPatterns.

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 fold `tensor.expand_shape` and
+/// `tensor.collapse_shape` into other ops.
+void populateReassociativeReshapeFoldingPatterns(RewritePatternSet &patterns);
+
 } // namespace tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_TENSOR_TRANSFORMS_TRANSFORMS_H

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

 add_mlir_dialect_library(MLIRTensorTransforms
   BufferizableOpInterfaceImpl.cpp
   Bufferize.cpp
   ExtractSliceFromReshapeUtils.cpp
   MergeConsecutiveInsertExtractSlicePatterns.cpp
+  ReshapePatterns.cpp
   SplitPaddingPatterns.cpp
   SwapExtractSliceWithProducerPatterns.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Tensor/Transforms
 
   DEPENDS
   MLIRTensorTransformsIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
   MLIRAffineUtils
   MLIRArithDialect
   MLIRBufferizationDialect
   MLIRBufferizationTransforms
   MLIRIR
   MLIRLinalgDialect
   MLIRMemRefDialect
   MLIRPass
   MLIRSCFDialect
   MLIRTensorDialect
   MLIRTilingInterface
   MLIRTransforms
-  )

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

+//===- RankReductionPatterns.cpp - Patterns related to rank reductions ----===//
+//
+// 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/PatternMatch.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "mlir-tensor-split-padding"
+
+using namespace mlir;
+using namespace mlir::tensor;
+
+namespace {
+/// Fold expand_shape(extract_slice) ops that cancel itself out.
+struct FoldExpandOfRankReducingExtract
+    : public OpRewritePattern<ExpandShapeOp> {
+  using OpRewritePattern<ExpandShapeOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ExpandShapeOp expandShapeOp,
+                                PatternRewriter &rewriter) const override {
+    RankedTensorType resultType = expandShapeOp.getResultType();
+    auto extractSliceOp =
+        expandShapeOp.getSrc().getDefiningOp<ExtractSliceOp>();
+    if (!extractSliceOp)
+      return failure();
+    RankedTensorType srcType = extractSliceOp.getSourceType();
+
+    // Only cases where the ExpandShapeOp can be folded away entirely are
+    // supported. Moreover, only simple cases where the resulting ExtractSliceOp
+    // has no rank-reduction anymore are supported at the moment.
+    RankedTensorType nonReducingExtractType = ExtractSliceOp::inferResultType(
+        srcType, extractSliceOp.getStaticOffsets(),
+        extractSliceOp.getStaticSizes(), extractSliceOp.getStaticStrides());
+    if (nonReducingExtractType != resultType)
+      return failure();
+
+    SmallVector<OpFoldResult> mixedOffsets = extractSliceOp.getMixedOffsets();
+    SmallVector<OpFoldResult> mixedSizes = extractSliceOp.getMixedSizes();
+    SmallVector<OpFoldResult> mixedStrides = extractSliceOp.getMixedStrides();
+    rewriter.replaceOpWithNewOp<tensor::ExtractSliceOp>(
+        expandShapeOp, extractSliceOp.getSource(), mixedOffsets, mixedSizes,
+        mixedStrides);
+    return success();
+  }
+};
+} // namespace
+
+void mlir::tensor::populateReassociativeReshapeFoldingPatterns(
+    RewritePatternSet &patterns) {
+  patterns.add<FoldExpandOfRankReducingExtract>(patterns.getContext());
+}

mlir/test/Dialect/Tensor/fold-reassociative-reshapes.mlir

+// RUN: mlir-opt -split-input-file -test-tensor-transform-patterns=test-reassociative-reshape-folding %s | FileCheck %s
+
+// CHECK-LABEL: func @expand_shape_of_rank_reducing_extract(
+//  CHECK-SAME:     %[[t:.*]]: tensor<?x?x?x?xf32>
+//   CHECK-DAG:   %[[extract1:.*]] = tensor.extract_slice %{{.*}}[0, 0, 0, 0] [%{{.*}}, 1, 1, 5] [1, 1, 1, 1] : tensor<?x?x?x?xf32> to tensor<?x1x1x5xf32>
+//   CHECK-DAG:   %[[extract2:.*]] = tensor.extract_slice %{{.*}}[0, 0, 0, 0] [%{{.*}}, 1, 1, 5] [1, 1, 1, 1] : tensor<?x?x?x?xf32> to tensor<?x1x1x5xf32>
+//       CHECK:   return %[[extract1]], %[[extract2]]
+func.func @expand_shape_of_rank_reducing_extract(
+    %t: tensor<?x?x?x?xf32>, %idx: index)
+  -> (tensor<?x1x1x5xf32>, tensor<?x1x1x5xf32>)
+{
+  %0 = tensor.extract_slice %t[0, 0, 0, 0][%idx, 1, 1, 5][1, 1, 1, 1]
+      : tensor<?x?x?x?xf32> to tensor<?x1x5xf32>
+  %1 = tensor.expand_shape %0 [[0], [1, 2], [3]]
+      : tensor<?x1x5xf32> into tensor<?x1x1x5xf32>
+  %2 = tensor.expand_shape %0 [[0, 1], [2], [3]]
+      : tensor<?x1x5xf32> into tensor<?x1x1x5xf32>
+  return %1, %2 : tensor<?x1x1x5xf32>, tensor<?x1x1x5xf32>
+}

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> testReassociativeReshapeFolding{
+      *this, "test-reassociative-reshape-folding",
+      llvm::cl::desc("Test folding of expand_shape/collapse_shape"),
+      llvm::cl::init(false)};
+
   Option<bool> useForeach{
       *this, "use-foreach",
       llvm::cl::desc(
           "Use the scf.foreach_thread operation when generating loop nests for "
           "the extract_slice of collapse_shape pattern"),
       llvm::cl::init(false)};
 };
 } // namespace
 
+static void applyReassociativeReshapeFoldingPatterns(Operation *rootOp) {
+  RewritePatternSet patterns(rootOp->getContext());
+  tensor::populateReassociativeReshapeFoldingPatterns(patterns);
+  (void)applyPatternsAndFoldGreedily(rootOp, std::move(patterns));
+}
+
 static void applySplitPaddingPatterns(Operation *rootOp) {
   RewritePatternSet patterns(rootOp->getContext());
   tensor::populateSplitPaddingPatterns(patterns);
   (void)applyPatternsAndFoldGreedily(rootOp, std::move(patterns));
 }
 
 static void applyFoldConstantExtractSlicePatterns(Operation *rootOp) {
   RewritePatternSet patterns(rootOp->getContext());

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