package(
    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
    default_visibility = [
        "//tensorflow/compiler/mlir/tensorflow:__pkg__",
        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
    ],
    licenses = ["notice"],
)

td_library(
    name = "tensorflow_tfrt_ops_td_files",
    srcs = [
        "tfrt_ops.td",
    ],
    compatible_with = get_compatible_with_portable(),

  TF_SessionOptions session_options;
  // true if async execution is enabled.
  bool async = false;
  TFE_ContextDevicePlacementPolicy device_placement_policy{
      TFE_DEVICE_PLACEMENT_SILENT};
  // If true, use TFRT backend
  bool use_tfrt = false;
  // Whether to run elementary eager ops wrapped in a call op.
  bool run_eager_op_as_function = false;
  // Whether to rewrite jit_compile functions.
  bool jit_compile_rewrite = false;
};

// RUN: tf-tfrt-opt -split-input-file -tf-mlrt-assign-op-key %s | FileCheck %s

// CHECK-LABEL: func @main
// CHECK: tf.AddV2
// CHECK-SAME: {__op_key = 0 : i32}

// CHECK: tf.AddV2
// CHECK-SAME: {__op_key = 1 : i32}

// CHECK: tf.AddV2
// CHECK-SAME: {__op_key = 2 : i32}

// CHECK: tf.AddV2
// CHECK-SAME: {__op_key = 3 : i32}

// CHECK: tf.Sub
// CHECK-SAME: {__op_key = 4 : i32}

// CHECK: tf.Sub

// are loaded from, since they retain pointers to the TensorHandles owned by the
// SavedModel, and the FunctionDef of the SavedModel.
//
// Note(bmzhao): This class is only TEMPORARILY virtual, as a way to unblock
// TFRT integration with TF Serving. Do not add more virtual implementations of
// this class. Eventually we want to remove this virtual base class indirection
// and have only a single implementation.
class ConcreteFunction {
 public:

// RUN: tf-tfrt-opt %s -tf-restore-merging | FileCheck %s

// CHECK-LABEL: func @single_restore_group
// CHECK-SAME:    (%[[ARG0:.*]]: {{.*}})
func.func @single_restore_group(%arg0: tensor<!tf_type.string>) -> (tensor<*xf32>, tensor<*xi32>) {
  %0 = "tf.Const"() {value = dense<"foo"> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
  %1 = "tf.Const"() {value = dense<""> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>

==============================================================================*/
#ifdef TF_MLRT_DIALECT
#else
#define TF_MLRT_DIALECT

include "mlir/Interfaces/SideEffectInterfaces.td"
include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.td"

// TODO(chky,rohitju): Unify this dialect with tfrt_fallback_sync dialect after
// vrooml is using the new interpreter.
def TensorflowMlrt_Dialect : Dialect {
  let name = "tf_mlrt";

    default_visibility = [
        "//learning/serving/contrib/tfrt/mlir/saved_model_analysis:__pkg__",
        "//tensorflow/compiler/mlir:__pkg__",
        "//tensorflow/compiler/mlir/tensorflow/transforms:__pkg__",
        "//tensorflow/compiler/mlir/tf2xla/api:__subpackages__",
        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
        "//tensorflow/compiler/tf2xla:__pkg__",

See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/mlir/tfrt/saved_model/saved_model.h"

#include <utility>

#include "llvm/ADT/STLFunctionalExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"

#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/kernels/ops_testutil.h"
#include "tensorflow/core/platform/refcount.h"
#include "tensorflow/core/tfrt/common/create_pjrt_client_util.h"
#include "tensorflow/core/tfrt/common/pjrt_util.h"
#include "tsl/framework/allocator.h"
#include "tsl/lib/core/status_test_util.h"
#include "tsl/platform/status.h"
#include "tsl/platform/statusor.h"

// 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