chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/tfrt/runtime_fallback/runtime_fallback_combine.cc

@@ -0,0 +1,80 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+//===----------------------------------------------------------------------===//
+//
+// This file implements a set of simple combiners for optimizing operations in
+// the Runtime Fallback dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/Matchers.h"
+#include "mlir/IR/PatternMatch.h"
+#include "tensorflow/compiler/mlir/tfrt/runtime_fallback/runtime_fallback_ops.h"
+
+// This optimizes the following scenario:
+// %tft0, %c2 = "tfd.move_dht_to_tft"(%dht0, %c1)
+//     : (!dht.host_tensor, !tfrt.chain) -> (!tfd.tf_tensor, !tfrt.chain)
+// %dht1, %c3 = "tfd.convert_tft_to_dht"(%tft0, %c2)
+//     : (!tfd.tf_tensor, !tfrt.chain) -> (!dht.host_tensor, !tfrt.chain)
+// some_op %dht1, %c3
+//
+// becomes
+// some_op %dht0, %c1
+
+struct SimplifyDoubleConversion
+    : public mlir::OpRewritePattern<mlir::tfd::ConvertTftToDhtOp> {
+  // We register this pattern to match every tfd.move_dht_to_tft op.
+  // The "benefit" is used by the framework to order the patterns and process
+  // them in order of profitability.
+  explicit SimplifyDoubleConversion(mlir::MLIRContext* context)
+      : mlir::OpRewritePattern<mlir::tfd::ConvertTftToDhtOp>(context,
+                                                             /*benefit=*/1) {}
+
+  // This method attempts to match a pattern and rewrite it. The rewriter
+  // argument is the orchestrator of the sequence of rewrites. The pattern is
+  // expected to interact with it to perform any changes to the IR from here.
+  mlir::LogicalResult matchAndRewrite(
+      mlir::tfd::ConvertTftToDhtOp op,
+      mlir::PatternRewriter& rewriter) const override {
+    // Look through the inputs of the ConvertTftToDhtOp.
+    mlir::Value convert_op_input_0 = op.getOperand(0);
+    mlir::Value convert_op_input_1 = op.getOperand(1);
+    mlir::tfd::MoveDhtToTftOp move_input_op_0 =
+        llvm::dyn_cast_or_null<mlir::tfd::MoveDhtToTftOp>(
+            convert_op_input_0.getDefiningOp());
+    mlir::tfd::MoveDhtToTftOp move_input_op_1 =
+        llvm::dyn_cast_or_null<mlir::tfd::MoveDhtToTftOp>(
+            convert_op_input_1.getDefiningOp());
+
+    // The inputs should be MoveDhtToTftOp.
+    if (!move_input_op_0 || !move_input_op_1) return mlir::failure();
+    // Both inputs are the same MoveDhtToTftOp.
+    if (move_input_op_0 != move_input_op_1) return mlir::failure();
+
+    // Use the rewriter to replace the ConvertTftToDhtOp's users with the
+    // operands of MoveDhtToTftOp.
+    rewriter.replaceOp(
+        op, {move_input_op_0.getOperand(0), move_input_op_0.getOperand(1)});
+    return mlir::success();
+  }
+};
+
+// Register rewrite pattern as "canonicalization" patterns on the MoveDhtToTftOp
+// so that they can be picked up by the Canonicalization framework.
+void mlir::tfd::ConvertTftToDhtOp::getCanonicalizationPatterns(
+    RewritePatternSet& results, MLIRContext* context) {
+  results.add<SimplifyDoubleConversion>(context);
+}

tensorflow/compiler/mlir/tfrt/runtime_fallback/runtime_fallback_executor.cc

@@ -0,0 +1,230 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/tfrt/runtime_fallback/runtime_fallback_executor.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/SourceMgr.h"
+#include "mlir/Parser/Parser.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tfrt/utils/host_context.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/platform/threadpool.h"
+#include "tensorflow/core/platform/threadpool_interface.h"
+#include "tensorflow/core/runtime_fallback/kernel/kernel_fallback_execute_compat_eager.h"
+#include "tensorflow/core/runtime_fallback/runtime/kernel_utils.h"
+#include "tensorflow/core/tfrt/utils/fallback_tensor.h"
+#include "tfrt/bef/bef_buffer.h"  // from @tf_runtime
+#include "tfrt/bef_converter/mlir_to_bef.h"  // from @tf_runtime
+#include "tfrt/bef_executor/bef_file.h"  // from @tf_runtime
+#include "tfrt/host_context/async_value.h"  // from @tf_runtime
+#include "tfrt/host_context/chain.h"  // from @tf_runtime
+#include "tfrt/host_context/execution_context.h"  // from @tf_runtime
+#include "tfrt/host_context/function.h"  // from @tf_runtime
+#include "tfrt/host_context/host_context.h"  // from @tf_runtime
+#include "tfrt/host_context/resource_context.h"  // from @tf_runtime
+#include "tfrt/support/ref_count.h"  // from @tf_runtime
+
+namespace tensorflow {
+
+using ::tfrt::AsyncValue;
+using ::tfrt::BEFFile;
+using ::tfrt::ExecutionContext;
+using ::tfrt::Function;
+using ::tfrt::MakeAvailableAsyncValueRef;
+using ::tfrt::RCReference;
+using ::tfrt::RequestContextBuilder;
+
+using ::tensorflow::Env;
+using ::tensorflow::thread::ThreadPool;
+using ::tensorflow::thread::ThreadPoolInterface;
+
+using ::tensorflow::tfrt_stub::FallbackTensor;
+
+// -------------------------------------------------------------------------- //
+// Run function via the TF->TFRT fallback lowering.
+// -------------------------------------------------------------------------- //
+
+namespace {
+// Thread pool for running `intra-op` tasks scheduled by the fallback kernels.
+class IntraOpThreadPool : public ThreadPoolInterface {
+ public:
+  explicit IntraOpThreadPool(int64_t num_threads)
+      : tpool_(Env::Default(), "intra-op",
+               std::max(1, static_cast<int32_t>(num_threads))) {}
+
+  void Schedule(std::function<void()> fn) override {
+    tpool_.Schedule(std::move(fn));
+  }
+
+  int NumThreads() const override { return tpool_.NumThreads(); }
+  int CurrentThreadId() const override { return tpool_.CurrentThreadId(); }
+  void Cancel() override {}
+
+ private:
+  ThreadPool tpool_;
+};
+}  // namespace
+
+RuntimeFallbackExecutor::RuntimeFallbackExecutor(int64_t num_threads)
+    : intra_op_(std::make_unique<IntraOpThreadPool>(num_threads)) {
+  // Create a HostContext for running TFRT functions. Concurrent work queue acts
+  // similar to the Tensorflow `inter-op` thread pool, so we'll match the size.
+  host_context_ = num_threads ? CreateMultiThreadedHostContext(num_threads)
+                              : CreateSingleThreadedHostContext();
+  tfrt::RegisterStaticKernels(host_context_->GetMutableRegistry());
+
+  // Build an ExecutionContext from the HostContext.
+  auto builder = RequestContextBuilder(host_context_.get(), &resource_context_);
+
+  // Get tensorflow::EagerContext for the kernel fallback.
+  auto* eager_context_resource =
+      resource_context_
+          .GetOrCreateResource<tensorflow::tfd::EagerContextResource>(
+              tensorflow::tfd::kEagerContextResourceName);
+  auto expected_eager_context = eager_context_resource->GetTFEagerContext();
+  auto* eager_context = expected_eager_context.get();
+
+  // Initialize fallback kernels state with a custom intra-op thread pool.
+  auto status = tensorflow::tfd::SetUpKernelFallbackCompatRequestContext(
+      &builder, /*runner_table=*/nullptr, eager_context, intra_op_.get());
+  CHECK(status.ok()) << "Failed to setup request context: " << status.message();
+
+  auto req_ctx = std::move(builder).build();
+  if (auto err = req_ctx.takeError())
+    LOG(FATAL) << "Failed to build a request context";
+
+  exec_ctx_ = std::make_unique<tfrt::ExecutionContext>(std::move(*req_ctx));
+}
+
+void RuntimeFallbackExecutor::Prepare(llvm::StringRef mlir_input) {
+  // We only support IR written in the Tensorflow dialect.
+  mlir::DialectRegistry registry;
+  mlir::RegisterAllTensorFlowDialects(registry);
+  mlir::MLIRContext context(registry);
+
+  llvm::SourceMgr source_mgr;
+  source_mgr.AddNewSourceBuffer(
+      llvm::MemoryBuffer::getMemBuffer(mlir_input, "test_ir"), llvm::SMLoc());
+
+  // Parse a kernel source code into the MLIR Module.
+  mlir::OwningOpRef<mlir::ModuleOp> module(
+      mlir::parseSourceFile<mlir::ModuleOp>(source_mgr, &context));
+  CHECK(module) << "failed to parse mlir module";
+
+  // Collect all diagnostics emitted while lowering parsed kernel module.
+  std::string diagnostic_str;
+  llvm::raw_string_ostream os(diagnostic_str);
+  mlir::SourceMgrDiagnosticHandler handler(source_mgr, module->getContext(),
+                                           os);
+
+  // Convert TF to TFRT fallback dialect.
+  TfrtPipelineOptions pipeline_opts;
+  pipeline_opts.default_device = kDefaultHostDeviceName;
+  pipeline_opts.hoist_invariant_ops = true;
+  pipeline_opts.sink_in_invariant_ops = false;
+  pipeline_opts.cost_threshold = 1024;
+  pipeline_opts.merge_inter_dependent_streams = true;
+
+  mlir::PassManager pm(module->getContext());
+  pm.addPass(CreateTfToTfrtConversionPass(pipeline_opts));
+
+  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(),
+                    host_context_->diag_handler(), host_context_->allocator());
+  CHECK(bef_file_) << "Failed to open BEF";
+
+  // Run TFRT initialization function to pre-instantiate fallback kernels.
+  RunTfrtInitializer();
+}
+
+llvm::SmallVector<Tensor> RuntimeFallbackExecutor::Execute(
+    llvm::StringRef function_name, llvm::ArrayRef<Tensor> arguments) {
+  // Get the kernel entrypoint function.
+  const Function* compute = bef_file_->GetFunction(function_name);
+  CHECK(compute) << "Entrypoint function not found";
+  CHECK_EQ(arguments.size() + 1, compute->num_arguments())
+      << "Wrong number of arguments for function " << function_name.str();
+
+  // Prepare function arguments from ready Chain and input Tensors.
+  llvm::SmallVector<tfrt::AsyncValue*> exec_arguments;
+  exec_arguments.reserve(compute->num_arguments());
+  exec_arguments.push_back(tfrt::GetReadyChain().release());
+  for (const Tensor& input_tensor : arguments) {
+    auto av = MakeAvailableAsyncValueRef<FallbackTensor>(input_tensor);
+    exec_arguments.push_back(av.release());
+  }
+
+  // Space for returned values.
+  llvm::SmallVector<RCReference<AsyncValue>> results(compute->num_results());
+
+  compute->Execute(*exec_ctx_, exec_arguments, results);
+
+  // Wait for the function execution to finish, as well as the side-effects.
+  host_context_->Await(results);
+
+  // Check that all results are available.
+  llvm::SmallVector<Tensor> ret_values;
+  for (unsigned i = 1; i < results.size(); ++i) {
+    if (auto* error = results[i]->GetErrorIfPresent())
+      LOG(FATAL) << "Failed to execute a function: " << error->message();
+    ret_values.push_back(results[i]->get<tfrt_stub::FallbackTensor>().tensor());
+  }
+
+  // Deallocate arguments.
+  for (auto* argument : exec_arguments) argument->DropRef();
+  return ret_values;
+}
+
+// Run TFRT fallback initialization function to instantiate all fallback
+// kernels ahead of executing the compute function.
+void RuntimeFallbackExecutor::RunTfrtInitializer() {
+  const Function* func = bef_file_->GetFunction("_tfrt_fallback_init");
+  CHECK(func) << "TFRT initialization function was not found";
+  CHECK_EQ(func->argument_types().size(), 1);
+
+  llvm::SmallVector<RCReference<AsyncValue>, 1> results;
+  results.resize(func->result_types().size());
+  CHECK_EQ(results.size(), 1);
+
+  func->Execute(*exec_ctx_, tfrt::GetReadyChain().GetAsyncValue(), results);
+
+  host_context_->Await(results);
+
+  CHECK(!results[0]->IsError()) << "Failed to run TFRT initialization function";
+}
+
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/runtime_fallback/runtime_fallback_executor.h

@@ -0,0 +1,63 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_RUNTIME_FALLBACK_RUNTIME_FALLBACK_EXECUTOR_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_RUNTIME_FALLBACK_RUNTIME_FALLBACK_EXECUTOR_H_
+
+#include <cstdint>
+#include <memory>
+
+#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
+#include "tfrt/host_context/resource_context.h"  // from @tf_runtime
+#include "tfrt/support/ref_count.h"  // from @tf_runtime
+
+namespace tensorflow {
+
+class RuntimeFallbackExecutor {
+ public:
+  explicit RuntimeFallbackExecutor(int64_t num_threads);
+
+  // Prepare() needs to be called once before calling Execute(). It sets up all
+  // things necessary to execute the given 'mlir_input' with the fallback to
+  // tensorflow.
+  void Prepare(llvm::StringRef mlir_input);
+
+  // Execute() can be called several times after the call to Prepare() (e.g. for
+  // benchmarking).
+  llvm::SmallVector<Tensor> Execute(llvm::StringRef function_name,
+                                    llvm::ArrayRef<Tensor> arguments);
+
+ private:
+  void RunTfrtInitializer();
+
+  std::unique_ptr<thread::ThreadPoolInterface> intra_op_;
+  std::unique_ptr<tfrt::HostContext> host_context_;
+  tfrt::ResourceContext resource_context_;
+  std::unique_ptr<tfrt::ExecutionContext> exec_ctx_;
+  tfrt::BefBuffer bef_buffer_;
+  tfrt::RCReference<tfrt::BEFFile> bef_file_;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_RUNTIME_FALLBACK_RUNTIME_FALLBACK_EXECUTOR_H_

tensorflow/compiler/mlir/tfrt/runtime_fallback/runtime_fallback_ops.cc

@@ -0,0 +1,47 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/tfrt/runtime_fallback/runtime_fallback_ops.h"
+
+#include "tfrt/basic_kernels/opdefs/types.h"  // from @tf_runtime
+
+namespace mlir {
+namespace tfd {
+
+//===----------------------------------------------------------------------===//
+// TfrtDelegate Dialect
+//===----------------------------------------------------------------------===//
+
+RuntimeFallbackDialect::RuntimeFallbackDialect(MLIRContext *context)
+    : Dialect(/*name=*/"tfd", context, TypeID::get<RuntimeFallbackDialect>()) {
+  allowUnknownTypes();
+
+  allowUnknownOperations();
+
+  addOperations<
+#define GET_OP_LIST
+#include "tensorflow/compiler/mlir/tfrt/runtime_fallback_ops.cc.inc"
+      >();
+}
+
+}  // namespace tfd
+}  // namespace mlir
+
+//===----------------------------------------------------------------------===//
+// TableGen'd op method definitions
+//===----------------------------------------------------------------------===//
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/runtime_fallback_ops.cc.inc"

tensorflow/compiler/mlir/tfrt/runtime_fallback/runtime_fallback_ops.h

@@ -0,0 +1,45 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This file defines the operations used in the Runtime Fallback dialect.
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_RUNTIME_FALLBACK_RUNTIME_FALLBACK_OPS_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_RUNTIME_FALLBACK_RUNTIME_FALLBACK_OPS_H_
+
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Dialect.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
+#include "tfrt/tensor/opdefs/tensor.h"  // from @tf_runtime
+
+namespace mlir {
+namespace tfd {
+
+// Dialect for TFRT delegate operations.
+class RuntimeFallbackDialect : public Dialect {
+ public:
+  explicit RuntimeFallbackDialect(MLIRContext* context);
+  static StringRef getDialectNamespace() { return "tfd"; }
+};
+
+}  // namespace tfd
+}  // namespace mlir
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/runtime_fallback_ops.h.inc"
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_RUNTIME_FALLBACK_RUNTIME_FALLBACK_OPS_H_

tensorflow/compiler/mlir/tfrt/runtime_fallback/runtime_fallback_ops.td

@@ -0,0 +1,158 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This is the definition file for the Runtime Fallback Dialect.
+
+#ifdef TFRT_DELEGATE_DIALECT
+#else
+#define TFRT_DELEGATE_DIALECT
+
+include "tfrt/tfrt_op_base.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+
+//===----------------------------------------------------------------------===//
+// Type definitions
+//===----------------------------------------------------------------------===//
+def TfTensorType : OpaqueType<"tfd", "tf_tensor", "!tfd.tf_tensor type">;
+
+//===----------------------------------------------------------------------===//
+// Runtime Fallback Dialect definitions
+//===----------------------------------------------------------------------===//
+
+def RuntimeFallback_Dialect : Dialect {
+  let name = "tfd";
+
+  let description = [{
+    The Runtime Fallback dialect.
+
+    This dialect contains operations to run existing TF kernels on TFRT by
+    invoking TF Eager API.
+  }];
+
+  let cppNamespace = "::mlir::tfd";
+}
+
+//===----------------------------------------------------------------------===//
+// Runtime Fallback Dialect Ops definitions
+//===----------------------------------------------------------------------===//
+
+// Base class for the operation in this dialect.
+class RuntimeFallbackDialect_Op<string mnemonic, list<Trait> traits = []> :
+    Op<RuntimeFallback_Dialect, mnemonic, traits> { }
+
+def InitEagerContextOp : RuntimeFallbackDialect_Op<"init_eager_context"> {
+  let summary = "eager context initialization operation";
+  let description = [{
+    The "tfd.init_eager_context" operation takes an input chain, creates and
+    initializes the TF EagerContext and returns an output chain.
+
+    Example:
+      %c1 = "tfd.init_eager_context"(%c0): (!tfrt.chain) -> !tfrt.chain
+  }];
+
+  let arguments = (ins TFRT_ChainType);
+  let results = (outs TFRT_ChainType);
+}
+
+def DelegateKernelOp : RuntimeFallbackDialect_Op<"delegate_kernel"> {
+  let summary = "delegate kernel operation";
+  let description = [{
+    The "tfd.delegate_kernel" operation takes an input chain, and arbitrary
+    number of input arguments, and runs a specified TF op via TFE C API. It
+    returns an output chain and variable number of outputs from the TF op.
+
+    The input arguments and attributes are passed to the TF op. The outputs are
+    outputs of the TF op.
+
+    Note that `_name` is a required attribute specifying the TF op to run.
+    TFRT attributes are sorted alphabetically, passed in as positional
+    attributes to the TFRT kernel, rather than as named attributes.
+
+    Example:
+      To run "tf.MatMul" op, which has two boolean attributes,
+        1. Set _name = "MatMul"
+        2. For each TF attribute, split it into two attributes, one for name of
+           the TF attribute, and the other for the type and value of the
+           attribute value. Attribute value is a string with the format of
+           "type$val", where type can be "bool", "string", "tfdtype", "tfshape",
+           "tftensor".
+           The value serialization format can be found in attr_util.h.
+
+      %out_c, %out_tensor = "tfd.delegate_kernel"(
+        %in_c, %in1_tensor, %in2_tensor) {
+        _name = "MatMul",
+        attr1_name = "transpose_a", attr1_value = "bool$false",
+        attr2_name = "transpose_b", attr2_value = "bool$false"
+      } : (!tfrt.chain, !tfd.tf_tensor, !tfd.tf_tensor) -> (
+        !tfrt.chain, !tfd.tf_tensor)
+  }];
+
+  let arguments = (ins TFRT_ChainType, Variadic<AnyType>);
+  let results = (outs TFRT_ChainType, Variadic<AnyType>);
+}
+
+def PrintTftOp : RuntimeFallbackDialect_Op<"print_tft"> {
+  let summary = "print TF tensor operation";
+  let description = [{
+    The "tfd.print_tft" operation prints the input TF tensor. It takes an input
+    TF tensor to be printed and an input chain, and returns an output chain.
+
+    Example:
+      %c1 = "tfd.print_tft"(%t, %c) : (!tfd.tf_tensor, !tfrt.chain) -> !tfrt.chain
+
+  }];
+
+  let arguments = (ins TfTensorType, TFRT_ChainType);
+  let results = (outs TFRT_ChainType);
+}
+
+def ConvertTftToDhtOp : RuntimeFallbackDialect_Op<"convert_tft_to_dht", [Pure]> {
+  let summary = "convert TF tensor to TFRT DHT tensor operation";
+  let description = [{
+    The "tfd.convert_tft_to_dht" operation converts a TF tensor to a TFRT
+    DenseHostTensor.
+
+    It takes as input a TF Tensor and an input chain, and returns a converted
+    TFRT DHT tensor and an output chain.
+
+    Example:
+      %dht, %c0 = "tfd.convert_tft_to_dht"(%tft, %c)
+      : (!tfd.tf_tensor, !tfrt.chain) -> (!dht.host_tensor, !tfrt.chain)
+  }];
+
+  let arguments = (ins TfTensorType, TFRT_ChainType);
+  // Enable registering canonicalization patterns with this operation.
+  let hasCanonicalizer = 1;
+  let results = (outs TensorType, TFRT_ChainType);
+}
+
+def MoveDhtToTftOp : RuntimeFallbackDialect_Op<"move_dht_to_tft", [Pure]> {
+  let summary = "convert TFRT DHT tensor to DHT tensor operation";
+  let description = [{
+    The "tfd.move_dht_to_tft" operation moves a TFRT tensor into a TF Tensor.
+
+    It takes as input a TFRT Tensor and an input chain, and returns a TF tensor
+    with the same underlying buffer and an output chain.
+
+    Example:
+      %dht, %c0 = "tfd.convert_tft_to_dht"(%tft, %c)
+        : (!tfd.tf_tensor, !tfrt.chain) -> (!dht.host_tensor, !tfrt.chain)
+  }];
+
+  let arguments = (ins TensorType, TFRT_ChainType);
+  let results = (outs TfTensorType, TFRT_ChainType);
+}
+
+#endif // TFRT_DELEGATE_DIALECT