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/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
index 09b0b4f3ddd..c8fce07709d 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
@@ -594,6 +594,8 @@ struct InnerMostDimsInfo {
   SmallVector<int64_t> dims;
   // the dimension sizes for the inner most dimensions.
   SmallVector<int64_t> dimSizes;
+  // dimSizes after being transposed
+  SmallVector<int64_t> transposedDimSizes;
   // permutation for the inner most dimensions.
   SmallVector<int64_t> permutation;
   // reverse permutation for the inner most dimensions.
@@ -628,6 +630,9 @@ void parseInnerMostDimsInfo(InnerMostDimsInfo& info, ShapedType sourceTy,
   for (int i = 0; i < tileLevelsVec.size(); ++i) {
     info.reversePermutation[info.permutation[i]] = i;
   }
+
+  info.transposedDimSizes =
+      disc_linalg_ext::interchange<int64_t>(info.dimSizes, info.permutation);
 }
 
 linalg::GenericOp makeTransposeOp(OpBuilder& b, Location loc, Value inputTensor,
@@ -718,9 +723,8 @@ LogicalResult readFromPackedValue(
         disc_linalg_ext::interchange<OpFoldResult>(newSizes, permutationVec);
   }
 
-  auto newShape = info.dimSizes;
-  newShape = disc_linalg_ext::interchange<int64_t>(newShape, info.permutation);
-  auto sliceTy = RankedTensorType::get(newShape, packedTy.getElementType());
+  auto sliceTy =
+      RankedTensorType::get(info.transposedDimSizes, packedTy.getElementType());
   auto sliceOp = b.create<tensor::ExtractSliceOp>(
       loc, sliceTy, packedValue, newOffsets, newSizes, newStrides);
   if (sliceTy != targetTy) {
@@ -880,6 +884,147 @@ DiagnosedSilenceableFailure CacheReadOp::apply(
   return DiagnosedSilenceableFailure(success());
 }
 
+//===---------------------------------------------------------------------===//
+// LowerMultiLevelPackToLoopOp
+//===---------------------------------------------------------------------===//
+
+namespace {
+
+SmallVector<Value> getDimValues(OpBuilder& b, Location loc, Value v) {
+  auto ty = v.getType().cast<RankedTensorType>();
+  SmallVector<Value> vs;
+  for (int64_t i = 0; i < ty.getRank(); ++i)
+    vs.push_back(b.create<tensor::DimOp>(loc, v, i));
+  return vs;
+}
+
+Value buildMin(OpBuilder& b, Location loc, Value lhs, Value rhs) {
+  SmallVector<OpFoldResult> vals = {lhs, rhs};
+  auto result = makeComposedFoldedAffineMin(
+      b, loc, AffineMap::getMultiDimIdentityMap(vals.size(), loc.getContext()),
+      vals);
+  return getValueOrCreateConstantIndexOp(b, loc, result);
+}
+
+}  // namespace
+
+void LowerMultiLevelPackToLoopOp::build(OpBuilder& builder,
+                                        OperationState& result, Value target) {
+  MLIRContext* ctx = builder.getContext();
+  result.addOperands(target);
+  result.addTypes({pdl::OperationType::get(ctx)});
+}
+
+DiagnosedSilenceableFailure LowerMultiLevelPackToLoopOp::applyToOne(
+    Operation* target, SmallVectorImpl<Operation*>& results,
+    transform::TransformState& state) {
+  auto multiLevelPackOp = dyn_cast<MultiLevelPackOp>(target);
+  if (!multiLevelPackOp) {
+    return mlir::emitDefiniteFailure(
+        target, "applies only to disc_linalg_ext::MultiLevelPackOp");
+  }
+
+  OpBuilder b(target);
+  Location loc = target->getLoc();
+  MLIRContext* ctx = target->getContext();
+  auto tileLevelsVec = multiLevelPackOp.getTileLevelsVec();
+  auto tileSizesVec = multiLevelPackOp.getTileSizesVec();
+  auto permutationVec = multiLevelPackOp.getPermutationVec();
+
+  Value src = multiLevelPackOp.getInput();
+  auto srcTy = src.getType().cast<RankedTensorType>();
+  int64_t srcRank = srcTy.getRank();
+  Value dst = multiLevelPackOp.getOutput();
+  auto dstTy = dst.getType().cast<RankedTensorType>();
+  int64_t dstRank = dstTy.getRank();
+
+  Value zero = b.create<arith::ConstantIndexOp>(loc, 0);
+  Value one = b.create<arith::ConstantIndexOp>(loc, 1);
+  SmallVector<Value> srcOffsets(srcRank, zero);
+  SmallVector<Value> srcSizes = getDimValues(b, loc, src);
+  SmallVector<Value> srcStrides(srcRank, one);
+  SmallVector<Value> dstOffsets(dstRank, zero);
+  SmallVector<Value> dstSizes = getDimValues(b, loc, dst);
+  SmallVector<Value> dstStrides(dstRank, one);
+
+  InnerMostDimsInfo innerMostDimsInfo;
+  parseInnerMostDimsInfo(innerMostDimsInfo, srcTy, tileLevelsVec, tileSizesVec,
+                         permutationVec);
+  auto logicalDim2SrcDim =
+      multiLevelPackOp.getOutputLogicalDimToInputDimMapping(tileLevelsVec,
+                                                            tileSizesVec);
+  auto logicalDim2TileSize =
+      multiLevelPackOp.getOutputLogicalDimToTileSizeMapping(tileLevelsVec,
+                                                            tileSizesVec);
+
+  Value loopInitValue = dst;
+  SmallVector<scf::ForOp> forOps;
+  SmallVector<Value> srcDimUppers = srcSizes;
+  for (int dstIdx = 0; dstIdx < dstRank; ++dstIdx) {
+    int logicalIdx = permutationVec[dstIdx];
+    if (innerMostDimsInfo.isInnerMostDim(logicalIdx)) continue;
+    Value step =
+        b.create<arith::ConstantIndexOp>(loc, logicalDim2TileSize[logicalIdx]);
+    int srcIdx = logicalDim2SrcDim[logicalIdx];
+    Value upper = srcDimUppers[srcIdx];
+    auto forOp =
+        b.create<scf::ForOp>(loc, zero, upper, step, ValueRange{loopInitValue});
+    b.setInsertionPoint(forOp.getBody(), forOp.getBody()->begin());
+    Value iv = forOp.getInductionVar();
+    srcOffsets[srcIdx] = b.create<arith::AddIOp>(loc, srcOffsets[srcIdx], iv);
+    Value remaining = b.create<arith::SubIOp>(loc, upper, iv);
+    srcDimUppers[srcIdx] = buildMin(b, loc, remaining, step);
+    dstOffsets[dstIdx] = b.create<arith::DivSIOp>(loc, iv, step);
+    dstSizes[dstIdx] = one;
+    forOps.push_back(forOp);
+    loopInitValue = *forOp.getRegionIterArgs().begin();
+  }
+
+  Value srcSlice = b.create<tensor::ExtractSliceOp>(loc, src, srcOffsets,
+                                                    srcDimUppers, srcStrides);
+  if (Value paddingValue = multiLevelPackOp.getPaddingValue()) {
+    SmallVector<OpFoldResult> lowIndices(srcRank,
+                                         IntegerAttr::get(b.getIndexType(), 0));
+    SmallVector<OpFoldResult> highIndices;
+    for (const auto& z : llvm::zip(srcDimUppers, innerMostDimsInfo.dimSizes)) {
+      Value paddedSize = b.create<arith::ConstantIndexOp>(loc, std::get<1>(z));
+      highIndices.push_back(
+          b.create<arith::SubIOp>(loc, paddedSize, std::get<0>(z)).getResult());
+    }
+    srcSlice = b.create<tensor::PadOp>(
+        loc,
+        RankedTensorType::get(innerMostDimsInfo.dimSizes,
+                              srcTy.getElementType()),
+        srcSlice, lowIndices, highIndices, paddingValue);
+  }
+
+  // transpose srcSlice when needed.
+  auto sortedPermutation = innerMostDimsInfo.permutation;
+  llvm::sort(sortedPermutation);
+  if (sortedPermutation != innerMostDimsInfo.permutation) {
+    Value transposeDst = b.create<tensor::EmptyOp>(
+        loc, innerMostDimsInfo.transposedDimSizes, dstTy.getElementType());
+    srcSlice = makeTransposeOp(b, loc, srcSlice, transposeDst,
+                               innerMostDimsInfo.permutation)
+                   ->getResult(0);
+  }
+
+  Value updateDst =
+      b.create<tensor::InsertSliceOp>(loc, srcSlice, loopInitValue, dstOffsets,
+                                      dstSizes, dstStrides)
+          ->getResult(0);
+  b.create<scf::YieldOp>(loc, updateDst);
+  for (int i = static_cast<int>(forOps.size()) - 2; i >= 0; --i) {
+    b.setInsertionPointAfter(forOps[i + 1]);
+    b.create<scf::YieldOp>(loc, forOps[i + 1]->getResult(0));
+  }
+  assert(forOps.size() > 0);
+  target->getResult(0).replaceAllUsesWith(forOps[0]->getResult(0));
+
+  results.assign({forOps[0]});
+  return DiagnosedSilenceableFailure(success());
+}
+
 }  // namespace transform_dialect
 
 void registerTransformDialectCommonExtension(DialectRegistry& registry) {
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
index a0584345112..95de7988c5b 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
@@ -234,4 +234,33 @@ def CacheReadOp : Op<Transform_Dialect, "disc.cache_read",
   ];
 }
 
+def LowerMultiLevelPackToLoopOp : Op<Transform_Dialect, "disc.lower_multi_level_pack_to_loop",
+    [FunctionalStyleTransformOpTrait,
+     MemoryEffectsOpInterface,
+     TransformEachOpTrait,
+     TransformOpInterface]> {
+  let description = [{
+    Replace a MultiLevelPackOp to its loop level equivalent.
+
+    Returns a handle to the outter most loop.
+  }];
+
+  let arguments = (ins PDL_Operation:$target);
+  let results = (outs PDL_Operation:$result);
+
+  let assemblyFormat = "$target attr-dict";
+  let cppNamespace = "::mlir::disc_ral::transform_dialect";
+
+  let builders = [
+    OpBuilder<(ins "Value":$target)>
+  ];
+
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::Operation *target,
+        ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+        ::mlir::transform::TransformState &state);
+  }];
+}
+
 #endif // DISC_TRANSFORM_OPS_EXT


Test patch:
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_nn_d_f32_large_schedule_4.mlir b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_nn_d_f32_large_schedule_4.mlir
new file mode 100644
index 00000000000..0349a84daf4
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_nn_d_f32_large_schedule_4.mlir
@@ -0,0 +1,68 @@
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %fill = transform.structured.match ops{["linalg.fill"]} in %arg1
+  %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1
+
+  %0:2 = transform.structured.tile_to_foreach_thread_op %matmul num_threads [1, 1]
+  transform.structured.fuse_into_containing_op %fill into %0#0
+
+  // first level tile and fuse matmul and fill op.
+  %1:3 = transform.structured.fuse %0#1 {tile_sizes = [288, 256, 0], tile_interchange = [0, 1, 2]}
+  // second level tile and fuse matmul and fill op.
+  %2:3 = transform.structured.fuse %1#0 {tile_sizes = [6, 16, 0], tile_interchange = [0, 1, 2]}
+
+  // gemm reduction axis tiling
+  %3:2 = transform.structured.tile %2#0 [0, 0, 1] {interchange=[0, 1, 2]}
+
+  // clean up
+  %func0 = transform.structured.match ops{["func.func"]} in %arg1
+  transform.disc.apply_patterns %func0 {canonicalization}
+  // fold two extract_slice ops generated by two-level tiling. It's needed to enable following
+  // pad and hosit schedule.
+  %weight_inner_slice = get_producer_of_operand %3#0[1] : (!pdl.operation) -> !pdl.operation
+  transform.disc.fold_producer_extract_slice %weight_inner_slice {max_repeat_num = 2}
+
+  // pad to match the requirement of hardware vector/tensor instruction.
+  %4 = transform.structured.pad %3#0 {padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32], padding_dimensions=[0, 1, 2], pack_paddings=[1, 1, 0], hoist_paddings=[4, 0, 0], transpose_paddings=[[1, 0], [0, 1], [0, 1]]}
+
+  %pad_for_input = get_producer_of_operand %4[0] : (!pdl.operation) -> !pdl.operation
+  %pad_for_weight = get_producer_of_operand %4[1] : (!pdl.operation) -> !pdl.operation
+  %foreach_op = transform.structured.match ops{["scf.foreach_thread"]} in %arg1
+  transform.disc.cache_read {padded} %pad_for_weight at %foreach_op with tile_levels = [1, 1] tile_sizes = [1, 16] permutation = [2, 0, 1, 3]
+
+  %pack_op = transform.structured.match ops{["disc_linalg_ext.multi_level_pack"]} in %arg1
+  transform.disc.lower_multi_level_pack_to_loop %pack_op
+
+  %func1 = transform.structured.match ops{["func.func"]} in %arg1
+  transform.disc.apply_patterns %func1 {canonicalization}
+
+  %func2 = transform.structured.match ops{["func.func"]} in %arg1
+  transform.structured.vectorize %func2 {vectorize_padding}
+
+  %func3 = transform.structured.match ops{["func.func"]} in %arg1
+  transform.disc.apply_patterns %func3 {canonicalization}
+
+  transform.disc.bufferize %arg1
+
+  transform.lower_vectors {
+    contraction_lowering = "outerproduct",
+    multireduction_lowering = "innerparallel",
+    split_transfers = "linalg-copy",
+    // stages = [0, 1, 2, 3, 4, 5, 6, 7],
+    stages = [0, 1, 2, 3],
+    transpose_avx2_lowering = false,
+    transpose_lowering = "eltwise",
+    unroll_vector_transfers = true
+  }
+
+  transform.lower_vectors {
+    contraction_lowering = "outerproduct",
+    multireduction_lowering = "innerparallel",
+    split_transfers = "linalg-copy",
+    // stages = [0, 1, 2, 3, 4, 5, 6, 7],
+    stages = [5, 6, 7],
+    transpose_avx2_lowering = false,
+    transpose_lowering = "eltwise",
+    unroll_vector_transfers = true
+  }
+}
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_nn_d_f32_large_schedule_5.mlir b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_nn_d_f32_large_schedule_5.mlir
new file mode 100644
index 00000000000..8a311c9a374
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_nn_d_f32_large_schedule_5.mlir
@@ -0,0 +1,70 @@
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %fill = transform.structured.match ops{["linalg.fill"]} in %arg1
+  %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1
+
+  %0:2 = transform.structured.tile_to_foreach_thread_op %matmul num_threads [1, 1]
+  transform.structured.fuse_into_containing_op %fill into %0#0
+
+  // first level tile and fuse matmul and fill op.
+  %1:3 = transform.structured.fuse %0#1 {tile_sizes = [288, 256, 0], tile_interchange = [0, 1, 2]}
+  // second level tile and fuse matmul and fill op.
+  %2:3 = transform.structured.fuse %1#0 {tile_sizes = [6, 16, 0], tile_interchange = [0, 1, 2]}
+
+  // gemm reduction axis tiling
+  %3:2 = transform.structured.tile %2#0 [0, 0, 1] {interchange=[0, 1, 2]}
+
+  // clean up
+  %func0 = transform.structured.match ops{["func.func"]} in %arg1
+  transform.disc.apply_patterns %func0 {canonicalization}
+  // fold two extract_slice ops generated by two-level tiling. It's needed to enable following
+  // pad and hosit schedule.
+  %weight_inner_slice = get_producer_of_operand %3#0[1] : (!pdl.operation) -> !pdl.operation
+  transform.disc.fold_producer_extract_slice %weight_inner_slice {max_repeat_num = 2}
+
+  // pad to match the requirement of hardware vector/tensor instruction.
+  %4 = transform.structured.pad %3#0 {padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32], padding_dimensions=[0, 1, 2], pack_paddings=[1, 1, 0], hoist_paddings=[0, 0, 0], transpose_paddings=[[1, 0], [0, 1], [0, 1]]}
+
+  %pad_for_input = get_producer_of_operand %4[0] : (!pdl.operation) -> !pdl.operation
+  %pad_for_weight = get_producer_of_operand %4[1] : (!pdl.operation) -> !pdl.operation
+  %foreach_op = transform.structured.match ops{["scf.foreach_thread"]} in %arg1
+  %loop_for_outter_most_n = transform.loop.get_parent_for %4 { num_loops = 4} : (!pdl.operation) -> !pdl.operation
+  transform.disc.cache_read {padded} %pad_for_input at %loop_for_outter_most_n with tile_levels = [1, 1] tile_sizes = [6, 1] permutation = [0, 2, 3, 1]
+  transform.disc.cache_read {padded} %pad_for_weight at %foreach_op with tile_levels = [1, 1] tile_sizes = [1, 16] permutation = [2, 0, 1, 3]
+
+  %pack_op = transform.structured.match ops{["disc_linalg_ext.multi_level_pack"]} in %arg1
+  transform.disc.lower_multi_level_pack_to_loop %pack_op
+
+  %func1 = transform.structured.match ops{["func.func"]} in %arg1
+  transform.disc.apply_patterns %func1 {canonicalization}
+
+  %func2 = transform.structured.match ops{["func.func"]} in %arg1
+  transform.structured.vectorize %func2 {vectorize_padding}
+
+  %func3 = transform.structured.match ops{["func.func"]} in %arg1
+  transform.disc.apply_patterns %func3 {canonicalization}
+
+  transform.disc.bufferize %arg1
+
+  transform.lower_vectors {
+    contraction_lowering = "outerproduct",
+    multireduction_lowering = "innerparallel",
+    split_transfers = "linalg-copy",
+    // stages = [0, 1, 2, 3, 4, 5, 6, 7],
+    stages = [0, 1, 2, 3],
+    transpose_avx2_lowering = false,
+    transpose_lowering = "eltwise",
+    unroll_vector_transfers = true
+  }
+
+  transform.lower_vectors {
+    contraction_lowering = "outerproduct",
+    multireduction_lowering = "innerparallel",
+    split_transfers = "linalg-copy",
+    // stages = [0, 1, 2, 3, 4, 5, 6, 7],
+    stages = [5, 6, 7],
+    transpose_avx2_lowering = false,
+    transpose_lowering = "eltwise",
+    unroll_vector_transfers = true
+  }
+}
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/matmul.cc b/tao_compiler/mlir/disc/tests/disc-transform/matmul.cc
index 536afb7c1a9..2183fcd0c7f 100644
--- a/tao_compiler/mlir/disc/tests/disc-transform/matmul.cc
+++ b/tao_compiler/mlir/disc/tests/disc-transform/matmul.cc
@@ -60,8 +60,8 @@ TEST(SimpleTest, MatMulF32_304x1024x256) {
       {"DISC_TRANSFORM_SCHEDULE_FILE",
        {c_ft_path + "matmul_nn_d_f32_large_schedule.mlir", false}},
       {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
-      {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", "false"}},
-      {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", "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_nn_d_f32.mlir",
@@ -80,8 +80,8 @@ TEST(SimpleTest, MatMulF32_1024x1024x1024) {
       {"DISC_TRANSFORM_SCHEDULE_FILE",
        {c_ft_path + "matmul_nn_d_f32_large_schedule.mlir", false}},
       {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
-      {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", "false"}},
-      {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", "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_nn_d_f32.mlir",
@@ -100,8 +100,8 @@ TEST(SimpleTest, MatMulF32_304x1024x256_2) {
       {"DISC_TRANSFORM_SCHEDULE_FILE",
        {c_ft_path + "matmul_nn_d_f32_large_schedule_2.mlir", false}},
       {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
-      {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", "false"}},
-      {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", "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_nn_d_f32.mlir",
@@ -120,8 +120,8 @@ TEST(SimpleTest, MatMulF32_1024x1024x1024_2) {
       {"DISC_TRANSFORM_SCHEDULE_FILE",
        {c_ft_path + "matmul_nn_d_f32_large_schedule_2.mlir", false}},
       {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
-      {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", "false"}},
-      {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", "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_nn_d_f32.mlir",
@@ -140,8 +140,8 @@ TEST(SimpleTest, MatMulF32_304x256x256_3) {
       {"DISC_TRANSFORM_SCHEDULE_FILE",
        {c_ft_path + "matmul_nn_d_f32_large_schedule_3.mlir", false}},
       {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
-      {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", "false"}},
-      {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", "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_nn_d_f32.mlir",
@@ -160,8 +160,8 @@ TEST(SimpleTest, MatMulF32_304x512x256_3) {
       {"DISC_TRANSFORM_SCHEDULE_FILE",
        {c_ft_path + "matmul_nn_d_f32_large_schedule_3.mlir", false}},
       {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
-      {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", "false"}},
-      {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", "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_nn_d_f32.mlir",
@@ -180,8 +180,8 @@ TEST(SimpleTest, MatMulF32_304x1024x256_3) {
       {"DISC_TRANSFORM_SCHEDULE_FILE",
        {c_ft_path + "matmul_nn_d_f32_large_schedule_3.mlir", false}},
       {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
-      {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", "false"}},
-      {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", "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_nn_d_f32.mlir",
@@ -200,8 +200,8 @@ TEST(SimpleTest, MatMulF32_304x1024x512_3) {
       {"DISC_TRANSFORM_SCHEDULE_FILE",
        {c_ft_path + "matmul_nn_d_f32_large_schedule_3.mlir", false}},
       {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
-      {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", "false"}},
-      {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", "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_nn_d_f32.mlir",
@@ -220,8 +220,8 @@ TEST(SimpleTest, MatMulF32_1024x1024x1024_3) {
       {"DISC_TRANSFORM_SCHEDULE_FILE",
        {c_ft_path + "matmul_nn_d_f32_large_schedule_3.mlir", false}},
       {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
-      {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", "false"}},
-      {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", "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_nn_d_f32.mlir",
@@ -235,4 +235,84 @@ TEST(SimpleTest, MatMulF32_1024x1024x1024_3) {
       /*profiling*/ true));
 }
 
+TEST(SimpleTest, MatMulF32_304x1024x512_4) {
+  EnvSetting setting = {
+      {"DISC_TRANSFORM_SCHEDULE_FILE",
+       {c_ft_path + "matmul_nn_d_f32_large_schedule_4.mlir", false}},
+      {"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_nn_d_f32.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 2,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"304x512xf32_X", "512x1024xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
+TEST(SimpleTest, MatMulF32_1024x1024x1024_4) {
+  EnvSetting setting = {
+      {"DISC_TRANSFORM_SCHEDULE_FILE",
+       {c_ft_path + "matmul_nn_d_f32_large_schedule_4.mlir", false}},
+      {"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_nn_d_f32.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 2,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"1024x1024xf32_X", "1024x1024xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
+TEST(SimpleTest, MatMulF32_1026x1024x1024_4) {
+  EnvSetting setting = {
+      {"DISC_TRANSFORM_SCHEDULE_FILE",
+       {c_ft_path + "matmul_nn_d_f32_large_schedule_4.mlir", false}},
+      {"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_nn_d_f32.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 2,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"1026x1024xf32_X", "1024x1024xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
+TEST(SimpleTest, MatMulF32_304x1024x512_5) {
+  EnvSetting setting = {
+      {"DISC_TRANSFORM_SCHEDULE_FILE",
+       {c_ft_path + "matmul_nn_d_f32_large_schedule_5.mlir", false}},
+      {"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_nn_d_f32.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 2,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"304x512xf32_X", "512x1024xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
 }  // namespace mlir_test
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/lower-multi-level-pack-to-loop.mlir b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/lower-multi-level-pack-to-loop.mlir
new file mode 100644
index 00000000000..fcd652b84a0
--- /dev/null
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/lower-multi-level-pack-to-loop.mlir
@@ -0,0 +1,58 @@
+// RUN: disc-opt --disc-transform-dialect-interpreter -split-input-file %s | FileCheck %s
+
+// CHECK-LABEL: @multi_level_pack
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)
+func.func @multi_level_pack(%arg0: tensor<?x?xf32>, %arg1: index, %arg2: index) -> tensor<?x?x32x32xf32> {
+  // CHECK: %[[INIT:.*]] = tensor.empty(%[[ARG1]], %[[ARG2]]) : tensor<?x?x32x32xf32>
+  // CHECK: %[[DIM0:.*]] = tensor.dim %[[ARG0]], %c0 : tensor<?x?xf32>
+  // CHECK: %[[DIM1:.*]] = tensor.dim %[[ARG0]], %c1 : tensor<?x?xf32>
+  // CHECK: %[[T0:.*]] = scf.for %[[IV0:.*]] = %c0 to %[[DIM0]] step %c32 iter_args(%[[V0:.*]] = %[[INIT]]) -> (tensor<?x?x32x32xf32>) {
+  // CHECK:   %[[T1:.*]] = scf.for %[[IV1:.*]] = %c0 to %[[DIM1]] step %c32 iter_args(%[[V1:.*]] = %[[V0]]) -> (tensor<?x?x32x32xf32>) {
+  // CHECK:     %[[SLICE:.*]] = tensor.extract_slice %[[ARG0]]
+  // CHECK-SAME: %[[IV0]], %[[IV1]]
+  // CHECK-SAME: 32, 32
+  // CHECK-SAME: 1, 1
+  // CHECK:     %[[UPDATE:.*]] = tensor.insert_slice %[[SLICE]] into %[[V1]]
+  // CHECK:     scf.yield %[[UPDATE]] : tensor<?x?x32x32xf32>
+  // CHECK:   scf.yield %[[T1]] : tensor<?x?x32x32xf32>
+  // CHECK: return %[[T0]]
+  %0 = tensor.empty(%arg1, %arg2) : tensor<?x?x32x32xf32>
+  %1 = disc_linalg_ext.multi_level_pack %arg0 with tile_levels = [1, 1] tile_sizes = [32, 32] permutation = [0, 2, 1, 3] into %0 : (tensor<?x?xf32> tensor<?x?x32x32xf32>) -> tensor<?x?x32x32xf32>
+  return %1 : tensor<?x?x32x32xf32>
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %pack_op = transform.structured.match ops{["disc_linalg_ext.multi_level_pack"]} in %arg1
+  transform.disc.lower_multi_level_pack_to_loop %pack_op
+}
+
+// -----
+
+// CHECK-LABEL: @multi_level_pack_with_pad
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)
+func.func @multi_level_pack_with_pad(%arg0: tensor<?x?xf32>, %arg1: index, %arg2: index) -> tensor<?x?x32x32xf32> {
+  %cst = arith.constant 1.0 : f32
+  // CHECK: %[[CST:.*]] = arith.constant 1.000000e+00 : f32
+  // CHECK: %[[INIT:.*]] = tensor.empty(%[[ARG1]], %[[ARG2]]) : tensor<?x?x32x32xf32>
+  // CHECK: %[[DIM0:.*]] = tensor.dim %[[ARG0]], %c0 : tensor<?x?xf32>
+  // CHECK: %[[DIM1:.*]] = tensor.dim %[[ARG0]], %c1 : tensor<?x?xf32>
+  // CHECK: %[[T0:.*]] = scf.for %[[IV0:.*]] = %c0 to %[[DIM0]] step %c32 iter_args(%[[V0:.*]] = %[[INIT]]) -> (tensor<?x?x32x32xf32>) {
+  // CHECK:   %[[T1:.*]] = scf.for %[[IV1:.*]] = %c0 to %[[DIM1]] step %c32 iter_args(%[[V1:.*]] = %[[V0]]) -> (tensor<?x?x32x32xf32>) {
+  // CHECK:     %[[SLICE:.*]] = tensor.extract_slice %[[ARG0]]
+  // CHECK-SAME: %[[IV0]], %[[IV1]]
+  // CHECK:     %[[PAD:.*]] = tensor.pad %[[SLICE]]
+  // CHECK:     %[[UPDATE:.*]] = tensor.insert_slice %[[PAD]] into %[[V1]]
+  // CHECK:     scf.yield %[[UPDATE]] : tensor<?x?x32x32xf32>
+  // CHECK:   scf.yield %[[T1]] : tensor<?x?x32x32xf32>
+  // CHECK: return %[[T0]]
+  %0 = tensor.empty(%arg1, %arg2) : tensor<?x?x32x32xf32>
+  %1 = disc_linalg_ext.multi_level_pack %arg0 with padding_value(%cst : f32) tile_levels = [1, 1] tile_sizes = [32, 32] permutation = [0, 2, 1, 3] into %0 : (tensor<?x?xf32> tensor<?x?x32x32xf32>) -> tensor<?x?x32x32xf32>
+  return %1 : tensor<?x?x32x32xf32>
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %pack_op = transform.structured.match ops{["disc_linalg_ext.multi_level_pack"]} in %arg1
+  transform.disc.lower_multi_level_pack_to_loop %pack_op
+}