chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/tf2xla/internal/utils/BUILD

@@ -0,0 +1,79 @@
+load("//tensorflow:tensorflow.bzl", "tf_cc_test")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        "//tensorflow/compiler/mlir/tf2xla/internal:__subpackages__",
+    ],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "dialect_detection_utils",
+    srcs = [
+        "dialect_detection_utils.cc",
+    ],
+    hdrs = [
+        "dialect_detection_utils.h",
+    ],
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/core:framework",
+        "//tensorflow/core/transforms/toposort:Pass",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:Transforms",
+    ],
+)
+
+tf_cc_test(
+    name = "dialect_detection_utils_test",
+    srcs = ["dialect_detection_utils_test.cc"],
+    deps = [
+        ":dialect_detection_utils",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@stablehlo//:chlo_ops",
+    ],
+)
+
+cc_library(
+    name = "test_metadata_config",
+    testonly = True,
+    srcs = ["test_metadata_config.cc"],
+    hdrs = ["test_metadata_config.h"],
+    visibility = [
+        "//tensorflow/compiler/mlir/tf2xla/api:__subpackages__",
+    ],
+    deps = [
+        "//tensorflow/compiler/jit",
+        "//tensorflow/compiler/jit:xla_tpu_device",
+        "//tensorflow/compiler/mlir:register_common_dialects",
+        "//tensorflow/compiler/mlir/tensorflow:convert_type",
+        "//tensorflow/compiler/mlir/tensorflow:deserialize_mlir_module_utils",
+        "//tensorflow/compiler/tf2xla:layout_util",
+        "//tensorflow/compiler/tf2xla:xla_compiler",
+        "//tensorflow/compiler/tf2xla:xla_helpers",
+        "//tensorflow/compiler/tf2xla:xla_tpu_backend_registration",
+        "//tensorflow/compiler/tf2xla/kernels:xla_ops",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
+        "//tensorflow/core/tpu/kernels/xla:host_compute_ops",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings:string_view",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@xla//xla:shape_util",
+        "@xla//xla/hlo/translate/mhlo_to_hlo:type_to_shape",
+        "@xla//xla/mlir_hlo:hlo_dialect_registration",
+        "@xla//xla/tsl/platform:errors",
+    ],
+)

tensorflow/compiler/mlir/tf2xla/internal/utils/dialect_detection_utils.cc

@@ -0,0 +1,45 @@
+/* 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/tf2xla/internal/utils/dialect_detection_utils.h"
+
+#include <set>
+#include <string>
+
+#include "mlir/IR/Dialect.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Visitors.h"  // from @llvm-project
+
+namespace tensorflow {
+namespace tf2xla {
+namespace internal {
+
+bool IsInBridgeAcceptableDialects(mlir::Operation* op) {
+  const std::set<std::string> kBuiltinNamespaces = {"func", "return",
+                                                    "builtin"};
+  const std::set<std::string> kBridgeAcceptableNamespaces = {"tf", "tf_device"};
+  bool isInDefaulNamespaces =
+      kBuiltinNamespaces.find(op->getDialect()->getNamespace().str()) !=
+      kBuiltinNamespaces.end();
+  bool isInBridgeAcceptableNamespaces =
+      kBridgeAcceptableNamespaces.find(
+          op->getDialect()->getNamespace().str()) !=
+      kBridgeAcceptableNamespaces.end();
+  return isInDefaulNamespaces || isInBridgeAcceptableNamespaces;
+}
+
+}  // namespace internal
+}  // namespace tf2xla
+}  // namespace tensorflow

tensorflow/compiler/mlir/tf2xla/internal/utils/dialect_detection_utils.h

@@ -0,0 +1,33 @@
+/* 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_TF2XLA_INTERNAL_UTILS_DIALECT_DETECTION_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_UTILS_DIALECT_DETECTION_UTILS_H_
+
+#include "mlir/IR/Operation.h"  // from @llvm-project
+
+namespace tensorflow {
+namespace tf2xla {
+namespace internal {
+
+// Returns true if the op has a valid namespace during clustering & tf dialect
+// to executor components of the Bridge.
+bool IsInBridgeAcceptableDialects(mlir::Operation* op);
+
+}  // namespace internal
+}  // namespace tf2xla
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_UTILS_DIALECT_DETECTION_UTILS_H_

tensorflow/compiler/mlir/tf2xla/internal/utils/dialect_detection_utils_test.cc

@@ -0,0 +1,76 @@
+/* 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/tf2xla/internal/utils/dialect_detection_utils.h"
+
+#include <gtest/gtest.h>
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "stablehlo/dialect/ChloOps.h"  // from @stablehlo
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
+
+namespace tensorflow {
+namespace tf2xla {
+namespace internal {
+
+namespace {
+
+using mlir::MLIRContext;
+using mlir::OpBuilder;
+using mlir::Operation;
+using mlir::OperationState;
+using mlir::UnknownLoc;
+using mlir::chlo::ChloDialect;
+using mlir::TF::TensorFlowDialect;
+using tensorflow::tf2xla::internal::IsInBridgeAcceptableDialects;
+
+class SharedUtilsTest : public ::testing::Test {};
+
+TEST_F(SharedUtilsTest, IsInFunctionalDialectPasses) {
+  MLIRContext context;
+  context.loadDialect<TensorFlowDialect>();
+  OpBuilder opBuilder(&context);
+  OperationState state(UnknownLoc::get(opBuilder.getContext()),
+                       /*OperationName=*/"tf.Const");
+  mlir::Operation* op = Operation::create(state);
+
+  bool result = IsInBridgeAcceptableDialects(op);
+
+  EXPECT_TRUE(result);
+  op->destroy();
+}
+
+TEST_F(SharedUtilsTest, IsInFunctionalDialectFails) {
+  MLIRContext context;
+  context.loadDialect<ChloDialect>();
+  OpBuilder opBuilder(&context);
+  OperationState state(UnknownLoc::get(opBuilder.getContext()),
+                       /*OperationName=*/"chlo.broadcast_add");
+  Operation* op = Operation::create(state);
+
+  bool result = IsInBridgeAcceptableDialects(op);
+
+  EXPECT_FALSE(result);
+  op->destroy();
+}
+
+}  // namespace
+}  // namespace internal
+}  // namespace tf2xla
+}  // namespace tensorflow

tensorflow/compiler/mlir/tf2xla/internal/utils/test_metadata_config.cc

@@ -0,0 +1,107 @@
+/* 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/tf2xla/internal/utils/test_metadata_config.h"
+
+#include <vector>
+
+#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/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/register_common_dialects.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/convert_type.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/deserialize_mlir_module_utils.h"
+#include "xla/hlo/translate/mhlo_to_hlo/type_to_shape.h"
+#include "xla/mlir_hlo/mhlo/IR/register.h"
+#include "xla/shape.h"
+#include "xla/tsl/platform/errors.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
+
+namespace tensorflow {
+namespace tf2xla {
+namespace internal {
+namespace {
+
+constexpr char kEntryFuncName[] = "main";
+
+absl::Status SetupArguments(mlir::ModuleOp module,
+                            std::vector<TensorShape>& arg_shapes,
+                            tpu::TPUCompileMetadataProto& metadata_proto) {
+  auto main_fn = module.lookupSymbol<mlir::func::FuncOp>(kEntryFuncName);
+  if (!main_fn) {
+    return absl::InternalError("Could not find main function in MLIR Module.");
+  }
+
+  mlir::FunctionType func_type = main_fn.getFunctionType();
+  for (auto input_type : func_type.getInputs()) {
+    tensorflow::TensorShape tensor_shape;
+    xla::Shape xla_shape = xla::TypeToShape(input_type);
+    TF_RETURN_IF_ERROR(tensorflow::TensorShape::BuildTensorShape(
+        xla_shape.dimensions(), &tensor_shape));
+    arg_shapes.emplace_back(tensor_shape);
+
+    DataType dtype;
+    TF_RETURN_IF_ERROR(ConvertToDataType(input_type, &dtype));
+
+    auto metadata_arg = metadata_proto.add_args();
+    metadata_arg->set_kind(tpu::TPUCompileMetadataProto::Arg::PARAMETER);
+    metadata_arg->set_dtype(dtype);
+  }
+
+  return absl::OkStatus();
+}
+
+absl::Status SetupReturnValues(mlir::ModuleOp module,
+                               tpu::TPUCompileMetadataProto& metadata_proto) {
+  auto main_fn = module.lookupSymbol<mlir::func::FuncOp>(kEntryFuncName);
+  if (!main_fn) {
+    return absl::InternalError("Could not find main function in MLIR Module.");
+  }
+
+  int func_results = main_fn.getFunctionType().getNumResults();
+  for (int i = 0; i < func_results; i++) {
+    metadata_proto.add_retvals();
+  }
+
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+absl::Status ConfigureMetadata(absl::string_view mlir_module_str,
+                               std::vector<TensorShape>& arg_shapes,
+                               tpu::TPUCompileMetadataProto& metadata_proto) {
+  mlir::DialectRegistry registry;
+  mlir::RegisterCommonToolingDialects(registry);
+  mlir::MLIRContext context(registry);
+  mlir::OwningOpRef<mlir::ModuleOp> mlir_module;
+
+  TF_RETURN_IF_ERROR(
+      DeserializeMlirModule(mlir_module_str, &context, &mlir_module));
+  TF_RETURN_IF_ERROR(SetupReturnValues(*mlir_module, metadata_proto));
+  TF_RETURN_IF_ERROR(SetupArguments(*mlir_module, arg_shapes, metadata_proto));
+
+  return absl::OkStatus();
+}
+
+}  // namespace internal
+}  // namespace tf2xla
+}  // namespace tensorflow

tensorflow/compiler/mlir/tf2xla/internal/utils/test_metadata_config.h

@@ -0,0 +1,41 @@
+/* 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_UTILS_TEST_METADATA_CONFIG_H_
+#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_UTILS_TEST_METADATA_CONFIG_H_
+
+#include <variant>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
+
+namespace tensorflow {
+namespace tf2xla {
+namespace internal {
+
+// Fills in arg_shapes and metadata_proto with appropriate values based on the
+// input mlir module.
+absl::Status ConfigureMetadata(absl::string_view mlir_module_str,
+                               std::vector<TensorShape>& arg_shapes,
+                               tpu::TPUCompileMetadataProto& metadata_proto);
+
+}  // namespace internal
+}  // namespace tf2xla
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_UTILS_TEST_METADATA_CONFIG_H_