// RUN: tf-tfrt-opt -optimize-tf-for-tfrt -split-input-file -verify-diagnostics %s | FileCheck %s

// CHECK-LABEL: @fold_device_index
func.func @fold_device_index() -> tensor<i32> {
  // CHECK-NOT: tf.DeviceIndex
  // CHECK: tf.Const
  // CHECK-SAME: value = dense<1> : tensor<i32>
  %0 = "tf.DeviceIndex"() {device = "/device:CPU:0", device_names = ["GPU", "CPU"]} : () -> tensor<i32>
  func.return %0 : tensor<i32>
}

// -----

syntax = "proto3";

package mlir.tfrt;

message CompatibilityAnalysisReportProto {
  bool unknown_dialect = 1;
  bool ref_variable = 2;
  bool incompatible_variable = 3;
  bool incompatible_attribute = 4;
  bool control_flow_v1 = 5;
  string method_name = 6;

  // TODO(chky): add more checks, eg. tensor datatypes.
}

message CompatibilityAnalysisProto {
  CompatibilityAnalysisReportProto summary = 1;

  message OpInfo {

  static mlir::PassPipelineRegistration<StandardPipelineOptions> pipeline(
      "tfrt-lower-cluster-to-runtime-ops-tpu",
      "Run all the passes involved after the clustering transformations from "
      "the TF2XLA Bridge. Takes as input a Module with tf_device.cluster ops "
      "and outputs TFRT runtime ops such as TPUCompile. This pipeline is for "
      "TPU.",

  std::string saved_model_dir(dirname);
  std::unique_ptr<tensorflow::SavedModelAPI> result;

  if (tensorflow::unwrap(ctx)->UsesTFRT()) {
    status->status = tensorflow::errors::Unimplemented(
        "TFRT SavedModel implementation will be added in the future");
  } else {
    std::unique_ptr<tensorflow::TFSavedModelAPI> saved_model;
    status->status = tensorflow::TFSavedModelAPI::Load(
        dirname, absl::nullopt,

// RUN: tf-opt %s -split-input-file -verify-diagnostics | FileCheck %s

// Tests for TensorFlow TFRT ops with custom verifiers.

//===--------------------------------------------------------------------===//
//  Test TF operations (tf.*)
//===--------------------------------------------------------------------===//

// CHECK-LABEL: func @testPwStreamResults
func.func @testPwStreamResults(%arg0: tensor<f32>, %arg1: tensor<f32>) {

// RUN: tf-tfrt-opt -tfrt-reorder-tf-assert %s | FileCheck %s

// CHECK-LABEL: @reorder_assert
func.func @reorder_assert(%key0: tensor<!tf_type.string>, %key1: tensor<!tf_type.string>) -> (tensor<i64>, tensor<i64>) {
  %error_message = "tf.Const"() {value = dense<"error"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
  %default = "tf.Const"() {value = dense<-1> : tensor<i64>} : () -> tensor<i64>

// RUN: tf-tfrt-opt -tfrt-convert-ref-variables -split-input-file -verify-diagnostics %s | FileCheck %s

// Test the basic cases where all uses of a ref variable can be converted.

// CHECK-LABEL: @init
func.func @init() {
  // CHECK-NOT: tf.VariableV2
  // CHECK-NOT: tf.Assign

  // CHECK: [[handle:%.*]] = "tf.VarHandleOp"
  // CHECK-SAME: shared_name = "x"
  // CHECK: "tf.AssignVariableOp"([[handle]], {{%.*}})

)

# 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",
        "@llvm-project//mlir:run_lit.sh",
    ],

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

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

namespace tensorflow {

BackendCompiler::~BackendCompiler() = default;

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

// CHECK-LABEL:   func.func @prune_unused_restore
func.func @prune_unused_restore() {
  %cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
  %cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>