chore: import upstream snapshot with attribution

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