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/BUILD b/tao_compiler/mlir/disc/BUILD
index c721dbd91bb..d4ca1861b9a 100644
--- a/tao_compiler/mlir/disc/BUILD
+++ b/tao_compiler/mlir/disc/BUILD
@@ -277,6 +277,14 @@ cc_library(
         "@llvm-project//mlir:AllPassesAndDialects",
         "@llvm-project//mlir:Analysis",
         "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:ArmSVEDialect",
+        "@llvm-project//mlir:ArmSVETransforms",
+        "@llvm-project//mlir:ArmNeonDialect",
+        "@llvm-project//mlir:X86VectorDialect",
+        "@llvm-project//mlir:X86VectorTransforms",
+        "@llvm-project//mlir:AMXDialect",
+        "@llvm-project//mlir:AMXTransforms",
+        "@llvm-project//mlir:VectorDialect",
         "@llvm-project//mlir:GPUDialect",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:LLVMDialect",
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/BUILD b/tao_compiler/mlir/disc/tools/disc-transform/BUILD
index ac2a66817a3..a8b1e94a7d2 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/BUILD
+++ b/tao_compiler/mlir/disc/tools/disc-transform/BUILD
@@ -48,6 +48,25 @@ gentbl_cc_library(
     ],
 )
 
+cc_library(
+    name = "disc_transform_utils",
+    srcs = [
+        "utils.cc",
+    ],
+    hdrs = [
+        "utils.h",
+    ],
+    deps = [
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:LinalgDialect",
+        "@llvm-project//mlir:MemRefDialect",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:Transforms",
+    ],
+    alwayslink = 1,
+)
+
 cc_library(
     name = "disc_transform_ops",
     srcs = [
@@ -59,6 +78,7 @@ cc_library(
         "TransformOps/TransformOpsExt.h",
     ],
     deps = [
+        ":disc_transform_utils",
         ":transform_ops_inc_gen",
         "@iree-dialects//:IREELinalgExtDialect",
         "@iree-dialects//:IREELinalgExtTransformOps",
@@ -204,6 +224,25 @@ cc_library(
     alwayslink = 1,
 )
 
+cc_library(
+    name = "memref_copy_to_linalg",
+    srcs = ["transforms/memref_copy_to_linalg.cc"],
+    deps = [
+        ":pass_details",
+        ":disc_transform_utils",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:LinalgDialect",
+        "@llvm-project//mlir:MemRefDialect",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:Transforms",
+    ],
+    alwayslink = 1,
+)
+
 cc_library(
     name = "all_passes",
     hdrs = [
@@ -215,6 +254,7 @@ cc_library(
     ],
     deps = [
         ":legalize_lmhlo_fusion_to_linalg",
+        ":memref_copy_to_linalg",
         ":rewrite-payload-ir-for-ral",
         ":transform_dialect_interpreter",
         "@llvm-project//mlir:Pass",
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 c12d1baf3e7..1e0680aeff6 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
@@ -33,6 +33,7 @@
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Pass/PassManager.h"
+#include "tensorflow/compiler/mlir/disc/tools/disc-transform/utils.h"
 
 namespace mlir {
 namespace disc_ral {
@@ -84,36 +85,6 @@ LogicalResult comprehensiveBufferizeDeallocationFn(OpBuilder& builder,
   return success();
 }
 
-/// Create a linalg::GenericOp version of an n-D copy that can further tile,
-/// lower to loops or vectorize, unlike the current implementation of
-/// memref::CopyOp.
-Operation* createLinalgCopyOp(OpBuilder& b, Location loc, Value from, Value to,
-                              ArrayRef<NamedAttribute> attributes = {}) {
-  auto memrefTypeFrom = from.getType().dyn_cast<MemRefType>();
-  auto memrefTypeTo = to.getType().dyn_cast<MemRefType>();
-  if (!memrefTypeFrom || !memrefTypeTo ||
-      memrefTypeFrom.getRank() != memrefTypeTo.getRank()) {
-    mlir::emitError(
-        loc, "unable to generate copy op within bufferization from type ")
-        << memrefTypeFrom << " to " << memrefTypeTo;
-    return nullptr;
-  }
-  AffineMap id =
-      AffineMap::getMultiDimIdentityMap(memrefTypeTo.getRank(), b.getContext());
-  SmallVector<StringRef> iteratorTypes(memrefTypeTo.getRank(),
-                                       getParallelIteratorTypeName());
-  return b.create<linalg::GenericOp>(
-      loc,
-      /*inputs=*/from,
-      /*outputs=*/to,
-      /*indexingMaps=*/llvm::makeArrayRef({id, id}),
-      /*iteratorTypes=*/iteratorTypes,
-      [](OpBuilder& b, Location loc, ValueRange args) {
-        b.create<linalg::YieldOp>(loc, args.front());
-      },
-      attributes);
-}
-
 LogicalResult comprehensiveBufferizeCopyFn(OpBuilder& builder, Location loc,
                                            Value from, Value to) {
   createLinalgCopyOp(builder, loc, from, to);
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/PassDetail.h b/tao_compiler/mlir/disc/tools/disc-transform/transforms/PassDetail.h
index be1143893a6..e5e584d0fdd 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/transforms/PassDetail.h
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/PassDetail.h
@@ -54,6 +54,10 @@ namespace arith {
 class ArithDialect;
 }
 
+namespace linalg {
+class LinalgDialect;
+}
+
 namespace disc_ral {
 
 #define GEN_PASS_CLASSES
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/memref_copy_to_linalg.cc b/tao_compiler/mlir/disc/tools/disc-transform/transforms/memref_copy_to_linalg.cc
new file mode 100644
index 00000000000..bb731e6aa8e
--- /dev/null
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/memref_copy_to_linalg.cc
@@ -0,0 +1,71 @@
+// 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 "llvm/Support/Debug.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "mlir/Transforms/Passes.h"
+#include "tensorflow/compiler/mlir/disc/tools/disc-transform/transforms/PassDetail.h"
+#include "tensorflow/compiler/mlir/disc/tools/disc-transform/utils.h"
+
+#define DEBUG_TYPE "disc-memref-copy-to-linalg"
+
+// This file implements the logic to convert a memref.copy op to its linalg
+// equivalent.
+
+namespace mlir {
+namespace disc_ral {
+namespace {
+
+using func::FuncOp;
+
+struct DiscMemrefCopyToLinalgPass
+    : public DiscMemrefCopyToLinalgPassBase<DiscMemrefCopyToLinalgPass> {
+  void runOnOperation() override;
+};
+
+struct MemrefCopyOpToLinalg : public OpRewritePattern<memref::CopyOp> {
+  using OpRewritePattern<memref::CopyOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(memref::CopyOp copyOp,
+                                PatternRewriter& rewriter) const override {
+    Operation* linalgCopy =
+        createLinalgCopyOp(rewriter, copyOp.getLoc(), copyOp.getSource(),
+                           copyOp.getTarget(), copyOp->getAttrs());
+    if (!linalgCopy) return failure();
+    rewriter.replaceOp(copyOp, linalgCopy->getResults());
+    return success();
+  }
+};
+
+void DiscMemrefCopyToLinalgPass::runOnOperation() {
+  MLIRContext* context = &getContext();
+  RewritePatternSet patterns(&getContext());
+  patterns.insert<MemrefCopyOpToLinalg>(context);
+  if (failed(
+          applyPatternsAndFoldGreedily(getOperation(), std::move(patterns)))) {
+    return signalPassFailure();
+  }
+}
+
+}  // namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>>
+createDiscMemrefCopyToLinalgPass() {
+  return std::make_unique<DiscMemrefCopyToLinalgPass>();
+}
+
+}  // namespace disc_ral
+}  // namespace mlir
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/passes.h b/tao_compiler/mlir/disc/tools/disc-transform/transforms/passes.h
index afd67f387f4..06332b4119c 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/transforms/passes.h
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/passes.h
@@ -21,12 +21,15 @@ limitations under the License.
 
 namespace mlir {
 
-class FuncOp;
 class FunctionPass;
 class ModuleOp;
 class Operation;
 class Pass;
 
+namespace func {
+class FuncOp;
+}
+
 template <typename T>
 class OperationPass;
 
@@ -46,6 +49,9 @@ createDiscTransformDialectInterpreterPass(const std::string& fileName = "",
 std::unique_ptr<OperationPass<ModuleOp>> createDiscRewritePayloadIRForRALPass(
     bool gpuEnabled = false);
 
+// Converts a memref.copy op to its linalg equivalent
+std::unique_ptr<OperationPass<func::FuncOp>> createDiscMemrefCopyToLinalgPass();
+
 }  // namespace disc_ral
 }  // namespace mlir
 
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/transform_passes.td b/tao_compiler/mlir/disc/tools/disc-transform/transforms/transform_passes.td
index e6ba060fd24..1bcaf54b258 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/transforms/transform_passes.td
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/transform_passes.td
@@ -39,3 +39,11 @@ def DiscRewritePayloadIRForRALPass : Pass<"disc-rewrite-payload-ir-for-ral", "Mo
             /*default=*/"false", "whether gpu is available.">,
   ];
 }
+
+def DiscMemrefCopyToLinalgPass : Pass<"disc-memref-copy-to-linalg", "func::FuncOp"> {
+  let summary = "Pass to convert memref.copy to linalg.";
+  let constructor = "createDiscMemrefCopyToLinalgPass()";
+  let dependentDialects = [
+      "linalg::LinalgDialect",
+  ];
+}
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/utils.cc b/tao_compiler/mlir/disc/tools/disc-transform/utils.cc
new file mode 100644
index 00000000000..035b1575c38
--- /dev/null
+++ b/tao_compiler/mlir/disc/tools/disc-transform/utils.cc
@@ -0,0 +1,53 @@
+// 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/tools/disc-transform/utils.h"
+
+#include "llvm/Support/Debug.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+
+#define DEBUG_TYPE "disc-transform-utils"
+
+namespace mlir {
+namespace disc_ral {
+
+/// Create a linalg::GenericOp version of an n-D copy that can further tile,
+/// lower to loops or vectorize, unlike the current implementation of
+/// memref::CopyOp.
+Operation* createLinalgCopyOp(OpBuilder& b, Location loc, Value from, Value to,
+                              ArrayRef<NamedAttribute> attributes) {
+  auto memrefTypeFrom = from.getType().dyn_cast<MemRefType>();
+  auto memrefTypeTo = to.getType().dyn_cast<MemRefType>();
+  if (!memrefTypeFrom || !memrefTypeTo ||
+      memrefTypeFrom.getRank() != memrefTypeTo.getRank()) {
+    mlir::emitError(
+        loc, "unable to generate copy op within bufferization from type ")
+        << memrefTypeFrom << " to " << memrefTypeTo;
+    return nullptr;
+  }
+  AffineMap id =
+      AffineMap::getMultiDimIdentityMap(memrefTypeTo.getRank(), b.getContext());
+  SmallVector<StringRef> iteratorTypes(memrefTypeTo.getRank(),
+                                       getParallelIteratorTypeName());
+  return b.create<linalg::GenericOp>(
+      loc,
+      /*inputs=*/from,
+      /*outputs=*/to,
+      /*indexingMaps=*/llvm::makeArrayRef({id, id}),
+      /*iteratorTypes=*/iteratorTypes,
+      [](OpBuilder& b, Location loc, ValueRange args) {
+        b.create<linalg::YieldOp>(loc, args.front());
+      },
+      attributes);
+}
+
+}  // namespace disc_ral
+}  // namespace mlir
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/utils.h b/tao_compiler/mlir/disc/tools/disc-transform/utils.h
new file mode 100644
index 00000000000..03be831ea1b
--- /dev/null
+++ b/tao_compiler/mlir/disc/tools/disc-transform/utils.h
@@ -0,0 +1,33 @@
+/* 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.
+==============================================================================*/
+
+#ifndef DISC_TOOLS_DISC_TRANSFORM_UTILS_H_
+#define DISC_TOOLS_DISC_TRANSFORM_UTILS_H_
+
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+
+namespace mlir {
+namespace disc_ral {
+
+/// Create a linalg::GenericOp version of an n-D copy that can further tile,
+/// lower to loops or vectorize, unlike the current implementation of
+/// memref::CopyOp.
+Operation* createLinalgCopyOp(OpBuilder& b, Location loc, Value from, Value to,
+                              ArrayRef<NamedAttribute> attributes = {});
+
+}  // namespace disc_ral
+}  // namespace mlir
+
+#endif  // DISC_TOOLS_DISC_TRANSFORM_UTILS_H_
diff --git a/tao_compiler/mlir/disc/transforms/disc_to_llvm.cc b/tao_compiler/mlir/disc/transforms/disc_to_llvm.cc
index fc5138e1da7..a6f346b3f55 100644
--- a/tao_compiler/mlir/disc/transforms/disc_to_llvm.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_to_llvm.cc
@@ -20,16 +20,26 @@
 #include "mlir/Conversion/LLVMCommon/Pattern.h"
 #include "mlir/Conversion/MathToLLVM/MathToLLVM.h"
 #include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"
+#include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
+#include "mlir/Dialect/AMX/AMXDialect.h"
+#include "mlir/Dialect/AMX/Transforms.h"
 #include "mlir/Dialect/Arith/Transforms/Passes.h"
+#include "mlir/Dialect/ArmNeon/ArmNeonDialect.h"
+#include "mlir/Dialect/ArmSVE/ArmSVEDialect.h"
+#include "mlir/Dialect/ArmSVE/Transforms.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
 #include "mlir/Dialect/Math/IR/Math.h"
+#include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
+#include "mlir/Dialect/X86Vector/Transforms.h"
+#include "mlir/Dialect/X86Vector/X86VectorDialect.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Transforms/DialectConversion.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "tensorflow/compiler/mlir/disc/IR/disc_ral_ops.h"
 #include "tensorflow/compiler/mlir/disc/IR/lhlo_disc_ops.h"
 #include "tensorflow/compiler/mlir/disc/disc_util.h"
@@ -1171,6 +1181,21 @@ class DiscToLLVMPass : public DiscToLLVMPassBase<DiscToLLVMPass> {
       return LLVM::LLVMPointerType::get(IntegerType::get(ctx, 8));
     });
 
+    // Run Vector -> Vector transformations ahead of conversion to LLVM.
+    {
+      RewritePatternSet patterns(&getContext());
+      vector::populateVectorToVectorCanonicalizationPatterns(patterns);
+      vector::populateVectorBroadcastLoweringPatterns(patterns);
+      vector::populateVectorContractLoweringPatterns(patterns);
+      vector::populateVectorMaskOpLoweringPatterns(patterns);
+      vector::populateVectorShapeCastLoweringPatterns(patterns);
+      vector::populateVectorTransposeLoweringPatterns(patterns);
+      if (failed(applyPatternsAndFoldGreedily(getOperation(),
+                                              std::move(patterns)))) {
+        return signalPassFailure();
+      }
+    }
+
     // Populate patterns.
     RewritePatternSet patterns(&getContext());
     mlir::arith::populateArithToLLVMConversionPatterns(type_converter,
@@ -1182,6 +1207,13 @@ class DiscToLLVMPass : public DiscToLLVMPassBase<DiscToLLVMPass> {
     cf::populateControlFlowToLLVMConversionPatterns(type_converter, patterns);
     populateDiscToLLVMConversionPatterns(&type_converter, &symbol_table,
                                          &patterns);
+    vector::populateVectorMaskMaterializationPatterns(
+        patterns, /*force32BitVectorIndices*/ true);
+    vector::populateVectorTransferLoweringPatterns(patterns);
+    populateVectorToLLVMMatrixConversionPatterns(type_converter, patterns);
+    populateVectorToLLVMConversionPatterns(type_converter, patterns,
+                                           /*reassociateFPReductions*/ true,
+                                           /*force32BitVectorIndices*/ true);
 
     // Set target.
     ConversionTarget target(*ctx);
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 ede2cb28c00..95760b4ac67 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
@@ -134,6 +134,7 @@ LogicalResult DiscTransformLegalizeToLoopPass::runTransformPipeline(
   pm.addPass(createDiscLegalizeLmhloFusionToLinalgPass());
   pm.addPass(createDiscTransformDialectInterpreterPass(transformFileName_,
                                                        enableExpensiveChecks_));
+  pm.addNestedPass<func::FuncOp>(createDiscMemrefCopyToLinalgPass());
   pm.addPass(createDiscRewritePayloadIRForRALPass(gpuEnabled_));
   return pm.run(m);
 }


Test patch:
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_multithread_nn_d_f32_large_schedule.mlir b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_multithread_nn_d_f32_large_schedule.mlir
new file mode 100644
index 00000000000..8b17586141d
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_multithread_nn_d_f32_large_schedule.mlir
@@ -0,0 +1,40 @@
+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
+
+  // TODO(wyzero): This actually disable multi-thread because packing for weights is not
+  // supported in muti-threading config a.t.m.
+  %0:2 = transform.structured.tile_to_foreach_thread_op %matmul tile_sizes [1024, 1024]
+  transform.structured.fuse_into_containing_op %fill into %0#0
+  %1:4 = transform.structured.tile %0#1 [6, 16, 1] {interchange=[0, 1, 2]}
+  %2 = transform.structured.pad %1#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]]}
+
+  %func = transform.structured.match ops{["func.func"]} in %arg1
+  transform.structured.vectorize %func {vectorize_padding}
+
+  transform.disc.bufferize %arg1
+
+  // TODO(wyzero): We can not merge two `lower_vectors` into one due to some wierd
+  // bugs in some tile configurations. For example, it'll not converge if the tile
+  // size of the k is one.
+  transform.lower_vectors {
+    contraction_lowering = "outerproduct",
+    multireduction_lowering = "innerparallel",
+    split_transfers = "linalg-copy",
+    stages = [0, 1, 2, 3, 4],
+    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 = [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.mlir b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_nn_d_f32_large_schedule.mlir
new file mode 100644
index 00000000000..b214772bf81
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_nn_d_f32_large_schedule.mlir
@@ -0,0 +1,38 @@
+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
+  %1:4 = transform.structured.tile %0#1 [6, 16, 1] {interchange=[0, 1, 2]}
+  %2 = transform.structured.pad %1#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]]}
+
+  %func = transform.structured.match ops{["func.func"]} in %arg1
+  transform.structured.vectorize %func {vectorize_padding}
+
+  transform.disc.bufferize %arg1
+
+  // TODO(wyzero): We can not merge two `lower_vectors` into one due to some wierd
+  // bugs in some tile configurations. For example, it'll not converge if the tile
+  // size of the k is one.
+  transform.lower_vectors {
+    contraction_lowering = "outerproduct",
+    multireduction_lowering = "innerparallel",
+    split_transfers = "linalg-copy",
+    stages = [0, 1, 2, 3, 4],
+    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 = [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 ffa6aa9fb80..dd57eb52c7a 100644
--- a/tao_compiler/mlir/disc/tests/disc-transform/matmul.cc
+++ b/tao_compiler/mlir/disc/tests/disc-transform/matmul.cc
@@ -50,4 +50,44 @@ TEST(SimpleTest, MatMulF32_111x131x121) {
       /*output_descriptors*/ {"f32_X"}));
 }
 
+TEST(SimpleTest, MatMulF32_304x1024x256) {
+  EnvSetting setting = {
+      {"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"}}};
+  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) {
+  EnvSetting setting = {
+      {"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"}}};
+  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/tests/disc-transform/matmul_multithread.cc b/tao_compiler/mlir/disc/tests/disc-transform/matmul_multithread.cc
index 7542c631591..4c830625e16 100644
--- a/tao_compiler/mlir/disc/tests/disc-transform/matmul_multithread.cc
+++ b/tao_compiler/mlir/disc/tests/disc-transform/matmul_multithread.cc
@@ -42,4 +42,19 @@ TEST(SimpleMTTest, MatMulF32_111x131x121_Thread_8) {
       /*output_descriptors*/ {"f32_X"}));
 }
 
+TEST(SimpleTest, MatMulF32_304x1024x256) {
+  EnvSetting setting = {
+      {"DISC_TRANSFORM_SCHEDULE_FILE",
+       {c_ft_path + "matmul_multithread_nn_d_f32_large_schedule.mlir", false}},
+      {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}}};
+  EnvSettingContext ctx(setting);
+  EXPECT_TRUE(feature_test_main(
+      /*mlir_file_path*/ c_ft_path + "matmul_multithread_nn_d_f32.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 2,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"304x256xf32_X", "256x1024xf32_X"},
+      /*output_descriptors*/ {"f32_X"}));
+}
+
 }  // namespace mlir_test
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/memref-copy-to-linalg.mlir b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/memref-copy-to-linalg.mlir
new file mode 100644
index 00000000000..bab6e8aaeaf
--- /dev/null
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/memref-copy-to-linalg.mlir
@@ -0,0 +1,9 @@
+// RUN: disc-opt --disc-memref-copy-to-linalg -split-input-file %s | FileCheck %s
+
+// CHECK-LABEL: @copy
+func.func @copy(%arg1: memref<?x?xf32, "cpu">, %arg2: memref<?x?xf32, "cpu">) -> memref<?x?xf32, "cpu"> {
+  // CHECK-NOT: memref.copy
+  // CHECK: linalg.generic
+  memref.copy %arg1, %arg2 : memref<?x?xf32, "cpu"> to memref<?x?xf32, "cpu">
+  return %arg2 : memref<?x?xf32, "cpu">
+}
\ No newline at end of file