#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/lite/tools/optimize/operator_property.h"

//===----------------------------------------------------------------------===//
// The prepare-quantize Pass for LSTM.
//
namespace mlir {
namespace TFL {

constexpr double power_of_two_scale = 32768.0;

// Same with the ordering of //tensorflow/compiler/mlir/lite/ir/tfl_ops.td

// RUN: stablehlo-quant-opt %s -split-input-file -stablehlo-quantize | FileCheck %s

// Test that hybrid quantized dot_general is produced when q/dq pair only exists
// for weight.

module attributes {tf_saved_model.semantics} {
  func.func private @quantize_dot_general_fn(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"} {
    %cst = stablehlo.constant dense<3.000000e-01> : tensor<2x3xf32>

  const std::unordered_set<std::string> tags_;

  const absl::flat_hash_map<std::string, tensorflow::SignatureDef>
      signature_def_map_;

  // Signature keys to identify the functions to load & quantize.
  const std::vector<std::string> signature_keys_;
};

// Runs passes to prepare the calibration model.
absl::Status RunCalibrationPasses(mlir::ModuleOp module_op, MLIRContext& ctx,

// The pass to decompose unregistered TF ops with the TFR compose function.
//
namespace mlir {
namespace TFR {

namespace {

// Quantize the float value based on given scale and zero point attributes.
IntegerAttr Quantize(float value, Attribute scale_attr, Attribute zp_attr,
                     OpBuilder builder) {
  double scale = mlir::cast<FloatAttr>(scale_attr).getValueAsDouble();

  pm.addPass(TFL::CreatePostQuantizeRemoveQDQPass());
  if (failed(pm.run(module.get()))) {
    const std::string err(statusHandler.ConsumeStatus().message());
    LOG(ERROR) << "Failed to quantize: " << err;
    return kTfLiteError;
  }

  // Export the results.
  tflite::FlatbufferExportOptions options;
  options.toco_flags.set_force_select_tf_ops(false);

// RUN: tf-quant-opt %s -split-input-file -quant-preprocess-op -quant-prepare-quantize-drq='enable-per-channel-quantization=true' | FileCheck %s

module {
  func.func @matmul(%arg0: tensor<1x2x2x3xf32>) -> (tensor<*xf32>) {
    %cst_0 = "tf.Const"() {value = dense<0.000000e+00> : tensor<2x1024xf32>} : () -> tensor<2x1024xf32>

            _QuantizationComponent.COMPONENT_ACTIVATION
        ].tensor_type
    )
    # Unlike the HISTOGRAM_PERCENTILE method, the HISTOGRAM_MSE method uses
    # num_bits because it actually quantizes and dequantizes values.
    if activation_tensor_type != _TensorType.TENSORTYPE_INT_8:
      raise ValueError(
          'Only TENSORTYPE_INT_8 is supported for HISTOGRAM_MSE calibration'

  if (dot_general_op == nullptr) return std::nullopt;
  const int64_t filter_rank =
      mlir::dyn_cast<ShapedType>(dot_general_op.getOperand(1).getType())
          .getRank();

  // To quantize rhs per-channel, we currently only consider the case where
  // `stablehlo.dot_general` is legalizable to `tfl.fully_connected`.
  const bool is_per_axis_quantizable =
      IsDotGeneralFullyConnected(dot_general_op).value();

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

module {
  func.func @matmul(%arg0: tensor<1x2x2x3xf32>) -> (tensor<*xf32>) {
    %cst_0 = "tf.Const"() {value = dense<0.000000e+00> : tensor<2x1024xf32>} : () -> tensor<2x1024xf32>

// RUN: stablehlo-quant-opt %s -split-input-file -stablehlo-prepare-quantize=enable-per-channel-quantized-weight=true -verify-diagnostics | FileCheck %s

// -----

module {
  // CHECK-LABEL: conv_with_bias_and_relu
  func.func private @conv_with_bias_and_relu(%arg0: tensor<1x3x2x3xf32>) -> tensor<1x2x2x2xf32> {
    %cst = "tf.Const"() {device = "", value = dense<[7.11401462, 7.05456924]> : tensor<2xf32>} : () -> tensor<2xf32>