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

// CHECK-LABEL: func @assign_variable
// CHECK-SAME: ([[in_chain:%.*]]: !tfrt.chain) -> !tfrt.chain
func.func @assign_variable() {
  // CHECK: [[ch1:%.*]], %results = tfrt_fallback_async.executeop.seq([[in_chain]]) key(0) cost({{.*}}) device("/device:CPU:0") "tf.VarHandleOp"

  %s = "tfrt_test.get_string"() { value = "Running @matmul_cpu" } : () -> !tfrt.string
  %ch4 = "tfrt_test.print_string"(%s, %ch3) : (!tfrt.string, !tfrt.chain) -> (!tfrt.chain)

  %ch5 = "tfrt_fallback_async.print_tensor"(%result, %ch4) : (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain

  tfrt.return %ch5, %result : !tfrt.chain, !tfrt_fallback.tf_tensor
}

#include "mlir/Interfaces/InferTypeOpInterface.h"  // from @llvm-project
#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
#include "tfrt/core_runtime/opdefs/traits.h"  // from @tf_runtime
#include "tfrt/tensor/opdefs/tensor.h"  // from @tf_runtime

using namespace mlir;  // NOLINT

namespace tfrt {
namespace fallback_sync {

// Dialect for fallback operations.
class FallbackSyncDialect : public Dialect {
 public:

#include "tfrt/init_tfrt_dialects.h"  // from @tf_runtime

namespace tfrt {
namespace fallback_async {
namespace {

TEST(SavedModelTest, MapFallbackArgs) {
  std::string saved_model_mlir_path = tensorflow::GetDataDependencyFilepath(
      "tensorflow/compiler/mlir/tfrt/tests/ir/testdata/test.mlir");

  mlir::DialectRegistry registry;
  RegisterTFRTDialects(registry);

  using Base::Base;
  static constexpr StringLiteral name = "tfrt.tf_tensor";
};

// The MLIR type represents a tensorflow::Allocator.
class TFAllocatorType
    : public Type::TypeBase<TFAllocatorType, Type, TypeStorage> {
 public:
  using Base::Base;
  static constexpr StringLiteral name = "tfrt.tf_allocator";
};

}  // namespace fallback
}  // namespace tfrt

#define GET_OP_CLASSES

        "//tensorflow/core/runtime_fallback/runtime:runtime_fallback_alwayslink",
        "//tensorflow/core/runtime_fallback/util:fallback_test_util",
        "//tensorflow/core/tfrt/fallback:op_kernel_runner",
        "//tensorflow/core/tfrt/runtime",
        "//tensorflow/core/tfrt/utils:thread_pool",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/strings",
        "@com_google_googletest//:gtest_main",

#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
#include "tensorflow/compiler/mlir/tfrt/transforms/passes.h"
#include "tensorflow/compiler/mlir/tfrt/transforms/tfrt_pipeline_options.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tfrt/bef/bef_buffer.h"  // from @tf_runtime
#include "tfrt/bef_converter/mlir_to_bef.h"  // from @tf_runtime

namespace tensorflow {

#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
#include "tfrt/compiler/opdefs/tfrt_op_interfaces.h"  // from @tf_runtime
#include "tfrt/compiler/opdefs/tfrt_traits.h"  // from @tf_runtime
#include "tfrt/core_runtime/opdefs/traits.h"  // from @tf_runtime

using namespace mlir;  // NOLINT

namespace tfrt {
namespace fallback_async {

// Dialect for fallback async operations.

// RUN: tf-tfrt-opt -tfrt-remove-device-attribute %s | FileCheck %s

func.func @test(%arg0: !tfrt.chain, %arg1: !corert.tensorhandle) -> (!tfrt.chain, !corert.tensorhandle) {
  %0 = corert.get_op_handler %arg0 "cpu"
  // CHECK: %[[RESULT:.*]] = corert.executeop(%[[ARG_0:.*]]) "tf.MatMul"(%[[ARG_1:.*]], %[[ARG_1]]) {T = f32, transpose_a = false, transpose_b = false} : 1

// RUN: tf-tfrt-opt -tf-executor-to-tfrt-pipeline="enable-optimizer=true tfrt-cost-threshold=1024" %s | FileCheck %s --dump-input=fail

// Check that unused While op results and the associated ops are removed.

module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 462 : i32}} {
  func.func @while_cond_lt9(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i1> {