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

// The XlaCompileOnDemandOp is an OpKernel that, when its Compute method is
// called, will generate an xla::Computation and run it asynchronously.

#ifndef TENSORFLOW_COMPILER_JIT_XLA_COMPILE_ON_DEMAND_OP_H_
#define TENSORFLOW_COMPILER_JIT_XLA_COMPILE_ON_DEMAND_OP_H_

#include <vector>

  func.return %0#0 : tensor<*xi32>
}

// RUN: tf-opt %s -split-input-file -tf-tpu-host-computation-expansion | FileCheck %s

// Tests expansion of a outside compiled ops at head/tail of TPU computation.

// CHECK-LABEL: func @identity_at_head_expanded
func.func @identity_at_head_expanded(%arg0: tensor<?xi32>) {
  // CHECK: "tf_device.cluster"
  // CHECK-NEXT: "tf.Identity"
  // CHECK-SAME: _xla_outside_compilation = ""
  "tf_device.cluster"() ({

       * to a valid final state ({@link #COMPLETED}, {@link #CANCELLED}, or {@link #INTERRUPTED}).
       *
       * @param v the value to set as the result of the computation.
       * @param t the exception to set as the result of the computation.
       * @param finalState the state to transition to.
       */
      private boolean complete(@Nullable V v, @Nullable Throwable t, int finalState) {

    $mlir_module is a serialized MLIR module with a `main` function that contains target computation.
    $metadata is a serialized TPUCompileMetadataProto describing the shapes and types of the inputs to the computation, as well as a mapping onto the TPU pod topology.

    %add_result = "tf.Add"(%computation, %computation) : (tensor<i32>, tensor<i32>) -> tensor<i32>
    // expected-error@+1 {{Input of TPUPartitionedOutputV2 must be in tpu computation.}}
    %partitioned_output:2 = "tf.TPUPartitionedOutputV2"(%add_result) {N = 2 : i64, partition_dims = []} : (tensor<i32>) -> (tensor<i32>, tensor<i32>)

  xla::ExecutableBuildOptions build_options = GetExecutableBuildOptions(
      options, result, client_->default_device_ordinal());
  TF_ASSIGN_OR_RETURN(
      auto executables,
      client_->Compile(*result.computation, argument_layouts, build_options));
  TF_RET_CHECK(executables.size() == 1);
  return std::move(executables[0]);
}

absl::StatusOr<std::string> XlaDeviceCompilerClient::SerializeExecutable(

namespace tf2xla {
namespace v2 {

// Legalizes the given mlir::Module into XLA HLO. If successful, returns the
// compiled XLA HLO. V1 of the tf2xla uses MLIR whereas V0 does not use MLIR.
//
// Inputs:
//  computation - The MLIR module op. It currently takes in
//  tpu::FunctionToHloArgs but this is deprecated. arg_shapes - The shapes of
//  the arguments in module_op. device_type - The device type to compile for.

	for len(p) >= chunk {
		// Can interlace the computation of w with the
		// rounds below if needed for speed.
		for i := 0; i < 16; i++ {
			j := i * 4
			w[i] = uint32(p[j])<<24 | uint32(p[j+1])<<16 | uint32(p[j+2])<<8 | uint32(p[j+3])
		}

		a, b, c, d, e := h0, h1, h2, h3, h4

		// Each of the four 20-iteration rounds
		// differs only in the computation of f and
		// the choice of K (_K0, _K1, etc).
		i := 0

  func.return %0#0 : tensor<*xi32>
}

// CHECK: func private @compute_0