// CHECK-SAME: {_quantization_method = "static_range_ptq { }"}
// CHECK-SAME: : (tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>, tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>) -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
// CHECK: %[[DCAST_0:.+]] = "quantfork.dcast"(%[[CALL_0]]) :  (tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>) -> tensor<1x3xf32>
// CHECK: return

include "mlir/Pass/PassBase.td"

def ConvertTFQuantOpsToMHLO : Pass<"quant-convert-tf-quant-ops-to-mhlo", "mlir::func::FuncOp"> {
  let summary = "Convert TF Quant ops to MHLO quantizated ops.";

  let description = [{
    Convert TF Quant ops to MHLO quant ops.
  }];

  let constructor = "mlir::quant::stablehlo::CreateConvertTFQuantOpsToMHLOPass()";

      %1 = stablehlo.uniform_quantize %0 : (tensor<1x3x!quant.uniform<i32:f32:1, {6.000000e+0,7.000000e+0,8.000000e+0}>>) -> tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>>
      return %1 : tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>>

// RUN: tf-quant-opt %s -quant-convert-fake-quant-to-qdq | FileCheck %s

func.func @fakeQuantArgs(%arg0: tensor<8x8x8x8xf32>) -> tensor<8x8x8x8xf32> {
  %0 = "tf.FakeQuantWithMinMaxArgs"(%arg0) {
    min = -0.1 : f32, max = 0.2 : f32, num_bits = 8
  } : (tensor<8x8x8x8xf32>) -> tensor<8x8x8x8xf32>
  func.return %0 : tensor<8x8x8x8xf32>
}
// CHECK: func @fakeQuantArgs

  %2 = func.call @quantization_test(%arg0, %0, %1) {tac.interface_name = "func1"} : (tensor<384x512x!quant.uniform<i8:f32, 0.1>>, tensor<128x512x!quant.uniform<i8<-127:127>:f32, 0.01>>, tensor<128x!quant.uniform<i32:f32, 0.7>>) -> tensor<384x128x!quant.uniform<i8:f32, 0.09:-4>>
  func.return %2 : tensor<384x128x!quant.uniform<i8:f32, 0.09:-4>>
}

// PARTIAL-LABEL: @quantization_test

func.func @remove_quantized_io(
  %arg0: tensor<1x10x!quant.uniform<i8:f32, 0.1:-128>>,
  %arg1: tensor<1x5xf32>) -> (tensor<1x10x!quant.uniform<i8:f32, 0.2:42>>, tensor<1x5xf32>) {
  %0 = "tf.MyRequantize"(%arg0) : (tensor<1x10x!quant.uniform<i8:f32, 0.1:-128>>) -> tensor<1x10x!quant.uniform<i8:f32, 0.2:42>>
  %1 = "tf.Intermediate"(%arg1) : (tensor<1x5xf32>) -> tensor<1x5xf32>
  func.return %0, %1 : tensor<1x10x!quant.uniform<i8:f32, 0.2:42>>, tensor<1x5xf32>

  %requant = stablehlo.uniform_quantize %conv : (tensor<1x3x4x2x!quant.uniform<i32:f32, 5.8949912267181218E-5>>) -> tensor<1x3x4x2x!quant.uniform<i8:f32, 0.045673100153605144:-62>>

    return "quant-convert-fake-quant-to-qdq";
  }

  StringRef getDescription() const final {
    // This is a brief description of the pass.
    return "Convert Fake Quant op to quant.qcast and quant.dcast pairs";
  }

  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<TF::TensorFlowDialect>();
    registry.insert<quant::QuantizationDialect>();

// CHECK: return %[[QUANT]]

// -----

// Tests that the pattern doesn't match when the input tensor's type is a
// quantized type.

func.func @uniform_quantize_op_quantized_input(%arg: tensor<2x2x!quant.uniform<i8:f32, 2.000000e+0:16>>) -> tensor<2x2x!quant.uniform<i8:f32, 3.000000e+0:127>> {

  // CHECK-SAME: quant.uniform<i8:f32, 0.023529411764705882:-128>
  // CHECK: %[[dq2:.*]] = "quantfork.dcast"(%[[q2]])
  // CHECK-SAME: quant.uniform<i8:f32, 0.023529411764705882:-128>
  // CHECK: %[[q3:.*]] = "quantfork.qcast"(%[[cst3]])
  // CHECK-SAME: quant.uniform<i8:f32, 3.9215686274509805E-9>
  // CHECK: %[[dq3:.*]] = "quantfork.dcast"(%[[q3]])
  // CHECK-SAME: quant.uniform<i8:f32, 3.9215686274509805E-9>