#include "tensorflow/core/platform/logging.h"

#define DEBUG_TYPE "tf-executor-tpu-v1-island-coarsening"

namespace mlir {
namespace tf_executor {

namespace {

constexpr llvm::StringRef kTpuStatusAttr = "_tpu_compilation_status";
constexpr llvm::StringRef kNoReplicationCluster = "__no_replication_cluster";

#define GEN_PASS_DEF_TPUV1BRIDGEEXECUTORISLANDCOARSENINGPASS

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

namespace mlir {
namespace quant {
namespace {

using ::stablehlo::quantization::CalibrationOptions;
using ::stablehlo::quantization::Method;

constexpr StringRef kQuantTraitAttrName = "_tfl_quant_trait";

// Whether the op is a call op to lifted composite function.
bool IsCallToQuantizableLiftedFunction(Operation *op) {
  if (!op) return false;

namespace {
#define GEN_PASS_DEF_LEGALIZETFPASS
#include "tensorflow/compiler/mlir/lite/transforms/passes.h.inc"

constexpr char kUnidirectionalSequenceLstm[] = "tf.UnidirectionalSequenceLstm";
constexpr char kUnidirectionalSequenceRnn[] = "tf.UnidirectionalSequenceRnn";
constexpr char kTfLiteInputIndices[] = "_tflite_input_indices";

// Legalize operations in functions.

//===----------------------------------------------------------------------===//
// ResourceAliasAnalysisInfo
//===----------------------------------------------------------------------===//

namespace {

constexpr char kResourceArgUniqueIdAttr[] = "tf._resource_arg_unique_id";

bool IsResourceAllocatingOp(Operation* op) {
  auto mem_interface = dyn_cast<MemoryEffectOpInterface>(op);
  if (!mem_interface) return false;

//   `kTfSavedModelInitializerInitType` may have arguments with
//   `tf_saved_model.bound_input`. This ordering avoids breaking the argument
//   ordering constraint.
constexpr std::array<StringRef, 2> kInitializerTypesByMergeOrder = {
    kTfSavedModelInitializerRestoreType, kTfSavedModelInitializerInitType};

// This pass moves all ops from initializer functions to the main function. A

#include "tensorflow/compiler/mlir/tensorflow/analysis/resource_alias_analysis.h"

namespace mlir {
namespace TF {
using ResourceId = int64_t;
inline constexpr ResourceId kUnknownResourceId =
    ResourceAliasAnalysis::Info::kUnknownResourceId;
static_assert(kUnknownResourceId < 0, "kUnknownResourceId must be < 0");

// Maps group IDs to branch IDs.

#include "tensorflow/core/tpu/tpu_defs.h"
#include "tensorflow/core/util/debug_data_dumper.h"
#include "tsl/platform/errors.h"

namespace tensorflow {
namespace {
constexpr absl::string_view kGroupSizeAttrName =
    "tf2xla.collective_info.group_size";
constexpr absl::string_view kGroupKeyAttrName =
    "tf2xla.collective_info.group_key";

// Extracts shape from XlaArgument as TensorShape. If shape is a xla::Shape,

  // between transfers. The particular value of '4' is chosen fairly
  // arbitrarily. It may be necessary to make this tunable via
  // XlaDevice::Options.
  static constexpr int kNumDeviceToDeviceStreams = 4;
  device_to_device_streams_.resize(kNumDeviceToDeviceStreams);
}

XlaDevice::~XlaDevice() {
  VLOG(1) << "Destroying XLA device " << jit_device_name_ << " " << this;

std::unique_ptr<OperationPass<func::FuncOp>>
CreateTFExecutorToControlDialectConversion();
}  // namespace mlir

namespace tensorflow {
namespace {
// Data layout supported by TFLite.
constexpr mlir::StringRef kTFLiteDataLayout = "NHWC";
}  // namespace

void AddQuantizationPasses(const mlir::TFL::PassConfig& pass_config,
                           mlir::OpPassManager& pass_manager) {

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "xla/client/sharding_builder.h"
#include "xla/service/hlo_parser.h"
#include "xla/xla_data.pb.h"

namespace tensorflow {
namespace {

constexpr char kNumSplitAttr[] = "num_split";

// Creates a tf::SplitOp that splits 'src_input' into 'num_splits' ways
// in 'split_dimension' dimension and returns the split values.
mlir::LogicalResult CreateSplitOp(const int num_split,