==============================================================================*/

#ifdef TFRT_FALLBACK_SYNC_OPS
#else
#define TFRT_FALLBACK_SYNC_OPS

include "tfrt/tfrt_op_base.td"
include "tfrt/core_runtime/opdefs/corert_base.td"
include "tfrt/core_runtime/opdefs/corert_traits.td"
include "tfrt_fallback.td"
include "mlir/Interfaces/SideEffectInterfaces.td"
include "tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td"

// RUN: tf-tfrt-opt -tf-to-tfrt %s | FileCheck %s --dump-input=fail --dump-input-filter=all
// RUN: tf-tfrt-opt -pass-pipeline='builtin.module(tf-to-tfrt{target-tpurt=true tpu-use-core-selector=false})' %s | FileCheck %s --dump-input=fail --dump-input-filter=all

// CHECK-LABEL: func @_tfrt_fallback_init
// CHECK-SAME: {{.*}} !tfrt.chain
// CHECK: tfrt_fallback_async.createop(%arg0) key(0) device("/device:CPU:0") "tf.ParseExampleV2"()

// RUN: tf-tfrt-opt -tfrt-lower-cluster-to-runtime-ops-non-tpu -split-input-file -verify-diagnostics %s | FileCheck %s

module attributes {tf.versions = {producer = 888 : i32}, 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:GPU:0"]} {

  // CHECK-LABEL: @converts_cluster
  func.func @converts_cluster() {
    // CHECK: "tf.XlaLaunch"()

    cc_api_version = 2,
    visibility = [
        "//learning/serving/contrib/tfrt/mlir/saved_model_analysis:__pkg__",
    ],
)

cc_library(
    name = "cluster_tf",
    srcs = ["cluster_tf.cc"],
    hdrs = ["cluster_tf.h"],
    visibility = [
        "//learning/serving/contrib/tfrt/mlir/saved_model_analysis:__pkg__",
        "//tensorflow/compiler/mlir/tensorflow/transforms:__pkg__",

    features = if_oss(["--path=org_tensorflow/tensorflow/compiler/mlir/tfrt"]),
    size_override = {
        "fallback.mlir": "medium",
    },
    test_file_exts = ["mlir"],
)

# Bundle together all of the test utilities that are used by tests.
filegroup(
    name = "test_utilities",
    testonly = True,
    data = [
        "//tensorflow/compiler/mlir/tfrt:tf-tfrt-opt",
        "@llvm-project//llvm:FileCheck",

// RUN: tf-tfrt-opt -split-input-file -tfrt-lower-tf-savedmodel="hoist-invariant-ops=true fuse-get-resource-ops=false" %s | FileCheck %s --dump-input=fail --dump-input-filter=all

module attributes {tf_saved_model.semantics} {

// Test hoisting hash table op.

// CHECK-LABEL: func @_tfrt_resource_init
// CHECK: [[handle:%.*]] = "tf.HashTableV2"()
// CHECK-SAME: shared_name = "x"

func.func @test(%ch: !tfrt.chain, %arg0: !corert.tensorhandle, %arg1_th: !corert.tensorhandle) {
  %cpu = corert.get_op_handler %ch "cpu"
  %0 = corert.executeop(%cpu) "tf.Relu"(%arg0) { T = f32 } : 1
  %arg1 = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor %arg1_th {_tfrt_cost = 1 : i64, device = "/CPU:0"} : (!corert.tensorhandle) -> (!tfrt_fallback.tf_tensor)
  %1 = tfrt_fallback_async.executeop key(0) cost(100) device("/CPU:0") "tf.Relu"(%arg1) { T = f32 } : 1

// RUN: tf-tfrt-opt -tf-to-tfrt %s | FileCheck %s --dump-input=fail

// CHECK-LABEL: func @device_test
func.func @device_test(
    %arg0: tensor<3x1xf32> {tf_saved_model.index_path = [0]},
    %arg1: tensor<1x3xf32> {tf_saved_model.index_path = [0]})
      -> (tensor<3x3xf32> {tf_saved_model.index_path = []}) {
  // CHECK: {{%.*}} = corert.get_op_handler %arg0 "/device:GPU:0"

#include "mlir/IR/OpDefinition.h"  // from @llvm-project
#include "mlir/IR/OpImplementation.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.h"

#define GET_OP_CLASSES
#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_ops.h.inc"

// RUN: not tf-tfrt-opt -tfrt-deduplicate-functions-invoked-by-batch-function %s 2>&1 | FileCheck %s

// This test verifies the error when two functions are different but invoked by
// the batch functions with same shared_name.

func.func private @batch_0(%arg0: tensor<?x?xi32>) -> tensor<*xi32> {