Repository: alibaba/BladeDISC (language: python)

Issue:
Support fusing "isSplat" constants
The problem is observed in swin-transformer, when pytorch is doing amp.
```
"lmhlo.constant"(%275) {value = dense<0.000000e+00> : tensor<64x784x768xf32>} : (memref<64x784x768xf32, "gpu">) -> ()

…

"lmhlo.fusion"() ( {
      …
      "lmhlo.multiply"(%1741, %275, %1742) {disc.device = "gpu"} : (memref<64x784x768xf32, "gpu">, memref<64x784x768xf32, "gpu">, memref<64x784x768xf32, "gpu">) -> ()
      …
      "lmhlo.terminator"() : () -> ()
 }) {disc.device = "gpu", disc.fusion.name = "main_kLoop_reshape__37_1_2" : () -> ()

"lmhlo.fusion"() ( {
      …
      "lmhlo.multiply"(%1888, %275, %1889) {disc.device = "gpu"} : (memref<64x784x768xf32, "gpu">, memref<64x784x768xf32, "gpu">, memref<64x784x768xf32, "gpu">) -> ()
      …
      "lmhlo.terminator"() : () -> ()
 }) {disc.device = "gpu", disc.fusion.name = "main_kLoop_reshape__37_1_3" : () -> ()
```
In general, the "splat" constants outside of a fusion kernel might cause severe performance issues. In swin-transformer, the performance degradation can be very severe. Please be aware that there might be multiple kernels consuming the splat constant.

Solution 1: mark "splat" constant as fusible in fusion pass; and add an additional fusion stage that allows to duplicate the producer according to some forms of rules, like the FusionMerger in XLA. 

Solution 2: add an additional FuseSplatConstPass after the regular fusion pass that specifically duplicate and fuse the Splat const into fusion kernels.

Both solutions need also to support the fusion codegen for "splat" constants. Solution 2 can be regarded as a shrink version of solution 1， which can not handle such cases: 
```
"lmhlo.constant"(%272) {value = dense<0.000000e+00> : tensor<64x784x768xf32>} : (memref<64x784x768xf32, "gpu">) -> ()
"lmhlo.add"(%272, %273, %275)  {disc.device = "gpu"} : (memref<64x784x768xf32, "gpu">, memref<64x784x768xf32, "gpu">, memref<64x784x768xf32, "gpu">) -> ()

…

"lmhlo.fusion"() ( {
      …
      "lmhlo.multiply"(%1741, %275, %1742) {disc.device = "gpu"} : (memref<64x784x768xf32, "gpu">, memref<64x784x768xf32, "gpu">, memref<64x784x768xf32, "gpu">) -> ()
      …
      "lmhlo.terminator"() : () -> ()
 }) {disc.device = "gpu", disc.fusion.name = "main_kLoop_reshape__37_1_2" : () -> ()

"lmhlo.fusion"() ( {
      …
      "lmhlo.multiply"(%1888, %275, %1889) {disc.device = "gpu"} : (memref<64x784x768xf32, "gpu">, memref<64x784x768xf32, "gpu">, memref<64x784x768xf32, "gpu">) -> ()
      …
      "lmhlo.terminator"() : () -> ()
 }) {disc.device = "gpu", disc.fusion.name = "main_kLoop_reshape__37_1_3" : () -> ()
```


Discussion / hints:
What about lowering large splat lmhlo const to ‘lmhlo scalar const + lmhlo.broadcastOp’ ?
> 

you will still have to have another pass, duplicate the lmhlo.broadcastOp and fuse it. Here lmhlo.broadcastOp has multiple consumers which cannot be fused together.
> > 
> 
> you will still have to have another pass, duplicate the lmhlo.broadcastOp and fuse it. Here lmhlo.broadcastOp has multiple consumers which cannot be fused together.

lmhlo.broadcastOp could just be copied for each consumer and thus cold be fused into their consumer. So, we don't need to worry about "multiple consumers" if we split the large splat ops.

Furthermore, splitting it into ‘lmhlo scalar const + lmhlo.broadcastOp’ could still help us to reduce memory consumption if large splat const op could not be fused for some reasons.
> "lmhlo.constant"(%272) {value = dense<0.000000e+00> : tensor<64x784x768xf32>} : (memref<64x784x768xf32, "gpu">) -> ()
"lmhlo.add"(%272, %273, %275)  {disc.device = "gpu"} : (memref<64x784x768xf32, "gpu">, memref<64x784x768xf32, "gpu">, memref<64x784x768xf32, "gpu">) -> ()

The above pattern could be constant folded.

Actually I am not sure why the constant outside a fusion will hurt the performance. Where is the constant placed after code generation?
> What about lowering large splat lmhlo const to ‘lmhlo scalar const + lmhlo.broadcastOp’ ?

This is a solution works for the case in this issue. I am curious about whether there are large constant tensors that have different values and cannot be represented with `scalar+broadcast`.
> > > 
> > 
> > 
> > you will still have to have another pass, duplicate the lmhlo.broadcastOp and fuse it. Here lmhlo.broadcastOp has multiple consumers which cannot be fused together.
> 
> lmhlo.broadcastOp could just be copied for each consumer and thus cold be fused into their consumer. So, we don't need to worry about "multiple consumers" if we split the large splat ops.
> 
> Furthermore, splitting it into ‘lmhlo scalar const + lmhlo.broadcastOp’ could still help us to reduce memory consumption if large splat const op could not be fused for some reasons.

OK，this might be a better workaround compared with Solution2, the advantage is that we don't need to additionaly handle splat const during codegen. The only side-effect i can think of is, it might bring more BroadcastInDimOp in the kernels (in swin-transformer, much more). I suspect that the overhead to calculate the indices can be unacceptable.
> "lmhlo.constant"(%272) {value = dense<0.000000e+00> : tensor<64x784x768xf32>} : (memref<64x784x768xf32, "gpu">) -> () "lmhlo.add"(%272, %273, %275) {disc.device = "gpu"} : (memref<64x784x768xf32, "gpu">, memref<64x784x768xf32, "gpu">, memref<64x784x768xf32, "gpu">) -> () Such pattern could be const folded.

no, it cannot be folded, the "add" has other operands.
> > "lmhlo.constant"(%272) {value = dense<0.000000e+00> : tensor<64x784x768xf32>} : (memref<64x784x768xf32, "gpu">) -> () "lmhlo.add"(%272, %273, %275) {disc.device = "gpu"} : (memref<64x784x768xf32, "gpu">, memref<64x784x768xf32, "gpu">, memref<64x784x768xf32, "gpu">) -> () Such pattern could be const folded.
> 
> no, it cannot be folded, the "add" has other operands.

Indeed. I thought it was `lmhlo.add"(%272, %272, %275)`.

> > > > 
> > > 
> > > 
> > > you will still have to have another pass, duplicate the lmhlo.broadcastOp and fuse it. Here lmhlo.broadcastOp has multiple consumers which cannot be fused together.
> > 
> > 
> > lmhlo.broadcastOp could just be copied for each consumer and thus cold be fused into their consumer. So, we don't need to worry about "multiple consumers" if we split the large splat ops.
> > Furthermore, splitting it into ‘lmhlo scalar const + lmhlo.broadcastOp’ could still help us to reduce memory consumption if large splat const op could not be fused for some reasons.
> 
> OK，this might be a better workaround compared with Solution2, the advantage is that we don't need to additionaly handle splat const during codegen. The only side-effect i can think of is, it might bring more BroadcastInDimOp in the kernels (in swin-transformer, much more). I suspect that the overhead to calculate the indices can be unacceptable.

the operand of bcast is scalar in `scalar + bcast` pattern,  so, we do not need to propagate the indices of the output of the `bcast` to its operand. It may be something else wrong if we observe extra index calculation. Maybe we need some further investigation.
> > > > > 
> > > > 
> > > > 
> > > > you will still have to have another pass, duplicate the lmhlo.broadcastOp and fuse it. Here lmhlo.broadcastOp has multiple consumers which cannot be fused together.
> > > 
> > > 
> > > lmhlo.broadcastOp could just be copied for each consumer and thus cold be fused into their consumer. So, we don't need to worry about "multiple consumers" if we split the large splat ops.
> > > Furthermore, splitting it into ‘lmhlo scalar const + lmhlo.broadcastOp’ could still help us to reduce memory consumption if large splat const op could not be fused for some reasons.
> > 
> > 
> > OK，this might be a better workaround compared with Solution2, the advantage is that we don't need to additionaly handle splat const during codegen. The only side-effect i can think of is, it might bring more BroadcastInDimOp in the kernels (in swin-transformer, much more). I suspect that the overhead to calculate the indices can be unacceptable.
> 
> the operand of bcast is scalar in `scalar + bcast` pattern, so, we do not need to propagate the indices of the output of the `bcast` to its operand. It may be something else wrong if we observe extra index calculation. Maybe we need some further investigation.

Oh, sorry. i just missed the extra index calculation was introduced by the splat const version in current implementation.   `scalar + bcast` splitting  should not introduce extra index calculation. But, I agree it's more like a workaround.

Fix patch:
diff --git a/tao_compiler/mlir/disc/BUILD b/tao_compiler/mlir/disc/BUILD
index 42298762684..4c4feb2812f 100644
--- a/tao_compiler/mlir/disc/BUILD
+++ b/tao_compiler/mlir/disc/BUILD
@@ -309,6 +309,22 @@ cc_library(
     alwayslink = 1,
 )
 
+cc_library(
+    name = "disc_fuse_splat_const",
+    srcs = ["transforms/disc_fuse_splat_const.cc"],
+    deps = [
+        ":pass_details",
+        "//tensorflow/compiler/mlir/hlo:LmhloPassIncGen",
+        "//tensorflow/compiler/mlir/hlo:lhlo",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:MemRefDialect",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+    ],
+    alwayslink = 1,
+)
+
 cc_library(
     name = "revise_args_for_static_rank",
     srcs = [
@@ -1464,6 +1480,7 @@ cc_library(
         ":disc_convert_const_to_ral",
         ":disc_cpu_map_parallel_loop",
         ":disc_flatten_memref_access",
+        ":disc_fuse_splat_const",
         ":disc_fusion",
         ":disc_hlo_legalize_to_lhlo",
         ":disc_hlo_to_std",
diff --git a/tao_compiler/mlir/disc/disc_compiler.cc b/tao_compiler/mlir/disc/disc_compiler.cc
index a8c1af34406..73b5f2a1567 100644
--- a/tao_compiler/mlir/disc/disc_compiler.cc
+++ b/tao_compiler/mlir/disc/disc_compiler.cc
@@ -311,6 +311,8 @@ LogicalResult LowerHLOToLLVM(ModuleOp m, const DISCLoweringOptions& options) {
   pm.addNestedPass<FuncOp>(disc_ral::createDiscFusionPass(
       gpu_enabled, enable_stitch ? "stitch" : "base"));
   if (gpu_enabled) {
+    // TODO: Support cpu stitch with splat const
+    pm.addNestedPass<FuncOp>(disc_ral::createDiscFuseSplatConstPass());
     auto& gpu_options = options.gpu_info;
     pm.addNestedPass<FuncOp>(
         disc_ral::createDiscSpecializeFusionWithSpeculationPass(
diff --git a/tao_compiler/mlir/disc/transforms/disc_fuse_splat_const.cc b/tao_compiler/mlir/disc/transforms/disc_fuse_splat_const.cc
new file mode 100644
index 00000000000..65ddae2d73b
--- /dev/null
+++ b/tao_compiler/mlir/disc/transforms/disc_fuse_splat_const.cc
@@ -0,0 +1,117 @@
+/* Copyright 2021 The TensorFlow 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.
+==============================================================================*/
+
+// This pass works after the DiscFusion pass, fuse the splat constant into its
+// consumers. The splat const will be duplicated if multiple consumers exist.
+// This pass can be regarded as a restricted version of the FusionMerger pass
+// of XLA, dealing with splat const as producer only.
+
+#include "mlir-hlo/Dialect/lhlo/IR/lhlo_ops.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Location.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/Pass/Pass.h"
+#include "tensorflow/compiler/mlir/disc/transforms/PassDetail.h"
+
+namespace mlir {
+namespace disc_ral {
+
+using namespace lmhlo;
+
+namespace {
+
+class DiscFuseSplatConstPass
+    : public DiscFuseSplatConstPassBase<DiscFuseSplatConstPass> {
+  void runOnOperation() override;
+
+ private:
+  void processSplatConst(
+      lmhlo::ConstOp constant,
+      DenseMap<lmhlo::FusionOp, SmallVector<Operation*, 4>> users);
+};
+
+// step 1, find all the unfused splat constants
+// step 2, fuse them in to the consumers
+// step 3, if not used anymore, erase the original
+void DiscFuseSplatConstPass::runOnOperation() {
+  FuncOp func = getOperation();
+  auto* context = &this->getContext();
+
+  SmallVector<lmhlo::ConstOp, 4> worklist;
+  func.walk([&](lmhlo::ConstOp constant) {
+    if (constant.value().isSplat() &&
+        (constant->getParentOfType<lmhlo::FusionOp>() == nullptr)) {
+      worklist.push_back(constant);
+    }
+  });
+  for (auto constant : worklist) {
+    Value memref = constant.output();
+    DenseMap<lmhlo::FusionOp, SmallVector<Operation*, 4>> users;
+    // TODO: Support patterns like const -> memref.cast -> lmhlo.xxxop
+    for (Operation* user : memref.getUsers()) {
+      if (user == constant.getOperation()) continue;
+      // TODO: support if users are non-fused lmhlo op
+      if (isa<lmhlo::LmhloOp>(user)) {
+        lmhlo::FusionOp fusion = user->getParentOfType<lmhlo::FusionOp>();
+        if (fusion != nullptr) {
+          users[fusion].emplace_back(user);
+        }
+      }
+    }
+    if (users.size() == 0) continue;
+    processSplatConst(constant, users);
+  }
+}
+
+void DiscFuseSplatConstPass::processSplatConst(
+    lmhlo::ConstOp constant,
+    DenseMap<lmhlo::FusionOp, SmallVector<Operation*, 4>> users) {
+  for (auto user_item : users) {
+    auto orig_alloc =
+        dyn_cast<memref::AllocOp>(constant.output().getDefiningOp());
+    if (orig_alloc == nullptr) continue;
+    OpBuilder builder(orig_alloc);
+    auto new_alloc = builder.clone(*orig_alloc.getOperation());
+    Value memref = new_alloc->getResult(0);
+    builder.setInsertionPointToStart(&user_item.first.region().front());
+    builder.create<lmhlo::ConstOp>(constant.getLoc(), constant.value(), memref);
+    for (Operation* user : user_item.second) {
+      user->replaceUsesOfWith(constant.output(), memref);
+    }
+  }
+
+  // Erase the original lmhlo.ConstOp if it has no other users
+  bool should_erase = true;
+  for (Operation* user : constant.output().getUsers()) {
+    if (user == constant.getOperation()) continue;
+    should_erase = false;
+    break;
+  }
+  if (should_erase) {
+    constant.erase();
+  }
+}
+
+}  // namespace
+
+std::unique_ptr<OperationPass<FuncOp>> createDiscFuseSplatConstPass() {
+  return std::make_unique<DiscFuseSplatConstPass>();
+}
+
+}  // 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 ebb48447284..a46445513a4 100644
--- a/tao_compiler/mlir/disc/transforms/disc_passes.td
+++ b/tao_compiler/mlir/disc/transforms/disc_passes.td
@@ -122,6 +122,11 @@ def DiscFusionPass : Pass<"disc-fusion", "FuncOp"> {
   ];
 }
 
+def DiscFuseSplatConstPass: Pass<"disc-fuse-splat-const", "FuncOp"> {
+  let summary = "Fuse splat constant to avoid memory access.";
+  let constructor = "createDiscFuseSplatConstPass()";
+}
+
 def DiscMarkShapeCalculationPass : Pass<"disc-mhlo-mark-shape-calc", "ModuleOp"> {
   let summary = "Identify and mark with an attribute operations performing shape calculation";
   let constructor = "createDiscMarkShapeCalcOpPass()";
diff --git a/tao_compiler/mlir/disc/transforms/fusion_utils.cc b/tao_compiler/mlir/disc/transforms/fusion_utils.cc
index 86db6a68249..af7be4a4fd1 100644
--- a/tao_compiler/mlir/disc/transforms/fusion_utils.cc
+++ b/tao_compiler/mlir/disc/transforms/fusion_utils.cc
@@ -365,9 +365,10 @@ bool isRank2ColReduction(Operation* op) {
 // engine.
 bool isFusible(Operation* op) {
   // Only scalar const are supported by the fusion codegen engine a.t.m.
-  if (dyn_cast<lmhlo::ConstOp>(op)) {
-    MemRefType type = op->getOperand(0).getType().cast<MemRefType>();
-    return (type.getRank() == 0);
+  if (isa<lmhlo::ConstOp>(op)) {
+    auto constant = cast<lmhlo::ConstOp>(op);
+    MemRefType type = constant.output().getType().cast<MemRefType>();
+    return (type.getRank() == 0 || constant.value().isSplat());
   }
 
   // All element ops are supported by the fusion codegen engine.
@@ -2765,4 +2766,4 @@ bool StitchCPUAnalysis::doCodeGeneration(OpBuilder& b, lmhlo::FusionOp fusion) {
 }
 
 }  // namespace disc_ral
-}  // namespace mlir
\ No newline at end of file
+}  // namespace mlir
diff --git a/tao_compiler/mlir/disc/transforms/input_inline_fusion_pattern.cc b/tao_compiler/mlir/disc/transforms/input_inline_fusion_pattern.cc
index a7330bebbe9..3d8ea1d79d4 100644
--- a/tao_compiler/mlir/disc/transforms/input_inline_fusion_pattern.cc
+++ b/tao_compiler/mlir/disc/transforms/input_inline_fusion_pattern.cc
@@ -81,14 +81,20 @@ bool miscFuseHelper<ConstOp>(PatternRewriter& rewriter, Operation* user,
                              const SmallVector<memref::LoadOp>& load_ops,
                              LowerConfig* lower_config) {
   if (!isa<ConstOp>(producer)) return false;
-  auto memref_type =
-      cast<LmhloOp>(producer).getResultBuffer().getType().cast<MemRefType>();
-  assert(memref_type.getRank() == 0 && "only scalar ConstOp can be fused");
+  auto constant = cast<lmhlo::ConstOp>(producer);
+  auto memref_type = constant.output().getType().cast<MemRefType>();
+  bool is_splat = constant.value().isSplat();
+  assert((memref_type.getRank() == 0 || is_splat) &&
+         "only scalar ConstOp can be fused");
   auto loc = user->getLoc();
   rewriter.setInsertionPoint(load_op);
-  Value inlined_result = rewriter.create<arith::ConstantOp>(
-      loc, memref_type.getElementType(),
-      cast<ConstOp>(producer).value().getValues<Attribute>()[{}]);
+  Value inlined_result = is_splat
+                             ? rewriter.create<arith::ConstantOp>(
+                                   loc, memref_type.getElementType(),
+                                   constant.value().getSplatValue<Attribute>())
+                             : rewriter.create<arith::ConstantOp>(
+                                   loc, memref_type.getElementType(),
+                                   constant.value().getValues<Attribute>()[{}]);
   for (memref::LoadOp to_be_replaced : load_ops)
     to_be_replaced.replaceAllUsesWith(inlined_result);
   return true;
@@ -274,4 +280,4 @@ LogicalResult InputInlineFusionPattern::inlineFuseLhloOp(
 }
 
 }  // namespace disc_ral
-}  // namespace mlir
\ No newline at end of file
+}  // namespace mlir
diff --git a/tao_compiler/mlir/disc/transforms/passes.h b/tao_compiler/mlir/disc/transforms/passes.h
index 61a15104122..00fdfaec91a 100644
--- a/tao_compiler/mlir/disc/transforms/passes.h
+++ b/tao_compiler/mlir/disc/transforms/passes.h
@@ -90,6 +90,9 @@ std::unique_ptr<OperationPass<FuncOp>> createParallelLoopTilingPass(
 std::unique_ptr<OperationPass<FuncOp>> createDiscFusionPass(
     bool gpu_enabled = true, const std::string& fusion_strategy = "base");
 
+// Fuse splat constant to avoid memory access
+std::unique_ptr<OperationPass<FuncOp>> createDiscFuseSplatConstPass();
+
 // Mark shape calculating Ops.
 std::unique_ptr<OperationPass<ModuleOp>> createDiscMarkShapeCalcOpPass();
 


Test patch:
diff --git a/tao_compiler/mlir/disc/tests/regression/data/kstitch_fusion_with_splat_const.mlir b/tao_compiler/mlir/disc/tests/regression/data/kstitch_fusion_with_splat_const.mlir
new file mode 100644
index 00000000000..7569ec86578
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/regression/data/kstitch_fusion_with_splat_const.mlir
@@ -0,0 +1,12 @@
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 0 : i32}} {
+  func @main(%arg0: tensor<?x?xf32>) -> (tensor<?xf32>, tensor<?xf32>) attributes {tf.entry_function = {inputs = "{{INPUTS}}", outputs = "{{OUTPUTS}}", input_placements="{{INPUT_PLACEMENTS}}", output_placements="{{OUTPUT_PLACEMENTS}}"}} {
+    %graph:2 = tf_executor.graph {
+      %1:2 = tf_executor.island wraps "tf.Const"() {value = dense<[1]> : tensor<1xi32>} : () -> tensor<1xi32>
+      %2:2 = tf_executor.island wraps "tf.Max"(%arg0, %1) : (tensor<?x?xf32>, tensor<1xi32>) -> tensor<?xf32>
+      %3:2 = tf_executor.island wraps "tf.Const"() {value = dense<1.234> : tensor<f32>} : () -> tensor<f32>
+      %4:2 = tf_executor.island wraps "tf.Add"(%2, %3) : (tensor<?xf32>, tensor<f32>) -> tensor<?xf32>
+      tf_executor.fetch %2, %4 : tensor<?xf32>, tensor<?xf32>
+    }
+    return %graph#0, %graph#1 : tensor<?xf32>, tensor<?xf32>
+  }
+}
diff --git a/tao_compiler/mlir/disc/tests/regression/kstitch_fusion.cc b/tao_compiler/mlir/disc/tests/regression/kstitch_fusion.cc
index 80aeac89229..6c1a6a0aa10 100644
--- a/tao_compiler/mlir/disc/tests/regression/kstitch_fusion.cc
+++ b/tao_compiler/mlir/disc/tests/regression/kstitch_fusion.cc
@@ -133,4 +133,19 @@ TEST(KStitchFusionGPUTest, KStitch3DF64) {
   unsetenv("DISC_EXPECTED_KERNELS_IN_UT");
 }
 
-}  // namespace mlir_test
\ No newline at end of file
+// With splat const in the kStitch fusion
+TEST(KStitchFusionGPUTest, KStitchSimpleWithSplatConstF32) {
+  setenv("DISC_ENABLE_STITCH", "true", 1);
+  setenv("DISC_EXPECTED_KERNELS_IN_UT", "1", 1);
+  EXPECT_TRUE(feature_test_main(
+      /*mlir_file_path*/ c_ft_path + "kstitch_fusion_with_splat_const.mlir",
+      /*backend_types*/ {BackendType::kCuda},
+      /*num_inputs*/ 1,
+      /*num_outputs*/ 2,
+      /*input_descriptors*/ {"11000x123xf32_X"},
+      /*output_descriptors*/ {"f32_X", "f32_X"}));
+  unsetenv("DISC_ENABLE_STITCH");
+  unsetenv("DISC_EXPECTED_KERNELS_IN_UT");
+}
+
+}  // namespace mlir_test
diff --git a/tao_compiler/mlir/disc/transforms/tests/disc-fuse-splat-const.mlir b/tao_compiler/mlir/disc/transforms/tests/disc-fuse-splat-const.mlir
new file mode 100644
index 00000000000..be4c7de3458
--- /dev/null
+++ b/tao_compiler/mlir/disc/transforms/tests/disc-fuse-splat-const.mlir
@@ -0,0 +1,56 @@
+// RUN: disc-opt -disc-fuse-splat-const %s | FileCheck %s
+
+// CHECK-LABEL: @fuse_without_erase
+func @fuse_without_erase(%arg0: memref<2x3xf32, "gpu">) -> memref<2x3xf32, "gpu"> {
+  // CHECK: memref.alloc
+  // CHECK: memref.alloc
+  // CHECK: memref.alloc
+  // CHECK: memref.alloc
+  // CHECK: memref.alloc
+  %0 = memref.alloc() : memref<2x3xf32, "gpu">
+  %1 = memref.alloc() : memref<2x3xf32, "gpu">
+  %2 = memref.alloc() : memref<2x3xf32, "gpu">
+  %3 = memref.alloc() : memref<2x3xf32, "gpu">
+  // CHECK: lmhlo.constant
+  "lmhlo.constant"(%0) {value = dense<0.000000e+00> : tensor<2x3xf32>} : (memref<2x3xf32, "gpu">) -> ()
+  // CHECK: lmhlo.fusion
+  "lmhlo.fusion"() ({
+    // CHECK: lmhlo.constant
+    "lmhlo.abs"(%arg0, %1) : (memref<2x3xf32, "gpu">, memref<2x3xf32, "gpu">) -> ()
+    "lmhlo.add"(%1, %0, %2) : (memref<2x3xf32, "gpu">, memref<2x3xf32, "gpu">, memref<2x3xf32, "gpu">) -> ()
+    "lmhlo.terminator"() : () -> ()
+  }) : () -> ()
+  "lmhlo.abs"(%0, %3) : (memref<2x3xf32, "gpu">, memref<2x3xf32, "gpu">) -> ()
+  return %3 : memref<2x3xf32, "gpu">
+}
+
+// CHECK-LABEL: @fuse_and_duplicate_with_erase
+func @fuse_and_duplicate_with_erase(%arg0: memref<2x3xf32, "gpu">) -> (memref<2x3xf32, "gpu">, memref<2x3xf32, "gpu">) {
+  // CHECK: memref.alloc
+  // CHECK: memref.alloc
+  // CHECK: memref.alloc
+  // CHECK: memref.alloc
+  // CHECK: memref.alloc
+  // CHECK: memref.alloc
+  %0 = memref.alloc() : memref<2x3xf32, "gpu">
+  %1 = memref.alloc() : memref<2x3xf32, "gpu">
+  %2 = memref.alloc() : memref<2x3xf32, "gpu">
+  %3 = memref.alloc() : memref<2x3xf32, "gpu">
+  // CHECK-NOT: lmhlo.constant
+  "lmhlo.constant"(%0) {value = dense<0.000000e+00> : tensor<2x3xf32>} : (memref<2x3xf32, "gpu">) -> ()
+  // CHECK: lmhlo.fusion
+  "lmhlo.fusion"() ({
+    // CHECK: lmhlo.constant
+    "lmhlo.abs"(%arg0, %1) : (memref<2x3xf32, "gpu">, memref<2x3xf32, "gpu">) -> ()
+    "lmhlo.add"(%1, %0, %2) : (memref<2x3xf32, "gpu">, memref<2x3xf32, "gpu">, memref<2x3xf32, "gpu">) -> ()
+    "lmhlo.terminator"() : () -> ()
+  }) : () -> ()
+  // CHECK: lmhlo.fusion
+  "lmhlo.fusion"() ({
+    // CHECK: lmhlo.constant
+    "lmhlo.abs"(%0, %1) : (memref<2x3xf32, "gpu">, memref<2x3xf32, "gpu">) -> ()
+    "lmhlo.add"(%1, %arg0, %3) : (memref<2x3xf32, "gpu">, memref<2x3xf32, "gpu">, memref<2x3xf32, "gpu">) -> ()
+    "lmhlo.terminator"() : () -> ()
+  }) : () -> ()
+  return %2, %3 : memref<2x3xf32, "gpu">, memref<2x3xf32, "gpu">
+}