#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tsl/platform/protobuf.h"  // IWYU pragma: keep

// To be used with LLVM_DEBUG.
#define DEBUG_TYPE "stablehlo_opt_quant_spec"

namespace mlir::quant::stablehlo {
namespace {

using ::mlir::stablehlo::DotGeneralOp;
using ::stablehlo::quantization::Method;

          llvm::dbgs() << "  " << *arg;
        }
      }
  };
  (void)debug_clustered_ops;

  LLVM_DEBUG({
    llvm::dbgs() << "\n\n--- Try to merge cluster:\n";
    debug_clustered_ops(src_root);
    llvm::dbgs() << "\n--- With cluster:\n";
    debug_clustered_ops(dst_root);
    LLVM_DEBUG(llvm::dbgs() << "\n--- Diagnostics:\n");
  });

      if (!ret) continue;

      LLVM_DEBUG(ret.print(llvm::dbgs() << "\ncomputed result = "));

      // If worklist is empty, then this is the root query op.
      if (worklist.empty()) {
        LLVM_DEBUG(llvm::dbgs() << "[root node]\n");
        if (auto dea = mlir::dyn_cast<DenseIntElementsAttr>(ret)) {
          if (dea.getNumElements() != 1) {
            LLVM_DEBUG(llvm::dbgs() << "Unexpected number of elements\n");

    if (!call_op.getF().getRootReference().getValue().starts_with(
            kNestedModule))
      return WalkResult::advance();
    // This is a call we need to inline!
    LLVM_DEBUG(llvm::dbgs()
               << "Found call to inline: " << *call_op.getOperation() << "\n");

    auto call_interface = cast<CallOpInterface>(call_op.getOperation());
    auto called_func =

  if (op_iterator_range.empty()) {
    LLVM_DEBUG(llvm::dbgs() << "Function does not have "
                            << GemmStyleOp::getOperationName() << " op.\n");
    return failure();
  }
  if (!isa<RankedTensorType>(
          (*op_iterator_range.begin()).getResult().getType())) {
    LLVM_DEBUG(llvm::dbgs() << GemmStyleOp::getOperationName()

  TextFormat::Printer printer;
  printer.SetSingleLineMode(true);

  std::string method_txtpb;
  if (!printer.PrintToString(method, &method_txtpb)) {
    LLVM_DEBUG(llvm::dbgs() << "Failed to convert Method to textproto\n.");
    return failure();
  }

  // Single line mode might have an extra space at the end, due to the internal
  // details of `Printer`.

      }
    }

    auto emplace_res = var_access_info.per_resource_info.try_emplace(
        resource, VariableAccessInfo());
    if (!emplace_res.second) {
      LLVM_DEBUG(llvm::dbgs()
                 << "Skipping execute that has multiple reads of a variable: "
                 << execute << "\n");
      var_access_info.per_resource_info.shrink_and_clear();
      return var_access_info;

  // Use the result's state if it is immutable and it is the only one result.
  if (op->getNumResults() == 1 && immutable_results_num == 1) {
    return immutable_states.back()->params;
  }

  LLVM_DEBUG(llvm::dbgs()
             << "Quantization parameters are not collected in an ideal place. "
                "Has to fallback values which might introduce errors.\n");

  if (!result.has_value()) return success();
  llvm::StringRef cluster_name = result.value();

  // We found a _replication_info, let's build an island for the full cluster!
  LLVM_DEBUG(llvm::dbgs() << "Processing candidate island: "
                          << *island.getOperation() << "\n");

  // Collect the islands to merge together in this new cluster starting with the
  // given island.