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:

Fix patch:
diff --git a/tao_compiler/decoupling/mlir_compiler.cc b/tao_compiler/decoupling/mlir_compiler.cc
index 792da4298a7..efc700a562d 100644
--- a/tao_compiler/decoupling/mlir_compiler.cc
+++ b/tao_compiler/decoupling/mlir_compiler.cc
@@ -347,7 +347,9 @@ Status CompilerMLIR::ConvertToMlir(const TaoCompilerInput& input,
   std::string graph_def_str = graph_def.SerializeAsString();
   VLOG(2) << "GraphDef str:\n" << graph_def.DebugString();
 
-  context_.reset(new mlir::MLIRContext);
+  mlir::DialectRegistry registry;
+  mlir::registerAllDialects(registry);
+  context_.reset(new mlir::MLIRContext(registry));
   auto& context = *context_;
   GraphDebugInfo debug_info;
   GraphImportConfig specs;
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/transform_dialect_interpreter.cc b/tao_compiler/mlir/disc/tools/disc-transform/transforms/transform_dialect_interpreter.cc
index ebfd463ea69..69b25f407d7 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/transforms/transform_dialect_interpreter.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/transform_dialect_interpreter.cc
@@ -118,9 +118,18 @@ void DiscTransformDialectInterpreterPass::runOnOperation() {
   if (transformFileName_.empty()) {
     llvm::errs() << "no transform file name specified, assuming the transform "
                     "module is embedded in the IR next to the top-level\n";
-    // parse transform ops from the module itself.
+    // Parse transform ops from the module itself.
+
+    // Note that We cloned the transform IR into a standalone module to avoid
+    // it's being modified during transformation.
+    auto transformModule = ModuleOp::create(module->getLoc());
+    auto b = OpBuilder::atBlockBegin(transformModule.getBody());
+    for (auto op : module.getBody()->getOps<transform::TransformOpInterface>())
+      b.clone(*op);
+
+    // Apply the transform IR.
     for (auto op :
-         module.getBody()->getOps<transform::TransformOpInterface>()) {
+         transformModule.getBody()->getOps<transform::TransformOpInterface>()) {
       if (failed(transform::applyTransforms(
               module, op,
               transform::TransformOptions().enableExpensiveChecks(
diff --git a/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc b/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc
index dd0ee704ba5..ee776bac657 100644
--- a/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc
@@ -234,6 +234,38 @@ LogicalResult aarch64GEMMDefaultScheduleFactory(PatternDescription& pd,
   b.setInsertionPointToStart(&bodyBlock);
   Value variant = bodyBlock.getArgument(0);
 
+  auto& fusionPattern = pd.getFusionPattern();
+  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);
+  if (lhsTranspose || rhsTranspose) {
+    return m->emitError() << "not support transposed gemm a.t.m.\n";
+  }
+  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];
+
   // %fill = transform.structured.match ops{["linalg.fill"]} in %variant
   Value fill = buildMatchOp(b, loc, variant, {"linalg.fill"});
   // %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1
@@ -248,16 +280,23 @@ LogicalResult aarch64GEMMDefaultScheduleFactory(PatternDescription& pd,
   auto fuseIntoContainingOp =
       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;
+
   // first level tile and fuse matmul and fill op.
-  auto fuseOp0 = buildFuseOp(b, loc, tiledMatmul, {288, 48, 0}, {0, 1, 2});
+  auto fuseOp0 = buildFuseOp(b, loc, tiledMatmul, {M0, N0, 0}, {0, 1, 2});
 
   // second level tile and fuse matmul and fill op.
   auto fuseOp1 =
-      buildFuseOp(b, loc, fuseOp0->getResult(0), {6, 16, 0}, {0, 1, 2});
+      buildFuseOp(b, loc, fuseOp0->getResult(0), {M1, N1, 0}, {0, 1, 2});
 
   // gemm reduction axis tiling
   auto tileOp =
-      buildTileOp(b, loc, fuseOp1->getResult(0), {0, 0, 1}, {0, 1, 2});
+      buildTileOp(b, loc, fuseOp1->getResult(0), {0, 0, K0}, {0, 1, 2});
 
   variant = buildRunCanonicalizer(b, loc, variant);
 
@@ -272,12 +311,65 @@ LogicalResult aarch64GEMMDefaultScheduleFactory(PatternDescription& pd,
 
   Value padForInput = buildGetProducerOfOperand(b, loc, padOp, 0);
   Value padForWeight = buildGetProducerOfOperand(b, loc, padOp, 1);
-  forEachThreadLoop = buildMatchOp(b, loc, variant, {"scf.foreach_thread"});
-  auto outterLoopForN = buildGetParentForOp(b, loc, padForInput, 4);
-  buildCacheRead(b, loc, padForWeight, forEachThreadLoop, {1, 1}, {1, 16}, true,
-                 {2, 0, 1, 3});
-  buildCacheRead(b, loc, padForInput, outterLoopForN, {1, 1}, {6, 1}, true,
-                 {0, 2, 3, 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;
+    buildCacheRead(b, loc, padForInput, loopN0, {1, 1}, {6, 1}, inputIsPadded,
+                   {0, 2, 3, 1});
+  }
+
+  // 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"});
+    buildCacheRead(b, loc, padForWeight, forEachThreadLoop, {1, 1}, {1, 16},
+                   weightIsPadded, {2, 0, 1, 3});
+  }
 
   variant = buildRunCanonicalizer(b, loc, variant);
 
diff --git a/tao_compiler/mlir/disc/transforms/disc_transform_schedule.h b/tao_compiler/mlir/disc/transforms/disc_transform_schedule.h
index 1b05d6c9795..0d601b8e6f2 100644
--- a/tao_compiler/mlir/disc/transforms/disc_transform_schedule.h
+++ b/tao_compiler/mlir/disc/transforms/disc_transform_schedule.h
@@ -52,6 +52,15 @@ class PatternDescription {
   // Returns the full pattern kind str + tag str.
   std::string getTaggedPatternStr() const;
 
+  // Returns the fusion op this descriptor holds.
+  lmhlo::FusionOp getFusionOp() { return op_; }
+
+  // Returns the fusion pattern this descriptor holds.
+  FusionPattern& getFusionPattern() { return fusionPattern_; }
+
+  // Returns the fusion pattern this descriptor holds.
+  ShapeAnalysis& shapeAnalysis() { return shapeAnalysis_; }
+
  private:
   lmhlo::FusionOp op_;
   FusionPattern& fusionPattern_;


Test patch:
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/data/default_schedule_matmul_nn_d_f32.mlir b/tao_compiler/mlir/disc/tests/disc-transform/data/default_schedule_matmul_nn_d_f32.mlir
new file mode 100644
index 00000000000..c00b1704361
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/data/default_schedule_matmul_nn_d_f32.mlir
@@ -0,0 +1,9 @@
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 0 : i32}} {
+  func.func @main(%arg0: tensor<?x?xf32>, %arg1: 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.MatMul"(%arg0, %arg1) {transpose_a = false, transpose_b = false} : (tensor<?x?xf32>, tensor<?x?xf32>) -> (tensor<?x?xf32>)
+      tf_executor.fetch %0 : 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/default_schedule_matmul_nn_p_2x768_f32.mlir b/tao_compiler/mlir/disc/tests/disc-transform/data/default_schedule_matmul_nn_p_2x768_f32.mlir
new file mode 100644
index 00000000000..49e6d985cd6
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/data/default_schedule_matmul_nn_p_2x768_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<2x768xf32>} : () -> tensor<2x768xf32>
+      %1:2 = tf_executor.island wraps "tf.MatMul"(%arg0, %0) {transpose_a = false, transpose_b = false} : (tensor<?x?xf32>, tensor<2x768xf32>) -> (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/data/default_schedule_matmul_nn_p_512x1024_f32.mlir b/tao_compiler/mlir/disc/tests/disc-transform/data/default_schedule_matmul_nn_p_512x1024_f32.mlir
new file mode 100644
index 00000000000..9de4365ac7e
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/data/default_schedule_matmul_nn_p_512x1024_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<512x1024xf32>} : () -> tensor<512x1024xf32>
+      %1:2 = tf_executor.island wraps "tf.MatMul"(%arg0, %0) {transpose_a = false, transpose_b = false} : (tensor<?x?xf32>, tensor<512x1024xf32>) -> (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/data/default_schedule_matmul_nn_p_768x2_f32.mlir b/tao_compiler/mlir/disc/tests/disc-transform/data/default_schedule_matmul_nn_p_768x2_f32.mlir
new file mode 100644
index 00000000000..02e110d8d76
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/data/default_schedule_matmul_nn_p_768x2_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<768x2xf32>} : () -> tensor<768x2xf32>
+      %1:2 = tf_executor.island wraps "tf.MatMul"(%arg0, %0) {transpose_a = false, transpose_b = false} : (tensor<?x?xf32>, tensor<768x2xf32>) -> (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/data/default_schedule_matmul_nn_s_24x24x64_f32.mlir b/tao_compiler/mlir/disc/tests/disc-transform/data/default_schedule_matmul_nn_s_24x24x64_f32.mlir
new file mode 100644
index 00000000000..0ba066912fe
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/data/default_schedule_matmul_nn_s_24x24x64_f32.mlir
@@ -0,0 +1,9 @@
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 0 : i32}} {
+  func.func @main(%arg0: tensor<24x64xf32>, %arg1: tensor<64x24xf32>) -> (tensor<24x24xf32>) 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.MatMul"(%arg0, %arg1) {transpose_a = false, transpose_b = false} : (tensor<24x64xf32>, tensor<64x24xf32>) -> (tensor<24x24xf32>)
+      tf_executor.fetch %0 : tensor<24x24xf32>
+    }
+    return %graph : tensor<24x24xf32>
+  }
+}
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/data/default_schedule_matmul_nn_s_24x64x24_f32.mlir b/tao_compiler/mlir/disc/tests/disc-transform/data/default_schedule_matmul_nn_s_24x64x24_f32.mlir
new file mode 100644
index 00000000000..fbbf061342d
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/data/default_schedule_matmul_nn_s_24x64x24_f32.mlir
@@ -0,0 +1,9 @@
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 0 : i32}} {
+  func.func @main(%arg0: tensor<24x24xf32>, %arg1: tensor<24x64xf32>) -> (tensor<24x64xf32>) 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.MatMul"(%arg0, %arg1) {transpose_a = false, transpose_b = false} : (tensor<24x24xf32>, tensor<24x64xf32>) -> (tensor<24x64xf32>)
+      tf_executor.fetch %0 : tensor<24x64xf32>
+    }
+    return %graph : tensor<24x64xf32>
+  }
+}
\ 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
new file mode 100644
index 00000000000..c9606e82409
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/default_schedule_matmul.cc
@@ -0,0 +1,188 @@
+/* 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 "tensorflow/compiler/mlir/disc/tests/mlir_feature_test.h"
+#include "tensorflow/compiler/mlir/disc/tests/mlir_test.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace mlir_test {
+
+const std::string c_ft_path =
+    "tensorflow/compiler/mlir/disc/tests/disc-transform/data/";
+
+static bool init_threads = []() {
+  setenv("OMP_NUM_THREADS", "1", 1);
+  return true;
+}();
+
+TEST(DynamicShapeMatmul, F32_24x24x64_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_d_f32.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 2,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"24x64xf32_X", "64x24xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
+TEST(DynamicShapeMatmul, F32_24x768x768_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_d_f32.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 2,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"24x768xf32_X", "768x768xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
+TEST(DynamicShapeMatmul, F32_24x3072x768_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_d_f32.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 2,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"24x768xf32_X", "768x3072xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
+TEST(DynamicShapeMatmul, F32_24x768x3072_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_d_f32.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 2,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"24x3072xf32_X", "3072x768xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
+TEST(StaticShapeMatmul, F32_24x24x64_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_s_24x24x64_f32.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 2,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"24x64xf32_X", "64x24xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
+TEST(StaticShapeMatmul, F32_24x64x24_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_s_24x64x24_f32.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 2,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"24x24xf32_X", "24x64xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
+TEST(PackedMatmul, F32_304x1024x512_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_512x1024_f32.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 1,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"304x512xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
+TEST(PackedMatmul, F32_2x768_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_2x768_f32.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 1,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"304x2xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
+TEST(PackedMatmul, F32_768x2_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_768x2_f32.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 1,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"304x768xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
+}  // namespace mlir_test
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/packed_matmul.cc b/tao_compiler/mlir/disc/tests/disc-transform/packed_matmul.cc
index 5cf780a866c..5af536ac586 100644
--- a/tao_compiler/mlir/disc/tests/disc-transform/packed_matmul.cc
+++ b/tao_compiler/mlir/disc/tests/disc-transform/packed_matmul.cc
@@ -48,21 +48,4 @@ TEST(PackedMatmul, F32_304x1024x512) {
       /*profiling*/ true));
 }
 
-TEST(PackedMatmul, F32_304x1024x512_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 + "packed_matmul_nn_p_512x1024_f32.mlir",
-      /*backend_types*/ {BackendType::kAArch64},
-      /*num_inputs*/ 1,
-      /*num_outputs*/ 1,
-      /*input_descriptors*/ {"304x512xf32_X"},
-      /*output_descriptors*/ {"f32_X"},
-      /*input_vals*/ {},
-      /*expected_output_vals*/ {},
-      /*profiling*/ true));
-}
-
 }  // namespace mlir_test