chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/tfrt/translate/mlrt/BUILD

@@ -0,0 +1,94 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load(
+    "//tensorflow:tensorflow.bzl",
+    "tf_cc_test",
+)
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        # copybara:uncomment "//learning/brain/experimental/tfrt:__subpackages__",
+        # copybara:uncomment "//smartass/brain/ops/tfrt_kernels:__subpackages__",
+        "//tensorflow/compiler/mlir/tfrt/transforms/mlrt:__subpackages__",
+        "//tensorflow/core/tfrt:__subpackages__",
+    ],
+)
+
+cc_library(
+    name = "mlir_to_bytecode",
+    srcs = ["mlir_to_bytecode.cc"],
+    hdrs = ["mlir_to_bytecode.h"],
+    deps = [
+        "//tensorflow/compiler/mlir/tfrt/ir/mlrt:mlrt_ops",
+        "//tensorflow/core/tfrt/mlrt/bytecode",
+        "//tensorflow/core/tfrt/mlrt/bytecode:executable",
+        "//tensorflow/core/tfrt/mlrt/bytecode:function",
+        "//tensorflow/core/tfrt/mlrt/bytecode:kernel",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+tf_cc_test(
+    name = "mlir_to_bytecode_test",
+    srcs = ["mlir_to_bytecode_test.cc"],
+    data = glob(["testdata/**"]),
+    deps = [
+        ":mlir_to_bytecode",
+        "//tensorflow/compiler/mlir/tfrt/ir/mlrt:mlrt_ops",
+        "//tensorflow/core/tfrt/mlrt/bytecode",
+        "//tensorflow/core/tfrt/mlrt/bytecode:executable",
+        "//tensorflow/core/tfrt/mlrt/interpreter:attribute_span",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Parser",
+        "@llvm-project//mlir:Support",
+        "@xla//xla/tsl/platform:resource_loader",
+    ],
+)
+
+cc_library(
+    name = "test_utils",
+    testonly = 1,
+    srcs = ["test_utils.cc"],
+    hdrs = ["test_utils.h"],
+    deps = [
+        # copybara:uncomment "//learning/brain/experimental/tfrt/native_lowering/stubs:tfrt_native_lowering_impl",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_absl//absl/types:span",
+        "@xla//xla/tsl/platform:errors",
+        "@xla//xla/tsl/platform:statusor",
+        "//tensorflow/core:core_cpu_base",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/tfrt/graph_executor:sync_resource_state",
+        "//tensorflow/core/tfrt/mlrt/attribute",
+        "//tensorflow/core/tfrt/mlrt/bytecode",
+        "//tensorflow/core/tfrt/mlrt/bytecode:kernel",
+        "//tensorflow/core/tfrt/mlrt/interpreter:context",
+        "//tensorflow/core/tfrt/mlrt/interpreter:interpreter_testutil",
+        "//tensorflow/core/tfrt/mlrt/interpreter:value",
+        "//tensorflow/core/tfrt/stubs:tfrt_native_lowering_stub",
+        "//tensorflow/core/tfrt/utils:tensor_util",
+        "@tf_runtime//:hostcontext",
+        "@tf_runtime//:support",
+        "@tf_runtime//:tensor",
+    ],
+)

tensorflow/compiler/mlir/tfrt/translate/mlrt/mlir_to_bytecode.cc

@@ -0,0 +1,526 @@
+/* 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/translate/mlrt/mlir_to_bytecode.h"
+
+#include <cstdint>
+#include <cstring>
+#include <iterator>
+#include <optional>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Region.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.h"
+#include "tensorflow/core/tfrt/mlrt/bytecode/bytecode.h"
+#include "tensorflow/core/tfrt/mlrt/bytecode/executable.h"
+#include "tensorflow/core/tfrt/mlrt/bytecode/function.h"
+#include "tensorflow/core/tfrt/mlrt/bytecode/kernel.h"
+
+namespace mlrt {
+namespace {
+
+// LINT.IfChange(mlrt_attributes)
+bool CanBeInlined(mlir::Attribute attr, absl::string_view data) {
+  // FlatSymbolRefAttr is a special case as we are emitting it as integer.
+  return mlir::isa<mlir::IntegerAttr, mlir::FloatAttr, mlir::FlatSymbolRefAttr>(
+             attr) &&
+         data.size() <= sizeof(uint32_t);
+}
+// LINT.ThenChange(../../../../../core/tfrt/mlrt/interpreter/attribute_span.h:mlrt_attributes)
+
+// Encode integer or float-point numbers as bytes.
+template <typename T>
+std::string EncodeIntegerOrFloat(T attr) {
+  std::string data(sizeof(attr), '\0');
+  std::memcpy(data.data(), &attr, sizeof(attr));
+  return data;
+}
+
+// Encode a list of I64 integers as bytes using bc::Vector<uint64_t>. The bytes
+// can be decoded directly using bc::Vector<uint64_t>. If `array` is not a list
+// I64 integers, a nullopt will be returned.
+
+template <typename T>
+std::optional<std::string> EncodeListOfInteger(mlir::ArrayAttr array) {
+  bc::Buffer buffer;
+  bc::Allocator allocator(&buffer);
+  auto ctor = bc::New<bc::Vector<T>>(&allocator, array.size());
+
+  mlir::Type type;
+
+  for (int i = 0; i < array.size(); ++i) {
+    if (auto integer_attr = mlir::dyn_cast<mlir::IntegerAttr>(array[i])) {
+      if (type && integer_attr.getType() != type) return std::nullopt;
+      type = integer_attr.getType();
+      llvm::APInt value = integer_attr.getValue();
+      if (value.getBitWidth() != sizeof(T) * 8) return std::nullopt;
+      ctor.ConstructAt(i, value.getZExtValue());
+    } else {
+      return std::nullopt;
+    }
+  }
+
+  return std::string(buffer.data(), buffer.size());
+}
+
+std::optional<std::string> EncodeListOfSymbolRef(
+    const ModuleEmitterContext& module_context, mlir::ArrayAttr array) {
+  bc::Buffer buffer;
+  bc::Allocator allocator(&buffer);
+  auto ctor = bc::New<bc::Vector<uint32_t>>(&allocator, array.size());
+
+  for (int i = 0; i < array.size(); ++i) {
+    if (auto symbol_ref = mlir::dyn_cast<mlir::FlatSymbolRefAttr>(array[i])) {
+      ctor.ConstructAt(i, module_context.GetFunctionId(symbol_ref.getValue()));
+    } else {
+      return std::nullopt;
+    }
+  }
+  return std::string(buffer.data(), buffer.size());
+}
+
+template <typename T>
+std::optional<std::string> EncodeDenseArray(llvm::ArrayRef<T> array) {
+  bc::Buffer buffer;
+  bc::Allocator allocator(&buffer);
+  auto ctor = bc::New<bc::Vector<T>>(&allocator, array.size());
+
+  if (!array.empty()) {
+    ctor.Place(reinterpret_cast<const char*>(array.data()),
+               array.size() * sizeof(T));
+  }
+
+  return std::string(buffer.data(), buffer.size());
+}
+
+// bool values has special encoding in MLIR. It occupies one bit in MLIR
+// but in bytecode it is one byte.
+std::optional<std::string> EncodeDenseBoolArray(llvm::ArrayRef<bool> array) {
+  bc::Buffer buffer;
+  bc::Allocator allocator(&buffer);
+  auto ctor = bc::New<bc::Vector<uint8_t>>(&allocator, array.size());
+
+  if (!array.empty()) {
+    std::vector<uint8_t> data(array.size());
+    int i = 0;
+    for (auto v : array) {
+      data[i++] = static_cast<uint8_t>(v);
+    }
+    ctor.Place(reinterpret_cast<const char*>(data.data()), data.size());
+  }
+  return std::string(buffer.data(), buffer.size());
+}
+
+// Encode a list of strings as bytes using bc::Vector<bc::String>. The bytes
+// can be decoded directly using bc::Vector<bc::String>. If `array` is not a
+// list of strings, a nullopt will be returned.
+std::optional<std::string> EncodeListOfString(mlir::ArrayAttr array) {
+  bc::Buffer buffer;
+  bc::Allocator allocator(&buffer);
+  auto ctor = bc::New<bc::Vector<bc::String>>(&allocator, array.size());
+
+  for (int i = 0; i < array.size(); ++i) {
+    if (auto string_attr = mlir::dyn_cast<mlir::StringAttr>(array[i])) {
+      ctor.ConstructAt(i, string_attr.getValue().str());
+    } else {
+      return std::nullopt;
+    }
+  }
+
+  return std::string(buffer.data(), buffer.size());
+}
+
+struct FunctionEmitterContext {
+  explicit FunctionEmitterContext(const ModuleEmitterContext* module_context)
+      : module_context(*module_context) {}
+
+  const ModuleEmitterContext& module_context;
+
+  struct RegInfo {
+    int num_uses = 0;
+    int id = -1;
+    bool persistent = false;  // True if the register should not be freed
+  };
+
+  int next_reg_id = 0;
+  llvm::DenseMap<mlir::Value, RegInfo> register_table;
+  std::vector<int> free_regs;
+
+  int AssignRegId(bool is_persistent) {
+    if (is_persistent) {
+      // Persistent types ALWAYS get a brand new ID.
+      return next_reg_id++;
+    }
+
+    // Non-persistent types can reuse from free_regs.
+    if (!free_regs.empty()) {
+      int id = free_regs.back();
+      free_regs.pop_back();
+      return id;
+    }
+    return next_reg_id++;
+  }
+
+  void FreeRegId(int id) { free_regs.push_back(id); }
+};
+
+// Emit the bytecode for a kernel. It uses the information in an MLIR operation
+// and populates the bytecode using bc::Kernel::Constructor. For a kernel's
+// bytecode format, please refer to kernel.h.
+void EmitKernel(FunctionEmitterContext& function_context,
+                bc::Kernel::Constructor& constructor, mlir::Operation& op,
+                std::vector<uint32_t>& function_output_regs,
+                std::vector<uint8_t>& function_output_last_uses) {
+  // Assign reg ids for results first to make sure results does not reuse reg
+  // ids freed from args in the same operation.
+  std::vector<uint32_t> results;
+  results.reserve(op.getNumResults());
+  for (auto result : op.getResults()) {
+    auto iter = function_context.register_table.find(result);
+    CHECK(iter != function_context.register_table.end());  // Crash Ok
+    CHECK_EQ(iter->second.id, -1);                         // Crash Ok
+    iter->second.id = function_context.AssignRegId(iter->second.persistent);
+    results.push_back(iter->second.id);
+  }
+  constructor.construct_results(results.size())
+      .Assign(results.begin(), results.end());
+
+  std::vector<uint32_t> arguments;
+  std::vector<uint8_t> last_uses;
+  arguments.reserve(op.getNumOperands());
+  last_uses.reserve(op.getNumOperands());
+  for (auto operand : op.getOperands()) {
+    auto iter = function_context.register_table.find(operand);
+    CHECK(iter != function_context.register_table.end());  // Crash Ok
+    int id = iter->second.id;
+    CHECK_NE(id, -1);  // Crash Ok
+    last_uses.push_back(0);
+    auto& reg_info = iter->second;
+    if (!reg_info.persistent) {
+      if (--reg_info.num_uses == 0) {
+        function_context.FreeRegId(id);
+        last_uses.back() = 1;
+      }
+    }
+    arguments.push_back(id);
+  }
+
+  constructor.construct_arguments(arguments.size())
+      .Assign(arguments.begin(), arguments.end());
+  constructor.construct_last_uses(last_uses.size())
+      .Assign(last_uses.begin(), last_uses.end());
+
+  std::vector<uint32_t> attributes;
+  attributes.reserve(op.getAttrs().size());
+  for (auto attr : op.getAttrs()) {
+    int attr_id =
+        function_context.module_context.GetAttributeId(attr.getValue());
+    absl::string_view attr_data =
+        function_context.module_context.attributes().at(attr_id);
+
+    if (CanBeInlined(attr.getValue(), attr_data)) {
+      uint32_t data = 0;
+      std::memcpy(&data, attr_data.data(), attr_data.size());
+      attributes.push_back(data);
+    } else {
+      attributes.push_back(attr_id);
+    }
+  }
+  constructor.construct_attributes(attributes.size())
+      .Assign(attributes.begin(), attributes.end());
+
+  if (llvm::isa<mlir::func::ReturnOp>(&op)) {
+    constructor.set_code(function_context.module_context.GetKernelId("return"));
+
+    function_output_regs = std::move(arguments);
+    function_output_last_uses = std::move(last_uses);
+
+  } else if (llvm::isa<mlir::func::CallOp>(&op)) {
+    constructor.set_code(function_context.module_context.GetKernelId("call"));
+  } else {
+    llvm::StringRef op_name = op.getName().getStringRef();
+    constructor.set_code(function_context.module_context.GetKernelId(op_name));
+  }
+}
+
+// Emit the bytecode for a function. It uses information in an MLIR function or
+// an MLIR region, and populates the bytecode using bc::Function::Constructor.
+// For a function's bytecode format, please refer to function.h.
+void EmitFunction(const ModuleEmitterContext& module_context,
+                  bc::Function::Constructor& constructor, llvm::StringRef name,
+                  mlir::Region& region) {
+  FunctionEmitterContext function_context(&module_context);
+
+  constructor.construct_name(name.str());
+
+  DCHECK(llvm::hasSingleElement(region)) << "should have a single block";
+
+  auto& block = region.front();
+
+  auto& register_table = function_context.register_table;
+
+  std::vector<uint32_t> input_regs;
+  input_regs.reserve(block.getNumArguments());
+  for (auto arg : block.getArguments()) {
+    bool persistent = mlir::isa<mlrt::compiler::AsyncHandleType>(arg.getType());
+    int id = function_context.AssignRegId(persistent);
+    input_regs.push_back(id);
+    register_table[arg] = {static_cast<int>(std::distance(arg.getUses().begin(),
+                                                          arg.getUses().end())),
+                           id, persistent};
+  }
+  constructor.construct_input_regs(input_regs);
+
+  for (auto& op : block) {
+    for (auto result : op.getResults()) {
+      bool persistent =
+          mlir::isa<mlrt::compiler::AsyncHandleType>(result.getType());
+      register_table[result] = {
+          static_cast<int>(
+              std::distance(result.getUses().begin(), result.getUses().end())),
+          -1, persistent};
+    }
+  }
+
+  auto kernels_constructor =
+      constructor.construct_kernels(block.getOperations().size());
+
+  std::vector<uint32_t> output_regs;
+  std::vector<uint8_t> output_last_uses;
+  for (const auto& iter : llvm::enumerate(block.getOperations())) {
+    int i = iter.index();
+    mlir::Operation& op = iter.value();
+    auto kernel_ctor = kernels_constructor.ConstructAt(i);
+    EmitKernel(function_context, kernel_ctor, op, output_regs,
+               output_last_uses);
+  }
+
+  constructor.set_num_regs(function_context.next_reg_id);
+  constructor.construct_output_regs(output_regs);
+  constructor.construct_output_last_uses(output_last_uses);
+}
+
+// Emit the bytecode for an executable. It converts attributes, kernels, and
+// functions in an MLIR module to bytecode using bc::Executable::Constructor.
+// For an executable's bytecode format, please refer to executable.h.
+absl::Status EmitExecutable(ModuleEmitterContext& module_context,
+                            bc::Executable::Constructor& constructor,
+                            mlir::ModuleOp module) {
+  module.walk(
+      [&](mlir::func::FuncOp func) { module_context.AddFunction(func); });
+
+  auto functions = module_context.functions();
+  for (auto func : functions) {
+    if (!llvm::hasSingleElement(func.getRegion())) {
+      return absl::InvalidArgumentError("function should have a single block.");
+    }
+    auto& block = func.getRegion().front();
+
+    for (auto& op : block) {
+      if (llvm::isa<mlir::func::CallOp>(&op)) {
+        // Canonicalize the MLIR builtin call op's name to "call".
+        module_context.AddKernelName("call");
+      } else if (llvm::isa<mlir::func::ReturnOp>(&op)) {
+        // Canonicalize the return op's name to "return".
+        if (op.getNumResults() != 0) {
+          return absl::InvalidArgumentError(
+              "Block terminator must be a return op.");
+        }
+        module_context.AddKernelName("return");
+      } else {
+        module_context.AddKernelName(op.getName().getStringRef().str());
+      }
+
+      for (auto attr : op.getAttrs()) {
+        if (auto status = module_context.AddAttribute(&op, attr.getValue());
+            !status.ok()) {
+          return status;
+        }
+      }
+
+      // TODO(chky): Support inline regions.
+    }
+  }
+
+  constructor.construct_kernel_names(module_context.kernels().size())
+      .Assign(module_context.kernels().begin(), module_context.kernels().end());
+
+  auto functions_constructor =
+      constructor.construct_functions(functions.size());
+  for (int i = 0; i < functions.size(); ++i) {
+    auto func = functions[i];
+    auto function_ctor = functions_constructor.ConstructAt(i);
+    EmitFunction(module_context, function_ctor, func.getSymName(),
+                 func.getRegion());
+  }
+
+  // Emit attributes after emitting functions as attributes might be large.
+  // Large attributes may result in large offsets that do not fit into a
+  // unit32_t integer. Since functions section should fit into 2GB size limit,
+  // so we emit functions first.
+  constructor.construct_attributes(module_context.attributes().size())
+      .Assign(module_context.attributes().begin(),
+              module_context.attributes().end());
+
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+absl::Status ModuleEmitterContext::AddAttribute(mlir::Operation* op,
+                                                mlir::Attribute attr) {
+  absl::StatusOr<std::string> attr_data;
+  if (auto* encoder = attribute_encoder_registry_.Get(
+          op->getName().getDialectNamespace())) {
+    attr_data = (*encoder)(*this, attr);
+  } else {
+    attr_data = DefaultEncodeAttribute(attr);
+  }
+  if (!attr_data.ok()) return std::move(attr_data).status();
+
+  int id = AddData(std::move(*attr_data), attributes_, attribute_data_id_map_);
+  attribute_id_map_[attr] = id;
+
+  return absl::OkStatus();
+}
+
+int ModuleEmitterContext::AddFunction(mlir::func::FuncOp func) {
+  int id = functions_.size();
+  functions_.push_back(func);
+  DCHECK(!function_name_id_map_.contains(func.getSymName()));
+  function_name_id_map_[func.getSymName()] = id;
+  return id;
+}
+
+std::optional<std::string> EncodeSimpleAttribute(
+    const ModuleEmitterContext& module_context, mlir::Attribute attr) {
+  return llvm::TypeSwitch<mlir::Attribute, std::optional<std::string>>(attr)
+      .Case<mlir::StringAttr>(
+          [](const auto& str_attr) { return str_attr.str(); })
+      .Case<mlir::IntegerAttr>(
+          [](const auto& integer_attr) -> std::optional<std::string> {
+            switch (llvm::APInt value = integer_attr.getValue();
+                    value.getBitWidth()) {
+              case 1:
+                return EncodeIntegerOrFloat<uint8_t>(value.getZExtValue());
+              case 32:
+                return EncodeIntegerOrFloat<uint32_t>(value.getZExtValue());
+              case 64:
+                return EncodeIntegerOrFloat<uint64_t>(value.getZExtValue());
+              default:
+                return std::nullopt;
+            }
+          })
+      .Case<mlir::FloatAttr>(
+          [](const auto& float_attr) -> std::optional<std::string> {
+            llvm::APFloat value = float_attr.getValue();
+            if (float_attr.getType().isF32()) {
+              return EncodeIntegerOrFloat<float>(value.convertToFloat());
+            }
+            return std::nullopt;
+          })
+      .Case<mlir::ArrayAttr>([&](const auto& array_attr)
+                                 -> std::optional<std::string> {
+        if (auto encoded_list_i32 = EncodeListOfInteger<uint32_t>(array_attr)) {
+          return std::move(*encoded_list_i32);
+        } else if (auto encoded_list_i64 =
+                       EncodeListOfInteger<uint64_t>(array_attr)) {
+          return std::move(*encoded_list_i64);
+        } else if (auto encoded_list_string = EncodeListOfString(array_attr)) {
+          return std::move(*encoded_list_string);
+        } else if (auto encoded_list_symbol_ref =
+                       EncodeListOfSymbolRef(module_context, array_attr)) {
+          return std::move(*encoded_list_symbol_ref);
+        } else {
+          return std::nullopt;
+        }
+      })
+      .Case<mlir::DenseI32ArrayAttr>(
+          [](const auto& dense_array_i32) -> std::optional<std::string> {
+            return EncodeDenseArray<int32_t>(dense_array_i32);
+          })
+      .Case<mlir::DenseI64ArrayAttr>(
+          [](const auto& dense_array_i64) -> std::optional<std::string> {
+            return EncodeDenseArray<int64_t>(dense_array_i64);
+          })
+      .Case<mlir::DenseBoolArrayAttr>(
+          [](const auto& dense_array_bool) -> std::optional<std::string> {
+            return EncodeDenseBoolArray(dense_array_bool.asArrayRef());
+          })
+      .Case<mlir::FlatSymbolRefAttr>([&](const auto& symbol_ref) {
+        return EncodeIntegerOrFloat<uint32_t>(
+            module_context.GetFunctionId(symbol_ref.getValue()));
+      })
+      .Default([](const auto& attr) { return std::nullopt; });
+}
+
+// Encode mlir attributes with a limited support such as I64, string and array
+// of I64. Returns an error if the attribute is not supported.
+absl::StatusOr<std::string> ModuleEmitterContext::DefaultEncodeAttribute(
+    mlir::Attribute attr) {
+  if (auto result = EncodeSimpleAttribute(*this, attr)) {
+    return std::move(*result);
+  }
+
+  // TODO(chky): Add a unit test for the error below. This requires we
+  // propagate the error all the way back to the entry point.
+  std ::string attr_str;
+  llvm::raw_string_ostream os(attr_str);
+  attr.print(os);
+
+  return absl::InvalidArgumentError(
+      absl::StrCat("Try to encode unsupported attribute: ", attr_str));
+}
+
+absl::StatusOr<bc::Buffer> EmitExecutable(
+    const AttributeEncoderRegistry& attribute_encoder_registry,
+    mlir::ModuleOp module) {
+  bc::Buffer buffer;
+  bc::Allocator allocator(&buffer);
+
+  ModuleEmitterContext module_context(&attribute_encoder_registry);
+
+  auto executable_ctor = bc::New<bc::Executable>(&allocator);
+
+  if (auto status = EmitExecutable(module_context, executable_ctor, module);
+      !status.ok()) {
+    return status;
+  }
+
+  buffer.shrink_to_fit();
+
+  return buffer;
+}
+
+}  // namespace mlrt

tensorflow/compiler/mlir/tfrt/translate/mlrt/mlir_to_bytecode.h

@@ -0,0 +1,135 @@
+/* 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_TRANSLATE_MLRT_MLIR_TO_BYTECODE_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSLATE_MLRT_MLIR_TO_BYTECODE_H_
+
+#include <functional>
+#include <optional>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/StringRef.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/core/tfrt/mlrt/bytecode/bytecode.h"
+
+namespace mlrt {
+
+class ModuleEmitterContext;
+
+// Defines a custom attribute encoding registry. Users can register custom
+// attribute encoding for their dialects in this registry. If no custom encoder
+// is registered for a dialect, the default encoding with a limited support, the
+// EncodeSimpleAttribute() below, will be used.
+class AttributeEncoderRegistry {
+ public:
+  using EncoderFn = std::function<absl::StatusOr<std::string>(
+      const ModuleEmitterContext&, mlir::Attribute)>;
+
+  void Register(absl::string_view dialect, EncoderFn encoder) {
+    encoders_[dialect] = std::move(encoder);
+  }
+
+  // Returns the encoder for the specified dialect. It can be nullptr if it is
+  // not registered for this dialect. The returned reference will be invalidated
+  // if Register() is called.
+  const EncoderFn* Get(absl::string_view dialect) const {
+    auto iter = encoders_.find(dialect);
+    if (iter != encoders_.end()) return &iter->second;
+    return nullptr;
+  }
+
+ private:
+  absl::flat_hash_map<std::string, EncoderFn> encoders_;
+};
+
+class ModuleEmitterContext {
+ public:
+  explicit ModuleEmitterContext(
+      const AttributeEncoderRegistry* attribute_encoder_registry)
+      : attribute_encoder_registry_(*attribute_encoder_registry) {}
+
+  void AddKernelName(std::string name) {
+    AddData(std::move(name), kernels_, kernel_id_map_);
+  }
+
+  int GetKernelId(llvm::StringRef name) const {
+    return kernel_id_map_.at(name);
+  }
+
+  absl::Status AddAttribute(mlir::Operation* op, mlir::Attribute attr);
+
+  int GetAttributeId(mlir::Attribute attr) const {
+    return attribute_id_map_.lookup(attr);
+  }
+
+  int AddFunction(mlir::func::FuncOp func);
+
+  int GetFunctionId(absl::string_view name) const {
+    return function_name_id_map_.at(name);
+  }
+
+  absl::Span<const std::string> kernels() const { return kernels_; }
+  absl::Span<const std::string> attributes() const { return attributes_; }
+  absl::Span<const mlir::func::FuncOp> functions() const { return functions_; }
+
+ private:
+  int AddData(std::string data, std::vector<std::string>& data_vector,
+              absl::flat_hash_map<std::string, int>& data_map) {
+    auto iter = data_map.find(data);
+    if (iter != data_map.end()) return iter->second;
+
+    int id = data_vector.size();
+    data_map[data] = id;
+    data_vector.push_back(std::move(data));
+    return id;
+  }
+
+  absl::StatusOr<std::string> DefaultEncodeAttribute(mlir::Attribute attr);
+
+  const AttributeEncoderRegistry& attribute_encoder_registry_;
+
+  std::vector<std::string> kernels_;
+  absl::flat_hash_map<std::string, int> kernel_id_map_;
+
+  std::vector<std::string> attributes_;
+  llvm::DenseMap<mlir::Attribute, int> attribute_id_map_;
+  absl::flat_hash_map<std::string, int> attribute_data_id_map_;
+
+  std::vector<mlir::func::FuncOp> functions_;
+  absl::flat_hash_map<std::string, int> function_name_id_map_;
+};
+
+// Encodes a few simple attributes. Users can use this function in their custom
+// attribute encoder.
+std::optional<std::string> EncodeSimpleAttribute(
+    const ModuleEmitterContext& module_context, mlir::Attribute attr);
+
+absl::StatusOr<bc::Buffer> EmitExecutable(
+    const AttributeEncoderRegistry& attribute_encoder_registry,
+    mlir::ModuleOp module);
+
+}  // namespace mlrt
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSLATE_MLRT_MLIR_TO_BYTECODE_H_

tensorflow/compiler/mlir/tfrt/translate/mlrt/mlir_to_bytecode_test.cc

@@ -0,0 +1,511 @@
+/* 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/translate/mlrt/mlir_to_bytecode.h"
+
+#include <cstdint>
+#include <cstring>
+#include <string>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/Parser/Parser.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.h"
+#include "xla/tsl/platform/resource_loader.h"
+#include "tensorflow/core/tfrt/mlrt/bytecode/bytecode.h"
+#include "tensorflow/core/tfrt/mlrt/bytecode/executable.h"
+#include "tensorflow/core/tfrt/mlrt/interpreter/attribute_span.h"
+
+namespace mlrt {
+namespace {
+
+using ::testing::ElementsAreArray;
+using ::testing::FloatEq;
+using ::testing::IsEmpty;
+
+TEST(MlirToByteCodeTest, Basic) {
+  constexpr char kBasicMlir[] =
+      "tensorflow/compiler/mlir/tfrt/translate/mlrt/testdata/basic.mlir";
+
+  mlir::DialectRegistry registry;
+  registry.insert<mlir::func::FuncDialect>();
+  mlir::MLIRContext mlir_context(registry);
+  mlir_context.allowUnregisteredDialects();
+  auto mlir_module = mlir::parseSourceFile<mlir::ModuleOp>(
+      tsl::GetDataDependencyFilepath(kBasicMlir), &mlir_context);
+
+  AttributeEncoderRegistry attribute_encoder_registry;
+  bc::Buffer buffer =
+      EmitExecutable(attribute_encoder_registry, mlir_module.get()).value();
+
+  bc::Executable executable(buffer.data());
+
+  auto kernel_names = executable.kernel_names();
+  EXPECT_THAT(kernel_names,
+              ElementsAreArray({"test_mlbc.add.i32", "test_mlbc.sub.i32",
+                                "call", "return"}));
+
+  auto functions = executable.functions();
+  ASSERT_GE(functions.size(), 1);
+
+  auto function = functions[0];
+  EXPECT_EQ(function.name().str(), "add_i32_10");
+  EXPECT_EQ(function.num_regs(), 5);
+  EXPECT_THAT(function.input_regs(), ElementsAreArray({0}));
+  EXPECT_THAT(function.output_regs(), ElementsAreArray({0, 2, 2}));
+  EXPECT_THAT(function.output_last_uses(),
+              ElementsAreArray({true, false, true}));
+
+  auto kernels = function.kernels();
+  ASSERT_EQ(kernels.size(), 11);
+
+  EXPECT_EQ(kernels[0].code(), 0);
+  EXPECT_THAT(kernels[0].arguments(), ElementsAreArray({0, 0}));
+  EXPECT_THAT(kernels[0].results(), ElementsAreArray({1}));
+  EXPECT_THAT(kernels[0].last_uses(), ElementsAreArray({0, 0}));
+
+  for (int i = 1; i < 9; i++) {
+    EXPECT_EQ(kernels[i].code(), i % 2);
+    EXPECT_THAT(kernels[i].arguments(), ElementsAreArray({(i - 1) % 2 + 1, 0}));
+    EXPECT_THAT(kernels[i].results(), ElementsAreArray({i % 2 + 1}));
+    EXPECT_THAT(kernels[i].last_uses(), ElementsAreArray({1, 0}));
+  }
+
+  EXPECT_EQ(kernels[9].code(), 2);
+  EXPECT_THAT(kernels[9].arguments(), ElementsAreArray({1}));
+  EXPECT_THAT(kernels[9].last_uses(), ElementsAreArray({true}));
+  EXPECT_THAT(kernels[9].results(), ElementsAreArray({2, 3, 4}));
+
+  EXPECT_EQ(kernels[10].code(), 3);
+  EXPECT_THAT(kernels[10].arguments(), ElementsAreArray({0, 2, 2}));
+  EXPECT_THAT(kernels[10].last_uses(), ElementsAreArray({true, false, true}));
+  EXPECT_TRUE(kernels[10].results().empty());
+}
+
+template <typename T>
+absl::StatusOr<T> DecodeAttribute(absl::string_view data) {
+  if (data.size() < sizeof(T))
+    return absl::InvalidArgumentError("Invalid data size for attribute.");
+
+  T value;
+  std::memcpy(&value, data.data(), sizeof(T));
+  return value;
+}
+
+TEST(MlirToByteCodeTest, BasicAttributes) {
+  constexpr char kBasicAttributesMlir[] =
+      "tensorflow/compiler/mlir/tfrt/translate/mlrt/testdata/"
+      "basic_attributes.mlir";
+
+  mlir::DialectRegistry registry;
+  registry.insert<mlir::func::FuncDialect>();
+  mlir::MLIRContext mlir_context(registry);
+  mlir_context.allowUnregisteredDialects();
+  auto mlir_module = mlir::parseSourceFile<mlir::ModuleOp>(
+      tsl::GetDataDependencyFilepath(kBasicAttributesMlir), &mlir_context);
+
+  AttributeEncoderRegistry attribute_encoder_registry;
+  bc::Buffer buffer =
+      EmitExecutable(attribute_encoder_registry, mlir_module.get()).value();
+
+  bc::Executable executable(buffer.data());
+
+  auto attributes = executable.attributes();
+
+  ASSERT_EQ(attributes.size(), 15);
+
+  auto attr_iter = attributes.begin();
+
+  EXPECT_EQ(*attr_iter, "test string");
+  ++attr_iter;
+
+  EXPECT_EQ(*attr_iter, "ts");
+  ++attr_iter;
+
+  EXPECT_THAT(DecodeAttribute<int32_t>(*attr_iter),
+              absl_testing::IsOkAndHolds(100));
+  ++attr_iter;
+
+  EXPECT_THAT(DecodeAttribute<int64_t>(*attr_iter),
+              absl_testing::IsOkAndHolds(200));
+  ++attr_iter;
+
+  EXPECT_THAT(DecodeAttribute<float>(*attr_iter),
+              absl_testing::IsOkAndHolds(FloatEq(3.0)));
+  ++attr_iter;
+
+  EXPECT_THAT(DecodeAttribute<uint8_t>(*attr_iter),
+              absl_testing::IsOkAndHolds(0));
+  ++attr_iter;
+
+  bc::Vector<int64_t> list_of_i64((*attr_iter).data());
+  EXPECT_THAT(list_of_i64, ElementsAreArray({0, 1, 2, 3, 4}));
+  ++attr_iter;
+
+  bc::Vector<int32_t> list_of_i32((*attr_iter).data());
+  EXPECT_THAT(list_of_i32, ElementsAreArray({0, 1, 2, 3}));
+  ++attr_iter;
+
+  bc::Vector<bc::String> list_of_str((*attr_iter).data());
+  EXPECT_THAT(list_of_str, ElementsAreArray({"string 0", "string 1"}));
+  ++attr_iter;
+
+  EXPECT_THAT(DecodeAttribute<uint32_t>(*attr_iter),
+              absl_testing::IsOkAndHolds(1));
+  EXPECT_EQ(executable.functions()[1].name().Get(), "callee");
+  ++attr_iter;
+
+  bc::Vector<int32_t> list_of_symbol_ref((*attr_iter).data());
+  EXPECT_EQ(executable.functions()[2].name().Get(), "callee0");
+  EXPECT_EQ(executable.functions()[3].name().Get(), "callee1");
+  EXPECT_THAT(list_of_symbol_ref, ElementsAreArray({2, 3}));
+  ++attr_iter;
+
+  bc::Vector<int32_t> dense_array_of_i32((*attr_iter).data());
+  EXPECT_THAT(dense_array_of_i32, ElementsAreArray({0, 1, 2}));
+  ++attr_iter;
+
+  bc::Vector<int64_t> dense_array_of_i64((*attr_iter).data());
+  EXPECT_THAT(dense_array_of_i64, ElementsAreArray({0, 1, 2}));
+  ++attr_iter;
+
+  bc::Vector<int32_t> empty_dense_array((*attr_iter).data());
+  EXPECT_TRUE(empty_dense_array.empty());
+  ++attr_iter;
+
+  bc::Vector<uint8_t> dense_array_of_bool((*attr_iter).data());
+  EXPECT_THAT(dense_array_of_bool, ElementsAreArray({true, false}));
+
+  auto kernels = executable.functions()[0].kernels();
+  ASSERT_EQ(kernels.size(), 16);
+  auto kernel_iter = kernels.begin();
+
+  auto attribute_span = [&](auto kernel_iter) {
+    return mlrt::AttributeSpan((*kernel_iter).attributes(), attributes);
+  };
+
+  EXPECT_EQ(attribute_span(kernel_iter).GetAs<bc::String>(0).Get(),
+            "test string");
+  ++kernel_iter;
+
+  EXPECT_EQ(attribute_span(kernel_iter).GetAs<bc::String>(0).Get(), "ts");
+  ++kernel_iter;
+
+  EXPECT_EQ(attribute_span(kernel_iter).GetAs<int32_t>(0), 100);
+  ++kernel_iter;
+
+  EXPECT_EQ(attribute_span(kernel_iter).GetAs<int64_t>(0), 200);
+  ++kernel_iter;
+
+  EXPECT_THAT(attribute_span(kernel_iter).GetAs<float>(0), FloatEq(3.0));
+  ++kernel_iter;
+
+  EXPECT_EQ(attribute_span(kernel_iter).GetAs<uint8_t>(0), false);
+  ++kernel_iter;
+
+  EXPECT_THAT(attribute_span(kernel_iter).GetAs<bc::Vector<int64_t>>(0),
+              ElementsAreArray({0, 1, 2, 3, 4}));
+  ++kernel_iter;
+
+  EXPECT_THAT(attribute_span(kernel_iter).GetAs<bc::Vector<int32_t>>(0),
+              ElementsAreArray({0, 1, 2, 3}));
+  ++kernel_iter;
+
+  EXPECT_THAT(attribute_span(kernel_iter).GetAs<bc::Vector<bc::String>>(0),
+              ElementsAreArray({"string 0", "string 1"}));
+  ++kernel_iter;
+
+  EXPECT_EQ(attribute_span(kernel_iter).GetAs<uint32_t>(0), 1);
+  ++kernel_iter;
+
+  EXPECT_THAT(attribute_span(kernel_iter).GetAs<bc::Vector<int32_t>>(0),
+              ElementsAreArray({2, 3}));
+  ++kernel_iter;
+
+  EXPECT_THAT(attribute_span(kernel_iter).GetAs<bc::Vector<int32_t>>(0),
+              ElementsAreArray({0, 1, 2}));
+  ++kernel_iter;
+
+  EXPECT_THAT(attribute_span(kernel_iter).GetAs<bc::Vector<int64_t>>(0),
+              ElementsAreArray({0, 1, 2}));
+  ++kernel_iter;
+
+  EXPECT_THAT(attribute_span(kernel_iter).GetAs<bc::Vector<int32_t>>(0),
+              IsEmpty());
+  ++kernel_iter;
+
+  EXPECT_THAT(attribute_span(kernel_iter).GetAs<bc::Vector<bool>>(0),
+              ElementsAreArray({true, false}));
+}
+
+TEST(MlirToByteCodeTest, UnsupportedAttributes) {
+  constexpr char kUnsupportedAttributesMlir[] =
+      "tensorflow/compiler/mlir/tfrt/translate/mlrt/testdata/"
+      "unsupported_attributes.mlir";
+
+  mlir::DialectRegistry registry;
+  registry.insert<mlir::func::FuncDialect>();
+  mlir::MLIRContext mlir_context(registry);
+  mlir_context.allowUnregisteredDialects();
+  auto mlir_module = mlir::parseSourceFile<mlir::ModuleOp>(
+      tsl::GetDataDependencyFilepath(kUnsupportedAttributesMlir),
+      &mlir_context);
+
+  AttributeEncoderRegistry attribute_encoder_registry;
+  EXPECT_THAT(
+      EmitExecutable(attribute_encoder_registry, mlir_module.get()),
+      absl_testing::StatusIs(absl::StatusCode::kInvalidArgument,
+                             "Try to encode unsupported attribute: unit"));
+}
+
+class CustomDense {
+ public:
+  struct StorageType {
+    using Self = StorageType;
+    DEFINE_BYTECODE_FIELD(bc::Vector<int64_t>, shape);
+    DEFINE_BYTECODE_FIELD(bc::Vector<uint32_t>, data);
+  };
+
+  class Constructor {
+   public:
+    Constructor(bc::Allocator* allocator, bc::BcAddr_t address)
+        : allocator_(allocator), address_(address) {}
+
+    template <typename... Args>
+    auto construct_shape(Args&&... args) {
+      return StorageType::construct_shape(allocator_, address_,
+                                          std::forward<Args>(args)...);
+    }
+    template <typename... Args>
+    auto construct_data(Args&&... args) {
+      return StorageType::construct_data(allocator_, address_,
+                                         std::forward<Args>(args)...);
+    }
+
+    bc::BcAddr_t address() const { return address_; }
+
+   private:
+    bc::Allocator* allocator_;
+    bc::BcAddr_t address_;
+  };
+  using NonTrivialConstructorType = Constructor;
+
+  explicit CustomDense(const char* p) : p_(p) {}
+
+  bc::Vector<int64_t> shape() const { return StorageType::read_shape(p_); }
+  bc::Vector<uint32_t> data() const { return StorageType::read_data(p_); }
+
+ private:
+  const char* p_ = nullptr;
+};
+
+absl::StatusOr<std::string> EncodeCustomDense(const ModuleEmitterContext&,
+                                              mlir::Attribute attr) {
+  auto dense_int_attr = mlir::dyn_cast<mlir::DenseIntElementsAttr>(attr);
+  if (!dense_int_attr)
+    return absl::InvalidArgumentError(
+        "The element of the custom dense attribute must be an integer.");
+
+  if (mlir::cast<mlir::IntegerType>(dense_int_attr.getElementType())
+          .getWidth() != 32) {
+    return absl::InvalidArgumentError(
+        "The element of the custom dense attribute must be an i32 integer.");
+  }
+
+  bc::Buffer buffer;
+  bc::Allocator allocator(&buffer);
+  auto custom_dense_ctor = bc::New<CustomDense>(&allocator);
+
+  auto shaped_type = dense_int_attr.getType();
+  std::vector<int64_t> shape(shaped_type.getShape().begin(),
+                             shaped_type.getShape().end());
+  custom_dense_ctor.construct_shape(shape);
+
+  custom_dense_ctor.construct_data(shaped_type.getNumElements())
+      .Place(dense_int_attr.getRawData().data(),
+             dense_int_attr.getRawData().size());
+
+  return std::string(buffer.data(), buffer.size());
+}
+
+TEST(MlirToByteCodeTest, CustomDense) {
+  constexpr char kCustomAttributesMlir[] =
+      "tensorflow/compiler/mlir/tfrt/translate/mlrt/testdata/"
+      "custom_attributes.mlir";
+
+  mlir::DialectRegistry registry;
+  registry.insert<mlir::func::FuncDialect>();
+  mlir::MLIRContext mlir_context(registry);
+  mlir_context.allowUnregisteredDialects();
+  auto mlir_module = mlir::parseSourceFile<mlir::ModuleOp>(
+      tsl::GetDataDependencyFilepath(kCustomAttributesMlir), &mlir_context);
+
+  AttributeEncoderRegistry attribute_encoder_registry;
+  attribute_encoder_registry.Register("test_custom", &EncodeCustomDense);
+  bc::Buffer buffer =
+      EmitExecutable(attribute_encoder_registry, mlir_module.get()).value();
+
+  bc::Executable executable(buffer.data());
+
+  auto attributes = executable.attributes();
+
+  ASSERT_EQ(attributes.size(), 10);
+  for (int i = 0; i < 10; ++i) {
+    bc::String attr_data = attributes[i];
+
+    CustomDense custom_dense(attr_data.data());
+    EXPECT_THAT(custom_dense.shape(), ElementsAreArray({1}));
+    EXPECT_THAT(custom_dense.data(), ElementsAreArray({i}));
+  }
+}
+
+TEST(MlirToByteCodeTest, AsyncNotFreed) {
+  constexpr char kAsyncMlir[] =
+      "tensorflow/compiler/mlir/tfrt/translate/mlrt/testdata/async.mlir";
+
+  mlir::DialectRegistry registry;
+  registry.insert<mlir::func::FuncDialect, mlrt::compiler::MlrtDialect>();
+  mlir::MLIRContext mlir_context(registry);
+  mlir_context.allowUnregisteredDialects();
+  auto mlir_module = mlir::parseSourceFile<mlir::ModuleOp>(
+      tsl::GetDataDependencyFilepath(kAsyncMlir), &mlir_context);
+
+  AttributeEncoderRegistry attribute_encoder_registry;
+  bc::Buffer buffer =
+      EmitExecutable(attribute_encoder_registry, mlir_module.get()).value();
+
+  bc::Executable executable(buffer.data());
+
+  auto kernel_names = executable.kernel_names();
+  EXPECT_THAT(kernel_names,
+              ElementsAreArray({"test_mlbc.add.i32", "return", "mlrt.async",
+                                "mlrt.await_handle"}));
+
+  auto functions = executable.functions();
+  ASSERT_EQ(functions.size(), 2);
+
+  auto function = functions[1];
+  EXPECT_EQ(function.name().str(), "main");
+  EXPECT_EQ(function.num_regs(), 4);
+  EXPECT_THAT(function.input_regs(), ElementsAreArray({0, 1}));
+  EXPECT_THAT(function.output_regs(), ElementsAreArray({1}));
+  EXPECT_THAT(function.output_last_uses(), ElementsAreArray({true}));
+
+  auto kernels = function.kernels();
+  ASSERT_EQ(kernels.size(), 5);
+
+  EXPECT_EQ(kernels[0].code(), 2);  // mlrt.async
+  EXPECT_THAT(kernels[0].arguments(), ElementsAreArray({0, 1}));
+  // The returned handle is in register 2, which is never used by other kernels.
+  EXPECT_THAT(kernels[0].results(), ElementsAreArray({2}));
+  EXPECT_THAT(kernels[0].last_uses(), ElementsAreArray({false, false}));
+
+  EXPECT_EQ(kernels[1].code(), 3);  // mlrt.await_handle
+  EXPECT_THAT(kernels[1].arguments(), ElementsAreArray({2}));
+  EXPECT_THAT(kernels[1].results(), IsEmpty());
+
+  EXPECT_EQ(kernels[2].code(), 0);  // test_mlbc.add.i32
+  EXPECT_THAT(kernels[2].arguments(), ElementsAreArray({0, 1}));
+  EXPECT_THAT(kernels[2].results(), ElementsAreArray({3}));
+  EXPECT_THAT(kernels[2].last_uses(), ElementsAreArray({true, true}));
+
+  EXPECT_EQ(kernels[3].code(), 0);  // test_mlbc.add.i32
+  EXPECT_THAT(kernels[3].arguments(), ElementsAreArray({3, 3}));
+  EXPECT_THAT(kernels[3].results(), ElementsAreArray({1}));
+  EXPECT_THAT(kernels[3].last_uses(), ElementsAreArray({false, true}));
+
+  EXPECT_EQ(kernels[4].code(), 1);  // return
+  EXPECT_THAT(kernels[4].arguments(), ElementsAreArray({1}));
+  EXPECT_THAT(kernels[4].results(), IsEmpty());
+  EXPECT_THAT(kernels[4].last_uses(), ElementsAreArray({true}));
+}
+
+TEST(MlirToByteCodeTest, AsyncUseNewId) {
+  constexpr char kAsyncMlir[] =
+      "tensorflow/compiler/mlir/tfrt/translate/mlrt/testdata/async2.mlir";
+
+  mlir::DialectRegistry registry;
+  registry.insert<mlir::func::FuncDialect, mlrt::compiler::MlrtDialect>();
+  mlir::MLIRContext mlir_context(registry);
+  mlir_context.allowUnregisteredDialects();
+  auto mlir_module = mlir::parseSourceFile<mlir::ModuleOp>(
+      tsl::GetDataDependencyFilepath(kAsyncMlir), &mlir_context);
+
+  AttributeEncoderRegistry attribute_encoder_registry;
+  bc::Buffer buffer =
+      EmitExecutable(attribute_encoder_registry, mlir_module.get()).value();
+
+  bc::Executable executable(buffer.data());
+
+  auto kernel_names = executable.kernel_names();
+  EXPECT_THAT(kernel_names,
+              ElementsAreArray({"test_mlbc.add.i32", "return", "mlrt.async",
+                                "mlrt.await_handle"}));
+
+  auto functions = executable.functions();
+  ASSERT_EQ(functions.size(), 2);
+
+  auto function = functions[1];
+  EXPECT_EQ(function.name().str(), "main");
+  EXPECT_EQ(function.num_regs(), 4);
+  EXPECT_THAT(function.input_regs(), ElementsAreArray({0, 1}));
+  EXPECT_THAT(function.output_regs(), ElementsAreArray({1}));
+  EXPECT_THAT(function.output_last_uses(), ElementsAreArray({true}));
+
+  auto kernels = function.kernels();
+  ASSERT_EQ(kernels.size(), 5);
+
+  EXPECT_EQ(kernels[0].code(), 0);  // test_mlbc.add.i32
+  EXPECT_THAT(kernels[0].arguments(), ElementsAreArray({0, 1}));
+  EXPECT_THAT(kernels[0].results(), ElementsAreArray({2}));
+  EXPECT_THAT(kernels[0].last_uses(), ElementsAreArray({true, true}));
+
+  EXPECT_EQ(kernels[1].code(), 2);  // mlrt.async
+  EXPECT_THAT(kernels[1].arguments(), ElementsAreArray({2, 2}));
+  // The returned handle is in register 3, which is never used by other kernels.
+  EXPECT_THAT(kernels[1].results(), ElementsAreArray({3}));
+  EXPECT_THAT(kernels[1].last_uses(), ElementsAreArray({false, false}));
+
+  EXPECT_EQ(kernels[2].code(), 3);  // mlrt.await_handle
+  EXPECT_THAT(kernels[2].arguments(), ElementsAreArray({3}));
+  EXPECT_THAT(kernels[2].results(), IsEmpty());
+  EXPECT_THAT(kernels[2].last_uses(), ElementsAreArray({false}));
+
+  EXPECT_EQ(kernels[3].code(), 0);  // test_mlbc.add.i32
+  EXPECT_THAT(kernels[3].arguments(), ElementsAreArray({2, 2}));
+  // AsyncHandle does not free its register. So this can only use 1.
+  EXPECT_THAT(kernels[3].results(), ElementsAreArray({1}));
+  EXPECT_THAT(kernels[3].last_uses(), ElementsAreArray({false, true}));
+
+  EXPECT_EQ(kernels[4].code(), 1);  // return
+  EXPECT_THAT(kernels[4].arguments(), ElementsAreArray({1}));
+  EXPECT_THAT(kernels[4].results(), IsEmpty());
+  EXPECT_THAT(kernels[4].last_uses(), ElementsAreArray({true}));
+}
+
+}  // namespace
+}  // namespace mlrt

tensorflow/compiler/mlir/tfrt/translate/mlrt/test_utils.cc

@@ -0,0 +1,183 @@
+/* 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/translate/mlrt/test_utils.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <cstring>
+#include <numeric>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/tfrt/mlrt/attribute/attribute.h"
+#include "tensorflow/core/tfrt/mlrt/bytecode/bytecode.h"
+#include "tensorflow/core/tfrt/mlrt/bytecode/kernel.h"
+#include "tensorflow/core/tfrt/mlrt/interpreter/context.h"
+#include "tensorflow/core/tfrt/mlrt/interpreter/interpreter_testutil.h"
+
+namespace mlrt {
+namespace testing {
+
+absl::StatusOr<std::string> EncodeAttribute(const tensorflow::AttrValue& attr) {
+  if (attr.has_b()) {
+    std::string result;
+    result.resize(sizeof(uint8_t));
+    uint8_t v = attr.b();
+    std::memcpy(result.data(), &v, sizeof(v));
+    return result;
+  }
+
+  if (attr.has_i()) {
+    std::string result;
+    result.resize(sizeof(int64_t));
+    int64_t v = attr.i();
+    std::memcpy(result.data(), &v, sizeof(v));
+    return result;
+  }
+
+  if (attr.has_f()) {
+    std::string result;
+    result.resize(sizeof(float));
+    float v = attr.f();
+    std::memcpy(result.data(), &v, sizeof(v));
+    return result;
+  }
+
+  if (attr.has_s()) {
+    return attr.s();
+  }
+
+  if (attr.has_list()) {
+    if (attr.list().s_size() > 0) {
+      mlrt::bc::Buffer buffer;
+      mlrt::bc::Allocator allocator(&buffer);
+      auto ctor = mlrt::bc::New<mlrt::bc::Vector<mlrt::bc::String>>(
+          &allocator, attr.list().s_size());
+
+      for (int i = 0; i < attr.list().s_size(); ++i) {
+        ctor.ConstructAt(i, attr.list().s(i));
+      }
+
+      return std::string(buffer.data(), buffer.size());
+    }
+  }
+
+  if (attr.has_tensor()) {
+    mlrt::bc::Buffer buffer;
+    mlrt::bc::Allocator allocator(&buffer);
+
+    tensorflow::Tensor tensor;
+    if (!tensor.FromProto(attr.tensor())) {
+      return absl::InvalidArgumentError("Invalid tensor proto.");
+    }
+
+    auto tensor_attr_ctor = mlrt::bc::New<tensorflow::tf_mlrt::TensorAttr>(
+        &allocator, tensor.dtype());
+
+    auto shape = tensor.shape().dim_sizes();
+
+    tensor_attr_ctor.construct_shape(shape.size())
+        .Assign(shape.begin(), shape.end());
+
+    auto tensor_data = tensor.tensor_data();
+    tensor_attr_ctor.construct_data(tensor_data.size())
+        .Place(tensor_data.data(), tensor_data.size());
+
+    return std::string(buffer.data(), buffer.size());
+  }
+
+  // TODO(chky,rohitju): Add more attribute support.
+
+  return absl::InvalidArgumentError("Unsupported attribute.");
+}
+
+namespace {
+
+bool CanBeInlined(const tensorflow::AttrValue& attr) {
+  return attr.has_b() || attr.has_f();
+}
+
+}  // namespace
+
+absl::Status EncodeAttributes(AttributeTable& attributes,
+                              const tensorflow::AttrValueMap& attr_map) {
+  std::vector<std::pair<std::string, tensorflow::AttrValue>> attrs(
+      attr_map.begin(), attr_map.end());
+  std::sort(attrs.begin(), attrs.end(),
+            [](const auto& x, const auto& y) { return x.first < y.first; });
+
+  for (int i = 0; i < attrs.size(); ++i) {
+    const tensorflow::AttrValue& attr = attrs[i].second;
+    TF_ASSIGN_OR_RETURN(auto attr_str, EncodeAttribute(attr));
+    if (CanBeInlined(attr)) {
+      attributes.AddInline(absl::StrCat(i), attr_str);
+    } else {
+      attributes.Add(absl::StrCat(i), attr_str);
+    }
+  }
+
+  return absl::OkStatus();
+}
+
+absl::StatusOr<std::pair<mlrt::bc::Kernel, mlrt::bc::Vector<mlrt::bc::String>>>
+CreateKernelAndAttrs(int num_inputs, int num_outputs,
+                     mlrt::ExecutionContext& exec_ctx, mlrt::bc::Buffer* buffer,
+                     const tensorflow::AttrValueMap& attrs) {
+  mlrt::bc::Allocator allocator(buffer);
+  auto attributes_ctor = mlrt::bc::New<mlrt::bc::Vector<mlrt::bc::String>>(
+      &allocator, attrs.size());
+  AttributeTable attribute_table(attributes_ctor);
+  TF_RETURN_IF_ERROR(EncodeAttributes(attribute_table, attrs));
+
+  auto kernel_ctor = mlrt::bc::New<mlrt::bc::Kernel>(&allocator);
+  kernel_ctor.set_code(0);
+
+  std::vector<int> input_indices(num_inputs);
+  std::iota(input_indices.begin(), input_indices.end(), 0);
+  kernel_ctor.construct_arguments(input_indices.size())
+      .Assign(input_indices.begin(), input_indices.end());
+
+  std::vector<int> output_indices(num_outputs);
+  std::iota(output_indices.begin(), output_indices.end(), num_inputs);
+  kernel_ctor.construct_results(output_indices.size())
+      .Assign(output_indices.begin(), output_indices.end());
+
+  std::vector<uint32_t> attr_indices;
+  attr_indices.reserve(attrs.size());
+  for (int i = 0; i < attrs.size(); ++i) {
+    attr_indices.push_back(attribute_table.GetHandle(absl::StrCat(i)));
+  }
+
+  kernel_ctor.construct_attributes(attr_indices.size())
+      .Assign(attr_indices.begin(), attr_indices.end());
+
+  mlrt::bc::Vector<mlrt::bc::String> attributes(
+      buffer->Get(attributes_ctor.address()));
+  mlrt::bc::Kernel kernel(buffer->Get(kernel_ctor.address()));
+
+  return std::make_pair(kernel, attributes);
+}
+
+}  // namespace testing
+}  // namespace mlrt

tensorflow/compiler/mlir/tfrt/translate/mlrt/test_utils.h

@@ -0,0 +1,119 @@
+/* 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_TRANSLATE_MLRT_TEST_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSLATE_MLRT_TEST_UTILS_H_
+
+#include <memory>
+#include <numeric>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/tfrt/graph_executor/sync_resource_state.h"
+#include "tensorflow/core/tfrt/mlrt/attribute/attribute.h"
+#include "tensorflow/core/tfrt/mlrt/bytecode/bytecode.h"
+#include "tensorflow/core/tfrt/mlrt/bytecode/kernel.h"
+#include "tensorflow/core/tfrt/mlrt/interpreter/context.h"
+#include "tensorflow/core/tfrt/mlrt/interpreter/interpreter_testutil.h"
+#include "tensorflow/core/tfrt/mlrt/interpreter/value.h"
+#include "tensorflow/core/tfrt/stubs/tfrt_native_lowering_stub.h"
+#include "tensorflow/core/tfrt/utils/tensor_util.h"
+#include "tfrt/host_context/concurrent_work_queue.h"  // from @tf_runtime
+#include "tfrt/host_context/execution_context.h"  // from @tf_runtime
+#include "tfrt/host_context/host_allocator.h"  // from @tf_runtime
+#include "tfrt/host_context/host_context.h"  // from @tf_runtime
+#include "tfrt/support/string_util.h"  // from @tf_runtime
+#include "tfrt/tensor/dense_host_tensor.h"  // from @tf_runtime
+#include "tfrt/tensor/dense_tensor_utils.h"  // from @tf_runtime
+
+namespace mlrt {
+namespace testing {
+
+absl::StatusOr<std::string> EncodeAttribute(const tensorflow::AttrValue& attr);
+
+absl::Status EncodeAttributes(AttributeTable& attributes,
+                              const tensorflow::AttrValueMap& attr_map);
+
+absl::StatusOr<std::pair<mlrt::bc::Kernel, mlrt::bc::Vector<mlrt::bc::String>>>
+CreateKernelAndAttrs(int num_inputs, int num_outputs,
+                     mlrt::ExecutionContext& exec_ctx, mlrt::bc::Buffer* buffer,
+                     const tensorflow::AttrValueMap& attrs = {});
+
+template <typename T>
+absl::Status TestMlrtKernel(
+    absl::string_view kernel_name, absl::Span<mlrt::Value> regs,
+    tfrt::HostContext* host, int num_inputs, int num_outputs,
+    absl::Span<const tensorflow::Tensor> expected_outputs,
+    mlrt::KernelRegistry* registry, bool approx_equal = false,
+    const tensorflow::AttrValueMap& attrs = {}) {
+  mlrt::ExecutionContext execution_context(nullptr);
+
+  mlrt::bc::Buffer buffer;
+  TF_ASSIGN_OR_RETURN(auto kernel_and_attrs,
+                      CreateKernelAndAttrs(num_inputs, num_outputs,
+                                           execution_context, &buffer, attrs));
+
+  tensorflow::tfrt_stub::SyncResourceState sync_resource_state;
+  tfrt::AddSyncContext(execution_context, *host, &sync_resource_state);
+
+  auto kernel_fn = registry->Get(kernel_name);
+  mlrt::KernelFrame::State state(regs, kernel_and_attrs.second,
+                                 &execution_context);
+  mlrt::KernelFrame frame(&state);
+  frame.set_kernel(kernel_and_attrs.first);
+
+  kernel_fn(frame);
+
+  TF_RETURN_IF_ERROR(execution_context.status());
+
+  for (int i = 0, j = num_inputs; i < expected_outputs.size(); ++i, ++j) {
+    const auto& expected_output = expected_outputs[i];
+    auto expected_dht = tfrt::ConvertTfTensorToDHT(expected_output);
+    if (!expected_dht) {
+      return absl::InternalError(tfrt::StrCat(expected_dht.takeError()));
+    }
+
+    if (!approx_equal) {
+      if (!tfrt::TensorEqual<T>(regs[j].Get<tfrt::DenseHostTensor>(),
+                                *expected_dht)) {
+        return absl::InternalError(
+            absl::StrCat("wrong result for ", kernel_name));
+      }
+    } else {
+      if (!tfrt::TensorApproxEqual<T>(regs[j].Get<tfrt::DenseHostTensor>(),
+                                      *expected_dht)) {
+        return absl::InternalError(
+            absl::StrCat("wrong result for ", kernel_name));
+      }
+    }
+  }
+
+  return absl::OkStatus();
+}
+
+}  // namespace testing
+}  // namespace mlrt
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSLATE_MLRT_TEST_UTILS_H_

tensorflow/compiler/mlir/tfrt/translate/mlrt/testdata/async.mlir

@@ -0,0 +1,30 @@
+// Copyright 2026 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.
+// ==============================================================================
+
+func.func @add_i32(%arg0: i32, %arg1: i32) -> i32 {
+  %0 = "test_mlbc.add.i32"(%arg0, %arg1) : (i32, i32) -> i32
+  func.return %0 : i32
+}
+
+func.func @main(%arg0: i32, %arg1: i32) -> i32 {
+  %handle = "mlrt.async"(%arg0, %arg1) {callee = @add_i32} : (i32, i32) -> !mlrt.async_handle
+
+  "mlrt.await_handle"(%handle) : (!mlrt.async_handle) -> () 
+
+  %c1 = "test_mlbc.add.i32"(%arg0, %arg1) : (i32, i32) -> i32
+  %c2 = "test_mlbc.add.i32"(%c1, %c1) : (i32, i32) -> i32
+
+  func.return %c2 : i32
+}

tensorflow/compiler/mlir/tfrt/translate/mlrt/testdata/async2.mlir

@@ -0,0 +1,31 @@
+// Copyright 2026 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.
+// ==============================================================================
+
+func.func @add_i32(%arg0: i32, %arg1: i32) -> i32 {
+  %0 = "test_mlbc.add.i32"(%arg0, %arg1) : (i32, i32) -> i32
+  func.return %0 : i32
+}
+
+func.func @main(%arg0: i32, %arg1: i32) -> i32 {
+  %c1 = "test_mlbc.add.i32"(%arg0, %arg1) : (i32, i32) -> i32
+ 
+  %handle = "mlrt.async"(%c1, %c1) {callee = @add_i32} : (i32, i32) -> !mlrt.async_handle
+
+  "mlrt.await_handle"(%handle) : (!mlrt.async_handle) -> () 
+
+  %c2 = "test_mlbc.add.i32"(%c1, %c1) : (i32, i32) -> i32
+
+  func.return %c2 : i32
+}

tensorflow/compiler/mlir/tfrt/translate/mlrt/testdata/basic.mlir

@@ -0,0 +1,28 @@
+// Copyright 2026 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.
+// ==============================================================================
+
+func.func @add_i32_10(%c0: i32) -> (i32, i32, i32) {
+  %c1 = "test_mlbc.add.i32"(%c0, %c0) : (i32, i32) -> i32
+  %c2 = "test_mlbc.sub.i32"(%c1, %c0) : (i32, i32) -> i32
+  %c3 = "test_mlbc.add.i32"(%c2, %c0) : (i32, i32) -> i32
+  %c4 = "test_mlbc.sub.i32"(%c3, %c0) : (i32, i32) -> i32
+  %c5 = "test_mlbc.add.i32"(%c4, %c0) : (i32, i32) -> i32
+  %c6 = "test_mlbc.sub.i32"(%c5, %c0) : (i32, i32) -> i32
+  %c7 = "test_mlbc.add.i32"(%c6, %c0) : (i32, i32) -> i32
+  %c8 = "test_mlbc.sub.i32"(%c7, %c0) : (i32, i32) -> i32
+  %c9 = "test_mlbc.add.i32"(%c8, %c0) : (i32, i32) -> i32
+  %c10, %c11, %c12 = call @add_i32_10(%c9) : (i32) -> (i32, i32, i32)
+  func.return %c0, %c10, %c10 : i32, i32, i32
+}

tensorflow/compiler/mlir/tfrt/translate/mlrt/testdata/basic_attributes.mlir

@@ -0,0 +1,45 @@
+// Copyright 2026 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.
+// ==============================================================================
+
+func.func @simple_attributes() {
+  "test_custom.attribute"() {value = "test string"} : () -> ()
+  "test_custom.attribute"() {value = "ts"} : () -> ()
+  "test_custom.attribute"() {value = 100 : i32} : () -> ()
+  "test_custom.attribute"() {value = 200 : i64} : () -> ()
+  "test_custom.attribute"() {value = 3.0 : f32} : () -> ()
+  "test_custom.attribute"() {value = false} : () -> ()
+  "test_custom.attribute"() {value = [0, 1, 2, 3, 4]} : () -> ()
+  "test_custom.attribute"() {value = [0 : i32, 1 : i32, 2 : i32, 3 : i32]} : () -> ()
+  "test_custom.attribute"() {value = ["string 0", "string 1"]} : () -> ()
+  "test_custom.attribute"() {value = @callee} : () -> ()
+  "test_custom.attribute"() {value = [@callee0, @callee1]} : () -> ()
+  "test_custom.attribute"() {value = array<i32: 0, 1, 2>} : () -> ()
+  "test_custom.attribute"() {value = array<i64: 0, 1, 2>} : () -> ()
+  "test_custom.attribute"() {value = array<i32>} : () -> ()
+  "test_custom.attribute"() {value = array<i1: true, false>} : () -> ()
+  func.return
+}
+
+func.func @callee() {
+  return
+}
+
+func.func @callee0() {
+  return
+}
+
+func.func @callee1() {
+  return
+}

tensorflow/compiler/mlir/tfrt/translate/mlrt/testdata/custom_attributes.mlir

@@ -0,0 +1,28 @@
+// Copyright 2026 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.
+// ==============================================================================
+
+func.func @add_const_custom_dense_i32_10(%c0: i32) -> i32 {
+  %c1 = "test_custom.add.const.i32"(%c0) {value = dense<[0]> : tensor<1xi32>} : (i32) -> i32
+  %c2 = "test_custom.add.const.i32"(%c1) {value = dense<[1]> : tensor<1xi32>} : (i32) -> i32
+  %c3 = "test_custom.add.const.i32"(%c2) {value = dense<[2]> : tensor<1xi32>} : (i32) -> i32
+  %c4 = "test_custom.add.const.i32"(%c3) {value = dense<[3]> : tensor<1xi32>} : (i32) -> i32
+  %c5 = "test_custom.add.const.i32"(%c4) {value = dense<[4]> : tensor<1xi32>} : (i32) -> i32
+  %c6 = "test_custom.add.const.i32"(%c5) {value = dense<[5]> : tensor<1xi32>} : (i32) -> i32
+  %c7 = "test_custom.add.const.i32"(%c6) {value = dense<[6]> : tensor<1xi32>} : (i32) -> i32
+  %c8 = "test_custom.add.const.i32"(%c7) {value = dense<[7]> : tensor<1xi32>} : (i32) -> i32
+  %c9 = "test_custom.add.const.i32"(%c8) {value = dense<[8]> : tensor<1xi32>} : (i32) -> i32
+  %c10 = "test_custom.add.const.i32"(%c9) {value = dense<[9]> : tensor<1xi32>} : (i32) -> i32
+  func.return %c10 : i32
+}

tensorflow/compiler/mlir/tfrt/translate/mlrt/testdata/unsupported_attributes.mlir

@@ -0,0 +1,20 @@
+// Copyright 2026 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.
+// ==============================================================================
+
+func.func @unsupported_attributes() {
+  "test_custom.attribute"() {unit} : () -> ()
+  func.return
+}
+