[mlir][Linalg] Add support for fusing linalg.tensor_reshape with linalg.generic operations.

[mravishankar]

Summary

Depends On D78463

[mravishankar]

Updating the logic to handle cases where some dimensions are dynamic.

[nicolasvasilache]

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

583

Adding the example from your test would be informative:

E.g.
%0 = op ... : tensor<?x?x4x5xf32>
with output index_map `affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>`

and reshape:
%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>

would be rewritten into:
%0 = op ... : tensor<?x?x4x5xf32>
with output index_map `affine_map<(d0, d1, d2, d3) -> (d0, d1 * 20 + d2 * 5 + d3)>`

644

Could you please add a:

// TODO: In the future this restriction may justify extending the linalg.generic to semi-affine maps.
// TODO: Alternatively, fusing across a reshape and pushing the reshape towards the boundary of the function could help too.

648

I wonder if we could slightly refactor the function exposed here https://reviews.llvm.org/D75575
and just use it instead of rewriting a slightly different version.

Basically the existing function would just need a starting dimPos = 0 by default.

You could then just use that to get:

for ... {
  AffineExpr stridedExpression = makeCanonicalStridedLayoutExpr(srcShape.slice(), context);
  if (!isPureAffine(stridedExpression))
    return AffineMap();
  results.push_back(stridedExpression);
}
return AffineMap.get(...);

And check whether AffineMap is empty to know whether it is fusible.

The downside is that in your current code structure you could be recomputing it twice but IMO this would largely be beneficial to save a few dozen lines of code that is similar to something we already have in the codebase.

If the recomputing bothers you, you could refactor a bit but I wouldn't be worried at this time, these are really trivial computations.

681

Can we simplify l680 - 695 with something along the lines of:

SmallVector<Attribute, 4> fusedIndexMaps(consumerIndexMaps.begin(), consumerIndexMaps.end());
fusedIndexMaps[consumerIdx] = linearize();

726

typo thee

729

Same here, create then update should be simpler.

783

If you adopt the style in the previous PR, this conditional would just fold into the Operation* based form above.

[mravishankar]

Addressing comments

[nicolasvasilache]

Thanks for the cleanup!

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

650

Could this whole loop be replaced by something like:

if (!useIndexMap.isIdentity()) return false;
for (auto linearizedExpr : linearizeCollapsedDims(useIndexMap, srcShape, reassociationMaps) )
  if (!isPureAffine(linearizedExpr)) return false;

[mravishankar]

Submitting some responses that I had forgetten to submit

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

583

THanks for the detailed comment.

648

It ended up being quite a significant change to makeCanonicalStridedLayoutExpr, but it is more general now, and it works well with the current use case as well. It does recompute twice, but thats OK by me. PTAL

681

Nice! Done

[antiagainst]

Cool stuff!

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

580

Linearizes

592

Neat: alignment :)

661

Can return WalkResult::interrupt() after this to avoid further walking.

[mravishankar]

Rebase

[mravishankar]

Addressing comments

[mravishankar]

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

650

I removed this in the checker and just added a check in the fuse operations to check that the resulting map is pure affine.

661

I dont think the walk on AffineExpr allows that. In any case not needed anymore

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

def Linalg_TensorReshapeOp : Linalg_ReshapeLikeOp<"tensor_reshape">,
   }];
   let extraClassDeclaration = commonExtraClassDeclaration # [{
     RankedTensorType getSrcType() {
       return src().getType().cast<RankedTensorType>();
     }
     RankedTensorType getResultType() {
       return result().getType().cast<RankedTensorType>();
     }
+    SmallVector<AffineMap, 4> getReassociationMaps() {
+      return llvm::to_vector<4>(llvm::map_range(reassociation(),
+        [](Attribute a) { return a.cast<AffineMapAttr>().getValue(); }));
+    }
   }];
 }
 
 def Linalg_SliceOp : Linalg_Op<"slice", [NoSideEffect]>,
     Arguments<(ins AnyStridedMemRef:$view,
                    Variadic<AnyTypeOf<[Range, Index]>>:$indexings)>,
     Results<(outs AnyStridedMemRef)> {
   let summary = "Produce a rank-reduced `subview` of a base `view`.";

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

 #ifndef MLIR_DIALECT_LINALG_PASSES_H_
 #define MLIR_DIALECT_LINALG_PASSES_H_
 
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/ArrayRef.h"
 
 namespace mlir {
 class FuncOp;
+class MLIRContext;
 class ModuleOp;
 template <typename T> class OperationPass;
+class OwningRewritePatternList;
 class Pass;
 
 std::unique_ptr<OperationPass<FuncOp>> createLinalgFusionPass();
 std::unique_ptr<Pass> createLinalgFusionOfTensorOpsPass();
 
 std::unique_ptr<OperationPass<FuncOp>>
 createLinalgTilingPass(ArrayRef<int64_t> tileSizes = {});
 

 /// std.load/std.store accesses.
 std::unique_ptr<OperationPass<FuncOp>> createConvertLinalgToParallelLoopsPass();
 
 /// Create a pass to convert Linalg operations to affine.for loops and
 /// affine_load/affine_store accesses.
 /// Placeholder for now, this is NYI.
 std::unique_ptr<OperationPass<FuncOp>> createConvertLinalgToAffineLoopsPass();
 
+/// Patterns for fusing linalg operation on tensors.
+void populateLinalgTensorOpsFusionPatterns(MLIRContext *context,
+                                           OwningRewritePatternList &patterns);
+
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_PASSES_H_

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

computeTensorReshapeCollapsedType(RankedTensorType type,
   newShape.reserve(reassociation.size());
 
   // Use the fact that reassociation is valid to simplify the logic: only use
   // each map's rank.
   assert(isReassociationValid(reassociation) && "invalid reassociation");
   unsigned currentDim = 0;
   for (AffineMap m : reassociation) {
     unsigned dim = m.getNumResults();
-    auto band = shape.drop_front(currentDim).take_front(dim);
+    auto band = shape.slice(currentDim, dim);
     int64_t size = 1;
     if (llvm::is_contained(band, ShapedType::kDynamicSize))
       size = ShapedType::kDynamicSize;
     else
       for (unsigned d = 0; d < dim; ++d)
         size *= shape[currentDim + d];
     newShape.push_back(size);
     currentDim += dim;

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

for (auto &op : producerBlock.getOperations()) {
       rewriter.clone(op, mapper);
     }
     for (auto &op : consumerBlock.getOperations())
       rewriter.clone(op, mapper);
   }
 };
 } // namespace
 
+/// Linearize the expressions in `sourceMap` based on the `reassociationMaps`
+/// provided, given the shape of the source tensor that corresponds to the
+/// `sourceMap`. Note that this implicitly assumes that the tensors dimensions
+/// are "row-major" ordered logically.
+static AffineMap linearizeCollapsedDims(AffineMap sourceMap,
+                                        ArrayRef<int64_t> sourceShape,
+                                        ArrayRef<AffineMap> reassociationMaps) {
+  SmallVector<AffineExpr, 4> resultExprs;
+  resultExprs.reserve(reassociationMaps.size());
+  ArrayRef<AffineExpr> sourceExprs = sourceMap.getResults();
+  MLIRContext *context = sourceMap.getContext();
+
+  // Generic function to linearize the collapsed dimensions of the source map.
+  auto linearizeFn = [&context](ArrayRef<AffineExpr> exprs,
+                                ArrayRef<int64_t> shape) -> AffineExpr {
+    assert(!exprs.empty() && exprs.size() == shape.size());
+    AffineExpr resultExpr = nullptr;
+    int64_t stride = 1;
+    for (auto it : llvm::zip(reverse(exprs), reverse(shape))) {
+      resultExpr =
+          (resultExpr
+               ? std::get<0>(it) * getAffineConstantExpr(stride, context) +
+                     resultExpr
+               : std::get<0>(it));
+      stride *= std::get<1>(it);
+    }
+    return resultExpr;
+  };
+
+  // Compute the result exprs based on the reassociation maps.
+  for (AffineMap map : reassociationMaps) {
+    ArrayRef<AffineExpr> collapsedDims = map.getResults();
+    // Assume that they are in-order and contiguous (already checked in
+    // verifier).
+    assert(!collapsedDims.empty());
+    unsigned startDim =
+        collapsedDims.front().cast<AffineDimExpr>().getPosition();
+    unsigned endDim = collapsedDims.back().cast<AffineDimExpr>().getPosition();
+    AffineExpr linearizedExpr =
+        linearizeFn(sourceExprs.slice(startDim, endDim - startDim + 1),
+                    sourceShape.slice(startDim, endDim - startDim + 1));
+    resultExprs.push_back(linearizedExpr);
+  }
+  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) {
+  RankedTensorType returnType = reshapeOp.getResultType();
+  RankedTensorType operandType = reshapeOp.getSrcType();
+  // Reshape is fusible 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()))
+    return false;
+  // For reshape to be fused, the collapsed (expanded) dimensions of the operand
+  // (result) must be statically shaped.
+  ArrayRef<int64_t> srcShape =
+      (asProducer ? returnType.getShape() : operandType.getShape());
+  auto reassociationMaps = reshapeOp.getReassociationMaps();
+  unsigned dim = 0;
+  for (auto collapsedDims : reassociationMaps) {
+    if (dim < srcShape.size() - 1 &&
+        llvm::is_contained(
+            srcShape.slice(dim + 1, collapsedDims.getNumResults() - 1),
+            ShapedType::kDynamicSize))
+      return false;
+    dim += collapsedDims.getNumResults();
+  }
+  // Currently only implement the case where the consumer is using an identity
+  // map for the indexing_map
+  return useIndexMap.isIdentity();
+}
+
+namespace {
+template <typename LinalgOpTy>
+struct FuseTensorReshapeOpAsProducer
+    : public FuseTensorOpsImpl<FuseTensorReshapeOpAsProducer<LinalgOpTy>,
+                               TensorReshapeOp, LinalgOpTy> {
+  static bool isFusible(TensorReshapeOp producer, LinalgOpTy consumer,
+                        unsigned consumerIdx) {
+    return isTensorReshapeOpFusible(
+        producer, consumer.getInputIndexingMap(consumerIdx), true);
+  }
+
+  static Operation *createFusedOp(PatternRewriter &rewriter,
+                                  ArrayRef<Value> fusedOperands,
+                                  TensorReshapeOp producer, LinalgOpTy consumer,
+                                  unsigned consumerIdx,
+                                  OperationFolder *folder = nullptr) {
+    // Compute indexing_maps for the fused operation. The indexing_maps for the
+    // operands of the consumers that arent fused are the same.
+    auto consumerIndexMaps = consumer.indexing_maps();
+    SmallVector<Attribute, 4> fusedIndexMaps;
+    fusedIndexMaps.reserve(consumerIndexMaps.size());
+    fusedIndexMaps.append(consumerIndexMaps.begin(),
+                          std::next(consumerIndexMaps.begin(), consumerIdx));
+
+    // Compute the indexing map to use for the operand of the producer.
+    AffineMap producerOperandIndexingMap = linearizeCollapsedDims(
+        consumer.getInputIndexingMap(consumerIdx),
+        producer.getResultType().getShape(), producer.getReassociationMaps());
+    fusedIndexMaps.push_back(AffineMapAttr::get(producerOperandIndexingMap));
+
+    // The rest of the indexing_maps need to be obtained from the remaining
+    // unfused consumer operands.
+    fusedIndexMaps.append(std::next(consumerIndexMaps.begin(), consumerIdx + 1),
+                          consumerIndexMaps.end());
+
+    auto fusedOp = rewriter.create<LinalgOpTy>(
+        rewriter.getUnknownLoc(), consumer.getResultTypes(), fusedOperands,
+        rewriter.getI64IntegerAttr(fusedOperands.size()),
+        rewriter.getI64IntegerAttr(consumer.getNumResults()),
+        rewriter.getArrayAttr(fusedIndexMaps), consumer.iterator_types(),
+        /*doc=*/nullptr,
+        /*library_call=*/nullptr);
+    auto &fusedRegion = fusedOp.region();
+    rewriter.cloneRegionBefore(consumer.region(), fusedRegion,
+                               fusedRegion.begin());
+    return fusedOp;
+  }
+};
+
+template <typename LinalgOpTy>
+struct FuseTensorReshapeOpAsConsumer
+    : public FuseTensorOpsImpl<FuseTensorReshapeOpAsConsumer<LinalgOpTy>,
+                               LinalgOpTy, TensorReshapeOp> {
+  static bool isFusible(LinalgOpTy producer, TensorReshapeOp consumer,
+                        unsigned consumerIdx) {
+    return isTensorReshapeOpFusible(consumer, producer.getOutputIndexingMap(0),
+                                    false);
+  }
+
+  static Operation *createFusedOp(PatternRewriter &rewriter,
+                                  ArrayRef<Value> fusedOperands,
+                                  LinalgOpTy producer, TensorReshapeOp consumer,
+                                  unsigned consumerIdx,
+                                  OperationFolder *folder = nullptr) {
+    // Thee indexing_maps for the operands that were originally operands from
+    // the producers are the same as before.
+    auto producerIndexMaps = producer.indexing_maps();
+    SmallVector<Attribute, 4> fusedIndexMaps;
+    fusedIndexMaps.reserve(producerIndexMaps.size());
+    fusedIndexMaps.append(producerIndexMaps.begin(),
+                          std::prev(producerIndexMaps.end()));
+
+    // Compute the indexing map to use for the operand of the producer.
+    AffineMap producerResultIndexingMap = linearizeCollapsedDims(
+        producer.getOutputIndexingMap(0), consumer.getSrcType().getShape(),
+        consumer.getReassociationMaps());
+    fusedIndexMaps.push_back(AffineMapAttr::get(producerResultIndexingMap));
+
+    auto fusedOp = rewriter.create<LinalgOpTy>(
+        rewriter.getUnknownLoc(), consumer.getResultType(), fusedOperands,
+        rewriter.getI64IntegerAttr(fusedOperands.size()),
+        rewriter.getI64IntegerAttr(1), rewriter.getArrayAttr(fusedIndexMaps),
+        producer.iterator_types(),
+        /*doc=*/nullptr,
+        /*library_call=*/nullptr);
+    auto &fusedRegion = fusedOp.region();
+    rewriter.cloneRegionBefore(producer.region(), fusedRegion,
+                               fusedRegion.begin());
+    return fusedOp;
+  }
+};
+} // namespace
+
 template <>
 Operation *mlir::linalg::fuseTensorOps<GenericOp>(PatternRewriter &rewriter,
                                                   GenericOp consumer,
                                                   unsigned consumerIdx,
                                                   OperationFolder *folder) {
   if (consumerIdx >= consumer.getNumInputs() || !consumer.hasTensorSemantics())
     return nullptr;
   Operation *producer = consumer.getInput(consumerIdx).getDefiningOp();
   if (!producer || producer->getNumResults() != 1)
     return nullptr;
   if (auto genericOp = dyn_cast<GenericOp>(producer)) {
     if (genericOp.hasTensorSemantics())
       return FuseGenericOpsOnTensors::fuse(rewriter, genericOp, consumer,
                                            consumerIdx);
+  } else if (auto reshapeOp = dyn_cast<TensorReshapeOp>(producer)) {
+    return FuseTensorReshapeOpAsProducer<GenericOp>::fuse(
+        rewriter, reshapeOp, consumer, consumerIdx);
+  }
+  return nullptr;
+}
+
+template <>
+Operation *mlir::linalg::fuseTensorOps<TensorReshapeOp>(
+    PatternRewriter &rewriter, TensorReshapeOp consumer, unsigned consumerIdx,
+    OperationFolder *folder) {
+  Operation *producer = consumer.getOperand().getDefiningOp();
+  if (!producer || producer->getNumResults() != 1)
+    return nullptr;
+  if (auto genericOp = dyn_cast<GenericOp>(producer)) {
+    if (genericOp.hasTensorSemantics())
+      return FuseTensorReshapeOpAsConsumer<GenericOp>::fuse(
+          rewriter, genericOp, consumer, consumerIdx);
   }
   return nullptr;
 }
 
 namespace {
 /// Patterns to fuse a generic op, with the producer of its operands.
 template <typename LinalgOpTy>
 struct FuseTensorOps : public OpRewritePattern<LinalgOpTy> {

 };
 
 /// Pass that fuses generic ops on tensors. Used only for testing.
 struct FusionOfTensorOpsPass
     : public LinalgFusionOfTensorOpsBase<FusionOfTensorOpsPass> {
   void runOnOperation() override {
     OwningRewritePatternList patterns;
     Operation *op = getOperation();
-    patterns.insert<FuseTensorOps<GenericOp>>(op->getContext());
+    populateLinalgTensorOpsFusionPatterns(op->getContext(), patterns);
     applyPatternsAndFoldGreedily(op->getRegions(), patterns);
   };
 };
 
 struct LinalgFusionPass : public LinalgFusionBase<LinalgFusionPass> {
   void runOnFunction() override { fuseLinalgOpsGreedily(getFunction()); }
 };
 } // namespace
 
+void mlir::populateLinalgTensorOpsFusionPatterns(
+    MLIRContext *context, OwningRewritePatternList &patterns) {
+  patterns.insert<FuseTensorOps<GenericOp>, FuseTensorOps<TensorReshapeOp>>(
+      context);
+}
+
 std::unique_ptr<OperationPass<FuncOp>> mlir::createLinalgFusionPass() {
   return std::make_unique<LinalgFusionPass>();
 }
 
 std::unique_ptr<Pass> mlir::createLinalgFusionOfTensorOpsPass() {
   return std::make_unique<FusionOfTensorOpsPass>();
 }

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

func @add_mul_scalar_fusion(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32>
   %1 = linalg.generic {args_in = 2 : i64, args_out = 1 : i64, indexing_maps = [#map0, #map0, #map0], iterator_types = []} %0, %arg2 {
     ^bb0(%arg3: f32, %arg4: f32):       // no predecessors
       %1 = mulf %arg3, %arg4 : f32
       linalg.yield %1 : f32
   }: tensor<f32>, tensor<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
+    {args_in = 2 : i64, args_out = 1 : i64,
+     indexing_maps = [#map0, #map0, #map0],
+     iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
+    %0, %arg1 {
+    ^bb0(%arg3: f32, %arg4: f32):       // no predecessors
+      %1 = mulf %arg3, %arg4 : f32
+      linalg.yield %1 : f32
+  }: tensor<?x?x4x?xf32>, tensor<?x?x4x?xf32> -> tensor<?x?x4x?xf32>
+  return %1 : tensor<?x?x4x?xf32>
+}
+
+// CHECK-LABEL: func @generic_op_reshape_producer_fusion
+//       CHECK: linalg.generic
+//  CHECK-SAME:   args_in = 2
+//  CHECK-SAME:   args_out = 1
+//  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
+    {args_in = 2 : i64, args_out = 1 : i64,
+     indexing_maps = [#map0, #map0, #map0],
+     iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
+    %arg0, %arg1 {
+    ^bb0(%arg3: f32, %arg4: f32):       // no predecessors
+      %1 = mulf %arg3, %arg4 : f32
+      linalg.yield %1 : f32
+  }: tensor<?x?x4x5xf32>, tensor<?x?x4x5xf32> -> 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:   args_in = 2
+//  CHECK-SAME:   args_out = 1
+//  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
+    {args_in = 2 : i64, args_out = 1 : i64,
+     indexing_maps = [#map0, #map0, #map0],
+     iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
+    %arg0, %arg1 {
+    ^bb0(%arg3: f32, %arg4: f32):       // no predecessors
+      %1 = mulf %arg3, %arg4 : f32
+      linalg.yield %1 : f32
+  }: tensor<?x?x?x5xf32>, tensor<?x?x?x5xf32> -> 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