Repository: alibaba/BladeDISC (language: python)

Issue:
Introduce MLIR transform dialect to BladeDISC
We'll start to explore using MLIR [transform dialect ](https://mlir.llvm.org/docs/Dialects/Transform) to do codegen for (fused) compute-intensive pattern. The initial target is to support gemm codegen on ARM platform to address the dynamic shape problem of Arm Compute Library.

The initial plan is:

- [x]  Step 1,  enhance the fusion decision pass. We’ll add a new fusion kind `kTransform` for the transform-based fusion pattern.
- [x]  Step 2,  lower the lmhlo fusion op to linalg on tensor.
- [x]  Step 3,  transform the linalg computation to loops using transform dialect.
- [x]  Step 4,  refined the transformed loop to make it suitable for BladeDISC runtime.
- [x]  Step 5, add a new pass to the disc pass pipeline to drive the above process.
- [x]   Step 6, weight pre-packing support
   - [x]  add `disc_linalg.multi_level_pack` op, used for doing packing.
   - [x]  add  `transform.disc.cache_read` transform op, relying on `disc_linalg.multi_level_pack` op.
   - [x]  add folding support for `disc_linalg.multi_level_pack`.
   - [x]  lower `disc_linalg.multi_level_pack` to loop if it can not be folded.
   - [x]  fuse const weight op into the `kTransform` fusion pattern, lower it to linalg and then schedule it.
- [x]  Step 7,  assign a default schedule for each `kTransform` pattern.  
- [x]  Step 8,  schedule selection logic injection
- [x]  Step 9,  initial model level testing:  bert (albert).
- [x]  Step 10, support nt, tn, tt format GEMM.
- [ ]  Step 11, support batch matmul
- [x]  Step 12, support GEMM epilogue fusion.
- [ ]  Step 13, performance optimization


Discussion / hints:
e2e model test on: Bert Base (TF) and Albert (PyTorch),  on g6r, using single thread. Note that we only have one default schedule for all shape and the schedule is known to be less performant when n or k is large, thus the initial performance is supposed to be improved when we support schedule selection logic.

Bert Base (TF) 

input | TF 2.8(s) | DISC-ACL(s) | DISC-Transform(s) | speedup (DISC-transform / DISC-ACL)
-- | -- | -- | -- | --
(1, 128) | 0.742 | 0.638 | 0.656 | 97.3%
(2, 128) | 1.41 | 1.24 | 1.27 | 97.6%
(4, 128) | 2.85 | 2.36 | 2.55 | 92.5%
(8, 128) | 5.84 | 4.68 | 5.07 | 92.3%
(16, 128) | 11.9 | 9.55 | 10.2 | 93.6%
----


Albert (PyTorch)

input | TorchScript | OnnxRuntime | DISC-ACL | DISC-Transform
-- | -- | -- | -- | --
(2, 12) | 0.197 | 0.140 | 0.117 | 0.139



Fix patch:
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/BUILD b/tao_compiler/mlir/disc/tools/disc-transform/BUILD
index 7c115765ae2..edd6fa6f124 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/BUILD
+++ b/tao_compiler/mlir/disc/tools/disc-transform/BUILD
@@ -61,6 +61,14 @@ gentbl_cc_library(
             ["--gen-op-defs"],
             "LinalgExt/LinalgExtOps.cc.inc",
         ),
+        (
+            ["-gen-enum-decls"],
+            "LinalgExt/LinalgExtEnums.h.inc",
+        ),
+        (
+            ["-gen-enum-defs"],
+            "LinalgExt/LinalgExtEnums.cc.inc",
+        ),
         (
             ["--gen-typedef-decls"],
             "LinalgExt/LinalgExtTypes.h.inc",
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.cc b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.cc
index e7aaea900b7..e85c187f314 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.cc
@@ -43,6 +43,7 @@
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Support/MathExtras.h"
 #include "mlir/disc/tools/disc-transform/LinalgExt/LinalgExtDialect.h"
+#include "mlir/disc/tools/disc-transform/LinalgExt/LinalgExtEnums.cc.inc"
 
 namespace mlir {
 namespace disc_ral {
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.h b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.h
index 1e68635d35e..01b4f10ffda 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.h
+++ b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.h
@@ -20,6 +20,7 @@
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/TilingInterface.h"
+#include "mlir/disc/tools/disc-transform/LinalgExt/LinalgExtEnums.h.inc"
 #include "mlir/disc/tools/disc-transform/LinalgExt/LinalgExtInterfaces.h"
 
 #define GET_OP_CLASSES
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.td b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.td
index 87fac36ab9c..a6d6b1316b6 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.td
+++ b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.td
@@ -16,6 +16,7 @@ include "external/org_tensorflow/tensorflow/compiler/mlir/disc/tools/disc-transf
 include "external/org_tensorflow/tensorflow/compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtInterfaces.td"
 
 include "mlir/IR/AttrTypeBase.td"
+include "mlir/IR/BuiltinAttributeInterfaces.td"
 include "mlir/IR/EnumAttr.td"
 include "mlir/IR/OpBase.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
@@ -72,6 +73,56 @@ def DISCLinalgExt_ConstantWrapperOp : Op<DISCLinalgExt_Dialect, "constant_wrappe
   let assemblyFormat = "attr-dict $value";
 }
 
+def DISCLinalgExt_Any: I64EnumAttrCase<"kAny", 0>; // padding value doesn't matter.
+def DISCLinalgExt_Exact: I64EnumAttrCase<"kExact", 1>; // padding value matters.
+def DISCLinalgExt_PaddingValueMode : I64EnumAttr<"PaddingValueModeEnum",
+    "Padding mode for padding related-ops",
+    [
+        DISCLinalgExt_Any,
+        DISCLinalgExt_Exact
+    ]> {
+  let summary = "Padding mode for padding related-ops";
+  let cppNamespace = "::mlir::disc_ral::disc_linalg_ext";
+}
+
+def DISCLinalgExt_PaddingValuePlaceholderOp : Op<DISCLinalgExt_Dialect, "padding_value_placeholder", [
+    Pure, AllTypesMatch<["value", "result"]>]>{
+  let summary = "Used as a placeholder for padding value when vectorize or padding linalg ops";
+  let description = [{
+    Wrapper the padding value whne vectorize or pad a linalg op.
+
+    These padding value may have different semantics.
+
+    Example #1, we do not care the exact value of padding value. Padding is just used to get a
+    static shape value which is suitable for hardware vector/tensor register.
+    ```
+    %padding_value = ... // padding value
+    %0 = vector.transfer_read %arg0[%c0, %c0], %padding_value {in_bounds = [false, true]} : tensor<?x16xf32>, vector<6x16xf32>
+    %1 = vector.transfer_read %arg1[...], %padding_value {in_bounds = [true]} : tensor<3072xf32>, vector<16xf32>
+    %2 = vector.broadcast %1 : vector<16xf32> to vector<6x16xf32>
+    %3 = arith.addf %0, %2 : vector<6x16xf32>
+    %4 = arith.maxf %3, %cst_0 : vector<6x16xf32>
+    %5 = vector.transfer_write %4, %arg0[%c0, %c0] {in_bounds = [false, true]} : vector<6x16xf32>, tensor<?x16xf32>
+    ```
+
+    Example #2: Padding value does matter, otherwise the result is wrong.
+    ```
+    %0 = vector.transfer_read %arg0[%c0, %c0], %padding_value {in_bounds = [false, true]} : tensor<?x?xf32>, vector<6x1xf32>
+    %1 = vector.transfer_read %arg1[...], %padding_value {in_bounds = [true, false]} : tensor<?x?xf32>, vector<1x16xf32>
+    %2 = vector.transfer_read %arg2[...], %padding_value {in_bounds = [false, false]} : tensor<?x?xf32>, vector<6x16xf32>
+    %3 = vector.contract %0, %1, %2 : vector<6x16xf32>
+    %4 = vector.transfer_write %3, %arg2[%c0, %c0] {in_bounds = [false, false]} : vector<6x16xf32>, tensor<?x?xf32>
+    ```
+  }];
+
+  let arguments = (ins TypedAttrInterface:$value,
+                   DefaultValuedAttr<DISCLinalgExt_PaddingValueMode,
+                    "::mlir::disc_ral::disc_linalg_ext::PaddingValueModeEnum::kAny">:$mode);
+  let results = (outs AnyType:$result);
+
+  let assemblyFormat = "attr-dict `padding_mode` `(` $mode `)` `,` `value` `(` $value `)`";
+}
+
 def DISCLinalgExt_MultiLevelPackOp : DISCLinalgExt_Op<"multi_level_pack", [
   DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>,
 ]>{
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
index f087cb1359c..670e5d2d1d1 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
@@ -34,6 +34,7 @@
 #include "mlir/Dialect/PDL/IR/PDLTypes.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/IR/Matchers.h"
 #include "mlir/Interfaces/DestinationStyleOpInterface.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Support/MathExtras.h"
@@ -487,6 +488,75 @@ struct FoldMultiLevelPackOfConstantWrapperPattern
   }
 };
 
+/// convert:
+///   %0 = disc_linalg_ext.constant_wrapper dense<0.000000e+00> ...
+///   %1 = tensor.extract %0 ...
+///   use(%1)
+/// to:
+///   %0 = arith.constant 0 : float
+///   use(%0)
+struct FoldTensorExtractOfConstantWrapperPattern
+    : public OpRewritePattern<tensor::ExtractOp> {
+  using OpRewritePattern<tensor::ExtractOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tensor::ExtractOp op,
+                                PatternRewriter& rewriter) const override {
+    auto constOp = op.getTensor().getDefiningOp<ConstantWrapperOp>();
+    if (!constOp) return failure();
+
+    // Collect the constant indices into the tensor.
+    SmallVector<uint64_t, 8> indices;
+    for (Value indice : llvm::drop_begin(op->getOperands(), 1)) {
+      auto constIndexOp = indice.getDefiningOp<arith::ConstantIndexOp>();
+      if (!constIndexOp) return failure();
+      indices.push_back(constIndexOp.getValue().cast<IntegerAttr>().getInt());
+    }
+
+    auto attr = constOp.getValue().getValues<Attribute>()[indices];
+    rewriter.replaceOpWithNewOp<arith::ConstantOp>(op, op.getResult().getType(),
+                                                   attr);
+    return success();
+  }
+};
+
+/// convert:
+///   %0 = ... : vector<16xf32>
+///   %1 = vector.transfer_write %0, %arg0[%c0, %c0]
+///   %2 = disc_linalg_ext. dense<0.000000e+00> ...
+///   %3 = vector.transfer_read %1[%c0, %c0], %2 ...
+///   use(%3)
+/// to:
+///   %0 = ... : vector<16xf32>
+///   use(%0)
+struct FoldXferReadOfXferWriterPattern
+    : public OpRewritePattern<vector::TransferReadOp> {
+  using OpRewritePattern<vector::TransferReadOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::TransferReadOp op,
+                                PatternRewriter& rewriter) const override {
+    auto writeOp = op.getSource().getDefiningOp<vector::TransferWriteOp>();
+    if (!writeOp || writeOp.getVector().getType() != op->getResult(0).getType())
+      return failure();
+
+    // check all the indices are equal
+    if (op.getIndices().size() != writeOp.getIndices().size()) return failure();
+    for (auto [idx0, idx1] : llvm::zip(op.getIndices(), writeOp.getIndices()))
+      if (idx0 != idx1) return failure();
+
+    // check padding value of xfer read op is padding_value_placeholder with
+    // kAny mode padding.
+    auto placeholderOp =
+        op.getPadding()
+            .getDefiningOp<disc_linalg_ext::PaddingValuePlaceholderOp>();
+    if (!placeholderOp ||
+        placeholderOp.getMode() != disc_linalg_ext::PaddingValueModeEnum::kAny)
+      return failure();
+
+    rewriter.replaceOp(op, writeOp.getVector());
+    return success();
+  }
+};
+
 static void addAllRegisteredCanonicalizationPatterns(
     RewritePatternSet& patterns) {
   MLIRContext* ctx = patterns.getContext();
@@ -502,6 +572,8 @@ static void addAllRegisteredCanonicalizationPatterns(
   patterns.insert<TransferReadOfFillOpPattern>(ctx);
   patterns.insert<TransferWriteOfFillOpPattern>(ctx);
   patterns.insert<FoldMultiLevelPackOfConstantWrapperPattern>(ctx);
+  patterns.insert<FoldTensorExtractOfConstantWrapperPattern>(ctx);
+  patterns.insert<FoldXferReadOfXferWriterPattern>(ctx);
 }
 
 }  // namespace
@@ -1642,6 +1714,71 @@ DiagnosedSilenceableFailure LinalgFuseProducersOp::apply(
   return DiagnosedSilenceableFailure(success());
 }
 
+//===---------------------------------------------------------------------===//
+// ReplaceConstPaddingValueOp
+//===---------------------------------------------------------------------===//
+
+void ReplaceConstPaddingValueOp::build(OpBuilder& builder,
+                                       OperationState& result, Value target,
+                                       StringRef mode) {
+  MLIRContext* ctx = builder.getContext();
+  result.addOperands(target);
+  result.addAttribute(ReplaceConstPaddingValueOp::getModeAttrName(result.name),
+                      builder.getStringAttr(mode));
+  result.addTypes({pdl::OperationType::get(ctx)});
+}
+
+DiagnosedSilenceableFailure ReplaceConstPaddingValueOp::applyToOne(
+    Operation* target, SmallVectorImpl<Operation*>& results,
+    transform::TransformState& state) {
+  auto padOp = dyn_cast<tensor::PadOp>(target);
+  if (!padOp) {
+    return mlir::emitDefiniteFailure(target, "applies only to tensor::PadOp");
+  }
+
+  auto mode =
+      disc_linalg_ext::symbolizeEnum<disc_linalg_ext::PaddingValueModeEnum>(
+          getMode());
+  if (!mode) {
+    return mlir::emitDefiniteFailure(target, "invalid mode attr");
+  }
+
+  OpBuilder b(target);
+  auto yieldOp = padOp.getBody()->getTerminator();
+  Attribute value;
+  if (!matchPattern(yieldOp->getOperand(0), m_Constant(&value))) {
+    return mlir::emitDefiniteFailure(target, "applies only to pad with const");
+  }
+  auto placeholder = b.create<disc_linalg_ext::PaddingValuePlaceholderOp>(
+      padOp.getLoc(), value, *mode);
+  yieldOp->replaceUsesOfWith(yieldOp->getOperand(0), placeholder->getResult(0));
+  results.assign({target});
+  return DiagnosedSilenceableFailure(success());
+}
+
+//===---------------------------------------------------------------------===//
+// ConvertPaddingPlaceholderToConstOp
+//===---------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure ConvertPaddingPlaceholderToConstOp::applyToOne(
+    Operation* target, SmallVectorImpl<Operation*>& results,
+    transform::TransformState& state) {
+  auto placeholderOp =
+      dyn_cast<disc_linalg_ext::PaddingValuePlaceholderOp>(target);
+  if (!placeholderOp) {
+    return mlir::emitDefiniteFailure(
+        target, "applies only to PaddingValuePlaceholderOp");
+  }
+
+  OpBuilder b(target);
+
+  auto constOp =
+      b.create<arith::ConstantOp>(target->getLoc(), placeholderOp.getValue());
+  target->replaceAllUsesWith(constOp->getResults());
+  results.assign({constOp.getOperation()});
+  return DiagnosedSilenceableFailure(success());
+}
+
 }  // namespace transform_dialect
 
 void registerTransformDialectCommonExtension(DialectRegistry& registry) {
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
index eef1a6f08a3..8cb9eb07ac9 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
@@ -481,4 +481,58 @@ def LinalgFuseProducersOp : Op<Transform_Dialect, "disc.linalg.fuse_producers",
   let cppNamespace = "::mlir::disc_ral::transform_dialect";
 }
 
+def ReplaceConstPaddingValueOp : Op<Transform_Dialect, "disc.replace_const_padding_value",
+    [FunctionalStyleTransformOpTrait,
+     MemoryEffectsOpInterface,
+     TransformEachOpTrait,
+     TransformOpInterface]> {
+  let description = [{
+    Given a tensor.PadOp, replace its padding content using `padding_value_placeholder`
+
+    Returns the updated tensor.PadOp op.
+  }];
+
+  let arguments = (ins PDL_Operation:$target,
+                       StrAttr:$mode);
+  let results = (outs PDL_Operation:$result);
+
+  let assemblyFormat = "$target attr-dict `mode` `(` $mode `)`";
+  let cppNamespace = "::mlir::disc_ral::transform_dialect";
+
+  let builders = [
+    OpBuilder<(ins "Value":$target, "StringRef":$mode)>
+  ];
+
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::Operation *target,
+        ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+        ::mlir::transform::TransformState &state);
+  }];
+}
+
+def ConvertPaddingPlaceholderToConstOp : Op<Transform_Dialect, "disc.convert_padding_placeholder_to_const",
+    [FunctionalStyleTransformOpTrait,
+     MemoryEffectsOpInterface,
+     TransformEachOpTrait,
+     TransformOpInterface]> {
+  let description = [{
+    Given a `padding_value_placeholder` op, convert to arith.const (ignore its padding mode).
+
+    Returns the result const op.
+  }];
+
+  let arguments = (ins PDL_Operation:$target);
+  let results = (outs PDL_Operation:$result);
+
+  let assemblyFormat = "$target attr-dict";
+  let cppNamespace = "::mlir::disc_ral::transform_dialect";
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::Operation *target,
+        ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+        ::mlir::transform::TransformState &state);
+  }];
+}
+
 #endif // DISC_TRANSFORM_OPS_EXT
diff --git a/tao_compiler/mlir/disc/transforms/disc_transform_legalize_to_loop.cc b/tao_compiler/mlir/disc/transforms/disc_transform_legalize_to_loop.cc
index c700a4966ba..16686285b88 100644
--- a/tao_compiler/mlir/disc/transforms/disc_transform_legalize_to_loop.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_transform_legalize_to_loop.cc
@@ -252,6 +252,7 @@ LogicalResult DiscTransformLegalizeToLoopPass::injectScheduleSelectionIR(
   }
   // Only one candidate, no need to inject selection logic.
   if (factories.size() == 1) {
+    mergeFusionTag(b, fusionOp, factories.back()->getTagSet());
     clonedOps.push_back(fusionOp);
     return success();
   }
diff --git a/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc b/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc
index d7e0d5671f1..e7a44c09b0c 100644
--- a/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc
@@ -172,10 +172,11 @@ transform_dialect::FoldProducerExtractSliceOp buildFoldProducerExtractSlice(
 }
 
 transform::PadOp buildPadOp(OpBuilder& b, Location& loc, Value target,
-                            ArrayRef<int64_t> paddingDimensions) {
+                            ArrayRef<int64_t> paddingDimensions,
+                            int64_t numOperands) {
   auto pdlType = pdl::OperationType::get(b.getContext());
   // TODO(wyzero): support other types.
-  SmallVector<Attribute> paddingAttrs(paddingDimensions.size(),
+  SmallVector<Attribute> paddingAttrs(numOperands,
                                       b.getZeroAttr(b.getF32Type()));
   return b.create<transform::PadOp>(loc, pdlType, target,
                                     b.getArrayAttr(paddingAttrs),
@@ -251,6 +252,28 @@ transform_dialect::DecomposeVectorsOp buildDecomposeVectors(
                                                          vectorSize);
 }
 
+transform_dialect::LinalgFuseProducersOp buildLinalgFuseProducersOp(
+    OpBuilder& b, Location& loc, Value target, ValueRange producers) {
+  auto pdlType = pdl::OperationType::get(b.getContext());
+  return b.create<transform_dialect::LinalgFuseProducersOp>(loc, pdlType,
+                                                            target, producers);
+}
+
+transform_dialect::ReplaceConstPaddingValueOp buildReplaceConstPaddingValueOp(
+    OpBuilder& b, Location& loc, Value target, StringRef mode) {
+  auto pdlType = pdl::OperationType::get(b.getContext());
+  return b.create<transform_dialect::ReplaceConstPaddingValueOp>(loc, pdlType,
+                                                                 target, mode);
+}
+
+transform_dialect::ConvertPaddingPlaceholderToConstOp
+buildConvertPaddingPlaceholderToConstOp(OpBuilder& b, Location& loc,
+                                        Value target) {
+  auto pdlType = pdl::OperationType::get(b.getContext());
+  return b.create<transform_dialect::ConvertPaddingPlaceholderToConstOp>(
+      loc, pdlType, target);
+}
+
 class ParsedFromFileScheduleFactory : public ScheduleFactoryWithNoGuard {
  public:
   explicit ParsedFromFileScheduleFactory(int64_t id, PatternKind kind,
@@ -281,9 +304,23 @@ LogicalResult ParsedFromFileScheduleFactory::assignSchedule(
 class Aarch64GEMMDefaultScheduleFactory : public ScheduleFactoryWithNoGuard {
  public:
   using ScheduleFactoryWithNoGuard::ScheduleFactoryWithNoGuard;
+  bool checkFusionPatternProperties(PatternDescription&) override;
   LogicalResult assignSchedule(PatternDescription&, ModuleOp) override;
 };
 
+// TODO(wyzero): merge default schedule and default with epilogue schedule.
+bool Aarch64GEMMDefaultScheduleFactory::checkFusionPatternProperties(
+    PatternDescription& pd) {
+  auto& fusionPattern = pd.getFusionPattern();
+  auto& rootOps = fusionPattern.getRootOps();
+  // Only support single output a.t.m.
+  if (rootOps.size() != 1) return false;
+
+  // This schedule not support epilogue fusion
+  auto dominantOp = fusionPattern.getDominantOp();
+  return rootOps[0] == dominantOp && isa<lmhlo::DotGeneralOp>(dominantOp);
+}
+
 LogicalResult Aarch64GEMMDefaultScheduleFactory::assignSchedule(
     PatternDescription& pd, ModuleOp m) {
   OpBuilder b(m);
@@ -371,7 +408,7 @@ LogicalResult Aarch64GEMMDefaultScheduleFactory::assignSchedule(
   buildFoldProducerExtractSlice(b, loc, weightInnerSlice, 2);
 
   // pad to match the requirement of hardware vector/tensor instruction.
-  auto padOp = buildPadOp(b, loc, tileOp->getResult(0), {0, 1, 2});
+  auto padOp = buildPadOp(b, loc, tileOp->getResult(0), {0, 1, 2}, 3);
 
   Value padForInput = buildGetProducerOfOperand(b, loc, padOp, 0);
   Value padForWeight = buildGetProducerOfOperand(b, loc, padOp, 1);
@@ -475,14 +512,271 @@ LogicalResult Aarch64GEMMDefaultScheduleFactory::assignSchedule(
   return success();
 }
 
+class Aarch64GEMMDefaultScheduleWithEpilogueFactory
+    : public ScheduleFactoryWithNoGuard {
+ public:
+  using ScheduleFactoryWithNoGuard::ScheduleFactoryWithNoGuard;
+  bool checkFusionPatternProperties(PatternDescription&) override;
+  LogicalResult assignSchedule(PatternDescription&, ModuleOp) override;
+};
+
+bool Aarch64GEMMDefaultScheduleWithEpilogueFactory::
+    checkFusionPatternProperties(PatternDescription& pd) {
+  auto& fusionPattern = pd.getFusionPattern();
+  auto& rootOps = fusionPattern.getRootOps();
+  // Only support single output a.t.m.
+  if (rootOps.size() != 1) return false;
+
+  // This schedule only works for gemm with epilogue fusion
+  auto dominantOp = fusionPattern.getDominantOp();
+  return rootOps[0] != dominantOp && isa<lmhlo::DotGeneralOp>(dominantOp);
+}
+
+LogicalResult Aarch64GEMMDefaultScheduleWithEpilogueFactory::assignSchedule(
+    PatternDescription& pd, ModuleOp m) {
+  OpBuilder b(m);
+  b.setInsertionPointToStart(&m.getBodyRegion().front());
+  Location loc = m.getLoc();
+  MLIRContext* ctx = m->getContext();
+  auto seqOp = b.create<transform_ext::CanonicalizedSequenceOp>(
+      loc, TypeRange{}, transform::FailurePropagationMode::Propagate, Value{});
+  seqOp.getBody().push_back(new Block);
+  auto& bodyBlock = seqOp.getBody().front();
+  auto pdlOpType = pdl::OperationType::get(ctx);
+  bodyBlock.addArgument(pdl::OperationType::get(ctx), loc);
+  b.setInsertionPointToStart(&bodyBlock);
+  Value variant = bodyBlock.getArgument(0);
+
+  auto& fusionPattern = pd.getFusionPattern();
+  auto nameMap = TransformNameAssigner(fusionPattern.getOpList()).getNameMap();
+  auto dotOp =
+      dyn_cast_or_null<lmhlo::DotGeneralOp>(fusionPattern.getDominantOp());
+  if (!dotOp) {
+    return m->emitError() << "expect dot_general op as dominant\n";
+  }
+  Value lhs = dotOp->getOperand(0);
+  Value rhs = dotOp->getOperand(1);
+  auto lhsTy = lhs.getType().cast<MemRefType>();
+  auto rhsTy = rhs.getType().cast<MemRefType>();
+  if (lhsTy.getRank() != 2 || rhsTy.getRank() != 2) {
+    return m->emitError() << "only support rank 2 GEMM a.t.m.\n";
+  }
+
+  auto dimNumbers = dotOp.getDotDimensionNumbers();
+  auto lhsCntractingDims = dimNumbers.getLhsContractingDimensions();
+  auto rhsCntractingDims = dimNumbers.getRhsContractingDimensions();
+  if (lhsCntractingDims.size() != 1 || rhsCntractingDims.size() != 1) {
+    return m->emitError() << "only support exactly 1 contract dim\n";
+  }
+  bool lhsTranspose = (lhsCntractingDims[0] == lhsTy.getRank() - 2);
+  bool rhsTranspose = (rhsCntractingDims[0] == rhsTy.getRank() - 1);
+  int64_t M = lhsTranspose ? lhsTy.getShape()[lhsTy.getRank() - 1]
+                           : lhsTy.getShape()[lhsTy.getRank() - 2];
+  int64_t K = lhsTranspose ? lhsTy.getShape()[lhsTy.getRank() - 2]
+                           : lhsTy.getShape()[lhsTy.getRank() - 1];
+  int64_t N = rhsTranspose ? rhsTy.getShape()[rhsTy.getRank() - 2]
+                           : rhsTy.getShape()[rhsTy.getRank() - 1];
+
+  if (fusionPattern.getRootOps().size() != 1) {
+    return m->emitError() << "only support single output a.t.m.\n";
+  }
+  Operation* rootOp = fusionPattern.getRootOps()[0];
+  Value rootHandle = buildMatchOp(b, loc, variant, {}, nameMap[rootOp]);
+
+  // merge elemwise ops in case there are many.
+  SmallVector<Value> otherElemOpHandles;
+  for (Operation* op : fusionPattern.getOpList()) {
+    if (op == rootOp || op == dotOp.getOperation()) continue;
+    otherElemOpHandles.push_back(
+        buildMatchOp(b, loc, variant, {}, nameMap[op]));
+  }
+  if (!otherElemOpHandles.empty()) {
+    buildLinalgFuseProducersOp(b, loc, rootHandle, otherElemOpHandles);
+    rootHandle = buildMatchOp(b, loc, variant, {}, nameMap[rootOp]);
+  }
+
+  // split root ops.
+  auto forEachThreadOp = buildTileToForEachThreadOp(b, loc, rootHandle, {1, 1});
+  Value forEachThreadLoop = forEachThreadOp->getResult(0);
+  Value tiledRoot = forEachThreadOp->getResult(1);
+
+  // build handle to target dot op.
+  Value fillAndMatmul = buildMatchOp(b, loc, variant, {}, nameMap[dotOp]);
+  auto matmulSplitOp = buildSplitHandlesOp(b, loc, fillAndMatmul, 2);
+  Value fill = matmulSplitOp->getResult(0);
+  Value matmul = matmulSplitOp->getResult(1);
+  matmul = buildFuseIntoContainingOp(b, loc, matmul, forEachThreadLoop);
+  fill = buildFuseIntoContainingOp(b, loc, fill, forEachThreadLoop);
+
+  // first/second level tile size for dimension m
+  int64_t M0 = 288, M1 = 6;
+  // first/second level tile size for dimension n
+  int64_t N0 = 48, N1 = 16;
+  // first level tile size for dimension k
+  int64_t K0 = 1;
+  // TODO(wyzero): query cpuinfo.
+  int64_t hardwareVectorSizeInBytes = 4;
+
+  // first level tile and fuse matmul and fill op.
+  auto tileOp1 = buildTileOp(b, loc, tiledRoot, {M0, N0}, {0, 1});
+  matmul = buildFuseIntoContainingOp(b, loc, matmul, tileOp1->getResult(1));
+  fill = buildFuseIntoContainingOp(b, loc, fill, tileOp1->getResult(1));
+
+  // second level tile and fuse matmul and fill op.
+  auto tileOp2 = buildTileOp(b, loc, tileOp1->getResult(0), {M1, N1}, {0, 1});
+  // pad root ops to register level size (static size).
+  int64_t numOperandsUpperBound = fusionPattern.getOperands().size() +
+                                  fusionPattern.getResults().size() +
+                                  fusionPattern.getInternalResults().size() + 1;
+  auto padRootOp =
+      buildPadOp(b, loc, tileOp2->getResult(0), {0, 1}, numOperandsUpperBound);
+  auto padOps = buildMatchOp(b, loc, variant, {"tensor.pad"}, {});
+  buildReplaceConstPaddingValueOp(b, loc, padOps, "kAny");
+  matmul = buildFuseIntoContainingOp(b, loc, matmul, tileOp2->getResult(1));
+  fill = buildFuseIntoContainingOp(b, loc, fill, tileOp2->getResult(1));
+
+  // gemm reduction axis tiling
+  auto tileOp3 = buildTileOp(b, loc, matmul, {0, 0, K0}, {0, 1, 2});
+
+  variant = buildRunCanonicalizer(b, loc, variant);
+
+  // fold two extract_slice ops generated by two-level tiling. It's needed to
+  // enable following pad and hosit schedule.
+  Value weightInnerSlice =
+      buildGetProducerOfOperand(b, loc, tileOp3->getResult(0), 1);
+  buildFoldProducerExtractSlice(b, loc, weightInnerSlice, 2);
+
+  // pad to match the requirement of hardware vector/tensor instruction.
+  auto padGEMMOp = buildPadOp(b, loc, tileOp3->getResult(0), {0, 1, 2}, 3);
+
+  Value padForInput = buildGetProducerOfOperand(b, loc, padGEMMOp, 0);
+  Value padForWeight = buildGetProducerOfOperand(b, loc, padGEMMOp, 1);
+
+  // Check if we need to pad dimension `m/n/k` if input or weight is packed
+  bool mIsPadded = (M1 != 1) && (M == ShapedType::kDynamicSize ||
+                                 (M > M0 && (M % M0 != 0 || M0 % M1 != 0)) ||
+                                 (M <= M0 && M > M1 && M % M1 != 0));
+  bool nIsPadded = (N1 != 1) && (N == ShapedType::kDynamicSize ||
+                                 (N > N0 && (N % N0 != 0 || N0 % N1 != 0)) ||
+                                 (N <= N0 && N > N1 && N % N1 != 0));
+  bool kIsPadded =
+      (K0 != 1) && (K == ShapedType::kDynamicSize || K > K0 && K % K0 != 0);
+
+  // Check if we need to pack the input:
+  bool packInput = ((M == ShapedType::kDynamicSize || M >= M1) &&
+                    (K == ShapedType::kDynamicSize || K > K0) &&
+                    (N == ShapedType::kDynamicSize || N > N0));
+  if (packInput) {
+    // supposed loop order:
+    //  loop_m0
+    //   loop_n0
+    //    loop_m1
+    //     loop_n1
+    //      loop_k0 {
+    //        inner_most_gemm
+    //      }
+    // We want to cache the packed A below loop_m0 and above loop_n0.
+    // Thus the initial loop_level is 4.
+    int loopLevel = 4;
+    // in case:
+    // - the size of dimension N <= N0, then loop_n0 will be folded.
+    loopLevel -= (N != ShapedType::kDynamicSize && N <= N0);
+    // - the size of dimension M <= M1, then loop_m1 will be folded.
+    loopLevel -= (M != ShapedType::kDynamicSize && M <= M1);
+    // - the size of dimension N <= N1, then loop_n1 will be folded.
+    loopLevel -= (N != ShapedType::kDynamicSize && N <= N1);
+    // - the size of dimension K <= K0, then loop_k0 will be folded.
+    loopLevel -= (K != ShapedType::kDynamicSize && K <= K0);
+
+    if (loopLevel <= 0) {
+      return m->emitError()
+             << "failed to cache the packed input due to loopLevel = "
+             << loopLevel << " is invalid\n";
+    }
+    auto loopN0 = buildGetParentForOp(b, loc, padForInput, loopLevel);
+    bool inputIsPadded = mIsPadded || kIsPadded;
+    SmallVector<int64_t> tileSizes;
+    SmallVector<int64_t> permutation;
+    if (lhsTranspose) {
+      tileSizes = {K0, M1};
+      permutation = {2, 0, 1, 3};
+    } else {
+      tileSizes = {M1, K0};
+      permutation = {0, 2, 3, 1};
+    }
+    buildCacheRead(b, loc, padForInput, loopN0, {1, 1}, tileSizes,
+                   inputIsPadded, permutation);
+  }
+
+  // Check if we need to pack the weight, one of the following conditions:
+  // - if M, N and K are both dynamic, we always pad input a.t.m.
+  // - if N is known and N >= N0 && N0 > N1
+  bool packWeight = ((K == ShapedType::kDynamicSize || K > K0) &&
+                     (N == ShapedType::kDynamicSize || N > N1));
+  if (packWeight) {
+    bool weightIsPadded = nIsPadded || kIsPadded;
+    forEachThreadLoop = buildMatchOp(b, loc, variant, {"scf.foreach_thread"});
+    SmallVector<int64_t> tileSizes;
+    SmallVector<int64_t> permutation;
+    if (rhsTranspose) {
+      tileSizes = {N1, K0};
+      permutation = {0, 2, 3, 1};
+    } else {
+      tileSizes = {K0, N1};
+      permutation = {2, 0, 1, 3};
+    }
+    buildCacheRead(b, loc, padForWeight, forEachThreadLoop, {1, 1}, tileSizes,
+                   weightIsPadded, permutation);
+  }
+
+  variant = buildRunCanonicalizer(b, loc, variant);
+
+  Value multiLevelPackOps =
+      buildMatchOp(b, loc, variant, {"disc_linalg_ext.multi_level_pack"});
+  buildLowerMultiLevelPackToLoop(b, loc, multiLevelPackOps);
+
+  variant = buildRunCanonicalizer(b, loc, variant);
+
+  Value func = buildMatchOp(b, loc, variant, {"func.func"});
+  buildVectorize(b, loc, func, true);
+
+  variant = buildRunCanonicalizer(b, loc, variant);
+  auto placeholderOps = buildMatchOp(
+      b, loc, variant, {"disc_linalg_ext.padding_value_placeholder"}, {});
+  buildConvertPaddingPlaceholderToConstOp(b, loc, placeholderOps);
+  variant = buildDISCBufferize(b, loc, variant);
+
+  buildLowerVectors(b, loc, {0, 1, 2, 3, 4}, "outerproduct", "innerparallel",
+                    "linalg-copy", true, "eltwise", false);
+  buildLowerVectors(b, loc, {5, 6, 7}, "outerproduct", "innerparallel",
+                    "linalg-copy", true, "eltwise", false);
+  // de-compose large size vector operations
+  variant = buildDecomposeVectors(b, loc, variant, hardwareVectorSizeInBytes);
+  b.create<transform::YieldOp>(loc);
+  return success();
+}
+
 class Aarch64GEMMLargeKScheduleFactory : public ScheduleFactory {
  public:
   using ScheduleFactory::ScheduleFactory;
+  bool checkFusionPatternProperties(PatternDescription&) override;
   LogicalResult assignSchedule(PatternDescription&, ModuleOp) override;
   LogicalResult buildGuardCondition(OpBuilder& b, Location loc,
                                     PatternDescription&, Value&) override;
 };
 
+bool Aarch64GEMMLargeKScheduleFactory::checkFusionPatternProperties(
+    PatternDescription& pd) {
+  auto& fusionPattern = pd.getFusionPattern();
+  auto& rootOps = fusionPattern.getRootOps();
+  // Only support single output a.t.m.
+  if (rootOps.size() != 1) return false;
+
+  // This schedule not support epilogue fusion
+  auto dominantOp = fusionPattern.getDominantOp();
+  return rootOps[0] == dominantOp && isa<lmhlo::DotGeneralOp>(dominantOp);
+}
+
 LogicalResult Aarch64GEMMLargeKScheduleFactory::buildGuardCondition(
     OpBuilder& b, Location loc, PatternDescription& pd, Value& pred) {
   auto& fusionPattern = pd.getFusionPattern();
@@ -590,7 +884,7 @@ LogicalResult Aarch64GEMMLargeKScheduleFactory::assignSchedule(
   buildFoldProducerExtractSlice(b, loc, weightInnerSlice, 2);
 
   // pad to match the requirement of hardware vector/tensor instruction.
-  auto padOp = buildPadOp(b, loc, tileOp1->getResult(0), {0, 1, 2});
+  auto padOp = buildPadOp(b, loc, tileOp1->getResult(0), {0, 1, 2}, 3);
 
   Value padForInput = buildGetProducerOfOperand(b, loc, padOp, 0);
   Value padForWeight = buildGetProducerOfOperand(b, loc, padOp, 1);
@@ -717,6 +1011,10 @@ DISC_TRANSFORM_SCHEDULE(PatternKind::kGEMM, kDefaultScheduleFactoryPriority,
                         Aarch64GEMMDefaultScheduleFactory,
                         ArrayRef<StringRef>{kDefaultScheduleFactoryTag});
 
+DISC_TRANSFORM_SCHEDULE(PatternKind::kGEMM, 10,
+                        Aarch64GEMMDefaultScheduleWithEpilogueFactory,
+                        ArrayRef<StringRef>{"default_epilogue"});
+
 DISC_TRANSFORM_SCHEDULE(PatternKind::kGEMM, 100,
                         Aarch64GEMMLargeKScheduleFactory,
                         ArrayRef<StringRef>{"large_k"});
diff --git a/tao_compiler/mlir/disc/transforms/disc_transform_schedule.h b/tao_compiler/mlir/disc/transforms/disc_transform_schedule.h
index 5ae9541557f..4431b450124 100644
--- a/tao_compiler/mlir/disc/transforms/disc_transform_schedule.h
+++ b/tao_compiler/mlir/disc/transforms/disc_transform_schedule.h
@@ -86,6 +86,7 @@ class ScheduleFactory {
  public:
   explicit ScheduleFactory(int64_t id, PatternKind kind,
                            ArrayRef<StringRef> tags);
+  virtual ~ScheduleFactory() = default;
 
   // Returns true if the factory accepts the pattern at compile time.
   // Note that it's only an conservative check, which means pattern rejected
diff --git a/tao_compiler/mlir/disc/transforms/fusion_utils_transform_based.cc b/tao_compiler/mlir/disc/transforms/fusion_utils_transform_based.cc
index d279463d3cc..3457bdd3509 100644
--- a/tao_compiler/mlir/disc/transforms/fusion_utils_transform_based.cc
+++ b/tao_compiler/mlir/disc/transforms/fusion_utils_transform_based.cc
@@ -98,6 +98,9 @@ bool TransformBasedCpuFusionStrategy::initFusionPattern(
     return true;
   }
 
+  // We only support single output right now
+  if (fusionPattern.getResults().size() != 1) return true;
+
   DenseSet<Value> dotWeights;
   DenseSet<Operation*> supportedDotOps;
   for (Operation* op : fusionPattern.getOpList()) {
@@ -111,8 +114,21 @@ bool TransformBasedCpuFusionStrategy::initFusionPattern(
     }
   }
 
-  // Only support one gemm a.t.m.
-  if (supportedDotOps.size() != 1) return true;
+  // Only support at most one gemm a.t.m.
+  if (supportedDotOps.size() > 1) return true;
+  if (supportedDotOps.empty()) {
+    // special case: for elem+bcast epilogue subgraph fusion
+    // 1, check no large const
+    if (llvm::all_of(fusionPattern.getOpList(), [&](Operation* op) {
+          return !isa<lmhlo::ConstantOp>(op) || isScalarConstOp(op);
+        })) {
+      fusionPattern.setFusionType(FusionType::kLoop);
+      fusionPattern.setDominantOp(*fusionPattern.getRootOps().begin());
+    }
+    return true;
+  }
+
+  // normal case: gemm + epilogue fusion
   Operation* dominantDotOp = *supportedDotOps.begin();
 
   // Only support fuse const ops that are not consumed by ops outside the fusion
@@ -141,8 +157,6 @@ bool TransformBasedCpuFusionStrategy::initFusionPattern(
     if (llvm::find(operands, operand) == operands.end()) return true;
   }
 
-  // We only support single output right now
-  if (fusionPattern.getResults().size() != 1) return true;
   // the shape of the output should be the same as the shape of result of
   // dominant op.
   for (Value result : fusionPattern.getResults()) {
@@ -160,7 +174,7 @@ bool TransformBasedCpuFusionStrategy::tryFuse(ShapeAnalysis& shapeAnalysis,
                                               FusionPattern& rhs,
                                               FusionPattern& target) {
   if (!initFusionPattern(shapeAnalysis, target)) return false;
-  return target.isTransformBasedFusion();
+  return target.isFusible();
 }
 
 }  // namespace disc_ral


Test patch:
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_epilogue_biasadd_p_768x3072_f32.mlir b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_epilogue_biasadd_p_768x3072_f32.mlir
new file mode 100644
index 00000000000..c4dacbad07b
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_epilogue_biasadd_p_768x3072_f32.mlir
@@ -0,0 +1,12 @@
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 0 : i32}} {
+  func.func @main(%arg0: tensor<?x?xf32>) -> (tensor<?x?xf32>) attributes {tf.entry_function = {inputs = "{{INPUTS}}", outputs = "{{OUTPUTS}}", input_placements="{{INPUT_PLACEMENTS}}", output_placements="{{OUTPUT_PLACEMENTS}}"}} {
+    %graph = tf_executor.graph {
+      %0:2 = tf_executor.island wraps "tf.Const"() {value = dense<-0.8> : tensor<768x3072xf32>} : () -> tensor<768x3072xf32>
+      %1:2 = tf_executor.island wraps "tf.MatMul"(%arg0, %0) {transpose_a = false, transpose_b = false} : (tensor<?x?xf32>, tensor<768x3072xf32>) -> (tensor<?x?xf32>)
+      %2:2 = tf_executor.island wraps "tf.Const"() {value = dense<-0.1> : tensor<3072xf32>} : () -> tensor<3072xf32>
+      %3:2 = tf_executor.island wraps "tf.BiasAdd"(%1, %2) : (tensor<?x?xf32>, tensor<3072xf32>) -> tensor<?x?xf32>
+      tf_executor.fetch %3 : tensor<?x?xf32>
+    }
+    return %graph : tensor<?x?xf32>
+  }
+}
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_epilogue_biasadd_relu_p_768_3072_f32.mlir b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_epilogue_biasadd_relu_p_768_3072_f32.mlir
new file mode 100644
index 00000000000..d2cf2045107
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_epilogue_biasadd_relu_p_768_3072_f32.mlir
@@ -0,0 +1,13 @@
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 0 : i32}} {
+  func.func @main(%arg0: tensor<?x?xf32>) -> (tensor<?x?xf32>) attributes {tf.entry_function = {inputs = "{{INPUTS}}", outputs = "{{OUTPUTS}}", input_placements="{{INPUT_PLACEMENTS}}", output_placements="{{OUTPUT_PLACEMENTS}}"}} {
+    %graph = tf_executor.graph {
+      %0:2 = tf_executor.island wraps "tf.Const"() {value = dense<-0.8> : tensor<768x3072xf32>} : () -> tensor<768x3072xf32>
+      %1:2 = tf_executor.island wraps "tf.MatMul"(%arg0, %0) {transpose_a = false, transpose_b = false} : (tensor<?x?xf32>, tensor<768x3072xf32>) -> (tensor<?x?xf32>)
+      %2:2 = tf_executor.island wraps "tf.Const"() {value = dense<-0.1> : tensor<3072xf32>} : () -> tensor<3072xf32>
+      %3:2 = tf_executor.island wraps "tf.BiasAdd"(%1, %2) : (tensor<?x?xf32>, tensor<3072xf32>) -> tensor<?x?xf32>
+      %4:2 = tf_executor.island wraps "tf.Relu"(%3) : (tensor<?x?xf32>) -> tensor<?x?xf32>
+      tf_executor.fetch %4 : tensor<?x?xf32>
+    }
+    return %graph : tensor<?x?xf32>
+  }
+}
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_epilogue_complex_test0.mlir b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_epilogue_complex_test0.mlir
new file mode 100644
index 00000000000..6d58e93d7d2
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_epilogue_complex_test0.mlir
@@ -0,0 +1,25 @@
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 0 : i32}} {
+  func.func @main(%arg0: tensor<?x?xf32>) -> (tensor<?x?xf32>) attributes {tf.entry_function = {inputs = "{{INPUTS}}", outputs = "{{OUTPUTS}}", input_placements="{{INPUT_PLACEMENTS}}", output_placements="{{OUTPUT_PLACEMENTS}}"}} {
+    %graph = tf_executor.graph {
+      %0:2 = tf_executor.island wraps "tf.Const"() {value = dense<-0.8> : tensor<768x3072xf32>} : () -> tensor<768x3072xf32>
+      %1:2 = tf_executor.island wraps "tf.MatMul"(%arg0, %0) {transpose_a = false, transpose_b = false} : (tensor<?x?xf32>, tensor<768x3072xf32>) -> (tensor<?x?xf32>)
+      %2:2 = tf_executor.island wraps "tf.Const"() {value = dense<-0.1> : tensor<3072xf32>} : () -> tensor<3072xf32>
+      %3:2 = tf_executor.island wraps "tf.BiasAdd"(%1, %2) : (tensor<?x?xf32>, tensor<3072xf32>) -> tensor<?x?xf32>
+      %4:2 = tf_executor.island wraps "tf.Square"(%3) : (tensor<?x?xf32>) -> tensor<?x?xf32>
+      %5:2 = tf_executor.island wraps "tf.Const"() {value = dense<4.471500e-02> : tensor<f32>} : () -> tensor<f32>
+      %6:2 = tf_executor.island wraps "tf.Mul"(%4, %5) : (tensor<?x?xf32>, tensor<f32>) -> tensor<?x?xf32>
+      %7:2 = tf_executor.island wraps "tf.Mul"(%6, %3) : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
+      %8:2 = tf_executor.island wraps "tf.AddV2"(%3, %7) : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
+      %9:2 = tf_executor.island wraps "tf.Const"() {value = dense<0.797884583> : tensor<f32>} : () -> tensor<f32>
+      %10:2 = tf_executor.island wraps "tf.Mul"(%8, %9) : (tensor<?x?xf32>, tensor<f32>) -> tensor<?x?xf32>
+      %11:2 = tf_executor.island wraps "tf.Tanh"(%10) : (tensor<?x?xf32>) -> tensor<?x?xf32>
+      %12:2 = tf_executor.island wraps "tf.Const"() {value = dense<1.797884583> : tensor<f32>} : () -> tensor<f32>
+      %13:2 = tf_executor.island wraps "tf.AddV2"(%11, %12) : (tensor<?x?xf32>, tensor<f32>) -> tensor<?x?xf32>
+      %14:2 = tf_executor.island wraps "tf.Const"() {value = dense<0.5> : tensor<f32>} : () -> tensor<f32>
+      %15:2 = tf_executor.island wraps "tf.Mul"(%3, %14) : (tensor<?x?xf32>, tensor<f32>) -> tensor<?x?xf32>
+      %16:2 = tf_executor.island wraps "tf.Mul"(%13, %15) : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
+      tf_executor.fetch %16 : tensor<?x?xf32>
+    }
+    return %graph : tensor<?x?xf32>
+  }
+}
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/matmul_epilogue.cc b/tao_compiler/mlir/disc/tests/disc-transform/matmul_epilogue.cc
new file mode 100644
index 00000000000..32a6d58d332
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/matmul_epilogue.cc
@@ -0,0 +1,83 @@
+/* Copyright 2022 The BladeDISC Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "mlir/disc/tests/mlir_feature_test.h"
+#include "mlir/disc/tests/mlir_test.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace mlir_test {
+
+const std::string c_ft_path = "mlir/disc/tests/disc-transform/data/";
+
+static bool init_threads = []() {
+  setenv("OMP_NUM_THREADS", "1", 1);
+  setenv("DISC_CPU_ENABLE_WEIGHT_PRE_PACKING", "1", 1);
+  return true;
+}();
+
+TEST(PackedMatmul, F32_biasadd_768x3072) {
+  EnvSetting setting = {{"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
+                        {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", false}},
+                        {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", false}}};
+  EnvSettingContext ctx(setting);
+  EXPECT_TRUE(feature_test_main(
+      /*mlir_file_path*/ c_ft_path +
+          "matmul_epilogue_biasadd_p_768x3072_f32.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 1,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"512x768xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
+TEST(PackedMatmul, F32_biasadd_relu_768x3072) {
+  EnvSetting setting = {{"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
+                        {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", false}},
+                        {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", false}}};
+  EnvSettingContext ctx(setting);
+  EXPECT_TRUE(feature_test_main(
+      /*mlir_file_path*/ c_ft_path +
+          "matmul_epilogue_biasadd_relu_p_768_3072_f32.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 1,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"512x768xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
+TEST(PackedMatmul, F32_complex_test0) {
+  EnvSetting setting = {{"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
+                        {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", false}},
+                        {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", false}}};
+  EnvSettingContext ctx(setting);
+  EXPECT_TRUE(feature_test_main(
+      /*mlir_file_path*/ c_ft_path + "matmul_epilogue_complex_test0.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 1,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"512x768xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
+}  // namespace mlir_test
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/tests/ops.mlir b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/tests/ops.mlir
index fdb8557492d..cf342ebf7a0 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/tests/ops.mlir
+++ b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/tests/ops.mlir
@@ -14,3 +14,13 @@ func.func @const() -> tensor<2x4xf32> {
   %0 = disc_linalg_ext.constant_wrapper dense<[[0.0, 1.0, 2.0, 3.0], [4.0, 5.0, 6.0, 7.0]]> : tensor<2x4xf32>
   return %0 : tensor<2x4xf32>
 }
+
+// -----
+
+// CHECK-LABEL: @padding_value_placeholder
+func.func @padding_value_placeholder() -> f32 {
+  %0 = disc_linalg_ext.padding_value_placeholder padding_mode(kAny), value(0.0 : f32)
+  %1 = disc_linalg_ext.padding_value_placeholder padding_mode(kExact), value(0.0 : f32)
+  %2 = arith.addf %0, %1 : f32
+  return %2 : f32
+}
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/convert-padding-placeholder-to-const.mlir b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/convert-padding-placeholder-to-const.mlir
new file mode 100644
index 00000000000..e08fa10d6d8
--- /dev/null
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/convert-padding-placeholder-to-const.mlir
@@ -0,0 +1,16 @@
+// RUN: disc-opt --disc-transform-dialect-interpreter -split-input-file %s | FileCheck %s
+
+// CHECK-LABEL: @convert_padding_placeholder_to_const
+func.func @convert_padding_placeholder_to_const() -> f32 {
+  // CHECK-NOT: disc_linalg_ext.padding_value_placeholder
+  // CHECK: %[[T0:.*]] = arith.constant 0.000000e+00 : f32
+  // CHECK: return %[[T0]] : f32
+  %0 = disc_linalg_ext.padding_value_placeholder padding_mode(kAny), value(0.000000e+00 : f32)
+  return %0 : f32
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation):
+  %placeholder = transform.structured.match ops{["disc_linalg_ext.padding_value_placeholder"]} in %arg0
+  transform.disc.convert_padding_placeholder_to_const %placeholder
+}
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/replace-const-padding-value.mlir b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/replace-const-padding-value.mlir
new file mode 100644
index 00000000000..f1a42986bde
--- /dev/null
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/replace-const-padding-value.mlir
@@ -0,0 +1,22 @@
+// RUN: disc-opt --disc-transform-dialect-interpreter -split-input-file %s | FileCheck %s
+
+// CHECK-LABEL: @replace_const_padding_value
+func.func @replace_const_padding_value(%arg0: tensor<?xf32>) -> tensor<?xf32> {
+  %c1 = arith.constant 1 : index
+  %cst = arith.constant 0.000000e+00 : f32
+  // CHECK: %[[T0:.*]] = disc_linalg_ext.padding_value_placeholder padding_mode(kAny), value(0.000000e+00 : f32)
+  // CHECK-NEXT: tensor.pad
+  // CHECK-NEXT: ^bb0
+  // CHECK-NEXT:   tensor.yield %[[T0]] : f32
+  %out = tensor.pad %arg0 low[%c1] high[%c1] {
+  ^bb0(%arg1: index):
+    tensor.yield %cst : f32
+  } : tensor<?xf32> to tensor<?xf32>
+  return %out : tensor<?xf32>
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation):
+  %pad = transform.structured.match ops{["tensor.pad"]} in %arg0
+  transform.disc.replace_const_padding_value %pad mode("kAny")
+}
\ No newline at end of file