#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/config.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/context.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/pass_pipeline.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/saved_model_export.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/saved_model_import.h"

#define GEN_PASS_DEF_LIFTQUANTIZABLESPOTSASFUNCTIONSPASS
#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"

namespace {

using ::stablehlo::quantization::FunctionNameMatcherSpec;
using ::stablehlo::quantization::Method;
using ::stablehlo::quantization::QuantizationSpec;
using ::stablehlo::quantization::QuantizationSpecs;
using ::tsl::protobuf::TextFormat;

// RUN: stablehlo-quant-opt %s -stablehlo-test-lift-quantizable-spots-as-functions-with-quantization-specs="quantization-specs=disable-all-dot-general" \
// RUN:   -split-input-file | FileCheck %s --check-prefix=DISABLE-ALL-DOT-GENERAL

// Tests that `composite_dot_general_fn_1` and its corresponding XlaCallModuleOp
// contains attributes required for quantization, including the
// `_quantization_method` attribute that contains textpb of `Method`.

#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"

namespace stablehlo::quantization::pywrap {

using ::mlir::quant::stablehlo::QuantizeStaticRangePtq;
using ::mlir::quant::stablehlo::QuantizeWeightOnlyPtq;
using ::tensorflow::SignatureDef;

        input_tensor->quantization->scale.clear();
        input_tensor->quantization->zero_point.clear();
        input_tensor->quantization->min.push_back(-1.0);
        input_tensor->quantization->max.push_back(1.0);

        auto& output_tensor = sub_graph->tensors[op->outputs[0]];
        auto shape = output_tensor->shape;
        output_tensor->quantization->scale.clear();

    compatible_with = get_compatible_with_portable(),
    deps = [
        "//tensorflow/compiler/mlir/quantization/common:attrs_and_constraints",
        "//tensorflow/compiler/mlir/quantization/common/quantization_lib",
        "//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_cc",
        "//tensorflow/compiler/mlir/quantization/tensorflow/ops:uniform_op_quant_spec",
        "//tensorflow/compiler/mlir/tensorflow:tensorflow_types",

#include "mlir/IR/PatternMatch.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h"

namespace mlir::quant {

LogicalResult FillAttributesForUniformQuantizedDotOp(
    PatternRewriter& rewriter, Operation* op,
    llvm::StringMap<Attribute>& identifier_to_attr,
    tensorflow::quantization::QuantizationMethod::PresetMethod
        quantization_method,

              "The minimum number of elements in a weights array required to apply quantization.">,
      Option<"enable_float16_quantization_",
              "enable-float16-quantization", "bool",
              "false", "Whether apply float16 quantization. If false, int8 quantization is applied.">,
      Option<"enable_custom_op_quantization_",
              "enable-custom-op-quantization", "std::string", "",

        "//tensorflow/compiler/mlir/lite/quantization/ir:QuantOps",
        "//tensorflow/compiler/mlir/quantization/common:attrs_and_constraints",
        "//tensorflow/compiler/mlir/quantization/common:func",
        "//tensorflow/compiler/mlir/quantization/common:lift_as_function_call",
        "//tensorflow/compiler/mlir/quantization/common:uniform_quantized_types",
        "//tensorflow/compiler/mlir/quantization/common/ir:QuantOps",

#include "tensorflow/compiler/mlir/quantization/tensorflow/cc/run_passes.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
#include "tsl/platform/errors.h"

namespace mlir::quant::stablehlo {

using ::stablehlo::quantization::QuantizationConfig;
using ::tensorflow::quantization::RunPasses;

PreCalibrationComponent::PreCalibrationComponent(
    absl::Nonnull<MLIRContext*> ctx)