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

// This is mostly safe, we can use i64 but it will break Darwinn case.
// TODO(b/199054442): Change to i64 when Darwinn supports i64.
def LowerUInt32AddV2 : Pat<
  (TF_AddV2Op $lhs, $rhs),
  (TF_CastOp
    (TF_AddV2Op
      (CreateTFCastOpI32 $lhs, /*truncate=*/ConstBoolAttrFalse),

func.func @quant_act_range(%arg0: !tfr.attr, %arg1: f32, %arg2: i64) -> !tfr.tensor {
  %0:2 = "tfr.quant_act_range"(%arg0, %arg1, %arg2) : (!tfr.attr,f32,i64) -> (!tfr.tensor,!tfr.tensor)
  func.return %0#0 : !tfr.tensor
}

// -----

// CHECK-LABEL: quant_rescale
func.func @quant_rescale(%arg0: !tfr.tensor, %arg1: !tfr.tensor, %arg2: i64) -> !tfr.tensor {

    %5 = "quantfork.qcast"(%4) {volatile} : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>

  %one = "tf.Const"() {value = dense<1> : tensor<4xi32>} : () -> tensor<4xi32>
  // CHECK: tf.StridedSlice
  %0 = "tf.StridedSlice"(%input, %zero, %one, %one) {begin_mask = 15 : i64, device = "", ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64} : (tensor<1024x1024x1024x1024xf32>, tensor<4xi32>, tensor<4xi32>, tensor<4xi32>) -> tensor<*xf32>

  func.return %0 : tensor<*xf32>
}

  %1:5 = "tf.While"(%cst_0, %max_iterations, %array_handle, %array_flow, %matrix) {body= @"map/while_body", cond = @"map/while_cond", is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>) ->  (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>)

%cst, %cst_0, %arg0, %cst_1) {T = [i32, i32, i32, f32, f32],_lower_using_switch_merge = true, _num_original_outputs = 5 : i64, _read_only_resource_inputs = [], body = @while_body, cond = @while_cond, device = "", is_stateless = true, output_shapes = [#tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<1x1024>, #tf_type.shape<1024x1024>], parallel_iterations = 10 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<i32>, tensor<1x1024xf32>, tensor<1024x1024xf32>) -> (tensor<i32>,...

          // CHECK-NEXT: "tf._TPUCompileMlir"()
          %compile:2 = "tf_device.launch"() ({
            %b2:2 = "tf._TPUCompileMlir"() {
              NumDynamicShapes = 0 : i64,
              // The metadata encodes 2 parameter and 2 return values.

    Example:

    ```mlir
    %3, %4 = tfr.quant_act_range(%2, %1, %0) :
        (tfr.attr, float, i64) -> (tfr.tensor, tfr.tensor)
    ```
  }];

  let arguments = (ins
      TFR_AttrType:$act,
      F32:$scale,
      I64:$zp);

  let results = (outs TFR_TensorType:$min, TFR_TensorType:$max);

  let assemblyFormat = [{

func (l *dlogger) i(x int) *dlogger {
	return l.i64(int64(x))
}

//go:nosplit
func (l *dlogger) i8(x int8) *dlogger {
	return l.i64(int64(x))
}

//go:nosplit
func (l *dlogger) i16(x int16) *dlogger {
	return l.i64(int64(x))
}

//go:nosplit
func (l *dlogger) i32(x int32) *dlogger {
	return l.i64(int64(x))
}

//go:nosplit
func (l *dlogger) i64(x int64) *dlogger {
	if !dlogEnabled {

          %outputs_2, %control_3 = tf_executor.island wraps "tf.XlaSpmdFullToShardShape"(%outputs_0) {dim = -1 : i64, manual_sharding = "\08\03\1A\02\02\01\22\02\00\01", unspecified_dims = []} : (tensor<2x2xf32>) -> tensor<1x2xf32>