namespace mlir::quant::stablehlo {

// Performs post-calibration graph transformation as part of post-training
// static-range quantization.
//
// The resulting `ModuleOp` contains quantized StableHLO ops serialized in
// `TF::XlaCallModuleOp`s. They are quantized using the statistics collected
// after the calibration step, corresponding to each `TF::CustomAggregatorOp`s
// in the input module op.

// Passes only used for testing purposes.

include "mlir/Pass/PassBase.td"

def TestPreCalibrationComponentPass : Pass<"stablehlo-test-pre-calibration-component", "mlir::ModuleOp"> {
  let summary = "Test-only pass to test the PreCalibrationComponent.";
  let description = [{
    Runs the pre calibration passes for post-training quantization with default
    configuration.
  }];
  let dependentDialects = [

    const ::stablehlo::quantization::DebuggerConfig& debugger_config);

// Adds passes for static-range quantization post-calibration. Utilizes tensor
// statistics collected from the calibration step and performs quantization.
void AddPostCalibrationPasses(
    OpPassManager& pm,
    const ::stablehlo::quantization::PipelineConfig& pipeline_config,
    const ::stablehlo::quantization::QuantizationSpecs& specs);

#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"

namespace stablehlo::quantization {

// Reads the calibration statistics from the given directory.
absl::StatusOr<absl::flat_hash_map<
    std::string, tensorflow::calibrator::CalibrationStatistics>>
ReadStatistics(absl::string_view calibration_data_dir);

#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"

namespace mlir::quant::stablehlo {

// Performs pre-calibration graph transformation as part of post-training
// static-range quantization.

// The resulting `ModuleOp` contains `TF::CustomAggregatorOp`s for collecting

  string unquantized_dump_model_path = 2;

  // Path to save debugger related logs. Defaults to '/tmp/dumps'.
  string log_dir_path = 3;
}

// Defines various calibration options.
// Next ID: 6
message CalibrationOptions {
  // Configurations for calibration methods.
  // Next ID: 7
  enum CalibrationMethod {
    CALIBRATION_METHOD_UNSPECIFIED = 0;
    // Use the min, max values of all sample datasets.

// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Calibration used to determine thresholds for using
// different algorithms.  Ideally, this would be converted
// to go generate to create thresholds.go

// This file prints execution times for the Mul benchmark
// given different Karatsuba thresholds. The result may be

#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_base.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"

namespace tensorflow {
namespace calibrator {

using ::stablehlo::quantization::CalibrationOptions;

// MinMax calibration calculates the global min and global max values.

// SavedModel for calibration. Similarly, `saved_model_dir` is required to
// access the assets of the original model. `saved_model_tags` uniquely
// identifies the `MetaGraphDef`. `quantization_config` determines the behavior
// of StableHLO Quantizer. `quantization_py_function_lib` contains python
// implementations of certain APIs that are required for calibration.

                  const tensorflow::quantization::PyFunctionLibrary*
                      quantization_py_function_library = nullptr);

// Quantize the model with calibration data. Throw errors if `fully_quantize`
// is specified by the calibration data are not sufficient to quantize the
// model.
PyObject* MlirQuantizeModel(PyObject* data, bool disable_per_channel,
                            bool fully_quantize, int inference_type,