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

namespace tensorflow {
namespace tf2xla {
namespace internal {

// This pass extracts a CPU computation cluster with `_xla_outside_compilation`
// annotation from the head or tail of a TPU cluster.

namespace {

using mlir::Block;
using mlir::BlockArgument;
using mlir::BoolAttr;
using mlir::ModuleOp;
using mlir::OpBuilder;

  func.return %0#0 : tensor<*xf32>...

  auto opts = builder.opts();
  Node* const0 = ops::SourceOp("InputFloatOp", opts);
  Node* uncompilable_op = ops::UnaryOp("MissingKernel", const0, opts);
  uncompilable_op->AddAttr("_xla_outside_compilation", "0");
  GraphDef graph_def;
  TF_EXPECT_OK(builder.ToGraphDef(&graph_def));

  auto* flib_runtime = GetFunctionLibraryRuntime();

  // Outside compiled ops are considered by default..

    }
  });

  for (auto parallel_execute_op : tmp_parallel_execute_ops) {
    parallel_execute_op.erase();
  }
  // Remove `_xla_outside_compilation` attribute from all ops.  These ops will
  // be outside of the device cluster. The `_xla_outside_compilation` attribute
  // on ops outside of tf_device.cluster don't have any meaning and can lead to
  // errors later on.  These ops were likely lifted out of the

    // CHECK: %2 = "tf.Reshape"(%1#0, %cst) {_xla_outside_compilation = "0"} : (tensor<3360x8xi32>, tensor<3xi32>) -> tensor<3x1120x8xi32>
    // CHECK: %3:3 = "tf.Split"(%cst_0, %2) {_xla_outside_compilation = "0"} : (tensor<i32>, tensor<3x1120x8xi32>) -> (tensor<1x1120x8xi32>, tensor<1x1120x8xi32>, tensor<1x1120x8xi32>)

    // CHECK: [[GATHER:%.+]] = "tf.GatherV2"([[READVAR]], [[INDEX]], [[ZERO]]) <{batch_dims = 0 : i64}> {_xla_outside_compilation = "0"} : (tensor<*xi32>, tensor<?xi32>, tensor<i64>) -> tensor<*xi32>
    // CHECK: return [[GATHER]]
    %1 = "tf.ResourceGather"(%resource, %indices) {_xla_outside_compilation = "0"} : (tensor<*x!tf_type.resource<tensor<*xi32>>>, tensor<?xi32>) -> (tensor<*xi32>)
    tf_device.return %1 : tensor<*xi32>

// attribute on the `tf_device.launch`.
std::unique_ptr<OperationPass<func::FuncOp>> CreateLaunchToDeviceAttributePass(
    bool legacy_graph_export = true);

// Creates a pass to ensure that the `_xla_outside_compilation` and
// tf_device.launch op no longer exist after Outside Compilation is complete.
std::unique_ptr<OperationPass<func::FuncOp>>
CreateVerifyNoOutsideCompilationMarkersPass();

  "tf._XlaHostComputeMlir"(%arg0) {send_key="", recv_key="", host_mlir_module="module  {\0A  func.func @host_func(%arg0: tensor<*xf32>) -> tensor<*xf32> {\0A    %0 = \22tf.Identity\22(%arg0) {_xla_outside_compilation = \22cluster1\22} : (tensor<*xf32>) -> tensor<*xf32> \0A    func.return %0 : tensor<*xf32> \0A  } \0A} \0A"} : (tensor<2xf32>) -> (tensor<2xf32>)
  func.return
}

// -----

  %0 = "tf.Pack"(%arg0) {axis = 0 : i64, _xla_outside_compilation = "1", device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<2x3xf32>) -> tensor<1x2x3xf32>
  func.return %0 : tensor<1x2x3xf32>
  // CHECK: %[[SHAPE:.*]] = "tf.Const"() <{value = dense<[1, 2, 3]> : tensor<3xi32>}> : () -> tensor<3xi32>

  for (Node* n : g.op_nodes()) {
    bool is_lifted_arg;
    string outside_compilation_attr;
    if (TryGetNodeAttr(n->def(), kXlaIsLiftedArgAttrName, &is_lifted_arg) &&
        TryGetNodeAttr(n->def(), "_xla_outside_compilation",
                       &outside_compilation_attr)) {
      TF_RET_CHECK(is_lifted_arg);
      TF_RET_CHECK(n->IsIdentity() || n->type_string() == "Placeholder");