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/internal/BUILD
+429
View File
@@ -0,0 +1,429 @@
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load(
"//tensorflow/core/platform:build_config.bzl",
"tf_proto_library",
)
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow/compiler/mlir/tensorflow/transforms:__pkg__",
"//tensorflow/compiler/mlir/tf2xla/api/v1:__subpackages__",
"//tensorflow/compiler/mlir/tf2xla/api/v2:__subpackages__",
"//tensorflow/core/tfrt/saved_model:__pkg__",
],
)
cc_library(
name = "compilation_timer",
hdrs = ["compilation_timer.h"],
deps = [
"//tensorflow/core/platform:profile_utils_cpu_utils",
],
)
tf_cc_test(
name = "compilation_timer_test",
srcs = ["compilation_timer_test.cc"],
deps = [
":compilation_timer",
"@com_google_absl//absl/time",
"@com_google_googletest//:gtest_main",
],
)
cc_library(
name = "test_matchers",
testonly = True,
hdrs = ["test_matchers.h"],
deps = [
"//tensorflow/compiler/mlir/tf2xla:compile_mlir_util",
"@com_google_absl//absl/status:statusor",
"@com_google_googletest//:gtest",
"@xla//xla/tsl/platform:statusor",
],
)
tf_cc_test(
name = "test_matchers_test",
srcs = ["test_matchers_test.cc"],
deps = [
":test_matchers",
"//tensorflow/compiler/mlir/tf2xla:compile_mlir_util",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/core:lib",
"//tensorflow/core/lib/monitoring:cell_reader",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_googletest//:gtest_main",
"@xla//xla/hlo/builder:xla_computation",
"@xla//xla/service:hlo_proto_cc",
"@xla//xla/tsl/platform:statusor",
],
)
cc_library(
name = "mlir_pass_instrumentation",
srcs = ["mlir_pass_instrumentation.cc"],
hdrs = ["mlir_pass_instrumentation.h"],
deps = [
"//tensorflow/core/platform:logging",
"@com_google_absl//absl/log",
"@llvm-project//mlir:Pass",
],
)
tf_cc_test(
name = "mlir_pass_instrumentation_test",
srcs = ["mlir_pass_instrumentation_test.cc"],
deps = [
":mlir_pass_instrumentation",
"//tensorflow/compiler/mlir/tf2xla/api/v1:compile_mlir_util_no_tf_dialect_passes",
"//tensorflow/core:test",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:string_view",
"@com_google_googletest//:gtest_main",
],
)
cc_library(
name = "legalize_tf_mlir",
srcs = ["legalize_tf_mlir.cc"],
hdrs = ["legalize_tf_mlir.h"],
deps = [
":compilation_timer",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:serialize_mlir_module_utils",
"//tensorflow/compiler/mlir/tensorflow/transforms:set_tpu_infeed_layout",
"//tensorflow/compiler/mlir/tf2xla:compile_mlir_util",
"//tensorflow/compiler/tf2xla:layout_util",
"//tensorflow/compiler/tf2xla:xla_compiler",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
"//tensorflow/core/tpu:tpu_compile",
"//tensorflow/core/tpu/kernels:tpu_compile_op_support",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@stablehlo//:register",
"@tsl//tsl/platform:error_logging",
"@xla//xla:shape_util",
"@xla//xla/mlir_hlo:hlo_dialect_registration",
"@xla//xla/tsl/platform:statusor",
],
)
tf_cc_test(
name = "legalize_tf_mlir_test",
srcs = ["legalize_tf_mlir_test.cc"],
deps = [
":legalize_tf_mlir",
":test_matchers",
"//tensorflow/compiler/jit",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:deserialize_mlir_module_utils",
"//tensorflow/compiler/tf2xla:xla_compiler",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/core:framework",
"//tensorflow/core:test_main",
"//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
"//tensorflow/core/tpu/kernels:tpu_compile_op_support",
"@com_google_absl//absl/status:statusor",
"@com_google_googletest//:gtest",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@stablehlo//:register",
"@xla//xla:shape_util",
"@xla//xla/mlir_hlo:hlo_dialect_registration",
"@xla//xla/tsl/platform:errors",
],
)
cc_library(
name = "legalize_tf_to_hlo",
srcs = ["legalize_tf_to_hlo.cc"],
hdrs = ["legalize_tf_to_hlo.h"],
deps = [
":legalize_tf_mlir",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:deserialize_mlir_module_utils",
"//tensorflow/compiler/mlir/tensorflow/transforms:set_tpu_infeed_layout",
"//tensorflow/compiler/mlir/tf2xla/api/v1:compile_tf_graph",
"//tensorflow/compiler/tf2xla:layout_util",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/compiler/tf2xla:xla_tpu_backend_registration",
"//tensorflow/core:framework",
"//tensorflow/core/platform:status",
"//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
"//tensorflow/core/tpu/kernels:tpu_compile_op_support",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:DataLayoutInterfaces",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@stablehlo//:register",
"@xla//xla:shape_util",
"@xla//xla/client:compile_only_client",
"@xla//xla/mlir_hlo:hlo_dialect_registration",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:statusor",
],
)
tf_cc_test(
name = "legalize_tf_to_hlo_test",
srcs = ["legalize_tf_to_hlo_test.cc"],
deps = [
":legalize_tf_to_hlo",
":test_matchers",
"//tensorflow/compiler/jit",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:deserialize_mlir_module_utils",
"//tensorflow/compiler/tf2xla:xla_compiler",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test_main",
"//tensorflow/core/lib/monitoring:cell_reader",
"//tensorflow/core/protobuf:for_core_protos_cc",
"//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
"//tensorflow/core/tpu/kernels:tpu_compile_op_support",
"@com_google_absl//absl/status:statusor",
"@com_google_googletest//:gtest",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@stablehlo//:register",
"@xla//xla:shape_util",
"@xla//xla/client:client_library",
"@xla//xla/mlir_hlo:hlo_dialect_registration",
"@xla//xla/stream_executor:platform",
"@xla//xla/stream_executor:platform_manager",
"@xla//xla/tsl/platform:errors",
],
)
cc_library(
name = "clustering_bridge_passes",
srcs = ["clustering_bridge_passes.cc"],
hdrs = ["clustering_bridge_passes.h"],
deps = [
"//tensorflow/compiler/jit:flags_headers",
"//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
"//tensorflow/compiler/mlir/tensorflow/transforms:verify_no_outside_compilation_markers_pass",
"//tensorflow/compiler/mlir/tensorflow/transforms/sparsecore:sparsecore_passes",
"//tensorflow/compiler/mlir/tf2xla/internal/passes:clustering_passes",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/log",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Transforms",
],
)
tf_cc_test(
name = "clustering_bridge_passes_test",
srcs = ["clustering_bridge_passes_test.cc"],
deps = [
":clustering_bridge_passes",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:Pass",
],
)
cc_library(
name = "logging_hooks",
srcs = ["logging_hooks.cc"],
hdrs = ["logging_hooks.h"],
deps = [
"//tensorflow/compiler/mlir/tensorflow:bridge_logger",
"//tensorflow/core:framework",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:Pass",
],
)
tf_cc_test(
name = "logging_hooks_test",
srcs = ["logging_hooks_test.cc"],
data = [
"testdata/dead_const.mlir",
],
deps = [
":logging_hooks",
"//tensorflow/compiler/mlir:register_common_dialects",
"//tensorflow/core:lib",
"//tensorflow/core:test",
"//tensorflow/core/platform:resource_loader",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Transforms",
"@xla//xla/tsl/lib/core:status_test_util",
"@xla//xla/tsl/platform:status",
],
)
cc_library(
name = "mlir_bridge_pass_util",
srcs = ["mlir_bridge_pass_util.cc"],
hdrs = ["mlir_bridge_pass_util.h"],
visibility = ["//tensorflow/compiler/tf2xla:__pkg__"],
deps = [
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/tf2xla:tf2xla_defs",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/common_runtime:function_body",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:string_view",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
"@xla//xla/tsl/platform:status",
],
alwayslink = 1,
)
tf_cc_test(
name = "mlir_bridge_pass_util_test",
srcs = ["mlir_bridge_pass_util_test.cc"],
deps = [
":mlir_bridge_pass_util",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"//tensorflow/cc:array_ops",
"//tensorflow/cc:function_ops",
"//tensorflow/cc:functional_ops",
"//tensorflow/cc:ops",
"//tensorflow/cc:scope",
"//tensorflow/cc:tpu_ops",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/tf2xla:tf2xla_defs",
"//tensorflow/core/protobuf:for_core_protos_cc",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core/platform:types",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test_main",
"//tensorflow/core:testlib",
"//tensorflow/core/platform:enable_tf2_utils",
# "//tensorflow/core/platform:resource_loader",
"@xla//xla/tsl/lib/core:status_test_util",
],
)
tf_proto_library(
name = "reproducer_proto",
srcs = ["reproducer.proto"],
protodeps = [
"//tensorflow/core:protos_all",
"//tensorflow/core/protobuf/tpu:compile_metadata_proto",
],
visibility = [
"//learning/brain/mlir/bridge:__pkg__",
"//tensorflow/compiler/mlir/tf2xla/api/v2:__pkg__",
],
)
cc_library(
name = "node_order",
srcs = ["node_order.cc"],
hdrs = ["node_order.h"],
deps = [
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:lib",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
],
)
tf_cc_test(
name = "node_order_test",
size = "small",
srcs = [
"node_order_test.cc",
],
deps = [
":node_order",
"//tensorflow/cc:cc_ops",
"//tensorflow/cc:cc_ops_internal",
"//tensorflow/cc:function_ops",
"//tensorflow/cc:sendrecv_ops",
"//tensorflow/core",
"//tensorflow/core:core_cpu",
"//tensorflow/core:core_cpu_internal",
"//tensorflow/core:direct_session_internal",
"//tensorflow/core:framework",
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:ops",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core:testlib",
"@com_google_absl//absl/log:check",
"@com_google_googletest//:gtest",
],
)
cc_library(
name = "graph_to_tf_executor_util",
srcs = ["graph_to_tf_executor_util.cc"],
hdrs = ["graph_to_tf_executor_util.h"],
deps = [
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:framework_types_hdr",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/common_runtime:function_body",
"//tensorflow/core/platform:enable_tf2_utils",
"//tensorflow/core/protobuf:for_core_protos_cc",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:string_view",
],
)
tf_cc_test(
name = "graph_to_tf_executor_util_test",
srcs = ["graph_to_tf_executor_util_test.cc"],
deps = [
":graph_to_tf_executor_util",
"//tensorflow/cc:array_ops",
"//tensorflow/cc:cc_ops",
"//tensorflow/cc:functional_ops",
"//tensorflow/cc:ops",
"//tensorflow/cc:scope",
"//tensorflow/cc:tpu_ops",
"//tensorflow/compiler/tf2xla/ops:xla_ops",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/framework:tensor_testutil",
"//tensorflow/core/platform:enable_tf2_utils",
"//tensorflow/core/platform:types",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest_main",
"@xla//xla/tsl/lib/core:status_test_util",
],
)
tensorflow/compiler/mlir/tf2xla/internal/clustering_bridge_passes.cc
+220
View File
@@ -0,0 +1,220 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/internal/clustering_bridge_passes.h"
#include <string>
#include "absl/log/log.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "mlir/Transforms/Passes.h" // from @llvm-project
#include "tensorflow/compiler/jit/flags.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/sparsecore/sparsecore_passes.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
using mlir::OpPassManager;
using mlir::func::FuncOp;
// Adds replicated Bridge clustering pipeline passes to the given pass_manager.
// Does not run them.
void AddReplicatedBridgeClusteringPipelinePasses(OpPassManager& pm,
llvm::StringRef module_name) {
// The following ops must be preserved regardless of reachability. Ideally,
// all graphs should have control dependencies to enforce this but this is
// currently not the case (see b/177478741).
const llvm::SmallVector<std::string, 4> ops_to_preserve = {
"tf.TPUReplicateMetadata", "tf.TPUCompilationResult",
"tf.TPUReplicatedOutput"};
bool strict_clusters =
tensorflow::GetMlirCommonFlags()->tf_mlir_enable_strict_clusters;
pm.addNestedPass<FuncOp>(
mlir::tf_executor::CreateTFExecutorGraphPruningPass(ops_to_preserve));
// It is assumed at this stage there are no V1 control flow ops as Graph
// functionalization is ran before import. Ops can be lifted out of
// tf_executor dialect islands/graphs.
pm.addNestedPass<FuncOp>(
mlir::CreateExecutorDialectToFunctionalConversionPass());
// Guarantee all functions have one use, which enables more exact shape
// inference.
pm.addPass(mlir::TF::CreateGuaranteeAllFuncsOneUsePass());
// Run shape inference so that tf_executor/tf_device ops created later will
// likely to inherit more concrete types.
pm.addPass(mlir::TF::CreateTFShapeInferencePass());
pm.addNestedPass<FuncOp>(mlir::TFTPU::CreateTPUPartitionedOpConversionPass());
pm.addNestedPass<FuncOp>(
mlir::TFTPU::CreateTPUReorderReplicateAndPartitionedInputsPass());
pm.addNestedPass<FuncOp>(mlir::TF::CreateDecomposeReduceDatasetPass());
// Only one of EmbeddingSequencing and EmbeddingPipelining will actually
// run and the logic is in EmbeddingPipeliningPass. If the pipelining pass
// runs, embedding attributes are stripped and the sequencing pass will have
// no effect. If the pipelining pass doesn't run, embedding attributes are
// preserved and the sequencing rewrite will trigger.
pm.addPass(mlir::TFDevice::CreateEmbeddingPipeliningPass());
pm.addPass(mlir::TFDevice::CreateEmbeddingSequencingPass());
pm.addPass(tensorflow::tf2xla::internal::CreateTPUClusterFormationPass(
strict_clusters));
// CreateEmbeddingPipeliningPass may have created more functions, but
// TPUClusterCleanup and OutsideCompiledToHostLaunch need every function to be
// only called from one cluster. Here, we choose to fix the all-funcs-one-use
// invariant right before it's needed, not after it's been broken.
pm.addPass(mlir::TF::CreateGuaranteeAllFuncsOneUsePass());
// Run TPU cluster cleanup attributes so ops with no outside compiled
// attribute have no host device attribute.
pm.addPass(mlir::TFTPU::CreateTPUClusterCleanupAttributesPass());
pm.addNestedPass<FuncOp>(mlir::TFDevice::CreateDeviceAttributeToLaunchPass());
// Running canonicalizer before decomposing resource ops in cluster helps the
// latter pass to converge faster as it does not have to spend time folding
// away dead ops.
pm.addNestedPass<FuncOp>(mlir::createCanonicalizerPass());
// Place DecomposeResourceOpsPass before TFExecutorConstantSinking pass
// because DecomposeResourceOpsPass uses pattern rewriter which hoists
// changed constants out of tf_device.Launch.
pm.addPass(mlir::TFDevice::CreateDecomposeResourceOpsInClusterPass());
// Encode this in its own scope so that func_pm is not mistakenly used
// later on.
{
OpPassManager& func_pm = pm.nest<FuncOp>();
func_pm.addPass(mlir::TFTPU::CreateTPUHostComputationExpansionPass());
func_pm.addPass(mlir::TFTPU::CreateTPUUpdateEmbeddingEnqueueOpInputsPass());
}
// TODO(b/173622615): Once OutsideCompilation is represented by launch op and
// the remaining passes including Inliner support it, remove this
// LaunchToDeviceAttributePass. This LaunchToDeviceAttribute pass needs to
// come before TPUClusterCleanupAttributes pass or else the device attribute
// will be removed from launch causing an error.
pm.addNestedPass<FuncOp>(mlir::TFDevice::CreateLaunchToDeviceAttributePass());
// Note that the region-based control-flow produced here still contains
// function call ops which get inlined by the subsequent inliner pass.
pm.addPass(mlir::TF::CreateTFFunctionalControlFlowToRegions());
pm.addPass(mlir::createInlinerPass());
pm.addNestedPass<FuncOp>(
mlir::TF::CreateDropWhileShapeInvariantInDeviceClusterPass());
// Run another shape inference pass because resource decomposition might have
// created new partial types. Also, after dropping `shape_invariant` attribute
// from While/WhileRegion ops within cluster would lead to more precise
// shapes.
pm.addPass(mlir::TF::CreateTFShapeInferencePass());
pm.addNestedPass<FuncOp>(mlir::createCanonicalizerPass());
pm.addPass(mlir::TFTPU::CreateTPUClusterCleanupAttributesPass());
pm.addPass(mlir::TFDevice::CreateResourceOpLiftingPass());
// Re-run the canonicalizer pass as some cleanup during resource op lifting
// pass opens up some opportunities for canonicalization of cluster ops.
// Specifically, we want to eliminate pass through results from the cluster
// op.
pm.addNestedPass<FuncOp>(mlir::createCanonicalizerPass());
pm.addNestedPass<FuncOp>(mlir::createCSEPass());
if (tensorflow::GetMlirCommonFlags()
->tf_mlir_enable_merge_control_flow_pass) {
pm.addPass(mlir::TFDevice::CreateMergeControlFlowPass());
}
pm.addPass(
tensorflow::tf2xla::internal::CreateMarkOpsForOutsideCompilationPass());
pm.addPass(tensorflow::tf2xla::internal::
CreateExtractHeadTailOutsideCompilationPass());
pm.addPass(
tensorflow::tf2xla::internal::CreateExtractOutsideCompilationPass());
pm.addNestedPass<FuncOp>(
mlir::TFDevice::CreateVerifyNoOutsideCompilationMarkersPass());
pm.addNestedPass<FuncOp>(mlir::TFDevice::CreateClusterConstantSinkingPass());
pm.addPass(mlir::TF::CreateResourceDeviceInferencePass());
pm.addNestedPass<FuncOp>(
tensorflow::tf2xla::internal::CreateHoistBroadcastReadPass());
pm.addNestedPass<FuncOp>(
tensorflow::tf2xla::internal::CreateXlaBroadcastPass());
pm.addPass(mlir::TFDevice::CreateClusterOutliningPass());
pm.addPass(mlir::TFTPU::CreateTPUResourceReadForWritePass());
pm.addPass(mlir::TFDevice::CreateMarkInputOutputAliasesPass());
pm.addPass(
tensorflow::tf2xla::internal::CreateTPUShardingIdentificationPass());
pm.addNestedPass<FuncOp>(
mlir::TFTPU::CreateTPUResourceReadsWritesPartitioningPass());
pm.addPass(mlir::TFDevice::CreateAnnotateParameterReplicationPass());
pm.addNestedPass<FuncOp>(mlir::TF::CreateRewriteTPUEmbeddingOpsPass());
pm.addPass(mlir::TFTPU::CreateTPUAnnotateDynamicShapeInputsPass());
pm.addNestedPass<FuncOp>(
mlir::TF::CreateHoistReplicateInvariantResourceWritesPass());
// Verifies clustering has conformed with the expected invariants
pm.addNestedPass<FuncOp>(
tensorflow::tf2xla::internal::CreateVerifyClusteringPass());
}
void NoCanonicalization(OpPassManager& pm) {}
// Same as above but for non-replicated Bridge.
void AddNonReplicatedBridgeClusteringPipelinePasses(OpPassManager& pm) {
// The following ops must be preserved regardless of reachability. Ideally,
// all graphs should have control dependencies to enforce this.
VLOG(2) << "Create TF XLA Bridge pipeline";
pm.addPass(mlir::TFDevice::CreateXlaValidateInputsPass());
pm.addNestedPass<FuncOp>(
mlir::TF::CreateCanonicalizeCompileAndReplicateAttributesPass());
const llvm::SmallVector<std::string, 4> ops_to_preserve = {};
pm.addNestedPass<FuncOp>(
mlir::tf_executor::CreateTFExecutorGraphPruningPass(ops_to_preserve));
// It is assumed at this stage there are no V1 control flow ops as Graph
// functionalization is ran before import. Ops can be lifted out of
// tf_executor dialect islands/graphs.
pm.addNestedPass<FuncOp>(
mlir::CreateExecutorDialectToFunctionalConversionPass());
// Guarantee all functions have one use, which enables more exact shape
// inference.
pm.addPass(mlir::TF::CreateGuaranteeAllFuncsOneUsePass());
pm.addPass(mlir::TF::CreateTFShapeInferencePass());
// Encapsulate PartitionedCall ops within a cluster so that the composite
// resource ops can be decomposed.
pm.addPass(tensorflow::tf2xla::internal::CreateXlaClusterFormationPass());
// Running canonicalizer before decomposing resource ops in cluster helps the
// latter pass to converge faster as it does not have to spend time folding
// away dead ops.
pm.addNestedPass<FuncOp>(mlir::createCanonicalizerPass());
// Decompose resource ops.
pm.addPass(mlir::TFDevice::CreateDecomposeResourceOpsInClusterPass());
// Run another shape inference pass because resource decomposition might have
// created new partial types. Also, after dropping `shape_invariant` attribute
// from While/WhileRegion ops within cluster would lead to more precise
// shapes.
pm.addPass(mlir::TF::CreateTFShapeInferencePass());
pm.addNestedPass<FuncOp>(mlir::createCanonicalizerPass());
// Inline all the function calls. Do not call canonicalizer to prevent it from
// moving the definition of any constant operand of ops within a cluster to
// its outside. This may cause the op to fail to verify after the cluster is
// outlined, as the constant operand is replaced by an argument.
pm.addPass(mlir::createInlinerPass({}, NoCanonicalization));
// Lift resource operations out of device computation. This step needs to be
// done after inlining.
pm.addPass(mlir::TFDevice::CreateResourceOpLiftingPass());
// Outline clusters into cluster functions.
pm.addPass(mlir::TFDevice::CreateClusterOutliningPass());
// Verifies clustering has conformed with the expected invariants
pm.addNestedPass<FuncOp>(
tensorflow::tf2xla::internal::CreateVerifyClusteringPass());
}
}; // namespace internal
}; // namespace tf2xla
}; // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/clustering_bridge_passes.h
+39
View File
@@ -0,0 +1,39 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_CLUSTERING_BRIDGE_PASSES_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_CLUSTERING_BRIDGE_PASSES_H_
#include "absl/base/attributes.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Pass/PassManager.h" // from @llvm-project
namespace tensorflow {
namespace tf2xla {
namespace internal {
// Given the pass manager, add Bridge passes to cluster the replicated input
// graphs.
void AddReplicatedBridgeClusteringPipelinePasses(
mlir::OpPassManager& pm, llvm::StringRef module_name = llvm::StringRef());
// Same as above but for non replicated graphs.
void AddNonReplicatedBridgeClusteringPipelinePasses(mlir::OpPassManager& pm);
}; // namespace internal
}; // namespace tf2xla
}; // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_CLUSTERING_BRIDGE_PASSES_H_
tensorflow/compiler/mlir/tf2xla/internal/clustering_bridge_passes_test.cc
+43
View File
@@ -0,0 +1,43 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/internal/clustering_bridge_passes.h"
#include <gtest/gtest.h>
#include "mlir/Pass/PassManager.h" // from @llvm-project
namespace tensorflow {
namespace tf2xla {
namespace internal {
using mlir::OpPassManager;
TEST(ClusteringBridgePassesTest, AddsBridgePasses) {
OpPassManager pass_manager;
AddReplicatedBridgeClusteringPipelinePasses(pass_manager);
EXPECT_EQ(pass_manager.size(), 45);
}
TEST(ClusteringBridgePassesTest, AddsNonTPUBridgePasses) {
OpPassManager pass_manager;
AddNonReplicatedBridgeClusteringPipelinePasses(pass_manager);
EXPECT_EQ(pass_manager.size(), 15);
}
}; // namespace internal
}; // namespace tf2xla
}; // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/compilation_timer.h
+43
View File
@@ -0,0 +1,43 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_COMPILATION_TIMER_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_COMPILATION_TIMER_H_
#include <chrono> // NOLINT(build/c++11)
#include "tensorflow/core/platform/profile_utils/cpu_utils.h"
// Time the execution of kernels (in CPU cycles). Meant to be used as RAII.
struct CompilationTimer {
uint64_t start_cycles =
tensorflow::profile_utils::CpuUtils::GetCurrentClockCycle();
uint64_t ElapsedCycles() {
return tensorflow::profile_utils::CpuUtils::GetCurrentClockCycle() -
start_cycles;
}
int64_t ElapsedCyclesInMilliseconds() {
std::chrono::duration<double> duration =
tensorflow::profile_utils::CpuUtils::ConvertClockCycleToTime(
ElapsedCycles());
return std::chrono::duration_cast<std::chrono::milliseconds>(duration)
.count();
}
};
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_COMPILATION_TIMER_H_
tensorflow/compiler/mlir/tf2xla/internal/compilation_timer_test.cc
+41
View File
@@ -0,0 +1,41 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/internal/compilation_timer.h"
#include <cstdint>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/time/clock.h"
#include "absl/time/time.h"
namespace {
TEST(CompilationTimer, MeasuresElapsedTime) {
uint64_t timer_result_in_milliseconds;
{
CompilationTimer timer;
absl::SleepFor(absl::Milliseconds(100));
timer_result_in_milliseconds = timer.ElapsedCyclesInMilliseconds();
}
ASSERT_THAT(timer_result_in_milliseconds, testing::Ne(0));
}
} // namespace
tensorflow/compiler/mlir/tf2xla/internal/graph_to_tf_executor_util.cc
+332
View File
@@ -0,0 +1,332 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/internal/graph_to_tf_executor_util.h"
#include <memory>
#include <optional>
#include <string>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "xla/tsl/platform/errors.h"
#include "tensorflow/core/common_runtime/function_body.h"
#include "tensorflow/core/common_runtime/function_def_utils.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/platform/enable_tf2_utils.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/util/device_name_utils.h"
namespace tensorflow {
namespace {
// Internal encapsulation of state for the MLIR bridge graph analyzer. Steps
// through the nodes in the graph and reachable functions, tracking whether
// each feature of interest is found.
//
// Tracks the presence of each feature of interest in the corresponding streamz
// metric. Note that the graph traversal does not terminate early so as to
// capture all of these features.
class MlirBridgeGraphAnalyzer {
public:
explicit MlirBridgeGraphAnalyzer(bool single_core_inference_mode)
: single_core_inference_mode_(single_core_inference_mode) {}
~MlirBridgeGraphAnalyzer() = default;
// Not copyable or movable.
MlirBridgeGraphAnalyzer(const MlirBridgeGraphAnalyzer&) = delete;
MlirBridgeGraphAnalyzer& operator=(const MlirBridgeGraphAnalyzer&) = delete;
// Analyzes whether the graph has features not guaranteed to be supported by
// the MLIR-based TF XLA bridge.
bool HasUnsupportedFeatures(const Graph& graph,
const FunctionLibraryDefinition* function_library,
std::optional<ConfigProto> config_proto,
tensorflow::TF2XLABridgeVersion bridge_version) {
// Non-ok status is considered as "unsupported" since this means something
// is wrong or unexpected with the graph itself.
invalid_graph_ =
invalid_graph_ || !AnalyzeGraphAndReachableFunctions(
graph, function_library, config_proto)
.ok();
// We conservatively consider the graph to be unsupported if it's not
// *known* to be TF2. That is, graphs that have kNotTracked construction
// context are considered unsupported, even though they might in fact be
// TF2 models.
auto construction_context = graph.GetConstructionContextInternal();
bool is_tf2 = construction_context == ConstructionContext::kEagerRuntime;
auto is_tf2_execution_enabled = tensorflow::tf2_execution_enabled();
auto has_unsupported_features = false;
auto is_v1_compat = bridge_version == TF2XLABridgeVersion::kV1Compat;
auto is_nominal_bridge = bridge_version == TF2XLABridgeVersion::kNominal;
auto is_tfrt_bridge = bridge_version == TF2XLABridgeVersion::kTFRTNominal;
is_eager_compliant_ = is_tf2_execution_enabled || is_tf2 ||
is_nominal_bridge || is_tfrt_bridge;
is_eager_compliant_ |= (is_v1_compat && contains_partitioned_call_);
has_unsupported_features = contains_ref_type_ || invalid_graph_;
// For non single core inference mode, checking conditions:
if (!single_core_inference_mode_) {
has_unsupported_features |=
!is_eager_compliant_ || uses_v1_control_flow_ ||
HasTpuReplicatedCoreUnsupportedFeature(is_nominal_bridge,
is_v1_compat, is_tfrt_bridge);
}
PrintGraphUnsupportedFeatures(is_tf2, is_tf2_execution_enabled,
is_v1_compat, is_tfrt_bridge,
is_nominal_bridge, has_unsupported_features);
// Determine whether or not the graph contains unsupported features.
return has_unsupported_features;
}
private:
static constexpr char kPartitionedCall[] = "TPUPartitionedCall";
bool HasTPUReplicatedCoreAttr(const Node& node) {
constexpr absl::string_view kTPUReplicatedCore = "TPU_REPLICATED_CORE";
const std::string& device = node.requested_device();
if (!device.empty()) {
DeviceNameUtils::ParsedName name;
if (DeviceNameUtils::ParseFullName(device, &name)) {
// The TPU_REPLICATED_CORE attrs is not relevant for single TPU core
// inference.
// TODO(b/201091475): this can be generalized to check
// num_cores_per_replica != 1, rather than being special cased for
// single core inference.
if (name.type == kTPUReplicatedCore && !single_core_inference_mode_) {
return true;
}
}
}
return false;
}
bool HasTpuReplicatedCoreUnsupportedFeature(bool is_nominal_bridge,
bool is_v1_compat,
bool is_tfrt_bridge) {
if (!has_tpu_replicated_core_) {
return false;
}
return has_infeed_dequeue_tuple_with_tpu_replicated_core_;
}
void PrintGraphUnsupportedFeatures(bool is_tf2, bool is_tf2_execution_enabled,
bool is_v1_compat, bool is_tfrt_bridge,
bool is_nominal_bridge,
bool has_unsupported_features) {
if (!has_unsupported_features) {
VLOG(1) << "Graph doesn't have unsupported features";
return;
}
LOG(INFO)
<< "Graph has unsupported features: " << (is_tf2 ? "" : "not is_tf2, ")
<< (is_tf2_execution_enabled ? "" : "not tf2_execution, ")
<< (is_nominal_bridge ? "" : "not nominal bridge, ")
<< (is_tfrt_bridge ? "" : "not tfrt bridge, ")
<< (is_v1_compat && contains_partitioned_call_
? "contains partitioned calls at v1 compat bridge call site, "
: "")
<< (contains_ref_type_ ? "contains ref variables, " : "")
<< (invalid_graph_ ? "Invalid graph, " : "")
<< (uses_v1_control_flow_ ? "uses control flow v1 " : "")
<< ((has_tpu_replicated_core_ &&
has_infeed_dequeue_tuple_with_tpu_replicated_core_)
? "InfeedDequeueTuple op with TPU_REPLICATED_CORE attr, "
: "");
}
// Traverses each node in the graph and gathers information about each of the
// features. Specifically, sets the relevant class variable to true when a
// feature is found.
void AnalyzeGraphNodes(const Graph& graph) {
constexpr absl::string_view kIdentityOp = "Identity";
constexpr absl::string_view kIdentityNOp = "IdentityN";
constexpr absl::string_view kCastOp = "Cast";
constexpr absl::string_view kInfeedDequeueTuple = "InfeedDequeueTuple";
constexpr absl::string_view kOutsideCompilationAttr =
"_xla_outside_compilation";
constexpr absl::string_view kAllowSoftPlacementAttr =
"allow_soft_placement";
constexpr absl::string_view kManualControlDepsAttr =
"_has_manual_control_dependencies";
auto has_ref_type = [](const DataTypeVector& types) {
for (const DataType& dtype : types)
if (IsRefType(dtype)) return true;
return false;
};
for (const Node* node : graph.nodes()) {
contains_ref_type_ =
(contains_ref_type_ || has_ref_type(node->input_types()) ||
has_ref_type(node->output_types()));
contains_partitioned_call_ = (contains_partitioned_call_ ||
node->type_string() == kPartitionedCall);
uses_v1_control_flow_ = (uses_v1_control_flow_ || node->IsControlFlow());
uses_outside_compilation_ =
(uses_outside_compilation_ ||
node->attrs().Find(kOutsideCompilationAttr) != nullptr);
has_manual_control_deps_ = (has_manual_control_deps_ ||
node->attrs().Find(kManualControlDepsAttr));
auto soft_placement_attr = node->attrs().Find(kAllowSoftPlacementAttr);
if (soft_placement_attr != nullptr) {
uses_outside_compilation_ =
(uses_outside_compilation_ || soft_placement_attr->b());
}
// TODO(b/187611527): Add support for the ops with explicit device
// assignment on the TPU_REPLICATED_CORE.
if (node->type_string() == kIdentityOp ||
node->type_string() == kCastOp ||
node->type_string() == kIdentityNOp) {
if (HasTPUReplicatedCoreAttr(*node)) {
has_tpu_replicated_core_ = true;
VLOG(2) << node->type_string()
<< " node has TPU_REPLICATED_CORE attribute.";
}
}
if (node->type_string() == kInfeedDequeueTuple &&
HasTPUReplicatedCoreAttr(*node)) {
has_infeed_dequeue_tuple_with_tpu_replicated_core_ = true;
}
}
}
// Analyze all functions from the flib_def if there are any that belong to
// the inference graph.
void AnalyzeInferenceGraphs(const FunctionLibraryDefinition& flib_def) {
if (contains_partitioned_call_) return;
for (const std::string& func_name : flib_def.ListFunctionNames()) {
const FunctionDef* func_def = flib_def.Find(func_name);
for (const NodeDef& node_def : func_def->node_def()) {
contains_partitioned_call_ = node_def.op() == kPartitionedCall;
if (contains_partitioned_call_) return;
}
}
}
// Checks any reachable functions from `graph_def` in `flib_def`
// for unsupported features in the MLIR-based bridge.
//
// Returns failure in the event that the FunctionDef fails to convert to
// FunctionBody. Otherwise returns success.
absl::Status AnalyzeReachableFunctions(
const GraphDef& graph_def, const FunctionLibraryDefinition& flib_def) {
// Check the inputs and outputs of a function for reference variables.
auto signature_contains_ref_type = [](const OpDef& signature) {
for (const auto& args : {signature.input_arg(), signature.output_arg()}) {
for (const auto& arg : args) {
if (IsRefType(arg.type())) return true;
}
}
return false;
};
for (const std::string& func_name :
flib_def.ReachableDefinitions(graph_def).ListFunctionNames()) {
const FunctionDef* func_def = flib_def.Find(func_name);
if (func_def->has_signature()) {
contains_ref_type_ = contains_ref_type_ ||
signature_contains_ref_type(func_def->signature());
}
// Check the function body.
std::unique_ptr<FunctionBody> func_body;
TF_RETURN_IF_ERROR(FunctionDefToBodyHelper(
*func_def, AttrSlice(&func_def->attr()), &flib_def, &func_body));
AnalyzeGraphNodes(*func_body->graph);
}
return absl::OkStatus();
}
// Checks the inputted graph for any features which aren't supported in the
// MLIR-based bridge, stepping through each node in the graph as well as any
// reachable functions (inputs, outputs, and function body).
//
// Note that this analysis does not terminate early because we care about
// collecting all of these metrics.
//
// Returns failure in the event that the FunctionDef fails to convert to
// FunctionBody. Otherwise returns success.
absl::Status AnalyzeGraphAndReachableFunctions(
const Graph& graph, const FunctionLibraryDefinition* function_library,
std::optional<ConfigProto> config_proto) {
// First, check whether soft placement is enabled. This means that auto
// outside compilation may be used.
uses_outside_compilation_ =
uses_outside_compilation_ ||
(config_proto.has_value() && config_proto->allow_soft_placement());
// Analyze each node in this graph.
AnalyzeGraphNodes(graph);
// Then check any associated functions in the graph
// FunctionLibraryDefinition.
GraphDef graph_def;
graph.ToGraphDef(&graph_def);
TF_RETURN_IF_ERROR(AnalyzeReachableFunctions(graph_def, graph.flib_def()));
// Analyze whether there is an inference graph, including non reachable
// from the `graph` itself. This happens when there is a sequence of
// TPUPartitionedCall()->main()->PartitionedCall() and only second part
// of the graph is processed by the MLIR bridge.
AnalyzeInferenceGraphs(graph.flib_def());
// Check any associated function in the graph defined in a separate
// FunctionLibraryDefinition.
if (function_library != nullptr) {
TF_RETURN_IF_ERROR(
AnalyzeReachableFunctions(graph_def, *function_library));
AnalyzeInferenceGraphs(*function_library);
}
return absl::OkStatus();
}
bool contains_partitioned_call_ = false;
bool contains_ref_type_ = false;
bool invalid_graph_ = false;
bool uses_outside_compilation_ = false;
bool uses_v1_control_flow_ = false;
bool has_manual_control_deps_ = false;
bool single_core_inference_mode_ = false;
bool is_eager_compliant_ = false;
bool has_tpu_replicated_core_ = false;
bool has_infeed_dequeue_tuple_with_tpu_replicated_core_ = false;
};
} // namespace
bool GraphHasUnsupportedFeaturesInMlirBridge(
const Graph& graph, const FunctionLibraryDefinition* function_library,
std::optional<ConfigProto> config_proto, TF2XLABridgeVersion bridge_version,
bool single_core_inference_mode) {
return MlirBridgeGraphAnalyzer(single_core_inference_mode)
.HasUnsupportedFeatures(graph, function_library, config_proto,
bridge_version);
}
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/graph_to_tf_executor_util.h
+65
View File
@@ -0,0 +1,65 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_GRAPH_TO_TF_EXECUTOR_UTIL_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_GRAPH_TO_TF_EXECUTOR_UTIL_H_
#include <optional>
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/protobuf/config.pb.h"
namespace tensorflow {
// These are used for grouping the recorded stats appropriately. Specifically,
// we're considering different entrypoints to the bridge as having potentially
// interesting differences at least in the domain of accepted graphs so we want
// to separately track graph features based on these unique entrypoints. One key
// example of this distinction is for TFRT which uses the "nominal" TPU bridge
// pipeline, but may potentially allow graphs with v1 control flow. This
// separate grouping will allow us to dig into these differences granularly.
enum class TF2XLABridgeVersion {
kNominal = 0,
kV1Compat,
kTFRTNominal,
kNotBridgeUseCase,
};
// Analyzes whether the graph has features not guaranteed to be supported by the
// MLIR-based TF XLA bridge for phase 1. If MLIR bridge phase 1 is not used,
// then MLIR bridge phase 2 will not be used. The optional `function_library`
// can be provided if it contains function definitions not including in the
// `graph` FunctionLibraryDefinition.
//
// Conservatively, during the initial rollout, we are not supporting graphs for
// which any of the following are true:
//
// - Not known to be TF2
// - Contains one or more reference variables
// - Contains one or more TPUPartitionedCall ops (which is a proxy for
// inference), but the graph is not v1 compat
// - Uses V1 control flow
// - Graph is invalid or otherwise encounters error during traversal
// If `single_core_inference_mode` is true, we skip some of check conditions
// because they are not applicable.
// TODO(b/241702857): remove single_core_inference_mode
bool GraphHasUnsupportedFeaturesInMlirBridge(
const Graph& graph, const FunctionLibraryDefinition* function_library,
std::optional<ConfigProto> config_proto, TF2XLABridgeVersion bridge_version,
bool single_core_inference_mode);
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_GRAPH_TO_TF_EXECUTOR_UTIL_H_
tensorflow/compiler/mlir/tf2xla/internal/graph_to_tf_executor_util_test.cc
+736
View File
@@ -0,0 +1,736 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/internal/graph_to_tf_executor_util.h"
#include <cstdint>
#include <initializer_list>
#include <string>
#include <vector>
#include <gtest/gtest.h>
#include "absl/log/check.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "tensorflow/cc/framework/ops.h"
#include "tensorflow/cc/framework/scope.h"
#include "tensorflow/cc/ops/array_ops.h"
#include "tensorflow/cc/ops/control_flow_ops.h"
#include "tensorflow/cc/ops/functional_ops.h"
#include "tensorflow/cc/ops/tpu_functional_ops.h"
#include "tensorflow/cc/ops/tpu_replication_ops.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/platform/enable_tf2_utils.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/config.pb.h"
namespace tensorflow {
namespace {
REGISTER_OP("OneRefOutput").Output("y: Ref(float)");
FunctionDef XTimesTwo() {
const Tensor kTwo = test::AsScalar<int64_t>(2);
return FunctionDefHelper::Define(
// Name
"XTimesTwo",
// Args
{"x: T"},
// Return values
{"y: T"},
// Attr def
{"T: {float, double, int32, int64}"},
// Nodes
{
{{"two"}, "Const", {}, {{"value", kTwo}, {"dtype", DT_INT64}}},
{{"scale"}, "Cast", {"two"}, {{"SrcT", DT_INT64}, {"DstT", "$T"}}},
{{"y"}, "Mul", {"x", "scale"}, {{"T", "$T"}}},
});
}
FunctionDef XTimesTwoFloat() {
const Tensor kTwo = test::AsScalar<int64_t>(2);
return FunctionDefHelper::Define(
// Name
"XTimesTwoFloat",
// Args
{"x: float"},
// Return values
{"y: float"},
// Attr def
{},
// Nodes
{
{{"two"}, "Const", {}, {{"value", kTwo}, {"dtype", DT_INT64}}},
{{"scale"},
"Cast",
{"two"},
{{"SrcT", DT_INT64}, {"DstT", DT_FLOAT}}},
{{"y"}, "Mul", {"x", "scale"}, {{"T", DT_FLOAT}}},
});
}
FunctionDef XTimesTwoFloatRef() {
const Tensor kTwo = test::AsScalar<int64_t>(2);
return FunctionDefHelper::Define(
// Name
"XTimesTwoFloatRef",
// Args
{"x: float"},
// Return values
{"y: float"},
// Attr def
{},
// Nodes
{
{{"two"}, "Const", {}, {{"value", kTwo}, {"dtype", DT_INT64_REF}}},
{{"scale"},
"Cast",
{"two"},
{{"SrcT", DT_INT64_REF}, {"DstT", DT_FLOAT}}},
{{"y"}, "Mul", {"x", "scale"}, {{"T", DT_FLOAT}}},
});
}
Node* FromNodeDef(absl::string_view name, absl::string_view node_type,
int num_inputs, DataType dt, Graph& graph) {
auto builder = NodeDefBuilder(name, node_type);
for (int i = 0; i < num_inputs; ++i) {
builder = builder.Input(absl::StrCat("node_", i), i, dt);
}
NodeDef node_def;
CHECK_OK(builder.Finalize(&node_def));
absl::Status s;
Node* node = graph.AddNode(node_def, &s);
CHECK_OK(s);
return node;
}
TEST(SupportedGraphTest, SupportedGraphReturnsFalse) {
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
auto input = tensorflow::ops::Placeholder(root.WithOpName("input"), DT_UINT8);
auto depth = tensorflow::ops::Placeholder(root.WithOpName("depth"), DT_INT32);
auto on = tensorflow::ops::Placeholder(root.WithOpName("on"), DT_UINT8);
auto off = tensorflow::ops::Placeholder(root.WithOpName("off"), DT_UINT8);
tensorflow::set_tf2_execution(true);
(void)tensorflow::ops::OneHot(root.WithOpName("output"), input, depth, on,
off);
Graph graph(OpRegistry::Global());
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/false));
}
TEST(InvalidGraphTest, InvalidFuncBodyReturnsTrue) {
tensorflow::set_tf2_execution(true);
FunctionDefLibrary flib;
*flib.add_function() = XTimesTwo();
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
NameAttrList f_name_attr;
f_name_attr.set_name("XTimesTwo");
ops::PartitionedCall f(root.WithOpName("f"), {x}, {DT_FLOAT}, f_name_attr);
TF_ASSERT_OK(root.ToGraph(&graph));
// The call to XTimesTwo is invalid (missing an attribute), so we expect the
// graph to be unsupported.
EXPECT_TRUE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/false));
}
TEST(RefVarTest, RefVariablesReturnsTrue) {
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
Output cond_a = ops::Placeholder(root.WithOpName("cond_a"), DT_BOOL);
Output cond_b = ops::Placeholder(root.WithOpName("cond_b"), DT_BOOL);
// Output value = ops::Placeholder(root.WithOpName("value"), DT_FLOAT);
tensorflow::set_tf2_execution(true);
const std::vector<int32_t> shape_array{2, 2};
auto shape = TensorShape();
TF_ASSERT_OK(TensorShapeUtils::MakeShape(shape_array, &shape));
Output value = Output(
FromNodeDef("value", "OneRefOutput", 0, DT_FLOAT_REF, *root.graph()));
Graph graph(OpRegistry::Global());
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_TRUE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/false));
}
TEST(RefVarTest, NoRefVariablesCalleeFuncReturnsFalse) {
tensorflow::set_tf2_execution(true);
FunctionDefLibrary flib;
*flib.add_function() = XTimesTwoFloat();
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
NameAttrList f_name_attr;
f_name_attr.set_name("XTimesTwoFloat");
ops::PartitionedCall f(root.WithOpName("f"), {x}, {DT_FLOAT}, f_name_attr);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/false));
}
TEST(RefVarTest, RefVariablesInCalleeFunctionReturnsTrue) {
tensorflow::set_tf2_execution(true);
FunctionDefLibrary flib;
*flib.add_function() = XTimesTwoFloatRef();
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
NameAttrList f_name_attr;
f_name_attr.set_name("XTimesTwoFloatRef");
ops::PartitionedCall f(root.WithOpName("f"), {x}, {DT_FLOAT}, f_name_attr);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_TRUE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/false));
}
TEST(RefVarTest, RefVariablesInExternalCalleeFunctionReturnsTrue) {
tensorflow::set_tf2_execution(true);
Graph graph(OpRegistry::Global());
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
FunctionDefLibrary flib;
*flib.add_function() = XTimesTwoFloatRef();
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
NameAttrList f_name_attr;
f_name_attr.set_name("XTimesTwoFloatRef");
ops::PartitionedCall f(root.WithOpName("f"), {x}, {DT_FLOAT}, f_name_attr);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_TRUE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/&flib_def, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/false));
}
TEST(InferenceTest, ContainsInferenceNodeEagerRuntimeReturnsTrue) {
tensorflow::set_tf2_execution(true);
FunctionDefLibrary flib;
*flib.add_function() = XTimesTwoFloat();
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
NameAttrList f_name_attr;
f_name_attr.set_name("XTimesTwoFloat");
ops::TPUPartitionedCall f(root.WithOpName("f"), {x}, /*device_ordinal=*/0,
{DT_FLOAT}, f_name_attr);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/false));
}
TEST(InferenceTest, ContainsInferenceNodeTFRTBridgeReturnsTrue) {
tensorflow::set_tf2_execution(true);
FunctionDefLibrary flib;
*flib.add_function() = XTimesTwoFloat();
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
NameAttrList f_name_attr;
f_name_attr.set_name("XTimesTwoFloat");
ops::TPUPartitionedCall f(root.WithOpName("f"), {x}, /*device_ordinal=*/0,
{DT_FLOAT}, f_name_attr);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kTFRTNominal,
/*single_core_inference_mode=*/false));
}
TEST(InferenceTest, ContainsInferenceNodeDirectSessionReturnsFalse) {
tensorflow::set_tf2_execution(true);
FunctionDefLibrary flib;
*flib.add_function() = XTimesTwoFloat();
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kDirectSession);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
NameAttrList f_name_attr;
f_name_attr.set_name("XTimesTwoFloat");
ops::TPUPartitionedCall f(root.WithOpName("f"), {x}, /*device_ordinal=*/0,
{DT_FLOAT}, f_name_attr);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kV1Compat,
/*single_core_inference_mode=*/false));
}
TEST(ControlFlowTest, ContainsV1ControlFlowReturnsTrue) {
tensorflow::set_tf2_execution(true);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
Output cond_a = ops::Placeholder(root.WithOpName("cond_a"), DT_BOOL);
Output cond_b = ops::Placeholder(root.WithOpName("cond_b"), DT_BOOL);
Output value = ops::Placeholder(root.WithOpName("value"), DT_FLOAT);
ops::Switch switch_a(root.WithOpName("switch_a"), value, cond_a);
ops::Switch switch_b(root.WithOpName("switch_b"), value, cond_b);
Graph graph(OpRegistry::Global());
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_TRUE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/false));
}
TEST(ControlFlowTest, TFRTContainsV1ControlFlowReturnsTrue) {
tensorflow::set_tf2_execution(true);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
Output cond_a = ops::Placeholder(root.WithOpName("cond_a"), DT_BOOL);
Output cond_b = ops::Placeholder(root.WithOpName("cond_b"), DT_BOOL);
Output value = ops::Placeholder(root.WithOpName("value"), DT_FLOAT);
ops::Switch switch_a(root.WithOpName("switch_a"), value, cond_a);
ops::Switch switch_b(root.WithOpName("switch_b"), value, cond_b);
Graph graph(OpRegistry::Global());
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_TRUE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kTFRTNominal,
/*single_core_inference_mode=*/false));
}
TEST(TFVersionTest, TF1ReturnsTrue) {
tensorflow::set_tf2_execution(false);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
auto input = tensorflow::ops::Placeholder(root.WithOpName("input"), DT_UINT8);
auto depth = tensorflow::ops::Placeholder(root.WithOpName("depth"), DT_INT32);
auto on = tensorflow::ops::Placeholder(root.WithOpName("on"), DT_UINT8);
auto off = tensorflow::ops::Placeholder(root.WithOpName("off"), DT_UINT8);
(void)tensorflow::ops::OneHot(root.WithOpName("output"), input, depth, on,
off);
Graph graph(OpRegistry::Global());
TF_ASSERT_OK(root.ToGraph(&graph));
graph.SetConstructionContext(ConstructionContext::kDirectSession);
EXPECT_TRUE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kV1Compat,
/*single_core_inference_mode=*/false));
}
TEST(TFVersionTest, TF2ExecutionFalseV1CompatBridgeReturnTrue) {
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
auto input = tensorflow::ops::Placeholder(root.WithOpName("input"), DT_UINT8);
auto depth = tensorflow::ops::Placeholder(root.WithOpName("depth"), DT_INT32);
auto on = tensorflow::ops::Placeholder(root.WithOpName("on"), DT_UINT8);
auto off = tensorflow::ops::Placeholder(root.WithOpName("off"), DT_UINT8);
(void)tensorflow::ops::OneHot(root.WithOpName("output"), input, depth, on,
off);
Graph graph(OpRegistry::Global());
TF_ASSERT_OK(root.ToGraph(&graph));
tensorflow::set_tf2_execution(false);
EXPECT_TRUE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kV1Compat,
/*single_core_inference_mode=*/false));
}
TEST(TFVersionTest, TF2ExecutionTrueV1CompatBridgeReturnFalse) {
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
auto input = tensorflow::ops::Placeholder(root.WithOpName("input"), DT_UINT8);
auto depth = tensorflow::ops::Placeholder(root.WithOpName("depth"), DT_INT32);
auto on = tensorflow::ops::Placeholder(root.WithOpName("on"), DT_UINT8);
auto off = tensorflow::ops::Placeholder(root.WithOpName("off"), DT_UINT8);
(void)tensorflow::ops::OneHot(root.WithOpName("output"), input, depth, on,
off);
Graph graph(OpRegistry::Global());
TF_ASSERT_OK(root.ToGraph(&graph));
tensorflow::set_tf2_execution(true);
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kV1Compat,
/*single_core_inference_mode=*/false));
}
TEST(TFVersionTest, TF2ExecutionFalseTfrtNominalBridgeReturnFalse) {
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
auto input = tensorflow::ops::Placeholder(root.WithOpName("input"), DT_UINT8);
auto depth = tensorflow::ops::Placeholder(root.WithOpName("depth"), DT_INT32);
auto on = tensorflow::ops::Placeholder(root.WithOpName("on"), DT_UINT8);
auto off = tensorflow::ops::Placeholder(root.WithOpName("off"), DT_UINT8);
(void)tensorflow::ops::OneHot(root.WithOpName("output"), input, depth, on,
off);
Graph graph(OpRegistry::Global());
TF_ASSERT_OK(root.ToGraph(&graph));
tensorflow::set_tf2_execution(false);
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kTFRTNominal,
/*single_core_inference_mode=*/false));
}
TEST(TFVersionTest, TF2ExecutionTrueTfrtNominalBridgeReturnFalse) {
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
auto input = tensorflow::ops::Placeholder(root.WithOpName("input"), DT_UINT8);
auto depth = tensorflow::ops::Placeholder(root.WithOpName("depth"), DT_INT32);
auto on = tensorflow::ops::Placeholder(root.WithOpName("on"), DT_UINT8);
auto off = tensorflow::ops::Placeholder(root.WithOpName("off"), DT_UINT8);
(void)tensorflow::ops::OneHot(root.WithOpName("output"), input, depth, on,
off);
Graph graph(OpRegistry::Global());
TF_ASSERT_OK(root.ToGraph(&graph));
tensorflow::set_tf2_execution(true);
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kTFRTNominal,
/*single_core_inference_mode=*/false));
}
TEST(TFVersionTest, TF2ExecutionFalseNominalBridgeReturnsFalse) {
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
auto input = tensorflow::ops::Placeholder(root.WithOpName("input"), DT_UINT8);
Graph graph(OpRegistry::Global());
TF_ASSERT_OK(root.ToGraph(&graph));
tensorflow::set_tf2_execution(false);
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/false));
}
TEST(TFVersionTest, TF2ExecutionTrueNominalBridgeReturnsFalse) {
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
auto input = tensorflow::ops::Placeholder(root.WithOpName("input"), DT_UINT8);
Graph graph(OpRegistry::Global());
TF_ASSERT_OK(root.ToGraph(&graph));
tensorflow::set_tf2_execution(true);
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/false));
}
TEST(UnsupportedOpTest,
InfeedDequeueTupleWithTPUReplicatedCoreAttrNotSupported) {
tensorflow::set_tf2_execution(true);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
auto input =
tensorflow::ops::Placeholder(root.WithOpName("node_0"), DT_FLOAT);
auto node = FromNodeDef("Identity", "Identity", 1, DT_FLOAT, *root.graph());
ASSERT_NE(node, nullptr);
node->set_requested_device("/device:TPU_REPLICATED_CORE:0");
// Build InfeedDequeueTuple node with TPU_REPLICATED_CORE Attr
auto builder = NodeDefBuilder("InfeedDequeueTuple", "InfeedDequeueTuple");
builder.Attr("dtypes", DT_FLOAT);
builder.Attr("shapes", 1);
NodeDef node_def;
CHECK_OK(builder.Finalize(&node_def));
absl::Status s;
Node* node_InfeedDequeueTuple = (*root.graph()).AddNode(node_def, &s);
node_InfeedDequeueTuple->set_requested_device(
"/device:TPU_REPLICATED_CORE:0");
CHECK_OK(s);
ASSERT_NE(node_InfeedDequeueTuple, nullptr);
Graph graph(OpRegistry::Global());
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_TRUE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/false));
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/true));
}
TEST(ManualControlDependencyTest,
TPUReplicatedCoreWithManualControlDependencyReturnsFalse) {
tensorflow::set_tf2_execution(true);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
auto input =
tensorflow::ops::Placeholder(root.WithOpName("node_0"), DT_FLOAT);
auto node = FromNodeDef("Identity", "Identity", 1, DT_FLOAT, *root.graph());
ASSERT_NE(node, nullptr);
node->set_requested_device("/device:TPU_REPLICATED_CORE:0");
auto metadata = tensorflow::ops::TPUReplicateMetadata(root, 2);
metadata.operation.node()->AddAttr("_has_manual_control_dependencies", true);
Graph graph(OpRegistry::Global());
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/false));
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/true));
}
TEST(InferenceTest,
ContainsInferenceNodeTPUReplicatedCoreDirectSessionReturnsFalse) {
tensorflow::set_tf2_execution(true);
FunctionDefLibrary flib;
*flib.add_function() = XTimesTwoFloat();
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kDirectSession);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
auto input =
tensorflow::ops::Placeholder(root.WithOpName("node_0"), DT_FLOAT);
auto node = FromNodeDef("Identity", "Identity", 1, DT_FLOAT, *root.graph());
ASSERT_NE(node, nullptr);
node->set_requested_device("/device:TPU_REPLICATED_CORE:0");
Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
NameAttrList f_name_attr;
f_name_attr.set_name("XTimesTwoFloat");
ops::TPUPartitionedCall f(root.WithOpName("f"), {x}, /*device_ordinal=*/0,
{DT_FLOAT}, f_name_attr);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kV1Compat,
/*single_core_inference_mode=*/false));
}
TEST(InferenceTest,
ContainsInferenceNodeTPUReplicatedCoreEagerRuntimeReturnsTrue) {
tensorflow::set_tf2_execution(true);
FunctionDefLibrary flib;
*flib.add_function() = XTimesTwoFloat();
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
auto input =
tensorflow::ops::Placeholder(root.WithOpName("node_0"), DT_FLOAT);
auto node = FromNodeDef("Identity", "Identity", 1, DT_FLOAT, *root.graph());
ASSERT_NE(node, nullptr);
node->set_requested_device("/device:TPU_REPLICATED_CORE:0");
Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
NameAttrList f_name_attr;
f_name_attr.set_name("XTimesTwoFloat");
ops::TPUPartitionedCall f(root.WithOpName("f"), {x}, /*device_ordinal=*/0,
{DT_FLOAT}, f_name_attr);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kNominal,
/*single_core_inference_mode=*/false));
}
TEST(InferenceTest, TF2ExecutionFalseV1CompatBridgeReturnFalse) {
tensorflow::set_tf2_execution(false);
FunctionDefLibrary flib;
*flib.add_function() = XTimesTwoFloat();
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kDirectSession);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
auto input =
tensorflow::ops::Placeholder(root.WithOpName("node_0"), DT_FLOAT);
auto node = FromNodeDef("Identity", "Identity", 1, DT_FLOAT, *root.graph());
ASSERT_NE(node, nullptr);
node->set_requested_device("/device:TPU_REPLICATED_CORE:0");
Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
NameAttrList f_name_attr;
f_name_attr.set_name("XTimesTwoFloat");
ops::TPUPartitionedCall f(root.WithOpName("f"), {x}, /*device_ordinal=*/0,
{DT_FLOAT}, f_name_attr);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_FALSE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kV1Compat,
/*single_core_inference_mode=*/false));
}
TEST(InferenceTest, V1CompatBridgeVariableRefReturnTrue) {
tensorflow::set_tf2_execution(false);
FunctionDefLibrary flib;
*flib.add_function() = XTimesTwoFloat();
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kDirectSession);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
auto input =
tensorflow::ops::Placeholder(root.WithOpName("node_0"), DT_FLOAT);
auto node = FromNodeDef("Identity", "Identity", 1, DT_FLOAT, *root.graph());
ASSERT_NE(node, nullptr);
node->set_requested_device("/device:TPU_REPLICATED_CORE:0");
Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
NameAttrList f_name_attr;
f_name_attr.set_name("XTimesTwoFloat");
ops::TPUPartitionedCall f(root.WithOpName("f"), {x}, /*device_ordinal=*/0,
{DT_FLOAT}, f_name_attr);
Output cond_a = ops::Placeholder(root.WithOpName("cond_a"), DT_BOOL);
Output cond_b = ops::Placeholder(root.WithOpName("cond_b"), DT_BOOL);
tensorflow::set_tf2_execution(true);
const std::vector<int32_t> shape_array{2, 2};
auto shape = TensorShape();
TF_ASSERT_OK(TensorShapeUtils::MakeShape(shape_array, &shape));
Output value = Output(
FromNodeDef("value", "OneRefOutput", 0, DT_FLOAT_REF, *root.graph()));
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_TRUE(GraphHasUnsupportedFeaturesInMlirBridge(
graph, /*function_library=*/nullptr, config,
/*bridge_version=*/tensorflow::TF2XLABridgeVersion::kV1Compat,
/*single_core_inference_mode=*/false));
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/hlo_post_processing/BUILD
+7
View File
@@ -0,0 +1,7 @@
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow/compiler/mlir/tf2xla/internal:__subpackages__",
],
licenses = ["notice"],
)
tensorflow/compiler/mlir/tf2xla/internal/inference/BUILD
+62
View File
@@ -0,0 +1,62 @@
load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library")
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//visibility:private",
],
)
gentbl_cc_library(
name = "inference_passes_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"inference_passes.h.inc": [
"-gen-pass-decls",
"-name=TF2XLA",
]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "inference_passes.td",
deps = [
"@llvm-project//mlir:PassBaseTdFiles",
],
)
cc_library(
name = "inference_metrics_pass",
srcs = ["inference_metrics_pass.cc"],
hdrs = ["inference_passes.h"],
visibility = ["//visibility:public"],
deps = [
":inference_passes_inc_gen",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/core:lib",
"@com_google_absl//absl/log",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:Dialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:FuncExtensions",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
],
)
tf_cc_test(
name = "inference_metrics_pass_test",
srcs = ["inference_metrics_pass_test.cc"],
deps = [
"inference_metrics_pass",
"//tensorflow/compiler/mlir/tf2xla/transforms:test_utils",
"//tensorflow/core/lib/monitoring:cell_reader",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@xla//xla/tsl/platform:statusor",
],
)
tensorflow/compiler/mlir/tf2xla/internal/inference/inference_metrics_pass.cc
+73
View File
@@ -0,0 +1,73 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <memory>
#include <string>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/core/lib/monitoring/counter.h"
namespace mlir {
namespace tf2xla {
namespace internal {
auto* has_tpu_partitioned_call_streamz =
tensorflow::monitoring::Counter<1>::New(
"/tensorflow/core/tf2xla/internal/inference/tpu_partitioned_call",
"Whether the model has TPUPartitionedCallOp.",
"has_tpu_partitioned_call");
namespace {
#define GEN_PASS_DEF_INFERENCEMETRICSPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/inference/inference_passes.h.inc"
class InferenceMetricsPass
: public impl::InferenceMetricsPassBase<InferenceMetricsPass> {
public:
void runOnOperation() override;
};
void InferenceMetricsPass::runOnOperation() {
bool has_tpu_partitioned_call = false;
ModuleOp module = getOperation();
for (auto func_op : module.getOps<func::FuncOp>()) {
func_op->walk(
[&](TF::TPUPartitionedCallOp op) { has_tpu_partitioned_call = true; });
if (has_tpu_partitioned_call) break;
}
std::string has_tpu_partitioned_call_str =
has_tpu_partitioned_call ? "true" : "false";
has_tpu_partitioned_call_streamz->GetCell(has_tpu_partitioned_call_str)
->IncrementBy(1);
}
} // namespace
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateInferenceMetricsPass() {
return std::make_unique<InferenceMetricsPass>();
}
} // namespace internal
} // namespace tf2xla
} // namespace mlir
tensorflow/compiler/mlir/tf2xla/internal/inference/inference_metrics_pass_test.cc
+105
View File
@@ -0,0 +1,105 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <cstdint>
#include <memory>
#include <string>
#include <gtest/gtest.h>
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tf2xla/internal/inference/inference_passes.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/test_utils.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/lib/monitoring/cell_reader.h"
namespace mlir {
namespace tf2xla {
namespace internal {
namespace {
using ::mlir::MLIRContext;
using ::mlir::ModuleOp;
using ::mlir::OwningOpRef;
using ::mlir::hlo::test::GetMlirModuleFromString;
using ::tensorflow::monitoring::testing::CellReader;
static constexpr char kHasTpuPartitionedCallStreamzName[] =
"/tensorflow/core/tf2xla/internal/inference/tpu_partitioned_call";
class InferenceMetricsPassTest : public ::testing::Test {
protected:
void CreateModule(const char* module_string) {
TF_ASSERT_OK_AND_ASSIGN(module_,
GetMlirModuleFromString(module_string, &context_));
pm_ = std::make_unique<mlir::PassManager>(&context_);
pm_->addPass(CreateInferenceMetricsPass());
}
mlir::LogicalResult Run() { return pm_->run(module_.get()); }
private:
MLIRContext context_;
OwningOpRef<ModuleOp> module_;
std::unique_ptr<mlir::PassManager> pm_;
};
TEST_F(InferenceMetricsPassTest, RecordsTrueForTPUPartitionedCallOp) {
static constexpr char kMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @tpu_partitioned_call_func(%arg0: tensor<?xi32>) -> (tensor<?xi32>) {
func.return %arg0 : tensor<?xi32>
}
func.func @main(%arg0: tensor<20xi32>, %arg1: tensor<?xi32>) -> tensor<*xi32> {
%2 = "tf.TPUPartitionedCall"(%arg0, %arg1) {f = @tpu_partitioned_call_func} : (tensor<20xi32>, tensor<?xi32>) -> tensor<*xi32>
func.return %2 : tensor<*xi32>
}
})";
CellReader<int64_t> error(kHasTpuPartitionedCallStreamzName);
CreateModule(kMlirModuleStr);
auto result = Run();
EXPECT_TRUE(result.succeeded());
EXPECT_EQ(error.Delta("true"), 1);
EXPECT_EQ(error.Delta("false"), 0);
}
TEST_F(InferenceMetricsPassTest, RecordsFalseForNonTPUPartitionedCallOp) {
static constexpr char kMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> tensor<1xi32> {
%0 = "tf.BadValue"() {value = dense<1000> : tensor<1xi32>} : () -> tensor<1xi32>
func.return %0 : tensor<1xi32>
}
})";
CellReader<int64_t> error(kHasTpuPartitionedCallStreamzName);
CreateModule(kMlirModuleStr);
auto result = Run();
EXPECT_TRUE(result.succeeded());
EXPECT_EQ(error.Delta("false"), 1);
EXPECT_EQ(error.Delta("true"), 0);
}
} // namespace
} // namespace internal
} // namespace tf2xla
} // namespace mlir
tensorflow/compiler/mlir/tf2xla/internal/inference/inference_passes.h
+39
View File
@@ -0,0 +1,39 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_INFERENCE_INFERENCE_PASSES_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_INFERENCE_INFERENCE_PASSES_H_
#include <memory>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
namespace mlir {
namespace tf2xla {
namespace internal {
std::unique_ptr<OperationPass<ModuleOp>> CreateInferenceMetricsPass();
#define GEN_PASS_REGISTRATION
#define GEN_PASS_DECL_INFERENCEMETRICSPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/inference/inference_passes.h.inc"
} // namespace internal
} // namespace tf2xla
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_INFERENCE_INFERENCE_PASSES_H_
tensorflow/compiler/mlir/tf2xla/internal/inference/inference_passes.td
+25
View File
@@ -0,0 +1,25 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
include "mlir/Pass/PassBase.td"
// Internal TF2XLA Bridge Passes used during interence only.
def InferenceMetricsPass : Pass<"tf2xla-inference-metrics-pass", "ModuleOp"> {
let summary = "A pass to go over a Module and collect various metrics";
let description = [{
Collects metrics about a Module during inference, such as which ops are used.
}];
let constructor = "internal::CreateInferenceMetricsPass()";
}
tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_mlir.cc
+98
View File
@@ -0,0 +1,98 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_mlir.h"
#include <memory>
#include <vector>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "stablehlo/dialect/Register.h" // from @stablehlo
#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/set_tpu_infeed_layout.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/serialize_mlir_module_utils.h"
#include "tensorflow/compiler/mlir/tf2xla/api/v1/compile_mlir_util.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/compilation_timer.h"
#include "tensorflow/compiler/tf2xla/layout_util.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/mlir_hlo/mhlo/IR/register.h"
#include "xla/shape.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/profile_utils/cpu_utils.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
#include "tensorflow/core/tpu/tpu_compile.h"
#include "tsl/platform/error_logging.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
// Name of component for error logging. This name is fixed and required to
// enable logging.
constexpr char kBridgeComponent[] = "TFXLABridge";
using tpu::ShardingAndIndex;
absl::Status CompileFromMlirToXlaHlo(
bool lower_to_xla_hlo, mlir::ModuleOp mlir_module_op,
const tpu::TPUCompileMetadataProto& metadata, llvm::StringRef device_type,
const XlaShapeLayoutHelpers::ShapeDeterminationFns& shape_determination_fns,
bool use_tuple_args, XlaCompiler::CompilationResult* compilation_result,
std::vector<std::unique_ptr<mlir::Pass>>& custom_legalization_passes,
const std::vector<TensorShape>& arg_shapes,
std::vector<ShardingAndIndex>* arg_core_mapping,
std::vector<std::vector<xla::Shape>>* per_core_arg_shapes) {
LOG_FIRST_N(INFO, 1)
<< "Compiling MLIR computation to XLA HLO using MLIR tf2xla bridge in "
"the op by op fallback mode. This is Phase 2 of the TF2XLA Bridge. "
"Old (non-MLIR) bridge may be used in case of unsupported feature "
"or compilation failure from the MLIR bridge (full fallback mode).";
llvm::SmallVector<TensorOrResourceShape, 4> tensor_or_resource_shapes;
tensor_or_resource_shapes.reserve(arg_shapes.size());
for (const auto& arg_shape : arg_shapes)
tensor_or_resource_shapes.push_back({arg_shape});
TF_RETURN_IF_ERROR(CompileMlirToXlaHlo(
mlir_module_op, tensor_or_resource_shapes, device_type, use_tuple_args,
/*enable_op_fallback=*/true, /*use_return_tuple=*/true,
/*use_resource_updates_for_aliases=*/false, shape_determination_fns,
compilation_result, custom_legalization_passes, metadata.module_name(),
lower_to_xla_hlo));
// Compute how arguments are shared across different cores.
auto sharding_result =
tpu::GetShardingInfo(metadata, arg_shapes, shape_determination_fns,
arg_core_mapping, per_core_arg_shapes);
if (!sharding_result.ok()) {
return sharding_result;
}
return absl::OkStatus();
}
}; // namespace internal
}; // namespace tf2xla
}; // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_mlir.h
+51
View File
@@ -0,0 +1,51 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_LEGALIZE_TF_MLIR_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_LEGALIZE_TF_MLIR_H_
#include <string>
#include <vector>
#include "absl/status/status.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
// Runs all the MLIR Bridge passes on the given MLIR module.
// If compile_to_xla_hlo is true then those passes include all the Legalization
// to XLA HLO which is returned in the compilation_result.
absl::Status CompileFromMlirToXlaHlo(
bool lower_to_xla_hlo, mlir::ModuleOp mlir_module_op,
const tpu::TPUCompileMetadataProto& metadata, llvm::StringRef device_type,
const XlaShapeLayoutHelpers::ShapeDeterminationFns& shape_determination_fns,
bool use_tuple_args, XlaCompiler::CompilationResult* compilation_result,
std::vector<std::unique_ptr<mlir::Pass>>& custom_legalization_passes,
const std::vector<TensorShape>& arg_shapes,
std::vector<tpu::ShardingAndIndex>* arg_core_mapping,
std::vector<std::vector<xla::Shape>>* per_core_arg_shapes);
}; // namespace internal
}; // namespace tf2xla
}; // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_LEGALIZE_TF_MLIR_H_
tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_mlir_test.cc
+96
View File
@@ -0,0 +1,96 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_mlir.h"
#include <memory>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/statusor.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/DialectRegistry.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "stablehlo/dialect/Register.h" // from @stablehlo
#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/deserialize_mlir_module_utils.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/test_matchers.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/mlir_hlo/mhlo/IR/register.h"
#include "xla/shape.h"
#include "xla/tsl/platform/errors.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {
using tpu::ShardingAndIndex;
using tpu::TPUCompileMetadataProto;
static constexpr char kMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> tensor<1xi32> {
%0 = "tf.Const"() {value = dense<1000> : tensor<1xi32>} : () -> tensor<1xi32>
func.return %0 : tensor<1xi32>
}
})";
absl::StatusOr<XlaCompilationResult> CompileMlirModule(bool compile_to_xla_hlo,
const char* module_str) {
XlaCompilationResult compilation_result;
mlir::DialectRegistry registry;
mlir::RegisterAllTensorFlowDialects(registry);
mlir::mhlo::registerAllMhloDialects(registry);
mlir::stablehlo::registerAllDialects(registry);
mlir::MLIRContext context(registry);
mlir::OwningOpRef<mlir::ModuleOp> mlir_module;
TF_RETURN_IF_ERROR(
tensorflow::DeserializeMlirModule(module_str, &context, &mlir_module));
std::vector<TensorShape> arg_shapes;
TPUCompileMetadataProto metadata_proto;
bool use_tuple_args = true;
std::vector<ShardingAndIndex> arg_core_mapping;
std::vector<std::vector<xla::Shape>> per_core_arg_shapes;
std::vector<std::unique_ptr<mlir::Pass>> custom_legalization_passes;
TF_RETURN_IF_ERROR(CompileFromMlirToXlaHlo(
compile_to_xla_hlo, mlir_module.get(), metadata_proto,
/*device_type=*/"XLA_TPU_JIT",
/*shape_determination_fns=*/{}, use_tuple_args, &compilation_result,
custom_legalization_passes, arg_shapes, &arg_core_mapping,
&per_core_arg_shapes));
return compilation_result;
}
TEST(CompileFromMlir, ReturnsLoweredModule) {
auto result = CompileMlirModule(/*compile_to_xla_hlo=*/true, kMlirModuleStr);
ASSERT_THAT(result, IsOkOrFiltered());
EXPECT_THAT(result, ComputationProtoContains("opcode: \"constant\""));
}
} // namespace
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo.cc
+151
View File
@@ -0,0 +1,151 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo.h"
#include <memory>
#include <vector>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/DialectRegistry.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "mlir/Interfaces/DataLayoutInterfaces.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "stablehlo/dialect/Register.h" // from @stablehlo
#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/set_tpu_infeed_layout.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/deserialize_mlir_module_utils.h"
#include "tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_mlir.h"
#include "tensorflow/compiler/tf2xla/layout_util.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/client/compile_only_client.h"
#include "xla/mlir_hlo/mhlo/IR/register.h"
#include "xla/shape.h"
#include "xla/tsl/framework/device_type.h"
#include "xla/tsl/platform/errors.h"
#include "tensorflow/core/framework/metrics.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {
using metrics::IncrementTfMlirBridgeSecondPhaseCounter;
using metrics::MlirBridgeSecondPhaseMetric;
using tpu::MlirToHloArgs;
absl::Status CheckAndIncrementCounter(absl::Status status,
MlirBridgeSecondPhaseMetric metric) {
if (!status.ok()) {
IncrementTfMlirBridgeSecondPhaseCounter(metric);
return status;
}
return absl::OkStatus();
}
absl::StatusOr<XlaCompilationResult> LegalizeTfToHlo(
mlir::ModuleOp mlir_module_op, const tpu::TPUCompileMetadataProto& metadata,
bool use_tuple_args, llvm::StringRef device_type,
XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,
const std::vector<tensorflow::TensorShape>& arg_shapes,
std::vector<tpu::ShardingAndIndex>* arg_core_mapping,
std::vector<std::vector<xla::Shape>>* per_core_arg_shapes,
std::vector<std::unique_ptr<mlir::Pass>>& custom_legalization_passes,
xla::CompileOnlyClient* client, XlaCompilationResult* compilation_result) {
if (!mlir::SetTPUInfeedLayout(mlir_module_op))
return absl::AbortedError("Failed to set layouts attribute");
absl::Status mlir_compilation = internal::CompileFromMlirToXlaHlo(
/*lower_to_xla_hlo=*/false, mlir_module_op, metadata, device_type,
shape_determination_fns, use_tuple_args, compilation_result,
custom_legalization_passes, arg_shapes, arg_core_mapping,
per_core_arg_shapes);
TF_RETURN_IF_ERROR(CheckAndIncrementCounter(
mlir_compilation, MlirBridgeSecondPhaseMetric::kMlirCombinedMlirFailure));
IncrementTfMlirBridgeSecondPhaseCounter(
MlirBridgeSecondPhaseMetric::kMlirCombinedMlirSuccess);
MlirToHloArgs mlir_to_hlo_args;
mlir_to_hlo_args.mlir_module_op = mlir_module_op;
absl::Status old_bridge_status = v1::CompileTensorflowGraphToHlo(
mlir_to_hlo_args, metadata, use_tuple_args, shape_determination_fns,
arg_shapes, tsl::DeviceType(device_type.str()), arg_core_mapping,
per_core_arg_shapes, client, compilation_result);
TF_RETURN_IF_ERROR(CheckAndIncrementCounter(
old_bridge_status, MlirBridgeSecondPhaseMetric::kMlirCombinedOldFailure));
IncrementTfMlirBridgeSecondPhaseCounter(
MlirBridgeSecondPhaseMetric::kMlirCombinedOldSuccess);
return *compilation_result;
}
} // namespace
absl::StatusOr<XlaCompilationResult> LegalizeTfToHlo(
const tpu::MlirToHloArgs& computation,
const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,
llvm::StringRef device_type,
XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,
const std::vector<tensorflow::TensorShape>& arg_shapes,
std::vector<tpu::ShardingAndIndex>* arg_core_mapping,
std::vector<std::vector<xla::Shape>>* per_core_arg_shapes,
std::vector<std::unique_ptr<mlir::Pass>>& custom_legalization_passes,
xla::CompileOnlyClient* client, XlaCompilationResult* compilation_result) {
LOG_FIRST_N(INFO, 1) << "Compiling MLIR computation to XLA HLO using the "
"Combined MLIR Tf2Xla Bridge.";
if (computation.mlir_module_op.has_value()) {
return LegalizeTfToHlo(computation.mlir_module_op.value(), metadata,
use_tuple_args, device_type, shape_determination_fns,
arg_shapes, arg_core_mapping, per_core_arg_shapes,
custom_legalization_passes, client,
compilation_result);
}
mlir::DialectRegistry registry;
mlir::RegisterAllTensorFlowDialects(registry);
mlir::mhlo::registerAllMhloDialects(registry);
mlir::stablehlo::registerAllDialects(registry);
mlir::MLIRContext context(registry);
mlir::OwningOpRef<mlir::ModuleOp> mlir_module;
absl::Status deserialization_status = tensorflow::DeserializeMlirModule(
computation.mlir_module, &context, &mlir_module);
TF_RETURN_IF_ERROR(CheckAndIncrementCounter(
deserialization_status,
MlirBridgeSecondPhaseMetric::kMlirCombinedMlirFailure));
return LegalizeTfToHlo(mlir_module.get(), metadata, use_tuple_args,
device_type, shape_determination_fns, arg_shapes,
arg_core_mapping, per_core_arg_shapes,
custom_legalization_passes, client,
compilation_result);
}
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo.h
+49
View File
@@ -0,0 +1,49 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_LEGALIZE_TF_TO_HLO_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_LEGALIZE_TF_TO_HLO_H_
#include "absl/status/statusor.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/client/compile_only_client.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
// Legalize the given MLIR module to XLA HLO using a combination of the MLIR
// Bridge and XlaBuilder
absl::StatusOr<XlaCompilationResult> LegalizeTfToHlo(
const tpu::MlirToHloArgs& computation,
const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,
llvm::StringRef device_type,
XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,
const std::vector<tensorflow::TensorShape>& arg_shapes,
std::vector<tpu::ShardingAndIndex>* arg_core_mapping,
std::vector<std::vector<xla::Shape>>* per_core_arg_shapes,
std::vector<std::unique_ptr<mlir::Pass>>& custom_legalization_passes,
xla::CompileOnlyClient* client, XlaCompilationResult* compilation_result);
}; // namespace internal
}; // namespace tf2xla
}; // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_LEGALIZE_TF_TO_HLO_H_
tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo_test.cc
+186
View File
@@ -0,0 +1,186 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo.h"
#include <cstdint>
#include <memory>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/statusor.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/DialectRegistry.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "stablehlo/dialect/Register.h" // from @stablehlo
#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/deserialize_mlir_module_utils.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/test_matchers.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/client/client_library.h"
#include "xla/mlir_hlo/mhlo/IR/register.h"
#include "xla/shape.h"
#include "xla/stream_executor/platform.h"
#include "xla/stream_executor/platform_manager.h"
#include "xla/tsl/platform/errors.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/lib/monitoring/cell_reader.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
using ::tensorflow::monitoring::testing::CellReader;
using tpu::MlirToHloArgs;
using tpu::ShardingAndIndex;
using tpu::TPUCompileMetadataProto;
static constexpr char kMlirLegalizeCount[] =
"/tensorflow/core/tf2xla/v1/mlir_failed_xla_legalize_tf_count";
static constexpr char kMlirLegalizeErrors[] =
"/tensorflow/core/tf2xla/v1/mlir_failed_xla_legalize_tf_pass_count";
static constexpr char kBridgeStatusCounter[] =
"/tensorflow/core/tf2xla/api/v2/phase2_compilation_status";
constexpr char kMlirCombinedMlirSuccess[] = "kMlirCombinedMlirSuccess";
constexpr char kMlirCombinedOldSuccess[] = "kMlirCombinedOldSuccess";
constexpr char kMlirCombinedOldFailure[] = "kMlirCombinedOldFailure";
static constexpr char kMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main(%arg0 : tensor<1xf32>) -> tensor<1xf32> {
%0 = "tf.Acos"(%arg0) : (tensor<1xf32>) -> tensor<1xf32>
func.return %0 : tensor<1xf32>
}
})";
static constexpr char kBadMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> tensor<1xi32> {
%0 = "tf.DoesntExist"() {value = dense<1000> : tensor<1xi32>} : () -> tensor<1xi32>
func.return %0 : tensor<1xi32>
}
})";
absl::StatusOr<XlaCompiler::CompilationResult> CompileMlirModule(
const char* module_str, bool compile_serialized = true) {
MlirToHloArgs mlir_to_hlo_args;
mlir_to_hlo_args.rollout_state =
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED;
mlir::DialectRegistry registry;
mlir::RegisterAllTensorFlowDialects(registry);
mlir::mhlo::registerAllMhloDialects(registry);
mlir::stablehlo::registerAllDialects(registry);
mlir::MLIRContext context(registry);
mlir::OwningOpRef<mlir::ModuleOp> mlir_module;
if (compile_serialized) {
mlir_to_hlo_args.mlir_module = module_str;
} else {
TF_RETURN_IF_ERROR(
tensorflow::DeserializeMlirModule(module_str, &context, &mlir_module));
mlir_to_hlo_args.mlir_module_op = mlir_module.get();
}
se::Platform* platform =
se::PlatformManager::PlatformWithName("Host").value();
auto client =
xla::ClientLibrary::GetOrCreateCompileOnlyClient(platform).value();
std::vector<TensorShape> arg_shapes = {{1}};
TPUCompileMetadataProto metadata_proto;
auto arg = metadata_proto.add_args();
arg->set_dtype(DataType::DT_FLOAT);
arg->set_kind(TPUCompileMetadataProto::Arg::PARAMETER);
metadata_proto.add_retvals();
bool use_tuple_args = true;
std::vector<ShardingAndIndex> arg_core_mapping;
std::vector<std::vector<xla::Shape>> per_core_arg_shapes;
std::vector<std::unique_ptr<mlir::Pass>> custom_legalization_passes;
auto compilation_result = std::make_unique<XlaCompilationResult>();
return LegalizeTfToHlo(mlir_to_hlo_args, metadata_proto, use_tuple_args,
/*device_type=*/"XLA_TPU_JIT",
/*shape_determination_fns=*/{}, arg_shapes,
&arg_core_mapping, &per_core_arg_shapes,
custom_legalization_passes, client,
compilation_result.get());
}
using LegalizeWithCombinedBridge = ::testing::TestWithParam<bool>;
TEST_P(LegalizeWithCombinedBridge, DoesNotUseMlirLowering) {
const bool compile_serialized = GetParam();
CellReader<int64_t> mlir_bridge_legalize_count(kMlirLegalizeCount);
CellReader<int64_t> counts(kBridgeStatusCounter);
auto result = CompileMlirModule(kMlirModuleStr, compile_serialized);
ASSERT_THAT(result, IsOkOrFiltered());
EXPECT_EQ(mlir_bridge_legalize_count.Delta("tf.Acos"), 0);
EXPECT_THAT(result,
IncrementedOrFiltered(counts.Delta(kMlirCombinedMlirSuccess), 1));
EXPECT_THAT(result,
IncrementedOrFiltered(counts.Delta(kMlirCombinedOldSuccess), 1));
}
TEST_P(LegalizeWithCombinedBridge,
CorrectlyCountsMlirBridgePassingAndGraphBridgeFailing) {
const bool compile_serialized = GetParam();
CellReader<int64_t> legalize_failure_count(kMlirLegalizeErrors);
CellReader<int64_t> counts(kBridgeStatusCounter);
auto result = CompileMlirModule(kBadMlirModuleStr, compile_serialized);
ASSERT_FALSE(result.ok());
// Never failed to legalize because it was never attempted
EXPECT_EQ(legalize_failure_count.Read("tf.DoesntExist", "Unknown"), 0);
EXPECT_THAT(result,
IncrementedOrFiltered(counts.Delta(kMlirCombinedMlirSuccess), 1));
EXPECT_THAT(result,
IncrementedOrFiltered(counts.Delta(kMlirCombinedOldFailure), 1));
}
TEST_P(LegalizeWithCombinedBridge, RecordsDynamicOps) {
const bool compile_serialized = GetParam();
static constexpr char kDynamismFunctionCounterStreamzName[] =
"/tensorflow/core/tf2xla/api/v2/dynamism_function_counter";
constexpr char kNotDynamicFunctionName[] = "kNotDynamicFunction";
CellReader<int64_t> dynamic_function_op_count(
kDynamismFunctionCounterStreamzName);
auto result = CompileMlirModule(kMlirModuleStr, compile_serialized);
ASSERT_TRUE(result.ok());
EXPECT_EQ(dynamic_function_op_count.Delta(kNotDynamicFunctionName), 1);
}
INSTANTIATE_TEST_SUITE_P(LegalizeWithCombinedBridge, LegalizeWithCombinedBridge,
::testing::Bool());
}; // namespace internal
}; // namespace tf2xla
}; // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/logging_hooks.cc
+52
View File
@@ -0,0 +1,52 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/internal/logging_hooks.h"
#include <memory>
#include <string>
#include "llvm/ADT/StringRef.h"
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/utils/data_dumper_logger_config.h"
#include "tensorflow/core/util/debug_data_dumper.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
using mlir::PassManager;
// Setup the input pass manager to enable IR dumping after each pass.
// Note a side effect of this method is that multi threading will be disabled.
void EnablePassIRPrinting(PassManager& pm, const std::string& dump_group_name,
llvm::StringRef module_name) {
// Print the whole module after each pass, which requires disabling
// multi-threading as well.
pm.getContext()->disableMultithreading();
pm.enableIRPrinting(std::make_unique<::tensorflow::DataDumperLoggerConfig>(
[module_name, dump_group_name](const std::string& pass_tag_name,
mlir::Operation* op) {
return DEBUG_DATA_DUMPER()->GetDumpFilename(
module_name.str(), dump_group_name, pass_tag_name);
},
/*pass_prefix=*/"",
/*print_module_scope=*/true));
pm.enableTiming();
}
}; // namespace internal
}; // namespace tf2xla
}; // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/logging_hooks.h
+38
View File
@@ -0,0 +1,38 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_LOGGING_HOOKS_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_LOGGING_HOOKS_H_
#include <string>
#include "llvm/ADT/StringRef.h"
#include "mlir/Pass/PassManager.h" // from @llvm-project
namespace tensorflow {
namespace tf2xla {
namespace internal {
// Setup the input pass manager to enable IR dumping after each pass.
// Note a side effect of this method is that multi threading will be disabled.
void EnablePassIRPrinting(mlir::PassManager& pm,
const std::string& dump_group_name,
llvm::StringRef module_name = llvm::StringRef());
}; // namespace internal
}; // namespace tf2xla
}; // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_LOGGING_HOOKS_H_
tensorflow/compiler/mlir/tf2xla/internal/logging_hooks_test.cc
+112
View File
@@ -0,0 +1,112 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/internal/logging_hooks.h"
#include <string>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/Parser/Parser.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "mlir/Transforms/Passes.h" // from @llvm-project
#include "tensorflow/compiler/mlir/register_common_dialects.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/platform/file_statistics.h"
#include "xla/tsl/platform/status.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/resource_loader.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {
using mlir::DialectRegistry;
using mlir::LogicalResult;
using mlir::MLIRContext;
using mlir::ModuleOp;
using mlir::OwningOpRef;
using mlir::PassManager;
using mlir::func::FuncOp;
std::string TestDataPath() {
return tensorflow::GetDataDependencyFilepath(
"tensorflow/compiler/mlir/tf2xla/internal/testdata/");
}
class LoggingHooksTest : public ::testing::Test {
public:
LoggingHooksTest() {
mlir::RegisterCommonToolingDialects(registry_);
context_.appendDialectRegistry(registry_);
context_.loadAllAvailableDialects();
env_ = Env::Default();
test_group_name_ = "TestGroup";
test_dir_ = testing::TmpDir();
setenv(/*name=*/"TF_DUMP_GRAPH_PREFIX", /*value=*/test_dir_.c_str(),
/*overwrite=*/1);
}
absl::Status CreateMlirModule(std::string mlir_module_filename) {
std::string mlir_module_path = TestDataPath() + mlir_module_filename;
mlir_module_ =
mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context_);
if (!mlir_module_) {
return absl::Status(
absl::StatusCode::kNotFound,
absl::StrCat("Could not find MLIR module at ", mlir_module_path));
}
return absl::OkStatus();
}
DialectRegistry registry_;
MLIRContext context_;
OwningOpRef<mlir::ModuleOp> mlir_module_;
Env* env_;
std::string test_dir_;
std::string test_group_name_;
};
TEST_F(LoggingHooksTest, DumpsPassData) {
std::vector<std::string> files;
TF_ASSERT_OK(env_->GetChildren(test_dir_, &files));
EXPECT_THAT(files, ::testing::IsEmpty());
TF_ASSERT_OK(CreateMlirModule("dead_const.mlir"));
PassManager pass_manager(&context_);
pass_manager.addNestedPass<FuncOp>(mlir::createCanonicalizerPass());
EnablePassIRPrinting(pass_manager, test_group_name_);
LogicalResult pass_status = pass_manager.run(mlir_module_.get());
EXPECT_TRUE(pass_status.succeeded());
TF_ASSERT_OK(env_->GetChildren(test_dir_, &files));
EXPECT_THAT(files, ::testing::SizeIs(2));
}
}; // namespace
}; // namespace internal
}; // namespace tf2xla
}; // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/mlir_bridge_pass_util.cc
+246
View File
@@ -0,0 +1,246 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/internal/mlir_bridge_pass_util.h"
#include <functional>
#include <memory>
#include <string>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/tf2xla/tf2xla_defs.h"
#include "xla/tsl/platform/status.h"
#include "tensorflow/core/common_runtime/function_body.h"
#include "tensorflow/core/common_runtime/function_def_utils.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/graph/graph.h"
namespace tensorflow {
using ::mlir::failure;
using ::mlir::LogicalResult;
using ::mlir::success;
namespace {
constexpr absl::string_view kPartitionedCall = "TPUPartitionedCall";
LogicalResult HasAttr(
const Graph& graph, const FunctionLibraryDefinition* function_library,
const std::function<bool(const Graph& graph)>& predicate) {
if (predicate(graph)) {
return success();
}
// Check if any reachable functions from the graph has the target attribute.
GraphDef graph_def;
graph.ToGraphDef(&graph_def);
if (!function_library) return failure();
for (const std::string& func_name :
function_library->ReachableDefinitions(graph_def).ListFunctionNames()) {
const FunctionDef* func_def = function_library->Find(func_name);
std::unique_ptr<FunctionBody> func_body;
absl::Status status = FunctionDefToBodyHelper(
*func_def, AttrSlice(&func_def->attr()), function_library, &func_body);
// This is not expected to happen in practice
if (!status.ok()) {
LOG(ERROR) << "Failed to parse " << func_name << ": "
<< absl::StatusMessageAsCStr(status) << "; this can "
<< "happen if the function contains unregistered ops.";
return failure();
}
if (predicate(*func_body->graph)) {
return success();
}
}
return failure();
}
// Check if the `graph` has parameter server jobs and resource variable
// arguments that are on parameter servers
bool HasPsWithResourceVariable(const Graph& graph) {
// Check parameter serverjobs and resource variable arguments that are
// on parameter servers.
const std::string jobType = "ps";
const std::string nodeType = "_Arg";
const std::string attrKey = "T";
for (const Node* node : graph.nodes()) {
if (node->type_string() == nodeType) {
auto device_name = node->assigned_device_name();
DeviceNameUtils::ParsedName device;
if (DeviceNameUtils::ParseFullName(device_name, &device) &&
device.has_job && device.job == jobType) {
for (const auto& attr : node->attrs()) {
auto attr_key = attr.first;
auto attr_value = attr.second;
if (attr_key == attrKey &&
attr_value.value_case() == AttrValue::kType &&
attr_value.type() == DT_RESOURCE) {
return true;
break;
}
}
}
}
}
return false;
}
bool IsNonReplicatedGraph(const Graph& graph,
const FunctionLibraryDefinition* function_library) {
auto predicate = [](const Graph& graph) {
const std::string kStatefulPartitionedCallOp = "StatefulPartitionedCall";
for (const Node* node : graph.nodes()) {
auto node_op = node->type_string();
if (node_op == kStatefulPartitionedCallOp) {
// Functions called by StatefulfulPartitionedCall ops with
// _XlaMustCompile=true are compiled by XLA.
auto attr = node->attrs().FindByString(std::string(kMustCompileAttr));
if (attr != nullptr && attr->b() == true) {
return true;
}
}
}
return false;
};
return HasAttr(graph, function_library, predicate).succeeded();
}
bool IsReplicatedGraph(const Graph& graph,
const FunctionLibraryDefinition* function_library) {
auto predicate = [](const Graph& graph) {
for (const Node* node : graph.nodes()) {
// _tpu_replicate is used in replicated TPU graphs. It will be converted
// to_replication_info and _xla_compile_device_type in phase 1 pipelines.
if (node->attrs().FindByString(std::string(kTpuReplicateAttr))) {
return true;
}
}
return false;
};
return HasAttr(graph, function_library, predicate).succeeded();
}
bool IsReplicatedGraph(mlir::ModuleOp module) {
auto walk_result = module.walk([&](mlir::Operation* op) {
// TODO(b/223677572): Once the scope for new compilation and replication
// markers is expanded beyond bridge we can remove this check for
// `kTPUReplicateAttr`, we will then always have a `kCompileDeviceTypeAttr`
// in such cases (see above).
// TODO(b/229028654): Remove string conversion once we have C++17.
const llvm::StringRef tpu_replicate_attr_name(kTpuReplicateAttr.data(),
kTpuReplicateAttr.size());
auto replicate_attr =
op->getAttrOfType<mlir::StringAttr>(tpu_replicate_attr_name);
if (replicate_attr) return mlir::WalkResult::interrupt();
return mlir::WalkResult::advance();
});
return walk_result.wasInterrupted();
}
// Traverses each node in the graph and check if any of them is
// TPUPartitionedCall. If so, return true. Otherwise, return false.
bool DoesGraphContainTPUPartitionedCall(const Graph& graph) {
for (const Node* node : graph.nodes()) {
if (node->type_string() == kPartitionedCall) return true;
}
return false;
}
// Checks any reachable functions from `graph_def` in `flib_def`
// for inference graphs.
bool DoReachableFuncsContainTPUPartitionedCall(
const GraphDef& graph_def, const FunctionLibraryDefinition& flib_def) {
for (const std::string& func_name :
flib_def.ReachableDefinitions(graph_def).ListFunctionNames()) {
const FunctionDef* func_def = flib_def.Find(func_name);
std::unique_ptr<FunctionBody> func_body;
if (!FunctionDefToBodyHelper(*func_def, AttrSlice(&func_def->attr()),
&flib_def, &func_body)
.ok())
return false;
if (DoesGraphContainTPUPartitionedCall(*func_body->graph)) return true;
}
return false;
}
// Iterate all functions from the flib_def if there are any that belong to
// the inference graph.
bool AreFunctionsFromFlibDefInference(
const FunctionLibraryDefinition& flib_def) {
for (const std::string& func_name : flib_def.ListFunctionNames()) {
const FunctionDef* func_def = flib_def.Find(func_name);
for (const NodeDef& node_def : func_def->node_def()) {
if (node_def.op() == kPartitionedCall) return true;
}
}
return false;
}
} // namespace
bool IsSupportedByNonReplicatedBridge(
const Graph& graph, const FunctionLibraryDefinition* function_library) {
return IsNonReplicatedGraph(graph, function_library) &&
HasPsWithResourceVariable(graph);
}
bool IsSupportedByReplicatedBridge(
const Graph& graph, const FunctionLibraryDefinition* function_library) {
return IsReplicatedGraph(graph, function_library);
}
bool IsSupportedByReplicatedBridge(mlir::ModuleOp module) {
return IsReplicatedGraph(module);
}
bool HasTPUPartitionedCallOpInModule(mlir::ModuleOp module) {
bool has_tpu_partitioned_call = false;
for (auto func_op : module.getOps<mlir::func::FuncOp>()) {
func_op->walk([&](mlir::TF::TPUPartitionedCallOp op) {
has_tpu_partitioned_call = true;
});
if (has_tpu_partitioned_call) break;
}
return has_tpu_partitioned_call;
}
bool IsInferenceGraph(const Graph& graph,
const FunctionLibraryDefinition* function_library) {
if (DoesGraphContainTPUPartitionedCall(graph)) return true;
GraphDef graph_def;
graph.ToGraphDef(&graph_def);
if (DoReachableFuncsContainTPUPartitionedCall(graph_def, graph.flib_def()))
return true;
if (AreFunctionsFromFlibDefInference(graph.flib_def())) return true;
if (function_library == nullptr) return false;
if (DoReachableFuncsContainTPUPartitionedCall(graph_def, *function_library))
return true;
if (AreFunctionsFromFlibDefInference(*function_library)) return true;
return false;
}
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/mlir_bridge_pass_util.h
+54
View File
@@ -0,0 +1,54 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_MLIR_BRIDGE_PASS_UTIL_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_MLIR_BRIDGE_PASS_UTIL_H_
#include <optional>
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "tensorflow/core/framework/function.h"
namespace tensorflow {
// Checks if a graph or reachable functions in the library have any
// StatefulPartitionedOps with _XlaMustCompile=true. The function library will
// be skipped if nullptr is provided.
bool IsSupportedByNonReplicatedBridge(
const Graph& graph, const FunctionLibraryDefinition* function_library);
// Checks if a graph or reachable functions in the library have any ops with
// _tpu_replicate or _xla_compile_device_type=TPU. The function library will be
// skipped if nullptr is provided.
bool IsSupportedByReplicatedBridge(
const Graph& graph, const FunctionLibraryDefinition* function_library);
// Check if an MLIR module has any ops with _tpu_replicate or
// _xla_compile_device_type=TPU.
bool IsSupportedByReplicatedBridge(mlir::ModuleOp module);
// Check if an MLIR module contains TPUPartitionedCall op. If so, we define
// such graph as an inference graph. Otherwise, it is non inference graph.
bool HasTPUPartitionedCallOpInModule(mlir::ModuleOp module);
// Check if a graph contains TPUPartitionedCall op, including its reachable
// functions. The function library is used to store the functions that are
// defined in a TensorFlow program
bool IsInferenceGraph(const Graph& graph,
const FunctionLibraryDefinition* function_library);
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_MLIR_BRIDGE_PASS_UTIL_H_
tensorflow/compiler/mlir/tf2xla/internal/mlir_bridge_pass_util_test.cc
+316
View File
@@ -0,0 +1,316 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/internal/mlir_bridge_pass_util.h"
#include <cstdint>
#include <string>
#include <vector>
#include <gtest/gtest.h>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "mlir/Parser/Parser.h" // from @llvm-project
#include "tensorflow/cc/framework/ops.h"
#include "tensorflow/cc/framework/scope.h"
#include "tensorflow/cc/ops/array_ops.h"
#include "tensorflow/cc/ops/function_ops.h"
#include "tensorflow/cc/ops/tpu_functional_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/tf2xla/tf2xla_defs.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/function_testlib.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/platform/enable_tf2_utils.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/config.pb.h"
namespace tensorflow {
namespace {
// Produce a valid graph with a resource-type input.
FunctionDef PassThroughResource() {
return FunctionDefHelper::Define(
/*function_name=*/"PassThroughResource",
/*arg_def=*/{"in: resource"},
/*ret_def=*/{"out: resource"},
/*attr_def=*/{},
/*node_def=*/
{{{"out"}, "Identity", {"in"}, {{"T", DataType::DT_RESOURCE}}}});
}
TEST(IsSupportedByNonReplicatedBridge, NonReplicatedGraph) {
const FunctionDef& fd = PassThroughResource();
FunctionDefLibrary flib;
*flib.add_function() = fd;
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
tensorflow::set_tf2_execution(true);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
Output a = ops::_Arg(root.WithOpName("A"), DT_RESOURCE, 0);
std::vector<NodeBuilder::NodeOut> inputs({NodeBuilder::NodeOut(a.node())});
Node* call;
NameAttrList f_name_attr;
f_name_attr.set_name(fd.signature().name());
TF_ASSERT_OK(
NodeBuilder("B", "StatefulPartitionedCall", &root.graph()->flib_def())
.Input(inputs)
.Attr("Tin", {DT_RESOURCE})
.Attr("Tout", {DT_RESOURCE})
.Attr("f", f_name_attr)
.Finalize(root.graph(), &call));
call->AddAttr(std::string(kMustCompileAttr), true);
TF_ASSERT_OK(root.ToGraph(&graph));
// Required for passing the PS server parameter check.
for (Node* node : graph.nodes()) {
node->set_assigned_device_name("/job:ps/replica:0/task:0/device:GPU:0");
}
EXPECT_TRUE(
IsSupportedByNonReplicatedBridge(graph, /*function_library=*/nullptr));
}
TEST(IsSupportedByReplicatedBridge, ReplicatedGraph) {
const FunctionDef& fd = test::function::XTimesTwo();
FunctionDefLibrary flib;
*flib.add_function() = fd;
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
tensorflow::set_tf2_execution(true);
ConfigProto config = ConfigProto();
Scope root = Scope::NewRootScope().ExitOnError();
Output a = ops::_Arg(root.WithOpName("A"), DT_FLOAT, 0);
std::vector<NodeBuilder::NodeOut> inputs({NodeBuilder::NodeOut(a.node())});
Node* call;
NameAttrList f_name_attr;
f_name_attr.set_name(fd.signature().name());
TF_ASSERT_OK(
NodeBuilder("B", "StatefulPartitionedCall", &root.graph()->flib_def())
.Input(inputs)
.Attr("Tin", {DT_FLOAT})
.Attr("Tout", {DT_FLOAT})
.Attr("f", f_name_attr)
.Finalize(root.graph(), &call));
call->AddAttr(std::string(kTpuReplicateAttr), "cluster");
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_TRUE(
IsSupportedByReplicatedBridge(graph, /*function_library=*/nullptr));
}
TEST(IsSupportedByReplicatedBridge, ReplicatedModule) {
const char* const code = R"mlir(
func.func @entry_func_1(%arg0: tensor<i32>) -> tensor<i32> attributes {tf.entry_function = {}} {
%0 = "tf.Identity"(%arg0) {_tpu_replicate = "cluster"} : (tensor<i32>) -> (tensor<i32>)
func.return %0 : tensor<i32>
}
)mlir";
mlir::MLIRContext context;
context.loadDialect<mlir::func::FuncDialect, mlir::TF::TensorFlowDialect>();
mlir::OwningOpRef<mlir::ModuleOp> module =
mlir::parseSourceString<mlir::ModuleOp>(code, &context);
ASSERT_TRUE(module);
EXPECT_TRUE(IsSupportedByReplicatedBridge(*module));
}
TEST(HasTPUPartitionedCallOpInModule, HasTPUPartitionedCallModule) {
const char* const code = R"mlir(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() {
%outputs_0 = "tf.TPUOrdinalSelector"() {device = ""} : () -> tensor<?xi32>
"tf.TPUPartitionedCall"(%outputs_0) {f = @reachable_func} : (tensor<?xi32>) -> ()
func.return
}
func.func @reachable_func() {
func.return
}
}
)mlir";
mlir::MLIRContext context;
context.loadDialect<mlir::func::FuncDialect, mlir::TF::TensorFlowDialect>();
mlir::OwningOpRef<mlir::ModuleOp> module =
mlir::parseSourceString<mlir::ModuleOp>(code, &context);
ASSERT_TRUE(module);
EXPECT_TRUE(HasTPUPartitionedCallOpInModule(*module));
}
TEST(HasTPUPartitionedCallOpInModule, HasNotTPUPartitionedCallModule) {
const char* const code = R"mlir(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() {
"tf.StatefulPartitionedCall"() {config = "", config_proto = "", executor_type = "", f = @reachable_func} : () -> ()
func.return
}
func.func @reachable_func() {
func.return
}
}
)mlir";
mlir::MLIRContext context;
context.loadDialect<mlir::func::FuncDialect, mlir::TF::TensorFlowDialect>();
mlir::OwningOpRef<mlir::ModuleOp> module =
mlir::parseSourceString<mlir::ModuleOp>(code, &context);
ASSERT_TRUE(module);
EXPECT_FALSE(HasTPUPartitionedCallOpInModule(*module));
}
TEST(IsInferenceGraph, GraphContrainsTPUPartitionedCall) {
FunctionDef fd = FunctionDefHelper::Define(
// Name
"XTimesTwoFloat",
// Args
{"x: float"},
// Return values
{"y: float"},
// Attr def
{},
// Nodes
{
{{"two"},
"Const",
{},
{{"value", test::AsScalar<int32_t>(2)}, {"dtype", DT_INT64}}},
{{"scale"},
"Cast",
{"two"},
{{"SrcT", DT_INT64}, {"DstT", DT_FLOAT}}},
{{"y"}, "Mul", {"x", "scale"}, {{"T", DT_FLOAT}}},
});
tensorflow::set_tf2_execution(true);
FunctionDefLibrary flib;
*flib.add_function() = fd;
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kDirectSession);
Scope root = Scope::NewRootScope().ExitOnError();
Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
NameAttrList f_name_attr;
f_name_attr.set_name("XTimesTwoFloat");
ops::TPUPartitionedCall f(root.WithOpName("f"), {x}, /*device_ordinal=*/0,
{DT_FLOAT}, f_name_attr);
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_TRUE(IsInferenceGraph(graph, /*function_library=*/nullptr));
}
TEST(IsInferenceGraph, GraphDoesNotContrainTPUPartitionedCall) {
FunctionDef fd = FunctionDefHelper::Define(
// Name
"XTimesTwoFloat",
// Args
{"x: float"},
// Return values
{"y: float"},
// Attr def
{},
// Nodes
{
{{"two"},
"Const",
{},
{{"value", test::AsScalar<int32_t>(2)}, {"dtype", DT_INT64}}},
{{"scale"},
"Cast",
{"two"},
{{"SrcT", DT_INT64}, {"DstT", DT_FLOAT}}},
{{"y"}, "Mul", {"x", "scale"}, {{"T", DT_FLOAT}}},
});
tensorflow::set_tf2_execution(true);
FunctionDefLibrary flib;
*flib.add_function() = fd;
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kDirectSession);
Scope root = Scope::NewRootScope().ExitOnError();
Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
NameAttrList f_name_attr;
f_name_attr.set_name("XTimesTwoFloat");
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_FALSE(IsInferenceGraph(graph, /*function_library=*/nullptr));
}
TEST(IsInferenceGraph, FlibDefIsNotNullptrAndContainsTPUPartitionedCall) {
FunctionDef fd = FunctionDefHelper::Define(
// Name
"XTimesTwoFloat",
// Args
{"x: float"},
// Return values
{"y: float"},
// Attr def
{},
// Nodes
{
{{"two"},
"Const",
{},
{{"value", test::AsScalar<int32_t>(2)}, {"dtype", DT_INT64}}},
{{"scale"},
"Cast",
{"two"},
{{"SrcT", DT_INT64}, {"DstT", DT_FLOAT}}},
{{"y"}, "Mul", {"x", "scale"}, {{"T", DT_FLOAT}}},
{{"tpu_op"}, "TPUPartitionedCall", {}, {{"Tout", DT_FLOAT}}},
});
tensorflow::set_tf2_execution(true);
FunctionDefLibrary flib;
*flib.add_function() = fd;
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
Graph graph(flib_def);
graph.SetConstructionContext(ConstructionContext::kDirectSession);
Scope root = Scope::NewRootScope().ExitOnError();
Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
NameAttrList f_name_attr;
f_name_attr.set_name("XTimesTwoFloat");
TF_ASSERT_OK(root.ToGraph(&graph));
EXPECT_TRUE(IsInferenceGraph(graph, /*function_library=*/&flib_def));
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/mlir_pass_instrumentation.cc
+67
View File
@@ -0,0 +1,67 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/internal/mlir_pass_instrumentation.h"
#include <algorithm>
#include <functional>
#include <iterator>
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>
#include "absl/log/log.h"
#include "tensorflow/core/platform/logging.h"
namespace mlir {
class MlirPassInstrumentationRegistry {
public:
static MlirPassInstrumentationRegistry& Instance() {
static MlirPassInstrumentationRegistry* r =
new MlirPassInstrumentationRegistry;
return *r;
}
std::unordered_map<std::string,
std::function<std::unique_ptr<PassInstrumentation>()>>
instrumentors_;
};
void RegisterPassInstrumentor(
const std::string& name,
std::function<std::unique_ptr<PassInstrumentation>()> creator) {
MlirPassInstrumentationRegistry& r =
MlirPassInstrumentationRegistry::Instance();
auto result = r.instrumentors_.emplace(name, creator);
if (!result.second) {
VLOG(1) << "Duplicate MLIR pass instrumentor registration";
}
}
std::vector<std::function<std::unique_ptr<PassInstrumentation>()>>
GetPassInstrumentors() {
MlirPassInstrumentationRegistry& r =
MlirPassInstrumentationRegistry::Instance();
std::vector<std::function<std::unique_ptr<PassInstrumentation>()>> result;
result.reserve(r.instrumentors_.size());
std::transform(r.instrumentors_.begin(), r.instrumentors_.end(),
std::back_inserter(result), [](auto v) { return v.second; });
return result;
}
} // namespace mlir
tensorflow/compiler/mlir/tf2xla/internal/mlir_pass_instrumentation.h
+36
View File
@@ -0,0 +1,36 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_MLIR_PASS_INSTRUMENTATION_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_MLIR_PASS_INSTRUMENTATION_H_
#include <functional>
#include <memory>
#include <string>
#include <vector>
#include "mlir/Pass/PassInstrumentation.h" // from @llvm-project
namespace mlir {
void RegisterPassInstrumentor(
const std::string& name,
std::function<std::unique_ptr<PassInstrumentation>()> creator);
std::vector<std::function<std::unique_ptr<PassInstrumentation>()>>
GetPassInstrumentors();
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_MLIR_PASS_INSTRUMENTATION_H_
tensorflow/compiler/mlir/tf2xla/internal/mlir_pass_instrumentation_test.cc
+112
View File
@@ -0,0 +1,112 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/internal/mlir_pass_instrumentation.h"
#include <cstddef>
#include <cstdint>
#include <memory>
#include <sstream>
#include <string>
#include <unordered_map>
#include <vector>
#include <gtest/gtest.h>
#include "absl/strings/match.h"
#include "absl/strings/string_view.h"
#include "tensorflow/compiler/mlir/tf2xla/api/v1/compile_mlir_util.h"
#include "tensorflow/core/lib/core/status_test_util.h"
namespace mlir {
namespace {
static const char* kTestInstrumentationName = "test-intrumentatron";
static const char* kTestInstrumentationSearch = "tf.Identity";
struct StringStream : public llvm::raw_ostream {
StringStream() { SetUnbuffered(); }
~StringStream() override = default;
uint64_t current_pos() const override { return 0; }
void write_impl(const char* ptr, size_t size) override {
ss.write(ptr, size);
}
std::stringstream ss;
};
class TestPassInstrumentation : public ::testing::Test {
public:
void SetPassThatChangedIdentity(absl::string_view pass_name) {
pass_that_changed_identity_ = pass_name;
}
absl::string_view GetPassThatChangedIdentity() {
return pass_that_changed_identity_;
}
private:
std::string pass_that_changed_identity_;
friend class TestInstrumentor;
};
class TestInstrumentor : public PassInstrumentation {
public:
explicit TestInstrumentor(TestPassInstrumentation* test) : test_(test) {}
private:
void runBeforePass(Pass* pass, Operation* op) override {
StringStream stream;
op->print(stream, mlir::OpPrintingFlags().useLocalScope());
ops_seen_by_pass_[pass] = stream.ss.str();
}
void runAfterPass(Pass* pass, Operation* op) override {
StringStream stream;
op->print(stream, mlir::OpPrintingFlags().useLocalScope());
if (!absl::StrContains(stream.ss.str(), kTestInstrumentationSearch) &&
absl::StrContains(ops_seen_by_pass_[pass],
kTestInstrumentationSearch)) {
test_->SetPassThatChangedIdentity(pass->getName().str());
}
}
private:
TestPassInstrumentation* test_;
std::unordered_map<mlir::Pass*, std::string> ops_seen_by_pass_;
};
TEST_F(TestPassInstrumentation, CreatedCalledAndSetsPassName) {
RegisterPassInstrumentor(kTestInstrumentationName, [&]() {
return std::make_unique<TestInstrumentor>(this);
});
constexpr char legalization[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main(%arg0: tensor<?xi32, #mhlo.type_extensions<bounds = [1]>>) -> tensor<?xi32, #mhlo.type_extensions<bounds = [1]>> {
%0 = "tf.Identity"(%arg0) : (tensor<?xi32, #mhlo.type_extensions<bounds = [1]>>) -> tensor<?xi32, #mhlo.type_extensions<bounds = [1]>>
func.return %0 : tensor<?xi32, #mhlo.type_extensions<bounds = [1]>>
}
})";
SetPassThatChangedIdentity("");
std::vector<::tensorflow::TensorShape> arg_shapes = {{1}};
auto compilation_result = tensorflow::XlaCompilationResult();
TF_EXPECT_OK(tensorflow::CompileSerializedMlirToXlaHlo(
legalization, arg_shapes, /*device_type=*/"XLA_TPU_JIT",
/*use_tuple_args=*/true, /*enable_op_fallback=*/false,
/*shape_determination_fns=*/{}, &compilation_result)
.status());
EXPECT_FALSE(GetPassThatChangedIdentity().empty());
}
} // namespace
} // namespace mlir
tensorflow/compiler/mlir/tf2xla/internal/node_order.cc
+122
View File
@@ -0,0 +1,122 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/internal/node_order.h"
#include <algorithm>
#include <cassert>
#include <functional>
#include <iterator>
#include <set>
#include <string>
#include <unordered_map>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "tensorflow/core/graph/algorithm.h"
#include "tensorflow/core/graph/graph.h"
namespace tensorflow {
void TopologicalOrdering(
const Graph& g, const std::function<void(Node*)>& emit,
const std::function<std::string(Node*)>& get_grouping_key) {
std::unordered_map<std::string, int> group_key_string_to_integer;
absl::flat_hash_map<Node*, int> node_to_group;
absl::flat_hash_map<Node*, int> remaining_incoming_nodes;
absl::flat_hash_map<Node*, int> node_to_position;
using Ready = std::vector<Node*>;
std::vector<Ready> group_members_that_are_ready;
std::set<int> groups_that_are_ready;
// Visit all nodes once, for initialization. It doesn't matter whether we use
// BFS or DFS.
int i = 0;
DFS(
g, [](Node*) {},
[&](Node* n) {
// Find which group this node belongs to.
std::string group_key_string = get_grouping_key(n);
auto entry = group_key_string_to_integer.try_emplace(
group_key_string, group_key_string_to_integer.size());
int group_key = entry.first->second;
node_to_position[n] = i++;
node_to_group[n] = group_key;
if (entry.second) {
group_members_that_are_ready.push_back({});
}
// Count the incoming nodes and store. Also remember nodes ("sources")
// that don't have any inputs.
auto in_nodes = n->in_nodes();
int num_incoming = std::distance(in_nodes.begin(), in_nodes.end());
remaining_incoming_nodes[n] = num_incoming;
if (num_incoming == 0) {
// NO_CDC: This array is max(group_key) + 1.
group_members_that_are_ready[group_key].push_back(n);
groups_that_are_ready.emplace(group_key);
}
},
[](const Node* n1, const Node* n2) { return n1->name() < n2->name(); });
assert(group_key_string_to_integer.size() ==
group_members_that_are_ready.size());
int num_nodes = remaining_incoming_nodes.size();
// We emit one node per step, thus we just run this as often as we have nodes.
int current_group = 0;
for (int i = 0; i < num_nodes; i++) {
if (groups_that_are_ready.find(current_group) ==
groups_that_are_ready.end()) {
current_group = *groups_that_are_ready.begin();
}
// NO_CDC: This array is max(group_key) + 1.
int size = group_members_that_are_ready[current_group].size();
assert(size);
// NO_CDC: This array is max(group_key) + 1.
Node* node = group_members_that_are_ready[current_group][--size];
// NO_CDC: This array is max(group_key) + 1.
group_members_that_are_ready[current_group].pop_back();
if (size == 0) {
groups_that_are_ready.erase(current_group);
}
// Emit the operation and make its results available.
emit(node);
auto out_nodes = node->out_nodes();
std::vector<Node*> nodes_sorted(out_nodes.begin(), out_nodes.end());
std::sort(nodes_sorted.begin(), nodes_sorted.end(), [&](Node* a, Node* b) {
return node_to_position[a] < node_to_position[b];
});
for (Node* out : nodes_sorted) {
remaining_incoming_nodes[out]--;
if (remaining_incoming_nodes[out] == 0) {
int group_key = node_to_group[out];
// NO_CDC: This array is max(group_key) + 1.
if (group_members_that_are_ready[group_key].empty()) {
groups_that_are_ready.emplace(group_key);
}
// NO_CDC: This array is max(group_key) + 1.
group_members_that_are_ready[group_key].push_back(out);
}
}
}
}
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/node_order.h
+51
View File
@@ -0,0 +1,51 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_NODE_ORDER_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_NODE_ORDER_H_
#include <functional>
#include <unordered_set>
#include <vector>
#include "absl/container/flat_hash_set.h"
#include "tensorflow/core/graph/algorithm.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/lib/gtl/array_slice.h"
namespace tensorflow {
struct GroupByDevice {
std::string operator()(const Node* node) const {
return node->requested_device();
}
};
// Performs a topological ordering of nodes.
// This has the property that any child node of a parent node p is emitted
// before p. A grouping function is used to break ties if multiple child nodes
// (of possibly different parents) are ready to be emitted at some point, which
// is when we prefer to stay in the current group. Remaining ties are broken by
// node name.
// The "emit" function is used for outputing the result, and is called once
// for each node.
// This algorithm is O(n * k * log k), with k the largest node degree.
void TopologicalOrdering(
const Graph& g, const std::function<void(Node*)>& emit,
const std::function<std::string(Node*)>& get_grouping_key);
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_NODE_ORDER_H_
tensorflow/compiler/mlir/tf2xla/internal/node_order_test.cc
+294
View File
@@ -0,0 +1,294 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/internal/node_order.h"
#include <string>
#include <utility>
#include <vector>
#include "xla/tsl/lib/core/status_test_util.h"
#include "tensorflow/core/common_runtime/graph_def_builder_util.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/graph/algorithm.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/graph/graph_def_builder.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace {
REGISTER_OP("TestParams").Output("o: float");
REGISTER_OP("TestInput").Output("a: float").Output("b: float");
REGISTER_OP("TestMul").Input("a: float").Input("b: float").Output("o: float");
REGISTER_OP("TestUnary").Input("a: float").Output("o: float");
REGISTER_OP("TestTwoOutputs").Output("a: float").Output("b: float");
REGISTER_OP("TestBinary")
.Input("a: float")
.Input("b: float")
.Output("o: float");
// Compares that the order of nodes in 'inputs' respects the
// pair orders described in 'ordered_pairs'.
bool ExpectBefore(
const std::vector<std::pair<std::string, std::string>>& ordered_pairs,
const std::vector<Node*>& inputs, std::string* error) {
for (const std::pair<std::string, std::string>& pair : ordered_pairs) {
const std::string& before_node = pair.first;
const std::string& after_node = pair.second;
bool seen_before = false;
bool seen_both = false;
for (const Node* node : inputs) {
if (!seen_before && after_node == node->name()) {
*error = std::string("Saw ") + after_node + std::string(" before ") +
before_node;
return false;
}
if (before_node == node->name()) {
seen_before = true;
} else if (after_node == node->name()) {
seen_both = seen_before;
break;
}
}
if (!seen_both) {
*error = std::string("didn't see either ") + before_node +
std::string(" or ") + after_node;
return false;
}
}
return true;
}
TEST(AlgorithmTest, TopologicalOrdering) {
GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
using namespace ::tensorflow::ops; // NOLINT
Node* n1 = SourceOp("TestParams", b.opts().WithName("n1"));
Node* n2 =
SourceOp("TestParams", b.opts().WithName("n2").WithControlInput(n1));
Node* n3 =
SourceOp("TestParams", b.opts().WithName("n3").WithControlInput(n2));
Node* n4 = BinaryOp("TestMul", n1, {n3, 0}, b.opts().WithName("n4"));
Node* n5 = BinaryOp("TestMul", n1, {n3, 0},
b.opts().WithName("n5").WithControlInput(n1));
Node* n6 = BinaryOp("TestMul", n2, {n3, 0}, b.opts().WithName("n6"));
n3->set_requested_device("a");
n4->set_requested_device("a");
n5->set_requested_device("b");
n6->set_requested_device("b");
Graph g(OpRegistry::Global());
TF_ASSERT_OK(GraphDefBuilderToGraph(b, &g));
std::vector<Node*> order;
TopologicalOrdering(g, [&](Node* n) { order.push_back(n); }, GroupByDevice());
std::vector<std::pair<std::string, std::string>> desired_order = {
{"n1", "n2"}, // because of control dependency
{"n2", "n3"}, // because of control dependency
{"n3", "n4"}, // because of NodeScorerDevice
{"n1", "n4"}, // data dependency
{"n1", "n5"}, // data dependency
{"n2", "n6"}, // data dependency
{"n3", "n4"}, // data dependency
{"n3", "n5"}, // data dependency
{"n3", "n6"}, // data dependency
};
std::string error;
EXPECT_TRUE(ExpectBefore(desired_order, order, &error)) << error;
}
TEST(AlgorithmTest, TopologicalOrderingOnShallowTree) {
GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
using namespace ::tensorflow::ops; // NOLINT
Node* n1 = SourceOp("TestParams", b.opts().WithName("n1").WithDevice("a"));
Node* n2 =
SourceOp("TestParams",
b.opts().WithName("n2").WithDevice("b").WithControlInput(n1));
Node* n3 =
SourceOp("TestParams",
b.opts().WithName("n3").WithDevice("c").WithControlInput(n2));
Node* n4 =
SourceOp("TestParams",
b.opts().WithName("n4").WithDevice("a").WithControlInput(n1));
Node* n5 =
SourceOp("TestParams",
b.opts().WithName("n5").WithDevice("b").WithControlInput(n2));
Node* n6 =
SourceOp("TestParams",
b.opts().WithName("n6").WithDevice("c").WithControlInput(n3));
Node* n7 =
SourceOp("TestParams",
b.opts().WithName("n7").WithDevice("a").WithControlInput(n4));
Node* n8 =
SourceOp("TestParams",
b.opts().WithName("n8").WithDevice("b").WithControlInput(n5));
Node* n9 =
SourceOp("TestParams",
b.opts().WithName("n9").WithDevice("c").WithControlInput(n6));
Graph g(OpRegistry::Global());
TF_ASSERT_OK(GraphDefBuilderToGraph(b, &g));
std::vector<Node*> order;
TopologicalOrdering(g, [&](Node* n) { order.push_back(n); }, GroupByDevice());
std::vector<Node*> desired_order = {
g.source_node(), n1, n4, n7, n2, n5, n8, n3, n6, n9, g.sink_node()};
for (int i = 0; i < desired_order.size(); i++) {
desired_order[i] = g.FindNodeId(desired_order[i]->id());
}
EXPECT_EQ(order, desired_order);
}
TEST(AlgorithmTest, TopologicalOrderingGivesTheSameResultIfCalledTwice) {
GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
using namespace ::tensorflow::ops; // NOLINT
SourceOp("TestParams", b.opts().WithName("n1"));
SourceOp("TestParams", b.opts().WithName("n2"));
SourceOp("TestParams", b.opts().WithName("n3"));
SourceOp("TestParams", b.opts().WithName("n4"));
SourceOp("TestParams", b.opts().WithName("n5"));
SourceOp("TestParams", b.opts().WithName("n6"));
SourceOp("TestParams", b.opts().WithName("n7"));
SourceOp("TestParams", b.opts().WithName("n8"));
SourceOp("TestParams", b.opts().WithName("n9"));
Graph g(OpRegistry::Global());
TF_ASSERT_OK(GraphDefBuilderToGraph(b, &g));
std::vector<Node*> order1;
std::vector<Node*> order2;
TopologicalOrdering(
g, [&](Node* n) { order1.push_back(n); },
[&](const Node* node) { return std::string("same"); });
TopologicalOrdering(
g, [&](Node* n) { order2.push_back(n); },
[&](const Node* node) { return std::string("same"); });
EXPECT_EQ(order1, order2);
}
TEST(AlgorithmTest, TopologicalOrderingOnChain) {
GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
using namespace ::tensorflow::ops; // NOLINT
Node* n1 = SourceOp("TestParams", b.opts().WithName("n1"));
Node* n2 = UnaryOp("TestUnary", n1, b.opts().WithName("n2"));
Node* n3 = UnaryOp("TestUnary", n2, b.opts().WithName("n3"));
Node* n4 = UnaryOp("TestUnary", n3, b.opts().WithName("n4"));
Node* n5 = UnaryOp("TestUnary", n4, b.opts().WithName("n5"));
Node* n6 = UnaryOp("TestUnary", n5, b.opts().WithName("n6"));
Graph g(OpRegistry::Global());
TF_ASSERT_OK(GraphDefBuilderToGraph(b, &g));
std::vector<Node*> order;
TopologicalOrdering(g, [&](Node* n) { order.push_back(n); }, GroupByDevice());
std::vector<Node*> desired_order = {g.source_node(), n1, n2, n3, n4, n5, n6,
g.sink_node()};
for (int i = 0; i < desired_order.size(); i++) {
desired_order[i] = g.FindNodeId(desired_order[i]->id());
}
EXPECT_EQ(order, desired_order);
}
TEST(AlgorithmTest, TopologicalOrderingOnMultipleOutputs) {
GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
using namespace ::tensorflow::ops; // NOLINT
Node* n1 = SourceOp("TestTwoOutputs", b.opts().WithName("n1"));
UnaryOp("TestUnary", {n1, 0}, b.opts().WithName("n2"));
UnaryOp("TestUnary", {n1, 1}, b.opts().WithName("n3"));
UnaryOp("TestUnary", {n1, 0}, b.opts().WithName("n4"));
UnaryOp("TestUnary", {n1, 1}, b.opts().WithName("n5"));
Graph g(OpRegistry::Global());
TF_ASSERT_OK(GraphDefBuilderToGraph(b, &g));
std::vector<Node*> order;
TopologicalOrdering(g, [&](Node* n) { order.push_back(n); }, GroupByDevice());
std::vector<std::pair<std::string, std::string>> desired_order = {
{"n1", "n2"},
{"n1", "n3"},
{"n1", "n4"},
{"n1", "n5"},
};
std::string error;
EXPECT_TRUE(ExpectBefore(desired_order, order, &error)) << error;
}
TEST(AlgorithmTest, TopologicalOrderingSameAsReversePostOrder) {
GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
using namespace ::tensorflow::ops; // NOLINT
Node* n = SourceOp("TestTwoOutputs", b.opts().WithName("n"));
Node* n0 = UnaryOp("TestUnary", {n, 0}, b.opts().WithName("n2"));
Node* n1 = UnaryOp("TestUnary", {n, 1}, b.opts().WithName("n1"));
UnaryOp("TestUnary", n0, b.opts().WithName("n1a"));
UnaryOp("TestUnary", n0, b.opts().WithName("n8a"));
UnaryOp("TestUnary", n0, b.opts().WithName("n2a"));
UnaryOp("TestUnary", n0, b.opts().WithName("n7a"));
UnaryOp("TestUnary", n1, b.opts().WithName("n1b"));
UnaryOp("TestUnary", n1, b.opts().WithName("n8b"));
UnaryOp("TestUnary", n1, b.opts().WithName("n2b"));
UnaryOp("TestUnary", n1, b.opts().WithName("n7b"));
UnaryOp("TestUnary", n0, b.opts().WithName("n3a"));
UnaryOp("TestUnary", n0, b.opts().WithName("n6a"));
UnaryOp("TestUnary", n0, b.opts().WithName("n4a"));
UnaryOp("TestUnary", n0, b.opts().WithName("n5a"));
UnaryOp("TestUnary", n1, b.opts().WithName("n3b"));
UnaryOp("TestUnary", n1, b.opts().WithName("n6b"));
UnaryOp("TestUnary", n1, b.opts().WithName("n4b"));
UnaryOp("TestUnary", n1, b.opts().WithName("n5b"));
Graph g(OpRegistry::Global());
TF_ASSERT_OK(GraphDefBuilderToGraph(b, &g));
std::vector<Node*> order;
TopologicalOrdering(g, [&](Node* n) { order.push_back(n); }, GroupByDevice());
std::vector<Node*> desired_order;
GetReversePostOrder(g, &desired_order, [](const Node* n1, const Node* n2) {
return n1->name() < n2->name();
});
EXPECT_EQ(desired_order, order);
}
TEST(AlgorithmTest, TopologicalOrderingWithEachDeviceUsedOnce) {
GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
using namespace ::tensorflow::ops; // NOLINT
SourceOp("TestParams", b.opts().WithName("n1").WithDevice("a"));
SourceOp("TestParams", b.opts().WithName("n2").WithDevice("b"));
SourceOp("TestParams", b.opts().WithName("n3").WithDevice("c"));
SourceOp("TestParams", b.opts().WithName("n4").WithDevice("d"));
Graph g(OpRegistry::Global());
TF_ASSERT_OK(GraphDefBuilderToGraph(b, &g));
int count = 0;
TopologicalOrdering(g, [&](Node* n) { count++; }, GroupByDevice());
EXPECT_EQ(count, 6);
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/passes/BUILD
+573
View File
@@ -0,0 +1,573 @@
load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library")
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow:tensorflow.default.bzl", "filegroup", "get_compatible_with_portable")
load("//tensorflow/compiler/mlir:glob_lit_test.bzl", "glob_lit_tests")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//learning/brain/tfrt/tpu/compiler/mlir:__pkg__",
"//tensorflow/compiler/mlir:__subpackages__",
"//tensorflow/compiler/mlir/tf2xla/api:__subpackages__",
"//tensorflow/compiler/mlir/tf2xla/internal:__subpackages__",
],
licenses = ["notice"],
)
cc_library(
name = "clustering_passes",
hdrs = [
"clustering_passes.h",
],
textual_hdrs = [
"clustering_passes.h.inc",
],
deps = [
":extract_head_tail_outside_compilation",
":extract_outside_compilation",
":hoist_broadcast_read",
":mark_ops_for_outside_compilation",
":tpu_cluster_formation",
":tpu_sharding_identification_pass",
":tpu_validate_inputs",
":tpu_validate_session_inputs",
":verify_clustering_pass",
":xla_broadcast",
":xla_cluster_formation",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
],
)
cc_library(
name = "verify_clustering_pass",
srcs = [
"verify_clustering_pass.cc",
],
deps = [
":clustering_passes_inc_gen",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:string_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_analysis",
"//tensorflow/compiler/mlir/tensorflow:tpu_rewrite_device_util",
"//tensorflow/compiler/mlir/tf2xla/internal/utils:dialect_detection_utils",
"//tensorflow/core:framework",
"//tensorflow/core/transforms/toposort:Pass",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:Transforms",
],
)
gentbl_cc_library(
name = "clustering_passes_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"clustering_passes.h.inc": [
"-gen-pass-decls",
"-name=TFXLABridgeClustering",
]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "clustering_passes.td",
deps = [
"@llvm-project//mlir:PassBaseTdFiles",
],
)
tf_cc_test(
name = "verify_clustering_pass_test",
srcs = ["verify_clustering_pass_test.cc"],
deps = [
":clustering_passes",
"//tensorflow/compiler/mlir/tf2xla/transforms:test_utils",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@xla//xla/tsl/platform:statusor",
],
)
glob_lit_tests(
name = "all_tests",
data = [":test_utilities"],
driver = "@llvm-project//mlir:run_lit.sh",
test_file_exts = [
"mlir",
],
)
# Bundle together all of the test utilities that are used by tests.
filegroup(
name = "test_utilities",
testonly = True,
data = [
"//tensorflow/compiler/mlir:tf-opt",
"@llvm-project//llvm:FileCheck",
],
)
cc_library(
name = "tpu_cluster_formation",
srcs = ["tpu_cluster_formation.cc"],
textual_hdrs = [
"clustering_passes.h.inc",
],
deps = [
":clustering_passes_inc_gen",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:string_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_analysis",
"//tensorflow/compiler/mlir/tensorflow:tpu_rewrite_device_util",
"//tensorflow/compiler/mlir/tensorflow/transforms:tf_pass_inc_gen",
"//tensorflow/core:framework",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
],
)
cc_library(
name = "extract_outside_compilation",
srcs = ["extract_outside_compilation.cc"],
textual_hdrs = [
"clustering_passes.h.inc",
],
deps = [
":clustering_passes_inc_gen",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:device_util",
"//tensorflow/compiler/mlir/tensorflow:serialize_mlir_module_utils",
"//tensorflow/compiler/mlir/tensorflow:string_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_analysis",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow:tpu_rewrite_device_util",
"//tensorflow/compiler/mlir/tensorflow/transforms:shape_inference_pass",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Rewrite",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
],
)
cc_library(
name = "extract_head_tail_outside_compilation",
srcs = ["extract_head_tail_outside_compilation.cc"],
textual_hdrs = [
"clustering_passes.h.inc",
],
deps = [
":clustering_passes_inc_gen",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:device_util",
"//tensorflow/compiler/mlir/tensorflow:string_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_analysis",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow:tpu_rewrite_device_util",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Rewrite",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
],
)
cc_library(
name = "mlir_to_graph_passes",
hdrs = [
"mlir_to_graph_passes.h",
],
textual_hdrs = [
"mlir_to_graph_passes.h.inc",
],
deps = [
":verify_input_dialect_to_executor_pass",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:Pass",
],
)
gentbl_cc_library(
name = "mlir_to_graph_passes_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"mlir_to_graph_passes.h.inc": [
"-gen-pass-decls",
"-name=TFXLABridgeMlirToGraph",
]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "mlir_to_graph_passes.td",
deps = [
"@llvm-project//mlir:PassBaseTdFiles",
],
)
cc_library(
name = "verify_input_dialect_to_executor_pass",
srcs = [
"verify_input_dialect_to_executor_pass.cc",
],
deps = [
":mlir_to_graph_passes_inc_gen",
"//tensorflow/compiler/mlir/tf2xla/internal/utils:dialect_detection_utils",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
],
)
cc_library(
name = "xla_cluster_formation",
srcs = ["xla_cluster_formation.cc"],
textual_hdrs = [
"clustering_passes.h.inc",
],
deps = [
":clustering_passes_inc_gen",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:call_graph_util",
"//tensorflow/compiler/mlir/tensorflow:cluster_util",
"//tensorflow/compiler/mlir/tensorflow:string_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_analysis",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow:tpu_rewrite_device_util",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core/platform:types",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
],
)
cc_library(
name = "mark_ops_for_outside_compilation",
srcs = ["mark_ops_for_outside_compilation.cc"],
textual_hdrs = [
"clustering_passes.h.inc",
],
deps = [
":clustering_passes_inc_gen",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:string_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_analysis",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow:tpu_rewrite_device_util",
"//tensorflow/compiler/mlir/tensorflow/transforms:lower_tf_lib",
"//tensorflow/compiler/mlir/tensorflow/transforms:tf_pass_inc_gen",
"//tensorflow/compiler/mlir/tensorflow/transforms:verify_no_outside_compilation_markers_pass",
"//tensorflow/compiler/mlir/tf2xla/transforms:legalization_op_config",
"//tensorflow/compiler/mlir/tf2xla/transforms:legalize_tf",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Rewrite",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
],
)
cc_library(
name = "tpu_sharding_identification_pass",
srcs = ["tpu_sharding_identification_pass.cc"],
textual_hdrs = [
"clustering_passes.h.inc",
],
deps = [
":clustering_passes_inc_gen",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:string_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_analysis",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_traits",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow:tpu_rewrite_device_util",
"//tensorflow/compiler/mlir/tensorflow:xla_sharding_util",
"//tensorflow/compiler/mlir/tensorflow/transforms:lower_tf_lib",
"//tensorflow/compiler/mlir/tensorflow/transforms:tf_pass_inc_gen",
"//tensorflow/compiler/mlir/tensorflow/transforms:verify_no_outside_compilation_markers_pass",
"//tensorflow/compiler/mlir/tf2xla/transforms:legalization_op_config",
"//tensorflow/compiler/mlir/tf2xla/transforms:legalize_tf",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Rewrite",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
"@xla//xla:xla_data_proto_cc",
"@xla//xla/hlo/builder:sharding_builder",
],
)
cc_library(
name = "hoist_broadcast_read",
srcs = ["hoist_broadcast_read.cc"],
textual_hdrs = [
"clustering_passes.h.inc",
],
deps = [
":clustering_passes_inc_gen",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:string_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_analysis",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow:tpu_rewrite_device_util",
"//tensorflow/compiler/mlir/tensorflow/transforms:lower_tf_lib",
"//tensorflow/compiler/mlir/tensorflow/transforms:tf_pass_inc_gen",
"//tensorflow/compiler/mlir/tensorflow/transforms:verify_no_outside_compilation_markers_pass",
"//tensorflow/compiler/mlir/tf2xla/transforms:legalization_op_config",
"//tensorflow/compiler/mlir/tf2xla/transforms:legalize_tf",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Rewrite",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
],
)
cc_library(
name = "xla_broadcast",
srcs = ["xla_broadcast.cc"],
textual_hdrs = [
"clustering_passes.h.inc",
],
deps = [
":clustering_passes_inc_gen",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:device_util",
"//tensorflow/compiler/mlir/tensorflow:string_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_analysis",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow:tpu_rewrite_device_util",
"//tensorflow/compiler/mlir/tensorflow:xla_rewrite_util",
"//tensorflow/compiler/mlir/tensorflow:xla_sharding_util",
"//tensorflow/compiler/mlir/tensorflow/transforms:lower_tf_lib",
"//tensorflow/compiler/mlir/tensorflow/transforms:tf_pass_inc_gen",
"//tensorflow/compiler/mlir/tensorflow/transforms:verify_no_outside_compilation_markers_pass",
"//tensorflow/compiler/mlir/tf2xla/transforms:legalization_op_config",
"//tensorflow/compiler/mlir/tf2xla/transforms:legalize_tf",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core/ir/types:Dialect",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Rewrite",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
],
)
cc_library(
name = "lowering_passes",
hdrs = [
"lowering_passes.h",
],
textual_hdrs = [
"lowering_passes.h.inc",
],
deps = [
":input_metrics_lowering_pass",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
],
)
gentbl_cc_library(
name = "lowering_passes_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"lowering_passes.h.inc": [
"-gen-pass-decls",
"-name=TFXLABridgeLowering",
]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "lowering_passes.td",
deps = [
"@llvm-project//mlir:PassBaseTdFiles",
],
)
cc_library(
name = "input_metrics_lowering_pass",
srcs = [
"input_lowering_metrics_pass.cc",
],
deps = [
":lowering_passes_inc_gen",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tf2xla/transforms:legalization_op_config",
"//tensorflow/core:lib",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
],
)
tf_cc_test(
name = "input_metrics_lowering_pass_test",
srcs = ["input_lowering_metrics_pass_test.cc"],
deps = [
":lowering_passes",
"//tensorflow/compiler/mlir/tf2xla/transforms:test_utils",
"//tensorflow/core/lib/monitoring:cell_reader",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@xla//xla/tsl/platform:statusor",
],
)
cc_library(
name = "tpu_validate_inputs_utils",
srcs = ["tpu_validate_inputs_utils.cc"],
hdrs = ["tpu_validate_inputs_utils.h"],
deps = [
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "tpu_validate_session_inputs",
srcs = ["tpu_validate_session_inputs.cc"],
textual_hdrs = [
"clustering_passes.h.inc",
],
deps = [
":clustering_passes_inc_gen",
":tpu_validate_inputs_utils",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
"@com_google_absl//absl/log",
"@com_google_absl//absl/strings",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
],
)
cc_library(
name = "tpu_validate_inputs",
srcs = ["tpu_validate_inputs.cc"],
textual_hdrs = [
"clustering_passes.h.inc",
],
deps = [
":clustering_passes_inc_gen",
":tpu_validate_inputs_utils",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
"//tensorflow/compiler/mlir/tensorflow:string_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_analysis",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow:tpu_rewrite_device_util",
"//tensorflow/compiler/mlir/tensorflow/transforms:tf_pass_inc_gen",
"//tensorflow/core:framework",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
"@xla//xla:xla_data_proto_cc",
"@xla//xla/hlo/builder:sharding_builder",
],
)
tf_cc_test(
name = "tpu_validate_inputs_utils_test",
srcs = ["tpu_validate_inputs_utils_test.cc"],
deps = [
":tpu_validate_inputs_utils",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tf2xla/transforms:test_utils",
"@com_google_googletest//:gtest_main",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
],
)
tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h
+93
View File
@@ -0,0 +1,93 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_PASSES_CLUSTERING_PASSES_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_PASSES_CLUSTERING_PASSES_H_
#include <memory>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
namespace tensorflow {
namespace tf2xla {
namespace internal {
// Verifies that all MLIR Ops have the expected attributes.
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateVerifyClusteringPass();
// Creates a pass that forms clusters from operations of the same
// `_replication_info` attribute.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateTPUClusterFormationPass(bool strict_clusters = false);
// Creates a pass that extracts outside compilation (Host ops inside device
// cluster) at head/tail of Device cluster to run before/after XLA computation.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateExtractHeadTailOutsideCompilationPass();
// Creates a pass that extract outside compilation (Host ops inside cevice
// cluster) ops to a separate parallel_execute region to run on CPU.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateExtractOutsideCompilationPass();
// Create a pass that encapsulates StatefulPartitionedCallOp within a cluster.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateXlaClusterFormationPass();
// Creates a pass that marks unsupported ops in device cluster for outside
// compilation.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateMarkOpsForOutsideCompilationPass();
// Creates a pass that hoists reads out of a replicate that are on a variable
// whose value is broacast to all replicas.
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateHoistBroadcastReadPass();
// Creates a pass that moves broadcasts from TF host ops to XLA code, encoded as
// XlaAllReduces. This enables use of the device network for broadcasts, which
// is faster.
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateXlaBroadcastPass();
// Creates a pass that identifies XLASharding ops in launch op for TPU
// computation.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateTPUShardingIdentificationPass();
// Creates a pass that validates the inputs to a TPU computation.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateTPUValidateSessionInputsPass();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateTPUValidateInputsPass();
#define GEN_PASS_REGISTRATION
#define GEN_PASS_DECL_MARKOPSFOROUTSIDECOMPILATIONPASS
#define GEN_PASS_DECL_TPUCLUSTERFORMATIONPASS
#define GEN_PASS_DECL_TPUEXTRACTHEADTAILOUTSIDECOMPILATIONPASS
#define GEN_PASS_DECL_TPUEXTRACTOUTSIDECOMPILATIONPASS
#define GEN_PASS_DECL_TPUSHARDINGIDENTIFICATIONPASS
#define GEN_PASS_DECL_TPUVALIDATEINPUTSPASS
#define GEN_PASS_DECL_TPUVALIDATESESSIONINPUTSPASS
#define GEN_PASS_DECL_VERIFYCLUSTERINGPASS
#define GEN_PASS_DECL_XLACLUSTERFORMATIONPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc"
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_PASSES_CLUSTERING_PASSES_H_
tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.td
+466
View File
@@ -0,0 +1,466 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
include "mlir/Pass/PassBase.td"
def VerifyClusteringPass : Pass<"verify-clustering-pass", "mlir::func::FuncOp"> {
let summary = "Verify that the Bridge output is correct and errors if verification fails.";
let description = [{
Verifies whether clustering has resulted in the expected invariants. These
include verifying that clusters have been created and have been outside
compiled, the result is device agnostic and in TF functional dialect &
that the device attribute exists.
}];
let constructor = "tensorflow::tf2xla::internal::CreateVerifyClusteringPass()";
}
def TPUClusterFormationPass : Pass<"tf-tpu-cluster-formation", "ModuleOp"> {
let summary = "Forms clusters from operations assigned to the same TPU computation";
let description = [{
TPU computations from the frontend are composed of a `tf.TPUReplicateMetadata`
op, a subgraph of ops (TensorFlow Dialect) each with a matching
`_replication_info` attribute relative to the associated
`tf.TPUReplicateMetadata` op, and optionally `tf.TPUReplicatedInput` and
`tf.TPUReplicatedOutput` ops feeding in inputs and outputs to and from a
replicated TPU computation. The number of times a TPU computation is
replicated is defined in the `tf.TPUReplicateMetadata` op (`num_replicas`
attribute) and operand and result sizes of `tf.TPUReplicatedInput` and
`tf.TPUReplicatedOutput` respectively must match, excluding packed tensors.
It is also assumed ops of the same TPU computation do not have ops outside
of the TPU computation that are both inputs and outputs to the same TPU
computation. Furthermore, we assume that every node has either none or both
of `_replication_info` and `_xla_compile_device_type` attributes defined.
This pass takes the TPU computation subgraph, moves them into a
`tf_device.cluster`, and copies over attributes from the associated
`tf.TPUReplicateMetadata` op to the newly created `tf_device.cluster`. If the
computation is replicated (`num_replicas` > 1), the `num_replicas` attribute is
not copied over but instead the `tf_device.cluster` is further wrapped with a
`tf_device.replicate`, and associated `tf.TPUReplicatedInput` and
`tf.TPUReplicatedOutput` ops are replaced as the `tf_device.replicate` operands
and results. Otherwise, the single operands and results of the associated
`tf.TPUReplicatedInput` and `tf.TPUReplicatedOutput` ops are simply forwarded to
the `tf_device.cluster`.
For example, the following non replicated computation:
```mlir
func @tpu_computation(%arg0: tensor<i32>) -> tensor<i32> {
// Metadata op for cluster `cluster` with 1 replica, 1 core per replica and
// with topology `<topology>`.
"tf.TPUReplicateMetadata"() {_xla_compile_device_type = "TPU", _replication_info = "cluster", num_relicas = 1, num_cores_per_replica = 1, topology = "<topology>", device_assignment = [], padding_map = []} : () -> ()
%replicated_input = "tf.TPUReplicatedInput"(%arg0) : (tensor<i32>) -> tensor<i32>
%identity = "tf.Identity"(%replicated_input) {_xla_compile_device_type = "TPU", _replication_info = "cluster"} : (tensor<i32>) -> tensor<i32>
%replicated_output = "tf.TPUReplicatedOutput(%identity) : (tensor<i32>) -> tensor<i32>
return %replicated_output : tensor<i32>
}
```
will be transformed into:
```mlir
func @tpu_computation(%arg0: tensor<i32>) -> tensor<i32> {
%cluster = "tf_device.cluster"() ( {
%identity = "tf.Identity"(%arg0) : (tensor<i32>) -> tensor<i32>
tf_device.return %identity : tensor<i32>
}) {_xla_compile_device_type = "TPU", _replication_info = "cluster", num_cores_per_replica = 1, topology = "topology", device_assignment = [], padding_map = []} : () -> (tensor<i32>)
return %cluster : tensor<i32>
}
```
The following replicated computation:
```mlir
func @tpu_computation(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
"tf.TPUReplicateMetadata"() {_xla_compile_device_type = "TPU", _replication_info = "cluster", num_relicas = 2, num_cores_per_replica = 1, topology = "topology", device_assignment = [], padding_map = []} : () -> ()
%replicated_input = "tf.TPUReplicatedInput"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
%identity = "tf.Identity"(%replicated_input) {_xla_compile_device_type = "TPU", _replication_info = "cluster"} : (tensor<i32>) -> tensor<i32>
%replicated_output:2 = "tf.TPUReplicatedOutput(%identity) : (tensor<i32>) -> (tensor<i32>, tensor<i32>)
return %replicated_output#0, %replicated_output#1 : tensor<i32>, tensor<i32>
}
```
will be transformed into:
```mlir
func @tpu_computation(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
%replicate:2 = tf_device.replicate([%arg0, %arg1] as %replicated_input) {n = 2 : i32} {
%cluster = "tf_device.cluster"() ( {
%identity = "tf.Identity"(%replicated_input) : (tensor<i32>) -> tensor<i32>
tf_device.return %identity : tensor<i32>
}) {_xla_compile_device_type = "TPU", _replication_info = "cluster", num_cores_per_replica = 1, topology = "topology", device_assignment = [], padding_map = []} : () -> (tensor<i32>)
tf_device.return %cluster : tensor<i32>
}
return %replicate#0, %replicate#1 : tensor<i32>, tensor<i32>
}
```
}];
let constructor = "tensorflow::tf2xla::internal::CreateTPUClusterFormationPass()";
}
def ExtractHeadTailOutsideCompilationPass : Pass<"tf-extract-head-tail-outside-compilation", "ModuleOp"> {
let summary = "Extracts head or tail outside compilation to separate host launches before/after device cluster.";
let description = [{
This pass extracts a CPU computation cluster with `_xla_outside_compilation`
annotation from the head or tail of a Device cluster.
For example:
```mlir
%cluster = "tf_device.cluster"() ( {
%a = "tf.A"(%arg0) {_xla_outside_compilation = "cluster1"} : (tensor<i32>) -> tensor<i32>
%b = "tf.B"(%a) : (tensor<i32>) -> tensor<i32>
%c = "tf.C"(%b) {_xla_outside_compilation = "cluster1"} : (tensor<i32>) -> tensor<i32>
tf_device.return %c : tensor<i32>
}) {num_cores_per_replica = 1, step_marker_location = "", padding_map = [], topology = "", device_assignment = []} : () -> tensor<i32>
return %cluster : tensor<i32>
```
becomes:
```mlir
%0 = "tf_device.launch"() ( {
%3 = "tf.A"(%arg0) : (tensor<i32>) -> tensor<i32>
tf_device.return %3 : tensor<i32>
}) {device = "/job:worker/replica:0/task:0/device:CPU:0"} : () -> tensor<i32>
%1 = "tf_device.cluster"() ( {
%3 = "tf.B"(%0) : (tensor<i32>) -> tensor<i32>
tf_device.return %3 : tensor<i32>
}) {device_assignment = [], num_cores_per_replica = 1 : i64, padding_map = [], step_marker_location = "", topology = ""} : () -> tensor<i32>
%2 = "tf_device.launch"() ( {
%3 = "tf.C"(%1) : (tensor<i32>) -> tensor<i32>
tf_device.return %3 : tensor<i32>
}) {device = "/job:worker/replica:0/task:0/device:CPU:0"} : () -> tensor<i32>
return %2 : tensor<i32>
```
}];
let constructor = "tensorflow::tf2xla::internal::CreateExtractHeadTailOutsideCompilationPass()";
}
def ExtractOutsideCompilationPass : Pass<"tf-extract-outside-compilation", "ModuleOp"> {
let summary = "Extracts device outside compilation computation to a separate tf_device.parallel_execute region.";
let description = [{
This pass extracts a CPU computation cluster with `_xla_outside_compilation`
annotation, which denotes ops that should be run on CPU/host, from a device cluster.
Each outside compilation cluster is moved to
a tf_device.parallel_execute region. The device cluster is also moved to a
tf_device.parallel_execute region. Communication ops between device and host are
added to pass inputs/outputs to/from the outside compiled region.
For example, the following tf_device.cluster with an op marked for `xla_outside_compilation`:
```mlir
func @outside_compilation() -> tensor<f32> {
%0 = "tf_device.cluster"() ( {
%1 = "tf.Const"() {_xla_outside_compilation = "0", value = dense<1.0> : tensor<f32>} : () -> (tensor<f32>)
%2 = "tf.Identity"(%1) {_xla_outside_compilation = "0"} : (tensor<f32>) -> (tensor<f32>)
%3 = "tf.AddV2"(%1, %2) : (tensor<f32>, tensor<f32>) -> (tensor<f32>)
tf_device.return %3 : tensor<f32>
}) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> tensor<f32>
return %0 : tensor<f32>
}
```
will become a tf_device.parallel_execute op with a CPU/host region and
a tf_device.cluster with communication ops to send data to/from device/host:
```mlir
func @outside_compilation() -> tensor<f32> {
%0 = "tf_device.parallel_execute"() ( {
"tf_device.launch"() ( {
%1 = "tf._XlaCompileMlirPlaceholderProgramKey"() : () -> tensor<3x!tf_type.string>
%2 = "tf._XlaRecvAtHost"(%1) {device_ordinal = 0 : i64, key = "host_compute_channel_0_0_args"} : (tensor<3x!tf_type.string>) -> tensor<f32>
%3 = "tf.Identity"(%2) : (tensor<f32>) -> tensor<f32>
"tf._XlaSendFromHost"(%3, %1) {device_ordinal = 0 : i64, key = "host_compute_channel_0_0_retvals"} : (tensor<f32>, tensor<3x!tf_type.string>) -> ()
tf_device.return
}) {device = "/job:worker/replica:0/task:0/device:CPU:0"} : () -> ()
tf_device.return
}, {
%1 = "tf_device.cluster"() ( {
%2 = "tf.Const"() {value = dense<1.000000e+00> : tensor<f32>} : () -> tensor<f32>
%3 = "tf._XlaHostComputeMlir"(%2) {recv_key = "host_compute_channel_0_0_retvals", send_key = "host_compute_channel_0_0_args", tpu_core = 0 : i64} : (tensor<f32>) -> tensor<f32>
%4 = "tf.AddV2"(%2, %3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
tf_device.return %4 : tensor<f32>
}) {device_assignment = [], num_cores_per_replica = 1 : i64, topology = ""} : () -> tensor<f32>
tf_device.return %1 : tensor<f32>
}) : () -> tensor<f32>
return %0 : tensor<f32>
}
```
}];
let constructor = "tensorflow::tf2xla::internal::CreateExtractOutsideCompilationPass()";
}
def XlaClusterFormationPass : Pass<"tf-xla-cluster-formation", "ModuleOp"> {
let summary = "Encapsulate partitioned calls within a Cluster op";
let description = [{
This pass clusters `tf.PartitionedCall` and `tf.StatefulPartitionedCall`
with `_xla_compile_device_type` attribute into a `tf_device.cluster`.
Notice this pass will only rewrite the outermost call if there are nested
calls to avoid nested `tf.XlaLaunch` operations from being created later.
For example, the following code
```mlir
func.func @main() -> tensor<i32> {
%0 = "tf.StatefulPartitionedCall"() {_xla_compile_device_type = "CPU", f = @stateful_pcall_func} : () -> (tensor<i32>)
func.return %0 : tensor<i32>
}
func.func @stateful_pcall_func() -> tensor<i32> {
%0 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
func.return %0 : tensor<i32>
}
```
will be transformed into,
```mlir
func.func @main() -> tensor<i32> {
%0 = "tf_device.cluster"() ({
%1 = "tf.StatefulPartitionedCall"() {_xla_compile_device_type = "CPU", f = @stateful_pcall_func} : () -> tensor<i32>
tf_device.return %1 : tensor<i32>
}) : () -> tensor<i32>
func.return %0 : tensor<i32>
}
func.func @stateful_pcall_func() -> tensor<i32> {
%0 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
func.return %0 : tensor<i32>
}
```
}];
let constructor = "tensorflow::tf2xla::internal::CreateXlaClusterFormationPass()";
let dependentDialects = ["mlir::tf_device::TensorFlowDeviceDialect"];
}
def MarkOpsForOutsideCompilationPass : Pass<"tf-mark-ops-for-outside-compilation", "ModuleOp"> {
let summary = "Marks ops in device cluster for outside compilation if they are unsupported on device.";
let description = [{
This pass marks unsupported ops in a device cluster with
`_xla_outside_compilation` attribute so the operations will run on the host
instead of the device. Unsupported ops are ops that can not be code
generated to run on the device for the cluster including:
1. String operations on TPUs.
2. Operations that don't have a kernel defined for the device.
This pass is conservative in that it will mark all ops for outside compilation
that can not be compiled for the device. Exceptions for this are added for ops
that will be rewritten or decomposed before compiling on device.
For example, tf_device.cluster op with an unsupported op, tf.UnsupportedOp:
```mlir
func @unsupported_op() -> tensor<i32> {
%0 = "tf_device.cluster"() ( {
%1 = "tf.UnsupportedOp"() : () -> tensor<i32>
%2 = "tf.Identity"(%1) : (tensor<i32>) -> tensor<i32>
tf_device.return %2 : tensor<i32>
}) {allow_soft_placement = true, num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> tensor<i32>
return %0 : tensor<i32>
}
```
will mark tf.UnsupportedOp with `_xla_outside_compilation` attribute:
```mlir
func @unsupported_op() -> tensor<i32> {
%0 = "tf_device.cluster"() ( {
%1 = "tf.UnsupportedOp"() {_xla_outside_compilation = "auto0"} : () -> tensor<i32>
%2 = "tf.Identity"(%1) : (tensor<i32>) -> tensor<i32>
tf_device.return %2 : tensor<i32>
}) {allow_soft_placement = true, device_assignment = [], num_cores_per_replica = 1 : i64, topology = ""} : () -> tensor<i32>
return %0 : tensor<i32>
}
```
}];
let constructor = "tensorflow::tf2xla::internal::CreateMarkOpsForOutsideCompilationPass()";
}
def HoistBroadcastReadPass : Pass<"tf-hoist-broadcast-read", "mlir::func::FuncOp"> {
let summary = "Hoist reads out of a replicate that are on a resource that is broacast to all replicas.";
let description = [{
Some `ReadVariableOp`s that are within a `tf_device.replicate` read the same
value across all replicas. These reads can be hoisted out of the
`tf_device.replicate` so there's one read for all replicas, and each replica
depends on the result of the read. This transform enables the
xla-broadcast-pass to optimize the broadcast value.
For example, the following:
```mlir
tf_device.replicate {n = 2 : i32} {
%0 = "tf.ReadVariableOp"(%arg0) : (tensor<*x!tf_type.resource<tensor<f32>>>) -> tensor<f32>
"tf.OpA"(%0) : (tensor<f32>) -> ()
}
```
will be transformed into:
``mlir
%0 = "tf.ReadVariableOp"(%arg0) : (tensor<*x!tf_type.resource<tensor<f32>>>) -> tensor<f32>
tf_device.replicate {n = 2 : i32} {
"tf.OpA"(%0) : (tensor<f32>) -> ()
}
```
We must ensure that there is a single underlying resource that not
distributed across replicas. There is a single underlying resource when the
resource device type is CPU, so we cautiously only apply in this case.
To be cautious we never hoist a read that comes after a write to the same
resource.
}];
let constructor = "tensorflow::tf2xla::internal::CreateHoistBroadcastReadPass()";
}
def XlaBroadcastPass : Pass<"tf-xla-broadcast", "mlir::func::FuncOp"> {
let summary = "Moves a broadcast from host into XLA, encoded as XlaAllReduce";
let description = [{
This pass moves brodcasts from host TF ops into XLA. This enables use of
the inter-device network, which is faster than the inter-host network.
Broadcasts in XLA are encoded as XlaAllReduce. An all_reduce with all 0
inputs except for one real input produces the same result as a broadcast.
For example, the following:
```mlir
func.func @main(%arg0: tensor<f32>) -> () {
tf_device.replicate {n = 2 : i32} {
"tf_device.cluster"() ({
"tf.OpA"(%arg0) : (tensor<f32>) -> ()
tf_device.return
}) : () -> ()
}
func.return
}
```
will be transformed into
```mlir
func.func @main(%arg0: tensor<f32>) {
%elem = "tf.Const"() <{value = dense<0.000000e+00> : tensor<f32>}> : () -> tensor<f32>
%shape = "tf.Const"() <{value = dense<> : tensor<0xi64>}> : () -> tensor<0xi64>
%zero = "tf.Fill"(%shape, %elem) : (tensor<0xi64>, tensor<f32>) -> tensor<f32>
tf_device.replicate([%arg0, %zero] as %arg1: tensor<f32>) {n = 2 : i32} {
%arg1_id = "tf_device.launch"() <{device = "TPU_REPLICATED_HOST_0"}> ({
%1 = "tf.Identity"(%arg1) : (tensor<f32>) -> tensor<f32>
tf_device.return %1 : tensor<f32>
}) : () -> tensor<f32>
"tf_device.cluster"() ({
%group = "tf.Const"() <{value = dense<[[0, 1]]> : tensor<1x2xi32>}> : () -> tensor<1x2xi32>
%arg1_reduced = "tf.XlaAllReduce"(%arg1_id, %group) <{mode = "CrossReplica", reduce_op = "Add"}> : (tensor<f32>, tensor<1x2xi32>) -> tensor<f32>
"tf.OpA"(%arg1_reduced) : (tensor<f32>) -> ()
tf_device.return
}) : () -> ()
tf_device.return
}
return
}
```
}];
let constructor = "tensorflow::tf2xla::internal::CreateXlaBroadcastPass()";
let dependentDialects = ["mlir::tf_device::TensorFlowDeviceDialect"];
}
def TPUShardingIdentificationPass : Pass<"tf-tpu-sharding-identification", "ModuleOp"> {
let summary = "Identifies and handles inputs/outputs of TPU computation that is "
"sharded across logical cores.";
let constructor = "tensorflow::tf2xla::internal::CreateTPUShardingIdentificationPass()";
let description = [{
Bubbles up sharding configuration from `cluster_func` regions into
the attributes of `cluster_func`. This is done by parsing the
`XlaSharding` / `TPUPartitionedOutput` / `TPUPartitionedInput` ops inside
`cluster_func`.
For example, given the following `cluster_func` wrapping `func`:
```mlir
func @test(%arg0: tensor<*xi32>) {
"tf_device.cluster_func"(%arg0) {
func = @func,
step_marker_location = ""} : (tensor<*xi32>) -> tensor<*xi32>
return
}
func @func(%arg0: tensor<*xi32>) -> tensor<*xi32> {
%0 = "tf.XlaSharding"(%arg0) {_XlaSharding = "\01\02\03",
sharding = "\01\02\03"} : (tensor<*xi32>) -> tensor<*xi32>
%1 = "tf.A"(%0) : (tensor<*xi32>) -> (tensor<*xi32>)
return %1 : tensor<*xi32>
}
```
Now, cluster_func receives the following `*_sharding_configuration`
attributes, and `func` receives the mhlo.sharding attribute:
```mlir
func @test(%arg0: tensor<*xi32>) {
%0 = "tf_device.cluster_func"(%arg0) {
func = @func,
input_sharding_configuration = ["\01\02\03"],
output_sharding_configuration = ["\08\01\1A\01\01\22\01\00"],
step_marker_location = ""} : (tensor<*xi32>) -> tensor<*xi32>
return
}
func @func(%arg0: tensor<*xi32> {mhlo.sharding = "\01\02\03"}) ->
(tensor<*xi32> {mhlo.sharding = "\08\01\1A\01\01\22\01\00"}) {
%0 = "tf.XlaSharding"(%arg0) {_XlaSharding = "\01\02\03", sharding = "\01\02\03"} : (tensor<*xi32>) -> tensor<*xi32>
%1 = "tf.A"(%0) : (tensor<*xi32>) -> tensor<*xi32>
return %1 : tensor<*xi32>
}
```
}];
}
def TPUValidateSessionInputsPass : Pass<"tf-tpu-validate-session-inputs", "mlir::ModuleOp"> {
let summary = "Validates inputs to the TPU TF/XLA bridge in session api";
let description = [{
This pass checks that the IR has valid input to TPU TF/XLA bridge in session api.
It checks the relations of multiple ops. Properties of single ops are
checked by the 'verify' method of ops.
}];
let constructor = "tensorflow::tf2xla::internal::CreateTPUValidateSessionInputsPass()";
}
def TPUValidateInputsPass : Pass<"tf-tpu-validate-inputs", "mlir::ModuleOp"> {
let summary = "Validates inputs to the TPU TF/XLA bridge";
let description = [{
This pass checks that the IR has valid input to TPU TF/XLA bridge.
It checks the relations of multiple ops. Properties of single ops are
checked by the 'verify' method of ops.
}];
let constructor = "tensorflow::tf2xla::internal::CreateTPUValidateInputsPass()";
}
tensorflow/compiler/mlir/tf2xla/internal/passes/extract_head_tail_outside_compilation.cc
+492
View File
@@ -0,0 +1,492 @@
/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <memory>
#include <string>
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Block.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Transforms/RegionUtils.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_structs.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/device_util.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/tpu_rewrite_device_util.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
// This pass extracts a CPU computation cluster with `_xla_outside_compilation`
// annotation from the head or tail of a TPU cluster.
namespace {
using mlir::Block;
using mlir::BlockArgument;
using mlir::BoolAttr;
using mlir::ModuleOp;
using mlir::OpBuilder;
using mlir::Operation;
using mlir::OperationPass;
using mlir::Region;
using mlir::StringAttr;
using mlir::Type;
using mlir::Value;
using mlir::WalkResult;
using mlir::func::FuncOp;
constexpr char kXlaMapOutsideCompilationAttr[] = "_xla_map_outside_compilation";
constexpr char kXlaOutsideCompilationAttr[] = "_xla_outside_compilation";
// Return true if `op` has attributes that say it can be outside compiled by
// this pass. This pass ignores _xla_map_outside_compilation, which will only be
// handled by extract_outside_compilation pass.
bool HasOutsideCompilationAttribute(Operation* op) {
return op->getAttrOfType<StringAttr>(kXlaOutsideCompilationAttr) != nullptr &&
!op->hasAttrOfType<BoolAttr>(kXlaMapOutsideCompilationAttr);
}
// Finds op that created a given value. If the value is a BlockArgument, this
// returns the owner of the Block.
Operation* GetOpOfValue(Value value) {
if (auto block_arg = mlir::dyn_cast<BlockArgument>(value))
return block_arg.getOwner()->getParentOp();
return value.getDefiningOp();
}
// Checks if `op` is nested in `block`.
bool OpInBlock(Operation* op, Block* block) {
Block* op_block = op->getBlock();
while (op_block) {
if (op_block == block) return true;
if (auto* parent_op = op_block->getParentOp()) {
op_block = parent_op->getBlock();
} else {
break;
}
}
return false;
}
// Wraps block in a Launch. External uses of ops in the block will be return
// values of the Launch and remapped to the Launch results. If `before` is set
// to true, the Launch is created before `op`. Otherwise the Launch is created
// after `op`.
mlir::tf_device::LaunchOp CreateLaunchForBlock(OpBuilder* builder,
Operation* op, bool before,
Block* launch_block,
llvm::StringRef host_device) {
// Find results and result types of ops in block that needs to returned.
llvm::SmallVector<Value, 4> launch_results;
llvm::SmallVector<Type, 4> launch_result_types;
for (Operation& head_outside_compiled_op : *launch_block) {
for (Value result : head_outside_compiled_op.getResults()) {
bool has_external_uses = false;
for (Operation* user : result.getUsers()) {
if (OpInBlock(user, launch_block)) continue;
has_external_uses = true;
break;
}
if (has_external_uses) {
launch_results.push_back(result);
launch_result_types.push_back(result.getType());
}
}
}
before ? builder->setInsertionPoint(op) : builder->setInsertionPointAfter(op);
auto launch = mlir::tf_device::LaunchOp::create(
*builder, op->getLoc(), builder->getStringAttr(host_device),
launch_result_types);
launch.getBody().push_back(launch_block);
builder->setInsertionPointToEnd(&launch.GetBody());
mlir::tf_device::ReturnOp::create(*builder, op->getLoc(), launch_results);
return launch;
}
// Checks if an operation is a supported TPU embedding op.
bool IsEmbeddingOp(Operation* op) {
return llvm::isa<mlir::TF::EnqueueTPUEmbeddingRaggedTensorBatchOp,
mlir::TF::EnqueueTPUEmbeddingSparseTensorBatchOp,
mlir::TF::EnqueueTPUEmbeddingArbitraryTensorBatchOp,
mlir::TF::RecvTPUEmbeddingActivationsOp,
mlir::TF::SendTPUEmbeddingGradientsOp>(op);
}
// Returns a set of ops that are outside compiled and can be extracted to before
// the TPU computation. These ops are either connected to the inputs of the TPU
// computation or other ops that can be extracted, and have no operands from
// other ops in the TPU computation that cannot be extracted.
llvm::SmallVector<Operation*, 4> FindOutsideCompiledOpsAtHead(
const mlir::TF::SideEffectAnalysis& side_effect_analysis,
mlir::tf_device::ClusterOp cluster) {
const auto& analysis = side_effect_analysis.GetAnalysisForFunc(
cluster->getParentOfType<mlir::func::FuncOp>());
Region* cluster_region = &cluster.getBody();
llvm::SmallSetVector<Operation*, 4> head_outside_compiled_ops;
auto cluster_ops = cluster.GetBody().without_terminator();
for (Operation& cluster_op : cluster_ops) {
if (!HasOutsideCompilationAttribute(&cluster_op)) continue;
// An outside compiled op can be extracted if its operands are not from
// other ops in the cluster that cannot be extracted.
// Check if the side effecting op right before this side effecting op, if
// it is side effecting, can be head extracted. Because of op ordering due
// to side effects, if this is not true, this op cannot be head extracted.
// TODO(lyandy): Remove special handling of embedding ops. Currently the IR
// is in a topological sort order and depending on that ordering, embedding
// ops may prevent other ops from being head extracted.
auto predecessors = analysis.DirectControlPredecessors(&cluster_op);
if (!predecessors.empty() && !IsEmbeddingOp(&cluster_op)) {
bool skip = false;
for (Operation* predecessor : llvm::reverse(predecessors)) {
if (IsEmbeddingOp(predecessor)) continue;
skip = !head_outside_compiled_ops.contains(predecessor);
break;
}
if (skip) continue;
}
auto walk_result = cluster_op.walk([&](Operation* op) {
for (Value operand : op->getOperands()) {
Operation* operand_op = GetOpOfValue(operand);
if (head_outside_compiled_ops.count(operand_op) ||
operand_op == &cluster_op)
continue;
if (operand_op->getParentRegion() == cluster_region)
return WalkResult::interrupt();
}
return WalkResult::advance();
});
if (!walk_result.wasInterrupted())
head_outside_compiled_ops.insert(&cluster_op);
}
return head_outside_compiled_ops.takeVector();
}
// Moves head outside compiled ops into its own `tf_device.LaunchOp`
// computation before the cluster.
void CreateHeadComputation(OpBuilder* builder,
mlir::tf_device::ClusterOp cluster,
llvm::ArrayRef<Operation*> head_outside_compiled_ops,
llvm::StringRef host_device) {
Block* launch_block = new Block;
for (Operation* head_outside_compiled_op : head_outside_compiled_ops) {
head_outside_compiled_op->removeAttr(kXlaOutsideCompilationAttr);
head_outside_compiled_op->moveBefore(launch_block, launch_block->end());
}
mlir::tf_device::LaunchOp launch = CreateLaunchForBlock(
builder, cluster, /*before=*/true, launch_block, host_device);
for (auto result : llvm::zip(launch.GetBody().getTerminator()->getOperands(),
launch.getResults()))
replaceAllUsesInRegionWith(std::get<0>(result), std::get<1>(result),
cluster.getBody());
}
// Extracts and move outside compiled ops that have no dependencies in the
// cluster to before the cluster.
mlir::LogicalResult LiftHeadOutsideCompiledOps(
OpBuilder* builder,
const mlir::TF::SideEffectAnalysis& side_effect_analysis,
const mlir::TF::RuntimeDevices& devices, mlir::tf_device::ClusterOp cluster,
std::string* host_device, bool* cluster_updated) {
llvm::SmallVector<Operation*, 4> head_outside_compiled_ops =
FindOutsideCompiledOpsAtHead(side_effect_analysis, cluster);
if (head_outside_compiled_ops.empty()) return mlir::success();
if (failed(tensorflow::GetHostDeviceOutsideComputation(devices, cluster,
host_device)))
return mlir::failure();
CreateHeadComputation(builder, cluster, head_outside_compiled_ops,
*host_device);
*cluster_updated = true;
return mlir::success();
}
// Fills `tail_outside_compiled_ops` with ops that are outside compiled and
// can be extracted to after the TPU computation, and `cluster_results` with new
// results of the cluster. These ops are either connected to the output of the
// TPU computation or other ops that can be extracted, and have no results used
// by other ops in the TPU computation that cannot be extracted.
void FindOutsideCompiledOpsAtTailAndClusterResults(
const mlir::TF::SideEffectAnalysis& side_effect_analysis,
mlir::tf_device::ClusterOp cluster,
llvm::SmallVectorImpl<Operation*>* tail_outside_compiled_ops,
llvm::SmallVectorImpl<Value>* cluster_results) {
const auto& analysis = side_effect_analysis.GetAnalysisForFunc(
cluster->getParentOfType<FuncOp>());
Region* cluster_region = &cluster.getBody();
llvm::SmallSetVector<Operation*, 4> tail_outside_compiled_ops_set;
Operation* terminator = cluster.GetBody().getTerminator();
llvm::SmallSetVector<Value, 4> cluster_results_set;
cluster_results_set.insert(terminator->getOperands().begin(),
terminator->getOperands().end());
auto cluster_ops = llvm::reverse(cluster.GetBody().without_terminator());
for (Operation& cluster_op : cluster_ops) {
if (!HasOutsideCompilationAttribute(&cluster_op)) continue;
// Check if the side effecting op right after this side effecting op, if
// it is side effecting, can be tail extracted. Because of op ordering due
// to side effects, if this is not true, this op cannot be tail extracted.
// TODO(lyandy): Remove special handling of embedding ops. Currently the IR
// is in a topological sort order and depending on that ordering, embedding
// ops may prevent other ops from being tail extracted.
auto successors = analysis.DirectControlSuccessors(
&cluster_op, [&terminator](Operation* op) { return op != terminator; });
if (!successors.empty() && !IsEmbeddingOp(&cluster_op)) {
bool skip = false;
for (Operation* successor : successors) {
if (IsEmbeddingOp(successor)) continue;
skip = !tail_outside_compiled_ops_set.contains(successor);
break;
}
if (skip) continue;
}
llvm::SmallVector<int, 4> results_to_forward;
bool can_be_extracted =
llvm::all_of(cluster_op.getUsers(), [&](Operation* op) {
return op == terminator || tail_outside_compiled_ops_set.count(op);
});
if (!can_be_extracted) continue;
// Collect operands of cluster op that are generated within the cluster.
// These values should be returned by the cluster.
cluster_op.walk([&](Operation* op) {
for (Value operand : op->getOperands()) {
Operation* operand_op = GetOpOfValue(operand);
if (operand_op->getParentRegion() == cluster_region)
cluster_results_set.insert(operand);
}
});
// Remove results of op to be extracted as there are no uses in the cluster.
for (Value result : cluster_op.getResults())
cluster_results_set.remove(result);
// Insert all ops including nested ops for checking outputs/side effects.
cluster_op.walk(
[&](Operation* op) { tail_outside_compiled_ops_set.insert(op); });
// Only add top level ops to output vector.
tail_outside_compiled_ops->push_back(&cluster_op);
}
*cluster_results = cluster_results_set.takeVector();
}
// Moves tail outside compiled ops into its own `tf_device.LaunchOp`
// computation after the cluster.
void CreateTailComputation(OpBuilder* builder,
mlir::tf_device::ClusterOp cluster,
llvm::ArrayRef<Operation*> tail_outside_compiled_ops,
llvm::StringRef host_device) {
Block* launch_block = new Block;
for (Operation* tail_outside_compiled_op : tail_outside_compiled_ops) {
tail_outside_compiled_op->removeAttr(kXlaOutsideCompilationAttr);
tail_outside_compiled_op->moveBefore(launch_block, launch_block->begin());
}
mlir::tf_device::LaunchOp launch = CreateLaunchForBlock(
builder, cluster, /*before=*/false, launch_block, host_device);
auto operand_not_in_launch = [&](mlir::OpOperand& operand) {
return !launch.getOperation()->isProperAncestor(operand.getOwner());
};
for (auto result : llvm::zip(launch.GetBody().getTerminator()->getOperands(),
launch.getResults()))
std::get<0>(result).replaceUsesWithIf(std::get<1>(result),
operand_not_in_launch);
}
// Updates cluster with updated cluster results after extracting tail outside
// compiled ops.
mlir::tf_device::ClusterOp UpdateClusterResults(
OpBuilder* builder, mlir::tf_device::ClusterOp cluster,
llvm::ArrayRef<Value> new_cluster_results) {
Operation* old_terminator = cluster.GetBody().getTerminator();
builder->setInsertionPoint(old_terminator);
mlir::tf_device::ReturnOp::create(*builder, old_terminator->getLoc(),
new_cluster_results);
old_terminator->erase();
builder->setInsertionPoint(cluster);
llvm::SmallVector<Type, 4> new_cluster_result_types;
new_cluster_result_types.reserve(new_cluster_results.size());
for (const auto& new_cluster_result : new_cluster_results)
new_cluster_result_types.push_back(new_cluster_result.getType());
auto new_cluster = mlir::tf_device::ClusterOp::create(
*builder, cluster.getLoc(), new_cluster_result_types,
/*operands=*/llvm::ArrayRef<Value>{}, cluster->getAttrs());
new_cluster.getBody().takeBody(cluster.getBody());
auto operand_not_in_cluster = [&](mlir::OpOperand& operand) {
return !new_cluster.getOperation()->isProperAncestor(operand.getOwner());
};
for (auto result :
llvm::zip(new_cluster.GetBody().getTerminator()->getOperands(),
new_cluster.getResults()))
std::get<0>(result).replaceUsesWithIf(std::get<1>(result),
operand_not_in_cluster);
cluster.erase();
return new_cluster;
}
// Extracts and move outside compiled ops that do not create dependencies in the
// cluster to after the cluster.
mlir::LogicalResult LiftTailOutsideCompiledOps(
OpBuilder* builder,
const mlir::TF::SideEffectAnalysis& side_effect_analysis,
const mlir::TF::RuntimeDevices& devices, std::string host_device,
mlir::tf_device::ClusterOp* cluster, bool* cluster_updated) {
llvm::SmallVector<Operation*, 4> tail_outside_compiled_ops;
llvm::SmallVector<Value, 4> cluster_results;
FindOutsideCompiledOpsAtTailAndClusterResults(side_effect_analysis, *cluster,
&tail_outside_compiled_ops,
&cluster_results);
if (tail_outside_compiled_ops.empty()) return mlir::success();
if (host_device.empty())
if (failed(tensorflow::GetHostDeviceOutsideComputation(devices, *cluster,
&host_device)))
return mlir::failure();
// Forward all results of cluster first. These results will be remapped once
// a new cluster is formed.
cluster->replaceAllUsesWith(
cluster->GetBody().getTerminator()->getOperands());
CreateTailComputation(builder, *cluster, tail_outside_compiled_ops,
host_device);
*cluster = UpdateClusterResults(builder, *cluster, cluster_results);
*cluster_updated = true;
return mlir::success();
}
// Removes aliased outputs in cluster from ops outside of cluster.
void RemoveClusterAliasedOutputs(OpBuilder* builder,
mlir::tf_device::ClusterOp cluster) {
llvm::SmallVector<Value, 4> used_old_cluster_results;
llvm::SmallVector<Value, 4> new_cluster_results;
llvm::SmallVector<Type, 4> new_cluster_result_types;
Operation* cluster_terminator = cluster.GetBody().getTerminator();
for (auto result :
llvm::zip(cluster_terminator->getOperands(), cluster.getResults())) {
Value cluster_terminator_operand = std::get<0>(result);
if (cluster_terminator_operand.getDefiningOp() &&
cluster.getOperation()->isProperAncestor(
cluster_terminator_operand.getDefiningOp())) {
new_cluster_results.push_back(cluster_terminator_operand);
new_cluster_result_types.push_back(cluster_terminator_operand.getType());
used_old_cluster_results.push_back(std::get<1>(result));
} else {
std::get<1>(result).replaceAllUsesWith(cluster_terminator_operand);
}
}
if (new_cluster_results.size() == cluster.getNumResults()) return;
builder->setInsertionPoint(cluster);
auto new_cluster = mlir::tf_device::ClusterOp::create(
*builder, cluster.getLoc(), new_cluster_result_types,
/*operands=*/llvm::ArrayRef<Value>{}, cluster->getAttrs());
new_cluster.getBody().takeBody(cluster.getBody());
new_cluster.GetBody().getTerminator()->setOperands(new_cluster_results);
for (auto result :
llvm::zip(used_old_cluster_results, new_cluster.getResults()))
std::get<0>(result).replaceAllUsesWith(std::get<1>(result));
cluster.erase();
}
#define GEN_PASS_DEF_EXTRACTHEADTAILOUTSIDECOMPILATIONPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc"
struct ExtractHeadTailOutsideCompilationPass
: public impl::ExtractHeadTailOutsideCompilationPassBase<
ExtractHeadTailOutsideCompilationPass> {
void runOnOperation() override;
};
void ExtractHeadTailOutsideCompilationPass::runOnOperation() {
auto& side_effect_analysis = getAnalysis<mlir::TF::SideEffectAnalysis>();
// Get runtime devices information from the closest parent module.
auto module = getOperation();
mlir::TF::RuntimeDevices devices;
if (failed(tensorflow::GetDevicesFromOp(module, &devices)))
return signalPassFailure();
OpBuilder builder(&getContext());
llvm::SmallVector<mlir::tf_device::ClusterOp, 4> clusters;
module.walk(
[&](mlir::tf_device::ClusterOp cluster) { clusters.push_back(cluster); });
for (mlir::tf_device::ClusterOp cluster : clusters) {
std::string host_device;
bool cluster_updated = false;
if (failed(LiftHeadOutsideCompiledOps(&builder, side_effect_analysis,
devices, cluster, &host_device,
&cluster_updated)) ||
failed(LiftTailOutsideCompiledOps(&builder, side_effect_analysis,
devices, host_device, &cluster,
&cluster_updated)))
return signalPassFailure();
if (cluster_updated) RemoveClusterAliasedOutputs(&builder, cluster);
}
}
} // anonymous namespace
std::unique_ptr<OperationPass<ModuleOp>>
CreateExtractHeadTailOutsideCompilationPass() {
return std::make_unique<ExtractHeadTailOutsideCompilationPass>();
}
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/passes/extract_outside_compilation.cc
+1593
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/mlir/tf2xla/internal/passes/hoist_broadcast_read.cc
+154
View File
@@ -0,0 +1,154 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <memory>
#include <string>
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/core/util/device_name_utils.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {
using mlir::BlockArgument;
using mlir::failure;
using mlir::LogicalResult;
using mlir::Operation;
using mlir::OperationPass;
using mlir::OpOperand;
using mlir::StringAttr;
using mlir::success;
using mlir::Value;
using mlir::WalkResult;
using mlir::func::FuncOp;
using mlir::TF::ReadVariableOp;
using mlir::tf_device::ReplicateOp;
#define GEN_PASS_DEF_HOISTBROADCASTREADPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc"
constexpr char kFuncDeviceAttr[] = "tf.device";
constexpr char kCpuDeviceType[] = "CPU";
struct HoistBroadcastRead
: public impl::HoistBroadcastReadPassBase<HoistBroadcastRead> {
void runOnOperation() override;
};
// Get the ancestor of `descendant` that is a direct child of `ancestor`.
Operation* GetAncestorBelow(Operation* descendant, Operation* ancestor) {
Operation* parent = descendant->getParentOp();
if (!parent) return nullptr;
if (parent == ancestor) return descendant;
return GetAncestorBelow(parent, ancestor);
}
// `is_cpu_read` is set to `true` iff `read` is on a resource with device type
// CPU.
LogicalResult IsCpuRead(FuncOp func, ReadVariableOp read, bool& is_cpu_read) {
if (auto arg = mlir::dyn_cast<BlockArgument>(read->getOperand(0))) {
if (arg.getOwner() != &(func.front())) {
is_cpu_read = false;
return success();
}
if (auto attr = func.getArgAttrOfType<StringAttr>(arg.getArgNumber(),
kFuncDeviceAttr)) {
std::string device = attr.getValue().str();
tensorflow::DeviceNameUtils::ParsedName parsed_name;
if (!tensorflow::DeviceNameUtils::ParseFullName(device, &parsed_name)) {
return read->emitOpError() << "invalid device '" << device << "'";
}
is_cpu_read = parsed_name.type == kCpuDeviceType;
return success();
}
}
is_cpu_read = false;
return success();
}
// Get the reads to hoist in the `replicate`.
LogicalResult GetReads(FuncOp func, ReplicateOp replicate,
llvm::SmallVector<ReadVariableOp, 4>& reads) {
for (Operation& op : replicate.getBody().front()) {
if (auto read = llvm::dyn_cast<ReadVariableOp>(&op)) {
bool is_cpu_read;
if (failed(IsCpuRead(func, read, is_cpu_read))) return failure();
if (is_cpu_read) reads.push_back(read);
}
}
return success();
}
// Move reads above the `replicate`. Skip reads that come after a write to the
// same resource.
void MoveReads(ReplicateOp replicate,
llvm::SmallVector<ReadVariableOp, 4>& reads) {
for (ReadVariableOp read : reads) {
Value res = read.getResource();
Operation* scope = res.getParentBlock()->getParentOp();
if (!scope->isProperAncestor(replicate)) continue;
bool has_conflicting_write = false;
for (OpOperand& use : res.getUses()) {
Operation* using_op = use.getOwner();
if (using_op == read) continue;
if (!replicate->isProperAncestor(using_op)) continue;
Operation* peer = GetAncestorBelow(using_op, replicate);
if (read->isBeforeInBlock(peer)) continue;
if (llvm::isa<ReadVariableOp>(peer)) continue;
has_conflicting_write = true;
}
if (has_conflicting_write) continue;
read->moveBefore(replicate);
}
}
// Hoist `ReadVariableOp`s above the `tf_device.replicate`s.
void HoistBroadcastRead::runOnOperation() {
FuncOp func = getOperation();
auto result = func.walk([&](ReplicateOp replicate) {
llvm::SmallVector<ReadVariableOp, 4> reads;
if (failed(GetReads(func, replicate, reads)))
return WalkResult::interrupt();
MoveReads(replicate, reads);
return WalkResult::advance();
});
if (result.wasInterrupted()) return signalPassFailure();
}
} // namespace
std::unique_ptr<OperationPass<FuncOp>> CreateHoistBroadcastReadPass() {
return std::make_unique<HoistBroadcastRead>();
}
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/passes/input_lowering_metrics_pass.cc
+88
View File
@@ -0,0 +1,88 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <memory>
#include "llvm/ADT/DenseSet.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/legalization_op_config.h"
#include "tensorflow/core/lib/monitoring/counter.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {
using mlir::Operation;
using mlir::WalkResult;
#define GEN_PASS_DEF_INPUTLOWERINGMETRICSPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/lowering_passes.h.inc"
auto* dynamism_op_counter = tensorflow::monitoring::Counter<1>::New(
"/tensorflow/core/tf2xla/api/v2/dynamism_op_counter",
"Counts how many ops are dynamic", "op_name");
auto* dynamism_function_counter = tensorflow::monitoring::Counter<1>::New(
"/tensorflow/core/tf2xla/api/v2/dynamism_function_counter",
"Counts how many functions are dynamic", "has_dynamism");
constexpr char kNotDynamicFunctionName[] = "kNotDynamicFunction";
constexpr char kDynamicFunctionName[] = "kDynamicFunction";
class InputMetricsLoweringPass
: public impl::InputLoweringMetricsPassBase<InputMetricsLoweringPass> {
public:
void runOnOperation() override;
};
void InputMetricsLoweringPass::runOnOperation() {
bool has_dynamic_op = false;
Operation* func_op = getOperation();
func_op->walk([&](Operation* op) {
auto abstractOp = op->getRegisteredInfo();
if (!abstractOp) return WalkResult::advance();
if (mlir::hlo::IsDynamicPadderOp(abstractOp->getTypeID())) {
has_dynamic_op = true;
dynamism_op_counter->GetCell(op->getName().getStringRef().str())
->IncrementBy(1);
}
return WalkResult::advance();
});
if (has_dynamic_op) {
dynamism_function_counter->GetCell(kDynamicFunctionName)->IncrementBy(1);
} else {
dynamism_function_counter->GetCell(kNotDynamicFunctionName)->IncrementBy(1);
}
}
} // namespace
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateInputLoweringMetricsPass() {
return std::make_unique<InputMetricsLoweringPass>();
}
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/passes/input_lowering_metrics_pass_test.cc
+124
View File
@@ -0,0 +1,124 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <cstdint>
#include <memory>
#include <gtest/gtest.h>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OperationSupport.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/lowering_passes.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/test_utils.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/lib/monitoring/cell_reader.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {
using ::mlir::LogicalResult;
using ::mlir::ModuleOp;
using ::mlir::hlo::test::GetMlirModuleFromString;
using ::tensorflow::monitoring::testing::CellReader;
constexpr char kNotDynamicFunctionName[] = "kNotDynamicFunction";
constexpr char kDynamicFunctionName[] = "kDynamicFunction";
static constexpr char kDynamismOpCounterStreamzName[] =
"/tensorflow/core/tf2xla/api/v2/dynamism_op_counter";
static constexpr char kDynamismFunctionCounterStreamzName[] =
"/tensorflow/core/tf2xla/api/v2/dynamism_function_counter";
class InputLoweringMetricsPassTest : public testing::Test {
protected:
void CreateModule(const char* module_string) {
TF_ASSERT_OK_AND_ASSIGN(module_,
GetMlirModuleFromString(module_string, &context_));
pm_ = std::make_unique<mlir::PassManager>(&context_);
pm_->addNestedPass<mlir::func::FuncOp>(CreateInputLoweringMetricsPass());
}
bool ModulesEqual(const ModuleOp& module_before,
const ModuleOp& module_after) {
return mlir::OperationEquivalence::isEquivalentTo(
module_before, module_after, mlir::OperationEquivalence::None);
}
mlir::LogicalResult Run() {
mlir::OwningOpRef<mlir::ModuleOp> module_before = module_->clone();
LogicalResult run_result = pm_->run(module_.get());
EXPECT_TRUE(ModulesEqual(*module_before, *module_));
return run_result;
}
private:
mlir::MLIRContext context_;
mlir::OwningOpRef<ModuleOp> module_;
std::unique_ptr<mlir::PassManager> pm_;
};
TEST_F(InputLoweringMetricsPassTest, CountsNoDynamicOps) {
static constexpr char kMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> tensor<1xi32> {
%0 = "tf.Const"() {value = dense<1000> : tensor<1xi32>} : () -> tensor<1xi32>
return %0 : tensor<1xi32>
}
})";
CellReader<int64_t> dynamism_op_counter(kDynamismOpCounterStreamzName);
CellReader<int64_t> dynamism_function_counter(
kDynamismFunctionCounterStreamzName);
CreateModule(kMlirModuleStr);
auto result = Run();
EXPECT_TRUE(result.succeeded());
EXPECT_EQ(dynamism_function_counter.Delta(kNotDynamicFunctionName), 1);
}
TEST_F(InputLoweringMetricsPassTest, CountsDynamicOps) {
static constexpr char kMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> () {
%cst0 = "tf.Const"(){ value = dense<0> : tensor<3x5xi1>} : () -> tensor<3x5xi1>
%0 = "tf.Where"(%cst0) : (tensor<3x5xi1>) -> tensor<?x2xi64>
func.return
}
})";
CellReader<int64_t> dynamism_counter(kDynamismOpCounterStreamzName);
CellReader<int64_t> dynamism_function_counter(
kDynamismFunctionCounterStreamzName);
CreateModule(kMlirModuleStr);
auto result = Run();
EXPECT_TRUE(result.succeeded());
EXPECT_EQ(dynamism_counter.Delta("tf.Where"), 1);
EXPECT_EQ(dynamism_function_counter.Delta(kDynamicFunctionName), 1);
}
} // namespace
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/passes/lowering_passes.h
+38
View File
@@ -0,0 +1,38 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_PASSES_LOWERING_PASSES_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_PASSES_LOWERING_PASSES_H_
#include <memory>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
namespace tensorflow {
namespace tf2xla {
namespace internal {
// Create a pass that just collects metrics about the input MLIR. Does not
// logically transform the program.
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateInputLoweringMetricsPass();
#define GEN_PASS_REGISTRATION
#define GEN_PASS_DECL_INPUTLOWERINGMETRICSPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/lowering_passes.h.inc"
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_PASSES_LOWERING_PASSES_H_
tensorflow/compiler/mlir/tf2xla/internal/passes/lowering_passes.td
+25
View File
@@ -0,0 +1,25 @@
/* Copyright 2024 TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
include "mlir/Pass/PassBase.td"
def InputLoweringMetricsPass : Pass<"input-lowering-metrics-pass", "mlir::func::FuncOp"> {
let summary = "Collects various metrics about the input to the lowering "
"portion of the bridge. This is a logical no-op.";
let description = [{
Gathers metrics about the input MLIR to Phase 2 of the TFXLA Bridge, which
does a strict semantic lowering from Tensorflow ops to XLA HLO.
}];
let constructor = "tensorflow::tf2xla::internal::CreateInputLoweringMetricsPass()";
}
tensorflow/compiler/mlir/tf2xla/internal/passes/mark_ops_for_outside_compilation.cc
+524
View File
@@ -0,0 +1,524 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <memory>
#include <optional>
#include <queue>
#include <string>
#include <utility>
#include "absl/log/log.h"
#include "absl/strings/str_join.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/FormatVariadic.h"
#include "mlir/IR/Block.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/OperationSupport.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/IR/TypeUtilities.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Rewrite/PatternApplicator.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Transforms/RegionUtils.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/lower_tf.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/string_util.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/legalization_op_config.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/passes.h"
#include "tensorflow/core/lib/monitoring/gauge.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {
using mlir::Block;
using mlir::BoolAttr;
using mlir::Dialect;
using mlir::LogicalResult;
using mlir::MLIRContext;
using mlir::ModuleOp;
using mlir::Operation;
using mlir::OperationName;
using mlir::OperationPass;
using mlir::Pattern;
using mlir::PatternApplicator;
using mlir::RewritePatternSet;
using mlir::StringAttr;
using mlir::TensorType;
using mlir::Type;
using mlir::Value;
using mlir::WalkResult;
auto* auto_outside_compilation_gauge =
tensorflow::monitoring::Gauge<bool, 0>::New(
"/tensorflow/core/use_auto_outside_compilation",
"Tracks if auto outside compilation is enabled");
#define GEN_PASS_DEF_MARKOPSFOROUTSIDECOMPILATIONPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc"
struct MarkOpsForOutsideCompilation
: public impl::MarkOpsForOutsideCompilationPassBase<
MarkOpsForOutsideCompilation> {
void runOnOperation() override;
};
// Adds any canonicalization patterns to list of supported `patterns`.
// TODO(b/161726307): Move or import the relevant patterns to LowerTF pass and
// remove this.
void AddCanonicalizationPatterns(MLIRContext* context,
RewritePatternSet* patterns) {
for (auto op : context->getRegisteredOperations())
op.getCanonicalizationPatterns(*patterns, context);
}
// Adds the list of ops that are supported on TPU through constant folding which
// may depend on the inputs shapes not known at this point. Such ops may not
// have any legalization or canonicalization patterns but shouldn't be marked
// for outside compilation.
//
// TODO(b/177523289): Remove manual handling once we support constant folding
// and shape inference through the computation on the host side.
void AddSupportedOpsUsingFolding(MLIRContext* context,
llvm::DenseSet<OperationName>* supported_ops) {
llvm::SmallDenseSet<OperationName, 8> allowlist_ops = {
OperationName(mlir::TF::BroadcastArgsOp::getOperationName(), context),
OperationName(mlir::TF::BroadcastGradientArgsOp::getOperationName(),
context),
OperationName(mlir::TF::ConcatOffsetOp::getOperationName(), context),
OperationName(mlir::TF::EmptyOp::getOperationName(), context),
OperationName(mlir::TF::ListDiffOp::getOperationName(), context),
OperationName(mlir::TF::RankOp::getOperationName(), context),
OperationName(mlir::TF::RangeOp::getOperationName(), context),
OperationName(mlir::TF::ShapeOp::getOperationName(), context),
OperationName(mlir::TF::ShapeNOp::getOperationName(), context),
OperationName(mlir::TF::SizeOp::getOperationName(), context),
};
supported_ops->insert(allowlist_ops.begin(), allowlist_ops.end());
}
// Adds the list of ops that are only supported in the old bridge.
// TODO(b/168036682): Remove bounded dynamism ops now that MLIR bridge supports
// bounded dynamism.
// TODO(b/257574556): Remove the need for this manual list by making use of old
// bridge phase 2 op list.
void AddOldBridgeOnlyOps(MLIRContext* context,
llvm::DenseSet<OperationName>* supported_ops) {
llvm::SmallDenseSet<OperationName, 8> allowlist_ops = {
OperationName(mlir::TF::DynamicPartitionOp::getOperationName(), context),
OperationName(mlir::TF::OutfeedEnqueueOp::getOperationName(), context),
OperationName(mlir::TF::WhereOp::getOperationName(), context),
OperationName(mlir::TF::UniqueOp::getOperationName(), context),
OperationName(mlir::TF::XlaSetDynamicDimensionSizeOp::getOperationName(),
context),
OperationName(mlir::TF::XlaSpmdFullToShardShapeOp::getOperationName(),
context),
OperationName(mlir::TF::XlaSpmdShardToFullShapeOp::getOperationName(),
context),
};
supported_ops->insert(allowlist_ops.begin(), allowlist_ops.end());
}
// TODO(b/159128666): Check the control flow legalization passes instead once
// added.
void AddSupportedFunctionalOps(MLIRContext* context,
llvm::DenseSet<OperationName>* supported_ops) {
supported_ops->insert(
OperationName(mlir::TF::CaseRegionOp::getOperationName(), context));
supported_ops->insert(
OperationName(mlir::TF::IfRegionOp::getOperationName(), context));
supported_ops->insert(
OperationName(mlir::TF::InplaceAddOp::getOperationName(), context));
supported_ops->insert(
OperationName(mlir::TF::WhileRegionOp::getOperationName(), context));
supported_ops->insert(
OperationName(mlir::TF::XlaCallModuleOp::getOperationName(), context));
supported_ops->insert(
OperationName(mlir::TF::XlaHostComputeOp::getOperationName(), context));
supported_ops->insert(
OperationName(mlir::TF::XlaReduceOp::getOperationName(), context));
supported_ops->insert(
OperationName(mlir::TF::XlaReduceWindowOp::getOperationName(), context));
supported_ops->insert(OperationName(
mlir::TF::XlaRngBitGeneratorOp::getOperationName(), context));
supported_ops->insert(
OperationName(mlir::TF::XlaScatterOp::getOperationName(), context));
supported_ops->insert(OperationName(
mlir::TF::XlaSelectAndScatterOp::getOperationName(), context));
supported_ops->insert(
OperationName(mlir::TF::SymbolicGradientOp::getOperationName(), context));
supported_ops->insert(OperationName(
mlir::TF::XlaVariadicReduceOp::getOperationName(), context));
supported_ops->insert(OperationName(
mlir::TF::XlaVariadicReduceV2Op::getOperationName(), context));
supported_ops->insert(
OperationName(mlir::TF::XlaVariadicSortOp::getOperationName(), context));
supported_ops->insert(
OperationName(mlir::TF::XlaReplicaIdOp::getOperationName(), context));
supported_ops->insert(
OperationName(mlir::TF::YieldOp::getOperationName(), context));
}
// These embedding ops are rewritten when running TPUCompileOp.
void AddRewrittenEmbeddingOps(MLIRContext* context,
llvm::DenseSet<OperationName>* supported_ops) {
supported_ops->insert(OperationName(
mlir::TF::RecvTPUEmbeddingActivationsOp::getOperationName(), context));
supported_ops->insert(OperationName(
mlir::TF::SendTPUEmbeddingGradientsOp::getOperationName(), context));
}
// Stack, TensorList and TensorArray ops are rewritten during the second phase
// of the bridge (compilation of TPUCompile op). They would not match any
// legalization/canonicalization pattern and have to be manually added to the
// list of supported ops.
void AddRewrittenCompositeOps(MLIRContext* context,
llvm::DenseSet<OperationName>* supported_ops) {
#define GET_OPERATION_NAME(op) OperationName(op::getOperationName(), context)
llvm::SmallDenseSet<OperationName, 32> allowlist_ops = {
// Stack ops.
GET_OPERATION_NAME(mlir::TF::StackV2Op),
GET_OPERATION_NAME(mlir::TF::StackPushV2Op),
GET_OPERATION_NAME(mlir::TF::StackPopV2Op),
// Tensor Array ops.
GET_OPERATION_NAME(mlir::TF::TensorArrayV3Op),
GET_OPERATION_NAME(mlir::TF::TensorArrayReadV3Op),
GET_OPERATION_NAME(mlir::TF::TensorArrayWriteV3Op),
GET_OPERATION_NAME(mlir::TF::TensorArrayConcatV3Op),
GET_OPERATION_NAME(mlir::TF::TensorArraySplitV3Op),
GET_OPERATION_NAME(mlir::TF::TensorArraySizeV3Op),
GET_OPERATION_NAME(mlir::TF::TensorArrayGradV3Op),
GET_OPERATION_NAME(mlir::TF::TensorArrayGatherV3Op),
GET_OPERATION_NAME(mlir::TF::TensorArrayScatterV3Op),
// Tensor List Ops.
GET_OPERATION_NAME(mlir::TF::EmptyTensorListOp),
GET_OPERATION_NAME(mlir::TF::TensorListReserveOp),
GET_OPERATION_NAME(mlir::TF::TensorListFromTensorOp),
GET_OPERATION_NAME(mlir::TF::TensorListPushBackOp),
GET_OPERATION_NAME(mlir::TF::TensorListPopBackOp),
GET_OPERATION_NAME(mlir::TF::TensorListGetItemOp),
GET_OPERATION_NAME(mlir::TF::TensorListSetItemOp),
GET_OPERATION_NAME(mlir::TF::TensorListLengthOp),
GET_OPERATION_NAME(mlir::TF::TensorListElementShapeOp),
GET_OPERATION_NAME(mlir::TF::TensorListGatherOp),
GET_OPERATION_NAME(mlir::TF::TensorListScatterIntoExistingListOp),
GET_OPERATION_NAME(mlir::TF::TensorListStackOp),
};
#undef GET_OPERATION_NAME
supported_ops->insert(allowlist_ops.begin(), allowlist_ops.end());
}
bool IsStringType(Type type) {
if (mlir::isa<mlir::TF::StringType>(type)) return true;
auto sub_type = mlir::dyn_cast<mlir::TF::TensorFlowTypeWithSubtype>(type);
if (!sub_type) return false;
bool has_string = llvm::any_of(sub_type.GetSubtypes(), [](TensorType type) {
return mlir::isa<mlir::TF::StringType>(type.getElementType());
});
return has_string;
}
bool HasStringOperand(Operation& op) {
for (auto operand : op.getOperands()) {
auto operand_type = getElementTypeOrSelf(operand);
if (IsStringType(operand_type)) return true;
}
return false;
}
bool HasStringResult(Operation& op) {
for (auto result : op.getResults()) {
auto result_type = getElementTypeOrSelf(result);
if (IsStringType(result_type)) return true;
}
return false;
}
bool MatchesPattern(Operation& op,
const llvm::DenseSet<OperationName>& supported_ops) {
return (supported_ops.contains(op.getName()));
}
// Checks if the op is supported inside of a device cluster. Ops not
// in `tf_dialect` are considered supported.
bool IsSupportedOp(Operation& op,
const llvm::DenseSet<OperationName>& supported_ops,
const Dialect* tf_dialect) {
if (op.getDialect() != tf_dialect) return true;
// Assert has a legalization that later removes it so we don't want to outside
// compile it ever for performance reasons.
if (llvm::isa<mlir::TF::AssertOp>(op)) return true;
if (HasStringOperand(op)) return false;
if (HasStringResult(op)) return false;
if (MatchesPattern(op, supported_ops)) return true;
auto abstractOp = op.getRegisteredInfo();
if (!abstractOp) return false;
return mlir::hlo::HasTf2XlaFallback(abstractOp->getTypeID());
}
bool IsVariant(Value value) {
return mlir::isa<mlir::TF::VariantType>(
getElementTypeOrSelf(value.getType()));
}
bool HasOutsideCompiledAncestor(Operation* op) {
Operation* parent = op->getParentOp();
while (parent) {
if (parent->getAttrOfType<StringAttr>(kXlaOutsideCompilationAttr))
return true;
parent = parent->getParentOp();
}
return false;
}
// If any tf.variants are inputs/outputs to the another outside compiled
// Operation, `op`, mark them for outside compilation unless they are already
// marks with outside compilation attribute.
void MarkVariantInputsOutputs(mlir::tf_device::ClusterOp tpu_cluster) {
std::queue<Operation*> outside_compiled_ops;
tpu_cluster.walk([&](Operation* op) {
if (op->hasAttrOfType<StringAttr>(kXlaOutsideCompilationAttr))
outside_compiled_ops.push(op);
});
while (!outside_compiled_ops.empty()) {
Operation* host_op = outside_compiled_ops.front();
outside_compiled_ops.pop();
host_op->walk([&](Operation* op) {
// Add any operations that provide variant inputs to the cluster.
for (auto value : op->getOperands()) {
Operation* input_defining_op = value.getDefiningOp();
if (IsVariant(value) && input_defining_op &&
!HasOutsideCompiledAncestor(input_defining_op) &&
!input_defining_op->hasAttrOfType<StringAttr>(
kXlaOutsideCompilationAttr)) {
input_defining_op->setAttr(
kXlaOutsideCompilationAttr,
StringAttr::get(input_defining_op->getContext(), "auto"));
outside_compiled_ops.push(input_defining_op);
}
}
// Mark for outside compilation any operations that consume variant
// outputs from an outside compiled operation.
for (auto value : op->getResults()) {
if (IsVariant(value)) {
for (auto user : value.getUsers()) {
if (!user->hasTrait<mlir::OpTrait::IsTerminator>() &&
!HasOutsideCompiledAncestor(user) &&
!user->getAttrOfType<StringAttr>(kXlaOutsideCompilationAttr)) {
user->setAttr(kXlaOutsideCompilationAttr,
StringAttr::get(user->getContext(), "auto"));
outside_compiled_ops.push(user);
}
}
}
}
});
}
}
// Marks uncompilable ops that are in `tf_dialect` for outside compilation.
LogicalResult MarkUncompilableOps(
const Dialect* tf_dialect, Block* block,
llvm::DenseSet<OperationName>& supported_ops) {
// Automatically marked ops for outside compilation have
// `_xla_outside_compilation` attribute value of "auto" plus
// an increasing counter. Manually marked ops for outside compilation only
// have an increasing counteri for the attribute value. Therefore there is no
// collision in
// `_xla_outside_compilation` attribute between automatically and manually
// marking ops.
int outside_compiled_cluster_counter = 0;
block->walk([&](Operation* op) {
if (!IsSupportedOp(*op, supported_ops, tf_dialect)) {
VLOG(3) << "Cloud TPU: Op " << op->getName().getStringRef().str()
<< " isn't compilable, adding outside_compilation attr. "
"This op will automatically be placed on CPU.";
op->setAttr(kXlaOutsideCompilationAttr,
StringAttr::get(
op->getContext(),
llvm::formatv("auto{0}", outside_compiled_cluster_counter)
.str()));
outside_compiled_cluster_counter++;
}
});
if (outside_compiled_cluster_counter > 0) {
auto_outside_compilation_gauge->GetCell()->Set(true);
}
return mlir::success();
}
// Check for uncompilable ops that are in `tf_dialect` and are not already
// marked for outside compilation.
bool ContainsUncompilableOps(const Dialect* tf_dialect, Block* block,
llvm::DenseSet<OperationName>& supported_ops) {
int uncompilable_op_count = 0;
// Check if op or any parent is already marked for outside compilation.
block->walk([&](Operation* op) {
Operation* iter_op = op;
while (iter_op && !llvm::isa<mlir::tf_device::ClusterOp>(iter_op)) {
if (iter_op->hasAttrOfType<StringAttr>(kXlaOutsideCompilationAttr)) {
return;
}
iter_op = iter_op->getParentOp();
}
if (!IsSupportedOp(*op, supported_ops, tf_dialect)) {
op->emitOpError() << "isn't compilable for TPU device. enable "
"soft_device_placement option to run on CPU";
++uncompilable_op_count;
}
});
return uncompilable_op_count > 0;
}
// Unmarks outside compilation for any op that has parents already
// marked for outside compilation since the child will be extracted
// anyways.
void UnmarkChildren(ModuleOp module) {
module->walk([&](Operation* op) {
if (!op->getAttrOfType<StringAttr>(kXlaOutsideCompilationAttr)) return;
Operation* iter_op = op;
bool remove_attr = false;
while (auto* parent_op = iter_op->getParentOp()) {
if (parent_op->getAttrOfType<StringAttr>(kXlaOutsideCompilationAttr)) {
remove_attr = true;
break;
}
iter_op = parent_op;
}
if (remove_attr) op->removeAttr(kXlaOutsideCompilationAttr);
});
}
constexpr int kTooManyOutsideCompileRegionThreshold = 32;
constexpr int kOpDetailCount = 8;
void WarnOnExcessOutsideCompilationOps(ModuleOp module) {
// Count the number of outside compilation ops. If it exceeds the reporting
// threshold, warn the user that their model may run slowly.
llvm::SmallVector<Operation*, 8> outside_compile_ops;
module->walk([&](Operation* op) {
if (op->getAttrOfType<StringAttr>(kXlaOutsideCompilationAttr)) {
outside_compile_ops.push_back(op);
}
});
if (outside_compile_ops.size() > kTooManyOutsideCompileRegionThreshold) {
llvm::SmallVector<std::string, kOpDetailCount> op_info;
for (int i = 0; i < kOpDetailCount; ++i) {
auto& op = outside_compile_ops[i];
op_info.push_back(tensorflow::OpAsString(*op));
}
LOG(WARNING) << outside_compile_ops.size() << " outside compilation "
<< "regions found while processing "
<< module->getName().getStringRef().str()
<< ". This may result in excessively slow model execution. "
<< "First " << op_info.size()
<< " ops: " << absl::StrJoin(op_info, "\n");
} else {
LOG(INFO) << "Found " << outside_compile_ops.size()
<< " outside compilation regions.";
}
}
void MarkOpsForOutsideCompilation::runOnOperation() {
auto module = getOperation();
const Dialect* tf_dialect = getContext().getLoadedDialect("tf");
if (!tf_dialect) {
getOperation().emitError() << "'tf' dialect is not registered";
return signalPassFailure();
}
RewritePatternSet patterns(&getContext());
mlir::hlo::PopulateLegalizeTfPatterns(module.getContext(), &patterns);
mlir::TF::PopulateTFLoweringBeforeHLOPatterns(module.getContext(), &patterns);
mlir::TF::PopulateLoweringQuantizedPatterns(module.getContext(), &patterns);
AddCanonicalizationPatterns(module.getContext(), &patterns);
// `supported_ops` contains the name of all of the ops that can potentially be
// lowered into HLO on the device. This doesn't always mean that the op can
// be lowered in the future passes but if the op is not in this set, it can't
// be lowered in a subsequent pass.
llvm::DenseSet<OperationName> supported_ops;
PatternApplicator(std::move(patterns))
.walkAllPatterns([&](const Pattern& pattern) {
std::optional<OperationName> root_kind = pattern.getRootKind();
if (root_kind.has_value()) supported_ops.insert(root_kind.value());
});
AddSupportedFunctionalOps(module.getContext(), &supported_ops);
AddSupportedOpsUsingFolding(module.getContext(), &supported_ops);
AddOldBridgeOnlyOps(module.getContext(), &supported_ops);
AddRewrittenEmbeddingOps(module.getContext(), &supported_ops);
AddRewrittenCompositeOps(module.getContext(), &supported_ops);
auto result = module.walk([&](mlir::tf_device::ClusterOp cluster) {
// Only if `allow_soft_placement` attribute is true should we mark ops
// for outside compilation.
auto soft_placement_attr =
cluster->getAttrOfType<BoolAttr>(mlir::TF::kAllowSoftPlacementAttr);
if ((soft_placement_attr && soft_placement_attr.getValue())) {
if (failed(MarkUncompilableOps(tf_dialect, &cluster.GetBody(),
supported_ops)))
return WalkResult::interrupt();
} else {
if (ContainsUncompilableOps(tf_dialect, &cluster.GetBody(),
supported_ops))
return WalkResult::interrupt();
}
MarkVariantInputsOutputs(cluster);
return WalkResult::advance();
});
if (result.wasInterrupted()) return signalPassFailure();
UnmarkChildren(module);
WarnOnExcessOutsideCompilationOps(module);
}
} // namespace
std::unique_ptr<OperationPass<ModuleOp>>
CreateMarkOpsForOutsideCompilationPass() {
return std::make_unique<MarkOpsForOutsideCompilation>();
}
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/passes/mlir_to_graph_passes.h
+35
View File
@@ -0,0 +1,35 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_PASSES_MLIR_TO_GRAPH_PASSES_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_PASSES_MLIR_TO_GRAPH_PASSES_H_
#include <memory>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
namespace tensorflow {
namespace tf2xla {
namespace internal {
// Verifies that Executor input is of the expected format.
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateVerifyInputDialectToExecutorPass();
#define GEN_PASS_REGISTRATION
#define GEN_PASS_DECL_VERIFYINPUTDIALECTTOEXECUTORPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/mlir_to_graph_passes.h.inc"
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_PASSES_MLIR_TO_GRAPH_PASSES_H_
tensorflow/compiler/mlir/tf2xla/internal/passes/mlir_to_graph_passes.td
+20
View File
@@ -0,0 +1,20 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
include "mlir/Pass/PassBase.td"
def VerifyInputDialectToExecutorPass : Pass<"verify-input-dialect-to-executor-pass", "mlir::func::FuncOp"> {
let summary = "Verify that TF dialect to executor converter receives the correct input.";
let description = [{
Verifies the input before exporting to TF executor. This includes checking whether the Ops are in TF functional, have device attributes & there are no tf_device.cluster_func ops.
}];
let constructor = "tensorflow::tf2xla::internal::CreateVerifyInputDialectToExecutorPass()";
}
tensorflow/compiler/mlir/tf2xla/internal/passes/tpu_cluster_formation.cc
+977
View File
@@ -0,0 +1,977 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <cassert>
#include <cstdint>
#include <iterator>
#include <memory>
#include <set>
#include <string>
#include <tuple>
#include <unordered_map>
#include <utility>
#include "absl/container/flat_hash_map.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringExtras.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/iterator_range.h"
#include "llvm/Support/Casting.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Diagnostics.h" // from @llvm-project
#include "mlir/IR/DialectRegistry.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/OperationSupport.h" // from @llvm-project
#include "mlir/IR/Region.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/IR/ValueRange.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Transforms/RegionUtils.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/string_util.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/tpu_rewrite_device_util.h"
#include "tensorflow/core/util/device_name_utils.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {
using mlir::Block;
using mlir::DialectRegistry;
using mlir::LogicalResult;
using mlir::ModuleOp;
using mlir::NamedAttribute;
using mlir::NamedAttrList;
using mlir::OpBuilder;
using mlir::Operation;
using mlir::OpResult;
using mlir::Region;
using mlir::StringAttr;
using mlir::success;
using mlir::Type;
using mlir::Value;
using mlir::ValueRange;
using mlir::WalkResult;
constexpr llvm::StringRef kDeviceAttr = "device";
constexpr llvm::StringRef kNameAttr = "name";
constexpr llvm::StringRef kNumCoresPerReplicaAttr = "num_cores_per_replica";
constexpr llvm::StringRef kNumReplicasAttr = "num_replicas";
constexpr llvm::StringRef kMirroredVariableIndicesAttr =
"_mirrored_variable_indices";
constexpr llvm::StringRef kBadReplicateInfoAttrMsg =
"requires '_replication_info' string attribute";
// Mapping for `_replication_info` attribute to TPUReplicateMetadata attributes.
using MetadataMap = llvm::SmallDenseMap<llvm::StringRef, NamedAttrList, 8>;
// A set of operations. We use a `SmallSetVector` in order to have deterministic
// traversal order (= insertion order), independent of the pointer keys.
using OpSetVector = llvm::SmallSetVector<Operation*, 8>;
// Mapping for `_replication_info` attribute to ops of a cluster.
using ClusterMap = llvm::SmallDenseMap<llvm::StringRef, OpSetVector, 8>;
#define GEN_PASS_DEF_TPUCLUSTERFORMATIONPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc"
class TPUClusterFormationPass
: public impl::TPUClusterFormationPassBase<TPUClusterFormationPass> {
public:
explicit TPUClusterFormationPass(bool strict_clusters)
: strict_clusters_(strict_clusters) {}
void getDependentDialects(DialectRegistry& registry) const override {
registry.insert<mlir::tf_device::TensorFlowDeviceDialect>();
}
void runOnOperation() override;
private:
bool strict_clusters_;
};
// Creates a mapping from the TPUReplicateMetadata ops `_replication_info`
// attribute to its attributes and removes the ops. If multiple
// TPUReplicateMetadata ops have the same `_replication_info` attribute, an
// error will be returned.
LogicalResult CollectMetadata(Block* block, MetadataMap* metadata_map) {
// Just look at top-level operations in the block (not nested ones)
for (Operation& op : llvm::make_early_inc_range(*block)) {
auto metadata_op = llvm::dyn_cast<mlir::TF::TPUReplicateMetadataOp>(op);
if (!metadata_op) continue;
NamedAttrList attrs(metadata_op->getAttrDictionary());
// Missing or bad `_replication_info` attribute.
auto replication_info_attr = attrs.get(mlir::TF::kReplicationInfoAttr);
if (!replication_info_attr)
return metadata_op.emitError() << kBadReplicateInfoAttrMsg;
auto replication_info_attr_str =
mlir::dyn_cast<StringAttr>(replication_info_attr);
if (!replication_info_attr_str ||
replication_info_attr_str.getValue().empty())
return metadata_op.emitError() << kBadReplicateInfoAttrMsg;
// Remove `name` attribute.
attrs.erase(StringAttr::get(metadata_op.getContext(), kNameAttr));
auto it = metadata_map->try_emplace(replication_info_attr_str.getValue(),
std::move(attrs));
// There are multiple TPUReplicateMetadata ops with the same
// `_replication_info` attribute.
if (!it.second) {
return metadata_op.emitError()
<< "multiple TPUReplicateMetadata ops with the same '"
<< mlir::TF::kReplicationInfoAttr << "' attribute '"
<< replication_info_attr_str.getValue() << "' found";
}
metadata_op.erase();
}
return success();
}
struct OpDevice {
Operation* op;
std::string device;
};
LogicalResult HasValidDeviceTypeAttribute(Block* block) {
// Use ordered set here to make error message below deterministic.
std::set<llvm::StringRef> device_types;
for (Operation& op : *block) {
// Collect device types which currently must be consistent per block
// (checked later).
if (auto device_type_attr =
op.getAttrOfType<StringAttr>(mlir::TF::kCompileDeviceTypeAttr)) {
// tf.StatefulPartitionedCall ops with and without
// _tpu_replicate attributes may exist in the same graph. Ops without
// the attribute but with _XlaMustCompile=true would have
// _xla_compile_device_type="" after
// CanonicalizeCompileAndReplicateAttributesPass. Skip empty value here.
if (!device_type_attr.getValue().empty()) {
device_types.insert(device_type_attr);
}
}
}
if (device_types.size() > 1) {
return block->getParentOp()->emitError()
<< "found different '" << mlir::TF::kCompileDeviceTypeAttr
<< "' attribute values (" << llvm::join(device_types, ",")
<< ") in same block which is not supported";
}
return success();
}
// Collects and clusters ops based on `_replication_info` attribute. Returns
// an error in case of invalid compilation or replication attribute(s).
LogicalResult CollectAndGroupClusterOps(Block* block, ClusterMap* clusters) {
LogicalResult result = HasValidDeviceTypeAttribute(block);
if (failed(result)) return result;
for (Operation& op : *block) {
LogicalResult result =
mlir::TF::HasValidCompilationAndReplicationAttributes(op);
if (failed(result)) return result;
// Skip ops with non-TPU device type, they are handled elsewhere.
auto device_type_attr =
op.getAttrOfType<StringAttr>(mlir::TF::kCompileDeviceTypeAttr);
if (device_type_attr) {
if (device_type_attr.getValue().empty()) continue;
if (device_type_attr.getValue() != mlir::TF::kTpuDevice) continue;
}
if (op.hasAttr(mlir::TF::kReplicationInfoAttr)) {
// For replicated case, borrow cluster structure from replication info.
// Following condition is already checked in
// `HasValidCompilationAndReplicationAttributes` above, assert here for
// documentation and to avoid breakage when that function is changed.
assert(op.hasAttr(mlir::TF::kCompileDeviceTypeAttr));
auto attr = op.getAttrOfType<StringAttr>(mlir::TF::kReplicationInfoAttr);
auto it = clusters->try_emplace(attr.getValue());
it.first->getSecond().insert(&op);
}
}
return success();
}
// Returns true iff `op` has a direct control dependency from (`incoming` ==
// true) or to (`incoming` == false) any op in `cluster_ops` or
// `cluster_dependent_ops`.
Operation* getOpClusterControlDependency(
Operation* op, bool incoming, const OpSetVector& cluster_ops,
const OpSetVector& cluster_dependent_ops,
const mlir::TF::SideEffectAnalysis::Info& side_effect_analysis) {
auto filter = [&](Operation* other_op) {
return cluster_ops.contains(other_op) ||
cluster_dependent_ops.contains(other_op);
};
auto other_ops =
incoming ? side_effect_analysis.DirectControlPredecessors(op, filter)
: side_effect_analysis.DirectControlSuccessors(op, filter);
if (other_ops.empty())
return nullptr;
else
return *other_ops.begin();
}
// Returns true iff `op` has a direct data dependency from (`incoming` == true
// or to (`incoming` == false) any op in `cluster_ops` or
// `cluster_dependent_ops`.
Operation* getOpClusterDataDependency(
Operation* op, bool incoming, const OpSetVector& cluster_ops,
const OpSetVector& cluster_dependent_ops) {
Operation* other_op = nullptr;
op->walk([&](Operation* inner_op) {
ValueRange values = incoming ? ValueRange(inner_op->getOperands())
: ValueRange(inner_op->getResults());
llvm::SmallVector<Operation*, 4> candidates;
for (Value value : values) {
if (incoming) {
candidates = {value.getDefiningOp()};
} else {
candidates.assign(value.getUsers().begin(), value.getUsers().end());
}
for (Operation* candidate_op : candidates) {
if (cluster_ops.contains(candidate_op) ||
cluster_dependent_ops.contains(candidate_op)) {
other_op = candidate_op;
return WalkResult::interrupt();
}
}
}
return WalkResult::advance();
});
return other_op;
}
// Collects ops that need to be moved behind the cluster due to data or control
// dependencies.
mlir::FailureOr<llvm::SmallSetVector<Operation*, 8>> CollectClusterSuccessorOps(
Block* block, const OpSetVector& cluster_ops,
const mlir::TF::SideEffectAnalysis::Info& side_effect_analysis,
bool strict_clusters) {
OpSetVector cluster_predecessor_ops;
OpSetVector cluster_successor_ops;
// Collect non-cluster ops that have a dependency to the cluster. For this
// traverse all ops from last to first cluster op and keep track of in-between
// non-cluster ops that have some outgoing (transitive) dependency to some
// cluster op (`cluster_predecessor_ops`).
auto rfront = Block::reverse_iterator(cluster_ops.front());
auto rback = Block::reverse_iterator(cluster_ops.back());
for (Operation& op : llvm::make_range(rback, rfront)) {
if (cluster_ops.contains(&op)) continue;
bool has_dependency_to_cluster =
getOpClusterDataDependency(&op, /*incoming=*/false, cluster_ops,
cluster_predecessor_ops) != nullptr ||
getOpClusterControlDependency(&op, /*incoming=*/false, cluster_ops,
cluster_predecessor_ops,
side_effect_analysis) != nullptr;
if (has_dependency_to_cluster) cluster_predecessor_ops.insert(&op);
}
// Collect non-cluster ops that have a dependency from the cluster. For this
// traverse all ops from first to last cluster op and keep track of in-between
// non-cluster ops that have some incoming (transitive) dependency from some
// cluster op (`cluster_successor_ops`).
auto front = Block::iterator(cluster_ops.front());
auto back = Block::iterator(cluster_ops.back());
for (Operation& op : llvm::make_range(front, back)) {
if (cluster_ops.contains(&op)) continue;
Operation* data_predecessor = getOpClusterDataDependency(
&op, /*incoming=*/true, cluster_ops, cluster_successor_ops);
Operation* control_predecessor = getOpClusterControlDependency(
&op, /*incoming=*/true, cluster_ops, cluster_successor_ops,
side_effect_analysis);
if (data_predecessor || control_predecessor) {
if (cluster_predecessor_ops.contains(&op)) {
// Op has a dependency from and to the cluster which is invalid. Instead
// of erroring out we don't add the op to `cluster_successor_ops` which
// is in line with previous behavior when certain control dependencies
// were not considered.
// TODO(b/216706460) Establish some contract here: Should we expect only
// valid clusters, or should we split clusters accordingly? The latter
// might have runtime impact for existing models.
// We should make this message an error once there is such a contract
// and once existing cases have been fixed.
mlir::InFlightDiagnostic error =
strict_clusters ? mlir::emitError(op.getLoc(), "")
: mlir::emitWarning(op.getLoc(), "");
error << "Op has cyclic dependency with a compilation cluster:\n";
error << "The cluster depends on\n";
error << op.getName() << "\n"
<< "which is outside of cluster, but itself depends ";
if (data_predecessor) {
error << "on\n" << data_predecessor->getName() << "\n";
} else if (control_predecessor) {
error << "(via control) on\n"
<< control_predecessor->getName() << "\n";
}
if (cluster_ops.contains(data_predecessor) ||
cluster_ops.contains(control_predecessor))
error << "which is inside the cluster.\n";
else
error << "which the cluster depends on.\n";
} else {
cluster_successor_ops.insert(&op);
}
}
}
return cluster_successor_ops;
}
// Collects results and associated types of the cluster that are used outside of
// the cluster. These results and types are used to create the clusters
// `tf_device.cluster` and associated terminator. Results that have no uses
// outside of the cluster (i.e. results of ops in the cluster are only consumed
// by other ops in the cluster) are pruned.
llvm::SmallVector<Value, 8> CollectClusterResults(
Block* block, const OpSetVector& cluster_ops) {
llvm::SmallVector<Value, 8> results;
for (Operation* op : cluster_ops) {
for (Value result : op->getResults()) {
for (Operation* user : result.getUsers()) {
// Check if user is not an op in the cluster.
if (cluster_ops.count(block->findAncestorOpInBlock(*user)) == 0) {
results.push_back(result);
break;
}
}
}
}
return results;
}
// Creates a `tf_device.cluster` to wrap cluster ops.
mlir::tf_device::ClusterOp CreateClusterOp(
Block* block, const OpSetVector& cluster_ops, llvm::ArrayRef<Value> results,
llvm::ArrayRef<Operation*> cluster_successor_ops) {
// `tf_device.cluster` will be placed at where the last op of the cluster is.
Operation* last_cluster_op = cluster_ops.back();
OpBuilder builder(last_cluster_op);
llvm::SmallVector<Type, 8> result_types;
for (Value result : results) result_types.push_back(result.getType());
auto cluster = mlir::tf_device::ClusterOp::create(
builder, last_cluster_op->getLoc(), result_types);
Block* body = new Block;
cluster.getBody().push_back(body);
// Move cluster ops to the cluster body. Also remove `_replication_info` and
// `device` attribute from ops in the cluster when that information is
// redundant will the `tf_device.cluster`. Do this for all ops including
// nested ops.
for (Operation* cluster_op : cluster_ops) {
cluster_op->moveBefore(body, body->end());
cluster_op->walk([&](Operation* inner_op) {
inner_op->removeAttr(mlir::TF::kReplicationInfoAttr);
inner_op->removeAttr(mlir::TF::kCompileDeviceTypeAttr);
if (auto attr = inner_op->getAttrOfType<StringAttr>(kDeviceAttr)) {
// Preserve device attribute if the op is placed on a replicated core
// device. Device attribute is used to infer the appropriate sharding
// within TPUs for this op.
// TODO(b/183598857): Use explicit sharding ops from the front-end.
// For example, dequeue ops generated by
// tensorflow/python/tpu/tpu_feed.py
if (!tensorflow::IsTPUReplicatedCore(attr.getValue())) {
inner_op->removeAttr(kDeviceAttr);
}
}
});
}
// Add terminator.
builder.setInsertionPointToEnd(body);
mlir::tf_device::ReturnOp::create(builder, last_cluster_op->getLoc(),
results);
// Replaces uses of cluster ops results outside of cluster with the associated
// `tf_device.cluster` results.
for (auto ret_vals : llvm::zip(results, cluster.getResults())) {
Value old_ret = std::get<0>(ret_vals);
Value new_ret = std::get<1>(ret_vals);
for (auto& use : llvm::make_early_inc_range(old_ret.getUses())) {
Operation* user = use.getOwner();
if (!body->findAncestorOpInBlock(*user)) use.set(new_ret);
}
}
// Move ops that depend on something in the cluster behind the cluster.
Operation* op_after_cluster = cluster.getOperation()->getNextNode();
for (Operation* op : cluster_successor_ops) op->moveBefore(op_after_cluster);
return cluster;
}
// Returns an op of the given type that uses the result, along with
// a list of identity ops along the way.
template <typename T>
std::tuple<T, llvm::SmallVector<mlir::TF::IdentityOp, 4>> GetSingleUserOfType(
OpResult result) {
llvm::SmallVector<mlir::TF::IdentityOp, 4> identity_ops;
do {
Operation* user = result.hasOneUse() ? *result.getUsers().begin() : nullptr;
if (auto t = llvm::dyn_cast_or_null<T>(user)) {
return std::make_tuple(t, identity_ops);
} else if (auto identity =
llvm::dyn_cast_or_null<mlir::TF::IdentityOp>(user)) {
identity_ops.emplace_back(identity);
result = identity->getResult(0);
} else {
result = OpResult(); // reset to stop iterating
}
} while (result);
return std::make_tuple(T(), identity_ops);
}
using PartitionedClusterOutputMap = absl::flat_hash_map<
uint64_t, llvm::SmallVector<mlir::TF::TPUPartitionedOutputV2Op, 8>>;
// Returns the partitioned output ops from the cluster if there are any,
// along with any single user identity ops between them. Not all outputs
// of a cluster must be partitioned, so the output is a map from cluster
// output ids to ops.
std::tuple<PartitionedClusterOutputMap,
llvm::SmallVector<mlir::TF::IdentityOp, 8>>
GetPartitionedOutputsAndIdentityOps(mlir::tf_device::ClusterOp cluster) {
PartitionedClusterOutputMap partitioned_outputs;
llvm::SmallVector<mlir::TF::IdentityOp, 8> erase_list;
for (auto [cluster_result_id, cluster_result] :
llvm::enumerate(cluster.getResults())) {
auto [replicated_output, _] =
GetSingleUserOfType<mlir::TF::TPUReplicatedOutputOp>(cluster_result);
if (replicated_output) {
for (OpResult per_replica_result : replicated_output->getResults()) {
auto [partitioned_output, id_ops] =
GetSingleUserOfType<mlir::TF::TPUPartitionedOutputV2Op>(
per_replica_result);
if (partitioned_output) {
erase_list.insert(erase_list.end(), id_ops.begin(), id_ops.end());
partitioned_outputs[cluster_result_id].emplace_back(
partitioned_output);
}
}
}
}
return std::forward_as_tuple(partitioned_outputs, erase_list);
}
// Inlines the partitioned output ops into the cluster, and updates
// their users to point to the replicate op instead.
Operation* BuildPartitionedOutputs(
OpBuilder& builder, mlir::tf_device::ClusterOp cluster,
mlir::tf_device::ReplicateOp replicate_op,
PartitionedClusterOutputMap& partitioned_outputs,
llvm::SmallVector<mlir::TF::IdentityOp, 8>& erase_list,
llvm::SmallVector<Type, 8>& result_types, int num_replicas) {
Operation* result_op;
llvm::SmallVector<Value, 8> results;
uint64_t num_results = cluster.getNumResults();
for (uint64_t result_id = 0; result_id < num_results; ++result_id) {
auto search = partitioned_outputs.find(result_id);
if (search == partitioned_outputs.end()) {
// If the output is not partitioned, directly pass it through.
results.emplace_back(cluster.getResult(result_id));
continue;
}
// Otherwise, "inline" the partitioned output ops by:
// - Building a new op within the cluster.
// - Replacing all the uses of the original ops with the cluster's outputs.
llvm::SmallVector<mlir::TF::TPUPartitionedOutputV2Op, 8>& ops =
search->second;
for (auto [replica_id, partitioned_output] : llvm::enumerate(ops)) {
for (auto [core_id, result] :
llvm::enumerate(partitioned_output->getResults())) {
// outputs from replicate op are interleaved:
// [(replica:0,core:0), (replica:1,core:0), ...,
// (replica:0,core:1), (replica:1,core:1), ...]
uint64_t output_id =
core_id * num_replicas + replica_id + results.size();
result.replaceAllUsesWith(replicate_op.getResult(output_id));
}
}
// Assume all the replicas have the same structure.
mlir::TF::TPUPartitionedOutputV2Op first_op = *(ops.begin());
mlir::ArrayAttr dims = first_op.getPartitionDimsAttr();
StringAttr sharding = first_op.get_XlaShardingAttr();
StringAttr sharding_v2 = first_op.get_XlaShardingV2Attr();
Operation::result_type_range output_types = first_op.getResultTypes();
result_op = mlir::TF::TPUPartitionedOutputV2Op::create(
builder, replicate_op.getLoc(), output_types,
cluster.getResult(result_id), dims, sharding, sharding_v2);
results.insert(results.end(), result_op->getResults().begin(),
result_op->getResults().end());
}
// Once we've accumulated all the cluster's results, build a return op.
mlir::tf_device::ReturnOp::create(builder, result_op->getLoc(), results);
// Then erase all the identity and partitioned output ops.
for (const auto& [_, ops] : partitioned_outputs) {
for (mlir::TF::TPUPartitionedOutputV2Op op : ops) {
op->erase();
}
}
for (mlir::TF::IdentityOp to_erase : erase_list) {
to_erase->erase();
}
return result_op;
}
// Return the cluster's per-replica result type, converting any full-shaped
// tensor types into sharded-shaped ones if they're partitioned.
llvm::SmallVector<Type, 8> GetClusterResultTypes(
mlir::tf_device::ClusterOp cluster,
const PartitionedClusterOutputMap& partitioned_outputs) {
llvm::SmallVector<Type, 8> result_types;
Operation::result_type_range cluster_result_types = cluster.getResultTypes();
if (partitioned_outputs.empty()) {
// Directly pass through the cluster's outputs if none are partitioned.
result_types.insert(result_types.end(), cluster_result_types.begin(),
cluster_result_types.end());
} else {
// For each output of the cluster...
for (auto [output_id, result_type] :
llvm::enumerate(cluster_result_types)) {
auto search = partitioned_outputs.find(output_id);
if (search == std::end(partitioned_outputs)) {
// If it's not partitioned, directly pass it through.
result_types.emplace_back(result_type);
} else {
// Otherwise, pass through the result shard types.
Operation::result_type_range partitioned_result_types =
(*search->second.begin())->getResultTypes();
result_types.insert(result_types.end(),
partitioned_result_types.begin(),
partitioned_result_types.end());
}
}
}
return result_types;
}
// Creates a `tf_device.replicate` to represent replication for the cluster, if
// necessary. Erases Identity ops between partitioned and replicated output ops.
LogicalResult ReplicateCluster(mlir::tf_device::ClusterOp cluster,
int num_replicas, int num_cores_per_replica) {
OpBuilder builder(cluster);
auto [partitioned_outputs, erase_list] =
GetPartitionedOutputsAndIdentityOps(cluster);
for (auto [_, ops] : partitioned_outputs) {
if (!(ops.empty() || ops.size() == num_replicas)) {
return (ops.begin())->emitOpError()
<< "expected zero or " << num_replicas
<< " 'TPUPartitionedOutput' op(s), instead got "
<< partitioned_outputs.size();
}
}
// No need to replicate.
if (num_replicas == 1) {
// Collapse all the Identity ops between the TRO and TPO ops.
if (!partitioned_outputs.empty()) {
for (mlir::TF::IdentityOp to_erase : erase_list) {
Value in = to_erase->getOperand(0);
OpResult out = to_erase->getResult(0);
out.replaceAllUsesWith(in);
to_erase->erase();
}
}
return success();
}
if (num_replicas < 1)
return cluster.emitError() << "requires '" << kNumReplicasAttr
<< "' int attribute to be at least 1";
LogicalResult status = success();
// Collect all used TPUReplicatedInput ops.
llvm::SmallVector<mlir::TF::TPUReplicatedInputOp, 8> replicated_input_ops;
llvm::SmallSet<mlir::TF::TPUReplicatedInputOp, 8> seen_ops;
mlir::visitUsedValuesDefinedAbove(
cluster.getBody(), cluster.getBody(), [&](mlir::OpOperand* operand) {
Operation* def = operand->get().getDefiningOp();
if (auto ri =
llvm::dyn_cast_or_null<mlir::TF::TPUReplicatedInputOp>(def)) {
if (!seen_ops.contains(ri)) {
seen_ops.insert(ri);
replicated_input_ops.push_back(ri);
}
}
// When model parallelism is used in conjunction with data parallelism
// for resource inputs, we need to collect the per replica resource
// inputs from input to `tf.TPUPartitionedInputV2` ops.
if (auto pi = llvm::dyn_cast_or_null<mlir::TF::TPUPartitionedInputV2Op>(
def)) {
if (pi->getNumOperands() != num_cores_per_replica)
status = pi.emitOpError()
<< "requires " << num_cores_per_replica
<< " operands but found " << pi->getNumOperands();
for (auto operand : pi.getInputs()) {
if (auto ri =
llvm::dyn_cast_or_null<mlir::TF::TPUReplicatedInputOp>(
operand.getDefiningOp())) {
if (!seen_ops.contains(ri)) {
seen_ops.insert(ri);
replicated_input_ops.push_back(ri);
}
}
}
}
});
if (failed(status)) return mlir::failure();
// Indices of the replicate op's arguments that are mirrored variables.
llvm::SmallVector<int64_t, 8> mirrored_variable_indices;
// Check if number of operands of each used TPUReplicatedInput op matches
// `num_replicas` or 1. Collect all their operands and associated type for
// creating the replicate op.
llvm::SmallVector<std::pair<ValueRange, Type>, 8> replicated_inputs;
llvm::SmallVector<Value, 8> packed_inputs;
llvm::SmallVector<mlir::TF::TPUReplicatedInputOp, 8> replicated_ops;
llvm::SmallVector<mlir::TF::TPUReplicatedInputOp, 8> packed_ops;
for (const auto& pos_and_input : llvm::enumerate(replicated_input_ops)) {
auto input = pos_and_input.value();
bool is_packed = input.getIsPacked();
const int num_operands = input->getNumOperands();
int num_inputs = is_packed ? 1 : num_replicas;
if (num_operands != num_inputs)
return input->emitOpError() << "requires " << num_inputs << " operands";
if (is_packed) {
packed_inputs.push_back(input->getOperand(0));
packed_ops.push_back(input);
} else {
replicated_inputs.push_back(
{input->getOperands(), input->getOperand(0).getType()});
replicated_ops.push_back(input);
}
}
// Create `ordered_tpu_replicate_inputs` which contains the final ordered
// replicate inputs. All packed arguments are moved to the end of the arg
// list.
llvm::SmallVector<mlir::TF::TPUReplicatedInputOp, 8>
ordered_tpu_replicate_inputs = replicated_ops;
ordered_tpu_replicate_inputs.append(packed_ops.begin(), packed_ops.end());
// Assign `mirrored_variable_indices` based on the ordered replicated inputs.
for (const auto& pos_and_input :
llvm::enumerate(ordered_tpu_replicate_inputs)) {
auto tpu_replicated_input = pos_and_input.value();
if (tpu_replicated_input.getIsMirroredVariable()) {
mirrored_variable_indices.push_back(pos_and_input.index());
}
}
// Create replicate op.
auto result_types = GetClusterResultTypes(cluster, partitioned_outputs);
auto replicate_op = mlir::tf_device::ReplicateOp::create(
builder, cluster.getLoc(), num_replicas,
llvm::SmallDenseMap<llvm::StringRef,
llvm::SmallVector<llvm::StringRef, 4>>(),
replicated_inputs, packed_inputs, result_types);
if (!mirrored_variable_indices.empty())
replicate_op->setAttr(kMirroredVariableIndicesAttr,
builder.getI64ArrayAttr(mirrored_variable_indices));
// Replace replicated cluster results with replicate op results.
uint64_t offset = 0;
for (auto [idx, result] : llvm::enumerate(cluster.getResults())) {
if (partitioned_outputs.contains(idx)) {
// Partitioned output propagation happens in BuildPartitionedOutputs.
offset += num_replicas * num_cores_per_replica;
continue;
}
auto replicate_outputs = llvm::make_range(
std::next(replicate_op.result_begin(), offset),
std::next(replicate_op.result_begin(), offset + num_replicas));
for (auto& use : llvm::make_early_inc_range(result.getUses())) {
Operation* def = use.getOwner();
if (!llvm::isa<mlir::TF::TPUReplicatedOutputOp>(def)) {
// If user is not a `tf.TPUReplicatedOutput`, simply forward the first
// replica output. Certain Graphs under V1 create `tf.Identity` users of
// replicated ops to pin the TPU computation for execution.
use.set(*replicate_outputs.begin());
continue;
}
const int def_num_results = def->getNumResults();
if (def_num_results != num_replicas)
return def->emitOpError() << "requires " << num_replicas << " results";
def->replaceAllUsesWith(replicate_outputs);
}
offset += num_replicas;
}
// Collect all `tf.TPUPartitionedInputV2` ops to be moved inside the
// `tf_device.replicate` later.
llvm::SmallSet<Operation*, 4> partitioned_inputs;
for (auto input_and_block_arg :
llvm::zip(ordered_tpu_replicate_inputs,
replicate_op.GetBody().getArguments())) {
mlir::TF::TPUReplicatedInputOp input = std::get<0>(input_and_block_arg);
Value block_arg = std::get<1>(input_and_block_arg);
mlir::replaceAllUsesInRegionWith(input->getResult(0), block_arg,
cluster.getBody());
// Update replicated input use in tf.TPUPartitionedInputV2 op.
for (auto& use : input->getUses()) {
auto pi =
llvm::dyn_cast<mlir::TF::TPUPartitionedInputV2Op>(use.getOwner());
if (pi) {
pi.setOperand(use.getOperandNumber(), block_arg);
partitioned_inputs.insert(pi.getOperation());
}
}
}
// Create terminator for replicate op and move `tf_device.cluster` and
// `tf.TPUPartitionedInputV2`(s) into replicate body.
builder.setInsertionPointToEnd(&replicate_op.GetBody());
Operation* result_op;
if (!partitioned_outputs.empty()) {
result_op = BuildPartitionedOutputs(builder, cluster, replicate_op,
partitioned_outputs, erase_list,
result_types, num_replicas);
} else {
result_op = mlir::tf_device::ReturnOp::create(
builder, replicate_op.getLoc(), cluster.getResults());
}
for (auto pi : partitioned_inputs) pi->moveBefore(result_op);
cluster.getOperation()->moveBefore(result_op);
return success();
}
// Forms clusters with ops of the same `_replication_info` attribute under a
// block.
//
// For a given block, clusters are formed via grouping ops by
// `_replication_info` attributes. For every cluster formed:
// 1. Find associated TPUReplicateMetadata attributes with the same
// `_replication_info` attribute.
// 2. Find users not in cluster that are interleaved between cluster ops.
// 3. Find external uses of cluster ops.
// 4. Create `tf_device.cluster` with results consisting of the external uses
// of cluster ops determined at 3.
// 5. Move cluster ops to `tf_device.cluster` body.
// 6. Replace external uses of cluster ops uses with `tf_device.cluster`
// results.
// 7. Move users from 2 to after the `tf_device.cluster`.
// 8. Wrap cluster (`tf_device.cluster`) in a `tf_device.replicate` if
// attribute `num_replicas` is greater than 1.
// 9. Copy over TPUReplicateMetadata attributes to `tf_device.cluster`.
LogicalResult FormClustersInBlock(
Block* block,
const mlir::TF::SideEffectAnalysis::Info& side_effect_analysis,
bool strict_clusters) {
MetadataMap metadata_map;
LogicalResult result = CollectMetadata(block, &metadata_map);
if (failed(result)) return result;
// If there is no TPUReplicateMetadata op in this block, process blocks in
// regions attached to the op's in the block.
if (metadata_map.empty()) {
for (Operation& op : *block) {
for (Region& region : op.getRegions()) {
if (!llvm::hasSingleElement(region))
return op.emitOpError("Expected single block region");
if (failed(FormClustersInBlock(&region.front(), side_effect_analysis,
strict_clusters)))
return mlir::failure();
}
}
return success();
}
ClusterMap clusters;
result = CollectAndGroupClusterOps(block, &clusters);
if (failed(result)) return result;
for (const auto& cluster_metadata_and_ops : clusters) {
const auto& cluster_ops = cluster_metadata_and_ops.getSecond();
auto cluster_metadata =
metadata_map.find(cluster_metadata_and_ops.getFirst());
// llvm::errs() << __func__ << "\n";
// No TPUReplicateMetadata for a `_replication_info` attribute.
if (cluster_metadata == metadata_map.end()) {
block->getParentOp()->emitWarning()
<< "TPUReplicateMetadata for associated '"
<< mlir::TF::kReplicationInfoAttr << "' attribute '"
<< cluster_metadata_and_ops.getFirst() << "' is missing";
continue;
}
auto status = CollectClusterSuccessorOps(
block, cluster_ops, side_effect_analysis, strict_clusters);
if (failed(status)) return status;
OpSetVector cluster_successor_ops = *status;
llvm::SmallVector<Value, 8> results =
CollectClusterResults(block, cluster_ops);
mlir::tf_device::ClusterOp cluster = CreateClusterOp(
block, cluster_ops, results, cluster_successor_ops.getArrayRef());
auto num_replicas = cluster_metadata->getSecond().get(kNumReplicasAttr);
if (!num_replicas || !mlir::isa<mlir::IntegerAttr>(num_replicas))
return cluster.emitError()
<< "requires '" << kNumReplicasAttr << "' int attribute";
int num_cores_per_replica = 1;
auto num_cores_per_replica_attr = mlir::dyn_cast_or_null<mlir::IntegerAttr>(
cluster_metadata->getSecond().get(kNumCoresPerReplicaAttr));
if (num_cores_per_replica_attr)
num_cores_per_replica = num_cores_per_replica_attr.getInt();
if (failed(ReplicateCluster(
cluster, mlir::cast<mlir::IntegerAttr>(num_replicas).getInt(),
num_cores_per_replica)))
return mlir::failure();
// Copy TPUReplicateMetadata attributes to `tf_device.cluster`.
cluster->setAttrs(
cluster_metadata->second.getDictionary(cluster.getContext()));
// Exclude `num_replicas` as cluster should be replicated if necessary.
cluster->removeAttr(kNumReplicasAttr);
}
return success();
}
LogicalResult FormClustersInFunction(
mlir::func::FuncOp func,
const mlir::TF::SideEffectAnalysis::Info& side_effect_analysis,
bool strict_clusters) {
if (!llvm::hasSingleElement(func))
return func.emitOpError("Expecting a single block function");
if (failed(FormClustersInBlock(&func.front(), side_effect_analysis,
strict_clusters)))
return mlir::failure();
// Remove TPUReplicatedInput and TPUReplicatedOutput nodes.
auto remove_result = func.walk([&](Operation* op) {
if (!llvm::isa<mlir::TF::TPUReplicatedInputOp,
mlir::TF::TPUReplicatedOutputOp>(op))
return WalkResult::advance();
// Forward operand to result. When `num_replicas` attribute is 1, no
// `tf_device.replicate` is created and replicated (1) operands/results are
// untouched.
if (op->getNumOperands() == 1 && op->getNumResults() == 1)
op->getResult(0).replaceAllUsesWith(op->getOperand(0));
// Leftover TPUReplicatedInput/TPUReplicatedOutput that are not of
// `num_replicas` to 1.
if (!op->use_empty()) {
op->emitOpError() << "is expected to have no uses, but it is operand#"
<< op->use_begin()->getOperandNumber() << " of "
<< *op->use_begin()->getOwner();
return WalkResult::interrupt();
}
op->erase();
return WalkResult::advance();
});
return mlir::failure(remove_result.wasInterrupted());
}
void TPUClusterFormationPass::runOnOperation() {
// Attributes on tf.Constant aren't reliable: CSE will merge ConstantLike ops
// with the same value (but different attributes!) into the same tf.Const
// definition, potentially leading to bogus _replication_info attributes. So
// we just scrub all tf.Constants of all extra attributes.
// TODO(kramm): Remove this once tf.Const's folder is aware of extra
// attributes.
auto value_str_attr = StringAttr::get(&getContext(), "value");
getOperation().walk([&](mlir::TF::ConstOp cst) {
auto dict = cst->getAttrDictionary();
if (dict.size() == 1) {
return; // Optimization. Assume the one attribute is "value".
}
// Recreate the attributes dictionary to only contain "value".
NamedAttrList attributes;
attributes.append(NamedAttribute(value_str_attr, cst->getAttr("value")));
cst->setAttrs(attributes.getDictionary(&getContext()));
});
auto& side_effect_analysis = getAnalysis<mlir::TF::SideEffectAnalysis>();
for (auto func : getOperation().getOps<mlir::func::FuncOp>())
if (!func.isExternal() &&
failed(FormClustersInFunction(
func, side_effect_analysis.GetAnalysisForFunc(func),
strict_clusters_)))
return signalPassFailure();
}
} // anonymous namespace
std::unique_ptr<mlir::OperationPass<ModuleOp>> CreateTPUClusterFormationPass(
bool strict_clusters) {
return std::make_unique<TPUClusterFormationPass>(strict_clusters);
}
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/passes/tpu_sharding_identification_pass.cc
+814
View File
@@ -0,0 +1,814 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <type_traits>
#include <utility>
#include <variant>
#include "absl/algorithm/container.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/FormatVariadic.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Block.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Interfaces/CallInterfaces.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_traits.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/xla_sharding_util.h"
#include "xla/hlo/builder/sharding_builder.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/statusor.h"
#include "xla/xla_data.pb.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {
using OpShardingVariant = std::variant<mlir::Operation*, llvm::StringRef>;
using OpShardingVector = llvm::SmallVector<OpShardingVariant, 8>;
using OptionalOpShardingVector =
llvm::SmallVector<std::optional<OpShardingVariant>, 8>;
using llvm::StringRef;
using mlir::Block;
using mlir::BlockArgument;
using mlir::BoolAttr;
using mlir::Builder;
using mlir::IntegerAttr;
using mlir::LogicalResult;
using mlir::ModuleOp;
using mlir::Operation;
using mlir::OpOperand;
using mlir::OpResult;
using mlir::RankedTensorType;
using mlir::StringAttr;
using mlir::Value;
using mlir::WalkResult;
constexpr char kReplicateSharding[] = "";
constexpr char kShardingAttr[] = "mhlo.sharding";
constexpr char kUseSpmdAttr[] = "use_spmd_for_xla_partitioning";
constexpr char kAliasingAttr[] = "tf.aliasing_output";
constexpr char kNumCoresPerReplicaAttr[] = "num_cores_per_replica";
const char kShardingAttribute[] = "_XlaSharding";
const char kShardingAttributeV2[] = "_XlaShardingV2";
#define GEN_PASS_DEF_TPUSHARDINGIDENTIFICATIONPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc"
struct TPUShardingIdentificationPass
: public impl::TPUShardingIdentificationPassBase<
TPUShardingIdentificationPass> {
void runOnOperation() final;
};
std::string CreateMissingAttributeMsg(llvm::StringRef attribute) {
return llvm::formatv("requires attribute '{0}'", attribute).str();
}
// Returns nullptr if the op does not have a sharding attribute.
template <typename PartitionedOp>
mlir::Operation* NullUnlessSharded(PartitionedOp op) {
return op.get_XlaSharding() ? op : nullptr;
}
// Returns true if unitary op has one of the traits that meets the requirements
// for sharding, otherwise returns false.
bool UnaryOpHasTraitsForSharding(Operation* op) {
// Trait "SameOperandsAndResultTypeResolveRef" for Cast, real/imag, etc.
if (op->hasTrait<mlir::OpTrait::TF::SameOperandsAndResultTypeResolveRef>())
return true;
// Trait "SameOperandsAndResultType" for Exp, ceil, etc.
if (op->hasTrait<mlir::OpTrait::SameOperandsAndResultType>()) return true;
// Trait "OperandsSameAsResultsTypeOrRef" for Identity.
if (op->hasTrait<mlir::OpTrait::TF::OperandsSameAsResultsTypeOrRef>())
return true;
return false;
}
// Returns true if binary op has one of the traits that meets the requirements
// for sharding, otherwise returns false.
bool BinaryOpHasTraitsForSharding(Operation* op) {
// Trait "CwiseBinary" and "SameOperandsAndResultElementTypeResolveRef" for
// AddV2, Sub, etc.
if (op->hasTrait<
mlir::OpTrait::TF::SameOperandsAndResultElementTypeResolveRef>() &&
op->hasTrait<mlir::OpTrait::TF::CwiseBinary>())
return true;
return false;
}
bool DoTypesHavePartialSameShape(Value value_0, Value value_1) {
auto shape_0 =
mlir::dyn_cast_or_null<mlir::RankedTensorType>(value_0.getType());
auto shape_1 =
mlir::dyn_cast_or_null<mlir::RankedTensorType>(value_1.getType());
if (shape_0 && shape_1) {
if (shape_0.hasStaticShape() && shape_1.hasStaticShape())
return shape_0.getShape() == shape_1.getShape();
int i = 0, j = 0;
while (i < shape_0.getShape().size() && j < shape_1.getShape().size()) {
if (shape_0.getShape()[i] != shape_1.getShape()[j] &&
!shape_0.isDynamicDim(i) && !shape_1.isDynamicDim(j)) {
return false;
}
if (shape_0.getShape()[i] == shape_1.getShape()[j]) {
i++;
j++;
} else {
if (shape_0.isDynamicDim(i)) {
i++;
}
if (shape_1.isDynamicDim(j)) {
j++;
}
}
}
return i == shape_0.getShape().size() && j == shape_1.getShape().size();
}
return false;
}
// Returns a TPUPartitionedInput op connected to a `tf_device.cluster_func`
// operand value if it has an XLA sharding. If value is a resource type then
// TPUPartitionedInput op will be connected to a ReadVariable op that feeds into
// a `tf_device.cluster_func`.
mlir::Operation* GetXlaShardingFromOperand(Value value) {
Value value_to_visit = value;
if (auto read_var = value_to_visit.getDefiningOp<mlir::TF::ReadVariableOp>())
value_to_visit = read_var.getResource();
if (auto partitioned_input =
value_to_visit.getDefiningOp<mlir::TF::TPUPartitionedInputV2Op>()) {
return NullUnlessSharded(partitioned_input);
}
return nullptr;
}
// Returns the op sharding attribute from a partitioned operator.
std::optional<StringRef> GetXlaShardingFromOperator(mlir::Operation* op) {
if (auto partitioned_output =
llvm::dyn_cast<mlir::TF::TPUPartitionedOutputV2Op>(op)) {
if (partitioned_output.get_XlaShardingV2().has_value()) {
return partitioned_output.get_XlaShardingV2();
}
return partitioned_output.get_XlaSharding();
} else if (auto partitioned_input =
llvm::dyn_cast<mlir::TF::TPUPartitionedInputV2Op>(op)) {
if (partitioned_input.get_XlaShardingV2().has_value()) {
return partitioned_input.get_XlaShardingV2();
}
return partitioned_input.get_XlaSharding();
}
return std::nullopt;
}
// Returns the sharding string from a op-sharding variant if it is available.
std::optional<StringRef> GetShardingStringFromVariant(
const OpShardingVariant& sharding_or_op) {
return std::visit(
[](auto&& sharding_or_op) -> std::optional<StringRef> {
using T = std::decay_t<decltype(sharding_or_op)>;
if constexpr (std::is_same_v<T, StringRef>) {
return sharding_or_op;
} else {
return GetXlaShardingFromOperator(sharding_or_op);
}
},
sharding_or_op);
}
// Returns the sharding from a op-sharding variant if it is available and valid.
std::optional<xla::OpSharding> GetShardingFromVariant(
const OpShardingVariant& sharding_or_op) {
xla::OpSharding sharding;
const auto sharding_string = GetShardingStringFromVariant(sharding_or_op);
if (!sharding_string) return std::nullopt;
if (tensorflow::DecodeShardingAttribute(sharding_string->str(), sharding,
false)
.failed()) {
return std::nullopt;
}
return sharding;
}
// Converts an op-sharding vector into a string attr using the builder.
mlir::ArrayAttr GetStrArrayAttr(Builder* builder,
const OpShardingVector& vect) {
llvm::SmallVector<mlir::Attribute, 8> strings;
for (const auto& sharding_or_op : vect) {
if (const auto sharding = GetShardingStringFromVariant(sharding_or_op)) {
strings.emplace_back(builder->getStringAttr(*sharding));
}
}
return builder->getArrayAttr(strings);
}
// Verify whether the given sharding can be applied to the given (tensor) type.
// (A bad sharding might mean failing tf.Split ops if the graph later executes
// on CPU)
// If the sharding is incorrect, return failure. If it's good, or if we can't
// verify it, return success.
LogicalResult VerifySharding(mlir::Type type,
const OpShardingVariant& sharding_or_op) {
auto* partitioned_op =
std::holds_alternative<mlir::Operation*>(sharding_or_op)
? std::get<mlir::Operation*>(sharding_or_op)
: nullptr;
const auto sharding = GetShardingFromVariant(sharding_or_op);
if (!sharding || sharding->type() != xla::OpSharding::OTHER) {
// Some test cases use \01\02\03 as sharding, to test propagation. Treat
// a non-proto sharding as valid, and don't verify further. We also only
// verify shardings that actually break a tensor apart.
return mlir::success();
}
if (RankedTensorType ranked_type = mlir::dyn_cast<RankedTensorType>(type)) {
const int64_t tensor_rank = ranked_type.getRank();
int tile_assignment_rank = sharding->tile_assignment_dimensions_size();
// When a tensor is partial or subgroup tiled, its tile assignment will
// have one or more dimension(s) than its rank; so, we subtract them to
// determine which rank the sharding is compatible with.
tile_assignment_rank -= (int)sharding->replicate_on_last_tile_dim();
tile_assignment_rank -= sharding->last_tile_dims_size();
if (tensor_rank < tile_assignment_rank) {
if (partitioned_op) {
partitioned_op->emitError()
<< "tensor of type " << ranked_type << " (rank=" << tensor_rank
<< ") sharded in " << (tile_assignment_rank - tensor_rank)
<< " extra dimension(s) by: " << sharding->DebugString();
}
return mlir::failure();
}
}
return mlir::success();
}
// Verify sharding for all arguments and return values.
LogicalResult VerifyShardings(mlir::func::FuncOp func,
const OpShardingVector& sharding_for_args,
const OpShardingVector& sharding_for_rets) {
Block& function_block = func.front();
for (auto sharding_and_arg :
llvm::zip(sharding_for_args, function_block.getArguments())) {
const auto& sharding = std::get<0>(sharding_and_arg);
BlockArgument arg = std::get<1>(sharding_and_arg);
if (failed(VerifySharding(arg.getType(), sharding))) return mlir::failure();
}
Operation* terminator = function_block.getTerminator();
for (auto sharding_and_retval :
llvm::zip(sharding_for_rets, terminator->getOpOperands())) {
const auto& sharding = std::get<0>(sharding_and_retval);
OpOperand& retval = std::get<1>(sharding_and_retval);
if (failed(VerifySharding(retval.get().getType(), sharding)))
return mlir::failure();
}
return mlir::success();
}
// Assign the logical device if an op has an attribute `TPU_REPLICATED_CORE:n`,
// the corresponding input sharding arg will be associated with
// logical device `n`.
std::optional<llvm::StringRef> AssignLogicalDeviceFromTPUReplicatedCoreAttr(
Operation* op, const llvm::SmallVector<std::string>& logical_device_vec) {
if (auto device = op->getAttrOfType<StringAttr>("device")) {
if (!device.getValue().empty() && !device.getValue().str().empty()) {
tensorflow::DeviceNameUtils::ParsedName name;
if (tensorflow::DeviceNameUtils::ParseFullName(device.str(), &name)) {
if (name.type == "TPU_REPLICATED_CORE") {
// TODO(hanxiongwang): Add check for out of bound of name.id
return llvm::StringRef(logical_device_vec[name.id]);
}
}
}
}
return std::nullopt;
}
absl::StatusOr<std::optional<llvm::StringRef>> GetXlaShardingFromShardingOp(
mlir::TF::XlaShardingOp sharding) {
TF_ASSIGN_OR_RETURN(auto attr, GetXlaShardingAttrFromShardingOp(sharding));
return attr ? ::std::optional<::llvm::StringRef>(attr.getValue())
: (::std::nullopt);
}
// Returns XLA sharding from a XlaSharding op connected to an argument value. If
// value is a resource type then XlaSharding op will be connected to a
// ReadVariable op. XlaSharding op may be direct user of inputs but it may also
// be followed by an Identity op and, in the case where bfloat16 type is used,
// Cast op may be added right after the input.
//
// TODO(hongjunchoi): Add logic to parse XlaSharding op inside control flow (If,
// Case) ops and Caller return values.
// TODO(hongjunchoi): Consider explicitly checking op patterns to detect sharded
// inputs.
absl::StatusOr<std::optional<llvm::StringRef>> GetXlaShardingFromArg(
Value value, const llvm::SmallVector<std::string>& logical_device_vec) {
llvm::SmallPtrSet<Value, 4> visited_values;
llvm::SmallVector<Value, 4> values_to_visit{value};
while (!values_to_visit.empty()) {
llvm::SmallVector<Value, 4> next_values_to_visit;
for (Value value_to_visit : values_to_visit) {
if (!visited_values.insert(value_to_visit).second) continue;
for (auto& use : value_to_visit.getUses()) {
Operation* owner = use.getOwner();
if (auto sharding = llvm::dyn_cast<mlir::TF::XlaShardingOp>(owner))
return GetXlaShardingFromShardingOp(sharding);
if (auto logical_device = AssignLogicalDeviceFromTPUReplicatedCoreAttr(
owner, logical_device_vec)) {
return logical_device;
}
if (auto while_op = llvm::dyn_cast<mlir::TF::WhileRegionOp>(owner)) {
const int operand_number = use.getOperandNumber();
next_values_to_visit.push_back(
while_op.getCond().front().getArgument(operand_number));
next_values_to_visit.push_back(
while_op.getBody().front().getArgument(operand_number));
continue;
}
if (UnaryOpHasTraitsForSharding(owner)) {
next_values_to_visit.push_back(use.getOwner()->getResult(0));
continue;
}
if (BinaryOpHasTraitsForSharding(owner)) {
if (DoTypesHavePartialSameShape(value_to_visit,
owner->getResult(0))) {
next_values_to_visit.push_back(use.getOwner()->getResult(0));
continue;
}
}
if (llvm::isa<mlir::TF::CastOp, mlir::TF::XlaAllReduceOp,
mlir::TF::ReadVariableOp>(owner)) {
next_values_to_visit.push_back(use.getOwner()->getResult(0));
continue;
}
if (auto call_op = llvm::dyn_cast<mlir::CallOpInterface>(owner)) {
mlir::func::FuncOp func =
llvm::dyn_cast<mlir::func::FuncOp>(call_op.resolveCallable());
if (!func) continue;
next_values_to_visit.push_back(
func.getArgument(use.getOperandNumber()));
}
}
}
values_to_visit.swap(next_values_to_visit);
}
return std::nullopt;
}
// Tries to extract sharding configurations for all inputs by parsing
// XlaSharding/ TPUPartitionedInput op connected to the operands/arguments. If
// argument to the `cluster_func` directly feeds into another function call op,
// then recursively walk the function definition to find the connected
// XlaSharding op.
absl::Status IdentifyXlaShardingForComputationInputs(
const llvm::SmallVector<std::string>& logical_device_vec,
bool infer_from_computation, mlir::tf_device::ClusterFuncOp cluster_func,
mlir::func::FuncOp func, Builder* builder,
OptionalOpShardingVector& sharding_for_args) {
// Look up function definition from module.
Block& function_block = func.front();
sharding_for_args.reserve(function_block.getNumArguments());
// Iterate through operands of `cluster_func`.
// The computation operand can either be:
// 1) a TPUPartitionedInput Op if the input has a non-resource type;
// 2) a ReadVariableOp else.
//
// Iterate through input arguments to the entry block of
// tf_device.ClusterFunc. For input ops, look for XlaSharding ops.
// XlaSharding ops can:
// 1) Directly follow the input argument if input argument has non-resource
// types.
// 2) Follow ReadVariableOp if the input type is of resource type.
// 3) Follow IdentityOp or CastOp after above cases (1), (2).
//
// Sharding configurations are added to the tf_device.ClusterFunc as an
// attribute and the function as an argument attribute.
for (auto operand_and_arg :
llvm::zip(cluster_func.getOperands(), function_block.getArguments())) {
Value operand = std::get<0>(operand_and_arg);
BlockArgument arg = std::get<1>(operand_and_arg);
if (auto operand_sharding = GetXlaShardingFromOperand(operand)) {
sharding_for_args.push_back(operand_sharding);
continue;
}
if (infer_from_computation) {
TF_ASSIGN_OR_RETURN(auto arg_sharding,
GetXlaShardingFromArg(arg, logical_device_vec));
if (arg_sharding) {
sharding_for_args.push_back(arg_sharding.value());
continue;
}
}
sharding_for_args.push_back(std::nullopt);
}
return absl::OkStatus();
}
// Returns a TPUPartitionedOutput or TPUPartitionedInput op with XLA sharding
// connected to a `tf_device.cluster_func` result value (via AssignVariableOp/
// resource write).
mlir::Operation* GetXlaShardingFromResult(Value value) {
if (!value.hasOneUse()) return nullptr;
Operation* user = *value.getUsers().begin();
if (auto partitioned_output =
llvm::dyn_cast<mlir::TF::TPUPartitionedOutputV2Op>(user))
return NullUnlessSharded(partitioned_output);
if (auto assign_var = llvm::dyn_cast<mlir::TF::AssignVariableOp>(user))
if (auto partitioned_input =
assign_var.getResource()
.getDefiningOp<mlir::TF::TPUPartitionedInputV2Op>())
return NullUnlessSharded(partitioned_input);
return nullptr;
}
absl::Status DetermineShardingFromAlias(
mlir::func::FuncOp func, OptionalOpShardingVector& input_shardings,
OptionalOpShardingVector& output_shardings) {
for (int arg_idx = 0; arg_idx < func.getNumArguments(); ++arg_idx) {
if (auto v =
func.getArgAttrOfType<mlir::IntegerAttr>(arg_idx, kAliasingAttr)) {
if (int retval_idx = v.getInt();
retval_idx >= 0 && retval_idx < func.getNumResults()) {
auto& input_sharding = input_shardings[arg_idx];
auto& output_sharding = output_shardings[retval_idx];
if (input_sharding.has_value() && output_sharding.has_value() &&
input_sharding.value() != output_sharding.value()) {
return absl::InvalidArgumentError(absl::StrCat(
"arg#", arg_idx, " is aliased to retval#", retval_idx,
" but their sharding configurations don't match."));
} else if (input_sharding.has_value() && !output_sharding.has_value()) {
output_sharding = input_sharding;
} else if (!input_sharding.has_value() && output_sharding.has_value()) {
input_sharding = output_sharding;
}
}
}
}
return absl::OkStatus();
}
// Returns XLA sharding from XlaSharding op connected to a result value.
// XlaSharding op may be directly connected to output but it may also be
// followed by Identity or simple arithmetic ops. In case where bfloat16 type is
// used, we might see a Cast op.
//
// TODO(hongjunchoi): Add logic to parse XlaSharding op inside control flow (If,
// Case) ops and Caller argument values.
// TODO(hongjunchoi): Consider explicitly checking op patterns to detect sharded
// inputs.
absl::StatusOr<std::optional<StringRef>> GetXlaShardingFromRetval(
Value value, const llvm::SmallVector<std::string>& logical_device_vec) {
llvm::SmallPtrSet<Value, 4> visited_values;
llvm::SmallVector<Value, 4> values_to_visit;
values_to_visit.push_back(value);
while (!values_to_visit.empty()) {
Value value_to_visit = values_to_visit.pop_back_val();
if (!visited_values.insert(value_to_visit).second) {
continue;
}
Operation* def = value_to_visit.getDefiningOp();
if (!def) {
continue;
}
if (auto sharding = llvm::dyn_cast_or_null<mlir::TF::XlaShardingOp>(def))
return GetXlaShardingFromShardingOp(sharding);
if (auto sharding = def->getAttrOfType<StringAttr>(kShardingAttributeV2)) {
return sharding.strref();
}
if (auto sharding = def->getAttrOfType<StringAttr>(kShardingAttribute)) {
return sharding.strref();
}
if (auto logical_device = AssignLogicalDeviceFromTPUReplicatedCoreAttr(
def, logical_device_vec)) {
return logical_device;
}
if (UnaryOpHasTraitsForSharding(def) || BinaryOpHasTraitsForSharding(def)) {
for (auto operand : def->getOperands()) {
values_to_visit.push_back(operand);
}
continue;
}
if (auto call_op = llvm::dyn_cast_or_null<mlir::CallOpInterface>(def)) {
mlir::func::FuncOp func =
llvm::dyn_cast<mlir::func::FuncOp>(call_op.resolveCallable());
if (!func) continue;
value_to_visit = func.front().getTerminator()->getOperand(
mlir::cast<OpResult>(value_to_visit).getResultNumber());
values_to_visit.push_back(value_to_visit);
continue;
}
if (auto while_op = llvm::dyn_cast<mlir::TF::WhileRegionOp>(def)) {
if (auto op_result = mlir::cast<OpResult>(value_to_visit)) {
int result_idx = op_result.getResultNumber();
if (auto yield_op = llvm::dyn_cast<mlir::TF::YieldOp>(
while_op.getBody().front().getTerminator())) {
values_to_visit.push_back(yield_op.getOperand(result_idx));
}
}
continue;
}
}
return std::nullopt;
}
// Tries to extract sharding configurations for all outputs by parsing
// XlaSharding/ TPUPartitionedOutput op connected to the retvals/results.
absl::Status IdentifyXlaShardingForComputationOutputs(
const llvm::SmallVector<std::string>& logical_device_vec,
bool infer_from_computation, mlir::tf_device::ClusterFuncOp cluster_func,
mlir::func::FuncOp func, Builder* builder,
OptionalOpShardingVector& sharding_for_rets) {
Block& function_block = func.front();
Operation* terminator = function_block.getTerminator();
sharding_for_rets.reserve(terminator->getNumOperands());
// Iterate through results of `cluster_func`. For output ops, look for
// TPUPartitionedOutput ops.
//
// Iterate through operands of the terminator. If the preceding op is
// XlaShardingOp, then the provided sharding configuration is added to the
// tf_device.ClusterFunc as an attribute and the function as a result
// attribute.
for (auto result_and_retval :
llvm::zip(cluster_func.getResults(), terminator->getOpOperands())) {
Value result = std::get<0>(result_and_retval);
OpOperand& retval = std::get<1>(result_and_retval);
if (auto result_sharding = GetXlaShardingFromResult(result)) {
sharding_for_rets.push_back(result_sharding);
continue;
}
if (infer_from_computation) {
TF_ASSIGN_OR_RETURN(
auto retval_sharding,
GetXlaShardingFromRetval(retval.get(), logical_device_vec));
if (retval_sharding) {
sharding_for_rets.push_back(retval_sharding.value());
continue;
}
}
sharding_for_rets.push_back(std::nullopt);
}
return absl::OkStatus();
}
void SetReplicatedOrMaximalShardingIfNoShardingFound(
const llvm::SmallVector<std::string>& logical_device_vec, bool use_spmd,
OptionalOpShardingVector& shardings) {
for (auto& sharding : shardings) {
if (sharding == std::nullopt) {
// If we haven't found sharding, default to either replicated or maximal
// sharding depending on whether XLA SPMD is enabled.
if (use_spmd) {
// If XLA SPMD is enabled, host variables or non-variable per-replica
// inputs, and outputs should take on replicate sharding, so that every
// device gets the whole tensor(s) (and can slice them up later eg.
// using dynamic-slice).
sharding = kReplicateSharding;
} else {
// Otherwise, default to maximal sharding core 0.
sharding = logical_device_vec[0];
}
}
}
}
// Moves shardings from `optional_shardings` to `shardings`.
absl::Status MoveSharding(OptionalOpShardingVector& optional_shardings,
OpShardingVector& shardings) {
shardings.clear();
for (auto& sharding : optional_shardings) {
if (!sharding) {
return absl::InternalError(
"Couldn't find/assign sharding for an input/output. All shardings "
"should have been identified by this point.");
}
shardings.push_back(std::move(sharding.value()));
}
return absl::OkStatus();
}
// Determines XlaSharding for inputs and outputs. If there are aliased
// inputs/outputs for which no sharding was found directly, the corresponding
// output/input sharding is used (if it exists). If we still don't find sharding
// for some inputs/outputs, we default to replicated or maximal sharding
// depending on `use_spmd`.
absl::Status IdentifyXlaShardingForInputsAndOutputs(
const llvm::SmallVector<std::string>& logical_device_vec, bool use_spmd,
bool infer_from_computation, mlir::tf_device::ClusterFuncOp cluster_func,
mlir::func::FuncOp func, Builder* builder, OpShardingVector& input_sharding,
OpShardingVector& output_sharding) {
OptionalOpShardingVector optional_input_sharding;
OptionalOpShardingVector optional_output_sharding;
TF_RETURN_IF_ERROR(IdentifyXlaShardingForComputationInputs(
logical_device_vec, infer_from_computation, cluster_func, func, builder,
optional_input_sharding));
TF_RETURN_IF_ERROR(IdentifyXlaShardingForComputationOutputs(
logical_device_vec, infer_from_computation, cluster_func, func, builder,
optional_output_sharding));
TF_RETURN_IF_ERROR(DetermineShardingFromAlias(func, optional_input_sharding,
optional_output_sharding));
SetReplicatedOrMaximalShardingIfNoShardingFound(logical_device_vec, use_spmd,
optional_input_sharding);
SetReplicatedOrMaximalShardingIfNoShardingFound(logical_device_vec, use_spmd,
optional_output_sharding);
TF_RETURN_IF_ERROR(MoveSharding(optional_input_sharding, input_sharding));
TF_RETURN_IF_ERROR(MoveSharding(optional_output_sharding, output_sharding));
return absl::OkStatus();
}
// Extracts input/output sharding configuration of `cluster_func` by parsing
// XlaSharding ops inside the `cluster_func`.
LogicalResult IdentifyXlaShardingForTPUComputation(
Builder* builder, mlir::tf_device::ClusterFuncOp cluster_func) {
// Look up function definition from module.
mlir::func::FuncOp func =
cluster_func->getParentOfType<ModuleOp>()
.lookupSymbol<mlir::func::FuncOp>(cluster_func.getFunc());
bool use_spmd = false;
if (auto use_spmd_attr = cluster_func->getAttrOfType<BoolAttr>(kUseSpmdAttr))
use_spmd = use_spmd_attr.getValue();
auto num_cores_per_replica_attr =
cluster_func->getAttrOfType<IntegerAttr>(kNumCoresPerReplicaAttr);
if (!num_cores_per_replica_attr)
return cluster_func.emitOpError(
CreateMissingAttributeMsg(kNumCoresPerReplicaAttr));
int num_cores_per_replica = num_cores_per_replica_attr.getInt();
llvm::SmallVector<std::string> logical_device_vec(num_cores_per_replica);
for (int idx = 0; idx < num_cores_per_replica; idx++) {
logical_device_vec[idx] =
xla::sharding_builder::SingleDevice(idx).SerializeAsString();
}
OpShardingVector sharding_for_args;
OpShardingVector sharding_for_rets;
if (auto status = IdentifyXlaShardingForInputsAndOutputs(
logical_device_vec, use_spmd,
/*infer_from_computation=*/true, cluster_func, func, builder,
sharding_for_args, sharding_for_rets);
!status.ok()) {
return cluster_func.emitOpError(status.message());
};
auto has_maximal_sharding =
[](const OpShardingVariant& sharding_or_op) -> bool {
const auto sharding = GetShardingFromVariant(sharding_or_op);
return sharding && sharding->type() == xla::OpSharding::MAXIMAL;
};
// XLA SPMD only supports cases where all inputs/outputs exist on every
// partition (sharded or replicated). If any of the inputs/outputs have
// maximal sharding, then fallback to MPMD. Also fall back if any of the
// shardings aren't compatible with the rank of their tensor.
if ((use_spmd && (absl::c_any_of(sharding_for_args, has_maximal_sharding) ||
absl::c_any_of(sharding_for_rets, has_maximal_sharding))) ||
failed(VerifyShardings(func, sharding_for_args, sharding_for_rets))) {
LOG(WARNING) << "XLA SPMD only supports cases where all inputs/outputs "
"exist on every partition (sharded or replicated). If any "
"of the inputs/outputs have maximal sharding, then "
"fallback to MPMD.";
sharding_for_args.clear();
sharding_for_rets.clear();
cluster_func->setAttr(kUseSpmdAttr, builder->getBoolAttr(false));
if (auto status = IdentifyXlaShardingForInputsAndOutputs(
logical_device_vec, /*use_spmd=*/false,
/*infer_from_computation=*/false, cluster_func, func, builder,
sharding_for_args, sharding_for_rets);
!status.ok()) {
return cluster_func.emitOpError(status.message());
}
}
// Update sharding on function arguments and returns.
Block& function_block = func.front();
for (auto sharding_and_arg :
llvm::zip(sharding_for_args, function_block.getArguments())) {
BlockArgument arg = std::get<1>(sharding_and_arg);
const auto& sharding_or_op = std::get<0>(sharding_and_arg);
if (auto sharding = GetShardingStringFromVariant(sharding_or_op)) {
func.setArgAttr(arg.getArgNumber(), kShardingAttr,
builder->getStringAttr(*sharding));
}
}
Operation* terminator = function_block.getTerminator();
for (auto sharding_and_retval :
llvm::zip(sharding_for_rets, terminator->getOpOperands())) {
OpOperand& retval = std::get<1>(sharding_and_retval);
const auto& sharding_or_op = std::get<0>(sharding_and_retval);
if (auto sharding = GetShardingStringFromVariant(sharding_or_op)) {
func.setResultAttr(retval.getOperandNumber(), kShardingAttr,
builder->getStringAttr(*sharding));
}
}
// Update input/output sharding attributes on tf_device.cluster_func op.
cluster_func->setAttr(tensorflow::kInputShardingAttr,
GetStrArrayAttr(builder, sharding_for_args));
cluster_func->setAttr(tensorflow::kOutputShardingAttr,
GetStrArrayAttr(builder, sharding_for_rets));
return mlir::success();
}
void TPUShardingIdentificationPass::runOnOperation() {
Builder builder(getOperation().getContext());
auto result =
getOperation().walk([&](mlir::tf_device::ClusterFuncOp cluster_func) {
if (failed(
IdentifyXlaShardingForTPUComputation(&builder, cluster_func))) {
return WalkResult::interrupt();
}
return WalkResult::advance();
});
if (result.wasInterrupted()) return signalPassFailure();
}
} // namespace
std::unique_ptr<mlir::OperationPass<ModuleOp>>
CreateTPUShardingIdentificationPass() {
return std::make_unique<TPUShardingIdentificationPass>();
}
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/passes/tpu_validate_inputs.cc
+645
View File
@@ -0,0 +1,645 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <functional>
#include <memory>
#include <optional>
#include <queue>
#include <string>
#include <unordered_map>
#include "absl/log/log.h"
#include "absl/strings/match.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/Support/Casting.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/TypeUtilities.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/tpu_rewrite_device_util.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/tpu_validate_inputs_utils.h"
#include "xla/xla_data.pb.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {
#define GEN_PASS_DEF_TPUVALIDATEINPUTSPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc"
constexpr char kXLAShardingAttr[] = "_XlaSharding";
constexpr char kShardingAttr[] = "sharding";
using mlir::dyn_cast;
using mlir::isa;
using mlir::ModuleOp;
using mlir::Operation;
using mlir::OperationPass;
using mlir::StringAttr;
using mlir::StringRef;
using mlir::Type;
using mlir::TypeID;
using mlir::func::FuncOp;
using mlir::func::ReturnOp;
using mlir::TF::AssertOp;
using mlir::TF::ConstOp;
using mlir::TF::kCompileDeviceTypeAttr;
using mlir::TF::kTpuReplicateAttr;
using mlir::TF::OutfeedEnqueueTupleOp;
using mlir::TF::PartitionedCallOp;
using mlir::TF::StatefulPartitionedCallOp;
using mlir::TF::TPUPartitionedCallOp;
using mlir::TF::TPUPartitionedInputOp;
using mlir::TF::TPUPartitionedInputV2Op;
using mlir::TF::TPUPartitionedOutputOp;
using mlir::TF::TPUPartitionedOutputV2Op;
using mlir::TF::TPUReplicatedInputOp;
using mlir::TF::TPUReplicatedOutputOp;
using mlir::TF::TPUReplicateMetadataOp;
using mlir::TF::WhileOp;
using mlir::TF::XlaSetDynamicDimensionSizeOp;
using mlir::tf_executor::FetchOp;
using mlir::tf_executor::GraphOp;
using mlir::tf_executor::IslandOp;
using mlir::tf_executor::YieldOp;
typedef std::unordered_map<std::string, TPUReplicateMetadataOp> MetadataMap;
struct TPUValidateInputsPass
: public impl::TPUValidateInputsPassBase<TPUValidateInputsPass> {
void runOnOperation() override;
};
bool IsTpuRegularOp(Operation* op) {
static auto* ops = [] {
llvm::SmallDenseSet<mlir::TypeID, 32>* ops_set =
new llvm::SmallDenseSet<mlir::TypeID, 32>{
TypeID::get<ModuleOp>(),
TypeID::get<GraphOp>(),
TypeID::get<ReturnOp>(),
TypeID::get<FuncOp>(),
TypeID::get<YieldOp>(),
TypeID::get<IslandOp>(),
TypeID::get<TPUReplicatedInputOp>(),
TypeID::get<TPUReplicatedOutputOp>(),
TypeID::get<TPUPartitionedInputOp>(),
TypeID::get<TPUPartitionedInputV2Op>(),
TypeID::get<TPUPartitionedOutputOp>(),
TypeID::get<TPUPartitionedOutputV2Op>(),
TypeID::get<TPUReplicateMetadataOp>(),
TypeID::get<FetchOp>(),
TypeID::get<OutfeedEnqueueTupleOp>(),
};
return ops_set;
}();
auto abstractOp = op->getRegisteredInfo();
if (!abstractOp) return true;
return ops->count(abstractOp->getTypeID()) == 0;
}
bool IsIntersectionXlaNonXlaOps(Operation* op) {
static auto* ops = [] {
llvm::SmallDenseSet<mlir::TypeID, 32>* ops_set =
new llvm::SmallDenseSet<mlir::TypeID, 32>{
TypeID::get<ConstOp>(),
TypeID::get<WhileOp>(),
TypeID::get<AssertOp>(),
TypeID::get<XlaSetDynamicDimensionSizeOp>(),
};
return ops_set;
}();
auto abstractOp = op->getRegisteredInfo();
if (!abstractOp) return true;
return ops->count(abstractOp->getTypeID()) == 0;
}
bool IsPartitionedOp(Operation* op) {
static auto* ops = [] {
llvm::SmallDenseSet<mlir::TypeID, 32>* ops_set =
new llvm::SmallDenseSet<mlir::TypeID, 32>{
TypeID::get<StatefulPartitionedCallOp>(),
TypeID::get<PartitionedCallOp>(),
TypeID::get<TPUPartitionedCallOp>(),
};
return ops_set;
}();
auto abstractOp = op->getRegisteredInfo();
if (!abstractOp) return false;
return ops->count(abstractOp->getTypeID()) != 0;
}
// Gets the successors of an op wrapped in a tf_executor.island.
llvm::SmallVector<Operation*> GetSuccessors(Operation* op) {
llvm::SmallVector<Operation*> successors;
for (auto result : op->getParentOp()->getOpResults()) {
for (auto& use : result.getUses()) {
auto succ = use.getOwner();
successors.push_back(succ);
}
}
return successors;
}
// Gets the predecessors of an op wrapped in tf_executor.island.
llvm::SmallVector<Operation*> GetPredecessors(Operation* op) {
llvm::SmallVector<Operation*> predecessors;
for (auto operand : op->getOperands()) {
if (Operation* pred = operand.getDefiningOp()) {
pred->walk([&](mlir::Operation* opinexecutor) {
predecessors.push_back(opinexecutor);
});
}
}
return predecessors;
}
bool CheckTpuReplicateAttr(Operation* op, StringAttr attr,
std::function<std::string()> errormsg) {
if (!op->hasAttr(kTpuReplicateAttr)) {
op->emitOpError("TF2XLA TPU bridge input check: " + errormsg() +
"missing _tpu_replicate attr");
return false;
}
auto opattr = op->getAttr(kTpuReplicateAttr);
if (opattr != attr) {
op->emitOpError("TF2XLA TPU bridge input check: " + errormsg() +
"invalid _tpu_replicate attr.")
<< " Expected attr: " << attr << ", Actual attr: " << opattr;
return false;
}
return true;
}
bool ValidateReplicatedInput(TPUReplicatedInputOp rep, int num_replicas,
StringAttr attr) {
int arity = rep.getInputs().size();
if (rep.getIsPacked() && arity != 1) {
rep.emitOpError(
"TF2XLA TPU bridge input check: packed with number of inputs not 1.")
<< " num_replicas=" << num_replicas << " no. of inputs=" << arity;
return false;
} else if (!rep.getIsPacked() && arity != num_replicas) {
rep.emitOpError(
"TF2XLA TPU bridge input check: number of inputs inconsistent.")
<< " num_replicas=" << num_replicas << " no. of inputs=" << arity;
return false;
}
for (auto& succ : GetSuccessors(rep)) {
if (!IsTpuRegularOp(succ)) continue;
auto errormsg = [&]() -> std::string {
return rep->getName().getStringRef().str() + " op has successor op " +
succ->getName().getStringRef().str() + " with error: ";
};
if (!CheckTpuReplicateAttr(succ, attr, errormsg)) return false;
}
return true;
}
bool ValidateReplicatedOutput(TPUReplicatedOutputOp rep, int num_replicas,
StringAttr attr) {
int arity = rep.getOutputs().size();
if (arity != num_replicas) {
rep.emitOpError(
"TF2XLA TPU bridge input check: number of outputs inconsistent.")
<< " num_replicas=" << num_replicas << " no. of outputs=" << arity;
return false;
}
for (auto& pred : GetPredecessors(rep)) {
if (!IsTpuRegularOp(pred)) continue;
auto errormsg = [&]() -> std::string {
return rep->getName().getStringRef().str() + " op has predecessor op " +
pred->getName().getStringRef().str() + " with error: ";
};
if (!CheckTpuReplicateAttr(pred, attr, errormsg)) return false;
}
return true;
}
bool ValidatePartitionedInput(TPUPartitionedInputOp rep,
int num_cores_per_replica) {
int arity = rep.getInputs().size();
if (arity != num_cores_per_replica) {
rep.emitOpError(
"TF2XLA TPU bridge input check: number of inputs inconsistent.")
<< " num_cores_per_replica=" << num_cores_per_replica
<< " no. of inputs=" << arity;
return false;
}
return true;
}
bool ValidatePartitionedInputV2(TPUPartitionedInputV2Op rep,
int num_cores_per_replica) {
int arity = rep.getInputs().size();
if (rep.getIsPacked() && arity != 1) {
rep.emitOpError(
"TF2XLA TPU bridge input check: packed with number of inputs not 1.")
<< " num_cores_per_replicas=" << num_cores_per_replica
<< " no. of inputs=" << arity;
return false;
} else if (!rep.getIsPacked() && arity != num_cores_per_replica) {
rep.emitOpError(
"TF2XLA TPU bridge input check: number of inputs inconsistent.")
<< " num_cores_per_replica=" << num_cores_per_replica
<< " no. of inputs=" << arity;
return false;
}
return true;
}
template <typename T>
bool ValidatePartitionedOutput(T rep, int num_cores_per_replica) {
int arity = rep.getOutput().size();
if (arity != num_cores_per_replica) {
rep.emitOpError(
"TF2XLA TPU bridge input check: number of outputs inconsistent.")
<< " num_cores_per_replica=" << num_cores_per_replica
<< " no. of outputs=" << arity;
return false;
}
return true;
}
bool CheckReplicatedIOOp(Operation* op, TPUReplicateMetadataOp metadata,
Operation* parent) {
int num_replicas = metadata.getNumReplicas();
int num_cores_per_replica = metadata.getNumCoresPerReplica();
StringAttr tpu_replicate_attr =
metadata->getAttrOfType<StringAttr>(kTpuReplicateAttr);
if (auto repinput = dyn_cast<TPUReplicatedInputOp>(op)) {
if (!ValidateReplicatedInput(repinput, num_replicas, tpu_replicate_attr))
return false;
}
if (auto repoutput = dyn_cast<TPUReplicatedOutputOp>(op)) {
if (!ValidateReplicatedOutput(repoutput, num_replicas, tpu_replicate_attr))
return false;
}
if (auto partinput = dyn_cast<TPUPartitionedInputOp>(op)) {
if (!ValidatePartitionedInput(partinput, num_cores_per_replica))
return false;
}
if (auto partinput = dyn_cast<TPUPartitionedInputV2Op>(op)) {
if (!ValidatePartitionedInputV2(partinput, num_cores_per_replica))
return false;
}
if (auto partoutput = dyn_cast<TPUPartitionedOutputOp>(op)) {
if (!ValidatePartitionedOutput(partoutput, num_cores_per_replica))
return false;
}
if (auto partoutput = dyn_cast<TPUPartitionedOutputV2Op>(op)) {
if (!ValidatePartitionedOutput(partoutput, num_cores_per_replica))
return false;
}
return true;
}
// Checking op which is successor to a cluster op.
bool CheckClusterSuccessors(Operation* op, std::string cluster,
Operation* parent, MetadataMap& metadata_map) {
std::string cluster_succ = "";
if (op->hasAttr(kTpuReplicateAttr)) {
cluster_succ = op->getAttrOfType<StringAttr>(kTpuReplicateAttr).str();
}
if (cluster_succ.empty()) {
// TODO (b/269195256#comment16): Change to error after resolving issue
// with test. Will fix it after the upstream code is fixed.
op->emitWarning("TF2XLA TPU bridge input check: cluster op = ")
<< parent->getName() << " with cluster = " << cluster
<< " has successor as non cluster op " << op->getName();
return true;
}
if (cluster != cluster_succ) {
op->emitOpError(
"TF2XLA TPU bridge input check: mismatch clusters tpu_replicate "
"attr. Parent op ")
<< parent->getName() << " with cluster = " << cluster
<< " has successor cluster op " << op->getName()
<< " with cluster = " << cluster_succ;
return false;
}
return true;
}
// Checking op which is a predecessor to a non-cluster op.
bool CheckNonClusterSuccessors(Operation* op, Operation* parent,
MetadataMap& metadata_map) {
if (!IsTpuRegularOp(op)) {
if (isa<TPUReplicatedOutputOp>(op)) {
op->emitOpError("TF2XLA TPU bridge input check: non-cluster op = ")
<< parent->getName()
<< " has invalid successor op = " << op->getName();
return false;
} else {
return true;
}
}
return true;
}
// Checking op which is a successor to a non-cluster op.
bool CheckNonClusterPredecessors(Operation* op, Operation* parent,
MetadataMap& metadata_map) {
if (!IsTpuRegularOp(op)) {
if (isa<TPUReplicatedInputOp>(op)) {
op->emitOpError("TF2XLA TPU bridge input check: non-cluster op = ")
<< parent->getName()
<< " has invalid predecessor op = " << op->getName();
return false;
} else {
return true;
}
}
return true;
}
bool CheckOpsClusterIO(Operation* op, MetadataMap& metadata_map) {
bool is_cluster_op = false;
std::string cluster = "";
if (op->hasAttr(kTpuReplicateAttr)) {
cluster = op->getAttrOfType<StringAttr>(kTpuReplicateAttr).str();
if (cluster.empty()) {
op->emitOpError("TF2XLA TPU bridge input check: empty _tpu_replicate")
<< " attr for op = " << op->getName();
return false;
}
is_cluster_op = true;
}
bool has_cluster_metadata =
(metadata_map.find(cluster) != metadata_map.end());
for (auto pred : GetPredecessors(op)) {
if (is_cluster_op && !IsTpuRegularOp(pred) && has_cluster_metadata) {
if (!CheckReplicatedIOOp(pred, metadata_map[cluster], op)) return false;
}
if (!is_cluster_op) {
if (!CheckNonClusterPredecessors(pred, op, metadata_map)) return false;
}
}
for (auto succ : GetSuccessors(op)) {
if (is_cluster_op && !IsTpuRegularOp(succ) && has_cluster_metadata) {
if (!CheckReplicatedIOOp(succ, metadata_map[cluster], op)) return false;
}
if (is_cluster_op && IsTpuRegularOp(succ)) {
if (!CheckClusterSuccessors(succ, cluster, op, metadata_map))
return false;
}
if (!is_cluster_op) {
if (!CheckNonClusterSuccessors(succ, op, metadata_map)) return false;
}
}
return true;
}
bool TypeMustBeNonXLA(const Type& type) {
const Type elem = getElementTypeOrSelf(type);
return !mlir::isa<mlir::TF::ResourceType>(elem) &&
!tensorflow::TypeValidForXLA(type);
}
// Check if the op cannot be XLA compiled. If the op does not satisfy this
// criteria, then it is possible for the op to be XLA and non-XLA. But this
// function specifically checks if the op must be non-xla.
bool IsMustNotBeXlaOp(Operation* op) {
for (auto& input : op->getOpOperands()) {
if (TypeMustBeNonXLA(input.get().getType())) return true;
}
for (auto output_types : op->getResultTypes()) {
if (TypeMustBeNonXLA(output_types)) return true;
}
return false;
}
// Check if the op must be compiled with XLA. If the op does not satisfy this
// critiria for "must be xla" then it is still possible for this op to be xla
// and non-xla as well. But below function specifically checks for the op to be
// only XLA op.
bool IsMustBeXlaOp(Operation* op, MetadataMap metadata_map) {
// All PartitionedCall are inlined-out before XLA.
// So MustBeXLA should return false
if (IsPartitionedOp(op)) return false;
if (!op->hasAttr(kTpuReplicateAttr)) return false;
auto cluster = op->getAttrOfType<StringAttr>(kTpuReplicateAttr).str();
if (metadata_map.find(cluster) == metadata_map.end()) return false;
auto metadata = metadata_map[cluster];
if (!metadata.getAllowSoftPlacement() &&
!op->hasAttr(mlir::TF::kXlaOutsideCompilationAttr))
return true;
std::string device = "";
if (op->hasAttr(mlir::TF::kDeviceAttr))
device = op->getAttrOfType<StringAttr>(mlir::TF::kDeviceAttr).str();
else
return false;
if (absl::StrContains(device, mlir::TF::kTpuDevice)) return true;
return false;
}
bool ValidateIntersectionXlaNonXlaOps(Operation* op, MetadataMap metadata_map) {
if (isa<TPUReplicateMetadataOp>(op) || isa<TPUReplicatedInputOp>(op) ||
isa<TPUReplicatedOutputOp>(op) || isa<TPUPartitionedInputOp>(op) ||
isa<TPUPartitionedInputV2Op>(op) || isa<TPUPartitionedOutputOp>(op) ||
isa<TPUPartitionedOutputV2Op>(op))
return true;
if (IsMustBeXlaOp(op, metadata_map) && IsMustNotBeXlaOp(op)) {
// TODO(b/269195256#comment19) change the warning for Identity op to error
// when issue with input graph is resolved. Possible issue with python layer
// inserting Identity op incorrectly.
if (isa<mlir::TF::IdentityOp>(op)) {
op->emitWarning("TF/XLA TPU bridge input check: found invalid op. ")
<< op->getName() << " can't be both xla and non-xla";
return true;
}
op->emitOpError("TF/XLA TPU bridge input check: found invalid op. ")
<< "Can't be both xla and non-xla";
return false;
}
return true;
}
void GetFlattenedShardings(llvm::SmallVector<xla::OpSharding>& shardings_result,
std::string shard_string) {
xla::OpSharding sharding;
if (shard_string.empty()) return;
if (!sharding.ParseFromString(shard_string)) return;
std::queue<xla::OpSharding> shardings;
shardings.push(sharding);
while (!shardings.empty()) {
auto sharding_next = shardings.front();
shardings.pop();
if (sharding_next.type() == xla::OpSharding::TUPLE) {
for (auto& child_sharding : sharding_next.tuple_shardings()) {
shardings.push(child_sharding);
}
} else {
shardings_result.push_back(sharding_next);
}
}
}
bool IsValidShardingTupleForArity(Operation* op) {
if (!op->hasAttr(kXLAShardingAttr) && !op->hasAttr(kShardingAttr)) {
return true;
}
std::string shard_string;
if (op->hasAttr(kXLAShardingAttr)) {
shard_string =
op->getAttrOfType<StringAttr>(kXLAShardingAttr).strref().str();
} else {
shard_string = op->getAttrOfType<StringAttr>(kShardingAttr).strref().str();
}
xla::OpSharding sharding;
if (!shard_string.empty() && sharding.ParseFromString(shard_string)) {
// Only checking op with TUPLE sharding
if (sharding.type() != xla::OpSharding::TUPLE) {
return true;
}
// Each output is expected to have a corresponding sharding given by
// tuple shardings. So, the no. of outputs (arity) should be same as
// number of tuple shardings.
if (sharding.tuple_shardings().size() != op->getNumResults()) {
op->emitOpError(
"TF2XLA TPU bridge input check: invalid no. of tuple shardings ")
<< sharding.tuple_shardings().size()
<< " for arity = " << op->getNumResults() << "\n The sharding is "
<< sharding.DebugString() << "\n";
return false;
}
}
return true;
}
bool IsValidMAXIMALSharding(Operation* op, MetadataMap& metadata_map) {
if (!op->hasAttr(kTpuReplicateAttr)) return true;
if (!op->hasAttr(kXLAShardingAttr) && !op->hasAttr(kShardingAttr)) {
return true;
}
int num_cores_per_replica;
// Assuming that the op is a IsTpuRegularOp and has a cluster metadata
// for it. These checks are already performed in CheckOpsClusterIO.
// Also assuming that if there is sharding, then there must be
// cluster and the metadata corresponding to it.
auto cluster = op->getAttrOfType<StringAttr>(kTpuReplicateAttr).str();
if (cluster.empty()) {
return true;
}
if (metadata_map.find(cluster) == metadata_map.end()) {
return true;
}
num_cores_per_replica = metadata_map[cluster].getNumCoresPerReplica();
std::optional<StringRef> shard_string;
if (op->hasAttr(kXLAShardingAttr)) {
shard_string = op->getAttrOfType<StringAttr>(kXLAShardingAttr).strref();
} else {
shard_string = op->getAttrOfType<StringAttr>(kShardingAttr).strref();
}
llvm::SmallVector<xla::OpSharding> shardings;
GetFlattenedShardings(shardings, shard_string.value().str());
for (auto& sharding : shardings) {
if (sharding.type() != xla::OpSharding::MAXIMAL) {
continue;
}
if (sharding.tile_assignment_devices_size() != 1) {
op->emitOpError("TF/XLA TPU bridge input check: There must be ")
<< "exactly 1 device for MAXIMAL sharding."
<< " Number of devices assigned are "
<< sharding.tile_assignment_devices_size() << "\n";
return false;
} else {
int sharding_device = sharding.tile_assignment_devices(0);
if (sharding_device >= num_cores_per_replica || sharding_device < 0) {
op->emitOpError(
"TF2XLA TPU bridge input check: invalid sharding device ")
<< sharding_device
<< " for num_cores_per_replica = " << num_cores_per_replica
<< "\n The sharding is " << sharding.DebugString() << "\n";
return false;
}
}
}
return true;
}
bool HasSingleCoreTpu(Operation* op) {
if (auto compilation_attr =
op->getAttrOfType<mlir::StringAttr>(kCompileDeviceTypeAttr)) {
if (compilation_attr.getValue().str() == mlir::TF::kTpuDevice) {
op->emitOpError(
"TF2XLA TPU bridge input check: found a single-core TPU graph");
return true;
}
}
return false;
}
void TPUValidateInputsPass::runOnOperation() {
ModuleOp module = getOperation();
bool success = true;
int num_metadata = 0;
TPUReplicateMetadataOp metadata;
MetadataMap metadata_map;
module.walk([&](TPUReplicateMetadataOp meta) {
++num_metadata;
metadata = meta;
metadata_map[meta->getAttrOfType<StringAttr>(kTpuReplicateAttr).str()] =
meta;
});
getOperation().walk([&](mlir::Operation* op) {
if (IsPotentialUnsupportedOp(op)) {
LOG(WARNING) << "Potential unsupported op: "
<< op->getName().getStringRef().str()
<< ". TF2XLA MLIR bridge does not guarantee to support it.";
}
if (IsTpuRegularOp(op)) {
success &= CheckOpsClusterIO(op, metadata_map);
}
if (IsIntersectionXlaNonXlaOps(op)) {
success &= ValidateIntersectionXlaNonXlaOps(op, metadata_map);
}
if (IsTpuRegularOp(op)) {
success &= IsValidMAXIMALSharding(op, metadata_map);
success &= IsValidShardingTupleForArity(op);
}
success &= !HasSingleCoreTpu(op);
if (!success) {
signalPassFailure();
}
});
module.walk([&](GraphOp graph) {
if (HasV1ControlFlow(graph)) {
LOG(WARNING) << "TF2XLA MLIR bridge does not support v1 control flow."
<< " Use at your own risk.";
}
if (!success) {
signalPassFailure();
}
});
}
} // anonymous namespace
std::unique_ptr<OperationPass<ModuleOp>> CreateTPUValidateInputsPass() {
return std::make_unique<TPUValidateInputsPass>();
}
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/passes/tpu_validate_inputs.mlir
+293
View File
@@ -0,0 +1,293 @@
// 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.
// ==============================================================================
// RUN: tf-opt %s -split-input-file -verify-diagnostics -tf-tpu-validate-inputs | FileCheck %s
// CHECK-LABEL: func @num_replicas_replicated
func.func @num_replicas_replicated(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
%0:2 = tf_executor.graph {
%control = tf_executor.island() wraps "tf.TPUReplicateMetadata"() {_tpu_replicate = "cluster", device = "/device:TPU:0", num_replicas = 2, topology = "topology"} : () -> ()
%ri, %c0 = tf_executor.island wraps "tf.TPUReplicatedInput"(%arg0, %arg1) {index = 1 : i64, is_mirrored_variable = false, is_packed = false} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%out, %c1 = tf_executor.island wraps "tf.opA"(%ri) {_tpu_replicate = "cluster"} : (tensor<i32>) -> tensor<i32>
%ro:2, %c2 = tf_executor.island wraps "tf.TPUReplicatedOutput"(%out) : (tensor<i32>) -> (tensor<i32>, tensor<i32>)
tf_executor.fetch %ro#0, %ro#1 : tensor<i32>, tensor<i32>
}
return %0#0, %0#1 : tensor<i32>, tensor<i32>
}
// -----
func.func @num_replicas_replicated_input(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
%0:2 = tf_executor.graph {
%control = tf_executor.island() wraps "tf.TPUReplicateMetadata"() {_tpu_replicate = "cluster", device = "/device:TPU:0", num_replicas = 2, topology = "topology"} : () -> ()
// expected-error @+1 {{'tf.TPUReplicatedInput' op TF2XLA TPU bridge input check: number of inputs inconsistent. num_replicas=2 no. of inputs=3}}
%ri, %c0 = tf_executor.island wraps "tf.TPUReplicatedInput"(%arg0, %arg1, %arg1) {index = 1 : i64, is_mirrored_variable = false, is_packed = false} : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<i32>
%out, %c1 = tf_executor.island wraps "tf.opA"(%ri) {_tpu_replicate = "cluster"} : (tensor<i32>) -> tensor<i32>
%ro:2, %c2 = tf_executor.island wraps "tf.TPUReplicatedOutput"(%out) : (tensor<i32>) -> (tensor<i32>, tensor<i32>)
tf_executor.fetch %ro#0, %ro#1 : tensor<i32>, tensor<i32>
}
return %0#0, %0#1 : tensor<i32>, tensor<i32>
}
// -----
func.func @num_replicas_replicated_input_packed(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
%0:2 = tf_executor.graph {
%control = tf_executor.island() wraps "tf.TPUReplicateMetadata"() {_tpu_replicate = "cluster", device = "/device:TPU:0", num_replicas = 2, topology = "topology"} : () -> ()
// expected-error @+1 {{'tf.TPUReplicatedInput' op TF2XLA TPU bridge input check: packed with number of inputs not 1. num_replicas=2 no. of inputs=2}}
%ri, %c0 = tf_executor.island wraps "tf.TPUReplicatedInput"(%arg0, %arg1) {index = 1 : i64, is_mirrored_variable = false, is_packed = true} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%out, %c1 = tf_executor.island wraps "tf.opA"(%ri) {_tpu_replicate = "cluster"} : (tensor<i32>) -> tensor<i32>
%ro:2, %c2 = tf_executor.island wraps "tf.TPUReplicatedOutput"(%out) : (tensor<i32>) -> (tensor<i32>, tensor<i32>)
tf_executor.fetch %ro#0, %ro#1 : tensor<i32>, tensor<i32>
}
return %0#0, %0#1 : tensor<i32>, tensor<i32>
}
// -----
func.func @num_replicas_replicated_output(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
%0:2 = tf_executor.graph {
%control = tf_executor.island() wraps "tf.TPUReplicateMetadata"() {_tpu_replicate = "cluster", device = "/device:TPU:0", num_replicas = 2, topology = "topology"} : () -> ()
%ri, %c0 = tf_executor.island wraps "tf.TPUReplicatedInput"(%arg0, %arg1) {index = 1 : i64, is_mirrored_variable = false, is_packed = false} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%out, %c1 = tf_executor.island wraps "tf.opA"(%ri) {_tpu_replicate = "cluster"} : (tensor<i32>) -> tensor<i32>
// expected-error @+1 {{'tf.TPUReplicatedOutput' op TF2XLA TPU bridge input check: number of outputs inconsistent. num_replicas=2 no. of outputs=3}}
%ro:3, %c2 = tf_executor.island wraps "tf.TPUReplicatedOutput"(%out) : (tensor<i32>) -> (tensor<i32>, tensor<i32>, tensor<i32>)
tf_executor.fetch %ro#0, %ro#1 : tensor<i32>, tensor<i32>
}
return %0#0, %0#1 : tensor<i32>, tensor<i32>
}
// -----
func.func @num_core_per_replica_partitioned_input(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
%0:2 = tf_executor.graph {
%control = tf_executor.island() wraps "tf.TPUReplicateMetadata"() {_tpu_replicate = "cluster", device = "/device:TPU:0", num_cores_per_replica = 2 : i64, num_replicas = 1 : i64, topology = "topology"} : () -> ()
// expected-error @+1 {{'tf.TPUPartitionedInput' op TF2XLA TPU bridge input check: number of inputs inconsistent. num_cores_per_replica=2 no. of inputs=3}}
%pi, %c0 = tf_executor.island wraps "tf.TPUPartitionedInput"(%arg0, %arg1, %arg1) {index = 1 : i64} : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<i32>
%out, %c1 = tf_executor.island wraps "tf.opA"(%pi) {_tpu_replicate = "cluster"} : (tensor<i32>) -> tensor<i32>
%po:2, %c2 = tf_executor.island wraps "tf.TPUPartitionedOutput"(%out) : (tensor<i32>) -> (tensor<i32>, tensor<i32>)
tf_executor.fetch %po#0, %po#1 : tensor<i32>, tensor<i32>
}
return %0#0, %0#1 : tensor<i32>, tensor<i32>
}
// -----
func.func @num_core_per_replica_partitioned_output(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
%0:2 = tf_executor.graph {
%control = tf_executor.island() wraps "tf.TPUReplicateMetadata"() {_tpu_replicate = "cluster", device = "/device:TPU:0", num_cores_per_replica = 2 : i64, num_replicas = 1 : i64, topology = "topology"} : () -> ()
%pi, %c0 = tf_executor.island wraps "tf.TPUPartitionedInput"(%arg0, %arg1) {index = 1 : i64} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%out, %c1 = tf_executor.island wraps "tf.opA"(%pi) {_tpu_replicate = "cluster"} : (tensor<i32>) -> tensor<i32>
// expected-error @+1 {{'tf.TPUPartitionedOutput' op TF2XLA TPU bridge input check: number of outputs inconsistent. num_cores_per_replica=2 no. of outputs=3}}
%po:3, %c2 = tf_executor.island wraps "tf.TPUPartitionedOutput"(%out) : (tensor<i32>) -> (tensor<i32>, tensor<i32>, tensor<i32>)
tf_executor.fetch %po#0, %po#1 : tensor<i32>, tensor<i32>
}
return %0#0, %0#1 : tensor<i32>, tensor<i32>
}
// -----
func.func @validate_tpu_replicate_no_attr(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
%0:2 = tf_executor.graph {
%control = tf_executor.island() wraps "tf.TPUReplicateMetadata"() {_tpu_replicate = "cluster", device = "/device:TPU:0", num_replicas = 2, topology = "topology"} : () -> ()
%ri, %c0 = tf_executor.island wraps "tf.TPUReplicatedInput"(%arg0, %arg1) {index = 1 : i64, is_mirrored_variable = false, is_packed = false} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%out, %c1 = tf_executor.island wraps "tf.opA"(%ri) {_tpu_replicate="cluster"}: (tensor<i32>) -> tensor<i32>
// expected-warning @+1 {{TF2XLA TPU bridge input check: cluster op = tf.opA with cluster = cluster has successor as non cluster op tf.opB}}
%out2, %c2 = tf_executor.island wraps "tf.opB"(%out) : (tensor<i32>) -> tensor<i32>
// expected-error @+1 {{tf.TPUReplicatedOutput' op TF2XLA TPU bridge input check: non-cluster op = tf.opB has invalid successor op = tf.TPUReplicatedOutput}}
%ro:2, %c4 = tf_executor.island wraps "tf.TPUReplicatedOutput"(%out2) : (tensor<i32>) -> (tensor<i32>, tensor<i32>)
tf_executor.fetch %ro#0, %ro#1 : tensor<i32>, tensor<i32>
}
return %0#0, %0#1 : tensor<i32>, tensor<i32>
}
// -----
func.func @validate_tpu_replicate_wrong_attr(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
%0:2 = tf_executor.graph {
%control = tf_executor.island() wraps "tf.TPUReplicateMetadata"() {_tpu_replicate = "cluster", device = "/device:TPU:0", num_replicas = 2, topology = "topology"} : () -> ()
%ri, %c0 = tf_executor.island wraps "tf.TPUReplicatedInput"(%arg0, %arg1) {index = 1 : i64, is_mirrored_variable = false, is_packed = false} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%out, %c1 = tf_executor.island wraps "tf.opA"(%ri) {_tpu_replicate = "cluster_wrong"}: (tensor<i32>) -> tensor<i32>
// expected-error @+1 {{'tf.opB' op TF2XLA TPU bridge input check: mismatch clusters tpu_replicate attr. Parent op tf.opA with cluster = cluster_wrong has successor cluster op tf.opB with cluster = cluster}}
%out2, %c2 = tf_executor.island wraps "tf.opB"(%out) {_tpu_replicate = "cluster"}: (tensor<i32>) -> tensor<i32>
%ro:2, %c3 = tf_executor.island wraps "tf.TPUReplicatedOutput"(%out2) : (tensor<i32>) -> (tensor<i32>, tensor<i32>)
tf_executor.fetch %ro#0, %ro#1 : tensor<i32>, tensor<i32>
}
return %0#0, %0#1 : tensor<i32>, tensor<i32>
}
// -----
func.func @valid_xla_nonxla(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
%0:2 = tf_executor.graph {
%control = tf_executor.island wraps "tf.TPUReplicateMetadata"() {_tpu_replicate = "cluster", device = "/device:TPU:0", num_replicas = 2, topology = "topology"} : () -> ()
%ri, %c0 = tf_executor.island wraps "tf.TPUReplicatedInput"(%arg0, %arg1) {index = 1 : i64, is_mirrored_variable = false, is_packed = false} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%out, %c1 = tf_executor.island wraps "tf.opA"(%ri) {_tpu_replicate = "cluster", device = "TPU"} : (tensor<i32>) -> tensor<i32>
%ro:2, %c2 = tf_executor.island wraps "tf.TPUReplicatedOutput"(%out) : (tensor<i32>) -> (tensor<i32>, tensor<i32>)
tf_executor.fetch %ro#0, %ro#1 : tensor<i32>, tensor<i32>
}
return %0#0, %0#1 : tensor<i32>, tensor<i32>
}
// -----
func.func @valid_xla_nonxla_warning(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<i32>) -> (tensor<*x!tf_type.string>, tensor<*x!tf_type.string>) {
%0:2 = tf_executor.graph {
%control = tf_executor.island wraps "tf.TPUReplicateMetadata"() {_tpu_replicate = "cluster", device = "/device:TPU:0", num_replicas = 2, topology = "topology"} : () -> ()
%ri, %c0 = tf_executor.island wraps "tf.TPUReplicatedInput"(%arg0, %arg1) {index = 1 : i64, is_mirrored_variable = false, is_packed = false} : (tensor<i32>, tensor<i32>) -> tensor<*x!tf_type.string>
// expected-warning @+1 {{TF/XLA TPU bridge input check: found invalid op. tf.Identity can't be both xla and non-xla}}
%out, %c1 = tf_executor.island(%c0) wraps "tf.Identity"(%ri) {_tpu_replicate = "cluster", device = ""} : (tensor<*x!tf_type.string>) -> tensor<*x!tf_type.string>
%ro:2, %c2 = tf_executor.island wraps "tf.TPUReplicatedOutput"(%out) : (tensor<*x!tf_type.string>) -> (tensor<*x!tf_type.string>, tensor<*x!tf_type.string>)
tf_executor.fetch %ro#0, %ro#1 : tensor<*x!tf_type.string>, tensor<*x!tf_type.string>
}
return %0#0, %0#1 : tensor<*x!tf_type.string>, tensor<*x!tf_type.string>
}
// -----
// Serialized string:
// "\08\01\1A\01\01\22\01\00"
// Proto debug string:
// type: MAXIMAL
// tile_assignment_dimensions: 1
// tile_assignment_devices: 0
func.func @valid_MAXIMAL_sharding_device(%arg0: tensor<i32>) -> tensor<i32> {
%0 = tf_executor.graph {
%control = tf_executor.island wraps "tf.TPUReplicateMetadata"() {_tpu_replicate = "cluster", device = "/device:TPU:0", num_cores_per_replica = 2 : i64, num_replicas = 1 : i64, topology = "topology"} : () -> ()
%0, %c = tf_executor.island wraps "tf.Identity"(%arg0) {_tpu_replicate = "cluster", _XlaSharding = "\08\01\1A\01\01\22\01\00"} : (tensor<i32>) -> tensor<i32>
tf_executor.fetch %0 : tensor<i32>
}
return %0 : tensor<i32>
}
// -----
// Serialized string:
// "\08\01\1A\01\01\22\01\02"
// Proto debug string:
// type: MAXIMAL
// tile_assignment_dimensions: 1
// tile_assignment_devices: 2
func.func @invalid_MAXIMAL_sharding_device(%arg0: tensor<i32>) -> tensor<i32> {
%0 = tf_executor.graph {
%control = tf_executor.island wraps "tf.TPUReplicateMetadata"() {_tpu_replicate = "cluster", device = "/device:TPU:0", num_cores_per_replica = 2 : i64, num_replicas = 1 : i64, topology = "topology"} : () -> ()
// expected-error @+1 {{'tf.Identity' op TF2XLA TPU bridge input check: invalid sharding device 2 for num_cores_per_replica = 2}}
%0, %c = tf_executor.island wraps "tf.Identity"(%arg0) {_tpu_replicate = "cluster", _XlaSharding = "\08\01\1A\01\01\22\01\02"} : (tensor<i32>) -> tensor<i32>
tf_executor.fetch %0 : tensor<i32>
}
return %0 : tensor<i32>
}
// -----
// Serialized string:
// "\08\02\2a\08\08\01\1a\01\01\22\01\00\2a\08\08\01\1a\01\01\22\01\01"
// Proto debug string:
// type: TUPLE
// tuple_shardings {
// type: MAXIMAL
// tile_assignment_dimensions: 1
// tile_assignment_devices: 0
// }
// tuple_shardings {
// type: MAXIMAL
// tile_assignment_dimensions: 1
// tile_assignment_devices: 1
// }
func.func @invalid_TUPLE_sharding_arity(%arg0: tensor<i32>) -> tensor<i32> {
%0 = tf_executor.graph {
%control = tf_executor.island wraps "tf.TPUReplicateMetadata"() {_tpu_replicate = "cluster", device = "/device:TPU:0", num_cores_per_replica = 2 : i64, num_replicas = 1 : i64, topology = "topology"} : () -> ()
// expected-error @+1 {{'tf.Identity' op TF2XLA TPU bridge input check: invalid no. of tuple shardings 2 for arity = 1}}
%0, %c = tf_executor.island wraps "tf.Identity"(%arg0) {_tpu_replicate = "cluster", _XlaSharding = "\08\02\2a\08\08\01\1a\01\01\22\01\00\2a\08\08\01\1a\01\01\22\01\01"} : (tensor<i32>) -> tensor<i32>
tf_executor.fetch %0 : tensor<i32>
}
return %0 : tensor<i32>
}
// -----
// Serialized string:
// "\08\02\2a\08\08\01\1a\01\01\22\01\00\2a\08\08\01\1a\01\01\22\01\01"
// Proto debug string:
// type: TUPLE
// tuple_shardings {
// type: MAXIMAL
// tile_assignment_dimensions: 1
// tile_assignment_devices: 0
// }
// tuple_shardings {
// type: MAXIMAL
// tile_assignment_dimensions: 1
// tile_assignment_devices: 1
// }
func.func @outfeed_enqueue_tuple_sharding_exception(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
%0 = tf_executor.graph {
%control = tf_executor.island wraps "tf.TPUReplicateMetadata"() {_tpu_replicate = "cluster", device = "/device:TPU:0", num_cores_per_replica = 2 : i64, num_replicas = 1 : i64, topology = "topology"} : () -> ()
%0, %c0 = tf_executor.island wraps "tf.AddV2"(%arg0, %arg1) {_tpu_replicate = "cluster"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%c1 = tf_executor.island wraps "tf.OutfeedEnqueueTuple"(%arg0, %arg1) {_tpu_replicate = "cluster", _XlaSharding = "\08\02\2a\08\08\01\1a\01\01\22\01\00\2a\08\08\01\1a\01\01\22\01\01"} : (tensor<i32>, tensor<i32>) -> ()
tf_executor.fetch %0 : tensor<i32>
}
return %0 : tensor<i32>
}
// -----
func.func @single_core_tpu(%arg0: tensor<i32>) -> () {
tf_executor.graph {
// expected-error @+1 {{found a single-core TPU graph}}
tf_executor.island wraps "tf.Identity"(%arg0) {_xla_compile_device_type = "TPU"} : (tensor<i32>) -> tensor<i32>
tf_executor.fetch
}
return
}
// -----
// CHECK-LABEL: func @num_replicas_1
func.func @num_replicas_1(%arg0: tensor<i32>) -> (tensor<i32>) {
%0 = tf_executor.graph {
%control = tf_executor.island() wraps "tf.TPUReplicateMetadata"() {_tpu_replicate = "cluster", device = "/device:TPU:0", num_replicas = 1, num_cores_per_replica = 1, topology = "topology"} : () -> ()
%ri, %c0 = tf_executor.island wraps "tf.TPUReplicatedInput"(%arg0) {index = 1 : i64, is_mirrored_variable = false, is_packed = false} : (tensor<i32>) -> tensor<i32>
%out, %c1 = tf_executor.island wraps "tf.opA"(%ri) {_tpu_replicate = "cluster"} : (tensor<i32>) -> tensor<i32>
%ro, %c2 = tf_executor.island wraps "tf.TPUReplicatedOutput"(%out) : (tensor<i32>) -> tensor<i32>
tf_executor.fetch %ro : tensor<i32>
}
return %0 : tensor<i32>
}
// -----
func.func @contians_InfeedDequeueTuple(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
%0:2 = tf_executor.graph {
%control = tf_executor.island() wraps "tf.TPUReplicateMetadata"() {_tpu_replicate = "cluster", device = "/device:TPU:0", num_replicas = 2, topology = "topology"} : () -> ()
%ri, %c0 = tf_executor.island wraps "tf.TPUReplicatedInput"(%arg0, %arg1) {index = 1 : i64, is_mirrored_variable = false, is_packed = false} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%out, %c1 = tf_executor.island wraps "tf.opA"(%ri) {_tpu_replicate = "cluster"} : (tensor<i32>) -> tensor<i32>
// expected-warning @+1 {{TPU_REPLICATED_CORE:0 device is not supported for op = tf.InfeedDequeueTuple in TF2XLA MLIR Bridge}}
%infeed_output:3, %c2 = tf_executor.island wraps "tf.InfeedDequeueTuple"() {device = "/device:TPU_REPLICATED_CORE:0"} : () -> (tensor<3xi32>, tensor<4x?xf32>, tensor<*xi16>)
%ro:2, %c3 = tf_executor.island wraps "tf.TPUReplicatedOutput"(%out) : (tensor<i32>) -> (tensor<i32>, tensor<i32>)
tf_executor.fetch %ro#0, %ro#1 : tensor<i32>, tensor<i32>
}
return %0#0, %0#1 : tensor<i32>, tensor<i32>
}
// -----
func.func @graph_contains_v1_control_flow() {
tf_executor.graph {
// expected-warning @+1 {{ is v1 control flow op which is not supported in TF2XLA MLIR Bridge.}}
%control = tf_executor.ControlTrigger {}
tf_executor.fetch
}
func.return
}
tensorflow/compiler/mlir/tf2xla/internal/passes/tpu_validate_inputs_utils.cc
+66
View File
@@ -0,0 +1,66 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/tpu_validate_inputs_utils.h"
#include <string>
#include "absl/strings/match.h"
#include "llvm/ADT/DenseSet.h"
#include "mlir/Support/TypeID.h" // from @llvm-project
namespace tensorflow {
namespace tf2xla {
namespace internal {
bool IsPotentialUnsupportedOp(Operation* op) {
static auto* ops = [] {
llvm::SmallDenseSet<mlir::TypeID, 32>* ops_set =
new llvm::SmallDenseSet<mlir::TypeID, 32>{
TypeID::get<InfeedDequeueTupleOp>(),
};
return ops_set;
}();
auto abstractOp = op->getRegisteredInfo();
if (!abstractOp) return false;
bool is_in_ops = ops->count(abstractOp->getTypeID()) != 0;
if (!is_in_ops) return false;
std::string device = "";
if (!op->hasAttr(kDeviceAttr)) return false;
device = op->getAttrOfType<StringAttr>(kDeviceAttr).str();
if (!absl::StrContains(device, kTpuReplicatedCoreZeroAttr)) return false;
op->emitWarning("TPU_REPLICATED_CORE:0 device is not supported for op = ")
<< op->getName() << " in TF2XLA MLIR Bridge";
return true;
}
bool HasV1ControlFlow(GraphOp graph) {
for (Operation& op : graph.GetBody().without_terminator()) {
auto island_op = llvm::dyn_cast<mlir::tf_executor::IslandOp>(op);
if (!island_op) {
op.emitWarning() << " is v1 control flow op which is not supported in "
"TF2XLA MLIR Bridge.";
return true;
}
}
return false;
}
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/passes/tpu_validate_inputs_utils.h
+49
View File
@@ -0,0 +1,49 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_PASSES_TPU_VALIDATE_INPUTS_UTILS_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_PASSES_TPU_VALIDATE_INPUTS_UTILS_H_
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
constexpr char kTpuReplicatedCoreZeroAttr[] = "TPU_REPLICATED_CORE:0";
using mlir::ModuleOp;
using mlir::Operation;
using mlir::StringAttr;
using mlir::TypeID;
using mlir::TF::InfeedDequeueTupleOp;
using mlir::TF::kDeviceAttr;
using mlir::tf_executor::GraphOp;
bool IsPotentialUnsupportedOp(Operation* op);
bool HasV1ControlFlow(GraphOp graph);
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_PASSES_TPU_VALIDATE_INPUTS_UTILS_H_
tensorflow/compiler/mlir/tf2xla/internal/passes/tpu_validate_inputs_utils_test.cc
+94
View File
@@ -0,0 +1,94 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/tpu_validate_inputs_utils.h"
#include <gtest/gtest.h>
#include "llvm/ADT/SmallVector.h"
#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/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/test_utils.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {
using mlir::hlo::test::GetMlirModuleFromString;
TEST(IsPotentialUnsupportedOp, ClusterOpReturnsFalse) {
mlir::MLIRContext context;
context.loadDialect<mlir::tf_device::TensorFlowDeviceDialect>();
mlir::OwningOpRef<mlir::ModuleOp> module_ref =
mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
mlir::OpBuilder builder(module_ref->getBodyRegion());
llvm::SmallVector<mlir::Type, 8> result_types;
auto cluster = mlir::tf_device::ClusterOp::create(
builder, mlir::UnknownLoc::get(&context), result_types);
cluster->dump();
EXPECT_FALSE(IsPotentialUnsupportedOp(cluster));
}
TEST(IsPotentialUnsupportedOp, InfeedDequeueTupleOpReturnsTrue) {
mlir::MLIRContext context;
context.loadDialect<mlir::TF::TensorFlowDialect>();
mlir::OwningOpRef<mlir::ModuleOp> module_ref =
mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
mlir::OpBuilder builder(module_ref->getBodyRegion());
llvm::SmallVector<mlir::Type, 8> result_types;
mlir::StringAttr _XlaSharding = mlir::StringAttr::get(&context, "");
mlir::ArrayAttr layouts = mlir::ArrayAttr::get(&context, {});
auto infeed_dequeue_tuple =
InfeedDequeueTupleOp::create(builder, mlir::UnknownLoc::get(&context),
result_types, _XlaSharding, layouts);
infeed_dequeue_tuple->setAttr(
kDeviceAttr, mlir::StringAttr::get(&context, kTpuReplicatedCoreZeroAttr));
EXPECT_TRUE(IsPotentialUnsupportedOp(infeed_dequeue_tuple));
}
TEST(HasV1ControlFlow, ReturnsTrue) {
static constexpr char kMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @graph_contains_v1_control_flow() {
tf_executor.graph {
%control = tf_executor.ControlTrigger {}
tf_executor.fetch
}
func.return
}
})";
mlir::MLIRContext context;
context.loadDialect<mlir::tf_executor::TensorFlowExecutorDialect>();
auto module = GetMlirModuleFromString(kMlirModuleStr, &context);
module->get().walk(
[&](GraphOp graph) { EXPECT_TRUE(HasV1ControlFlow(graph)); });
}
} // namespace
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/passes/tpu_validate_session_inputs.cc
+75
View File
@@ -0,0 +1,75 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <memory>
#include "absl/log/log.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/tpu_validate_inputs_utils.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {
#define GEN_PASS_DEF_TPUVALIDATESESSIONINPUTSPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc"
using mlir::ModuleOp;
struct TPUValidateSessionInputsPass
: public impl::TPUValidateSessionInputsPassBase<
TPUValidateSessionInputsPass> {
void runOnOperation() override;
};
void TPUValidateSessionInputsPass::runOnOperation() {
ModuleOp module = getOperation();
bool success = true;
module.walk([&](mlir::Operation* op) {
if (IsPotentialUnsupportedOp(op)) {
LOG(WARNING) << "Potential unsupported op: "
<< op->getName().getStringRef().str()
<< ". TF2XLA MLIR bridge does not guarantee to support it.";
}
if (!success) {
signalPassFailure();
}
});
module.walk([&](GraphOp graph) {
if (HasV1ControlFlow(graph)) {
LOG(WARNING) << "TF2XLA MLIR bridge does not support v1 control flow."
<< " Use at your own risk.";
}
if (!success) {
signalPassFailure();
}
});
}
} // anonymous namespace
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateTPUValidateSessionInputsPass() {
return std::make_unique<TPUValidateSessionInputsPass>();
}
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/passes/tpu_validate_session_inputs.mlir
+51
View File
@@ -0,0 +1,51 @@
// 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.
// ==============================================================================
// RUN: tf-opt %s -split-input-file -verify-diagnostics -tf-tpu-validate-session-inputs | FileCheck %s
// CHECK-LABEL: func @does_not_contian_InfeedDequeueTuple
func.func @does_not_contian_InfeedDequeueTuple(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
%0:2 = tf_executor.graph {
%control = tf_executor.island() wraps "tf.TPUReplicateMetadata"() {_tpu_replicate = "cluster", device = "/device:TPU:0", num_replicas = 2, topology = "topology"} : () -> ()
%ri, %c0 = tf_executor.island wraps "tf.TPUReplicatedInput"(%arg0, %arg1) {index = 1 : i64, is_mirrored_variable = false, is_packed = false} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%out, %c1 = tf_executor.island wraps "tf.opA"(%ri) {_tpu_replicate = "cluster"} : (tensor<i32>) -> tensor<i32>
%ro:2, %c2 = tf_executor.island wraps "tf.TPUReplicatedOutput"(%out) : (tensor<i32>) -> (tensor<i32>, tensor<i32>)
tf_executor.fetch %ro#0, %ro#1 : tensor<i32>, tensor<i32>
}
return %0#0, %0#1 : tensor<i32>, tensor<i32>
}
// -----
func.func @contians_InfeedDequeueTuple(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
%0:2 = tf_executor.graph {
%control = tf_executor.island() wraps "tf.TPUReplicateMetadata"() {_tpu_replicate = "cluster", device = "/device:TPU:0", num_replicas = 2, topology = "topology"} : () -> ()
%ri, %c0 = tf_executor.island wraps "tf.TPUReplicatedInput"(%arg0, %arg1) {index = 1 : i64, is_mirrored_variable = false, is_packed = false} : (tensor<i32>, tensor<i32>) -> tensor<i32>
%out, %c1 = tf_executor.island wraps "tf.opA"(%ri) {_tpu_replicate = "cluster"} : (tensor<i32>) -> tensor<i32>
// expected-warning @+1 {{TPU_REPLICATED_CORE:0 device is not supported for op = tf.InfeedDequeueTuple in TF2XLA MLIR Bridge}}
%infeed_output:3, %c2 = tf_executor.island wraps "tf.InfeedDequeueTuple"() {device = "/device:TPU_REPLICATED_CORE:0"} : () -> (tensor<3xi32>, tensor<4x?xf32>, tensor<*xi16>)
%ro:2, %c3 = tf_executor.island wraps "tf.TPUReplicatedOutput"(%out) : (tensor<i32>) -> (tensor<i32>, tensor<i32>)
tf_executor.fetch %ro#0, %ro#1 : tensor<i32>, tensor<i32>
}
return %0#0, %0#1 : tensor<i32>, tensor<i32>
}
// -----
func.func @graph_contains_v1_control_flow() {
tf_executor.graph {
// expected-warning @+1 {{ is v1 control flow op which is not supported in TF2XLA MLIR Bridge.}}
%control = tf_executor.ControlTrigger {}
tf_executor.fetch
}
func.return
}
tensorflow/compiler/mlir/tf2xla/internal/passes/verify_clustering_pass.cc
+80
View File
@@ -0,0 +1,80 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <memory>
#include <string>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/utils/dialect_detection_utils.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {
#define GEN_PASS_DEF_VERIFYCLUSTERINGPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc"
using mlir::Operation;
using mlir::WalkResult;
class VerifyClusteringPass
: public impl::VerifyClusteringPassBase<VerifyClusteringPass> {
public:
void runOnOperation() override;
};
void VerifyClusteringPass::runOnOperation() {
Operation* func_op = getOperation();
auto walk_result = func_op->walk([&](Operation* op) {
if (!tensorflow::tf2xla::internal::IsInBridgeAcceptableDialects(op)) {
std::string error = "op is in dialect " +
op->getDialect()->getNamespace().str() +
" not in tf functional dialect";
op->emitError() << error;
return WalkResult::interrupt();
}
if (op->hasAttr(mlir::TF::kXlaOutsideCompilationAttr)) {
std::string error =
"op has outside compilation attribute _xla_outside_compilation which "
"is not allowed after clustering";
op->emitError() << error;
return mlir::WalkResult::interrupt();
}
return WalkResult::advance();
});
if (walk_result.wasInterrupted()) {
signalPassFailure();
}
}
} // namespace
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateVerifyClusteringPass() {
return std::make_unique<VerifyClusteringPass>();
}
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/passes/verify_clustering_pass_test.cc
+87
View File
@@ -0,0 +1,87 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <memory>
#include <gtest/gtest.h>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/test_utils.h"
#include "xla/tsl/platform/statusor.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {
using mlir::hlo::test::GetMlirModuleFromString;
class VerifyClusteringPassTest : public testing::Test {
protected:
void CreateModule(const char* module_string) {
TF_ASSERT_OK_AND_ASSIGN(module_,
GetMlirModuleFromString(module_string, &context_));
pm_ = std::make_unique<mlir::PassManager>(&context_);
pm_->addNestedPass<mlir::func::FuncOp>(CreateVerifyClusteringPass());
}
mlir::LogicalResult Run() { return pm_->run(module_.get()); }
private:
mlir::MLIRContext context_;
mlir::OwningOpRef<mlir::ModuleOp> module_;
std::unique_ptr<mlir::PassManager> pm_;
};
TEST_F(VerifyClusteringPassTest, OnlyTfFunctionalPasses) {
static constexpr char kMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> tensor<1xi32> {
%0 = "tf.Const"() {value = dense<1000> : tensor<1xi32>} : () -> tensor<1xi32>
return %0 : tensor<1xi32>
}
})";
CreateModule(kMlirModuleStr);
auto result = Run();
EXPECT_TRUE(result.succeeded());
}
TEST_F(VerifyClusteringPassTest, NotTfFunctionalFails) {
static constexpr char kMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> tensor<3x32x32x3xf32> {
%0 = mhlo.constant dense<2.550000e+02> : tensor<3x32x32x3xf32>
return %0 : tensor<3x32x32x3xf32>
}
})";
CreateModule(kMlirModuleStr);
auto result = Run();
EXPECT_TRUE(result.failed());
}
} // namespace
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/passes/verify_clustering_pass_test.mlir
+40
View File
@@ -0,0 +1,40 @@
// 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.
// ==============================================================================
// RUN: tf-opt -verify-clustering-pass -split-input-file -verify-diagnostics %s | FileCheck %s
// Tests the VerifyClusteringPass Pass, ensures that an error is thrown when validation fails.
func.func @testNotTfDialect(%arg0: tensor<1x32x10x32xi32>, %arg1: tensor<32xi32>) -> tensor<1x32x10x32xi32> {
// expected-error@below {{op is in dialect chlo not in tf functional dialect}}
%0 = "chlo.broadcast_add"(%arg0, %arg1) {broadcast_dimensions = array<i64: 3>} : (tensor<1x32x10x32xi32>, tensor<32xi32>) -> tensor<1x32x10x32xi32>
func.return %0 : tensor<1x32x10x32xi32>
}
// -----
// CHECK-LABEL: func @testTFDialect
func.func @testTFDialect(%arg0: tensor<4x?x!tf_type.stringref>) -> tensor<4x2x!tf_type.string> {
%0 = "tf.Identity"(%arg0) : (tensor<4x?x!tf_type.stringref>) -> tensor<4x2x!tf_type.string>
func.return %0 : tensor<4x2x!tf_type.string>
}
// -----
func.func @testTFDialect(%arg0: tensor<4x?x!tf_type.stringref>) -> tensor<4x2x!tf_type.string> {
// expected-error@below {{op has outside compilation attribute _xla_outside_compilation which is not allowed after clustering}}
%0 = "tf.Identity"(%arg0) {_xla_outside_compilation = "cluster1"}: (tensor<4x?x!tf_type.stringref>) -> tensor<4x2x!tf_type.string>
func.return %0 : tensor<4x2x!tf_type.string>
}
tensorflow/compiler/mlir/tf2xla/internal/passes/verify_input_dialect_to_executor_pass.cc
+84
View File
@@ -0,0 +1,84 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <memory>
#include <string>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tf2xla/internal/utils/dialect_detection_utils.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {
using mlir::Operation;
using mlir::OperationPass;
using mlir::WalkResult;
using mlir::func::FuncOp;
#define GEN_PASS_DEF_VERIFYINPUTDIALECTTOEXECUTORPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/mlir_to_graph_passes.h.inc"
class VerifyInputDialectToExecutorPass
: public impl::VerifyInputDialectToExecutorPassBase<
VerifyInputDialectToExecutorPass> {
public:
void runOnOperation() override;
};
bool IsTfDeviceClusterFuncOp(Operation* op) {
std::string kClusterFuncOpName = "tf_device.cluster_func";
return op->getName().getStringRef().str() == kClusterFuncOpName;
}
void VerifyInputDialectToExecutorPass::runOnOperation() {
Operation* func_op = getOperation();
auto walk_result = func_op->walk([&](Operation* op) {
if (!tensorflow::tf2xla::internal::IsInBridgeAcceptableDialects(op)) {
std::string error = "op is in dialect " +
op->getDialect()->getNamespace().str() +
" which is not an accepted dialect";
op->emitError() << error;
return WalkResult::interrupt();
}
if (IsTfDeviceClusterFuncOp(op)) {
std::string error =
"failed TF functional to executor validation, op "
"tf_device.cluster_func is not allowed";
op->emitError() << error;
return WalkResult::interrupt();
}
return WalkResult::advance();
});
if (walk_result.wasInterrupted()) {
signalPassFailure();
}
}
} // namespace
std::unique_ptr<OperationPass<FuncOp>>
CreateVerifyInputDialectToExecutorPass() {
return std::make_unique<VerifyInputDialectToExecutorPass>();
}
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/passes/verify_input_dialect_to_executor_pass_test.mlir
+52
View File
@@ -0,0 +1,52 @@
// 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.
// ==============================================================================
// RUN: tf-opt -verify-input-dialect-to-executor-pass -split-input-file -verify-diagnostics %s | FileCheck %s
// Tests the VerifyClusteringPass Pass, ensures that an error is thrown when validation fails.
// -----
// CHECK-LABEL: func @testNoClusterFuncOpPasses
func.func @testNoClusterFuncOpPasses(%arg0: tensor<4x?x!tf_type.stringref>) -> tensor<4x2x!tf_type.string> {
%0 = "tf.Identity"(%arg0) : (tensor<4x?x!tf_type.stringref>) -> tensor<4x2x!tf_type.string>
func.return %0 : tensor<4x2x!tf_type.string>
}
// -----
func.func @_func(%arg0: tensor<i32>) -> tensor<i32> {
func.return %arg0 : tensor<i32>
}
func.func @testClusterFuncOpFails(%arg0: tensor<i32>) -> tensor<i32> {
// expected-error@below {{failed TF functional to executor validation, op tf_device.cluster_func is not allowed}}
%cluster = "tf_device.cluster_func"(%arg0) {func = @_func} : (tensor<i32>) -> tensor<i32>
func.return %cluster : tensor<i32>
}
// -----
// CHECK-LABEL: func @testTFDialect
func.func @testTFDialect(%arg0: tensor<4x?x!tf_type.stringref>) -> tensor<4x2x!tf_type.string> {
%0 = "tf.Identity"(%arg0) : (tensor<4x?x!tf_type.stringref>) -> tensor<4x2x!tf_type.string>
func.return %0 : tensor<4x2x!tf_type.string>
}
// -----
func.func @testNotTfDialect(%arg0: tensor<1x32x10x32xi32>, %arg1: tensor<32xi32>) -> tensor<1x32x10x32xi32> {
// expected-error@below {{op is in dialect chlo which is not an accepted dialect}}
%0 = "chlo.broadcast_add"(%arg0, %arg1) {broadcast_dimensions = array<i64: 3>} : (tensor<1x32x10x32xi32>, tensor<32xi32>) -> tensor<1x32x10x32xi32>
func.return %0 : tensor<1x32x10x32xi32>
}
tensorflow/compiler/mlir/tf2xla/internal/passes/xla_broadcast.cc
+360
View File
@@ -0,0 +1,360 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <cstdint>
#include <memory>
#include <string>
#include <unordered_map>
#include "absl/log/log.h"
#include "absl/status/statusor.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/FormatVariadic.h"
#include "llvm/Support/raw_ostream.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/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/TypeUtilities.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/IR/ValueRange.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_structs.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/device_util.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/tpu_rewrite_device_util.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/xla_rewrite_util.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/xla_sharding_util.h"
#include "tensorflow/core/ir/types/dialect.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {
using llvm::dyn_cast;
using mlir::Block;
using mlir::BlockArgument;
using mlir::DenseIntElementsAttr;
using mlir::failure;
using mlir::Location;
using mlir::LogicalResult;
using mlir::OpBuilder;
using mlir::Operation;
using mlir::OperationPass;
using mlir::OpOperand;
using mlir::RankedTensorType;
using mlir::StringAttr;
using mlir::success;
using mlir::Type;
using mlir::Value;
using mlir::ValueRange;
using mlir::WalkResult;
using mlir::func::FuncOp;
using mlir::TF::ConstOp;
using mlir::TF::IdentityOp;
using mlir::TF::ShapeAttr;
using mlir::TF::TPUDummyInputOp;
using mlir::TF::XlaAllReduceOp;
using mlir::tf_device::ClusterOp;
using mlir::tf_device::LaunchOp;
using mlir::tf_device::ParallelExecuteOp;
using mlir::tf_device::ReplicateOp;
#define GEN_PASS_DEF_XLABROADCASTPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc"
struct XlaBroadcast : public impl::XlaBroadcastPassBase<XlaBroadcast> {
void runOnOperation() override;
};
// Returns true iff the broadcast val can be substituted with an XlaAllReduce.
// Sets `zero_type` and `shape` as params for the dummy zeros that will be
// created.
bool GetDummyParams(OpBuilder& builder, Value val_bcast, Type& zero_type,
ShapeAttr& shape) {
Type type = val_bcast.getType();
Type elem_type = getElementTypeOrSelf(type);
// Xla's all_reduce legalizer bitcasts to 32 bits, so only
// element types size <= 4 bytes are supported.
if (elem_type.isBF16() || elem_type.isF16() || elem_type.isTF32() ||
elem_type.isF32()) {
zero_type = type;
} else {
return false;
}
if (auto ranked_type = dyn_cast<RankedTensorType>(type)) {
llvm::ArrayRef<int64_t> type_shape = ranked_type.getShape();
for (int64_t i : type_shape) {
if (i < 0) return false;
}
shape = ShapeAttr::get(builder.getContext(), type_shape);
} else {
return false;
}
return true;
}
// Create a dummy zero to be fed locally from the host to the TPUExecute.
Value CreateZeroInput(Location loc, OpBuilder& builder, Type zero_type,
ShapeAttr shape_attr) {
TPUDummyInputOp tpu_dummy_input =
TPUDummyInputOp::create(builder, loc, zero_type, shape_attr);
tpu_dummy_input->setAttr(kICIWeightDistributionMlirBridgeMarker,
builder.getBoolAttr(true));
return tpu_dummy_input;
}
// Add parallel collection of inputs to the replicated op.
LogicalResult AppendReplicatedInput(OpBuilder& builder, ReplicateOp replicate,
ValueRange inputs, Type type,
BlockArgument& block_arg) {
// check that inputs length is same as num_replicas.
if (inputs.size() != replicate.getN()) {
return replicate.emitError()
<< "Expected numper of inputs (" << inputs.size()
<< ") to append to replicate to be num_replicas ("
<< replicate.getN() << ")";
}
// add block arg to region. This is in $body.
unsigned block_arg_idx = replicate.GetNumReplicatedBlockArguments();
Block& block = replicate.GetBody();
block_arg = block.insertArgument(block_arg_idx, type, replicate.getLoc());
// add to $replicated_inputs. This also updates OperandSegmentSizeAttr.
replicate.getReplicatedInputsMutable().append(inputs);
return success();
}
// Insert an XlaAllReduce.
Value CreateAllReduce(ReplicateOp replicate, OpBuilder& builder,
Value block_arg) {
// This group_assignment is a list of all replicas. This says that the
// reduction sent to each replica is over all replicas.
uint32_t num_replicas = replicate.getN();
llvm::SmallVector<int32_t, 4> group_assignment_val;
for (int i = 0; i < num_replicas; ++i) group_assignment_val.push_back(i);
Value group_assignment = ConstOp::create(
builder, block_arg.getLoc(),
DenseIntElementsAttr::get(
RankedTensorType::get({1, num_replicas}, builder.getIntegerType(32)),
group_assignment_val));
StringAttr reduce_op = builder.getStringAttr("Add");
StringAttr mode = builder.getStringAttr("CrossReplica");
return XlaAllReduceOp::create(builder, block_arg.getLoc(),
block_arg.getType(), block_arg,
group_assignment, reduce_op, mode);
}
// Creates a missing attribute error message.
std::string CreateMissingAttributeMsg(llvm::StringRef attribute) {
return llvm::formatv("requires attribute '{0}'", attribute).str();
}
LogicalResult GetTpuDeviceAssignment(
ClusterOp cluster, ReplicateOp replicate, mlir::ModuleOp module,
absl::StatusOr<TPUDeviceAssignment>& status_or_tpu_device_assignment) {
mlir::TF::RuntimeDevices devices;
if (failed(tensorflow::GetDevicesFromOp(module, &devices))) return failure();
uint32_t num_replicas = replicate.getN();
auto num_cores_per_replica_attr = cluster->getAttrOfType<mlir::IntegerAttr>(
tensorflow::kNumCoresPerReplicaAttr);
if (!num_cores_per_replica_attr)
return cluster.emitOpError(
CreateMissingAttributeMsg(tensorflow::kNumCoresPerReplicaAttr));
int num_cores_per_replica = num_cores_per_replica_attr.getInt();
auto topology_attr = cluster->getAttrOfType<StringAttr>("topology");
if (!topology_attr)
return cluster.emitOpError(
CreateMissingAttributeMsg(tensorflow::kTopologyAttr));
auto device_assignment_attr = cluster->getAttrOfType<mlir::ArrayAttr>(
tensorflow::kDeviceAssignmentAttr);
if (!device_assignment_attr)
return cluster.emitOpError(llvm::formatv("requires attribute '{0}'",
tensorflow::kDeviceAssignmentAttr)
.str());
auto status_or_device_coodinates =
tensorflow::GetDeviceCoordinates(device_assignment_attr);
if (!status_or_device_coodinates.ok())
return cluster.emitError()
<< "error in fetching tpu device coordinates: "
<< status_or_device_coodinates.status().message();
status_or_tpu_device_assignment =
tensorflow::GetTPUCompilationAndExecutionDevices(
devices.device_names(), num_replicas, num_cores_per_replica,
topology_attr.getValue(), status_or_device_coodinates.value());
return success();
}
// Move a broadcast into the XLA cluster, converting it to an XlaAllReduce. This
// skips if the element type is not known to be valid for XlaAllReduce.
LogicalResult MoveBroadcastToCluster(OpBuilder& builder,
OpBuilder& inner_builder,
ClusterOp cluster, ReplicateOp replicate,
llvm::DenseMap<Value, Value>& orig_to_new,
Value val_bcast, mlir::ModuleOp module) {
Type zero_type;
ShapeAttr shape_attr;
if (!GetDummyParams(builder, val_bcast, zero_type, shape_attr))
return success();
llvm::SmallVector<Value, 4> inputs;
inputs.push_back(val_bcast);
uint32_t num_replicas = replicate.getN();
absl::StatusOr<TPUDeviceAssignment> status_or_tpu_device_assignment;
if (failed(GetTpuDeviceAssignment(cluster, replicate, module,
status_or_tpu_device_assignment))) {
return failure();
}
if (!status_or_tpu_device_assignment.ok())
return cluster.emitError()
<< "error in fetching TPU compilation/execution devices: "
<< status_or_tpu_device_assignment.status().message();
llvm::ArrayRef<llvm::SmallVector<tensorflow::TPUDeviceAndHost, 8>>
tpu_devices = status_or_tpu_device_assignment.value().tpu_devices;
std::unordered_map<std::string, Value> host_to_fill;
for (uint32_t replica = 1; replica < num_replicas; ++replica) {
std::string host = tpu_devices[replica][0].host;
if (host_to_fill.find(host) == host_to_fill.end()) {
host_to_fill[host] =
CreateZeroInput(val_bcast.getLoc(), builder, zero_type, shape_attr);
}
inputs.push_back(host_to_fill[host]);
}
BlockArgument block_arg;
if (failed(AppendReplicatedInput(builder, replicate, inputs,
val_bcast.getType(), block_arg))) {
return failure();
}
OpBuilder before_cluster_builder(cluster);
IdentityOp assigned_id =
IdentityOp::create(before_cluster_builder, val_bcast.getLoc(),
block_arg.getType(), block_arg);
assigned_id->setAttr(kICIWeightDistributionMlirBridgeMarker,
before_cluster_builder.getBoolAttr(true));
std::string device = tensorflow::GetDeviceAliasForHostOfLogicalCore(0);
LaunchOp launch = tensorflow::WrapOpInLaunch(
&before_cluster_builder, val_bcast.getLoc(), assigned_id, device);
launch->setAttr(kICIWeightDistributionMlirBridgeMarker,
before_cluster_builder.getBoolAttr(true));
Value all_reduce =
CreateAllReduce(replicate, inner_builder, launch.getResult(0));
orig_to_new[val_bcast] = all_reduce;
return success();
}
// Move all suitable broadcasts across replicas to the `cluster` into the
// `cluster`.
LogicalResult MoveAllBroadcastsToCluster(ClusterOp cluster,
ReplicateOp replicate,
mlir::ModuleOp module) {
// TODO(b/325153657): Fix tpu_rewrite_pass so the parallel_execute case does
// not need to be skipped.
if (cluster->getParentOfType<ParallelExecuteOp>()) return success();
auto num_cores_per_replica_attr = cluster->getAttrOfType<mlir::IntegerAttr>(
tensorflow::kNumCoresPerReplicaAttr);
if (!num_cores_per_replica_attr)
return cluster.emitOpError(
CreateMissingAttributeMsg(tensorflow::kNumCoresPerReplicaAttr));
llvm::SetVector<Value> bcasts;
cluster->walk([&](Operation* op) {
if (op == cluster) return WalkResult::advance();
for (auto operand : op->getOperands()) {
Operation* scope = operand.getParentBlock()->getParentOp();
if (scope->isProperAncestor(replicate)) {
bcasts.insert(operand);
}
}
return WalkResult::advance();
});
OpBuilder builder(replicate);
OpBuilder inner_builder = OpBuilder::atBlockBegin(&cluster.getBody().front());
// mapping from original operand to new XlaAllReduce op.
llvm::DenseMap<Value, Value> orig_to_new;
for (Value bcast : bcasts) {
if (failed(MoveBroadcastToCluster(builder, inner_builder, cluster,
replicate, orig_to_new, bcast, module))) {
return failure();
}
}
cluster->walk([&](Operation* op) {
if (op == cluster) return WalkResult::advance();
for (OpOperand& operand : op->getOpOperands()) {
if (orig_to_new.count(operand.get())) {
operand.assign(orig_to_new[operand.get()]);
}
}
return WalkResult::advance();
});
return success();
}
void XlaBroadcast::runOnOperation() {
FuncOp func = getOperation();
mlir::ModuleOp module = func->getParentOfType<mlir::ModuleOp>();
if (!module) return signalPassFailure();
func.walk([&](ClusterOp cluster) {
if (auto replicate = cluster->getParentOfType<ReplicateOp>()) {
if (failed(MoveAllBroadcastsToCluster(cluster, replicate, module))) {
return signalPassFailure();
}
}
});
}
} // namespace
std::unique_ptr<OperationPass<FuncOp>> CreateXlaBroadcastPass() {
return std::make_unique<XlaBroadcast>();
}
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/passes/xla_cluster_formation.cc
+145
View File
@@ -0,0 +1,145 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <functional>
#include <memory>
#include "absl/strings/str_cat.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/Block.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/SymbolTable.h" // from @llvm-project
#include "mlir/Interfaces/CallInterfaces.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/call_graph_util.h"
#include "tensorflow/core/common_runtime/inline_function_utils.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
using mlir::Block;
using mlir::CallInterfaceCallable;
using mlir::CallOpInterface;
using mlir::ModuleOp;
using mlir::OpBuilder;
using mlir::Operation;
using mlir::OperationPass;
using mlir::SymbolTable;
using mlir::SymbolTableCollection;
using mlir::SymbolUserOpInterface;
using mlir::func::FuncOp;
#define GEN_PASS_DEF_XLACLUSTERFORMATIONPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc"
// Outlines partitioned call ops with `_XlaMustCompile` to device clusters.
struct XlaClusterFormationPass
: public impl::XlaClusterFormationPassBase<XlaClusterFormationPass> {
void runOnOperation() override;
};
void EncapsulatePartitionedCall(Operation *call_op,
mlir::StringAttr callee_name) {
OpBuilder builder(call_op);
auto cluster = mlir::tf_device::ClusterOp::create(builder, call_op->getLoc(),
call_op->getResultTypes());
cluster.getBody().push_back(new Block);
call_op->replaceAllUsesWith(cluster.getResults());
call_op->moveBefore(&cluster.GetBody(), cluster.GetBody().end());
builder.setInsertionPointToEnd(&cluster.GetBody());
mlir::tf_device::ReturnOp::create(builder, call_op->getLoc(),
call_op->getResults());
// Propagate necessary attributes to the cluster so that when it's outlined,
// the function will have correct attributes.
mlir::TF::CopyDeviceAndUnderscoredAttributes(call_op, cluster);
cluster->setAttr(mlir::TF::kClusterOutlinedFunctionNameAttr, callee_name);
cluster->setAttr(mlir::TF::kAllowSoftPlacementAttr,
builder.getBoolAttr(true));
}
// Encapsulate the first partitioned call that can be reached from
// `func` and is with compilation markers in a device cluster. For nested calls,
// if the outermost one has the markers, encapsulates the outermost call and
// returns. Otherwise, we'll keep going through inner calls until we found one.
mlir::LogicalResult EncapsulateFirstXlaCompilablePartitionedCalls(
FuncOp func, SymbolTableCollection &symbol_table_collection,
SymbolTable &symtab) {
auto has_no_compile_device_type = [](SymbolUserOpInterface op) {
return !op->hasAttr(mlir::TF::kCompileDeviceTypeAttr);
};
mlir::OpBuilder builder(func.getContext());
auto noinline_attr_name = absl::StrCat("tf.", tensorflow::kNoInlineAttr);
llvm::SmallVector<SymbolUserOpInterface> noinline_pcall_ops,
outermost_pcall_ops;
if (mlir::failed(
mlir::GetOpsOfTypeUntilMiss<mlir::TF::StatefulPartitionedCallOp,
mlir::TF::PartitionedCallOp>(
func, symtab, /*predicate*/ has_no_compile_device_type,
/*hits*/ noinline_pcall_ops,
/*first_misses*/ outermost_pcall_ops))) {
return mlir::failure();
}
// Cluster outermost partitioned calls with _xla_compile_device_type
// attribute.
for (auto &pcall_op : outermost_pcall_ops) {
auto call = llvm::cast<CallOpInterface>(pcall_op.getOperation());
CallInterfaceCallable callable = call.getCallableForCallee();
auto sym = callable.get<mlir::SymbolRefAttr>();
EncapsulatePartitionedCall(pcall_op, sym.getRootReference());
}
// Partitioned calls are executed asynchronous. The calls outside of
// device clusters therefore should not be inlined to perserve run-time
// performance.
for (auto &pcall_op : noinline_pcall_ops) {
auto call = llvm::cast<CallOpInterface>(pcall_op.getOperation());
auto callee = llvm::cast<FuncOp>(
call.resolveCallableInTable(&symbol_table_collection));
callee->setAttr(noinline_attr_name, builder.getBoolAttr(true));
}
return mlir::success();
}
void XlaClusterFormationPass::runOnOperation() {
ModuleOp module = getOperation();
SymbolTableCollection symbol_table_collection;
SymbolTable symtab = symbol_table_collection.getSymbolTable(module);
llvm::SmallVector<FuncOp> entry_funcs = GetEntryFunctions(module);
for (auto &entry_func : entry_funcs) {
if (EncapsulateFirstXlaCompilablePartitionedCalls(
entry_func, symbol_table_collection, symtab)
.failed()) {
return signalPassFailure();
}
}
}
std::unique_ptr<OperationPass<ModuleOp>> CreateXlaClusterFormationPass() {
return std::make_unique<XlaClusterFormationPass>();
}
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/reproducer.proto
+38
View File
@@ -0,0 +1,38 @@
// Copyright 2024 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
syntax = "proto3";
package tensorflow.mlir.tf2xla.internal;
import "tensorflow/core/framework/function.proto";
import "tensorflow/core/framework/tensor_shape.proto";
import "tensorflow/core/protobuf/tpu/compile_metadata.proto";
option cc_enable_arenas = true;
// The input to LegalizeMlirToHlo to reproduce a legalization failure.
message LegalizeMlirToHloReproducer {
// The TPUCompileMetadataProto used during legalization.
tensorflow.tpu.TPUCompileMetadataProto compile_metadata = 1;
// The input shapes used during legalization.
repeated tensorflow.TensorShapeProto input_shapes = 2;
oneof computation {
// The MLIR module to be legalized.
string mlir_module = 3;
// The FunctionDefLibrary to be legalized.
tensorflow.FunctionDefLibrary function_def_library = 4;
}
}
tensorflow/compiler/mlir/tf2xla/internal/test_matchers.h
+91
View File
@@ -0,0 +1,91 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_TEST_MATCHERS_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_TEST_MATCHERS_H_
#include <gmock/gmock.h>
#include "absl/status/statusor.h"
#include "tensorflow/compiler/mlir/tf2xla/api/v1/compile_mlir_util.h"
#include "xla/tsl/platform/statusor.h"
template <typename T>
bool WasGraphAnalysisFailure(const absl::StatusOr<T>& status) {
return (status.status() ==
tensorflow::CompileToHloGraphAnalysisFailedError());
}
/* The third party version of the Graph Analysis always returns disabled so
* these matchers short circuit on that error. */
MATCHER(IsOkOrFiltered,
"Status was OK or equal to the Graph Analysis failure") {
bool is_ok = arg.ok();
auto graph_analysis_failure = WasGraphAnalysisFailure(arg);
return testing::ExplainMatchResult(
testing::IsTrue(), is_ok || graph_analysis_failure, result_listener);
}
MATCHER_P2(IncrementedOrFiltered, metric, value,
"Metric was incremented by value or Status equal to the Graph "
"Analysis failure") {
auto graph_analysis_failure = WasGraphAnalysisFailure(arg);
if (graph_analysis_failure) {
return testing::ExplainMatchResult(testing::IsTrue(),
graph_analysis_failure, result_listener);
}
return testing::ExplainMatchResult(testing::Eq(metric), value,
result_listener);
}
MATCHER_P(ComputationProtoContains, regex,
"If not a Graph Analysis failure then matches the computation result "
"with the regex") {
auto graph_analysis_failure = WasGraphAnalysisFailure(arg);
if (graph_analysis_failure) {
return testing::ExplainMatchResult(testing::IsTrue(),
graph_analysis_failure, result_listener);
}
auto proto = arg.value().computation->proto().DebugString();
return testing::ExplainMatchResult(testing::ContainsRegex(regex), proto,
result_listener);
}
MATCHER_P(XlaComputationProtoContains, regex,
"If not a Graph Analysis failure then matches the computation result "
"with the regex") {
auto graph_analysis_failure = WasGraphAnalysisFailure(arg);
if (graph_analysis_failure) {
return testing::ExplainMatchResult(testing::IsTrue(),
graph_analysis_failure, result_listener);
}
auto proto = arg.value().proto().DebugString();
return testing::ExplainMatchResult(testing::ContainsRegex(regex), proto,
result_listener);
}
MATCHER_P(
HasMlirModuleWith, expected,
"If not a Graph Analysis failure then matches the mlir module result") {
auto graph_analysis_failure = WasGraphAnalysisFailure(arg);
if (graph_analysis_failure) {
return testing::ExplainMatchResult(testing::IsTrue(),
graph_analysis_failure, result_listener);
}
auto actual = arg.value();
return testing::ExplainMatchResult(testing::ContainsRegex(expected), actual,
result_listener);
}
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_TEST_MATCHERS_H_
tensorflow/compiler/mlir/tf2xla/internal/test_matchers_test.cc
+136
View File
@@ -0,0 +1,136 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/internal/test_matchers.h"
#include <cstdint>
#include <memory>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "tensorflow/compiler/mlir/tf2xla/api/v1/compile_mlir_util.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/hlo/builder/xla_computation.h"
#include "xla/service/hlo.pb.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/lib/monitoring/cell_reader.h"
#include "tensorflow/core/lib/monitoring/counter.h"
namespace {
using ::tensorflow::monitoring::testing::CellReader;
using ::testing::Not;
constexpr char kMetric[] = "/tensorflow/metric";
auto* counter =
tensorflow::monitoring::Counter<1>::New(kMetric, "description", "status");
constexpr char kOkStatus[] = "ok";
const int kArbitraryIntResult = 37;
template <typename T>
tsl::StatusOr<T> success(T t) {
return t;
}
absl::StatusOr<int> success() { return kArbitraryIntResult; }
template <typename T>
tsl::StatusOr<T> filtered(T t) {
return tsl::StatusOr<T>(tensorflow::CompileToHloGraphAnalysisFailedError());
}
absl::StatusOr<int> filtered() { return filtered(kArbitraryIntResult); }
absl::StatusOr<int> failed() {
return absl::StatusOr<int>(absl::InternalError("fail"));
}
TEST(TestUtil, MatchesOk) { ASSERT_THAT(success(), IsOkOrFiltered()); }
TEST(TestUtil, DoesntMatchesFailure) {
ASSERT_THAT(failed(), Not(IsOkOrFiltered()));
}
TEST(TestUtil, MatchesFiltered) { ASSERT_THAT(filtered(), IsOkOrFiltered()); }
TEST(TestUtil, IncrementsOk) {
CellReader<int64_t> reader(kMetric);
counter->GetCell(kOkStatus)->IncrementBy(1);
ASSERT_THAT(success(), IncrementedOrFiltered(reader.Delta(kOkStatus), 1));
}
TEST(TestUtil, FilteredDoesntIncrementsOk) {
CellReader<int64_t> reader(kMetric);
ASSERT_THAT(filtered(), IncrementedOrFiltered(reader.Delta(kOkStatus), 1));
}
TEST(TestUtil, FailureDoesntMatchIncrement) {
CellReader<int64_t> reader(kMetric);
ASSERT_THAT(failed(), Not(IncrementedOrFiltered(reader.Delta(kOkStatus), 1)));
}
tensorflow::XlaCompilationResult CreateXlaComputationResult(
const char* hlo_name) {
auto result = tensorflow::XlaCompilationResult();
xla::HloModuleProto hlo;
hlo.set_name(hlo_name);
result.computation = std::make_shared<xla::XlaComputation>(hlo);
return result;
}
TEST(TestUtil, ComputationContainsOk) {
constexpr char arbitrary_hlo[] = "arbitrary_hlo";
auto result = CreateXlaComputationResult(arbitrary_hlo);
ASSERT_THAT(success(result), ComputationProtoContains(arbitrary_hlo));
}
TEST(TestUtil, ComputationDoesNotContain) {
constexpr char arbitrary_hlo[] = "arbitrary_hlo";
constexpr char bad_hlo[] = "bad_hlo";
auto result = CreateXlaComputationResult(arbitrary_hlo);
ASSERT_THAT(success(result), Not(ComputationProtoContains(bad_hlo)));
}
TEST(TestUtil, ComputationDoesNotContainFiltered) {
constexpr char arbitrary_hlo[] = "arbitrary_hlo";
constexpr char bad_hlo[] = "bad_hlo";
auto result = CreateXlaComputationResult(arbitrary_hlo);
ASSERT_THAT(filtered(result), ComputationProtoContains(bad_hlo));
}
TEST(TestUtil, MlirModuleHas) {
constexpr char arbirary_mlir[] = "arbirary_mlir";
ASSERT_THAT(success(arbirary_mlir), HasMlirModuleWith(arbirary_mlir));
}
TEST(TestUtil, MlirModuleDoesNotHave) {
constexpr char arbirary_mlir[] = "arbirary_mlir";
constexpr char bad_mlir[] = "bad_mlir";
ASSERT_THAT(success(arbirary_mlir), Not(HasMlirModuleWith(bad_mlir)));
}
TEST(TestUtil, MlirModuleDoesNotHaveFiltered) {
constexpr char arbirary_mlir[] = "arbirary_mlir";
constexpr char bad_mlir[] = "bad_mlir";
ASSERT_THAT(filtered(arbirary_mlir), HasMlirModuleWith(bad_mlir));
}
} // namespace
tensorflow/compiler/mlir/tf2xla/internal/testdata/dead_const.mlir
Vendored
+21
View File
@@ -0,0 +1,21 @@
// 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() -> () {
%s0 = "tf.Const"() {value = dense<[4, 1]> : tensor<2xi32>} : () -> tensor<2xi32>
func.return
}
}
tensorflow/compiler/mlir/tf2xla/internal/utils/BUILD
+79
View File
@@ -0,0 +1,79 @@
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow/compiler/mlir/tf2xla/internal:__subpackages__",
],
licenses = ["notice"],
)
cc_library(
name = "dialect_detection_utils",
srcs = [
"dialect_detection_utils.cc",
],
hdrs = [
"dialect_detection_utils.h",
],
deps = [
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/core:framework",
"//tensorflow/core/transforms/toposort:Pass",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:Transforms",
],
)
tf_cc_test(
name = "dialect_detection_utils_test",
srcs = ["dialect_detection_utils_test.cc"],
deps = [
":dialect_detection_utils",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@stablehlo//:chlo_ops",
],
)
cc_library(
name = "test_metadata_config",
testonly = True,
srcs = ["test_metadata_config.cc"],
hdrs = ["test_metadata_config.h"],
visibility = [
"//tensorflow/compiler/mlir/tf2xla/api:__subpackages__",
],
deps = [
"//tensorflow/compiler/jit",
"//tensorflow/compiler/jit:xla_tpu_device",
"//tensorflow/compiler/mlir:register_common_dialects",
"//tensorflow/compiler/mlir/tensorflow:convert_type",
"//tensorflow/compiler/mlir/tensorflow:deserialize_mlir_module_utils",
"//tensorflow/compiler/tf2xla:layout_util",
"//tensorflow/compiler/tf2xla:xla_compiler",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/compiler/tf2xla:xla_tpu_backend_registration",
"//tensorflow/compiler/tf2xla/kernels:xla_ops",
"//tensorflow/core:framework",
"//tensorflow/core:test_main",
"//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
"//tensorflow/core/tpu/kernels/xla:host_compute_ops",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:string_view",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@xla//xla:shape_util",
"@xla//xla/hlo/translate/mhlo_to_hlo:type_to_shape",
"@xla//xla/mlir_hlo:hlo_dialect_registration",
"@xla//xla/tsl/platform:errors",
],
)
tensorflow/compiler/mlir/tf2xla/internal/utils/dialect_detection_utils.cc
+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.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/internal/utils/dialect_detection_utils.h"
#include <set>
#include <string>
#include "mlir/IR/Dialect.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
namespace tensorflow {
namespace tf2xla {
namespace internal {
bool IsInBridgeAcceptableDialects(mlir::Operation* op) {
const std::set<std::string> kBuiltinNamespaces = {"func", "return",
"builtin"};
const std::set<std::string> kBridgeAcceptableNamespaces = {"tf", "tf_device"};
bool isInDefaulNamespaces =
kBuiltinNamespaces.find(op->getDialect()->getNamespace().str()) !=
kBuiltinNamespaces.end();
bool isInBridgeAcceptableNamespaces =
kBridgeAcceptableNamespaces.find(
op->getDialect()->getNamespace().str()) !=
kBridgeAcceptableNamespaces.end();
return isInDefaulNamespaces || isInBridgeAcceptableNamespaces;
}
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/utils/dialect_detection_utils.h
+33
View File
@@ -0,0 +1,33 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_UTILS_DIALECT_DETECTION_UTILS_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_UTILS_DIALECT_DETECTION_UTILS_H_
#include "mlir/IR/Operation.h" // from @llvm-project
namespace tensorflow {
namespace tf2xla {
namespace internal {
// Returns true if the op has a valid namespace during clustering & tf dialect
// to executor components of the Bridge.
bool IsInBridgeAcceptableDialects(mlir::Operation* op);
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_UTILS_DIALECT_DETECTION_UTILS_H_
tensorflow/compiler/mlir/tf2xla/internal/utils/dialect_detection_utils_test.cc
+76
View File
@@ -0,0 +1,76 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/internal/utils/dialect_detection_utils.h"
#include <gtest/gtest.h>
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/OperationSupport.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "stablehlo/dialect/ChloOps.h" // from @stablehlo
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {
using mlir::MLIRContext;
using mlir::OpBuilder;
using mlir::Operation;
using mlir::OperationState;
using mlir::UnknownLoc;
using mlir::chlo::ChloDialect;
using mlir::TF::TensorFlowDialect;
using tensorflow::tf2xla::internal::IsInBridgeAcceptableDialects;
class SharedUtilsTest : public ::testing::Test {};
TEST_F(SharedUtilsTest, IsInFunctionalDialectPasses) {
MLIRContext context;
context.loadDialect<TensorFlowDialect>();
OpBuilder opBuilder(&context);
OperationState state(UnknownLoc::get(opBuilder.getContext()),
/*OperationName=*/"tf.Const");
mlir::Operation* op = Operation::create(state);
bool result = IsInBridgeAcceptableDialects(op);
EXPECT_TRUE(result);
op->destroy();
}
TEST_F(SharedUtilsTest, IsInFunctionalDialectFails) {
MLIRContext context;
context.loadDialect<ChloDialect>();
OpBuilder opBuilder(&context);
OperationState state(UnknownLoc::get(opBuilder.getContext()),
/*OperationName=*/"chlo.broadcast_add");
Operation* op = Operation::create(state);
bool result = IsInBridgeAcceptableDialects(op);
EXPECT_FALSE(result);
op->destroy();
}
} // namespace
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/utils/test_metadata_config.cc
+107
View File
@@ -0,0 +1,107 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/internal/utils/test_metadata_config.h"
#include <vector>
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/DialectRegistry.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "tensorflow/compiler/mlir/register_common_dialects.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/convert_type.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/deserialize_mlir_module_utils.h"
#include "xla/hlo/translate/mhlo_to_hlo/type_to_shape.h"
#include "xla/mlir_hlo/mhlo/IR/register.h"
#include "xla/shape.h"
#include "xla/tsl/platform/errors.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {
constexpr char kEntryFuncName[] = "main";
absl::Status SetupArguments(mlir::ModuleOp module,
std::vector<TensorShape>& arg_shapes,
tpu::TPUCompileMetadataProto& metadata_proto) {
auto main_fn = module.lookupSymbol<mlir::func::FuncOp>(kEntryFuncName);
if (!main_fn) {
return absl::InternalError("Could not find main function in MLIR Module.");
}
mlir::FunctionType func_type = main_fn.getFunctionType();
for (auto input_type : func_type.getInputs()) {
tensorflow::TensorShape tensor_shape;
xla::Shape xla_shape = xla::TypeToShape(input_type);
TF_RETURN_IF_ERROR(tensorflow::TensorShape::BuildTensorShape(
xla_shape.dimensions(), &tensor_shape));
arg_shapes.emplace_back(tensor_shape);
DataType dtype;
TF_RETURN_IF_ERROR(ConvertToDataType(input_type, &dtype));
auto metadata_arg = metadata_proto.add_args();
metadata_arg->set_kind(tpu::TPUCompileMetadataProto::Arg::PARAMETER);
metadata_arg->set_dtype(dtype);
}
return absl::OkStatus();
}
absl::Status SetupReturnValues(mlir::ModuleOp module,
tpu::TPUCompileMetadataProto& metadata_proto) {
auto main_fn = module.lookupSymbol<mlir::func::FuncOp>(kEntryFuncName);
if (!main_fn) {
return absl::InternalError("Could not find main function in MLIR Module.");
}
int func_results = main_fn.getFunctionType().getNumResults();
for (int i = 0; i < func_results; i++) {
metadata_proto.add_retvals();
}
return absl::OkStatus();
}
} // namespace
absl::Status ConfigureMetadata(absl::string_view mlir_module_str,
std::vector<TensorShape>& arg_shapes,
tpu::TPUCompileMetadataProto& metadata_proto) {
mlir::DialectRegistry registry;
mlir::RegisterCommonToolingDialects(registry);
mlir::MLIRContext context(registry);
mlir::OwningOpRef<mlir::ModuleOp> mlir_module;
TF_RETURN_IF_ERROR(
DeserializeMlirModule(mlir_module_str, &context, &mlir_module));
TF_RETURN_IF_ERROR(SetupReturnValues(*mlir_module, metadata_proto));
TF_RETURN_IF_ERROR(SetupArguments(*mlir_module, arg_shapes, metadata_proto));
return absl::OkStatus();
}
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/internal/utils/test_metadata_config.h
+41
View File
@@ -0,0 +1,41 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_UTILS_TEST_METADATA_CONFIG_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_UTILS_TEST_METADATA_CONFIG_H_
#include <variant>
#include <vector>
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
namespace tensorflow {
namespace tf2xla {
namespace internal {
// Fills in arg_shapes and metadata_proto with appropriate values based on the
// input mlir module.
absl::Status ConfigureMetadata(absl::string_view mlir_module_str,
std::vector<TensorShape>& arg_shapes,
tpu::TPUCompileMetadataProto& metadata_proto);
} // namespace internal
} // namespace tf2xla
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_UTILS_TEST_METADATA_CONFIG_H_