#define GEN_PASS_DECL_LAYOUTASSIGNMENTPASS
#define GEN_PASS_DECL_LEGALIZEHLOTOTFPASS
#define GEN_PASS_DECL_LEGALIZETFGTOTFPASS
#define GEN_PASS_DECL_LOCALIZEVARHANDLESPASS
#define GEN_PASS_DECL_LOWERQUANTIZEDPASS
#define GEN_PASS_DECL_MARKINPUTOUTPUTALIASESPASS
#define GEN_PASS_DECL_MATERIALIZEPASSTHROUGHOP
#define GEN_PASS_DECL_MERGECONTROLFLOWPASS
#define GEN_PASS_DECL_MOVETRANSPOSESPASS
#define GEN_PASS_DECL_ORDERBYDIALECTPASS

limitations under the License.
==============================================================================*/

#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V1_COMPILE_MLIR_UTIL_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V1_COMPILE_MLIR_UTIL_H_

#include <memory>

#include "absl/base/attributes.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/StringRef.h"

  if (cc_ctx->op_device_context() == nullptr) {  // CPU Device
    status->status = absl::FailedPreconditionError(
        "Accessing device stream is not supported for a CPU device.");
    return nullptr;
  } else if (!cc_ctx->op_device_context()->IsPluggableDevice()) {
    status->status = absl::FailedPreconditionError(
        "Accessing device stream is only supported for pluggable devices.");
    return nullptr;

  let description = [{
    The ExecuteOp executes an operation on the specified device using a custom device.
  }];

  let arguments = (ins
    TFDeviceType:$device,
    Variadic<TFTensorType>:$args,
    StrAttr:$node_def,
    I32Attr:$op_key
  );

  let results = (outs
    Variadic<TFTensorType>:$results
  );

                        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);

  assigned_id->setAttr(kICIWeightDistributionMlirBridgeMarker,
                       before_cluster_builder.getBoolAttr(true));
  std::string device = tensorflow::GetDeviceAliasForHostOfLogicalCore(0);
  LaunchOp launch = tensorflow::WrapOpInLaunch(
      &before_cluster_builder, val_bcast.getLoc(), assigned_id, device);

  Value all_reduce =
      CreateAllReduce(replicate, inner_builder, launch.getResult(0));

#include "xla/mlir_hlo/mhlo/transforms/rewriters.h"
#include "xla/mlir_hlo/mhlo/utils/type_conversion.h"
#include "tensorflow/core/lib/monitoring/counter.h"

namespace mlir {
namespace mhlo {
namespace {

#define GEN_PASS_DEF_LEGALIZETF
#include "tensorflow/compiler/mlir/tf2xla/transforms/xla_legalize_tf_passes.h.inc"

auto *mlir_legalization_count = tensorflow::monitoring::Counter<1>::New(

      task_and_device = {task, device};
    }
  }

  return topology;
}

// Determine execution devices when topology and device assignment are defined.
// With a topology device coordinate to task and device mapping, device
// assignment device coordinates can then be mapped to task and device for TPU
// devices. The device assignment array is also validated.
//

                                   TF_Status* status) {
    params->device->hardware_name = hardware_name;
    params->device->device_vendor = vendor;
    params->device->pci_bus_id = pci_bus_id;
  };

  device_fns_.get_numa_node = [](const SP_Device* device) { return 123; };
  device_fns_.get_memory_bandwidth = [](const SP_Device* device) -> int64_t {
    return 54;
  };

// Belows are the attributes in TFE.
// TFE Arguments and Results (Got from "_Arg",
// "_Retval", .etc)
//  NodeDef.device <-> "tf.device"
//  NodeDef.attr <-> "tf."
//
// TFE general operations
//  NodeDef.device <-> "device"
//
// The following two functions are only used for mapping/excluding attributes
// which are inconsistent between TFG and TFE.
//