wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 3 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/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