Repository: alibaba/BladeDISC (language: python)

Issue:
Issues encountered when adding a new fusion type for `lmhlo_disc.where` op
With some hacks, we have successfully run a small UT for adding a new fusion type for `lmhlo_disc.where`(kWhere), which fuse the input element-wise operations into where op's schedule.
However there exists some problems to be fixed:
- [x] 1. `DiscMarkShapeCalculationPass` mark  `lmhlo_disc.where` and other related ops(not i32 data types in these ops) as shape calc ops, thus cannot be fused by fusion pass(This is currently disable for UT to success)
- [ ] 2. We have a `All outputs of a fusion pattern should have compatible shape` constrains for fusion patterns. However when `lmhlo_disc.where` is involved, it's output cannot satisfy these constrains, maybe we need to add special constrains for `kWhere` fusion.(skip the `compatible shape` check in UT)
- [ ] 3. I think it is because `lmhlo_disc.where` has 2 outputs, but we still use some of the single output fusion utils, something  unexpected has come up -- we get 2 same root ops(both `lmhlo_disc.where`) for `kWhere` fusion. We need to handle 
- [ ] 4. After lower where op, the fusion op still contains a `lmhlo_disc.where` op, which can not be elimated by `cleanUnusedLhloOps` util. I think it has something to do with the `2 same root ops` issue.


Discussion / hints:
IR for issue 1:
```
// -----// IR Dump After DiscMarkShapeCalculationPass (disc-mhlo-mark-shape-calc) //----- //
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 561 : i32}} {
  func.func @main(%arg0: tensor<?x?x?xi64, [@S0, @S1, @S2]>) -> tensor<?x3xi64, [@S4, @C3]> attributes {tf.entry_function = {input_placements = "cpu", inputs = "input0", output_placements = "cpu", outputs = "output0"}} {
    %c0 = arith.constant 0 : index
    %c1 = arith.constant 1 : index
    %c2 = arith.constant 2 : index
    %c3 = arith.constant 3 : index
    %0 = mhlo.constant {disc.shape_op = true} dense<0> : tensor<i64>
    %cst = arith.constant dense<1> : tensor<2xindex>
    %cst_0 = arith.constant dense<0> : tensor<2xindex>
    %dim = tensor.dim %arg0, %c0 : tensor<?x?x?xi64, [@S0, @S1, @S2]>
    %dim_1 = tensor.dim %arg0, %c1 : tensor<?x?x?xi64, [@S0, @S1, @S2]>
    %dim_2 = tensor.dim %arg0, %c2 : tensor<?x?x?xi64, [@S0, @S1, @S2]>
    %from_elements = tensor.from_elements %dim, %dim_1, %dim_2 {disc.shape_op = true} : tensor<3xindex>
    %1 = "mhlo.dynamic_broadcast_in_dim"(%0, %from_elements) {broadcast_dimensions = dense<> : tensor<0xi64>, disc.shape_op = true} : (tensor<i64>, tensor<3xindex>) -> tensor<?x?x?xi64, [@S0, @S1, @S2]>
    %2 = mhlo.compare  NE, %arg0, %1 {disc.shape_op = true} : (tensor<?x?x?xi64, [@S0, @S1, @S2]>, tensor<?x?x?xi64, [@S0, @S1, @S2]>) -> tensor<?x?x?xi1, [@S0, @S1, @S2]>
    %index, %num_output_elements = "mhlo_disc.where"(%2) {disc.shape_op = true} : (tensor<?x?x?xi1, [@S0, @S1, @S2]>) -> (tensor<?x3xi64, [@S3, @C3]>, tensor<1xi64>)
    %extracted = tensor.extract %num_output_elements[%c0] {disc.shape_op = true} : tensor<1xi64>
    %3 = arith.index_cast %extracted : i64 to index
    %from_elements_3 = tensor.from_elements %3, %c3 {disc.shape_op = true} : tensor<2xindex>
    %4 = mhlo.real_dynamic_slice %index, %cst_0, %from_elements_3, %cst {kDiscSliceOpStaticKnownInfo = {limit_indices = dense<[-2, -1]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}} : (tensor<?x3xi64, [@S3, @C3]>, tensor<2xindex>, tensor<2xindex>, tensor<2xindex>) -> tensor<?x3xi64, [@S4, @C3]>
    return %4 : tensor<?x3xi64, [@S4, @C3]>
  }
  "disc_shape.SymbolicDim"() {knownNegativeOne = false, knownNonNegative = true, knownNonSizeOne = false, knownNonSizeZero = false, sym_name = "S0", value = -1 : i64} : () -> ()
  "disc_shape.SymbolicDim"() {knownNegativeOne = false, knownNonNegative = true, knownNonSizeOne = false, knownNonSizeZero = false, sym_name = "S1", value = -1 : i64} : () -> ()
  "disc_shape.SymbolicDim"() {knownNegativeOne = false, knownNonNegative = true, knownNonSizeOne = false, knownNonSizeZero = false, sym_name = "S2", value = -1 : i64} : () -> ()
  "disc_shape.SymbolicDim"() {knownNegativeOne = false, knownNonNegative = true, knownNonSizeOne = true, knownNonSizeZero = true, sym_name = "C3", value = 3 : i64} : () -> ()
  "disc_shape.SymbolicDim"() {knownNegativeOne = false, knownNonNegative = true, knownNonSizeOne = false, knownNonSizeZero = false, sym_name = "S3", value = -1 : i64} : () -> ()
  "disc_shape.SymbolicDim"() {knownNegativeOne = false, knownNonNegative = true, knownNonSizeOne = false, knownNonSizeZero = false, sym_name = "S4", value = -1 : i64} : () -> ()
  func.func @shape_constraint_graph() {
    return
  }
```
1. it seems to be a by-design behavior. since the `num_output_elements` is finally used as shape input of `read_dynamic_slice`, maybe an extra rule to exclude where will work. 
2. FusionStrategy of `kWhere` can override `tryFuse` method in base class?

> 1. it seems to be a by-design behavior. since the `num_output_elements` is finally used as shape input of `read_dynamic_slice`, maybe an extra rule to exclude where will work.

Trying  to add a rule for `tensor.extract` when its input is from `lmhlo_disc.where` 

> 2. FusionStrategy of `kWhere` can override `tryFuse` method in base class?

Currently, I modify the logic of `BaseCPUFusionStrategy` to support `kWhere`



Fix patch:
diff --git a/tao_compiler/mlir/disc/transforms/fusion_utils.cc b/tao_compiler/mlir/disc/transforms/fusion_utils.cc
index 3989d73f4f3..29cefb0ccaf 100644
--- a/tao_compiler/mlir/disc/transforms/fusion_utils.cc
+++ b/tao_compiler/mlir/disc/transforms/fusion_utils.cc
@@ -46,6 +46,7 @@ using placement_utils::kDiscPlaceAssignment;
 using placement_utils::kDiscShapeCalcAttr;
 
 void dumpFusionPattern(FusionPattern& pattern) {
+  llvm::dbgs() << "FusionType: " << pattern.getFusionTypeStr() << "\n";
   for (Operation* subOp : pattern.getOpList()) {
     llvm::dbgs() << "  " << *subOp << "\n";
   }
@@ -108,7 +109,8 @@ DenseSet<Operation*> NoLoaderUser(SmallVectorImpl<Operation*>& ops) {
 void cleanUnusedLhloOps(Block* parent, PatternRewriter* rewriter) {
   SmallVector<Operation*, 4> lhlo_ops;
   for (Operation& op : parent->getOperations()) {
-    if (op.getDialect() == op.getContext()->getLoadedDialect("lmhlo") &&
+    if ((op.getDialect() == op.getContext()->getLoadedDialect("lmhlo") ||
+         op.getDialect() == op.getContext()->getLoadedDialect("lmhlo_disc")) &&
         (!isa<lmhlo::TerminatorOp>(op)))
       lhlo_ops.push_back(&op);
   }
@@ -187,6 +189,8 @@ StringRef fusionTypeToString(FusionType ft) {
       return "kLargeConcat";
     case FusionType::kDot:
       return "kDot";
+    case FusionType::kWhere:
+      return "kWhere";
     default:
       assert(false && "unknown fusion type");
       return "";
@@ -211,7 +215,10 @@ FusionType fusionTypeFromString(StringRef ft) {
     return FusionType::kLargeConcat;
   } else if (ft == "kDot") {
     return FusionType::kDot;
+  } else if (ft == "kWhere") {
+    return FusionType::kWhere;
   }
+
   assert(false && "unknown fusion type");
   return FusionType::kNone;
 }
@@ -437,6 +444,10 @@ bool isFusible(Operation* op) {
   // All element ops are supported by the fusion codegen engine.
   if (isElementWise(op)) return true;
 
+  // lmhlo_disc.where is supported for fusion, fuse input element-wise ops like
+  // kInput fusion.
+  if (isa<lmhlo_disc::WhereOp>(op)) return true;
+
   // clang-format off
   return isa<
     lmhlo::BroadcastInDimOp,
@@ -624,7 +635,10 @@ void FusionPatternBase::calculateOperandsAndResults() {
 
       if (has_external_user) {
         results_.push_back(v);
-        root_ops_.push_back(op);
+        // For op with multi-output
+        if (!alreadyInRootOps(op)) {
+          root_ops_.push_back(op);
+        }
         if (!has_internal_user) {
           external_only_results_.push_back(v);
         }
@@ -1004,6 +1018,7 @@ bool FusionStrategy::tryFuseInplace(ShapeAnalysis& shapeAnalysis,
     return false;
   }
   lhs = result;
+
   return true;
 }
 
@@ -1074,6 +1089,9 @@ bool BaseFusionStrategy::tryFuse(ShapeAnalysis& shapeAnalysis,
   // Here `compatible` means:
   // - if `to` and `from` are both kInput fusion, all output should have same
   // shape.
+  // - if `to` or `from` is kWhere fusion, there should be no constrains for
+  // output shape, since lmhlo_disc.where's output shape is determined by values
+  // from its input, which cannot be decided at compile-time.
   // - otherwise, all output should have same number of elements.
 
   // No outside users, these ops may be eliminated. We fused it here and let
@@ -1082,14 +1100,26 @@ bool BaseFusionStrategy::tryFuse(ShapeAnalysis& shapeAnalysis,
     return true;
   }
 
-  Value ref_shape = getEffectiveShape(target, results[0]);
-  if (!llvm::all_of(results, [&](Value result) {
-        Value shape = getEffectiveShape(target, result);
-        return checkSameShape(lhs, rhs, target)
-                   ? shapeAnalysis.isShapeEqual(ref_shape, shape)
-                   : shapeAnalysis.isSameNumElements(ref_shape, shape);
-      })) {
-    return false;
+  if (rhs.getFusionType() != FusionType::kWhere &&
+      lhs.getFusionType() != FusionType::kWhere) {
+    Value ref_shape = getEffectiveShape(target, results[0]);
+    if (!llvm::all_of(results, [&](Value result) {
+          Value shape = getEffectiveShape(target, result);
+          return checkSameShape(lhs, rhs, target)
+                     ? shapeAnalysis.isShapeEqual(ref_shape, shape)
+                     : shapeAnalysis.isSameNumElements(ref_shape, shape);
+        })) {
+      return false;
+    }
+  } else {
+    // TODO(lanbo.llb): make a more general rule for kWhere fusion
+    if ((target.getRootOps().size() != 1) ||
+        (target.getRootOps().size() == 1 &&
+         !isa<lmhlo_disc::WhereOp>(target.getRootOps()[0]))) {
+      // Currently we only support input-fusion like kWhere fusion
+      // with one root op, which should be mhlo_disc.where
+      return false;
+    }
   }
   LLVM_DEBUG(llvm::dbgs() << "BaseFusionStrategy::tryFuse success()\n");
   return true;
@@ -1152,11 +1182,17 @@ bool BaseCpuFusionStrategy::initFusionPattern(ShapeAnalysis& shapeAnalysis,
       // the fusion type to kLoop and dominant op to current op. This supposes
       // that the last op inside the block is a valid candidate dominant op if
       // the fusion pattern is a kLoop.
-      if (inferredFusionType == FusionType::kNone ||
-          inferredFusionType == FusionType::kLoop) {
-        inferredFusionType = FusionType::kLoop;
+      if (!isa<lmhlo_disc::WhereOp>(op)) {
+        if (inferredFusionType == FusionType::kNone ||
+            inferredFusionType == FusionType::kLoop) {
+          inferredFusionType = FusionType::kLoop;
+          inferredDominantOp = op;
+        }
+      } else {
+        inferredFusionType = FusionType::kWhere;
         inferredDominantOp = op;
       }
+
     } else if (!isa<lmhlo::TerminatorOp>(op)) {
       // Not a supported fusionOp, early stop.
       inferredFusionType = FusionType::kNone;
diff --git a/tao_compiler/mlir/disc/transforms/fusion_utils.h b/tao_compiler/mlir/disc/transforms/fusion_utils.h
index a20d6ac539f..77c85cfadb4 100644
--- a/tao_compiler/mlir/disc/transforms/fusion_utils.h
+++ b/tao_compiler/mlir/disc/transforms/fusion_utils.h
@@ -115,7 +115,9 @@ enum FusionType {
   // A schedule for concat op having many operands.
   kLargeConcat,
   // Dot fusion.
-  kDot
+  kDot,
+  // Where op fusion, maybe need a more general name?
+  kWhere,
 };
 
 FusionType getFusionType(Operation* op);
@@ -239,6 +241,15 @@ class FusionPatternBase {
     last_writer_[value] = op;
   }
 
+  bool alreadyInRootOps(Operation* new_op) {
+    for (Operation* op : root_ops_) {
+      if (new_op == op) {
+        return true;
+      }
+    }
+    return false;
+  }
+
  protected:
   FusionPatternBase(SmallVectorImpl<Operation*>& op_list);
 
diff --git a/tao_compiler/mlir/disc/transforms/lhlo_legalize_roots_to_loops.cc b/tao_compiler/mlir/disc/transforms/lhlo_legalize_roots_to_loops.cc
index 2511e2398d5..4ac74c1fc2e 100644
--- a/tao_compiler/mlir/disc/transforms/lhlo_legalize_roots_to_loops.cc
+++ b/tao_compiler/mlir/disc/transforms/lhlo_legalize_roots_to_loops.cc
@@ -4761,8 +4761,6 @@ LogicalResult lowerWithScheduleWhereOpCPU(
   auto loc = where.getLoc();
   OpBuilder b(root_ops.back());
 
-  b.setInsertionPoint(where.getOperation());
-
   Value input = where.getInput();
   Value index = where.getIndex();
   Value num_output_elements = where.getNumOutputElements();
@@ -4778,26 +4776,35 @@ LogicalResult lowerWithScheduleWhereOpCPU(
       b.create<arith::ConstantOp>(loc, b.getIntegerAttr(b.getIndexType(), 1));
   Value one_for_acc = b.create<arith::ConstantOp>(
       loc, b.getIntegerAttr(b.getIntegerType(64), 1));
-  SmallVector<Value, 1> num_output_load_index = {
+  SmallVector<Value, 1> zero_load_index = {
       b.create<arith::ConstantIndexOp>(loc, 0)};
-  // memset num_output_elements to 0
+  // stack temp buffer for num_output_elements to 0
+  auto alloc = b.create<memref::AllocaOp>(
+      loc, MemRefType::get(
+               {}, b.getIntegerType(64), MemRefLayoutAttrInterface(),
+               StringAttr::get(where->getContext(), placement_utils::kCpu)));
+  auto temp_count = alloc.getResult();
   b.create<memref::StoreOp>(loc,
                             b.create<arith::ConstantOp>(
                                 loc, b.getIntegerAttr(b.getIntegerType(64), 0)),
-                            num_output_elements, num_output_load_index);
+                            temp_count);
 
   // for (i = 0; i < num_element; ++i)
-  llvm::SmallVector<scf::ForOp, 4> for_ops;
+  llvm::SmallVector<Value, 2> lower, upper, step;
   llvm::SmallVector<Value, 4> multidim_index;
   llvm::SmallVector<Value, 4> multidim_index_casted;
   for (int i = 0; i < input_rank; ++i) {
-    Value dim = b.create<memref::DimOp>(loc, input, i);
-    auto for_op = b.create<scf::ForOp>(loc, /* lowerBound */ zero,
-                                       /* upperBound */ dim, /* step */ one);
-    for_ops.push_back(for_op);
-    for_op.getBody()->clear();
-    b.setInsertionPointToStart(for_op.getBody());
-    Value loop_index = for_op.getInductionVar();
+    lower.push_back(zero);
+    upper.push_back(b.create<memref::DimOp>(loc, input, i));
+    step.push_back(one);
+  }
+  auto for_op =
+      b.create<scf::ParallelOp>(loc, /* lowerBound */ lower,
+                                /* upperBound */ upper, /* step */ step);
+  Value output_index_count;
+  b.setInsertionPointToStart(for_op.getBody());
+  for (int i = 0; i < input_rank; ++i) {
+    Value loop_index = for_op.getInductionVars()[i];
     multidim_index.push_back(loop_index);
     multidim_index_casted.push_back(
         b.create<arith::IndexCastOp>(loc, output_elem_type, loop_index));
@@ -4827,8 +4834,7 @@ LogicalResult lowerWithScheduleWhereOpCPU(
                                         /*hasElseRegion*/ false);
   if_is_zero.getThenRegion().front().clear();
   b.setInsertionPointToStart(&if_is_zero.getThenRegion().front());
-  Value output_index_count =
-      b.create<memref::LoadOp>(loc, num_output_elements, num_output_load_index);
+  output_index_count = b.create<memref::LoadOp>(loc, temp_count);
   Value output_index_casted =
       b.create<arith::IndexCastOp>(loc, b.getIndexType(), output_index_count);
   for (int i = 0; i < input_rank; ++i) {
@@ -4837,20 +4843,23 @@ LogicalResult lowerWithScheduleWhereOpCPU(
     b.create<memref::StoreOp>(loc, multidim_index_casted[i], index,
                               store_index);
   }
-  ArrayRef<Value> num_output_store_index = num_output_load_index;
   output_index_count =
       b.create<arith::AddIOp>(loc, output_index_count, one_for_acc);
-  b.create<memref::StoreOp>(loc, output_index_count, num_output_elements,
-                            num_output_store_index);
+  b.create<memref::StoreOp>(loc, output_index_count, temp_count);
+
   b.create<scf::YieldOp>(loc, ValueRange({}));  // if (input[i] != 0)
-  for (int i = input_rank - 1; i >= 0; --i) {
-    b.setInsertionPointToEnd(for_ops[i].getBody());
-    b.create<scf::YieldOp>(loc, ValueRange({}));
-  }  // for (i = 0; i < num_element; ++i)
-  b.setInsertionPoint(where.getOperation());
+  b.setInsertionPoint(
+      where.getOperation());  // for (i = 0; i < num_element; ++i)
 
-  // TODO: Support fusion
-  for (Operation* root_op : root_ops) root_op->erase();
+  b.create<memref::StoreOp>(loc, b.create<memref::LoadOp>(loc, temp_count),
+                            num_output_elements, zero_load_index);
+
+  if (non_fusion) {
+    for (Operation* root_op : root_ops) root_op->erase();
+  } else {
+    assert(parent != nullptr && "Parent must be provided for fusion lowering");
+    cleanUnusedLhloOps(parent);
+  }
   return success();
 }
 
@@ -5103,6 +5112,12 @@ LogicalResult HandleCpuFusionOp(OpBuilder& b, Operation* fusion,
         return dominant_op->emitError() << "failed to lower to loops";
       }
       break;
+    case FusionType::kWhere:
+      if (failed(lowerWithScheduleWhereOpCPU(root_ops, dominant_op, fused_block,
+                                             /*non_fusion*/ false))) {
+        return dominant_op->emitError() << "failed to lower to loops";
+      }
+      break;
     default:
       return dominant_op->emitError() << "Unknown fusion type";
   }
diff --git a/tao_compiler/mlir/disc/transforms/mhlo_mark_shape_calc.cc b/tao_compiler/mlir/disc/transforms/mhlo_mark_shape_calc.cc
index fe53cac3aa2..0a9a0fc036d 100644
--- a/tao_compiler/mlir/disc/transforms/mhlo_mark_shape_calc.cc
+++ b/tao_compiler/mlir/disc/transforms/mhlo_mark_shape_calc.cc
@@ -170,7 +170,9 @@ void DiscMarkShapeCalc::inferOperands(func::FuncOp func,
         if (ty && ty.hasStaticShape() && ty.getNumElements() > 64) {
           continue;
         }
-        marked_ops.insert(operand);
+        if (!withInputFromWhereOp(&op)) {
+          marked_ops.insert(operand);
+        }
       }
     }
   };


Test patch:
diff --git a/tao_compiler/mlir/disc/tests/regression/data/where_fusion_d_3d_i64.mlir b/tao_compiler/mlir/disc/tests/regression/data/where_fusion_d_3d_i64.mlir
new file mode 100644
index 00000000000..bc49eeb2da0
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/regression/data/where_fusion_d_3d_i64.mlir
@@ -0,0 +1,11 @@
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 561 : i32}} {
+  func.func @main(%arg0: tensor<?x?x?xi64>) -> tensor<?x3xi64> attributes {tf.entry_function = {inputs = "{{INPUTS}}", outputs = "{{OUTPUTS}}", input_placements="{{INPUT_PLACEMENTS}}", output_placements="{{OUTPUT_PLACEMENTS}}"}} {
+    %0 = tf_executor.graph {
+      %c_0:2 = tf_executor.island wraps "tf.Const"() {device = "", value = dense<0> : tensor<i64>} : () -> tensor<i64>
+      %1:2 = tf_executor.island wraps "tf.NotEqual"(%arg0, %c_0) {device = "", incompatible_shape_error = true} : (tensor<?x?x?xi64>, tensor<i64>) -> tensor<?x?x?xi1>
+      %2:2 = tf_executor.island wraps "tf.Where"(%1) {device = ""} : (tensor<?x?x?xi1>) -> tensor<?x3xi64>
+      tf_executor.fetch %2: tensor<?x3xi64>
+    }
+    return %0 : tensor<?x3xi64>
+  }
+}
diff --git a/tao_compiler/mlir/disc/tests/regression/where_fusion.cc b/tao_compiler/mlir/disc/tests/regression/where_fusion.cc
new file mode 100644
index 00000000000..d50252f9872
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/regression/where_fusion.cc
@@ -0,0 +1,35 @@
+// 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/regression/data/";
+
+TEST(WhereFusionTest, WhereFusionDynamicShape3DI64Test) {
+  setenv("DISC_EXPECTED_KERNELS_IN_UT", "1", 1);
+  EXPECT_TRUE(feature_test_main(
+      /*mlir_file_path*/ c_ft_path + "where_fusion_d_3d_i64.mlir",
+      /*backend_types*/ {kSupportedCPUBackendList},
+      /*num_inputs*/ 1,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"2x3x4xi64_X"},
+      /*output_descriptors*/ {"i64_X"},
+      /*input_vals*/ {{0, 6, 0, 0, 8, 4, 4, 8, 0, 7, 6, 9,
+                       6, 0, 8, 9, 0, 0, 0, 0, 0, 0, 9, 4}}));
+  unsetenv("DISC_EXPECTED_KERNELS_IN_UT");
+}
+
+}  // namespace mlir_test
diff --git a/tao_compiler/mlir/disc/transforms/tests/cpu-only-where.mlir b/tao_compiler/mlir/disc/transforms/tests/cpu-only-where.mlir
index fbacc0bbc5a..7a3f6d47737 100644
--- a/tao_compiler/mlir/disc/transforms/tests/cpu-only-where.mlir
+++ b/tao_compiler/mlir/disc/transforms/tests/cpu-only-where.mlir
@@ -11,24 +11,56 @@ func.func @where(
     memref<1xi64, "cpu">
   ) {
   // CHECK-NOT lmhlo
-  // CHECK: scf.for %[[ARG3:.*]] = %{{.*}} to %{{.*}} step %{{.*}} {
-  // CHECK:   %[[V1:.*]] = arith.index_cast %[[ARG3]] : index to i64
-  // CHECK:   %[[V2:.*]] = memref.dim %{{.*}}, %{{.*}} : memref<?x?xf32, "cpu">
-  // CHECK:   scf.for %[[ARG4:.*]] = %{{.*}} to %[[V2]] step %{{.*}} {
-  // CHECK:     %[[V3:.*]] = arith.index_cast %[[ARG4]] : index to i64
-  // CHECK:     %[[V4:.*]] = memref.load %{{.*}}[%[[ARG3]], %[[ARG4]]] : memref<?x?xf32, "cpu">
-  // CHECK:     %[[V5:.*]] = arith.cmpf one, %[[V4]], %{{.*}} : f32
-  // CHECK:     scf.if %[[V5]] {
-  // CHECK:        %[[V6:.*]] = memref.load %{{.*}}[%{{.*}}] : memref<1xi64, "cpu">
-  // CHECK:        %[[V7:.*]] = arith.index_cast %[[V6]] : i64 to index
-  // CHECK:        memref.store %[[V1]], %{{.*}}[%[[V7]], %{{.*}}] : memref<?x?xi64, "cpu">
-  // CHECK:        memref.store %[[V3]], %{{.*}}[%[[V7]], %{{.*}}] : memref<?x?xi64, "cpu">
-  // CHECK:        %[[V8:.*]] = arith.addi %[[V6]], %{{.*}} : i64
-  // CHECK:        memref.store %[[V8]], %{{.*}}[%{{.*}}] : memref<1xi64, "cpu">
-  // CHECK:     }
+  // CHECK: %[[D:.*]] = memref.dim %{{.*}}, %{{.*}} : memref<?x?xf32, "cpu">
+  // CHECK: %[[D_0:.*]] = memref.dim %{{.*}}, %{{.*}} : memref<?x?xf32, "cpu">
+  // CHECK: scf.parallel (%[[ARG3:.*]], %[[ARG4:.*]]) = (%{{.*}}, %{{.*}}) to (%[[D]], %[[D_0]]) step (%{{.*}}, %{{.*}}) {
+  // CHECK:   %[[V0:.*]] = arith.index_cast %[[ARG3]] : index to i64
+  // CHECK:   %[[V1:.*]] = arith.index_cast %[[ARG4]] : index to i64
+  // CHECK:   %[[V2:.*]] = memref.load %{{.*}}[%[[ARG3]], %[[ARG4]]] : memref<?x?xf32, "cpu">
+  // CHECK:   %[[V3:.*]] = arith.cmpf one, %[[V2]], %{{.*}} : f32
+  // CHECK:   scf.if %[[V3]] {
+  // CHECK:     %[[V4:.*]] = memref.load %{{.*}}[] : memref<i64, "cpu">
+  // CHECK:     %[[V5:.*]] = arith.index_cast %[[V4]] : i64 to index
+  // CHECK:     memref.store %[[V0]], %{{.*}}[%[[V5]], %{{.*}}] : memref<?x?xi64, "cpu">
+  // CHECK:     memref.store %[[V1]], %{{.*}}[%[[V5]], %{{.*}}] : memref<?x?xi64, "cpu">
+  // CHECK:     %[[V6:.*]] = arith.addi %[[V4]], %{{.*}} : i64
+  // CHECK:     memref.store %[[V6]], %{{.*}}[] : memref<i64, "cpu">
   // CHECK:   }
+  // CHECK:   scf.yield
   // CHECK: }
   "lmhlo_disc.where"(%input1, %out1, %out2) : (memref<?x?xf32, "cpu">, memref<?x?xi64, "cpu">, memref<1xi64, "cpu">) -> ()
   // CHECK: return %[[OUT1]], %[[OUT2]] : memref<?x?xi64, "cpu">, memref<1xi64, "cpu">
   return %out1, %out2 : memref<?x?xi64, "cpu">, memref<1xi64, "cpu">
 }
+
+// -----
+
+// CHECK-LABEL: @where_fusion
+// CHECK-SAME: (%[[INPUT1:.*]]: memref<i64, "cpu">, %[[INPUT2:.*]]: memref<3xindex, "cpu">, %[[INPUT3:.*]]: memref<?x?x?xi64, "cpu">, %[[INPUT4:.*]]: memref<?x?x?xi64, "cpu">, %[[INPUT5:.*]]: memref<?x?x?xi1, "cpu">, %[[OUT1:.*]]: memref<?x3xi64, "cpu">, %[[OUT2:.*]]: memref<1xi64, "cpu">) -> (memref<?x3xi64, "cpu">, memref<1xi64, "cpu">)
+func.func @where_fusion(
+    %input1: memref<i64, "cpu">,
+    %input2: memref<3xindex, "cpu">,
+    %input3: memref<?x?x?xi64, "cpu">,
+    %input4: memref<?x?x?xi64, "cpu">,
+    %input5: memref<?x?x?xi1, "cpu">,
+    %out1: memref<?x3xi64, "cpu">,
+    %out2: memref<1xi64, "cpu">
+    ) -> (
+    memref<?x3xi64, "cpu">,
+    memref<1xi64, "cpu">
+  ) {
+  // CHECK lmhlo.fusion
+  // CHECK lmhlo.constant
+  // CHECK lmhlo.dynamic_broadcast_in_dim
+  // CHECK lmhlo.compare
+  // CHECK-NOT lmhlo.where
+  "lmhlo.fusion"() ({
+    "lmhlo.constant"(%input1) {disc.device = "cpu", value = dense<0> : tensor<i64>} : (memref<i64, "cpu">) -> ()
+    "lmhlo.dynamic_broadcast_in_dim"(%input1, %input2, %input3) {broadcast_dimensions = dense<> : tensor<0xi64>, disc.device = "cpu"} : (memref<i64, "cpu">, memref<3xindex, "cpu">, memref<?x?x?xi64, "cpu">) -> ()
+    "lmhlo.compare"(%input4, %input3, %input5) {comparison_direction = #mhlo<comparison_direction NE>, disc.device = "cpu"} : (memref<?x?x?xi64, "cpu">, memref<?x?x?xi64, "cpu">, memref<?x?x?xi1, "cpu">) -> ()
+    "lmhlo_disc.where"(%input5, %out1, %out2) {disc.device = "cpu"} : (memref<?x?x?xi1, "cpu">, memref<?x3xi64, "cpu">, memref<1xi64, "cpu">) -> ()
+    "lmhlo.terminator"() : () -> ()
+  }) {disc.device = "cpu", disc.fusion.name = "main_kWhere_where__4_2_0", disc.fusion_type = "kWhere"} : () -> ()
+  // CHECK: return %[[OUT1]], %[[OUT2]] : memref<?x3xi64, "cpu">, memref<1xi64, "cpu">
+  return %out1, %out2 : memref<?x3xi64, "cpu">, memref<1xi64, "cpu">
+}