Add prepare pattern to rewrite TF_Select to TF_SelectV2 in the case of the inputs not all being static and equal.

TF_Select does not support broadcasting whereas TF_SelectV2 does. Broadcasting might be needed at runtime if stapes are not fully known ahead of time so we need TF_SelectV2. See tensorflow/ir/tf_ops.td for more op details.

PiperOrigin-RevId: 534260994

Author: suxinguo

tensorflow/compiler/mlir/lite/tests/prepare-tf.mlir

@@ -737,30 +737,6 @@ func.func @UnsupportedGroupConv_DynamicDimAtInputDimThree(%arg0: tensor<?x1x26x?
   // CHECK: "tf.Conv2D"
 }
 
-func.func @Select_SameStaticShapeUnchanged(%arg0: tensor<2xi1>, %arg1: tensor<2xf32>, %arg2: tensor<2xf32>) -> (tensor<2xf32>) {
-  %0 = "tf.Select"(%arg0, %arg1, %arg2) : (tensor<2xi1>, tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>
-  func.return %0 : tensor<2xf32>
-  // CHECK-LABEL: Select_SameStaticShapeUnchanged
-  // CHECK-NOT: "tf.SelectV2"
-  // CHECK: "tf.Select"
-}
-
-func.func @Select_SameStaticShapeUnchangedWithBroadcastedCond(%arg0: tensor<2xi1>, %arg1: tensor<2x2xf32>, %arg2: tensor<2x2xf32>) -> (tensor<2x2xf32>) {
-  %0 = "tf.Select"(%arg0, %arg1, %arg2) : (tensor<2xi1>, tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
-  func.return %0 : tensor<2x2xf32>
-  // CHECK-LABEL: SameStaticShapeUnchangedWithBroadcastedCond
-  // CHECK-NOT: "tf.SelectV2"
-  // CHECK: "tf.Select"
-}
-
-func.func @Select_NotSameStaticShapeRewritesToSelectV2(%arg0: tensor<*xi1>, %arg1: tensor<*xf32>, %arg2: tensor<*xf32>) -> (tensor<*xf32>) {
-  %0 = "tf.Select"(%arg0, %arg1, %arg2) : (tensor<*xi1>, tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
-  func.return %0 : tensor<*xf32>
-  // CHECK-LABEL: Select_NotSameStaticShapeRewritesToSelectV2
-  // CHECK-NOT: "tf.Select"
-  // CHECK: "tf.SelectV2"
-}
-
 func.func @RedundantShapeOp(%shape: tensor<?xi64>, %fill: tensor<f32>) -> (tensor<?xi64>) {
   %0 = "tf.Fill"(%shape, %fill) : (tensor<?xi64>, tensor<f32>) -> (tensor<*xf32>)
   %1 = "tf.Shape"(%0) : (tensor<*xf32>) -> (tensor<?xi64>)

tensorflow/compiler/mlir/lite/transforms/prepare_patterns.td

@@ -15,7 +15,6 @@ limitations under the License.
 
 include "tensorflow/compiler/mlir/tensorflow/transforms/optimize.td"
 include "tensorflow/compiler/mlir/lite/ir/tfl_ops.td"
-include "tensorflow/compiler/mlir/lite/utils/utils.td"
 
 def FalseBoolAttr : AttrConstraint<CPred<"!$_self.getValue()">>;
 
@@ -183,22 +182,6 @@ def LowerTensorScatterAdd: Pat<
   (TF_TensorScatterUpdateOp $input, $indices,
      (TF_AddOp $updates, (CreateGatherNdOp $updates, $input, $indices)))>;
 
-//===----------------------------------------------------------------------===//
-// Select op patterns.
-//===----------------------------------------------------------------------===//
-
-def NotOperandOneAndTwoHaveSameStaticShape : Constraint<Neg<CPred<"OpHasSameStaticShapes($0.getDefiningOp(), ArrayRef({1, 2}))">>, "">;
-
-// If a TF_Select op does not have same static shape for it's value operands, its arguments may
-// need to be broadcasted at runtime. Only V2 supports this so we rewrite
-// before legalization.
-// TODO(b/270755115) This is a temporary workaround, investigate if this
-// pattern is needed in all situations.
-def ConvertNonStaticSelectToSelectV2 : Pat<
-  (TF_SelectOp:$src $cond, $x, $y),
-  (TF_SelectV2Op $cond, $x, $y),
-  [(NotOperandOneAndTwoHaveSameStaticShape  $src)]>;
-
 //===----------------------------------------------------------------------===//
 // AddV2 op patterns.
 //===----------------------------------------------------------------------===//

tensorflow/compiler/mlir/lite/utils/utils.h

@@ -16,8 +16,6 @@ limitations under the License.
 #ifndef TENSORFLOW_COMPILER_MLIR_LITE_UTILS_UTILS_H_
 #define TENSORFLOW_COMPILER_MLIR_LITE_UTILS_UTILS_H_
 
-#include <numeric>
-
 #include "llvm/ADT/ArrayRef.h"
 #include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
 #include "mlir/IR/Operation.h"  // from @llvm-project
@@ -30,32 +28,27 @@ using mlir::Operation;
 using mlir::ShapedType;
 using mlir::Value;
 
-// Returns true if each Operand at given indices have the same static shape.
-inline bool OpHasSameStaticShapes(Operation* op,
-                                  llvm::ArrayRef<int> operand_idxs) {
-  if (op->getNumOperands() == 0 || operand_idxs.empty()) return true;
-  const int first_opr_idx = operand_idxs[0];
-  ArrayRef<int64_t> shape =
-      op->getOperand(first_opr_idx).getType().dyn_cast<ShapedType>().getShape();
-  for (int opr_idx : operand_idxs) {
-    Value operand = op->getOperand(opr_idx);
-    auto shaped_type = operand.getType().dyn_cast<ShapedType>();
+// Returns true if all tensor value in `values` has static shape and same shape.
+inline bool OpHasSameStaticShapes(Operation* op) {
+  auto values = op->getOperands();
+  int operand_num = 0;
+  ArrayRef<int64_t> shape;
+  for (Value value : values) {
+    auto shaped_type = value.getType().dyn_cast<ShapedType>();
     if (!shaped_type || !shaped_type.hasStaticShape()) {
       return false;
     }
-    if (shape != shaped_type.getShape()) {
-      return false;
+    if (operand_num == 0) {
+      shape = shaped_type.getShape();
+    } else {
+      if (shape != shaped_type.getShape()) {
+        return false;
+      }
     }
+    ++operand_num;
   }
   return true;
 }
-
-// Returns true if each Operand has the same static shape.
-inline bool OpHasSameStaticShapes(Operation* op) {
-  llvm::OwningArrayRef<int> operand_idxs(op->getNumOperands());
-  std::iota(operand_idxs.begin(), operand_idxs.end(), 0);
-  return OpHasSameStaticShapes(op, operand_idxs);
-}
 }  // namespace TFL
 }  // namespace mlir