Repository: alibaba/BladeDISC (language: python)

Issue:
Introduce MLIR transform dialect to BladeDISC
We'll start to explore using MLIR [transform dialect ](https://mlir.llvm.org/docs/Dialects/Transform) to do codegen for (fused) compute-intensive pattern. The initial target is to support gemm codegen on ARM platform to address the dynamic shape problem of Arm Compute Library.

The initial plan is:

- [x]  Step 1,  enhance the fusion decision pass. We’ll add a new fusion kind `kTransform` for the transform-based fusion pattern.
- [x]  Step 2,  lower the lmhlo fusion op to linalg on tensor.
- [x]  Step 3,  transform the linalg computation to loops using transform dialect.
- [x]  Step 4,  refined the transformed loop to make it suitable for BladeDISC runtime.
- [x]  Step 5, add a new pass to the disc pass pipeline to drive the above process.
- [x]   Step 6, weight pre-packing support
   - [x]  add `disc_linalg.multi_level_pack` op, used for doing packing.
   - [x]  add  `transform.disc.cache_read` transform op, relying on `disc_linalg.multi_level_pack` op.
   - [x]  add folding support for `disc_linalg.multi_level_pack`.
   - [x]  lower `disc_linalg.multi_level_pack` to loop if it can not be folded.
   - [x]  fuse const weight op into the `kTransform` fusion pattern, lower it to linalg and then schedule it.
- [x]  Step 7,  assign a default schedule for each `kTransform` pattern.  
- [x]  Step 8,  schedule selection logic injection
- [x]  Step 9,  initial model level testing:  bert (albert).
- [x]  Step 10, support nt, tn, tt format GEMM.
- [ ]  Step 11, support batch matmul
- [x]  Step 12, support GEMM epilogue fusion.
- [ ]  Step 13, performance optimization


Discussion / hints:
e2e model test on: Bert Base (TF) and Albert (PyTorch),  on g6r, using single thread. Note that we only have one default schedule for all shape and the schedule is known to be less performant when n or k is large, thus the initial performance is supposed to be improved when we support schedule selection logic.

Bert Base (TF) 

input | TF 2.8(s) | DISC-ACL(s) | DISC-Transform(s) | speedup (DISC-transform / DISC-ACL)
-- | -- | -- | -- | --
(1, 128) | 0.742 | 0.638 | 0.656 | 97.3%
(2, 128) | 1.41 | 1.24 | 1.27 | 97.6%
(4, 128) | 2.85 | 2.36 | 2.55 | 92.5%
(8, 128) | 5.84 | 4.68 | 5.07 | 92.3%
(16, 128) | 11.9 | 9.55 | 10.2 | 93.6%
----


Albert (PyTorch)

input | TorchScript | OnnxRuntime | DISC-ACL | DISC-Transform
-- | -- | -- | -- | --
(2, 12) | 0.197 | 0.140 | 0.117 | 0.139



Fix patch:
diff --git a/tao_compiler/mlir/disc/disc_compiler.cc b/tao_compiler/mlir/disc/disc_compiler.cc
index 0254846abd2..56da2688cad 100644
--- a/tao_compiler/mlir/disc/disc_compiler.cc
+++ b/tao_compiler/mlir/disc/disc_compiler.cc
@@ -63,6 +63,7 @@ limitations under the License.
 #include "mlir/Target/LLVMIR/Export.h"  // from @llvm-project
 #include "mlir/Transforms/Passes.h"
 #include "tensorflow/compiler/mlir/disc/disc_util.h"
+#include "tensorflow/compiler/mlir/disc/tools/disc-transform/transforms/passes.h"
 #include "tensorflow/compiler/mlir/disc/transforms/codegen_utils.h"
 #include "tensorflow/compiler/mlir/disc/transforms/fusion_utils.h"
 #include "tensorflow/compiler/mlir/disc/transforms/passes.h"
@@ -586,6 +587,8 @@ LogicalResult LowerHLOToLLVM(ModuleOp m, const DISCLoweringOptions& options) {
     pm.addPass(createGpuKernelOutliningPass());
     pm.addPass(disc_ral::createDiscAssignKernelNamePass());
   } else {
+    pm.addNestedPass<FuncOp>(
+        disc_ral::createDiscConvertForeachThreadOpToParallelOpPass());
     if (options.cpu_options.target_multi_threading) {
       pm.addNestedPass<FuncOp>(disc_ral::createDiscCpuMapParallelLoopPass());
       pm.addPass(disc_ral::createDiscOutlineCpuKernelPass());
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/BUILD b/tao_compiler/mlir/disc/tools/disc-transform/BUILD
index 8f73306d93d..686e8b284ee 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/BUILD
+++ b/tao_compiler/mlir/disc/tools/disc-transform/BUILD
@@ -371,6 +371,23 @@ cc_library(
     alwayslink = 1,
 )
 
+cc_library(
+    name = "convert-foreach-thread-op-to-parallel-op",
+    srcs = ["transforms/convert_foreach_thread_op_to_parallel_op.cc"],
+    deps = [
+        ":pass_details",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:SCFDialect",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:Transforms",
+    ],
+    alwayslink = 1,
+)
+
 cc_library(
     name = "all_passes",
     hdrs = [
@@ -381,6 +398,7 @@ cc_library(
         ":friends",
     ],
     deps = [
+        ":convert-foreach-thread-op-to-parallel-op",
         ":legalize_lmhlo_fusion_to_linalg",
         ":memref_copy_to_linalg",
         ":rewrite-payload-ir-for-ral",
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/convert_foreach_thread_op_to_parallel_op.cc b/tao_compiler/mlir/disc/tools/disc-transform/transforms/convert_foreach_thread_op_to_parallel_op.cc
new file mode 100644
index 00000000000..25c38173e5d
--- /dev/null
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/convert_foreach_thread_op_to_parallel_op.cc
@@ -0,0 +1,88 @@
+// Copyright 2023 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/SCF/IR/SCF.h"
+#include "mlir/IR/BlockAndValueMapping.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"
+
+#define DEBUG_TYPE "disc-convert-foreach-thread-op-to-parallel-op"
+
+// This file implements the logic to convert scf.foreach_thread to scf.parallel
+
+namespace mlir {
+namespace disc_ral {
+namespace {
+
+using func::FuncOp;
+using scf::ForeachThreadOp;
+using scf::ParallelOp;
+
+struct ForeachThreadToParallel : public OpRewritePattern<ForeachThreadOp> {
+  using OpRewritePattern<ForeachThreadOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ForeachThreadOp foreachThreadOp,
+                                PatternRewriter& rewriter) const override {
+    if (foreachThreadOp->getNumResults() != 0) return failure();
+    Location loc = foreachThreadOp.getLoc();
+    int64_t rank = foreachThreadOp.getRank();
+    Value zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
+    Value one = rewriter.create<arith::ConstantIndexOp>(loc, 1);
+    SmallVector<Value> lowerBounds(rank, zero);
+    SmallVector<Value> upperBounds = foreachThreadOp.getNumThreads();
+    SmallVector<Value> steps(rank, one);
+
+    auto parallelOp =
+        rewriter.create<scf::ParallelOp>(loc, lowerBounds, upperBounds, steps);
+    BlockAndValueMapping mapping;
+    for (const auto& z : llvm::zip(foreachThreadOp.getThreadIndices(),
+                                   parallelOp.getInductionVars()))
+      mapping.map(std::get<0>(z), std::get<1>(z));
+    rewriter.setInsertionPointToStart(parallelOp.getBody());
+    for (auto& nestedOp : foreachThreadOp.getBody()->without_terminator()) {
+      rewriter.clone(nestedOp, mapping);
+    }
+    rewriter.replaceOp(foreachThreadOp, parallelOp->getResults());
+    return success();
+  }
+};
+
+struct DiscConvertForeachThreadOpToParallelOpPass
+    : public DiscConvertForeachThreadOpToParallelOpPassBase<
+          DiscConvertForeachThreadOpToParallelOpPass> {
+  void runOnOperation() override;
+};
+
+void DiscConvertForeachThreadOpToParallelOpPass::runOnOperation() {
+  MLIRContext* context = &getContext();
+  RewritePatternSet patterns(&getContext());
+  patterns.insert<ForeachThreadToParallel>(context);
+  if (failed(
+          applyPatternsAndFoldGreedily(getOperation(), std::move(patterns)))) {
+    return signalPassFailure();
+  }
+}
+
+}  // namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>>
+createDiscConvertForeachThreadOpToParallelOpPass() {
+  return std::make_unique<DiscConvertForeachThreadOpToParallelOpPass>();
+}
+
+}  // 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 98076247412..cd49806c8fa 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/transforms/passes.h
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/passes.h
@@ -60,6 +60,10 @@ std::unique_ptr<OperationPass<ModuleOp>> createDiscRewritePayloadIRForRALPass(
 // Converts a memref.copy op to its linalg equivalent
 std::unique_ptr<OperationPass<func::FuncOp>> createDiscMemrefCopyToLinalgPass();
 
+// Converts scf.foreach_thread to scf.parallel
+std::unique_ptr<OperationPass<func::FuncOp>>
+createDiscConvertForeachThreadOpToParallelOpPass();
+
 }  // namespace disc_ral
 }  // namespace mlir
 
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/rewrite_payload_ir_for_ral.cc b/tao_compiler/mlir/disc/tools/disc-transform/transforms/rewrite_payload_ir_for_ral.cc
index dd6e671961d..3c783ca8ac4 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/transforms/rewrite_payload_ir_for_ral.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/rewrite_payload_ir_for_ral.cc
@@ -46,61 +46,12 @@ struct DiscRewritePayloadIRForRALPass
   }
   void runOnOperation() override;
 
-  // replace scf::foreach_thread op with scf::parallel op
-  LogicalResult convertForeachThreadToParallelOp();
-  LogicalResult funcLevelConvertForeachThreadToParallelOp(FuncOp funcOp);
-
   // assign placement info for each memref value, e.g. memref<f32> ->
   // memref<f32, "cpu">
   LogicalResult assignPlacement();
   LogicalResult assignPlacementForFuncOp(FuncOp funcOp);
 };
 
-LogicalResult
-DiscRewritePayloadIRForRALPass::funcLevelConvertForeachThreadToParallelOp(
-    FuncOp funcOp) {
-  SmallVector<ForeachThreadOp> forOps;
-  funcOp.walk([&](ForeachThreadOp op) { forOps.push_back(op); });
-
-  OpBuilder b(funcOp);
-  for (ForeachThreadOp foreachThreadOp : forOps) {
-    if (foreachThreadOp.getOutputs().size() != 0)
-      return foreachThreadOp->emitError()
-             << "Not support ForeachThreadOp with outputs a.t.m.\n";
-
-    b.setInsertionPoint(foreachThreadOp);
-    Location loc = foreachThreadOp.getLoc();
-    int64_t rank = foreachThreadOp.getRank();
-    Value zero = b.create<arith::ConstantIndexOp>(loc, 0);
-    Value one = b.create<arith::ConstantIndexOp>(loc, 1);
-    SmallVector<Value> lowerBounds(rank, zero);
-    SmallVector<Value> upperBounds = foreachThreadOp.getNumThreads();
-    SmallVector<Value> steps(rank, one);
-
-    auto parallelOp =
-        b.create<ParallelOp>(loc, lowerBounds, upperBounds, steps);
-    BlockAndValueMapping mapping;
-    for (const auto& z : llvm::zip(foreachThreadOp.getThreadIndices(),
-                                   parallelOp.getInductionVars()))
-      mapping.map(std::get<0>(z), std::get<1>(z));
-    b.setInsertionPointToStart(parallelOp.getBody());
-    for (auto& nestedOp : foreachThreadOp.getBody()->without_terminator()) {
-      Operation* cloned = b.clone(nestedOp, mapping);
-    }
-    foreachThreadOp->erase();
-  }
-  return success();
-}
-
-LogicalResult
-DiscRewritePayloadIRForRALPass::convertForeachThreadToParallelOp() {
-  for (auto funcOp : getOperation().getOps<FuncOp>()) {
-    if (failed(funcLevelConvertForeachThreadToParallelOp(funcOp)))
-      return failure();
-  }
-  return success();
-}
-
 LogicalResult DiscRewritePayloadIRForRALPass::assignPlacementForFuncOp(
     FuncOp funcOp) {
   auto maybeConvertType = [&](Type ty) -> Type {
@@ -170,14 +121,7 @@ LogicalResult DiscRewritePayloadIRForRALPass::assignPlacement() {
 }
 
 void DiscRewritePayloadIRForRALPass::runOnOperation() {
-  // 1, rewrite scf.foreach_thread to scf.parallel
-  if (failed(convertForeachThreadToParallelOp())) {
-    return signalPassFailure();
-  }
-  LLVM_DEBUG(llvm::dbgs() << "After ForeachThreadOp -> ParallelOp:\n"
-                          << getOperation() << "\n");
-
-  // 2, assign placement info for each memref value.
+  // assign placement info for each memref value.
   if (failed(assignPlacement())) {
     return signalPassFailure();
   }
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 203e36a2620..d402df8929d 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
@@ -55,3 +55,8 @@ def DiscMemrefCopyToLinalgPass : Pass<"disc-memref-copy-to-linalg", "func::FuncO
       "linalg::LinalgDialect",
   ];
 }
+
+def DiscConvertForeachThreadOpToParallelOpPass : Pass<"disc-convert-foreach-thread-op-to-parallel-op", "func::FuncOp"> {
+  let summary = "Pass to convert memref.copy to linalg.";
+  let constructor = "createDiscConvertForeachThreadOpToParallelOpPass()";
+}
diff --git a/tao_compiler/mlir/disc/transforms/disc_outline_cpu_kernel.cc b/tao_compiler/mlir/disc/transforms/disc_outline_cpu_kernel.cc
index 2fd4a8bf596..47d9222fb58 100644
--- a/tao_compiler/mlir/disc/transforms/disc_outline_cpu_kernel.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_outline_cpu_kernel.cc
@@ -221,34 +221,39 @@ static void convertToLaunchFuncOp(Operation* launchOp, func::FuncOp kernelFunc,
   launchOp->erase();
 }
 
+/// Checks whether the given op can be hoisted by checking that
+/// - the op and none of its contained operations depend on values inside of the
+///   loop (by means of calling definedOutside).
+/// - the op is not terminator.
+bool canBeHoisted(Operation* op,
+                  llvm::function_ref<bool(Value)> definedOutside) {
+  // Do not move terminators.
+  if (op->hasTrait<OpTrait::IsTerminator>()) return false;
+
+  // Walk the nested operations and check that all used values are either
+  // defined outside of the loop or in a nested region, but not at the level of
+  // the loop body.
+  auto walkFn = [&](Operation* child) {
+    for (Value operand : child->getOperands()) {
+      // Ignore values defined in a nested region.
+      if (op->isAncestor(operand.getParentRegion()->getParentOp())) continue;
+      if (!definedOutside(operand)) return WalkResult::interrupt();
+    }
+    return WalkResult::advance();
+  };
+  return !op->walk(walkFn).wasInterrupted();
+}
+
 LogicalResult cloneAndMoveTo(OpBuilder& b, scf::ParallelOp parallelOp,
                              ValueRange lower, ValueRange upper,
                              ValueRange step, Block* block) {
-  auto cloned = cast<scf::ParallelOp>(b.clone(*parallelOp.getOperation()));
+  scf::ParallelOp cloned =
+      cast<scf::ParallelOp>(b.clone(*parallelOp.getOperation()));
   cloned->setOperands(0, lower.size(), lower);
   cloned->setOperands(lower.size(), upper.size(), upper);
   cloned->setOperands(lower.size() + upper.size(), step.size(), step);
   cloned->moveBefore(block, block->begin());
 
-  // Move out memref.alloc/dealloc op from the loop op.
-  // They are used to manage the tile buffers. We can re-use these buffers cross
-  // iterations in the same thread.
-  SmallVector<Operation*> allocOps;
-  SmallVector<Operation*> deallocOps;
-  for (Operation& op : cloned.getLoopBody().front()) {
-    if (isa<memref::AllocOp>(&op)) {
-      allocOps.push_back(&op);
-    } else if (isa<memref::DeallocOp>(&op)) {
-      deallocOps.push_back(&op);
-    }
-  }
-  for (Operation* op : allocOps) {
-    op->moveBefore(cloned);
-  }
-  for (Operation* op : llvm::reverse(deallocOps)) {
-    op->moveAfter(cloned);
-  }
-
   if (cloned
           .walk([&](LoopLikeOpInterface loopLike) {
             moveLoopInvariantCode(loopLike);
@@ -258,7 +263,35 @@ LogicalResult cloneAndMoveTo(OpBuilder& b, scf::ParallelOp parallelOp,
     return failure();
   }
 
-  moveLoopInvariantCode(cloned);
+  // Move out memref.alloc/dealloc op from the loop op.
+  // They are used to manage the tile buffers. We can re-use these buffers cross
+  // iterations in the same thread.
+  auto definedOutside = [&](Value value) {
+    auto loopLike = dyn_cast<LoopLikeOpInterface>(cloned.getOperation());
+    assert(loopLike);
+    return loopLike.isDefinedOutsideOfLoop(value);
+  };
+
+  struct AllocDeallocPair {
+    memref::AllocOp allocOp;
+    memref::DeallocOp deallocOp;
+  };
+  DenseMap<Value, AllocDeallocPair> allocAndDeallocPairs;
+  for (Operation& op : cloned.getLoopBody().front()) {
+    if (auto allocOp = dyn_cast<memref::AllocOp>(&op)) {
+      if (canBeHoisted(&op, definedOutside))
+        allocAndDeallocPairs[op.getResult(0)].allocOp = allocOp;
+    } else if (auto deallocOp = dyn_cast<memref::DeallocOp>(&op)) {
+      auto it = allocAndDeallocPairs.find(op.getOperand(0));
+      if (it != allocAndDeallocPairs.end()) it->second.deallocOp = deallocOp;
+    }
+  }
+  for (auto& pair : allocAndDeallocPairs) {
+    if (!pair.second.allocOp || !pair.second.deallocOp) continue;
+    pair.second.allocOp->moveBefore(cloned);
+    pair.second.deallocOp->moveAfter(cloned);
+  }
+
   return success();
 }
 
diff --git a/tao_compiler/mlir/disc/transforms/fusion_utils.cc b/tao_compiler/mlir/disc/transforms/fusion_utils.cc
index e1dfe8c48a1..901749dfd65 100644
--- a/tao_compiler/mlir/disc/transforms/fusion_utils.cc
+++ b/tao_compiler/mlir/disc/transforms/fusion_utils.cc
@@ -311,6 +311,34 @@ std::string generateSignatureForFusion(FusionPattern& fusionPattern) {
     sig.push_back('_');
   }
 
+  if (fusionPattern.isTransformBasedFusion()) {
+    if (auto dotOp = dyn_cast_or_null<lmhlo::DotGeneralOp>(
+            fusionPattern.getDominantOp())) {
+      auto lhsTy = dotOp->getOperand(0).getType().cast<MemRefType>();
+      auto rhsTy = dotOp->getOperand(1).getType().cast<MemRefType>();
+      auto appendShapeToStr = [&](ArrayRef<int64_t> shape) {
+        for (const auto& en : llvm::enumerate(shape)) {
+          if (en.index() > 0) sig.append("x");
+          if (en.value() == ShapedType::kDynamicSize) {
+            sig.append("u");
+          } else {
+            sig.append(Twine(en.value()).str());
+          }
+        }
+      };
+      sig.append("_LS_");
+      appendShapeToStr(lhsTy.getShape());
+      sig.append("_RS_");
+      appendShapeToStr(rhsTy.getShape());
+      auto dimNumbers = dotOp.getDotDimensionNumbers();
+      sig.append("_LC_");
+      appendShapeToStr(dimNumbers.getLhsContractingDimensions());
+      sig.append("_RC_");
+      appendShapeToStr(dimNumbers.getRhsContractingDimensions());
+      sig.append("_");
+    }
+  }
+
   sig.append(
       ("_" + Twine(num_ops) + "_" + Twine(fusionPattern.getResults().size()))
           .str());


Test patch:
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/convert-foreach-thread-to-parallel.mlir b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/convert-foreach-thread-to-parallel.mlir
new file mode 100644
index 00000000000..5ab44c34ea3
--- /dev/null
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/convert-foreach-thread-to-parallel.mlir
@@ -0,0 +1,35 @@
+// RUN: disc-opt --disc-convert-foreach-thread-op-to-parallel-op -split-input-file %s | FileCheck %s
+
+#map = affine_map<()[s0] -> (s0 ceildiv 6)>
+#map1 = affine_map<()[s0] -> (s0 ceildiv 16)>
+#map2 = affine_map<(d0)[s0] -> (d0 * -6 + s0, 6)>
+#map3 = affine_map<(d0)[s0] -> (d0 * -16 + s0, 16)>
+#map4 = affine_map<(d0) -> (d0 * 6)>
+#map5 = affine_map<(d0) -> (d0 * 16)>
+module {
+  // CHECK-LABEL: @matmul_nn
+  func.func @matmul_nn(%arg0: memref<?x?xf32>, %arg1: memref<?x?xf32>, %arg2: memref<?x?xf32>) -> memref<?x?xf32> attributes {test = true} {
+    %cst = arith.constant 0.000000e+00 : f32
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %dim = memref.dim %arg0, %c0 : memref<?x?xf32>
+    %dim_0 = memref.dim %arg1, %c1 : memref<?x?xf32>
+    %0 = affine.apply #map()[%dim]
+    %1 = affine.apply #map1()[%dim_0]
+    %dim_1 = memref.dim %arg0, %c1 : memref<?x?xf32>
+    // CHECK-NOT: scf.foreach_thread
+    // CHECK: scf.parallel
+    scf.foreach_thread (%arg3, %arg4) in (%0, %1) {
+      %2 = affine.min #map2(%arg3)[%dim]
+      %3 = affine.min #map3(%arg4)[%dim_0]
+      %4 = affine.apply #map4(%arg3)
+      %5 = affine.apply #map5(%arg4)
+      %subview = memref.subview %arg0[%4, 0] [%2, %dim_1] [1, 1] : memref<?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>>
+      %subview_2 = memref.subview %arg1[0, %5] [%dim_1, %3] [1, 1] : memref<?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>>
+      %subview_3 = memref.subview %arg2[%4, %5] [%2, %3] [1, 1] : memref<?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>>
+      linalg.fill ins(%cst : f32) outs(%subview_3 : memref<?x?xf32, strided<[?, 1], offset: ?>>)
+      linalg.matmul ins(%subview, %subview_2 : memref<?x?xf32, strided<[?, 1], offset: ?>>, memref<?x?xf32, strided<[?, 1], offset: ?>>) outs(%subview_3 : memref<?x?xf32, strided<[?, 1], offset: ?>>)
+    } {thread_dim_mapping = []}
+    return %arg2 : memref<?x?xf32>
+  }
+}
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/rewrite-payload-ir-for-ral.mlir b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/rewrite-payload-ir-for-ral.mlir
index eee3fbeb357..866aa419c58 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/rewrite-payload-ir-for-ral.mlir
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/rewrite-payload-ir-for-ral.mlir
@@ -18,8 +18,6 @@ module {
     %0 = affine.apply #map()[%dim]
     %1 = affine.apply #map1()[%dim_0]
     %dim_1 = memref.dim %arg0, %c1 : memref<?x?xf32>
-    // CHECK-NOT: scf.foreach_thread
-    // CHECK: scf.parallel
     scf.foreach_thread (%arg3, %arg4) in (%0, %1) {
       %2 = affine.min #map2(%arg3)[%dim]
       %3 = affine.min #map3(%arg4)[%dim_0]