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 7e13e19cf13..c721dbd91bb 100644
--- a/tao_compiler/mlir/disc/BUILD
+++ b/tao_compiler/mlir/disc/BUILD
@@ -2012,6 +2012,36 @@ disc_cc_library(
     alwayslink = 1,
 )
 
+cc_library(
+    name = "disc_transform_legalize_to_loop",
+    srcs = ["transforms/disc_transform_legalize_to_loop.cc"],
+    deps = [
+        ":codegen_utils",
+        ":disc_lhlo_elemental_utils",
+        ":disc_util",
+        ":fusion_utils",
+        ":pass_details",
+        ":placement_utils",
+        "//tensorflow/compiler/xla/mlir_hlo:lmhlo_pass_inc_gen",
+        "//tensorflow/compiler/xla/mlir_hlo:mlir_hlo",
+        "//tensorflow/compiler/xla/mlir_hlo:lhlo",
+        "//tensorflow/compiler/mlir/disc/tools/disc-transform:all_passes",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:AffineDialect",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:LinalgDialect",
+        "@llvm-project//mlir:MemRefDialect",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:SCFDialect",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TensorDialect",
+        "@llvm-project//mlir:Transforms",
+        "@llvm-project//mlir:VectorDialect",
+    ],
+    alwayslink = 1,
+)
+
 cc_library(
     name = "all_passes",
     hdrs = [
@@ -2067,6 +2097,7 @@ cc_library(
         ":disc_strip_shape_constraint_ops",
         ":disc_tensor_to_std",
         ":disc_to_llvm",
+        ":disc_transform_legalize_to_loop",
         ":disc_transpose_simplifier",
         ":dot_rewriter",
         ":element_type_converter",
diff --git a/tao_compiler/mlir/disc/disc_compiler.cc b/tao_compiler/mlir/disc/disc_compiler.cc
index 5afa15e027d..cb2041176ea 100644
--- a/tao_compiler/mlir/disc/disc_compiler.cc
+++ b/tao_compiler/mlir/disc/disc_compiler.cc
@@ -43,6 +43,7 @@ limitations under the License.
 #include "mlir/Dialect/Func/Transforms/Passes.h"
 #include "mlir/Dialect/GPU/Transforms/Passes.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/Dialect/Shape/Transforms/Passes.h"
@@ -444,6 +445,14 @@ LogicalResult LowerHLOToLLVM(ModuleOp m, const DISCLoweringOptions& options) {
     pm.addNestedPass<FuncOp>(disc_ral::createDiscStitchFusionPass());
   }
 
+  if (useTransformSchedule()) {
+    std::string transform_schedule;
+    tensorflow::ReadStringFromEnvVar("DISC_TRANSFORM_SCHEDULE_FILE", "",
+                                     &transform_schedule);
+    pm.addNestedPass<FuncOp>(disc_ral::createDiscTransformLegalizeToLoopPass(
+        gpu_enabled, transform_schedule));
+  }
+
   pm.addNestedPass<FuncOp>(createCanonicalizerPass());
   pm.addNestedPass<FuncOp>(createCSEPass());
   pm.addNestedPass<FuncOp>(createCanonicalizerPass());
@@ -604,6 +613,7 @@ LogicalResult LowerHLOToLLVM(ModuleOp m, const DISCLoweringOptions& options) {
 
   pm.addNestedPass<FuncOp>(disc_ral::createDiscRemoveDeadBufferPass());
 
+  pm.addNestedPass<FuncOp>(::mlir::createConvertLinalgToLoopsPass());
   pm.addPass(createConvertSCFToCFPass());
   pm.addPass(createLowerAffinePass());
   pm.addNestedPass<FuncOp>(createCanonicalizerPass());
diff --git a/tao_compiler/mlir/disc/disc_util.cc b/tao_compiler/mlir/disc/disc_util.cc
index 7bac3b62e29..921c1794b9d 100644
--- a/tao_compiler/mlir/disc/disc_util.cc
+++ b/tao_compiler/mlir/disc/disc_util.cc
@@ -218,6 +218,17 @@ bool useHorizontalFusion() {
   return enabled;
 }
 
+// Returns true if `DISC_ENABLE_TRANSFORM_SCHEDULE` is true
+bool useTransformSchedule() {
+  static bool enabled = []() {
+    bool enabled = false;
+    tensorflow::ReadBoolFromEnvVar("DISC_ENABLE_TRANSFORM_SCHEDULE", enabled,
+                                   &enabled);
+    return enabled;
+  }();
+  return enabled;
+}
+
 bool lowerFakeQuantToQuantAndDequant() {
   static bool enabled = []() {
     bool enabled = false;
diff --git a/tao_compiler/mlir/disc/disc_util.h b/tao_compiler/mlir/disc/disc_util.h
index 1b8b9980091..f7326a3daee 100644
--- a/tao_compiler/mlir/disc/disc_util.h
+++ b/tao_compiler/mlir/disc/disc_util.h
@@ -114,6 +114,9 @@ bool useShapeConstraintIR();
 // Returns true if `DISC_ENABLE_HORIZONTAL_FUSION` is true.
 bool useHorizontalFusion();
 
+// Returns true if `DISC_ENABLE_TRANSFORM_SCHEDULE` is true.
+bool useTransformSchedule();
+
 // Returns true if `DISC_FAKE_QUANT_TO_QUANT_AND_DEQUANT` is true
 bool lowerFakeQuantToQuantAndDequant();
 
diff --git a/tao_compiler/mlir/disc/transforms/PassDetail.h b/tao_compiler/mlir/disc/transforms/PassDetail.h
index f3759a530c4..fc4fefc4fdb 100644
--- a/tao_compiler/mlir/disc/transforms/PassDetail.h
+++ b/tao_compiler/mlir/disc/transforms/PassDetail.h
@@ -54,6 +54,18 @@ namespace arith {
 class ArithDialect;
 }
 
+namespace linalg {
+class LinalgDialect;
+}
+
+namespace vector {
+class VectorDialect;
+}
+
+namespace tensor {
+class TensorDialect;
+}
+
 namespace disc_ral {
 
 #define GEN_PASS_CLASSES
diff --git a/tao_compiler/mlir/disc/transforms/disc_passes.td b/tao_compiler/mlir/disc/transforms/disc_passes.td
index 925cafe14e7..2e30830d759 100644
--- a/tao_compiler/mlir/disc/transforms/disc_passes.td
+++ b/tao_compiler/mlir/disc/transforms/disc_passes.td
@@ -553,3 +553,25 @@ def DiscGPUSourceToLibPass : Pass<"disc-gpu-source-to-lib", "ModuleOp"> {
             /*default=*/"0", "gpu sm cc_minor.">,
   ];
 }
+
+def DiscTransformLegalizeToLoopPass : Pass<"disc-transform-legalize-to-loop", "mlir::func::FuncOp"> {
+  let summary = "Legalize transfrom fusion pattern to loop.";
+  let constructor = "createDiscTransformLegalizeToLoopPass()";
+  let options = [
+    Option<"gpuEnabled_", "gpu-enabled", "bool",
+            /*default=*/"false", "whether gpu is available.">,
+    Option<"transformFileName_", "transform-file-name", "std::string",
+           /*default=*/"\"\"", "Filename of the transform schedule.">,
+    Option<"enableExpensiveChecks_", "enable-expensive-checks", "bool",
+           /*default=*/"false", "perform expensive checks to better report errors in the transform IR.">,
+  ];
+  let dependentDialects = [
+      "AffineDialect",
+      "LLVM::LLVMDialect",
+      "arith::ArithDialect",
+      "linalg::LinalgDialect",
+      "scf::SCFDialect",
+      "tensor::TensorDialect",
+      "vector::VectorDialect",
+  ];
+}
diff --git a/tao_compiler/mlir/disc/transforms/disc_stitch_fusion.cc b/tao_compiler/mlir/disc/transforms/disc_stitch_fusion.cc
index c87c6cd23b4..afad7bef0dd 100644
--- a/tao_compiler/mlir/disc/transforms/disc_stitch_fusion.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_stitch_fusion.cc
@@ -39,26 +39,6 @@ namespace disc_ral {
 
 namespace {
 
-bool isOnGpu(Operation* op) {
-  auto attr =
-      op->getAttrOfType<StringAttr>(placement_utils::kDiscPlaceAssignment);
-  if (attr) {
-    return (attr.getValue() == placement_utils::kGpu);
-  }
-
-  // Fusion should have been set a placement attribute in lhlo_fusion pass.
-  // Leave a default placement here just in case there are fusion ops not
-  // generated by the fusion pass (e.g. in the file-check based ut).
-  if (isa<lmhlo::FusionOp>(op)) return true;
-
-  assert(isa<lmhlo::LmhloOp>(op) && "Unexpected usage of isOnGpu");
-  auto result_memref = cast<lmhlo::LmhloOp>(op).getResultBuffer();
-  auto memory_space =
-      result_memref.getType().cast<MemRefType>().getMemorySpace();
-  return memory_space && memory_space.isa<StringAttr>() &&
-         memory_space.cast<StringAttr>().getValue() == placement_utils::kGpu;
-}
-
 LogicalResult HandleCpuFusionOp(OpBuilder& b, Operation* fusion,
                                 ShapeAnalysis& shapeAnalysis) {
   auto fusionOp = cast<lmhlo::FusionOp>(fusion);
@@ -93,7 +73,7 @@ struct DiscStitchFusion : public DiscStitchFusionBase<DiscStitchFusion> {
     SmallVector<Operation*, 4> gpu_fusion_worklist;
     SmallVector<Operation*, 4> cpu_fusion_worklist;
     func.walk([&](lmhlo::FusionOp op) {
-      if (isOnGpu(op))
+      if (placement_utils::isGpuLmhlo(op))
         gpu_fusion_worklist.push_back(op);
       else
         cpu_fusion_worklist.push_back(op);
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
new file mode 100644
index 00000000000..ede2cb28c00
--- /dev/null
+++ b/tao_compiler/mlir/disc/transforms/disc_transform_legalize_to_loop.cc
@@ -0,0 +1,242 @@
+/* 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-hlo/Dialect/lhlo/IR/lhlo_ops.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/BlockAndValueMapping.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Location.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Support/LogicalResult.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/PassDetail.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/lhlo_elemental_utils.h"
+#include "tensorflow/compiler/mlir/disc/transforms/placement_utils.h"
+
+// This file implements logic to legalize transform fusion pattern to loop.
+
+namespace mlir {
+namespace disc_ral {
+
+namespace {
+
+BlockAndValueMapping buildValueMapping(FusionPattern& fusionPattern,
+                                       func::FuncOp funcOp,
+                                       bool reverse = false) {
+  int argPosition = 0;
+  BlockAndValueMapping mapping;
+  auto addMapping = [&](Value v) {
+    if (reverse)
+      mapping.map(funcOp.getArgument(argPosition++), v);
+    else
+      mapping.map(v, funcOp.getArgument(argPosition++));
+  };
+
+  for (Value v : fusionPattern.getOperands()) addMapping(v);
+  for (Value v : fusionPattern.getInternalResults()) addMapping(v);
+  for (Value v : fusionPattern.getResults()) addMapping(v);
+  return mapping;
+}
+
+struct DiscTransformLegalizeToLoopPass
+    : public DiscTransformLegalizeToLoopPassBase<
+          DiscTransformLegalizeToLoopPass> {
+  explicit DiscTransformLegalizeToLoopPass(bool gpuEnabled,
+                                           const std::string& transformFileName,
+                                           bool enableExpensiveChecks)
+      : DiscTransformLegalizeToLoopPassBase<DiscTransformLegalizeToLoopPass>::
+            DiscTransformLegalizeToLoopPassBase() {
+    this->gpuEnabled_ = gpuEnabled;
+    this->transformFileName_ = transformFileName;
+    this->enableExpensiveChecks_ = enableExpensiveChecks;
+  }
+
+  void runOnOperation() override;
+
+  LogicalResult handleCpuFusionOp(OpBuilder& b, Operation* fusion,
+                                  ShapeAnalysis& shapeAnalysis);
+  // Outlines the fusion op to a standalone module op.
+  LogicalResult outlineFusionOp(lmhlo::FusionOp fusionOp,
+                                FusionPattern& fusionPattern,
+                                OwningOpRef<ModuleOp>& m);
+  // Builds a nested pass pipeline to legalize the outlined fusion op.
+  LogicalResult runTransformPipeline(ModuleOp m);
+  // Inlines the lowered IR into the orignal module.
+  LogicalResult inlineTransformedModule(OpBuilder& b, Operation* fusion,
+                                        FusionPattern& fusionPattern,
+                                        ModuleOp m);
+};
+
+LogicalResult DiscTransformLegalizeToLoopPass::outlineFusionOp(
+    lmhlo::FusionOp fusionOp, FusionPattern& fusionPattern,
+    OwningOpRef<ModuleOp>& m) {
+  Location loc = fusionOp->getLoc();
+  m = ModuleOp::create(loc);
+  auto b = OpBuilder::atBlockBegin(m->getBody());
+
+  SmallVector<Type> inputTypes;
+  SmallVector<Type> outputTypes;
+  for (Value v : fusionPattern.getOperands()) {
+    inputTypes.push_back(v.getType());
+  }
+  for (Value v : fusionPattern.getInternalResults()) {
+    inputTypes.push_back(v.getType());
+  }
+  for (Value v : fusionPattern.getResults()) {
+    inputTypes.push_back(v.getType());
+    outputTypes.push_back(v.getType());
+  }
+  auto funcType =
+      FunctionType::get(fusionOp->getContext(), inputTypes, outputTypes);
+  auto funcOp =
+      b.create<func::FuncOp>(loc, getFusionFullName(fusionOp), funcType);
+  Block* entryBlock = funcOp.addEntryBlock();
+  auto mapping = buildValueMapping(fusionPattern, funcOp);
+  b.setInsertionPoint(entryBlock, entryBlock->begin());
+  b.clone(*fusionOp.getOperation(), mapping);
+  b.create<func::ReturnOp>(loc, funcOp.getArguments().drop_front(
+                                    fusionPattern.getOperands().size() +
+                                    fusionPattern.getInternalResults().size()));
+  return success();
+}
+
+LogicalResult DiscTransformLegalizeToLoopPass::runTransformPipeline(
+    ModuleOp m) {
+  PassManager pm(m.getContext());
+  pm.addPass(createDiscLegalizeLmhloFusionToLinalgPass());
+  pm.addPass(createDiscTransformDialectInterpreterPass(transformFileName_,
+                                                       enableExpensiveChecks_));
+  pm.addPass(createDiscRewritePayloadIRForRALPass(gpuEnabled_));
+  return pm.run(m);
+}
+
+LogicalResult DiscTransformLegalizeToLoopPass::inlineTransformedModule(
+    OpBuilder& b, Operation* fusion, FusionPattern& fusionPattern, ModuleOp m) {
+  auto funcOps = llvm::to_vector<4>(m.getOps<func::FuncOp>());
+  if (funcOps.size() != 1)
+    return fusion->emitError() << "failed to inline the transformed module "
+                                  "with multiple functions\n";
+  if (funcOps[0].getBody().getBlocks().size() != 1)
+    return fusion->emitError()
+           << "failed to inline the transformed func with multiple blocks\n";
+
+  Block* body = cast<lmhlo::FusionOp>(fusion).getBody();
+  for (auto& nestedOp : llvm::make_early_inc_range(body->without_terminator()))
+    nestedOp.erase();
+  b.setInsertionPoint(body, body->begin());
+  auto mapping = buildValueMapping(fusionPattern, funcOps[0], true);
+  for (auto& nestedOp : funcOps[0].getBody().front().without_terminator()) {
+    Operation* cloned = b.clone(nestedOp, mapping);
+  }
+  return success();
+}
+
+LogicalResult DiscTransformLegalizeToLoopPass::handleCpuFusionOp(
+    OpBuilder& b, Operation* fusion, ShapeAnalysis& shapeAnalysis) {
+  auto fusionOp = cast<lmhlo::FusionOp>(fusion);
+  assert(fusionOp);
+  FusionPattern fusionPattern(fusionOp, &shapeAnalysis);
+  if (!fusionPattern.isTransformBasedFusion()) {
+    // skip non-transform-based fusion pattern.
+    return success();
+  }
+
+  // 1, Outline the fusion to a standalone module op.
+  OwningOpRef<ModuleOp> m;
+  if (failed(outlineFusionOp(fusionOp, fusionPattern, m))) return failure();
+  LLVM_DEBUG(llvm::dbgs() << "After outline fusion op:\n" << m.get() << "\n");
+
+  // 2, TODO(wyzero): assign a default schedule for each pattern here.
+
+  // 3, Build a nested pass pipeline to legalize the outlined fusion op.
+  if (failed(runTransformPipeline(m.get()))) return failure();
+  LLVM_DEBUG(llvm::dbgs() << "After run transform pipeline:\n"
+                          << m.get() << "\n");
+
+  // 4, Inline the lowered IR into the orignal module.
+  if (failed(inlineTransformedModule(b, fusion, fusionPattern, m.get())))
+    return failure();
+  LLVM_DEBUG(llvm::dbgs() << "After inline transformed module:\n"
+                          << *fusion << "\n");
+  return success();
+}
+
+void DiscTransformLegalizeToLoopPass::runOnOperation() {
+  func::FuncOp func = getOperation();
+  OpBuilder b(func);
+  SmallVector<Operation*, 4> gpu_fusion_worklist;
+  SmallVector<Operation*, 4> cpu_fusion_worklist;
+  func.walk([&](lmhlo::FusionOp op) {
+    if (placement_utils::isGpuLmhlo(op))
+      gpu_fusion_worklist.push_back(op);
+    else
+      cpu_fusion_worklist.push_back(op);
+  });
+
+  std::unique_ptr<ShapeAnalysis> shapeAnalysisPtr;
+  if (!gpu_fusion_worklist.empty() || !cpu_fusion_worklist.empty()) {
+    if (useShapeConstraintIR()) {
+      shapeAnalysisPtr.reset(new ShapeConstraintIRAnalysis(func));
+    }
+  }
+
+  for (Operation* fusion : gpu_fusion_worklist) {
+    // TODO(disc): handling stitch fusion on GPU.
+    return signalPassFailure();
+  }
+
+  for (Operation* fusion : cpu_fusion_worklist) {
+    if (!useShapeConstraintIR()) {
+      // TODO: use FuncOp that contains `fusionOp` to construct
+      // shape-analysis, which will use global information for shape equality
+      // and decomposition analysis.
+      shapeAnalysisPtr.reset(new ShapeAnalysisDeprecated{fusion});
+    }
+
+    // Error message should be emitted inside the function.
+    if (failed(handleCpuFusionOp(b, fusion, *shapeAnalysisPtr))) {
+      return signalPassFailure();
+    }
+  }
+}
+
+}  // namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>>
+createDiscTransformLegalizeToLoopPass(bool gpuEnabled,
+                                      const std::string& filename,
+                                      bool expensiveCheck) {
+  return std::make_unique<DiscTransformLegalizeToLoopPass>(gpuEnabled, filename,
+                                                           expensiveCheck);
+}
+
+}  // namespace disc_ral
+}  // namespace mlir
diff --git a/tao_compiler/mlir/disc/transforms/lhlo_fusion.cc b/tao_compiler/mlir/disc/transforms/lhlo_fusion.cc
index 9e84c4c0ac0..0ec8a69d4df 100644
--- a/tao_compiler/mlir/disc/transforms/lhlo_fusion.cc
+++ b/tao_compiler/mlir/disc/transforms/lhlo_fusion.cc
@@ -555,7 +555,7 @@ struct DiscFusionPass : public DiscFusionPassBase<DiscFusionPass> {
         pipeline.emplace_back(
             makeNewPlacementAwareFusionStrategy(gpu_enabled_, "stitch"));
       } else {
-        if (isCompIntensFusionEnabled()) {
+        if (useTransformSchedule()) {
           pipeline.emplace_back(makeNewPlacementAwareFusionStrategy(
               gpu_enabled_, "transform_based"));
         }
diff --git a/tao_compiler/mlir/disc/transforms/passes.h b/tao_compiler/mlir/disc/transforms/passes.h
index f0b3f29e4e8..be942f124b1 100644
--- a/tao_compiler/mlir/disc/transforms/passes.h
+++ b/tao_compiler/mlir/disc/transforms/passes.h
@@ -284,6 +284,12 @@ createDiscCompIntensFusionToCUDASourcePass(int cc_major = 8, int cc_minor = 0);
 std::unique_ptr<OperationPass<ModuleOp>> createDiscGPUSourceToLibPass(
     int cc_major = 8, int cc_minor = 0);
 
+// Legalizes transform-based fusion pattern to loop.
+std::unique_ptr<OperationPass<func::FuncOp>>
+createDiscTransformLegalizeToLoopPass(bool gpuEnabled = false,
+                                      const std::string& filename = "",
+                                      bool expensiveCheck = false);
+
 }  // namespace disc_ral
 }  // namespace mlir
 
diff --git a/tao_compiler/mlir/disc/transforms/placement_utils.cc b/tao_compiler/mlir/disc/transforms/placement_utils.cc
index 329aebd205a..198a6dcd361 100644
--- a/tao_compiler/mlir/disc/transforms/placement_utils.cc
+++ b/tao_compiler/mlir/disc/transforms/placement_utils.cc
@@ -30,6 +30,26 @@ bool isGpuMhlo(Operation* op) {
   return ((!attr) || ((attr != nullptr) && (attr.getValue() == kGpu)));
 }
 
+bool isGpuLmhlo(Operation* op) {
+  auto attr =
+      op->getAttrOfType<StringAttr>(placement_utils::kDiscPlaceAssignment);
+  if (attr) {
+    return (attr.getValue() == placement_utils::kGpu);
+  }
+
+  // Fusion should have been set a placement attribute in lhlo_fusion pass.
+  // Leave a default placement here just in case there are fusion ops not
+  // generated by the fusion pass (e.g. in the file-check based ut).
+  if (isa<lmhlo::FusionOp>(op)) return true;
+
+  assert(isa<lmhlo::LmhloOp>(op) && "Unexpected usage of isOnGpu");
+  auto result_memref = cast<lmhlo::LmhloOp>(op).getResultBuffer();
+  auto memory_space =
+      result_memref.getType().cast<MemRefType>().getMemorySpace();
+  return memory_space && memory_space.isa<StringAttr>() &&
+         memory_space.cast<StringAttr>().getValue() == placement_utils::kGpu;
+}
+
 // Return true if the memref is on GPU
 bool isGpuMemRef(Value memref) {
   assert(memref.getType().isa<MemRefType>());
diff --git a/tao_compiler/mlir/disc/transforms/placement_utils.h b/tao_compiler/mlir/disc/transforms/placement_utils.h
index d73a237d7fe..1dd7aa7319a 100644
--- a/tao_compiler/mlir/disc/transforms/placement_utils.h
+++ b/tao_compiler/mlir/disc/transforms/placement_utils.h
@@ -60,13 +60,14 @@ LogicalResult parseEntryFunctionOutputPlacements(
     func::FuncOp main, bool default_on_gpu,
     SmallVectorImpl<PlacementType>& out);
 
-// If Op is placed on GPU
-bool OnGpu(Operation* op);
-
 // Return true if the Operation is placed on GPU
 // The typical usage is for mhlo ops on tensor layer
 bool isGpuMhlo(Operation* op);
 
+// Return true if the Operation is placed on GPU
+// The typical usage is for lmhlo ops.
+bool isGpuLmhlo(Operation* op);
+
 // Return true if the MemRef is placed on GPU
 bool isGpuMemRef(Value memref);
 


Test patch:
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/BUILD b/tao_compiler/mlir/disc/tests/disc-transform/BUILD
new file mode 100644
index 00000000000..c06d07c0a81
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/BUILD
@@ -0,0 +1,7 @@
+load("//tensorflow/compiler/mlir/disc/tests:glob_op_test.bzl", "glob_op_tests")
+
+glob_op_tests(
+    data = [],
+    driver = "@llvm-project//mlir:run_lit.sh",
+    test_file_exts = ["cc"],
+)
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_multithread_nn_d_f32.mlir b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_multithread_nn_d_f32.mlir
new file mode 100644
index 00000000000..c00b1704361
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_multithread_nn_d_f32.mlir
@@ -0,0 +1,9 @@
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 0 : i32}} {
+  func.func @main(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> (tensor<?x?xf32>) attributes {tf.entry_function = {inputs = "{{INPUTS}}", outputs = "{{OUTPUTS}}", input_placements="{{INPUT_PLACEMENTS}}", output_placements="{{OUTPUT_PLACEMENTS}}"}} {
+    %graph = tf_executor.graph {
+      %0:2 = tf_executor.island wraps "tf.MatMul"(%arg0, %arg1) {transpose_a = false, transpose_b = false} : (tensor<?x?xf32>, tensor<?x?xf32>) -> (tensor<?x?xf32>)
+      tf_executor.fetch %0 : tensor<?x?xf32>
+    }
+    return %graph : tensor<?x?xf32>
+  }
+}
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_multithread_nn_d_f32_schedule.mlir b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_multithread_nn_d_f32_schedule.mlir
new file mode 100644
index 00000000000..38d22143b22
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_multithread_nn_d_f32_schedule.mlir
@@ -0,0 +1,10 @@
+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 tile_sizes [6, 16]
+  transform.structured.fuse_into_containing_op %fill into %0#0
+
+  transform.disc.bufferize %arg1
+}
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_nn_d_f32.mlir b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_nn_d_f32.mlir
new file mode 100644
index 00000000000..c00b1704361
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_nn_d_f32.mlir
@@ -0,0 +1,9 @@
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 0 : i32}} {
+  func.func @main(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> (tensor<?x?xf32>) attributes {tf.entry_function = {inputs = "{{INPUTS}}", outputs = "{{OUTPUTS}}", input_placements="{{INPUT_PLACEMENTS}}", output_placements="{{OUTPUT_PLACEMENTS}}"}} {
+    %graph = tf_executor.graph {
+      %0:2 = tf_executor.island wraps "tf.MatMul"(%arg0, %arg1) {transpose_a = false, transpose_b = false} : (tensor<?x?xf32>, tensor<?x?xf32>) -> (tensor<?x?xf32>)
+      tf_executor.fetch %0 : tensor<?x?xf32>
+    }
+    return %graph : tensor<?x?xf32>
+  }
+}
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_nn_d_f32_schedule.mlir b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_nn_d_f32_schedule.mlir
new file mode 100644
index 00000000000..38d22143b22
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_nn_d_f32_schedule.mlir
@@ -0,0 +1,10 @@
+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 tile_sizes [6, 16]
+  transform.structured.fuse_into_containing_op %fill into %0#0
+
+  transform.disc.bufferize %arg1
+}
\ 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
new file mode 100644
index 00000000000..ffa6aa9fb80
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/matmul.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/tests/mlir_feature_test.h"
+#include "tensorflow/compiler/mlir/disc/tests/mlir_test.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace mlir_test {
+
+const std::string c_ft_path =
+    "tensorflow/compiler/mlir/disc/tests/disc-transform/data/";
+
+TEST(SimpleTest, MatMulF32_11x13x12) {
+  EnvSetting setting = {{"DISC_TRANSFORM_SCHEDULE_FILE",
+                         {c_ft_path + "matmul_nn_d_f32_schedule.mlir", false}},
+                        {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}}};
+  EnvSettingContext ctx(setting);
+  EXPECT_TRUE(feature_test_main(
+      /*mlir_file_path*/ c_ft_path + "matmul_nn_d_f32.mlir",
+      /*backend_types*/ kSupportedCPUBackendList,
+      /*num_inputs*/ 2,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"11x12xf32_X", "12x13xf32_X"},
+      /*output_descriptors*/ {"f32_X"}));
+}
+
+TEST(SimpleTest, MatMulF32_111x131x121) {
+  EnvSetting setting = {{"DISC_TRANSFORM_SCHEDULE_FILE",
+                         {c_ft_path + "matmul_nn_d_f32_schedule.mlir", false}},
+                        {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}}};
+  EnvSettingContext ctx(setting);
+  EXPECT_TRUE(feature_test_main(
+      /*mlir_file_path*/ c_ft_path + "matmul_nn_d_f32.mlir",
+      /*backend_types*/ kSupportedCPUBackendList,
+      /*num_inputs*/ 2,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"111x121xf32_X", "121x131xf32_X"},
+      /*output_descriptors*/ {"f32_X"}));
+}
+
+}  // 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
new file mode 100644
index 00000000000..7542c631591
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/matmul_multithread.cc
@@ -0,0 +1,45 @@
+/* Copyright 2022 The BladeDISC Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/disc/tests/mlir_feature_test.h"
+#include "tensorflow/compiler/mlir/disc/tests/mlir_test.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace mlir_test {
+
+const std::string c_ft_path =
+    "tensorflow/compiler/mlir/disc/tests/disc-transform/data/";
+
+static bool init_threads = []() {
+  setenv("OMP_NUM_THREADS", "8", 1);
+  return true;
+}();
+
+TEST(SimpleMTTest, MatMulF32_111x131x121_Thread_8) {
+  EnvSetting setting = {
+      {"DISC_TRANSFORM_SCHEDULE_FILE",
+       {c_ft_path + "matmul_multithread_nn_d_f32_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*/ kSupportedCPUBackendList,
+      /*num_inputs*/ 2,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"111x121xf32_X", "121x131xf32_X"},
+      /*output_descriptors*/ {"f32_X"}));
+}
+
+}  // namespace mlir_test
diff --git a/tao_compiler/mlir/disc/transforms/tests/cpu-only-lmhlo-fusion.mlir b/tao_compiler/mlir/disc/transforms/tests/cpu-only-lmhlo-fusion.mlir
index 02c60544242..db384f52cf8 100644
--- a/tao_compiler/mlir/disc/transforms/tests/cpu-only-lmhlo-fusion.mlir
+++ b/tao_compiler/mlir/disc/transforms/tests/cpu-only-lmhlo-fusion.mlir
@@ -1,5 +1,5 @@
 // RUN: disc-opt -pass-pipeline='func.func(disc-fusion{gpu-enabled=false fusion-strategy=base})' -split-input-file %s -o - | FileCheck %s --check-prefix=BASE
-// RUN: DISC_ENABLE_COMPUTE_INTENSIVE_FUSE=1 disc-opt -pass-pipeline='func.func(disc-fusion{gpu-enabled=false fusion-strategy=stitch})' -split-input-file %s -o - | FileCheck %s --check-prefix=TRANSFORM
+// RUN: DISC_ENABLE_TRANSFORM_SCHEDULE=1 disc-opt -pass-pipeline='func.func(disc-fusion{gpu-enabled=false fusion-strategy=stitch})' -split-input-file %s -o - | FileCheck %s --check-prefix=TRANSFORM
 
 // BASE-LABEL: @custom_call_op
 // BASE-SAME: (%[[ARG0:.*]]: memref<?x?xf32, "cpu">, %[[ARG1:.*]]: memref<?x?xf32, "cpu">, %[[ARG2:.*]]: memref<?x?xf32, "cpu">, %[[ARG3:.*]]: memref<?x?xf32, "cpu">) -> memref<?x?xf32, "cpu">