// RUN: tf-quant-opt %s -quant-insert-quantized-functions='quantization-method=weight_only target-opset=XLA' | FileCheck %s

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

// CHECK-NOT: func private @internal_dequantize_f32
// CHECK-NOT: func private @internal_conv3d_fn
// CHECK-NOT: func private @internal_batch_matmul_fn

include "tensorflow/compiler/mlir/lite/ir/tfl_ops.td"

// Quantize attribute $0 by using quantization parameter from %1.
def QuantizeByQuantizedType : NativeCodeCall<"quant::Quantize($0, $1.getValue())">;
def F32ElementsAttr : ElementsAttrBase<
  CPred<"$_self.cast<ElementsAttr>().getShapedType().getElementType().isF32()">, "float constant tensor">;

// Squash tfl.dequantize and tfl.quantize pairs.

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

  Option<QuantMethod> quantization_method_{
      *this, "quantization-method",
      llvm::cl::init(tensorflow::quantization::QuantizationMethod::
                         METHOD_STATIC_RANGE_INT8),

    "tf-custom-opdefs", llvm::cl::desc("List of custom opdefs when importing "
                                       "graphdef"));

// Quantize and Dequantize ops pair can be optionally emitted before and after
// the quantized model as the adaptors to receive and produce floating point
// type data with the quantized model. Set this to `false` if the model input is
// integer types.
// NOLINTNEXTLINE
opt<bool> emit_quant_adaptor_ops(

  static bool IsQuantizableCustomOp(Operation* op,
                                    const quant::CustomOpMap& custom_op_map) {
    // In some cases, ops may need to be quantized even though their op trait is
    // not quantizable. For example, for the case of custom op various ops can
    // be categorized as cusom ops despite each of them may require different

  // The following two passes find specific uniform quantization patterns in
  // StableHLO and converts them to TFLite ops that accept or produce uniform
  // quantized types. They only target a specific set of models that contain
  // "decomposed" quantized ops produced from the framework level. This is why
  // they are placed right after the `LegalizeTFXlaCallModuleToStablehloPass`

      != _PresetMethod.METHOD_STATIC_RANGE_WEIGHT_ONLY_INT8
  ):
    raise ValueError(
        'StableHLO quantized opset currently only supports static range'
        ' quantization and weight-only quantizationvia TF Quantizer.'
    )

  # Set `force_graph_mode_calibration` to True to avoid skipping op execution,
  # which are not connected to return ops, during calibration execution.

    return success();
  }
};

// UniformDequantizeOp takes TF quantized types as input which would have been
// converted to the mhlo quantized types. Use OpConversionPattern in order to
// retrieve the operand type *after* conversion, using OpAdaptor operand
// accessor.
// Same for other Uniform Quant Ops that take TF quantized types as input.
class ConvertUniformDequantizeOp

}

// CHECK-NOT: stablehlo.transpose
// CHECK: %[[CONV:.+]] = stablehlo.convolution
// CHECK-SAME{LITERAL}: [b, f, 0, 1]x[o, i, 0, 1]->[b, 0, 1, f]
// CHECK-NOT: stablehlo.transpose

// -----

// Tests that a quantized convolution does not match. No conversion occurs.

// CHECK-LABEL: quantized_convolution

        per_axis_type.getStorageTypeMin(), per_axis_type.getStorageTypeMax());
  }

  auto quantize = builder.create<quantfork::QuantizeCastOp>(
      q_op.getLoc(), new_value_type.clone(new_qtype), new_value);
  auto dequantize = builder.create<quantfork::DequantizeCastOp>(
      dq_op.getLoc(), new_value_type, quantize.getResult());
  return dequantize.getResult();
}