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/mhlo_mark_shape_calc.cc b/tao_compiler/mlir/disc/transforms/mhlo_mark_shape_calc.cc
index d18ea721c31..fe53cac3aa2 100644
--- a/tao_compiler/mlir/disc/transforms/mhlo_mark_shape_calc.cc
+++ b/tao_compiler/mlir/disc/transforms/mhlo_mark_shape_calc.cc
@@ -90,6 +90,22 @@ void DiscMarkShapeCalc::MarkShapeCalcOps() {
   }
 }
 
+// NOTE(lanbo.llb): mhlo.disc op will produce 2 outputs, the 2nd output is i64
+// tensor with single element represent number of valid elements in the 1st
+// output, this will be used to calc the final output shape. It is indeed a
+// shape calc, but where op should not be considered as shape calc op. By
+// marking it as shape calc op will cause mhlo_disc.where impossible to fuse in
+// later fusion pass.
+bool withInputFromWhereOp(Operation* op) {
+  for (Value operand : op->getOperands()) {
+    auto op = operand.getDefiningOp();
+    if (op != nullptr && isa<mhlo_disc::WhereOp>(op)) {
+      return true;
+    }
+  }
+  return false;
+}
+
 void DiscMarkShapeCalc::markShapeCalculationOps(
     func::FuncOp func, llvm::DenseSet<Operation*>& marked_ops) {
   auto& block = func.getBlocks().front();
@@ -100,7 +116,8 @@ void DiscMarkShapeCalc::markShapeCalculationOps(
     if (!marked_ops.contains(&op)) {
       // Mark following Ops into shape calculation set
       if (isa<mhlo::GetDimensionSizeOp, tensor::FromElementsOp,
-              tensor::ExtractOp>(&op)) {
+              tensor::ExtractOp>(&op) &&
+          !withInputFromWhereOp(&op)) {
         marked_ops.insert(&op);
         continue;
       }


Test patch:
diff --git a/tao_compiler/mlir/disc/transforms/tests/mhlo_mark_shape_calc_op.mlir b/tao_compiler/mlir/disc/transforms/tests/mhlo_mark_shape_calc_op.mlir
index 28001c9819e..234f0d2bd44 100644
--- a/tao_compiler/mlir/disc/transforms/tests/mhlo_mark_shape_calc_op.mlir
+++ b/tao_compiler/mlir/disc/transforms/tests/mhlo_mark_shape_calc_op.mlir
@@ -139,3 +139,29 @@ func.func @main(%arg0: tensor<?x?xf32>, %arg1: tensor<2xi32>) -> tensor<?x?xf32>
   } : (tensor<?x?xf32>, tensor<2xi32>) -> tensor<?x?xf32>
   return %1 : tensor<?x?xf32>
 }
+
+// -----
+
+// CHECK-LABEL: @main
+func.func @main(%arg0: tensor<?x?x?xi64>) -> tensor<?x3xi64> {
+  // CHECK: %[[V1:.*]], %[[V1:.*]] = "mhlo_disc.where"(%{{.*}}) : (tensor<?x?x?xi1>) -> (tensor<?x3xi64>, tensor<1xi64>)
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %c3 = arith.constant 3 : index
+  %0 = mhlo.constant 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>
+  %dim_1 = tensor.dim %arg0, %c1 : tensor<?x?x?xi64>
+  %dim_2 = tensor.dim %arg0, %c2 : tensor<?x?x?xi64>
+  %from_elements = tensor.from_elements %dim, %dim_1, %dim_2 : tensor<3xindex>
+  %1 = "mhlo.dynamic_broadcast_in_dim"(%0, %from_elements) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<i64>, tensor<3xindex>) -> tensor<?x?x?xi64>
+  %2 = mhlo.compare  NE, %arg0, %1 : (tensor<?x?x?xi64>, tensor<?x?x?xi64>) -> tensor<?x?x?xi1>
+  %index, %num_output_elements = "mhlo_disc.where"(%2) : (tensor<?x?x?xi1>) -> (tensor<?x3xi64>, tensor<1xi64>)
+  %extracted = tensor.extract %num_output_elements[%c0] : tensor<1xi64>
+  %3 = arith.index_cast %extracted : i64 to index
+  %from_elements_3 = tensor.from_elements %3, %c3 : tensor<2xindex>
+  %4 = mhlo.real_dynamic_slice %index, %cst_0, %from_elements_3, %cst : (tensor<?x3xi64>, tensor<2xindex>, tensor<2xindex>, tensor<2xindex>) -> tensor<?x3xi64>
+  return %4 : tensor<?x3xi64>
+}