[mlir][Linalg] Mostly NFC: Change fusion on tensor to use RewritePatterns.

[mravishankar]

Summary

Some of the fusions used in the LinalgFusionOfTensorOpsPass can be
treated as regular folding patterns. To allow for this, refactor the
fusion logic to execute within different patterns. These patterns can
be later moved into folding/canonicalize patterns, but for now just
provide a hook to add these patterns to the pattern list.

Depends on D89002

[nicolasvasilache]

Looks good overall.
Noted a bunch of places where it should be possible to significantly reduce code size / legibility with some mechanical refactorings.

Thanks @mravishankar !

mlir/include/mlir/Dialect/Linalg/Passes.h

59

s/an/a

mlir/lib/Dialect/Linalg/Transforms/FusionOnTensors.cpp

595

could you please put some pseudo-IR snippet to doc this?

624

Can we refactor this loop into a condition using llvm::all_of that captures the index?
Then we can just exit early here and remove most of the code from under the loop.

It was confusing to me to see that we have a loop but we always at most execute only 1 iteration.

628

getInputIndexingMap would not assume that you know the internal storage order of maps.

640

llvm::to_vector<4>(op.indexing_maps().getAsValueRange<AffineMapAttr>());
should work?

662

getAsValueRange

666

Can we update in place and simplify the logic around this or do you expect it to be more complex?
Maybe that could let us drop createLinalgOpOfSameType too.
This seems to be more generally applicable and maybe I can also drop clone and create from the LinalgStructuredOpInterface.

680

This seems like something you should let DCE do?

682

Same comment about restructuring with llvm::all_of that captures the first element you want and dropping the nesting under this "execute at most once"-loop.

702

same comment re DCE.

726–727

getAsValueRange

747

Same Q re in-place update.

761

DCE?

806

Same remark re loop/getAsValueRange/DCE as above.
Same Q on DCE as above.

890

Please refactor at most one loop.

900

DCE.

[mravishankar]

Addressing comments.

[mravishankar]

I addressed some of the comments. Others I will address in a follow up CL to not make these stack of changes already larger. Thanks for the review!

mlir/lib/Dialect/Linalg/Transforms/FusionOnTensors.cpp

624

This part is just a refactor of the logic that existed earlier. I can fix this up in a later CL. These set of changes are already quite big.

640

Thanks!

662

The base container is not an ArrayAttr. But I can use rewriter.getAffineMapArrayAttr at the create method. So I can remove these lines.

666

Good idea. I will try it out.

680

Is there a DCE pass? Would be painful to write tests without it. For now leaving it as is. Ill do another follow up with addressing these comments.

682

Sure thing. WIll address this in a follow up (though its simpler here).

702

Please see above.

747

Here the linalg.generic/indexed_generic is the producer. Not sure its a good idea to replace the operation in place. Its better to create a new fused operation cause the uses of the consumer linalg.tensor_reshape is being replaced with the new op.

806

The loop question, ill do in a follow up. Used getAsValueRange. Maybe inplace update is possible, but I need to signature convert the region since an operand is dropped. I think its easier to do by cloning. So leaving as is.

[mravishankar]

Has been folded into commit rGde2568aab819: [mlir][Linalg] Rethink fusion of linalg ops with reshape ops. (authored by mravishankar)

mlir/include/mlir/Dialect/Linalg/Passes.h

 #define MLIR_DIALECT_LINALG_PASSES_H_
 
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 std::unique_ptr<OperationPass<FuncOp>> createLinalgFoldUnitExtentDimsPass();
 
 std::unique_ptr<OperationPass<FuncOp>> createLinalgFusionPass();
-std::unique_ptr<Pass>
-createLinalgFusionOfTensorOpsPass(bool useReshapeFusionByExpansion = false);
+std::unique_ptr<Pass> createLinalgFusionOfTensorOpsPass();
+std::unique_ptr<Pass> createFoldReshapeOpsByLinearizationPass();
 
 std::unique_ptr<OperationPass<FuncOp>>
 createLinalgTilingPass(ArrayRef<int64_t> tileSizes = {});
 
 std::unique_ptr<OperationPass<FuncOp>>
 createLinalgTilingToParallelLoopsPass(ArrayRef<int64_t> tileSizes = {});
 
 std::unique_ptr<OperationPass<FuncOp>>

 /// Placeholder for now, this is NYI.
 std::unique_ptr<OperationPass<FuncOp>> createConvertLinalgToAffineLoopsPass();
 
 /// Create a pass to convert Linalg operations which work on tensors to use
 /// buffers instead.
 std::unique_ptr<OperationPass<ModuleOp>>
 createConvertLinalgOnTensorsToBuffersPass();
 
+/// Patterns to fold an expanding (collapsing) tensor_reshape operation with its
+/// producer (consumer) generic operation by expanding the dimensionality of the
+/// loop in the generic op.
+void populateFoldReshapeOpsByExpansionPatterns(
+    MLIRContext *context, OwningRewritePatternList &patterns);
+
+/// Patterns to fold a collapsing (expanding) tensor_reshape operation with its
+/// producer (consumer) generic/indexed_generic operation by linearizing the
+/// indexing map used to access the source (target) of the reshape operation in
+/// the generic/indexed_generic operation.
+void populateFoldReshapeOpsByLinearizationPatterns(
+    MLIRContext *context, OwningRewritePatternList &patterns);
+
 /// Patterns for fusing linalg operation on tensors.
-void populateLinalgTensorOpsFusionPatterns(
-    MLIRContext *context, OwningRewritePatternList &patterns,
-    bool useReshapeFusionByExpansion = false);
+void populateLinalgTensorOpsFusionPatterns(MLIRContext *context,
+                                           OwningRewritePatternList &patterns);
 
 /// Patterns to fold unit-extent dimensions in operands/results of linalg ops on
 /// tensors.
 void populateLinalgFoldUnitExtentDimsPatterns(
     MLIRContext *context, OwningRewritePatternList &patterns);
 
 //===----------------------------------------------------------------------===//
 // Registration

mlir/include/mlir/Dialect/Linalg/Passes.td

 }
 
 def LinalgFusionOfTensorOps : Pass<"linalg-fusion-for-tensor-ops"> {
   let summary = "Fuse operations on RankedTensorType in linalg dialect";
   let constructor = "mlir::createLinalgFusionOfTensorOpsPass()";
   let dependentDialects = ["linalg::LinalgDialect", "AffineDialect"];
 }
 
+def LinalgFoldReshapeOpsByLinearization :
+  Pass<"linalg-fold-reshape-ops-by-linearization"> {
+  let summary = "Fold TensorReshapeOps with generic/indexed generic ops by "
+                "linearization";
+  let constructor = "mlir::createFoldReshapeOpsByLinearizationPass()";
+  let dependentDialects = ["AffineDialect"];
+}
+
 def LinalgLowerToAffineLoops : FunctionPass<"convert-linalg-to-affine-loops"> {
   let summary = "Lower the operations from the linalg dialect into affine "
                 "loops";
   let constructor = "mlir::createConvertLinalgToAffineLoopsPass()";
   let dependentDialects = ["linalg::LinalgDialect", "AffineDialect"];
 }
 
 def LinalgLowerToLoops : FunctionPass<"convert-linalg-to-loops"> {

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

Optional<FusionInfo> fuseProducerOf(OpBuilder &b, LinalgOp consumer,
                                     unsigned consumerIdx,
                                     const LinalgDependenceGraph &graph,
                                     OperationFolder *folder = nullptr);
 
 /// Fuse linalg operation on tensors, with the producer of the operand at
 /// position `consumerIdx` of the consumer.
 Optional<SmallVector<Value, 1>>
 fuseTensorOps(PatternRewriter &rewriter, Operation *consumer,
-              unsigned consumerIdx, OperationFolder *folder = nullptr,
-              bool useReshapeFusionByExpansion = false);
+              unsigned consumerIdx, OperationFolder *folder = nullptr);
 
 /// Returns the linearized list of all shape dimensions in a `linalgOp`.
 /// Applying the inverse, concatenated loopToOperandRangeMaps to this list
 /// allows the derivation of loop ranges for any linalgOp.
 SmallVector<Value, 8> getShape(OpBuilder &builder, LinalgOp linalgOp);
 template <typename ConcreteOpTy>
 SmallVector<Value, 8> getShape(OpBuilder &builder, ConcreteOpTy linalgOp) {
   return getShape(builder, cast<linalg::LinalgOp>(linalgOp.getOperation()));

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

if (areReshapeOpsFoldable(reshapeOp.getResultType(), reshapeOp.getSrcType(),
       rewriter.replaceOpWithNewOp<ReshapeOpTy>(
           reshapeOp, reshapeOp.getResultType(), srcReshapeOp.src(),
           collapseReassociationMaps(reshapeOp.getReassociationMaps(),
                                     srcReshapeOp.getReassociationMaps(),
                                     rewriter.getContext()));
       return success();
     }
     // Check if producer and consumer are both collapsing dims.
-    else if (areReshapeOpsFoldable(srcReshapeOp.getSrcType(),
-                                   reshapeOp.getSrcType(),
-                                   reshapeOp.getResultType())) {
+    if (areReshapeOpsFoldable(srcReshapeOp.getSrcType(), reshapeOp.getSrcType(),
       rewriter.replaceOpWithNewOp<ReshapeOpTy>(
           reshapeOp, reshapeOp.getResultType(), srcReshapeOp.src(),
           collapseReassociationMaps(srcReshapeOp.getReassociationMaps(),
                                     reshapeOp.getReassociationMaps(),
                                     rewriter.getContext()));
       return success();
     }
     return failure();

mlir/lib/Dialect/Linalg/Transforms/FusionOnTensors.cpp

static AffineMap linearizeCollapsedDims(AffineMap sourceMap,
   }
   return AffineMap::get(sourceMap.getNumDims(), sourceMap.getNumSymbols(),
                         resultExprs, context);
 }
 
 /// Checks if the `reshapeOp` can be fused with it consumer (if `asProducer` is
 /// true) or its producer (if `asProducer` is false) given the indexing map at
 /// its use.
-static bool isTensorReshapeOpFusableByLinearization(TensorReshapeOp reshapeOp,
+static bool isTensorReshapeOpFoldableByLinearization(TensorReshapeOp reshapeOp,
                                                      AffineMap useIndexMap,
                                                      bool asProducer) {
   RankedTensorType returnType = reshapeOp.getResultType();
   RankedTensorType operandType = reshapeOp.getSrcType();
   // Reshape is fusable with its consumer (i.e. reshape as a producer) when its
   // operand is of lesser rank than the result. Fusing when operand has higher
   // rank will require use of mods and divs in the indexing maps of the fused op
   // which would make it non-invertible. Similarly reshape is fused with its
   // producer (i.e. reshape as consumer) only if the return type has lesser
   // rank.
static LinalgOp createLinalgOpOfSameType(LinalgOp op, PatternRewriter &rewriter,
   if (isa<IndexedGenericOp>(op.getOperation()))
     return cast<LinalgOp>(
         rewriter.create<IndexedGenericOp>(args...).getOperation());
   llvm_unreachable(
       "expected only linalg.generic or linalg.indexed_generic ops");
   return nullptr;
 }
 
-/// Conditions for fusing a generic/indexed-generic operation with a reshape op
-/// by expanding the iteration space dimensionality. These are preconditions
-/// assumed by `fusWithReshapeByDimExpansion` which implements the following
-/// fusion pattern.
+/// Conditions for folding a generic/indexed-generic operation with a reshape op
+/// by expanding the iteration space dimensionality for tensor operations. These
+/// are preconditions assumed by `foldReshapeByDimExpansion` which implements
+/// the following fusion pattern.
 ///
 ///  Consider
 ///
 ///  %c = linalg.generic ins(%a, %b : memref<?x?x?xf32>, memref<?x?xf32>)
 ///         indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d0, d2)>,
 ///                          affine_map<(d0, d1, d2) -> (d1, d2)>,
 ///                          affine_map<(d0, d1, d2) -> (d0, d2, d1)>]
 ///  %d = linalg.tensor_reshape %c

 ///       : tensor<?x?x?xf32> into tensor<?x?x?x?xf32>
 ///
 ///  The added reshapes are again expanding patterns, so they will get fused
 ///  with its producers if possible.
 static bool isFusableWithReshapeByDimExpansion(LinalgOp linalgOp,
                                                unsigned fusedTensorIndex) {
   // Is fusable only if:
   // - The linalgOp is a generic op.
-  // - The linalgOp has tensor semantics with a single output.
   // - All the indexing maps for operands in linalgOp are projected
   //   permutations.
   // - The indexing map at the position representing the fused tensor is a
   //   permutation.
   // - All the loops in linalgOp are parallel loops.
   return isa<GenericOp>(linalgOp.getOperation()) &&
-         linalgOp.hasTensorSemantics() && linalgOp.getNumOutputs() == 1 &&
+         linalgOp.hasTensorSemantics() &&
          llvm::all_of(linalgOp.indexing_maps().getValue().take_front(
                           linalgOp.getNumInputs()),
                       [](Attribute attr) {
                         return attr.cast<AffineMapAttr>()
                             .getValue()
                             .isProjectedPermutation();
                       }) &&
          linalgOp.getIndexingMap(fusedTensorIndex).isPermutation() &&
          llvm::all_of(linalgOp.iterator_types(), [](Attribute attr) {
            return attr.cast<StringAttr>().getValue() ==
                   getParallelIteratorTypeName();
          });
 }
 
+/// Implements the fusion of a tensor_reshape op and a generic/indexed_generic
+/// op as explained in `isFusableWithReshapeByExpansion`. Assumes that those
+/// conditions have been satisfied.
 static Optional<SmallVector<Value, 1>>
 fuseWithReshapeByExpansion(LinalgOp linalgOp, TensorReshapeOp reshapeOp,
                            unsigned fusedTensorIndex, PatternRewriter &rewriter,
                            OperationFolder *folder = nullptr) {
+  assert(isFusableWithReshapeByDimExpansion(linalgOp, fusedTensorIndex) &&
+         "preconditions for fuse operation failed");
   // Check if reshape is expanding or collapsing.
   bool isExpanding =
       reshapeOp.getSrcType().getRank() < reshapeOp.getResultType().getRank();
   RankedTensorType expandedType =
       isExpanding ? reshapeOp.getResultType() : reshapeOp.getSrcType();
   RankedTensorType foldedType =
       isExpanding ? reshapeOp.getSrcType() : reshapeOp.getResultType();
   AffineMap fusedIndexMap = linalgOp.getIndexingMap(fusedTensorIndex);
for (auto result : llvm::enumerate(linalgOp.getOperation()->getResults())) {
     }
   }
   // Assuming a single result.
   return resultVals;
 }
 
 namespace {
 
-/// Implementation of fusion on tensor ops when producer is a TensorReshapeOp.
-struct FuseTensorReshapeOpAsProducerByLinearization {
-  static bool isFusable(TensorReshapeOp producer, LinalgOp consumer,
-                        unsigned consumerIdx) {
-    return isa<GenericOp, IndexedGenericOp>(consumer.getOperation()) &&
-           consumer.hasTensorSemantics() &&
-           isTensorReshapeOpFusableByLinearization(
-               producer, consumer.getInputIndexingMap(consumerIdx),
-               /*asProducer=*/true);
-  }
-
-  static Optional<SmallVector<Value, 1>>
-  fuse(TensorReshapeOp producer, LinalgOp consumer, unsigned consumerIdx,
-       PatternRewriter &rewriter, OperationFolder *folder = nullptr) {
-    if (producer.src().getDefiningOp<ConstantOp>())
-      return llvm::None;
+/// Pattern to fold tensor_reshape op with its consumer by using the source of
+/// the reshape op as the operand in the consumer (instead of the result of the
+/// tensor_reshapeop) when the tensor_reshape op is collapsing. The
+/// corresponding index map in the consumer needs to be modified to linearize
+/// the folded dimension.
+///
+/// For example,
+///
+/// #map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+/// %0 = linalg.tensor_reshape %arg0
+///        [affine_map<(i, j, k, l) -> (i)>, affine_map<(i, j, k, l) -> (j, k)>,
+///         affine_map<(i, j, k, l) -> (l)>]
+///      tensor<?x?x?xf32> into tensor<?x?x4x?xf32>
+/// %1 = linalg.generic { indexing_maps = [#map0, #map0, #map0], ... }
+///        ins(%0, %arg1 : tensor<?x?x4x?xf32>, tensor<?x?x4x?xf32>) ...
+///        -> tensor<?x?x4x?xf32>
+///
+/// can be folded into
+///
+/// #map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1 * 4 + d2, d3)>
+/// #map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+/// %0 = linalg.generic { indexing_maps = [#map0, #map1, #map1] ... }
+///        ins(%arg0, %arg1 : tensor<?x?x?xf32>, tensor<?x?x4x?xf32>) ...
+///        -> tensor<?x?x4x?xf32>
+template <typename LinalgOpTy>
+struct FoldProducerReshapeOpByLinearization
+    : public OpRewritePattern<LinalgOpTy> {
+  using OpRewritePattern<LinalgOpTy>::OpRewritePattern;
 
-    if (!isFusable(producer, consumer, consumerIdx))
-      return llvm::None;
+  LogicalResult matchAndRewrite(LinalgOpTy op,
+                                PatternRewriter &rewriter) const override {
+    if (!op.hasTensorSemantics())
+      return failure();
+    LinalgOp linalgOp = cast<LinalgOp>(op.getOperation());
+    for (auto operand : llvm::enumerate(linalgOp.getInputs())) {
+      TensorReshapeOp reshapeOp =
+          operand.value().getDefiningOp<TensorReshapeOp>();
+      if (!reshapeOp ||
+          !isTensorReshapeOpFoldableByLinearization(
+              reshapeOp, linalgOp.getInputIndexingMap(operand.index()),
+              /*asProducer =*/true))
+        continue;
 
       // Compute the fused operands list,
-    SmallVector<Value, 2> fusedOperands(consumer.getInputs());
-    fusedOperands[consumerIdx] = producer.src();
+      SmallVector<Value, 2> fusedOperands(linalgOp.getInputs());
+      fusedOperands[operand.index()] = reshapeOp.src();
 
-    // Compute indexing_maps for the fused operation. The indexing_maps for the
-    // operands of the consumers that arent fused are the same.
-    SmallVector<AffineMap, 4> fusedIndexMaps =
-        llvm::to_vector<4>(llvm::map_range(
-            consumer.indexing_maps(), [](Attribute attr) -> AffineMap {
-              return attr.cast<AffineMapAttr>().getValue();
-            }));
+      // Compute indexing_maps for the fused operation. The indexing_maps for
+      // the operands of the consumers that arent fused are the same.
+      SmallVector<AffineMap, 4> fusedIndexMaps = llvm::to_vector<4>(
+          op.indexing_maps().template getAsValueRange<AffineMapAttr>());
 
       // Accepted consumer maps are either identity or permutation.
-    auto invMap = inversePermutation(fusedIndexMaps[consumerIdx]);
+      auto invMap = inversePermutation(fusedIndexMaps[operand.index()]);
 
-    // Compute the indexing map to use for the operand of the producer.
+      // Compute the indexing map to use for the result of the producer.
       AffineMap modifiedMap =
-        linearizeCollapsedDims(invMap, producer.getResultType().getShape(),
-                               producer.getReassociationMaps());
+          linearizeCollapsedDims(invMap, reshapeOp.getResultType().getShape(),
+                                 reshapeOp.getReassociationMaps());
       for (AffineExpr expr : modifiedMap.getResults()) {
         if (!expr.isPureAffine())
-        return llvm::None;
+          return failure();
       }
-    fusedIndexMaps[consumerIdx] = modifiedMap;
+      fusedIndexMaps[operand.index()] = modifiedMap;
 
       // Further check that the resulting index maps can be fused and
       // inverted. Without this the resultant op is not legal.
       if (!inversePermutation(concatAffineMaps(fusedIndexMaps)))
-      return llvm::None;
+        return op.emitRemark("fused op loop bound computation failed");
 
-    SmallVector<Attribute, 4> indexMapAttrs = llvm::to_vector<4>(
-        llvm::map_range(fusedIndexMaps, [](AffineMap map) -> Attribute {
-          return AffineMapAttr::get(map);
-        }));
-    LinalgOp fusedOp = createLinalgOpOfSameType(
-        consumer, rewriter, rewriter.getUnknownLoc(),
-        consumer.getOperation()->getResultTypes(),
-        /*inputs=*/fusedOperands,
-        /*outputBuffers=*/ValueRange{},
-        /*initTensors=*/ValueRange{}, // no init tensors for now.
-        rewriter.getArrayAttr(indexMapAttrs), consumer.iterator_types(),
-        /*doc=*/nullptr,
-        /*library_call=*/nullptr,
-        /*symbol_source=*/nullptr);
-    auto &fusedRegion = fusedOp.getOperation()->getRegion(0);
-    rewriter.cloneRegionBefore(consumer.getOperation()->getRegion(0),
-                               fusedRegion, fusedRegion.begin());
-    return SmallVector<Value, 1>(fusedOp.getOperation()->getResults());
+      rewriter.startRootUpdate(op);
+      op.getOperation()->setOperands(fusedOperands);
+      op.indexing_mapsAttr(rewriter.getAffineMapArrayAttr(fusedIndexMaps));
+      rewriter.finalizeRootUpdate(op);
+      if (reshapeOp.use_empty())
+        rewriter.eraseOp(reshapeOp);
+      return success();
+    }
+    return op.emitRemark("no fusion candidates found");
   }
 };
 
-struct FuseTensorReshapeOpAsProducerByExpansion {
-  static bool isFusable(TensorReshapeOp producer, LinalgOp consumer,
-                        unsigned consumerIdx) {
-    // Fuse only if
+/// Pattern to fuse a tensor_reshape op with its consumer generic op, when the
+/// reshape op is collapsing dimensions. The dimensionality of the loop in the
+/// consumer generic op is expanded.
+struct FoldWithProducerReshapeOpByExpansion
+    : public OpRewritePattern<GenericOp> {
+  using OpRewritePattern<GenericOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(GenericOp genericOp,
+                                PatternRewriter &rewriter) const override {
+    LinalgOp linalgOp = cast<LinalgOp>(genericOp.getOperation());
+    for (auto operand : llvm::enumerate(linalgOp.getInputs())) {
+      TensorReshapeOp reshapeOp =
+          operand.value().getDefiningOp<TensorReshapeOp>();
+      if (!reshapeOp)
+        continue;
+
+      // Fold only if
       // - The tensor reshape op is folding.
       // - All constraints of fusing with reshape by expansion are met.
-    return producer.getSrcType().getRank() >
-               producer.getResultType().getRank() &&
-           isFusableWithReshapeByDimExpansion(consumer, consumerIdx);
-  }
+      if (reshapeOp.getSrcType().getRank() <
+              reshapeOp.getResultType().getRank() ||
+          !isFusableWithReshapeByDimExpansion(linalgOp, operand.index()))
+        continue;
 
-  static Optional<SmallVector<Value, 1>>
-  fuse(TensorReshapeOp producer, LinalgOp consumer, unsigned consumerIdx,
-       PatternRewriter &rewriter, OperationFolder *folder = nullptr) {
-    if (!isFusable(producer, consumer, consumerIdx))
-      return llvm::None;
-    return fuseWithReshapeByExpansion(consumer, producer, consumerIdx,
-                                      rewriter);
+      Optional<SmallVector<Value, 1>> replacementValues =
+          fuseWithReshapeByExpansion(linalgOp, reshapeOp, operand.index(),
+      if (!replacementValues)
+        return failure();
+      rewriter.replaceOp(genericOp, replacementValues.getValue());
+      if (reshapeOp.use_empty())
+        rewriter.eraseOp(reshapeOp);
+      return success();
+    }
+    return failure();
   }
 };
 
-struct FuseTensorReshapeOpAsConsumerByLinearization {
-  static bool isFusable(LinalgOp producer, TensorReshapeOp consumer,
-                        unsigned consumerIdx) {
-    return isa<GenericOp, IndexedGenericOp>(producer.getOperation()) &&
-           producer.hasTensorSemantics() &&
-           isTensorReshapeOpFusableByLinearization(
-               consumer, producer.getOutputIndexingMap(0),
-               /*asProducer=*/false);
-  }
+/// Pattern to fold tensor_reshape op with its producer. The corresponding index
+/// map in the consumer needs to be modified to linearize the folded dimension.
+struct FoldConsumerReshapeOpByLinearization
+    : public OpRewritePattern<TensorReshapeOp> {
+  using OpRewritePattern<TensorReshapeOp>::OpRewritePattern;
 
-  static Optional<SmallVector<Value, 1>>
-  fuse(LinalgOp producer, TensorReshapeOp consumer, unsigned consumerIdx,
-       PatternRewriter &rewriter, OperationFolder *folder = nullptr) {
-    if (!isFusable(producer, consumer, consumerIdx))
-      return llvm::None;
+  LogicalResult matchAndRewrite(TensorReshapeOp reshapeOp,
+                                PatternRewriter &rewriter) const override {
+    LinalgOp producer = reshapeOp.src().getDefiningOp<LinalgOp>();
+    if (!producer ||
+        !isa<GenericOp, IndexedGenericOp>(producer.getOperation()) ||
+        !producer.hasTensorSemantics() || producer.getNumOutputs() != 1 ||
+        !isTensorReshapeOpFoldableByLinearization(
+            reshapeOp, producer.getOutputIndexingMap(0), /*asProducer =*/false))
+      return failure();
     // The indexing_maps for the operands of the fused operation are same as
     // those for the operands of the producer.
-    SmallVector<AffineMap, 4> fusedIndexMaps =
-        llvm::to_vector<4>(llvm::map_range(
-            producer.indexing_maps(), [](Attribute attr) -> AffineMap {
-              return attr.cast<AffineMapAttr>().getValue();
-            }));
+    SmallVector<AffineMap, 4> fusedIndexMaps = llvm::to_vector<4>(
+        producer.indexing_maps().getAsValueRange<AffineMapAttr>());
 
     auto invMap = inversePermutation(producer.getOutputIndexingMap(0));
 
     // Compute the indexing map to use for the operand of the producer.
     AffineMap modifiedMap =
-        linearizeCollapsedDims(invMap, consumer.getSrcType().getShape(),
-                               consumer.getReassociationMaps());
+        linearizeCollapsedDims(invMap, reshapeOp.getSrcType().getShape(),
+                               reshapeOp.getReassociationMaps());
     for (AffineExpr expr : modifiedMap.getResults()) {
       if (!expr.isPureAffine())
-        return llvm::None;
+        return reshapeOp.emitRemark("fused op indexing map is not affine");
     }
     fusedIndexMaps.back() = modifiedMap;
 
     // Further check that the resulting index maps can be fused and
     // inverted. Without this the resultant op is not legal.
     if (!inversePermutation(concatAffineMaps(fusedIndexMaps)))
-      return llvm::None;
-
-    SmallVector<Attribute, 4> indexMapAttrs = llvm::to_vector<4>(
-        llvm::map_range(fusedIndexMaps, [](AffineMap map) -> Attribute {
-          return AffineMapAttr::get(map);
-        }));
+      return reshapeOp.emitRemark("fused op loop bound computation failed");
 
     LinalgOp fusedOp = createLinalgOpOfSameType(
-        producer, rewriter, rewriter.getUnknownLoc(), consumer.getResultType(),
+        producer, rewriter, rewriter.getUnknownLoc(), reshapeOp.getResultType(),
         /*inputs=*/producer.getInputs(),
         /*outputBuffers=*/ValueRange{},
         /*initTensors=*/ValueRange{}, // no init tensors for now.
-        rewriter.getArrayAttr(indexMapAttrs), producer.iterator_types(),
+        rewriter.getAffineMapArrayAttr(fusedIndexMaps),
+        producer.iterator_types(),
         /*doc=*/nullptr,
         /*library_call=*/nullptr,
         /*symbol_source=*/nullptr);
     auto &fusedRegion = fusedOp.getOperation()->getRegion(0);
     rewriter.cloneRegionBefore(producer.getOperation()->getRegion(0),
                                fusedRegion, fusedRegion.begin());
-    return SmallVector<Value, 1>(fusedOp.getOperation()->getResults());
+    rewriter.replaceOp(reshapeOp, fusedOp.getOperation()->getResults());
+    if (producer.use_empty())
+      rewriter.eraseOp(producer);
+    return success();
   }
 };
 
-/// Implementation of fusion on tensor ops when consumer is a TensorReshapeOp.
-struct FuseTensorReshapeOpAsConsumerByExpansion {
-  static bool isFusable(LinalgOp producer, TensorReshapeOp consumer,
-                        unsigned consumerIdx) {
-    // Fuse only if
+/// Pattern to fold a tensor_reshape op with its producer generic op if the
+/// tensor_reshape op is expanding, by expanding the dimensionality of the loop
+/// in the producer op.
+struct FoldReshapeWithGenericOpByExpansion
+    : public OpRewritePattern<TensorReshapeOp> {
+  using OpRewritePattern<TensorReshapeOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(TensorReshapeOp reshapeOp,
+                                PatternRewriter &rewriter) const override {
+    // Fold only if
     // - The tensor reshape op is a expanding case.
     // - All constraints of fusing with reshape by expansion are met.
-    return consumer.getSrcType().getRank() <
-               consumer.getResultType().getRank() &&
-           isFusableWithReshapeByDimExpansion(producer,
-                                              producer.getNumInputs());
-  }
-
-  static Optional<SmallVector<Value, 1>>
-  fuse(LinalgOp producer, TensorReshapeOp consumer, unsigned consumerIdx,
-       PatternRewriter &rewriter, OperationFolder *folder = nullptr) {
-    if (!isFusable(producer, consumer, consumerIdx))
-      return llvm::None;
-    return fuseWithReshapeByExpansion(
-        producer, consumer, producer.getNumInputs(), rewriter, folder);
+    if (reshapeOp.getSrcType().getRank() > reshapeOp.getResultType().getRank())
+      return failure();
+    LinalgOp producer = reshapeOp.src().getDefiningOp<LinalgOp>();
+    if (!producer || producer.getNumOutputs() != 1 ||
+        !isFusableWithReshapeByDimExpansion(producer, producer.getNumInputs()))
+      return failure();
+    Optional<SmallVector<Value, 1>> replacementValues =
+        fuseWithReshapeByExpansion(producer, reshapeOp, producer.getNumInputs(),
+                                   rewriter);
+    if (!replacementValues)
+      return failure();
+    rewriter.replaceOp(reshapeOp, replacementValues.getValue());
+    if (producer.use_empty())
+      rewriter.eraseOp(producer);
+    return success();
   }
 };
 
-/// Implementation of fusion on tensor ops when producer is a splat constant.
-struct FuseConstantOpAsProducer {
-  static bool isFusable(ConstantOp producer, LinalgOp consumer,
-                        unsigned consumerIdx) {
-    return isa<GenericOp, IndexedGenericOp>(consumer.getOperation()) &&
-           consumer.hasTensorSemantics() &&
-           producer.getResult().getType().isa<RankedTensorType>() &&
-           producer.value().cast<DenseElementsAttr>().isSplat();
-  }
+/// Pattern to fold a GenericOp/IndexedGenericOp with a splat constant.
+template <typename LinalgOpTy>
+struct FoldSplatConstants : public OpRewritePattern<LinalgOpTy> {
+  using OpRewritePattern<LinalgOpTy>::OpRewritePattern;
 
-  static Optional<SmallVector<Value, 1>>
-  fuse(ConstantOp producer, LinalgOp consumer, unsigned consumerIdx,
-       PatternRewriter &rewriter, OperationFolder *folder = nullptr) {
-    if (!isFusable(producer, consumer, consumerIdx))
-      return llvm::None;
+  LogicalResult matchAndRewrite(LinalgOpTy op,
+                                PatternRewriter &rewriter) const override {
+    if (!op.hasTensorSemantics())
+      return failure();
+    LinalgOp linalgOp = cast<LinalgOp>(op.getOperation());
+    for (auto operand : llvm::enumerate(linalgOp.getInputs())) {
+      ConstantOp constantOp = operand.value().getDefiningOp<ConstantOp>();
+      if (!constantOp ||
+          !constantOp.value().cast<DenseElementsAttr>().isSplat())
+        continue;
 
       // The indexing_maps for the operands of the fused operation are same as
-    // those for the operands of the consumer without the indexing map at
-    // consumerIdx
-    SmallVector<AffineMap, 4> fusedIndexMaps =
-        llvm::to_vector<4>(llvm::map_range(
-            consumer.indexing_maps(), [](Attribute attr) -> AffineMap {
-              return attr.cast<AffineMapAttr>().getValue();
-            }));
-    fusedIndexMaps.erase(std::next(fusedIndexMaps.begin(), consumerIdx));
-
-    // The operands list is same as the consumer with the argument for constant
-    // index dropped.
-    SmallVector<Value, 4> fusedOperands(consumer.getInputs());
-    fusedOperands.erase(std::next(fusedOperands.begin(), consumerIdx));
+      // those for the operands of the linalgOp without the indexing map at
+      // operand.index()
+      SmallVector<AffineMap, 4> fusedIndexMaps = llvm::to_vector<4>(
+          linalgOp.indexing_maps().getAsValueRange<AffineMapAttr>());
+      fusedIndexMaps.erase(std::next(fusedIndexMaps.begin(), operand.index()));
+
+      // The operands list is same as the linalgOp with the argument for
+      // constant index dropped.
+      SmallVector<Value, 4> fusedOperands(linalgOp.getInputs());
+      fusedOperands.erase(std::next(fusedOperands.begin(), operand.index()));
 
       // Create a constant scalar value from the splat constant.
       Value scalarConstant = rewriter.create<ConstantOp>(
-        producer.getLoc(),
-        producer.value().cast<DenseElementsAttr>().getSplatValue());
+          constantOp.getLoc(),
+          constantOp.value().cast<DenseElementsAttr>().getSplatValue());
 
       LinalgOp fusedOp = createLinalgOpOfSameType(
-        consumer, rewriter, rewriter.getUnknownLoc(),
-        consumer.getOperation()->getResultTypes(),
+          linalgOp, rewriter, rewriter.getUnknownLoc(),
+          linalgOp.getOperation()->getResultTypes(),
           /*inputs=*/fusedOperands,
           /*outputBuffers=*/ValueRange{},
           /*initTensors=*/ValueRange{}, // no init tensors for now.
           rewriter.getAffineMapArrayAttr(fusedIndexMaps),
-        consumer.iterator_types(),
+          linalgOp.iterator_types(),
           /*doc=*/nullptr,
           /*library_call=*/nullptr,
           /*symbol_source=*/nullptr);
 
       // Map the block argument corresponding to the replaced argument with the
       // scalar constant.
-    Region &consumerRegion = consumer.getOperation()->getRegion(0);
-    Block &entryBlock = *consumerRegion.begin();
-    unsigned argIndex =
-        entryBlock.getNumArguments() - consumer.getNumInputs() + consumerIdx;
+      Region &linalgOpRegion = linalgOp.getOperation()->getRegion(0);
+      Block &entryBlock = *linalgOpRegion.begin();
+      unsigned argIndex = entryBlock.getNumArguments() -
+                          linalgOp.getNumInputs() + operand.index();
       BlockAndValueMapping mapping;
       mapping.map(entryBlock.getArgument(argIndex), scalarConstant);
       Region &fusedRegion = fusedOp.getOperation()->getRegion(0);
-    rewriter.cloneRegionBefore(consumerRegion, fusedRegion, fusedRegion.begin(),
-                               mapping);
-    return SmallVector<Value, 1>(fusedOp.getOperation()->getResults());
+      rewriter.cloneRegionBefore(linalgOpRegion, fusedRegion,
+                                 fusedRegion.begin(), mapping);
+      rewriter.replaceOp(linalgOp, fusedOp.getOperation()->getResults());
+      if (constantOp.use_empty())
+        rewriter.eraseOp(constantOp);
+      return success();
+    }
+    return failure();
   }
 };
 } // namespace
 
 Optional<SmallVector<Value, 1>>
 mlir::linalg::fuseTensorOps(PatternRewriter &rewriter, Operation *consumer,
-                            unsigned consumerIdx, OperationFolder *folder,
-                            bool useReshapeFusionByExpansion) {
+                            unsigned consumerIdx, OperationFolder *folder) {
   if (consumerIdx >= consumer->getNumOperands())
     return llvm::None;
   Operation *producer = consumer->getOperand(consumerIdx).getDefiningOp();
   if (!producer || producer->getNumResults() != 1)
     return llvm::None;
 
   // Fuse when consumer is GenericOp or IndexedGenericOp.
-  if (isa<GenericOp, IndexedGenericOp>(consumer)) {
-    if (isa<GenericOp, IndexedGenericOp>(producer))
-      return FuseGenericOpsOnTensors::fuse(cast<LinalgOp>(producer),
-                                           cast<LinalgOp>(consumer),
-                                           consumerIdx, rewriter, folder);
-    if (auto reshapeOpProducer = dyn_cast<TensorReshapeOp>(producer)) {
-      if (useReshapeFusionByExpansion) {
-        return FuseTensorReshapeOpAsProducerByExpansion::fuse(
-            reshapeOpProducer, cast<LinalgOp>(consumer), consumerIdx, rewriter,
-            folder);
-      }
-      return FuseTensorReshapeOpAsProducerByLinearization::fuse(
-          reshapeOpProducer, cast<LinalgOp>(consumer), consumerIdx, rewriter,
-          folder);
-    }
-    if (auto constantOpProducer = dyn_cast<ConstantOp>(producer))
-      return FuseConstantOpAsProducer::fuse(constantOpProducer,
-                                            cast<LinalgOp>(consumer),
-                                            consumerIdx, rewriter, folder);
+  if (!isa<GenericOp, IndexedGenericOp>(consumer) ||
+      !isa<GenericOp, IndexedGenericOp>(producer))
     return llvm::None;
-  }
 
-  if (isa<GenericOp, IndexedGenericOp>(producer)) {
-    // Fuse when consumer is a TensorReshapeOp.
-    if (TensorReshapeOp reshapeOp = dyn_cast<TensorReshapeOp>(consumer)) {
-      if (useReshapeFusionByExpansion) {
-        return FuseTensorReshapeOpAsConsumerByExpansion::fuse(
-            cast<LinalgOp>(producer), reshapeOp, consumerIdx, rewriter, folder);
-      }
-      return FuseTensorReshapeOpAsConsumerByLinearization::fuse(
-          cast<LinalgOp>(producer), reshapeOp, consumerIdx, rewriter, folder);
-    }
-  }
-
-  return llvm::None;
+  return FuseGenericOpsOnTensors::fuse(cast<LinalgOp>(producer),
+                                       cast<LinalgOp>(consumer), consumerIdx,
+                                       rewriter, folder);
 }
 
 namespace {
 /// Patterns to fuse a generic op, with the producer of its operands.
 template <typename LinalgOpTy>
 struct FuseTensorOps : public OpRewritePattern<LinalgOpTy> {
-  FuseTensorOps(MLIRContext *context, bool useReshapeFusionByExpansion = false)
-      : OpRewritePattern<LinalgOpTy>::OpRewritePattern(context),
-        useReshapeFusionByExpansion(useReshapeFusionByExpansion) {}
+  using OpRewritePattern<LinalgOpTy>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(LinalgOpTy op,
                                 PatternRewriter &rewriter) const override {
     // Find the first operand that is defined by another generic op on tensors.
     for (auto operandNum :
          llvm::seq<unsigned>(0, op.getOperation()->getNumOperands())) {
       Operation *producer =
           op.getOperation()->getOperand(operandNum).getDefiningOp();
-      Optional<SmallVector<Value, 1>> fusedOpResults = fuseTensorOps(
-          rewriter, op, operandNum, nullptr, useReshapeFusionByExpansion);
+      if (!producer)
+        continue;
+      Optional<SmallVector<Value, 1>> fusedOpResults =
+          fuseTensorOps(rewriter, op, operandNum);
       if (fusedOpResults) {
         rewriter.replaceOp(op, *fusedOpResults);
-        if (producer && llvm::all_of(producer->getResults(),
-                                     [](Value val) { return val.use_empty(); }))
+        if (producer->use_empty())
           rewriter.eraseOp(producer);
         return success();
       }
     }
     return failure();
   }
-
-private:
-  bool useReshapeFusionByExpansion;
 };
 
 /// Pass that fuses generic ops on tensors. Used only for testing.
 struct FusionOfTensorOpsPass
     : public LinalgFusionOfTensorOpsBase<FusionOfTensorOpsPass> {
-  FusionOfTensorOpsPass() = default;
-  FusionOfTensorOpsPass(const FusionOfTensorOpsPass &) {}
-  FusionOfTensorOpsPass(bool vUseReshapeFusionByExpansion) {
-    this->useReshapeFusionByExpansion = vUseReshapeFusionByExpansion;
+  void runOnOperation() override {
+    OwningRewritePatternList patterns;
+    Operation *op = getOperation();
+    populateLinalgTensorOpsFusionPatterns(op->getContext(), patterns);
+    applyPatternsAndFoldGreedily(op->getRegions(), patterns);
   }
+};
 
+/// Pass to test folding of reshape op with generic/indexed_generic ops by
+/// linearization.
+struct FoldReshapeOpsByLinearizationPass
+    : public LinalgFoldReshapeOpsByLinearizationBase<
+          FoldReshapeOpsByLinearizationPass> {
   void runOnOperation() override {
     OwningRewritePatternList patterns;
     Operation *op = getOperation();
-    populateLinalgTensorOpsFusionPatterns(op->getContext(), patterns,
-                                          useReshapeFusionByExpansion);
+    populateFoldReshapeOpsByLinearizationPatterns(op->getContext(), patterns);
     applyPatternsAndFoldGreedily(op->getRegions(), patterns);
-  };
+  }
 
-  Option<bool> useReshapeFusionByExpansion{
-      *this, "use-reshape-fusion-by-expansion", llvm::cl::init(true),
-      llvm::cl::desc("Enable use of reshape fusion with producer/consumer by "
-                     "expanding the dimensionality of the producer/consumer")};
-};
 } // namespace
 
+void mlir::populateFoldReshapeOpsByLinearizationPatterns(
+    MLIRContext *context, OwningRewritePatternList &patterns) {
+  patterns.insert<FoldProducerReshapeOpByLinearization<GenericOp>,
+                  FoldProducerReshapeOpByLinearization<IndexedGenericOp>,
+                  FoldConsumerReshapeOpByLinearization>(context);
+}
+
+void mlir::populateFoldReshapeOpsByExpansionPatterns(
+    MLIRContext *context, OwningRewritePatternList &patterns) {
+  patterns.insert<FoldReshapeWithGenericOpByExpansion,
+                  FoldWithProducerReshapeOpByExpansion>(context);
+}
+
 void mlir::populateLinalgTensorOpsFusionPatterns(
-    MLIRContext *context, OwningRewritePatternList &patterns,
-    bool useReshapeFusionByExpansion) {
+    MLIRContext *context, OwningRewritePatternList &patterns) {
   patterns.insert<FuseTensorOps<GenericOp>, FuseTensorOps<IndexedGenericOp>,
-                  FuseTensorOps<TensorReshapeOp>>(context,
-                                                  useReshapeFusionByExpansion);
+                  FoldSplatConstants<GenericOp>,
+                  FoldSplatConstants<IndexedGenericOp>>(context);
+  populateFoldReshapeOpsByExpansionPatterns(context, patterns);
+  GenericOp::getCanonicalizationPatterns(patterns, context);
+  IndexedGenericOp::getCanonicalizationPatterns(patterns, context);
+  TensorReshapeOp::getCanonicalizationPatterns(patterns, context);
+}
+
+std::unique_ptr<Pass> mlir::createLinalgFusionOfTensorOpsPass() {
+  return std::make_unique<FusionOfTensorOpsPass>();
 }
 
-std::unique_ptr<Pass>
-mlir::createLinalgFusionOfTensorOpsPass(bool useReshapeFusionByExpansion) {
-  return std::make_unique<FusionOfTensorOpsPass>(useReshapeFusionByExpansion);
+std::unique_ptr<Pass> mlir::createFoldReshapeOpsByLinearizationPass() {
+  return std::make_unique<FoldReshapeOpsByLinearizationPass>();
 }

mlir/test/Dialect/Linalg/fusion-tensor.mlir

^bb0(%arg3: f32, %arg4: f32):       // no predecessors
       linalg.yield %1 : f32
   } -> tensor<f32>
 
   return %1 : tensor<f32>
 }
 
 // -----
 
-// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1 * 4 + d2, d3)>
-// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
-
-#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
-func @generic_op_reshape_producer_fusion(%arg0 : tensor<?x?x?xf32>,
-                                         %arg1 : tensor<?x?x4x?xf32>) ->
-                                         tensor<?x?x4x?xf32>
-{
-  %0 = linalg.tensor_reshape %arg0 [affine_map<(i, j, k, l) -> (i)>,
-                                    affine_map<(i, j, k, l) -> (j, k)>,
-                                    affine_map<(i, j, k, l) -> (l)>] :
-    tensor<?x?x?xf32> into tensor<?x?x4x?xf32>
-  %1 = linalg.generic {
-     indexing_maps = [#map0, #map0, #map0],
-     iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
-      ins(%0, %arg1 : tensor<?x?x4x?xf32>, tensor<?x?x4x?xf32>) {
-    ^bb0(%arg3: f32, %arg4: f32):       // no predecessors
-      %1 = mulf %arg3, %arg4 : f32
-      linalg.yield %1 : f32
-  } -> tensor<?x?x4x?xf32>
-  return %1 : tensor<?x?x4x?xf32>
-}
-
-// CHECK-LABEL: func @generic_op_reshape_producer_fusion
-//       CHECK: linalg.generic
-//  CHECK-SAME:   indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]]
-//   CHECK-NOT: linalg.generic
-
-
-// -----
-
-// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
-// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1 * 20 + d2 * 5 + d3)>
-
-#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
-func @generic_op_reshape_consumer_fusion(%arg0 : tensor<?x?x4x5xf32>,
-                                         %arg1 : tensor<?x?x4x5xf32>) ->
-                                         tensor<?x?xf32>
-{
-  %0 = linalg.generic {
-     indexing_maps = [#map0, #map0, #map0],
-     iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
-      ins(%arg0, %arg1 : tensor<?x?x4x5xf32>, tensor<?x?x4x5xf32>) {
-    ^bb0(%arg3: f32, %arg4: f32):       // no predecessors
-      %1 = mulf %arg3, %arg4 : f32
-      linalg.yield %1 : f32
-  } -> tensor<?x?x4x5xf32>
-  %1 = linalg.tensor_reshape %0 [affine_map<(i, j, k, l) -> (i)>,
-                                 affine_map<(i, j, k, l) -> (j, k, l)>] :
-    tensor<?x?x4x5xf32> into tensor<?x?xf32>
-  return %1 : tensor<?x?xf32>
-}
-
-// CHECK-LABEL: func @generic_op_reshape_consumer_fusion
-//       CHECK: linalg.generic
-//  CHECK-SAME:   indexing_maps = [#[[$MAP0]], #[[$MAP0]], #[[$MAP1]]]
-//   CHECK-NOT: linalg.generic
-
-// -----
-
-#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
-func @generic_op_reshape_consumer_nofusion(%arg0 : tensor<?x?x?x5xf32>,
-                                           %arg1 : tensor<?x?x?x5xf32>) ->
-                                           tensor<?x?xf32>
-{
-  %0 = linalg.generic {
-     indexing_maps = [#map0, #map0, #map0],
-     iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
-      ins(%arg0, %arg1 : tensor<?x?x?x5xf32>, tensor<?x?x?x5xf32>) {
-    ^bb0(%arg3: f32, %arg4: f32):       // no predecessors
-      %1 = mulf %arg3, %arg4 : f32
-      linalg.yield %1 : f32
-  } -> tensor<?x?x?x5xf32>
-  %1 = linalg.tensor_reshape %0 [affine_map<(i, j, k, l) -> (i)>,
-                                 affine_map<(i, j, k, l) -> (j, k, l)>] :
-    tensor<?x?x?x5xf32> into tensor<?x?xf32>
-  return %1 : tensor<?x?xf32>
-}
-
-// CHECK-LABEL: func @generic_op_reshape_consumer_nofusion
-//       CHECK: linalg.tensor_reshape
-
-// -----
-
 #map0 = affine_map<(d0, d1, d2) -> (d0)>
 #map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 func @generic_op_constant_fusion(%arg0 : tensor<5x?x?xf32>) -> tensor<5x?x?xf32>
 {
   %0 = constant dense<42.0> : tensor<5xf32>
   %1 = linalg.generic {
          indexing_maps = [#map0, #map1, #map1],
          iterator_types = ["parallel", "parallel", "parallel"]}

 //      CHECK:   %[[VAL1:.+]] = addi %[[ARG2]], %[[ADD_OPERAND1]] : i32
 //      CHECK:   %[[VAL2:.+]] = subi %[[VAL1]], %[[SUB_OPERAND1]] : i32
 //      CHECK:   %[[ADD_OPERAND2:.+]] = index_cast %[[ARG0]] : index to i32
 //      CHECK:   %[[SUB_OPERAND2:.+]] = index_cast %[[ARG1]] : index to i32
 //      CHECK:   %[[VAL3:.+]] = addi %[[VAL2]], %[[ADD_OPERAND2]] : i32
 //      CHECK:   %[[VAL4:.+]] = subi %[[VAL3]], %[[SUB_OPERAND2]] : i32
 //      CHECK:   linalg.yield %[[VAL4]] : i32
 //   CHECK-NOT: linalg.indexed_generic
-
-// -----
-
-// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1 * 4 + d2, d3)>
-// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
-
-#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
-func @indexed_generic_op_reshape_producer_fusion(%arg0 : tensor<?x?x?xi32>)
-  -> tensor<?x?x4x?xi32> {
-  %0 = linalg.tensor_reshape %arg0 [affine_map<(i, j, k, l) -> (i)>,
-                                    affine_map<(i, j, k, l) -> (j, k)>,
-                                    affine_map<(i, j, k, l) -> (l)>] :
-    tensor<?x?x?xi32> into tensor<?x?x4x?xi32>
-  %1 = linalg.indexed_generic {
-    indexing_maps = [#map0, #map0],
-    iterator_types = ["parallel", "parallel", "parallel", "parallel"] }
-    ins(%0 : tensor<?x?x4x?xi32>) {
-  ^bb0(%arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: i32):       // no predecessors
-    %2 = index_cast %arg2 : index to i32
-    %3 = addi %arg6, %2 : i32
-    linalg.yield %3 : i32
-  } -> tensor<?x?x4x?xi32>
-  return %1 : tensor<?x?x4x?xi32>
-}
-
-// CHECK-LABEL: func @indexed_generic_op_reshape_producer_fusion
-//   CHECK-NOT: linalg.tensor_reshape
-//       CHECK: linalg.indexed_generic
-//  CHECK-SAME:   indexing_maps = [#[[$MAP0]], #[[$MAP1]]]
-//   CHECK-NOT: linalg.tensor_reshape
-
-// -----
-
-// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
-// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1 * 20 + d2 * 5 + d3)>
-
-#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
-func @indexed_generic_op_reshape_consumer_fusion(%arg0 : tensor<?x?x4x5xi32>)
-  -> tensor<?x?xi32> {
-  %0 = linalg.indexed_generic {
-    indexing_maps = [#map0, #map0],
-    iterator_types = ["parallel", "parallel", "parallel", "parallel"] }
-    ins(%arg0 : tensor<?x?x4x5xi32>) {
-  ^bb0(%arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: i32):       // no predecessors
-    %2 = index_cast %arg2 : index to i32
-    %3 = addi %arg6, %2 : i32
-    linalg.yield %3 : i32
-  } -> tensor<?x?x4x5xi32>
-  %1 = linalg.tensor_reshape %0 [affine_map<(i, j, k, l) -> (i)>,
-                                 affine_map<(i, j, k, l) -> (j, k, l)>] :
-    tensor<?x?x4x5xi32> into tensor<?x?xi32>
-  return %1 : tensor<?x?xi32>
-}
-
-// CHECK-LABEL: func @indexed_generic_op_reshape_consumer_fusion
-//   CHECK-NOT: linalg.tensor_reshape
-//       CHECK: linalg.indexed_generic
-//  CHECK-SAME:   indexing_maps = [#[[$MAP0]], #[[$MAP1]]]
-//   CHECK-NOT: linalg.tensor_reshape
-
-// -----
-
-// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1 + d2 * 7)>
-// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
-
-#map0 = affine_map<(d0, d1, d2) -> (d0)>
-#map1 = affine_map<(d0, d1, d2) -> (d1, d2)>
-#map2 = affine_map<(d0, d1, d2) -> (d0, d2, d1)>
-#map3 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
-func @generic_op_021_permultation_reshape_producer_fusion(%arg0 : tensor<3x35xf32>) -> tensor<3x7x5xf32> {
-  %0 = linalg.tensor_reshape %arg0 [#map0, #map1] : tensor<3x35xf32> into tensor<3x5x7xf32>
-  %1 = linalg.generic {indexing_maps = [#map2, #map3], iterator_types = ["parallel", "parallel", "parallel"]} ins(%0 : tensor<3x5x7xf32>) {
-    ^bb0(%arg2: f32):  // no predecessors
-      linalg.yield %arg2 : f32
-    } -> tensor<3x7x5xf32>
-    return %1 : tensor<3x7x5xf32>
-}
-
-// CHECK-LABEL: func @generic_op_021_permultation_reshape_producer_fusion
-//   CHECK-NOT: linalg.tensor_reshape
-//       CHECK: linalg.generic
-//  CHECK-SAME:   indexing_maps = [#[[$MAP0]], #[[$MAP1]]]
-//   CHECK-NOT: linalg.tensor_reshape
-
-// -----
-
-// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d2, d0 * 7 + d1)>
-// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
-
-#map0 = affine_map<(d0, d1, d2) -> (d0)>
-#map1 = affine_map<(d0, d1, d2) -> (d1, d2)>
-#map2 = affine_map<(d0, d1, d2) -> (d1, d2, d0)>
-#map3 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
-func @generic_op_120_permultation_reshape_producer_fusion(%arg0 : tensor<3x35xf32>) -> tensor<5x7x3xf32> {
-  %0 = linalg.tensor_reshape %arg0 [#map0, #map1] : tensor<3x35xf32> into tensor<3x5x7xf32>
-  %1 = linalg.generic {indexing_maps = [#map2, #map3], iterator_types = ["parallel", "parallel", "parallel"]} ins(%0 : tensor<3x5x7xf32>) {
-    ^bb0(%arg2: f32):  // no predecessors
-      linalg.yield %arg2 : f32
-    } -> tensor<5x7x3xf32>
-    return %1 : tensor<5x7x3xf32>
-}
-
-// CHECK-LABEL: func @generic_op_120_permultation_reshape_producer_fusion
-//   CHECK-NOT: linalg.tensor_reshape
-//       CHECK: linalg.generic
-//  CHECK-SAME:   indexing_maps = [#[[$MAP0]], #[[$MAP1]]]
-//   CHECK-NOT: linalg.tensor_reshape
-
-// -----
-
-// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d1, d0 * 7 + d2)>
-// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
-
-#map0 = affine_map<(d0, d1, d2) -> (d0)>
-#map1 = affine_map<(d0, d1, d2) -> (d1, d2)>
-#map2 = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
-#map3 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
-func @generic_op_102_permultation_reshape_producer_fusion(%arg0 : tensor<3x35xf32>) -> tensor<5x3x7xf32> {
-  %0 = linalg.tensor_reshape %arg0 [#map0, #map1] : tensor<3x35xf32> into tensor<3x5x7xf32>
-  %1 = linalg.generic {indexing_maps = [#map2, #map3], iterator_types = ["parallel", "parallel", "parallel"]} ins(%0 : tensor<3x5x7xf32>) {
-    ^bb0(%arg2: f32):  // no predecessors
-      linalg.yield %arg2 : f32
-    } -> tensor<5x3x7xf32>
-    return %1 : tensor<5x3x7xf32>
-}
-
-// CHECK-LABEL: func @generic_op_102_permultation_reshape_producer_fusion
-//   CHECK-NOT: linalg.tensor_reshape
-//       CHECK: linalg.generic
-//  CHECK-SAME:   indexing_maps = [#[[$MAP0]], #[[$MAP1]]]
-//   CHECK-NOT: linalg.tensor_reshape
-
-// -----
-
-// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
-// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d1, d0 * 7 + d2)>
-
-
-#map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
-#map1 = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
-#map2 = affine_map<(d0, d1, d2) -> (d0)>
-#map3 = affine_map<(d0, d1, d2) -> (d1, d2)>
-func @generic_op_102_permultation_reshape_consumer_fusion(%arg0 : tensor<3x5x7xf32>) -> tensor<5x21xf32> {
-  %0 = linalg.generic {indexing_maps = [#map0, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg0 : tensor<3x5x7xf32>) {
-    ^bb0(%arg2: f32):  // no predecessors
-      linalg.yield %arg2 : f32
-  } -> tensor<5x3x7xf32>
-  %1 = linalg.tensor_reshape %0 [#map2, #map3] : tensor<5x3x7xf32> into tensor<5x21xf32>
-  return %1 : tensor<5x21xf32>
-}
-
-// CHECK-LABEL: func @generic_op_102_permultation_reshape_consumer_fusion
-//   CHECK-NOT: linalg.tensor_reshape
-//       CHECK: linalg.generic
-//  CHECK-SAME:   indexing_maps = [#[[$MAP0]], #[[$MAP1]]]
-//   CHECK-NOT: linalg.tensor_reshape

mlir/test/Dialect/Linalg/reshape_fusion.mlir

-// RUN: mlir-opt %s -linalg-fusion-for-tensor-ops=use-reshape-fusion-by-expansion -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -linalg-fusion-for-tensor-ops -split-input-file | FileCheck %s
 
 #map0 = affine_map<(d0, d1, d2) -> (d2, d0, d1)>
 #map1 = affine_map<(d0, d1, d2) -> (d1, d2, d0)>
 func @generic_op_reshape_producer_fusion(%arg0 : tensor<?x?x?x?xf32>,
                                          %arg1 : tensor<?x?x?xf32>) ->
                                          tensor<?x?x?xf32>
 {
   %0 = linalg.tensor_reshape %arg0 [affine_map<(i, j, k, l) -> (i)>,

mlir/test/Dialect/Linalg/reshape_linearization_fusion.mlir

+// RUN: mlir-opt -split-input-file -linalg-fold-reshape-ops-by-linearization %s | FileCheck %s
+
+
+// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1 * 4 + d2, d3)>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+
+#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+func @generic_op_reshape_producer_fusion(%arg0 : tensor<?x?x?xf32>,
+                                         %arg1 : tensor<?x?x4x?xf32>) ->
+                                         tensor<?x?x4x?xf32>
+{
+  %0 = linalg.tensor_reshape %arg0 [affine_map<(i, j, k, l) -> (i)>,
+                                    affine_map<(i, j, k, l) -> (j, k)>,
+                                    affine_map<(i, j, k, l) -> (l)>] :
+    tensor<?x?x?xf32> into tensor<?x?x4x?xf32>
+  %1 = linalg.generic {
+     indexing_maps = [#map0, #map0, #map0],
+     iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
+      ins(%0, %arg1 : tensor<?x?x4x?xf32>, tensor<?x?x4x?xf32>) {
+    ^bb0(%arg3: f32, %arg4: f32):       // no predecessors
+      %1 = mulf %arg3, %arg4 : f32
+      linalg.yield %1 : f32
+  } -> tensor<?x?x4x?xf32>
+  return %1 : tensor<?x?x4x?xf32>
+}
+
+// CHECK-LABEL: func @generic_op_reshape_producer_fusion
+//       CHECK: linalg.generic
+//  CHECK-SAME:   indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]]
+//   CHECK-NOT: linalg.generic
+
+
+// -----
+
+// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1 * 20 + d2 * 5 + d3)>
+
+#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+func @generic_op_reshape_consumer_fusion(%arg0 : tensor<?x?x4x5xf32>,
+                                         %arg1 : tensor<?x?x4x5xf32>) ->
+                                         tensor<?x?xf32>
+{
+  %0 = linalg.generic {
+     indexing_maps = [#map0, #map0, #map0],
+     iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
+      ins(%arg0, %arg1 : tensor<?x?x4x5xf32>, tensor<?x?x4x5xf32>) {
+    ^bb0(%arg3: f32, %arg4: f32):       // no predecessors
+      %1 = mulf %arg3, %arg4 : f32
+      linalg.yield %1 : f32
+  } -> tensor<?x?x4x5xf32>
+  %1 = linalg.tensor_reshape %0 [affine_map<(i, j, k, l) -> (i)>,
+                                 affine_map<(i, j, k, l) -> (j, k, l)>] :
+    tensor<?x?x4x5xf32> into tensor<?x?xf32>
+  return %1 : tensor<?x?xf32>
+}
+
+// CHECK-LABEL: func @generic_op_reshape_consumer_fusion
+//       CHECK: linalg.generic
+//  CHECK-SAME:   indexing_maps = [#[[$MAP0]], #[[$MAP0]], #[[$MAP1]]]
+//   CHECK-NOT: linalg.generic
+
+// -----
+
+#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+func @generic_op_reshape_consumer_nofusion(%arg0 : tensor<?x?x?x5xf32>,
+                                           %arg1 : tensor<?x?x?x5xf32>) ->
+                                           tensor<?x?xf32>
+{
+  %0 = linalg.generic {
+     indexing_maps = [#map0, #map0, #map0],
+     iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
+      ins(%arg0, %arg1 : tensor<?x?x?x5xf32>, tensor<?x?x?x5xf32>) {
+    ^bb0(%arg3: f32, %arg4: f32):       // no predecessors
+      %1 = mulf %arg3, %arg4 : f32
+      linalg.yield %1 : f32
+  } -> tensor<?x?x?x5xf32>
+  %1 = linalg.tensor_reshape %0 [affine_map<(i, j, k, l) -> (i)>,
+                                 affine_map<(i, j, k, l) -> (j, k, l)>] :
+    tensor<?x?x?x5xf32> into tensor<?x?xf32>
+  return %1 : tensor<?x?xf32>
+}
+
+// CHECK-LABEL: func @generic_op_reshape_consumer_nofusion
+//       CHECK: linalg.tensor_reshape
+
+// -----
+
+
+// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1 * 4 + d2, d3)>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+
+#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+func @indexed_generic_op_reshape_producer_fusion(%arg0 : tensor<?x?x?xi32>)
+  -> tensor<?x?x4x?xi32> {
+  %0 = linalg.tensor_reshape %arg0 [affine_map<(i, j, k, l) -> (i)>,
+                                    affine_map<(i, j, k, l) -> (j, k)>,
+                                    affine_map<(i, j, k, l) -> (l)>] :
+    tensor<?x?x?xi32> into tensor<?x?x4x?xi32>
+  %1 = linalg.indexed_generic {
+    indexing_maps = [#map0, #map0],
+    iterator_types = ["parallel", "parallel", "parallel", "parallel"] }
+    ins(%0 : tensor<?x?x4x?xi32>) {
+  ^bb0(%arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: i32):       // no predecessors
+    %2 = index_cast %arg2 : index to i32
+    %3 = addi %arg6, %2 : i32
+    linalg.yield %3 : i32
+  } -> tensor<?x?x4x?xi32>
+  return %1 : tensor<?x?x4x?xi32>
+}
+
+// CHECK-LABEL: func @indexed_generic_op_reshape_producer_fusion
+//   CHECK-NOT: linalg.tensor_reshape
+//       CHECK: linalg.indexed_generic
+//  CHECK-SAME:   indexing_maps = [#[[$MAP0]], #[[$MAP1]]]
+//   CHECK-NOT: linalg.tensor_reshape
+
+// -----
+
+// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1 * 20 + d2 * 5 + d3)>
+
+#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+func @indexed_generic_op_reshape_consumer_fusion(%arg0 : tensor<?x?x4x5xi32>)
+  -> tensor<?x?xi32> {
+  %0 = linalg.indexed_generic {
+    indexing_maps = [#map0, #map0],
+    iterator_types = ["parallel", "parallel", "parallel", "parallel"] }
+    ins(%arg0 : tensor<?x?x4x5xi32>) {
+  ^bb0(%arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: i32):       // no predecessors
+    %2 = index_cast %arg2 : index to i32
+    %3 = addi %arg6, %2 : i32
+    linalg.yield %3 : i32
+  } -> tensor<?x?x4x5xi32>
+  %1 = linalg.tensor_reshape %0 [affine_map<(i, j, k, l) -> (i)>,
+                                 affine_map<(i, j, k, l) -> (j, k, l)>] :
+    tensor<?x?x4x5xi32> into tensor<?x?xi32>
+  return %1 : tensor<?x?xi32>
+}
+
+// CHECK-LABEL: func @indexed_generic_op_reshape_consumer_fusion
+//   CHECK-NOT: linalg.tensor_reshape
+//       CHECK: linalg.indexed_generic
+//  CHECK-SAME:   indexing_maps = [#[[$MAP0]], #[[$MAP1]]]
+//   CHECK-NOT: linalg.tensor_reshape
+
+// -----
+
+// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1 + d2 * 7)>
+// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+
+#map0 = affine_map<(d0, d1, d2) -> (d0)>
+#map1 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d2, d1)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+func @generic_op_021_permultation_reshape_producer_fusion(%arg0 : tensor<3x35xf32>) -> tensor<3x7x5xf32> {
+  %0 = linalg.tensor_reshape %arg0 [#map0, #map1] : tensor<3x35xf32> into tensor<3x5x7xf32>
+  %1 = linalg.generic {indexing_maps = [#map2, #map3], iterator_types = ["parallel", "parallel", "parallel"]} ins(%0 : tensor<3x5x7xf32>) {
+    ^bb0(%arg2: f32):  // no predecessors
+      linalg.yield %arg2 : f32
+    } -> tensor<3x7x5xf32>
+    return %1 : tensor<3x7x5xf32>
+}
+
+// CHECK-LABEL: func @generic_op_021_permultation_reshape_producer_fusion
+//   CHECK-NOT: linalg.tensor_reshape
+//       CHECK: linalg.generic
+//  CHECK-SAME:   indexing_maps = [#[[$MAP0]], #[[$MAP1]]]
+//   CHECK-NOT: linalg.tensor_reshape
+
+// -----
+
+// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d2, d0 * 7 + d1)>
+// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+
+#map0 = affine_map<(d0, d1, d2) -> (d0)>
+#map1 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d1, d2, d0)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+func @generic_op_120_permultation_reshape_producer_fusion(%arg0 : tensor<3x35xf32>) -> tensor<5x7x3xf32> {
+  %0 = linalg.tensor_reshape %arg0 [#map0, #map1] : tensor<3x35xf32> into tensor<3x5x7xf32>
+  %1 = linalg.generic {indexing_maps = [#map2, #map3], iterator_types = ["parallel", "parallel", "parallel"]} ins(%0 : tensor<3x5x7xf32>) {
+    ^bb0(%arg2: f32):  // no predecessors
+      linalg.yield %arg2 : f32
+    } -> tensor<5x7x3xf32>
+    return %1 : tensor<5x7x3xf32>
+}
+
+// CHECK-LABEL: func @generic_op_120_permultation_reshape_producer_fusion
+//   CHECK-NOT: linalg.tensor_reshape
+//       CHECK: linalg.generic
+//  CHECK-SAME:   indexing_maps = [#[[$MAP0]], #[[$MAP1]]]
+//   CHECK-NOT: linalg.tensor_reshape
+
+// -----
+
+// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d1, d0 * 7 + d2)>
+// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+
+#map0 = affine_map<(d0, d1, d2) -> (d0)>
+#map1 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+func @generic_op_102_permultation_reshape_producer_fusion(%arg0 : tensor<3x35xf32>) -> tensor<5x3x7xf32> {
+  %0 = linalg.tensor_reshape %arg0 [#map0, #map1] : tensor<3x35xf32> into tensor<3x5x7xf32>
+  %1 = linalg.generic {indexing_maps = [#map2, #map3], iterator_types = ["parallel", "parallel", "parallel"]} ins(%0 : tensor<3x5x7xf32>) {
+    ^bb0(%arg2: f32):  // no predecessors
+      linalg.yield %arg2 : f32
+    } -> tensor<5x3x7xf32>
+    return %1 : tensor<5x3x7xf32>
+}
+
+// CHECK-LABEL: func @generic_op_102_permultation_reshape_producer_fusion
+//   CHECK-NOT: linalg.tensor_reshape
+//       CHECK: linalg.generic
+//  CHECK-SAME:   indexing_maps = [#[[$MAP0]], #[[$MAP1]]]
+//   CHECK-NOT: linalg.tensor_reshape
+
+// -----
+
+// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d1, d0 * 7 + d2)>
+
+
+#map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d0)>
+#map3 = affine_map<(d0, d1, d2) -> (d1, d2)>
+func @generic_op_102_permultation_reshape_consumer_fusion(%arg0 : tensor<3x5x7xf32>) -> tensor<5x21xf32> {
+  %0 = linalg.generic {indexing_maps = [#map0, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg0 : tensor<3x5x7xf32>) {
+    ^bb0(%arg2: f32):  // no predecessors
+      linalg.yield %arg2 : f32
+  } -> tensor<5x3x7xf32>
+  %1 = linalg.tensor_reshape %0 [#map2, #map3] : tensor<5x3x7xf32> into tensor<5x21xf32>
+  return %1 : tensor<5x21xf32>
+}
+
+// CHECK-LABEL: func @generic_op_102_permultation_reshape_consumer_fusion
+//   CHECK-NOT: linalg.tensor_reshape
+//       CHECK: linalg.generic
+//  CHECK-SAME:   indexing_maps = [#[[$MAP0]], #[[$MAP1]]]
+//   CHECK-NOT: linalg.tensor_reshape