chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/tfrt/ir/BUILD

@@ -0,0 +1,235 @@
+# BUILD targets for runtime fallback opdefs.
+
+load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+
+# Note: keep the following lines separate due to the way copybara works
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "tfrt_fallback_opdefs",
+    srcs = [
+        "tfrt_fallback.cc",
+    ],
+    hdrs = ["tfrt_fallback.h"],
+    compatible_with = get_compatible_with_portable(),  # copybara: comment
+    visibility = [
+        # copybara:uncomment "//learning/brain/experimental/tfrt:__subpackages__",
+        # copybara:uncomment "//learning/brain/tfrt/tpu/compiler:__subpackages__",
+        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
+        "//tensorflow/core/runtime_fallback:internal",
+        "//tensorflow/core/tfrt/saved_model:friends",
+        # copybara:uncomment "//third_party/tf_runtime_google:__subpackages__",
+    ],
+    deps = [
+        ":tfrt_fallback_opdefs_inc_gen",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:SideEffectInterfaces",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+cc_library(
+    name = "tfrt_fallback_async_opdefs",
+    srcs = [
+        "tfrt_fallback_async.cc",
+    ],
+    hdrs = ["tfrt_fallback_async.h"],
+    compatible_with = get_compatible_with_portable(),  # copybara: comment
+    visibility = [
+        # copybara:uncomment "//learning/brain/experimental/tfrt:__subpackages__",
+        # copybara:uncomment "//learning/brain/tfrt/tpu/compiler/mlir:__subpackages__",
+        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
+        "//tensorflow/core/tfrt/saved_model:friends",
+    ],
+    deps = [
+        ":tfrt_fallback_async_opdefs_inc_gen",
+        ":tfrt_fallback_common",
+        ":tfrt_fallback_opdefs",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:BytecodeOpInterface",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:InliningUtils",
+        "@llvm-project//mlir:SideEffectInterfaces",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+        "@llvm-project//mlir:Transforms",
+        "@tf_runtime//:basic_kernels_opdefs",
+        "@tf_runtime//:compiler_tfrt_op_interfaces",
+        "@tf_runtime//:compiler_tfrt_traits",
+        "@tf_runtime//:core_runtime_opdefs",
+        "@tf_runtime//:core_runtime_sync_opdefs",
+    ],
+)
+
+cc_library(
+    name = "tfrt_fallback_sync_opdefs",
+    srcs = [
+        "tfrt_fallback_sync.cc",
+    ],
+    hdrs = ["tfrt_fallback_sync.h"],
+    visibility = [
+        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
+        "//tensorflow/core/runtime_fallback:internal",
+    ],
+    deps = [
+        ":tfrt_fallback_common",
+        ":tfrt_fallback_opdefs",
+        ":tfrt_fallback_sync_opdefs_inc_gen",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:InferTypeOpInterface",
+        "@llvm-project//mlir:SideEffectInterfaces",
+        "@tf_runtime//:basic_kernels_opdefs",
+        "@tf_runtime//:core_runtime_opdefs",
+        "@tf_runtime//:tensor_opdefs",
+    ],
+)
+
+cc_library(
+    name = "tfrt_fallback_common",
+    srcs = ["tfrt_fallback_common.cc"],
+    hdrs = ["tfrt_fallback_common.h"],
+    compatible_with = get_compatible_with_portable(),  # copybara: comment
+    deps = [
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@tf_runtime//:basic_kernels_opdefs",
+    ],
+)
+
+cc_library(
+    name = "tfrt_fallback_util",
+    srcs = ["tfrt_fallback_util.cc"],
+    hdrs = ["tfrt_fallback_util.h"],
+    visibility = ["//tensorflow/compiler/mlir/tfrt:__subpackages__"],
+    deps = [
+        ":tfrt_fallback_async_opdefs",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+    ],
+)
+
+td_library(
+    name = "tfrt_fallback_td_files",
+    srcs = [
+        "tfrt_fallback.td",
+        "tfrt_fallback_async.td",
+        "tfrt_fallback_sync.td",
+    ],
+    compatible_with = get_compatible_with_portable(),  # copybara: comment
+    includes = ["."],
+    visibility = [
+        "//learning/brain/experimental/tfrt:__subpackages__",
+        "//learning/brain/tfrt/tpu/compiler:__subpackages__",
+        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
+    ],
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops_td_files",
+        "@llvm-project//mlir:SideEffectInterfacesTdFiles",
+        "@tf_runtime//:CoreRTTdFiles",
+        "@tf_runtime//:OpBaseTdFiles",
+        "@tf_runtime//:compiler_td_files",
+    ],
+)
+
+gentbl_cc_library(
+    name = "tfrt_fallback_opdefs_inc_gen",
+    compatible_with = get_compatible_with_portable(),  # copybara: comment
+    tbl_outs = {
+        "tfrt_fallback.h.inc": ["-gen-op-decls"],
+        "tfrt_fallback.cpp.inc": ["-gen-op-defs"],
+    },
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "tfrt_fallback.td",
+    deps = [":tfrt_fallback_td_files"],
+)
+
+gentbl_cc_library(
+    name = "tfrt_fallback_async_opdefs_inc_gen",
+    compatible_with = get_compatible_with_portable(),  # copybara: comment
+    tbl_outs = {
+        "tfrt_fallback_async.h.inc": ["-gen-op-decls"],
+        "tfrt_fallback_async.cpp.inc": ["-gen-op-defs"],
+    },
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "tfrt_fallback_async.td",
+    deps = [":tfrt_fallback_td_files"],
+)
+
+gentbl_cc_library(
+    name = "tfrt_fallback_sync_opdefs_inc_gen",
+    tbl_outs = {
+        "tfrt_fallback_sync.h.inc": ["-gen-op-decls"],
+        "tfrt_fallback_sync.cpp.inc": ["-gen-op-defs"],
+        "tfrt_fallback_sync_dialect.h.inc": [
+            "-gen-dialect-decls",
+            "-dialect=tfrt_fallback_sync",
+        ],
+    },
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "tfrt_fallback_sync.td",
+    test = True,
+    visibility = [
+        "//learning/brain/experimental/tfrt:__subpackages__",
+    ],
+    deps = [":tfrt_fallback_td_files"],
+)
+
+td_library(
+    name = "gpu_ops_td_file",
+    srcs = [
+        "gpu_ops.td",
+    ],
+    includes = ["."],
+    visibility = [
+        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
+    ],
+    deps = [
+        ":tfrt_fallback_td_files",
+        "@tf_runtime//:OpBaseTdFiles",
+    ],
+)
+
+gentbl_cc_library(
+    name = "tfrt_gpu_opdefs_inc_gen",
+    tbl_outs = {
+        "gpu_ops.h.inc": ["-gen-op-decls"],
+        "gpu_ops.cpp.inc": ["-gen-op-defs"],
+        "gpurt_dialect.h.inc": [
+            "-gen-dialect-decls",
+            "-dialect=gpurt",
+        ],
+    },
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "gpu_ops.td",
+    test = True,
+    visibility = [
+        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
+    ],
+    deps = [":gpu_ops_td_file"],
+)
+
+cc_library(
+    name = "tfrt_gpu_opdefs",
+    srcs = [
+        "gpu_ops.cc",
+    ],
+    hdrs = ["gpu_ops.h"],
+    visibility = [
+        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
+    ],
+    deps = [
+        ":tfrt_fallback_opdefs",
+        ":tfrt_gpu_opdefs_inc_gen",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:BytecodeOpInterface",
+        "@llvm-project//mlir:IR",
+    ],
+)

tensorflow/compiler/mlir/tfrt/ir/gpu_ops.cc

@@ -0,0 +1,41 @@
+/* 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/ir/gpu_ops.h"
+
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback.h"
+
+namespace tfrt {
+namespace gpu {
+
+GpuRuntimeDialect::GpuRuntimeDialect(MLIRContext *context)
+    : Dialect(/*name=*/"gpurt", context, TypeID::get<GpuRuntimeDialect>()) {
+  addOperations<
+#define GET_OP_LIST
+#include "tensorflow/compiler/mlir/tfrt/ir/gpu_ops.cpp.inc"
+      >();
+}
+
+}  // namespace gpu
+}  // namespace tfrt
+
+//===----------------------------------------------------------------------===//
+// TableGen'd op method definitions
+//===----------------------------------------------------------------------===//
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/gpu_ops.cpp.inc"

tensorflow/compiler/mlir/tfrt/ir/gpu_ops.h

@@ -0,0 +1,41 @@
+/* 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_IR_GPU_OPS_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_IR_GPU_OPS_H_
+
+#include "mlir/Bytecode/BytecodeOpInterface.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
+
+using namespace mlir;  // NOLINT
+
+namespace tfrt {
+namespace gpu {
+
+// Dialect for TFRT GPU operations.
+class GpuRuntimeDialect : public Dialect {
+ public:
+  explicit GpuRuntimeDialect(MLIRContext *context);
+  static StringRef getDialectNamespace() { return "gpurt"; }
+};
+
+}  // namespace gpu
+}  // namespace tfrt
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/gpu_ops.h.inc"
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_IR_GPU_OPS_H_

tensorflow/compiler/mlir/tfrt/ir/gpu_ops.td

@@ -0,0 +1,114 @@
+/* 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.
+==============================================================================*/
+
+#ifdef TFRT_GPU_OPS
+#else
+#define TFRT_GPU_OPS
+
+include "tfrt/tfrt_op_base.td"
+include "tfrt/compiler/opdefs/tfrt_op_interfaces.td"
+include "tfrt/compiler/opdefs/tfrt_traits.td"
+include "tfrt_fallback.td"
+
+def TFRT_GPU_Dialect : Dialect {
+  let name = "gpurt";
+
+  let description = [{
+    The TFRT GPU Dialect.
+  }];
+
+  let cppNamespace = "::tfrt::gpu";
+}
+
+class Gpu_Op<string mnemonic, list<Trait> traits = []> :
+    Op<TFRT_GPU_Dialect, mnemonic, traits> {
+}
+
+// TODO(b/260267885): We may add a device argument when we want to support
+// GPU MIG.
+def TransferToDeviceOp: Gpu_Op<"transfer_to_device"> {
+  let summary = "Transfer a CPU tensor to device.";
+
+  let description = [{
+    Transfer a CPU tensor to device.
+
+    Example:
+      %device_tensor = gpurt.transfer_to_device %cpu_tensor
+  }];
+
+  let arguments = (ins TFTensorType);
+  let results = (outs TFTensorType);
+  let assemblyFormat = "operands attr-dict";
+}
+
+// TODO(b/260267885): We may add a device argument when we want to support
+// GPU MIG.
+def TransferFromDeviceOp: Gpu_Op<"transfer_from_device"> {
+  let summary = "Transfer a tensor from device.";
+
+  let description = [{
+    Transfer a tensor from device.
+
+    Example:
+      %cpu_tensor = gpurt.transfer_from_device %device_tensor
+  }];
+
+  let arguments = (ins TFTensorType);
+  let results = (outs TFTensorType);
+  let assemblyFormat = "operands attr-dict";
+}
+
+// TODO(b/260267885): We may add a device argument when we want to support
+// GPU MIG.
+def MaybeTransferVariableOp: Gpu_Op<"maybe_transfer_variable"> {
+  let summary = "Transfer a CPU variable tensor to device.";
+  let description = [{
+    Transfer a CPU variable tensor to device if the variable has not been
+    transferred before.
+
+    Example:
+      %device_var = gpurt.maybe_transfer_variable %cpu_var
+  }];
+
+  let arguments = (ins TFTensorType);
+  let results = (outs TFTensorType);
+  let assemblyFormat = "operands attr-dict";
+}
+
+def CompileAndExecuteOp: Gpu_Op<"compile_and_execute"> {
+  let summary = "GPU compile and execute operation.";
+  let description = [{
+    The op compiles and executes a GPU cluster function.
+
+    func_name is the name of the function to be executed on GPU.
+    resource_indices are the indices of inputs that are resources.
+    used_output_indices are the indices of outputs that have users.
+
+    Example:
+      %results = gpurt.compile_and_execute {func_name = "xla_func_0", resource_indices = [1] ...}
+  }];
+
+  let arguments = (ins
+    Variadic<TFTensorType>:$operands,
+    StrAttr:$func_name,
+    I64ArrayAttr:$resource_indices,
+    I64ArrayAttr:$used_output_indices
+  );
+  let results = (outs
+    Variadic<TFTensorType>:$results
+  );
+}
+
+#endif // TFRT_GPU_OPS

tensorflow/compiler/mlir/tfrt/ir/mlrt/BUILD

@@ -0,0 +1,157 @@
+load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+td_library(
+    name = "mlrt_td_files",
+    srcs = [
+        "mlrt_dialect.td",
+        "mlrt_ops.td",
+    ],
+    includes = ["."],
+    visibility = ["//visibility:private"],
+    deps = [
+        "@llvm-project//mlir:OpBaseTdFiles",
+        "@llvm-project//mlir:SideEffectInterfacesTdFiles",
+    ],
+)
+
+gentbl_cc_library(
+    name = "mlrt_ops_inc_gen",
+    tbl_outs = {
+        "mlrt_ops.h.inc": ["-gen-op-decls"],
+        "mlrt_ops.cpp.inc": ["-gen-op-defs"],
+    },
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "mlrt_ops.td",
+    deps = [":mlrt_td_files"],
+)
+
+cc_library(
+    name = "mlrt_ops",
+    srcs = [
+        "mlrt_dialect.cc",
+        "mlrt_ops.cc",
+    ],
+    hdrs = [
+        "mlrt_dialect.h",
+        "mlrt_ops.h",
+    ],
+    visibility = ["//tensorflow/compiler/mlir/tfrt:__subpackages__"],
+    deps = [
+        ":mlrt_ops_inc_gen",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:InliningUtils",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+td_library(
+    name = "tf_mlrt_td_files",
+    srcs = [
+        "tf_mlrt_dialect.td",
+        "tf_mlrt_ops.td",
+        "tf_ops.td",
+    ],
+    includes = ["."],
+    deps = [
+        ":mlrt_td_files",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops_td_files",
+        "@llvm-project//mlir:OpBaseTdFiles",
+        "@llvm-project//mlir:SideEffectInterfacesTdFiles",
+        "@tf_runtime//:compiler_td_files",
+    ],
+)
+
+td_library(
+    name = "tf_mlrt_tpu_td_files",
+    srcs = [
+        "tf_mlrt_tpu_ops.td",
+    ],
+    includes = ["."],
+    visibility = ["//visibility:private"],
+    deps = [
+        ":mlrt_td_files",
+        ":tf_mlrt_td_files",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops_td_files",
+        "@llvm-project//mlir:OpBaseTdFiles",
+        "@llvm-project//mlir:SideEffectInterfacesTdFiles",
+    ],
+)
+
+gentbl_cc_library(
+    name = "tf_mlrt_ops_inc_gen",
+    tbl_outs = {
+        "tf_mlrt_ops.h.inc": ["-gen-op-decls"],
+        "tf_mlrt_ops.cpp.inc": ["-gen-op-defs"],
+    },
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "tf_mlrt_ops.td",
+    deps = [":tf_mlrt_td_files"],
+)
+
+gentbl_cc_library(
+    name = "tf_mlrt_tpu_ops_inc_gen",
+    tbl_outs = {
+        "tf_mlrt_tpu_ops.h.inc": ["-gen-op-decls"],
+        "tf_mlrt_tpu_ops.cpp.inc": ["-gen-op-defs"],
+    },
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "tf_mlrt_tpu_ops.td",
+    deps = [":tf_mlrt_tpu_td_files"],
+)
+
+gentbl_cc_library(
+    name = "tf_ops_inc_gen",
+    tbl_outs = {
+        "tf_ops.h.inc": ["-gen-op-decls"],
+        "tf_ops.cpp.inc": ["-gen-op-defs"],
+    },
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "tf_ops.td",
+    deps = [":tf_mlrt_td_files"],
+)
+
+cc_library(
+    name = "tf_mlrt_ops",
+    srcs = ["tf_mlrt_ops.cc"],
+    hdrs = ["tf_mlrt_ops.h"],
+    visibility = [
+        # copybara:uncomment "//learning/brain/experimental/tfrt/mlrt/application/tensorflow/tests:__subpackages__",
+        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
+    ],
+    deps = [
+        ":mlrt_ops",
+        ":tf_mlrt_ops_inc_gen",
+        ":tf_ops_inc_gen",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_side_effects",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:InliningUtils",
+        "@llvm-project//mlir:SideEffectInterfaces",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:Transforms",
+        "@tf_runtime//:compiler_tfrt_op_interfaces",
+        "@tf_runtime//:compiler_tfrt_traits",
+    ],
+)
+
+cc_library(
+    name = "tf_mlrt_tpu_ops",
+    srcs = ["tf_mlrt_tpu_ops.cc"],
+    hdrs = ["tf_mlrt_tpu_ops.h"],
+    visibility = ["//tensorflow/compiler/mlir/tfrt/transforms/mlrt:__subpackages__"],
+    deps = [
+        ":mlrt_ops",
+        ":tf_mlrt_ops",
+        ":tf_mlrt_tpu_ops_inc_gen",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:SideEffectInterfaces",
+    ],
+)

tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.cc

@@ -0,0 +1,96 @@
+/* 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/ir/mlrt/mlrt_dialect.h"
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/DialectImplementation.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Region.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Transforms/InliningUtils.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_ops.h"
+
+namespace mlrt {
+namespace compiler {
+
+namespace {
+
+struct MlrtInlinerInterface : public mlir::DialectInlinerInterface {
+  using DialectInlinerInterface::DialectInlinerInterface;
+  bool isLegalToInline(mlir::Operation *op, mlir::Region *dest,
+                       bool would_be_cloned,
+                       mlir::IRMapping &mapping) const final {
+    // All mlrt dialect ops can be inlined.
+    return true;
+  }
+};
+
+}  // namespace
+
+MlrtDialect::MlrtDialect(mlir::MLIRContext *context)
+    : mlir::Dialect(/*name=*/"mlrt", context,
+                    mlir::TypeID::get<MlrtDialect>()) {
+  addTypes<FutureType>();
+  addTypes<PromiseType>();
+  addTypes<AsyncHandleType>();
+  addInterfaces<MlrtInlinerInterface>();
+
+  addOperations<
+#define GET_OP_LIST
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_ops.cpp.inc"
+      >();
+}
+
+// Parse a type registered to this dialect.
+mlir::Type MlrtDialect::parseType(mlir::DialectAsmParser &parser) const {
+  llvm::StringRef keyword;
+  if (parser.parseKeyword(&keyword)) return mlir::Type();
+
+  if (keyword == "future") return FutureType::get(getContext());
+  if (keyword == "promise") return PromiseType::get(getContext());
+  if (keyword == "async_handle") return AsyncHandleType::get(getContext());
+
+  parser.emitError(parser.getNameLoc(), "unknown type: ") << keyword;
+  return mlir::Type();
+}
+
+// Print a type registered to this dialect.
+void MlrtDialect::printType(mlir::Type type,
+                            mlir::DialectAsmPrinter &os) const {
+  if (mlir::isa<FutureType>(type)) {
+    os << "future";
+    return;
+  }
+
+  if (mlir::isa<PromiseType>(type)) {
+    os << "promise";
+    return;
+  }
+
+  if (mlir::isa<AsyncHandleType>(type)) {
+    os << "async_handle";
+    return;
+  }
+
+  llvm_unreachable("unexpected mlrt type kind");
+}
+
+}  // namespace compiler
+}  // namespace mlrt

tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.h

@@ -0,0 +1,62 @@
+/* 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.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_IR_MLRT_MLRT_DIALECT_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_IR_MLRT_MLRT_DIALECT_H_
+
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Dialect.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/OpImplementation.h"  // from @llvm-project
+
+namespace mlrt {
+namespace compiler {
+
+class MlrtDialect : public mlir::Dialect {
+ public:
+  explicit MlrtDialect(mlir::MLIRContext *context);
+  static llvm::StringRef getDialectNamespace() { return "mlrt"; }
+
+  mlir::Type parseType(mlir::DialectAsmParser &parser) const override;
+  void printType(mlir::Type type, mlir::DialectAsmPrinter &os) const override;
+};
+
+// The MLIR type represents a C++ mlrt::Future.
+class FutureType
+    : public mlir::Type::TypeBase<FutureType, mlir::Type, mlir::TypeStorage> {
+ public:
+  using Base::Base;
+  static constexpr mlir::StringLiteral name = "mlrt.compiler.future";
+};
+
+// The MLIR type represents a C++ mlrt::Promise.
+class PromiseType
+    : public mlir::Type::TypeBase<PromiseType, mlir::Type, mlir::TypeStorage> {
+ public:
+  using Base::Base;
+  static constexpr mlir::StringLiteral name = "mlrt.compiler.promise";
+};
+
+// The MLIR type represents a C++ mlrt::AsyncHandle.
+class AsyncHandleType : public mlir::Type::TypeBase<AsyncHandleType, mlir::Type,
+                                                    mlir::TypeStorage> {
+ public:
+  using Base::Base;
+  static constexpr mlir::StringLiteral name = "mlrt.compiler.async_handle";
+};
+
+}  // namespace compiler
+}  // namespace mlrt
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_IR_MLRT_MLRT_DIALECT_H_

tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.td

@@ -0,0 +1,55 @@
+/* Copyright 2019 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.
+==============================================================================*/
+#ifdef MLRT_DIALECT
+#else
+#define MLRT_DIALECT
+
+include "mlir/IR/OpBase.td"
+
+def Mlrt_Dialect : Dialect {
+  let name = "mlrt";
+
+  let description = [{
+    The MLRT Dialect.
+  }];
+
+  let cppNamespace = "::mlrt::compiler";
+}
+
+def MlrtFutureType : DialectType<Mlrt_Dialect,
+    CPred<"::llvm::isa<::mlrt::compiler::FutureType>($_self)">, "!mlrt.future type">,
+    BuildableType<"$_builder.getType<::mlrt::compiler::FutureType>()"> {
+  let description = [{
+    `!mlrt.future type` represents a C++ mlrt::Future.
+  }];
+}
+
+def MlrtPromiseType : DialectType<Mlrt_Dialect,
+    CPred<"::llvm::isa<::mlrt::compiler::PromiseType>($_self)">, "!mlrt.promise type">,
+    BuildableType<"$_builder.getType<::mlrt::compiler::PromiseType>()"> {
+  let description = [{
+    `!mlrt.promise type` represents a C++ mlrt::Promise.
+  }];
+}
+
+def MlrtAsyncHandleType : DialectType<Mlrt_Dialect,
+    CPred<"::llvm::isa<::mlrt::compiler::AsyncHandleType>($_self)">, "!mlrt.async_handle type">,
+    BuildableType<"$_builder.getType<::mlrt::compiler::AsyncHandleType>()"> {
+  let description = [{
+    `!mlrt.async_handle type` represents a C++ mlrt::AsyncHandle.
+  }];
+}
+
+#endif

tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_ops.cc

@@ -0,0 +1,28 @@
+/* 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/ir/mlrt/mlrt_ops.h"
+
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/DialectImplementation.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+
+//===----------------------------------------------------------------------===//
+// TableGen'd op method definitions
+//===----------------------------------------------------------------------===//
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_ops.cpp.inc"

tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_ops.h

@@ -0,0 +1,27 @@
+/* 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.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_IR_MLRT_MLRT_OPS_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_IR_MLRT_MLRT_OPS_H_
+
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Dialect.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/OpImplementation.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.h"
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_ops.h.inc"
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_IR_MLRT_MLRT_OPS_H_

tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_ops.td

@@ -0,0 +1,240 @@
+/* Copyright 2019 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.
+==============================================================================*/
+#ifdef MLRT_OPS
+#else
+#define MLRT_OPS
+
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.td"
+
+class Mlrt_Op<string mnemonic, list<Trait> traits = []> :
+    Op<Mlrt_Dialect, mnemonic, !listconcat(traits, [IsolatedFromAbove])> {
+}
+
+def CondOp: Mlrt_Op<"cond", []> {
+  let summary = "mlrt.cond op";
+
+  let description = [{
+    Execute $a_true_fn with $args if $cond is true; otherwise, %b_false_fn is
+    executed.
+  }];
+
+  let arguments = (ins
+    I1:$cond,
+    Variadic<AnyType>:$args,
+    SymbolRefAttr:$a_true_fn,
+    SymbolRefAttr:$b_false_fn
+  );
+
+  let results = (outs
+    Variadic<AnyType>:$results
+  );
+
+  let assemblyFormat = [{
+    $cond $a_true_fn $b_false_fn `(` $args `)` attr-dict `:` `(` type($args) `)` `->` `(` type($results) `)`
+  }];
+}
+
+def AsyncOp: Mlrt_Op<"async", []> {
+  let summary = "Launches a function asynchronously.";
+
+  let description = [{
+    Launch a function asynchronously.
+
+    $args: a list of arguments to be passed.
+    $callee: The function to be launched. Its return op must not have operands.
+
+    $handle: This op returns a handle object that manages the context of the async execution.
+  }];
+
+  let arguments = (ins
+    Variadic<AnyType>:$args,
+    SymbolRefAttr:$callee
+  );
+
+  let results = (outs
+    MlrtAsyncHandleType:$handle
+  );
+
+  let assemblyFormat = "`(` $args `)` attr-dict `:` functional-type($args, $handle)";
+}
+
+def AwaitHandleOp: Mlrt_Op<"await_handle", []> {
+  let summary = "Awaits an async execution ";
+
+  let description = [{
+    Awaits an async execution.
+
+    $handle: The handle returned by mlrt.async op.
+  }];
+
+  let arguments = (ins
+    MlrtAsyncHandleType:$handle
+  );
+
+  let assemblyFormat = "operands attr-dict";
+}
+
+def AwaitAllHandleOp: Mlrt_Op<"await_all_handle", []> {
+  let summary = "Awaits multiple async executions";
+
+  let description = [{
+    Awaits multiple async execution.
+
+    $handles: A list of handles returned by mlrt.async ops.
+  }];
+
+  let arguments = (ins
+    Variadic<MlrtAsyncHandleType>:$handles
+  );
+
+  let assemblyFormat = "operands attr-dict `:` type($handles)";
+}
+
+def AwaitControlOp: Mlrt_Op<"await_control", []> {
+  let summary = "Await a signal from a future";
+
+  let description = [{
+    Await a signal, instead of a value, from a future.
+
+    $future: A value of !mlrt.future type.
+  }];
+
+  let arguments = (ins
+    MlrtFutureType:$future
+  );
+
+  let assemblyFormat = "operands attr-dict";
+}
+
+def AwaitAllControlOp: Mlrt_Op<"await_all_control", []> {
+  let summary = "Awaits multiple signals";
+
+  let description = [{
+    Awaits multiple signals
+
+    $futures: A list of !mlrt.futures
+  }];
+
+  let arguments = (ins
+    Variadic<MlrtFutureType>:$futures
+  );
+
+  let assemblyFormat = "operands attr-dict `:` type($futures)";
+}
+
+def PromiseControlOp: Mlrt_Op<"promise_control", []> {
+  let summary = "Set a control promise";
+
+  let description = [{
+    Set a control promise.
+
+    $promise: A value of !mlrt.promise type.
+  }];
+
+  let arguments = (ins
+    MlrtPromiseType:$promise
+  );
+
+  let assemblyFormat = "operands attr-dict";
+}
+
+def CaseOp : Mlrt_Op<"case"> {
+  let summary = "An n-way switch statement which calls a single branch function.";
+  let description = [{
+    An n-way switch statement, implementing the following:
+        ```
+        switch (branch_index) {
+          case 0:
+            outputs = branches[0](inputs);
+            break;
+          case 1:
+            outputs = branches[1](inputs);
+            break;
+          ...
+          case [[nbranches-1]]:
+          default:
+            outputs = branches[nbranches-1](inputs);
+            break;
+        }
+        ```
+    Example: %res = mlrt.case %branch_idx [@branch0, @branch1] (%arg0, %arg1) : (i32, i32) -> (i32)
+  }];
+
+  let arguments = (ins I32:$branch_index,
+                       ConfinedAttr<SymbolRefArrayAttr, [ArrayMinCount<1>]>:$branches,
+                       Variadic<AnyType>:$branch_operands);
+
+  let results = (outs Variadic<AnyType>:$branch_outputs);
+  let assemblyFormat = [{
+    $branch_index $branches `(` $branch_operands `)` attr-dict `:` `(` type($branch_operands) `)` `->` `(` type($branch_outputs) `)`
+  }];
+}
+
+def AllocateControlFuturesOp: Mlrt_Op<"allocate_control_futures", [AttrSizedResultSegments]> {
+  let summary = "Allocate futures and corresponding promises";
+
+  let description = [{
+    Allocate futures and corresponding promises.
+
+    $num: The number of futures to be allocated.
+
+    $promises: There are $num promises, and promises[i] shares the state with futures[i].
+    $futures: There are $num futures, and futures[i] shares the state with promises[i].
+  }];
+
+  let arguments = (ins
+    I32Attr:$num
+  );
+
+  let results = (outs
+    Variadic<MlrtPromiseType>:$promises,
+    Variadic<MlrtFutureType>:$futures
+  );
+}
+
+def WhileOp : Mlrt_Op<"while", []> {
+  let summary = "while operation";
+  let description = [{
+    cond: The boolean to control whether the first iteration should be
+      executed.
+    operands: The arguments to the first iteration.
+    results: The results of the last iteration. The number and types of results
+      excluding the last one are the same as the number and types of operands. The
+      last element of results is an I1 value that is false.
+    body_fn: The body function that takes the arguments and returns the results
+      that includes an I1 value to indicate whether next iteration should be executed.
+
+    The pseudo code:
+
+    while(cond) {
+      results = body_fn(operands)
+      cond = results#1
+    }
+    return results
+
+  }];
+
+  let arguments = (ins I1:$cond,
+                       Variadic<AnyType>:$arguments,
+                       FlatSymbolRefAttr:$body_fn);
+
+  let results = (outs Variadic<AnyType>);
+
+  let assemblyFormat = [{
+    $cond $body_fn `(` $arguments `)` attr-dict `:` `(` type($arguments) `)` `->` `(` type(results) `)`
+  }];
+}
+#endif

tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_dialect.td

@@ -0,0 +1,56 @@
+/* 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.
+==============================================================================*/
+#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";
+
+  let description = [{
+    The TF MLRT Dialect.
+  }];
+
+  let cppNamespace = "::tensorflow::tf_mlrt";
+}
+
+class TensorflowMlrt_Op<string mnemonic, list<Trait> traits = []> :
+    Op<TensorflowMlrt_Dialect, mnemonic, !listconcat(traits, [IsolatedFromAbove])> {
+}
+
+// This corresponds to tensorflow::Tensor.
+def TFTensorType : DialectType<TensorflowMlrt_Dialect,
+    CPred<"::llvm::isa<::tensorflow::tf_mlrt::TFTensorType>($_self)">, "!tf_mlrt.tensor type">,
+    BuildableType<"$_builder.getType<::tensorflow::tf_mlrt::TFTensorType>()"> {
+  let description = [{
+    `!tf_mlrt.tensor type` represents a tensorflow::Tensor.
+  }];
+}
+
+// This corresponds to tensorflow::Device* .
+def TFDeviceType : DialectType<TensorflowMlrt_Dialect,
+    CPred<"::llvm::isa<::tensorflow::tf_mlrt::TFDeviceType>($_self)">, "!tf_mlrt.device type">,
+    BuildableType<"$_builder.getType<::tensorflow::tf_mlrt::TFDeviceType>()"> {
+  let description = [{
+    `!tf_mlrt.device type` represents a tensorflow::device.
+  }];
+}
+
+#endif

tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_ops.cc

@@ -0,0 +1,96 @@
+/* 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/ir/mlrt/tf_mlrt_ops.h"
+
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/DialectImplementation.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Transforms/InliningUtils.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.h"
+
+namespace tensorflow {
+namespace tf_mlrt {
+
+namespace {
+
+struct TensorflowMlrtInlinerInterface : public mlir::DialectInlinerInterface {
+  using DialectInlinerInterface::DialectInlinerInterface;
+  bool isLegalToInline(mlir::Operation *op, mlir::Region *dest,
+                       bool would_be_cloned,
+                       mlir::IRMapping &mapping) const final {
+    // All tf_mlrt dialect ops can be inlined.
+    return true;
+  }
+  // Note that CallOp and ReturnOp are handled by func; so need to implement
+  // handleTerminator.
+};
+
+}  // namespace
+
+TensorflowMlrtDialect::TensorflowMlrtDialect(mlir::MLIRContext *context)
+    : mlir::Dialect(/*name=*/"tf_mlrt", context,
+                    mlir::TypeID::get<TensorflowMlrtDialect>()) {
+  addTypes<TFTensorType, TFDeviceType>();
+  addInterfaces<TensorflowMlrtInlinerInterface>();
+
+  addOperations<
+#define GET_OP_LIST
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_ops.cpp.inc"
+      >();
+
+  addOperations<
+#define GET_OP_LIST
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_ops.cpp.inc"
+      >();
+}
+
+// Parse a type registered to this dialect.
+mlir::Type TensorflowMlrtDialect::parseType(
+    mlir::DialectAsmParser &parser) const {
+  llvm::StringRef keyword;
+  if (parser.parseKeyword(&keyword)) return mlir::Type();
+
+  if (keyword == "tensor") return TFTensorType::get(getContext());
+
+  parser.emitError(parser.getNameLoc(), "unknown type: ") << keyword;
+  return mlir::Type();
+}
+
+// Print a type registered to this dialect.
+void TensorflowMlrtDialect::printType(mlir::Type type,
+                                      mlir::DialectAsmPrinter &os) const {
+  if (mlir::isa<TFTensorType>(type)) {
+    os << "tensor";
+    return;
+  }
+
+  llvm_unreachable("unexpected tf_mlrt type kind");
+}
+
+}  // namespace tf_mlrt
+}  // namespace tensorflow
+
+//===----------------------------------------------------------------------===//
+// TableGen'd op method definitions
+//===----------------------------------------------------------------------===//
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_ops.cpp.inc"
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_ops.cpp.inc"

tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_ops.h

@@ -0,0 +1,63 @@
+/* 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.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_IR_MLRT_TF_MLRT_OPS_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_IR_MLRT_TF_MLRT_OPS_H_
+
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Dialect.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/OpImplementation.h"  // from @llvm-project
+#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_side_effects.h"
+#include "tfrt/compiler/opdefs/tfrt_op_interfaces.h"  // from @tf_runtime
+#include "tfrt/compiler/opdefs/tfrt_traits.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace tf_mlrt {
+
+class TensorflowMlrtDialect : public mlir::Dialect {
+ public:
+  explicit TensorflowMlrtDialect(mlir::MLIRContext *context);
+  static llvm::StringRef getDialectNamespace() { return "tf_mlrt"; }
+
+  mlir::Type parseType(mlir::DialectAsmParser &parser) const override;
+  void printType(mlir::Type type, mlir::DialectAsmPrinter &os) const override;
+};
+
+// The MLIR type represents a tensorflow::Tensor.
+class TFTensorType
+    : public mlir::Type::TypeBase<TFTensorType, mlir::Type, mlir::TypeStorage> {
+ public:
+  using Base::Base;
+  static constexpr mlir::StringLiteral name = "tensorflow.tf_mlrt.tf_tensor";
+};
+
+// The MLIR type represents a tensorflow::Device*
+class TFDeviceType
+    : public mlir::Type::TypeBase<TFDeviceType, mlir::Type, mlir::TypeStorage> {
+ public:
+  using Base::Base;
+  static constexpr mlir::StringLiteral name = "tensorflow.tf_mlirt.tf_device";
+};
+
+}  // namespace tf_mlrt
+}  // namespace tensorflow
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_ops.h.inc"
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_ops.h.inc"
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_IR_MLRT_TF_MLRT_OPS_H_

tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_ops.td

@@ -0,0 +1,577 @@
+/* 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.
+==============================================================================*/
+#ifdef TF_MLRT_OPS
+#else
+#define TF_MLRT_OPS
+
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_dialect.td"
+include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.td"
+
+def CreateOp: TensorflowMlrt_Op<"createop", []> {
+  let summary = "The Fallback CreateOp";
+
+  let description = [{
+    The CreateOp creates the tensorflow::OpKernel in the fallback context.
+  }];
+
+  let arguments = (ins
+    StrAttr:$node_def,
+    I32Attr:$op_key
+  );
+
+  let assemblyFormat = "attr-dict";
+}
+
+def ConstOp: TensorflowMlrt_Op<"constop", []> {
+  let summary = "The tf_mlrt ConstOp";
+
+  let description = [{
+    The ConstOp creates a constant tensorflow::Tensor from serialized proto.
+  }];
+
+  let arguments = (ins
+    StrAttr:$tensor_proto
+  );
+
+  let results = (outs
+    TFTensorType:$result
+  );
+
+  let assemblyFormat = "attr-dict";
+}
+
+def ExecuteOp : TensorflowMlrt_Op<"executeop", []> {
+  let summary = "The Fallback ExecuteOp";
+  let description = [{
+    The ExecuteOp executes an operation on the specified device.
+  }];
+
+  let arguments = (ins
+    Variadic<TFTensorType>:$args,
+    StrAttr:$node_def,
+    I32Attr:$op_key
+  );
+
+  let results = (outs
+    Variadic<TFTensorType>:$results
+  );
+
+  let assemblyFormat = "`(` $args `)` attr-dict `:` functional-type($args, $results)";
+}
+
+def ExecuteOpWithDevice: TensorflowMlrt_Op<"executeop.device", []> {
+  let summary = "The Fallback ExecuteOp with custom device";
+  let description = [{
+    The ExecuteOp executes an operation on the specified device using a custom device.
+  }];
+
+  let arguments = (ins
+    TFDeviceType:$device,
+    Variadic<TFTensorType>:$args,
+    StrAttr:$node_def,
+    I32Attr:$op_key
+  );
+
+  let results = (outs
+    Variadic<TFTensorType>:$results
+  );
+
+  let assemblyFormat = "`(` $device`)` `(` $args `)` attr-dict `:` functional-type($args, $results)";
+}
+
+def AsyncExecuteOp : TensorflowMlrt_Op<"async_executeop", []> {
+  let summary = "The Fallback ExecuteOp for tensorflow::AsyncOpKernel";
+  let description = [{
+    The ExecuteOp executes an operation on the specified device asynchronously.
+  }];
+
+  let arguments = (ins
+    Variadic<TFTensorType>:$args,
+    StrAttr:$node_def,
+    I32Attr:$op_key
+  );
+
+  let results = (outs
+    Variadic<MlrtFutureType>:$results
+  );
+
+  let assemblyFormat = "`(` $args `)` attr-dict `:` functional-type($args, $results)";
+}
+
+def AsyncExecuteOpWithDevice : TensorflowMlrt_Op<"async_executeop.device", []> {
+  let summary = "The Fallback ExecuteOp for tensorflow::AsyncOpKernel";
+  let description = [{
+    The ExecuteOp executes an operation on the specified device asynchronously.
+  }];
+
+  let arguments = (ins
+    TFDeviceType:$device,
+    Variadic<TFTensorType>:$args,
+    StrAttr:$node_def,
+    I32Attr:$op_key
+  );
+
+  let results = (outs
+    Variadic<MlrtFutureType>:$results
+  );
+
+  let assemblyFormat = "`(` $device`)` `(` $args `)` attr-dict `:` functional-type($args, $results)";
+}
+
+def SetResourceOp : TensorflowMlrt_Op<"set_resource", []> {
+  let summary = "Set a tensor in resource array";
+
+  let description = [{
+    Set a tensor in resource array.
+
+    arg: the tensor to be set in the resource array.
+    index: the index in the resource array
+  }];
+
+  let arguments = (ins
+    TFTensorType:$arg,
+    I64Attr:$index
+  );
+
+  let results = (outs);
+
+  let assemblyFormat = "operands attr-dict";
+}
+
+def GetResourceOp : TensorflowMlrt_Op<"get_resource", []> {
+  let summary = "get a tensor in resource array";
+
+  let description = [{
+    Get a tensor in resource array.
+
+    indices: the indices in the resource array.
+    results: the tensor values for the corresponding indices.
+  }];
+
+  let arguments = (ins
+    I64ArrayAttr:$indices
+  );
+
+  let results = (outs
+    Variadic<TFTensorType>:$results
+  );
+
+  let assemblyFormat = "attr-dict `:` type($results)";
+}
+
+def AwaitOp: TensorflowMlrt_Op<"await", [Pure]> {
+  let summary = "Await a tensor from a !mlrt.future";
+
+  let description = [{
+    Await a tensor from a !mlrt.future.
+
+    $future: A value of type !mlrt.future. The underlying value must be a tensorflow tensor.
+
+    $result: a tensorflow tensor.
+  }];
+
+  let arguments = (ins
+    MlrtFutureType:$future
+  );
+
+  let results = (outs
+    TFTensorType:$result
+  );
+
+  let assemblyFormat = "operands attr-dict";
+}
+
+def AwaitAllOp: TensorflowMlrt_Op<"await_all", [Pure]> {
+  let summary = "Await tensors from a list of !mlrt.future";
+
+  let description = [{
+    Await tensors from a list of !mlrt.future.
+
+    $futures: A list of !mlrt.future. The underlying value must be tensorflow tensors.
+
+    $results: A list of tensorflow tensors.
+  }];
+
+  let arguments = (ins
+    Variadic<MlrtFutureType>:$futures
+  );
+
+  let results = (outs
+    Variadic<TFTensorType>:$results
+  );
+
+  let assemblyFormat = "$futures attr-dict `:` type($results)";
+}
+
+def PromiseOp: TensorflowMlrt_Op<"promise", []> {
+  let summary = "Set a tensor in a promise";
+
+  let description = [{
+    Set a tensor in a promise.
+
+    $promise: A value of type !mlrt.promise. The underlying value must be a tensorflow tensor.
+    $tensor: A tensorflow tensor.
+  }];
+
+  let arguments = (ins
+    MlrtPromiseType:$promise,
+    TFTensorType:$tensor
+  );
+
+  let assemblyFormat = "operands attr-dict";
+}
+
+def PromiseFutureOp: TensorflowMlrt_Op<"promise_future", []> {
+  let summary = "Set a tensor future in a promise";
+
+  let description = [{
+    Set a tensor future in a promise.
+
+    $promise: A value of type !mlrt.promise. The underlying value must be a tensorflow tensor.
+    $future: A value of type !mlrt.future. Must represents a tensorflow tensor.
+  }];
+
+  let arguments = (ins
+    MlrtPromiseType:$promise,
+    MlrtFutureType:$tensor
+  );
+
+  let assemblyFormat = "operands attr-dict";
+}
+
+def PromiseReturnOp: TensorflowMlrt_Op<"promise_return", [Terminator]> {
+  let summary = "Set a tensor in a promise and return from a function with no output";
+
+  let description = [{
+    Set a tensor in a promise and return from a function with no output.
+
+    $promise: A value of type !mlrt.promise. The underlying value must be a tensorflow tensor.
+    $tensor: A tensorflow tensor.
+  }];
+
+  let arguments = (ins
+    MlrtPromiseType:$promise,
+    TFTensorType:$tensor
+  );
+
+  let assemblyFormat = "operands attr-dict";
+}
+
+def AllocateFuturesOp: TensorflowMlrt_Op<"allocate_futures", [AttrSizedResultSegments]> {
+  let summary = "Allocate futures and promsies for tensorflow tensors";
+
+  let description = [{
+    Allocate futures and promsies for tensorflow tensors.
+
+    $num_futures: The number of futures to be allocated.
+
+    $promises: There are $num_futures promises. promises[i] shares the state with futures[i].
+    $futures: There are $num_futures futures. futures[i] shares the state with promises[i].
+  }];
+
+  let arguments = (ins
+    I32Attr:$num_futures
+  );
+
+  let results = (outs
+    Variadic<MlrtPromiseType>:$promises,
+    Variadic<MlrtFutureType>:$futures
+  );
+}
+
+def TensorToIntOp : TensorflowMlrt_Op<"tensor_to_int32", [Pure]> {
+  let summary = "Cast a Tensor to int32.";
+  let description = [{
+    Cast a Tensor to int32.
+
+    Example:
+      %one = tf_mlrt.tensor_to_int32 %src_tenosr
+  }];
+
+  let arguments = (ins TFTensorType:$src);
+  let results = (outs I32:$result);
+  let assemblyFormat = "operands attr-dict";
+}
+
+def BoolToTensorOp : TensorflowMlrt_Op<"bool_to_tensor", [Pure]> {
+  let summary = "Cast a boolean to a tensor.";
+  let description = [{
+    Cast a boolean to a tensor.
+
+    Example:
+      %one = tf_mlrt.bool_to_tensor %src_bool
+  }];
+
+  let arguments = (ins I1: $src);
+  let results = (outs TFTensorType:$result);
+
+  let assemblyFormat = "operands attr-dict";
+}
+
+def PredicateOp : TensorflowMlrt_Op<"predicate", [Pure]> {
+  let summary = "Converts a fallback tensor to a bool";
+
+  let description = [{
+    Note: this kernel is used for CPU tensors.
+
+    Converts a fallback tensor to a bool with the following rules:
+
+    - For 0D tensors, truthiness is determined by comparing against a "zero"
+      value. For numerical types it is the obvious zero. For strings it is the
+      empty string.
+
+    - For >0D tensors, truthiness is determined by looking at the number of
+      elements. If has zero elements, then the result is false. Otherwise the
+      result is true.
+
+    input: a fallback tensor representing the condition.
+    device: the name of the tensorflow device that is associated with the
+      input fallback tensor.
+
+    output: the converted bool.
+  }];
+
+  let arguments = (ins
+    TFTensorType:$input
+  );
+
+  let results = (outs
+    I1:$output
+  );
+
+  let assemblyFormat = "$input attr-dict";
+}
+
+def BatchFunctionOp : TensorflowMlrt_Op<"batch_function", [Pure]> {
+  let summary = "Fallback ExecuteOp specialized for tf.BatchFunction.";
+
+  let description = [{
+    This kernel executes a variant tf.BatchFunction kernel that supports having
+    the `f` attribute as a bytecode function.
+
+    Example:
+      %res = tf_mlrt.batch_function(%input, %captured_input)  {
+          device = "/device:CPU:0",
+          f = @batch_function,
+          node_def = "..."
+        } : (!tf_mlrt.tensor,!tf_mlrt.tensor) -> (!tf_mlrt.tensor)
+
+    Note that the trailing number indicates the number of results.
+  }];
+
+  let arguments = (ins
+    Variadic<TFTensorType>:$args,
+    StrAttr:$device,
+    SymbolRefAttr:$f,
+    StrAttr:$node_def
+  );
+
+  let results = (outs
+    Variadic<MlrtFutureType>:$results
+  );
+
+  let assemblyFormat = "`(` $args `)` attr-dict `:` functional-type($args, $results)";
+}
+
+def BatchFunctionWithDeviceOp : TensorflowMlrt_Op<"batch_function.device", [Pure]> {
+  let summary = "Fallback ExecuteOp specialized for tf.BatchFunction with custom device.";
+
+  let description = [{
+    This kernel executes a variant tf.BatchFunction kernel that supports having
+    the `f` attribute as a bytecode function with a custom device.
+
+    Example:
+      %res = tf_mlrt.batch_function.device(%custom_device, %input, %captured_input)  {
+          device = "/device:CPU:0",
+          f = @batch_function,
+          node_def = "..."
+        } : (!tf_mlrt.tensor,!tf_mlrt.tensor) -> (!tf_mlrt.tensor)
+
+    Note that the trailing number indicates the number of results.
+    Also note that the custom device argument is not the same as the device attribute. 
+    The latter is just the name in which the enclosure graph is executed.
+  }];
+
+  let arguments = (ins
+    TFDeviceType:$custom_device,
+    Variadic<TFTensorType>:$args,
+    StrAttr:$device,
+    SymbolRefAttr:$f,
+    StrAttr:$node_def
+  );
+
+  let results = (outs
+    Variadic<MlrtFutureType>:$results
+  );
+
+  let assemblyFormat = " `(` $custom_device`)` `(` $args `)` attr-dict `:` functional-type($args, $results)";
+}
+
+
+def CancelOp: TensorflowMlrt_Op<"cancel", []> {
+  let summary = "Handle cancellation request.";
+
+  let description = [{
+    This kernel will early terminate the program upon cancellation request (e.g. time out).
+  }];
+}
+
+def MapFnOp : TensorflowMlrt_Op<"map_fn", [AttrSizedOperandSegments, Pure]> {
+  let summary = "The Parallel Map for tf_mlrt dialect";
+  let description = [{
+    The Pmap executes body function in parallel for all ranges up to $max_iterations.
+
+    The pseudo code:
+      for(int i = 0; i < $max_iterations; i++) {
+        body_fn(MlrtFture($tensor_list_or_flow_in[i]),
+                MlrtPromise($tensor_list_or_flow_in[i+1]),
+                i, i, $invariant_args);
+      }
+
+      return $tensor_list_or_flow_in[$max_iterations]
+  }];
+
+  let arguments = (ins
+    TFTensorType:$max_iterations,
+    Variadic<TFTensorType>:$tensor_list_or_flow_in,
+    Variadic<TFTensorType>:$invariant_args,
+    FlatSymbolRefAttr:$body_fn,
+    I32Attr:$num_tensor_list_or_flow_in
+  );
+
+  let results = (outs
+     Variadic<TFTensorType>:$result
+  );
+
+  let assemblyFormat = "`(`$max_iterations`,` $tensor_list_or_flow_in`,` $invariant_args `)` attr-dict `:` functional-type(operands, results)";
+}
+
+def AsyncWhileOp : TensorflowMlrt_Op<"async_while", [Pure]> {
+  let summary = "Asynchronously execution of while op for tf_mlrt";
+  let description = [{
+    cond: The boolean to control whether the first iteration should be executed.
+    arguments: the last $immutable_size elements are invariants between iterations.
+    results: a list of futures.
+    body_fn: its input are [predicate_promise, arg0_future, arg0_promise, arg1_future, arg1_promise, ..., invariant_args] and has no returned results.
+  }];
+
+  let arguments = (ins TFTensorType:$cond,
+                       Variadic<TFTensorType>:$arguments,
+                       FlatSymbolRefAttr:$body_fn,
+                       I32Attr:$invariant_size);
+
+  let results = (outs Variadic<MlrtFutureType>);
+
+  let assemblyFormat = [{
+    $body_fn `(` $cond `,` $arguments `)` attr-dict `:` `(` type($cond) `,` type($arguments) `)` `->` `(` type(results) `)`
+  }];
+}
+
+def IfrtLoadVariableOp: TensorflowMlrt_Op<"ifrt_load_variable", [Pure]> {
+  let summary = "Loads a variable tensor as an IFRT array for mlrt";
+
+  let description = [{
+    This op loads a restored variable tensor as a tensor future. It is a
+    replacement of `tf.ReadVariableOp`.
+
+    This op returns a scalar string tensor containing the restored variable name, which can be
+    used as a key within the runtime, as well as a future for the tensor.
+
+    The `tf.IfrtCall` kernel uses the output $array_key.
+    Other ops executed by TFRT may make use of $tensor_future.
+  }];
+
+  let arguments = (ins
+    TFTensorType:$variable,
+    DefaultValuedAttr<BoolAttr, "false">:$used_by_host
+  );
+
+  let results = (outs
+    TFTensorType:$array_key,
+    MlrtFutureType: $tensor_future
+  );
+}
+
+def IfrtRestoreVariableOp: TensorflowMlrt_Op<"ifrt_restore_variable", []> {
+  let summary = "Restore variable tensors";
+  let description = [{
+    This is the MLRT version of tf.IfrtRestoreVariableOp.
+
+    This op restores tensors from a checkpoint asynchronously, similar to tf.RestoreV2.
+    It fuses variable assignment: tensors that are consumed by tf.AssignVariableOp
+    are handled internally via the $var_handles operands and do not appear in
+    the results.
+
+    Tensors that are consumed by other operations (e.g. tf.LookupTableImportV2)
+    will be returned as results in $tensors. Consequently, the number of
+    results is dynamic, containing only those tensors not assigned to variables
+    but consumed directly by other operations. A restored tensor with no
+    consumer will cause an error in the kernel.
+    This Op is specific to MLRT runtime and is not a stable interface for
+    serialization.
+
+  }];
+
+  let arguments = (ins
+    TFTensorType:$prefix,
+    TFTensorType:$tensor_names,
+    TFTensorType:$shape_and_slices,
+    Variadic<TFTensorType>:$var_handles,
+    TypeArrayAttr: $restored_dtypes,
+    DenseBoolArrayAttr: $truncate_in_cast,
+     // The names of the returned tensors. The order of names corresponds to the
+    // order of the returned tensors.
+    StrArrayAttr: $returned_tensor_names
+  );
+
+  let results = (outs
+    Res<Variadic<MlrtFutureType>, [{The restored tensor futures as requested by the $returned_Tensor_names attribute.}]>:$tensors
+  );
+}
+
+def MlrtIfrtResourceDeserializeOp: TensorflowMlrt_Op<"ifrt_resource_deserialize", []> {
+  let summary = "Deserialize resource vars.";
+  let description = [{
+    This is the MLRT version of the IfrtResourceDeserialize op.
+  }];
+
+  let arguments = (ins
+    TFTensorType:$resource_var,
+    TFTensorType:$input_dir,
+    StrAttr:$tensor_name
+  );
+}
+
+def AsyncIfrtCallOp : TensorflowMlrt_Op<"async_ifrt_call", []> {
+  let summary = "IFRT execution op for mlrt";
+  let description = [{
+    This op executes an IFRT program asynchronously.
+  }];
+
+  let arguments = (ins
+    I64Attr:$program_id,
+    I32ArrayAttr:$variable_arg_indices,
+    Variadic<TFTensorType>:$args
+  );
+
+  let results = (outs
+    Variadic<MlrtFutureType>:$results
+  );
+
+  let assemblyFormat = "`(` $args `)` attr-dict `:` functional-type($args, $results)";
+}
+
+#endif

tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_tpu_ops.cc

@@ -0,0 +1,46 @@
+/* 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/ir/mlrt/tf_mlrt_tpu_ops.h"
+
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/DialectImplementation.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_ops.h"
+
+namespace tensorflow {
+namespace tf_mlrt_tpu {
+
+TensorflowMlrtTpuDialect::TensorflowMlrtTpuDialect(mlir::MLIRContext *context)
+    : mlir::Dialect(/*name=*/"tf_mlrt_tpu", context,
+                    mlir::TypeID::get<TensorflowMlrtTpuDialect>()) {
+  addOperations<
+#define GET_OP_LIST
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_tpu_ops.cpp.inc"
+      >();
+}
+
+}  // namespace tf_mlrt_tpu
+}  // namespace tensorflow
+
+//===----------------------------------------------------------------------===//
+// TableGen'd op method definitions
+//===----------------------------------------------------------------------===//
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_tpu_ops.cpp.inc"

tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_tpu_ops.h

@@ -0,0 +1,39 @@
+/* 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.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_IR_MLRT_TF_MLRT_TPU_OPS_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_IR_MLRT_TF_MLRT_TPU_OPS_H_
+
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Dialect.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/OpImplementation.h"  // from @llvm-project
+#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
+
+namespace tensorflow {
+namespace tf_mlrt_tpu {
+
+class TensorflowMlrtTpuDialect : public mlir::Dialect {
+ public:
+  explicit TensorflowMlrtTpuDialect(mlir::MLIRContext *context);
+  static llvm::StringRef getDialectNamespace() { return "tf_mlrt_tpu"; }
+};
+
+}  // namespace tf_mlrt_tpu
+}  // namespace tensorflow
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_tpu_ops.h.inc"
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_IR_MLRT_TF_MLRT_TPU_OPS_H_

tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_tpu_ops.td

@@ -0,0 +1,82 @@
+/* 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.
+==============================================================================*/
+#ifdef TF_MLRT_TPU_OPS
+#else
+#define TF_MLRT_TPU_OPS
+
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_dialect.td"
+include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.td"
+
+def TensorflowMlrtTpu_Dialect : Dialect {
+  let name = "tf_mlrt_tpu";
+
+  let description = [{
+    The TF MLRT TPU Dialect.
+  }];
+
+  let cppNamespace = "::tensorflow::tf_mlrt_tpu";
+}
+
+class TensorflowMlrtTpu_Op<string mnemonic, list<Trait> traits = []> :
+    Op<TensorflowMlrtTpu_Dialect, mnemonic, !listconcat(traits, [IsolatedFromAbove])> {
+}
+
+def GetTpuHostDeviceOp : TensorflowMlrtTpu_Op<"get_tpu_host_device", [Pure]> {
+  let summary = "get the tpu host allocator that implements tensorflow::Device";
+
+  let results = (outs
+    TFDeviceType:$device
+  );
+
+  let assemblyFormat = "attr-dict";
+}
+
+def CompileAndExecuteOp : TensorflowMlrtTpu_Op<"compile_and_execute"> {
+  let summary = "tpu compile and execute operation";
+  let description = [{
+    tf_mlrt_tpu.compile_and_execute compiles a mlir tpu program and executes the compiled tpu program.
+
+    $mlir_module is a serialized MLIR module with a `main` function that contains target computation.
+    $metadata is a serialized TPUCompileMetadataProto describing the shapes and types of the inputs to the computation, as well as a mapping onto the TPU pod topology.
+    $constant_operand_indices are the indices of the inputs that are constant to the TPU program (e.g. weights in inference), the rest of the inputs are input tensors.
+    constant_operand_indices is sorted in ascending order.
+    $operands_with_static_shape are indices of operands that are tagged with a maximum static shape.
+    $producer_name is a string describing the name of the framework that added support for running this portion of the model on TPUs.
+
+    Example:
+      %rendezvous_key_base, %result = tf_mlrt_tpu.compile_and_execute (%operands) constant_operand_indices = [1, 3] metadata = "metadata..." mlir_module = "mlir_module..."
+  }];
+  let arguments = (ins
+    Variadic<TFTensorType>:$operands_and_static_shapes,
+    DenseI32ArrayAttr:$constant_operand_indices,
+    StrAttr:$metadata,
+    StrAttr:$mlir_module,
+    UI32Attr:$num_operands,
+    DenseI32ArrayAttr:$operands_with_static_shape,
+    StrAttr:$producer_name
+  );
+
+  let results = (outs
+    TFTensorType:$rendezvous_key_base,
+    Variadic<MlrtFutureType>:$results
+  );
+
+  let assemblyFormat = [{
+    `(` $operands_and_static_shapes `)` attr-dict `:` functional-type($operands_and_static_shapes, results)
+  }];
+}
+
+#endif

tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_ops.td

@@ -0,0 +1,182 @@
+/* Copyright 2015 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.
+==============================================================================*/
+#ifdef MLRT_TF_OPS
+#else
+#define MLRT_TF_OPS
+
+include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_dialect.td"
+include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.td"
+include "tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td"
+include "tfrt/compiler/opdefs/tfrt_op_interfaces.td"  // from @tf_runtime
+include "tfrt/compiler/opdefs/tfrt_traits.td"  // from @tf_runtime
+
+// tf_mlrt.tf_await returns a tensorflow Tensor. It is a fake op that is only
+// used during parallelization and has no runtime implementation.
+def TFAwaitOp: TensorflowMlrt_Op<"tf_await", [Pure, TFRT_CostFunctionInterface, TFRT_FixedCost<1>]> {
+  let summary = "Await a tensor from a !mlrt.future";
+
+  let description = [{
+    Await a tensor from a !mlrt.future.
+
+    $future: A value of type !mlrt.future. The underlying value must be a tensorflow tensor.
+
+    $result: a tensorflow tensor.
+  }];
+
+  let arguments = (ins
+    MlrtFutureType:$future
+  );
+
+  let results = (outs
+    TF_Tensor:$result
+  );
+}
+
+// tf_mlrt.tf_promise takes a tensorflow Tensor. It is a fake op that is only
+// used during parallelization and has no runtime implementation.
+def TFPromiseOp: TensorflowMlrt_Op<"tf_promise", [TF_MustExecute, TFRT_CostFunctionInterface, TFRT_FixedCost<1>]> {
+  let summary = "Set a tensor in a promise";
+
+  let description = [{
+    Set a tensor in a promise.
+
+    $promise: A value of type !mlrt.promise. The underlying value will always be a tensorflow tensor.
+    $tensor: A tensorflow tensor.
+  }];
+
+  let arguments = (ins
+    MlrtPromiseType:$promise,
+    TF_Tensor:$tensor
+  );
+}
+
+def TFMapFnOp : TensorflowMlrt_Op<"tf_map_fn", [AttrSizedOperandSegments, Pure]> {
+  let summary = "The Parallel Map for tf_mlrt dialect";
+  let description = [{
+    The Pmap executes body function in parallel for all ranges up to $max_iterations.
+
+    The pseudo code:
+      for(int i = 0; i < $max_iterations; i++) {
+        body_fn(MlrtFture($tensor_list_or_flow_in[i]),
+                MlrtPromise($tensor_list_or_flow_in[i+1]),
+                i, i, $invariant_args);
+      }
+
+      return $tensor_list_or_flow_in[$max_iterations]
+  }];
+
+  let arguments = (ins
+    TF_Tensor:$max_iterations,
+    Variadic<TF_Tensor>:$tensor_list_or_flow_in,
+    Variadic<TF_Tensor>:$invariant_args,
+    FlatSymbolRefAttr:$body_fn,
+    I32Attr:$num_tensor_list_or_flow_in
+  );
+
+  let results = (outs
+     Variadic<TF_Tensor>:$result
+  );
+
+  let assemblyFormat = "`(`$max_iterations`,` $tensor_list_or_flow_in`,` $invariant_args `)` attr-dict `:` functional-type(operands, results)";
+}
+
+def TFAsyncWhileOp : TensorflowMlrt_Op<"tf_async_while", [Pure]> {
+  let summary = "Asynchronously execution of while op for tf_mlrt";
+  let description = [{
+    cond: The boolean to control whether the first iteration should be executed.
+    arguments: The last $invariant_size elements are invariants between iterations.
+    results: a list of futures.
+    body_fn: input is [predicate_promise, arg0_future, arg0_promise, arg1_future, arg1_promise, ..., invariant_args] and has no returned results.
+  }];
+
+  let arguments = (ins TF_Tensor:$cond,
+                       Variadic<TF_Tensor>:$arguments,
+                       FlatSymbolRefAttr:$body_fn,
+                       I32Attr:$invariant_size);
+
+  let results = (outs Variadic<MlrtFutureType>);
+
+  let assemblyFormat = [{
+    $body_fn `(` $cond `,` $arguments `)` attr-dict `:` `(` type($cond) `,` type($arguments) `)` `->` `(` type(results) `)`
+  }];
+}
+
+def TFTPUCompileAndExecuteOp : TensorflowMlrt_Op<"tf_tpu_compile_and_execute", [TF_MustExecute]> {
+  let summary = "tpu compile and execute operation";
+  let description = [{
+    tf_mlrt_tpu.compile_and_execute compiles a mlir tpu program and executes the compiled tpu program.
+
+    $mlir_module is a serialized MLIR module with a `main` function that contains target computation.
+    $metadata is a serialized TPUCompileMetadataProto describing the shapes and types of the inputs to the computation, as well as a mapping onto the TPU pod topology.
+    $constant_operand_indices are the indices of the inputs that are constant to the TPU program (e.g. weights in inference), the rest of the inputs are input tensors.
+    constant_operand_indices is sorted in ascending order.
+    $operands_with_static_shape are indices of operands that are tagged with a maximum static shape.
+    $producer_name is a string describing the name of the framework that added support for running this portion of the model on TPUs.
+
+    Example:
+      %rendezvous_key_base, %result = tf_mlrt_tpu.compile_and_execute (%operands) constant_operand_indices = [1, 3] metadata = "metadata..." mlir_module = "mlir_module..."
+  }];
+  let arguments = (ins
+    Variadic<TF_Tensor>:$operands_and_static_shapes,
+    DenseI32ArrayAttr:$constant_operand_indices,
+    StrAttr:$metadata,
+    StrAttr:$mlir_module,
+    UI32Attr:$num_operands,
+    DenseI32ArrayAttr:$operands_with_static_shape,
+    StrAttr:$producer_name
+  );
+
+  let results = (outs
+    TF_Tensor:$rendezvous_key_base,
+    Variadic<TF_Tensor>:$results
+  );
+
+  let assemblyFormat = [{
+    `(` $operands_and_static_shapes `)` attr-dict `:` functional-type($operands_and_static_shapes, results)
+  }];
+}
+
+def TFIfrtLoadVariableOp: TensorflowMlrt_Op<"tf_ifrt_load_variable", [Pure]> {
+  let summary = "Loads a variable tensor as an IFRT array for mlrt";
+
+  let description = [{
+    This is the MLRT version of tf.IfrtLoadVariableOp.
+
+    This op loads a variable tensor as an IFRT array and binds it with the specified name.
+
+    This op is an replacement of `tf.ReadVariableOp` in the case that a constant
+    variable tensor is an input to the tpu program invoked by `tf.IfrtCall`.
+
+    After a `tf.ReadVariableOp` is lowered into `tf.IfrtLoadVariableOp`, the `tf.IfrtCall` kernel
+    will bind the loaded IFRT array by name with the tpu program's input.
+
+    This op returns a scalar string tensor as a key for user to look for the loaded array
+    and a future containing the restored tensor.
+  }];
+
+  let arguments = (ins
+    TF_Tensor:$variable,
+    DefaultValuedAttr<BoolAttr, "false">:$used_by_host
+  );
+
+  let results = (outs
+    TF_Tensor:$array_key,
+    MlrtFutureType: $tensor_future
+  );
+}
+
+
+
+#endif

tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback.cc

@@ -0,0 +1,72 @@
+/* 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/ir/tfrt_fallback.h"
+
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/DialectImplementation.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace tfrt {
+namespace fallback {
+
+FallbackDialect::FallbackDialect(MLIRContext *context)
+    : Dialect(/*name=*/"tfrt_fallback", context,
+              TypeID::get<FallbackDialect>()) {
+  addTypes<TFTensorType, TFAllocatorType>();
+
+  addOperations<
+#define GET_OP_LIST
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback.cpp.inc"
+      >();
+}
+
+/// Parse a type registered to this dialect.
+Type FallbackDialect::parseType(DialectAsmParser &parser) const {
+  StringRef keyword;
+  if (parser.parseKeyword(&keyword)) return Type();
+
+  if (keyword == "tf_tensor") return TFTensorType::get(getContext());
+  if (keyword == "tf_allocator") return TFAllocatorType::get(getContext());
+
+  parser.emitError(parser.getNameLoc(), "unknown type: ") << keyword;
+  return Type();
+}
+
+/// Print a type registered to this dialect.
+void FallbackDialect::printType(Type type, DialectAsmPrinter &os) const {
+  if (mlir::isa<TFTensorType>(type)) {
+    os << "tf_tensor";
+    return;
+  }
+
+  if (mlir::isa<TFAllocatorType>(type)) {
+    os << "tf_allocator";
+    return;
+  }
+
+  llvm_unreachable("unexpected fallback type kind");
+}
+
+}  // namespace fallback
+}  // namespace tfrt
+
+//===----------------------------------------------------------------------===//
+// TableGen'd op method definitions
+//===----------------------------------------------------------------------===//
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback.cpp.inc"

tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback.h

@@ -0,0 +1,60 @@
+/* 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.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_H_
+
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Dialect.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/OpImplementation.h"  // from @llvm-project
+#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
+
+using namespace mlir;  // NOLINT
+
+namespace tfrt {
+namespace fallback {
+
+// Dialect for fallback operations.
+class FallbackDialect : public Dialect {
+ public:
+  explicit FallbackDialect(MLIRContext *context);
+  static StringRef getDialectNamespace() { return "tfrt_fallback"; }
+
+  Type parseType(DialectAsmParser &parser) const override;
+  void printType(Type type, DialectAsmPrinter &os) const override;
+};
+
+// The MLIR type represents a tensorflow::Tensor.
+class TFTensorType : public Type::TypeBase<TFTensorType, Type, TypeStorage> {
+ public:
+  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
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback.h.inc"
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_H_

tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback.td

@@ -0,0 +1,50 @@
+/* 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.
+==============================================================================*/
+
+#ifdef TFRT_FALLBACK_OPS
+#else
+#define TFRT_FALLBACK_OPS
+
+include "mlir/Interfaces/SideEffectInterfaces.td"
+
+def Fallback_Dialect : Dialect {
+  let name = "tfrt_fallback";
+
+  let description = [{
+    The TFRT Fallback Dialect.
+  }];
+
+  let cppNamespace = "::tfrt::fallback";
+}
+
+// This corresponds to tensorflow::Tensor.
+def TFTensorType : DialectType<Fallback_Dialect,
+    CPred<"::llvm::isa<::tfrt::fallback::TFTensorType>($_self)">, "!tfrt_fallback.tf_tensor type">,
+    BuildableType<"$_builder.getType<::tfrt::fallback::TFTensorType>()"> {
+  let description = [{
+    `!tfrt_fallback.tf_tensor type` represents a tensorflow::Tensor.
+  }];
+}
+
+// This corresponds to tensorflow::Allocator.
+def TFAllocatorType : DialectType<Fallback_Dialect,
+    CPred<"::llvm::isa<::tfrt::fallback::TFAllocatorType>($_self)">, "!tfrt_fallback.tf_allocator type">,
+    BuildableType<"$_builder.getType<::tfrt::fallback::TFAllocatorType>()"> {
+  let description = [{
+    `!tfrt_fallback.tf_tensor type` represents a tensorflow::Tensor.
+  }];
+}
+
+#endif

tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_async.cc

@@ -0,0 +1,407 @@
+/* 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/ir/tfrt_fallback_async.h"
+
+#include <utility>
+
+#include "llvm/ADT/STLExtras.h"
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/DialectImplementation.h"  // from @llvm-project
+#include "mlir/IR/Matchers.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/OpImplementation.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/InliningUtils.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_common.h"
+#include "tfrt/basic_kernels/opdefs/tfrt_base.h"  // from @tf_runtime
+#include "tfrt/basic_kernels/opdefs/types.h"  // from @tf_runtime
+#include "tfrt/core_runtime/opdefs/attributes.h"  // from @tf_runtime
+#include "tfrt/core_runtime/opdefs/core_runtime.h"  // from @tf_runtime
+#include "tfrt/core_runtime/opdefs/sync/core_runtime.h"  // from @tf_runtime
+#include "tfrt/core_runtime/opdefs/types.h"  // from @tf_runtime
+
+namespace tfrt {
+namespace fallback_async {
+
+namespace {
+
+struct FallbackInlinerInterface : public mlir::DialectInlinerInterface {
+  using DialectInlinerInterface::DialectInlinerInterface;
+
+  bool isLegalToInline(Operation *op, Region *dest, bool would_be_cloned,
+                       IRMapping &) const final {
+    return true;
+  }
+};
+
+}  // namespace
+
+FallbackAsyncDialect::FallbackAsyncDialect(MLIRContext *context)
+    : Dialect(/*name=*/"tfrt_fallback_async", context,
+              TypeID::get<FallbackAsyncDialect>()) {
+  context->getOrLoadDialect<tfrt::fallback::FallbackDialect>();
+  context->getOrLoadDialect<compiler::TFRTDialect>();
+  context->getOrLoadDialect<corert::CoreRTDialect>();
+
+  allowUnknownTypes();
+
+  allowUnknownOperations();
+
+  addInterfaces<FallbackInlinerInterface>();
+
+  addOperations<
+#define GET_OP_LIST
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_async.cpp.inc"
+      >();
+}
+
+static Type GetChainType(Builder *builder) {
+  return builder->getType<compiler::ChainType>();
+}
+
+LogicalResult CreateOp::verify() {
+  return fallback_common::VerifyFallbackExecuteOp(*this);
+}
+LogicalResult ExecuteOp::verify() {
+  return fallback_common::VerifyFallbackExecuteOp(*this);
+}
+LogicalResult ExecuteOpSeq::verify() {
+  return fallback_common::VerifyFallbackExecuteOp(*this);
+}
+LogicalResult ExecuteOpWithAllocator::verify() {
+  return fallback_common::VerifyExecuteOpCommon(*this);
+}
+LogicalResult ExecuteOpSeqWithAllocator::verify() {
+  return fallback_common::VerifyExecuteOpCommon(*this);
+}
+LogicalResult BatchFunctionOp::verify() {
+  return fallback_common::VerifyExecuteOpCommon(*this);
+}
+
+ParseResult CreateOp::parse(OpAsmParser &parser, OperationState &result) {
+  fallback_common::ParseExecuteOpOptions parse_options;
+  parse_options.has_chain = true;
+  parse_options.has_key = true;
+  parse_options.has_device = true;
+  parse_options.has_func_attr = true;
+  parse_options.has_cost = false;
+
+  auto &builder = parser.getBuilder();
+  if (mlir::failed(fallback_common::ParseExecuteOpCommon(
+          parser, builder, result, builder.getType<fallback::TFTensorType>(),
+          parse_options)))
+    return mlir::failure();
+
+  mlir::IntegerAttr num_args;
+  if (parser.parseKeyword("num_args") || parser.parseLParen() ||
+      parser.parseAttribute(num_args, "num_args", result.attributes) ||
+      parser.parseRParen())
+    return mlir::failure();
+
+  return mlir::success();
+}
+ParseResult ExecuteOp::parse(OpAsmParser &parser, OperationState &result) {
+  fallback_common::ParseExecuteOpOptions parse_options;
+  parse_options.has_chain = false;
+  parse_options.has_key = true;
+  parse_options.has_device = true;
+  parse_options.has_func_attr = true;
+  parse_options.has_cost = true;
+
+  auto &builder = parser.getBuilder();
+  return fallback_common::ParseExecuteOpCommon(
+      parser, builder, result, builder.getType<fallback::TFTensorType>(),
+      parse_options);
+}
+ParseResult ExecuteOpSeq::parse(OpAsmParser &parser, OperationState &result) {
+  fallback_common::ParseExecuteOpOptions parse_options;
+  parse_options.has_chain = true;
+  parse_options.has_key = true;
+  parse_options.has_device = true;
+  parse_options.has_func_attr = true;
+  parse_options.has_cost = true;
+
+  auto &builder = parser.getBuilder();
+  return fallback_common::ParseExecuteOpCommon(
+      parser, builder, result, builder.getType<fallback::TFTensorType>(),
+      parse_options);
+}
+ParseResult ExecuteOpWithAllocator::parse(OpAsmParser &parser,
+                                          OperationState &result) {
+  auto &builder = parser.getBuilder();
+  llvm::SmallVector<mlir::OpAsmParser::UnresolvedOperand, 1> allocator;
+  if (parser.parseOperandList(allocator,
+                              /*requiredOperandCount=*/1,
+                              mlir::OpAsmParser::Delimiter::Paren))
+    return mlir::failure();
+
+  if (parser.resolveOperand(allocator.front(),
+                            builder.getType<fallback::TFAllocatorType>(),
+                            result.operands))
+    return mlir::failure();
+
+  fallback_common::ParseExecuteOpOptions parse_options;
+  parse_options.has_chain = false;
+  parse_options.has_key = true;
+  parse_options.has_device = true;
+  parse_options.has_func_attr = true;
+  parse_options.has_cost = true;
+
+  return fallback_common::ParseExecuteOpCommon(
+      parser, builder, result, builder.getType<fallback::TFTensorType>(),
+      parse_options);
+}
+ParseResult ExecuteOpSeqWithAllocator::parse(OpAsmParser &parser,
+                                             OperationState &result) {
+  auto &builder = parser.getBuilder();
+  llvm::SmallVector<mlir::OpAsmParser::UnresolvedOperand, 2>
+      chain_and_allocator;
+  if (parser.parseOperandList(chain_and_allocator,
+                              /*requiredOperandCount=*/2,
+                              mlir::OpAsmParser::Delimiter::Paren))
+    return mlir::failure();
+
+  auto &chain = chain_and_allocator[0];
+  auto &allocator = chain_and_allocator[1];
+
+  if (parser.resolveOperand(chain, builder.getType<compiler::ChainType>(),
+                            result.operands))
+    return mlir::failure();
+
+  if (parser.resolveOperand(allocator,
+                            builder.getType<fallback::TFAllocatorType>(),
+                            result.operands))
+    return mlir::failure();
+
+  // The first result is a chain.
+  result.types.push_back(builder.getType<compiler::ChainType>());
+
+  fallback_common::ParseExecuteOpOptions parse_options;
+  parse_options.has_chain = false;
+  parse_options.has_key = true;
+  parse_options.has_device = true;
+  parse_options.has_func_attr = true;
+  parse_options.has_cost = true;
+
+  return fallback_common::ParseExecuteOpCommon(
+      parser, builder, result, builder.getType<fallback::TFTensorType>(),
+      parse_options);
+}
+
+ParseResult BatchFunctionOp::parse(OpAsmParser &parser,
+                                   OperationState &result) {
+  fallback_common::ParseExecuteOpOptions parse_options;
+  parse_options.has_chain = false;
+  parse_options.has_key = false;
+  parse_options.has_device = true;
+  parse_options.has_func_attr = false;
+  parse_options.has_cost = false;
+  parse_options.has_op_name = false;
+  parse_options.has_symbol_ref = true;
+
+  auto &builder = parser.getBuilder();
+  return fallback_common::ParseExecuteOpCommon(
+      parser, builder, result, builder.getType<fallback::TFTensorType>(),
+      parse_options);
+}
+
+void CreateOp::print(OpAsmPrinter &p) {
+  CreateOp op = *this;
+  p << "(" << op.getInCh() << ") key("
+    << op->getAttrOfType<mlir::IntegerAttr>("op_key").getInt() << ") device("
+    << op->getAttr("device") << ") " << op->getAttr("op_name") << "()";
+
+  fallback_common::PrintExecuteOpCommon(p, op);
+  fallback_common::PrintExecuteOpFuncAttribute(p, op);
+
+  p << " num_args(" << op->getAttrOfType<mlir::IntegerAttr>("num_args").getInt()
+    << ')';
+}
+
+void ExecuteOp::print(OpAsmPrinter &p) {
+  ExecuteOp op = *this;
+  p << " key(" << op->getAttrOfType<mlir::IntegerAttr>("op_key").getInt()
+    << ") cost(" << op->getAttrOfType<mlir::IntegerAttr>("_tfrt_cost").getInt()
+    << ") device(" << op->getAttr("device") << ") " << op->getAttr("op_name")
+    << '(' << op.getArgs() << ')';
+
+  fallback_common::PrintExecuteOpCommon(p, op);
+  fallback_common::PrintExecuteOpFuncAttribute(p, op);
+  if (!op.getResults().empty()) p << " : " << op.getResults().size();
+}
+
+void ExecuteOpSeq::print(OpAsmPrinter &p) {
+  ExecuteOpSeq op = *this;
+  p << "(" << op.getInOpChain() << ") key("
+    << op->getAttrOfType<mlir::IntegerAttr>("op_key").getInt() << ") cost("
+    << op->getAttrOfType<mlir::IntegerAttr>("_tfrt_cost").getInt()
+    << ") device(" << op->getAttr("device") << ") " << op->getAttr("op_name")
+    << '(' << op.getArgs() << ')';
+
+  fallback_common::PrintExecuteOpCommon(p, op);
+  fallback_common::PrintExecuteOpFuncAttribute(p, op);
+  if (!op.getResults().empty()) p << " : " << op.getResults().size();
+}
+
+void ExecuteOpWithAllocator::print(OpAsmPrinter &p) {
+  ExecuteOpWithAllocator op = *this;
+  p << "(" << op.getAllocator() << ") key("
+    << op->getAttrOfType<mlir::IntegerAttr>("op_key").getInt() << ") cost("
+    << op->getAttrOfType<mlir::IntegerAttr>("_tfrt_cost").getInt()
+    << ") device(" << op->getAttr("device") << ") " << op->getAttr("op_name")
+    << '(' << op.getArgs() << ')';
+
+  fallback_common::PrintExecuteOpCommon(p, op);
+  fallback_common::PrintExecuteOpFuncAttribute(p, op);
+  if (!op.getResults().empty()) p << " : " << op.getResults().size();
+}
+
+void ExecuteOpSeqWithAllocator::print(OpAsmPrinter &p) {
+  ExecuteOpSeqWithAllocator op = *this;
+  p << "(" << op.getInOpChain() << ", " << op.getAllocator() << ") key("
+    << op->getAttrOfType<mlir::IntegerAttr>("op_key").getInt() << ") cost("
+    << op->getAttrOfType<mlir::IntegerAttr>("_tfrt_cost").getInt()
+    << ") device(" << op->getAttr("device") << ") " << op->getAttr("op_name")
+    << '(' << op.getArgs() << ')';
+
+  fallback_common::PrintExecuteOpCommon(p, op);
+  fallback_common::PrintExecuteOpFuncAttribute(p, op);
+  if (!op.getResults().empty()) p << " : " << op.getResults().size();
+}
+
+void BatchFunctionOp::print(OpAsmPrinter &p) {
+  BatchFunctionOp op = *this;
+  p << " device(" << op->getAttr("device") << ") " << op->getAttr("f") << " ("
+    << op.getArgs() << ")";
+
+  fallback_common::PrintExecuteOpCommon(p, op);
+}
+
+void ExecuteOp::getOpAttrs(
+    SmallVectorImpl<std::pair<StringRef, Attribute>> *op_attrs) {
+  fallback_common::GetExecuteOpAttrsCommon(
+      this->getContext(), this->getOpAttrs().getValue(), op_attrs);
+}
+
+//===----------------------------------------------------------------------===//
+// ConstDenseTensorOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult ConstDenseTensorOp::fold(FoldAdaptor) { return getValue(); }
+
+//===----------------------------------------------------------------------===//
+// CoreRTTensorHandleToFallbackTensorOp
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+// Simplifies pattern containing a corert const tensor op followed by a
+// `tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor` op to a single
+// tfrt_fallback_async const tensor.
+struct ConstCoreRTTensorHandleToFallbackTensorCanonicalization
+    : public OpRewritePattern<CoreRTTensorHandleToFallbackTensorOp> {
+  using OpRewritePattern<
+      CoreRTTensorHandleToFallbackTensorOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(CoreRTTensorHandleToFallbackTensorOp op,
+                                PatternRewriter &rewriter) const override {
+    SmallVector<Value, 1> new_values;
+    bool should_rewrite = false;
+    for (auto operand : op.getArgs()) {
+      if (auto corert_const_dense_tensor_op =
+              operand.getDefiningOp<corert::ConstDenseTensorOp>()) {
+        new_values.push_back(fallback_async::ConstDenseTensorOp::create(
+            rewriter, op.getLoc(), rewriter.getType<fallback::TFTensorType>(),
+            corert_const_dense_tensor_op.getValue()));
+        should_rewrite = true;
+        continue;
+      }
+      if (auto corert_const_string_tensor_op =
+              operand.getDefiningOp<corert::ConstStringTensorOp>()) {
+        new_values.push_back(fallback_async::ConstStringTensorOp::create(
+            rewriter, op.getLoc(), rewriter.getType<fallback::TFTensorType>(),
+            corert_const_string_tensor_op.getShape(),
+            corert_const_string_tensor_op.getValue()));
+        should_rewrite = true;
+        continue;
+      }
+      // To guarantee that the new values are in the same order as the old
+      // ones, we create individual ops for the non-canonicalizable operands.
+      // For simplicity, we don't consolidate these ops when all the
+      // non-canonicalizable operands are adjacent.
+      new_values.push_back(
+          fallback_async::CoreRTTensorHandleToFallbackTensorOp::create(
+              rewriter, op.getLoc(), rewriter.getType<fallback::TFTensorType>(),
+              operand, op->getAttrOfType<mlir::StringAttr>("device"))
+              .getResult(0));
+    }
+
+    if (!should_rewrite) return failure();
+    rewriter.replaceOp(op, new_values);
+    return success();
+  }
+};
+
+// Removes the following double tensor conversion:
+//  %1 = tfrt_fallback_async.fallback_tensor_to_corert_tensorhandle %0
+//  %2 = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor %1
+struct RemoveDoubleTensorConversion
+    : mlir::OpRewritePattern<CoreRTTensorHandleToFallbackTensorOp> {
+  using OpRewritePattern<
+      CoreRTTensorHandleToFallbackTensorOp>::OpRewritePattern;
+
+  mlir::LogicalResult matchAndRewrite(
+      CoreRTTensorHandleToFallbackTensorOp op,
+      mlir::PatternRewriter &rewriter) const override {
+    // Currently only handles the case where there is only one value in the
+    // conversion op. This should be enough for most of the cases.
+    if (op.getNumOperands() > 1) return mlir::failure();
+
+    auto def =
+        op.getOperand(0).getDefiningOp<FallbackTensorToCoreRTTensorHandleOp>();
+    if (!def) return mlir::failure();
+
+    if (def.getNumResults() > 1) return mlir::failure();
+
+    rewriter.replaceOp(op, def.getOperand(0));
+
+    return mlir::success();
+  }
+};
+
+}  // namespace
+
+void CoreRTTensorHandleToFallbackTensorOp::getCanonicalizationPatterns(
+    RewritePatternSet &results, MLIRContext *context) {
+  results.add<ConstCoreRTTensorHandleToFallbackTensorCanonicalization,
+              RemoveDoubleTensorConversion>(context);
+}
+
+}  // namespace fallback_async
+}  // namespace tfrt
+
+//===----------------------------------------------------------------------===//
+// TableGen'd op method definitions
+//===----------------------------------------------------------------------===//
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_async.cpp.inc"

tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_async.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.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_ASYNC_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_ASYNC_H_
+
+#include "mlir/Bytecode/BytecodeOpInterface.h"  // from @llvm-project
+#include "mlir/IR/Dialect.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/OpImplementation.h"  // from @llvm-project
+#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.
+class FallbackAsyncDialect : public Dialect {
+ public:
+  explicit FallbackAsyncDialect(MLIRContext *context);
+  static StringRef getDialectNamespace() { return "tfrt_fallback_async"; }
+};
+
+}  // namespace fallback_async
+}  // namespace tfrt
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_async.h.inc"
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_ASYNC_H_

tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_async.td

@@ -0,0 +1,531 @@
+/* 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.
+==============================================================================*/
+
+#ifdef TFRT_FALLBACK_ASYNC_OPS
+#else
+#define TFRT_FALLBACK_ASYNC_OPS
+
+include "tfrt/tfrt_op_base.td"
+include "tfrt/compiler/opdefs/tfrt_op_interfaces.td"
+include "tfrt/compiler/opdefs/tfrt_traits.td"
+include "tfrt/core_runtime/opdefs/corert_base.td"
+include "tfrt/core_runtime/opdefs/corert_traits.td"
+include "tfrt_fallback.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+
+def FallbackAsync_Dialect : Dialect {
+  let name = "tfrt_fallback_async";
+
+  let description = [{
+    The TFRT Fallback Async Dialect.
+  }];
+
+  let cppNamespace = "::tfrt::fallback_async";
+}
+
+class FallbackAsync_Op<string mnemonic, list<Trait> traits = []> :
+    Op<FallbackAsync_Dialect, mnemonic, !listconcat(traits, [IsolatedFromAbove])> {
+}
+
+def SetResourceOp : FallbackAsync_Op<"set_resource", [CoreRT_TypedAttributeTrait]> {
+  let summary = "Set a tensor in resource array";
+
+  let description = [{
+    Set a tensor in resource array.
+
+    arg: the tensor to be set in the resource array.
+    index: the index in the resource array
+
+    in_ch and out_ch are for control dependencies.
+  }];
+
+  let arguments = (ins
+    TFRT_ChainType:$in_ch,
+    TFTensorType:$arg,
+    StrAttr:$device,
+    I64Attr:$index
+  );
+
+  let results = (outs
+    TFRT_ChainType:$out_ch
+  );
+
+  let assemblyFormat = "operands attr-dict";
+}
+
+def GetResourceOp : FallbackAsync_Op<"get_resource",
+    [CoreRT_TypedAttributeTrait,
+     TFRT_CostFunctionInterface, TFRT_FixedCost<1>]> {
+  let summary = "get a tensor in resource array";
+
+  let description = [{
+    Get a tensor in resource array.
+
+    indices: the indices in the resource array.
+    results: the tensor values for the corresponding indices.
+
+    in_ch and out_ch are for control dependencies.
+  }];
+
+  let arguments = (ins
+    TFRT_ChainType:$in_ch,
+    StrAttr:$device,
+    I64ArrayAttr:$indices
+  );
+
+  let results = (outs
+    TFRT_ChainType:$out_ch,
+    Variadic<TFTensorType>:$results
+  );
+
+  let assemblyFormat = "operands attr-dict `:` `(` type($in_ch) `)` `->` `(` type($out_ch) `,` type($results) `)`";
+}
+
+def ConstTensorProtoOp : FallbackAsync_Op<"const_tensor_proto",
+    [Pure, CoreRT_TypedAttributeTrait,
+     TFRT_CostFunctionInterface, TFRT_FixedCost<1>]> {
+  let summary = "Create a tensorflow tensor from serialized tensor proto";
+
+  let description = [{
+    Create a tensorflow tensor from serilaized tensor proto.
+
+    Example:
+      %tensor = tfrt_fallback_async.const_tensor_proto "<serialized proto content>"
+  }];
+
+  let arguments = (ins
+    StrAttr:$serialized_tensor_proto
+  );
+
+  let results = (outs
+    TFTensorType:$result
+  );
+
+  let assemblyFormat = "$serialized_tensor_proto attr-dict";
+}
+
+def ConstDenseTensorOp : FallbackAsync_Op<"const_dense_tensor",
+    [ConstantLike, Pure, CoreRT_TypedAttributeTrait,
+     TFRT_CostFunctionInterface, TFRT_FixedCost<1>]> {
+  let summary = "Create a constant dense tensorflow tensor.";
+
+  let description = [{
+    ConstDenseTensorOp creates a constant dense tensorflow tensor.
+
+    Example:
+      %tensor = tfrt_fallback_async.const_dense_tensor dense<[0, 1, 2]> : tensor<3xi32>
+  }];
+
+  let arguments = (ins
+    ElementsAttr:$value
+  );
+
+  let results = (outs
+    TFTensorType:$result
+  );
+
+  let assemblyFormat = "$value attr-dict";
+
+  let hasFolder = 1;
+}
+
+// TODO(b/181257202): This should be `ConstantLike`.
+def ConstStringTensorOp : FallbackAsync_Op<"const_string_tensor",
+    [Pure, CoreRT_TypedAttributeTrait,
+     TFRT_CostFunctionInterface, TFRT_FixedCost<1>]> {
+  let summary = "Create a constant string tensorflow tensor.";
+  let description = [{
+    ConstStringTensorOp creates a constant string tensorflow tensor.
+
+    Example:
+      %tensor = tfrt_fallback_async.const_string_tensor {shape = [1, 2], value = ["const", "string"]}
+
+    Note: if the entries in `value` are all the same, we can just keep one copy.
+      In other words, the following two lines are effectively the same.
+      %tensor = tfrt_fallback_async.const_string_tensor {shape = [1, 2], value = ["string", "string"]}
+      %tensor = tfrt_fallback_async.const_string_tensor {shape = [1, 2], value = ["string"]}
+  }];
+
+  let arguments = (ins
+    I64ArrayAttr:$shape,
+    StrArrayAttr:$value
+  );
+
+  let results = (outs
+    TFTensorType:$result
+  );
+
+  let assemblyFormat = "attr-dict";
+}
+
+def FallbackTensorToCoreRTTensorHandleOp : FallbackAsync_Op<
+                              "fallback_tensor_to_corert_tensorhandle",
+                              [Pure, CoreRT_TypedAttributeTrait,
+                               TFRT_CostFunctionInterface, TFRT_FixedCost<1>]> {
+  let summary = "Convert fallback tensor to corert tensorhandle";
+
+  let description = [{
+    Convert the fallback tensor to corert tensorhandle.
+
+    args: a list of fallback tensors.
+    device: the name of the tensorflow device that is associated with the
+      input fallback tensors.
+
+    results: the converted corert tensorhandles.
+  }];
+
+  let arguments = (ins
+    Variadic<TFTensorType>:$args,
+    StrAttr:$device
+  );
+
+  let results = (outs
+    Variadic<CoreRT_TensorHandleType>:$results
+  );
+
+  let assemblyFormat = "operands attr-dict `:` `(` type($args) `)` `->` `(` type($results) `)`";
+}
+
+def CoreRTTensorHandleToFallbackTensorOp : FallbackAsync_Op<
+                              "corert_tensorhandle_to_fallback_tensor",
+                              [Pure, CoreRT_TypedAttributeTrait,
+                               TFRT_CostFunctionInterface, TFRT_FixedCost<1>]> {
+  let summary = "Convert corert tensorhandle to fallback tensor";
+
+  let description = [{
+    Convert the corert tensorhandle to the fallback tensor.
+
+    args: a list of corert tensorhandles.
+    device: the name of the tensorflow device that is associated with the
+      output fallback tensor.
+
+    results: the converted fallback tensors.
+  }];
+
+  let arguments = (ins
+    Variadic<CoreRT_TensorHandleType>:$args,
+    StrAttr:$device
+  );
+
+  let results = (outs
+    Variadic<TFTensorType>:$results
+  );
+
+  let assemblyFormat = "operands attr-dict `:` `(` type($args) `)` `->` `(` type($results) `)`";
+
+  let hasCanonicalizer = 1;
+}
+
+def PredicateOp : FallbackAsync_Op<"predicate",
+    [Pure, CoreRT_TypedAttributeTrait,
+     TFRT_CostFunctionInterface, TFRT_FixedCost<1>]> {
+  let summary = "Converts a fallback tensor to a bool";
+
+  let description = [{
+    Note: this kernel is used for CPU tensors.
+
+    Converts a fallback tensor to a bool with the following rules:
+
+    - For 0D tensors, truthiness is determined by comparing against a "zero"
+      value. For numerical types it is the obvious zero. For strings it is the
+      empty string.
+
+    - For >0D tensors, truthiness is determined by looking at the number of
+      elements. If has zero elements, then the result is false. Otherwise the
+      result is true.
+
+    input: a fallback tensor representing the condition.
+    device: the name of the tensorflow device that is associated with the
+      input fallback tensor.
+
+    output: the converted bool.
+  }];
+
+  let arguments = (ins
+    TFTensorType:$input
+  );
+
+  let results = (outs
+    I1:$output
+  );
+
+  let assemblyFormat = "$input attr-dict";
+}
+
+def CopyIfSmallOp : FallbackAsync_Op<"copy_if_small",
+    [Pure, CoreRT_TypedAttributeTrait,
+     TFRT_CostFunctionInterface, TFRT_FixedCost<1>]> {
+  let summary = "Copy the CPU fallback tensor if it is small";
+
+  let description = [{
+    This kernel performs deep copy on the input tensor if it is small, to avoid
+    atomic contention on its refcount.
+
+    Note that this kernel always create a new AsyncValue for each result to
+    avoid atomic contention on AsyncValue's refcount.
+  }];
+
+  let arguments = (ins
+    TFTensorType:$operand
+  );
+
+  let results = (outs
+    Variadic<TFTensorType>:$results
+  );
+
+  let assemblyFormat = "$operand attr-dict `:` `(` type($operand) `)` `->` `(` type($results) `)`";
+}
+
+def CreateOp: FallbackAsync_Op<"createop", [CoreRT_TypedAttributeTrait]> {
+  let summary = "The Fallback CreateOp";
+
+  let description = [{
+    The CreateOp creates the tensorflow::OpKernel in the fallback context.
+
+    op_key: The unique key to identify this op.
+    num_args: The number of inputs.
+    device: The tensorflow device. eg. "/CPU:0"
+    op_attrs: The tensorflow attributes excluding the func attrs.
+    op_func_attrs: The func attrs.
+    op_name: The tensorflow op name. eg. "tf.AddV2"
+
+    Example:
+      %out_ch = tfrt_fallback_async.createop(%in_ch) key(0) device("/CPU:0")
+        "some.op"() {attr1 = value, attr2 = value2} {f_attr1 = "fn1", f_attr2 = "fn2"} num_args(1)
+  }];
+
+  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;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def ExecuteOp : FallbackAsync_Op<"executeop",
+    [Pure, CoreRT_TypedAttributeTrait,
+     TFRT_CostFunctionInterface, TFRT_AttrCostTrait]> {
+  let summary = "The Fallback ExecuteOp";
+  let description = [{
+    The ExecuteOp executes an operation on the specified device.
+
+    Example:
+      %res = tfrt_fallback_async.executeop key(0) cost(100) device("/CPU:0") "some.op"(%arg) : 1
+
+    Note that the trailing number indicates the number of results.
+
+    The func attrs and op attrs are placed in two separate regions.
+
+    Example:
+      %res = tfrt_fallback_async.executeop key(0) device("/CPU:0") "some.op"(%arg)
+          {attr1 = value, attr2 = value2} {f_attr1 = "fn1", f_attr2 = "fn2"} : 1
+
+    Note that in this case, the first bracket contains op attributes other than
+    function attributes, and the second braket contains function attributes,
+    which represent function names.
+
+    Refer to b/172597167 for approaches considered for supporting function
+    attributes.
+  }];
+
+  let arguments = (ins
+    Variadic<TFTensorType>:$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,
+    I64Attr:$_tfrt_cost
+  );
+
+  let results = (outs
+    Variadic<TFTensorType>:$results
+  );
+
+  let extraClassDeclaration = [{
+    void getOpAttrs(SmallVectorImpl<std::pair<::llvm::StringRef, ::mlir::Attribute>>* op_attrs);
+  }];
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def ExecuteOpSeq : FallbackAsync_Op<"executeop.seq",
+    [CoreRT_TypedAttributeTrait,
+     TFRT_CostFunctionInterface, TFRT_AttrCostTrait]> {
+  let summary = "The sequenced version of Fallback ExecuteOp";
+  let description = [{
+    The ExecuteOpSeq executes an operation on the specified device, similar to
+    ExecuteOp. In addition, it takes an in_op_chain and produces an out_op_chain
+    that can be used to sequence op executions. ExecuteOpSeq is nonstrict on
+    in_op_chain.
+
+    Example:
+      %op_ch_out, %res = tfrt_fallback_async.executeop.seq(%op_ch_in)
+        key(0) cost(100) device("/CPU:0") "some.op"(%arg) : 1
+
+    Note that the trailing number indicates the number of results.
+
+    The func attrs and op attrs are placed in two separate regions.
+
+    Example:
+      %op_ch_out, %res = tfrt_fallback_async.executeop.seq(%op_ch_in)
+        key(0) cost(100) device("/CPU:0") "some.op"(%arg)
+        {fn = function_1} {attr1 = value, attr2 = value2} : 1
+
+    Note that in this case, the first braket contains op attributes other than
+    function attributes, and the second braket contains function attributes,
+    which represent function names.
+
+    Refer to b/172597167 for approaches considered for supporting function
+    attributes.
+  }];
+
+  let arguments = (ins
+    TFRT_ChainType:$in_op_chain,
+    Variadic<TFTensorType>:$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,
+    I64Attr:$_tfrt_cost
+  );
+
+  let results = (outs
+    TFRT_ChainType:$out_op_chain,
+    Variadic<TFTensorType>:$results
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def BatchFunctionOp : FallbackAsync_Op<"batch_function",
+    [Pure, CoreRT_TypedAttributeTrait]> {
+  let summary = "Fallback ExecuteOp specialized for tf.BatchFunction.";
+  let description = [{
+    This kernel executes a variant tf.BatchFunction kernel that supports having
+    the `f` attribute as a BEF function.
+
+    Example:
+      %res = tfrt_fallback_async.batch_function device("/CPU:0") @f(%input, %captured_input)  {
+          num_batch_threads = 1,
+          max_batch_size = 4,
+          allowed_batch_sizes = [2 , 4],
+          batch_timeout_micros = 50,
+          container = "container",
+          shared_name = "shared_name",
+          batching_queue = "batching_queue",
+          enable_large_batch_splitting = false,
+          Tin = [f32],
+          Tcaptured = [f32],
+          Tout = [f32]} : 1
+
+    Note that the trailing number indicates the number of results.
+  }];
+
+  let arguments = (ins
+    Variadic<TFTensorType>:$args,
+    StrAttr:$device,
+    SymbolRefAttr:$f,
+    ArrayAttr:$op_attrs
+  );
+
+  let results = (outs
+    Variadic<TFTensorType>:$results
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def ExecuteOpWithAllocator : FallbackAsync_Op<"executeop.allocator",
+    [Pure, CoreRT_TypedAttributeTrait, TFRT_CostFunctionInterface, TFRT_AttrCostTrait]> {
+  let summary = "The Fallback ExecuteOp with custom allocator";
+  let description = [{
+    Similar to ExecuteOp but takes a custom allocator for allocating output tensors.
+  }];
+
+  let arguments = (ins
+    TFAllocatorType:$allocator,
+    Variadic<TFTensorType>:$args,
+    StrAttr:$device,
+    ArrayAttr:$op_attrs,
+    ArrayAttr:$op_func_attrs,
+    I64Attr:$op_key,
+    StrAttr:$op_name,
+    I64Attr:$_tfrt_cost
+  );
+
+  let results = (outs
+    Variadic<TFTensorType>:$results
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def ExecuteOpSeqWithAllocator : FallbackAsync_Op<"executeop.seq.allocator",
+    [CoreRT_TypedAttributeTrait, TFRT_CostFunctionInterface, TFRT_AttrCostTrait]> {
+  let summary = "The sequenced version of Fallback ExecuteOp with custom allocator";
+  let description = [{
+    Similar to ExecuteOpSeq but takes a custom allocator for allocating output tensors.
+
+    Example:
+      %op_ch_out, %res = tfrt_fallback_async.executeop.seq.allocator(%op_ch_in, %allocator)
+        key(0) cost(100) device("/CPU:0") "some.op"(%arg) : 1
+  }];
+
+  let arguments = (ins
+    TFRT_ChainType:$in_op_chain,
+    TFAllocatorType:$allocator,
+    Variadic<TFTensorType>:$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,
+    I64Attr:$_tfrt_cost
+  );
+
+  let results = (outs
+    TFRT_ChainType:$out_op_chain,
+    Variadic<TFTensorType>:$results
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+#endif

tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_common.cc

@@ -0,0 +1,145 @@
+/* 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/ir/tfrt_fallback_common.h"
+
+#include <cassert>
+#include <cstdint>
+#include <utility>
+
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace tfrt {
+namespace fallback_common {
+
+void GetExecuteOpAttrsCommon(
+    mlir::MLIRContext *context, llvm::ArrayRef<mlir::Attribute> op_attr_array,
+    llvm::SmallVectorImpl<std::pair<llvm::StringRef, mlir::Attribute>>
+        *op_attrs) {
+  assert(op_attrs);
+  op_attrs->clear();
+
+  mlir::Builder builder(context);
+  for (auto iter : op_attr_array) {
+    auto key_value = mlir::cast<mlir::ArrayAttr>(iter).getValue();
+    llvm::StringRef key = mlir::cast<mlir::StringAttr>(key_value[0]).getValue();
+    mlir::Attribute value = key_value[1];
+    op_attrs->push_back({key, value});
+  }
+}
+
+mlir::ParseResult ParseExecuteOpCommon(mlir::OpAsmParser &parser,
+                                       mlir::Builder &builder,
+                                       mlir::OperationState &result,
+                                       mlir::Type tensor_type,
+                                       const ParseExecuteOpOptions &options) {
+  auto chain_type = builder.getType<compiler::ChainType>();
+
+  mlir::IntegerAttr op_key;
+  mlir::IntegerAttr cost;
+  mlir::StringAttr device;
+  mlir::StringAttr op_name;
+  mlir::SymbolRefAttr f;
+  llvm::SmallVector<mlir::OpAsmParser::UnresolvedOperand, 4> in_chains;
+  llvm::SmallVector<mlir::OpAsmParser::UnresolvedOperand, 4> operands;
+  mlir::NamedAttrList op_attrs;
+  mlir::NamedAttrList op_func_attrs;
+  auto loc = parser.getNameLoc();
+
+  if (options.has_chain &&
+      parser.parseOperandList(in_chains,
+                              /*requiredOperandCount=*/1,
+                              mlir::OpAsmParser::Delimiter::Paren))
+    return mlir::failure();
+
+  if (options.has_key &&
+      (parser.parseKeyword("key") || parser.parseLParen() ||
+       parser.parseAttribute(op_key, "op_key", result.attributes) ||
+       parser.parseRParen()))
+    return mlir::failure();
+
+  if (options.has_cost &&
+      (parser.parseKeyword("cost") || parser.parseLParen() ||
+       parser.parseAttribute(cost, "_tfrt_cost", result.attributes) ||
+       parser.parseRParen()))
+    return mlir::failure();
+
+  if (options.has_device &&
+      (parser.parseKeyword("device") || parser.parseLParen() ||
+       parser.parseAttribute(device, "device", result.attributes) ||
+       parser.parseRParen()))
+    return mlir::failure();
+
+  if (options.has_op_name &&
+      parser.parseAttribute(op_name, "op_name", result.attributes))
+    return mlir::failure();
+
+  if (options.has_symbol_ref &&
+      parser.parseAttribute(f, "f", result.attributes))
+    return mlir::failure();
+
+  if (parser.parseOperandList(operands, mlir::OpAsmParser::Delimiter::Paren) ||
+      parser.parseOptionalAttrDict(op_attrs) ||
+      parser.parseOptionalAttrDict(op_func_attrs))
+    return mlir::failure();
+
+  int64_t num_results = 0;
+  if (succeeded(parser.parseOptionalColon())) {
+    mlir::IntegerAttr attr;
+    mlir::NamedAttrList attrs;
+    if (failed(parser.parseAttribute(attr, "num_results", attrs)))
+      return mlir::failure();
+    num_results = attr.getValue().getSExtValue();
+  }
+  if (num_results < 0) return mlir::failure();
+
+  llvm::SmallVector<mlir::Type, 4> operand_types;
+  if (options.has_chain) operand_types.push_back(chain_type);
+  if (parser.resolveOperands(in_chains, operand_types, loc, result.operands) ||
+      parser.resolveOperands(operands, tensor_type, result.operands))
+    return mlir::failure();
+
+  if (options.has_chain) result.types.push_back(chain_type);
+  result.types.append(num_results, tensor_type);
+
+  llvm::SmallVector<mlir::Attribute, 4> op_attr_array;
+  for (const auto &key_value : op_attrs) {
+    auto key = key_value.getName();
+    auto value = key_value.getValue();
+    op_attr_array.push_back(builder.getArrayAttr({key, value}));
+  }
+
+  result.attributes.push_back(
+      builder.getNamedAttr("op_attrs", builder.getArrayAttr(op_attr_array)));
+
+  // TODO(tfrt-devs): support func attributes in tfrt_fallback_sync.
+  if (options.has_func_attr) {
+    llvm::SmallVector<mlir::Attribute, 4> op_func_attr_array;
+    for (const auto &key_value : op_func_attrs) {
+      auto key = key_value.getName();
+      auto value = key_value.getValue();
+      op_func_attr_array.push_back(builder.getArrayAttr({key, value}));
+    }
+
+    result.attributes.push_back(builder.getNamedAttr(
+        "op_func_attrs", builder.getArrayAttr(op_func_attr_array)));
+  }
+
+  return mlir::success();
+}
+
+}  // namespace fallback_common
+}  // namespace tfrt

tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_common.h

@@ -0,0 +1,127 @@
+/* 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.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_COMMON_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_COMMON_H_
+
+#include <utility>
+
+#include "llvm/ADT/STLExtras.h"
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/OpImplementation.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tfrt/basic_kernels/opdefs/types.h"  // from @tf_runtime
+
+namespace tfrt {
+namespace fallback_common {
+
+template <typename OpTy>
+mlir::LogicalResult VerifyExecuteOpCommon(OpTy op) {
+  auto op_attr_array = op.getOpAttrs().getValue();
+  for (auto op_attr : op_attr_array) {
+    auto key_value = mlir::dyn_cast<mlir::ArrayAttr>(op_attr);
+    if (!key_value || key_value.getValue().size() != 2 ||
+        !mlir::isa<mlir::StringAttr>(key_value.getValue()[0]))
+      return op.emitOpError() << "each op_attr should be a key-value pair, "
+                                 "where the key is a string";
+  }
+  return mlir::success();
+}
+
+template <typename OpTy>
+mlir::LogicalResult VerifyFallbackExecuteOp(OpTy op) {
+  auto result = VerifyExecuteOpCommon(op);
+  if (failed(result)) return result;
+
+  // Verify function attributes.
+  auto op_func_attr_array = op.getOpFuncAttrs().getValue();
+  for (auto op_attr : op_func_attr_array) {
+    auto key_value = mlir::dyn_cast<mlir::ArrayAttr>(op_attr);
+    if (!key_value || key_value.getValue().size() != 2 ||
+        !mlir::isa<mlir::StringAttr>(key_value.getValue()[0]) ||
+        !mlir::isa<mlir::StringAttr>(key_value.getValue()[1]))
+      return op.emitOpError() << "each op_func_attr should be a key-value "
+                                 "pair, where both the key and the value are "
+                                 "strings";
+  }
+  return mlir::success();
+}
+
+template <typename OpTy>
+void PrintExecuteOpFuncAttribute(mlir::OpAsmPrinter &p, OpTy op) {
+  auto op_func_attrs = op.getOpFuncAttrs();
+  if (!op_func_attrs.empty()) {
+    auto print_key_value = [&](mlir::Attribute attr) {
+      auto key_value = mlir::cast<mlir::ArrayAttr>(attr).getValue();
+      auto key = key_value[0];
+      auto value = key_value[1];
+
+      p << mlir::cast<mlir::StringAttr>(key).getValue();
+      p << " = ";
+      p << value;
+    };
+
+    auto op_func_attr_array = op_func_attrs.getValue();
+    p << " {";
+    llvm::interleaveComma(op_func_attr_array, p, print_key_value);
+    p << '}';
+  }
+}
+
+template <typename OpTy>
+void PrintExecuteOpCommon(mlir::OpAsmPrinter &p, OpTy op) {
+  auto op_attrs = op.getOpAttrs();
+  if (!op_attrs.empty()) {
+    auto print_key_value = [&](mlir::Attribute attr) {
+      auto key_value = mlir::cast<mlir::ArrayAttr>(attr).getValue();
+      auto key = key_value[0];
+      auto value = key_value[1];
+
+      p << mlir::cast<mlir::StringAttr>(key).getValue();
+      p << " = ";
+      p << value;
+    };
+
+    auto op_attr_array = op_attrs.getValue();
+    p << " {";
+    llvm::interleaveComma(op_attr_array, p, print_key_value);
+    p << '}';
+  }
+}
+
+void GetExecuteOpAttrsCommon(
+    mlir::MLIRContext *context, llvm::ArrayRef<mlir::Attribute> op_attr_array,
+    llvm::SmallVectorImpl<std::pair<llvm::StringRef, mlir::Attribute>>
+        *op_attrs);
+
+struct ParseExecuteOpOptions {
+  bool has_chain = false;
+  bool has_key = false;
+  bool has_device = false;
+  bool has_func_attr = false;
+  bool has_cost = false;
+  bool has_op_name = true;
+  bool has_symbol_ref = false;
+};
+
+mlir::ParseResult ParseExecuteOpCommon(mlir::OpAsmParser &parser,
+                                       mlir::Builder &builder,
+                                       mlir::OperationState &result,
+                                       mlir::Type tensor_type,
+                                       const ParseExecuteOpOptions &options);
+}  // namespace fallback_common
+}  // namespace tfrt
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_COMMON_H_

tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.cc

@@ -0,0 +1,65 @@
+/* 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/ir/tfrt_fallback_sync.h"
+
+#include "llvm/ADT/STLExtras.h"
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/DialectImplementation.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/OpImplementation.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_common.h"
+#include "tfrt/basic_kernels/opdefs/tfrt_base.h"  // from @tf_runtime
+#include "tfrt/basic_kernels/opdefs/types.h"  // from @tf_runtime
+#include "tfrt/core_runtime/opdefs/attributes.h"  // from @tf_runtime
+#include "tfrt/core_runtime/opdefs/core_runtime.h"  // from @tf_runtime
+#include "tfrt/core_runtime/opdefs/types.h"  // from @tf_runtime
+#include "tfrt/tensor/opdefs/tensor.h"  // from @tf_runtime
+
+namespace tfrt {
+namespace fallback_sync {
+
+FallbackSyncDialect::FallbackSyncDialect(MLIRContext *context)
+    : Dialect(/*name=*/"tfrt_fallback_sync", context,
+              TypeID::get<FallbackSyncDialect>()) {
+  context->getOrLoadDialect<tfrt::fallback::FallbackDialect>();
+  context->getOrLoadDialect<compiler::TFRTDialect>();
+  context->getOrLoadDialect<corert::CoreRTDialect>();
+
+  allowUnknownTypes();
+
+  addOperations<
+#define GET_OP_LIST
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.cpp.inc"
+      >();
+}
+
+static Type GetTensorType(Builder *builder) {
+  return tfrt::tfrt_tensor::TensorType::get(builder->getContext());
+}
+
+}  // namespace fallback_sync
+}  // namespace tfrt
+
+//===----------------------------------------------------------------------===//
+// TableGen'd op method definitions
+//===----------------------------------------------------------------------===//
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.cpp.inc"

tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.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.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_SYNC_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_SYNC_H_
+
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Dialect.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/OpImplementation.h"  // from @llvm-project
+#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:
+  explicit FallbackSyncDialect(MLIRContext *context);
+  static StringRef getDialectNamespace() { return "tfrt_fallback_sync"; }
+};
+
+}  // namespace fallback_sync
+}  // namespace tfrt
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.h.inc"
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_SYNC_H_

tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.td

@@ -0,0 +1,286 @@
+/* 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.
+==============================================================================*/
+
+#ifdef TFRT_FALLBACK_SYNC_OPS
+#else
+#define TFRT_FALLBACK_SYNC_OPS
+
+include "tfrt/tfrt_op_base.td"
+include "tfrt/core_runtime/opdefs/corert_base.td"
+include "tfrt/core_runtime/opdefs/corert_traits.td"
+include "tfrt_fallback.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td"
+
+def FallbackSync_Dialect : Dialect {
+  let name = "tfrt_fallback_sync";
+
+  let description = [{
+    The TFRT Fallback Sync Dialect.
+  }];
+
+  let cppNamespace = "::tfrt::fallback_sync";
+}
+
+class FallbackSync_Op<string mnemonic, list<Trait> traits = []> :
+    Op<FallbackSync_Dialect, mnemonic, !listconcat(traits, [IsolatedFromAbove])> {
+}
+
+// TODO(rohitju): Add a verifier if there is anything that should be verified
+// compile time.
+def SetResourceOp : FallbackSync_Op<"set_resource", [CoreRT_TypedAttributeTrait]> {
+  let summary = "Set a tensor in resource array";
+
+  let description = [{
+    Set a tensor in resource array.
+
+    arg: the tensor to be set in the resource array.
+    index: the index in the resource array
+  }];
+
+  let arguments = (ins
+    TFTensorType:$arg,
+    I64Attr:$index
+  );
+
+  let results = (outs);
+
+  let assemblyFormat = "operands attr-dict";
+}
+
+def SetResourceDhtOp : FallbackSync_Op<"set_resource_dht", [CoreRT_TypedAttributeTrait]> {
+  let summary = "Set a DHT in resource array";
+
+  let description = [{
+    Set a DHT in resource array.
+
+    arg: the tensor to be set in the resource array.
+    index: the index in the resource array
+  }];
+
+  let arguments = (ins
+    TensorType:$arg,
+    I64Attr:$index
+  );
+
+  let results = (outs);
+
+  let assemblyFormat = "operands attr-dict";
+}
+
+def GetResourceOp : FallbackSync_Op<"get_resource",
+    [CoreRT_TypedAttributeTrait]> {
+  let summary = "get a tensor in resource array";
+
+  let description = [{
+    Get a tensor in resource array.
+
+    indices: the indices in the resource array.
+    results: the tensor values for the corresponding indices.
+  }];
+
+  let arguments = (ins
+    I64ArrayAttr:$indices
+  );
+
+  let results = (outs
+    Variadic<TFTensorType>:$results
+  );
+
+  let assemblyFormat = "attr-dict `:` type($results)";
+}
+
+def GetResourceDhtOp : FallbackSync_Op<"get_resource_dht",
+    [CoreRT_TypedAttributeTrait]> {
+  let summary = "get a DHT in resource array";
+
+  let description = [{
+    Get a tensor in resource array.
+
+    indices: the indices in the resource array.
+    results: the tensor values for the corresponding indices.
+  }];
+
+  let arguments = (ins
+    I64ArrayAttr:$indices
+  );
+
+  let results = (outs
+    Variadic<TensorType>:$results
+  );
+
+  let assemblyFormat = "attr-dict `:` type($results)";
+}
+
+def CreateOp: FallbackSync_Op<"createop", []> {
+  let summary = "The Fallback CreateOp";
+
+  let description = [{
+    The CreateOp creates the tensorflow::OpKernel in the fallback context.
+
+    op_key: The unique key to identify this op.
+    num_args: The number of inputs.
+    device: The tensorflow device. eg. "/CPU:0"
+    op_attrs: The tensorflow attributes excluding the func attrs.
+    op_name: The tensorflow op name. eg. "tf.AddV2"
+
+    Example:
+      tfrt_fallback_sync.createop() key(0) device("/CPU:0")
+        "some.op"() {attr1 = value, attr2 = value2} {f_attr1 = "fn1", f_attr2 = "fn2"} num_args(1)
+  }];
+
+  let arguments = (ins
+    StrAttr:$node_def,
+    I32Attr:$op_key
+  );
+
+  let results = (outs);
+
+  let assemblyFormat = "attr-dict";
+}
+
+def SyncExecuteOp : FallbackSync_Op<"executeop", [Pure]> {
+  let summary = "The Fallback Sync ExecuteOp";
+  let description = [{
+    The ExecuteOp executes an operation on the specified device.
+
+    Example:
+      %res = tfrt_fallback_sync.executeop "some.op"(%arg) : 1
+
+    Note that the trailing number indicates the number of results.
+  }];
+
+  let arguments = (ins
+    Variadic<TFTensorType>,
+    StrAttr:$node_def,
+    I32Attr:$op_key
+  );
+
+  let results = (outs Variadic<TFTensorType>);
+
+  let assemblyFormat = "`(`operands`)` attr-dict `:` functional-type(operands, results)";
+}
+
+
+def ConvertDhtToFallbackTensorOp : FallbackSync_Op<"convert_dht_to_fallback_tensor",
+    [CoreRT_TypedAttributeTrait]> {
+  let summary = "convert a DHT to a fallback tensor";
+
+  let description = [{
+    Convert a TFRT DHT to a fallback tensor.
+
+    dht: the input DHT.
+    fallback_tensor: The output fallback tensor
+  }];
+
+  let arguments = (ins
+    TensorType:$dht
+  );
+
+  let results = (outs
+    TFTensorType:$fallback_tensor
+  );
+
+  let assemblyFormat = "operands attr-dict `:` type($fallback_tensor)";
+}
+
+def ConvertShtToFallbackTensorOp : FallbackSync_Op<"convert_sht_to_fallback_tensor",
+    [CoreRT_TypedAttributeTrait]> {
+  let summary = "convert a SHT to a fallback tensor";
+
+  let description = [{
+    Convert a TFRT SHT to a fallback tensor.
+
+    sht: the input SHT.
+    fallback_tensor: The output fallback tensor
+  }];
+
+  let arguments = (ins
+    TensorType:$sht
+  );
+
+  let results = (outs
+    TFTensorType:$fallback_tensor
+  );
+
+  let assemblyFormat = "operands attr-dict `:` type($fallback_tensor)";
+}
+
+def ConvertFallbackTensorToDhtOp : FallbackSync_Op<"convert_fallback_tensor_to_dht",
+    [CoreRT_TypedAttributeTrait]> {
+  let summary = "convert fallback tensor to a DHT";
+
+  let description = [{
+    Convert a fallback tensor to a TFRT DHT.
+
+    fallback_tensor: The input fallback tensor
+    dht: the output DHT.
+  }];
+
+  let arguments = (ins
+    TFTensorType:$fallback_tensor
+  );
+
+  let results = (outs
+    TensorType:$dht
+  );
+
+  let assemblyFormat = "operands attr-dict `:` type($dht)";
+}
+
+// TODO(rohitju): This is Ads specific, need to find an appropriate home for it.
+def SetSparseMatrixResourceOp : FallbackSync_Op<"set_sparse_matrix_resource", [CoreRT_TypedAttributeTrait]> {
+  let summary = "Set a Sparse matrix in resource array";
+
+  let description = [{
+    Set a sparse matrix in resource array.
+
+    arg: the matrix to be set in the resource array.
+    index: the index in the resource array
+  }];
+
+  let arguments = (ins
+    TFTensorType:$arg,
+    I64Attr:$index
+  );
+
+  let results = (outs);
+
+  let assemblyFormat = "operands attr-dict";
+}
+
+def GetSparseMatrixResourceOp : FallbackSync_Op<"get_sparse_matrix_resource",
+    [CoreRT_TypedAttributeTrait]> {
+  let summary = "get a sparse matrix from resource array";
+
+  let description = [{
+    Get a sparse matrix from resource array.
+
+    indices: the indices in the resource array.
+    results: the tensor values for the corresponding indices.
+  }];
+
+  let arguments = (ins
+    I64ArrayAttr:$indices
+  );
+
+  let results = (outs
+    Variadic<TFTensorType>:$results
+  );
+
+  let assemblyFormat = "attr-dict `:` type($results)";
+}
+
+#endif

tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_util.cc

@@ -0,0 +1,56 @@
+/* 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/ir/tfrt_fallback_util.h"
+
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_async.h"
+
+namespace tfrt {
+namespace fallback_async {
+
+bool IsArgConsumedByFallback(mlir::func::FuncOp func, int arg_index) {
+  auto arg = func.getArgument(arg_index);
+
+  // Return true if any user is a fallback op. It is more interesting to know
+  // whether it is consumed by any fallback op than whether it is only consumed
+  // by fallback ops. For example, the conversion from a DenseHostTensor to a
+  // fallback tensor is more costly than the conversion from fallback tensor to
+  // a DenseHostTensor. So as long as one of the users of a captured variable is
+  // a fallback op, we should keep this variable as a fallback tensor.
+  for (mlir::Operation *user : arg.getUsers()) {
+    if (llvm::isa<FallbackAsyncDialect>(user->getDialect())) return true;
+  }
+  return false;
+}
+
+void ForEachArgConsumedByFallback(
+    mlir::func::FuncOp func, llvm::function_ref<void(int arg_index)> action) {
+  for (int arg_index = 0; arg_index < func.getNumArguments(); ++arg_index) {
+    if (IsArgConsumedByFallback(func, arg_index)) action(arg_index);
+  }
+}
+
+void ForEachArgConsumedByFallback(
+    mlir::ModuleOp module,
+    llvm::function_ref<void(llvm::StringRef func_name, int arg_index)> action) {
+  for (auto func : module.getOps<mlir::func::FuncOp>()) {
+    ForEachArgConsumedByFallback(
+        func, [func_name = func.getName(), action](int arg_index) {
+          action(func_name, arg_index);
+        });
+  }
+}
+
+}  // namespace fallback_async
+}  // namespace tfrt

tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_util.h

@@ -0,0 +1,36 @@
+/* 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.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_UTIL_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_UTIL_H_
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+
+namespace tfrt {
+namespace fallback_async {
+
+bool IsArgConsumedByFallback(mlir::func::FuncOp func, int arg_index);
+
+void ForEachArgConsumedByFallback(
+    mlir::func::FuncOp func, llvm::function_ref<void(int arg_index)> action);
+
+void ForEachArgConsumedByFallback(
+    mlir::ModuleOp module,
+    llvm::function_ref<void(llvm::StringRef func_name, int arg_index)> action);
+
+}  // namespace fallback_async
+}  // namespace tfrt
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_UTIL_H_