[mlir][Linalg] Enable TileAndFusePattern to work with tensors.

[mravishankar]

Summary

Depends On D93086

[hanchung]

mlir/lib/Dialect/Linalg/Analysis/DependenceAnalysis.cpp

164 ↗

(On Diff #316154)

auto

165 ↗

(On Diff #316154)

I think src won't be nullptr, so if (definingOp == src) is enough?

176–177 ↗

(On Diff #316154)

s/mixe/mixed

btw, is the assertion correct? After reading the message, I would expect either src.hasTensorSemantics() && dst.hasBufferSemantics() or a vice versa check.

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

97

Do you intend to test matmul + add + matmul? I only see matmul + add in the function.

[mravishankar]

mlir/lib/Dialect/Linalg/Analysis/DependenceAnalysis.cpp

165 ↗

(On Diff #316154)

I think so, but I prefer to leave it this way.

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

97

Both operands of add are the result of the matmul

[hanchung]

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

97

Oh I see, thanks!

[mravishankar]

Rebase on top of enhancements to LinalgDependenceGraph allowing
dependence tracking through results.

[mravishankar]

Removing errant function.

[mravishankar]

Rebase

[mravishankar]

Rebase and fix test

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

for (auto dependence : fusableDependences.lookup(op)) {
         return {};
       std::swap(candidates, fusableLoops);
     }
   }
 
   return fusableLoops;
 }
 
-// /// For `consumer` with tensor semantics, find the Linalg operation on
-// tensors
-// /// producer the operand at position `consumerIdx`. This is a simple use-def
-// /// chain using the SSA value, but returned as an element of the
-// /// `LinalgDependenceGraphElem` to use the same analysis for both tensors and
-// /// buffers.
-// static Optional<LinalgDependenceGraph::LinalgDependenceGraphElem>
-// findFusableProducerForTensorOp(OpOperand &consumerOpOperand) {
-//   // For now only looking for cases where the operand is produced by another
-//   // Linalg structured operation.
-//   LinalgOp consumer = cast<LinalgOp>(consumerOpOperand.getOwner());
-//   if (!consumer || !consumer.hasTensorSemantics())
-//     return llvm::None;
-//   unsigned consumerIdx = consumerOpOperand.getOperandNumber();
-//   Value value = consumerOpOperand.get();
-//   if (auto linalgOp = value.getDefiningOp<LinalgOp>()) {
-//     return LinalgDependenceGraph::LinalgDependenceGraphElem{
-//         &(linalgOp
-//               .getOutputOpOperands()[value.cast<OpResult>().getResultNumber()]),
-//         &(consumer.getInputOpOperands()[consumerIdx]),
-//         LinalgDependenceGraph::DependenceType::RAW};
-//   }
-//   return llvm::None;
-// }
-
-// static Optional<LinalgDependenceGraph::LinalgDependenceGraphElem>
-// findFusableProducer(OpOperand &consumerOpOperand,
-//                     const LinalgDependenceGraph &dependenceGraph) {
-//   LinalgOp consumer = cast<LinalgOp>(consumerOpOperand.getOwner());
-//   if (!consumer)
-//     return llvm::None;
-//   if (consumer.hasBufferSemantics())
-//     return findFusableProducerForBufferOp(consumerOpOperand,
-//     dependenceGraph);
-//   if (consumer.hasTensorSemantics())
-//     return findFusableProducerForTensorOp(consumerOpOperand);
-//   return llvm::None;
-// }
-
 /// Find all dependences that are fusable.
 FusableOpDependencesTy mlir::linalg::findAllFusableDependences(
     ArrayRef<LinalgOp> ops, const LinalgDependenceGraph &dependenceGraph) {
   FusableOpDependencesTy fusableDependences;
   DenseMap<Operation *, SmallVector<AffineMap, 1>> fusedProducerIndexingMap;
   for (LinalgOp op : reverse(ops)) {
     for (OpOperand &opOperand : op.getShapedOpOperands()) {
       Optional<LinalgDependenceGraph::LinalgDependenceGraphElem>

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

if (failed(rewriteAsPaddedOp(rewriter, *res, options))) {
     return failure();
   }
 
   // Do not perform replacement of `linalgOp`, let the derived patterns
   // do this as they see fit, from the resulting TiledLinalgOp.
   return success();
 }
 
+static ValueRange getTiledOpResult(TiledLinalgOp tiledOp) {
+  if (tiledOp.loops.empty())
+    return tiledOp.op.getOperation()->getResults();
+  return tiledOp.loops.front()->getResults();
+}
+
+static ValueRange
+getTiledAndFusedOpResult(TiledAndFusedLinalgOps tiledAndFusedOp) {
+  if (tiledAndFusedOp.fusedLoops.empty())
+    return tiledAndFusedOp.op.getOperation()->getResults();
+  return tiledAndFusedOp.fusedLoops.front()->getResults();
+}
+
 mlir::linalg::LinalgBaseTileAndFusePattern::LinalgBaseTileAndFusePattern(
     StringRef opName, MLIRContext *context,
     const LinalgDependenceGraph &dependenceGraph,
     LinalgTilingOptions tilingOptions, LinalgFusionOptions fusionOptions,
     LinalgTransformationFilter marker, LinalgTransformationFilter fusedOpMarker,
     LinalgTransformationFilter originalOpMarker, PatternBenefit benefit)
     : RewritePattern(opName, {}, benefit, context),
       dependenceGraph(dependenceGraph), tilingOptions(tilingOptions),
       fusionOptions(fusionOptions), marker(marker),
       fusedOpMarker(fusedOpMarker), originalOpMarker(originalOpMarker) {}
 
 LogicalResult mlir::linalg::LinalgBaseTileAndFusePattern::matchAndRewrite(
     Operation *op, PatternRewriter &rewriter) const {
   LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
   if (!linalgOp)
     return failure();
   if (failed(marker.checkAndNotify(rewriter, linalgOp)))
     return failure();
-  if (!linalgOp.hasBufferSemantics())
-    return failure();
 
   DenseSet<Operation *> producers;
   producers.insert(linalgOp);
   for (auto dependence : dependenceGraph.getDependentOperationsInto(linalgOp)) {
     Optional<unsigned> operandNumber = dependence.getIndexingOpViewOperandNum();
     // When looking at dependences into, indexingOp is always OpOperand. We
     // could assert, but continue if this is not the case.
     if (!operandNumber)
if (llvm::any_of(unfusedLoopTileSizes, [](Value val) {
         return true;
       })) {
     LinalgTilingOptions unfusedTilingOptions = tilingOptions;
     unfusedTilingOptions.setTileSizes(unfusedLoopTileSizes);
     Optional<TiledLinalgOp> unfusedTiledOp =
         tileLinalgOp(rewriter, tiledAndFusedOps->op, unfusedTilingOptions);
     if (!unfusedTiledOp)
       return failure();
-    rewriter.eraseOp(tiledAndFusedOps->op);
+    rewriter.replaceOp(tiledAndFusedOps->op,
+                       getTiledOpResult(unfusedTiledOp.getValue()));
     tiledAndFusedOps->op = unfusedTiledOp->op;
   }
+  op->replaceAllUsesWith(getTiledAndFusedOpResult(tiledAndFusedOps.getValue()));
 
   marker.replaceLinalgTransformationFilter(rewriter,
                                            tiledAndFusedOps->op.getOperation());
   for (auto fusedOp : tiledAndFusedOps->fusedProducers) {
     fusedOpMarker.replaceLinalgTransformationFilter(rewriter,
                                                     fusedOp.getOperation());
   }
   for (auto origProducerOp : ArrayRef<LinalgOp>(fusionOps).drop_back()) {

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

+// RUN: mlir-opt %s -test-linalg-tensor-fusion-transform-patterns -canonicalize -cse -split-input-file -verify-diagnostics | FileCheck %s
+
+module {
+  func @matmul_fusion(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>,
+                      %arg2: tensor<?x?xf32>, %arg3: tensor<?x?xf32>,
+                      %arg4: tensor<?x?xf32>) -> tensor<?x?xf32> {
+    %0 = linalg.matmul ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
+      outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>   // <MxN1> <N1xN2>
+    %1 = linalg.matmul {__internal_linalg_transform__ = "lhs_fusion"}
+      ins(%0, %arg3 : tensor<?x?xf32>, tensor<?x?xf32>)
+      outs(%arg4 : tensor<?x?xf32>) -> tensor<?x?xf32>   // <MxN2> <N2xN3>
+    return %1 : tensor<?x?xf32>
+  }
+}
+//  CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0)[s0] -> (32, -d0 + s0)>
+//  CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1) -> (32, d0 - d1)>
+//  CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0)[s0] -> (16, -d0 + s0)>
+//  CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0)[s0] -> (64, -d0 + s0)>
+//  CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1) -> (64, d0 - d1)>
+//      CHECK: func @matmul_fusion
+// CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
+// CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
+// CHECK-SAME:   %[[ARG2:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
+// CHECK-SAME:   %[[ARG3:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
+// CHECK-SAME:   %[[ARG4:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
+//  CHECK-DAG:   %[[C0:.+]] = constant 0 : index
+//  CHECK-DAG:   %[[C1:.+]] = constant 1 : index
+//  CHECK-DAG:   %[[C32:.+]] = constant 32 : index
+//  CHECK-DAG:   %[[C64:.+]] = constant 64 : index
+//  CHECK-DAG:   %[[C16:.+]] = constant 16 : index
+//  CHECK-DAG:   %[[M:.+]] = dim %[[ARG0]], %[[C0]]
+//      CHECK:   %[[RESULT:.+]] = scf.for %[[IV0:[a-zA-Z0-9]+]] =
+// CHECK-SAME:     %[[C0]] to %[[M]] step %[[C32]]
+// CHECK-SAME:     iter_args(%[[ARG6:.+]] = %[[ARG4]]) -> (tensor<?x?xf32>) {
+//      CHECK:     %[[TILE_M:.+]] = affine.min #[[MAP0]](%[[IV0]])[%[[M]]]
+//      CHECK:     %[[M_2:.+]] = dim %[[ARG6]], %[[C0]]
+//      CHECK:     %[[TILE_M_2:.+]] = affine.min #[[MAP1]](%[[M_2]], %[[IV0]])
+//      CHECK:     %[[N3:.+]] = dim %[[ARG6]], %[[C1]]
+//      CHECK:     %[[ST_ARG6:.+]] = subtensor %[[ARG6]][%[[IV0]], 0]
+// CHECK-SAME:       [%[[TILE_M_2]], %[[N3]]]
+//      CHECK:     %[[N2:.+]] = dim %[[ARG1]], %[[C1]]
+//      CHECK:     %[[N1:.+]] = dim %[[ARG0]], %[[C1]]
+//      CHECK:     %[[ST_ARG0:.+]] = subtensor %[[ARG0]][%[[IV0]], 0]
+// CHECK-SAME:       [%[[TILE_M]], %[[N1]]]
+//      CHECK:     %[[ST_ARG1:.+]] = subtensor %[[ARG1]][0, 0]
+// CHECK-SAME:       [%[[N1]], %[[N2]]]
+//      CHECK:     %[[ST_ARG2:.+]] = subtensor %[[ARG2]][%[[IV0]], 0]
+// CHECK-SAME:       [%[[TILE_M]], %[[N2]]]
+//      CHECK:     %[[LHS:.+]] = linalg.matmul
+// CHECK-SAME:       __internal_linalg_transform__ = "after_lhs_fusion_producer"
+// CHECK-SAME:       ins(%[[ST_ARG0]], %[[ST_ARG1]] : tensor<?x?xf32>, tensor<?x?xf32>)
+// CHECK-SAME:       outs(%[[ST_ARG2]] : tensor<?x?xf32>)
+//      CHECK:     %[[N3_2:.+]] = dim %[[ARG3]], %[[C1]]
+//      CHECK:     %[[YIELD0:.+]] = scf.for %[[IV1:[a-zA-Z0-9]+]] =
+// CHECK-SAME:       %[[C0]] to %[[N3_2]] step %[[C64]]
+// CHECK-SAME:       iter_args(%[[ARG8:.+]] = %[[ST_ARG6]]) -> (tensor<?x?xf32>) {
+//      CHECK:       %[[YIELD1:.+]] = scf.for %[[IV2:[a-zA-Z0-9]+]] =
+// CHECK-SAME:         %[[C0]] to %[[N2]] step %[[C16]]
+// CHECK-SAME:         iter_args(%[[ARG10:.+]] = %[[ARG8]]) -> (tensor<?x?xf32>) {
+//      CHECK:         %[[TILE_N2:.+]] = affine.min #[[MAP2]](%[[IV2]])[%[[N2]]]
+//      CHECK:         %[[ST_LHS:.+]] = subtensor %[[LHS]][0, %[[IV2]]]
+// CHECK-SAME:           [%[[TILE_M]], %[[TILE_N2]]]
+//      CHECK:         %[[N2_3:.+]] = dim %[[ARG3]], %[[C0]]
+//      CHECK:         %[[TILE_N2_2:.+]] = affine.min #[[MAP2]](%[[IV2]])[%[[N2_3]]]
+//      CHECK:         %[[TILE_N3:.+]] = affine.min #[[MAP3]](%[[IV1]])[%[[N3_2]]]
+//      CHECK:         %[[ST_ARG3:.+]] = subtensor %[[ARG3]][%[[IV2]], %[[IV1]]]
+// CHECK-SAME:           [%[[TILE_N2_2]], %[[TILE_N3]]]
+//      CHECK:         %[[M_4:.+]] = dim %[[ARG10]], %[[C0]]
+//      CHECK:         %[[N3_3:.+]] = dim %[[ARG10]], %[[C1]]
+//      CHECK:         %[[TILE_N3_2:.+]] = affine.min #[[MAP4]](%[[N3_3]], %[[IV1]])
+//      CHECK:         %[[ST_ARG4:.+]] = subtensor %[[ARG10]][0, %[[IV1]]]
+// CHECK-SAME:           [%[[M_4]], %[[TILE_N3_2]]]
+//      CHECK:         %[[ST_RESULT:.+]] = linalg.matmul
+// CHECK-SAME:           __internal_linalg_transform__ = "after_lhs_fusion"
+// CHECK-SAME:           ins(%[[ST_LHS]], %[[ST_ARG3]]
+// CHECK-SAME:             : tensor<?x?xf32>, tensor<?x?xf32>)
+// CHECK-SAME:           outs(%[[ST_ARG4]] : tensor<?x?xf32>)
+//      CHECK:         %[[UPDATE1:.+]] = subtensor_insert %[[ST_RESULT]]
+// CHECK-SAME:           into %[[ARG10]][0, %[[IV1]]] [%[[M_4]], %[[TILE_N3_2]]]
+//      CHECK:         scf.yield %[[UPDATE1]]
+//      CHECK:       }
+//      CHECK:       scf.yield %[[YIELD1]]
+//      CHECK:     }
+//      CHECK:     %[[UPDATE0:.+]] = subtensor_insert %[[YIELD0]] into
+// CHECK-SAME:       %[[ARG6]][%[[IV0]], 0] [%[[TILE_M_2]], %[[N3]]]
+//      CHECK:     scf.yield %[[UPDATE0]]
+//      CHECK:   }
+//      CHECK:   return %[[RESULT]]
+
+// -----
+
+module {
+  func @matmul_plus_matmul(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>,
+                           %arg2: tensor<?x?xf32>) -> tensor<?x?xf32>{
+    %c0 = constant 0 : index
+    %c1 = constant 1 : index
+    %0 = dim %arg2, %c0 : tensor<?x?xf32>
+    %1 = dim %arg2, %c1 : tensor<?x?xf32>
+    %2 = linalg.matmul ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
+      outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
+    %3 = dim %2, %c0 : tensor<?x?xf32>
+    %4 = dim %2, %c1 : tensor<?x?xf32>
+    %5 = linalg.init_tensor [%3, %4] : tensor<?x?xf32>
+    %6 = linalg.generic
+      {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+                        affine_map<(d0, d1) -> (d0, d1)>,
+                        affine_map<(d0, d1) -> (d0, d1)>],
+       iterator_types = ["parallel", "parallel"],
+       __internal_linalg_transform__ = "transpose_fusion"}
+      ins(%2, %2 : tensor<?x?xf32>, tensor<?x?xf32>)
+      outs(%5 : tensor<?x?xf32>) {
+      ^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32) :
+        %7 = addf %arg3, %arg4 : f32
+        linalg.yield %7 : f32
+      } -> tensor<?x?xf32>
+    return %6 : tensor<?x?xf32>
+  }
+}
+//       CHECK: func @matmul_plus_matmul
+//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
+//  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
+//  CHECK-SAME:   %[[ARG2:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
+//       CHECK:   %[[RESULT:.+]] = scf.for %[[IV0:[a-zA-Z0-9_]+]]
+//  CHECK-SAME:     iter_args(%[[ARG4:.+]] = %{{[a-zA-Z0-9_]+}})
+//       CHECK:     %[[YIELD:.+]] = scf.for %[[IV1:[a-zA-Z0-9_]+]]
+//  CHECK-SAME:       iter_args(%[[ARG6:.+]] = %[[ARG4]])
+//       CHECK:       %[[ST_ARG6:.+]] = subtensor %[[ARG6]][%[[IV0]], %[[IV1]]]
+//       CHECK:       %[[ST_ARG0:.+]] = subtensor %[[ARG0]][%[[IV0]], 0]
+//       CHECK:       %[[ST_ARG1:.+]] = subtensor %[[ARG1]][0, %[[IV1]]]
+//       CHECK:       %[[ST_ARG2:.+]] = subtensor %[[ARG2]][%[[IV0]], %[[IV1]]]
+//       CHECK:       %[[LHS:.+]] = linalg.matmul
+//  CHECK-SAME:         ins(%[[ST_ARG0]], %[[ST_ARG1]]
+//  CHECK-SAME:           : tensor<?x?xf32>, tensor<?x?xf32>)
+//  CHECK-SAME:         outs(%[[ST_ARG2]] : tensor<?x?xf32>)
+//       CHECK:       %[[ST_RESULT:.+]] = linalg.generic
+//  CHECK-SAME:         ins(%[[LHS]] : tensor<?x?xf32>)
+//  CHECK-SAME:         outs(%[[ST_ARG6]] : tensor<?x?xf32>)
+//       CHECK:       %[[UPDATE:.+]] = subtensor_insert %[[ST_RESULT]]
+//  CHECK-SAME:         into %[[ARG6]][%[[IV0]], %[[IV1]]]
+//       CHECK:       scf.yield %[[UPDATE]]
+//       CHECK:     scf.yield %[[YIELD]]
+//       CHECK:   return %[[RESULT]]

mlir/test/lib/Transforms/TestLinalgFusionTransforms.cpp

 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/Passes.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 
-namespace {
-struct TestLinalgFusionTransforms
-    : public PassWrapper<TestLinalgFusionTransforms, FunctionPass> {
-  TestLinalgFusionTransforms() = default;
-  TestLinalgFusionTransforms(const TestLinalgFusionTransforms &pass) {}
-
-  void getDependentDialects(DialectRegistry &registry) const override {
-    registry.insert<AffineDialect, linalg::LinalgDialect, scf::SCFDialect,
-                    StandardOpsDialect>();
-  }
-
-  void runOnFunction() override;
-};
-} // namespace
-
+template <LinalgTilingLoopType LoopType>
 static void fillFusionPatterns(MLIRContext *context,
                                const LinalgDependenceGraph &dependenceGraph,
                                OwningRewritePatternList &patterns) {
   patterns.insert<LinalgTileAndFusePattern<MatmulOp>,
                   LinalgTileAndFusePattern<ConvOp>>(
       context, dependenceGraph,
-      LinalgTilingOptions()
-          .setTileSizes({32, 64, 16})
-          .setLoopType(LinalgTilingLoopType::ParallelLoops),
+      LinalgTilingOptions().setTileSizes({32, 64, 16}).setLoopType(LoopType),
       LinalgFusionOptions().setIndicesToFuse({2}),
       LinalgTransformationFilter(
           Identifier::get("basic_fusion", context),
           Identifier::get("after_basic_fusion", context)),
       LinalgTransformationFilter(
           ArrayRef<Identifier>(),
           Identifier::get("after_basic_fusion_producer", context)),
       LinalgTransformationFilter(
           ArrayRef<Identifier>(),
           Identifier::get("after_basic_fusion_original", context)));
 
   patterns.insert<LinalgTileAndFusePattern<MatmulOp>>(
       context, dependenceGraph,
-      LinalgTilingOptions()
-          .setTileSizes({32, 64, 16})
-          .setLoopType(LinalgTilingLoopType::ParallelLoops),
+      LinalgTilingOptions().setTileSizes({32, 64, 16}).setLoopType(LoopType),
       LinalgFusionOptions().setIndicesToFuse({0}),
       LinalgTransformationFilter(Identifier::get("lhs_fusion", context),
                                  Identifier::get("after_lhs_fusion", context)),
       LinalgTransformationFilter(
           ArrayRef<Identifier>(),
           Identifier::get("after_lhs_fusion_producer", context)),
       LinalgTransformationFilter(
           ArrayRef<Identifier>(),
           Identifier::get("after_lhs_fusion_original", context)));
 
   patterns.insert<LinalgTileAndFusePattern<MatmulOp>>(
       context, dependenceGraph,
-      LinalgTilingOptions()
-          .setTileSizes({32, 64, 16})
-          .setLoopType(LinalgTilingLoopType::ParallelLoops),
+      LinalgTilingOptions().setTileSizes({32, 64, 16}).setLoopType(LoopType),
       LinalgFusionOptions().setIndicesToFuse({1}),
       LinalgTransformationFilter(Identifier::get("rhs_fusion", context),
                                  Identifier::get("after_rhs_fusion", context)),
       LinalgTransformationFilter(
           ArrayRef<Identifier>(),
           Identifier::get("after_rhs_fusion_producer", context)),
       LinalgTransformationFilter(
           ArrayRef<Identifier>(),
           Identifier::get("after_rhs_fusion_original", context)));
 
   patterns.insert<LinalgTileAndFusePattern<MatmulOp>>(
       context, dependenceGraph,
-      LinalgTilingOptions()
-          .setTileSizes({32, 64, 16})
-          .setLoopType(LinalgTilingLoopType::ParallelLoops),
+      LinalgTilingOptions().setTileSizes({32, 64, 16}).setLoopType(LoopType),
       LinalgFusionOptions().setIndicesToFuse({0, 2}),
       LinalgTransformationFilter(
           Identifier::get("two_operand_fusion", context),
           Identifier::get("after_two_operand_fusion", context)),
       LinalgTransformationFilter(
           ArrayRef<Identifier>(),
           Identifier::get("after_two_operand_fusion_producer", context)),
       LinalgTransformationFilter(
           ArrayRef<Identifier>(),
           Identifier::get("after_two_operand_fusion_original", context)));
 
   patterns.insert<LinalgTileAndFusePattern<GenericOp>>(
       context, dependenceGraph,
-      LinalgTilingOptions().setTileSizes({32, 64}).setLoopType(
-          LinalgTilingLoopType::ParallelLoops),
+      LinalgTilingOptions().setTileSizes({32, 64}).setLoopType(LoopType),
       LinalgFusionOptions().setIndicesToFuse({0, 1}),
       LinalgTransformationFilter(
           Identifier::get("transpose_fusion", context),
           Identifier::get("after_transpose_fusion", context)),
       LinalgTransformationFilter(
           ArrayRef<Identifier>(),
           Identifier::get("after_transpose_fusion_producer", context)),
       LinalgTransformationFilter(
           ArrayRef<Identifier>(),
           Identifier::get("after_transpose_fusion_original", context)));
 }
 
-static void applyFusionPatterns(MLIRContext *context, FuncOp funcOp) {
+namespace {
+template <LinalgTilingLoopType LoopType = LinalgTilingLoopType::ParallelLoops>
+struct TestLinalgFusionTransforms
+    : public PassWrapper<TestLinalgFusionTransforms<LoopType>, FunctionPass> {
+  TestLinalgFusionTransforms() = default;
+  TestLinalgFusionTransforms(const TestLinalgFusionTransforms &pass) {}
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<AffineDialect, linalg::LinalgDialect, scf::SCFDialect,
+                    StandardOpsDialect>();
+  }
+
+  void runOnFunction() override {
+    MLIRContext *context = &this->getContext();
+    FuncOp funcOp = this->getFunction();
     OwningRewritePatternList fusionPatterns;
     Aliases alias;
     LinalgDependenceGraph dependenceGraph =
         LinalgDependenceGraph::buildDependenceGraph(alias, funcOp);
-  fillFusionPatterns(context, dependenceGraph, fusionPatterns);
+    fillFusionPatterns<LoopType>(context, dependenceGraph, fusionPatterns);
     applyPatternsAndFoldGreedily(funcOp, std::move(fusionPatterns));
   }
-
-void TestLinalgFusionTransforms::runOnFunction() {
-  applyFusionPatterns(&getContext(), getFunction());
-}
+};
+} // namespace
 
 static LogicalResult fuseLinalgOpsGreedily(FuncOp f) {
   OpBuilder b(f);
   DenseSet<Operation *> eraseSet;
 
   // Save original Linalg ops, we only want to make a pass over those.
   SmallVector<LinalgOp, 8> linalgOps;
   f.walk([&](LinalgOp op) {
if (!llvm::hasSingleElement(blocks)) {
       return;
     }
     SmallVector<LinalgOp, 2> linalgOps =
         llvm::to_vector<2>(blocks.front().getOps<LinalgOp>());
     Aliases aliases;
     LinalgDependenceGraph dependenceGraph(aliases, linalgOps);
     OpBuilder builder(funcOp.getContext());
     linalg::LinalgTilingLoopType loopType = LinalgTilingLoopType::ParallelLoops;
-    if (llvm::all_of(linalgOps, [](LinalgOp linalgOp) {
+    if (llvm::any_of(linalgOps, [](LinalgOp linalgOp) {
           return linalgOp.hasTensorSemantics();
         }))
       loopType = LinalgTilingLoopType::Loops;
     Optional<TiledAndFusedLinalgOps> tileAndFuseOps = tileAndFuseLinalgOps(
         builder, linalgOps, dependenceGraph,
         LinalgTilingOptions().setTileSizes(tileSizes).setLoopType(loopType));
     if (!tileAndFuseOps)
       return signalPassFailure();
     if (linalgOps.back().hasTensorSemantics()) {
       linalgOps.back().getOperation()->replaceAllUsesWith(
           tileAndFuseOps->fusedLoops.front());
     }
     for (auto op : linalgOps)
       if (op.hasBufferSemantics())
         op.erase();
   }
 };
 } // namespace
 
 namespace mlir {
 namespace test {
 void registerTestLinalgFusionTransforms() {
-  PassRegistration<TestLinalgFusionTransforms> testFusionTransformsPass(
+  PassRegistration<TestLinalgFusionTransforms<>> testFusionTransformsPass(
       "test-linalg-fusion-transform-patterns",
       "Test Linalg fusion transformation patterns by applying them greedily.");
 }
+void registerTestLinalgTensorFusionTransforms() {
+  PassRegistration<TestLinalgFusionTransforms<LinalgTilingLoopType::Loops>>
+      testTensorFusionTransformsPass(
+          "test-linalg-tensor-fusion-transform-patterns",
+          "Test Linalg on tensor fusion transformation "
+          "patterns by applying them greedily.");
+}
 void registerTestLinalgGreedyFusion() {
   PassRegistration<TestLinalgGreedyFusion> testFusionTransformsPass(
       "test-linalg-greedy-fusion",
       "Test Linalg fusion by applying a greedy test transformation.");
 }
 void registerTestLinalgTileAndFuseSequencePass() {
   PassRegistration<TestLinalgTileAndFuseSequencePass>
       testTileAndFuseSequencePass(
           "test-linalg-tile-and-fuse",
           "Test Linalg tiling and fusion of a sequence of Linalg operations.");
 }
 
 } // namespace test
 } // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestDialect(DialectRegistry &);
 void registerTestDominancePass();
 void registerTestDynamicPipelinePass();
 void registerTestExpandTanhPass();
 void registerTestGpuParallelLoopMappingPass();
 void registerTestInterfaces();
 void registerTestLinalgCodegenStrategy();
 void registerTestLinalgFusionTransforms();
+void registerTestLinalgTensorFusionTransforms();
 void registerTestLinalgGreedyFusion();
 void registerTestLinalgHoisting();
 void registerTestLinalgTileAndFuseSequencePass();
 void registerTestLinalgTransforms();
 void registerTestLivenessPass();
 void registerTestLoopFusion();
 void registerTestLoopMappingPass();
 void registerTestLoopUnrollingPass();
#endif
   test::registerTestDecomposeCallGraphTypes();
   test::registerTestDominancePass();
   test::registerTestDynamicPipelinePass();
   test::registerTestExpandTanhPass();
   test::registerTestGpuParallelLoopMappingPass();
   test::registerTestInterfaces();
   test::registerTestLinalgCodegenStrategy();
   test::registerTestLinalgFusionTransforms();
+  test::registerTestLinalgTensorFusionTransforms();
   test::registerTestLinalgGreedyFusion();
   test::registerTestLinalgHoisting();
   test::registerTestLinalgTileAndFuseSequencePass();
   test::registerTestLinalgTransforms();
   test::registerTestLivenessPass();
   test::registerTestLoopFusion();
   test::registerTestLoopMappingPass();
   test::registerTestLoopUnrollingPass();