#include "mlir/Support/LLVM.h"  // from @llvm-project

namespace mlir {
namespace quant {
namespace {

// Prefix and suffix to the QuantizationUnit string representation.
constexpr std::string_view kQuantizationUnitPrefix = "QuantizationUnit(";
constexpr std::string_view kQuantizationUnitSuffix = ")";

// Concatenates node name and func name with a "@" separator.

  } else {
    for (Location child_loc : locations) {
      FindQuantizationUnitsRecursively(child_loc, units);
    }
  }
}

// Finds the QuantizationUnit from location.
std::optional<QuantizationUnit> FindQuantizationUnit(Operation* op) {
  SmallVector<QuantizationUnit> quant_units;
  FindQuantizationUnitsRecursively(op->getLoc(), quant_units);

  if (quant_units.size() == 1) {
    return *quant_units.begin();
  }

}

// `DumpTensor` op saves entire value of input to as a tensor proto into a
// specified directory and filename. When enabled is set to false, op is
// disabled and won't save any value. It also creates `QuantizationUnit` proto
// with `func_name` and `node_name` to identify the op.
REGISTER_OP("DumpTensor")
    .Input("tensor_data: T")
    .Attr("log_dir_path: string")
    .Attr("file_name: string")
    .Attr("T: type")

//===----------------------------------------------------------------------===//
// The prepare-quantize-drq Pass.
//
namespace mlir {
namespace quant {

namespace {

using QuantizationUnit = std::pair<Operation*, int>;
using QuantizationUnits = llvm::SetVector<QuantizationUnit>;
using ::tensorflow::quantization::OpSet;

// Applies prepare quantization on the model in TF dialect for dynamic range
// quantization case.
class PrepareQuantizeDRQPass

    %1 = "tf.PartitionedCall"(%arg0, %cst, %cst_0) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_1} : (tensor<1x2x2x3xf32>, tensor<2x2x3x2xf32>, tensor<2xf32>) -> tensor<*xf32> loc(callsite("test@conv"("Conv2D_1") at "QuantizationUnit(\12\08Conv2D_1\1a\04conv)"))

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"

namespace mlir {
namespace quant {
namespace {

using QuantizationUnit =
    ::tensorflow::quantization::UnitWiseQuantizationSpec::QuantizationUnit;
using ::tensorflow::quantization::OpSet;
using ::tensorflow::quantization::QuantizationComponentSpec;
using ::tensorflow::quantization::QuantizationMethod;

//
namespace mlir {
namespace quant {

namespace {

using QuantMethod =
    ::tensorflow::quantization::QuantizationMethod::PresetMethod;
using QuantizationUnit = std::pair<Operation*, int>;
using QuantizationUnits = llvm::SetVector<QuantizationUnit>;
using ::tensorflow::quantization::OpSet;

// Preprocesses ops to allow multi-axis quantization, prior to quantization

// contains the scope of the unit, and the quantization method for each unit.
// NEXT ID: 7
message UnitWiseQuantizationSpec {
  // Quantization unit granularity.
  // NEXT ID: 4
  message QuantizationUnit {
    // Type of the op, ex: Conv2D, MatMul, Einsum... The node_name field can
    // be omitted if it is intended to match all nodes with this type.
    string op_type = 1;

// filled and node_name is always set to "_empty_node".
std::pair<std::string, std::string> GetFuncNameAndNodeName(
    TF::PartitionedCallOp call_op, const FlatSymbolRefAttr &f_attr) {
  std::optional<QuantizationUnitLoc::QuantizationUnit> quant_unit =
      FindQuantizationUnitFromLoc(call_op->getLoc());
  return std::make_pair(quant_unit->func_name(), quant_unit->node_name());
}

std::pair<std::string, std::string> GetFuncNameAndNodeName(

      return %0 : tensor<2x2xf32>
    }
    func.func private @composite_fn_1(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf32>) -> tensor<2x2xf32> attributes {tf_quant.composite_function} {