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/transforms/BUILD
+532
View File
@@ -0,0 +1,532 @@
# Description:
# TF2XLA Bridge transforms
load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library")
load("@xla//xla/tsl/platform:build_config_root.bzl", "if_static")
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:public"],
licenses = ["notice"],
)
gentbl_cc_library(
name = "legalize_tf_patterns_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"generated_legalize_tf.inc": ["-gen-rewriters"]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "legalize_tf_patterns.td",
deps = [
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops_td_files",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncTdFiles",
"@llvm-project//mlir:TensorOpsTdFiles",
"@stablehlo//:chlo_ops_td_files",
"@stablehlo//:stablehlo_ops_td_files",
"@xla//xla/mlir_hlo:hlo_ops_td_files",
],
)
gentbl_cc_library(
name = "xla_legalize_tf_passes_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"xla_legalize_tf_passes.h.inc": [
"-gen-pass-decls",
"-name=LegalizeTf",
]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "xla_legalize_tf_passes.td",
deps = [
"@llvm-project//mlir:PassBaseTdFiles",
],
)
gentbl_cc_library(
name = "tf_xla_passes_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"tf_xla_passes.h.inc": [
"-gen-pass-decls",
"-name=TfXla",
]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "tf_xla_passes.td",
deps = [
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops_td_files",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncTdFiles",
"@llvm-project//mlir:PassBaseTdFiles",
"@llvm-project//mlir:SparseTensorDialect",
"@llvm-project//mlir:TensorOpsTdFiles",
"@xla//xla/mlir_hlo:hlo_ops_td_files",
],
)
cc_library(
name = "tf_xla_passes",
srcs = [
"xla_legalize_tf_passes.h.inc",
],
hdrs = [
"passes.h",
],
deps = [
":tf_xla_passes_inc_gen",
":xla_legalize_tf",
":xla_legalize_tf_passes_inc_gen",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:MemRefDialect",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:SparseTensorDialect",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
"@xla//xla/mlir_hlo",
],
)
cc_library(
name = "legalize_utils",
srcs = ["utils.cc"],
hdrs = ["utils.h"],
deps = [
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@xla//xla/mlir_hlo",
],
)
cc_library(
name = "test_utils",
testonly = True,
srcs = ["test_utils.cc"],
hdrs = ["test_utils.h"],
deps = [
"//tensorflow/compiler/mlir:register_common_dialects",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:deserialize_mlir_module_utils",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest_main",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@xla//xla/tsl/platform:statusor",
],
)
cc_library(
name = "legalize_tf",
srcs = [
"generated_legalize_tf.inc",
"legalize_tf.cc",
],
hdrs = [
"passes.h",
],
# DEPRECATED: use v2/legalize_tf.h::LegalizeMlirToHlo instead.
visibility = [
"//tensorflow/compiler/mlir/tensorflow/transforms:__pkg__",
"//tensorflow/compiler/mlir/tf2xla/internal/passes:__pkg__",
],
deps = [
":legalize_tf_patterns_inc_gen",
":legalize_utils",
":tf_xla_passes_inc_gen",
":xla_legalize_tf_passes_inc_gen",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:dynamic_shape_utils",
"//tensorflow/compiler/mlir/tensorflow:xla_sharding_util",
"//tensorflow/core:framework",
"//tensorflow/core/kernels:conv_grad_shape_utils",
"@com_google_absl//absl/status",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:Dialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:MemRefDialect",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:ShapeDialect",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TensorDialect",
"@llvm-project//mlir:TransformUtils",
"@stablehlo//:chlo_ops",
"@stablehlo//:stablehlo_ops",
"@stablehlo//:stablehlo_pass_utils",
"@tsl//tsl/platform:bfloat16",
"@tsl//tsl/platform:tensor_float_32_hdr_lib",
"@xla//xla:xla_data_proto_cc",
"@xla//xla/hlo/builder:padding",
"@xla//xla/hlo/builder:sharding_builder",
"@xla//xla/hlo/builder/lib:conv_grad_size_util",
"@xla//xla/hlo/translate/hlo_to_mhlo:attribute_importer",
"@xla//xla/mlir_hlo",
"@xla//xla/mlir_hlo:convert_op_folder",
"@xla//xla/tsl/platform:status",
] + if_static(["@tsl//tsl/platform:tensor_float_32_utils"]),
)
cc_library(
name = "xla_legalize_targets",
srcs = [
"xla_legalize_targets.cc",
],
hdrs = [
"xla_legalize_targets.h",
],
deps = [
"//tensorflow/compiler/mlir/tensorflow",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:ShapeDialect",
"@llvm-project//mlir:TensorDialect",
"@llvm-project//mlir:TransformUtils",
"@stablehlo//:chlo_ops",
"@stablehlo//:stablehlo_ops",
"@xla//xla/mlir_hlo",
],
)
tf_cc_test(
name = "xla_legalize_targets_test",
srcs = ["xla_legalize_targets_test.cc"],
deps = [
":xla_legalize_targets",
"//tensorflow/compiler/mlir/tensorflow",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:ShapeDialect",
"@llvm-project//mlir:TensorDialect",
"@llvm-project//mlir:TransformUtils",
"@stablehlo//:chlo_ops",
],
)
tf_cc_test(
name = "verify_tfxla_legalization_test",
srcs = ["verify_tfxla_legalization_test.cc"],
deps = [
":legalize_tf",
":test_utils",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:serialize_mlir_module_utils",
"//tensorflow/core:test",
"//tensorflow/core/lib/monitoring:cell_reader",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest_main",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@xla//xla/tsl/platform:statusor",
],
)
cc_library(
name = "xla_legalize_tf",
srcs = [
"infeed_ops_xla_adjust_layout.cc",
"legalize_tf_collective.cc",
"legalize_tf_communication.cc",
"tf_xla_passes.h.inc",
"tfxla_device_specific_transforms.cc",
"verify_tfxla_legalization.cc",
"xla_legalize_tf.cc",
"xla_legalize_tf_passes.h.inc",
],
hdrs = [
"passes.h",
],
deps = [
":legalization_op_config",
":legalize_tf",
":legalize_utils",
":tf_xla_passes_inc_gen",
":xla_legalize_targets",
":xla_legalize_tf_passes_inc_gen",
":xla_legalize_tf_with_tf2xla",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:mangling_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow/transforms:lower_tf_lib",
"//tensorflow/compiler/mlir/tensorflow/transforms:set_tpu_infeed_layout",
"//tensorflow/compiler/tf2xla/kernels:rng_converter_utils",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/util/quantization:uniform_quant_ops_params",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/log",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/synchronization",
"@com_google_absl//absl/types:span",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:FuncExtensions",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:MemRefDialect",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:QuantOps",
"@llvm-project//mlir:ShapeDialect",
"@llvm-project//mlir:SparseTensorDialect",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TensorDialect",
"@llvm-project//mlir:TransformUtils",
"@llvm-project//mlir:Transforms",
"@stablehlo//:base",
"@stablehlo//:chlo_ops",
"@stablehlo//:stablehlo_ops",
"@stablehlo//:stablehlo_passes",
"@xla//xla:shape_util",
"@xla//xla:side_effect_util",
"@xla//xla:xla_data_proto_cc",
"@xla//xla/hlo/builder:padding",
"@xla//xla/hlo/builder:sharding_builder",
"@xla//xla/hlo/translate/hlo_to_mhlo:attribute_importer",
"@xla//xla/hlo/translate/mhlo_to_hlo:type_to_shape",
"@xla//xla/mlir_hlo",
"@xla//xla/mlir_hlo:convert_op_folder",
"@xla//xla/mlir_hlo:mhlo_passes",
"@xla//xla/mlir_hlo:type_conversion",
"@xla//xla/tpu:c_api_conversions",
"@xla//xla/tpu:tpu_api",
],
)
cc_library(
name = "tf2xla_rewriter",
srcs = [
"tf2xla_rewriter.cc",
],
hdrs = [
"tf2xla_rewriter.h",
],
visibility = ["//visibility:private"],
deps = [
":legalize_tf",
"//tensorflow/compiler/mlir:op_or_arg_name_mapper",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:convert_tensor",
"//tensorflow/compiler/mlir/tensorflow:convert_type",
"//tensorflow/compiler/mlir/tensorflow:tpu_embedding_ops_registry",
"//tensorflow/compiler/mlir/tensorflow:translate_utils",
"//tensorflow/compiler/tf2xla:xla_compilation_device",
"//tensorflow/compiler/tf2xla:xla_context",
"//tensorflow/compiler/tf2xla:xla_expression",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/compiler/tf2xla:xla_op_registry",
"//tensorflow/core:core_cpu",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:framework_types_hdr",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/framework:allocator",
"//tensorflow/core/protobuf:for_core_protos_cc",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:span",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:SparseTensorDialect",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TensorDialect",
"@llvm-project//mlir:TransformUtils",
"@stablehlo//:base",
"@stablehlo//:stablehlo_ops",
"@xla//xla:xla_data_proto_cc",
"@xla//xla:xla_proto_cc",
"@xla//xla/hlo/builder:xla_builder",
"@xla//xla/hlo/builder:xla_computation",
"@xla//xla/hlo/ir:hlo",
"@xla//xla/hlo/translate/hlo_to_mhlo:hlo_function_importer",
"@xla//xla/hlo/translate/hlo_to_mhlo:hlo_to_mlir_hlo",
"@xla//xla/hlo/translate/mhlo_to_hlo:type_to_shape",
"@xla//xla/service:hlo_proto_cc",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:status",
"@xla//xla/tsl/platform:statusor",
],
)
tf_cc_test(
name = "tf2xla_rewriter_test",
srcs = [
"tf2xla_rewriter_test.cc",
],
deps = [
":test_utils",
":tf2xla_rewriter",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/tf2xla:xla_op_registry",
"//tensorflow/compiler/tf2xla/kernels:xla_ops",
"//tensorflow/core:ops",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_googletest//:gtest_main",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
"@stablehlo//:stablehlo_ops",
"@xla//xla:shape_util",
"@xla//xla:xla_data_proto_cc",
"@xla//xla/hlo/builder:xla_builder",
"@xla//xla/hlo/builder:xla_computation",
"@xla//xla/tsl/lib/core:status_test_util",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:status",
"@xla//xla/tsl/platform:statusor",
],
)
cc_library(
name = "xla_legalize_tf_with_tf2xla",
srcs = [
"legalize_tf_with_tf2xla.cc",
],
hdrs = [
"legalize_tf_with_tf2xla_passes.h",
"passes.h",
],
deps = [
":legalization_op_config",
":tf2xla_rewriter",
":tf_xla_passes_inc_gen",
":xla_legalize_tf_passes_inc_gen",
"//tensorflow/compiler/mlir:op_or_arg_name_mapper",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops",
"//tensorflow/compiler/mlir/tensorflow:tpu_embedding_ops_registry",
"//tensorflow/compiler/tf2xla:xla_compilation_device",
"//tensorflow/compiler/tf2xla:xla_context",
"//tensorflow/compiler/tf2xla:xla_expression",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/compiler/tf2xla:xla_op_registry",
"//tensorflow/core:core_cpu_lib",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/container:inlined_vector",
"@com_google_absl//absl/memory",
"@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:TensorDialect",
"@llvm-project//mlir:TransformUtils",
"@stablehlo//:base",
],
)
tf_cc_test(
name = "xla_legalize_tf_test",
srcs = [
"xla_legalize_tf_test.cc",
],
deps = [
":tf_xla_passes",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:deserialize_mlir_module_utils",
"//tensorflow/compiler/tf2xla:xla_compilation_device",
"//tensorflow/compiler/tf2xla:xla_context",
"//tensorflow/compiler/tf2xla:xla_expression",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/compiler/tf2xla:xla_op_registry",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core/framework:allocator",
"//tensorflow/core/lib/monitoring:cell_reader",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest_main",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:TransformUtils",
"@xla//xla/tsl/platform:statusor",
],
)
cc_library(
name = "legalization_op_config",
srcs = ["legalization_op_config.cc"],
hdrs = ["legalization_op_config.h"],
visibility = [
"//tensorflow/compiler/mlir/tensorflow/transforms:__pkg__",
"//tensorflow/compiler/mlir/tf2xla/internal:__subpackages__",
],
deps = [
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:tpu_embedding_ops_registry",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:Support",
],
)
tf_cc_test(
name = "legalization_op_config_test",
srcs = ["legalization_op_config_test.cc"],
deps = [
":legalization_op_config",
":legalize_tf",
"//tensorflow/compiler/jit",
"//tensorflow/compiler/mlir:register_common_dialects",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/tf2xla:xla_op_registry",
"//tensorflow/compiler/tf2xla/kernels:xla_ops",
"//tensorflow/core:protos_all_cc",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "split_into_island_per_op_pass",
srcs = ["split_into_island_per_op_pass.cc"],
hdrs = [
"split_into_island_per_op_pass.h",
],
deps = [
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_executor_inc_gen",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow/transforms:tf_pass_inc_gen",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:Dialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:InferTypeOpInterface",
"@llvm-project//mlir:Pass",
],
)
tf_cc_test(
name = "split_into_island_per_op_pass_test",
srcs = ["split_into_island_per_op_pass_test.cc"],
deps = [
":split_into_island_per_op_pass",
"//tensorflow/compiler/mlir:register_common_dialects",
"//tensorflow/compiler/mlir/tensorflow",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
)
tensorflow/compiler/mlir/tf2xla/transforms/infeed_ops_xla_adjust_layout.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 "llvm/ADT/StringRef.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/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Dialect.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Types.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/transforms/set_tpu_infeed_layout.h"
#include "xla/hlo/translate/mhlo_to_hlo/type_to_shape.h"
#include "xla/layout.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
#include "xla/shape.h"
namespace mlir {
namespace mhlo {
namespace {
#define GEN_PASS_DEF_INFEEDOPSXLAADJUSTLAYOUT
#include "tensorflow/compiler/mlir/tf2xla/transforms/xla_legalize_tf_passes.h.inc"
class InfeedsOpsXlaAdjustLayout
: public impl::InfeedOpsXlaAdjustLayoutBase<InfeedsOpsXlaAdjustLayout> {
public:
void runOnOperation() override;
private:
static void runOnInfeedOp(::mlir::mhlo::InfeedOp op) {
OpBuilder builder(op.getContext());
SmallVector<Type> result_types(op.getResultTypes().begin(),
op.getResultTypes().end());
if (!op->getAttr("layout")) {
auto layout = mlir::GetTPUInfeedLayout(result_types, builder);
if (failed(layout)) return;
op->setAttr("layout", layout.value());
}
}
};
void InfeedsOpsXlaAdjustLayout::runOnOperation() {
getOperation().walk(runOnInfeedOp);
}
} // anonymous namespace
std::unique_ptr<mlir::OperationPass<func::FuncOp>>
CreateInfeedsOpsXlaAdjustLayoutPass() {
return std::make_unique<InfeedsOpsXlaAdjustLayout>();
}
} // namespace mhlo
} // namespace mlir
tensorflow/compiler/mlir/tf2xla/transforms/legalization_op_config.cc
+564
View File
@@ -0,0 +1,564 @@
/* 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/transforms/legalization_op_config.h"
#include "llvm/ADT/DenseSet.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tpu_embedding_ops_registry.h"
namespace mlir {
namespace hlo {
namespace {
// Returns ops that should use MLIR legalization.
// All other ops not in this list should use XlaOpKernel.
const llvm::DenseSet<mlir::TypeID>& MlirAlwaysOps() {
// The static variable is a pointer in order to avoid destruction upon thread
// termination.
static const llvm::DenseSet<mlir::TypeID>* ops = new llvm::DenseSet<
mlir::TypeID>{
// Ops that should always use the MLIR legalization.
TypeID::get<TF::FusedBatchNormV3Op>(),
TypeID::get<TF::FusedBatchNormGradV3Op>(),
TypeID::get<TF::XlaReduceScatterOp>(),
TypeID::get<TF::ModOp>(),
// MatrixDiagPartV3 should use the MLIR implementation due to performance.
TypeID::get<TF::MatrixDiagPartV3Op>(),
// Ops that are legalized in the old bridge using MlirXlaOpKernel
TypeID::get<TF::AbsOp>(),
TypeID::get<TF::AtanOp>(),
TypeID::get<TF::AvgPool3DOp>(),
TypeID::get<TF::BiasAddGradOp>(),
TypeID::get<TF::CeilOp>(),
TypeID::get<TF::CheckNumericsOp>(),
TypeID::get<TF::CosOp>(),
TypeID::get<TF::TanOp>(),
TypeID::get<TF::DiagPartOp>(),
TypeID::get<TF::EinsumOp>(),
TypeID::get<TF::ExpOp>(),
TypeID::get<TF::Expm1Op>(),
TypeID::get<TF::FakeQuantWithMinMaxArgsOp>(),
TypeID::get<TF::FloorOp>(),
TypeID::get<TF::IFFTOp>(),
TypeID::get<TF::ImagOp>(),
TypeID::get<TF::IsFiniteOp>(),
TypeID::get<TF::IsInfOp>(),
TypeID::get<TF::IsNanOp>(),
TypeID::get<TF::LgammaOp>(),
TypeID::get<TF::Log1pOp>(),
TypeID::get<TF::LogSoftmaxOp>(),
TypeID::get<TF::MatrixBandPartOp>(),
TypeID::get<TF::MaxPool3DGradOp>(),
TypeID::get<TF::PreventGradientOp>(),
TypeID::get<TF::RandomShuffleOp>(),
TypeID::get<TF::RealOp>(),
TypeID::get<TF::ReciprocalOp>(),
TypeID::get<TF::ReluOp>(),
TypeID::get<TF::Relu6Op>(),
TypeID::get<TF::ReluGradOp>(),
TypeID::get<TF::RsqrtOp>(),
TypeID::get<TF::SelectOp>(),
TypeID::get<TF::SigmoidOp>(),
TypeID::get<TF::SignOp>(),
TypeID::get<TF::SoftmaxOp>(),
TypeID::get<TF::SqrtOp>(),
TypeID::get<TF::TanhOp>(),
TypeID::get<TF::XlaConvV2Op>(),
TypeID::get<TF::XlaDotOp>(),
TypeID::get<TF::XlaDotV2Op>(),
TypeID::get<TF::XlaDynamicSliceOp>(),
TypeID::get<TF::XlaEinsumOp>(),
TypeID::get<TF::XlaReduceWindowOp>(),
TypeID::get<TF::XlaReplicaIdOp>(),
TypeID::get<TF::XlaRngBitGeneratorOp>(),
TypeID::get<TF::XlaSelectAndScatterOp>(),
TypeID::get<TF::XlaSortOp>(),
TypeID::get<TF::XlaVariadicReduceV2Op>(),
TypeID::get<TF::XlaVariadicSortOp>(),
// Ops that have no XlaOpKernel.
TypeID::get<TF::RiscAddOp>(),
TypeID::get<TF::RiscDotOp>(),
// Const op has a simple legalization and it is much more efficient to
// lower
// within MLIR.
TypeID::get<TF::ConstOp>(),
// AssertOp with string types are not supported by the fallback.
TypeID::get<TF::AssertOp>(),
// TF2XLA fallback pattern doesn't support these op as MLIR hlo builder
// doesn't override the necessary builder methods. These ops have simple
// lowering pattern so this should be safe.
TypeID::get<TF::CrossReplicaSumOp>(),
TypeID::get<TF::InfeedDequeueTupleOp>(),
TypeID::get<TF::OutfeedEnqueueTupleOp>(),
TypeID::get<TF::XlaShardingOp>(),
// These ops have undetermined bugs, may not be legalizable with
// XlaOpKernel
// legalization in TF2XLA fallback. By legalization with MLIR, we can fix
// the bug. b/195583695 describes the motivation of this change.
// See b/216355804 how to reproduce the bug regarding tf.RandomUniform Op
// Conditional ops
TypeID::get<TF::IfRegionOp>(),
TypeID::get<TF::WhileRegionOp>(),
TypeID::get<TF::CaseRegionOp>(),
TypeID::get<TF::YieldOp>(),
};
return *ops;
}
bool IsOpTypeAllowedTf2XlaFallback(const TypeID& type_id) {
// Allowlisted TensorFlow ops are known to have well behaved tf2xla kernels
// building valid MLIR using MlirHloBuilder.
// TODO(hinsu): Drop explicit allowlist when MLIR based bridge is enabled for
// all tf2xla kernels.
// Use a pointer for the static set, so the set is not destructed upon thread
// end, which would not be thread safe.
static auto* ops = [] {
llvm::SmallDenseSet<mlir::TypeID, 512>* ops_set = new llvm::SmallDenseSet<
mlir::TypeID, 512>{
TypeID::get<TF::AcoshOp>(),
TypeID::get<TF::AcosOp>(),
TypeID::get<TF::AddNOp>(),
TypeID::get<TF::AddV2Op>(),
TypeID::get<TF::AngleOp>(),
TypeID::get<TF::AdjustContrastv2Op>(),
TypeID::get<TF::AdjustHueOp>(),
TypeID::get<TF::AdjustSaturationOp>(),
TypeID::get<TF::ApproximateEqualOp>(),
TypeID::get<TF::ApproxTopKOp>(),
TypeID::get<TF::ArgMaxOp>(),
TypeID::get<TF::ArgMinOp>(),
TypeID::get<TF::AsinhOp>(),
TypeID::get<TF::AsinOp>(),
TypeID::get<TF::Atan2Op>(),
TypeID::get<TF::AtanhOp>(),
TypeID::get<TF::BatchMatMulV2Op>(),
TypeID::get<TF::BatchMatMulV3Op>(),
TypeID::get<TF::BatchToSpaceOp>(),
TypeID::get<TF::BesselI0eOp>(),
TypeID::get<TF::BesselI1eOp>(),
TypeID::get<TF::BetaincOp>(),
TypeID::get<TF::BiasAddOp>(),
TypeID::get<TF::BitwiseAndOp>(),
TypeID::get<TF::BitwiseOrOp>(),
TypeID::get<TF::BitwiseXorOp>(),
TypeID::get<TF::BucketizeOp>(),
// CaseOp isn't actually supported but is enabled for testing to
// make sure ops with symbol ref attributes are filtered out.
TypeID::get<TF::CaseOp>(),
TypeID::get<TF::CastOp>(),
TypeID::get<TF::ClipByValueOp>(),
TypeID::get<TF::CholeskyOp>(),
TypeID::get<TF::CollectiveReduceV2Op>(),
TypeID::get<TF::ComplexAbsOp>(),
TypeID::get<TF::ConjugateTransposeOp>(),
TypeID::get<TF::ConcatV2Op>(),
TypeID::get<TF::ConvOp>(),
TypeID::get<TF::CoshOp>(),
TypeID::get<TF::CrossOp>(),
TypeID::get<TF::CumulativeLogsumexpOp>(),
TypeID::get<TF::DataFormatDimMapOp>(),
TypeID::get<TF::DataFormatVecPermuteOp>(),
TypeID::get<TF::DepthToSpaceOp>(),
TypeID::get<TF::DepthwiseConv2dNativeBackpropFilterOp>(),
TypeID::get<TF::DepthwiseConv2dNativeBackpropInputOp>(),
TypeID::get<TF::DiagOp>(),
TypeID::get<TF::DigammaOp>(),
TypeID::get<TF::DivNoNanOp>(),
TypeID::get<TF::DynamicPartitionOp>(),
TypeID::get<TF::EluGradOp>(),
TypeID::get<TF::EluOp>(),
TypeID::get<TF::EnsureShapeOp>(),
TypeID::get<TF::EqualOp>(),
TypeID::get<TF::ErfcOp>(),
TypeID::get<TF::ErfinvOp>(),
TypeID::get<TF::ErfOp>(),
TypeID::get<TF::ExtractImagePatchesOp>(),
TypeID::get<TF::FFT2DOp>(),
TypeID::get<TF::FFT3DOp>(),
TypeID::get<TF::FFTOp>(),
TypeID::get<TF::FakeParamOp>(),
TypeID::get<TF::FakeQuantWithMinMaxArgsGradientOp>(),
TypeID::get<TF::FakeQuantWithMinMaxVarsGradientOp>(),
TypeID::get<TF::FakeQuantWithMinMaxVarsPerChannelOp>(),
TypeID::get<TF::FakeQuantWithMinMaxVarsPerChannelGradientOp>(),
TypeID::get<TF::FloorDivOp>(),
TypeID::get<TF::FloorModOp>(),
TypeID::get<TF::GetMinibatchesInCsrWithPhysicalReplicaOp>(),
TypeID::get<TF::GetMinibatchSplitsWithPhysicalReplicaOp>(),
TypeID::get<TF::GreaterOp>(),
TypeID::get<TF::HSVToRGBOp>(),
TypeID::get<TF::IFFT2DOp>(),
TypeID::get<TF::IFFT3DOp>(),
TypeID::get<TF::IRFFT2DOp>(),
TypeID::get<TF::IRFFT3DOp>(),
TypeID::get<TF::IgammaOp>(),
TypeID::get<TF::IgammacOp>(),
TypeID::get<TF::IgammaGradAOp>(),
TypeID::get<TF::InplaceAddOp>(),
TypeID::get<TF::InTopKV2Op>(),
TypeID::get<TF::InvertOp>(),
TypeID::get<TF::InvOp>(),
TypeID::get<TF::KthOrderStatisticOp>(),
TypeID::get<TF::LRNOp>(),
TypeID::get<TF::LRNGradOp>(),
TypeID::get<TF::LeakyReluGradOp>(),
TypeID::get<TF::LeakyReluOp>(),
TypeID::get<TF::LeftShiftOp>(),
TypeID::get<TF::LessOp>(),
TypeID::get<TF::ListDiffOp>(),
TypeID::get<TF::LogicalAndOp>(),
TypeID::get<TF::LogicalNotOp>(),
TypeID::get<TF::LogOp>(),
TypeID::get<TF::LowerBoundOp>(),
TypeID::get<TF::MakeUniqueOp>(),
TypeID::get<TF::MatMulOp>(),
TypeID::get<TF::MatrixDiagV3Op>(),
TypeID::get<TF::MatrixInverseOp>(),
TypeID::get<TF::MatrixSetDiagV3Op>(),
TypeID::get<TF::MatrixSolveOp>(),
TypeID::get<TF::MatrixTriangularSolveOp>(),
TypeID::get<TF::MaxPool3DGradGradOp>(),
TypeID::get<TF::MaxPoolGradOp>(),
TypeID::get<TF::MaxPoolGradGradOp>(),
TypeID::get<TF::MirrorPadOp>(),
TypeID::get<TF::MirrorPadGradOp>(),
TypeID::get<TF::MulOp>(),
TypeID::get<TF::MultinomialOp>(),
TypeID::get<TF::NdtriOp>(),
TypeID::get<TF::NegOp>(),
TypeID::get<TF::NextAfterOp>(),
TypeID::get<TF::NonMaxSuppressionV4Op>(),
TypeID::get<TF::NotEqualOp>(),
TypeID::get<TF::PadOp>(),
TypeID::get<TF::ParameterizedTruncatedNormalOp>(),
TypeID::get<TF::PlaceholderWithDefaultOp>(),
TypeID::get<TF::PolygammaOp>(),
TypeID::get<TF::PopulationCountOp>(),
TypeID::get<TF::PowOp>(),
TypeID::get<TF::QrOp>(),
// TODO(hinsu): Canonicalize QuantizeAndDequantize and
// QuantizeAndDequantizeV2 to QuantizeAndDequantizeV3 by converting
// attributes to operands.
TypeID::get<TF::QuantizeAndDequantizeOp>(),
TypeID::get<TF::QuantizeAndDequantizeV2Op>(),
TypeID::get<TF::QuantizeAndDequantizeV3Op>(),
TypeID::get<TF::QuantizeAndDequantizeV4Op>(),
TypeID::get<TF::RFFT2DOp>(),
TypeID::get<TF::RFFT3DOp>(),
TypeID::get<TF::RGBToHSVOp>(),
TypeID::get<TF::RandomUniformIntOp>(),
TypeID::get<TF::RandomUniformOp>(),
TypeID::get<TF::RealDivOp>(),
TypeID::get<TF::ReciprocalGradOp>(),
TypeID::get<TF::Relu6GradOp>(),
TypeID::get<TF::ResizeBilinearOp>(),
TypeID::get<TF::ResizeBilinearGradOp>(),
TypeID::get<TF::ResizeNearestNeighborOp>(),
TypeID::get<TF::ResizeNearestNeighborGradOp>(),
TypeID::get<TF::ReverseSequenceOp>(),
TypeID::get<TF::RightShiftOp>(),
TypeID::get<TF::RintOp>(),
TypeID::get<TF::RollOp>(),
TypeID::get<TF::RoundOp>(),
TypeID::get<TF::SegmentSumV2Op>(),
TypeID::get<TF::SegmentProdV2Op>(),
TypeID::get<TF::SegmentMinV2Op>(),
TypeID::get<TF::SegmentMaxV2Op>(),
TypeID::get<TF::SelectV2Op>(),
TypeID::get<TF::SelfAdjointEigV2Op>(),
TypeID::get<TF::SeluGradOp>(),
TypeID::get<TF::SeluOp>(),
TypeID::get<TF::SigmoidGradOp>(),
TypeID::get<TF::SinOp>(),
TypeID::get<TF::SliceOp>(),
TypeID::get<TF::SoftplusGradOp>(),
TypeID::get<TF::SoftsignGradOp>(),
TypeID::get<TF::SoftsignOp>(),
TypeID::get<TF::SpaceToBatchNDOp>(),
TypeID::get<TF::SpaceToBatchOp>(),
TypeID::get<TF::SpaceToDepthOp>(),
TypeID::get<TF::SparseToDenseOp>(),
TypeID::get<TF::SquareOp>(),
TypeID::get<TF::StatelessMultinomialOp>(),
TypeID::get<TF::StatelessParameterizedTruncatedNormalOp>(),
TypeID::get<TF::StatelessRandomGetAlgOp>(),
TypeID::get<TF::StatelessRandomGetKeyCounterOp>(),
TypeID::get<TF::StatelessRandomGetKeyCounterAlgOp>(),
TypeID::get<TF::StatelessRandomNormalOp>(),
TypeID::get<TF::StatelessRandomNormalV2Op>(),
TypeID::get<TF::StatelessRandomUniformOp>(),
TypeID::get<TF::StatelessRandomUniformFullIntOp>(),
TypeID::get<TF::StatelessRandomUniformFullIntV2Op>(),
TypeID::get<TF::StatelessRandomUniformV2Op>(),
TypeID::get<TF::StatelessRandomUniformIntOp>(),
TypeID::get<TF::StatelessRandomUniformIntV2Op>(),
TypeID::get<TF::StatelessTruncatedNormalOp>(),
TypeID::get<TF::StatelessTruncatedNormalV2Op>(),
TypeID::get<TF::StoreMinibatchStatisticsInFdoOp>(),
TypeID::get<TF::StridedSliceOp>(),
TypeID::get<TF::SubOp>(),
TypeID::get<TF::SvdOp>(),
TypeID::get<TF::TanOp>(),
TypeID::get<TF::TensorScatterAddOp>(),
TypeID::get<TF::TensorScatterSubOp>(),
TypeID::get<TF::TPUEmbeddingActivationsOp>(),
TypeID::get<TF::TopKUniqueOp>(),
TypeID::get<TF::TopKWithUniqueOp>(),
TypeID::get<TF::TransposeOp>(),
TypeID::get<TF::TridiagonalSolveOp>(),
TypeID::get<TF::TridiagonalMatMulOp>(),
TypeID::get<TF::TruncateDivOp>(),
TypeID::get<TF::TruncatedNormalOp>(),
TypeID::get<TF::TruncateModOp>(),
TypeID::get<TF::UniqueOp>(),
TypeID::get<TF::UnpackOp>(),
TypeID::get<TF::UpperBoundOp>(),
TypeID::get<TF::WhereOp>(),
TypeID::get<TF::XlaSendTPUEmbeddingGradientsOp>(),
TypeID::get<TF::XlaBroadcastHelperOp>(),
TypeID::get<TF::XlaCallModuleOp>(),
TypeID::get<TF::XlaCustomCallV2Op>(),
TypeID::get<TF::XlaDynamicUpdateSliceOp>(),
TypeID::get<TF::XlaKeyValueSortOp>(),
TypeID::get<TF::XlaPadOp>(),
TypeID::get<TF::XlaSetBoundOp>(),
TypeID::get<TF::XlaSetDynamicDimensionSizeOp>(),
TypeID::get<TF::XlaSparseActivationsUnstackOp>(),
TypeID::get<TF::XlaSparseCoreAdagradMomentumOp>(),
TypeID::get<TF::XlaSparseCoreAdagradOp>(),
TypeID::get<TF::XlaSparseCoreAdamOp>(),
TypeID::get<TF::XlaSparseCoreFtrlOp>(),
TypeID::get<TF::XlaSparseCoreSgdOp>(),
TypeID::get<TF::XlaSparseDenseMatmulGradWithAdagradAndCsrInputOp>(),
TypeID::get<
TF::XlaSparseDenseMatmulGradWithAdagradMomentumAndCsrInputOp>(),
TypeID::get<TF::XlaSparseDenseMatmulGradWithAdamAndCsrInputOp>(),
TypeID::get<TF::XlaSparseDenseMatmulGradWithFtrlAndCsrInputOp>(),
TypeID::get<TF::XlaSparseDenseMatmulGradWithSgdAndCsrInputOp>(),
TypeID::get<TF::XlaSparseDenseMatmulWithCsrInputOp>(),
TypeID::get<TF::XlaSparseDenseMatmulCustomCombinerOnTcWithCsrInputOp>(),
TypeID::get<TF::XlaSparseDenseMatmulWithStaticBufferSizeOp>(),
TypeID::get<
TF::XlaSparseDenseMatmulGradWithAdagradAndStaticBufferSizeOp>(),
TypeID::get<
TF::XlaSparseDenseMatmulGradWithAdagradMomentumAndStaticBufferSizeOp>(), // NOLINT
TypeID::get<
TF::XlaSparseDenseMatmulGradWithAdamAndStaticBufferSizeOp>(),
TypeID::get<
TF::XlaSparseDenseMatmulGradWithFtrlAndStaticBufferSizeOp>(),
TypeID::get<
TF::XlaSparseDenseMatmulGradWithSgdAndStaticBufferSizeOp>(), // NOLINT
TypeID::get<TF::XlaSparseDenseMatmulGradWithCsrInputOp>(),
TypeID::get<
TF::XlaSparseDenseMatmulCustomCombinerOnTcGradWithSgdAndCsrInputOp>(), // NOLINT
TypeID::get<
TF::XlaSparseDenseMatmulCustomCombinerOnTcGradWithAdagradAndCsrInputOp>(), // NOLINT
TypeID::get<
TF::XlaSparseDenseMatmulCustomCombinerOnTcGradWithAdagradMomentumAndCsrInputOp>(), // NOLINT
TypeID::get<
TF::XlaSparseDenseMatmulCustomCombinerOnTcGradWithAdamAndCsrInputOp>(), // NOLINT
TypeID::get<
TF::XlaSparseDenseMatmulCustomCombinerOnTcGradWithFtrlAndCsrInputOp>(), // NOLINT
TypeID::get<
TF::XlaSparseDenseMatmulCustomCombinerOnTcGradWithCsrInputOp>(),
TypeID::get<TF::XlaLocalSparseDenseMatmulOp>(), // NOLINT
TypeID::get<TF::XlaSparseGradientsStackOp>(),
TypeID::get<TF::XlaSpmdFullToShardShapeOp>(),
TypeID::get<TF::XlaSpmdShardToFullShapeOp>(),
TypeID::get<TF::XlaSvdOp>(),
};
// Add the ops from the TPUEmbeddingOpsRegistry.
for (auto op_type_id :
TF::TPUEmbeddingOpsRegistry::Global().GetOpsTypeIds()) {
ops_set->insert(op_type_id);
}
return ops_set;
}();
return ops->count(type_id);
}
/// List of ops that should use XlaOpKernel legalization only in the case of
/// prefer_tf2xla. All other ops not in this list should use MLIR legalization
/// only or not be legalized by the new bridge.
bool IsOpTypeAllowedTf2XlaPreferred(const TypeID& type_id) {
// Use a pointer for the static set, so the set is not destructed upon thread
// end, which would not be thread safe.
// clang-format off
static auto* ops =
new llvm::SmallDenseSet<mlir::TypeID, 512>{
TypeID::get<TF::AllOp>(),
TypeID::get<TF::AllToAllOp>(),
TypeID::get<TF::AnyOp>(),
TypeID::get<TF::AvgPoolOp>(),
TypeID::get<TF::AvgPool3DGradOp>(),
TypeID::get<TF::AvgPoolGradOp>(),
TypeID::get<TF::BatchToSpaceNDOp>(),
TypeID::get<TF::BitcastOp>(),
TypeID::get<TF::BroadcastToOp>(),
TypeID::get<TF::CollectivePermuteOp>(),
TypeID::get<TF::ComplexOp>(),
TypeID::get<TF::ConcatV2Op>(),
TypeID::get<TF::ConjOp>(),
TypeID::get<TF::Conv2DOp>(),
TypeID::get<TF::Conv2DBackpropFilterOp>(),
TypeID::get<TF::Conv2DBackpropInputOp>(),
TypeID::get<TF::Conv3DOp>(),
TypeID::get<TF::Conv3DBackpropFilterV2Op>(),
TypeID::get<TF::Conv3DBackpropInputV2Op>(),
TypeID::get<TF::CumprodOp>(),
TypeID::get<TF::CumsumOp>(),
TypeID::get<TF::DepthwiseConv2dNativeOp>(),
TypeID::get<TF::DivOp>(),
TypeID::get<TF::DynamicStitchOp>(),
TypeID::get<TF::_EagerConstOp>(),
TypeID::get<TF::EmptyOp>(),
TypeID::get<TF::ExpandDimsOp>(),
TypeID::get<TF::FakeQuantWithMinMaxVarsOp>(),
TypeID::get<TF::FillOp>(),
TypeID::get<TF::FusedBatchNormOp>(),
TypeID::get<TF::FusedBatchNormGradOp>(),
TypeID::get<TF::FusedBatchNormGradV2Op>(),
TypeID::get<TF::FusedBatchNormV2Op>(),
TypeID::get<TF::_FusedConv2DOp>(),
TypeID::get<TF::GatherNdOp>(),
TypeID::get<TF::GatherV2Op>(),
TypeID::get<TF::GreaterEqualOp>(),
TypeID::get<TF::IdentityOp>(),
TypeID::get<TF::IdentityNOp>(),
TypeID::get<TF::InplaceUpdateOp>(),
TypeID::get<TF::InvertPermutationOp>(),
TypeID::get<TF::IRFFTOp>(),
TypeID::get<TF::L2LossOp>(),
TypeID::get<TF::LegacyCallOp>(),
TypeID::get<TF::LessEqualOp>(),
TypeID::get<TF::LinSpaceOp>(),
TypeID::get<TF::LogicalOrOp>(),
TypeID::get<TF::MaxOp>(),
TypeID::get<TF::MaximumOp>(),
TypeID::get<TF::MaxPoolOp>(),
TypeID::get<TF::MaxPool3DOp>(),
TypeID::get<TF::MeanOp>(),
TypeID::get<TF::MinOp>(),
TypeID::get<TF::MinimumOp>(),
TypeID::get<TF::MulNoNanOp>(),
TypeID::get<TF::OneHotOp>(),
TypeID::get<TF::OnesLikeOp>(),
TypeID::get<TF::PackOp>(),
TypeID::get<TF::PadV2Op>(),
TypeID::get<TF::ParallelDynamicStitchOp>(),
TypeID::get<TF::PartitionedCallOp>(),
TypeID::get<TF::ProdOp>(),
TypeID::get<TF::QrOp>(),
TypeID::get<TF::RandomStandardNormalOp>(),
TypeID::get<TF::RandomUniformOp>(),
TypeID::get<TF::RangeOp>(),
TypeID::get<TF::ReshapeOp>(),
TypeID::get<TF::ReverseV2Op>(),
TypeID::get<TF::RFFTOp>(),
TypeID::get<TF::RsqrtGradOp>(),
TypeID::get<TF::ScatterNdOp>(),
TypeID::get<TF::ShapeOp>(),
TypeID::get<TF::SinhOp>(),
TypeID::get<TF::SizeOp>(),
TypeID::get<TF::SliceOp>(),
TypeID::get<TF::SoftmaxCrossEntropyWithLogitsOp>(),
TypeID::get<TF::SoftplusOp>(),
TypeID::get<TF::SparseMatMulOp>(),
TypeID::get<TF::SparseSoftmaxCrossEntropyWithLogitsOp>(),
TypeID::get<TF::SplitOp>(),
TypeID::get<TF::SplitVOp>(),
TypeID::get<TF::SqrtGradOp>(),
TypeID::get<TF::SquaredDifferenceOp>(),
TypeID::get<TF::SqueezeOp>(),
TypeID::get<TF::StatelessParameterizedTruncatedNormalOp>(),
TypeID::get<TF::StatefulPartitionedCallOp>(),
TypeID::get<TF::StopGradientOp>(),
TypeID::get<TF::StridedSliceOp>(),
TypeID::get<TF::StridedSliceGradOp>(),
TypeID::get<TF::SumOp>(),
TypeID::get<TF::TanhGradOp>(),
TypeID::get<TF::TensorScatterUpdateOp>(),
TypeID::get<TF::TileOp>(),
TypeID::get<TF::TopKV2Op>(),
TypeID::get<TF::_UnaryOpsCompositionOp>(),
TypeID::get<TF::UnsortedSegmentMaxOp>(),
TypeID::get<TF::UnsortedSegmentMinOp>(),
TypeID::get<TF::UnsortedSegmentProdOp>(),
TypeID::get<TF::UnsortedSegmentSumOp>(),
TypeID::get<TF::XdivyOp>(),
TypeID::get<TF::XlaSendTPUEmbeddingGradientsOp>(),
TypeID::get<TF::XlaAllReduceOp>(),
TypeID::get<TF::XlaGatherOp>(),
TypeID::get<TF::Xlog1pyOp>(),
TypeID::get<TF::XlogyOp>(),
TypeID::get<TF::ZerosLikeOp>(),
TypeID::get<TF::ZetaOp>(),
};
// clang-format on
return ops->contains(type_id);
}
const llvm::DenseSet<mlir::TypeID>& DynamicTensorflowOps() {
// The static variable is a pointer in order to avoid destruction upon thread
// termination.
static const llvm::DenseSet<mlir::TypeID>* ops =
new llvm::DenseSet<mlir::TypeID>{
TypeID::get<mlir::TF::DynamicPartitionOp>(),
TypeID::get<mlir::TF::UniqueOp>(),
TypeID::get<mlir::TF::WhereOp>(),
TypeID::get<mlir::TF::XlaSetDynamicDimensionSizeOp>(),
};
return *ops;
}
} // namespace
bool HasTf2XlaFallback(const TypeID& type_id) {
return IsOpTypeAllowedTf2XlaFallback(type_id) ||
IsOpTypeAllowedTf2XlaPreferred(type_id);
}
bool IsTypeLegalizedWithMlir(const TypeID& type_id) {
return MlirAlwaysOps().contains(type_id);
}
bool IsOpAllowedTf2xlaFallback(const TypeID& type_id) {
return IsOpTypeAllowedTf2XlaFallback(type_id);
}
bool IsOpAllowedTf2xlaPreferred(const TypeID& type_id) {
return IsOpTypeAllowedTf2XlaPreferred(type_id);
}
bool IsDynamicPadderOp(const TypeID& type_id) {
return DynamicTensorflowOps().contains(type_id);
}
} // namespace hlo
} // namespace mlir
tensorflow/compiler/mlir/tf2xla/transforms/legalization_op_config.h
+45
View File
@@ -0,0 +1,45 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_TRANSFORMS_LEGALIZATION_OP_CONFIG_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_TRANSFORMS_LEGALIZATION_OP_CONFIG_H_
#include "mlir/Support/TypeID.h" // from @llvm-project
namespace mlir {
namespace hlo {
// Given the type ID, check if it's legalized with MLIR.
bool IsTypeLegalizedWithMlir(const TypeID& type_id);
// Returns true if the op is considered a dynamic padder op.
bool IsDynamicPadderOp(const TypeID& type_id);
// Returns True if this op has a Tf2XLA fallback. Currently, this is not the
// inverse of the !IsOpLegalizedWithMlir, but it should be.
bool HasTf2XlaFallback(const TypeID& type_id);
// Whether this type is allowed to have a TF2XLA fallback.
bool IsOpAllowedTf2xlaFallback(const TypeID& type_id);
// Whether this type is Preferred to use a TF2XLA fallback kernel when using
// the MLIR bridge. If this is true, then the TF2XLA fallback kernel will be
// used over the MLIR lowering.
bool IsOpAllowedTf2xlaPreferred(const TypeID& type_id);
} // namespace hlo
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_TRANSFORMS_LEGALIZATION_OP_CONFIG_H_
tensorflow/compiler/mlir/tf2xla/transforms/legalization_op_config_test.cc
+183
View File
@@ -0,0 +1,183 @@
/* 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/transforms/legalization_op_config.h"
#include <optional>
#include <set>
#include <string>
#include <vector>
#include <gtest/gtest.h>
#include "mlir/IR/DialectRegistry.h" // from @llvm-project
#include "mlir/IR/OperationSupport.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "tensorflow/compiler/mlir/register_common_dialects.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/passes.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "tensorflow/core/framework/kernel_def.pb.h"
namespace mlir {
namespace hlo {
TEST(LegalizationOpConfigTest, ExpectsTrueForMlirTypeID) {
EXPECT_TRUE(IsTypeLegalizedWithMlir(TypeID::get<TF::ModOp>()));
EXPECT_FALSE(HasTf2XlaFallback(TypeID::get<TF::ModOp>()));
EXPECT_FALSE(IsOpAllowedTf2xlaFallback(TypeID::get<TF::ModOp>()));
EXPECT_FALSE(IsOpAllowedTf2xlaPreferred(TypeID::get<TF::ModOp>()));
}
TEST(LegalizationOpConfigTest, ExpectsTrueForTF2XLATypeID) {
EXPECT_TRUE(HasTf2XlaFallback(TypeID::get<TF::AllOp>()));
EXPECT_TRUE(IsOpAllowedTf2xlaPreferred(TypeID::get<TF::AllOp>()));
EXPECT_FALSE(IsTypeLegalizedWithMlir(TypeID::get<TF::AllOp>()));
}
TEST(LegalizationOpConfigTest, ChecksDynamicPadderOps) {
EXPECT_TRUE(
IsDynamicPadderOp(TypeID::get<TF::XlaSetDynamicDimensionSizeOp>()));
EXPECT_FALSE(IsDynamicPadderOp(TypeID::get<TF::ConstOp>()));
}
// This test is kind of odd. We go through all the Tensorflow types and check
// whether they are legalized with MLIR, TF2XLA, or both. Ideally the sets are
// disjoint, but until that happens, this tests ensures that the set doesn't
// grow.
TEST(LegalizationOpConfigTest, CountLoweringsSet) {
int mlir_lowering_count = 0;
int tf2xla_fallback_count = 0;
int non_categorized_count = 0;
DialectRegistry dialect_registry;
dialect_registry.insert<mlir::TF::TensorFlowDialect>();
MLIRContext context(dialect_registry);
context.loadAllAvailableDialects();
for (auto operation : context.getRegisteredOperations()) {
if (IsTypeLegalizedWithMlir(operation.getTypeID())) {
mlir_lowering_count++;
} else if (HasTf2XlaFallback(operation.getTypeID())) {
tf2xla_fallback_count++;
} else {
non_categorized_count++;
}
}
// If an op moves from one lowering implementation to a different one (e.g.
// from MLIR to TF2XLA), these numbers should change. Or if TF Dialect adds
// a new op, we should expect these to change too.
EXPECT_EQ(mlir_lowering_count, 67);
EXPECT_EQ(tf2xla_fallback_count, 333);
EXPECT_EQ(non_categorized_count, 437);
}
// Just a counter test to see which ops have duplicate lowerings. This isn't a
// correctness test versus a test to easily ensure the op registry is kept
// in sync.
TEST(LegalizationOpConfigTest, CountTypesWhichHaveBothMlirAndTf2xlaFallback) {
int double_lowering_count = 0;
DialectRegistry dialect_registry;
dialect_registry.insert<mlir::TF::TensorFlowDialect>();
MLIRContext context(dialect_registry);
context.loadAllAvailableDialects();
for (auto operation : context.getRegisteredOperations()) {
if (IsTypeLegalizedWithMlir(operation.getTypeID()) &&
HasTf2XlaFallback(operation.getTypeID())) {
double_lowering_count++;
}
}
// TODO(b/288876609): This should get to zero.
EXPECT_EQ(double_lowering_count, 1);
}
// Counts which ops have MLIR only lowerings. This isn't a
// correctness test versus a test to easily ensure the op registry is kept
// in sync.
TEST(LegalizationOpConfigTest, CountAllMlirLoweringPatterns) {
DialectRegistry dialect_registry;
mlir::RegisterCommonToolingDialects(dialect_registry);
MLIRContext context(dialect_registry);
context.loadAllAvailableDialects();
RewritePatternSet mlir_legalize_lower_patterns(&context);
hlo::PopulateLegalizeTfPatterns(&context, &mlir_legalize_lower_patterns);
int mlir_only_patterns = 0;
for (auto& pattern : mlir_legalize_lower_patterns.getNativePatterns()) {
std::optional<OperationName> pat_op_name = pattern->getRootKind();
if (!pat_op_name) {
continue;
}
if (!HasTf2XlaFallback(pat_op_name->getRegisteredInfo()->getTypeID())) {
mlir_only_patterns++;
}
}
EXPECT_EQ(mlir_only_patterns, 63);
}
// Counts which ops have lowerings without XlaOpKernels. This isn't a
// correctness test versus a test to easily ensure the op registry is kept
// in sync.
TEST(LegalizationOpConfigTest, MlirLoweringWithoutXlaKernel) {
tensorflow::XlaOpRegistry::RegisterCompilationKernels();
std::vector<const tensorflow::KernelDef*> kernel_defs =
tensorflow::XlaOpRegistry::DeviceKernels(
tensorflow::DEVICE_CPU_XLA_JIT,
/*include_compilation_only_kernels=*/true);
std::set<std::string> xla_op_kernels;
for (auto kernel_def : kernel_defs) {
std::string tf_name = "tf." + kernel_def->op();
xla_op_kernels.insert(tf_name);
}
DialectRegistry dialect_registry;
mlir::RegisterCommonToolingDialects(dialect_registry);
MLIRContext context(dialect_registry);
context.loadAllAvailableDialects();
RewritePatternSet mlir_legalize_lower_patterns(&context);
hlo::PopulateLegalizeTfPatterns(&context, &mlir_legalize_lower_patterns);
int mlir_without_xla_count = 0;
for (auto& pattern : mlir_legalize_lower_patterns.getNativePatterns()) {
std::optional<OperationName> pat_op_name = pattern->getRootKind();
if (!pat_op_name) {
continue;
}
if (xla_op_kernels.find(pat_op_name->getStringRef().str()) ==
xla_op_kernels.end()) {
mlir_without_xla_count++;
}
}
EXPECT_EQ(mlir_without_xla_count, 13);
}
} // namespace hlo
} // namespace mlir
tensorflow/compiler/mlir/tf2xla/transforms/legalize_tf.cc
+7049
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/mlir/tf2xla/transforms/legalize_tf_collective.cc
+413
View File
@@ -0,0 +1,413 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
// This file implements logic for lowering TensorFlow dialect's collective
// ops (TF/XLA) to the HLO dialect.
#include <cstdint>
#include <memory>
#include <numeric>
#include <utility>
#include "absl/strings/string_view.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Dialect.h" // from @llvm-project
#include "mlir/IR/Matchers.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Support/DebugStringHelper.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "stablehlo/dialect/ChloOps.h" // from @stablehlo
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/utils.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
#include "xla/mlir_hlo/utils/convert_op_folder.h"
#include "xla/mlir_hlo/utils/hlo_utils.h"
#include "xla/xla_data.pb.h"
namespace mlir {
namespace mhlo {
namespace {
constexpr absl::string_view kGroupSizeAttrName =
"tf2xla.collective_info.group_size";
constexpr absl::string_view kGroupKeyAttrName =
"tf2xla.collective_info.group_key";
#define GEN_PASS_DEF_LEGALIZETFCOLLECTIVE
#include "tensorflow/compiler/mlir/tf2xla/transforms/xla_legalize_tf_passes.h.inc"
class LegalizeTFCollective
: public impl::LegalizeTFCollectiveBase<LegalizeTFCollective> {
public:
void runOnOperation() override;
};
LogicalResult SetOnceModuleAttribute(StringRef attr_name,
IntegerAttr attr_value, Operation* op,
ModuleOp& module) {
const auto ex_attr_value = module->getAttrOfType<IntegerAttr>(attr_name);
if (ex_attr_value == nullptr) {
module->setAttr(attr_name, attr_value);
return success();
}
if (ex_attr_value == attr_value) {
return success();
}
return op->emitOpError() << "module already contains an attribute "
<< attr_name << "=" << ex_attr_value.getInt()
<< ", overwritting to a new value "
<< attr_value.getInt() << " is not allowed.";
}
LogicalResult SetCollectiveInfo(IntegerAttr group_size, IntegerAttr group_key,
Operation* op) {
ModuleOp module = op->getParentOfType<ModuleOp>();
// The StringRef cast is necessary before cxx14.
if (failed(SetOnceModuleAttribute(
StringRef(kGroupSizeAttrName.data(), kGroupSizeAttrName.size()),
group_size, op, module))) {
return failure();
}
if (failed(SetOnceModuleAttribute(
StringRef(kGroupKeyAttrName.data(), kGroupKeyAttrName.size()),
group_key, op, module))) {
return failure();
}
return success();
}
LogicalResult SetCollectiveInfo(OpBuilder& builder,
DenseIntElementsAttr replica_groups,
Operation* op) {
// Use special group_key 0 to represent "all available devices". This
// shall resolve to a DeviceAssignment that includes all devices intended for
// replica_groups.
IntegerAttr group_size = builder.getI32IntegerAttr(replica_groups.size());
IntegerAttr group_key = builder.getI32IntegerAttr(0);
return SetCollectiveInfo(group_size, group_key, op);
}
LogicalResult ConvertReplicaGroups(OpBuilder& builder,
Value group_assignment_value,
DenseIntElementsAttr& replica_groups,
Operation* op) {
DenseIntElementsAttr group_assignment;
if (!matchPattern(group_assignment_value, m_Constant(&group_assignment))) {
return op->emitOpError() << "expects constant group_assignment";
}
replica_groups = mlir::cast<DenseIntElementsAttr>(
hlo::convertElementsAttr(group_assignment, builder.getIntegerType(64)));
if (replica_groups.getType().getRank() != 2) {
return op->emitOpError() << "group_assignment should have rank 2, got "
<< replica_groups.getType().getRank();
}
return success();
}
ChannelHandleAttr ConvertChannel(OpBuilder& builder, int64_t channel_id,
StringRef mode) {
if (mode == "CrossReplica") {
return ChannelHandleAttr();
}
return ChannelHandleAttr::get(builder.getContext(),
/*handle=*/channel_id,
/*type=*/xla::ChannelHandle::DEVICE_TO_DEVICE);
}
LogicalResult ConvertAllReduce(OpBuilder& builder, int64_t channel_id,
TensorType result_type,
DenseIntElementsAttr replica_groups,
StringRef mode, Value input, StringRef merge_op,
StringRef final_op, Operation* op) {
builder.setInsertionPoint(op);
ChannelHandleAttr channel_handle = ConvertChannel(builder, channel_id, mode);
Location loc = op->getLoc();
Type element_type = getElementTypeOrSelf(input.getType());
auto all_reduce =
AllReduceOp::create(builder, loc, result_type, input, replica_groups,
channel_handle, nullptr);
if (all_reduce.getNumResults() != 1) {
return op->emitOpError()
<< "AllReduceOp must have one result: " << *all_reduce;
}
if (merge_op == "Add") {
BuildReduceBody<AddOp>(element_type, &all_reduce.getComputation(),
&builder);
} else if (merge_op == "Mul") {
BuildReduceBody<MulOp>(element_type, &all_reduce.getComputation(),
&builder);
} else if (merge_op == "Min") {
BuildReduceBody<MinOp>(element_type, &all_reduce.getComputation(),
&builder);
} else if (merge_op == "Max") {
BuildReduceBody<MaxOp>(element_type, &all_reduce.getComputation(),
&builder);
} else {
return op->emitOpError() << "invalid merge_op " << merge_op
<< ", want one of [Add, Mul, Min, Max]";
}
Operation* result = all_reduce;
// For "Div" final op, divide the merge result by group size.
if (final_op == "Div") {
int64_t replica_group_size = replica_groups.getType().getDimSize(1);
if (replica_group_size == 0) {
op->emitOpError()
<< "Div final_op requires a non-empty replica_groups argument.";
}
auto divisor =
GetScalarConstOfType(element_type, loc, replica_group_size, &builder);
auto broadcast_dims = builder.getDenseI64ArrayAttr({});
result = chlo::BroadcastDivOp::create(builder, loc, all_reduce.getResult(0),
divisor.getResult(), broadcast_dims);
} else if (final_op != "Id") {
return op->emitOpError()
<< "invalid final_op " << final_op << ", want one of [Id, Div]";
}
op->replaceAllUsesWith(result);
op->erase();
return success();
}
template <typename T>
class CollectiveRewritePattern : public OpRewritePattern<T> {
public:
// Does not take any ownership. Caller must ensure channel_id is valid during
// life-cylce of this object.
CollectiveRewritePattern(MLIRContext* context, int64_t* channel_id)
: OpRewritePattern<T>(context), channel_id_(*channel_id) {}
protected:
int64_t& channel_id_; // A unique channel_id shared by all rewrite patterns
// in this pass. Not thread-safe.
};
// Converts XlaAllReduce. Not thread-safe.
class ConvertXlaAllReduce
: public CollectiveRewritePattern<TF::XlaAllReduceOp> {
public:
using CollectiveRewritePattern::CollectiveRewritePattern;
LogicalResult matchAndRewrite(TF::XlaAllReduceOp all_reduce,
PatternRewriter& rewriter) const override {
DenseIntElementsAttr replica_groups;
if (failed(ConvertReplicaGroups(rewriter, all_reduce.getGroupAssignment(),
replica_groups, all_reduce))) {
return failure();
}
// TODO(b/226201111): Stop emitting CollectiveInfo when it is no longer
// needed.
if (failed(SetCollectiveInfo(rewriter, replica_groups, all_reduce))) {
return failure();
}
StringRef reduce_op = all_reduce.getReduceOp();
StringRef merge_op, final_op;
if (reduce_op == "Add") {
merge_op = "Add";
final_op = "Id";
} else if (reduce_op == "Mul") {
merge_op = "Mul";
final_op = "Id";
} else if (reduce_op == "Min") {
merge_op = "Min";
final_op = "Id";
} else if (reduce_op == "Max") {
merge_op = "Max";
final_op = "Id";
} else if (reduce_op == "Mean") {
merge_op = "Add";
final_op = "Div";
} else {
return all_reduce->emitOpError()
<< "invalid reduce_op " << reduce_op
<< ", want one of [Add, Mul, Min, Max, Mean]";
}
int64_t channel_id = channel_id_++;
return ConvertAllReduce(rewriter, channel_id, all_reduce.getType(),
replica_groups, all_reduce.getMode(),
all_reduce.getInput(), merge_op, final_op,
all_reduce);
}
};
// Converts CollectiveReduceV2, with or without a preceding
// CollectiveAssignGroupV2. Not thread-safe.
class ConvertCollectiveReduceV2
: public CollectiveRewritePattern<TF::CollectiveReduceV2Op> {
public:
using CollectiveRewritePattern::CollectiveRewritePattern;
LogicalResult matchAndRewrite(TF::CollectiveReduceV2Op all_reduce,
PatternRewriter& rewriter) const override {
TF::CollectiveAssignGroupV2Op assign_group =
all_reduce.getGroupSize()
.getDefiningOp<TF::CollectiveAssignGroupV2Op>();
if (assign_group) {
// Found a group assignment. Use replica_groups to represent group
// assignment.
if (assign_group != all_reduce.getGroupKey()
.getDefiningOp<TF::CollectiveAssignGroupV2Op>()) {
return all_reduce->emitOpError()
<< "group_size and group_key are not from the "
"same CollectiveAssignGroupV2Op";
}
DenseIntElementsAttr replica_groups;
if (failed(ConvertReplicaGroups(rewriter,
assign_group.getGroupAssignment(),
replica_groups, all_reduce))) {
return failure();
}
// TODO(b/226201111): Stop emitting CollectiveInfo when it is no longer
// needed.
if (failed(SetCollectiveInfo(rewriter, replica_groups, all_reduce))) {
return failure();
}
int64_t channel_id = channel_id_++;
// FIXME(b/226139061): Mode should be set to CrossReplicaAndPartition
// in order to use XLA:GPU for more than one workers.
// The mode is set to use CrossReplica to keep the
// behavior on the primary user of this optimized path, because
// CrossReplicaAndPartition triggers a conflict with the channel_id
// allocation in the communication lowering, and the user uses both set of
// ops are used.
return ConvertAllReduce(rewriter, channel_id, all_reduce.getType(),
replica_groups, /* mode=*/"CrossReplica",
all_reduce.getInput(), all_reduce.getMergeOp(),
all_reduce.getFinalOp(), all_reduce);
}
// No group assignment, use separate channels per group_key.
DenseIntElementsAttr group_size_attr;
if (!matchPattern(all_reduce.getGroupSize(),
m_Constant(&group_size_attr))) {
return all_reduce.emitOpError()
<< "group_size must be a compile time constant";
}
if (!group_size_attr.isSplat() || group_size_attr.size() != 1) {
return all_reduce.emitOpError() << "group_size must be a scalar";
}
const auto group_size = group_size_attr.getSplatValue<IntegerAttr>();
// Create a full group assignment. Empty group assignment errors when
// final_op = "Div"
llvm::SmallVector<int64_t> indices(group_size.getInt());
std::iota(indices.begin(), indices.end(), 0);
auto replica_groups = mlir::DenseIntElementsAttr::get(
mlir::RankedTensorType::get({1, group_size.getInt()},
rewriter.getI64Type()),
indices);
{
// TODO(b/226201111): Stop emitting CollectiveInfo when it is no longer
// needed.
DenseIntElementsAttr group_key_attr;
if (!matchPattern(all_reduce.getGroupKey(),
m_Constant(&group_key_attr))) {
return all_reduce.emitOpError()
<< "group_key must be a compile time constant";
}
if (failed(SetCollectiveInfo(
/* group_size=*/group_size,
/* group_key=*/group_key_attr.getSplatValue<IntegerAttr>(),
all_reduce))) {
return failure();
}
}
// CrossReplicaAndPartition:
// Even though TF2XLA will setup the device assignment to include
// devices in this group as replicas before launching this module,
// "CrossReplica" mode (no channel) produces a deadlock when
// not using XLA SPMD expansion.
int64_t channel_id = channel_id_++;
return ConvertAllReduce(
rewriter, channel_id, all_reduce.getType(), replica_groups,
/* mode= */ "CrossReplicaAndPartition", all_reduce.getInput(),
all_reduce.getMergeOp(), all_reduce.getFinalOp(), all_reduce);
}
};
class ConvertCollectiveAssignGroupV2
: public CollectiveRewritePattern<TF::CollectiveAssignGroupV2Op> {
public:
using CollectiveRewritePattern::CollectiveRewritePattern;
LogicalResult matchAndRewrite(TF::CollectiveAssignGroupV2Op assign_group,
PatternRewriter& rewriter) const override {
DenseIntElementsAttr replica_groups;
if (failed(ConvertReplicaGroups(rewriter, assign_group.getGroupAssignment(),
replica_groups, assign_group))) {
return failure();
}
IntegerAttr group_size = rewriter.getI32IntegerAttr(replica_groups.size());
IntegerAttr group_key = rewriter.getI32IntegerAttr(0);
auto const_group_size =
TF::ConstOp::create(rewriter, assign_group->getLoc(),
assign_group.getResult(0).getType(), group_size);
auto const_group_key =
TF::ConstOp::create(rewriter, assign_group->getLoc(),
assign_group.getResult(1).getType(), group_key);
rewriter.replaceAllUsesWith(assign_group.getResult(0), const_group_size);
rewriter.replaceAllUsesWith(assign_group.getResult(1), const_group_key);
rewriter.eraseOp(assign_group);
return success();
}
};
void LegalizeTFCollective::runOnOperation() {
// FIXME(b/226139061): Figure out a way to share the channel_id with
// send/recv Ops. For now, start with a different range to avoid collision.
int64_t channel_id = 10000;
auto module = getOperation();
MLIRContext* context = module->getContext();
RewritePatternSet patterns(context);
patterns.insert<ConvertCollectiveAssignGroupV2>(context, &channel_id);
patterns.insert<ConvertCollectiveReduceV2>(context, &channel_id);
patterns.insert<ConvertXlaAllReduce>(context, &channel_id);
if (failed(applyPatternsGreedily(module, std::move(patterns)))) {
signalPassFailure();
}
}
} // namespace
std::unique_ptr<OperationPass<ModuleOp>> CreateLegalizeTFCollectivePass() {
return std::make_unique<LegalizeTFCollective>();
}
} // namespace mhlo
} // namespace mlir
tensorflow/compiler/mlir/tf2xla/transforms/legalize_tf_communication.cc
+994
View File
@@ -0,0 +1,994 @@
/* 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.
==============================================================================*/
// This file implements logic for lowering TensorFlow dialect's communication
// ops (TF/XLA) to the HLO dialect.
#include <atomic>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/Sequence.h"
#include "llvm/Support/ErrorHandling.h"
#include "llvm/Support/FormatVariadic.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/SparseTensor/IR/SparseTensor.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/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Location.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/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "xla/hlo/builder/sharding_builder.h"
#include "xla/hlo/translate/mhlo_to_hlo/type_to_shape.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
#include "xla/primitive_util.h"
#include "xla/side_effect_util.h"
#include "xla/xla_data.pb.h"
namespace mlir {
using func::FuncOp;
namespace mhlo {
namespace {
constexpr char kShardingAttr[] = "mhlo.sharding";
constexpr char kFrontendAttributesAttr[] = "mhlo.frontend_attributes";
// TPU core that sends to and receives from host.
constexpr int64_t kShardingTpuCore = 0;
#define GEN_PASS_DEF_LEGALIZETFCOMMUNICATIONPASS
#include "tensorflow/compiler/mlir/tf2xla/transforms/tf_xla_passes.h.inc"
// A pass that legalizes TF/XLA communication ops, propagate their respective
// tokens (for ordering), and rewrite their respective functions and control
// flow ops when necessary.
// Note, this currently does not handle nested modules/functions or region based
// ops other than certain control flow ops (`mhlo.if`, `mhlo.while`).
class LegalizeTFCommunication
: public impl::LegalizeTFCommunicationPassBase<LegalizeTFCommunication> {
void runOnOperation() override;
};
// A generator to serve out unique channel ids.
class ChannelIdGenerator {
public:
ChannelIdGenerator() = default;
ChannelIdGenerator(const ChannelIdGenerator&) = delete;
ChannelIdGenerator& operator=(const ChannelIdGenerator&) = delete;
ChannelIdGenerator(ChannelIdGenerator&&) = delete;
ChannelIdGenerator& operator=(ChannelIdGenerator&&) = delete;
int64_t operator++(int) { return next(); }
int64_t next() { return channel_id_.fetch_add(1, std::memory_order_relaxed); }
private:
// All usage code expects positive int64_t values so we can't use uint64_t
// and will just have to limit ourselves to half the number space.
std::atomic<int64_t> channel_id_ = 1;
};
int64_t GetNextChannelId() {
static ChannelIdGenerator* channel_id = new ChannelIdGenerator();
return channel_id->next();
}
// Checks if an op is a TF/XLA communication op.
bool IsCommunicationOp(Operation* op) {
return isa<TF::_XlaHostComputeMlirOp, TF::XlaSendToHostOp,
TF::XlaRecvFromHostOp>(op);
}
// Checks if an op is a supported HLO control flow op.
bool IsControlFlowOp(Operation* op) { return isa<IfOp, WhileOp>(op); }
// Collects control flow op ancestors of a given op, up until FuncOp. If any
// ancestor is not a control flow op or a FuncOp, or of a single block region,
// an error will be returned.
LogicalResult GetControlFlowAncestors(
Operation* op, llvm::SmallPtrSetImpl<Operation*>& control_flow_ops,
llvm::SmallPtrSetImpl<Block*>& control_flow_blocks) {
Block* block = op->getBlock();
Operation* parent = block->getParentOp();
while (block && parent && !isa<func::FuncOp>(parent)) {
if (!IsControlFlowOp(parent))
return op->emitOpError()
<< "expects ancestor(s) to be of ['" << IfOp::getOperationName()
<< "', '" << func::FuncOp::getOperationName() << "']";
if (!llvm::hasSingleElement(block->getParent()->getBlocks()))
return op->emitOpError() << "expects single block region ancestor(s)";
control_flow_ops.insert(parent);
control_flow_blocks.insert(block);
parent = block->getParentOp();
block = parent->getBlock();
}
return success();
}
// Finds communication ops in a function. `control_flow_ops` and
// `control_flow_blocks` will be populated with control flow op ancestors for
// every communication op.
LogicalResult FindCommunicationOps(
func::FuncOp func, llvm::SmallPtrSetImpl<Operation*>& control_flow_ops,
llvm::SmallPtrSetImpl<Block*>& control_flow_blocks,
bool& has_communication_ops) {
auto result = func.walk([&](Operation* op) {
if (!IsCommunicationOp(op)) return WalkResult::advance();
has_communication_ops = true;
if (failed(
GetControlFlowAncestors(op, control_flow_ops, control_flow_blocks)))
return WalkResult::interrupt();
return WalkResult::advance();
});
return failure(result.wasInterrupted());
}
// Helper struct holding a function to be rewritten, it's control flow ops that
// lead to a communication op or function call with a communication op
// (transitively), and an optional clone of itself. If `clone` is set, function
// calls to `original` will be replaced with `clone`.
struct FuncToRewrite {
func::FuncOp original;
llvm::SmallPtrSet<Operation*, 4> control_flow_ops;
llvm::SmallPtrSet<Block*, 4> control_flow_blocks;
func::FuncOp clone;
};
// Finds all functions that need to be rewritten with communication ops and
// and associated tokens.
LogicalResult GetFunctionsToRewrite(
ModuleOp module,
llvm::SmallDenseMap<StringRef, FuncToRewrite>& funcs_to_rewrite) {
// Find functions containing communication ops.
SmallVector<func::FuncOp, 4> funcs_to_visit;
for (func::FuncOp func : module.getOps<func::FuncOp>()) {
FuncToRewrite func_to_rewrite{/*original=*/func, /*control_flow_ops=*/{},
/*control_flow_blocks=*/{},
/*clone=*/nullptr};
bool has_communication_ops = false;
if (failed(FindCommunicationOps(func, func_to_rewrite.control_flow_ops,
func_to_rewrite.control_flow_blocks,
has_communication_ops)))
return failure();
if (!has_communication_ops) continue;
funcs_to_rewrite.insert({func.getName(), func_to_rewrite});
funcs_to_visit.push_back(func);
}
// Find functions that call functions with communication ops, transitively.
while (!funcs_to_visit.empty()) {
SmallVector<func::FuncOp, 4> new_funcs_to_visit;
for (func::FuncOp& func : funcs_to_visit) {
auto uses = func.getSymbolUses(module);
if (!uses) continue;
for (auto& use : *uses) {
// Only `mlir::func::CallOp` is supported as this requires knowing how
// to rewrite arguments and results to a function.
if (!isa<mlir::func::CallOp>(use.getUser())) continue;
auto caller_parent_func =
use.getUser()->getParentOfType<func::FuncOp>();
if (!caller_parent_func) continue;
FuncToRewrite func_to_rewrite{/*original=*/caller_parent_func,
/*control_flow_ops=*/{},
/*control_flow_blocks=*/{},
/*clone=*/nullptr};
if (failed(GetControlFlowAncestors(
use.getUser(), func_to_rewrite.control_flow_ops,
func_to_rewrite.control_flow_blocks)))
return failure();
auto it = funcs_to_rewrite.insert(
{caller_parent_func.getName(), func_to_rewrite});
if (it.second) {
new_funcs_to_visit.push_back(caller_parent_func);
} else {
it.first->getSecond().control_flow_ops.insert(
func_to_rewrite.control_flow_ops.begin(),
func_to_rewrite.control_flow_ops.end());
it.first->getSecond().control_flow_blocks.insert(
func_to_rewrite.control_flow_blocks.begin(),
func_to_rewrite.control_flow_blocks.end());
}
}
}
funcs_to_visit.swap(new_funcs_to_visit);
}
// Clone public functions that need to be rewritten. Function calls to this
// function will be replaced with the cloned function.
SymbolTable symbol_table(module);
for (auto& func : funcs_to_rewrite) {
if (func.getSecond().original.isPublic() &&
!func.getSecond().original.symbolKnownUseEmpty(module)) {
auto clone = func.getSecond().original.clone();
clone.setPrivate();
symbol_table.insert(clone);
func.getSecond().clone = clone;
}
}
return success();
}
// Assigns either MAXIMAL or MANUAL sharding. The MAXIMAL sharding sends/recvs
// one message from core `kShardingTpuCore` with the full tensor. MANUAL
// sharding sends/recvs one message for each core with the core's shard.
void SetOpSharding(Operation* op, bool manual_sharding) {
xla::OpSharding sharding =
manual_sharding ? ::xla::sharding_builder::Manual()
: ::xla::sharding_builder::SingleDevice(kShardingTpuCore);
op->setAttr(kShardingAttr,
StringAttr::get(op->getContext(), sharding.SerializeAsString()));
}
// Assigns frontend attributes holding information about data type and
// TensorFlow rendezvous channel name. The TensorFlow rendezvous channel name is
// handled differently as individual names are used per data send and receive.
void SetFrontendAttributes(Operation* op, int32_t index, StringRef key,
Type type, bool device_to_host,
StringRef host_handler_name) {
MLIRContext* context = op->getContext();
std::string formatted_key =
device_to_host ? llvm::formatv("{0}_dtoh_{1}", key, index).str()
: llvm::formatv("{0}_htod_{1}", key, index).str();
auto rendezvous_name = StringAttr::get(context, formatted_key);
auto rendezvous_name_attr = NamedAttribute(
StringAttr::get(context, xla::kXlaHostTransferRendezvousNameAttr),
rendezvous_name);
auto host_handler_name_value =
StringAttr::get(context, host_handler_name.str());
auto host_handler_name_attr = NamedAttribute(
StringAttr::get(context, xla::kXlaHostTransferHandlerNameAttr),
host_handler_name_value);
auto frontend_attributes = DictionaryAttr::get(
context,
ArrayRef<NamedAttribute>{rendezvous_name_attr, host_handler_name_attr});
op->setAttr(kFrontendAttributesAttr, frontend_attributes);
}
// Creates a `mhlo.send` op for sending value `operand`.
Value CreateSendOp(OpBuilder& builder, Location loc, Value operand,
StringRef key, size_t index, Value token,
StringRef host_handler_name, bool manual_sharding) {
// type 2 == DEVICE_TO_HOST
auto channel_handle = ChannelHandleAttr::get(builder.getContext(),
/*handle=*/GetNextChannelId(),
/*type=*/2);
auto empty_source_target_pairs = builder.getI64TensorAttr({});
auto send = SendOp::create(builder, loc, token.getType(), operand, token,
channel_handle,
/*is_host_transfer=*/builder.getBoolAttr(true),
/*source_target_pairs=*/empty_source_target_pairs);
SetFrontendAttributes(send, index, key, operand.getType(),
/*device_to_host=*/true, host_handler_name);
SetOpSharding(send, manual_sharding);
return send.getResult();
}
// Creates a `mhlo.recv` op for receiving a value.
Value CreateRecvOp(OpBuilder& builder, Location loc, Value result,
StringRef key, size_t index, Value token,
StringRef host_handler_name, bool manual_sharding) {
// type 3 == HOST_TO_DEVICE
auto channel_handle = ChannelHandleAttr::get(builder.getContext(),
/*handle=*/GetNextChannelId(),
/*type=*/3);
auto result_type = result.getType();
SmallVector<Type, 2> recv_result_type = {result_type, token.getType()};
auto recv =
RecvOp::create(builder, loc, recv_result_type, token, channel_handle,
/*is_host_transfer=*/builder.getBoolAttr(true),
/*source_target_pairs=*/builder.getI64TensorAttr({}));
SetFrontendAttributes(recv, index, key, result_type,
/*device_to_host=*/false, host_handler_name);
SetOpSharding(recv, manual_sharding);
result.replaceAllUsesWith(recv.getResult(0));
return recv.getResult(1);
}
// Creates a new token if necessary, acting as a sink to previous tokens. If
// there is only one token in `tokens`, the only token is returned. If `tokens`
// is empty, `original_token` is returned instead.
Value CreateSinkToken(OpBuilder& builder, Location loc, ArrayRef<Value> tokens,
Value original_token) {
if (tokens.empty()) {
return original_token;
} else if (llvm::hasSingleElement(tokens)) {
return tokens[0];
} else {
return AfterAllOp::create(builder, loc, original_token.getType(), tokens)
.getResult();
}
}
// Replaces `tf._XlaHostComputeMlir` with individual `mhlo.send` and `mhlo.recv`
// ops per operand and result. Unique Channel IDs are assigned per transfer.
// Sink tokens are created across all `mhlo.send` ops first and then by
// all `mhlo.recv` ops.
Value RewriteHostComputeOp(OpBuilder& builder,
TF::_XlaHostComputeMlirOp host_compute,
Value token) {
builder.setInsertionPoint(host_compute);
Location loc = host_compute.getLoc();
bool manual_sharding = host_compute.getManualSharding();
SmallVector<Value, 4> send_tokens;
for (auto operand : llvm::enumerate(host_compute.getInputs())) {
auto send_token = CreateSendOp(
builder, loc, operand.value(), host_compute.getSendKey(),
operand.index(), token, xla::kXlaHostTransferTfRendezvousHandlerName,
manual_sharding);
send_tokens.push_back(send_token);
}
token = CreateSinkToken(builder, loc, send_tokens, token);
SmallVector<Value, 4> recv_tokens;
for (auto result : llvm::enumerate(host_compute.getOutputs())) {
auto recv_token = CreateRecvOp(
builder, loc, result.value(), host_compute.getRecvKey(), result.index(),
token, xla::kXlaHostTransferTfRendezvousHandlerName, manual_sharding);
recv_tokens.push_back(recv_token);
}
token = CreateSinkToken(builder, loc, recv_tokens, token);
host_compute.erase();
return token;
}
// Replaces `tf.XlaSendToHost` with a `mhlo.send`.
Value RewriteSendToHostOp(OpBuilder& builder, TF::XlaSendToHostOp send_to_host,
Value token) {
builder.setInsertionPoint(send_to_host);
token = CreateSendOp(builder, send_to_host.getLoc(), send_to_host.getInput(),
send_to_host.getKey(),
/*index=*/0, token,
xla::kXlaHostTransferTfRendezvousHandlerName,
/*manual_sharding=*/false);
send_to_host.erase();
return token;
}
// Replaces `tf.XlaRecvFromHost` with a `mhlo.recv`.
Value RewriteRecvFromHostOp(OpBuilder& builder,
TF::XlaRecvFromHostOp recv_from_host, Value token) {
builder.setInsertionPoint(recv_from_host);
token = CreateRecvOp(builder, recv_from_host.getLoc(),
recv_from_host.getOutput(), recv_from_host.getKey(),
/*index=*/0, token,
xla::kXlaHostTransferTfRendezvousHandlerName,
/*manual_sharding=*/false);
recv_from_host.erase();
return token;
}
// Replaces a `mlir::func::CallOp` with one that has an extra `!mhlo.token`
// operand and `!mhlo.token` result. If `new_symbol` is set, the new call will
// be updated to call the `new_symbol` instead.
Value RewriteCallOp(OpBuilder& builder, func::CallOp call,
const std::optional<StringRef>& new_symbol, Value token) {
builder.setInsertionPoint(call);
auto new_operands = llvm::to_vector(call.getArgOperands());
new_operands.push_back(token);
auto new_result_types = llvm::to_vector(call.getResultTypes());
new_result_types.push_back(token.getType());
auto new_call = func::CallOp::create(
builder, call.getLoc(), new_result_types,
new_symbol ? *new_symbol : call.getCallee(), new_operands);
for (auto results : llvm::zip(call.getResults(), new_call.getResults()))
std::get<0>(results).replaceAllUsesWith(std::get<1>(results));
call.erase();
return new_call.getResults().back();
}
// Helper struct holding state of which op to visit to next. If `op` is in a
// control flow op region, `region_idx` will be set with the respective region
// index. `token` will be current token from the last communication op/control
// flow op transitive communication ops.
struct OpVisitorState {
std::optional<unsigned> region_idx;
Value token;
Operation* op;
};
// Creates a tuple from a sequence of values.
Value CreateTuple(OpBuilder& builder, Location loc, ArrayRef<Value> operands) {
return TupleOp::create(builder, loc, operands).getResult();
}
// Extends `values` with the value `token` attached. If `flatten_tuple` is
// false, `values` will have a single element, say `value`. If `value` is not a
// tuple, a new tuple is formed with `token`. If `values` is a tuple, it is
// extended instead. New tuple values created are cached.
SmallVector<Value> GetValueWithToken(
OpBuilder& builder, ArrayRef<Value> values, Value token,
llvm::SmallDenseMap<Value, Value>& rewritten_values, bool flatten_tuple) {
if (flatten_tuple) {
auto operands = llvm::to_vector(values);
operands.push_back(token);
return operands;
}
auto value = values[0];
// If value with token already exists, reuse it.
auto it = rewritten_values.find(value);
if (it != rewritten_values.end()) return {it->getSecond()};
auto create_tuple = [&](ArrayRef<Value> operands) {
auto new_result = CreateTuple(builder, value.getLoc(), operands);
rewritten_values.insert({value, new_result});
return new_result;
};
auto tuple_type = mlir::dyn_cast<TupleType>(value.getType());
// `value` is not a tuple, create a new tuple.
if (!tuple_type) return {create_tuple({value, token})};
// Extend tuple if `value` is a tuple.
// If `value` is an op result and the owner is a `mhlo.tuple`, simply unpack
// the tuple.
if (auto tuple_op = value.getDefiningOp<TupleOp>()) {
auto tuple_operands = llvm::to_vector(tuple_op.getOperands());
tuple_operands.push_back(token);
return {create_tuple(tuple_operands)};
}
// `value` is not created via a `mhlo.tuple` directly, unpack individual
// elements directly with `mhlo.get_tuple_element`.
SmallVector<Value, 4> tuple_operands;
for (auto idx : llvm::seq<int32_t>(0, tuple_type.getTypes().size()))
tuple_operands.push_back(
GetTupleElementOp::create(builder, value.getLoc(), value, idx)
.getResult());
tuple_operands.push_back(token);
return {create_tuple(tuple_operands)};
}
// Extends the 'types' to include a `mhlo.token` type. If `flatten_tuple` is
// false, `types` will have a single element, say `type`. If `type` is not a
// tuple type, a new tuple type with `type` and `mhlo.token` type is created
// instead.
SmallVector<Type> GetTypeWithToken(OpBuilder& builder, ArrayRef<Type> types,
bool flatten_tuple) {
SmallVector<Type> new_result_types;
auto token_type = TokenType::get(builder.getContext());
if (flatten_tuple) {
auto result_types = llvm::to_vector(types);
result_types.push_back(token_type);
return result_types;
}
auto type = types[0];
if (auto tuple_type = mlir::dyn_cast<TupleType>(type)) {
auto result_types = llvm::to_vector(tuple_type.getTypes());
result_types.push_back(token_type);
return {builder.getTupleType(result_types)};
}
return {builder.getTupleType({type, token_type})};
}
// Creates a slice of a tuple `value` with `mhlo.get_tuple_element` from index 0
// to `end`, exclusive.
Value CreateSubTuple(OpBuilder& builder, Value value, size_t end) {
SmallVector<Value, 4> tuple_operands;
for (auto idx : llvm::seq<int32_t>(0, end))
tuple_operands.push_back(
GetTupleElementOp::create(builder, value.getLoc(), value, idx)
.getResult());
return CreateTuple(builder, value.getLoc(), tuple_operands);
}
// Replaces uses of `values` with `replacements`. If `flatten_tuple` is false,
// `values` will have a single element, say `value`. If `value` is not a tuple
// type, an explicit `mhlo.get_tuple_element` is created to unpack the tuple and
// return the first element. Otherwise, `mhlo.get_tuple_element` users are
// simply updated with `replacement`, and all other users are updated with a
// slice of `replacement`.
void ReplaceWithTupleResult(OpBuilder& builder, ValueRange values,
ValueRange replacements, bool flatten_tuple) {
if (flatten_tuple) {
for (size_t result_index = 0; result_index < values.size(); result_index++)
values[result_index].replaceAllUsesWith(replacements[result_index]);
return;
}
auto value = values[0];
auto replacement = replacements[0];
auto tuple_type = mlir::dyn_cast<TupleType>(value.getType());
if (!tuple_type) {
if (!value.use_empty()) {
auto new_element = GetTupleElementOp::create(
builder, replacement.getLoc(), replacement, 0);
value.replaceAllUsesWith(new_element.getResult());
}
return;
}
Value sub_tuple;
for (auto& use : llvm::make_early_inc_range(value.getUses())) {
if (isa<GetTupleElementOp>(use.getOwner())) {
use.set(replacement);
continue;
}
if (!sub_tuple)
sub_tuple = CreateSubTuple(builder, replacement, tuple_type.size());
use.set(sub_tuple);
}
}
// Replaces control flow op block arguments with new block arguments
// of types `types`. The last element of the new block argument (token) is
// returned.
Value UpdateControlFlowBlockArgWithToken(OpBuilder& builder, Block& block,
ArrayRef<Type> types) {
builder.setInsertionPointToStart(&block);
auto old_args_size = block.getNumArguments();
block.addArguments(
types, SmallVector<Location>(types.size(), block.getParent()->getLoc()));
ValueRange old_args = block.getArguments().take_front(old_args_size);
ValueRange new_args = block.getArguments().drop_front(old_args_size);
assert(!new_args.empty());
ReplaceWithTupleResult(builder, old_args, new_args, /*flatten_tuple=*/true);
auto new_arg = new_args[new_args.size() - 1];
block.eraseArguments(0, old_args_size);
return new_arg;
}
// Updates control flow op terminator with an extra element `token`.
void RewriteControlFlowTerminator(OpBuilder& builder, Operation* terminator,
Value token, bool flatten_tuple) {
assert(flatten_tuple || terminator->getNumOperands() == 1);
assert(flatten_tuple || terminator->getBlock()->getNumArguments() == 1);
// `mhlo.while` cond terminator does not need to be rewritten as it always
// returns a tensor<i1> predicate value.
if (auto while_parent = dyn_cast_or_null<WhileOp>(terminator->getParentOp()))
if (terminator->getParentRegion() == &while_parent.getCond()) return;
builder.setInsertionPoint(terminator);
llvm::SmallDenseMap<Value, Value> rewritten_operands;
auto new_results =
GetValueWithToken(builder, llvm::to_vector(terminator->getOperands()),
token, rewritten_operands, flatten_tuple);
terminator->setOperands(new_results);
}
// Rewrites a `mhlo.if` op to receive and forward a `mhlo.token`. As If op does
// not have any operands other than the predicate, hence we implicitly capture
// the parent token. Also we use the same implicit token for use in the If op's
// regions.
void RewriteRegionIfOp(OpBuilder& builder, IfOp region_if,
SmallVectorImpl<OpVisitorState>& ops_to_visit,
Value token) {
llvm::SmallDenseMap<Value, Value> rewritten_operands;
auto new_result_types =
GetTypeWithToken(builder, llvm::to_vector(region_if.getResultTypes()),
/*flatten_tuple=*/true);
// Create new `mhlo.if` op with extra token operands and result.
auto new_if = IfOp::create(builder, region_if.getLoc(), new_result_types,
region_if.getPred());
// Move all regions from the old `mhlo.if` op to its replacement.
new_if.getTrueBranch().takeBody(region_if.getTrueBranch());
new_if.getFalseBranch().takeBody(region_if.getFalseBranch());
// Forward result from old `mhlo.if` with replacement.
SmallVector<Value> old_if_results = region_if.getResults();
SmallVector<Value> new_if_results = new_if.getResults();
ReplaceWithTupleResult(builder, old_if_results, new_if_results,
/*flatten_tuple=*/true);
// auto new_token = new_if_results[new_if_results.size() - 1];
region_if.erase();
// Next op to visit. The replacement is visited but at its first region.
// The new region use the same implicit token used by the If op.
ops_to_visit.push_back({/*region_idx=*/0, token, new_if});
}
// Rewrites a `mhlo.if`/`mhlo.while` region to receive and forward a
// `mhlo.token`. The block argument is updated to have an extra `mhlo.token`
// element. If the region block is to be rewritten, the next op to visit is set
// to the first op in the block. Otherwise the terminator is updated to forward
// `token`.
void RewriteControlFlowOpRegion(
OpBuilder& builder, Operation* region_op, unsigned region_idx,
ArrayRef<Type> block_arg_types,
SmallVectorImpl<OpVisitorState>& ops_to_visit,
const llvm::SmallPtrSetImpl<Block*>& control_flow_blocks, Value token) {
ops_to_visit.push_back({region_idx + 1, token, region_op});
Region& region = region_op->getRegion(region_idx);
assert(llvm::hasSingleElement(region));
auto block_token = UpdateControlFlowBlockArgWithToken(builder, region.front(),
block_arg_types);
if (control_flow_blocks.contains(&region.front())) {
ops_to_visit.push_back(
{/*region_idx=*/std::nullopt, block_token, &region.front().front()});
return;
}
RewriteControlFlowTerminator(builder, region.front().getTerminator(),
block_token, /*flatten_tuple=*/true);
}
// For mlir::IfOp or mlir::CaseOp, replace the use of their region's block
// argument (of type token) with 'implicit_operand'.
void ReplaceBlockArgumentsWithImplicitOperands(mlir::Operation* op,
unsigned region_idx,
Value implicit_operand) {
assert((mlir::dyn_cast<mlir::mhlo::IfOp>(*op) ||
mlir::dyn_cast<mlir::mhlo::CaseOp>(*op)) &&
"Unexpected mlir op in "
"HloFunctionImporter::ReplaceBlockArgumentsWithImplicitOperands!");
auto& region = op->getRegion(region_idx);
region.getArgument(0).replaceAllUsesWith(implicit_operand);
region.front().eraseArguments(0, region.getNumArguments());
}
// Rewrites an `mhlo.if` op or its region. If `region_idx` is not set, the op
// operands and results are rewritten. If `region_idx` is set, region
// `region_idx` is rewritten to take in and return an additional token. Returns
// true if the op or its region was rewritten.
bool ProcessRegionIfOp(OpBuilder& builder, IfOp region_if,
std::optional<unsigned> region_idx,
SmallVectorImpl<OpVisitorState>& ops_to_visit,
const llvm::SmallPtrSetImpl<Block*>& control_flow_blocks,
Value token) {
builder.setInsertionPoint(region_if);
if (!region_idx) {
RewriteRegionIfOp(builder, region_if, ops_to_visit, token);
return true;
}
if (*region_idx < region_if.getNumRegions()) {
// For the region-blocks of If op, we create a dummy token argument. Later
// we replace that block-argument's uses with the same (implicitly captured)
// token 'token', used for If op, and erase the argument.
// Note that 'RewriteControlFlowOpRegion' sets the token, used for the first
// operation of region_idx'th region, to the dummy block-argument. As we
// erase that argument, we also need to make sure that the token used for
// the next operation is set to 'token'.
RewriteControlFlowOpRegion(builder, region_if, *region_idx,
{token.getType()}, ops_to_visit,
control_flow_blocks, token);
ReplaceBlockArgumentsWithImplicitOperands(region_if.getOperation(),
*region_idx, token);
auto next_visitor_state = ops_to_visit.back();
next_visitor_state.token = token;
ops_to_visit.pop_back();
ops_to_visit.push_back(next_visitor_state);
return true;
}
return false;
}
// Rewrites a `mhlo.while` op to receive and forward a `mhlo.token`. Operands to
// the op for all of its regions are extended to have an extra operand `token`.
void RewriteRegionWhileOp(OpBuilder& builder, WhileOp region_while,
SmallVectorImpl<OpVisitorState>& ops_to_visit,
Value token) {
llvm::SmallDenseMap<Value, Value> rewritten_operands;
// Rewrite region operand to have an extra operand `token`.
auto new_val_operands =
GetValueWithToken(builder, llvm::to_vector(region_while.getOperands()),
token, rewritten_operands,
/*flatten_tuple=*/true);
auto new_result_types =
GetTypeWithToken(builder, llvm::to_vector(region_while.getResultTypes()),
/*flatten_tuple*/ true);
// Create new `mhlo.while` op with extra token operand and result.
auto new_while = WhileOp::create(builder, region_while.getLoc(),
new_result_types, new_val_operands);
// Move all regions from the old `mhlo.while` op to its replacement.
new_while.getCond().takeBody(region_while.getCond());
new_while.getBody().takeBody(region_while.getBody());
// Forward result from old `mhlo.while` with replacement.
SmallVector<Value> old_while_results = region_while.getResults();
SmallVector<Value> new_while_results = new_while.getResults();
ReplaceWithTupleResult(builder, old_while_results, new_while_results,
/*flatten_tuple*/ true);
auto new_token = new_while_results[new_while_results.size() - 1];
region_while.erase();
// Next op to visit. The replacement is visited but at its first region. The
// token result of the new region if is propagated.
ops_to_visit.push_back({/*region_idx=*/0, new_token, new_while});
}
// Rewrites an `mhlo.while` op or its region. If `region_idx` is not set, the op
// operands and results are rewritten. If `region_idx` is set, region
// `region_idx` is rewritten to take in and return an additional token. Returns
// true if the op or its region was rewritten.
bool ProcessRegionWhileOp(
OpBuilder& builder, WhileOp region_while,
std::optional<unsigned> region_idx,
SmallVectorImpl<OpVisitorState>& ops_to_visit,
const llvm::SmallPtrSetImpl<Block*>& control_flow_blocks, Value token) {
builder.setInsertionPoint(region_while);
if (!region_idx) {
RewriteRegionWhileOp(builder, region_while, ops_to_visit, token);
return true;
}
if (*region_idx < region_while.getNumRegions()) {
SmallVector<Type> operand_types;
for (auto operand : region_while.getOperand())
operand_types.push_back(operand.getType());
RewriteControlFlowOpRegion(builder, region_while, *region_idx,
operand_types, ops_to_visit, control_flow_blocks,
token);
return true;
}
return false;
}
// Updates function type based on current function body block arguments and
// terminator operand types.
void UpdateFunctionType(OpBuilder& builder, func::FuncOp func,
Block& func_body) {
auto new_argument_types = llvm::to_vector(func_body.getArgumentTypes());
auto new_result_types =
llvm::to_vector(func_body.getTerminator()->getOperandTypes());
func.setType(FunctionType::get(builder.getContext(), new_argument_types,
new_result_types));
}
// Replaces a function terminator `return` with another `return` that has an
// extra `mhlo.token` operand.
void RewriteFunctionTerminator(OpBuilder& builder,
mlir::func::ReturnOp terminator, Value token) {
auto new_results = llvm::to_vector(terminator.getOperands());
new_results.push_back(token);
builder.setInsertionPoint(terminator);
mlir::func::ReturnOp::create(builder, terminator.getLoc(), new_results);
terminator.erase();
}
// Rewrites a function body and communication ops inside. Region control flow
// are updated when necessary, to propagate tokens. The function may either be
// rewritten to create a token or take in and return a token, depending on its
// visibility and if there are any callers.
LogicalResult RewriteFunction(
OpBuilder& builder, ModuleOp module, FuncOp func,
const llvm::SmallDenseMap<StringRef, FuncToRewrite>& funcs,
const llvm::SmallPtrSetImpl<Operation*>& control_flow_ops,
const llvm::SmallPtrSetImpl<Block*>& control_flow_blocks, bool is_clone) {
MLIRContext* context = module.getContext();
if (!llvm::hasSingleElement(func.getBody()))
return func.emitError()
<< "'" << FuncOp::getOperationName()
<< "' ops with more than one block are not supported";
bool rewrite_block =
is_clone || (!func.isPublic() && !func.symbolKnownUseEmpty(module));
Block& func_body = func.front();
builder.setInsertionPointToStart(&func_body);
auto token_type = TokenType::get(context);
// If a function is public, it's signature should not be modified, and instead
// a token will be created. Otherwise a token block argument is inserted.
Value init_token =
rewrite_block ? func_body.addArgument(token_type, func.getLoc())
: CreateTokenOp::create(builder, func.getLoc(), token_type)
.getResult();
// Stack to keep track of region based control flow op nesting and current
// op to visit.
SmallVector<OpVisitorState, 4> ops_to_visit{
{/*region_idx=*/std::nullopt, init_token, &func_body.front()}};
while (!ops_to_visit.empty()) {
OpVisitorState op_to_visit = ops_to_visit.pop_back_val();
Operation* curr_op = op_to_visit.op;
Value token = op_to_visit.token;
// Ops may be removed, so the next op is kept track of beforehand.
Operation* next_op = curr_op->getNextNode();
if (auto host_compute = dyn_cast<TF::_XlaHostComputeMlirOp>(curr_op)) {
token = RewriteHostComputeOp(builder, host_compute, token);
} else if (auto send_to_host = dyn_cast<TF::XlaSendToHostOp>(curr_op)) {
token = RewriteSendToHostOp(builder, send_to_host, token);
} else if (auto recv_from_host = dyn_cast<TF::XlaRecvFromHostOp>(curr_op)) {
token = RewriteRecvFromHostOp(builder, recv_from_host, token);
} else if (auto call = dyn_cast<mlir::func::CallOp>(curr_op)) {
// Only `mlir::func::CallOp` is supported as this requires knowing how to
// rewrite arguments and results to a function.
auto it = funcs.find(call.getCallee());
if (it != funcs.end()) {
func::FuncOp clone = it->getSecond().clone;
std::optional<StringRef> symbol_name =
clone ? std::optional<StringRef>(clone.getName()) : std::nullopt;
// If the function being called is to be cloned, update the call to also
// point to the cloned function.
token = RewriteCallOp(builder, call, symbol_name, token);
}
} else if (auto region_if = dyn_cast<IfOp>(curr_op)) {
if (op_to_visit.region_idx || control_flow_ops.contains(region_if)) {
auto exist_unprocessed_region =
ProcessRegionIfOp(builder, region_if, op_to_visit.region_idx,
ops_to_visit, control_flow_blocks, token);
// Once all the IfOp regions are processed (i.e.
// 'exist_unprocessed_region' == false), select returned token-value
// from IfOp as the token to be used for the following op.
if (!exist_unprocessed_region) {
token = curr_op->getResult(curr_op->getNumResults() - 1);
} else {
continue;
}
}
} else if (auto region_while = dyn_cast<WhileOp>(curr_op)) {
if (op_to_visit.region_idx || control_flow_ops.contains(region_while))
if (ProcessRegionWhileOp(builder, region_while, op_to_visit.region_idx,
ops_to_visit, control_flow_blocks, token))
continue;
} else if (auto region_terminator = dyn_cast<mhlo::ReturnOp>(curr_op)) {
bool flatten_tuple = isa<mhlo::WhileOp, mhlo::IfOp, mhlo::CaseOp>(
region_terminator->getParentOp());
RewriteControlFlowTerminator(builder, region_terminator, token,
flatten_tuple);
// There is no next op after the control flow op terminator, simply let
// stack have one less element.
continue;
} else if (auto func_terminator = dyn_cast<mlir::func::ReturnOp>(curr_op)) {
if (rewrite_block)
RewriteFunctionTerminator(builder, func_terminator, token);
// There is no next op after the function terminator, simply let stack
// have one less element/be empty.
continue;
}
// Visit next op.
ops_to_visit.push_back({/*region_idx=*/std::nullopt, token, next_op});
}
if (rewrite_block) UpdateFunctionType(builder, func, func_body);
return success();
}
// Checks if a function call is pointing to a function with communication ops.
bool IsFunctionCallWithCommunication(
Operation* op,
const llvm::SmallDenseMap<StringRef, FuncToRewrite>& funcs_to_rewrite) {
if (auto call = dyn_cast<mlir::func::CallOp>(op))
return funcs_to_rewrite.count(call.getCallee());
return false;
}
// Collects all control flow op ancestors of communication ops or function calls
// with communication ops (transitively).
void GetCommunicationControlFlowOps(
func::FuncOp func,
const llvm::SmallDenseMap<StringRef, FuncToRewrite>& funcs_to_rewrite,
llvm::SmallPtrSetImpl<Operation*>& control_flow_ops,
llvm::SmallPtrSetImpl<Block*>& control_flow_blocks) {
func.walk([&](Operation* op) {
if (IsCommunicationOp(op) ||
IsFunctionCallWithCommunication(op, funcs_to_rewrite))
if (failed(GetControlFlowAncestors(op, control_flow_ops,
control_flow_blocks)))
llvm_unreachable(
"checking original function for control flow ancestors should have "
"errored first");
});
}
void LegalizeTFCommunication::runOnOperation() {
auto module = getOperation();
llvm::SmallDenseMap<StringRef, FuncToRewrite> funcs_to_rewrite;
if (failed(GetFunctionsToRewrite(module, funcs_to_rewrite)))
return signalPassFailure();
OpBuilder builder(&getContext());
for (const auto& func_and_name : funcs_to_rewrite) {
const auto& func_to_rewrite = func_and_name.getSecond();
func::FuncOp func = func_to_rewrite.original;
if (failed(RewriteFunction(builder, module, func, funcs_to_rewrite,
func_to_rewrite.control_flow_ops,
func_to_rewrite.control_flow_blocks,
/*is_clone=*/false)))
return signalPassFailure();
func::FuncOp clone = func_and_name.getSecond().clone;
if (!clone) continue;
llvm::SmallPtrSet<Operation*, 4> clone_control_flow_ops;
llvm::SmallPtrSet<Block*, 4> clone_control_flow_blocks;
GetCommunicationControlFlowOps(clone, funcs_to_rewrite,
clone_control_flow_ops,
clone_control_flow_blocks);
if (failed(RewriteFunction(builder, module, clone, funcs_to_rewrite,
clone_control_flow_ops,
clone_control_flow_blocks,
/*is_clone=*/true)))
llvm_unreachable(
"rewriting of original function should have errored first");
}
}
} // namespace
std::unique_ptr<OperationPass<ModuleOp>> CreateLegalizeTFCommunicationPass() {
return std::make_unique<LegalizeTFCommunication>();
}
} // namespace mhlo
} // namespace mlir
tensorflow/compiler/mlir/tf2xla/transforms/legalize_tf_patterns.td
+827
View File
@@ -0,0 +1,827 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
// This is the legalization pattern definition file for TF to XLA.
include "mlir/IR/OpBase.td"
include "mlir/Dialect/Shape/IR/ShapeOps.td"
include "mlir/Dialect/Func/IR/FuncOps.td"
include "mlir/Dialect/Tensor/IR/TensorOps.td"
include "stablehlo/dialect/ChloOps.td"
include "stablehlo/dialect/StablehloOps.td"
include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
include "mhlo/IR/hlo_ops.td" // for hlo_utils.td
def SignedIntTensor : TensorOf<[I1, I8, I16, I32, I64]>;
def UnsignedIntTensor : TensorOf<[UI8, UI16, UI32, UI64]>;
// IEEE compliant floating point tensors.
def IEEEFloatTensor : TensorOf<[F16, F32, F64]>;
//===----------------------------------------------------------------------===//
// BatchNorm op patterns.
//===----------------------------------------------------------------------===//
def FalseBoolAttr : AttrConstraint<CPred<"!llvm::cast<BoolAttr>($_self).getValue()">>;
def TrueBoolAttr : AttrConstraint<CPred<"llvm::cast<BoolAttr>($_self).getValue()">>;
def CastValueToI64: NativeCodeCall<
"CastValueToI64($0.getLoc(), $1, &$_builder)">;
def CastValueToElementType: NativeCodeCall<
"stablehlo::ConvertOp::create($_builder, $0.getLoc(), $1, "
"getElementTypeOrSelf($2.getType()))">;
// Here, $0 is an ElementsAttr with exactly one element of type integer. $1 is
// the corresponding value of ranked tensor type whose axis is referred in $0.
def GetHLOAxisFromTFAxis : NativeCodeCall<
"GetHLOAxisFromTFAxis("
"$0, llvm::cast<RankedTensorType>($1.getType()).getRank(), &$_builder)">;
// Same as the above but with $1 of type operand_range from variadic TensorFlow
// input.
def GetHLOAxisFromTFAxisVariadic : NativeCodeCall<
"GetHLOAxisFromTFAxis("
"$0, llvm::cast<RankedTensorType>((*$1.begin()).getType()).getRank(), "
"&$_builder)">;
def CastElementsToI64Elements : NativeCodeCall<[{
llvm::cast<mlir::DenseIntElementsAttr>(hlo::convertElementsAttr(
llvm::cast<ElementsAttr>($0), $_builder.getIntegerType(64)))
}]>;
def CastElementsToI64Array : NativeCodeCall<[{
ToDenseI64ArrayAttr(
llvm::cast<mlir::DenseIntElementsAttr>(hlo::convertElementsAttr(
llvm::cast<ElementsAttr>($0), $_builder.getIntegerType(64))), &$_builder)
}]>;
def EmptyDotAlgorithmAttr : NativeCodeCall<"mlir::stablehlo::DotAlgorithmAttr{}">;
def ConstDefaultResultAccuracyAttr :
ConstantAttr<StableHLO_ResultAccuracyAttr, "::mlir::stablehlo::ResultAccuracyMode::DEFAULT">;
//===----------------------------------------------------------------------===//
// ApproximateEqual op pattern.
//===----------------------------------------------------------------------===//
class StableHLO_ComparisonDirectionValue<string enumStr> :
ConstantAttr<StableHLO_ComparisonDirectionAttr, "::mlir::stablehlo::ComparisonDirection::" # enumStr>;
class CHLO_ComparisonDirectionValue<string enumStr> :
ConstantAttr<CHLO_ComparisonDirectionAttr, "::mlir::chlo::ComparisonDirection::" # enumStr>;
// TODO(b/228291745): Assert that $x and $y have the same shape.
def : Pat<(TF_ApproximateEqualOp:$result $x, $y, $tolerance),
(CHLO_BroadcastCompareOp
(StableHLO_AbsOp:$abs (StableHLO_SubtractOp $x, $y)),
(CastValueToElementType $result, (StableHLO_ConstantOp $tolerance), $abs),
(NullDenseI64ArrayAttr),
CHLO_ComparisonDirectionValue<"LT">,
(CHLO_DEFAULT_COMPARISON_TYPE))>;
//===----------------------------------------------------------------------===//
// Assert op pattern.
//===----------------------------------------------------------------------===//
// HLO and XLA doesn't support Assertions.
def LowerAssert : Pattern<(TF_AssertOp $condition, $data, $summarize), []>;
//===----------------------------------------------------------------------===//
// Binary op patterns.
//===----------------------------------------------------------------------===//
// Check that two values can be broadcasted together
def AreBroadcastCompatible : Constraint<CPred<"AreBroadcastCompatible($0, $1)">,
"types must be broadcastable">;
class DirectBinaryPat<Op FromOp, Op ToOp>
: Pat<(FromOp AnyTensor:$l, AnyTensor:$r),
(ToOp $l, $r, (BinBroadcastDimensions $l, $r))>;
foreach fromToBinPair = [[TF_AddV2Op, CHLO_BroadcastAddOp],
[TF_Atan2Op, CHLO_BroadcastAtan2Op],
[TF_ComplexOp, CHLO_BroadcastComplexOp],
[TF_DivOp, CHLO_BroadcastDivOp],
[TF_LeftShiftOp, CHLO_BroadcastShiftLeftOp],
[TF_MaximumOp, CHLO_BroadcastMaxOp],
[TF_MinimumOp, CHLO_BroadcastMinOp],
[TF_ModOp, CHLO_BroadcastRemOp],
[TF_MulOp, CHLO_BroadcastMulOp],
[TF_NextAfterOp, CHLO_BroadcastNextAfterOp],
[TF_PolygammaOp, CHLO_BroadcastPolygammaOp],
[TF_PowOp, CHLO_BroadcastPowOp],
[TF_RealDivOp, CHLO_BroadcastDivOp],
[TF_SubOp, CHLO_BroadcastSubOp],
[TF_ZetaOp, CHLO_BroadcastZetaOp]] in
def : DirectBinaryPat<fromToBinPair[0], fromToBinPair[1]>;
def LowerRightShiftSigned :
Pat<(TF_RightShiftOp AnyTensor:$l, AnyTensor:$r),
(CHLO_BroadcastShiftRightArithmeticOp $l, $r,
(BinBroadcastDimensions $l, $r)),
[(SignedIntTensor $r)]>;
def LowerRightShiftUnsigned :
Pat<(TF_RightShiftOp AnyTensor:$l, AnyTensor:$r),
(CHLO_BroadcastShiftRightLogicalOp $l, $r,
(BinBroadcastDimensions $l, $r)),
[(UnsignedIntTensor $r)]>;
// Performs a substitution of FloorDiv, pseudo code below:
//
// return floor(div(x, y))
def : Pat<(TF_FloorDivOp AnyTensor:$l, AnyTensor:$r),
(StableHLO_FloorOp
(CHLO_BroadcastDivOp $l, $r, (BinBroadcastDimensions $l, $r))),
[(IEEEFloatTensor $l)]>;
// Performs a substitution of FloorDiv for integer tensors, which required
// additional correction for a negative numerator / denominator. Equivalent
// pseudocode is shown below:
//
// T z = x / y
// return (z * y != x && (x < 0) != (y < 0)) ? z - 1 : z
//
// BroadcastToDimensions is used to compute the broadcast attr to higher
// dimensions. This computes the broadcast of 'l' to broadcast('l', 'r')
// without returning the broadcast of 'r' to broadcast('l', 'r').
def : Pat<(TF_FloorDivOp AnyTensor:$l, AnyTensor:$r),
(StableHLO_SelectOp
(CHLO_BroadcastAndOp
(CHLO_BroadcastCompareOp
(CHLO_BroadcastMulOp:$mul
(CHLO_BroadcastDivOp:$div $l, $r,
(BinBroadcastDimensions $l, $r)),
$r, (BinBroadcastDimensions $div, $r)),
$l, (BinBroadcastDimensions $mul, $l), CHLO_ComparisonDirectionValue<"NE">,
(CHLO_DEFAULT_COMPARISON_TYPE)),
(CHLO_BroadcastCompareOp
(CHLO_BroadcastCompareOp:$l_cmp $l,
(StableHLO_ConstantOp:$l_zeros (GetScalarOfType<0> $l)),
(NullDenseI64ArrayAttr), CHLO_ComparisonDirectionValue<"LT">,
(CHLO_DEFAULT_COMPARISON_TYPE)),
(CHLO_BroadcastCompareOp:$r_cmp $r,
(StableHLO_ConstantOp:$r_zeros (GetScalarOfType<0> $r)),
(NullDenseI64ArrayAttr), CHLO_ComparisonDirectionValue<"LT">,
(CHLO_DEFAULT_COMPARISON_TYPE)),
(BinBroadcastDimensions $l_cmp, $r_cmp), CHLO_ComparisonDirectionValue<"NE">,
(CHLO_DEFAULT_COMPARISON_TYPE)),
(NullDenseI64ArrayAttr)),
(CHLO_BroadcastSubOp $div,
(StableHLO_ConstantOp:$ones (GetScalarOfType<1> $div)),
(NullDenseI64ArrayAttr)), $div),
[(SignedIntTensor $l)]>;
// FloorDiv of unsigned is just div.
def : Pat<(TF_FloorDivOp AnyTensor:$l, AnyTensor:$r),
(CHLO_BroadcastDivOp $l, $r, (BinBroadcastDimensions $l, $r)),
[(UnsignedIntTensor $l)]>;
// Performs a substitution of FloorMod designed to correct for possibly negative
// values. Pseudocode shown below:
//
// T trunc_mod = std::fmod(x, y);
// return trunc_mod != 0 && (y < 0 != trunc_mod < 0) ? trunc_mod + y
// : trunc_mod
def : Pat<(TF_FloorModOp AnyTensor:$l, AnyTensor:$r),
(StableHLO_SelectOp
(CHLO_BroadcastAndOp
(CHLO_BroadcastCompareOp
(CHLO_BroadcastRemOp:$rem $l, $r, (BinBroadcastDimensions $l, $r)),
(StableHLO_ConstantOp:$l_zeros (GetScalarOfType<0> $l)),
(NullDenseI64ArrayAttr), CHLO_ComparisonDirectionValue<"NE">,
(CHLO_DEFAULT_COMPARISON_TYPE)),
(CHLO_BroadcastCompareOp
(CHLO_BroadcastCompareOp:$r_cmp $r,
(StableHLO_ConstantOp:$r_zeros (GetScalarOfType<0> $r)),
(NullDenseI64ArrayAttr), CHLO_ComparisonDirectionValue<"LT">,
(CHLO_DEFAULT_COMPARISON_TYPE)),
(CHLO_BroadcastCompareOp:$rem_cmp $rem, $r_zeros,
(BinBroadcastDimensions $rem, $r_zeros), CHLO_ComparisonDirectionValue<"LT">,
(CHLO_DEFAULT_COMPARISON_TYPE)),
(BinBroadcastDimensions $r_cmp, $rem_cmp), CHLO_ComparisonDirectionValue<"NE">,
(CHLO_DEFAULT_COMPARISON_TYPE)),
(NullDenseI64ArrayAttr)),
(CHLO_BroadcastAddOp $r,
$rem, (BinBroadcastDimensions $r, $rem)), $rem),
[(TensorOf<[I8, I16, I32, I64, F16, F32, F64]> $l)]>;
// FloorMod of unsigned is just mod.
def : Pat<(TF_FloorModOp AnyTensor:$l, AnyTensor:$r),
(CHLO_BroadcastRemOp $l, $r, (BinBroadcastDimensions $l, $r)),
[(UnsignedIntTensor $l)]>;
def Get2DTransposePerm: NativeCodeCall<
"Get2DTransposePerm($0, &$_builder)">;
def : Pat<(TF_RiscAddOp $l, $r), (StableHLO_AddOp $l, $r)>;
def : Pat<(TF_RiscDotOp $a, $b, $transpose_a, $transpose_b),
(StableHLO_DotOp
(TF_TransposeOp $a, (TF_ConstOp (Get2DTransposePerm $transpose_a))),
(TF_TransposeOp $b, (TF_ConstOp (Get2DTransposePerm $transpose_b))),
/*precision_config=*/(NullArrayAttr))>;
//===----------------------------------------------------------------------===//
// Logical & bitwise binary op patterns.
//===----------------------------------------------------------------------===//
class DirectLogicalBinaryPat<Op FromOp, Op ToOp>
: Pat<(FromOp AnyTensor:$l, AnyTensor:$r),
(ToOp $l, $r, (BinBroadcastDimensions $l, $r)),
[(AnyTypeOf<[SignedIntTensor, UnsignedIntTensor]> $l)]>;
foreach fromToBinPair = [[TF_LogicalAndOp, CHLO_BroadcastAndOp],
[TF_LogicalOrOp, CHLO_BroadcastOrOp],
[TF_BitwiseAndOp, CHLO_BroadcastAndOp],
[TF_BitwiseOrOp, CHLO_BroadcastOrOp],
[TF_BitwiseXorOp, CHLO_BroadcastXorOp]] in
def : DirectLogicalBinaryPat<fromToBinPair[0], fromToBinPair[1]>;
//===----------------------------------------------------------------------===//
// Compare op patterns.
//===----------------------------------------------------------------------===//
class DirectComparePat<Op FromOp, CHLO_ComparisonDirectionValue direction>
: Pat<(FromOp AnyTensor:$l, AnyTensor:$r),
(CHLO_BroadcastCompareOp
$l, $r, (BinBroadcastDimensions $l, $r), direction,
(CHLO_DEFAULT_COMPARISON_TYPE))>;
def : DirectComparePat<TF_GreaterOp, CHLO_ComparisonDirectionValue<"GT">>;
def : DirectComparePat<TF_GreaterEqualOp, CHLO_ComparisonDirectionValue<"GE">>;
def : DirectComparePat<TF_LessOp, CHLO_ComparisonDirectionValue<"LT">>;
def : DirectComparePat<TF_LessEqualOp, CHLO_ComparisonDirectionValue<"LE">>;
class EqualityPat<Op FromOp, CHLO_ComparisonDirectionValue direction>
: Pat<(FromOp AnyTensor:$l, AnyTensor:$r,
TrueBoolAttr:$incompatible_shape_error),
(CHLO_BroadcastCompareOp
$l, $r, (BinBroadcastDimensions $l, $r), direction,
(CHLO_DEFAULT_COMPARISON_TYPE)),
[(HLO_Tensor $l)]>;
def : EqualityPat<TF_EqualOp, CHLO_ComparisonDirectionValue<"EQ">>;
def : EqualityPat<TF_NotEqualOp, CHLO_ComparisonDirectionValue<"NE">>;
//===----------------------------------------------------------------------===//
// Concat op patterns.
//===----------------------------------------------------------------------===//
def OneElementAttrPred
: CPred<"llvm::cast<ElementsAttr>($_self).getShapedType().getNumElements() == 1">;
def OneElementAttr
: ElementsAttrBase<And<[ElementsAttr.predicate, OneElementAttrPred]>,
"Scalar ElementsAttr">;
def HasRankedFirstOperand
: Constraint<CPred<"llvm::isa<RankedTensorType>((*$0.begin()).getType())">>;
def IsShapedTensor
: Constraint<CPred<"llvm::isa<RankedTensorType>($0.getType())">>;
// This pattern converts TensorFlow axis format to HLO axis format which
// doesn't wrap around like TensorFlow and is always positive. For this
// conversion, use the first input to get inputs rank. Other inputs need not be
// ranked.
// Defining op for `axis` is TensorFlow constant op in the pattern as during
// the conversion, original Concat op operands still refers to the old ops even
// if HLO constant op is introduced as an replacement for the TensorFlow
// Constant op.
def : Pat<(TF_ConcatV2Op $inputs, (ConstantLikeMatcher OneElementAttr:$axis)),
(StableHLO_ConcatenateOp $inputs,
(GetHLOAxisFromTFAxisVariadic $axis, $inputs)),
[(HasRankedFirstOperand $inputs)]>;
//===----------------------------------------------------------------------===//
// CollectivePermute op patterns.
//===----------------------------------------------------------------------===//
def : Pat<(TF_CollectivePermuteOp $input, (ConstantLikeMatcher ElementsAttr:$source_target_pairs)),
(StableHLO_CollectivePermuteOp $input,
(CastElementsToI64Elements $source_target_pairs),
(StableHLO_NullChannelHandleAttr))>;
//===----------------------------------------------------------------------===//
// CrossReplicaSum op patterns.
//===----------------------------------------------------------------------===//
def : Pat<(TF_CrossReplicaSumOp $input, (ConstantLikeMatcher ElementsAttr:$group_assignment)),
(StableHLO_CrossReplicaSumOp $input,
(CastElementsToI64Elements $group_assignment))>;
//===----------------------------------------------------------------------===//
// All2All op patterns.
//===----------------------------------------------------------------------===//
def ValueToVariadic: NativeCodeCall<"SmallVector<Value, 1>{$0}">;
def : Pat<(TF_AllToAllOp AnyRankedTensor:$input, (ConstantLikeMatcher ElementsAttr:$group_assignment), I64Attr:$concat_dimension, $split_dimension, $split_count),
(StableHLO_AllToAllOp (ValueToVariadic $input), $split_dimension, $concat_dimension, $split_count, (CastElementsToI64Elements $group_assignment), (StableHLO_NullChannelHandleAttr))>;
//===----------------------------------------------------------------------===//
// FFT op patterns.
//===----------------------------------------------------------------------===//
class StableHLO_FftTypeValue<string enumStr> :
ConstantAttr<StableHLO_FftTypeAttr, "::mlir::stablehlo::FftType::" # enumStr>;
def GetInnerDimFromValue : NativeCodeCall<
"GetInnerDimFromValue(llvm::cast<ShapedType>($0.getType()), &$_builder)">;
def CheckInnerDimStatic
: Constraint<CPred<"CheckInnerDimStatic(llvm::cast<ShapedType>($0.getType()), &$_builder)">>;
def : Pat<(TF_FFTOp:$res $input),
(StableHLO_FftOp $input, StableHLO_FftTypeValue<"FFT">, (GetInnerDimFromValue $res)),
[(CheckInnerDimStatic $input)]>;
def : Pat<(TF_IFFTOp:$res $input),
(StableHLO_FftOp $input, StableHLO_FftTypeValue<"IFFT">, (GetInnerDimFromValue $res)),
[(CheckInnerDimStatic $input)]>;
//===----------------------------------------------------------------------===//
// GatherV2 op patterns.
//===----------------------------------------------------------------------===//
// Here, $params and $indices needs to be ranked so that $axis and $batch_dims
// attributes can be converted from TensorFlow axis format supporting negative
// indexing to the HLO format.
def LegalizeGatherV2 :
Pat<(TF_GatherV2Op AnyRankedTensor:$params, AnyRankedTensor:$indices,
(ConstantLikeMatcher ElementsAttr:$axis), $batch_dims),
(StableHLO_TorchIndexSelectOp $params, $indices,
(GetHLOAxisFromTFAxis $axis, $params),
(GetHLOAxisFromTFAxis $batch_dims, $indices))>;
//===----------------------------------------------------------------------===//
// Pad op patterns.
//===----------------------------------------------------------------------===//
class SliceDenseIntElementsAttrColumn2D<string column> : NativeCodeCall<
"SliceDenseIntElementsAttrColumn2D(llvm::cast<ElementsAttr>($0), " # column # " )">;
class SliceDenseIntElementsAttr<string index, string axis> : NativeCodeCall<
"SliceDenseIntElementsAttr(llvm::cast<ElementsAttr>($0), " # index # ", " # axis # ")">;
// Interior padding attribute based on the TF padding.
def GetInteriorPadding : NativeCodeCall<
"GetInteriorPadding(llvm::cast<ElementsAttr>($0))">;
def : Pat<(TF_PadV2Op $input, (ConstantLikeMatcher ElementsAttr:$padding), $c),
(StableHLO_PadOp $input, $c,
(SliceDenseIntElementsAttrColumn2D<"0"> $padding),
(SliceDenseIntElementsAttrColumn2D<"1"> $padding),
(GetInteriorPadding $padding))>;
//===----------------------------------------------------------------------===//
// Identity op patterns.
//===----------------------------------------------------------------------===//
foreach src = [TF_IdentityOp, TF_StopGradientOp, TF__EagerConstOp] in
def : Pat<(src $op), (replaceWithValue $op)>;
// TODO(b/32223192): Support CheckNumerics in HLO.
foreach src = [TF_PreventGradientOp, TF_CheckNumericsOp] in
def : Pat<(src $op, $msg), (replaceWithValue $op)>;
//===----------------------------------------------------------------------===//
// MatMul op patterns.
//===----------------------------------------------------------------------===//
def StableHLO_GetPrecisionConfig: NativeCodeCall<
"GetPrecisionConfig(&$_builder)">;
def : Pat<(TF_MatMulOp $a, $b, $transpose_a, $transpose_b, $grad_a, $grad_b),
(StableHLO_DotOp
(TF_TransposeOp $a, (TF_ConstOp (Get2DTransposePerm $transpose_a))),
(TF_TransposeOp $b, (TF_ConstOp (Get2DTransposePerm $transpose_b))),
/*precision_config=*/(StableHLO_GetPrecisionConfig))>;
//===----------------------------------------------------------------------===//
// Lower `tf.ZerosLike`
//===----------------------------------------------------------------------===//
def : Pat<(TF_ZerosLikeOp AnyTensor:$arg),
(StableHLO_ConstantLike<"0"> $arg)>;
//===----------------------------------------------------------------------===//
// Lower `tf.OnesLike`
//===----------------------------------------------------------------------===//
def : Pat<(TF_OnesLikeOp AnyTensor:$arg),
(StableHLO_ConstantLike<"1"> $arg)>;
//===----------------------------------------------------------------------===//
// Elu op patterns.
//===----------------------------------------------------------------------===//
def : Pat<(TF_EluOp AnyTensor:$features),
(StableHLO_SelectOp
(StableHLO_CompareOp
$features,
(StableHLO_ConstantLike<"0">:$zero $features),
StableHLO_ComparisonDirectionValue<"GT">, (STABLEHLO_DEFAULT_COMPARISON_TYPE)),
$features,
(StableHLO_Expm1Op $features, ConstDefaultResultAccuracyAttr))>;
def : Pat<(TF_EluGradOp AnyStaticShapeTensor:$gradients, AnyRankedTensor:$features),
(StableHLO_SelectOp
(CHLO_BroadcastCompareOp
$features,
(StableHLO_ConstantOp:$zero (GetScalarOfType<0> $features)),
(BinBroadcastDimensions $zero, $features),
CHLO_ComparisonDirectionValue<"GT">, (CHLO_DEFAULT_COMPARISON_TYPE)),
$gradients,
(StableHLO_MulOp
$gradients,
(CHLO_BroadcastAddOp
$features,
(StableHLO_ConstantOp:$one (GetScalarOfType<1> $features)),
(BinBroadcastDimensions $one, $features))))>;
//===----------------------------------------------------------------------===//
// Relu op patterns.
//===----------------------------------------------------------------------===//
// TODO(hinsu): Make these patterns to TF to TF lowering. Relu6 lowering will
// require HLO canonicalization of min and max on a tensor to ClampOp.
// TODO(hinsu): Lower quantized types after supporting them in GetScalarOfType.
def : Pat<(TF_ReluOp AnyTensor:$input),
(CHLO_BroadcastMaxOp
(StableHLO_ConstantOp:$zero (GetScalarOfType<0> $input)), $input,
(BinBroadcastDimensions $zero, $input)),
[(TF_IntOrFpTensor $input)]>;
// TODO(hinsu): Lower quantized types after supporting them in GetScalarOfType.
def : Pat<(TF_Relu6Op AnyRankedTensor:$input),
(StableHLO_ClampOp (StableHLO_ConstantOp (GetScalarOfType<0> $input)), $input,
(StableHLO_ConstantOp (GetScalarOfType<6> $input))),
[(TF_IntOrFpTensor $input)]>;
// ReluGrad(gradients, features) = gradients * (features > 0)
// The condition that $gradients and $features need to have the same shape is
// implicitly enforced: $zero is created to have the same shape as $features,
// StableHLO_SelectOp enforces that $gradients and $zero have the same shape.
def : Pat<(TF_ReluGradOp AnyTensor:$gradients, AnyTensor:$features),
(StableHLO_SelectOp
(StableHLO_CompareOp $features, (StableHLO_ConstantLike<"0">:$zero $features),
StableHLO_ComparisonDirectionValue<"GT">, (STABLEHLO_DEFAULT_COMPARISON_TYPE)),
$gradients, $zero)>;
//===----------------------------------------------------------------------===//
// Softsign op patterns.
//===----------------------------------------------------------------------===//
/// Converts a TF::SoftsignOp to HLO.
/// Softsign(features) = features / (1 + abs(features))
def : Pat<(TF_SoftsignOp AnyTensor:$input),
(StableHLO_DivOp
$input,
(StableHLO_AddOp (StableHLO_ConstantLike<"1"> $input), (StableHLO_AbsOp $input))
)
>;
/// Converts a TF::SoftsignGradOp to HLO.
/// SoftsignGrad(gradient, features) = gradient / ((1 + abs(features)) ^ 2)
def : Pattern<
(TF_SoftsignGradOp AnyRankedTensor:$gradients, AnyRankedTensor:$features),
[(CHLO_BroadcastAddOp:$add
(StableHLO_ConstantOp:$one (GetScalarOfType<1> $features)), (StableHLO_AbsOp $features),
(BinBroadcastDimensions $one, $features)
),
(CHLO_BroadcastDivOp
$gradients,
(StableHLO_MulOp $add, $add),
(BinBroadcastDimensions $gradients, $add)
)
]>;
//===----------------------------------------------------------------------===//
// Slice op patterns.
//===----------------------------------------------------------------------===//
def UnpackStartingIndices: NativeCodeCall<
"UnpackTensorAlongZeroDim($0.getLoc(), $1, &$_builder).getOutput()">;
def CanBeTranslatedToDynamicSlice : Constraint<CPred<
"CanBeTranslatedToDynamicSlice($0, $1, llvm::cast<DenseIntElementsAttr>($2))">>;
def TFSliceSizes2HLOSliceSizes : NativeCodeCall<
"TFSliceSizes2HLOSliceSizes($0, $1, llvm::cast<DenseIntElementsAttr>($2),"
"&$_builder)">;
def : Pat<(TF_SliceOp:$op HLO_Tensor:$input, HLO_Tensor:$starting_indices,
(ConstantLikeMatcher AnyAttr:$slice_sizes)),
(StableHLO_DynamicSliceOp $input,
(UnpackStartingIndices $op, $starting_indices),
(TFSliceSizes2HLOSliceSizes $input, $starting_indices, $slice_sizes)),
[(CanBeTranslatedToDynamicSlice $input, $starting_indices,
$slice_sizes)]>;
//===----------------------------------------------------------------------===//
// Select op patterns.
//===----------------------------------------------------------------------===//
def : Pat<(TF_SelectV2Op HLO_Tensor:$pred, HLO_Tensor:$on_true,
HLO_Tensor:$on_false),
(CHLO_BroadcastSelectOp $pred, $on_true, $on_false)>;
//===----------------------------------------------------------------------===//
// PartitionedCall and LegacyCall op patterns.
//===----------------------------------------------------------------------===//
def ArgTypesMatchCallee : Constraint<
// $0 is a resultset (possibly empty), and $_op isn't assigned. So retrieve
// the op using the builder.
CPred<"ArgTypesMatchCallee(&*$_builder.getInsertionPoint(), $1, $2)">>;
foreach callOp = [TF_PartitionedCallOp, TF_StatefulPartitionedCallOp] in {
def : Pat<(callOp:$op $args, $args_attrs, $res_attrs, FlatSymbolRefAttr:$f,
$config, $config_proto, $executor_type),
(CallOp $f, $args, $args_attrs, $res_attrs,
ConstantAttr<UnitAttr, "false">),
[(ArgTypesMatchCallee $op, $args, $f)]>;
}
// The extra attr on this op is _disable_call_shape_inference, which we ignore
// in the bridge.
def : Pat<(TF_LegacyCallOp:$op $args, $args_attrs, $res_attrs,
FlatSymbolRefAttr:$f, $attr),
(CallOp $f, $args, $args_attrs, $res_attrs,
ConstantAttr<UnitAttr, "false">),
[(ArgTypesMatchCallee $op, $args, $f)]>;
//===----------------------------------------------------------------------===//
// Reverse op patterns.
//===----------------------------------------------------------------------===//
// Handles axis conversion for TF reverse.
def ConvertAxisAttr : NativeCodeCall<"ConvertAxisAttr($0, llvm::cast<ElementsAttr>($1), &$_builder)">;
def : Pat<(TF_ReverseV2Op AnyRankedTensor:$values, (ConstantLikeMatcher ElementsAttr:$axis)),
(StableHLO_ReverseOp $values, (ConvertAxisAttr $values, $axis))>;
//===----------------------------------------------------------------------===//
// Unary op patterns.
//===----------------------------------------------------------------------===//
foreach Mapping = [
[TF_AbsOp, StableHLO_AbsOp],
[TF_CeilOp, StableHLO_CeilOp],
[TF_ComplexAbsOp, StableHLO_AbsOp],
[TF_ErfOp, CHLO_ErfOp],
[TF_FloorOp, StableHLO_FloorOp],
[TF_ImagOp, StableHLO_ImagOp],
[TF_InvertOp, StableHLO_NotOp],
[TF_IsFiniteOp, StableHLO_IsFiniteOp],
[TF_LogicalNotOp, StableHLO_NotOp],
[TF_NegOp, StableHLO_NegOp],
[TF_RealOp, StableHLO_RealOp],
] in {
def : Pat<(Mapping[0] HLO_Tensor:$input),
(Mapping[1] $input)>;
}
foreach Mapping = [
[TF_CosOp, StableHLO_CosineOp],
[TF_ExpOp, StableHLO_ExpOp],
[TF_Expm1Op, StableHLO_Expm1Op],
[TF_LogOp, StableHLO_LogOp],
[TF_Log1pOp, StableHLO_Log1pOp],
[TF_RsqrtOp, StableHLO_RsqrtOp],
[TF_SigmoidOp, StableHLO_LogisticOp],
[TF_SinOp, StableHLO_SineOp],
[TF_SqrtOp, StableHLO_SqrtOp],
[TF_TanhOp, StableHLO_TanhOp],
[TF_TanOp, StableHLO_TanOp]
] in {
def : Pat<(Mapping[0] HLO_Tensor:$input),
(Mapping[1] $input, ConstDefaultResultAccuracyAttr)>;
}
foreach Mapping = [
[TF_AcosOp, CHLO_AcosOp],
[TF_AcoshOp, CHLO_AcoshOp],
[TF_AsinOp, CHLO_AsinOp],
[TF_AsinhOp, CHLO_AsinhOp],
[TF_AtanOp, CHLO_AtanOp],
[TF_AtanhOp, CHLO_AtanhOp],
[TF_CoshOp, CHLO_CoshOp],
[TF_ConjOp, CHLO_ConjOp],
[TF_DigammaOp, CHLO_DigammaOp],
[TF_ErfcOp, CHLO_ErfcOp],
[TF_IsInfOp, CHLO_IsInfOp],
[TF_LgammaOp, CHLO_LgammaOp],
[TF_SinhOp, CHLO_SinhOp],
] in {
def : Pat<(Mapping[0] HLO_AnyTensor:$input),
(Mapping[1] $input)>;
}
def : Pat<(TF_AngleOp $x), (StableHLO_Atan2Op (StableHLO_ImagOp $x), (StableHLO_RealOp $x))>;
// TODO(bixia): Lower with Truncate=True for floating point value conversions.
def : Pat<(TF_CastOp $arg, ConstBoolAttrFalse), (StableHLO_ConvertOp $arg)>;
def : Pat<(TF_TransposeOp:$res $arg, (ConstantLikeMatcher ElementsAttr:$permutation)),
(StableHLO_TransposeOp $arg, (CastElementsToI64Array $permutation))>;
// Lowering these ops with static shape to stablehlo.reshape
foreach TfOp = [TF_ExpandDimsOp, TF_ReshapeOp, TF_SqueezeOp, ] in {
def : Pat<(TfOp:$res HLO_Tensor:$arg, $ignored),
(StableHLO_ReshapeOp $arg), [(AnyStaticShapeTensor $res)], [],
(addBenefit 2)>;
}
// Returns NaN if x is NaN, 0 if x is 0, -1 if x < 0 and 1 if x > 0.
def : Pat<(TF_SignOp $x), (StableHLO_SignOp $x)>;
def BothElementTypesSameWidthIntOrFloat : Constraint<CPred<
"getElementTypeOrSelf($0.getType()).isIntOrFloat() && "
"getElementTypeOrSelf($1.getType()).isIntOrFloat()">,
"element types must be integers or floats">;
// TODO(mgester): Due to restrictions of xla::BitcastConvertType we currently
// only lower if both input and output types are int or float and have same width
def : Pat<(TF_BitcastOp:$res HLO_Tensor:$arg),
(StableHLO_BitcastConvertOp $arg),
[(BothElementTypesSameWidthIntOrFloat $res, $arg)]>;
// TODO(jpienaar): Lower constant like to constant to broadcast if dynamic
// and going to StableHLO.
//===----------------------------------------------------------------------===//
// Random ops.
//===----------------------------------------------------------------------===//
// TODO(b/148269299): handle random number generator seeds/states correctly.
class StableHLO_RngDistributionValue<string enumStr> :
ConstantAttr<StableHLO_RngDistributionAttr, "::mlir::stablehlo::RngDistribution::" # enumStr>;
def : Pat<(TF_RandomUniformOp:$old $shape, $seed, $seed2),
(StableHLO_RngOp
(StableHLO_ConstantOp
(NativeCodeCall<"$_builder.getFloatAttr(old.getDtype(), 0.0)">)),
(StableHLO_ConstantOp
(NativeCodeCall<"$_builder.getFloatAttr(old.getDtype(), 1.0)">)),
(CastValueToI64 $old, $shape),
StableHLO_RngDistributionValue<"UNIFORM">),
[(IsShapedTensor $shape)]>;
def : Pat<(TF_RandomStandardNormalOp:$old $shape, $seed, $seed2),
(StableHLO_RngOp
(StableHLO_ConstantOp
(NativeCodeCall<"$_builder.getFloatAttr(old.getDtype(), 0.0)">)),
(StableHLO_ConstantOp
(NativeCodeCall<"$_builder.getFloatAttr(old.getDtype(), 1.0)">)),
(CastValueToI64 $old, $shape),
StableHLO_RngDistributionValue<"NORMAL">),
[(IsShapedTensor $shape)]>;
//===----------------------------------------------------------------------===//
// Sigmoid grad op.
//===----------------------------------------------------------------------===//
// Only handle static shape here, dynamic shape is handled by
// ConvertSigmoidGradOpDynamic
def HasStaticShape : Constraint<
CPred<"::llvm::dyn_cast<ShapedType>($0.getType()).hasStaticShape()">>;
def : Pat<(TF_SigmoidGradOp AnyRankedTensor:$l, AnyRankedTensor:$r),
(StableHLO_MulOp
(StableHLO_MulOp $r, $l),
(StableHLO_SubtractOp (StableHLO_ConstantOp (ConstantSplat<"1"> $l)), $l)),
[(HasStaticShape $l)]>;
//===----------------------------------------------------------------------===//
// Softplus op.
//===----------------------------------------------------------------------===//
def EpsilonValue : NativeCodeCall<"GetEpsilonValue($0.getType())">;
def : Pattern<(TF_SoftplusOp AnyTensor:$features),
[
(StableHLO_ExpOp:$features_exp $features, ConstDefaultResultAccuracyAttr),
(CHLO_BroadcastAddOp:$threshold
(StableHLO_LogOp (StableHLO_ConstantOp (EpsilonValue $features)), ConstDefaultResultAccuracyAttr),
(StableHLO_ConstantOp (GetScalarOfType<2> $features)),
(NullDenseI64ArrayAttr)
),
(StableHLO_SelectOp:$output
(CHLO_BroadcastCompareOp
$features,
(StableHLO_NegOp $threshold),
(NullDenseI64ArrayAttr),
CHLO_ComparisonDirectionValue<"GT">,
(CHLO_DEFAULT_COMPARISON_TYPE)
),
$features,
(StableHLO_SelectOp
(CHLO_BroadcastCompareOp
$features,
$threshold,
(NullDenseI64ArrayAttr),
CHLO_ComparisonDirectionValue<"LT">,
(CHLO_DEFAULT_COMPARISON_TYPE)
),
$features_exp,
(StableHLO_Log1pOp $features_exp, ConstDefaultResultAccuracyAttr)
)
),
(replaceWithValue $output)
]>;
//===----------------------------------------------------------------------===//
// XlaReplicaId op.
//===----------------------------------------------------------------------===//
def : Pat<(TF_XlaReplicaIdOp),
(TF_CastOp (StableHLO_ReplicaIdOp), /*truncate=*/ConstBoolAttrFalse)>;
//===----------------------------------------------------------------------===//
// XlaGather op.
//===----------------------------------------------------------------------===//
def ToGatherDimNumsAttr : NativeCodeCall<"GetGatherDimNumsAttr($0, &$_builder)">;
def HasValidGatherDims : Constraint<CPred<"HasValidGatherDims($0)">>;
def : Pat<(TF_XlaGatherOp $operand, $start_indices, (ConstantLikeMatcher ElementsAttr:$slice_sizes),
$dimension_numbers, $indices_are_sorted),
(StableHLO_GatherOp $operand, $start_indices,
(ToGatherDimNumsAttr $dimension_numbers),
(CastElementsToI64Array $slice_sizes),
$indices_are_sorted),
[(HasValidGatherDims $dimension_numbers)]>;
//===----------------------------------------------------------------------===//
// XlaDotOp op.
//===----------------------------------------------------------------------===//
def ToDotDimNumsAttr : NativeCodeCall<"GetDotDimNumsAttr($0, &$_builder)">;
def ToPrecisionConfigsAttr : NativeCodeCall<"GetPrecisionConfigAttr($0, &$_builder)">;
def HasValidDotDims : Constraint<CPred<"HasValidDotDims($0)">>;
def HasValidPrecisionConfig : Constraint<CPred<"HasValidPrecisionConfig($0)">>;
def : Pat<(TF_XlaDotOp $lhs, $rhs, $dimension_numbers, $precision_config),
(StableHLO_DotGeneralOp $lhs, $rhs,
(ToDotDimNumsAttr $dimension_numbers),
(ToPrecisionConfigsAttr $precision_config),
(EmptyDotAlgorithmAttr)),
[(HasValidDotDims $dimension_numbers), (HasValidPrecisionConfig $precision_config)]>;
//===----------------------------------------------------------------------===//
// XlaDotV2Op op.
//===----------------------------------------------------------------------===//
def : Pat<(TF_XlaDotV2Op $lhs, $rhs, $dimension_numbers, $precision_config),
(StableHLO_DotGeneralOp $lhs, $rhs,
(ToDotDimNumsAttr $dimension_numbers),
(ToPrecisionConfigsAttr $precision_config),
(EmptyDotAlgorithmAttr)),
[(HasValidDotDims $dimension_numbers), (HasValidPrecisionConfig $precision_config)]>;
//===----------------------------------------------------------------------===//
// XlaDynamicSlice op.
//===----------------------------------------------------------------------===//
def : Pat<(TF_XlaDynamicSliceOp:$op HLO_Tensor:$input, HLO_Tensor:$starting_indices,
(ConstantLikeMatcher AnyAttr:$slice_sizes)),
(StableHLO_DynamicSliceOp $input,
(UnpackStartingIndices $op, $starting_indices),
(TFSliceSizes2HLOSliceSizes $input, $starting_indices, $slice_sizes))>;
//===----------------------------------------------------------------------===//
// XlaEisumOp op.
//===----------------------------------------------------------------------===//
def : Pat<(TF_XlaEinsumOp $lhs, $rhs, $equation),
(StableHLO_EinsumOp $lhs, $rhs, $equation)>;
//===----------------------------------------------------------------------===//
// XlaOptimizationBarrierOp op.
//===----------------------------------------------------------------------===//
def : Pat<(TF_XlaOptimizationBarrierOp $args),
(StableHLO_OptimizationBarrierOp $args)>;
tensorflow/compiler/mlir/tf2xla/transforms/legalize_tf_with_tf2xla.cc
+240
View File
@@ -0,0 +1,240 @@
/* 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 <string>
#include <utility>
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/Tensor/IR/Tensor.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/BuiltinTypeInterfaces.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Diagnostics.h" // from @llvm-project
#include "mlir/IR/IRMapping.h" // from @llvm-project
#include "mlir/IR/Location.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "stablehlo/dialect/Base.h" // from @stablehlo
#include "tensorflow/compiler/mlir/op_or_arg_name_mapper.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tpu_embedding_ops_registry.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/legalization_op_config.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/legalize_tf_with_tf2xla_passes.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/passes.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/tf2xla_rewriter.h"
#include "tensorflow/compiler/tf2xla/xla_compilation_device.h"
#include "tensorflow/compiler/tf2xla/xla_context.h"
#include "tensorflow/compiler/tf2xla/xla_expression.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "xla/tsl/platform/env.h"
#include "xla/tsl/platform/status.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/common_runtime/device_factory.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/common_runtime/process_function_library_runtime.h"
#include "tensorflow/core/framework/allocator.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/node_properties.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
namespace mlir {
namespace mhlo {
namespace {
// Returns true if the given type is a ranked tensor type with static or bounded
// dimensions.
bool IsBounded(Type ty) {
auto ranked_ty = mlir::dyn_cast<RankedTensorType>(ty);
if (!ranked_ty) return false;
if (ranked_ty.hasStaticShape()) return true;
auto bounds = hlo::encodingToBounds(ranked_ty.getEncoding());
if (bounds.empty()) return false;
for (int i = 0; i < ranked_ty.getRank(); ++i) {
if (ranked_ty.isDynamicDim(i) && bounds[i] == ShapedType::kDynamic) {
return false;
}
}
return true;
}
bool HasSymbolRefAttr(Operation* op) {
for (const auto& attr : op->getAttrs()) {
Attribute attr_value = attr.getValue();
if (mlir::isa<SymbolRefAttr>(attr_value)) {
return true;
} else if (auto array_attr = mlir::dyn_cast<ArrayAttr>(attr_value)) {
if (!array_attr.empty() &&
mlir::isa<SymbolRefAttr>(*array_attr.begin())) {
return true;
}
}
}
return false;
}
class Tf2XlaRewritePattern : public ConversionPattern {
public:
explicit Tf2XlaRewritePattern(MLIRContext* ctx, TypeConverter& converter,
const std::string& device_type,
bool prefer_tf2xla)
: ConversionPattern(converter, MatchAnyOpTypeTag(), /*benefit=*/1, ctx),
device_type_(device_type),
prefer_tf2xla_(prefer_tf2xla) {}
LogicalResult matchAndRewrite(
Operation* op, ArrayRef<Value> operands,
ConversionPatternRewriter& rewriter) const override {
// This pattern is a conversion pattern because we want to specify a type
// converter. However, this pattern still uses the original op's operands
// while creating the ops so make sure there aren't any type changes between
// the original op operands and the operands during the conversion.
for (auto&& [old_val, new_val] : llvm::zip(op->getOperands(), operands)) {
if (old_val.getType() != new_val.getType()) return failure();
}
auto abstractOp = op->getRegisteredInfo();
if (!abstractOp) return failure();
if (!(hlo::IsOpAllowedTf2xlaFallback(abstractOp->getTypeID()) ||
(prefer_tf2xla_ &&
hlo::IsOpAllowedTf2xlaPreferred(abstractOp->getTypeID())))) {
return failure();
}
return hlo::Tf2XlaRewriter::RewriteOp(op, rewriter, device_type_);
}
private:
std::string device_type_;
bool prefer_tf2xla_;
bool use_tf2xla_hlo_importer_;
};
bool ShouldRefineTypeTo(Type original_ty, Type updated_ty) {
auto updated = mlir::dyn_cast<ShapedType>(updated_ty);
auto original = mlir::dyn_cast<ShapedType>(original_ty);
// Both types must be shaped types.
if (!original || !updated) return false;
// Element types must match.
if (original.getElementType() != updated.getElementType()) return false;
// If the updated type doesn't have a rank, then it can't be a more refined
// type.
if (!updated.hasRank()) return false;
// If the original type doesn't have a rank, then refine as the updated type
// has a rank.
if (!original.hasRank()) return true;
// Both types must have the same rank.
if (original.getRank() != updated.getRank()) return false;
// Refine if the updated type is bounded.
return IsBounded(updated);
}
// Propagates more refined type by cloning op using the new operands. This
// allows all rewrite patterns that requires refined types to work without
// requiring a rewrite to the conversion pattern. Declarative rewrite pattern
// (DRR) doesn't even support conversion patterns with TableGen.
class TypePropagator : public ConversionPattern {
public:
explicit TypePropagator(MLIRContext* ctx)
: ConversionPattern(MatchAnyOpTypeTag(), /*benefit=*/1, ctx) {}
LogicalResult matchAndRewrite(
Operation* op, ArrayRef<Value> operands,
ConversionPatternRewriter& rewriter) const override {
// This could be generalized to other ops as needs arise. We could even
// remove this restriction altogether except for the terminators that
// require function signature change and shouldn't be
if (op->getName().getDialectNamespace() !=
TF::TensorFlowDialect::getDialectNamespace())
return failure();
// Refining types may have implications to the attached regions or symbol
// references so do not update such ops.
if (!op->getRegions().empty() || HasSymbolRefAttr(op)) return failure();
IRMapping mapper;
bool has_type_change = false;
for (auto [original, updated] : llvm::zip(op->getOperands(), operands)) {
Type original_ty = original.getType();
Type updated_ty = updated.getType();
if (original_ty != updated_ty) has_type_change = true;
if (!ShouldRefineTypeTo(original_ty, updated_ty)) return failure();
mapper.map(original, updated);
}
if (!has_type_change) return failure();
Operation* cloned_op = rewriter.clone(*op, mapper);
rewriter.replaceOp(op, cloned_op->getResults());
return success();
}
};
} // end namespace
Tf2XlaTypeConverter::Tf2XlaTypeConverter() {
// Currently, we don't do any type conversions. Any TensorFlow op with a type
// that is not supported in MHLO will fail conversion. Quantized types are
// going to handled separately so we don't need to handle those.
addConversion([](Type ty) { return ty; });
// This materialization is helpful in cases where we have more refined types
// after conversion to mhlo compared to the original type in TF. For example,
// a TF op with result type tensor<*xf32> will have a bounded type after
// fallback legalization.
auto cast_value = [&](OpBuilder& builder, Type result_type, ValueRange inputs,
Location loc) -> Value {
return mlir::tensor::CastOp::create(builder, loc, result_type,
inputs.front());
};
addSourceMaterialization(cast_value);
}
void PopulateLegalizeTfWithTf2XlaPatterns(llvm::StringRef device_type,
RewritePatternSet& patterns,
MLIRContext* ctx,
Tf2XlaTypeConverter& converter,
bool prefer_tf2xla) {
patterns.add<TypePropagator>(ctx);
patterns.add<Tf2XlaRewritePattern>(ctx, converter, device_type.str(),
prefer_tf2xla);
}
} // end namespace mhlo
} // end namespace mlir
tensorflow/compiler/mlir/tf2xla/transforms/legalize_tf_with_tf2xla_passes.h
+63
View File
@@ -0,0 +1,63 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_TRANSFORMS_LEGALIZE_TF_WITH_TF2XLA_PASSES_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_TRANSFORMS_LEGALIZE_TF_WITH_TF2XLA_PASSES_H_
#include <memory>
#include <optional>
#include "llvm/ADT/StringRef.h"
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Transforms/DialectConversion.h" // from @llvm-project
namespace mlir {
namespace func {
class FuncOp;
}
class ModuleOp;
class Operation;
template <typename T>
class OperationPass;
class Pass;
namespace mhlo {
/// Converter to be used along with the fallback Tf2Xla patterns below.
class Tf2XlaTypeConverter : public TypeConverter {
public:
Tf2XlaTypeConverter();
};
/// Adds the TF to XLA via TF2XLA rewrite patterns to the pattern list.
/// `prefer_tf2xla` means an op will be included iff it is not in
/// `MlirLegalizedUnderPreferTf2XlaSet`. `!prefer_tf2xla` mean an op will be
/// included if there is no native MLIR legalization for the op.
void PopulateLegalizeTfWithTf2XlaPatterns(llvm::StringRef device_type,
RewritePatternSet& patterns,
MLIRContext* ctx,
Tf2XlaTypeConverter& converter,
bool prefer_tf2xla = false);
} // namespace mhlo
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_TRANSFORMS_LEGALIZE_TF_WITH_TF2XLA_PASSES_H_
tensorflow/compiler/mlir/tf2xla/transforms/passes.h
+124
View File
@@ -0,0 +1,124 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_TRANSFORMS_PASSES_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_TRANSFORMS_PASSES_H_
#include <memory>
#include <optional>
#include "llvm/ADT/StringRef.h"
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Transforms/DialectConversion.h" // from @llvm-project
namespace mlir {
namespace func {
class FuncOp;
}
class ModuleOp;
class Operation;
template <typename T>
class OperationPass;
class Pass;
namespace hlo {
// Verifies that the TF/XLA ops have all been lowered to MHLO.
std::unique_ptr<OperationPass<func::FuncOp>> CreateVerifyTFXLALegalizationPass(
bool legalize_chlo = true);
/// Adds the TF to TF lowerings and TF to XLA rewrite patterns to the pattern
/// list.
void PopulateLegalizeTfPatterns(MLIRContext* context,
RewritePatternSet* patterns);
} // namespace hlo
namespace mhlo {
/// Lowers from TF dialect to HLO dialect. When allow_partial_conversion is
/// false, emits an error if there is any operation that can't be legalized.
/// When `tf2xla_fallback_device_type` is not `None`, also uses legalization
/// patterns from TF2XLA fallback for provided device type (see
/// legalize_tf_with_tf2xla.cc for details). By default, TF2XLA fallback is not
/// used.
/// Note: This is a module pass because when legalizing with TF2XLA fallback,
/// functions are imported into the module. Importing functions into a
/// module is not thread safe.
std::unique_ptr<OperationPass<ModuleOp>> createLegalizeTFPass(
bool legalize_chlo = true,
std::optional<StringRef> tf2xla_fallback_device_type = std::nullopt,
bool prefer_tf2xla = false);
// Populates TF to MHLO legalization for some of the quantization ops.
//
// TODO(hinsu): Remove this once we combine quantized and non quantized op
// legalization in the ODML conversion pipeline.
void PopulateLegalizeTfQuantizationPatterns(MLIRContext* context,
RewritePatternSet* patterns);
/// Converts the provided Operation as well as all nested operations into HLO
/// dialect using the conversion patterns registered by the HLO dialect. When
/// allow_partial_conversion is false, emits an error if there is any operation
/// that can't be legalized.
/// When `tf2xla_fallback_device_type` is not `None`, also uses legalization
/// patterns from TF2XLA fallback for provided device type (see
/// legalize_tf_with_tf2xla.cc for details). By default, TF2XLA fallback is not
/// used.
LogicalResult legalizeTF(
Operation* op, bool allow_partial_conversion = false,
bool legalize_chlo = true,
std::optional<StringRef> tf2xla_fallback_device_type = std::nullopt,
bool prefer_tf2xla = false);
// Legalizes TF/XLA communication ops (TF dialect) to HLO dialect communication
// ops.
std::unique_ptr<OperationPass<ModuleOp>> CreateLegalizeTFCommunicationPass();
// Legalizes TF/XLA collective ops (TF dialect) to HLO dialect collective
// ops.
std::unique_ptr<OperationPass<ModuleOp>> CreateLegalizeTFCollectivePass();
// Transforms TFXLA Device specific ops into device independent ops.
std::unique_ptr<OperationPass<func::FuncOp>>
CreateTFXLADeviceSpecificTransformsPass(
std::optional<StringRef> tf2xla_fallback_device_type = std::nullopt);
// Adjusts XLA layout for Infeed ops.
std::unique_ptr<OperationPass<func::FuncOp>>
CreateInfeedsOpsXlaAdjustLayoutPass();
#define GEN_PASS_REGISTRATION
#define GEN_PASS_DECL_INFEEDSOPSXLAADJUSTLAYOUT
#define GEN_PASS_DECL_LEGALIZETF
#define GEN_PASS_DECL_LEGALIZETFCOLLECTIVE
#define GEN_PASS_DECL_LEGALIZETFMODULEPASS
#define GEN_PASS_DECL_LEGALIZETFTYPESPASS
#define GEN_PASS_DECL_TFXLADEVICESPECIFICTRANSFORMS
#define GEN_PASS_DECL_VERIFYTFXLALEGALIZATION
#include "tensorflow/compiler/mlir/tf2xla/transforms/xla_legalize_tf_passes.h.inc"
#define GEN_PASS_REGISTRATION
#define GEN_PASS_DECL_LEGALIZETFCOMMUNICATIONPASS
#include "tensorflow/compiler/mlir/tf2xla/transforms/tf_xla_passes.h.inc"
} // namespace mhlo
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_TRANSFORMS_PASSES_H_
tensorflow/compiler/mlir/tf2xla/transforms/split_into_island_per_op_pass.cc
+192
View File
@@ -0,0 +1,192 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/transforms/split_into_island_per_op_pass.h"
#include <memory>
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
// This pass is used in preparation for Graph export.
// The GraphDef exporter expects each op to be in its own island.
// This pass puts the IR in that form.
//
// We do this as an IR->IR transform to keep the Graph exporter as simple as
// possible.
namespace mlir {
namespace TF {
namespace {
#define GEN_PASS_DEF_SPLITINTOISLANDPEROPPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"
class SplitIntoIslandPerOpPass
: public impl::SplitIntoIslandPerOpPassBase<SplitIntoIslandPerOpPass> {
public:
void runOnOperation() override;
};
void SplitIntoIslandPerOpPass::runOnOperation() {
func::FuncOp func = getOperation();
if (func.isExternal()) {
// Just ignore the op if this is an external func with no body.
return;
}
tf_executor::GraphOp graph_op;
if (llvm::hasSingleElement(func.front().without_terminator())) {
graph_op = dyn_cast<tf_executor::GraphOp>(func.front().front());
}
if (!graph_op) {
func.emitError("expected function to contain only a graph_op");
signalPassFailure();
return;
}
if (!(llvm::hasSingleElement(graph_op.GetBody().without_terminator()) &&
llvm::isa<tf_executor::IslandOp>(graph_op.GetBody().front()))) {
graph_op.emitError(
"expected graph op to contain only a single island_op and a single "
"fetch_op");
signalPassFailure();
return;
}
tf_executor::IslandOp island_op =
dyn_cast<tf_executor::IslandOp>(graph_op.GetBody().front());
// We don't need to honor *any* control deps that are already placed on
// islands. Drop them now in this pass - a following pass will use side
// effect analysis to completely explain and apply correct control deps.
island_op.getControl().dropAllUses();
// Break up all islands by simply creating a new island wrapping each
// individual sub op. Do not create any control dependencies between the
// newly created islands.
SplitIsland(island_op, tf_executor::ControlType::get(&getContext()));
// None of the originally given control deps are necessary.
tf_executor::FetchOp fetch_op = graph_op.GetFetch();
int num_control_fetches =
fetch_op.getNumOperands() - graph_op.getNumResults();
if (num_control_fetches > 0) {
fetch_op.getFetchesMutable().erase(graph_op.getNumResults(),
num_control_fetches);
}
}
} // namespace
// Populates an empty IslandOp and with a NoOp or Identity/IdentityN depending
// on if there are any data results.
void PopulateEmptyIsland(tf_executor::IslandOp island) {
OpBuilder builder(&island.GetBody(), island.GetBody().begin());
tf_executor::YieldOp yield = island.GetYield();
if (yield.getNumOperands() == 0) {
TF::NoOp::create(builder, island.getLoc(), TypeRange{}, ValueRange{});
} else if (yield.getNumOperands() == 1) {
Value operand = yield.getOperand(0);
auto identity = TF::IdentityOp::create(builder, island.getLoc(),
operand.getType(), operand);
yield.setOperand(0, identity.getOutput());
} else {
auto identity_n = TF::IdentityNOp::create(
builder, island.getLoc(), yield.getOperandTypes(), yield.getOperands());
for (const auto& it : llvm::enumerate(identity_n.getResults()))
yield.setOperand(it.index(), it.value());
}
}
// Helper that creates an island that `sub_op` will be moved to.
tf_executor::IslandOp CreateIsland(TypeRange result_types,
const tf_executor::ControlType& control_type,
const Location& loc, Operation& sub_op,
tf_executor::IslandOp original_island) {
OpBuilder builder(original_island);
auto island = tf_executor::IslandOp::create(builder, loc, result_types,
control_type, mlir::ValueRange{});
island.getBody().push_back(new Block);
Block* block = &island.getBody().back();
OpBuilder island_builder(original_island);
island_builder.setInsertionPointToEnd(block);
sub_op.replaceAllUsesWith(island.getOutputs());
sub_op.moveBefore(block, block->begin());
tf_executor::YieldOp::create(island_builder, loc, sub_op.getResults());
return island;
}
// Converts a single island into multiple islands (one for each op).
void SplitIsland(mlir::tf_executor::IslandOp island_op,
mlir::tf_executor::ControlType control_type) {
auto island_body = island_op.GetBody().without_terminator();
// Populate islands that are empty (only yield).
if (island_body.empty()) {
PopulateEmptyIsland(island_op);
return;
}
// Skip islands that are already only a single op.
if (island_op.WrapsSingleOp()) return;
// For each operation in the island, construct a new island to wrap the op,
// yield all the results, and replace all the usages with the results of the
// new island.
for (auto& sub_op : llvm::make_early_inc_range(island_body)) {
CreateIsland(sub_op.getResultTypes(), control_type, sub_op.getLoc(), sub_op,
island_op);
}
// Ensure that consumers of the outputs of the original island now depend
// directly on the new island wrapping the original output op instead of
// depending on the original island's output since we're about to erase the
// original island op.
for (auto item :
llvm::zip(island_op.getOutputs(), island_op.GetYield().getFetches()))
std::get<0>(item).replaceAllUsesWith(std::get<1>(item));
auto graph_op = island_op->getParentOfType<mlir::tf_executor::GraphOp>();
// Dropping all uses of an island op's control dep using
// `island_op.getControl().dropAllUses();` of a control dep that's only used
// in a graph's fetch, immediately leads to a segfault. Turns out we need to
// drop its uses manually so that we don't leave dangling controls.
for (const auto& fetch : llvm::enumerate(graph_op.GetFetch().getFetches())) {
if (fetch.value() == island_op.getControl()) {
graph_op.GetFetch().getFetchesMutable().erase(fetch.index(), 1);
break;
}
}
island_op.erase();
}
std::unique_ptr<OperationPass<func::FuncOp>> CreateSplitIntoIslandPerOpPass() {
return std::make_unique<TF::SplitIntoIslandPerOpPass>();
}
} // namespace TF
} // namespace mlir
tensorflow/compiler/mlir/tf2xla/transforms/split_into_island_per_op_pass.h
+31
View File
@@ -0,0 +1,31 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_TRANSFORMS_SPLIT_INTO_ISLAND_PER_OP_PASS_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_TRANSFORMS_SPLIT_INTO_ISLAND_PER_OP_PASS_H_
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
namespace mlir {
namespace TF {
// Converts a single island into multiple islands (one for each op).
void SplitIsland(mlir::tf_executor::IslandOp island_op,
mlir::tf_executor::ControlType control_type);
} // namespace TF
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_TRANSFORMS_SPLIT_INTO_ISLAND_PER_OP_PASS_H_
tensorflow/compiler/mlir/tf2xla/transforms/split_into_island_per_op_pass_test.cc
+150
View File
@@ -0,0 +1,150 @@
/* 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/transforms/split_into_island_per_op_pass.h"
#include <set>
#include <string>
#include <vector>
#include <gtest/gtest.h>
#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/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OpDefinition.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/IR/TypeRange.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/register_common_dialects.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
namespace mlir {
namespace TF {
using mlir::DialectRegistry;
using mlir::MLIRContext;
class SplitIntoIslandPerOpPass : public ::testing::Test {
public:
SplitIntoIslandPerOpPass() : op_builder_(&context_) {
mlir::RegisterCommonToolingDialects(registry_);
context_.appendDialectRegistry(registry_);
context_.loadAllAvailableDialects();
context_.allowUnregisteredDialects(true);
}
mlir::tf_executor::IslandOp GenerateIslandOp() {
mlir::OperationState op_state(mlir::UnknownLoc::get(&context_),
StringRef("test_island"));
llvm::SmallVector<mlir::Type, 1> island_result_types;
island_result_types.push_back(op_builder_.getF64Type());
mlir::Operation* yield_op = mlir::tf_executor::YieldOp::create(
op_builder_, op_state.location, mlir::ValueRange{});
mlir::tf_executor::IslandOp island_op = mlir::tf_executor::IslandOp::create(
op_builder_, op_state.location, island_result_types, mlir::ValueRange{},
mlir::ArrayRef<mlir::NamedAttribute>{});
island_op.getBody().push_back(new mlir::Block);
island_op.getBody().back().push_back(yield_op);
return island_op;
}
mlir::Operation* GenerateOp(StringRef name, bool add_return_type = false) {
mlir::OperationState state(mlir::UnknownLoc::get(&context_), name);
if (add_return_type) {
state.addTypes(ArrayRef<mlir::FloatType>{op_builder_.getF64Type()});
}
return mlir::Operation::create(state);
}
std::vector<std::string> GetOpNames(mlir::Operation* op) {
std::vector<std::string> op_names;
for (auto& op : op->getBlock()->getOperations()) {
op_names.push_back(op.getName().getStringRef().str());
}
return op_names;
}
DialectRegistry registry_;
MLIRContext context_;
OpBuilder op_builder_;
};
TEST_F(SplitIntoIslandPerOpPass, EmptyIslandPoulatesNoOp) {
mlir::tf_executor::ControlType control_type;
mlir::tf_executor::IslandOp islandOp = GenerateIslandOp();
SplitIsland(islandOp, control_type);
std::set<std::string> actual_op_names;
for (auto& op : islandOp.getBody().getOps()) {
actual_op_names.insert(op.getName().getStringRef().str());
}
std::set<std::string> expected_op_names = {"tf_executor.yield", "tf.NoOp"};
ASSERT_EQ(actual_op_names, expected_op_names);
islandOp.erase();
}
TEST_F(SplitIntoIslandPerOpPass, IslandOpSingleOpLeftUnchanged) {
mlir::tf_executor::ControlType control_type;
mlir::tf_executor::IslandOp islandOp = GenerateIslandOp();
mlir::Operation* inner_op = GenerateOp("inner_op");
islandOp.getBody().front().push_front(inner_op);
SplitIsland(islandOp, control_type);
std::set<std::string> actual_op_names;
for (auto& op : islandOp.getBody().getOps()) {
actual_op_names.insert(op.getName().getStringRef().str());
}
std::set<std::string> expected_op_names = {"inner_op", "tf_executor.yield"};
ASSERT_EQ(actual_op_names, expected_op_names);
islandOp.erase();
}
TEST_F(SplitIntoIslandPerOpPass, IslandOpTwoOpsSplitsIntoTwoIslands) {
auto control_type = mlir::tf_executor::ControlType::get(&context_);
mlir::tf_executor::IslandOp islandOp = GenerateIslandOp();
mlir::Operation* inner_op_1 = GenerateOp("inner_op_1", true);
mlir::Operation* inner_op_2 = GenerateOp("inner_op_2", true);
islandOp.getBody().front().push_front(inner_op_1);
islandOp.getBody().back().push_front(inner_op_2);
// Code relies on a parent with a fetch op containing the island op.
mlir::tf_executor::GraphOp parent_graph_op =
mlir::tf_executor::GraphOp::create(
op_builder_, mlir::UnknownLoc::get(&context_),
mlir::TypeRange{op_builder_.getF64Type()});
parent_graph_op.getRegion().push_back(new mlir::Block);
parent_graph_op.push_back(islandOp);
mlir::tf_executor::FetchOp fetch_op =
mlir::tf_executor::FetchOp::create(op_builder_, parent_graph_op.getLoc());
parent_graph_op.GetBody().push_back(fetch_op);
SplitIsland(islandOp, control_type);
std::vector<std::string> actual_op_names;
for (auto& op : parent_graph_op.getBody().getOps()) {
actual_op_names.push_back(op.getName().getStringRef().str());
}
std::vector<std::string> expected_op_names = {
"tf_executor.island", "tf_executor.island", "tf_executor.fetch"};
ASSERT_EQ(actual_op_names, expected_op_names);
parent_graph_op.erase();
}
} // namespace TF
} // namespace mlir
tensorflow/compiler/mlir/tf2xla/transforms/test_utils.cc
+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.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/transforms/test_utils.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "tensorflow/compiler/mlir/register_common_dialects.h"
#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/deserialize_mlir_module_utils.h"
#include "xla/tsl/platform/statusor.h"
namespace mlir {
namespace hlo {
namespace test {
using ::mlir::DialectRegistry;
using ::mlir::MLIRContext;
using ::mlir::ModuleOp;
using ::mlir::OwningOpRef;
using ::tsl::StatusOr;
absl::StatusOr<OwningOpRef<ModuleOp>> GetMlirModuleFromString(
absl::string_view module_string, MLIRContext* context) {
DialectRegistry mlir_registry;
RegisterCommonToolingDialects(mlir_registry);
context->appendDialectRegistry(mlir_registry);
OwningOpRef<ModuleOp> mlir_module;
auto status =
tensorflow::DeserializeMlirModule(module_string, context, &mlir_module);
if (!status.ok()) {
return status;
}
return mlir_module;
}
} // namespace test
} // namespace hlo
} // namespace mlir
tensorflow/compiler/mlir/tf2xla/transforms/test_utils.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_TRANSFORMS_TEST_UTILS_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_TRANSFORMS_TEST_UTILS_H_
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/utils/deserialize_mlir_module_utils.h"
#include "xla/tsl/platform/statusor.h"
namespace mlir {
namespace hlo {
namespace test {
// Given a raw string, return a ModuleOp that can be used with the given
// MLIRContext.
absl::StatusOr<OwningOpRef<ModuleOp>> GetMlirModuleFromString(
absl::string_view module_string, MLIRContext* mlir_context);
} // namespace test
} // namespace hlo
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_TRANSFORMS_TEST_UTILS_H_
tensorflow/compiler/mlir/tf2xla/transforms/tf2xla_rewriter.cc
+526
View File
@@ -0,0 +1,526 @@
/* 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/transforms/tf2xla_rewriter.h"
#include <cstddef>
#include <cstdint>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/memory/memory.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_replace.h"
#include "absl/types/span.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h" // from @llvm-project
#include "mlir/Dialect/Tensor/IR/Tensor.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/BuiltinTypeInterfaces.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Diagnostics.h" // from @llvm-project
#include "mlir/IR/IRMapping.h" // from @llvm-project
#include "mlir/IR/Location.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "mlir/IR/Value.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/DialectConversion.h" // from @llvm-project
#include "stablehlo/dialect/Base.h" // from @stablehlo
#include "stablehlo/dialect/StablehloOps.h" // from @stablehlo
#include "tensorflow/compiler/mlir/op_or_arg_name_mapper.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tpu_embedding_ops_registry.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/convert_tensor.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/convert_type.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/translate_utils.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/passes.h"
#include "tensorflow/compiler/tf2xla/xla_compilation_device.h"
#include "tensorflow/compiler/tf2xla/xla_context.h"
#include "tensorflow/compiler/tf2xla/xla_expression.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "xla/hlo/builder/xla_builder.h"
#include "xla/hlo/builder/xla_computation.h"
#include "xla/hlo/ir/hlo_instruction.h"
#include "xla/hlo/ir/hlo_opcode.h"
#include "xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.h"
#include "xla/hlo/translate/hlo_to_mhlo/hlo_to_mlir_hlo.h"
#include "xla/hlo/translate/mhlo_to_hlo/type_to_shape.h"
#include "xla/service/hlo.pb.h"
#include "xla/tsl/platform/env.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/status.h"
#include "xla/tsl/platform/statusor.h"
#include "xla/xla.pb.h"
#include "xla/xla_data.pb.h"
#include "tensorflow/core/common_runtime/process_function_library_runtime.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/public/session_options.h"
namespace mlir {
namespace hlo {
namespace {
using ::mlir::ModuleOp;
using ::tensorflow::Tensor;
using ::tsl::StatusOr;
using ::xla::XlaComputation;
// The OpOrArgLocNameMapper adds invalid characters to the name of the op when
// concatenating locations. This version removes those characters to make the
// name valid for NodeDef.
class OpOrArgLocNameMapperWithoutInvalidCharacters
: public tensorflow::OpOrArgLocNameMapper {
public:
OpOrArgLocNameMapperWithoutInvalidCharacters() = default;
~OpOrArgLocNameMapperWithoutInvalidCharacters() override = default;
protected:
std::string GetName(tensorflow::OpOrVal op_or_val) override {
std::string name = OpOrArgLocNameMapper::GetName(op_or_val);
return absl::StrReplaceAll(name, {{";", "."}});
}
};
static std::unique_ptr<tensorflow::StaticDeviceMgr> CreateDeviceMgr(
const std::string& device_type) {
// Register compilation kernels for all registered XLA backends.
tensorflow::XlaOpRegistry::RegisterCompilationKernels();
auto device = std::make_unique<tensorflow::XlaCompilationDevice>(
tensorflow::SessionOptions(), tensorflow::DeviceType(device_type));
return std::make_unique<tensorflow::StaticDeviceMgr>(std::move(device));
}
bool RootInstructionIsTuple(const xla::HloModule& hlo_module) {
xla::HloInstruction* root_instruction =
hlo_module.entry_computation()->root_instruction();
return root_instruction->opcode() == xla::HloOpcode::kTuple;
}
}; // namespace
LogicalResult Tf2XlaRewriter::RewriteOp(Operation* op,
PatternRewriter& rewriter,
const std::string& device_type) {
Tf2XlaRewriter tf2xla_rewriter(op, rewriter, device_type);
return tf2xla_rewriter.LegalizeOp();
}
Tf2XlaRewriter::Tf2XlaRewriter(Operation* op, PatternRewriter& rewriter,
const std::string& device_type)
: op_(op),
device_type_(device_type),
rewriter_(rewriter),
name_mapper_(
std::make_unique<OpOrArgLocNameMapperWithoutInvalidCharacters>()),
context_(nullptr),
xla_builder_(op_->getName().getStringRef().str()) {}
Tf2XlaRewriter::~Tf2XlaRewriter() {
if (context_) context_->Unref();
}
absl::StatusOr<stablehlo::TupleOp> Tf2XlaRewriter::ImportXlaComputation(
XlaComputation& computation) {
xla::DebugOptions debug_options;
TF_ASSIGN_OR_RETURN(auto hlo_module_config,
xla::HloModule::CreateModuleConfigFromProto(
computation.proto(), debug_options));
TF_ASSIGN_OR_RETURN(
std::unique_ptr<xla::HloModule> hlo_module,
xla::HloModule::CreateFromProto(computation.proto(), hlo_module_config));
if (!RootInstructionIsTuple(*hlo_module)) {
return absl::InvalidArgumentError("Imported XLA Root is not a tuple op");
}
if (op_->getNumOperands() !=
hlo_module->entry_computation()->num_parameters()) {
return absl::InvalidArgumentError(
"Entry computation does not have equal number of parameters to op "
"operands");
}
ModuleOp mlir_module = op_->getParentOfType<ModuleOp>();
mlir::OpBuilder builder(op_);
mlir::SymbolTable symbol_table(mlir_module);
llvm::SmallVector<mlir::Value> arguments;
for (int i = 0; i < op_->getNumOperands(); i++) {
arguments.push_back(op_->getOperand(i));
}
// Ideally we could use the Function Importer but it increases compilation
// time when we have a model with thousands of tf2xla op fallbacks. At time
// of writing, this caused compilation time to be greater than 2x slower.
// So we have to directly import these instructions.
TF_ASSIGN_OR_RETURN(
mlir::Value root_value,
xla::HloFunctionImporter::ImportInstructions(
*hlo_module->entry_computation(), arguments, symbol_table, &builder));
stablehlo::TupleOp root_tuple =
mlir::dyn_cast_or_null<stablehlo::TupleOp>(root_value.getDefiningOp());
if (!root_tuple) {
return absl::InvalidArgumentError(
"Imported XLA Root Value is not a tuple op");
}
return root_tuple;
}
LogicalResult Tf2XlaRewriter::PrepareParams() {
// XlaCompiler within the context is only used by the functional ops to
// compile functions. We are not handling those at the moment so
// XlaCompiler is not required.
context_ = new tensorflow::XlaContext(/*compiler=*/nullptr, &xla_builder_,
/*graph=*/nullptr);
context_->Ref();
device_mgr_ = CreateDeviceMgr(device_type_);
if (!device_mgr_) return failure();
// Type of params_.device is DeviceBase* so store it as Device* to access
// derived class method.
device_ = device_mgr_->ListDevices().front();
params_.device = device_;
params_.resource_manager = device_->resource_manager();
// Resources are cleared at the time of device manager destruction so pass
// no-op cleanup function.
auto cleanup = [](const std::string& name) {};
// Use step_id zero as we only have a single context concurrently and
// concurrently running each of the MLIR functions create a new device.
step_container_ = std::make_unique<tensorflow::ScopedStepContainer>(
/*step_id=*/0, cleanup);
absl::Status status = step_container_->Create(
device_->resource_manager(),
tensorflow::XlaContext::kXlaContextResourceName, context_);
if (!status.ok()) {
return emitRemark(op_->getLoc())
<< "failed to create XlaContext resource: " << status.ToString();
}
params_.step_container = step_container_.get();
absl::StatusOr<int64_t> version_or = tensorflow::GetTfGraphProducerVersion(
op_->getParentOfType<mlir::ModuleOp>());
if (!version_or.ok()) {
return emitError(op_->getLoc()) << version_or.status().ToString();
}
flib_def_ = std::make_unique<tensorflow::FunctionLibraryDefinition>(
tensorflow::OpRegistry::Global(), tensorflow::FunctionDefLibrary());
pflr_ = std::make_unique<tensorflow::ProcessFunctionLibraryRuntime>(
device_mgr_.get(), tensorflow::Env::Default(), /*config=*/nullptr,
version_or.value(), flib_def_.get(), tensorflow::OptimizerOptions());
params_.function_library = pflr_->GetFLR(device_->name());
return success();
}
// Returns true if the given type is a ranked tensor type with static or
// bounded dimensions.
bool IsBounded(Type ty) {
auto ranked_ty = mlir::dyn_cast<RankedTensorType>(ty);
if (!ranked_ty) return false;
if (ranked_ty.hasStaticShape()) return true;
ArrayRef<int64_t> bounds = hlo::encodingToBounds(ranked_ty.getEncoding());
if (bounds.empty()) return false;
for (int i = 0; i < ranked_ty.getRank(); ++i) {
if (ranked_ty.isDynamicDim(i) && bounds[i] == ShapedType::kDynamic) {
return false;
}
}
return true;
}
bool HasSymbolRefAttr(Operation* op) {
for (const auto& attr : op->getAttrs()) {
Attribute attr_value = attr.getValue();
if (mlir::isa<SymbolRefAttr>(attr_value)) {
return true;
} else if (auto array_attr = mlir::dyn_cast<ArrayAttr>(attr_value)) {
if (!array_attr.empty() &&
mlir::isa<SymbolRefAttr>(*array_attr.begin())) {
return true;
}
}
}
return false;
}
LogicalResult Tf2XlaRewriter::PrepareKernelInputs(
const llvm::SmallDenseSet<int>& required_consts,
std::vector<tensorflow::XlaExpression>& expressions,
std::vector<tensorflow::Tensor>& tensors,
std::vector<tensorflow::TensorValue>& inputs) {
// Prepare the list of Tensor inputs for the kernel.
for (auto it : llvm::enumerate(op_->getOperands())) {
Value operand = it.value();
size_t idx = it.index();
tensorflow::XlaExpression expr = GetExprForOperand(operand, op_, idx);
tensorflow::XlaExpression::Kind kind = expr.kind();
if (kind == tensorflow::XlaExpression::Kind::kInvalid) return failure();
expressions.push_back(expr);
if (!tensorflow::DataTypeCanUseMemcpy(expr.dtype())) {
return op_->emitRemark()
<< "skipping legalization due to unsupported type "
<< operand.getType();
}
auto shape_or = expr.GetShape();
if (!shape_or.ok()) {
return op_->emitRemark()
<< "failed to get shape for expression. " << expr.HumanString();
}
tensors.emplace_back(
device_->GetAllocator(tensorflow::AllocatorAttributes()), expr.dtype(),
shape_or.value());
tensorflow::Tensor& tensor = tensors.back();
tensorflow::XlaExpression::AssignExpressionToTensor(expr, &tensor);
inputs.emplace_back(&tensor);
}
return success();
}
LogicalResult Tf2XlaRewriter::LegalizeOp() {
for (Type ty : op_->getOperandTypes()) {
auto ranked_ty = mlir::dyn_cast<ShapedType>(ty);
// Only bounded operands are supported in the XLA builders.
if (!IsBounded(ranked_ty)) {
return op_->emitRemark()
<< "lowering requires bounded tensor operands " << ranked_ty;
}
}
if (HasSymbolRefAttr(op_)) {
return op_->emitRemark() << "ops with symbol references are not supported";
}
auto nodedef_or = tensorflow::ConvertTFDialectOpToNodeDef(
op_, name_mapper_->GetUniqueName(op_),
/*ignore_unregistered_attrs=*/true);
if (!nodedef_or.ok()) {
return op_->emitRemark() << "failed to convert op to NodeDef: "
<< nodedef_or.status().ToString();
}
if (failed(PrepareParams())) return failure();
std::shared_ptr<const tensorflow::NodeProperties> props;
absl::Status status = tensorflow::NodeProperties::CreateFromNodeDef(
*nodedef_or.value(),
params_.function_library->GetFunctionLibraryDefinition(), &props);
if (!status.ok()) {
return op_->emitRemark()
<< "failed to create NodeProperties: " << status.ToString();
}
tensorflow::OpKernel* op_kernel_raw;
status = params_.function_library->CreateKernel(props, &op_kernel_raw);
if (!status.ok()) {
return op_->emitRemark()
<< "failed to create tf2xla kernel: " << status.ToString();
}
// Transfer ownership of the kernel to a local smart pointer.
auto op_kernel = absl::WrapUnique(op_kernel_raw);
std::vector<int> required_constants;
status = tensorflow::XlaOpRegistry::CompileTimeConstantInputs(
*op_kernel, &required_constants);
if (!status.ok()) {
return op_->emitRemark()
<< "failed to compute required constants: " << status.ToString();
}
llvm::SmallDenseSet<int> required_consts;
required_consts.insert(required_constants.begin(), required_constants.end());
// TensorValue in inputs are backed by tensors which in turn depend on
// expressions. So, pre-allocate them to the required size. Subtle note:
// Since these are assigned to params_, these have to live past the kernel
// compilation.
std::vector<tensorflow::XlaExpression> expressions;
std::vector<tensorflow::Tensor> tensors;
std::vector<tensorflow::TensorValue> inputs;
expressions.reserve(op_->getNumOperands());
tensors.reserve(op_->getNumOperands());
inputs.reserve(op_->getNumOperands());
if (failed(
PrepareKernelInputs(required_consts, expressions, tensors, inputs)))
return failure();
params_.inputs = inputs;
params_.op_kernel = op_kernel.get();
llvm::SmallVector<tensorflow::AllocatorAttributes, 4> output_attr(
op_->getNumResults());
params_.output_attr_array = output_attr.data();
tensorflow::OpKernelContext op_context(&params_, op_->getNumResults());
device_->Compute(params_.op_kernel, &op_context);
status = op_context.status();
if (!status.ok()) {
return op_->emitRemark()
<< "compilation to HLO failed: " << status.ToString();
}
if (failed(VerifyOpResults(op_context))) return failure();
absl::StatusOr<stablehlo::TupleOp> tuple_result_or_status =
CompileWithHloImporter(op_context);
if (!tuple_result_or_status.ok()) {
return op_->emitRemark() << tuple_result_or_status.status().ToString();
}
stablehlo::TupleOp tuple_result = tuple_result_or_status.value();
llvm::SmallVector<Value> output_values;
if (failed(GetKernelOutputs(op_context, tuple_result, output_values))) {
return failure();
}
rewriter_.replaceOp(op_, output_values);
return success();
}
absl::StatusOr<stablehlo::TupleOp> Tf2XlaRewriter::CompileWithHloImporter(
tensorflow::OpKernelContext& op_context) {
// XLA can only return a single value. Wrap all output op return values
// in a Tuple op that gets unpacked later.
std::vector<xla::XlaOp> output_values;
for (int i = 0, e = op_->getNumResults(); i < e; i++) {
tensorflow::Tensor* output = op_context.mutable_output(i);
const tensorflow::XlaExpression* expr =
tensorflow::XlaExpression::CastExpressionFromTensor(*output);
output_values.push_back(expr->AsXlaOp(&xla_builder_));
}
absl::Span<const xla::XlaOp> return_values(output_values);
xla::XlaOp root_value = xla::Tuple(&xla_builder_, return_values);
TF_ASSIGN_OR_RETURN(XlaComputation computation,
xla_builder_.Build(root_value,
/*remove_dynamic_dimensions=*/false));
return ImportXlaComputation(computation);
}
mlir::LogicalResult Tf2XlaRewriter::VerifyOpResults(
tensorflow::OpKernelContext& op_context) {
for (int i = 0, e = op_->getNumResults(); i < e; i++) {
tensorflow::Tensor* output = op_context.mutable_output(i);
const tensorflow::XlaExpression* expr =
tensorflow::XlaExpression::CastExpressionFromTensor(*output);
if (expr->kind() != tensorflow::XlaExpression::Kind::kXlaOp &&
expr->kind() != tensorflow::XlaExpression::Kind::kConstant) {
return op_->emitRemark(absl::StrCat(
"expects XlaExpression of kind kXlaOp or kConstant in compiled "
"output index ",
i));
}
}
return success();
}
// XLA computations can only return a single value, but TF ops can return
// multiple values. We get around this by returning a tuple as an XLA op. We
// then unpack it here to return the multiple values instead.
mlir::LogicalResult Tf2XlaRewriter::UnpackTupleResults(
stablehlo::TupleOp tuple_result, llvm::SmallVector<Value>& outputs) {
if (tuple_result->getNumOperands() != op_->getNumResults()) {
return op_->emitRemark() << "Translated TF2XLA tuple has different "
"number of results than original op";
}
for (int i = 0; i < tuple_result->getNumOperands(); i++) {
outputs.push_back(tuple_result->getOperand(i));
}
tuple_result.getOperation()->erase();
return success();
}
mlir::LogicalResult Tf2XlaRewriter::GetKernelOutputs(
tensorflow::OpKernelContext& op_context, stablehlo::TupleOp tuple_results,
llvm::SmallVector<Value>& outputs) {
outputs.reserve(op_->getNumResults());
return UnpackTupleResults(tuple_results, outputs);
}
tensorflow::XlaExpression Tf2XlaRewriter::GetExprForOperand(
Value operand, Operation* op, int64_t operand_index) {
ElementsAttr const_attr;
auto defining_op = operand.getDefiningOp();
::xla::XlaOp xla_op = xla::Parameter(&xla_builder_, operand_index,
xla::TypeToShape(operand.getType()),
std::to_string(operand_index));
if (defining_op && matchPattern(defining_op, m_Constant(&const_attr))) {
tensorflow::Tensor tensor;
auto status = tensorflow::ConvertToTensor(const_attr, &tensor);
if (!status.ok()) {
op->emitRemark() << "skipping legalization due to failed const conversion"
<< status.ToString();
return tensorflow::XlaExpression::Invalid();
}
return tensorflow::XlaExpression::Constant(tensor);
}
tensorflow::DataType dtype;
auto status = tensorflow::ConvertToDataType(operand.getType(), &dtype);
if (!status.ok()) {
op->emitRemark() << "skipping legalization due to " << status.ToString();
return tensorflow::XlaExpression::Invalid();
}
return tensorflow::XlaExpression::XlaOp(xla_op, dtype);
}
} // namespace hlo
} // namespace mlir
tensorflow/compiler/mlir/tf2xla/transforms/tf2xla_rewriter.h
+127
View File
@@ -0,0 +1,127 @@
/* 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_TRANSFORMS_TF2XLA_REWRITER_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_TRANSFORMS_TF2XLA_REWRITER_H_
#include <cstdint>
#include <memory>
#include <string>
#include <vector>
#include "absl/status/statusor.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "stablehlo/dialect/StablehloOps.h" // from @stablehlo
#include "tensorflow/compiler/mlir/op_or_arg_name_mapper.h"
#include "tensorflow/compiler/tf2xla/xla_context.h"
#include "tensorflow/compiler/tf2xla/xla_expression.h"
#include "xla/hlo/builder/xla_builder.h"
#include "xla/hlo/builder/xla_computation.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/framework/op_kernel.h"
namespace mlir {
namespace hlo {
class Tf2XlaRewriterTestPeer;
class Tf2XlaRewriter {
public:
static mlir::LogicalResult RewriteOp(mlir::Operation* op,
mlir::PatternRewriter& rewriter,
const std::string& device_type);
private:
friend class Tf2XlaRewriterTestPeer;
Tf2XlaRewriter(mlir::Operation* op, mlir::PatternRewriter& rewriter,
const std::string& device_type);
~Tf2XlaRewriter();
// Compiles the given Operation with XlaBuilder and imports the generated HLO
// via the HLO -> MHLO importer.
absl::StatusOr<stablehlo::TupleOp> CompileWithHloImporter(
tensorflow::OpKernelContext& op_context);
// Import the given XlaComputation into the parent module. Returns the given
// generated function.
absl::StatusOr<stablehlo::TupleOp> ImportXlaComputation(
xla::XlaComputation& computation);
// Prepares OpKernelContext params common to all the ops.
// Emits an error on failure.
mlir::LogicalResult PrepareParams();
// Given the required_consts, it will fill the 3 output vectors with
// their respective data.
// Expressions: Output XLA expressions as required by the compiled kernel.
// Tensors: Vector of tensors that back the TensorValue inputs
// Inputs: Vector of inputs that are backed by tensors.
mlir::LogicalResult PrepareKernelInputs(
const llvm::SmallDenseSet<int>& required_consts,
std::vector<tensorflow::XlaExpression>& expressions,
std::vector<tensorflow::Tensor>& tensors,
std::vector<tensorflow::TensorValue>& inputs);
mlir::LogicalResult VerifyOpResults(tensorflow::OpKernelContext& op_context);
mlir::LogicalResult GetKernelOutputs(tensorflow::OpKernelContext& op_context,
stablehlo::TupleOp tuple_results,
llvm::SmallVector<Value>& outputs);
// Given a translated function with a single return value, unpack the tuple
// results.
mlir::LogicalResult UnpackTupleResults(stablehlo::TupleOp tuple_result,
llvm::SmallVector<Value>& outputs);
// Tries to legalize the specified TensorFlow op, if supported.
//
// Emits an error and returns failure if an error is encountered during
// conversion. Note that success return value doesn't mean successful
// legalization.
mlir::LogicalResult LegalizeOp();
// Converts the given operand to expression of kind kConstant or kXlaOp.
// Emits a remark and returns expression of kind kInvalid on failure.
tensorflow::XlaExpression GetExprForOperand(mlir::Value operand,
mlir::Operation* op,
int64_t operand_index);
mlir::Operation* op_;
std::string device_type_;
mlir::PatternRewriter& rewriter_;
std::unique_ptr<tensorflow::OpOrArgLocNameMapper> name_mapper_;
tensorflow::XlaContext* context_; // Ref-counted.
std::unique_ptr<tensorflow::StaticDeviceMgr> device_mgr_;
tensorflow::Device* device_; // Owned by device_mgr_;
std::unique_ptr<tensorflow::ScopedStepContainer> step_container_;
std::unique_ptr<tensorflow::FunctionLibraryDefinition> flib_def_;
std::unique_ptr<tensorflow::ProcessFunctionLibraryRuntime> pflr_;
tensorflow::OpKernelContext::Params params_;
xla::XlaBuilder xla_builder_;
};
} // namespace hlo
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_TRANSFORMS_TF2XLA_REWRITER_H_
tensorflow/compiler/mlir/tf2xla/transforms/tf2xla_rewriter_test.cc
+366
View File
@@ -0,0 +1,366 @@
/* 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/transforms/tf2xla_rewriter.h"
#include <memory>
#include <string>
#include <vector>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/Diagnostics.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Transforms/DialectConversion.h" // from @llvm-project
#include "stablehlo/dialect/StablehloOps.h" // from @stablehlo
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/test_utils.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "xla/hlo/builder/xla_builder.h"
#include "xla/hlo/builder/xla_computation.h"
#include "xla/shape_util.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/status.h"
#include "xla/tsl/platform/statusor.h"
#include "xla/xla_data.pb.h"
namespace mlir {
namespace hlo {
using ::mlir::LogicalResult;
using ::mlir::ModuleOp;
using ::mlir::OpBuilder;
using ::mlir::Operation;
using ::mlir::func::FuncOp;
using ::tsl::Status;
using ::tsl::StatusOr;
using ::xla::ReplicaGroup;
using ::xla::ShapeUtil;
using ::xla::XlaBuilder;
using ::xla::XlaComputation;
using ::xla::XlaOp;
static constexpr char kMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1442 : i32}} {
func.func @main(%arg0: tensor<3xi64> {tf._user_specified_name = "resource", tf.aliasing_output = 3 : i64}) -> () attributes {tf.entry_function = {control_outputs = "stateful_normal/RngReadAndSkip,stateful_uniform/RngReadAndSkip,stateful_uniform_full_int/RngReadAndSkip", inputs = "stateful_normal_rngreadandskip_resource", outputs = "identity_RetVal,identity_1_RetVal,identity_2_RetVal"}} {
%0:3 = "tf.Unpack"(%arg0) {axis = 0 : i64} : (tensor<3xi64>) -> (tensor<i64>, tensor<i64>, tensor<i64>)
return
}
})";
XlaComputation GetTestXlaComputation() {
XlaBuilder xla_builder("test");
auto param =
Parameter(&xla_builder, 0, ShapeUtil::MakeScalarShape(xla::F32), "a");
XlaOp add = xla::Add(param, xla::ConstantR0<float>(&xla_builder, 2.0));
std::vector<XlaOp> tuple_values;
tuple_values.push_back(add);
xla::Tuple(&xla_builder, tuple_values);
return xla_builder.Build().value();
}
class EmptyPatternRewriter : public mlir::PatternRewriter {
public:
explicit EmptyPatternRewriter(const OpBuilder& other_builder)
: mlir::PatternRewriter(other_builder) {}
~EmptyPatternRewriter() override = default;
};
class Tf2XlaRewriterTestPeer {
public:
explicit Tf2XlaRewriterTestPeer() = delete;
explicit Tf2XlaRewriterTestPeer(mlir::Operation* op)
: op_builder_(op),
empty_rewriter_(op_builder_),
tf2xla_rewriter_(op, empty_rewriter_,
/*device_type=*/"XLA_CPU_JIT") {}
absl::StatusOr<stablehlo::TupleOp> ImportXlaComputationIntoModule(
XlaComputation& computation) {
return tf2xla_rewriter_.ImportXlaComputation(computation);
}
private:
OpBuilder op_builder_;
EmptyPatternRewriter empty_rewriter_;
Tf2XlaRewriter tf2xla_rewriter_;
};
// This should only have unit tests. End to end tests should be done with
// FILECHECK and MLIR tests.
class Tf2XlaRewriterTest : public ::testing::Test {
public:
void SetUp() override {
tensorflow::XlaOpRegistry::RegisterCompilationKernels();
}
Status CreateMlirModule(std::string module_string = kMlirModuleStr) {
TF_ASSIGN_OR_RETURN(
module_, hlo::test::GetMlirModuleFromString(module_string, &context_));
context_.loadAllAvailableDialects();
return absl::OkStatus();
}
Status LegalizeSingleOp(Operation& op) {
SourceMgrDiagnosticHandler sourceMgrHandler(source_manager_, &context_);
OpBuilder op_builder(&op);
EmptyPatternRewriter pattern_rewriter(op_builder);
LogicalResult result =
Tf2XlaRewriter::RewriteOp(&op, pattern_rewriter,
/*device_type=*/"XLA_CPU_JIT");
if (!result.succeeded()) {
return absl::InternalError("Failed to rewrite op");
}
return absl::OkStatus();
}
Status LegalizeModule(std::string module_string = kMlirModuleStr) {
TF_EXPECT_OK(CreateMlirModule(module_string));
FuncOp main = module_->lookupSymbol<mlir::func::FuncOp>("main");
if (!main) {
return absl::InvalidArgumentError("Could not find a main function");
}
WalkResult walk_result = main.walk([&](Operation* op) {
if (op->getDialect()->getNamespace() !=
TF::TensorFlowDialect::getDialectNamespace()) {
return WalkResult::advance();
}
if (!LegalizeSingleOp(*op).ok()) {
return WalkResult::interrupt();
}
return WalkResult::advance();
});
if (walk_result.wasInterrupted()) {
return absl::InternalError("Could not legalize all ops");
}
return absl::OkStatus();
}
mlir::func::FuncOp GetMainFunc() {
func::FuncOp main_func = module_->lookupSymbol<mlir::func::FuncOp>("main");
EXPECT_TRUE(main_func);
return main_func;
}
mlir::Operation& GetFirstOpFromMain() {
mlir::func::FuncOp main_func = GetMainFunc();
return main_func.getBody().front().front();
}
absl::StatusOr<stablehlo::TupleOp> ImportXlaComputationIntoModule(
XlaComputation& computation) {
SourceMgrDiagnosticHandler sourceMgrHandler(source_manager_, &context_);
mlir::Operation& first_op = GetFirstOpFromMain();
Tf2XlaRewriterTestPeer test_peer(&first_op);
return test_peer.ImportXlaComputationIntoModule(computation);
}
protected:
MLIRContext context_;
OwningOpRef<ModuleOp> module_;
llvm::SourceMgr source_manager_;
};
TEST_F(Tf2XlaRewriterTest, LegalizesOpWithTf2xlaHloImporter) {
TF_EXPECT_OK(LegalizeModule());
int num_tuple_ops = 0;
module_->walk(
[&num_tuple_ops](stablehlo::TupleOp tuple_op) { num_tuple_ops += 1; });
EXPECT_EQ(num_tuple_ops, 0);
}
TEST_F(Tf2XlaRewriterTest, ImportsXlaComputationIntoModule) {
TF_ASSERT_OK(CreateMlirModule());
XlaComputation computation = GetTestXlaComputation();
TF_ASSERT_OK_AND_ASSIGN(stablehlo::TupleOp root_tuple,
ImportXlaComputationIntoModule(computation));
ModuleOp parent_module =
root_tuple.getOperation()->getParentOfType<ModuleOp>();
EXPECT_EQ(parent_module, *module_);
}
TEST_F(Tf2XlaRewriterTest, FailsWithoutRootTuple) {
TF_ASSERT_OK(CreateMlirModule());
XlaBuilder xla_builder("test_fail");
xla::Add(xla::ConstantR0<float>(&xla_builder, 1.0),
xla::ConstantR0<float>(&xla_builder, 2.0));
XlaComputation bad_computation = xla_builder.Build().value();
EXPECT_FALSE(ImportXlaComputationIntoModule(bad_computation).ok());
}
TEST_F(Tf2XlaRewriterTest, ImportsSingleComputation) {
XlaBuilder builder("test_builder");
XlaComputation to_apply;
{
auto sub_builder = builder.CreateSubBuilder("add");
auto arg0 = Parameter(sub_builder.get(), 0,
ShapeUtil::MakeScalarShape(xla::F32), "x");
auto arg1 = Parameter(sub_builder.get(), 1,
ShapeUtil::MakeScalarShape(xla::F32), "y");
Add(arg0, arg1);
TF_ASSERT_OK_AND_ASSIGN(to_apply, sub_builder->Build());
}
auto x = Parameter(&builder, 0, ShapeUtil::MakeShape(xla::F32, {4, 16}), "x");
ReplicaGroup group;
group.add_replica_ids(0);
group.add_replica_ids(1);
XlaOp reduce_scatter =
ReduceScatter(x, to_apply, /*scatter_dimension=*/1, /*shard_count=*/2,
/*replica_groups=*/{group});
std::vector<XlaOp> tuple_values;
tuple_values.push_back(reduce_scatter);
xla::Tuple(&builder, tuple_values);
TF_ASSERT_OK_AND_ASSIGN(XlaComputation computation, builder.Build());
EXPECT_EQ(computation.proto().computations_size(), 2);
TF_ASSERT_OK(CreateMlirModule());
TF_ASSERT_OK_AND_ASSIGN(stablehlo::TupleOp root_tuple,
ImportXlaComputationIntoModule(computation));
EXPECT_TRUE(root_tuple);
int num_func_ops = 0;
module_->walk([&num_func_ops](func::FuncOp func_op) { num_func_ops++; });
// Ensure that only a single computation was imported.
EXPECT_EQ(num_func_ops, 1);
}
TEST_F(Tf2XlaRewriterTest, InsertsConstantParameters) {
static constexpr char kModuleWithConstParam[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1442 : i32}} {
func.func @main(%arg0: tensor<2xf32>) -> tensor<2xf32> {
%0 = "tf.Const"() {value = dense<1.42> : tensor<2xf32>} : () -> tensor<2xf32>
%1 = "tf.Atan2"(%arg0, %0) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>
func.return %0 : tensor<2xf32>
}
})";
TF_ASSERT_OK(LegalizeModule(kModuleWithConstParam));
}
TEST_F(Tf2XlaRewriterTest, DoesntEnforceCompileTimeConstantCheck) {
static constexpr char kModuleWithNonConstParam[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1610 : i32}} {
func.func @main(%arg0: tensor<3x3x10xbf16>, %arg1: tensor<3xi32>) -> tensor<1x?x4xbf16> attributes {allow_soft_placement = false, tf.entry_function = {control_outputs = "", inputs = "_arg0,_arg1,_arg2", outputs = "_retval0"}} {
%cst = "tf.Const"() {value = dense<[1, -1, 4]> : tensor<3xi32>} : () -> tensor<3xi32>
%0 = "tf.Slice"(%arg0, %arg1, %cst) {_XlaHasReferenceVars = false, _xla_inferred_shapes = [#tf_type.shape<1x?x4>], device = "/job:localhost/replica:0/task:0/device:TPU:0"} : (tensor<3x3x10xbf16>, tensor<3xi32>, tensor<3xi32>) -> tensor<1x?x4xbf16>
return %0 : tensor<1x?x4xbf16>
}
})";
TF_ASSERT_OK(LegalizeModule(kModuleWithNonConstParam));
}
TEST_F(Tf2XlaRewriterTest, CreatesDefaultValues) {
// If a TF op has default value attributes and the mlir is missing them then
// the LegalizeOp should insert the default values when converting the dialect
// op to a node def.
// TF.RandomUniform would fail without the seeds being set if they were not
// automatically inserted with the default values.
static constexpr char kModuleWithOpWithoutValuesThatShouldBeDefaulted[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1610 : i32}} {
func.func @main() -> tensor<1x2x3x4xf32> attributes {allow_soft_placement = false, tf.entry_function = {control_outputs = "", inputs = "_arg0,_arg1,_arg2", outputs = "_retval0"}} {
%cst = "tf.Const"() {value = dense<[1, 2, 3, 4]> : tensor<4xi32>} : () -> tensor<4xi32>
%0 = "tf.RandomUniform"(%cst) : (tensor<4xi32>) -> tensor<1x2x3x4xf32>
return %0 : tensor<1x2x3x4xf32>
}
})";
TF_ASSERT_OK(LegalizeModule(kModuleWithOpWithoutValuesThatShouldBeDefaulted));
}
TEST_F(Tf2XlaRewriterTest, OpWithLocationDoesntBreakNodeDefName) {
// A named location 'Name(Source)' causes the GetNameFromLoc method to append
// all the other locations to the name with a ';' separator. This test ensures
// that the name used for the NodeDef does not contain that invalid character.
static constexpr char kModuleWithOpWithoutValuesThatShouldBeDefaulted[] =
R"mlir(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1610 : i32}} {
func.func @main(%arg0: tensor<2xf32>) -> tensor<2xf32> {
%0 = "tf.Exp"(%arg0) : (tensor<2xf32>) -> tensor<2xf32> loc(fused["exp"("exp"), "exp"])
func.return %0 : tensor<2xf32>
}
})mlir";
TF_ASSERT_OK(LegalizeModule(kModuleWithOpWithoutValuesThatShouldBeDefaulted));
}
TEST_F(Tf2XlaRewriterTest, ErrorsWithInvalidNumberOfParametersToArgs) {
XlaBuilder builder("test_builder");
XlaComputation to_apply;
{
auto sub_builder = builder.CreateSubBuilder("add");
auto arg0 = Parameter(sub_builder.get(), 0,
ShapeUtil::MakeScalarShape(xla::F32), "x");
auto arg1 = Parameter(sub_builder.get(), 1,
ShapeUtil::MakeScalarShape(xla::F32), "y");
Add(arg0, arg1);
TF_ASSERT_OK_AND_ASSIGN(to_apply, sub_builder->Build());
}
auto a = Parameter(&builder, 0, ShapeUtil::MakeScalarShape(xla::F32), "a");
auto b = Parameter(&builder, 1, ShapeUtil::MakeScalarShape(xla::F32), "b");
XlaOp call_op = xla::Call(&builder, to_apply, {a, b});
std::vector<XlaOp> tuple_values;
tuple_values.push_back(call_op);
xla::Tuple(&builder, tuple_values);
TF_ASSERT_OK_AND_ASSIGN(XlaComputation computation, builder.Build());
EXPECT_EQ(computation.proto().computations_size(), 2);
TF_ASSERT_OK(CreateMlirModule());
absl::StatusOr<stablehlo::TupleOp> status_or_tuple_op =
ImportXlaComputationIntoModule(computation);
EXPECT_FALSE(status_or_tuple_op.ok());
}
} // namespace hlo
} // namespace mlir
tensorflow/compiler/mlir/tf2xla/transforms/tf_xla_passes.td
+47
View File
@@ -0,0 +1,47 @@
/* Copyright 2021 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 LegalizeTFCommunicationPass : Pass<"xla-legalize-tf-communication", "ModuleOp"> {
let summary = "Legalize TF/XLA communication ops (TensorFlow dialect) to the HLO "
"dialect";
let dependentDialects = ["mhlo::MhloDialect", "sparse_tensor::SparseTensorDialect"];
let constructor = "::mlir::mhlo::CreateLegalizeTFCommunicationPass()";
let description = [{
A pass that legalizes TF/XLA communication ops, propagates their respective
tokens (for ordering), and rewrites their respective functions and control
flow ops when necessary.
For example, given the program
```mlir
func @send_to_host(%arg0: tensor<i32>) {
"tf.XlaSendToHost"(%arg0) {key = "send_key"} : (tensor<i32>) -> ()
return
}
```
This might be legalized like this:
```mlir
func @send_to_host(%arg0: tensor<i32>) {
%0 = "mhlo.create_token"() : () -> !mhlo.token
%1 = "mhlo.send"(%arg0, %0) {...}} : (tensor<i32>, !mhlo.token) -> !mhlo.token
return
}
'''
}];
}
tensorflow/compiler/mlir/tf2xla/transforms/tfxla_device_specific_transforms.cc
+94
View File
@@ -0,0 +1,94 @@
/* 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 "mlir/IR/BuiltinOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/PatternMatch.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/mlir/tf2xla/transforms/passes.h"
#include "tensorflow/compiler/tf2xla/kernels/rng_converter_utils.h"
#include "xla/xla_data.pb.h"
namespace mlir {
namespace mhlo {
namespace {
#define GEN_PASS_DEF_TFXLADEVICESPECIFICTRANSFORMS
#include "tensorflow/compiler/mlir/tf2xla/transforms/xla_legalize_tf_passes.h.inc"
class TFXLADeviceSpecificTransforms
: public impl::TFXLADeviceSpecificTransformsBase<
TFXLADeviceSpecificTransforms> {
public:
explicit TFXLADeviceSpecificTransforms(std::optional<StringRef> device_type) {
if (device_type.has_value()) {
device_type_ = device_type.value().str();
}
}
void runOnOperation() override;
private:
LogicalResult ConvertGetAlgOp(TF::StatelessRandomGetAlgOp get_alg_op);
};
LogicalResult TFXLADeviceSpecificTransforms::ConvertGetAlgOp(
TF::StatelessRandomGetAlgOp get_alg_op) {
if (!device_type_.hasValue()) return failure();
xla::RandomAlgorithm xla_rng =
tensorflow::DefaultRngAlgForDeviceType(device_type_);
tensorflow::Algorithm tensorflow_rng =
tensorflow::ToTensorflowAlgorithm(xla_rng);
OpBuilder opbuilder(get_alg_op);
auto tf_const =
TF::ConstOp::create(opbuilder, get_alg_op->getLoc(),
opbuilder.getI32IntegerAttr((int)tensorflow_rng));
get_alg_op->replaceAllUsesWith(tf_const);
get_alg_op->erase();
return success();
}
void TFXLADeviceSpecificTransforms::runOnOperation() {
if (!device_type_.hasValue()) return;
auto func_op = getOperation();
auto walk_result = func_op->walk([&](TF::StatelessRandomGetAlgOp op) {
if (failed(ConvertGetAlgOp(op))) {
op->emitOpError(
"Could not convert and remove Device specific information");
return WalkResult::interrupt();
}
return WalkResult::advance();
});
if (walk_result.wasInterrupted()) signalPassFailure();
}
} // namespace
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateTFXLADeviceSpecificTransformsPass(std::optional<StringRef> device_type) {
return std::make_unique<TFXLADeviceSpecificTransforms>(device_type);
}
} // namespace mhlo
} // namespace mlir
tensorflow/compiler/mlir/tf2xla/transforms/utils.cc
+49
View File
@@ -0,0 +1,49 @@
/* Copyright 2021 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/transforms/utils.h"
#include <cstdint>
#include "xla/mlir_hlo/utils/hlo_utils.h"
namespace mlir {
namespace mhlo {
ConstantOp GetScalarConstOfType(Type ty, Location loc, int64_t raw_value,
OpBuilder* builder) {
return ConstantOp::create(*builder, loc, hlo::getScalarOfType(ty, raw_value));
}
ConstantOp GetScalarNegZeroOfType(Type ty, Location loc, OpBuilder* builder) {
return ConstantOp::create(*builder, loc, hlo::getScalarNegZeroOfType(ty));
}
DenseIntElementsAttr GetI64ElementsAttr(ArrayAttr attr) {
RankedTensorType ty =
RankedTensorType::get(static_cast<int64_t>(attr.size()),
IntegerType::get(attr.getContext(), 64));
return DenseIntElementsAttr::get(ty, attr.getValue());
}
DenseIntElementsAttr GetI64ElementsAttr(ArrayRef<int64_t> values,
Builder* builder) {
RankedTensorType ty = RankedTensorType::get(
{static_cast<int64_t>(values.size())}, builder->getIntegerType(64));
return DenseIntElementsAttr::get(ty, values);
}
} // namespace mhlo
} // namespace mlir
tensorflow/compiler/mlir/tf2xla/transforms/utils.h
+62
View File
@@ -0,0 +1,62 @@
/* Copyright 2021 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_TRANSFORMS_UTILS_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_TRANSFORMS_UTILS_H_
#include <cstdint>
#include "llvm/ADT/ArrayRef.h"
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Location.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
namespace mlir {
namespace mhlo {
// Builds body for reduce op by using the template binary op as the
// reducer op.
template <typename Op>
void BuildReduceBody(Type element_type, Region* body, OpBuilder* builder) {
OpBuilder::InsertionGuard guard(*builder);
Block* block = builder->createBlock(body);
// Block arguments are scalars of the given element type.
Type type = RankedTensorType::get(/*shape=*/{}, element_type);
Location loc = body->getLoc();
block->addArguments({type, type}, SmallVector<Location, 2>(2, loc));
auto reducer =
Op::create(*builder, loc, block->getArgument(0), block->getArgument(1));
ReturnOp::create(*builder, loc, reducer.getResult());
}
ConstantOp GetScalarConstOfType(Type ty, Location loc, int64_t raw_value,
OpBuilder* builder);
ConstantOp GetScalarNegZeroOfType(Type ty, Location loc, OpBuilder* builder);
// Converts an ArrayAttr to a 1D 64-bit dense elements attribute.
DenseIntElementsAttr GetI64ElementsAttr(ArrayAttr attr);
DenseIntElementsAttr GetI64ElementsAttr(llvm::ArrayRef<int64_t> values,
Builder* builder);
} // namespace mhlo
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_TRANSFORMS_UTILS_H_
tensorflow/compiler/mlir/tf2xla/transforms/verify_tfxla_legalization.cc
+180
View File
@@ -0,0 +1,180 @@
/* 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 "mlir/IR/BuiltinOps.h"
#include "absl/strings/str_cat.h"
#include "llvm/ADT/StringExtras.h"
#include "llvm/Support/Error.h"
#include "llvm/Support/FormatVariadic.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Diagnostics.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Transforms/DialectConversion.h" // from @llvm-project
#include "stablehlo/dialect/Base.h" // from @stablehlo
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/passes.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/xla_legalize_targets.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
#include "tensorflow/core/lib/monitoring/counter.h"
#include "tensorflow/core/platform/errors.h"
namespace mlir {
namespace mhlo {
namespace {
#define GEN_PASS_DEF_VERIFYTFXLALEGALIZATION
#include "tensorflow/compiler/mlir/tf2xla/transforms/xla_legalize_tf_passes.h.inc"
auto* mlir_failed_legalization_op_count =
tensorflow::monitoring::Counter<1>::New(
"/tensorflow/core/tf2xla/"
"mlir_second_phase_failed_legalization_op_count",
"Counts which op fails to legalize", "op_name");
auto* mlir_non_static_op_count = tensorflow::monitoring::Counter<1>::New(
"/tensorflow/core/tf2xla/"
"mlir_second_phase_non_static_op_count",
"Counts which ops do not have static results", "op_name");
auto* mlir_non_static_op_skip_count = tensorflow::monitoring::Counter<1>::New(
"/tensorflow/core/tf2xla/"
"mlir_second_phase_non_static_op_skip_count",
"Counts skipped ops which do not have static results", "op_name");
static const char* kMustBeConstantError =
"must have compile-time constant inputs and outputs.\n\n"
"XLA compilation requires that operator arguments that represent shapes or "
"dimensions be evaluated to concrete values at compile time. This error "
"means that a shape or dimension argument could not be evaluated at "
"compile time, usually because the value of the argument depends on a "
"parameter to the computation, on a variable, or on a stateful operation "
"such as a random number generator.";
// TODO(b/282188914) remove the operations to skip once tests are fixed.
static const DenseSet<mlir::TypeID>* operations_to_skip =
new DenseSet<mlir::TypeID>{mlir::TypeID::get<mhlo::EinsumOp>()};
class VerifyTFXLALegalization
: public impl::VerifyTFXLALegalizationBase<VerifyTFXLALegalization> {
public:
explicit VerifyTFXLALegalization(bool legalize_chlo) {
legalize_chlo_ = legalize_chlo;
}
void runOnOperation() override;
};
static void IncrementCounterFor(tensorflow::monitoring::Counter<1>* counter,
Operation* op) {
counter->GetCell(op->getName().getStringRef().str())->IncrementBy(1);
}
bool HasBounds(RankedTensorType type) {
auto bounds = hlo::encodingToBounds(type.getEncoding());
return !bounds.empty();
}
bool HasStaticShapeOrBounded(Value val) {
auto type = val.getType();
if (mlir::isa<UnrankedTensorType>(type)) {
return false;
}
if (mlir::isa<RankedTensorType>(type)) {
auto ranked_tensor = mlir::dyn_cast<RankedTensorType>(type);
if (ranked_tensor.hasStaticShape()) {
return true;
}
return HasBounds(ranked_tensor);
}
return true;
}
bool EmitMustBeConstantError(Operation* op) {
if (operations_to_skip->contains(op->getRegisteredInfo()->getTypeID())) {
IncrementCounterFor(mlir_non_static_op_skip_count, op);
return true;
}
emitError(op->getLoc()) << absl::StrCat(
"Node `", op->getName().getStringRef().str(), "` ", kMustBeConstantError);
return false;
}
bool IsStaticOperation(Operation* op) {
for (auto o : op->getResults()) {
if (!HasStaticShapeOrBounded(o)) {
return EmitMustBeConstantError(op);
}
}
return true;
}
bool IsMhloAndStatic(Operation* op) {
if (!llvm::isa<mlir::mhlo::MhloDialect>(op->getDialect())) {
// Skip this op if it isn't an mhlo op.
return true;
}
return IsStaticOperation(op);
}
bool IsDefaultConversionLegal(
Operation* op, const ConversionTarget& default_conversion_target) {
if (!default_conversion_target.isLegal(op)) {
emitError(op->getLoc()) << "Could not legalize op: " << op->getName();
return false;
}
return true;
}
void VerifyTFXLALegalization::runOnOperation() {
Operation* func_op = getOperation();
ConversionTarget default_conversion_target =
hlo::GetDefaultLegalConversionTargets(getContext(), legalize_chlo_);
bool has_invalid_ops = false;
func_op->walk([&](Operation* op) {
if (!IsMhloAndStatic(op)) {
has_invalid_ops = true;
IncrementCounterFor(mlir_non_static_op_count, op);
return WalkResult::interrupt();
}
if (!IsDefaultConversionLegal(op, default_conversion_target)) {
has_invalid_ops = true;
IncrementCounterFor(mlir_failed_legalization_op_count, op);
}
return WalkResult::advance();
});
if (has_invalid_ops) signalPassFailure();
}
} // namespace
} // namespace mhlo
namespace hlo {
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateVerifyTFXLALegalizationPass(bool legalize_chlo) {
return std::make_unique<mhlo::VerifyTFXLALegalization>(legalize_chlo);
}
} // namespace hlo
} // namespace mlir
tensorflow/compiler/mlir/tf2xla/transforms/verify_tfxla_legalization_test.cc
+180
View File
@@ -0,0 +1,180 @@
/* 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 "absl/strings/string_view.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/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/serialize_mlir_module_utils.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/passes.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/test_utils.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/monitoring/cell_reader.h"
namespace tensorflow {
namespace {
using ::mlir::MLIRContext;
using ::mlir::ModuleOp;
using ::mlir::OwningOpRef;
using ::mlir::hlo::test::GetMlirModuleFromString;
using ::tensorflow::monitoring::testing::CellReader;
static constexpr char kFailedLegalizationStreamz[] =
"/tensorflow/core/tf2xla/mlir_second_phase_failed_legalization_op_count";
static constexpr char kNonStaticOpStreamz[] =
"/tensorflow/core/tf2xla/mlir_second_phase_non_static_op_count";
static constexpr char kNonStaticOpSkipStreamz[] =
"/tensorflow/core/tf2xla/mlir_second_phase_non_static_op_skip_count";
class VerifyTfxlaLegalizationTest : 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>(
mlir::hlo::CreateVerifyTFXLALegalizationPass(/*legalize_chlo=*/false));
}
mlir::LogicalResult Run() { return pm_->run(module_.get()); }
private:
MLIRContext context_;
OwningOpRef<ModuleOp> module_;
std::unique_ptr<mlir::PassManager> pm_;
};
TEST_F(VerifyTfxlaLegalizationTest, RecordsStreamzFailedVerification) {
// Using a string constant here instead of testdata to make this compatible
// with open source.
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(kFailedLegalizationStreamz);
CreateModule(kMlirModuleStr);
auto result = Run();
EXPECT_TRUE(result.failed());
EXPECT_EQ(error.Delta("tf.BadValue"), 1);
}
TEST_F(VerifyTfxlaLegalizationTest, ErrorsNonStaticInputs) {
// Using a string constant here instead of testdata to make this compatible
// with open source.
static constexpr char kNonStaticFailure[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1504 : i32}} {
func.func @main() -> tensor<?xi32> attributes {tf.entry_function = {control_outputs = "", inputs = "i,j", outputs = "identity_RetVal"}} {
%0 = mhlo.constant dense<1.000000e+00> : tensor<f64>
%1 = mhlo.convert %0 : (tensor<f64>) -> tensor<i64>
%2 = mhlo.reshape %1 : (tensor<i64>) -> tensor<1xi64>
%3 = "mhlo.dynamic_iota"(%2) {iota_dimension = 0 : i64} : (tensor<1xi64>) -> tensor<?xi32>
%4 = mhlo.multiply %3, %3 : tensor<?xi32>
return %4 : tensor<?xi32>
}
})";
CellReader<int64_t> legal_error(kFailedLegalizationStreamz);
CellReader<int64_t> static_error(kNonStaticOpStreamz);
CreateModule(kNonStaticFailure);
auto result = Run();
EXPECT_TRUE(result.failed());
EXPECT_EQ(legal_error.Delta("mhlo.dynamic_iota"), 0);
EXPECT_EQ(static_error.Delta("mhlo.dynamic_iota"), 1);
}
TEST_F(VerifyTfxlaLegalizationTest, SkipsSpecificNonStaticInputs) {
// Using a string constant here instead of testdata to make this compatible
// with open source.
static constexpr char kNonStaticFailure[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1504 : i32}} {
func.func @main(%a : tensor<5x14x1xf32>, %b : tensor<1x14x32xf32>) -> tensor<?x?x?xf32> attributes {tf.entry_function = {control_outputs = "", inputs = "i,j", outputs = "identity_RetVal"}} {
%c = "mhlo.einsum"(%a, %b) {einsum_config = "bji,bjk->bik"} : (tensor<5x14x1xf32>, tensor<1x14x32xf32>) -> tensor<?x?x?xf32>
return %c : tensor<?x?x?xf32>
}
})";
CellReader<int64_t> static_error(kNonStaticOpStreamz);
CellReader<int64_t> skipped(kNonStaticOpSkipStreamz);
CreateModule(kNonStaticFailure);
auto result = Run();
EXPECT_TRUE(result.succeeded());
EXPECT_EQ(static_error.Delta("mhlo.einsum"), 0);
EXPECT_EQ(skipped.Delta("mhlo.einsum"), 1);
}
TEST_F(VerifyTfxlaLegalizationTest, SkipsNonStaticInputsWithBounds) {
// Using a string constant here instead of testdata to make this compatible
// with open source.
static constexpr char kNonStaticWithBoundsSuccess[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1504 : i32}} {
func.func @main() -> tensor<?xi32, #mhlo.type_extensions<bounds = [4]>> attributes {tf.entry_function = {control_outputs = "", inputs = "i,j", outputs = "identity_RetVal"}} {
%0 = mhlo.constant dense<1.000000e+00> : tensor<f64>
%1 = mhlo.convert %0 : (tensor<f64>) -> tensor<i64>
%2 = mhlo.reshape %1 : (tensor<i64>) -> tensor<1xi64>
%3 = "mhlo.dynamic_iota"(%2) {iota_dimension = 0 : i64} : (tensor<1xi64>) -> tensor<?xi32, #mhlo.type_extensions<bounds = [4]>>
%4 = mhlo.multiply %3, %3 : tensor<?xi32, #mhlo.type_extensions<bounds = [4]>>
return %4 : tensor<?xi32, #mhlo.type_extensions<bounds = [4]>>
}
})";
CellReader<int64_t> legal_error(kFailedLegalizationStreamz);
CellReader<int64_t> static_error(kNonStaticOpStreamz);
CreateModule(kNonStaticWithBoundsSuccess);
auto result = Run();
EXPECT_TRUE(result.succeeded());
EXPECT_EQ(legal_error.Delta("mhlo.multiply"), 0);
EXPECT_EQ(static_error.Delta("mhlo.multiply"), 0);
}
TEST_F(VerifyTfxlaLegalizationTest, RecordsMultipleFailures) {
// Using a string constant here instead of testdata to make this compatible
// with open source.
static constexpr char kMultipleFailures[] = 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>
%1 = "tf.AlsoBad"() {value = dense<10> : tensor<1xi32>} : () -> tensor<1xi32>
func.return %0 : tensor<1xi32>
}
})";
CellReader<int64_t> error(kFailedLegalizationStreamz);
CreateModule(kMultipleFailures);
auto result = Run();
EXPECT_TRUE(result.failed());
EXPECT_EQ(error.Delta("tf.BadValue"), 1);
EXPECT_EQ(error.Delta("tf.AlsoBad"), 1);
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/transforms/xla_legalize_targets.cc
+58
View File
@@ -0,0 +1,58 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tf2xla/transforms/xla_legalize_targets.h"
#include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/Shape/IR/Shape.h" // from @llvm-project
#include "mlir/Dialect/Tensor/IR/Tensor.h" // from @llvm-project
#include "mlir/Transforms/DialectConversion.h" // from @llvm-project
#include "stablehlo/dialect/ChloOps.h" // from @stablehlo
#include "stablehlo/dialect/StablehloOps.h" // from @stablehlo
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
namespace mlir {
namespace hlo {
ConversionTarget GetDefaultLegalConversionTargets(MLIRContext& mlir_context,
bool legalize_chlo) {
ConversionTarget target(mlir_context);
if (legalize_chlo) {
target.addIllegalDialect<chlo::ChloDialect>();
target.addIllegalDialect<stablehlo::StablehloDialect>();
} else {
target.addLegalDialect<chlo::ChloDialect>();
}
target.addLegalDialect<mhlo::MhloDialect>();
target.addLegalDialect<arith::ArithDialect>();
target.addLegalDialect<func::FuncDialect>();
target.addLegalDialect<tensor::TensorDialect>();
target.addLegalDialect<shape::ShapeDialect>();
target.addLegalOp<func::CallOp>();
// These ops are legalized in LegalizeTFCommunication after this and that pass
// only operates on MHLO control flow ops.
target.addLegalOp<TF::_XlaHostComputeMlirOp, TF::XlaSendToHostOp,
TF::XlaRecvFromHostOp>();
return target;
}
} // namespace hlo
} // namespace mlir
tensorflow/compiler/mlir/tf2xla/transforms/xla_legalize_targets.h
+34
View File
@@ -0,0 +1,34 @@
/* 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_TRANSFORMS_XLA_LEGALIZE_TARGETS_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_TRANSFORMS_XLA_LEGALIZE_TARGETS_H_
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/Transforms/DialectConversion.h" // from @llvm-project
namespace mlir {
namespace hlo {
// Returns a ConversionTarget that includes default legalized MLIR dialects
// for conversion to XLA.
// If legalize_chlo is true, the resulting conversion target cannot have CHLO.
mlir::ConversionTarget GetDefaultLegalConversionTargets(
MLIRContext& mlir_context, bool legalize_chlo);
} // namespace hlo
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_TRANSFORMS_XLA_LEGALIZE_TARGETS_H_
tensorflow/compiler/mlir/tf2xla/transforms/xla_legalize_targets_test.cc
+95
View File
@@ -0,0 +1,95 @@
/* 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/transforms/xla_legalize_targets.h"
#include <gtest/gtest.h>
#include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/Shape/IR/Shape.h" // from @llvm-project
#include "mlir/Dialect/Tensor/IR/Tensor.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/DialectRegistry.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "mlir/Transforms/DialectConversion.h" // from @llvm-project
#include "stablehlo/dialect/ChloOps.h" // from @stablehlo
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
namespace mlir {
namespace hlo {
namespace {
mlir::DialectRegistry GetDefaultDialectRegistry() {
mlir::DialectRegistry registry;
registry.insert<arith::ArithDialect>();
registry.insert<func::FuncDialect>();
registry.insert<tensor::TensorDialect>();
registry.insert<shape::ShapeDialect>();
registry.insert<TF::TensorFlowDialect>();
registry.insert<chlo::ChloDialect>();
return registry;
}
class XlaLegalizeTargetsTest : public testing::Test {
public:
XlaLegalizeTargetsTest()
: context_(GetDefaultDialectRegistry()),
module_(mlir::ModuleOp::create(mlir::UnknownLoc::get(&context_))),
builder_(&module_->getBodyRegion()) {
context_.loadAllAvailableDialects();
}
protected:
mlir::MLIRContext context_;
mlir::OwningOpRef<mlir::ModuleOp> module_;
mlir::OpBuilder builder_;
};
TEST_F(XlaLegalizeTargetsTest, CreatesConversionTargets) {
auto const_int = mlir::arith::ConstantIntOp::create(
builder_, builder_.getUnknownLoc(), builder_.getI32Type(), /*value=*/10);
ConversionTarget target =
GetDefaultLegalConversionTargets(context_, /*legalize_chlo=*/false);
EXPECT_TRUE(target.isLegal(const_int));
}
TEST_F(XlaLegalizeTargetsTest, AllowsCHLODialect) {
auto const_int = chlo::ConstantOp::create(builder_, builder_.getUnknownLoc(),
builder_.getI32TensorAttr({42}));
ConversionTarget target =
GetDefaultLegalConversionTargets(context_, /*legalize_chlo=*/true);
EXPECT_TRUE(target.isIllegal(const_int));
}
TEST_F(XlaLegalizeTargetsTest, DontAllowCHLODialect) {
auto const_int = chlo::ConstantOp::create(builder_, builder_.getUnknownLoc(),
builder_.getI32TensorAttr({42}));
ConversionTarget target =
GetDefaultLegalConversionTargets(context_, /*legalize_chlo=*/false);
EXPECT_TRUE(target.isLegal(const_int));
}
} // namespace
} // namespace hlo
} // namespace mlir
tensorflow/compiler/mlir/tf2xla/transforms/xla_legalize_tf.cc
+276
View File
@@ -0,0 +1,276 @@
/* Copyright 2021 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 <string>
#include <utility>
#include "mhlo/transforms/rewriters.h"
#include "absl/log/log.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/FormatVariadic.h"
#include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project
#include "mlir/Dialect/Func/Extensions/AllExtensions.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/Quant/IR/Quant.h" // from @llvm-project
#include "mlir/Dialect/Shape/IR/Shape.h" // from @llvm-project
#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h" // from @llvm-project
#include "mlir/Dialect/Tensor/IR/Tensor.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributeInterfaces.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Support/DebugStringHelper.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Transforms/DialectConversion.h" // from @llvm-project
#include "stablehlo/dialect/ChloOps.h" // from @stablehlo
#include "stablehlo/dialect/StablehloOps.h" // from @stablehlo
#include "stablehlo/transforms/Passes.h" // from @stablehlo
#include "tensorflow/compiler/mlir/tensorflow/transforms/lower_tf.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/legalization_op_config.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/legalize_tf_with_tf2xla_passes.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/passes.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/xla_legalize_targets.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h" // IWYU pragma: keep, dependent dialect
#include "xla/mlir_hlo/mhlo/transforms/passes.h"
#include "xla/mlir_hlo/mhlo/transforms/rewriters.h"
#include "xla/mlir_hlo/mhlo/utils/type_conversion.h"
#include "tensorflow/core/lib/monitoring/counter.h"
namespace mlir {
namespace mhlo {
namespace {
#define GEN_PASS_DEF_LEGALIZETF
#include "tensorflow/compiler/mlir/tf2xla/transforms/xla_legalize_tf_passes.h.inc"
auto *mlir_legalization_count = tensorflow::monitoring::Counter<1>::New(
"/tensorflow/core/tf2xla/v1/mlir_failed_xla_legalize_tf_count",
"Counts the attempts of legalization of ops", "op_name");
auto *mlir_failed_legalization_count = tensorflow::monitoring::Counter<2>::New(
"/tensorflow/core/tf2xla/v1/mlir_failed_xla_legalize_tf_pass_count",
"Counts the failure of legalization of ops", "op_name", "legality");
class LegalizeTF : public impl::LegalizeTFBase<LegalizeTF> {
public:
explicit LegalizeTF(bool legalize_chlo,
std::optional<StringRef> tf2xla_fallback_device_type,
bool prefer_tf2xla) {
legalize_chlo_ = legalize_chlo;
prefer_tf2xla_ = prefer_tf2xla;
use_tf2xla_fallback_ = tf2xla_fallback_device_type.has_value();
if (tf2xla_fallback_device_type.has_value()) {
device_type_ = tf2xla_fallback_device_type.value().str();
}
}
/// Performs the lowering to XLA dialect.
void runOnOperation() override;
};
#define GEN_PASS_DEF_LEGALIZETFMODULEPASS
#include "tensorflow/compiler/mlir/tf2xla/transforms/xla_legalize_tf_passes.h.inc"
// Patterns whose root op is in the set `include_ops` are moved from the set
// `from` to the returned set. This is used to partition patterns by op so they
// can be cleanly migrated from the old bridge to the MLIR bridge.
RewritePatternSet PatternsIncludeOps(RewritePatternSet &from) {
RewritePatternSet to(from.getContext());
// Filter NativePatterns.
for (auto &pattern : from.getNativePatterns()) {
std::optional<OperationName> pat_op_name = pattern->getRootKind();
// If the pattern does not have a specific operation, always include it,
// If the pattern is in include_ops then include it.
bool include =
!pat_op_name || hlo::IsTypeLegalizedWithMlir(
pat_op_name->getRegisteredInfo()->getTypeID());
if (include) to.add(std::move(pattern));
}
// Don't filter PDLPatterns.
to.add(std::move(from.getPDLPatterns()));
return to;
}
std::string OperationLegalityString(Operation *op,
const ConversionTarget &target) {
auto op_name = op->getName();
auto action = target.getOpAction(op_name);
if (!action.has_value()) {
return "Unknown";
}
switch (action.value_or(ConversionTarget::LegalizationAction::Legal)) {
case ConversionTarget::LegalizationAction::Legal:
return "Legal";
case ConversionTarget::LegalizationAction::Dynamic:
return "Dynamic";
case ConversionTarget::LegalizationAction::Illegal:
return "Illegal";
default:
return "Invalid";
}
}
void IncrementFailedLegalizationCount(Operation *op,
const ConversionTarget &target) {
auto op_name = op->getName();
auto name_string = op_name.getStringRef().str();
auto op_legality = OperationLegalityString(op, target);
mlir_failed_legalization_count->GetCell(name_string, op_legality)
->IncrementBy(1);
}
mlir::LogicalResult ApplyPatterns(Operation *op, RewritePatternSet &patterns,
bool legalize_chlo) {
ConversionTarget target =
hlo::GetDefaultLegalConversionTargets(*op->getContext(), legalize_chlo);
DenseSet<Operation *> unconverted_ops;
ConversionConfig config;
config.unlegalizedOps = &unconverted_ops;
auto result = applyPartialConversion(op, target, std::move(patterns), config);
if (failed(result)) {
IncrementFailedLegalizationCount(op, target);
}
for (const auto &unconverted_op : unconverted_ops) {
IncrementFailedLegalizationCount(unconverted_op, target);
}
return result;
}
mlir::LogicalResult StablehloToMhlo(Operation *op) {
ConversionTarget target(*op->getContext());
stablehlo::setupStablehloToHloConversionTarget(target);
RewritePatternSet patterns(op->getContext());
stablehlo::StablehloToHloTypeConverter shlo_converter;
stablehlo::populateStablehloToHloPatterns(&patterns, &shlo_converter,
patterns.getContext());
stablehlo::registerFuncOpsForTypeConversion(target, patterns, shlo_converter);
if (failed(applyPartialConversion(op, target, std::move(patterns)))) {
return op->emitError("TF2XLA failed to convert StableHLO to MHLO");
}
return success();
}
/// When `tf2xla_fallback_device_type` is not `None`, also uses legalization
/// patterns from TF2XLA fallback for provided device type (see
/// legalize_tf_with_tf2xla.cc for details). By default, TF2XLA fallback is
/// not used.
LogicalResult legalizeTF(Operation *op, bool legalize_chlo,
std::optional<StringRef> tf2xla_fallback_device_type,
bool prefer_tf2xla) {
MLIRContext *context = op->getContext();
RewritePatternSet legalize_lower_patterns(context);
// Note that the `OperationConverter` orders patterns lexicographically by:
// 1) Ascending legalization depth (i.e., minimum number of patterns
// necessary
// to arrive at conversion target). This requires relevant patterns to
// specify the list of ops generated by it which most of patterns
// implemented in C++ don't do so this comparison doesn't work in those
// cases.
// 2) Descending pattern benefit.
// 3) Op specific patterns over patterns with MatchAnyOpTypeTag.
// 4) Order of patterns in `RewritePatternSet`.
// Add TF->HLO legalization patterns.
hlo::PopulateLegalizeTfPatterns(context, &legalize_lower_patterns);
// Add TF->TF lowering patterns.
TF::PopulateTFLoweringBeforeHLOPatterns(context, &legalize_lower_patterns);
if (tf2xla_fallback_device_type && prefer_tf2xla) {
VLOG(1) << "TF to XLA legalization patterns are partitioned by op into "
"either native MLIR legalization, or TF2XLA fallback "
"legalzation, with a preference toward TF2XLA.";
} else if (tf2xla_fallback_device_type) {
VLOG(1) << "TF to XLA legalization patterns include all native patterns "
"and TF2XLA fallback patterns.";
} else {
VLOG(1) << "TF to XLA legalization patterns are native patterns only.";
}
// Set patterns to legalize_lower_patterns to check whether they should use
// MLIR or TF2XLA lowering patterns.
RewritePatternSet patterns = (tf2xla_fallback_device_type && prefer_tf2xla)
? PatternsIncludeOps(legalize_lower_patterns)
: std::move(legalize_lower_patterns);
Tf2XlaTypeConverter converter;
if (tf2xla_fallback_device_type) {
// Add TF->HLO legalization patterns via TF2XLA fallback.
PopulateLegalizeTfWithTf2XlaPatterns(tf2xla_fallback_device_type.value(),
patterns, context, converter,
prefer_tf2xla);
}
// Populate with CHLO->HLO lowerings to account for TF ops legalized to
// CHLO first.
stablehlo::StablehloToHloTypeConverter hlo_converter;
stablehlo::populateStablehloToHloPatterns(&patterns, &hlo_converter, context);
if (legalize_chlo) {
chlo::populateChloToHighLevelMhloOpPatterns(context, &patterns);
stablehlo::populateChloToStablehloPatterns(context, &patterns);
}
// ConstantLike op is convenient to create splat constants, but is
// canonicalized to plain HLO constant if statically shaped. Add the
// canonicalization pattern to pattern list to enable multi-hop lowering.
chlo::ConstantLikeOp::getCanonicalizationPatterns(patterns, context);
if (failed(ApplyPatterns(op, patterns, legalize_chlo))) {
return failure();
}
// HLO->MLIR raises to StableHLO, but users of this pass expect MHLO.
return StablehloToMhlo(op);
}
// Performs the lowering to XLA dialect.
void LegalizeTF::runOnOperation() {
auto op = getOperation();
VLOG(3) << "LegalizeTF(legalize_chlo=" << legalize_chlo_
<< ", prefer_tf2xla=" << prefer_tf2xla_ << ") on module:\n"
<< mlir::debugString(*op);
auto op_name = op->getName().getStringRef().str();
mlir_legalization_count->GetCell(op_name)->IncrementBy(1);
std::optional<StringRef> tf2xla_fallback_device_type = std::nullopt;
if (use_tf2xla_fallback_) {
tf2xla_fallback_device_type = device_type_;
}
if (failed(legalizeTF(op, legalize_chlo_, tf2xla_fallback_device_type,
prefer_tf2xla_))) {
signalPassFailure();
}
}
} // end namespace
std::unique_ptr<OperationPass<ModuleOp>> createLegalizeTFPass(
bool legalize_chlo, std::optional<StringRef> tf2xla_fallback_device_type,
bool prefer_tf2xla) {
return std::make_unique<LegalizeTF>(
legalize_chlo, tf2xla_fallback_device_type, prefer_tf2xla);
}
} // end namespace mhlo
} // end namespace mlir
tensorflow/compiler/mlir/tf2xla/transforms/xla_legalize_tf_passes.td
+110
View File
@@ -0,0 +1,110 @@
/* Copyright 2021 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.
==============================================================================*/
// Declare passes used in xla_legalize_tf.
include "mlir/Pass/PassBase.td"
def LegalizeTF : Pass<"xla-legalize-tf", "ModuleOp"> {
let summary = "Legalize from TF dialect's or HLO dialect's control flow.";
let description = [{
Legalizes from TF dialect to HLO dialect. When allow_partial_conversion is
false, emits an error if there is any operation that can't be legalized.
When `tf2xla_fallback_device_type` is not `None`, also uses legalization
patterns from TF2XLA fallback for provided device type (see
legalize_tf_with_tf2xla.cc for details). By default, TF2XLA fallback is not
used.
}];
let options = [
Option<"legalize_chlo_", "legalize-chlo", "bool", /*default=*/"true",
"Legalizes intermediate chlo ops to hlo">,
Option<"use_tf2xla_fallback_", "use-tf2xla-fallback", "bool",
/*default=*/"false",
"Use TF2XLA fallback for legalization">,
Option<"device_type_", "device-type", "std::string",
/*default=*/"\"INVALID_DEVICE_TYPE\"",
"The device type used by TF2XLA fallback. Must be specified if "
"use-tf2xla-fallback is true, otherwise not used">,
Option<"prefer_tf2xla_", "prefer-tf2xla", "bool",
/*default=*/"false",
"Prioritize tf2xla fallback legalization over MLIR legalization "
"patterns">,
];
let constructor = "mlir::mhlo::createLegalizeTFPass()";
let dependentDialects = [
"arith::ArithDialect",
"chlo::ChloDialect",
"func::FuncDialect",
"mhlo::MhloDialect",
"quant::QuantDialect",
"shape::ShapeDialect",
"sparse_tensor::SparseTensorDialect",
"stablehlo::StablehloDialect"
];
}
def LegalizeTFCollective : Pass<"xla-legalize-tf-collective", "ModuleOp"> {
let summary = "Legalize TF/XLA collective ops (TensorFlow dialect) to the HLO dialect";
let constructor = "mlir::mhlo::CreateLegalizeTFCollectivePass()";
let dependentDialects = ["mhlo::MhloDialect", "sparse_tensor::SparseTensorDialect"];
}
def VerifyTFXLALegalization : Pass<"tfxla-verify-legalization", "mlir::func::FuncOp"> {
let summary = "Verifies that all TF ops have been legalized to XLA.";
let description = [{"Ensures that all Tensorflow ops have been legalized to "
"XLA and reports an error about which op has not been"
"legalized. This pass does not transform any ops and is just"
" a verification pass to ensure invariants are true."}];
let options = [
Option<"legalize_chlo_", "legalize-chlo", "bool", /*default=*/"true",
"Legalizes intermediate chlo ops to hlo">
];
let constructor = "mlir::hlo::CreateVerifyTFXLALegalizationPass()";
}
def TFXLADeviceSpecificTransforms : Pass<"tfxla-device-specific-transforms",
"mlir::func::FuncOp"> {
let summary = "Transforms ops that require device context into device independent TF Ops.";
let description = [{"Transforms device specific ops into device independent"
"ops."}];
let options = [
Option<"device_type_", "device-type", "std::string",
/*default=*/"\"INVALID_DEVICE_TYPE\"",
"The device type being targeted.">,
];
let constructor = "mlir::mhlo::CreateTFXLADeviceSpecificTransformsPass()";
}
def InfeedOpsXlaAdjustLayout : Pass<"infeed-ops-xla-adjust-layout",
"mlir::func::FuncOp"> {
let summary = "Adjusts Infeed ops layout for XLA.";
let description = [{"Adjust layouts so infeed send & receive use the same "
"format."}];
let constructor = "mlir::mhlo::CreateInfeedsOpsXlaAdjustLayoutPass()";
let dependentDialects = ["mhlo::MhloDialect"];
}
tensorflow/compiler/mlir/tf2xla/transforms/xla_legalize_tf_test.cc
+117
View File
@@ -0,0 +1,117 @@
/* 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 <functional>
#include <memory>
#include <gtest/gtest.h>
#include "absl/status/statusor.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/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/PassManager.h" // from @llvm-project
#include "mlir/Transforms/DialectConversion.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/deserialize_mlir_module_utils.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/passes.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/lib/monitoring/cell_reader.h"
namespace tensorflow {
namespace {
using ::mlir::MLIRContext;
using ::mlir::ModuleOp;
using ::mlir::OwningOpRef;
using ::mlir::PassManager;
using ::tensorflow::monitoring::testing::CellReader;
absl::StatusOr<OwningOpRef<ModuleOp>> GetMlirModuleFromString(
absl::string_view module_string, MLIRContext* context) {
mlir::DialectRegistry mlir_registry;
RegisterAllTensorFlowDialects(mlir_registry);
context->appendDialectRegistry(mlir_registry);
OwningOpRef<ModuleOp> mlir_module;
auto status =
tensorflow::DeserializeMlirModule(module_string, context, &mlir_module);
if (!status.ok()) {
return status;
}
return mlir_module;
}
bool BuildAndRunPipeline(absl::string_view module_string,
const std::function<void(PassManager*)>& passes) {
mlir::registerPassManagerCLOptions();
MLIRContext context;
OwningOpRef<ModuleOp> module =
GetMlirModuleFromString(module_string, &context).value();
PassManager pm(&context);
if (mlir::failed(mlir::applyPassManagerCLOptions(pm))) return false;
passes(&pm);
return pm.run(module.get()).succeeded();
}
std::function<void(PassManager*)> legalizeTFPasses() {
return [](PassManager* pm) {
pm->addPass(mlir::mhlo::createLegalizeTFPass(
/* legalize_chlo=*/true, llvm::StringRef("gpu/xpu"),
/* prefer_tf2xla=*/false));
};
}
TEST(XlaLegalizeTest, IllegalOp) {
constexpr char kMlirIllegalOpStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
func.func @main() -> tensor<1xi32> {
%0 = "tf.DoesntExist"() : () -> tensor<1xi32>
func.return %0 : tensor<1xi32>
}
})";
CellReader<int64_t> legalize_failure_count(
"/tensorflow/core/tf2xla/v1/mlir_failed_xla_legalize_tf_pass_count");
auto status = BuildAndRunPipeline(kMlirIllegalOpStr, legalizeTFPasses());
EXPECT_TRUE(status);
EXPECT_EQ(legalize_failure_count.Read("tf.DoesntExist", "Unknown"), 1);
}
TEST(XlaLegalizeTest, LegalOp) {
// We expect legalization to fail for legal op with dynamic shapes:
static constexpr char kMlirLegalOpStr[] = R"(
func.func @infeed_dequeue_tuple_dynamic_error() -> (tensor<3x3xf32>, tensor<4x?xf32>) {
%0:2 = "tf.InfeedDequeueTuple"() : () -> (tensor<3x3xf32>, tensor<4x?xf32>) func.return %0#0, %0#1 : tensor<3x3xf32>, tensor<4x?xf32>
})";
CellReader<int64_t> legalize_failure_count(
"/tensorflow/core/tf2xla/v1/mlir_failed_xla_legalize_tf_pass_count");
auto status = BuildAndRunPipeline(kMlirLegalOpStr, legalizeTFPasses());
EXPECT_TRUE(status);
EXPECT_EQ(legalize_failure_count.Read("tf.InfeedDequeueTuple", "Unknown"), 1);
}
} // namespace
} // namespace tensorflow