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 {
    PresetQuantizationMethod preset_quantization_method = 1;
    CustomQuantizationMethod custom_quantization_method = 2;
  }
}

// Preset model quantization method for optimization.
//
// Common quantization methods are defined as preset methods in this message.
// Note that the quantization specs (ex: bit width) for preset quantization
// methods are fixed. To use a different quantization spec for a given method,
// use CustomQuantizationMethod.
// NEXT ID: 2
message PresetQuantizationMethod {
  // Preset quantization methods that are supported as a stable API.
  // NEXT ID: 5
  enum PresetMethod {
    // TODO(b/266173150): Update preset methods after redefining quantization
    // pattern matching in DarwiNN.
    // This should never be used. Using this will generally result in an error.
    METHOD_UNSPECIFIED = 0;  // go/do-include-enum-unspecified

    // Apply default weight-only quantization. Weights are quantized during
    // conversion, then dequantized during inference.
    // Activation: f32, Weight: qi8, Bias: f32
    WEIGHT_ONLY = 1;

    // Apply default dynamic range quantization. Quantized tensor value's
    // ranges are determined during graph runtime.
    // Activation: f32, Weight: qi8, Bias: f32
    POST_TRAINING_QUANTIZATION_DYNAMIC_RANGE = 2;

    // Apply float16 quantization to all the weights. Quantized weights will be
    // dequantized before running inference.
    // Activation: f32, Weight: f16, Bias: f16
    FLOAT16 = 3;

    // Apply static range quantization. The quantization range is determined
    // via calibration phase and quantized during conversion.
    // Activation: qi8, Weight: qi8, Bias: qi32
    POST_TRAINING_QUANTIZATION_STATIC_RANGE = 4;
  }
  PresetMethod preset_method = 1;
}

// Custom option for specifying quantization spec details.
// If the selected quantization option is not available, StableHLO quantizer
// will raise an error.
// NEXT ID: 2
message CustomQuantizationMethod {
  // Specify component name, bit width, and other specs for all compoenents
  // intended to be quantized.
  repeated QuantizationComponentSpec quantization_component_spec = 1;
}

// Quantization spec per each component designated to be quantized.
// Components whose QuantizationComponentSpec is not specified will not be
// quantized, and remain f32.
// NEXT ID: 7
message QuantizationComponentSpec {
  // NEXT ID: 4
  enum QuantizationComponent {
    COMPONENT_UNSPECIFIED = 0;
    COMPONENT_ACTIVATION = 1;
    COMPONENT_WEIGHT = 2;
    COMPONENT_BIAS = 3;
  }

  // NEXT ID: 5
  enum BitWidth {
    BIT_WIDTH_UNSPECIFIED = 0;
    BIT_WIDTH_4 = 1;
    BIT_WIDTH_8 = 2;
    BIT_WIDTH_16 = 3;
    BIT_WIDTH_32 = 4;
  }

  // NEXT ID: 4
  enum BitType {
    BIT_TYPE_UNSPECIFIED = 0;
    BIT_TYPE_INT = 1;
    BIT_TYPE_FLOAT = 2;
    BIT_TYPE_BFLOAT = 3;
  }

  QuantizationComponent quantization_component = 1;

  // Defines the target bit of the data.
  BitWidth bit_width = 2;

  // Defines the type of data of the quantized component.
  BitType bit_type = 3;

  // Defines whether quantization is done in narrow range.
  bool enable_narrow_range = 4;

  // Defines whether quantiation is done per-channel.
  bool enable_per_channel_quantization = 5;

  // Defines whether quantization is done symmetrically.
  bool enable_symmetric = 6;
}
