[mlir][Linalg] Rethink fusion ot linalg ops with reshape ops.

[mravishankar]

Summary

The current fusion on tensors fuses reshape ops with generic ops by
linearizing the indexing maps of the fused tensor in the generic
op. This has some limitations

• It only works for static shapes
• The resulting indexing map has a linearization that would be potentially prevent fusion later on (for ex. tile + fuse).

Instead, try to fuse the reshape consumer (producer) with generic op
producer (consumer) by expanding the dimensionality of the generic op
when the reshape is expanding (folding).
This approach conflicts with the linearization approach. The use of
the linearization approach is still kept as the default, but the
approach added in this change will be made the default after further
experimentation.

[asaadaldien]

First pass...

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

325–326

nit: explicitly return if(!reshapeOp.getSrcType().hasStaticShape()) return false; to simplify condition below.

385

I wonder the case when reshape is a collapsing dims, how do we avoid linearization ?

398

Very nice! and the assumption here pushing reshapes up in the chain the top most reshape will be folded away (e.g IREE or in general reshape of arguments)

447

nit: s/folded/collapsed

[nicolasvasilache]

I am unclear why this is implemented as part of fusion ?
My thinking was that reshapes that expand are pulled upwards, reshapes that contract are pushed downwards.
Applying this reshape folding onto producer and consumer ops enables tiling and fusion in higher-D.

Inside the tiles, depending on the tile sizes and the contiguity aspects, we may or may not be able to just "reshape-without-copy" to convert a matmul in 4-D form to a regular 2-D matmul (think fusing the output of the last conv layer with the following linear layer).

Can we revisit this to just introduce 2 patterns for each case and keep fusion separate from this?

[mravishankar]

In D89002#2318711, @nicolasvasilache wrote:

I am unclear why this is implemented as part of fusion ?
My thinking was that reshapes that expand are pulled upwards, reshapes that contract are pushed downwards.
Applying this reshape folding onto producer and consumer ops enables tiling and fusion in higher-D.

Inside the tiles, depending on the tile sizes and the contiguity aspects, we may or may not be able to just "reshape-without-copy" to convert a matmul in 4-D form to a regular 2-D matmul (think fusing the output of the last conv layer with the following linear layer).

Can we revisit this to just introduce 2 patterns for each case and keep fusion separate from this?

We can do this as folding, but the "fusion" added here conflicts with the fusion that uses "linearization" by folding a collapsing reshape with its producer by linearizing the collapsed dimensions in the indexing map of the fused tensor in the producer. This is opposite to the pattern here where collapsing reshapes get fused with its consumer. (In the previous two statements you can replace collapsing with expanding and producer with consumer and the statement still holds). So I think it makes sense to drop the linearization approach. Just trying to stage this here. Here is what I plan

1. Land this change in its current state
2. The reshape + (generic/indexed-generic) op fusion can be replace with patterns. Both for the expansion approach here and the linearization approach. The "expansion' pattern can be moved into a folding pattern with the linearization patterns dropped from the fusion pass here. I just need to verify that this works in IREE (which is probably the main user of these fusions). So far it seems to be the case, but the pain will be lesser if the switch is staged a bit

WDYT?

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

385

They conflict. So eventually want to deprecate linearization. No one is really using it, so maybe I will refactor these into separate patterns (they kind of already are). Will wait to get resolution on Nicolas' comment.

398

Thats right. The "collapsing" reshapes get pushed down the function to the returns, and the "expanding" reshape get pushed all the way to the function arguments.

[asaadaldien]

L(very)GTM, thanks Mahesh!. I will leave the pattern splitting/refactoring discussion between you and Nicolas.

[nicolasvasilache]

SGTM, I am fine with evolving after this CL.

[nicolasvasilache]

Just added a small comment.
I am tight on time this week and will not have time to review more deeply.
I'll take a deeper pass when this gets rewritten as independent patterns.

Thanks for pushing this!

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

553

There is a clone in the interface LinalgStructuredOpsInterface.td.
Can we (improve and) reuse that ?

[mravishankar]

Rebase

[mravishankar]

In D89002#2321837, @nicolasvasilache wrote:

Just added a small comment.
I am tight on time this week and will not have time to review more deeply.
I'll take a deeper pass when this gets rewritten as independent patterns.

Thanks for pushing this!

I added a couple of patches on top of this that does the refactoring to use patterns (particularly : https://reviews.llvm.org/D89201). PTAL. I would like to squash and submit them in one go.

[nicolasvasilache]

Please go ahead and squash merge once child comments are addressed.
Thanks for these improvements!

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

let builders = [
     OpBuilder<"OpBuilder &b, OperationState &result, Type resultType, "
               "Value src, ArrayRef<ReassociationExprs> reassociation, "
               "ArrayRef<NamedAttribute> attrs = {}">,
     OpBuilder<"OpBuilder &b, OperationState &result, Type resultType, "
               "Value src, ArrayRef<ReassociationIndices> reassociation, "
               "ArrayRef<NamedAttribute> attrs = {}", [{
       auto reassociationMaps =
           convertReassociationIndicesToMaps(b, reassociation);
-      build(b, result, src, reassociationMaps, attrs);
+      build(b, result, resultType, src, reassociationMaps, attrs);
     }]>
   ];
 
   code commonExtraClassDeclaration = [{
     static StringRef getReassociationAttrName() { return "reassociation"; }
     SmallVector<AffineMap, 4> getReassociationMaps() {
       return llvm::to_vector<4>(llvm::map_range(reassociation(), [
       ](Attribute a) { return a.cast<AffineMapAttr>().getValue(); }));

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();
+std::unique_ptr<Pass>
+createLinalgFusionOfTensorOpsPass(bool useReshapeFusionByExpansion = false);
 
 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>>

 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 for fusing linalg operation on tensors.
-void populateLinalgTensorOpsFusionPatterns(MLIRContext *context,
-                                           OwningRewritePatternList &patterns);
+void populateLinalgTensorOpsFusionPatterns(
+    MLIRContext *context, OwningRewritePatternList &patterns,
+    bool useReshapeFusionByExpansion = false);
 
 /// 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/Utils/Utils.h

 /// method is called.
 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.
-Operation *fuseTensorOps(PatternRewriter &rewriter, Operation *consumer,
-                         unsigned consumerIdx,
-                         OperationFolder *folder = nullptr);
+Optional<SmallVector<Value, 1>>
+fuseTensorOps(PatternRewriter &rewriter, Operation *consumer,
+              unsigned consumerIdx, OperationFolder *folder = nullptr,
+              bool useReshapeFusionByExpansion = false);
 
 /// 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/Transforms/FusionOnTensors.cpp

 
 using namespace mlir;
 using namespace mlir::linalg;
 
 namespace {
 
 /// Implementation of fusion of generic ops and indexed_generic ops.
 struct FuseGenericOpsOnTensors {
-  static bool isFusible(LinalgOp producer, LinalgOp consumer,
+  static bool isFusable(LinalgOp producer, LinalgOp consumer,
                         unsigned consumerIdx) {
     // Producer and consumer must have tensor semantics.
     if (!producer.hasTensorSemantics() || !consumer.hasTensorSemantics())
       return false;
 
     // Verify that
     // - the producer has all "parallel" iterator type.
     if (producer.getNumParallelLoops() != producer.getNumLoops())
       return false;
 
     // Get the consumer index map. The number of results of the consumer index
     // map must match the number of loops of the producer.
     AffineMap consumerIndexMap = consumer.getIndexingMap(consumerIdx);
     if (consumerIndexMap.getNumResults() != producer.getNumLoops())
       return false;
 
     // Finally the index_map for the result must be invertible. For now just
     // verify it is a permutation.
     AffineMap producerResultIndexMap = producer.getOutputIndexingMap(0);
     return producerResultIndexMap.isPermutation();
   }
 
-  static LinalgOp fuse(LinalgOp producer, LinalgOp consumer,
-                       unsigned consumerIdx, PatternRewriter &rewriter,
-                       OperationFolder *folder = nullptr) {
-    if (!isFusible(producer, consumer, consumerIdx))
-      return nullptr;
+  static Optional<SmallVector<Value, 1>>
+  fuse(LinalgOp producer, LinalgOp consumer, unsigned consumerIdx,
+       PatternRewriter &rewriter, OperationFolder *folder = nullptr) {
+    if (!isFusable(producer, consumer, consumerIdx))
+      return llvm::None;
 
-    unsigned numFusedOperands = producer.getOperation()->getNumOperands() +
-                                consumer.getOperation()->getNumOperands() - 1;
+    unsigned numFusedOperands =
+        producer.getNumInputs() + consumer.getNumInputs() - 1;
 
     // Compute the fused operands list,
     SmallVector<Value, 2> fusedOperands;
     fusedOperands.reserve(numFusedOperands);
-    auto consumerOperands = consumer.getOperation()->getOperands();
-    auto producerOperands = producer.getOperation()->getOperands();
+    auto consumerOperands = consumer.getInputs();
+    auto producerOperands = producer.getInputs();
     fusedOperands.assign(consumerOperands.begin(),
                          std::next(consumerOperands.begin(), consumerIdx));
     fusedOperands.append(producerOperands.begin(), producerOperands.end());
     fusedOperands.append(std::next(consumerOperands.begin(), consumerIdx + 1),
                          consumerOperands.end());
 
     // Compute indexing_maps for the fused operation. The indexing_maps for the
     // operands of the consumers that arent fused are the same. The
     // indexing_maps for the producers need to be computed based on the
     // indexing_map of the operand at consumerIdx in the consumer.
     SmallVector<Attribute, 4> fusedIndexMaps;
     auto consumerIndexMaps = consumer.indexing_maps();
-    fusedIndexMaps.reserve(fusedOperands.size() +
-                           consumer.getOperation()->getNumResults());
+    fusedIndexMaps.reserve(fusedOperands.size() + consumer.getNumOutputs());
     fusedIndexMaps.assign(consumerIndexMaps.begin(),
                           std::next(consumerIndexMaps.begin(), consumerIdx));
     // Compute indexing maps for the producer args in the fused operation.
     computeProducerOperandIndex(
         producer, consumer.getInputIndexingMap(consumerIdx), fusedIndexMaps);
 
     // Append the indexing maps for the remaining consumer operands.
     fusedIndexMaps.append(std::next(consumerIndexMaps.begin(), consumerIdx + 1),
assert(invProducerResultIndexMap &&
            "expected producer result indexig map to be invertible");
     // consumer loop -> producer loop
     AffineMap consumerToProducerLoopsMap =
         invProducerResultIndexMap.compose(consumerResultIndexMap);
 
     generateFusedRegion(rewriter, fusedOp, producer, consumer,
                         consumerToProducerLoopsMap, consumerIdx,
                         consumer.getNumLoops());
-    return fusedOp;
+    return SmallVector<Value, 1>(fusedOp.getOperation()->getResults());
   }
 
 private:
   /// Append to `fusedOpIndexingMapAttrs` the indexing maps for the operands of
   /// the `producer` to use in the fused operation given the indexing map of the
   /// result of the producer in the consumer.
   static void computeProducerOperandIndex(
       LinalgOp producer, AffineMap fusedConsumerArgIndexMap,
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 isTensorReshapeOpFusible(TensorReshapeOp reshapeOp,
-                                     AffineMap useIndexMap, bool asProducer) {
+static bool isTensorReshapeOpFusableByLinearization(TensorReshapeOp reshapeOp,
+                                                    AffineMap useIndexMap,
+                                                    bool asProducer) {
   RankedTensorType returnType = reshapeOp.getResultType();
   RankedTensorType operandType = reshapeOp.getSrcType();
-  // Reshape is fusible with its consumer (i.e. reshape as a producer) when its
+  // 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.
-  if ((asProducer && returnType.getRank() < operandType.getRank()) ||
-      (!asProducer && operandType.getRank() < returnType.getRank()))
+  if ((asProducer && reshapeOp.getSrcType().hasStaticShape() &&
+       returnType.getRank() < operandType.getRank()) ||
+      (!asProducer && reshapeOp.getResultType().hasStaticShape() &&
+       operandType.getRank() < returnType.getRank()))
     return false;
   return useIndexMap.isPermutation();
 }
 
 /// Based on the type of `op` create a linalg op of the same type, i.e. if `op`
 /// is a linalg.generic operation, the create a `linalg.generic` operation with
 /// the given `args`. Expects `op` to be `linalg.generic` or
 /// `linalg.indexed_generic`.
 template <typename... Args>
 static LinalgOp createLinalgOpOfSameType(LinalgOp op, PatternRewriter &rewriter,
                                          Args... args) {
   if (isa<GenericOp>(op.getOperation()))
     return cast<LinalgOp>(rewriter.create<GenericOp>(args...).getOperation());
   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.
+///
+///  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
+///         [affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1)>,
+///          affine_map<(d0, d1, d2, d3, d4, d5) -> (d2)>,
+///          affine_map<(d0, d1, d2, d3, d4, d5) -> (d3, d4, d5)>]
+///       : tensor<?x?x?xf32> into tensor<?x?x?x?x?x?xf32>
+///
+///  The reshape can be folded into the `linalgOp` if the
+///  generic/indexed-generic op loop dimensionality is increased to match the
+///  result (operand) of the tensor_reshape when the reshape is expanding
+///  (folding). The indexing_map of the fused tensor in the `linalgOp` and the
+///  reassociation map helps compute the indexing maps of the modified op. For
+///  the above example, based on the reassociation map it can be concluded that
+///
+///  - The loop used to access the first dimension of the fused tensor is split
+///    into two.
+///  - The loop used to access the second dimension of the fused tensor is kept
+///    as is.
+///  - The loop used to access the third dimension of the fused tensor is split
+///    into three.
+///
+///  i.e. (e0, e1, e2, e3, e4) is the domain of the indexing map of the modified
+///  op, then
+///
+///   d0 -> e0, e1
+///   d1 -> e2, e3, e4
+///   d2 -> e5
+///
+///  substituting this, the generic op can be rewritten as
+///
+///  %d = linalg.generic ins(%0, %1 : )
+///        indexing_maps =
+///         [affine_map<(e0, e1, e2, e3, e4, e5) -> (e2, e3, e4, e0, e1, e5)>,
+///          affine_map<(e0, e1, e2, e3, e4, e5) -> (e2, e3, e4, e5)>,
+///          affine_map<(e0, e1, e2, e3, e4, e5) -> (e0, e1, e5, e2, e3, e4)>]
+///
+///  Since operands to the linalg generic are now 5D, reshapes can be introduced
+///  to make it consistent
+///
+///  %0 = linalg.tensor_reshape %a
+///         [affine_map<(e0, e1, e2, e3, e4, e5) -> (e0, e1, e2),
+///          affine_map<(e0, e1, e2, e3, e4, e5) -> (e3, e4),
+///          affine_map<(e0, e1, e2, e3, e4, e5) -> (e5)]
+///       : tensor<?x?x?xf32> into tensor<?x?x?x?x?x?xf32>
+///  %1 = linalg.tensor_reshape %b
+///         [affine_map<(e0, e1, e2, e3) -> (e0, e1, e2),
+///          affine_map<(e0, e1, e2, e3) -> (e3)]
+///       : 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 &&
+         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();
+         });
+}
+
+static Optional<SmallVector<Value, 1>>
+fuseWithReshapeByExpansion(LinalgOp linalgOp, TensorReshapeOp reshapeOp,
+                           unsigned fusedTensorIndex, PatternRewriter &rewriter,
+                           OperationFolder *folder = nullptr) {
+  // 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);
+
+  // The reshape is folding/expanding consecutive dimensions. Given the indexing
+  // map of the fused tensor find the number of dimensions each of the loops of
+  // the original op is expanded into. Also record the shape of the expanded
+  // dimensions.
+  ArrayRef<int64_t> expandedShape = expandedType.getShape();
+  SmallVector<unsigned, 4> numFoldedDims(foldedType.getRank(), 0);
+  SmallVector<SmallVector<int64_t, 4>, 4> expandedDimsShape(
+      expandedType.getRank());
+  auto reassociationMaps = reshapeOp.getReassociationMaps();
+  for (auto resultExpr : llvm::enumerate(fusedIndexMap.getResults())) {
+    unsigned pos = resultExpr.value().cast<AffineDimExpr>().getPosition();
+    AffineMap foldedDims = reassociationMaps[resultExpr.index()];
+    numFoldedDims[pos] = foldedDims.getNumResults();
+    ArrayRef<int64_t> shape = expandedShape.slice(
+        foldedDims.getResult(0).cast<AffineDimExpr>().getPosition(),
+        numFoldedDims[pos]);
+    expandedDimsShape[pos].assign(shape.begin(), shape.end());
+  }
+
+  // The remapping of the indices is then the prefix sum (inclusive) of the
+  // numFoldedDims.
+  SmallVector<unsigned, 4> remapping(numFoldedDims.size() + 1, 0);
+  unsigned sum = 0;
+  for (auto numFoldedDim : llvm::enumerate(numFoldedDims)) {
+    sum += numFoldedDim.value();
+    remapping[numFoldedDim.index() + 1] = sum;
+  }
+
+  SmallVector<AffineMap, 4> expandedOpIndexingMaps;
+  // Compute the modified indexing maps by replacing every loop (AffineDimExpr)
+  // in the original indexing map with the sequence of loops that it is expanded
+  // to.
+  for (AffineMap indexingMap : linalgOp.getIndexingMaps()) {
+    SmallVector<AffineExpr, 4> newExprs;
+    for (AffineExpr expr : indexingMap.getResults()) {
+      unsigned pos = expr.cast<AffineDimExpr>().getPosition();
+      for (unsigned newPos :
+           llvm::seq<unsigned>(remapping[pos], remapping[pos + 1])) {
+        newExprs.push_back(rewriter.getAffineDimExpr(newPos));
+      }
+    }
+    expandedOpIndexingMaps.push_back(
+        AffineMap::get(remapping.back(), indexingMap.getNumSymbols(), newExprs,
+                       rewriter.getContext()));
+  }
+
+  // The operands of the expanded op are computed by reshaping the original
+  // operands. The reshape depends on the ordering of the loop used to access
+  // the tensor in the original operation, and are expanded into as many
+  // dimensions as the loop is expanded into (as computed by `remapping`).
+  auto getReshapeInfo =
+      [&](AffineMap operandIndexingMap,
+          SmallVectorImpl<ReassociationIndices> &reassociation,
+          SmallVectorImpl<int64_t> &expandedOpOperandShape) {
+        unsigned reshapeDims = 0;
+        for (AffineExpr expr : operandIndexingMap.getResults()) {
+          unsigned origDim = expr.cast<AffineDimExpr>().getPosition();
+          auto foldedDims = llvm::seq<int64_t>(
+              reshapeDims, reshapeDims + numFoldedDims[origDim]);
+          reassociation.emplace_back(foldedDims.begin(), foldedDims.end());
+          expandedOpOperandShape.append(expandedDimsShape[origDim].begin(),
+                                        expandedDimsShape[origDim].end());
+          reshapeDims += numFoldedDims[origDim];
+        }
+      };
+  SmallVector<Value, 4> expandedOpOperands;
+  for (auto operand : llvm::enumerate(linalgOp.getInputs())) {
+    if (operand.index() == fusedTensorIndex) {
+      expandedOpOperands.push_back(reshapeOp.src());
+      continue;
+    }
+    AffineMap indexingMap = linalgOp.getIndexingMap(operand.index());
+    SmallVector<ReassociationIndices, 4> reassociation;
+    SmallVector<int64_t, 4> expandedOperandShape;
+    getReshapeInfo(indexingMap, reassociation, expandedOperandShape);
+    Type expandedOperandType = RankedTensorType::get(
+        expandedOperandShape,
+        operand.value().getType().cast<ShapedType>().getElementType());
+    if (expandedOperandType != operand.value().getType()) {
+      expandedOpOperands.push_back(rewriter.create<TensorReshapeOp>(
+          linalgOp.getLoc(), expandedOperandType, operand.value(),
+          reassociation));
+    } else {
+      expandedOpOperands.push_back(operand.value());
+    }
+  }
+  SmallVector<Type, 1> resultTypes;
+  SmallVector<SmallVector<ReassociationIndices, 4>, 1> resultReassociation;
+  for (auto result : llvm::enumerate(linalgOp.getOperation()->getResults())) {
+    AffineMap indexingMap =
+        linalgOp.getIndexingMap(linalgOp.getNumInputs() + result.index());
+    SmallVector<ReassociationIndices, 4> reassociation;
+    SmallVector<int64_t, 4> expandedResultShape;
+    getReshapeInfo(indexingMap, reassociation, expandedResultShape);
+    resultTypes.push_back(RankedTensorType::get(
+        expandedResultShape,
+        result.value().getType().cast<ShapedType>().getElementType()));
+    resultReassociation.emplace_back(std::move(reassociation));
+  }
+
+  // The iterator types of the expanded op are all parallel.
+  SmallVector<StringRef, 4> iteratorTypes(remapping.back(),
+                                          getParallelIteratorTypeName());
+
+  LinalgOp fusedOp = createLinalgOpOfSameType(
+      linalgOp, rewriter, linalgOp.getLoc(), resultTypes,
+      /*inputs=*/expandedOpOperands,
+      /*outputBuffers=*/ValueRange{},
+      /*initTensors=*/ValueRange{}, expandedOpIndexingMaps, iteratorTypes);
+  Region &fusedRegion = fusedOp.getOperation()->getRegion(0);
+  // TODO: Add support for indexed generic op, which would need mapping the
+  // expanded dimensions to the original dimension arguments.
+  rewriter.cloneRegionBefore(linalgOp.getOperation()->getRegion(0), fusedRegion,
+                             fusedRegion.begin());
+
+  // Reshape the result values to their original shape if this is a collapsing
+  // reshape folded into its consumer.
+  SmallVector<Value, 1> resultVals;
+  for (auto result : llvm::enumerate(linalgOp.getOperation()->getResults())) {
+    if (!isExpanding &&
+        resultTypes[result.index()] != result.value().getType()) {
+      resultVals.push_back(rewriter.create<TensorReshapeOp>(
+          linalgOp.getLoc(), result.value().getType(),
+          fusedOp.getOperation()->getResult(result.index()),
+          resultReassociation[result.index()]));
+    } else {
+      resultVals.push_back(fusedOp.getOperation()->getResult(result.index()));
+    }
+  }
+  // Assuming a single result.
+  return resultVals;
+}
+
 namespace {
 
 /// Implementation of fusion on tensor ops when producer is a TensorReshapeOp.
-struct FuseTensorReshapeOpAsProducer {
-  static bool isFusible(TensorReshapeOp producer, LinalgOp consumer,
+struct FuseTensorReshapeOpAsProducerByLinearization {
+  static bool isFusable(TensorReshapeOp producer, LinalgOp consumer,
                         unsigned consumerIdx) {
     return isa<GenericOp, IndexedGenericOp>(consumer.getOperation()) &&
            consumer.hasTensorSemantics() &&
-           isTensorReshapeOpFusible(producer,
-                                    consumer.getInputIndexingMap(consumerIdx),
-                                    /*asProducer=*/true);
+           isTensorReshapeOpFusableByLinearization(
+               producer, consumer.getInputIndexingMap(consumerIdx),
   }
 
-  static LinalgOp fuse(TensorReshapeOp producer, LinalgOp consumer,
-                       unsigned consumerIdx, PatternRewriter &rewriter,
-                       OperationFolder *folder = nullptr) {
+  static Optional<SmallVector<Value, 1>>
+  fuse(TensorReshapeOp producer, LinalgOp consumer, unsigned consumerIdx,
+       PatternRewriter &rewriter, OperationFolder *folder = nullptr) {
     if (producer.src().getDefiningOp<ConstantOp>())
-      return nullptr;
+      return llvm::None;
 
-    if (!isFusible(producer, consumer, consumerIdx))
-      return nullptr;
+    if (!isFusable(producer, consumer, consumerIdx))
+      return llvm::None;
 
     // Compute the fused operands list,
-    Operation *consumerOp = consumer.getOperation();
-    SmallVector<Value, 2> fusedOperands(consumerOp->getOperands());
+    SmallVector<Value, 2> fusedOperands(consumer.getInputs());
     fusedOperands[consumerIdx] = producer.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();
             }));
 
     // Accepted consumer maps are either identity or permutation.
     auto invMap = inversePermutation(fusedIndexMaps[consumerIdx]);
 
     // Compute the indexing map to use for the operand of the producer.
     AffineMap modifiedMap =
         linearizeCollapsedDims(invMap, producer.getResultType().getShape(),
                                producer.getReassociationMaps());
     for (AffineExpr expr : modifiedMap.getResults()) {
       if (!expr.isPureAffine())
-        return nullptr;
+        return llvm::None;
     }
     fusedIndexMaps[consumerIdx] = 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 nullptr;
+      return llvm::None;
 
     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(),
-        consumerOp->getResultTypes(),
+        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(consumerOp->getRegion(0), fusedRegion,
-                               fusedRegion.begin());
-    return fusedOp;
+    rewriter.cloneRegionBefore(consumer.getOperation()->getRegion(0),
+                               fusedRegion, fusedRegion.begin());
+    return SmallVector<Value, 1>(fusedOp.getOperation()->getResults());
   }
 };
 
-/// Implementation of fusion on tensor ops when consumer is a TensorReshapeOp.
-struct FuseTensorReshapeOpAsConsumer {
-  static bool isCollapsingAndFusible(LinalgOp producer,
-                                     TensorReshapeOp consumer,
-                                     unsigned consumerIdx) {
+struct FuseTensorReshapeOpAsProducerByExpansion {
+  static bool isFusable(TensorReshapeOp producer, LinalgOp consumer,
+                        unsigned consumerIdx) {
+    // Fuse 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);
+  }
+
+  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);
+  }
+};
+
+struct FuseTensorReshapeOpAsConsumerByLinearization {
+  static bool isFusable(LinalgOp producer, TensorReshapeOp consumer,
     return isa<GenericOp, IndexedGenericOp>(producer.getOperation()) &&
            producer.hasTensorSemantics() &&
-           isTensorReshapeOpFusible(consumer, producer.getOutputIndexingMap(0),
-                                    /*asProducer=*/false);
+           isTensorReshapeOpFusableByLinearization(
+               consumer, producer.getOutputIndexingMap(0),
   }
 
-  static LinalgOp fuseCollapsingCase(LinalgOp producer,
-                                     TensorReshapeOp consumer,
-                                     unsigned consumerIdx,
-                                     PatternRewriter &rewriter) {
+  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;
     // 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();
             }));
 
     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());
     for (AffineExpr expr : modifiedMap.getResults()) {
       if (!expr.isPureAffine())
-        return nullptr;
+        return llvm::None;
     }
     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 nullptr;
+      return llvm::None;
 
     SmallVector<Attribute, 4> indexMapAttrs = llvm::to_vector<4>(
         llvm::map_range(fusedIndexMaps, [](AffineMap map) -> Attribute {
           return AffineMapAttr::get(map);
         }));
 
-    Operation *producerOp = producer.getOperation();
     LinalgOp fusedOp = createLinalgOpOfSameType(
         producer, rewriter, rewriter.getUnknownLoc(), consumer.getResultType(),
-        /*inputs=*/producerOp->getOperands(),
+        /*inputs=*/producer.getInputs(),
         /*outputBuffers=*/ValueRange{},
         /*initTensors=*/ValueRange{}, // no init tensors for now.
         rewriter.getArrayAttr(indexMapAttrs), producer.iterator_types(),
         /*doc=*/nullptr,
         /*library_call=*/nullptr,
         /*symbol_source=*/nullptr);
     auto &fusedRegion = fusedOp.getOperation()->getRegion(0);
-    rewriter.cloneRegionBefore(producerOp->getRegion(0), fusedRegion,
-                               fusedRegion.begin());
-    return fusedOp;
+    rewriter.cloneRegionBefore(producer.getOperation()->getRegion(0),
+                               fusedRegion, fusedRegion.begin());
+    return SmallVector<Value, 1>(fusedOp.getOperation()->getResults());
   }
+};
 
-  static bool isExpandingAndFusible(LinalgOp producer, TensorReshapeOp consumer,
-                                    unsigned consumerIdx) {
-    // Is fusible only if:
-    //   1) The producer is a generic op.
-    //   2) The producer has tensor semantics.
-    //   3) The tensor reshape op is a expanding case.
-    //   4) All the shapes are the same for the generic op.
-    //   5) All the indexing maps in producer are identity.
-    //   6) All the loops in producer are parallel loops.
-    //   7) The producer has a single user.
-    auto types = producer.getInputOutputShapedTypes();
-    assert(!types.empty());
-    return isa<GenericOp>(producer.getOperation()) &&
-           producer.hasTensorSemantics() &&
-           consumer.getSrcType().getRank() <
+/// Implementation of fusion on tensor ops when consumer is a TensorReshapeOp.
+struct FuseTensorReshapeOpAsConsumerByExpansion {
+  static bool isFusable(LinalgOp producer, TensorReshapeOp consumer,
+    // Fuse 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() &&
-           std::equal(types.begin() + 1, types.end(), types.begin()) &&
-           llvm::all_of(producer.getIndexingMaps(),
-                        [](AffineMap map) { return map.isIdentity(); }) &&
-           llvm::all_of(producer.iterator_types(),
-                        [](Attribute attr) {
-                          return attr.cast<StringAttr>().getValue() ==
-                                 getParallelIteratorTypeName();
-                        }) &&
-           producer.getOperation()->hasOneUse();
-  }
-
-  static LinalgOp fuseExpandingCase(LinalgOp producer, TensorReshapeOp consumer,
-                                    unsigned consumerIdx,
-                                    PatternRewriter &rewriter) {
-    Location loc = producer.getLoc();
-    auto dstShape = consumer.getResultType().cast<ShapedType>().getShape();
-    SmallVector<Value, 4> args;
-    for (auto arg : producer.getOperation()->getOperands()) {
-      auto type = RankedTensorType::get(
-          dstShape, arg.getType().cast<ShapedType>().getElementType());
-      args.push_back(rewriter.createOrFold<linalg::TensorReshapeOp>(
-          loc, type, arg, consumer.reassociation()));
-    }
-
-    SmallVector<Type, 4> resultTypes;
-    for (auto t : producer.getOutputTensorTypes()) {
-      Type type = RankedTensorType::get(dstShape,
-                                        t.cast<ShapedType>().getElementType());
-      resultTypes.push_back(type);
-    }
-
-    int rank = dstShape.size();
-    auto genericOp = rewriter.create<linalg::GenericOp>(
-        loc, resultTypes, /*inputs=*/args,
-        /*outputBuffers=*/ValueRange{},
-        /*initTensors=*/ValueRange{},
-        SmallVector<AffineMap, 3>(args.size() + resultTypes.size(),
-                                  rewriter.getMultiDimIdentityMap(rank)),
-        SmallVector<StringRef, 3>(rank, getParallelIteratorTypeName()));
-    Region &region = genericOp.getRegion();
-    rewriter.cloneRegionBefore(producer.getOperation()->getRegion(0), region,
-                               region.begin());
-    return cast<LinalgOp>(genericOp.getOperation());
+           isFusableWithReshapeByDimExpansion(producer,
+                                              producer.getNumInputs());
   }
 
-  static LinalgOp fuse(LinalgOp producer, TensorReshapeOp consumer,
-                       unsigned consumerIdx, PatternRewriter &rewriter,
-                       OperationFolder *folder = nullptr) {
-    if (isCollapsingAndFusible(producer, consumer, consumerIdx))
-      return fuseCollapsingCase(producer, consumer, consumerIdx, rewriter);
-    if (isExpandingAndFusible(producer, consumer, consumerIdx))
-      return fuseExpandingCase(producer, consumer, consumerIdx, rewriter);
-    return nullptr;
+  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);
   }
 };
 
 /// Implementation of fusion on tensor ops when producer is a splat constant.
 struct FuseConstantOpAsProducer {
-  static bool isFusible(ConstantOp producer, LinalgOp consumer,
+  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();
   }
 
-  static LinalgOp fuse(ConstantOp producer, LinalgOp consumer,
-                       unsigned consumerIdx, PatternRewriter &rewriter,
-                       OperationFolder *folder = nullptr) {
-    if (!isFusible(producer, consumer, consumerIdx))
-      return nullptr;
+  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;
 
     // 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.
-    Operation *consumerOp = consumer.getOperation();
-    SmallVector<Value, 4> fusedOperands(consumerOp->getOperands());
+    SmallVector<Value, 4> fusedOperands(consumer.getInputs());
     fusedOperands.erase(std::next(fusedOperands.begin(), consumerIdx));
 
     // Create a constant scalar value from the splat constant.
     Value scalarConstant = rewriter.create<ConstantOp>(
         producer.getLoc(),
         producer.value().cast<DenseElementsAttr>().getSplatValue());
 
     LinalgOp fusedOp = createLinalgOpOfSameType(
         consumer, rewriter, rewriter.getUnknownLoc(),
-        consumerOp->getResultTypes(),
+        consumer.getOperation()->getResultTypes(),
         /*inputs=*/fusedOperands,
         /*outputBuffers=*/ValueRange{},
         /*initTensors=*/ValueRange{}, // no init tensors for now.
         rewriter.getAffineMapArrayAttr(fusedIndexMaps),
         consumer.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 = consumerOp->getRegion(0);
+    Region &consumerRegion = consumer.getOperation()->getRegion(0);
     Block &entryBlock = *consumerRegion.begin();
-    unsigned argIndex = entryBlock.getNumArguments() -
-                        consumerOp->getNumOperands() + consumerIdx;
+    unsigned argIndex =
+        entryBlock.getNumArguments() - consumer.getNumInputs() + consumerIdx;
     BlockAndValueMapping mapping;
     mapping.map(entryBlock.getArgument(argIndex), scalarConstant);
     Region &fusedRegion = fusedOp.getOperation()->getRegion(0);
     rewriter.cloneRegionBefore(consumerRegion, fusedRegion, fusedRegion.begin(),
                                mapping);
-    return fusedOp;
+    return SmallVector<Value, 1>(fusedOp.getOperation()->getResults());
   }
 };
 } // namespace
 
-Operation *mlir::linalg::fuseTensorOps(PatternRewriter &rewriter,
-                                       Operation *consumer,
-                                       unsigned consumerIdx,
-                                       OperationFolder *folder) {
+Optional<SmallVector<Value, 1>>
+mlir::linalg::fuseTensorOps(PatternRewriter &rewriter, Operation *consumer,
+                            unsigned consumerIdx, OperationFolder *folder,
+                            bool useReshapeFusionByExpansion) {
   if (consumerIdx >= consumer->getNumOperands())
-    return nullptr;
+    return llvm::None;
   Operation *producer = consumer->getOperand(consumerIdx).getDefiningOp();
   if (!producer || producer->getNumResults() != 1)
-    return nullptr;
+    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))
-      return FuseTensorReshapeOpAsProducer::fuse(reshapeOpProducer,
-                                                 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);
-    return nullptr;
+    return llvm::None;
   }
 
   if (isa<GenericOp, IndexedGenericOp>(producer)) {
     // Fuse when consumer is a TensorReshapeOp.
     if (TensorReshapeOp reshapeOp = dyn_cast<TensorReshapeOp>(consumer)) {
-      return FuseTensorReshapeOpAsConsumer::fuse(
+      if (useReshapeFusionByExpansion) {
+        return FuseTensorReshapeOpAsConsumerByExpansion::fuse(
+            cast<LinalgOp>(producer), reshapeOp, consumerIdx, rewriter, folder);
+      }
+      return FuseTensorReshapeOpAsConsumerByLinearization::fuse(
           cast<LinalgOp>(producer), reshapeOp, consumerIdx, rewriter, folder);
     }
   }
 
-  return nullptr;
+  return llvm::None;
 }
 
 namespace {
 /// Patterns to fuse a generic op, with the producer of its operands.
 template <typename LinalgOpTy>
 struct FuseTensorOps : public OpRewritePattern<LinalgOpTy> {
-  using OpRewritePattern<LinalgOpTy>::OpRewritePattern;
+  FuseTensorOps(MLIRContext *context, bool useReshapeFusionByExpansion = false)
+      : OpRewritePattern<LinalgOpTy>::OpRewritePattern(context),
+        useReshapeFusionByExpansion(useReshapeFusionByExpansion) {}
 
   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();
-      if (Operation *fusedOp = fuseTensorOps(rewriter, op, operandNum)) {
-        rewriter.replaceOp(op, fusedOp->getResults());
+      Optional<SmallVector<Value, 1>> fusedOpResults = fuseTensorOps(
+          rewriter, op, operandNum, nullptr, useReshapeFusionByExpansion);
+      if (fusedOpResults) {
+        rewriter.replaceOp(op, *fusedOpResults);
         if (producer && llvm::all_of(producer->getResults(),
                                      [](Value val) { return val.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);
+    populateLinalgTensorOpsFusionPatterns(op->getContext(), patterns,
+                                          useReshapeFusionByExpansion);
     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::populateLinalgTensorOpsFusionPatterns(
-    MLIRContext *context, OwningRewritePatternList &patterns) {
+    MLIRContext *context, OwningRewritePatternList &patterns,
+    bool useReshapeFusionByExpansion) {
   patterns.insert<FuseTensorOps<GenericOp>, FuseTensorOps<IndexedGenericOp>,
-                  FuseTensorOps<TensorReshapeOp>>(context);
+                  FuseTensorOps<TensorReshapeOp>>(context,
+                                                  useReshapeFusionByExpansion);
 }
 
-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);
 }

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

func @generic_op_reshape_consumer_nofusion(%arg0 : tensor<?x?x?x5xf32>,
   return %1 : tensor<?x?xf32>
 }
 
 // CHECK-LABEL: func @generic_op_reshape_consumer_nofusion
 //       CHECK: linalg.tensor_reshape
 
 // -----
 
-#map0 = affine_map<(d0, d1) -> (d0, d1)>
-#map1 = affine_map<(d0, d1, d2) -> (d0, d1)>
-#map2 = affine_map<(d0, d1, d2) -> (d2)>
-
-func @generic_op_reshape_consumer_expanding(%arg0: tensor<264x4xf32>)
-                                            -> tensor<8x33x4xf32> {
-  %cst = constant dense<2.000000e+00> : tensor<264x4xf32>
-  %0 = linalg.generic {
-     indexing_maps = [#map0, #map0, #map0],
-     iterator_types = ["parallel", "parallel"]}
-      ins(%arg0, %cst : tensor<264x4xf32>, tensor<264x4xf32>) {
-    ^bb0(%arg1: f32, %arg2: f32):  // no predecessors
-      %2 = mulf %arg1, %arg2 : f32
-      linalg.yield %2 : f32
-    } -> tensor<264x4xf32>
-  %1 = linalg.tensor_reshape %0 [#map1, #map2] :
-    tensor<264x4xf32> into tensor<8x33x4xf32>
-  return %1 : tensor<8x33x4xf32>
-}
-
-// The reshape op in `%arg0` is folded into the indexing map of generic op.
-//   CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0 * 33 + d1, d2)>
-//   CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
-//       CHECK: func @generic_op_reshape_consumer_expanding
-//   CHECK-NOT:   linalg.tensor_reshape
-//       CHECK:   %[[CST:.*]] = constant {{.*}} : f32
-//       CHECK:   linalg.generic
-//  CHECK-SAME:     indexing_maps = [#[[MAP0]], #[[MAP1]]]
-//  CHECK-SAME:   tensor<264x4xf32>
-//       CHECK:   -> tensor<8x33x4xf32>
-//   CHECK-NOT:   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"]}

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
+
+#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)>,
+                                    affine_map<(i, j, k, l) -> (j, k)>,
+                                    affine_map<(i, j, k, l) -> (l)>] :
+    tensor<?x?x?x?xf32> into tensor<?x?x?xf32>
+  %1 = linalg.generic {
+     indexing_maps = [#map0, #map1, #map1],
+     iterator_types = ["parallel", "parallel", "parallel"]}
+      ins(%0, %arg1 : tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
+    ^bb0(%arg3: f32, %arg4: f32):       // no predecessors
+      %1 = mulf %arg3, %arg4 : f32
+      linalg.yield %1 : f32
+  } -> tensor<?x?x?xf32>
+  return %1 : tensor<?x?x?xf32>
+}
+
+//  CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0)>
+//  CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d1)>
+//  CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d2, d3)>
+//  CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d0, d1, d2)>
+//  CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1, d2, d3) -> (d2, d3, d0, d1)>
+//      CHECK: func @generic_op_reshape_producer_fusion
+// CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?x?xf32>
+// CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
+//      CHECK:   %[[T0:.+]] = linalg.tensor_reshape %[[ARG1]]
+// CHECK-SAME:     [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
+// CHECK-SAME:     tensor<?x?x?xf32> into tensor<?x?x?x?xf32>
+//      CHECK:   %[[T1:.+]] = linalg.generic
+// CHECK-SAME:     indexing_maps = [#[[MAP3]], #[[MAP4]], #[[MAP4]]]
+// CHECK-SAME:     ["parallel", "parallel", "parallel", "parallel"]
+// CHECK-SAME:     ins(%[[ARG0]], %[[T0]] : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>)
+//      CHECK:   %[[T2:.+]] = linalg.tensor_reshape
+// CHECK-SAME:     [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
+// CHECK-SAME:     tensor<?x?x?x?xf32> into tensor<?x?x?xf32>
+//      CHECK:   return %[[T2]]
+
+// -----
+
+#map0 = affine_map<(d0, d1) -> (d0, d1)>
+func @generic_op_reshape_consumer_fusion(%arg0 : tensor<?x?xf32>,
+                                         %arg1 : tensor<?x?xf32>) ->
+                                         tensor<?x?x4x5xf32>
+{
+  %0 = linalg.generic {
+     indexing_maps = [#map0, #map0, #map0],
+     iterator_types = ["parallel", "parallel"]}
+      ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) {
+    ^bb0(%arg3: f32, %arg4: f32):       // no predecessors
+      %1 = mulf %arg3, %arg4 : f32
+      linalg.yield %1 : f32
+  } -> tensor<?x?xf32>
+  %1 = linalg.tensor_reshape %0 [affine_map<(i, j, k, l) -> (i)>,
+                                 affine_map<(i, j, k, l) -> (j, k, l)>] :
+    tensor<?x?xf32> into tensor<?x?x4x5xf32>
+  return %1 : tensor<?x?x4x5xf32>
+}
+
+//  CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0)>
+//  CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d2, d3)>
+//  CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+//      CHECK: func @generic_op_reshape_consumer_fusion
+// CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
+// CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
+//      CHECK:   %[[T0:.+]] = linalg.tensor_reshape %[[ARG0]]
+// CHECK-SAME:     [#[[MAP0]], #[[MAP1]]]
+// CHECK-SAME:     tensor<?x?xf32> into tensor<?x?x4x5xf32>
+//      CHECK:   %[[T1:.+]] = linalg.tensor_reshape %[[ARG1]]
+// CHECK-SAME:     [#[[MAP0]], #[[MAP1]]]
+// CHECK-SAME:     tensor<?x?xf32> into tensor<?x?x4x5xf32>
+//      CHECK:   %[[T2:.+]] = linalg.generic
+// CHECK-SAME:     indexing_maps = [#[[MAP2]], #[[MAP2]], #[[MAP2]]]
+// CHECK-SAME:     ["parallel", "parallel", "parallel", "parallel"]
+// CHECK-SAME:     ins(%[[T0]], %[[T1]] : tensor<?x?x4x5xf32>, tensor<?x?x4x5xf32>)
+//      CHECK:   return %[[T2]] : tensor<?x?x4x5xf32>
+
+
+// -----
+
+func @reshape_as_consumer_permutation
+  (%a : tensor<?x?x?xf32>, %b : tensor<?x?xf32>)
+    -> tensor<?x?x?x?x?x?xf32> {
+  %c = linalg.generic {
+         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)>],
+         iterator_types = ["parallel", "parallel", "parallel"]}
+         ins(%a, %b : tensor<?x?x?xf32>, tensor<?x?xf32>) {
+       ^bb0(%arg0 : f32, %arg1: f32):
+         %1 = addf %arg0, %arg1 : f32
+         linalg.yield %1 : f32
+       } -> tensor<?x?x?xf32>
+  %d = linalg.tensor_reshape %c
+         [affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1)>,
+          affine_map<(d0, d1, d2, d3, d4, d5) -> (d2)>,
+          affine_map<(d0, d1, d2, d3, d4, d5) -> (d3, d4, d5)>]
+       : tensor<?x?x?xf32> into tensor<?x?x?x?x?x?xf32>
+  return %d : tensor<?x?x?x?x?x?xf32>
+}
+//  CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2)>
+//  CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d3, d4)>
+//  CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d5)>
+//  CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+//  CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1, d2, d3) -> (d3)>
+//  CHECK-DAG: #[[MAP5:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d3, d4, d0, d1, d5)>
+//  CHECK-DAG: #[[MAP6:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d3, d4, d5)>
+//  CHECK-DAG: #[[MAP7:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d5, d2, d3, d4)>
+//      CHECK: func @reshape_as_consumer_permutation
+// CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
+// CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
+//      CHECK:   %[[T0:.+]] = linalg.tensor_reshape %[[ARG0]]
+// CHECK-SAME:     [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
+// CHECK-SAME:     tensor<?x?x?xf32> into tensor<?x?x?x?x?x?xf32>
+//      CHECK:   %[[T1:.+]] = linalg.tensor_reshape %[[ARG1]]
+// CHECK-SAME:     [#[[MAP3]], #[[MAP4]]]
+// CHECK-SAME:     tensor<?x?xf32> into tensor<?x?x?x?xf32>
+//      CHECK:   %[[T2:.+]] = linalg.generic
+// CHECK-SAME:     indexing_maps = [#[[MAP5]], #[[MAP6]], #[[MAP7]]]
+// CHECK-SAME:     ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]
+// CHECK-SAME:     ins(%[[T0]], %[[T1]] : tensor<?x?x?x?x?x?xf32>, tensor<?x?x?x?xf32>)
+//      CHECK:   return %[[T2]] : tensor<?x?x?x?x?x?xf32>
+
+// -----
+
+#map0 = affine_map<(d0, d1) -> (d0, d1)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d2)>
+
+func @generic_op_reshape_consumer_static(%arg0: tensor<264x4xf32>)
+                                            -> tensor<8x33x4xf32> {
+  %cst = constant dense<2.000000e+00> : tensor<264x4xf32>
+  %0 = linalg.generic {
+     indexing_maps = [#map0, #map0, #map0],
+     iterator_types = ["parallel", "parallel"]}
+      ins(%arg0, %cst : tensor<264x4xf32>, tensor<264x4xf32>) {
+    ^bb0(%arg1: f32, %arg2: f32):  // no predecessors
+      %2 = mulf %arg1, %arg2 : f32
+      linalg.yield %2 : f32
+    } -> tensor<264x4xf32>
+  %1 = linalg.tensor_reshape %0 [#map1, #map2] :
+    tensor<264x4xf32> into tensor<8x33x4xf32>
+  return %1 : tensor<8x33x4xf32>
+}
+
+//  CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+//  CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d2)>
+//  CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+//      CHECK: func @generic_op_reshape_consumer_static
+// CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<264x4xf32>
+//      CHECK:   %[[T0:.+]] = linalg.tensor_reshape %[[ARG0]]
+// CHECK-SAME:     [#[[MAP0]], #[[MAP1]]]
+// CHECK-SAME:     tensor<264x4xf32> into tensor<8x33x4xf32>
+//      CHECK:   %[[T1:.+]] = linalg.generic
+// CHECK-SAME:     indexing_maps = [#[[MAP2]], #[[MAP2]]]
+// CHECK-SAME:     ["parallel", "parallel", "parallel"]
+// CHECK-SAME:     ins(%[[T0]] : tensor<8x33x4xf32>)
+//      CHECK:   return %[[T1]] : tensor<8x33x4xf32>
+
+// -----
+
+func @scalar_reshape(%arg0 : tensor<1x10xf32>, %arg1 : tensor<1xf32>)
+                     -> tensor<1x10xf32> {
+  %0 = linalg.tensor_reshape %arg1 [] : tensor<1xf32> into tensor<f32>
+  %1 = linalg.generic
+    {indexing_maps = [affine_map<(d0) -> ()>, affine_map<(d0) -> (d0)>],
+     iterator_types = ["parallel"]} ins(%0 : tensor<f32>) {
+  ^bb0(%arg2: f32):  // no predecessors
+    linalg.yield %arg2 : f32
+  } -> tensor<10xf32>
+  %2 = linalg.tensor_reshape %1 [affine_map<(d0, d1) -> (d0, d1)>]
+    : tensor<10xf32> into tensor<1x10xf32>
+  return %2 : tensor<1x10xf32>
+}
+
+//  CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> ()>
+//  CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1) -> (d0, d1)>
+//      CHECK: func @scalar_reshape
+// CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<1x10xf32>
+// CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<1xf32>
+//      CHECK:   %[[T0:.+]] = linalg.tensor_reshape %[[ARG1]] []
+// CHECK-SAME:     tensor<1xf32> into tensor<f32>
+//      CHECK:   %[[T1:.+]] = linalg.generic
+// CHECK-SAME:     indexing_maps = [#[[MAP0]], #[[MAP1]]]
+// CHECK-SAME:     iterator_types = ["parallel", "parallel"]
+// CHECK-SAME:     ins(%[[T0]] : tensor<f32>)
+//      CHECK:   return %[[T1]] : tensor<1x10xf32>