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/transforms/disc_transform_legalize_to_loop.cc b/tao_compiler/mlir/disc/transforms/disc_transform_legalize_to_loop.cc
index 21b4bf65640..418e8913a43 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
@@ -94,6 +94,12 @@ struct DiscTransformLegalizeToLoopPass
   LogicalResult handleCpuFusionOp(OpBuilder& b, Operation* fusion,
                                   ShapeAnalysis& shapeAnalysis,
                                   ScheduleDispatcher& scheduleDispatcher);
+
+  // Inject schedule selection logic
+  LogicalResult injectScheduleSelectionIR(
+      OpBuilder& b, PatternDescription& pd,
+      SmallVectorImpl<Operation*>& clonedOps);
+
   // Outlines the fusion op to a standalone module op.
   LogicalResult outlineFusionOp(lmhlo::FusionOp fusionOp,
                                 FusionPattern& fusionPattern,
@@ -206,9 +212,52 @@ LogicalResult DiscTransformLegalizeToLoopPass::inlineTransformedModule(
   return success();
 }
 
+LogicalResult DiscTransformLegalizeToLoopPass::injectScheduleSelectionIR(
+    OpBuilder& b, PatternDescription& pd,
+    SmallVectorImpl<Operation*>& clonedOps) {
+  auto fusionOp = pd.getFusionOp();
+  auto factories =
+      ScheduleFactoryRegistry::get().getAllCandidateScheduleFactories(pd);
+  if (factories.empty()) {
+    return fusionOp->emitError() << "failed to find candidate schedule.\n";
+  }
+  // The returned candidate factories are sorted by priority.
+  // The last schedule should always have `noGuardCondition`.
+  if (!factories.back()->noGuardCondition(pd)) {
+    return fusionOp->emitError()
+           << "There are no canidate schedules for some shapes of "
+           << pd.getTaggedPatternStr() << "\n";
+  }
+  // Only one candidate, no need to inject selection logic.
+  if (factories.size() == 1) {
+    clonedOps.push_back(fusionOp);
+    return success();
+  }
+  for (size_t i = 0; i < factories.size() - 1; ++i) {
+    auto cloned = cast<lmhlo::FusionOp>(b.clone(*fusionOp.getOperation()));
+    mergeFusionTag(b, cloned, factories[i]->getTagSet());
+    FusionPattern fusionPattern(cloned, &pd.getShapeAnalysis());
+    PatternDescription clonedPd(cloned, fusionPattern, pd.getShapeAnalysis());
+    Value pred;
+    if (failed(factories[i]->buildGuardCondition(b, cloned->getLoc(), clonedPd,
+                                                 pred)))
+      return cloned->emitError() << "faield to build guard IR\n";
+
+    auto ifOp = b.create<scf::IfOp>(cloned->getLoc(), llvm::None, pred, true);
+    cloned->moveBefore(ifOp.thenBlock(), ifOp.thenBlock()->begin());
+    fusionOp->moveBefore(ifOp.elseBlock(), ifOp.elseBlock()->begin());
+    clonedOps.push_back(cloned);
+    b.setInsertionPointToStart(ifOp.elseBlock());
+  }
+  mergeFusionTag(b, fusionOp, factories.back()->getTagSet());
+  clonedOps.push_back(fusionOp);
+  return success();
+}
+
 LogicalResult DiscTransformLegalizeToLoopPass::handleCpuFusionOp(
     OpBuilder& b, Operation* fusion, ShapeAnalysis& shapeAnalysis,
     ScheduleDispatcher& scheduleDispatcher) {
+  b.setInsertionPoint(fusion);
   auto fusionOp = cast<lmhlo::FusionOp>(fusion);
   assert(fusionOp);
   FusionPattern fusionPattern(fusionOp, &shapeAnalysis);
@@ -217,34 +266,51 @@ LogicalResult DiscTransformLegalizeToLoopPass::handleCpuFusionOp(
     return success();
   }
 
-  // 1, Outline the fusion to a standalone module op.
-  OwningOpRef<ModuleOp> m;
-  if (failed(outlineFusionOp(fusionOp, fusionPattern, m))) {
-    return fusionOp->emitError() << "failed to outlineFusionOp\n";
-  }
-  LLVM_DEBUG(llvm::dbgs() << "After outline fusion op:\n" << m.get() << "\n");
-
-  // 2, assign a default schedule for each pattern here.
-  PatternDescription patternDescription(fusionOp, fusionPattern, shapeAnalysis);
-  if (failed(scheduleDispatcher.dispatch(patternDescription, m.get()))) {
-    return fusionOp->emitError() << "failed to assignSchedule\n";
+  // 0, inject schedule selection logic
+  // clone the fusion op, each for one candidate schedule.
+  SmallVector<Operation*> clonedFusionOps;
+  PatternDescription pd(fusionOp, fusionPattern, shapeAnalysis);
+  if (failed(injectScheduleSelectionIR(b, pd, clonedFusionOps))) {
+    return fusionOp->emitError() << "failed to injectScheduleSelectionIR\n";
   }
-  LLVM_DEBUG(llvm::dbgs() << "After assign schedule for fusion op:\n"
+  LLVM_DEBUG(llvm::dbgs() << "After injectScheduleSelectionIR:\n"
                           << m.get() << "\n");
 
-  // 3, Build a nested pass pipeline to legalize the outlined fusion op.
-  if (failed(runTransformPipeline(m.get()))) {
-    return fusionOp->emitError() << "failed to run runTransformPipeline\n";
-  }
-  LLVM_DEBUG(llvm::dbgs() << "After run transform pipeline:\n"
-                          << m.get() << "\n");
+  for (auto fusion : clonedFusionOps) {
+    b.setInsertionPoint(fusion);
+    auto fusionOp = cast<lmhlo::FusionOp>(fusion);
+    FusionPattern fusionPattern(fusionOp, &shapeAnalysis);
+    // 1, Outline the fusion to a standalone module op.
+    OwningOpRef<ModuleOp> m;
+    if (failed(outlineFusionOp(fusionOp, fusionPattern, m))) {
+      return fusionOp->emitError() << "failed to outlineFusionOp\n";
+    }
+    LLVM_DEBUG(llvm::dbgs() << "After outline fusion op:\n" << m.get() << "\n");
 
-  // 4, Inline the lowered IR into the orignal module.
-  if (failed(inlineTransformedModule(b, fusion, fusionPattern, m.get()))) {
-    return fusion->emitError() << "failed to inlineTransformedModule\n";
+    // 2, assign a default schedule for each pattern here.
+    PatternDescription patternDescription(fusionOp, fusionPattern,
+                                          shapeAnalysis);
+    if (failed(scheduleDispatcher.dispatch(patternDescription, m.get()))) {
+      return fusionOp->emitError() << "failed to assignSchedule\n";
+    }
+    LLVM_DEBUG(llvm::dbgs() << "After assign schedule for fusion op:\n"
+                            << m.get() << "\n");
+
+    // 3, Build a nested pass pipeline to legalize the outlined fusion op.
+    if (failed(runTransformPipeline(m.get()))) {
+      return fusionOp->emitError() << "failed to run runTransformPipeline\n";
+    }
+    LLVM_DEBUG(llvm::dbgs() << "After run transform pipeline:\n"
+                            << m.get() << "\n");
+
+    // 4, Inline the lowered IR into the orignal module.
+    if (failed(inlineTransformedModule(b, fusion, fusionPattern, m.get()))) {
+      return fusion->emitError() << "failed to inlineTransformedModule\n";
+    }
+    LLVM_DEBUG(llvm::dbgs() << "After inline transformed module:\n"
+                            << *fusion << "\n");
   }
-  LLVM_DEBUG(llvm::dbgs() << "After inline transformed module:\n"
-                          << *fusion << "\n");
+
   return success();
 }
 
@@ -269,6 +335,11 @@ void DiscTransformLegalizeToLoopPass::runOnOperation() {
 
   // Assign a transform schedule for the given fusion pattern.
   ScheduleDispatcher scheduleDispatcher{transformFileName_};
+  if (failed(scheduleDispatcher.parseModuleFromFile(b.getContext()))) {
+    func->emitError() << "failed to parse transform module form "
+                      << transformFileName_ << " .\n";
+    return signalPassFailure();
+  }
 
   for (Operation* fusion : gpu_fusion_worklist) {
     // TODO(disc): handling stitch fusion on GPU.
diff --git a/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc b/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc
index 07db691c052..633cd53ab3a 100644
--- a/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc
@@ -236,8 +236,49 @@ SplitHandlesOp buildSplitHandlesOp(OpBuilder& b, Location& loc, Value target,
   return b.create<SplitHandlesOp>(loc, pdlTypes, target, num_result_handles);
 }
 
-LogicalResult aarch64GEMMDefaultScheduleFactory(PatternDescription& pd,
-                                                ModuleOp m) {
+transform_dialect::InlineReductionInitializerOp
+buildInlineReductionInitializerOp(OpBuilder& b, Location& loc, Value initOp,
+                                  Value loopOp, Value readerOp) {
+  auto pdlType = pdl::OperationType::get(b.getContext());
+  return b.create<transform_dialect::InlineReductionInitializerOp>(
+      loc, pdlType, initOp, loopOp, readerOp);
+}
+
+class ParsedFromFileScheduleFactory : public ScheduleFactoryWithNoGuard {
+ public:
+  explicit ParsedFromFileScheduleFactory(int64_t id, PatternKind kind,
+                                         ArrayRef<StringRef> tags,
+                                         ModuleOp transformModule);
+  LogicalResult assignSchedule(PatternDescription&, ModuleOp) override;
+
+ private:
+  ModuleOp transformModule_;
+};
+
+ParsedFromFileScheduleFactory::ParsedFromFileScheduleFactory(
+    int64_t id, PatternKind kind, ArrayRef<StringRef> tags,
+    ModuleOp transformModule)
+    : ScheduleFactoryWithNoGuard(id, kind, tags),
+      transformModule_(transformModule) {}
+
+LogicalResult ParsedFromFileScheduleFactory::assignSchedule(
+    PatternDescription& pd, ModuleOp m) {
+  OpBuilder b(m);
+  for (auto& op : transformModule_.getBody()->getOperations()) {
+    if (!isa<transform::TransformOpInterface>(&op)) continue;
+    m.push_back(b.clone(op));
+  }
+  return success();
+}
+
+class Aarch64GEMMDefaultScheduleFactory : public ScheduleFactoryWithNoGuard {
+ public:
+  using ScheduleFactoryWithNoGuard::ScheduleFactoryWithNoGuard;
+  LogicalResult assignSchedule(PatternDescription&, ModuleOp) override;
+};
+
+LogicalResult Aarch64GEMMDefaultScheduleFactory::assignSchedule(
+    PatternDescription& pd, ModuleOp m) {
   OpBuilder b(m);
   b.setInsertionPointToStart(&m.getBodyRegion().front());
   Location loc = m.getLoc();
@@ -427,11 +468,250 @@ LogicalResult aarch64GEMMDefaultScheduleFactory(PatternDescription& pd,
   return success();
 }
 
-DISC_TRANSFORM_SCHEDULE(PatternKind::kGEMM, "",
-                        aarch64GEMMDefaultScheduleFactory);
+class Aarch64GEMMLargeKScheduleFactory : public ScheduleFactory {
+ public:
+  using ScheduleFactory::ScheduleFactory;
+  LogicalResult assignSchedule(PatternDescription&, ModuleOp) override;
+  LogicalResult buildGuardCondition(OpBuilder& b, Location loc,
+                                    PatternDescription&, Value&) override;
+};
+
+LogicalResult Aarch64GEMMLargeKScheduleFactory::buildGuardCondition(
+    OpBuilder& b, Location loc, PatternDescription& pd, Value& pred) {
+  auto& fusionPattern = pd.getFusionPattern();
+  auto dotOp =
+      dyn_cast_or_null<lmhlo::DotGeneralOp>(fusionPattern.getDominantOp());
+  if (!dotOp) {
+    llvm::dbgs() << "expect dot_general op as dominant\n";
+    return failure();
+  }
+  Value lhs = dotOp->getOperand(0);
+  auto lhsTy = lhs.getType().cast<MemRefType>();
+  if (lhsTy.getRank() != 2) {
+    return dotOp->emitError() << "only support rank 2 GEMM a.t.m.\n";
+  }
+
+  auto dimNumbers = dotOp.getDotDimensionNumbers();
+  auto lhsContractingDims = dimNumbers.getLhsContractingDimensions();
+  if (lhsContractingDims.size() != 1) {
+    return dotOp->emitError() << "only support exactly 1 contract dim\n";
+  }
+  bool lhsTranspose = (lhsContractingDims[0] == lhsTy.getRank() - 2);
+  Value dimK = b.create<memref::DimOp>(loc, lhs, lhsTranspose ? 0 : 1);
+  Value largeK = b.create<arith::ConstantIndexOp>(loc, 1024);
+  pred = b.create<arith::CmpIOp>(loc, arith::CmpIPredicate::sgt, dimK, largeK);
+  return success();
+}
+
+LogicalResult Aarch64GEMMLargeKScheduleFactory::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];
+
+  // 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);
+
+  // transform.structured.tile_to_foreach_thread_op %matmul num_threads [1, 1]
+  auto forEachThreadOp = buildTileToForEachThreadOp(b, loc, matmul, {1, 1});
+  Value forEachThreadLoop = forEachThreadOp->getResult(0);
+  Value tiledMatmul = forEachThreadOp->getResult(1);
+
+  // transform.structured.fuse_into_containing_op %fill into %0#0
+  auto fuseIntoContainingOp =
+      buildFuseIntoContainingOp(b, loc, fill, forEachThreadLoop);
+
+  // first/second level tile size for dimension m
+  int64_t M0 = 6;
+  // first/second level tile size for dimension n
+  int64_t N0 = 208, N1 = 16;
+  // first level tile size for dimension k
+  int64_t K0 = 512, K1 = 1;
+
+  auto tileOp0 = buildTileOp(b, loc, tiledMatmul, {0, N0, K0}, {0, 2, 1});
+  auto tileOp1 =
+      buildTileOp(b, loc, tileOp0->getResult(0), {M0, N1, K1}, {0, 1, 2});
+
+  // fold extract_slice ops generated by two-level tiling. It's needed to
+  // enable following pad and cache_read schedule.
+  Value weightInnerSlice =
+      buildGetProducerOfOperand(b, loc, tileOp1->getResult(0), 1);
+  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});
+
+  Value padForInput = buildGetProducerOfOperand(b, loc, padOp, 0);
+  Value padForWeight = buildGetProducerOfOperand(b, loc, padOp, 1);
+
+  // Check if we need to pad dimension `m/n/k` if input or weight is packed
+  bool mIsPadded =
+      (M0 != 1) && (M == ShapedType::kDynamicSize || M > M0 && M % M0 != 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 >= M0) &&
+                    (K == ShapedType::kDynamicSize || K > K0) &&
+                    (N == ShapedType::kDynamicSize || N > N0));
+  // supposed loop order:
+  //  loop_k0
+  //   loop_n0
+  //    loop_m0
+  //     loop_n1
+  //      loop_k1 {
+  //        inner_most_gemm
+  //      }
+  // in case:
+  // - 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 m0Skipped = (M != ShapedType::kDynamicSize && M <= M0);
+  // - 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);
+  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 - m0Skipped - 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, M0};
+      permutation = {2, 0, 1, 3};
+    } else {
+      tileSizes = {M0, 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 func = buildMatchOp(b, loc, variant, {"func.func"});
+  buildVectorize(b, loc, func, true);
+
+  variant = buildRunCanonicalizer(b, loc, variant);
+  variant = buildDISCBufferize(b, loc, variant);
+
+  if (!k0Skipped && !k1Skipped) {
+    Value leftFillOp = buildMatchOp(b, loc, variant, {}, nameMap[dotOp]);
+    Value contractOp = buildMatchOp(b, loc, variant, {"vector.contract"});
+    int outterMostLoopLevel =
+        5 - k0Skipped - n0Skipped - m0Skipped - n1Skipped - k1Skipped;
+    auto loop0 = buildGetParentForOp(b, loc, contractOp, outterMostLoopLevel);
+    auto loop1 = buildGetParentForOp(b, loc, contractOp, 2);
+    auto readers = buildMatchOp(b, loc, loop1, {"vector.transfer_read"});
+    auto splitedReaders = buildSplitHandlesOp(b, loc, readers, 3);
+    buildInlineReductionInitializerOp(b, loc, leftFillOp, loop0,
+                                      splitedReaders->getResult(0));
+    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);
+  b.create<transform::YieldOp>(loc);
+  return success();
+}
+
+DISC_TRANSFORM_SCHEDULE(PatternKind::kGEMM, kDefaultScheduleFactoryPriority,
+                        Aarch64GEMMDefaultScheduleFactory,
+                        ArrayRef<StringRef>{kDefaultScheduleFactoryTag});
+
+DISC_TRANSFORM_SCHEDULE(PatternKind::kGEMM, 100,
+                        Aarch64GEMMLargeKScheduleFactory,
+                        ArrayRef<StringRef>{"large_k"});
 
 }  // namespace
 
+const char* kDefaultScheduleFactoryTag = "default";
+
 std::string patternKindToString(PatternKind kind) {
   auto& map = getPatternKindToStringMap();
   auto it = map.find(kind);
@@ -451,7 +731,10 @@ PatternKind patternKindFromString(const std::string& str) {
 PatternDescription::PatternDescription(lmhlo::FusionOp op,
                                        FusionPattern& fusionPattern,
                                        ShapeAnalysis& shapeAnalysis)
-    : op_(op), fusionPattern_(fusionPattern), shapeAnalysis_(shapeAnalysis) {
+    : op_(op),
+      fusionPattern_(fusionPattern),
+      shapeAnalysis_(shapeAnalysis),
+      tagSet_(parsefusionTagSetFromStr(getFusionTagStr(op))) {
   // TODO(wyzero): select the pattern kind according to the `fusionPattern`.
   patternKind_ = PatternKind::kGEMM;
 }
@@ -462,33 +745,105 @@ std::string PatternDescription::getPatternTagStr() const {
   return getFusionTagStr(op_).str();
 }
 
+const std::set<std::string>& PatternDescription::getPatternTagSet() const {
+  return tagSet_;
+}
+
 std::string PatternDescription::getTaggedPatternStr() const {
   return patternKindToString(patternKind_) + "@" + getPatternTagStr();
 }
 
+ScheduleFactory::ScheduleFactory(int64_t id, PatternKind kind,
+                                 ArrayRef<StringRef> tags)
+    : id_(id), kind_(kind) {
+  tagSet_.insert(Twine(id).str());
+  for (auto tag : tags) {
+    tagSet_.insert(tag.str());
+  }
+}
+
+bool ScheduleFactory::accept(PatternDescription& pattern) {
+  return checkKindAndTags(pattern) && checkFusionPatternProperties(pattern);
+}
+
+bool ScheduleFactory::checkKindAndTags(PatternDescription& pattern) {
+  if (pattern.getPatternKind() != kind_) return false;
+  for (auto& tag : pattern.getPatternTagSet())
+    if (tagSet_.find(tag) == tagSet_.end()) return false;
+  return true;
+}
+
+bool ScheduleFactory::checkFusionPatternProperties(PatternDescription&) {
+  return true;
+}
+
+LogicalResult ScheduleFactory::buildGuardCondition(OpBuilder& b, Location loc,
+                                                   PatternDescription&,
+                                                   Value&) {
+  return failure();
+}
+
+LogicalResult ScheduleFactory::assignSchedule(PatternDescription&, ModuleOp) {
+  return failure();
+}
+
+bool ScheduleFactory::noGuardCondition(PatternDescription&) { return false; }
+
 /* static */ ScheduleFactoryRegistry& ScheduleFactoryRegistry::get() {
   static ScheduleFactoryRegistry instance;
   return instance;
 }
 
-bool ScheduleFactoryRegistry::registerScheduleFactory(PatternKind kind,
-                                                      const std::string& tag,
-                                                      ScheduleFactory factory) {
-  return factoryMap_[kind].emplace(tag, factory).second;
+/* static */ int64_t ScheduleFactoryRegistry::getNextUniqueId() {
+  static std::atomic<int64_t> nextIdx{0};
+  return nextIdx++;
+}
+
+bool ScheduleFactoryRegistry::registerScheduleFactory(
+    PatternKind kind, int priority, ScheduleFactoryPtr factory) {
+  return patternMap_[kind].emplace(priority, std::move(factory)).second;
+}
+
+void ScheduleFactoryRegistry::unregisterScheduleFactory(PatternKind kind,
+                                                        int priority) {
+  patternMap_[kind].erase(priority);
+}
+
+ScheduleFactory* ScheduleFactoryRegistry::getScheduleFactoryWithHighestPriority(
+    PatternDescription& pd) {
+  auto& factoryMap = patternMap_[pd.getPatternKind()];
+  for (auto it = factoryMap.rbegin(); it != factoryMap.rend(); ++it) {
+    if (it->second->accept(pd)) return it->second.get();
+  }
+  return nullptr;
 }
 
-ScheduleFactory ScheduleFactoryRegistry::getScheduleFactory(
-    PatternKind kind, const std::string& tag) {
-  auto it = factoryMap_.find(kind);
-  if (it == factoryMap_.end()) return nullptr;
-  auto factoryIt = it->second.find(tag);
-  if (factoryIt == it->second.end()) return nullptr;
-  return factoryIt->second;
+SmallVector<ScheduleFactory*>
+ScheduleFactoryRegistry::getAllCandidateScheduleFactories(
+    PatternDescription& pd) {
+  SmallVector<ScheduleFactory*> factories;
+  auto& factoryMap = patternMap_[pd.getPatternKind()];
+  for (auto it = factoryMap.rbegin(); it != factoryMap.rend(); ++it) {
+    if (it->second->accept(pd)) {
+      factories.push_back(it->second.get());
+      // early stop
+      if (it->second->noGuardCondition(pd)) break;
+    }
+  }
+  return factories;
 }
 
 ScheduleDispatcher::ScheduleDispatcher(const std::string& transformFileName)
     : transformFileName_(transformFileName) {}
 
+ScheduleDispatcher::~ScheduleDispatcher() {
+  int priority = kParsedFromFileScheduleFactoryStartPriority;
+  for (const auto& outter : parsedModuleMap_)
+    for (const auto& inner : outter.second)
+      ScheduleFactoryRegistry::get().unregisterScheduleFactory(outter.first,
+                                                               priority++);
+}
+
 LogicalResult ScheduleDispatcher::parseModuleFromFile(MLIRContext* ctx) {
   if (transformFileName_.empty() || !parsedModuleMap_.empty()) return success();
   std::string expectedFormatStr =
@@ -497,6 +852,7 @@ LogicalResult ScheduleDispatcher::parseModuleFromFile(MLIRContext* ctx) {
   SmallVector<StringRef> patternSettings;
   StringRef(transformFileName_)
       .split(patternSettings, ';', /*MaxSplit=*/-1, /*KeepEmpty=*/false);
+  int priority = kParsedFromFileScheduleFactoryStartPriority;
   for (auto& patternSetting : patternSettings) {
     SmallVector<StringRef> items;
     patternSetting.split(items, ":", /*MaxSplit=*/-1, /*KeepEmpty=*/true);
@@ -512,44 +868,35 @@ LogicalResult ScheduleDispatcher::parseModuleFromFile(MLIRContext* ctx) {
       return failure();
     }
 
-    if (failed(parseTransformModuleFromFile(
-            ctx, items[2], parsedModuleMap_[kind][items[1].str()]))) {
+    auto& transformModule = parsedModuleMap_[kind][items[1].str()];
+    if (failed(parseTransformModuleFromFile(ctx, items[2], transformModule))) {
       llvm::dbgs()
           << "illegal transform file setting, unable to load module from: "
           << items[2] << "\n";
       return failure();
     }
+    SmallVector<StringRef> tags;
+    items[1].split(tags, kFusionTagSeparator, /*MaxSplit*/ -1,
+                   /*KeepEmpty*/ false);
+    ScheduleFactoryRegistry::get().registerScheduleFactory(
+        kind, priority++,
+        std::make_unique<ParsedFromFileScheduleFactory>(
+            ScheduleFactoryRegistry::get().getNextUniqueId(), kind, tags,
+            transformModule.get()));
   }
   return success();
 }
 
-bool ScheduleDispatcher::tryToApplyScheduleFromParsedFile(
-    PatternDescription& pd, ModuleOp m) {
-  auto it = parsedModuleMap_.find(pd.getPatternKind());
-  if (it == parsedModuleMap_.end()) return false;
-  auto tagIt = it->second.find(pd.getPatternTagStr());
-  if (tagIt == it->second.end()) return false;
-  OpBuilder b(m);
-  for (auto& op : tagIt->second.get().getBody()->getOperations()) {
-    if (!isa<transform::TransformOpInterface>(&op)) continue;
-    m.push_back(b.clone(op));
-  }
-  return true;
-}
-
 LogicalResult ScheduleDispatcher::dispatch(PatternDescription& pd, ModuleOp m) {
-  if (failed(parseModuleFromFile(m.getContext()))) return failure();
-  if (tryToApplyScheduleFromParsedFile(pd, m)) return success();
-
-  auto factory = ScheduleFactoryRegistry::get().getScheduleFactory(
-      pd.getPatternKind(), pd.getPatternTagStr());
+  auto factory =
+      ScheduleFactoryRegistry::get().getScheduleFactoryWithHighestPriority(pd);
   if (!factory) {
     llvm::dbgs() << "no default schedule for pattern: "
                  << pd.getTaggedPatternStr() << "\n";
     return failure();
   }
 
-  return factory(pd, m);
+  return factory->assignSchedule(pd, m);
 }
 
 }  // namespace disc_ral
diff --git a/tao_compiler/mlir/disc/transforms/disc_transform_schedule.h b/tao_compiler/mlir/disc/transforms/disc_transform_schedule.h
index 0d601b8e6f2..f3ddc1e5eb6 100644
--- a/tao_compiler/mlir/disc/transforms/disc_transform_schedule.h
+++ b/tao_compiler/mlir/disc/transforms/disc_transform_schedule.h
@@ -16,6 +16,9 @@ limitations under the License.
 #ifndef DISC_TRANSFORMS_TRANSFORM_SCHEDULE_H_
 #define DISC_TRANSFORMS_TRANSFORM_SCHEDULE_H_
 
+#include <map>
+#include <memory>
+#include <set>
 #include <string>
 #include <unordered_map>
 
@@ -52,6 +55,8 @@ class PatternDescription {
   // Returns the full pattern kind str + tag str.
   std::string getTaggedPatternStr() const;
 
+  const std::set<std::string>& getPatternTagSet() const;
+
   // Returns the fusion op this descriptor holds.
   lmhlo::FusionOp getFusionOp() { return op_; }
 
@@ -59,18 +64,77 @@ class PatternDescription {
   FusionPattern& getFusionPattern() { return fusionPattern_; }
 
   // Returns the fusion pattern this descriptor holds.
-  ShapeAnalysis& shapeAnalysis() { return shapeAnalysis_; }
+  ShapeAnalysis& getShapeAnalysis() { return shapeAnalysis_; }
 
  private:
   lmhlo::FusionOp op_;
   FusionPattern& fusionPattern_;
   ShapeAnalysis& shapeAnalysis_;
   PatternKind patternKind_;
+  std::set<std::string> tagSet_;
 };
 
+// The name of the default schedule factory for a pattern kind.
+// The default schedule should always has the lowest priority and generates
+// workable code for any candidate pattern.
+extern const char* kDefaultScheduleFactoryTag;
+constexpr const int kDefaultScheduleFactoryPriority = 0;
+constexpr const int kParsedFromFileScheduleFactoryStartPriority = 10000;
+
 // Factory used to assign specific schedule for the given PatternDescription
-using ScheduleFactory =
-    std::function<LogicalResult(PatternDescription&, ModuleOp)>;
+class ScheduleFactory {
+ public:
+  explicit ScheduleFactory(int64_t id, PatternKind kind,
+                           ArrayRef<StringRef> tags);
+
+  // Returns true if the factory accepts the pattern at compile time.
+  // Note that it's only an conservative check, which means pattern rejected
+  // by this check is definitely not supported by this factory. We need to
+  // inject runtime check logic to further protect the schedule due to missing
+  // some information at compile time.
+  virtual bool accept(PatternDescription&);
+
+  // Returns true if the schedule can always generate workable code for the
+  // given pattern once it passes the compile-time check.
+  virtual bool noGuardCondition(PatternDescription&);
+
+  // Builds the runtime guard ir to protect the shedule generated by this
+  // factory.
+  virtual LogicalResult buildGuardCondition(OpBuilder& b, Location loc,
+                                            PatternDescription&, Value&);
+
+  // Assign the transform schedule and attach it into the module op.
+  // The pattern should be accepted by this factory and the guard condition
+  // should be emitted before successfully.
+  virtual LogicalResult assignSchedule(PatternDescription&, ModuleOp);
+
+  // Returns the id this factory has.
+  int64_t getId() { return id_; }
+
+  // Returns the pattern kind this factory has.
+  PatternKind getPatternKind() { return kind_; }
+
+  // Returns the tag set this factory has.
+  const std::set<std::string>& getTagSet() { return tagSet_; }
+
+ protected:
+  // these are called by `accept`.
+  virtual bool checkKindAndTags(PatternDescription&);
+  virtual bool checkFusionPatternProperties(PatternDescription&);
+
+ protected:
+  int64_t id_;
+  PatternKind kind_;
+  std::set<std::string> tagSet_;
+};
+
+class ScheduleFactoryWithNoGuard : public ScheduleFactory {
+ public:
+  using ScheduleFactory::ScheduleFactory;
+  bool noGuardCondition(PatternDescription&) override { return true; };
+};
+
+using ScheduleFactoryPtr = std::unique_ptr<ScheduleFactory>;
 
 // A registry for different schedule factories.
 class ScheduleFactoryRegistry {
@@ -78,33 +142,53 @@ class ScheduleFactoryRegistry {
   // Returns the singleton
   static ScheduleFactoryRegistry& get();
 
-  // Inserts the new `ScheduleFactory`. Returns true if inserted, otherwise
-  // false.
-  bool registerScheduleFactory(PatternKind kind, const std::string& tag,
-                               ScheduleFactory);
+  // Returns next available id for schedule factory.
+  static int64_t getNextUniqueId();
 
-  // Returns the schedule factory factor according to `kind` and `tag`.
+  // Inserts the new `ScheduleFactory`. Returns true if inserted, otherwise
+  // false. The larger the `prioirty`, the larger the chance the factory being
+  // choosed. Note that we do not allow to assign the same priority for two
+  // different factories with the same pattern kind.
+  bool registerScheduleFactory(PatternKind kind, int priority,
+                               ScheduleFactoryPtr factory);
+  void unregisterScheduleFactory(PatternKind kind, int priority);
+
+  // Returns the schedule factory with the highes priority for `pd`.
   // Returns nullptr if not found.
-  ScheduleFactory getScheduleFactory(PatternKind kind, const std::string& tag);
+  ScheduleFactory* getScheduleFactoryWithHighestPriority(
+      PatternDescription& pd);
+
+  // Returns all suitable schedule factories for `pd`. The returned factory list
+  // is sorted by `priority`. The first one has the highest priority.
+  SmallVector<ScheduleFactory*> getAllCandidateScheduleFactories(
+      PatternDescription& pd);
 
  private:
   ScheduleFactoryRegistry() = default;
-  std::unordered_map<PatternKind,
-                     std::unordered_map<std::string, ScheduleFactory>>
-      factoryMap_;
+  std::unordered_map<PatternKind, std::map<int, ScheduleFactoryPtr>>
+      patternMap_;
 };
 
 // Macros used to define disc transform schedule factory.
-#define DISC_TRANSFORM_SCHEDULE(kind, tag, ...) \
-  DISC_TRANSFORM_SCHEDULE_UNIQ_HELPER(kind, tag, __COUNTER__, __VA_ARGS__)
-
-#define DISC_TRANSFORM_SCHEDULE_UNIQ_HELPER(kind, tag, ctr, ...) \
-  DISC_TRANSFORM_SCHEDULE_UNIQ(kind, tag, ctr, __VA_ARGS__)
-
-#define DISC_TRANSFORM_SCHEDULE_UNIQ(kind, tag, ctr, ...) \
-  static bool unused_ret_val_##ctr =                      \
-      ::mlir::disc_ral::ScheduleFactoryRegistry::get()    \
-          .registerScheduleFactory(kind, tag, __VA_ARGS__);
+#define DISC_TRANSFORM_SCHEDULE(kind, priority, T, ...)               \
+  DISC_TRANSFORM_SCHEDULE_UNIQ_HELPER(__COUNTER__, kind, priority, T, \
+                                      __VA_ARGS__)
+
+#define DISC_TRANSFORM_SCHEDULE_UNIQ_HELPER(ctr, kind, priority, T, ...) \
+  DISC_TRANSFORM_SCHEDULE_UNIQ(ctr, kind, priority, T, __VA_ARGS__)
+
+#define DISC_TRANSFORM_SCHEDULE_UNIQ(ctr, kind, priority, T, ...)             \
+  static bool unused_ret_val_##ctr = []() {                                   \
+    bool ret = ::mlir::disc_ral::ScheduleFactoryRegistry::get()               \
+                   .registerScheduleFactory(                                  \
+                       kind, priority,                                        \
+                       std::make_unique<T>(                                   \
+                           ::mlir::disc_ral::ScheduleFactoryRegistry::        \
+                               getNextUniqueId(),                             \
+                           kind, __VA_ARGS__));                               \
+    if (!ret) ::llvm::dbgs() << "failed to register a new scheduleFactory\n"; \
+    return ret;                                                               \
+  }();
 
 // Assign schedule for the given PatternDescription according to its kind and
 // tag.
@@ -115,15 +199,13 @@ class ScheduleDispatcher {
   // Format of `transformFileName`:
   //  "<kind-0>:<tag-str-0>:<filename-0>;<kind-1>:<tag-str-1>:<filename-1>;"
   explicit ScheduleDispatcher(const std::string& transformFileName);
+  ~ScheduleDispatcher();
 
   // Attaches a schedule for the given pattern description.
   LogicalResult dispatch(PatternDescription& pd, ModuleOp m);
 
- private:
   // Parses schedule modules from the given files.
   LogicalResult parseModuleFromFile(MLIRContext* ctx);
-  // Returns true when applied, otherwise false.
-  bool tryToApplyScheduleFromParsedFile(PatternDescription& pd, ModuleOp m);
 
  private:
   std::string transformFileName_;
diff --git a/tao_compiler/mlir/disc/transforms/fusion_utils.cc b/tao_compiler/mlir/disc/transforms/fusion_utils.cc
index fead6aa7832..e1dfe8c48a1 100644
--- a/tao_compiler/mlir/disc/transforms/fusion_utils.cc
+++ b/tao_compiler/mlir/disc/transforms/fusion_utils.cc
@@ -257,16 +257,38 @@ void addFusionTag(OpBuilder& b, lmhlo::FusionOp op, StringRef tag) {
   std::string oldTag;
   auto attr = op->getAttrOfType<StringAttr>(kFusionOpTagAttr);
   if (attr && attr.getValue().size()) {
-    oldTag = (Twine(attr.getValue()) + "X").str();
+    oldTag = (Twine(attr.getValue()) + kFusionTagSeparator).str();
   }
   op->setAttr(kFusionOpTagAttr, b.getStringAttr((Twine(oldTag) + tag).str()));
 }
 
+void mergeFusionTag(OpBuilder& b, lmhlo::FusionOp op,
+                    const std::set<std::string>& tagSet) {
+  auto opTagSet = parsefusionTagSetFromStr(getFusionTagStr(op));
+  for (const auto& tag : tagSet) opTagSet.insert(tag);
+  auto newTagStr = llvm::join(opTagSet, kFusionTagSeparator);
+  op->setAttr(kFusionOpTagAttr, b.getStringAttr(newTagStr));
+}
+
 StringRef getFusionTagStr(lmhlo::FusionOp op) {
   auto attr = op->getAttrOfType<StringAttr>(kFusionOpTagAttr);
   return attr ? attr.getValue() : "";
 }
 
+std::string fusionTagSetToStr(const std::set<std::string>& tagSet) {
+  SmallVector<StringRef> tags;
+  for (const auto& tag : tagSet) tags.push_back(tag);
+  return llvm::join(tags, kFusionTagSeparator);
+}
+
+std::set<std::string> parsefusionTagSetFromStr(StringRef tagStr) {
+  SmallVector<StringRef> tags;
+  std::set<std::string> tagSet;
+  tagStr.split(tags, kFusionTagSeparator, /*MaxSplit*/ -1, /*KeepEmpty*/ false);
+  for (auto tag : tags) tagSet.insert(tag.str());
+  return tagSet;
+}
+
 // Returns the full name of the fusion op
 // Here full name is composed of the name and tag of the fusion op.
 std::string getFusionFullName(lmhlo::FusionOp op) {
diff --git a/tao_compiler/mlir/disc/transforms/fusion_utils.h b/tao_compiler/mlir/disc/transforms/fusion_utils.h
index 61a9af0d80f..d89cf7cf5bb 100644
--- a/tao_compiler/mlir/disc/transforms/fusion_utils.h
+++ b/tao_compiler/mlir/disc/transforms/fusion_utils.h
@@ -17,6 +17,8 @@ limitations under the License.
 #define TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_TRANSFORMS_FUSION_UTILS_H_
 
 #include <memory>
+#include <set>
+#include <string>
 #include <vector>
 
 #include "llvm/ADT/EquivalenceClasses.h"
@@ -78,6 +80,7 @@ bool isOnGpu(Operation* op);
 constexpr const char* kDiscFusionTypeAttrName = "disc.fusion_type";
 constexpr StringRef kFusionOpNameAttr = "disc.fusion.name";
 constexpr StringRef kFusionOpTagAttr = "disc.fusion.tag";
+constexpr const char* kFusionTagSeparator = "X";
 
 // kLoop fusion template satisfies:
 //   - all ops in the fusion pattern are element-wise.
@@ -426,9 +429,19 @@ void setFusionName(OpBuilder& b, lmhlo::FusionOp op, StringRef name);
 // Here different tags is mapping to different variants of the fusion op.
 void addFusionTag(OpBuilder& b, lmhlo::FusionOp op, StringRef tag);
 
+// Merge the tags of the op and `tagSet`, assign the new tag set to `op`.
+void mergeFusionTag(OpBuilder& b, lmhlo::FusionOp op,
+                    const std::set<std::string>& tagSet);
+
 // Returns the tag string attached to the fusion op.
 StringRef getFusionTagStr(lmhlo::FusionOp op);
 
+// Returns the unified tag string for the whole tagSet.
+std::string fusionTagSetToStr(const std::set<std::string>& tagSet);
+
+// Returns the parsed tag set from `tagStr`.
+std::set<std::string> parsefusionTagSetFromStr(StringRef tagStr);
+
 // Returns the full name of the fusion op
 // Here full name is composed of the name and tag of the fusion op.
 std::string getFusionFullName(lmhlo::FusionOp op);


Test patch:
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/data/default_schedule_matmul_nn_p_3072x768_f32.mlir b/tao_compiler/mlir/disc/tests/disc-transform/data/default_schedule_matmul_nn_p_3072x768_f32.mlir
new file mode 100644
index 00000000000..7007357436a
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/data/default_schedule_matmul_nn_p_3072x768_f32.mlir
@@ -0,0 +1,10 @@
+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<3072x768xf32>} : () -> tensor<3072x768xf32>
+      %1:2 = tf_executor.island wraps "tf.MatMul"(%arg0, %0) {transpose_a = false, transpose_b = false} : (tensor<?x?xf32>, tensor<3072x768xf32>) -> (tensor<?x?xf32>)
+      tf_executor.fetch %1 : tensor<?x?xf32>
+    }
+    return %graph : tensor<?x?xf32>
+  }
+}
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/default_schedule_matmul.cc b/tao_compiler/mlir/disc/tests/disc-transform/default_schedule_matmul.cc
index 0d34bf98501..a99a9be1e84 100644
--- a/tao_compiler/mlir/disc/tests/disc-transform/default_schedule_matmul.cc
+++ b/tao_compiler/mlir/disc/tests/disc-transform/default_schedule_matmul.cc
@@ -237,4 +237,22 @@ TEST(PackedMatmul, F32_768x2_Using_Default_Schedule) {
       /*profiling*/ true));
 }
 
+TEST(PackedMatmul, F32_768x3072_Using_Default_Schedule) {
+  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 +
+          "default_schedule_matmul_nn_p_3072x768_f32.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 1,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"24x3072xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
 }  // namespace mlir_test