wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 6 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/compiler/mlir/tf2xla/BUILD
+37
View File
@@ -0,0 +1,37 @@
# Description:
# TF2XLA Bridge and related components.
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package_group(
name = "tensorflow_mlir_tf2xla",
packages = [
"//tensorflow/compiler/mlir/tf2xla/...",
],
)
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//visibility:public"],
licenses = ["notice"],
)
cc_library(
name = "mlir_bridge_rollout_policy",
srcs = ["mlir_bridge_rollout_policy.cc"],
hdrs = ["mlir_bridge_rollout_policy.h"],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/compiler/jit:flags",
"//tensorflow/core:framework",
"//tensorflow/core:graph",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/types:optional",
"@llvm-project//mlir:IR",
],
)
alias(
name = "compile_mlir_util",
actual = "//tensorflow/compiler/mlir/tf2xla/api/v1:compile_mlir_util_no_tf_dialect_passes",
)
tensorflow/compiler/mlir/tf2xla/api/v1/BUILD
+336
View File
@@ -0,0 +1,336 @@
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow/core/platform:build_config.bzl", "tf_additional_all_protos", "tf_proto_library")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow/compiler/mlir/tf2xla/api:__subpackages__",
"//tensorflow/compiler/mlir/tf2xla/internal:__subpackages__",
],
)
cc_library(
name = "compile_mlir_util_no_tf_dialect_passes",
srcs = ["compile_mlir_util.cc"],
hdrs = ["compile_mlir_util.h"],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/compiler/mlir/quantization/stablehlo:bridge_passes",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:bridge_logger",
"//tensorflow/compiler/mlir/tensorflow:convert_tensor",
"//tensorflow/compiler/mlir/tensorflow:convert_type",
"//tensorflow/compiler/mlir/tensorflow:deserialize_mlir_module_utils",
"//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
"//tensorflow/compiler/mlir/tensorflow:dynamic_shape_utils",
"//tensorflow/compiler/mlir/tensorflow:error_util",
"//tensorflow/compiler/mlir/tensorflow:import_model",
"//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_flags",
"//tensorflow/compiler/mlir/tensorflow:serialize_mlir_module_utils",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow:translate_utils",
"//tensorflow/compiler/mlir/tensorflow:xla_sharding_util",
"//tensorflow/compiler/mlir/tensorflow/transforms:shape_inference_pass",
"//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
"//tensorflow/compiler/mlir/tf2xla:mlir_bridge_rollout_policy",
"//tensorflow/compiler/mlir/tf2xla/api/v2:graph_to_tf_executor",
"//tensorflow/compiler/mlir/tf2xla/internal:mlir_pass_instrumentation",
"//tensorflow/compiler/mlir/tf2xla/internal/passes:lowering_passes",
"//tensorflow/compiler/mlir/tf2xla/transforms:xla_legalize_tf_with_tf2xla",
"//tensorflow/compiler/tf2xla:common",
"//tensorflow/compiler/tf2xla:layout_util",
"//tensorflow/compiler/tf2xla:xla_argument",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/common_runtime:core_cpu_internal",
"//tensorflow/core/platform:error_payloads",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:logging",
"//tensorflow/core/platform:status",
"//tensorflow/core/tpu:tpu_defs",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/log:vlog_is_on",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:string_view",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:ShapeDialect",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TensorDialect",
"@llvm-project//mlir:Transforms",
"@stablehlo//:base",
"@xla//xla:shape_util",
"@xla//xla:xla_data_proto_cc",
"@xla//xla/hlo/builder:xla_computation",
"@xla//xla/hlo/ir:hlo",
"@xla//xla/hlo/translate:register",
"@xla//xla/hlo/translate/mhlo_to_hlo:layout_util",
"@xla//xla/hlo/translate/mhlo_to_hlo:mlir_hlo_to_hlo",
"@xla//xla/hlo/translate/mhlo_to_hlo:type_to_shape",
"@xla//xla/mlir_hlo",
"@xla//xla/mlir_hlo:mhlo_passes",
"@xla//xla/mlir_hlo:stablehlo_extension_passes",
"@xla//xla/service:hlo_proto_cc",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:statusor",
],
)
tf_cc_test(
name = "compile_mlir_util_test",
srcs = ["compile_mlir_util_test.cc"],
deps = [
":compile_mlir_util_no_tf_dialect_passes",
"//tensorflow/compiler/jit:xla_compile_util",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:serialize_mlir_module_utils",
"//tensorflow/compiler/mlir/tf2xla/internal:test_matchers",
"//tensorflow/compiler/tf2xla:xla_compiler",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/framework:fake_input",
"//tensorflow/core/lib/monitoring:cell_reader",
"//tensorflow/core/platform:types",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest_main",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@xla//xla:shape_util",
"@xla//xla:xla_data_proto_cc",
"@xla//xla/hlo/builder:xla_builder",
"@xla//xla/hlo/builder:xla_computation",
"@xla//xla/tsl/lib/core:status_test_util",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:statusor",
],
)
cc_library(
name = "compile_tf_graph",
srcs = ["compile_tf_graph.cc"],
hdrs = ["compile_tf_graph.h"],
deps = [
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:deserialize_mlir_module_utils",
"//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
"//tensorflow/compiler/mlir/tensorflow:error_util",
"//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_flags",
"//tensorflow/compiler/mlir/tensorflow:translate_utils",
"//tensorflow/compiler/mlir/tensorflow/transforms:set_tpu_infeed_layout",
"//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
"//tensorflow/compiler/mlir/tf2xla/api/v2:tf_executor_to_graph",
"//tensorflow/compiler/mlir/tf2xla/internal:logging_hooks",
"//tensorflow/compiler/tf2xla:layout_util",
"//tensorflow/compiler/tf2xla:xla_compiler",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
"//tensorflow/core/tpu:tpu_compile",
"//tensorflow/core/tpu/kernels:tpu_compile_op_support",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:vlog_is_on",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:variant",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@xla//xla:shape_util",
"@xla//xla:status_macros",
"@xla//xla/client:compile_only_client",
"@xla//xla/hlo/ir:hlo",
"@xla//xla/mlir_hlo:hlo_dialect_registration",
"@xla//xla/pjrt/proto:compile_options_proto_cc",
"@xla//xla/service:hlo_proto_cc",
],
)
tf_cc_test(
name = "compile_tf_graph_test",
testonly = 1,
srcs = ["compile_tf_graph_test.cc"],
data = [
"testdata/prepare_to_library.mlir",
],
linkstatic = 1,
deps = [
":compile_tf_graph",
"//tensorflow/compiler/jit",
"//tensorflow/compiler/jit:xla_tpu_device",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:deserialize_mlir_module_utils",
"//tensorflow/compiler/mlir/tf2xla/internal:test_matchers",
"//tensorflow/compiler/mlir/tf2xla/internal/utils:test_metadata_config",
"//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:protos_all_cc",
"//tensorflow/core:test_main",
"//tensorflow/core/lib/monitoring:cell_reader",
"//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
"//tensorflow/core/tpu/kernels:tpu_compile_op_support",
"//tensorflow/core/tpu/kernels/xla:host_compute_ops",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_googletest//:gtest",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@xla//xla:shape_util",
"@xla//xla/client:client_library",
"@xla//xla/hlo/translate/mhlo_to_hlo:type_to_shape",
"@xla//xla/mlir_hlo:hlo_dialect_registration",
"@xla//xla/stream_executor:platform",
"@xla//xla/stream_executor:platform_manager",
"@xla//xla/tsl/lib/core:status_test_util",
"@xla//xla/tsl/lib/monitoring:test_utils",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:statusor",
],
)
cc_library(
name = "cluster_tf",
srcs = ["cluster_tf.cc"],
hdrs = ["cluster_tf.h"],
visibility = [
"//tensorflow/compiler/tf2xla:__pkg__",
],
deps = [
":tf_dialect_to_executor",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:bridge_logger",
"//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
"//tensorflow/compiler/mlir/tensorflow:error_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow/transforms:verify_no_outside_compilation_markers_pass",
"//tensorflow/compiler/mlir/tensorflow/transforms/host_runtime:lower_cluster_to_runtime_ops",
"//tensorflow/compiler/mlir/tf2xla/internal:clustering_bridge_passes",
"//tensorflow/compiler/mlir/tf2xla/internal:logging_hooks",
"//tensorflow/compiler/mlir/tf2xla/internal/inference:inference_metrics_pass",
"//tensorflow/compiler/mlir/tf2xla/internal/passes:clustering_passes",
"//tensorflow/core:framework",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:stacktrace",
"//tensorflow/core/platform:status",
"//tensorflow/core/tpu:tpu_defs",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:vlog_is_on",
"@com_google_absl//absl/status",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@tsl//tsl/platform:error_logging",
"@xla//xla/tsl/platform:errors",
],
)
tf_cc_test(
name = "cluster_tf_test",
srcs = ["cluster_tf_test.cc"],
data = [
"testdata/empty_func.mlir",
"testdata/invalid_executor.mlir",
"testdata/multiple_submodules.mlir",
],
deps = [
":cluster_tf",
"//tensorflow/compiler/mlir:register_common_dialects",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:tf_dialect_lib",
"//tensorflow/core/lib/monitoring:cell_reader",
"//tensorflow/core/platform:resource_loader",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"@xla//xla/tsl/lib/core:status_test_util",
"@xla//xla/tsl/platform:status",
],
)
cc_library(
name = "tf_dialect_to_executor",
srcs = ["tf_dialect_to_executor.cc"],
hdrs = ["tf_dialect_to_executor.h"],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/compiler/jit:flags_headers",
"//tensorflow/compiler/mlir/tensorflow:bridge_logger",
"//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
"//tensorflow/compiler/mlir/tensorflow:error_util",
"//tensorflow/compiler/mlir/tensorflow/transforms:verify_no_outside_compilation_markers_pass",
"//tensorflow/compiler/mlir/tf2xla/internal:logging_hooks",
"//tensorflow/core:framework",
"//tensorflow/core/platform:error_payloads",
"//tensorflow/core/platform:status",
"//tensorflow/core/protobuf:for_core_protos_cc",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:vlog_is_on",
"@com_google_absl//absl/status",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:Transforms",
"@tsl//tsl/platform:error_logging",
"@xla//xla/tsl/lib/monitoring:counter",
"@xla//xla/tsl/platform:status",
],
)
tf_cc_test(
name = "tf_dialect_to_executor_test",
srcs = ["tf_dialect_to_executor_test.cc"],
data = [
"testdata/empty_func.mlir",
"testdata/invalid_executor.mlir",
],
deps = [
":tf_dialect_to_executor",
"//tensorflow/compiler/mlir:register_common_dialects",
"//tensorflow/core/lib/monitoring:cell_reader",
"//tensorflow/core/platform:resource_loader",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"@xla//xla/tsl/lib/core:status_test_util",
"@xla//xla/tsl/platform:status",
],
)
tf_proto_library(
name = "mlir_bridge_config_v1_proto",
srcs = ["mlir_bridge_config_v1.proto"],
protodeps = tf_additional_all_protos(),
visibility = ["//visibility:public"],
)
tensorflow/compiler/mlir/tf2xla/api/v1/cluster_tf.cc
+245
View File
@@ -0,0 +1,245 @@
/* 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/api/v1/cluster_tf.h"
#include <memory>
#include <string>
#include "absl/log/log.h"
#include "absl/log/vlog_is_on.h"
#include "absl/status/status.h"
#include "llvm/ADT/STLFunctionalExtras.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.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/utils/attribute_utils.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/data_dumper_logger_config.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.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/internal/clustering_bridge_passes.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/inference/inference_passes.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/logging_hooks.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h"
#include "xla/tsl/framework/device_type.h"
#include "xla/tsl/platform/errors.h"
#include "tensorflow/core/framework/metrics.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/stacktrace.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/tpu/tpu_defs.h"
#include "tensorflow/core/util/debug_data_dumper.h"
#include "tsl/platform/error_logging.h"
namespace tensorflow {
namespace tf2xla {
namespace v1 {
using mlir::LogicalResult;
using mlir::ModuleOp;
using mlir::OpPassManager;
using mlir::PassManager;
using mlir::func::FuncOp;
namespace {
void CreateReplicatedBridgePipelineV1(OpPassManager &pm) {
pm.addPass(
tensorflow::tf2xla::internal::CreateTPUValidateSessionInputsPass());
pm.addPass(mlir::tf2xla::internal::CreateInferenceMetricsPass());
// Convert to unified compilation and replication attributes.
pm.addNestedPass<FuncOp>(
mlir::TF::CreateCanonicalizeCompileAndReplicateAttributesPass());
// Guarantee all functions have one use, which enables more exact shape
// inference.
pm.addPass(mlir::TF::CreateGuaranteeAllFuncsOneUsePass());
pm.addPass(mlir::TF::CreateTFShapeInferencePass());
// For V1 compatibility, we process a module where the graph does not have
// feeds and fetched. We extract first the TPU computation in a submodule,
// where it'll be in a function with args and returned values, much more
// like a TF v2 module. We can then run the usual pipeline on this nested
// module. Afterward we inline back in the parent module and delete the
// nested one.
pm.addPass(mlir::tf_executor::CreateTFExecutorTPUV1IslandCoarseningPass());
pm.addPass(mlir::tf_executor::CreateTFExecutorTPUV1IslandOutliningPass());
}
// Run the TF XLA Bridge based on the input pipeline, which can be either TPU
// bridge pipeline or non TPU bridge pipeline.
absl::Status RunTFXLABridge(
ModuleOp module,
llvm::function_ref<void(OpPassManager &pm)> pipeline_builder,
llvm::StringRef module_name = llvm::StringRef(),
llvm::StringRef dump_prefix = "tf_xla_bridge_v1") {
// Explicitly check that the TensorFlow dialect can constant fold ops.
// Constant folding is essential for the bridge. Without this check, the
// bridge may fail with an error that is difficult to understand and not
// actionable.
if (!mlir::TF::TensorFlowDialect::HasConstantFoldHook()) {
return absl::InternalError(
"TensorFlow dialect missing constant fold hook in TFXLA bridge phase "
"1; this could happen if the binary doesn't link the constant fold "
"hook registration library.");
}
PassManager bridge(module.getContext());
bridge.enableVerifier();
::tensorflow::applyTensorflowAndCLOptions(bridge);
// Populate a passmanager with the list of passes that implement the bridge.
pipeline_builder(bridge);
mlir::StatusScopedDiagnosticHandler diag_handler(
module.getContext(), /*propagate=*/false,
/*filter_stack=*/!VLOG_IS_ON(1));
if (VLOG_IS_ON(1) ||
DEBUG_DATA_DUMPER()->ShouldDump(module_name.str(), kDebugGroupMain)) {
::tensorflow::DumpMlirOpToFile(DEBUG_DATA_DUMPER()->GetDumpFilename(
module_name.str(), kDebugGroupMain,
"_" + dump_prefix.str() + "_before"),
module, llvm::StringRef(), &bridge);
}
if (VLOG_IS_ON(2) ||
DEBUG_DATA_DUMPER()->ShouldDump(module_name.str(),
kDebugGroupBridgePhase1Clustering)) {
internal::EnablePassIRPrinting(bridge, kDebugGroupBridgePhase1Clustering,
module_name);
}
LogicalResult result = bridge.run(module);
(void)result;
if (VLOG_IS_ON(1) ||
DEBUG_DATA_DUMPER()->ShouldDump(module_name.str(), kDebugGroupMain)) {
::tensorflow::DumpMlirOpToFile(DEBUG_DATA_DUMPER()->GetDumpFilename(
module_name.str(), kDebugGroupMain,
"_" + dump_prefix.str() + "_after"),
module, llvm::StringRef(), &bridge);
}
return diag_handler.ConsumeStatus();
}
} // namespace
absl::Status RecordStatusIfError(const std::string error_prefix,
bool is_in_fallback_enabled_mode,
absl::Status status) {
if (status.ok()) {
return status;
}
tensorflow::metrics::UpdateTfMlirBridgeFirstPhaseCounter(
/*bridge_type=*/mlir::TF::kMlirPh1BridgeCounterReplicated,
/*bridge_version=*/mlir::TF::kMlirPh1BridgeCounterV1,
/*device_type*/ mlir::TF::kMlirPh1BridgeCounterTpu,
/*fallback_enabled=*/is_in_fallback_enabled_mode,
/*result=*/"failure");
tsl::error_logging::Log(mlir::TF::kBridgeComponent,
"TFXLA_PHASE_ONE_MLIR_TPU_V1_COMPAT_BRIDGE",
status.ToString())
.IgnoreError();
return status;
}
// V1 Compat Bridge takes a TF Executor dialect and extracts the TF2 portion
// and inserts it into a submodule. We just want to run the clustering
// portion of the pipeline on just the single submodule.
absl::Status RunClusteringPipelineOnSubmodule(
ModuleOp parent_module, bool is_in_fallback_enabled_mode) {
int num_submodules = 0;
absl::Status clustering_pipeline_status;
parent_module.walk([&](ModuleOp submodule) {
if (submodule == parent_module) return mlir::WalkResult::advance();
num_submodules++;
clustering_pipeline_status = RunTFXLABridge(
submodule,
[](OpPassManager &pm) {
internal::AddReplicatedBridgeClusteringPipelinePasses(pm);
},
/*module_name=*/"", /*dump_prefix=*/"tf_xla_clustering_bridge_v1");
if (num_submodules > 1) {
return mlir::WalkResult::interrupt();
}
return mlir::WalkResult::advance();
});
if (num_submodules > 1) {
auto num_submodules_error = absl::InternalError(
"V1 Compat Bridge has more than one submodule. Erroring out.");
TF_RETURN_IF_ERROR(RecordStatusIfError(
/*error_prefix=*/"Bridge has more than one submodule:",
is_in_fallback_enabled_mode, num_submodules_error));
}
if (!clustering_pipeline_status.ok()) {
TF_RETURN_IF_ERROR(RecordStatusIfError(
/*error_prefix=*/"Bridge Errored running clustering pipeline:",
is_in_fallback_enabled_mode, clustering_pipeline_status));
}
return absl::OkStatus();
}
absl::Status RunSessionTf2xlaClusteringBridge(
ModuleOp module, bool is_in_fallback_enabled_mode) {
VLOG(2) << "TPU Sessions Bridge called stack trace is "
<< "(NOTE: this is not an error; rather the stack trace for "
"debugging) : "
<< tensorflow::CurrentStackTrace();
absl::Status functional_import_status = RunTFXLABridge(
module, [](OpPassManager &pm) { CreateReplicatedBridgePipelineV1(pm); },
/*module_name=*/"", /*dump_prefix=*/"tf_xla_functional_import_bridge_v1");
TF_RETURN_IF_ERROR(RecordStatusIfError(
/*error_prefix=*/"Bridge Function Import V1", is_in_fallback_enabled_mode,
functional_import_status));
TF_RETURN_IF_ERROR(
RunClusteringPipelineOnSubmodule(module, is_in_fallback_enabled_mode));
tensorflow::metrics::UpdateTfMlirBridgeFirstPhaseCounter(
/*bridge_type=*/mlir::TF::kMlirPh1BridgeCounterReplicated,
/*bridge_version=*/mlir::TF::kMlirPh1BridgeCounterV1,
/*device_type*/ mlir::TF::kMlirPh1BridgeCounterTpu,
/*n_fallback_enabled*/ is_in_fallback_enabled_mode,
/*result=*/"success");
return absl::OkStatus();
}
// Registers a pipeline builder function for TF TPU V1 bridge.
mlir::PassPipelineRegistration<> replicated_clustering_bridge_v1(
"tf-replicated-clustering-bridge-v1",
"Run all the passes involved in transforming a TensorFlow V1 graph before "
"execution so that it is suitable for targeting devices.",
CreateReplicatedBridgePipelineV1);
} // namespace v1
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v1/cluster_tf.h
+45
View File
@@ -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.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V1_CLUSTER_TF_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V1_CLUSTER_TF_H_
#include "absl/status/status.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "tensorflow/core/lib/core/status.h"
namespace tensorflow {
namespace tf2xla {
namespace v1 {
// Run all the passes involved in transforming the graph before execution so
// that it is suitable for targeting devices when called via the TF1 Session
// API.
// These transformations take as input a Tensorflow Graph as an MLIR Module
// and transforms the module in place to cluster the given ops for compilation
// that is compatible with the given device_type. The MLIR should be in the TF
// Executor Dialect for graph nodes and edges or TF Functional. It will convert
// to TF Functional internally. Individual Op inside a node should be the
// Tensorflow Dialect. The output MLIR is in the TF Functional Dialect. The
// input MLIR should not have infeed and outfeed ops, which are unsupported via
// this API. Returns OkStatus if passed, otherwise an error.
absl::Status RunSessionTf2xlaClusteringBridge(mlir::ModuleOp module,
bool is_in_fallback_enabled_mode);
} // namespace v1
} // namespace tf2xla
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V1_CLUSTER_TF_H_
tensorflow/compiler/mlir/tf2xla/api/v1/cluster_tf_test.cc
+115
View File
@@ -0,0 +1,115 @@
/* 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/api/v1/cluster_tf.h"
#include <cstdint>
#include <string>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "mlir/IR/BuiltinOps.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 "mlir/Parser/Parser.h" // from @llvm-project
#include "tensorflow/compiler/mlir/register_common_dialects.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/platform/status.h"
#include "tensorflow/core/lib/monitoring/cell_reader.h"
#include "tensorflow/core/platform/resource_loader.h"
namespace tensorflow {
namespace tf2xla {
namespace v1 {
namespace {
using ::mlir::DialectRegistry;
using ::mlir::MLIRContext;
using ::mlir::ModuleOp;
using ::mlir::OwningOpRef;
using ::tensorflow::monitoring::testing::CellReader;
static constexpr char kCompilationStreamz[] =
"/tensorflow/core/tf_mlir_bridge_first_phase_v2_count";
std::string TestDataPath() {
return tensorflow::GetDataDependencyFilepath(
"tensorflow/compiler/mlir/tf2xla/api/v1/testdata/");
}
class SessionClusterTensorflowDialectTest : public ::testing::Test {
public:
SessionClusterTensorflowDialectTest() {
mlir::RegisterCommonToolingDialects(registry_);
context_.appendDialectRegistry(registry_);
context_.loadAllAvailableDialects();
}
absl::Status CreateMlirModule(std::string mlir_module_filename) {
std::string mlir_module_path = TestDataPath() + mlir_module_filename;
mlir_module_ =
mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context_);
if (!mlir_module_) {
return absl::Status(
absl::StatusCode::kNotFound,
absl::StrCat("Could not find MLIR module at ", mlir_module_path));
}
return absl::OkStatus();
}
DialectRegistry registry_;
MLIRContext context_;
OwningOpRef<mlir::ModuleOp> mlir_module_;
};
TEST_F(SessionClusterTensorflowDialectTest, ClustersTf) {
CellReader<int64_t> compilation_status(kCompilationStreamz);
TF_ASSERT_OK(CreateMlirModule("empty_func.mlir"));
TF_EXPECT_OK(
RunSessionTf2xlaClusteringBridge(*mlir_module_,
/*is_in_fallback_enabled_mode=*/false));
EXPECT_EQ(compilation_status.Delta(mlir::TF::kMlirPh1BridgeCounterReplicated,
mlir::TF::kMlirPh1BridgeCounterV1,
mlir::TF::kMlirPh1BridgeCounterTpu,
"fallback_disabled", "success"),
1);
}
TEST_F(SessionClusterTensorflowDialectTest, FailsWithMultipleSubmodules) {
CellReader<int64_t> compilation_status(kCompilationStreamz);
TF_ASSERT_OK(CreateMlirModule("multiple_submodules.mlir"));
EXPECT_FALSE(
RunSessionTf2xlaClusteringBridge(*mlir_module_,
/*is_in_fallback_enabled_mode=*/false)
.ok());
EXPECT_EQ(compilation_status.Delta(mlir::TF::kMlirPh1BridgeCounterReplicated,
mlir::TF::kMlirPh1BridgeCounterV1,
mlir::TF::kMlirPh1BridgeCounterTpu,
"fallback_disabled", "failure"),
1);
}
} // namespace
} // namespace v1
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v1/compile_mlir_util.cc
+1139
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/mlir/tf2xla/api/v1/compile_mlir_util.h
+242
View File
@@ -0,0 +1,242 @@
/* 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_API_V1_COMPILE_MLIR_UTIL_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V1_COMPILE_MLIR_UTIL_H_
#include <memory>
#include <string>
#include <vector>
#include "absl/base/attributes.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "tensorflow/compiler/tf2xla/layout_util.h"
#include "tensorflow/compiler/tf2xla/xla_argument.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/hlo/builder/xla_computation.h"
#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/framework/tensor_shape.h"
namespace tensorflow {
// Lowers MLIR module to XLA HLO inside an XlaComputation. The input module
// should only contain operations in tf dialect. If the input module contains
// operation in the tf_executor dialect, for example, returns an error.
// Exception to this are tf_executor dialect ops that are optimized away through
// canonicalization.
//
// Operations in tf dialect are lowered to XLA HLO through the following steps:
// . Legalizes control flow operations.
// . Decomposes compound resource operations so that the only remaining
// operations on resource variables are resource reads/writes..
// . Replaces resource reads/writes with function inputs/outputs and
// eliminates the use of resource variables.
// . Legalizes the operations to XLA HLO operations.
// . Canonicalizes the XLA HLO operations.
//
// device_type: XLA JIT device to use for compilation such as "XLA_CPU_JIT",
// "XLA_GPU_JIT" or "XLA_TPU_JIT".
// use_tuple_args: when this is true, always create a tuple argument for the
// entry computation.
// enable_op_fallback: when this is true, prefer tf2xla fallback kernels over
// MLIR
// native kernels for legalization to HLO.
// return_tuple: when this is true, always create a tuple result for the
// entry computation.
// shape_determination_fns: Contains layout preference fn and shape
// representation fn. The two functions are used to determine argument and
// result shapes.
// custom_legalization_passes: passes to run before the default TF legalization
// passes for backend-specific ops.
ABSL_DEPRECATED("Use v2/legalize_tf.h::LegalizeMlirToHlo instead.")
absl::Status ConvertMLIRToXlaComputation(
mlir::ModuleOp module_op, llvm::StringRef device_type,
xla::XlaComputation* xla_computation, bool use_tuple_args,
bool enable_op_fallback, bool return_tuple,
XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns = {},
llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
custom_legalization_passes = {},
llvm::StringRef module_name = llvm::StringRef());
// Creates a MLIR pipeline that lowers MLIR module to MHLO dialect. The input
// module should only contain operations in tf dialect. For example, if the
// input module contains operation in the tf_executor dialect, the pass raises
// an error unless the tf_executor dialect ops are optimized away by
// canonicalization.
//
// The pipeline is used in ConvertMLIRToXlaComputation. And it generally has the
// following pass structure:
// - TensorFlow passes
// - Legalization passes
// - MHLO passes
//
// device_type: XLA JIT device to use for compilation such as "XLA_CPU_JIT",
// "XLA_GPU_JIT" or "XLA_TPU_JIT".
// enable_op_fallback: when this is true, prefer tf2xla fallback kernels over
// MLIR
// native kernels for legalization to HLO.
// custom_legalization_passes: passes to run before the default TF legalization
// passes for backend-specific ops.
// lower_to_xla_hlo: Temporary parameter to be removed in imminent update. If
// true, includes legalization and MHLO lowering passes.
// allow_partial_conversion: when this is true, allow operations that can't be
// legalized.
ABSL_DEPRECATED("Use v2/legalize_tf.h::LegalizeMlirToHlo instead.")
void CreateConvertMlirToXlaHloPipeline(
mlir::OpPassManager& pm, llvm::StringRef device_type,
bool enable_op_fallback,
llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
custom_legalization_passes,
bool lower_to_xla_hlo = true, bool allow_partial_conversion = false);
// Helper struct representing argument tensor or resource handle shapes.
struct TensorOrResourceShape {
TensorShape shape;
bool is_resource = false;
};
// Refine MLIR types based on new shape information.
ABSL_DEPRECATED("Not meant to be used directly and should be a util.")
absl::Status RefineShapes(llvm::ArrayRef<TensorOrResourceShape> arg_shapes,
mlir::ModuleOp module);
// Lower TF to MHLO and insert HLO into the XlaBuilder. xla_params are HLO-level
// inputs to module_op that have already been added to the XlaBuilder. returns
// are the returned XlaOps.
ABSL_DEPRECATED("Use v2/legalize_tf.h::LegalizeMlirToHlo instead.")
absl::Status BuildHloFromTf(mlir::ModuleOp module_op, xla::XlaBuilder& builder,
llvm::ArrayRef<xla::XlaOp> xla_params,
std::vector<xla::XlaOp>& returns,
llvm::ArrayRef<TensorOrResourceShape> arg_shapes,
llvm::StringRef device_type,
llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
custom_legalization_passes);
// Apply shape, description, and resource information to inputs and outputs
// in the XlaCompilationResult. This should be called after
// compilation_result->computation was set.
ABSL_DEPRECATED("Not meant to be used directly and should be a util.")
absl::Status PopulateResultIOInfo(
mlir::ModuleOp module_op, llvm::ArrayRef<TensorOrResourceShape> arg_shapes,
bool use_tuple_args, bool use_resource_updates_for_aliases,
XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,
XlaCompilationResult* compilation_result);
// Runs MLIR Bridge on an MLIR module.
//
// If lower_to_xla_hlo is true then compiles down into XLA HLO, generates all
// accompanying metadata and stores them in CompilationResult.
//
// If enable_op_fallback is set to false, graph is legalized only if the graph
// analysis for the graph is successful. Otherwise, an error is returned.
//
// Running the MLIR Bridge performs many transformations on the input module
// which is modified in place.
ABSL_DEPRECATED("Use v2/legalize_tf.h::LegalizeMlirToHlo instead.")
absl::Status CompileMlirToXlaHlo(
mlir::ModuleOp module_op, llvm::ArrayRef<TensorOrResourceShape> arg_shapes,
llvm::StringRef device_type, bool use_tuple_args, bool enable_op_fallback,
bool use_return_tuple, bool use_resource_updates_for_aliases,
XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,
XlaCompilationResult* compilation_result,
llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
custom_legalization_passes,
llvm::StringRef module_name = llvm::StringRef(),
bool lower_to_xla_hlo = true);
// Runs MLIR Bridge on a MLIR module.
//
// If lower_to_xla_hlo is true then compiles down into XLA HLO, generates all
// accompanying metadata and stores them in CompilationResult.
//
// If enable_op_fallback is set to false, graph is legalized only if the graph
// analysis for the graph is successful. Otherwise, an error is returned.
//
// On success, returns the serialized MLIR module.
ABSL_DEPRECATED("Use v2/legalize_tf.h::LegalizeMlirToHlo instead.")
absl::StatusOr<std::string> CompileMlirToXlaHloAndSerialize(
mlir::ModuleOp module_op, llvm::ArrayRef<TensorOrResourceShape> arg_shapes,
llvm::StringRef device_type, bool use_tuple_args, bool enable_op_fallback,
bool use_return_tuple, bool use_resource_updates_for_aliases,
XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,
XlaCompilationResult* compilation_result,
llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
custom_legalization_passes,
llvm::StringRef module_name = llvm::StringRef(),
bool lower_to_xla_hlo = true);
// Runs MLIR Bridge on a serialized MLIR module.
//
// If lower_to_xla_hlo is true then compiles down into XLA HLO, generates all
// accompanying metadata and stores them in CompilationResult.
ABSL_DEPRECATED("Use v2/legalize_tf.h::LegalizeMlirToHlo instead.")
absl::StatusOr<std::string> CompileSerializedMlirToXlaHlo(
llvm::StringRef mlir_module_string, llvm::ArrayRef<TensorShape> arg_shapes,
llvm::StringRef device_type, bool use_tuple_args, bool enable_op_fallback,
XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,
XlaCompilationResult* compilation_result,
llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
custom_legalization_passes = {},
llvm::StringRef module_name = llvm::StringRef(),
bool lower_to_xla_hlo = true);
// Compiles a TensorFlow Graph (already converted to MLIR, imported with
// tf_executor dialect still present) into XLA HLO, generates all accompanying
// metadata and stores them in CompilationResult. This will rewrite arguments
// and run the TensorFlow standard pipeline prior to invoking
// `CompileMlirToXlaHlo`.
ABSL_DEPRECATED("Use v2/legalize_tf.h::LegalizeMlirToHlo instead.")
absl::Status CompileGraphToXlaHlo(
mlir::ModuleOp module_op, llvm::ArrayRef<XlaArgument> args,
llvm::StringRef device_type, bool use_tuple_args, bool enable_op_fallback,
bool use_return_tuple,
XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,
XlaCompilationResult* compilation_result,
llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
custom_legalization_passes);
// Compiles a Graph from TF to HLO and adds the resulting HLO to the
// XlaBuilder. This function adds HLO to a larger HLO computation, so
// HLO-level inputs are supplied, and HLO-level outputs are produced.
// xla_params is the HLO-level inputs and returns is the HLO-level outputs.
// If unconditionally_use_output_shapes is true then the unregistered
// attribute _output_shapes is always used to set the output shapes of the ops.
ABSL_DEPRECATED(
"Use v1/compile_tf_graph.h::CompileTensorflowGraphToHlo instead.")
absl::Status BuildHloFromGraph(
const Graph& graph, xla::XlaBuilder& builder,
mlir::MLIRContext& mlir_context, llvm::ArrayRef<xla::XlaOp> xla_params,
std::vector<xla::XlaOp>& returns, bool unconditionally_use_output_shapes,
llvm::ArrayRef<XlaArgument> args, llvm::ArrayRef<std::string> control_rets,
llvm::StringRef device_type, const FunctionLibraryDefinition& flib_def);
static inline absl::Status CompileToHloGraphAnalysisFailedError() {
return absl::InternalError("disabled after graph analysis");
}
// Register a convenient pipeline for invoking TF/XLA lowering from the command
// line.
void RegisterConvertMlirToXlaHloPipelineWithDefaults();
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V1_COMPILE_MLIR_UTIL_H_
tensorflow/compiler/mlir/tf2xla/api/v1/compile_mlir_util_test.cc
+316
View File
@@ -0,0 +1,316 @@
/* 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/api/v1/compile_mlir_util.h"
#include <cstdint>
#include <initializer_list>
#include <memory>
#include <string>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/IR/DialectRegistry.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "tensorflow/compiler/jit/xla_compile_util.h"
#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/serialize_mlir_module_utils.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/test_matchers.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/hlo/builder/xla_builder.h"
#include "xla/hlo/builder/xla_computation.h"
#include "xla/shape.h"
#include "xla/shape_util.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/statusor.h"
#include "xla/xla_data.pb.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/fake_input.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/lib/monitoring/cell_reader.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace {
using ::mlir::OpPassManager;
using ::tensorflow::monitoring::testing::CellReader;
using ::testing::HasSubstr;
static constexpr char kMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> tensor<1xi32> {
%0 = "tf.Const"() {value = dense<1000> : tensor<1xi32>} : () -> tensor<1xi32>
func.return %0 : tensor<1xi32>
}
})";
TEST(LegalizeMlirTest, LegalizesModule) {
mlir::DialectRegistry mlir_registry;
RegisterAllTensorFlowDialects(mlir_registry);
std::vector<tensorflow::TensorShape> arg_shapes;
XlaCompilationResult compilation_result;
auto status = CompileSerializedMlirToXlaHlo(
kMlirModuleStr, arg_shapes, /*device_type=*/"XLA_TPU_JIT",
/*use_tuple_args=*/true, /*enable_op_fallback=*/false,
/*shape_determination_fns=*/{}, &compilation_result);
EXPECT_TRUE(status.ok());
EXPECT_THAT(status.value(), HasSubstr("stablehlo.constant"));
}
TEST(LegalizeMlirTest, FailsLegalizesModule) {
constexpr char failed_legalization[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> tensor<1xi32> {
%0 = "tf.DoesntExist"() : () -> tensor<1xi32>
func.return %0 : tensor<1xi32>
}
})";
CellReader<int64_t> count(
"/tensorflow/core/tf2xla/v1/mlir_failed_xla_legalize_tf_pass_count");
std::vector<tensorflow::TensorShape> arg_shapes;
XlaCompilationResult compilation_result;
auto status = CompileSerializedMlirToXlaHlo(
failed_legalization, arg_shapes, /*device_type=*/"XLA_TPU_JIT",
/*use_tuple_args=*/true, /*enable_op_fallback=*/false,
/*shape_determination_fns=*/{}, &compilation_result);
EXPECT_FALSE(status.ok());
EXPECT_EQ(count.Delta("tf.DoesntExist", "Unknown"), 1);
}
TEST(CompileMlirUtil, CreatesPipeline) {
OpPassManager pass_manager;
llvm::StringRef device_type = "XLA_CPU_JIT";
CreateConvertMlirToXlaHloPipeline(pass_manager, device_type,
/*enable_op_fallback=*/false,
/*custom_legalization_passes*/ {});
EXPECT_FALSE(pass_manager.getPasses().empty());
}
TEST(CompileMlirUtil, HasLegalizationPass) {
OpPassManager pass_manager;
llvm::StringRef device_type = "XLA_CPU_JIT";
absl::string_view kLegalizeTfPass = "xla-legalize-tf";
CreateConvertMlirToXlaHloPipeline(pass_manager, device_type,
/*enable_op_fallback=*/true,
/*custom_legalization_passes*/ {});
std::string pass_description;
llvm::raw_string_ostream raw_stream(pass_description);
pass_manager.printAsTextualPipeline(raw_stream);
EXPECT_THAT(pass_description, HasSubstr(kLegalizeTfPass));
}
TEST(CompileMlirUtil, DoesNotHaveLegalizationPass) {
OpPassManager pass_manager;
llvm::StringRef device_type = "XLA_CPU_JIT";
absl::string_view kLegalizeTfPass = "xla-legalize-tf";
CreateConvertMlirToXlaHloPipeline(pass_manager, device_type,
/*enable_op_fallback=*/false,
/*custom_legalization_passes*/ {},
/*lower_to_xla_hlo=*/false);
std::string pass_description;
llvm::raw_string_ostream raw_stream(pass_description);
pass_manager.printAsTextualPipeline(raw_stream);
EXPECT_THAT(pass_description, Not(HasSubstr(kLegalizeTfPass)));
}
TEST(CompileMlirUtil, DoesNotLowerWhenTold) {
mlir::DialectRegistry mlir_registry;
RegisterAllTensorFlowDialects(mlir_registry);
std::vector<tensorflow::TensorShape> arg_shapes;
XlaCompilationResult compilation_result;
auto status = CompileSerializedMlirToXlaHlo(
kMlirModuleStr, arg_shapes, /*device_type=*/"XLA_TPU_JIT",
/*use_tuple_args=*/true, /*enable_op_fallback=*/false,
/*shape_determination_fns=*/{}, &compilation_result,
/*custom_legalization_passes=*/{},
/*module_name=*/"",
/*lower_to_xla_hlo=*/false);
EXPECT_TRUE(status.ok());
EXPECT_THAT(status.value(), HasSubstr("tf.Const"));
}
TEST(CompileMlirUtil, CanonicalizationIsExplicitDuringInlining) {
OpPassManager pass_manager;
llvm::StringRef device_type = "XLA_CPU_JIT";
absl::string_view kInlinePass =
"inline{default-pipeline=canonicalize "
"inlining-threshold=4294967295 max-iterations=4 }";
CreateConvertMlirToXlaHloPipeline(pass_manager, device_type,
/*enable_op_fallback=*/true,
/*custom_legalization_passes*/ {});
std::string pass_description;
llvm::raw_string_ostream raw_stream(pass_description);
pass_manager.printAsTextualPipeline(raw_stream);
EXPECT_THAT(pass_description, HasSubstr(kInlinePass));
}
TEST(LegalizeMlirTest, LegalizesModuleWithDynamicShape) {
constexpr char legalization[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main(%arg0: tensor<?xi32, #mhlo.type_extensions<bounds = [1]>>) -> tensor<?xi32, #mhlo.type_extensions<bounds = [1]>> {
%0 = "tf.Identity"(%arg0) : (tensor<?xi32, #mhlo.type_extensions<bounds = [1]>>) -> tensor<?xi32, #mhlo.type_extensions<bounds = [1]>>
func.return %0 : tensor<?xi32, #mhlo.type_extensions<bounds = [1]>>
}
})";
std::vector<tensorflow::TensorShape> arg_shapes = {{1}};
XlaCompilationResult compilation_result;
auto status = CompileSerializedMlirToXlaHlo(
legalization, arg_shapes, /*device_type=*/"XLA_TPU_JIT",
/*use_tuple_args=*/true, /*enable_op_fallback=*/false,
/*shape_determination_fns=*/{}, &compilation_result);
EXPECT_TRUE(status.ok());
}
absl::StatusOr<std::unique_ptr<Graph>> BuildConstOpGraphWithOutputShapes() {
DataType data_type = DT_INT32;
std::initializer_list<int64_t> dims = {2, 3, 4, 5};
Tensor tensor(data_type, TensorShape(dims));
for (int i = 0; i < 2 * 3 * 4 * 5; ++i) {
tensor.flat<int32_t>()(i) = i;
}
NodeDef node;
auto builder = NodeDefBuilder("some_node", "Const")
.Attr("dtype", data_type)
.Attr("value", tensor);
// Create a bad output shape attr.
AttrValue shape_attr;
TensorShapeProto* shape_proto = shape_attr.mutable_list()->add_shape();
shape_proto->add_dim()->set_size(1);
builder.Attr("_output_shapes", shape_attr);
TF_RETURN_IF_ERROR(builder.Finalize(&node));
return CreateSingleOpGraph(node, {}, {data_type});
}
absl::StatusOr<std::unique_ptr<Graph>> BuildEmptyOpGraph(
std::vector<XlaCompiler::Argument>& xla_args) {
DataType data_type = DT_INT32;
XlaCompiler::Argument arg;
arg.type = DT_INT32;
arg.shape = xla::ShapeUtil::MakeShape(xla::S32, {});
arg.name = "arg0";
arg.kind = XlaCompiler::Argument::kParameter;
xla_args.push_back(arg);
NodeDef node;
auto builder = NodeDefBuilder("some_node", "Empty")
.Input(FakeInput(DT_INT32))
.Attr("dtype", data_type);
TF_RETURN_IF_ERROR(builder.Finalize(&node));
return CreateSingleOpGraph(node, xla_args, {data_type});
}
absl::StatusOr<xla::XlaComputation> BuildHloFromGraph(
Graph& graph, std::vector<XlaCompiler::Argument>& xla_args,
bool use_output_shapes) {
xla::XlaBuilder builder(
::testing::UnitTest::GetInstance()->current_test_info()->name());
mlir::MLIRContext mlir_context;
llvm::SmallVector<xla::XlaOp, 4> xla_params;
for (int i = 0; i < xla_args.size(); ++i) {
xla_params.push_back(Parameter(&builder, i,
std::get<xla::Shape>(xla_args[i].shape),
"arg" + std::to_string(i)));
}
std::vector<xla::XlaOp> returns(1);
TF_RETURN_IF_ERROR(
BuildHloFromGraph(graph, builder, mlir_context, xla_params, returns,
use_output_shapes, xla_args,
/*control_rets=*/{}, DEVICE_TPU,
FunctionLibraryDefinition(OpRegistry::Global())));
return builder.Build();
}
TEST(CompileMlirUtil, UsesCorrectOriginalShapeWithoutOutputShapes) {
std::vector<XlaCompiler::Argument> xla_args;
// Build a graph with an op that is supported by the MLIR lowerings.
TF_ASSERT_OK_AND_ASSIGN(auto graph, BuildConstOpGraphWithOutputShapes());
auto build_result =
BuildHloFromGraph(*graph, xla_args, /*use_output_shapes=*/false);
TF_ASSERT_OK(build_result);
EXPECT_THAT(build_result,
XlaComputationProtoContains("opcode: \"constant\""));
}
TEST(CompileMlirUtil, UsesIncorrectOutputShapesWhenPresent) {
std::vector<XlaCompiler::Argument> xla_args;
TF_ASSERT_OK_AND_ASSIGN(auto graph, BuildConstOpGraphWithOutputShapes());
auto build_result =
BuildHloFromGraph(*graph, xla_args, /*use_output_shapes=*/true);
ASSERT_FALSE(build_result.ok());
EXPECT_THAT(build_result.status().message(),
HasSubstr("op operand type 'tensor<2x3x4x5xi32>' and result type "
"'tensor<1xi32>' are cast incompatible"));
}
TEST(CompileMlirUtil, DoesNotLowerFallbackOps) {
std::vector<XlaCompiler::Argument> xla_args;
// Build a graph with an op that is not supported by the MLIR lowerings.
TF_ASSERT_OK_AND_ASSIGN(auto graph, BuildEmptyOpGraph(xla_args));
auto build_result =
BuildHloFromGraph(*graph, xla_args, /*use_output_shapes=*/true);
ASSERT_FALSE(build_result.ok());
EXPECT_THAT(build_result.status().message(),
HasSubstr("'tf.Empty' op unsupported op"));
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph.cc
+381
View File
@@ -0,0 +1,381 @@
/* 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/api/v1/compile_tf_graph.h"
#include <chrono>
#include <cstdint>
#include <memory>
#include <string>
#include <variant>
#include <vector>
#include "absl/container/flat_hash_set.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 "llvm/ADT/DenseMap.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/LogicalResult.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.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/dialect_registration.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/set_tpu_infeed_layout.h"
#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/deserialize_mlir_module_utils.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/translate_utils.h"
#include "tensorflow/compiler/mlir/tf2xla/api/v2/tf_executor_to_graph.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/logging_hooks.h"
#include "tensorflow/compiler/tf2xla/layout_util.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/client/compile_only_client.h"
#include "xla/hlo/ir/hlo_input_output_alias_config.h"
#include "xla/mlir_hlo/mhlo/IR/register.h"
#include "xla/service/hlo.pb.h"
#include "xla/shape.h"
#include "xla/shape_util.h"
#include "xla/status_macros.h"
#include "xla/tsl/framework/device_type.h"
#include "xla/tsl/lib/monitoring/sampler.h"
#include "xla/tsl/platform/errors.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/versions.pb.h"
#include "tensorflow/core/lib/monitoring/counter.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/profile_utils/cpu_utils.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
#include "tensorflow/core/tpu/tpu_compile.h"
#include "tensorflow/core/util/debug_data_dumper.h"
namespace tensorflow {
namespace tf2xla {
namespace v1 {
using ::tensorflow::tpu::FunctionToHloArgs;
using ::tensorflow::tpu::GuaranteedConsts;
using ::tensorflow::tpu::MlirToHloArgs;
using ::tensorflow::tpu::ShardingAndIndex;
auto* phase2_bridge_compilation_status =
tensorflow::monitoring::Counter<1>::New(
"/tensorflow/core/tf2xla/api/v1/"
"phase2_compilation_status", /*metric_name*/
"Tracks the compilation status of the non-mlir bridge",
/* metric description */ "status" /* metric label */);
auto* phase2_bridge_compilation_time = tsl::monitoring::Sampler<1>::New(
{"/tensorflow/core/tf2xla/api/v1/phase2_compilation_time",
"The wall-clock time spent on executing graphs in milliseconds.",
"configuration"},
// Power of 1.5 with bucket count 45 (> 23 hours)
{tsl::monitoring::Buckets::Exponential(1, 1.5, 45)});
// There were no MLIR ops so the old bridge was called successfully.
constexpr char kOldBridgeNoMlirSuccess[] = "kOldBridgeNoMlirSuccess";
// There were no MLIR ops so the old bridge was called but it failed.
constexpr char kOldBridgeNoMlirFailure[] = "kOldBridgeNoMlirFailure";
namespace {
// Time the execution of kernels (in CPU cycles). Meant to be used as RAII.
struct CompilationTimer {
uint64_t start_cycles = profile_utils::CpuUtils::GetCurrentClockCycle();
uint64_t ElapsedCycles() {
return profile_utils::CpuUtils::GetCurrentClockCycle() - start_cycles;
}
int64_t ElapsedCyclesInMilliseconds() {
std::chrono::duration<double> duration =
profile_utils::CpuUtils::ConvertClockCycleToTime(ElapsedCycles());
return std::chrono::duration_cast<std::chrono::milliseconds>(duration)
.count();
}
};
// Populates input_output_alias field in the HLO Module proto.
absl::Status PopulateInputOutputAliasing(
mlir::func::FuncOp main_fn,
XlaCompiler::CompilationResult* compilation_result, bool use_tuple_args) {
constexpr char kAliasingAttr[] = "tf.aliasing_output";
llvm::SmallDenseMap<unsigned, unsigned> output_to_input_alias;
unsigned num_arguments = main_fn.getNumArguments();
for (unsigned arg_index = 0; arg_index < num_arguments; ++arg_index) {
if (auto aliasing_output = main_fn.getArgAttrOfType<mlir::IntegerAttr>(
arg_index, kAliasingAttr))
output_to_input_alias[aliasing_output.getInt()] = arg_index;
}
if (output_to_input_alias.empty()) return absl::OkStatus();
xla::HloModuleProto* module_proto =
compilation_result->computation->mutable_proto();
absl::StatusOr<xla::ProgramShape> program_shape_or_status =
compilation_result->computation->GetProgramShape();
TF_RET_CHECK(program_shape_or_status.ok());
xla::ProgramShape& program_shape = program_shape_or_status.value();
if (!program_shape.result().IsTuple())
return absl::InternalError("Expect result to have tuple shape");
xla::HloInputOutputAliasConfig config(program_shape.result());
for (auto alias : output_to_input_alias) {
if (use_tuple_args) {
TF_RETURN_IF_ERROR(config.SetUpAlias(
xla::ShapeIndex({alias.first}), 0, xla::ShapeIndex({alias.second}),
xla::HloInputOutputAliasConfig::AliasKind::kMayAlias));
} else {
TF_RETURN_IF_ERROR(config.SetUpAlias(
xla::ShapeIndex({alias.first}), alias.second, xla::ShapeIndex({}),
xla::HloInputOutputAliasConfig::AliasKind::kMayAlias));
}
}
*module_proto->mutable_input_output_alias() = config.ToProto();
return absl::OkStatus();
}
bool failed(const absl::Status& status) { return !status.ok(); }
// Transforms the given module to be suitable for export to TensorFlow GraphDef
// and then exports all functions to the given library.
absl::Status PrepareAndExportToLibrary(mlir::ModuleOp module,
FunctionLibraryDefinition* flib_def) {
// Pass pipeline is defined here instead of leveraging the phase one export
// pipeline because only the functional to executor dialect conversion and
// breakup islands passes are common between the export pipeline and here.
// Reconsider this if there is more commonality in the future with more
// passes.
mlir::PassManager manager(module.getContext());
applyTensorflowAndCLOptions(manager);
manager.addPass(mlir::TF::CreatePrepareTpuComputationForTfExportPass());
manager.addPass(mlir::TF::CreateTFRegionControlFlowToFunctional());
manager.addPass(mlir::TF::CreateTFShapeInferencePass());
manager.addNestedPass<mlir::func::FuncOp>(
mlir::CreateFunctionalToExecutorDialectConversionPass());
manager.addPass(mlir::CreateBreakUpIslandsPass());
mlir::StatusScopedDiagnosticHandler diag_handler(module.getContext());
if (VLOG_IS_ON(2)) {
llvm::StringRef module_name = llvm::StringRef();
constexpr const char* kDebugGroupBridgePhase2 =
"v1_prepare_and_export_to_library";
internal::EnablePassIRPrinting(manager, kDebugGroupBridgePhase2,
module_name);
}
auto prepare_status = manager.run(module);
auto diag_handler_status = diag_handler.ConsumeStatus();
// There are cases where the scoped diagnostic handler catches a failure that
// the running of the passes does not. That causes the handler to throw if
// it is not consumed.
if (failed(prepare_status) || failed(diag_handler_status)) {
return diag_handler_status;
}
GraphExportConfig config;
config.export_entry_func_to_flib = true;
absl::flat_hash_set<Node*> control_ret_nodes;
return tensorflow::tf2xla::v2::ConvertTfExecutorToGraph(
module, config, /*graph=*/nullptr, flib_def, &control_ret_nodes);
}
absl::Status CompileTFFunctionWithoutMlir(
FunctionToHloArgs function_computation,
const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,
const XlaShapeLayoutHelpers::ShapeDeterminationFns
shape_determination_funcs,
const std::vector<tensorflow::TensorShape>& arg_shapes,
const DeviceType& device_type,
std::vector<tpu::ShardingAndIndex>* arg_core_mapping,
std::vector<std::vector<xla::Shape>>* per_core_arg_shapes,
xla::CompileOnlyClient* client,
XlaCompiler::CompilationResult* compilation_result) {
absl::Status comp_status = CompileTFFunctionToHlo(
*function_computation.flib_def, function_computation.graph_def_version,
shape_determination_funcs, arg_shapes, device_type,
function_computation.guaranteed_constants, *function_computation.function,
metadata, client, arg_core_mapping, per_core_arg_shapes, use_tuple_args,
compilation_result);
if (comp_status.ok()) {
phase2_bridge_compilation_status->GetCell(kOldBridgeNoMlirSuccess)
->IncrementBy(1);
} else {
phase2_bridge_compilation_status->GetCell(kOldBridgeNoMlirFailure)
->IncrementBy(1);
}
return comp_status;
}
absl::Status CompileMLIRTFFunction(
mlir::ModuleOp mlir_module, const tpu::TPUCompileMetadataProto& metadata,
bool use_tuple_args,
const XlaShapeLayoutHelpers::ShapeDeterminationFns
shape_determination_funcs,
const std::vector<tensorflow::TensorShape>& arg_shapes,
const DeviceType& device_type,
std::vector<tpu::ShardingAndIndex>* arg_core_mapping,
std::vector<std::vector<xla::Shape>>* per_core_arg_shapes,
xla::CompileOnlyClient* client,
XlaCompiler::CompilationResult* compilation_result) {
if (!mlir::SetTPUInfeedLayout(mlir_module))
return absl::InternalError("Failed to set layouts attribute");
if (VLOG_IS_ON(2)) {
tensorflow::DumpMlirOpToFile("legalize_with_old_bridge", mlir_module);
}
constexpr char kEntryFuncName[] = "main";
auto main_fn = mlir_module.lookupSymbol<mlir::func::FuncOp>(kEntryFuncName);
if (!main_fn) {
return absl::InternalError(
"TPU compile op requires module with a entry function main");
}
// Export functions to the library.
auto flib_def = std::make_unique<FunctionLibraryDefinition>(
OpRegistry::Global(), FunctionDefLibrary());
TF_RETURN_IF_ERROR(PrepareAndExportToLibrary(mlir_module, flib_def.get()));
if (VLOG_IS_ON(2)) {
tensorflow::DumpMlirOpToFile("legalize_with_old_bridge_post_transform",
mlir_module);
}
VersionDef versions;
if (mlir::failed(ExtractTfVersions(mlir_module, &versions))) {
return absl::InternalError(
"module attribute in _TPUCompileMlir op is missing tf versions.");
}
NameAttrList func;
func.set_name(kEntryFuncName);
GuaranteedConsts consts;
*compilation_result = {};
TF_RETURN_IF_ERROR(CompileTFFunctionToHlo(
*flib_def, versions.producer(), shape_determination_funcs, arg_shapes,
device_type, consts, func, metadata, client, arg_core_mapping,
per_core_arg_shapes, use_tuple_args, compilation_result));
if (compilation_result != nullptr &&
compilation_result->computation != nullptr) {
if (auto module_name = mlir_module.getName()) {
compilation_result->computation->mutable_proto()->set_name(
module_name->str());
}
}
return PopulateInputOutputAliasing(main_fn, compilation_result,
use_tuple_args);
}
absl::Status CompileMLIRTFFunction(
tpu::MlirToHloArgs mlir_computation,
const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,
const XlaShapeLayoutHelpers::ShapeDeterminationFns
shape_determination_funcs,
const std::vector<tensorflow::TensorShape>& arg_shapes,
const DeviceType& device_type,
std::vector<tpu::ShardingAndIndex>* arg_core_mapping,
std::vector<std::vector<xla::Shape>>* per_core_arg_shapes,
xla::CompileOnlyClient* client,
XlaCompiler::CompilationResult* compilation_result) {
if (mlir_computation.mlir_module_op.has_value()) {
return CompileMLIRTFFunction(
mlir_computation.mlir_module_op.value(), metadata, use_tuple_args,
shape_determination_funcs, arg_shapes, device_type, arg_core_mapping,
per_core_arg_shapes, client, compilation_result);
}
mlir::DialectRegistry registry;
mlir::RegisterAllTensorFlowDialects(registry);
mlir::mhlo::registerAllMhloDialects(registry);
mlir::MLIRContext context(registry);
mlir::OwningOpRef<mlir::ModuleOp> mlir_module;
TF_RETURN_IF_ERROR(DeserializeMlirModule(mlir_computation.mlir_module,
&context, &mlir_module));
return CompileMLIRTFFunction(*mlir_module, metadata, use_tuple_args,
shape_determination_funcs, arg_shapes,
device_type, arg_core_mapping,
per_core_arg_shapes, client, compilation_result);
}
} // namespace
absl::Status CompileTensorflowGraphToHlo(
const std::variant<tpu::MlirToHloArgs, tpu::FunctionToHloArgs>& computation,
const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,
const XlaShapeLayoutHelpers::ShapeDeterminationFns
shape_determination_funcs,
const std::vector<tensorflow::TensorShape>& arg_shapes,
tsl::DeviceType device_type,
std::vector<tpu::ShardingAndIndex>* arg_core_mapping,
std::vector<std::vector<xla::Shape>>* per_core_arg_shapes,
xla::CompileOnlyClient* client,
XlaCompiler::CompilationResult* compilation_result) {
LOG_FIRST_N(INFO, 1) << "Compiling MLIR computation to XLA HLO using the "
"old (non-MLIR) tf2xla bridge";
CompilationTimer timer;
*compilation_result = {};
bool has_mlir = computation.index() == 0;
std::string mlir_string = has_mlir ? "has_mlir" : "has_function_to_hlo";
const std::string kBridgePhase2Config =
absl::StrCat("graph_old_bridge_", mlir_string);
if (has_mlir) {
TF_RETURN_IF_ERROR(CompileMLIRTFFunction(
std::get<0>(computation), metadata, use_tuple_args,
shape_determination_funcs, arg_shapes, device_type, arg_core_mapping,
per_core_arg_shapes, client, compilation_result));
} else {
FunctionToHloArgs function_computation = std::get<1>(computation);
TF_RETURN_IF_ERROR(CompileTFFunctionWithoutMlir(
function_computation, metadata, use_tuple_args,
shape_determination_funcs, arg_shapes, device_type, arg_core_mapping,
per_core_arg_shapes, client, compilation_result));
}
phase2_bridge_compilation_time->GetCell(kBridgePhase2Config)
->Add(timer.ElapsedCyclesInMilliseconds());
return absl::OkStatus();
}
}; // namespace v1
}; // namespace tf2xla
}; // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph.h
+55
View File
@@ -0,0 +1,55 @@
/* 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_API_V1_COMPILE_TF_GRAPH_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V1_COMPILE_TF_GRAPH_H_
#include <variant>
#include <vector>
#include "absl/status/status.h"
#include "absl/types/variant.h"
#include "tensorflow/compiler/tf2xla/layout_util.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/client/compile_only_client.h"
#include "xla/pjrt/proto/compile_options.pb.h"
#include "xla/shape.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
namespace tensorflow {
namespace tf2xla {
namespace v1 {
// Compiles the given Tensorflow graph into xla::HLO. The result is in
// compilation_result. If the input computation is in MLIR, it will be
// converted to a Tensorflow graph. Otherwise, the graph compiler will be run.
absl::Status CompileTensorflowGraphToHlo(
const std::variant<tpu::MlirToHloArgs, tpu::FunctionToHloArgs>& computation,
const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,
XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_funcs,
const std::vector<tensorflow::TensorShape>& arg_shapes,
tsl::DeviceType device_type,
std::vector<tpu::ShardingAndIndex>* arg_core_mapping,
std::vector<std::vector<xla::Shape>>* per_core_arg_shapes,
xla::CompileOnlyClient* client,
XlaCompiler::CompilationResult* compilation_result);
} // namespace v1
} // namespace tf2xla
}; // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V1_COMPILE_TF_GRAPH_H_
tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph_test.cc
+272
View File
@@ -0,0 +1,272 @@
/* 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/api/v1/compile_tf_graph.h"
#include <cstdint>
#include <string>
#include <variant>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/deserialize_mlir_module_utils.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/test_matchers.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/utils/test_metadata_config.h"
#include "tensorflow/compiler/tf2xla/layout_util.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/client/client_library.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/stream_executor/platform.h"
#include "xla/stream_executor/platform_manager.h"
#include "xla/tsl/framework/device_type.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/lib/monitoring/test_utils.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/function.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/versions.pb.h"
#include "tensorflow/core/lib/monitoring/cell_reader.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
#include "tensorflow/core/util/debug_data_dumper.h"
namespace tensorflow {
namespace tf2xla {
namespace v1 {
namespace {
using ::tensorflow::monitoring::testing::CellReader;
using ::tensorflow::tpu::FunctionToHloArgs;
using ::tensorflow::tpu::MlirToHloArgs;
using ::tensorflow::tpu::ShardingAndIndex;
using ::tsl::monitoring::testing::Histogram;
static constexpr char kCompilationTimeStreamzName[] =
"/tensorflow/core/tf2xla/api/v1/phase2_compilation_time";
static constexpr char kCompilationStatusStreamzName[] =
"/tensorflow/core/tf2xla/api/v1/phase2_compilation_status";
static constexpr char kPlatformName[] = "Host";
constexpr char kEntryFuncName[] = "main";
static constexpr char kMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> () {
func.return
}
})";
MlirToHloArgs CreateTestMlirToHloArgs(const char* module_str = kMlirModuleStr) {
MlirToHloArgs mlir_to_hlo_args;
mlir_to_hlo_args.rollout_state =
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_DISABLED;
mlir_to_hlo_args.mlir_module = module_str;
return mlir_to_hlo_args;
}
absl::StatusOr<XlaCompilationResult> CompileWithComputation(
std::variant<tpu::MlirToHloArgs, tpu::FunctionToHloArgs> computation,
bool use_serialized_mlir) {
XlaCompilationResult compilation_result;
se::Platform* platform =
se::PlatformManager::PlatformWithName(kPlatformName).value();
auto client =
xla::ClientLibrary::GetOrCreateCompileOnlyClient(platform).value();
bool use_tuple_args = true;
std::vector<ShardingAndIndex> arg_core_mapping;
std::vector<std::vector<xla::Shape>> per_core_arg_shapes;
tpu::TPUCompileMetadataProto metadata_proto;
std::vector<TensorShape> arg_shapes;
if (computation.index() == 0) {
TF_RETURN_IF_ERROR(tensorflow::tf2xla::internal::ConfigureMetadata(
std::get<0>(computation).mlir_module, arg_shapes, metadata_proto));
}
XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns;
mlir::DialectRegistry registry;
mlir::RegisterAllTensorFlowDialects(registry);
mlir::mhlo::registerAllMhloDialects(registry);
mlir::MLIRContext context(registry);
mlir::OwningOpRef<mlir::ModuleOp> mlir_module;
if (std::holds_alternative<MlirToHloArgs>(computation) &&
!use_serialized_mlir) {
TF_RETURN_IF_ERROR(DeserializeMlirModule(
std::get<0>(computation).mlir_module, &context, &mlir_module));
std::get<0>(computation).mlir_module_op = mlir_module.get();
std::get<0>(computation).mlir_module = "";
}
absl::Status compilation_status =
tensorflow::tf2xla::v1::CompileTensorflowGraphToHlo(
computation, metadata_proto, use_tuple_args, shape_determination_fns,
arg_shapes, tsl::DeviceType("XLA_TPU_JIT"), &arg_core_mapping,
&per_core_arg_shapes, client, &compilation_result);
if (!compilation_status.ok()) return compilation_status;
return compilation_result;
}
using CompileTFGraphTest = ::testing::TestWithParam<bool>;
TEST_P(CompileTFGraphTest, RecordsStreamzForMlirFallback) {
bool use_serialized_mlir = GetParam();
CellReader<Histogram> compilation_time(kCompilationTimeStreamzName);
MlirToHloArgs mlir_to_hlo_args = CreateTestMlirToHloArgs();
TF_EXPECT_OK(
CompileWithComputation(mlir_to_hlo_args, use_serialized_mlir).status());
Histogram histogram = compilation_time.Delta("graph_old_bridge_has_mlir");
EXPECT_EQ(histogram.num(), 1);
}
TEST_P(CompileTFGraphTest, RecordsStreamzForFunctionToHlo) {
bool use_serialized_mlir = GetParam();
CellReader<Histogram> compilation_time(kCompilationTimeStreamzName);
CellReader<int64_t> compilation_status(kCompilationStatusStreamzName);
FunctionDef empty_function =
tensorflow::FunctionDefHelper::Create("empty", {}, {}, {}, {}, {});
tensorflow::FunctionDefLibrary fdef;
*(fdef.add_function()) = empty_function;
tensorflow::FunctionLibraryDefinition flib_def(
tensorflow::OpRegistry::Global(), fdef);
OpInputList guaranteed_constants;
NameAttrList function;
function.set_name("empty");
FunctionToHloArgs function_to_hlo_args = {&function,
&flib_def,
/*graph_def_version=*/0,
{&guaranteed_constants}};
TF_EXPECT_OK(CompileWithComputation(function_to_hlo_args, use_serialized_mlir)
.status());
Histogram histogram =
compilation_time.Delta("graph_old_bridge_has_function_to_hlo");
EXPECT_EQ(histogram.num(), 1);
EXPECT_EQ(compilation_status.Delta("kOldBridgeNoMlirSuccess"), 1);
}
TEST_P(CompileTFGraphTest, SuccessfullyCompilesWithManualSharding) {
bool use_serialized_mlir = GetParam();
// MLIR module from failing test.
constexpr char kSupportedManualSharding[] = R"(
module @module___inference_tpu_function_41 attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1617 : i32}} {
func.func @main(%arg0: tensor<2x2xf32>) -> (tensor<2x2xf32> {mhlo.sharding = "\08\03\1A\02\02\01\22\02\00\01"}) {
%0 = tf_executor.graph {
%outputs, %control = tf_executor.island wraps "tf.XlaSharding"(%arg0) {_XlaSharding = "\08\03\1A\02\02\01\22\02\00\01", sharding = "\08\03\1A\02\02\01\22\02\00\01"} : (tensor<2x2xf32>) -> tensor<2x2xf32>
%outputs_0, %control_1 = tf_executor.island wraps "tf.XlaSharding"(%outputs) {_XlaSharding = "\08\03\1A\02\02\01\22\02\00\01", sharding = "\08\03\1A\02\02\01\22\02\00\01", unspecified_dims = []} : (tensor<2x2xf32>) -> tensor<2x2xf32>
%outputs_2, %control_3 = tf_executor.island wraps "tf.XlaSpmdFullToShardShape"(%outputs_0) {dim = -1 : i64, manual_sharding = "\08\03\1A\02\02\01\22\02\00\01", unspecified_dims = []} : (tensor<2x2xf32>) -> tensor<1x2xf32>
%control_4 = tf_executor.island wraps "tf._XlaHostComputeMlir"(%outputs_2) {host_mlir_module = "", manual_sharding = true, recv_key = "host_compute_channel_0_retvals", send_key = "host_compute_channel_0_args"} : (tensor<1x2xf32>) -> ()
%outputs_5, %control_6 = tf_executor.island(%control_4) wraps "tf._XlaHostComputeMlir"() {host_mlir_module = "module {\0A func.func @host_func() -> tensor<1x2xf32> {\0A %0 = \22tf.Const\22() {value = dense<0.1> : tensor<1x2xf32>} : () -> tensor<1x2xf32> \0A return %0 : tensor<1x2xf32>}}", manual_sharding = true, recv_key = "host_compute_channel_1_retvals", send_key = "host_compute_channel_1_args"} : () -> tensor<1x2xf32>
%outputs_7, %control_8 = tf_executor.island wraps "tf.XlaSpmdShardToFullShape"(%outputs_5) {dim = -1 : i64, full_shape = #tf_type.shape<2x2>, manual_sharding = "\08\03\1A\02\02\01\22\02\00\01", unspecified_dims = []} : (tensor<1x2xf32>) -> tensor<2x2xf32>
%outputs_9, %control_10 = tf_executor.island wraps "tf.XlaSharding"(%outputs_7) {_XlaSharding = "\08\03\1A\02\02\01\22\02\00\01", sharding = "\08\03\1A\02\02\01\22\02\00\01", unspecified_dims = []} : (tensor<2x2xf32>) -> tensor<2x2xf32>
tf_executor.fetch %outputs_9 : tensor<2x2xf32>
}
return %0 : tensor<2x2xf32>
}
}
)";
auto mlir_to_hlo_args = CreateTestMlirToHloArgs(kSupportedManualSharding);
auto result = CompileWithComputation(mlir_to_hlo_args, use_serialized_mlir);
EXPECT_TRUE(result.ok());
}
TEST_P(CompileTFGraphTest, DoesNotInlineStatelessRandomOps) {
bool use_serialized_mlir = GetParam();
static constexpr char kHasReturnValues[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> (tensor<32x64xf32> {mhlo.sharding = "\08\01\1A\01\01\22\01\00"}) {
%cst = "tf.Const"() {value = dense<[524170, 523952]> : tensor<2xi32>} : () -> tensor<2xi32>
%cst_0 = "tf.Const"() {value = dense<[32, 64]> : tensor<2xi32>} : () -> tensor<2xi32>
%0 = "tf.StatelessRandomNormal"(%cst_0, %cst) : (tensor<2xi32>, tensor<2xi32>) -> tensor<32x64xf32>
return %0 : tensor<32x64xf32>
}
})";
auto compilation_result = CompileWithComputation(
CreateTestMlirToHloArgs(kHasReturnValues), use_serialized_mlir);
EXPECT_TRUE(compilation_result.ok());
// The StatelessRandomNormal must not be replaced by a literal tensor.
EXPECT_THAT(compilation_result,
ComputationProtoContains("tf.StatelessRandomNormal"));
}
TEST_P(CompileTFGraphTest, TestRunsShapeInference) {
bool use_serialized_mlir = GetParam();
static constexpr char kShapeInferenceModule[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> () {
%0 = "tf.Const"() <{value = dense<-1> : tensor<3360x8xi32>}> : () -> tensor<3360x8xi32>
%cst_33 = "tf.Const"() <{value = dense<[1120, -1]> : tensor<2xi32>}> : () -> tensor<2xi32>
%cst_34 = "tf.Const"() <{value = dense<[3, 1120, -1]> : tensor<3xi32>}> : () -> tensor<3xi32>
%cst_63 = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>
%1965:4 = "tf._XlaHostComputeMlir"(%0, %cst_34, %cst_63, %cst_33) <{host_mlir_module = "#loc1 = loc(\22Reshape:\22)\0A#loc2 = loc(\22Reshape_4\22)\0A#loc3 = loc(\22Reshape\22)\0A#loc9 = loc(fused[#loc1, #loc2, #loc3])\0Amodule {\0A func.func @host_func(%arg0: tensor<3360x?xi32> loc(fused[#loc1, #loc2, #loc3]), %arg1: tensor<3xi32> loc(fused[#loc1, #loc2, #loc3]), %arg2: tensor<i32> loc(fused[#loc1, #loc2, #loc3]), %arg3: tensor<2xi32> loc(fused[#loc1, #loc2, #loc3])) -> (tensor<1x1120x?xi32>, tensor<1x1120x?xi32>, tensor<1120x?xi32>, tensor<2xi32>) {\0A %0 = \22tf.Reshape\22(%arg0, %arg1) {_xla_outside_compilation = \220\22} : (tensor<3360x?xi32>, tensor<3xi32>) -> tensor<3x1120x?xi32> loc(#loc9)\0A %1:3 = \22tf.Split\22(%arg2, %0) {_xla_outside_compilation = \220\22} : (tensor<i32>, tensor<3x1120x?xi32>) -> (tensor<1x1120x?xi32>, tensor<1x1120x?xi32>, tensor<1x1120x?xi32>) loc(#loc10)\0A %2 = \22tf.Reshape\22(%1#0, %arg3) {_xla_outside_compilation = \220\22} : (tensor<1x1120x?xi32>, tensor<2xi32>) -> tensor<1120x?xi32> loc(#loc11)\0A %3 = \22tf.Shape\22(%2) {_xla_outside_compilation = \220\22} : (tensor<1120x?xi32>) -> tensor<2xi32> loc(#loc12)\0A return %1#1, %1#2, %2, %3 : tensor<1x1120x?xi32>, tensor<1x1120x?xi32>, tensor<1120x?xi32>, tensor<2xi32> loc(#loc9)\0A } loc(#loc9)\0A} loc(#loc)\0A#loc = loc(unknown)\0A#loc4 = loc(\22Split:\22)\0A#loc5 = loc(\22split\22)\0A#loc6 = loc(\22Reshape_5\22)\0A#loc7 = loc(\22Shape:\22)\0A#loc8 = loc(\22Shape_4\22)\0A#loc10 = loc(fused[#loc4, #loc5])\0A#loc11 = loc(fused[#loc1, #loc6])\0A#loc12 = loc(fused[#loc7, #loc8])\0A", recv_key = "host_compute_channel_0_retvals", send_key = "host_compute_channel_0_args"}> : (tensor<3360x8xi32>, tensor<3xi32>, tensor<i32>, tensor<2xi32>) -> (tensor<1x1120x?xi32>, tensor<1x1120x?xi32>, tensor<1120x?xi32>, tensor<2xi32>)
return
}
}
)";
auto compilation_result = CompileWithComputation(
CreateTestMlirToHloArgs(kShapeInferenceModule), use_serialized_mlir);
EXPECT_TRUE(compilation_result.ok());
}
INSTANTIATE_TEST_SUITE_P(CompileTFGraphTest, CompileTFGraphTest,
::testing::Bool());
} // namespace
} // namespace v1
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v1/mlir_bridge_config_v1.proto
+26
View File
@@ -0,0 +1,26 @@
// 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.
// ==============================================================================
// This proto defines messages for TF2XLA MLIR bridge configuration.
syntax = "proto2";
package tensorflow.tf2xla.v1;
message MlirBridgeConfig {
// Whether to enable tf2xla mlir bridge in compiling SavedModel. By default,
// it is enabled.
optional bool enable_tf2xla_mlir_bridge = 1 [default = true];
}
tensorflow/compiler/mlir/tf2xla/api/v1/testdata/empty_func.mlir
Vendored
+20
View File
@@ -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.
// ==============================================================================
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> () {
func.return
}
}
tensorflow/compiler/mlir/tf2xla/api/v1/testdata/invalid_executor.mlir
Vendored
+32
View File
@@ -0,0 +1,32 @@
// 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.
// ==============================================================================
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() {
tf_executor.graph {
%control = tf_executor.island {
tf_executor.yield
}
tf_executor.fetch %control : !tf_executor.control
}
tf_executor.graph {
%control = tf_executor.island {
tf_executor.yield
}
tf_executor.fetch %control : !tf_executor.control
}
return
}
}
tensorflow/compiler/mlir/tf2xla/api/v1/testdata/multiple_submodules.mlir
Vendored
+24
View File
@@ -0,0 +1,24 @@
// 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.
// ==============================================================================
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> () {
func.return
}
module @submodule attributes {tf.devices = {"/job:localhost/replica:0/task:0/device:CPU:0", "/job:localhost/replica:0/task:0/device:TPU:0", "/job:localhost/replica:0/task:0/device:TPU:1", "/job:localhost/replica:0/task:0/device:TPU_SYSTEM:0"}, tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1647 : i32}} {
}
}
tensorflow/compiler/mlir/tf2xla/api/v1/testdata/prepare_to_library.mlir
Vendored
+25
View File
@@ -0,0 +1,25 @@
// 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.
// ==============================================================================
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> () {
%0 = "tf.Const"() <{value = dense<-1> : tensor<3360x8xi32>}> : () -> tensor<3360x8xi32>
%cst_33 = "tf.Const"() <{value = dense<[1120, -1]> : tensor<2xi32>}> : () -> tensor<2xi32>
%cst_34 = "tf.Const"() <{value = dense<[3, 1120, -1]> : tensor<3xi32>}> : () -> tensor<3xi32>
%cst_63 = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>
%1965:4 = "tf._XlaHostComputeMlir"(%0, %cst_34, %cst_63, %cst_33) <{host_mlir_module = "#loc1 = loc(\22Reshape:\22)\0A#loc2 = loc(\22Reshape_4\22)\0A#loc3 = loc(\22Reshape\22)\0A#loc9 = loc(fused[#loc1, #loc2, #loc3])\0Amodule {\0A func.func @host_func(%arg0: tensor<3360x?xi32> loc(fused[#loc1, #loc2, #loc3]), %arg1: tensor<3xi32> loc(fused[#loc1, #loc2, #loc3]), %arg2: tensor<i32> loc(fused[#loc1, #loc2, #loc3]), %arg3: tensor<2xi32> loc(fused[#loc1, #loc2, #loc3])) -> (tensor<1x1120x?xi32>, tensor<1x1120x?xi32>, tensor<1120x?xi32>, tensor<2xi32>) {\0A %0 = \22tf.Reshape\22(%arg0, %arg1) {_xla_outside_compilation = \220\22} : (tensor<3360x?xi32>, tensor<3xi32>) -> tensor<3x1120x?xi32> loc(#loc9)\0A %1:3 = \22tf.Split\22(%arg2, %0) {_xla_outside_compilation = \220\22} : (tensor<i32>, tensor<3x1120x?xi32>) -> (tensor<1x1120x?xi32>, tensor<1x1120x?xi32>, tensor<1x1120x?xi32>) loc(#loc10)\0A %2 = \22tf.Reshape\22(%1#0, %arg3) {_xla_outside_compilation = \220\22} : (tensor<1x1120x?xi32>, tensor<2xi32>) -> tensor<1120x?xi32> loc(#loc11)\0A %3 = \22tf.Shape\22(%2) {_xla_outside_compilation = \220\22} : (tensor<1120x?xi32>) -> tensor<2xi32> loc(#loc12)\0A return %1#1, %1#2, %2, %3 : tensor<1x1120x?xi32>, tensor<1x1120x?xi32>, tensor<1120x?xi32>, tensor<2xi32> loc(#loc9)\0A } loc(#loc9)\0A} loc(#loc)\0A#loc = loc(unknown)\0A#loc4 = loc(\22Split:\22)\0A#loc5 = loc(\22split\22)\0A#loc6 = loc(\22Reshape_5\22)\0A#loc7 = loc(\22Shape:\22)\0A#loc8 = loc(\22Shape_4\22)\0A#loc10 = loc(fused[#loc4, #loc5])\0A#loc11 = loc(fused[#loc1, #loc6])\0A#loc12 = loc(fused[#loc7, #loc8])\0A", recv_key = "host_compute_channel_0_retvals", send_key = "host_compute_channel_0_args"}> : (tensor<3360x8xi32>, tensor<3xi32>, tensor<i32>, tensor<2xi32>) -> (tensor<1x1120x?xi32>, tensor<1x1120x?xi32>, tensor<1120x?xi32>, tensor<2xi32>)
return
}
}
tensorflow/compiler/mlir/tf2xla/api/v1/tf_dialect_to_executor.cc
+193
View File
@@ -0,0 +1,193 @@
/* 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/api/v1/tf_dialect_to_executor.h"
#include <memory>
#include <string>
#include <utility>
#include "absl/log/log.h"
#include "absl/log/vlog_is_on.h"
#include "absl/status/status.h"
#include "llvm/ADT/StringRef.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/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Transforms/Passes.h" // from @llvm-project
#include "tensorflow/compiler/jit/flags.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/data_dumper_logger_config.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/logging_hooks.h"
#include "xla/tsl/lib/monitoring/counter.h"
#include "xla/tsl/platform/status.h"
#include "tensorflow/core/platform/error_payloads.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/core_platform_payloads.pb.h"
#include "tensorflow/core/util/debug_data_dumper.h"
#include "tsl/platform/error_logging.h"
namespace tensorflow {
namespace tf2xla {
namespace v1 {
using mlir::LogicalResult;
using mlir::ModuleOp;
using mlir::OpPassManager;
using mlir::Pass;
using mlir::PassManager;
using mlir::func::FuncOp;
auto *tf_dialect_to_executor_dialect_status = tsl::monitoring::Counter<1>::New(
"/tensorflow/core/tf2xla/api/v1/tf_dialect_to_executor_dialect_status",
"Counts how often a successful export from TF Dialect to Executor Dialect "
"is",
"status");
constexpr char kExportSuccess[] = "success";
constexpr char kExportFailed[] = "failed";
namespace {
void AddTfDialectToExecutorPasses(OpPassManager &pm) {
auto add_pass = [&](std::unique_ptr<Pass> pass) {
pm.addNestedPass<FuncOp>(std::move(pass));
pm.addPass(mlir::CreateBreakUpIslandsPass());
};
pm.addPass(mlir::tf_executor::CreateTFExecutorTPUV1IslandInliningPass());
// There are cases where we don't consume all compilation and
// replication attributes like we do for the V2 pipeline, so we need to
// convert them from unified to legacy attributes before they get
// exposed to outside of the bridge.
pm.addNestedPass<FuncOp>(
mlir::TFTPU::CreateConvertToLegacyCompileAndReplicateAttributesPass());
pm.addPass(mlir::TF::CreateTFRegionControlFlowToFunctional());
add_pass(mlir::CreateFunctionalToExecutorDialectConversionPass());
add_pass(mlir::TFDevice::CreateReplicateToIslandPass(
/*legacy_graph_export=*/true));
add_pass(mlir::TFDevice::CreateReplicaIDToDeviceOrdinalPass());
add_pass(mlir::TFDevice::CreateParallelExecuteToIslandsPass(
/*legacy_graph_export=*/true));
add_pass(mlir::TFDevice::CreateLaunchToDeviceAttributePass(
/*legacy_graph_export=*/true));
pm.addNestedPass<FuncOp>(mlir::TFTPU::CreateTPUDevicePropagationPass());
pm.addNestedPass<FuncOp>(mlir::TFTPU::CreateTPUColocateSplitsPass());
pm.addPass(mlir::createSymbolDCEPass());
if (tensorflow::GetMlirCommonFlags()
->tf_mlir_enable_convert_control_to_data_outputs_pass) {
bool composite_tpuexecute_side_effects =
tensorflow::GetMlirCommonFlags()
->tf_mlir_enable_composite_tpuexecute_side_effects;
pm.addPass(
mlir::tf_executor::CreateTFExecutorConvertControlToDataOutputsPass(
composite_tpuexecute_side_effects));
}
pm.addPass(mlir::TF::CreateVerifySuitableForExportPass());
}
absl::Status RecordStatusIfError(absl::Status status) {
if (status.ok()) {
return absl::OkStatus();
}
VLOG(1) << "Failed to export from TF Dialect to TF Executor Dialect. "
<< status;
tf_dialect_to_executor_dialect_status->GetCell(kExportFailed)->IncrementBy(1);
constexpr char bridge_subcomponent[] =
"TFXLA_TF_FUNCTIONAL_TO_EXECUTOR_EXPORT_v1";
constexpr char kBridgeComponent[] = "TFXLABridge";
tsl::OkOrSetErrorCounterPayload(
tensorflow::core::platform::ErrorSourceProto::MLIR_BRIDGE_PHASE_1,
status);
tsl::error_logging::Log(kBridgeComponent, bridge_subcomponent,
status.ToString())
.IgnoreError();
return status;
}
} // namespace
absl::Status ExportFromTensorflowDialectToExecutor(
ModuleOp module, llvm::StringRef module_name) {
PassManager tf_to_executor(module.getContext());
::tensorflow::applyTensorflowAndCLOptions(tf_to_executor);
tf_to_executor.enableVerifier();
AddTfDialectToExecutorPasses(tf_to_executor);
mlir::StatusScopedDiagnosticHandler diag_handler(
module.getContext(), /*propagate=*/false,
/*filter_stack=*/!VLOG_IS_ON(1));
if (VLOG_IS_ON(1) ||
DEBUG_DATA_DUMPER()->ShouldDump(module_name.str(), kDebugGroupMain)) {
::tensorflow::DumpMlirOpToFile(
DEBUG_DATA_DUMPER()->GetDumpFilename(
module_name.str(), kDebugGroupMain,
"tfxla_bridge_v1_tfdialect_to_executor_before"),
module, llvm::StringRef(), &tf_to_executor);
if (VLOG_IS_ON(2) ||
DEBUG_DATA_DUMPER()->ShouldDump(
module_name.str(), kDebugGroupBridgePhase1ExecutorExport)) {
internal::EnablePassIRPrinting(
tf_to_executor, kDebugGroupBridgePhase1ExecutorExport, module_name);
}
}
LogicalResult result = tf_to_executor.run(module);
(void)result; // Ignore the error since it's captured by the diag_handler.
if (VLOG_IS_ON(1) ||
DEBUG_DATA_DUMPER()->ShouldDump(module_name.str(), kDebugGroupMain)) {
::tensorflow::DumpMlirOpToFile(
DEBUG_DATA_DUMPER()->GetDumpFilename(
module_name.str(), kDebugGroupMain,
"tfxla_bridge_v1_tfdialect_to_executor_after"),
module, llvm::StringRef(), &tf_to_executor);
}
if (result.failed()) {
return RecordStatusIfError(diag_handler.ConsumeStatus());
}
tf_dialect_to_executor_dialect_status->GetCell(kExportSuccess)
->IncrementBy(1);
return diag_handler.ConsumeStatus();
}
// Registers a pipeline builder function for TF Graph export. Should be
// the same as ExportFromTensorflowDialectToExecutor just in PassRegistration
// form.
mlir::PassPipelineRegistration<> tf_dialect_to_executor_pipeline(
"tf-dialect-to-executor-v1",
"Run passes to convert from TF Dialect to Executor in preparation for "
"exporting module back to TF Graph.",
AddTfDialectToExecutorPasses);
} // namespace v1
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v1/tf_dialect_to_executor.h
+58
View File
@@ -0,0 +1,58 @@
/* 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_API_V1_TF_DIALECT_TO_EXECUTOR_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V1_TF_DIALECT_TO_EXECUTOR_H_
#include "absl/base/attributes.h"
#include "absl/status/status.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
namespace tf2xla {
namespace v1 {
// Given the input Module op that's in the Tensorflow Dialect, convert the MLIR
// module in place to the Tensorflow Executor Dialect. Returns an OK Status if
// success, otherwise failure with an error message.
// The Tensorflow Executor Dialect is required to export an MLIR module to a
// Tensorflow GraphDef. This API will add control dependencies and verify that
// the conversion was successful. This version adds extra control dependencies
// for replication and parallel execution ops, which may slow performance.
// Prefer to use the v2 of this API.
//
// This also converts the Tensorflow Dialect MLIR into the Tensorflow Executor
// dialect that is suitable to be exported to GraphDef. Graph -> MLIR -> Graph
// is not perfectly round trippable, so this API will attempt to make the module
// exportable and verify some properties of the Tensorflow Executor MLIR that
// are required by Graph Export. It will return an error if it cannot.
//
// Input: A MLIR Module in the Tensorflow Dialect with no
// `tf_device.cluster_func` ops.
// Output: A MLIR module in the Tensorflow Executor Dialect.
ABSL_DEPRECATED(
"Use v2/tf_dialect_to_executor.h::ExportFromTensorflowDialectToExecutor "
"instead.")
absl::Status ExportFromTensorflowDialectToExecutor(
mlir::ModuleOp module, llvm::StringRef module_name = llvm::StringRef());
} // namespace v1
} // namespace tf2xla
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V1_TF_DIALECT_TO_EXECUTOR_H_
tensorflow/compiler/mlir/tf2xla/api/v1/tf_dialect_to_executor_test.cc
+104
View File
@@ -0,0 +1,104 @@
/* 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/api/v1/tf_dialect_to_executor.h"
#include <cstdint>
#include <string>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "mlir/IR/BuiltinOps.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 "mlir/Parser/Parser.h" // from @llvm-project
#include "tensorflow/compiler/mlir/register_common_dialects.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/platform/status.h"
#include "tensorflow/core/lib/monitoring/cell_reader.h"
#include "tensorflow/core/platform/resource_loader.h"
namespace tensorflow {
namespace tf2xla {
namespace v1 {
namespace {
using mlir::ModuleOp;
using mlir::DialectRegistry;
using mlir::MLIRContext;
using mlir::ModuleOp;
using mlir::OwningOpRef;
using tensorflow::monitoring::testing::CellReader;
static constexpr char kCompilationStreamz[] =
"/tensorflow/core/tf2xla/api/v1/tf_dialect_to_executor_dialect_status";
std::string TestDataPath() {
return tensorflow::GetDataDependencyFilepath(
"tensorflow/compiler/mlir/tf2xla/api/v1/testdata/");
}
class TensorflowDialectToExecutorTest : public ::testing::Test {
public:
TensorflowDialectToExecutorTest() {
mlir::RegisterCommonToolingDialects(registry_);
context_.appendDialectRegistry(registry_);
context_.loadAllAvailableDialects();
}
absl::Status CreateMlirModule(std::string mlir_module_filename) {
std::string mlir_module_path = TestDataPath() + mlir_module_filename;
mlir_module_ =
mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context_);
if (!mlir_module_) {
return absl::Status(
absl::StatusCode::kNotFound,
absl::StrCat("Could not find MLIR module at ", mlir_module_path));
}
return absl::OkStatus();
}
DialectRegistry registry_;
MLIRContext context_;
OwningOpRef<mlir::ModuleOp> mlir_module_;
};
TEST_F(TensorflowDialectToExecutorTest, ConvertsToExecutor) {
CellReader<int64_t> compilation_status(kCompilationStreamz);
TF_ASSERT_OK(CreateMlirModule("empty_func.mlir"));
TF_EXPECT_OK(ExportFromTensorflowDialectToExecutor(*mlir_module_));
EXPECT_EQ(compilation_status.Delta("success"), 1);
}
TEST_F(TensorflowDialectToExecutorTest, ErrorsWhenCannotConvert) {
CellReader<int64_t> compilation_status(kCompilationStreamz);
TF_ASSERT_OK(CreateMlirModule("invalid_executor.mlir"));
EXPECT_FALSE(ExportFromTensorflowDialectToExecutor(*mlir_module_).ok());
EXPECT_EQ(compilation_status.Delta("failed"), 1);
}
} // namespace
} // namespace v1
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v2/BUILD
+412
View File
@@ -0,0 +1,412 @@
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow/core/platform:build_config.bzl", "tf_proto_library")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
# Please reach out to tf-bridge-team@ before using the TF2XLA bridge.
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
":__subpackages__",
],
)
cc_library(
name = "legalize_tf",
srcs = ["legalize_tf.cc"],
hdrs = ["legalize_tf.h"],
visibility = [
"//learning/brain/google/xla:__pkg__",
"//learning/brain/mlir/bridge:__pkg__",
"//tensorflow/compiler/mlir/quantization/stablehlo:__pkg__",
"//tensorflow/compiler/mlir/tf2xla/api/v2/testing:__pkg__",
"//tensorflow/compiler/mlir/tfrt/transforms/ifrt:__pkg__",
],
deps = [
":device_type_proto_cc",
"//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
"//tensorflow/compiler/mlir/tensorflow:serialize_mlir_module_utils",
"//tensorflow/compiler/mlir/tf2xla/api/v1:compile_tf_graph",
"//tensorflow/compiler/mlir/tf2xla/internal:compilation_timer",
"//tensorflow/compiler/mlir/tf2xla/internal:legalize_tf_to_hlo",
"//tensorflow/compiler/mlir/tf2xla/internal:reproducer_proto_cc",
"//tensorflow/compiler/tf2xla:layout_util",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core/tpu/kernels:tpu_compile_op_support",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:string_view",
"@com_google_absl//absl/types:variant",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@tsl//tsl/platform:protobuf",
"@xla//xla:shape_util",
"@xla//xla:xla_proto_cc",
"@xla//xla/client:compile_only_client",
"@xla//xla/hlo/ir:hlo",
"@xla//xla/pjrt/proto:compile_options_proto_cc",
"@xla//xla/tsl/platform:statusor",
],
)
tf_cc_test(
name = "legalize_tf_test",
srcs = ["legalize_tf_test.cc"],
deps = [
":legalize_tf",
"//tensorflow/compiler/jit",
"//tensorflow/compiler/mlir/tf2xla/api/v2/testing:compile_mlir",
"//tensorflow/compiler/mlir/tf2xla/internal:test_matchers",
"//tensorflow/compiler/tf2xla:xla_compiler",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/lib/monitoring:cell_reader",
"//tensorflow/core/protobuf:for_core_protos_cc",
"//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
"//tensorflow/core/tpu/kernels:tpu_compile_op_support",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:str_format",
"@com_google_googletest//:gtest",
"@llvm-project//mlir:Pass",
"@xla//xla:shape_util",
"@xla//xla/client:client_library",
"@xla//xla/stream_executor:platform",
"@xla//xla/stream_executor:platform_manager",
"@xla//xla/tsl/lib/monitoring:test_utils",
"@xla//xla/tsl/platform:statusor",
],
)
tf_cc_test(
name = "legalize_tf_test_gpu",
srcs = ["legalize_tf_test_gpu.cc"],
tags = [
"config-cuda-only",
"no_oss", # This test only runs with GPU.
"requires-gpu-nvidia",
],
deps = [
"//tensorflow/compiler/jit",
"//tensorflow/compiler/mlir/tf2xla/api/v2/testing:compile_mlir",
"//tensorflow/compiler/tf2xla:xla_compiler",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/protobuf:for_core_protos_cc",
"//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_matchers",
"@com_google_absl//absl/status:statusor",
"@com_google_googletest//:gtest",
"@xla//xla/tsl/platform:statusor",
],
)
tf_proto_library(
name = "device_type_proto",
srcs = ["device_type.proto"],
visibility = [
"//learning/serving/contrib/tfrt/mlir/saved_model_analysis:__pkg__",
],
)
cc_library(
name = "cluster_tf",
srcs = ["cluster_tf.cc"],
hdrs = ["cluster_tf.h"],
visibility = [
"//learning/serving/contrib/tfrt/mlir/saved_model_analysis:__pkg__",
"//tensorflow/compiler/mlir/tfrt:__pkg__",
"//tensorflow/compiler/mlir/tfrt/transforms/ifrt:__pkg__",
"//tensorflow/compiler/tf2xla:__pkg__",
],
deps = [
":device_type_proto_cc",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
"//tensorflow/compiler/mlir/tensorflow:error_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow/transforms:verify_no_outside_compilation_markers_pass",
"//tensorflow/compiler/mlir/tf2xla/internal:clustering_bridge_passes",
"//tensorflow/compiler/mlir/tf2xla/internal:logging_hooks",
"//tensorflow/compiler/mlir/tf2xla/internal/passes:clustering_passes",
"//tensorflow/core:framework",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core/platform:error_payloads",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:stacktrace",
"//tensorflow/core/platform:status",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@tsl//tsl/platform:error_logging",
"@xla//xla/tsl/platform:errors",
],
)
tf_cc_test(
name = "cluster_tf_test",
srcs = ["cluster_tf_test.cc"],
data = [
"testdata/empty_func.mlir",
"testdata/invalid_executor.mlir",
"testdata/outside_compilation.mlir",
],
deps = [
":cluster_tf",
"//tensorflow/compiler/mlir:register_common_dialects",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:tf_dialect_lib",
"//tensorflow/compiler/mlir/tf2xla/api/v2/testing:utils",
"//tensorflow/core/lib/monitoring:cell_reader",
"//tensorflow/core/platform:resource_loader",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"@xla//xla/tsl/lib/core:status_test_util",
"@xla//xla/tsl/platform:status",
],
)
cc_library(
name = "tf_dialect_to_executor",
srcs = ["tf_dialect_to_executor.cc"],
hdrs = ["tf_dialect_to_executor.h"],
visibility = [
"//learning/serving/contrib/tfrt/mlir/saved_model_analysis:__pkg__",
"//tensorflow/compiler/mlir/tfrt:__pkg__",
"//tensorflow/compiler/tf2xla:__pkg__",
],
deps = [
"//tensorflow/compiler/jit:flags_headers",
"//tensorflow/compiler/mlir/tensorflow:bridge_logger",
"//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
"//tensorflow/compiler/mlir/tensorflow/transforms:verify_no_outside_compilation_markers_pass",
"//tensorflow/compiler/mlir/tf2xla/internal:logging_hooks",
"//tensorflow/core:framework",
"//tensorflow/core/platform:error_payloads",
"//tensorflow/core/platform:status",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:Transforms",
"@tsl//tsl/platform:error_logging",
"@xla//xla/tsl/lib/monitoring:counter",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:status",
],
)
tf_cc_test(
name = "tf_dialect_to_executor_test",
srcs = ["tf_dialect_to_executor_test.cc"],
data = [
"testdata/empty_func.mlir",
"testdata/func_with_dead_ops.mlir",
"testdata/invalid_executor.mlir",
],
deps = [
":tf_dialect_to_executor",
"//tensorflow/compiler/mlir:register_common_dialects",
"//tensorflow/compiler/mlir/tf2xla/api/v2/testing:utils",
"//tensorflow/core/lib/monitoring:cell_reader",
"//tensorflow/core/platform:resource_loader",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest_main",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"@xla//xla/tsl/lib/core:status_test_util",
],
)
cc_library(
name = "tf_executor_to_graph",
srcs = [
"tf_executor_to_graph.cc",
],
hdrs = [
"tf_executor_to_graph.h",
],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/compiler/mlir:op_or_arg_name_mapper",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:convert_type",
"//tensorflow/compiler/mlir/tensorflow:error_util",
"//tensorflow/compiler/mlir/tensorflow:export_utils",
"//tensorflow/compiler/mlir/tensorflow:translate_utils",
"//tensorflow/compiler/mlir/tensorflow:verify_suitable_for_graph_export",
"//tensorflow/compiler/mlir/tensorflow/translate:mlir_roundtrip_flags",
"//tensorflow/compiler/mlir/utils:name_utils",
"//tensorflow/core:core_cpu",
"//tensorflow/core:framework",
"//tensorflow/core:framework_internal",
"//tensorflow/core:graph",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/graph/regularization:util",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/container:inlined_vector",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:optional",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@xla//xla:status_macros",
],
)
tf_cc_test(
name = "tf_executor_to_graph_test",
srcs = ["tf_executor_to_graph_test.cc"],
data = [
"testdata/valid_executor.mlir",
"testdata/valid_graph.txt",
],
deps = [
":tf_executor_to_graph",
"//tensorflow/compiler/jit",
"//tensorflow/compiler/mlir:register_common_dialects",
"//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_flags",
"//tensorflow/compiler/tf2xla:xla_compiler",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test_main",
"//tensorflow/core/platform:resource_loader",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/status:statusor",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"@riegeli//riegeli/bytes:fd_reader",
"@riegeli//riegeli/bytes:read_all",
"@tsl//tsl/platform:protobuf",
"@xla//xla/tsl/lib/core:status_test_util",
],
)
cc_library(
name = "graph_to_tf_executor",
srcs = [
"graph_to_tf_executor.cc",
],
hdrs = [
"graph_to_tf_executor.h",
],
visibility = ["//visibility:public"],
deps = [
":mlir_roundtrip_flags",
"//tensorflow/compiler/jit:shape_inference_helpers",
"//tensorflow/compiler/mlir:op_or_arg_name_mapper",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:convert_attr",
"//tensorflow/compiler/mlir/tensorflow:convert_tensor",
"//tensorflow/compiler/mlir/tensorflow:convert_type",
"//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
"//tensorflow/compiler/mlir/tensorflow:dynamic_shape_utils",
"//tensorflow/compiler/mlir/tensorflow:error_util",
"//tensorflow/compiler/mlir/tensorflow:mangling_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_attributes",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow:translate_utils",
"//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
"//tensorflow/compiler/mlir/tensorflow/translate:mlir_roundtrip_flags",
"//tensorflow/compiler/mlir/tf2xla/internal:graph_to_tf_executor_util",
"//tensorflow/compiler/mlir/tf2xla/internal:node_order",
"//tensorflow/compiler/tf2xla:functionalize_control_flow",
"//tensorflow/compiler/tf2xla:functionalize_control_flow_util",
"//tensorflow/compiler/tf2xla:tf2xla_defs",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/common_runtime:function_body",
"//tensorflow/core/platform:types",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/container:inlined_vector",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:DerivedAttributeOpInterface",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@xla//xla:status_macros",
],
)
tf_cc_test(
name = "graph_to_tf_executor_test",
srcs = ["graph_to_tf_executor_test.cc"],
data = [
"testdata/graph_with_flib_def.txt",
"testdata/valid_graph.txt",
],
deps = [
":graph_to_tf_executor",
"//tensorflow/compiler/jit",
"//tensorflow/compiler/mlir:register_common_dialects",
"//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_flags",
"//tensorflow/compiler/mlir/tf2xla/internal:graph_to_tf_executor_util",
"//tensorflow/compiler/tf2xla:xla_compiler",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test_main",
"//tensorflow/core/platform:resource_loader",
"@com_google_absl//absl/status:statusor",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@riegeli//riegeli/bytes:fd_reader",
"@riegeli//riegeli/bytes:read_all",
"@tsl//tsl/platform:protobuf",
"@xla//xla/tsl/lib/core:status_test_util",
],
)
cc_library(
name = "mlir_roundtrip_flags",
hdrs = ["mlir_roundtrip_flags.h"],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
],
)
tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf.cc
+214
View File
@@ -0,0 +1,214 @@
/* 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/api/v2/cluster_tf.h"
#include <string>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "llvm/ADT/STLFunctionalExtras.h"
#include "llvm/ADT/StringRef.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/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
#include "tensorflow/compiler/mlir/tf2xla/api/v2/device_type.pb.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/clustering_bridge_passes.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/logging_hooks.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h"
#include "xla/tsl/platform/errors.h"
#include "tensorflow/core/framework/metrics.h"
#include "tensorflow/core/platform/error_payloads.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/stacktrace.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/util/debug_data_dumper.h"
#include "tsl/platform/error_logging.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
using mlir::LogicalResult;
using mlir::ModuleOp;
using mlir::OpPassManager;
using mlir::PassManager;
using mlir::func::FuncOp;
// Run the TF XLA Bridge based on the input pipeline, which can be either TPU
// bridge pipeline or non TPU bridge pipeline.
absl::Status RunTFXLABridge(
ModuleOp module,
llvm::function_ref<void(OpPassManager &pm)> pipeline_builder,
llvm::StringRef module_name = llvm::StringRef(),
llvm::StringRef dump_prefix = "tf_xla_bridge_v2") {
// Explicitly check that the TensorFlow dialect can constant fold ops.
// Constant folding is essential for the bridge. Without this check, the
// bridge may fail with an error that is difficult to understand and not
// actionable.
if (!mlir::TF::TensorFlowDialect::HasConstantFoldHook()) {
return absl::InternalError(
"TensorFlow dialect missing constant fold hook in TFXLA bridge phase "
"1; this could happen if the binary doesn't link the constant fold "
"hook registration library.");
}
PassManager bridge(module.getContext());
bridge.enableVerifier();
::tensorflow::applyTensorflowAndCLOptions(bridge);
// Populate a passmanager with the list of passes that implement the bridge.
pipeline_builder(bridge);
mlir::StatusScopedDiagnosticHandler diag_handler(
module.getContext(), /*propagate=*/false,
/*filter_stack=*/!VLOG_IS_ON(1));
if (VLOG_IS_ON(1) ||
DEBUG_DATA_DUMPER()->ShouldDump(module_name.str(), kDebugGroupMain)) {
::tensorflow::DumpMlirOpToFile(
DEBUG_DATA_DUMPER()->GetDumpFilename(module_name.str(), kDebugGroupMain,
dump_prefix.str() + "_before"),
module, llvm::StringRef(), &bridge);
}
if (VLOG_IS_ON(2) ||
DEBUG_DATA_DUMPER()->ShouldDump(module_name.str(),
kDebugGroupBridgePhase1Clustering)) {
::tensorflow::tf2xla::internal::EnablePassIRPrinting(
bridge, kDebugGroupBridgePhase1Clustering, module_name);
}
LogicalResult result = bridge.run(module);
(void)result;
if (VLOG_IS_ON(1) ||
DEBUG_DATA_DUMPER()->ShouldDump(module_name.str(), kDebugGroupMain)) {
::tensorflow::DumpMlirOpToFile(
DEBUG_DATA_DUMPER()->GetDumpFilename(module_name.str(), kDebugGroupMain,
dump_prefix.str() + "_after"),
module, llvm::StringRef(), &bridge);
}
return diag_handler.ConsumeStatus();
}
absl::Status RecordIfErrorStatus(const std::string error_prefix,
bool fallback_enabled, std::string bridge_type,
std::string device_type, absl::Status status) {
if (status.ok()) {
return status;
}
VLOG(2) << error_prefix << " " << status;
tensorflow::metrics::UpdateTfMlirBridgeFirstPhaseCounter(
/*bridge_type*/ bridge_type, /*bridge_version=*/"v2", device_type,
/*fallback_enabled=*/fallback_enabled,
/*result=*/"failure");
tsl::OkOrSetErrorCounterPayload(
tensorflow::core::platform::ErrorSourceProto::MLIR_BRIDGE_PHASE_1,
status);
std::string bridge_subcomponent = "TFXLA_PHASE_ONE_MLIR_TPU_BRIDGE";
if (device_type != "tpu") {
bridge_subcomponent = "TFXLA_PHASE_ONE_MLIR_CPU/GPU_BRIDGE";
}
tsl::error_logging::Log(mlir::TF::kBridgeComponent, bridge_subcomponent,
status.ToString())
.IgnoreError();
return status;
}
void CreateReplicatedClusteringPipeline(OpPassManager &pm,
llvm::StringRef module_name) {
// Since the internal bridge clustering passes are shared among TF1/TF2
// TF2-only passes should go here. However, this should be very rare and
// new passes generally should go into the internal
// AddReplicatedBridgeClusteringPipelinePasses.
pm.addPass(tensorflow::tf2xla::internal::CreateTPUValidateInputsPass());
pm.addNestedPass<FuncOp>(
mlir::TF::CreateCanonicalizeCompileAndReplicateAttributesPass());
tensorflow::tf2xla::internal::AddReplicatedBridgeClusteringPipelinePasses(
pm, module_name);
}
void CreateReplicatedClusteringPipelineV2(OpPassManager &pm) {
CreateReplicatedClusteringPipeline(pm, /*module_name=*/"");
}
absl::Status RunFunctionTf2xlaClusteringBridge(
ModuleOp module, bool is_supported_by_replicated_brige,
bool is_in_fallback_enabled_mode, llvm::StringRef module_name) {
std::string device_type = is_supported_by_replicated_brige
? mlir::TF::kMlirPh1BridgeCounterTpu
: mlir::TF::kMlirPh1BridgeCounterNonTpu;
VLOG(2)
<< (is_supported_by_replicated_brige ? "Replicated" : "NonReplicated")
<< " Bridge called stack trace is "
<< "(NOTE: this is not an error; rather the stack trace for debugging) : "
<< tensorflow::CurrentStackTrace();
absl::Status clustering_status =
is_supported_by_replicated_brige
? RunTFXLABridge(
module,
[module_name](OpPassManager &pm) {
CreateReplicatedClusteringPipeline(pm, module_name);
},
module_name, /*dump_prefix=*/"tf_xla_bridge_v2_replicated")
: RunTFXLABridge(
module,
[](OpPassManager &pm) {
tensorflow::tf2xla::internal::
AddNonReplicatedBridgeClusteringPipelinePasses(pm);
},
module_name, /*dump_prefix=*/"tf_xla_bridge_v2_nonreplicated");
std::string bridge_type = is_supported_by_replicated_brige
? mlir::TF::kMlirPh1BridgeCounterReplicated
: mlir::TF::kMlirPh1BridgeCounterNonReplicated;
// TODO(b/317798386): add is_supported_by_replicated_brige as a filter.
TF_RETURN_IF_ERROR(RecordIfErrorStatus(
/*error_prefix=*/"clustering_v2", is_in_fallback_enabled_mode,
bridge_type, device_type, clustering_status));
// TODO(b/317798386): add is_supported_by_replicated_brige as a filter.
tensorflow::metrics::UpdateTfMlirBridgeFirstPhaseCounter(
bridge_type, /*bridge_version=*/"v2", device_type,
/*fallback_enabled=*/is_in_fallback_enabled_mode,
/*result=*/"success");
return absl::OkStatus();
}
mlir::PassPipelineRegistration<> replicated_clustering_bridge_v2(
"tf-replicated-clustering-bridge-v2",
"Run all the passes involved in transforming a TensorFlow 2 graph before "
"execution so that it is suitable for targeting devices.",
CreateReplicatedClusteringPipelineV2);
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf.h
+61
View File
@@ -0,0 +1,61 @@
/* 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_API_V2_CLUSTER_TF_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_CLUSTER_TF_H_
#include "llvm/ADT/StringRef.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tf2xla/api/v2/device_type.pb.h"
#include "tensorflow/core/lib/core/status.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
// Run all the passes involved in transforming the graph before execution so
// that it is suitable for targeting devices when called with the TF 2 Function
// API. Users that need clustering with the Session API should use the v1 Bridge
// API. These transformations take as input a Tensorflow Graph as an MLIR Module
// and transforms the module in place to cluster the given ops for compilation
// that is compatible with the given device_type. The MLIR should be in the TF
// Executor Dialect for graph nodes and edges or be in TF Functional already.
// Individual Op inside a node should be the Tensorflow Functional Dialect. The
// output MLIR is in the TF Functional Dialect. Returns OkStatus if passed,
// otherwise an error.
//
// Inputs:
// module - The MLIR Module that will be clustered. Expected to be in TF
// Executor Dialect or TF Functional Dialect. Will convert to TF Functional.
// is_supported_by_replicated_brige - If the graph targets the replicated
// bridge. Set it to true for replicated/partitioned graphs. e.g. replicated
// and single-core TPU graphs. Set this to false if the graph is not
// replicated, e.g. CPU/GPU graphs. is_in_fallback_enabled_mode - Whether this
// was called with fallback to the non-MLIR Bridge. This is just for logging
// purposes and doesn't affect logic. module_name - What the input module name
// is for debugging help.
//
// Output: Modifies the input module in place with clustered operations.
// status - Whether the transformation to cluster the input MLIR module was
// successful.
absl::Status RunFunctionTf2xlaClusteringBridge(
mlir::ModuleOp module, bool is_supported_by_replicated_brige,
bool is_in_fallback_enabled_mode,
llvm::StringRef module_name = llvm::StringRef());
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_CLUSTER_TF_H_
tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf_test.cc
+168
View File
@@ -0,0 +1,168 @@
/* 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/api/v2/cluster_tf.h"
#include <cstdint>
#include <string>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.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 "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Parser/Parser.h" // from @llvm-project
#include "tensorflow/compiler/mlir/register_common_dialects.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h"
#include "tensorflow/compiler/mlir/tf2xla/api/v2/testing/utils.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/platform/status.h"
#include "tensorflow/core/lib/monitoring/cell_reader.h"
#include "tensorflow/core/platform/resource_loader.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
namespace {
using ::mlir::DialectRegistry;
using ::mlir::MLIRContext;
using ::mlir::ModuleOp;
using ::mlir::OwningOpRef;
using ::mlir::WalkResult;
using ::mlir::func::FuncOp;
using ::tensorflow::monitoring::testing::CellReader;
using ::tensorflow::tf2xla::v2::testing::TestDataPath;
static constexpr char kCompilationStreamz[] =
"/tensorflow/core/tf_mlir_bridge_first_phase_v2_count";
class FunctionClusterTensorflowDialectTest : public ::testing::Test {
public:
FunctionClusterTensorflowDialectTest() {
mlir::RegisterCommonToolingDialects(registry_);
context_.appendDialectRegistry(registry_);
context_.loadAllAvailableDialects();
}
absl::Status CreateMlirModule(std::string mlir_module_filename) {
std::string mlir_module_path = TestDataPath() + mlir_module_filename;
mlir_module_ =
mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context_);
if (!mlir_module_) {
return absl::Status(
absl::StatusCode::kNotFound,
absl::StrCat("Could not find MLIR module at ", mlir_module_path));
}
return absl::OkStatus();
}
DialectRegistry registry_;
MLIRContext context_;
OwningOpRef<mlir::ModuleOp> mlir_module_;
};
TEST_F(FunctionClusterTensorflowDialectTest, ClustersTfReplicatedBridge) {
CellReader<int64_t> compilation_status(kCompilationStreamz);
TF_ASSERT_OK(CreateMlirModule("empty_func.mlir"));
TF_EXPECT_OK(RunFunctionTf2xlaClusteringBridge(
*mlir_module_, /*is_supported_by_replicated_brige*/ true,
/*is_in_fallback_enabled_mode=*/false));
FuncOp main = mlir_module_->lookupSymbol<mlir::func::FuncOp>("main");
ASSERT_TRUE(main);
EXPECT_EQ(compilation_status.Delta(mlir::TF::kMlirPh1BridgeCounterReplicated,
mlir::TF::kMlirPh1BridgeCounterV2,
mlir::TF::kMlirPh1BridgeCounterTpu,
"fallback_disabled", "success"),
1);
}
TEST_F(FunctionClusterTensorflowDialectTest,
RunsOutsideCompilationReplicatedBridge) {
CellReader<int64_t> compilation_status(kCompilationStreamz);
TF_ASSERT_OK(CreateMlirModule("outside_compilation.mlir"));
TF_EXPECT_OK(RunFunctionTf2xlaClusteringBridge(
*mlir_module_, /*is_supported_by_replicated_brige*/ true,
/*is_in_fallback_enabled_mode=*/false));
FuncOp main = mlir_module_->lookupSymbol<mlir::func::FuncOp>("main");
ASSERT_TRUE(main);
bool has_cluster_op = false;
main.walk([&](mlir::tf_device::ClusterFuncOp cluster_op) {
has_cluster_op = true;
return WalkResult::advance();
});
EXPECT_TRUE(has_cluster_op);
EXPECT_EQ(compilation_status.Delta(mlir::TF::kMlirPh1BridgeCounterReplicated,
mlir::TF::kMlirPh1BridgeCounterV2,
mlir::TF::kMlirPh1BridgeCounterTpu,
"fallback_disabled", "success"),
1);
}
TEST_F(FunctionClusterTensorflowDialectTest, ClustersTFNonReplicatedBridge) {
CellReader<int64_t> compilation_status(kCompilationStreamz);
TF_ASSERT_OK(CreateMlirModule("empty_func.mlir"));
TF_EXPECT_OK(RunFunctionTf2xlaClusteringBridge(
*mlir_module_, /*is_supported_by_replicated_brige*/ false,
/*is_in_fallback_enabled_mode=*/false));
FuncOp main = mlir_module_->lookupSymbol<mlir::func::FuncOp>("main");
ASSERT_TRUE(main);
EXPECT_EQ(
compilation_status.Delta(mlir::TF::kMlirPh1BridgeCounterNonReplicated,
mlir::TF::kMlirPh1BridgeCounterV2,
mlir::TF::kMlirPh1BridgeCounterNonTpu,
"fallback_disabled", "success"),
1);
}
TEST_F(FunctionClusterTensorflowDialectTest, LogsFallbackMode) {
CellReader<int64_t> compilation_status(kCompilationStreamz);
TF_ASSERT_OK(CreateMlirModule("empty_func.mlir"));
TF_EXPECT_OK(RunFunctionTf2xlaClusteringBridge(
*mlir_module_, /*is_supported_by_replicated_brige*/ true,
/*is_in_fallback_enabled_mode=*/true));
EXPECT_EQ(compilation_status.Delta(mlir::TF::kMlirPh1BridgeCounterReplicated,
mlir::TF::kMlirPh1BridgeCounterV2,
mlir::TF::kMlirPh1BridgeCounterTpu,
"fallback_enabled", "success"),
1);
}
} // namespace
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v2/device_type.proto
+26
View File
@@ -0,0 +1,26 @@
// 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.
// ==============================================================================
syntax = "proto2";
package tensorflow.tf2xla.v2;
// The requested device type to compile for.
enum DeviceType {
DEVICE_TYPE_UNSPECIFIED = 0;
XLA_TPU_JIT = 1;
XLA_CPU_JIT = 2;
XLA_GPU_JIT = 3;
}
tensorflow/compiler/mlir/tf2xla/api/v2/graph_to_tf_executor.cc
+2844
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/mlir/tf2xla/api/v2/graph_to_tf_executor.h
+54
View File
@@ -0,0 +1,54 @@
/* 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_API_V2_GRAPH_TO_TF_EXECUTOR_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_GRAPH_TO_TF_EXECUTOR_H_
#include <optional>
#include <string>
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/graph_to_tf_executor_util.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/graph_debug_info.pb.h"
#include "tensorflow/core/graph/graph.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
inline constexpr absl::string_view kImportModelDefaultGraphFuncName = "main";
// Given a Graph, returns a MLIR module containing the graph, expressed with
// tf_executor dialect.
absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> ConvertGraphToTfExecutor(
const Graph& graph, const GraphDebugInfo& debug_info,
const FunctionLibraryDefinition& flib_def, const GraphImportConfig& specs,
mlir::MLIRContext* context,
std::unordered_map<std::string, std::string>* tf_name_to_mlir_name =
nullptr,
const ConfigProto& config_proto = {},
tensorflow::TF2XLABridgeVersion bridge_version =
tensorflow::TF2XLABridgeVersion::kNotBridgeUseCase,
std::optional<absl::string_view> module_name = std::nullopt);
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_GRAPH_TO_TF_EXECUTOR_H_
tensorflow/compiler/mlir/tf2xla/api/v2/graph_to_tf_executor_test.cc
+205
View File
@@ -0,0 +1,205 @@
/* 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/api/v2/graph_to_tf_executor.h"
#include <stdlib.h>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <gtest/gtest.h>
#include "absl/status/statusor.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.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 "riegeli/bytes/fd_reader.h" // from @riegeli
#include "riegeli/bytes/read_all.h" // from @riegeli
#include "tensorflow/compiler/mlir/register_common_dialects.h"
#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/graph_to_tf_executor_util.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/platform/resource_loader.h"
#include "tsl/platform/protobuf.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
namespace {
using mlir::DialectRegistry;
using mlir::MLIRContext;
constexpr char kGraphWithFlibDefFileName[] = "graph_with_flib_def.txt";
std::string TestDataPath() {
return tensorflow::GetDataDependencyFilepath(
"tensorflow/compiler/mlir/tf2xla/api/v2/testdata/");
}
class GraphToTfExecutorTest : public ::testing::Test {
public:
GraphToTfExecutorTest() {
mlir::RegisterCommonToolingDialects(registry_);
context_.appendDialectRegistry(registry_);
context_.loadAllAvailableDialects();
}
GraphDef CreateGraphDef(std::string graphdef_filename) {
std::string file_path = TestDataPath() + graphdef_filename;
std::string contents;
GraphDef graph_def;
auto status = riegeli::ReadAll(riegeli::FdReader(file_path), contents);
if (!status.ok()) {
return graph_def;
}
tsl::protobuf::TextFormat::ParseFromString(contents, &graph_def);
return graph_def;
}
int CountNumberOfFunctionsInModule(mlir::ModuleOp module) {
int count = 0;
for (auto unused : module.getOps<mlir::func::FuncOp>()) {
(void)unused; // Avoid unused variable warning
count++;
}
return count;
}
DialectRegistry registry_;
MLIRContext context_;
};
TEST_F(GraphToTfExecutorTest, BasicConvertGraphToTfExecutorPasses) {
Graph graph(OpRegistry::Global());
GraphDebugInfo debug_info;
FunctionLibraryDefinition flib_def(OpRegistry::Global(),
FunctionDefLibrary());
GraphImportConfig specs;
GraphDef graph_def = CreateGraphDef("valid_graph.txt");
GraphConstructorOptions opts;
TF_ASSERT_OK(ConvertGraphDefToGraph(opts, graph_def, &graph));
TF_ASSERT_OK(
ConvertGraphToTfExecutor(graph, debug_info, flib_def, specs, &context_));
}
TEST_F(
GraphToTfExecutorTest,
ConvertGraphToTfExecutorConvertAllFunctionsTrueConvertsAllFunctionsInFlibDef) {
Graph graph(OpRegistry::Global());
GraphDebugInfo debug_info;
FunctionLibraryDefinition flib_def(OpRegistry::Global(),
FunctionDefLibrary());
GraphDef graph_def = CreateGraphDef(kGraphWithFlibDefFileName);
GraphConstructorOptions opts;
TF_ASSERT_OK(ConvertGraphDefToGraph(opts, graph_def, &graph));
GraphImportConfig specs;
specs.convert_all_functions_to_mlir = true;
absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> result =
ConvertGraphToTfExecutor(graph, debug_info, graph.flib_def(), specs,
&context_);
// should equal main + 4 functions in flib_def
ASSERT_EQ(CountNumberOfFunctionsInModule(result->get()), 5);
}
TEST_F(
GraphToTfExecutorTest,
ConvertGraphToTfExecutorConvertAllFunctionsFalseOnlyConvertsFunctionsReferencedInGraph) {
Graph graph(OpRegistry::Global());
GraphDebugInfo debug_info;
FunctionLibraryDefinition flib_def(OpRegistry::Global(),
FunctionDefLibrary());
GraphDef graph_def = CreateGraphDef(kGraphWithFlibDefFileName);
GraphConstructorOptions opts;
TF_ASSERT_OK(ConvertGraphDefToGraph(opts, graph_def, &graph));
GraphImportConfig specs;
specs.convert_all_functions_to_mlir = false;
absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> result =
ConvertGraphToTfExecutor(graph, debug_info, graph.flib_def(), specs,
&context_);
// should equal main + 2 functions referenced by nodes in the graph via the
// "f" attr.
ASSERT_EQ(CountNumberOfFunctionsInModule(result->get()), 3);
}
TEST_F(GraphToTfExecutorTest,
ConvertGraphToTfExecutorPopulatesTfNameToMlirNameMap) {
Graph graph(OpRegistry::Global());
GraphDebugInfo debug_info;
FunctionLibraryDefinition flib_def(OpRegistry::Global(),
FunctionDefLibrary());
GraphDef graph_def = CreateGraphDef(kGraphWithFlibDefFileName);
GraphConstructorOptions opts;
TF_ASSERT_OK(ConvertGraphDefToGraph(opts, graph_def, &graph));
GraphImportConfig specs;
std::unordered_map<std::string, std::string> tf_name_to_mlir_name;
TF_ASSERT_OK(ConvertGraphToTfExecutor(graph, debug_info, graph.flib_def(),
specs, &context_,
&tf_name_to_mlir_name));
std::unordered_set<std::string> result_set;
for (const auto& pair : tf_name_to_mlir_name) {
result_set.insert(pair.first);
}
// Converted functions referenced by nodes in the graph via the
// "f" attr. These are before they are converted to their corresponding MLIR
// name.
std::unordered_set<std::string> expected_set = {
"__inference__traced_save_45", "__inference__traced_restore_57"};
EXPECT_EQ(result_set, expected_set);
}
TEST_F(GraphToTfExecutorTest, ConvertGraphToTfExecutorSetsModuleName) {
Graph graph(OpRegistry::Global());
GraphDebugInfo debug_info;
FunctionLibraryDefinition flib_def(OpRegistry::Global(),
FunctionDefLibrary());
GraphImportConfig specs;
GraphDef graph_def = CreateGraphDef("valid_graph.txt");
GraphConstructorOptions opts;
TF_ASSERT_OK(ConvertGraphDefToGraph(opts, graph_def, &graph));
const std::string kModuleName = "my_module_name";
absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> result =
ConvertGraphToTfExecutor(
graph, debug_info, flib_def, specs, &context_,
/*tf_name_to_mlir_name=*/nullptr,
/*config_proto=*/{},
/*bridge_version=*/
tensorflow::TF2XLABridgeVersion::kNotBridgeUseCase,
/*module_name=*/kModuleName);
TF_ASSERT_OK(result.status());
EXPECT_EQ(result->get().getName().value_or(""), kModuleName);
}
} // namespace
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.cc
+206
View File
@@ -0,0 +1,206 @@
/* 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/api/v2/legalize_tf.h"
#include <memory>
#include <string>
#include <variant>
#include <vector>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "absl/types/variant.h"
#include "llvm/ADT/ScopeExit.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/serialize_mlir_module_utils.h"
#include "tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/compilation_timer.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/reproducer.pb.h"
#include "tensorflow/compiler/tf2xla/layout_util.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/client/compile_only_client.h"
#include "xla/hlo/ir/hlo_module.h"
#include "xla/shape.h"
#include "xla/tsl/framework/device_type.h"
#include "xla/tsl/lib/monitoring/sampler.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/statusor.h"
#include "xla/xla.pb.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
#include "tensorflow/core/util/debug_data_dumper.h"
#include "tsl/platform/protobuf.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
using tpu::FunctionToHloArgs;
using tpu::MlirToHloArgs;
using tpu::ShardingAndIndex;
auto* phase2_bridge_compilation_time = tsl::monitoring::Sampler<1>::New(
{"/tensorflow/core/tf2xla/api/v2/phase2_compilation_time",
"The wall-clock time spent on executing graphs in milliseconds.",
"configuration"},
// Power of 1.5 with bucket count 45 (> 23 hours)
{tsl::monitoring::Buckets::Exponential(1, 1.5, 45)});
// Name of component for error logging. This name is fixed and required to
// enable logging.
constexpr char kBridgeComponent[] = "TFXLABridge";
constexpr char kFullBridge[] = "full_bridge";
namespace {
bool ShouldFallbackToGraphCompiler(
const std::variant<MlirToHloArgs, FunctionToHloArgs>& computation) {
if (computation.index() == 1) return true;
return std::get<0>(computation).rollout_state ==
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_DISABLED;
}
std::string ComputationToString(const MlirToHloArgs computation) {
if (computation.mlir_module_op.has_value()) {
return SerializeMlirModule(computation.mlir_module_op.value());
}
return std::string(computation.mlir_module);
}
void DumpComputationInput(
const tpu::TPUCompileMetadataProto& metadata,
const std::vector<tensorflow::TensorShape>& arg_shapes,
const std::variant<tpu::MlirToHloArgs, tpu::FunctionToHloArgs>
computation) {
if (!VLOG_IS_ON(2)) {
return;
}
tensorflow::mlir::tf2xla::internal::LegalizeMlirToHloReproducer reproducer;
*reproducer.mutable_compile_metadata() = metadata;
for (const auto& shape : arg_shapes) {
shape.AsProto(reproducer.add_input_shapes());
}
switch (computation.index()) {
case 0:
reproducer.set_mlir_module(ComputationToString(std::get<0>(computation)));
break;
case 1: {
auto input = std::get<1>(computation);
*reproducer.mutable_function_def_library() = input.flib_def->ToProto();
} break;
default:
VLOG(2) << "LegalizeMlirToHlo computation input: unknown";
break;
}
std::string string_reproducer;
tensorflow::protobuf::TextFormat::PrintToString(reproducer,
&string_reproducer);
DumpRawStringToFile("legalize_tf_reproducer.textproto", string_reproducer);
}
absl::Status DumpHloCompilationResult(
absl::string_view name, XlaCompilationResult* compilation_result) {
if (!VLOG_IS_ON(2) &&
!DEBUG_DATA_DUMPER()->ShouldDump(std::string(name), kDebugGroupMain)) {
return absl::OkStatus();
}
TF_ASSIGN_OR_RETURN(
auto hlo_module_config,
xla::HloModule::CreateModuleConfigFromProto(
compilation_result->computation->proto(), xla::DebugOptions()));
TF_ASSIGN_OR_RETURN(
std::unique_ptr<xla::HloModule> hlo_module,
xla::HloModule::CreateFromProto(compilation_result->computation->proto(),
hlo_module_config));
std::string all_computations;
for (auto computation : hlo_module->computations()) {
all_computations += computation->ToString() + "\n\n";
}
tensorflow::DumpRawStringToFile(name, all_computations);
return absl::OkStatus();
}
} // namespace
absl::StatusOr<tensorflow::XlaCompilationResult> LegalizeMlirToHlo(
const std::variant<tpu::MlirToHloArgs, tpu::FunctionToHloArgs>& computation,
const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,
llvm::StringRef device_type,
std::vector<std::unique_ptr<::mlir::Pass>>& custom_legalization_passes,
XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,
const std::vector<tensorflow::TensorShape>& arg_shapes,
std::vector<tpu::ShardingAndIndex>* arg_core_mapping,
std::vector<std::vector<xla::Shape>>* per_core_arg_shapes,
xla::CompileOnlyClient* client) {
CompilationTimer timer;
auto record_time = llvm::make_scope_exit([&timer] {
phase2_bridge_compilation_time->GetCell(kFullBridge)
->Add(timer.ElapsedCyclesInMilliseconds());
});
auto compilation_result = std::make_unique<XlaCompilationResult>();
DumpComputationInput(metadata, arg_shapes, computation);
// If there are no MLIR args, compile the given function in the library.
if (ShouldFallbackToGraphCompiler(computation)) {
TF_RETURN_IF_ERROR(tf2xla::v1::CompileTensorflowGraphToHlo(
computation, metadata, use_tuple_args, shape_determination_fns,
arg_shapes, tsl::DeviceType(device_type.str()), arg_core_mapping,
per_core_arg_shapes, client, compilation_result.get()));
DumpHloCompilationResult("legalize_tf_fallback.hlo",
compilation_result.get())
.IgnoreError();
return *compilation_result;
}
auto combined_bridge_status = internal::LegalizeTfToHlo(
std::get<0>(computation), metadata, use_tuple_args, device_type,
shape_determination_fns, arg_shapes, arg_core_mapping,
per_core_arg_shapes, custom_legalization_passes, client,
compilation_result.get());
if (combined_bridge_status.ok()) {
VLOG(1) << "Successfully compiled MLIR computation to XLA HLO using "
"Combined MLIR and XlaBuilder Bridge.";
DumpHloCompilationResult("legalize_tf_combined_bridge.hlo",
compilation_result.get())
.IgnoreError();
return *compilation_result;
}
return combined_bridge_status.status();
}
}; // namespace v2
}; // namespace tf2xla
}; // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.h
+67
View File
@@ -0,0 +1,67 @@
/* 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_API_V2_LEGALIZE_TF_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_LEGALIZE_TF_H_
#include <memory>
#include <variant>
#include <vector>
#include "absl/status/status.h"
#include "absl/types/variant.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tf2xla/api/v2/device_type.pb.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/client/compile_only_client.h"
#include "xla/pjrt/proto/compile_options.pb.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
// Legalizes the given mlir::Module into XLA HLO. If successful, returns the
// compiled XLA HLO. V1 of the tf2xla uses MLIR whereas V0 does not use MLIR.
//
// Inputs:
// computation - The MLIR module op. It currently takes in
// tpu::FunctionToHloArgs but this is deprecated. arg_shapes - The shapes of
// the arguments in module_op. device_type - The device type to compile for.
// use_tuple_args - Pack the incoming arg shapes into a single tuple.
// custom_legalization_passes - Extra passes to lower from TF -> MHLO.
// arg_shapes - The shapes of the args.
// arg_core_mapping - Which args go on which cores.
// per_core_arg_shapes - For each core, the shapes for each argument.
// client - The Xla Compilation client.
absl::StatusOr<tensorflow::XlaCompilationResult> LegalizeMlirToHlo(
const std::variant<tpu::MlirToHloArgs, tpu::FunctionToHloArgs>& computation,
const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,
llvm::StringRef device_type,
std::vector<std::unique_ptr<mlir::Pass>>& custom_legalization_passes,
XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,
const std::vector<tensorflow::TensorShape>& arg_shapes,
std::vector<tpu::ShardingAndIndex>* arg_core_mapping,
std::vector<std::vector<xla::Shape>>* per_core_arg_shapes,
xla::CompileOnlyClient* client);
}; // namespace v2
}; // namespace tf2xla
}; // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_LEGALIZE_TF_H_
tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf_test.cc
+367
View File
@@ -0,0 +1,367 @@
/* 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/api/v2/legalize_tf.h"
#include <cstdint>
#include <memory>
#include <string>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/statusor.h"
#include "absl/strings/str_format.h"
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tf2xla/api/v2/testing/compile_mlir.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/test_matchers.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/client/client_library.h"
#include "xla/shape.h"
#include "xla/stream_executor/platform.h"
#include "xla/stream_executor/platform_manager.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/lib/monitoring/test_utils.h"
#include "xla/tsl/platform/statusor.h"
#include "xla/tsl/platform/test.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/lib/monitoring/cell_reader.h"
#include "tensorflow/core/lib/monitoring/test_utils.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
#include "tensorflow/core/util/debug_data_dumper.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
using ::tensorflow::monitoring::testing::CellReader;
using tensorflow::tf2xla::v2::testing::CompileMlirModule;
using ::testing::Not;
using ::testing::TestWithParam;
using tpu::FunctionToHloArgs;
using tpu::ShardingAndIndex;
using tpu::TPUCompileMetadataProto;
using tsl::testing::TmpDir;
static constexpr char kCompilationTimeStreamzName[] =
"/tensorflow/core/tf2xla/api/v2/phase2_compilation_time";
static constexpr char kFullBridge[] = "full_bridge";
static constexpr char kCompilationStatusStreamzName[] =
"/tensorflow/core/tf2xla/api/v2/phase2_compilation_status";
static const char kMlirWithFallbackModeSuccess[] =
"kMlirWithFallbackModeSuccess";
static const char kMlirWithFallbackModeFailure[] =
"kMlirWithFallbackModeFailure";
static const char kOldBridgeMlirFilteredFailure[] =
"kOldBridgeMlirFilteredFailure";
static const char kOldBridgeWithFallbackModeFailure[] =
"kOldBridgeWithFallbackModeFailure";
static const char kOldBridgeMlirFilteredSuccess[] =
"kOldBridgeMlirFilteredSuccess";
static const char kOldBridgeWithFallbackModeSuccess[] =
"kOldBridgeWithFallbackModeSuccess";
static const char kMlirCombinedMlirSuccess[] = "kMlirCombinedMlirSuccess";
static const char kMlirCombinedMlirFailure[] = "kMlirCombinedMlirFailure";
static const char kMlirCombinedOldSuccess[] = "kMlirCombinedOldSuccess";
static const char kMlirCombinedOldFailure[] = "kMlirCombinedOldFailure";
static constexpr char kMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> () {
func.return
}
})";
// MLIR which should not legalize at all
static constexpr char kBadMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> () {
%0 = tf.Unknown() -> ()
func.return %0
}
})";
// MLIR which should be filtered by the MLIR bridge but fully legalize with the
// combined bridge.
static constexpr char kUnsupportedMlirBridgeModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> () {
%cst0 = "tf.Const"(){ value = dense<0> : tensor<3x5xi1>} : () -> tensor<3x5xi1>
%0 = "tf.Where"(%cst0) : (tensor<3x5xi1>) -> tensor<?x2xi64>
func.return
}
})";
TEST(LegalizeTFTest, RecordsStreamzForSuccessfulLegalizeWithMlirBridge) {
CellReader<int64_t> compilation_status(kCompilationStatusStreamzName);
TF_ASSERT_OK_AND_ASSIGN(
XlaCompiler::CompilationResult result,
CompileMlirModule(
kMlirModuleStr,
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED));
// May have been filtered so check for lack of failure instead of success.
EXPECT_EQ(compilation_status.Delta(kMlirWithFallbackModeFailure), 0);
}
TEST(LegalizeTFTest, MatMul) {
static constexpr char kMatMulModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> (tensor<5x11xf32>) {
%arg0 = "tf.Const"() {value = dense<-3.0> : tensor<5x7xf32>} : () -> tensor<5x7xf32>
%arg1 = "tf.Const"() {value = dense<-3.0> : tensor<11x7xf32>} : () -> tensor<11x7xf32>
%1 = "tf.MatMul"(%arg0, %arg1) {transpose_a = false, transpose_b = true} : (tensor<5x7xf32>, tensor<11x7xf32>) -> tensor<5x11xf32>
func.return %1 : tensor<5x11xf32>
}
})";
TF_ASSERT_OK_AND_ASSIGN(
XlaCompiler::CompilationResult result,
CompileMlirModule(
kMatMulModuleStr,
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED));
}
struct MatMulTestCase {
std::string mat_mul_method;
};
using BatchMatMulTest = TestWithParam<MatMulTestCase>;
TEST_P(BatchMatMulTest, BatchMatMul) {
const MatMulTestCase& test_case = GetParam();
static constexpr char kMatMulModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> (tensor<1x4x4xf32>) {
%%arg0 = "tf.Const"() {value = dense<-3.0> : tensor<1x4x2xf32>} : () -> tensor<1x4x2xf32>
%%arg1 = "tf.Const"() {value = dense<-3.0> : tensor<1x2x4xf32>} : () -> tensor<1x2x4xf32>
%%1 = "tf.%s"(%%arg0, %%arg1) {T = f32, adj_x = false, adj_y = false, grad_x = false, grad_y = false, device = ""} : (tensor<1x4x2xf32>, tensor<1x2x4xf32>) -> tensor<1x4x4xf32>
func.return %%1 : tensor<1x4x4xf32>
}
})";
std::string mat_mul_method =
absl::StrFormat(kMatMulModuleStr, test_case.mat_mul_method);
TF_ASSERT_OK_AND_ASSIGN(
XlaCompiler::CompilationResult result,
CompileMlirModule(
mat_mul_method.c_str(),
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED));
}
INSTANTIATE_TEST_SUITE_P(
BatchMatMulTest, BatchMatMulTest,
::testing::ValuesIn<MatMulTestCase>({
{"BatchMatMul"},
{"BatchMatMulV2"},
{"BatchMatMulV3"},
}),
[](const ::testing::TestParamInfo<BatchMatMulTest::ParamType>& info) {
return info.param.mat_mul_method;
});
TEST(LegalizeTFTest, DumpsProducedHLO) {
Env* env = Env::Default();
std::string test_dir = TmpDir();
setenv("TF_DUMP_GRAPH_PREFIX", test_dir.c_str(), /*overwrite=*/1);
setenv("TF_DUMP_GRAPH_NAME_FILTER", "*", 1);
DEBUG_DATA_DUMPER()->LoadEnvvars();
std::vector<std::string> files;
TF_ASSERT_OK(env->GetChildren(test_dir, &files));
int original_files_size = files.size();
TF_ASSERT_OK_AND_ASSIGN(
XlaCompiler::CompilationResult result,
CompileMlirModule(
kMlirModuleStr,
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED));
// Due to the shared test of this infrastructure, we just need to make sure
// that the dumped file size is greater than what was originally inside
// the test directory.
TF_ASSERT_OK(env->GetChildren(test_dir, &files));
EXPECT_THAT(files.size(), ::testing::Gt(original_files_size));
setenv("TF_DUMP_GRAPH_PREFIX", test_dir.c_str(), /*overwrite=*/0);
}
TEST(LegalizeTFTest, RecordsStreamzForFailedLegalizeWithMlirBridge) {
CellReader<int64_t> compilation_status(kCompilationStatusStreamzName);
auto result = CompileMlirModule(
kBadMlirModuleStr,
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED);
EXPECT_FALSE(result.ok());
EXPECT_EQ(compilation_status.Delta(kMlirCombinedMlirFailure), 1);
}
TEST(LegalizeTFTest, RecordsStreamzForSuccessWithCombinedBridge) {
CellReader<int64_t> compilation_status(kCompilationStatusStreamzName);
auto result = CompileMlirModule(
kUnsupportedMlirBridgeModuleStr,
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED);
// MLIR Bridge will filter this unsupported MLIR, Combined will succeed.
EXPECT_TRUE(result.ok());
EXPECT_EQ(compilation_status.Delta(kMlirCombinedMlirSuccess), 1);
EXPECT_EQ(compilation_status.Delta(kMlirCombinedMlirFailure), 0);
EXPECT_EQ(compilation_status.Delta(kMlirCombinedOldSuccess), 1);
EXPECT_EQ(compilation_status.Delta(kMlirCombinedOldFailure), 0);
// Old bridge should never be called at all.
EXPECT_EQ(compilation_status.Delta(kOldBridgeMlirFilteredFailure), 0);
EXPECT_EQ(compilation_status.Delta(kOldBridgeWithFallbackModeFailure), 0);
EXPECT_EQ(compilation_status.Delta(kOldBridgeMlirFilteredSuccess), 0);
EXPECT_EQ(compilation_status.Delta(kOldBridgeWithFallbackModeSuccess), 0);
}
TEST(LegalizeTFTest, RecordsStreamzForNoMlirFallback) {
FunctionDef my_func =
tensorflow::FunctionDefHelper::Create("empty", {}, {}, {}, {}, {});
tensorflow::FunctionDefLibrary fdef;
*(fdef.add_function()) = my_func;
tensorflow::FunctionLibraryDefinition flib_def(
tensorflow::OpRegistry::Global(), fdef);
OpInputList guaranteed_constants;
NameAttrList function;
FunctionToHloArgs function_to_hlo_args{&function,
&flib_def,
/*graph_def_version=*/0,
{&guaranteed_constants}};
se::Platform* cpu_platform =
se::PlatformManager::PlatformWithName("Host").value();
auto client =
xla::ClientLibrary::GetOrCreateCompileOnlyClient(cpu_platform).value();
std::vector<TensorShape> arg_shapes;
TPUCompileMetadataProto metadata_proto;
bool use_tuple_args = true;
std::vector<ShardingAndIndex> arg_core_mapping;
std::vector<std::vector<xla::Shape>> per_core_arg_shapes;
std::vector<std::unique_ptr<mlir::Pass>> custom_legalization_passes;
// This doesn't actually compile correctly.
absl::StatusOr<XlaCompiler::CompilationResult> compile_result =
LegalizeMlirToHlo(function_to_hlo_args, metadata_proto, use_tuple_args,
/*device_type=*/"XLA_CPU_JIT",
custom_legalization_passes,
/*shape_determination_fns=*/{}, arg_shapes,
&arg_core_mapping, &per_core_arg_shapes, client);
EXPECT_FALSE(compile_result.ok());
}
TEST(LegalizeTFTest, RecordsCompilationTimeForSuccessfulCompilation) {
CellReader<monitoring::testing::Histogram> compilation_time(
kCompilationTimeStreamzName);
TF_ASSERT_OK_AND_ASSIGN(
XlaCompiler::CompilationResult result,
CompileMlirModule(
kMlirModuleStr,
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_ENABLED));
// Compilation time should have been updated.
EXPECT_GT(compilation_time.Delta(kFullBridge).num(), 0);
}
TEST(LegalizeTFTest, SuccessfullyCompilesModulesWithReturnValues) {
static constexpr char kHasReturnValuesAndNoMetadataRetvals[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> (tensor<2xi32>) {
%cst = "tf.Const"() {value = dense<[524170, 523952]> : tensor<2xi32>} : () -> tensor<2xi32>
return %cst : tensor<2xi32>
}
})";
auto compilation_result = CompileMlirModule(
kHasReturnValuesAndNoMetadataRetvals,
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED);
EXPECT_TRUE(compilation_result.ok());
// Ensure that the compilation result contains a constant.
EXPECT_THAT(compilation_result,
ComputationProtoContains("opcode:.*constant"));
}
TEST(LegalizeTFTest, SkipsTensorListSetItemIfDimensionsTooLarge) {
static constexpr char kTensorListSetItemDimensionTooLarge[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> tensor<!tf_type.variant<tensor<64x1xbf16>>> {
// unknown rank
%elem_shape = "tf.Const"() <{value = dense<-1> : tensor<i32>}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> tensor<i32>
// zero reserved elements
%num_elements = "tf.Const"() <{value = dense<0> : tensor<i32>}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> tensor<i32>
%list = "tf.TensorListReserve"(%elem_shape, %num_elements) : (tensor<i32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<64x1xbf16>>>
%index = "tf.Const"() <{value = dense<0> : tensor<i32>}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> tensor<i32>
%element = "tf.Const"() <{value = dense<0.0> : tensor<64x1xbf16>}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> tensor<64x1xbf16>
// Results in a bad mismatch of shapes.
%updated_list = "tf.TensorListSetItem"(%list, %index, %element) : (tensor<!tf_type.variant<tensor<64x1xbf16>>>, tensor<i32>, tensor<64x1xbf16>) -> tensor<!tf_type.variant<tensor<64x1xbf16>>>
return %updated_list : tensor<!tf_type.variant<tensor<64x1xbf16>>>
}
})";
auto compilation_result = CompileMlirModule(
kTensorListSetItemDimensionTooLarge,
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED);
// Ensure that it compile
ASSERT_TRUE(compilation_result.ok());
// Assert that the tensor list operation is lowered to something.
ASSERT_THAT(compilation_result,
Not(ComputationProtoContains("%.*= \"tf.TensorListSetItem")));
// Assert that the tensor list operation is lowered to something that doesn't
// get stuck on a broken dynamic update slice.
ASSERT_THAT(compilation_result,
Not(ComputationProtoContains("%.*=.*DynamicUpdateSlice")));
}
TEST(LegalizeTFTest, LegalizesFunctionWithBoundedDynamicArg) {
static constexpr char kMlirModuleWithBoundedDynamicArgStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main(%arg0: tensor<?xi32, #mhlo.type_extensions<bounds = [3]>> ) -> (tensor<?xi32, #mhlo.type_extensions<bounds = [3]>>) {
func.return %arg0 : tensor<?xi32, #mhlo.type_extensions<bounds = [3]>>
}
})";
auto compilation_result = CompileMlirModule(
kMlirModuleWithBoundedDynamicArgStr,
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED);
ASSERT_TRUE(compilation_result.ok());
EXPECT_THAT(compilation_result,
ComputationProtoContains("element_type:.S32\n.*dimensions: 3"));
}
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf_test_gpu.cc
+82
View File
@@ -0,0 +1,82 @@
/* 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 <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "absl/status/status_matchers.h"
#include "absl/status/statusor.h"
#include "tensorflow/compiler/mlir/tf2xla/api/v2/testing/compile_mlir.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
// These test are in a separate file because they requires separate build
// environments (--config=cuda).
using tensorflow::tf2xla::v2::testing::CompileMlirModule;
// MLIR which is legalize only with the right device type.
// The mlir is generated by running
// tensorflow/compiler/mlir/tf-opt -tf-xla-call-module-serialization
//
// module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32,
// producer = 268 : i32}} {
// func.func private @_jit_sin(%arg0: tensor<f32>) -> tensor<f32> {
// %0 = stablehlo.sine %arg0 : tensor<f32>
// return %0 : tensor<f32>
// }
// func.func @main(%arg0: tensor<f32>) -> tensor<*xf32> {
// %0 = "tf.XlaCallModule"(%arg0) {Sout = [#tf_type.shape<*>], device =
// "", dim_args_spec = [], _entry_function = @_jit_sin,
// _stablehlo_version = "1.0.0", module = "", platforms = ["CUDA"],
// version = 6 : i64} : (tensor<f32>) -> tensor<*xf32>
// func.return %0 : tensor<*xf32>
// }
// }
//
static constexpr char kGpuMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main(%arg0: tensor<f32>) -> tensor<*xf32> {
%0 = "tf.XlaCallModule"(%arg0) {Sout = [#tf_type.shape<*>], device = "", dim_args_spec = [], module = "ML\EFR\03StableHLOv1.0.0\00\01\17\05\01\05\01\03\05\03\07\07\t\0B\03K5\07\01\1B\07\0B\13\0B3\0B\0B\0B\0B\0F\0B\13\0B\03\1B\0F\1B\0B\0B\0B\0B\0B\0F\13\0B\0B\0B\0B\03\07\0F\17\07\02\A7\1F\05\0D\03\03\03\07\05\0F\03\0B\0B\1B\0D'\0F)\031\113\05\11\05\13\05\15\05\17\1D\15\17\05\19\17\19\EF\01\05\1B\03\03\1D\0D\05\1F!#%\1D\1D\1D\1F\1D!\1D##\03\03\03+\0D\03-/\1D%\1D'\1D)\1D+)\01\05\11\03\01\03\01\t\04A\05\01\11\01\05\07\03\01\05\03\11\01\t\05\03\05\0B\03\01\01\05\06\13\03\01\03\01\07\04\01\03\03\06\03\01\05\01\00\9A\04-\0F\0B\03!\1B\1D\05\1B\83/\1F\15\1D\15\11\13\15\11\11\0F\0B\11builtin\00vhlo\00module\00func_v1\00sine_v1\00return_v1\00sym_name\00jit_sin\00arg_attrs\00function_type\00res_attrs\00sym_visibility\00jit(sin)/jit(main)/sin\00third_party/py/jax/experimental/jax2tf/tests/back_compat_test.py\00jax.arg_info\00x\00mhlo.sharding\00{replicated}\00jax.result_info\00\00main\00public\00", platforms = ["CUDA"], version = 6 : i64} : (tensor<f32>) -> tensor<*xf32>
func.return %0 : tensor<*xf32>
}
})";
TEST(LegalizeTFTest, RightDeviceTypeShallPass) {
TF_ASSERT_OK_AND_ASSIGN(
XlaCompiler::CompilationResult result,
CompileMlirModule(
kGpuMlirModuleStr,
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED,
"XLA_GPU_JIT"));
}
TEST(LegalizeTFTest, WrongDeviceTypeShallFail) {
absl::StatusOr<XlaCompiler::CompilationResult> result = CompileMlirModule(
kGpuMlirModuleStr,
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED,
"XLA_TPU_JIT");
EXPECT_THAT(result, absl_testing::StatusIs(absl::StatusCode::kNotFound));
}
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v2/mlir_roundtrip_flags.h
+119
View File
@@ -0,0 +1,119 @@
/* Copyright 2019 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_API_V2_MLIR_ROUNDTRIP_FLAGS_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_MLIR_ROUNDTRIP_FLAGS_H_
#include <string>
#include <vector>
#include "absl/strings/string_view.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/StringMap.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
namespace tensorflow {
struct ArrayInfoBase {
// The node type when the input node is imported. Typically needs to be
// specified when passing arbitrary nodes (some node attributes are removed).
DataType imported_dtype;
// Node "shape" attribute value.
TensorShapeProto shape;
};
struct ArrayInfo : public ArrayInfoBase {
using SubTypeInfo = ArrayInfoBase;
// DT_RESOURCE and DT_VARIANT have subtypes
std::vector<SubTypeInfo> subtypes;
};
struct GraphImportConfig {
// Returns string representation of config.
std::string str() const;
using InputArrays =
llvm::MapVector<std::string, ArrayInfo, llvm::StringMap<unsigned>>;
// The name assigned to the function which is the import result of the given
// graph. If empty, a default one will be used.
std::string graph_func_name;
// Maps input node names to node data types and shapes.
InputArrays inputs;
// name:index strings for the data outputs.
std::vector<std::string> outputs;
// name strings for the control outputs.
std::vector<std::string> control_outputs;
// Setting prune_unused_nodes to true, would prune unreachable nodes if
// output_arrays is specified.
bool prune_unused_nodes = false;
// If true, inputs of type LegacyFedInput are replaced with Placeholder ops.
// LegacyFedInput ops have two outputs unlike Placeholder which has only one
// output, so if both outputs of the LegacyFedInput ops are used then returns
// an error.
bool convert_legacy_fed_inputs = false;
// If true, the main graph will be treated as a function.
bool graph_as_function = false;
// If true, upgrade legacy features of the graph (for instance, functionalize
// control-flow).
bool upgrade_legacy = false;
// If true, functionalization is restricted to nodes that will be
// XLA-compiled. This is only needed if
// - `upgrade_legacy` is true
// - upgrading legacy features of the graph (which includes functionalization)
// runs before compilation cluster extraction (as for MLIR-based TPU bridge)
// - session runtime is used (session runtime has issues with function names
// rewritten by functionalization).
// Otherwise, this parameter should be set to false.
bool restrict_functionalization_to_compiled_nodes = false;
// If true, enables shape inference on input.
// TODO(jpienaar): This will be removed shortly.
bool enable_shape_inference = true;
// _output_shapes is an unregistered attribute which is used during
// GraphConstructor::ConvertGraph to override shapes. It is unfortunately
// not always set correctly (which is undesirable and should be addressed)
// so make it opt-in to consider it unconditionally also when importing the
// graph.
bool unconditionally_use_set_output_shapes = false;
// If set, use the value as the device type and mark the function graph for
// XLA compilation.
std::string xla_compile_device_type;
// If true, enables moving ops to different devices or moving unsupported ops
// out of a compilation cluster.
bool enable_soft_placement = false;
// If true, a function attribute, `tf._original_func_name`, will be set in
// functions which contains the corresponding original TF function name.
bool set_original_tf_func_name = false;
// If true, all functions in the graph will be converted to MLIR regardless of
// whether the functions are referenced by the nodes. This is needed if
// aliases and saved model object graph function matching is needed.
bool convert_all_functions_to_mlir = false;
};
struct GraphExportConfig {
// Whether to export the entry function to function library instead of the
// graph.
bool export_entry_func_to_flib = false;
// Whether to export functions using the name set in the attribute
// `tf._original_func_name` if it exists.
bool export_original_tf_func_name = false;
};
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_MLIR_ROUNDTRIP_FLAGS_H_
tensorflow/compiler/mlir/tf2xla/api/v2/testdata/empty_func.mlir
Vendored
+20
View File
@@ -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.
// ==============================================================================
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> () {
func.return
}
}
tensorflow/compiler/mlir/tf2xla/api/v2/testdata/func_with_dead_ops.mlir
Vendored
+77
View File
@@ -0,0 +1,77 @@
// 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.
// ==============================================================================
module attributes {tf.devices = {"/job:tpu_host_worker/replica:0/task:0/device:CPU:0", "/job:tpu_host_worker/replica:0/task:0/device:TPU:0", "/job:tpu_host_worker/replica:0/task:0/device:TPU:1", "/job:tpu_host_worker/replica:0/task:0/device:TPU_SYSTEM:0", "/job:tpu_host_worker/replica:0/task:1/device:CPU:0", "/job:tpu_host_worker/replica:0/task:1/device:TPU:0", "/job:tpu_host_worker/replica:0/task:1/device:TPU:1", "/job:tpu_host_worker/replica:0/task:1/device:TPU_SYSTEM:0", "/job:tpu_host_worker/replica:0/task:2/device:CPU:0", "/job:tpu_host_worker/replica:0/task:2/device:TPU:0", "/job:tpu_host_worker/replica:0/task:2/device:TPU:1", "/job:tpu_host_worker/replica:0/task:2/device:TPU_SYSTEM:0", "/job:tpu_host_worker/replica:0/task:3/device:CPU:0", "/job:tpu_host_worker/replica:0/task:3/device:TPU:0", "/job:tpu_host_worker/replica:0/task:3/device:TPU:1", "/job:tpu_host_worker/replica:0/task:3/device:TPU_SYSTEM:0"}, tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1847 : i32}} {
func.func @main(%arg0: tensor<i32> {tf._user_specified_name = "steps", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg1: tensor<*x!tf_type.resource<tensor<i64>>> {tf._user_specified_name = "899", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg2: tensor<*x!tf_type.resource<tensor<i64>>> {tf._user_specified_name = "901", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg3: tensor<*x!tf_type.resource<tensor<128x1024xf32>>> {tf._user_specified_name = "903", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg4: tensor<*x!tf_type.resource<tensor<1024xf32>>> {tf._user_specified_name = "905", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg5: tensor<*x!tf_type.resource<tensor<1024x1xf32>>> {tf._user_specified_name = "907", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg6: tensor<*x!tf_type.resource<tensor<i64>>> {tf._user_specified_name = "909", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg7: tensor<*x!tf_type.resource<tensor<25001x64xf32>>> {tf._user_specified_name = "911", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg8: tensor<*x!tf_type.resource<tensor<25001x64xf32>>> {tf._user_specified_name = "913", tf.device = "/job:tpu_host_worker/replica:0/task:1/device:CPU:0"}, %arg9: tensor<*x!tf_type.resource<tensor<25001x64xf32>>> {tf._user_specified_name = "915", tf.device = "/job:tpu_host_worker/replica:0/task:2/device:CPU:0"}, %arg10: tensor<*x!tf_type.resource<tensor<25001x64xf32>>> {tf._user_specified_name = "917", tf.device = "/job:tpu_host_worker/replica:0/task:3/device:CPU:0"}, %arg11: tensor<*x!tf_type.resource<tensor<25001x32xf32>>> {tf._user_specified_name = "919", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg12: tensor<*x!tf_type.resource<tensor<25001x32xf32>>> {tf._user_specified_name = "921", tf.device = "/job:tpu_host_worker/replica:0/task:1/device:CPU:0"}, %arg13: tensor<*x!tf_type.resource<tensor<25001x32xf32>>> {tf._user_specified_name = "923", tf.device = "/job:tpu_host_worker/replica:0/task:2/device:CPU:0"}, %arg14: tensor<*x!tf_type.resource<tensor<25001x32xf32>>> {tf._user_specified_name = "925", tf.device = "/job:tpu_host_worker/replica:0/task:3/device:CPU:0"}, %arg15: tensor<*x!tf_type.resource<tensor<6x32xf32>>> {tf._user_specified_name = "927", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg16: tensor<*x!tf_type.resource<tensor<6x32xf32>>> {tf._user_specified_name = "929", tf.device = "/job:tpu_host_worker/replica:0/task:1/device:CPU:0"}, %arg17: tensor<*x!tf_type.resource<tensor<6x32xf32>>> {tf._user_specified_name = "931", tf.device = "/job:tpu_host_worker/replica:0/task:2/device:CPU:0"}, %arg18: tensor<*x!tf_type.resource<tensor<6x32xf32>>> {tf._user_specified_name = "933", tf.device = "/job:tpu_host_worker/replica:0/task:3/device:CPU:0"}, %arg19: tensor<*x!tf_type.resource<tensor<128x1024xf32>>> {tf._user_specified_name = "935", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg20: tensor<*x!tf_type.resource<tensor<1024xf32>>> {tf._user_specified_name = "937", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg21: tensor<*x!tf_type.resource<tensor<1024x1xf32>>> {tf._user_specified_name = "939", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}) -> tensor<i64> attributes {allow_soft_placement = false, tf.entry_function = {control_outputs = "", inputs = "steps,unknown,unknown_0,unknown_1,unknown_2,unknown_3,unknown_4,unknown_5,unknown_6,unknown_7,unknown_8,unknown_9,unknown_10,unknown_11,unknown_12,unknown_13,unknown_14,unknown_15,unknown_16,unknown_17,unknown_18,unknown_19", outputs = "statefulpartitionedcall_RetVal"}} {
%0 = "tf.ReadVariableOp"(%arg19) : (tensor<*x!tf_type.resource<tensor<128x1024xf32>>>) -> tensor<128x1024xf32>
%1 = "tf.ReadVariableOp"(%arg1) : (tensor<*x!tf_type.resource<tensor<i64>>>) -> tensor<i64>
%2 = "tf.ReadVariableOp"(%arg2) : (tensor<*x!tf_type.resource<tensor<i64>>>) -> tensor<i64>
%3 = "tf.ReadVariableOp"(%arg4) : (tensor<*x!tf_type.resource<tensor<1024xf32>>>) -> tensor<1024xf32>
%4 = "tf.ReadVariableOp"(%arg3) : (tensor<*x!tf_type.resource<tensor<128x1024xf32>>>) -> tensor<128x1024xf32>
%5 = "tf.ReadVariableOp"(%arg5) : (tensor<*x!tf_type.resource<tensor<1024x1xf32>>>) -> tensor<1024x1xf32>
%6 = "tf.ReadVariableOp"(%arg20) : (tensor<*x!tf_type.resource<tensor<1024xf32>>>) -> tensor<1024xf32>
%7 = "tf.ReadVariableOp"(%arg21) : (tensor<*x!tf_type.resource<tensor<1024x1xf32>>>) -> tensor<1024x1xf32>
%8 = "tf.ReadVariableOp"(%arg6) : (tensor<*x!tf_type.resource<tensor<i64>>>) -> tensor<i64>
%9 = "tf.Const"() <{value = dense<"test"> : tensor<3x!tf_type.string>}> : () -> tensor<3x!tf_type.string>
%cst = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>
%11:4 = "tf.Split"(%cst, %0) {num_split = 4 : i32} : (tensor<i32>, tensor<128x1024xf32>) -> (tensor<32x1024xf32>, tensor<32x1024xf32>, tensor<32x1024xf32>, tensor<32x1024xf32>)
%cst_0 = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>
%12:4 = "tf.Split"(%cst_0, %4) {num_split = 4 : i32} : (tensor<i32>, tensor<128x1024xf32>) -> (tensor<32x1024xf32>, tensor<32x1024xf32>, tensor<32x1024xf32>, tensor<32x1024xf32>)
%cst_1 = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>
%cst_2 = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>
%13:20 = tf_device.replicate {devices = {TPU_REPLICATED_CORE_0 = ["/job:tpu_host_worker/replica:0/task:0/device:TPU:0", "/job:tpu_host_worker/replica:0/task:2/device:TPU:0"], TPU_REPLICATED_CORE_1 = ["/job:tpu_host_worker/replica:0/task:0/device:TPU:1", "/job:tpu_host_worker/replica:0/task:2/device:TPU:1"], TPU_REPLICATED_CORE_2 = ["/job:tpu_host_worker/replica:0/task:1/device:TPU:0", "/job:tpu_host_worker/replica:0/task:3/device:TPU:0"], TPU_REPLICATED_CORE_3 = ["/job:tpu_host_worker/replica:0/task:1/device:TPU:1", "/job:tpu_host_worker/replica:0/task:3/device:TPU:1"], TPU_REPLICATED_HOST_0 = ["/job:tpu_host_worker/replica:0/task:0/device:CPU:0", "/job:tpu_host_worker/replica:0/task:2/device:CPU:0"], TPU_REPLICATED_HOST_1 = ["/job:tpu_host_worker/replica:0/task:0/device:CPU:0", "/job:tpu_host_worker/replica:0/task:2/device:CPU:0"], TPU_REPLICATED_HOST_2 = ["/job:tpu_host_worker/replica:0/task:1/device:CPU:0", "/job:tpu_host_worker/replica:0/task:3/device:CPU:0"], TPU_REPLICATED_HOST_3 = ["/job:tpu_host_worker/replica:0/task:1/device:CPU:0", "/job:tpu_host_worker/replica:0/task:3/device:CPU:0"]}, n = 2 : i32} {
%16:40 = "tf_device.parallel_execute"() ({
%19:10 = "tf_device.launch"() <{device = "TPU_REPLICATED_CORE_0"}> ({
%20:10 = "tf.TPUExecute"(%arg0, %11#0, %1, %2, %3, %12#0, %5, %6, %7, %8, %9) : (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>, tensor<3x!tf_type.string>) -> (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>)
tf_device.return %20#0, %20#1, %20#2, %20#3, %20#4, %20#5, %20#6, %20#7, %20#8, %20#9 : tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>
}) : () -> (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>)
tf_device.return %19#0, %19#1, %19#2, %19#3, %19#4, %19#5, %19#6, %19#7, %19#8, %19#9 : tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>
}, {
%19:10 = "tf_device.launch"() <{device = "TPU_REPLICATED_CORE_1"}> ({
%20:10 = "tf.TPUExecute"(%arg0, %11#1, %1, %2, %3, %12#1, %5, %6, %7, %8, %9) : (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>, tensor<3x!tf_type.string>) -> (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>)
tf_device.return %20#0, %20#1, %20#2, %20#3, %20#4, %20#5, %20#6, %20#7, %20#8, %20#9 : tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>
}) : () -> (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>)
tf_device.return %19#0, %19#1, %19#2, %19#3, %19#4, %19#5, %19#6, %19#7, %19#8, %19#9 : tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>
}, {
%19:10 = "tf_device.launch"() <{device = "TPU_REPLICATED_CORE_2"}> ({
%20:10 = "tf.TPUExecute"(%arg0, %11#2, %1, %2, %3, %12#2, %5, %6, %7, %8, %9) : (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>, tensor<3x!tf_type.string>) -> (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>)
tf_device.return %20#0, %20#1, %20#2, %20#3, %20#4, %20#5, %20#6, %20#7, %20#8, %20#9 : tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>
}) : () -> (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>)
tf_device.return %19#0, %19#1, %19#2, %19#3, %19#4, %19#5, %19#6, %19#7, %19#8, %19#9 : tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>
}, {
%19:10 = "tf_device.launch"() <{device = "TPU_REPLICATED_CORE_3"}> ({
%20:10 = "tf.TPUExecute"(%arg0, %11#3, %1, %2, %3, %12#3, %5, %6, %7, %8, %9) : (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>, tensor<3x!tf_type.string>) -> (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>)
tf_device.return %20#0, %20#1, %20#2, %20#3, %20#4, %20#5, %20#6, %20#7, %20#8, %20#9 : tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>
}) : () -> (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>)
tf_device.return %19#0, %19#1, %19#2, %19#3, %19#4, %19#5, %19#6, %19#7, %19#8, %19#9 : tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>
}) : () -> (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>, tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>, tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>, tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>)
%17 = "tf.Concat"(%cst_1, %16#5, %16#15, %16#25, %16#35) : (tensor<i32>, tensor<32x1024xf32>, tensor<32x1024xf32>, tensor<32x1024xf32>, tensor<32x1024xf32>) -> tensor<128x1024xf32>
%18 = "tf.Concat"(%cst_2, %16#1, %16#11, %16#21, %16#31) : (tensor<i32>, tensor<32x1024xf32>, tensor<32x1024xf32>, tensor<32x1024xf32>, tensor<32x1024xf32>) -> tensor<128x1024xf32>
tf_device.return %16#0, %16#9, %16#8, %16#7, %16#6, %17, %16#4, %16#3, %16#2, %18 : tensor<i32>, tensor<i64>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<128x1024xf32>, tensor<1024xf32>, tensor<i64>, tensor<i64>, tensor<128x1024xf32>
}
"tf.AssignVariableOp"(%arg19, %13#18) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<128x1024xf32>>>, tensor<128x1024xf32>) -> ()
"tf.AssignVariableOp"(%arg1, %13#16) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<i64>>>, tensor<i64>) -> ()
"tf.AssignVariableOp"(%arg2, %13#14) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<i64>>>, tensor<i64>) -> ()
"tf.AssignVariableOp"(%arg4, %13#12) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<1024xf32>>>, tensor<1024xf32>) -> ()
"tf.AssignVariableOp"(%arg3, %13#10) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<128x1024xf32>>>, tensor<128x1024xf32>) -> ()
"tf.AssignVariableOp"(%arg5, %13#8) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<1024x1xf32>>>, tensor<1024x1xf32>) -> ()
"tf.AssignVariableOp"(%arg20, %13#6) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<1024xf32>>>, tensor<1024xf32>) -> ()
"tf.AssignVariableOp"(%arg21, %13#4) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<1024x1xf32>>>, tensor<1024x1xf32>) -> ()
"tf.AssignVariableOp"(%arg6, %13#2) <{validate_shape = true}> {_has_manual_control_dependencies = true} : (tensor<*x!tf_type.resource<tensor<i64>>>, tensor<i64>) -> ()
%14 = "tf.ReadVariableOp"(%arg2) {device = ""} : (tensor<*x!tf_type.resource<tensor<i64>>>) -> tensor<i64>
%15 = "tf.Identity"(%14) {device = ""} : (tensor<i64>) -> tensor<i64>
return %15 : tensor<i64>
}
}
tensorflow/compiler/mlir/tf2xla/api/v2/testdata/graph_with_flib_def.txt
Vendored
+1535
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/mlir/tf2xla/api/v2/testdata/invalid_executor.mlir
Vendored
+32
View File
@@ -0,0 +1,32 @@
// 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.
// ==============================================================================
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() {
tf_executor.graph {
%control = tf_executor.island {
tf_executor.yield
}
tf_executor.fetch %control : !tf_executor.control
}
tf_executor.graph {
%control = tf_executor.island {
tf_executor.yield
}
tf_executor.fetch %control : !tf_executor.control
}
return
}
}
tensorflow/compiler/mlir/tf2xla/api/v2/testdata/outside_compilation.mlir
Vendored
+150
View File
@@ -0,0 +1,150 @@
// 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.
// ==============================================================================
module attributes {tf.devices = {"/job:localhost/replica:0/task:0/device:CPU:0", "/job:localhost/replica:0/task:0/device:TPU:0", "/job:localhost/replica:0/task:0/device:TPU:1", "/job:localhost/replica:0/task:0/device:TPU_SYSTEM:0"}, tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1654 : i32}} {
func.func @main(%arg0: tensor<*x!tf_type.resource> {tf._user_specified_name = "input_1", tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}) attributes {allow_soft_placement = true, tf.entry_function = {control_outputs = "while,image_sample/write_summary/summary_cond", inputs = "image_sample_write_summary_summary_cond_input_1", outputs = ""}} {
tf_executor.graph {
%outputs, %control = tf_executor.island wraps "tf.Const"() {device = "", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%outputs_0, %control_1 = tf_executor.island wraps "tf.Const"() {device = "", value = dense<1> : tensor<i32>} : () -> tensor<i32>
%outputs_2, %control_3 = tf_executor.island wraps "tf.TPUReplicatedInput"(%outputs, %outputs_0) {device = "", index = -1 : i64, is_mirrored_variable = false, is_packed = false} : (tensor<i32>, tensor<i32>) -> tensor<*xi32>
%outputs_4, %control_5 = tf_executor.island wraps "tf.Const"() {device = "", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%outputs_6, %control_7 = tf_executor.island wraps "tf.Const"() {_post_device_rewrite = true, device = "", value = dense<> : tensor<0xi32>} : () -> tensor<0xi32>
%outputs_8, %control_9 = tf_executor.island wraps "tf.Const"() {_post_device_rewrite = true, device = "", value = dense<> : tensor<0xi32>} : () -> tensor<0xi32>
%outputs_10, %control_11 = tf_executor.island wraps "tf.Pack"(%outputs_6, %outputs_8) {axis = 0 : i64, device = ""} : (tensor<0xi32>, tensor<0xi32>) -> tensor<*xi32>
%outputs_12, %control_13 = tf_executor.island wraps "tf.Max"(%outputs_10, %outputs_4) {device = "", keep_dims = false} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
%control_14 = tf_executor.island wraps "tf.NoOp"() {_pivot_for_cluster = "cluster_sample_sequence", device = ""} : () -> ()
%control_15 = tf_executor.island(%control_14) wraps "tf.NoOp"() {_has_manual_control_dependencies = true, _tpu_replicate = "cluster_sample_sequence", device = ""} : () -> ()
%control_16 = tf_executor.island(%control_15) wraps "tf.NoOp"() {device = ""} : () -> ()
%control_17 = tf_executor.island(%control_15) wraps "tf.NoOp"() {device = ""} : () -> ()
%control_18 = tf_executor.island(%control_14) wraps "tf.TPUReplicateMetadata"() {_has_manual_control_dependencies = true, _tpu_replicate = "cluster_sample_sequence", allow_soft_placement = true, computation_shape = [], device = "", device_assignment = [], host_compute_core = [], num_cores_per_replica = 1 : i64, num_replicas = 2 : i64, padding_map = [], step_marker_location = "STEP_MARK_AT_ENTRY", topology = "", tpu_compile_options_proto = "", use_spmd_for_xla_partitioning = false, use_tpu = true} : () -> ()
%outputs_19, %control_20 = tf_executor.island(%control_18) wraps "tf.Const"() {_tpu_replicate = "cluster_sample_sequence", device = "", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%outputs_21, %control_22 = tf_executor.island(%control_18) wraps "tf.Const"() {_tpu_replicate = "cluster_sample_sequence", device = "", value = dense<5> : tensor<i64>} : () -> tensor<i64>
%outputs_23, %control_24 = tf_executor.island(%control_18) wraps "tf.Const"() {_tpu_replicate = "cluster_sample_sequence", device = "", value = dense<[3, 32, 32, 1]> : tensor<4xi32>} : () -> tensor<4xi32>
%outputs_25, %control_26 = tf_executor.island(%control_18) wraps "tf.TPUCompilationResult"() {_tpu_compilation_status = "cluster_sample_sequence", device = ""} : () -> tensor<!tf_type.string>
%outputs_27, %control_28 = tf_executor.island(%control_18) wraps "tf.Const"() {_tpu_replicate = "cluster_sample_sequence", device = "", value = dense<3> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_29, %control_30 = tf_executor.island(%control_18) wraps "tf.Const"() {_tpu_replicate = "cluster_sample_sequence", device = "", value = dense<[1, 1, 1, 3]> : tensor<4xi32>} : () -> tensor<4xi32>
%outputs_31, %control_32 = tf_executor.island(%control_18) wraps "tf.Const"() {_tpu_replicate = "cluster_sample_sequence", device = "", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%outputs_33, %control_34 = tf_executor.island(%control_18) wraps "tf.Const"() {_tpu_replicate = "cluster_sample_sequence", device = "", value = dense<true> : tensor<i1>} : () -> tensor<i1>
%outputs_35, %control_36 = tf_executor.island(%control_18) wraps "tf.Identity"(%outputs_2) {_tpu_input_identity = true, _tpu_replicate = "cluster_sample_sequence", device = ""} : (tensor<*xi32>) -> tensor<*xi32>
%outputs_37, %control_38 = tf_executor.island wraps "tf.Equal"(%outputs_35, %outputs_31) {_tpu_replicate = "cluster_sample_sequence", device = "", incompatible_shape_error = true} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi1>
%outputs_39, %control_40 = tf_executor.island wraps "tf.LogicalAnd"(%outputs_37, %outputs_33) {_tpu_replicate = "cluster_sample_sequence", device = ""} : (tensor<*xi1>, tensor<i1>) -> tensor<*xi1>
%outputs_41, %control_42 = tf_executor.island(%control_18) wraps "tf.Const"() {_tpu_replicate = "cluster_sample_sequence", device = "", value = dense<3> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_43, %control_44 = tf_executor.island(%control_18) wraps "tf.Const"() {_tpu_replicate = "cluster_sample_sequence", device = "", value = dense<1024> : tensor<i32>} : () -> tensor<i32>
%outputs_45, %control_46 = tf_executor.island wraps "tf.TensorListReserve"(%outputs_41, %outputs_43) {_tpu_replicate = "cluster_sample_sequence", device = ""} : (tensor<1xi32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
%outputs_47, %control_48 = tf_executor.island(%control_18) wraps "tf.Const"() {_tpu_replicate = "cluster_sample_sequence", device = "", value = dense<[1, 0]> : tensor<2xi32>} : () -> tensor<2xi32>
%outputs_49, %control_50 = tf_executor.island(%control_18) wraps "tf.Const"() {_tpu_replicate = "cluster_sample_sequence", device = "", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%outputs_51, %control_52 = tf_executor.island(%control_18) wraps "tf.Const"() {_tpu_replicate = "cluster_sample_sequence", device = "", value = dense<-1> : tensor<i32>} : () -> tensor<i32>
%outputs_53:4, %control_54 = tf_executor.island wraps "tf.While"(%outputs_49, %outputs_51, %outputs_19, %outputs_45) {_num_original_outputs = 4 : i64, _read_only_resource_inputs = [], _tpu_replicate = "cluster_sample_sequence", _xla_propagate_compile_time_consts = true, body = @while_body_260, cond = @while_cond_250, device = "", is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>) -> (tensor<i32>, tensor<i32>, tensor<i32>, tensor<!tf_type.variant>)
%outputs_55, %control_56 = tf_executor.island wraps "tf.TensorListStack"(%outputs_53#3, %outputs_27) {_tpu_replicate = "cluster_sample_sequence", device = "", num_elements = 1024 : i64} : (tensor<!tf_type.variant>, tensor<1xi32>) -> tensor<*xf32>
%outputs_57, %control_58 = tf_executor.island wraps "tf.Transpose"(%outputs_55, %outputs_47) {_tpu_replicate = "cluster_sample_sequence", device = ""} : (tensor<*xf32>, tensor<2xi32>) -> tensor<*xf32>
%outputs_59, %control_60 = tf_executor.island wraps "tf.Reshape"(%outputs_57, %outputs_23) {_tpu_replicate = "cluster_sample_sequence", device = ""} : (tensor<*xf32>, tensor<4xi32>) -> tensor<*xf32>
%outputs_61, %control_62 = tf_executor.island wraps "tf.Tile"(%outputs_59, %outputs_29) {_tpu_replicate = "cluster_sample_sequence", device = ""} : (tensor<*xf32>, tensor<4xi32>) -> tensor<*xf32>
%outputs_63, %control_64 = tf_executor.island wraps "tf.If"(%outputs_39, %outputs_61, %arg0, %outputs_21) {_read_only_resource_inputs = [], _tpu_replicate = "cluster_sample_sequence", _xla_propagate_compile_time_consts = true, device = "", else_branch = @image_sample_write_summary_summary_cond_false_710, is_stateless = false, then_branch = @image_sample_write_summary_summary_cond_true_700} : (tensor<*xi1>, tensor<*xf32>, tensor<*x!tf_type.resource>, tensor<i64>) -> tensor<*xi1>
%outputs_65, %control_66 = tf_executor.island wraps "tf.Identity"(%outputs_63) {_tpu_replicate = "cluster_sample_sequence", device = ""} : (tensor<*xi1>) -> tensor<*xi1>
tf_executor.fetch %control_54, %control_64 : !tf_executor.control, !tf_executor.control
}
return
}
func.func private @while_body_260(%arg0: tensor<i32> {tf._user_specified_name = "while/loop_counter"}, %arg1: tensor<i32> {tf._user_specified_name = "while/maximum_iterations"}, %arg2: tensor<i32>, %arg3: tensor<!tf_type.variant>) -> (tensor<*xi32>, tensor<*xi32>, tensor<*xi32>, tensor<*x!tf_type.variant>) attributes {tf._construction_context = "kEagerRuntime", tf.signature.is_stateful} {
%0:4 = tf_executor.graph {
%outputs, %control = tf_executor.island wraps "tf.Const"() {device = "", value = dense<1> : tensor<i32>} : () -> tensor<i32>
%outputs_0, %control_1 = tf_executor.island wraps "tf.Const"() {device = "", value = dense<1> : tensor<i32>} : () -> tensor<i32>
%outputs_2, %control_3 = tf_executor.island wraps "tf.Const"() {device = "", value = dense<3> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_4, %control_5 = tf_executor.island wraps "tf.RandomUniform"(%outputs_2) {device = "", seed = 87654321 : i64, seed2 = 0 : i64} : (tensor<1xi32>) -> tensor<*xf32>
%outputs_6, %control_7 = tf_executor.island wraps "tf.AddV2"(%arg2, %outputs) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<*xi32>
%outputs_8, %control_9 = tf_executor.island wraps "tf.Identity"(%outputs_6) {device = ""} : (tensor<*xi32>) -> tensor<*xi32>
%outputs_10, %control_11 = tf_executor.island wraps "tf.TensorListSetItem"(%arg3, %arg2, %outputs_4) {device = "", resize_if_index_out_of_bounds = false} : (tensor<!tf_type.variant>, tensor<i32>, tensor<*xf32>) -> tensor<*x!tf_type.variant>
%outputs_12, %control_13 = tf_executor.island wraps "tf.Identity"(%outputs_10) {device = ""} : (tensor<*x!tf_type.variant>) -> tensor<*x!tf_type.variant>
%outputs_14, %control_15 = tf_executor.island wraps "tf.AddV2"(%arg0, %outputs_0) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<*xi32>
%outputs_16, %control_17 = tf_executor.island wraps "tf.Identity"(%outputs_14) {device = ""} : (tensor<*xi32>) -> tensor<*xi32>
%outputs_18, %control_19 = tf_executor.island wraps "tf.Identity"(%arg1) {device = ""} : (tensor<i32>) -> tensor<*xi32>
tf_executor.fetch %outputs_16, %outputs_18, %outputs_8, %outputs_12 : tensor<*xi32>, tensor<*xi32>, tensor<*xi32>, tensor<*x!tf_type.variant>
}
return %0#0, %0#1, %0#2, %0#3 : tensor<*xi32>, tensor<*xi32>, tensor<*xi32>, tensor<*x!tf_type.variant>
}
func.func private @while_cond_250(%arg0: tensor<i32> {tf._user_specified_name = "while/loop_counter"}, %arg1: tensor<i32> {tf._user_specified_name = "while/maximum_iterations"}, %arg2: tensor<i32>, %arg3: tensor<!tf_type.variant>) -> tensor<*xi1> attributes {tf._construction_context = "kEagerRuntime"} {
%0 = tf_executor.graph {
%outputs, %control = tf_executor.island wraps "tf.Const"() {device = "", value = dense<1024> : tensor<i32>} : () -> tensor<i32>
%outputs_0, %control_1 = tf_executor.island wraps "tf.Less"(%arg2, %outputs) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<*xi1>
%outputs_2, %control_3 = tf_executor.island wraps "tf.Identity"(%outputs_0) {device = ""} : (tensor<*xi1>) -> tensor<*xi1>
tf_executor.fetch %outputs_2 : tensor<*xi1>
}
return %0 : tensor<*xi1>
}
func.func private @image_sample_write_summary_summary_cond_false_710(%arg0: tensor<3x32x32x3xf32>, %arg1: tensor<*x!tf_type.resource>, %arg2: tensor<i64>) -> tensor<*xi1> attributes {tf._construction_context = "kEagerRuntime"} {
%0 = tf_executor.graph {
%outputs, %control = tf_executor.island wraps "tf.Const"() {device = "", value = dense<false> : tensor<i1>} : () -> tensor<i1>
%outputs_0, %control_1 = tf_executor.island wraps "tf.Identity"(%outputs) {device = ""} : (tensor<i1>) -> tensor<*xi1>
tf_executor.fetch %outputs_0 : tensor<*xi1>
}
return %0 : tensor<*xi1>
}
func.func private @image_sample_write_summary_summary_cond_true_700(%arg0: tensor<3x32x32x3xf32> {tf._user_specified_name = "Tile"}, %arg1: tensor<*x!tf_type.resource> {tf._user_specified_name = "writer"}, %arg2: tensor<i64> {tf._user_specified_name = "Const_3"}) -> tensor<*xi1> attributes {tf._construction_context = "kEagerRuntime", tf.signature.is_stateful} {
%0 = tf_executor.graph {
%outputs, %control = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<[3, 32, 32, 3]> : tensor<4xi32>} : () -> tensor<4xi32>
%outputs_0, %control_1 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%outputs_2, %control_3 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<""> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
%outputs_4, %control_5 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<"Condition x >= 0 did not hold element-wise:"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
%outputs_6, %control_7 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<"x (image_sample/write_summary/summary_cond/assert_non_negative/x:0) = "> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
%outputs_8, %control_9 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<""> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
%outputs_10, %control_11 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<"Condition x >= 0 did not hold element-wise:"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
%outputs_12, %control_13 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<"x (image_sample/write_summary/summary_cond/assert_non_negative/x:0) = "> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
%outputs_14, %control_15 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%outputs_16, %control_17 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
%outputs_18, %control_19 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%outputs_20, %control_21 = tf_executor.island wraps "tf.Range"(%outputs_18, %outputs_14, %outputs_16) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<*xi32>
%outputs_22, %control_23 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<3> : tensor<i32>} : () -> tensor<i32>
%outputs_24, %control_25 = tf_executor.island wraps "tf.LessEqual"(%outputs_0, %outputs_22) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<*xi1>
%outputs_26, %control_27 = tf_executor.island wraps "tf.All"(%outputs_24, %outputs_20) {device = "/device:CPU:0", keep_dims = false} : (tensor<*xi1>, tensor<*xi32>) -> tensor<*xi1>
%control_28 = tf_executor.island wraps "tf.Assert"(%outputs_26, %outputs_2, %outputs_4, %outputs_6, %outputs_22) {device = "/device:CPU:0", summarize = 3 : i64} : (tensor<*xi1>, tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<i32>) -> ()
%outputs_29, %control_30 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<[3, 32, 32, 3]> : tensor<4xi32>} : () -> tensor<4xi32>
%control_31 = tf_executor.island wraps "tf.NoOp"() {device = "/device:CPU:0"} : () -> ()
%outputs_32, %control_33 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<4> : tensor<i32>} : () -> tensor<i32>
%control_34 = tf_executor.island wraps "tf.NoOp"() {device = "/device:CPU:0"} : () -> ()
%outputs_35, %control_36 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%outputs_37, %control_38 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<0.000000e+00> : tensor<f32>} : () -> tensor<f32>
%outputs_39, %control_40 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<2.550000e+02> : tensor<f32>} : () -> tensor<f32>
%outputs_41, %control_42 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<2.555000e+02> : tensor<f32>} : () -> tensor<f32>
%outputs_43, %control_44 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_45, %control_46 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<3> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_47, %control_48 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_49, %control_50 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_51, %control_52 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<3> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_53, %control_54 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_55, %control_56 = tf_executor.island wraps "tf.StridedSlice"(%outputs, %outputs_49, %outputs_51, %outputs_53) {begin_mask = 0 : i64, device = "/device:CPU:0", ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<*xi32>
%outputs_57, %control_58 = tf_executor.island wraps "tf.AsString"(%outputs_55) {device = "/device:CPU:0", fill = "", precision = -1 : i64, scientific = false, shortest = false, width = -1 : i64} : (tensor<*xi32>) -> tensor<*x!tf_type.string>
%outputs_59, %control_60 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_61, %control_62 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_63, %control_64 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_65, %control_66 = tf_executor.island wraps "tf.StridedSlice"(%outputs, %outputs_59, %outputs_61, %outputs_63) {begin_mask = 0 : i64, device = "/device:CPU:0", ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<*xi32>
%outputs_67, %control_68 = tf_executor.island wraps "tf.AsString"(%outputs_65) {device = "/device:CPU:0", fill = "", precision = -1 : i64, scientific = false, shortest = false, width = -1 : i64} : (tensor<*xi32>) -> tensor<*x!tf_type.string>
%outputs_69, %control_70 = tf_executor.island wraps "tf.Pack"(%outputs_57, %outputs_67) {axis = 0 : i64, device = "/device:CPU:0"} : (tensor<*x!tf_type.string>, tensor<*x!tf_type.string>) -> tensor<*x!tf_type.string>
%outputs_71, %control_72 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<"\0A\08\0A\06images"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
%outputs_73, %control_74 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<"image_sample"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
%outputs_75, %control_76 = tf_executor.island wraps "tf.Mul"(%arg0, %outputs_41) {device = "/device:CPU:0"} : (tensor<3x32x32x3xf32>, tensor<f32>) -> tensor<*xf32>
%outputs_77, %control_78 = tf_executor.island wraps "tf.ClipByValue"(%outputs_75, %outputs_37, %outputs_39) {device = "/device:CPU:0"} : (tensor<*xf32>, tensor<f32>, tensor<f32>) -> tensor<*xf32>
%outputs_79, %control_80 = tf_executor.island wraps "tf.Cast"(%outputs_77) {Truncate = false, device = "/device:CPU:0"} : (tensor<*xf32>) -> tensor<*xui8>
%outputs_81, %control_82 = tf_executor.island wraps "tf.StridedSlice"(%outputs_79, %outputs_43, %outputs_45, %outputs_47) {begin_mask = 1 : i64, device = "/device:CPU:0", ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64} : (tensor<*xui8>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<*xui8>
%outputs_83, %control_84 = tf_executor.island wraps "tf.EncodePng"(%outputs_81) {compression = -1 : i64, device = "/device:CPU:0"} : (tensor<*xui8>) -> tensor<*x!tf_type.string>
%outputs_85, %control_86 = tf_executor.island wraps "tf.ConcatV2"(%outputs_69, %outputs_83, %outputs_35) {device = "/device:CPU:0"} : (tensor<*x!tf_type.string>, tensor<*x!tf_type.string>, tensor<i32>) -> tensor<*x!tf_type.string>
%control_87 = tf_executor.island wraps "tf.WriteSummary"(%arg1, %arg2, %outputs_85, %outputs_73, %outputs_71) {_has_manual_control_dependencies = true, device = "/device:CPU:0"} : (tensor<*x!tf_type.resource>, tensor<i64>, tensor<*x!tf_type.string>, tensor<!tf_type.string>, tensor<!tf_type.string>) -> ()
%outputs_88, %control_89 = tf_executor.island(%control_87) wraps "tf.Const"() {device = "/device:CPU:0", value = dense<true> : tensor<i1>} : () -> tensor<i1>
%control_90 = tf_executor.island(%control_28, %control_87) wraps "tf.NoOp"() {device = ""} : () -> ()
%outputs_91, %control_92 = tf_executor.island(%control_90) wraps "tf.Identity"(%outputs_88) {device = ""} : (tensor<i1>) -> tensor<*xi1>
tf_executor.fetch %outputs_91, %control_28, %control_87 : tensor<*xi1>, !tf_executor.control, !tf_executor.control
}
return %0 : tensor<*xi1>
}
}
tensorflow/compiler/mlir/tf2xla/api/v2/testdata/valid_executor.mlir
Vendored
+24
View File
@@ -0,0 +1,24 @@
// 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.
// ==============================================================================
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() {
tf_executor.graph {
%0:2 = tf_executor.island wraps "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:TPU:0", dtype = "tfdtype$DT_INT32", value = #tf_type<tensor_proto : "0x746674656E736F722464747970653A2044545F494E5433320A74656E736F725F7368617065207B0A202064696D207B0A2020202073697A653A20320A20207D0A7D0A74656E736F725F636F6E74656E743A20225C3230305C3030305C3030305C3030305C3230305C3030305C3030305C303030220A"> : tensor<2xi32>} : () -> tensor<2xi32> loc("Empty/shape")
tf_executor.fetch
}
func.return
}
}
tensorflow/compiler/mlir/tf2xla/api/v2/testdata/valid_graph.txt
Vendored
+44
View File
@@ -0,0 +1,44 @@
node {
name: "Empty/shape"
op: "Const"
device: "/job:localhost/replica:0/task:0/device:TPU:0"
attr {
key: "_output_shapes"
value {
list {
shape {
dim {
size: 2
}
}
}
}
}
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
dim {
size: 2
}
}
tensor_content: "\200\000\000\000\200\000\000\000"
}
}
}
experimental_debug_info {
}
}
library {
}
versions {
producer: 268
}
tensorflow/compiler/mlir/tf2xla/api/v2/testing/BUILD
+53
View File
@@ -0,0 +1,53 @@
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/api/v2:__subpackages__",
],
)
cc_library(
name = "compile_mlir",
testonly = True,
srcs = ["compile_mlir.cc"],
hdrs = ["compile_mlir.h"],
deps = [
"//tensorflow/compiler/jit",
"//tensorflow/compiler/mlir/tf2xla/api/v2:legalize_tf",
"//tensorflow/compiler/mlir/tf2xla/internal/utils:test_metadata_config",
"//tensorflow/compiler/tf2xla:xla_compiler",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:test",
"//tensorflow/core/protobuf:for_core_protos_cc",
"//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
"//tensorflow/core/tpu/kernels:tpu_compile_op_support",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:string_view",
"@llvm-project//mlir:Pass",
"@xla//xla:shape_util",
"@xla//xla/client:client_library",
"@xla//xla/stream_executor:platform",
"@xla//xla/stream_executor:platform_manager",
],
)
cc_library(
name = "utils",
testonly = True,
srcs = ["utils.cc"],
hdrs = ["utils.h"],
deps = ["//tensorflow/core/platform:resource_loader"],
)
tf_cc_test(
name = "utils_test",
srcs = ["utils_test.cc"],
deps = [
":utils",
"@com_google_googletest//:gtest_main",
],
)
tensorflow/compiler/mlir/tf2xla/api/v2/testing/compile_mlir.cc
+81
View File
@@ -0,0 +1,81 @@
/* 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/api/v2/testing/compile_mlir.h"
#include <memory>
#include <vector>
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/utils/test_metadata_config.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/client/client_library.h"
#include "xla/shape.h"
#include "xla/stream_executor/platform.h"
#include "xla/stream_executor/platform_manager.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
namespace testing {
using tpu::MlirToHloArgs;
using tpu::ShardingAndIndex;
using tpu::TPUCompileMetadataProto;
absl::StatusOr<XlaCompiler::CompilationResult> CompileMlirModule(
const char* mlir_module_str,
ConfigProto::Experimental::MlirBridgeRollout rollout_state,
absl::string_view device_type) {
MlirToHloArgs mlir_to_hlo_args;
mlir_to_hlo_args.rollout_state = rollout_state;
mlir_to_hlo_args.mlir_module = mlir_module_str;
TF_ASSIGN_OR_RETURN(se::Platform * platform,
se::PlatformManager::PlatformWithName("Host"));
TF_ASSIGN_OR_RETURN(
auto client, xla::ClientLibrary::GetOrCreateCompileOnlyClient(platform));
std::vector<TensorShape> arg_shapes;
TPUCompileMetadataProto metadata_proto;
// Configure metadata requires parsing the module and if we are testing a
// failure, we ignore this particular set up error assuming we'll not get
// far enough to need valid metadata.
tensorflow::tf2xla::internal::ConfigureMetadata(mlir_module_str, arg_shapes,
metadata_proto)
.IgnoreError();
bool use_tuple_args = true;
std::vector<ShardingAndIndex> arg_core_mapping;
std::vector<std::vector<xla::Shape>> per_core_arg_shapes;
std::vector<std::unique_ptr<mlir::Pass>> custom_legalization_passes;
return LegalizeMlirToHlo(mlir_to_hlo_args, metadata_proto, use_tuple_args,
device_type, custom_legalization_passes,
/*shape_determination_fns=*/{}, arg_shapes,
&arg_core_mapping, &per_core_arg_shapes, client);
}
} // namespace testing
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v2/testing/compile_mlir.h
+40
View File
@@ -0,0 +1,40 @@
/* 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_API_V2_TESTING_COMPILE_MLIR_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_TESTING_COMPILE_MLIR_H_
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
namespace testing {
// Compiles the given MLIR module to XLA HLO.
absl::StatusOr<XlaCompiler::CompilationResult> CompileMlirModule(
const char* mlir_module_str,
ConfigProto::Experimental::MlirBridgeRollout rollout_state,
absl::string_view device_type = "XLA_TPU_JIT");
} // namespace testing
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_TESTING_COMPILE_MLIR_H_
tensorflow/compiler/mlir/tf2xla/api/v2/testing/utils.cc
+35
View File
@@ -0,0 +1,35 @@
/* 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/api/v2/testing/utils.h"
#include <string>
#include "tensorflow/core/platform/resource_loader.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
namespace testing {
std::string TestDataPath() {
return tensorflow::GetDataDependencyFilepath(
"tensorflow/compiler/mlir/tf2xla/api/v2/testdata/");
}
} // namespace testing
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v2/testing/utils.h
+34
View File
@@ -0,0 +1,34 @@
/* 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_API_V2_TESTING_UTILS_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_TESTING_UTILS_H_
#include <string>
namespace tensorflow {
namespace tf2xla {
namespace v2 {
namespace testing {
// Returns the path to the testdata directory.
std::string TestDataPath();
} // namespace testing
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_TESTING_UTILS_H_
tensorflow/compiler/mlir/tf2xla/api/v2/testing/utils_test.cc
+47
View File
@@ -0,0 +1,47 @@
/* 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/api/v2/testing/utils.h"
#include <stdlib.h>
#include <string>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
namespace tensorflow {
namespace tf2xla {
namespace v2 {
namespace testing {
namespace {
class UtilsTest : public ::testing::Test {};
TEST_F(UtilsTest, TestDataPathSucceeds) {
std::string expected_test_data_path_regex =
".*tensorflow/compiler/mlir/tf2xla/api/v2/testdata/";
std::string result_test_data_path = TestDataPath();
EXPECT_THAT(result_test_data_path,
::testing::ContainsRegex(expected_test_data_path_regex));
}
} // namespace
} // namespace testing
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor.cc
+184
View File
@@ -0,0 +1,184 @@
/* 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/api/v2/tf_dialect_to_executor.h"
#include <memory>
#include <string>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "llvm/ADT/StringRef.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/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Transforms/Passes.h" // from @llvm-project
#include "tensorflow/compiler/jit/flags.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/data_dumper_logger_config.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/logging_hooks.h"
#include "xla/tsl/lib/monitoring/counter.h"
#include "xla/tsl/platform/status.h"
#include "tensorflow/core/platform/error_payloads.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/util/debug_data_dumper.h"
#include "tsl/platform/error_logging.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
using mlir::LogicalResult;
using mlir::ModuleOp;
using mlir::OpPassManager;
using mlir::PassManager;
using mlir::func::FuncOp;
auto *tf_dialect_to_executor_dialect_status = tsl::monitoring::Counter<1>::New(
"/tensorflow/core/tf2xla/api/v2/tf_dialect_to_executor_dialect_status",
"Counts how often a successful export from TF Dialect to Executor Dialect "
"is",
"status");
constexpr char kExportSuccess[] = "success";
constexpr char kExportFailed[] = "failed";
namespace {
void AddTfDialectToExecutorPasses(OpPassManager &pm) {
pm.addPass(mlir::TF::CreateTFRegionControlFlowToFunctional());
// First, we need to convert from functional, to executor dialect.
pm.addNestedPass<FuncOp>(
mlir::CreateFunctionalToExecutorDialectConversionPass());
// Do a single pass to split the graph's single island op into an island per
// op as expected by the following passes.
pm.addNestedPass<FuncOp>(mlir::TF::CreateSplitIntoIslandPerOpPass());
pm.addNestedPass<FuncOp>(mlir::TFDevice::CreateReplicateToIslandPass(
/*legacy_graph_export=*/false));
pm.addNestedPass<FuncOp>(
mlir::TFDevice::CreateReplicaIDToDeviceOrdinalPass());
pm.addNestedPass<FuncOp>(mlir::TFDevice::CreateParallelExecuteToIslandsPass(
/*legacy_graph_export=*/false));
pm.addNestedPass<FuncOp>(mlir::TFDevice::CreateLaunchToDeviceAttributePass(
/*legacy_graph_export=*/false));
// Do a single pass to encode necessary control deps in the IR according to
// the results of side effect analysis.
pm.addPass(
mlir::tf_executor::CreateTFExecutorUpdateControlDependenciesPass());
pm.addNestedPass<FuncOp>(mlir::TFTPU::CreateTPUDevicePropagationPass());
pm.addNestedPass<FuncOp>(mlir::TFTPU::CreateTPUColocateSplitsPass());
pm.addPass(mlir::createSymbolDCEPass());
pm.addNestedPass<FuncOp>(
mlir::tf_executor::CreateTFExecutorGraphPruningPass());
if (tensorflow::GetMlirCommonFlags()
->tf_mlir_enable_convert_control_to_data_outputs_pass) {
bool composite_tpuexecute_side_effects =
tensorflow::GetMlirCommonFlags()
->tf_mlir_enable_composite_tpuexecute_side_effects;
pm.addPass(
mlir::tf_executor::CreateTFExecutorConvertControlToDataOutputsPass(
composite_tpuexecute_side_effects));
}
pm.addPass(mlir::TF::CreateVerifySuitableForExportPass());
}
absl::Status RecordStatusIfError(absl::Status status) {
if (status.ok()) {
return absl::OkStatus();
}
tf_dialect_to_executor_dialect_status->GetCell(kExportFailed)->IncrementBy(1);
VLOG(1) << "Failed to export from TF Dialect to TF Executor Dialect. "
<< status;
constexpr char bridge_subcomponent[] =
"TFXLA_TF_FUNCTIONAL_TO_EXECUTOR_EXPORT_v2";
constexpr char kBridgeComponent[] = "TFXLABridge";
tsl::OkOrSetErrorCounterPayload(
tensorflow::core::platform::ErrorSourceProto::MLIR_BRIDGE_PHASE_1,
status);
tsl::error_logging::Log(kBridgeComponent, bridge_subcomponent,
status.ToString())
.IgnoreError();
return status;
}
} // namespace
absl::Status ExportFromTensorflowDialectToExecutor(
ModuleOp module, llvm::StringRef module_name) {
PassManager tf_to_executor(module.getContext());
::tensorflow::applyTensorflowAndCLOptions(tf_to_executor);
tf_to_executor.enableVerifier();
AddTfDialectToExecutorPasses(tf_to_executor);
if (VLOG_IS_ON(1) ||
DEBUG_DATA_DUMPER()->ShouldDump(module_name.str(), kDebugGroupMain)) {
::tensorflow::DumpMlirOpToFile(
DEBUG_DATA_DUMPER()->GetDumpFilename(
module_name.str(), kDebugGroupMain,
"tfxla_bridge_v2_tfdialect_to_executor_before"),
module, llvm::StringRef(), &tf_to_executor);
if (VLOG_IS_ON(2) ||
DEBUG_DATA_DUMPER()->ShouldDump(
module_name.str(), kDebugGroupBridgePhase1ExecutorExport)) {
internal::EnablePassIRPrinting(
tf_to_executor, kDebugGroupBridgePhase1ExecutorExport, module_name);
}
}
LogicalResult result = tf_to_executor.run(module);
if (VLOG_IS_ON(1) ||
DEBUG_DATA_DUMPER()->ShouldDump(module_name.str(), kDebugGroupMain)) {
::tensorflow::DumpMlirOpToFile(
DEBUG_DATA_DUMPER()->GetDumpFilename(
module_name.str(), kDebugGroupMain,
"tfxla_bridge_v2_tfdialect_to_executor_after"),
module, llvm::StringRef(), &tf_to_executor);
}
if (result.failed()) {
return RecordStatusIfError(
absl::InternalError("Failed to export from TF Dialect to TF Executor "
"Dialect. Read LLVM Pipeline Error"));
}
tf_dialect_to_executor_dialect_status->GetCell(kExportSuccess)
->IncrementBy(1);
return absl::OkStatus();
}
mlir::PassPipelineRegistration<> tf_dialect_to_executor_pipeline(
"tf-dialect-to-executor-v2",
"Run passes to convert from TF Dialect to Executor in preparation for "
"exporting module back to TF Graph.",
AddTfDialectToExecutorPasses);
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor.h
+51
View File
@@ -0,0 +1,51 @@
/* 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_API_V2_TF_DIALECT_TO_EXECUTOR_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_TF_DIALECT_TO_EXECUTOR_H_
#include "llvm/ADT/StringRef.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
// Given the input Module op that's in the Tensorflow Dialect, convert the MLIR
// module in place to the Tensorflow Executor Dialect. Returns an OK Status if
// success, otherwise failure with an error message.
// The Tensorflow Executor Dialect is required to export an MLIR module to a
// Tensorflow GraphDef. This API will add control dependencies and verify that
// the conversion was successful.
//
// This also converts the Tensorflow Dialect MLIR into the Tensorflow Executor
// dialect that is suitable to be exported to GraphDef. Graph -> MLIR -> Graph
// is not perfectly round trippable, so this API will attempt to make the module
// exportable and verify some properties of the Tensorflow Executor MLIR that
// are required by Graph Export. It will return an error if it cannot.
//
// Input: A MLIR Module in the Tensorflow Dialect with no
// `tf_device.cluster_func` ops.
// Output: A MLIR module in the Tensorflow Executor Dialect.
ABSL_DEPRECATED("Use tensorflow::tf2xla::v2::ConvertGraphToTfExecutor instead.")
absl::Status ExportFromTensorflowDialectToExecutor(
mlir::ModuleOp module, llvm::StringRef module_name = llvm::StringRef());
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_TF_DIALECT_TO_EXECUTOR_H_
tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor_test.cc
+133
View File
@@ -0,0 +1,133 @@
/* 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/api/v2/tf_dialect_to_executor.h"
#include <stdlib.h>
#include <cstdint>
#include <string>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/IR/BuiltinOps.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 "mlir/Parser/Parser.h" // from @llvm-project
#include "tensorflow/compiler/mlir/register_common_dialects.h"
#include "tensorflow/compiler/mlir/tf2xla/api/v2/testing/utils.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "tensorflow/core/lib/monitoring/cell_reader.h"
#include "tensorflow/core/platform/resource_loader.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
namespace {
constexpr char kExportStreamzName[] =
"/tensorflow/core/tf2xla/api/v2/tf_dialect_to_executor_dialect_status";
constexpr char kExportSuccess[] = "success";
constexpr char kExportFailed[] = "failed";
using mlir::DialectRegistry;
using mlir::MLIRContext;
using mlir::ModuleOp;
using mlir::OwningOpRef;
using ::tensorflow::monitoring::testing::CellReader;
using ::tensorflow::tf2xla::v2::testing::TestDataPath;
size_t CountSubstring(absl::string_view str, absl::string_view substr) {
size_t count = 0;
size_t idx = str.find(substr);
while (idx != std::string::npos) {
count++;
idx = str.find(substr, idx + 1);
}
return count;
}
class TensorflowDialectToExecutorTest : public ::testing::Test {
public:
TensorflowDialectToExecutorTest() {
mlir::RegisterCommonToolingDialects(registry_);
context_.appendDialectRegistry(registry_);
context_.loadAllAvailableDialects();
}
absl::Status CreateMlirModule(std::string mlir_module_filename) {
std::string mlir_module_path = TestDataPath() + mlir_module_filename;
mlir_module_ =
mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context_);
if (!mlir_module_) {
return absl::Status(
absl::StatusCode::kNotFound,
absl::StrCat("Could not find MLIR module at ", mlir_module_path));
}
return absl::OkStatus();
}
DialectRegistry registry_;
MLIRContext context_;
OwningOpRef<mlir::ModuleOp> mlir_module_;
};
TEST_F(TensorflowDialectToExecutorTest, ConvertsToExecutor) {
CellReader<int64_t> compilation_status(kExportStreamzName);
TF_ASSERT_OK(CreateMlirModule("empty_func.mlir"));
TF_EXPECT_OK(ExportFromTensorflowDialectToExecutor(*mlir_module_));
EXPECT_EQ(compilation_status.Delta(kExportSuccess), 1);
EXPECT_EQ(compilation_status.Delta(kExportFailed), 0);
}
TEST_F(TensorflowDialectToExecutorTest, ErrorsWhenCannotConvert) {
CellReader<int64_t> compilation_status(kExportStreamzName);
TF_ASSERT_OK(CreateMlirModule("invalid_executor.mlir"));
EXPECT_FALSE(ExportFromTensorflowDialectToExecutor(*mlir_module_).ok());
EXPECT_EQ(compilation_status.Delta(kExportSuccess), 0);
EXPECT_EQ(compilation_status.Delta(kExportFailed), 1);
}
TEST_F(TensorflowDialectToExecutorTest, PrunesDeadOps) {
CellReader<int64_t> compilation_status(kExportStreamzName);
TF_ASSERT_OK(CreateMlirModule("func_with_dead_ops.mlir"));
TF_EXPECT_OK(ExportFromTensorflowDialectToExecutor(*mlir_module_));
std::string module_dump;
llvm::raw_string_ostream raw_stream(module_dump);
mlir_module_->print(raw_stream);
EXPECT_EQ(compilation_status.Delta(kExportSuccess), 1);
EXPECT_EQ(compilation_status.Delta(kExportFailed), 0);
EXPECT_EQ(
CountSubstring(module_dump, "tf_executor.island wraps \"tf.Concat\""), 2);
}
} // namespace
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v2/tf_executor_to_graph.cc
+842
View File
@@ -0,0 +1,842 @@
/* Copyright 2019 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/api/v2/tf_executor_to_graph.h"
#include <string>
#include <utility>
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/container/inlined_vector.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/types/optional.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#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/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Location.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/SymbolTable.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "mlir/Support/DebugStringHelper.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/op_or_arg_name_mapper.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/convert_type.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/export_utils.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/translate_utils.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/verify_suitable_for_graph_export.h"
#include "tensorflow/compiler/mlir/utils/name_utils.h"
#include "xla/status_macros.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/graph_to_functiondef.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/framework/versions.pb.h"
#include "tensorflow/core/graph/algorithm.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/graph/regularization/util.h"
#include "tensorflow/core/graph/tensor_id.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/status.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
using mlir::BlockArgument;
using mlir::Dialect;
using mlir::Operation;
using mlir::SymbolTable;
using mlir::Value;
using mlir::func::FuncOp;
using tsl::StatusOr;
namespace {
constexpr char kDeviceAttr[] = "tf.device";
constexpr char kResourceArgUniqueIdAttr[] = "tf._resource_arg_unique_id";
constexpr char kEntryFuncAttr[] = "tf.entry_function";
constexpr char kAliasingAttr[] = "tf.aliasing_output";
// OpOrArgLocNameMapper that legalizes the returned name.
class LegalizedOpOrValLocNameMapper : public OpOrArgLocNameMapper {
private:
std::string GetName(OpOrVal op_or_val) override {
std::string name = OpOrArgLocNameMapper::GetName(op_or_val);
assert(!name.empty() && "expected non-empty name");
mlir::LegalizeNodeName(name);
return name;
}
};
// Finds first inner op if `op` is a tf_executor.island. Otherwise `op` is
// returned.
Operation* GetIslandInnerOpOrSelf(mlir::Operation* op) {
auto island = llvm::dyn_cast<mlir::tf_executor::IslandOp>(op);
if (island) return &island.GetBody().front();
return op;
}
// Stateful helper class to export a function into a Graph.
class Exporter {
public:
// Converts the given Module to a Graph. The given module should only contain
// one entry function, which is identified by name "main". This entry function
// is converted to the base of the graph graph. The rest of the functions are
// converted to the library functions in that graph.
static absl::Status Convert(mlir::ModuleOp module,
const GraphExportConfig& configs,
std::unique_ptr<Graph>* graph,
FunctionLibraryDefinition* flib_def,
absl::flat_hash_set<Node*>* control_ret_nodes);
// Converts a given FuncOp to a FunctionDef and adds it to the function
// definition library
static absl::Status ConvertLibFunction(
const GraphExportConfig& configs, const Dialect* tf_dialect,
const SymbolTable& symbol_table, FuncOp function,
FunctionLibraryDefinition* flib_def,
llvm::SmallDenseSet<FuncOp>& visited_functions);
// Converts the given FuncOp to a Graph. The arguments and returns of
// function are added to the graph with special op names kArgOp and kRetOp.
// Later on, this graph can be converted a function definition and added to
// another graph.
static absl::StatusOr<std::unique_ptr<Graph>> Convert(
const GraphExportConfig& configs, const Dialect* tf_dialect,
const SymbolTable& symbol_table, FuncOp function,
FunctionLibraryDefinition* flib_def,
llvm::SmallDenseSet<FuncOp>& visited_functions,
absl::flat_hash_set<Node*>* control_ret_nodes);
private:
explicit Exporter(const GraphExportConfig* configs, Graph* graph,
const Dialect* tf_dialect, const SymbolTable* symbol_table)
: configs_(*configs),
graph_(graph),
tf_dialect_(tf_dialect),
symbol_table_(*symbol_table) {
graph_->ToGraphDef(&graphdef_);
}
absl::Status AddArgumentNode(BlockArgument arg, unsigned index,
llvm::StringRef name);
absl::Status AddFetchNode(FuncOp function, mlir::tf_executor::FetchOp fetch,
llvm::ArrayRef<llvm::StringRef> names);
absl::Status AddInstructionNode(Operation* inst);
void UseOriginalFunctionNames(NodeDef& node_def);
absl::Status AddEdge(Operation* inst);
absl::StatusOr<std::unique_ptr<NodeDef>> GetArgumentNode(
BlockArgument arg, unsigned index, llvm::StringRef name);
absl::StatusOr<std::unique_ptr<NodeDef>> GetReturnNode(FuncOp function,
Value operand,
unsigned index,
llvm::StringRef name);
absl::Status GetControlRetNodes(
mlir::tf_executor::FetchOp fetch,
absl::flat_hash_set<Node*>* control_ret_nodes);
// Adds one edge between src_node and dst_node. If it is not a control edge,
// an index is used to find out the right operand of the dst_node.
absl::Status AddEdgeBetweenNodes(Value src, Node* dst_node,
unsigned dst_index);
const GraphExportConfig& configs_;
Graph* graph_;
GraphDef graphdef_;
LegalizedOpOrValLocNameMapper op_to_name_;
absl::flat_hash_map<Operation*, Node*> nodes_;
llvm::DenseMap<BlockArgument, Node*> args_;
// One single return operation can return multiple results, and each of them
// will be converted to one node in the graph.
typedef absl::InlinedVector<Node*, 4> NodeVector;
absl::flat_hash_map<Operation*, NodeVector> returns_;
const mlir::Dialect* tf_dialect_;
const SymbolTable& symbol_table_;
};
std::string FindFunctionName(const GraphExportConfig& configs, FuncOp func) {
if (auto original_func_name =
func->getAttrOfType<mlir::StringAttr>("tf._original_func_name");
configs.export_original_tf_func_name && original_func_name) {
return original_func_name.str();
}
return func.getName().str();
}
absl::StatusOr<std::unique_ptr<NodeDef>> Exporter::GetArgumentNode(
BlockArgument arg, unsigned index, llvm::StringRef name) {
auto func = arg.getParentRegion()->getParentOfType<FuncOp>();
auto node_def = std::make_unique<NodeDef>();
if (!name.empty())
node_def->set_name(ParseTensorName(name.str()).node());
else
node_def->set_name(
std::string(op_to_name_.GetUniqueName(func.getName().str())));
node_def->set_op(FunctionLibraryDefinition::kArgOp);
mlir::TensorType arg_type = mlir::cast<mlir::TensorType>(arg.getType());
if (auto resource_type =
mlir::dyn_cast<mlir::TF::ResourceType>(arg_type.getElementType())) {
llvm::ArrayRef<mlir::TensorType> subtypes = resource_type.getSubtypes();
if (!subtypes.empty()) {
AttrValue handle_dtypes_attr;
AttrValue handle_shapes_attr;
for (mlir::TensorType subtype : subtypes) {
DataType dtype;
TF_RETURN_IF_ERROR(ConvertToDataType(subtype.getElementType(), &dtype));
handle_dtypes_attr.mutable_list()->add_type(dtype);
SetTensorShapeProto(subtype,
handle_shapes_attr.mutable_list()->add_shape());
}
(*node_def->mutable_attr())["_handle_dtypes"] = handle_dtypes_attr;
(*node_def->mutable_attr())["_handle_shapes"] = handle_shapes_attr;
}
}
TF_RETURN_IF_ERROR(
SetShapeAttribute("_output_shapes", arg_type, node_def->mutable_attr()));
DataType dtype;
TF_RETURN_IF_ERROR(ConvertToDataType(arg_type.getElementType(), &dtype));
AttrValue type_attr;
type_attr.set_type(dtype);
(*node_def->mutable_attr())["T"] = type_attr;
AttrValue index_attr;
index_attr.set_i(index);
(*node_def->mutable_attr())["index"] = index_attr;
if (auto device_attr =
func.getArgAttrOfType<mlir::StringAttr>(index, kDeviceAttr))
*node_def->mutable_device() = device_attr.getValue().str();
llvm::ArrayRef<mlir::NamedAttribute> func_arg_i_attrs =
mlir::function_interface_impl::getArgAttrs(func, index);
absl::flat_hash_set<absl::string_view> attrs_to_ignore = {kDeviceAttr,
kAliasingAttr};
TF_RETURN_IF_ERROR(ConvertAttributes(func_arg_i_attrs, attrs_to_ignore,
/*remove_ref_type=*/false,
node_def->mutable_attr()));
return node_def;
}
absl::StatusOr<std::unique_ptr<NodeDef>> Exporter::GetReturnNode(
FuncOp function, Value operand, unsigned index, llvm::StringRef name) {
auto node_def = std::make_unique<NodeDef>();
if (!name.empty())
node_def->set_name(ParseTensorName(name.str()).node());
else
node_def->set_name(
std::string(op_to_name_.GetUniqueName(function.getName().str())));
node_def->set_op(FunctionLibraryDefinition::kRetOp);
DataType dtype;
TF_RETURN_IF_ERROR(ConvertToDataType(
mlir::cast<mlir::TensorType>(operand.getType()).getElementType(),
&dtype));
AttrValue type_attr;
type_attr.set_type(dtype);
(*node_def->mutable_attr())["T"] = type_attr;
AttrValue index_attr;
index_attr.set_i(index);
(*node_def->mutable_attr())["index"] = index_attr;
if (auto device_attr =
function.getResultAttrOfType<mlir::StringAttr>(index, kDeviceAttr))
*node_def->mutable_device() = device_attr.getValue().str();
llvm::ArrayRef<mlir::NamedAttribute> func_res_i_attrs =
function.getResultAttrs(index);
absl::flat_hash_set<absl::string_view> attrs_to_ignore = {kDeviceAttr};
TF_RETURN_IF_ERROR(ConvertAttributes(func_res_i_attrs, attrs_to_ignore,
/*remove_ref_type=*/false,
node_def->mutable_attr()));
return node_def;
}
absl::Status Exporter::AddEdgeBetweenNodes(Value src, Node* dst_node,
unsigned dst_index) {
if (auto input_result = mlir::dyn_cast<mlir::OpResult>(src)) {
auto* input_inst = GetIslandInnerOpOrSelf(input_result.getOwner());
// Replaces the input node with NextIteration sink if it is a NextIteration
// source.
if (auto next_iter_source =
llvm::dyn_cast<mlir::tf_executor::NextIterationSourceOp>(
input_inst))
input_inst = next_iter_source.GetSink();
auto node_it = nodes_.find(input_inst);
TF_RET_CHECK(node_it != nodes_.end())
<< "Use of OpResult encountered before def!";
if (mlir::isa<mlir::tf_executor::ControlType>(input_result.getType())) {
graph_->AddControlEdge(node_it->second, dst_node,
/*allow_duplicates=*/true);
} else {
graph_->AddEdge(node_it->second, input_result.getResultNumber(), dst_node,
dst_index);
}
return absl::OkStatus();
}
auto input_arg = mlir::cast<BlockArgument>(src);
auto input_node_it = args_.find(input_arg);
TF_RET_CHECK(input_node_it != args_.end())
<< "Use of BlockArgument encounted before def!";
// For argument, there is only one result output, so the index is always 0.
graph_->AddEdge(input_node_it->second, 0, dst_node, dst_index);
return absl::OkStatus();
}
absl::Status Exporter::AddEdge(Operation* inst) {
// For tf_executor.fetch, add only its data edges. Control edges are captured
// later.
if (auto fetch = llvm::dyn_cast<mlir::tf_executor::FetchOp>(inst)) {
for (auto operand_and_idx : llvm::enumerate(fetch.getOperands())) {
Value operand = operand_and_idx.value();
if (mlir::isa<mlir::tf_executor::ControlType>(operand.getType())) break;
auto* dst_node = returns_[fetch][operand_and_idx.index()];
TF_RETURN_IF_ERROR(AddEdgeBetweenNodes(operand, dst_node, 0));
}
return absl::OkStatus();
}
// For tf_executor.NextIteration.Sink, skip its token operand and add data and
// control edges with their index offset by 1.
if (auto next_iter_sink =
llvm::dyn_cast<mlir::tf_executor::NextIterationSinkOp>(inst)) {
auto* dst_node = nodes_[inst];
TF_RETURN_IF_ERROR(
AddEdgeBetweenNodes(next_iter_sink.getInput(), dst_node, 0));
for (auto control_and_idx :
llvm::enumerate(next_iter_sink.getControlInputs()))
TF_RETURN_IF_ERROR(AddEdgeBetweenNodes(control_and_idx.value(), dst_node,
control_and_idx.index() + 1));
return absl::OkStatus();
}
// For tf_executor.NextIteration.Source, op can be skipped as it is assumed
// there are no operands.
if (llvm::isa<mlir::tf_executor::NextIterationSourceOp>(inst)) {
assert(inst->getNumOperands() == 0);
return absl::OkStatus();
}
Operation* op = GetIslandInnerOpOrSelf(inst);
auto* dst_node = nodes_[op];
int operand_offset = 0;
// For tf_executor.island, add data edges from its wrapped op before control
// edges.
if (auto island = llvm::dyn_cast<mlir::tf_executor::IslandOp>(inst)) {
for (auto operand_and_idx : llvm::enumerate(op->getOperands()))
TF_RETURN_IF_ERROR(AddEdgeBetweenNodes(operand_and_idx.value(), dst_node,
operand_and_idx.index()));
operand_offset = op->getNumOperands();
}
// For all other ops (including tf_executor.island), add remaining edges.
for (auto operand_and_idx : llvm::enumerate(inst->getOperands()))
TF_RETURN_IF_ERROR(
AddEdgeBetweenNodes(operand_and_idx.value(), dst_node,
operand_and_idx.index() + operand_offset));
return absl::OkStatus();
}
void Exporter::UseOriginalFunctionNames(NodeDef& node_def) {
if (!configs_.export_original_tf_func_name) return;
auto& attrs = *node_def.mutable_attr();
auto try_use_original_func_name = [this](std::string* name) {
if (auto func = symbol_table_.lookup<FuncOp>(*name)) {
if (auto original_func_name =
func->getAttrOfType<mlir::StringAttr>("tf._original_func_name")) {
*name = original_func_name.str();
}
}
};
// Change its op name if it is a legacy call.
try_use_original_func_name(node_def.mutable_op());
// Change any function attributes in the attrs.
for (auto& iter : attrs) {
auto& attr = iter.second;
if (attr.has_func()) {
try_use_original_func_name(attr.mutable_func()->mutable_name());
} else if (attr.has_list()) {
for (auto& func_attr : *attr.mutable_list()->mutable_func()) {
try_use_original_func_name(func_attr.mutable_name());
}
}
}
}
absl::Status Exporter::AddInstructionNode(Operation* inst) {
std::unique_ptr<NodeDef> node_def;
int graph_hash_value = graph_regularization::ComputeHash(graphdef_);
auto name = op_to_name_.GetUniqueName(inst, graph_hash_value);
// Convert registered TF ops to NodeDef. Only registered ops are handled to
// ensure that PopulateDerivedAttrs adds the correct attributes.
TF_ASSIGN_OR_RETURN(node_def,
ConvertTFDialectOpToNodeDef(
inst, name, /*ignore_unregistered_attrs=*/false));
UseOriginalFunctionNames(*node_def);
TF_ASSIGN_OR_RETURN(Node * node, graph_->AddNode(std::move(*node_def)));
DCHECK(node != nullptr);
nodes_[inst] = node;
return absl::OkStatus();
}
bool IsEntryFunctionArg(BlockArgument arg) {
return arg.getParentRegion()->getParentOfType<FuncOp>().getName() == "main";
}
// Creates argument nodes from Block argument. If a name is supplied, that
// name will be used instead of generating a unique name.
absl::Status Exporter::AddArgumentNode(BlockArgument arg, unsigned index,
llvm::StringRef name) {
TF_ASSIGN_OR_RETURN(auto node_def, GetArgumentNode(arg, index, name));
TF_ASSIGN_OR_RETURN(Node * node, graph_->AddNode(std::move(*node_def)));
args_[arg] = node;
return absl::OkStatus();
}
// Creates return nodes per operand of a FetchOp. If names is supplied, those
// names will be used per node in order instead of generating a unique name.
absl::Status Exporter::AddFetchNode(FuncOp function,
mlir::tf_executor::FetchOp fetch,
llvm::ArrayRef<llvm::StringRef> names) {
auto& return_nodes = returns_[fetch];
for (auto operand_and_idx : llvm::enumerate(fetch.getOperands())) {
if (mlir::isa<mlir::tf_executor::ControlType>(
operand_and_idx.value().getType()))
break;
TF_ASSIGN_OR_RETURN(
auto node_def,
GetReturnNode(function, operand_and_idx.value(),
operand_and_idx.index(),
names.empty() ? "" : names[operand_and_idx.index()]));
TF_ASSIGN_OR_RETURN(Node * node, graph_->AddNode(std::move(*node_def)));
return_nodes.push_back(node);
}
return absl::OkStatus();
}
// Collects control ret Nodes based on tf_executor.graph's associated
// tf_executor.fetch control inputs.
absl::Status Exporter::GetControlRetNodes(
mlir::tf_executor::FetchOp fetch,
absl::flat_hash_set<Node*>* control_ret_nodes) {
for (Value fetch_operand : fetch.getOperands()) {
if (mlir::isa<mlir::tf_executor::ControlType>(fetch_operand.getType())) {
Operation* defining_op =
GetIslandInnerOpOrSelf(fetch_operand.getDefiningOp());
auto node_it = nodes_.find(defining_op);
TF_RET_CHECK(node_it != nodes_.end());
control_ret_nodes->insert(node_it->second);
}
}
return absl::OkStatus();
}
// After conversion from MLIR the input names are all blank which causes
// graph compilation to fail. This uses the edges to fix up the input names.
void FixupInputNamesFromEdges(Graph* graph) {
for (Node* n : graph->nodes()) {
if (n->IsOp()) {
NodeDef* node_def = n->mutable_def();
node_def->clear_input();
for (const Edge* e : n->in_edges()) {
Node* src = e->src();
if (src->IsOp()) {
Graph::AddInput(node_def, src->name(), e->src_output());
}
}
}
}
}
absl::StatusOr<std::unique_ptr<Graph>> Exporter::Convert(
const GraphExportConfig& configs, const Dialect* tf_dialect,
const SymbolTable& symbol_table, FuncOp function,
FunctionLibraryDefinition* flib_def,
llvm::SmallDenseSet<FuncOp>& visited_functions,
absl::flat_hash_set<Node*>* control_ret_nodes) {
mlir::Block& block = function.front();
// Extract input & output names if set.
llvm::SmallVector<llvm::StringRef, 2> input_names;
llvm::SmallVector<llvm::StringRef, 2> output_names;
llvm::SmallVector<llvm::StringRef, 2> unique_output_names;
auto dict_attr =
function->getAttrOfType<mlir::DictionaryAttr>(kEntryFuncAttr);
if (dict_attr) {
TF_RET_CHECK(mlir::isa<mlir::StringAttr>(dict_attr.get("inputs")))
<< "inputs missing in entry function attribute";
TF_RET_CHECK(mlir::isa<mlir::StringAttr>(dict_attr.get("outputs")))
<< "outputs missing in entry function attribute";
mlir::cast<mlir::StringAttr>(dict_attr.get("inputs"))
.getValue()
.split(input_names, ',', /*MaxSplit=*/-1, /*KeepEmpty=*/false);
mlir::cast<mlir::StringAttr>(dict_attr.get("outputs"))
.getValue()
.split(output_names, ',', /*MaxSplit=*/-1, /*KeepEmpty=*/false);
}
auto graph = std::make_unique<Graph>(OpRegistry::Global());
// Extract version info.
VersionDef versions;
auto module = function->getParentOfType<mlir::ModuleOp>();
if (mlir::succeeded(ExtractTfVersions(module, &versions))) {
graph->set_versions(versions);
}
Exporter exporter(&configs, graph.get(), tf_dialect, &symbol_table);
auto graph_op = llvm::cast<mlir::tf_executor::GraphOp>(block.front());
// Set input and output names and increment the use counter for them to help
// generate unique names.
if (!output_names.empty()) {
const int num_data_results = graph_op.getNumResults();
const int64_t output_names_size = output_names.size();
TF_RET_CHECK(output_names_size == num_data_results)
<< "output names (" << output_names.size()
<< ") != terminator operands (" << num_data_results << ")";
llvm::DenseMap<Operation*, llvm::StringRef> output_op_to_name;
llvm::StringMap<Operation*> name_to_op;
for (const auto& it : llvm::enumerate(graph_op.GetFetch().getOperands())) {
// Skip control rets.
const int64_t index = it.index();
if (index >= num_data_results) break;
// TODO(jpienaar): If there is a result index specified, ensure only one
// and that it matches the result index of the op.
std::string name(output_names[index]);
auto tensor_id = ParseTensorName(name);
std::string tensor_id_node(tensor_id.node());
assert(!tensor_id_node.empty() && "expected non-empty name");
mlir::LegalizeNodeName(tensor_id_node);
// Ensure name does not get reused.
unique_output_names.push_back(
exporter.op_to_name_.GetUniqueName(tensor_id_node));
}
}
if (!input_names.empty()) {
TF_RET_CHECK(input_names.size() == block.getNumArguments());
for (const auto& it : llvm::enumerate(function.getArguments())) {
// TODO(lyandy): Update when changing feed/fetch import.
std::string name(input_names[it.index()]);
assert(!name.empty() && "expected non-empty name");
mlir::LegalizeNodeName(name);
auto tensor_id = ParseTensorName(name);
TF_RET_CHECK(tensor_id.index() == 0)
<< "input port designation not supported";
// Only assign user of argument the input name if the main graph did not
// have its _Arg nodes lifted into the functions arguments.
// Ensure name does not get reused.
(void)exporter.op_to_name_.GetUniqueName(name);
}
}
// Adds nodes for basic block (function) arguments.
for (auto it : llvm::enumerate(block.getArguments())) {
int index = it.index();
auto arg = it.value();
mlir::Type type = arg.getType();
if (!mlir::isa<mlir::TensorType>(type)) {
return absl::InvalidArgumentError(absl::StrCat(
"FuncOps arguments must have tensor types. Found ",
mlir::debugString(type), " in function ", function.getName().str()));
}
TF_RETURN_IF_ERROR(exporter.AddArgumentNode(
arg, index, !input_names.empty() ? input_names[index] : ""));
}
auto convert_called_function = [&](llvm::StringRef name) {
auto func = symbol_table.lookup<FuncOp>(name);
if (func != nullptr) {
TF_RETURN_IF_ERROR(ConvertLibFunction(configs, tf_dialect, symbol_table,
func, flib_def, visited_functions));
// TODO(prakalps): Optimize to only add the requested function to graph
// library rather than the all the functions exported so far.
TF_RETURN_IF_ERROR(graph->mutable_flib_def()->AddLibrary(*flib_def));
}
return absl::OkStatus();
};
// Adds nodes for operations.
for (Operation& inst : graph_op.GetBody()) {
for (auto type : inst.getResultTypes())
if (!mlir::isa<mlir::TensorType, mlir::tf_executor::ControlType,
mlir::tf_executor::TokenType>(type))
return absl::InvalidArgumentError(absl::StrCat(
"Values must be of tensor type, TensorFlow control type, or "
"TensorFlow token type. Found ",
mlir::debugString(type)));
if (llvm::isa<mlir::tf_executor::NextIterationSourceOp>(inst)) {
// Skip tf_executor.NextIteration.Source as associated
// tf_executor.NextIteration.Sink will be used instead.
continue;
} else if (auto fetch = llvm::dyn_cast<mlir::tf_executor::FetchOp>(inst)) {
TF_RETURN_IF_ERROR(
exporter.AddFetchNode(function, fetch, unique_output_names));
} else if (auto island =
llvm::dyn_cast<mlir::tf_executor::IslandOp>(inst)) {
Operation& inner_op = island.GetBody().front();
auto op_name = GetTensorFlowOpName(inner_op.getName().getStringRef());
if (llvm::isa<FuncOp>(inner_op) && op_name.ok()) {
// If it is TF Control dialect specific op, look up custom operation
// in the module and first convert that, then add it to function
// definition library
// TODO(prakalps): If two functions have cyclic dependence, this will
// introduce an infinite loop.
TF_RETURN_IF_ERROR(convert_called_function(op_name.value().str()));
}
if (IsLegacyCallInstruction(&inner_op)) {
TF_RETURN_IF_ERROR(convert_called_function(
inner_op.getAttrOfType<mlir::SymbolRefAttr>("f")
.getLeafReference()
.getValue()));
}
TF_RETURN_IF_ERROR(exporter.AddInstructionNode(&inner_op));
} else {
TF_RETURN_IF_ERROR(exporter.AddInstructionNode(&inst));
}
}
// Adds edges between the argument, operation and return nodes.
for (Operation& inst : graph_op.GetBody()) {
TF_RETURN_IF_ERROR(exporter.AddEdge(&inst));
}
// Fixes the edges between the inserted nodes and special "_SOURCE" and
// "_SINK".
FixupSourceAndSinkEdges(graph.get());
FixupInputNamesFromEdges(graph.get());
TF_RETURN_IF_ERROR(
exporter.GetControlRetNodes(graph_op.GetFetch(), control_ret_nodes));
return graph;
}
absl::Status Exporter::ConvertLibFunction(
const GraphExportConfig& configs, const Dialect* tf_dialect,
const SymbolTable& symbol_table, FuncOp function,
FunctionLibraryDefinition* flib_def,
llvm::SmallDenseSet<FuncOp>& visited_functions) {
// Return early if the function has already been exported.
bool is_new_function = visited_functions.insert(function).second;
if (!is_new_function) return absl::OkStatus();
auto function_name = FindFunctionName(configs, function);
// TODO(fengliuai): use a small flib_def to reduce overhead
absl::flat_hash_set<Node*> control_ret_nodes;
TF_ASSIGN_OR_RETURN(
auto sub_graph,
Exporter::Convert(configs, tf_dialect, symbol_table, function, flib_def,
visited_functions, &control_ret_nodes));
const auto control_ret = [&](const Node* n) -> std::optional<std::string> {
return control_ret_nodes.contains(n)
? std::make_optional<std::string>(n->name())
: std::nullopt;
};
FunctionDef func_def;
TF_RETURN_IF_ERROR(
GraphToFunctionDef(*sub_graph, function_name, control_ret, &func_def));
// Checks for gradient attribute. If present converts the gradient function
// and populates the GradientDef.
auto grad_string = mlir::TF::TensorFlowDialect::GetGradientAttrName();
if (auto attr =
function->getAttrOfType<mlir::FlatSymbolRefAttr>(grad_string)) {
auto grad_func = symbol_table.lookup<FuncOp>(attr.getValue());
TF_RETURN_IF_ERROR(ConvertLibFunction(configs, tf_dialect, symbol_table,
grad_func, flib_def,
visited_functions));
GradientDef grad;
grad.set_function_name(function_name);
grad.set_gradient_func(grad_func.getName().str());
TF_RETURN_IF_ERROR(flib_def->AddGradientDef(grad));
}
auto stateful_string = mlir::TF::TensorFlowDialect::GetStatefulAttrName();
if (auto attr = function->getAttrOfType<mlir::UnitAttr>(stateful_string)) {
func_def.mutable_signature()->set_is_stateful(true);
}
// Ignore the gradient and is_stateful attribute on the function as they have
// been handled above. Ignore the entry func attribute as it is an MLIR
// metadata attribute and is not required in the function definition.
absl::flat_hash_set<absl::string_view> attrs_to_ignore = {
grad_string.data(), stateful_string.data(), kEntryFuncAttr};
llvm::SmallVector<mlir::NamedAttribute, 8> funcAttrs(
function->getDialectAttrs());
TF_RETURN_IF_ERROR(ConvertAttributes(funcAttrs, attrs_to_ignore,
/*remove_ref_type=*/false,
func_def.mutable_attr()));
for (int i = 0, e = function.getNumArguments(); i < e; ++i) {
if (auto resource_arg_unique_id_attr =
function.getArgAttrOfType<mlir::IntegerAttr>(
i, kResourceArgUniqueIdAttr)) {
(*func_def.mutable_resource_arg_unique_id())[i] =
resource_arg_unique_id_attr.getInt();
}
}
return flib_def->AddFunctionDef(std::move(func_def));
}
absl::Status Exporter::Convert(mlir::ModuleOp module,
const GraphExportConfig& configs,
std::unique_ptr<Graph>* graph,
FunctionLibraryDefinition* flib_def,
absl::flat_hash_set<Node*>* control_ret_nodes) {
mlir::StringAttr entry_func_id =
mlir::StringAttr::get(module.getContext(), "main");
std::optional<FuncOp> entry_func;
FunctionLibraryDefinition temp_flib_def(OpRegistry::Global(),
FunctionDefLibrary());
llvm::SmallDenseSet<FuncOp> visited_functions;
auto tf_dialect = module.getContext()->getLoadedDialect("tf");
// Construct SymbolTable to enable cheap function lookups. The cost
// of constructing the table is offset by the number of queries.
SymbolTable symbol_table(module);
for (auto function : module.getOps<FuncOp>()) {
if (function.isExternal())
return absl::FailedPreconditionError("External functions not supported");
if (function.getName() == entry_func_id &&
!configs.export_entry_func_to_flib) {
entry_func.emplace(function);
} else {
TF_RETURN_IF_ERROR(ConvertLibFunction(configs, tf_dialect, symbol_table,
function, &temp_flib_def,
visited_functions));
}
}
if (flib_def != nullptr) {
TF_RETURN_IF_ERROR(flib_def->AddLibrary(temp_flib_def));
}
if (!configs.export_entry_func_to_flib) {
if (!entry_func.has_value())
return absl::FailedPreconditionError(
"entry function `main` must be present");
// Updates the graph and the function library definition.
TF_ASSIGN_OR_RETURN(
*graph, Exporter::Convert(configs, tf_dialect, symbol_table,
entry_func.value(), &temp_flib_def,
visited_functions, control_ret_nodes));
// Add FunctionDefs and GradientDefs of MLIR functions to graph's function
// library. If duplicate FunctionDefs already exist (can happen if exporter
// had already added some FunctionDefs to the library to support legacy
// calls), they are ignored.
TF_RETURN_IF_ERROR(
graph->get()->mutable_flib_def()->AddLibrary(temp_flib_def));
} else if (graph != nullptr) {
TF_RETURN_IF_ERROR(
graph->get()->mutable_flib_def()->AddLibrary(std::move(*flib_def)));
}
return absl::OkStatus();
}
} // namespace
absl::Status ConvertTfExecutorToGraph(
mlir::ModuleOp module, const GraphExportConfig& configs,
std::unique_ptr<Graph>* graph, FunctionLibraryDefinition* flib_def,
absl::flat_hash_set<Node*>* control_ret_nodes) {
mlir::StatusScopedDiagnosticHandler sh(module.getContext());
if (failed(VerifyExportSuitable(module))) return sh.ConsumeStatus();
return sh.Combine(
Exporter::Convert(module, configs, graph, flib_def, control_ret_nodes));
}
absl::Status ConvertMlirFunctionToFunctionLibraryDef(
FuncOp func, const GraphExportConfig& configs, FunctionDef* function_def) {
Dialect* tf_dialect = func.getContext()->getLoadedDialect("tf");
FunctionLibraryDefinition flib_def(OpRegistry::Global(),
FunctionDefLibrary());
llvm::SmallDenseSet<FuncOp> visited_functions;
// Construct SymbolTable to enable cheap function lookups. The cost
// of constructing the table is offset by the number of queries. Even
// though this only converts one function in theory, this function
// may have gradient associated which would result in a lookup. This
// could be made lazy if we find this to be broad.
SymbolTable symbol_table(func->getParentOfType<mlir::ModuleOp>());
TF_RETURN_IF_ERROR(Exporter::ConvertLibFunction(
configs, tf_dialect, symbol_table, func, &flib_def, visited_functions));
auto name = FindFunctionName(configs, func);
const FunctionDef* func_def = flib_def.Find(name);
if (func_def != nullptr) {
*function_def = *func_def;
return absl::OkStatus();
}
return absl::InvalidArgumentError(
absl::StrCat("Function '", name,
"' couldn't be found in the FunctionDefLibrary after "
"converting from MLIR"));
}
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v2/tf_executor_to_graph.h
+54
View File
@@ -0,0 +1,54 @@
/* Copyright 2019 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_API_V2_TF_EXECUTOR_TO_GRAPH_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_TF_EXECUTOR_TO_GRAPH_H_
#include "absl/base/attributes.h"
#include "absl/container/flat_hash_set.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/graph/graph.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
// Converts an MLIR module to TensorFlow graph and FunctionLibraryDefinition.
// The "main" function of the module is stored in the graph and the rest of
// functions are stored in the library. Control ret nodes are stored separately
// in `control_ret_nodes`.
absl::Status ConvertTfExecutorToGraph(
mlir::ModuleOp module, const GraphExportConfig& configs,
std::unique_ptr<Graph>* graph, FunctionLibraryDefinition* flib_def,
absl::flat_hash_set<Node*>* control_ret_nodes);
// Converts an MLIR function and adds it to a FunctionLibraryDefinition.
absl::Status ConvertMlirFunctionToFunctionLibraryDef(
mlir::func::FuncOp func, const GraphExportConfig& configs,
FunctionDef* function_def);
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_TF_EXECUTOR_TO_GRAPH_H_
tensorflow/compiler/mlir/tf2xla/api/v2/tf_executor_to_graph_test.cc
+110
View File
@@ -0,0 +1,110 @@
/* 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/api/v2/tf_executor_to_graph.h"
#include <stdlib.h>
#include <memory>
#include <string>
#include <gtest/gtest.h>
#include "absl/container/flat_hash_set.h"
#include "absl/status/statusor.h"
#include "mlir/IR/BuiltinOps.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 "mlir/Parser/Parser.h" // from @llvm-project
#include "riegeli/bytes/fd_reader.h" // from @riegeli
#include "riegeli/bytes/read_all.h" // from @riegeli
#include "tensorflow/compiler/mlir/register_common_dialects.h"
#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/platform/resource_loader.h"
#include "tsl/platform/protobuf.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
namespace {
using mlir::DialectRegistry;
using mlir::MLIRContext;
using mlir::ModuleOp;
using mlir::OwningOpRef;
std::string TestDataPath() {
return tensorflow::GetDataDependencyFilepath(
"tensorflow/compiler/mlir/tf2xla/api/v2/testdata/");
}
class TfExecutorToGraphTest : public ::testing::Test {
public:
TfExecutorToGraphTest() {
mlir::RegisterCommonToolingDialects(registry_);
context_.appendDialectRegistry(registry_);
context_.loadAllAvailableDialects();
}
absl::StatusOr<OwningOpRef<mlir::ModuleOp>> CreateMlirModule(
std::string mlir_module_filename) {
std::string mlir_module_path = TestDataPath() + mlir_module_filename;
return mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context_);
}
GraphDef CreateGraphDef(std::string graphdef_filename) {
std::string file_path = TestDataPath() + graphdef_filename;
std::string contents;
GraphDef graph_def;
auto status = riegeli::ReadAll(riegeli::FdReader(file_path), contents);
if (!status.ok()) {
return graph_def;
}
tsl::protobuf::TextFormat::ParseFromString(contents, &graph_def);
return graph_def;
}
DialectRegistry registry_;
MLIRContext context_;
OwningOpRef<mlir::ModuleOp> mlir_module_;
};
TEST_F(TfExecutorToGraphTest, ConvertMlirToGraphSucceeds) {
auto valid_executor_module = CreateMlirModule("valid_executor.mlir");
GraphExportConfig confs;
absl::flat_hash_set<Node*> control_ret_nodes;
FunctionLibraryDefinition flib_def(OpRegistry::Global(),
FunctionDefLibrary());
auto result_graph = std::make_unique<Graph>(flib_def);
TF_ASSERT_OK(ConvertTfExecutorToGraph(valid_executor_module.value().get(),
confs, &result_graph, &flib_def,
&control_ret_nodes));
GraphDef result_graphdef;
result_graph->ToGraphDef(&result_graphdef);
GraphDef expected_graphdef = CreateGraphDef("valid_graph.txt");
EXPECT_EQ(result_graphdef.DebugString(), expected_graphdef.DebugString());
}
} // namespace
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/BUILD
+429
View File
@@ -0,0 +1,429 @@
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load(
"//tensorflow/core/platform:build_config.bzl",
"tf_proto_library",
)
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow/compiler/mlir/tensorflow/transforms:__pkg__",
"//tensorflow/compiler/mlir/tf2xla/api/v1:__subpackages__",
"//tensorflow/compiler/mlir/tf2xla/api/v2:__subpackages__",
"//tensorflow/core/tfrt/saved_model:__pkg__",
],
)
cc_library(
name = "compilation_timer",
hdrs = ["compilation_timer.h"],
deps = [
"//tensorflow/core/platform:profile_utils_cpu_utils",
],
)
tf_cc_test(
name = "compilation_timer_test",
srcs = ["compilation_timer_test.cc"],
deps = [
":compilation_timer",
"@com_google_absl//absl/time",
"@com_google_googletest//:gtest_main",
],
)
cc_library(
name = "test_matchers",
testonly = True,
hdrs = ["test_matchers.h"],
deps = [
"//tensorflow/compiler/mlir/tf2xla:compile_mlir_util",
"@com_google_absl//absl/status:statusor",
"@com_google_googletest//:gtest",
"@xla//xla/tsl/platform:statusor",
],
)
tf_cc_test(
name = "test_matchers_test",
srcs = ["test_matchers_test.cc"],
deps = [
":test_matchers",
"//tensorflow/compiler/mlir/tf2xla:compile_mlir_util",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/core:lib",
"//tensorflow/core/lib/monitoring:cell_reader",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_googletest//:gtest_main",
"@xla//xla/hlo/builder:xla_computation",
"@xla//xla/service:hlo_proto_cc",
"@xla//xla/tsl/platform:statusor",
],
)
cc_library(
name = "mlir_pass_instrumentation",
srcs = ["mlir_pass_instrumentation.cc"],
hdrs = ["mlir_pass_instrumentation.h"],
deps = [
"//tensorflow/core/platform:logging",
"@com_google_absl//absl/log",
"@llvm-project//mlir:Pass",
],
)
tf_cc_test(
name = "mlir_pass_instrumentation_test",
srcs = ["mlir_pass_instrumentation_test.cc"],
deps = [
":mlir_pass_instrumentation",
"//tensorflow/compiler/mlir/tf2xla/api/v1:compile_mlir_util_no_tf_dialect_passes",
"//tensorflow/core:test",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:string_view",
"@com_google_googletest//:gtest_main",
],
)
cc_library(
name = "legalize_tf_mlir",
srcs = ["legalize_tf_mlir.cc"],
hdrs = ["legalize_tf_mlir.h"],
deps = [
":compilation_timer",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:serialize_mlir_module_utils",
"//tensorflow/compiler/mlir/tensorflow/transforms:set_tpu_infeed_layout",
"//tensorflow/compiler/mlir/tf2xla:compile_mlir_util",
"//tensorflow/compiler/tf2xla:layout_util",
"//tensorflow/compiler/tf2xla:xla_compiler",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
"//tensorflow/core/tpu:tpu_compile",
"//tensorflow/core/tpu/kernels:tpu_compile_op_support",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@stablehlo//:register",
"@tsl//tsl/platform:error_logging",
"@xla//xla:shape_util",
"@xla//xla/mlir_hlo:hlo_dialect_registration",
"@xla//xla/tsl/platform:statusor",
],
)
tf_cc_test(
name = "legalize_tf_mlir_test",
srcs = ["legalize_tf_mlir_test.cc"],
deps = [
":legalize_tf_mlir",
":test_matchers",
"//tensorflow/compiler/jit",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:deserialize_mlir_module_utils",
"//tensorflow/compiler/tf2xla:xla_compiler",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/core:framework",
"//tensorflow/core:test_main",
"//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
"//tensorflow/core/tpu/kernels:tpu_compile_op_support",
"@com_google_absl//absl/status:statusor",
"@com_google_googletest//:gtest",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@stablehlo//:register",
"@xla//xla:shape_util",
"@xla//xla/mlir_hlo:hlo_dialect_registration",
"@xla//xla/tsl/platform:errors",
],
)
cc_library(
name = "legalize_tf_to_hlo",
srcs = ["legalize_tf_to_hlo.cc"],
hdrs = ["legalize_tf_to_hlo.h"],
deps = [
":legalize_tf_mlir",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:deserialize_mlir_module_utils",
"//tensorflow/compiler/mlir/tensorflow/transforms:set_tpu_infeed_layout",
"//tensorflow/compiler/mlir/tf2xla/api/v1:compile_tf_graph",
"//tensorflow/compiler/tf2xla:layout_util",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/compiler/tf2xla:xla_tpu_backend_registration",
"//tensorflow/core:framework",
"//tensorflow/core/platform:status",
"//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
"//tensorflow/core/tpu/kernels:tpu_compile_op_support",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:DataLayoutInterfaces",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@stablehlo//:register",
"@xla//xla:shape_util",
"@xla//xla/client:compile_only_client",
"@xla//xla/mlir_hlo:hlo_dialect_registration",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:statusor",
],
)
tf_cc_test(
name = "legalize_tf_to_hlo_test",
srcs = ["legalize_tf_to_hlo_test.cc"],
deps = [
":legalize_tf_to_hlo",
":test_matchers",
"//tensorflow/compiler/jit",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:deserialize_mlir_module_utils",
"//tensorflow/compiler/tf2xla:xla_compiler",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test_main",
"//tensorflow/core/lib/monitoring:cell_reader",
"//tensorflow/core/protobuf:for_core_protos_cc",
"//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
"//tensorflow/core/tpu/kernels:tpu_compile_op_support",
"@com_google_absl//absl/status:statusor",
"@com_google_googletest//:gtest",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@stablehlo//:register",
"@xla//xla:shape_util",
"@xla//xla/client:client_library",
"@xla//xla/mlir_hlo:hlo_dialect_registration",
"@xla//xla/stream_executor:platform",
"@xla//xla/stream_executor:platform_manager",
"@xla//xla/tsl/platform:errors",
],
)
cc_library(
name = "clustering_bridge_passes",
srcs = ["clustering_bridge_passes.cc"],
hdrs = ["clustering_bridge_passes.h"],
deps = [
"//tensorflow/compiler/jit:flags_headers",
"//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
"//tensorflow/compiler/mlir/tensorflow/transforms:verify_no_outside_compilation_markers_pass",
"//tensorflow/compiler/mlir/tensorflow/transforms/sparsecore:sparsecore_passes",
"//tensorflow/compiler/mlir/tf2xla/internal/passes:clustering_passes",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/log",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Transforms",
],
)
tf_cc_test(
name = "clustering_bridge_passes_test",
srcs = ["clustering_bridge_passes_test.cc"],
deps = [
":clustering_bridge_passes",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:Pass",
],
)
cc_library(
name = "logging_hooks",
srcs = ["logging_hooks.cc"],
hdrs = ["logging_hooks.h"],
deps = [
"//tensorflow/compiler/mlir/tensorflow:bridge_logger",
"//tensorflow/core:framework",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:Pass",
],
)
tf_cc_test(
name = "logging_hooks_test",
srcs = ["logging_hooks_test.cc"],
data = [
"testdata/dead_const.mlir",
],
deps = [
":logging_hooks",
"//tensorflow/compiler/mlir:register_common_dialects",
"//tensorflow/core:lib",
"//tensorflow/core:test",
"//tensorflow/core/platform:resource_loader",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Transforms",
"@xla//xla/tsl/lib/core:status_test_util",
"@xla//xla/tsl/platform:status",
],
)
cc_library(
name = "mlir_bridge_pass_util",
srcs = ["mlir_bridge_pass_util.cc"],
hdrs = ["mlir_bridge_pass_util.h"],
visibility = ["//tensorflow/compiler/tf2xla:__pkg__"],
deps = [
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/tf2xla:tf2xla_defs",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/common_runtime:function_body",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:string_view",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
"@xla//xla/tsl/platform:status",
],
alwayslink = 1,
)
tf_cc_test(
name = "mlir_bridge_pass_util_test",
srcs = ["mlir_bridge_pass_util_test.cc"],
deps = [
":mlir_bridge_pass_util",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"//tensorflow/cc:array_ops",
"//tensorflow/cc:function_ops",
"//tensorflow/cc:functional_ops",
"//tensorflow/cc:ops",
"//tensorflow/cc:scope",
"//tensorflow/cc:tpu_ops",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/tf2xla:tf2xla_defs",
"//tensorflow/core/protobuf:for_core_protos_cc",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core/platform:types",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test_main",
"//tensorflow/core:testlib",
"//tensorflow/core/platform:enable_tf2_utils",
# "//tensorflow/core/platform:resource_loader",
"@xla//xla/tsl/lib/core:status_test_util",
],
)
tf_proto_library(
name = "reproducer_proto",
srcs = ["reproducer.proto"],
protodeps = [
"//tensorflow/core:protos_all",
"//tensorflow/core/protobuf/tpu:compile_metadata_proto",
],
visibility = [
"//learning/brain/mlir/bridge:__pkg__",
"//tensorflow/compiler/mlir/tf2xla/api/v2:__pkg__",
],
)
cc_library(
name = "node_order",
srcs = ["node_order.cc"],
hdrs = ["node_order.h"],
deps = [
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:lib",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
],
)
tf_cc_test(
name = "node_order_test",
size = "small",
srcs = [
"node_order_test.cc",
],
deps = [
":node_order",
"//tensorflow/cc:cc_ops",
"//tensorflow/cc:cc_ops_internal",
"//tensorflow/cc:function_ops",
"//tensorflow/cc:sendrecv_ops",
"//tensorflow/core",
"//tensorflow/core:core_cpu",
"//tensorflow/core:core_cpu_internal",
"//tensorflow/core:direct_session_internal",
"//tensorflow/core:framework",
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:ops",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core:testlib",
"@com_google_absl//absl/log:check",
"@com_google_googletest//:gtest",
],
)
cc_library(
name = "graph_to_tf_executor_util",
srcs = ["graph_to_tf_executor_util.cc"],
hdrs = ["graph_to_tf_executor_util.h"],
deps = [
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:framework_types_hdr",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/common_runtime:function_body",
"//tensorflow/core/platform:enable_tf2_utils",
"//tensorflow/core/protobuf:for_core_protos_cc",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:string_view",
],
)
tf_cc_test(
name = "graph_to_tf_executor_util_test",
srcs = ["graph_to_tf_executor_util_test.cc"],
deps = [
":graph_to_tf_executor_util",
"//tensorflow/cc:array_ops",
"//tensorflow/cc:cc_ops",
"//tensorflow/cc:functional_ops",
"//tensorflow/cc:ops",
"//tensorflow/cc:scope",
"//tensorflow/cc:tpu_ops",
"//tensorflow/compiler/tf2xla/ops:xla_ops",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/framework:tensor_testutil",
"//tensorflow/core/platform:enable_tf2_utils",
"//tensorflow/core/platform:types",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest_main",
"@xla//xla/tsl/lib/core:status_test_util",
],
)
tensorflow/compiler/mlir/tf2xla/internal/clustering_bridge_passes.cc
+220
View File
@@ -0,0 +1,220 @@
/* 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/clustering_bridge_passes.h"
#include <string>
#include "absl/log/log.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "mlir/Transforms/Passes.h" // from @llvm-project
#include "tensorflow/compiler/jit/flags.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/sparsecore/sparsecore_passes.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
using mlir::OpPassManager;
using mlir::func::FuncOp;
// Adds replicated Bridge clustering pipeline passes to the given pass_manager.
// Does not run them.
void AddReplicatedBridgeClusteringPipelinePasses(OpPassManager& pm,
llvm::StringRef module_name) {
// The following ops must be preserved regardless of reachability. Ideally,
// all graphs should have control dependencies to enforce this but this is
// currently not the case (see b/177478741).
const llvm::SmallVector<std::string, 4> ops_to_preserve = {
"tf.TPUReplicateMetadata", "tf.TPUCompilationResult",
"tf.TPUReplicatedOutput"};
bool strict_clusters =
tensorflow::GetMlirCommonFlags()->tf_mlir_enable_strict_clusters;
pm.addNestedPass<FuncOp>(
mlir::tf_executor::CreateTFExecutorGraphPruningPass(ops_to_preserve));
// It is assumed at this stage there are no V1 control flow ops as Graph
// functionalization is ran before import. Ops can be lifted out of
// tf_executor dialect islands/graphs.
pm.addNestedPass<FuncOp>(
mlir::CreateExecutorDialectToFunctionalConversionPass());
// Guarantee all functions have one use, which enables more exact shape
// inference.
pm.addPass(mlir::TF::CreateGuaranteeAllFuncsOneUsePass());
// Run shape inference so that tf_executor/tf_device ops created later will
// likely to inherit more concrete types.
pm.addPass(mlir::TF::CreateTFShapeInferencePass());
pm.addNestedPass<FuncOp>(mlir::TFTPU::CreateTPUPartitionedOpConversionPass());
pm.addNestedPass<FuncOp>(
mlir::TFTPU::CreateTPUReorderReplicateAndPartitionedInputsPass());
pm.addNestedPass<FuncOp>(mlir::TF::CreateDecomposeReduceDatasetPass());
// Only one of EmbeddingSequencing and EmbeddingPipelining will actually
// run and the logic is in EmbeddingPipeliningPass. If the pipelining pass
// runs, embedding attributes are stripped and the sequencing pass will have
// no effect. If the pipelining pass doesn't run, embedding attributes are
// preserved and the sequencing rewrite will trigger.
pm.addPass(mlir::TFDevice::CreateEmbeddingPipeliningPass());
pm.addPass(mlir::TFDevice::CreateEmbeddingSequencingPass());
pm.addPass(tensorflow::tf2xla::internal::CreateTPUClusterFormationPass(
strict_clusters));
// CreateEmbeddingPipeliningPass may have created more functions, but
// TPUClusterCleanup and OutsideCompiledToHostLaunch need every function to be
// only called from one cluster. Here, we choose to fix the all-funcs-one-use
// invariant right before it's needed, not after it's been broken.
pm.addPass(mlir::TF::CreateGuaranteeAllFuncsOneUsePass());
// Run TPU cluster cleanup attributes so ops with no outside compiled
// attribute have no host device attribute.
pm.addPass(mlir::TFTPU::CreateTPUClusterCleanupAttributesPass());
pm.addNestedPass<FuncOp>(mlir::TFDevice::CreateDeviceAttributeToLaunchPass());
// Running canonicalizer before decomposing resource ops in cluster helps the
// latter pass to converge faster as it does not have to spend time folding
// away dead ops.
pm.addNestedPass<FuncOp>(mlir::createCanonicalizerPass());
// Place DecomposeResourceOpsPass before TFExecutorConstantSinking pass
// because DecomposeResourceOpsPass uses pattern rewriter which hoists
// changed constants out of tf_device.Launch.
pm.addPass(mlir::TFDevice::CreateDecomposeResourceOpsInClusterPass());
// Encode this in its own scope so that func_pm is not mistakenly used
// later on.
{
OpPassManager& func_pm = pm.nest<FuncOp>();
func_pm.addPass(mlir::TFTPU::CreateTPUHostComputationExpansionPass());
func_pm.addPass(mlir::TFTPU::CreateTPUUpdateEmbeddingEnqueueOpInputsPass());
}
// TODO(b/173622615): Once OutsideCompilation is represented by launch op and
// the remaining passes including Inliner support it, remove this
// LaunchToDeviceAttributePass. This LaunchToDeviceAttribute pass needs to
// come before TPUClusterCleanupAttributes pass or else the device attribute
// will be removed from launch causing an error.
pm.addNestedPass<FuncOp>(mlir::TFDevice::CreateLaunchToDeviceAttributePass());
// Note that the region-based control-flow produced here still contains
// function call ops which get inlined by the subsequent inliner pass.
pm.addPass(mlir::TF::CreateTFFunctionalControlFlowToRegions());
pm.addPass(mlir::createInlinerPass());
pm.addNestedPass<FuncOp>(
mlir::TF::CreateDropWhileShapeInvariantInDeviceClusterPass());
// Run another shape inference pass because resource decomposition might have
// created new partial types. Also, after dropping `shape_invariant` attribute
// from While/WhileRegion ops within cluster would lead to more precise
// shapes.
pm.addPass(mlir::TF::CreateTFShapeInferencePass());
pm.addNestedPass<FuncOp>(mlir::createCanonicalizerPass());
pm.addPass(mlir::TFTPU::CreateTPUClusterCleanupAttributesPass());
pm.addPass(mlir::TFDevice::CreateResourceOpLiftingPass());
// Re-run the canonicalizer pass as some cleanup during resource op lifting
// pass opens up some opportunities for canonicalization of cluster ops.
// Specifically, we want to eliminate pass through results from the cluster
// op.
pm.addNestedPass<FuncOp>(mlir::createCanonicalizerPass());
pm.addNestedPass<FuncOp>(mlir::createCSEPass());
if (tensorflow::GetMlirCommonFlags()
->tf_mlir_enable_merge_control_flow_pass) {
pm.addPass(mlir::TFDevice::CreateMergeControlFlowPass());
}
pm.addPass(
tensorflow::tf2xla::internal::CreateMarkOpsForOutsideCompilationPass());
pm.addPass(tensorflow::tf2xla::internal::
CreateExtractHeadTailOutsideCompilationPass());
pm.addPass(
tensorflow::tf2xla::internal::CreateExtractOutsideCompilationPass());
pm.addNestedPass<FuncOp>(
mlir::TFDevice::CreateVerifyNoOutsideCompilationMarkersPass());
pm.addNestedPass<FuncOp>(mlir::TFDevice::CreateClusterConstantSinkingPass());
pm.addPass(mlir::TF::CreateResourceDeviceInferencePass());
pm.addNestedPass<FuncOp>(
tensorflow::tf2xla::internal::CreateHoistBroadcastReadPass());
pm.addNestedPass<FuncOp>(
tensorflow::tf2xla::internal::CreateXlaBroadcastPass());
pm.addPass(mlir::TFDevice::CreateClusterOutliningPass());
pm.addPass(mlir::TFTPU::CreateTPUResourceReadForWritePass());
pm.addPass(mlir::TFDevice::CreateMarkInputOutputAliasesPass());
pm.addPass(
tensorflow::tf2xla::internal::CreateTPUShardingIdentificationPass());
pm.addNestedPass<FuncOp>(
mlir::TFTPU::CreateTPUResourceReadsWritesPartitioningPass());
pm.addPass(mlir::TFDevice::CreateAnnotateParameterReplicationPass());
pm.addNestedPass<FuncOp>(mlir::TF::CreateRewriteTPUEmbeddingOpsPass());
pm.addPass(mlir::TFTPU::CreateTPUAnnotateDynamicShapeInputsPass());
pm.addNestedPass<FuncOp>(
mlir::TF::CreateHoistReplicateInvariantResourceWritesPass());
// Verifies clustering has conformed with the expected invariants
pm.addNestedPass<FuncOp>(
tensorflow::tf2xla::internal::CreateVerifyClusteringPass());
}
void NoCanonicalization(OpPassManager& pm) {}
// Same as above but for non-replicated Bridge.
void AddNonReplicatedBridgeClusteringPipelinePasses(OpPassManager& pm) {
// The following ops must be preserved regardless of reachability. Ideally,
// all graphs should have control dependencies to enforce this.
VLOG(2) << "Create TF XLA Bridge pipeline";
pm.addPass(mlir::TFDevice::CreateXlaValidateInputsPass());
pm.addNestedPass<FuncOp>(
mlir::TF::CreateCanonicalizeCompileAndReplicateAttributesPass());
const llvm::SmallVector<std::string, 4> ops_to_preserve = {};
pm.addNestedPass<FuncOp>(
mlir::tf_executor::CreateTFExecutorGraphPruningPass(ops_to_preserve));
// It is assumed at this stage there are no V1 control flow ops as Graph
// functionalization is ran before import. Ops can be lifted out of
// tf_executor dialect islands/graphs.
pm.addNestedPass<FuncOp>(
mlir::CreateExecutorDialectToFunctionalConversionPass());
// Guarantee all functions have one use, which enables more exact shape
// inference.
pm.addPass(mlir::TF::CreateGuaranteeAllFuncsOneUsePass());
pm.addPass(mlir::TF::CreateTFShapeInferencePass());
// Encapsulate PartitionedCall ops within a cluster so that the composite
// resource ops can be decomposed.
pm.addPass(tensorflow::tf2xla::internal::CreateXlaClusterFormationPass());
// Running canonicalizer before decomposing resource ops in cluster helps the
// latter pass to converge faster as it does not have to spend time folding
// away dead ops.
pm.addNestedPass<FuncOp>(mlir::createCanonicalizerPass());
// Decompose resource ops.
pm.addPass(mlir::TFDevice::CreateDecomposeResourceOpsInClusterPass());
// Run another shape inference pass because resource decomposition might have
// created new partial types. Also, after dropping `shape_invariant` attribute
// from While/WhileRegion ops within cluster would lead to more precise
// shapes.
pm.addPass(mlir::TF::CreateTFShapeInferencePass());
pm.addNestedPass<FuncOp>(mlir::createCanonicalizerPass());
// Inline all the function calls. Do not call canonicalizer to prevent it from
// moving the definition of any constant operand of ops within a cluster to
// its outside. This may cause the op to fail to verify after the cluster is
// outlined, as the constant operand is replaced by an argument.
pm.addPass(mlir::createInlinerPass({}, NoCanonicalization));
// Lift resource operations out of device computation. This step needs to be
// done after inlining.
pm.addPass(mlir::TFDevice::CreateResourceOpLiftingPass());
// Outline clusters into cluster functions.
pm.addPass(mlir::TFDevice::CreateClusterOutliningPass());
// Verifies clustering has conformed with the expected invariants
pm.addNestedPass<FuncOp>(
tensorflow::tf2xla::internal::CreateVerifyClusteringPass());
}
}; // namespace internal
}; // namespace tf2xla
}; // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/clustering_bridge_passes.h
+39
View File
@@ -0,0 +1,39 @@
/* 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_CLUSTERING_BRIDGE_PASSES_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_CLUSTERING_BRIDGE_PASSES_H_
#include "absl/base/attributes.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Pass/PassManager.h" // from @llvm-project
namespace tensorflow {
namespace tf2xla {
namespace internal {
// Given the pass manager, add Bridge passes to cluster the replicated input
// graphs.
void AddReplicatedBridgeClusteringPipelinePasses(
mlir::OpPassManager& pm, llvm::StringRef module_name = llvm::StringRef());
// Same as above but for non replicated graphs.
void AddNonReplicatedBridgeClusteringPipelinePasses(mlir::OpPassManager& pm);
}; // namespace internal
}; // namespace tf2xla
}; // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_CLUSTERING_BRIDGE_PASSES_H_
tensorflow/compiler/mlir/tf2xla/internal/clustering_bridge_passes_test.cc
+43
View File
@@ -0,0 +1,43 @@
/* 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/clustering_bridge_passes.h"
#include <gtest/gtest.h>
#include "mlir/Pass/PassManager.h" // from @llvm-project
namespace tensorflow {
namespace tf2xla {
namespace internal {
using mlir::OpPassManager;
TEST(ClusteringBridgePassesTest, AddsBridgePasses) {
OpPassManager pass_manager;
AddReplicatedBridgeClusteringPipelinePasses(pass_manager);
EXPECT_EQ(pass_manager.size(), 45);
}
TEST(ClusteringBridgePassesTest, AddsNonTPUBridgePasses) {
OpPassManager pass_manager;
AddNonReplicatedBridgeClusteringPipelinePasses(pass_manager);
EXPECT_EQ(pass_manager.size(), 15);
}
}; // namespace internal
}; // namespace tf2xla
}; // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/compilation_timer.h
+43
View File
@@ -0,0 +1,43 @@
/* 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_COMPILATION_TIMER_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_COMPILATION_TIMER_H_
#include <chrono> // NOLINT(build/c++11)
#include "tensorflow/core/platform/profile_utils/cpu_utils.h"
// Time the execution of kernels (in CPU cycles). Meant to be used as RAII.
struct CompilationTimer {
uint64_t start_cycles =
tensorflow::profile_utils::CpuUtils::GetCurrentClockCycle();
uint64_t ElapsedCycles() {
return tensorflow::profile_utils::CpuUtils::GetCurrentClockCycle() -
start_cycles;
}
int64_t ElapsedCyclesInMilliseconds() {
std::chrono::duration<double> duration =
tensorflow::profile_utils::CpuUtils::ConvertClockCycleToTime(
ElapsedCycles());
return std::chrono::duration_cast<std::chrono::milliseconds>(duration)
.count();
}
};
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_COMPILATION_TIMER_H_
tensorflow/compiler/mlir/tf2xla/internal/compilation_timer_test.cc
+41
View File
@@ -0,0 +1,41 @@
/* 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/compilation_timer.h"
#include <cstdint>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/time/clock.h"
#include "absl/time/time.h"
namespace {
TEST(CompilationTimer, MeasuresElapsedTime) {
uint64_t timer_result_in_milliseconds;
{
CompilationTimer timer;
absl::SleepFor(absl::Milliseconds(100));
timer_result_in_milliseconds = timer.ElapsedCyclesInMilliseconds();
}
ASSERT_THAT(timer_result_in_milliseconds, testing::Ne(0));
}
} // namespace
tensorflow/compiler/mlir/tf2xla/internal/graph_to_tf_executor_util.cc
+332
View File
@@ -0,0 +1,332 @@
/* 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/graph_to_tf_executor_util.h"
#include <memory>
#include <optional>
#include <string>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "xla/tsl/platform/errors.h"
#include "tensorflow/core/common_runtime/function_body.h"
#include "tensorflow/core/common_runtime/function_def_utils.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/platform/enable_tf2_utils.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/util/device_name_utils.h"
namespace tensorflow {
namespace {
// Internal encapsulation of state for the MLIR bridge graph analyzer. Steps
// through the nodes in the graph and reachable functions, tracking whether
// each feature of interest is found.
//
// Tracks the presence of each feature of interest in the corresponding streamz
// metric. Note that the graph traversal does not terminate early so as to
// capture all of these features.
class MlirBridgeGraphAnalyzer {
public:
explicit MlirBridgeGraphAnalyzer(bool single_core_inference_mode)
: single_core_inference_mode_(single_core_inference_mode) {}
~MlirBridgeGraphAnalyzer() = default;
// Not copyable or movable.
MlirBridgeGraphAnalyzer(const MlirBridgeGraphAnalyzer&) = delete;
MlirBridgeGraphAnalyzer& operator=(const MlirBridgeGraphAnalyzer&) = delete;
// Analyzes whether the graph has features not guaranteed to be supported by
// the MLIR-based TF XLA bridge.
bool HasUnsupportedFeatures(const Graph& graph,
const FunctionLibraryDefinition* function_library,
std::optional<ConfigProto> config_proto,
tensorflow::TF2XLABridgeVersion bridge_version) {
// Non-ok status is considered as "unsupported" since this means something
// is wrong or unexpected with the graph itself.
invalid_graph_ =
invalid_graph_ || !AnalyzeGraphAndReachableFunctions(
graph, function_library, config_proto)
.ok();
// We conservatively consider the graph to be unsupported if it's not
// *known* to be TF2. That is, graphs that have kNotTracked construction
// context are considered unsupported, even though they might in fact be
// TF2 models.
auto construction_context = graph.GetConstructionContextInternal();
bool is_tf2 = construction_context == ConstructionContext::kEagerRuntime;
auto is_tf2_execution_enabled = tensorflow::tf2_execution_enabled();
auto has_unsupported_features = false;
auto is_v1_compat = bridge_version == TF2XLABridgeVersion::kV1Compat;
auto is_nominal_bridge = bridge_version == TF2XLABridgeVersion::kNominal;
auto is_tfrt_bridge = bridge_version == TF2XLABridgeVersion::kTFRTNominal;
is_eager_compliant_ = is_tf2_execution_enabled || is_tf2 ||
is_nominal_bridge || is_tfrt_bridge;
is_eager_compliant_ |= (is_v1_compat && contains_partitioned_call_);
has_unsupported_features = contains_ref_type_ || invalid_graph_;
// For non single core inference mode, checking conditions:
if (!single_core_inference_mode_) {
has_unsupported_features |=
!is_eager_compliant_ || uses_v1_control_flow_ ||
HasTpuReplicatedCoreUnsupportedFeature(is_nominal_bridge,
is_v1_compat, is_tfrt_bridge);
}
PrintGraphUnsupportedFeatures(is_tf2, is_tf2_execution_enabled,
is_v1_compat, is_tfrt_bridge,
is_nominal_bridge, has_unsupported_features);
// Determine whether or not the graph contains unsupported features.
return has_unsupported_features;
}
private:
static constexpr char kPartitionedCall[] = "TPUPartitionedCall";
bool HasTPUReplicatedCoreAttr(const Node& node) {
constexpr absl::string_view kTPUReplicatedCore = "TPU_REPLICATED_CORE";
const std::string& device = node.requested_device();
if (!device.empty()) {
DeviceNameUtils::ParsedName name;
if (DeviceNameUtils::ParseFullName(device, &name)) {
// The TPU_REPLICATED_CORE attrs is not relevant for single TPU core
// inference.
// TODO(b/201091475): this can be generalized to check
// num_cores_per_replica != 1, rather than being special cased for
// single core inference.
if (name.type == kTPUReplicatedCore && !single_core_inference_mode_) {
return true;
}
}
}
return false;
}
bool HasTpuReplicatedCoreUnsupportedFeature(bool is_nominal_bridge,
bool is_v1_compat,
bool is_tfrt_bridge) {
if (!has_tpu_replicated_core_) {
return false;
}
return has_infeed_dequeue_tuple_with_tpu_replicated_core_;
}
void PrintGraphUnsupportedFeatures(bool is_tf2, bool is_tf2_execution_enabled,
bool is_v1_compat, bool is_tfrt_bridge,
bool is_nominal_bridge,
bool has_unsupported_features) {
if (!has_unsupported_features) {
VLOG(1) << "Graph doesn't have unsupported features";
return;
}
LOG(INFO)
<< "Graph has unsupported features: " << (is_tf2 ? "" : "not is_tf2, ")
<< (is_tf2_execution_enabled ? "" : "not tf2_execution, ")
<< (is_nominal_bridge ? "" : "not nominal bridge, ")
<< (is_tfrt_bridge ? "" : "not tfrt bridge, ")
<< (is_v1_compat && contains_partitioned_call_
? "contains partitioned calls at v1 compat bridge call site, "
: "")
<< (contains_ref_type_ ? "contains ref variables, " : "")
<< (invalid_graph_ ? "Invalid graph, " : "")
<< (uses_v1_control_flow_ ? "uses control flow v1 " : "")
<< ((has_tpu_replicated_core_ &&
has_infeed_dequeue_tuple_with_tpu_replicated_core_)
? "InfeedDequeueTuple op with TPU_REPLICATED_CORE attr, "
: "");
}
// Traverses each node in the graph and gathers information about each of the
// features. Specifically, sets the relevant class variable to true when a
// feature is found.
void AnalyzeGraphNodes(const Graph& graph) {
constexpr absl::string_view kIdentityOp = "Identity";
constexpr absl::string_view kIdentityNOp = "IdentityN";
constexpr absl::string_view kCastOp = "Cast";
constexpr absl::string_view kInfeedDequeueTuple = "InfeedDequeueTuple";
constexpr absl::string_view kOutsideCompilationAttr =
"_xla_outside_compilation";
constexpr absl::string_view kAllowSoftPlacementAttr =
"allow_soft_placement";
constexpr absl::string_view kManualControlDepsAttr =
"_has_manual_control_dependencies";
auto has_ref_type = [](const DataTypeVector& types) {
for (const DataType& dtype : types)
if (IsRefType(dtype)) return true;
return false;
};
for (const Node* node : graph.nodes()) {
contains_ref_type_ =
(contains_ref_type_ || has_ref_type(node->input_types()) ||
has_ref_type(node->output_types()));
contains_partitioned_call_ = (contains_partitioned_call_ ||
node->type_string() == kPartitionedCall);
uses_v1_control_flow_ = (uses_v1_control_flow_ || node->IsControlFlow());
uses_outside_compilation_ =
(uses_outside_compilation_ ||
node->attrs().Find(kOutsideCompilationAttr) != nullptr);
has_manual_control_deps_ = (has_manual_control_deps_ ||
node->attrs().Find(kManualControlDepsAttr));
auto soft_placement_attr = node->attrs().Find(kAllowSoftPlacementAttr);
if (soft_placement_attr != nullptr) {
uses_outside_compilation_ =
(uses_outside_compilation_ || soft_placement_attr->b());
}
// TODO(b/187611527): Add support for the ops with explicit device
// assignment on the TPU_REPLICATED_CORE.
if (node->type_string() == kIdentityOp ||
node->type_string() == kCastOp ||
node->type_string() == kIdentityNOp) {
if (HasTPUReplicatedCoreAttr(*node)) {
has_tpu_replicated_core_ = true;
VLOG(2) << node->type_string()
<< " node has TPU_REPLICATED_CORE attribute.";
}
}
if (node->type_string() == kInfeedDequeueTuple &&
HasTPUReplicatedCoreAttr(*node)) {
has_infeed_dequeue_tuple_with_tpu_replicated_core_ = true;
}
}
}
// Analyze all functions from the flib_def if there are any that belong to
// the inference graph.
void AnalyzeInferenceGraphs(const FunctionLibraryDefinition& flib_def) {
if (contains_partitioned_call_) return;
for (const std::string& func_name : flib_def.ListFunctionNames()) {
const FunctionDef* func_def = flib_def.Find(func_name);
for (const NodeDef& node_def : func_def->node_def()) {
contains_partitioned_call_ = node_def.op() == kPartitionedCall;
if (contains_partitioned_call_) return;
}
}
}
// Checks any reachable functions from `graph_def` in `flib_def`
// for unsupported features in the MLIR-based bridge.
//
// Returns failure in the event that the FunctionDef fails to convert to
// FunctionBody. Otherwise returns success.
absl::Status AnalyzeReachableFunctions(
const GraphDef& graph_def, const FunctionLibraryDefinition& flib_def) {
// Check the inputs and outputs of a function for reference variables.
auto signature_contains_ref_type = [](const OpDef& signature) {
for (const auto& args : {signature.input_arg(), signature.output_arg()}) {
for (const auto& arg : args) {
if (IsRefType(arg.type())) return true;
}
}
return false;
};
for (const std::string& func_name :
flib_def.ReachableDefinitions(graph_def).ListFunctionNames()) {
const FunctionDef* func_def = flib_def.Find(func_name);
if (func_def->has_signature()) {
contains_ref_type_ = contains_ref_type_ ||
signature_contains_ref_type(func_def->signature());
}
// Check the function body.
std::unique_ptr<FunctionBody> func_body;
TF_RETURN_IF_ERROR(FunctionDefToBodyHelper(
*func_def, AttrSlice(&func_def->attr()), &flib_def, &func_body));
AnalyzeGraphNodes(*func_body->graph);
}
return absl::OkStatus();
}
// Checks the inputted graph for any features which aren't supported in the
// MLIR-based bridge, stepping through each node in the graph as well as any
// reachable functions (inputs, outputs, and function body).
//
// Note that this analysis does not terminate early because we care about
// collecting all of these metrics.
//
// Returns failure in the event that the FunctionDef fails to convert to
// FunctionBody. Otherwise returns success.
absl::Status AnalyzeGraphAndReachableFunctions(
const Graph& graph, const FunctionLibraryDefinition* function_library,
std::optional<ConfigProto> config_proto) {
// First, check whether soft placement is enabled. This means that auto
// outside compilation may be used.
uses_outside_compilation_ =
uses_outside_compilation_ ||
(config_proto.has_value() && config_proto->allow_soft_placement());
// Analyze each node in this graph.
AnalyzeGraphNodes(graph);
// Then check any associated functions in the graph
// FunctionLibraryDefinition.
GraphDef graph_def;
graph.ToGraphDef(&graph_def);
TF_RETURN_IF_ERROR(AnalyzeReachableFunctions(graph_def, graph.flib_def()));
// Analyze whether there is an inference graph, including non reachable
// from the `graph` itself. This happens when there is a sequence of
// TPUPartitionedCall()->main()->PartitionedCall() and only second part
// of the graph is processed by the MLIR bridge.
AnalyzeInferenceGraphs(graph.flib_def());
// Check any associated function in the graph defined in a separate
// FunctionLibraryDefinition.
if (function_library != nullptr) {
TF_RETURN_IF_ERROR(
AnalyzeReachableFunctions(graph_def, *function_library));
AnalyzeInferenceGraphs(*function_library);
}
return absl::OkStatus();
}
bool contains_partitioned_call_ = false;
bool contains_ref_type_ = false;
bool invalid_graph_ = false;
bool uses_outside_compilation_ = false;
bool uses_v1_control_flow_ = false;
bool has_manual_control_deps_ = false;
bool single_core_inference_mode_ = false;
bool is_eager_compliant_ = false;
bool has_tpu_replicated_core_ = false;
bool has_infeed_dequeue_tuple_with_tpu_replicated_core_ = false;
};
} // namespace
bool GraphHasUnsupportedFeaturesInMlirBridge(
const Graph& graph, const FunctionLibraryDefinition* function_library,
std::optional<ConfigProto> config_proto, TF2XLABridgeVersion bridge_version,
bool single_core_inference_mode) {
return MlirBridgeGraphAnalyzer(single_core_inference_mode)
.HasUnsupportedFeatures(graph, function_library, config_proto,
bridge_version);
}
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/graph_to_tf_executor_util.h
+65
View File
@@ -0,0 +1,65 @@
/* 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_GRAPH_TO_TF_EXECUTOR_UTIL_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_GRAPH_TO_TF_EXECUTOR_UTIL_H_
#include <optional>
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/protobuf/config.pb.h"
namespace tensorflow {
// These are used for grouping the recorded stats appropriately. Specifically,
// we're considering different entrypoints to the bridge as having potentially
// interesting differences at least in the domain of accepted graphs so we want
// to separately track graph features based on these unique entrypoints. One key
// example of this distinction is for TFRT which uses the "nominal" TPU bridge
// pipeline, but may potentially allow graphs with v1 control flow. This
// separate grouping will allow us to dig into these differences granularly.
enum class TF2XLABridgeVersion {
kNominal = 0,
kV1Compat,
kTFRTNominal,
kNotBridgeUseCase,
};
// Analyzes whether the graph has features not guaranteed to be supported by the
// MLIR-based TF XLA bridge for phase 1. If MLIR bridge phase 1 is not used,
// then MLIR bridge phase 2 will not be used. The optional `function_library`
// can be provided if it contains function definitions not including in the
// `graph` FunctionLibraryDefinition.
//
// Conservatively, during the initial rollout, we are not supporting graphs for
// which any of the following are true:
//
// - Not known to be TF2
// - Contains one or more reference variables
// - Contains one or more TPUPartitionedCall ops (which is a proxy for
// inference), but the graph is not v1 compat
// - Uses V1 control flow
// - Graph is invalid or otherwise encounters error during traversal
// If `single_core_inference_mode` is true, we skip some of check conditions
// because they are not applicable.
// TODO(b/241702857): remove single_core_inference_mode
bool GraphHasUnsupportedFeaturesInMlirBridge(
const Graph& graph, const FunctionLibraryDefinition* function_library,
std::optional<ConfigProto> config_proto, TF2XLABridgeVersion bridge_version,
bool single_core_inference_mode);
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_GRAPH_TO_TF_EXECUTOR_UTIL_H_
tensorflow/compiler/mlir/tf2xla/internal/graph_to_tf_executor_util_test.cc
+736
View File
@@ -0,0 +1,736 @@
/* 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/graph_to_tf_executor_util.h"
#include <cstdint>
#include <initializer_list>
#include <string>
#include <vector>
#include <gtest/gtest.h>
#include "absl/log/check.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "tensorflow/cc/framework/ops.h"
#include "tensorflow/cc/framework/scope.h"
#include "tensorflow/cc/ops/array_ops.h"
#include "tensorflow/cc/ops/control_flow_ops.h"
#include "tensorflow/cc/ops/functional_ops.h"
#include "tensorflow/cc/ops/tpu_functional_ops.h"
#include "tensorflow/cc/ops/tpu_replication_ops.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/platform/enable_tf2_utils.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/config.pb.h"
namespace tensorflow {
namespace {
REGISTER_OP("OneRefOutput").Output("y: Ref(float)");
FunctionDef XTimesTwo() {
const Tensor kTwo = test::AsScalar<int64_t>(2);
return FunctionDefHelper::Define(
// Name
"XTimesTwo",
// Args
{"x: T"},
// Return values
{"y: T"},
// Attr def
{"T: {float, double, int32, int64}"},
// Nodes
{
{{"two"}, "Const", {}, {{"value", kTwo}, {"dtype", DT_INT64}}},
{{"scale"}, "Cast", {"two"}, {{"SrcT", DT_INT64}, {"DstT", "$T"}}},
{{"y"}, "Mul", {"x", "scale"}, {{"T", "$T"}}},
});
}
FunctionDef XTimesTwoFloat() {
const Tensor kTwo = test::AsScalar<int64_t>(2);
return FunctionDefHelper::Define(
// Name
"XTimesTwoFloat",
// Args
{"x: float"},
// Return values
{"y: float"},
// Attr def
{},
// Nodes
{
{{"two"}, "Const", {}, {{"value", kTwo}, {"dtype", DT_INT64}}},
{{"scale"},
"Cast",
{"two"},
{{"SrcT", DT_INT64}, {"DstT", DT_FLOAT}}},
{{"y"}, "Mul", {"x", "scale"}, {{"T", DT_FLOAT}}},
});
}
FunctionDef XTimesTwoFloatRef() {
const Tensor kTwo = test::AsScalar<int64_t>(2);
return FunctionDefHelper::Define(
// Name
"XTimesTwoFloatRef",
// Args
{"x: float"},
// Return values
{"y: float"},
// Attr def
{},
// Nodes
{
{{"two"}, "Const", {}, {{"value", kTwo}, {"dtype", DT_INT64_REF}}},
{{"scale"},
"Cast",
{"two"},
{{"SrcT", DT_INT64_REF}, {"DstT", DT_FLOAT}}},
{{"y"}, "Mul", {"x", "scale"}, {{"T", DT_FLOAT}}},
});
}
Node* FromNodeDef(absl::string_view name, absl::string_view node_type,
int num_inputs, DataType dt, Graph& graph) {
auto builder = NodeDefBuilder(name, node_type);
for (int i = 0; i < num_inputs; ++i) {
builder = builder.Input(absl::StrCat("node_", i), i, dt);
}
NodeDef node_def;
CHECK_OK(builder.Finalize(&node_def));
absl::Status s;
Node* node = graph.AddNode(node_def, &s);
CHECK_OK(s);
return node;
}
TEST(SupportedGraphTest, SupportedGraphReturnsFalse) {
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
auto input = tensorflow::ops::Placeholder(root.WithOpName("input"), DT_UINT8);
auto depth = tensorflow::ops::Placeholder(root.WithOpName("depth"), DT_INT32);
auto on = tensorflow::ops::Placeholder(root.WithOpName("on"), DT_UINT8);
auto off = tensorflow::ops::Placeholder(root.WithOpName("off"), DT_UINT8);
tensorflow::set_tf2_execution(true);
(void)tensorflow::ops::OneHot(root.WithOpName("output"), input, depth, on,
off);
Graph graph(OpRegistry::Global());
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/false));
}
TEST(InvalidGraphTest, InvalidFuncBodyReturnsTrue) {
tensorflow::set_tf2_execution(true);
FunctionDefLibrary flib;
*flib.add_function() = XTimesTwo();
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
NameAttrList f_name_attr;
f_name_attr.set_name("XTimesTwo");
ops::PartitionedCall f(root.WithOpName("f"), {x}, {DT_FLOAT}, f_name_attr);
TF_ASSERT_OK(root.ToGraph(&graph));
// The call to XTimesTwo is invalid (missing an attribute), so we expect the
// graph to be unsupported.
EXPECT_TRUE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/false));
}
TEST(RefVarTest, RefVariablesReturnsTrue) {
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
Output cond_a = ops::Placeholder(root.WithOpName("cond_a"), DT_BOOL);
Output cond_b = ops::Placeholder(root.WithOpName("cond_b"), DT_BOOL);
// Output value = ops::Placeholder(root.WithOpName("value"), DT_FLOAT);
tensorflow::set_tf2_execution(true);
const std::vector<int32_t> shape_array{2, 2};
auto shape = TensorShape();
TF_ASSERT_OK(TensorShapeUtils::MakeShape(shape_array, &shape));
Output value = Output(
FromNodeDef("value", "OneRefOutput", 0, DT_FLOAT_REF, *root.graph()));
Graph graph(OpRegistry::Global());
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_TRUE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/false));
}
TEST(RefVarTest, NoRefVariablesCalleeFuncReturnsFalse) {
tensorflow::set_tf2_execution(true);
FunctionDefLibrary flib;
*flib.add_function() = XTimesTwoFloat();
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
NameAttrList f_name_attr;
f_name_attr.set_name("XTimesTwoFloat");
ops::PartitionedCall f(root.WithOpName("f"), {x}, {DT_FLOAT}, f_name_attr);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/false));
}
TEST(RefVarTest, RefVariablesInCalleeFunctionReturnsTrue) {
tensorflow::set_tf2_execution(true);
FunctionDefLibrary flib;
*flib.add_function() = XTimesTwoFloatRef();
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
NameAttrList f_name_attr;
f_name_attr.set_name("XTimesTwoFloatRef");
ops::PartitionedCall f(root.WithOpName("f"), {x}, {DT_FLOAT}, f_name_attr);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_TRUE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/false));
}
TEST(RefVarTest, RefVariablesInExternalCalleeFunctionReturnsTrue) {
tensorflow::set_tf2_execution(true);
Graph graph(OpRegistry::Global());
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
FunctionDefLibrary flib;
*flib.add_function() = XTimesTwoFloatRef();
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
NameAttrList f_name_attr;
f_name_attr.set_name("XTimesTwoFloatRef");
ops::PartitionedCall f(root.WithOpName("f"), {x}, {DT_FLOAT}, f_name_attr);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_TRUE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/&flib_def, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/false));
}
TEST(InferenceTest, ContainsInferenceNodeEagerRuntimeReturnsTrue) {
tensorflow::set_tf2_execution(true);
FunctionDefLibrary flib;
*flib.add_function() = XTimesTwoFloat();
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
NameAttrList f_name_attr;
f_name_attr.set_name("XTimesTwoFloat");
ops::TPUPartitionedCall f(root.WithOpName("f"), {x}, /*device_ordinal=*/0,
{DT_FLOAT}, f_name_attr);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/false));
}
TEST(InferenceTest, ContainsInferenceNodeTFRTBridgeReturnsTrue) {
tensorflow::set_tf2_execution(true);
FunctionDefLibrary flib;
*flib.add_function() = XTimesTwoFloat();
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
NameAttrList f_name_attr;
f_name_attr.set_name("XTimesTwoFloat");
ops::TPUPartitionedCall f(root.WithOpName("f"), {x}, /*device_ordinal=*/0,
{DT_FLOAT}, f_name_attr);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kTFRTNominal,
/*single_core_inference_mode=*/false));
}
TEST(InferenceTest, ContainsInferenceNodeDirectSessionReturnsFalse) {
tensorflow::set_tf2_execution(true);
FunctionDefLibrary flib;
*flib.add_function() = XTimesTwoFloat();
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kDirectSession);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
NameAttrList f_name_attr;
f_name_attr.set_name("XTimesTwoFloat");
ops::TPUPartitionedCall f(root.WithOpName("f"), {x}, /*device_ordinal=*/0,
{DT_FLOAT}, f_name_attr);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kV1Compat,
/*single_core_inference_mode=*/false));
}
TEST(ControlFlowTest, ContainsV1ControlFlowReturnsTrue) {
tensorflow::set_tf2_execution(true);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
Output cond_a = ops::Placeholder(root.WithOpName("cond_a"), DT_BOOL);
Output cond_b = ops::Placeholder(root.WithOpName("cond_b"), DT_BOOL);
Output value = ops::Placeholder(root.WithOpName("value"), DT_FLOAT);
ops::Switch switch_a(root.WithOpName("switch_a"), value, cond_a);
ops::Switch switch_b(root.WithOpName("switch_b"), value, cond_b);
Graph graph(OpRegistry::Global());
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_TRUE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/false));
}
TEST(ControlFlowTest, TFRTContainsV1ControlFlowReturnsTrue) {
tensorflow::set_tf2_execution(true);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
Output cond_a = ops::Placeholder(root.WithOpName("cond_a"), DT_BOOL);
Output cond_b = ops::Placeholder(root.WithOpName("cond_b"), DT_BOOL);
Output value = ops::Placeholder(root.WithOpName("value"), DT_FLOAT);
ops::Switch switch_a(root.WithOpName("switch_a"), value, cond_a);
ops::Switch switch_b(root.WithOpName("switch_b"), value, cond_b);
Graph graph(OpRegistry::Global());
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_TRUE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kTFRTNominal,
/*single_core_inference_mode=*/false));
}
TEST(TFVersionTest, TF1ReturnsTrue) {
tensorflow::set_tf2_execution(false);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
auto input = tensorflow::ops::Placeholder(root.WithOpName("input"), DT_UINT8);
auto depth = tensorflow::ops::Placeholder(root.WithOpName("depth"), DT_INT32);
auto on = tensorflow::ops::Placeholder(root.WithOpName("on"), DT_UINT8);
auto off = tensorflow::ops::Placeholder(root.WithOpName("off"), DT_UINT8);
(void)tensorflow::ops::OneHot(root.WithOpName("output"), input, depth, on,
off);
Graph graph(OpRegistry::Global());
TF_ASSERT_OK(root.ToGraph(&graph));
graph.SetConstructionContext(ConstructionContext::kDirectSession);
EXPECT_TRUE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kV1Compat,
/*single_core_inference_mode=*/false));
}
TEST(TFVersionTest, TF2ExecutionFalseV1CompatBridgeReturnTrue) {
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
auto input = tensorflow::ops::Placeholder(root.WithOpName("input"), DT_UINT8);
auto depth = tensorflow::ops::Placeholder(root.WithOpName("depth"), DT_INT32);
auto on = tensorflow::ops::Placeholder(root.WithOpName("on"), DT_UINT8);
auto off = tensorflow::ops::Placeholder(root.WithOpName("off"), DT_UINT8);
(void)tensorflow::ops::OneHot(root.WithOpName("output"), input, depth, on,
off);
Graph graph(OpRegistry::Global());
TF_ASSERT_OK(root.ToGraph(&graph));
tensorflow::set_tf2_execution(false);
EXPECT_TRUE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kV1Compat,
/*single_core_inference_mode=*/false));
}
TEST(TFVersionTest, TF2ExecutionTrueV1CompatBridgeReturnFalse) {
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
auto input = tensorflow::ops::Placeholder(root.WithOpName("input"), DT_UINT8);
auto depth = tensorflow::ops::Placeholder(root.WithOpName("depth"), DT_INT32);
auto on = tensorflow::ops::Placeholder(root.WithOpName("on"), DT_UINT8);
auto off = tensorflow::ops::Placeholder(root.WithOpName("off"), DT_UINT8);
(void)tensorflow::ops::OneHot(root.WithOpName("output"), input, depth, on,
off);
Graph graph(OpRegistry::Global());
TF_ASSERT_OK(root.ToGraph(&graph));
tensorflow::set_tf2_execution(true);
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kV1Compat,
/*single_core_inference_mode=*/false));
}
TEST(TFVersionTest, TF2ExecutionFalseTfrtNominalBridgeReturnFalse) {
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
auto input = tensorflow::ops::Placeholder(root.WithOpName("input"), DT_UINT8);
auto depth = tensorflow::ops::Placeholder(root.WithOpName("depth"), DT_INT32);
auto on = tensorflow::ops::Placeholder(root.WithOpName("on"), DT_UINT8);
auto off = tensorflow::ops::Placeholder(root.WithOpName("off"), DT_UINT8);
(void)tensorflow::ops::OneHot(root.WithOpName("output"), input, depth, on,
off);
Graph graph(OpRegistry::Global());
TF_ASSERT_OK(root.ToGraph(&graph));
tensorflow::set_tf2_execution(false);
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kTFRTNominal,
/*single_core_inference_mode=*/false));
}
TEST(TFVersionTest, TF2ExecutionTrueTfrtNominalBridgeReturnFalse) {
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
auto input = tensorflow::ops::Placeholder(root.WithOpName("input"), DT_UINT8);
auto depth = tensorflow::ops::Placeholder(root.WithOpName("depth"), DT_INT32);
auto on = tensorflow::ops::Placeholder(root.WithOpName("on"), DT_UINT8);
auto off = tensorflow::ops::Placeholder(root.WithOpName("off"), DT_UINT8);
(void)tensorflow::ops::OneHot(root.WithOpName("output"), input, depth, on,
off);
Graph graph(OpRegistry::Global());
TF_ASSERT_OK(root.ToGraph(&graph));
tensorflow::set_tf2_execution(true);
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kTFRTNominal,
/*single_core_inference_mode=*/false));
}
TEST(TFVersionTest, TF2ExecutionFalseNominalBridgeReturnsFalse) {
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
auto input = tensorflow::ops::Placeholder(root.WithOpName("input"), DT_UINT8);
Graph graph(OpRegistry::Global());
TF_ASSERT_OK(root.ToGraph(&graph));
tensorflow::set_tf2_execution(false);
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/false));
}
TEST(TFVersionTest, TF2ExecutionTrueNominalBridgeReturnsFalse) {
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
auto input = tensorflow::ops::Placeholder(root.WithOpName("input"), DT_UINT8);
Graph graph(OpRegistry::Global());
TF_ASSERT_OK(root.ToGraph(&graph));
tensorflow::set_tf2_execution(true);
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/false));
}
TEST(UnsupportedOpTest,
InfeedDequeueTupleWithTPUReplicatedCoreAttrNotSupported) {
tensorflow::set_tf2_execution(true);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
auto input =
tensorflow::ops::Placeholder(root.WithOpName("node_0"), DT_FLOAT);
auto node = FromNodeDef("Identity", "Identity", 1, DT_FLOAT, *root.graph());
ASSERT_NE(node, nullptr);
node->set_requested_device("/device:TPU_REPLICATED_CORE:0");
// Build InfeedDequeueTuple node with TPU_REPLICATED_CORE Attr
auto builder = NodeDefBuilder("InfeedDequeueTuple", "InfeedDequeueTuple");
builder.Attr("dtypes", DT_FLOAT);
builder.Attr("shapes", 1);
NodeDef node_def;
CHECK_OK(builder.Finalize(&node_def));
absl::Status s;
Node* node_InfeedDequeueTuple = (*root.graph()).AddNode(node_def, &s);
node_InfeedDequeueTuple->set_requested_device(
"/device:TPU_REPLICATED_CORE:0");
CHECK_OK(s);
ASSERT_NE(node_InfeedDequeueTuple, nullptr);
Graph graph(OpRegistry::Global());
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_TRUE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/false));
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/true));
}
TEST(ManualControlDependencyTest,
TPUReplicatedCoreWithManualControlDependencyReturnsFalse) {
tensorflow::set_tf2_execution(true);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
auto input =
tensorflow::ops::Placeholder(root.WithOpName("node_0"), DT_FLOAT);
auto node = FromNodeDef("Identity", "Identity", 1, DT_FLOAT, *root.graph());
ASSERT_NE(node, nullptr);
node->set_requested_device("/device:TPU_REPLICATED_CORE:0");
auto metadata = tensorflow::ops::TPUReplicateMetadata(root, 2);
metadata.operation.node()->AddAttr("_has_manual_control_dependencies", true);
Graph graph(OpRegistry::Global());
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/false));
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/true));
}
TEST(InferenceTest,
ContainsInferenceNodeTPUReplicatedCoreDirectSessionReturnsFalse) {
tensorflow::set_tf2_execution(true);
FunctionDefLibrary flib;
*flib.add_function() = XTimesTwoFloat();
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kDirectSession);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
auto input =
tensorflow::ops::Placeholder(root.WithOpName("node_0"), DT_FLOAT);
auto node = FromNodeDef("Identity", "Identity", 1, DT_FLOAT, *root.graph());
ASSERT_NE(node, nullptr);
node->set_requested_device("/device:TPU_REPLICATED_CORE:0");
Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
NameAttrList f_name_attr;
f_name_attr.set_name("XTimesTwoFloat");
ops::TPUPartitionedCall f(root.WithOpName("f"), {x}, /*device_ordinal=*/0,
{DT_FLOAT}, f_name_attr);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kV1Compat,
/*single_core_inference_mode=*/false));
}
TEST(InferenceTest,
ContainsInferenceNodeTPUReplicatedCoreEagerRuntimeReturnsTrue) {
tensorflow::set_tf2_execution(true);
FunctionDefLibrary flib;
*flib.add_function() = XTimesTwoFloat();
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
auto input =
tensorflow::ops::Placeholder(root.WithOpName("node_0"), DT_FLOAT);
auto node = FromNodeDef("Identity", "Identity", 1, DT_FLOAT, *root.graph());
ASSERT_NE(node, nullptr);
node->set_requested_device("/device:TPU_REPLICATED_CORE:0");
Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
NameAttrList f_name_attr;
f_name_attr.set_name("XTimesTwoFloat");
ops::TPUPartitionedCall f(root.WithOpName("f"), {x}, /*device_ordinal=*/0,
{DT_FLOAT}, f_name_attr);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/false));
}
TEST(InferenceTest, TF2ExecutionFalseV1CompatBridgeReturnFalse) {
tensorflow::set_tf2_execution(false);
FunctionDefLibrary flib;
*flib.add_function() = XTimesTwoFloat();
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kDirectSession);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
auto input =
tensorflow::ops::Placeholder(root.WithOpName("node_0"), DT_FLOAT);
auto node = FromNodeDef("Identity", "Identity", 1, DT_FLOAT, *root.graph());
ASSERT_NE(node, nullptr);
node->set_requested_device("/device:TPU_REPLICATED_CORE:0");
Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
NameAttrList f_name_attr;
f_name_attr.set_name("XTimesTwoFloat");
ops::TPUPartitionedCall f(root.WithOpName("f"), {x}, /*device_ordinal=*/0,
{DT_FLOAT}, f_name_attr);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kV1Compat,
/*single_core_inference_mode=*/false));
}
TEST(InferenceTest, V1CompatBridgeVariableRefReturnTrue) {
tensorflow::set_tf2_execution(false);
FunctionDefLibrary flib;
*flib.add_function() = XTimesTwoFloat();
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kDirectSession);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
auto input =
tensorflow::ops::Placeholder(root.WithOpName("node_0"), DT_FLOAT);
auto node = FromNodeDef("Identity", "Identity", 1, DT_FLOAT, *root.graph());
ASSERT_NE(node, nullptr);
node->set_requested_device("/device:TPU_REPLICATED_CORE:0");
Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
NameAttrList f_name_attr;
f_name_attr.set_name("XTimesTwoFloat");
ops::TPUPartitionedCall f(root.WithOpName("f"), {x}, /*device_ordinal=*/0,
{DT_FLOAT}, f_name_attr);
Output cond_a = ops::Placeholder(root.WithOpName("cond_a"), DT_BOOL);
Output cond_b = ops::Placeholder(root.WithOpName("cond_b"), DT_BOOL);
tensorflow::set_tf2_execution(true);
const std::vector<int32_t> shape_array{2, 2};
auto shape = TensorShape();
TF_ASSERT_OK(TensorShapeUtils::MakeShape(shape_array, &shape));
Output value = Output(
FromNodeDef("value", "OneRefOutput", 0, DT_FLOAT_REF, *root.graph()));
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_TRUE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kV1Compat,
/*single_core_inference_mode=*/false));
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/hlo_post_processing/BUILD
+7
View File
@@ -0,0 +1,7 @@
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow/compiler/mlir/tf2xla/internal:__subpackages__",
],
licenses = ["notice"],
)
tensorflow/compiler/mlir/tf2xla/internal/inference/BUILD
+62
View File
@@ -0,0 +1,62 @@
load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library")
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//visibility:private",
],
)
gentbl_cc_library(
name = "inference_passes_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"inference_passes.h.inc": [
"-gen-pass-decls",
"-name=TF2XLA",
]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "inference_passes.td",
deps = [
"@llvm-project//mlir:PassBaseTdFiles",
],
)
cc_library(
name = "inference_metrics_pass",
srcs = ["inference_metrics_pass.cc"],
hdrs = ["inference_passes.h"],
visibility = ["//visibility:public"],
deps = [
":inference_passes_inc_gen",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/core:lib",
"@com_google_absl//absl/log",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:Dialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:FuncExtensions",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
],
)
tf_cc_test(
name = "inference_metrics_pass_test",
srcs = ["inference_metrics_pass_test.cc"],
deps = [
"inference_metrics_pass",
"//tensorflow/compiler/mlir/tf2xla/transforms:test_utils",
"//tensorflow/core/lib/monitoring:cell_reader",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@xla//xla/tsl/platform:statusor",
],
)
tensorflow/compiler/mlir/tf2xla/internal/inference/inference_metrics_pass.cc
+73
View File
@@ -0,0 +1,73 @@
/* 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 <memory>
#include <string>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/core/lib/monitoring/counter.h"
namespace mlir {
namespace tf2xla {
namespace internal {
auto* has_tpu_partitioned_call_streamz =
tensorflow::monitoring::Counter<1>::New(
"/tensorflow/core/tf2xla/internal/inference/tpu_partitioned_call",
"Whether the model has TPUPartitionedCallOp.",
"has_tpu_partitioned_call");
namespace {
#define GEN_PASS_DEF_INFERENCEMETRICSPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/inference/inference_passes.h.inc"
class InferenceMetricsPass
: public impl::InferenceMetricsPassBase<InferenceMetricsPass> {
public:
void runOnOperation() override;
};
void InferenceMetricsPass::runOnOperation() {
bool has_tpu_partitioned_call = false;
ModuleOp module = getOperation();
for (auto func_op : module.getOps<func::FuncOp>()) {
func_op->walk(
[&](TF::TPUPartitionedCallOp op) { has_tpu_partitioned_call = true; });
if (has_tpu_partitioned_call) break;
}
std::string has_tpu_partitioned_call_str =
has_tpu_partitioned_call ? "true" : "false";
has_tpu_partitioned_call_streamz->GetCell(has_tpu_partitioned_call_str)
->IncrementBy(1);
}
} // namespace
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateInferenceMetricsPass() {
return std::make_unique<InferenceMetricsPass>();
}
} // namespace internal
} // namespace tf2xla
} // namespace mlir
tensorflow/compiler/mlir/tf2xla/internal/inference/inference_metrics_pass_test.cc
+105
View File
@@ -0,0 +1,105 @@
/* 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 <cstdint>
#include <memory>
#include <string>
#include <gtest/gtest.h>
#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/tf2xla/internal/inference/inference_passes.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/test_utils.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/lib/monitoring/cell_reader.h"
namespace mlir {
namespace tf2xla {
namespace internal {
namespace {
using ::mlir::MLIRContext;
using ::mlir::ModuleOp;
using ::mlir::OwningOpRef;
using ::mlir::hlo::test::GetMlirModuleFromString;
using ::tensorflow::monitoring::testing::CellReader;
static constexpr char kHasTpuPartitionedCallStreamzName[] =
"/tensorflow/core/tf2xla/internal/inference/tpu_partitioned_call";
class InferenceMetricsPassTest : public ::testing::Test {
protected:
void CreateModule(const char* module_string) {
TF_ASSERT_OK_AND_ASSIGN(module_,
GetMlirModuleFromString(module_string, &context_));
pm_ = std::make_unique<mlir::PassManager>(&context_);
pm_->addPass(CreateInferenceMetricsPass());
}
mlir::LogicalResult Run() { return pm_->run(module_.get()); }
private:
MLIRContext context_;
OwningOpRef<ModuleOp> module_;
std::unique_ptr<mlir::PassManager> pm_;
};
TEST_F(InferenceMetricsPassTest, RecordsTrueForTPUPartitionedCallOp) {
static constexpr char kMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @tpu_partitioned_call_func(%arg0: tensor<?xi32>) -> (tensor<?xi32>) {
func.return %arg0 : tensor<?xi32>
}
func.func @main(%arg0: tensor<20xi32>, %arg1: tensor<?xi32>) -> tensor<*xi32> {
%2 = "tf.TPUPartitionedCall"(%arg0, %arg1) {f = @tpu_partitioned_call_func} : (tensor<20xi32>, tensor<?xi32>) -> tensor<*xi32>
func.return %2 : tensor<*xi32>
}
})";
CellReader<int64_t> error(kHasTpuPartitionedCallStreamzName);
CreateModule(kMlirModuleStr);
auto result = Run();
EXPECT_TRUE(result.succeeded());
EXPECT_EQ(error.Delta("true"), 1);
EXPECT_EQ(error.Delta("false"), 0);
}
TEST_F(InferenceMetricsPassTest, RecordsFalseForNonTPUPartitionedCallOp) {
static constexpr char kMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> tensor<1xi32> {
%0 = "tf.BadValue"() {value = dense<1000> : tensor<1xi32>} : () -> tensor<1xi32>
func.return %0 : tensor<1xi32>
}
})";
CellReader<int64_t> error(kHasTpuPartitionedCallStreamzName);
CreateModule(kMlirModuleStr);
auto result = Run();
EXPECT_TRUE(result.succeeded());
EXPECT_EQ(error.Delta("false"), 1);
EXPECT_EQ(error.Delta("true"), 0);
}
} // namespace
} // namespace internal
} // namespace tf2xla
} // namespace mlir
tensorflow/compiler/mlir/tf2xla/internal/inference/inference_passes.h
+39
View File
@@ -0,0 +1,39 @@
/* 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_INFERENCE_INFERENCE_PASSES_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_INFERENCE_INFERENCE_PASSES_H_
#include <memory>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
namespace mlir {
namespace tf2xla {
namespace internal {
std::unique_ptr<OperationPass<ModuleOp>> CreateInferenceMetricsPass();
#define GEN_PASS_REGISTRATION
#define GEN_PASS_DECL_INFERENCEMETRICSPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/inference/inference_passes.h.inc"
} // namespace internal
} // namespace tf2xla
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_INFERENCE_INFERENCE_PASSES_H_
tensorflow/compiler/mlir/tf2xla/internal/inference/inference_passes.td
+25
View File
@@ -0,0 +1,25 @@
/* 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 "mlir/Pass/PassBase.td"
// Internal TF2XLA Bridge Passes used during interence only.
def InferenceMetricsPass : Pass<"tf2xla-inference-metrics-pass", "ModuleOp"> {
let summary = "A pass to go over a Module and collect various metrics";
let description = [{
Collects metrics about a Module during inference, such as which ops are used.
}];
let constructor = "internal::CreateInferenceMetricsPass()";
}
tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_mlir.cc
+98
View File
@@ -0,0 +1,98 @@
/* 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/legalize_tf_mlir.h"
#include <memory>
#include <vector>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "stablehlo/dialect/Register.h" // from @stablehlo
#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/set_tpu_infeed_layout.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/serialize_mlir_module_utils.h"
#include "tensorflow/compiler/mlir/tf2xla/api/v1/compile_mlir_util.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/compilation_timer.h"
#include "tensorflow/compiler/tf2xla/layout_util.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/mlir_hlo/mhlo/IR/register.h"
#include "xla/shape.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/profile_utils/cpu_utils.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
#include "tensorflow/core/tpu/tpu_compile.h"
#include "tsl/platform/error_logging.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
// Name of component for error logging. This name is fixed and required to
// enable logging.
constexpr char kBridgeComponent[] = "TFXLABridge";
using tpu::ShardingAndIndex;
absl::Status CompileFromMlirToXlaHlo(
bool lower_to_xla_hlo, mlir::ModuleOp mlir_module_op,
const tpu::TPUCompileMetadataProto& metadata, llvm::StringRef device_type,
const XlaShapeLayoutHelpers::ShapeDeterminationFns& shape_determination_fns,
bool use_tuple_args, XlaCompiler::CompilationResult* compilation_result,
std::vector<std::unique_ptr<mlir::Pass>>& custom_legalization_passes,
const std::vector<TensorShape>& arg_shapes,
std::vector<ShardingAndIndex>* arg_core_mapping,
std::vector<std::vector<xla::Shape>>* per_core_arg_shapes) {
LOG_FIRST_N(INFO, 1)
<< "Compiling MLIR computation to XLA HLO using MLIR tf2xla bridge in "
"the op by op fallback mode. This is Phase 2 of the TF2XLA Bridge. "
"Old (non-MLIR) bridge may be used in case of unsupported feature "
"or compilation failure from the MLIR bridge (full fallback mode).";
llvm::SmallVector<TensorOrResourceShape, 4> tensor_or_resource_shapes;
tensor_or_resource_shapes.reserve(arg_shapes.size());
for (const auto& arg_shape : arg_shapes)
tensor_or_resource_shapes.push_back({arg_shape});
TF_RETURN_IF_ERROR(CompileMlirToXlaHlo(
mlir_module_op, tensor_or_resource_shapes, device_type, use_tuple_args,
/*enable_op_fallback=*/true, /*use_return_tuple=*/true,
/*use_resource_updates_for_aliases=*/false, shape_determination_fns,
compilation_result, custom_legalization_passes, metadata.module_name(),
lower_to_xla_hlo));
// Compute how arguments are shared across different cores.
auto sharding_result =
tpu::GetShardingInfo(metadata, arg_shapes, shape_determination_fns,
arg_core_mapping, per_core_arg_shapes);
if (!sharding_result.ok()) {
return sharding_result;
}
return absl::OkStatus();
}
}; // namespace internal
}; // namespace tf2xla
}; // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_mlir.h
+51
View File
@@ -0,0 +1,51 @@
/* 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_LEGALIZE_TF_MLIR_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_LEGALIZE_TF_MLIR_H_
#include <string>
#include <vector>
#include "absl/status/status.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
// Runs all the MLIR Bridge passes on the given MLIR module.
// If compile_to_xla_hlo is true then those passes include all the Legalization
// to XLA HLO which is returned in the compilation_result.
absl::Status CompileFromMlirToXlaHlo(
bool lower_to_xla_hlo, mlir::ModuleOp mlir_module_op,
const tpu::TPUCompileMetadataProto& metadata, llvm::StringRef device_type,
const XlaShapeLayoutHelpers::ShapeDeterminationFns& shape_determination_fns,
bool use_tuple_args, XlaCompiler::CompilationResult* compilation_result,
std::vector<std::unique_ptr<mlir::Pass>>& custom_legalization_passes,
const std::vector<TensorShape>& arg_shapes,
std::vector<tpu::ShardingAndIndex>* arg_core_mapping,
std::vector<std::vector<xla::Shape>>* per_core_arg_shapes);
}; // namespace internal
}; // namespace tf2xla
}; // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_LEGALIZE_TF_MLIR_H_
tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_mlir_test.cc
+96
View File
@@ -0,0 +1,96 @@
/* 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/legalize_tf_mlir.h"
#include <memory>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/statusor.h"
#include "mlir/IR/BuiltinOps.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 "mlir/Pass/Pass.h" // from @llvm-project
#include "stablehlo/dialect/Register.h" // from @stablehlo
#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/deserialize_mlir_module_utils.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/test_matchers.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.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"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {
using tpu::ShardingAndIndex;
using tpu::TPUCompileMetadataProto;
static constexpr char kMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> tensor<1xi32> {
%0 = "tf.Const"() {value = dense<1000> : tensor<1xi32>} : () -> tensor<1xi32>
func.return %0 : tensor<1xi32>
}
})";
absl::StatusOr<XlaCompilationResult> CompileMlirModule(bool compile_to_xla_hlo,
const char* module_str) {
XlaCompilationResult compilation_result;
mlir::DialectRegistry registry;
mlir::RegisterAllTensorFlowDialects(registry);
mlir::mhlo::registerAllMhloDialects(registry);
mlir::stablehlo::registerAllDialects(registry);
mlir::MLIRContext context(registry);
mlir::OwningOpRef<mlir::ModuleOp> mlir_module;
TF_RETURN_IF_ERROR(
tensorflow::DeserializeMlirModule(module_str, &context, &mlir_module));
std::vector<TensorShape> arg_shapes;
TPUCompileMetadataProto metadata_proto;
bool use_tuple_args = true;
std::vector<ShardingAndIndex> arg_core_mapping;
std::vector<std::vector<xla::Shape>> per_core_arg_shapes;
std::vector<std::unique_ptr<mlir::Pass>> custom_legalization_passes;
TF_RETURN_IF_ERROR(CompileFromMlirToXlaHlo(
compile_to_xla_hlo, mlir_module.get(), metadata_proto,
/*device_type=*/"XLA_TPU_JIT",
/*shape_determination_fns=*/{}, use_tuple_args, &compilation_result,
custom_legalization_passes, arg_shapes, &arg_core_mapping,
&per_core_arg_shapes));
return compilation_result;
}
TEST(CompileFromMlir, ReturnsLoweredModule) {
auto result = CompileMlirModule(/*compile_to_xla_hlo=*/true, kMlirModuleStr);
ASSERT_THAT(result, IsOkOrFiltered());
EXPECT_THAT(result, ComputationProtoContains("opcode: \"constant\""));
}
} // namespace
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo.cc
+151
View File
@@ -0,0 +1,151 @@
/* 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/legalize_tf_to_hlo.h"
#include <memory>
#include <vector>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/IR/BuiltinOps.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 "mlir/Interfaces/DataLayoutInterfaces.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "stablehlo/dialect/Register.h" // from @stablehlo
#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/set_tpu_infeed_layout.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/deserialize_mlir_module_utils.h"
#include "tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_mlir.h"
#include "tensorflow/compiler/tf2xla/layout_util.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/client/compile_only_client.h"
#include "xla/mlir_hlo/mhlo/IR/register.h"
#include "xla/shape.h"
#include "xla/tsl/framework/device_type.h"
#include "xla/tsl/platform/errors.h"
#include "tensorflow/core/framework/metrics.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {
using metrics::IncrementTfMlirBridgeSecondPhaseCounter;
using metrics::MlirBridgeSecondPhaseMetric;
using tpu::MlirToHloArgs;
absl::Status CheckAndIncrementCounter(absl::Status status,
MlirBridgeSecondPhaseMetric metric) {
if (!status.ok()) {
IncrementTfMlirBridgeSecondPhaseCounter(metric);
return status;
}
return absl::OkStatus();
}
absl::StatusOr<XlaCompilationResult> LegalizeTfToHlo(
mlir::ModuleOp mlir_module_op, const tpu::TPUCompileMetadataProto& metadata,
bool use_tuple_args, llvm::StringRef device_type,
XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,
const std::vector<tensorflow::TensorShape>& arg_shapes,
std::vector<tpu::ShardingAndIndex>* arg_core_mapping,
std::vector<std::vector<xla::Shape>>* per_core_arg_shapes,
std::vector<std::unique_ptr<mlir::Pass>>& custom_legalization_passes,
xla::CompileOnlyClient* client, XlaCompilationResult* compilation_result) {
if (!mlir::SetTPUInfeedLayout(mlir_module_op))
return absl::AbortedError("Failed to set layouts attribute");
absl::Status mlir_compilation = internal::CompileFromMlirToXlaHlo(
/*lower_to_xla_hlo=*/false, mlir_module_op, metadata, device_type,
shape_determination_fns, use_tuple_args, compilation_result,
custom_legalization_passes, arg_shapes, arg_core_mapping,
per_core_arg_shapes);
TF_RETURN_IF_ERROR(CheckAndIncrementCounter(
mlir_compilation, MlirBridgeSecondPhaseMetric::kMlirCombinedMlirFailure));
IncrementTfMlirBridgeSecondPhaseCounter(
MlirBridgeSecondPhaseMetric::kMlirCombinedMlirSuccess);
MlirToHloArgs mlir_to_hlo_args;
mlir_to_hlo_args.mlir_module_op = mlir_module_op;
absl::Status old_bridge_status = v1::CompileTensorflowGraphToHlo(
mlir_to_hlo_args, metadata, use_tuple_args, shape_determination_fns,
arg_shapes, tsl::DeviceType(device_type.str()), arg_core_mapping,
per_core_arg_shapes, client, compilation_result);
TF_RETURN_IF_ERROR(CheckAndIncrementCounter(
old_bridge_status, MlirBridgeSecondPhaseMetric::kMlirCombinedOldFailure));
IncrementTfMlirBridgeSecondPhaseCounter(
MlirBridgeSecondPhaseMetric::kMlirCombinedOldSuccess);
return *compilation_result;
}
} // namespace
absl::StatusOr<XlaCompilationResult> LegalizeTfToHlo(
const tpu::MlirToHloArgs& computation,
const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,
llvm::StringRef device_type,
XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,
const std::vector<tensorflow::TensorShape>& arg_shapes,
std::vector<tpu::ShardingAndIndex>* arg_core_mapping,
std::vector<std::vector<xla::Shape>>* per_core_arg_shapes,
std::vector<std::unique_ptr<mlir::Pass>>& custom_legalization_passes,
xla::CompileOnlyClient* client, XlaCompilationResult* compilation_result) {
LOG_FIRST_N(INFO, 1) << "Compiling MLIR computation to XLA HLO using the "
"Combined MLIR Tf2Xla Bridge.";
if (computation.mlir_module_op.has_value()) {
return LegalizeTfToHlo(computation.mlir_module_op.value(), metadata,
use_tuple_args, device_type, shape_determination_fns,
arg_shapes, arg_core_mapping, per_core_arg_shapes,
custom_legalization_passes, client,
compilation_result);
}
mlir::DialectRegistry registry;
mlir::RegisterAllTensorFlowDialects(registry);
mlir::mhlo::registerAllMhloDialects(registry);
mlir::stablehlo::registerAllDialects(registry);
mlir::MLIRContext context(registry);
mlir::OwningOpRef<mlir::ModuleOp> mlir_module;
absl::Status deserialization_status = tensorflow::DeserializeMlirModule(
computation.mlir_module, &context, &mlir_module);
TF_RETURN_IF_ERROR(CheckAndIncrementCounter(
deserialization_status,
MlirBridgeSecondPhaseMetric::kMlirCombinedMlirFailure));
return LegalizeTfToHlo(mlir_module.get(), metadata, use_tuple_args,
device_type, shape_determination_fns, arg_shapes,
arg_core_mapping, per_core_arg_shapes,
custom_legalization_passes, client,
compilation_result);
}
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo.h
+49
View File
@@ -0,0 +1,49 @@
/* 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_LEGALIZE_TF_TO_HLO_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_LEGALIZE_TF_TO_HLO_H_
#include "absl/status/statusor.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/client/compile_only_client.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
// Legalize the given MLIR module to XLA HLO using a combination of the MLIR
// Bridge and XlaBuilder
absl::StatusOr<XlaCompilationResult> LegalizeTfToHlo(
const tpu::MlirToHloArgs& computation,
const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,
llvm::StringRef device_type,
XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,
const std::vector<tensorflow::TensorShape>& arg_shapes,
std::vector<tpu::ShardingAndIndex>* arg_core_mapping,
std::vector<std::vector<xla::Shape>>* per_core_arg_shapes,
std::vector<std::unique_ptr<mlir::Pass>>& custom_legalization_passes,
xla::CompileOnlyClient* client, XlaCompilationResult* compilation_result);
}; // namespace internal
}; // namespace tf2xla
}; // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_LEGALIZE_TF_TO_HLO_H_
tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo_test.cc
+186
View File
@@ -0,0 +1,186 @@
/* 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/legalize_tf_to_hlo.h"
#include <cstdint>
#include <memory>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/statusor.h"
#include "mlir/IR/BuiltinOps.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 "mlir/Pass/Pass.h" // from @llvm-project
#include "stablehlo/dialect/Register.h" // from @stablehlo
#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/deserialize_mlir_module_utils.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/test_matchers.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/client/client_library.h"
#include "xla/mlir_hlo/mhlo/IR/register.h"
#include "xla/shape.h"
#include "xla/stream_executor/platform.h"
#include "xla/stream_executor/platform_manager.h"
#include "xla/tsl/platform/errors.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/lib/monitoring/cell_reader.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
using ::tensorflow::monitoring::testing::CellReader;
using tpu::MlirToHloArgs;
using tpu::ShardingAndIndex;
using tpu::TPUCompileMetadataProto;
static constexpr char kMlirLegalizeCount[] =
"/tensorflow/core/tf2xla/v1/mlir_failed_xla_legalize_tf_count";
static constexpr char kMlirLegalizeErrors[] =
"/tensorflow/core/tf2xla/v1/mlir_failed_xla_legalize_tf_pass_count";
static constexpr char kBridgeStatusCounter[] =
"/tensorflow/core/tf2xla/api/v2/phase2_compilation_status";
constexpr char kMlirCombinedMlirSuccess[] = "kMlirCombinedMlirSuccess";
constexpr char kMlirCombinedOldSuccess[] = "kMlirCombinedOldSuccess";
constexpr char kMlirCombinedOldFailure[] = "kMlirCombinedOldFailure";
static constexpr char kMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main(%arg0 : tensor<1xf32>) -> tensor<1xf32> {
%0 = "tf.Acos"(%arg0) : (tensor<1xf32>) -> tensor<1xf32>
func.return %0 : tensor<1xf32>
}
})";
static constexpr char kBadMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> tensor<1xi32> {
%0 = "tf.DoesntExist"() {value = dense<1000> : tensor<1xi32>} : () -> tensor<1xi32>
func.return %0 : tensor<1xi32>
}
})";
absl::StatusOr<XlaCompiler::CompilationResult> CompileMlirModule(
const char* module_str, bool compile_serialized = true) {
MlirToHloArgs mlir_to_hlo_args;
mlir_to_hlo_args.rollout_state =
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED;
mlir::DialectRegistry registry;
mlir::RegisterAllTensorFlowDialects(registry);
mlir::mhlo::registerAllMhloDialects(registry);
mlir::stablehlo::registerAllDialects(registry);
mlir::MLIRContext context(registry);
mlir::OwningOpRef<mlir::ModuleOp> mlir_module;
if (compile_serialized) {
mlir_to_hlo_args.mlir_module = module_str;
} else {
TF_RETURN_IF_ERROR(
tensorflow::DeserializeMlirModule(module_str, &context, &mlir_module));
mlir_to_hlo_args.mlir_module_op = mlir_module.get();
}
se::Platform* platform =
se::PlatformManager::PlatformWithName("Host").value();
auto client =
xla::ClientLibrary::GetOrCreateCompileOnlyClient(platform).value();
std::vector<TensorShape> arg_shapes = {{1}};
TPUCompileMetadataProto metadata_proto;
auto arg = metadata_proto.add_args();
arg->set_dtype(DataType::DT_FLOAT);
arg->set_kind(TPUCompileMetadataProto::Arg::PARAMETER);
metadata_proto.add_retvals();
bool use_tuple_args = true;
std::vector<ShardingAndIndex> arg_core_mapping;
std::vector<std::vector<xla::Shape>> per_core_arg_shapes;
std::vector<std::unique_ptr<mlir::Pass>> custom_legalization_passes;
auto compilation_result = std::make_unique<XlaCompilationResult>();
return LegalizeTfToHlo(mlir_to_hlo_args, metadata_proto, use_tuple_args,
/*device_type=*/"XLA_TPU_JIT",
/*shape_determination_fns=*/{}, arg_shapes,
&arg_core_mapping, &per_core_arg_shapes,
custom_legalization_passes, client,
compilation_result.get());
}
using LegalizeWithCombinedBridge = ::testing::TestWithParam<bool>;
TEST_P(LegalizeWithCombinedBridge, DoesNotUseMlirLowering) {
const bool compile_serialized = GetParam();
CellReader<int64_t> mlir_bridge_legalize_count(kMlirLegalizeCount);
CellReader<int64_t> counts(kBridgeStatusCounter);
auto result = CompileMlirModule(kMlirModuleStr, compile_serialized);
ASSERT_THAT(result, IsOkOrFiltered());
EXPECT_EQ(mlir_bridge_legalize_count.Delta("tf.Acos"), 0);
EXPECT_THAT(result,
IncrementedOrFiltered(counts.Delta(kMlirCombinedMlirSuccess), 1));
EXPECT_THAT(result,
IncrementedOrFiltered(counts.Delta(kMlirCombinedOldSuccess), 1));
}
TEST_P(LegalizeWithCombinedBridge,
CorrectlyCountsMlirBridgePassingAndGraphBridgeFailing) {
const bool compile_serialized = GetParam();
CellReader<int64_t> legalize_failure_count(kMlirLegalizeErrors);
CellReader<int64_t> counts(kBridgeStatusCounter);
auto result = CompileMlirModule(kBadMlirModuleStr, compile_serialized);
ASSERT_FALSE(result.ok());
// Never failed to legalize because it was never attempted
EXPECT_EQ(legalize_failure_count.Read("tf.DoesntExist", "Unknown"), 0);
EXPECT_THAT(result,
IncrementedOrFiltered(counts.Delta(kMlirCombinedMlirSuccess), 1));
EXPECT_THAT(result,
IncrementedOrFiltered(counts.Delta(kMlirCombinedOldFailure), 1));
}
TEST_P(LegalizeWithCombinedBridge, RecordsDynamicOps) {
const bool compile_serialized = GetParam();
static constexpr char kDynamismFunctionCounterStreamzName[] =
"/tensorflow/core/tf2xla/api/v2/dynamism_function_counter";
constexpr char kNotDynamicFunctionName[] = "kNotDynamicFunction";
CellReader<int64_t> dynamic_function_op_count(
kDynamismFunctionCounterStreamzName);
auto result = CompileMlirModule(kMlirModuleStr, compile_serialized);
ASSERT_TRUE(result.ok());
EXPECT_EQ(dynamic_function_op_count.Delta(kNotDynamicFunctionName), 1);
}
INSTANTIATE_TEST_SUITE_P(LegalizeWithCombinedBridge, LegalizeWithCombinedBridge,
::testing::Bool());
}; // namespace internal
}; // namespace tf2xla
}; // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/logging_hooks.cc
+52
View File
@@ -0,0 +1,52 @@
/* 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/logging_hooks.h"
#include <memory>
#include <string>
#include "llvm/ADT/StringRef.h"
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/utils/data_dumper_logger_config.h"
#include "tensorflow/core/util/debug_data_dumper.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
using mlir::PassManager;
// Setup the input pass manager to enable IR dumping after each pass.
// Note a side effect of this method is that multi threading will be disabled.
void EnablePassIRPrinting(PassManager& pm, const std::string& dump_group_name,
llvm::StringRef module_name) {
// Print the whole module after each pass, which requires disabling
// multi-threading as well.
pm.getContext()->disableMultithreading();
pm.enableIRPrinting(std::make_unique<::tensorflow::DataDumperLoggerConfig>(
[module_name, dump_group_name](const std::string& pass_tag_name,
mlir::Operation* op) {
return DEBUG_DATA_DUMPER()->GetDumpFilename(
module_name.str(), dump_group_name, pass_tag_name);
},
/*pass_prefix=*/"",
/*print_module_scope=*/true));
pm.enableTiming();
}
}; // namespace internal
}; // namespace tf2xla
}; // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/logging_hooks.h
+38
View File
@@ -0,0 +1,38 @@
/* 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_LOGGING_HOOKS_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_LOGGING_HOOKS_H_
#include <string>
#include "llvm/ADT/StringRef.h"
#include "mlir/Pass/PassManager.h" // from @llvm-project
namespace tensorflow {
namespace tf2xla {
namespace internal {
// Setup the input pass manager to enable IR dumping after each pass.
// Note a side effect of this method is that multi threading will be disabled.
void EnablePassIRPrinting(mlir::PassManager& pm,
const std::string& dump_group_name,
llvm::StringRef module_name = llvm::StringRef());
}; // namespace internal
}; // namespace tf2xla
}; // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_LOGGING_HOOKS_H_
tensorflow/compiler/mlir/tf2xla/internal/logging_hooks_test.cc
+112
View File
@@ -0,0 +1,112 @@
/* 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/logging_hooks.h"
#include <string>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/Parser/Parser.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "mlir/Transforms/Passes.h" // from @llvm-project
#include "tensorflow/compiler/mlir/register_common_dialects.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/platform/file_statistics.h"
#include "xla/tsl/platform/status.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/resource_loader.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {
using mlir::DialectRegistry;
using mlir::LogicalResult;
using mlir::MLIRContext;
using mlir::ModuleOp;
using mlir::OwningOpRef;
using mlir::PassManager;
using mlir::func::FuncOp;
std::string TestDataPath() {
return tensorflow::GetDataDependencyFilepath(
"tensorflow/compiler/mlir/tf2xla/internal/testdata/");
}
class LoggingHooksTest : public ::testing::Test {
public:
LoggingHooksTest() {
mlir::RegisterCommonToolingDialects(registry_);
context_.appendDialectRegistry(registry_);
context_.loadAllAvailableDialects();
env_ = Env::Default();
test_group_name_ = "TestGroup";
test_dir_ = testing::TmpDir();
setenv(/*name=*/"TF_DUMP_GRAPH_PREFIX", /*value=*/test_dir_.c_str(),
/*overwrite=*/1);
}
absl::Status CreateMlirModule(std::string mlir_module_filename) {
std::string mlir_module_path = TestDataPath() + mlir_module_filename;
mlir_module_ =
mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context_);
if (!mlir_module_) {
return absl::Status(
absl::StatusCode::kNotFound,
absl::StrCat("Could not find MLIR module at ", mlir_module_path));
}
return absl::OkStatus();
}
DialectRegistry registry_;
MLIRContext context_;
OwningOpRef<mlir::ModuleOp> mlir_module_;
Env* env_;
std::string test_dir_;
std::string test_group_name_;
};
TEST_F(LoggingHooksTest, DumpsPassData) {
std::vector<std::string> files;
TF_ASSERT_OK(env_->GetChildren(test_dir_, &files));
EXPECT_THAT(files, ::testing::IsEmpty());
TF_ASSERT_OK(CreateMlirModule("dead_const.mlir"));
PassManager pass_manager(&context_);
pass_manager.addNestedPass<FuncOp>(mlir::createCanonicalizerPass());
EnablePassIRPrinting(pass_manager, test_group_name_);
LogicalResult pass_status = pass_manager.run(mlir_module_.get());
EXPECT_TRUE(pass_status.succeeded());
TF_ASSERT_OK(env_->GetChildren(test_dir_, &files));
EXPECT_THAT(files, ::testing::SizeIs(2));
}
}; // namespace
}; // namespace internal
}; // namespace tf2xla
}; // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/mlir_bridge_pass_util.cc
+246
View File
@@ -0,0 +1,246 @@
/* 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/mlir_bridge_pass_util.h"
#include <functional>
#include <memory>
#include <string>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/tf2xla/tf2xla_defs.h"
#include "xla/tsl/platform/status.h"
#include "tensorflow/core/common_runtime/function_body.h"
#include "tensorflow/core/common_runtime/function_def_utils.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/graph/graph.h"
namespace tensorflow {
using ::mlir::failure;
using ::mlir::LogicalResult;
using ::mlir::success;
namespace {
constexpr absl::string_view kPartitionedCall = "TPUPartitionedCall";
LogicalResult HasAttr(
const Graph& graph, const FunctionLibraryDefinition* function_library,
const std::function<bool(const Graph& graph)>& predicate) {
if (predicate(graph)) {
return success();
}
// Check if any reachable functions from the graph has the target attribute.
GraphDef graph_def;
graph.ToGraphDef(&graph_def);
if (!function_library) return failure();
for (const std::string& func_name :
function_library->ReachableDefinitions(graph_def).ListFunctionNames()) {
const FunctionDef* func_def = function_library->Find(func_name);
std::unique_ptr<FunctionBody> func_body;
absl::Status status = FunctionDefToBodyHelper(
*func_def, AttrSlice(&func_def->attr()), function_library, &func_body);
// This is not expected to happen in practice
if (!status.ok()) {
LOG(ERROR) << "Failed to parse " << func_name << ": "
<< absl::StatusMessageAsCStr(status) << "; this can "
<< "happen if the function contains unregistered ops.";
return failure();
}
if (predicate(*func_body->graph)) {
return success();
}
}
return failure();
}
// Check if the `graph` has parameter server jobs and resource variable
// arguments that are on parameter servers
bool HasPsWithResourceVariable(const Graph& graph) {
// Check parameter serverjobs and resource variable arguments that are
// on parameter servers.
const std::string jobType = "ps";
const std::string nodeType = "_Arg";
const std::string attrKey = "T";
for (const Node* node : graph.nodes()) {
if (node->type_string() == nodeType) {
auto device_name = node->assigned_device_name();
DeviceNameUtils::ParsedName device;
if (DeviceNameUtils::ParseFullName(device_name, &device) &&
device.has_job && device.job == jobType) {
for (const auto& attr : node->attrs()) {
auto attr_key = attr.first;
auto attr_value = attr.second;
if (attr_key == attrKey &&
attr_value.value_case() == AttrValue::kType &&
attr_value.type() == DT_RESOURCE) {
return true;
break;
}
}
}
}
}
return false;
}
bool IsNonReplicatedGraph(const Graph& graph,
const FunctionLibraryDefinition* function_library) {
auto predicate = [](const Graph& graph) {
const std::string kStatefulPartitionedCallOp = "StatefulPartitionedCall";
for (const Node* node : graph.nodes()) {
auto node_op = node->type_string();
if (node_op == kStatefulPartitionedCallOp) {
// Functions called by StatefulfulPartitionedCall ops with
// _XlaMustCompile=true are compiled by XLA.
auto attr = node->attrs().FindByString(std::string(kMustCompileAttr));
if (attr != nullptr && attr->b() == true) {
return true;
}
}
}
return false;
};
return HasAttr(graph, function_library, predicate).succeeded();
}
bool IsReplicatedGraph(const Graph& graph,
const FunctionLibraryDefinition* function_library) {
auto predicate = [](const Graph& graph) {
for (const Node* node : graph.nodes()) {
// _tpu_replicate is used in replicated TPU graphs. It will be converted
// to_replication_info and _xla_compile_device_type in phase 1 pipelines.
if (node->attrs().FindByString(std::string(kTpuReplicateAttr))) {
return true;
}
}
return false;
};
return HasAttr(graph, function_library, predicate).succeeded();
}
bool IsReplicatedGraph(mlir::ModuleOp module) {
auto walk_result = module.walk([&](mlir::Operation* op) {
// TODO(b/223677572): Once the scope for new compilation and replication
// markers is expanded beyond bridge we can remove this check for
// `kTPUReplicateAttr`, we will then always have a `kCompileDeviceTypeAttr`
// in such cases (see above).
// TODO(b/229028654): Remove string conversion once we have C++17.
const llvm::StringRef tpu_replicate_attr_name(kTpuReplicateAttr.data(),
kTpuReplicateAttr.size());
auto replicate_attr =
op->getAttrOfType<mlir::StringAttr>(tpu_replicate_attr_name);
if (replicate_attr) return mlir::WalkResult::interrupt();
return mlir::WalkResult::advance();
});
return walk_result.wasInterrupted();
}
// Traverses each node in the graph and check if any of them is
// TPUPartitionedCall. If so, return true. Otherwise, return false.
bool DoesGraphContainTPUPartitionedCall(const Graph& graph) {
for (const Node* node : graph.nodes()) {
if (node->type_string() == kPartitionedCall) return true;
}
return false;
}
// Checks any reachable functions from `graph_def` in `flib_def`
// for inference graphs.
bool DoReachableFuncsContainTPUPartitionedCall(
const GraphDef& graph_def, const FunctionLibraryDefinition& flib_def) {
for (const std::string& func_name :
flib_def.ReachableDefinitions(graph_def).ListFunctionNames()) {
const FunctionDef* func_def = flib_def.Find(func_name);
std::unique_ptr<FunctionBody> func_body;
if (!FunctionDefToBodyHelper(*func_def, AttrSlice(&func_def->attr()),
&flib_def, &func_body)
.ok())
return false;
if (DoesGraphContainTPUPartitionedCall(*func_body->graph)) return true;
}
return false;
}
// Iterate all functions from the flib_def if there are any that belong to
// the inference graph.
bool AreFunctionsFromFlibDefInference(
const FunctionLibraryDefinition& flib_def) {
for (const std::string& func_name : flib_def.ListFunctionNames()) {
const FunctionDef* func_def = flib_def.Find(func_name);
for (const NodeDef& node_def : func_def->node_def()) {
if (node_def.op() == kPartitionedCall) return true;
}
}
return false;
}
} // namespace
bool IsSupportedByNonReplicatedBridge(
const Graph& graph, const FunctionLibraryDefinition* function_library) {
return IsNonReplicatedGraph(graph, function_library) &&
HasPsWithResourceVariable(graph);
}
bool IsSupportedByReplicatedBridge(
const Graph& graph, const FunctionLibraryDefinition* function_library) {
return IsReplicatedGraph(graph, function_library);
}
bool IsSupportedByReplicatedBridge(mlir::ModuleOp module) {
return IsReplicatedGraph(module);
}
bool HasTPUPartitionedCallOpInModule(mlir::ModuleOp module) {
bool has_tpu_partitioned_call = false;
for (auto func_op : module.getOps<mlir::func::FuncOp>()) {
func_op->walk([&](mlir::TF::TPUPartitionedCallOp op) {
has_tpu_partitioned_call = true;
});
if (has_tpu_partitioned_call) break;
}
return has_tpu_partitioned_call;
}
bool IsInferenceGraph(const Graph& graph,
const FunctionLibraryDefinition* function_library) {
if (DoesGraphContainTPUPartitionedCall(graph)) return true;
GraphDef graph_def;
graph.ToGraphDef(&graph_def);
if (DoReachableFuncsContainTPUPartitionedCall(graph_def, graph.flib_def()))
return true;
if (AreFunctionsFromFlibDefInference(graph.flib_def())) return true;
if (function_library == nullptr) return false;
if (DoReachableFuncsContainTPUPartitionedCall(graph_def, *function_library))
return true;
if (AreFunctionsFromFlibDefInference(*function_library)) return true;
return false;
}
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/mlir_bridge_pass_util.h
+54
View File
@@ -0,0 +1,54 @@
/* 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_MLIR_BRIDGE_PASS_UTIL_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_MLIR_BRIDGE_PASS_UTIL_H_
#include <optional>
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "tensorflow/core/framework/function.h"
namespace tensorflow {
// Checks if a graph or reachable functions in the library have any
// StatefulPartitionedOps with _XlaMustCompile=true. The function library will
// be skipped if nullptr is provided.
bool IsSupportedByNonReplicatedBridge(
const Graph& graph, const FunctionLibraryDefinition* function_library);
// Checks if a graph or reachable functions in the library have any ops with
// _tpu_replicate or _xla_compile_device_type=TPU. The function library will be
// skipped if nullptr is provided.
bool IsSupportedByReplicatedBridge(
const Graph& graph, const FunctionLibraryDefinition* function_library);
// Check if an MLIR module has any ops with _tpu_replicate or
// _xla_compile_device_type=TPU.
bool IsSupportedByReplicatedBridge(mlir::ModuleOp module);
// Check if an MLIR module contains TPUPartitionedCall op. If so, we define
// such graph as an inference graph. Otherwise, it is non inference graph.
bool HasTPUPartitionedCallOpInModule(mlir::ModuleOp module);
// Check if a graph contains TPUPartitionedCall op, including its reachable
// functions. The function library is used to store the functions that are
// defined in a TensorFlow program
bool IsInferenceGraph(const Graph& graph,
const FunctionLibraryDefinition* function_library);
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_MLIR_BRIDGE_PASS_UTIL_H_
tensorflow/compiler/mlir/tf2xla/internal/mlir_bridge_pass_util_test.cc
+316
View File
@@ -0,0 +1,316 @@
/* 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/mlir_bridge_pass_util.h"
#include <cstdint>
#include <string>
#include <vector>
#include <gtest/gtest.h>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "mlir/Parser/Parser.h" // from @llvm-project
#include "tensorflow/cc/framework/ops.h"
#include "tensorflow/cc/framework/scope.h"
#include "tensorflow/cc/ops/array_ops.h"
#include "tensorflow/cc/ops/function_ops.h"
#include "tensorflow/cc/ops/tpu_functional_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/tf2xla/tf2xla_defs.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/function_testlib.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/platform/enable_tf2_utils.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/config.pb.h"
namespace tensorflow {
namespace {
// Produce a valid graph with a resource-type input.
FunctionDef PassThroughResource() {
return FunctionDefHelper::Define(
/*function_name=*/"PassThroughResource",
/*arg_def=*/{"in: resource"},
/*ret_def=*/{"out: resource"},
/*attr_def=*/{},
/*node_def=*/
{{{"out"}, "Identity", {"in"}, {{"T", DataType::DT_RESOURCE}}}});
}
TEST(IsSupportedByNonReplicatedBridge, NonReplicatedGraph) {
const FunctionDef& fd = PassThroughResource();
FunctionDefLibrary flib;
*flib.add_function() = fd;
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
tensorflow::set_tf2_execution(true);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
Output a = ops::_Arg(root.WithOpName("A"), DT_RESOURCE, 0);
std::vector<NodeBuilder::NodeOut> inputs({NodeBuilder::NodeOut(a.node())});
Node* call;
NameAttrList f_name_attr;
f_name_attr.set_name(fd.signature().name());
TF_ASSERT_OK(
NodeBuilder("B", "StatefulPartitionedCall", &root.graph()->flib_def())
.Input(inputs)
.Attr("Tin", {DT_RESOURCE})
.Attr("Tout", {DT_RESOURCE})
.Attr("f", f_name_attr)
.Finalize(root.graph(), &call));
call->AddAttr(std::string(kMustCompileAttr), true);
TF_ASSERT_OK(root.ToGraph(&graph));
// Required for passing the PS server parameter check.
for (Node* node : graph.nodes()) {
node->set_assigned_device_name("/job:ps/replica:0/task:0/device:GPU:0");
}
EXPECT_TRUE(
IsSupportedByNonReplicatedBridge(graph, /*function_library=*/nullptr));
}
TEST(IsSupportedByReplicatedBridge, ReplicatedGraph) {
const FunctionDef& fd = test::function::XTimesTwo();
FunctionDefLibrary flib;
*flib.add_function() = fd;
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
tensorflow::set_tf2_execution(true);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
Output a = ops::_Arg(root.WithOpName("A"), DT_FLOAT, 0);
std::vector<NodeBuilder::NodeOut> inputs({NodeBuilder::NodeOut(a.node())});
Node* call;
NameAttrList f_name_attr;
f_name_attr.set_name(fd.signature().name());
TF_ASSERT_OK(
NodeBuilder("B", "StatefulPartitionedCall", &root.graph()->flib_def())
.Input(inputs)
.Attr("Tin", {DT_FLOAT})
.Attr("Tout", {DT_FLOAT})
.Attr("f", f_name_attr)
.Finalize(root.graph(), &call));
call->AddAttr(std::string(kTpuReplicateAttr), "cluster");
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_TRUE(
IsSupportedByReplicatedBridge(graph, /*function_library=*/nullptr));
}
TEST(IsSupportedByReplicatedBridge, ReplicatedModule) {
const char* const code = R"mlir(
func.func @entry_func_1(%arg0: tensor<i32>) -> tensor<i32> attributes {tf.entry_function = {}} {
%0 = "tf.Identity"(%arg0) {_tpu_replicate = "cluster"} : (tensor<i32>) -> (tensor<i32>)
func.return %0 : tensor<i32>
}
)mlir";
mlir::MLIRContext context;
context.loadDialect<mlir::func::FuncDialect, mlir::TF::TensorFlowDialect>();
mlir::OwningOpRef<mlir::ModuleOp> module =
mlir::parseSourceString<mlir::ModuleOp>(code, &context);
ASSERT_TRUE(module);
EXPECT_TRUE(IsSupportedByReplicatedBridge(*module));
}
TEST(HasTPUPartitionedCallOpInModule, HasTPUPartitionedCallModule) {
const char* const code = R"mlir(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() {
%outputs_0 = "tf.TPUOrdinalSelector"() {device = ""} : () -> tensor<?xi32>
"tf.TPUPartitionedCall"(%outputs_0) {f = @reachable_func} : (tensor<?xi32>) -> ()
func.return
}
func.func @reachable_func() {
func.return
}
}
)mlir";
mlir::MLIRContext context;
context.loadDialect<mlir::func::FuncDialect, mlir::TF::TensorFlowDialect>();
mlir::OwningOpRef<mlir::ModuleOp> module =
mlir::parseSourceString<mlir::ModuleOp>(code, &context);
ASSERT_TRUE(module);
EXPECT_TRUE(HasTPUPartitionedCallOpInModule(*module));
}
TEST(HasTPUPartitionedCallOpInModule, HasNotTPUPartitionedCallModule) {
const char* const code = R"mlir(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() {
"tf.StatefulPartitionedCall"() {config = "", config_proto = "", executor_type = "", f = @reachable_func} : () -> ()
func.return
}
func.func @reachable_func() {
func.return
}
}
)mlir";
mlir::MLIRContext context;
context.loadDialect<mlir::func::FuncDialect, mlir::TF::TensorFlowDialect>();
mlir::OwningOpRef<mlir::ModuleOp> module =
mlir::parseSourceString<mlir::ModuleOp>(code, &context);
ASSERT_TRUE(module);
EXPECT_FALSE(HasTPUPartitionedCallOpInModule(*module));
}
TEST(IsInferenceGraph, GraphContrainsTPUPartitionedCall) {
FunctionDef fd = FunctionDefHelper::Define(
// Name
"XTimesTwoFloat",
// Args
{"x: float"},
// Return values
{"y: float"},
// Attr def
{},
// Nodes
{
{{"two"},
"Const",
{},
{{"value", test::AsScalar<int32_t>(2)}, {"dtype", DT_INT64}}},
{{"scale"},
"Cast",
{"two"},
{{"SrcT", DT_INT64}, {"DstT", DT_FLOAT}}},
{{"y"}, "Mul", {"x", "scale"}, {{"T", DT_FLOAT}}},
});
tensorflow::set_tf2_execution(true);
FunctionDefLibrary flib;
*flib.add_function() = fd;
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kDirectSession);
Scope root = Scope::NewRootScope().ExitOnError();
Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
NameAttrList f_name_attr;
f_name_attr.set_name("XTimesTwoFloat");
ops::TPUPartitionedCall f(root.WithOpName("f"), {x}, /*device_ordinal=*/0,
{DT_FLOAT}, f_name_attr);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_TRUE(IsInferenceGraph(graph, /*function_library=*/nullptr));
}
TEST(IsInferenceGraph, GraphDoesNotContrainTPUPartitionedCall) {
FunctionDef fd = FunctionDefHelper::Define(
// Name
"XTimesTwoFloat",
// Args
{"x: float"},
// Return values
{"y: float"},
// Attr def
{},
// Nodes
{
{{"two"},
"Const",
{},
{{"value", test::AsScalar<int32_t>(2)}, {"dtype", DT_INT64}}},
{{"scale"},
"Cast",
{"two"},
{{"SrcT", DT_INT64}, {"DstT", DT_FLOAT}}},
{{"y"}, "Mul", {"x", "scale"}, {{"T", DT_FLOAT}}},
});
tensorflow::set_tf2_execution(true);
FunctionDefLibrary flib;
*flib.add_function() = fd;
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kDirectSession);
Scope root = Scope::NewRootScope().ExitOnError();
Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
NameAttrList f_name_attr;
f_name_attr.set_name("XTimesTwoFloat");
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_FALSE(IsInferenceGraph(graph, /*function_library=*/nullptr));
}
TEST(IsInferenceGraph, FlibDefIsNotNullptrAndContainsTPUPartitionedCall) {
FunctionDef fd = FunctionDefHelper::Define(
// Name
"XTimesTwoFloat",
// Args
{"x: float"},
// Return values
{"y: float"},
// Attr def
{},
// Nodes
{
{{"two"},
"Const",
{},
{{"value", test::AsScalar<int32_t>(2)}, {"dtype", DT_INT64}}},
{{"scale"},
"Cast",
{"two"},
{{"SrcT", DT_INT64}, {"DstT", DT_FLOAT}}},
{{"y"}, "Mul", {"x", "scale"}, {{"T", DT_FLOAT}}},
{{"tpu_op"}, "TPUPartitionedCall", {}, {{"Tout", DT_FLOAT}}},
});
tensorflow::set_tf2_execution(true);
FunctionDefLibrary flib;
*flib.add_function() = fd;
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kDirectSession);
Scope root = Scope::NewRootScope().ExitOnError();
Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
NameAttrList f_name_attr;
f_name_attr.set_name("XTimesTwoFloat");
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_TRUE(IsInferenceGraph(graph, /*function_library=*/&flib_def));
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/mlir_pass_instrumentation.cc
+67
View File
@@ -0,0 +1,67 @@
/* 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/mlir_pass_instrumentation.h"
#include <algorithm>
#include <functional>
#include <iterator>
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>
#include "absl/log/log.h"
#include "tensorflow/core/platform/logging.h"
namespace mlir {
class MlirPassInstrumentationRegistry {
public:
static MlirPassInstrumentationRegistry& Instance() {
static MlirPassInstrumentationRegistry* r =
new MlirPassInstrumentationRegistry;
return *r;
}
std::unordered_map<std::string,
std::function<std::unique_ptr<PassInstrumentation>()>>
instrumentors_;
};
void RegisterPassInstrumentor(
const std::string& name,
std::function<std::unique_ptr<PassInstrumentation>()> creator) {
MlirPassInstrumentationRegistry& r =
MlirPassInstrumentationRegistry::Instance();
auto result = r.instrumentors_.emplace(name, creator);
if (!result.second) {
VLOG(1) << "Duplicate MLIR pass instrumentor registration";
}
}
std::vector<std::function<std::unique_ptr<PassInstrumentation>()>>
GetPassInstrumentors() {
MlirPassInstrumentationRegistry& r =
MlirPassInstrumentationRegistry::Instance();
std::vector<std::function<std::unique_ptr<PassInstrumentation>()>> result;
result.reserve(r.instrumentors_.size());
std::transform(r.instrumentors_.begin(), r.instrumentors_.end(),
std::back_inserter(result), [](auto v) { return v.second; });
return result;
}
} // namespace mlir
tensorflow/compiler/mlir/tf2xla/internal/mlir_pass_instrumentation.h
+36
View File
@@ -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_TF2XLA_INTERNAL_MLIR_PASS_INSTRUMENTATION_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_MLIR_PASS_INSTRUMENTATION_H_
#include <functional>
#include <memory>
#include <string>
#include <vector>
#include "mlir/Pass/PassInstrumentation.h" // from @llvm-project
namespace mlir {
void RegisterPassInstrumentor(
const std::string& name,
std::function<std::unique_ptr<PassInstrumentation>()> creator);
std::vector<std::function<std::unique_ptr<PassInstrumentation>()>>
GetPassInstrumentors();
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_MLIR_PASS_INSTRUMENTATION_H_
tensorflow/compiler/mlir/tf2xla/internal/mlir_pass_instrumentation_test.cc
+112
View File
@@ -0,0 +1,112 @@
/* 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/mlir_pass_instrumentation.h"
#include <cstddef>
#include <cstdint>
#include <memory>
#include <sstream>
#include <string>
#include <unordered_map>
#include <vector>
#include <gtest/gtest.h>
#include "absl/strings/match.h"
#include "absl/strings/string_view.h"
#include "tensorflow/compiler/mlir/tf2xla/api/v1/compile_mlir_util.h"
#include "tensorflow/core/lib/core/status_test_util.h"
namespace mlir {
namespace {
static const char* kTestInstrumentationName = "test-intrumentatron";
static const char* kTestInstrumentationSearch = "tf.Identity";
struct StringStream : public llvm::raw_ostream {
StringStream() { SetUnbuffered(); }
~StringStream() override = default;
uint64_t current_pos() const override { return 0; }
void write_impl(const char* ptr, size_t size) override {
ss.write(ptr, size);
}
std::stringstream ss;
};
class TestPassInstrumentation : public ::testing::Test {
public:
void SetPassThatChangedIdentity(absl::string_view pass_name) {
pass_that_changed_identity_ = pass_name;
}
absl::string_view GetPassThatChangedIdentity() {
return pass_that_changed_identity_;
}
private:
std::string pass_that_changed_identity_;
friend class TestInstrumentor;
};
class TestInstrumentor : public PassInstrumentation {
public:
explicit TestInstrumentor(TestPassInstrumentation* test) : test_(test) {}
private:
void runBeforePass(Pass* pass, Operation* op) override {
StringStream stream;
op->print(stream, mlir::OpPrintingFlags().useLocalScope());
ops_seen_by_pass_[pass] = stream.ss.str();
}
void runAfterPass(Pass* pass, Operation* op) override {
StringStream stream;
op->print(stream, mlir::OpPrintingFlags().useLocalScope());
if (!absl::StrContains(stream.ss.str(), kTestInstrumentationSearch) &&
absl::StrContains(ops_seen_by_pass_[pass],
kTestInstrumentationSearch)) {
test_->SetPassThatChangedIdentity(pass->getName().str());
}
}
private:
TestPassInstrumentation* test_;
std::unordered_map<mlir::Pass*, std::string> ops_seen_by_pass_;
};
TEST_F(TestPassInstrumentation, CreatedCalledAndSetsPassName) {
RegisterPassInstrumentor(kTestInstrumentationName, [&]() {
return std::make_unique<TestInstrumentor>(this);
});
constexpr char legalization[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main(%arg0: tensor<?xi32, #mhlo.type_extensions<bounds = [1]>>) -> tensor<?xi32, #mhlo.type_extensions<bounds = [1]>> {
%0 = "tf.Identity"(%arg0) : (tensor<?xi32, #mhlo.type_extensions<bounds = [1]>>) -> tensor<?xi32, #mhlo.type_extensions<bounds = [1]>>
func.return %0 : tensor<?xi32, #mhlo.type_extensions<bounds = [1]>>
}
})";
SetPassThatChangedIdentity("");
std::vector<::tensorflow::TensorShape> arg_shapes = {{1}};
auto compilation_result = tensorflow::XlaCompilationResult();
TF_EXPECT_OK(tensorflow::CompileSerializedMlirToXlaHlo(
legalization, arg_shapes, /*device_type=*/"XLA_TPU_JIT",
/*use_tuple_args=*/true, /*enable_op_fallback=*/false,
/*shape_determination_fns=*/{}, &compilation_result)
.status());
EXPECT_FALSE(GetPassThatChangedIdentity().empty());
}
} // namespace
} // namespace mlir
tensorflow/compiler/mlir/tf2xla/internal/node_order.cc
+122
View File
@@ -0,0 +1,122 @@
/* 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/node_order.h"
#include <algorithm>
#include <cassert>
#include <functional>
#include <iterator>
#include <set>
#include <string>
#include <unordered_map>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "tensorflow/core/graph/algorithm.h"
#include "tensorflow/core/graph/graph.h"
namespace tensorflow {
void TopologicalOrdering(
const Graph& g, const std::function<void(Node*)>& emit,
const std::function<std::string(Node*)>& get_grouping_key) {
std::unordered_map<std::string, int> group_key_string_to_integer;
absl::flat_hash_map<Node*, int> node_to_group;
absl::flat_hash_map<Node*, int> remaining_incoming_nodes;
absl::flat_hash_map<Node*, int> node_to_position;
using Ready = std::vector<Node*>;
std::vector<Ready> group_members_that_are_ready;
std::set<int> groups_that_are_ready;
// Visit all nodes once, for initialization. It doesn't matter whether we use
// BFS or DFS.
int i = 0;
DFS(
g, [](Node*) {},
[&](Node* n) {
// Find which group this node belongs to.
std::string group_key_string = get_grouping_key(n);
auto entry = group_key_string_to_integer.try_emplace(
group_key_string, group_key_string_to_integer.size());
int group_key = entry.first->second;
node_to_position[n] = i++;
node_to_group[n] = group_key;
if (entry.second) {
group_members_that_are_ready.push_back({});
}
// Count the incoming nodes and store. Also remember nodes ("sources")
// that don't have any inputs.
auto in_nodes = n->in_nodes();
int num_incoming = std::distance(in_nodes.begin(), in_nodes.end());
remaining_incoming_nodes[n] = num_incoming;
if (num_incoming == 0) {
// NO_CDC: This array is max(group_key) + 1.
group_members_that_are_ready[group_key].push_back(n);
groups_that_are_ready.emplace(group_key);
}
},
[](const Node* n1, const Node* n2) { return n1->name() < n2->name(); });
assert(group_key_string_to_integer.size() ==
group_members_that_are_ready.size());
int num_nodes = remaining_incoming_nodes.size();
// We emit one node per step, thus we just run this as often as we have nodes.
int current_group = 0;
for (int i = 0; i < num_nodes; i++) {
if (groups_that_are_ready.find(current_group) ==
groups_that_are_ready.end()) {
current_group = *groups_that_are_ready.begin();
}
// NO_CDC: This array is max(group_key) + 1.
int size = group_members_that_are_ready[current_group].size();
assert(size);
// NO_CDC: This array is max(group_key) + 1.
Node* node = group_members_that_are_ready[current_group][--size];
// NO_CDC: This array is max(group_key) + 1.
group_members_that_are_ready[current_group].pop_back();
if (size == 0) {
groups_that_are_ready.erase(current_group);
}
// Emit the operation and make its results available.
emit(node);
auto out_nodes = node->out_nodes();
std::vector<Node*> nodes_sorted(out_nodes.begin(), out_nodes.end());
std::sort(nodes_sorted.begin(), nodes_sorted.end(), [&](Node* a, Node* b) {
return node_to_position[a] < node_to_position[b];
});
for (Node* out : nodes_sorted) {
remaining_incoming_nodes[out]--;
if (remaining_incoming_nodes[out] == 0) {
int group_key = node_to_group[out];
// NO_CDC: This array is max(group_key) + 1.
if (group_members_that_are_ready[group_key].empty()) {
groups_that_are_ready.emplace(group_key);
}
// NO_CDC: This array is max(group_key) + 1.
group_members_that_are_ready[group_key].push_back(out);
}
}
}
}
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/node_order.h
+51
View File
@@ -0,0 +1,51 @@
/* 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_NODE_ORDER_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_NODE_ORDER_H_
#include <functional>
#include <unordered_set>
#include <vector>
#include "absl/container/flat_hash_set.h"
#include "tensorflow/core/graph/algorithm.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/lib/gtl/array_slice.h"
namespace tensorflow {
struct GroupByDevice {
std::string operator()(const Node* node) const {
return node->requested_device();
}
};
// Performs a topological ordering of nodes.
// This has the property that any child node of a parent node p is emitted
// before p. A grouping function is used to break ties if multiple child nodes
// (of possibly different parents) are ready to be emitted at some point, which
// is when we prefer to stay in the current group. Remaining ties are broken by
// node name.
// The "emit" function is used for outputing the result, and is called once
// for each node.
// This algorithm is O(n * k * log k), with k the largest node degree.
void TopologicalOrdering(
const Graph& g, const std::function<void(Node*)>& emit,
const std::function<std::string(Node*)>& get_grouping_key);
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_NODE_ORDER_H_
tensorflow/compiler/mlir/tf2xla/internal/node_order_test.cc
+294
View File
@@ -0,0 +1,294 @@
/* 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/node_order.h"
#include <string>
#include <utility>
#include <vector>
#include "xla/tsl/lib/core/status_test_util.h"
#include "tensorflow/core/common_runtime/graph_def_builder_util.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/graph/algorithm.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/graph/graph_def_builder.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace {
REGISTER_OP("TestParams").Output("o: float");
REGISTER_OP("TestInput").Output("a: float").Output("b: float");
REGISTER_OP("TestMul").Input("a: float").Input("b: float").Output("o: float");
REGISTER_OP("TestUnary").Input("a: float").Output("o: float");
REGISTER_OP("TestTwoOutputs").Output("a: float").Output("b: float");
REGISTER_OP("TestBinary")
.Input("a: float")
.Input("b: float")
.Output("o: float");
// Compares that the order of nodes in 'inputs' respects the
// pair orders described in 'ordered_pairs'.
bool ExpectBefore(
const std::vector<std::pair<std::string, std::string>>& ordered_pairs,
const std::vector<Node*>& inputs, std::string* error) {
for (const std::pair<std::string, std::string>& pair : ordered_pairs) {
const std::string& before_node = pair.first;
const std::string& after_node = pair.second;
bool seen_before = false;
bool seen_both = false;
for (const Node* node : inputs) {
if (!seen_before && after_node == node->name()) {
*error = std::string("Saw ") + after_node + std::string(" before ") +
before_node;
return false;
}
if (before_node == node->name()) {
seen_before = true;
} else if (after_node == node->name()) {
seen_both = seen_before;
break;
}
}
if (!seen_both) {
*error = std::string("didn't see either ") + before_node +
std::string(" or ") + after_node;
return false;
}
}
return true;
}
TEST(AlgorithmTest, TopologicalOrdering) {
GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
using namespace ::tensorflow::ops; // NOLINT
Node* n1 = SourceOp("TestParams", b.opts().WithName("n1"));
Node* n2 =
SourceOp("TestParams", b.opts().WithName("n2").WithControlInput(n1));
Node* n3 =
SourceOp("TestParams", b.opts().WithName("n3").WithControlInput(n2));
Node* n4 = BinaryOp("TestMul", n1, {n3, 0}, b.opts().WithName("n4"));
Node* n5 = BinaryOp("TestMul", n1, {n3, 0},
b.opts().WithName("n5").WithControlInput(n1));
Node* n6 = BinaryOp("TestMul", n2, {n3, 0}, b.opts().WithName("n6"));
n3->set_requested_device("a");
n4->set_requested_device("a");
n5->set_requested_device("b");
n6->set_requested_device("b");
Graph g(OpRegistry::Global());
TF_ASSERT_OK(GraphDefBuilderToGraph(b, &g));
std::vector<Node*> order;
TopologicalOrdering(g, [&](Node* n) { order.push_back(n); }, GroupByDevice());
std::vector<std::pair<std::string, std::string>> desired_order = {
{"n1", "n2"}, // because of control dependency
{"n2", "n3"}, // because of control dependency
{"n3", "n4"}, // because of NodeScorerDevice
{"n1", "n4"}, // data dependency
{"n1", "n5"}, // data dependency
{"n2", "n6"}, // data dependency
{"n3", "n4"}, // data dependency
{"n3", "n5"}, // data dependency
{"n3", "n6"}, // data dependency
};
std::string error;
EXPECT_TRUE(ExpectBefore(desired_order, order, &error)) << error;
}
TEST(AlgorithmTest, TopologicalOrderingOnShallowTree) {
GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
using namespace ::tensorflow::ops; // NOLINT
Node* n1 = SourceOp("TestParams", b.opts().WithName("n1").WithDevice("a"));
Node* n2 =
SourceOp("TestParams",
b.opts().WithName("n2").WithDevice("b").WithControlInput(n1));
Node* n3 =
SourceOp("TestParams",
b.opts().WithName("n3").WithDevice("c").WithControlInput(n2));
Node* n4 =
SourceOp("TestParams",
b.opts().WithName("n4").WithDevice("a").WithControlInput(n1));
Node* n5 =
SourceOp("TestParams",
b.opts().WithName("n5").WithDevice("b").WithControlInput(n2));
Node* n6 =
SourceOp("TestParams",
b.opts().WithName("n6").WithDevice("c").WithControlInput(n3));
Node* n7 =
SourceOp("TestParams",
b.opts().WithName("n7").WithDevice("a").WithControlInput(n4));
Node* n8 =
SourceOp("TestParams",
b.opts().WithName("n8").WithDevice("b").WithControlInput(n5));
Node* n9 =
SourceOp("TestParams",
b.opts().WithName("n9").WithDevice("c").WithControlInput(n6));
Graph g(OpRegistry::Global());
TF_ASSERT_OK(GraphDefBuilderToGraph(b, &g));
std::vector<Node*> order;
TopologicalOrdering(g, [&](Node* n) { order.push_back(n); }, GroupByDevice());
std::vector<Node*> desired_order = {
g.source_node(), n1, n4, n7, n2, n5, n8, n3, n6, n9, g.sink_node()};
for (int i = 0; i < desired_order.size(); i++) {
desired_order[i] = g.FindNodeId(desired_order[i]->id());
}
EXPECT_EQ(order, desired_order);
}
TEST(AlgorithmTest, TopologicalOrderingGivesTheSameResultIfCalledTwice) {
GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
using namespace ::tensorflow::ops; // NOLINT
SourceOp("TestParams", b.opts().WithName("n1"));
SourceOp("TestParams", b.opts().WithName("n2"));
SourceOp("TestParams", b.opts().WithName("n3"));
SourceOp("TestParams", b.opts().WithName("n4"));
SourceOp("TestParams", b.opts().WithName("n5"));
SourceOp("TestParams", b.opts().WithName("n6"));
SourceOp("TestParams", b.opts().WithName("n7"));
SourceOp("TestParams", b.opts().WithName("n8"));
SourceOp("TestParams", b.opts().WithName("n9"));
Graph g(OpRegistry::Global());
TF_ASSERT_OK(GraphDefBuilderToGraph(b, &g));
std::vector<Node*> order1;
std::vector<Node*> order2;
TopologicalOrdering(
g, [&](Node* n) { order1.push_back(n); },
[&](const Node* node) { return std::string("same"); });
TopologicalOrdering(
g, [&](Node* n) { order2.push_back(n); },
[&](const Node* node) { return std::string("same"); });
EXPECT_EQ(order1, order2);
}
TEST(AlgorithmTest, TopologicalOrderingOnChain) {
GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
using namespace ::tensorflow::ops; // NOLINT
Node* n1 = SourceOp("TestParams", b.opts().WithName("n1"));
Node* n2 = UnaryOp("TestUnary", n1, b.opts().WithName("n2"));
Node* n3 = UnaryOp("TestUnary", n2, b.opts().WithName("n3"));
Node* n4 = UnaryOp("TestUnary", n3, b.opts().WithName("n4"));
Node* n5 = UnaryOp("TestUnary", n4, b.opts().WithName("n5"));
Node* n6 = UnaryOp("TestUnary", n5, b.opts().WithName("n6"));
Graph g(OpRegistry::Global());
TF_ASSERT_OK(GraphDefBuilderToGraph(b, &g));
std::vector<Node*> order;
TopologicalOrdering(g, [&](Node* n) { order.push_back(n); }, GroupByDevice());
std::vector<Node*> desired_order = {g.source_node(), n1, n2, n3, n4, n5, n6,
g.sink_node()};
for (int i = 0; i < desired_order.size(); i++) {
desired_order[i] = g.FindNodeId(desired_order[i]->id());
}
EXPECT_EQ(order, desired_order);
}
TEST(AlgorithmTest, TopologicalOrderingOnMultipleOutputs) {
GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
using namespace ::tensorflow::ops; // NOLINT
Node* n1 = SourceOp("TestTwoOutputs", b.opts().WithName("n1"));
UnaryOp("TestUnary", {n1, 0}, b.opts().WithName("n2"));
UnaryOp("TestUnary", {n1, 1}, b.opts().WithName("n3"));
UnaryOp("TestUnary", {n1, 0}, b.opts().WithName("n4"));
UnaryOp("TestUnary", {n1, 1}, b.opts().WithName("n5"));
Graph g(OpRegistry::Global());
TF_ASSERT_OK(GraphDefBuilderToGraph(b, &g));
std::vector<Node*> order;
TopologicalOrdering(g, [&](Node* n) { order.push_back(n); }, GroupByDevice());
std::vector<std::pair<std::string, std::string>> desired_order = {
{"n1", "n2"},
{"n1", "n3"},
{"n1", "n4"},
{"n1", "n5"},
};
std::string error;
EXPECT_TRUE(ExpectBefore(desired_order, order, &error)) << error;
}
TEST(AlgorithmTest, TopologicalOrderingSameAsReversePostOrder) {
GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
using namespace ::tensorflow::ops; // NOLINT
Node* n = SourceOp("TestTwoOutputs", b.opts().WithName("n"));
Node* n0 = UnaryOp("TestUnary", {n, 0}, b.opts().WithName("n2"));
Node* n1 = UnaryOp("TestUnary", {n, 1}, b.opts().WithName("n1"));
UnaryOp("TestUnary", n0, b.opts().WithName("n1a"));
UnaryOp("TestUnary", n0, b.opts().WithName("n8a"));
UnaryOp("TestUnary", n0, b.opts().WithName("n2a"));
UnaryOp("TestUnary", n0, b.opts().WithName("n7a"));
UnaryOp("TestUnary", n1, b.opts().WithName("n1b"));
UnaryOp("TestUnary", n1, b.opts().WithName("n8b"));
UnaryOp("TestUnary", n1, b.opts().WithName("n2b"));
UnaryOp("TestUnary", n1, b.opts().WithName("n7b"));
UnaryOp("TestUnary", n0, b.opts().WithName("n3a"));
UnaryOp("TestUnary", n0, b.opts().WithName("n6a"));
UnaryOp("TestUnary", n0, b.opts().WithName("n4a"));
UnaryOp("TestUnary", n0, b.opts().WithName("n5a"));
UnaryOp("TestUnary", n1, b.opts().WithName("n3b"));
UnaryOp("TestUnary", n1, b.opts().WithName("n6b"));
UnaryOp("TestUnary", n1, b.opts().WithName("n4b"));
UnaryOp("TestUnary", n1, b.opts().WithName("n5b"));
Graph g(OpRegistry::Global());
TF_ASSERT_OK(GraphDefBuilderToGraph(b, &g));
std::vector<Node*> order;
TopologicalOrdering(g, [&](Node* n) { order.push_back(n); }, GroupByDevice());
std::vector<Node*> desired_order;
GetReversePostOrder(g, &desired_order, [](const Node* n1, const Node* n2) {
return n1->name() < n2->name();
});
EXPECT_EQ(desired_order, order);
}
TEST(AlgorithmTest, TopologicalOrderingWithEachDeviceUsedOnce) {
GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
using namespace ::tensorflow::ops; // NOLINT
SourceOp("TestParams", b.opts().WithName("n1").WithDevice("a"));
SourceOp("TestParams", b.opts().WithName("n2").WithDevice("b"));
SourceOp("TestParams", b.opts().WithName("n3").WithDevice("c"));
SourceOp("TestParams", b.opts().WithName("n4").WithDevice("d"));
Graph g(OpRegistry::Global());
TF_ASSERT_OK(GraphDefBuilderToGraph(b, &g));
int count = 0;
TopologicalOrdering(g, [&](Node* n) { count++; }, GroupByDevice());
EXPECT_EQ(count, 6);
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/passes/BUILD
+573
View File
@@ -0,0 +1,573 @@
load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library")
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow:tensorflow.default.bzl", "filegroup", "get_compatible_with_portable")
load("//tensorflow/compiler/mlir:glob_lit_test.bzl", "glob_lit_tests")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//learning/brain/tfrt/tpu/compiler/mlir:__pkg__",
"//tensorflow/compiler/mlir:__subpackages__",
"//tensorflow/compiler/mlir/tf2xla/api:__subpackages__",
"//tensorflow/compiler/mlir/tf2xla/internal:__subpackages__",
],
licenses = ["notice"],
)
cc_library(
name = "clustering_passes",
hdrs = [
"clustering_passes.h",
],
textual_hdrs = [
"clustering_passes.h.inc",
],
deps = [
":extract_head_tail_outside_compilation",
":extract_outside_compilation",
":hoist_broadcast_read",
":mark_ops_for_outside_compilation",
":tpu_cluster_formation",
":tpu_sharding_identification_pass",
":tpu_validate_inputs",
":tpu_validate_session_inputs",
":verify_clustering_pass",
":xla_broadcast",
":xla_cluster_formation",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
],
)
cc_library(
name = "verify_clustering_pass",
srcs = [
"verify_clustering_pass.cc",
],
deps = [
":clustering_passes_inc_gen",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:string_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_analysis",
"//tensorflow/compiler/mlir/tensorflow:tpu_rewrite_device_util",
"//tensorflow/compiler/mlir/tf2xla/internal/utils:dialect_detection_utils",
"//tensorflow/core:framework",
"//tensorflow/core/transforms/toposort:Pass",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log",
"@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",
],
)
gentbl_cc_library(
name = "clustering_passes_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"clustering_passes.h.inc": [
"-gen-pass-decls",
"-name=TFXLABridgeClustering",
]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "clustering_passes.td",
deps = [
"@llvm-project//mlir:PassBaseTdFiles",
],
)
tf_cc_test(
name = "verify_clustering_pass_test",
srcs = ["verify_clustering_pass_test.cc"],
deps = [
":clustering_passes",
"//tensorflow/compiler/mlir/tf2xla/transforms:test_utils",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@xla//xla/tsl/platform:statusor",
],
)
glob_lit_tests(
name = "all_tests",
data = [":test_utilities"],
driver = "@llvm-project//mlir:run_lit.sh",
test_file_exts = [
"mlir",
],
)
# Bundle together all of the test utilities that are used by tests.
filegroup(
name = "test_utilities",
testonly = True,
data = [
"//tensorflow/compiler/mlir:tf-opt",
"@llvm-project//llvm:FileCheck",
],
)
cc_library(
name = "tpu_cluster_formation",
srcs = ["tpu_cluster_formation.cc"],
textual_hdrs = [
"clustering_passes.h.inc",
],
deps = [
":clustering_passes_inc_gen",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:string_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_analysis",
"//tensorflow/compiler/mlir/tensorflow:tpu_rewrite_device_util",
"//tensorflow/compiler/mlir/tensorflow/transforms:tf_pass_inc_gen",
"//tensorflow/core:framework",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log",
"@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:TransformUtils",
],
)
cc_library(
name = "extract_outside_compilation",
srcs = ["extract_outside_compilation.cc"],
textual_hdrs = [
"clustering_passes.h.inc",
],
deps = [
":clustering_passes_inc_gen",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:device_util",
"//tensorflow/compiler/mlir/tensorflow:serialize_mlir_module_utils",
"//tensorflow/compiler/mlir/tensorflow:string_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_analysis",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow:tpu_rewrite_device_util",
"//tensorflow/compiler/mlir/tensorflow/transforms:shape_inference_pass",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Rewrite",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
],
)
cc_library(
name = "extract_head_tail_outside_compilation",
srcs = ["extract_head_tail_outside_compilation.cc"],
textual_hdrs = [
"clustering_passes.h.inc",
],
deps = [
":clustering_passes_inc_gen",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:device_util",
"//tensorflow/compiler/mlir/tensorflow:string_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_analysis",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow:tpu_rewrite_device_util",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Rewrite",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
],
)
cc_library(
name = "mlir_to_graph_passes",
hdrs = [
"mlir_to_graph_passes.h",
],
textual_hdrs = [
"mlir_to_graph_passes.h.inc",
],
deps = [
":verify_input_dialect_to_executor_pass",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:Pass",
],
)
gentbl_cc_library(
name = "mlir_to_graph_passes_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"mlir_to_graph_passes.h.inc": [
"-gen-pass-decls",
"-name=TFXLABridgeMlirToGraph",
]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "mlir_to_graph_passes.td",
deps = [
"@llvm-project//mlir:PassBaseTdFiles",
],
)
cc_library(
name = "verify_input_dialect_to_executor_pass",
srcs = [
"verify_input_dialect_to_executor_pass.cc",
],
deps = [
":mlir_to_graph_passes_inc_gen",
"//tensorflow/compiler/mlir/tf2xla/internal/utils:dialect_detection_utils",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
],
)
cc_library(
name = "xla_cluster_formation",
srcs = ["xla_cluster_formation.cc"],
textual_hdrs = [
"clustering_passes.h.inc",
],
deps = [
":clustering_passes_inc_gen",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:call_graph_util",
"//tensorflow/compiler/mlir/tensorflow:cluster_util",
"//tensorflow/compiler/mlir/tensorflow:string_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_analysis",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow:tpu_rewrite_device_util",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core/platform:types",
"@com_google_absl//absl/container:flat_hash_map",
"@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:TransformUtils",
],
)
cc_library(
name = "mark_ops_for_outside_compilation",
srcs = ["mark_ops_for_outside_compilation.cc"],
textual_hdrs = [
"clustering_passes.h.inc",
],
deps = [
":clustering_passes_inc_gen",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:string_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_analysis",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow:tpu_rewrite_device_util",
"//tensorflow/compiler/mlir/tensorflow/transforms:lower_tf_lib",
"//tensorflow/compiler/mlir/tensorflow/transforms:tf_pass_inc_gen",
"//tensorflow/compiler/mlir/tensorflow/transforms:verify_no_outside_compilation_markers_pass",
"//tensorflow/compiler/mlir/tf2xla/transforms:legalization_op_config",
"//tensorflow/compiler/mlir/tf2xla/transforms:legalize_tf",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Rewrite",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
],
)
cc_library(
name = "tpu_sharding_identification_pass",
srcs = ["tpu_sharding_identification_pass.cc"],
textual_hdrs = [
"clustering_passes.h.inc",
],
deps = [
":clustering_passes_inc_gen",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:string_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_analysis",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_traits",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow:tpu_rewrite_device_util",
"//tensorflow/compiler/mlir/tensorflow:xla_sharding_util",
"//tensorflow/compiler/mlir/tensorflow/transforms:lower_tf_lib",
"//tensorflow/compiler/mlir/tensorflow/transforms:tf_pass_inc_gen",
"//tensorflow/compiler/mlir/tensorflow/transforms:verify_no_outside_compilation_markers_pass",
"//tensorflow/compiler/mlir/tf2xla/transforms:legalization_op_config",
"//tensorflow/compiler/mlir/tf2xla/transforms:legalize_tf",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Rewrite",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
"@xla//xla:xla_data_proto_cc",
"@xla//xla/hlo/builder:sharding_builder",
],
)
cc_library(
name = "hoist_broadcast_read",
srcs = ["hoist_broadcast_read.cc"],
textual_hdrs = [
"clustering_passes.h.inc",
],
deps = [
":clustering_passes_inc_gen",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:string_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_analysis",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow:tpu_rewrite_device_util",
"//tensorflow/compiler/mlir/tensorflow/transforms:lower_tf_lib",
"//tensorflow/compiler/mlir/tensorflow/transforms:tf_pass_inc_gen",
"//tensorflow/compiler/mlir/tensorflow/transforms:verify_no_outside_compilation_markers_pass",
"//tensorflow/compiler/mlir/tf2xla/transforms:legalization_op_config",
"//tensorflow/compiler/mlir/tf2xla/transforms:legalize_tf",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Rewrite",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
],
)
cc_library(
name = "xla_broadcast",
srcs = ["xla_broadcast.cc"],
textual_hdrs = [
"clustering_passes.h.inc",
],
deps = [
":clustering_passes_inc_gen",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:device_util",
"//tensorflow/compiler/mlir/tensorflow:string_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_analysis",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow:tpu_rewrite_device_util",
"//tensorflow/compiler/mlir/tensorflow:xla_rewrite_util",
"//tensorflow/compiler/mlir/tensorflow:xla_sharding_util",
"//tensorflow/compiler/mlir/tensorflow/transforms:lower_tf_lib",
"//tensorflow/compiler/mlir/tensorflow/transforms:tf_pass_inc_gen",
"//tensorflow/compiler/mlir/tensorflow/transforms:verify_no_outside_compilation_markers_pass",
"//tensorflow/compiler/mlir/tf2xla/transforms:legalization_op_config",
"//tensorflow/compiler/mlir/tf2xla/transforms:legalize_tf",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core/ir/types:Dialect",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Rewrite",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
],
)
cc_library(
name = "lowering_passes",
hdrs = [
"lowering_passes.h",
],
textual_hdrs = [
"lowering_passes.h.inc",
],
deps = [
":input_metrics_lowering_pass",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
],
)
gentbl_cc_library(
name = "lowering_passes_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"lowering_passes.h.inc": [
"-gen-pass-decls",
"-name=TFXLABridgeLowering",
]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "lowering_passes.td",
deps = [
"@llvm-project//mlir:PassBaseTdFiles",
],
)
cc_library(
name = "input_metrics_lowering_pass",
srcs = [
"input_lowering_metrics_pass.cc",
],
deps = [
":lowering_passes_inc_gen",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tf2xla/transforms:legalization_op_config",
"//tensorflow/core:lib",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
],
)
tf_cc_test(
name = "input_metrics_lowering_pass_test",
srcs = ["input_lowering_metrics_pass_test.cc"],
deps = [
":lowering_passes",
"//tensorflow/compiler/mlir/tf2xla/transforms:test_utils",
"//tensorflow/core/lib/monitoring:cell_reader",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@xla//xla/tsl/platform:statusor",
],
)
cc_library(
name = "tpu_validate_inputs_utils",
srcs = ["tpu_validate_inputs_utils.cc"],
hdrs = ["tpu_validate_inputs_utils.h"],
deps = [
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "tpu_validate_session_inputs",
srcs = ["tpu_validate_session_inputs.cc"],
textual_hdrs = [
"clustering_passes.h.inc",
],
deps = [
":clustering_passes_inc_gen",
":tpu_validate_inputs_utils",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
"@com_google_absl//absl/log",
"@com_google_absl//absl/strings",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
],
)
cc_library(
name = "tpu_validate_inputs",
srcs = ["tpu_validate_inputs.cc"],
textual_hdrs = [
"clustering_passes.h.inc",
],
deps = [
":clustering_passes_inc_gen",
":tpu_validate_inputs_utils",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
"//tensorflow/compiler/mlir/tensorflow:string_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_analysis",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow:tpu_rewrite_device_util",
"//tensorflow/compiler/mlir/tensorflow/transforms:tf_pass_inc_gen",
"//tensorflow/core:framework",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log",
"@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:TransformUtils",
"@xla//xla:xla_data_proto_cc",
"@xla//xla/hlo/builder:sharding_builder",
],
)
tf_cc_test(
name = "tpu_validate_inputs_utils_test",
srcs = ["tpu_validate_inputs_utils_test.cc"],
deps = [
":tpu_validate_inputs_utils",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tf2xla/transforms:test_utils",
"@com_google_googletest//:gtest_main",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
],
)
tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h
+93
View File
@@ -0,0 +1,93 @@
/* 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_PASSES_CLUSTERING_PASSES_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_PASSES_CLUSTERING_PASSES_H_
#include <memory>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
namespace tensorflow {
namespace tf2xla {
namespace internal {
// Verifies that all MLIR Ops have the expected attributes.
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateVerifyClusteringPass();
// Creates a pass that forms clusters from operations of the same
// `_replication_info` attribute.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateTPUClusterFormationPass(bool strict_clusters = false);
// Creates a pass that extracts outside compilation (Host ops inside device
// cluster) at head/tail of Device cluster to run before/after XLA computation.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateExtractHeadTailOutsideCompilationPass();
// Creates a pass that extract outside compilation (Host ops inside cevice
// cluster) ops to a separate parallel_execute region to run on CPU.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateExtractOutsideCompilationPass();
// Create a pass that encapsulates StatefulPartitionedCallOp within a cluster.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateXlaClusterFormationPass();
// Creates a pass that marks unsupported ops in device cluster for outside
// compilation.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateMarkOpsForOutsideCompilationPass();
// Creates a pass that hoists reads out of a replicate that are on a variable
// whose value is broacast to all replicas.
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateHoistBroadcastReadPass();
// Creates a pass that moves broadcasts from TF host ops to XLA code, encoded as
// XlaAllReduces. This enables use of the device network for broadcasts, which
// is faster.
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateXlaBroadcastPass();
// Creates a pass that identifies XLASharding ops in launch op for TPU
// computation.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateTPUShardingIdentificationPass();
// Creates a pass that validates the inputs to a TPU computation.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateTPUValidateSessionInputsPass();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateTPUValidateInputsPass();
#define GEN_PASS_REGISTRATION
#define GEN_PASS_DECL_MARKOPSFOROUTSIDECOMPILATIONPASS
#define GEN_PASS_DECL_TPUCLUSTERFORMATIONPASS
#define GEN_PASS_DECL_TPUEXTRACTHEADTAILOUTSIDECOMPILATIONPASS
#define GEN_PASS_DECL_TPUEXTRACTOUTSIDECOMPILATIONPASS
#define GEN_PASS_DECL_TPUSHARDINGIDENTIFICATIONPASS
#define GEN_PASS_DECL_TPUVALIDATEINPUTSPASS
#define GEN_PASS_DECL_TPUVALIDATESESSIONINPUTSPASS
#define GEN_PASS_DECL_VERIFYCLUSTERINGPASS
#define GEN_PASS_DECL_XLACLUSTERFORMATIONPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc"
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_PASSES_CLUSTERING_PASSES_H_
tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.td
+466
View File
@@ -0,0 +1,466 @@
/* 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 "mlir/Pass/PassBase.td"
def VerifyClusteringPass : Pass<"verify-clustering-pass", "mlir::func::FuncOp"> {
let summary = "Verify that the Bridge output is correct and errors if verification fails.";
let description = [{
Verifies whether clustering has resulted in the expected invariants. These
include verifying that clusters have been created and have been outside
compiled, the result is device agnostic and in TF functional dialect &
that the device attribute exists.
}];
let constructor = "tensorflow::tf2xla::internal::CreateVerifyClusteringPass()";
}
def TPUClusterFormationPass : Pass<"tf-tpu-cluster-formation", "ModuleOp"> {
let summary = "Forms clusters from operations assigned to the same TPU computation";
let description = [{
TPU computations from the frontend are composed of a `tf.TPUReplicateMetadata`
op, a subgraph of ops (TensorFlow Dialect) each with a matching
`_replication_info` attribute relative to the associated
`tf.TPUReplicateMetadata` op, and optionally `tf.TPUReplicatedInput` and
`tf.TPUReplicatedOutput` ops feeding in inputs and outputs to and from a
replicated TPU computation. The number of times a TPU computation is
replicated is defined in the `tf.TPUReplicateMetadata` op (`num_replicas`
attribute) and operand and result sizes of `tf.TPUReplicatedInput` and
`tf.TPUReplicatedOutput` respectively must match, excluding packed tensors.
It is also assumed ops of the same TPU computation do not have ops outside
of the TPU computation that are both inputs and outputs to the same TPU
computation. Furthermore, we assume that every node has either none or both
of `_replication_info` and `_xla_compile_device_type` attributes defined.
This pass takes the TPU computation subgraph, moves them into a
`tf_device.cluster`, and copies over attributes from the associated
`tf.TPUReplicateMetadata` op to the newly created `tf_device.cluster`. If the
computation is replicated (`num_replicas` > 1), the `num_replicas` attribute is
not copied over but instead the `tf_device.cluster` is further wrapped with a
`tf_device.replicate`, and associated `tf.TPUReplicatedInput` and
`tf.TPUReplicatedOutput` ops are replaced as the `tf_device.replicate` operands
and results. Otherwise, the single operands and results of the associated
`tf.TPUReplicatedInput` and `tf.TPUReplicatedOutput` ops are simply forwarded to
the `tf_device.cluster`.
For example, the following non replicated computation:
```mlir
func @tpu_computation(%arg0: tensor<i32>) -> tensor<i32> {
// Metadata op for cluster `cluster` with 1 replica, 1 core per replica and
// with topology `<topology>`.
"tf.TPUReplicateMetadata"() {_xla_compile_device_type = "TPU", _replication_info = "cluster", num_relicas = 1, num_cores_per_replica = 1, topology = "<topology>", device_assignment = [], padding_map = []} : () -> ()
%replicated_input = "tf.TPUReplicatedInput"(%arg0) : (tensor<i32>) -> tensor<i32>
%identity = "tf.Identity"(%replicated_input) {_xla_compile_device_type = "TPU", _replication_info = "cluster"} : (tensor<i32>) -> tensor<i32>
%replicated_output = "tf.TPUReplicatedOutput(%identity) : (tensor<i32>) -> tensor<i32>
return %replicated_output : tensor<i32>
}
```
will be transformed into:
```mlir
func @tpu_computation(%arg0: tensor<i32>) -> tensor<i32> {
%cluster = "tf_device.cluster"() ( {
%identity = "tf.Identity"(%arg0) : (tensor<i32>) -> tensor<i32>
tf_device.return %identity : tensor<i32>
}) {_xla_compile_device_type = "TPU", _replication_info = "cluster", num_cores_per_replica = 1, topology = "topology", device_assignment = [], padding_map = []} : () -> (tensor<i32>)
return %cluster : tensor<i32>
}
```
The following replicated computation:
```mlir
func @tpu_computation(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
"tf.TPUReplicateMetadata"() {_xla_compile_device_type = "TPU", _replication_info = "cluster", num_relicas = 2, num_cores_per_replica = 1, topology = "topology", device_assignment = [], padding_map = []} : () -> ()
%replicated_input = "tf.TPUReplicatedInput"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%identity = "tf.Identity"(%replicated_input) {_xla_compile_device_type = "TPU", _replication_info = "cluster"} : (tensor<i32>) -> tensor<i32>
%replicated_output:2 = "tf.TPUReplicatedOutput(%identity) : (tensor<i32>) -> (tensor<i32>, tensor<i32>)
return %replicated_output#0, %replicated_output#1 : tensor<i32>, tensor<i32>
}
```
will be transformed into:
```mlir
func @tpu_computation(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
%replicate:2 = tf_device.replicate([%arg0, %arg1] as %replicated_input) {n = 2 : i32} {
%cluster = "tf_device.cluster"() ( {
%identity = "tf.Identity"(%replicated_input) : (tensor<i32>) -> tensor<i32>
tf_device.return %identity : tensor<i32>
}) {_xla_compile_device_type = "TPU", _replication_info = "cluster", num_cores_per_replica = 1, topology = "topology", device_assignment = [], padding_map = []} : () -> (tensor<i32>)
tf_device.return %cluster : tensor<i32>
}
return %replicate#0, %replicate#1 : tensor<i32>, tensor<i32>
}
```
}];
let constructor = "tensorflow::tf2xla::internal::CreateTPUClusterFormationPass()";
}
def ExtractHeadTailOutsideCompilationPass : Pass<"tf-extract-head-tail-outside-compilation", "ModuleOp"> {
let summary = "Extracts head or tail outside compilation to separate host launches before/after device cluster.";
let description = [{
This pass extracts a CPU computation cluster with `_xla_outside_compilation`
annotation from the head or tail of a Device cluster.
For example:
```mlir
%cluster = "tf_device.cluster"() ( {
%a = "tf.A"(%arg0) {_xla_outside_compilation = "cluster1"} : (tensor<i32>) -> tensor<i32>
%b = "tf.B"(%a) : (tensor<i32>) -> tensor<i32>
%c = "tf.C"(%b) {_xla_outside_compilation = "cluster1"} : (tensor<i32>) -> tensor<i32>
tf_device.return %c : tensor<i32>
}) {num_cores_per_replica = 1, step_marker_location = "", padding_map = [], topology = "", device_assignment = []} : () -> tensor<i32>
return %cluster : tensor<i32>
```
becomes:
```mlir
%0 = "tf_device.launch"() ( {
%3 = "tf.A"(%arg0) : (tensor<i32>) -> tensor<i32>
tf_device.return %3 : tensor<i32>
}) {device = "/job:worker/replica:0/task:0/device:CPU:0"} : () -> tensor<i32>
%1 = "tf_device.cluster"() ( {
%3 = "tf.B"(%0) : (tensor<i32>) -> tensor<i32>
tf_device.return %3 : tensor<i32>
}) {device_assignment = [], num_cores_per_replica = 1 : i64, padding_map = [], step_marker_location = "", topology = ""} : () -> tensor<i32>
%2 = "tf_device.launch"() ( {
%3 = "tf.C"(%1) : (tensor<i32>) -> tensor<i32>
tf_device.return %3 : tensor<i32>
}) {device = "/job:worker/replica:0/task:0/device:CPU:0"} : () -> tensor<i32>
return %2 : tensor<i32>
```
}];
let constructor = "tensorflow::tf2xla::internal::CreateExtractHeadTailOutsideCompilationPass()";
}
def ExtractOutsideCompilationPass : Pass<"tf-extract-outside-compilation", "ModuleOp"> {
let summary = "Extracts device outside compilation computation to a separate tf_device.parallel_execute region.";
let description = [{
This pass extracts a CPU computation cluster with `_xla_outside_compilation`
annotation, which denotes ops that should be run on CPU/host, from a device cluster.
Each outside compilation cluster is moved to
a tf_device.parallel_execute region. The device cluster is also moved to a
tf_device.parallel_execute region. Communication ops between device and host are
added to pass inputs/outputs to/from the outside compiled region.
For example, the following tf_device.cluster with an op marked for `xla_outside_compilation`:
```mlir
func @outside_compilation() -> tensor<f32> {
%0 = "tf_device.cluster"() ( {
%1 = "tf.Const"() {_xla_outside_compilation = "0", value = dense<1.0> : tensor<f32>} : () -> (tensor<f32>)
%2 = "tf.Identity"(%1) {_xla_outside_compilation = "0"} : (tensor<f32>) -> (tensor<f32>)
%3 = "tf.AddV2"(%1, %2) : (tensor<f32>, tensor<f32>) -> (tensor<f32>)
tf_device.return %3 : tensor<f32>
}) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> tensor<f32>
return %0 : tensor<f32>
}
```
will become a tf_device.parallel_execute op with a CPU/host region and
a tf_device.cluster with communication ops to send data to/from device/host:
```mlir
func @outside_compilation() -> tensor<f32> {
%0 = "tf_device.parallel_execute"() ( {
"tf_device.launch"() ( {
%1 = "tf._XlaCompileMlirPlaceholderProgramKey"() : () -> tensor<3x!tf_type.string>
%2 = "tf._XlaRecvAtHost"(%1) {device_ordinal = 0 : i64, key = "host_compute_channel_0_0_args"} : (tensor<3x!tf_type.string>) -> tensor<f32>
%3 = "tf.Identity"(%2) : (tensor<f32>) -> tensor<f32>
"tf._XlaSendFromHost"(%3, %1) {device_ordinal = 0 : i64, key = "host_compute_channel_0_0_retvals"} : (tensor<f32>, tensor<3x!tf_type.string>) -> ()
tf_device.return
}) {device = "/job:worker/replica:0/task:0/device:CPU:0"} : () -> ()
tf_device.return
}, {
%1 = "tf_device.cluster"() ( {
%2 = "tf.Const"() {value = dense<1.000000e+00> : tensor<f32>} : () -> tensor<f32>
%3 = "tf._XlaHostComputeMlir"(%2) {recv_key = "host_compute_channel_0_0_retvals", send_key = "host_compute_channel_0_0_args", tpu_core = 0 : i64} : (tensor<f32>) -> tensor<f32>
%4 = "tf.AddV2"(%2, %3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
tf_device.return %4 : tensor<f32>
}) {device_assignment = [], num_cores_per_replica = 1 : i64, topology = ""} : () -> tensor<f32>
tf_device.return %1 : tensor<f32>
}) : () -> tensor<f32>
return %0 : tensor<f32>
}
```
}];
let constructor = "tensorflow::tf2xla::internal::CreateExtractOutsideCompilationPass()";
}
def XlaClusterFormationPass : Pass<"tf-xla-cluster-formation", "ModuleOp"> {
let summary = "Encapsulate partitioned calls within a Cluster op";
let description = [{
This pass clusters `tf.PartitionedCall` and `tf.StatefulPartitionedCall`
with `_xla_compile_device_type` attribute into a `tf_device.cluster`.
Notice this pass will only rewrite the outermost call if there are nested
calls to avoid nested `tf.XlaLaunch` operations from being created later.
For example, the following code
```mlir
func.func @main() -> tensor<i32> {
%0 = "tf.StatefulPartitionedCall"() {_xla_compile_device_type = "CPU", f = @stateful_pcall_func} : () -> (tensor<i32>)
func.return %0 : tensor<i32>
}
func.func @stateful_pcall_func() -> tensor<i32> {
%0 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
func.return %0 : tensor<i32>
}
```
will be transformed into,
```mlir
func.func @main() -> tensor<i32> {
%0 = "tf_device.cluster"() ({
%1 = "tf.StatefulPartitionedCall"() {_xla_compile_device_type = "CPU", f = @stateful_pcall_func} : () -> tensor<i32>
tf_device.return %1 : tensor<i32>
}) : () -> tensor<i32>
func.return %0 : tensor<i32>
}
func.func @stateful_pcall_func() -> tensor<i32> {
%0 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
func.return %0 : tensor<i32>
}
```
}];
let constructor = "tensorflow::tf2xla::internal::CreateXlaClusterFormationPass()";
let dependentDialects = ["mlir::tf_device::TensorFlowDeviceDialect"];
}
def MarkOpsForOutsideCompilationPass : Pass<"tf-mark-ops-for-outside-compilation", "ModuleOp"> {
let summary = "Marks ops in device cluster for outside compilation if they are unsupported on device.";
let description = [{
This pass marks unsupported ops in a device cluster with
`_xla_outside_compilation` attribute so the operations will run on the host
instead of the device. Unsupported ops are ops that can not be code
generated to run on the device for the cluster including:
1. String operations on TPUs.
2. Operations that don't have a kernel defined for the device.
This pass is conservative in that it will mark all ops for outside compilation
that can not be compiled for the device. Exceptions for this are added for ops
that will be rewritten or decomposed before compiling on device.
For example, tf_device.cluster op with an unsupported op, tf.UnsupportedOp:
```mlir
func @unsupported_op() -> tensor<i32> {
%0 = "tf_device.cluster"() ( {
%1 = "tf.UnsupportedOp"() : () -> tensor<i32>
%2 = "tf.Identity"(%1) : (tensor<i32>) -> tensor<i32>
tf_device.return %2 : tensor<i32>
}) {allow_soft_placement = true, num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> tensor<i32>
return %0 : tensor<i32>
}
```
will mark tf.UnsupportedOp with `_xla_outside_compilation` attribute:
```mlir
func @unsupported_op() -> tensor<i32> {
%0 = "tf_device.cluster"() ( {
%1 = "tf.UnsupportedOp"() {_xla_outside_compilation = "auto0"} : () -> tensor<i32>
%2 = "tf.Identity"(%1) : (tensor<i32>) -> tensor<i32>
tf_device.return %2 : tensor<i32>
}) {allow_soft_placement = true, device_assignment = [], num_cores_per_replica = 1 : i64, topology = ""} : () -> tensor<i32>
return %0 : tensor<i32>
}
```
}];
let constructor = "tensorflow::tf2xla::internal::CreateMarkOpsForOutsideCompilationPass()";
}
def HoistBroadcastReadPass : Pass<"tf-hoist-broadcast-read", "mlir::func::FuncOp"> {
let summary = "Hoist reads out of a replicate that are on a resource that is broacast to all replicas.";
let description = [{
Some `ReadVariableOp`s that are within a `tf_device.replicate` read the same
value across all replicas. These reads can be hoisted out of the
`tf_device.replicate` so there's one read for all replicas, and each replica
depends on the result of the read. This transform enables the
xla-broadcast-pass to optimize the broadcast value.
For example, the following:
```mlir
tf_device.replicate {n = 2 : i32} {
%0 = "tf.ReadVariableOp"(%arg0) : (tensor<*x!tf_type.resource<tensor<f32>>>) -> tensor<f32>
"tf.OpA"(%0) : (tensor<f32>) -> ()
}
```
will be transformed into:
``mlir
%0 = "tf.ReadVariableOp"(%arg0) : (tensor<*x!tf_type.resource<tensor<f32>>>) -> tensor<f32>
tf_device.replicate {n = 2 : i32} {
"tf.OpA"(%0) : (tensor<f32>) -> ()
}
```
We must ensure that there is a single underlying resource that not
distributed across replicas. There is a single underlying resource when the
resource device type is CPU, so we cautiously only apply in this case.
To be cautious we never hoist a read that comes after a write to the same
resource.
}];
let constructor = "tensorflow::tf2xla::internal::CreateHoistBroadcastReadPass()";
}
def XlaBroadcastPass : Pass<"tf-xla-broadcast", "mlir::func::FuncOp"> {
let summary = "Moves a broadcast from host into XLA, encoded as XlaAllReduce";
let description = [{
This pass moves brodcasts from host TF ops into XLA. This enables use of
the inter-device network, which is faster than the inter-host network.
Broadcasts in XLA are encoded as XlaAllReduce. An all_reduce with all 0
inputs except for one real input produces the same result as a broadcast.
For example, the following:
```mlir
func.func @main(%arg0: tensor<f32>) -> () {
tf_device.replicate {n = 2 : i32} {
"tf_device.cluster"() ({
"tf.OpA"(%arg0) : (tensor<f32>) -> ()
tf_device.return
}) : () -> ()
}
func.return
}
```
will be transformed into
```mlir
func.func @main(%arg0: tensor<f32>) {
%elem = "tf.Const"() <{value = dense<0.000000e+00> : tensor<f32>}> : () -> tensor<f32>
%shape = "tf.Const"() <{value = dense<> : tensor<0xi64>}> : () -> tensor<0xi64>
%zero = "tf.Fill"(%shape, %elem) : (tensor<0xi64>, tensor<f32>) -> tensor<f32>
tf_device.replicate([%arg0, %zero] as %arg1: tensor<f32>) {n = 2 : i32} {
%arg1_id = "tf_device.launch"() <{device = "TPU_REPLICATED_HOST_0"}> ({
%1 = "tf.Identity"(%arg1) : (tensor<f32>) -> tensor<f32>
tf_device.return %1 : tensor<f32>
}) : () -> tensor<f32>
"tf_device.cluster"() ({
%group = "tf.Const"() <{value = dense<[[0, 1]]> : tensor<1x2xi32>}> : () -> tensor<1x2xi32>
%arg1_reduced = "tf.XlaAllReduce"(%arg1_id, %group) <{mode = "CrossReplica", reduce_op = "Add"}> : (tensor<f32>, tensor<1x2xi32>) -> tensor<f32>
"tf.OpA"(%arg1_reduced) : (tensor<f32>) -> ()
tf_device.return
}) : () -> ()
tf_device.return
}
return
}
```
}];
let constructor = "tensorflow::tf2xla::internal::CreateXlaBroadcastPass()";
let dependentDialects = ["mlir::tf_device::TensorFlowDeviceDialect"];
}
def TPUShardingIdentificationPass : Pass<"tf-tpu-sharding-identification", "ModuleOp"> {
let summary = "Identifies and handles inputs/outputs of TPU computation that is "
"sharded across logical cores.";
let constructor = "tensorflow::tf2xla::internal::CreateTPUShardingIdentificationPass()";
let description = [{
Bubbles up sharding configuration from `cluster_func` regions into
the attributes of `cluster_func`. This is done by parsing the
`XlaSharding` / `TPUPartitionedOutput` / `TPUPartitionedInput` ops inside
`cluster_func`.
For example, given the following `cluster_func` wrapping `func`:
```mlir
func @test(%arg0: tensor<*xi32>) {
"tf_device.cluster_func"(%arg0) {
func = @func,
step_marker_location = ""} : (tensor<*xi32>) -> tensor<*xi32>
return
}
func @func(%arg0: tensor<*xi32>) -> tensor<*xi32> {
%0 = "tf.XlaSharding"(%arg0) {_XlaSharding = "\01\02\03",
sharding = "\01\02\03"} : (tensor<*xi32>) -> tensor<*xi32>
%1 = "tf.A"(%0) : (tensor<*xi32>) -> (tensor<*xi32>)
return %1 : tensor<*xi32>
}
```
Now, cluster_func receives the following `*_sharding_configuration`
attributes, and `func` receives the mhlo.sharding attribute:
```mlir
func @test(%arg0: tensor<*xi32>) {
%0 = "tf_device.cluster_func"(%arg0) {
func = @func,
input_sharding_configuration = ["\01\02\03"],
output_sharding_configuration = ["\08\01\1A\01\01\22\01\00"],
step_marker_location = ""} : (tensor<*xi32>) -> tensor<*xi32>
return
}
func @func(%arg0: tensor<*xi32> {mhlo.sharding = "\01\02\03"}) ->
(tensor<*xi32> {mhlo.sharding = "\08\01\1A\01\01\22\01\00"}) {
%0 = "tf.XlaSharding"(%arg0) {_XlaSharding = "\01\02\03", sharding = "\01\02\03"} : (tensor<*xi32>) -> tensor<*xi32>
%1 = "tf.A"(%0) : (tensor<*xi32>) -> tensor<*xi32>
return %1 : tensor<*xi32>
}
```
}];
}
def TPUValidateSessionInputsPass : Pass<"tf-tpu-validate-session-inputs", "mlir::ModuleOp"> {
let summary = "Validates inputs to the TPU TF/XLA bridge in session api";
let description = [{
This pass checks that the IR has valid input to TPU TF/XLA bridge in session api.
It checks the relations of multiple ops. Properties of single ops are
checked by the 'verify' method of ops.
}];
let constructor = "tensorflow::tf2xla::internal::CreateTPUValidateSessionInputsPass()";
}
def TPUValidateInputsPass : Pass<"tf-tpu-validate-inputs", "mlir::ModuleOp"> {
let summary = "Validates inputs to the TPU TF/XLA bridge";
let description = [{
This pass checks that the IR has valid input to TPU TF/XLA bridge.
It checks the relations of multiple ops. Properties of single ops are
checked by the 'verify' method of ops.
}];
let constructor = "tensorflow::tf2xla::internal::CreateTPUValidateInputsPass()";
}
tensorflow/compiler/mlir/tf2xla/internal/passes/extract_head_tail_outside_compilation.cc
+492
View File
@@ -0,0 +1,492 @@
/* Copyright 2020 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 <memory>
#include <string>
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Block.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Transforms/RegionUtils.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_structs.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/device_util.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/tpu_rewrite_device_util.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
// This pass extracts a CPU computation cluster with `_xla_outside_compilation`
// annotation from the head or tail of a TPU cluster.
namespace {
using mlir::Block;
using mlir::BlockArgument;
using mlir::BoolAttr;
using mlir::ModuleOp;
using mlir::OpBuilder;
using mlir::Operation;
using mlir::OperationPass;
using mlir::Region;
using mlir::StringAttr;
using mlir::Type;
using mlir::Value;
using mlir::WalkResult;
using mlir::func::FuncOp;
constexpr char kXlaMapOutsideCompilationAttr[] = "_xla_map_outside_compilation";
constexpr char kXlaOutsideCompilationAttr[] = "_xla_outside_compilation";
// Return true if `op` has attributes that say it can be outside compiled by
// this pass. This pass ignores _xla_map_outside_compilation, which will only be
// handled by extract_outside_compilation pass.
bool HasOutsideCompilationAttribute(Operation* op) {
return op->getAttrOfType<StringAttr>(kXlaOutsideCompilationAttr) != nullptr &&
!op->hasAttrOfType<BoolAttr>(kXlaMapOutsideCompilationAttr);
}
// Finds op that created a given value. If the value is a BlockArgument, this
// returns the owner of the Block.
Operation* GetOpOfValue(Value value) {
if (auto block_arg = mlir::dyn_cast<BlockArgument>(value))
return block_arg.getOwner()->getParentOp();
return value.getDefiningOp();
}
// Checks if `op` is nested in `block`.
bool OpInBlock(Operation* op, Block* block) {
Block* op_block = op->getBlock();
while (op_block) {
if (op_block == block) return true;
if (auto* parent_op = op_block->getParentOp()) {
op_block = parent_op->getBlock();
} else {
break;
}
}
return false;
}
// Wraps block in a Launch. External uses of ops in the block will be return
// values of the Launch and remapped to the Launch results. If `before` is set
// to true, the Launch is created before `op`. Otherwise the Launch is created
// after `op`.
mlir::tf_device::LaunchOp CreateLaunchForBlock(OpBuilder* builder,
Operation* op, bool before,
Block* launch_block,
llvm::StringRef host_device) {
// Find results and result types of ops in block that needs to returned.
llvm::SmallVector<Value, 4> launch_results;
llvm::SmallVector<Type, 4> launch_result_types;
for (Operation& head_outside_compiled_op : *launch_block) {
for (Value result : head_outside_compiled_op.getResults()) {
bool has_external_uses = false;
for (Operation* user : result.getUsers()) {
if (OpInBlock(user, launch_block)) continue;
has_external_uses = true;
break;
}
if (has_external_uses) {
launch_results.push_back(result);
launch_result_types.push_back(result.getType());
}
}
}
before ? builder->setInsertionPoint(op) : builder->setInsertionPointAfter(op);
auto launch = mlir::tf_device::LaunchOp::create(
*builder, op->getLoc(), builder->getStringAttr(host_device),
launch_result_types);
launch.getBody().push_back(launch_block);
builder->setInsertionPointToEnd(&launch.GetBody());
mlir::tf_device::ReturnOp::create(*builder, op->getLoc(), launch_results);
return launch;
}
// Checks if an operation is a supported TPU embedding op.
bool IsEmbeddingOp(Operation* op) {
return llvm::isa<mlir::TF::EnqueueTPUEmbeddingRaggedTensorBatchOp,
mlir::TF::EnqueueTPUEmbeddingSparseTensorBatchOp,
mlir::TF::EnqueueTPUEmbeddingArbitraryTensorBatchOp,
mlir::TF::RecvTPUEmbeddingActivationsOp,
mlir::TF::SendTPUEmbeddingGradientsOp>(op);
}
// Returns a set of ops that are outside compiled and can be extracted to before
// the TPU computation. These ops are either connected to the inputs of the TPU
// computation or other ops that can be extracted, and have no operands from
// other ops in the TPU computation that cannot be extracted.
llvm::SmallVector<Operation*, 4> FindOutsideCompiledOpsAtHead(
const mlir::TF::SideEffectAnalysis& side_effect_analysis,
mlir::tf_device::ClusterOp cluster) {
const auto& analysis = side_effect_analysis.GetAnalysisForFunc(
cluster->getParentOfType<mlir::func::FuncOp>());
Region* cluster_region = &cluster.getBody();
llvm::SmallSetVector<Operation*, 4> head_outside_compiled_ops;
auto cluster_ops = cluster.GetBody().without_terminator();
for (Operation& cluster_op : cluster_ops) {
if (!HasOutsideCompilationAttribute(&cluster_op)) continue;
// An outside compiled op can be extracted if its operands are not from
// other ops in the cluster that cannot be extracted.
// Check if the side effecting op right before this side effecting op, if
// it is side effecting, can be head extracted. Because of op ordering due
// to side effects, if this is not true, this op cannot be head extracted.
// TODO(lyandy): Remove special handling of embedding ops. Currently the IR
// is in a topological sort order and depending on that ordering, embedding
// ops may prevent other ops from being head extracted.
auto predecessors = analysis.DirectControlPredecessors(&cluster_op);
if (!predecessors.empty() && !IsEmbeddingOp(&cluster_op)) {
bool skip = false;
for (Operation* predecessor : llvm::reverse(predecessors)) {
if (IsEmbeddingOp(predecessor)) continue;
skip = !head_outside_compiled_ops.contains(predecessor);
break;
}
if (skip) continue;
}
auto walk_result = cluster_op.walk([&](Operation* op) {
for (Value operand : op->getOperands()) {
Operation* operand_op = GetOpOfValue(operand);
if (head_outside_compiled_ops.count(operand_op) ||
operand_op == &cluster_op)
continue;
if (operand_op->getParentRegion() == cluster_region)
return WalkResult::interrupt();
}
return WalkResult::advance();
});
if (!walk_result.wasInterrupted())
head_outside_compiled_ops.insert(&cluster_op);
}
return head_outside_compiled_ops.takeVector();
}
// Moves head outside compiled ops into its own `tf_device.LaunchOp`
// computation before the cluster.
void CreateHeadComputation(OpBuilder* builder,
mlir::tf_device::ClusterOp cluster,
llvm::ArrayRef<Operation*> head_outside_compiled_ops,
llvm::StringRef host_device) {
Block* launch_block = new Block;
for (Operation* head_outside_compiled_op : head_outside_compiled_ops) {
head_outside_compiled_op->removeAttr(kXlaOutsideCompilationAttr);
head_outside_compiled_op->moveBefore(launch_block, launch_block->end());
}
mlir::tf_device::LaunchOp launch = CreateLaunchForBlock(
builder, cluster, /*before=*/true, launch_block, host_device);
for (auto result : llvm::zip(launch.GetBody().getTerminator()->getOperands(),
launch.getResults()))
replaceAllUsesInRegionWith(std::get<0>(result), std::get<1>(result),
cluster.getBody());
}
// Extracts and move outside compiled ops that have no dependencies in the
// cluster to before the cluster.
mlir::LogicalResult LiftHeadOutsideCompiledOps(
OpBuilder* builder,
const mlir::TF::SideEffectAnalysis& side_effect_analysis,
const mlir::TF::RuntimeDevices& devices, mlir::tf_device::ClusterOp cluster,
std::string* host_device, bool* cluster_updated) {
llvm::SmallVector<Operation*, 4> head_outside_compiled_ops =
FindOutsideCompiledOpsAtHead(side_effect_analysis, cluster);
if (head_outside_compiled_ops.empty()) return mlir::success();
if (failed(tensorflow::GetHostDeviceOutsideComputation(devices, cluster,
host_device)))
return mlir::failure();
CreateHeadComputation(builder, cluster, head_outside_compiled_ops,
*host_device);
*cluster_updated = true;
return mlir::success();
}
// Fills `tail_outside_compiled_ops` with ops that are outside compiled and
// can be extracted to after the TPU computation, and `cluster_results` with new
// results of the cluster. These ops are either connected to the output of the
// TPU computation or other ops that can be extracted, and have no results used
// by other ops in the TPU computation that cannot be extracted.
void FindOutsideCompiledOpsAtTailAndClusterResults(
const mlir::TF::SideEffectAnalysis& side_effect_analysis,
mlir::tf_device::ClusterOp cluster,
llvm::SmallVectorImpl<Operation*>* tail_outside_compiled_ops,
llvm::SmallVectorImpl<Value>* cluster_results) {
const auto& analysis = side_effect_analysis.GetAnalysisForFunc(
cluster->getParentOfType<FuncOp>());
Region* cluster_region = &cluster.getBody();
llvm::SmallSetVector<Operation*, 4> tail_outside_compiled_ops_set;
Operation* terminator = cluster.GetBody().getTerminator();
llvm::SmallSetVector<Value, 4> cluster_results_set;
cluster_results_set.insert(terminator->getOperands().begin(),
terminator->getOperands().end());
auto cluster_ops = llvm::reverse(cluster.GetBody().without_terminator());
for (Operation& cluster_op : cluster_ops) {
if (!HasOutsideCompilationAttribute(&cluster_op)) continue;
// Check if the side effecting op right after this side effecting op, if
// it is side effecting, can be tail extracted. Because of op ordering due
// to side effects, if this is not true, this op cannot be tail extracted.
// TODO(lyandy): Remove special handling of embedding ops. Currently the IR
// is in a topological sort order and depending on that ordering, embedding
// ops may prevent other ops from being tail extracted.
auto successors = analysis.DirectControlSuccessors(
&cluster_op, [&terminator](Operation* op) { return op != terminator; });
if (!successors.empty() && !IsEmbeddingOp(&cluster_op)) {
bool skip = false;
for (Operation* successor : successors) {
if (IsEmbeddingOp(successor)) continue;
skip = !tail_outside_compiled_ops_set.contains(successor);
break;
}
if (skip) continue;
}
llvm::SmallVector<int, 4> results_to_forward;
bool can_be_extracted =
llvm::all_of(cluster_op.getUsers(), [&](Operation* op) {
return op == terminator || tail_outside_compiled_ops_set.count(op);
});
if (!can_be_extracted) continue;
// Collect operands of cluster op that are generated within the cluster.
// These values should be returned by the cluster.
cluster_op.walk([&](Operation* op) {
for (Value operand : op->getOperands()) {
Operation* operand_op = GetOpOfValue(operand);
if (operand_op->getParentRegion() == cluster_region)
cluster_results_set.insert(operand);
}
});
// Remove results of op to be extracted as there are no uses in the cluster.
for (Value result : cluster_op.getResults())
cluster_results_set.remove(result);
// Insert all ops including nested ops for checking outputs/side effects.
cluster_op.walk(
[&](Operation* op) { tail_outside_compiled_ops_set.insert(op); });
// Only add top level ops to output vector.
tail_outside_compiled_ops->push_back(&cluster_op);
}
*cluster_results = cluster_results_set.takeVector();
}
// Moves tail outside compiled ops into its own `tf_device.LaunchOp`
// computation after the cluster.
void CreateTailComputation(OpBuilder* builder,
mlir::tf_device::ClusterOp cluster,
llvm::ArrayRef<Operation*> tail_outside_compiled_ops,
llvm::StringRef host_device) {
Block* launch_block = new Block;
for (Operation* tail_outside_compiled_op : tail_outside_compiled_ops) {
tail_outside_compiled_op->removeAttr(kXlaOutsideCompilationAttr);
tail_outside_compiled_op->moveBefore(launch_block, launch_block->begin());
}
mlir::tf_device::LaunchOp launch = CreateLaunchForBlock(
builder, cluster, /*before=*/false, launch_block, host_device);
auto operand_not_in_launch = [&](mlir::OpOperand& operand) {
return !launch.getOperation()->isProperAncestor(operand.getOwner());
};
for (auto result : llvm::zip(launch.GetBody().getTerminator()->getOperands(),
launch.getResults()))
std::get<0>(result).replaceUsesWithIf(std::get<1>(result),
operand_not_in_launch);
}
// Updates cluster with updated cluster results after extracting tail outside
// compiled ops.
mlir::tf_device::ClusterOp UpdateClusterResults(
OpBuilder* builder, mlir::tf_device::ClusterOp cluster,
llvm::ArrayRef<Value> new_cluster_results) {
Operation* old_terminator = cluster.GetBody().getTerminator();
builder->setInsertionPoint(old_terminator);
mlir::tf_device::ReturnOp::create(*builder, old_terminator->getLoc(),
new_cluster_results);
old_terminator->erase();
builder->setInsertionPoint(cluster);
llvm::SmallVector<Type, 4> new_cluster_result_types;
new_cluster_result_types.reserve(new_cluster_results.size());
for (const auto& new_cluster_result : new_cluster_results)
new_cluster_result_types.push_back(new_cluster_result.getType());
auto new_cluster = mlir::tf_device::ClusterOp::create(
*builder, cluster.getLoc(), new_cluster_result_types,
/*operands=*/llvm::ArrayRef<Value>{}, cluster->getAttrs());
new_cluster.getBody().takeBody(cluster.getBody());
auto operand_not_in_cluster = [&](mlir::OpOperand& operand) {
return !new_cluster.getOperation()->isProperAncestor(operand.getOwner());
};
for (auto result :
llvm::zip(new_cluster.GetBody().getTerminator()->getOperands(),
new_cluster.getResults()))
std::get<0>(result).replaceUsesWithIf(std::get<1>(result),
operand_not_in_cluster);
cluster.erase();
return new_cluster;
}
// Extracts and move outside compiled ops that do not create dependencies in the
// cluster to after the cluster.
mlir::LogicalResult LiftTailOutsideCompiledOps(
OpBuilder* builder,
const mlir::TF::SideEffectAnalysis& side_effect_analysis,
const mlir::TF::RuntimeDevices& devices, std::string host_device,
mlir::tf_device::ClusterOp* cluster, bool* cluster_updated) {
llvm::SmallVector<Operation*, 4> tail_outside_compiled_ops;
llvm::SmallVector<Value, 4> cluster_results;
FindOutsideCompiledOpsAtTailAndClusterResults(side_effect_analysis, *cluster,
&tail_outside_compiled_ops,
&cluster_results);
if (tail_outside_compiled_ops.empty()) return mlir::success();
if (host_device.empty())
if (failed(tensorflow::GetHostDeviceOutsideComputation(devices, *cluster,
&host_device)))
return mlir::failure();
// Forward all results of cluster first. These results will be remapped once
// a new cluster is formed.
cluster->replaceAllUsesWith(
cluster->GetBody().getTerminator()->getOperands());
CreateTailComputation(builder, *cluster, tail_outside_compiled_ops,
host_device);
*cluster = UpdateClusterResults(builder, *cluster, cluster_results);
*cluster_updated = true;
return mlir::success();
}
// Removes aliased outputs in cluster from ops outside of cluster.
void RemoveClusterAliasedOutputs(OpBuilder* builder,
mlir::tf_device::ClusterOp cluster) {
llvm::SmallVector<Value, 4> used_old_cluster_results;
llvm::SmallVector<Value, 4> new_cluster_results;
llvm::SmallVector<Type, 4> new_cluster_result_types;
Operation* cluster_terminator = cluster.GetBody().getTerminator();
for (auto result :
llvm::zip(cluster_terminator->getOperands(), cluster.getResults())) {
Value cluster_terminator_operand = std::get<0>(result);
if (cluster_terminator_operand.getDefiningOp() &&
cluster.getOperation()->isProperAncestor(
cluster_terminator_operand.getDefiningOp())) {
new_cluster_results.push_back(cluster_terminator_operand);
new_cluster_result_types.push_back(cluster_terminator_operand.getType());
used_old_cluster_results.push_back(std::get<1>(result));
} else {
std::get<1>(result).replaceAllUsesWith(cluster_terminator_operand);
}
}
if (new_cluster_results.size() == cluster.getNumResults()) return;
builder->setInsertionPoint(cluster);
auto new_cluster = mlir::tf_device::ClusterOp::create(
*builder, cluster.getLoc(), new_cluster_result_types,
/*operands=*/llvm::ArrayRef<Value>{}, cluster->getAttrs());
new_cluster.getBody().takeBody(cluster.getBody());
new_cluster.GetBody().getTerminator()->setOperands(new_cluster_results);
for (auto result :
llvm::zip(used_old_cluster_results, new_cluster.getResults()))
std::get<0>(result).replaceAllUsesWith(std::get<1>(result));
cluster.erase();
}
#define GEN_PASS_DEF_EXTRACTHEADTAILOUTSIDECOMPILATIONPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc"
struct ExtractHeadTailOutsideCompilationPass
: public impl::ExtractHeadTailOutsideCompilationPassBase<
ExtractHeadTailOutsideCompilationPass> {
void runOnOperation() override;
};
void ExtractHeadTailOutsideCompilationPass::runOnOperation() {
auto& side_effect_analysis = getAnalysis<mlir::TF::SideEffectAnalysis>();
// Get runtime devices information from the closest parent module.
auto module = getOperation();
mlir::TF::RuntimeDevices devices;
if (failed(tensorflow::GetDevicesFromOp(module, &devices)))
return signalPassFailure();
OpBuilder builder(&getContext());
llvm::SmallVector<mlir::tf_device::ClusterOp, 4> clusters;
module.walk(
[&](mlir::tf_device::ClusterOp cluster) { clusters.push_back(cluster); });
for (mlir::tf_device::ClusterOp cluster : clusters) {
std::string host_device;
bool cluster_updated = false;
if (failed(LiftHeadOutsideCompiledOps(&builder, side_effect_analysis,
devices, cluster, &host_device,
&cluster_updated)) ||
failed(LiftTailOutsideCompiledOps(&builder, side_effect_analysis,
devices, host_device, &cluster,
&cluster_updated)))
return signalPassFailure();
if (cluster_updated) RemoveClusterAliasedOutputs(&builder, cluster);
}
}
} // anonymous namespace
std::unique_ptr<OperationPass<ModuleOp>>
CreateExtractHeadTailOutsideCompilationPass() {
return std::make_unique<ExtractHeadTailOutsideCompilationPass>();
}
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/passes/extract_outside_compilation.cc
+1593
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/mlir/tf2xla/internal/passes/hoist_broadcast_read.cc
+154
View File
@@ -0,0 +1,154 @@
/* 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 <memory>
#include <string>
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Pass/Pass.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_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/core/util/device_name_utils.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {
using mlir::BlockArgument;
using mlir::failure;
using mlir::LogicalResult;
using mlir::Operation;
using mlir::OperationPass;
using mlir::OpOperand;
using mlir::StringAttr;
using mlir::success;
using mlir::Value;
using mlir::WalkResult;
using mlir::func::FuncOp;
using mlir::TF::ReadVariableOp;
using mlir::tf_device::ReplicateOp;
#define GEN_PASS_DEF_HOISTBROADCASTREADPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc"
constexpr char kFuncDeviceAttr[] = "tf.device";
constexpr char kCpuDeviceType[] = "CPU";
struct HoistBroadcastRead
: public impl::HoistBroadcastReadPassBase<HoistBroadcastRead> {
void runOnOperation() override;
};
// Get the ancestor of `descendant` that is a direct child of `ancestor`.
Operation* GetAncestorBelow(Operation* descendant, Operation* ancestor) {
Operation* parent = descendant->getParentOp();
if (!parent) return nullptr;
if (parent == ancestor) return descendant;
return GetAncestorBelow(parent, ancestor);
}
// `is_cpu_read` is set to `true` iff `read` is on a resource with device type
// CPU.
LogicalResult IsCpuRead(FuncOp func, ReadVariableOp read, bool& is_cpu_read) {
if (auto arg = mlir::dyn_cast<BlockArgument>(read->getOperand(0))) {
if (arg.getOwner() != &(func.front())) {
is_cpu_read = false;
return success();
}
if (auto attr = func.getArgAttrOfType<StringAttr>(arg.getArgNumber(),
kFuncDeviceAttr)) {
std::string device = attr.getValue().str();
tensorflow::DeviceNameUtils::ParsedName parsed_name;
if (!tensorflow::DeviceNameUtils::ParseFullName(device, &parsed_name)) {
return read->emitOpError() << "invalid device '" << device << "'";
}
is_cpu_read = parsed_name.type == kCpuDeviceType;
return success();
}
}
is_cpu_read = false;
return success();
}
// Get the reads to hoist in the `replicate`.
LogicalResult GetReads(FuncOp func, ReplicateOp replicate,
llvm::SmallVector<ReadVariableOp, 4>& reads) {
for (Operation& op : replicate.getBody().front()) {
if (auto read = llvm::dyn_cast<ReadVariableOp>(&op)) {
bool is_cpu_read;
if (failed(IsCpuRead(func, read, is_cpu_read))) return failure();
if (is_cpu_read) reads.push_back(read);
}
}
return success();
}
// Move reads above the `replicate`. Skip reads that come after a write to the
// same resource.
void MoveReads(ReplicateOp replicate,
llvm::SmallVector<ReadVariableOp, 4>& reads) {
for (ReadVariableOp read : reads) {
Value res = read.getResource();
Operation* scope = res.getParentBlock()->getParentOp();
if (!scope->isProperAncestor(replicate)) continue;
bool has_conflicting_write = false;
for (OpOperand& use : res.getUses()) {
Operation* using_op = use.getOwner();
if (using_op == read) continue;
if (!replicate->isProperAncestor(using_op)) continue;
Operation* peer = GetAncestorBelow(using_op, replicate);
if (read->isBeforeInBlock(peer)) continue;
if (llvm::isa<ReadVariableOp>(peer)) continue;
has_conflicting_write = true;
}
if (has_conflicting_write) continue;
read->moveBefore(replicate);
}
}
// Hoist `ReadVariableOp`s above the `tf_device.replicate`s.
void HoistBroadcastRead::runOnOperation() {
FuncOp func = getOperation();
auto result = func.walk([&](ReplicateOp replicate) {
llvm::SmallVector<ReadVariableOp, 4> reads;
if (failed(GetReads(func, replicate, reads)))
return WalkResult::interrupt();
MoveReads(replicate, reads);
return WalkResult::advance();
});
if (result.wasInterrupted()) return signalPassFailure();
}
} // namespace
std::unique_ptr<OperationPass<FuncOp>> CreateHoistBroadcastReadPass() {
return std::make_unique<HoistBroadcastRead>();
}
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/passes/input_lowering_metrics_pass.cc
+88
View File
@@ -0,0 +1,88 @@
/* 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 <memory>
#include "llvm/ADT/DenseSet.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/legalization_op_config.h"
#include "tensorflow/core/lib/monitoring/counter.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {
using mlir::Operation;
using mlir::WalkResult;
#define GEN_PASS_DEF_INPUTLOWERINGMETRICSPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/lowering_passes.h.inc"
auto* dynamism_op_counter = tensorflow::monitoring::Counter<1>::New(
"/tensorflow/core/tf2xla/api/v2/dynamism_op_counter",
"Counts how many ops are dynamic", "op_name");
auto* dynamism_function_counter = tensorflow::monitoring::Counter<1>::New(
"/tensorflow/core/tf2xla/api/v2/dynamism_function_counter",
"Counts how many functions are dynamic", "has_dynamism");
constexpr char kNotDynamicFunctionName[] = "kNotDynamicFunction";
constexpr char kDynamicFunctionName[] = "kDynamicFunction";
class InputMetricsLoweringPass
: public impl::InputLoweringMetricsPassBase<InputMetricsLoweringPass> {
public:
void runOnOperation() override;
};
void InputMetricsLoweringPass::runOnOperation() {
bool has_dynamic_op = false;
Operation* func_op = getOperation();
func_op->walk([&](Operation* op) {
auto abstractOp = op->getRegisteredInfo();
if (!abstractOp) return WalkResult::advance();
if (mlir::hlo::IsDynamicPadderOp(abstractOp->getTypeID())) {
has_dynamic_op = true;
dynamism_op_counter->GetCell(op->getName().getStringRef().str())
->IncrementBy(1);
}
return WalkResult::advance();
});
if (has_dynamic_op) {
dynamism_function_counter->GetCell(kDynamicFunctionName)->IncrementBy(1);
} else {
dynamism_function_counter->GetCell(kNotDynamicFunctionName)->IncrementBy(1);
}
}
} // namespace
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateInputLoweringMetricsPass() {
return std::make_unique<InputMetricsLoweringPass>();
}
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/passes/input_lowering_metrics_pass_test.cc
+124
View File
@@ -0,0 +1,124 @@
/* 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 <cstdint>
#include <memory>
#include <gtest/gtest.h>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.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 "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/lowering_passes.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/test_utils.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/lib/monitoring/cell_reader.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {
using ::mlir::LogicalResult;
using ::mlir::ModuleOp;
using ::mlir::hlo::test::GetMlirModuleFromString;
using ::tensorflow::monitoring::testing::CellReader;
constexpr char kNotDynamicFunctionName[] = "kNotDynamicFunction";
constexpr char kDynamicFunctionName[] = "kDynamicFunction";
static constexpr char kDynamismOpCounterStreamzName[] =
"/tensorflow/core/tf2xla/api/v2/dynamism_op_counter";
static constexpr char kDynamismFunctionCounterStreamzName[] =
"/tensorflow/core/tf2xla/api/v2/dynamism_function_counter";
class InputLoweringMetricsPassTest : public testing::Test {
protected:
void CreateModule(const char* module_string) {
TF_ASSERT_OK_AND_ASSIGN(module_,
GetMlirModuleFromString(module_string, &context_));
pm_ = std::make_unique<mlir::PassManager>(&context_);
pm_->addNestedPass<mlir::func::FuncOp>(CreateInputLoweringMetricsPass());
}
bool ModulesEqual(const ModuleOp& module_before,
const ModuleOp& module_after) {
return mlir::OperationEquivalence::isEquivalentTo(
module_before, module_after, mlir::OperationEquivalence::None);
}
mlir::LogicalResult Run() {
mlir::OwningOpRef<mlir::ModuleOp> module_before = module_->clone();
LogicalResult run_result = pm_->run(module_.get());
EXPECT_TRUE(ModulesEqual(*module_before, *module_));
return run_result;
}
private:
mlir::MLIRContext context_;
mlir::OwningOpRef<ModuleOp> module_;
std::unique_ptr<mlir::PassManager> pm_;
};
TEST_F(InputLoweringMetricsPassTest, CountsNoDynamicOps) {
static constexpr char kMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> tensor<1xi32> {
%0 = "tf.Const"() {value = dense<1000> : tensor<1xi32>} : () -> tensor<1xi32>
return %0 : tensor<1xi32>
}
})";
CellReader<int64_t> dynamism_op_counter(kDynamismOpCounterStreamzName);
CellReader<int64_t> dynamism_function_counter(
kDynamismFunctionCounterStreamzName);
CreateModule(kMlirModuleStr);
auto result = Run();
EXPECT_TRUE(result.succeeded());
EXPECT_EQ(dynamism_function_counter.Delta(kNotDynamicFunctionName), 1);
}
TEST_F(InputLoweringMetricsPassTest, CountsDynamicOps) {
static constexpr char kMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> () {
%cst0 = "tf.Const"(){ value = dense<0> : tensor<3x5xi1>} : () -> tensor<3x5xi1>
%0 = "tf.Where"(%cst0) : (tensor<3x5xi1>) -> tensor<?x2xi64>
func.return
}
})";
CellReader<int64_t> dynamism_counter(kDynamismOpCounterStreamzName);
CellReader<int64_t> dynamism_function_counter(
kDynamismFunctionCounterStreamzName);
CreateModule(kMlirModuleStr);
auto result = Run();
EXPECT_TRUE(result.succeeded());
EXPECT_EQ(dynamism_counter.Delta("tf.Where"), 1);
EXPECT_EQ(dynamism_function_counter.Delta(kDynamicFunctionName), 1);
}
} // namespace
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/passes/lowering_passes.h
+38
View File
@@ -0,0 +1,38 @@
/* 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_PASSES_LOWERING_PASSES_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_PASSES_LOWERING_PASSES_H_
#include <memory>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
namespace tensorflow {
namespace tf2xla {
namespace internal {
// Create a pass that just collects metrics about the input MLIR. Does not
// logically transform the program.
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateInputLoweringMetricsPass();
#define GEN_PASS_REGISTRATION
#define GEN_PASS_DECL_INPUTLOWERINGMETRICSPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/lowering_passes.h.inc"
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_PASSES_LOWERING_PASSES_H_
tensorflow/compiler/mlir/tf2xla/internal/passes/lowering_passes.td
+25
View File
@@ -0,0 +1,25 @@
/* Copyright 2024 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 "mlir/Pass/PassBase.td"
def InputLoweringMetricsPass : Pass<"input-lowering-metrics-pass", "mlir::func::FuncOp"> {
let summary = "Collects various metrics about the input to the lowering "
"portion of the bridge. This is a logical no-op.";
let description = [{
Gathers metrics about the input MLIR to Phase 2 of the TFXLA Bridge, which
does a strict semantic lowering from Tensorflow ops to XLA HLO.
}];
let constructor = "tensorflow::tf2xla::internal::CreateInputLoweringMetricsPass()";
}
tensorflow/compiler/mlir/tf2xla/internal/passes/mark_ops_for_outside_compilation.cc
+524
View File
@@ -0,0 +1,524 @@
/* 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 <memory>
#include <optional>
#include <queue>
#include <string>
#include <utility>
#include "absl/log/log.h"
#include "absl/strings/str_join.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/FormatVariadic.h"
#include "mlir/IR/Block.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/OperationSupport.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/IR/TypeUtilities.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Rewrite/PatternApplicator.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Transforms/RegionUtils.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/lower_tf.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/string_util.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/legalization_op_config.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/passes.h"
#include "tensorflow/core/lib/monitoring/gauge.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {
using mlir::Block;
using mlir::BoolAttr;
using mlir::Dialect;
using mlir::LogicalResult;
using mlir::MLIRContext;
using mlir::ModuleOp;
using mlir::Operation;
using mlir::OperationName;
using mlir::OperationPass;
using mlir::Pattern;
using mlir::PatternApplicator;
using mlir::RewritePatternSet;
using mlir::StringAttr;
using mlir::TensorType;
using mlir::Type;
using mlir::Value;
using mlir::WalkResult;
auto* auto_outside_compilation_gauge =
tensorflow::monitoring::Gauge<bool, 0>::New(
"/tensorflow/core/use_auto_outside_compilation",
"Tracks if auto outside compilation is enabled");
#define GEN_PASS_DEF_MARKOPSFOROUTSIDECOMPILATIONPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc"
struct MarkOpsForOutsideCompilation
: public impl::MarkOpsForOutsideCompilationPassBase<
MarkOpsForOutsideCompilation> {
void runOnOperation() override;
};
// Adds any canonicalization patterns to list of supported `patterns`.
// TODO(b/161726307): Move or import the relevant patterns to LowerTF pass and
// remove this.
void AddCanonicalizationPatterns(MLIRContext* context,
RewritePatternSet* patterns) {
for (auto op : context->getRegisteredOperations())
op.getCanonicalizationPatterns(*patterns, context);
}
// Adds the list of ops that are supported on TPU through constant folding which
// may depend on the inputs shapes not known at this point. Such ops may not
// have any legalization or canonicalization patterns but shouldn't be marked
// for outside compilation.
//
// TODO(b/177523289): Remove manual handling once we support constant folding
// and shape inference through the computation on the host side.
void AddSupportedOpsUsingFolding(MLIRContext* context,
llvm::DenseSet<OperationName>* supported_ops) {
llvm::SmallDenseSet<OperationName, 8> allowlist_ops = {
OperationName(mlir::TF::BroadcastArgsOp::getOperationName(), context),
OperationName(mlir::TF::BroadcastGradientArgsOp::getOperationName(),
context),
OperationName(mlir::TF::ConcatOffsetOp::getOperationName(), context),
OperationName(mlir::TF::EmptyOp::getOperationName(), context),
OperationName(mlir::TF::ListDiffOp::getOperationName(), context),
OperationName(mlir::TF::RankOp::getOperationName(), context),
OperationName(mlir::TF::RangeOp::getOperationName(), context),
OperationName(mlir::TF::ShapeOp::getOperationName(), context),
OperationName(mlir::TF::ShapeNOp::getOperationName(), context),
OperationName(mlir::TF::SizeOp::getOperationName(), context),
};
supported_ops->insert(allowlist_ops.begin(), allowlist_ops.end());
}
// Adds the list of ops that are only supported in the old bridge.
// TODO(b/168036682): Remove bounded dynamism ops now that MLIR bridge supports
// bounded dynamism.
// TODO(b/257574556): Remove the need for this manual list by making use of old
// bridge phase 2 op list.
void AddOldBridgeOnlyOps(MLIRContext* context,
llvm::DenseSet<OperationName>* supported_ops) {
llvm::SmallDenseSet<OperationName, 8> allowlist_ops = {
OperationName(mlir::TF::DynamicPartitionOp::getOperationName(), context),
OperationName(mlir::TF::OutfeedEnqueueOp::getOperationName(), context),
OperationName(mlir::TF::WhereOp::getOperationName(), context),
OperationName(mlir::TF::UniqueOp::getOperationName(), context),
OperationName(mlir::TF::XlaSetDynamicDimensionSizeOp::getOperationName(),
context),
OperationName(mlir::TF::XlaSpmdFullToShardShapeOp::getOperationName(),
context),
OperationName(mlir::TF::XlaSpmdShardToFullShapeOp::getOperationName(),
context),
};
supported_ops->insert(allowlist_ops.begin(), allowlist_ops.end());
}
// TODO(b/159128666): Check the control flow legalization passes instead once
// added.
void AddSupportedFunctionalOps(MLIRContext* context,
llvm::DenseSet<OperationName>* supported_ops) {
supported_ops->insert(
OperationName(mlir::TF::CaseRegionOp::getOperationName(), context));
supported_ops->insert(
OperationName(mlir::TF::IfRegionOp::getOperationName(), context));
supported_ops->insert(
OperationName(mlir::TF::InplaceAddOp::getOperationName(), context));
supported_ops->insert(
OperationName(mlir::TF::WhileRegionOp::getOperationName(), context));
supported_ops->insert(
OperationName(mlir::TF::XlaCallModuleOp::getOperationName(), context));
supported_ops->insert(
OperationName(mlir::TF::XlaHostComputeOp::getOperationName(), context));
supported_ops->insert(
OperationName(mlir::TF::XlaReduceOp::getOperationName(), context));
supported_ops->insert(
OperationName(mlir::TF::XlaReduceWindowOp::getOperationName(), context));
supported_ops->insert(OperationName(
mlir::TF::XlaRngBitGeneratorOp::getOperationName(), context));
supported_ops->insert(
OperationName(mlir::TF::XlaScatterOp::getOperationName(), context));
supported_ops->insert(OperationName(
mlir::TF::XlaSelectAndScatterOp::getOperationName(), context));
supported_ops->insert(
OperationName(mlir::TF::SymbolicGradientOp::getOperationName(), context));
supported_ops->insert(OperationName(
mlir::TF::XlaVariadicReduceOp::getOperationName(), context));
supported_ops->insert(OperationName(
mlir::TF::XlaVariadicReduceV2Op::getOperationName(), context));
supported_ops->insert(
OperationName(mlir::TF::XlaVariadicSortOp::getOperationName(), context));
supported_ops->insert(
OperationName(mlir::TF::XlaReplicaIdOp::getOperationName(), context));
supported_ops->insert(
OperationName(mlir::TF::YieldOp::getOperationName(), context));
}
// These embedding ops are rewritten when running TPUCompileOp.
void AddRewrittenEmbeddingOps(MLIRContext* context,
llvm::DenseSet<OperationName>* supported_ops) {
supported_ops->insert(OperationName(
mlir::TF::RecvTPUEmbeddingActivationsOp::getOperationName(), context));
supported_ops->insert(OperationName(
mlir::TF::SendTPUEmbeddingGradientsOp::getOperationName(), context));
}
// Stack, TensorList and TensorArray ops are rewritten during the second phase
// of the bridge (compilation of TPUCompile op). They would not match any
// legalization/canonicalization pattern and have to be manually added to the
// list of supported ops.
void AddRewrittenCompositeOps(MLIRContext* context,
llvm::DenseSet<OperationName>* supported_ops) {
#define GET_OPERATION_NAME(op) OperationName(op::getOperationName(), context)
llvm::SmallDenseSet<OperationName, 32> allowlist_ops = {
// Stack ops.
GET_OPERATION_NAME(mlir::TF::StackV2Op),
GET_OPERATION_NAME(mlir::TF::StackPushV2Op),
GET_OPERATION_NAME(mlir::TF::StackPopV2Op),
// Tensor Array ops.
GET_OPERATION_NAME(mlir::TF::TensorArrayV3Op),
GET_OPERATION_NAME(mlir::TF::TensorArrayReadV3Op),
GET_OPERATION_NAME(mlir::TF::TensorArrayWriteV3Op),
GET_OPERATION_NAME(mlir::TF::TensorArrayConcatV3Op),
GET_OPERATION_NAME(mlir::TF::TensorArraySplitV3Op),
GET_OPERATION_NAME(mlir::TF::TensorArraySizeV3Op),
GET_OPERATION_NAME(mlir::TF::TensorArrayGradV3Op),
GET_OPERATION_NAME(mlir::TF::TensorArrayGatherV3Op),
GET_OPERATION_NAME(mlir::TF::TensorArrayScatterV3Op),
// Tensor List Ops.
GET_OPERATION_NAME(mlir::TF::EmptyTensorListOp),
GET_OPERATION_NAME(mlir::TF::TensorListReserveOp),
GET_OPERATION_NAME(mlir::TF::TensorListFromTensorOp),
GET_OPERATION_NAME(mlir::TF::TensorListPushBackOp),
GET_OPERATION_NAME(mlir::TF::TensorListPopBackOp),
GET_OPERATION_NAME(mlir::TF::TensorListGetItemOp),
GET_OPERATION_NAME(mlir::TF::TensorListSetItemOp),
GET_OPERATION_NAME(mlir::TF::TensorListLengthOp),
GET_OPERATION_NAME(mlir::TF::TensorListElementShapeOp),
GET_OPERATION_NAME(mlir::TF::TensorListGatherOp),
GET_OPERATION_NAME(mlir::TF::TensorListScatterIntoExistingListOp),
GET_OPERATION_NAME(mlir::TF::TensorListStackOp),
};
#undef GET_OPERATION_NAME
supported_ops->insert(allowlist_ops.begin(), allowlist_ops.end());
}
bool IsStringType(Type type) {
if (mlir::isa<mlir::TF::StringType>(type)) return true;
auto sub_type = mlir::dyn_cast<mlir::TF::TensorFlowTypeWithSubtype>(type);
if (!sub_type) return false;
bool has_string = llvm::any_of(sub_type.GetSubtypes(), [](TensorType type) {
return mlir::isa<mlir::TF::StringType>(type.getElementType());
});
return has_string;
}
bool HasStringOperand(Operation& op) {
for (auto operand : op.getOperands()) {
auto operand_type = getElementTypeOrSelf(operand);
if (IsStringType(operand_type)) return true;
}
return false;
}
bool HasStringResult(Operation& op) {
for (auto result : op.getResults()) {
auto result_type = getElementTypeOrSelf(result);
if (IsStringType(result_type)) return true;
}
return false;
}
bool MatchesPattern(Operation& op,
const llvm::DenseSet<OperationName>& supported_ops) {
return (supported_ops.contains(op.getName()));
}
// Checks if the op is supported inside of a device cluster. Ops not
// in `tf_dialect` are considered supported.
bool IsSupportedOp(Operation& op,
const llvm::DenseSet<OperationName>& supported_ops,
const Dialect* tf_dialect) {
if (op.getDialect() != tf_dialect) return true;
// Assert has a legalization that later removes it so we don't want to outside
// compile it ever for performance reasons.
if (llvm::isa<mlir::TF::AssertOp>(op)) return true;
if (HasStringOperand(op)) return false;
if (HasStringResult(op)) return false;
if (MatchesPattern(op, supported_ops)) return true;
auto abstractOp = op.getRegisteredInfo();
if (!abstractOp) return false;
return mlir::hlo::HasTf2XlaFallback(abstractOp->getTypeID());
}
bool IsVariant(Value value) {
return mlir::isa<mlir::TF::VariantType>(
getElementTypeOrSelf(value.getType()));
}
bool HasOutsideCompiledAncestor(Operation* op) {
Operation* parent = op->getParentOp();
while (parent) {
if (parent->getAttrOfType<StringAttr>(kXlaOutsideCompilationAttr))
return true;
parent = parent->getParentOp();
}
return false;
}
// If any tf.variants are inputs/outputs to the another outside compiled
// Operation, `op`, mark them for outside compilation unless they are already
// marks with outside compilation attribute.
void MarkVariantInputsOutputs(mlir::tf_device::ClusterOp tpu_cluster) {
std::queue<Operation*> outside_compiled_ops;
tpu_cluster.walk([&](Operation* op) {
if (op->hasAttrOfType<StringAttr>(kXlaOutsideCompilationAttr))
outside_compiled_ops.push(op);
});
while (!outside_compiled_ops.empty()) {
Operation* host_op = outside_compiled_ops.front();
outside_compiled_ops.pop();
host_op->walk([&](Operation* op) {
// Add any operations that provide variant inputs to the cluster.
for (auto value : op->getOperands()) {
Operation* input_defining_op = value.getDefiningOp();
if (IsVariant(value) && input_defining_op &&
!HasOutsideCompiledAncestor(input_defining_op) &&
!input_defining_op->hasAttrOfType<StringAttr>(
kXlaOutsideCompilationAttr)) {
input_defining_op->setAttr(
kXlaOutsideCompilationAttr,
StringAttr::get(input_defining_op->getContext(), "auto"));
outside_compiled_ops.push(input_defining_op);
}
}
// Mark for outside compilation any operations that consume variant
// outputs from an outside compiled operation.
for (auto value : op->getResults()) {
if (IsVariant(value)) {
for (auto user : value.getUsers()) {
if (!user->hasTrait<mlir::OpTrait::IsTerminator>() &&
!HasOutsideCompiledAncestor(user) &&
!user->getAttrOfType<StringAttr>(kXlaOutsideCompilationAttr)) {
user->setAttr(kXlaOutsideCompilationAttr,
StringAttr::get(user->getContext(), "auto"));
outside_compiled_ops.push(user);
}
}
}
}
});
}
}
// Marks uncompilable ops that are in `tf_dialect` for outside compilation.
LogicalResult MarkUncompilableOps(
const Dialect* tf_dialect, Block* block,
llvm::DenseSet<OperationName>& supported_ops) {
// Automatically marked ops for outside compilation have
// `_xla_outside_compilation` attribute value of "auto" plus
// an increasing counter. Manually marked ops for outside compilation only
// have an increasing counteri for the attribute value. Therefore there is no
// collision in
// `_xla_outside_compilation` attribute between automatically and manually
// marking ops.
int outside_compiled_cluster_counter = 0;
block->walk([&](Operation* op) {
if (!IsSupportedOp(*op, supported_ops, tf_dialect)) {
VLOG(3) << "Cloud TPU: Op " << op->getName().getStringRef().str()
<< " isn't compilable, adding outside_compilation attr. "
"This op will automatically be placed on CPU.";
op->setAttr(kXlaOutsideCompilationAttr,
StringAttr::get(
op->getContext(),
llvm::formatv("auto{0}", outside_compiled_cluster_counter)
.str()));
outside_compiled_cluster_counter++;
}
});
if (outside_compiled_cluster_counter > 0) {
auto_outside_compilation_gauge->GetCell()->Set(true);
}
return mlir::success();
}
// Check for uncompilable ops that are in `tf_dialect` and are not already
// marked for outside compilation.
bool ContainsUncompilableOps(const Dialect* tf_dialect, Block* block,
llvm::DenseSet<OperationName>& supported_ops) {
int uncompilable_op_count = 0;
// Check if op or any parent is already marked for outside compilation.
block->walk([&](Operation* op) {
Operation* iter_op = op;
while (iter_op && !llvm::isa<mlir::tf_device::ClusterOp>(iter_op)) {
if (iter_op->hasAttrOfType<StringAttr>(kXlaOutsideCompilationAttr)) {
return;
}
iter_op = iter_op->getParentOp();
}
if (!IsSupportedOp(*op, supported_ops, tf_dialect)) {
op->emitOpError() << "isn't compilable for TPU device. enable "
"soft_device_placement option to run on CPU";
++uncompilable_op_count;
}
});
return uncompilable_op_count > 0;
}
// Unmarks outside compilation for any op that has parents already
// marked for outside compilation since the child will be extracted
// anyways.
void UnmarkChildren(ModuleOp module) {
module->walk([&](Operation* op) {
if (!op->getAttrOfType<StringAttr>(kXlaOutsideCompilationAttr)) return;
Operation* iter_op = op;
bool remove_attr = false;
while (auto* parent_op = iter_op->getParentOp()) {
if (parent_op->getAttrOfType<StringAttr>(kXlaOutsideCompilationAttr)) {
remove_attr = true;
break;
}
iter_op = parent_op;
}
if (remove_attr) op->removeAttr(kXlaOutsideCompilationAttr);
});
}
constexpr int kTooManyOutsideCompileRegionThreshold = 32;
constexpr int kOpDetailCount = 8;
void WarnOnExcessOutsideCompilationOps(ModuleOp module) {
// Count the number of outside compilation ops. If it exceeds the reporting
// threshold, warn the user that their model may run slowly.
llvm::SmallVector<Operation*, 8> outside_compile_ops;
module->walk([&](Operation* op) {
if (op->getAttrOfType<StringAttr>(kXlaOutsideCompilationAttr)) {
outside_compile_ops.push_back(op);
}
});
if (outside_compile_ops.size() > kTooManyOutsideCompileRegionThreshold) {
llvm::SmallVector<std::string, kOpDetailCount> op_info;
for (int i = 0; i < kOpDetailCount; ++i) {
auto& op = outside_compile_ops[i];
op_info.push_back(tensorflow::OpAsString(*op));
}
LOG(WARNING) << outside_compile_ops.size() << " outside compilation "
<< "regions found while processing "
<< module->getName().getStringRef().str()
<< ". This may result in excessively slow model execution. "
<< "First " << op_info.size()
<< " ops: " << absl::StrJoin(op_info, "\n");
} else {
LOG(INFO) << "Found " << outside_compile_ops.size()
<< " outside compilation regions.";
}
}
void MarkOpsForOutsideCompilation::runOnOperation() {
auto module = getOperation();
const Dialect* tf_dialect = getContext().getLoadedDialect("tf");
if (!tf_dialect) {
getOperation().emitError() << "'tf' dialect is not registered";
return signalPassFailure();
}
RewritePatternSet patterns(&getContext());
mlir::hlo::PopulateLegalizeTfPatterns(module.getContext(), &patterns);
mlir::TF::PopulateTFLoweringBeforeHLOPatterns(module.getContext(), &patterns);
mlir::TF::PopulateLoweringQuantizedPatterns(module.getContext(), &patterns);
AddCanonicalizationPatterns(module.getContext(), &patterns);
// `supported_ops` contains the name of all of the ops that can potentially be
// lowered into HLO on the device. This doesn't always mean that the op can
// be lowered in the future passes but if the op is not in this set, it can't
// be lowered in a subsequent pass.
llvm::DenseSet<OperationName> supported_ops;
PatternApplicator(std::move(patterns))
.walkAllPatterns([&](const Pattern& pattern) {
std::optional<OperationName> root_kind = pattern.getRootKind();
if (root_kind.has_value()) supported_ops.insert(root_kind.value());
});
AddSupportedFunctionalOps(module.getContext(), &supported_ops);
AddSupportedOpsUsingFolding(module.getContext(), &supported_ops);
AddOldBridgeOnlyOps(module.getContext(), &supported_ops);
AddRewrittenEmbeddingOps(module.getContext(), &supported_ops);
AddRewrittenCompositeOps(module.getContext(), &supported_ops);
auto result = module.walk([&](mlir::tf_device::ClusterOp cluster) {
// Only if `allow_soft_placement` attribute is true should we mark ops
// for outside compilation.
auto soft_placement_attr =
cluster->getAttrOfType<BoolAttr>(mlir::TF::kAllowSoftPlacementAttr);
if ((soft_placement_attr && soft_placement_attr.getValue())) {
if (failed(MarkUncompilableOps(tf_dialect, &cluster.GetBody(),
supported_ops)))
return WalkResult::interrupt();
} else {
if (ContainsUncompilableOps(tf_dialect, &cluster.GetBody(),
supported_ops))
return WalkResult::interrupt();
}
MarkVariantInputsOutputs(cluster);
return WalkResult::advance();
});
if (result.wasInterrupted()) return signalPassFailure();
UnmarkChildren(module);
WarnOnExcessOutsideCompilationOps(module);
}
} // namespace
std::unique_ptr<OperationPass<ModuleOp>>
CreateMarkOpsForOutsideCompilationPass() {
return std::make_unique<MarkOpsForOutsideCompilation>();
}
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/passes/mlir_to_graph_passes.h
+35
View File
@@ -0,0 +1,35 @@
/* 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_PASSES_MLIR_TO_GRAPH_PASSES_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_PASSES_MLIR_TO_GRAPH_PASSES_H_
#include <memory>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
namespace tensorflow {
namespace tf2xla {
namespace internal {
// Verifies that Executor input is of the expected format.
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateVerifyInputDialectToExecutorPass();
#define GEN_PASS_REGISTRATION
#define GEN_PASS_DECL_VERIFYINPUTDIALECTTOEXECUTORPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/mlir_to_graph_passes.h.inc"
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_PASSES_MLIR_TO_GRAPH_PASSES_H_
tensorflow/compiler/mlir/tf2xla/internal/passes/mlir_to_graph_passes.td
+20
View File
@@ -0,0 +1,20 @@
/* 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 "mlir/Pass/PassBase.td"
def VerifyInputDialectToExecutorPass : Pass<"verify-input-dialect-to-executor-pass", "mlir::func::FuncOp"> {
let summary = "Verify that TF dialect to executor converter receives the correct input.";
let description = [{
Verifies the input before exporting to TF executor. This includes checking whether the Ops are in TF functional, have device attributes & there are no tf_device.cluster_func ops.
}];
let constructor = "tensorflow::tf2xla::internal::CreateVerifyInputDialectToExecutorPass()";
}
tensorflow/compiler/mlir/tf2xla/internal/passes/tpu_cluster_formation.cc
+977
View File
@@ -0,0 +1,977 @@
/* 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 <cassert>
#include <cstdint>
#include <iterator>
#include <memory>
#include <set>
#include <string>
#include <tuple>
#include <unordered_map>
#include <utility>
#include "absl/container/flat_hash_map.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringExtras.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/iterator_range.h"
#include "llvm/Support/Casting.h"
#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/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Diagnostics.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/OperationSupport.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/IR/ValueRange.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Transforms/RegionUtils.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/string_util.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/tpu_rewrite_device_util.h"
#include "tensorflow/core/util/device_name_utils.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {
using mlir::Block;
using mlir::DialectRegistry;
using mlir::LogicalResult;
using mlir::ModuleOp;
using mlir::NamedAttribute;
using mlir::NamedAttrList;
using mlir::OpBuilder;
using mlir::Operation;
using mlir::OpResult;
using mlir::Region;
using mlir::StringAttr;
using mlir::success;
using mlir::Type;
using mlir::Value;
using mlir::ValueRange;
using mlir::WalkResult;
constexpr llvm::StringRef kDeviceAttr = "device";
constexpr llvm::StringRef kNameAttr = "name";
constexpr llvm::StringRef kNumCoresPerReplicaAttr = "num_cores_per_replica";
constexpr llvm::StringRef kNumReplicasAttr = "num_replicas";
constexpr llvm::StringRef kMirroredVariableIndicesAttr =
"_mirrored_variable_indices";
constexpr llvm::StringRef kBadReplicateInfoAttrMsg =
"requires '_replication_info' string attribute";
// Mapping for `_replication_info` attribute to TPUReplicateMetadata attributes.
using MetadataMap = llvm::SmallDenseMap<llvm::StringRef, NamedAttrList, 8>;
// A set of operations. We use a `SmallSetVector` in order to have deterministic
// traversal order (= insertion order), independent of the pointer keys.
using OpSetVector = llvm::SmallSetVector<Operation*, 8>;
// Mapping for `_replication_info` attribute to ops of a cluster.
using ClusterMap = llvm::SmallDenseMap<llvm::StringRef, OpSetVector, 8>;
#define GEN_PASS_DEF_TPUCLUSTERFORMATIONPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc"
class TPUClusterFormationPass
: public impl::TPUClusterFormationPassBase<TPUClusterFormationPass> {
public:
explicit TPUClusterFormationPass(bool strict_clusters)
: strict_clusters_(strict_clusters) {}
void getDependentDialects(DialectRegistry& registry) const override {
registry.insert<mlir::tf_device::TensorFlowDeviceDialect>();
}
void runOnOperation() override;
private:
bool strict_clusters_;
};
// Creates a mapping from the TPUReplicateMetadata ops `_replication_info`
// attribute to its attributes and removes the ops. If multiple
// TPUReplicateMetadata ops have the same `_replication_info` attribute, an
// error will be returned.
LogicalResult CollectMetadata(Block* block, MetadataMap* metadata_map) {
// Just look at top-level operations in the block (not nested ones)
for (Operation& op : llvm::make_early_inc_range(*block)) {
auto metadata_op = llvm::dyn_cast<mlir::TF::TPUReplicateMetadataOp>(op);
if (!metadata_op) continue;
NamedAttrList attrs(metadata_op->getAttrDictionary());
// Missing or bad `_replication_info` attribute.
auto replication_info_attr = attrs.get(mlir::TF::kReplicationInfoAttr);
if (!replication_info_attr)
return metadata_op.emitError() << kBadReplicateInfoAttrMsg;
auto replication_info_attr_str =
mlir::dyn_cast<StringAttr>(replication_info_attr);
if (!replication_info_attr_str ||
replication_info_attr_str.getValue().empty())
return metadata_op.emitError() << kBadReplicateInfoAttrMsg;
// Remove `name` attribute.
attrs.erase(StringAttr::get(metadata_op.getContext(), kNameAttr));
auto it = metadata_map->try_emplace(replication_info_attr_str.getValue(),
std::move(attrs));
// There are multiple TPUReplicateMetadata ops with the same
// `_replication_info` attribute.
if (!it.second) {
return metadata_op.emitError()
<< "multiple TPUReplicateMetadata ops with the same '"
<< mlir::TF::kReplicationInfoAttr << "' attribute '"
<< replication_info_attr_str.getValue() << "' found";
}
metadata_op.erase();
}
return success();
}
struct OpDevice {
Operation* op;
std::string device;
};
LogicalResult HasValidDeviceTypeAttribute(Block* block) {
// Use ordered set here to make error message below deterministic.
std::set<llvm::StringRef> device_types;
for (Operation& op : *block) {
// Collect device types which currently must be consistent per block
// (checked later).
if (auto device_type_attr =
op.getAttrOfType<StringAttr>(mlir::TF::kCompileDeviceTypeAttr)) {
// tf.StatefulPartitionedCall ops with and without
// _tpu_replicate attributes may exist in the same graph. Ops without
// the attribute but with _XlaMustCompile=true would have
// _xla_compile_device_type="" after
// CanonicalizeCompileAndReplicateAttributesPass. Skip empty value here.
if (!device_type_attr.getValue().empty()) {
device_types.insert(device_type_attr);
}
}
}
if (device_types.size() > 1) {
return block->getParentOp()->emitError()
<< "found different '" << mlir::TF::kCompileDeviceTypeAttr
<< "' attribute values (" << llvm::join(device_types, ",")
<< ") in same block which is not supported";
}
return success();
}
// Collects and clusters ops based on `_replication_info` attribute. Returns
// an error in case of invalid compilation or replication attribute(s).
LogicalResult CollectAndGroupClusterOps(Block* block, ClusterMap* clusters) {
LogicalResult result = HasValidDeviceTypeAttribute(block);
if (failed(result)) return result;
for (Operation& op : *block) {
LogicalResult result =
mlir::TF::HasValidCompilationAndReplicationAttributes(op);
if (failed(result)) return result;
// Skip ops with non-TPU device type, they are handled elsewhere.
auto device_type_attr =
op.getAttrOfType<StringAttr>(mlir::TF::kCompileDeviceTypeAttr);
if (device_type_attr) {
if (device_type_attr.getValue().empty()) continue;
if (device_type_attr.getValue() != mlir::TF::kTpuDevice) continue;
}
if (op.hasAttr(mlir::TF::kReplicationInfoAttr)) {
// For replicated case, borrow cluster structure from replication info.
// Following condition is already checked in
// `HasValidCompilationAndReplicationAttributes` above, assert here for
// documentation and to avoid breakage when that function is changed.
assert(op.hasAttr(mlir::TF::kCompileDeviceTypeAttr));
auto attr = op.getAttrOfType<StringAttr>(mlir::TF::kReplicationInfoAttr);
auto it = clusters->try_emplace(attr.getValue());
it.first->getSecond().insert(&op);
}
}
return success();
}
// Returns true iff `op` has a direct control dependency from (`incoming` ==
// true) or to (`incoming` == false) any op in `cluster_ops` or
// `cluster_dependent_ops`.
Operation* getOpClusterControlDependency(
Operation* op, bool incoming, const OpSetVector& cluster_ops,
const OpSetVector& cluster_dependent_ops,
const mlir::TF::SideEffectAnalysis::Info& side_effect_analysis) {
auto filter = [&](Operation* other_op) {
return cluster_ops.contains(other_op) ||
cluster_dependent_ops.contains(other_op);
};
auto other_ops =
incoming ? side_effect_analysis.DirectControlPredecessors(op, filter)
: side_effect_analysis.DirectControlSuccessors(op, filter);
if (other_ops.empty())
return nullptr;
else
return *other_ops.begin();
}
// Returns true iff `op` has a direct data dependency from (`incoming` == true
// or to (`incoming` == false) any op in `cluster_ops` or
// `cluster_dependent_ops`.
Operation* getOpClusterDataDependency(
Operation* op, bool incoming, const OpSetVector& cluster_ops,
const OpSetVector& cluster_dependent_ops) {
Operation* other_op = nullptr;
op->walk([&](Operation* inner_op) {
ValueRange values = incoming ? ValueRange(inner_op->getOperands())
: ValueRange(inner_op->getResults());
llvm::SmallVector<Operation*, 4> candidates;
for (Value value : values) {
if (incoming) {
candidates = {value.getDefiningOp()};
} else {
candidates.assign(value.getUsers().begin(), value.getUsers().end());
}
for (Operation* candidate_op : candidates) {
if (cluster_ops.contains(candidate_op) ||
cluster_dependent_ops.contains(candidate_op)) {
other_op = candidate_op;
return WalkResult::interrupt();
}
}
}
return WalkResult::advance();
});
return other_op;
}
// Collects ops that need to be moved behind the cluster due to data or control
// dependencies.
mlir::FailureOr<llvm::SmallSetVector<Operation*, 8>> CollectClusterSuccessorOps(
Block* block, const OpSetVector& cluster_ops,
const mlir::TF::SideEffectAnalysis::Info& side_effect_analysis,
bool strict_clusters) {
OpSetVector cluster_predecessor_ops;
OpSetVector cluster_successor_ops;
// Collect non-cluster ops that have a dependency to the cluster. For this
// traverse all ops from last to first cluster op and keep track of in-between
// non-cluster ops that have some outgoing (transitive) dependency to some
// cluster op (`cluster_predecessor_ops`).
auto rfront = Block::reverse_iterator(cluster_ops.front());
auto rback = Block::reverse_iterator(cluster_ops.back());
for (Operation& op : llvm::make_range(rback, rfront)) {
if (cluster_ops.contains(&op)) continue;
bool has_dependency_to_cluster =
getOpClusterDataDependency(&op, /*incoming=*/false, cluster_ops,
cluster_predecessor_ops) != nullptr ||
getOpClusterControlDependency(&op, /*incoming=*/false, cluster_ops,
cluster_predecessor_ops,
side_effect_analysis) != nullptr;
if (has_dependency_to_cluster) cluster_predecessor_ops.insert(&op);
}
// Collect non-cluster ops that have a dependency from the cluster. For this
// traverse all ops from first to last cluster op and keep track of in-between
// non-cluster ops that have some incoming (transitive) dependency from some
// cluster op (`cluster_successor_ops`).
auto front = Block::iterator(cluster_ops.front());
auto back = Block::iterator(cluster_ops.back());
for (Operation& op : llvm::make_range(front, back)) {
if (cluster_ops.contains(&op)) continue;
Operation* data_predecessor = getOpClusterDataDependency(
&op, /*incoming=*/true, cluster_ops, cluster_successor_ops);
Operation* control_predecessor = getOpClusterControlDependency(
&op, /*incoming=*/true, cluster_ops, cluster_successor_ops,
side_effect_analysis);
if (data_predecessor || control_predecessor) {
if (cluster_predecessor_ops.contains(&op)) {
// Op has a dependency from and to the cluster which is invalid. Instead
// of erroring out we don't add the op to `cluster_successor_ops` which
// is in line with previous behavior when certain control dependencies
// were not considered.
// TODO(b/216706460) Establish some contract here: Should we expect only
// valid clusters, or should we split clusters accordingly? The latter
// might have runtime impact for existing models.
// We should make this message an error once there is such a contract
// and once existing cases have been fixed.
mlir::InFlightDiagnostic error =
strict_clusters ? mlir::emitError(op.getLoc(), "")
: mlir::emitWarning(op.getLoc(), "");
error << "Op has cyclic dependency with a compilation cluster:\n";
error << "The cluster depends on\n";
error << op.getName() << "\n"
<< "which is outside of cluster, but itself depends ";
if (data_predecessor) {
error << "on\n" << data_predecessor->getName() << "\n";
} else if (control_predecessor) {
error << "(via control) on\n"
<< control_predecessor->getName() << "\n";
}
if (cluster_ops.contains(data_predecessor) ||
cluster_ops.contains(control_predecessor))
error << "which is inside the cluster.\n";
else
error << "which the cluster depends on.\n";
} else {
cluster_successor_ops.insert(&op);
}
}
}
return cluster_successor_ops;
}
// Collects results and associated types of the cluster that are used outside of
// the cluster. These results and types are used to create the clusters
// `tf_device.cluster` and associated terminator. Results that have no uses
// outside of the cluster (i.e. results of ops in the cluster are only consumed
// by other ops in the cluster) are pruned.
llvm::SmallVector<Value, 8> CollectClusterResults(
Block* block, const OpSetVector& cluster_ops) {
llvm::SmallVector<Value, 8> results;
for (Operation* op : cluster_ops) {
for (Value result : op->getResults()) {
for (Operation* user : result.getUsers()) {
// Check if user is not an op in the cluster.
if (cluster_ops.count(block->findAncestorOpInBlock(*user)) == 0) {
results.push_back(result);
break;
}
}
}
}
return results;
}
// Creates a `tf_device.cluster` to wrap cluster ops.
mlir::tf_device::ClusterOp CreateClusterOp(
Block* block, const OpSetVector& cluster_ops, llvm::ArrayRef<Value> results,
llvm::ArrayRef<Operation*> cluster_successor_ops) {
// `tf_device.cluster` will be placed at where the last op of the cluster is.
Operation* last_cluster_op = cluster_ops.back();
OpBuilder builder(last_cluster_op);
llvm::SmallVector<Type, 8> result_types;
for (Value result : results) result_types.push_back(result.getType());
auto cluster = mlir::tf_device::ClusterOp::create(
builder, last_cluster_op->getLoc(), result_types);
Block* body = new Block;
cluster.getBody().push_back(body);
// Move cluster ops to the cluster body. Also remove `_replication_info` and
// `device` attribute from ops in the cluster when that information is
// redundant will the `tf_device.cluster`. Do this for all ops including
// nested ops.
for (Operation* cluster_op : cluster_ops) {
cluster_op->moveBefore(body, body->end());
cluster_op->walk([&](Operation* inner_op) {
inner_op->removeAttr(mlir::TF::kReplicationInfoAttr);
inner_op->removeAttr(mlir::TF::kCompileDeviceTypeAttr);
if (auto attr = inner_op->getAttrOfType<StringAttr>(kDeviceAttr)) {
// Preserve device attribute if the op is placed on a replicated core
// device. Device attribute is used to infer the appropriate sharding
// within TPUs for this op.
// TODO(b/183598857): Use explicit sharding ops from the front-end.
// For example, dequeue ops generated by
// tensorflow/python/tpu/tpu_feed.py
if (!tensorflow::IsTPUReplicatedCore(attr.getValue())) {
inner_op->removeAttr(kDeviceAttr);
}
}
});
}
// Add terminator.
builder.setInsertionPointToEnd(body);
mlir::tf_device::ReturnOp::create(builder, last_cluster_op->getLoc(),
results);
// Replaces uses of cluster ops results outside of cluster with the associated
// `tf_device.cluster` results.
for (auto ret_vals : llvm::zip(results, cluster.getResults())) {
Value old_ret = std::get<0>(ret_vals);
Value new_ret = std::get<1>(ret_vals);
for (auto& use : llvm::make_early_inc_range(old_ret.getUses())) {
Operation* user = use.getOwner();
if (!body->findAncestorOpInBlock(*user)) use.set(new_ret);
}
}
// Move ops that depend on something in the cluster behind the cluster.
Operation* op_after_cluster = cluster.getOperation()->getNextNode();
for (Operation* op : cluster_successor_ops) op->moveBefore(op_after_cluster);
return cluster;
}
// Returns an op of the given type that uses the result, along with
// a list of identity ops along the way.
template <typename T>
std::tuple<T, llvm::SmallVector<mlir::TF::IdentityOp, 4>> GetSingleUserOfType(
OpResult result) {
llvm::SmallVector<mlir::TF::IdentityOp, 4> identity_ops;
do {
Operation* user = result.hasOneUse() ? *result.getUsers().begin() : nullptr;
if (auto t = llvm::dyn_cast_or_null<T>(user)) {
return std::make_tuple(t, identity_ops);
} else if (auto identity =
llvm::dyn_cast_or_null<mlir::TF::IdentityOp>(user)) {
identity_ops.emplace_back(identity);
result = identity->getResult(0);
} else {
result = OpResult(); // reset to stop iterating
}
} while (result);
return std::make_tuple(T(), identity_ops);
}
using PartitionedClusterOutputMap = absl::flat_hash_map<
uint64_t, llvm::SmallVector<mlir::TF::TPUPartitionedOutputV2Op, 8>>;
// Returns the partitioned output ops from the cluster if there are any,
// along with any single user identity ops between them. Not all outputs
// of a cluster must be partitioned, so the output is a map from cluster
// output ids to ops.
std::tuple<PartitionedClusterOutputMap,
llvm::SmallVector<mlir::TF::IdentityOp, 8>>
GetPartitionedOutputsAndIdentityOps(mlir::tf_device::ClusterOp cluster) {
PartitionedClusterOutputMap partitioned_outputs;
llvm::SmallVector<mlir::TF::IdentityOp, 8> erase_list;
for (auto [cluster_result_id, cluster_result] :
llvm::enumerate(cluster.getResults())) {
auto [replicated_output, _] =
GetSingleUserOfType<mlir::TF::TPUReplicatedOutputOp>(cluster_result);
if (replicated_output) {
for (OpResult per_replica_result : replicated_output->getResults()) {
auto [partitioned_output, id_ops] =
GetSingleUserOfType<mlir::TF::TPUPartitionedOutputV2Op>(
per_replica_result);
if (partitioned_output) {
erase_list.insert(erase_list.end(), id_ops.begin(), id_ops.end());
partitioned_outputs[cluster_result_id].emplace_back(
partitioned_output);
}
}
}
}
return std::forward_as_tuple(partitioned_outputs, erase_list);
}
// Inlines the partitioned output ops into the cluster, and updates
// their users to point to the replicate op instead.
Operation* BuildPartitionedOutputs(
OpBuilder& builder, mlir::tf_device::ClusterOp cluster,
mlir::tf_device::ReplicateOp replicate_op,
PartitionedClusterOutputMap& partitioned_outputs,
llvm::SmallVector<mlir::TF::IdentityOp, 8>& erase_list,
llvm::SmallVector<Type, 8>& result_types, int num_replicas) {
Operation* result_op;
llvm::SmallVector<Value, 8> results;
uint64_t num_results = cluster.getNumResults();
for (uint64_t result_id = 0; result_id < num_results; ++result_id) {
auto search = partitioned_outputs.find(result_id);
if (search == partitioned_outputs.end()) {
// If the output is not partitioned, directly pass it through.
results.emplace_back(cluster.getResult(result_id));
continue;
}
// Otherwise, "inline" the partitioned output ops by:
// - Building a new op within the cluster.
// - Replacing all the uses of the original ops with the cluster's outputs.
llvm::SmallVector<mlir::TF::TPUPartitionedOutputV2Op, 8>& ops =
search->second;
for (auto [replica_id, partitioned_output] : llvm::enumerate(ops)) {
for (auto [core_id, result] :
llvm::enumerate(partitioned_output->getResults())) {
// outputs from replicate op are interleaved:
// [(replica:0,core:0), (replica:1,core:0), ...,
// (replica:0,core:1), (replica:1,core:1), ...]
uint64_t output_id =
core_id * num_replicas + replica_id + results.size();
result.replaceAllUsesWith(replicate_op.getResult(output_id));
}
}
// Assume all the replicas have the same structure.
mlir::TF::TPUPartitionedOutputV2Op first_op = *(ops.begin());
mlir::ArrayAttr dims = first_op.getPartitionDimsAttr();
StringAttr sharding = first_op.get_XlaShardingAttr();
StringAttr sharding_v2 = first_op.get_XlaShardingV2Attr();
Operation::result_type_range output_types = first_op.getResultTypes();
result_op = mlir::TF::TPUPartitionedOutputV2Op::create(
builder, replicate_op.getLoc(), output_types,
cluster.getResult(result_id), dims, sharding, sharding_v2);
results.insert(results.end(), result_op->getResults().begin(),
result_op->getResults().end());
}
// Once we've accumulated all the cluster's results, build a return op.
mlir::tf_device::ReturnOp::create(builder, result_op->getLoc(), results);
// Then erase all the identity and partitioned output ops.
for (const auto& [_, ops] : partitioned_outputs) {
for (mlir::TF::TPUPartitionedOutputV2Op op : ops) {
op->erase();
}
}
for (mlir::TF::IdentityOp to_erase : erase_list) {
to_erase->erase();
}
return result_op;
}
// Return the cluster's per-replica result type, converting any full-shaped
// tensor types into sharded-shaped ones if they're partitioned.
llvm::SmallVector<Type, 8> GetClusterResultTypes(
mlir::tf_device::ClusterOp cluster,
const PartitionedClusterOutputMap& partitioned_outputs) {
llvm::SmallVector<Type, 8> result_types;
Operation::result_type_range cluster_result_types = cluster.getResultTypes();
if (partitioned_outputs.empty()) {
// Directly pass through the cluster's outputs if none are partitioned.
result_types.insert(result_types.end(), cluster_result_types.begin(),
cluster_result_types.end());
} else {
// For each output of the cluster...
for (auto [output_id, result_type] :
llvm::enumerate(cluster_result_types)) {
auto search = partitioned_outputs.find(output_id);
if (search == std::end(partitioned_outputs)) {
// If it's not partitioned, directly pass it through.
result_types.emplace_back(result_type);
} else {
// Otherwise, pass through the result shard types.
Operation::result_type_range partitioned_result_types =
(*search->second.begin())->getResultTypes();
result_types.insert(result_types.end(),
partitioned_result_types.begin(),
partitioned_result_types.end());
}
}
}
return result_types;
}
// Creates a `tf_device.replicate` to represent replication for the cluster, if
// necessary. Erases Identity ops between partitioned and replicated output ops.
LogicalResult ReplicateCluster(mlir::tf_device::ClusterOp cluster,
int num_replicas, int num_cores_per_replica) {
OpBuilder builder(cluster);
auto [partitioned_outputs, erase_list] =
GetPartitionedOutputsAndIdentityOps(cluster);
for (auto [_, ops] : partitioned_outputs) {
if (!(ops.empty() || ops.size() == num_replicas)) {
return (ops.begin())->emitOpError()
<< "expected zero or " << num_replicas
<< " 'TPUPartitionedOutput' op(s), instead got "
<< partitioned_outputs.size();
}
}
// No need to replicate.
if (num_replicas == 1) {
// Collapse all the Identity ops between the TRO and TPO ops.
if (!partitioned_outputs.empty()) {
for (mlir::TF::IdentityOp to_erase : erase_list) {
Value in = to_erase->getOperand(0);
OpResult out = to_erase->getResult(0);
out.replaceAllUsesWith(in);
to_erase->erase();
}
}
return success();
}
if (num_replicas < 1)
return cluster.emitError() << "requires '" << kNumReplicasAttr
<< "' int attribute to be at least 1";
LogicalResult status = success();
// Collect all used TPUReplicatedInput ops.
llvm::SmallVector<mlir::TF::TPUReplicatedInputOp, 8> replicated_input_ops;
llvm::SmallSet<mlir::TF::TPUReplicatedInputOp, 8> seen_ops;
mlir::visitUsedValuesDefinedAbove(
cluster.getBody(), cluster.getBody(), [&](mlir::OpOperand* operand) {
Operation* def = operand->get().getDefiningOp();
if (auto ri =
llvm::dyn_cast_or_null<mlir::TF::TPUReplicatedInputOp>(def)) {
if (!seen_ops.contains(ri)) {
seen_ops.insert(ri);
replicated_input_ops.push_back(ri);
}
}
// When model parallelism is used in conjunction with data parallelism
// for resource inputs, we need to collect the per replica resource
// inputs from input to `tf.TPUPartitionedInputV2` ops.
if (auto pi = llvm::dyn_cast_or_null<mlir::TF::TPUPartitionedInputV2Op>(
def)) {
if (pi->getNumOperands() != num_cores_per_replica)
status = pi.emitOpError()
<< "requires " << num_cores_per_replica
<< " operands but found " << pi->getNumOperands();
for (auto operand : pi.getInputs()) {
if (auto ri =
llvm::dyn_cast_or_null<mlir::TF::TPUReplicatedInputOp>(
operand.getDefiningOp())) {
if (!seen_ops.contains(ri)) {
seen_ops.insert(ri);
replicated_input_ops.push_back(ri);
}
}
}
}
});
if (failed(status)) return mlir::failure();
// Indices of the replicate op's arguments that are mirrored variables.
llvm::SmallVector<int64_t, 8> mirrored_variable_indices;
// Check if number of operands of each used TPUReplicatedInput op matches
// `num_replicas` or 1. Collect all their operands and associated type for
// creating the replicate op.
llvm::SmallVector<std::pair<ValueRange, Type>, 8> replicated_inputs;
llvm::SmallVector<Value, 8> packed_inputs;
llvm::SmallVector<mlir::TF::TPUReplicatedInputOp, 8> replicated_ops;
llvm::SmallVector<mlir::TF::TPUReplicatedInputOp, 8> packed_ops;
for (const auto& pos_and_input : llvm::enumerate(replicated_input_ops)) {
auto input = pos_and_input.value();
bool is_packed = input.getIsPacked();
const int num_operands = input->getNumOperands();
int num_inputs = is_packed ? 1 : num_replicas;
if (num_operands != num_inputs)
return input->emitOpError() << "requires " << num_inputs << " operands";
if (is_packed) {
packed_inputs.push_back(input->getOperand(0));
packed_ops.push_back(input);
} else {
replicated_inputs.push_back(
{input->getOperands(), input->getOperand(0).getType()});
replicated_ops.push_back(input);
}
}
// Create `ordered_tpu_replicate_inputs` which contains the final ordered
// replicate inputs. All packed arguments are moved to the end of the arg
// list.
llvm::SmallVector<mlir::TF::TPUReplicatedInputOp, 8>
ordered_tpu_replicate_inputs = replicated_ops;
ordered_tpu_replicate_inputs.append(packed_ops.begin(), packed_ops.end());
// Assign `mirrored_variable_indices` based on the ordered replicated inputs.
for (const auto& pos_and_input :
llvm::enumerate(ordered_tpu_replicate_inputs)) {
auto tpu_replicated_input = pos_and_input.value();
if (tpu_replicated_input.getIsMirroredVariable()) {
mirrored_variable_indices.push_back(pos_and_input.index());
}
}
// Create replicate op.
auto result_types = GetClusterResultTypes(cluster, partitioned_outputs);
auto replicate_op = mlir::tf_device::ReplicateOp::create(
builder, cluster.getLoc(), num_replicas,
llvm::SmallDenseMap<llvm::StringRef,
llvm::SmallVector<llvm::StringRef, 4>>(),
replicated_inputs, packed_inputs, result_types);
if (!mirrored_variable_indices.empty())
replicate_op->setAttr(kMirroredVariableIndicesAttr,
builder.getI64ArrayAttr(mirrored_variable_indices));
// Replace replicated cluster results with replicate op results.
uint64_t offset = 0;
for (auto [idx, result] : llvm::enumerate(cluster.getResults())) {
if (partitioned_outputs.contains(idx)) {
// Partitioned output propagation happens in BuildPartitionedOutputs.
offset += num_replicas * num_cores_per_replica;
continue;
}
auto replicate_outputs = llvm::make_range(
std::next(replicate_op.result_begin(), offset),
std::next(replicate_op.result_begin(), offset + num_replicas));
for (auto& use : llvm::make_early_inc_range(result.getUses())) {
Operation* def = use.getOwner();
if (!llvm::isa<mlir::TF::TPUReplicatedOutputOp>(def)) {
// If user is not a `tf.TPUReplicatedOutput`, simply forward the first
// replica output. Certain Graphs under V1 create `tf.Identity` users of
// replicated ops to pin the TPU computation for execution.
use.set(*replicate_outputs.begin());
continue;
}
const int def_num_results = def->getNumResults();
if (def_num_results != num_replicas)
return def->emitOpError() << "requires " << num_replicas << " results";
def->replaceAllUsesWith(replicate_outputs);
}
offset += num_replicas;
}
// Collect all `tf.TPUPartitionedInputV2` ops to be moved inside the
// `tf_device.replicate` later.
llvm::SmallSet<Operation*, 4> partitioned_inputs;
for (auto input_and_block_arg :
llvm::zip(ordered_tpu_replicate_inputs,
replicate_op.GetBody().getArguments())) {
mlir::TF::TPUReplicatedInputOp input = std::get<0>(input_and_block_arg);
Value block_arg = std::get<1>(input_and_block_arg);
mlir::replaceAllUsesInRegionWith(input->getResult(0), block_arg,
cluster.getBody());
// Update replicated input use in tf.TPUPartitionedInputV2 op.
for (auto& use : input->getUses()) {
auto pi =
llvm::dyn_cast<mlir::TF::TPUPartitionedInputV2Op>(use.getOwner());
if (pi) {
pi.setOperand(use.getOperandNumber(), block_arg);
partitioned_inputs.insert(pi.getOperation());
}
}
}
// Create terminator for replicate op and move `tf_device.cluster` and
// `tf.TPUPartitionedInputV2`(s) into replicate body.
builder.setInsertionPointToEnd(&replicate_op.GetBody());
Operation* result_op;
if (!partitioned_outputs.empty()) {
result_op = BuildPartitionedOutputs(builder, cluster, replicate_op,
partitioned_outputs, erase_list,
result_types, num_replicas);
} else {
result_op = mlir::tf_device::ReturnOp::create(
builder, replicate_op.getLoc(), cluster.getResults());
}
for (auto pi : partitioned_inputs) pi->moveBefore(result_op);
cluster.getOperation()->moveBefore(result_op);
return success();
}
// Forms clusters with ops of the same `_replication_info` attribute under a
// block.
//
// For a given block, clusters are formed via grouping ops by
// `_replication_info` attributes. For every cluster formed:
// 1. Find associated TPUReplicateMetadata attributes with the same
// `_replication_info` attribute.
// 2. Find users not in cluster that are interleaved between cluster ops.
// 3. Find external uses of cluster ops.
// 4. Create `tf_device.cluster` with results consisting of the external uses
// of cluster ops determined at 3.
// 5. Move cluster ops to `tf_device.cluster` body.
// 6. Replace external uses of cluster ops uses with `tf_device.cluster`
// results.
// 7. Move users from 2 to after the `tf_device.cluster`.
// 8. Wrap cluster (`tf_device.cluster`) in a `tf_device.replicate` if
// attribute `num_replicas` is greater than 1.
// 9. Copy over TPUReplicateMetadata attributes to `tf_device.cluster`.
LogicalResult FormClustersInBlock(
Block* block,
const mlir::TF::SideEffectAnalysis::Info& side_effect_analysis,
bool strict_clusters) {
MetadataMap metadata_map;
LogicalResult result = CollectMetadata(block, &metadata_map);
if (failed(result)) return result;
// If there is no TPUReplicateMetadata op in this block, process blocks in
// regions attached to the op's in the block.
if (metadata_map.empty()) {
for (Operation& op : *block) {
for (Region& region : op.getRegions()) {
if (!llvm::hasSingleElement(region))
return op.emitOpError("Expected single block region");
if (failed(FormClustersInBlock(&region.front(), side_effect_analysis,
strict_clusters)))
return mlir::failure();
}
}
return success();
}
ClusterMap clusters;
result = CollectAndGroupClusterOps(block, &clusters);
if (failed(result)) return result;
for (const auto& cluster_metadata_and_ops : clusters) {
const auto& cluster_ops = cluster_metadata_and_ops.getSecond();
auto cluster_metadata =
metadata_map.find(cluster_metadata_and_ops.getFirst());
// llvm::errs() << __func__ << "\n";
// No TPUReplicateMetadata for a `_replication_info` attribute.
if (cluster_metadata == metadata_map.end()) {
block->getParentOp()->emitWarning()
<< "TPUReplicateMetadata for associated '"
<< mlir::TF::kReplicationInfoAttr << "' attribute '"
<< cluster_metadata_and_ops.getFirst() << "' is missing";
continue;
}
auto status = CollectClusterSuccessorOps(
block, cluster_ops, side_effect_analysis, strict_clusters);
if (failed(status)) return status;
OpSetVector cluster_successor_ops = *status;
llvm::SmallVector<Value, 8> results =
CollectClusterResults(block, cluster_ops);
mlir::tf_device::ClusterOp cluster = CreateClusterOp(
block, cluster_ops, results, cluster_successor_ops.getArrayRef());
auto num_replicas = cluster_metadata->getSecond().get(kNumReplicasAttr);
if (!num_replicas || !mlir::isa<mlir::IntegerAttr>(num_replicas))
return cluster.emitError()
<< "requires '" << kNumReplicasAttr << "' int attribute";
int num_cores_per_replica = 1;
auto num_cores_per_replica_attr = mlir::dyn_cast_or_null<mlir::IntegerAttr>(
cluster_metadata->getSecond().get(kNumCoresPerReplicaAttr));
if (num_cores_per_replica_attr)
num_cores_per_replica = num_cores_per_replica_attr.getInt();
if (failed(ReplicateCluster(
cluster, mlir::cast<mlir::IntegerAttr>(num_replicas).getInt(),
num_cores_per_replica)))
return mlir::failure();
// Copy TPUReplicateMetadata attributes to `tf_device.cluster`.
cluster->setAttrs(
cluster_metadata->second.getDictionary(cluster.getContext()));
// Exclude `num_replicas` as cluster should be replicated if necessary.
cluster->removeAttr(kNumReplicasAttr);
}
return success();
}
LogicalResult FormClustersInFunction(
mlir::func::FuncOp func,
const mlir::TF::SideEffectAnalysis::Info& side_effect_analysis,
bool strict_clusters) {
if (!llvm::hasSingleElement(func))
return func.emitOpError("Expecting a single block function");
if (failed(FormClustersInBlock(&func.front(), side_effect_analysis,
strict_clusters)))
return mlir::failure();
// Remove TPUReplicatedInput and TPUReplicatedOutput nodes.
auto remove_result = func.walk([&](Operation* op) {
if (!llvm::isa<mlir::TF::TPUReplicatedInputOp,
mlir::TF::TPUReplicatedOutputOp>(op))
return WalkResult::advance();
// Forward operand to result. When `num_replicas` attribute is 1, no
// `tf_device.replicate` is created and replicated (1) operands/results are
// untouched.
if (op->getNumOperands() == 1 && op->getNumResults() == 1)
op->getResult(0).replaceAllUsesWith(op->getOperand(0));
// Leftover TPUReplicatedInput/TPUReplicatedOutput that are not of
// `num_replicas` to 1.
if (!op->use_empty()) {
op->emitOpError() << "is expected to have no uses, but it is operand#"
<< op->use_begin()->getOperandNumber() << " of "
<< *op->use_begin()->getOwner();
return WalkResult::interrupt();
}
op->erase();
return WalkResult::advance();
});
return mlir::failure(remove_result.wasInterrupted());
}
void TPUClusterFormationPass::runOnOperation() {
// Attributes on tf.Constant aren't reliable: CSE will merge ConstantLike ops
// with the same value (but different attributes!) into the same tf.Const
// definition, potentially leading to bogus _replication_info attributes. So
// we just scrub all tf.Constants of all extra attributes.
// TODO(kramm): Remove this once tf.Const's folder is aware of extra
// attributes.
auto value_str_attr = StringAttr::get(&getContext(), "value");
getOperation().walk([&](mlir::TF::ConstOp cst) {
auto dict = cst->getAttrDictionary();
if (dict.size() == 1) {
return; // Optimization. Assume the one attribute is "value".
}
// Recreate the attributes dictionary to only contain "value".
NamedAttrList attributes;
attributes.append(NamedAttribute(value_str_attr, cst->getAttr("value")));
cst->setAttrs(attributes.getDictionary(&getContext()));
});
auto& side_effect_analysis = getAnalysis<mlir::TF::SideEffectAnalysis>();
for (auto func : getOperation().getOps<mlir::func::FuncOp>())
if (!func.isExternal() &&
failed(FormClustersInFunction(
func, side_effect_analysis.GetAnalysisForFunc(func),
strict_clusters_)))
return signalPassFailure();
}
} // anonymous namespace
std::unique_ptr<mlir::OperationPass<ModuleOp>> CreateTPUClusterFormationPass(
bool strict_clusters) {
return std::make_unique<TPUClusterFormationPass>(strict_clusters);
}
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/passes/tpu_sharding_identification_pass.cc
+814
View File
@@ -0,0 +1,814 @@
/* 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 <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <type_traits>
#include <utility>
#include <variant>
#include "absl/algorithm/container.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/FormatVariadic.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Block.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Interfaces/CallInterfaces.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.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_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_traits.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/xla_sharding_util.h"
#include "xla/hlo/builder/sharding_builder.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/statusor.h"
#include "xla/xla_data.pb.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {
using OpShardingVariant = std::variant<mlir::Operation*, llvm::StringRef>;
using OpShardingVector = llvm::SmallVector<OpShardingVariant, 8>;
using OptionalOpShardingVector =
llvm::SmallVector<std::optional<OpShardingVariant>, 8>;
using llvm::StringRef;
using mlir::Block;
using mlir::BlockArgument;
using mlir::BoolAttr;
using mlir::Builder;
using mlir::IntegerAttr;
using mlir::LogicalResult;
using mlir::ModuleOp;
using mlir::Operation;
using mlir::OpOperand;
using mlir::OpResult;
using mlir::RankedTensorType;
using mlir::StringAttr;
using mlir::Value;
using mlir::WalkResult;
constexpr char kReplicateSharding[] = "";
constexpr char kShardingAttr[] = "mhlo.sharding";
constexpr char kUseSpmdAttr[] = "use_spmd_for_xla_partitioning";
constexpr char kAliasingAttr[] = "tf.aliasing_output";
constexpr char kNumCoresPerReplicaAttr[] = "num_cores_per_replica";
const char kShardingAttribute[] = "_XlaSharding";
const char kShardingAttributeV2[] = "_XlaShardingV2";
#define GEN_PASS_DEF_TPUSHARDINGIDENTIFICATIONPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc"
struct TPUShardingIdentificationPass
: public impl::TPUShardingIdentificationPassBase<
TPUShardingIdentificationPass> {
void runOnOperation() final;
};
std::string CreateMissingAttributeMsg(llvm::StringRef attribute) {
return llvm::formatv("requires attribute '{0}'", attribute).str();
}
// Returns nullptr if the op does not have a sharding attribute.
template <typename PartitionedOp>
mlir::Operation* NullUnlessSharded(PartitionedOp op) {
return op.get_XlaSharding() ? op : nullptr;
}
// Returns true if unitary op has one of the traits that meets the requirements
// for sharding, otherwise returns false.
bool UnaryOpHasTraitsForSharding(Operation* op) {
// Trait "SameOperandsAndResultTypeResolveRef" for Cast, real/imag, etc.
if (op->hasTrait<mlir::OpTrait::TF::SameOperandsAndResultTypeResolveRef>())
return true;
// Trait "SameOperandsAndResultType" for Exp, ceil, etc.
if (op->hasTrait<mlir::OpTrait::SameOperandsAndResultType>()) return true;
// Trait "OperandsSameAsResultsTypeOrRef" for Identity.
if (op->hasTrait<mlir::OpTrait::TF::OperandsSameAsResultsTypeOrRef>())
return true;
return false;
}
// Returns true if binary op has one of the traits that meets the requirements
// for sharding, otherwise returns false.
bool BinaryOpHasTraitsForSharding(Operation* op) {
// Trait "CwiseBinary" and "SameOperandsAndResultElementTypeResolveRef" for
// AddV2, Sub, etc.
if (op->hasTrait<
mlir::OpTrait::TF::SameOperandsAndResultElementTypeResolveRef>() &&
op->hasTrait<mlir::OpTrait::TF::CwiseBinary>())
return true;
return false;
}
bool DoTypesHavePartialSameShape(Value value_0, Value value_1) {
auto shape_0 =
mlir::dyn_cast_or_null<mlir::RankedTensorType>(value_0.getType());
auto shape_1 =
mlir::dyn_cast_or_null<mlir::RankedTensorType>(value_1.getType());
if (shape_0 && shape_1) {
if (shape_0.hasStaticShape() && shape_1.hasStaticShape())
return shape_0.getShape() == shape_1.getShape();
int i = 0, j = 0;
while (i < shape_0.getShape().size() && j < shape_1.getShape().size()) {
if (shape_0.getShape()[i] != shape_1.getShape()[j] &&
!shape_0.isDynamicDim(i) && !shape_1.isDynamicDim(j)) {
return false;
}
if (shape_0.getShape()[i] == shape_1.getShape()[j]) {
i++;
j++;
} else {
if (shape_0.isDynamicDim(i)) {
i++;
}
if (shape_1.isDynamicDim(j)) {
j++;
}
}
}
return i == shape_0.getShape().size() && j == shape_1.getShape().size();
}
return false;
}
// Returns a TPUPartitionedInput op connected to a `tf_device.cluster_func`
// operand value if it has an XLA sharding. If value is a resource type then
// TPUPartitionedInput op will be connected to a ReadVariable op that feeds into
// a `tf_device.cluster_func`.
mlir::Operation* GetXlaShardingFromOperand(Value value) {
Value value_to_visit = value;
if (auto read_var = value_to_visit.getDefiningOp<mlir::TF::ReadVariableOp>())
value_to_visit = read_var.getResource();
if (auto partitioned_input =
value_to_visit.getDefiningOp<mlir::TF::TPUPartitionedInputV2Op>()) {
return NullUnlessSharded(partitioned_input);
}
return nullptr;
}
// Returns the op sharding attribute from a partitioned operator.
std::optional<StringRef> GetXlaShardingFromOperator(mlir::Operation* op) {
if (auto partitioned_output =
llvm::dyn_cast<mlir::TF::TPUPartitionedOutputV2Op>(op)) {
if (partitioned_output.get_XlaShardingV2().has_value()) {
return partitioned_output.get_XlaShardingV2();
}
return partitioned_output.get_XlaSharding();
} else if (auto partitioned_input =
llvm::dyn_cast<mlir::TF::TPUPartitionedInputV2Op>(op)) {
if (partitioned_input.get_XlaShardingV2().has_value()) {
return partitioned_input.get_XlaShardingV2();
}
return partitioned_input.get_XlaSharding();
}
return std::nullopt;
}
// Returns the sharding string from a op-sharding variant if it is available.
std::optional<StringRef> GetShardingStringFromVariant(
const OpShardingVariant& sharding_or_op) {
return std::visit(
[](auto&& sharding_or_op) -> std::optional<StringRef> {
using T = std::decay_t<decltype(sharding_or_op)>;
if constexpr (std::is_same_v<T, StringRef>) {
return sharding_or_op;
} else {
return GetXlaShardingFromOperator(sharding_or_op);
}
},
sharding_or_op);
}
// Returns the sharding from a op-sharding variant if it is available and valid.
std::optional<xla::OpSharding> GetShardingFromVariant(
const OpShardingVariant& sharding_or_op) {
xla::OpSharding sharding;
const auto sharding_string = GetShardingStringFromVariant(sharding_or_op);
if (!sharding_string) return std::nullopt;
if (tensorflow::DecodeShardingAttribute(sharding_string->str(), sharding,
false)
.failed()) {
return std::nullopt;
}
return sharding;
}
// Converts an op-sharding vector into a string attr using the builder.
mlir::ArrayAttr GetStrArrayAttr(Builder* builder,
const OpShardingVector& vect) {
llvm::SmallVector<mlir::Attribute, 8> strings;
for (const auto& sharding_or_op : vect) {
if (const auto sharding = GetShardingStringFromVariant(sharding_or_op)) {
strings.emplace_back(builder->getStringAttr(*sharding));
}
}
return builder->getArrayAttr(strings);
}
// Verify whether the given sharding can be applied to the given (tensor) type.
// (A bad sharding might mean failing tf.Split ops if the graph later executes
// on CPU)
// If the sharding is incorrect, return failure. If it's good, or if we can't
// verify it, return success.
LogicalResult VerifySharding(mlir::Type type,
const OpShardingVariant& sharding_or_op) {
auto* partitioned_op =
std::holds_alternative<mlir::Operation*>(sharding_or_op)
? std::get<mlir::Operation*>(sharding_or_op)
: nullptr;
const auto sharding = GetShardingFromVariant(sharding_or_op);
if (!sharding || sharding->type() != xla::OpSharding::OTHER) {
// Some test cases use \01\02\03 as sharding, to test propagation. Treat
// a non-proto sharding as valid, and don't verify further. We also only
// verify shardings that actually break a tensor apart.
return mlir::success();
}
if (RankedTensorType ranked_type = mlir::dyn_cast<RankedTensorType>(type)) {
const int64_t tensor_rank = ranked_type.getRank();
int tile_assignment_rank = sharding->tile_assignment_dimensions_size();
// When a tensor is partial or subgroup tiled, its tile assignment will
// have one or more dimension(s) than its rank; so, we subtract them to
// determine which rank the sharding is compatible with.
tile_assignment_rank -= (int)sharding->replicate_on_last_tile_dim();
tile_assignment_rank -= sharding->last_tile_dims_size();
if (tensor_rank < tile_assignment_rank) {
if (partitioned_op) {
partitioned_op->emitError()
<< "tensor of type " << ranked_type << " (rank=" << tensor_rank
<< ") sharded in " << (tile_assignment_rank - tensor_rank)
<< " extra dimension(s) by: " << sharding->DebugString();
}
return mlir::failure();
}
}
return mlir::success();
}
// Verify sharding for all arguments and return values.
LogicalResult VerifyShardings(mlir::func::FuncOp func,
const OpShardingVector& sharding_for_args,
const OpShardingVector& sharding_for_rets) {
Block& function_block = func.front();
for (auto sharding_and_arg :
llvm::zip(sharding_for_args, function_block.getArguments())) {
const auto& sharding = std::get<0>(sharding_and_arg);
BlockArgument arg = std::get<1>(sharding_and_arg);
if (failed(VerifySharding(arg.getType(), sharding))) return mlir::failure();
}
Operation* terminator = function_block.getTerminator();
for (auto sharding_and_retval :
llvm::zip(sharding_for_rets, terminator->getOpOperands())) {
const auto& sharding = std::get<0>(sharding_and_retval);
OpOperand& retval = std::get<1>(sharding_and_retval);
if (failed(VerifySharding(retval.get().getType(), sharding)))
return mlir::failure();
}
return mlir::success();
}
// Assign the logical device if an op has an attribute `TPU_REPLICATED_CORE:n`,
// the corresponding input sharding arg will be associated with
// logical device `n`.
std::optional<llvm::StringRef> AssignLogicalDeviceFromTPUReplicatedCoreAttr(
Operation* op, const llvm::SmallVector<std::string>& logical_device_vec) {
if (auto device = op->getAttrOfType<StringAttr>("device")) {
if (!device.getValue().empty() && !device.getValue().str().empty()) {
tensorflow::DeviceNameUtils::ParsedName name;
if (tensorflow::DeviceNameUtils::ParseFullName(device.str(), &name)) {
if (name.type == "TPU_REPLICATED_CORE") {
// TODO(hanxiongwang): Add check for out of bound of name.id
return llvm::StringRef(logical_device_vec[name.id]);
}
}
}
}
return std::nullopt;
}
absl::StatusOr<std::optional<llvm::StringRef>> GetXlaShardingFromShardingOp(
mlir::TF::XlaShardingOp sharding) {
TF_ASSIGN_OR_RETURN(auto attr, GetXlaShardingAttrFromShardingOp(sharding));
return attr ? ::std::optional<::llvm::StringRef>(attr.getValue())
: (::std::nullopt);
}
// Returns XLA sharding from a XlaSharding op connected to an argument value. If
// value is a resource type then XlaSharding op will be connected to a
// ReadVariable op. XlaSharding op may be direct user of inputs but it may also
// be followed by an Identity op and, in the case where bfloat16 type is used,
// Cast op may be added right after the input.
//
// TODO(hongjunchoi): Add logic to parse XlaSharding op inside control flow (If,
// Case) ops and Caller return values.
// TODO(hongjunchoi): Consider explicitly checking op patterns to detect sharded
// inputs.
absl::StatusOr<std::optional<llvm::StringRef>> GetXlaShardingFromArg(
Value value, const llvm::SmallVector<std::string>& logical_device_vec) {
llvm::SmallPtrSet<Value, 4> visited_values;
llvm::SmallVector<Value, 4> values_to_visit{value};
while (!values_to_visit.empty()) {
llvm::SmallVector<Value, 4> next_values_to_visit;
for (Value value_to_visit : values_to_visit) {
if (!visited_values.insert(value_to_visit).second) continue;
for (auto& use : value_to_visit.getUses()) {
Operation* owner = use.getOwner();
if (auto sharding = llvm::dyn_cast<mlir::TF::XlaShardingOp>(owner))
return GetXlaShardingFromShardingOp(sharding);
if (auto logical_device = AssignLogicalDeviceFromTPUReplicatedCoreAttr(
owner, logical_device_vec)) {
return logical_device;
}
if (auto while_op = llvm::dyn_cast<mlir::TF::WhileRegionOp>(owner)) {
const int operand_number = use.getOperandNumber();
next_values_to_visit.push_back(
while_op.getCond().front().getArgument(operand_number));
next_values_to_visit.push_back(
while_op.getBody().front().getArgument(operand_number));
continue;
}
if (UnaryOpHasTraitsForSharding(owner)) {
next_values_to_visit.push_back(use.getOwner()->getResult(0));
continue;
}
if (BinaryOpHasTraitsForSharding(owner)) {
if (DoTypesHavePartialSameShape(value_to_visit,
owner->getResult(0))) {
next_values_to_visit.push_back(use.getOwner()->getResult(0));
continue;
}
}
if (llvm::isa<mlir::TF::CastOp, mlir::TF::XlaAllReduceOp,
mlir::TF::ReadVariableOp>(owner)) {
next_values_to_visit.push_back(use.getOwner()->getResult(0));
continue;
}
if (auto call_op = llvm::dyn_cast<mlir::CallOpInterface>(owner)) {
mlir::func::FuncOp func =
llvm::dyn_cast<mlir::func::FuncOp>(call_op.resolveCallable());
if (!func) continue;
next_values_to_visit.push_back(
func.getArgument(use.getOperandNumber()));
}
}
}
values_to_visit.swap(next_values_to_visit);
}
return std::nullopt;
}
// Tries to extract sharding configurations for all inputs by parsing
// XlaSharding/ TPUPartitionedInput op connected to the operands/arguments. If
// argument to the `cluster_func` directly feeds into another function call op,
// then recursively walk the function definition to find the connected
// XlaSharding op.
absl::Status IdentifyXlaShardingForComputationInputs(
const llvm::SmallVector<std::string>& logical_device_vec,
bool infer_from_computation, mlir::tf_device::ClusterFuncOp cluster_func,
mlir::func::FuncOp func, Builder* builder,
OptionalOpShardingVector& sharding_for_args) {
// Look up function definition from module.
Block& function_block = func.front();
sharding_for_args.reserve(function_block.getNumArguments());
// Iterate through operands of `cluster_func`.
// The computation operand can either be:
// 1) a TPUPartitionedInput Op if the input has a non-resource type;
// 2) a ReadVariableOp else.
//
// Iterate through input arguments to the entry block of
// tf_device.ClusterFunc. For input ops, look for XlaSharding ops.
// XlaSharding ops can:
// 1) Directly follow the input argument if input argument has non-resource
// types.
// 2) Follow ReadVariableOp if the input type is of resource type.
// 3) Follow IdentityOp or CastOp after above cases (1), (2).
//
// Sharding configurations are added to the tf_device.ClusterFunc as an
// attribute and the function as an argument attribute.
for (auto operand_and_arg :
llvm::zip(cluster_func.getOperands(), function_block.getArguments())) {
Value operand = std::get<0>(operand_and_arg);
BlockArgument arg = std::get<1>(operand_and_arg);
if (auto operand_sharding = GetXlaShardingFromOperand(operand)) {
sharding_for_args.push_back(operand_sharding);
continue;
}
if (infer_from_computation) {
TF_ASSIGN_OR_RETURN(auto arg_sharding,
GetXlaShardingFromArg(arg, logical_device_vec));
if (arg_sharding) {
sharding_for_args.push_back(arg_sharding.value());
continue;
}
}
sharding_for_args.push_back(std::nullopt);
}
return absl::OkStatus();
}
// Returns a TPUPartitionedOutput or TPUPartitionedInput op with XLA sharding
// connected to a `tf_device.cluster_func` result value (via AssignVariableOp/
// resource write).
mlir::Operation* GetXlaShardingFromResult(Value value) {
if (!value.hasOneUse()) return nullptr;
Operation* user = *value.getUsers().begin();
if (auto partitioned_output =
llvm::dyn_cast<mlir::TF::TPUPartitionedOutputV2Op>(user))
return NullUnlessSharded(partitioned_output);
if (auto assign_var = llvm::dyn_cast<mlir::TF::AssignVariableOp>(user))
if (auto partitioned_input =
assign_var.getResource()
.getDefiningOp<mlir::TF::TPUPartitionedInputV2Op>())
return NullUnlessSharded(partitioned_input);
return nullptr;
}
absl::Status DetermineShardingFromAlias(
mlir::func::FuncOp func, OptionalOpShardingVector& input_shardings,
OptionalOpShardingVector& output_shardings) {
for (int arg_idx = 0; arg_idx < func.getNumArguments(); ++arg_idx) {
if (auto v =
func.getArgAttrOfType<mlir::IntegerAttr>(arg_idx, kAliasingAttr)) {
if (int retval_idx = v.getInt();
retval_idx >= 0 && retval_idx < func.getNumResults()) {
auto& input_sharding = input_shardings[arg_idx];
auto& output_sharding = output_shardings[retval_idx];
if (input_sharding.has_value() && output_sharding.has_value() &&
input_sharding.value() != output_sharding.value()) {
return absl::InvalidArgumentError(absl::StrCat(
"arg#", arg_idx, " is aliased to retval#", retval_idx,
" but their sharding configurations don't match."));
} else if (input_sharding.has_value() && !output_sharding.has_value()) {
output_sharding = input_sharding;
} else if (!input_sharding.has_value() && output_sharding.has_value()) {
input_sharding = output_sharding;
}
}
}
}
return absl::OkStatus();
}
// Returns XLA sharding from XlaSharding op connected to a result value.
// XlaSharding op may be directly connected to output but it may also be
// followed by Identity or simple arithmetic ops. In case where bfloat16 type is
// used, we might see a Cast op.
//
// TODO(hongjunchoi): Add logic to parse XlaSharding op inside control flow (If,
// Case) ops and Caller argument values.
// TODO(hongjunchoi): Consider explicitly checking op patterns to detect sharded
// inputs.
absl::StatusOr<std::optional<StringRef>> GetXlaShardingFromRetval(
Value value, const llvm::SmallVector<std::string>& logical_device_vec) {
llvm::SmallPtrSet<Value, 4> visited_values;
llvm::SmallVector<Value, 4> values_to_visit;
values_to_visit.push_back(value);
while (!values_to_visit.empty()) {
Value value_to_visit = values_to_visit.pop_back_val();
if (!visited_values.insert(value_to_visit).second) {
continue;
}
Operation* def = value_to_visit.getDefiningOp();
if (!def) {
continue;
}
if (auto sharding = llvm::dyn_cast_or_null<mlir::TF::XlaShardingOp>(def))
return GetXlaShardingFromShardingOp(sharding);
if (auto sharding = def->getAttrOfType<StringAttr>(kShardingAttributeV2)) {
return sharding.strref();
}
if (auto sharding = def->getAttrOfType<StringAttr>(kShardingAttribute)) {
return sharding.strref();
}
if (auto logical_device = AssignLogicalDeviceFromTPUReplicatedCoreAttr(
def, logical_device_vec)) {
return logical_device;
}
if (UnaryOpHasTraitsForSharding(def) || BinaryOpHasTraitsForSharding(def)) {
for (auto operand : def->getOperands()) {
values_to_visit.push_back(operand);
}
continue;
}
if (auto call_op = llvm::dyn_cast_or_null<mlir::CallOpInterface>(def)) {
mlir::func::FuncOp func =
llvm::dyn_cast<mlir::func::FuncOp>(call_op.resolveCallable());
if (!func) continue;
value_to_visit = func.front().getTerminator()->getOperand(
mlir::cast<OpResult>(value_to_visit).getResultNumber());
values_to_visit.push_back(value_to_visit);
continue;
}
if (auto while_op = llvm::dyn_cast<mlir::TF::WhileRegionOp>(def)) {
if (auto op_result = mlir::cast<OpResult>(value_to_visit)) {
int result_idx = op_result.getResultNumber();
if (auto yield_op = llvm::dyn_cast<mlir::TF::YieldOp>(
while_op.getBody().front().getTerminator())) {
values_to_visit.push_back(yield_op.getOperand(result_idx));
}
}
continue;
}
}
return std::nullopt;
}
// Tries to extract sharding configurations for all outputs by parsing
// XlaSharding/ TPUPartitionedOutput op connected to the retvals/results.
absl::Status IdentifyXlaShardingForComputationOutputs(
const llvm::SmallVector<std::string>& logical_device_vec,
bool infer_from_computation, mlir::tf_device::ClusterFuncOp cluster_func,
mlir::func::FuncOp func, Builder* builder,
OptionalOpShardingVector& sharding_for_rets) {
Block& function_block = func.front();
Operation* terminator = function_block.getTerminator();
sharding_for_rets.reserve(terminator->getNumOperands());
// Iterate through results of `cluster_func`. For output ops, look for
// TPUPartitionedOutput ops.
//
// Iterate through operands of the terminator. If the preceding op is
// XlaShardingOp, then the provided sharding configuration is added to the
// tf_device.ClusterFunc as an attribute and the function as a result
// attribute.
for (auto result_and_retval :
llvm::zip(cluster_func.getResults(), terminator->getOpOperands())) {
Value result = std::get<0>(result_and_retval);
OpOperand& retval = std::get<1>(result_and_retval);
if (auto result_sharding = GetXlaShardingFromResult(result)) {
sharding_for_rets.push_back(result_sharding);
continue;
}
if (infer_from_computation) {
TF_ASSIGN_OR_RETURN(
auto retval_sharding,
GetXlaShardingFromRetval(retval.get(), logical_device_vec));
if (retval_sharding) {
sharding_for_rets.push_back(retval_sharding.value());
continue;
}
}
sharding_for_rets.push_back(std::nullopt);
}
return absl::OkStatus();
}
void SetReplicatedOrMaximalShardingIfNoShardingFound(
const llvm::SmallVector<std::string>& logical_device_vec, bool use_spmd,
OptionalOpShardingVector& shardings) {
for (auto& sharding : shardings) {
if (sharding == std::nullopt) {
// If we haven't found sharding, default to either replicated or maximal
// sharding depending on whether XLA SPMD is enabled.
if (use_spmd) {
// If XLA SPMD is enabled, host variables or non-variable per-replica
// inputs, and outputs should take on replicate sharding, so that every
// device gets the whole tensor(s) (and can slice them up later eg.
// using dynamic-slice).
sharding = kReplicateSharding;
} else {
// Otherwise, default to maximal sharding core 0.
sharding = logical_device_vec[0];
}
}
}
}
// Moves shardings from `optional_shardings` to `shardings`.
absl::Status MoveSharding(OptionalOpShardingVector& optional_shardings,
OpShardingVector& shardings) {
shardings.clear();
for (auto& sharding : optional_shardings) {
if (!sharding) {
return absl::InternalError(
"Couldn't find/assign sharding for an input/output. All shardings "
"should have been identified by this point.");
}
shardings.push_back(std::move(sharding.value()));
}
return absl::OkStatus();
}
// Determines XlaSharding for inputs and outputs. If there are aliased
// inputs/outputs for which no sharding was found directly, the corresponding
// output/input sharding is used (if it exists). If we still don't find sharding
// for some inputs/outputs, we default to replicated or maximal sharding
// depending on `use_spmd`.
absl::Status IdentifyXlaShardingForInputsAndOutputs(
const llvm::SmallVector<std::string>& logical_device_vec, bool use_spmd,
bool infer_from_computation, mlir::tf_device::ClusterFuncOp cluster_func,
mlir::func::FuncOp func, Builder* builder, OpShardingVector& input_sharding,
OpShardingVector& output_sharding) {
OptionalOpShardingVector optional_input_sharding;
OptionalOpShardingVector optional_output_sharding;
TF_RETURN_IF_ERROR(IdentifyXlaShardingForComputationInputs(
logical_device_vec, infer_from_computation, cluster_func, func, builder,
optional_input_sharding));
TF_RETURN_IF_ERROR(IdentifyXlaShardingForComputationOutputs(
logical_device_vec, infer_from_computation, cluster_func, func, builder,
optional_output_sharding));
TF_RETURN_IF_ERROR(DetermineShardingFromAlias(func, optional_input_sharding,
optional_output_sharding));
SetReplicatedOrMaximalShardingIfNoShardingFound(logical_device_vec, use_spmd,
optional_input_sharding);
SetReplicatedOrMaximalShardingIfNoShardingFound(logical_device_vec, use_spmd,
optional_output_sharding);
TF_RETURN_IF_ERROR(MoveSharding(optional_input_sharding, input_sharding));
TF_RETURN_IF_ERROR(MoveSharding(optional_output_sharding, output_sharding));
return absl::OkStatus();
}
// Extracts input/output sharding configuration of `cluster_func` by parsing
// XlaSharding ops inside the `cluster_func`.
LogicalResult IdentifyXlaShardingForTPUComputation(
Builder* builder, mlir::tf_device::ClusterFuncOp cluster_func) {
// Look up function definition from module.
mlir::func::FuncOp func =
cluster_func->getParentOfType<ModuleOp>()
.lookupSymbol<mlir::func::FuncOp>(cluster_func.getFunc());
bool use_spmd = false;
if (auto use_spmd_attr = cluster_func->getAttrOfType<BoolAttr>(kUseSpmdAttr))
use_spmd = use_spmd_attr.getValue();
auto num_cores_per_replica_attr =
cluster_func->getAttrOfType<IntegerAttr>(kNumCoresPerReplicaAttr);
if (!num_cores_per_replica_attr)
return cluster_func.emitOpError(
CreateMissingAttributeMsg(kNumCoresPerReplicaAttr));
int num_cores_per_replica = num_cores_per_replica_attr.getInt();
llvm::SmallVector<std::string> logical_device_vec(num_cores_per_replica);
for (int idx = 0; idx < num_cores_per_replica; idx++) {
logical_device_vec[idx] =
xla::sharding_builder::SingleDevice(idx).SerializeAsString();
}
OpShardingVector sharding_for_args;
OpShardingVector sharding_for_rets;
if (auto status = IdentifyXlaShardingForInputsAndOutputs(
logical_device_vec, use_spmd,
/*infer_from_computation=*/true, cluster_func, func, builder,
sharding_for_args, sharding_for_rets);
!status.ok()) {
return cluster_func.emitOpError(status.message());
};
auto has_maximal_sharding =
[](const OpShardingVariant& sharding_or_op) -> bool {
const auto sharding = GetShardingFromVariant(sharding_or_op);
return sharding && sharding->type() == xla::OpSharding::MAXIMAL;
};
// XLA SPMD only supports cases where all inputs/outputs exist on every
// partition (sharded or replicated). If any of the inputs/outputs have
// maximal sharding, then fallback to MPMD. Also fall back if any of the
// shardings aren't compatible with the rank of their tensor.
if ((use_spmd && (absl::c_any_of(sharding_for_args, has_maximal_sharding) ||
absl::c_any_of(sharding_for_rets, has_maximal_sharding))) ||
failed(VerifyShardings(func, sharding_for_args, sharding_for_rets))) {
LOG(WARNING) << "XLA SPMD only supports cases where all inputs/outputs "
"exist on every partition (sharded or replicated). If any "
"of the inputs/outputs have maximal sharding, then "
"fallback to MPMD.";
sharding_for_args.clear();
sharding_for_rets.clear();
cluster_func->setAttr(kUseSpmdAttr, builder->getBoolAttr(false));
if (auto status = IdentifyXlaShardingForInputsAndOutputs(
logical_device_vec, /*use_spmd=*/false,
/*infer_from_computation=*/false, cluster_func, func, builder,
sharding_for_args, sharding_for_rets);
!status.ok()) {
return cluster_func.emitOpError(status.message());
}
}
// Update sharding on function arguments and returns.
Block& function_block = func.front();
for (auto sharding_and_arg :
llvm::zip(sharding_for_args, function_block.getArguments())) {
BlockArgument arg = std::get<1>(sharding_and_arg);
const auto& sharding_or_op = std::get<0>(sharding_and_arg);
if (auto sharding = GetShardingStringFromVariant(sharding_or_op)) {
func.setArgAttr(arg.getArgNumber(), kShardingAttr,
builder->getStringAttr(*sharding));
}
}
Operation* terminator = function_block.getTerminator();
for (auto sharding_and_retval :
llvm::zip(sharding_for_rets, terminator->getOpOperands())) {
OpOperand& retval = std::get<1>(sharding_and_retval);
const auto& sharding_or_op = std::get<0>(sharding_and_retval);
if (auto sharding = GetShardingStringFromVariant(sharding_or_op)) {
func.setResultAttr(retval.getOperandNumber(), kShardingAttr,
builder->getStringAttr(*sharding));
}
}
// Update input/output sharding attributes on tf_device.cluster_func op.
cluster_func->setAttr(tensorflow::kInputShardingAttr,
GetStrArrayAttr(builder, sharding_for_args));
cluster_func->setAttr(tensorflow::kOutputShardingAttr,
GetStrArrayAttr(builder, sharding_for_rets));
return mlir::success();
}
void TPUShardingIdentificationPass::runOnOperation() {
Builder builder(getOperation().getContext());
auto result =
getOperation().walk([&](mlir::tf_device::ClusterFuncOp cluster_func) {
if (failed(
IdentifyXlaShardingForTPUComputation(&builder, cluster_func))) {
return WalkResult::interrupt();
}
return WalkResult::advance();
});
if (result.wasInterrupted()) return signalPassFailure();
}
} // namespace
std::unique_ptr<mlir::OperationPass<ModuleOp>>
CreateTPUShardingIdentificationPass() {
return std::make_unique<TPUShardingIdentificationPass>();
}
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/passes/tpu_validate_inputs.cc
+645
View File
@@ -0,0 +1,645 @@
/* 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 <functional>
#include <memory>
#include <optional>
#include <queue>
#include <string>
#include <unordered_map>
#include "absl/log/log.h"
#include "absl/strings/match.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/Support/Casting.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/TypeUtilities.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/tpu_rewrite_device_util.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/tpu_validate_inputs_utils.h"
#include "xla/xla_data.pb.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {
#define GEN_PASS_DEF_TPUVALIDATEINPUTSPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc"
constexpr char kXLAShardingAttr[] = "_XlaSharding";
constexpr char kShardingAttr[] = "sharding";
using mlir::dyn_cast;
using mlir::isa;
using mlir::ModuleOp;
using mlir::Operation;
using mlir::OperationPass;
using mlir::StringAttr;
using mlir::StringRef;
using mlir::Type;
using mlir::TypeID;
using mlir::func::FuncOp;
using mlir::func::ReturnOp;
using mlir::TF::AssertOp;
using mlir::TF::ConstOp;
using mlir::TF::kCompileDeviceTypeAttr;
using mlir::TF::kTpuReplicateAttr;
using mlir::TF::OutfeedEnqueueTupleOp;
using mlir::TF::PartitionedCallOp;
using mlir::TF::StatefulPartitionedCallOp;
using mlir::TF::TPUPartitionedCallOp;
using mlir::TF::TPUPartitionedInputOp;
using mlir::TF::TPUPartitionedInputV2Op;
using mlir::TF::TPUPartitionedOutputOp;
using mlir::TF::TPUPartitionedOutputV2Op;
using mlir::TF::TPUReplicatedInputOp;
using mlir::TF::TPUReplicatedOutputOp;
using mlir::TF::TPUReplicateMetadataOp;
using mlir::TF::WhileOp;
using mlir::TF::XlaSetDynamicDimensionSizeOp;
using mlir::tf_executor::FetchOp;
using mlir::tf_executor::GraphOp;
using mlir::tf_executor::IslandOp;
using mlir::tf_executor::YieldOp;
typedef std::unordered_map<std::string, TPUReplicateMetadataOp> MetadataMap;
struct TPUValidateInputsPass
: public impl::TPUValidateInputsPassBase<TPUValidateInputsPass> {
void runOnOperation() override;
};
bool IsTpuRegularOp(Operation* op) {
static auto* ops = [] {
llvm::SmallDenseSet<mlir::TypeID, 32>* ops_set =
new llvm::SmallDenseSet<mlir::TypeID, 32>{
TypeID::get<ModuleOp>(),
TypeID::get<GraphOp>(),
TypeID::get<ReturnOp>(),
TypeID::get<FuncOp>(),
TypeID::get<YieldOp>(),
TypeID::get<IslandOp>(),
TypeID::get<TPUReplicatedInputOp>(),
TypeID::get<TPUReplicatedOutputOp>(),
TypeID::get<TPUPartitionedInputOp>(),
TypeID::get<TPUPartitionedInputV2Op>(),
TypeID::get<TPUPartitionedOutputOp>(),
TypeID::get<TPUPartitionedOutputV2Op>(),
TypeID::get<TPUReplicateMetadataOp>(),
TypeID::get<FetchOp>(),
TypeID::get<OutfeedEnqueueTupleOp>(),
};
return ops_set;
}();
auto abstractOp = op->getRegisteredInfo();
if (!abstractOp) return true;
return ops->count(abstractOp->getTypeID()) == 0;
}
bool IsIntersectionXlaNonXlaOps(Operation* op) {
static auto* ops = [] {
llvm::SmallDenseSet<mlir::TypeID, 32>* ops_set =
new llvm::SmallDenseSet<mlir::TypeID, 32>{
TypeID::get<ConstOp>(),
TypeID::get<WhileOp>(),
TypeID::get<AssertOp>(),
TypeID::get<XlaSetDynamicDimensionSizeOp>(),
};
return ops_set;
}();
auto abstractOp = op->getRegisteredInfo();
if (!abstractOp) return true;
return ops->count(abstractOp->getTypeID()) == 0;
}
bool IsPartitionedOp(Operation* op) {
static auto* ops = [] {
llvm::SmallDenseSet<mlir::TypeID, 32>* ops_set =
new llvm::SmallDenseSet<mlir::TypeID, 32>{
TypeID::get<StatefulPartitionedCallOp>(),
TypeID::get<PartitionedCallOp>(),
TypeID::get<TPUPartitionedCallOp>(),
};
return ops_set;
}();
auto abstractOp = op->getRegisteredInfo();
if (!abstractOp) return false;
return ops->count(abstractOp->getTypeID()) != 0;
}
// Gets the successors of an op wrapped in a tf_executor.island.
llvm::SmallVector<Operation*> GetSuccessors(Operation* op) {
llvm::SmallVector<Operation*> successors;
for (auto result : op->getParentOp()->getOpResults()) {
for (auto& use : result.getUses()) {
auto succ = use.getOwner();
successors.push_back(succ);
}
}
return successors;
}
// Gets the predecessors of an op wrapped in tf_executor.island.
llvm::SmallVector<Operation*> GetPredecessors(Operation* op) {
llvm::SmallVector<Operation*> predecessors;
for (auto operand : op->getOperands()) {
if (Operation* pred = operand.getDefiningOp()) {
pred->walk([&](mlir::Operation* opinexecutor) {
predecessors.push_back(opinexecutor);
});
}
}
return predecessors;
}
bool CheckTpuReplicateAttr(Operation* op, StringAttr attr,
std::function<std::string()> errormsg) {
if (!op->hasAttr(kTpuReplicateAttr)) {
op->emitOpError("TF2XLA TPU bridge input check: " + errormsg() +
"missing _tpu_replicate attr");
return false;
}
auto opattr = op->getAttr(kTpuReplicateAttr);
if (opattr != attr) {
op->emitOpError("TF2XLA TPU bridge input check: " + errormsg() +
"invalid _tpu_replicate attr.")
<< " Expected attr: " << attr << ", Actual attr: " << opattr;
return false;
}
return true;
}
bool ValidateReplicatedInput(TPUReplicatedInputOp rep, int num_replicas,
StringAttr attr) {
int arity = rep.getInputs().size();
if (rep.getIsPacked() && arity != 1) {
rep.emitOpError(
"TF2XLA TPU bridge input check: packed with number of inputs not 1.")
<< " num_replicas=" << num_replicas << " no. of inputs=" << arity;
return false;
} else if (!rep.getIsPacked() && arity != num_replicas) {
rep.emitOpError(
"TF2XLA TPU bridge input check: number of inputs inconsistent.")
<< " num_replicas=" << num_replicas << " no. of inputs=" << arity;
return false;
}
for (auto& succ : GetSuccessors(rep)) {
if (!IsTpuRegularOp(succ)) continue;
auto errormsg = [&]() -> std::string {
return rep->getName().getStringRef().str() + " op has successor op " +
succ->getName().getStringRef().str() + " with error: ";
};
if (!CheckTpuReplicateAttr(succ, attr, errormsg)) return false;
}
return true;
}
bool ValidateReplicatedOutput(TPUReplicatedOutputOp rep, int num_replicas,
StringAttr attr) {
int arity = rep.getOutputs().size();
if (arity != num_replicas) {
rep.emitOpError(
"TF2XLA TPU bridge input check: number of outputs inconsistent.")
<< " num_replicas=" << num_replicas << " no. of outputs=" << arity;
return false;
}
for (auto& pred : GetPredecessors(rep)) {
if (!IsTpuRegularOp(pred)) continue;
auto errormsg = [&]() -> std::string {
return rep->getName().getStringRef().str() + " op has predecessor op " +
pred->getName().getStringRef().str() + " with error: ";
};
if (!CheckTpuReplicateAttr(pred, attr, errormsg)) return false;
}
return true;
}
bool ValidatePartitionedInput(TPUPartitionedInputOp rep,
int num_cores_per_replica) {
int arity = rep.getInputs().size();
if (arity != num_cores_per_replica) {
rep.emitOpError(
"TF2XLA TPU bridge input check: number of inputs inconsistent.")
<< " num_cores_per_replica=" << num_cores_per_replica
<< " no. of inputs=" << arity;
return false;
}
return true;
}
bool ValidatePartitionedInputV2(TPUPartitionedInputV2Op rep,
int num_cores_per_replica) {
int arity = rep.getInputs().size();
if (rep.getIsPacked() && arity != 1) {
rep.emitOpError(
"TF2XLA TPU bridge input check: packed with number of inputs not 1.")
<< " num_cores_per_replicas=" << num_cores_per_replica
<< " no. of inputs=" << arity;
return false;
} else if (!rep.getIsPacked() && arity != num_cores_per_replica) {
rep.emitOpError(
"TF2XLA TPU bridge input check: number of inputs inconsistent.")
<< " num_cores_per_replica=" << num_cores_per_replica
<< " no. of inputs=" << arity;
return false;
}
return true;
}
template <typename T>
bool ValidatePartitionedOutput(T rep, int num_cores_per_replica) {
int arity = rep.getOutput().size();
if (arity != num_cores_per_replica) {
rep.emitOpError(
"TF2XLA TPU bridge input check: number of outputs inconsistent.")
<< " num_cores_per_replica=" << num_cores_per_replica
<< " no. of outputs=" << arity;
return false;
}
return true;
}
bool CheckReplicatedIOOp(Operation* op, TPUReplicateMetadataOp metadata,
Operation* parent) {
int num_replicas = metadata.getNumReplicas();
int num_cores_per_replica = metadata.getNumCoresPerReplica();
StringAttr tpu_replicate_attr =
metadata->getAttrOfType<StringAttr>(kTpuReplicateAttr);
if (auto repinput = dyn_cast<TPUReplicatedInputOp>(op)) {
if (!ValidateReplicatedInput(repinput, num_replicas, tpu_replicate_attr))
return false;
}
if (auto repoutput = dyn_cast<TPUReplicatedOutputOp>(op)) {
if (!ValidateReplicatedOutput(repoutput, num_replicas, tpu_replicate_attr))
return false;
}
if (auto partinput = dyn_cast<TPUPartitionedInputOp>(op)) {
if (!ValidatePartitionedInput(partinput, num_cores_per_replica))
return false;
}
if (auto partinput = dyn_cast<TPUPartitionedInputV2Op>(op)) {
if (!ValidatePartitionedInputV2(partinput, num_cores_per_replica))
return false;
}
if (auto partoutput = dyn_cast<TPUPartitionedOutputOp>(op)) {
if (!ValidatePartitionedOutput(partoutput, num_cores_per_replica))
return false;
}
if (auto partoutput = dyn_cast<TPUPartitionedOutputV2Op>(op)) {
if (!ValidatePartitionedOutput(partoutput, num_cores_per_replica))
return false;
}
return true;
}
// Checking op which is successor to a cluster op.
bool CheckClusterSuccessors(Operation* op, std::string cluster,
Operation* parent, MetadataMap& metadata_map) {
std::string cluster_succ = "";
if (op->hasAttr(kTpuReplicateAttr)) {
cluster_succ = op->getAttrOfType<StringAttr>(kTpuReplicateAttr).str();
}
if (cluster_succ.empty()) {
// TODO (b/269195256#comment16): Change to error after resolving issue
// with test. Will fix it after the upstream code is fixed.
op->emitWarning("TF2XLA TPU bridge input check: cluster op = ")
<< parent->getName() << " with cluster = " << cluster
<< " has successor as non cluster op " << op->getName();
return true;
}
if (cluster != cluster_succ) {
op->emitOpError(
"TF2XLA TPU bridge input check: mismatch clusters tpu_replicate "
"attr. Parent op ")
<< parent->getName() << " with cluster = " << cluster
<< " has successor cluster op " << op->getName()
<< " with cluster = " << cluster_succ;
return false;
}
return true;
}
// Checking op which is a predecessor to a non-cluster op.
bool CheckNonClusterSuccessors(Operation* op, Operation* parent,
MetadataMap& metadata_map) {
if (!IsTpuRegularOp(op)) {
if (isa<TPUReplicatedOutputOp>(op)) {
op->emitOpError("TF2XLA TPU bridge input check: non-cluster op = ")
<< parent->getName()
<< " has invalid successor op = " << op->getName();
return false;
} else {
return true;
}
}
return true;
}
// Checking op which is a successor to a non-cluster op.
bool CheckNonClusterPredecessors(Operation* op, Operation* parent,
MetadataMap& metadata_map) {
if (!IsTpuRegularOp(op)) {
if (isa<TPUReplicatedInputOp>(op)) {
op->emitOpError("TF2XLA TPU bridge input check: non-cluster op = ")
<< parent->getName()
<< " has invalid predecessor op = " << op->getName();
return false;
} else {
return true;
}
}
return true;
}
bool CheckOpsClusterIO(Operation* op, MetadataMap& metadata_map) {
bool is_cluster_op = false;
std::string cluster = "";
if (op->hasAttr(kTpuReplicateAttr)) {
cluster = op->getAttrOfType<StringAttr>(kTpuReplicateAttr).str();
if (cluster.empty()) {
op->emitOpError("TF2XLA TPU bridge input check: empty _tpu_replicate")
<< " attr for op = " << op->getName();
return false;
}
is_cluster_op = true;
}
bool has_cluster_metadata =
(metadata_map.find(cluster) != metadata_map.end());
for (auto pred : GetPredecessors(op)) {
if (is_cluster_op && !IsTpuRegularOp(pred) && has_cluster_metadata) {
if (!CheckReplicatedIOOp(pred, metadata_map[cluster], op)) return false;
}
if (!is_cluster_op) {
if (!CheckNonClusterPredecessors(pred, op, metadata_map)) return false;
}
}
for (auto succ : GetSuccessors(op)) {
if (is_cluster_op && !IsTpuRegularOp(succ) && has_cluster_metadata) {
if (!CheckReplicatedIOOp(succ, metadata_map[cluster], op)) return false;
}
if (is_cluster_op && IsTpuRegularOp(succ)) {
if (!CheckClusterSuccessors(succ, cluster, op, metadata_map))
return false;
}
if (!is_cluster_op) {
if (!CheckNonClusterSuccessors(succ, op, metadata_map)) return false;
}
}
return true;
}
bool TypeMustBeNonXLA(const Type& type) {
const Type elem = getElementTypeOrSelf(type);
return !mlir::isa<mlir::TF::ResourceType>(elem) &&
!tensorflow::TypeValidForXLA(type);
}
// Check if the op cannot be XLA compiled. If the op does not satisfy this
// criteria, then it is possible for the op to be XLA and non-XLA. But this
// function specifically checks if the op must be non-xla.
bool IsMustNotBeXlaOp(Operation* op) {
for (auto& input : op->getOpOperands()) {
if (TypeMustBeNonXLA(input.get().getType())) return true;
}
for (auto output_types : op->getResultTypes()) {
if (TypeMustBeNonXLA(output_types)) return true;
}
return false;
}
// Check if the op must be compiled with XLA. If the op does not satisfy this
// critiria for "must be xla" then it is still possible for this op to be xla
// and non-xla as well. But below function specifically checks for the op to be
// only XLA op.
bool IsMustBeXlaOp(Operation* op, MetadataMap metadata_map) {
// All PartitionedCall are inlined-out before XLA.
// So MustBeXLA should return false
if (IsPartitionedOp(op)) return false;
if (!op->hasAttr(kTpuReplicateAttr)) return false;
auto cluster = op->getAttrOfType<StringAttr>(kTpuReplicateAttr).str();
if (metadata_map.find(cluster) == metadata_map.end()) return false;
auto metadata = metadata_map[cluster];
if (!metadata.getAllowSoftPlacement() &&
!op->hasAttr(mlir::TF::kXlaOutsideCompilationAttr))
return true;
std::string device = "";
if (op->hasAttr(mlir::TF::kDeviceAttr))
device = op->getAttrOfType<StringAttr>(mlir::TF::kDeviceAttr).str();
else
return false;
if (absl::StrContains(device, mlir::TF::kTpuDevice)) return true;
return false;
}
bool ValidateIntersectionXlaNonXlaOps(Operation* op, MetadataMap metadata_map) {
if (isa<TPUReplicateMetadataOp>(op) || isa<TPUReplicatedInputOp>(op) ||
isa<TPUReplicatedOutputOp>(op) || isa<TPUPartitionedInputOp>(op) ||
isa<TPUPartitionedInputV2Op>(op) || isa<TPUPartitionedOutputOp>(op) ||
isa<TPUPartitionedOutputV2Op>(op))
return true;
if (IsMustBeXlaOp(op, metadata_map) && IsMustNotBeXlaOp(op)) {
// TODO(b/269195256#comment19) change the warning for Identity op to error
// when issue with input graph is resolved. Possible issue with python layer
// inserting Identity op incorrectly.
if (isa<mlir::TF::IdentityOp>(op)) {
op->emitWarning("TF/XLA TPU bridge input check: found invalid op. ")
<< op->getName() << " can't be both xla and non-xla";
return true;
}
op->emitOpError("TF/XLA TPU bridge input check: found invalid op. ")
<< "Can't be both xla and non-xla";
return false;
}
return true;
}
void GetFlattenedShardings(llvm::SmallVector<xla::OpSharding>& shardings_result,
std::string shard_string) {
xla::OpSharding sharding;
if (shard_string.empty()) return;
if (!sharding.ParseFromString(shard_string)) return;
std::queue<xla::OpSharding> shardings;
shardings.push(sharding);
while (!shardings.empty()) {
auto sharding_next = shardings.front();
shardings.pop();
if (sharding_next.type() == xla::OpSharding::TUPLE) {
for (auto& child_sharding : sharding_next.tuple_shardings()) {
shardings.push(child_sharding);
}
} else {
shardings_result.push_back(sharding_next);
}
}
}
bool IsValidShardingTupleForArity(Operation* op) {
if (!op->hasAttr(kXLAShardingAttr) && !op->hasAttr(kShardingAttr)) {
return true;
}
std::string shard_string;
if (op->hasAttr(kXLAShardingAttr)) {
shard_string =
op->getAttrOfType<StringAttr>(kXLAShardingAttr).strref().str();
} else {
shard_string = op->getAttrOfType<StringAttr>(kShardingAttr).strref().str();
}
xla::OpSharding sharding;
if (!shard_string.empty() && sharding.ParseFromString(shard_string)) {
// Only checking op with TUPLE sharding
if (sharding.type() != xla::OpSharding::TUPLE) {
return true;
}
// Each output is expected to have a corresponding sharding given by
// tuple shardings. So, the no. of outputs (arity) should be same as
// number of tuple shardings.
if (sharding.tuple_shardings().size() != op->getNumResults()) {
op->emitOpError(
"TF2XLA TPU bridge input check: invalid no. of tuple shardings ")
<< sharding.tuple_shardings().size()
<< " for arity = " << op->getNumResults() << "\n The sharding is "
<< sharding.DebugString() << "\n";
return false;
}
}
return true;
}
bool IsValidMAXIMALSharding(Operation* op, MetadataMap& metadata_map) {
if (!op->hasAttr(kTpuReplicateAttr)) return true;
if (!op->hasAttr(kXLAShardingAttr) && !op->hasAttr(kShardingAttr)) {
return true;
}
int num_cores_per_replica;
// Assuming that the op is a IsTpuRegularOp and has a cluster metadata
// for it. These checks are already performed in CheckOpsClusterIO.
// Also assuming that if there is sharding, then there must be
// cluster and the metadata corresponding to it.
auto cluster = op->getAttrOfType<StringAttr>(kTpuReplicateAttr).str();
if (cluster.empty()) {
return true;
}
if (metadata_map.find(cluster) == metadata_map.end()) {
return true;
}
num_cores_per_replica = metadata_map[cluster].getNumCoresPerReplica();
std::optional<StringRef> shard_string;
if (op->hasAttr(kXLAShardingAttr)) {
shard_string = op->getAttrOfType<StringAttr>(kXLAShardingAttr).strref();
} else {
shard_string = op->getAttrOfType<StringAttr>(kShardingAttr).strref();
}
llvm::SmallVector<xla::OpSharding> shardings;
GetFlattenedShardings(shardings, shard_string.value().str());
for (auto& sharding : shardings) {
if (sharding.type() != xla::OpSharding::MAXIMAL) {
continue;
}
if (sharding.tile_assignment_devices_size() != 1) {
op->emitOpError("TF/XLA TPU bridge input check: There must be ")
<< "exactly 1 device for MAXIMAL sharding."
<< " Number of devices assigned are "
<< sharding.tile_assignment_devices_size() << "\n";
return false;
} else {
int sharding_device = sharding.tile_assignment_devices(0);
if (sharding_device >= num_cores_per_replica || sharding_device < 0) {
op->emitOpError(
"TF2XLA TPU bridge input check: invalid sharding device ")
<< sharding_device
<< " for num_cores_per_replica = " << num_cores_per_replica
<< "\n The sharding is " << sharding.DebugString() << "\n";
return false;
}
}
}
return true;
}
bool HasSingleCoreTpu(Operation* op) {
if (auto compilation_attr =
op->getAttrOfType<mlir::StringAttr>(kCompileDeviceTypeAttr)) {
if (compilation_attr.getValue().str() == mlir::TF::kTpuDevice) {
op->emitOpError(
"TF2XLA TPU bridge input check: found a single-core TPU graph");
return true;
}
}
return false;
}
void TPUValidateInputsPass::runOnOperation() {
ModuleOp module = getOperation();
bool success = true;
int num_metadata = 0;
TPUReplicateMetadataOp metadata;
MetadataMap metadata_map;
module.walk([&](TPUReplicateMetadataOp meta) {
++num_metadata;
metadata = meta;
metadata_map[meta->getAttrOfType<StringAttr>(kTpuReplicateAttr).str()] =
meta;
});
getOperation().walk([&](mlir::Operation* op) {
if (IsPotentialUnsupportedOp(op)) {
LOG(WARNING) << "Potential unsupported op: "
<< op->getName().getStringRef().str()
<< ". TF2XLA MLIR bridge does not guarantee to support it.";
}
if (IsTpuRegularOp(op)) {
success &= CheckOpsClusterIO(op, metadata_map);
}
if (IsIntersectionXlaNonXlaOps(op)) {
success &= ValidateIntersectionXlaNonXlaOps(op, metadata_map);
}
if (IsTpuRegularOp(op)) {
success &= IsValidMAXIMALSharding(op, metadata_map);
success &= IsValidShardingTupleForArity(op);
}
success &= !HasSingleCoreTpu(op);
if (!success) {
signalPassFailure();
}
});
module.walk([&](GraphOp graph) {
if (HasV1ControlFlow(graph)) {
LOG(WARNING) << "TF2XLA MLIR bridge does not support v1 control flow."
<< " Use at your own risk.";
}
if (!success) {
signalPassFailure();
}
});
}
} // anonymous namespace
std::unique_ptr<OperationPass<ModuleOp>> CreateTPUValidateInputsPass() {
return std::make_unique<TPUValidateInputsPass>();
}
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow

Some files were not shown because too many files have changed in this diff Show More