#include "tensorflow/core/tfrt/utils/fallback_tensor.h"
#include "tfrt/bef/bef_buffer.h"  // from @tf_runtime
#include "tfrt/bef_executor/bef_file.h"  // from @tf_runtime
#include "tfrt/core_runtime/core_runtime.h"  // from @tf_runtime
#include "tfrt/host_context/async_value.h"  // from @tf_runtime
#include "tfrt/host_context/async_value_ref.h"  // from @tf_runtime
#include "tfrt/host_context/chain.h"  // from @tf_runtime

// RUN: tf-tfrt-opt %s -inline | FileCheck %s

func.func @_tfrt_fallback_init(%arg0: !tfrt.chain) -> !tfrt.chain {
  %0 = tfrt_fallback_async.createop(%arg0) key(0) device("/device:CPU:0") "tf.Less"() {T = i32} num_args(2)
  tfrt.return %0 : !tfrt.chain
}

func.func @callee(%ch: !tfrt.chain, %arg: !tfrt_fallback.tf_tensor) -> (!tfrt.chain, !tfrt_fallback.tf_tensor) {
  %const = tfrt_fallback_async.const_dense_tensor dense<9> : tensor<i32> {_tfrt_cost = 1 : i64}

// CHECK: tfrt.return [[pred_res]]#0, [[body_res]]#1, [[pred_res]]#1 : !tfrt.chain, !tfrt_fallback.tf_tensor, i1

// CHECK: func @"while_cond_lt9/tfrt_predicate"([[ch:%.*]]: !tfrt.chain, [[arg:%.*]]: !tfrt_fallback.tf_tensor) -> (!tfrt.chain, i1)

    visibility = [
        # copybara:uncomment "//learning/brain/experimental/tfrt:__subpackages__",
        # copybara:uncomment "//learning/brain/tfrt/tpu/compiler:__subpackages__",
        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
        "//tensorflow/core/runtime_fallback:internal",
        "//tensorflow/core/tfrt/saved_model:friends",
        # copybara:uncomment "//third_party/tf_runtime_google:__subpackages__",

  // CHECK: [[ch:%.*]] = tfrt_fallback_async.createop(%{{.*}}) key(100) device("cpu") "tf.AddV2"() {T = i32} num_args(2)
  %out_ch = tfrt_fallback_async.createop(%in_ch) key(100) device("cpu") "tf.AddV2"() {T = i32} num_args(2)

  // CHECK: tfrt.return [[ch]]
  tfrt.return %out_ch: !tfrt.chain
}

// CHECK-LABEL: func @fallback_resource
func.func @fallback_resource(%ch0: !tfrt.chain) -> !tfrt.chain {

#include "tfrt/host_context/function.h"  // from @tf_runtime
#include "tfrt/host_context/host_context.h"  // from @tf_runtime
#include "tfrt/host_context/resource_context.h"  // from @tf_runtime
#include "tfrt/support/ref_count.h"  // from @tf_runtime

namespace tensorflow {

using ::tfrt::AsyncValue;
using ::tfrt::BEFFile;
using ::tfrt::ExecutionContext;
using ::tfrt::Function;
using ::tfrt::MakeAvailableAsyncValueRef;
using ::tfrt::RCReference;

#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback.h"
#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_common.h"
#include "tfrt/basic_kernels/opdefs/tfrt_base.h"  // from @tf_runtime
#include "tfrt/basic_kernels/opdefs/types.h"  // from @tf_runtime
#include "tfrt/core_runtime/opdefs/attributes.h"  // from @tf_runtime
#include "tfrt/core_runtime/opdefs/core_runtime.h"  // from @tf_runtime

        "//tensorflow/core/tfrt/ifrt/...",
        "//tensorflow/core/tfrt/kernels/...",
        "//tensorflow/core/tfrt/mlrt/...",
        "//tensorflow/core/tfrt/saved_model/...",
        "//tensorflow/core/tfrt/tfrt_session/...",
    ] + if_google([
        "//learning/brain/experimental/mlir/tflite/tfmrt/...",
        "//learning/brain/experimental/mlir/tfrt_compiler/...",
        "//learning/brain/experimental/tfrt/...",

#include "tensorflow/core/platform/threadpool_interface.h"
#include "tfrt/bef/bef_buffer.h"  // from @tf_runtime
#include "tfrt/bef_executor/bef_file.h"  // from @tf_runtime
#include "tfrt/host_context/execution_context.h"  // from @tf_runtime
#include "tfrt/host_context/host_context.h"  // from @tf_runtime
#include "tfrt/host_context/resource_context.h"  // from @tf_runtime
#include "tfrt/support/ref_count.h"  // from @tf_runtime

namespace tensorflow {

  %4 = corert.executeop(%1) "tf.AddV2"(%3, %3) {T = f32, device = "/device:GPU:0"} : 1
  tfrt.return
}

// CHECK: func @test_transfer_func_arg(%[[ARG_0:.*]]: !tfrt.chain, %[[ARG_1:.*]]: !corert.tensorhandle
func.func @test_transfer_func_arg(%arg0: !tfrt.chain, %arg1: !corert.tensorhandle {tfrt.device = "/device:CPU:0"}) -> () {
  // CHECK-NEXT: %[[RESULT_0:.*]] = corert.get_op_handler %[[ARG_0]] "gpu"