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/TransformOps/TransformOpsExt.cc b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
index 7b6557340ea..2d96b944e15 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
@@ -2335,12 +2335,6 @@ DiagnosedSilenceableFailure ReductionOutputFuseOp::apply(
   auto yieldOp = cast<scf::YieldOp>(newForOp.getBody()->getTerminator());
   auto insertSliceOp =
       yieldOp->getOperand(forResultIdx).getDefiningOp<tensor::InsertSliceOp>();
-  if (!insertSliceOp) {
-    return mlir::emitDefiniteFailure(
-        this->getOperation(),
-        "insert_slice op not found for target result of for op\n");
-  }
-
   b.setInsertionPoint(yieldOp);
   for (auto [idx, val] : llvm::enumerate(yieldOp->getOperands())) {
     auto it = forResultToLinalgInitOperandMap.find(idx);
@@ -2348,22 +2342,24 @@ DiagnosedSilenceableFailure ReductionOutputFuseOp::apply(
       mapping.map(forOp->getResult(idx), val);
       continue;
     }
-    auto insertOp = val.getDefiningOp<tensor::InsertSliceOp>();
-    if (!insertOp)
-      return mlir::emitDefiniteFailure(
-          this->getOperation(),
-          "insert_slice op not found for target result of for op\n");
-    mapping.map(forOp->getResult(idx), insertOp->getOperand(0));
+    if (auto insertOp = val.getDefiningOp<tensor::InsertSliceOp>()) {
+      mapping.map(forOp->getResult(idx), insertOp->getOperand(0));
+    } else {
+      mapping.map(forOp->getResult(idx), val);
+    }
   }
   auto clonedTarget = b.clone(*target, mapping);
-  auto tileableTarget = cast<TilingInterface>(clonedTarget);
-  FailureOr<Value> tiledTarget = tileableTarget.generateResultTileValue(
-      b, 0, insertSliceOp.getMixedOffsets(), insertSliceOp.getMixedSizes());
-  if (failed(tiledTarget)) {
-    return mlir::emitDefiniteFailure(this->getOperation(),
-                                     "failed to tile target op\n");
+  Operation* tiledOp = clonedTarget;
+  if (insertSliceOp) {
+    auto tileableTarget = cast<TilingInterface>(clonedTarget);
+    FailureOr<Value> tiledTarget = tileableTarget.generateResultTileValue(
+        b, 0, insertSliceOp.getMixedOffsets(), insertSliceOp.getMixedSizes());
+    if (failed(tiledTarget)) {
+      return mlir::emitDefiniteFailure(this->getOperation(),
+                                       "failed to tile target op\n");
+    }
+    tiledOp = tiledTarget->getDefiningOp();
   }
-  Operation* tiledOp = tiledTarget->getDefiningOp();
   Value iv = newForOp.getInductionVar();
   Value nextIv =
       b.create<arith::AddIOp>(tiledOp->getLoc(), iv, newForOp.getStep());
@@ -2383,8 +2379,12 @@ DiagnosedSilenceableFailure ReductionOutputFuseOp::apply(
       newResults.push_back(val);
       continue;
     }
-    Value initValue = newForOp.getRegionIterArg(idx);
     Value updateValue = conditionalGenericOp->getResult(it->second);
+    if (!insertSliceOp) {
+      newResults.push_back(updateValue);
+      continue;
+    }
+    Value initValue = newForOp.getRegionIterArg(idx);
     newResults.push_back(b.create<tensor::InsertSliceOp>(
                               tiledOp->getLoc(), updateValue, initValue,
                               insertSliceOp.getMixedOffsets(),
@@ -2397,6 +2397,10 @@ DiagnosedSilenceableFailure ReductionOutputFuseOp::apply(
        llvm::enumerate(conditionalGenericOp->getResults())) {
     auto it = linalgInitOperandToForResultMap.find(idx);
     if (it != linalgInitOperandToForResultMap.end()) continue;
+    if (!insertSliceOp) {
+      newResults.push_back(updateValue);
+      continue;
+    }
     Value initValue = newForOp.getRegionIterArgs()
                           .drop_front(forOp.getRegionIterArgs().size() +
                                       nonFusedInitOperandIdx++)
@@ -2486,29 +2490,38 @@ DiagnosedSilenceableFailure ReductionInputFuseOp::apply(
         }
         return false;
       });
-  if (numCandidateSliceOp != 1) {
+  if (numCandidateSliceOp > 1) {
     return mlir::emitDefiniteFailure(
-        this->getOperation(), (Twine("only support one candidate extract_slice "
-                                     "of target in the loop but got") +
-                               Twine(numCandidateSliceOp))
-                                  .str());
+        this->getOperation(),
+        (Twine("only support at most one candidate extract_slice "
+               "of target in the loop but got") +
+         Twine(numCandidateSliceOp))
+            .str());
   }
 
-  OpBuilder b(sliceOp);
+  OpBuilder b(target->getContext());
+  if (sliceOp)
+    b.setInsertionPoint(sliceOp);
+  else
+    b.setInsertionPoint(forOp.getBody(), forOp.getBody()->begin());
+
   forOp->setOperand(
       forOp.getNumControlOperands() + (argIt - iterOperands.begin()),
       linalgOp.getDpsInitOperands().front()->get());
   auto clonedTarget = b.clone(*target);
   clonedTarget->replaceUsesOfWith(linalgOp.getDpsInitOperands().front()->get(),
                                   iterArg);
-  auto tileableTarget = cast<TilingInterface>(clonedTarget);
-  FailureOr<Value> tiledTarget = tileableTarget.generateResultTileValue(
-      b, 0, sliceOp.getMixedOffsets(), sliceOp.getMixedSizes());
-  if (failed(tiledTarget)) {
-    return mlir::emitDefiniteFailure(this->getOperation(),
-                                     "failed to tile target op\n");
+  Operation* tiledOp = clonedTarget;
+  if (sliceOp) {
+    auto tileableTarget = cast<TilingInterface>(clonedTarget);
+    FailureOr<Value> tiledTarget = tileableTarget.generateResultTileValue(
+        b, 0, sliceOp.getMixedOffsets(), sliceOp.getMixedSizes());
+    if (failed(tiledTarget)) {
+      return mlir::emitDefiniteFailure(this->getOperation(),
+                                       "failed to tile target op\n");
+    }
+    tiledOp = tiledTarget->getDefiningOp();
   }
-  Operation* tiledOp = tiledTarget->getDefiningOp();
   Value pred =
       b.create<arith::CmpIOp>(tiledOp->getLoc(), arith::CmpIPredicate::eq,
                               forOp.getInductionVar(), forOp.getLowerBound());
@@ -2518,7 +2531,12 @@ DiagnosedSilenceableFailure ReductionInputFuseOp::apply(
     return mlir::emitDefiniteFailure(
         this->getOperation(), "failed to build conditional_generic op\n");
   }
-  sliceOp->replaceAllUsesWith(conditionalGenericOp->getResults());
+  if (sliceOp) {
+    sliceOp->replaceAllUsesWith(conditionalGenericOp->getResults());
+  } else {
+    iterArg.replaceAllUsesExcept(conditionalGenericOp->getResult(0),
+                                 {conditionalGenericOp});
+  }
   results.set(getTiledTarget().cast<OpResult>(), {conditionalGenericOp});
   results.set(getFusedLoop().cast<OpResult>(), {forOp});
   return DiagnosedSilenceableFailure(success());
diff --git a/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc b/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc
index e7a44c09b0c..aace3f60e1a 100644
--- a/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc
@@ -173,11 +173,15 @@ transform_dialect::FoldProducerExtractSliceOp buildFoldProducerExtractSlice(
 
 transform::PadOp buildPadOp(OpBuilder& b, Location& loc, Value target,
                             ArrayRef<int64_t> paddingDimensions,
-                            int64_t numOperands) {
+                            int64_t numOperands,
+                            ArrayRef<Type> paddingTypes = {}) {
   auto pdlType = pdl::OperationType::get(b.getContext());
   // TODO(wyzero): support other types.
   SmallVector<Attribute> paddingAttrs(numOperands,
                                       b.getZeroAttr(b.getF32Type()));
+  for (auto [idx, type] : llvm::enumerate(paddingTypes)) {
+    paddingAttrs[idx] = b.getZeroAttr(type);
+  }
   return b.create<transform::PadOp>(loc, pdlType, target,
                                     b.getArrayAttr(paddingAttrs),
                                     b.getI64ArrayAttr(paddingDimensions),
@@ -274,6 +278,44 @@ buildConvertPaddingPlaceholderToConstOp(OpBuilder& b, Location& loc,
       loc, pdlType, target);
 }
 
+transform_dialect::LinalgEagerlyBackwardInitTensorOp
+buildLinalgEagerlyBackwardInitTensorOp(OpBuilder& b, Location& loc,
+                                       Value target) {
+  auto pdlType = pdl::OperationType::get(b.getContext());
+  return b.create<transform_dialect::LinalgEagerlyBackwardInitTensorOp>(
+      loc, pdlType, target);
+}
+
+transform_dialect::DISCFuseIntoContainingOp buildDISCFuseIntoContainingOp(
+    OpBuilder& b, Location& loc, Value target, Value anchor) {
+  auto pdlType = pdl::OperationType::get(b.getContext());
+  return b.create<transform_dialect::DISCFuseIntoContainingOp>(loc, pdlType,
+                                                               target, anchor);
+}
+
+transform_dialect::ReductionOutputFuseOp buildReductionOutputFuseOp(
+    OpBuilder& b, Location& loc, Value target, Value loop) {
+  SmallVector<Type> pdlTypes(2, pdl::OperationType::get(b.getContext()));
+  return b.create<transform_dialect::ReductionOutputFuseOp>(loc, pdlTypes,
+                                                            target, loop);
+}
+
+transform_dialect::ReductionInputFuseOp buildReductionInputFuseOp(OpBuilder& b,
+                                                                  Location& loc,
+                                                                  Value target,
+                                                                  Value loop) {
+  SmallVector<Type> pdlTypes(2, pdl::OperationType::get(b.getContext()));
+  return b.create<transform_dialect::ReductionInputFuseOp>(loc, pdlTypes,
+                                                           target, loop);
+}
+
+transform_dialect::VectorizeConditionalGenericOp
+buildVectorizeConditionalGenericOp(OpBuilder& b, Location& loc, Value target) {
+  auto pdlType = pdl::OperationType::get(b.getContext());
+  return b.create<transform_dialect::VectorizeConditionalGenericOp>(
+      loc, pdlType, target);
+}
+
 class ParsedFromFileScheduleFactory : public ScheduleFactoryWithNoGuard {
  public:
   explicit ParsedFromFileScheduleFactory(int64_t id, PatternKind kind,
@@ -595,12 +637,10 @@ LogicalResult Aarch64GEMMDefaultScheduleWithEpilogueFactory::assignSchedule(
     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);
+  rootHandle = 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);
@@ -617,13 +657,45 @@ LogicalResult Aarch64GEMMDefaultScheduleWithEpilogueFactory::assignSchedule(
   // TODO(wyzero): query cpuinfo.
   int64_t hardwareVectorSizeInBytes = 4;
 
+  // supposed loop order:
+  //  loop_m0
+  //   loop_n0
+  //    loop_m1
+  //     loop_n1
+  //      loop_k0 {
+  //        inner_most_gemm
+  //      }
+  bool m0Skipped = (M != ShapedType::kDynamicSize && M <= M0);
+  bool n0Skipped = (N != ShapedType::kDynamicSize && N <= N0);
+  bool m1Skipped = (M != ShapedType::kDynamicSize && M <= M1);
+  bool n1Skipped = (N != ShapedType::kDynamicSize && N <= N1);
+  bool k0Skipped = (K != ShapedType::kDynamicSize && K <= K0);
+
+  // for very small m or n
+  if (m1Skipped || n1Skipped) {
+    // TODO(wyzero): finetune the schedule for small m or n
+    buildTileOp(b, loc, matmul, {1, 1, 1}, {0, 2, 1});
+    variant = buildRunCanonicalizer(b, loc, variant);
+    variant = buildDISCBufferize(b, loc, variant);
+    b.create<transform::YieldOp>(loc);
+    return success();
+  }
+
   // 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));
+  if (!m0Skipped) {
+    auto tileOp1 = buildTileOp(b, loc, rootHandle, {M0, N0}, {0, 1});
+    rootHandle = tileOp1->getResult(0);
+    matmul = buildFuseIntoContainingOp(b, loc, matmul, tileOp1->getResult(1));
+    fill = buildFuseIntoContainingOp(b, loc, fill, tileOp1->getResult(1));
+  } else if (!n0Skipped) {
+    auto tileOp1 = buildTileOp(b, loc, rootHandle, {M0, N0}, {0, 1});
+    rootHandle = tileOp1->getResult(0);
+    matmul = buildFuseIntoContainingOp(b, loc, matmul, tileOp1->getResult(2));
+    fill = buildFuseIntoContainingOp(b, loc, fill, tileOp1->getResult(2));
+  }
 
   // second level tile and fuse matmul and fill op.
-  auto tileOp2 = buildTileOp(b, loc, tileOp1->getResult(0), {M1, N1}, {0, 1});
+  auto tileOp2 = buildTileOp(b, loc, rootHandle, {M1, N1}, {0, 1});
   // pad root ops to register level size (static size).
   int64_t numOperandsUpperBound = fusionPattern.getOperands().size() +
                                   fusionPattern.getResults().size() +
@@ -667,27 +739,9 @@ LogicalResult Aarch64GEMMDefaultScheduleWithEpilogueFactory::assignSchedule(
                     (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);
-
+    int loopLevel = 4 - n0Skipped - m1Skipped - n1Skipped - k0Skipped;
     if (loopLevel <= 0) {
       return m->emitError()
              << "failed to cache the packed input due to loopLevel = "
@@ -1007,6 +1061,295 @@ LogicalResult Aarch64GEMMLargeKScheduleFactory::assignSchedule(
   return success();
 }
 
+class Aarch64GEMMLargeKScheduleWithEpilogueFactory
+    : public Aarch64GEMMLargeKScheduleFactory {
+ public:
+  using Aarch64GEMMLargeKScheduleFactory::Aarch64GEMMLargeKScheduleFactory;
+  bool checkFusionPatternProperties(PatternDescription&) override;
+  LogicalResult assignSchedule(PatternDescription&, ModuleOp) override;
+};
+
+bool Aarch64GEMMLargeKScheduleWithEpilogueFactory::checkFusionPatternProperties(
+    PatternDescription& pd) {
+  auto& fusionPattern = pd.getFusionPattern();
+  auto& rootOps = fusionPattern.getRootOps();
+  // Only support single output a.t.m.
+  if (rootOps.size() != 1) return false;
+
+  auto dominantOp = fusionPattern.getDominantOp();
+  return rootOps[0] != dominantOp && isa<lmhlo::DotGeneralOp>(dominantOp);
+}
+
+LogicalResult Aarch64GEMMLargeKScheduleWithEpilogueFactory::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 lhsContractingDims = dimNumbers.getLhsContractingDimensions();
+  auto rhsContractingDims = dimNumbers.getRhsContractingDimensions();
+  if (lhsContractingDims.size() != 1 || rhsContractingDims.size() != 1) {
+    return m->emitError() << "only support exactly 1 contract dim\n";
+  }
+  bool lhsTranspose = (lhsContractingDims[0] == lhsTy.getRank() - 2);
+  bool rhsTranspose = (rhsContractingDims[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]);
+  }
+  rootHandle = buildLinalgEagerlyBackwardInitTensorOp(b, loc, rootHandle);
+
+  auto forEachThreadOp = buildTileToForEachThreadOp(b, loc, rootHandle, {1, 1});
+  Value forEachThreadLoop = forEachThreadOp->getResult(0);
+  rootHandle = 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 level tile size for dimension m
+  int64_t M0 = 288, M1 = 8;
+  // first/second level tile size for dimension n
+  int64_t N0 = 204, N1 = 12;
+  // first/second level tile size for dimension k
+  int64_t K0 = 512, K1 = 1;
+  // TODO(wyzero): query cpuinfo.
+  int64_t hardwareVectorSizeInBytes = 4;
+
+  // 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 = (K1 != 1) && (K == ShapedType::kDynamicSize ||
+                                 (K > K0 && (K % K0 != 0 || K0 % K1 != 0)) ||
+                                 (K <= K0 && K > K1 && K % K1 != 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));
+  // supposed loop order:
+  //  loop_m0
+  //   loop_k0
+  //    loop_n0
+  //     loop_m1
+  //      loop_n1
+  //       loop_k1 {
+  //         inner_most_gemm
+  //       }
+  // in case:
+  // - the size of dimension K <= K0, then loop_k0 will be folded.
+  bool m0Skipped = (M != ShapedType::kDynamicSize && M <= M0);
+  // - the size of dimension K <= K0, then loop_k0 will be folded.
+  bool k0Skipped = (K != ShapedType::kDynamicSize && K <= K0);
+  // - the size of dimension N <= N0, then loop_n0 will be folded.
+  bool n0Skipped = (N != ShapedType::kDynamicSize && N <= N0);
+  // - the size of dimension M <= M1, then loop_m1 will be folded.
+  bool m1Skipped = (M != ShapedType::kDynamicSize && M <= M1);
+  // - the size of dimension N <= N1, then loop_n1 will be folded.
+  bool n1Skipped = (N != ShapedType::kDynamicSize && N <= N1);
+  // - the size of dimension K <= K0, then loop_k0 will be folded.
+  bool k1Skipped = (K != ShapedType::kDynamicSize && K <= K1);
+
+  // for very small m and n
+  if (m1Skipped || n1Skipped || k1Skipped) {
+    buildTileOp(b, loc, matmul, {1, 1, 1}, {0, 2, 1});
+    variant = buildRunCanonicalizer(b, loc, variant);
+    variant = buildDISCBufferize(b, loc, variant);
+    b.create<transform::YieldOp>(loc);
+    return success();
+  }
+
+  if (!m0Skipped) {
+    auto tileM0 = buildTileOp(b, loc, rootHandle, {M0}, {0, 1});
+    rootHandle = tileM0->getResult(0);
+    Value loopM0 = tileM0->getResult(1);
+    variant = buildRunCanonicalizer(b, loc, variant);
+    matmul = buildDISCFuseIntoContainingOp(b, loc, matmul, loopM0);
+    fill = buildDISCFuseIntoContainingOp(b, loc, fill, loopM0);
+  }
+
+  if (!k0Skipped) {
+    auto tileK0 = buildTileOp(b, loc, matmul, {0, 0, K0}, {0, 1, 2});
+    matmul = tileK0->getResult(0);
+    auto outputFuseOp =
+        buildReductionOutputFuseOp(b, loc, rootHandle, tileK0->getResult(1));
+    rootHandle = outputFuseOp->getResult(0);
+    auto inputFuseOp =
+        buildReductionInputFuseOp(b, loc, fill, outputFuseOp->getResult(1));
+    fill = inputFuseOp->getResult(0);
+    variant = buildRunCanonicalizer(b, loc, variant);
+  }
+
+  if (!n0Skipped) {
+    auto tileN0 = buildTileOp(b, loc, rootHandle, {0, N0}, {0, 1});
+    rootHandle = tileN0->getResult(0);
+    Value loopN0 = tileN0->getResult(1);
+    variant = buildRunCanonicalizer(b, loc, variant);
+    matmul = buildDISCFuseIntoContainingOp(b, loc, matmul, loopN0);
+    fill = buildDISCFuseIntoContainingOp(b, loc, fill, loopN0);
+  }
+
+  auto tileM1 = buildTileOp(b, loc, rootHandle, {M1, 0}, {0, 1});
+  rootHandle = tileM1->getResult(0);
+  Value loopM1 = tileM1->getResult(1);
+  variant = buildRunCanonicalizer(b, loc, variant);
+  matmul = buildDISCFuseIntoContainingOp(b, loc, matmul, loopM1);
+  fill = buildDISCFuseIntoContainingOp(b, loc, fill, loopM1);
+
+  auto tileN1 = buildTileOp(b, loc, rootHandle, {0, N1}, {0, 1});
+  rootHandle = tileN1->getResult(0);
+  Value loopN1 = tileN1->getResult(1);
+  variant = buildRunCanonicalizer(b, loc, variant);
+  // pad root ops to register level size (static size).
+  int64_t numOperandsUpperBound = fusionPattern.getOperands().size() +
+                                  fusionPattern.getResults().size() +
+                                  fusionPattern.getInternalResults().size() + 2;
+  SmallVector<Type> paddingTypes;
+  if (!k0Skipped) paddingTypes.push_back(b.getIntegerType(1));
+  rootHandle = buildPadOp(b, loc, rootHandle, {0, 1}, numOperandsUpperBound,
+                          paddingTypes);
+  auto padForRootOp = buildMatchOp(b, loc, variant, {"tensor.pad"});
+  buildReplaceConstPaddingValueOp(b, loc, padForRootOp, "kAny");
+  matmul = buildDISCFuseIntoContainingOp(b, loc, matmul, loopN1);
+  fill = buildDISCFuseIntoContainingOp(b, loc, fill, loopN1);
+
+  auto tileK1 = buildTileOp(b, loc, matmul, {0, 0, K1}, {0, 1, 2});
+  matmul = tileK1->getResult(0);
+  variant = buildRunCanonicalizer(b, loc, variant);
+  Value weightSlice = buildGetProducerOfOperand(b, loc, matmul, 1);
+  buildFoldProducerExtractSlice(b, loc, weightSlice, 2);
+  // pad to match the requirement of hardware vector/tensor instruction.
+  matmul = buildPadOp(b, loc, matmul, {0, 1, 2}, 3);
+  fill = buildPadOp(b, loc, fill, {0, 1}, 3, paddingTypes);
+  Value padForInput = buildGetProducerOfOperand(b, loc, matmul, 0);
+  Value padForWeight = buildGetProducerOfOperand(b, loc, matmul, 1);
+
+  if (packInput) {
+    // We want to cache the packed A below loop_k0 and above loop_n0.
+    // Thus the initial loop_level is 4.
+    int loopLevel = 4 - n0Skipped - m1Skipped - n1Skipped - k1Skipped;
+    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 = {K1, M1};
+      permutation = {2, 0, 1, 3};
+    } else {
+      tileSizes = {M1, K1};
+      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 > K1) &&
+                     (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, K1};
+      permutation = {0, 2, 3, 1};
+    } else {
+      tileSizes = {K1, 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 conditionalOps =
+      buildMatchOp(b, loc, variant, {"disc_linalg_ext.conditional_generic"});
+  buildVectorizeConditionalGenericOp(b, loc, conditionalOps);
+  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 = buildRunCanonicalizer(b, loc, variant);
+
+  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);
+  variant = buildDecomposeVectors(b, loc, variant, hardwareVectorSizeInBytes);
+  b.create<transform::YieldOp>(loc);
+  return success();
+}
+
 DISC_TRANSFORM_SCHEDULE(PatternKind::kGEMM, kDefaultScheduleFactoryPriority,
                         Aarch64GEMMDefaultScheduleFactory,
                         ArrayRef<StringRef>{kDefaultScheduleFactoryTag});
@@ -1019,6 +1362,10 @@ DISC_TRANSFORM_SCHEDULE(PatternKind::kGEMM, 100,
                         Aarch64GEMMLargeKScheduleFactory,
                         ArrayRef<StringRef>{"large_k"});
 
+DISC_TRANSFORM_SCHEDULE(PatternKind::kGEMM, 110,
+                        Aarch64GEMMLargeKScheduleWithEpilogueFactory,
+                        ArrayRef<StringRef>{"large_k_epilogue"});
+
 }  // namespace
 
 const char* kDefaultScheduleFactoryTag = "default";


Test patch:
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_epilogue_small_static_shape_test0.mlir b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_epilogue_small_static_shape_test0.mlir
new file mode 100644
index 00000000000..01b6aa8a858
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_epilogue_small_static_shape_test0.mlir
@@ -0,0 +1,25 @@
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 0 : i32}} {
+  func.func @main(%arg0: tensor<5x?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<768x7xf32>} : () -> tensor<768x7xf32>
+      %1:2 = tf_executor.island wraps "tf.MatMul"(%arg0, %0) {transpose_a = false, transpose_b = false} : (tensor<5x?xf32>, tensor<768x7xf32>) -> (tensor<?x?xf32>)
+      %2:2 = tf_executor.island wraps "tf.Const"() {value = dense<-0.1> : tensor<7xf32>} : () -> tensor<7xf32>
+      %3:2 = tf_executor.island wraps "tf.BiasAdd"(%1, %2) : (tensor<?x?xf32>, tensor<7xf32>) -> 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/data/matmul_epilogue_small_static_shape_test1.mlir b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_epilogue_small_static_shape_test1.mlir
new file mode 100644
index 00000000000..fbacdacfafd
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_epilogue_small_static_shape_test1.mlir
@@ -0,0 +1,25 @@
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 0 : i32}} {
+  func.func @main(%arg0: tensor<20x?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<768x34xf32>} : () -> tensor<768x34xf32>
+      %1:2 = tf_executor.island wraps "tf.MatMul"(%arg0, %0) {transpose_a = false, transpose_b = false} : (tensor<20x?xf32>, tensor<768x34xf32>) -> (tensor<?x?xf32>)
+      %2:2 = tf_executor.island wraps "tf.Const"() {value = dense<-0.1> : tensor<34xf32>} : () -> tensor<34xf32>
+      %3:2 = tf_executor.island wraps "tf.BiasAdd"(%1, %2) : (tensor<?x?xf32>, tensor<34xf32>) -> 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/data/matmul_epilogue_small_static_shape_test2.mlir b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_epilogue_small_static_shape_test2.mlir
new file mode 100644
index 00000000000..02b6d953912
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_epilogue_small_static_shape_test2.mlir
@@ -0,0 +1,25 @@
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 0 : i32}} {
+  func.func @main(%arg0: tensor<500x?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<768x34xf32>} : () -> tensor<768x34xf32>
+      %1:2 = tf_executor.island wraps "tf.MatMul"(%arg0, %0) {transpose_a = false, transpose_b = false} : (tensor<500x?xf32>, tensor<768x34xf32>) -> (tensor<?x?xf32>)
+      %2:2 = tf_executor.island wraps "tf.Const"() {value = dense<-0.1> : tensor<34xf32>} : () -> tensor<34xf32>
+      %3:2 = tf_executor.island wraps "tf.BiasAdd"(%1, %2) : (tensor<?x?xf32>, tensor<34xf32>) -> 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/data/matmul_epilogue_small_static_shape_test3.mlir b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_epilogue_small_static_shape_test3.mlir
new file mode 100644
index 00000000000..ceab6154a3a
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_epilogue_small_static_shape_test3.mlir
@@ -0,0 +1,25 @@
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 0 : i32}} {
+  func.func @main(%arg0: tensor<20x?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<768x340xf32>} : () -> tensor<768x340xf32>
+      %1:2 = tf_executor.island wraps "tf.MatMul"(%arg0, %0) {transpose_a = false, transpose_b = false} : (tensor<20x?xf32>, tensor<768x340xf32>) -> (tensor<?x?xf32>)
+      %2:2 = tf_executor.island wraps "tf.Const"() {value = dense<-0.1> : tensor<340xf32>} : () -> tensor<340xf32>
+      %3:2 = tf_executor.island wraps "tf.BiasAdd"(%1, %2) : (tensor<?x?xf32>, tensor<340xf32>) -> 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
index 32a6d58d332..8c91f3c3e87 100644
--- a/tao_compiler/mlir/disc/tests/disc-transform/matmul_epilogue.cc
+++ b/tao_compiler/mlir/disc/tests/disc-transform/matmul_epilogue.cc
@@ -80,4 +80,76 @@ TEST(PackedMatmul, F32_complex_test0) {
       /*profiling*/ true));
 }
 
+TEST(PackedMatmul, F32_SmallStaticShape_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_small_static_shape_test0.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 1,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"5x768xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
+TEST(PackedMatmul, F32_SmallStaticShape_test1) {
+  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_small_static_shape_test1.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 1,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"20x768xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
+TEST(PackedMatmul, F32_SmallStaticShape_test2) {
+  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_small_static_shape_test2.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 1,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"500x768xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
+TEST(PackedMatmul, F32_SmallStaticShape_test3) {
+  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_small_static_shape_test3.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 1,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"20x768xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
 }  // namespace mlir_test