chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/tfrt/translate/import_model.cc

@@ -0,0 +1,407 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/tfrt/translate/import_model.h"
+
+#include <cstdint>
+#include <deque>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/functional/function_ref.h"
+#include "absl/log/log.h"
+#include "absl/log/vlog_is_on.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "mlir/Dialect/Func/Extensions/AllExtensions.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/transforms/host_runtime/lower_cluster_to_runtime_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_saved_model_asset_sinking_pass.h"
+#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/serialize_mlir_module_utils.h"
+#include "tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf.h"
+#include "tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor.h"
+#include "tensorflow/compiler/mlir/tf2xla/api/v2/tf_executor_to_graph.h"
+#include "tensorflow/compiler/mlir/tfrt/backend_compiler.h"
+#include "tensorflow/compiler/mlir/tfrt/function/function.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/tfrt_pipeline_options.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/tpu_passes.h"
+#include "tensorflow/compiler/mlir/tfrt/translate/tfrt_compile_options.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/tsl/framework/device_type.h"
+#include "xla/tsl/platform/errors.h"
+#include "tensorflow/core/common_runtime/function_def_utils.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/tfrt/fallback/fallback_state.h"
+#include "tensorflow/core/tfrt/runtime/runtime.h"
+#include "tensorflow/core/tpu/tpu_defs.h"
+#include "tsl/platform/env.h"
+#include "tsl/platform/errors.h"
+#include "tsl/platform/statusor.h"
+#include "tfrt/bef/bef_buffer.h"  // from @tf_runtime
+#include "tfrt/bef_converter/mlir_to_bef.h"  // from @tf_runtime
+#include "tfrt/support/forward_decls.h"  // from @tf_runtime
+
+namespace tensorflow {
+
+namespace {
+
+// Exports all XLA functions in the form of XlaLaunch, and their nested
+// functions.
+absl::StatusOr<std::vector<FunctionDef>> ExportXlaFunctions(
+    mlir::ModuleOp module, std::vector<std::string>* added_xla_function_names) {
+  // Find all XLA functions.
+  std::vector<std::string> xla_functions;
+  module.walk([&](mlir::TF::XlaLaunchOp xla_launch_op) {
+    std::string func_name =
+        xla_launch_op.getFunctionAttr().getRootReference().str();
+    xla_functions.push_back(func_name);
+    if (added_xla_function_names != nullptr) {
+      added_xla_function_names->push_back(func_name);
+    }
+  });
+
+  // Convert all XLA functions and their nested functions.
+  std::deque<std::string> queue;
+  for (const std::string& func : xla_functions) {
+    queue.push_back(func);
+  }
+
+  const mlir::SymbolTable symbol_table(module);
+  absl::flat_hash_set<std::string> visited;
+  std::vector<FunctionDef> xla_func_defs;
+  while (!queue.empty()) {
+    const std::string func_name = queue.front();
+    queue.pop_front();
+
+    if (visited.contains(func_name)) continue;
+
+    const auto func_op = symbol_table.lookup<mlir::func::FuncOp>(func_name);
+    if (!func_op) {
+      return absl::InternalError(
+          absl::StrCat("Function ", func_name, " is not found."));
+    }
+    FunctionDef func_def;
+    TF_RETURN_IF_ERROR(
+        tensorflow::tf2xla::v2::ConvertMlirFunctionToFunctionLibraryDef(
+            func_op, GraphExportConfig(), &func_def));
+    xla_func_defs.push_back(func_def);
+
+    // Visit each op in the function and find out referenced functions from the
+    // attributes.
+    func_op->walk([&](mlir::Operation* op) {
+      for (const mlir::NamedAttribute& attr : op->getAttrs()) {
+        if (const auto sym =
+                mlir::dyn_cast<mlir::FlatSymbolRefAttr>(attr.getValue())) {
+          mlir::Operation* func =
+              mlir::SymbolTable::lookupNearestSymbolFrom(op, sym);
+          if (func) {
+            queue.push_back(sym.getValue().str());
+          }
+        }
+      }
+    });
+
+    // Remove the function from the module, as it will be handled by XLA.
+    // It is safe to remove the function, i.e., the function won't be invoked on
+    // CPU. This is because bridge guarantees that each function has only one
+    // use. We don't replace the uses of the function, because we iterate from
+    // the root caller and hence its uses should have been removed.
+    func_op->erase();
+
+    visited.insert(func_name);
+  }
+  return xla_func_defs;
+}
+
+}  // namespace
+
+absl::Status ConvertTfMlirToRuntimeExecutable(
+    const TfrtCompileOptions& options, mlir::ModuleOp module,
+    absl::FunctionRef<
+        absl::Status(mlir::PassManager&, mlir::ModuleOp,
+                     const tensorflow::TfrtPipelineOptions& options)>
+        emit_executable,
+    tfrt_stub::ModelRuntimeContext& model_context,
+    tfrt_stub::FallbackState* fallback_state,
+    std::vector<std::string>* added_xla_function_names) {
+  mlir::StatusScopedDiagnosticHandler diag_handler(module.getContext());
+  {
+    mlir::PassManager pm(module.getContext());
+    pm.addNestedPass<mlir::func::FuncOp>(
+        mlir::tf_executor::CreateTFExecutorGraphPruningPass());
+    pm.addNestedPass<mlir::func::FuncOp>(
+        mlir::CreateExecutorDialectToFunctionalConversionPass());
+    if (!options.saved_model_dir.empty()) {
+      pm.addPass(mlir::tf_saved_model::CreateAssetSinkingPass(
+          options.saved_model_dir));
+    }
+    if (mlir::failed(pm.run(module))) {
+      return diag_handler.Combine(absl::InternalError(
+          "Failed to sinking assets into initialization graphs."));
+    }
+  }
+
+  if (options.backend_compiler != nullptr) {
+    if (VLOG_IS_ON(1)) {
+      tensorflow::DumpMlirOpToFile("tf_dialect_before_backend_compile", module);
+    }
+    TF_RETURN_IF_ERROR(
+        options.backend_compiler->CompileTensorflow(model_context, module));
+  } else if (options.device_target == TfrtDeviceInfraTarget::kTpurt) {
+    VLOG(1) << "Running MLIR TPU bridge for tpurt";
+    if (VLOG_IS_ON(1)) {
+      tensorflow::DumpMlirOpToFile("tpu_bct_conversion_before", module);
+    }
+
+    TfrtTpuCompileOptions tpu_compile_options;
+    tpu_compile_options.move_resource_gather_to_host =
+        options.tpu_move_resource_gather_to_host;
+    tpu_compile_options.gather_table_width_threshold_bytes =
+        options.tpu_gather_table_width_threshold_bytes;
+
+    auto backward_compat_result =
+        tensorflow::RunTPUBackwardCompatConversion(module, tpu_compile_options);
+    if (mlir::failed(backward_compat_result)) {
+      return diag_handler.Combine(
+          absl::InternalError("Failed to handle legacy TPU Ops"));
+    }
+
+    if (VLOG_IS_ON(1)) {
+      tensorflow::DumpMlirOpToFile("tpu_bct_conversion_after", module);
+    }
+
+    TF_RETURN_IF_ERROR(
+        tensorflow::tf2xla::v2::RunFunctionTf2xlaClusteringBridge(
+            module, /*is_supported_by_replicated_brige*/ true,
+            /*is_in_fallback_enabled_mode=*/VLOG_IS_ON(1)));
+    if (VLOG_IS_ON(1)) {
+      tensorflow::DumpMlirOpToFile("after_tf2xla_clustering_bridge", module);
+    }
+
+    TF_RETURN_IF_ERROR(
+        tensorflow::tfrt_compiler::RunLowerClusterToRuntimeOpsPassPipeline(
+            module, tsl::DeviceType(DEVICE_TPU_XLA_JIT)));
+
+    if (VLOG_IS_ON(1)) {
+      tensorflow::DumpMlirOpToFile("after_lower_cluster_to_runtime_ops",
+                                   module);
+    }
+    TF_RETURN_IF_ERROR(
+        tensorflow::tf2xla::v2::ExportFromTensorflowDialectToExecutor(module));
+  } else if (options.device_target == TfrtDeviceInfraTarget::kTfFallback) {
+    auto tpu_partitioned_call_fallback_compat_result =
+        tensorflow::RunTPUPartitionedCallFallbackCompatConversion(module);
+    if (mlir::failed(tpu_partitioned_call_fallback_compat_result)) {
+      return diag_handler.Combine(absl::InternalError(
+          "Failed to process TPUPartitionedCallOp for fallback execution"));
+    }
+  } else if (options.device_target == TfrtDeviceInfraTarget::kGpu) {
+    TF_RETURN_IF_ERROR(
+        tensorflow::tf2xla::v2::RunFunctionTf2xlaClusteringBridge(
+            module, /*is_supported_by_replicated_brige*/ false,
+            /*is_in_fallback_enabled_mode=*/false));
+
+    TF_RETURN_IF_ERROR(
+        tensorflow::tfrt_compiler::RunLowerClusterToRuntimeOpsPassPipeline(
+            module, tsl::DeviceType(DEVICE_GPU_XLA_JIT)));
+
+    TF_RETURN_IF_ERROR(
+        tensorflow::tf2xla::v2::ExportFromTensorflowDialectToExecutor(module));
+
+    if (options.serialize_mlir_module_to_aot_packages) {
+      const std::string mlir_string = SerializeMlirModule(module);
+      TF_RETURN_IF_ERROR(WriteStringToFile(
+          tsl::Env::Default(), options.aot_mlir_module_file, mlir_string));
+    }
+
+    // GPU XLA clusters are wrapped in functions, which could be transformed by
+    // bridge. Hence, the MLIR functions for XLA clusters are exported and added
+    // to the function library.
+    TF_RETURN_IF_ERROR(
+        AddXlaFunctions(fallback_state, module, added_xla_function_names));
+  }
+
+  if (VLOG_IS_ON(1)) {
+    tensorflow::DumpMlirOpToFile("tf_dialect", module);
+  }
+
+  mlir::DialectRegistry registry;
+  mlir::func::registerAllExtensions(registry);
+  module.getContext()->appendDialectRegistry(registry);
+
+  // Lower MLIR TF Dialect to MLIR TFRT CoreRT dialect.
+  mlir::PassManager pm(module.getContext());
+
+  auto pipeline_options = GetTfrtPipelineOptions(options);
+
+  TF_RETURN_IF_ERROR(
+      tensorflow::CreateTFExecutorToTFPreInvariantOptimizationPipeline(
+          pm, *pipeline_options));
+
+  auto status = emit_executable(pm, module, *pipeline_options);
+
+  if (VLOG_IS_ON(1)) {
+    tensorflow::DumpMlirOpToFile("tfrt_dialect", module);
+  }
+
+  return status;
+}
+
+absl::Status ConvertTfMlirToBef(
+    const TfrtCompileOptions& options, mlir::ModuleOp module,
+    tfrt::BefBuffer* bef_buffer, tfrt_stub::ModelRuntimeContext& model_context,
+    tfrt_stub::FallbackState* fallback_state,
+    std::vector<std::string>* added_xla_function_names) {
+  return ConvertTfMlirToRuntimeExecutable(
+      options, module,
+      [bef_buffer](mlir::PassManager& pm, mlir::ModuleOp module,
+                   const tensorflow::TfrtPipelineOptions& options) {
+        mlir::StatusScopedDiagnosticHandler diag_handler(module.getContext());
+        tensorflow::CreateTFInvariantOptimizationPipelineHelper(pm, options);
+        tensorflow::CreateTfToTfrtPipeline(pm, options);
+
+        if (mlir::failed(pm.run(module))) {
+          if (VLOG_IS_ON(1)) {
+            tensorflow::DumpMlirOpToFile("tf_to_corert_failure", module);
+          }
+          return diag_handler.Combine(absl::InternalError(
+              "failed to lower TF Dialect to CoreRT dialect."));
+        }
+
+        *bef_buffer =
+            tfrt::ConvertMLIRToBEF(module, /*disable_optional_sections=*/true);
+        if (bef_buffer->empty())
+          return diag_handler.Combine(
+              absl::InternalError("failed to convert MLIR to BEF."));
+
+        bef_buffer->shrink_to_fit();
+        return absl::OkStatus();
+      },
+      model_context, fallback_state, added_xla_function_names);
+}
+
+std::unique_ptr<tensorflow::TfrtPipelineOptions> GetTfrtPipelineOptions(
+    const TfrtCompileOptions& options) {
+  auto pipeline_options = std::make_unique<tensorflow::TfrtPipelineOptions>();
+
+  pipeline_options->saved_model_dir = options.saved_model_dir;
+
+  if (!options.default_device.empty()) {
+    pipeline_options->default_device = options.default_device;
+  }
+  if (!options.force_data_format.empty()) {
+    pipeline_options->force_data_format = options.force_data_format;
+  }
+
+  // TODO(b/187991150): Consider only decomposing read-only resource variable
+  // ops.
+  pipeline_options->decompose_resource_ops = options.decompose_resource_ops;
+  pipeline_options->enable_optimizer = options.enable_optimizer;
+  pipeline_options->target_tpurt =
+      (options.device_target == TfrtDeviceInfraTarget::kTpurt);
+  pipeline_options->target_gpu =
+      (options.device_target == TfrtDeviceInfraTarget::kGpu);
+  pipeline_options->use_gpu_compile_and_execute_op =
+      options.use_gpu_compile_and_execute_op;
+  pipeline_options->tpu_fuse_ops = options.tpu_fuse_ops;
+  pipeline_options->use_tpu_host_allocator_for_inputs =
+      options.use_tpu_host_allocator_for_inputs;
+  pipeline_options->tpu_allow_unpadded_batch = options.tpu_allow_unpadded_batch;
+  pipeline_options->sink_in_invariant_ops = options.sink_in_invariant_ops;
+  pipeline_options->hoist_invariant_ops = options.hoist_invariant_ops;
+  pipeline_options->fuse_get_resource_ops_in_hoisting =
+      options.fuse_get_resource_ops_in_hoisting;
+  pipeline_options->enable_while_parallel_iterations =
+      options.enable_while_parallel_iterations;
+  pipeline_options->cost_threshold = options.cost_threshold;
+  pipeline_options->min_num_batch_threads = options.min_num_batch_threads;
+  pipeline_options->min_max_enqueued_batches = options.min_max_enqueued_batches;
+  pipeline_options->batch_queue_global_prioritization_num_threads =
+      options.batch_queue_global_prioritization_num_threads;
+  pipeline_options->enable_priority_aware_batch_scheduler =
+      options.enable_priority_aware_batch_scheduler;
+  pipeline_options->enable_priority_aware_batch_scheduler_resplit =
+      options.enable_priority_aware_batch_scheduler_resplit;
+  pipeline_options->enable_batching_task_lazy_cancellation =
+      options.enable_batching_task_lazy_cancellation;
+  pipeline_options->batch_padding_policy = options.batch_padding_policy;
+  pipeline_options->num_batch_threads =
+      options.batch_options.num_batch_threads();
+  pipeline_options->max_batch_size = options.batch_options.max_batch_size();
+  pipeline_options->batch_timeout_micros =
+      options.batch_options.batch_timeout_micros();
+  pipeline_options->allowed_batch_sizes = llvm::ArrayRef<int64_t>(
+      std::vector<int64_t>(options.batch_options.allowed_batch_sizes().begin(),
+                           options.batch_options.allowed_batch_sizes().end()));
+  pipeline_options->max_enqueued_batches =
+      options.batch_options.max_enqueued_batches();
+  pipeline_options->low_priority_max_batch_size =
+      options.batch_options.low_priority_max_batch_size();
+  pipeline_options->low_priority_batch_timeout_micros =
+      options.batch_options.low_priority_batch_timeout_micros();
+  pipeline_options->low_priority_allowed_batch_sizes =
+      llvm::ArrayRef<int64_t>(std::vector<int64_t>(
+          options.batch_options.low_priority_allowed_batch_sizes().begin(),
+          options.batch_options.low_priority_allowed_batch_sizes().end()));
+  pipeline_options->low_priority_max_enqueued_batches =
+      options.batch_options.low_priority_max_enqueued_batches();
+  pipeline_options->num_warmup_batch_threads =
+      options.batch_options.num_warmup_batch_threads();
+  pipeline_options->enable_large_batch_splitting =
+      options.batch_options.enable_large_batch_splitting();
+  pipeline_options->mixed_priority_batching_policy =
+      options.batch_options.mixed_priority_batching_policy();
+  pipeline_options->merge_inter_dependent_streams =
+      options.merge_inter_dependent_streams;
+  pipeline_options->allow_xla_cpu = options.allow_xla_cpu;
+
+  return pipeline_options;
+}
+
+absl::Status AddXlaFunctions(
+    tfrt_stub::FallbackState* fallback_state, mlir::ModuleOp mlir_module,
+    std::vector<std::string>* added_xla_function_names) {
+  if (fallback_state != nullptr) {
+    TF_ASSIGN_OR_RETURN(
+        const std::vector<FunctionDef> xla_func_defs,
+        ExportXlaFunctions(mlir_module, added_xla_function_names));
+    for (const auto& func_def : xla_func_defs) {
+      TF_RETURN_IF_ERROR(fallback_state->AddFunctionDef(func_def));
+    }
+  }
+
+  return absl::OkStatus();
+}
+
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/translate/import_model.h

@@ -0,0 +1,74 @@
+/* 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_IMPORT_MODEL_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSLATE_IMPORT_MODEL_H_
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "absl/functional/function_ref.h"
+#include "absl/status/status.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/function/function.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/tfrt_pipeline_options.h"
+#include "tensorflow/compiler/mlir/tfrt/translate/tfrt_compile_options.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/tfrt/fallback/fallback_state.h"
+#include "tensorflow/core/tfrt/runtime/runtime.h"
+#include "tfrt/bef/bef_buffer.h"  // from @tf_runtime
+
+namespace tensorflow {
+
+struct FunctionBody;
+
+// Converts an MLIR `module` in TF dialect to TFRT's Binary Executable Format.
+// If `fallback_state` is not null, the MLIR functions for XLA clusters in
+// the form of XlaLaunch will be exported and added to the function library when
+// needed. The nested functions will also be exported. If
+// `added_xla_function_names` is not null, it will be populated with the names
+// of the added XLA functions.
+absl::Status ConvertTfMlirToBef(
+    const TfrtCompileOptions& options, mlir::ModuleOp module,
+    tfrt::BefBuffer* bef_buffer, tfrt_stub::ModelRuntimeContext& model_context,
+    tfrt_stub::FallbackState* fallback_state = nullptr,
+    std::vector<std::string>* added_xla_function_names = nullptr);
+
+absl::Status ConvertTfMlirToRuntimeExecutable(
+    const TfrtCompileOptions& options, mlir::ModuleOp module,
+    absl::FunctionRef<
+        absl::Status(mlir::PassManager&, mlir::ModuleOp,
+                     const tensorflow::TfrtPipelineOptions& options)>
+        emit_executable,
+    tfrt_stub::ModelRuntimeContext& model_context,
+    tfrt_stub::FallbackState* fallback_state = nullptr,
+    std::vector<std::string>* added_xla_function_names = nullptr);
+
+std::unique_ptr<tensorflow::TfrtPipelineOptions> GetTfrtPipelineOptions(
+    const TfrtCompileOptions& options);
+
+// Adds MLIR functions for XLA clusters to the function library.
+absl::Status AddXlaFunctions(
+    tfrt_stub::FallbackState* fallback_state, mlir::ModuleOp mlir_module,
+    std::vector<std::string>* added_xla_function_names = nullptr);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSLATE_IMPORT_MODEL_H_

tensorflow/compiler/mlir/tfrt/translate/import_model_test.cc

@@ -0,0 +1,123 @@
+/* Copyright 2024 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/import_model.h"
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/compiler/mlir/tfrt/translate/tfrt_compile_options.h"
+
+namespace tensorflow {
+namespace {
+
+using ::testing::SizeIs;
+
+TEST(GetTfrtPipelineOptions, BatchPaddingPolicy) {
+  tensorflow::TfrtCompileOptions options;
+  options.batch_padding_policy = "PAD_TEST_OPTION";
+  auto pipeline_options = GetTfrtPipelineOptions(options);
+  EXPECT_EQ(pipeline_options->batch_padding_policy, "PAD_TEST_OPTION");
+}
+
+TEST(GetTfrtPipelineOptions, NumBatchThreads) {
+  tensorflow::TfrtCompileOptions options;
+  options.batch_options.set_num_batch_threads(2);
+  auto pipeline_options = GetTfrtPipelineOptions(options);
+  EXPECT_EQ(pipeline_options->num_batch_threads, 2);
+}
+
+TEST(GetTfrtPipelineOptions, MaxBatchSize) {
+  tensorflow::TfrtCompileOptions options;
+  options.batch_options.set_max_batch_size(8);
+  auto pipeline_options = GetTfrtPipelineOptions(options);
+  EXPECT_EQ(pipeline_options->max_batch_size, 8);
+}
+
+TEST(GetTfrtPipelineOptions, BatchTimeoutMicros) {
+  tensorflow::TfrtCompileOptions options;
+  options.batch_options.set_batch_timeout_micros(5000);
+  auto pipeline_options = GetTfrtPipelineOptions(options);
+  EXPECT_EQ(pipeline_options->batch_timeout_micros, 5000);
+}
+
+TEST(GetTfrtPipelineOptions, AllowedBatchSizes) {
+  tensorflow::TfrtCompileOptions options;
+  options.batch_options.add_allowed_batch_sizes(2);
+  options.batch_options.add_allowed_batch_sizes(4);
+  options.batch_options.add_allowed_batch_sizes(8);
+  auto pipeline_options = GetTfrtPipelineOptions(options);
+  EXPECT_THAT(pipeline_options->allowed_batch_sizes, SizeIs(3));
+}
+
+TEST(GetTfrtPipelineOptions, MaxEnqueuedBatches) {
+  tensorflow::TfrtCompileOptions options;
+  options.batch_options.set_max_enqueued_batches(250);
+  auto pipeline_options = GetTfrtPipelineOptions(options);
+  EXPECT_EQ(pipeline_options->max_enqueued_batches, 250);
+}
+
+TEST(GetTfrtPipelineOptions, LowPriorityMaxBatchSize) {
+  tensorflow::TfrtCompileOptions options;
+  options.batch_options.set_low_priority_max_batch_size(16);
+  auto pipeline_options = GetTfrtPipelineOptions(options);
+  EXPECT_EQ(pipeline_options->low_priority_max_batch_size, 16);
+}
+
+TEST(GetTfrtPipelineOptions, LowPriorityBatchTimeoutMicros) {
+  tensorflow::TfrtCompileOptions options;
+  options.batch_options.set_low_priority_batch_timeout_micros(10000);
+  auto pipeline_options = GetTfrtPipelineOptions(options);
+  EXPECT_EQ(pipeline_options->low_priority_batch_timeout_micros, 10000);
+}
+
+TEST(GetTfrtPipelineOptions, LowPriorityAllowedBatchSizes) {
+  tensorflow::TfrtCompileOptions options;
+  options.batch_options.add_low_priority_allowed_batch_sizes(4);
+  options.batch_options.add_low_priority_allowed_batch_sizes(8);
+  options.batch_options.add_low_priority_allowed_batch_sizes(16);
+  auto pipeline_options = GetTfrtPipelineOptions(options);
+  EXPECT_THAT(pipeline_options->low_priority_allowed_batch_sizes, SizeIs(3));
+}
+
+TEST(GetTfrtPipelineOptions, LowPriorityMaxEnqueuedBatches) {
+  tensorflow::TfrtCompileOptions options;
+  options.batch_options.set_low_priority_max_enqueued_batches(100);
+  auto pipeline_options = GetTfrtPipelineOptions(options);
+  EXPECT_EQ(pipeline_options->low_priority_max_enqueued_batches, 100);
+}
+
+TEST(GetTfrtPipelineOptions, NumWarmupBatchThreads) {
+  tensorflow::TfrtCompileOptions options;
+  options.batch_options.set_num_warmup_batch_threads(4);
+  auto pipeline_options = GetTfrtPipelineOptions(options);
+  EXPECT_EQ(pipeline_options->num_warmup_batch_threads, 4);
+}
+
+TEST(GetTfrtPipelineOptions, EnableLargeBatchSplitting) {
+  tensorflow::TfrtCompileOptions options;
+  options.batch_options.set_enable_large_batch_splitting(true);
+  auto pipeline_options = GetTfrtPipelineOptions(options);
+  EXPECT_TRUE(pipeline_options->enable_large_batch_splitting);
+}
+
+TEST(GetTfrtPipelineOptions, MixedPriorityBatchingPolicy) {
+  tensorflow::TfrtCompileOptions options;
+  options.batch_options.set_mixed_priority_batching_policy("priority_merge");
+  auto pipeline_options = GetTfrtPipelineOptions(options);
+  EXPECT_EQ(pipeline_options->mixed_priority_batching_policy, "priority_merge");
+}
+
+}  // namespace
+}  // namespace tensorflow

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
+}
+

tensorflow/compiler/mlir/tfrt/translate/tfrt_compile_options.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/translate/tfrt_compile_options.h"
+
+#include <ostream>
+
+#include "absl/strings/str_join.h"
+
+namespace tensorflow {
+
+std::ostream& operator<<(std::ostream& os,
+                         TfrtDeviceInfraTarget device_target) {
+  switch (device_target) {
+    case TfrtDeviceInfraTarget::kCpu:
+      return os << "Cpu";
+    case TfrtDeviceInfraTarget::kTpurt:
+      return os << "Tpurt";
+    case TfrtDeviceInfraTarget::kTfFallback:
+      return os << "TfFallback";
+    case TfrtDeviceInfraTarget::kBridgeFallback:
+      return os << "BridgeFallback";
+    case TfrtDeviceInfraTarget::kGpu:
+      return os << "Gpu";
+  }
+}
+
+std::ostream& operator<<(std::ostream& os, const TfrtCompileOptions& options) {
+  return os << "{" << "variable_device = " << options.variable_device
+            << ", default_device = " << options.default_device
+            << ", enable_optimizer = " << options.enable_optimizer
+            << ", enable_grappler = " << options.enable_grappler
+            << ", force_data_format = " << options.force_data_format
+            << ", device_target = " << options.device_target
+            << ", tpu_fuse_ops = " << options.tpu_fuse_ops
+            << ", tpu_move_resource_gather_to_host = "
+            << options.tpu_move_resource_gather_to_host
+            << ", tpu_gather_table_width_threshold_bytes = "
+            << options.tpu_gather_table_width_threshold_bytes
+            << ", use_tpu_host_allocator_for_inputs = "
+            << options.use_tpu_host_allocator_for_inputs
+            << ", hoist_invariant_ops = " << options.hoist_invariant_ops
+            << ", enable_while_parallel_iterations = "
+            << options.enable_while_parallel_iterations
+            << ", cost_threshold = " << options.cost_threshold
+            << ", min_num_batch_threads = " << options.min_num_batch_threads
+            << ", min_max_enqueued_batches = "
+            << options.min_max_enqueued_batches
+            << ", batch_padding_policy = " << options.batch_padding_policy
+            << ", num_batch_threads = "
+            << options.batch_options.num_batch_threads()
+            << ", max_batch_size = " << options.batch_options.max_batch_size()
+            << ", batch_timeout_micros = "
+            << options.batch_options.batch_timeout_micros()
+            << ", allowed_batch_sizes = "
+            << absl::StrJoin(options.batch_options.allowed_batch_sizes(), ",")
+            << ", max_enqueued_batches = "
+            << options.batch_options.max_enqueued_batches()
+            << ", low_priority_max_batch_size = "
+            << options.batch_options.low_priority_max_batch_size()
+            << ", low_priority_batch_timeout_micros = "
+            << options.batch_options.low_priority_batch_timeout_micros()
+            << ", low_priority_allowed_batch_sizes = "
+            << absl::StrJoin(
+                   options.batch_options.low_priority_allowed_batch_sizes(),
+                   ",")
+            << ", low_priority_max_enqueued_batches = "
+            << options.batch_options.low_priority_max_enqueued_batches()
+            << ", num_warmup_batch_threads = "
+            << options.batch_options.num_warmup_batch_threads()
+            << ", enable_large_batch_splitting = "
+            << options.batch_options.enable_large_batch_splitting()
+            << ", mixed_priority_batching_policy = "
+            << options.batch_options.mixed_priority_batching_policy()
+            << ", merge_inter_dependent_streams = "
+            << options.merge_inter_dependent_streams
+            << ", decompose_resource_ops = " << options.decompose_resource_ops
+            << ", compile_to_sync_tfrt_dialect = "
+            << options.compile_to_sync_tfrt_dialect
+            << ", batch_queue_global_prioritization_num_threads = "
+            << options.batch_queue_global_prioritization_num_threads << "}";
+}
+
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/translate/tfrt_compile_options.h

@@ -0,0 +1,210 @@
+/* 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_TFRT_COMPILE_OPTIONS_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSLATE_TFRT_COMPILE_OPTIONS_H_
+
+#include <cstdint>
+#include <iosfwd>
+#include <ostream>
+#include <string>
+#include <vector>
+
+#include "tensorflow/core/protobuf/config.pb.h"
+
+namespace tensorflow {
+
+class BackendCompiler;
+
+enum class TfrtDeviceInfraTarget {
+  kCpu,             // CPU only, no device support.
+  kTpurt,           // Target TPURT dialect and kernels.
+  kTfFallback,      // Target TPU kernels in TF Fallback.
+  kBridgeFallback,  // TPU support but choose kTpurt or kTfFallback depending on
+                    // whether the graph has unsupported feature in Bridge.
+  kGpu,             // Target GPU specific compiler passes and runtime
+                    // initializations.
+};
+
+std::ostream& operator<<(std::ostream& os, TfrtDeviceInfraTarget device_target);
+
+struct TfrtCompileOptions {
+  std::string saved_model_dir;
+  // TODO(tfrt-devs): Ideally, compiler should make the decision where
+  // to place the variable.
+  std::string variable_device = "/job:localhost/replica:0/task:0/device:CPU:0";
+  std::string default_device = "/job:localhost/replica:0/task:0/device:CPU:0";
+
+  // Enable compiler optimization in TFRT dialect.
+  bool enable_optimizer = true;
+
+  // If true, run grappler passes before compiling.
+  bool enable_grappler = true;
+
+  // Graph rewrite options that will be applied on GraphDef before converting to
+  // MLIR.
+  GraphOptions graph_options;
+
+  // Force data format for all layout sensitive operations, eg. setting it to
+  // "NHWC" will changes all data format in the graph to "NHWC" by inserting
+  // or removing related tf.Transpose op. Currently the supported formats are
+  // "NHWC" and "NCHW".
+  //
+  // TODO(tfrt-devs): Ideally compiler should figure out whether the
+  // data format should be changed, instead of controlled by users.
+  std::string force_data_format;
+
+  // The target device infrastructure to use. This will trigger target specific
+  // compiler passes and runtime initialization.
+  TfrtDeviceInfraTarget device_target = TfrtDeviceInfraTarget::kCpu;
+
+  // The custom compiler for device compilation. Instead of using the enum above
+  // to choose predefined device target, users can use this `backend_compiler`
+  // to inject their customized implementation.
+  BackendCompiler* backend_compiler = nullptr;
+
+  // If true, use the fused TPU compile_and_execute kernel, which performs all
+  // TPU inference related operations, e.g. core selection, h2d/d2h transfers,
+  // compile and execute.
+  bool tpu_fuse_ops = false;
+
+  // If true, resource gather ops in the device graph are moved to host graphs
+  // in order to saved TPU memory usage. This option is experimental.
+  bool tpu_move_resource_gather_to_host = false;
+
+  // The threshold in bytes that controls whether a resource gather op on TPU
+  // should be moved to host. A negative value means there is no threshold. This
+  // option is experimental.
+  int64_t tpu_gather_table_width_threshold_bytes = -1;
+
+  // If true, fallback executeops that produce inputs to tpu program will use
+  // tpu host allocator. This options is experimental.
+  bool use_tpu_host_allocator_for_inputs = false;
+
+  // To allow unpadded batch for TPU execution.
+  enum class TpuAllowUnpaddedBatch {
+    // Disable this feature.
+    kDisabled,
+    // Enable this feature when in-graph batching is detected.
+    kAuto,
+    // Force to enable this feature.
+    kEnforced,
+  };
+  TpuAllowUnpaddedBatch tpu_allow_unpadded_batch =
+      TpuAllowUnpaddedBatch::kDisabled;
+
+  // If true, the compiler will try to hoist invariant ops (e.g., const ops and
+  // their non-side-effecting consumers) to loading phase, which avoids the
+  // runtime cost during later running.
+  // TODO(tfrt-devs): Set the default value to true after testing as it is
+  // supposed to be turned on by default.
+  bool hoist_invariant_ops = false;
+
+  // If true, get_resource_op will be fused during hoisting.
+  bool fuse_get_resource_ops_in_hoisting = true;
+
+  // If true, the compiler will try to sink in the invariant ops (e.g. const
+  // ops, var handle ops, etc.) to the nested function (e.g. batch function) to
+  // facilitate invariant ops hoisting.
+  // TODO(tfrt-devs): Set the default value to true after testing as it is
+  // supposed to be turned on by default.
+  bool sink_in_invariant_ops = false;
+
+  // This flag behaves differently for TFRT and MLRT.
+  // For TFRT, if true, tf.While's iterations will be parallelized on a
+  // best-effort basis. This is currently experimental. MLRT attempts to convert
+  // tf.while to tf_mlrt.map_fn regardless of this flag. For tf.While that
+  // cannot be converted tf_mlrt.map_fn, MLRT try to parallelize tf.while's
+  // iterations on a best-effort basis.
+  bool enable_while_parallel_iterations = false;
+
+  // The cost threshold to decide whether a sequence of operations is cheap, and
+  // then whether it can be executed inline. If the cost is smaller than the
+  // threshold, it will be considered as cheap operations. Since the cost must
+  // be positive integers, setting the threshold to 1 makes all operations
+  // expensive.
+  uint64_t cost_threshold = 1;
+
+  // The minimum number of batch threads. This number provides a lower bound on
+  // the number of batch threads on top of what is specified in the model. If
+  // the number of batch threads is too small (e.g. smaller than the number of
+  // parallel hardware accelerator available), it can lead to under utilization
+  // of resources.
+  int64_t min_num_batch_threads = 1;
+
+  // The minimum of the maximum number of enqueued batches. This number provides
+  // a lower bound on top of what is specified in the model. If the number of
+  // max_enqueued_batches is too small, it can lead to under utilization of
+  // resources.
+  int64_t min_max_enqueued_batches = 1;
+
+  // If non-zero, all models on this server are switched to use a prioritized
+  // batching function using this number of global threads.
+  int64_t batch_queue_global_prioritization_num_threads = 0;
+
+  // If true, the queue implementation will have a separate subqueue for each
+  // criticality.
+  bool enable_priority_aware_batch_scheduler = false;
+
+  // If true, the queue implementation will resplit tasks for priority aware
+  // scheduling.
+  bool enable_priority_aware_batch_scheduler_resplit = false;
+
+  // If true, enable lazy cancellation filtering in the priority-aware batch
+  // scheduler.
+  bool enable_batching_task_lazy_cancellation = false;
+
+  // The policy used by a BatchScheduler to pad (or split) batches.
+  std::string batch_padding_policy;
+
+  // Batching parameters to be rewritten in the existing BatchFunction ops.
+  BatchingOptions batch_options;
+
+  // If true, streams with inter data dependencies will be preferred to be
+  // merged for inline execution.
+  bool merge_inter_dependent_streams = true;
+
+  // Whether to enable the DecomposeResourceOpsPass.
+  bool decompose_resource_ops = true;
+
+  // Whether to compile to sync TFRT dialect.
+  bool compile_to_sync_tfrt_dialect = false;
+
+  // Whether to use gpurt.compile_and_execute for GPU.
+  // TODO(b/294895431): Remove the flag and default to the fused op.
+  bool use_gpu_compile_and_execute_op = false;
+
+  // If true, MLIR module will be serialized to aot_packages.
+  bool serialize_mlir_module_to_aot_packages = false;
+
+  // Serialized MLIR module file under aot_packages.
+  std::string aot_mlir_module_file;
+
+  // If true, BEF will be serialized to aot_packages.
+  bool serialize_bef_to_aot_packages = false;
+
+  // Serialized BEF file under aot_packages.
+  std::string aot_bef_file;
+  // If true, use XLA:CPU for CPU computations.
+  bool allow_xla_cpu = true;
+  // If true, enable asynchronous IFRT execution.
+  bool enable_async_ifrt = false;
+};
+
+std::ostream& operator<<(std::ostream& os, const TfrtCompileOptions& options);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSLATE_TFRT_COMPILE_OPTIONS_H_