#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"

namespace mlir {
namespace TFDevice {
namespace {
constexpr char kDeviceAttr[] = "device";
constexpr char kMiniBatchSplitsAttr[] = "mini_batch_splits";
constexpr char kMiniBatchCsrAttr[] = "mini_batch_in_csr";

#define GEN_PASS_DEF_EMBEDDINGPROGRAMKEYPASS

#include "tensorflow/compiler/mlir/tensorflow/utils/device_util.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/tpu_rewrite_device_util.h"

namespace mlir {
namespace TFDevice {
namespace {
constexpr char kDeviceAttr[] = "device";
constexpr char kReplicaIdAttr[] = "_xla_replica_id";
constexpr char kDeviceOrdinalAttr[] = "device_ordinal";

#define GEN_PASS_DEF_REPLICATETOISLANDPASS

#include "tensorflow/core/platform/random.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"

namespace mlir {
namespace TFTPU {

namespace {

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

  if (!metadata.getAllowSoftPlacement() &&
      !op->hasAttr(TF::kXlaOutsideCompilationAttr))
    return true;
  std::string device = "";
  if (op->hasAttr(TF::kDeviceAttr))
    device = op->getAttrOfType<StringAttr>(TF::kDeviceAttr).str();
  else
    return false;
  if (absl::StrContains(device, TF::kTpuDevice)) return true;
  return false;
}

inline constexpr llvm::StringRef kTpuDevice = "TPU";
// _xla_outside_compilation
inline constexpr llvm::StringRef kXlaOutsideCompilationAttr =
    "_xla_outside_compilation";
// device attr
inline constexpr llvm::StringRef kDeviceAttr = "device";
// Function attribute to signal that a function should be skipped from TPU
// island outlining. The attribute is set in
// `TpuV1BridgeExecutorIslandCoarsening` and removed in the subsequent

#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
#include "tensorflow/compiler/mlir/tensorflow/transforms/host_runtime/runtime_passes.h.inc"

using mlir::OpResult;
using mlir::Region;
using mlir::StringAttr;
using mlir::success;
using mlir::Type;
using mlir::Value;
using mlir::ValueRange;
using mlir::WalkResult;

constexpr llvm::StringRef kDeviceAttr = "device";
constexpr llvm::StringRef kNameAttr = "name";
constexpr llvm::StringRef kNumCoresPerReplicaAttr = "num_cores_per_replica";
constexpr llvm::StringRef kNumReplicasAttr = "num_replicas";

#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/util/device_name_utils.h"

namespace mlir {
namespace TFTPU {

namespace {

constexpr char kDeviceAttr[] = "device";
typedef std::pair<TF::Conv2DOp, int64_t> Conv2DWithBlockSize;

struct BlockArgumentInfo {
  unsigned arg_num;
  unsigned num_users;
};

#define GEN_PASS_DEF_TPUSPACETODEPTHPASS

                                                      kAnyDevice};
constexpr ResourceAndDevice kUnknownResourceAndDevice{kUnknownResourceId,
                                                      kAnyDevice};

constexpr char kDeviceAttr[] = "device";

#define GEN_PASS_DEF_EXECUTORCONVERTCONTROLTODATAOUTPUTSPASS
#define GEN_PASS_DECL_EXECUTORCONVERTCONTROLTODATAOUTPUTSPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

#include "tensorflow/compiler/mlir/tensorflow/utils/mangling_util.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"

namespace mlir {

namespace {

constexpr char kDeviceAttr[] = "device";

#define GEN_PASS_DEF_RESOURCEOPLIFTINGPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_device_passes.h.inc"

// Lift resource operations out of device computation.