syntax = "proto3";

package stablehlo.quantization;

option cc_enable_arenas = true;

// Defines arious options to specify and control the behavior of the
// StableHLO quantizer.
// NEXT ID: 2
message QuantizationOptions {
  QuantizationMethod quantization_method = 1;
}

// NEXT ID: 3
message QuantizationMethod {
  // Quantization Method can be either preset or custom.
  oneof quantization_method {

// None, the default configuration is "do not quantize the model".
// 2) A set of supported operations.
// 3) Unit wise quantization precision.
// 4) Target hardware name.
// NEXT ID: 18
message QuantizationOptions {
  // The default quantization configuration for the model. If the below
  // unit-wise configuration does not exist, we use this quantization
  // configuration for the entire model. For each method, default configuration

void AddQuantizePtqDynamicRangePasses(
    mlir::OpPassManager &pm, const QuantizationOptions &quantization_options,
    std::optional<const absl::string_view> mlir_dump_file_prefix =
        std::nullopt);

void AddQuantizeWeightOnlyPasses(
    mlir::OpPassManager &pm, const QuantizationOptions &quantization_options,
    std::optional<const absl::string_view> mlir_dump_file_prefix =
        std::nullopt);

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

namespace stablehlo {
namespace quantization {

void AddQuantizationPasses(mlir::PassManager& pass_manager,
                           const QuantizationOptions& quantization_options) {
  QuantizationOptions quantization_options_ = quantization_options;

    const QuantizationOptions& quantization_options);

// Applies post-training dynamic-range quantization to the model.
absl::StatusOr<ExportedModel> QuantizeDynamicRangePtq(
    absl::string_view saved_model_path,
    const std::vector<std::string>& signature_keys,
    const std::unordered_set<std::string>& tags,
    const QuantizationOptions& quantization_options);

    default:
      break;
  }
  *quantization_options_.mutable_quantization_method()
       ->mutable_custom_quantization_method() = custom_method;
  return quantization_options_;
}

LogicalResult GetActivationBitWidth(QuantizationOptions quantization_options,
                                    int* bit_width) {
  CustomQuantizationMethod custom_method =
      quantization_options.quantization_method().custom_quantization_method();

QuantizationOptions FillPresetQuantizationOptions(
    QuantizationOptions quantization_options);

// Returns LogicalResult depending on the look up of activation bit width in the
// custom quantization method. If such information exists, returns success,
// otherwise, returns false.
LogicalResult GetActivationBitWidth(QuantizationOptions quantization_options,
                                    int* bit_width);

    const QuantizationComponentSpec expected_bias_component) {
  QuantizationOptions quantization_options;
  quantization_options.mutable_quantization_method()
      ->mutable_preset_quantization_method()
      ->set_preset_method(preset_quantization_options);
  QuantizationOptions filled_quantization_options =
      quant::stablehlo::FillPresetQuantizationOptions(quantization_options);
  for (QuantizationComponentSpec component :

        signature_keys=['serving_default'],
    )

    model = quantize_model.quantize(
        temp_path,
        quantization_options=quantization_options,
        representative_dataset=data_gen(),
    )
    return model

  @test_util.run_in_graph_and_eager_modes
  def test_multiple_conversion_jobs_with_calibration(self):
    # Ensure that multiple conversion jobs with calibration won't encounter any
    # concurrency issue.

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

namespace stablehlo {
namespace quantization {
// Adds passes for quantization of individual quantizable components.
// (i.e. activation, weight, bias)
void AddQuantizationPasses(mlir::PassManager& pass_manager,
                           const QuantizationOptions& quantization_options);

}  // namespace quantization
}  // namespace stablehlo