// RUN: tf-quant-opt %s -quant-insert-quantized-functions='quantization-method=drq' | FileCheck %s
// RUN: tf-quant-opt %s -quant-insert-quantized-functions='quantization-method=drq target-opset=UNIFORM_QUANTIZED' | FileCheck --check-prefix=UQ-CHECK %s

// Empty module
module {
  func.func @simple_fn(%arg0: tensor<*xf32>) -> tensor<*xf32> {
    func.return %arg0 : tensor<*xf32>
  }
}

// CHECK-NOT: func private @internal_calculate_quant_params

      %3 = "tf.Const"() {value = dense<1.0> : tensor<f32>} : () -> tensor<f32>
      %4 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>

bool IsTFQintType(Type type);

// Convert qint type to the corresponding int type. Return original type if it
// is not qint type.
Type GetIntTypeFromTFQint(Type type);

// Check if an op is TF UniformQuantized op.
bool IsTFUniformQuantizedOp(Operation* op);

}  // namespace mlir::quant::tensorflow

def ConvertTFQuantTypes : Pass<"convert-tf-quant-types", "mlir::func::FuncOp"> {
  let summary = "Replace TensorFlow qint types with int types.";

  let description = [{
    Converts TF ops with qint types to int types. Some UniformQuantized ops
    argument/result allow qint type only. For such cases, add qint <-> int
    tf.Cast around the ops so that they are still valid.
  }];

    std::optional<const absl::string_view> mlir_dump_file_prefix) {
  pm.addNestedPass<mlir::func::FuncOp>(
      mlir::quant::CreateConvertFakeQuantToQdqPass());
  if (quantization_options.op_set() == OpSet::UNIFORM_QUANTIZED) {
    pm.addNestedPass<mlir::func::FuncOp>(
        mlir::TF::CreateUnrollBatchMatMulPassPass());
  }
  pm.addPass(mlir::TF::CreateTFShapeInferencePass());

// RUN: tf-quant-opt %s -quant-insert-quantized-functions | FileCheck %s
// RUN: tf-quant-opt %s -quant-insert-quantized-functions='quantization-method=ptq target-opset=UNIFORM_QUANTIZED' --mlir-print-ir-after-all | FileCheck --check-prefix=UQ-CHECK %s

// Empty module
module {
  func.func @simple_fn(%arg0: tensor<*xf32>) -> tensor<*xf32> {
    func.return %arg0 : tensor<*xf32>
  }
}

// CHECK-NOT: func private @internal_rescale_fn

          clEnumValN(OpSet::TF, "TF",
                     "Uses TF ops that mimic quantization behavior"),
          clEnumValN(OpSet::XLA, "XLA", "Uses TF XLA ops"),
          clEnumValN(OpSet::UNIFORM_QUANTIZED, "UNIFORM_QUANTIZED",
                     "Uses TF Uniform Quantized ops"))};

  Option<bool> enable_per_channel_quantization_{
      *this, "enable-per-channel-quantization", llvm::cl::init(false),

// RUN: tf-quant-opt %s -split-input-file -quant-insert-quantized-functions='quantization-method=drq target-opset=UNIFORM_QUANTIZED' -quant-quantize-composite-functions='quantization-method=drq target-opset=UNIFORM_QUANTIZED' -symbol-dce | FileCheck %s

module {
  // TODO(b/260020937): Support transpose_a, transpose_b for matmul.
  func.func @matmul(%arg0: tensor<2x12xf32>) -> (tensor<*xf32>) {

          clEnumValN(OpSet::TF, "TF",
                     "Uses TF ops that mimic quantization behavior"),
          clEnumValN(OpSet::XLA, "XLA", "Uses TF XLA ops"),
          clEnumValN(OpSet::UNIFORM_QUANTIZED, "UNIFORM_QUANTIZED",
                     "Uses TF Uniform Quantized ops"))};

  // Initialize for tests.
  void initializeForTest() {
    if (!test_mode_) return;

// RUN: tf-quant-opt %s -split-input-file -quant-preprocess-op | FileCheck %s

module {
  // For UniformQuantized depthwise convolution, tensor shape should have
  // transformed from [H,W,C,M] to [H,W,1,CxM],
  func.func @depthwise_conv(%arg0: tensor<1x3x4x3xf32>) -> (tensor<*xf32>) {
    %cst_0 = "tf.Const"() {value = dense<0.000000e+00> : tensor<6xf32>} : () -> tensor<6xf32>