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: 534142327

Author: LukeBoyer

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

@@ -737,4 +737,28 @@ 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"
+}
+
 }

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

@@ -15,6 +15,7 @@ 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()">>;
 
@@ -182,6 +183,22 @@ 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,6 +16,8 @@ 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
@@ -28,27 +30,32 @@ using mlir::Operation;
 using mlir::ShapedType;
 using mlir::Value;
 
-// 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>();
+// 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>();
     if (!shaped_type || !shaped_type.hasStaticShape()) {
       return false;
     }
-    if (operand_num == 0) {
-      shape = shaped_type.getShape();
-    } else {
-      if (shape != shaped_type.getShape()) {
-        return false;
-      }
+    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