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
tensorflow/compiler/mlir/tf2xla/api/v2/BUILD
+412
View File
@@ -0,0 +1,412 @@
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow/core/platform:build_config.bzl", "tf_proto_library")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
# Please reach out to tf-bridge-team@ before using the TF2XLA bridge.
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
":__subpackages__",
],
)
cc_library(
name = "legalize_tf",
srcs = ["legalize_tf.cc"],
hdrs = ["legalize_tf.h"],
visibility = [
"//learning/brain/google/xla:__pkg__",
"//learning/brain/mlir/bridge:__pkg__",
"//tensorflow/compiler/mlir/quantization/stablehlo:__pkg__",
"//tensorflow/compiler/mlir/tf2xla/api/v2/testing:__pkg__",
"//tensorflow/compiler/mlir/tfrt/transforms/ifrt:__pkg__",
],
deps = [
":device_type_proto_cc",
"//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
"//tensorflow/compiler/mlir/tensorflow:serialize_mlir_module_utils",
"//tensorflow/compiler/mlir/tf2xla/api/v1:compile_tf_graph",
"//tensorflow/compiler/mlir/tf2xla/internal:compilation_timer",
"//tensorflow/compiler/mlir/tf2xla/internal:legalize_tf_to_hlo",
"//tensorflow/compiler/mlir/tf2xla/internal:reproducer_proto_cc",
"//tensorflow/compiler/tf2xla:layout_util",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core/tpu/kernels:tpu_compile_op_support",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:string_view",
"@com_google_absl//absl/types:variant",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@tsl//tsl/platform:protobuf",
"@xla//xla:shape_util",
"@xla//xla:xla_proto_cc",
"@xla//xla/client:compile_only_client",
"@xla//xla/hlo/ir:hlo",
"@xla//xla/pjrt/proto:compile_options_proto_cc",
"@xla//xla/tsl/platform:statusor",
],
)
tf_cc_test(
name = "legalize_tf_test",
srcs = ["legalize_tf_test.cc"],
deps = [
":legalize_tf",
"//tensorflow/compiler/jit",
"//tensorflow/compiler/mlir/tf2xla/api/v2/testing:compile_mlir",
"//tensorflow/compiler/mlir/tf2xla/internal:test_matchers",
"//tensorflow/compiler/tf2xla:xla_compiler",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/lib/monitoring:cell_reader",
"//tensorflow/core/protobuf:for_core_protos_cc",
"//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
"//tensorflow/core/tpu/kernels:tpu_compile_op_support",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:str_format",
"@com_google_googletest//:gtest",
"@llvm-project//mlir:Pass",
"@xla//xla:shape_util",
"@xla//xla/client:client_library",
"@xla//xla/stream_executor:platform",
"@xla//xla/stream_executor:platform_manager",
"@xla//xla/tsl/lib/monitoring:test_utils",
"@xla//xla/tsl/platform:statusor",
],
)
tf_cc_test(
name = "legalize_tf_test_gpu",
srcs = ["legalize_tf_test_gpu.cc"],
tags = [
"config-cuda-only",
"no_oss", # This test only runs with GPU.
"requires-gpu-nvidia",
],
deps = [
"//tensorflow/compiler/jit",
"//tensorflow/compiler/mlir/tf2xla/api/v2/testing:compile_mlir",
"//tensorflow/compiler/tf2xla:xla_compiler",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/protobuf:for_core_protos_cc",
"//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_matchers",
"@com_google_absl//absl/status:statusor",
"@com_google_googletest//:gtest",
"@xla//xla/tsl/platform:statusor",
],
)
tf_proto_library(
name = "device_type_proto",
srcs = ["device_type.proto"],
visibility = [
"//learning/serving/contrib/tfrt/mlir/saved_model_analysis:__pkg__",
],
)
cc_library(
name = "cluster_tf",
srcs = ["cluster_tf.cc"],
hdrs = ["cluster_tf.h"],
visibility = [
"//learning/serving/contrib/tfrt/mlir/saved_model_analysis:__pkg__",
"//tensorflow/compiler/mlir/tfrt:__pkg__",
"//tensorflow/compiler/mlir/tfrt/transforms/ifrt:__pkg__",
"//tensorflow/compiler/tf2xla:__pkg__",
],
deps = [
":device_type_proto_cc",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
"//tensorflow/compiler/mlir/tensorflow:error_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow/transforms:verify_no_outside_compilation_markers_pass",
"//tensorflow/compiler/mlir/tf2xla/internal:clustering_bridge_passes",
"//tensorflow/compiler/mlir/tf2xla/internal:logging_hooks",
"//tensorflow/compiler/mlir/tf2xla/internal/passes:clustering_passes",
"//tensorflow/core:framework",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core/platform:error_payloads",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:stacktrace",
"//tensorflow/core/platform:status",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@tsl//tsl/platform:error_logging",
"@xla//xla/tsl/platform:errors",
],
)
tf_cc_test(
name = "cluster_tf_test",
srcs = ["cluster_tf_test.cc"],
data = [
"testdata/empty_func.mlir",
"testdata/invalid_executor.mlir",
"testdata/outside_compilation.mlir",
],
deps = [
":cluster_tf",
"//tensorflow/compiler/mlir:register_common_dialects",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:tf_dialect_lib",
"//tensorflow/compiler/mlir/tf2xla/api/v2/testing:utils",
"//tensorflow/core/lib/monitoring:cell_reader",
"//tensorflow/core/platform:resource_loader",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"@xla//xla/tsl/lib/core:status_test_util",
"@xla//xla/tsl/platform:status",
],
)
cc_library(
name = "tf_dialect_to_executor",
srcs = ["tf_dialect_to_executor.cc"],
hdrs = ["tf_dialect_to_executor.h"],
visibility = [
"//learning/serving/contrib/tfrt/mlir/saved_model_analysis:__pkg__",
"//tensorflow/compiler/mlir/tfrt:__pkg__",
"//tensorflow/compiler/tf2xla:__pkg__",
],
deps = [
"//tensorflow/compiler/jit:flags_headers",
"//tensorflow/compiler/mlir/tensorflow:bridge_logger",
"//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
"//tensorflow/compiler/mlir/tensorflow/transforms:verify_no_outside_compilation_markers_pass",
"//tensorflow/compiler/mlir/tf2xla/internal:logging_hooks",
"//tensorflow/core:framework",
"//tensorflow/core/platform:error_payloads",
"//tensorflow/core/platform:status",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:Transforms",
"@tsl//tsl/platform:error_logging",
"@xla//xla/tsl/lib/monitoring:counter",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:status",
],
)
tf_cc_test(
name = "tf_dialect_to_executor_test",
srcs = ["tf_dialect_to_executor_test.cc"],
data = [
"testdata/empty_func.mlir",
"testdata/func_with_dead_ops.mlir",
"testdata/invalid_executor.mlir",
],
deps = [
":tf_dialect_to_executor",
"//tensorflow/compiler/mlir:register_common_dialects",
"//tensorflow/compiler/mlir/tf2xla/api/v2/testing:utils",
"//tensorflow/core/lib/monitoring:cell_reader",
"//tensorflow/core/platform:resource_loader",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest_main",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"@xla//xla/tsl/lib/core:status_test_util",
],
)
cc_library(
name = "tf_executor_to_graph",
srcs = [
"tf_executor_to_graph.cc",
],
hdrs = [
"tf_executor_to_graph.h",
],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/compiler/mlir:op_or_arg_name_mapper",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:convert_type",
"//tensorflow/compiler/mlir/tensorflow:error_util",
"//tensorflow/compiler/mlir/tensorflow:export_utils",
"//tensorflow/compiler/mlir/tensorflow:translate_utils",
"//tensorflow/compiler/mlir/tensorflow:verify_suitable_for_graph_export",
"//tensorflow/compiler/mlir/tensorflow/translate:mlir_roundtrip_flags",
"//tensorflow/compiler/mlir/utils:name_utils",
"//tensorflow/core:core_cpu",
"//tensorflow/core:framework",
"//tensorflow/core:framework_internal",
"//tensorflow/core:graph",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/graph/regularization:util",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/container:inlined_vector",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:optional",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@xla//xla:status_macros",
],
)
tf_cc_test(
name = "tf_executor_to_graph_test",
srcs = ["tf_executor_to_graph_test.cc"],
data = [
"testdata/valid_executor.mlir",
"testdata/valid_graph.txt",
],
deps = [
":tf_executor_to_graph",
"//tensorflow/compiler/jit",
"//tensorflow/compiler/mlir:register_common_dialects",
"//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_flags",
"//tensorflow/compiler/tf2xla:xla_compiler",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test_main",
"//tensorflow/core/platform:resource_loader",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/status:statusor",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"@riegeli//riegeli/bytes:fd_reader",
"@riegeli//riegeli/bytes:read_all",
"@tsl//tsl/platform:protobuf",
"@xla//xla/tsl/lib/core:status_test_util",
],
)
cc_library(
name = "graph_to_tf_executor",
srcs = [
"graph_to_tf_executor.cc",
],
hdrs = [
"graph_to_tf_executor.h",
],
visibility = ["//visibility:public"],
deps = [
":mlir_roundtrip_flags",
"//tensorflow/compiler/jit:shape_inference_helpers",
"//tensorflow/compiler/mlir:op_or_arg_name_mapper",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:convert_attr",
"//tensorflow/compiler/mlir/tensorflow:convert_tensor",
"//tensorflow/compiler/mlir/tensorflow:convert_type",
"//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
"//tensorflow/compiler/mlir/tensorflow:dynamic_shape_utils",
"//tensorflow/compiler/mlir/tensorflow:error_util",
"//tensorflow/compiler/mlir/tensorflow:mangling_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_attributes",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow:translate_utils",
"//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
"//tensorflow/compiler/mlir/tensorflow/translate:mlir_roundtrip_flags",
"//tensorflow/compiler/mlir/tf2xla/internal:graph_to_tf_executor_util",
"//tensorflow/compiler/mlir/tf2xla/internal:node_order",
"//tensorflow/compiler/tf2xla:functionalize_control_flow",
"//tensorflow/compiler/tf2xla:functionalize_control_flow_util",
"//tensorflow/compiler/tf2xla:tf2xla_defs",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/common_runtime:function_body",
"//tensorflow/core/platform:types",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/container:inlined_vector",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:DerivedAttributeOpInterface",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@xla//xla:status_macros",
],
)
tf_cc_test(
name = "graph_to_tf_executor_test",
srcs = ["graph_to_tf_executor_test.cc"],
data = [
"testdata/graph_with_flib_def.txt",
"testdata/valid_graph.txt",
],
deps = [
":graph_to_tf_executor",
"//tensorflow/compiler/jit",
"//tensorflow/compiler/mlir:register_common_dialects",
"//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_flags",
"//tensorflow/compiler/mlir/tf2xla/internal:graph_to_tf_executor_util",
"//tensorflow/compiler/tf2xla:xla_compiler",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test_main",
"//tensorflow/core/platform:resource_loader",
"@com_google_absl//absl/status:statusor",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@riegeli//riegeli/bytes:fd_reader",
"@riegeli//riegeli/bytes:read_all",
"@tsl//tsl/platform:protobuf",
"@xla//xla/tsl/lib/core:status_test_util",
],
)
cc_library(
name = "mlir_roundtrip_flags",
hdrs = ["mlir_roundtrip_flags.h"],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
],
)
tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf.cc
+214
View File
@@ -0,0 +1,214 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf.h"
#include <string>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "llvm/ADT/STLFunctionalExtras.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
#include "tensorflow/compiler/mlir/tf2xla/api/v2/device_type.pb.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/clustering_bridge_passes.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/logging_hooks.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h"
#include "xla/tsl/platform/errors.h"
#include "tensorflow/core/framework/metrics.h"
#include "tensorflow/core/platform/error_payloads.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/stacktrace.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/util/debug_data_dumper.h"
#include "tsl/platform/error_logging.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
using mlir::LogicalResult;
using mlir::ModuleOp;
using mlir::OpPassManager;
using mlir::PassManager;
using mlir::func::FuncOp;
// Run the TF XLA Bridge based on the input pipeline, which can be either TPU
// bridge pipeline or non TPU bridge pipeline.
absl::Status RunTFXLABridge(
ModuleOp module,
llvm::function_ref<void(OpPassManager &pm)> pipeline_builder,
llvm::StringRef module_name = llvm::StringRef(),
llvm::StringRef dump_prefix = "tf_xla_bridge_v2") {
// Explicitly check that the TensorFlow dialect can constant fold ops.
// Constant folding is essential for the bridge. Without this check, the
// bridge may fail with an error that is difficult to understand and not
// actionable.
if (!mlir::TF::TensorFlowDialect::HasConstantFoldHook()) {
return absl::InternalError(
"TensorFlow dialect missing constant fold hook in TFXLA bridge phase "
"1; this could happen if the binary doesn't link the constant fold "
"hook registration library.");
}
PassManager bridge(module.getContext());
bridge.enableVerifier();
::tensorflow::applyTensorflowAndCLOptions(bridge);
// Populate a passmanager with the list of passes that implement the bridge.
pipeline_builder(bridge);
mlir::StatusScopedDiagnosticHandler diag_handler(
module.getContext(), /*propagate=*/false,
/*filter_stack=*/!VLOG_IS_ON(1));
if (VLOG_IS_ON(1) ||
DEBUG_DATA_DUMPER()->ShouldDump(module_name.str(), kDebugGroupMain)) {
::tensorflow::DumpMlirOpToFile(
DEBUG_DATA_DUMPER()->GetDumpFilename(module_name.str(), kDebugGroupMain,
dump_prefix.str() + "_before"),
module, llvm::StringRef(), &bridge);
}
if (VLOG_IS_ON(2) ||
DEBUG_DATA_DUMPER()->ShouldDump(module_name.str(),
kDebugGroupBridgePhase1Clustering)) {
::tensorflow::tf2xla::internal::EnablePassIRPrinting(
bridge, kDebugGroupBridgePhase1Clustering, module_name);
}
LogicalResult result = bridge.run(module);
(void)result;
if (VLOG_IS_ON(1) ||
DEBUG_DATA_DUMPER()->ShouldDump(module_name.str(), kDebugGroupMain)) {
::tensorflow::DumpMlirOpToFile(
DEBUG_DATA_DUMPER()->GetDumpFilename(module_name.str(), kDebugGroupMain,
dump_prefix.str() + "_after"),
module, llvm::StringRef(), &bridge);
}
return diag_handler.ConsumeStatus();
}
absl::Status RecordIfErrorStatus(const std::string error_prefix,
bool fallback_enabled, std::string bridge_type,
std::string device_type, absl::Status status) {
if (status.ok()) {
return status;
}
VLOG(2) << error_prefix << " " << status;
tensorflow::metrics::UpdateTfMlirBridgeFirstPhaseCounter(
/*bridge_type*/ bridge_type, /*bridge_version=*/"v2", device_type,
/*fallback_enabled=*/fallback_enabled,
/*result=*/"failure");
tsl::OkOrSetErrorCounterPayload(
tensorflow::core::platform::ErrorSourceProto::MLIR_BRIDGE_PHASE_1,
status);
std::string bridge_subcomponent = "TFXLA_PHASE_ONE_MLIR_TPU_BRIDGE";
if (device_type != "tpu") {
bridge_subcomponent = "TFXLA_PHASE_ONE_MLIR_CPU/GPU_BRIDGE";
}
tsl::error_logging::Log(mlir::TF::kBridgeComponent, bridge_subcomponent,
status.ToString())
.IgnoreError();
return status;
}
void CreateReplicatedClusteringPipeline(OpPassManager &pm,
llvm::StringRef module_name) {
// Since the internal bridge clustering passes are shared among TF1/TF2
// TF2-only passes should go here. However, this should be very rare and
// new passes generally should go into the internal
// AddReplicatedBridgeClusteringPipelinePasses.
pm.addPass(tensorflow::tf2xla::internal::CreateTPUValidateInputsPass());
pm.addNestedPass<FuncOp>(
mlir::TF::CreateCanonicalizeCompileAndReplicateAttributesPass());
tensorflow::tf2xla::internal::AddReplicatedBridgeClusteringPipelinePasses(
pm, module_name);
}
void CreateReplicatedClusteringPipelineV2(OpPassManager &pm) {
CreateReplicatedClusteringPipeline(pm, /*module_name=*/"");
}
absl::Status RunFunctionTf2xlaClusteringBridge(
ModuleOp module, bool is_supported_by_replicated_brige,
bool is_in_fallback_enabled_mode, llvm::StringRef module_name) {
std::string device_type = is_supported_by_replicated_brige
? mlir::TF::kMlirPh1BridgeCounterTpu
: mlir::TF::kMlirPh1BridgeCounterNonTpu;
VLOG(2)
<< (is_supported_by_replicated_brige ? "Replicated" : "NonReplicated")
<< " Bridge called stack trace is "
<< "(NOTE: this is not an error; rather the stack trace for debugging) : "
<< tensorflow::CurrentStackTrace();
absl::Status clustering_status =
is_supported_by_replicated_brige
? RunTFXLABridge(
module,
[module_name](OpPassManager &pm) {
CreateReplicatedClusteringPipeline(pm, module_name);
},
module_name, /*dump_prefix=*/"tf_xla_bridge_v2_replicated")
: RunTFXLABridge(
module,
[](OpPassManager &pm) {
tensorflow::tf2xla::internal::
AddNonReplicatedBridgeClusteringPipelinePasses(pm);
},
module_name, /*dump_prefix=*/"tf_xla_bridge_v2_nonreplicated");
std::string bridge_type = is_supported_by_replicated_brige
? mlir::TF::kMlirPh1BridgeCounterReplicated
: mlir::TF::kMlirPh1BridgeCounterNonReplicated;
// TODO(b/317798386): add is_supported_by_replicated_brige as a filter.
TF_RETURN_IF_ERROR(RecordIfErrorStatus(
/*error_prefix=*/"clustering_v2", is_in_fallback_enabled_mode,
bridge_type, device_type, clustering_status));
// TODO(b/317798386): add is_supported_by_replicated_brige as a filter.
tensorflow::metrics::UpdateTfMlirBridgeFirstPhaseCounter(
bridge_type, /*bridge_version=*/"v2", device_type,
/*fallback_enabled=*/is_in_fallback_enabled_mode,
/*result=*/"success");
return absl::OkStatus();
}
mlir::PassPipelineRegistration<> replicated_clustering_bridge_v2(
"tf-replicated-clustering-bridge-v2",
"Run all the passes involved in transforming a TensorFlow 2 graph before "
"execution so that it is suitable for targeting devices.",
CreateReplicatedClusteringPipelineV2);
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf.h
+61
View File
@@ -0,0 +1,61 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_CLUSTER_TF_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_CLUSTER_TF_H_
#include "llvm/ADT/StringRef.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tf2xla/api/v2/device_type.pb.h"
#include "tensorflow/core/lib/core/status.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
// Run all the passes involved in transforming the graph before execution so
// that it is suitable for targeting devices when called with the TF 2 Function
// API. Users that need clustering with the Session API should use the v1 Bridge
// API. These transformations take as input a Tensorflow Graph as an MLIR Module
// and transforms the module in place to cluster the given ops for compilation
// that is compatible with the given device_type. The MLIR should be in the TF
// Executor Dialect for graph nodes and edges or be in TF Functional already.
// Individual Op inside a node should be the Tensorflow Functional Dialect. The
// output MLIR is in the TF Functional Dialect. Returns OkStatus if passed,
// otherwise an error.
//
// Inputs:
// module - The MLIR Module that will be clustered. Expected to be in TF
// Executor Dialect or TF Functional Dialect. Will convert to TF Functional.
// is_supported_by_replicated_brige - If the graph targets the replicated
// bridge. Set it to true for replicated/partitioned graphs. e.g. replicated
// and single-core TPU graphs. Set this to false if the graph is not
// replicated, e.g. CPU/GPU graphs. is_in_fallback_enabled_mode - Whether this
// was called with fallback to the non-MLIR Bridge. This is just for logging
// purposes and doesn't affect logic. module_name - What the input module name
// is for debugging help.
//
// Output: Modifies the input module in place with clustered operations.
// status - Whether the transformation to cluster the input MLIR module was
// successful.
absl::Status RunFunctionTf2xlaClusteringBridge(
mlir::ModuleOp module, bool is_supported_by_replicated_brige,
bool is_in_fallback_enabled_mode,
llvm::StringRef module_name = llvm::StringRef());
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_CLUSTER_TF_H_
tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf_test.cc
+168
View File
@@ -0,0 +1,168 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf.h"
#include <cstdint>
#include <string>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/DialectRegistry.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Parser/Parser.h" // from @llvm-project
#include "tensorflow/compiler/mlir/register_common_dialects.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h"
#include "tensorflow/compiler/mlir/tf2xla/api/v2/testing/utils.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/platform/status.h"
#include "tensorflow/core/lib/monitoring/cell_reader.h"
#include "tensorflow/core/platform/resource_loader.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
namespace {
using ::mlir::DialectRegistry;
using ::mlir::MLIRContext;
using ::mlir::ModuleOp;
using ::mlir::OwningOpRef;
using ::mlir::WalkResult;
using ::mlir::func::FuncOp;
using ::tensorflow::monitoring::testing::CellReader;
using ::tensorflow::tf2xla::v2::testing::TestDataPath;
static constexpr char kCompilationStreamz[] =
"/tensorflow/core/tf_mlir_bridge_first_phase_v2_count";
class FunctionClusterTensorflowDialectTest : public ::testing::Test {
public:
FunctionClusterTensorflowDialectTest() {
mlir::RegisterCommonToolingDialects(registry_);
context_.appendDialectRegistry(registry_);
context_.loadAllAvailableDialects();
}
absl::Status CreateMlirModule(std::string mlir_module_filename) {
std::string mlir_module_path = TestDataPath() + mlir_module_filename;
mlir_module_ =
mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context_);
if (!mlir_module_) {
return absl::Status(
absl::StatusCode::kNotFound,
absl::StrCat("Could not find MLIR module at ", mlir_module_path));
}
return absl::OkStatus();
}
DialectRegistry registry_;
MLIRContext context_;
OwningOpRef<mlir::ModuleOp> mlir_module_;
};
TEST_F(FunctionClusterTensorflowDialectTest, ClustersTfReplicatedBridge) {
CellReader<int64_t> compilation_status(kCompilationStreamz);
TF_ASSERT_OK(CreateMlirModule("empty_func.mlir"));
TF_EXPECT_OK(RunFunctionTf2xlaClusteringBridge(
*mlir_module_, /*is_supported_by_replicated_brige*/ true,
/*is_in_fallback_enabled_mode=*/false));
FuncOp main = mlir_module_->lookupSymbol<mlir::func::FuncOp>("main");
ASSERT_TRUE(main);
EXPECT_EQ(compilation_status.Delta(mlir::TF::kMlirPh1BridgeCounterReplicated,
mlir::TF::kMlirPh1BridgeCounterV2,
mlir::TF::kMlirPh1BridgeCounterTpu,
"fallback_disabled", "success"),
1);
}
TEST_F(FunctionClusterTensorflowDialectTest,
RunsOutsideCompilationReplicatedBridge) {
CellReader<int64_t> compilation_status(kCompilationStreamz);
TF_ASSERT_OK(CreateMlirModule("outside_compilation.mlir"));
TF_EXPECT_OK(RunFunctionTf2xlaClusteringBridge(
*mlir_module_, /*is_supported_by_replicated_brige*/ true,
/*is_in_fallback_enabled_mode=*/false));
FuncOp main = mlir_module_->lookupSymbol<mlir::func::FuncOp>("main");
ASSERT_TRUE(main);
bool has_cluster_op = false;
main.walk([&](mlir::tf_device::ClusterFuncOp cluster_op) {
has_cluster_op = true;
return WalkResult::advance();
});
EXPECT_TRUE(has_cluster_op);
EXPECT_EQ(compilation_status.Delta(mlir::TF::kMlirPh1BridgeCounterReplicated,
mlir::TF::kMlirPh1BridgeCounterV2,
mlir::TF::kMlirPh1BridgeCounterTpu,
"fallback_disabled", "success"),
1);
}
TEST_F(FunctionClusterTensorflowDialectTest, ClustersTFNonReplicatedBridge) {
CellReader<int64_t> compilation_status(kCompilationStreamz);
TF_ASSERT_OK(CreateMlirModule("empty_func.mlir"));
TF_EXPECT_OK(RunFunctionTf2xlaClusteringBridge(
*mlir_module_, /*is_supported_by_replicated_brige*/ false,
/*is_in_fallback_enabled_mode=*/false));
FuncOp main = mlir_module_->lookupSymbol<mlir::func::FuncOp>("main");
ASSERT_TRUE(main);
EXPECT_EQ(
compilation_status.Delta(mlir::TF::kMlirPh1BridgeCounterNonReplicated,
mlir::TF::kMlirPh1BridgeCounterV2,
mlir::TF::kMlirPh1BridgeCounterNonTpu,
"fallback_disabled", "success"),
1);
}
TEST_F(FunctionClusterTensorflowDialectTest, LogsFallbackMode) {
CellReader<int64_t> compilation_status(kCompilationStreamz);
TF_ASSERT_OK(CreateMlirModule("empty_func.mlir"));
TF_EXPECT_OK(RunFunctionTf2xlaClusteringBridge(
*mlir_module_, /*is_supported_by_replicated_brige*/ true,
/*is_in_fallback_enabled_mode=*/true));
EXPECT_EQ(compilation_status.Delta(mlir::TF::kMlirPh1BridgeCounterReplicated,
mlir::TF::kMlirPh1BridgeCounterV2,
mlir::TF::kMlirPh1BridgeCounterTpu,
"fallback_enabled", "success"),
1);
}
} // namespace
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v2/device_type.proto
+26
View File
@@ -0,0 +1,26 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto2";
package tensorflow.tf2xla.v2;
// The requested device type to compile for.
enum DeviceType {
DEVICE_TYPE_UNSPECIFIED = 0;
XLA_TPU_JIT = 1;
XLA_CPU_JIT = 2;
XLA_GPU_JIT = 3;
}
tensorflow/compiler/mlir/tf2xla/api/v2/graph_to_tf_executor.cc
+2844
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/mlir/tf2xla/api/v2/graph_to_tf_executor.h
+54
View File
@@ -0,0 +1,54 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_GRAPH_TO_TF_EXECUTOR_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_GRAPH_TO_TF_EXECUTOR_H_
#include <optional>
#include <string>
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/graph_to_tf_executor_util.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/graph_debug_info.pb.h"
#include "tensorflow/core/graph/graph.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
inline constexpr absl::string_view kImportModelDefaultGraphFuncName = "main";
// Given a Graph, returns a MLIR module containing the graph, expressed with
// tf_executor dialect.
absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> ConvertGraphToTfExecutor(
const Graph& graph, const GraphDebugInfo& debug_info,
const FunctionLibraryDefinition& flib_def, const GraphImportConfig& specs,
mlir::MLIRContext* context,
std::unordered_map<std::string, std::string>* tf_name_to_mlir_name =
nullptr,
const ConfigProto& config_proto = {},
tensorflow::TF2XLABridgeVersion bridge_version =
tensorflow::TF2XLABridgeVersion::kNotBridgeUseCase,
std::optional<absl::string_view> module_name = std::nullopt);
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_GRAPH_TO_TF_EXECUTOR_H_
tensorflow/compiler/mlir/tf2xla/api/v2/graph_to_tf_executor_test.cc
+205
View File
@@ -0,0 +1,205 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/api/v2/graph_to_tf_executor.h"
#include <stdlib.h>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <gtest/gtest.h>
#include "absl/status/statusor.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/DialectRegistry.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "riegeli/bytes/fd_reader.h" // from @riegeli
#include "riegeli/bytes/read_all.h" // from @riegeli
#include "tensorflow/compiler/mlir/register_common_dialects.h"
#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/graph_to_tf_executor_util.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/platform/resource_loader.h"
#include "tsl/platform/protobuf.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
namespace {
using mlir::DialectRegistry;
using mlir::MLIRContext;
constexpr char kGraphWithFlibDefFileName[] = "graph_with_flib_def.txt";
std::string TestDataPath() {
return tensorflow::GetDataDependencyFilepath(
"tensorflow/compiler/mlir/tf2xla/api/v2/testdata/");
}
class GraphToTfExecutorTest : public ::testing::Test {
public:
GraphToTfExecutorTest() {
mlir::RegisterCommonToolingDialects(registry_);
context_.appendDialectRegistry(registry_);
context_.loadAllAvailableDialects();
}
GraphDef CreateGraphDef(std::string graphdef_filename) {
std::string file_path = TestDataPath() + graphdef_filename;
std::string contents;
GraphDef graph_def;
auto status = riegeli::ReadAll(riegeli::FdReader(file_path), contents);
if (!status.ok()) {
return graph_def;
}
tsl::protobuf::TextFormat::ParseFromString(contents, &graph_def);
return graph_def;
}
int CountNumberOfFunctionsInModule(mlir::ModuleOp module) {
int count = 0;
for (auto unused : module.getOps<mlir::func::FuncOp>()) {
(void)unused; // Avoid unused variable warning
count++;
}
return count;
}
DialectRegistry registry_;
MLIRContext context_;
};
TEST_F(GraphToTfExecutorTest, BasicConvertGraphToTfExecutorPasses) {
Graph graph(OpRegistry::Global());
GraphDebugInfo debug_info;
FunctionLibraryDefinition flib_def(OpRegistry::Global(),
FunctionDefLibrary());
GraphImportConfig specs;
GraphDef graph_def = CreateGraphDef("valid_graph.txt");
GraphConstructorOptions opts;
TF_ASSERT_OK(ConvertGraphDefToGraph(opts, graph_def, &graph));
TF_ASSERT_OK(
ConvertGraphToTfExecutor(graph, debug_info, flib_def, specs, &context_));
}
TEST_F(
GraphToTfExecutorTest,
ConvertGraphToTfExecutorConvertAllFunctionsTrueConvertsAllFunctionsInFlibDef) {
Graph graph(OpRegistry::Global());
GraphDebugInfo debug_info;
FunctionLibraryDefinition flib_def(OpRegistry::Global(),
FunctionDefLibrary());
GraphDef graph_def = CreateGraphDef(kGraphWithFlibDefFileName);
GraphConstructorOptions opts;
TF_ASSERT_OK(ConvertGraphDefToGraph(opts, graph_def, &graph));
GraphImportConfig specs;
specs.convert_all_functions_to_mlir = true;
absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> result =
ConvertGraphToTfExecutor(graph, debug_info, graph.flib_def(), specs,
&context_);
// should equal main + 4 functions in flib_def
ASSERT_EQ(CountNumberOfFunctionsInModule(result->get()), 5);
}
TEST_F(
GraphToTfExecutorTest,
ConvertGraphToTfExecutorConvertAllFunctionsFalseOnlyConvertsFunctionsReferencedInGraph) {
Graph graph(OpRegistry::Global());
GraphDebugInfo debug_info;
FunctionLibraryDefinition flib_def(OpRegistry::Global(),
FunctionDefLibrary());
GraphDef graph_def = CreateGraphDef(kGraphWithFlibDefFileName);
GraphConstructorOptions opts;
TF_ASSERT_OK(ConvertGraphDefToGraph(opts, graph_def, &graph));
GraphImportConfig specs;
specs.convert_all_functions_to_mlir = false;
absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> result =
ConvertGraphToTfExecutor(graph, debug_info, graph.flib_def(), specs,
&context_);
// should equal main + 2 functions referenced by nodes in the graph via the
// "f" attr.
ASSERT_EQ(CountNumberOfFunctionsInModule(result->get()), 3);
}
TEST_F(GraphToTfExecutorTest,
ConvertGraphToTfExecutorPopulatesTfNameToMlirNameMap) {
Graph graph(OpRegistry::Global());
GraphDebugInfo debug_info;
FunctionLibraryDefinition flib_def(OpRegistry::Global(),
FunctionDefLibrary());
GraphDef graph_def = CreateGraphDef(kGraphWithFlibDefFileName);
GraphConstructorOptions opts;
TF_ASSERT_OK(ConvertGraphDefToGraph(opts, graph_def, &graph));
GraphImportConfig specs;
std::unordered_map<std::string, std::string> tf_name_to_mlir_name;
TF_ASSERT_OK(ConvertGraphToTfExecutor(graph, debug_info, graph.flib_def(),
specs, &context_,
&tf_name_to_mlir_name));
std::unordered_set<std::string> result_set;
for (const auto& pair : tf_name_to_mlir_name) {
result_set.insert(pair.first);
}
// Converted functions referenced by nodes in the graph via the
// "f" attr. These are before they are converted to their corresponding MLIR
// name.
std::unordered_set<std::string> expected_set = {
"__inference__traced_save_45", "__inference__traced_restore_57"};
EXPECT_EQ(result_set, expected_set);
}
TEST_F(GraphToTfExecutorTest, ConvertGraphToTfExecutorSetsModuleName) {
Graph graph(OpRegistry::Global());
GraphDebugInfo debug_info;
FunctionLibraryDefinition flib_def(OpRegistry::Global(),
FunctionDefLibrary());
GraphImportConfig specs;
GraphDef graph_def = CreateGraphDef("valid_graph.txt");
GraphConstructorOptions opts;
TF_ASSERT_OK(ConvertGraphDefToGraph(opts, graph_def, &graph));
const std::string kModuleName = "my_module_name";
absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> result =
ConvertGraphToTfExecutor(
graph, debug_info, flib_def, specs, &context_,
/*tf_name_to_mlir_name=*/nullptr,
/*config_proto=*/{},
/*bridge_version=*/
tensorflow::TF2XLABridgeVersion::kNotBridgeUseCase,
/*module_name=*/kModuleName);
TF_ASSERT_OK(result.status());
EXPECT_EQ(result->get().getName().value_or(""), kModuleName);
}
} // namespace
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.cc
+206
View File
@@ -0,0 +1,206 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.h"
#include <memory>
#include <string>
#include <variant>
#include <vector>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "absl/types/variant.h"
#include "llvm/ADT/ScopeExit.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/serialize_mlir_module_utils.h"
#include "tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/compilation_timer.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/reproducer.pb.h"
#include "tensorflow/compiler/tf2xla/layout_util.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/client/compile_only_client.h"
#include "xla/hlo/ir/hlo_module.h"
#include "xla/shape.h"
#include "xla/tsl/framework/device_type.h"
#include "xla/tsl/lib/monitoring/sampler.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/statusor.h"
#include "xla/xla.pb.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
#include "tensorflow/core/util/debug_data_dumper.h"
#include "tsl/platform/protobuf.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
using tpu::FunctionToHloArgs;
using tpu::MlirToHloArgs;
using tpu::ShardingAndIndex;
auto* phase2_bridge_compilation_time = tsl::monitoring::Sampler<1>::New(
{"/tensorflow/core/tf2xla/api/v2/phase2_compilation_time",
"The wall-clock time spent on executing graphs in milliseconds.",
"configuration"},
// Power of 1.5 with bucket count 45 (> 23 hours)
{tsl::monitoring::Buckets::Exponential(1, 1.5, 45)});
// Name of component for error logging. This name is fixed and required to
// enable logging.
constexpr char kBridgeComponent[] = "TFXLABridge";
constexpr char kFullBridge[] = "full_bridge";
namespace {
bool ShouldFallbackToGraphCompiler(
const std::variant<MlirToHloArgs, FunctionToHloArgs>& computation) {
if (computation.index() == 1) return true;
return std::get<0>(computation).rollout_state ==
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_DISABLED;
}
std::string ComputationToString(const MlirToHloArgs computation) {
if (computation.mlir_module_op.has_value()) {
return SerializeMlirModule(computation.mlir_module_op.value());
}
return std::string(computation.mlir_module);
}
void DumpComputationInput(
const tpu::TPUCompileMetadataProto& metadata,
const std::vector<tensorflow::TensorShape>& arg_shapes,
const std::variant<tpu::MlirToHloArgs, tpu::FunctionToHloArgs>
computation) {
if (!VLOG_IS_ON(2)) {
return;
}
tensorflow::mlir::tf2xla::internal::LegalizeMlirToHloReproducer reproducer;
*reproducer.mutable_compile_metadata() = metadata;
for (const auto& shape : arg_shapes) {
shape.AsProto(reproducer.add_input_shapes());
}
switch (computation.index()) {
case 0:
reproducer.set_mlir_module(ComputationToString(std::get<0>(computation)));
break;
case 1: {
auto input = std::get<1>(computation);
*reproducer.mutable_function_def_library() = input.flib_def->ToProto();
} break;
default:
VLOG(2) << "LegalizeMlirToHlo computation input: unknown";
break;
}
std::string string_reproducer;
tensorflow::protobuf::TextFormat::PrintToString(reproducer,
&string_reproducer);
DumpRawStringToFile("legalize_tf_reproducer.textproto", string_reproducer);
}
absl::Status DumpHloCompilationResult(
absl::string_view name, XlaCompilationResult* compilation_result) {
if (!VLOG_IS_ON(2) &&
!DEBUG_DATA_DUMPER()->ShouldDump(std::string(name), kDebugGroupMain)) {
return absl::OkStatus();
}
TF_ASSIGN_OR_RETURN(
auto hlo_module_config,
xla::HloModule::CreateModuleConfigFromProto(
compilation_result->computation->proto(), xla::DebugOptions()));
TF_ASSIGN_OR_RETURN(
std::unique_ptr<xla::HloModule> hlo_module,
xla::HloModule::CreateFromProto(compilation_result->computation->proto(),
hlo_module_config));
std::string all_computations;
for (auto computation : hlo_module->computations()) {
all_computations += computation->ToString() + "\n\n";
}
tensorflow::DumpRawStringToFile(name, all_computations);
return absl::OkStatus();
}
} // namespace
absl::StatusOr<tensorflow::XlaCompilationResult> LegalizeMlirToHlo(
const std::variant<tpu::MlirToHloArgs, tpu::FunctionToHloArgs>& computation,
const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,
llvm::StringRef device_type,
std::vector<std::unique_ptr<::mlir::Pass>>& custom_legalization_passes,
XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,
const std::vector<tensorflow::TensorShape>& arg_shapes,
std::vector<tpu::ShardingAndIndex>* arg_core_mapping,
std::vector<std::vector<xla::Shape>>* per_core_arg_shapes,
xla::CompileOnlyClient* client) {
CompilationTimer timer;
auto record_time = llvm::make_scope_exit([&timer] {
phase2_bridge_compilation_time->GetCell(kFullBridge)
->Add(timer.ElapsedCyclesInMilliseconds());
});
auto compilation_result = std::make_unique<XlaCompilationResult>();
DumpComputationInput(metadata, arg_shapes, computation);
// If there are no MLIR args, compile the given function in the library.
if (ShouldFallbackToGraphCompiler(computation)) {
TF_RETURN_IF_ERROR(tf2xla::v1::CompileTensorflowGraphToHlo(
computation, metadata, use_tuple_args, shape_determination_fns,
arg_shapes, tsl::DeviceType(device_type.str()), arg_core_mapping,
per_core_arg_shapes, client, compilation_result.get()));
DumpHloCompilationResult("legalize_tf_fallback.hlo",
compilation_result.get())
.IgnoreError();
return *compilation_result;
}
auto combined_bridge_status = internal::LegalizeTfToHlo(
std::get<0>(computation), metadata, use_tuple_args, device_type,
shape_determination_fns, arg_shapes, arg_core_mapping,
per_core_arg_shapes, custom_legalization_passes, client,
compilation_result.get());
if (combined_bridge_status.ok()) {
VLOG(1) << "Successfully compiled MLIR computation to XLA HLO using "
"Combined MLIR and XlaBuilder Bridge.";
DumpHloCompilationResult("legalize_tf_combined_bridge.hlo",
compilation_result.get())
.IgnoreError();
return *compilation_result;
}
return combined_bridge_status.status();
}
}; // namespace v2
}; // namespace tf2xla
}; // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.h
+67
View File
@@ -0,0 +1,67 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_LEGALIZE_TF_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_LEGALIZE_TF_H_
#include <memory>
#include <variant>
#include <vector>
#include "absl/status/status.h"
#include "absl/types/variant.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tf2xla/api/v2/device_type.pb.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/client/compile_only_client.h"
#include "xla/pjrt/proto/compile_options.pb.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
// Legalizes the given mlir::Module into XLA HLO. If successful, returns the
// compiled XLA HLO. V1 of the tf2xla uses MLIR whereas V0 does not use MLIR.
//
// Inputs:
// computation - The MLIR module op. It currently takes in
// tpu::FunctionToHloArgs but this is deprecated. arg_shapes - The shapes of
// the arguments in module_op. device_type - The device type to compile for.
// use_tuple_args - Pack the incoming arg shapes into a single tuple.
// custom_legalization_passes - Extra passes to lower from TF -> MHLO.
// arg_shapes - The shapes of the args.
// arg_core_mapping - Which args go on which cores.
// per_core_arg_shapes - For each core, the shapes for each argument.
// client - The Xla Compilation client.
absl::StatusOr<tensorflow::XlaCompilationResult> LegalizeMlirToHlo(
const std::variant<tpu::MlirToHloArgs, tpu::FunctionToHloArgs>& computation,
const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,
llvm::StringRef device_type,
std::vector<std::unique_ptr<mlir::Pass>>& custom_legalization_passes,
XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,
const std::vector<tensorflow::TensorShape>& arg_shapes,
std::vector<tpu::ShardingAndIndex>* arg_core_mapping,
std::vector<std::vector<xla::Shape>>* per_core_arg_shapes,
xla::CompileOnlyClient* client);
}; // namespace v2
}; // namespace tf2xla
}; // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_LEGALIZE_TF_H_
tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf_test.cc
+367
View File
@@ -0,0 +1,367 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.h"
#include <cstdint>
#include <memory>
#include <string>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/statusor.h"
#include "absl/strings/str_format.h"
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tf2xla/api/v2/testing/compile_mlir.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/test_matchers.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/client/client_library.h"
#include "xla/shape.h"
#include "xla/stream_executor/platform.h"
#include "xla/stream_executor/platform_manager.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/lib/monitoring/test_utils.h"
#include "xla/tsl/platform/statusor.h"
#include "xla/tsl/platform/test.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/lib/monitoring/cell_reader.h"
#include "tensorflow/core/lib/monitoring/test_utils.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
#include "tensorflow/core/util/debug_data_dumper.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
using ::tensorflow::monitoring::testing::CellReader;
using tensorflow::tf2xla::v2::testing::CompileMlirModule;
using ::testing::Not;
using ::testing::TestWithParam;
using tpu::FunctionToHloArgs;
using tpu::ShardingAndIndex;
using tpu::TPUCompileMetadataProto;
using tsl::testing::TmpDir;
static constexpr char kCompilationTimeStreamzName[] =
"/tensorflow/core/tf2xla/api/v2/phase2_compilation_time";
static constexpr char kFullBridge[] = "full_bridge";
static constexpr char kCompilationStatusStreamzName[] =
"/tensorflow/core/tf2xla/api/v2/phase2_compilation_status";
static const char kMlirWithFallbackModeSuccess[] =
"kMlirWithFallbackModeSuccess";
static const char kMlirWithFallbackModeFailure[] =
"kMlirWithFallbackModeFailure";
static const char kOldBridgeMlirFilteredFailure[] =
"kOldBridgeMlirFilteredFailure";
static const char kOldBridgeWithFallbackModeFailure[] =
"kOldBridgeWithFallbackModeFailure";
static const char kOldBridgeMlirFilteredSuccess[] =
"kOldBridgeMlirFilteredSuccess";
static const char kOldBridgeWithFallbackModeSuccess[] =
"kOldBridgeWithFallbackModeSuccess";
static const char kMlirCombinedMlirSuccess[] = "kMlirCombinedMlirSuccess";
static const char kMlirCombinedMlirFailure[] = "kMlirCombinedMlirFailure";
static const char kMlirCombinedOldSuccess[] = "kMlirCombinedOldSuccess";
static const char kMlirCombinedOldFailure[] = "kMlirCombinedOldFailure";
static constexpr char kMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> () {
func.return
}
})";
// MLIR which should not legalize at all
static constexpr char kBadMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> () {
%0 = tf.Unknown() -> ()
func.return %0
}
})";
// MLIR which should be filtered by the MLIR bridge but fully legalize with the
// combined bridge.
static constexpr char kUnsupportedMlirBridgeModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> () {
%cst0 = "tf.Const"(){ value = dense<0> : tensor<3x5xi1>} : () -> tensor<3x5xi1>
%0 = "tf.Where"(%cst0) : (tensor<3x5xi1>) -> tensor<?x2xi64>
func.return
}
})";
TEST(LegalizeTFTest, RecordsStreamzForSuccessfulLegalizeWithMlirBridge) {
CellReader<int64_t> compilation_status(kCompilationStatusStreamzName);
TF_ASSERT_OK_AND_ASSIGN(
XlaCompiler::CompilationResult result,
CompileMlirModule(
kMlirModuleStr,
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED));
// May have been filtered so check for lack of failure instead of success.
EXPECT_EQ(compilation_status.Delta(kMlirWithFallbackModeFailure), 0);
}
TEST(LegalizeTFTest, MatMul) {
static constexpr char kMatMulModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> (tensor<5x11xf32>) {
%arg0 = "tf.Const"() {value = dense<-3.0> : tensor<5x7xf32>} : () -> tensor<5x7xf32>
%arg1 = "tf.Const"() {value = dense<-3.0> : tensor<11x7xf32>} : () -> tensor<11x7xf32>
%1 = "tf.MatMul"(%arg0, %arg1) {transpose_a = false, transpose_b = true} : (tensor<5x7xf32>, tensor<11x7xf32>) -> tensor<5x11xf32>
func.return %1 : tensor<5x11xf32>
}
})";
TF_ASSERT_OK_AND_ASSIGN(
XlaCompiler::CompilationResult result,
CompileMlirModule(
kMatMulModuleStr,
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED));
}
struct MatMulTestCase {
std::string mat_mul_method;
};
using BatchMatMulTest = TestWithParam<MatMulTestCase>;
TEST_P(BatchMatMulTest, BatchMatMul) {
const MatMulTestCase& test_case = GetParam();
static constexpr char kMatMulModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> (tensor<1x4x4xf32>) {
%%arg0 = "tf.Const"() {value = dense<-3.0> : tensor<1x4x2xf32>} : () -> tensor<1x4x2xf32>
%%arg1 = "tf.Const"() {value = dense<-3.0> : tensor<1x2x4xf32>} : () -> tensor<1x2x4xf32>
%%1 = "tf.%s"(%%arg0, %%arg1) {T = f32, adj_x = false, adj_y = false, grad_x = false, grad_y = false, device = ""} : (tensor<1x4x2xf32>, tensor<1x2x4xf32>) -> tensor<1x4x4xf32>
func.return %%1 : tensor<1x4x4xf32>
}
})";
std::string mat_mul_method =
absl::StrFormat(kMatMulModuleStr, test_case.mat_mul_method);
TF_ASSERT_OK_AND_ASSIGN(
XlaCompiler::CompilationResult result,
CompileMlirModule(
mat_mul_method.c_str(),
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED));
}
INSTANTIATE_TEST_SUITE_P(
BatchMatMulTest, BatchMatMulTest,
::testing::ValuesIn<MatMulTestCase>({
{"BatchMatMul"},
{"BatchMatMulV2"},
{"BatchMatMulV3"},
}),
[](const ::testing::TestParamInfo<BatchMatMulTest::ParamType>& info) {
return info.param.mat_mul_method;
});
TEST(LegalizeTFTest, DumpsProducedHLO) {
Env* env = Env::Default();
std::string test_dir = TmpDir();
setenv("TF_DUMP_GRAPH_PREFIX", test_dir.c_str(), /*overwrite=*/1);
setenv("TF_DUMP_GRAPH_NAME_FILTER", "*", 1);
DEBUG_DATA_DUMPER()->LoadEnvvars();
std::vector<std::string> files;
TF_ASSERT_OK(env->GetChildren(test_dir, &files));
int original_files_size = files.size();
TF_ASSERT_OK_AND_ASSIGN(
XlaCompiler::CompilationResult result,
CompileMlirModule(
kMlirModuleStr,
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED));
// Due to the shared test of this infrastructure, we just need to make sure
// that the dumped file size is greater than what was originally inside
// the test directory.
TF_ASSERT_OK(env->GetChildren(test_dir, &files));
EXPECT_THAT(files.size(), ::testing::Gt(original_files_size));
setenv("TF_DUMP_GRAPH_PREFIX", test_dir.c_str(), /*overwrite=*/0);
}
TEST(LegalizeTFTest, RecordsStreamzForFailedLegalizeWithMlirBridge) {
CellReader<int64_t> compilation_status(kCompilationStatusStreamzName);
auto result = CompileMlirModule(
kBadMlirModuleStr,
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED);
EXPECT_FALSE(result.ok());
EXPECT_EQ(compilation_status.Delta(kMlirCombinedMlirFailure), 1);
}
TEST(LegalizeTFTest, RecordsStreamzForSuccessWithCombinedBridge) {
CellReader<int64_t> compilation_status(kCompilationStatusStreamzName);
auto result = CompileMlirModule(
kUnsupportedMlirBridgeModuleStr,
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED);
// MLIR Bridge will filter this unsupported MLIR, Combined will succeed.
EXPECT_TRUE(result.ok());
EXPECT_EQ(compilation_status.Delta(kMlirCombinedMlirSuccess), 1);
EXPECT_EQ(compilation_status.Delta(kMlirCombinedMlirFailure), 0);
EXPECT_EQ(compilation_status.Delta(kMlirCombinedOldSuccess), 1);
EXPECT_EQ(compilation_status.Delta(kMlirCombinedOldFailure), 0);
// Old bridge should never be called at all.
EXPECT_EQ(compilation_status.Delta(kOldBridgeMlirFilteredFailure), 0);
EXPECT_EQ(compilation_status.Delta(kOldBridgeWithFallbackModeFailure), 0);
EXPECT_EQ(compilation_status.Delta(kOldBridgeMlirFilteredSuccess), 0);
EXPECT_EQ(compilation_status.Delta(kOldBridgeWithFallbackModeSuccess), 0);
}
TEST(LegalizeTFTest, RecordsStreamzForNoMlirFallback) {
FunctionDef my_func =
tensorflow::FunctionDefHelper::Create("empty", {}, {}, {}, {}, {});
tensorflow::FunctionDefLibrary fdef;
*(fdef.add_function()) = my_func;
tensorflow::FunctionLibraryDefinition flib_def(
tensorflow::OpRegistry::Global(), fdef);
OpInputList guaranteed_constants;
NameAttrList function;
FunctionToHloArgs function_to_hlo_args{&function,
&flib_def,
/*graph_def_version=*/0,
{&guaranteed_constants}};
se::Platform* cpu_platform =
se::PlatformManager::PlatformWithName("Host").value();
auto client =
xla::ClientLibrary::GetOrCreateCompileOnlyClient(cpu_platform).value();
std::vector<TensorShape> arg_shapes;
TPUCompileMetadataProto metadata_proto;
bool use_tuple_args = true;
std::vector<ShardingAndIndex> arg_core_mapping;
std::vector<std::vector<xla::Shape>> per_core_arg_shapes;
std::vector<std::unique_ptr<mlir::Pass>> custom_legalization_passes;
// This doesn't actually compile correctly.
absl::StatusOr<XlaCompiler::CompilationResult> compile_result =
LegalizeMlirToHlo(function_to_hlo_args, metadata_proto, use_tuple_args,
/*device_type=*/"XLA_CPU_JIT",
custom_legalization_passes,
/*shape_determination_fns=*/{}, arg_shapes,
&arg_core_mapping, &per_core_arg_shapes, client);
EXPECT_FALSE(compile_result.ok());
}
TEST(LegalizeTFTest, RecordsCompilationTimeForSuccessfulCompilation) {
CellReader<monitoring::testing::Histogram> compilation_time(
kCompilationTimeStreamzName);
TF_ASSERT_OK_AND_ASSIGN(
XlaCompiler::CompilationResult result,
CompileMlirModule(
kMlirModuleStr,
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_ENABLED));
// Compilation time should have been updated.
EXPECT_GT(compilation_time.Delta(kFullBridge).num(), 0);
}
TEST(LegalizeTFTest, SuccessfullyCompilesModulesWithReturnValues) {
static constexpr char kHasReturnValuesAndNoMetadataRetvals[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> (tensor<2xi32>) {
%cst = "tf.Const"() {value = dense<[524170, 523952]> : tensor<2xi32>} : () -> tensor<2xi32>
return %cst : tensor<2xi32>
}
})";
auto compilation_result = CompileMlirModule(
kHasReturnValuesAndNoMetadataRetvals,
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED);
EXPECT_TRUE(compilation_result.ok());
// Ensure that the compilation result contains a constant.
EXPECT_THAT(compilation_result,
ComputationProtoContains("opcode:.*constant"));
}
TEST(LegalizeTFTest, SkipsTensorListSetItemIfDimensionsTooLarge) {
static constexpr char kTensorListSetItemDimensionTooLarge[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> tensor<!tf_type.variant<tensor<64x1xbf16>>> {
// unknown rank
%elem_shape = "tf.Const"() <{value = dense<-1> : tensor<i32>}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> tensor<i32>
// zero reserved elements
%num_elements = "tf.Const"() <{value = dense<0> : tensor<i32>}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> tensor<i32>
%list = "tf.TensorListReserve"(%elem_shape, %num_elements) : (tensor<i32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<64x1xbf16>>>
%index = "tf.Const"() <{value = dense<0> : tensor<i32>}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> tensor<i32>
%element = "tf.Const"() <{value = dense<0.0> : tensor<64x1xbf16>}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> tensor<64x1xbf16>
// Results in a bad mismatch of shapes.
%updated_list = "tf.TensorListSetItem"(%list, %index, %element) : (tensor<!tf_type.variant<tensor<64x1xbf16>>>, tensor<i32>, tensor<64x1xbf16>) -> tensor<!tf_type.variant<tensor<64x1xbf16>>>
return %updated_list : tensor<!tf_type.variant<tensor<64x1xbf16>>>
}
})";
auto compilation_result = CompileMlirModule(
kTensorListSetItemDimensionTooLarge,
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED);
// Ensure that it compile
ASSERT_TRUE(compilation_result.ok());
// Assert that the tensor list operation is lowered to something.
ASSERT_THAT(compilation_result,
Not(ComputationProtoContains("%.*= \"tf.TensorListSetItem")));
// Assert that the tensor list operation is lowered to something that doesn't
// get stuck on a broken dynamic update slice.
ASSERT_THAT(compilation_result,
Not(ComputationProtoContains("%.*=.*DynamicUpdateSlice")));
}
TEST(LegalizeTFTest, LegalizesFunctionWithBoundedDynamicArg) {
static constexpr char kMlirModuleWithBoundedDynamicArgStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main(%arg0: tensor<?xi32, #mhlo.type_extensions<bounds = [3]>> ) -> (tensor<?xi32, #mhlo.type_extensions<bounds = [3]>>) {
func.return %arg0 : tensor<?xi32, #mhlo.type_extensions<bounds = [3]>>
}
})";
auto compilation_result = CompileMlirModule(
kMlirModuleWithBoundedDynamicArgStr,
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED);
ASSERT_TRUE(compilation_result.ok());
EXPECT_THAT(compilation_result,
ComputationProtoContains("element_type:.S32\n.*dimensions: 3"));
}
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf_test_gpu.cc
+82
View File
@@ -0,0 +1,82 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "absl/status/status_matchers.h"
#include "absl/status/statusor.h"
#include "tensorflow/compiler/mlir/tf2xla/api/v2/testing/compile_mlir.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
// These test are in a separate file because they requires separate build
// environments (--config=cuda).
using tensorflow::tf2xla::v2::testing::CompileMlirModule;
// MLIR which is legalize only with the right device type.
// The mlir is generated by running
// tensorflow/compiler/mlir/tf-opt -tf-xla-call-module-serialization
//
// module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32,
// producer = 268 : i32}} {
// func.func private @_jit_sin(%arg0: tensor<f32>) -> tensor<f32> {
// %0 = stablehlo.sine %arg0 : tensor<f32>
// return %0 : tensor<f32>
// }
// func.func @main(%arg0: tensor<f32>) -> tensor<*xf32> {
// %0 = "tf.XlaCallModule"(%arg0) {Sout = [#tf_type.shape<*>], device =
// "", dim_args_spec = [], _entry_function = @_jit_sin,
// _stablehlo_version = "1.0.0", module = "", platforms = ["CUDA"],
// version = 6 : i64} : (tensor<f32>) -> tensor<*xf32>
// func.return %0 : tensor<*xf32>
// }
// }
//
static constexpr char kGpuMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main(%arg0: tensor<f32>) -> tensor<*xf32> {
%0 = "tf.XlaCallModule"(%arg0) {Sout = [#tf_type.shape<*>], device = "", dim_args_spec = [], module = "ML\EFR\03StableHLOv1.0.0\00\01\17\05\01\05\01\03\05\03\07\07\t\0B\03K5\07\01\1B\07\0B\13\0B3\0B\0B\0B\0B\0F\0B\13\0B\03\1B\0F\1B\0B\0B\0B\0B\0B\0F\13\0B\0B\0B\0B\03\07\0F\17\07\02\A7\1F\05\0D\03\03\03\07\05\0F\03\0B\0B\1B\0D'\0F)\031\113\05\11\05\13\05\15\05\17\1D\15\17\05\19\17\19\EF\01\05\1B\03\03\1D\0D\05\1F!#%\1D\1D\1D\1F\1D!\1D##\03\03\03+\0D\03-/\1D%\1D'\1D)\1D+)\01\05\11\03\01\03\01\t\04A\05\01\11\01\05\07\03\01\05\03\11\01\t\05\03\05\0B\03\01\01\05\06\13\03\01\03\01\07\04\01\03\03\06\03\01\05\01\00\9A\04-\0F\0B\03!\1B\1D\05\1B\83/\1F\15\1D\15\11\13\15\11\11\0F\0B\11builtin\00vhlo\00module\00func_v1\00sine_v1\00return_v1\00sym_name\00jit_sin\00arg_attrs\00function_type\00res_attrs\00sym_visibility\00jit(sin)/jit(main)/sin\00third_party/py/jax/experimental/jax2tf/tests/back_compat_test.py\00jax.arg_info\00x\00mhlo.sharding\00{replicated}\00jax.result_info\00\00main\00public\00", platforms = ["CUDA"], version = 6 : i64} : (tensor<f32>) -> tensor<*xf32>
func.return %0 : tensor<*xf32>
}
})";
TEST(LegalizeTFTest, RightDeviceTypeShallPass) {
TF_ASSERT_OK_AND_ASSIGN(
XlaCompiler::CompilationResult result,
CompileMlirModule(
kGpuMlirModuleStr,
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED,
"XLA_GPU_JIT"));
}
TEST(LegalizeTFTest, WrongDeviceTypeShallFail) {
absl::StatusOr<XlaCompiler::CompilationResult> result = CompileMlirModule(
kGpuMlirModuleStr,
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED,
"XLA_TPU_JIT");
EXPECT_THAT(result, absl_testing::StatusIs(absl::StatusCode::kNotFound));
}
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v2/mlir_roundtrip_flags.h
+119
View File
@@ -0,0 +1,119 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_MLIR_ROUNDTRIP_FLAGS_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_MLIR_ROUNDTRIP_FLAGS_H_
#include <string>
#include <vector>
#include "absl/strings/string_view.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/StringMap.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
namespace tensorflow {
struct ArrayInfoBase {
// The node type when the input node is imported. Typically needs to be
// specified when passing arbitrary nodes (some node attributes are removed).
DataType imported_dtype;
// Node "shape" attribute value.
TensorShapeProto shape;
};
struct ArrayInfo : public ArrayInfoBase {
using SubTypeInfo = ArrayInfoBase;
// DT_RESOURCE and DT_VARIANT have subtypes
std::vector<SubTypeInfo> subtypes;
};
struct GraphImportConfig {
// Returns string representation of config.
std::string str() const;
using InputArrays =
llvm::MapVector<std::string, ArrayInfo, llvm::StringMap<unsigned>>;
// The name assigned to the function which is the import result of the given
// graph. If empty, a default one will be used.
std::string graph_func_name;
// Maps input node names to node data types and shapes.
InputArrays inputs;
// name:index strings for the data outputs.
std::vector<std::string> outputs;
// name strings for the control outputs.
std::vector<std::string> control_outputs;
// Setting prune_unused_nodes to true, would prune unreachable nodes if
// output_arrays is specified.
bool prune_unused_nodes = false;
// If true, inputs of type LegacyFedInput are replaced with Placeholder ops.
// LegacyFedInput ops have two outputs unlike Placeholder which has only one
// output, so if both outputs of the LegacyFedInput ops are used then returns
// an error.
bool convert_legacy_fed_inputs = false;
// If true, the main graph will be treated as a function.
bool graph_as_function = false;
// If true, upgrade legacy features of the graph (for instance, functionalize
// control-flow).
bool upgrade_legacy = false;
// If true, functionalization is restricted to nodes that will be
// XLA-compiled. This is only needed if
// - `upgrade_legacy` is true
// - upgrading legacy features of the graph (which includes functionalization)
// runs before compilation cluster extraction (as for MLIR-based TPU bridge)
// - session runtime is used (session runtime has issues with function names
// rewritten by functionalization).
// Otherwise, this parameter should be set to false.
bool restrict_functionalization_to_compiled_nodes = false;
// If true, enables shape inference on input.
// TODO(jpienaar): This will be removed shortly.
bool enable_shape_inference = true;
// _output_shapes is an unregistered attribute which is used during
// GraphConstructor::ConvertGraph to override shapes. It is unfortunately
// not always set correctly (which is undesirable and should be addressed)
// so make it opt-in to consider it unconditionally also when importing the
// graph.
bool unconditionally_use_set_output_shapes = false;
// If set, use the value as the device type and mark the function graph for
// XLA compilation.
std::string xla_compile_device_type;
// If true, enables moving ops to different devices or moving unsupported ops
// out of a compilation cluster.
bool enable_soft_placement = false;
// If true, a function attribute, `tf._original_func_name`, will be set in
// functions which contains the corresponding original TF function name.
bool set_original_tf_func_name = false;
// If true, all functions in the graph will be converted to MLIR regardless of
// whether the functions are referenced by the nodes. This is needed if
// aliases and saved model object graph function matching is needed.
bool convert_all_functions_to_mlir = false;
};
struct GraphExportConfig {
// Whether to export the entry function to function library instead of the
// graph.
bool export_entry_func_to_flib = false;
// Whether to export functions using the name set in the attribute
// `tf._original_func_name` if it exists.
bool export_original_tf_func_name = false;
};
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_MLIR_ROUNDTRIP_FLAGS_H_
tensorflow/compiler/mlir/tf2xla/api/v2/testdata/empty_func.mlir
Vendored
+20
View File
@@ -0,0 +1,20 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> () {
func.return
}
}
tensorflow/compiler/mlir/tf2xla/api/v2/testdata/func_with_dead_ops.mlir
Vendored
+77
View File
@@ -0,0 +1,77 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
module attributes {tf.devices = {"/job:tpu_host_worker/replica:0/task:0/device:CPU:0", "/job:tpu_host_worker/replica:0/task:0/device:TPU:0", "/job:tpu_host_worker/replica:0/task:0/device:TPU:1", "/job:tpu_host_worker/replica:0/task:0/device:TPU_SYSTEM:0", "/job:tpu_host_worker/replica:0/task:1/device:CPU:0", "/job:tpu_host_worker/replica:0/task:1/device:TPU:0", "/job:tpu_host_worker/replica:0/task:1/device:TPU:1", "/job:tpu_host_worker/replica:0/task:1/device:TPU_SYSTEM:0", "/job:tpu_host_worker/replica:0/task:2/device:CPU:0", "/job:tpu_host_worker/replica:0/task:2/device:TPU:0", "/job:tpu_host_worker/replica:0/task:2/device:TPU:1", "/job:tpu_host_worker/replica:0/task:2/device:TPU_SYSTEM:0", "/job:tpu_host_worker/replica:0/task:3/device:CPU:0", "/job:tpu_host_worker/replica:0/task:3/device:TPU:0", "/job:tpu_host_worker/replica:0/task:3/device:TPU:1", "/job:tpu_host_worker/replica:0/task:3/device:TPU_SYSTEM:0"}, tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1847 : i32}} {
func.func @main(%arg0: tensor<i32> {tf._user_specified_name = "steps", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg1: tensor<*x!tf_type.resource<tensor<i64>>> {tf._user_specified_name = "899", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg2: tensor<*x!tf_type.resource<tensor<i64>>> {tf._user_specified_name = "901", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg3: tensor<*x!tf_type.resource<tensor<128x1024xf32>>> {tf._user_specified_name = "903", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg4: tensor<*x!tf_type.resource<tensor<1024xf32>>> {tf._user_specified_name = "905", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg5: tensor<*x!tf_type.resource<tensor<1024x1xf32>>> {tf._user_specified_name = "907", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg6: tensor<*x!tf_type.resource<tensor<i64>>> {tf._user_specified_name = "909", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg7: tensor<*x!tf_type.resource<tensor<25001x64xf32>>> {tf._user_specified_name = "911", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg8: tensor<*x!tf_type.resource<tensor<25001x64xf32>>> {tf._user_specified_name = "913", tf.device = "/job:tpu_host_worker/replica:0/task:1/device:CPU:0"}, %arg9: tensor<*x!tf_type.resource<tensor<25001x64xf32>>> {tf._user_specified_name = "915", tf.device = "/job:tpu_host_worker/replica:0/task:2/device:CPU:0"}, %arg10: tensor<*x!tf_type.resource<tensor<25001x64xf32>>> {tf._user_specified_name = "917", tf.device = "/job:tpu_host_worker/replica:0/task:3/device:CPU:0"}, %arg11: tensor<*x!tf_type.resource<tensor<25001x32xf32>>> {tf._user_specified_name = "919", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg12: tensor<*x!tf_type.resource<tensor<25001x32xf32>>> {tf._user_specified_name = "921", tf.device = "/job:tpu_host_worker/replica:0/task:1/device:CPU:0"}, %arg13: tensor<*x!tf_type.resource<tensor<25001x32xf32>>> {tf._user_specified_name = "923", tf.device = "/job:tpu_host_worker/replica:0/task:2/device:CPU:0"}, %arg14: tensor<*x!tf_type.resource<tensor<25001x32xf32>>> {tf._user_specified_name = "925", tf.device = "/job:tpu_host_worker/replica:0/task:3/device:CPU:0"}, %arg15: tensor<*x!tf_type.resource<tensor<6x32xf32>>> {tf._user_specified_name = "927", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg16: tensor<*x!tf_type.resource<tensor<6x32xf32>>> {tf._user_specified_name = "929", tf.device = "/job:tpu_host_worker/replica:0/task:1/device:CPU:0"}, %arg17: tensor<*x!tf_type.resource<tensor<6x32xf32>>> {tf._user_specified_name = "931", tf.device = "/job:tpu_host_worker/replica:0/task:2/device:CPU:0"}, %arg18: tensor<*x!tf_type.resource<tensor<6x32xf32>>> {tf._user_specified_name = "933", tf.device = "/job:tpu_host_worker/replica:0/task:3/device:CPU:0"}, %arg19: tensor<*x!tf_type.resource<tensor<128x1024xf32>>> {tf._user_specified_name = "935", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg20: tensor<*x!tf_type.resource<tensor<1024xf32>>> {tf._user_specified_name = "937", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg21: tensor<*x!tf_type.resource<tensor<1024x1xf32>>> {tf._user_specified_name = "939", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}) -> tensor<i64> attributes {allow_soft_placement = false, tf.entry_function = {control_outputs = "", inputs = "steps,unknown,unknown_0,unknown_1,unknown_2,unknown_3,unknown_4,unknown_5,unknown_6,unknown_7,unknown_8,unknown_9,unknown_10,unknown_11,unknown_12,unknown_13,unknown_14,unknown_15,unknown_16,unknown_17,unknown_18,unknown_19", outputs = "statefulpartitionedcall_RetVal"}} {
%0 = "tf.ReadVariableOp"(%arg19) : (tensor<*x!tf_type.resource<tensor<128x1024xf32>>>) -> tensor<128x1024xf32>
%1 = "tf.ReadVariableOp"(%arg1) : (tensor<*x!tf_type.resource<tensor<i64>>>) -> tensor<i64>
%2 = "tf.ReadVariableOp"(%arg2) : (tensor<*x!tf_type.resource<tensor<i64>>>) -> tensor<i64>
%3 = "tf.ReadVariableOp"(%arg4) : (tensor<*x!tf_type.resource<tensor<1024xf32>>>) -> tensor<1024xf32>
%4 = "tf.ReadVariableOp"(%arg3) : (tensor<*x!tf_type.resource<tensor<128x1024xf32>>>) -> tensor<128x1024xf32>
%5 = "tf.ReadVariableOp"(%arg5) : (tensor<*x!tf_type.resource<tensor<1024x1xf32>>>) -> tensor<1024x1xf32>
%6 = "tf.ReadVariableOp"(%arg20) : (tensor<*x!tf_type.resource<tensor<1024xf32>>>) -> tensor<1024xf32>
%7 = "tf.ReadVariableOp"(%arg21) : (tensor<*x!tf_type.resource<tensor<1024x1xf32>>>) -> tensor<1024x1xf32>
%8 = "tf.ReadVariableOp"(%arg6) : (tensor<*x!tf_type.resource<tensor<i64>>>) -> tensor<i64>
%9 = "tf.Const"() <{value = dense<"test"> : tensor<3x!tf_type.string>}> : () -> tensor<3x!tf_type.string>
%cst = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>
%11:4 = "tf.Split"(%cst, %0) {num_split = 4 : i32} : (tensor<i32>, tensor<128x1024xf32>) -> (tensor<32x1024xf32>, tensor<32x1024xf32>, tensor<32x1024xf32>, tensor<32x1024xf32>)
%cst_0 = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>
%12:4 = "tf.Split"(%cst_0, %4) {num_split = 4 : i32} : (tensor<i32>, tensor<128x1024xf32>) -> (tensor<32x1024xf32>, tensor<32x1024xf32>, tensor<32x1024xf32>, tensor<32x1024xf32>)
%cst_1 = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>
%cst_2 = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>
%13:20 = tf_device.replicate {devices = {TPU_REPLICATED_CORE_0 = ["/job:tpu_host_worker/replica:0/task:0/device:TPU:0", "/job:tpu_host_worker/replica:0/task:2/device:TPU:0"], TPU_REPLICATED_CORE_1 = ["/job:tpu_host_worker/replica:0/task:0/device:TPU:1", "/job:tpu_host_worker/replica:0/task:2/device:TPU:1"], TPU_REPLICATED_CORE_2 = ["/job:tpu_host_worker/replica:0/task:1/device:TPU:0", "/job:tpu_host_worker/replica:0/task:3/device:TPU:0"], TPU_REPLICATED_CORE_3 = ["/job:tpu_host_worker/replica:0/task:1/device:TPU:1", "/job:tpu_host_worker/replica:0/task:3/device:TPU:1"], TPU_REPLICATED_HOST_0 = ["/job:tpu_host_worker/replica:0/task:0/device:CPU:0", "/job:tpu_host_worker/replica:0/task:2/device:CPU:0"], TPU_REPLICATED_HOST_1 = ["/job:tpu_host_worker/replica:0/task:0/device:CPU:0", "/job:tpu_host_worker/replica:0/task:2/device:CPU:0"], TPU_REPLICATED_HOST_2 = ["/job:tpu_host_worker/replica:0/task:1/device:CPU:0", "/job:tpu_host_worker/replica:0/task:3/device:CPU:0"], TPU_REPLICATED_HOST_3 = ["/job:tpu_host_worker/replica:0/task:1/device:CPU:0", "/job:tpu_host_worker/replica:0/task:3/device:CPU:0"]}, n = 2 : i32} {
%16:40 = "tf_device.parallel_execute"() ({
%19:10 = "tf_device.launch"() <{device = "TPU_REPLICATED_CORE_0"}> ({
%20:10 = "tf.TPUExecute"(%arg0, %11#0, %1, %2, %3, %12#0, %5, %6, %7, %8, %9) : (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>, tensor<3x!tf_type.string>) -> (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>)
tf_device.return %20#0, %20#1, %20#2, %20#3, %20#4, %20#5, %20#6, %20#7, %20#8, %20#9 : tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>
}) : () -> (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>)
tf_device.return %19#0, %19#1, %19#2, %19#3, %19#4, %19#5, %19#6, %19#7, %19#8, %19#9 : tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>
}, {
%19:10 = "tf_device.launch"() <{device = "TPU_REPLICATED_CORE_1"}> ({
%20:10 = "tf.TPUExecute"(%arg0, %11#1, %1, %2, %3, %12#1, %5, %6, %7, %8, %9) : (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>, tensor<3x!tf_type.string>) -> (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>)
tf_device.return %20#0, %20#1, %20#2, %20#3, %20#4, %20#5, %20#6, %20#7, %20#8, %20#9 : tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>
}) : () -> (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>)
tf_device.return %19#0, %19#1, %19#2, %19#3, %19#4, %19#5, %19#6, %19#7, %19#8, %19#9 : tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>
}, {
%19:10 = "tf_device.launch"() <{device = "TPU_REPLICATED_CORE_2"}> ({
%20:10 = "tf.TPUExecute"(%arg0, %11#2, %1, %2, %3, %12#2, %5, %6, %7, %8, %9) : (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>, tensor<3x!tf_type.string>) -> (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>)
tf_device.return %20#0, %20#1, %20#2, %20#3, %20#4, %20#5, %20#6, %20#7, %20#8, %20#9 : tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>
}) : () -> (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>)
tf_device.return %19#0, %19#1, %19#2, %19#3, %19#4, %19#5, %19#6, %19#7, %19#8, %19#9 : tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>
}, {
%19:10 = "tf_device.launch"() <{device = "TPU_REPLICATED_CORE_3"}> ({
%20:10 = "tf.TPUExecute"(%arg0, %11#3, %1, %2, %3, %12#3, %5, %6, %7, %8, %9) : (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>, tensor<3x!tf_type.string>) -> (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>)
tf_device.return %20#0, %20#1, %20#2, %20#3, %20#4, %20#5, %20#6, %20#7, %20#8, %20#9 : tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>
}) : () -> (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>)
tf_device.return %19#0, %19#1, %19#2, %19#3, %19#4, %19#5, %19#6, %19#7, %19#8, %19#9 : tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>
}) : () -> (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>, tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>, tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>, tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>)
%17 = "tf.Concat"(%cst_1, %16#5, %16#15, %16#25, %16#35) : (tensor<i32>, tensor<32x1024xf32>, tensor<32x1024xf32>, tensor<32x1024xf32>, tensor<32x1024xf32>) -> tensor<128x1024xf32>
%18 = "tf.Concat"(%cst_2, %16#1, %16#11, %16#21, %16#31) : (tensor<i32>, tensor<32x1024xf32>, tensor<32x1024xf32>, tensor<32x1024xf32>, tensor<32x1024xf32>) -> tensor<128x1024xf32>
tf_device.return %16#0, %16#9, %16#8, %16#7, %16#6, %17, %16#4, %16#3, %16#2, %18 : tensor<i32>, tensor<i64>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<128x1024xf32>, tensor<1024xf32>, tensor<i64>, tensor<i64>, tensor<128x1024xf32>
}
"tf.AssignVariableOp"(%arg19, %13#18) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<128x1024xf32>>>, tensor<128x1024xf32>) -> ()
"tf.AssignVariableOp"(%arg1, %13#16) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<i64>>>, tensor<i64>) -> ()
"tf.AssignVariableOp"(%arg2, %13#14) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<i64>>>, tensor<i64>) -> ()
"tf.AssignVariableOp"(%arg4, %13#12) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<1024xf32>>>, tensor<1024xf32>) -> ()
"tf.AssignVariableOp"(%arg3, %13#10) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<128x1024xf32>>>, tensor<128x1024xf32>) -> ()
"tf.AssignVariableOp"(%arg5, %13#8) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<1024x1xf32>>>, tensor<1024x1xf32>) -> ()
"tf.AssignVariableOp"(%arg20, %13#6) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<1024xf32>>>, tensor<1024xf32>) -> ()
"tf.AssignVariableOp"(%arg21, %13#4) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<1024x1xf32>>>, tensor<1024x1xf32>) -> ()
"tf.AssignVariableOp"(%arg6, %13#2) <{validate_shape = true}> {_has_manual_control_dependencies = true} : (tensor<*x!tf_type.resource<tensor<i64>>>, tensor<i64>) -> ()
%14 = "tf.ReadVariableOp"(%arg2) {device = ""} : (tensor<*x!tf_type.resource<tensor<i64>>>) -> tensor<i64>
%15 = "tf.Identity"(%14) {device = ""} : (tensor<i64>) -> tensor<i64>
return %15 : tensor<i64>
}
}
tensorflow/compiler/mlir/tf2xla/api/v2/testdata/graph_with_flib_def.txt
Vendored
+1535
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/mlir/tf2xla/api/v2/testdata/invalid_executor.mlir
Vendored
+32
View File
@@ -0,0 +1,32 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() {
tf_executor.graph {
%control = tf_executor.island {
tf_executor.yield
}
tf_executor.fetch %control : !tf_executor.control
}
tf_executor.graph {
%control = tf_executor.island {
tf_executor.yield
}
tf_executor.fetch %control : !tf_executor.control
}
return
}
}
tensorflow/compiler/mlir/tf2xla/api/v2/testdata/outside_compilation.mlir
Vendored
+150
View File
@@ -0,0 +1,150 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
module attributes {tf.devices = {"/job:localhost/replica:0/task:0/device:CPU:0", "/job:localhost/replica:0/task:0/device:TPU:0", "/job:localhost/replica:0/task:0/device:TPU:1", "/job:localhost/replica:0/task:0/device:TPU_SYSTEM:0"}, tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1654 : i32}} {
func.func @main(%arg0: tensor<*x!tf_type.resource> {tf._user_specified_name = "input_1", tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}) attributes {allow_soft_placement = true, tf.entry_function = {control_outputs = "while,image_sample/write_summary/summary_cond", inputs = "image_sample_write_summary_summary_cond_input_1", outputs = ""}} {
tf_executor.graph {
%outputs, %control = tf_executor.island wraps "tf.Const"() {device = "", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%outputs_0, %control_1 = tf_executor.island wraps "tf.Const"() {device = "", value = dense<1> : tensor<i32>} : () -> tensor<i32>
%outputs_2, %control_3 = tf_executor.island wraps "tf.TPUReplicatedInput"(%outputs, %outputs_0) {device = "", index = -1 : i64, is_mirrored_variable = false, is_packed = false} : (tensor<i32>, tensor<i32>) -> tensor<*xi32>
%outputs_4, %control_5 = tf_executor.island wraps "tf.Const"() {device = "", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%outputs_6, %control_7 = tf_executor.island wraps "tf.Const"() {_post_device_rewrite = true, device = "", value = dense<> : tensor<0xi32>} : () -> tensor<0xi32>
%outputs_8, %control_9 = tf_executor.island wraps "tf.Const"() {_post_device_rewrite = true, device = "", value = dense<> : tensor<0xi32>} : () -> tensor<0xi32>
%outputs_10, %control_11 = tf_executor.island wraps "tf.Pack"(%outputs_6, %outputs_8) {axis = 0 : i64, device = ""} : (tensor<0xi32>, tensor<0xi32>) -> tensor<*xi32>
%outputs_12, %control_13 = tf_executor.island wraps "tf.Max"(%outputs_10, %outputs_4) {device = "", keep_dims = false} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
%control_14 = tf_executor.island wraps "tf.NoOp"() {_pivot_for_cluster = "cluster_sample_sequence", device = ""} : () -> ()
%control_15 = tf_executor.island(%control_14) wraps "tf.NoOp"() {_has_manual_control_dependencies = true, _tpu_replicate = "cluster_sample_sequence", device = ""} : () -> ()
%control_16 = tf_executor.island(%control_15) wraps "tf.NoOp"() {device = ""} : () -> ()
%control_17 = tf_executor.island(%control_15) wraps "tf.NoOp"() {device = ""} : () -> ()
%control_18 = tf_executor.island(%control_14) wraps "tf.TPUReplicateMetadata"() {_has_manual_control_dependencies = true, _tpu_replicate = "cluster_sample_sequence", allow_soft_placement = true, computation_shape = [], device = "", device_assignment = [], host_compute_core = [], num_cores_per_replica = 1 : i64, num_replicas = 2 : i64, padding_map = [], step_marker_location = "STEP_MARK_AT_ENTRY", topology = "", tpu_compile_options_proto = "", use_spmd_for_xla_partitioning = false, use_tpu = true} : () -> ()
%outputs_19, %control_20 = tf_executor.island(%control_18) wraps "tf.Const"() {_tpu_replicate = "cluster_sample_sequence", device = "", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%outputs_21, %control_22 = tf_executor.island(%control_18) wraps "tf.Const"() {_tpu_replicate = "cluster_sample_sequence", device = "", value = dense<5> : tensor<i64>} : () -> tensor<i64>
%outputs_23, %control_24 = tf_executor.island(%control_18) wraps "tf.Const"() {_tpu_replicate = "cluster_sample_sequence", device = "", value = dense<[3, 32, 32, 1]> : tensor<4xi32>} : () -> tensor<4xi32>
%outputs_25, %control_26 = tf_executor.island(%control_18) wraps "tf.TPUCompilationResult"() {_tpu_compilation_status = "cluster_sample_sequence", device = ""} : () -> tensor<!tf_type.string>
%outputs_27, %control_28 = tf_executor.island(%control_18) wraps "tf.Const"() {_tpu_replicate = "cluster_sample_sequence", device = "", value = dense<3> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_29, %control_30 = tf_executor.island(%control_18) wraps "tf.Const"() {_tpu_replicate = "cluster_sample_sequence", device = "", value = dense<[1, 1, 1, 3]> : tensor<4xi32>} : () -> tensor<4xi32>
%outputs_31, %control_32 = tf_executor.island(%control_18) wraps "tf.Const"() {_tpu_replicate = "cluster_sample_sequence", device = "", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%outputs_33, %control_34 = tf_executor.island(%control_18) wraps "tf.Const"() {_tpu_replicate = "cluster_sample_sequence", device = "", value = dense<true> : tensor<i1>} : () -> tensor<i1>
%outputs_35, %control_36 = tf_executor.island(%control_18) wraps "tf.Identity"(%outputs_2) {_tpu_input_identity = true, _tpu_replicate = "cluster_sample_sequence", device = ""} : (tensor<*xi32>) -> tensor<*xi32>
%outputs_37, %control_38 = tf_executor.island wraps "tf.Equal"(%outputs_35, %outputs_31) {_tpu_replicate = "cluster_sample_sequence", device = "", incompatible_shape_error = true} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi1>
%outputs_39, %control_40 = tf_executor.island wraps "tf.LogicalAnd"(%outputs_37, %outputs_33) {_tpu_replicate = "cluster_sample_sequence", device = ""} : (tensor<*xi1>, tensor<i1>) -> tensor<*xi1>
%outputs_41, %control_42 = tf_executor.island(%control_18) wraps "tf.Const"() {_tpu_replicate = "cluster_sample_sequence", device = "", value = dense<3> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_43, %control_44 = tf_executor.island(%control_18) wraps "tf.Const"() {_tpu_replicate = "cluster_sample_sequence", device = "", value = dense<1024> : tensor<i32>} : () -> tensor<i32>
%outputs_45, %control_46 = tf_executor.island wraps "tf.TensorListReserve"(%outputs_41, %outputs_43) {_tpu_replicate = "cluster_sample_sequence", device = ""} : (tensor<1xi32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
%outputs_47, %control_48 = tf_executor.island(%control_18) wraps "tf.Const"() {_tpu_replicate = "cluster_sample_sequence", device = "", value = dense<[1, 0]> : tensor<2xi32>} : () -> tensor<2xi32>
%outputs_49, %control_50 = tf_executor.island(%control_18) wraps "tf.Const"() {_tpu_replicate = "cluster_sample_sequence", device = "", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%outputs_51, %control_52 = tf_executor.island(%control_18) wraps "tf.Const"() {_tpu_replicate = "cluster_sample_sequence", device = "", value = dense<-1> : tensor<i32>} : () -> tensor<i32>
%outputs_53:4, %control_54 = tf_executor.island wraps "tf.While"(%outputs_49, %outputs_51, %outputs_19, %outputs_45) {_num_original_outputs = 4 : i64, _read_only_resource_inputs = [], _tpu_replicate = "cluster_sample_sequence", _xla_propagate_compile_time_consts = true, body = @while_body_260, cond = @while_cond_250, device = "", is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>) -> (tensor<i32>, tensor<i32>, tensor<i32>, tensor<!tf_type.variant>)
%outputs_55, %control_56 = tf_executor.island wraps "tf.TensorListStack"(%outputs_53#3, %outputs_27) {_tpu_replicate = "cluster_sample_sequence", device = "", num_elements = 1024 : i64} : (tensor<!tf_type.variant>, tensor<1xi32>) -> tensor<*xf32>
%outputs_57, %control_58 = tf_executor.island wraps "tf.Transpose"(%outputs_55, %outputs_47) {_tpu_replicate = "cluster_sample_sequence", device = ""} : (tensor<*xf32>, tensor<2xi32>) -> tensor<*xf32>
%outputs_59, %control_60 = tf_executor.island wraps "tf.Reshape"(%outputs_57, %outputs_23) {_tpu_replicate = "cluster_sample_sequence", device = ""} : (tensor<*xf32>, tensor<4xi32>) -> tensor<*xf32>
%outputs_61, %control_62 = tf_executor.island wraps "tf.Tile"(%outputs_59, %outputs_29) {_tpu_replicate = "cluster_sample_sequence", device = ""} : (tensor<*xf32>, tensor<4xi32>) -> tensor<*xf32>
%outputs_63, %control_64 = tf_executor.island wraps "tf.If"(%outputs_39, %outputs_61, %arg0, %outputs_21) {_read_only_resource_inputs = [], _tpu_replicate = "cluster_sample_sequence", _xla_propagate_compile_time_consts = true, device = "", else_branch = @image_sample_write_summary_summary_cond_false_710, is_stateless = false, then_branch = @image_sample_write_summary_summary_cond_true_700} : (tensor<*xi1>, tensor<*xf32>, tensor<*x!tf_type.resource>, tensor<i64>) -> tensor<*xi1>
%outputs_65, %control_66 = tf_executor.island wraps "tf.Identity"(%outputs_63) {_tpu_replicate = "cluster_sample_sequence", device = ""} : (tensor<*xi1>) -> tensor<*xi1>
tf_executor.fetch %control_54, %control_64 : !tf_executor.control, !tf_executor.control
}
return
}
func.func private @while_body_260(%arg0: tensor<i32> {tf._user_specified_name = "while/loop_counter"}, %arg1: tensor<i32> {tf._user_specified_name = "while/maximum_iterations"}, %arg2: tensor<i32>, %arg3: tensor<!tf_type.variant>) -> (tensor<*xi32>, tensor<*xi32>, tensor<*xi32>, tensor<*x!tf_type.variant>) attributes {tf._construction_context = "kEagerRuntime", tf.signature.is_stateful} {
%0:4 = tf_executor.graph {
%outputs, %control = tf_executor.island wraps "tf.Const"() {device = "", value = dense<1> : tensor<i32>} : () -> tensor<i32>
%outputs_0, %control_1 = tf_executor.island wraps "tf.Const"() {device = "", value = dense<1> : tensor<i32>} : () -> tensor<i32>
%outputs_2, %control_3 = tf_executor.island wraps "tf.Const"() {device = "", value = dense<3> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_4, %control_5 = tf_executor.island wraps "tf.RandomUniform"(%outputs_2) {device = "", seed = 87654321 : i64, seed2 = 0 : i64} : (tensor<1xi32>) -> tensor<*xf32>
%outputs_6, %control_7 = tf_executor.island wraps "tf.AddV2"(%arg2, %outputs) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<*xi32>
%outputs_8, %control_9 = tf_executor.island wraps "tf.Identity"(%outputs_6) {device = ""} : (tensor<*xi32>) -> tensor<*xi32>
%outputs_10, %control_11 = tf_executor.island wraps "tf.TensorListSetItem"(%arg3, %arg2, %outputs_4) {device = "", resize_if_index_out_of_bounds = false} : (tensor<!tf_type.variant>, tensor<i32>, tensor<*xf32>) -> tensor<*x!tf_type.variant>
%outputs_12, %control_13 = tf_executor.island wraps "tf.Identity"(%outputs_10) {device = ""} : (tensor<*x!tf_type.variant>) -> tensor<*x!tf_type.variant>
%outputs_14, %control_15 = tf_executor.island wraps "tf.AddV2"(%arg0, %outputs_0) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<*xi32>
%outputs_16, %control_17 = tf_executor.island wraps "tf.Identity"(%outputs_14) {device = ""} : (tensor<*xi32>) -> tensor<*xi32>
%outputs_18, %control_19 = tf_executor.island wraps "tf.Identity"(%arg1) {device = ""} : (tensor<i32>) -> tensor<*xi32>
tf_executor.fetch %outputs_16, %outputs_18, %outputs_8, %outputs_12 : tensor<*xi32>, tensor<*xi32>, tensor<*xi32>, tensor<*x!tf_type.variant>
}
return %0#0, %0#1, %0#2, %0#3 : tensor<*xi32>, tensor<*xi32>, tensor<*xi32>, tensor<*x!tf_type.variant>
}
func.func private @while_cond_250(%arg0: tensor<i32> {tf._user_specified_name = "while/loop_counter"}, %arg1: tensor<i32> {tf._user_specified_name = "while/maximum_iterations"}, %arg2: tensor<i32>, %arg3: tensor<!tf_type.variant>) -> tensor<*xi1> attributes {tf._construction_context = "kEagerRuntime"} {
%0 = tf_executor.graph {
%outputs, %control = tf_executor.island wraps "tf.Const"() {device = "", value = dense<1024> : tensor<i32>} : () -> tensor<i32>
%outputs_0, %control_1 = tf_executor.island wraps "tf.Less"(%arg2, %outputs) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<*xi1>
%outputs_2, %control_3 = tf_executor.island wraps "tf.Identity"(%outputs_0) {device = ""} : (tensor<*xi1>) -> tensor<*xi1>
tf_executor.fetch %outputs_2 : tensor<*xi1>
}
return %0 : tensor<*xi1>
}
func.func private @image_sample_write_summary_summary_cond_false_710(%arg0: tensor<3x32x32x3xf32>, %arg1: tensor<*x!tf_type.resource>, %arg2: tensor<i64>) -> tensor<*xi1> attributes {tf._construction_context = "kEagerRuntime"} {
%0 = tf_executor.graph {
%outputs, %control = tf_executor.island wraps "tf.Const"() {device = "", value = dense<false> : tensor<i1>} : () -> tensor<i1>
%outputs_0, %control_1 = tf_executor.island wraps "tf.Identity"(%outputs) {device = ""} : (tensor<i1>) -> tensor<*xi1>
tf_executor.fetch %outputs_0 : tensor<*xi1>
}
return %0 : tensor<*xi1>
}
func.func private @image_sample_write_summary_summary_cond_true_700(%arg0: tensor<3x32x32x3xf32> {tf._user_specified_name = "Tile"}, %arg1: tensor<*x!tf_type.resource> {tf._user_specified_name = "writer"}, %arg2: tensor<i64> {tf._user_specified_name = "Const_3"}) -> tensor<*xi1> attributes {tf._construction_context = "kEagerRuntime", tf.signature.is_stateful} {
%0 = tf_executor.graph {
%outputs, %control = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<[3, 32, 32, 3]> : tensor<4xi32>} : () -> tensor<4xi32>
%outputs_0, %control_1 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%outputs_2, %control_3 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<""> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
%outputs_4, %control_5 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<"Condition x >= 0 did not hold element-wise:"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
%outputs_6, %control_7 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<"x (image_sample/write_summary/summary_cond/assert_non_negative/x:0) = "> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
%outputs_8, %control_9 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<""> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
%outputs_10, %control_11 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<"Condition x >= 0 did not hold element-wise:"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
%outputs_12, %control_13 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<"x (image_sample/write_summary/summary_cond/assert_non_negative/x:0) = "> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
%outputs_14, %control_15 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%outputs_16, %control_17 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
%outputs_18, %control_19 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%outputs_20, %control_21 = tf_executor.island wraps "tf.Range"(%outputs_18, %outputs_14, %outputs_16) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<*xi32>
%outputs_22, %control_23 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<3> : tensor<i32>} : () -> tensor<i32>
%outputs_24, %control_25 = tf_executor.island wraps "tf.LessEqual"(%outputs_0, %outputs_22) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<*xi1>
%outputs_26, %control_27 = tf_executor.island wraps "tf.All"(%outputs_24, %outputs_20) {device = "/device:CPU:0", keep_dims = false} : (tensor<*xi1>, tensor<*xi32>) -> tensor<*xi1>
%control_28 = tf_executor.island wraps "tf.Assert"(%outputs_26, %outputs_2, %outputs_4, %outputs_6, %outputs_22) {device = "/device:CPU:0", summarize = 3 : i64} : (tensor<*xi1>, tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<i32>) -> ()
%outputs_29, %control_30 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<[3, 32, 32, 3]> : tensor<4xi32>} : () -> tensor<4xi32>
%control_31 = tf_executor.island wraps "tf.NoOp"() {device = "/device:CPU:0"} : () -> ()
%outputs_32, %control_33 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<4> : tensor<i32>} : () -> tensor<i32>
%control_34 = tf_executor.island wraps "tf.NoOp"() {device = "/device:CPU:0"} : () -> ()
%outputs_35, %control_36 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
%outputs_37, %control_38 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<0.000000e+00> : tensor<f32>} : () -> tensor<f32>
%outputs_39, %control_40 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<2.550000e+02> : tensor<f32>} : () -> tensor<f32>
%outputs_41, %control_42 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<2.555000e+02> : tensor<f32>} : () -> tensor<f32>
%outputs_43, %control_44 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_45, %control_46 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<3> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_47, %control_48 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_49, %control_50 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_51, %control_52 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<3> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_53, %control_54 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_55, %control_56 = tf_executor.island wraps "tf.StridedSlice"(%outputs, %outputs_49, %outputs_51, %outputs_53) {begin_mask = 0 : i64, device = "/device:CPU:0", ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<*xi32>
%outputs_57, %control_58 = tf_executor.island wraps "tf.AsString"(%outputs_55) {device = "/device:CPU:0", fill = "", precision = -1 : i64, scientific = false, shortest = false, width = -1 : i64} : (tensor<*xi32>) -> tensor<*x!tf_type.string>
%outputs_59, %control_60 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_61, %control_62 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_63, %control_64 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32>
%outputs_65, %control_66 = tf_executor.island wraps "tf.StridedSlice"(%outputs, %outputs_59, %outputs_61, %outputs_63) {begin_mask = 0 : i64, device = "/device:CPU:0", ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<*xi32>
%outputs_67, %control_68 = tf_executor.island wraps "tf.AsString"(%outputs_65) {device = "/device:CPU:0", fill = "", precision = -1 : i64, scientific = false, shortest = false, width = -1 : i64} : (tensor<*xi32>) -> tensor<*x!tf_type.string>
%outputs_69, %control_70 = tf_executor.island wraps "tf.Pack"(%outputs_57, %outputs_67) {axis = 0 : i64, device = "/device:CPU:0"} : (tensor<*x!tf_type.string>, tensor<*x!tf_type.string>) -> tensor<*x!tf_type.string>
%outputs_71, %control_72 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<"\0A\08\0A\06images"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
%outputs_73, %control_74 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<"image_sample"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
%outputs_75, %control_76 = tf_executor.island wraps "tf.Mul"(%arg0, %outputs_41) {device = "/device:CPU:0"} : (tensor<3x32x32x3xf32>, tensor<f32>) -> tensor<*xf32>
%outputs_77, %control_78 = tf_executor.island wraps "tf.ClipByValue"(%outputs_75, %outputs_37, %outputs_39) {device = "/device:CPU:0"} : (tensor<*xf32>, tensor<f32>, tensor<f32>) -> tensor<*xf32>
%outputs_79, %control_80 = tf_executor.island wraps "tf.Cast"(%outputs_77) {Truncate = false, device = "/device:CPU:0"} : (tensor<*xf32>) -> tensor<*xui8>
%outputs_81, %control_82 = tf_executor.island wraps "tf.StridedSlice"(%outputs_79, %outputs_43, %outputs_45, %outputs_47) {begin_mask = 1 : i64, device = "/device:CPU:0", ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64} : (tensor<*xui8>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<*xui8>
%outputs_83, %control_84 = tf_executor.island wraps "tf.EncodePng"(%outputs_81) {compression = -1 : i64, device = "/device:CPU:0"} : (tensor<*xui8>) -> tensor<*x!tf_type.string>
%outputs_85, %control_86 = tf_executor.island wraps "tf.ConcatV2"(%outputs_69, %outputs_83, %outputs_35) {device = "/device:CPU:0"} : (tensor<*x!tf_type.string>, tensor<*x!tf_type.string>, tensor<i32>) -> tensor<*x!tf_type.string>
%control_87 = tf_executor.island wraps "tf.WriteSummary"(%arg1, %arg2, %outputs_85, %outputs_73, %outputs_71) {_has_manual_control_dependencies = true, device = "/device:CPU:0"} : (tensor<*x!tf_type.resource>, tensor<i64>, tensor<*x!tf_type.string>, tensor<!tf_type.string>, tensor<!tf_type.string>) -> ()
%outputs_88, %control_89 = tf_executor.island(%control_87) wraps "tf.Const"() {device = "/device:CPU:0", value = dense<true> : tensor<i1>} : () -> tensor<i1>
%control_90 = tf_executor.island(%control_28, %control_87) wraps "tf.NoOp"() {device = ""} : () -> ()
%outputs_91, %control_92 = tf_executor.island(%control_90) wraps "tf.Identity"(%outputs_88) {device = ""} : (tensor<i1>) -> tensor<*xi1>
tf_executor.fetch %outputs_91, %control_28, %control_87 : tensor<*xi1>, !tf_executor.control, !tf_executor.control
}
return %0 : tensor<*xi1>
}
}
tensorflow/compiler/mlir/tf2xla/api/v2/testdata/valid_executor.mlir
Vendored
+24
View File
@@ -0,0 +1,24 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() {
tf_executor.graph {
%0:2 = tf_executor.island wraps "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:TPU:0", dtype = "tfdtype$DT_INT32", value = #tf_type<tensor_proto : "0x746674656E736F722464747970653A2044545F494E5433320A74656E736F725F7368617065207B0A202064696D207B0A2020202073697A653A20320A20207D0A7D0A74656E736F725F636F6E74656E743A20225C3230305C3030305C3030305C3030305C3230305C3030305C3030305C303030220A"> : tensor<2xi32>} : () -> tensor<2xi32> loc("Empty/shape")
tf_executor.fetch
}
func.return
}
}
tensorflow/compiler/mlir/tf2xla/api/v2/testdata/valid_graph.txt
Vendored
+44
View File
@@ -0,0 +1,44 @@
node {
name: "Empty/shape"
op: "Const"
device: "/job:localhost/replica:0/task:0/device:TPU:0"
attr {
key: "_output_shapes"
value {
list {
shape {
dim {
size: 2
}
}
}
}
}
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
dim {
size: 2
}
}
tensor_content: "\200\000\000\000\200\000\000\000"
}
}
}
experimental_debug_info {
}
}
library {
}
versions {
producer: 268
}
tensorflow/compiler/mlir/tf2xla/api/v2/testing/BUILD
+53
View File
@@ -0,0 +1,53 @@
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow/compiler/mlir/tf2xla/api/v2:__subpackages__",
],
)
cc_library(
name = "compile_mlir",
testonly = True,
srcs = ["compile_mlir.cc"],
hdrs = ["compile_mlir.h"],
deps = [
"//tensorflow/compiler/jit",
"//tensorflow/compiler/mlir/tf2xla/api/v2:legalize_tf",
"//tensorflow/compiler/mlir/tf2xla/internal/utils:test_metadata_config",
"//tensorflow/compiler/tf2xla:xla_compiler",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:test",
"//tensorflow/core/protobuf:for_core_protos_cc",
"//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
"//tensorflow/core/tpu/kernels:tpu_compile_op_support",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:string_view",
"@llvm-project//mlir:Pass",
"@xla//xla:shape_util",
"@xla//xla/client:client_library",
"@xla//xla/stream_executor:platform",
"@xla//xla/stream_executor:platform_manager",
],
)
cc_library(
name = "utils",
testonly = True,
srcs = ["utils.cc"],
hdrs = ["utils.h"],
deps = ["//tensorflow/core/platform:resource_loader"],
)
tf_cc_test(
name = "utils_test",
srcs = ["utils_test.cc"],
deps = [
":utils",
"@com_google_googletest//:gtest_main",
],
)
tensorflow/compiler/mlir/tf2xla/api/v2/testing/compile_mlir.cc
+81
View File
@@ -0,0 +1,81 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/api/v2/testing/compile_mlir.h"
#include <memory>
#include <vector>
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/utils/test_metadata_config.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/client/client_library.h"
#include "xla/shape.h"
#include "xla/stream_executor/platform.h"
#include "xla/stream_executor/platform_manager.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
namespace testing {
using tpu::MlirToHloArgs;
using tpu::ShardingAndIndex;
using tpu::TPUCompileMetadataProto;
absl::StatusOr<XlaCompiler::CompilationResult> CompileMlirModule(
const char* mlir_module_str,
ConfigProto::Experimental::MlirBridgeRollout rollout_state,
absl::string_view device_type) {
MlirToHloArgs mlir_to_hlo_args;
mlir_to_hlo_args.rollout_state = rollout_state;
mlir_to_hlo_args.mlir_module = mlir_module_str;
TF_ASSIGN_OR_RETURN(se::Platform * platform,
se::PlatformManager::PlatformWithName("Host"));
TF_ASSIGN_OR_RETURN(
auto client, xla::ClientLibrary::GetOrCreateCompileOnlyClient(platform));
std::vector<TensorShape> arg_shapes;
TPUCompileMetadataProto metadata_proto;
// Configure metadata requires parsing the module and if we are testing a
// failure, we ignore this particular set up error assuming we'll not get
// far enough to need valid metadata.
tensorflow::tf2xla::internal::ConfigureMetadata(mlir_module_str, arg_shapes,
metadata_proto)
.IgnoreError();
bool use_tuple_args = true;
std::vector<ShardingAndIndex> arg_core_mapping;
std::vector<std::vector<xla::Shape>> per_core_arg_shapes;
std::vector<std::unique_ptr<mlir::Pass>> custom_legalization_passes;
return LegalizeMlirToHlo(mlir_to_hlo_args, metadata_proto, use_tuple_args,
device_type, custom_legalization_passes,
/*shape_determination_fns=*/{}, arg_shapes,
&arg_core_mapping, &per_core_arg_shapes, client);
}
} // namespace testing
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v2/testing/compile_mlir.h
+40
View File
@@ -0,0 +1,40 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_TESTING_COMPILE_MLIR_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_TESTING_COMPILE_MLIR_H_
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
namespace testing {
// Compiles the given MLIR module to XLA HLO.
absl::StatusOr<XlaCompiler::CompilationResult> CompileMlirModule(
const char* mlir_module_str,
ConfigProto::Experimental::MlirBridgeRollout rollout_state,
absl::string_view device_type = "XLA_TPU_JIT");
} // namespace testing
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_TESTING_COMPILE_MLIR_H_
tensorflow/compiler/mlir/tf2xla/api/v2/testing/utils.cc
+35
View File
@@ -0,0 +1,35 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/api/v2/testing/utils.h"
#include <string>
#include "tensorflow/core/platform/resource_loader.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
namespace testing {
std::string TestDataPath() {
return tensorflow::GetDataDependencyFilepath(
"tensorflow/compiler/mlir/tf2xla/api/v2/testdata/");
}
} // namespace testing
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v2/testing/utils.h
+34
View File
@@ -0,0 +1,34 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_TESTING_UTILS_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_TESTING_UTILS_H_
#include <string>
namespace tensorflow {
namespace tf2xla {
namespace v2 {
namespace testing {
// Returns the path to the testdata directory.
std::string TestDataPath();
} // namespace testing
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_TESTING_UTILS_H_
tensorflow/compiler/mlir/tf2xla/api/v2/testing/utils_test.cc
+47
View File
@@ -0,0 +1,47 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/api/v2/testing/utils.h"
#include <stdlib.h>
#include <string>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
namespace tensorflow {
namespace tf2xla {
namespace v2 {
namespace testing {
namespace {
class UtilsTest : public ::testing::Test {};
TEST_F(UtilsTest, TestDataPathSucceeds) {
std::string expected_test_data_path_regex =
".*tensorflow/compiler/mlir/tf2xla/api/v2/testdata/";
std::string result_test_data_path = TestDataPath();
EXPECT_THAT(result_test_data_path,
::testing::ContainsRegex(expected_test_data_path_regex));
}
} // namespace
} // namespace testing
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor.cc
+184
View File
@@ -0,0 +1,184 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor.h"
#include <memory>
#include <string>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Transforms/Passes.h" // from @llvm-project
#include "tensorflow/compiler/jit/flags.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/data_dumper_logger_config.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/logging_hooks.h"
#include "xla/tsl/lib/monitoring/counter.h"
#include "xla/tsl/platform/status.h"
#include "tensorflow/core/platform/error_payloads.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/util/debug_data_dumper.h"
#include "tsl/platform/error_logging.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
using mlir::LogicalResult;
using mlir::ModuleOp;
using mlir::OpPassManager;
using mlir::PassManager;
using mlir::func::FuncOp;
auto *tf_dialect_to_executor_dialect_status = tsl::monitoring::Counter<1>::New(
"/tensorflow/core/tf2xla/api/v2/tf_dialect_to_executor_dialect_status",
"Counts how often a successful export from TF Dialect to Executor Dialect "
"is",
"status");
constexpr char kExportSuccess[] = "success";
constexpr char kExportFailed[] = "failed";
namespace {
void AddTfDialectToExecutorPasses(OpPassManager &pm) {
pm.addPass(mlir::TF::CreateTFRegionControlFlowToFunctional());
// First, we need to convert from functional, to executor dialect.
pm.addNestedPass<FuncOp>(
mlir::CreateFunctionalToExecutorDialectConversionPass());
// Do a single pass to split the graph's single island op into an island per
// op as expected by the following passes.
pm.addNestedPass<FuncOp>(mlir::TF::CreateSplitIntoIslandPerOpPass());
pm.addNestedPass<FuncOp>(mlir::TFDevice::CreateReplicateToIslandPass(
/*legacy_graph_export=*/false));
pm.addNestedPass<FuncOp>(
mlir::TFDevice::CreateReplicaIDToDeviceOrdinalPass());
pm.addNestedPass<FuncOp>(mlir::TFDevice::CreateParallelExecuteToIslandsPass(
/*legacy_graph_export=*/false));
pm.addNestedPass<FuncOp>(mlir::TFDevice::CreateLaunchToDeviceAttributePass(
/*legacy_graph_export=*/false));
// Do a single pass to encode necessary control deps in the IR according to
// the results of side effect analysis.
pm.addPass(
mlir::tf_executor::CreateTFExecutorUpdateControlDependenciesPass());
pm.addNestedPass<FuncOp>(mlir::TFTPU::CreateTPUDevicePropagationPass());
pm.addNestedPass<FuncOp>(mlir::TFTPU::CreateTPUColocateSplitsPass());
pm.addPass(mlir::createSymbolDCEPass());
pm.addNestedPass<FuncOp>(
mlir::tf_executor::CreateTFExecutorGraphPruningPass());
if (tensorflow::GetMlirCommonFlags()
->tf_mlir_enable_convert_control_to_data_outputs_pass) {
bool composite_tpuexecute_side_effects =
tensorflow::GetMlirCommonFlags()
->tf_mlir_enable_composite_tpuexecute_side_effects;
pm.addPass(
mlir::tf_executor::CreateTFExecutorConvertControlToDataOutputsPass(
composite_tpuexecute_side_effects));
}
pm.addPass(mlir::TF::CreateVerifySuitableForExportPass());
}
absl::Status RecordStatusIfError(absl::Status status) {
if (status.ok()) {
return absl::OkStatus();
}
tf_dialect_to_executor_dialect_status->GetCell(kExportFailed)->IncrementBy(1);
VLOG(1) << "Failed to export from TF Dialect to TF Executor Dialect. "
<< status;
constexpr char bridge_subcomponent[] =
"TFXLA_TF_FUNCTIONAL_TO_EXECUTOR_EXPORT_v2";
constexpr char kBridgeComponent[] = "TFXLABridge";
tsl::OkOrSetErrorCounterPayload(
tensorflow::core::platform::ErrorSourceProto::MLIR_BRIDGE_PHASE_1,
status);
tsl::error_logging::Log(kBridgeComponent, bridge_subcomponent,
status.ToString())
.IgnoreError();
return status;
}
} // namespace
absl::Status ExportFromTensorflowDialectToExecutor(
ModuleOp module, llvm::StringRef module_name) {
PassManager tf_to_executor(module.getContext());
::tensorflow::applyTensorflowAndCLOptions(tf_to_executor);
tf_to_executor.enableVerifier();
AddTfDialectToExecutorPasses(tf_to_executor);
if (VLOG_IS_ON(1) ||
DEBUG_DATA_DUMPER()->ShouldDump(module_name.str(), kDebugGroupMain)) {
::tensorflow::DumpMlirOpToFile(
DEBUG_DATA_DUMPER()->GetDumpFilename(
module_name.str(), kDebugGroupMain,
"tfxla_bridge_v2_tfdialect_to_executor_before"),
module, llvm::StringRef(), &tf_to_executor);
if (VLOG_IS_ON(2) ||
DEBUG_DATA_DUMPER()->ShouldDump(
module_name.str(), kDebugGroupBridgePhase1ExecutorExport)) {
internal::EnablePassIRPrinting(
tf_to_executor, kDebugGroupBridgePhase1ExecutorExport, module_name);
}
}
LogicalResult result = tf_to_executor.run(module);
if (VLOG_IS_ON(1) ||
DEBUG_DATA_DUMPER()->ShouldDump(module_name.str(), kDebugGroupMain)) {
::tensorflow::DumpMlirOpToFile(
DEBUG_DATA_DUMPER()->GetDumpFilename(
module_name.str(), kDebugGroupMain,
"tfxla_bridge_v2_tfdialect_to_executor_after"),
module, llvm::StringRef(), &tf_to_executor);
}
if (result.failed()) {
return RecordStatusIfError(
absl::InternalError("Failed to export from TF Dialect to TF Executor "
"Dialect. Read LLVM Pipeline Error"));
}
tf_dialect_to_executor_dialect_status->GetCell(kExportSuccess)
->IncrementBy(1);
return absl::OkStatus();
}
mlir::PassPipelineRegistration<> tf_dialect_to_executor_pipeline(
"tf-dialect-to-executor-v2",
"Run passes to convert from TF Dialect to Executor in preparation for "
"exporting module back to TF Graph.",
AddTfDialectToExecutorPasses);
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor.h
+51
View File
@@ -0,0 +1,51 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_TF_DIALECT_TO_EXECUTOR_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_TF_DIALECT_TO_EXECUTOR_H_
#include "llvm/ADT/StringRef.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
// Given the input Module op that's in the Tensorflow Dialect, convert the MLIR
// module in place to the Tensorflow Executor Dialect. Returns an OK Status if
// success, otherwise failure with an error message.
// The Tensorflow Executor Dialect is required to export an MLIR module to a
// Tensorflow GraphDef. This API will add control dependencies and verify that
// the conversion was successful.
//
// This also converts the Tensorflow Dialect MLIR into the Tensorflow Executor
// dialect that is suitable to be exported to GraphDef. Graph -> MLIR -> Graph
// is not perfectly round trippable, so this API will attempt to make the module
// exportable and verify some properties of the Tensorflow Executor MLIR that
// are required by Graph Export. It will return an error if it cannot.
//
// Input: A MLIR Module in the Tensorflow Dialect with no
// `tf_device.cluster_func` ops.
// Output: A MLIR module in the Tensorflow Executor Dialect.
ABSL_DEPRECATED("Use tensorflow::tf2xla::v2::ConvertGraphToTfExecutor instead.")
absl::Status ExportFromTensorflowDialectToExecutor(
mlir::ModuleOp module, llvm::StringRef module_name = llvm::StringRef());
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_TF_DIALECT_TO_EXECUTOR_H_
tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor_test.cc
+133
View File
@@ -0,0 +1,133 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor.h"
#include <stdlib.h>
#include <cstdint>
#include <string>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/DialectRegistry.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "mlir/Parser/Parser.h" // from @llvm-project
#include "tensorflow/compiler/mlir/register_common_dialects.h"
#include "tensorflow/compiler/mlir/tf2xla/api/v2/testing/utils.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "tensorflow/core/lib/monitoring/cell_reader.h"
#include "tensorflow/core/platform/resource_loader.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
namespace {
constexpr char kExportStreamzName[] =
"/tensorflow/core/tf2xla/api/v2/tf_dialect_to_executor_dialect_status";
constexpr char kExportSuccess[] = "success";
constexpr char kExportFailed[] = "failed";
using mlir::DialectRegistry;
using mlir::MLIRContext;
using mlir::ModuleOp;
using mlir::OwningOpRef;
using ::tensorflow::monitoring::testing::CellReader;
using ::tensorflow::tf2xla::v2::testing::TestDataPath;
size_t CountSubstring(absl::string_view str, absl::string_view substr) {
size_t count = 0;
size_t idx = str.find(substr);
while (idx != std::string::npos) {
count++;
idx = str.find(substr, idx + 1);
}
return count;
}
class TensorflowDialectToExecutorTest : public ::testing::Test {
public:
TensorflowDialectToExecutorTest() {
mlir::RegisterCommonToolingDialects(registry_);
context_.appendDialectRegistry(registry_);
context_.loadAllAvailableDialects();
}
absl::Status CreateMlirModule(std::string mlir_module_filename) {
std::string mlir_module_path = TestDataPath() + mlir_module_filename;
mlir_module_ =
mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context_);
if (!mlir_module_) {
return absl::Status(
absl::StatusCode::kNotFound,
absl::StrCat("Could not find MLIR module at ", mlir_module_path));
}
return absl::OkStatus();
}
DialectRegistry registry_;
MLIRContext context_;
OwningOpRef<mlir::ModuleOp> mlir_module_;
};
TEST_F(TensorflowDialectToExecutorTest, ConvertsToExecutor) {
CellReader<int64_t> compilation_status(kExportStreamzName);
TF_ASSERT_OK(CreateMlirModule("empty_func.mlir"));
TF_EXPECT_OK(ExportFromTensorflowDialectToExecutor(*mlir_module_));
EXPECT_EQ(compilation_status.Delta(kExportSuccess), 1);
EXPECT_EQ(compilation_status.Delta(kExportFailed), 0);
}
TEST_F(TensorflowDialectToExecutorTest, ErrorsWhenCannotConvert) {
CellReader<int64_t> compilation_status(kExportStreamzName);
TF_ASSERT_OK(CreateMlirModule("invalid_executor.mlir"));
EXPECT_FALSE(ExportFromTensorflowDialectToExecutor(*mlir_module_).ok());
EXPECT_EQ(compilation_status.Delta(kExportSuccess), 0);
EXPECT_EQ(compilation_status.Delta(kExportFailed), 1);
}
TEST_F(TensorflowDialectToExecutorTest, PrunesDeadOps) {
CellReader<int64_t> compilation_status(kExportStreamzName);
TF_ASSERT_OK(CreateMlirModule("func_with_dead_ops.mlir"));
TF_EXPECT_OK(ExportFromTensorflowDialectToExecutor(*mlir_module_));
std::string module_dump;
llvm::raw_string_ostream raw_stream(module_dump);
mlir_module_->print(raw_stream);
EXPECT_EQ(compilation_status.Delta(kExportSuccess), 1);
EXPECT_EQ(compilation_status.Delta(kExportFailed), 0);
EXPECT_EQ(
CountSubstring(module_dump, "tf_executor.island wraps \"tf.Concat\""), 2);
}
} // namespace
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v2/tf_executor_to_graph.cc
+842
View File
@@ -0,0 +1,842 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/api/v2/tf_executor_to_graph.h"
#include <string>
#include <utility>
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/container/inlined_vector.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/types/optional.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Location.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/SymbolTable.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "mlir/Support/DebugStringHelper.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/op_or_arg_name_mapper.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/convert_type.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/export_utils.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/translate_utils.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/verify_suitable_for_graph_export.h"
#include "tensorflow/compiler/mlir/utils/name_utils.h"
#include "xla/status_macros.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/graph_to_functiondef.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/framework/versions.pb.h"
#include "tensorflow/core/graph/algorithm.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/graph/regularization/util.h"
#include "tensorflow/core/graph/tensor_id.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/status.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
using mlir::BlockArgument;
using mlir::Dialect;
using mlir::Operation;
using mlir::SymbolTable;
using mlir::Value;
using mlir::func::FuncOp;
using tsl::StatusOr;
namespace {
constexpr char kDeviceAttr[] = "tf.device";
constexpr char kResourceArgUniqueIdAttr[] = "tf._resource_arg_unique_id";
constexpr char kEntryFuncAttr[] = "tf.entry_function";
constexpr char kAliasingAttr[] = "tf.aliasing_output";
// OpOrArgLocNameMapper that legalizes the returned name.
class LegalizedOpOrValLocNameMapper : public OpOrArgLocNameMapper {
private:
std::string GetName(OpOrVal op_or_val) override {
std::string name = OpOrArgLocNameMapper::GetName(op_or_val);
assert(!name.empty() && "expected non-empty name");
mlir::LegalizeNodeName(name);
return name;
}
};
// Finds first inner op if `op` is a tf_executor.island. Otherwise `op` is
// returned.
Operation* GetIslandInnerOpOrSelf(mlir::Operation* op) {
auto island = llvm::dyn_cast<mlir::tf_executor::IslandOp>(op);
if (island) return &island.GetBody().front();
return op;
}
// Stateful helper class to export a function into a Graph.
class Exporter {
public:
// Converts the given Module to a Graph. The given module should only contain
// one entry function, which is identified by name "main". This entry function
// is converted to the base of the graph graph. The rest of the functions are
// converted to the library functions in that graph.
static absl::Status Convert(mlir::ModuleOp module,
const GraphExportConfig& configs,
std::unique_ptr<Graph>* graph,
FunctionLibraryDefinition* flib_def,
absl::flat_hash_set<Node*>* control_ret_nodes);
// Converts a given FuncOp to a FunctionDef and adds it to the function
// definition library
static absl::Status ConvertLibFunction(
const GraphExportConfig& configs, const Dialect* tf_dialect,
const SymbolTable& symbol_table, FuncOp function,
FunctionLibraryDefinition* flib_def,
llvm::SmallDenseSet<FuncOp>& visited_functions);
// Converts the given FuncOp to a Graph. The arguments and returns of
// function are added to the graph with special op names kArgOp and kRetOp.
// Later on, this graph can be converted a function definition and added to
// another graph.
static absl::StatusOr<std::unique_ptr<Graph>> Convert(
const GraphExportConfig& configs, const Dialect* tf_dialect,
const SymbolTable& symbol_table, FuncOp function,
FunctionLibraryDefinition* flib_def,
llvm::SmallDenseSet<FuncOp>& visited_functions,
absl::flat_hash_set<Node*>* control_ret_nodes);
private:
explicit Exporter(const GraphExportConfig* configs, Graph* graph,
const Dialect* tf_dialect, const SymbolTable* symbol_table)
: configs_(*configs),
graph_(graph),
tf_dialect_(tf_dialect),
symbol_table_(*symbol_table) {
graph_->ToGraphDef(&graphdef_);
}
absl::Status AddArgumentNode(BlockArgument arg, unsigned index,
llvm::StringRef name);
absl::Status AddFetchNode(FuncOp function, mlir::tf_executor::FetchOp fetch,
llvm::ArrayRef<llvm::StringRef> names);
absl::Status AddInstructionNode(Operation* inst);
void UseOriginalFunctionNames(NodeDef& node_def);
absl::Status AddEdge(Operation* inst);
absl::StatusOr<std::unique_ptr<NodeDef>> GetArgumentNode(
BlockArgument arg, unsigned index, llvm::StringRef name);
absl::StatusOr<std::unique_ptr<NodeDef>> GetReturnNode(FuncOp function,
Value operand,
unsigned index,
llvm::StringRef name);
absl::Status GetControlRetNodes(
mlir::tf_executor::FetchOp fetch,
absl::flat_hash_set<Node*>* control_ret_nodes);
// Adds one edge between src_node and dst_node. If it is not a control edge,
// an index is used to find out the right operand of the dst_node.
absl::Status AddEdgeBetweenNodes(Value src, Node* dst_node,
unsigned dst_index);
const GraphExportConfig& configs_;
Graph* graph_;
GraphDef graphdef_;
LegalizedOpOrValLocNameMapper op_to_name_;
absl::flat_hash_map<Operation*, Node*> nodes_;
llvm::DenseMap<BlockArgument, Node*> args_;
// One single return operation can return multiple results, and each of them
// will be converted to one node in the graph.
typedef absl::InlinedVector<Node*, 4> NodeVector;
absl::flat_hash_map<Operation*, NodeVector> returns_;
const mlir::Dialect* tf_dialect_;
const SymbolTable& symbol_table_;
};
std::string FindFunctionName(const GraphExportConfig& configs, FuncOp func) {
if (auto original_func_name =
func->getAttrOfType<mlir::StringAttr>("tf._original_func_name");
configs.export_original_tf_func_name && original_func_name) {
return original_func_name.str();
}
return func.getName().str();
}
absl::StatusOr<std::unique_ptr<NodeDef>> Exporter::GetArgumentNode(
BlockArgument arg, unsigned index, llvm::StringRef name) {
auto func = arg.getParentRegion()->getParentOfType<FuncOp>();
auto node_def = std::make_unique<NodeDef>();
if (!name.empty())
node_def->set_name(ParseTensorName(name.str()).node());
else
node_def->set_name(
std::string(op_to_name_.GetUniqueName(func.getName().str())));
node_def->set_op(FunctionLibraryDefinition::kArgOp);
mlir::TensorType arg_type = mlir::cast<mlir::TensorType>(arg.getType());
if (auto resource_type =
mlir::dyn_cast<mlir::TF::ResourceType>(arg_type.getElementType())) {
llvm::ArrayRef<mlir::TensorType> subtypes = resource_type.getSubtypes();
if (!subtypes.empty()) {
AttrValue handle_dtypes_attr;
AttrValue handle_shapes_attr;
for (mlir::TensorType subtype : subtypes) {
DataType dtype;
TF_RETURN_IF_ERROR(ConvertToDataType(subtype.getElementType(), &dtype));
handle_dtypes_attr.mutable_list()->add_type(dtype);
SetTensorShapeProto(subtype,
handle_shapes_attr.mutable_list()->add_shape());
}
(*node_def->mutable_attr())["_handle_dtypes"] = handle_dtypes_attr;
(*node_def->mutable_attr())["_handle_shapes"] = handle_shapes_attr;
}
}
TF_RETURN_IF_ERROR(
SetShapeAttribute("_output_shapes", arg_type, node_def->mutable_attr()));
DataType dtype;
TF_RETURN_IF_ERROR(ConvertToDataType(arg_type.getElementType(), &dtype));
AttrValue type_attr;
type_attr.set_type(dtype);
(*node_def->mutable_attr())["T"] = type_attr;
AttrValue index_attr;
index_attr.set_i(index);
(*node_def->mutable_attr())["index"] = index_attr;
if (auto device_attr =
func.getArgAttrOfType<mlir::StringAttr>(index, kDeviceAttr))
*node_def->mutable_device() = device_attr.getValue().str();
llvm::ArrayRef<mlir::NamedAttribute> func_arg_i_attrs =
mlir::function_interface_impl::getArgAttrs(func, index);
absl::flat_hash_set<absl::string_view> attrs_to_ignore = {kDeviceAttr,
kAliasingAttr};
TF_RETURN_IF_ERROR(ConvertAttributes(func_arg_i_attrs, attrs_to_ignore,
/*remove_ref_type=*/false,
node_def->mutable_attr()));
return node_def;
}
absl::StatusOr<std::unique_ptr<NodeDef>> Exporter::GetReturnNode(
FuncOp function, Value operand, unsigned index, llvm::StringRef name) {
auto node_def = std::make_unique<NodeDef>();
if (!name.empty())
node_def->set_name(ParseTensorName(name.str()).node());
else
node_def->set_name(
std::string(op_to_name_.GetUniqueName(function.getName().str())));
node_def->set_op(FunctionLibraryDefinition::kRetOp);
DataType dtype;
TF_RETURN_IF_ERROR(ConvertToDataType(
mlir::cast<mlir::TensorType>(operand.getType()).getElementType(),
&dtype));
AttrValue type_attr;
type_attr.set_type(dtype);
(*node_def->mutable_attr())["T"] = type_attr;
AttrValue index_attr;
index_attr.set_i(index);
(*node_def->mutable_attr())["index"] = index_attr;
if (auto device_attr =
function.getResultAttrOfType<mlir::StringAttr>(index, kDeviceAttr))
*node_def->mutable_device() = device_attr.getValue().str();
llvm::ArrayRef<mlir::NamedAttribute> func_res_i_attrs =
function.getResultAttrs(index);
absl::flat_hash_set<absl::string_view> attrs_to_ignore = {kDeviceAttr};
TF_RETURN_IF_ERROR(ConvertAttributes(func_res_i_attrs, attrs_to_ignore,
/*remove_ref_type=*/false,
node_def->mutable_attr()));
return node_def;
}
absl::Status Exporter::AddEdgeBetweenNodes(Value src, Node* dst_node,
unsigned dst_index) {
if (auto input_result = mlir::dyn_cast<mlir::OpResult>(src)) {
auto* input_inst = GetIslandInnerOpOrSelf(input_result.getOwner());
// Replaces the input node with NextIteration sink if it is a NextIteration
// source.
if (auto next_iter_source =
llvm::dyn_cast<mlir::tf_executor::NextIterationSourceOp>(
input_inst))
input_inst = next_iter_source.GetSink();
auto node_it = nodes_.find(input_inst);
TF_RET_CHECK(node_it != nodes_.end())
<< "Use of OpResult encountered before def!";
if (mlir::isa<mlir::tf_executor::ControlType>(input_result.getType())) {
graph_->AddControlEdge(node_it->second, dst_node,
/*allow_duplicates=*/true);
} else {
graph_->AddEdge(node_it->second, input_result.getResultNumber(), dst_node,
dst_index);
}
return absl::OkStatus();
}
auto input_arg = mlir::cast<BlockArgument>(src);
auto input_node_it = args_.find(input_arg);
TF_RET_CHECK(input_node_it != args_.end())
<< "Use of BlockArgument encounted before def!";
// For argument, there is only one result output, so the index is always 0.
graph_->AddEdge(input_node_it->second, 0, dst_node, dst_index);
return absl::OkStatus();
}
absl::Status Exporter::AddEdge(Operation* inst) {
// For tf_executor.fetch, add only its data edges. Control edges are captured
// later.
if (auto fetch = llvm::dyn_cast<mlir::tf_executor::FetchOp>(inst)) {
for (auto operand_and_idx : llvm::enumerate(fetch.getOperands())) {
Value operand = operand_and_idx.value();
if (mlir::isa<mlir::tf_executor::ControlType>(operand.getType())) break;
auto* dst_node = returns_[fetch][operand_and_idx.index()];
TF_RETURN_IF_ERROR(AddEdgeBetweenNodes(operand, dst_node, 0));
}
return absl::OkStatus();
}
// For tf_executor.NextIteration.Sink, skip its token operand and add data and
// control edges with their index offset by 1.
if (auto next_iter_sink =
llvm::dyn_cast<mlir::tf_executor::NextIterationSinkOp>(inst)) {
auto* dst_node = nodes_[inst];
TF_RETURN_IF_ERROR(
AddEdgeBetweenNodes(next_iter_sink.getInput(), dst_node, 0));
for (auto control_and_idx :
llvm::enumerate(next_iter_sink.getControlInputs()))
TF_RETURN_IF_ERROR(AddEdgeBetweenNodes(control_and_idx.value(), dst_node,
control_and_idx.index() + 1));
return absl::OkStatus();
}
// For tf_executor.NextIteration.Source, op can be skipped as it is assumed
// there are no operands.
if (llvm::isa<mlir::tf_executor::NextIterationSourceOp>(inst)) {
assert(inst->getNumOperands() == 0);
return absl::OkStatus();
}
Operation* op = GetIslandInnerOpOrSelf(inst);
auto* dst_node = nodes_[op];
int operand_offset = 0;
// For tf_executor.island, add data edges from its wrapped op before control
// edges.
if (auto island = llvm::dyn_cast<mlir::tf_executor::IslandOp>(inst)) {
for (auto operand_and_idx : llvm::enumerate(op->getOperands()))
TF_RETURN_IF_ERROR(AddEdgeBetweenNodes(operand_and_idx.value(), dst_node,
operand_and_idx.index()));
operand_offset = op->getNumOperands();
}
// For all other ops (including tf_executor.island), add remaining edges.
for (auto operand_and_idx : llvm::enumerate(inst->getOperands()))
TF_RETURN_IF_ERROR(
AddEdgeBetweenNodes(operand_and_idx.value(), dst_node,
operand_and_idx.index() + operand_offset));
return absl::OkStatus();
}
void Exporter::UseOriginalFunctionNames(NodeDef& node_def) {
if (!configs_.export_original_tf_func_name) return;
auto& attrs = *node_def.mutable_attr();
auto try_use_original_func_name = [this](std::string* name) {
if (auto func = symbol_table_.lookup<FuncOp>(*name)) {
if (auto original_func_name =
func->getAttrOfType<mlir::StringAttr>("tf._original_func_name")) {
*name = original_func_name.str();
}
}
};
// Change its op name if it is a legacy call.
try_use_original_func_name(node_def.mutable_op());
// Change any function attributes in the attrs.
for (auto& iter : attrs) {
auto& attr = iter.second;
if (attr.has_func()) {
try_use_original_func_name(attr.mutable_func()->mutable_name());
} else if (attr.has_list()) {
for (auto& func_attr : *attr.mutable_list()->mutable_func()) {
try_use_original_func_name(func_attr.mutable_name());
}
}
}
}
absl::Status Exporter::AddInstructionNode(Operation* inst) {
std::unique_ptr<NodeDef> node_def;
int graph_hash_value = graph_regularization::ComputeHash(graphdef_);
auto name = op_to_name_.GetUniqueName(inst, graph_hash_value);
// Convert registered TF ops to NodeDef. Only registered ops are handled to
// ensure that PopulateDerivedAttrs adds the correct attributes.
TF_ASSIGN_OR_RETURN(node_def,
ConvertTFDialectOpToNodeDef(
inst, name, /*ignore_unregistered_attrs=*/false));
UseOriginalFunctionNames(*node_def);
TF_ASSIGN_OR_RETURN(Node * node, graph_->AddNode(std::move(*node_def)));
DCHECK(node != nullptr);
nodes_[inst] = node;
return absl::OkStatus();
}
bool IsEntryFunctionArg(BlockArgument arg) {
return arg.getParentRegion()->getParentOfType<FuncOp>().getName() == "main";
}
// Creates argument nodes from Block argument. If a name is supplied, that
// name will be used instead of generating a unique name.
absl::Status Exporter::AddArgumentNode(BlockArgument arg, unsigned index,
llvm::StringRef name) {
TF_ASSIGN_OR_RETURN(auto node_def, GetArgumentNode(arg, index, name));
TF_ASSIGN_OR_RETURN(Node * node, graph_->AddNode(std::move(*node_def)));
args_[arg] = node;
return absl::OkStatus();
}
// Creates return nodes per operand of a FetchOp. If names is supplied, those
// names will be used per node in order instead of generating a unique name.
absl::Status Exporter::AddFetchNode(FuncOp function,
mlir::tf_executor::FetchOp fetch,
llvm::ArrayRef<llvm::StringRef> names) {
auto& return_nodes = returns_[fetch];
for (auto operand_and_idx : llvm::enumerate(fetch.getOperands())) {
if (mlir::isa<mlir::tf_executor::ControlType>(
operand_and_idx.value().getType()))
break;
TF_ASSIGN_OR_RETURN(
auto node_def,
GetReturnNode(function, operand_and_idx.value(),
operand_and_idx.index(),
names.empty() ? "" : names[operand_and_idx.index()]));
TF_ASSIGN_OR_RETURN(Node * node, graph_->AddNode(std::move(*node_def)));
return_nodes.push_back(node);
}
return absl::OkStatus();
}
// Collects control ret Nodes based on tf_executor.graph's associated
// tf_executor.fetch control inputs.
absl::Status Exporter::GetControlRetNodes(
mlir::tf_executor::FetchOp fetch,
absl::flat_hash_set<Node*>* control_ret_nodes) {
for (Value fetch_operand : fetch.getOperands()) {
if (mlir::isa<mlir::tf_executor::ControlType>(fetch_operand.getType())) {
Operation* defining_op =
GetIslandInnerOpOrSelf(fetch_operand.getDefiningOp());
auto node_it = nodes_.find(defining_op);
TF_RET_CHECK(node_it != nodes_.end());
control_ret_nodes->insert(node_it->second);
}
}
return absl::OkStatus();
}
// After conversion from MLIR the input names are all blank which causes
// graph compilation to fail. This uses the edges to fix up the input names.
void FixupInputNamesFromEdges(Graph* graph) {
for (Node* n : graph->nodes()) {
if (n->IsOp()) {
NodeDef* node_def = n->mutable_def();
node_def->clear_input();
for (const Edge* e : n->in_edges()) {
Node* src = e->src();
if (src->IsOp()) {
Graph::AddInput(node_def, src->name(), e->src_output());
}
}
}
}
}
absl::StatusOr<std::unique_ptr<Graph>> Exporter::Convert(
const GraphExportConfig& configs, const Dialect* tf_dialect,
const SymbolTable& symbol_table, FuncOp function,
FunctionLibraryDefinition* flib_def,
llvm::SmallDenseSet<FuncOp>& visited_functions,
absl::flat_hash_set<Node*>* control_ret_nodes) {
mlir::Block& block = function.front();
// Extract input & output names if set.
llvm::SmallVector<llvm::StringRef, 2> input_names;
llvm::SmallVector<llvm::StringRef, 2> output_names;
llvm::SmallVector<llvm::StringRef, 2> unique_output_names;
auto dict_attr =
function->getAttrOfType<mlir::DictionaryAttr>(kEntryFuncAttr);
if (dict_attr) {
TF_RET_CHECK(mlir::isa<mlir::StringAttr>(dict_attr.get("inputs")))
<< "inputs missing in entry function attribute";
TF_RET_CHECK(mlir::isa<mlir::StringAttr>(dict_attr.get("outputs")))
<< "outputs missing in entry function attribute";
mlir::cast<mlir::StringAttr>(dict_attr.get("inputs"))
.getValue()
.split(input_names, ',', /*MaxSplit=*/-1, /*KeepEmpty=*/false);
mlir::cast<mlir::StringAttr>(dict_attr.get("outputs"))
.getValue()
.split(output_names, ',', /*MaxSplit=*/-1, /*KeepEmpty=*/false);
}
auto graph = std::make_unique<Graph>(OpRegistry::Global());
// Extract version info.
VersionDef versions;
auto module = function->getParentOfType<mlir::ModuleOp>();
if (mlir::succeeded(ExtractTfVersions(module, &versions))) {
graph->set_versions(versions);
}
Exporter exporter(&configs, graph.get(), tf_dialect, &symbol_table);
auto graph_op = llvm::cast<mlir::tf_executor::GraphOp>(block.front());
// Set input and output names and increment the use counter for them to help
// generate unique names.
if (!output_names.empty()) {
const int num_data_results = graph_op.getNumResults();
const int64_t output_names_size = output_names.size();
TF_RET_CHECK(output_names_size == num_data_results)
<< "output names (" << output_names.size()
<< ") != terminator operands (" << num_data_results << ")";
llvm::DenseMap<Operation*, llvm::StringRef> output_op_to_name;
llvm::StringMap<Operation*> name_to_op;
for (const auto& it : llvm::enumerate(graph_op.GetFetch().getOperands())) {
// Skip control rets.
const int64_t index = it.index();
if (index >= num_data_results) break;
// TODO(jpienaar): If there is a result index specified, ensure only one
// and that it matches the result index of the op.
std::string name(output_names[index]);
auto tensor_id = ParseTensorName(name);
std::string tensor_id_node(tensor_id.node());
assert(!tensor_id_node.empty() && "expected non-empty name");
mlir::LegalizeNodeName(tensor_id_node);
// Ensure name does not get reused.
unique_output_names.push_back(
exporter.op_to_name_.GetUniqueName(tensor_id_node));
}
}
if (!input_names.empty()) {
TF_RET_CHECK(input_names.size() == block.getNumArguments());
for (const auto& it : llvm::enumerate(function.getArguments())) {
// TODO(lyandy): Update when changing feed/fetch import.
std::string name(input_names[it.index()]);
assert(!name.empty() && "expected non-empty name");
mlir::LegalizeNodeName(name);
auto tensor_id = ParseTensorName(name);
TF_RET_CHECK(tensor_id.index() == 0)
<< "input port designation not supported";
// Only assign user of argument the input name if the main graph did not
// have its _Arg nodes lifted into the functions arguments.
// Ensure name does not get reused.
(void)exporter.op_to_name_.GetUniqueName(name);
}
}
// Adds nodes for basic block (function) arguments.
for (auto it : llvm::enumerate(block.getArguments())) {
int index = it.index();
auto arg = it.value();
mlir::Type type = arg.getType();
if (!mlir::isa<mlir::TensorType>(type)) {
return absl::InvalidArgumentError(absl::StrCat(
"FuncOps arguments must have tensor types. Found ",
mlir::debugString(type), " in function ", function.getName().str()));
}
TF_RETURN_IF_ERROR(exporter.AddArgumentNode(
arg, index, !input_names.empty() ? input_names[index] : ""));
}
auto convert_called_function = [&](llvm::StringRef name) {
auto func = symbol_table.lookup<FuncOp>(name);
if (func != nullptr) {
TF_RETURN_IF_ERROR(ConvertLibFunction(configs, tf_dialect, symbol_table,
func, flib_def, visited_functions));
// TODO(prakalps): Optimize to only add the requested function to graph
// library rather than the all the functions exported so far.
TF_RETURN_IF_ERROR(graph->mutable_flib_def()->AddLibrary(*flib_def));
}
return absl::OkStatus();
};
// Adds nodes for operations.
for (Operation& inst : graph_op.GetBody()) {
for (auto type : inst.getResultTypes())
if (!mlir::isa<mlir::TensorType, mlir::tf_executor::ControlType,
mlir::tf_executor::TokenType>(type))
return absl::InvalidArgumentError(absl::StrCat(
"Values must be of tensor type, TensorFlow control type, or "
"TensorFlow token type. Found ",
mlir::debugString(type)));
if (llvm::isa<mlir::tf_executor::NextIterationSourceOp>(inst)) {
// Skip tf_executor.NextIteration.Source as associated
// tf_executor.NextIteration.Sink will be used instead.
continue;
} else if (auto fetch = llvm::dyn_cast<mlir::tf_executor::FetchOp>(inst)) {
TF_RETURN_IF_ERROR(
exporter.AddFetchNode(function, fetch, unique_output_names));
} else if (auto island =
llvm::dyn_cast<mlir::tf_executor::IslandOp>(inst)) {
Operation& inner_op = island.GetBody().front();
auto op_name = GetTensorFlowOpName(inner_op.getName().getStringRef());
if (llvm::isa<FuncOp>(inner_op) && op_name.ok()) {
// If it is TF Control dialect specific op, look up custom operation
// in the module and first convert that, then add it to function
// definition library
// TODO(prakalps): If two functions have cyclic dependence, this will
// introduce an infinite loop.
TF_RETURN_IF_ERROR(convert_called_function(op_name.value().str()));
}
if (IsLegacyCallInstruction(&inner_op)) {
TF_RETURN_IF_ERROR(convert_called_function(
inner_op.getAttrOfType<mlir::SymbolRefAttr>("f")
.getLeafReference()
.getValue()));
}
TF_RETURN_IF_ERROR(exporter.AddInstructionNode(&inner_op));
} else {
TF_RETURN_IF_ERROR(exporter.AddInstructionNode(&inst));
}
}
// Adds edges between the argument, operation and return nodes.
for (Operation& inst : graph_op.GetBody()) {
TF_RETURN_IF_ERROR(exporter.AddEdge(&inst));
}
// Fixes the edges between the inserted nodes and special "_SOURCE" and
// "_SINK".
FixupSourceAndSinkEdges(graph.get());
FixupInputNamesFromEdges(graph.get());
TF_RETURN_IF_ERROR(
exporter.GetControlRetNodes(graph_op.GetFetch(), control_ret_nodes));
return graph;
}
absl::Status Exporter::ConvertLibFunction(
const GraphExportConfig& configs, const Dialect* tf_dialect,
const SymbolTable& symbol_table, FuncOp function,
FunctionLibraryDefinition* flib_def,
llvm::SmallDenseSet<FuncOp>& visited_functions) {
// Return early if the function has already been exported.
bool is_new_function = visited_functions.insert(function).second;
if (!is_new_function) return absl::OkStatus();
auto function_name = FindFunctionName(configs, function);
// TODO(fengliuai): use a small flib_def to reduce overhead
absl::flat_hash_set<Node*> control_ret_nodes;
TF_ASSIGN_OR_RETURN(
auto sub_graph,
Exporter::Convert(configs, tf_dialect, symbol_table, function, flib_def,
visited_functions, &control_ret_nodes));
const auto control_ret = [&](const Node* n) -> std::optional<std::string> {
return control_ret_nodes.contains(n)
? std::make_optional<std::string>(n->name())
: std::nullopt;
};
FunctionDef func_def;
TF_RETURN_IF_ERROR(
GraphToFunctionDef(*sub_graph, function_name, control_ret, &func_def));
// Checks for gradient attribute. If present converts the gradient function
// and populates the GradientDef.
auto grad_string = mlir::TF::TensorFlowDialect::GetGradientAttrName();
if (auto attr =
function->getAttrOfType<mlir::FlatSymbolRefAttr>(grad_string)) {
auto grad_func = symbol_table.lookup<FuncOp>(attr.getValue());
TF_RETURN_IF_ERROR(ConvertLibFunction(configs, tf_dialect, symbol_table,
grad_func, flib_def,
visited_functions));
GradientDef grad;
grad.set_function_name(function_name);
grad.set_gradient_func(grad_func.getName().str());
TF_RETURN_IF_ERROR(flib_def->AddGradientDef(grad));
}
auto stateful_string = mlir::TF::TensorFlowDialect::GetStatefulAttrName();
if (auto attr = function->getAttrOfType<mlir::UnitAttr>(stateful_string)) {
func_def.mutable_signature()->set_is_stateful(true);
}
// Ignore the gradient and is_stateful attribute on the function as they have
// been handled above. Ignore the entry func attribute as it is an MLIR
// metadata attribute and is not required in the function definition.
absl::flat_hash_set<absl::string_view> attrs_to_ignore = {
grad_string.data(), stateful_string.data(), kEntryFuncAttr};
llvm::SmallVector<mlir::NamedAttribute, 8> funcAttrs(
function->getDialectAttrs());
TF_RETURN_IF_ERROR(ConvertAttributes(funcAttrs, attrs_to_ignore,
/*remove_ref_type=*/false,
func_def.mutable_attr()));
for (int i = 0, e = function.getNumArguments(); i < e; ++i) {
if (auto resource_arg_unique_id_attr =
function.getArgAttrOfType<mlir::IntegerAttr>(
i, kResourceArgUniqueIdAttr)) {
(*func_def.mutable_resource_arg_unique_id())[i] =
resource_arg_unique_id_attr.getInt();
}
}
return flib_def->AddFunctionDef(std::move(func_def));
}
absl::Status Exporter::Convert(mlir::ModuleOp module,
const GraphExportConfig& configs,
std::unique_ptr<Graph>* graph,
FunctionLibraryDefinition* flib_def,
absl::flat_hash_set<Node*>* control_ret_nodes) {
mlir::StringAttr entry_func_id =
mlir::StringAttr::get(module.getContext(), "main");
std::optional<FuncOp> entry_func;
FunctionLibraryDefinition temp_flib_def(OpRegistry::Global(),
FunctionDefLibrary());
llvm::SmallDenseSet<FuncOp> visited_functions;
auto tf_dialect = module.getContext()->getLoadedDialect("tf");
// Construct SymbolTable to enable cheap function lookups. The cost
// of constructing the table is offset by the number of queries.
SymbolTable symbol_table(module);
for (auto function : module.getOps<FuncOp>()) {
if (function.isExternal())
return absl::FailedPreconditionError("External functions not supported");
if (function.getName() == entry_func_id &&
!configs.export_entry_func_to_flib) {
entry_func.emplace(function);
} else {
TF_RETURN_IF_ERROR(ConvertLibFunction(configs, tf_dialect, symbol_table,
function, &temp_flib_def,
visited_functions));
}
}
if (flib_def != nullptr) {
TF_RETURN_IF_ERROR(flib_def->AddLibrary(temp_flib_def));
}
if (!configs.export_entry_func_to_flib) {
if (!entry_func.has_value())
return absl::FailedPreconditionError(
"entry function `main` must be present");
// Updates the graph and the function library definition.
TF_ASSIGN_OR_RETURN(
*graph, Exporter::Convert(configs, tf_dialect, symbol_table,
entry_func.value(), &temp_flib_def,
visited_functions, control_ret_nodes));
// Add FunctionDefs and GradientDefs of MLIR functions to graph's function
// library. If duplicate FunctionDefs already exist (can happen if exporter
// had already added some FunctionDefs to the library to support legacy
// calls), they are ignored.
TF_RETURN_IF_ERROR(
graph->get()->mutable_flib_def()->AddLibrary(temp_flib_def));
} else if (graph != nullptr) {
TF_RETURN_IF_ERROR(
graph->get()->mutable_flib_def()->AddLibrary(std::move(*flib_def)));
}
return absl::OkStatus();
}
} // namespace
absl::Status ConvertTfExecutorToGraph(
mlir::ModuleOp module, const GraphExportConfig& configs,
std::unique_ptr<Graph>* graph, FunctionLibraryDefinition* flib_def,
absl::flat_hash_set<Node*>* control_ret_nodes) {
mlir::StatusScopedDiagnosticHandler sh(module.getContext());
if (failed(VerifyExportSuitable(module))) return sh.ConsumeStatus();
return sh.Combine(
Exporter::Convert(module, configs, graph, flib_def, control_ret_nodes));
}
absl::Status ConvertMlirFunctionToFunctionLibraryDef(
FuncOp func, const GraphExportConfig& configs, FunctionDef* function_def) {
Dialect* tf_dialect = func.getContext()->getLoadedDialect("tf");
FunctionLibraryDefinition flib_def(OpRegistry::Global(),
FunctionDefLibrary());
llvm::SmallDenseSet<FuncOp> visited_functions;
// Construct SymbolTable to enable cheap function lookups. The cost
// of constructing the table is offset by the number of queries. Even
// though this only converts one function in theory, this function
// may have gradient associated which would result in a lookup. This
// could be made lazy if we find this to be broad.
SymbolTable symbol_table(func->getParentOfType<mlir::ModuleOp>());
TF_RETURN_IF_ERROR(Exporter::ConvertLibFunction(
configs, tf_dialect, symbol_table, func, &flib_def, visited_functions));
auto name = FindFunctionName(configs, func);
const FunctionDef* func_def = flib_def.Find(name);
if (func_def != nullptr) {
*function_def = *func_def;
return absl::OkStatus();
}
return absl::InvalidArgumentError(
absl::StrCat("Function '", name,
"' couldn't be found in the FunctionDefLibrary after "
"converting from MLIR"));
}
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/api/v2/tf_executor_to_graph.h
+54
View File
@@ -0,0 +1,54 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_TF_EXECUTOR_TO_GRAPH_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_TF_EXECUTOR_TO_GRAPH_H_
#include "absl/base/attributes.h"
#include "absl/container/flat_hash_set.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/graph/graph.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
// Converts an MLIR module to TensorFlow graph and FunctionLibraryDefinition.
// The "main" function of the module is stored in the graph and the rest of
// functions are stored in the library. Control ret nodes are stored separately
// in `control_ret_nodes`.
absl::Status ConvertTfExecutorToGraph(
mlir::ModuleOp module, const GraphExportConfig& configs,
std::unique_ptr<Graph>* graph, FunctionLibraryDefinition* flib_def,
absl::flat_hash_set<Node*>* control_ret_nodes);
// Converts an MLIR function and adds it to a FunctionLibraryDefinition.
absl::Status ConvertMlirFunctionToFunctionLibraryDef(
mlir::func::FuncOp func, const GraphExportConfig& configs,
FunctionDef* function_def);
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_TF_EXECUTOR_TO_GRAPH_H_
tensorflow/compiler/mlir/tf2xla/api/v2/tf_executor_to_graph_test.cc
+110
View File
@@ -0,0 +1,110 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/api/v2/tf_executor_to_graph.h"
#include <stdlib.h>
#include <memory>
#include <string>
#include <gtest/gtest.h>
#include "absl/container/flat_hash_set.h"
#include "absl/status/statusor.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/DialectRegistry.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "mlir/Parser/Parser.h" // from @llvm-project
#include "riegeli/bytes/fd_reader.h" // from @riegeli
#include "riegeli/bytes/read_all.h" // from @riegeli
#include "tensorflow/compiler/mlir/register_common_dialects.h"
#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/platform/resource_loader.h"
#include "tsl/platform/protobuf.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
namespace {
using mlir::DialectRegistry;
using mlir::MLIRContext;
using mlir::ModuleOp;
using mlir::OwningOpRef;
std::string TestDataPath() {
return tensorflow::GetDataDependencyFilepath(
"tensorflow/compiler/mlir/tf2xla/api/v2/testdata/");
}
class TfExecutorToGraphTest : public ::testing::Test {
public:
TfExecutorToGraphTest() {
mlir::RegisterCommonToolingDialects(registry_);
context_.appendDialectRegistry(registry_);
context_.loadAllAvailableDialects();
}
absl::StatusOr<OwningOpRef<mlir::ModuleOp>> CreateMlirModule(
std::string mlir_module_filename) {
std::string mlir_module_path = TestDataPath() + mlir_module_filename;
return mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context_);
}
GraphDef CreateGraphDef(std::string graphdef_filename) {
std::string file_path = TestDataPath() + graphdef_filename;
std::string contents;
GraphDef graph_def;
auto status = riegeli::ReadAll(riegeli::FdReader(file_path), contents);
if (!status.ok()) {
return graph_def;
}
tsl::protobuf::TextFormat::ParseFromString(contents, &graph_def);
return graph_def;
}
DialectRegistry registry_;
MLIRContext context_;
OwningOpRef<mlir::ModuleOp> mlir_module_;
};
TEST_F(TfExecutorToGraphTest, ConvertMlirToGraphSucceeds) {
auto valid_executor_module = CreateMlirModule("valid_executor.mlir");
GraphExportConfig confs;
absl::flat_hash_set<Node*> control_ret_nodes;
FunctionLibraryDefinition flib_def(OpRegistry::Global(),
FunctionDefLibrary());
auto result_graph = std::make_unique<Graph>(flib_def);
TF_ASSERT_OK(ConvertTfExecutorToGraph(valid_executor_module.value().get(),
confs, &result_graph, &flib_def,
&control_ret_nodes));
GraphDef result_graphdef;
result_graph->ToGraphDef(&result_graphdef);
GraphDef expected_graphdef = CreateGraphDef("valid_graph.txt");
EXPECT_EQ(result_graphdef.DebugString(), expected_graphdef.DebugString());
}
} // namespace
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow