#     ],
# )
# copybara:uncomment_end

tf_cc_shared_test(
    name = "kernel_fallback_compat_test",
    srcs = ["kernel_fallback_compat_test.cc"],
    data = [
        "create_op.mlir.bef",
        "custom_thread_pool.mlir.bef",
    ],
    tags = ["no_oss"],
    deps = [
        "//tensorflow/compiler/mlir/tfrt/ir:tfrt_fallback_async_opdefs",
        "//tensorflow/core:all_kernels",
        "//tensorflow/core:lib",

  host->Await(results);

  for (auto& result : results) {
    EXPECT_FALSE(result->IsError());
  }
}

// Runs a BEF function that batches a function that does nothing just to measure
// the runtime overhead. The BEF function signature is adapted from a real model
// and is useful for benchmarking ops with large attributes and many
// input/output.

  CHECK(mlir::succeeded(pm.run(*module)))
      << "Failed to lower module to TFRT: " << os.str();

  // Convert module to BEF.
  bef_buffer_ =
      tfrt::ConvertMLIRToBEF(*module, /*disable_optional_sections=*/false);
  CHECK(!bef_buffer_.empty()) << "Failed to convert module to BEF";

  bef_file_ =
      BEFFile::Open(bef_buffer_, host_context_->GetKernelRegistry(),

#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 {

Status CompileTFMLIRToBEF(const TfrtFunctionCompileOptions& options,

#include "absl/strings/string_view.h"
#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tfrt/translate/tfrt_compile_options.h"
#include "tensorflow/core/platform/status.h"
#include "tfrt/bef/bef_buffer.h"  // from @tf_runtime
#include "tfrt/core_runtime/tensor_handle.h"  // from @tf_runtime

namespace tfrt {
class CoreRuntime;
}

namespace mlir {
class ModuleOp;
}

namespace tensorflow {

  }];

  let arguments = (ins
    TFRT_ChainType:$in_ch,
    I64Attr:$num_args,
    StrAttr:$device,
    ArrayAttr:$op_attrs,
    // TODO(b/173025975): consider using DictionaryAttr after we support
    // BEF conversion for this type.
    ArrayAttr:$op_func_attrs,
    I64Attr:$op_key,
    StrAttr:$op_name
  );

  let results = (outs
    TFRT_ChainType:$out_ch
  );

  let hasVerifier = 1;

// RUN: tfrt_fallback_translate -mlir-to-bef %s | tf_bef_executor --work_queue_type=mstd 2>&1 | FileCheck %s

func.func @test_async_op_kernel_thread() -> !tfrt.chain {
  %ch0 = tfrt.new.chain

  %ch1 = tfrt_fallback_async.createop(%ch0) key(0) device("/CPU:0") "tf.Const"()
         { dtype = i32, value = dense<[2]> : tensor<1xi32> } num_args(0)
  %ch2 = tfrt_fallback_async.createop(%ch1) key(1) device("/CPU:0")

        "@llvm-project//llvm:Support",
        "@llvm-project//mlir:Parser",
        "@llvm-project//mlir:Pass",
        "@tf_runtime//:basic_kernels_alwayslink",
        "@tf_runtime//:bef",
        "@tf_runtime//:befexecutor",
        "@tf_runtime//:core_runtime_alwayslink",
        "@tf_runtime//:hostcontext",
        "@tf_runtime//:mlirtobef",
        "@tf_runtime//:support",
    ],
)

// RUN: tfrt_fallback_translate --mlir-to-bef %s | tf_bef_executor --work_queue_type=mstd:1,1 | FileCheck %s
// RUN: tfrt_fallback_translate --mlir-to-bef %s | tf_bef_executor --work_queue_type=mstd:8 | FileCheck %s

func.func @matmul_cpu(%ch: !tfrt.chain, %a: !tfrt_fallback.tf_tensor, %b: !tfrt_fallback.tf_tensor) -> (!tfrt.chain, !tfrt_fallback.tf_tensor) {
  // Enqueue a sleep onto blocking work queue, %ch0 is fulfilled when sleeping is done.
  %us = tfrt.constant.i32 1000

#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "tensorflow/core/framework/tensor.h"
#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