chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/tfrt/utils/export.cc

@@ -0,0 +1,75 @@
+/* Copyright 2023 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/utils/export.h"
+
+#include <memory>
+#include <utility>
+
+#include "absl/functional/any_invocable.h"
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
+#include "tensorflow/compiler/mlir/tf2xla/api/v1/tf_dialect_to_executor.h"
+#include "tensorflow/compiler/mlir/tf2xla/api/v2/tf_executor_to_graph.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tsl/platform/errors.h"
+#include "tsl/profiler/lib/traceme.h"
+
+namespace tensorflow {
+
+absl::Status ExportFunctionDefs(
+    mlir::ModuleOp module,
+    absl::AnyInvocable<absl::Status(tensorflow::FunctionDef)> callback,
+    bool export_tf_original_func_name) {
+  tsl::profiler::TraceMe traceme([&]() {
+    return tsl::profiler::TraceMeEncode(
+        "ExportFunctionDefs",
+        {{"module_name", absl::string_view(module.getName().value_or("?"))}});
+  });
+
+  TF_RETURN_IF_ERROR(
+      tensorflow::tf2xla::v1::ExportFromTensorflowDialectToExecutor(module));
+
+  {
+    mlir::StatusScopedDiagnosticHandler diag_handler(module.getContext());
+
+    mlir::PassManager pm(module.getContext());
+    pm.addPass(mlir::CreateBreakUpIslandsPass());
+
+    if (mlir::failed(pm.run(module))) {
+      return diag_handler.ConsumeStatus();
+    }
+  }
+  tensorflow::GraphExportConfig configs;
+  configs.export_original_tf_func_name = export_tf_original_func_name;
+
+  for (auto func : module.getOps<mlir::func::FuncOp>()) {
+    tensorflow::FunctionDef function_def;
+    TF_RETURN_IF_ERROR(
+        tensorflow::tf2xla::v2::ConvertMlirFunctionToFunctionLibraryDef(
+            func, configs, &function_def));
+    TF_RETURN_IF_ERROR(callback(std::move(function_def)));
+  }
+
+  return absl::OkStatus();
+}
+
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/utils/export.h

@@ -0,0 +1,36 @@
+/* Copyright 2023 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_UTILS_EXPORT_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_UTILS_EXPORT_H_
+
+
+#include "absl/functional/any_invocable.h"
+#include "absl/status/status.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "tensorflow/core/framework/function.pb.h"
+
+namespace tensorflow {
+
+// Exports every function in `module` into `tensorflow.FunctionDef` and calls
+// `callback` for each `tensorflow.FunctionDef`. Modifies `module` in place to
+// be suitable for FunctionDef export.
+absl::Status ExportFunctionDefs(
+    mlir::ModuleOp module,
+    absl::AnyInvocable<absl::Status(tensorflow::FunctionDef)> callback,
+    bool export_tf_original_func_name = true);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_UTILS_EXPORT_H_

tensorflow/compiler/mlir/tfrt/utils/host_context.cc

@@ -0,0 +1,56 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/tfrt/utils/host_context.h"
+
+#include <cstdint>
+#include <memory>
+
+#include "absl/log/log.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tfrt/host_context/concurrent_work_queue.h"  // from @tf_runtime
+#include "tfrt/host_context/diagnostic.h"  // from @tf_runtime
+#include "tfrt/host_context/host_allocator.h"  // from @tf_runtime
+#include "tfrt/host_context/host_context.h"  // from @tf_runtime
+
+namespace tensorflow {
+
+using ::tfrt::HostContext;
+
+const char* const kDefaultHostDeviceName =
+    "/job:localhost/replica:0/task:0/device:CPU:0";
+
+std::unique_ptr<HostContext> CreateSingleThreadedHostContext() {
+  return std::make_unique<HostContext>(
+      [](const tfrt::DecodedDiagnostic& diag) {
+        LOG(FATAL) << "Runtime error: " << diag.message() << "\n";
+      },
+      tfrt::CreateMallocAllocator(), tfrt::CreateSingleThreadedWorkQueue(),
+      kDefaultHostDeviceName);
+}
+
+std::unique_ptr<HostContext> CreateMultiThreadedHostContext(
+    int64_t num_threads) {
+  return std::make_unique<HostContext>(
+      [](const tfrt::DecodedDiagnostic& diag) {
+        LOG(FATAL) << "Runtime error: " << diag.message() << "\n";
+      },
+      tfrt::CreateMallocAllocator(),
+      tfrt::CreateMultiThreadedWorkQueue(num_threads,
+                                         /*num_blocking_threads=*/1),
+      kDefaultHostDeviceName);
+}
+
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/utils/host_context.h

@@ -0,0 +1,36 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_UTILS_HOST_CONTEXT_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_UTILS_HOST_CONTEXT_H_
+
+#include <cstdint>
+#include <memory>
+
+#include "absl/base/attributes.h"
+#include "tfrt/host_context/host_context.h"  // from @tf_runtime
+
+namespace tensorflow {
+
+// The name of the default host device for running fallback kernels.
+ABSL_CONST_INIT extern const char* const kDefaultHostDeviceName;
+
+std::unique_ptr<tfrt::HostContext> CreateSingleThreadedHostContext();
+std::unique_ptr<tfrt::HostContext> CreateMultiThreadedHostContext(
+    int64_t num_threads);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_UTILS_HOST_CONTEXT_H_