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

[springerm]

Summary

Depends On D139118

[nicolasvasilache]

Thanks!

I think replacing expand_shape by extract_slice where possible is a great canonicalization.
I think unit-dim-like patterns that create expand_shape however need to be updated too.

[mravishankar]

My thumb rule is if an op is not a simple op -> op change it doesnt pay for itself as a canonicalization. I am happy to discuss if indeed it is a canonicalization, but forcing everything to be rank-reducing when canonicalization is invoked is a very opinionated solution. It should be opt-in (ftr, I am in favor of using rank-reduce slices everywhere, but not in favor of forcing that view on everyone)

[springerm]

address comments

[springerm]

Patterns are now no longer canonicalizations and must be added via a populate... function.

[mravishankar]

Echoing my comment from the other patch

This is yet another instance of a patch being submitted when it has still changes requested marked. Whats the hurry, and AFAIK is not in keeping with developer practices. When patches land this way it drops off the dashboard and I dont see the resolution of the issue. That aside, specifically here, https://reviews.llvm.org/D139118 was tagged as something that conflicted with this. Does it or does it not. I am interested to know whats happening here (and hence I requested changes). Please do not land patches with request changes unless its been a while from the reviewer.

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