//
// is lowered to:
//
//   %sum0 = "tf.AddV2"(%0, %1)
//   %result = "tf.AddV2"(%sum0, %2)
//
// While
//
//   %result = "tf.AddN"(%0, %1, %2, %3, %4)
//
// is lowered to:
//
//   %sum0 = "tf.AddV2"(%0, %1)
//   %sum1 = "tf.AddV2"(%2, %3)
//   %sum2 = "tf.AddV2"(%sum0, %sum1)
//   %result = "tf.AddV2"(%sum2, %4)
//
class LowerAddNOp : public RewritePattern {
 public:

  auto rhs_attr = mlir::dyn_cast_or_null<DenseElementsAttr>(operands[1]);
  if (!rhs_attr && !(is_commutative && lhs_attr)) return {};

  // Mul and Div ops have identity value one while AddV2 and SubOp have identity
  // value zero.
  const int identity =
      (std::is_same<OpT, MulOp>::value || std::is_same<OpT, DivOp>::value ||
       std::is_same<OpT, RealDivOp>::value)
          ? 1
          : 0;

int64_t InferLookupTableFindV2Cost(const CostContext& context,
                                   mlir::TF::LookupTableFindV2Op op) {
  // tf.LookupTableFindV2 ops are usually more costly than tf.AddV2 with the
  // same input size, as it involves more operations like hashing, map lookup,
  // etc.
  constexpr int64_t kLookupTableFindCostScale = 8;
  constexpr int64_t kLookupTableFindStringKeyCostScale = 16;

        %1 = "tf.Const"() {_xla_outside_compilation = "0", value = dense<1.0> : tensor<f32>} : () -> (tensor<f32>)
        %2 = "tf.Identity"(%1) {_xla_outside_compilation = "0"} : (tensor<f32>) -> (tensor<f32>)
        %3 = "tf.AddV2"(%1, %2) : (tensor<f32>, tensor<f32>) -> (tensor<f32>)
        tf_device.return %3 : tensor<f32>
      }) {num_cores_per_replica = 1, topology =  "", device_assignment =  []} : () -> tensor<f32>
      return %0 : tensor<f32>

  %28 = "tf.AddV2"(%27, %3) {device = ""} : (tensor<2x16xi32>, tensor<i32>) -> tensor<2x16xi32>
  %29 = "tf.AddV2"(%17, %28) {device = ""} : (tensor<2x16xi32>, tensor<2x16xi32>) -> tensor<2x16xi32>
  %30 = "tf.GatherV2"(%9, %29, %5) {batch_dims = 0 : i64, device = ""} : (tensor<32x1xf32>, tensor<2x16xi32>, tensor<i32>) -> tensor<2x16x1xf32>

  (TF_TensorScatterUpdateOp $input, $indices,
     (TF_AddOp $updates, (CreateGatherNdOp $updates, $input, $indices)))>;

//===----------------------------------------------------------------------===//
// AddV2 op patterns.
//===----------------------------------------------------------------------===//

// This is mostly safe, we can use i64 but it will break Darwinn case.

    %10 = "tf.Cast"(%arg3) {Truncate = false, device = ""} : (tensor<i32>) -> tensor<f32>
    %11 = "tf.AddV2"(%arg3, %9) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32>
    %12 = "tf.StackPushV2"(%arg4, %10) {device = "", swap_memory = false} : (tensor<!tf_type.resource>, tensor<f32>) -> tensor<f32>
    %13 = "tf.AddV2"(%arg1, %9) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32>

  %0 = "tf.AddV2"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
  func.return %0 : tensor<ui32>

  // CHECK-LABEL: add_v2_uint32
  // CHECK:  %[[CAST:.*]] = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<ui32>) -> tensor<i32>
  // CHECK:  %[[CAST1:.*]] = "tf.Cast"(%arg1) <{Truncate = false}> : (tensor<ui32>) -> tensor<i32>

    device: The tensorflow device. eg. "/CPU:0"
    op_attrs: The tensorflow attributes excluding the func attrs.
    op_func_attrs: The func attrs.
    op_name: The tensorflow op name. eg. "tf.AddV2"

    Example:
      %out_ch = tfrt_fallback_async.createop(%in_ch) key(0) device("/CPU:0")
        "some.op"() {attr1 = value, attr2 = value2} {f_attr1 = "fn1", f_attr2 = "fn2"} num_args(1)
  }];

    %0, %c0 = tf_executor.island wraps "tf.AddV2"(%arg0, %arg1) {_tpu_replicate = "cluster"} : (tensor<i32>, tensor<i32>) -> tensor<i32>