[Transform][Fusion] Support fusion for K-tiling deep matmul (#217)

Fix several issues found in llama2 MLP workload:

1. Support `if-else` control block.
2. Cover `GenericOp` and `LinalgXOp`.
3. Enhance use-def check for residual pattern.

Author: Yun-Fly

include/gc/Transforms/Passes.td

@@ -59,25 +59,18 @@ def IterativeTilingAndFusion : Pass<"iterative-tiling-and-fusion",
       3. Multi-consumer and multi-producer support.
       4. Multiple level of nest loops and candidates.
       5. Flexible option to control the boundary of iterative process.
-      6. Cost-model to determine whether to fuse or not.
-
-    It intends to control the granularity of fusion by `fusion-level`, E.g.
-    * `0`: disable any fusion.
-    * `1`:[Default] enable both producer and consumer fusion, covering any tilable operation including tensor.pack/tensor.fill/linalg.reduce etc but excluding branches forked by multiple uses.
-    * `2`: `LEVEL 1` + extend to any topology including branches.
+      6. Default tiling when no op is tiled before fusion.
+      7. Cost-model to determine whether to fuse or not.
   }];
   let dependentDialects = ["func::FuncDialect", "linalg::LinalgDialect", "scf::SCFDialect",
                            "tensor::TensorDialect"];
 
   let options = [
-    Option<"fusionLevel", "fusion-level", "int64_t",
-           /*default=*/"1",
-           "Control the granularity of fusion.">,
     Option<"useCostModel", "use-cost-model", "bool",
            /*default=*/"false",
            "Decide if enable cost model to control iterative fusion.">,
     ListOption<"defaultTileSize", "default-tile-size", "std::string",
-           "Set default TileSize for the certain type of op, saying matmul:{32,32}">,
+           "Set default TileSize for the certain type of op, saying `matmul:{32,32}`">,
   ];
 }
 

lib/gc/Transforms/IterativeTilingAndFusion.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "gc/Analysis/TargetDescriptionAnalysis.h"
+#include "gc/Dialect/Linalgx/LinalgxOps.h"
 #include "gc/Transforms/Passes.h"
 #include "mlir/Analysis/TopologicalSortUtils.h"
 #include "mlir/Dialect/DLTI/Traits.h"
@@ -68,17 +69,22 @@ getClosestInsertSliceOfResult(OpResult result) {
       sliceOp =
           dyn_cast<OffsetSizeAndStrideOpInterface>(useOfResult.getOwner());
     } else if (auto yieldOp = dyn_cast<scf::YieldOp>(useOfResult.getOwner())) {
-      if (auto loop = dyn_cast<LoopLikeOpInterface>(yieldOp->getParentOp()))
+      if (isa<LoopLikeOpInterface, RegionBranchOpInterface>(
+              yieldOp->getParentOp())) {
         return getClosestInsertSliceOfResult(
-            loop->getResult(useOfResult.getOperandNumber()));
+            yieldOp->getParentOp()->getResult(useOfResult.getOperandNumber()));
+      }
+    } else if (isa<tensor::ExpandShapeOp, tensor::CollapseShapeOp>(
+                   useOfResult.getOwner())) {
+      return getClosestInsertSliceOfResult(
+          useOfResult.getOwner()->getResult(useOfResult.getOperandNumber()));
     }
   }
 
   if (!llvm::detail::isPresent(sliceOp))
     return failure();
-  else {
-    return sliceOp;
-  }
+
+  return sliceOp;
 }
 
 struct CandidateDefOrUse {
@@ -155,12 +161,12 @@ exactTilingOnPackUnPackFilter(RewriterBase &rewriter,
   if (auto packOp = dyn_cast<tensor::PackOp>(defOrUse.ownerOp)) {
     // tileSize comes from OpResult
     if (defOrUse.isDef()) {
-      targetInnerTileSizes = packOp.getInnerTiles();
+      targetInnerTileSizes = packOp.getMixedTiles();
       targetTileSizes = llvm::to_vector(
           ArrayRef(tileSizes).take_back(targetInnerTileSizes.size()));
     } else {
       // tileSize comes from OpOperand
-      targetTileSizes = llvm::to_vector(tileSizes);
+      targetTileSizes = tileSizes;
       DenseMap<int64_t, OpFoldResult> dimAndTileMapping =
           packOp.getDimAndTileMapping();
       targetInnerTileSizes.resize(dimAndTileMapping.size());
@@ -171,16 +177,18 @@ exactTilingOnPackUnPackFilter(RewriterBase &rewriter,
   } else if (auto unPackOp = dyn_cast<tensor::UnPackOp>(defOrUse.ownerOp)) {
     // tileSize comes from OpResult
     if (defOrUse.isDef()) {
-      targetTileSizes = llvm::to_vector(tileSizes);
+      targetTileSizes = tileSizes;
       DenseMap<int64_t, OpFoldResult> dimAndTileMapping =
           unPackOp.getDimAndTileMapping();
+      if (dimAndTileMapping.empty())
+        return failure();
       targetInnerTileSizes.resize(dimAndTileMapping.size());
       for (const auto &dimAndTile : dimAndTileMapping) {
         targetInnerTileSizes[dimAndTile.first] = dimAndTile.second;
       }
     } else {
       // tileSize comes from OpOperand
-      targetInnerTileSizes = unPackOp.getInnerTiles();
+      targetInnerTileSizes = unPackOp.getMixedTiles();
       targetTileSizes = llvm::to_vector(
           ArrayRef(tileSizes).take_back(targetInnerTileSizes.size()));
     }
@@ -481,21 +489,10 @@ tileAndFuseConsumerOfOpResult(RewriterBase &rewriter, OpResult result,
 
     if (succeeded(fusedResult)) {
       fusedResultList.push_back(*fusedResult);
-      auto whileProducerOutOfLoopBlock =
-          [&fusedResult](LoopLikeOpInterface loop) -> LogicalResult {
-        Block &body = loop->getRegion(0).front();
-        return failure(fusedResult.value().tiledOps[0]->getBlock() == &body);
-      };
-      SmallVector<LoopLikeOpInterface> outerLoops =
-          scfX::getOuterNestLoopsWhile(
-              (*bestCandidate)->getParentOfType<LoopLikeOpInterface>(),
-              whileProducerOutOfLoopBlock);
-      // g. Manually run cse on region which contains top-level loop of
-      // candidate slice in avoid of conflict with subsequent
-      // `tileAndFuseConsumerOfSlice` get nest loops between next candidate
-      // sliceOp and tiled producer.
-      (void)mlir::simplifyRegions(rewriter,
-                                  {*outerLoops.front()->getParentRegion()});
+      // f. Manually run cse on region which contains original consumer op in
+      // avoid of conflict with subsequent `tileAndFuseConsumerOfSlice` get nest
+      // loops between next candidate sliceOp and tiled producer.
+      (void)mlir::simplifyRegions(rewriter, {*consumer->getParentRegion()});
     }
   }
   if (fusedResultList.empty())
@@ -543,7 +540,13 @@ LogicalResult iterativelyFuseProducerAndConsumerOfTiledOp(
   return success(numTiledOps > 1);
 }
 
-/// What is self tiled op compared with other fused op?
+/// This is a workaround to deal with LinalgXOp
+static bool isTilableLinalgXOp(Operation *op) {
+  return isa<linalgx::BatchReduceMatmulVnniOp, linalgx::MultiBatchMatmulOp,
+             linalgx::Mm2DVnniOp, linalgx::Mm4DVnniOp>(op);
+}
+
+/// Check if tiled op inside a loop?
 /// E.g.
 /// %1 = scf.for(){
 ///   %2 = scf.for(){
@@ -553,25 +556,17 @@ LogicalResult iterativelyFuseProducerAndConsumerOfTiledOp(
 ///       yield %5
 ///   }
 /// }
-static LogicalResult isSelfTiledOp(Operation *targetOp) {
-  // 0. check tilable
-  if (!isa<TilingInterface>(targetOp))
+static LogicalResult isTiledOpInLoop(Operation *targetOp) {
+  // 1. check tilable
+  if (!isa<TilingInterface>(targetOp) && !isTilableLinalgXOp(targetOp))
     return failure();
-  // 1. check parentOp
+  // 2. check parentOp
   auto forOp = targetOp->getParentOfType<LoopLikeOpInterface>();
   if (!forOp)
     return failure();
-  // 2. check single one tiling interface in loop body
-  auto walkResult = forOp->walk([&targetOp](TilingInterface op) {
-    // some special op maybe already deal with in template
-    if (isa<linalg::FillOp, linalg::CopyOp>(op))
-      return WalkResult::skip();
-    return op != targetOp ? WalkResult::interrupt() : WalkResult::advance();
-  });
-  if (walkResult.wasInterrupted())
-    return failure();
+
   // 3. check whether has either extract or insert slice op
-  walkResult = forOp->walk(
+  auto walkResult = forOp->walk(
       [](tensor::ExtractSliceOp) { return WalkResult::interrupt(); });
   if (walkResult.wasInterrupted())
     return success();
@@ -583,11 +578,12 @@ static LogicalResult isSelfTiledOp(Operation *targetOp) {
 using OpTileSizeMap = std::unordered_map<std::string, SmallVector<int64_t>>;
 
 template <typename OpTy>
-static bool defaultTilingOfType(RewriterBase &rewriter, Operation *op,
-                                const OpTileSizeMap &tsMap) {
+static FailureOr<scf::SCFTilingResult>
+defaultTilingOfType(RewriterBase &rewriter, Operation *op,
+                    const OpTileSizeMap &tsMap) {
   // a. Check <OpTy>
   if (!isa<TilingInterface>(op) || !isa<OpTy>(op))
-    return false;
+    return failure();
   auto tilingInterfaceOp = cast<TilingInterface>(op);
 
   scf::SCFTilingOptions options;
@@ -618,26 +614,29 @@ static bool defaultTilingOfType(RewriterBase &rewriter, Operation *op,
   }
   // If the tile sizes are all zero, no tiling would happen.
   if (llvm::all_of(defaultTileSize, isZeroIndex))
-    return false;
+    return failure();
 
   options.setTileSizes(defaultTileSize);
   // c. Set loop type
   options.setLoopType(scf::SCFTilingOptions::LoopType::ForallOp);
   // d. Use builtin tiling interface
   FailureOr<scf::SCFTilingResult> tilingResult =
       scf::tileUsingSCF(rewriter, tilingInterfaceOp, options);
-  if (succeeded(tilingResult)) {
-    rewriter.replaceOp(op, tilingResult->replacements);
-    return true;
-  }
-  return false;
+
+  if (failed(tilingResult))
+    return failure();
+
+  return tilingResult;
 }
 
+using DefaultTilingFn = std::function<FailureOr<scf::SCFTilingResult>(
+    RewriterBase &, Operation *, const OpTileSizeMap &)>;
+
 void iterativeTilingAndFusionUntilExhaustion(
     RewriterBase &rewriter, func::FuncOp &f,
     const CandidateSliceOptions &sliceOptions, const OpTileSizeMap &tsMap) {
   // Collect untiled and tiled ops respectively
-  llvm::SetVector<Operation *> selfTiledOp, unTiledOps;
+  llvm::SetVector<Operation *> tiledOps, unTiledOps;
 
   auto collectUnTiledOps = [&f, &unTiledOps]() -> bool {
     // Reset
@@ -648,67 +647,64 @@ void iterativeTilingAndFusionUntilExhaustion(
         return WalkResult::skip();
       if (isa<TilingInterface>(op)) {
         auto parentLoop = op->getParentOfType<LoopLikeOpInterface>();
-        if (!parentLoop.getOperation())
+        auto parentGeneric = op->getParentOfType<linalg::GenericOp>();
+        if (!llvm::detail::isPresent(parentLoop) &&
+            !llvm::detail::isPresent(parentGeneric))
           unTiledOps.insert(op);
       }
       return WalkResult::advance();
     });
     return !unTiledOps.empty();
   };
 
-  auto collectSelfTiledOp = [&f, &selfTiledOp]() -> bool {
-    // Reset
-    selfTiledOp.clear();
-    // Walk through funcOp
-    f->walk([&selfTiledOp](Operation *op) {
-      // Target at certain kind of tiled op, such as matmul/conv implemented
-      // by multiple level of nest loops and candidate slices for better
-      // utilization of parallelism and memory hierarchy.
-      if (succeeded(isSelfTiledOp(op))) {
-        selfTiledOp.insert(op);
-      }
-    });
-    return !selfTiledOp.empty();
-  };
+  // Walk through funcOp
+  f->walk([&tiledOps](Operation *op) {
+    if (succeeded(isTiledOpInLoop(op))) {
+      tiledOps.insert(op);
+    }
+  });
 
   // Iterative tiling and fusion until exhaustion.
   while (collectUnTiledOps()) {
     // If existing tiled op before tiling.
-    if (collectSelfTiledOp()) {
+    if (!tiledOps.empty()) {
       // Sort by topology
-      mlir::topologicalSort(selfTiledOp);
+      mlir::topologicalSort(tiledOps);
       // Record if any fusion happens
       bool changed = false;
       // Iteratively fuse in forward and backward fashion.
-      llvm::for_each(selfTiledOp, [&rewriter, &sliceOptions,
-                                   &changed](Operation *tiledOp) {
-        changed |= succeeded(iterativelyFuseProducerAndConsumerOfTiledOp(
-            rewriter, tiledOp, sliceOptions));
-      });
+      llvm::for_each(
+          tiledOps, [&rewriter, &sliceOptions, &changed](Operation *tiledOp) {
+            changed |= succeeded(iterativelyFuseProducerAndConsumerOfTiledOp(
+                rewriter, tiledOp, sliceOptions));
+          });
+      tiledOps.clear();
       if (changed)
-        (void)mlir::simplifyRegions(rewriter, {f.getRegion()});
+        (void)mlir::simplifyRegions(rewriter, f->getRegions());
     } else {
       // Auto tiling with default tile size if no tiled op found. Follow tiling
       // priority based on OpTy: `Contraction`->`Reduction`->`Elementwise`.
-      SmallVector<std::function<bool(RewriterBase &, Operation *,
-                                     const OpTileSizeMap &)>>
-          priorityTilingPipeLine = {
-              defaultTilingOfType<mlir::linalg::ContractionOpInterface>,
-              defaultTilingOfType<mlir::linalg::ReduceOp>,
-              defaultTilingOfType<mlir::linalg::LinalgOp>};
-      if (llvm::all_of(priorityTilingPipeLine,
-                       [&rewriter, &tsMap, &unTiledOps](
-                           function_ref<bool(RewriterBase &, Operation *,
-                                             const OpTileSizeMap &)>
-                               tilingFn) {
-                         return !llvm::any_of(
-                             unTiledOps, std::bind(tilingFn, std::ref(rewriter),
-                                                   std::placeholders::_1,
-                                                   std::cref(tsMap)));
-                       })) {
-        // If no op can be tiled
-        return;
+      SmallVector<DefaultTilingFn> priorityTilingPipeLine = {
+          defaultTilingOfType<mlir::linalg::ContractionOpInterface>,
+          defaultTilingOfType<mlir::linalg::ReduceOp>,
+          defaultTilingOfType<TilingInterface>};
+
+      for (auto &tilingFn : priorityTilingPipeLine) {
+        for (auto &op : unTiledOps) {
+          FailureOr<scf::SCFTilingResult> tilingResult =
+              tilingFn(rewriter, op, tsMap);
+          if (succeeded(tilingResult)) {
+            tiledOps.insert(tilingResult->tiledOps[0]);
+            rewriter.replaceOp(op, tilingResult->replacements);
+            break;
+          }
+        }
+        if (!tiledOps.empty())
+          break;
       }
+      // If no op can be tiled
+      if (tiledOps.empty())
+        return;
     }
   }
 }