#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"

#define DEBUG_TYPE "tf-tpu-merge-variables-with-execute"

namespace mlir {
namespace TFTPU {

namespace {
constexpr char kAliasingAttr[] = "tf.aliasing_output";
constexpr char kDeviceAttr[] = "device";
constexpr char kFuncDeviceAttr[] = "tf.device";

#define GEN_PASS_DEF_TPUMERGEVARIABLESWITHEXECUTEPASS

//
//      // Set functions to create a new optimizer.
//      params->optimizer->struct_size = TP_OPTIMIZER_STRUCT_SIZE;
//      params->optimizer->create_func = (My_optimizer::create_func);
//    }

#define GO_MAJOR 0
#define GO_MINOR 0
#define GO_PATCH 1

#ifdef __cplusplus
extern "C" {
#endif

// TF_TriState is the C API typedef for tri-state.
typedef enum TF_TriState {
  TF_TriState_Default = 0,
  TF_TriState_Off,

      llvm::dyn_cast<tf_device::LaunchOp>(preprocess_op->getParentOp());
  // Device of original launch looked up before moving the preprocessing ops
  // around.
  StringAttr device =
      original_launch_op->getAttrOfType<StringAttr>(kDeviceAttr);

  if (!device) {
    return original_launch_op->emitOpError()
           << "Launch op has an invalid device attribute.";
  }

  // Find the TPUCompile successor.

// Controls how to act when we try to run an operation on a given device but
// some input tensors are not on that device.
// LINT.IfChange
// Note: Keep in sync with internal copy of enum in eager/context.h.
typedef enum TFE_ContextDevicePlacementPolicy {
  // Running operations with input tensors on the wrong device will fail.
  TFE_DEVICE_PLACEMENT_EXPLICIT = 0,
  // Copy the tensor to the right device but log a warning.
  TFE_DEVICE_PLACEMENT_WARN = 1,

    return success();
  }
};

const char kDeviceAttr[] = "device";
const char kDeviceGpu[] = "GPU";

std::optional<std::string> GetDevice(mlir::Operation *op) {
  mlir::StringAttr device = op->getAttrOfType<mlir::StringAttr>(kDeviceAttr);
  if (!device || device.getValue().empty()) {
    return std::nullopt;
  }
  const std::string device_name = device.str();
  tensorflow::DeviceNameUtils::ParsedName parsed_name;

  Scope WithNoControlDependencies() const;

  /// Return a new scope. All ops created within the returned scope will have
  /// the device field set to 'device'.
  Scope WithDevice(const string& device) const;

  /// Returns a new scope.  All ops created within the returned scope will have
  /// their assigned device set to `assigned_device`.
  Scope WithAssignedDevice(const string& assigned_device) const;

namespace TFTPU {

namespace {

constexpr char kDeviceAttr[] = "device";
constexpr char kFuncDeviceAttr[] = "tf.device";
constexpr char kDefaultShardingValue[] = "";
constexpr char kMirroredVariableIndicesAttr[] = "_mirrored_variable_indices";

std::string GetRandomStateVariableName() {
  return absl::StrCat("VariablesFormatState_", tensorflow::random::New64());
}

#define GEN_PASS_DEF_TPUVARIABLERUNTIMEREFORMATTINGPASS

    // TODO(jpienaar): Remove this later.
    if (auto device = res->getAttrOfType<StringAttr>("device")) {
      if (!device.getValue().empty())
        result_id->setAttr("device", device);
      else
        result_id->setAttr("device", compile_device_op);
    } else if (compile_device_op) {
      result_id->setAttr("device", compile_device_op);
    }

                        XlaRunOp);
REGISTER_KERNEL_BUILDER(
    Name("_XlaRun").Device(DEVICE_DEFAULT).HostMemory("key"), XlaRunOp);

REGISTER_KERNEL_BUILDER(Name("_XlaMerge").Device(DEVICE_CPU), XlaMergeOp);
REGISTER_KERNEL_BUILDER(Name("_XlaMerge").Device(DEVICE_GPU), XlaMergeOp);
REGISTER_KERNEL_BUILDER(Name("_XlaMerge").Device(DEVICE_DEFAULT), XlaMergeOp);

        return failure();
      }
      read_use.set(new_reads[0]);
    }
  }
  return success();
}

namespace {
constexpr char kDeviceAttr[] = "device";
constexpr char kFuncDeviceAttr[] = "tf.device";
mlir::Attribute GetDeviceOfResource(mlir::func::FuncOp func,
                                    mlir::Value resource) {
  if (auto* resource_op = resource.getDefiningOp()) {