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 1e0680aeff6..f95658c0eb4 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
@@ -16,6 +16,7 @@
 #include "iree-dialects/Dialect/LinalgTransform/StructuredTransformOpsExt.h"
 #include "iree-dialects/Transforms/ListenerGreedyPatternRewriteDriver.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/Debug.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Utils/Utils.h"
@@ -28,10 +29,12 @@
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/Dialect/PDL/IR/PDLTypes.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/Interfaces/DestinationStyleOpInterface.h"
 #include "mlir/Pass/PassManager.h"
 #include "tensorflow/compiler/mlir/disc/tools/disc-transform/utils.h"
 
@@ -191,6 +194,312 @@ DiagnosedSilenceableFailure DISCBufferizeOp::apply(
   return DiagnosedSilenceableFailure::success();
 }
 
+//===---------------------------------------------------------------------===//
+// ApplyPatternsOp
+//===---------------------------------------------------------------------===//
+
+namespace {
+
+bool allEquals(ArrayAttr attrs, int val) {
+  return llvm::all_of(attrs, [&](Attribute attr) {
+    return attr.cast<IntegerAttr>().getInt() == val;
+  });
+}
+
+bool allConstIndexValues(ValueRange vals, int val) {
+  return llvm::all_of(vals, [&](Value indexVal) {
+    auto constOp = indexVal.getDefiningOp<arith::ConstantIndexOp>();
+    return constOp && constOp.getValue().cast<IntegerAttr>().getInt() == val;
+  });
+}
+
+bool offsetAllZeros(tensor::ExtractSliceOp slice) {
+  return allEquals(slice.getStaticOffsets(), 0);
+}
+
+bool offsetAllZeros(tensor::InsertSliceOp slice) {
+  return allEquals(slice.getStaticOffsets(), 0);
+}
+
+bool strideAllOnes(tensor::ExtractSliceOp slice) {
+  auto strides = slice.getStaticStrides();
+  return allEquals(strides, 1);
+}
+
+bool strideAllOnes(tensor::InsertSliceOp slice) {
+  return allEquals(slice.getStaticStrides(), 1);
+}
+
+bool hasSameShape(ArrayRef<OpFoldResult> lhs, ArrayRef<OpFoldResult> rhs) {
+  for (const auto& z : llvm::zip(lhs, rhs)) {
+    auto lhsAttr = std::get<0>(z).dyn_cast<Attribute>();
+    auto rhsAttr = std::get<1>(z).dyn_cast<Attribute>();
+    if ((lhsAttr == nullptr) != (rhsAttr == nullptr)) return false;
+
+    if (lhsAttr) {
+      if (lhsAttr.cast<IntegerAttr>().getInt() !=
+          rhsAttr.cast<IntegerAttr>().getInt())
+        return false;
+    } else {
+      if (std::get<0>(z).dyn_cast<Value>() != std::get<0>(z).dyn_cast<Value>())
+        return false;
+    }
+  }
+  return true;
+}
+
+bool hasSameShape(tensor::ExtractSliceOp lhs, tensor::ExtractSliceOp rhs) {
+  return hasSameShape(lhs.getMixedSizes(), rhs.getMixedSizes());
+}
+
+Value mapResultToInitOperandOfDestinationStyleOp(Value result) {
+  auto dstOp = result.getDefiningOp<DestinationStyleOpInterface>();
+  if (!dstOp) return nullptr;
+  int resultNumber = result.cast<OpResult>().getResultNumber();
+  return dstOp.getDpsInitOperand(resultNumber)->get();
+}
+
+tensor::ExtractSliceOp getSliceProducer(Value result) {
+  if (!result) return nullptr;
+  if (auto slice = result.getDefiningOp<tensor::ExtractSliceOp>()) return slice;
+  return getSliceProducer(mapResultToInitOperandOfDestinationStyleOp(result));
+}
+
+/// Pattern to fold tensor.extract_slice of tensor.extract_slice.
+/// convert:
+///   %0 = tensor.extract_slice %arg0[%0, %t1][%s0, %s1][%c1, %c1]
+///   %1 = tensor.extract_slice %0[%t2, %0][%s2, %s3][%c1, %c1]
+/// to:
+///   %1 = tensor.extract_slice %0[%t2, %t1][%s2, %s3][%c1, %c1]
+struct FoldExtractSliceOfExtractSlicePattern
+    : public OpRewritePattern<tensor::ExtractSliceOp> {
+  using OpRewritePattern<tensor::ExtractSliceOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tensor::ExtractSliceOp op,
+                                PatternRewriter& rewriter) const override {
+    auto producerOp = op.getSource().getDefiningOp<tensor::ExtractSliceOp>();
+    if (!strideAllOnes(op) || !producerOp || !strideAllOnes(producerOp))
+      return failure();
+
+    SmallVector<OpFoldResult> newOffsets;
+    for (const auto& z :
+         llvm::zip(op.getMixedOffsets(), producerOp.getMixedOffsets())) {
+      auto lhsAttr = std::get<0>(z).dyn_cast<Attribute>();
+      auto rhsAttr = std::get<1>(z).dyn_cast<Attribute>();
+      if (!lhsAttr && !rhsAttr) return failure();
+      IntegerAttr intLhsAttr, intRhsAttr;
+      if (lhsAttr) intLhsAttr = lhsAttr.cast<IntegerAttr>();
+      if (rhsAttr) intRhsAttr = rhsAttr.cast<IntegerAttr>();
+      if (intLhsAttr && intRhsAttr) {
+        newOffsets.push_back(IntegerAttr::get(
+            intLhsAttr.getType(), intLhsAttr.getInt() + intRhsAttr.getInt()));
+      } else if (intLhsAttr && intLhsAttr.getInt() == 0) {
+        newOffsets.push_back(std::get<1>(z));
+      } else if (intRhsAttr && intRhsAttr.getInt() == 0) {
+        newOffsets.push_back(std::get<0>(z));
+      } else {
+        return failure();
+      }
+    }
+
+    rewriter.replaceOpWithNewOp<tensor::ExtractSliceOp>(
+        op, producerOp.getSource(), newOffsets, op.getMixedSizes(),
+        op.getMixedStrides());
+    return success();
+  }
+};
+
+/// Fold a tensor.extract_slice op if we can know that it's a full select slice.
+///
+/// convert:
+///   %0 = tensor.extract_slice %arg0[%o0, %o1][%s0, %s1][1, 1]
+///   %1 = destination_style_op(%0)
+///   %2 = tensor.extract_slice %1[0, 0][%s0, %s1][1, 1] // full select
+///   use(%2)
+/// to:
+///   %0 = tensor.extract_slice %arg0[%o0, %o1][%s0, %s1][1, 1]
+///   %1 = destination_style_op(%0)
+///   use(%1)
+struct FoldFullSelectedExtractSlicePattern
+    : public OpRewritePattern<tensor::ExtractSliceOp> {
+  using OpRewritePattern<tensor::ExtractSliceOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tensor::ExtractSliceOp op,
+                                PatternRewriter& rewriter) const override {
+    if (!offsetAllZeros(op) || !strideAllOnes(op)) return failure();
+
+    auto sliceOp = getSliceProducer(op.getSource());
+    if (!sliceOp || !hasSameShape(op, sliceOp)) return failure();
+    rewriter.replaceOp(op, op.getSource());
+    return success();
+  }
+};
+
+/// Fold a self assigned tensor.insert_slice op.
+///
+/// convert:
+///   %0 = tensor.extract_slice %arg0[%o0, %o1][%s0, %s1][1, 1]
+///   %1 = destination_style_op(%0)
+///   %2 = vector.transfer_write %vec, %1[%c0, %c0]
+///   %3 = tensor.insert_slice %2 into %1[0, 0] [%s0, %s1] [1, 1]
+///   use(%3)
+/// to:
+///   %0 = tensor.extract_slice %arg0[%o0, %o1][%s0, %s1][1, 1]
+///   %1 = destination_style_op(%0)
+///   %2 = vector.transfer_write %vec, %1[%c0, %c0]
+///   use(%2)
+struct FoldSelfInsertSlicePattern
+    : public OpRewritePattern<tensor::InsertSliceOp> {
+  using OpRewritePattern<tensor::InsertSliceOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tensor::InsertSliceOp op,
+                                PatternRewriter& rewriter) const override {
+    if (!offsetAllZeros(op) || !strideAllOnes(op)) return failure();
+
+    auto sourceSliceOp = getSliceProducer(op.getSource());
+    auto destSliceOp = getSliceProducer(op.getDest());
+    if (!sourceSliceOp || !destSliceOp) return failure();
+    // 1, dest & source have the same shape
+    if (!hasSameShape(sourceSliceOp, destSliceOp)) return failure();
+    // 2, the inserted tile from source and the dest have the same shape.
+    if (!hasSameShape(sourceSliceOp.getMixedSizes(), op.getMixedSizes()))
+      return failure();
+
+    rewriter.replaceOp(op, op.getSource());
+    return success();
+  }
+};
+
+/// convert:
+///   %0 = tensor.extract_slice %arg0[0, 0] [%s0, %s1] [1, 1] : tensor<?x?xf32>
+///   to tensor<?x?xf32> %1 = linalg.fill ins(%cst : f32) outs(%0 :
+///   tensor<?x?xf32>) -> tensor<?x?xf32>
+///   %2 = vector.transfer_read %1[0, 0],
+///            %cst : tensor<?x?xf32>, vector<6x16xf32>
+/// to:
+///   %2 = vector.splat %cst : vector<6x16xf32>
+struct TransferReadOfFillOpPattern
+    : public OpRewritePattern<vector::TransferReadOp> {
+  using OpRewritePattern<vector::TransferReadOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::TransferReadOp op,
+                                PatternRewriter& rewriter) const override {
+    // offsets are zeros.
+    if (!allConstIndexValues(op.getIndices(), 0)) return failure();
+
+    auto fillOp = op.getSource().getDefiningOp<linalg::FillOp>();
+    if (!fillOp || fillOp->getOperand(0) != op.getPadding()) return failure();
+    rewriter.replaceOpWithNewOp<vector::SplatOp>(op, op.getType(),
+                                                 op.getPadding());
+    return success();
+  }
+};
+
+/// convert:
+///   %0 = tensor.extract_slice %arg0[0, 0] [%s0, %s1] [1, 1] : tensor<?x?xf32>
+///   to tensor<?x?xf32> %1 = linalg.fill ins(%cst : f32) outs(%0 :
+///   tensor<?x?xf32>) -> tensor<?x?xf32>
+///   %2 = vector.transfer_write %result, %1[0, 0]
+///                     : vector<6x16xf32>, tensor<?x?xf32>
+/// to:
+///   %2 = vector.transfer_write %result, %0[0, 0] : vector<6x16xf32>
+/// iff:
+///   %1 has only one user
+///   %s0 < dim 0 of %result and %s1 < dim 1 of %result
+struct TransferWriteOfFillOpPattern
+    : public OpRewritePattern<vector::TransferWriteOp> {
+  using OpRewritePattern<vector::TransferWriteOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::TransferWriteOp op,
+                                PatternRewriter& rewriter) const override {
+    // offsets are zeros.
+    if (!allConstIndexValues(op.getIndices(), 0)) return failure();
+
+    auto fillOp = op.getSource().getDefiningOp<linalg::FillOp>();
+    if (!fillOp || !fillOp.result().hasOneUse()) return failure();
+
+    auto sliceOp = fillOp.output().getDefiningOp<tensor::ExtractSliceOp>();
+    if (!sliceOp || !strideAllOnes(sliceOp)) return failure();
+
+    auto type = op.getVector().getType().cast<VectorType>();
+    for (const auto& z : llvm::zip(sliceOp.getMixedSizes(), type.getShape())) {
+      if (std::get<1>(z) == ShapedType::kDynamicSize) return failure();
+      if (auto attr = std::get<0>(z).dyn_cast<Attribute>()) {
+        if (attr.cast<IntegerAttr>().getInt() > std::get<1>(z))
+          return failure();
+      } else {
+        auto bound = linalg::getConstantUpperBoundForIndex(
+            std::get<0>(z).dyn_cast<Value>());
+        if (failed(bound) || *bound > std::get<1>(z)) return failure();
+      }
+    }
+
+    auto cloned = rewriter.clone(*op.getOperation());
+    cloned->setOperand(1, fillOp.output());
+    rewriter.replaceOp(op, cloned->getResults());
+    return success();
+  }
+};
+
+static void addAllRegisteredCanonicalizationPatterns(
+    RewritePatternSet& patterns) {
+  MLIRContext* ctx = patterns.getContext();
+  for (Dialect* dialect : ctx->getLoadedDialects())
+    dialect->getCanonicalizationPatterns(patterns);
+  for (RegisteredOperationName op : ctx->getRegisteredOperations())
+    op.getCanonicalizationPatterns(patterns, ctx);
+  memref::populateResolveShapedTypeResultDimsPatterns(patterns);
+  memref::populateResolveRankedShapeTypeResultDimsPatterns(patterns);
+  patterns.insert<FoldExtractSliceOfExtractSlicePattern>(ctx);
+  patterns.insert<FoldFullSelectedExtractSlicePattern>(ctx);
+  patterns.insert<FoldSelfInsertSlicePattern>(ctx);
+  patterns.insert<TransferReadOfFillOpPattern>(ctx);
+  patterns.insert<TransferWriteOfFillOpPattern>(ctx);
+}
+
+}  // namespace
+
+void ApplyPatternsOp::build(OpBuilder& builder, OperationState& result,
+                            Value target, bool canonicalization) {
+  MLIRContext* ctx = builder.getContext();
+  result.addOperands(target);
+  if (canonicalization) {
+    result.addAttribute(
+        ApplyPatternsOp::getCanonicalizationAttrName(result.name),
+        builder.getUnitAttr());
+  }
+  result.addTypes({pdl::OperationType::get(ctx)});
+}
+
+DiagnosedSilenceableFailure ApplyPatternsOp::applyToOne(
+    Operation* target, SmallVectorImpl<Operation*>& results,
+    transform::TransformState& state) {
+  if (!target->hasTrait<OpTrait::IsIsolatedFromAbove>()) {
+    return mlir::emitDefiniteFailure(
+        target,
+        "applies only to isolated-from-above targets because it needs to apply "
+        "patterns greedily");
+  }
+  MLIRContext* ctx = target->getContext();
+  RewritePatternSet patterns(ctx);
+  if (getCanonicalization()) addAllRegisteredCanonicalizationPatterns(patterns);
+
+  TrackingListener listener(state);
+  GreedyRewriteConfig config;
+  LogicalResult result = applyPatternsAndFoldGreedily(
+      target, std::move(patterns), config, &listener);
+  LogicalResult listenerResult = listener.checkErrorState();
+  if (failed(result)) {
+    return mlir::emitDefiniteFailure(target,
+                                     "greedy pattern application failed");
+  }
+  if (failed(listenerResult))
+    return mlir::emitDefiniteFailure(target, "listener tracking failed");
+
+  results.assign({target});
+  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 260d398e0f5..ee824fdd512 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
@@ -49,4 +49,56 @@ def DISCBufferizeOp : Op<Transform_Dialect, "disc.bufferize",
   let cppNamespace = "::mlir::disc_ral::transform_dialect";
 }
 
+def ApplyPatternsOp : Op<Transform_Dialect, "disc.apply_patterns",
+    [FunctionalStyleTransformOpTrait,
+     MemoryEffectsOpInterface,
+     TransformEachOpTrait,
+     TransformOpInterface]> {
+  let description = [{
+    Greedily applies patterns as specified by its attributes.
+
+    Must be applied to an op with trait IsolatedFromAbove since the
+    GreedyPatternRewriter asserts those.
+
+    Returns the IsolatedFromAbove op whose content it has modified for better
+    chaining APIs.
+
+    The following additive attributes can be set, they add patterns in an
+    unspecified order:
+      - canonicalization: adds all the canonicalization patterns of all
+      registered dialects and ops.
+
+    Return modes:
+    =============
+    This operation applies a number of patterns to rewrite vector IR into
+    distributed warp form. To apply these patterns, this operation must target
+    an operation that is isolated from above, otherwise the transform definitely
+    fails.
+
+    If the pattern application fails, or if the underlying listener fails to
+    capture op handles, the transformation definitely fails.
+
+    Otherwise the transformation is successful and no result is returned.
+  }];
+
+  let arguments = (ins PDL_Operation:$target,
+                       UnitAttr:$canonicalization);
+  let results = (outs PDL_Operation:$result);
+
+  let assemblyFormat = "$target attr-dict";
+  let cppNamespace = "::mlir::disc_ral::transform_dialect";
+
+  let builders = [
+    // TODO: Some bitvector to scale better than n-bools.
+    OpBuilder<(ins "Value":$target, "bool":$canonicalization)>
+  ];
+
+  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_2.mlir b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_nn_d_f32_large_schedule_2.mlir
new file mode 100644
index 00000000000..1cb4d7d016b
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_nn_d_f32_large_schedule_2.mlir
@@ -0,0 +1,51 @@
+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
+
+  // further tile and fuse matmul and fill op.
+  %1:3 = transform.structured.fuse %0#1 {tile_sizes = [6, 16, 0], tile_interchange = [0, 1, 2]}
+
+  // gemm reduction axis tiling
+  %2:2 = transform.structured.tile %1#0 [0, 0, 1] {interchange=[0, 1, 2]}
+
+  // do some basic cleanup.
+  %func0 = transform.structured.match ops{["func.func"]} in %arg1
+  transform.disc.apply_patterns %func0 {canonicalization}
+
+  // pad to match the requirement of hardware vector/tensor instruction.
+  %3 = transform.structured.pad %2#0 {padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32], padding_dimensions=[0, 1, 2], pack_paddings=[1, 1, 0], hoist_paddings=[2, 3, 0], transpose_paddings=[[1, 0], [0, 1], [0, 1]]}
+
+  %func1 = transform.structured.match ops{["func.func"]} in %arg1
+  transform.structured.vectorize %func1 {vectorize_padding}
+
+  %func2 = transform.structured.match ops{["func.func"]} in %arg1
+  transform.disc.apply_patterns %func2 {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 dd57eb52c7a..6a3cc3f3208 100644
--- a/tao_compiler/mlir/disc/tests/disc-transform/matmul.cc
+++ b/tao_compiler/mlir/disc/tests/disc-transform/matmul.cc
@@ -22,6 +22,11 @@ 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(SimpleTest, MatMulF32_11x13x12) {
   EnvSetting setting = {{"DISC_TRANSFORM_SCHEDULE_FILE",
                          {c_ft_path + "matmul_nn_d_f32_schedule.mlir", false}},
@@ -90,4 +95,44 @@ TEST(SimpleTest, MatMulF32_1024x1024x1024) {
       /*profiling*/ true));
 }
 
+TEST(SimpleTest, MatMulF32_304x1024x256_2) {
+  EnvSetting setting = {
+      {"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"}}};
+  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*/ {"304x256xf32_X", "256x1024xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
+TEST(SimpleTest, MatMulF32_1024x1024x1024_2) {
+  EnvSetting setting = {
+      {"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"}}};
+  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));
+}
+
 }  // namespace mlir_test
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/apply-patterns-canonicalization.mlir b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/apply-patterns-canonicalization.mlir
new file mode 100644
index 00000000000..67e69b03a1f
--- /dev/null
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/apply-patterns-canonicalization.mlir
@@ -0,0 +1,111 @@
+// RUN: disc-opt --disc-transform-dialect-interpreter -split-input-file %s | FileCheck %s
+
+// CHECK-LABEL: @fold_extracted_slice
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index)
+func.func @fold_extracted_slice(%arg0: tensor<?x?xf32>, %arg1: index, %arg2: index, %arg3: index, %arg4: index) -> tensor<?x?xf32> {
+  // CHECK: %[[RET:.*]] = tensor.extract_slice %[[ARG0]]
+  // CHECK-SAME: %[[ARG2]], %[[ARG1]]
+  // CHECK-SAME: %[[ARG3]], %[[ARG4]]
+  // CHECK-SAME: 1, 1
+  // CHECK-NOT: tensor.extract_slice
+  // CHECK: return %[[RET]]
+  %0 = tensor.extract_slice %arg0[0, %arg1][%arg3, %arg4][1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+  %1 = tensor.extract_slice %0[%arg2, 0][%arg3, %arg4][1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+  return %1 : tensor<?x?xf32>
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  transform.disc.apply_patterns %arg1 {canonicalization}
+}
+
+// -----
+
+// CHECK-LABEL: @full_selected_slice
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index)
+func.func @full_selected_slice(%arg0: tensor<?x?xf32>, %arg1: index, %arg2: index, %arg3: index, %arg4: index) -> tensor<?x?xf32> {
+  // CHECK: %[[T0:.*]] = tensor.extract_slice
+  // CHECK: %[[T1:.*]] = linalg.fill
+  // CHECK-SAME: outs(%[[T0]] : tensor<?x?xf32>)
+  // CHECK-NOT: tensor.extract_slice
+  // CHECK: return %[[T1]]
+  %cst = arith.constant 0.000000e+00 : f32
+  %0 = tensor.extract_slice %arg0[%arg1, %arg2][%arg3, %arg4][1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+  %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<?x?xf32>) -> tensor<?x?xf32>
+  %2 = tensor.extract_slice %1[0, 0][%arg3, %arg4][1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+  return %2 : tensor<?x?xf32>
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  transform.disc.apply_patterns %arg1 {canonicalization}
+}
+
+// -----
+
+// CHECK-LABEL: @self_assigned_slice
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index, %[[ARG5:.*]]: vector<6x16xf32>)
+func.func @self_assigned_slice(%arg0: tensor<?x?xf32>, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %vec: vector<6x16xf32>) -> tensor<?x?xf32> {
+  // CHECK: %[[T0:.*]] = tensor.extract_slice
+  // CHECK: %[[T1:.*]] = linalg.fill
+  // CHECK-SAME: outs(%[[T0]] : tensor<?x?xf32>)
+  // CHECK: %[[T2:.*]] = vector.transfer_write %[[ARG5]], %[[T1]]
+  // CHECK: return %[[T2]]
+  %cst = arith.constant 0.000000e+00 : f32
+  %c0 = arith.constant 0 : index
+  %0 = tensor.extract_slice %arg0[%arg1, %arg2][%arg3, %arg4][1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+  %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<?x?xf32>) -> tensor<?x?xf32>
+  %2 = vector.transfer_write %vec, %1[%c0, %c0] : vector<6x16xf32>, tensor<?x?xf32>
+  %3 = tensor.insert_slice %2 into %1[0, 0] [%arg3, %arg4] [1, 1] : tensor<?x?xf32> into tensor<?x?xf32>
+  return %3 : tensor<?x?xf32>
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  transform.disc.apply_patterns %arg1 {canonicalization}
+}
+
+// -----
+
+// CHECK-LABEL: @transfer_read_of_fill
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index, %[[ARG5:.*]]: vector<6x16xf32>)
+func.func @transfer_read_of_fill(%arg0: tensor<?x?xf32>, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %vec: vector<6x16xf32>) -> vector<6x16xf32> {
+  // CHECK: %[[RET:.*]] = arith.constant dense<0.000000e+00> : vector<6x16xf32>
+  // CHECK: return %[[RET]]
+  %cst = arith.constant 0.000000e+00 : f32
+  %c0 = arith.constant 0 : index
+  %0 = tensor.extract_slice %arg0[0, 0][%arg3, %arg4][1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+  %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<?x?xf32>) -> tensor<?x?xf32>
+  %2 = vector.transfer_read %1[%c0, %c0], %cst : tensor<?x?xf32>, vector<6x16xf32>
+  return %2 : vector<6x16xf32>
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  transform.disc.apply_patterns %arg1 {canonicalization}
+}
+
+// -----
+
+// CHECK-LABEL: @transfer_write_of_fill
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index, %[[ARG5:.*]]: vector<6x16xf32>)
+#map1 = affine_map<(d0)[s0] -> (-d0 + s0, 6)>
+#map2 = affine_map<(d0)[s0] -> (-d0 + s0, 16)>
+func.func @transfer_write_of_fill(%arg0: tensor<?x?xf32>, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %vec: vector<6x16xf32>) -> tensor<?x?xf32> {
+  // CHECK: %[[T0:.*]] = tensor.extract_slice %[[ARG0]]
+  // CHECK: %[[T1:.*]] = vector.transfer_write %[[ARG5]], %[[T0]]
+  // CHECK: return %[[T1]]
+  %cst = arith.constant 0.000000e+00 : f32
+  %c0 = arith.constant 0 : index
+  %d0 = affine.min #map1(%arg1)[%arg2]
+  %d1 = affine.min #map2(%arg3)[%arg4]
+  %0 = tensor.extract_slice %arg0[0, 0][%d0, %d1][1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+  %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<?x?xf32>) -> tensor<?x?xf32>
+  %2 = vector.transfer_write %vec, %1[%c0, %c0] : vector<6x16xf32>, tensor<?x?xf32>
+  return %2 : tensor<?x?xf32>
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  transform.disc.apply_patterns %arg1 {canonicalization}
+}