// exported as a TF SavedModel.
void AddCallModuleSerializationPasses(OpPassManager& pm);

// Passes for unpacking quantized ops to int valued StableHLO ops. This is
// useful when uniform quantized types are suboptimal for the hardware. It goes
// through a StableHLO <-> MHLO roundtrip to utilize the MHLOQuantToInt pass.
void AddStablehloQuantToIntPasses(OpPassManager& pm);

// limitations under the License.

// Quantization as a function library with Uniform Quantized Ops for Dynamic
// PTQ
//
// Internal functions should be marked as private. They will be inlined and
// deleted in `InsertQuantizedFunctionsPass`.
//
// For Uniform Quantized op case, attributes are generated during quantize
// composite pass. Therefore, attr_map is set to an empty string.

module {

  bool reduce_type_precision = false;
  // Whether to consider this model a quantized model with quantize/dequantize
  // ops and to convert kernels to quantized kernels wherever appropriate.
  quant::QDQConversionMode qdq_conversion_mode =
      quant::QDQConversionMode::kQDQNone;

  // When set to true, StableHLO Quantizer is run. The full configuration for
  // the quantizer is at `TocoFlags::quantization_config`.

    if (!mlir::isa<mlir::IntegerType>(raw_elem_type)) {
      return absl::InvalidArgumentError(
          "Quantized tensors must be stored as integers");
    }
    storage_type = mlir::cast<mlir::IntegerType>(raw_elem_type);
  }

  // TFlite uses narrow-range [u]int8 for constant buffers of quantized weights.
  // Since we don't know which ones are weights, we represent this optimization

std::unique_ptr<OperationPass<func::FuncOp>> createConvertSimulatedQuantPass();

/// Creates a pass that converts constants followed by a qbarrier to a
/// constant whose value is quantized. This is typically one of the last
/// passes done when lowering to express actual quantized arithmetic in a
/// low level representation. Because it modifies the constant, it is
/// destructive and cannot be undone.

        "float and quantized types"),
    llvm::cl::init(""));

// NOLINTNEXTLINE
opt<std::string> min_values(
    "tf-input-min-values",
    llvm::cl::desc(
        "Sets the lower bound of the input data. Separated by ','; Each entry "
        "in the list should match an entry in -tf-input-arrays. This is "
        "used when -tf-inference-type is a quantized type."),
    llvm::cl::Optional, llvm::cl::init(""));

    SymbolTable symbol_table(call_op->getParentOfType<ModuleOp>());
    auto func_op =
        dyn_cast_or_null<func::FuncOp>(symbol_table.lookup(func_name));
    if (!func_op) return failure();
    // The quantized fusion should have requantize and return ops at the end.
    auto return_op = dyn_cast_or_null<func::ReturnOp>(
        func_op.getRegion().getBlocks().front().getTerminator());
    if (!return_op) return failure();

    %mul = "tf.Mul"(%cast, %scale) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
    func.return %mul : tensor<*xf32>
  }

  // Requantizes and clips to the range of quantized type if there is no specific activation.
  func.func private @internal_requantize_no_activation_fn(%accumulation : tensor<*xi32>,
                         %input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,

  let summary = "Converts constants followed by qbarrier to actual quantized "
                "values";
  let constructor = "mlir::quantfork::createConvertConstPass()";
}

def QuantConvertSimulatedQuant
    : Pass<"quant-convert-simulated-quantization", "func::FuncOp"> {
  let summary = "Converts training-time simulated quantization ops to "
                "corresponding quantize/dequantize casts";

// of "apply<Name of the Quantization Algorithm>Quantization".
// After applying the function, a quantize/dequantize functions are created
// where the body of each function contains a specific quantization algorithm.
// The input of the quantize function has one operand of
// IsValueWithQuantizablePrecision and the output is a tensor with supported
// quantized precision (like int8). For dequantize function, it is the other way
// around.