(!mlrt.promise, !mlrt.promise, !mlrt.promise,
     !mlrt.future, !mlrt.future, !mlrt.future)

  // CHECK: [[handle_0:%.*]] = mlrt.async([[input0]], [[promises]]#0)
  // CHECK-SAME: callee = @main_stream_0
  %handle_0 = mlrt.async(%input0, %promise_b)
    {callee = @main_stream_0} :
    (tensor<i32>, !mlrt.promise) -> !mlrt.async_handle

def Mlrt_Dialect : Dialect {
  let name = "mlrt";

  let description = [{
    The MLRT Dialect.
  }];

  let cppNamespace = "::mlrt::compiler";
}

def MlrtFutureType : DialectType<Mlrt_Dialect,
    CPred<"$_self.isa<::mlrt::compiler::FutureType>()">, "!mlrt.future type">,
    BuildableType<"$_builder.getType<::mlrt::compiler::FutureType>()"> {
  let description = [{
    `!mlrt.future type` represents a C++ mlrt::Future.
  }];

};

// 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:

// CHECK: mlrt.async([[FUTURES]]#0) {callee = @main_stream_1
// CHECK: mlrt.await_handle
// CHECK: mlrt.await_handle
// CHECK: mlrt.await_handle
// CHECK: mlrt.await_handle

func.func @main() {
  "tf.DummySideEffecting"() {id = 1} : () -> ()
  "tf.DummySideEffecting"() {id = 2} : () -> ()
  "tf.DummySideEffecting"() {id = 3} : () -> ()

// CHECK-LABEL: @main
// CHECK-SAME: ([[f0:%.*]]: !mlrt.async_handle, [[f1:%.*]]: !mlrt.async_handle, [[f2:%.*]]: !mlrt.async_handle)
func.func @main(%f0: !mlrt.async_handle, %f1: !mlrt.async_handle, %f2: !mlrt.async_handle) -> () {
  // CHECK-NEXT: mlrt.await_all_handle [[f0]], [[f1]], [[f2]]
  // CHECK-NOT: mlrt.await_handle
  // CHECK-NEXT: return
  mlrt.await_handle %f0
  mlrt.await_handle %f1
  mlrt.await_handle %f2
  func.return
}

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

#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

  %1:9 = "tf.While"(%cst_0, %max_iterations, %array_handle, %array_flow, %matrix , %array_handle, %array_flow, %matrix, %bound) {body= @"sort_map/while_body", cond = @"sort_map/while_cond",...

  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

  // Predicate should be inlined.
  // CHECK-NEXT: tf_mlrt.predicate
  // CHECK-NEXT: mlrt.cond
  // CHECK-NEXT: tf_mlrt.predicate

  // CHECK-NEXT: mlrt.while