// "replicate_on_last_tile_dim" and "last_tile_dims" should be deducted from the
// real Tensor dimensions when tiled.
// For example:
// f32[8,512](sharding={devices=[1,1,2]0,1 last_tile_dims={REPLICATED})
// also means a replicated tensor over all devices.
//
// See xla_data.proto for detailed explanations on the fields.
int GetDimsFromXLAShardingTiled(const xla::OpSharding& xla_sharding);

// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

#include "textflag.h"

// Linux/ARM atomic operations.

// Because there is so much variation in ARM devices,
// the Linux kernel provides an appropriate compare-and-swap
// implementation at address 0xffff0fc0.  Caller sets:
//	R0 = old value
//	R1 = new value
//	R2 = addr
//	LR = return address

      "SmallVector<unsigned, 4>", "GetLayoutDependentResults", (ins)
    >,
    InterfaceMethod<
      [{Returns the optimal data layout based on the available devices.}],
      "StringRef", "GetOptimalLayout", (ins "const RuntimeDevices&":$devices)
    >,
    InterfaceMethod<
      [{Updates operation attributes and operands to account for the updated
        data format. If data format is not supported, must return failure.}],

	PodUIDs sets.Set[string]

	// ResourceHandles is a list of opaque resource data for processing by a specific kubelet plugin
	ResourceHandles []resourcev1alpha2.ResourceHandle

	// CDIDevices is a map of DriverName --> CDI devices returned by the
	// GRPC API call NodePrepareResource
	CDIDevices map[string][]string
}

// ClaimInfoStateWithoutResourceHandles is an old implementation of the ClaimInfoState
// TODO: remove in Beta

// RUN: tf-opt %s -split-input-file -tf-xla-broadcast | FileCheck %s

    %2:3 = "tf.UnpackHyp"(%arg1) {_tpu_replicate = "cluster_Fn", _xla_outside_compilation = "0", device = "", max_seq_length = 16 : i64} : (tensor<*x!tf_type.string>) -> (tensor<*xi32>, tensor<*xi32>, tensor<*xf32>)
    %3 = "tf.Reshape"(%2#2, %1) {_tpu_replicate = "cluster_Fn", _xla_outside_compilation = "0", device = ""} : (tensor<*xf32>, tensor<*xi32>) -> tensor<*xf32>

module attributes {tf.devices = ["/job:worker/replica:0/task:0/device:CPU:0", "/job:worker/replica:0/task:0/device:TPU_SYSTEM:0", "/job:worker/replica:0/task:0/device:TPU:0"]} {
   func.func @main(
      %arg0: tensor<2048xi64> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"},
      %arg1: tensor<2048xi64> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}) -> tensor<2048xi32> {

    %1 = "tf.Rank"(%0) {_tpu_replicate = "cluster", device = ""} : (tensor<*xi1>) -> tensor<*xi32>
    %2 = "tf.Range"(%cst_0, %1, %cst_1) {_tpu_replicate = "cluster", _xla_outside_compilation = "0", device = ""} : (tensor<i32>, tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
    %3 = "tf.All"(%0, %2) {_tpu_replicate = "cluster", _xla_outside_compilation = "0", device = "", keep_dims = false} : (tensor<*xi1>, tensor<*xi32>) -> tensor<*xi1>

  "func.func"() ({
  ^bb0(%arg0: tensor<i32>):
    %0 = "tf.Less"(%arg0, %arg0) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<*xi1>
    %1 = "tf.Less"(%arg0, %arg0) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<*xi1>
    %2 = "tf.LogicalAnd"(%0, %1) {device = ""} : (tensor<*xi1>, tensor<*xi1>) -> tensor<*xi1>
    "func.return"(%2) : (tensor<*xi1>) -> ()

  void CopyCPUTensorToDevice(const Tensor* cpu_tensor, Device* device,
                             Tensor* device_tensor, StatusCallback done,
                             bool sync_dst_compute) const override;
  void CopyDeviceTensorToCPU(const Tensor* device_tensor,
                             absl::string_view tensor_name, Device* device,
                             Tensor* cpu_tensor, StatusCallback done) override;