// RUN: tf-opt %s -tf-replicated-clustering-bridge-v2 -tfrt-lower-cluster-to-runtime-ops-tpu 2>&1 | FileCheck %s

// This test verifies that the tail extraction is not terminated prematurely
// due to the outside compilation attribute could be removed in
// canonicalization of Reshape ops.

// Reshape should not be executed on TPU as all are marked by outside
// compilation. And there should be no host-device communication.
// CHECK: tf._TPUCompileMlir

    %control = tf_executor.island wraps "tf.TPUReplicateMetadata"() {_tpu_replicate = "cluster", allow_soft_placement = false, computation_shape = [], device = "", device_assignment = [], host_compute_core...

        // Handle replicated TPU case.
        islands_to_outline.push_back(island_op);
        break;
      }
      auto device_type =
          op.getAttrOfType<StringAttr>(TF::kCompileDeviceTypeAttr);
      if (device_type && device_type.getValue() == TF::kTpuDevice &&
          !op.hasAttrOfType<StringAttr>(TF::kReplicationInfoAttr)) {
        // Handle single-core TPU case (no `TPUReplicateMetadataOp`).

// RUN: tf-opt %s -tf-replicated-clustering-bridge-v2 -tfrt-lower-cluster-to-runtime-ops-tpu 2>&1 | FileCheck %s

// This test verifies that the tail extraction is not terminated prematurely
// in handling tf.If op which would end up with excessive host-device
// communication.

// In this test, all ops other than tf.Rank are marked with outside_compilation
// . So the TPU program should contain tf.Rank op and there should be no
// host-device communication.

#include "tensorflow/compiler/mlir/tensorflow/transforms/set_tpu_infeed_layout.h"
#include "xla/layout.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
#include "xla/shape.h"
#include "xla/stream_executor/tpu/c_api_conversions.h"
#include "xla/stream_executor/tpu/tpu_api.h"
#include "xla/translate/mhlo_to_hlo/type_to_shape.h"

namespace mlir {
namespace mhlo {
namespace {

#define GEN_PASS_DEF_INFEEDOPSXLAADJUSTLAYOUT

  func.func @missing_num_cores_per_replica() {
    // expected-error@+1 {{requires attribute 'num_cores_per_replica'}}

	if err != nil && err != io.ErrUnexpectedEOF {
		return err
	}

	return nil
}

func osCPUInfoName() string {
	modelName := ""
	cpuMHz := ""

	// The 512-byte buffer is enough to hold the contents of CPU0
	buf := make([]byte, 512)
	err := readLinuxProcCPUInfo(buf)
	if err != nil {
		return ""
	}

	scanner := bufio.NewScanner(bytes.NewReader(buf))
	for scanner.Scan() {

          "tf.opA"() : () -> ()
          tf_device.return
        }) {device = "/TPU:2"} : () -> ()
        tf_device.return
      }
      tf_executor.yield
    }
    tf_executor.fetch
  }
  func.return
}

// CHECK: "tf_device.launch"
// CHECK: device = "/TPU:2"
// CHECK: "tf.opA"
// CHECK: "tf_device.launch"
// CHECK: device = "/TPU:2"
// CHECK: "tf.opA"

inline constexpr llvm::StringRef kReplicationInfoAttr = "_replication_info";
// Marks a node for XLA-TPU compilation. The attribute value indicates the
// associated compilation cluster and replication metadata op.
inline constexpr llvm::StringRef kTpuReplicateAttr = "_tpu_replicate";
// Device types.
inline constexpr llvm::StringRef kTpuDevice = "TPU";
// _xla_outside_compilation
inline constexpr llvm::StringRef kXlaOutsideCompilationAttr =

#include "mlir/IR/Value.h"  // from @llvm-project
#include "mlir/Support/LogicalResult.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "xla/xla_data.pb.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"

namespace tensorflow {

inline constexpr llvm::StringRef kInputShardingAttr =
    "input_sharding_configuration";
inline constexpr llvm::StringRef kOutputShardingAttr =