[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.

799

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/include/mlir/IR/StandardTypes.h

 /// Given MemRef `sizes` that are either static or dynamic, returns the
 /// canonical "contiguous" strides AffineExpr. Strides are multiplicative and
 /// once a dynamic dimension is encountered, all canonical strides become
 /// dynamic and need to be encoded with a different symbol.
 /// For canonical strides expressions, the offset is always 0 and and fastest
 /// varying stride is always `1`.
 ///
 /// Examples:
-///   - memref<3x4x5xf32> has canonical stride expression `20*d0 + 5*d1 + d2`.
-///   - memref<3x?x5xf32> has canonical stride expression `s0*d0 + 5*d1 + d2`.
-///   - memref<3x4x?xf32> has canonical stride expression `s1*d0 + s0*d1 + d2`.
+///   - memref<3x4x5xf32> has canonical stride expression
+///         `20*exprs[0] + 5*exprs[1] + exprs[2]`.
+///   - memref<3x?x5xf32> has canonical stride expression
+///         `s0*exprs[0] + 5*exprs[1] + exprs[2]`.
+///   - memref<3x4x?xf32> has canonical stride expression
+///         `s1*exprs[0] + s0*exprs[1] + exprs[2]`.
+AffineExpr makeCanonicalStridedLayoutExpr(ArrayRef<int64_t> sizes,
+                                          ArrayRef<AffineExpr> exprs,
+                                          MLIRContext *context);
+
+/// Return the result of makeCanonicalStrudedLayoutExpr for the common case
+/// where `exprs` is {d0, d1, .., d_(sizes.size()-1)}
 AffineExpr makeCanonicalStridedLayoutExpr(ArrayRef<int64_t> sizes,
                                           MLIRContext *context);
 
 /// Return true if the layout for `t` is compatible with strided semantics.
 bool isStrided(MemRefType t);
 
 } // end namespace mlir
 
 #endif // MLIR_IR_STANDARDTYPES_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.
+///
+/// For example:
+///
+/// %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)>`
+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();
+
+  // 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();
+    AffineExpr linearizedExpr = makeCanonicalStridedLayoutExpr(
+        sourceShape.slice(startDim, collapsedDims.size()),
+        sourceExprs.slice(startDim, collapsedDims.size()), context);
+    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.
+  // 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.
+  ArrayRef<int64_t> srcShape =
+      (asProducer ? returnType.getShape() : operandType.getShape());
+  ArrayRef<AffineMap> reassociationMaps = reshapeOp.getReassociationMaps();
+  unsigned dim = 0;
+  for (AffineMap map : reassociationMaps) {
+    ArrayRef<AffineExpr> collapsedDims = map.getResults();
+    // If the linearized expr has symbols, then disable fusion since indexing
+    // maps cannot have symbols right now.
+    AffineExpr linearizedExpr = makeCanonicalStridedLayoutExpr(
+        srcShape.slice(dim, collapsedDims.size()),
+        useIndexMap.getResults().slice(dim, collapsedDims.size()),
+        reshapeOp.getContext());
+    bool hasSymbol = false;
+    linearizedExpr.walk([&hasSymbol](AffineExpr d) {
+      if (d.isa<AffineSymbolExpr>())
+        hasSymbol = true;
+    });
+    if (hasSymbol)
+      return false;
+    dim += collapsedDims.size();
+  }
+  // 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.
+    SmallVector<AffineMap, 4> fusedIndexMaps =
+        llvm::to_vector<4>(llvm::map_range(
+            consumer.indexing_maps(), [](Attribute attr) -> AffineMap {
+              return attr.cast<AffineMapAttr>().getValue();
+            }));
+    // Compute the indexing map to use for the operand of the producer.
+    fusedIndexMaps[consumerIdx] = linearizeCollapsedDims(
+        fusedIndexMaps[consumerIdx], producer.getResultType().getShape(),
+        producer.getReassociationMaps());
+
+    // 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;
+
+    SmallVector<Attribute, 4> indexMapAttrs = llvm::to_vector<4>(
+        llvm::map_range(fusedIndexMaps, [](AffineMap map) -> Attribute {
+          return AffineMapAttr::get(map);
+        }));
+    auto fusedOp = rewriter.create<LinalgOpTy>(
+        rewriter.getUnknownLoc(), consumer.getResultTypes(), fusedOperands,
+        rewriter.getI64IntegerAttr(fusedOperands.size()),
+        rewriter.getI64IntegerAttr(consumer.getNumResults()),
+        rewriter.getArrayAttr(indexMapAttrs), 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) {
+    // 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();
+            }));
+    // Compute the indexing map to use for the operand of the producer.
+    fusedIndexMaps.back() = linearizeCollapsedDims(
+        producer.getOutputIndexingMap(0), consumer.getSrcType().getShape(),
+        consumer.getReassociationMaps());
+
+    // 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;
+
+    SmallVector<Attribute, 4> indexMapAttrs = llvm::to_vector<4>(
+        llvm::map_range(fusedIndexMaps, [](AffineMap map) -> Attribute {
+          return AffineMapAttr::get(map);
+        }));
+
+    auto fusedOp = rewriter.create<LinalgOpTy>(
+        rewriter.getUnknownLoc(), consumer.getResultType(), fusedOperands,
+        rewriter.getI64IntegerAttr(fusedOperands.size()),
+        rewriter.getI64IntegerAttr(1), rewriter.getArrayAttr(indexMapAttrs),
+        producer.iterator_types(),
+        /*doc=*/nullptr,
+        /*library_call=*/nullptr);
+    auto &fusedRegion = fusedOp.region();
+    rewriter.cloneRegionBefore(producer.region(), fusedRegion,
+                               fusedRegion.begin());
+    return fusedOp;
+  }
+};
+} // namespace
+
 Operation *mlir::linalg::fuseTensorOps(PatternRewriter &rewriter,
                                        Operation *consumer,
                                        unsigned consumerIdx,
                                        OperationFolder *folder) {
   if (consumerIdx >= consumer->getNumOperands())
     return nullptr;
   Operation *producer = consumer->getOperand(consumerIdx).getDefiningOp();
   if (!producer || producer->getNumResults() != 1)
     return nullptr;
 
+  // Fuse when consumer is GenericOp.
   if (GenericOp genericOp = dyn_cast<GenericOp>(consumer)) {
     if (!genericOp.hasTensorSemantics())
       return nullptr;
     if (auto genericOpProducer = dyn_cast<GenericOp>(producer)) {
       if (genericOpProducer.hasTensorSemantics())
         return FuseGenericOpsOnTensors::fuse(rewriter, genericOpProducer,
                                              genericOp, consumerIdx);
-    }
+    } else if (auto reshapeOpProducer = dyn_cast<TensorReshapeOp>(producer)) {
+      return FuseTensorReshapeOpAsProducer<GenericOp>::fuse(
+          rewriter, reshapeOpProducer, genericOp, consumerIdx);
+    }
+    return nullptr;
+
+  // Fuse when consumer is a TensorReshapeOp.
+  if (TensorReshapeOp reshapeOp = dyn_cast<TensorReshapeOp>(consumer)) {
+    if (auto genericOpProducer = dyn_cast<GenericOp>(producer)) {
+      if (genericOpProducer.hasTensorSemantics())
+        return FuseTensorReshapeOpAsConsumer<GenericOp>::fuse(
+            rewriter, genericOpProducer, reshapeOp, consumerIdx);
+    }
+    return nullptr;
   }
   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/lib/IR/StandardTypes.cpp

auto simplifiedLayoutExpr =
       simplifyAffineExpr(m.getResult(0), m.getNumDims(), m.getNumSymbols());
   if (expr != simplifiedLayoutExpr)
     return MemRefType::Builder(t).setAffineMaps({AffineMap::get(
         m.getNumDims(), m.getNumSymbols(), simplifiedLayoutExpr)});
   return MemRefType::Builder(t).setAffineMaps({});
 }
 
 AffineExpr mlir::makeCanonicalStridedLayoutExpr(ArrayRef<int64_t> sizes,
+                                                ArrayRef<AffineExpr> exprs,
                                                 MLIRContext *context) {
   AffineExpr expr;
   bool dynamicPoisonBit = false;
+  unsigned numDims = 0;
   unsigned nSymbols = 0;
+  // Compute the number of symbols and dimensions of the passed exprs.
+  for (AffineExpr expr : exprs) {
+    expr.walk([&numDims, &nSymbols](AffineExpr d) {
+      if (AffineDimExpr dim = d.dyn_cast<AffineDimExpr>())
+        numDims = std::max(numDims, dim.getPosition() + 1);
+      else if (AffineSymbolExpr symbol = d.dyn_cast<AffineSymbolExpr>())
+        nSymbols = std::max(nSymbols, symbol.getPosition() + 1);
+    });
+  }
   int64_t runningSize = 1;
-  unsigned rank = sizes.size();
-  for (auto en : llvm::enumerate(llvm::reverse(sizes))) {
-    auto size = en.value();
-    auto position = rank - 1 - en.index();
+  for (auto en : llvm::zip(llvm::reverse(exprs), llvm::reverse(sizes))) {
+    int64_t size = std::get<1>(en);
     // Degenerate case, no size =-> no stride
     if (size == 0)
       continue;
-    auto d = getAffineDimExpr(position, context);
-    // Static case: stride = runningSize and runningSize *= size.
-    if (!dynamicPoisonBit) {
-      auto cst = getAffineConstantExpr(runningSize, context);
-      expr = expr ? expr + cst * d : cst * d;
-      if (size > 0)
-        runningSize *= size;
-      else
-        // From now on bail into dynamic mode.
-        dynamicPoisonBit = true;
-      continue;
-    }
-    // Dynamic case, new symbol for each new stride.
-    auto sym = getAffineSymbolExpr(nSymbols++, context);
-    expr = expr ? expr + d * sym : d * sym;
-  }
-  return simplifyAffineExpr(expr, rank, nSymbols);
+    AffineExpr dimExpr = std::get<0>(en);
+    AffineExpr stride = dynamicPoisonBit
+                            ? getAffineSymbolExpr(nSymbols++, context)
+                            : getAffineConstantExpr(runningSize, context);
+    expr = expr ? expr + dimExpr * stride : dimExpr * stride;
+  return simplifyAffineExpr(expr, numDims, nSymbols);
+
+AffineExpr mlir::makeCanonicalStridedLayoutExpr(ArrayRef<int64_t> sizes,
+                                                MLIRContext *context) {
+  SmallVector<AffineExpr, 4> exprs;
+  exprs.reserve(sizes.size());
+  for (auto dim : llvm::seq<unsigned>(0, sizes.size()))
+    exprs.push_back(getAffineDimExpr(dim, context));
+  return makeCanonicalStridedLayoutExpr(sizes, exprs, context);
 }
 
 /// Return true if the layout for `t` is compatible with strided semantics.
 bool mlir::isStrided(MemRefType t) {
   int64_t offset;
   SmallVector<int64_t, 4> stridesAndOffset;
   auto res = getStridesAndOffset(t, stridesAndOffset, offset);
   return succeeded(res);
 }

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