#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"

//===----------------------------------------------------------------------===//
// The post-quantize Passes.
//
namespace mlir {
namespace quant {
namespace {

// Applies all the clean up steps after quantization.
class PostQuantizePass

#include "tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_op_quant_spec.h"

#include <gtest/gtest.h>
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"

namespace mlir::quant {
namespace {

using QuantizationOptions = tensorflow::quantization::QuantizationOptions;
using QuantizationComponentSpec =
    tensorflow::quantization::QuantizationComponentSpec;

  // CHECK-NEXT:      buffer: 1,
  // CHECK-NEXT:      name: "arg0",
  // CHECK-NEXT:      quantization: {
  // CHECK-EMPTY:
  // CHECK-NEXT:      },
  // CHECK-NEXT:      has_rank: true
  // CHECK-NEXT:    }, {
  // CHECK-NEXT:      shape: [ 32 ],
  // CHECK-NEXT:      buffer: 2,
  // CHECK-NEXT:      name: "Const",
  // CHECK-NEXT:      quantization: {
  // CHECK-EMPTY:
  // CHECK-NEXT:      },
  // CHECK-NEXT:      has_rank: true

// CHECK-NEXT:      buffer: 1,
// CHECK-NEXT:      name: "arg0",
// CHECK-NEXT:      quantization: {
// CHECK-EMPTY:
// CHECK-NEXT:      },
// CHECK-NEXT:      has_rank: true
// CHECK-NEXT:    }, {
// CHECK-NEXT:      shape: [ 4 ],
// CHECK-NEXT:      buffer: 2,
// CHECK-NEXT:      name: "Const",
// CHECK-NEXT:      quantization: {
// CHECK-EMPTY:
// CHECK-NEXT:      },
// CHECK-NEXT:      has_rank: true

  // CHECK-NEXT:      buffer: 1,
  // CHECK-NEXT:      name: "arg0",
  // CHECK-NEXT:      quantization: {
  // CHECK-EMPTY:
  // CHECK-NEXT:      },
  // CHECK-NEXT:      has_rank: true
  // CHECK-NEXT:    }, {
  // CHECK-NEXT:      shape: [ 40, 37 ],
  // CHECK-NEXT:      buffer: 2,
  // CHECK-NEXT:      name: "arg1",
  // CHECK-NEXT:      quantization: {
  // CHECK-EMPTY:
  // CHECK-NEXT:      },
  // CHECK-NEXT:      has_rank: true

  // CHECK-NEXT:      buffer: 1,
  // CHECK-NEXT:      name: "arg0",
  // CHECK-NEXT:      quantization: {
  // CHECK-EMPTY:
  // CHECK-NEXT:      },
  // CHECK-NEXT:      has_rank: true
  // CHECK-NEXT:    }, {
  // CHECK-NEXT:      shape: [ 40, 37 ],
  // CHECK-NEXT:      buffer: 2,
  // CHECK-NEXT:      name: "arg1",
  // CHECK-NEXT:      quantization: {
  // CHECK-EMPTY:
  // CHECK-NEXT:      },
  // CHECK-NEXT:      has_rank: true

#include "mlir/Pass/PassInstrumentation.h"  // from @llvm-project

namespace mlir::quant::stablehlo {

// A `PassInstrumentation` that saves quantization report to file after
// `QuantizeCompositeFunctionsPass` is run. It inspects the `ModuleOp` after
// quantization and analyzes the quantizable units and quantization methods
// used. The report file will be saved at the `file_path`. The report file

// Functions for quantization specifications of Uniform Quantized ops.

#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_OPS_UNIFORM_OP_QUANT_SPEC_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_OPS_UNIFORM_OP_QUANT_SPEC_H_

#include <memory>

#include "mlir/IR/Operation.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"

// CHECK-NEXT:         name: "arg0",
// CHECK-NEXT:         quantization: {
// CHECK-EMPTY:
// CHECK-NEXT:         },
// CHECK-NEXT:         has_rank: true
// CHECK-NEXT:       }, {
// CHECK-NEXT:         shape: [ 4 ],
// CHECK-NEXT:         type: INT8,
// CHECK-NEXT:         buffer: 2,
// CHECK-NEXT:         name: "arg1",
// CHECK-NEXT:         quantization: {
// CHECK-EMPTY:
// CHECK-NEXT:         },

// CHECK-NEXT:      quantization: {
// CHECK-EMPTY:
// CHECK-NEXT:      },
// CHECK-NEXT:      shape_signature: [ -1, 3 ],
// CHECK-NEXT:      has_rank: true
// CHECK-NEXT:    }, {
// CHECK-NEXT:      shape: [ 1, 3 ],
// CHECK-NEXT:      buffer: 2,
// CHECK-NEXT:      name: "serving_default_input1:0",
// CHECK-NEXT:      quantization: {
// CHECK-EMPTY:
// CHECK-NEXT:      },