extend LinalgBlockPackMatmul pass to support mmt4d op

Author: yifeizh2

lib/gc/Transforms/PropagateLayout.cpp

@@ -272,12 +272,37 @@ LogicalResult namedOpLayoutPropagation(MLIRContext *ctx, mlir::Operation *graph,
 void PropagateLayoutOnNamedOps::runOnOperation() {
   MLIRContext *ctx = &getContext();
   mlir::Operation *graph = getOperation();
-  ControlPackNamedOpsFn controlFn =
+  // stage1:
+  RewritePatternSet patterns(&getContext());
+  mlir::linalg::ControlBlockPackMatmulFn packMatmulControlFn =
+      [&](linalg::LinalgOp op) -> mlir::linalg::BlockPackMatmulOptions {
+    mlir::linalg::BlockPackMatmulOptions options;
+    auto &layoutAnalysisResult = getAnalysis<GlobalAnalysis>();
+    auto matmulLayout = *(layoutAnalysisResult.getOpLayout(op));
+    TensorLayout LHSLayout = matmulLayout.getSupportedInputLayouts()[0];
+    TensorLayout RHSLayout = matmulLayout.getSupportedInputLayouts()[1];
+    // hardcode to mmt4d format
+    options.rhsTransposeOuterBlocks = true;
+    options.rhsTransposeInnerBlocks = true;
+    options.blockFactors.push_back(
+        *getConstantIntValue(LHSLayout.getTileSizes()[0]));
+    options.blockFactors.push_back(
+        *getConstantIntValue(LHSLayout.getTileSizes()[1]));
+    options.blockFactors.push_back(
+        *getConstantIntValue(RHSLayout.getTileSizes()[1]));
+    return options;
+  };
+  linalg::populateBlockPackMatmulPatterns(patterns, packMatmulControlFn);
+  if (failed(applyPatternsAndFoldGreedily(graph, std::move(patterns))))
+    return signalPassFailure();
+
+  // stage3: propagate layout on other namsed ops
+  ControlPackNamedOpsFn layoutControlFn =
       [&](Operation *op) -> FailureOr<OperatorLayout> {
     auto &layoutAnalysisResult = getAnalysis<GlobalAnalysis>();
     return layoutAnalysisResult.getOpLayout(op);
   };
-  if (failed(namedOpLayoutPropagation(ctx, graph, controlFn)))
+  if (failed(namedOpLayoutPropagation(ctx, graph, layoutControlFn)))
     return signalPassFailure();
 }
 

packMatmul.patch

@@ -0,0 +1,32 @@
+diff --git a/mlir/lib/Dialect/Linalg/Transforms/BlockPackMatmul.cpp b/mlir/lib/Dialect/Linalg/Transforms/BlockPackMatmul.cpp
+index 91d4efa3372b..f3f61ff92140 100644
+--- a/mlir/lib/Dialect/Linalg/Transforms/BlockPackMatmul.cpp
++++ b/mlir/lib/Dialect/Linalg/Transforms/BlockPackMatmul.cpp
+@@ -210,6 +210,19 @@ linalg::blockPackMatmul(RewriterBase &rewriter, linalg::LinalgOp linalgOp,
+   packedMatmul->packOps[1] = packedRhs->transposedPackOp;
+   packedMatmul->packedLinalgOp = packedRhs->transposedLinalgOp;
+ 
++  // rewrite generic to mmt4d
++  if (!options->lhsTransposeOuterBlocks && !options->lhsTransposeInnerBlocks &&
++      options->rhsTransposeOuterBlocks && options->rhsTransposeInnerBlocks &&
++      options->mnkOrder == SmallVector<int64_t>{0, 1, 2}) {
++    auto originalLinalgOp = packedMatmul->packedLinalgOp;
++    rewriter.setInsertionPoint(originalLinalgOp);
++    auto mmt4d = rewriter.create<linalg::Mmt4DOp>(
++        originalLinalgOp.getLoc(), originalLinalgOp.getDpsInits().getTypes(),
++        originalLinalgOp.getDpsInputs(), originalLinalgOp.getDpsInits());
++    rewriter.replaceOp(originalLinalgOp, mmt4d);
++    packedMatmul->packedLinalgOp = mmt4d;
++  }
++
+   return packedMatmul;
+ }
+ 
+@@ -307,6 +320,7 @@ struct LinalgBlockPackMatmul
+ };
+ } // namespace
+ 
++// extend to transform to mmt4d or batch_mmt4d
+ void linalg::populateBlockPackMatmulPatterns(
+     RewritePatternSet &patterns, const ControlBlockPackMatmulFn &controlFn) {
+   patterns.add<BlockPackMatmul<linalg::GenericOp>,