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/BUILD b/tao_compiler/mlir/disc/BUILD
index e887360ebbb..3acf84f450f 100644
--- a/tao_compiler/mlir/disc/BUILD
+++ b/tao_compiler/mlir/disc/BUILD
@@ -2132,6 +2132,28 @@ cc_library(
     alwayslink = 1,
 )
 
+cc_library(
+    name = "disc_duplicate_computation_after_fusion",
+    srcs = ["transforms/disc_duplicate_computation_after_fusion.cc"],
+    deps = [
+        ":disc_shape_optimization_utils",
+        ":disc_util",
+        ":fusion_utils",
+        ":pass_details",
+        ":placement_utils",
+        "@org_tensorflow//tensorflow/compiler/xla/mlir_hlo:lhlo",
+        "@llvm-project//llvm:Core",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:ShapeDialect",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+    ],
+    alwayslink = 1,
+)
+
 cc_library(
     name = "all_passes",
     hdrs = [
@@ -2156,6 +2178,7 @@ cc_library(
         ":disc_convert_fake_quant_op",
         ":disc_custom_call_rewriter",
         ":disc_cpu_map_parallel_loop",
+        ":disc_duplicate_computation_after_fusion",
         ":disc_duplicate_computation_for_fusion",
         ":disc_dynamic_slice_converter",
         ":disc_sparse_op_rewriter",
@@ -2171,7 +2194,7 @@ cc_library(
         ":disc_lower_tf",
         ":disc_lower_quantize_and_dequantize",
         ":disc_lower_to_library_call",
-	":disc_mhlo_cse",
+        ":disc_mhlo_cse",
         ":disc_math_approximation",
         ":disc_memref_canonicalizer",
         ":disc_outline_cpu_kernel",
diff --git a/tao_compiler/mlir/disc/disc_compiler.cc b/tao_compiler/mlir/disc/disc_compiler.cc
index d6d5b106065..4ec7561d1f0 100644
--- a/tao_compiler/mlir/disc/disc_compiler.cc
+++ b/tao_compiler/mlir/disc/disc_compiler.cc
@@ -486,6 +486,8 @@ LogicalResult LowerHLOToLLVM(ModuleOp m, const DISCLoweringOptions& options) {
         disc_ral::createDiscSpecializeFusionWithSpeculationPass(
             gpu_options.sm_count, gpu_options.max_threads_per_sm));
   } else {
+    pm.addNestedPass<FuncOp>(
+        disc_ral::createDiscDuplicateComputationAfterFusionPass());
     pm.addNestedPass<FuncOp>(
         disc_ral::createDiscSpecializeFusionWithSpeculationPass());
   }
diff --git a/tao_compiler/mlir/disc/disc_util.cc b/tao_compiler/mlir/disc/disc_util.cc
index dce9a5a0062..c9797b1d47d 100644
--- a/tao_compiler/mlir/disc/disc_util.cc
+++ b/tao_compiler/mlir/disc/disc_util.cc
@@ -229,6 +229,17 @@ bool useTransformSchedule() {
   return enabled;
 }
 
+// Returns true if `DISC_ENABLE_TRANSFORM_GEMM_EPILOGUE_FUSION` is true.
+bool useTransformGEMMEpilogueFusionSchedule() {
+  static bool enabled = []() {
+    bool enabled = true;
+    tensorflow::ReadBoolFromEnvVar("DISC_ENABLE_TRANSFORM_GEMM_EPILOGUE_FUSION",
+                                   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 db1efcce8e4..99c35e2b81c 100644
--- a/tao_compiler/mlir/disc/disc_util.h
+++ b/tao_compiler/mlir/disc/disc_util.h
@@ -126,6 +126,9 @@ bool useHorizontalFusion();
 // Returns true if `DISC_ENABLE_TRANSFORM_SCHEDULE` is true.
 bool useTransformSchedule();
 
+// Returns true if `DISC_ENABLE_TRANSFORM_GEMM_EPILOGUE_FUSION` is true.
+bool useTransformGEMMEpilogueFusionSchedule();
+
 // Returns true if `DISC_FAKE_QUANT_TO_QUANT_AND_DEQUANT` is true
 bool lowerFakeQuantToQuantAndDequant();
 
diff --git a/tao_compiler/mlir/disc/tools/disc-opt/disc-opt.cc b/tao_compiler/mlir/disc/tools/disc-opt/disc-opt.cc
index 878927a201b..6a134a6c954 100644
--- a/tao_compiler/mlir/disc/tools/disc-opt/disc-opt.cc
+++ b/tao_compiler/mlir/disc/tools/disc-opt/disc-opt.cc
@@ -70,6 +70,8 @@ int main(int argc, char** argv) {
       transform_ext::StructuredTransformOpsExtension>();
   mlir::disc_ral::disc_linalg_ext::registerTilingInterfaceExternalModels(
       registry);
+  mlir::disc_ral::disc_linalg_ext::
+      registerBufferizableOpInterfaceExternalModels(registry);
 
   return failed(mlir::MlirOptMain(argc, argv, "MLIR HLO pass driver\n",
                                   registry,
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.cc b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.cc
index b08b6a354d9..2a6e0ce80af 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.cc
@@ -22,6 +22,8 @@
 #include "mlir/Dialect/Affine/Utils.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Utils/Utils.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
@@ -832,6 +834,144 @@ void registerTilingInterfaceExternalModels(DialectRegistry& registry) {
   });
 }
 
+//===----------------------------------------------------------------------===//
+// External Model for implementing `BufferizableOpInterface` for `LinalgOp`s.
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+/// Generic conversion for any DestinationStyleOpInterface on tensors.
+static LogicalResult bufferizeDestinationStyleOpInterface(
+    RewriterBase& rewriter, DestinationStyleOpInterface op,
+    const bufferization::BufferizationOptions& options) {
+  // Take a guard before anything else.
+  OpBuilder::InsertionGuard g(rewriter);
+  rewriter.setInsertionPoint(op);
+
+  // Nothing to do. This op is already bufferized.
+  if (op.hasBufferSemantics()) return success();
+
+  // Ensure op has only tensors. Allow mixed tensor-buffer mode on a per-need
+  // basis.
+  if (!op.hasTensorSemantics())
+    return op->emitError() << "op does not have tensor semantics";
+
+  // New input operands for the cloned op.
+  SmallVector<Value> newInputBuffers;
+  newInputBuffers.reserve(op.getNumDpsInputs());
+  for (OpOperand* opOperand : op.getDpsInputOperands()) {
+    if (op.isScalar(opOperand)) {
+      newInputBuffers.push_back(opOperand->get());
+      continue;
+    }
+    FailureOr<Value> buffer = getBuffer(rewriter, opOperand->get(), options);
+    if (failed(buffer)) return failure();
+    newInputBuffers.push_back(*buffer);
+  }
+
+  // New output operands for the cloned op.
+  SmallVector<Value> newOutputBuffers;
+  for (OpResult opResult : op->getOpResults()) {
+    OpOperand* opOperand = op.getDpsInitOperand(opResult.getResultNumber());
+    FailureOr<Value> resultBuffer =
+        getBuffer(rewriter, opOperand->get(), options);
+    if (failed(resultBuffer)) return failure();
+    newOutputBuffers.push_back(*resultBuffer);
+  }
+
+  // Merge input/output operands.
+  SmallVector<Value> newOperands = newInputBuffers;
+  newOperands.append(newOutputBuffers.begin(), newOutputBuffers.end());
+
+  // Set insertion point now that potential alloc/dealloc are introduced.
+  rewriter.setInsertionPoint(op);
+  // Clone the op, but use the new operands. Move the existing block into the
+  // new op. Since the new op does not have any tensor results, it does not
+  // return anything.
+  assert(op->getNumRegions() == 1 && "expected that op has 1 region");
+  auto newOp = cast<DestinationStyleOpInterface>(op.cloneWithoutRegions(
+      rewriter, op.getLoc(), /*resultTypes=*/TypeRange{}, newOperands));
+  rewriter.inlineRegionBefore(op->getRegion(0), newOp->getRegion(0),
+                              newOp->getRegion(0).begin());
+
+  // Replace the results of the old op with the new output buffers.
+  bufferization::replaceOpWithBufferizedValues(rewriter, op, newOutputBuffers);
+
+  return success();
+}
+
+/// Bufferization of linalg.generic. Replace with a new linalg.generic that
+/// operates entirely on memrefs.
+template <typename OpTy>
+struct LinalgOpInterface
+    : public bufferization::BufferizableOpInterface::ExternalModel<
+          LinalgOpInterface<OpTy>, OpTy> {
+  bool bufferizesToMemoryRead(Operation* op, OpOperand& opOperand,
+                              const bufferization::AnalysisState& state) const {
+    // Operand is read if it is used in the computation.
+    auto genericOp = cast<linalg::LinalgOp>(op);
+    return genericOp.payloadUsesValueFromOperand(&opOperand);
+  }
+
+  bool bufferizesToMemoryWrite(
+      Operation* op, OpOperand& opOperand,
+      const bufferization::AnalysisState& state) const {
+    // Operand is written to if it has an aliasing OpResult.
+    auto bufferizableOp = cast<bufferization::BufferizableOpInterface>(op);
+    return !bufferizableOp.getAliasingOpResult(opOperand, state).empty();
+  }
+
+  SmallVector<OpOperand*> getAliasingOpOperand(
+      Operation* op, OpResult opResult,
+      const bufferization::AnalysisState& state) const {
+    auto genericOp = cast<DestinationStyleOpInterface>(op);
+
+    // The i-th OpResult may alias with the i-th "out" tensor.
+    return {genericOp.getDpsInitOperand(opResult.getResultNumber())};
+  }
+
+  SmallVector<OpResult> getAliasingOpResult(
+      Operation* op, OpOperand& opOperand,
+      const bufferization::AnalysisState& state) const {
+    auto genericOp = cast<DestinationStyleOpInterface>(op);
+
+    // The i-th "out" tensor may alias with the i-th OpResult.
+    if (genericOp.isDpsInit(&opOperand))
+      return {genericOp.getTiedOpResult(&opOperand)};
+    return {};
+  }
+
+  bufferization::BufferRelation bufferRelation(
+      Operation* op, OpResult opResult,
+      const bufferization::AnalysisState& state) const {
+    return bufferization::BufferRelation::Equivalent;
+  }
+
+  LogicalResult bufferize(
+      Operation* op, RewriterBase& rewriter,
+      const bufferization::BufferizationOptions& options) const {
+    return bufferizeDestinationStyleOpInterface(
+        rewriter, cast<DestinationStyleOpInterface>(op), options);
+  }
+};
+
+/// Helper structure that iterates over all LinalgOps in `OpTys` and registers
+/// the `BufferizableOpInterface` with each of them.
+template <typename... Ops>
+struct LinalgOpInterfaceHelper {
+  static void registerOpInterface(MLIRContext* ctx) {
+    (Ops::template attachInterface<LinalgOpInterface<Ops>>(*ctx), ...);
+  }
+};
+
+}  // namespace
+
+void registerBufferizableOpInterfaceExternalModels(DialectRegistry& registry) {
+  registry.addExtension(+[](MLIRContext* ctx, DISCLinalgExtDialect* dialect) {
+    LinalgOpInterfaceHelper<ConditionalGenericOp>::registerOpInterface(ctx);
+  });
+}
+
 //===----------------------------------------------------------------------===//
 // YieldOp
 //===----------------------------------------------------------------------===//
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.h b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.h
index d94a116cbc4..4c28153611a 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.h
+++ b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.h
@@ -59,6 +59,7 @@ SmallVector<T> interchange(ArrayRef<T> elements,
 }
 
 void registerTilingInterfaceExternalModels(DialectRegistry& registry);
+void registerBufferizableOpInterfaceExternalModels(DialectRegistry& registry);
 
 }  // namespace disc_linalg_ext
 }  // namespace disc_ral
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 2d96b944e15..c5eeb4aa419 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
@@ -30,12 +30,14 @@
 #include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
+#include "mlir/Dialect/Math/IR/Math.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/IR/Matchers.h"
 #include "mlir/Interfaces/DestinationStyleOpInterface.h"
+#include "mlir/Interfaces/VectorInterfaces.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Support/MathExtras.h"
 #include "mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.h"
@@ -613,6 +615,255 @@ struct FoldXferReadOfXferWriterWithSelectPattern
   }
 };
 
+Operation* getAutomaticAllocationScope(Operation* op) {
+  // Find the closest surrounding allocation scope that is not a known looping
+  // construct (putting alloca's in loops doesn't always lower to deallocation
+  // until the end of the loop).
+  Operation* scope = nullptr;
+  for (Operation* parent = op->getParentOp(); parent != nullptr;
+       parent = parent->getParentOp()) {
+    if (parent->hasTrait<OpTrait::AutomaticAllocationScope>()) scope = parent;
+    if (scope && !isa<scf::ForOp, AffineForOp>(scope)) break;
+  }
+  return scope;
+}
+
+/// Given two MemRefTypes `aT` and `bT`, return a MemRefType to which both can
+/// be cast. If the MemRefTypes don't have the same rank or are not strided,
+/// return null; otherwise:
+///   1. if `aT` and `bT` are cast-compatible, return `aT`.
+///   2. else return a new MemRefType obtained by iterating over the shape and
+///   strides and:
+///     a. keeping the ones that are static and equal across `aT` and `bT`.
+///     b. using a dynamic shape and/or stride for the dimensions that don't
+///        agree.
+MemRefType getCastCompatibleMemRefType(MemRefType aT, MemRefType bT) {
+  if (memref::CastOp::areCastCompatible(aT, bT)) return aT;
+  if (aT.getRank() != bT.getRank()) return MemRefType();
+  int64_t aOffset, bOffset;
+  SmallVector<int64_t, 4> aStrides, bStrides;
+  if (failed(getStridesAndOffset(aT, aStrides, aOffset)) ||
+      failed(getStridesAndOffset(bT, bStrides, bOffset)) ||
+      aStrides.size() != bStrides.size())
+    return MemRefType();
+
+  ArrayRef<int64_t> aShape = aT.getShape(), bShape = bT.getShape();
+  int64_t resOffset;
+  SmallVector<int64_t, 4> resShape(aT.getRank(), 0),
+      resStrides(bT.getRank(), 0);
+  for (int64_t idx = 0, e = aT.getRank(); idx < e; ++idx) {
+    resShape[idx] =
+        (aShape[idx] == bShape[idx]) ? aShape[idx] : ShapedType::kDynamicSize;
+    resStrides[idx] = (aStrides[idx] == bStrides[idx])
+                          ? aStrides[idx]
+                          : ShapedType::kDynamicStrideOrOffset;
+  }
+  resOffset =
+      (aOffset == bOffset) ? aOffset : ShapedType::kDynamicStrideOrOffset;
+  return MemRefType::get(
+      resShape, aT.getElementType(),
+      StridedLayoutAttr::get(aT.getContext(), resOffset, resStrides));
+}
+
+/// convert:
+///   %0 = vector.transfer_read %src, %cst : vector<8x12xf32>
+///   %1 = arith.select %pred, %0, %vector_cst : vector<8x12xf32>
+///   use(%1)
+/// to:
+///   %0 = memref.alloca() : memref<8x12xf32>
+///   %buffer = scf.if %pred -> memref<?x?xf32> {
+///     scf.yield %src
+///   else {
+///     scf.yield %0
+///   }
+///   %1 = vector.transfer_read %buffer, %cst : vector<8x12xf32>
+///   use(%1)
+///
+/// Note that we need this pattern because we observed that lowered assemble
+/// code of arith.select is very slow.
+struct SelectOfXferReadAndConstToFillPattern
+    : public OpRewritePattern<arith::SelectOp> {
+  using OpRewritePattern<arith::SelectOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(arith::SelectOp selectOp,
+                                PatternRewriter& rewriter) const override {
+    auto vectorTy = selectOp->getResult(0).getType().dyn_cast<VectorType>();
+    if (!vectorTy) return failure();
+
+    // find const op and xfer read op
+    auto xferReadOp =
+        selectOp->getOperand(1).getDefiningOp<vector::TransferReadOp>();
+    auto vectorConstOp =
+        selectOp->getOperand(2).getDefiningOp<arith::ConstantOp>();
+    bool thenIsConst = false;
+    if (!vectorConstOp || !xferReadOp) {
+      xferReadOp =
+          selectOp->getOperand(2).getDefiningOp<vector::TransferReadOp>();
+      vectorConstOp =
+          selectOp->getOperand(1).getDefiningOp<arith::ConstantOp>();
+      thenIsConst = true;
+    }
+    if (!vectorConstOp || !xferReadOp) return failure();
+
+    // check type is compatible
+    auto srcTy = xferReadOp.getSource().getType().dyn_cast<MemRefType>();
+    if (!srcTy) return failure();
+    auto memrefTy =
+        MemRefType::get(vectorTy.getShape(), vectorTy.getElementType());
+    auto castTy = getCastCompatibleMemRefType(srcTy, memrefTy);
+    if (!castTy) return failure();
+
+    auto denseAttr = vectorConstOp.getValue().cast<DenseElementsAttr>();
+    if (!denseAttr.isSplat()) return failure();
+
+    // check the padding value for xfer read op is a const
+    auto scalarConstOp =
+        xferReadOp.getPadding().getDefiningOp<arith::ConstantOp>();
+    if (!scalarConstOp) return failure();
+
+    if (*denseAttr.getValues<Attribute>().begin() != scalarConstOp.getValue())
+      return failure();
+
+    Operation* scope = getAutomaticAllocationScope(selectOp);
+    if (!scope) return failure();
+
+    Location loc = selectOp.getLoc();
+    RewriterBase::InsertionGuard guard(rewriter);
+    rewriter.setInsertionPointToStart(&scope->getRegion(0).front());
+    Value newInitBuffer = rewriter.create<memref::AllocaOp>(
+        loc, memrefTy, ValueRange{}, rewriter.getI64IntegerAttr(64));
+    rewriter.create<linalg::FillOp>(loc, ValueRange{xferReadOp.getPadding()},
+                                    ValueRange{newInitBuffer});
+
+    rewriter.setInsertionPoint(selectOp.getOperation());
+    auto ifOp = rewriter.create<scf::IfOp>(loc, TypeRange{castTy},
+                                           selectOp->getOperand(0), true);
+    rewriter.setInsertionPointToStart(ifOp.thenBlock());
+    Value thenResult =
+        rewriter
+            .create<memref::CastOp>(
+                loc, castTy,
+                thenIsConst ? newInitBuffer : xferReadOp.getSource())
+            ->getResult(0);
+    rewriter.create<scf::YieldOp>(loc, thenResult);
+    rewriter.setInsertionPointToStart(ifOp.elseBlock());
+    Value elseResult =
+        rewriter
+            .create<memref::CastOp>(
+                loc, castTy,
+                !thenIsConst ? newInitBuffer : xferReadOp.getSource())
+            ->getResult(0);
+    rewriter.create<scf::YieldOp>(loc, elseResult);
+
+    rewriter.setInsertionPointAfter(ifOp);
+    BlockAndValueMapping mapping;
+    mapping.map(xferReadOp.getSource(), ifOp->getResult(0));
+    auto clonedXferReadOp = rewriter.clone(*xferReadOp.getOperation(), mapping);
+
+    rewriter.replaceOp(selectOp, clonedXferReadOp->getResults());
+    return success();
+  }
+};
+
+template <typename T>
+SmallVector<int64_t> getVectorSliceOpOffset(T op) {
+  SmallVector<int64_t> offsets;
+  if (op->hasAttr("offsets")) {
+    for (Attribute attr : op->template getAttrOfType<ArrayAttr>("offsets")) {
+      offsets.push_back(attr.cast<IntegerAttr>().getInt());
+    }
+  } else {
+    assert(op->hasAttr("position"));
+    for (Attribute attr : op->template getAttrOfType<ArrayAttr>("position")) {
+      offsets.push_back(attr.cast<IntegerAttr>().getInt());
+    }
+  }
+  return offsets;
+}
+
+bool checkMajorDimsAllOnes(VectorType ty) {
+  if (ty.getRank() <= 0) return false;
+  for (int64_t d = 0; d < ty.getRank() - 1; ++d) {
+    if (ty.getShape()[d] != 1) return false;
+  }
+  return true;
+}
+
+struct VectorExtractStridedSliceOfInsertStridedSliceFold
+    : public OpRewritePattern<vector::ExtractStridedSliceOp> {
+  using OpRewritePattern<vector::ExtractStridedSliceOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::ExtractStridedSliceOp extractStridedOp,
+                                PatternRewriter& rewriter) const override {
+    // check only trivial stirdes
+    if (extractStridedOp.hasNonUnitStrides()) return failure();
+
+    // check the result vector shape match pattern: vector<1x1x...xVecxelemType>
+    auto vectorTy = extractStridedOp->getResult(0).getType().cast<VectorType>();
+    if (!checkMajorDimsAllOnes(vectorTy)) return failure();
+
+    // check the offset for the last dim is divisible by the result vector size.
+    auto extractOffsets = getVectorSliceOpOffset(extractStridedOp);
+    if (extractOffsets.back() % vectorTy.getShape()[vectorTy.getRank() - 1] !=
+        0)
+      return failure();
+
+    Value targetValue;
+    SmallVector<int64_t> targetOffsets = extractOffsets;
+    Operation* sourceOp = extractStridedOp.getVector().getDefiningOp();
+    while (sourceOp && !targetValue) {
+      if (auto insertStridedOp =
+              dyn_cast<vector::InsertStridedSliceOp>(sourceOp)) {
+        if (insertStridedOp.hasNonUnitStrides()) return failure();
+        if (targetOffsets == getVectorSliceOpOffset(insertStridedOp)) {
+          auto ty = insertStridedOp.getSourceVectorType();
+          if (!checkMajorDimsAllOnes(ty)) return failure();
+          if (ty.getShape()[ty.getRank() - 1] !=
+              vectorTy.getShape()[vectorTy.getRank() - 1])
+            return failure();
+          targetValue = insertStridedOp.getSource();
+        } else {
+          sourceOp = insertStridedOp.getDest().getDefiningOp();
+        }
+      } else if (auto insertOp = dyn_cast<vector::InsertOp>(sourceOp)) {
+        bool match = true;
+        auto insertOffsets = getVectorSliceOpOffset(insertOp);
+        SmallVector<int64_t> newTargetOffsets;
+        for (int64_t i = 0; match && i < targetOffsets.size(); ++i) {
+          if (i < insertOffsets.size()) {
+            match &= (insertOffsets[i] == targetOffsets[i]);
+          } else {
+            newTargetOffsets.push_back(targetOffsets[i]);
+          }
+        }
+
+        if (match) {
+          sourceOp = insertOp.getSource().getDefiningOp();
+          std::swap(targetOffsets, newTargetOffsets);
+        } else {
+          sourceOp = insertOp.getDest().getDefiningOp();
+        }
+      } else if (auto extractOp = dyn_cast<vector::ExtractOp>(sourceOp)) {
+        auto newTargetOffsets = getVectorSliceOpOffset(extractOp);
+        newTargetOffsets.append(targetOffsets);
+        sourceOp = extractOp.getVector().getDefiningOp();
+        std::swap(newTargetOffsets, targetOffsets);
+      } else {
+        return failure();
+      }
+    }
+
+    if (!targetValue) return failure();
+    if (targetValue.getType() != vectorTy) {
+      targetValue = rewriter.create<vector::ShapeCastOp>(
+          extractStridedOp.getLoc(), vectorTy, targetValue);
+    }
+    rewriter.replaceOp(extractStridedOp, targetValue);
+
+    return success();
+  }
+};
+
 static void addAllRegisteredCanonicalizationPatterns(
     RewritePatternSet& patterns) {
   MLIRContext* ctx = patterns.getContext();
@@ -631,6 +882,8 @@ static void addAllRegisteredCanonicalizationPatterns(
   patterns.insert<FoldTensorExtractOfConstantWrapperPattern>(ctx);
   patterns.insert<FoldXferReadOfXferWriterPattern>(ctx);
   patterns.insert<FoldXferReadOfXferWriterWithSelectPattern>(ctx);
+  patterns.insert<SelectOfXferReadAndConstToFillPattern>(ctx);
+  patterns.insert<VectorExtractStridedSliceOfInsertStridedSliceFold>(ctx);
 }
 
 }  // namespace
@@ -1215,47 +1468,6 @@ DiagnosedSilenceableFailure LowerMultiLevelPackToLoopOp::applyToOne(
 // InlineReductionInitializer
 //===---------------------------------------------------------------------===//
 
-namespace {
-
-/// Given two MemRefTypes `aT` and `bT`, return a MemRefType to which both can
-/// be cast. If the MemRefTypes don't have the same rank or are not strided,
-/// return null; otherwise:
-///   1. if `aT` and `bT` are cast-compatible, return `aT`.
-///   2. else return a new MemRefType obtained by iterating over the shape and
-///   strides and:
-///     a. keeping the ones that are static and equal across `aT` and `bT`.
-///     b. using a dynamic shape and/or stride for the dimensions that don't
-///        agree.
-MemRefType getCastCompatibleMemRefType(MemRefType aT, MemRefType bT) {
-  if (memref::CastOp::areCastCompatible(aT, bT)) return aT;
-  if (aT.getRank() != bT.getRank()) return MemRefType();
-  int64_t aOffset, bOffset;
-  SmallVector<int64_t, 4> aStrides, bStrides;
-  if (failed(getStridesAndOffset(aT, aStrides, aOffset)) ||
-      failed(getStridesAndOffset(bT, bStrides, bOffset)) ||
-      aStrides.size() != bStrides.size())
-    return MemRefType();
-
-  ArrayRef<int64_t> aShape = aT.getShape(), bShape = bT.getShape();
-  int64_t resOffset;
-  SmallVector<int64_t, 4> resShape(aT.getRank(), 0),
-      resStrides(bT.getRank(), 0);
-  for (int64_t idx = 0, e = aT.getRank(); idx < e; ++idx) {
-    resShape[idx] =
-        (aShape[idx] == bShape[idx]) ? aShape[idx] : ShapedType::kDynamicSize;
-    resStrides[idx] = (aStrides[idx] == bStrides[idx])
-                          ? aStrides[idx]
-                          : ShapedType::kDynamicStrideOrOffset;
-  }
-  resOffset =
-      (aOffset == bOffset) ? aOffset : ShapedType::kDynamicStrideOrOffset;
-  return MemRefType::get(
-      resShape, aT.getElementType(),
-      StridedLayoutAttr::get(aT.getContext(), resOffset, resStrides));
-}
-
-}  // namespace
-
 void InlineReductionInitializerOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance>& effects) {
   transform::consumesHandle({this->getOperation()->getOperand(0)}, effects);
@@ -1556,6 +1768,114 @@ struct DecomposeVectorInOutsOfForOp : public OpRewritePattern<scf::ForOp> {
   vector::UnrollVectorOptions options_;
 };
 
+struct DecomposeVectorInOutsOfIfOp : public OpRewritePattern<scf::IfOp> {
+  explicit DecomposeVectorInOutsOfIfOp(
+      MLIRContext* context, const vector::UnrollVectorOptions& options,
+      PatternBenefit benefit = 1)
+      : mlir::OpRewritePattern<scf::IfOp>(context, benefit),
+        options_(options) {}
+
+  struct VectorDecomposeInfo {
+    VectorType dstVecType;
+    int64_t startIdx;
+    int64_t numSubVectors;
+    SmallVector<SmallVector<int64_t>> strides;
+    SmallVector<SmallVector<int64_t>> offsets;
+  };
+
+  LogicalResult matchAndRewrite(scf::IfOp ifOp,
+                                PatternRewriter& rewriter) const final {
+    Location loc = ifOp.getLoc();
+    if (failed(options_.filterConstraint(ifOp))) return failure();
+    auto maybeTargetShape = options_.nativeShape(ifOp);
+    if (!maybeTargetShape) return failure();
+    auto& targetShape = *maybeTargetShape;
+
+    int numNewOuts = 0;
+    VectorType targetVectorType;
+    DenseMap<int, VectorDecomposeInfo> candidateValueMap;
+    SmallVector<Type> newResultTypes;
+    for (const auto& en : llvm::enumerate(ifOp->getResults())) {
+      auto ty = en.value().getType().dyn_cast<VectorType>();
+      Optional<SmallVector<int64_t>> maybeShapeRatio;
+      if (!ty ||
+          !(maybeShapeRatio = computeShapeRatio(ty.getShape(), targetShape)) ||
+          llvm::all_of(*maybeShapeRatio, [](int64_t v) { return v == 1; })) {
+        ++numNewOuts;
+        newResultTypes.push_back(en.value().getType());
+        continue;
+      }
+      // TODO(wyzero): support multiple iter args with different vector type.
+      if (targetVectorType && targetVectorType != ty) return failure();
+      targetVectorType = ty;
+      auto& item = candidateValueMap[en.index()];
+      item.dstVecType = ty;
+      item.numSubVectors = computeMaxLinearIndex(*maybeShapeRatio);
+      item.startIdx = numNewOuts;
+      SmallVector<int64_t> strides(targetShape.size(), 1);
+      for (int i = 0; i < item.numSubVectors; ++i) {
+        auto offsets = getVectorOffset(ty.getShape(), targetShape, i);
+        item.strides.push_back(strides);
+        item.offsets.push_back(offsets);
+        newResultTypes.push_back(
+            VectorType::get(targetShape, ty.getElementType()));
+        ++numNewOuts;
+      }
+    }
+
+    if (candidateValueMap.empty() || ifOp.getElseRegion().empty())
+      return failure();
+    scf::IfOp newIfOp = rewriter.create<scf::IfOp>(loc, newResultTypes,
+                                                   ifOp.getCondition(), true);
+
+    auto cloneAndUpdateRegion = [&](Region& src, Region& dst) {
+      dst.takeBody(src);
+      SmallVector<Value> newResults;
+      auto yieldOp = dst.front().getTerminator();
+      rewriter.setInsertionPoint(yieldOp);
+      for (auto [idx, val] : llvm::enumerate(yieldOp->getOperands())) {
+        auto it = candidateValueMap.find(idx);
+        if (it == candidateValueMap.end()) newResults.push_back(val);
+        auto& item = it->second;
+        for (int i = 0; i < item.numSubVectors; ++i) {
+          newResults.push_back(rewriter.create<vector::ExtractStridedSliceOp>(
+              loc, val, item.offsets[i], targetShape, item.strides[i]));
+        }
+      }
+      yieldOp->setOperands(newResults);
+    };
+
+    cloneAndUpdateRegion(ifOp.getThenRegion(), newIfOp.getThenRegion());
+    cloneAndUpdateRegion(ifOp.getElseRegion(), newIfOp.getElseRegion());
+
+    // merge the return values of new if op.
+    rewriter.setInsertionPointAfter(newIfOp);
+    size_t resultIdx = 0;
+    SmallVector<Value> newResults;
+    for (auto [idx, val] : llvm::enumerate(ifOp->getResults())) {
+      auto it = candidateValueMap.find(idx);
+      if (it == candidateValueMap.end()) {
+        newResults.push_back(newIfOp->getResult(resultIdx++));
+        continue;
+      }
+      auto& item = it->second;
+      Value mergedVec = rewriter.create<arith::ConstantOp>(
+          loc, item.dstVecType, rewriter.getZeroAttr(item.dstVecType));
+      for (int subIdx = 0; subIdx < item.numSubVectors; ++subIdx) {
+        mergedVec = rewriter.create<vector::InsertStridedSliceOp>(
+            loc, newIfOp->getResult(resultIdx++), mergedVec,
+            item.offsets[subIdx], item.strides[subIdx]);
+      }
+      newResults.push_back(mergedVec);
+    }
+    rewriter.replaceOp(ifOp, newResults);
+    return success();
+  }
+
+ private:
+  vector::UnrollVectorOptions options_;
+};
+
 }  // namespace
 
 void DecomposeVectorsOp::build(OpBuilder& builder, OperationState& result,
@@ -1591,8 +1911,8 @@ DiagnosedSilenceableFailure DecomposeVectorsOp::applyToOne(
                    this->getVectorSize() ==
                0;
   };
-  auto getVectorTypeOfForOp = [&](Operation* op) -> VectorType {
-    if (!isa<scf::ForOp>(op)) return nullptr;
+  auto getVectorTypeOfForOrIfOp = [&](Operation* op) -> VectorType {
+    if (!isa<scf::ForOp, scf::IfOp>(op)) return nullptr;
     VectorType vectorTy;
     for (auto ty : op->getResultTypes()) {
       if (!isTargetType(ty)) continue;
@@ -1602,19 +1922,20 @@ DiagnosedSilenceableFailure DecomposeVectorsOp::applyToOne(
     return vectorTy;
   };
   options.setFilterConstraint([&](Operation* op) {
-    if (getVectorTypeOfForOp(op)) return success();
+    if (getVectorTypeOfForOrIfOp(op)) return success();
     if (isa<vector::TransferReadOp, vector::TransferWriteOp>(op))
       return success();
     if (op->getNumResults() != 1) return failure();
     if (op->getDialect()->getTypeID() != TypeID::get<vector::VectorDialect>() &&
-        op->getDialect()->getTypeID() != TypeID::get<arith::ArithDialect>())
+        op->getDialect()->getTypeID() != TypeID::get<arith::ArithDialect>() &&
+        op->getDialect()->getTypeID() != TypeID::get<math::MathDialect>())
       return failure();
     return success(isTargetType(op->getResult(0).getType()));
   });
   vector::UnrollVectorOptions::NativeShapeFnType nativeShapeFn =
       [&](Operation* op) -> Optional<SmallVector<int64_t, 4>> {
     VectorType targetVectorTy;
-    if (auto vectorTy = getVectorTypeOfForOp(op)) {
+    if (auto vectorTy = getVectorTypeOfForOrIfOp(op)) {
       targetVectorTy = vectorTy;
     } else if (isa<vector::TransferWriteOp>(op)) {
       targetVectorTy = op->getOperand(0).getType().cast<VectorType>();
@@ -1628,6 +1949,7 @@ DiagnosedSilenceableFailure DecomposeVectorsOp::applyToOne(
   options.setNativeShapeFn(nativeShapeFn);
   vector::populateVectorUnrollPatterns(patterns, options);
   patterns.insert<DecomposeVectorInOutsOfForOp>(ctx, options);
+  patterns.insert<DecomposeVectorInOutsOfIfOp>(ctx, options);
 
   // some clean up patterns.
   vector::populateVectorToVectorCanonicalizationPatterns(patterns);
@@ -2732,6 +3054,80 @@ DiagnosedSilenceableFailure VectorizeConditionalGenericOp::applyToOne(
   return DiagnosedSilenceableFailure(success());
 }
 
+//===---------------------------------------------------------------------===//
+// SplitVectorTransferIntoFullAndPartialOp
+//===---------------------------------------------------------------------===//
+
+namespace {
+
+static LogicalResult splitFullAndPartialTransferPrecondition(
+    VectorTransferOpInterface xferOp) {
+  // TODO: support 0-d corner case.
+  if (xferOp.getTransferRank() == 0) return failure();
+
+  // TODO: expand support to these 2 cases.
+  if (!xferOp.permutation_map().isMinorIdentity()) return failure();
+  // Must have some out-of-bounds dimension to be a candidate for splitting.
+  if (!xferOp.hasOutOfBoundsDim()) return failure();
+  return success();
+}
+
+}  // namespace
+
+DiagnosedSilenceableFailure SplitVectorTransferIntoFullAndPartialOp::applyToOne(
+    Operation* target, SmallVectorImpl<Operation*>& results,
+    transform::TransformState& state) {
+  auto xferOp = dyn_cast<VectorTransferOpInterface>(target);
+  if (!xferOp) {
+    return mlir::emitDefiniteFailure(
+        target, "applies only to vector transfer like ops.");
+  }
+
+  if (failed(splitFullAndPartialTransferPrecondition(xferOp))) {
+    results.assign({xferOp.getOperation()});
+    return DiagnosedSilenceableFailure(success());
+  }
+
+  MLIRContext* ctx = target->getContext();
+  IRRewriter rewriter(ctx);
+  scf::IfOp ifOp;
+  vector::VectorTransformsOptions options;
+  options.setVectorTransferSplit(vector::VectorTransferSplit::LinalgCopy);
+
+  if (failed(vector::splitFullAndPartialTransfer(rewriter, xferOp, options,
+                                                 &ifOp))) {
+    return mlir::emitDefiniteFailure(target, "failed to split xfer ops.");
+  }
+  results.assign({ifOp.getOperation()});
+  return DiagnosedSilenceableFailure(success());
+}
+
+//===---------------------------------------------------------------------===//
+// LowerConditionalGenericOp
+//===---------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure LowerConditionalGenericOp::applyToOne(
+    Operation* target, SmallVectorImpl<Operation*>& results,
+    transform::TransformState& state) {
+  auto cgOp = dyn_cast<disc_linalg_ext::ConditionalGenericOp>(target);
+  if (!cgOp) {
+    return mlir::emitDefiniteFailure(
+        target, "applies only to disc_linalg_ext.conditional_generic.");
+  }
+
+  MLIRContext* ctx = target->getContext();
+  IRRewriter rewriter(ctx);
+  rewriter.setInsertionPoint(target);
+  scf::IfOp ifOp =
+      rewriter.create<scf::IfOp>(cgOp->getLoc(), cgOp->getOperand(0), false);
+  rewriter.setInsertionPointToStart(ifOp.thenBlock());
+  auto genericOp = convertConditionalGenericOpToGenericOp(
+      rewriter, dyn_cast<linalg::LinalgOp>(target));
+  rewriter.eraseOp(target);
+  results.assign({genericOp});
+  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 cec879893a2..5ee30c0167e 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
@@ -779,4 +779,88 @@ def VectorizeConditionalGenericOp : Op<Transform_Dialect, "disc.vectorize_condit
   }];
 }
 
+def SplitVectorTransferIntoFullAndPartialOp : Op<Transform_Dialect, "disc.split_vector_transfer_into_full_and_partial",
+    [FunctionalStyleTransformOpTrait,
+     MemoryEffectsOpInterface,
+     TransformEachOpTrait,
+     TransformOpInterface]> {
+  let description = [{
+    Split a vector.transfer operation into an in-bounds (i.e., no out-of-bounds masking) fastpath and a slowpath.
+
+    Returns the xfer op if no needs to split, otherwise returns the warpper if op for full version and partial version.
+
+    /// Example (a 2-D vector.transfer_read):
+    /// ```
+    ///    %1 = vector.transfer_read %0[...], %pad : memref<A...>, vector<...>
+    /// ```
+    /// is transformed into:
+    /// ```
+    ///    %1:3 = scf.if (%inBounds) {
+    ///      // fastpath, direct cast
+    ///      memref.cast %A: memref<A...> to compatibleMemRefType
+    ///      scf.yield %view : compatibleMemRefType, index, index
+    ///    } else {
+    ///      // slowpath, not in-bounds vector.transfer or linalg.copy.
+    ///      memref.cast %alloc: memref<B...> to compatibleMemRefType
+    ///      scf.yield %4 : compatibleMemRefType, index, index
+    //     }
+    ///    %0 = vector.transfer_read %1#0[%1#1, %1#2] {in_bounds = [true ... true]}
+    /// ```
+    /// where `alloc` is a top of the function alloca'ed buffer of one vector.
+
+  }];
+
+  let arguments = (ins PDL_Operation:$target);
+  let results = (outs PDL_Operation:$result);
+
+  let assemblyFormat = "$target attr-dict";
+  let cppNamespace = "::mlir::disc_ral::transform_dialect";
+
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::Operation *target,
+        ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+        ::mlir::transform::TransformState &state);
+  }];
+}
+
+def LowerConditionalGenericOp : Op<Transform_Dialect, "disc.lower_conditional_generic",
+    [FunctionalStyleTransformOpTrait,
+     MemoryEffectsOpInterface,
+     TransformEachOpTrait,
+     TransformOpInterface]> {
+  let description = [{
+    Lower the conditional_generic op to a scf.if + linalg.generic op.
+
+    Note the this transform primitive supposes the conditional_generic has been bufferized.
+
+    Returns the generated linalg.generic op.
+
+    /// Example:
+    /// ```
+    ///    %1 = disc_linalg_ext.conditional_generic ins(%pred, ...), outs(...)
+    /// ```
+    /// is transformed into:
+    /// ```
+    ///    scf.if (%pred) {
+    ///      linalg.generic ins(...) outs(...)
+    ///      scf.yield %view
+    ///    }
+    /// ```
+  }];
+
+  let arguments = (ins PDL_Operation:$target);
+  let results = (outs PDL_Operation:$result);
+
+  let assemblyFormat = "$target attr-dict";
+  let cppNamespace = "::mlir::disc_ral::transform_dialect";
+
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::Operation *target,
+        ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+        ::mlir::transform::TransformState &state);
+  }];
+}
+
 #endif // DISC_TRANSFORM_OPS_EXT
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/transform_dialect_interpreter.cc b/tao_compiler/mlir/disc/tools/disc-transform/transforms/transform_dialect_interpreter.cc
index d16d00ef0d4..96bef12901a 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/transforms/transform_dialect_interpreter.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/transform_dialect_interpreter.cc
@@ -92,6 +92,7 @@ void addTransformDialectDependentDialects(DialectRegistry& registry) {
   scf::registerTransformDialectExtension(registry);
   registerTransformDialectCommonExtension(registry);
   disc_linalg_ext::registerTilingInterfaceExternalModels(registry);
+  disc_linalg_ext::registerBufferizableOpInterfaceExternalModels(registry);
 }
 
 namespace {
diff --git a/tao_compiler/mlir/disc/transforms/disc_duplicate_computation_after_fusion.cc b/tao_compiler/mlir/disc/transforms/disc_duplicate_computation_after_fusion.cc
new file mode 100644
index 00000000000..e7dcaf71bfe
--- /dev/null
+++ b/tao_compiler/mlir/disc/transforms/disc_duplicate_computation_after_fusion.cc
@@ -0,0 +1,115 @@
+/* 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 "mlir/IR/MLIRContext.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "mlir/disc/disc_util.h"
+#include "mlir/disc/transforms/PassDetail.h"
+#include "mlir/disc/transforms/fusion_utils.h"
+
+// This file implements the logic to duplicate some lmhlo operations in order
+// to enable more opportunities for fusion and reduce memory footprint.
+
+namespace mlir {
+namespace disc_ral {
+
+namespace {
+
+using func::FuncOp;
+using namespace lmhlo;
+
+lmhlo::ConstantOp getConstProducer(Value value) {
+  for (Operation* user : value.getUsers()) {
+    auto constOp = dyn_cast<lmhlo::ConstantOp>(user);
+    if (constOp) return constOp;
+  }
+  return nullptr;
+}
+
+struct DiscDuplicateComputationAfterFusionPass
+    : public DiscDuplicateComputationAfterFusionPassBase<
+          DiscDuplicateComputationAfterFusionPass> {
+  using DiscDuplicateComputationAfterFusionPassBase<
+      DiscDuplicateComputationAfterFusionPass>::
+      DiscDuplicateComputationAfterFusionPassBase;
+
+  void runOnOperation() override {
+    FuncOp func = getOperation();
+    // skip shape constraint graph
+    if (func.getName() == SymbolicDimMgr::getShapeConstraintGraphFunctionName())
+      return;
+    if (useTransformSchedule()) {
+      if (failed(duplicateForTransformBaseFusion(func))) {
+        return signalPassFailure();
+      }
+    }
+  }
+
+  LogicalResult duplicateForTransformBaseFusion(FuncOp func);
+};
+
+LogicalResult
+DiscDuplicateComputationAfterFusionPass::duplicateForTransformBaseFusion(
+    FuncOp func) {
+  SmallVector<lmhlo::FusionOp> fusionOps;
+  func->walk([&](lmhlo::FusionOp fusionOp) { fusionOps.push_back(fusionOp); });
+
+  std::unique_ptr<ShapeAnalysis> shapeAnalysisPtr;
+  if (useShapeConstraintIR()) {
+    shapeAnalysisPtr.reset(new ShapeConstraintIRAnalysis(func));
+  } else {
+    shapeAnalysisPtr.reset(new ShapeAnalysisDeprecated{func});
+    if (failed(static_cast<ShapeAnalysisDeprecated*>(shapeAnalysisPtr.get())
+                   ->run())) {
+      return failure();
+    }
+  }
+  for (FusionOp fusionOp : fusionOps) {
+    FusionPattern fusionPattern(fusionOp, shapeAnalysisPtr.get());
+    if (!fusionPattern.isTransformBasedFusion()) continue;
+    for (Value operand : fusionPattern.getOperands()) {
+      auto constOp = getConstProducer(operand);
+      if (!constOp ||
+          constOp.getValue().cast<ElementsAttr>().getNumElements() != 1)
+        continue;
+      OpBuilder b(fusionOp);
+      auto newBuffer = b.create<memref::AllocOp>(
+          constOp->getLoc(),
+          constOp->getOperand(0).getType().cast<MemRefType>());
+      b.setInsertionPointToStart(fusionOp.getBody());
+      auto newConstOp = b.clone(*constOp);
+      newConstOp->setOperand(0, newBuffer);
+      for (Operation* op : fusionPattern.getOpList()) {
+        auto operands = op->getOperands();
+        if (llvm::find(operands, constOp->getOperand(0)) == operands.end())
+          continue;
+        op->replaceUsesOfWith(constOp->getOperand(0), newBuffer);
+      }
+    }
+  }
+
+  return success();
+}
+
+}  // namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>>
+createDiscDuplicateComputationAfterFusionPass() {
+  return std::make_unique<DiscDuplicateComputationAfterFusionPass>();
+}
+
+}  // namespace disc_ral
+}  // namespace mlir
diff --git a/tao_compiler/mlir/disc/transforms/disc_passes.td b/tao_compiler/mlir/disc/transforms/disc_passes.td
index 75e65f8a962..62d256f06aa 100644
--- a/tao_compiler/mlir/disc/transforms/disc_passes.td
+++ b/tao_compiler/mlir/disc/transforms/disc_passes.td
@@ -601,3 +601,8 @@ def DiscSparseOpRewriterPass : Pass<"disc-sparse-op-rewriter", "mlir::func::Func
   let summary = "Rewrite sparse op for possible fusion";
   let constructor = "createDiscSparseOpRewriterPass()";
 }
+
+def DiscDuplicateComputationAfterFusionPass : Pass<"disc-duplicate-computation-after-fusion", "mlir::func::FuncOp"> {
+  let summary = "Duplicate and fuse some computations into the consumer fusion op in order to remove memory footprint.";
+  let constructor = "createDiscDuplicateComputationAfterFusionPass()";
+}
diff --git a/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc b/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc
index aace3f60e1a..bd0cb9c7611 100644
--- a/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc
@@ -316,6 +316,21 @@ buildVectorizeConditionalGenericOp(OpBuilder& b, Location& loc, Value target) {
       loc, pdlType, target);
 }
 
+transform_dialect::SplitVectorTransferIntoFullAndPartialOp
+buildSplitVectorTransferIntoFullAndPartialOp(OpBuilder& b, Location& loc,
+                                             Value target) {
+  auto pdlType = pdl::OperationType::get(b.getContext());
+  return b.create<transform_dialect::SplitVectorTransferIntoFullAndPartialOp>(
+      loc, pdlType, target);
+}
+
+transform_dialect::LowerConditionalGenericOp buildLowerConditionalGenericOp(
+    OpBuilder& b, Location& loc, Value target) {
+  auto pdlType = pdl::OperationType::get(b.getContext());
+  return b.create<transform_dialect::LowerConditionalGenericOp>(loc, pdlType,
+                                                                target);
+}
+
 class ParsedFromFileScheduleFactory : public ScheduleFactoryWithNoGuard {
  public:
   explicit ParsedFromFileScheduleFactory(int64_t id, PatternKind kind,
@@ -1242,28 +1257,16 @@ LogicalResult Aarch64GEMMLargeKScheduleWithEpilogueFactory::assignSchedule(
     fill = buildDISCFuseIntoContainingOp(b, loc, fill, loopN0);
   }
 
-  auto tileM1 = buildTileOp(b, loc, rootHandle, {M1, 0}, {0, 1});
-  rootHandle = tileM1->getResult(0);
+  auto tileM1 = buildTileOp(b, loc, matmul, {M1, 0}, {0, 1});
+  matmul = tileM1->getResult(0);
   Value loopM1 = tileM1->getResult(1);
   variant = buildRunCanonicalizer(b, loc, variant);
-  matmul = buildDISCFuseIntoContainingOp(b, loc, matmul, loopM1);
   fill = buildDISCFuseIntoContainingOp(b, loc, fill, loopM1);
 
-  auto tileN1 = buildTileOp(b, loc, rootHandle, {0, N1}, {0, 1});
-  rootHandle = tileN1->getResult(0);
+  auto tileN1 = buildTileOp(b, loc, matmul, {0, N1}, {0, 1});
+  matmul = tileN1->getResult(0);
   Value loopN1 = tileN1->getResult(1);
   variant = buildRunCanonicalizer(b, loc, variant);
-  // pad root ops to register level size (static size).
-  int64_t numOperandsUpperBound = fusionPattern.getOperands().size() +
-                                  fusionPattern.getResults().size() +
-                                  fusionPattern.getInternalResults().size() + 2;
-  SmallVector<Type> paddingTypes;
-  if (!k0Skipped) paddingTypes.push_back(b.getIntegerType(1));
-  rootHandle = buildPadOp(b, loc, rootHandle, {0, 1}, numOperandsUpperBound,
-                          paddingTypes);
-  auto padForRootOp = buildMatchOp(b, loc, variant, {"tensor.pad"});
-  buildReplaceConstPaddingValueOp(b, loc, padForRootOp, "kAny");
-  matmul = buildDISCFuseIntoContainingOp(b, loc, matmul, loopN1);
   fill = buildDISCFuseIntoContainingOp(b, loc, fill, loopN1);
 
   auto tileK1 = buildTileOp(b, loc, matmul, {0, 0, K1}, {0, 1, 2});
@@ -1273,6 +1276,8 @@ LogicalResult Aarch64GEMMLargeKScheduleWithEpilogueFactory::assignSchedule(
   buildFoldProducerExtractSlice(b, loc, weightSlice, 2);
   // pad to match the requirement of hardware vector/tensor instruction.
   matmul = buildPadOp(b, loc, matmul, {0, 1, 2}, 3);
+  SmallVector<Type> paddingTypes;
+  if (!k0Skipped) paddingTypes.push_back(b.getIntegerType(1));
   fill = buildPadOp(b, loc, fill, {0, 1}, 3, paddingTypes);
   Value padForInput = buildGetProducerOfOperand(b, loc, matmul, 0);
   Value padForWeight = buildGetProducerOfOperand(b, loc, matmul, 1);
@@ -1328,24 +1333,31 @@ LogicalResult Aarch64GEMMLargeKScheduleWithEpilogueFactory::assignSchedule(
   buildLowerMultiLevelPackToLoop(b, loc, multiLevelPackOps);
   variant = buildRunCanonicalizer(b, loc, variant);
 
-  Value conditionalOps =
-      buildMatchOp(b, loc, variant, {"disc_linalg_ext.conditional_generic"});
-  buildVectorizeConditionalGenericOp(b, loc, conditionalOps);
+  Value fillConditionalOps = buildMatchOp(
+      b, loc, variant, {"disc_linalg_ext.conditional_generic"}, nameMap[dotOp]);
+  buildVectorizeConditionalGenericOp(b, loc, fillConditionalOps);
   variant = buildRunCanonicalizer(b, loc, variant);
   Value func = buildMatchOp(b, loc, variant, {"func.func"});
   buildVectorize(b, loc, func, true);
   variant = buildRunCanonicalizer(b, loc, variant);
-  auto placeholderOps = buildMatchOp(
-      b, loc, variant, {"disc_linalg_ext.padding_value_placeholder"}, {});
-  buildConvertPaddingPlaceholderToConstOp(b, loc, placeholderOps);
-  variant = buildRunCanonicalizer(b, loc, variant);
 
   variant = buildDISCBufferize(b, loc, variant);
+  variant = buildRunCanonicalizer(b, loc, variant);
+
+  Value conditionalOps =
+      buildMatchOp(b, loc, variant, {"disc_linalg_ext.conditional_generic"});
+  buildLowerConditionalGenericOp(b, loc, conditionalOps);
+
   buildLowerVectors(b, loc, {0, 1, 2, 3, 4}, "outerproduct", "innerparallel",
                     "linalg-copy", true, "eltwise", false);
   buildLowerVectors(b, loc, {5, 6, 7}, "outerproduct", "innerparallel",
                     "linalg-copy", true, "eltwise", false);
   variant = buildDecomposeVectors(b, loc, variant, hardwareVectorSizeInBytes);
+  variant = buildRunCanonicalizer(b, loc, variant);
+  buildLowerVectors(b, loc, {0, 1, 2, 3, 4}, "outerproduct", "innerparallel",
+                    "linalg-copy", true, "eltwise", false);
+  buildLowerVectors(b, loc, {5, 6, 7}, "outerproduct", "innerparallel",
+                    "linalg-copy", true, "eltwise", false);
   b.create<transform::YieldOp>(loc);
   return success();
 }
diff --git a/tao_compiler/mlir/disc/transforms/fusion_utils_transform_based.cc b/tao_compiler/mlir/disc/transforms/fusion_utils_transform_based.cc
index 3457bdd3509..fcb7c28fa0a 100644
--- a/tao_compiler/mlir/disc/transforms/fusion_utils_transform_based.cc
+++ b/tao_compiler/mlir/disc/transforms/fusion_utils_transform_based.cc
@@ -76,6 +76,9 @@ bool isSupportedBcast(Operation* op, ShapeAnalysis& shapeAnalysisBase) {
 }
 
 bool TransformBasedCpuFusionStrategy::isFusible(Operation* op) {
+  if (!useTransformGEMMEpilogueFusionSchedule()) {
+    return isSupportedDot(op) || isa<lmhlo::ConstantOp>(op);
+  }
   return isSupportedDot(op) || isElementWise(op) || isBcastOp(op) ||
          isa<lmhlo::ConstantOp>(op);
 }
diff --git a/tao_compiler/mlir/disc/transforms/lhlo_fusion.cc b/tao_compiler/mlir/disc/transforms/lhlo_fusion.cc
index f8d910c46dd..a576b922237 100644
--- a/tao_compiler/mlir/disc/transforms/lhlo_fusion.cc
+++ b/tao_compiler/mlir/disc/transforms/lhlo_fusion.cc
@@ -737,10 +737,6 @@ struct DiscFusionPass : public DiscFusionPassBase<DiscFusionPass> {
         pipeline.emplace_back(
             makeNewPlacementAwareFusionStrategy(gpu_enabled_, "stitch"));
       } else {
-        if (useTransformSchedule()) {
-          pipeline.emplace_back(makeNewPlacementAwareFusionStrategy(
-              gpu_enabled_, "transform_based"));
-        }
         pipeline.emplace_back(
             makeNewPlacementAwareFusionStrategy(gpu_enabled_, "sparse_base"));
         // Do some basic fusion first.
@@ -750,6 +746,10 @@ struct DiscFusionPass : public DiscFusionPassBase<DiscFusionPass> {
             makeNewPlacementAwareFusionStrategy(gpu_enabled_, "stitch"));
         pipeline.emplace_back(
             makeNewPlacementAwareFusionStrategy(gpu_enabled_, "base"));
+        if (useTransformSchedule()) {
+          pipeline.emplace_back(makeNewPlacementAwareFusionStrategy(
+              gpu_enabled_, "transform_based"));
+        }
       }
     }
     return pipeline;
diff --git a/tao_compiler/mlir/disc/transforms/passes.h b/tao_compiler/mlir/disc/transforms/passes.h
index bcdd192c792..b2070c6ecc9 100644
--- a/tao_compiler/mlir/disc/transforms/passes.h
+++ b/tao_compiler/mlir/disc/transforms/passes.h
@@ -309,6 +309,11 @@ createDiscTransformLegalizeToLoopPass(bool gpuEnabled = false,
                                       const std::string& filename = "",
                                       bool expensiveCheck = false);
 
+// Duplicate and fuse some computation into their fusion consumer to reduce
+// memory footprint.
+std::unique_ptr<OperationPass<func::FuncOp>>
+createDiscDuplicateComputationAfterFusionPass();
+
 }  // namespace disc_ral
 }  // namespace mlir
 


Test patch:
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/matmul_epilogue.cc b/tao_compiler/mlir/disc/tests/disc-transform/matmul_epilogue.cc
index 8c91f3c3e87..3490b7fb558 100644
--- a/tao_compiler/mlir/disc/tests/disc-transform/matmul_epilogue.cc
+++ b/tao_compiler/mlir/disc/tests/disc-transform/matmul_epilogue.cc
@@ -45,6 +45,26 @@ TEST(PackedMatmul, F32_biasadd_768x3072) {
       /*profiling*/ true));
 }
 
+TEST(PackedMatmul, F32_biasadd_768x3072_2) {
+  EnvSetting setting = {
+      {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
+      {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", false}},
+      {"DISC_ENABLE_TRANSFORM_GEMM_EPILOGUE_FUSION", {"0", false}},
+      {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", false}}};
+  EnvSettingContext ctx(setting);
+  EXPECT_TRUE(feature_test_main(
+      /*mlir_file_path*/ c_ft_path +
+          "matmul_epilogue_biasadd_p_768x3072_f32.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 1,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"512x768xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
 TEST(PackedMatmul, F32_biasadd_relu_768x3072) {
   EnvSetting setting = {{"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
                         {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", false}},
@@ -80,6 +100,25 @@ TEST(PackedMatmul, F32_complex_test0) {
       /*profiling*/ true));
 }
 
+TEST(PackedMatmul, F32_complex_test0_2) {
+  EnvSetting setting = {
+      {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
+      {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", false}},
+      {"DISC_ENABLE_TRANSFORM_GEMM_EPILOGUE_FUSION", {"0", false}},
+      {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", false}}};
+  EnvSettingContext ctx(setting);
+  EXPECT_TRUE(feature_test_main(
+      /*mlir_file_path*/ c_ft_path + "matmul_epilogue_complex_test0.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 1,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"512x768xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
 TEST(PackedMatmul, F32_SmallStaticShape_test0) {
   EnvSetting setting = {{"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
                         {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", false}},
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/decompose-vectors.mlir b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/decompose-vectors.mlir
index a89712d61d4..d5ba0e5925b 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/decompose-vectors.mlir
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/decompose-vectors.mlir
@@ -128,3 +128,38 @@ transform.structured.canonicalized_sequence failures(propagate) {
 ^bb0(%arg0: !pdl.operation):
   transform.disc.decompose_vectors %arg0 {vector_size = 4}
 }
+
+// -----
+
+// CHECK-LABEL: @test_if
+// CHECK-SAME: (%[[ARG0:.*]]: i1, %[[ARG1:.*]]: vector<8xf32>, %[[ARG2:.*]]: vector<8xf32>)
+func.func @test_if(%arg0: i1, %arg1: vector<8xf32>, %arg2: vector<8xf32>) -> vector<8xf32> {
+  // CHECK: scf.if %[[ARG0]] -> (vector<4xf32>, vector<4xf32>)
+  %out = scf.if %arg0 -> vector<8xf32> {
+    // CHECK: %[[T0:.*]] = vector.extract_strided_slice %[[ARG1]] {offsets = [0], sizes = [4], strides = [1]} : vector<8xf32> to vector<4xf32>
+    // CHECK: %[[T1:.*]] = vector.extract_strided_slice %[[ARG2]] {offsets = [0], sizes = [4], strides = [1]} : vector<8xf32> to vector<4xf32>
+    // CHECK: %[[T2:.*]] = arith.addf %[[T0]], %[[T1]] : vector<4xf32>
+    // CHECK: %[[T3:.*]] = vector.extract_strided_slice %[[ARG1]] {offsets = [4], sizes = [4], strides = [1]} : vector<8xf32> to vector<4xf32>
+    // CHECK: %[[T4:.*]] = vector.extract_strided_slice %[[ARG2]] {offsets = [4], sizes = [4], strides = [1]} : vector<8xf32> to vector<4xf32>
+    // CHECK: %[[T5:.*]] = arith.addf %[[T3]], %[[T4]] : vector<4xf32>
+    // CHECK: scf.yield %[[T2]], %[[T5]] : vector<4xf32>, vector<4xf32>
+    %0 = arith.addf %arg1, %arg2 : vector<8xf32>
+    scf.yield %0 : vector<8xf32>
+  } else {
+    // CHECK: %[[T10:.*]] = vector.extract_strided_slice %[[ARG1]] {offsets = [0], sizes = [4], strides = [1]} : vector<8xf32> to vector<4xf32>
+    // CHECK: %[[T11:.*]] = vector.extract_strided_slice %[[ARG2]] {offsets = [0], sizes = [4], strides = [1]} : vector<8xf32> to vector<4xf32>
+    // CHECK: %[[T12:.*]] = arith.subf %[[T10]], %[[T11]] : vector<4xf32>
+    // CHECK: %[[T13:.*]] = vector.extract_strided_slice %[[ARG1]] {offsets = [4], sizes = [4], strides = [1]} : vector<8xf32> to vector<4xf32>
+    // CHECK: %[[T14:.*]] = vector.extract_strided_slice %[[ARG2]] {offsets = [4], sizes = [4], strides = [1]} : vector<8xf32> to vector<4xf32>
+    // CHECK: %[[T15:.*]] = arith.subf %[[T13]], %[[T14]] : vector<4xf32>
+    // CHECK: scf.yield %[[T12]], %[[T15]] : vector<4xf32>, vector<4xf32>
+    %1 = arith.subf %arg1, %arg2 : vector<8xf32>
+    scf.yield %1 : vector<8xf32>
+  }
+  return %out: vector<8xf32>
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation):
+  transform.disc.decompose_vectors %arg0 {vector_size = 4}
+}
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/disc-bufferize.mlir b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/disc-bufferize.mlir
index b81bdac7500..26c2ef3be54 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/disc-bufferize.mlir
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/disc-bufferize.mlir
@@ -104,6 +104,34 @@ func.func @bufferize_constant_wrapper() -> tensor<512x1024xf32> {
   return %0 : tensor<512x1024xf32>
 }
 
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  transform.disc.bufferize %arg1
+}
+
+// -----
+
+#map0 = affine_map<(d0, d1) -> ()>
+#map1 = affine_map<(d0, d1) -> (d1)>
+#map2 = affine_map<(d0, d1) -> (d0, d1)>
+
+// CHECK-LABEL: @bufferize_conditional_generic
+// CHECK-SAME: (%[[ARG0:.*]]: i1, %[[ARG1:.*]]: memref<f32>, %[[ARG2:.*]]: memref<12xf32>, %[[ARG3:.*]]: memref<8x12xf32>)
+func.func @bufferize_conditional_generic(
+    %pred : i1, %arg0: tensor<f32>, %arg1 : tensor<12xf32>, %arg2 : tensor<8x12xf32>) -> tensor<8x12xf32> {
+  // CHECK: disc_linalg_ext.conditional_generic
+  // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]], %[[ARG2]] : i1, memref<f32>, memref<12xf32>)
+  // CHECK-SAME: outs(%[[ARG3]] : memref<8x12xf32>)
+  %out = disc_linalg_ext.conditional_generic {indexing_maps = [#map0, #map0, #map1, #map2], iterator_types = ["parallel", "parallel"]}
+      ins(%pred, %arg0, %arg1 : i1, tensor<f32>, tensor<12xf32>) outs(%arg2 : tensor<8x12xf32>) attrs =  {disc.device = "cpu", disc.transform.name = "maximum"} {
+  ^bb0(%in: i1, %in_1: f32, %in_2: f32, %out: f32):
+    %t0 = arith.addf %out, %in_2 : f32
+    %t1 = arith.maxf %in_1, %t0 : f32
+    disc_linalg_ext.yield %t1 : f32
+  } -> tensor<8x12xf32>
+  return %out : tensor<8x12xf32>
+}
+
 transform.structured.canonicalized_sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   transform.disc.bufferize %arg1
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/lower-conditional-generic.mlir b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/lower-conditional-generic.mlir
new file mode 100644
index 00000000000..0661e4a4c2a
--- /dev/null
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/lower-conditional-generic.mlir
@@ -0,0 +1,29 @@
+// RUN: disc-opt --disc-transform-dialect-interpreter -split-input-file %s | FileCheck %s
+
+#map0 = affine_map<(d0, d1) -> ()>
+#map1 = affine_map<(d0, d1) -> (d1)>
+#map2 = affine_map<(d0, d1) -> (d0, d1)>
+
+// CHECK-LABEL: @lower_conditional_generic
+// CHECK-SAME: (%[[ARG0:.*]]: i1, %[[ARG1:.*]]: memref<f32>, %[[ARG2:.*]]: memref<12xf32>, %[[ARG3:.*]]: memref<8x12xf32>)
+func.func @lower_conditional_generic(
+    %pred : i1, %arg0: memref<f32>, %arg1 : memref<12xf32>, %arg2 : memref<8x12xf32>) -> memref<8x12xf32> {
+  // CHECK: scf.if %[[ARG0]] {
+  // CHECK-NEXT: linalg.generic
+  // CHECK-SAME: ins(%[[ARG1]], %[[ARG2]] : memref<f32>, memref<12xf32>)
+  // CHECK-SAME: outs(%[[ARG3]] : memref<8x12xf32>)
+  disc_linalg_ext.conditional_generic {indexing_maps = [#map0, #map0, #map1, #map2], iterator_types = ["parallel", "parallel"]}
+      ins(%pred, %arg0, %arg1 : i1, memref<f32>, memref<12xf32>) outs(%arg2 : memref<8x12xf32>) {
+  ^bb0(%in: i1, %in_1: f32, %in_2: f32, %out: f32):
+    %t0 = arith.addf %out, %in_2 : f32
+    %t1 = arith.maxf %in_1, %t0 : f32
+    disc_linalg_ext.yield %t1 : f32
+  }
+  return %arg2 : memref<8x12xf32>
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation):
+  %0 = transform.structured.match ops{["disc_linalg_ext.conditional_generic"]} in %arg0
+  transform.disc.lower_conditional_generic %0
+}
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/transforms/tests/disc-duplicate-computation-after-fusion.mlir b/tao_compiler/mlir/disc/transforms/tests/disc-duplicate-computation-after-fusion.mlir
new file mode 100644
index 00000000000..1166895a144
--- /dev/null
+++ b/tao_compiler/mlir/disc/transforms/tests/disc-duplicate-computation-after-fusion.mlir
@@ -0,0 +1,19 @@
+// RUN: DISC_ENABLE_TRANSFORM_SCHEDULE=1 disc-opt -split-input-file --disc-duplicate-computation-after-fusion %s -o - | FileCheck %s --check-prefix=TRANSFORM
+
+// TRANSFORM-LABEL: @duplicate_scalar_weight_for_ktransform
+func.func @duplicate_scalar_weight_for_ktransform(
+    %arg0 : memref<128x128xf32, "cpu">, %arg1 : memref<128x128xf32, "cpu">,
+    %arg2 : memref<128x128xf32, "cpu">, %arg3: memref<f32, "cpu">, %arg4: memref<2xindex, "cpu">,
+    %arg5 : memref<128x128xf32, "cpu">, %arg6 : memref<128x128xf32, "cpu">) -> (memref<128x128xf32, "cpu">) {
+  "lmhlo.constant"(%arg3) {disc.device = "cpu", value = dense<2.000000e-01> : tensor<f32>} : (memref<f32, "cpu">) -> ()
+  // TRANSFORM: "lmhlo.fusion"
+  // TRANSFORM-NEXT: "lmhlo.constant"(%[[T:.*]]) {disc.device = "cpu", value = dense<2.000000e-01> : tensor<f32>} : (memref<f32, "cpu">) -> ()
+  "lmhlo.fusion"() ({
+    "lmhlo.constant"(%arg0) {disc.device = "cpu", value = dense<0.1> : tensor<128x128xf32>} : (memref<128x128xf32, "cpu">) -> ()
+    "lmhlo.dot_general"(%arg1, %arg0, %arg2) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>} : (memref<128x128xf32, "cpu">, memref<128x128xf32, "cpu">, memref<128x128xf32, "cpu">) -> ()
+    "lmhlo.dynamic_broadcast_in_dim"(%arg3, %arg4, %arg5) {broadcast_dimensions = dense<[]> : tensor<0xi64>} : (memref<f32, "cpu">, memref<2xindex, "cpu">, memref<128x128xf32, "cpu">) -> ()
+    "lmhlo.add"(%arg2, %arg5, %arg6) : (memref<128x128xf32, "cpu">, memref<128x128xf32, "cpu">, memref<128x128xf32, "cpu">) -> ()
+    "lmhlo.terminator"() : () -> ()
+  }) {disc.device = "cpu", disc.fusion.name = "xxx0", disc.fusion_type = "kTransform"} : () -> ()
+  return %arg6 :  memref<128x128xf32, "cpu">
+}
\ No newline at end of file