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/dtensor/BUILD
+29
View File
@@ -0,0 +1,29 @@
# DTensor is a TensorFlow extension that enables large-scale modeling with
# minimal changes to user code.
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
package_group(
name = "dtensor-users",
packages = [
"//tensorflow_models/google/...",
"//third_party/aimee/clara2_labs/...",
"//third_party/py/jax_tpu_embedding/...",
"//third_party/py/tf_keras/dtensor/...",
],
)
package_group(
name = "dtensor-internal",
packages = [
"//learning/brain/experimental/jax_tpu_embedding/...",
"//learning/brain/public/...",
"//tensorflow/core/...",
"//tensorflow/dtensor/...",
"//tensorflow/python/...",
"//tensorflow_federated/...", # Needs access to C++ API for dtensor device
],
)
tensorflow/dtensor/build_defs.bzl
+174
View File
@@ -0,0 +1,174 @@
"""Helpers for defining multi-platform DTensor test targets."""
load("//tensorflow:tensorflow.bzl", "py_test")
# LINT.IfChange
ALL_BACKENDS = [
"cpu", # 1 physical CPU,
"gpu", # 1 physical GPU,
"tpu", # 2 physical TPU devices
]
TPU_V3_DONUT_BACKEND = "tpu_v3_2x2" # 8 TPU devices
TPU_V4_DONUT_BACKEND = "tpu_v4_2x2" # 8 TPU devices for non-Megacore targets and 4 for Megacore targets
GPU_2DEVS_BACKEND = "2gpus" # 2 Physical GPUs.
# LINT.ThenChange(
# python/tests/test_backend_name.py:backend_name,
# python/tests/test_backend_name.oss.py:backend_name
# )
# FIXME(feyu): Gradually increase the coverage of OSS tests.
# LINT.IfChange
def _get_configurations(
disable,
enable,
disable_tfrt,
backend_tags,
backend_deps,
additional_backends, # buildifier: disable=unused-variable
shard_count):
"""Returns a list of dtensor_test configurations."""
disabled_tags = ["manual", "disabled"]
disabled_tfrt_configs = [d + "_tfrt" for d in disable_tfrt]
disabled_backends = [backend for backend in disable if backend not in enable]
backend_variant_deps = {
"gpu": [],
"tpu": [
],
TPU_V3_DONUT_BACKEND: [
],
TPU_V4_DONUT_BACKEND: [
],
}
configurations = [
dict(suffix = "cpu", backend = "cpu", tags = [], flags = [], env = {}, deps = []),
dict(
suffix = "gpu",
backend = "gpu",
tags = ["requires-gpu", "gpu"],
flags = [],
env = {},
deps = [],
),
]
if GPU_2DEVS_BACKEND in additional_backends:
configurations = configurations + [
dict(
suffix = GPU_2DEVS_BACKEND,
backend = GPU_2DEVS_BACKEND,
tags = ["requires-gpu:2", "gpu"],
flags = [],
env = {
},
deps = [],
),
]
# Post processing configurations.
for config in configurations:
config["tags"] = config["tags"] + backend_tags.get(config["backend"], [])
config["env"]["DTENSOR_TEST_UTIL_BACKEND"] = config["suffix"]
if config["backend"] in disabled_backends or config["suffix"] in disabled_tfrt_configs:
config["tags"] += disabled_tags
config["deps"] = (
config["deps"] +
backend_variant_deps.get(config["backend"], []) +
backend_deps.get(config["backend"], [])
)
config["shard_count"] = shard_count.get(config["backend"], None) if shard_count else None
return configurations
# LINT.ThenChange(build_defs.bzl)
def dtensor_test(
name,
srcs,
deps = [],
args = [],
env = {},
disable = [],
enable = [],
disable_tfrt = [],
data = [],
tags = [],
backend_tags = {},
backend_deps = {},
additional_backends = [],
main = None,
shard_count = None,
size = None,
get_configurations = _get_configurations,
test_rule = py_test):
"""Defines a set of per-platform DTensor test targets.
Generates test targets named:
:name # test suite that tests all backends
:name_cpu
:name_cpu_tfrt
:name_gpu # must run with --config=cuda
:name_tpu # recommend to be run with -c opt
Args:
name: test name
srcs: source files
deps: dependencies
args: arguments to pass to the test
env: environment variables to set when the test is executed
disable: list of backends on which the test should be disabled, e.g., ["cpu"]
enable: list of specific configs on which the test should be enabled,
e.g., ["tpu"]. This overrides 'disable'.
disable_tfrt: list of backends that are disabled for tfrt. This overrides 'enable'.
data: data dependencies
tags: test tags
backend_tags: a dictionary keyed by backend name of per-backend tags.
backend_deps: a dictionary keyed by backend name of per-backend deps.
additional_backends: list of backends in addition to common cpu/tpu/gpu.
main: the Python main file.
shard_count: a dictionary keyed by backend name of per-backend shard counts.
size: the test size.
get_configurations: a function that returns the list of configurations. Used to generate non-OSS test targets.
test_rule: test rule
"""
configurations = get_configurations(
disable = disable,
enable = enable,
disable_tfrt = disable_tfrt,
backend_tags = backend_tags,
backend_deps = backend_deps,
additional_backends = additional_backends,
shard_count = shard_count,
)
if main == None:
if len(srcs) == 1:
main = srcs[0]
else:
fail("Only one test source file is currently supported.")
all_tests = []
for config in configurations:
config_name = name + "_" + config["suffix"]
all_tests.append(config_name)
python_version = "PY3"
test_env = {}
test_env.update(config["env"])
test_env.update(env)
test_rule(
env = test_env,
name = config_name,
main = main,
srcs = srcs,
data = data,
args = config["flags"] + args,
deps = config["deps"] + deps,
tags = config["tags"] + tags,
python_version = python_version,
shard_count = config["shard_count"],
size = size,
)
native.test_suite(name = name, tests = all_tests, tags = ["-manual"])
tensorflow/dtensor/cc/BUILD
+435
View File
@@ -0,0 +1,435 @@
#include "third_party/absl/strings/str_cat.h"
#DTensor C++ runtime and libraries.
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow:tensorflow.default.bzl", "tf_kernel_library")
load(
"//tensorflow/core/platform:build_config.bzl",
"tf_dtensor_tpu_dependencies",
)
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
default_visibility = [
"//smartass/brain/jax/embeddings:__pkg__",
"//tensorflow/dtensor:dtensor-internal",
"//tensorflow:__pkg__",
]
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = default_visibility,
licenses = ["notice"],
)
cc_library(
name = "constants",
hdrs = ["constants.h"],
)
cc_library(
name = "dstatus",
hdrs = ["dstatus.h"],
deps = [
"//tensorflow/core:lib",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/types:optional",
],
)
cc_library(
name = "dtensor_utils",
srcs = ["dtensor_utils.cc"],
hdrs = ["dtensor_utils.h"],
deps = [
"//tensorflow/core:lib",
"@com_google_absl//absl/strings",
"@xla//xla/tsl/platform:status",
"@xla//xla/tsl/util:env_var",
],
)
cc_library(
name = "tensor_layout",
srcs = ["tensor_layout.cc"],
hdrs = ["tensor_layout.h"],
deps = [
":dstatus",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/dtensor/proto:layout_proto_cc",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/container:inlined_vector",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:optional",
"@llvm-project//mlir:IR",
],
)
cc_library(
name = "small_constant_optimization",
srcs = ["small_constant_optimization.cc"],
hdrs = ["small_constant_optimization.h"],
deps = [
":constants",
":tensor_layout",
"//tensorflow/c:c_api",
"//tensorflow/c:c_api_experimental",
"//tensorflow/c:tf_tensor_internal",
"//tensorflow/c/eager:c_api",
"//tensorflow/c/eager:c_api_experimental",
"//tensorflow/core:framework",
"//tensorflow/core:framework_lite",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/platform:tstring",
"//tensorflow/dtensor/proto:layout_proto_cc",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:optional",
],
)
cc_library(
name = "tensor_with_layout",
srcs = ["tensor_with_layout.cc"],
hdrs = ["tensor_with_layout.h"],
deps = [
":constants",
":tensor_layout",
"//tensorflow/c/eager:c_api",
"//tensorflow/core:framework",
"//tensorflow/core/platform:fingerprint",
"@llvm-project//llvm:Support",
],
)
cc_library(
name = "dtensor_operation",
hdrs = ["dtensor_operation.h"],
deps = [
":tensor_layout",
"//tensorflow/c/eager:c_api",
],
)
cc_library(
name = "dtensor_device_util",
srcs = ["dtensor_device_util.cc"],
hdrs = ["dtensor_device_util.h"],
deps = [
":constants",
":dstatus",
":dtensor_operation",
":dtensor_utils",
":small_constant_optimization",
":tensor_layout",
":tensor_with_layout",
"//tensorflow/c:safe_ptr",
"//tensorflow/c:tf_datatype",
"//tensorflow/c:tf_status_headers",
"//tensorflow/c:tf_status_helper",
"//tensorflow/c:tf_tensor",
"//tensorflow/c:tf_tensor_internal",
"//tensorflow/c/eager:c_api",
"//tensorflow/c/eager:c_api_experimental",
"//tensorflow/c/eager:c_api_internal",
"//tensorflow/c/eager:tfe_context_internal",
"//tensorflow/c/eager:tfe_tensorhandle_internal",
"//tensorflow/c/eager/parallel_device:parallel_device_lib",
"//tensorflow/core:core_cpu",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/common_runtime/eager:context",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:span",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@tsl//tsl/platform:refcount",
"@xla//xla:status_macros",
"@xla//xla/tsl/platform:logging",
"@xla//xla/tsl/platform:status",
],
)
cc_library(
name = "dtensor_ops",
srcs = [
"dtensor_ops.cc",
],
deps = [
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
],
alwayslink = 1,
)
# These ops are created only by DTensor MLIR passes and never by users, so they don't need Python wrappers.
cc_library(
name = "dtensor_meta_ops",
srcs = [
"dtensor_meta_ops.cc",
],
deps = [
":tensor_layout",
"//tensorflow/core:framework",
],
alwayslink = 1,
)
cc_library(
name = "tpu_system_interface",
srcs = ["tpu_system_interface.cc"],
hdrs = ["tpu_system_interface.h"],
deps = [
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"@com_google_absl//absl/time",
],
)
cc_library(
name = "save_restore_util",
srcs = ["save_restore_util.cc"],
hdrs = ["save_restore_util.h"],
deps = [
":dstatus",
":tensor_layout",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/dtensor/mlir:value_utils",
"@com_google_absl//absl/container:flat_hash_map",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "dtensor_tpu_ops",
srcs = ["dtensor_tpu_ops.cc"],
deps = [
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
],
alwayslink = 1,
)
tf_kernel_library(
name = "dtensor_tpu_kernels",
srcs = [
"dtensor_tpu_kernels.cc",
],
tags = [
"no_rocm",
"tpu",
], # Disable building of TPU kernels on non-TPU platforms.
deps = [
":tpu_system_interface",
"//tensorflow/core:framework",
"//tensorflow/core:portable_gif_internal",
"//tensorflow/core/tpu:tpu_configuration",
"//tensorflow/core/tpu/kernels:tpu_configuration_ops",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:vlog_is_on",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/time",
"@xla//xla/tpu:status_helper",
"@xla//xla/tpu:tpu_api",
"@xla//xla/tpu:tpu_ops_c_api_hdrs",
],
alwayslink = 1,
)
cc_library(
name = "dtensor_graph_to_mlir_pass",
srcs = ["dtensor_graph_to_mlir_pass.cc"],
hdrs = ["dtensor_graph_to_mlir_pass.h"],
deps = [
":constants",
":dtensor_utils",
"//tensorflow/compiler/jit:flags_headers",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:convert_type",
"//tensorflow/compiler/mlir/tensorflow:device_util",
"//tensorflow/compiler/mlir/tensorflow:error_util",
"//tensorflow/compiler/mlir/tf2xla/api/v2:graph_to_tf_executor",
"//tensorflow/core:core_cpu",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/common_runtime:device_set",
"//tensorflow/dtensor/mlir:dtensor_mlir_passes",
"//tensorflow/dtensor/mlir:tf_dtensor_dialect",
"//tensorflow/dtensor/mlir/dtensor_dialect:Dialect",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/strings",
"@llvm-project//mlir:AllExtensions",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:RegisterAllExtensions",
"@llvm-project//mlir:Support",
"@xla//xla:status_macros",
"@xla//xla/tsl/platform:status",
],
)
cc_library(
name = "default_parallel_executor_lib",
deps = [":default_parallel_executor"],
)
cc_library(
name = "default_parallel_executor",
srcs = ["default_parallel_executor.cc"],
deps = [
":parallel_executor_interface",
"//tensorflow/core/platform:logging",
"@xla//xla/tsl/platform:statusor",
],
)
cc_library(
name = "parallel_executor_interface",
hdrs = ["parallel_executor.h"],
deps = [
":tensor_layout",
":tensor_with_layout",
"//tensorflow/c/eager:c_api_internal",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@xla//xla:future",
],
)
cc_library(
name = "dtensor_device_cc",
srcs = ["dtensor_device.cc"],
hdrs = ["dtensor_device.h"],
deps = [
":constants",
":default_parallel_executor_lib",
":dstatus",
":dtensor_device_util",
":dtensor_graph_to_mlir_pass",
":dtensor_meta_ops",
":dtensor_operation",
":dtensor_ops",
":dtensor_tpu_ops",
":dtensor_utils",
":parallel_executor_interface",
":small_constant_optimization",
":tensor_layout",
":tensor_with_layout",
":tpu_system_interface",
"//tensorflow/c:c_api",
"//tensorflow/c:c_api_experimental",
"//tensorflow/c:tf_status_helper",
"//tensorflow/c:tf_tensor_internal",
"//tensorflow/c/eager:abstract_tensor_handle",
"//tensorflow/c/eager:c_api",
"//tensorflow/c/eager:c_api_experimental",
"//tensorflow/c/eager:immediate_execution_tensor_handle",
"//tensorflow/c/eager:tfe_context_internal",
"//tensorflow/c/eager:tfe_op_attrs_internal",
"//tensorflow/c/eager:tfe_tensorhandle_internal",
"//tensorflow/c/eager/parallel_device:parallel_device_lib",
"//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_flags",
"//tensorflow/compiler/mlir/tf2xla/api/v2:mlir_roundtrip_flags",
"//tensorflow/compiler/mlir/tf2xla/api/v2:tf_executor_to_graph",
"//tensorflow/core:core_cpu",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/common_runtime:device_set",
"//tensorflow/core/common_runtime/eager:context",
"//tensorflow/core/common_runtime/eager:eager_executor",
"//tensorflow/core/common_runtime/eager:eager_operation",
"//tensorflow/core/common_runtime/eager:tensor_handle",
"//tensorflow/core/platform:refcount",
"//tensorflow/core/profiler/lib:traceme",
"//tensorflow/dtensor/mlir:layout_parsing",
"//tensorflow/dtensor/mlir:op_utils",
"//tensorflow/dtensor/mlir:spmd_expander",
"//tensorflow/dtensor/proto:layout_proto_cc",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/container:fixed_array",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/log:vlog_is_on",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:span",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@tsl//tsl/platform:refcount",
"@tsl//tsl/profiler/lib:traceme",
"@xla//xla:status_macros",
"@xla//xla/tpu:c_api_decl",
"@xla//xla/tpu:tpu_topology_external",
"@xla//xla/tsl/platform:status",
"@xla//xla/tsl/util:env_var",
] + tf_dtensor_tpu_dependencies(),
)
cc_library(
name = "layout_to_xla_sharding",
srcs = ["xla_spmd/layout_to_xla_sharding.cc"],
hdrs = ["xla_spmd/layout_to_xla_sharding.h"],
deps = [
":dstatus",
":tensor_layout",
"//tensorflow/core:portable_gif_internal",
"//tensorflow/core/platform:status",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/types:span",
"@llvm-project//llvm:Support",
"@xla//xla:status_macros",
"@xla//xla:xla_data_proto_cc",
],
)
cc_library(
name = "mesh_type",
hdrs = ["mesh_type.h"],
deps = [
":tensor_layout",
"//tensorflow/c:conversion_macros",
],
)
cc_library(
name = "slice_util",
srcs = ["slice_util.cc"],
hdrs = ["slice_util.h"],
deps = [
":tensor_layout",
"//tensorflow/core:lib",
"@llvm-project//mlir:IR",
],
)
tf_cc_test(
name = "small_constant_optimization_test",
srcs = ["small_constant_optimization_test.cc"],
deps = [
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/dtensor/cc:small_constant_optimization",
"//tensorflow/dtensor/cc:tensor_layout",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
],
)
tensorflow/dtensor/cc/constants.h
+170
View File
@@ -0,0 +1,170 @@
/* Copyright 2022 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_DTENSOR_CC_CONSTANTS_H_
#define TENSORFLOW_DTENSOR_CC_CONSTANTS_H_
namespace tensorflow {
namespace dtensor {
// Constants used within dtensor scope.
// Qualified attribute without `_` prefix.
// Used in Ops attribute registration.
static constexpr char kQualifiedLayoutAttr[] = "layout";
// Internal attribute to DTensor MLIR passes and Graph nodes.
// Prefixed with `_` so that it doesn't require op attribute registration.
static constexpr char kLayoutAttr[] = "_layout";
// Indicates a non-binding layout hint provided by the user.
// `tf` prefix attached in MLIR importer for dialect requirements.
static constexpr char kCustomDefaultLayoutAttr[] = "tf._default_layout";
// Indicates a non-binding layout hint provided by the user.
static constexpr char kDefaultLayoutAttr[] = "_default_layout";
// Attribute carries layout information from Custom Device Arguments.
// `tf` prefix attached in MLIR importer for dialect requirements.
static constexpr char kCustomDeviceAttr[] = "tf._layout";
// Indicates a default mesh provided by the user as fallback during mesh
// propagation. `tf` prefix attached in MLIR importer for dialect requirements.
static constexpr char kCustomDefaultMeshAttr[] = "tf._default_mesh";
// Attribute attached on _Arg node for the mesh config.
static constexpr char kMeshAttr[] = "_mesh";
// Attribute carries mesh information from Custom Device Arguments.
// `tf` prefix attached in MLIR importer for dialect requirements.
static constexpr char kCustomDeviceMeshAttr[] = "tf._mesh";
// Attribute carries argument indices for newly inferred layout of resource
// handle.
static constexpr char kNewResourceLayoutIndices[] =
"_inferred_resource_indices";
// Attribute carries layout for newly inferred layout of resource handle.
static constexpr char kNewResourceArgLayouts[] = "_inferred_resource_layouts";
static constexpr char kNumLocalOutputsAttr[] = "_num_local_outputs";
// Attribute carries input layout information for shape op.
static constexpr char kShapeOpInputLayout[] = "_shape_input_layout";
// Attribute carries input layout index for shape op. This forms a 1 -> 1
// mapping for kShapeOpInputLayout above.
static constexpr char kShapeOpInputLayoutIndices[] = "_shape_input_indices";
// Attribute that carries global shape of operation. Used to preserve global
// shape to be used during SPMD expansion.
static constexpr char kGlobalShape[] = "_global_shape";
// Global shape attribute with `tf.` dialect to be used for annotating func op
// arguments/return values.
static constexpr char kGlobalShapeDialectAttr[] = "tf._global_shape";
// Attribute attached to resource-type function arguments containing the local
// shape of the tensor that is being assigned to it.
static constexpr char kAssignedResourceLocalShape[] =
"tf._assigned_resource_local_shape";
// Tensor handles smaller than this is considered as small tensor. We perform
// some optimizations around it. For example, will be transformed into constant
// values during graph building, instead of being passed as inputs. In addition,
// we allow automatical broadcasting small non-DTensor to DTensor device, which
// is very useful for shape/axis info tensor in eager mode (eliminating the need
// forcing users to do explicit copy-to-mesh).
static constexpr int kSmallTensorThreshold = 20;
// Contains a serialized mesh. Will be attached to a FloorMod op to denote which
// mesh the output of the FloorMod op is giving coordinates for.
static constexpr char kMeshCoordinatesAttr[] = "_mesh_coordinates";
// Attribute used to determine if a module pass should log long form information
// such as IR dumps etc.
static constexpr char kDoNotLog[] = "dtensor.do_not_log";
// Attribute used to record the name of the eager operation triggered the
// DTensor rewrites.
static constexpr char kEagerOperationName[] = "dtensor.eager_operation_name";
// The number of TPU cores in a donut.
static constexpr int kTpuDonutSize = 8;
// An attribute used to cache the computation of device seeds, so that we don't
// constantly recompute device seeds in a cluster for a given layout.
static constexpr char kDeviceSeedForMeshDims[] =
"dtensor.device_seed_for_mesh_dims";
// Attribute that determines whether to skip XlA compilation. There are some ops
// that run on a TPU mesh but are not expected to be compiled by XLA, e.g.
// VarHandleOp, DestroyResourceOp, etc. For such an case, set this attribute
// to true on the StatefulPartitionedCallOp generated by MLIR lowering.
static constexpr char kSkipXlaCompilation[] = "_skip_xla_compilation";
// An attribute which stores the cache_key for the graph in the module. Used
// to uniquely name functions.
static constexpr char kCacheKey[] = "dtensor.cache_key";
// An attribute on Const nodes to record which argument it was originally
// from.
static constexpr char kFromArgIndex[] = "dtensor.from_arg_index";
// To record the target layout of a DTensorSend, which is computed after
// layout propagation.
static constexpr char kTargetLayoutAttr[] = "target_layout";
// To record the source layout of a DTensorRecv, which is computed after
// layout propagation.
static constexpr char kSourceLayoutAttr[] = "source_layout";
// An attribute that determines whether a tensor is a sparse tensor. If this
// attribute exists in a tensor, then this tensor is a sparse tensor.
static constexpr char kSparseValue[] = "tf._sparse";
// Attribute which stores the layouts to be applied to the elements returned by
// calling IteratorGetNextOp on a tf.data iterator.
static constexpr char kIteratorElementLayouts[] = "tf._element_layouts";
// Attribute used in tf.data ops which stores the shapes of the output elements.
static constexpr char kIteratorOutputShapes[] = "output_shapes";
// The number of list of regular tensors used to represent sparse tensors.
static constexpr int kSparseTensorNum = 3;
// Attribute which stores the environment variable value for all_reduce
// optimization group size: DTENSOR_ALLREDUCE_COMBINE_OPTIMIZATION_GROUP_SIZE.
// This represents the maximum number of AllReduce ops to merge into one op. It
// is a determining factor used during dtensor_allreduce_combine_optimization.
static constexpr char kAllReduceNumOpsInGroup[] =
"dtensor.all_reduce_combiner.num_ops_in_group";
// Attribute which stores the environment variable value for whether
// multi-device expansion is enabled: DTENSOR_ENABLE_MULTI_DEVICE_EXPANSION.
static constexpr char kEnableMultiDeviceMode[] =
"dtensor.enable_multi_device_mode";
// Attribute which stores the environment variable value for all_reduce
// optimization group size: DTENSOR_ALLREDUCE_COMBINE_OPTIMIZATION_GROUP_SIZE.
// This represents the maximum distance between two AllReduce on the compute
// graph in terms of topological level. It is a determining factor used during
// dtensor_allreduce_combine_optimization.
static constexpr char kAllReduceTopologicalDistance[] =
"dtensor.all_reduce_combiner.topological_distance";
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_CC_CONSTANTS_H_
tensorflow/dtensor/cc/default_parallel_executor.cc
+31
View File
@@ -0,0 +1,31 @@
/* Copyright 2022 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 "xla/tsl/platform/statusor.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/dtensor/cc/parallel_executor.h"
namespace tensorflow {
namespace dtensor {
StatusOr<std::unique_ptr<ParallelExecutor>> CreateDefaultParallelExecutor() {
LOG(ERROR) << __func__ << " not implemented.";
return std::unique_ptr<ParallelExecutor>();
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/cc/dstatus.h
+102
View File
@@ -0,0 +1,102 @@
/* Copyright 2022 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_DTENSOR_CC_DSTATUS_H_
#define TENSORFLOW_DTENSOR_CC_DSTATUS_H_
#include <vector>
#include "absl/strings/match.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_format.h"
#include "absl/strings/str_join.h"
#include "absl/types/optional.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
namespace dtensor {
template <typename T>
using StatusOr = tsl::StatusOr<T>;
inline absl::Status WithContext(const absl::Status& ds, absl::string_view file,
int line_number,
absl::string_view context = "") {
if (ds.ok()) {
return ds;
}
return absl::Status(ds.code(), absl::StrCat(ds.message(), "\n", file, ":",
line_number, " :: ", context));
}
template <class T>
inline StatusOr<T> WithContext(StatusOr<T>&& ds, absl::string_view file,
int line_number,
absl::string_view context = "") {
if (ds.ok()) {
return ds;
}
return absl::Status(ds.status().code(),
absl::StrCat(ds.status().message(), "\n", file, ":",
line_number, " :: ", context));
}
#define DT_CTX(dstatus, ...) \
::tensorflow::dtensor::WithContext(dstatus, __FILE__, __LINE__, #__VA_ARGS__);
#undef TF_RETURN_IF_ERROR
#define TF_RETURN_IF_ERROR(...) \
do { \
::tensorflow::Status _status = (__VA_ARGS__); \
if (!_status.ok()) { \
return ::tensorflow::dtensor::WithContext(_status, __FILE__, __LINE__); \
} \
} while (0);
#undef TF_RETURN_WITH_CONTEXT
#define TF_RETURN_WITH_CONTEXT(status, ...) \
do { \
::tensorflow::Status _status = (status); \
if (!_status.ok()) { \
return ::tensorflow::dtensor::WithContext(_status, __FILE__, __LINE__, \
##__VA_ARGS__); \
} \
} while (0);
#define DT_STATUS_MACROS_CONCAT_NAME(x, y) DT_STATUS_MACROS_CONCAT_IMPL(x, y)
#define DT_STATUS_MACROS_CONCAT_IMPL(x, y) x##y
#define DT_ASSIGN_OR_RETURN_IMPL(statusor, lhs, rexpr, ...) \
auto statusor = (rexpr); \
if (!statusor.ok()) { \
return ::tensorflow::dtensor::WithContext(statusor.status(), __FILE__, \
__LINE__, ##__VA_ARGS__); \
} \
lhs = std::move(statusor.value())
#undef TF_ASSIGN_OR_RETURN
#define TF_ASSIGN_OR_RETURN(lhs, rexpr, ...) \
DT_ASSIGN_OR_RETURN_IMPL( \
DT_STATUS_MACROS_CONCAT_NAME(_status_or_value, __COUNTER__), lhs, rexpr, \
##__VA_ARGS__)
// Undefine TF status macros to ensure users use the context macros instead
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_CC_DSTATUS_H_
tensorflow/dtensor/cc/dtensor_device.cc
+2896
View File
File diff suppressed because it is too large Load Diff
tensorflow/dtensor/cc/dtensor_device.h
+147
View File
@@ -0,0 +1,147 @@
/* Copyright 2022 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_DTENSOR_CC_DTENSOR_DEVICE_H_
#define TENSORFLOW_DTENSOR_CC_DTENSOR_DEVICE_H_
#include <string>
#include <unordered_map>
#include <vector>
#include "absl/strings/string_view.h"
#include "tensorflow/c/eager/c_api_experimental.h"
#include "tensorflow/c/tf_status.h"
namespace tensorflow {
namespace dtensor {
// Configure a custom device which runs dtensor while executing
// operations on `underlying_devices`. Allocates `device_info` and fills
// `device`, which should then be passed to
// TFE_RegisterCustomDevice. This only affects eager execution.
//
// `device_name` arg should match the `device_name` argument to
// TFE_RegisterCustomDevice, and is the name of the custom device itself
// (e.g. pass it to `tf.device` to place operations on it from Python).
// TODO(b/268241383): Remove the `status = nullptr` overload.
void AllocateDTensorDevice(absl::string_view device_name,
TFE_CustomDevice* device, void** device_info,
bool is_async, int in_flight_nodes_limit,
TF_Status* status = nullptr);
// Add a mesh to the `DTensorDevice` indicated by `device_info`.
//
// `serialized_mesh` is a serialized Mesh proto.
//
// If `is_async` is true, it indicates the DTensor operations on this mesh will
// return immediately (with "non-ready" handles), otherwise block until
// executed. This is exposed as an option for ease of debugging, and will
// typically be on.
//
// `is_host_mesh` indicates this is a CPU mesh used only for sea-of-donuts-style
// host collectives.
//
// in_flight_nodes_limit throttles the number of inflight nodes in the eager
// async executors used by DTensor. The throttling bounds the memory usage
// of an eager training loop. Python API sets this value to 8 by default.
void AddMesh(const std::string& serialized_mesh, void* device_info,
bool is_host_mesh, TF_Status* status);
// Sets a requested layout for outputs of all operations.
void ExperimentalSetDefaultLayout(const std::string& serialized_layout,
void* device_info, TF_Status* status);
void ExperimentalClearDefaultLayout(void* device_info, TF_Status* status);
// TODO(b/175928457): remove once the bug is fixed.
// Sets a requested default mesh.
void ExperimentalSetDefaultMesh(const std::string& serialized_mesh,
void* device_info, TF_Status* status);
void ExperimentalClearDefaultMesh(void* device_info, TF_Status* status);
// Determines whether tensors with a shape previously associated with only one
// layout use that layout if nothing else can be inferred.
void SetSameShapePolicy(void* device_info, bool enabled);
// Sets the global device ID-to-core ID mapping for a mesh. Global device IDs
// are equal to XLA replica IDs for the single XLA computation used by DTensor.
//
// See the comment above Mesh::tpu_core_ids() for some nuances.
void SetTPUCoreIDs(const std::string& mesh_name,
const std::vector<int>& tpu_core_ids, void* device_info,
TF_Status* status);
// TODO(b/187112276): Delete once we have the TPUCoreIDs live with Device.
void ClearTPUCoreIDs(void* device_info);
// Returns TPU core locations when given a list of TPU core IDs.
std::vector<std::vector<int>> TPUCoreIDsToLocations(
TFE_Context* context, const std::vector<int>& tpu_core_ids,
void* device_info);
// Returns TPU core IDs when given a list of TPU core locations.
std::vector<int> TPUCoreLocationsToIDs(
TFE_Context* context,
const std::vector<std::vector<int>>& tpu_core_locations, void* device_info);
// Pack `inputs` tensors into a single parallel tensor handle.
TFE_TensorHandle* Pack(TFE_Context* context, int num_inputs,
TFE_TensorHandle** inputs,
const std::string& string_layout, void* device_info,
TF_Status* status);
// Returns the raw components placed on each device of `inputs`'s mesh.
std::vector<TFE_TensorHandle*> Unpack(TFE_Context* context,
TFE_TensorHandle* input,
void* device_info, TF_Status* status);
// Returns the layout of the dtensor 'input'.
std::string FetchLayout(TFE_Context* context, TFE_TensorHandle* input,
void* device_info, TF_Status* status);
// Returns whether `input` is a dtensor.
bool IsDTensor(TFE_Context* context, TFE_TensorHandle* input, void* device_info,
TF_Status* status);
// Pack `indices`, `values`, `shapes` tensors into a SparseTensorWithLayout.
TFE_TensorHandle* SparsePack(TFE_Context* context, int num_inputs,
TFE_TensorHandle** indices,
TFE_TensorHandle** values,
TFE_TensorHandle** shapes,
const std::string& string_layout,
void* device_info, TF_Status* status);
// Returns whether `input` is a sparse dtensor. Used in `Unpack` at the python
// level to determine whether we should wrap component tensors back into a
// SparseTensor.
bool IsSparseDTensor(TFE_Context* context, TFE_TensorHandle* input,
void* device_info, TF_Status* status);
// Returns a dictionary with cache stats.
// 'hit': cache hit count,
// 'miss': cache miss count,
// 'size': number of entries in the cache.
std::unordered_map<std::string, int> GetStats(TFE_Context* context,
void* device_info,
TF_Status* status);
// Sets the layouts for the elements emitted by an iterator resource tensor.
void SetIteratorElementLayouts(TFE_Context* context, TFE_TensorHandle* input,
const std::vector<std::string>& string_layouts,
void* device_info, TF_Status* status);
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_CC_DTENSOR_DEVICE_H_
tensorflow/dtensor/cc/dtensor_device_util.cc
+1187
View File
File diff suppressed because it is too large Load Diff
tensorflow/dtensor/cc/dtensor_device_util.h
+840
View File
@@ -0,0 +1,840 @@
/* Copyright 2022 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_DTENSOR_CC_DTENSOR_DEVICE_UTIL_H_
#define TENSORFLOW_DTENSOR_CC_DTENSOR_DEVICE_UTIL_H_
#include <atomic>
#include <cstdint>
#include <map>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>
#include "absl/memory/memory.h"
#include "absl/status/status.h"
#include "llvm/Support/ExtensibleRTTI.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "tensorflow/c/eager/c_api.h"
#include "tensorflow/c/eager/parallel_device/parallel_device_lib.h"
#include "tensorflow/c/eager/tfe_context_internal.h"
#include "tensorflow/c/safe_ptr.h"
#include "tensorflow/c/tf_status.h"
#include "tensorflow/core/common_runtime/eager/context.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/lib/strings/proto_serialization.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/fingerprint.h"
#include "tensorflow/dtensor/cc/constants.h"
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/dtensor_operation.h"
#include "tensorflow/dtensor/cc/small_constant_optimization.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/cc/tensor_with_layout.h"
#include "tsl/platform/fingerprint.h"
#include "tsl/platform/refcount.h"
namespace tensorflow {
namespace dtensor {
using TensorHandlePtr = tensorflow::Safe_TFE_TensorHandlePtr;
#define RETURN_STATUS(status, code, message) \
{ \
TF_SetStatus((status), (code), (message)); \
return; \
}
#define RETURN_C_STATUS_IF_NOT_OK(cpp_status, c_status) \
{ \
auto return_if_not_ok_status = (cpp_status); \
if (!return_if_not_ok_status.ok()) { \
RETURN_STATUS((c_status), \
static_cast<TF_Code>(return_if_not_ok_status.code()), \
absl::StatusMessageAsCStr(return_if_not_ok_status)); \
} \
}
// Using a counter to uniquify instead of a new block allows `var` to declare a
// new variable.
#define ASSIGN_OR_RETURN_C_STATUS(var, cpp_status, c_status) \
ASSIGN_OR_RETURN_C_STATUS_IMPL( \
TF_STATUS_MACROS_CONCAT_NAME(_dtensor_status_or_value, __COUNTER__), \
var, cpp_status, c_status)
#define ASSIGN_OR_RETURN_C_STATUS_IMPL(statusor, var, cpp_status, c_status) \
auto statusor = (cpp_status); \
RETURN_C_STATUS_IF_NOT_OK(statusor.status(), (c_status)); \
var = std::move(statusor.value());
struct TranslatedFunction {
// Mesh for which specified function will run.
Mesh function_mesh;
// StatefulPartitionedCall op to run the mesh function.
const Node* node_to_execute = nullptr;
// Maps i-th local input index to input index in global graph.
std::vector<int> input_index_map;
// Maps i-th local output to output index of global graph.
std::vector<int> output_index_map;
// Original function name in the graph.
std::string function_name;
// Translated function name to be called.
std::string translated_function_name;
// For resource ops, layouts of resource handles are inferred lazily
// during SPMD expansion of resource assign ops. In that case,
// inferred layouts of resource handles are attached to arg nodes
// of the returned graph.
std::map<int, Layout> resource_input_layouts;
// Record some metadata for output of a shape op. This would help recover
// local shape on future operations over the Tensor.
std::map<int, Layout> shape_output_metadata;
std::vector<Layout> output_layouts;
// Local shapes inferred for function outputs; these may be partially known.
std::vector<PartialTensorShape> local_output_shapes;
// Output data types.
std::vector<TF_DataType> output_dtypes;
// Number of local outputs for each layout.
std::vector<std::int64_t> num_local_outputs;
};
struct ExecutionFunctions {
// Stores information about all functions to execute for provided computation.
std::vector<TranslatedFunction> function_list;
// Number of device ids args added to translated functions.
// During translation, we insert one device id arg node per mesh.
// For a single mesh function, it equals 1.
// For a multi-mesh function (e.g. pipelining), it equals the number of
// meshes.
int num_device_ids;
// Mesh fingerprint of function_list. Set only when ExecutionFunctions refers
// to a function for performance reason, since an eager op doesn't use it.
uint64_t function_mesh_fingerprint = 0;
};
class TensorWithLayoutTf
: public llvm::RTTIExtends<TensorWithLayoutTf, TensorWithLayout> {
public:
// Broadcast a single non-parallel tensor onto `mesh` with a fully replicated
// sharding spec. Does not take ownership of `tensor`. The tensor must not
// already be on a DTensorDevice.
static std::unique_ptr<TensorWithLayoutTf> Broadcast(
TFE_Context* context, TFE_TensorHandle* tensor, const Mesh& target_mesh,
TF_Status* status);
// Given an already-parallel tensor, wraps it with a mesh and a layout.
static StatusOr<std::unique_ptr<TensorWithLayoutTf>> Wrap(
std::vector<TensorHandlePtr>&& tensors, const Layout& layout,
std::optional<std::vector<int64_t>>&& shape);
// Given a single tensor, wraps it with a single device layout.
static std::unique_ptr<TensorWithLayoutTf> Wrap(TensorHandlePtr single_tensor,
const Layout& layout,
TF_Status* status);
// Creates a dummy TensorWithLayoutTf without holding a ParallelTensor.
static std::unique_ptr<TensorWithLayoutTf> Dummy(
const std::vector<int64_t>& local_shape, TF_DataType dtype,
const Layout& layout);
~TensorWithLayoutTf() override = default;
const Layout& layout() const override { return layout_; }
TensorType tensor_type() const override { return TensorType::kDense; }
TF_DataType dtype() const override {
return dtype_.has_value() ? *dtype_
: TFE_TensorHandleDataType(tensors_[0].get());
}
// Encodes the NodeDef via provided builder, if applicable.
void EncodeAttributes(tensorflow::NodeDefBuilder& builder) const override {}
tensorflow::Fprint128 CacheKey() const override;
TFE_TensorHandle* get_tensor(size_t index) const override {
return tensors_[index].get();
}
size_t num_tensors() const override {
return layout_.IsSingleDevice() ? 1 : tensors_.size();
}
std::vector<TFE_TensorHandle*> tensors() const {
std::vector<TFE_TensorHandle*> result;
result.reserve(tensors_.size());
for (const TensorHandlePtr& tensor : tensors_) {
result.emplace_back(tensor.get());
}
return result;
}
TFE_TensorHandle* single_tensor() const {
return layout_.IsSingleDevice() ? get_tensor(0) : nullptr;
}
std::string SummarizeValue() const override;
std::string DebugString() const override;
std::vector<int64_t> global_shape() const override {
return layout_.GlobalShapeFromLocalShape(local_shape_, &local_shapes_);
}
ConstValueNode* const_value_node() const override {
return const_value_node_.get();
}
// llvm::RTTIExtends ID.
static char ID; // NOLINT
protected:
TensorWithLayoutTf(std::vector<TensorHandlePtr>&& tensors,
const Layout& layout,
const std::vector<int64_t>& local_shape,
const std::vector<std::vector<int64_t>>& local_shapes,
std::optional<TF_DataType> dtype = std::nullopt,
std::optional<NodeDef> const_value = std::nullopt)
: tensors_(std::move(tensors)),
layout_(layout),
local_shape_(local_shape),
local_shapes_(std::move(local_shapes)),
dtype_(dtype) {
const_value_node_ = std::make_unique<ConstValueNode>(const_value);
}
TensorWithLayoutTf(TensorHandlePtr&& single_tensor, const Layout& layout,
const std::vector<int64_t>& local_shape,
std::optional<TF_DataType> dtype = std::nullopt,
std::optional<NodeDef> const_value = std::nullopt)
: tensors_([&single_tensor] {
std::vector<TensorHandlePtr> result;
result.emplace_back(std::move(single_tensor));
return result;
}()),
layout_(layout),
local_shape_(local_shape),
dtype_(dtype) {
const_value_node_ = std::make_unique<ConstValueNode>(const_value);
}
std::vector<TensorHandlePtr> tensors_;
Layout layout_;
// The local shape of tensors placed on each of `tensor_`'s component devices.
std::vector<int64_t> local_shape_;
// The local shape of each individual tensor in `tensors_`.
// Initialized only when there is dynamic shape.
std::vector<std::vector<int64_t>> local_shapes_;
// dtype of tensor_. Empty if the layout is Single Device.
std::optional<TF_DataType> dtype_;
std::unique_ptr<ConstValueNode> const_value_node_;
};
// Extension of TensorWithLayout which holds resource handle with layout.
//
// The major differences are
// 1. The layout, shape, dtype are lazily set as they are unavailable upon
// creation.
// 2. Small const optimization should be disabled.
class ResourceHandleWithLayout
: public llvm::RTTIExtends<ResourceHandleWithLayout, TensorWithLayoutTf> {
public:
// Similar to `Wrap` in `TensorWithLayoutTf` but for resource handle.
static StatusOr<std::unique_ptr<ResourceHandleWithLayout>> Wrap(
std::vector<TensorHandlePtr>&& tensors, const Layout& layout,
std::optional<std::vector<int64_t>>&& shape);
// Similar to `Dummy` in `TensorWithLayoutTf` but for resource handle.
static std::unique_ptr<ResourceHandleWithLayout> Dummy(
const std::vector<int64_t>& local_shape, const Layout& layout);
// The layout of uninitialized resource tensors, or the layout of the tensor
// contained in an initialized resource.
const Layout& layout() const override {
return dereferenced_layout_.has_value() ? dereferenced_layout_.value()
: layout_;
}
TensorType tensor_type() const override { return TensorType::kResource; }
TF_DataType dtype() const override {
return dtype_.has_value() ? *dtype_
: TFE_TensorHandleDataType(tensors_[0].get());
}
void EncodeAttributes(tensorflow::NodeDefBuilder& builder) const override;
tensorflow::Fprint128 CacheKey() const override;
// Updates the layout for the tensors.
absl::Status UpdateLayout(const Layout& new_layout);
// Updates the element layouts for the tensors.
absl::Status UpdateElementLayouts(const std::vector<Layout>& layouts) {
dereferenced_element_layouts_.emplace(layouts);
return absl::OkStatus();
}
// Updates the local shape and dtype of the tensors.
absl::Status UpdateShapeAndDType(const TensorShapeProto& shape,
const DataType& dtype) {
set_dereferenced_shape(shape);
set_dereferenced_dtype(dtype);
return absl::OkStatus();
}
ConstValueNode* const_value_node() const override { return nullptr; }
void set_dereferenced_shape(const TensorShapeProto& shape) {
dereferenced_shape_.emplace(shape);
}
void set_dereferenced_dtype(const DataType& dtype) {
dereferenced_dtype_.emplace(dtype);
}
const std::optional<std::vector<Layout>>& dereferenced_element_layouts()
const {
return dereferenced_element_layouts_;
}
const std::optional<TensorShapeProto>& dereferenced_shape() const {
return dereferenced_shape_;
}
const std::optional<DataType>& dereferenced_dtype() const {
return dereferenced_dtype_;
}
// llvm::RTTIExtends ID.
static char ID; // NOLINT
private:
ResourceHandleWithLayout(std::vector<TensorHandlePtr>&& tensors,
const Layout& layout,
const std::vector<int64_t>& local_shape)
: llvm::RTTIExtends<ResourceHandleWithLayout, TensorWithLayoutTf>(
std::move(tensors), layout, local_shape, {}, TF_RESOURCE) {}
// The layout of the tensor pointed to by this handle, if any.
std::optional<Layout> dereferenced_layout_;
// The layouts of the tensors emitted by this resource handle if it is an
// iterator resource.
std::optional<std::vector<Layout>> dereferenced_element_layouts_;
// The shape and dtype of the tensor pointed to by this resource tensor.
std::optional<TensorShapeProto> dereferenced_shape_;
std::optional<DataType> dereferenced_dtype_;
};
// TensorWithLayout for SparseTensors.
//
// The main difference between this and TensorWithLayout is this
// contains 3 lists of tensors as opposed to one (values, indices, shapes).
// The shapes of the SparseTensors will always be the dense view of the shapes,
// and thus will have no difference with the TensorWithLayout in terms of
// shapes.
class SparseTensorWithLayout
: public llvm::RTTIExtends<SparseTensorWithLayout, TensorWithLayoutTf> {
public:
static StatusOr<std::unique_ptr<SparseTensorWithLayout>> Wrap(
std::unique_ptr<parallel_device::ParallelTensor> indices_tensor,
std::unique_ptr<parallel_device::ParallelTensor> values_tensor,
std::unique_ptr<parallel_device::ParallelTensor> shapes_tensor,
const Layout& layout, const std::vector<int64_t>& local_shape);
// A dummy TensorWithLayout without holding a ParallelTensor.
static std::unique_ptr<SparseTensorWithLayout> Dummy(
const std::vector<int64_t>& local_shape, const Layout& layout) {
return absl::WrapUnique(new SparseTensorWithLayout(
/*indices=*/nullptr, /*values=*/nullptr, /*dense_shapes=*/nullptr,
layout, local_shape));
}
// Add attribute '_sparse' to the NodeDefBuilder so that the mlir::Value
// that originate from SparseTensorWithLayout are marked as '_sparse'.
void EncodeAttributes(tensorflow::NodeDefBuilder& builder) const override {
builder.Attr("_sparse", true);
}
TensorType tensor_type() const override { return TensorType::kSparse; }
size_t num_tensors() const override {
return kSparseTensorNum * indices_->num_tensors();
}
TFE_TensorHandle* get_tensor(size_t index) const override;
std::string SummarizeValue() const override;
std::string DebugString() const override;
TF_DataType dtype() const override;
parallel_device::ParallelTensor* indices() const { return indices_.get(); }
parallel_device::ParallelTensor* values() const { return values_.get(); }
parallel_device::ParallelTensor* dense_shapes() const {
return dense_shapes_.get();
}
ConstValueNode* const_value_node() const override { return nullptr; }
// llvm::RTTIExtends ID.
static char ID; // NOLINT
private:
SparseTensorWithLayout(
std::unique_ptr<parallel_device::ParallelTensor> indices,
std::unique_ptr<parallel_device::ParallelTensor> values,
std::unique_ptr<parallel_device::ParallelTensor> dense_shapes,
const Layout& layout, const std::vector<int64_t>& local_shape,
std::optional<TF_DataType> dtype = std::nullopt,
std::optional<NodeDef> const_value = std::nullopt)
: llvm::RTTIExtends<SparseTensorWithLayout, TensorWithLayoutTf>(
std::vector<TensorHandlePtr>(), layout, local_shape, {}),
indices_(std::move(indices)),
values_(std::move(values)),
dense_shapes_(std::move(dense_shapes)) {}
std::unique_ptr<parallel_device::ParallelTensor> indices_;
std::unique_ptr<parallel_device::ParallelTensor> values_;
std::unique_ptr<parallel_device::ParallelTensor> dense_shapes_;
};
std::unique_ptr<TensorWithLayoutTf> CreateDummyTensorWithLayout(
const std::vector<int64_t>& local_shape, TF_DataType dtype,
const Layout& layout);
// Creates a DTensor from one or more tensor handles and a compatible
// layout. Optionally accepts a `shape` argument that overrides the
// actual shape of the underlying tensors; this argument should be
// provided when there's a possibility of the inferred shape from
// differing from the actual shape (like when it is dynamic).
StatusOr<std::unique_ptr<TensorWithLayoutTf>> CreateTensorWithLayout(
std::vector<TensorHandlePtr>&& tensor, const Layout& layout,
std::optional<std::vector<int64_t>>&& shape = std::nullopt);
template <typename T>
std::string ShapeToDebugString(const std::vector<T> shape_vector) {
std::vector<int64_t> cast_shape(shape_vector.begin(), shape_vector.end());
tensorflow::PartialTensorShape shape;
if (!tensorflow::PartialTensorShape::MakePartialShape(
cast_shape.data(), cast_shape.size(), &shape)
.ok()) {
return "<error displaying shape>";
} else {
return shape.DebugString();
}
}
// Internal class with shared functions for every ExecutableManager<T>.
class ExecutableManagerImpl {
template <typename T>
friend class ExecutableManager;
public:
absl::flat_hash_map<int, NodeDef> GetConstantFoldableTensors(
const std::vector<TensorWithLayout*>& inputs);
// Cache key for dtensor operation name, which includes the op name
// and the input shapes. This is needed as a higher level cache for constant
// folding.
tensorflow::Fprint128 CacheKeyForDTensorOperation(
const DTensorOperation& doperation) const;
private:
ExecutableManagerImpl() = default;
};
struct ExecutionManagerStats {
int64_t hits; // number of hits.
int64_t misses; // number of misses.
int64_t size; // size of cache (number of entries).
};
// Template Class that holds information about DTensor executable ran, including
// cached lowered executable and constant folding input information per
// function.
//
//
// The caching policy for constant folded inputs is the following:
// In the first call to a function, we assume that all the inputs that
// are constant foldable are constant folded and save these values. In the
// next call to the same function call, we compare the values of constant
// folded inputs to the previous constant folded inputs. We disable constant
// folding for the changed values, and save these new inputs.
// TODO(b/169348205) Support cache eviction if the cache gets bloated.
template <typename T>
class ExecutableManager : public tsl::core::WeakRefCounted {
public:
ExecutableManager() = default;
// Caches the executable with ParallelExecutable.
const T* AddCachedExecutable(tensorflow::Fprint128 cache_key, T executable);
// Removes the executable.
void Remove(tensorflow::Fprint128 cache_key);
// Returns the cache key and the cached lowered executable for the function.
// Returns a nullptr for the lowered executable if there is a cache miss.
// Upon a cache miss, this will save some metadata about the function
// and the small inputs to keep track of information for constant folding.
StatusOr<std::pair<tensorflow::Fprint128, const T*>> GetCachedExecutable(
const DTensorOperation& doperation, const NameAttrList& attributes,
const std::vector<TensorWithLayout*>& inputs,
const std::vector<const Layout*>& output_layouts);
// Returns the cached lowered graph for the function.
// Returns a nullptr for the lowered graph if there is a cache miss.
// This Get operation has no side effect.
const T* GetCachedExecutableSimple(tensorflow::Fprint128 cache_key);
// Returns whether the input at `input_index` should be constant
// folded into function `doperation`. An input is not constant folded if we
// have ran this function at least twice and the small input value changed
// across separate runs.
StatusOr<bool> ShouldFoldInput(const DTensorOperation& doperation,
const std::vector<TensorWithLayout*>& inputs,
int input_index) const;
// Returns the current Stats of the execution manager.
// The result is a snapshot at the moment of the call.
ExecutionManagerStats GetStats() const {
ExecutionManagerStats stats;
stats.hits = stats_.hits;
stats.misses = stats_.misses;
// A reader Lock is probably more suitable, but this code branch is
// barely executed.
mutex_lock lock(mu_);
stats.size = function_cache_.size();
return stats;
}
private:
// Generates a cache key for the graph, including its attributes,
// inputs, and outputs.
StatusOr<tensorflow::Fprint128> CacheKeyForGraph(
const DTensorOperation& doperation, const NameAttrList& attributes,
const std::vector<TensorWithLayout*>& inputs,
const std::vector<const Layout*>& output_layouts);
// Returns true for a missing entry in the small inputs cache.
bool UpdateDTensorOpAndSmallInputsCache(
const DTensorOperation& doperation,
const std::vector<TensorWithLayout*>& inputs);
mutable mutex mu_;
mutable mutex dtensor_op_and_small_inputs_mu_;
// Maps the hash of a graph with the lowered graph.
absl::flat_hash_map<tensorflow::Fprint128, T, tensorflow::Fprint128Hasher>
function_cache_ TF_GUARDED_BY(mu_);
// Maps the hash of dtensor_operation and its input shapes to a map
// representing the small constant indices and values to the function. The
// small constant indices are saved to make faster comparisons for constant
// folding validation.
absl::flat_hash_map<tensorflow::Fprint128, absl::flat_hash_map<int, NodeDef>,
tensorflow::Fprint128Hasher>
dtensor_op_and_small_inputs_
TF_GUARDED_BY(dtensor_op_and_small_inputs_mu_);
ExecutableManagerImpl executable_manager_impl_;
struct {
std::atomic<int64_t> hits = 0;
std::atomic<int64_t> misses = 0;
} stats_;
};
// Returns the shape of a given tensor.
StatusOr<std::vector<int64_t>> GetTensorShapeAsVector(
const tensorflow::PartialTensorShape& shape);
// Returns the shape of a given tensor.
StatusOr<std::vector<int64_t>> GetTensorShapeAsVector(TFE_TensorHandle* tensor);
absl::Status InferOutputLayouts(const DTensorOperation& doperation,
const NameAttrList& attributes,
const std::optional<Layout>& default_layout,
tensorflow::Graph* graph,
std::vector<const Layout*>* output_layouts);
// Creates a Graph with _Arg and _Retval nodes surrounding an
// `operation_name`-type node.
absl::Status PrepareGraphForMlir(
const ExecutableManager<mlir::OwningOpRef<mlir::ModuleOp>>& module_manager,
const std::vector<TensorWithLayout*>& inputs,
const DTensorOperation& doperation,
const tensorflow::FunctionLibraryDefinition& flib_def,
const NameAttrList& attributes,
const std::vector<const Layout*>& output_layouts, tensorflow::Graph* graph,
std::vector<PartialTensorShape>* global_output_shapes);
// Returns set of functions to run to execute DTensor computation.
StatusOr<ExecutionFunctions> IdentifyAllFunctionsToExecute(
const tensorflow::Graph& graph,
const std::vector<PartialTensorShape>& global_output_shapes);
// For functions with control outputs, add identity nodes between
// StatefulPartitionedCall and _Retvals, in order to preserve control output
// dependencies after StatefulPartitionedCall is inlined at runtime.
// Consider calling this in PrepareGraphForMlir, once the identity nodes won't
// be dropped during MLIR lowering.
// TODO(b/171265131): fix the underlying issue to avoid inserting identity
// nodes.
absl::Status MaybeInsertIdentityNodes(const FunctionDef* function_def,
Graph* graph);
// Add DTensor specific function attributes to be compatible with eager runtime.
void AddDTensorFunctionAttr(FunctionDef& function_def);
////////////////////////////////////////////////////////////////////////////////
// Implementation details for ExecutableManager<T>
// Thread safe method.
// Generates a cache key for the graph, including its attributes,
// inputs, and outputs.
// Cache key computation should consider all features of an op that affects
// the SPMD lowering. The cache keys of two ops must be different if the
// translated functions are different.
// - op name and attr
// - input shapes and layouts
// - default layout of outputs.
// - default mesh.
// - values of constant foldable inputs.
template <typename T>
StatusOr<tensorflow::Fprint128> ExecutableManager<T>::CacheKeyForGraph(
const DTensorOperation& doperation, const NameAttrList& attributes,
const std::vector<TensorWithLayout*>& inputs,
const std::vector<const Layout*>& output_layouts) {
tensorflow::Fprint128 cache_key = tensorflow::Fingerprint128(doperation.name);
std::string serialized;
SerializeToStringDeterministic(attributes, &serialized);
cache_key =
FingerprintCat128(cache_key, tensorflow::Fingerprint128(serialized));
cache_key = FingerprintCat128(
cache_key,
tensorflow::Fingerprint128(doperation.default_mesh.ToString()));
// Higher level cache based on operation name and input shapes.
for (int i = 0; i < inputs.size(); ++i) {
TF_ASSIGN_OR_RETURN(bool should_fold_input,
ShouldFoldInput(doperation, inputs, i));
if (!should_fold_input && inputs[i]->const_value_node()) {
inputs[i]->const_value_node()->reset_const_value();
}
cache_key = FingerprintCat128(
cache_key, tensorflow::Fingerprint128(absl::StrFormat("%x", i)));
cache_key = FingerprintCat128(cache_key, inputs[i]->CacheKey());
}
for (int output_index = 0; output_index < output_layouts.size();
++output_index) {
if (output_layouts[output_index]) {
cache_key = FingerprintCat128(cache_key, output_index);
cache_key = FingerprintCat128(
cache_key,
tensorflow::Fingerprint128(output_layouts[output_index]->ToString()));
}
}
return cache_key;
}
// Thread-safe method.
template <typename T>
StatusOr<std::pair<tensorflow::Fprint128, const T*>>
ExecutableManager<T>::GetCachedExecutable(
const DTensorOperation& doperation, const NameAttrList& attributes,
const std::vector<TensorWithLayout*>& inputs,
const std::vector<const Layout*>& output_layouts) {
TF_ASSIGN_OR_RETURN(
tensorflow::Fprint128 cache_key,
CacheKeyForGraph(doperation, attributes, inputs, output_layouts));
{
mutex_lock lock(mu_);
// Early return if we have a cache hit.
if (auto iter = function_cache_.find(cache_key);
iter != function_cache_.end()) {
stats_.hits++;
return {{cache_key, &iter->second}};
}
}
bool missed = UpdateDTensorOpAndSmallInputsCache(doperation, inputs);
if (missed) {
stats_.misses++;
return {{cache_key, nullptr}};
}
// Generate a new cache key since we updated small const inputs which change
// the cache key.
TF_ASSIGN_OR_RETURN(cache_key, CacheKeyForGraph(doperation, attributes,
inputs, output_layouts));
stats_.misses++;
return {{cache_key, nullptr}};
}
template <typename T>
bool ExecutableManager<T>::UpdateDTensorOpAndSmallInputsCache(
const DTensorOperation& doperation,
const std::vector<TensorWithLayout*>& inputs) {
const tensorflow::Fprint128 doperation_hash =
executable_manager_impl_.CacheKeyForDTensorOperation(doperation);
mutex_lock lock(dtensor_op_and_small_inputs_mu_);
// Save the constant folded inputs to this doperation if we have not seen
// this before. This is needed so that in the next call to this operation,
// we can compare these inputs to confirm which one is indeed a constant.
auto doperation_iter = dtensor_op_and_small_inputs_.find(doperation_hash);
if (doperation_iter == dtensor_op_and_small_inputs_.end()) {
dtensor_op_and_small_inputs_.insert(
{doperation_hash,
executable_manager_impl_.GetConstantFoldableTensors(inputs)});
return true;
}
// If we are here, then we have ran this function before but constant folded
// some input(s) when it was not a constant input i.e. one of the small
// value to this function input changed. So mark those changed values as
// non-constant.
absl::flat_hash_map<int, NodeDef>& previous_small_inputs =
doperation_iter->second;
std::vector<int> non_constant_indices;
for (auto const& [index, previous_small_input] : previous_small_inputs) {
// Some Ops the number of inputs can vary. We'll just skip updating them.
if (index >= inputs.size()) continue;
auto* const_value_node = inputs[index]->const_value_node();
if (const_value_node == nullptr) {
continue;
}
if (const_value_node->const_value().has_value()) {
if (NodeDefsHaveDifferentTensorProto(
previous_small_input, const_value_node->const_value().value())) {
non_constant_indices.push_back(index);
}
}
}
for (int non_constant_index : non_constant_indices) {
previous_small_inputs.erase(non_constant_index);
}
return false;
}
// Thread-safe method.
template <typename T>
const T* ExecutableManager<T>::GetCachedExecutableSimple(
tensorflow::Fprint128 cache_key) {
mutex_lock lock(mu_);
auto iter = function_cache_.find(cache_key);
if (iter == function_cache_.end()) {
stats_.misses++;
return nullptr;
}
stats_.hits++;
return &iter->second;
}
template <typename T>
const T* ExecutableManager<T>::AddCachedExecutable(
tensorflow::Fprint128 cache_key, T executable) {
mutex_lock lock(mu_);
return &function_cache_.insert({cache_key, std::move(executable)})
.first->second;
}
template <typename T>
void ExecutableManager<T>::Remove(tensorflow::Fprint128 cache_key) {
mutex_lock lock(mu_);
auto iter = function_cache_.find(cache_key);
if (iter != function_cache_.end()) {
function_cache_.erase(iter);
}
}
template <typename T>
StatusOr<bool> ExecutableManager<T>::ShouldFoldInput(
const DTensorOperation& doperation,
const std::vector<TensorWithLayout*>& inputs, const int input_index) const {
const auto input = inputs[input_index];
const bool can_fold = input->const_value_node() &&
input->const_value_node()->const_value().has_value();
// For eager ops, assume the inputs are constant foldable.
if (!doperation.is_func()) {
// Fold if we are in a function or if a special eager op.
// TODO(b/270762002): Think about how to generalize this so it does not
// depend on operation_name. For example, we can check the max abs value of
// the tensor value.
if (doperation.name == absl::string_view("StatelessRandomUniform") ||
doperation.name == absl::string_view("StatelessRandomUniformFullInt") ||
doperation.name == absl::string_view("StatelessRandomNormal") ||
doperation.name == absl::string_view("StatelessTruncatedNormal")) {
// For all stateless rng ops, we avoid fold seed (input_index==1) in
// graph. This is an important optimization to avoid unnecessary MLIR SPMD
// lowering and TPU compilation during model parameters initialization
// process. which typically have the same shape for rng ops but different
// seeds.
return can_fold && (input_index != 1);
}
// Certain Ops we shall never fold in their inputs. Enable caching to reduce
// sizes of the graphs. This list is incomplete.
// FIXME(b/270762002): We only need constant folding for args that are
// matched against Constants in MLIR.
if (doperation.name != absl::string_view("Identity") &&
doperation.name != absl::string_view("DivNoNan") &&
doperation.name != absl::string_view("CopyToMesh") &&
doperation.name != absl::string_view("CopyToMeshGrad") &&
doperation.name != absl::string_view("Relayout") &&
doperation.name != absl::string_view("RelayoutGrad")) {
return can_fold;
}
}
const tensorflow::Fprint128 doperation_hash =
executable_manager_impl_.CacheKeyForDTensorOperation(doperation);
mutex_lock lock(dtensor_op_and_small_inputs_mu_);
// If we didn't see this doperation before then optimisticly assume this is
// foldable. The input at `input_index` is foldable only if it is one of the
// indices we have saved as the small inputs.
auto doperation_iter = dtensor_op_and_small_inputs_.find(doperation_hash);
return can_fold && (doperation_iter == dtensor_op_and_small_inputs_.end() ||
doperation_iter->second.contains(input_index));
}
// ExecutionFunctions manager can not check if the input is foldable.
template <>
StatusOr<bool> ExecutableManager<ExecutionFunctions>::ShouldFoldInput(
const DTensorOperation& doperation,
const std::vector<TensorWithLayout*>& inputs, int input_index) const;
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_CC_DTENSOR_DEVICE_UTIL_H_
tensorflow/dtensor/cc/dtensor_graph_to_mlir_pass.cc
+158
View File
@@ -0,0 +1,158 @@
/* Copyright 2022 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/dtensor/cc/dtensor_graph_to_mlir_pass.h"
#include <memory>
#include <utility>
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Dialect.h" // from @llvm-project
#include "mlir/IR/SymbolTable.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "mlir/InitAllExtensions.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/jit/flags.h"
#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/convert_type.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/device_util.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
#include "tensorflow/compiler/mlir/tf2xla/api/v2/graph_to_tf_executor.h"
#include "xla/status_macros.h"
#include "xla/tsl/platform/status.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/common_runtime/optimization_registry.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/graph/algorithm.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/util/dump_graph.h"
#include "tensorflow/dtensor/cc/constants.h"
#include "tensorflow/dtensor/cc/dtensor_utils.h"
#include "tensorflow/dtensor/mlir/dtensor_dialect/ir/dialect.h"
#include "tensorflow/dtensor/mlir/dtensor_mlir_passes.h"
#include "tensorflow/dtensor/mlir/ir/tf_dtensor.h"
namespace tensorflow {
DTensorMlirPassRunner::DTensorMlirPassRunner()
: pass_manager_(&context_), logging_enabled_(false) {
mlir::DialectRegistry registry;
mlir::registerAllExtensions(registry);
mlir::RegisterAllTensorFlowDialects(registry);
context_.appendDialectRegistry(registry);
logging_enabled_ = dtensor::MaybeEnableLogging(&pass_manager_);
if (logging_enabled_) pass_manager_.getContext()->enableMultithreading();
// TODO(hinsu, hongjunchoi): Figure out a better place to explicitly enable
// the MLIR bridge.
// Explicitly enable MLIR bridge as DTensor introduces some ops like
// XlaAllReduce are only supported in MLIR.
GetMlirCommonFlags()->tf_mlir_enable_mlir_bridge =
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_ENABLED;
// Creates a pipeline that include each DTensor related passes.
mlir::TF::StandardPipelineOptions pipeline_options;
dtensor::CreateDTensorMLIRPass(pipeline_options, &pass_manager_);
}
absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>>
DTensorMlirPassRunner::ImportGraphToMlir(
const DeviceSet& device_set, absl::string_view name, bool is_func,
const dtensor::Mesh& default_mesh,
const FunctionLibraryDefinition& flib_def, const Graph& graph,
Fprint128 cache_key) {
GraphDebugInfo debug_info;
GraphImportConfig import_config;
import_config.graph_as_function = true;
// DTensor relies on importing with shape_inference to work properly ATM.
// Make it explicit so that we're not affected by potential flipping of the
// flag.
import_config.enable_shape_inference = true;
// Graph pruning will prune away an op (may be side effecting) if the op is
// not reachable from a fetch/result or target/control ret. With how the entry
// function/Graph is created, it is possible if the op has no data results. To
// make sure this op does not get pruned away, the op is defined as a
// target/control ret.
import_config.control_outputs = {"eager_operation"};
// Imports GraphDef to TF MLIR.
absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> module_ref =
tensorflow::tf2xla::v2::ConvertGraphToTfExecutor(
graph, debug_info, flib_def, import_config, &context_);
// Adds DTensor attributes to ModuleOp.
mlir::ModuleOp module = module_ref.value().get();
AddDevicesToOp(module, &device_set);
module->setAttr(dtensor::kCustomDefaultMeshAttr,
mlir::StringAttr::get(&context_, default_mesh.ToString()));
// Tag the module for logging or not depending on flag.
if (!is_func && !dtensor::LogOpByOp(name))
module->setAttr(dtensor::kDoNotLog, mlir::UnitAttr::get(&context_));
module->setAttr(dtensor::kEagerOperationName,
mlir::StringAttr::get(&context_, name));
// Set the cache key for the module as an attribute. This attribute will be
// used to rename all private functions in the module (by appending the
// cache key) so they have unique names.
module->setAttr(
dtensor::kCacheKey,
mlir::StringAttr::get(&context_, absl::StrCat("_", cache_key.low64, "_",
cache_key.high64)));
// Set the all_reduce_combine_optimization environment variable as module
// attribute
int group_size = dtensor::AllReduceCombineOptimizationGroupSize();
module->setAttr(
dtensor::kAllReduceNumOpsInGroup,
mlir::IntegerAttr::get(mlir::IntegerType::get(&context_, /*width=*/64),
group_size));
int topo_dist = dtensor::AllReduceCombineOptimizationTopologicalDistance();
module->setAttr(
dtensor::kAllReduceTopologicalDistance,
mlir::IntegerAttr::get(mlir::IntegerType::get(&context_, /*width=*/64),
topo_dist));
if (dtensor::EnableMultiDeviceMode()) {
module->setAttr(dtensor::kEnableMultiDeviceMode,
mlir::BoolAttr::get(&context_, true));
}
return module_ref;
}
absl::Status DTensorMlirPassRunner::Run(mlir::ModuleOp module) {
// Executes and collects results from the passes.
mlir::StatusScopedDiagnosticHandler diag_handler(&context_);
if (logging_enabled_ && !module->hasAttr(dtensor::kDoNotLog))
pass_manager_.getContext()->disableMultithreading();
mlir::LogicalResult result = pass_manager_.run(module);
(void)result;
TF_RETURN_IF_ERROR(diag_handler.ConsumeStatus());
if (logging_enabled_) pass_manager_.getContext()->enableMultithreading();
return absl::OkStatus();
}
} // namespace tensorflow
tensorflow/dtensor/cc/dtensor_graph_to_mlir_pass.h
+59
View File
@@ -0,0 +1,59 @@
/* Copyright 2022 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_DTENSOR_CC_DTENSOR_GRAPH_TO_MLIR_PASS_H_
#define TENSORFLOW_DTENSOR_CC_DTENSOR_GRAPH_TO_MLIR_PASS_H_
#include <memory>
#include "absl/container/flat_hash_set.h"
#include "absl/strings/string_view.h"
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "tensorflow/core/common_runtime/device_set.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/platform/fingerprint.h"
#include "tensorflow/dtensor/mlir/ir/tf_dtensor.h"
namespace tensorflow {
class DTensorMlirPassRunner {
public:
DTensorMlirPassRunner();
// Imports Graph to MLIR module in tf_execute Dialect with DTensor attributes.
absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> ImportGraphToMlir(
const DeviceSet& device_set, absl::string_view name, bool is_func,
const dtensor::Mesh& default_mesh,
const FunctionLibraryDefinition& flib_def, const Graph& graph,
Fprint128 cache_key);
// Transforms input MLIR module with DTensor Pass pipeline.
absl::Status Run(mlir::ModuleOp module);
private:
// N.B. op_registration_ must be initialized before context/pass-manager to
// ensure DTensor operations are available during optimization passes.
bool op_registration_ = mlir::TF::RegisterDTensorTFOps();
mlir::MLIRContext context_;
mlir::PassManager pass_manager_;
bool logging_enabled_;
};
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_CC_DTENSOR_GRAPH_TO_MLIR_PASS_H_
tensorflow/dtensor/cc/dtensor_meta_ops.cc
+225
View File
@@ -0,0 +1,225 @@
/* Copyright 2022 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 "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
namespace tensorflow {
namespace dtensor {
REGISTER_OP("DTensorAllReduce")
.Input("input: T")
.Input("group_assignment: int32")
.Output("output: T")
.Attr(
"T: {half, bfloat16, float, float64, int8, uint8, int32, uint32, "
"int64, uint64, bool}")
.Attr("reduce_op: {'Min', 'Max', 'Mul', 'Add', 'Mean', 'Any', 'All'}")
.Attr("device_type: string") // e.g. "/device:TPU"
.SetShapeFn(shape_inference::UnchangedShape);
REGISTER_OP("DTensorReduceScatter")
.Input("input: T")
.Input("group_assignment: int32")
.Input("scatter_dimension: int32")
.Output("output: T")
.Attr("T: {half, bfloat16, float, int32, uint32, int64, bool}")
.Attr("reduce_op: {'Min', 'Max', 'Mul', 'Add', 'Mean', 'Any', 'All'}")
.Attr("device_type: string") // e.g. "/device:TPU"
.SetShapeFn(shape_inference::ReduceScatterShape);
REGISTER_OP("DTensorAllScatter")
.Input("input: T")
.Output("output: T")
.Attr(
"T: {half, bfloat16, float, float64, int8, uint8, int32, uint32, "
"int64, uint64, bool, string}")
.Attr("input_layout: string")
.Attr("output_layout: string")
.SetShapeFn([](shape_inference::InferenceContext* c) -> absl::Status {
shape_inference::ShapeHandle in = c->input(0);
if (!c->RankKnown(in)) {
// Input shape unknown, so set unknown output shape.
c->set_output(0, in);
return absl::OkStatus();
}
std::string input_layout_string;
std::string output_layout_string;
TF_RETURN_IF_ERROR(c->GetAttr("input_layout", &input_layout_string));
TF_RETURN_IF_ERROR(c->GetAttr("output_layout", &output_layout_string));
TF_ASSIGN_OR_RETURN(Layout input_layout,
Layout::FromString(input_layout_string));
TF_ASSIGN_OR_RETURN(Layout output_layout,
Layout::FromString(output_layout_string));
if (c->Rank(in) != input_layout.rank() ||
c->Rank(in) != output_layout.rank()) {
return absl::InvalidArgumentError(absl::StrCat(
"Input tensor rank and layout ranks do not agree: input rank ",
c->Rank(in), " input layout rank ", input_layout.rank(),
" output "
"layout rank ",
output_layout.rank()));
}
const std::vector<int32_t> output_sharding = output_layout.num_shards();
std::vector<shape_inference::DimensionHandle> out_dims;
out_dims.reserve(c->Rank(in));
for (int i = 0; i < c->Rank(in); ++i) {
shape_inference::DimensionHandle dim = c->Dim(in, i);
if (!c->ValueKnown(dim) ||
input_layout.sharding_spec(i) == output_layout.sharding_spec(i)) {
out_dims.emplace_back(dim);
} else if (Layout::IsUnshardedDimension(
input_layout.sharding_spec(i))) {
shape_inference::DimensionHandle out_dim;
TF_RETURN_IF_ERROR(c->Divide(dim, output_sharding[i],
/*evenly_divisible=*/true, &out_dim));
out_dims.push_back(out_dim);
} else {
return absl::InvalidArgumentError(absl::StrCat(
"DTensorAllScatter only supports output layouts which are more "
"sharded than input layouts. Received input sharding spec ",
input_layout.sharding_spec(i), " and output sharding spec ",
output_layout.sharding_spec(i), " for dimension ", i, "."));
}
}
c->set_output(0, c->MakeShape(out_dims));
return absl::OkStatus();
});
REGISTER_OP("DTensorAllGather")
.Input("input: T")
.Output("output: T")
.Attr(
"T: {half, bfloat16, float, float64, int8, uint8, int32, uint32, "
"int64, uint64, "
"bool}")
.Attr("input_layout: string")
.Attr("output_layout: string")
.SetShapeFn([](shape_inference::InferenceContext* c) -> absl::Status {
shape_inference::ShapeHandle in = c->input(0);
if (!c->RankKnown(in)) {
// Input shape unknown, so set unknown output shape.
c->set_output(0, in);
return absl::OkStatus();
}
std::string input_layout_string;
std::string output_layout_string;
TF_RETURN_IF_ERROR(c->GetAttr("input_layout", &input_layout_string));
TF_RETURN_IF_ERROR(c->GetAttr("output_layout", &output_layout_string));
TF_ASSIGN_OR_RETURN(Layout input_layout,
Layout::FromString(input_layout_string));
TF_ASSIGN_OR_RETURN(Layout output_layout,
Layout::FromString(output_layout_string));
if (c->Rank(in) != input_layout.rank() ||
c->Rank(in) != output_layout.rank()) {
return absl::InvalidArgumentError(absl::StrCat(
"Input tensor rank and layout ranks do not agree: input rank ",
c->Rank(in), " input layout rank ", input_layout.rank(),
" output "
"layout rank ",
output_layout.rank()));
}
const std::vector<int32_t> input_sharding = input_layout.num_shards();
std::vector<shape_inference::DimensionHandle> out_dims;
out_dims.reserve(c->Rank(in));
for (int32_t i = 0; i < c->Rank(in); ++i) {
shape_inference::DimensionHandle dim = c->Dim(in, i);
if (!c->ValueKnown(dim) ||
input_layout.sharding_spec(i) == output_layout.sharding_spec(i)) {
out_dims.emplace_back(dim);
} else if (Layout::IsUnshardedDimension(
output_layout.sharding_spec(i))) {
shape_inference::DimensionHandle out_dim;
TF_RETURN_IF_ERROR(c->Multiply(dim, input_sharding[i], &out_dim));
out_dims.push_back(out_dim);
} else {
return absl::InvalidArgumentError(absl::StrCat(
"DTensorAllGatherr only supports input layouts which are more "
"sharded than output layouts. Received input sharding spec ",
input_layout.sharding_spec(i), " and output sharding spec ",
output_layout.sharding_spec(i), " for dimension ", i, "."));
}
}
c->set_output(0, c->MakeShape(out_dims));
return absl::OkStatus();
});
REGISTER_OP("DTensorAllToAll")
.Input("input: T")
.Output("output: T")
.Attr("T: {half, bfloat16, float, float64, int32, uint32, int64, bool}")
.Attr("input_layout: string")
.Attr("output_layout: string")
.SetShapeFn([](shape_inference::InferenceContext* c) -> absl::Status {
shape_inference::ShapeHandle in = c->input(0);
if (!c->RankKnown(in)) {
// Input shape unknown, so set unknown output shape.
c->set_output(0, in);
return absl::OkStatus();
}
std::string input_layout_string;
std::string output_layout_string;
TF_RETURN_IF_ERROR(c->GetAttr("input_layout", &input_layout_string));
TF_RETURN_IF_ERROR(c->GetAttr("output_layout", &output_layout_string));
TF_ASSIGN_OR_RETURN(Layout input_layout,
Layout::FromString(input_layout_string));
TF_ASSIGN_OR_RETURN(Layout output_layout,
Layout::FromString(output_layout_string));
if (c->Rank(in) != input_layout.rank() ||
c->Rank(in) != output_layout.rank()) {
return absl::InvalidArgumentError(absl::StrCat(
"Input tensor rank and layout ranks do not agree: input rank ",
c->Rank(in), " input layout rank ", input_layout.rank(),
" output "
"layout rank ",
output_layout.rank()));
}
const std::vector<int32_t> input_sharding = input_layout.num_shards();
const std::vector<int32_t> output_sharding = output_layout.num_shards();
std::vector<shape_inference::DimensionHandle> out_dims;
out_dims.reserve(c->Rank(in));
for (int i = 0; i < c->Rank(in); ++i) {
shape_inference::DimensionHandle dim = c->Dim(in, i);
if (!c->ValueKnown(dim) ||
input_layout.sharding_spec(i) == output_layout.sharding_spec(i)) {
out_dims.emplace_back(dim);
} else if (Layout::IsUnshardedDimension(
input_layout.sharding_spec(i)) &&
Layout::IsShardedDimension(output_layout.sharding_spec(i))) {
shape_inference::DimensionHandle out_dim;
TF_RETURN_IF_ERROR(c->Divide(dim, output_sharding[i],
/*evenly_divisible=*/true, &out_dim));
out_dims.push_back(out_dim);
} else if (Layout::IsShardedDimension(input_layout.sharding_spec(i)) &&
Layout::IsUnshardedDimension(
output_layout.sharding_spec(i))) {
shape_inference::DimensionHandle out_dim;
TF_RETURN_IF_ERROR(c->Multiply(dim, input_sharding[i], &out_dim));
out_dims.push_back(out_dim);
}
}
c->set_output(0, c->MakeShape(out_dims));
return absl::OkStatus();
});
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/cc/dtensor_operation.h
+59
View File
@@ -0,0 +1,59 @@
/* 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_DTENSOR_CC_DTENSOR_OPERATION_H_
#define TENSORFLOW_DTENSOR_CC_DTENSOR_OPERATION_H_
#include "tensorflow/c/eager/c_api.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
namespace tensorflow {
namespace dtensor {
// Captures the properties of an Operation currently being executed by DTensor.
struct DTensorOperation {
// For all by-ref fields: not owned. lifetime covers the whole usage.
const char* name;
const FunctionDef* function_def;
// Default mesh is used when Mesh Propagation does not identify a mesh
// otherwise.
const Mesh default_mesh;
const StackTracesMap& stack_traces;
bool is_func() const { return function_def != nullptr; }
// Returns True if the op has no side effects.
// Side effects include global side effect marked by IsStateful and
// Input or output of Resources.
// This definition is correct for all DTensor support Ops.
// Some odder TF Ops (e.g. Queue) do not mark themselve as stateful, but are
// still stateful. DTensor doesn't support them.
bool is_pure() const {
if (is_func()) {
// FIXME(feyu): some functions can still be pure, but we just don't yet
// handle the case, and treat all functions as non-pure.
return false;
}
const OpDef* op_def = nullptr;
absl::Status status = OpRegistry::Global()->LookUpOpDef(name, &op_def);
DCHECK_OK(status); // Not found. This really shouldn't happen.
if (!status.ok()) {
return false;
}
return !op_def->is_stateful();
}
};
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_CC_DTENSOR_OPERATION_H_
tensorflow/dtensor/cc/dtensor_ops.cc
+108
View File
@@ -0,0 +1,108 @@
/* Copyright 2022 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 <vector>
#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/framework/tensor_slice.h"
#include "tensorflow/core/util/saved_tensor_slice_util.h"
namespace tensorflow {
namespace dtensor {
using shape_inference::InferenceContext;
using shape_inference::ShapeHandle;
using shape_inference::UnchangedShape;
// Change layout of input to target layout inside the same mesh cluster.
REGISTER_OP("Relayout")
.Input("input: T")
.Output("output: T")
.Attr("layout: string")
.Attr("T: type")
.SetShapeFn(UnchangedShape);
// Relayout the input according to the layout of layout_input.
REGISTER_OP("RelayoutLike")
.Input("input: T")
.Input("layout_input: U") // To infer the output mesh.
.Output("output: T")
.Attr("T: type")
.Attr("U: type")
.SetShapeFn(UnchangedShape);
// FIXME(b/271292250): Add DTensor suffix to signal this is a meta Op
// Op. Or remove this altogether, if there is no use for it.
// Copy `input` to the given mesh and layout.
REGISTER_OP("CopyToMesh")
.Input("input: T")
.Output("output: T")
.Attr("mesh: string")
.Attr("T: type")
.SetShapeFn(UnchangedShape);
// FIXME(b/271292250): Remove this Op It is no longer used.
// Gradient of CopyToMesh.
REGISTER_OP("CopyToMeshGrad")
.Input("input: T")
.Input("forward_input: T") // To infer the output mesh.
.Output("output: T")
.Attr("T: type")
.SetShapeFn(UnchangedShape);
// DTensorRestoreV2 that is pretty much RestoreV2 but with extra global shapes
// and layouts.
REGISTER_OP("DTensorRestoreV2")
.Input("prefix: string")
.Input("tensor_names: string")
.Input("shape_and_slices: string")
.Output("tensors: dtypes")
.Attr("input_shapes: list(shape)")
.Attr("input_layouts: list(string)")
.Attr("dtypes: list(type)")
.SetIsStateful()
.SetShapeFn([](InferenceContext* c) {
ShapeHandle shape0, shape1, shape2;
TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &shape0));
TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 1, &shape1));
TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 1, &shape2));
TF_RETURN_IF_ERROR(c->Merge(shape1, shape2, &shape0));
std::vector<PartialTensorShape> input_shapes;
TF_RETURN_IF_ERROR(c->GetAttr("input_shapes", &input_shapes));
std::vector<std::string> input_layouts;
TF_RETURN_IF_ERROR(c->GetAttr("input_layouts", &input_layouts));
if (input_shapes.size() != input_layouts.size()) {
return absl::InvalidArgumentError(absl::StrCat(
"Size of input_shapes and input_layouts is expected to match, but "
"got ",
input_shapes.size(), " for input_shapes and ", input_layouts.size(),
" for input_layouts"));
}
// TODO(hthu): We should be able to infer from layout and global_shape
// field.
return UnknownShape(c);
});
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/cc/dtensor_tpu_kernels.cc
+175
View File
@@ -0,0 +1,175 @@
/* Copyright 2022 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 <cstddef>
#include <cstdint>
#include <vector>
#include "absl/log/log.h"
#include "absl/log/vlog_is_on.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/time/time.h"
#include "xla/tpu/status_helper.h"
#include "xla/tpu/tpu_api.h"
#include "xla/tpu/tpu_ops_c_api.h"
#include "tensorflow/core/framework/collective.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/op_requires.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/platform/tstring.h"
#include "tensorflow/core/tpu/tpu_configuration.h"
#include "tensorflow/dtensor/cc/tpu_system_interface.h"
// Timeout for waiting for TPU devices to appear.
const absl::Duration dtensor_tpu_init_retry_timeout = absl::Seconds(30);
namespace tensorflow {
namespace dtensor {
// Attempt to delete resource_name from resource_manager's default_container.
// Returns OK if the deletion succeeded, or if the resource was not found. Else
// return the deletion error.
template <class ResourceT>
absl::Status DeleteIfExists(ResourceMgr* resource_manager,
const char* resource_name) {
VLOG(1) << "Removing resource " << resource_name << " if it exists";
absl::Status status = resource_manager->Delete<ResourceT>(
resource_manager->default_container(), resource_name);
if (status.ok()) {
VLOG(1) << "Removed existing resource " << resource_name;
return absl::OkStatus();
}
if (status.code() == error::NOT_FOUND) {
VLOG(1) << "No resource " << resource_name << " to remove";
return absl::OkStatus();
}
VLOG(1) << "Error removing resource " << resource_name << " : " << status;
return status;
}
class ConfigureAndInitializeGlobalTPUOpKernel : public OpKernel {
public:
explicit ConfigureAndInitializeGlobalTPUOpKernel(OpKernelConstruction* ctx)
: OpKernel(ctx) {}
void Compute(OpKernelContext* ctx) override {
LOG(INFO) << "ConfigureAndInitializeGlobalTPUOpKernel op";
ResourceMgr* rmgr = GetTPUConfigResourceMgr();
std::vector<int32_t> core_id_output_vec;
auto retry_timeout = dtensor_tpu_init_retry_timeout;
VLOG(1) << "Initializing the TPU system.";
TpuSystemInterface* tpu_system = GetPreferredTpuSystem();
OP_REQUIRES(ctx, tpu_system != nullptr,
absl::FailedPreconditionError("TPU system not initialized"));
OP_REQUIRES_OK(ctx, tpu_system->Initialize(ctx, rmgr, retry_timeout,
&core_id_output_vec));
if (VLOG_IS_ON(1)) {
LOG(INFO) << "core_id_output_vec";
for (auto i : core_id_output_vec) {
LOG(INFO) << i;
}
}
// Set output using local core ID vector.
Tensor* ctx_output;
auto core_id_output_vec_size = core_id_output_vec.size();
OP_REQUIRES_OK(
ctx,
ctx->allocate_output(
0, TensorShape({static_cast<long long>(core_id_output_vec_size)}),
&ctx_output));
for (size_t i = 0; i < core_id_output_vec_size; ++i) {
ctx_output->flat<int32_t>()(i) = core_id_output_vec[i];
}
LOG(INFO) << "ConfigureAndInitializeGlobalTPUOpKernel done";
}
~ConfigureAndInitializeGlobalTPUOpKernel() override = default;
private:
// ConfigureAndInitializeGlobalTPUOpKernel is neither copyable nor movable.
ConfigureAndInitializeGlobalTPUOpKernel(
const ConfigureAndInitializeGlobalTPUOpKernel&) = delete;
ConfigureAndInitializeGlobalTPUOpKernel& operator=(
const ConfigureAndInitializeGlobalTPUOpKernel&) = delete;
};
class ShutdownTPUSystemOpKernel : public OpKernel {
public:
explicit ShutdownTPUSystemOpKernel(OpKernelConstruction* ctx)
: OpKernel(ctx) {}
void Compute(OpKernelContext* ctx) override {
LOG(INFO) << "ShutdownTPUSystemOpKernel op";
absl::Status status;
TpuSystemInterface* tpu_system = GetPreferredTpuSystem();
OP_REQUIRES(ctx, tpu_system != nullptr,
absl::FailedPreconditionError("TPU system not initialized."));
VLOG(1) << "Shutting down the TPU system.";
status = tpu_system->Shutdown();
Tensor* output_tensor;
OP_REQUIRES_OK(ctx,
ctx->allocate_output(0, TensorShape({1}), &output_tensor));
if (status.ok()) {
output_tensor->flat<bool>()(0) = true;
} else {
output_tensor->flat<bool>()(0) = false;
}
}
};
class SetGlobalTPUArrayOpKernel : public OpKernel {
public:
explicit SetGlobalTPUArrayOpKernel(OpKernelConstruction* ctx)
: OpKernel(ctx) {}
void Compute(OpKernelContext* ctx) override {
VLOG(1) << "SetGlobalTPUArrayOpKernel op";
OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(ctx->input(0).shape()),
absl::InvalidArgumentError(
absl::StrCat("Expected argument 0 to be a scalar. Received",
ctx->input(0).DebugString())));
auto tpu_topology = ctx->input(0).scalar<tstring>()();
StatusHelper status;
stream_executor::tpu::OpsApiFn()->SetGlobalTPUArrayOp_DoWorkFn(
tpu_topology.size(), tpu_topology.data(), status.c_status);
OP_REQUIRES_OK(ctx, status.status());
VLOG(1) << "SetGlobalTPUArrayOpKernel done";
}
};
REGISTER_KERNEL_BUILDER(Name("ConfigureAndInitializeGlobalTPU")
.Device(DEVICE_TPU_SYSTEM)
.HostMemory("output"),
ConfigureAndInitializeGlobalTPUOpKernel);
REGISTER_KERNEL_BUILDER(Name("ShutdownTPUSystem").Device(DEVICE_TPU_SYSTEM),
ShutdownTPUSystemOpKernel);
REGISTER_KERNEL_BUILDER(Name("DTensorSetGlobalTPUArray")
.Device(DEVICE_TPU_SYSTEM)
.HostMemory("topology"),
SetGlobalTPUArrayOpKernel);
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/cc/dtensor_tpu_ops.cc
+65
View File
@@ -0,0 +1,65 @@
/* Copyright 2022 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 <vector>
#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
namespace tensorflow {
namespace dtensor {
using shape_inference::InferenceContext;
using shape_inference::ShapeHandle;
// Initializes global TPU's for mutli-client execution.
//
// This op does the work of both ConfigureDistributedTpuOp and
// InitializeHostForDistributedTpuOp, and outputs the latter's result.
REGISTER_OP("ConfigureAndInitializeGlobalTPU")
.Output("output: int32")
.Attr("use_tfrt_host_runtime: bool = true")
.SetIsStateful()
.SetShapeFn([](InferenceContext* c) {
ShapeHandle input;
// Validate that all the inputs are scalars.
for (int i = 0; i < c->num_inputs(); ++i) {
TF_RETURN_IF_ERROR(c->WithRank(c->input(i), 0, &input));
}
c->set_output(0, c->Vector(c->UnknownDim()));
return absl::OkStatus();
});
REGISTER_OP("ShutdownTPUSystem")
.SetIsStateful()
.Output("success: bool")
.SetShapeFn(shape_inference::ScalarShape);
REGISTER_OP("DTensorSetGlobalTPUArray")
.Input("topology: string")
.SetIsStateful()
.SetShapeFn([](InferenceContext* c) {
ShapeHandle input;
TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &input));
return absl::OkStatus();
});
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/cc/dtensor_utils.cc
+187
View File
@@ -0,0 +1,187 @@
/* Copyright 2022 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/dtensor/cc/dtensor_utils.h"
#include <algorithm>
#include <cstdlib>
#include <string>
#include <vector>
#include "absl/strings/ascii.h"
#include "absl/strings/numbers.h"
#include "absl/strings/str_split.h"
#include "xla/tsl/platform/status.h"
#include "xla/tsl/util/env_var.h"
#include "tensorflow/core/platform/logging.h"
namespace tensorflow {
namespace dtensor {
// LINT.IfChange
int ClientId() {
char* client_id_str = std::getenv("DTENSOR_CLIENT_ID");
if (client_id_str == nullptr) return 0;
int client_id;
if (absl::SimpleAtoi(client_id_str, &client_id)) return client_id;
LOG(WARNING) << "Invalid DTENSOR_CLIENT_ID, using the default value 0.";
return 0;
}
// LINT.ThenChange(//tensorflow/dtensor/python/dtensor_device.py)
// LINT.IfChange
int NumClients() {
char* num_clients_str = std::getenv("DTENSOR_NUM_CLIENTS");
if (num_clients_str == nullptr) return 1;
int num_clients;
if (absl::SimpleAtoi(num_clients_str, &num_clients)) return num_clients;
LOG(WARNING) << "Invalid DTENSOR_NUM_CLIENTS, using the default value 1.";
return 1;
}
// LINT.ThenChange(//tensorflow/dtensor/python/dtensor_device.py)
bool LogOnAllTasks() {
char* dtensor_log_on_all_tasks_str = std::getenv("DTENSOR_LOG_ON_ALL_TASKS");
if (dtensor_log_on_all_tasks_str == nullptr) return false;
return true;
}
bool LogOpByOp(absl::string_view op_name) {
char* op_list_str = std::getenv("DTENSOR_LOG_OP_BY_OP");
if (op_list_str == nullptr) return false;
if (!strcmp(op_list_str, "*")) return true;
std::vector<absl::string_view> op_list = absl::StrSplit(op_list_str, ',');
if (std::find(op_list.begin(), op_list.end(), op_name) != op_list.end()) {
return true;
}
return false;
}
int LayoutPropagationMaxSteps() {
char* dtensor_layout_propagation_max_steps_str =
std::getenv("DTENSOR_LAYOUT_PROPAGATION_MAX_STEPS");
if (dtensor_layout_propagation_max_steps_str == nullptr) return 500;
int dtensor_layout_propagation_max_steps;
if (absl::SimpleAtoi(dtensor_layout_propagation_max_steps_str,
&dtensor_layout_propagation_max_steps))
return dtensor_layout_propagation_max_steps;
LOG(WARNING) << "Invalid DTENSOR_LAYOUT_PROPAGATION_MAX_STEPS, using "
"the default value 500.";
return 500;
}
bool EnableMixedPrecisionReduce() {
char* dtensor_enable_mixed_precision_reduce_str =
std::getenv("DTENSOR_ENABLE_MIXED_PRECISION_REDUCE");
if (dtensor_enable_mixed_precision_reduce_str == nullptr) return false;
return true;
}
bool DoNotFuseReduceScatter() {
char* dtensor_do_not_fuse_reduce_scatter_str =
std::getenv("DTENSOR_DO_NOT_FUSE_REDUCE_SCATTER");
if (dtensor_do_not_fuse_reduce_scatter_str == nullptr) return false;
return true;
}
int ReduceInBfloat16MaxGroupSize() {
char* dtensor_reduce_in_bfloat16_max_group_size_str =
std::getenv("DTENSOR_REDUCE_IN_BFLOAT16_MAX_GROUP_SIZE");
if (dtensor_reduce_in_bfloat16_max_group_size_str == nullptr) return 8;
int dtensor_reduce_in_bfloat16_max_group_size;
if (absl::SimpleAtoi(dtensor_reduce_in_bfloat16_max_group_size_str,
&dtensor_reduce_in_bfloat16_max_group_size))
return dtensor_reduce_in_bfloat16_max_group_size;
LOG(WARNING) << "Invalid DTENSOR_REDUCE_IN_BFLOAT16_MAX_GROUP_SIZE, using "
"the default value 8.";
return 8;
}
bool LowerCollectiveGatherToCollectiveGatherV2() {
// We lower DTensorGather to CollectiveReduceV2 ops instead of
// CollectiveGatherV2, since we do not observe a performance gain with Gather
// lowering and ReduceV2 is agnostic of the rank order.
//
// If LOWER_DTENSOR_GATHER_TO_COLLECTIVE_GATHER_V2 environment variable is set
// to '1', it is reduced to collective
char* use_collective_gather =
std::getenv("LOWER_DTENSOR_GATHER_TO_COLLECTIVE_GATHER_V2");
if (use_collective_gather == nullptr) return false;
return true;
}
bool EnableReplicatedSpmdAsDefault(const std::string& op_name) {
// These environment variables enroll MLIR ops of the given name for default
// replicated SPMD expansion. No expanders are registered for these Ops,
// and without enrolling to the default replicated behavior, SPMD expansion
// raises an error for these Op.
//
// For example, to enroll tf.Mod, set
// DTENSOR_ENABLE_REPLICATED_SPMD_AS_DEFAULT_TF.MOD = 1
std::string env_name = "DTENSOR_ENABLE_REPLICATED_SPMD_AS_DEFAULT_" +
absl::AsciiStrToUpper(op_name);
char* dtensor_enable_replicated_spmd_as_default =
std::getenv(env_name.c_str());
return dtensor_enable_replicated_spmd_as_default != nullptr;
}
bool EnableAllToAllForRelayout() {
// Whether to use all-to-all collective for relayout when possible.
static bool is_enabled = [] {
bool ret = true;
TF_CHECK_OK(tsl::ReadBoolFromEnvVar("DTENSOR_USE_ALL_TO_ALL_RELAYOUT",
/*default_val=*/true, &ret));
return ret;
}();
return is_enabled;
}
int AllReduceCombineOptimizationGroupSize() {
char* group_size_str =
std::getenv("DTENSOR_ALLREDUCE_COMBINE_OPTIMIZATION_GROUP_SIZE");
if (group_size_str == nullptr) return 0;
int group_size;
if (absl::SimpleAtoi(group_size_str, &group_size)) return group_size;
LOG(WARNING) << "Invalid DTENSOR_ALLREDUCE_COMBINE_OPTIMIZATION_GROUP_SIZE, "
"using the default value 0.";
return 0;
}
int AllReduceCombineOptimizationTopologicalDistance() {
int64_t topo_dist;
absl::Status status = tsl::ReadInt64FromEnvVar(
"DTENSOR_ALLREDUCE_COMBINE_OPTIMIZATION_TOPOLOGICAL_DISTANCE",
/*default_val=*/0, &topo_dist);
if (!status.ok()) {
LOG(WARNING) << "Invalid DTENSOR_ALLREDUCE_COMBINE_OPTIMIZATION_TOPOLOGICAL"
"_DISTANCE, using the default value 0.";
return 0;
} else if (topo_dist < 0) {
LOG(WARNING) << "Invalid DTENSOR_ALLREDUCE_COMBINE_OPTIMIZATION_TOPOLOGICAL"
"_DISTANCE, value must be a positive integer, using the "
"default value 0.";
return 0;
}
return topo_dist;
}
bool EnableMultiDeviceMode() {
bool multi_device_mode;
absl::Status status = tsl::ReadBoolFromEnvVar(
"DTENSOR_ENABLE_MULTI_DEVICE_EXPANSION", false, &multi_device_mode);
return status.ok() && multi_device_mode;
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/cc/dtensor_utils.h
+92
View File
@@ -0,0 +1,92 @@
/* Copyright 2022 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_DTENSOR_CC_DTENSOR_UTILS_H_
#define TENSORFLOW_DTENSOR_CC_DTENSOR_UTILS_H_
#include <string>
#include "absl/strings/string_view.h"
namespace tensorflow {
namespace dtensor {
// Returns the DTensor client ID of this process, usually equal to the TF task
// ID on this host.
int ClientId();
// Returns the total number of DTensor clients, usually equal to the total
// number of TF tasks.
int NumClients();
// Returns whether to enable logging for passes and layouts on all passes.
bool LogOnAllTasks();
// Returns whether to log op-by-op execution in addition to function execution
// when logging is enabled.
bool LogOpByOp(absl::string_view op_name);
// Returns the maximum number of steps to run layout propagation. If the number
// of steps exceeds this amount, layout propagation will fail.
int LayoutPropagationMaxSteps();
// Returns whether to upcast bfloat16 reduction inputs to float32 for
// sufficient reduction group size.
bool EnableMixedPrecisionReduce();
// Returns whether *not* to fuse AllReduce + AllScatter into ReduceScatter op,
// which can be more efficiently implemented.
bool DoNotFuseReduceScatter();
// Returns the maximum reduction group size for bfloat16 reduction. If the
// group size exceeds this, then tensors are upcasted to float32 before the
// reduce op.
int ReduceInBfloat16MaxGroupSize();
// Returns whether to lower DTensorAllGather to CollectiveReduceV2. If false,
// lowers it to CollectiveReduceV2 for GPU and CPU for supported data types.
bool LowerCollectiveGatherToCollectiveGatherV2();
// Returns whether to enable defaulting TF ops that do not have SPMD
// implementation to default to the ReplicatedOpSpmdExpander.
bool EnableReplicatedSpmdAsDefault(const std::string& op_name);
// Returns whether to use all-to-all collective for relayout when possible.
bool EnableAllToAllForRelayout();
// Returns the maximum number of AllReduce ops to merge into a group. This value
// determines the AllReduce grouping in dtensor_allreduce_combine_optimization.
// The input value should be in range of [0, INT_MAX]. It is advised to pick
// a value based on knowledge of the total number of AllReduces. When the value
// is too big, the behaviour will act as aggressive grouping. When the value is
// too small, the behaviour will act as having no extended grouping.
int AllReduceCombineOptimizationGroupSize();
// Returns the maximum topological distance between two AllReduce ops to merge
// into a single AllReduce. This value is used to determine AllReduce grouping
// in dtensor_allreduce_combine_optimization. The input value should be in range
// of [0, INT_MAX]. However, it is advised to select a value based on knowledge
// of the compute graph, such as the minimum distance between two model layers.
// When the input value is too big, the behaviour will act as aggressive group-
// ing. When the input value is too small, the behaviour will act as having no
// extended grouping.
int AllReduceCombineOptimizationTopologicalDistance();
// Returns whether to perform multi-device expansion.
bool EnableMultiDeviceMode();
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_CC_DTENSOR_UTILS_H_
tensorflow/dtensor/cc/mesh_type.h
+33
View File
@@ -0,0 +1,33 @@
/* Copyright 2022 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_DTENSOR_CC_MESH_TYPE_H_
#define TENSORFLOW_DTENSOR_CC_MESH_TYPE_H_
#include "tensorflow/c/conversion_macros.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
namespace tensorflow {
extern "C" {
typedef struct TF_Mesh TF_Mesh;
}
DEFINE_CONVERSION_FUNCTIONS(dtensor::Mesh, TF_Mesh);
typedef struct TF_Layout TF_Layout;
DEFINE_CONVERSION_FUNCTIONS(dtensor::Layout, TF_Layout);
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_CC_MESH_TYPE_H_
tensorflow/dtensor/cc/parallel_executor.h
+80
View File
@@ -0,0 +1,80 @@
/* Copyright 2022 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_DTENSOR_CC_PARALLEL_EXECUTOR_H_
#define TENSORFLOW_DTENSOR_CC_PARALLEL_EXECUTOR_H_
#include <memory>
#include <optional>
#include <vector>
#include "llvm/ADT/StringRef.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "tensorflow/c/eager/c_api_experimental.h"
#include "xla/future.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/cc/tensor_with_layout.h"
namespace tensorflow {
namespace dtensor {
template <typename T = void>
using Future = ::xla::Future<T>;
// ParallelExecutor Interface
// Note: The interface is under development and APIs are subject to change.
class ParallelExecutor {
public:
virtual ~ParallelExecutor() = default;
// Broadcasts `tensor` to `mesh` using replicated sharding and returns a
// DTensor representation. `mesh` can be a single device mesh and in that case
// `const_value` is useless.
virtual StatusOr<std::unique_ptr<TensorWithLayout>> Broadcast(
const Tensor& tensor, const Mesh& mesh,
std::optional<NodeDef> const_value) = 0;
// Takes input TensorWithLayouts and a MLIR module.
// The MLIR module should have `main` as its entry function name.
// Attributes are forwarded to executed operations unmodified.
// The execute is non-blocking and returns a Future of output TensorWithLayout
// raw pointers.
// The client is responsible for the ownership of the outputs.
struct ExecutionResult {
Future<> status;
// The pointed data of `outputs` are filled after `status` future resolves
// as ok.
std::vector<TensorWithLayout*> outputs;
};
virtual StatusOr<ExecutionResult> Execute(
TFE_Context* context, const std::vector<TensorWithLayout*>& inputs,
mlir::ModuleOp module, const TFE_OpAttrs* attributes) const = 0;
// Disassembles `t` into multiple TensorWithLayouts. `t` may or may not be
// valid to use afterwards.
virtual StatusOr<std::vector<std::unique_ptr<TensorWithLayout>>> Disassemble(
TensorWithLayout* t) = 0;
// Returns a tensor copied from `t` when `t` contains only a single device.
virtual Future<Tensor> ToHostBuffer(TensorWithLayout* t) = 0;
};
// Factory method for Default ParallelExecutor instance.
StatusOr<std::unique_ptr<ParallelExecutor>> CreateDefaultParallelExecutor();
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_CC_PARALLEL_EXECUTOR_H_
tensorflow/dtensor/cc/save_restore_util.cc
+201
View File
@@ -0,0 +1,201 @@
/* Copyright 2022 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/dtensor/cc/save_restore_util.h"
#include "llvm/ADT/SmallVector.h"
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/dtensor/mlir/value_utils.h"
namespace tensorflow {
namespace dtensor {
namespace {
// A map that is keyed by the index of tensor_name.
// For example, {2 : <"spec_a", "spec_b"> } means that the
// save_v2.tensor_names[2] should have "spec_a" and "spec_b" saved.
using SliceSpecByName = absl::flat_hash_map<int64_t, std::vector<std::string>>;
// Builds a map from tensor slice spec to saving device_id for the given Tensor
// and layout. The output would record the saving device and the slices it needs
// to save.
//
// For each sharded Tensor, each device would hold a slice of the Tensor - but
// it isn't necessary a unique copy. For a 2 way sharded Tensor in a (2,4) mesh
// on the first dimension, device [0-3] and device [4-7] will hold the same
// slice data. To avoid saving duplicated copies of the Tensor slice, the map
// would only contain the min(device_id) that occupies the slice and save from
// there.
//
// Furthermore, to save a Tensor that isn't on CPU mesh, send/recv is necessary
// from saving device to its corresponding host(CPU) devices. Since we don't
// have multi-mesh execution yet, this isn't implemented yet.
StatusOr<SliceSpecByName> BuildSliceSpecDeviceMap(
absl::Span<const int64_t> global_shape, Layout layout) {
if (!layout.mesh().is_cpu_mesh())
return absl::UnimplementedError(
"Saving tensors on non CPU mesh needs explicit send/receive and isn't "
"implemented yet");
// Result map that records the minimum device_id that occupies the unique
// copy.
// Note that llvm::SmallDenseMap won't accept std::string as a key.
absl::flat_hash_map<std::string, int64_t> min_device_for_slice_spec;
// Records the map of device_ids and a list of slice_spec that it needs to
// save.
SliceSpecByName device_slices;
const auto& mesh = layout.mesh();
// Construct SliceSpec for each device in the mesh.
for (int device_id = 0; device_id < mesh.size(); ++device_id) {
TF_ASSIGN_OR_RETURN(const DeviceLocation& coords,
mesh.device_location(device_id));
// Prefill with full spec on each dim.
TF_ASSIGN_OR_RETURN(std::vector<std::string> slice_specs,
SliceSpecOnDevice(layout, mesh, coords, global_shape));
// Build the real slice_spec from string pieces.
std::string slice_spec = absl::StrJoin(slice_specs, ":");
// Get local shape from the global shape.
std::string shape_spec = absl::StrJoin(global_shape, " ");
// Concat shape spec and slice spec to form a complete shape_and_slice.
std::string shape_and_slice = absl::StrCat(shape_spec, " ", slice_spec);
// Only record the min device_id for the unique slice_spec on a given
// Tensor.
if (min_device_for_slice_spec.find(shape_and_slice) ==
min_device_for_slice_spec.end() ||
device_id < min_device_for_slice_spec[shape_and_slice]) {
min_device_for_slice_spec[shape_and_slice] = device_id;
}
}
// Constructs device_id keyed map for future save operation conditioned on
// device_ids.
for (const auto& spec_and_id : min_device_for_slice_spec) {
device_slices[spec_and_id.second].push_back(spec_and_id.first);
}
return device_slices;
}
} // namespace
// Example is _dev-02-of-16.
std::string DeviceSuffix(int device_id, int total_devices) {
return absl::StrFormat("_dev-%0*d-of-%d", absl::StrCat(total_devices).size(),
device_id, total_devices);
}
StatusOr<absl::flat_hash_map<
int64_t, absl::flat_hash_map<int64_t, std::vector<std::string>>>>
BuildSavingSpec(absl::Span<const SavingTensorMetadata> tensor_metadatas) {
absl::flat_hash_map<int64_t,
absl::flat_hash_map<int64_t, std::vector<std::string>>>
saving_specs;
for (const SavingTensorMetadata& tensor_metadata : tensor_metadatas) {
// We use index to select the tensor names and shape_and_slices from the
// inputs. This is generic regardless whether the inputs are constants or
// just arguments.
int index = tensor_metadata.tensor_index;
const Layout& layout = tensor_metadata.layout;
absl::Span<const int64_t> tensor_shape = tensor_metadata.shape;
if (layout.IsFullyReplicated()) {
// Push a fully replicated save on device 0, where slice_spec is simply
// empty string.
saving_specs[0][index].push_back("");
} else {
// Calculate shape_and_slices for sharded case here.
TF_ASSIGN_OR_RETURN(const auto& slice_specs,
BuildSliceSpecDeviceMap(tensor_shape, layout));
// Push specs for each device into the global map.
for (const auto& slice_spec : slice_specs) {
int64_t saving_device_id = slice_spec.first;
for (const std::string& slice : slice_spec.second) {
saving_specs[saving_device_id][index].push_back(slice);
}
}
}
}
return saving_specs;
}
SaveOpSpecs BuildPerDeviceSave(
mlir::OpBuilder& builder,
const absl::flat_hash_map<int64_t, std::vector<std::string>>& saving_spec,
int device_id, mlir::Value prefix, int total_devices) {
std::vector<mlir::Value> new_prefixes;
std::vector<std::vector<int>> tensor_indices;
std::vector<std::vector<std::string>> shape_and_slice_specs;
for (const auto& tensor_name_index_and_slice_specs : saving_spec) {
int tensor_index = tensor_name_index_and_slice_specs.first;
const std::vector<std::string> specs =
tensor_name_index_and_slice_specs.second;
// For each tensor_name, we save its first slice_spec in the first
// save_op, second slice_spec in the second save op, etc.
// This allows us to group save ops together without running into
// duplicated tensor_names (which save_v2 op doesn't support).
for (int save_op_index = 0; save_op_index < specs.size(); ++save_op_index) {
if (save_op_index >= tensor_indices.size()) {
tensor_indices.push_back({});
shape_and_slice_specs.push_back({});
mlir::Value new_prefix =
mlir::TF::AddOp::create(
builder, prefix.getLoc(),
mlir::dyn_cast<mlir::RankedTensorType>(prefix.getType()),
prefix,
StringScalarConst(builder, prefix.getLoc(),
DeviceSuffix(device_id, total_devices)))
.getZ();
// Generate new prefix based on device_id and save op index, only when
// we need a new save_op.
new_prefixes.push_back(new_prefix);
}
tensor_indices[save_op_index].push_back(tensor_index);
shape_and_slice_specs[save_op_index].push_back(specs[save_op_index]);
}
}
return SaveOpSpecs(new_prefixes, tensor_indices, shape_and_slice_specs);
}
StatusOr<std::vector<std::string>> SliceSpecOnDevice(
const Layout& layout, const Mesh& mesh, const DeviceLocation& device_coords,
absl::Span<const int64_t> global_shape) {
// Prefill the slice with replicated layouts.
std::vector<std::string> slice_specs(global_shape.size(), "-");
const std::vector<std::string>& sharding_spec_strs =
layout.sharding_spec_strs();
for (int tensor_dim_index = 0; tensor_dim_index < sharding_spec_strs.size();
++tensor_dim_index) {
const std::string& mesh_dim = sharding_spec_strs[tensor_dim_index];
if (layout.IsShardedDimension(mesh_dim)) {
TF_ASSIGN_OR_RETURN(int mesh_dim_index, mesh.idx_for_dim(mesh_dim));
TF_ASSIGN_OR_RETURN(int64_t dim_size, mesh.dim_size(mesh_dim));
int64_t per_slice_size = global_shape[tensor_dim_index] / dim_size;
int start = device_coords[mesh_dim_index] * per_slice_size;
slice_specs[tensor_dim_index] = absl::StrCat(start, ",", per_slice_size);
}
}
return slice_specs;
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/cc/save_restore_util.h
+156
View File
@@ -0,0 +1,156 @@
/* Copyright 2022 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_DTENSOR_CC_SAVE_RESTORE_UTIL_H_
#define TENSORFLOW_DTENSOR_CC_SAVE_RESTORE_UTIL_H_
#include <string>
#include <utility>
#include "absl/container/flat_hash_map.h"
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
namespace tensorflow {
namespace dtensor {
// Defines an Metadata entry when saving a Tensor.
struct SavingTensorMetadata {
// Tracks index from the original save op.
int64_t tensor_index;
// The global shape of the saving tensor.
std::vector<int64_t> shape;
// The layout of the saving tensor.
Layout layout;
SavingTensorMetadata(int64_t index, std::vector<int64_t> global_shape,
Layout tensor_layout)
: tensor_index(index),
shape(std::move(global_shape)),
layout(std::move(tensor_layout)) {}
};
// Tracks a complete specification for a particular save op.
// The users would build out multiple save ops using the following manner for
// the given fields:
//
// save_op[i] = tf.SaveV2(
// prefix = new_prefixes[i],
// tensor_indices = tensor_indies[i],
// shape_and_slices = shape_and_slice_spec[i])
struct SaveOpSpecs {
std::vector<mlir::Value> new_prefixes;
std::vector<std::vector<int>> tensor_indices;
std::vector<std::vector<std::string>> shape_and_slice_spec;
SaveOpSpecs(std::vector<mlir::Value> prefixes,
std::vector<std::vector<int>> indices,
std::vector<std::vector<std::string>> specs)
: new_prefixes(std::move(prefixes)),
tensor_indices(std::move(indices)),
shape_and_slice_spec(std::move(specs)) {}
};
// Returns a device suffix with printf formatting.
std::string DeviceSuffix(int device_id, int total_devices);
// Builds a complete saving specification for each device on the mesh.
//
// The returned map contains a map of <device_id, SavingSpec>.
// Device_id is where the saving should happen, and SavingSpec is a
// mapping of <tensor_index -> shape_and_slices>. e.g.,
//
// A map of {device_id : 0 -> {
// 0 : "2 0,1",
// 1 : ""
// }
// }
//
// Means that device_0 is responsible for saving tensor 0 and 1 from the passed
// in tensors list. For tensor[0], it saves the only the first element in that
// 1d vector with 2 elements. For tensor[1], it saves all elements.
//
// We accept another map as input, that records the mapping of
// <tensor_index -> (tensor_global_shape, tensor_layout)>.
//
// (tensor_global_shape, tensor_layout & tensor_layout.mesh) defines which
// device saves what slices of the Tensor.
//
// For a complete definition of shape_and_slices field, please see:
// third_party/tensorflow/core/framework/tensor_slice.h
StatusOr<absl::flat_hash_map<
int64_t, absl::flat_hash_map<int64_t, std::vector<std::string>>>>
BuildSavingSpec(absl::Span<const SavingTensorMetadata> tensor_metadatas);
// For a given per device saving spec, find out the counts of SaveV2 ops
// needed and their corresponding inputs.
//
// Current SaveV2 op requires tensor_names to be unique in the list, which is a
// contract that distributed saving would break. For example, if the saving spec
// decides that device 0 is responsible for saving two slices of tensor[a], then
// a single SaveV2 op can't fufill. The setup is very likely to happen when
// saving on TPU - where 8 cores maps to 1 host. In that case, the CPU host will
// be responsible for saving slices on the same tensor across 8 TPU cores.
// TODO(b/179126981): Investigate whether we can make TF core API run with
// different slice spec on a same tensor key.
//
// That said, building one SaveV2 op for each save is wasteful, when a single
// SaveV2 op is capable of saving different tensors. Instead, we simply need to
// break the SaveV2 op to be able to track the longest saving specs for a single
// tensor happening on the device, e.g.,
//
// For given saving specs:
//
// { 'tensor_name_a' : <"spec_a", "spec_a_2"> }
// { 'tensor_name_b' : <"spec_b"> }
//
// would result into two save ops, where:
//
// SaveOp1 (tensor_names = <"tensor_name_a", tensor_name_b">,
// slice_spec = <"spec_a", "spec_b">)
//
// SaveOp2 (tensor_names = "<tensor_name_a>", slice_spec = <"spec_a_2">.
//
// The output vectors tracks the new SaveV2 op parameters and they must agree on
// size and indexing for saving tensors.
//
// tensor_indices trackes a list of indices of tensors that are being saved for
// each Save op, e.g.,
//
// tensor_indices[0] is a list of tensors (in index form) that needs to be saved
// on the first SaveV2 op.
//
// shape_and_slice_specs tracks a list of shape_and_slice_specs being saved for
// each Save op, e.g.,
//
// shape_and_slice_spec[0] is a list of shape_and_slices parameters for SaveV2
// op.
SaveOpSpecs BuildPerDeviceSave(
mlir::OpBuilder& builder,
const absl::flat_hash_map<int64_t, std::vector<std::string>>& saving_spec,
int device_id, mlir::Value prefix, int total_devices);
// Figures out the tensor slice_spec for a given layout and mesh device
// location.
StatusOr<std::vector<std::string>> SliceSpecOnDevice(
const Layout& layout, const Mesh& mesh, const DeviceLocation& device_coords,
absl::Span<const int64_t> global_shape);
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_CC_SAVE_RESTORE_UTIL_H_
tensorflow/dtensor/cc/slice_util.cc
+213
View File
@@ -0,0 +1,213 @@
/* 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/dtensor/cc/slice_util.h"
#include <optional>
#include <string>
#include <vector>
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
namespace tensorflow {
namespace dtensor {
namespace slice_util {
namespace {
// Computes the size of the ellipsis and the output rank.
StatusOr<int64_t> GetEllipsisSize(int64_t input_rank,
const std::vector<Token>& tokens,
int64_t* output_rank) {
bool found = false;
int64_t regular_axis = 0;
int64_t new_axis = 0;
int64_t shrink_axis = 0;
for (const auto& token : tokens) {
switch (token.token_type) {
case Token::ELLIPSIS:
if (found) {
return absl::InvalidArgumentError(
"More than one ellipsis was found.");
}
found = true;
break;
case Token::NEW_AXIS:
++new_axis;
break;
case Token::SHRINK_AXIS:
++shrink_axis;
break;
case Token::REGULAR:
++regular_axis;
break;
}
}
int64_t ellipsis_size = input_rank - (regular_axis + shrink_axis);
if (found && ellipsis_size < 0) {
return absl::InvalidArgumentError(absl::StrCat(
"Ellipsis was found, but there is no remaining axis for it.",
" input_rank=", input_rank, " regular_axis=", regular_axis,
" shrink_axis=", shrink_axis));
}
*output_rank = regular_axis + ellipsis_size + new_axis;
return ellipsis_size;
}
} // namespace
Token Token::normalize(int64_t dim_size) const {
if (dynamic_mask) {
return *this;
}
int64_t new_begin = begin;
int dir = (stride > 0) ? 1 : -1;
if (begin_mask) {
if (dir > 0) {
new_begin = 0;
} else {
new_begin = dim_size - 1;
}
}
int64_t new_end = end;
if (end_mask) {
if (dir > 0) {
new_end = dim_size;
} else {
new_end = -1;
}
}
// Shift begin and end by same number of periods to distinguish full cycle
// from empty.
int64_t shift = (new_begin - new_begin % dim_size);
new_begin -= shift;
new_end -= shift;
int64_t n = dir * (new_end - new_begin + stride - dir) / (dir * stride);
// Round end by cycle size to ensure `(end - begin) / strides` is the
// number of result elements. To support cases like begin=0, end=-1.
if (n < 0) {
new_end = new_end + dir * dim_size;
}
n = dir * (new_end - new_begin + stride - dir) / (dir * stride);
new_end = new_begin + n * stride;
Token r = *this;
r.begin = new_begin;
r.end = new_end;
return r;
}
// Returns a Token for local slicing if no relayout along this axis
// is needed. If no such local slicing is possible, returns nullopt.
std::optional<Token> Token::GetLocalToken(int64_t dim_size,
int64_t num_shards) const {
Token token = normalize(dim_size);
VLOG(5) << "Compute: "
<< "dim_size=" << dim_size << " num_shards=" << num_shards
<< " token.begin=" << token.begin << " token.end=" << token.end
<< " token.stride=" << token.stride;
if (token.begin_mask && token.end_mask) return token;
if (token.dynamic_mask) return std::nullopt;
if (token.stride < 0) return std::nullopt;
int64_t shard_dim_size = dim_size / num_shards;
if (shard_dim_size % token.stride == 0) {
// Simple striped slicing, where every 1 out of stride items
// are selected can remain sharded the same way.
if (token.begin >= 0 && token.begin < token.stride &&
token.end >= dim_size && token.end < dim_size + token.stride) {
token.end = shard_dim_size + (token.end - dim_size);
return token;
}
}
return std::nullopt;
}
absl::Status TokenProcessor::Run(const std::vector<Token>& tokens) {
int64_t input_rank = input_rank_;
int64_t output_rank;
TF_ASSIGN_OR_RETURN(int64_t ellipsis_size,
GetEllipsisSize(input_rank, tokens, &output_rank));
PrepareResults(tokens.size(), input_rank, output_rank);
bool out_of_bound = false;
int64_t input_index = 0;
int64_t output_index = 0;
for (const auto& token : tokens) {
switch (token.token_type) {
case Token::ELLIPSIS:
VisitEllipsisAxis(token);
out_of_bound = VisitLoop(input_rank, output_rank, ellipsis_size,
&input_index, &output_index);
ellipsis_size = 0;
break;
case Token::SHRINK_AXIS:
VisitShrinkAxis(token, input_index, output_index);
++input_index;
break;
case Token::NEW_AXIS:
VisitNewAxis(token, input_index, output_index);
++output_index;
break;
case Token::REGULAR:
if (input_index >= input_rank) {
out_of_bound = true;
break;
}
VisitRegularAxis(token, input_index, output_index);
++input_index;
++output_index;
break;
}
if (out_of_bound) {
break;
}
}
if (ellipsis_size > 0) {
out_of_bound = VisitLoop(input_rank, output_rank, ellipsis_size,
&input_index, &output_index);
}
if (out_of_bound) {
return absl::InvalidArgumentError(
"Reading axis beyond the input tensor's rank. "
"The slicing token is incorrect.");
}
return FinalizeResults(input_rank, output_rank);
}
bool TokenProcessor::VisitLoop(int64_t input_rank, int64_t output_rank,
int64_t ellipsis_size, int64_t* input_index,
int64_t* output_index) {
for (int64_t k = 0; k < ellipsis_size; ++k) {
if (*input_index >= input_rank) {
return true;
}
VisitImplicitAxis(*input_index, *output_index);
++*input_index;
++*output_index;
}
return false;
}
} // namespace slice_util
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/cc/slice_util.h
+315
View File
@@ -0,0 +1,315 @@
/* 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_DTENSOR_CC_SLICE_UTIL_H_
#define TENSORFLOW_DTENSOR_CC_SLICE_UTIL_H_
#include <optional>
#include <string>
#include <vector>
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
namespace tensorflow {
namespace dtensor {
namespace slice_util {
// Defines a token of the strided slicing mini-language.
// Refer to the definition of StridedSlice Op for the informal definition of
// the language. During slicing, axes of the input tensor are processed one
// by one according to the tokens of the slicing spec vector.
struct Token {
enum TokenType {
REGULAR, // Slice the current axis by begin/end/begin_mask/end_mask and
// stride.
NEW_AXIS, // Add a new axis at the current location to the output.
ELLIPSIS, // Copy over following axes to the output till the ellipsis ends.
SHRINK_AXIS // Like a regular axis, but sequeeze this axis from output
// after slicing.
} token_type;
int64_t begin = 0; // Begin of the slice.
int64_t end = 0; // End of the slice.
int64_t stride = 0; // Stride of the slice.
bool dynamic_mask = false; // If begin, end, or stride is a dynamic value.
bool begin_mask = false; // True if the begin is maximal.
bool end_mask = false; // True if the end is maximal.
Token() = default;
Token(TokenType token_type, int64_t begin, int64_t end, int64_t stride,
bool dynamic_mask = false, bool begin_mask = false,
bool end_mask = false)
: token_type(token_type),
begin(begin),
end(end),
stride(stride),
dynamic_mask(dynamic_mask),
begin_mask(begin_mask),
end_mask(end_mask) {}
// Normalizes the token such that (end - begin) is evenly divided by stride,
// and the result equals the total elements after the slicing.
Token normalize(int64_t dim_size) const;
std::optional<Token> GetLocalToken(int64_t dim_size,
int64_t num_shards) const;
};
// TODO(feyu): is there a C++ way to do vari args and templates move this out
// of this class?
template <typename T, typename... Types>
StatusOr<T> CreateAndRun(const std::vector<Token>& tokens, Types... args) {
T visitor(args...);
TF_RETURN_IF_ERROR(visitor.Run(tokens));
return visitor;
}
class TokenProcessor {
public:
explicit TokenProcessor(int64_t input_rank) : input_rank_(input_rank) {}
virtual ~TokenProcessor() = default;
absl::Status Run(const std::vector<Token>& tokens);
protected:
// Loop for an ellipsis or the unconsumed axes in the end.
bool VisitLoop(int64_t input_rank, int64_t output_rank, int64_t ellipsis_size,
int64_t* input_index, int64_t* output_index);
virtual void VisitImplicitAxis(int64_t input_index, int64_t output_index) = 0;
virtual void VisitEllipsisAxis(const Token& token) = 0;
virtual void VisitShrinkAxis(const Token& token, int64_t input_index,
int64_t output_index) = 0;
virtual void VisitNewAxis(const Token& token, int64_t input_index,
int64_t output_index) = 0;
virtual void VisitRegularAxis(const Token& token, int64_t input_index,
int64_t output_index) = 0;
virtual void PrepareResults(int64_t spec_rank, int64_t input_rank,
int64_t output_rank) = 0;
virtual absl::Status FinalizeResults(int64_t input_rank,
int64_t output_rank) = 0;
private:
const int64_t input_rank_;
};
// Forward layout inference of from a StridedSlice token vector.
//
// For value_layout = StridedSlice(input_layout, tokens)
//
// The inference consumes input_layout, and produces:
// - a planned expander_input_layout that is suitable for SPMD expansion.
// - a planned expander_value_layout that is suitable for SPMD expansion.
// - a local_tokens vector for the arguments of the post-SPMD StridedSliceOp.
// expander_input_layout and expander_value_layout are consistent with
// local_tokens.
class ForwardLayoutInference : public TokenProcessor {
public:
ForwardLayoutInference(const Layout& input_layout,
const llvm::ArrayRef<int64_t> input_shape)
: TokenProcessor(input_shape.size()),
input_layout_(input_layout),
input_shape_(input_shape),
input_sharding_(input_layout.sharding_spec_strs()) {}
const Layout& expander_value_layout() const { return expander_value_layout_; }
const Layout& expander_input_layout() const { return expander_input_layout_; }
const std::vector<Token>& local_tokens() const { return local_tokens_; }
protected:
void VisitEllipsisAxis(const Token& token) override {
local_tokens_.push_back(token);
}
void VisitImplicitAxis(int64_t input_index, int64_t output_index) override {
expander_input_sharding_.push_back(input_sharding_[output_index]);
expander_value_sharding_.push_back(input_sharding_[output_index]);
}
void VisitShrinkAxis(const Token& token, int64_t input_index,
int64_t output_index) override {
local_tokens_.push_back(token);
expander_input_sharding_.push_back(Layout::kUnshardedDim);
// Skips this axis from values, since it will be removed from the inputs.
}
void VisitNewAxis(const Token& token, int64_t input_index,
int64_t output_index) override {
local_tokens_.push_back(token);
expander_value_sharding_.push_back(Layout::kUnshardedDim);
}
void VisitRegularAxis(const Token& token, int64_t input_index,
int64_t output_index) override {
auto local_token = token.GetLocalToken(
/*dim_size=*/input_shape_[input_index],
/*num_shards*/ input_layout_.num_shards_for_dim(input_index));
std::string sharding = input_sharding_[input_index];
if (local_token.has_value()) {
local_tokens_.push_back(*local_token);
} else {
sharding = Layout::kUnshardedDim;
local_tokens_.push_back(token);
}
expander_value_sharding_.push_back(sharding);
expander_input_sharding_.push_back(sharding);
}
void PrepareResults(int64_t spec_rank, int64_t input_rank,
int64_t output_rank) override {
local_tokens_.reserve(spec_rank);
expander_input_sharding_.reserve(input_rank);
expander_value_sharding_.reserve(output_rank);
}
absl::Status FinalizeResults(int64_t input_rank,
int64_t output_rank) override {
DCHECK_EQ(expander_input_sharding_.size(), input_rank);
DCHECK_EQ(expander_value_sharding_.size(), output_rank);
TF_ASSIGN_OR_RETURN(
expander_input_layout_,
Layout::GetLayout(expander_input_sharding_, input_layout_.mesh()));
TF_ASSIGN_OR_RETURN(
expander_value_layout_,
Layout::GetLayout(expander_value_sharding_, input_layout_.mesh()));
return absl::OkStatus();
}
private:
const Layout& input_layout_;
const llvm::ArrayRef<int64_t> input_shape_;
std::vector<std::string> input_sharding_;
std::vector<std::string> expander_value_sharding_;
std::vector<std::string> expander_input_sharding_;
// Outputs
Layout expander_value_layout_;
Layout expander_input_layout_;
std::vector<Token> local_tokens_;
};
// Backward layout inference for a StridedSlice token vector.
//
// For value_layout = StridedSlice(input_layout, tokens)
//
// The inference consumes value_layout, and produces:
// - a planned expander_input_layout that is suitable for SPMD expansion.
// - a planned expander_value_layout that is suitable for SPMD expansion.
// - a local_tokens vector for the arguments of the post-SPMD StridedSliceOp.
// expander_input_layout and expander_value_layout are consistent with
// local_tokens.
class BackwardLayoutInference : public TokenProcessor {
public:
BackwardLayoutInference(const Layout& value_layout,
const llvm::ArrayRef<int64_t> input_shape)
: TokenProcessor(input_shape.size()),
value_layout_(value_layout),
input_shape_(input_shape),
value_sharding_(value_layout.sharding_spec_strs()) {}
const Layout& expander_input_layout() const { return expander_input_layout_; }
const Layout& expander_value_layout() const { return expander_value_layout_; }
const std::vector<Token>& local_tokens() const { return local_tokens_; }
protected:
void VisitEllipsisAxis(const Token& token) override {
local_tokens_.push_back(token);
}
void VisitImplicitAxis(int64_t input_index, int64_t output_index) override {
expander_input_sharding_.push_back(value_sharding_[output_index]);
expander_value_sharding_.push_back(value_sharding_[output_index]);
}
void VisitShrinkAxis(const Token& token, int64_t input_index,
int64_t output_index) override {
local_tokens_.push_back(token);
// There is no constraint on the input sharding, but we prefer to keep it
// unsharded to avoid inserting relayout toward the internal input layout.
expander_input_sharding_.push_back(Layout::kUnshardedDim);
}
void VisitNewAxis(const Token& token, int64_t input_index,
int64_t output_index) override {
local_tokens_.push_back(token);
// No corresponding input axis.
expander_value_sharding_.push_back(Layout::kUnshardedDim);
}
void VisitRegularAxis(const Token& token, int64_t input_index,
int64_t output_index) override {
auto local_token = token.GetLocalToken(
/*dim_size=*/input_shape_[input_index],
/*num_shards*/ value_layout_.num_shards_for_dim(output_index));
if (local_token.has_value()) {
std::string sharding = value_sharding_[output_index];
local_tokens_.push_back(*local_token);
expander_input_sharding_.push_back(sharding);
expander_value_sharding_.push_back(sharding);
} else {
local_tokens_.push_back(token);
// There is no constraint on the input sharding, but we prefer to keep it
// unsharded to avoid inserting relayout toward the internal input layout.
expander_input_sharding_.push_back(Layout::kUnshardedDim);
expander_value_sharding_.push_back(Layout::kUnshardedDim);
}
}
void PrepareResults(int64_t spec_rank, int64_t input_rank,
int64_t output_rank) override {
local_tokens_.reserve(spec_rank);
expander_input_sharding_.reserve(input_rank);
expander_value_sharding_.reserve(output_rank);
}
absl::Status FinalizeResults(int64_t input_rank,
int64_t output_rank) override {
DCHECK_EQ(expander_input_sharding_.size(), input_rank);
DCHECK_EQ(expander_value_sharding_.size(), output_rank);
TF_ASSIGN_OR_RETURN(
expander_input_layout_,
Layout::GetLayout(expander_input_sharding_, value_layout_.mesh()));
TF_ASSIGN_OR_RETURN(
expander_value_layout_,
Layout::GetLayout(expander_value_sharding_, value_layout_.mesh()));
return absl::OkStatus();
}
private:
const Layout& value_layout_;
const llvm::ArrayRef<int64_t> input_shape_;
std::vector<std::string> value_sharding_;
std::vector<std::string> expander_input_sharding_;
std::vector<std::string> expander_value_sharding_;
// Outputs
Layout expander_input_layout_;
Layout expander_value_layout_;
std::vector<Token> local_tokens_;
};
} // namespace slice_util
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_CC_SLICE_UTIL_H_
tensorflow/dtensor/cc/small_constant_optimization.cc
+169
View File
@@ -0,0 +1,169 @@
/* Copyright 2022 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/dtensor/cc/small_constant_optimization.h"
#include <cstdint>
#include <memory>
#include <optional>
#include <utility>
#include <vector>
#include "absl/algorithm/container.h"
#include "absl/strings/string_view.h"
#include "tensorflow/c/eager/c_api.h"
#include "tensorflow/c/eager/c_api_experimental.h"
#include "tensorflow/c/tf_status.h"
#include "tensorflow/c/tf_tensor_internal.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/platform/ctstring_internal.h"
#include "tensorflow/core/platform/protobuf.h"
#include "tensorflow/dtensor/cc/constants.h"
#include "tensorflow/dtensor/proto/layout.pb.h"
namespace tensorflow {
namespace dtensor {
namespace {
constexpr TF_DataType kAllowedDataType[] = {TF_INT32, TF_INT64, TF_FLOAT,
TF_STRING};
void AppendIntValues(const int num_of_elements, const int* int_values,
TensorProto* proto) {
for (int i = 0; i < num_of_elements; ++i) {
proto->add_int_val(int_values[i]);
}
}
void AppendInt64Values(const int num_of_elements, const int64_t* int64_values,
TensorProto* proto) {
for (int i = 0; i < num_of_elements; ++i) {
proto->add_int64_val(int64_values[i]);
}
}
void AppendStringValues(const int num_of_elements,
const TF_TString* string_values, TensorProto* proto) {
for (int i = 0; i < num_of_elements; ++i) {
proto->add_string_val(
std::string(TF_TString_GetDataPointer(&string_values[i]),
TF_TString_GetSize(&string_values[i])));
}
}
void AppendFloatValues(const int num_of_elements, const float* float_values,
TensorProto* proto) {
for (int i = 0; i < num_of_elements; ++i) {
proto->add_float_val(float_values[i]);
}
}
} // namespace
std::optional<NodeDef> ExtractSmallTensorValue(TFE_Context* context,
TFE_TensorHandle* tensor,
const Layout& layout,
TF_Status* status) {
if (!layout.IsFullyReplicated()) return std::nullopt;
auto num_elements = TFE_TensorHandleNumElements(tensor, status);
if (TF_GetCode(status) != TF_OK) return std::nullopt;
if (num_elements >= kSmallTensorThreshold) return std::nullopt;
// Check the DType before attempting to resolve the tensor so we don't try to
// copy resource-dtype tensors off the DTensor device. Currently we only
// extract small int32/int64_t tensors, primarily to catch shapes and axes,
// and tf_string tensors that are mostly used in save/restore ops.
const auto& dtype = TFE_TensorHandleDataType(tensor);
if (absl::c_find(kAllowedDataType, dtype) == std::end(kAllowedDataType)) {
return std::nullopt;
}
// This is the enum from protobuf, or the following AddNodeAttr will always
// set the integer field.
const auto& datatype = static_cast<DataType>(dtype);
std::unique_ptr<TF_Tensor, decltype(&TF_DeleteTensor)> value_tensor(
TFE_TensorHandleResolve(tensor, status), TF_DeleteTensor);
if (TF_GetCode(status) != TF_OK) return std::nullopt;
NodeDef node_def;
node_def.set_op("Const");
AddNodeAttr("dtype", datatype, &node_def);
TensorProto tensor_proto;
tensor_proto.set_dtype(datatype);
void* raw_data = TF_TensorData(value_tensor.get());
if (!raw_data) {
TF_SetStatus(status, TF_INTERNAL, "TF_TensorData returned nullptr.");
return std::nullopt;
}
switch (dtype) {
case TF_INT32:
AppendIntValues(num_elements, static_cast<const int*>(raw_data),
&tensor_proto);
break;
case TF_INT64:
AppendInt64Values(num_elements, static_cast<const int64_t*>(raw_data),
&tensor_proto);
break;
case TF_STRING:
AppendStringValues(num_elements, static_cast<const TF_TString*>(raw_data),
&tensor_proto);
break;
case TF_FLOAT:
AppendFloatValues(num_elements, static_cast<const float*>(raw_data),
&tensor_proto);
break;
default:
TF_SetStatus(status, TF_INTERNAL,
absl::StrCat("dtype: ", dtype,
" fell through the supported extraction list. "
"This should not happen.")
.c_str());
return std::nullopt;
}
std::vector<int64_t> dim_list;
int num_dims = value_tensor->tensor->NumDims();
dim_list.reserve(num_dims);
for (int i = 0; i < num_dims; ++i) {
dim_list.push_back(value_tensor->tensor->Dim(i));
}
TensorShape shape(std::move(dim_list));
shape.AsProto(tensor_proto.mutable_tensor_shape());
AddNodeAttr("value", tensor_proto, &node_def);
AddNodeAttr(kLayoutAttr, {layout.ToString()}, &node_def);
AddNodeAttr(kMeshAttr, layout.mesh().ToString(), &node_def);
return node_def;
}
bool NodeDefsHaveDifferentTensorProto(const NodeDef& a, const NodeDef& b) {
const TensorProto* tensor_proto_a;
bool read_a_tensor_proto = TryGetNodeAttr(a, "value", &tensor_proto_a);
if (!read_a_tensor_proto) return true;
const TensorProto* tensor_proto_b;
bool read_b_tensor_proto = TryGetNodeAttr(b, "value", &tensor_proto_b);
if (!read_b_tensor_proto) return true;
return !protobuf::util::MessageDifferencer::Equals(*tensor_proto_a,
*tensor_proto_b);
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/cc/small_constant_optimization.h
+46
View File
@@ -0,0 +1,46 @@
/* Copyright 2022 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_DTENSOR_CC_SMALL_CONSTANT_OPTIMIZATION_H_
#define TENSORFLOW_DTENSOR_CC_SMALL_CONSTANT_OPTIMIZATION_H_
#include <optional>
#include "tensorflow/c/eager/c_api.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
namespace tensorflow {
namespace dtensor {
// Attempt to convert small constant tensors into a constant NodeDef operation.
// This constant value will be available for constant propagation in DTensor and
// MLIR.
// This conversion is currently required for some DTensor operations. In
// particular, reductions require access to the axis argument at compilation
// time. While this is not strictly necessary, it greatly simplifies SPMD code
// generation and is generally available.
std::optional<NodeDef> ExtractSmallTensorValue(TFE_Context* context,
TFE_TensorHandle* tensor,
const Layout& layout,
TF_Status* status);
// Returns true if the tensor proto of a and b are different.
bool NodeDefsHaveDifferentTensorProto(const NodeDef& a, const NodeDef& b);
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_CC_SMALL_CONSTANT_OPTIMIZATION_H_
tensorflow/dtensor/cc/small_constant_optimization_test.cc
+41
View File
@@ -0,0 +1,41 @@
// Copyright 2025 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/dtensor/cc/small_constant_optimization.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
namespace tensorflow {
namespace dtensor {
namespace {
TEST(SmallConstantOptimization, NullTensorReturnsNullopt) {
Layout layout = Layout::Empty();
TF_Status* tf_status = TF_NewStatus();
std::optional<NodeDef> result = ExtractSmallTensorValue(
/*context=*/nullptr,
/*tensor=*/nullptr, layout, tf_status);
EXPECT_FALSE(result.has_value());
EXPECT_EQ(TF_GetCode(tf_status), TF_INVALID_ARGUMENT);
TF_DeleteStatus(tf_status);
}
} // namespace
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/cc/tensor_layout.cc
+1491
View File
File diff suppressed because it is too large Load Diff
tensorflow/dtensor/cc/tensor_layout.h
+465
View File
@@ -0,0 +1,465 @@
/* Copyright 2022 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_DTENSOR_CC_TENSOR_LAYOUT_H_
#define TENSORFLOW_DTENSOR_CC_TENSOR_LAYOUT_H_
#include <algorithm>
#include <cstdint>
#include <iostream>
#include <string>
#include <utility>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/container/inlined_vector.h"
#include "absl/strings/string_view.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/proto/layout.pb.h"
// Definitions for DTensor mesh & layout.
//
// A mesh describes how a set of devices is partitioned.
// A layout describes how a distributed tensor is partitioned across a mesh (and
// thus across devices). Defining tensor layouts in terms of mesh dimensions
// allows us to efficiently determine the communication required when computing
// an operation with tensors of different layouts.
namespace tensorflow {
namespace dtensor {
// Returns true if `size` is a dynamic size based on either MLIR and TF
// standards.
bool IsDynamicSize(int64_t size);
// Returns true if `shape` is a dynamic shape based on either MLIR and TF
// standards.
bool IsDynamicShape(absl::Span<const int64_t> shape);
// The location of a device in a mesh.
//
// Each device has a unique location in the mesh, which is indicated by the
// offset in each mesh dimension. e.g. a mesh:
//
// [x:4, y:3, z:2]
//
// Must consist of 24 devices placed densely into the corresponding 3D space.
using DeviceLocation = absl::InlinedVector<int64_t, 4>;
// A shard refers to a partition of a tensor. Shards are arranged in
// ShardVectors that contains a list of Shards and a list of integers
// representing the number of shards in each dimension.
//
// Example: layout = sharding_specs:x,y, mesh:|x=2,y=2|. This can be represented
// with a ShardVector:
// - shards = (1,1), (1,2), (2,1), (2,2)
// - num_shards_per_dim = (2,2).
//
// The number of elements in each shard matches the tensor rank.
using Shard = std::vector<int>;
struct ShardVector {
bool operator==(const ShardVector& other) const;
bool operator!=(const ShardVector& other) const { return !(*this == other); }
std::string ToString() const;
bool ContainsShard(const Shard& shard) const;
std::vector<Shard> shards;
std::vector<int> num_shards_per_dim;
};
struct MeshDimension {
MeshDimension(const std::string& name, int64_t size)
: name(std::move(name)), size(size) {}
MeshDimension() = default;
std::string name;
int64_t size;
};
class Mesh {
public:
// Failed serialized strings are represented with an empty string, therefore
// we use this string representation of an empty mesh instead to avoid
// confusion.
static constexpr const char* kEmptyMeshString = "empty_mesh";
static constexpr const char* kUseXLASPMDString = "use_xla_spmd";
static constexpr bool kUseXLASPMD = false;
enum class MeshType {
kTile,
kSingleDevice,
};
static Mesh Empty();
bool IsEmpty() const;
Mesh() { mesh_type_ = MeshType::kTile; }
bool IsTile() const { return mesh_type_ == MeshType::kTile; }
bool IsSingleDevice() const { return mesh_type_ == MeshType::kSingleDevice; }
// Creates fully defined mesh.
//
// When `use_xla_spmd` is true, all ops running on this mesh will use XLA SPMD
// instead of DTensor SPMD.
static Mesh CreateMesh(const std::string& mesh_name,
const std::vector<std::string>& dim_names,
const std::vector<std::int64_t>& mesh_shape,
const std::vector<std::int64_t>& global_device_ids,
const std::vector<std::string>& global_devices_str,
const std::vector<std::int64_t>& local_device_ids,
const std::vector<std::string>& local_devices_str,
bool use_xla_spmd = Mesh::kUseXLASPMD);
// Parses from MeshProto.
static StatusOr<Mesh> ParseFromProto(const MeshProto& proto);
// Parses from a human readable string version of the mesh, currently used
// to represent meshes in MLIR:
// mesh = <name|List[MeshDim]|List[GlobalId]|List[LocalId]|List[Devices]>
//
// Example:
// mesh =
// <name|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3>
static StatusOr<Mesh> FromString(absl::string_view str);
std::string ToString() const;
StatusOr<MeshProto> ToProto() const;
// Creates mesh without specific devices associated to it (aka abstract mesh).
// This is an experimental API. Use only if strictly needed.
static StatusOr<Mesh> GetAbstractMesh(
const std::string& name, const std::vector<MeshDimension>& mesh_dims);
// Creates fully defined mesh.
static StatusOr<Mesh> GetMesh(
const std::string& name, const std::vector<MeshDimension>& mesh_dims,
const std::vector<std::int64_t>& global_device_ids,
const std::vector<std::int64_t>& local_device_ids,
const std::vector<std::string>& local_devices,
const std::vector<std::string>& global_devices,
bool use_xla_spmd = Mesh::kUseXLASPMD);
static StatusOr<Mesh> GetSingleDeviceMesh(absl::string_view single_device);
bool is_cpu_mesh() const { return device_type() == "CPU"; }
bool is_epu_mesh() const { return device_type() == "EPU"; }
bool is_gpu_mesh() const { return device_type() == "GPU"; }
bool is_tpu_mesh() const { return device_type() == "TPU"; }
// Returns whether the mesh is a remote mesh.
bool is_remote() const {
return local_device_ids_.empty() && !global_device_ids_.empty();
}
StatusOr<Mesh> host_mesh() const { return ToDeviceType("CPU"); }
// Device information methods.
std::string device_type() const;
// Takes an index in the flattened list of devices and returns a location
// in the mesh.
StatusOr<const DeviceLocation> device_location(int offset) const;
int64_t num_devices() const;
absl::Span<const std::string> local_devices() const { return local_devices_; }
absl::Span<const int64_t> local_device_ids() const {
return local_device_ids_;
}
// Parses names of local_devices according to TF's Device Name Utils.
StatusOr<const std::vector<DeviceNameUtils::ParsedName>> ParsedDevices()
const;
// Convert to given device type.
StatusOr<Mesh> ToDeviceType(const std::string& device_type) const;
std::vector<std::string> hosts() const;
// Consumes a location in the mesh and returns its corresponding index in
// the flattened list of devices.
int64_t GetFlattenedCoordinate(const DeviceLocation& loc) const;
const MeshDimension& dim(int64_t index) const { return mesh_dims_[index]; }
std::vector<MeshDimension> dims() const { return mesh_dims_; }
// Returns size of mesh dimension.
StatusOr<int64_t> dim_size(absl::string_view name) const;
// Returns list of mesh dimension sizes.
std::vector<int64_t> dim_sizes() const;
const std::string& dim_name(int64_t index) const {
return mesh_dims_[index].name;
}
int64_t min_global_device_id() const {
DCHECK(!global_device_ids_.empty());
return *std::min_element(global_device_ids_.begin(),
global_device_ids_.end());
}
int64_t num_local_devices() const { return local_devices_.size(); }
absl::Span<const int64_t> global_device_ids() const {
return global_device_ids_;
}
const std::vector<std::string>& global_devices() const {
return global_devices_;
}
// Returns index of given dim_name in the mesh.
StatusOr<int32_t> idx_for_dim(absl::string_view dim_name) const;
// Returns the index of MeshDimension in mesh where the mesh dimension name is
// `mesh_name`.
int GetMeshDimIndexWithName(const std::string& mesh_name) const;
bool IsMeshDim(const std::string& dim_name) const;
std::vector<std::string> MeshDimNames() const;
int64_t rank() const;
int64_t size() const;
bool use_xla_spmd() const { return use_xla_spmd_; }
const std::string& name() const { return name_; }
absl::string_view single_device() const { return single_device_; }
// Global unique fingerprint. Same on different workers.
uint64_t GlobalFingerprint() const;
// Uses proto to compare the equality. If any conversion to proto fails,
// returns false.
bool operator==(const Mesh& b) const;
bool operator!=(const Mesh& b) const { return !((*this) == b); }
bool operator<(const Mesh& b) const {
return this->ToString() < b.ToString();
}
template <typename H>
friend H AbslHashValue(H h, const Mesh& m) {
return H::combine(std::move(h), m.ToString());
}
// A map from mesh names to their corresponding core ID mappings. The core ID
// mapping is stored as a vector. The i-th element in the vector is the ID of
// the core represented by global device ID of i in this mesh.
//
// The entry stored under the empty name key (the so-called "default mapping"
// in some comments) is special. It is always set at the end of TPU
// initialization. It represents the mapping for any mesh whose global device
// IDs follow TF task-device ordinals. Legacy and test meshes created without
// using the `create_tpu_mesh` helper follow that rule and can use this entry.
static std::map<std::string, std::vector<int>>& tpu_core_ids();
// The host mesh associated with any user-defined TPU mesh.
static std::string& tpu_host_mesh();
private:
MeshType mesh_type_;
std::string name_;
// The following fields store the information for tile sharding. Usable only
// when the mesh has type `kTile`.
std::vector<MeshDimension> mesh_dims_;
std::vector<std::string> local_devices_;
std::vector<int64_t> local_device_ids_;
std::vector<int64_t> global_device_ids_;
std::vector<std::string> global_devices_;
bool use_xla_spmd_ = Mesh::kUseXLASPMD;
// Stores the device when mesh is used for representing the state of a tensor
// on one device. Usable only when the mesh has type `kSingleDevice`.
std::string single_device_;
};
// Obtain all possible forms of indexing a mesh.
//
// e.g. given a mesh with dimensions [x=2, y=3], returns {
// [0, 0], [0, 1], [0, 2],
// [1, 0], [1, 1], [1, 2]
// }
std::vector<DeviceLocation> ComputeDeviceLocations(const Mesh& mesh);
class Layout {
public:
enum class LayoutType {
kEmpty,
kStatic,
kSingleDevice,
kParted,
};
static constexpr const char* kPartedPrefix = "parted:";
static constexpr const char* kStaticPrefix = "sharding_specs:";
static constexpr const char* kSingleDevicePrefix = "maximal:";
static constexpr const char* kUnshardedDim = "unsharded";
// This spec should only be used to express no preferred sharding in the
// Layout propagation algorithm.
static constexpr const char* kAny = "any";
// Failed serialized strings are represented with an empty string, therefore
// we use this string representation of an empty layout instead to avoid
// confusion.
static constexpr const char* kEmptyLayoutString = "empty_layout";
// Used for the relayout operation, to allow relayout act as an identity on
// the layout for the given dimension.
static constexpr const char* kMatch = "match";
Layout() = default;
Layout(const Layout& other) = default;
bool IsSingleDevice() const { return mesh_.IsSingleDevice(); }
// Returns empty layout.
static Layout Empty();
// Parses from LayoutProto.
static StatusOr<Layout> FromProto(const LayoutProto& proto);
// Parses from a human readable string version of the layout, currently used
// to represent layouts in MLIR:
// layout = <sharding_specs:List[specs] mesh:name|List[MeshDim]|
// List[GlobalId]|List[LocalId]|List[Devices]>
//
// Example:
// layout = <sharding_specs:x,not_sharded mesh:name|x=2,y=2|0,1,2,3|0,1,2,3|
// /job:localhost/task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1,
// /job:localhost/task:0/device:CPU:2,/job:localhost/task:0/device:CPU:3>
static StatusOr<Layout> FromString(absl::string_view layout_str);
// Creates human readable string version of a layout.
std::string ToString() const;
StatusOr<LayoutProto> ToProto() const;
LayoutType type() const { return type_; }
const Mesh& mesh() const { return mesh_; }
static Layout ReplicatedOnMesh(const Mesh& mesh, int rank);
static Layout BatchShardedOnMesh(const Mesh& mesh, int rank,
const std::string& mesh_dim, int axis = 0);
static Layout ReplicatedLike(const Layout& layout);
static Layout BatchShardedLike(const Layout& layout,
const std::string& mesh_dim, int axis = 0);
static Layout AnyOnMesh(const Mesh& mesh, int rank);
// Creates a mesh of unique shards.
Mesh ReducedMesh() const;
// Deprecated: Replace calls with GetLayout that creates a new instance.
void set_mesh(Mesh mesh) { mesh_ = mesh; }
// Returns a layout for the transposed matrix for given layout. This assumes
// that only the last two dimensions are used for matrix computation and all
// dimensions before are batch dimensions.
static StatusOr<Layout> Transposed2D(const Layout& layout);
static bool IsUnshardedDimension(const absl::string_view name) {
return name == kUnshardedDim;
}
static bool IsShardedDimension(const absl::string_view name) {
return !IsUnshardedDimension(name);
}
static StatusOr<Layout> GetLayout(
LayoutType type, const std::vector<std::string>& sharding_spec_strs,
const Mesh& mesh);
// Deprecated: Update all call sites to GetLayout(LayoutType::kStatic, ...);
static StatusOr<Layout> GetLayout(
const std::vector<std::string>& sharding_spec_strs, const Mesh& mesh) {
return GetLayout(LayoutType::kStatic, sharding_spec_strs, mesh);
}
// Deprecated: Update all call sites to GetLayout(LayoutType::kSingleDevice,
// {}, ...);
static StatusOr<Layout> GetSingleDeviceLayout(const Mesh& mesh) {
return GetLayout(LayoutType::kSingleDevice, {}, mesh);
}
// Makes a new layout from this one dropping the given dimensions.
// If keep_dims is true, the dimensions are replicated rather than
// deleted.
StatusOr<Layout> GetLayoutWithReducedDims(
const absl::flat_hash_set<int>& reduced_dims, bool keep_dims) const;
// Converts the Layout to Parted.
StatusOr<Layout> ToParted() const {
return GetLayout(LayoutType::kParted, sharding_specs_, mesh_);
}
// Truncates a layout at the front or back, depending on the value of end.
// end = false returns the layout up to the split point,
// end = true returns the layout from the split point.
Layout Truncate(int64_t split_point, bool end = false) const;
// Left or right pad the layout to a max rank.
Layout LeftPad(int64_t rank) const;
// Minimally pads replicated axes on the left, or removes axes on the right,
// such that the result layout has the provided rank.
StatusOr<Layout> EnsureRank(int64_t rank) const;
bool IsFullyReplicated() const;
bool IsLastDimReplicated() const;
// Checks that the last N-1 dimensions are replicated
bool IsBatchParallel() const;
// Checks that the dimensions from [-non_batch_rank, end) are replicated.
bool IsBatchParallel(int non_batch_rank) const;
bool IsEmpty() const;
// Compute global shape using the layout and provided local_shape.
// Optionally take a second parameter `local_shapes` that represents the shape
// of all local tensors.
std::vector<int64_t> GlobalShapeFromLocalShape(
absl::Span<const int64_t> local_shape,
const std::vector<std::vector<int64_t>>* local_shapes = nullptr) const;
std::vector<int64_t> LocalShapeFromGlobalShape(
absl::Span<const int64_t> global_shape) const;
PartialTensorShape LocalShapeFromGlobalShape(
const PartialTensorShape& global_shape) const;
int64_t rank() const { return sharding_specs_.size(); }
size_t num_shards_for_dim(int) const;
std::vector<int32_t> num_shards() const;
// Computes the corresponding shard vector to this layout.
ShardVector GetShardVector() const;
// Returns sharding specs in string form.
std::vector<std::string> sharding_spec_strs() const;
int64_t num_devices() const { return mesh_.num_devices(); }
// Map hosts to shards.
std::map<std::string, ShardVector> HostShardMap() const;
const std::string& sharding_spec(int idx) const;
// Similar to IsEquivalentIgnoringType, but also verifies the layout type are
// equal.
bool IsEquivalent(const Layout& b) const;
// Two layouts are equivalent if they would result in the same sharding for
// the tensor. E.g. if one is unsharded and the other is sharded on a mesh
// dimension of size 1.
bool IsEquivalentIgnoringType(const Layout& b) const;
// Uses proto to compare the equality. If any conversion to proto fails,
// returns false.
bool operator==(const Layout& b) const;
bool operator!=(const Layout& b) const { return !((*this) == b); }
bool operator<(const Layout& b) const {
return this->ToString() < b.ToString();
}
private:
std::vector<std::string> sharding_specs_;
LayoutType type_;
Mesh mesh_;
};
// Takes two layouts and concatenates their TensorDimensions. If the meshes for
// the two layouts are different or both layouts are using the same mesh
// dimension returns an error rather than a layout.
StatusOr<Layout> ConcatenateLayouts(const Layout& layout_a,
const Layout& layout_b);
StatusOr<Layout> GetMostShardedLayout(const std::vector<Layout>& layouts);
StatusOr<Layout> GetLeastShardedLayout(const std::vector<Layout>& layouts);
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_CC_TENSOR_LAYOUT_H_
tensorflow/dtensor/cc/tensor_with_layout.cc
+24
View File
@@ -0,0 +1,24 @@
/* 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/dtensor/cc/tensor_with_layout.h"
namespace tensorflow {
namespace dtensor {
char TensorWithLayout::ID = 0;
}
} // namespace tensorflow
tensorflow/dtensor/cc/tensor_with_layout.h
+147
View File
@@ -0,0 +1,147 @@
/* 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_DTENSOR_CC_TENSOR_WITH_LAYOUT_H_
#define TENSORFLOW_DTENSOR_CC_TENSOR_WITH_LAYOUT_H_
#include <optional>
#include <string>
#include <vector>
#include "llvm/Support/ExtensibleRTTI.h"
#include "tensorflow/c/eager/c_api.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/platform/fingerprint.h"
#include "tensorflow/dtensor/cc/constants.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
namespace tensorflow {
namespace dtensor {
enum TensorType {
kDense = 0,
kResource = 1,
kSparse = 2,
};
class ConstValueNode {
public:
explicit ConstValueNode(std::optional<NodeDef> const_value)
: const_value_(const_value),
input_layout_for_shape_op_result_(std::nullopt) {}
// Small constant value optimization for non-resource-handle tensors.
void set_const_value(NodeDef& const_node) {
// If we extracted a constant value from the tensor, check if this
// value was the output from `tf.shape`. In this case, we need to
// forward the kShapeOpInputLayout attribute to the new node def. This
// is needed for layout propagation when running in op-by-op mode.
//
// TODO(b/162747667): Improve the presentation for Shape input Op
// layout.
if (input_layout_for_shape_op_result_.has_value()) {
AddNodeAttr(kShapeOpInputLayout,
{input_layout_for_shape_op_result_->ToString()},
&(const_node));
}
const_value_.emplace(const_node);
}
// Clears the cached const value if present.
void reset_const_value() { const_value_.reset(); }
const std::optional<NodeDef>& const_value() const { return const_value_; }
void set_input_layout_for_shape_op_result(const Layout& layout) {
input_layout_for_shape_op_result_.emplace(layout);
}
const std::optional<Layout>& shape_metadata_layout() const {
return input_layout_for_shape_op_result_;
}
private:
// The value of a small, non-resource tensor. Small constants
// are directly folded into the SPMD graph instead of being passed as inputs.
// This provides extra information to the layout propagation and SPMD passes
// during op-by-op execution. (For example, the reduction indices for Sum,
// target shapes for Rng/Reshape, etc).
std::optional<NodeDef> const_value_;
// The original input layout for a shape Op returned Tensor.
// This is used to preserve information for a shape op output so that future
// uses could recover local shape.
std::optional<Layout> input_layout_for_shape_op_result_ = std::nullopt;
};
// The representation of tensors transferred to underlying devices and the
// layout for the tensors.
class TensorWithLayout
: public llvm::RTTIExtends<TensorWithLayout, llvm::RTTIRoot> {
public:
// Gets the layout for the tensors.
virtual const Layout& layout() const = 0;
// Gets the tensor type which indicates whether the tensors are dense,
// resource or sparse.
virtual TensorType tensor_type() const = 0;
// Gets the data type of tensors.
virtual TF_DataType dtype() const = 0;
// Encodes the NodeDef via provided builder, if applicable.
virtual void EncodeAttributes(tensorflow::NodeDefBuilder& builder) const = 0;
// Generates a key which can be used for SPMD lowering.
virtual tensorflow::Fprint128 CacheKey() const = 0;
// Gets the tensor handle at position `index`. This makes sense only when the
// implementation owns a list of tensor handles. Otherwise this returns
// `nullptr`.
virtual TFE_TensorHandle* get_tensor(size_t index) const = 0;
// Gets the number of tensors.
virtual size_t num_tensors() const = 0;
// Returns a string which includes just the value and layout of the tensors.
virtual std::string SummarizeValue() const = 0;
// Returns a string which includes `SummarizeValue` along with shape and type
// information.
virtual std::string DebugString() const = 0;
// Gets the mesh for the tensors.
const Mesh& mesh() const { return layout().mesh(); }
// Computes global shape from layout & local tensor shape.
//
// For replicated layout tensors, global shape is simply the shape of local
// tensors on each device. For sharded tensor, this is the global shape
// encodes layout & local shape on each device.
virtual std::vector<int64_t> global_shape() const = 0;
// Gets a `ConstValueNode` which can operate on a `NodeDef` representing a
// small const tensor. If it is not null, it can be used in the SPMD
// expansion, regardless of which runtime is being used.
virtual ConstValueNode* const_value_node() const = 0;
// llvm::RTTIExtends ID.
static char ID; // NOLINT
};
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_CC_TENSOR_WITH_LAYOUT_H_
tensorflow/dtensor/cc/tpu_system_interface.cc
+37
View File
@@ -0,0 +1,37 @@
/* Copyright 2022 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/dtensor/cc/tpu_system_interface.h"
namespace tensorflow {
namespace dtensor {
namespace {
static TpuSystemInterface* preferred_tpu_system = nullptr;
} // namespace
void SetPreferredTpuSystem(TpuSystemInterface* tpu_system) {
if (preferred_tpu_system != nullptr) {
delete preferred_tpu_system;
}
preferred_tpu_system = tpu_system;
}
TpuSystemInterface* GetPreferredTpuSystem() { return preferred_tpu_system; }
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/cc/tpu_system_interface.h
+64
View File
@@ -0,0 +1,64 @@
/* Copyright 2022 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_DTENSOR_CC_TPU_SYSTEM_INTERFACE_H_
#define TENSORFLOW_DTENSOR_CC_TPU_SYSTEM_INTERFACE_H_
#include <vector>
#include "absl/time/time.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/status.h"
// Forward declare TFE_Context to avoid interface depending on c_api.
typedef struct TFE_Context TFE_Context;
namespace tensorflow {
namespace dtensor {
// DTensor TPU ops by default use the stream_executor-based TPU runtime.
// This class defines what an alternative runtime (e.g. TFRT) needs to be
// capable of to replace the default runtime.
class TpuSystemInterface {
public:
virtual ~TpuSystemInterface() = default;
virtual absl::Status Initialize(OpKernelContext* ctx, ResourceMgr* rmgr,
absl::Duration retry_timeout,
std::vector<int32_t>* core_id_output_vec) = 0;
virtual absl::Status Shutdown() = 0;
virtual std::vector<std::vector<int>> TPUCoreIDsToLocations(
TFE_Context* context, const std::vector<int>& tpu_core_ids) = 0;
virtual std::vector<int> TPUCoreLocationsToIDs(
TFE_Context* context,
const std::vector<std::vector<int>>& tpu_core_locations) = 0;
};
// Sets a TPU system for DTensor to initialize and shut down the TPU mesh.
// This function takes over the ownership of `tpu_system`.
void SetPreferredTpuSystem(TpuSystemInterface* tpu_system);
// Returns the currently set preferred TPU system, nullptr if none.
TpuSystemInterface* GetPreferredTpuSystem();
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_CC_TPU_SYSTEM_INTERFACE_H_
tensorflow/dtensor/cc/xla_spmd/layout_to_xla_sharding.cc
+162
View File
@@ -0,0 +1,162 @@
/* Copyright 2022 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/dtensor/cc/xla_spmd/layout_to_xla_sharding.h"
#include <cstdint>
#include <string>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/types/span.h"
#include "llvm/ADT/STLExtras.h"
#include "xla/status_macros.h"
#include "xla/xla_data.pb.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
namespace tensorflow {
namespace dtensor {
namespace {
// Produces the flat list from a slice of MajorToMinor::permutation to
// `out_devices`.
//
// This function runs recursively to expand `permutation` from major to minor.
// `sizes` is the size of mesh dimensions before the permutaiton.
// `cum_sizes` is the accumulative product of the element in `sizes`.
// `base` is the start device id of this slice of `permutation`.
void PopulateDevices(absl::Span<const int64_t> permutation,
absl::Span<const int64_t> sizes,
absl::Span<const int64_t> cum_sizes,
std::vector<int64_t>* out_devices, int64_t base = 0) {
int expanding_dim = permutation[0];
int expanding_dim_size = sizes[expanding_dim];
int expanding_cum_dim_size = cum_sizes[expanding_dim];
for (int i = 0; i < expanding_dim_size; ++i) {
if (permutation.size() == 1) {
// This is the last dimension. Fill `out_devices` with the device id.
out_devices->push_back(base + i * expanding_cum_dim_size);
} else {
// Recursively call the function to process the truncated `permutation`.
PopulateDevices(permutation.subspan(1), sizes, cum_sizes, out_devices,
base + i * expanding_cum_dim_size);
}
}
}
} // namespace
std::vector<int64_t> MeshMajorToMinor::ToDeviceList() {
std::vector<int64_t> cum_sizes(sizes.size());
int64_t cum_size = 1;
for (int i = sizes.size() - 1; i >= 0; --i) {
cum_sizes[i] = cum_size;
cum_size *= sizes[i];
}
std::vector<int64_t> devices;
devices.reserve(cum_size * sizes[0]);
PopulateDevices(permutation, sizes, cum_sizes, &devices);
return devices;
}
StatusOr<MeshMajorToMinor> ConvertMeshMajorToMinor(const Layout& layout,
const Mesh& mesh) {
MeshMajorToMinor major_to_minor;
major_to_minor.permutation.reserve(mesh.dims().size());
major_to_minor.sizes.reserve(mesh.dims().size());
absl::flat_hash_map<std::string, int64_t> dim_name_to_index_map;
// Populate dim sizes according to the order in mesh.
for (const auto& [index, mesh_dim] : llvm::enumerate(mesh.dims())) {
major_to_minor.sizes.push_back(mesh_dim.size);
dim_name_to_index_map[mesh_dim.name] = index;
}
// Sharded dims appear at the beginning of permutations according to the
// order in layout.
for (const auto& spec : layout.sharding_spec_strs()) {
if (mesh.IsMeshDim(spec)) {
const auto it = dim_name_to_index_map.find(spec);
TF_RET_CHECK(it != dim_name_to_index_map.end());
const auto& dimension_index = it->second;
major_to_minor.permutation.push_back(dimension_index);
dim_name_to_index_map.erase(it);
}
}
// Replicated dims (dims not in layout) appear at the end of permutations
// according to the order in mesh. The order here doesn't matter
// mathematically.
for (const auto& [name, unused_size] : mesh.dims()) {
if (const auto it = dim_name_to_index_map.find(name);
it != dim_name_to_index_map.end()) {
const auto& dimension_index = it->second;
major_to_minor.permutation.push_back(dimension_index);
}
}
TF_RET_CHECK(major_to_minor.permutation.size() ==
major_to_minor.sizes.size());
return major_to_minor;
}
StatusOr<::xla::OpSharding> ConvertLayoutToXlaOpSharding(const Layout& layout) {
::xla::OpSharding xla_sharding;
if (layout.IsSingleDevice()) {
xla_sharding.set_type(::xla::OpSharding::MAXIMAL);
return xla_sharding;
} else if (layout.IsFullyReplicated()) {
xla_sharding.set_type(::xla::OpSharding::REPLICATED);
return xla_sharding;
}
// If not replicated, then this is tile sharded, aka OpSharding::OTHER.
xla_sharding.set_type(::xla::OpSharding::OTHER);
const Mesh& mesh = layout.mesh();
// Compute tile_assignment_dimensions.
{
// Set Tile Assignment Dimensions by handling both partially sharded and
// fully sharded.
int32_t product_of_sharded_dimensions = 1;
for (int32_t dim_size : layout.num_shards()) {
product_of_sharded_dimensions *= dim_size;
xla_sharding.add_tile_assignment_dimensions(dim_size);
}
// Add the (n+1)th dimension representing the replicated group size. This
// only happens for partially sharded layouts.
if (product_of_sharded_dimensions != mesh.num_devices()) {
xla_sharding.add_tile_assignment_dimensions(
mesh.num_devices() / product_of_sharded_dimensions);
xla_sharding.set_replicate_on_last_tile_dim(true);
}
}
// Compute tile_assignment_devices.
TF_ASSIGN_OR_RETURN(auto major_to_minor,
ConvertMeshMajorToMinor(layout, mesh));
std::vector<int64_t> tile_assignment_devices = major_to_minor.ToDeviceList();
*(xla_sharding.mutable_tile_assignment_devices()) = {
tile_assignment_devices.begin(), tile_assignment_devices.end()};
return xla_sharding;
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/cc/xla_spmd/layout_to_xla_sharding.h
+57
View File
@@ -0,0 +1,57 @@
/* Copyright 2022 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_DTENSOR_CC_XLA_SPMD_LAYOUT_TO_XLA_SHARDING_H_
#define TENSORFLOW_DTENSOR_CC_XLA_SPMD_LAYOUT_TO_XLA_SHARDING_H_
#include <cstdint>
#include <vector>
#include "xla/xla_data.pb.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
namespace tensorflow {
namespace dtensor {
// Mhlo sharding string attribute, used for setting hlo sharding on ops, inputs,
// and outputs of a function for XLA SPMD.
constexpr char kXlaShardingAttr[] = "mhlo.sharding";
// Represents a permutation of DTensor Layout mesh dimensions from major to
// minor.
//
// Sizes of `permutation` and `sizes` must be equal.
struct MeshMajorToMinor {
// A permutation of range [0...n].
std::vector<int64_t> permutation;
// The size of mesh dimensions before the permutation.
std::vector<int64_t> sizes;
// Produces a flat list of device ids according to the permutation.
std::vector<int64_t> ToDeviceList();
};
// Get the mesh dimensions from major to minor.
StatusOr<MeshMajorToMinor> ConvertMeshMajorToMinor(const Layout& layout,
const Mesh& mesh);
// Returns an ::xla::OpSharding protobuf from `layout`.
StatusOr<::xla::OpSharding> ConvertLayoutToXlaOpSharding(const Layout& layout);
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_CC_XLA_SPMD_LAYOUT_TO_XLA_SHARDING_H_
tensorflow/dtensor/mlir/BUILD
+592
View File
@@ -0,0 +1,592 @@
load("@bazel_skylib//rules:build_test.bzl", "build_test")
load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library")
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
# MLIR passes for DTensor support.
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow/dtensor:dtensor-internal",
# Allow visibility from the mlir language server.
"//learning/brain/mlir/mlir_lsp_server:__pkg__",
"//third_party/py/tensorflow_addons/custom_ops/dtensor:__subpackages__",
],
licenses = ["notice"],
)
gentbl_cc_library(
name = "tensorflow_dtensor_ops_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {
"ir/tf_dtensor.h.inc": ["-gen-op-decls"],
"ir/tf_dtensor.cc.inc": ["-gen-op-defs"],
},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "ir/tf_dtensor.td",
td_srcs = [
"//tensorflow/compiler/mlir/tensorflow:ir/tf_op_base.td",
"//tensorflow/compiler/mlir/tensorflow:ir/tf_op_interfaces.td",
],
deps = [
"@llvm-project//mlir:CallInterfacesTdFiles",
"@llvm-project//mlir:FuncTdFiles",
"@llvm-project//mlir:InferTypeOpInterfaceTdFiles",
"@llvm-project//mlir:OpBaseTdFiles",
"@llvm-project//mlir:SideEffectInterfacesTdFiles",
],
)
gentbl_cc_library(
name = "dtensor_passes_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"dtensor_passes.h.inc": [
"-gen-pass-decls",
"-name=DTensor",
]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "Passes.td",
deps = ["@llvm-project//mlir:PassBaseTdFiles"],
)
cc_library(
name = "tf_dtensor_dialect",
srcs = ["ir/tf_dtensor.cc"],
hdrs = ["ir/tf_dtensor.h"],
includes = ["include"],
deps = [
":tensorflow_dtensor_ops_inc_gen",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_attributes",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_op_interfaces",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_side_effects",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_traits",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/dtensor/mlir/dtensor_dialect:ir/dtensor_attributes",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:DerivedAttributeOpInterface",
"@llvm-project//mlir:Dialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:InferTypeOpInterface",
"@llvm-project//mlir:SideEffectInterfaces",
"@llvm-project//mlir:Support",
],
alwayslink = 1,
)
cc_library(
name = "collectives",
srcs = ["collectives.cc"],
hdrs = ["collectives.h"],
deps = [
":collectives_common",
":dtensor_location",
":layout_parsing",
":shape_utils",
":sparse_expander_common",
":spmd_expander_common",
":tf_dtensor_dialect",
":value_utils",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
"//tensorflow/core:lib",
"//tensorflow/dtensor/cc:dstatus",
"//tensorflow/dtensor/cc:dtensor_utils",
"//tensorflow/dtensor/cc:tensor_layout",
"//tensorflow/dtensor/mlir/dtensor_dialect:ir/dtensor_attributes",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "collectives_common",
srcs = ["collectives_common.cc"],
hdrs = ["collectives_common.h"],
deps = [
"//tensorflow/core:portable_gif_internal",
"//tensorflow/core/platform:errors",
"//tensorflow/dtensor/cc:dstatus",
"//tensorflow/dtensor/cc:tensor_layout",
"@com_google_absl//absl/container:flat_hash_set",
],
)
cc_library(
name = "device_utils",
srcs = ["device_utils.cc"],
hdrs = ["device_utils.h"],
deps = [
"//tensorflow/core/platform:errors",
"//tensorflow/dtensor/cc:dstatus",
"//tensorflow/dtensor/cc:tensor_layout",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
alwayslink = True,
)
cc_library(
name = "dtensor_location",
srcs = ["dtensor_location.cc"],
hdrs = ["dtensor_location.h"],
deps = [
"//tensorflow/compiler/mlir/utils:name_utils",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "create_dtensor_mlir_passes",
hdrs = [
"create_dtensor_mlir_passes.h",
],
deps = [
":device_utils",
":dtensor_passes_inc_gen",
":dtensor_send_recv",
":layout_parsing",
":op_utils",
":shape_utils",
":sparse_expander",
":spmd_expander",
":spmd_expander_common",
":tf_dtensor_dialect",
":value_utils",
"//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:lib",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
],
alwayslink = 1,
)
cc_library(
name = "dtensor_mlir_passes",
srcs = [
"annotate_global_shape.cc",
"cluster_function_conversion.cc",
"constant_folding.cc",
"dce.cc",
"designate_resource_handle_mesh.cc",
"device_mesh_cluster_coarsening.cc",
"dtensor_allreduce_combine_optimization.cc",
"dtensor_allreduce_scatter_optimization.cc",
"dtensor_allreduce_sum_optimization.cc",
"dtensor_collective_type_lowering.cc",
"dtensor_layout_to_xla_sharding_op.cc",
"dtensor_mixed_precision_reduce.cc",
"dtensor_mlir_passes.cc",
"dtensor_multi_device_expansion.cc",
"dtensor_remove_dtensorlayout.cc",
"dtensor_replace_auxiliary_layout_op.cc",
"dtensor_replace_relayout_with_identity.cc",
"dtensor_set_hlo_sharding.cc",
"function_renaming.cc",
"handle_cross_cluster_dependencies.cc",
"handle_sparsetensors.cc",
"layout_propagation_v2.cc",
"lower_send_recv.cc",
"merge_clusters.cc",
"mesh_propagation.cc",
"move_compilation_to_host.cc",
"op_to_device_cluster.cc",
"propagate_default_layout.cc",
"propagate_device_id_to_function_args.cc",
"restore_shape_inference.cc",
"set_default_sharding.cc",
"sparse_expansion.cc",
"spmd_expansion.cc",
"tpu_add_resource_device_attribute.cc",
"tpu_integration.cc",
"undo_merge_const_across_mesh.cc",
],
hdrs = ["dtensor_mlir_passes.h"],
deps = [
":collectives_common",
":create_dtensor_mlir_passes",
":device_utils",
":dtensor_location",
":dtensor_passes_inc_gen",
":dtensor_send_recv",
":layout_parsing",
":op_utils",
":shape_utils",
":sparse_expander",
":spmd_expander",
":spmd_expander_common",
":tf_dtensor_dialect",
":topological_iterator",
":value_utils",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:bridge_logger",
"//tensorflow/compiler/mlir/tensorflow:convert_tensor",
"//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_analysis",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_attributes",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow:tpu_rewrite_device_util",
"//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
"//tensorflow/compiler/mlir/tensorflow/transforms/host_runtime:runtime_passes",
"//tensorflow/compiler/mlir/utils:name_utils",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/dtensor/cc:constants",
"//tensorflow/dtensor/cc:dstatus",
"//tensorflow/dtensor/cc:dtensor_utils",
"//tensorflow/dtensor/cc:layout_to_xla_sharding",
"//tensorflow/dtensor/cc:tensor_layout",
"//tensorflow/dtensor/mlir/dtensor_dialect:ir/dtensor_attributes",
"//tensorflow/dtensor/mlir/utils:dtensor_mlir_passes_internal",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/log:vlog_is_on",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/types:optional",
"@com_google_absl//absl/types:span",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:Analysis",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:FunctionInterfaces",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:SideEffectInterfaces",
"@llvm-project//mlir:SparseTensorDialect",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
"@llvm-project//mlir:Transforms",
"@xla//xla:xla_data_proto_cc",
"@xla//xla/hlo/builder:sharding_builder",
"@xla//xla/tsl/platform:status",
],
alwayslink = 1,
)
cc_library(
name = "dtensor_send_recv",
srcs = ["dtensor_send_recv.cc"],
hdrs = ["dtensor_send_recv.h"],
deps = [
":collectives",
":device_utils",
":layout_parsing",
":op_utils",
":shape_utils",
":spmd_expander_common",
":tf_dtensor_dialect",
":value_utils",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:convert_tensor",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_attributes",
"//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
"//tensorflow/core/platform:errors",
"//tensorflow/dtensor/cc:constants",
"//tensorflow/dtensor/cc:dstatus",
"//tensorflow/dtensor/cc:tensor_layout",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/status",
"@com_google_absl//absl/types:span",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
alwayslink = True,
)
cc_library(
name = "layout_parsing",
srcs = [
"layout_parsing.cc",
],
hdrs = ["layout_parsing.h"],
deps = [
":tf_dtensor_dialect",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/core:lib",
"//tensorflow/dtensor/cc:constants",
"//tensorflow/dtensor/cc:dstatus",
"//tensorflow/dtensor/cc:tensor_layout",
"//tensorflow/dtensor/proto:layout_proto_cc",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/types:optional",
"@com_google_absl//absl/types:span",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
alwayslink = 1,
)
cc_library(
name = "op_utils",
srcs = ["op_utils.cc"],
hdrs = ["op_utils.h"],
deps = [
":tf_dtensor_dialect",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/dtensor/cc:constants",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:CallOpInterfaces",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
alwayslink = True,
)
cc_library(
name = "shape_utils",
srcs = ["shape_utils.cc"],
hdrs = ["shape_utils.h"],
deps = [
":tf_dtensor_dialect",
":value_utils",
"//tensorflow/compiler/mlir/tensorflow:shape_inference_utils",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_attributes",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/core:framework",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:status",
"//tensorflow/dtensor/cc:constants",
"//tensorflow/dtensor/cc:dstatus",
"//tensorflow/dtensor/cc:tensor_layout",
"@com_google_absl//absl/status",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:DerivedAttributeOpInterface",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:InferTypeOpInterface",
"@llvm-project//mlir:Support",
],
alwayslink = True,
)
cc_library(
name = "sparse_expander",
srcs = [
"sparse_expanders.cc",
] + glob([
"*sparse_expander.cc",
"sparse_expansions/*sparse_expander.cc",
]),
hdrs = glob([
"*sparse_expander.h",
"sparse_expansions/*sparse_expander.h",
]),
deps = [
":op_utils",
":sparse_expander_common",
":tf_dtensor_dialect",
":value_utils",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops",
"//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
"//tensorflow/core:framework",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:status",
"//tensorflow/dtensor/cc:dstatus",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
alwayslink = 1,
)
cc_library(
name = "sparse_expander_common",
srcs = ["sparse_expander_common.cc"],
hdrs = ["sparse_expander_common.h"],
deps = [
":tf_dtensor_dialect",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops",
"//tensorflow/core/platform:errors",
"//tensorflow/dtensor/cc:dstatus",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/types:optional",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:BytecodeOpInterface",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
],
alwayslink = True,
)
cc_library(
name = "spmd_expander_common",
srcs = ["spmd_expander_common.cc"],
hdrs = ["spmd_expander_common.h"],
deps = [
":device_utils",
":layout_parsing",
":op_utils",
":shape_utils",
":tf_dtensor_dialect",
":value_utils",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/core:lib",
"//tensorflow/dtensor/cc:constants",
"//tensorflow/dtensor/cc:dstatus",
"//tensorflow/dtensor/cc:tensor_layout",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:Dialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
alwayslink = True,
)
cc_library(
name = "spmd_expander",
srcs = [
"spmd_expanders.cc",
] + glob([
"*spmd_expander.cc",
"expansions/*spmd_expander.cc",
"expansions/*spmd_expander_v2.cc",
]),
hdrs = glob([
"*spmd_expander.h",
"expansions/*spmd_expander.h",
"expansions/*spmd_expander_v2.h",
]),
deps = [
":collectives",
":device_utils",
":dtensor_location",
":dtensor_send_recv",
":layout_parsing",
":op_utils",
":shape_utils",
":spmd_expander_common",
":tf_dtensor_dialect",
":value_utils",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_attributes",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/dtensor/cc:constants",
"//tensorflow/dtensor/cc:dstatus",
"//tensorflow/dtensor/cc:dtensor_utils",
"//tensorflow/dtensor/cc:save_restore_util",
"//tensorflow/dtensor/cc:slice_util",
"//tensorflow/dtensor/cc:tensor_layout",
"//tensorflow/dtensor/mlir/dtensor_dialect:ir/dtensor_attributes",
"//tensorflow/dtensor/proto:layout_proto_cc",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/numeric:bits",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:optional",
"@com_google_absl//absl/types:span",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:BytecodeOpInterface",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
"@xla//xla/mlir_hlo:convert_op_folder",
],
alwayslink = 1,
)
cc_library(
name = "value_utils",
srcs = ["value_utils.cc"],
hdrs = ["value_utils.h"],
deps = [
":op_utils",
":tf_dtensor_dialect",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:dynamic_shape_utils",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
"//tensorflow/core:lib",
"//tensorflow/dtensor/cc:dstatus",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
"@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
],
alwayslink = True,
)
cc_library(
name = "topological_iterator",
srcs = ["topological_iterator.cc"],
hdrs = ["topological_iterator.h"],
deps = [
":device_utils",
":layout_parsing",
":op_utils",
":shape_utils",
":tf_dtensor_dialect",
":value_utils",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
],
)
tf_cc_test(
name = "dtensor_location_test",
srcs = ["dtensor_location_test.cc"],
deps = [
":dtensor_location",
"//tensorflow/compiler/mlir/utils:name_utils",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
)
build_test(
name = "mlir_build_test",
targets = [
":tf_dtensor_dialect",
":tensorflow_dtensor_ops_inc_gen",
],
)
tensorflow/dtensor/mlir/Passes.td
+507
View File
@@ -0,0 +1,507 @@
/* Copyright 2022 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 TF_DTENSOR_PASSES
#define TF_DTENSOR_PASSES
include "mlir/Pass/PassBase.td"
def DTensorOpToDeviceCluster
: Pass<"dtensor-op-to-device-cluster", "mlir::func::FuncOp"> {
let summary = "Creates and wraps tf_device.cluster op for all TF ops";
let constructor = "CreateDTensorOpToDeviceClusterPass()";
let dependentDialects = [
];
}
def DTensorDeviceMeshClusterCoarsening
: Pass<"dtensor-device-mesh-cluster-coarsening", "mlir::func::FuncOp"> {
let summary = "Merges tf_device.cluster op with same mesh attribute.";
let constructor = "CreateDTensorDeviceMeshClusterCoarsening()";
let dependentDialects = [
];
}
def DTensorConstantFolding
: Pass<"dtensor-constant-folding", "mlir::func::FuncOp"> {
let summary = "Folds constants operations.";
let constructor = "CreateDTensorConstantFolding()";
let dependentDialects = [
];
}
def DTensorCollectiveTypeLoweringPass
: Pass<"dtensor-collective-type-lowering", "mlir::func::FuncOp"> {
let summary = "Lowers collectives from unsupported types (e.g. complex) to supported types.";
let constructor = "CreateDTensorCollectiveTypeLoweringPass()";
let dependentDialects = [
];
}
def DTensorAllReduceSumOptimization
: Pass<"dtensor-allreduce-sum-optimization", "mlir::func::FuncOp"> {
let summary = "Changes order of add/allreduce to minimize all reduce operations.";
let constructor = "CreateDTensorAllReduceSumOptimization()";
let dependentDialects = [
];
}
def DTensorAllReduceScatterOptimization
: Pass<"dtensor-allreduce-scatter-optimization", "mlir::func::FuncOp"> {
let summary = "Combines allreduce and scatter to reducescatter.";
let constructor = "CreateDTensorAllReduceScatterOptimization()";
let dependentDialects = [
];
}
def DTensorAllReduceCombineOptimization
: Pass<"dtensor-allreduce-combine-optimization", "mlir::func::FuncOp"> {
let summary = "Combine independent all reduce operations.";
let constructor = "CreateDTensorAllReduceCombineOptimization()";
let dependentDialects = [
];
}
def DTensorMixedPrecisionReduce
: Pass<"dtensor-mixed-precision-reduce", "mlir::func::FuncOp"> {
let summary = "Upcast tensors to higher precision type for reduction ops.";
let constructor = "CreateDTensorMixedPrecisionReducePass()";
let dependentDialects = [
];
}
def DTensorDCE
: Pass<"dtensor-dce", "mlir::func::FuncOp"> {
let summary = "Removes unused ops from graph.";
let constructor = "CreateDTensorDCE()";
let dependentDialects = [
];
}
def DTensorUndoMergeConstAcrossMesh
: Pass<"dtensor-undo-merge-const-across-mesh", "mlir::func::FuncOp"> {
let summary = "Undo constant merging across meshes";
let constructor = "CreateDTensorUndoMergeConstAcrossMesh()";
let dependentDialects = [
];
}
def DTensorElideIdentityBeforeCopyToMesh
: Pass<"dtensor-elide-identity-before-copy-to-mesh", "mlir::func::FuncOp"> {
let summary = "Elide IdentityOp right before CopyToMesh style Ops";
let constructor = "CreateDTensorElideIdentityBeforeCopyToMesh()";
let dependentDialects = [
];
}
def DTensorSetDefaultSharding
: Pass<"dtensor-set-default-sharding", "mlir::func::FuncOp"> {
let summary = "Sets default sharding of TPU computation inputs/outputs to maximal.";
let constructor = "CreateDTensorSetDefaultSharding()";
let dependentDialects = [
];
}
def DTensorDesignateResourceHandleMesh
: Pass<"dtensor-designate-resource-handle-mesh", "mlir::func::FuncOp"> {
let summary = "Sets empty mesh attributes for device cluster that creates or destroys resource handles.";
let constructor = "CreateDTensorDesignateResourceHandleMesh()";
let dependentDialects = [
];
}
def DTensorPropagateDefaultLayout
: Pass<"dtensor-propagate-default-layout", "mlir::func::FuncOp"> {
let summary = "Converts layout attributes added by end users to DTensorLayout op.";
let constructor = "CreateDTensorPropagateDefaultLayout()";
let dependentDialects = [
];
}
def DTensorHandleCrossClusterDependencies
: Pass<"dtensor-handle_cross_cluster_dependences", "mlir::ModuleOp"> {
let summary = "Lowers cross mesh cluster data depedences as send/recvs.";
let constructor = "CreateDTensorHandleCrossClusterDependencies()";
let dependentDialects = [
];
}
def DTensorAnnotateGlobalShape
: Pass<"dtensor-annotate-global-shape", "mlir::ModuleOp"> {
let summary = "Mark all operations and function arguments with `_global_shape` attribute to be used during SPMD expansion.";
let constructor = "CreateDTensorAnnotateGlobalShape()";
let dependentDialects = [
];
}
def DTensorLayoutPropagationV2
: Pass<"dtensor-layout-propagation-v2", "mlir::ModuleOp"> {
let summary = "Propagates layout information for all the TF Ops.";
let constructor = "CreateDTensorLayoutPropagationPassV2()";
let dependentDialects = [
];
}
def DTensorMeshPropagation
: Pass<"dtensor-mesh-propagation", "mlir::ModuleOp"> {
let summary = "Propagates mesh information to all clusters.";
let constructor = "CreateDTensorMeshPropagationPass()";
let dependentDialects = [
];
}
def DTensorSPMDExpansion
: Pass<"dtensor-spmd-expansion", "mlir::ModuleOp"> {
let summary = "Converts ops into SPMD expanded form.";
let constructor = "CreateDTensorSPMDExpansion()";
let dependentDialects = [
];
}
def DTensorAllReduceLowering
: Pass<"dtensor-all-reduce-lowering", "mlir::ModuleOp"> {
let summary = "Converts logical AllReduce ops into physical AllReduce ops.";
let constructor = "CreateDTensorAllReduceLoweringPass()";
let dependentDialects = [
"::mlir::dtensor::DTensorDialect"
];
}
def DTensorReduceScatterLowering
: Pass<"dtensor-reduce-scatter-lowering", "mlir::ModuleOp"> {
let summary = "Converts logical ReduceScatter ops into physical ReduceScatter ops.";
let constructor = "CreateDTensorReduceScatterLoweringPass()";
let dependentDialects = [
"::mlir::dtensor::DTensorDialect"
];
}
def DTensorAllGatherLowering
: Pass<"dtensor-all-gather-lowering", "mlir::ModuleOp"> {
let summary = "Converts logical AllGather ops into physical AllGather ops.";
let constructor = "CreateDTensorAllGatherLoweringPass()";
let dependentDialects = [
"::mlir::dtensor::DTensorDialect"
];
}
def DTensorAllScatterLowering
: Pass<"dtensor-all-scatter-lowering", "mlir::ModuleOp"> {
let summary = "Converts logical AllScatter ops into physical Split ops.";
let constructor = "CreateDTensorAllScatterLoweringPass()";
let dependentDialects = [
"::mlir::dtensor::DTensorDialect"
];
}
def DTensorAllToAllLowering
: Pass<"dtensor-all-to-all-lowering", "mlir::ModuleOp"> {
let summary = "Converts logical AllToAll ops into physical AllToAll ops.";
let constructor = "CreateDTensorAllToAllLoweringPass()";
let dependentDialects = [
"::mlir::dtensor::DTensorDialect"
];
}
def DTensorClusterFunctionConversion
: Pass<"dtensor-cluster-function-conversion", "mlir::ModuleOp"> {
let summary = "Converts tf_device.cluster_func ops into TF StatefulPartitioned call op with mesh attribute.";
let constructor = "CreateDTensorClusterFunctionConversion()";
let dependentDialects = [
];
}
def DTensorPropagateDeviceIdToFunctionArgs
: Pass<"dtensor-propagate-device-id-to-function-args", "mlir::ModuleOp"> {
let summary = "Adds device id as arguments to all private function in graph.";
let constructor = "CreateDTensorPropagateDeviceIdToFunctionArgs()";
let dependentDialects = [
];
}
def DTensorTPUIntegration
: Pass<"dtensor-tpu-integration", "mlir::ModuleOp"> {
let summary = "Adds TPUReplicateMetadata and converts ops that run on TPU's to a single tf_device.cluster to be compatible with following TF2XLA MLIR passes.";
let constructor = "CreateDTensorTPUIntegration()";
let dependentDialects = [
];
}
def DTensorTpuAddResourceDeviceAttribute
: Pass<"dtensor-tpu-add-resource-device-attribute", "mlir::ModuleOp"> {
let summary = "Adds placeholder device attributes to resources accessed by TPU computation to enable buffer aliasing.";
let constructor = "CreateDTensorTpuAddResourceDeviceAttribute()";
let dependentDialects = [
];
}
def DTensorUpdateTPUMetadata
: Pass<"dtensor-update-tpu-metadata", "mlir::ModuleOp"> {
let summary = "Changes metadata on TPU specific ops such as device placement.";
let constructor = "CreateDTensorUpdateTPUMetadata()";
let dependentDialects = [
];
}
def DTensorFunctionRenaming
: Pass<"dtensor-function-renaming", "mlir::ModuleOp"> {
let summary = "Renames private functions by appending an id to each name. This is used to make private function names unique across modules.";
let constructor = "CreateFunctionRenamingPass()";
let dependentDialects = [
];
}
def DTensorMultiDeviceExpansion
: Pass<"dtensor-multi-device-expansion", "mlir::ModuleOp"> {
let summary = "Expands a per-device, post-SPMD graph for multiple devices.";
let constructor = "CreateDTensorMultiDeviceExpansionPass()";
let dependentDialects = [
];
}
def DTensorMergeClusters
: Pass<"dtensor-merge-clusters", "mlir::ModuleOp"> {
let summary = "Merges tf_device.Clusters ops with same mesh specification.";
let constructor = "CreateDTensorMergeClustersPass()";
let dependentDialects = [
];
}
def DTensorDecomposeControlflow
: Pass<"dtensor-decompose-controlflow", "mlir::ModuleOp"> {
let summary = "Decompose control flow ops to different meshes";
let constructor = "CreateDTensorDecomposeControlflowPass()";
let dependentDialects = [
];
}
def DTensorLowerSendRecv
: Pass<"dtensor-lower-send-recv", "mlir::ModuleOp"> {
let summary = "Lowers DTensorSend/DTensorRecv ops to send/recv ops.";
let constructor = "CreateDTensorLowerSendRecv()";
let dependentDialects = [
];
}
def DTensorMoveCompilationToHost
: Pass<"dtensor-move-compilation-to-host", "mlir::ModuleOp"> {
let summary = "Moves XLA compilation ops to host computation.";
let constructor = "CreateDTensorMoveCompilationToHost()";
let dependentDialects = [
];
}
def DTensorSparseTensorToDenseTensor
: Pass<"dtensor-sparse-tensor-to-dense-tensor", "mlir::ModuleOp"> {
let summary = "Converts SparseTensor inputs to its component tensors inputs and emits a SparseToDenseOp for every op that consumes a SparseTensor.";
let constructor = "CreateDTensorSparseTensorToDenseTensor()";
let dependentDialects = [
];
}
def DTensorSparseExpansion
: Pass<"dtensor-sparse-expansion", "mlir::ModuleOp"> {
let summary = "Convert ops that take in SparseTensor input to its corresponding Sparse or Dense ops.";
let constructor = "CreateDTensorSparseExpansion()";
let dependentDialects = [
];
}
def DTensorInferShapesForRestoreV2Op
: Pass<"dtensor-infer-shapes-for-restorev2-op", "mlir::ModuleOp"> {
let summary = "Infer shapes of the outputs of tf.RestoreV2Op from the AssignVariableOps that consume those outputs. This is used for DTensor integration with TF Checkpoint.";
let constructor = "CreateDTensorInferShapesForRestoreV2Op()";
let dependentDialects = [
];
}
def DTensorSetHloShardingPass : Pass<"dtensor-set-hlo-sharding", "mlir::ModuleOp"> {
let summary = "Set `mhlo.sharding` attribute for function inputs and ops.";
let description = [{
For each `tf.DTensorLayout` op, the pass sets `mhlo.sharding` attributes for
related Ops and the related outer layer function's inputs and outputs.
By default the operation is applied on every DTensorLayout op.
If `check_layout_use_xla_spmd`is set to true, then the pass checks every
DTensorLayout must have Mesh config with use_xla_spmd.
For example, by default the follow code
```
func.func @main(%arg0) -> tensor<8x8xi32> {
%1 = "tf.DTensorLayout"(%arg0)
%2 = "tf.Identity"(%1)
%3 = "tf.DTensorLayout"(%2)
return %3 : tensor<8x8xi32>
}
```
will be converted to
```
func.func @main(%arg0 {mhlo.sharding = ""}) -> (tensor<8x8xi32> {mhlo.sharding = ""})
%1 = "tf.DTensorLayout"(%arg0)
%2 = "tf.Identity"(%1) {mhlo.sharding = ""}
%3 = "tf.DTensorLayout"(%2)
return %3 : tensor<8x8xi32>
}
```
When `check_layout_use_xla_spmd` is set to true, the pass throws an exception for the above example.
Meanwhile, if `check_layout_use_xla_spmd` is set to true, the follow code
```
func.func @main(%arg0) -> tensor<8x8xi32> {
%1 = "tf.DTensorLayout"(%0) {mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0|use_xla_spmd>}
%2 = "tf.Identity"(%1)
%3 = "tf.DTensorLayout"(%2){mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0|use_xla_spmd>}
return %3 : tensor<8x8xi32>
}
```
will be converted to
```
func.func @main(%arg0 {mhlo.sharding = ""}) -> (tensor<8x8xi32> {mhlo.sharding = ""})
%1 = "tf.DTensorLayout"(%0) {mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0|use_xla_spmd>}
%2 = "tf.Identity"(%1)
%3 = "tf.DTensorLayout"(%2){mesh:|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:0|use_xla_spmd>}
return %3 : tensor<8x8xi32>
}
```
}];
let constructor = "CreateDTensorSetHloShardingPass()";
let dependentDialects = [
"::mlir::dtensor::DTensorDialect"
];
let options = [
Option<"check_layout_use_xla_spmd_", "check_layout_use_xla_spmd", "bool",
/*default=*/"false",
"If true, the pass checks every DTensorLayout must have Mesh config with use_xla_spmd."
"Otherwise, the check is disabled.">
];
}
def DTensorLayoutToXlaShardingOpPass : Pass<"dtensor-layout-to-xla-sharding-op", "mlir::func::FuncOp"> {
let summary = "Replace `tf.DTensorLayout` op with `tf.XlaSharding` op.";
let description = [{
Provides sharding guide to XLA based on DTensor layout propagation result.
tf2xla bridge will further lower `tf.XlaSharding` to `hlo.custom_call`.
For example:
```
%1 = tf.DTensorLayout(%0)
```
will be converted to
```
%1 = tf.XlaSharding(%0)
```
When the DTensorLayout's operand is produced by a constant, the
DTensorLayout will be removed and no XlaSharding is inserted. This is
because tf2xla requires some op operands to be constant. Inserting
XlaSharding will break them.
A side note: "mhlo.sharding" attributes on ops except FuncOp won't be used
by XLA.
}];
let constructor = "CreateDTensorLayoutToXlaShardingOpPass()";
}
def DTensorReplaceAuxiliaryDTensorLayoutOpPass
: Pass<"dtensor-replace-auxiliary-layout-op", "mlir::ModuleOp"> {
let summary = "Replace auxiliary `tf.DTensorLayout` op with `tf.Identity`.";
let description = [{
Canonicalizer and DCE transformation passes may remove ops in the graph and
result in multiple consecutive DTensorLayout ops. The pass detects all such
cases and replaces unnecessary DTensorLayout ops with Identity ops.
Removes `tf.DTensorLayouts` and inserts a `tf.Identity`.
For example:
```
%0 = tf.DTensorLayout(arg0)
%1 = tf.DTensorLayout(%0)
%2 = tf.Add(%1, %1)
```
will be converted to:
```
%0 = tf.Identity(arg0)
%1 = tf.DTensorLayout(%0)
%2 = tf.Add(%1, %1)
```
}];
let constructor = "CreateDTensorReplaceAuxiliaryDTensorLayoutOpPass()";
}
def DTensorRemoveDTensorLayoutPass : Pass<"dtensor-remove-dtensorlayout", "mlir::ModuleOp"> {
let summary = "Remove DTensor `tf.DTensorLayout` ops.";
let description = [{
The pass removes DTensor `tf.DTensorLayout` ops.
For example,
```mlir
%1 = tf.DTensorLayout(%0)
%2 = tf.SomeOp(%1)
```
will be converted to
```mlir
%2 = tf.SomeOp(%0)
```
}];
let constructor = "CreateDTensorRemoveDTensorLayoutPass()";
let dependentDialects = [
"::mlir::dtensor::DTensorDialect"
];
}
def DTensorReplaceRelayoutWithIdentityPass : Pass<"dtensor-replace-relayout-with-identity", "mlir::func::FuncOp"> {
let summary = "Replace `tf.Relayout` op with `tf.Identity` op.";
let description = [{
This pass replaces `tf.Relayout` op with `tf.Identity` op.
It replaces all usages with the inputs. For example, the following code
```mlir
%1 = tf.Relayout(%0)
%2 = tf.SomeOp(%1)
```
will be replaced by
```mlir
%1 = tf.Identity(%0)
%2 = tf.SomeOp(%1)
```
}];
let constructor = "CreateDTensorReplaceRelayoutWithIdentityPass()";
}
#endif // TF_DTENSOR_PASSES
tensorflow/dtensor/mlir/annotate_global_shape.cc
+135
View File
@@ -0,0 +1,135 @@
/* Copyright 2022 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/ArrayRef.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Diagnostics.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/TypeUtilities.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/convert_tensor.h"
#include "tensorflow/dtensor/cc/constants.h"
#include "tensorflow/dtensor/mlir/ir/tf_dtensor.h"
#include "tensorflow/dtensor/mlir/value_utils.h"
namespace tensorflow {
namespace dtensor {
namespace {
#define GEN_PASS_DEF_DTENSORANNOTATEGLOBALSHAPE
#include "tensorflow/dtensor/mlir/dtensor_passes.h.inc"
// Sets `_global_shape` attributes to argument/return values of `function`.
void AnnotateFunctionArgRetvalGlobalShapes(mlir::func::FuncOp function,
mlir::OpBuilder* builder) {
for (const auto& argument_type_and_index :
llvm::enumerate(function.getArgumentTypes())) {
const int index = argument_type_and_index.index();
const auto& argument_type = argument_type_and_index.value();
// Extract TensorType from element of resource type to allow setting proper
// global shape of resource types.
if (auto resource_type = mlir::dyn_cast<mlir::TF::ResourceType>(
mlir::getElementTypeOrSelf(argument_type))) {
auto subtype = resource_type.getSubtypes();
if (subtype.size() == 1) {
// subtype returns a Array of TensorType -- if it contains more than one
// Tensor type, we give up extracting the single TensorType inside the
// subtype.
function.setArgAttr(index, kGlobalShapeDialectAttr,
ConvertTypeToTensorShapeAttr(subtype[0]));
}
} else {
function.setArgAttr(index, kGlobalShapeDialectAttr,
ConvertTypeToTensorShapeAttr(argument_type));
}
}
for (const auto& retval_type_and_index :
llvm::enumerate(function.getFunctionType().getResults())) {
const int index = retval_type_and_index.index();
const auto& retval_type = retval_type_and_index.value();
function.setResultAttr(index, kGlobalShapeDialectAttr,
ConvertTypeToTensorShapeAttr(retval_type));
}
}
// Sets `_global_shape` attribute of an `op` with array of ShapeAttr of
// `outputs.
void AnnotateOperationGlobalShape(mlir::Operation* op,
mlir::OpBuilder* builder) {
llvm::SmallVector<mlir::Attribute, 4> op_global_shape;
op_global_shape.reserve(op->getNumResults());
for (const auto& result_type : op->getResultTypes())
op_global_shape.emplace_back(ConvertTypeToTensorShapeAttr(result_type));
if (auto layout_op = mlir::dyn_cast<mlir::TF::DTensorLayout>(op)) {
// Shape of Resource type is incorrect when it is a variable.
// The global shape is undefined in this case; and usually we are supposed
// to propagate the value shape due to how resource variable layout is
// currently represented in DTensor.
if (!IsResourceType(op->getResult(0))) {
layout_op.setGlobalShapeAttr(op_global_shape[0]);
}
} else {
op->setAttr(kGlobalShape, builder->getArrayAttr(op_global_shape));
}
}
// Pass that annotates function argument/return values and all operation with
// `_global_shape` attribute. This will be used during SPMD expansion to
// preserve original global shape of operations in graph after shape has been
// modified to local shape.
struct DTensorAnnotateGlobalShape
: public impl::DTensorAnnotateGlobalShapeBase<DTensorAnnotateGlobalShape> {
void runOnOperation() override {
mlir::MLIRContext& context = getContext();
mlir::OpBuilder builder(&context);
auto module = getOperation();
module.walk([&](mlir::func::FuncOp function) {
if (function.empty()) return;
auto* terminator = function.getBody().front().getTerminator();
AnnotateFunctionArgRetvalGlobalShapes(function, &builder);
function.getBody().walk([&](mlir::Operation* op) {
if (op == terminator) return;
AnnotateOperationGlobalShape(op, &builder);
});
});
}
};
} // namespace
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorAnnotateGlobalShape() {
return std::make_unique<DTensorAnnotateGlobalShape>();
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/cluster_function_conversion.cc
+203
View File
@@ -0,0 +1,203 @@
/* Copyright 2022 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 <iterator>
#include <memory>
#include <optional>
#include <utility>
#include "absl/types/optional.h"
#include "absl/types/span.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/FormatVariadic.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Diagnostics.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/SymbolTable.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h"
#include "tensorflow/dtensor/cc/constants.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/layout_parsing.h"
#include "tensorflow/dtensor/mlir/op_utils.h"
#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
namespace tensorflow {
namespace dtensor {
namespace {
#define GEN_PASS_DEF_DTENSORCLUSTERFUNCTIONCONVERSION
#include "tensorflow/dtensor/mlir/dtensor_passes.h.inc"
// Attach layouts for all the returned values so that custom device could get
// layouts for the handles.
mlir::LogicalResult AttachRetvalLayouts(
mlir::OpBuilder* builder, mlir::TF::StatefulPartitionedCallOp sp_call_op) {
// Find the FuncOp that the StatefulPartitionedCallOp is invoking.
mlir::SymbolRefAttr sym =
sp_call_op.getCallableForCallee().dyn_cast<mlir::SymbolRefAttr>();
if (!sym)
return sp_call_op.emitOpError(
"has no symbolRef for given StatefulPartitionedCallOp");
auto func = mlir::dyn_cast<mlir::func::FuncOp>(
mlir::SymbolTable::lookupNearestSymbolFrom(sp_call_op, sym));
if (!func)
return sp_call_op.emitOpError() << "found no FuncOp for symbol " << sym;
llvm::SmallVector<std::optional<Layout>, 8> retvals_layouts;
retvals_layouts.reserve(func.getNumResults());
for (auto operand : func.front().getTerminator()->getOperands()) {
auto result_layout_or_status = ExtractLayoutFromOperand(operand);
if (!result_layout_or_status.ok()) {
return func.emitOpError("error while parsing result layout for function");
}
auto result_layout = result_layout_or_status.value();
// When function returns its arguments directly, layout information for the
// return value of `func` may be only obtainable by looking at it's callsite
// operations. In that case, query the input layouts for function callsite
// operations for layout information.
if (!result_layout) {
if (auto block_arg = mlir::dyn_cast<mlir::BlockArgument>(operand)) {
auto layout_or_status = ExtractLayoutFromOperand(
sp_call_op.getOperand(block_arg.getArgNumber()));
if (!layout_or_status.ok())
return func.emitOpError(
"error while parsing result layout for function");
result_layout = std::move(layout_or_status.value());
}
if (!result_layout)
return func.emitOpError(
llvm::formatv("missing result layout attribute for function. All "
"DTensor functions "
"must have layouts for its results."));
}
retvals_layouts.emplace_back(result_layout.value());
}
// Note that we set this unconditionally - retvals_layout could be empty, but
// that is fine and we will have an empty _layout for the
// StatefulPartitionedCallOp. This is fine as for op without return values,
// all we need is a placeholder layout so that no special case is needed in
// dtensor_device.
SetLayoutOnOp(sp_call_op,
absl::Span<const absl::optional<Layout>>(
retvals_layouts.data(), retvals_layouts.size()));
return mlir::success();
}
// Add an anotation to skip xla compilation for VarHandleOp and
// DestroyResourceOp.
void MaybeSkipXlaCompilation(mlir::OpBuilder* builder,
mlir::Operation* call_op) {
auto function = MaybeFindFunction(call_op);
const auto& body_ops = function->getBody().front().without_terminator();
// VarHandleOp and DestroyResourceOp run on op-by-op mode, so there is only
// one op in the function body.
if (std::distance(std::begin(body_ops), std::end(body_ops)) == 1 &&
llvm::isa<mlir::TF::VarHandleOp, mlir::TF::DestroyResourceOp>(
body_ops.begin())) {
call_op->setAttr(kSkipXlaCompilation, builder->getBoolAttr(true));
}
}
mlir::LogicalResult ReplaceClusterWithPartitionCallOp(
mlir::OpBuilder* builder, mlir::tf_device::ClusterFuncOp cluster_func) {
auto mesh_attr = cluster_func->getAttrOfType<mlir::StringAttr>(kMeshAttr);
if (!mesh_attr)
return cluster_func.emitOpError()
<< "requires " << llvm::StringRef(kMeshAttr) << " attribute";
llvm::SmallVector<mlir::Type, 8> output_types{
cluster_func.getResultTypes().begin(),
cluster_func.getResultTypes().end()};
llvm::StringRef function_name = cluster_func.getFunc();
builder->setInsertionPoint(cluster_func);
auto call_op = mlir::TF::StatefulPartitionedCallOp::create(
*builder, cluster_func.getLoc(), output_types, cluster_func.getOperands(),
/*args_attrs=*/nullptr, /*res_attrs=*/nullptr, function_name, mesh_attr,
/*config_proto=*/builder->getStringAttr(""),
/*executor_type=*/builder->getStringAttr(""));
MaybeSkipXlaCompilation(builder, call_op);
if (mlir::failed(ValidateMetadataAttributes(cluster_func)))
return mlir::failure();
// All attributes beginning with `_` is validate, perform copy.
mlir::TF::CopyUnderscoredAttributes(cluster_func, call_op);
cluster_func.replaceAllUsesWith(call_op.getResults());
cluster_func.erase();
return AttachRetvalLayouts(builder, call_op);
}
// MLIR pass that converts tf_device.cluster_func to TF partitioned call
// op with device mesh config added to `config` attribute.
struct DTensorClusterFunctionConversion
: public impl::DTensorClusterFunctionConversionBase<
DTensorClusterFunctionConversion> {
void runOnOperation() override {
mlir::MLIRContext& context = getContext();
// Find all tf_device.ClusterFunc ops and visit them in post order. This
// order guarantees that ops in function definition is visited before
// function call site operations. When python graph includes tf.functions
// this leads to nested tf_device.ClusterFunc ops. As we infer the layout
// of function call operations with layout attached to return values in the
// function definition, ClusterFunc op in nested/inner functions must be
// visited before ClusterFunc op in outer functions.
llvm::SmallVector<mlir::tf_device::ClusterFuncOp, 8> clusters;
getOperation().walk([&](mlir::tf_device::ClusterFuncOp cluster_func) {
clusters.emplace_back(cluster_func);
});
mlir::OpBuilder op_builder(&context);
for (auto cluster_func : llvm::reverse(clusters)) {
if (mlir::failed(
ReplaceClusterWithPartitionCallOp(&op_builder, cluster_func))) {
return signalPassFailure();
}
}
};
};
} // namespace
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorClusterFunctionConversion() {
return std::make_unique<DTensorClusterFunctionConversion>();
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/collectives.cc
+753
View File
@@ -0,0 +1,753 @@
/* Copyright 2022 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/dtensor/mlir/collectives.h"
#include <cstdint>
#include <iterator>
#include <string>
#include <vector>
#include "absl/container/flat_hash_set.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/FormatVariadic.h"
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinTypeInterfaces.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/IR/ValueRange.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/collection_ops_util.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/dtensor_utils.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/collectives_common.h"
#include "tensorflow/dtensor/mlir/dtensor_dialect/ir/dtensor_attributes.h"
#include "tensorflow/dtensor/mlir/dtensor_location.h"
#include "tensorflow/dtensor/mlir/ir/tf_dtensor.h"
#include "tensorflow/dtensor/mlir/layout_parsing.h"
#include "tensorflow/dtensor/mlir/shape_utils.h"
#include "tensorflow/dtensor/mlir/sparse_expander_common.h"
#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
#include "tensorflow/dtensor/mlir/value_utils.h"
namespace tensorflow {
namespace dtensor {
namespace {
namespace ops_util = ::mlir::TF::collection_ops_util;
} // namespace
StatusOr<mlir::Value> EmitAllGather(
mlir::OpBuilder& builder, mlir::Value input,
const dtensor::Layout& src_layout, const dtensor::Layout& tgt_layout,
llvm::SmallPtrSet<mlir::Operation*, 4>* newly_created_ops) {
if (src_layout.IsEquivalent(tgt_layout)) return input;
if (src_layout.rank() != tgt_layout.rank()) {
return absl::InvalidArgumentError(absl::StrCat(
"Expected source and target layout to have the same rank, got ",
src_layout.rank(), " vs ", tgt_layout.rank()));
}
// Check that the tgt_layout is less sharded then src_layout.
for (int i = 0; i < src_layout.rank(); ++i) {
if (src_layout.sharding_spec(i) != tgt_layout.sharding_spec(i) &&
Layout::IsShardedDimension(tgt_layout.sharding_spec(i))) {
return absl::InvalidArgumentError(
absl::StrCat("source layout (", src_layout.ToString(),
") for all gather is not less sharded "
"than the target layout (",
tgt_layout.ToString()));
}
}
// For convenience, operate on explicit input shapes. This isn't necessary,
// as we could instead generate operations on top of the dynamic shape.
const mlir::TensorType input_type =
mlir::dyn_cast<mlir::TensorType>(input.getType());
if (!input_type) {
return absl::InternalError(
llvm::formatv(
"Cannot cast input_type : {0} to TensorType. Shape must be "
" statically known before emitting AllGather. This should not "
"happen as we already cast it when getting its shape.",
input.getType())
.str());
}
TF_ASSIGN_OR_RETURN(mlir::TensorType global_type,
GlobalTypeFromLocalType(src_layout, input_type));
TF_ASSIGN_OR_RETURN(mlir::TensorType output_type,
LocalTypeFromGlobalType(tgt_layout, global_type));
mlir::Location loc = DT_LOC2(input.getLoc(), "DTensorAllGatherOp");
mlir::TF::DTensorAllGatherOp all_gather =
mlir::TF::DTensorAllGatherOp::create(
builder, loc, output_type, input,
mlir::dtensor::LayoutAttr::get(builder.getContext(), src_layout),
mlir::dtensor::LayoutAttr::get(builder.getContext(), tgt_layout));
SetSingleLayoutOnOp(all_gather, tgt_layout);
if (newly_created_ops != nullptr) newly_created_ops->insert(all_gather);
return all_gather.getOutput();
}
StatusOr<const mlir::Value> EmitAllScatter(
mlir::OpBuilder& builder, const mlir::Value& original_value,
const Layout& original_layout, const Layout& desired_layout,
llvm::SmallPtrSet<mlir::Operation*, 4>* newly_created_ops) {
if (original_layout.IsEquivalent(desired_layout)) return original_value;
// Have an early return if desired layout is not more sharded then the
// original_layout.
if (original_layout.rank() != desired_layout.rank()) {
return absl::InvalidArgumentError(absl::StrCat(
"Rank mismatch for original layout (", original_layout.ToString(),
") and desired layout (", desired_layout.ToString(), ")"));
}
for (int i = 0; i < original_layout.rank(); ++i) {
if (original_layout.sharding_spec(i) != desired_layout.sharding_spec(i) &&
Layout::IsShardedDimension(original_layout.sharding_spec(i))) {
return absl::InvalidArgumentError(
absl::StrCat("EmitAllScatter was passed a desired_layout ",
desired_layout.ToString(),
" which was not more sharded than the original_layout ",
original_layout.ToString()));
}
}
const mlir::TensorType input_type =
mlir::dyn_cast<mlir::TensorType>(original_value.getType());
if (!input_type)
return absl::InvalidArgumentError(
"input to EmitAllScatter does not have a TensorType");
TF_ASSIGN_OR_RETURN(const mlir::TensorType global_type,
GlobalTypeFromLocalType(original_layout, input_type));
TF_ASSIGN_OR_RETURN(const mlir::TensorType output_type,
LocalTypeFromGlobalType(desired_layout, global_type));
mlir::Location loc = DT_LOC2(original_value.getLoc(), "DTensorAllScatterOp");
mlir::TF::DTensorAllScatterOp all_scatter =
mlir::TF::DTensorAllScatterOp::create(
builder, loc, output_type, original_value,
mlir::dtensor::LayoutAttr::get(builder.getContext(), original_layout),
mlir::dtensor::LayoutAttr::get(builder.getContext(), desired_layout));
SetSingleLayoutOnOp(all_scatter, desired_layout);
if (newly_created_ops != nullptr) newly_created_ops->insert(all_scatter);
return all_scatter.getOutput();
}
bool CanUseAllToAll(const dtensor::Layout& src_layout,
const dtensor::Layout& tgt_layout) {
// All-to-all can be used for relayout if one dimension is becoming more
// sharded while another is becoming less sharded, for example x,unsharded ->
// unsharded,x.
// TODO(trevor-m): There may be more types of relayouts which can utilize
// all-to-all in addition to these which can be supported later.
int num_split_dims = 0;
int num_concat_dims = 0;
std::string split_spec;
std::string concat_spec;
for (int i = 0; i < src_layout.rank(); ++i) {
if (src_layout.sharding_spec(i) == tgt_layout.sharding_spec(i)) continue;
if (Layout::IsUnshardedDimension(src_layout.sharding_spec(i)) &&
Layout::IsShardedDimension(tgt_layout.sharding_spec(i))) {
num_split_dims++;
split_spec = tgt_layout.sharding_spec(i);
} else if (Layout::IsShardedDimension(src_layout.sharding_spec(i)) &&
Layout::IsUnshardedDimension(tgt_layout.sharding_spec(i))) {
num_concat_dims++;
concat_spec = src_layout.sharding_spec(i);
}
}
return num_split_dims == 1 && num_concat_dims == 1 &&
split_spec == concat_spec;
}
StatusOr<mlir::Value> EmitAllToAll(
mlir::OpBuilder& builder, mlir::Value input,
const dtensor::Layout& src_layout, const dtensor::Layout& tgt_layout,
llvm::SmallPtrSet<mlir::Operation*, 4>* newly_created_ops) {
if (src_layout.IsEquivalent(tgt_layout)) return input;
if (src_layout.rank() != tgt_layout.rank()) {
return absl::InvalidArgumentError(absl::StrCat(
"Expected source and target layout to have the same rank, got ",
src_layout.rank(), " vs ", tgt_layout.rank()));
}
// Assume valid because CanUseAllToAll returned true.
// For convenience, operate on explicit input shapes. This isn't necessary,
// as we could instead generate operations on top of the dynamic shape.
const mlir::TensorType input_type =
mlir::dyn_cast<mlir::TensorType>(input.getType());
if (!input_type) {
return absl::InternalError(
llvm::formatv(
"Cannot cast input_type : {0} to TensorType. Shape must be "
" statically known before emitting AllToAll. This should not "
"happen as we already cast it when getting its shape.",
input.getType())
.str());
}
TF_ASSIGN_OR_RETURN(mlir::TensorType global_type,
GlobalTypeFromLocalType(src_layout, input_type));
TF_ASSIGN_OR_RETURN(mlir::TensorType output_type,
LocalTypeFromGlobalType(tgt_layout, global_type));
mlir::Location loc = DT_LOC2(input.getLoc(), "DTensorAllToAllOp");
mlir::TF::DTensorAllToAllOp all_to_all = mlir::TF::DTensorAllToAllOp::create(
builder, loc, output_type, input,
mlir::dtensor::LayoutAttr::get(builder.getContext(), src_layout),
mlir::dtensor::LayoutAttr::get(builder.getContext(), tgt_layout));
SetSingleLayoutOnOp(all_to_all, tgt_layout);
if (newly_created_ops != nullptr) newly_created_ops->insert(all_to_all);
return all_to_all.getOutput();
}
StatusOr<mlir::Value> EmitDenseToSparseToDense(
mlir::OpBuilder& builder, mlir::Value input,
llvm::SmallPtrSet<mlir::Operation*, 4>* newly_created_ops) {
// First create a Dense To Sparse Op. Since there is no DenseToSparseOp,
// we do it manually by creating the indices, values, and shapes tensor
// through various ops.
//
// indices tensor = tf.where(tf.not_equal(input, tf.zeros_like(tensor)))
// values tensor = tf.gather_nd(input, indices)
// shape tensor = tf.shape(input)
mlir::TF::ZerosLikeOp zeros_like =
mlir::TF::ZerosLikeOp::create(builder, input.getLoc(), input);
mlir::TF::NotEqualOp not_equal =
mlir::TF::NotEqualOp::create(builder, zeros_like.getLoc(), input,
zeros_like, builder.getBoolAttr(false));
mlir::TF::WhereOp indices = mlir::TF::WhereOp::create(
builder, not_equal.getLoc(),
mlir::RankedTensorType::get(GetShapeOfValue(not_equal).value(),
builder.getI64Type()),
not_equal);
mlir::TF::GatherNdOp values = mlir::TF::GatherNdOp::create(
builder, input.getLoc(), input.getType(), input, indices);
auto shape = mlir::TF::ShapeOp::create(builder, input.getLoc(), input,
builder.getBoolAttr(false));
// Emit a SparseToDenseOp and replace the SparseTensor with the result of
// this new op.
TF_ASSIGN_OR_RETURN(
mlir::Value zero_scalar,
CreateZeroScalarConst(
builder, input.getLoc(),
mlir::cast<mlir::TensorType>(input.getType()).getElementType()));
auto dense = mlir::TF::SparseToDenseOp::create(
builder, input.getLoc(), input.getType(),
mlir::ValueRange({indices, shape, values, zero_scalar}));
if (newly_created_ops != nullptr) {
for (auto new_op : {dense.getOperation(), shape.getOperation(),
values.getOperation(), indices.getOperation(),
not_equal.getOperation(), zeros_like.getOperation()}) {
newly_created_ops->insert(new_op);
}
}
return dense.getResult();
}
StatusOr<mlir::Value> EmitRelayout(
mlir::Value input, const dtensor::Layout& src_layout,
const dtensor::Layout& tgt_layout,
llvm::SmallPtrSet<mlir::Operation*, 4>* newly_created_ops) {
// EmitRelayout is performed by doing a split, an AllGather and another split.
// The first split oppertunistically splits input tensor dimension i on mesh
// mesh axis x if:
// 1. tgt_layout contains x at position i
// 2. src_layout is unsharded at position i.
// 3. src_layout does not contain mesh axis x.
// This produces intermediate layout 1.
// Next an all concat is performed on any axis in the intermediate layout 1
// that does not agree with the sharding on the output axis.
// This produces intermediate layout 2.
// A split is performed from intermediate layout 2 to the tgt layout.
if (src_layout == tgt_layout) {
return input;
}
mlir::OpBuilder builder(input.getContext());
TF_RETURN_IF_ERROR(SetBuilderInsertionAfterValue(input, builder));
// If two layouts are the same, or the only difference is layout type, then
// there is no need to actually relayout data.
if (src_layout.IsEquivalentIgnoringType(tgt_layout)) {
mlir::TF::IdentityOp op = mlir::TF::IdentityOp::create(
builder, input.getLoc(), input.getType(), input);
if (newly_created_ops != nullptr) newly_created_ops->insert(op);
return op.getOutput();
}
// Save whether the input is from a SparseToDenseOp. If it is, then we will
// emit a DenseToSparse and a SparseToDense op.
bool is_sparse = IsSparseValue(input);
if (!mlir::isa<mlir::RankedTensorType>(input.getType()))
return absl::InternalError(
"attempting to relayout a tensor that does not "
"have a rank");
if (src_layout.mesh() != tgt_layout.mesh()) {
return absl::InternalError(
absl::StrCat("Attempted to relayout to a different "
" mesh. Source Mesh = (",
src_layout.mesh().ToString(),
"). Target Mesh = ", tgt_layout.mesh().ToString(), ")."));
}
if (src_layout.rank() != tgt_layout.rank()) {
return absl::InternalError(
"Attempted to relayout to a different global shape.");
}
if (EnableAllToAllForRelayout() && !is_sparse &&
CanUseAllToAll(src_layout, tgt_layout)) {
// TODO(tmorris): support sparse case
TF_ASSIGN_OR_RETURN(mlir::Value all_to_all_result,
EmitAllToAll(builder, input, src_layout, tgt_layout,
newly_created_ops));
return all_to_all_result;
}
absl::flat_hash_set<std::string> src_sharding_dims;
for (int i = 0; i < src_layout.rank(); ++i)
src_sharding_dims.emplace(src_layout.sharding_spec(i));
std::vector<std::string> intermediate_specs_1(src_layout.rank());
for (int i = 0; i < src_layout.rank(); ++i) {
if (Layout::IsShardedDimension(tgt_layout.sharding_spec(i)) &&
!Layout::IsShardedDimension(src_layout.sharding_spec(i)) &&
!src_sharding_dims.contains(tgt_layout.sharding_spec(i)))
intermediate_specs_1[i] = tgt_layout.sharding_spec(i);
else
intermediate_specs_1[i] = src_layout.sharding_spec(i);
}
TF_ASSIGN_OR_RETURN(Layout intermediate_layout_1,
Layout::GetLayout(tgt_layout.type(), intermediate_specs_1,
src_layout.mesh()));
llvm::SmallPtrSet<mlir::Operation*, 4> local_newly_created_ops;
TF_ASSIGN_OR_RETURN(mlir::Value split_result,
EmitAllScatter(builder, input, src_layout,
intermediate_layout_1, newly_created_ops));
std::vector<std::string> intermediate_specs_2(src_layout.rank());
for (int i = 0; i < src_layout.rank(); ++i) {
if (Layout::IsShardedDimension(intermediate_specs_1[i]) &&
intermediate_specs_1[i] != tgt_layout.sharding_spec(i))
intermediate_specs_2[i] = Layout::kUnshardedDim;
else
intermediate_specs_2[i] = intermediate_specs_1[i];
}
TF_ASSIGN_OR_RETURN(Layout intermediate_layout_2,
Layout::GetLayout(tgt_layout.type(), intermediate_specs_2,
src_layout.mesh()));
TF_ASSIGN_OR_RETURN(
mlir::Value concat_result,
EmitAllGather(builder, split_result, intermediate_layout_1,
intermediate_layout_2, newly_created_ops));
auto all_scatter =
EmitAllScatter(builder, concat_result, intermediate_layout_2, tgt_layout,
newly_created_ops);
if (!is_sparse) return all_scatter;
if (!all_scatter.ok()) return all_scatter;
return EmitDenseToSparseToDense(builder, all_scatter.value(),
newly_created_ops);
}
mlir::Operation* EmitTransposeOp(mlir::OpBuilder& builder,
const mlir::Location& loc, mlir::Value input,
std::vector<int64_t>& perm_arr) {
auto tr_input_type = mlir::cast<mlir::ShapedType>(input.getType());
auto shape = tr_input_type.getShape();
auto perm_type = mlir::RankedTensorType::get(
{static_cast<int64_t>(perm_arr.size())}, builder.getIntegerType(64));
auto constant_attr = builder.getI64TensorAttr(perm_arr);
auto perm_op =
mlir::TF::ConstOp::create(builder, loc, perm_type, constant_attr);
std::vector<int64_t> transposed_shape(shape.begin(), shape.end());
for (int i = 0; i < shape.size(); i++) {
transposed_shape[i] = shape[perm_arr[i]];
}
auto transposed_type = mlir::RankedTensorType::get(
transposed_shape, tr_input_type.getElementType());
return mlir::TF::TransposeOp::create(builder, loc, transposed_type, input,
perm_op);
}
StatusOr<mlir::Operation*> EmitBarrierWithConstValue(mlir::OpBuilder& builder,
mlir::Location loc,
const Mesh& mesh,
int32_t value) {
absl::flat_hash_set<std::string> reduce_dims;
for (const MeshDimension& mesh_dim : mesh.dims()) {
reduce_dims.insert(mesh_dim.name);
}
return EmitAllReduce(
builder, Layout::ReplicatedOnMesh(mesh, /*rank=*/1), reduce_dims,
IntConst(builder, loc, std::vector<int32_t>{value}).getDefiningOp(),
kReduceOpAdd);
}
StatusOr<mlir::Operation*> EmitAllReduce(
mlir::OpBuilder& builder, const dtensor::Layout& output_layout,
const absl::flat_hash_set<std::string>& reduced_dims,
mlir::Operation* input, absl::string_view reduce_op) {
TF_ASSIGN_OR_RETURN(auto partitions, GetAllReducePartitionsFromReducedDims(
output_layout, reduced_dims));
const int32_t num_partitions = partitions.size();
// If every device lives in its own partition, we don't need to emit a
// collective.
if (num_partitions == output_layout.num_devices()) {
return InferSPMDExpandedLocalShape(input);
}
// Construct a flattened list of reduce partitions. This will be converted
// into a 2-D const tensor for the DTensorAllReduce op.
std::vector<int32_t> partitions_flat;
for (auto& p : partitions) {
if (p.second.size() != partitions.begin()->second.size()) {
return absl::InvalidArgumentError(
"AllReduce partitions had different sizes -- this is not supported "
"in MLIR.");
}
partitions_flat.insert(partitions_flat.end(), p.second.begin(),
p.second.end());
}
int32_t partition_size = partitions.begin()->second.size();
auto shaped_type = mlir::RankedTensorType::get(
{num_partitions, partition_size},
mlir::IntegerType::get(builder.getContext(), 32));
auto group_assignment =
mlir::DenseIntElementsAttr::get(shaped_type, partitions_flat);
TF_ASSIGN_OR_RETURN(std::string device_type,
DeviceTypeFromMesh(output_layout.mesh()));
mlir::Location loc = DT_LOC2(input->getLoc(), "DTensorAllReduceOp");
auto all_reduce = mlir::TF::DTensorAllReduceOp::create(
builder, loc, input->getResultTypes()[0], input->getOpResult(0),
mlir::TF::ConstOp::create(builder, DT_LOC2(loc, "group_assignment"),
group_assignment),
builder.getStringAttr(std::string(reduce_op)),
builder.getStringAttr(device_type));
SetSingleLayoutOnOp(all_reduce, output_layout);
input->getOpResult(0).replaceAllUsesExcept(
all_reduce.getResult(),
llvm::SmallPtrSet<mlir::Operation*, 1>{all_reduce});
return all_reduce.getOperation();
}
namespace {
// Returns a offset multiplier to calculate device id / mesh coordinate.
int GetMeshDimensionOffsetWithNeighbor(const Mesh& mesh,
const std::string& mesh_dim) {
const int index = mesh.GetMeshDimIndexWithName(mesh_dim);
const std::vector<int64_t> mesh_dim_sizes = mesh.dim_sizes();
int offset = 1;
for (int i = index + 1; i < mesh_dim_sizes.size(); ++i) {
offset = offset * mesh_dim_sizes[i];
}
return offset;
}
// Returns a mesh coordinate of mesh index with `mesh_dim_name` given
// `device_id`.
StatusOr<int> GetMeshCoordinateIndex(const Mesh& mesh,
const std::string& mesh_dim_name,
int device_id) {
const int offset = GetMeshDimensionOffsetWithNeighbor(mesh, mesh_dim_name);
TF_ASSIGN_OR_RETURN(int64_t mesh_dim_size, mesh.dim_size(mesh_dim_name));
return (device_id / offset) % mesh_dim_size;
}
// Returns a 2D tensor array of size [N, 2] that specifies source target pair
// to be used for halo exchange.
StatusOr<mlir::Value> CreateConstSrcTargetPair(const Mesh& mesh,
const std::string& mesh_dim_name,
bool shift_left,
mlir::Location location,
mlir::OpBuilder& builder) {
const int mesh_dim_index = mesh.GetMeshDimIndexWithName(mesh_dim_name);
const std::vector<MeshDimension> mesh_dimensions = mesh.dims();
llvm::SmallVector<int, 4> src_target_pair_flat;
src_target_pair_flat.reserve(mesh.local_device_ids().size() * 2);
for (const int local_device_id : mesh.local_device_ids()) {
// Calculate the mesh coordinate of the current local device id.
llvm::SmallVector<int, 4> mesh_coordinate_for_device_id;
for (const MeshDimension& mesh_dim : mesh_dimensions) {
TF_ASSIGN_OR_RETURN(
const int coordinate,
GetMeshCoordinateIndex(mesh, mesh_dim.name, local_device_id));
mesh_coordinate_for_device_id.push_back(coordinate);
}
// If mesh coordinate is on the left/right edge, then we conduct halo
// exchange with a processor which executes input block which represent
// `wrapped around` block.
const int mesh_coordinate = mesh_coordinate_for_device_id[mesh_dim_index];
TF_ASSIGN_OR_RETURN(const int dim_size, mesh.dim_size(mesh_dim_name));
// For tensor requiring halo exchange, we use collective permute.
const int src_device_id = local_device_id;
int target_device_id = 0;
for (const auto& data : llvm::enumerate(mesh_dimensions)) {
const MeshDimension& mesh_dim = data.value();
const int index = data.index();
int target_mesh_coordinate = 1;
if (mesh_dim.name == mesh_dim_name) {
target_mesh_coordinate =
shift_left ? mesh_coordinate - 1 : mesh_coordinate + 1;
// For processors executing input tensor on the left/right edges, target
// processor is the processor that executes wrapped around input block.
if (target_mesh_coordinate < 0 || target_mesh_coordinate >= dim_size)
target_mesh_coordinate =
(target_mesh_coordinate + dim_size) % dim_size;
} else {
target_mesh_coordinate = mesh_coordinate_for_device_id[index];
}
target_device_id +=
target_mesh_coordinate *
GetMeshDimensionOffsetWithNeighbor(mesh, mesh_dim.name);
}
src_target_pair_flat.push_back(src_device_id);
src_target_pair_flat.push_back(target_device_id);
}
const int num_pairs = src_target_pair_flat.size() / 2;
auto shaped_type = mlir::RankedTensorType::get(
{num_pairs, 2}, mlir::IntegerType::get(builder.getContext(), 32));
auto src_target_attr =
mlir::DenseIntElementsAttr::get(shaped_type, src_target_pair_flat);
mlir::Value src_target_pair_tensor =
mlir::TF::ConstOp::create(builder, location, src_target_attr);
return src_target_pair_tensor;
}
} // namespace
StatusOr<mlir::Value> EmitHaloExchange(mlir::OpBuilder& builder, int halo_size,
const std::string& mesh_dim,
const Layout& layout,
mlir::Value mesh_coordinates,
mlir::tf_device::ClusterOp cluster,
mlir::Location location,
mlir::Value tensor) {
const Mesh& mesh = layout.mesh();
// Check mesh dimension requirements for halo exchange.
if (!mesh.IsMeshDim(mesh_dim))
return absl::InvalidArgumentError(
"Requested halo exchange on unknown mesh dim");
// TODO(b/261485237): Add support for halo exchange for GPU/CPU.
if (!mesh.is_tpu_mesh())
return absl::InvalidArgumentError(
"Halo exchange is only supported on TPU.");
auto input_tensor_type =
mlir::dyn_cast<mlir::RankedTensorType>(tensor.getType());
if (!input_tensor_type || !input_tensor_type.hasStaticShape())
return absl::InvalidArgumentError(
"Static shape of input tensor must be known for halo exchange.");
llvm::ArrayRef<int64_t> input_tensor_shape = input_tensor_type.getShape();
const std::vector<std::string> sharding_specs = layout.sharding_spec_strs();
const int split_dim_index = std::distance(
sharding_specs.begin(), llvm::find(sharding_specs, mesh_dim));
if (input_tensor_shape[split_dim_index] < halo_size)
return absl::InvalidArgumentError(
"For halo exhange, input shard tensor size of each processor must be "
"greater than halo size");
TF_ASSIGN_OR_RETURN(const int mesh_dim_index, mesh.idx_for_dim(mesh_dim));
TF_ASSIGN_OR_RETURN(mlir::Value scalar_mesh_coordinate,
SelectScalarValueFromArray(builder, mesh_dim_index,
location, mesh_coordinates));
llvm::SmallVector<int64_t, 4> halo_exchange_tensor_shape;
for (const auto& size_and_index : llvm::enumerate(input_tensor_shape)) {
const int index = size_and_index.index();
const int size = size_and_index.value();
halo_exchange_tensor_shape.push_back(index == split_dim_index ? halo_size
: size);
}
// Find the halo tensor value to pad on the `left` side. Note that halo
// exchange can happen on top/bottom/left/right sides of a spatially
// partitioned tensor. However, we use `left`/`right` as the
// direction is implicit based on mesh dimension.
//
// For example, if mesh dimension splits the input tensor by its height
// dimension, then `left` actually means tensor to pad on the top side.
mlir::Value is_on_left_edge = mlir::TF::EqualOp::create(
builder, location,
CreateIntScalarConst(0, builder, location, /*use_int64=*/false),
scalar_mesh_coordinate, builder.getBoolAttr(true));
TF_ASSIGN_OR_RETURN(const int mesh_dim_size, mesh.dim_size(mesh_dim));
mlir::Value is_on_right_edge = mlir::TF::EqualOp::create(
builder, location,
CreateIntScalarConst(mesh_dim_size - 1, builder, location,
/*use_int64=*/false),
scalar_mesh_coordinate, builder.getBoolAttr(true));
// Create zero ghost tensor to pad on left side.
mlir::RankedTensorType halo_tensor_type = mlir::RankedTensorType::get(
halo_exchange_tensor_shape, input_tensor_type.getElementType());
auto halo_type = mlir::RankedTensorType::get(
halo_tensor_type.getShape(), input_tensor_type.getElementType());
mlir::Attribute const_attr;
if (halo_type.getElementType().isIntOrIndex()) {
const_attr =
mlir::DenseIntElementsAttr::get(halo_type, llvm::SmallVector<int>{0});
} else {
const_attr =
mlir::DenseFPElementsAttr::get(halo_type, llvm::SmallVector<float>{0});
}
mlir::Value ghost_tensor_left =
mlir::TF::ConstOp::create(builder, location, const_attr).getResult();
// Get the right side slice of the input tensor to pad on left side.
llvm::SmallVector<int64_t, 4> begin_left(layout.rank(), 0);
begin_left[split_dim_index] = input_tensor_shape[split_dim_index] - halo_size;
mlir::Value begin_tensor_left =
ops_util::GetR1Const(begin_left, builder, location);
llvm::SmallVector<int64_t, 4> size(input_tensor_shape.begin(),
input_tensor_shape.end());
size[split_dim_index] = halo_size;
mlir::Value size_tensor_left = ops_util::GetR1Const(size, builder, location);
mlir::Value sliced_tensor_left =
mlir::TF::SliceOp::create(builder, location, halo_type, tensor,
begin_tensor_left, size_tensor_left);
mlir::Value halo_tensor_left =
mlir::TF::SelectV2Op::create(builder, location, is_on_right_edge,
ghost_tensor_left, sliced_tensor_left);
// Invoke collective permute to receive the tensor from neighboring processor.
// Halo slices from the left neighbor are received on each processor (they
// are shifted right).
TF_ASSIGN_OR_RETURN(
mlir::Value src_target_pair_left,
CreateConstSrcTargetPair(mesh, mesh_dim, /*shift_left=*/false, location,
builder));
mlir::Value left_concat_value = mlir::TF::CollectivePermuteOp::create(
builder, location, sliced_tensor_left.getType(), halo_tensor_left,
src_target_pair_left);
mlir::Value ghost_tensor_right =
mlir::TF::ConstOp::create(builder, location, const_attr).getResult();
// Else, values to pad is tensor from different processor. We use collective
// permute to access tensor slice from another device.
// Get the left side slice of the input tensor.
llvm::SmallVector<int64_t, 4> begin_right(layout.rank(), 0);
mlir::Value begin_tensor_right =
ops_util::GetR1Const(begin_right, builder, location);
mlir::Value size_tensor_right = ops_util::GetR1Const(size, builder, location);
mlir::Value sliced_tensor_right =
mlir::TF::SliceOp::create(builder, location, halo_type, tensor,
begin_tensor_right, size_tensor_right);
// Find the halo tensor value to pad on the `right` side.
// If input block is on the right edge, we use zero ghost tensor instead.
mlir::Value halo_tensor_right =
mlir::TF::SelectV2Op::create(builder, location, is_on_left_edge,
ghost_tensor_right, sliced_tensor_right);
// Invoke collective permute to receive the tensor from neighboring processor.
// Halo slices from the right neighbor are received on each processor (they
// are shifted left).
TF_ASSIGN_OR_RETURN(
mlir::Value src_target_pair_right,
CreateConstSrcTargetPair(mesh, mesh_dim, /*shift_left=*/true, location,
builder));
mlir::Value right_concat_value = mlir::TF::CollectivePermuteOp::create(
builder, location, sliced_tensor_right.getType(), halo_tensor_right,
src_target_pair_right);
// Final halo exchanged value is concatenated value of left_concat_value,
// tensor, and right_concat_value in the mesh_dimension.
llvm::SmallVector<int64_t, 4> final_shape(input_tensor_shape.begin(),
input_tensor_shape.end());
final_shape[split_dim_index] = final_shape[split_dim_index] + 2 * halo_size;
auto final_type = mlir::RankedTensorType::get(
final_shape, input_tensor_type.getElementType());
mlir::Value concat_axis =
CreateIntScalarConst(split_dim_index, builder, location);
mlir::Value final_value = mlir::TF::ConcatV2Op::create(
builder, location, final_type,
llvm::SmallVector<mlir::Value, 4>{left_concat_value, tensor,
right_concat_value},
concat_axis);
return final_value;
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/collectives.h
+116
View File
@@ -0,0 +1,116 @@
/* Copyright 2022 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_DTENSOR_MLIR_COLLECTIVES_H_
#define TENSORFLOW_DTENSOR_MLIR_COLLECTIVES_H_
#include <string>
#include <vector>
#include "absl/container/flat_hash_set.h"
#include "absl/strings/string_view.h"
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/Location.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
namespace tensorflow {
namespace dtensor {
// Emits collective ops to convert `input` from `src_layout` to `tgt_layout`.
// `src_layout` and `tgt_layout` must have the same rank. For each dimension,
// it can only go from sharded to replicated. `input` must have static shapes.
StatusOr<mlir::Value> EmitAllGather(
mlir::OpBuilder& builder, mlir::Value input,
const dtensor::Layout& src_layout, const dtensor::Layout& tgt_layout,
llvm::SmallPtrSet<mlir::Operation*, 4>* newly_created_ops = nullptr);
// Given an input layout and a desired layout, inserts the necessary slice to
// slice the original value based on the device id. All ops created by this
// function are added to new_created_ops.
//
// Note that the newly created ops are inserted `after` original_value.
StatusOr<const mlir::Value> EmitAllScatter(
mlir::OpBuilder& builder, const mlir::Value& original_value,
const Layout& original_layout, const Layout& desired_layout,
llvm::SmallPtrSet<mlir::Operation*, 4>* newly_created_ops = nullptr);
// Emits splits and calls EmitAllGather (once) to relayout from the src layout
// to the tgt layout on a single mesh.
// Shape of input is expected to be the local shape for src_layout.
StatusOr<mlir::Value> EmitRelayout(
mlir::Value input, const dtensor::Layout& src_layout,
const dtensor::Layout& tgt_layout,
llvm::SmallPtrSet<mlir::Operation*, 4>* newly_created_ops = nullptr);
// Emits TransposeOp that permutes the input shape.
mlir::Operation* EmitTransposeOp(mlir::OpBuilder& builder,
const mlir::Location& loc, mlir::Value input,
std::vector<int64_t>& perm_arr);
// Emits collective ops to reduce `input` over `reduced_dims`.
StatusOr<mlir::Operation*> EmitAllReduce(
mlir::OpBuilder& builder, const dtensor::Layout& output_layout,
const absl::flat_hash_set<std::string>& reduced_dims,
mlir::Operation* input, absl::string_view reduce_op);
// Emits a barrier used for synchronization purposes and returns
// a R1 const value using `value`. More precisely, this barrier
// guarantees that
// 1. Side-effect Ops before this barrier are complete before this op begins.
// 2. Side-effect Ops after this barrier start after this barrier completes.
//
// Note that the returned operation must be used in the graph. If it is not
// used, then this op will be removed from the graph from various compiler
// passes and thus there will be no barrier.
//
// Used for introducing a barrier before every Merge op during checkpointing.
StatusOr<mlir::Operation*> EmitBarrierWithConstValue(mlir::OpBuilder& builder,
mlir::Location loc,
const Mesh& mesh,
int32_t value);
// Given input `tensor` that is sharded across spatial dimensions, conduct
// halo exchange such that each spatially sharded input blocks exchange
// `halo_size` slice with its neighboring processors.
// If the input block is at the left/right/top/bottom edge, then ghost halo
// tensor (zero) are padded instead. `mesh_dim` specifies the dimension which
// halo exchange will be conducted. For example, if we consider a 4D Tensor
// (batch, height, width, channel) that has layout (*, h, w, *). Then,
// `mesh_dim` == "w" would mean that halo exchange will occur along the width
// dimension. That is halo tensors with right/left neighbors will be exchanged.
StatusOr<mlir::Value> EmitHaloExchange(mlir::OpBuilder& builder, int halo_size,
const std::string& mesh_dim,
const Layout& layout,
mlir::Value mesh_coordinates,
mlir::tf_device::ClusterOp cluster,
mlir::Location location,
mlir::Value tensor);
// Emits a DenseToSparse op followed by a SparseToDenseOp.
// This is useful for emitting a Relayout on a SparseTensor.
// One usage of this is in EmitRelayout when the input is a SparseTensor.
StatusOr<mlir::Value> EmitDenseToSparseToDense(
mlir::OpBuilder& builder, mlir::Value input,
llvm::SmallPtrSet<mlir::Operation*, 4>* newly_created_ops = nullptr);
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_MLIR_COLLECTIVES_H_
tensorflow/dtensor/mlir/collectives_common.cc
+98
View File
@@ -0,0 +1,98 @@
/* Copyright 2022 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/dtensor/mlir/collectives_common.h"
#include <cstddef>
#include <cstdint>
#include <map>
#include <string>
#include <vector>
#include "absl/container/flat_hash_set.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
namespace tensorflow {
namespace dtensor {
// A map from a unique set of kept mesh dimension values (a partition) to
// IDs of devices in that partition.
//
// Users will typically ignore the key, but use the map values as the group
// assignment for collective operations. This is intentionally a
// std::map instead of absl::flat_hash_map to guarantee all hosts in
// a multi-host cluster will generate the same grouping, and therefore the same
// XLA program fingerprint, independently. std::map guarantees the same
// iteration order.
using AllReducePartitions = std::map<DeviceLocation, std::vector<int32_t>>;
// Computes AllReduce partitions using reduced mesh dimension names.
//
// Reduction groups are formed across all _non_-reduced dimensions. For example,
// in the following scenario:
//
// output_layout.dims() = [a, b]
// output_layout.mesh() = [(x, 8), (y, 4)]
// reduced_dims = `x`
//
// We first reduce over `a` locally on each device, producing 32 local
// reductions. We then AllReduce within each of the 4 partitions. Each partition
// corresponds to one unique value of `y` and has 8 devices. The end result is
// sharded over the y mesh dimension and replicated 8 times.
//
// The returned map should have four entries with key values from [0] to [3]
// (unique values of `y`). Each key maps to IDs of devices with that `y` value.
StatusOr<AllReducePartitions> GetAllReducePartitionsFromReducedDims(
const dtensor::Layout& output_layout,
const absl::flat_hash_set<std::string>& reduced_dims) {
AllReducePartitions partitions;
for (int64_t device = 0; device < output_layout.num_devices(); ++device) {
TF_ASSIGN_OR_RETURN(const DeviceLocation device_loc,
output_layout.mesh().device_location(device));
DeviceLocation kept_dims;
for (int64_t dim_idx = 0; dim_idx < device_loc.size(); ++dim_idx) {
if (!reduced_dims.contains(output_layout.mesh().dim_name(dim_idx))) {
kept_dims.push_back(device_loc[dim_idx]);
}
}
partitions[kept_dims].push_back(device);
}
return partitions;
}
// Use the first device in the mesh to extract the device name. For example:
//
// device_path = "/job:localhost/replica:0/task:0/device:TPU:0"
// device_type = "/job:localhost/replica:0/task:0/device:TPU"
// device_id = 0
//
// The device ID can be obtained through DeviceId as a runtime input. We may
// need it in the future to enable device ID-based branch divergence.
StatusOr<std::string> DeviceTypeFromMesh(const Mesh& mesh) {
std::string device_path =
mesh.is_remote() ? mesh.global_devices()[0] : mesh.local_devices()[0];
size_t device_path_pos = device_path.find_last_of(':');
if (device_path_pos == std::string::npos) {
return absl::InvalidArgumentError(
absl::StrCat("Unexpected device path: ", device_path));
}
return device_path.substr(0, device_path_pos);
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/collectives_common.h
+43
View File
@@ -0,0 +1,43 @@
/* Copyright 2022 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_DTENSOR_MLIR_COLLECTIVES_COMMON_H_
#define TENSORFLOW_DTENSOR_MLIR_COLLECTIVES_COMMON_H_
#include <map>
#include <string>
#include <vector>
#include "absl/container/flat_hash_set.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
namespace tensorflow {
namespace dtensor {
// Computes AllReduce partitions using reduced mesh dimension names.
StatusOr<std::map<DeviceLocation, std::vector<int32_t>>>
GetAllReducePartitionsFromReducedDims(
const dtensor::Layout& output_layout,
const absl::flat_hash_set<std::string>& reduced_dims);
// Use the first device in the mesh to extract the device name.
StatusOr<std::string> DeviceTypeFromMesh(const Mesh& mesh);
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_MLIR_COLLECTIVES_COMMON_H_
tensorflow/dtensor/mlir/constant_folding.cc
+92
View File
@@ -0,0 +1,92 @@
/* Copyright 2022 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/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/LogicalResult.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/Interfaces/SideEffectInterfaces.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Transforms/FoldUtils.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
namespace tensorflow {
namespace dtensor {
namespace {
#define GEN_PASS_DEF_DTENSORCONSTANTFOLDING
#include "tensorflow/dtensor/mlir/dtensor_passes.h.inc"
constexpr int kMaxIteration = 10;
mlir::LogicalResult FoldConstantOp(mlir::OperationFolder& folder,
mlir::TF::ConstOp op) {
bool changed = false;
int i = 0;
// Iterate until convergence or until maxIterations. Deletion of the op as
// a result of being dead or folded is convergence.
do {
changed = false;
// If the operation is trivially dead - remove it.
if (isOpTriviallyDead(op)) {
op->erase();
return mlir::success();
}
// Try to fold this op.
bool inPlaceUpdate;
if (succeeded(folder.tryToFold(op, &inPlaceUpdate))) {
changed = true;
if (!inPlaceUpdate) {
return mlir::success();
}
}
} while (changed && ++i < kMaxIteration);
return mlir::success();
}
// MLIR pass that folds constants that can be removed or deduplicated away.
struct DTensorConstantFolding
: public impl::DTensorConstantFoldingBase<DTensorConstantFolding> {
void runOnOperation() override {
mlir::MLIRContext& context = getContext();
mlir::OperationFolder helper(&context);
// Collect and fold the operations within the function.
llvm::SmallVector<mlir::TF::ConstOp, 8> const_ops;
getOperation().walk([&](mlir::TF::ConstOp op) { const_ops.push_back(op); });
// Attempt to fold the specified operation, including handling unused or
// duplicated constants.
for (mlir::TF::ConstOp op : llvm::reverse(const_ops))
if (mlir::failed(FoldConstantOp(helper, op))) return signalPassFailure();
}
};
} // namespace
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateDTensorConstantFolding() {
return std::make_unique<DTensorConstantFolding>();
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/create_dtensor_mlir_passes.h
+169
View File
@@ -0,0 +1,169 @@
/* Copyright 2022 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_DTENSOR_MLIR_CREATE_DTENSOR_MLIR_PASSES_H_
#define TENSORFLOW_DTENSOR_MLIR_CREATE_DTENSOR_MLIR_PASSES_H_
#include <memory>
#include <optional>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "tensorflow/core/lib/core/status.h"
namespace tensorflow {
namespace dtensor {
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateDTensorOpToDeviceClusterPass();
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateDTensorDeviceMeshClusterCoarsening();
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>> CreateDTensorDCE();
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateDTensorUndoMergeConstAcrossMesh();
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateDTensorElideIdentityBeforeCopyToMesh();
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateDTensorConstantFolding();
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateDTensorAllReduceSumOptimization();
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateDTensorAllReduceScatterOptimization();
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateDTensorAllReduceCombineOptimization();
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateDTensorMixedPrecisionReducePass();
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateDTensorSetDefaultSharding();
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateDTensorDesignateResourceHandleMesh();
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateDTensorPropagateDefaultLayout();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorHandleCrossClusterDependencies();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorAnnotateGlobalShape();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorLayoutPropagationPassV2();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorMeshPropagationPass();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorSPMDExpansion();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorClusterFunctionConversion();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorPropagateDeviceIdToFunctionArgs();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorTPUIntegration();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorTpuAddResourceDeviceAttribute();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorUpdateTPUMetadata();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateFunctionRenamingPass();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorMultiDeviceExpansionPass();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorAllReduceLoweringPass();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorAllToAllLoweringPass();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorReduceScatterLoweringPass();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorAllGatherLoweringPass();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorAllScatterLoweringPass();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorMergeClustersPass();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorDecomposeControlflowPass();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorLowerSendRecv();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorMoveCompilationToHost();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorSparseTensorToDenseTensor();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorSparseExpansion();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorInferShapesForRestoreV2Op();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorSetHloShardingPass(
std::optional<bool> check_layout_use_xla_spmd = std::optional<bool>(false));
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateDTensorLayoutToXlaShardingOpPass();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorReplaceAuxiliaryDTensorLayoutOpPass();
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorRemoveDTensorLayoutPass();
// Creates a pass that replaces `tf.Relayout` with `tf.Identity`.
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateDTensorReplaceRelayoutWithIdentityPass();
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateDTensorCollectiveTypeLoweringPass();
// Generate the code for registering passes.
#define GEN_PASS_REGISTRATION
#include "tensorflow/dtensor/mlir/dtensor_passes.h.inc"
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_MLIR_CREATE_DTENSOR_MLIR_PASSES_H_
tensorflow/dtensor/mlir/dce.cc
+50
View File
@@ -0,0 +1,50 @@
/* Copyright 2022 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/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/Interfaces/SideEffectInterfaces.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
namespace tensorflow {
namespace dtensor {
namespace {
#define GEN_PASS_DEF_DTENSORDCE
#include "tensorflow/dtensor/mlir/dtensor_passes.h.inc"
// MLIR pass that removes trivially unused operations in graph.
struct DTensorDCE : public impl::DTensorDCEBase<DTensorDCE> {
void runOnOperation() override {
mlir::MLIRContext& context = getContext();
mlir::OpBuilder builder(&context);
getOperation().walk([](mlir::Operation* op) {
if (mlir::isOpTriviallyDead(op)) op->erase();
});
}
};
} // namespace
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>> CreateDTensorDCE() {
return std::make_unique<DTensorDCE>();
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/designate_resource_handle_mesh.cc
+82
View File
@@ -0,0 +1,82 @@
/* Copyright 2022 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/Support/Casting.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/dtensor/cc/constants.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
namespace tensorflow {
namespace dtensor {
namespace {
#define GEN_PASS_DEF_DTENSORDESIGNATERESOURCEHANDLEMESH
#include "tensorflow/dtensor/mlir/dtensor_passes.h.inc"
mlir::LogicalResult SetMeshForResourceCreatingCluster(
mlir::tf_device::ClusterOp cluster, mlir::OpBuilder* builder) {
auto result = cluster.walk([](mlir::Operation* op) {
if (llvm::isa<mlir::TF::VarHandleOp, mlir::TF::DestroyResourceOp>(op))
return mlir::WalkResult::interrupt();
return mlir::WalkResult::advance();
});
if (!result.wasInterrupted()) return mlir::success();
if (!cluster->hasAttr(kMeshAttr)) {
cluster->setAttr(kMeshAttr, builder->getStringAttr(Mesh::kEmptyMeshString));
}
return mlir::success();
}
struct DTensorDesignateResourceHandleMesh
: public impl::DTensorDesignateResourceHandleMeshBase<
DTensorDesignateResourceHandleMesh> {
void runOnOperation() override {
mlir::MLIRContext& context = getContext();
mlir::OpBuilder builder(&context);
auto walk_result =
getOperation().walk([&](mlir::tf_device::ClusterOp cluster) {
if (mlir::failed(
SetMeshForResourceCreatingCluster(cluster, &builder)))
return mlir::WalkResult::interrupt();
return mlir::WalkResult::advance();
});
if (walk_result.wasInterrupted()) return signalPassFailure();
}
};
} // namespace
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateDTensorDesignateResourceHandleMesh() {
return std::make_unique<DTensorDesignateResourceHandleMesh>();
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/device_mesh_cluster_coarsening.cc
+301
View File
@@ -0,0 +1,301 @@
/* Copyright 2022 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 <iterator>
#include <memory>
#include <optional>
#include <utility>
#include "absl/status/status.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/LogicalResult.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Block.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/Diagnostics.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "mlir/IR/ValueRange.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/dtensor/cc/constants.h"
#include "tensorflow/dtensor/mlir/layout_parsing.h"
namespace tensorflow {
namespace dtensor {
namespace {
#define GEN_PASS_DEF_DTENSORDEVICEMESHCLUSTERCOARSENING
#include "tensorflow/dtensor/mlir/dtensor_passes.h.inc"
constexpr char kMissingMeshAttributeErrorMessage[] =
"failed to merge mesh cluster as cluster does not have mesh attribute. "
"This is likely due to problem in mesh propagation.";
// Determines whether two adjoining clusters should be merged.
mlir::LogicalResult ShouldMergeClusters(mlir::tf_device::ClusterOp cluster_a,
mlir::tf_device::ClusterOp cluster_b,
bool* should_merge) {
if (cluster_a->getParentRegion() != cluster_b->getParentRegion()) {
*should_merge = false;
return mlir::success();
}
auto mesh_a_or_status = ExtractDeviceMeshFromOp(cluster_a.getOperation());
if (!mesh_a_or_status.ok())
return cluster_a.emitOpError(
absl::StatusMessageAsCStr(mesh_a_or_status.status()));
auto mesh_b_or_status = ExtractDeviceMeshFromOp(cluster_b.getOperation());
if (!mesh_b_or_status.ok())
return cluster_b.emitOpError(
absl::StatusMessageAsCStr(mesh_b_or_status.status()));
auto mesh_a = mesh_a_or_status.value();
auto mesh_b = mesh_b_or_status.value();
if (!mesh_a || !mesh_b) {
return !mesh_a ? cluster_a.emitOpError(kMissingMeshAttributeErrorMessage)
: cluster_b.emitOpError(kMissingMeshAttributeErrorMessage);
}
*should_merge = mesh_a == mesh_b;
return mlir::success();
}
// Moves all ops (except tf_device.return op) inside `src_cluster` to
// block inside `target_cluster`. Ops are moved before the `exit_op`
// inside the `target_cluster`.
void MoveOpsInsideCluster(mlir::tf_device::ClusterOp src_cluster,
mlir::tf_device::ClusterOp target_cluster,
mlir::Operation* exit_op) {
auto& cluster_body = src_cluster.GetBody().getOperations();
target_cluster.GetBody().getOperations().splice(
exit_op->getIterator(), cluster_body, cluster_body.begin(),
std::prev(cluster_body.end()));
}
// Returns a list of pair of mlir Values that represent <return values of ops
// inside the merged_cluster, output values of merged cluster>.
//
// If outputs of `current_cluster` is used as operands to ops in
// `merging_cluster`, then make sure to replace operands such that
// results values from the inner ops of `current_cluster` is used instead.
//
// For example,
// %0 = "tf_device.cluster"() ({
// %1 = "tf.A"() : () -> tensor<i32>
// "tf_device.return"(%1) : (tensor<i32>) -> ()
// }) { mesh = "mesh_config: cpu[1, 1]"} : () -> (tensor<i32>)
//
// %2 = "tf_device.cluster"() ({
// %3 = "tf.B"(%0) : (tenosr<i32>) -> tensor<f32>
// "tf_device.return"(%3) : (tensor<f32>) -> ()
// }) { mesh = "mesh_config: cpu[1, 1]"} : () -> (tensor<f32>)
//
// will be:
// %0 = "tf_device.cluster"() ({
// %1 = "tf.A"() : () -> tensor<i32>
//
// # NOTE: `tf.B` op now takes operand directly from
// # `tf.A` instead of `tf_dtensor.cluster op.
// %2 = "tf.B"(%1) : (tenosr<i32>) -> tensor<f32>
// "tf_device.return"(%1, %2) : (tensor<i32>, tensor<f32>)) -> ()
// }) {mesh = "mesh_config: cpu[1, 1]"} : () -> (tensor<i32>, tensor<f32>)
llvm::SmallVector<std::pair<mlir::Value, mlir::Value>, 8>
GetMergedMeshClusterResults(mlir::tf_device::ClusterOp current_cluster,
mlir::tf_device::ClusterOp merging_cluster) {
llvm::SmallVector<std::pair<mlir::Value, mlir::Value>, 8>
merged_cluster_results;
merged_cluster_results.reserve(current_cluster.getNumResults() +
merging_cluster.getNumResults());
auto current_cluster_return_op = current_cluster.GetBody().getTerminator();
for (auto result : llvm::zip(current_cluster_return_op->getOpOperands(),
current_cluster.getResults())) {
mlir::Value inner_op_result = std::get<0>(result).get();
mlir::Value outer_op_result = std::get<1>(result);
// If the output value of `current_cluster` is only used by ops
// inside the `merged_cluster`, do not add the value as a return
// value for newly created tf_device.cluster op.
bool result_only_used_by_merging_cluster = true;
for (auto& use : llvm::make_early_inc_range(outer_op_result.getUses())) {
if (merging_cluster.GetBody().findAncestorOpInBlock(*use.getOwner())) {
use.set(inner_op_result);
} else {
result_only_used_by_merging_cluster = false;
}
}
if (!result_only_used_by_merging_cluster) {
merged_cluster_results.emplace_back(inner_op_result, outer_op_result);
}
}
auto merging_cluster_return_op = merging_cluster.GetBody().getTerminator();
for (auto result : llvm::zip(merging_cluster_return_op->getOpOperands(),
merging_cluster.getResults())) {
mlir::Value inner_op_result = std::get<0>(result).get();
mlir::Value outer_op_result = std::get<1>(result);
if (!outer_op_result.getUses().empty())
merged_cluster_results.emplace_back(inner_op_result, outer_op_result);
}
return merged_cluster_results;
}
// Updates the users of `merging_cluster` so that they use values
// from `merged_cluster` instead.
void ReplaceOperandUsagesWithMergedClusterOutputs(
mlir::ValueRange values_to_replace,
mlir::tf_device::ClusterOp merged_cluster) {
for (auto result :
llvm::zip(values_to_replace, merged_cluster.getResults())) {
std::get<0>(result).replaceAllUsesWith(std::get<1>(result));
}
}
// Creates a new tf_device.cluster op that merges
// `current_cluster` and `merging_cluster`.
mlir::LogicalResult CreateMergedMeshCluster(
mlir::OpBuilder* builder, mlir::tf_device::ClusterOp current_cluster,
mlir::tf_device::ClusterOp merging_cluster,
mlir::tf_device::ClusterOp* merged_cluster) {
auto return_values =
GetMergedMeshClusterResults(current_cluster, merging_cluster);
llvm::SmallVector<mlir::Type, 8> merged_cluster_output_types;
llvm::SmallVector<mlir::Value, 8> merged_cluster_output_values;
llvm::SmallVector<mlir::Value, 8> output_values_to_replace;
merged_cluster_output_types.reserve(return_values.size());
merged_cluster_output_values.reserve(return_values.size());
output_values_to_replace.reserve(return_values.size());
for (auto cluster_return_value : return_values) {
auto inner_op_return_value = std::get<0>(cluster_return_value);
merged_cluster_output_types.emplace_back(inner_op_return_value.getType());
merged_cluster_output_values.emplace_back(inner_op_return_value);
output_values_to_replace.emplace_back(std::get<1>(cluster_return_value));
}
*merged_cluster = mlir::tf_device::ClusterOp::create(
*builder, current_cluster.getLoc(), merged_cluster_output_types);
auto mesh_attr = current_cluster->getAttrOfType<mlir::StringAttr>(kMeshAttr);
if (!mesh_attr)
return current_cluster.emitOpError(kMissingMeshAttributeErrorMessage);
(*merged_cluster)->setAttr(kMeshAttr, mesh_attr);
// Create a terminator op that returns all return values from
// `current_cluster` and `merging_cluster`.
merged_cluster->getBody().push_back(new mlir::Block);
builder->setInsertionPointToEnd(&merged_cluster->GetBody());
mlir::tf_device::ReturnOp::create(*builder, merged_cluster->getLoc(),
merged_cluster_output_values);
// Make sure to replace usages of tf_device.cluster ops to be merged-away with
// newly created tf_device.cluster op.
ReplaceOperandUsagesWithMergedClusterOutputs(output_values_to_replace,
*merged_cluster);
return mlir::success();
}
// Merges `current_cluster` and `merging_cluster` and returns a new merged
// tf_device.cluster.
mlir::LogicalResult MergeClusters(mlir::OpBuilder* builder,
mlir::tf_device::ClusterOp current_cluster,
mlir::tf_device::ClusterOp merging_cluster,
mlir::tf_device::ClusterOp* merged_cluster) {
builder->setInsertionPoint(current_cluster);
// Create new tf_device.cluster op that outputs results of both
// `current_cluster` and `merging_cluster`.
if (mlir::failed(CreateMergedMeshCluster(builder, current_cluster,
merging_cluster, merged_cluster)))
return mlir::failure();
// Move all ops to newly created merged cluster.
auto exit_op = merged_cluster->GetBody().getTerminator();
MoveOpsInsideCluster(current_cluster, *merged_cluster, exit_op);
MoveOpsInsideCluster(merging_cluster, *merged_cluster, exit_op);
// Remove mesh clusters as they are now merged to a new cluster.
current_cluster.erase();
merging_cluster.erase();
return mlir::success();
}
// Loops through tf_device.Cluster ops and merge clusters with same execution
// device set.
mlir::LogicalResult ClusterDeviceClusterOpsInBlock(mlir::OpBuilder* builder,
mlir::Block* block) {
llvm::SmallVector<mlir::tf_device::ClusterOp, 4> block_ops;
block->walk([&](mlir::Operation* op) {
if (auto cluster = llvm::dyn_cast<mlir::tf_device::ClusterOp>(op))
block_ops.emplace_back(cluster);
});
std::optional<mlir::tf_device::ClusterOp> current_cluster;
for (mlir::tf_device::ClusterOp cluster :
llvm::make_early_inc_range(block_ops)) {
if (!current_cluster.has_value()) {
current_cluster = cluster;
continue;
}
bool should_merge;
if (failed(ShouldMergeClusters(*current_cluster, cluster, &should_merge)))
return mlir::failure();
if (should_merge) {
mlir::tf_device::ClusterOp new_cluster;
if (mlir::failed(
MergeClusters(builder, *current_cluster, cluster, &new_cluster)))
return mlir::failure();
current_cluster.emplace(new_cluster);
} else {
current_cluster.emplace(cluster);
}
}
return mlir::success();
}
} // namespace
// MLIR pass that merges cluster ops with the same mesh attribute.
struct DTensorDeviceMeshClusterCoarsening
: public impl::DTensorDeviceMeshClusterCoarseningBase<
DTensorDeviceMeshClusterCoarsening> {
void runOnOperation() override {
mlir::MLIRContext& context = getContext();
mlir::OpBuilder builder(&context);
for (mlir::Block& block : getOperation())
if (mlir::failed(ClusterDeviceClusterOpsInBlock(&builder, &block)))
return signalPassFailure();
}
};
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateDTensorDeviceMeshClusterCoarsening() {
return std::make_unique<DTensorDeviceMeshClusterCoarsening>();
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/device_utils.cc
+70
View File
@@ -0,0 +1,70 @@
/* Copyright 2022 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/dtensor/mlir/device_utils.h"
#include "llvm/Support/Casting.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/dtensor/cc/dstatus.h"
namespace tensorflow {
namespace dtensor {
// Returns an MLIR value representing the current device ID.
StatusOr<mlir::Value> DeviceId(mlir::Operation* op) {
mlir::func::FuncOp function = llvm::dyn_cast<mlir::func::FuncOp>(op);
if (!function) {
// Device ID is the 0th argument of the enclosing function.
function = op->getParentOfType<mlir::func::FuncOp>();
if (!function)
return absl::InvalidArgumentError(
"operation must be enclosed inside a function.");
}
if (function.getNumArguments() == 0)
return absl::InvalidArgumentError(
"enclosing function must contain device id as argument");
auto device_id = function.getArgument(0);
auto device_id_type =
mlir::dyn_cast<mlir::RankedTensorType>(device_id.getType());
if (!device_id_type ||
!mlir::isa<mlir::IntegerType>(device_id_type.getElementType()))
return absl::InvalidArgumentError(
"0-th argument of the enclosing function should be integer device id.");
return device_id;
}
StatusOr<mlir::Value> DeviceId(mlir::Value val) {
if (auto block_arg = mlir::dyn_cast<mlir::BlockArgument>(val)) {
auto device_id = block_arg.getOwner()->getArgument(0);
auto device_id_type =
mlir::dyn_cast<mlir::RankedTensorType>(device_id.getType());
if (!device_id_type ||
!mlir::isa<mlir::IntegerType>(device_id_type.getElementType()))
return absl::InvalidArgumentError(
"0-th argument of the enclosing block should be integer device id.");
return device_id;
}
return DeviceId(val.getDefiningOp());
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/device_utils.h
+49
View File
@@ -0,0 +1,49 @@
/* Copyright 2022 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_DTENSOR_MLIR_DEVICE_UTILS_H_
#define TENSORFLOW_DTENSOR_MLIR_DEVICE_UTILS_H_
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Location.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
namespace tensorflow {
namespace dtensor {
// Returns an MLIR value representing the current device ID. Passing op is
// assumed to be wrapped by a function graph from dtensor_device, and the first
// argument of the it need to be DeviceId.
StatusOr<mlir::Value> DeviceId(mlir::Operation* op);
StatusOr<mlir::Value> DeviceId(mlir::Value val);
// Returns the device ordinal of the device executing a function. Device ordinal
// is the index of a TF device among all TF devices of the same type attached to
// the same TF task.
StatusOr<mlir::Value> GetDeviceOrdinal(const Mesh& mesh,
const mlir::Location& loc,
mlir::func::FuncOp function,
mlir::OpBuilder* builder,
bool return_int64_type = true);
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_MLIR_DEVICE_UTILS_H_
tensorflow/dtensor/mlir/dtensor_allreduce_combine_optimization.cc
+807
View File
@@ -0,0 +1,807 @@
/* Copyright 2022 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 <algorithm>
#include <cmath>
#include <cstdint>
#include <cstdlib>
#include <memory>
#include <optional>
#include <set>
#include <string>
#include <vector>
#include "absl/log/check.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/FormatVariadic.h"
#include "mlir/Analysis/TopologicalSortUtils.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Block.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.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/Matchers.h" // from @llvm-project
#include "mlir/IR/Types.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/analysis/side_effect_analysis.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/collection_ops_util.h"
#include "tensorflow/compiler/mlir/utils/name_utils.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/strcat.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/dtensor/cc/constants.h"
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/mlir/dtensor_location.h"
#include "tensorflow/dtensor/mlir/ir/tf_dtensor.h"
#include "tensorflow/dtensor/mlir/layout_parsing.h"
namespace tensorflow {
namespace dtensor {
namespace {
#define GEN_PASS_DEF_DTENSORALLREDUCECOMBINEOPTIMIZATION
#include "tensorflow/dtensor/mlir/dtensor_passes.h.inc"
namespace ops_util = ::mlir::TF::collection_ops_util;
// Pad the merged tensor shape to multiples of 1024B, so delinearization
// skipping optimization in XLA can get activated.
constexpr int32_t kAllReducePadding = 1024;
// Returns true if `successor` depends on `predecessor`.
// TODO(jiawenhao): Repeatedly computing dependency sets for a large cluster can
// get expensive when the number of all-reduces is high. Consider building a
// cluster-scope op dependency graph ahead of time to amortize the cost.
bool DependsOn(mlir::Operation* successor, mlir::Operation* predecessor,
const mlir::TF::detail::SideEffectAnalysisInfo& info) {
llvm::SmallVector<mlir::Operation*, 4> to_visit;
llvm::SmallPtrSet<mlir::Operation*, 4> visited;
to_visit.push_back(predecessor);
while (!to_visit.empty()) {
mlir::Operation* producer = to_visit.pop_back_val();
if (visited.contains(producer)) continue;
visited.insert(producer);
if (successor == producer) return true;
for (mlir::Operation* user : producer->getUsers()) {
if (visited.contains(user)) continue;
to_visit.push_back(user);
}
// Include indirectly dependent ops from side effects
for (mlir::Operation* user : info.DirectControlSuccessors(producer)) {
if (visited.contains(user)) continue;
to_visit.push_back(user);
}
}
return false;
}
// Merge all-reduces in the group into one all-reduce.
//
// Requirements:
// - The group should have at least two all-reduces.
// - They should be located next to each other in the parent block.
// - They should all have the same element type.
// - They should all have the same group assignment.
//
// The merged all-reduce operates on a 1D tensor, whose size is the sum of all
// merged all-reduce tensors padded to 1024B. (The padding is necessary for the
// XLA delinearization skipping logic.) Each to-be-merged all-reduce flattens
// its input tensor and writes the resulting 1D tensor into the corresponding
// offset in the merged 1D tensor. After the merged all-reduce is done, the
// reverse happens: results are sliced out and reshaped to the original shape.
mlir::LogicalResult MergeAllReduceGroup(
std::vector<mlir::TF::DTensorAllReduceOp>& all_reduce_group) {
// Create the initial all-zero merged tensor.
// The merged tensor's size is the sum of all individual all-reduces' sizes.
int num_all_reduces = all_reduce_group.size();
DCHECK(num_all_reduces > 1)
<< "All reduce group size expected to be greater than 1.";
int total_num_elements = 0;
std::vector<llvm::ArrayRef<int64_t>> all_reduce_shapes;
all_reduce_shapes.reserve(num_all_reduces);
for (mlir::TF::DTensorAllReduceOp& all_reduce : all_reduce_group) {
auto all_reduce_ranked_type =
mlir::dyn_cast<mlir::RankedTensorType>(all_reduce.getType());
if (!all_reduce_ranked_type || !all_reduce_ranked_type.hasStaticShape()) {
return all_reduce.emitOpError(llvm::formatv(
"requires static shape for DTensorAllReduceOp, but got : {0}",
all_reduce_ranked_type));
}
int num_elements = all_reduce_ranked_type.getNumElements();
total_num_elements += num_elements;
all_reduce_shapes.push_back(all_reduce_ranked_type.getShape());
}
// Pad the merged tensor shape to multiples of 1024B, so delinearization
// skipping optimization in XLA can get activated.
if (total_num_elements % kAllReducePadding != 0) {
total_num_elements =
total_num_elements / kAllReducePadding * kAllReducePadding +
kAllReducePadding;
}
// Fill the merged tensor with 0 initially.
mlir::OpBuilder builder(all_reduce_group[0]);
mlir::Location loc = all_reduce_group[0].getLoc();
mlir::Type elem_type = all_reduce_group[0].getType().getElementType();
auto zero_scalar = ops_util::CreateScalarConst(0, builder, loc);
auto zero_scalar_elem_type = mlir::TF::CastOp::create(
builder, loc, mlir::RankedTensorType::get({}, elem_type), zero_scalar);
auto merged = mlir::TF::FillOp::create(
builder, loc, ops_util::GetR1Const({total_num_elements}, builder, loc),
zero_scalar_elem_type);
// Store every all-reduce's input at an offset location in the merged tensor,
// as a 1D tensor.
int offset_num_elements = 0;
std::vector<mlir::Type> flattened_types;
flattened_types.reserve(num_all_reduces);
mlir::Value updated;
for (int i = 0; i < all_reduce_group.size(); ++i) {
mlir::TF::DTensorAllReduceOp& all_reduce = all_reduce_group[i];
mlir::Location loc = all_reduce.getLoc();
auto all_reduce_ranked_type =
mlir::dyn_cast<mlir::RankedTensorType>(all_reduce.getType());
if (!all_reduce_ranked_type || !all_reduce_ranked_type.hasStaticShape()) {
return all_reduce.emitOpError(llvm::formatv(
"requires static shape for DTensorAllReduceOp, but got : {0}",
all_reduce_ranked_type));
}
int num_elements = all_reduce_ranked_type.getNumElements();
auto flattened = mlir::TF::ReshapeOp::create(
builder, DT_LOC2(loc, "CombinedReduceFlatten"), all_reduce.getInput(),
ops_util::GetR1Const({num_elements}, builder, loc));
flattened_types.push_back(flattened.getType());
auto indices = ops_util::GetR1Const({offset_num_elements}, builder, loc);
if (all_reduce.getDeviceType().contains("TPU")) {
updated = mlir::TF::XlaDynamicUpdateSliceOp::create(
builder, DT_LOC2(loc, "CombinedReduceUpdateSlice"), merged.getType(),
/*input=*/i == 0 ? merged.getResult() : updated,
/*update=*/flattened, indices);
} else {
auto end = ops_util::GetR1Const({offset_num_elements + num_elements},
builder, loc);
auto strides = ops_util::GetR1Const({1}, builder, loc);
updated = mlir::TF::TensorStridedSliceUpdateOp::create(
builder, DT_LOC2(loc, "CombinedReduceUpdateSlice"), merged.getType(),
/*input=*/i == 0 ? merged.getResult() : updated, indices, end,
strides,
/*value=*/flattened);
}
offset_num_elements += num_elements;
}
// All-reduce the updated merged tensor.
auto merged_all_reduce = mlir::TF::DTensorAllReduceOp::create(
builder, all_reduce_group[0].getLoc(), updated.getType(), updated,
all_reduce_group[0].getGroupAssignment(),
all_reduce_group[0].getReduceOp(), all_reduce_group[0].getDeviceType());
SetSingleLayoutOnOp(
merged_all_reduce,
ExtractSingleLayoutFromOp(all_reduce_group[0]).value().value());
// Slice out the original all-reduces, and reshape back to the original shape.
offset_num_elements = 0;
std::vector<mlir::TF::ReshapeOp> replacements;
replacements.reserve(num_all_reduces);
for (int i = 0; i < all_reduce_group.size(); ++i) {
mlir::TF::DTensorAllReduceOp& all_reduce = all_reduce_group[i];
mlir::Location loc = all_reduce.getLoc();
auto all_reduce_ranked_type =
mlir::dyn_cast<mlir::RankedTensorType>(all_reduce.getType());
if (!all_reduce_ranked_type || !all_reduce_ranked_type.hasStaticShape()) {
return all_reduce.emitOpError(llvm::formatv(
"requires static shape for DTensorAllReduceOp, but got : {0}",
all_reduce_ranked_type));
}
int num_elements = all_reduce_ranked_type.getNumElements();
auto slice = mlir::TF::SliceOp::create(
builder, DT_LOC2(loc, "PostCombinedReduceSlice"), flattened_types[i],
/*input=*/merged_all_reduce,
/*begin=*/ops_util::GetR1Const({offset_num_elements}, builder, loc),
/*size=*/ops_util::GetR1Const({num_elements}, builder, loc));
auto replacement = mlir::TF::ReshapeOp::create(
builder, DT_LOC2(loc, "PostCombinedReduceReshape"), slice.getResult(),
ops_util::GetR1Const(all_reduce_shapes[i], builder, loc));
replacements.push_back(replacement);
offset_num_elements += num_elements;
}
// Replace usages and clean up.
for (int i = 0; i < all_reduce_group.size(); ++i) {
mlir::TF::DTensorAllReduceOp& all_reduce = all_reduce_group[i];
mlir::TF::ReshapeOp& replacement = replacements[i];
all_reduce.replaceAllUsesWith(replacement.getResult());
all_reduce.erase();
}
return mlir::success();
}
// Dump the dependencies between AllReduce ops as a DOT graph.
std::string DrawAllReduceDependencies(
std::vector<mlir::TF::DTensorAllReduceOp> all_reduces,
const mlir::TF::detail::SideEffectAnalysisInfo& info) {
std::vector<std::vector<int>> dependents(all_reduces.size(),
std::vector<int>());
for (int j = 0; j < all_reduces.size(); ++j) {
mlir::TF::DTensorAllReduceOp later = all_reduces[j];
for (int i = 0; i < j; ++i) {
mlir::TF::DTensorAllReduceOp earlier = all_reduces[i];
DCHECK(!DependsOn(earlier, later, info));
if (earlier->getBlock() != later->getBlock() ||
DependsOn(later, earlier, info)) {
dependents[i].push_back(j);
}
}
}
std::string output = "digraph all_reduces {\n";
for (int i = 0; i < dependents.size(); i++) {
absl::StrAppend(&output, i);
absl::StrAppend(&output, "\n");
}
for (int i = 0; i < dependents.size(); i++) {
for (int j : dependents[i]) {
absl::StrAppend(&output, i, " -> ", j, "\n");
}
}
output += "}";
return output;
}
// Combine cross-slice DTensorAllReduce ops of the same element type and group
// assignment into as few groups as possible. Only independent ops can be
// combined together.
//
// For example, this program:
//
// clang-format off
// NOLINTBEGIN(whitespace/line_length)
// %0 = "tf_device.cluster"() ({
// %1 = "tf.Const"() {value = dense<0.0> : tensor<4x4xf32>} : () -> tensor<4x4xf32>
// %2 = "tf.Const"() {value = dense<[[0, 1], [2, 3]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
// %3 = "tf.DTensorAllReduce"(%1, %2) {reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
// %4 = "tf.Const"() {value = dense<0.0> : tensor<4x4xf32>} : () -> tensor<4x4xf32>
// %5 = "tf.Const"() {value = dense<[[0, 1], [2, 3]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
// %6 = "tf.DTensorAllReduce"(%4, %5) {reduce_op = "Add"} : (tensor<4x4xf32>, tensor<2x2xi32>) -> tensor<4x4xf32>
// %7 = "tf.Add"(%3, %6) : (tensor<4x4xf32>, tensor<4x4xf32>) -> tensor<4x4xf32>
// "tf_device.return"(%7) : (tensor<4x4xf32>) -> ()
// }) : () -> tensor<4x4xf32>
// NOLINTEND
// clang-format on
//
// will become this:
//
// clang-format off
// NOLINTBEGIN(whitespace/line_length)
// %0 = "tf_device.cluster"() ( {
// %cst = "tf.Const"() {value = dense<0.000000e+00> : tensor<4x4xf32>} : () -> tensor<4x4xf32>
// %cst_0 = "tf.Const"() {value = dense<[[0, 1], [2, 3]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
// %cst_1 = "tf.Const"() {value = dense<0.000000e+00> : tensor<4x4xf32>} : () -> tensor<4x4xf32>
// %cst_2 = "tf.Const"() {value = dense<[[0, 1], [2, 3]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
// %cst_3 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
// %1 = "tf.Cast"(%cst_3) {Truncate = false} : (tensor<i32>) -> tensor<f32>
// %cst_4 = "tf.Const"() {value = dense<1024> : tensor<1xi32>} : () -> tensor<1xi32>
// %2 = "tf.Fill"(%cst_4, %1) : (tensor<1xi32>, tensor<f32>) -> tensor<1024xf32>
// %cst_5 = "tf.Const"() {value = dense<16> : tensor<1xi32>} : () -> tensor<1xi32>
// %3 = "tf.Reshape"(%cst, %cst_5) : (tensor<4x4xf32>, tensor<1xi32>) -> tensor<16xf32>
// %cst_6 = "tf.Const"() {value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32>
// %4 = "tf.XlaDynamicUpdateSlice"(%2, %3, %cst_6) : (tensor<1024xf32>, tensor<16xf32>, tensor<1xi32>) -> tensor<1024xf32>
// %cst_7 = "tf.Const"() {value = dense<16> : tensor<1xi32>} : () -> tensor<1xi32>
// %5 = "tf.Reshape"(%cst_1, %cst_7) : (tensor<4x4xf32>, tensor<1xi32>) -> tensor<16xf32>
// %cst_8 = "tf.Const"() {value = dense<16> : tensor<1xi32>} : () -> tensor<1xi32>
// %6 = "tf.XlaDynamicUpdateSlice"(%4, %5, %cst_8) : (tensor<1024xf32>, tensor<16xf32>, tensor<1xi32>) -> tensor<1024xf32>
// %7 = "tf.DTensorAllReduce"(%6, %cst_0) {reduce_op = "Add"} : (tensor<1024xf32>, tensor<2x2xi32>) -> tensor<1024xf32>
// %cst_9 = "tf.Const"() {value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32>
// %cst_10 = "tf.Const"() {value = dense<16> : tensor<1xi32>} : () -> tensor<1xi32>
// %8 = "tf.Slice"(%7, %cst_9, %cst_10) : (tensor<1024xf32>, tensor<1xi32>, tensor<1xi32>) -> tensor<16xf32>
// %cst_11 = "tf.Const"() {value = dense<4> : tensor<2xi32>} : () -> tensor<2xi32>
// %9 = "tf.Reshape"(%8, %cst_11) : (tensor<16xf32>, tensor<2xi32>) -> tensor<4x4xf32>
// %cst_12 = "tf.Const"() {value = dense<16> : tensor<1xi32>} : () -> tensor<1xi32>
// %cst_13 = "tf.Const"() {value = dense<16> : tensor<1xi32>} : () -> tensor<1xi32>
// %10 = "tf.Slice"(%7, %cst_12, %cst_13) : (tensor<1024xf32>, tensor<1xi32>, tensor<1xi32>) -> tensor<16xf32>
// %cst_14 = "tf.Const"() {value = dense<4> : tensor<2xi32>} : () -> tensor<2xi32>
// %11 = "tf.Reshape"(%10, %cst_14) : (tensor<16xf32>, tensor<2xi32>) -> tensor<4x4xf32>
// %12 = "tf.Add"(%9, %11) : (tensor<4x4xf32>, tensor<4x4xf32>) -> tensor<4x4xf32>
// tf_device.return %12 : tensor<4x4xf32>
// }) : () -> tensor<4x4xf32>
// NOLINTEND
// clang-format on
mlir::LogicalResult CombineAllReduceOps(
mlir::tf_device::ClusterOp cluster,
std::vector<mlir::TF::DTensorAllReduceOp>& all_reduces) {
// A single op has nothing to combine with.
int num_all_reduces = all_reduces.size();
if (num_all_reduces <= 1) return mlir::success();
// Move all-reduces in the same group together and combine them.
auto& all_reduce_group = all_reduces;
mlir::TF::DTensorAllReduceOp final_all_reduce =
all_reduce_group[num_all_reduces - 1];
for (int i = num_all_reduces - 2; i >= 0; --i) {
mlir::TF::DTensorAllReduceOp all_reduce = all_reduce_group[i];
all_reduce->moveBefore(final_all_reduce);
}
auto merge_result = MergeAllReduceGroup(all_reduce_group);
if (merge_result.failed()) return merge_result;
return mlir::success();
}
// Returns true if both group assignments are constant and equal.
bool same_group_assignments(mlir::Value group_assignment_a,
mlir::Value group_assignment_b) {
if (group_assignment_a == group_assignment_b) {
return true;
}
mlir::DenseIntElementsAttr attr_a;
if (!matchPattern(group_assignment_a, m_Constant(&attr_a))) {
return false;
}
mlir::DenseIntElementsAttr attr_b;
if (!matchPattern(group_assignment_b, m_Constant(&attr_b))) {
return false;
}
if (attr_a.getType().getShape() != attr_b.getType().getShape()) {
return false;
}
// Group assignment should not be empty.
DCHECK(!attr_a.empty() && !attr_b.empty());
return std::equal(attr_a.begin(), attr_a.end(), attr_b.begin(), attr_b.end());
}
std::vector<std::vector<mlir::TF::DTensorAllReduceOp>>
createIndependentReduceOpsGroups(
const std::vector<mlir::TF::DTensorAllReduceOp>& ordered_all_reduces,
const mlir::TF::detail::SideEffectAnalysisInfo& info) {
// Build a reverse adjacency matrix from node to its dependents.
std::vector<std::vector<int>> dependents(ordered_all_reduces.size(),
std::vector<int>());
auto num_all_reduces = ordered_all_reduces.size();
for (int i = 0; i < num_all_reduces - 1; ++i) {
mlir::TF::DTensorAllReduceOp requirement = ordered_all_reduces[i];
for (int j = i + 1; j < num_all_reduces; ++j) {
mlir::TF::DTensorAllReduceOp dependent = ordered_all_reduces[j];
DCHECK(!DependsOn(requirement, dependent,
info)); // guaranteed by program order
// In this example, all three DTensorAllReduce ops are independent
// from each other according to MLIR value use-def chains considered
// by DependsOn. However, moving all three to after the WhileRegion
// and combine them would break the program.
//
// %3 = tf.DTensorAllReduce(%1, %2)
// %4 = tf.WhileRegion(%1) ({
// ^bb0(%arg):
// %5 = tf.TooBool(%arg)
// tf.Yield(%5)
// }, {
// %6 = tf.DTensorAllReduce(%1, %2)
// tf.Yield(%5)
// })
// %7 = tf.DTensorAllReduce(%1, %2)
//
// Therefore, in addition to DependsOn, we also check if two
// DTensorAllReduceOps belong to different blocks. If they do, since
// they exist in the same ClusterOp, one or both of them must be
// inside a control flow region block. We treat them as if there is
// a dependency between them.
//
// In the example above, the second DTensorAllReduceOp would "depend
// on" the first one, and the third on the second. This effectively
// prevents any two DTensorAllReduce from merging together.
if (requirement->getBlock() != dependent->getBlock() ||
DependsOn(dependent, requirement, info)) {
dependents[i].push_back(j);
}
}
}
// Traverse the adjacency matrix layer by layer from last op to find
// combination groups.
std::vector<std::vector<mlir::TF::DTensorAllReduceOp>> all_reduce_groups;
std::set<int> fulfilled;
while (fulfilled.size() < ordered_all_reduces.size()) {
std::vector<mlir::TF::DTensorAllReduceOp> group;
std::vector<int64_t> group_ids;
for (int i = dependents.size() - 1; i >= 0; i--) {
if (fulfilled.count(i) > 0) continue; // Already added op
bool all_deps_added = true;
for (int j = dependents[i].size() - 1; j >= 0; j--) {
if (fulfilled.count(dependents[i][j]) == 0) {
all_deps_added = false;
break;
}
}
if (all_deps_added) {
// Node with no dependents/already captured dependents degrees.
group_ids.push_back(i);
}
}
std::sort(group_ids.begin(), group_ids.end(),
[](const int64_t lhs, const int64_t rhs) { return lhs < rhs; });
for (auto x : group_ids) {
group.push_back(ordered_all_reduces[x]);
fulfilled.insert(x);
}
all_reduce_groups.push_back(group);
}
// Export the all reduces as a DOT graph.
VLOG(4) << "Visualizing AllReduce dependencies:\n"
<< DrawAllReduceDependencies(ordered_all_reduces, info);
return all_reduce_groups;
}
std::vector<std::vector<mlir::TF::DTensorAllReduceOp>>
createSubgroupsByElemType(
std::vector<std::vector<mlir::TF::DTensorAllReduceOp>> all_reduce_groups) {
std::vector<std::vector<mlir::TF::DTensorAllReduceOp>> all_reduce_new_groups;
// Combine all-reduces of the same element type.
for (const auto& all_reduce_group : all_reduce_groups) {
llvm::DenseMap<mlir::Type, std::vector<mlir::TF::DTensorAllReduceOp>>
all_reduces_by_elem_type;
for (auto all_reduce : all_reduce_group) {
mlir::Type elem_type = all_reduce.getType().getElementType();
all_reduces_by_elem_type[elem_type].push_back(all_reduce);
}
for (const auto& elem_type_pair : all_reduces_by_elem_type) {
all_reduce_new_groups.push_back(elem_type_pair.second);
}
}
VLOG(4) << "current number of groups: " << all_reduce_new_groups.size()
<< " after grouping by element type.";
return all_reduce_new_groups;
}
std::vector<std::vector<mlir::TF::DTensorAllReduceOp>>
createSubgroupsByReductionAttr(
std::vector<std::vector<mlir::TF::DTensorAllReduceOp>> all_reduce_groups) {
std::vector<std::vector<mlir::TF::DTensorAllReduceOp>> all_reduce_new_groups;
// Combine all-reduces of the same reduction attribute.
for (const auto& all_reduce_group : all_reduce_groups) {
llvm::DenseMap<llvm::StringRef, std::vector<mlir::TF::DTensorAllReduceOp>>
all_reduces_by_attr_reduce_op;
for (mlir::TF::DTensorAllReduceOp all_reduce : all_reduce_group) {
llvm::StringRef attr_reduce_op = all_reduce.getReduceOp();
all_reduces_by_attr_reduce_op[attr_reduce_op].push_back(all_reduce);
}
for (const auto& all_reduces_for_reduce_op_attr :
all_reduces_by_attr_reduce_op) {
all_reduce_new_groups.push_back(all_reduces_for_reduce_op_attr.second);
}
}
VLOG(4) << "current number of groups: " << all_reduce_new_groups.size()
<< " after grouping by reduction attribute.";
return all_reduce_new_groups;
}
std::vector<std::vector<mlir::TF::DTensorAllReduceOp>>
createSubgroupsByGroupAssignment(
std::vector<std::vector<mlir::TF::DTensorAllReduceOp>> all_reduce_groups) {
std::vector<std::vector<mlir::TF::DTensorAllReduceOp>> all_reduce_new_groups;
// Combine all-reduces of the group assignment.
for (const auto& all_reduce_group : all_reduce_groups) {
std::vector<mlir::Value> group_assignments;
llvm::DenseMap<mlir::Value, std::vector<mlir::TF::DTensorAllReduceOp>>
all_reduces_by_group_assignment;
for (mlir::TF::DTensorAllReduceOp all_reduce : all_reduce_group) {
mlir::Value group_assignment = all_reduce.getGroupAssignment();
bool seen = false;
for (mlir::Value seen_group_assignment : group_assignments) {
if (same_group_assignments(group_assignment, seen_group_assignment)) {
group_assignment = seen_group_assignment;
seen = true;
break;
}
}
if (!seen) group_assignments.push_back(group_assignment);
all_reduces_by_group_assignment[group_assignment].push_back(all_reduce);
}
for (const auto& all_reduce_group_to_merge :
all_reduces_by_group_assignment) {
all_reduce_new_groups.push_back(all_reduce_group_to_merge.second);
}
}
VLOG(4) << "current number of groups: " << all_reduce_new_groups.size()
<< " after grouping by group assignment.";
return all_reduce_new_groups;
}
// Experimental extended grouping logics to avoid aggressive grouping.
// This function performs the same grouping method as tf.distribute, which group
// all reduce ops by user defined group size (number of ops) in the input order.
// Note that group_size will be in range of [0, INT_MAX]. It is advised to pick
// a value based on knowledge of the total number of AllReduces. When group_size
// is too big, the function will act as aggressive grouping. When group_size is
// too small, the function will act as having no extended grouping.
std::vector<std::vector<mlir::TF::DTensorAllReduceOp>>
createSubgroupsByExtendedNumOps(
std::vector<std::vector<mlir::TF::DTensorAllReduceOp>> all_reduce_groups,
int group_size) {
VLOG(4) << "max number of ops in a all-reduce group: " << group_size;
// Disable extended grouping if group size is set to zero
if (group_size <= 0) return all_reduce_groups;
std::vector<std::vector<mlir::TF::DTensorAllReduceOp>> all_reduce_new_groups;
// Further break down the current all_reduced_groups by extended group size
for (const auto& all_reduce_group : all_reduce_groups) {
if (all_reduce_group.size() <= group_size) {
all_reduce_new_groups.push_back(all_reduce_group);
continue;
}
// Safe to "assume" num_groups would be greater or equal to two, because the
// above condition check guarantees case of zero or one would be skipped
int num_groups = (all_reduce_group.size() + group_size - 1) / group_size;
VLOG(4) << all_reduce_group.size() << " all_reduce ops in the current group"
<< ", able to split into " << num_groups << " groups\n";
for (int i = 0; i < num_groups - 1; i++) {
all_reduce_new_groups.push_back(std::vector<mlir::TF::DTensorAllReduceOp>(
all_reduce_group.begin() + i * group_size,
all_reduce_group.begin() + (i + 1) * group_size));
}
// Handle the last sub-group
all_reduce_new_groups.push_back(std::vector<mlir::TF::DTensorAllReduceOp>(
all_reduce_group.begin() + (num_groups - 1) * group_size,
all_reduce_group.end()));
}
VLOG(4) << "current number of groups: " << all_reduce_new_groups.size()
<< " after grouping by extended num ops size.";
return all_reduce_new_groups;
}
// Experimental grouping logics to optimize from aggressive grouping.
// This function first sort by topological level, then create AllReduce sub-
// groups by accessing each topological distance from its previous AllReduce.
// Note that topo_dist will be in range of [0, INT_MAX]. It is advised to select
// a value based on knowledge of the compute graph, such as the minimum distance
// between two model layers. When topo_dist is too big, the function will act
// as aggressive grouping. When topo_dist is too small, the function will act as
// having no extended grouping.
StatusOr<std::vector<std::vector<mlir::TF::DTensorAllReduceOp>>>
createSubgroupsByTopoDist(
std::vector<std::vector<mlir::TF::DTensorAllReduceOp>> all_reduce_groups,
llvm::DenseMap<mlir::TF::DTensorAllReduceOp, int> all_reduce_topo,
int topo_dist) {
// Disable extended grouping if topological distance is set to zero or less
if (topo_dist <= 0) return all_reduce_groups;
std::vector<std::vector<mlir::TF::DTensorAllReduceOp>> all_reduce_new_groups;
// Further break down the current all_reduced_groups by topological distance
// between two ops
for (auto& all_reduce_group : all_reduce_groups) {
std::vector<mlir::TF::DTensorAllReduceOp> new_group;
absl::Status status = absl::OkStatus();
// Sort AllReduces by topological level as the input order may not reflect
// their dependencies on the operands in the compute graph.
std::sort(all_reduce_group.begin(), all_reduce_group.end(),
[&all_reduce_topo, &status](mlir::TF::DTensorAllReduceOp& lhs,
mlir::TF::DTensorAllReduceOp& rhs) {
if ((all_reduce_topo.find(lhs) == all_reduce_topo.end()) ||
(all_reduce_topo.find(rhs) == all_reduce_topo.end())) {
status = absl::InternalError(
"Error: encounter AllReduce op with no topological level"
" assignment.");
return false;
}
return all_reduce_topo[lhs] < all_reduce_topo[rhs];
});
// Unable to sort AllReduces based on topological level due to error. Return
// directly as we are not able to group based on incorrect/partial topology.
if (!status.ok()) return status;
// Form AllReduce groups based on the topological distance between ops
DCHECK(!all_reduce_group.empty());
int prev_topo_level = all_reduce_topo[all_reduce_group[0]];
for (const auto& all_reduce : all_reduce_group) {
DCHECK(all_reduce_topo.find(all_reduce) != all_reduce_topo.end());
int cur_topo_level = all_reduce_topo[all_reduce];
if (abs(cur_topo_level - prev_topo_level) <= topo_dist) {
new_group.push_back(all_reduce);
} else {
// Start a new group
all_reduce_new_groups.push_back(
std::vector<mlir::TF::DTensorAllReduceOp>(new_group.begin(),
new_group.end()));
new_group.clear();
new_group.push_back(all_reduce);
}
prev_topo_level = cur_topo_level;
}
all_reduce_new_groups.push_back(new_group);
}
VLOG(4) << "current number of groups: " << all_reduce_new_groups.size()
<< " after grouping by topological distance.";
return all_reduce_new_groups;
}
// Compute the topological level for each AllReduce op in a cluster. The level
// is defined as 1 + max operands' depth in the compute graph. If an op do not
// depend on any input/operand, then it is level 0.
llvm::DenseMap<mlir::TF::DTensorAllReduceOp, int> computeAllReduceTopoLevel(
mlir::tf_device::ClusterOp cluster) {
llvm::DenseMap<mlir::Operation*, int> op_topo_level;
llvm::DenseMap<mlir::TF::DTensorAllReduceOp, int> all_reduce_topo;
// Compute topological level for each op.
cluster.getBody().walk([&](mlir::Operation* op) {
int max_depth = 0;
for (mlir::Value operand : op->getOperands()) {
if (mlir::Operation* operand_op = operand.getDefiningOp()) {
if (op_topo_level.find(operand_op) != op_topo_level.end()) {
max_depth = fmax(max_depth, op_topo_level[operand_op]);
}
}
}
op_topo_level[op] = max_depth + 1;
// Save the AllReduce topological level
mlir::TF::DTensorAllReduceOp all_reduce =
llvm::dyn_cast<mlir::TF::DTensorAllReduceOp>(op);
if (all_reduce && !all_reduce.getDeviceType().contains("TPU")) {
all_reduce_topo[all_reduce] = op_topo_level[op];
}
});
return all_reduce_topo;
}
struct DTensorAllReduceCombineOptimization
: public impl::DTensorAllReduceCombineOptimizationBase<
DTensorAllReduceCombineOptimization> {
void runOnOperation() override {
mlir::func::FuncOp function = getOperation();
auto module = function->getParentOfType<mlir::ModuleOp>();
function.walk([&](mlir::tf_device::ClusterOp cluster) {
std::vector<mlir::TF::DTensorAllReduceOp> ordered_all_reduces;
std::vector<mlir::Block*> ordered_blocks;
llvm::DenseSet<mlir::Block*> blocks;
cluster.GetBody().walk([&](mlir::TF::DTensorAllReduceOp all_reduce) {
if (!all_reduce.getDeviceType().contains("TPU")) {
// Only combine all reduces for GPU and CPU
mlir::RankedTensorType all_reduce_ranked_type =
mlir::dyn_cast<mlir::RankedTensorType>(all_reduce.getType());
if (all_reduce_ranked_type &&
all_reduce_ranked_type.hasStaticShape()) {
// Static known shape is required to merge all reduces. If shape is
// not known skip merging.
ordered_all_reduces.push_back(all_reduce);
blocks.insert(all_reduce->getBlock());
}
}
});
if (ordered_all_reduces.size() > 1) {
VLOG(2) << ordered_all_reduces.size()
<< " all-reduce ops eligible for combine optimization.";
// Build side effect analysis to identify indirect dependencies between
// all eligible all_reduce operations
mlir::TF::SideEffectAnalysis side_effect_analysis(module);
const mlir::TF::detail::SideEffectAnalysisInfo& info =
side_effect_analysis.GetAnalysisForFunc(function);
// Create dependency graph for all eligible all_reduce operations,
// so that independent ops can be merged
auto all_reduce_groups =
createIndependentReduceOpsGroups(ordered_all_reduces, info);
all_reduce_groups = createSubgroupsByElemType(all_reduce_groups);
all_reduce_groups = createSubgroupsByReductionAttr(all_reduce_groups);
all_reduce_groups = createSubgroupsByGroupAssignment(all_reduce_groups);
// Experimental extended grouping: topological distance
if (module->hasAttrOfType<mlir::IntegerAttr>(
kAllReduceTopologicalDistance)) {
llvm::DenseMap<mlir::TF::DTensorAllReduceOp, int> all_reduce_topo =
computeAllReduceTopoLevel(cluster);
StatusOr<std::vector<std::vector<mlir::TF::DTensorAllReduceOp>>>
group = createSubgroupsByTopoDist(
all_reduce_groups, all_reduce_topo,
module
->getAttrOfType<mlir::IntegerAttr>(
kAllReduceTopologicalDistance)
.getInt());
if (!group.ok()) {
// This is a non-fatal error since topological level distance is one
// of the optimizations in this combiner pass. Output an error and
// continue with the rest of the grouping optimization.
LOG(WARNING) << "Failed to create subgroups using topological "
<< "level distance: " << group.status();
} else {
all_reduce_groups = group.value();
}
}
// Experimental extended grouping: fixed number of AllReduce ops
if (module->hasAttrOfType<mlir::IntegerAttr>(kAllReduceNumOpsInGroup)) {
all_reduce_groups = createSubgroupsByExtendedNumOps(
all_reduce_groups,
module->getAttrOfType<mlir::IntegerAttr>(kAllReduceNumOpsInGroup)
.getInt());
}
// Maintain relative order of ALLReduces within the block.
std::sort(all_reduce_groups.begin(), all_reduce_groups.end(),
[](std::vector<mlir::TF::DTensorAllReduceOp> lhs,
std::vector<mlir::TF::DTensorAllReduceOp> rhs) {
// Prefer groups that are not empty.
if (lhs.empty() && !rhs.empty()) return false;
if (!lhs.empty() && rhs.empty()) return true;
// Then prefer groups that are in earlier-in-memory blocks,
// this part just needs to be consistent for strict weak
// ordering purposes.
if (lhs[0]->getBlock() != rhs[0]->getBlock()) {
return lhs[0]->getBlock() < rhs[0]->getBlock();
}
// Within the block, use the group's actual sorting.
return lhs[0]->isBeforeInBlock(rhs[0]);
});
VLOG(2) << ordered_all_reduces.size() << " all-reduce ops in "
<< all_reduce_groups.size() << " groups";
for (auto& reduce_group : all_reduce_groups) {
if (reduce_group.size() > 1) {
VLOG(4) << "Combining following reduce ops into one: ------------";
for (auto reduce_op : reduce_group) {
VLOG(4) << mlir::GetNameFromLoc(reduce_op.getLoc());
}
VLOG(4) << "-----------------------------------------------------";
}
if (mlir::failed(CombineAllReduceOps(cluster, reduce_group))) {
return signalPassFailure();
}
}
for (auto* b : blocks) {
mlir::sortTopologically(b);
}
}
});
}
};
} // namespace
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateDTensorAllReduceCombineOptimization() {
return std::make_unique<DTensorAllReduceCombineOptimization>();
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/dtensor_allreduce_scatter_optimization.cc
+189
View File
@@ -0,0 +1,189 @@
/* Copyright 2022 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 <algorithm>
#include <cstdint>
#include <memory>
#include <string>
#include <vector>
#include "absl/container/flat_hash_set.h"
#include "absl/log/log.h"
#include "absl/log/vlog_is_on.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/Matchers.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/UseDefLists.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 "tensorflow/core/platform/types.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/collectives_common.h"
#include "tensorflow/dtensor/mlir/ir/tf_dtensor.h"
#include "tensorflow/dtensor/mlir/layout_parsing.h"
namespace tensorflow {
namespace dtensor {
namespace {
#define GEN_PASS_DEF_DTENSORALLREDUCESCATTEROPTIMIZATION
#include "tensorflow/dtensor/mlir/dtensor_passes.h.inc"
// Returns true if both group assignments are constant and equal.
bool same_group_assignments(mlir::DenseIntElementsAttr attr_a,
mlir::DenseIntElementsAttr attr_b) {
if (attr_a.getType().getShape() != attr_b.getType().getShape()) {
return false;
}
return std::equal(attr_a.begin(), attr_a.end(), attr_b.begin(), attr_b.end());
}
mlir::DenseIntElementsAttr GetScatterGroupAssignment(
mlir::TF::DTensorAllScatterOp all_scatter, int scatter_dim) {
const Layout original_layout = all_scatter.getInputLayout();
const Layout desired_layout = all_scatter.getOutputLayout();
absl::flat_hash_set<std::string> scattered_dims;
scattered_dims.insert(desired_layout.sharding_spec(scatter_dim));
auto partitions =
GetAllReducePartitionsFromReducedDims(original_layout, scattered_dims)
.value();
const int32_t num_partitions = partitions.size();
// Construct a flattened list of scatter partitions.
std::vector<int32_t> partitions_flat;
for (auto& p : partitions) {
partitions_flat.insert(partitions_flat.end(), p.second.begin(),
p.second.end());
}
int32_t partition_size = partitions.begin()->second.size();
mlir::OpBuilder builder(all_scatter);
auto group_shaped_type = mlir::RankedTensorType::get(
{num_partitions, partition_size},
mlir::IntegerType::get(builder.getContext(), 32));
return mlir::DenseIntElementsAttr::get(group_shaped_type, partitions_flat);
}
mlir::LogicalResult ApplyOptimization(mlir::func::FuncOp function) {
std::vector<mlir::Operation*> ops_to_delete;
function.walk([&](mlir::TF::DTensorAllReduceOp all_reduce) {
if (all_reduce->hasOneUse()) {
if (auto all_scatter = mlir::dyn_cast<mlir::TF::DTensorAllScatterOp>(
*all_reduce->getUsers().begin())) {
VLOG(2) << "Found potential AllReduce+AllScatter to fuse.";
if (VLOG_IS_ON(2)) all_reduce.dump();
if (VLOG_IS_ON(2)) all_scatter.dump();
const Layout original_layout = all_scatter.getInputLayout();
const Layout desired_layout = all_scatter.getOutputLayout();
// Find all potential scatter dimensions.
std::vector<int> scatter_dims;
for (int i = 0; i < original_layout.rank(); ++i) {
if (original_layout.sharding_spec(i) !=
desired_layout.sharding_spec(i)) {
scatter_dims.push_back(i);
}
}
if (scatter_dims.empty()) return mlir::WalkResult::advance();
if (scatter_dims.size() > 1) {
VLOG(2) << "Multiple dimensions are scatter. This is unsupported "
"for AllReduce+Scatter fusion.";
return mlir::WalkResult::advance();
}
int scatter_dim = scatter_dims[0];
VLOG(2) << "Scatter_dim: " << scatter_dim;
// Check that the all-reduce and all-scatter group assignments are the
// same.
mlir::DenseIntElementsAttr all_reduce_group_assignment_attr;
if (!matchPattern(all_reduce.getGroupAssignment(),
m_Constant(&all_reduce_group_assignment_attr))) {
all_reduce.emitOpError("group_assignment should be a constant");
return mlir::WalkResult::interrupt();
}
mlir::DenseIntElementsAttr all_scatter_group_assignment_attr =
GetScatterGroupAssignment(all_scatter, scatter_dim);
VLOG(2) << "All scatter group assignment: ";
if (VLOG_IS_ON(2)) all_scatter_group_assignment_attr.dump();
bool same_group =
same_group_assignments(all_reduce_group_assignment_attr,
all_scatter_group_assignment_attr);
if (!same_group) return mlir::WalkResult::advance();
VLOG(2) << "Fuse reduce scatter with scatter_dim: " << scatter_dim;
mlir::OpBuilder builder(all_reduce);
auto scatter_dim_const_op = mlir::TF::ConstOp::create(
builder, all_reduce.getLoc(),
mlir::DenseIntElementsAttr::get(
mlir::RankedTensorType::get({}, builder.getI32Type()),
{scatter_dim}));
auto reduce_scatter = mlir::TF::DTensorReduceScatterOp::create(
builder, all_reduce.getLoc(), all_scatter->getResultTypes(),
all_reduce.getOperand(0), all_reduce.getGroupAssignment(),
scatter_dim_const_op, all_reduce.getReduceOp(),
all_reduce.getDeviceType());
SetSingleLayoutOnOp(reduce_scatter, desired_layout);
all_scatter->replaceAllUsesWith(reduce_scatter);
ops_to_delete.push_back(all_scatter);
ops_to_delete.push_back(all_reduce);
}
}
return mlir::WalkResult::advance();
});
for (mlir::Operation* op : ops_to_delete) {
op->erase();
}
return mlir::success();
}
// MLIR pass that combines AllReduce and AllScatter to ReduceScatter.
struct DTensorAllReduceScatterOptimization
: public impl::DTensorAllReduceScatterOptimizationBase<
DTensorAllReduceScatterOptimization> {
void runOnOperation() override {
mlir::func::FuncOp function = getOperation();
if (mlir::failed(ApplyOptimization(function))) return signalPassFailure();
}
};
} // namespace
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateDTensorAllReduceScatterOptimization() {
return std::make_unique<DTensorAllReduceScatterOptimization>();
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/dtensor_allreduce_sum_optimization.cc
+529
View File
@@ -0,0 +1,529 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <memory>
#include <string>
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/FormatVariadic.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Matchers.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/UseDefLists.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 "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/dtensor/mlir/ir/tf_dtensor.h"
#include "tensorflow/dtensor/mlir/layout_parsing.h"
#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
namespace tensorflow {
namespace dtensor {
namespace {
#define GEN_PASS_DEF_DTENSORALLREDUCESUMOPTIMIZATION
#include "tensorflow/dtensor/mlir/dtensor_passes.h.inc"
constexpr int kMaxIteration = 10;
mlir::Value GetIdentitySkippedInputs(mlir::Value val) {
mlir::Value input = val;
while (auto identity = llvm::dyn_cast_or_null<mlir::TF::IdentityOp>(
input.getDefiningOp())) {
input = identity.getInput();
}
return input;
}
bool IsZeroConstant(mlir::Value val) {
auto const_input = llvm::dyn_cast_or_null<mlir::TF::ConstOp>(
GetIdentitySkippedInputs(val).getDefiningOp());
if (!const_input) return false;
mlir::DenseFPElementsAttr attr =
mlir::dyn_cast<mlir::DenseFPElementsAttr>(const_input.getValue());
// This uses the fact that constant Attrs becomes splats, so we only need to
// check one value.
if (!attr || !attr.isSplat()) return false;
return attr.getSplatValue<mlir::FloatAttr>().getValue().isZero();
}
// Extracts inputs/ops required for optimization and checks whether graph
// meets the criteria for reduction + sum optimization. The criterion are:
// a) All DTensorAllReduce operations must be sum operations.
// b) Group assignment of DTensorAllReduceOp must be the same
// c) All operands of Add op must be DTensorAllReduce operations.
mlir::LogicalResult CheckReduceAndSumOptimizationCriteria(
mlir::Operation* add_op,
llvm::SmallVectorImpl<mlir::Value>* reduction_inputs,
llvm::SmallVectorImpl<mlir::TF::DTensorAllReduceOp>* reduction_ops,
bool* can_be_reordered) {
for (mlir::Value operand : add_op->getOperands()) {
if (IsZeroConstant(operand)) {
reduction_inputs->emplace_back(operand);
continue;
}
auto reduction_op = llvm::dyn_cast_or_null<mlir::TF::DTensorAllReduceOp>(
operand.getDefiningOp());
if (!reduction_op) {
*can_be_reordered = false;
return mlir::success();
}
reduction_ops->emplace_back(reduction_op);
}
llvm::SmallDenseSet<mlir::Attribute> reduction_group_assignments;
for (mlir::TF::DTensorAllReduceOp reduction : *reduction_ops) {
if (reduction.getReduceOp().str() != kReduceOpAdd) {
*can_be_reordered = false;
return mlir::success();
}
mlir::DenseIntElementsAttr group_assignment;
if (!matchPattern(reduction.getGroupAssignment(),
m_Constant(&group_assignment))) {
*can_be_reordered = false;
return mlir::success();
}
reduction_group_assignments.insert(group_assignment);
reduction_inputs->emplace_back(reduction.getInput());
}
*can_be_reordered = (reduction_group_assignments.size() == 1);
return mlir::success();
}
// Applies optimization that reorders AllReduce + Add operations.
// For example:
// %3 = DTensorAllReduce(%0)
// %4 = DTensorAllReduce(%1)
// %5 = Add(%3, %4)
//
// Is transformed to:
// %2 = Add(%0, %1)
// %3 = DTensorAllReduce(%2)
//
// Therefore reducing the number of Reduction/cross device communication.
mlir::LogicalResult OptimizeAllReduceAndSum(mlir::Operation* op,
bool* changed) {
bool can_be_reordered;
llvm::SmallVector<mlir::TF::DTensorAllReduceOp, 4> reduction_ops;
llvm::SmallVector<mlir::Value, 4> reduction_op_inputs;
if (mlir::failed(CheckReduceAndSumOptimizationCriteria(
op, &reduction_op_inputs, &reduction_ops, &can_be_reordered)))
return mlir::failure();
if (!can_be_reordered || reduction_ops.empty()) return mlir::success();
// Forward the inputs from the DTensorAllReduce to the add op. Calling
// getOperand(i).getDefiningOp() since CheckReduceAndSumOptimizationCriteria
// checks that each input is fed by a DTensorAllReduce or a Zero constant.
for (int i = 0; i < op->getNumOperands(); ++i) {
if (mlir::isa<mlir::TF::DTensorAllReduceOp>(
op->getOperand(i).getDefiningOp()))
op->setOperand(i, op->getOperand(i).getDefiningOp()->getOperand(0));
}
mlir::TF::DTensorAllReduceOp first_reduction_op = reduction_ops.front();
// Invoke reduction operation on locally added tensor once.
// From above check `CheckOptimizationCriteria()`, we know that all reduction
// operations that are fused reused the same group assignment value.
// 1) Get mlir::Value that represents group assignment used for reduction.
mlir::Value group_assignment = first_reduction_op.getGroupAssignment();
// Create a singe reduction operation that reduces the result of the locally
// added tensor.
mlir::OpBuilder builder(op);
builder.setInsertionPointAfterValue(op->getResult(0));
mlir::TF::DTensorAllReduceOp all_reduce =
mlir::TF::DTensorAllReduceOp::create(
builder, op->getLoc(), op->getResult(0).getType(), op->getResult(0),
group_assignment, builder.getStringAttr(std::string(kReduceOpAdd)),
builder.getStringAttr(first_reduction_op.getDeviceType()));
const auto layout_or_status = ExtractSingleLayoutFromOp(first_reduction_op);
if (!layout_or_status.ok())
return first_reduction_op->emitOpError(llvm::formatv(
"Malformed layout specification for DTensorAllReduce op found: {0}",
layout_or_status.status().message()));
if (!layout_or_status->has_value())
return first_reduction_op->emitOpError(
"DTensorAllReduce op must have layout specification.");
// Set target layout that is equivalent to original DTensorReduction op in
// the graph. This is used during later optimization passes.
SetSingleLayoutOnOp(all_reduce, layout_or_status->value());
// Replace usages of original tf.Add op with newly created output of
// `all_reduce`.
op->getResult(0).replaceAllUsesExcept(
all_reduce.getOutput(),
llvm::SmallPtrSet<mlir::Operation*, 1>{all_reduce.getOperation()});
// TODO(hongjunchoi, bfontain): Consider adding optimization for the case when
// `tree` of Add operations with DTensorAllReduce op as inputs exists.
// Remove original tf.Add `op` and if reduction operation inputs to original
// `op` is only used by the `op`, then remove the DTensorAllReduce op as well.
for (mlir::Operation* original_reduction_op : reduction_ops) {
if (original_reduction_op->use_empty()) original_reduction_op->erase();
}
*changed = true;
return mlir::success();
}
mlir::Value SkipIdentityLikeOpsOutputs(mlir::Value val) {
while (val.hasOneUse() &&
llvm::isa<mlir::TF::CastOp, mlir::TF::ReshapeOp, mlir::TF::IdentityOp>(
*val.user_begin())) {
val = val.user_begin()->getResult(0);
}
return val;
}
// TODO(hongjunchoi): Consider using tracing algorithm to virtually transform
// the IR and only apply optimizations when total number of DTensorAllReduce in
// the graph is reduced.
bool MayRemoveAllReduce(mlir::Operation* op) {
mlir::Value op_output = op->getResult(0);
mlir::Value value_after_identity_like_ops =
SkipIdentityLikeOpsOutputs(op_output);
if (value_after_identity_like_ops.hasOneUse() &&
llvm::isa<mlir::TF::AddNOp, mlir::TF::AddV2Op, mlir::TF::AddOp>(
*value_after_identity_like_ops.user_begin()))
return true;
return false;
}
// Moves DTensorAllReduce ops after IdentityLike Operations if the operation is
// connected to Add operation which may led to optimization.
// For example:
//
// %0 = "tf.Const"() {value = dense<0> : tensor<2x64xi32>}
// %2 = "tf.Const"() {value = dense<0.0> : tensor<8192x916xbf16>}
// %4= "tf.DTensorAllReduce"(%2, %0) {reduce_op = "Add"}
// %5 = "tf.Cast"(%4){Truncate = false, device = ""}
// %6 = "tf.Identity"(%5){Truncate = false, device = ""}
// %7 = "tf.Const"() {value = dense<[916,8192]> : tensor<2xi32>}
// %8 = "tf.Reshape"(%6, %7)
//
// Becomes :
//
// %0 = "tf.Const"()
// %2 = "tf.Const"()
// %3 = "tf.Cast"(%2)
// %4 = "tf.Identity"(%3)
// %7 = "tf.Const"()
// %8 = "tf.Reshape"(%4, %7)
// %9 = "tf.DTensorAllReduce"(%8, %0) {reduce_op = "Add"}
void OptimizeIdentityLikeOps(mlir::Operation* op, bool* changed) {
auto dtensor_all_reduce =
llvm::dyn_cast_or_null<mlir::TF::DTensorAllReduceOp>(
op->getOperand(0).getDefiningOp());
if (!dtensor_all_reduce) return;
// TODO(hongjunchoi, bfontain): Consider allowing pushing DTensorAllReduce op
// with multiple usages if it can lead to performance optimization.
if (!dtensor_all_reduce->hasOneUse()) return;
if (!MayRemoveAllReduce(op)) return;
dtensor_all_reduce->moveAfter(op);
mlir::Value input = dtensor_all_reduce.getInput();
op->setOperand(0, input);
mlir::Value op_output = op->getResult(0);
dtensor_all_reduce.setOperand(0, op_output);
dtensor_all_reduce.getInput().setType(
mlir::cast<mlir::TensorType>(op_output.getType()));
dtensor_all_reduce.getOutput().setType(
mlir::cast<mlir::TensorType>(op_output.getType()));
llvm::SmallPtrSet<mlir::Operation*, 4> exceptions{dtensor_all_reduce};
op_output.replaceAllUsesExcept(dtensor_all_reduce.getOutput(), exceptions);
*changed = true;
}
bool CheckWhileLoopOptimizationCriteria(
const int index, mlir::TF::WhileRegionOp while_op, mlir::Value while_output,
mlir::Operation** add_op, mlir::TF::DTensorAllReduceOp* all_reduce_op,
mlir::OpOperand** add_input) {
// Loop variant input that is being optimized should not be used in loop
// condition.
mlir::Value loop_condition_input = while_op.getCond().getArgument(index);
if (!loop_condition_input.use_empty()) return false;
// While loop output should be connected to add op.
// If operand to while loop body terminator if from Identity op,
// skip through the input identity operations.
mlir::Value output_value = GetIdentitySkippedInputs(while_output);
mlir::Operation* output_defining_op = output_value.getDefiningOp();
if (!output_defining_op) return false;
// TODO(hongjunchoi): Handle AddN op as well.
if (!output_defining_op ||
!llvm::isa<mlir::TF::AddV2Op, mlir::TF::AddOp>(output_defining_op)) {
return false;
}
// Input operand of add operation should be
// 1) DTensorAllReduce
// 2) from block argument of while loop
mlir::OpOperand& first_operand = output_defining_op->getOpOperand(0);
mlir::OpOperand& second_operand = output_defining_op->getOpOperand(1);
mlir::BlockArgument block_arg;
mlir::TF::DTensorAllReduceOp all_reduce =
llvm::dyn_cast_or_null<mlir::TF::DTensorAllReduceOp>(
first_operand.get().getDefiningOp());
if (all_reduce) {
block_arg = mlir::dyn_cast<mlir::BlockArgument>(second_operand.get());
*add_input = &second_operand;
} else {
all_reduce = llvm::dyn_cast_or_null<mlir::TF::DTensorAllReduceOp>(
second_operand.get().getDefiningOp());
block_arg = mlir::dyn_cast<mlir::BlockArgument>(first_operand.get());
*add_input = &first_operand;
}
if (!block_arg || !all_reduce) return false;
// DTensorAllReduce should calculate sum across devices and group assignment
// must be statically known.
mlir::Operation* group_assignment =
all_reduce.getGroupAssignment().getDefiningOp();
if (!group_assignment || !llvm::isa<mlir::TF::ConstOp>(group_assignment))
return false;
if (all_reduce.getReduceOp().str() != kReduceOpAdd) return false;
// While loop block argument input connected to Add op should be
// connected to constant operations with zero value.
const int block_arg_index = block_arg.getArgNumber();
mlir::OpOperand& while_input = while_op->getOpOperand(block_arg_index);
if (!IsZeroConstant(while_input.get())) return false;
// TODO(hongjunchoi): Handle the case when input is from DTensorAllReduce op.
// If group assignment is the same, then the input DTensorAllReduce op can
// also be optimized away.
*add_op = output_defining_op;
*all_reduce_op = all_reduce;
return true;
}
// Extracts out DTensorAllReduce operation from while op if
// a) While op contains DTensorAllReduce op followed by an Add Operation
// b) Remaining operand of Add operation is a loop variant input of the while
// operation with zero initial value.
//
// For example:
//
// %0 = "tf.Const"() {value = dense<0> : tensor<2x64xi32>}
// %2 = "tf.Const"() {value = dense<0.0> : tensor<8192x916xbf16>}
// WhileRegionOp(%2) {
// %0 = "tf.A"(%2)
// "tf.Yield"(%0)
// }, {
// ^bb0(%barg0: tensor<8192x916xbf16>):
// ...
// %0 = "tf.Const"()
// %1 = "tf.Const"()
// %2 = "tf.DTensorAllReduce"(%1, %0) {reduce_op = "Add"}
// %3 = "tf.Add"(%2, %barg0)
// "tf.Yield"(%3)
// })
//
// Becomes :
//
// %0 = "tf.Const"() {value = dense<0> : tensor<2x64xi32>}
// %2 = "tf.Const"() {value = dense<0.0> : tensor<8192x916xbf16>}
// %4 = WhileRegionOp(%2) {
// %0 = "tf.A"(%2)
// "tf.Yield"(%0)
// }, {
// ^bb0(%barg0: tensor<8192x916xbf16>):
// ...
// %0 = "tf.Const"()
// %1 = "tf.Const"()
// %3 = "tf.Add"(%1, %barg0)
// "tf.Yield"(%3)
// })
// "tf.DTensorAllReduce"(%4, %0) {reduce_op = "Add"}
mlir::LogicalResult ExtractAllReduceFromWhileOp(
const int output_index, mlir::TF::DTensorAllReduceOp all_reduce,
mlir::TF::WhileRegionOp while_op, mlir::OpOperand& add_input,
mlir::Operation* add_op, bool* changed) {
// Set add input to input of all reduce.
mlir::Value all_reduce_input = all_reduce.getInput();
const int replacement_add_input_index =
add_input.getOperandNumber() == 0 ? 1 : 0;
add_op->setOperand(replacement_add_input_index, all_reduce_input);
mlir::OpBuilder builder(while_op);
builder.setInsertionPointAfter(while_op);
mlir::Value while_output = while_op.getResult(output_index);
mlir::Operation* group_assignment_const =
all_reduce.getGroupAssignment().getDefiningOp();
mlir::Operation* cloned_group_assignment =
builder.clone(*group_assignment_const);
// Create a singe reduction operation that reduces the result of the locally
// added tensor.
auto new_all_reduce = mlir::TF::DTensorAllReduceOp::create(
builder, all_reduce.getLoc(), while_output.getType(), while_output,
cloned_group_assignment->getResult(0),
builder.getStringAttr(std::string(kReduceOpAdd)),
builder.getStringAttr(all_reduce.getDeviceType()));
const auto layout_or_status = ExtractSingleLayoutFromOp(all_reduce);
if (!layout_or_status.ok())
return all_reduce->emitOpError(llvm::formatv(
"Malformed layout specification for DTensorAllReduce op found: {0}",
layout_or_status.status().message()));
if (!layout_or_status->has_value())
return all_reduce->emitOpError(
"DTensorAllReduce op must have layout specification.");
// Set target layout that is equivalent to original DTensorReduction op in
// the graph. This is used during later optimization passes.
SetSingleLayoutOnOp(new_all_reduce, layout_or_status->value());
llvm::SmallPtrSet<mlir::Operation*, 4> exceptions;
exceptions.insert(new_all_reduce.getOperation());
while_output.replaceAllUsesExcept(new_all_reduce.getOutput(), exceptions);
if (all_reduce.use_empty()) all_reduce.erase();
*changed = true;
return mlir::success();
}
mlir::LogicalResult OptimizeWhileLoopLazyAllReduce(
mlir::TF::WhileRegionOp while_op, bool* changed) {
mlir::Operation* while_body_terminator =
while_op.getBody().front().getTerminator();
for (const auto& data :
llvm::enumerate(while_body_terminator->getOpOperands())) {
const int index = data.index();
mlir::OpOperand& operand = data.value();
mlir::Operation* add_op = nullptr;
mlir::TF::DTensorAllReduceOp all_reduce;
mlir::OpOperand* add_input = nullptr;
if (!CheckWhileLoopOptimizationCriteria(index, while_op, operand.get(),
&add_op, &all_reduce, &add_input))
continue;
// Perform while loop lazy all reduce optimization.
if (mlir::failed(ExtractAllReduceFromWhileOp(index, all_reduce, while_op,
*add_input, add_op, changed)))
return mlir::failure();
}
return mlir::success();
}
mlir::LogicalResult ApplyOptimization(
mlir::func::FuncOp function,
const llvm::SmallVectorImpl<mlir::Operation*>& identity_like_ops,
const llvm::SmallVectorImpl<mlir::TF::WhileRegionOp>& while_ops,
const llvm::SmallVectorImpl<mlir::Operation*>& add_ops, bool* changed) {
// Collect and fold the reduction operations within the function.
for (mlir::Operation* add_op : add_ops)
if (mlir::failed(OptimizeAllReduceAndSum(add_op, changed)))
return mlir::failure();
for (mlir::Operation* op : identity_like_ops)
OptimizeIdentityLikeOps(op, changed);
for (mlir::TF::WhileRegionOp op : while_ops)
if (mlir::failed(OptimizeWhileLoopLazyAllReduce(op, changed)))
return mlir::failure();
return mlir::success();
}
// Finds all potential ops that could lead to all reduce optimizations. Those
// are:
// a) Identity like ops (e.g. Identity/Reshape/Cast) ops.
// b) WhileRegion op
// c) Add operations.
void CollectOptimizationCandidates(
mlir::func::FuncOp func,
llvm::SmallVectorImpl<mlir::Operation*>* identity_like_ops,
llvm::SmallVectorImpl<mlir::Operation*>* add_ops,
llvm::SmallVectorImpl<mlir::TF::WhileRegionOp>* while_ops) {
func.walk([&](mlir::Operation* op) {
if (llvm::isa<mlir::TF::IdentityOp, mlir::TF::CastOp, mlir::TF::ReshapeOp>(
op))
identity_like_ops->emplace_back(op);
if (auto while_op = llvm::dyn_cast<mlir::TF::WhileRegionOp>(op))
while_ops->emplace_back(while_op);
if (llvm::isa<mlir::TF::AddOp, mlir::TF::AddV2Op, mlir::TF::AddNOp>(op))
add_ops->emplace_back(op);
});
}
// MLIR pass that folds constants that can be removed or deduplicated away.
struct DTensorAllReduceSumOptimization
: public impl::DTensorAllReduceSumOptimizationBase<
DTensorAllReduceSumOptimization> {
void runOnOperation() override {
mlir::func::FuncOp function = getOperation();
bool changed = true;
int iteration = 0;
llvm::SmallVector<mlir::Operation*, 4> identity_like_ops;
llvm::SmallVector<mlir::Operation*, 4> add_ops;
llvm::SmallVector<mlir::TF::WhileRegionOp, 4> while_ops;
CollectOptimizationCandidates(function, &identity_like_ops, &add_ops,
&while_ops);
bool is_optimized = false;
while (changed && iteration < kMaxIteration) {
changed = false;
if (mlir::failed(ApplyOptimization(function, identity_like_ops, while_ops,
add_ops, &changed)))
return signalPassFailure();
iteration++;
if (changed) is_optimized = true;
}
}
};
} // namespace
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateDTensorAllReduceSumOptimization() {
return std::make_unique<DTensorAllReduceSumOptimization>();
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/dtensor_collective_type_lowering.cc
+339
View File
@@ -0,0 +1,339 @@
/* Copyright 2022 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 "llvm/ADT/SmallVector.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#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/MLIRContext.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/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 "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/ir/tf_dtensor.h"
#include "tensorflow/dtensor/mlir/layout_parsing.h"
#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
namespace tensorflow {
namespace dtensor {
namespace {
#define GEN_PASS_DEF_DTENSORCOLLECTIVETYPELOWERINGPASS
#include "tensorflow/dtensor/mlir/dtensor_passes.h.inc"
mlir::LogicalResult WrapOpWithCasts(const mlir::RankedTensorType& input_type,
const mlir::RankedTensorType& output_type,
mlir::Operation* reduce_op) {
mlir::OpBuilder builder(reduce_op);
auto intermediate_type = mlir::RankedTensorType::get(
output_type.getShape(), input_type.getElementType());
const mlir::Location loc = reduce_op->getLoc();
mlir::TF::CastOp cast_to_long = mlir::TF::CastOp::create(
builder, loc, input_type, reduce_op->getOperand(0));
reduce_op->setOperand(0, cast_to_long.getY());
reduce_op->getResult(0).setType(intermediate_type);
mlir::Value result = reduce_op->getResult(0);
builder.setInsertionPointAfter(reduce_op);
mlir::TF::CastOp cast_to_original =
mlir::TF::CastOp::create(builder, loc, output_type, result);
StatusOr<Layout> result_layout =
ExtractRequiredSingleLayoutFromOp(result.getDefiningOp());
if (!result_layout.ok()) {
return reduce_op->emitOpError(result_layout.status().message());
}
SetSingleLayoutOnOp(cast_to_original, *result_layout);
reduce_op->getResult(0).replaceAllUsesExcept(cast_to_original.getY(),
cast_to_original);
return mlir::success();
}
template <class ReduceOpType>
mlir::LogicalResult ConvertShortIntReduce(ReduceOpType reduce_op) {
mlir::OpBuilder builder(reduce_op);
StatusOr<Layout> output_layout = ExtractRequiredSingleLayoutFromOp(reduce_op);
if (!output_layout.ok()) {
return reduce_op.emitOpError(output_layout.status().message());
}
const mlir::Type output_type = reduce_op.getResult().getType();
const mlir::Type input_type = reduce_op.getOperand(0).getType();
// Handle bools by first casting to int32 and swapping All/Any for Min/Max.
const mlir::TensorType& tensor_input_type =
mlir::dyn_cast<mlir::TensorType>(input_type);
const mlir::TensorType& tensor_output_type =
mlir::dyn_cast<mlir::TensorType>(output_type);
if (!tensor_input_type) return mlir::success();
if (!tensor_output_type) return mlir::success();
if (tensor_input_type.getElementType().isInteger(1)) {
if (reduce_op.getReduceOpAttr().getValue().str() == kReduceOpAll)
reduce_op.setReduceOpAttr(
builder.getStringAttr(std::string(kReduceOpMin)));
else if (reduce_op.getReduceOpAttr().getValue().str() == kReduceOpAny)
reduce_op.setReduceOpAttr(
builder.getStringAttr(std::string(kReduceOpMax)));
else if (reduce_op.getReduceOpAttr().getValue().str() != kReduceOpMax &&
reduce_op.getReduceOpAttr().getValue().str() != kReduceOpMin)
return reduce_op.emitOpError()
<< "reduce for boolean only supports 'All'/'Min' or 'Any'/'Max' "
"reduction. "
<< "Received '" << reduce_op.getReduceOpAttr().getValue().str()
<< "'";
}
if (auto integer_type = mlir::dyn_cast<mlir::IntegerType>(
tensor_input_type.getElementType())) {
int32_t min_width = 64;
if (output_layout->mesh().is_tpu_mesh()) {
min_width = 32;
}
if (integer_type.getWidth() >= min_width) {
return mlir::success();
}
auto input_type = mlir::RankedTensorType::get(
tensor_input_type.getShape(), builder.getIntegerType(min_width));
auto output_type = mlir::RankedTensorType::get(
tensor_output_type.getShape(), integer_type);
return WrapOpWithCasts(input_type, output_type, reduce_op);
}
if (mlir::isa<mlir::BFloat16Type>(tensor_input_type.getElementType())) {
if (output_layout->mesh().is_tpu_mesh()) {
return mlir::success();
}
auto input_type = mlir::RankedTensorType::get(tensor_input_type.getShape(),
builder.getF32Type());
auto output_type = mlir::RankedTensorType::get(
tensor_output_type.getShape(), tensor_input_type.getElementType());
return WrapOpWithCasts(input_type, output_type, reduce_op);
}
return mlir::success();
}
// Complex for AllReduce and ReduceScatter
template <class ReduceOpType>
mlir::LogicalResult ConvertComplexReduce(ReduceOpType reduce_op) {
ReduceOpType real_reduce_op;
ReduceOpType imag_reduce_op;
mlir::OpBuilder builder(reduce_op);
StatusOr<Layout> output_layout = ExtractRequiredSingleLayoutFromOp(reduce_op);
if (!output_layout.ok()) {
return reduce_op.emitOpError(output_layout.status().message());
}
const mlir::Value tensor_input = reduce_op.getInput();
const mlir::Value tensor_result = reduce_op.getResult();
const mlir::TensorType complex_input_tensor_type =
mlir::dyn_cast<mlir::TensorType>(tensor_input.getType());
if (!complex_input_tensor_type) {
return mlir::success();
}
const mlir::TensorType complex_result_tensor_type =
mlir::dyn_cast<mlir::TensorType>(tensor_result.getType());
if (!complex_result_tensor_type) {
return mlir::success();
}
auto input_element_type = mlir::dyn_cast<mlir::ComplexType>(
complex_input_tensor_type.getElementType());
if (!input_element_type) {
return mlir::success();
}
auto real_input_tensor_type =
mlir::RankedTensorType::get(complex_input_tensor_type.getShape(),
input_element_type.getElementType());
auto real_result_tensor_type =
mlir::RankedTensorType::get(complex_result_tensor_type.getShape(),
input_element_type.getElementType());
const mlir::Value tensor_temp_real = mlir::TF::RealOp::create(
builder, reduce_op.getLoc(), real_input_tensor_type, tensor_input);
const mlir::Value tensor_temp_imag = mlir::TF::ImagOp::create(
builder, reduce_op.getLoc(), real_input_tensor_type, tensor_input);
real_reduce_op = mlir::dyn_cast<ReduceOpType>(builder.clone(*reduce_op));
real_reduce_op->setOperand(0, tensor_temp_real);
real_reduce_op->getResult(0).setType(real_result_tensor_type);
imag_reduce_op = mlir::dyn_cast<ReduceOpType>(builder.clone(*reduce_op));
imag_reduce_op->setOperand(0, tensor_temp_imag);
imag_reduce_op->getResult(0).setType(real_result_tensor_type);
const mlir::Type output_type = reduce_op.getResult().getType();
auto complex_reduce_op = mlir::TF::ComplexOp::create(
builder, reduce_op->getLoc(), output_type, real_reduce_op.getResult(),
imag_reduce_op.getResult());
StatusOr<Layout> desired_layout =
ExtractRequiredSingleLayoutFromOp(reduce_op);
SetSingleLayoutOnOp(complex_reduce_op, *desired_layout);
reduce_op.getOutput().replaceAllUsesWith(complex_reduce_op.getResult());
reduce_op.erase();
return mlir::success();
}
// Complex for AllToAll, AllGather, and AllScatter
template <class CollectiveType>
mlir::LogicalResult ConvertComplexCollectives(CollectiveType op) {
CollectiveType real_op;
CollectiveType imag_op;
mlir::OpBuilder builder(op);
StatusOr<Layout> output_layout = ExtractRequiredSingleLayoutFromOp(op);
if (!output_layout.ok()) {
return op.emitOpError(output_layout.status().message());
}
const mlir::Value tensor_input = op.getInput();
const mlir::Value tensor_result = op.getResult();
const mlir::TensorType complex_input_tensor_type =
mlir::dyn_cast<mlir::TensorType>(tensor_input.getType());
if (!complex_input_tensor_type) {
return mlir::success();
}
const mlir::TensorType& complex_result_tensor_type =
mlir::dyn_cast<mlir::TensorType>(tensor_result.getType());
if (!complex_result_tensor_type) {
return mlir::success();
}
auto input_element_type = mlir::dyn_cast<mlir::ComplexType>(
complex_input_tensor_type.getElementType());
if (!input_element_type) {
return mlir::success();
}
auto real_input_tensor_type =
mlir::RankedTensorType::get(complex_input_tensor_type.getShape(),
input_element_type.getElementType());
auto real_result_tensor_type =
mlir::RankedTensorType::get(complex_result_tensor_type.getShape(),
input_element_type.getElementType());
const mlir::Value tensor_temp_real = mlir::TF::RealOp::create(
builder, op.getLoc(), real_input_tensor_type, tensor_input);
const mlir::Value tensor_temp_imag = mlir::TF::ImagOp::create(
builder, op.getLoc(), real_input_tensor_type, tensor_input);
real_op = mlir::dyn_cast<CollectiveType>(builder.clone(*op));
real_op->setOperand(0, tensor_temp_real);
real_op->getResult(0).setType(real_result_tensor_type);
imag_op = mlir::dyn_cast<CollectiveType>(builder.clone(*op));
imag_op->setOperand(0, tensor_temp_imag);
imag_op->getResult(0).setType(real_result_tensor_type);
const mlir::Type output_type = op.getResult().getType();
auto complex_op =
mlir::TF::ComplexOp::create(builder, op.getLoc(), output_type,
real_op.getResult(), imag_op.getResult());
const Layout desired_layout = op.getOutputLayout();
SetSingleLayoutOnOp(complex_op, desired_layout);
op.getOutput().replaceAllUsesWith(complex_op.getResult());
op.erase();
return mlir::success();
}
// A Walk that allows mutation inside parent.
template <typename FuncT, typename OpT = mlir::detail::first_argument<FuncT>>
mlir::LogicalResult MutatingWalk(mlir::Operation* parent, FuncT func) {
llvm::SmallVector<OpT, 4> ops;
parent->walk([&](OpT op) { ops.push_back(op); });
for (auto op : ops) {
if (mlir::failed(func(op))) {
return mlir::LogicalResult::failure();
}
}
return mlir::LogicalResult::success();
}
class DTensorCollectiveTypeLoweringPass
: public impl::DTensorCollectiveTypeLoweringPassBase<
DTensorCollectiveTypeLoweringPass> {
public:
void runOnOperation() override {
mlir::func::FuncOp func = getOperation();
if (mlir::failed(
MutatingWalk(func, [&](mlir::TF::DTensorAllReduceOp all_reduce) {
// Lower integer type all reduce
return ConvertComplexReduce(all_reduce);
}))) {
signalPassFailure();
}
if (mlir::failed(
MutatingWalk(func, [&](mlir::TF::DTensorAllScatterOp all_scatter) {
// Lower complex type all scatter
return ConvertComplexCollectives(all_scatter);
}))) {
signalPassFailure();
}
if (mlir::failed(
MutatingWalk(func, [&](mlir::TF::DTensorAllGatherOp all_gather) {
// Lower complex type all gather.
return ConvertComplexCollectives(all_gather);
}))) {
signalPassFailure();
}
if (mlir::failed(
MutatingWalk(func, [&](mlir::TF::DTensorAllToAllOp all_to_all) {
// Lower complex type all to all
return ConvertComplexCollectives(all_to_all);
}))) {
signalPassFailure();
}
if (mlir::failed(MutatingWalk(
func, [&](mlir::TF::DTensorReduceScatterOp reduce_scatter) {
// Lower complex type reduce scatter.
return ConvertComplexReduce(reduce_scatter);
}))) {
signalPassFailure();
}
if (mlir::failed(
MutatingWalk(func, [&](mlir::TF::DTensorAllReduceOp all_reduce) {
// Lower integer type all reduce
return ConvertShortIntReduce(all_reduce);
}))) {
signalPassFailure();
}
if (mlir::failed(MutatingWalk(
func, [&](mlir::TF::DTensorReduceScatterOp reduce_scatter) {
// Lower integer type all reduce
return ConvertShortIntReduce(reduce_scatter);
}))) {
signalPassFailure();
}
}
};
} // namespace
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateDTensorCollectiveTypeLoweringPass() {
return std::make_unique<DTensorCollectiveTypeLoweringPass>();
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/dtensor_dialect/BUILD
+80
View File
@@ -0,0 +1,80 @@
load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
# DTensor MLIR dialect.
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow/dtensor:dtensor-internal",
# Allow visibility from the mlir language server.
"//learning/brain/mlir/mlir_lsp_server:__pkg__",
],
licenses = ["notice"],
)
# ODS (https://mlir.llvm.org/docs/OpDefinitions/) generation for op and dialect files to include.
td_library(
name = "dtensor_td_files",
srcs = [
"ir/dtensor_dialect.td",
"ir/dtensor_ops.td",
],
compatible_with = get_compatible_with_portable(),
deps = [
"@llvm-project//mlir:FuncTdFiles",
"@llvm-project//mlir:InferTypeOpInterfaceTdFiles",
"@llvm-project//mlir:OpBaseTdFiles",
],
)
gentbl_cc_library(
name = "DialectIncGen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {
"ir/ops.h.inc": ["-gen-op-decls"],
"ir/ops.cc.inc": ["-gen-op-defs"],
"ir/dialect.h.inc": ["-gen-dialect-decls"],
"ir/dialect.cc.inc": ["-gen-dialect-defs"],
},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "ir/dtensor_ops.td",
deps = [":dtensor_td_files"],
)
cc_library(
name = "Dialect",
srcs = ["ir/ops.cc"],
hdrs = [
"ir/dialect.h",
"ir/ops.h",
],
deps = [
":DialectIncGen",
":ir/dtensor_attributes",
"//tensorflow/dtensor/cc:dstatus",
"//tensorflow/dtensor/cc:tensor_layout",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:InferTypeOpInterface",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "ir/dtensor_attributes",
srcs = ["ir/dtensor_attributes.cc"],
hdrs = [
"ir/dialect.h",
"ir/dtensor_attributes.h",
],
deps = [
":DialectIncGen",
"//tensorflow/dtensor/cc:tensor_layout",
"//tensorflow/dtensor/proto:layout_proto_cc",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
)
tensorflow/dtensor/mlir/dtensor_dialect/ir/dialect.h
+37
View File
@@ -0,0 +1,37 @@
/* Copyright 2022 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_DTENSOR_MLIR_DTENSOR_DIALECT_IR_DIALECT_H_
#define TENSORFLOW_DTENSOR_MLIR_DTENSOR_DIALECT_IR_DIALECT_H_
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Dialect.h" // from @llvm-project
// Dialect main class is defined in ODS, we include it here. The
// constructor and the printing/parsing of dialect types are manually
// implemented (see ops.cpp).
#include "tensorflow/dtensor/mlir/dtensor_dialect/ir/dialect.h.inc"
namespace mlir {
namespace dtensor {
//===----------------------------------------------------------------------===//
// DTENSOR dialect types.
//===----------------------------------------------------------------------===//
} // namespace dtensor
} // namespace mlir
#endif // TENSORFLOW_DTENSOR_MLIR_DTENSOR_DIALECT_IR_DIALECT_H_
tensorflow/dtensor/mlir/dtensor_dialect/ir/dtensor_attributes.cc
+97
View File
@@ -0,0 +1,97 @@
/* Copyright 2022 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/dtensor/mlir/dtensor_dialect/ir/dtensor_attributes.h"
#include <utility>
#include "llvm/ADT/Hashing.h"
#include "mlir/IR/AttributeSupport.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/dtensor_dialect/ir/dialect.h"
namespace mlir {
namespace dtensor {
// Storage class for MeshAttr.
namespace detail {
struct MeshAttrStorage : public AttributeStorage {
using Mesh = tensorflow::dtensor::Mesh;
using KeyTy = Mesh;
explicit MeshAttrStorage(Mesh mesh) : mesh(std::move(mesh)) {}
bool operator==(const KeyTy& key) const { return key == KeyTy(mesh); }
static llvm::hash_code hashKey(const KeyTy& key) {
const Mesh& mesh = key;
return llvm::hash_value(mesh.ToString());
}
static MeshAttrStorage* construct(mlir::AttributeStorageAllocator& allocator,
const KeyTy& key) {
return new (allocator.allocate<MeshAttrStorage>()) MeshAttrStorage(key);
}
Mesh mesh;
};
} // namespace detail
MeshAttr MeshAttr::get(MLIRContext* context, const Mesh& mesh) {
return Base::get(context, mesh);
}
const MeshAttr::Mesh& MeshAttr::getValue() const { return getImpl()->mesh; }
// The storage class for LayoutAttr.
namespace detail {
struct LayoutAttrStorage : public AttributeStorage {
using Layout = tensorflow::dtensor::Layout;
using KeyTy = Layout;
explicit LayoutAttrStorage(Layout layout) : layout(std::move(layout)) {}
bool operator==(const KeyTy& key) const { return key == KeyTy(layout); }
static llvm::hash_code hashKey(const KeyTy& key) {
const Layout& layout = key;
return llvm::hash_value(layout.ToString());
}
static LayoutAttrStorage* construct(
mlir::AttributeStorageAllocator& allocator, const KeyTy& key) {
const Layout& layout = key;
return new (allocator.allocate<LayoutAttrStorage>())
LayoutAttrStorage(layout);
}
Layout layout;
};
} // namespace detail
LayoutAttr LayoutAttr::get(mlir::MLIRContext* context,
tensorflow::dtensor::Layout layout) {
return Base::get(context, std::move(layout));
}
const LayoutAttr::Layout& LayoutAttr::getValue() const {
return getImpl()->layout;
}
void DTensorDialect::registerAttributes() {
addAttributes<dtensor::MeshAttr, dtensor::LayoutAttr>();
}
} // namespace dtensor
} // namespace mlir
tensorflow/dtensor/mlir/dtensor_dialect/ir/dtensor_attributes.h
+71
View File
@@ -0,0 +1,71 @@
/* Copyright 2022 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 defines attributes for DTensor.
#ifndef TENSORFLOW_DTENSOR_MLIR_DTENSOR_DIALECT_IR_DTENSOR_ATTRIBUTES_H_
#define TENSORFLOW_DTENSOR_MLIR_DTENSOR_DIALECT_IR_DTENSOR_ATTRIBUTES_H_
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/proto/layout.pb.h"
namespace mlir {
namespace dtensor {
namespace detail {
struct LayoutAttrStorage;
struct MeshAttrStorage;
} // namespace detail
// Attribute to keep track of a mesh.
class MeshAttr
: public Attribute::AttrBase<MeshAttr, Attribute, detail::MeshAttrStorage> {
public:
using Base::Base;
using Mesh = tensorflow::dtensor::Mesh;
static constexpr StringLiteral name = "dtensor.mesh";
// Constructor of attribute
static MeshAttr get(MLIRContext* context, const Mesh& mesh);
// Returns Mesh
const Mesh& getValue() const;
};
// Custom attribute to keep track of dtensor layouts.
class LayoutAttr : public Attribute::AttrBase<LayoutAttr, Attribute,
detail::LayoutAttrStorage> {
public:
using Base::Base;
using Layout = tensorflow::dtensor::Layout;
using Mesh = tensorflow::dtensor::Mesh;
static constexpr StringLiteral name = "dtensor.layout";
// Create a layout attribute.
static LayoutAttr get(MLIRContext* context, Layout layout);
// Get layout.
const Layout& getValue() const;
};
} // namespace dtensor
} // namespace mlir
#endif // TENSORFLOW_DTENSOR_MLIR_DTENSOR_DIALECT_IR_DTENSOR_ATTRIBUTES_H_
tensorflow/dtensor/mlir/dtensor_dialect/ir/dtensor_dialect.td
+62
View File
@@ -0,0 +1,62 @@
/* Copyright 2022 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 DTENSOR_DIALECT
#define DTENSOR_DIALECT
include "mlir/IR/OpBase.td"
include "mlir/Interfaces/InferTypeOpInterface.td"
// ODS Definition for the dialect, see https://mlir.llvm.org/docs/OpDefinitions/
// for more information.
//===----------------------------------------------------------------------===//
// DTensor dialect definitions
//===----------------------------------------------------------------------===//
def DTensorDialect : Dialect {
let name = "dtensor";
let summary = "Dialect for dtensor";
let description = [{
This dialect contains key operations for DTensors. Currently it holds two
attributes; one that represents layouts and one for representing a mesh.
}];
let cppNamespace = "::mlir::dtensor";
let extraClassDeclaration = [{
// Parses an instance of an attribute registered to the dialect.
Attribute parseAttribute(DialectAsmParser &parser, Type type) const override;
// Prints an instance of an attribute registered to the dialect.
void printAttribute(Attribute attr, DialectAsmPrinter &os) const override;
private:
// Register the attributes of this dialect.
void registerAttributes();
public:
}];
}
//===----------------------------------------------------------------------===//
// DTENSOR dialect types definitions
//===----------------------------------------------------------------------===//
#endif // DTENSOR_DIALECT
tensorflow/dtensor/mlir/dtensor_dialect/ir/dtensor_ops.td
+30
View File
@@ -0,0 +1,30 @@
/* Copyright 2022 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.
==============================================================================*/
#ifdef DTENSOR_OPS
#else
#define DTENSOR_OPS
include "dtensor_dialect.td"
//===----------------------------------------------------------------------===//
// DTensor op definitions
//===----------------------------------------------------------------------===//
// Base class for the operations in this dialect.
class DTensor_Op<string mnemonic, list<Trait> traits = []> :
Op<DTensorDialect, mnemonic, traits>;
#endif // DTENSOR_OPS
tensorflow/dtensor/mlir/dtensor_dialect/ir/ops.cc
+161
View File
@@ -0,0 +1,161 @@
/* Copyright 2022 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/dtensor/mlir/dtensor_dialect/ir/ops.h"
#include <string>
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Diagnostics.h" // from @llvm-project
#include "mlir/IR/Dialect.h" // from @llvm-project
#include "mlir/IR/DialectImplementation.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/dtensor_dialect/ir/dialect.h"
// Generated dialect defs.
#include "tensorflow/dtensor/mlir/dtensor_dialect/ir/dialect.cc.inc"
#include "tensorflow/dtensor/mlir/dtensor_dialect/ir/dtensor_attributes.h"
namespace mlir {
namespace dtensor {
// Dialect construction: there is one instance per context and it registers its
// operations, types, and interfaces here.
void DTensorDialect::initialize() {
addOperations<
#define GET_OP_LIST
#include "tensorflow/dtensor/mlir/dtensor_dialect/ir/ops.cc.inc"
>();
registerAttributes();
}
// Parses a #dtensor.mesh attribute of the following format:
//
// #dtensor.mesh<serializedMesh>
//
// where the first element is a SymbolRefAttr and the second element is the
// location.
static MeshAttr ParseMeshAttr(MLIRContext *context, StringRef spec,
Location loc) {
// Define error function.
auto emit_error = [&](std::string text) {
emitError(loc, "invalid TensorFlow Mesh attribute ") << text;
return nullptr;
};
// Check correct format and consume prefix, otherwise throw error.
if (!spec.consume_front("mesh<"))
return emit_error("Unexpected start to mesh specification");
// Consume back from ">".
if (!spec.consume_back(">"))
return emit_error("Unexpected closing of mesh specification");
// Cast from StringRef to string.
std::string mesh_str = spec.str();
// Check if serializedMesh is correct.
using Mesh = tensorflow::dtensor::Mesh;
using MeshOr = tensorflow::dtensor::StatusOr<Mesh>;
MeshOr mesh_or = Mesh::FromString(mesh_str);
if (!mesh_or.ok()) {
std::string status_msg = mesh_or.status().ToString();
return emit_error("parsing serialized string. More details: " + status_msg);
}
return MeshAttr::get(context, mesh_or.value());
}
// Parses a #dtensor.layout attribute of the following format:
//
// #dtensor.layout<serializedLayout>
static LayoutAttr ParseLayoutAttr(MLIRContext *context, StringRef spec,
Location loc) {
// Define error function.
auto emit_error = [&](std::string text) {
emitError(loc, "invalid TensorFlow Mesh attribute ") << text;
return nullptr;
};
// Check correct format and consume prefix, otherwise throw error.
if (!spec.consume_front("layout<"))
return emit_error("Unexpected start to layout specification");
// Consume back from "\">".
if (!spec.consume_back(">"))
return emit_error("Unexpected closing of layout specification");
// Cast into string
std::string layout_str = spec.str();
// Check if serializedMesh is correct, else error from line 37.
using Layout = tensorflow::dtensor::Layout;
using LayoutOr = tensorflow::dtensor::StatusOr<Layout>;
LayoutOr layout_or = Layout::FromString(layout_str);
if (!layout_or.ok()) {
std::string status_msg = layout_or.status().ToString();
return emit_error("parsing serialized string. More details: " + status_msg);
}
// Extract layout.
Layout layout = layout_or.value();
return LayoutAttr::get(context, layout);
}
Attribute DTensorDialect::parseAttribute(DialectAsmParser &parser,
Type type) const {
StringRef spec = parser.getFullSymbolSpec();
Location loc = parser.getEncodedSourceLoc(parser.getNameLoc());
if (spec.starts_with("mesh")) return ParseMeshAttr(getContext(), spec, loc);
if (spec.starts_with("layout"))
return ParseLayoutAttr(getContext(), spec, loc);
return (emitError(loc, "unknown DTensor attribute: " + spec), nullptr);
}
// Print a type registered to this dialect.
// Prints a #dtensor.dtensor attribute of the following format:
//
// #dtensor.mesh<mesh>
static void printMeshAttr(MeshAttr attr, DialectAsmPrinter &os) {
os << "mesh<" << attr.getValue().ToString() << ">";
}
// Prints a #dtensor.dtensor attribute of the following format:
//
// #dtensor.layout<layout>
static void printLayoutAttr(LayoutAttr attr, DialectAsmPrinter &os) {
os << "layout<" << attr.getValue().ToString() << ">";
}
// Override general virtual function
void DTensorDialect::printAttribute(Attribute attr,
DialectAsmPrinter &os) const {
// Cast into correct attribute and print
if (auto mesh_attr = mlir::dyn_cast<MeshAttr>(attr))
printMeshAttr(mesh_attr, os);
if (auto layout_attr = mlir::dyn_cast<LayoutAttr>(attr))
printLayoutAttr(layout_attr, os);
}
} // namespace dtensor
} // namespace mlir
// Ops definition from ODS
#define GET_OP_CLASSES
#include "tensorflow/dtensor/mlir/dtensor_dialect/ir/ops.cc.inc"
tensorflow/dtensor/mlir/dtensor_dialect/ir/ops.h
+33
View File
@@ -0,0 +1,33 @@
/* Copyright 2022 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_DTENSOR_MLIR_DTENSOR_DIALECT_IR_OPS_H_
#define TENSORFLOW_DTENSOR_MLIR_DTENSOR_DIALECT_IR_OPS_H_
#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/Dialect.h" // from @llvm-project
#include "mlir/IR/Matchers.h" // from @llvm-project
#include "mlir/IR/OpImplementation.h" // from @llvm-project
#include "mlir/IR/TypeUtilities.h" // from @llvm-project
#include "mlir/Interfaces/InferTypeOpInterface.h" // from @llvm-project
#include "tensorflow/dtensor/mlir/dtensor_dialect/ir/dtensor_attributes.h"
#define GET_OP_CLASSES
#include "tensorflow/dtensor/mlir/dtensor_dialect/ir/ops.h.inc"
#endif // TENSORFLOW_DTENSOR_MLIR_DTENSOR_DIALECT_IR_OPS_H_
tensorflow/dtensor/mlir/dtensor_layout_to_xla_sharding_op.cc
+139
View File
@@ -0,0 +1,139 @@
/* 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 <utility>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "xla/xla_data.pb.h"
#include "tensorflow/dtensor/cc/xla_spmd/layout_to_xla_sharding.h"
#include "tensorflow/dtensor/mlir/ir/tf_dtensor.h"
namespace tensorflow {
namespace dtensor {
namespace {
#define GEN_PASS_DECL_DTENSORLAYOUTTOXLASHARDINGOPPASS
#define GEN_PASS_DEF_DTENSORLAYOUTTOXLASHARDINGOPPASS
#include "tensorflow/dtensor/mlir/dtensor_passes.h.inc"
using mlir::TF::DTensorLayout;
class RemoveDTensorLayoutAfterConstOrBlockArgPattern
: public mlir::OpRewritePattern<DTensorLayout> {
public:
using OpRewritePattern::OpRewritePattern;
mlir::LogicalResult matchAndRewrite(
DTensorLayout layout_op, mlir::PatternRewriter& rewriter) const override {
if (match(layout_op).failed()) {
return mlir::failure();
}
rewriter.replaceAllUsesWith(layout_op, layout_op.getInput());
rewriter.eraseOp(layout_op);
return mlir::success();
}
private:
mlir::LogicalResult match(DTensorLayout layout_op) const {
auto input = layout_op.getInput();
if (mlir::isa<mlir::BlockArgument>(input)) {
return mlir::success();
}
mlir::Operation* input_op = input.getDefiningOp();
if (input_op != nullptr) {
return mlir::success(input_op->hasTrait<mlir::OpTrait::ConstantLike>());
} else {
return layout_op->emitOpError() << "Can't find defining op for " << input;
}
}
};
class DTensorLayoutToXlaShardingOpPass
: public impl::DTensorLayoutToXlaShardingOpPassBase<
DTensorLayoutToXlaShardingOpPass> {
public:
void runOnOperation() override;
};
void DTensorLayoutToXlaShardingOpPass::runOnOperation() {
mlir::RewritePatternSet patterns(&getContext());
// Some patterns in tf2xla requires operands to be ConstantLike.
// Inserting tf.XlaSharding between them will fail the pattern match.
// We remove all tf.DTensorLayout after constants so no tf.XlaSharding is
// inserted in the above case. XLA will figure out the sharding of constants
// without DTensor guidance.
//
// For BlockArgument, the sharding is already attached to function attribute
// by DTensorSetHloShardingPass. No additional tf.XlaSharding is needed.
patterns.add<RemoveDTensorLayoutAfterConstOrBlockArgPattern>(&getContext());
if (mlir::failed(
mlir::applyPatternsGreedily(getOperation(), std::move(patterns)))) {
signalPassFailure();
}
auto result =
getOperation().walk([](DTensorLayout layout_op) -> mlir::WalkResult {
Layout layout = layout_op.getLayout();
StatusOr<xla::OpSharding> sharding =
ConvertLayoutToXlaOpSharding(layout);
if (!sharding.ok()) {
return layout_op.emitOpError()
<< "Failed to convert layout to sharding for "
<< layout.ToString() << ": " << sharding.status().message();
}
mlir::OpBuilder builder(layout_op);
auto sharding_attr =
builder.getStringAttr(sharding->SerializeAsString());
// TODO(b/414807890): It seems that the dtensor path later on clear up
// the V1 sharding attr, so set V2 sharding to "" here. It may be better
// to set the V2 sharding attr here and then removed it when V1 is
// removed.
auto sharding_op = mlir::TF::XlaShardingOp::create(
builder, layout_op.getLoc(), layout_op.getOutput().getType(),
layout_op.getInput(),
/*sharding=*/builder.getStringAttr(""), // Not used by tf2xla.
/*_xlaSharding=*/sharding_attr,
/*_XlaShardingV2=*/builder.getStringAttr(""));
layout_op.getOutput().replaceAllUsesWith(sharding_op);
layout_op.erase();
return mlir::WalkResult::advance();
});
if (result.wasInterrupted()) {
signalPassFailure();
}
}
} // namespace
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateDTensorLayoutToXlaShardingOpPass() {
return std::make_unique<DTensorLayoutToXlaShardingOpPass>();
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/dtensor_location.cc
+89
View File
@@ -0,0 +1,89 @@
/* Copyright 2022 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/dtensor/mlir/dtensor_location.h"
#include <algorithm>
#include <queue>
#include <string>
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/FormatVariadic.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/utils/name_utils.h"
namespace tensorflow {
namespace dtensor {
namespace {
std::string CreateLocalLocationString(mlir::FileLineColLoc loc) {
return llvm::formatv(">> {0}:{1}:{2}", loc.getFilename(), loc.getLine(),
loc.getColumn())
.str();
}
} // namespace
mlir::Location DTensorLocation(mlir::Location loc, llvm::StringRef file,
unsigned int line, llvm::StringRef name) {
// Strip dirname.
auto split = file.rsplit("/");
if (!split.second.empty()) file = split.second;
mlir::Location callee_loc =
mlir::FileLineColLoc::get(loc.getContext(), file, line, 0);
std::string new_name = GetNameFromLoc(loc);
if (!new_name.empty()) {
if (!name.empty()) {
new_name = llvm::formatv("{0}/{1}", new_name, name).str();
}
callee_loc = mlir::NameLoc::get(
mlir::StringAttr::get(loc.getContext(), new_name), callee_loc);
}
return mlir::CallSiteLoc::get(/*callee=*/callee_loc, /*caller=*/loc);
}
mlir::Location DTensorLocation(mlir::Operation* op, llvm::StringRef file,
unsigned int line, llvm::StringRef name) {
return DTensorLocation(op->getLoc(), file, line, name);
}
std::string DTensorLocationToString(mlir::Location loc) {
llvm::SmallVector<std::string, 4> stack;
std::queue<mlir::Location> queue;
queue.push(loc);
while (!queue.empty()) {
mlir::Location& front = queue.front();
if (auto name_loc = mlir::dyn_cast<mlir::NameLoc>(front)) {
queue.push(name_loc.getChildLoc());
} else if (auto callsite_loc = mlir::dyn_cast<mlir::CallSiteLoc>(front)) {
queue.push(callsite_loc.getCallee());
queue.push(callsite_loc.getCaller());
} else if (auto line_loc = mlir::dyn_cast<mlir::FileLineColLoc>(front)) {
stack.push_back(CreateLocalLocationString(line_loc));
}
queue.pop();
}
std::reverse(stack.begin(), stack.end());
std::string s;
llvm::raw_string_ostream ss(s);
llvm::interleave(stack, ss, "\n");
return ss.str();
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/dtensor_location.h
+61
View File
@@ -0,0 +1,61 @@
/* Copyright 2022 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_DTENSOR_MLIR_DTENSOR_LOCATION_H_
#define TENSORFLOW_DTENSOR_MLIR_DTENSOR_LOCATION_H_
#include <string>
#include "mlir/IR/Location.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
// mlir::Location utilities for DTensor. `DTensorLocation` augments a location
// object with the current file and line _of the C++ code creating an
// operation_. This simplifies tracking down the creator of an invalid operation
// while debugging.
namespace tensorflow {
namespace dtensor {
mlir::Location DTensorLocation(mlir::Location loc, llvm::StringRef file,
unsigned int line, llvm::StringRef name = "");
mlir::Location DTensorLocation(mlir::Operation* op, llvm::StringRef file,
unsigned int line, llvm::StringRef name = "");
// Creates a string from a location of the following format:
// >> pass_file_1:line1:col1
// >> pass_file_2:line2:col2
//
// DTensor location format overloads the filename value to encode pass
// information.
// original_file
// >> pass_file_1:line1:col1
// >> pass_file_2:line2:col2
// original_line:original_col
std::string DTensorLocationToString(mlir::Location loc);
} // namespace dtensor
} // namespace tensorflow
// Creates a location, reusing the current name scope.
#define DT_LOC(loc) \
::tensorflow::dtensor::DTensorLocation(loc, __FILE__, __LINE__)
// Creates a location, recording a new nested name scope.
#define DT_LOC2(loc, name) \
::tensorflow::dtensor::DTensorLocation(loc, __FILE__, __LINE__, name)
#endif // TENSORFLOW_DTENSOR_MLIR_DTENSOR_LOCATION_H_
tensorflow/dtensor/mlir/dtensor_location_test.cc
+114
View File
@@ -0,0 +1,114 @@
/* Copyright 2022 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/dtensor/mlir/dtensor_location.h"
#include "mlir/IR/Location.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/utils/name_utils.h"
#include "tensorflow/core/platform/test.h"
namespace {
void CheckFileLineColLocation(mlir::Location loc, unsigned line,
unsigned column) {
ASSERT_TRUE(mlir::isa<mlir::FileLineColLoc>(loc));
auto file_line_col_loc = mlir::cast<mlir::FileLineColLoc>(loc);
EXPECT_EQ(file_line_col_loc.getFilename(), "test.cc");
EXPECT_EQ(file_line_col_loc.getLine(), line);
EXPECT_EQ(file_line_col_loc.getColumn(), column);
}
TEST(DTensorLocationTest, HandlesEmptyLocation) {
mlir::MLIRContext ctx;
mlir::Location loc = mlir::FileLineColLoc::get(&ctx, "test.cc", 10, 20);
loc = tensorflow::dtensor::DTensorLocation(loc, "test.cc", 21);
ASSERT_TRUE(mlir::isa<mlir::CallSiteLoc>(loc));
auto callsite_loc = mlir::cast<mlir::CallSiteLoc>(loc);
CheckFileLineColLocation(callsite_loc.getCallee(), 21, 0);
CheckFileLineColLocation(callsite_loc.getCaller(), 10, 20);
constexpr char stack[] = R"stack(>> test.cc:10:20
>> test.cc:21:0)stack";
EXPECT_EQ(tensorflow::dtensor::DTensorLocationToString(loc), stack);
}
TEST(DTensorLocationTest, HandlesMultipleCalls) {
mlir::MLIRContext ctx;
mlir::Location test_loc = mlir::FileLineColLoc::get(&ctx, "test.cc", 10, 20);
test_loc = tensorflow::dtensor::DTensorLocation(test_loc, "test.cc", 21);
test_loc = tensorflow::dtensor::DTensorLocation(test_loc, "test.cc", 22);
test_loc = tensorflow::dtensor::DTensorLocation(test_loc, "test.cc", 23);
test_loc = tensorflow::dtensor::DTensorLocation(test_loc, "test.cc", 24);
auto verify_loc = test_loc;
for (int i = 0; i < 4; ++i) {
ASSERT_TRUE(mlir::isa<mlir::CallSiteLoc>(verify_loc));
auto callsite_loc = mlir::cast<mlir::CallSiteLoc>(verify_loc);
auto callee_loc = callsite_loc.getCallee();
CheckFileLineColLocation(callee_loc, 24 - i, 0);
verify_loc = callsite_loc.getCaller();
}
CheckFileLineColLocation(verify_loc, 10, 20);
constexpr char stack[] = R"stack(>> test.cc:10:20
>> test.cc:21:0
>> test.cc:22:0
>> test.cc:23:0
>> test.cc:24:0)stack";
EXPECT_EQ(tensorflow::dtensor::DTensorLocationToString(test_loc), stack);
}
TEST(DTensorLocationTest, HandlesNameLoc) {
mlir::MLIRContext ctx;
mlir::Location test_loc =
mlir::NameLoc::get(mlir::StringAttr::get(&ctx, "op@"),
mlir::FileLineColLoc::get(&ctx, "test.cc", 10, 20));
test_loc = tensorflow::dtensor::DTensorLocation(test_loc, "test.cc", 21);
ASSERT_EQ(mlir::GetNameFromLoc(test_loc), "op");
ASSERT_TRUE(mlir::isa<mlir::CallSiteLoc>(test_loc));
auto callsite_loc = mlir::cast<mlir::CallSiteLoc>(test_loc);
mlir::Location caller_loc =
mlir::cast<mlir::CallSiteLoc>(test_loc).getCaller();
ASSERT_TRUE(mlir::isa<mlir::NameLoc>(caller_loc));
CheckFileLineColLocation(mlir::cast<mlir::NameLoc>(caller_loc).getChildLoc(),
10, 20);
mlir::Location callee_loc = callsite_loc.getCallee();
ASSERT_TRUE(mlir::isa<mlir::NameLoc>(callee_loc));
CheckFileLineColLocation(mlir::cast<mlir::NameLoc>(callee_loc).getChildLoc(),
21, 0);
constexpr char stack[] = R"stack(>> test.cc:10:20
>> test.cc:21:0)stack";
EXPECT_EQ(tensorflow::dtensor::DTensorLocationToString(test_loc), stack);
}
TEST(DTensorLocationTest, HandlesNameLocWithName) {
mlir::MLIRContext ctx;
mlir::Location test_loc =
mlir::NameLoc::get(mlir::StringAttr::get(&ctx, "op@"),
mlir::FileLineColLoc::get(&ctx, "test.cc", 10, 20));
test_loc =
tensorflow::dtensor::DTensorLocation(test_loc, "test.cc", 21, "nested");
EXPECT_EQ(mlir::GetNameFromLoc(test_loc), "op/nested");
constexpr char stack[] = R"stack(>> test.cc:10:20
>> test.cc:21:0)stack";
EXPECT_EQ(tensorflow::dtensor::DTensorLocationToString(test_loc), stack);
}
} // namespace
tensorflow/dtensor/mlir/dtensor_mixed_precision_reduce.cc
+171
View File
@@ -0,0 +1,171 @@
/* Copyright 2022 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 "absl/log/log.h"
#include "absl/strings/string_view.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/Visitors.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Transforms/Passes.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/dtensor/cc/dtensor_utils.h"
#include "tensorflow/dtensor/mlir/dtensor_mlir_passes.h"
#include "tensorflow/dtensor/mlir/ir/tf_dtensor.h"
#include "tensorflow/dtensor/mlir/layout_parsing.h"
#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
namespace tensorflow {
namespace dtensor {
namespace {
#define GEN_PASS_DEF_DTENSORMIXEDPRECISIONREDUCE
#include "tensorflow/dtensor/mlir/dtensor_passes.h.inc"
// Extracts the reduction group size from the group_assignment operand of the
// reduce op. group_assignment is a 2-dimensional array where each element is
// the list of devices that are a part of the same reduction group.
template <class ReduceOpType>
mlir::LogicalResult GetAllReduceGroupSize(ReduceOpType reduce_op,
int32_t* group_size) {
mlir::DenseIntElementsAttr group_assignment_attr;
if (!matchPattern(reduce_op.getGroupAssignment(),
m_Constant(&group_assignment_attr)))
return mlir::emitError(reduce_op.getLoc(),
"group_assigment must be a constant.");
if (group_assignment_attr.getType().getRank() != 2)
return mlir::emitError(reduce_op.getLoc(),
"group_assignment should have two dimensions.");
*group_size = group_assignment_attr.getType().getShape()[1];
return mlir::success();
}
// For large enough reduction groups, we compute reductions in a higher
// precision type to ensure accuracy is not lost with sequential addition
// of large numbers in a lower precision type. If the given reduce op meets the
// following criteria:
// - the tensors being reduced are of type bfloat16,
// - the reduction group is at least as large as the configurable env var
// DTENSOR_REDUCE_IN_BFLOAT16_MAX_GROUP_SIZE,
// then the tensors are upcasted to float32 for the reduction before being
// downcasted again.
template <class ReduceOpType>
mlir::LogicalResult MaybeUpcastForReduction(ReduceOpType reduce_op,
bool* changed) {
const mlir::RankedTensorType& input_type =
mlir::dyn_cast<mlir::RankedTensorType>(reduce_op.getInput().getType());
if (!input_type.getElementType().isBF16()) {
// Upcast only applies for bfloat16 input.
return mlir::success();
}
mlir::OpBuilder builder(reduce_op);
const mlir::Location loc = reduce_op.getLoc();
int32_t group_size;
if (mlir::failed(GetAllReduceGroupSize(reduce_op, &group_size)))
return mlir::failure();
if (group_size <= ReduceInBfloat16MaxGroupSize())
// Reduce group size is not sufficient, so we do not modify the ops.
return mlir::success();
const auto reduce_layout = ExtractRequiredSingleLayoutFromOp(reduce_op);
if (!reduce_layout.ok())
return reduce_op.emitOpError(llvm::formatv(
"Malformed layout specification for DTensor reduce op found: {0}",
reduce_layout.status().message()));
// The original output tensor type that would have been used by all users of
// the reduce op.
const mlir::RankedTensorType& output_type =
mlir::dyn_cast<mlir::RankedTensorType>(reduce_op.getOutput().getType());
mlir::TF::CastOp upcast = mlir::TF::CastOp::create(
builder, loc,
mlir::RankedTensorType::get(input_type.getShape(), builder.getF32Type()),
reduce_op.getInput());
reduce_op->setOperand(0, upcast.getY());
reduce_op.getOutput().setType(upcast.getY().getType());
builder.setInsertionPointAfter(reduce_op);
mlir::TF::CastOp downcast = mlir::TF::CastOp::create(
builder, loc,
mlir::RankedTensorType::get(output_type.getShape(),
output_type.getElementType()),
reduce_op);
// Match the layout of the downcast with the reduce op, this is required for
// the later passes.
SetSingleLayoutOnOp(downcast, *reduce_layout);
reduce_op.getOutput().replaceAllUsesExcept(downcast.getY(), downcast);
*changed = true;
return mlir::success();
}
template <class ReduceOpType>
mlir::LogicalResult TryMixedPrecisionReduce(mlir::func::FuncOp function,
absl::string_view opName) {
int32_t reduceOpsCounter = 0;
int32_t changedReduceOpsCounter = 0;
mlir::WalkResult walk_result = function.walk([&](ReduceOpType reduce_op) {
if (reduce_op.getReduceOp().str() == kReduceOpAdd) {
reduceOpsCounter += 1;
bool changed = false;
if (mlir::failed(MaybeUpcastForReduction(reduce_op, &changed)))
return mlir::WalkResult::interrupt();
if (changed) changedReduceOpsCounter += 1;
}
return mlir::WalkResult::advance();
});
if (walk_result.wasInterrupted()) return mlir::failure();
VLOG(2) << "Applied mixed precision to " << changedReduceOpsCounter << " of "
<< reduceOpsCounter << " Add " << opName << " ops.";
return mlir::success();
}
// MLIR pass that enables tensor upcasting within mixed-precision reduction.
struct DTensorMixedPrecisionReducePass
: public impl::DTensorMixedPrecisionReduceBase<
DTensorMixedPrecisionReducePass> {
void runOnOperation() override {
mlir::func::FuncOp function = getOperation();
if (mlir::failed(TryMixedPrecisionReduce<mlir::TF::DTensorAllReduceOp>(
function, "DTensorAllReduce")))
return signalPassFailure();
if (mlir::failed(TryMixedPrecisionReduce<mlir::TF::DTensorReduceScatterOp>(
function, "DTensorReduceScatter")))
return signalPassFailure();
}
};
} // namespace
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateDTensorMixedPrecisionReducePass() {
return std::make_unique<DTensorMixedPrecisionReducePass>();
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/dtensor_mlir_passes.cc
+417
View File
@@ -0,0 +1,417 @@
/* Copyright 2022 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/dtensor/mlir/dtensor_mlir_passes.h"
#include <functional>
#include <memory>
#include <string>
#include "absl/log/vlog_is_on.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_replace.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/OperationSupport.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "mlir/Transforms/Passes.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/transforms/host_runtime/runtime_passes.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/data_dumper_logger_config.h"
#include "tensorflow/core/util/debug_data_dumper.h"
#include "tensorflow/dtensor/cc/constants.h"
#include "tensorflow/dtensor/cc/dtensor_utils.h"
#include "tensorflow/dtensor/mlir/create_dtensor_mlir_passes.h"
#include "tensorflow/dtensor/mlir/op_utils.h"
#include "tensorflow/dtensor/mlir/utils/dtensor_mlir_passes_internal.h"
namespace tensorflow {
namespace dtensor {
namespace {
class ConditionalPrinter : public DataDumperLoggerConfig {
private:
bool do_not_print_;
public:
explicit ConditionalPrinter(
std::function<std::string(const std::string &, mlir::Operation *op)>
get_filename,
bool print_module_scope = false, bool print_after_only_on_change = true,
mlir::OpPrintingFlags op_printing_flags = mlir::OpPrintingFlags())
: DataDumperLoggerConfig(get_filename,
/*pass_prefix=*/"", print_module_scope,
print_after_only_on_change, op_printing_flags) {
do_not_print_ = !(LogOnAllTasks() || (ClientId() == 0));
}
void printBeforeIfEnabled(mlir::Pass *pass, mlir::Operation *operation,
PrintCallbackFn print_callback) override {}
void printAfterIfEnabled(mlir::Pass *pass, mlir::Operation *operation,
PrintCallbackFn print_callback) override {
// NOTE(b/284312504): Disable dumping of
// FunctionalToExecutorDialectConversionPass as it tends to get very large
// being a nested pass on FuncOp before inliner.
if (pass->getName() == "ExecutorDialectToFunctionalPass") {
return;
}
if (pass->getName() == "FunctionalToExecutorDialectConversionPass") {
return;
}
mlir::ModuleOp module = mlir::dyn_cast<mlir::ModuleOp>(operation);
if (!module) module = operation->getParentOfType<mlir::ModuleOp>();
if (module && !module->hasAttr(dtensor::kDoNotLog) && !do_not_print_)
DataDumperLoggerConfig::printAfterIfEnabled(pass, operation,
print_callback);
}
};
} // namespace
// Adds logger to DTensor transformation passmanager.
bool MaybeEnableLogging(mlir::PassManager *pm) {
if (VLOG_IS_ON(1) ||
DEBUG_DATA_DUMPER()->ShouldDump("", kDebugGroupDTensorMlir)) {
// Print the whole module after each pass, which requires disabling
// multi-threading as well.
pm->getContext()->disableMultithreading();
mlir::OpPrintingFlags flags;
if (VLOG_IS_ON(5)) {
// Enable debug information, which includes the call stack of each op.
// This might generate a huge MLIR graph dump, so put it under VLOG(5).
flags = flags.enableDebugInfo(true, true).useLocalScope();
}
pm->enableIRPrinting(std::make_unique<ConditionalPrinter>(
[](const std::string &tag, mlir::Operation *op) {
// As we don't have a good way to pass down the DOperation name, use
// a dummy string.
auto module = mlir::dyn_cast<mlir::ModuleOp>(op);
if (!module) {
module = op->getParentOfType<mlir::ModuleOp>();
}
std::string operation_name = GetOperationName(module);
return DEBUG_DATA_DUMPER()->GetDumpFilename(
"dtensor", kDebugGroupDTensorMlir,
absl::StrReplaceAll(absl::StrCat(tag, ".", operation_name),
{{" ", "_"}}));
},
/*print_module_scope=*/true, /*print_after_only_on_change=*/true,
flags));
return true;
}
return false;
}
void CreateDTensorMLIRPass(const mlir::TF::StandardPipelineOptions &options,
mlir::OpPassManager *pm) {
// Remove ops that cannot be reached from the sink node.
pm->addNestedPass<mlir::func::FuncOp>(
mlir::tf_executor::CreateTFExecutorGraphPruningPass());
// Remove graph-def executor dialect and represent IR as a flattened list of
// TF ops in functions.
pm->addNestedPass<mlir::func::FuncOp>(
mlir::CreateExecutorDialectToFunctionalConversionPass());
// This does not guarantee that shape are inferred for all ops. For ops with
// dynamic shapes, shape information may still be missing.
pm->addPass(mlir::TF::CreateTFShapeInferencePass());
// If V2 layout propagation algorithm, layouts are expressed as DTensorLayout
// op and Canonicalize and Inliner passes will not lose layout information.
pm->addNestedPass<mlir::func::FuncOp>(CreateDTensorPropagateDefaultLayout());
pm->addPass(mlir::createSCCPPass());
pm->addPass(mlir::createCanonicalizerPass());
pm->addPass(mlir::TF::CreateTFFunctionalControlFlowToRegions());
pm->addPass(mlir::createInlinerPass());
// An additional shape inference to catch any newly created constants
// from canonicalizer.
pm->addPass(mlir::TF::CreateTFShapeInferencePass());
// Ensure that all functions have `device_id` as 0th argument.
pm->addPass(CreateDTensorPropagateDeviceIdToFunctionArgs());
// Ensure that all functions with SparseTensor input is converted to its
// three component tensors and SparseToDenseOps are emitted for every usage
// of a SparseTensor.
pm->addPass(CreateDTensorSparseTensorToDenseTensor());
// After shape inference, there may be unused constants ops added when
// propagating caller-callee constants. As DTensor mesh/layout propagation
// passes assumes that there are no unreachable ops, removes trivial unused
// ops. Note that `Canonicalizer` pass in TF includes similar optimization.
// However, canonicalizer pass also rewrites some ops and may remove `_layout`
// or `_mesh` attributes in the re-written TF ops.
// TODO(hongjunchoi): Remove this pass once shape inference pass no longer
// creates unnecessary constants ops.
pm->addNestedPass<mlir::func::FuncOp>(CreateDTensorDCE());
// Canonicalization will merge tf.ConstOp from different DTensorLayout
// annotations, causing problem during mesh propagation. Undo the merge
// before creating clusters.
pm->addNestedPass<mlir::func::FuncOp>(
CreateDTensorUndoMergeConstAcrossMesh());
// Backward functions insert tf.IdentityOp before CopyToMesh's gradient Ops.
// These tf.IdentityOp are semantically no-op to DTensor, but stops the
// backward mesh propagation into the originating tf.ConstOp. Elide the no-op
// tf.IdentityOp to workaround this.
pm->addNestedPass<mlir::func::FuncOp>(
CreateDTensorElideIdentityBeforeCopyToMesh());
// Propagate mesh cluster config and cluster ops by mesh cluster so that
// SPMD expansion can be isolated to a single device mesh.
pm->addNestedPass<mlir::func::FuncOp>(CreateDTensorOpToDeviceClusterPass());
pm->addPass(CreateDTensorMeshPropagationPass());
{
mlir::OpPassManager &func_pm = pm->nest<mlir::func::FuncOp>();
func_pm.addPass(CreateDTensorDeviceMeshClusterCoarsening());
// Set empty layout to cluster wrapping `tf.VarHandleOp`. VarHandle op
// always runs in the default device where client program executes.
func_pm.addPass(CreateDTensorDesignateResourceHandleMesh());
}
// Clone Control Flow.
pm->addPass(CreateDTensorDecomposeControlflowPass());
// Validates that all cross mesh data transfers are expressed via
// DTensorLayout operation and lowers it to send/recvs.
pm->addPass(CreateDTensorHandleCrossClusterDependencies());
// Merge Clusters
pm->addPass(CreateDTensorMergeClustersPass());
////////
// Propagate layout to all ops in graph.
// For DTensor Checkpoint V2, the outputs of tf.RestoreV2 ops
// do not have shape information. We can infer the shapes of these
// outputs from the tf.AssignVariableOps that consume these outputs.
// This pass fills in all missing shapes caused by tf.RestoreV2 ops.
pm->addPass(CreateDTensorInferShapesForRestoreV2Op());
// Mark all ops and functions with global shape attribute to preserve global
// shape information as it is needed during Layout Propagation and SPMD
// expansion.
pm->addPass(CreateDTensorAnnotateGlobalShape());
pm->addPass(CreateDTensorLayoutPropagationPassV2());
// Expand graph to SPMD form given layouts are annotated to all ops.
// Remove all DTensorLayout ops after the expansion is done.
pm->addPass(CreateDTensorSPMDExpansion());
// Expand all ops that consume SparseTensors to possibly new ops.
// Remove any unused SparseToDense, Layout, and Const Ops after
// the expansion is done.
//
// Note that this pass assumes that SparseTensor operands is represented
// as an operand from the output of a SparseToDenseOp. Thus, this pass
// must happen after SparseTensorToDenseTensor pass and after
// the SPMD Expansion pass.
pm->addPass(CreateDTensorSparseExpansion());
// Do a round of CSE: this helps reduce the number of consts in the graph now
// that SPMD expansion is done. We had replicated all Consts (so that each
// const only had one usage) as part of layout propagation.
pm->addPass(mlir::createCSEPass());
// Lower the AllGather collectives. This has to happen before the all reduce
// optimizations and AllGather may emit an AllReduce.
pm->addPass(CreateDTensorAllGatherLoweringPass());
// Fuses AllReduce and AllScatter into ReduceScatter.
if (!DoNotFuseReduceScatter()) {
pm->addNestedPass<mlir::func::FuncOp>(
CreateDTensorAllReduceScatterOptimization());
}
// Changes order of DTensorAllReduce + Add to Add + DTensorAllReduce to
// minimize number of all reduce operations.
pm->addNestedPass<mlir::func::FuncOp>(
CreateDTensorAllReduceSumOptimization());
AddDTensorAllReduceCombineOptimization(pm);
// Lowers complex and other unsupported types to supported types.
pm->addNestedPass<mlir::func::FuncOp>(
CreateDTensorCollectiveTypeLoweringPass());
// DTensorReduceScatter lowering should come before DTensorAllReduce
// and DTensorAllScatter lowerings since for some devices DTensorReduceScatter
// will be decomposed into a DTensorAllReduce+DTensorScatter.
pm->addPass(CreateDTensorReduceScatterLoweringPass());
// For large enough reduction groups in reduction ops, upcast the input
// tensors to higher precision type (e.g. bfloat16 -> float32).
if (EnableMixedPrecisionReduce()) {
pm->addNestedPass<mlir::func::FuncOp>(
CreateDTensorMixedPrecisionReducePass());
}
// Lower device-agnostic logical AllReduce ops into device-specific physical
// AllReduce ops.
//
// First, find DTensor collective ops such as DTensorAllReduce, which are
// generated by SPMD expansion. Lower them into device-specific forms. For
// most devices, there is a one-to-one mapping: DTensorAllReduce becomes
// CollectiveReduce on CPUs/GPUs and XlaAllReduce on TPU pods.
// Optionally, for special topologies, DTensorAllReduce
// could become a chain of collectives running on different devices:
// XlaAllReduce on each donut followed by CollectiveReduce on the hosts. Those
// collective ops running on hosts will have their _mesh attribute set to
// empty by this pass. The other ops continue to have no _mesh attributes,
// which means they run on the cluster mesh.
pm->addPass(CreateDTensorAllReduceLoweringPass());
pm->addPass(CreateDTensorAllScatterLoweringPass());
pm->addPass(CreateDTensorAllToAllLoweringPass());
// Group together multiple device clusters assigned to the same mesh. Repeat
// this for every mesh to support multi-mesh. Collective lowering may have
// created multiple CPU mesh clusters for executing collective operations on
// CPUs.
// As so, we merge newly created CPU clusters after collective lowering
// especially for special topologies.
pm->addPass(CreateDTensorMergeClustersPass());
pm->addPass(CreateDTensorLowerSendRecv());
// Convert tf_device.cluster into a function call op.
pm->addPass(mlir::TFDevice::CreateClusterOutliningPass());
pm->addPass(CreateDTensorClusterFunctionConversion());
// During layout propagation, we clone all constants with multiple consumers
// for easier analaysis.
// This may create multiple same constants ops. Apply constant folding on
// duplicated constant operations to reduce graph size.
pm->addNestedPass<mlir::func::FuncOp>(CreateDTensorConstantFolding());
// DTensor SPMD lowering passes may have created auxiliary operations that are
// no longer used. Add additional DCE pass to remove unused non-side effecting
// ops.
pm->addNestedPass<mlir::func::FuncOp>(CreateDTensorDCE());
// DTensor SPMD Expansion may have caused multiple control flows and
// duplicate ops to calculate device ordinal. Re-run SCCP and merge
// controlflows if possible.
pm->addNestedPass<mlir::func::FuncOp>(mlir::createSCCPPass());
pm->addNestedPass<mlir::func::FuncOp>(mlir::createCanonicalizerPass());
pm->addPass(mlir::TFDevice::CreateMergeControlFlowPass());
// TF2XLA Integration
{
// Make sure clusters that run on TPU's are correct metadata ops and
// attributes attached to be compatible with later TPU specific optimization
// passes.
pm->addPass(CreateDTensorTPUIntegration());
pm->addNestedPass<mlir::func::FuncOp>(
mlir::TFDevice::CreateDecomposeResourceOpsPass());
// Sink constant ops into cluster region as DecomposeResourceOpsPass() could
// lift constant out due to folding.
pm->addNestedPass<mlir::func::FuncOp>(
mlir::TFDevice::CreateClusterConstantSinkingPass());
// Run another shape inference pass (and following DCE pass) because
// resource decomposition might have created new partial types.
pm->addPass(mlir::TF::CreateTFShapeInferencePass());
pm->addNestedPass<mlir::func::FuncOp>(CreateDTensorDCE());
pm->addPass(mlir::TFDevice::CreateResourceOpLiftingPass());
pm->addPass(mlir::TFDevice::CreateClusterOutliningPass());
// Prepare for XLA SPMD integration for XLA SPMD mesh. If there are layout
// operations on XLA SPMD mesh, then convert all of them to appropriate
// XLA sharding attributes.
pm->addPass(CreateDTensorSetHloShardingPass(
/*check_layout_use_xla_spmd=*/true));
pm->addPass(CreateDTensorReplaceAuxiliaryDTensorLayoutOpPass());
pm->addNestedPass<mlir::func::FuncOp>(
CreateDTensorLayoutToXlaShardingOpPass());
// We lower all remaining Relayout to Identity here to make XLA happy.
// Under XLA SPMD the RelayoutOp is not expanded by DTensor's SPMD expander.
// Note that we do not lower much earlier because
// canonicalization / const folding may produce chains of
// DTensorLayout that confuses DTensorReplaceAuxiliaryDTensorLayoutOpPass.
pm->addNestedPass<mlir::func::FuncOp>(
CreateDTensorReplaceRelayoutWithIdentityPass());
// Rename functions with unique names, to avoid collisions in the function
// library.
pm->addPass(CreateFunctionRenamingPass());
// As DTensor SPMD expansion handles sharded inputs for model
// parallelism, we set input/output sharding to maximal sharding
// for inputs/outputs of the TPU computation.
pm->addNestedPass<mlir::func::FuncOp>(CreateDTensorSetDefaultSharding());
// Creates a pass that marks TPU cluster input-output pairs reading and
// writing to same resource variable as aliases.
pm->addPass(mlir::TFDevice::CreateMarkInputOutputAliasesPass());
// Convert compilation and replication attributes to unified attributes
// expected by TPURewritePass.
pm->addNestedPass<mlir::func::FuncOp>(
mlir::TF::CreateCanonicalizeCompileAndReplicateAttributesPass());
// Rewrite RecvTPUEmbeddingActivationsOp and SendTPUEmbeddingGradients ops
// to internal variants by introducing XlaRecvTPUEmbeddingDeduplicationData
// op.
pm->addNestedPass<mlir::func::FuncOp>(
mlir::TF::CreateRewriteTPUEmbeddingOpsPass());
// Create TPU Compile and TPU Execute ops for each TPU devices.
pm->addPass(mlir::TFTPU::CreateTPURewritePass());
// Convert unified compilation and replication attributes back to legacy
// attributes for subsequent passes.
pm->addNestedPass<mlir::func::FuncOp>(
mlir::TFTPU::CreateConvertToLegacyCompileAndReplicateAttributesPass());
// Add placeholder device attributes to resource arguments of TPU
// computation. This ensures the following
// CreateTPUMergeVariablesWithExecutePass correctly merges resource
// operations with TPUExecute op.
pm->addPass(CreateDTensorTpuAddResourceDeviceAttribute());
// Translate TPUExecute op to TPUExecuteAndUpdateVariable op to enable
// buffer aliasing.
pm->addPass(mlir::TFTPU::CreateTPUMergeVariablesWithExecutePass());
pm->addPass(CreateDTensorUpdateTPUMetadata());
// If send/recv exists between TPU and CPU, then TPU Compilation program key
// is used as input for recv op in host computation as well as TPUExecute op
// in device computation. As so, move TPUCompile logic to host computation
// and transfer program key using send/recv operations.
pm->addPass(CreateDTensorMoveCompilationToHost());
pm->addPass(mlir::createSymbolDCEPass());
// Expands the DTensor call ops across devices within a "multi-device" main.
pm->addPass(CreateDTensorMultiDeviceExpansionPass());
}
pm->addPass(mlir::TF::CreateTFRegionControlFlowToFunctional());
// Convert graph into graph executor dialect so that transformed graph can be
// exported back to Graphdef.
pm->addNestedPass<mlir::func::FuncOp>(
mlir::CreateFunctionalToExecutorDialectConversionPass());
pm->addPass(mlir::CreateBreakUpIslandsPass());
pm->addNestedPass<mlir::func::FuncOp>(
mlir::TFDevice::CreateParallelExecuteToIslandsPass());
pm->addNestedPass<mlir::func::FuncOp>(
mlir::TFDevice::CreateLaunchToDeviceAttributePass());
// Add additional BreakUpIslandPass as LaunchToDeviceAttribute pass may have
// created additional islands.
pm->addPass(mlir::CreateBreakUpIslandsPass());
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/dtensor_mlir_passes.h
+40
View File
@@ -0,0 +1,40 @@
/* Copyright 2022 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_DTENSOR_MLIR_DTENSOR_MLIR_PASSES_H_
#define TENSORFLOW_DTENSOR_MLIR_DTENSOR_MLIR_PASSES_H_
#include <memory>
#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "tensorflow/core/lib/core/status.h"
namespace tensorflow {
namespace dtensor {
bool MaybeEnableLogging(mlir::PassManager* pm);
// Adds MLIR passes to `pm`.
void CreateDTensorMLIRPass(const mlir::TF::StandardPipelineOptions& options,
mlir::OpPassManager* pm);
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_MLIR_DTENSOR_MLIR_PASSES_H_
tensorflow/dtensor/mlir/dtensor_multi_device_expansion.cc
+927
View File
@@ -0,0 +1,927 @@
/* Copyright 2022 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 <algorithm>
#include <cstddef>
#include <cstdint>
#include <iterator>
#include <memory>
#include <optional>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_format.h"
#include "absl/strings/str_join.h"
#include "absl/types/span.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/iterator_range.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/FormatVariadic.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/IRMapping.h" // from @llvm-project
#include "mlir/IR/SymbolTable.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "xla/tsl/platform/status.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/dtensor/cc/constants.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/layout_parsing.h"
#include "tensorflow/dtensor/mlir/op_utils.h"
namespace tensorflow {
namespace dtensor {
namespace {
#define GEN_PASS_DEF_DTENSORMULTIDEVICEEXPANSION
#include "tensorflow/dtensor/mlir/dtensor_passes.h.inc"
constexpr char kDeviceAttr[] = "device";
constexpr char kFuncDeviceAttr[] = "tf.device";
constexpr char kEntryFuncAttr[] = "tf.entry_function";
constexpr char kMainFuncName[] = "main";
constexpr char kIsStatelessAttr[] = "is_stateless";
constexpr int kDeviceIDArgumentNumber = 0;
// This is a map from argument numbers and meshes to per-device values.
// Most arguments will only be expanded on one mesh (the one given by its
// "tf._layout" attribute); however, the device id may be expanded across
// multiple meshes. For example, when main functions have both cpu and tpu
// mesh partitioned calls.
using ExpandedArgumentMap =
absl::flat_hash_map<int,
absl::flat_hash_map<Mesh, std::vector<mlir::Value>>>;
struct ExpandedResults {
std::optional<Layout> layout;
std::vector<mlir::Value> results;
template <typename Value>
void insert(Value&& value) {
using T = std::decay_t<Value>;
if constexpr (std::is_same_v<T, mlir::Value>) {
results.emplace_back(std::forward<Value>(value));
} else {
results.insert(results.end(), value.begin(), value.end());
}
}
};
mlir::BlockArgument InsertArgumentForDevice(mlir::OpBuilder& builder,
mlir::func::FuncOp func,
mlir::Type arg_type,
const std::string& device) {
const int arg_index = func.getNumArguments();
std::vector<mlir::NamedAttribute> named_attrs = {builder.getNamedAttr(
kFuncDeviceAttr, builder.getStringAttr(llvm::StringRef(device)))};
llvm::ArrayRef<mlir::NamedAttribute> named_array_ref(named_attrs);
mlir::DictionaryAttr dic_attr = builder.getDictionaryAttr(named_array_ref);
(void)func.insertArgument(arg_index, arg_type, dic_attr, func.getLoc());
return func.getArgument(arg_index);
}
// Returns the user of all the ops in the span iff it is a single return op.
// Otherwise, returns nullptr; for example, if there are multiple return ops.
template <typename UserType, typename Operations>
mlir::LogicalResult GetUniqueUserFromOps(const Operations& ops,
UserType* result) {
for (mlir::Operation* op : ops) {
for (mlir::Operation* user : op->getUsers()) {
// TODO(twelve): Determine whether we should follow identity ops.
UserType typed;
if constexpr (std::is_same_v<UserType, mlir::Operation*>) {
typed = user;
} else {
typed = llvm::dyn_cast_or_null<UserType>(user);
}
if (typed) {
if (*result == nullptr) {
*result = typed;
} else if (*result != typed) {
return mlir::failure();
}
} else {
return mlir::failure();
}
}
}
return mlir::success();
}
// Returns the devices for a given mesh.
absl::Span<const std::string> GetDevices(const Mesh& mesh) {
const std::vector<std::string>& devices = mesh.global_devices();
if (devices.empty()) {
return mesh.local_devices();
} else {
return devices;
}
}
StatusOr<absl::Span<mlir::Value>> GetExpandedArguments(
mlir::OpBuilder& builder, mlir::func::FuncOp target_func,
ExpandedArgumentMap& expanded_arguments, mlir::BlockArgument argument,
const Mesh* target_mesh = nullptr, bool has_device_id = true);
StatusOr<std::optional<std::vector<Layout>>> GetResourceLayouts(
mlir::Operation* op) {
if (op->hasAttr(kNewResourceArgLayouts)) {
auto attrs = op->getAttrOfType<mlir::ArrayAttr>(kNewResourceArgLayouts);
std::vector<Layout> layouts;
layouts.reserve(attrs.size());
for (mlir::Attribute attr : attrs) {
auto string_attr = mlir::cast<mlir::StringAttr>(attr);
auto layout = Layout::FromString(string_attr.str());
if (layout.ok()) {
layouts.emplace_back(std::move(layout.value()));
} else {
return layout.status();
}
}
return layouts;
} else {
return std::nullopt;
}
}
bool IsResource(mlir::Value value) {
return mlir::isa<mlir::TF::ResourceType>(
getElementTypeOrSelf(value.getType()));
}
StatusOr<std::optional<Layout>> FindResourceLayout(mlir::BlockArgument arg) {
uint32_t arg_num = arg.getArgNumber();
for (mlir::Operation* user : arg.getUsers()) {
auto resource_layouts = GetResourceLayouts(user);
if (resource_layouts.ok()) {
const auto& opt = resource_layouts.value();
if (!opt || opt->empty()) {
continue;
}
} else {
return resource_layouts.status();
}
auto resource_indices = user->getAttrOfType<mlir::DenseIntElementsAttr>(
kNewResourceLayoutIndices);
if (!resource_indices) {
return absl::InvalidArgumentError(
absl::StrCat("missing ", kNewResourceLayoutIndices));
}
for (auto [i, index] : llvm::enumerate(resource_indices)) {
int64_t index_value = index.getSExtValue();
if (index_value == arg_num) {
return (resource_layouts.value())->at(i);
}
}
}
return std::nullopt;
}
template <typename Results>
mlir::FunctionType GetFunctionType(mlir::OpBuilder& builder,
mlir::func::FuncOp func, Results results) {
std::vector<mlir::Type> input_types, result_types;
for (mlir::BlockArgument input : func.getArguments()) {
input_types.emplace_back(input.getType());
}
for (const auto result : results) {
result_types.emplace_back(result.getType());
}
return builder.getFunctionType(input_types, result_types);
}
mlir::LogicalResult MakeParallelExecute(
mlir::OpBuilder& builder, mlir::tf_device::LaunchOp launch_op,
absl::Span<const std::string> devices,
std::vector<mlir::IRMapping>& mappings) {
const unsigned num_devices = devices.size();
const unsigned num_results = launch_op->getNumResults();
std::vector<mlir::Type> result_types;
result_types.reserve(num_results * num_devices);
// Expand the result types across all the devices.
for (int i = 0; i < num_devices; ++i) {
for (mlir::Type result_type : launch_op.getResultTypes()) {
result_types.emplace_back(result_type);
}
}
// Build the ParallelExecuteOp.
mlir::Location loc = launch_op.getLoc();
mlir::tf_device::ParallelExecuteOp parallel_execute_op =
mlir::tf_device::ParallelExecuteOp::create(builder, loc, num_devices,
result_types);
// Clone the LaunchOp for each of the devices.
for (unsigned dev_idx = 0; dev_idx < num_devices; ++dev_idx) {
mlir::IRMapping& mapping = mappings[dev_idx];
mlir::Block& block = parallel_execute_op.GetRegionBlockWithIndex(dev_idx);
builder.setInsertionPointToStart(&block);
mlir::Operation* clone =
builder.clone(*((mlir::Operation*)launch_op), mapping);
clone->setAttr(kDeviceAttr, builder.getStringAttr(devices[dev_idx]));
mlir::tf_device::ReturnOp::create(builder, loc, clone->getResults());
}
// Map the results of the LaunchOp to the ParallelExecuteOp.
for (unsigned dev_idx = 0; dev_idx < num_devices; ++dev_idx) {
mlir::IRMapping& mapping = mappings[dev_idx];
for (unsigned res_idx = 0; res_idx < num_results; ++res_idx) {
unsigned expanded_idx = res_idx + dev_idx * num_results;
mapping.map(launch_op->getResult(res_idx),
parallel_execute_op->getResult(expanded_idx));
}
}
return mlir::success();
}
// Rewrites a launch to a ParallelExecuteOp with per-device launches.
mlir::LogicalResult RewriteTPUFunction(mlir::func::FuncOp func,
mlir::tf_device::LaunchOp launch_op,
const Mesh& target_mesh,
mlir::FunctionType& func_type) {
mlir::OpBuilder builder = mlir::OpBuilder::atBlockBegin(func->getBlock());
builder.setInsertionPointAfter(launch_op);
const absl::Span<const std::string> devices = GetDevices(target_mesh);
const size_t num_devices = devices.size();
// Maps values from original to per-device representations.
std::vector<mlir::IRMapping> mappings(num_devices);
// Expand each of the function's inputs across all the devices.
// TODO(twelve): This logic assumes all the function's inputs should be
// expanded; however, when might that not be true?
ExpandedArgumentMap expanded_arguments_map;
const unsigned num_arguments = func.getNumArguments();
for (unsigned arg_idx = 0; arg_idx < num_arguments; ++arg_idx) {
mlir::BlockArgument argument = func.getArgument(arg_idx);
StatusOr<absl::Span<mlir::Value>> expanded_arguments = GetExpandedArguments(
builder, func, expanded_arguments_map, argument, &target_mesh, false);
if (expanded_arguments.ok()) {
// Map the original arguments to their per-device counterparts.
for (unsigned dev_idx = 0; dev_idx < num_devices; ++dev_idx) {
mappings[dev_idx].map(argument, expanded_arguments->at(dev_idx));
}
} else {
func->emitOpError(absl::StatusMessageAsCStr(expanded_arguments.status()));
return mlir::failure();
}
}
// Collect all the original ops.
llvm::iterator_range<mlir::Region::OpIterator> op_range =
func.getFunctionBody().getOps();
std::vector<mlir::Operation*> ops;
ops.reserve(std::distance(op_range.begin(), op_range.end()));
for (mlir::Operation& op : op_range) {
ops.emplace_back(&op);
}
// Clone all the original ops.
std::vector<mlir::Value> results;
for (mlir::Operation* op : ops) {
builder.setInsertionPointAfter(op);
if (llvm::isa<mlir::tf_device::LaunchOp>(op)) {
if (op == (mlir::Operation*)launch_op) {
// Build the parallel execute operation.
mlir::LogicalResult result =
MakeParallelExecute(builder, launch_op, devices, mappings);
if (mlir::failed(result)) {
return result;
}
}
} else if (llvm::isa<mlir::func::ReturnOp>(op)) {
// Collect the per-device results.
results.reserve(op->getNumOperands() * num_devices);
for (mlir::Value operand : op->getOperands()) {
for (unsigned dev_idx = 0; dev_idx < num_devices; ++dev_idx) {
results.emplace_back(mappings[dev_idx].lookup(operand));
}
}
mlir::func::ReturnOp::create(builder, op->getLoc(), results);
} else {
// Clone the operation across all the devices.
for (unsigned dev_idx = 0; dev_idx < num_devices; ++dev_idx) {
mlir::Operation* clone = builder.clone(*op, mappings[dev_idx]);
clone->setAttr(kDeviceAttr, builder.getStringAttr(devices[dev_idx]));
}
}
}
// Erase the original non-launch ops (go backwards to remove users first).
for (int op_idx = (int)ops.size() - 1; op_idx >= 0; --op_idx) {
mlir::Operation* op = ops[op_idx];
if (!llvm::isa<mlir::tf_device::LaunchOp>(op) ||
(op == (mlir::Operation*)launch_op)) {
op->erase();
}
}
// Erase the function's original arguments.
for (unsigned arg_idx = 0; arg_idx < num_arguments; ++arg_idx) {
if (failed(func.eraseArgument(0))) return mlir::failure();
}
// Update the function's type.
func_type = GetFunctionType(builder, func, results);
func.setFunctionType(func_type);
return mlir::success();
}
// Rewrites a call-like op into an equivalent op for each device;
// de/multiplexes the per-device inputs/outputs for each "expanded" op.
// Usable when there are not any device launch ops (typically on CPU/GPU).
template <typename OperationType>
mlir::LogicalResult ExpandOperation(
mlir::func::FuncOp target_func, mlir::func::ReturnOp return_op,
ExpandedArgumentMap& expanded_arguments,
std::vector<ExpandedResults>& expanded_results, const Mesh& target_mesh,
OperationType op);
// Rewrites a call-like op targeting a single device launch op into
// a parallel execute.
template <typename OperationType>
mlir::LogicalResult ExpandTPUOperation(
mlir::func::FuncOp target_func, mlir::func::ReturnOp return_op,
ExpandedArgumentMap& expanded_arguments,
std::vector<ExpandedResults>& expanded_results, const Mesh& target_mesh,
OperationType op) {
mlir::FunctionType func_type;
mlir::OpBuilder builder(target_func.getBody());
mlir::ModuleOp module = op->template getParentOfType<mlir::ModuleOp>();
mlir::func::FuncOp func = module.lookupSymbol<mlir::func::FuncOp>(op.getF());
mlir::tf_device::LaunchOp launch_op;
func.walk([&](mlir::tf_device::LaunchOp op) {
(op.getBody()).walk([&](mlir::Operation* child) {
if (llvm::isa<mlir::TF::TPUExecuteOp>(child) ||
llvm::isa<mlir::TF::TPUExecuteAndUpdateVariablesOp>(child)) {
launch_op = op;
return mlir::WalkResult::interrupt();
} else {
return mlir::WalkResult::advance();
}
});
});
if (launch_op) {
// Found a device launch with TPUExecute, try expanding it.
if (mlir::failed(
RewriteTPUFunction(func, launch_op, target_mesh, func_type))) {
return mlir::failure();
}
} else {
// There were not any TPUExecute ops, fallback to the conventional path.
return ExpandOperation(target_func, return_op, expanded_arguments,
expanded_results, target_mesh, op);
}
llvm::SmallVector<mlir::Value, 8> operands;
for (const mlir::Value& operand : op->getOperands()) {
if (const auto arg = mlir::dyn_cast_or_null<mlir::BlockArgument>(operand)) {
const StatusOr<absl::Span<mlir::Value>> new_args = GetExpandedArguments(
builder, target_func, expanded_arguments, arg, &target_mesh);
if (!new_args.ok()) {
op->emitOpError(absl::StatusMessageAsCStr(new_args.status()));
return mlir::failure();
} else if (new_args->empty()) {
operands.push_back(operand);
} else {
operands.insert(operands.end(), new_args->begin(), new_args->end());
}
} else {
operands.push_back(operand);
}
}
auto call_op =
OperationType::create(builder, op->getLoc(), func_type.getResults(),
operands, /*args_attrs=*/nullptr,
/*res_attrs=*/nullptr, op.getFAttr(),
/*config=*/builder.getStringAttr(""),
/*config_proto=*/builder.getStringAttr(""),
/*executor_type=*/builder.getStringAttr(""));
const absl::Span<const std::string> devices = GetDevices(target_mesh);
const size_t num_devices = devices.size();
if (return_op) {
mlir::Operation::operand_range operands = return_op->getOperands();
// Expand each operand of the return operation.
for (const auto [result_index, operand] : llvm::enumerate(operands)) {
// (All the operands should originate from the original call op.)
if (op == operand.getDefiningOp()) {
const mlir::Operation::result_range results = op->getResults();
const mlir::Operation::result_range::iterator search =
llvm::find(results, operand);
const size_t i = search - results.begin();
// For each device--
for (int j = 0; j < num_devices; j++) {
// Find the corresponding expanded output within the new op.
mlir::Value result = call_op->getResult(i * num_devices + j);
auto identity_op = mlir::TF::IdentityOp::create(
builder, result.getLoc(), result.getType(), result);
expanded_results[result_index].insert(
(mlir::Value)identity_op.getResult());
}
}
}
}
return mlir::success();
}
template <typename Operation>
mlir::LogicalResult ExpandOperation(
mlir::func::FuncOp target_func, mlir::func::ReturnOp return_op,
ExpandedArgumentMap& expanded_arguments,
std::vector<ExpandedResults>& expanded_results, const Mesh& target_mesh,
Operation op) {
mlir::OpBuilder builder(target_func.getBody());
const absl::Span<const std::string> devices = GetDevices(target_mesh);
const size_t num_devices = devices.size();
llvm::SmallVector<Operation> replications;
for (size_t i = 0; i < num_devices; ++i) {
llvm::SmallVector<mlir::Value, 8> operands;
for (const mlir::Value& operand : op->getOperands()) {
if (const auto arg =
mlir::dyn_cast_or_null<mlir::BlockArgument>(operand)) {
const StatusOr<absl::Span<mlir::Value>> new_args = GetExpandedArguments(
builder, target_func, expanded_arguments, arg, &target_mesh);
if (!new_args.ok()) {
op->emitOpError(absl::StatusMessageAsCStr(new_args.status()));
return mlir::failure();
} else if (new_args->empty()) {
operands.push_back(operand);
} else {
operands.push_back((*new_args)[i]);
}
} else {
operands.push_back(operand);
}
}
auto new_op =
Operation::create(builder, op->getLoc(), op->getResultTypes(), operands,
/*args_attrs=*/nullptr,
/*res_attrs=*/nullptr, op.getFAttr(),
/*config=*/builder.getStringAttr(""),
/*config_proto=*/builder.getStringAttr(""),
/*executor_type=*/builder.getStringAttr(""));
// Set the "is_stateless" attribute to ensure that side-effect analysis
// does not set the per-device call ops to depend on one another (see
// `CanHaveSideEffects`), which would cause collectives to hang.
new_op->setAttr(kIsStatelessAttr, builder.getBoolAttr(true));
new_op->setAttr(kDeviceAttr, builder.getStringAttr(devices[i]));
replications.emplace_back(new_op);
}
if (return_op) {
mlir::Operation::operand_range operands = return_op->getOperands();
for (const auto [i, operand] : llvm::enumerate(operands)) {
if (op == operand.getDefiningOp()) {
const mlir::Operation::result_range results = op->getResults();
const mlir::Operation::result_range::iterator search =
llvm::find(results, operand);
const size_t result_number = search - results.begin();
for (const Operation& replication : replications) {
expanded_results[i].insert(
(mlir::Value)replication->getResult(result_number));
}
}
}
}
return mlir::success();
}
// Generates a `NamedAttr` whose value is a comma-separated list of
// elements, given by applying the format function to each integer in
// the range 0..len(types).
template <typename Format>
mlir::NamedAttribute GetNamesAttr(mlir::OpBuilder& builder, const char* name,
llvm::ArrayRef<mlir::Type> types,
const Format& fmt) {
llvm::SmallVector<std::string, 8> names;
for (int i = 0; i < types.size(); ++i) {
names.push_back(fmt(i));
}
return builder.getNamedAttr(name,
builder.getStringAttr(absl::StrJoin(names, ",")));
}
// Updates the entry function's attributes to reflect its inputs/outputs.
void UpdateEntryFuncAttr(mlir::OpBuilder& builder, mlir::func::FuncOp func) {
const mlir::FunctionType func_type = func.getFunctionType();
const mlir::NamedAttribute inputs =
GetNamesAttr(builder, "inputs", func_type.getInputs(),
[](int i) { return absl::StrFormat("input_%d", i); });
const mlir::NamedAttribute outputs =
GetNamesAttr(builder, "outputs", func_type.getResults(),
[](int i) { return absl::StrFormat("output_%d", i); });
llvm::SmallVector<mlir::NamedAttribute, 2> named_attrs = {inputs, outputs};
llvm::ArrayRef<mlir::NamedAttribute> named_array_ref(named_attrs);
mlir::DictionaryAttr dic_attr = builder.getDictionaryAttr(named_array_ref);
func->setAttr(kEntryFuncAttr, dic_attr);
}
StatusOr<absl::Span<mlir::Value>> GetExpandedArguments(
mlir::OpBuilder& builder, mlir::func::FuncOp target_func,
ExpandedArgumentMap& expanded_arguments, mlir::BlockArgument arg,
const Mesh* target_mesh, bool has_device_id) {
std::optional<Mesh> mesh;
unsigned int argument_number = arg.getArgNumber();
if (!has_device_id || (argument_number == kDeviceIDArgumentNumber)) {
if (target_mesh) {
mesh = *target_mesh;
}
} else {
TF_ASSIGN_OR_RETURN(std::optional<Layout> layout,
ExtractLayoutFromOperand(arg));
if (layout) {
mesh = layout->mesh();
if (mesh->IsEmpty()) {
if (target_mesh) {
mesh = *target_mesh;
} else if (IsResource(arg)) {
TF_ASSIGN_OR_RETURN(layout, FindResourceLayout(arg));
if (layout) {
mesh = layout->mesh();
} else {
return absl::InvalidArgumentError(
absl::StrCat("Could not find resource layout for %arg",
arg.getArgNumber(), "!"));
}
}
}
}
}
if (mesh.has_value()) {
std::vector<mlir::Value>& replications =
expanded_arguments[argument_number][*mesh];
if (replications.empty()) {
const absl::Span<const std::string> devices = GetDevices(*mesh);
const size_t num_devices = devices.size();
replications.reserve(num_devices);
if (has_device_id && argument_number == kDeviceIDArgumentNumber) {
for (int i = 0; i < num_devices; ++i) {
const auto value_attr = mlir::DenseIntElementsAttr::get<int>(
mlir::RankedTensorType::get({}, builder.getI32Type()), {i});
replications.emplace_back(
mlir::TF::ConstOp::create(builder, arg.getLoc(), value_attr));
}
} else {
mlir::TensorType tensor_type =
mlir::dyn_cast_or_null<mlir::TensorType>(arg.getType());
if (!tensor_type) {
return absl::InvalidArgumentError("Could not determine tensor type.");
}
for (int i = 0; i < num_devices; ++i) {
replications.emplace_back(InsertArgumentForDevice(
builder, target_func, tensor_type, devices[i]));
}
}
}
return absl::Span<mlir::Value>(replications);
} else {
return absl::Span<mlir::Value>(); // no per-device arguments necessary
}
}
struct InferredResourceAttributes {
mlir::Attribute layouts;
mlir::Attribute indices;
InferredResourceAttributes(mlir::Attribute layouts_, mlir::Attribute indices_)
: layouts(layouts_), indices(indices_) {}
};
template <typename Operations>
mlir::LogicalResult GetInferredResourceAttributes(
mlir::OpBuilder& builder, const Operations& call_ops,
std::optional<InferredResourceAttributes>* resource_attrs) {
llvm::SmallVector<mlir::Attribute, 8> resource_layouts;
llvm::SmallVector<int32_t, 8> resource_indices;
for (mlir::Operation* call_op : call_ops) {
const auto resource_layouts_attr =
call_op->getAttrOfType<mlir::ArrayAttr>(kNewResourceArgLayouts);
const auto resource_indices_attr =
call_op->getAttrOfType<mlir::DenseIntElementsAttr>(
kNewResourceLayoutIndices);
if (resource_indices_attr && resource_layouts_attr) {
for (auto [index, layout] :
llvm::zip(resource_indices_attr, resource_layouts_attr)) {
// Build up the lists of resource indices and layouts.
resource_indices.emplace_back(index.getSExtValue());
resource_layouts.emplace_back(layout);
}
}
}
if (!resource_layouts.empty()) {
resource_attrs->emplace(builder.getArrayAttr(resource_layouts),
builder.getI32VectorAttr(resource_indices));
}
return mlir::success();
}
// Build a new main function that calls the multi-device/translated function.
template <typename Operations>
mlir::LogicalResult BuildOuterMainFunc(
mlir::ModuleOp module, mlir::func::FuncOp old_main_func,
mlir::func::FuncOp translated_func, mlir::func::ReturnOp return_op,
const std::vector<ExpandedResults>& expanded_results,
mlir::ArrayAttr num_local_outputs_attr, Operations&& call_ops) {
using CallOp = typename std::decay_t<Operations>::value_type;
mlir::SymbolTable symbol_table(module);
mlir::Block* module_body = module.getBody();
mlir::OpBuilder builder = mlir::OpBuilder::atBlockBegin(module_body);
// Build a new main function with no initial attributes/return type.
mlir::func::FuncOp main_func = mlir::func::FuncOp::create(
old_main_func.getLoc(), "main", builder.getFunctionType({}, {}));
mlir::Block* entry_block = main_func.addEntryBlock();
builder.setInsertionPointToEnd(entry_block);
// Copy the arguments from the translated function to the new main function.
std::vector<mlir::Value> inputs;
for (auto [arg_index, arg] :
llvm::enumerate(translated_func.getArguments())) {
if (failed((main_func.insertArgument(
arg_index, arg.getType(), translated_func.getArgAttrDict(arg_index),
old_main_func.getLoc())))) {
return mlir::failure();
}
inputs.emplace_back(main_func.getArgument(arg_index));
}
// Get the type of the translated function.
mlir::FunctionType func_type = translated_func.getFunctionType();
// Then build a call op targeting it (reflecting its result types)
auto expanded_call_op = CallOp::create(
builder, call_ops[0].getLoc(), func_type.getResults(), inputs,
/*args_attrs=*/nullptr, /*res_attrs=*/nullptr,
translated_func.getSymName(),
/*config=*/builder.getStringAttr(""),
/*config_proto=*/builder.getStringAttr(""),
/*executor_type=*/builder.getStringAttr(""));
// Set the output layout attribute on the new call op.
std::vector<std::optional<Layout>> output_layouts;
std::transform(expanded_results.begin(), expanded_results.end(),
std::back_inserter(output_layouts),
[](const ExpandedResults& result) { return result.layout; });
SetLayoutOnOp(expanded_call_op, builder, output_layouts);
expanded_call_op->setAttr(kNumLocalOutputsAttr, num_local_outputs_attr);
std::optional<InferredResourceAttributes> resource_attrs;
if (failed(
GetInferredResourceAttributes(builder, call_ops, &resource_attrs))) {
return mlir::failure();
}
if (resource_attrs) {
expanded_call_op->setAttr(kNewResourceArgLayouts, resource_attrs->layouts);
expanded_call_op->setAttr(kNewResourceLayoutIndices,
resource_attrs->indices);
}
// Return all the values from the new call op.
mlir::Operation::result_range outputs = expanded_call_op.getResults();
if (return_op || outputs.empty()) {
mlir::Location loc = return_op ? return_op.getLoc() : main_func.getLoc();
mlir::func::ReturnOp::create(builder, loc, outputs);
} else {
call_ops[0]->emitOpError("Call had results, but they were not used.");
return mlir::failure();
}
// Update the function's type based on the arguments and return values.
main_func.setFunctionType(GetFunctionType(builder, main_func, outputs));
UpdateEntryFuncAttr(builder, main_func);
// Erase the original main func.
symbol_table.remove(old_main_func);
old_main_func.erase();
// Add the new main function to the module's symbol table, ensuring that it's
// located before all the other functions with the module.
symbol_table.insert(main_func, module_body->begin());
return mlir::success();
}
absl::Status ExtractResultLayouts(
mlir::Operation* op, mlir::func::ReturnOp return_op,
std::vector<ExpandedResults>& expanded_results) {
if (!return_op || (return_op.getNumOperands() == 0)) {
return absl::OkStatus();
}
TF_ASSIGN_OR_RETURN(std::vector<std::optional<Layout>> layouts,
ExtractLayoutFromOp(op));
mlir::Operation::operand_range operands = return_op.getOperands();
for (auto [layout_index, result] : llvm::enumerate(op->getResults())) {
auto search = std::find(operands.begin(), operands.end(), result);
if (search == operands.end()) {
continue;
}
size_t result_index = std::distance(operands.begin(), search);
expanded_results[result_index].layout = layouts[layout_index];
}
return absl::OkStatus();
}
struct DTensorMultiDeviceExpansion
: public impl::DTensorMultiDeviceExpansionBase<
DTensorMultiDeviceExpansion> {
void runOnOperation() override {
mlir::ModuleOp module = getOperation();
auto multi_device_mode =
module->getAttrOfType<mlir::BoolAttr>(dtensor::kEnableMultiDeviceMode);
if (!multi_device_mode || !multi_device_mode.getValue()) {
return; // Skip modules for whom multi-device mode is disabled.
}
mlir::SymbolTable symbol_table(module);
mlir::func::FuncOp main_func =
module.lookupSymbol<mlir::func::FuncOp>(kMainFuncName);
if (!main_func) {
return;
}
std::string translated_func_name =
llvm::formatv("_multi_device_func_{0}_{1}", OpHash(module),
OpHash(main_func))
.str();
mlir::OpBuilder builder = mlir::OpBuilder::atBlockEnd(module.getBody());
mlir::func::FuncOp translated_func =
mlir::func::FuncOp::create(main_func.getLoc(), translated_func_name,
builder.getFunctionType({}, {}));
// build the entry block and return op of the translated function
builder.setInsertionPointToEnd(translated_func.addEntryBlock());
auto translated_terminator_op =
mlir::func::ReturnOp::create(builder, main_func.getLoc());
// so the function has a "terminator" and we can insert it into the module
translated_func.setVisibility(mlir::SymbolTable::Visibility::Private);
symbol_table.insert(translated_func);
ExpandedArgumentMap expanded_arguments_map;
for (unsigned i = 1; i < main_func.getNumArguments(); ++i) {
// Expand all the arguments (in case they're unused).
StatusOr<absl::Span<mlir::Value>> expanded_arguments =
GetExpandedArguments(builder, translated_func, expanded_arguments_map,
main_func.getArgument(i));
if (!expanded_arguments.ok()) {
main_func->emitOpError(
absl::StatusMessageAsCStr(expanded_arguments.status()));
return;
}
}
// Note, we cannot simultaneously walk through the call ops and expand
// them since we'd be creating and removing ops as we walk through them.
llvm::SmallVector<mlir::TF::StatefulPartitionedCallOp, 8> stateful_call_ops;
main_func.walk([&](mlir::Operation* op) {
if (const auto stateful_call_op =
llvm::dyn_cast_or_null<mlir::TF::StatefulPartitionedCallOp>(op)) {
if (stateful_call_op->hasAttr(kLayoutAttr) &&
stateful_call_op->hasAttr(kMeshAttr)) {
stateful_call_ops.emplace_back(stateful_call_op);
}
}
});
// Ensure that all the call ops return results via the same op.
mlir::func::ReturnOp return_op;
if (GetUniqueUserFromOps(stateful_call_ops, &return_op).failed()) {
stateful_call_ops[0]->emitOpError(
"Calls must be used by exactly one return op.");
return;
}
std::vector<ExpandedResults> expanded_results(
return_op ? return_op->getNumOperands() : 0);
for (const mlir::TF::StatefulPartitionedCallOp& stateful_call_op :
stateful_call_ops) {
const absl::Status status =
ExtractResultLayouts(stateful_call_op, return_op, expanded_results);
const StatusOr<std::optional<Mesh>> mesh =
status.ok() ? ExtractDeviceMeshFromOp(stateful_call_op) : status;
if (!(mesh.ok() && *mesh)) {
stateful_call_op->emitOpError("Failed to retrieve op mesh or layout.");
return;
}
const Mesh& target_mesh = **mesh;
if (target_mesh.IsSingleDevice()) {
stateful_call_op->emitOpError(
"Unimplemented, single-device expansion support.");
return;
} else if (target_mesh.is_tpu_mesh()) {
if (mlir::failed(ExpandTPUOperation(
translated_func, return_op, expanded_arguments_map,
expanded_results, target_mesh, stateful_call_op))) {
return;
}
} else {
if (mlir::failed(ExpandOperation(
translated_func, return_op, expanded_arguments_map,
expanded_results, target_mesh, stateful_call_op))) {
return;
}
}
}
std::vector<mlir::Value> results;
llvm::SmallVector<mlir::Attribute, 8> num_local_outputs;
if (return_op) {
for (unsigned i = 0; i < return_op->getNumOperands(); ++i) {
std::vector<mlir::Value>& values = expanded_results[i].results;
int num_outputs;
if (values.empty()) {
results.emplace_back(return_op->getOperand(i));
num_outputs = 1;
} else {
results.insert(results.end(), values.begin(), values.end());
num_outputs = values.size();
}
num_local_outputs.emplace_back(builder.getI64IntegerAttr(num_outputs));
}
}
mlir::ArrayAttr num_local_outputs_attr =
builder.getArrayAttr(num_local_outputs);
// update the operands of the translated return op
translated_terminator_op->setOperands(results);
// and, update the function's type accordingly
translated_func.setFunctionType(GetFunctionType(
builder, translated_func, absl::Span<mlir::Value>(results)));
UpdateEntryFuncAttr(builder, translated_func);
mlir::LogicalResult status = BuildOuterMainFunc(
module, main_func, translated_func, return_op, expanded_results,
num_local_outputs_attr, stateful_call_ops);
if (mlir::failed(status)) {
return;
}
}
};
} // namespace
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorMultiDeviceExpansionPass() {
return std::make_unique<DTensorMultiDeviceExpansion>();
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/dtensor_remove_dtensorlayout.cc
+47
View File
@@ -0,0 +1,47 @@
/* Copyright 2022 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" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "tensorflow/dtensor/mlir/dtensor_dialect/ir/dialect.h"
#include "tensorflow/dtensor/mlir/op_utils.h"
namespace tensorflow {
namespace dtensor {
namespace {
#define GEN_PASS_DEF_DTENSORREMOVEDTENSORLAYOUTPASS
#include "tensorflow/dtensor/mlir/dtensor_passes.h.inc"
class DTensorRemoveDTensorLayoutPass
: public impl::DTensorRemoveDTensorLayoutPassBase<
DTensorRemoveDTensorLayoutPass> {
public:
void runOnOperation() override {
RemoveDTensorLayoutOps(getOperation(), /*remove_xla_spmd_layouts=*/true);
}
};
} // namespace
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorRemoveDTensorLayoutPass() {
return std::make_unique<DTensorRemoveDTensorLayoutPass>();
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/dtensor_replace_auxiliary_layout_op.cc
+47
View File
@@ -0,0 +1,47 @@
/* Copyright 2022 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" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "tensorflow/dtensor/mlir/op_utils.h"
namespace tensorflow {
namespace dtensor {
namespace {
#define GEN_PASS_DEF_DTENSORREPLACEAUXILIARYDTENSORLAYOUTOPPASS
#include "tensorflow/dtensor/mlir/dtensor_passes.h.inc"
class DTensorReplaceAuxiliaryDTensorLayoutOpPass
: public impl::DTensorReplaceAuxiliaryDTensorLayoutOpPassBase<
DTensorReplaceAuxiliaryDTensorLayoutOpPass> {
public:
void runOnOperation() override {
mlir::ModuleOp module = getOperation();
if (mlir::failed(ReplaceAuxiliaryDTensorLayoutOpsWithIdentity(module)))
return signalPassFailure();
}
};
} // namespace
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorReplaceAuxiliaryDTensorLayoutOpPass() {
return std::make_unique<DTensorReplaceAuxiliaryDTensorLayoutOpPass>();
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/dtensor_replace_relayout_with_identity.cc
+57
View File
@@ -0,0 +1,57 @@
/* Copyright 2022 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/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "tensorflow/dtensor/mlir/ir/tf_dtensor.h"
namespace tensorflow {
namespace dtensor {
namespace {
#define GEN_PASS_DEF_DTENSORREPLACERELAYOUTWITHIDENTITYPASS
#include "tensorflow/dtensor/mlir/dtensor_passes.h.inc"
class DTensorReplaceRelayoutWithIdentityPass
: public impl::DTensorReplaceRelayoutWithIdentityPassBase<
DTensorReplaceRelayoutWithIdentityPass> {
public:
void runOnOperation() override {
mlir::func::FuncOp function = getOperation();
function.walk([&](mlir::TF::RelayoutOp relayout_op) {
mlir::OpBuilder builder(relayout_op);
// Inserts an IdentityOp at the position of the relayout_op with the same
// attributes as the relayout_op.
auto new_identity = mlir::TF::IdentityOp::create(
builder, relayout_op->getLoc(), relayout_op.getType(),
relayout_op.getInput(), relayout_op->getAttrs());
relayout_op.getOutput().replaceAllUsesWith(new_identity.getOutput());
relayout_op.erase();
});
}
};
} // namespace
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateDTensorReplaceRelayoutWithIdentityPass() {
return std::make_unique<DTensorReplaceRelayoutWithIdentityPass>();
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/dtensor_send_recv.cc
+924
View File
@@ -0,0 +1,924 @@
/* Copyright 2022 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/dtensor/mlir/dtensor_send_recv.h"
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <optional>
#include <string>
#include <vector>
#include "absl/container/flat_hash_set.h"
#include "absl/status/status.h"
#include "absl/types/span.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/FormatVariadic.h"
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/Support/DebugStringHelper.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_attributes.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/collection_ops_util.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/convert_tensor.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/dtensor/cc/constants.h"
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/collectives.h"
#include "tensorflow/dtensor/mlir/device_utils.h"
#include "tensorflow/dtensor/mlir/ir/tf_dtensor.h"
#include "tensorflow/dtensor/mlir/layout_parsing.h"
#include "tensorflow/dtensor/mlir/op_utils.h"
#include "tensorflow/dtensor/mlir/shape_utils.h"
#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
#include "tensorflow/dtensor/mlir/value_utils.h"
namespace tensorflow {
namespace dtensor {
namespace {
bool IsStringType(mlir::Type type) {
if (mlir::isa<mlir::TF::StringType>(type)) return true;
auto sub_type = mlir::dyn_cast<mlir::TF::TensorFlowTypeWithSubtype>(type);
if (!sub_type) return false;
bool has_string =
llvm::any_of(sub_type.GetSubtypes(), [](mlir::TensorType type) {
return mlir::isa<mlir::TF::StringType>(type.getElementType());
});
return has_string;
}
// Returns compilation key placeholder. This placeholder will be replaced with
// output of TPUCompile op during TPURewrite pass. Program key (output of
// TPUCompile op) is used to differentiate TPU computation from which to receive
// data.
mlir::Value GetOrCreateCompilationKey(mlir::Operation* op) {
mlir::Value key;
auto cluster = op->getParentOfType<mlir::tf_device::ClusterOp>();
assert(cluster);
cluster.walk(
[&](mlir::TF::_XlaCompileMlirPlaceholderProgramKeyOp compilation_key) {
key = compilation_key.getProgram();
});
if (key) return key;
mlir::OpBuilder builder(&cluster.GetBody().front());
auto result_type =
mlir::RankedTensorType::get({3}, builder.getType<mlir::TF::StringType>());
auto new_compilation_key =
mlir::TF::_XlaCompileMlirPlaceholderProgramKeyOp::create(
builder, cluster.getLoc(), /*program=*/result_type,
llvm::ArrayRef<mlir::Value>{});
return new_compilation_key.getProgram();
}
} // namespace
StatusOr<mlir::Value> GetDeviceOrdinal(const Mesh& mesh,
const mlir::Location& loc,
mlir::func::FuncOp function,
mlir::OpBuilder* builder,
bool return_int64_type) {
// Create as many entries as the number of devices in the entire mesh.
llvm::SmallVector<int32_t, 4> device_id_to_ordinal(mesh.num_devices(), 0);
// Only fill in entries with indices equal to local device IDs. For TPUs,
// there are usually 8 local devices.
for (int i = 0; i < mesh.local_device_ids().size(); ++i) {
device_id_to_ordinal[mesh.local_device_ids()[i]] = i;
}
// Slice out the device ordinal using the device ID as index.
TF_ASSIGN_OR_RETURN(mlir::Value device_id, DeviceId(function));
mlir::TF::SliceOp device_ordinal = mlir::TF::SliceOp::create(
*builder, loc,
/*output=*/EffectivelyScalarR1Type(builder->getIntegerType(32)),
/*input=*/IntConst(*builder, loc, device_id_to_ordinal),
/*begin=*/
mlir::TF::collection_ops_util::ReshapeScalarToSizeType(*builder,
device_id, loc),
/*size=*/IntConst(*builder, loc, {1}));
mlir::Value device_ordinal_scalar =
ReshapeSizeTypeToScalar(*builder, loc, device_ordinal);
if (return_int64_type) {
device_ordinal_scalar = mlir::TF::CastOp::create(
*builder, loc, mlir::RankedTensorType::get({}, builder->getI64Type()),
device_ordinal_scalar);
}
return device_ordinal_scalar;
}
StatusOr<mlir::Operation*> LowerDTensorSendToTFOp(
const Layout& send_input_layout, mlir::Value send_input,
mlir::TF::DTensorSend dtensor_send) {
mlir::OpBuilder builder(dtensor_send);
builder.setInsertionPointAfter(send_input.getDefiningOp());
std::string tensor_name = dtensor_send.getKey().str();
Mesh target_mesh = dtensor_send.getTargetMesh();
absl::Span<const std::string> sending_devices =
send_input_layout.mesh().local_devices();
absl::Span<const std::string> receiving_devices = target_mesh.local_devices();
mlir::Operation* lowered_send_op;
lowered_send_op = mlir::TF::_HostSendOp::create(
builder, send_input.getLoc(), send_input, tensor_name, sending_devices[0],
/*send_device_incarnation=*/0, receiving_devices[0],
/*client_terminated=*/false);
dtensor_send.erase();
return lowered_send_op;
}
// Lowers DTensorSend Op to either one of XlaSendFromHost op or XlaSendToHost,
// depending on the src mesh cluster.
StatusOr<mlir::Operation*> LowerDTensorSendToXlaOp(
const Layout& send_input_layout, mlir::Value send_input,
mlir::TF::DTensorSend dtensor_send, bool send_from_device_zero) {
const bool send_from_cpu = !send_input_layout.mesh().is_tpu_mesh();
mlir::OpBuilder builder(dtensor_send);
mlir::Location loc = dtensor_send.getLoc();
mlir::Operation* lowered_send_op;
if (send_from_cpu) {
llvm::SmallVector<mlir::Value, 4> value_to_send{send_input};
mlir::OpBuilder::InsertPoint insertion_point = builder.saveInsertionPoint();
mlir::Value program_key = GetOrCreateCompilationKey(dtensor_send);
builder.restoreInsertionPoint(insertion_point);
mlir::Value device_ordinal;
if (send_from_device_zero) {
// For CopyToMesh, we currently only support sending from host device 0
// to target TPUs.
device_ordinal = CreateIntScalarConst(0, builder, loc);
} else {
// For special topologies, always send from CPU device i to TPU device i.
auto send_cluster =
dtensor_send->getParentOfType<mlir::tf_device::ClusterOp>();
if (!send_cluster) {
return absl::InvalidArgumentError(
"DTensorSend is not inside a ClusterOp");
}
auto send_func = send_cluster->getParentOfType<mlir::func::FuncOp>();
if (!send_func) {
return absl::InvalidArgumentError("DTensorSend is not inside a FuncOp");
}
TF_ASSIGN_OR_RETURN(
device_ordinal,
GetDeviceOrdinal(send_input_layout.mesh(), loc, send_func, &builder));
}
// Create XlaSendFromHostV2 op
lowered_send_op = mlir::TF::_XlaSendFromHostV2Op::create(
builder, loc, value_to_send, program_key, device_ordinal,
dtensor_send.getKey());
} else {
// Note that for ops running in XLA/TPU, device ordinal input is not needed.
lowered_send_op = mlir::TF::XlaSendToHostOp::create(
builder, loc, send_input, dtensor_send.getKey());
}
dtensor_send.erase();
return lowered_send_op;
}
// Creates a shape attribute of the local shape version of RecvFromOp's result.
StatusOr<mlir::TF::ShapeAttr> GetDTensorRecvLocalShapeAttr(
mlir::TF::DTensorRecv dtensor_recv) {
if (dtensor_recv->getNumResults() != 1) {
return absl::InvalidArgumentError(
"XlaRecvFromHostOp must have exactly one result.");
}
TF_ASSIGN_OR_RETURN(std::vector<Layout> layouts,
ExtractRequiredLayoutFromOp(dtensor_recv));
if (layouts.empty() || layouts.size() > 1) {
return absl::InvalidArgumentError(
"invalid layout for XlaRecvFromHostOp specified");
}
auto result_type = dtensor_recv->getResult(0);
TF_ASSIGN_OR_RETURN(auto result_shape, GetShapeOfValue(result_type));
auto local_shape = layouts[0].LocalShapeFromGlobalShape(result_shape);
return mlir::TF::ShapeAttr::get(result_type.getContext(), local_shape);
}
// Lowers DTensorRecv op to either one of XlaRecvAtHost or XlaRecvFromHost,
// depending on src mesh cluster configuration. `output_type` can be set to the
// specific local tensor type needed, if different from the Recv op output type.
StatusOr<mlir::Operation*> LowerDTensorRecvToXlaOp(
mlir::TF::DTensorRecv dtensor_recv, mlir::Type output_type) {
const bool recv_at_cpu = dtensor_recv.getMesh().is_cpu_mesh();
mlir::Operation* recv_xla_op = nullptr;
mlir::OpBuilder builder(dtensor_recv);
if (recv_at_cpu) {
// Create XlaRecvAtHostV2 op.
llvm::SmallVector<mlir::Type, 4> output_types{output_type};
auto recv_cluster =
dtensor_recv->getParentOfType<mlir::tf_device::ClusterOp>();
TF_ASSIGN_OR_RETURN(std::optional<Mesh> mesh,
ExtractDeviceMeshFromOp(recv_cluster));
if (!mesh.has_value())
return absl::InvalidArgumentError(
"failed to get device ordinal as mesh for operation is not "
"specified.");
mlir::OpBuilder builder(&recv_cluster.GetBody().front());
TF_ASSIGN_OR_RETURN(
mlir::Value device_ordinal,
GetDeviceOrdinal(*mesh, recv_cluster.getLoc(),
recv_cluster->getParentOfType<mlir::func::FuncOp>(),
&builder));
auto program_key = GetOrCreateCompilationKey(dtensor_recv);
builder.setInsertionPoint(dtensor_recv);
recv_xla_op = mlir::TF::_XlaRecvAtHostV2Op::create(
builder, dtensor_recv.getLoc(), output_types,
/*dynamic_key=*/program_key, device_ordinal, dtensor_recv.getKeyAttr());
} else {
TF_ASSIGN_OR_RETURN(auto local_shape_attr,
GetDTensorRecvLocalShapeAttr(dtensor_recv));
// Create XlaRecvFromHost op.
recv_xla_op = mlir::TF::XlaRecvFromHostOp::create(
builder, dtensor_recv.getLoc(), output_type, local_shape_attr,
dtensor_recv.getKeyAttr());
}
assert(recv_xla_op);
// TODO(hongjunchoi): After receiving tensor, convert tensor to requested
// layout with EmitRelayout.
return recv_xla_op;
}
// Lowers DTensorRecv op to either one of XlaRecvAtHost or XlaRecvFromHost,
// depending on src mesh cluster configuration.
StatusOr<mlir::Operation*> LowerDTensorRecvToXlaOp(
mlir::TF::DTensorRecv dtensor_recv) {
return LowerDTensorRecvToXlaOp(dtensor_recv, dtensor_recv.getType());
}
StatusOr<mlir::Operation*> LowerDTensorRecvToXlaOp(
const Mesh&, mlir::TF::DTensorRecv dtensor_recv, mlir::Type output_type) {
return LowerDTensorRecvToXlaOp(dtensor_recv, output_type);
}
// Lowers a DTensorSend Op from a CPU to a TF Send op.
StatusOr<mlir::Operation*> LowerDTensorSendFromCPUToTFOp(
const Layout& send_input_layout, mlir::Value send_input,
mlir::TF::DTensorSend dtensor_send) {
mlir::OpBuilder builder(dtensor_send);
builder.setInsertionPointAfter(send_input.getDefiningOp());
llvm::SmallVector<mlir::Value, 4> value_to_send{send_input};
// Create multiple send from host. There should be #number of local
// devices(in target mesh) number of sends.
absl::Span<const std::string> sending_devices =
send_input_layout.mesh().local_devices();
Mesh target_mesh = dtensor_send.getTargetMesh();
absl::Span<const std::string> receiving_devices = target_mesh.local_devices();
std::string tensor_name = dtensor_send.getKey().str();
mlir::Operation* lowered_send_op;
for (size_t i = 0; i < receiving_devices.size(); ++i)
lowered_send_op = mlir::TF::_HostSendOp::create(
builder, send_input.getLoc(), dtensor_send.getInput(), tensor_name,
sending_devices[0],
/*send_device_incarnation=*/0, receiving_devices[i]);
dtensor_send.erase();
return lowered_send_op;
}
// Lowers DTensorRecv op to TF Recv Op.
StatusOr<mlir::Operation*> LowerDTensorRecvFromCPUToTFOp(
const Mesh& send_mesh, mlir::TF::DTensorRecv dtensor_recv) {
const Mesh& recv_mesh = dtensor_recv.getMesh();
auto recv_cluster =
dtensor_recv->getParentOfType<mlir::tf_device::ClusterOp>();
mlir::OpBuilder builder(&recv_cluster.GetBody().front());
llvm::SmallVector<mlir::Type, 4> output_types{dtensor_recv.getType()};
builder.setInsertionPoint(dtensor_recv);
std::string tensor_name = dtensor_recv.getKey().str();
absl::Span<const std::string> sending_devices = send_mesh.local_devices();
absl::Span<const std::string> receiving_devices = recv_mesh.local_devices();
mlir::Operation* lowered_recv_op;
mlir::Location loc = dtensor_recv.getLoc();
for (size_t i = 0; i < receiving_devices.size(); ++i)
lowered_recv_op = mlir::TF::_HostRecvOp::create(
builder, loc, dtensor_recv.getType(), tensor_name, sending_devices[0],
/*send_device_incarnation=*/0, receiving_devices[i]);
// Replace dtensor_recv with newly created recv op and remove DTensorRecv op.
assert(lowered_recv_op);
dtensor_recv.replaceAllUsesWith(lowered_recv_op);
dtensor_recv.erase();
return lowered_recv_op;
}
StatusOr<mlir::Operation*> LowerDTensorRecvToTFOp(
const Mesh& send_mesh, mlir::TF::DTensorRecv dtensor_recv,
mlir::Type output_type) {
const Mesh& recv_mesh = dtensor_recv.getMesh();
auto recv_cluster =
dtensor_recv->getParentOfType<mlir::tf_device::ClusterOp>();
mlir::OpBuilder builder(&recv_cluster.GetBody().front());
builder.setInsertionPoint(dtensor_recv);
std::string tensor_name = dtensor_recv.getKey().str();
absl::Span<const std::string> sending_devices = send_mesh.local_devices();
absl::Span<const std::string> receiving_devices = recv_mesh.local_devices();
mlir::Location loc = dtensor_recv.getLoc();
mlir::Operation* lowered_recv_op = mlir::TF::_HostRecvOp::create(
builder, loc, output_type, tensor_name, sending_devices[0],
/*send_device_incarnation=*/0, receiving_devices[0]);
return lowered_recv_op;
}
namespace {
template <typename It, typename Fn>
llvm::SmallVector<mlir::Attribute, 4> GenerateBranches(
mlir::Operation* op, mlir::SymbolTable& symbol_table,
llvm::ArrayRef<mlir::Type> result_types, const char* fmt, const It& values,
const Fn& fn) {
llvm::SmallVector<mlir::Attribute, 4> branches;
for (const auto& it : llvm::enumerate(values)) {
// Builds the restore op on device_id.
mlir::OpBuilder builder(op);
auto func_type = mlir::FunctionType::get(
builder.getContext(), op->getOperandTypes(), result_types);
mlir::Location location = op->getLoc();
mlir::func::FuncOp func_op = mlir::func::FuncOp::create(
location, llvm::formatv(fmt, OpName(op), OpHash(op), it.index()).str(),
func_type, llvm::ArrayRef<mlir::NamedAttribute>{});
func_op.setVisibility(mlir::SymbolTable::Visibility::Private);
symbol_table.insert(func_op);
mlir::Block* fn_block = func_op.addEntryBlock();
mlir::OpBuilder fn_builder = mlir::OpBuilder::atBlockBegin(fn_block);
mlir::BlockArgument arg = (func_op.getNumArguments() > 0)
? func_op.getArgument(0)
: mlir::BlockArgument{};
auto branch_op = fn(fn_builder, location, arg, it.value());
mlir::func::ReturnOp::create(fn_builder, location, branch_op->getResults());
branches.push_back(mlir::SymbolRefAttr::get(func_op));
}
return branches;
}
} // namespace
StatusOr<mlir::Operation*> LowerOneToOneDTensorSendToTFHostSend(
const Layout& send_layout, const Mesh& recv_mesh,
mlir::TF::DTensorSend dtensor_send) {
const auto& send_mesh = send_layout.mesh();
bool i32_copy =
dtensor_send.getInput().getType().getElementType().isInteger(32);
auto module = dtensor_send->getParentOfType<mlir::ModuleOp>();
mlir::SymbolTable symbol_table(module);
auto device_pairs =
llvm::zip(send_mesh.local_devices(), recv_mesh.local_devices());
mlir::OpBuilder builder(dtensor_send);
auto send_cluster =
dtensor_send->getParentOfType<mlir::tf_device::ClusterOp>();
auto send_fn = send_cluster->getParentOfType<mlir::func::FuncOp>();
TF_ASSIGN_OR_RETURN(std::optional<Mesh> mesh,
ExtractDeviceMeshFromOp(send_cluster));
TF_ASSIGN_OR_RETURN(
mlir::Value device_ordinal,
GetDeviceOrdinal(*mesh, dtensor_send.getLoc(), send_fn, &builder,
/*return_int64_type=*/false));
mlir::StringAttr tensor_name =
builder.getStringAttr(dtensor_send.getKey().str());
auto branches = GenerateBranches(
dtensor_send, symbol_table, llvm::ArrayRef<mlir::Type>{},
"{0}_send_{1}_{2}", device_pairs,
[&](mlir::OpBuilder& op_builder, auto& loc, mlir::BlockArgument arg,
auto device_pair) -> mlir::Operation* {
auto func_op = mlir::dyn_cast_or_null<mlir::func::FuncOp>(
arg.getOwner()->getParentOp());
func_op.setArgAttr(arg.getArgNumber(), kCustomDeviceAttr,
op_builder.getStringAttr(send_layout.ToString()));
mlir::Value val = arg;
if (i32_copy) {
auto val_type = mlir::cast<mlir::TensorType>(val.getType());
val = mlir::TF::CastOp::create(
op_builder, loc,
mlir::RankedTensorType::get(val_type.getShape(),
op_builder.getIntegerType(64)),
val)
->getResult(0);
}
return mlir::TF::_HostSendOp::create(
op_builder, loc, val, tensor_name, std::get<0>(device_pair),
/*send_device_incarnation=*/0, std::get<1>(device_pair));
});
mlir::Operation* case_op =
mlir::TF::CaseOp::create(builder, dtensor_send.getLoc(),
/*output=*/llvm::ArrayRef<mlir::Type>{},
/*branch_index=*/device_ordinal,
/*input=*/dtensor_send->getOperands(),
/*branches=*/builder.getArrayAttr(branches),
/*is_stateless=*/builder.getBoolAttr(false));
// erase the send op here iff targeting a gpu
// otherwise there will be 'op not within cluster' error(s)
if (recv_mesh.device_type() == "GPU") {
dtensor_send.erase();
}
return case_op;
}
StatusOr<mlir::Operation*> LowerOneToOneDTensorRecvToTFHostRecv(
const Mesh& send_mesh, const Layout& recv_layout,
mlir::TF::DTensorRecv dtensor_recv) {
auto module = dtensor_recv->getParentOfType<mlir::ModuleOp>();
const auto& recv_mesh = recv_layout.mesh();
mlir::SymbolTable symbol_table(module);
auto device_pairs =
llvm::zip(send_mesh.local_devices(), recv_mesh.local_devices());
mlir::OpBuilder builder(dtensor_recv);
auto recv_cluster =
dtensor_recv->getParentOfType<mlir::tf_device::ClusterOp>();
auto recv_fn = recv_cluster->getParentOfType<mlir::func::FuncOp>();
TF_ASSIGN_OR_RETURN(std::optional<Mesh> mesh,
ExtractDeviceMeshFromOp(recv_cluster));
TF_ASSIGN_OR_RETURN(
mlir::Value device_ordinal,
GetDeviceOrdinal(*mesh, recv_cluster.getLoc(), recv_fn, &builder,
/*return_int64_type=*/false));
mlir::TensorType recv_type = dtensor_recv.getType();
bool i32_copy = recv_type.getElementType().isInteger(32);
TF_ASSIGN_OR_RETURN(mlir::TensorType local_recv_type,
LocalTypeFromGlobalType(recv_layout, recv_type));
mlir::TensorType local_output_type =
i32_copy ? mlir::RankedTensorType::get(local_recv_type.getShape(),
builder.getIntegerType(64))
: local_recv_type;
mlir::StringAttr tensor_name =
builder.getStringAttr(dtensor_recv.getKey().str());
auto branches = GenerateBranches(
dtensor_recv, symbol_table, llvm::ArrayRef<mlir::Type>{local_output_type},
"{0}_receive_{1}_{2}", device_pairs,
[&](mlir::OpBuilder& op_builder, auto& loc, auto _,
auto device_pair) -> mlir::Operation* {
auto recv_op = mlir::TF::_HostRecvOp::create(
op_builder, loc, local_output_type, tensor_name,
std::get<0>(device_pair),
/*send_device_incarnation=*/0, std::get<1>(device_pair));
SetSingleLayoutOnOp(recv_op, recv_layout);
return recv_op;
});
mlir::Operation* case_op = mlir::TF::CaseOp::create(
builder, dtensor_recv.getLoc(),
/*output=*/llvm::ArrayRef<mlir::Type>{local_output_type},
/*branch_index=*/device_ordinal,
/*input=*/dtensor_recv->getOperands(),
/*branches=*/builder.getArrayAttr(branches),
/*is_stateless=*/builder.getBoolAttr(false));
mlir::Operation* lowered_recv;
if (i32_copy) {
lowered_recv = mlir::TF::CastOp::create(
builder, dtensor_recv.getLoc(), local_recv_type, case_op->getResult(0));
} else {
lowered_recv = case_op;
}
dtensor_recv.getOutput().replaceAllUsesWith(lowered_recv->getResult(0));
dtensor_recv.erase();
return lowered_recv;
}
namespace {
bool IsTpuToHostMeshTransfer(const Mesh& send_mesh, const Mesh& recv_mesh) {
// Check tensor is being transferred between CPU <-> TPU.
if (!(send_mesh.is_tpu_mesh() && recv_mesh.is_cpu_mesh()) &&
!(recv_mesh.is_tpu_mesh() && send_mesh.is_cpu_mesh()))
return false;
// Check tensor transfer is happening between TPU and its host mesh.
return ((send_mesh.is_tpu_mesh() &&
send_mesh.tpu_host_mesh() == recv_mesh.ToString()) ||
(recv_mesh.is_tpu_mesh() &&
recv_mesh.tpu_host_mesh() == send_mesh.ToString()));
}
bool IsGpuToHostMeshTransfer(const Mesh& send_mesh, const Mesh& recv_mesh) {
return ((send_mesh.device_type() == "GPU") && recv_mesh.is_cpu_mesh()) ||
((recv_mesh.device_type() == "GPU") && send_mesh.is_cpu_mesh());
}
// Returns whether send/recv layout represents send/recv of tensor between
// i-th TPU device and i-th device of the host mesh. Host mesh represents the
// CPU devices that are 1-to-1 mapped with the TPU mesh devices, having the same
// global and local device IDs.
bool IsOneToOneMeshTransfer(const Layout& send_layout,
const Layout& recv_layout) {
const Mesh& send_mesh = send_layout.mesh();
const Mesh& recv_mesh = recv_layout.mesh();
// Check local device IDs are fully matching so that there is no cross-host
// transfer.
if (send_mesh.local_device_ids() != recv_mesh.local_device_ids())
return false;
return send_layout.GetShardVector() == recv_layout.GetShardVector();
}
// Returns whether to lower DTensorSend/DTensorRecv op to xla backend ops.
// Xla backend ops are used when either sending/receiving device uses XLA
// compiler.
bool SendRecvOpUsesXla(const Mesh& send_mesh, const Mesh& recv_mesh) {
assert(!(send_mesh.is_tpu_mesh() && recv_mesh.is_tpu_mesh()));
return (send_mesh.is_tpu_mesh() || recv_mesh.is_tpu_mesh());
}
} // namespace
// FIXME(b/271292250): Remove the recv_op argument.
StatusOr<mlir::Operation*> LowerDTensorSend(mlir::Operation* send_op,
mlir::Operation* recv_op) {
auto dtensor_send = llvm::cast<mlir::TF::DTensorSend>(send_op);
TF_ASSIGN_OR_RETURN(
const Layout input_layout,
ExtractRequiredLayoutFromOperand(dtensor_send.getInput()));
const Mesh& input_mesh = input_layout.mesh();
const Mesh& target_mesh = dtensor_send.getTargetMesh();
auto layout_attr =
dtensor_send->getAttrOfType<mlir::dtensor::LayoutAttr>(kTargetLayoutAttr);
if (!layout_attr) {
return absl::InvalidArgumentError("target_layout is not found");
}
const Layout& recv_layout = layout_attr.getValue();
bool one_to_one = IsOneToOneMeshTransfer(input_layout, recv_layout);
// Force string type to not use the allreduce/broadcast optimization as there
// is no string type allreduce.
bool is_string_type =
IsStringType(dtensor_send.getInput().getType().getElementType());
// Is tensor transfer is from TPU mesh to host mesh and send layout and recv
// layout is identical, then tensor from each source device is sent to
// target device asynchronously.
mlir::Operation* lowered_send;
if (IsTpuToHostMeshTransfer(input_mesh, target_mesh) && one_to_one) {
TF_ASSIGN_OR_RETURN(lowered_send,
LowerDTensorSendToXlaOp(
input_layout, dtensor_send.getInput(), dtensor_send,
/*send_from_device_zero=*/false));
} else if (IsGpuToHostMeshTransfer(input_mesh, target_mesh) &&
(one_to_one &&
(!recv_layout.IsFullyReplicated() || is_string_type))) {
TF_ASSIGN_OR_RETURN(lowered_send,
LowerOneToOneDTensorSendToTFHostSend(
input_layout, target_mesh, dtensor_send));
} else {
// Calculate input tensor layout of data to send and target fully replicated
// layout. For now, we ensure that all data transfer happen with fully
// replicated tensors.
const int rank = ValueRank(dtensor_send.getInput());
const Layout target_layout = Layout::ReplicatedOnMesh(input_mesh, rank);
// Convert tensor to send to replicated layout.
mlir::OpBuilder builder(dtensor_send);
TF_ASSIGN_OR_RETURN(mlir::Value send_input,
EmitAllGather(builder, dtensor_send.getInput(),
input_layout, target_layout));
// Insert control flow such that only device with device ordinal == 0 sends
// the tensor data across mesh.
auto send_cluster =
dtensor_send->getParentOfType<mlir::tf_device::ClusterOp>();
TF_ASSIGN_OR_RETURN(std::optional<Mesh> mesh,
ExtractDeviceMeshFromOp(send_cluster));
if (!mesh.has_value()) {
return absl::InvalidArgumentError(
"failed to lower DTensor CopyToMesh op as sending side mesh is not "
"specified.");
}
mlir::Location loc = dtensor_send.getLoc();
TF_ASSIGN_OR_RETURN(
mlir::Value device_ordinal,
GetDeviceOrdinal(*mesh, loc,
send_cluster->getParentOfType<mlir::func::FuncOp>(),
&builder));
mlir::Value predicate = mlir::TF::EqualOp::create(
builder, loc, device_ordinal, CreateIntScalarConst(0, builder, loc),
/*incompatible_shape_error=*/builder.getBoolAttr(true));
auto send_if = mlir::TF::IfRegionOp::create(
builder, loc, llvm::SmallVector<mlir::Type, 4>{}, predicate,
/*is_stateless=*/builder.getBoolAttr(true),
GetUniqueControlflowFnName("copy_to_mesh_send_if_then", builder),
GetUniqueControlflowFnName("copy_to_mesh_send_if_else", builder));
// Create empty else branch region.
auto& else_branch = send_if.getElseBranch();
else_branch.push_back(new mlir::Block);
builder.setInsertionPointToEnd(&else_branch.front());
mlir::TF::YieldOp::create(builder, loc,
/*operands=*/llvm::ArrayRef<mlir::Value>{});
// Create then branch region with DTensorSend op.
auto& then_branch = send_if.getThenBranch();
then_branch.push_back(new mlir::Block);
builder.setInsertionPointToEnd(&then_branch.front());
auto yield = mlir::TF::YieldOp::create(
builder, loc, /*operands=*/llvm::ArrayRef<mlir::Value>{});
dtensor_send->moveBefore(yield);
// Lower DTensorSend op to actual TF op.
const Mesh& recv_mesh = recv_layout.mesh();
if (SendRecvOpUsesXla(input_layout.mesh(), recv_mesh)) {
// Lower DTensorSend op to Xla Send ops.
TF_ASSIGN_OR_RETURN(
lowered_send,
LowerDTensorSendToXlaOp(input_layout, send_input, dtensor_send,
/*send_from_device_zero=*/true));
} else if (input_layout.mesh().is_cpu_mesh() && recv_mesh.is_cpu_mesh()) {
// Lower DTensorSend op to TF Host Send op.
TF_ASSIGN_OR_RETURN(
lowered_send, LowerDTensorSendFromCPUToTFOp(input_layout, send_input,
dtensor_send));
} else {
mlir::TensorType send_type =
mlir::cast<mlir::TensorType>(send_input.getType());
if (!recv_mesh.is_cpu_mesh() &&
send_type.getElementType().isInteger(32)) {
builder.setInsertionPointAfter(send_input.getDefiningOp());
auto cast_to_int64 = mlir::TF::CastOp::create(
builder, send_input.getLoc(),
mlir::RankedTensorType::get(send_type.getShape(),
builder.getIntegerType(64)),
send_input);
send_input = cast_to_int64->getResult(0);
}
TF_ASSIGN_OR_RETURN(
lowered_send,
LowerDTensorSendToTFOp(input_layout, send_input, dtensor_send));
}
}
return lowered_send;
}
// FIXME(b/271292250): Remove the send_op argument.
StatusOr<mlir::Operation*> LowerDTensorRecv(mlir::Operation* send_op,
mlir::Operation* recv_op) {
auto dtensor_recv = llvm::cast<mlir::TF::DTensorRecv>(recv_op);
TF_ASSIGN_OR_RETURN(const Layout output_layout,
ExtractRequiredSingleLayoutFromOp(recv_op));
mlir::Operation* lowered_recv;
auto layout_attr =
dtensor_recv->getAttrOfType<mlir::dtensor::LayoutAttr>(kSourceLayoutAttr);
if (!layout_attr) {
return absl::InvalidArgumentError("source_layout is not found");
}
const Layout send_layout = layout_attr.getValue();
const Mesh send_mesh = send_layout.mesh();
const Mesh& recv_mesh = dtensor_recv.getMesh();
const Layout& recv_layout = output_layout;
mlir::OpBuilder builder(dtensor_recv);
bool cpu_to_cpu = recv_mesh.is_cpu_mesh() && send_mesh.is_cpu_mesh();
bool one_to_one = IsOneToOneMeshTransfer(send_layout, recv_layout);
bool send_recv_xla = SendRecvOpUsesXla(send_mesh, recv_mesh);
// Force string type to not use the allreduce/broadcast optimization as there
// is no string type allreduce.
bool is_string_type = IsStringType(dtensor_recv.getType().getElementType());
if (IsGpuToHostMeshTransfer(send_mesh, recv_mesh) &&
(one_to_one && (!recv_layout.IsFullyReplicated() || is_string_type))) {
TF_ASSIGN_OR_RETURN(lowered_recv,
LowerOneToOneDTensorRecvToTFHostRecv(
send_mesh, recv_layout, dtensor_recv));
// erase the send op here iff not targeting a gpu
if (recv_mesh.device_type() != "GPU") {
send_op->erase();
}
return lowered_recv;
} else if (cpu_to_cpu) {
// Lower DTensorRecv op to TF Host Recv op.
TF_ASSIGN_OR_RETURN(lowered_recv,
LowerDTensorRecvFromCPUToTFOp(send_mesh, dtensor_recv));
} else if ((IsTpuToHostMeshTransfer(send_mesh, recv_mesh) && one_to_one) ||
(send_recv_xla && recv_mesh.is_cpu_mesh())) {
// Recv can be lowered directly for a 1-to-1 transfer between host and
// device (*for XLA/TPUs).
TF_ASSIGN_OR_RETURN(
mlir::TensorType local_output_type,
LocalTypeFromGlobalType(
recv_layout, mlir::cast<mlir::TensorType>(dtensor_recv.getType())));
TF_ASSIGN_OR_RETURN(
lowered_recv, LowerDTensorRecvToXlaOp(dtensor_recv, local_output_type));
dtensor_recv->replaceAllUsesWith(lowered_recv);
dtensor_recv.erase();
} else {
// Choose which receive lowering function to use.
auto lower_fn =
send_recv_xla
? (decltype(&LowerDTensorRecvToTFOp))LowerDTensorRecvToXlaOp
: LowerDTensorRecvToTFOp;
// For other send/recv layouts, the tensor needs to be replicated.
if (!recv_layout.IsFullyReplicated()) {
return absl::InvalidArgumentError(
"CopyToMesh where target mesh is GPU/TPU requires a replicated "
"target layout.");
}
// For Receiving at GPU/TPU, only device 0 (ordinal) receives from the
// host, then it shares the tensor with its peers.
auto recv_cluster =
dtensor_recv->getParentOfType<mlir::tf_device::ClusterOp>();
mlir::Location loc = dtensor_recv.getLoc();
TF_ASSIGN_OR_RETURN(
mlir::Value device_ordinal,
GetDeviceOrdinal(recv_mesh, loc,
recv_cluster->getParentOfType<mlir::func::FuncOp>(),
&builder));
mlir::Value predicate = mlir::TF::EqualOp::create(
builder, loc, device_ordinal, CreateIntScalarConst(0, builder, loc),
/*incompatible_shape_error=*/builder.getBoolAttr(true));
mlir::TensorType recv_type = dtensor_recv.getType();
bool i32_copy = recv_type.getElementType().isInteger(32);
bool need_i32_to_i64_upcast =
i32_copy && !(recv_mesh.is_cpu_mesh() || send_recv_xla);
mlir::TensorType output_type =
need_i32_to_i64_upcast
? mlir::RankedTensorType::get(recv_type.getShape(),
builder.getIntegerType(64))
: recv_type;
auto recv_if = mlir::TF::IfRegionOp::create(
builder, loc, llvm::SmallVector<mlir::Type, 4>{output_type}, predicate,
/*is_stateless=*/builder.getBoolAttr(true),
GetUniqueControlflowFnName("copy_to_mesh_recv_if_then", builder),
GetUniqueControlflowFnName("copy_to_mesh_recv_if_else", builder));
// Create empty else branch region that outputs zeros.
auto& else_branch = recv_if.getElseBranch();
else_branch.push_back(new mlir::Block);
builder.setInsertionPointToEnd(&else_branch.front());
// Create a zero constant.
mlir::Attribute const_attr;
auto output_element_type = output_type.getElementType();
if (output_element_type.isIntOrIndex()) {
if (output_element_type.isInteger(64)) {
const_attr = mlir::DenseIntElementsAttr::get(
output_type, llvm::SmallVector<int64_t>{0});
} else {
const_attr = mlir::DenseIntElementsAttr::get(
output_type, llvm::SmallVector<int32_t>{0});
}
} else if (output_element_type.isBF16()) {
mlir::FloatAttr zero = mlir::FloatAttr::get(output_element_type, 0.);
const_attr = mlir::DenseElementsAttr::get(
output_type, llvm::SmallVector<mlir::Attribute>{zero});
} else if (output_element_type.isF16() || output_element_type.isF32()) {
const_attr = mlir::DenseFPElementsAttr::get(
output_type, llvm::SmallVector<float>{0.0});
} else if (output_element_type.isF64()) {
const_attr = mlir::DenseFPElementsAttr::get(
output_type, llvm::SmallVector<double>{0.0});
} else {
return absl::InvalidArgumentError("unsupported output type");
}
mlir::Value zeros = mlir::TF::ConstOp::create(builder, loc, const_attr);
mlir::TF::YieldOp::create(builder, loc,
/*operands=*/llvm::ArrayRef<mlir::Value>{zeros});
// Create then branch region with DTensorRecv op.
auto& then_branch = recv_if.getThenBranch();
then_branch.push_back(new mlir::Block);
builder.setInsertionPointToEnd(&then_branch.front());
dtensor_recv->moveBefore(&then_branch.front(), then_branch.front().end());
TF_ASSIGN_OR_RETURN(mlir::Operation * xla_recv,
lower_fn(send_mesh, dtensor_recv, output_type));
mlir::TF::YieldOp::create(
builder, loc,
/*operands=*/llvm::ArrayRef<mlir::Value>{xla_recv->getResult(0)});
// Broadcast the received output to all GPU/TPU devices.
mlir::Value if_output = recv_if->getResult(0);
builder.setInsertionPointAfterValue(if_output);
absl::flat_hash_set<std::string> reduced_dims;
for (const auto& mesh_dim : recv_mesh.dims())
reduced_dims.insert(mesh_dim.name);
TF_ASSIGN_OR_RETURN(
lowered_recv, EmitAllReduce(builder, recv_layout, reduced_dims, recv_if,
kReduceOpAdd));
if (need_i32_to_i64_upcast) {
lowered_recv = mlir::TF::CastOp::create(builder, loc, recv_type,
lowered_recv->getResult(0));
}
// Replaces usages of DTensorRecv op with the broadcasted value.
dtensor_recv.getOutput().replaceUsesWithIf(
lowered_recv->getResult(0), [&](mlir::OpOperand& operand) {
return !recv_if->isProperAncestor(operand.getOwner());
});
dtensor_recv.erase();
}
llvm::SmallPtrSet<mlir::Operation*, 4> newly_created_ops;
builder.setInsertionPointAfter(lowered_recv);
TF_ASSIGN_OR_RETURN(
mlir::Value recv_output,
EmitAllScatter(builder, lowered_recv->getResult(0), recv_layout,
output_layout, &newly_created_ops));
lowered_recv->getResult(0).replaceAllUsesExcept(recv_output,
newly_created_ops);
return recv_output.getDefiningOp();
}
StatusOr<mlir::Operation*> LowerDTensorSendAndRecv(mlir::Operation* send_op,
mlir::Operation* recv_op) {
auto dtensor_send = llvm::cast<mlir::TF::DTensorSend>(send_op);
auto dtensor_recv = llvm::dyn_cast<mlir::TF::DTensorRecv>(recv_op);
const Mesh recv_mesh = dtensor_recv.getMesh();
TF_ASSIGN_OR_RETURN(
std::optional<Mesh> send_mesh,
ExtractDeviceMeshFromOp(
send_op->getParentOfType<mlir::tf_device::ClusterOp>()));
if (!send_mesh.has_value())
return absl::InvalidArgumentError(
"failed to get device ordinal as mesh for operation is not "
"specified.");
if (!send_mesh->is_tpu_mesh() && !recv_mesh.is_tpu_mesh()) {
return absl::UnimplementedError(
"Multi-mesh tensor transfer between non-xla devices are not yet "
"supported.");
}
const Layout recv_layout =
Layout::ReplicatedOnMesh(recv_mesh, ValueRank(dtensor_recv.getOutput()));
const Layout send_input_layout =
Layout::ReplicatedOnMesh(*send_mesh, ValueRank(dtensor_send.getInput()));
TF_ASSIGN_OR_RETURN(mlir::Operation * lowered_recv,
LowerDTensorRecvToXlaOp(dtensor_recv));
dtensor_recv->replaceAllUsesWith(lowered_recv);
dtensor_recv.erase();
return LowerDTensorSendToXlaOp(send_input_layout, dtensor_send.getInput(),
dtensor_send,
/*send_from_device_zero=*/false);
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/dtensor_send_recv.h
+127
View File
@@ -0,0 +1,127 @@
/* Copyright 2022 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_DTENSOR_MLIR_DTENSOR_SEND_RECV_H_
#define TENSORFLOW_DTENSOR_MLIR_DTENSOR_SEND_RECV_H_
#include "absl/status/status.h"
#include "llvm/Support/Casting.h"
#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/Value.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/ir/tf_dtensor.h"
namespace tensorflow {
namespace dtensor {
// Given DTensorSend or DTensorRecv op, returns the corresponding DTensorRecv
// or DTensorSend op with the same key.
template <typename DTensorOp>
StatusOr<mlir::Operation*> GetCorrespondingDTensorSendRecvOp(
mlir::ModuleOp module, DTensorOp dtensor_op) {
mlir::Operation* corresponding_op = nullptr;
if (std::is_same_v<DTensorOp, mlir::TF::DTensorSend>) {
module.walk([&](mlir::Operation* op) {
if (auto xla_recv_tpu = llvm::dyn_cast<mlir::TF::XlaRecvFromHostOp>(op)) {
if (dtensor_op.getKey() == xla_recv_tpu.getKey()) {
corresponding_op = op;
return mlir::WalkResult::interrupt();
}
} else if (auto xla_recv_cpu =
llvm::dyn_cast<mlir::TF::_XlaRecvAtHostV2Op>(op)) {
if (dtensor_op.getKey() == xla_recv_cpu.getKey()) {
corresponding_op = op;
return mlir::WalkResult::interrupt();
}
} else if (auto dtensor_recv =
llvm::dyn_cast<mlir::TF::DTensorRecv>(op)) {
if (dtensor_op.getKey() == dtensor_recv.getKey()) {
corresponding_op = op;
return mlir::WalkResult::interrupt();
}
} else if (auto host_recv = llvm::dyn_cast<mlir::TF::_HostRecvOp>(op)) {
if (dtensor_op.getKey() == host_recv.getTensorName()) {
corresponding_op = op;
return mlir::WalkResult::interrupt();
}
}
return mlir::WalkResult::advance();
});
} else {
const bool is_recv = std::is_same_v<DTensorOp, mlir::TF::DTensorRecv>;
if (!is_recv) {
return absl::InternalError(
"Error checking if is same for DTensorOp and DTensorRecv.");
}
module.walk([&](mlir::Operation* op) {
if (auto xla_send_tpu = llvm::dyn_cast<mlir::TF::XlaSendToHostOp>(op)) {
if (dtensor_op.getKey() == xla_send_tpu.getKey()) {
corresponding_op = op;
return mlir::WalkResult::interrupt();
}
} else if (auto xla_send_cpu =
llvm::dyn_cast<mlir::TF::_XlaSendFromHostV2Op>(op)) {
if (dtensor_op.getKey() == xla_send_cpu.getKey()) {
corresponding_op = op;
return mlir::WalkResult::interrupt();
}
} else if (auto dtensor_send =
llvm::dyn_cast<mlir::TF::DTensorSend>(op)) {
if (dtensor_op.getKey() == dtensor_send.getKey()) {
corresponding_op = op;
return mlir::WalkResult::interrupt();
}
} else if (auto host_send = llvm::dyn_cast<mlir::TF::_HostSendOp>(op)) {
if (dtensor_op.getKey() == host_send.getTensorName()) {
corresponding_op = op;
return mlir::WalkResult::interrupt();
}
}
return mlir::WalkResult::advance();
});
}
if (!corresponding_op)
return absl::InvalidArgumentError(
"DTensorSend/DTensorRecv op must have corresponding "
"DTensorRecv/DTensorSend op.");
return corresponding_op;
}
// Lowers DTensorSend to a number of different device-specific ops:
// _HostSend, XlaSendFromHost, XlaSendToHost, etc.
StatusOr<mlir::Operation*> LowerDTensorRecv(mlir::Operation* send_op,
mlir::Operation* recv_op);
// Lowers DTensorRecv to a number of different device-specific ops:
// _HostRecv, XlaRecvAtHost, XlaRecvFromHost, etc.
StatusOr<mlir::Operation*> LowerDTensorSend(mlir::Operation* send_op,
mlir::Operation* recv_op);
// Lowers a DTensorSend and DTensorRecv pair to XLA ops
StatusOr<mlir::Operation*> LowerDTensorSendAndRecv(mlir::Operation* send_op,
mlir::Operation* recv_op);
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_MLIR_DTENSOR_SEND_RECV_H_
tensorflow/dtensor/mlir/dtensor_set_hlo_sharding.cc
+146
View File
@@ -0,0 +1,146 @@
/* Copyright 2022 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 "llvm/Support/Casting.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "xla/xla_data.pb.h"
#include "tensorflow/dtensor/cc/xla_spmd/layout_to_xla_sharding.h"
#include "tensorflow/dtensor/mlir/dtensor_dialect/ir/dialect.h"
#include "tensorflow/dtensor/mlir/ir/tf_dtensor.h"
namespace tensorflow {
namespace dtensor {
namespace {
#define GEN_PASS_DECL_DTENSORSETHLOSHARDINGPASS
#define GEN_PASS_DEF_DTENSORSETHLOSHARDINGPASS
#include "tensorflow/dtensor/mlir/dtensor_passes.h.inc"
mlir::LogicalResult SetHloShardingForInputs(mlir::ModuleOp module,
mlir::OpBuilder builder,
bool check_layout_use_xla_spmd) {
auto result = module.walk([&](mlir::func::FuncOp func_op)
-> mlir::WalkResult {
for (int arg_index = 0; arg_index < func_op.getNumArguments();
++arg_index) {
mlir::BlockArgument arg = func_op.getArgument(arg_index);
for (const auto* user : arg.getUsers()) {
auto layout_op = llvm::dyn_cast<mlir::TF::DTensorLayout>(*user);
if (!layout_op) {
continue;
}
if (check_layout_use_xla_spmd &&
!layout_op.getLayout().mesh().use_xla_spmd()) {
return layout_op.emitOpError(
"Found a layout operation that is not on XLA SPMD mesh during "
"XLA SPMD integration.");
}
StatusOr<xla::OpSharding> xla_sharding =
ConvertLayoutToXlaOpSharding(layout_op.getLayout());
if (!xla_sharding.ok()) {
return layout_op.emitError(xla_sharding.status().message());
}
func_op.setArgAttr(
arg_index, kXlaShardingAttr,
builder.getStringAttr(xla_sharding->SerializeAsString()));
break;
}
}
return mlir::WalkResult::advance();
});
return mlir::failure(result.wasInterrupted());
}
mlir::LogicalResult SetHloShardingForOutputs(mlir::ModuleOp module,
mlir::OpBuilder builder) {
// Set output attributes
auto result =
module.walk([&](mlir::func::ReturnOp return_op) -> mlir::WalkResult {
for (int return_index = 0; return_index < return_op.getNumOperands();
++return_index) {
auto layout_op = llvm::dyn_cast_or_null<mlir::TF::DTensorLayout>(
return_op->getOperand(return_index).getDefiningOp());
if (!layout_op) continue;
StatusOr<xla::OpSharding> xla_sharding =
ConvertLayoutToXlaOpSharding(layout_op.getLayout());
if (!xla_sharding.ok()) {
return module.emitError(xla_sharding.status().message());
}
mlir::func::FuncOp func_op =
layout_op->getParentOfType<mlir::func::FuncOp>();
if (!func_op) {
return module.emitError(
"Error finding surrounding FuncOp during "
"DTensorXlaSpmdIntegration.");
}
func_op.setResultAttr(
return_index, kXlaShardingAttr,
builder.getStringAttr(xla_sharding->SerializeAsString()));
}
return mlir::WalkResult::advance();
});
return mlir::failure(result.wasInterrupted());
}
class DTensorSetHloShardingPass
: public impl::DTensorSetHloShardingPassBase<DTensorSetHloShardingPass> {
public:
using DTensorSetHloShardingPassBase::DTensorSetHloShardingPassBase;
explicit DTensorSetHloShardingPass(
std::optional<bool> check_layout_use_xla_spmd) {
if (check_layout_use_xla_spmd.has_value()) {
check_layout_use_xla_spmd_ = *check_layout_use_xla_spmd;
}
}
void runOnOperation() override {
mlir::MLIRContext& context = getContext();
mlir::OpBuilder builder(&context);
mlir::ModuleOp module = getOperation();
if (mlir::failed(SetHloShardingForInputs(
module, builder, check_layout_use_xla_spmd_.getValue()))) {
return signalPassFailure();
}
if (mlir::failed(SetHloShardingForOutputs(module, builder))) {
return signalPassFailure();
}
}
};
} // namespace
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateDTensorSetHloShardingPass(std::optional<bool> check_layout_use_xla_spmd) {
return std::make_unique<DTensorSetHloShardingPass>(check_layout_use_xla_spmd);
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/expansions/argmax_spmd_expander.cc
+175
View File
@@ -0,0 +1,175 @@
/* Copyright 2022 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/dtensor/mlir/expansions/argmax_spmd_expander.h"
#include <cstdint>
#include <string>
#include <vector>
#include "llvm/ADT/DenseMap.h"
#include "llvm/Support/Casting.h"
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/collectives.h"
#include "tensorflow/dtensor/mlir/layout_parsing.h"
#include "tensorflow/dtensor/mlir/op_utils.h"
#include "tensorflow/dtensor/mlir/shape_utils.h"
#include "tensorflow/dtensor/mlir/value_utils.h"
#include "tensorflow/dtensor/proto/layout.pb.h"
namespace tensorflow {
namespace dtensor {
namespace {
StatusOr<Layout> ComputeResultLayout(mlir::Operation* op,
const Layout& input_layout) {
if (!mlir::isa<mlir::TF::ArgMaxOp>(op))
return absl::UnimplementedError(absl::StrCat(
"SPMD expansion for op type: ", OpName(op), " not yet implemented."));
auto argmax_op = llvm::cast<mlir::TF::ArgMaxOp>(op);
const auto input_rank = ValueRank(argmax_op.getInput());
TF_ASSIGN_OR_RETURN(int64_t axis,
ExtractConstIntFromValue(argmax_op.getDimension()));
if (axis < 0) axis += input_rank;
std::vector<std::string> sharding_specs;
for (int i = 0; i < input_rank; ++i) {
if (i != axis) sharding_specs.push_back(input_layout.sharding_spec(i));
}
return Layout::GetLayout(sharding_specs, input_layout.mesh());
}
} // namespace
StatusOr<mlir::Operation*> ArgMaxSPMDExpander::ExpandOp(mlir::Operation* op) {
auto argmax_op = llvm::cast<mlir::TF::ArgMaxOp>(op);
TF_ASSIGN_OR_RETURN(int64_t axis,
ExtractConstIntFromValue(argmax_op.getDimension()));
TF_ASSIGN_OR_RETURN(auto input_layout,
ExtractLayoutFromOperand(argmax_op.getInput()));
TF_ASSIGN_OR_RETURN(auto output_layout, ExtractSingleLayoutFromOp(argmax_op));
if (!input_layout || !output_layout)
return absl::InvalidArgumentError(
absl::StrCat(OpName(op), " is missing layouts during SPMD Expansion."));
mlir::Value input = argmax_op.getInput();
const auto input_rank = ValueRank(input);
TF_ASSIGN_OR_RETURN(auto input_shape, GetShapeOfValue(input));
if (input_rank == -1)
return absl::UnimplementedError("missing rank for input.");
if (axis < 0) axis += input_rank;
mlir::OpBuilder builder(op);
{
std::vector<std::string> tgt_input_sharding_specs;
tgt_input_sharding_specs.reserve(input_shape.size());
Mesh mesh = input_layout->mesh();
for (int i = 0; i < input_shape.size(); ++i) {
// const auto dim_name
if (i == axis) {
// Set replicated for `axis` dim.
tgt_input_sharding_specs.push_back(Layout::kUnshardedDim);
} else {
// Keep the rest dimension.
tgt_input_sharding_specs.push_back(input_layout->sharding_spec(i));
}
}
if (!Layout::IsUnshardedDimension(input_layout->sharding_spec(axis))) {
TF_ASSIGN_OR_RETURN(Layout layout,
Layout::GetLayout(input_layout->type(),
tgt_input_sharding_specs, mesh));
TF_ASSIGN_OR_RETURN(input,
EmitAllGather(builder, input, *input_layout, layout));
}
}
auto new_argmax = mlir::TF::ArgMaxOp::create(builder, argmax_op.getLoc(),
argmax_op.getResult().getType(),
input, argmax_op.getDimension());
op->getResult(0).replaceAllUsesWith(new_argmax.getOutput());
op->erase();
return InferSPMDExpandedLocalShape(new_argmax);
}
StatusOr<llvm::DenseMap<int, Layout>> ArgMaxSPMDExpander::ComputeLayoutForward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
// If the input layout is missing, don't return an output layout.
if (input_layouts.find(0) == input_layouts.end())
return llvm::DenseMap<int, Layout>();
const Layout& input_layout = input_layouts.lookup(0);
TF_ASSIGN_OR_RETURN(auto result_layout,
ComputeResultLayout(op, input_layout));
if (result_layout.rank() != input_layout.rank() - 1)
return absl::FailedPreconditionError(absl::StrCat(
OpName(op), " derived output layout rank is ", result_layout.rank(),
" not ", input_layout.rank() - 1, " as expected."));
return llvm::DenseMap<int, Layout>({{0, result_layout}});
}
StatusOr<llvm::DenseMap<int, Layout>> ArgMaxSPMDExpander::ComputeLayoutBackward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
// If no output layout, then do not infer any operand layouts.
if (output_layouts.find(0) == output_layouts.end())
return llvm::DenseMap<int, Layout>();
auto argmax_op = llvm::cast<mlir::TF::ArgMaxOp>(op);
TF_ASSIGN_OR_RETURN(int64_t axis,
ExtractConstIntFromValue(argmax_op.getDimension()));
auto input = argmax_op.getInput();
const auto input_rank = ValueRank(input);
// Handle the case of negative axis.
if (axis < 0) axis += input_rank;
const Layout& output_layout = output_layouts.lookup(0);
TF_ASSIGN_OR_RETURN(auto input_shape, GetShapeOfValue(input));
std::vector<std::string> layout_sharding;
int output_dim = 0;
for (int i = 0; i < input_shape.size(); ++i) {
if (i == axis) {
layout_sharding.emplace_back(Layout::kUnshardedDim);
} else {
layout_sharding.emplace_back(output_layout.sharding_spec(output_dim));
output_dim += 1;
}
}
// Add Layout for first input attribute, while the second one is axis as a
// scalar, we don't need to set its layout.
TF_ASSIGN_OR_RETURN(const Layout result_layout,
Layout::GetLayout(layout_sharding, output_layout.mesh()));
return llvm::DenseMap<int, Layout>({{0, result_layout}});
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/expansions/argmax_spmd_expander.h
+45
View File
@@ -0,0 +1,45 @@
/* Copyright 2022 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_DTENSOR_MLIR_EXPANSIONS_ARGMAX_SPMD_EXPANDER_H_
#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_ARGMAX_SPMD_EXPANDER_H_
#include "llvm/ADT/DenseMap.h"
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/spmd_expander.h"
namespace tensorflow {
namespace dtensor {
class ArgMaxSPMDExpander : public SPMDExpanderBase {
public:
StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& input_layouts) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& output_layouts) override;
};
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_ARGMAX_SPMD_EXPANDER_H_
tensorflow/dtensor/mlir/expansions/bias_add_spmd_expander.cc
+173
View File
@@ -0,0 +1,173 @@
/* Copyright 2022 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/dtensor/mlir/expansions/bias_add_spmd_expander.h"
#include <algorithm>
#include <cassert>
#include <cstdint>
#include <string>
#include <vector>
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.h"
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/collectives.h"
#include "tensorflow/dtensor/mlir/layout_parsing.h"
#include "tensorflow/dtensor/mlir/shape_utils.h"
namespace tensorflow {
namespace dtensor {
namespace {
int get_c_dimension_idx(const Layout& layout, llvm::StringRef data_format) {
// If format is "N...C", the bias is added to the last dimension.
int c_dim_idx = layout.sharding_spec_strs().size() - 1;
if (data_format.starts_with("NC")) {
// If format is "NC...", the bias is added to the 'C' dimension.
c_dim_idx = layout.sharding_spec_strs().size() - 3;
}
return c_dim_idx;
}
} // namespace
StatusOr<mlir::Operation*> BiasAddExpander::ExpandOp(mlir::Operation* op) {
TF_ASSIGN_OR_RETURN(auto output_layout,
ExtractRequiredSingleLayoutFromOp(op));
mlir::TF::BiasAddOp bias_add_op = llvm::cast<mlir::TF::BiasAddOp>(op);
const llvm::StringRef data_format = bias_add_op.getDataFormat();
const int c_dim_idx = get_c_dimension_idx(output_layout, data_format);
// Bias add op has 2 inputs: value and bias.
assert(op->getOpOperands().size() == 2);
mlir::OpOperand& input = op->getOpOperand(0);
TF_ASSIGN_OR_RETURN(Layout input_layout,
ExtractRequiredLayoutFromOperand(input.get()));
mlir::OpOperand& bias = op->getOpOperand(1);
TF_ASSIGN_OR_RETURN(const Layout bias_layout,
ExtractRequiredLayoutFromOperand(bias.get()));
// Check if output is sharded more, change input layout to match output
// layout.
int64_t num_input_shards = input_layout.num_shards_for_dim(c_dim_idx);
int64_t num_output_shards = output_layout.num_shards_for_dim(c_dim_idx);
if (num_input_shards < num_output_shards) {
mlir::Value output;
std::vector<std::string> input_new_specs =
output_layout.sharding_spec_strs();
TF_ASSIGN_OR_RETURN(
const Layout new_input_layout,
Layout::GetLayout(input_new_specs, input_layout.mesh()));
TF_ASSIGN_OR_RETURN(
output, EmitRelayout(input.get(), input_layout, new_input_layout));
input.set(output);
input_layout = new_input_layout;
}
// Map bias layout sharding to match sharding for 'c' dimension of input, if
// not same already.
if (bias_layout.sharding_spec(0) != input_layout.sharding_spec(c_dim_idx)) {
mlir::Value output;
std::vector<std::string> bias_new_specs = {
input_layout.sharding_spec_strs()[c_dim_idx]};
TF_ASSIGN_OR_RETURN(const Layout new_bias_layout,
Layout::GetLayout(bias_new_specs, bias_layout.mesh()));
TF_ASSIGN_OR_RETURN(output,
EmitRelayout(bias.get(), bias_layout, new_bias_layout));
bias.set(output);
}
// Perform SPMD operation locally
mlir::Operation* new_local_op = InferSPMDExpandedLocalShape(op);
// Convert result layout to output layout.
llvm::SmallPtrSet<mlir::Operation*, 4> newly_created_ops;
TF_ASSIGN_OR_RETURN(mlir::Value relayout_output,
EmitRelayout(new_local_op->getOpResult(0), input_layout,
output_layout, &newly_created_ops));
op->getResult(0).replaceAllUsesExcept(relayout_output, newly_created_ops);
return relayout_output.getDefiningOp();
}
StatusOr<llvm::DenseMap<int, Layout>> BiasAddExpander::ComputeLayoutForward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
// If we do not have an input layout then do not infer an output layout.
if (input_layouts.find(0) == input_layouts.end())
return llvm::DenseMap<int, Layout>();
Layout input_layout = input_layouts.lookup(0);
mlir::TF::BiasAddOp bias_add_op = llvm::cast<mlir::TF::BiasAddOp>(op);
llvm::StringRef data_format = bias_add_op.getDataFormat();
int c_dim_idx = get_c_dimension_idx(input_layout, data_format);
std::vector<std::string> new_output_layout_specs =
input_layout.sharding_spec_strs();
if (Layout::IsUnshardedDimension(new_output_layout_specs[c_dim_idx]) &&
input_layouts.find(1) != input_layouts.end()) {
// Shard c_dim using bias sharding as long as the sharding spec is not
// already used in input for some other dimension.
Layout bias_layout = input_layouts.lookup(1);
std::string bias_sharding = bias_layout.sharding_spec(0);
if (std::find(new_output_layout_specs.begin(),
new_output_layout_specs.end(),
bias_sharding) == new_output_layout_specs.end()) {
new_output_layout_specs[c_dim_idx] = bias_layout.sharding_spec(0);
}
}
TF_ASSIGN_OR_RETURN(
Layout new_output_layout,
Layout::GetLayout(new_output_layout_specs, input_layout.mesh()));
return llvm::DenseMap<int, Layout>({{0, new_output_layout}});
}
StatusOr<llvm::DenseMap<int, Layout>> BiasAddExpander::ComputeLayoutBackward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
if (output_layouts.find(0) == output_layouts.end())
return llvm::DenseMap<int, Layout>();
llvm::DenseMap<int, Layout> input_layouts;
// If output layout is given, match input_layout and bias layout to match
// it.
Layout output_layout = output_layouts.lookup(0);
// Bias layout should match 'C' dimension of input layout.
mlir::TF::BiasAddOp bias_add_op = llvm::cast<mlir::TF::BiasAddOp>(op);
llvm::StringRef data_format = bias_add_op.getDataFormat();
const int c_dim_idx = get_c_dimension_idx(output_layout, data_format);
std::vector<std::string> bias_new_specs = {
output_layout.sharding_spec_strs()[c_dim_idx]};
TF_ASSIGN_OR_RETURN(Layout new_bias_layout,
Layout::GetLayout(bias_new_specs, output_layout.mesh()));
return llvm::DenseMap<int, Layout>(
{{0, output_layout}, {1, new_bias_layout}});
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/expansions/bias_add_spmd_expander.h
+48
View File
@@ -0,0 +1,48 @@
/* Copyright 2022 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_DTENSOR_MLIR_EXPANSIONS_BIAS_ADD_SPMD_EXPANDER_H_
#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_BIAS_ADD_SPMD_EXPANDER_H_
#include "llvm/ADT/DenseMap.h"
#include "mlir/IR/Operation.h" // from @llvm-project
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/expansions/elementwise_spmd_expander.h"
#include "tensorflow/dtensor/mlir/spmd_expander.h"
namespace tensorflow {
namespace dtensor {
// Expander for `BiasAdd` Op.
// `BiasAdd` uses non-standard broadcasting rules if 'NC...' data_format is
// used, thus we can't just reuse ElementwiseSPMDExpander as it is.
class BiasAddExpander : public SPMDExpanderBase {
protected:
StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& input_layouts) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& output_layouts) override;
};
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_BIAS_ADD_SPMD_EXPANDER_H_
tensorflow/dtensor/mlir/expansions/broadcast_to_spmd_expander.cc
+220
View File
@@ -0,0 +1,220 @@
/* Copyright 2022 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/dtensor/mlir/expansions/broadcast_to_spmd_expander.h"
#include <cstdint>
#include <string>
#include <vector>
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/collectives.h"
#include "tensorflow/dtensor/mlir/layout_parsing.h"
#include "tensorflow/dtensor/mlir/shape_utils.h"
#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
#include "tensorflow/dtensor/mlir/value_utils.h"
#include "tensorflow/dtensor/proto/layout.pb.h"
namespace tensorflow {
namespace dtensor {
StatusOr<mlir::Operation*> BroadcastToSPMDExpander::ExpandOp(
mlir::Operation* op) {
auto broadcast_op = llvm::cast<mlir::TF::BroadcastToOp>(op);
TF_ASSIGN_OR_RETURN(
const Layout shape_layout,
ExtractRequiredLayoutFromOperand(broadcast_op.getShape()));
if (!shape_layout.IsFullyReplicated()) {
return absl::InvalidArgumentError(
"Error during BroadcastOp SPMD Expansion. Shape input of broadcast op "
"must be fully replicated.");
}
TF_ASSIGN_OR_RETURN(
const Layout input_layout,
ExtractRequiredLayoutFromOperand(broadcast_op.getInput()));
TF_ASSIGN_OR_RETURN(const Layout output_layout,
ExtractRequiredSingleLayoutFromOp(broadcast_op));
TF_ASSIGN_OR_RETURN(
llvm::ArrayRef<int64_t> input_global_size,
GetGlobalShapeOfValueFromDTensorLayout(broadcast_op.getInput()));
llvm::SmallVector<int64_t, 4> broadcast_to_shape;
TF_RETURN_IF_ERROR(ExtractConstVectorFromValue(
GetForwardedDTensorLayoutInput(broadcast_op.getShape()),
&broadcast_to_shape));
// Input to BroadcastTo op requires all to all if non-broadcasted-dimensions
// are not same.
const int broadcasted_dimensions = output_layout.rank() - input_layout.rank();
bool requires_all_to_all = false;
const auto output_num_shards = output_layout.num_shards();
for (int i = 0; i < input_layout.rank(); ++i) {
const int output_dim_index = i + broadcasted_dimensions;
const std::string& output_layout_dim =
output_layout.sharding_spec(output_dim_index);
if (input_global_size[i] > 1 &&
input_layout.sharding_spec(i) != output_layout_dim) {
requires_all_to_all = true;
}
if (output_layout_dim != Layout::kUnshardedDim) {
broadcast_to_shape[output_dim_index] /=
output_num_shards[output_dim_index];
}
}
// Insert all-to-all operations just before Broadcast op to ensure all inputs
// in correct local values.
mlir::OpBuilder builder(op);
mlir::Value input_data = broadcast_op.getInput();
TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
const Layout all_to_all_input_layout =
Layout::ReplicatedOnMesh(mesh, input_layout.rank());
if (requires_all_to_all) {
TF_ASSIGN_OR_RETURN(auto input_data,
EmitAllGather(builder, input_data, input_layout,
all_to_all_input_layout));
op->setOperand(0, input_data);
} else {
// When all-to-all is not needed, output of BroadcastTo operation may be
// sharded. In that case, we must ensure that `shape` input of BroadcastTo
// op has correct local sharded shape.
// Note that we include the sharding on the first
for (int i = 0; i < broadcasted_dimensions; ++i)
if (output_layout.sharding_spec(i) != Layout::kUnshardedDim)
broadcast_to_shape[i] /= output_num_shards[i];
mlir::Value new_broadcast_to_shape =
Int64Const(builder, op->getLoc(), broadcast_to_shape);
op->setOperand(1, new_broadcast_to_shape);
}
op = InferSPMDExpandedLocalShape(op);
if (!requires_all_to_all) return op;
// If we all-to-all'ed, we may need to split after the local BroadcastTo op
// has been created in graph.
llvm::SmallPtrSet<mlir::Operation*, 4> newly_created_ops;
builder.setInsertionPointAfter(op);
TF_ASSIGN_OR_RETURN(
auto final_output,
EmitAllScatter(builder, op->getOpResult(0),
all_to_all_input_layout.LeftPad(output_layout.rank()),
output_layout, &newly_created_ops));
op->getOpResult(0).replaceAllUsesExcept(final_output, newly_created_ops);
return final_output.getDefiningOp();
}
StatusOr<llvm::DenseMap<int, Layout>>
BroadcastToSPMDExpander::ComputeLayoutForward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
// If we do not have an input layout then do not infer an output layout.
if (input_layouts.find(0) == input_layouts.end())
return llvm::DenseMap<int, Layout>();
TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op));
auto broadcast_op = llvm::cast<mlir::TF::BroadcastToOp>(op);
TF_ASSIGN_OR_RETURN(
const auto broadcasted_output_shape,
GetShapeOfValue(broadcast_op.getOutput(), /*fail_on_dynamic=*/true));
TF_ASSIGN_OR_RETURN(
const auto input_shape,
GetShapeOfValue(broadcast_op.getInput(), /*fail_on_dynamic=*/true));
// Broadcasting works from trailing dimensions and dimensions are broadcasted
// in forward direction.
const int output_shape_rank = broadcasted_output_shape.size();
const int input_shape_rank = input_shape.size();
const int broadcasted_dimensions = output_shape_rank - input_shape_rank;
if (broadcasted_dimensions < 0)
return absl::FailedPreconditionError(
"Broadcasted dimension was less than 0.");
Layout input_layout = input_layouts.lookup(0);
std::vector<std::string> layout_sharding;
for (int i = 0; i < output_shape_rank; ++i) {
if (i < broadcasted_dimensions) {
layout_sharding.push_back(Layout::kUnshardedDim);
} else {
layout_sharding.push_back(
input_layout.sharding_spec(i - broadcasted_dimensions));
}
}
TF_ASSIGN_OR_RETURN(Layout inferred_output_layout,
Layout::GetLayout(layout_sharding, mesh));
return llvm::DenseMap<int, Layout>({{0, inferred_output_layout}});
}
StatusOr<llvm::DenseMap<int, Layout>>
BroadcastToSPMDExpander::ComputeLayoutBackward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op));
// If output layout is not set, then we can only infer the `shape` input
// which should always be replicated.
if (output_layouts.find(0) == output_layouts.end())
return llvm::DenseMap<int, Layout>(
{{1, Layout::ReplicatedOnMesh(mesh, ValueRank(op->getOperand(1)))}});
auto output_layout = output_layouts.lookup(0);
auto broadcast_op = llvm::cast<mlir::TF::BroadcastToOp>(op);
TF_ASSIGN_OR_RETURN(
const auto broadcasted_output_shape,
GetShapeOfValue(broadcast_op.getOutput(), /*fail_on_dynamic=*/true));
TF_ASSIGN_OR_RETURN(
const auto input_shape,
GetShapeOfValue(broadcast_op.getInput(), /*fail_on_dynamic=*/true));
// Broadcasting works from trailing dimensions and dimensions are broadcasted
// in forward direction.
const int output_shape_rank = broadcasted_output_shape.size();
const int input_shape_rank = input_shape.size();
const int broadcasted_dimensions = output_shape_rank - input_shape_rank;
std::vector<std::string> sharding_specs;
for (int i = 0; i < input_shape_rank; ++i) {
if (input_shape[i] == 1) {
sharding_specs.push_back(Layout::kUnshardedDim);
} else {
sharding_specs.push_back(
output_layout.sharding_spec(i + broadcasted_dimensions));
}
}
TF_ASSIGN_OR_RETURN(Layout inferred_operand_layout,
Layout::GetLayout(sharding_specs, mesh));
// `shape` input of BroadcastTo is always set as replicated.
return llvm::DenseMap<int, Layout>(
{{0, inferred_operand_layout},
{1, Layout::ReplicatedOnMesh(mesh, ValueRank(op->getOperand(1)))}});
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/expansions/broadcast_to_spmd_expander.h
+44
View File
@@ -0,0 +1,44 @@
/* Copyright 2022 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_DTENSOR_MLIR_EXPANSIONS_BROADCAST_TO_SPMD_EXPANDER_H_
#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_BROADCAST_TO_SPMD_EXPANDER_H_
#include "llvm/ADT/DenseMap.h"
#include "mlir/IR/Operation.h" // from @llvm-project
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/spmd_expander.h"
namespace tensorflow {
namespace dtensor {
class BroadcastToSPMDExpander : public SPMDExpanderBase {
public:
StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& input_layouts) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& output_layouts) override;
};
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_BROADCAST_TO_SPMD_EXPANDER_H_
tensorflow/dtensor/mlir/expansions/concat_spmd_expander.cc
+169
View File
@@ -0,0 +1,169 @@
/* Copyright 2022 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/dtensor/mlir/expansions/concat_spmd_expander.h"
#include <cstdint>
#include <optional>
#include "absl/status/status.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/Casting.h"
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/collectives.h"
#include "tensorflow/dtensor/mlir/layout_parsing.h"
#include "tensorflow/dtensor/mlir/shape_utils.h"
#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
#include "tensorflow/dtensor/mlir/value_utils.h"
namespace tensorflow {
namespace dtensor {
namespace {
absl::Status VerifyConcatLayout(mlir::Value concat_dim_operand,
const Layout& concat_layout) {
TF_ASSIGN_OR_RETURN(int64_t concat_dim_value,
ExtractConstIntFromValue(concat_dim_operand));
for (const auto& shard_and_dimension :
llvm::enumerate(concat_layout.num_shards())) {
if (shard_and_dimension.index() == concat_dim_value &&
shard_and_dimension.value() > 1) {
return absl::InvalidArgumentError(
"Concat op SPMD with concat dimension in sharded dimension is "
"not supported.");
}
}
return absl::OkStatus();
}
StatusOr<Layout> ReduceForConcatOutputLayout(mlir::Value concat_dim_operand,
const Layout& layout) {
TF_ASSIGN_OR_RETURN(int64_t concat_dim_value,
ExtractConstIntFromValue(concat_dim_operand));
// Set concatenated dimension to replicated.
return layout.GetLayoutWithReducedDims({concat_dim_value},
/*keep_dims=*/true);
}
} // namespace
StatusOr<mlir::Operation*> ConcatSPMDExpander::ExpandOp(mlir::Operation* op) {
if (!llvm::isa<mlir::TF::ConcatOp, mlir::TF::ConcatV2Op>(op))
return absl::InvalidArgumentError(
"Requested SPMD Expansion for op that is not Concat or ConcatV2.");
// Extract the concat dim. ConcatOp and ConcatV2Op define the dim at
// different position.
bool is_concat_v1 = llvm::isa<mlir::TF::ConcatOp>(op);
const int concat_dim_operand_idx =
is_concat_v1 ? 0 : op->getNumOperands() - 1;
mlir::Value concat_dim = op->getOperand(concat_dim_operand_idx);
// Ensure that Concat op is not sharded on concat-dimension.
TF_ASSIGN_OR_RETURN(auto concat_output_layout,
ExtractRequiredSingleLayoutFromOp(op));
TF_RETURN_IF_ERROR(VerifyConcatLayout(concat_dim, concat_output_layout));
// Relayout all inputs to match output layout before concating.
// This will ensure that a local concat is the correct thing to do.
for (int i = 0; i < op->getNumOperands(); ++i) {
// Skip relayout for the concat dim operand.
if (i == concat_dim_operand_idx) continue;
TF_ASSIGN_OR_RETURN(Layout layout,
ExtractRequiredLayoutFromOperand(op->getOperand(i)));
TF_ASSIGN_OR_RETURN(
mlir::Value new_input,
EmitRelayout(op->getOperand(i), layout, concat_output_layout));
op->setOperand(i, new_input);
}
return InferSPMDExpandedLocalShape(op);
}
StatusOr<llvm::DenseMap<int, Layout>> ConcatSPMDExpander::ComputeLayoutForward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
// Verify operand layouts to ensure that there are no conflicting concat
// operand layouts.
const bool is_concat_v1 = llvm::isa<mlir::TF::ConcatOp>(op);
auto begin_idx = is_concat_v1 ? 1 : 0;
auto end_idx =
is_concat_v1 ? op->getNumOperands() - 1 : op->getNumOperands() - 2;
llvm::DenseMap<int, Layout> concat_operands_layouts;
for (auto idx = begin_idx; idx <= end_idx; ++idx) {
if (input_layouts.find(idx) != input_layouts.end())
concat_operands_layouts[idx] = input_layouts.lookup(idx);
}
TF_ASSIGN_OR_RETURN(std::optional<Layout> concat_operand_layout,
GetMergedOperandLayout(concat_operands_layouts, op));
// Concat/ConcatV2 has different operand index for concat dim. Retrieve the
// correct concat dim value.
const int concat_dim_operand_idx =
is_concat_v1 ? 0 : op->getNumOperands() - 1;
mlir::Value concat_dim = op->getOperand(concat_dim_operand_idx);
// If consistent operand layout exists, propagate the operand layout to output
// layout.
if (concat_operand_layout) {
TF_ASSIGN_OR_RETURN(
const Layout reduced_concat_layout,
ReduceForConcatOutputLayout(concat_dim, *concat_operand_layout));
return llvm::DenseMap<int, Layout>({{0, reduced_concat_layout}});
} else {
return llvm::DenseMap<int, Layout>();
}
}
StatusOr<llvm::DenseMap<int, Layout>> ConcatSPMDExpander::ComputeLayoutBackward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
// If output layout does not exist then do not infer any operand layouts.
if (output_layouts.find(0) == output_layouts.end())
return llvm::DenseMap<int, Layout>();
Layout output_layout = output_layouts.lookup(0);
const bool is_concat_v1 = llvm::isa<mlir::TF::ConcatOp>(op);
// Concat/ConcatV2 has different operand index for concat dim. Retrieve the
// correct concat dim value.
const int concat_dim_operand_idx =
is_concat_v1 ? 0 : op->getNumOperands() - 1;
mlir::Value concat_dim = op->getOperand(concat_dim_operand_idx);
// If suggested output layout exists, verify that concatenated dimension is
// replicated to ensure no cross device communication is needed, then
// propagate the output layout to all concat tensor operands' layout.
// Set concatenated dimension to replicated.
TF_ASSIGN_OR_RETURN(const Layout inferred_input_layout,
ReduceForConcatOutputLayout(concat_dim, output_layout));
llvm::DenseMap<int, Layout> operand_layouts;
for (auto index = 0; index < op->getNumOperands(); ++index) {
if (index == concat_dim_operand_idx) continue;
operand_layouts[index] = inferred_input_layout;
}
return operand_layouts;
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/expansions/concat_spmd_expander.h
+45
View File
@@ -0,0 +1,45 @@
/* Copyright 2022 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_DTENSOR_MLIR_EXPANSIONS_CONCAT_SPMD_EXPANDER_H_
#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_CONCAT_SPMD_EXPANDER_H_
#include "llvm/ADT/DenseMap.h"
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/spmd_expander.h"
namespace tensorflow {
namespace dtensor {
class ConcatSPMDExpander : public SPMDExpanderBase {
public:
StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& input_layouts) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& output_layouts) override;
};
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_CONCAT_SPMD_EXPANDER_H_
tensorflow/dtensor/mlir/expansions/control_flow_spmd_expander.cc
+123
View File
@@ -0,0 +1,123 @@
/* Copyright 2022 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/dtensor/mlir/expansions/control_flow_spmd_expander.h"
#include <cassert>
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/Casting.h"
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/ir/tf_dtensor.h"
#include "tensorflow/dtensor/mlir/layout_parsing.h"
#include "tensorflow/dtensor/mlir/value_utils.h"
namespace tensorflow {
namespace dtensor {
StatusOr<mlir::Operation*> WhileRegionSPMDExpander::ExpandOp(
mlir::Operation* op) {
assert(op->getNumOperands() == op->getNumResults());
// Set the type for the results of the WhileRegion explicitly.
//
// Normally we would use InferSPMDExpandedLocalShape for this, but that
// function requires the op to either have a type inference interface
// (which WhileRegion does not) or a TensorFlow ShapeFn (WhileRegion is not
// a TensorFlow op). During the standard MLIR shape inference pass this op
// is handled by a special case in InferShapeForSingleOperation.
for (int i = 0; i < op->getNumOperands(); ++i)
op->getResult(i).setType(op->getOperand(i).getType());
auto while_op = llvm::cast<mlir::TF::WhileRegionOp>(op);
for (const auto& data :
llvm::enumerate(llvm::zip(while_op.getCond().front().getArguments(),
while_op.getBody().front().getArguments()))) {
const int index = data.index();
mlir::BlockArgument cond_arg = std::get<0>(data.value());
mlir::BlockArgument body_arg = std::get<1>(data.value());
cond_arg.setType(while_op.getOperand(index).getType());
body_arg.setType(while_op.getOperand(index).getType());
}
return op;
}
StatusOr<llvm::DenseMap<int, Layout>>
WhileRegionSPMDExpander::ComputeLayoutForward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
return absl::UnimplementedError(
"WhileRegion does not support compute layouts. This should not be "
"called.");
}
StatusOr<llvm::DenseMap<int, Layout>>
WhileRegionSPMDExpander::ComputeLayoutBackward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
return absl::UnimplementedError(
"WhileRegion does not support compute layouts. This should not be "
"called.");
}
StatusOr<mlir::Operation*> IfRegionSPMDExpander::ExpandOp(mlir::Operation* op) {
auto if_op = llvm::cast<mlir::TF::IfRegionOp>(op);
for (mlir::Value result : if_op->getResults()) {
auto result_layout_op = llvm::dyn_cast_or_null<mlir::TF::DTensorLayout>(
*result.getUsers().begin());
if (!result_layout_op)
return absl::InvalidArgumentError(
"Missing layout of If op result during SPMD expansion.");
const Layout layout = result_layout_op.getLayout();
if (!layout.IsFullyReplicated()) {
const auto global_shape = result_layout_op.getGlobalShape();
if (!global_shape)
return absl::InvalidArgumentError(
"Shape of If op must be statically known for SPMD expansion.");
result.setType(mlir::RankedTensorType::get(
layout.LocalShapeFromGlobalShape(*global_shape),
mlir::cast<mlir::TensorType>(result.getType()).getElementType()));
}
}
return op;
}
StatusOr<llvm::DenseMap<int, Layout>>
IfRegionSPMDExpander::ComputeLayoutForward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
// No-op for forward propagation.
return llvm::DenseMap<int, Layout>();
}
StatusOr<llvm::DenseMap<int, Layout>>
IfRegionSPMDExpander::ComputeLayoutBackward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
// Layout propagation for TF::IfRegion op is no-op. Actual layout
// propagation logic depends on layout propgation of ops inside the
// then/else regions of the IfRegion op.
TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
return llvm::DenseMap<int, Layout>(
{{0, Layout::ReplicatedOnMesh(mesh, ValueRank(op->getOperand(0)))}});
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/expansions/control_flow_spmd_expander.h
+63
View File
@@ -0,0 +1,63 @@
/* Copyright 2022 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_DTENSOR_MLIR_EXPANSIONS_CONTROL_FLOW_SPMD_EXPANDER_H_
#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_CONTROL_FLOW_SPMD_EXPANDER_H_
#include "llvm/ADT/DenseMap.h"
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/layout_parsing.h"
#include "tensorflow/dtensor/mlir/spmd_expander.h"
#include "tensorflow/dtensor/mlir/value_utils.h"
namespace tensorflow {
namespace dtensor {
class WhileRegionSPMDExpander : public SPMDExpanderBase {
public:
StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& input_layouts) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& output_layouts) override;
};
class IfRegionSPMDExpander : public SPMDExpanderBase {
public:
StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& input_layouts) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& output_layouts) override;
};
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_CONTROL_FLOW_SPMD_EXPANDER_H_
tensorflow/dtensor/mlir/expansions/conv_spmd_expander.cc
+815
View File
@@ -0,0 +1,815 @@
/* Copyright 2022 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/dtensor/mlir/expansions/conv_spmd_expander.h"
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <string>
#include <vector>
#include "absl/status/status.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/FormatVariadic.h"
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/IR/ValueRange.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_attributes.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/collectives.h"
#include "tensorflow/dtensor/mlir/dtensor_dialect/ir/dtensor_attributes.h"
#include "tensorflow/dtensor/mlir/ir/tf_dtensor.h"
#include "tensorflow/dtensor/mlir/layout_parsing.h"
#include "tensorflow/dtensor/mlir/op_utils.h"
#include "tensorflow/dtensor/mlir/shape_utils.h"
#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
#include "tensorflow/dtensor/mlir/value_utils.h"
namespace tensorflow {
namespace dtensor {
namespace {
template <typename ConvOp>
absl::Status VerifyConvLayout(const Layout& input_layout,
const Layout& filter_layout, ConvOp conv_op) {
if (!filter_layout.IsFullyReplicated())
return absl::InvalidArgumentError(
"Filter for convolution must have fully replicated layout.");
// Data format "NCHW" or "NCDHW".
int channel_dim = 1;
if (conv_op.getDataFormat() == "NHWC")
channel_dim = 3;
else if (conv_op.getDataFormat() == "NDHWC")
channel_dim = 4;
if (input_layout.sharding_spec(channel_dim) != Layout::kUnshardedDim)
return absl::InvalidArgumentError(
"Conv input's channel dimension must be replicated.");
if (input_layout.IsBatchParallel() || input_layout.IsFullyReplicated())
// No further checks needed for replicated case.
return absl::OkStatus();
if (conv_op.getPadding() == "EXPLICIT")
return absl::InvalidArgumentError(
"Explicit padding not supported for convolution with spatial "
"partitions.");
const int num_non_default_dilations =
llvm::count_if(conv_op.getDilations(), [](mlir::Attribute dilation) {
return mlir::cast<mlir::IntegerAttr>(dilation).getInt() != 1;
});
if (num_non_default_dilations > 0)
return absl::InvalidArgumentError(
"Only dilation rate 1 is supported for convolution with spatial "
"partitions.");
// TODO(b/208700444): support convolution with strides greater than 1.
const int num_non_default_strides =
llvm::count_if(conv_op.getStrides(), [](mlir::Attribute stride) {
return mlir::cast<mlir::IntegerAttr>(stride).getInt() != 1;
});
if (num_non_default_strides > 0)
return absl::InvalidArgumentError(
"Only stride 1 is supported for convolution with spatial partitions.");
mlir::Value input = conv_op.getInput();
auto input_type = mlir::dyn_cast<mlir::RankedTensorType>(input.getType());
if (!input_type || !input_type.hasStaticShape())
return absl::InvalidArgumentError(
"Input must have static shapes for convolution with spatial "
"partitions.");
mlir::Value filter = conv_op.getFilter();
auto filter_type = mlir::dyn_cast<mlir::RankedTensorType>(filter.getType());
if (!filter_type || !filter_type.hasStaticShape())
return absl::InvalidArgumentError(
"Filter must have static shapes for convolution with spatial "
"partitions.");
llvm::ArrayRef<int64_t> filter_shape = filter_type.getShape();
for (auto it = filter_shape.begin(); it != filter_shape.end() - 2; ++it) {
if (*it % 2 != 1)
return absl::InvalidArgumentError(
"Filter dimensions must be odd numbers for convolution with "
"spatial partitions.");
}
return absl::OkStatus();
}
mlir::Value PadInputOnUnshardedDim(mlir::OpBuilder& builder,
mlir::Location location,
mlir::Value input_tensor, int curr_input_dim,
int64_t curr_filter_dim_size) {
auto input_tensor_type =
mlir::dyn_cast<mlir::RankedTensorType>(input_tensor.getType());
auto input_tensor_shape = input_tensor_type.getShape();
const size_t paddings_flat_length = input_tensor_type.getRank() * 2;
llvm::SmallVector<int64_t, 4> paddings_flat_vec(paddings_flat_length, 0);
int64_t padding_size = curr_filter_dim_size - 1;
paddings_flat_vec[2 * curr_input_dim] = padding_size / 2;
paddings_flat_vec[2 * curr_input_dim + 1] = padding_size / 2;
llvm::SmallVector<int64_t, 4> paddings_shape(input_tensor_shape.begin(),
input_tensor_shape.end());
paddings_shape[curr_input_dim] += padding_size;
mlir::Value paddings_flat = Int64Const(builder, location, paddings_flat_vec);
mlir::RankedTensorType paddings_type = mlir::RankedTensorType::get(
paddings_shape, input_tensor_type.getElementType());
mlir::Value paddings = mlir::TF::ReshapeOp::create(
builder, location, paddings_flat,
Int64Const(builder, location, {input_tensor_type.getRank(), 2}));
return mlir::TF::PadOp::create(builder, location, paddings_type, input_tensor,
paddings);
}
template <typename ConvOp>
StatusOr<mlir::Operation*> HandleConv(ConvOp conv_op) {
mlir::OpBuilder builder(conv_op);
TF_ASSIGN_OR_RETURN(const Layout input_layout,
ExtractRequiredLayoutFromOperand(conv_op.getInput()));
TF_ASSIGN_OR_RETURN(const Layout filter_layout,
ExtractRequiredLayoutFromOperand(conv_op.getFilter()));
TF_ASSIGN_OR_RETURN(const Layout output_layout,
ExtractRequiredSingleLayoutFromOp(conv_op));
TF_RETURN_IF_ERROR(VerifyConvLayout(input_layout, filter_layout, conv_op));
if (input_layout.IsBatchParallel() || input_layout.IsFullyReplicated())
// No special handling needed for replicated case.
return InferSPMDExpandedLocalShape(conv_op);
mlir::tf_device::ClusterOp cluster =
conv_op->template getParentOfType<mlir::tf_device::ClusterOp>();
TF_ASSIGN_OR_RETURN(mlir::Value mesh_coordinates,
GetMeshCoordinatesFromCluster(cluster));
const Mesh& mesh = input_layout.mesh();
mlir::Location location = conv_op->getLoc();
const std::vector<std::string> input_sharding_spec =
input_layout.sharding_spec_strs();
const std::vector<std::string> output_sharding_spec =
output_layout.sharding_spec_strs();
llvm::StringRef format = conv_op.getDataFormat();
llvm::StringRef padding = conv_op.getPadding();
const auto input_num_shards = input_layout.num_shards();
const auto output_num_shards = output_layout.num_shards();
auto filter_type =
mlir::dyn_cast<mlir::RankedTensorType>(conv_op.getFilter().getType());
auto filter_shape = filter_type.getShape();
int begin_input_dim = -1, end_input_dim = -1;
if (format == "NCHW") {
begin_input_dim = 2;
end_input_dim = 3;
} else if (format == "NHWC") {
begin_input_dim = 1;
end_input_dim = 2;
} else if (format == "NCDHW") {
begin_input_dim = 2;
end_input_dim = 4;
} else if (format == "NDHWC") {
begin_input_dim = 1;
end_input_dim = 3;
}
// For non-batch, non-channel dimension sharding, conduct halo exchange.
for (int curr_input_dim = begin_input_dim; curr_input_dim <= end_input_dim;
++curr_input_dim) {
int curr_filter_dim = curr_input_dim - begin_input_dim;
auto input_type =
mlir::dyn_cast<mlir::RankedTensorType>(conv_op.getInput().getType());
auto input_shape = input_type.getShape();
if (input_sharding_spec[curr_input_dim] == Layout::kUnshardedDim) {
if (padding == "SAME") {
// Since we always emit a Conv op with "VALID" padding, we need to
// manually pad the input tensor.
conv_op->setOperand(
0, PadInputOnUnshardedDim(builder, location, conv_op.getInput(),
curr_input_dim,
filter_shape[curr_filter_dim]));
}
// No halo exchange is needed for unsharded dims.
continue;
}
TF_ASSIGN_OR_RETURN(const int mesh_dim_index,
mesh.idx_for_dim(input_sharding_spec[curr_input_dim]));
TF_ASSIGN_OR_RETURN(mlir::Value scalar_mesh_coordinate,
SelectScalarValueFromArray(builder, mesh_dim_index,
location, mesh_coordinates));
int halo_size;
if (padding == "SAME") {
halo_size = std::floor(filter_shape[curr_filter_dim] / 2);
} else if (padding == "VALID") {
int input_local_size = input_shape[curr_input_dim];
int input_size = input_local_size * input_num_shards[curr_input_dim];
int output_size = input_size - (filter_shape[curr_filter_dim] - 1);
int output_local_size = output_size / output_num_shards[curr_input_dim];
halo_size = output_local_size + (filter_shape[curr_filter_dim] - 1) -
input_local_size;
} else {
return absl::UnimplementedError(
absl::StrCat("Spatially partitioned convolution with padding \"",
padding.str(), "\" is not supported."));
}
if (halo_size == 0)
// No exchange is needed for empty halos.
continue;
builder.setInsertionPoint(conv_op);
TF_ASSIGN_OR_RETURN(
mlir::Value halo_exchanged_input,
EmitHaloExchange(builder, halo_size,
input_sharding_spec[curr_input_dim], input_layout,
mesh_coordinates, cluster, location,
conv_op.getInput()));
if (padding == "SAME") {
conv_op->setOperand(0, halo_exchanged_input);
} else if (padding == "VALID") {
// Slice the halo exchanged tensor to the desired size based on the index
// of the shard on the current dimension.
llvm::SmallVector<int32_t, 4> halo_sizes(input_layout.rank(), 0);
halo_sizes[curr_input_dim] = halo_size;
mlir::Value halo_sizes_const = IntConst(builder, location, halo_sizes);
llvm::SmallVector<int32_t, 4> halo_increments(input_layout.rank(), 0);
halo_increments[curr_input_dim] =
halo_size / (input_num_shards[curr_input_dim] - 1);
mlir::Value halo_increments_const =
IntConst(builder, location, halo_increments);
mlir::Value offset = mlir::TF::MulOp::create(
builder, location, halo_increments_const.getType(),
scalar_mesh_coordinate, halo_increments_const);
mlir::Value slice_begin =
mlir::TF::SubOp::create(builder, location, halo_sizes_const, offset);
llvm::SmallVector<int64_t, 4> slice_size(input_shape.begin(),
input_shape.end());
slice_size[curr_input_dim] += halo_size;
mlir::Value slice_size_const = Int64Const(builder, location, slice_size);
// slice_size_const and slize_begin_int64 has to be same type.
mlir::Value slice_begin_int64 = mlir::TF::CastOp::create(
builder, location,
mlir::RankedTensorType::get({input_layout.rank()},
builder.getI64Type()),
slice_begin);
mlir::RankedTensorType sliced_input_type =
mlir::RankedTensorType::get(slice_size, input_type.getElementType());
mlir::Value sliced_input = mlir::TF::SliceOp::create(
builder, location, sliced_input_type, /*input=*/halo_exchanged_input,
/*begin=*/slice_begin_int64, /*size=*/slice_size_const);
conv_op->setOperand(0, sliced_input);
}
// Spatially partitioned convolution always uses VALID padding after halo
// exchange.
conv_op.setPaddingAttr(builder.getStringAttr("VALID"));
}
return InferSPMDExpandedLocalShape(conv_op);
}
template <typename ConvBackpropInputOp>
StatusOr<mlir::Operation*> HandleConvBackpropInput(
const Layout& output_layout, ConvBackpropInputOp conv_op) {
TF_ASSIGN_OR_RETURN(std::vector<Layout> input_layouts,
ExtractRequiredLayoutFromOperands(conv_op));
const Layout& input_shape_layout = input_layouts[0];
const Layout& filter_layout = input_layouts[1];
const Layout& grad_layout = input_layouts[2];
// We only support batch sharding for these. In this case the output and input
// gradient must both be batch sharded. The filter input must be replicated.
if (!(output_layout.IsBatchParallel() || output_layout.IsFullyReplicated()) ||
!(grad_layout.IsBatchParallel() || grad_layout.IsFullyReplicated())) {
return errors::InvalidArgument("{0} only supports batch parallel layouts.",
conv_op->getName().getStringRef().str());
}
if (!filter_layout.IsFullyReplicated()) {
return errors::InvalidArgument("{0} only supports replicated filters.",
conv_op->getName().getStringRef().str());
}
if (!input_shape_layout.IsFullyReplicated()) {
return errors::InvalidArgument(
"Layout of the input shape (parameter 0) of {0} must be replicated.",
conv_op->getName().getStringRef().str());
}
llvm::SmallVector<int64_t, 4> global_shape;
absl::Status extract_status =
ExtractConstVectorFromValue(conv_op.getInputSizes(), &global_shape);
// If the input is dynamic size, we expect the output is all so dynamic size
// since they should roughly be the same shape. Don't support this for right
// now. The easy way to support this is to all gather the gradient input and
// compute this as a large local convolution and then slice to the output
// layout.
if (!extract_status.ok()) {
return errors::InvalidArgument("{0} requires static shape for input size.",
conv_op->getName().getStringRef().str());
}
// Compute the 'true' input/output layout of the operation. E.g. batch sharded
// vs non-batch sharded. If at least one of the the input gradient or output
// gradient is batch sharded, use that dimension.
std::string batch_sharding_dimension = grad_layout.sharding_spec(0);
if (batch_sharding_dimension == Layout::kUnshardedDim) {
batch_sharding_dimension = output_layout.sharding_spec(0);
} else if ((output_layout.sharding_spec(0) != Layout::kUnshardedDim) &&
(batch_sharding_dimension != output_layout.sharding_spec(0))) {
return errors::InvalidArgument(
"Input and output layout to {2} have incompatible sharding dimensions: "
"\"{0}\" and \"{1}\".",
grad_layout.sharding_spec(0), output_layout.sharding_spec(0),
conv_op->getName().getStringRef().str());
}
const Layout desired_input_gradient_layout =
Layout::BatchShardedLike(grad_layout, batch_sharding_dimension);
const Layout desired_output_gradient_layout =
Layout::BatchShardedLike(output_layout, batch_sharding_dimension);
const Layout desired_input_layout = Layout::BatchShardedOnMesh(
grad_layout.mesh(), global_shape.size(), batch_sharding_dimension);
const std::vector<int64_t> local_shape =
desired_input_layout.LocalShapeFromGlobalShape(global_shape);
mlir::OpBuilder builder(conv_op.getOperation());
mlir::Value new_const = IntConst(
builder, conv_op->getLoc(),
llvm::SmallVector<int32_t, 4>(local_shape.begin(), local_shape.end()));
conv_op.getInputSizesMutable().assign(new_const);
TF_ASSIGN_OR_RETURN(mlir::Value local_input_gradient,
EmitRelayout(conv_op.getOutBackprop(), grad_layout,
desired_input_gradient_layout));
conv_op.getOutBackpropMutable().assign(local_input_gradient);
InferSPMDExpandedLocalShape(conv_op);
llvm::SmallPtrSet<mlir::Operation*, 4> newly_created_ops;
TF_ASSIGN_OR_RETURN(
mlir::Value local_output_gradient,
EmitRelayout(conv_op.getOutput(), desired_output_gradient_layout,
output_layout, &newly_created_ops));
conv_op.getOutput().replaceAllUsesExcept(local_output_gradient,
newly_created_ops);
return local_output_gradient.getDefiningOp();
}
// This expands backprop ops which take tensor inputs into those which take
// sizes. We first convert the (currently local) input shape to global and use
// the const rather than input. The shape will be converted back to local in
// HandleConvBackpropInput, but this is the correct behavior as
// HandleConvBackpropInput will decided how it wants to expand the op.
template <typename To, typename From>
StatusOr<mlir::Operation*> HandleConvBackpropInputTensor(
const Layout& output_layout, From conv_op) {
TF_ASSIGN_OR_RETURN(llvm::SmallVector<int64_t> local_shape,
GetTFShapeFromType(conv_op.getInput().getType()));
TF_ASSIGN_OR_RETURN(const Layout input_layout,
ExtractRequiredLayoutFromOperand(conv_op.getInput()));
const std::vector<int64_t> global_shape =
input_layout.GlobalShapeFromLocalShape(local_shape);
mlir::OpBuilder builder(conv_op);
mlir::Value global_input_shape = IntConst(
builder, conv_op->getLoc(),
llvm::SmallVector<int32_t, 4>(global_shape.begin(), global_shape.end()));
// Insert a replicated layout along this edge, so that we can call
// HandleConvBackpropInput which expects there to be a layout here.
mlir::TF::ShapeAttr global_input_shape_shape = mlir::TF::ShapeAttr::get(
builder.getContext(),
mlir::cast<mlir::TensorType>(global_input_shape.getType()));
mlir::TF::DTensorLayout global_input_shape_with_layout =
mlir::TF::DTensorLayout::create(
builder, conv_op->getLoc(), global_input_shape,
mlir::dtensor::LayoutAttr::get(
builder.getContext(),
Layout::ReplicatedOnMesh(input_layout.mesh(), 1)),
global_input_shape_shape);
To new_conv = To::create(
builder, conv_op->getLoc(), conv_op->getResultTypes(),
mlir::ValueRange({global_input_shape_with_layout, conv_op.getFilter(),
conv_op.getOutBackprop()}),
conv_op->getAttrs());
conv_op.getOutput().replaceAllUsesWith(new_conv.getOutput());
conv_op.erase();
return HandleConvBackpropInput(output_layout, new_conv);
}
template <typename ConvBackpropFilterOp>
StatusOr<mlir::Operation*> HandleConvBackpropFilter(
const Layout& output_layout, ConvBackpropFilterOp conv_op) {
TF_ASSIGN_OR_RETURN(std::vector<Layout> input_layouts,
ExtractRequiredLayoutFromOperands(conv_op));
const Layout& input_layout = input_layouts[0];
const Layout& filter_shape_layout = input_layouts[1];
const Layout& grad_layout = input_layouts[2];
// We only support batch sharding for these. In this case the input
// activations and input gradient should both be batch sharded and
// the output (the filter gradient) should be replicated.
if (!(output_layout.IsBatchParallel() || output_layout.IsFullyReplicated()) ||
!(grad_layout.IsBatchParallel() || grad_layout.IsFullyReplicated())) {
return errors::InvalidArgument("{0} only supports batch parallel layouts.",
conv_op->getName().getStringRef().str());
}
if (!output_layout.IsFullyReplicated()) {
return errors::InvalidArgument("{0} only supports replicated filters.",
conv_op->getName().getStringRef().str());
}
if (!filter_shape_layout.IsFullyReplicated()) {
return errors::InvalidArgument(
"Filter shape input (parameter 1) for {0} must have replicated layout.",
conv_op->getName().getStringRef().str());
}
// Compute the 'true' input layouts of the operation. E.g. batch sharded
// vs non-batch sharded. Basically we get the batch sharding dimension from
// one of the inputs and check that the other is potentially sharded on the
// same dimension.
// TODO(b/262417847): if batch_sharding_dimension is Layout::kUnsharded, then
// we should consider sharding the input here. It may be faster to spread
// the convolution out and then all reduce after vs running it all locally.
std::string batch_sharding_dimension = input_layout.sharding_spec(0);
if (batch_sharding_dimension == Layout::kUnshardedDim) {
batch_sharding_dimension = grad_layout.sharding_spec(0);
} else if ((grad_layout.sharding_spec(0) != Layout::kUnshardedDim) &&
(batch_sharding_dimension != grad_layout.sharding_spec(0))) {
return errors::InvalidArgument(
"Input and gradient layouts for {2} have incompatible batch sharding "
"dimensions: \"{0}\" and \"{1}\".",
input_layouts[0].sharding_spec(0), input_layouts[0].sharding_spec(0),
conv_op->getName().getStringRef().str());
}
const Layout desired_input_activation_layout =
Layout::BatchShardedLike(input_layout, batch_sharding_dimension);
const Layout desired_input_gradient_layout =
Layout::BatchShardedLike(grad_layout, batch_sharding_dimension);
TF_ASSIGN_OR_RETURN(mlir::Value local_input_activation,
EmitRelayout(conv_op.getInput(), input_layout,
desired_input_activation_layout));
conv_op.getInputMutable().assign(local_input_activation);
TF_ASSIGN_OR_RETURN(mlir::Value local_input_gradient,
EmitRelayout(conv_op.getOutBackprop(), grad_layout,
desired_input_gradient_layout));
conv_op.getOutBackpropMutable().assign(local_input_gradient);
InferSPMDExpandedLocalShape(conv_op);
// Output shall be replicated. If we were batch sharded, we need to
// all-reduce the partial results.
if (batch_sharding_dimension != Layout::kUnshardedDim) {
mlir::OpBuilder builder(conv_op.getOperation());
builder.setInsertionPointAfter(conv_op);
return DT_CTX(EmitAllReduce(builder, output_layout,
{batch_sharding_dimension}, conv_op,
kReduceOpAdd));
}
return conv_op.getOperation();
}
// This expands backprop ops which take tensor inputs into those which take
// sizes. We check that the filter input shape is global and then make that a
// const, and replace the op with the version taking shapes.
template <typename To, typename From>
StatusOr<mlir::Operation*> HandleConvBackpropFilterTensor(
const Layout& output_layout, From conv_op) {
TF_ASSIGN_OR_RETURN(const Layout filter_layout,
ExtractRequiredLayoutFromOperand(conv_op.getFilter()));
if (!filter_layout.IsFullyReplicated()) {
return absl::InvalidArgumentError(
"Convolution backpropation ops only support replicated filters.");
}
TF_ASSIGN_OR_RETURN(llvm::SmallVector<int64_t> global_filter_shape,
GetTFShapeFromType(conv_op.getFilter().getType()));
mlir::OpBuilder builder(conv_op);
mlir::Value global_filter_shape_const =
IntConst(builder, conv_op->getLoc(),
llvm::SmallVector<int32_t, 4>(global_filter_shape.begin(),
global_filter_shape.end()));
// Insert a replicated layout along this edge, so that we can call
// HandleConvBackpropInput which expects there to be a layout here.
mlir::TF::ShapeAttr global_filter_shape_shape = mlir::TF::ShapeAttr::get(
builder.getContext(),
mlir::cast<mlir::TensorType>(global_filter_shape_const.getType()));
mlir::TF::DTensorLayout global_filter_shape_with_layout =
mlir::TF::DTensorLayout::create(
builder, conv_op->getLoc(), global_filter_shape_const,
mlir::dtensor::LayoutAttr::get(
builder.getContext(),
Layout::ReplicatedOnMesh(filter_layout.mesh(), 1)),
global_filter_shape_shape);
To new_conv = To::create(
builder, conv_op->getLoc(), conv_op->getResultTypes(),
mlir::ValueRange({conv_op.getInput(), global_filter_shape_with_layout,
conv_op.getOutBackprop()}),
conv_op->getAttrs());
conv_op.getOutput().replaceAllUsesWith(new_conv.getOutput());
conv_op.erase();
return HandleConvBackpropFilter(output_layout, new_conv);
}
StatusOr<mlir::Operation*> HandleMaxPoolGradOp(
const Layout& output_layout, mlir::TF::MaxPoolGradOp max_pool_grad_op) {
// MaxPoolGrad has 3 inputs: Original Input to MaxPool, Output of MaxPool and
// Gradients.
assert(max_pool_grad_op->getOpOperands().size() == 3);
// Relayout gradient input to match layout of output of maxpool.
mlir::OpOperand& max_pool_output = max_pool_grad_op->getOpOperand(1);
TF_ASSIGN_OR_RETURN(Layout max_pool_output_layout,
ExtractRequiredLayoutFromOperand(max_pool_output.get()));
mlir::OpOperand& grad_input = max_pool_grad_op->getOpOperand(2);
TF_ASSIGN_OR_RETURN(Layout grad_input_layout,
ExtractRequiredLayoutFromOperand(grad_input.get()));
TF_ASSIGN_OR_RETURN(mlir::Value new_grad_input,
EmitRelayout(grad_input.get(), grad_input_layout,
max_pool_output_layout));
grad_input.set(new_grad_input);
return InferSPMDExpandedLocalShape(max_pool_grad_op);
}
} // namespace
StatusOr<mlir::Operation*> ConvSPMDExpander::ExpandOp(mlir::Operation* op) {
// The first argument to Conv2DBackpropInputOp is the shape of the input we
// are generating. Since this is almost always the output of a call to
// `shape`, we lose the ability to infer the original input layout. (c.f if
// Conv2DBackpropInput accepted the input _tensor_ instead of the shape).
// Since in eager execution, we cannot look ahead at consumer operations, we
// instead attach the original input layout as a secondary attribute on the
// output of the shape operation, and use this to infer the desired layout for
// this op.
TF_ASSIGN_OR_RETURN(const auto output_layout, ExtractSingleLayoutFromOp(op));
// Forward prop ops.
if (llvm::isa<mlir::TF::Conv2DOp>(op))
return HandleConv<>(llvm::cast<mlir::TF::Conv2DOp>(op));
if (llvm::isa<mlir::TF::Conv3DOp>(op))
return HandleConv<>(llvm::cast<mlir::TF::Conv3DOp>(op));
// Backward prop input ops.
if (llvm::isa<mlir::TF::Conv2DBackpropInputOp>(op))
return HandleConvBackpropInput<>(
*output_layout, llvm::cast<mlir::TF::Conv2DBackpropInputOp>(op));
if (auto conv_op = llvm::dyn_cast<mlir::TF::Conv2DBackpropInputV2Op>(op))
return HandleConvBackpropInputTensor<mlir::TF::Conv2DBackpropInputOp>(
*output_layout, conv_op);
if (auto conv_op = llvm::dyn_cast<mlir::TF::Conv3DBackpropInputOp>(op))
return HandleConvBackpropInputTensor<mlir::TF::Conv3DBackpropInputV2Op>(
*output_layout, conv_op);
if (llvm::isa<mlir::TF::Conv3DBackpropInputV2Op>(op))
return HandleConvBackpropInput<>(
*output_layout, llvm::cast<mlir::TF::Conv3DBackpropInputV2Op>(op));
// Backward prop filter ops.
if (llvm::isa<mlir::TF::Conv2DBackpropFilterOp>(op))
return HandleConvBackpropFilter<>(
*output_layout, llvm::cast<mlir::TF::Conv2DBackpropFilterOp>(op));
if (auto conv_op = llvm::dyn_cast<mlir::TF::Conv2DBackpropFilterV2Op>(op))
return HandleConvBackpropFilterTensor<mlir::TF::Conv2DBackpropFilterOp>(
*output_layout, conv_op);
if (auto conv_op = llvm::dyn_cast<mlir::TF::Conv3DBackpropFilterOp>(op))
return HandleConvBackpropFilterTensor<mlir::TF::Conv3DBackpropFilterV2Op>(
*output_layout, conv_op);
if (llvm::isa<mlir::TF::Conv3DBackpropFilterV2Op>(op))
return HandleConvBackpropFilter<>(
*output_layout, llvm::cast<mlir::TF::Conv3DBackpropFilterV2Op>(op));
// For all other ops, only batch sharded or fully replicated sharding is
// supported for now.
if (!output_layout->IsFullyReplicated() && !output_layout->IsBatchParallel())
return absl::UnimplementedError(
llvm::formatv(
"Only replicated or batch parallel layout is supported in "
"expansion of {0}, but got output layout: {1}",
op->getName().getStringRef().str(), output_layout->ToString())
.str());
if (auto max_pool_grad = mlir::dyn_cast<mlir::TF::MaxPoolGradOp>(op))
return HandleMaxPoolGradOp(*output_layout, max_pool_grad);
// Local expansion only for all other ops.
return InferSPMDExpandedLocalShape(op);
}
StatusOr<llvm::DenseMap<int, Layout>> ConvSPMDExpander::ComputeLayoutForward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op));
llvm::DenseMap<int, Layout> output_layouts(op->getNumResults());
if (llvm::isa<mlir::TF::Conv2DOp, mlir::TF::Conv3DOp, mlir::TF::MaxPoolOp,
mlir::TF::MaxPoolGradOp>(op)) {
// Conv2d/Conv3d and MaxPool ops are grouped together as they all try to
// propagate layout from input image (operand 0).
// If requested 'input' layout exist, try to request same layout for output.
if (input_layouts.find(0) != input_layouts.end()) {
output_layouts[0] = input_layouts.lookup(0);
} else {
// For MaxPoolGrad, request same layout as 'orig_output' or 'grad'
// whatever is present.
if (llvm::isa<mlir::TF::MaxPoolGradOp>(op)) {
if (input_layouts.find(1) != input_layouts.end())
output_layouts[0] = input_layouts.lookup(1);
else if (input_layouts.find(2) != input_layouts.end())
output_layouts[0] = input_layouts.lookup(2);
}
}
} else if (llvm::isa<mlir::TF::Conv2DBackpropInputOp,
mlir::TF::Conv2DBackpropInputV2Op,
mlir::TF::Conv3DBackpropInputOp,
mlir::TF::Conv3DBackpropInputV2Op>(op)) {
if (llvm::isa<mlir::TF::Conv2DBackpropInputOp,
mlir::TF::Conv3DBackpropInputV2Op>(op)) {
if (input_layouts.find(2) != input_layouts.end()) {
// The propagate the gradient layout to the new gradient, e.g. respect
// the spatial partitioning of the input gradient.
output_layouts[0] = input_layouts.lookup(2);
}
} else {
if (input_layouts.find(0) != input_layouts.end()) {
// The propagate the gradient layout to the new gradient, e.g. respect
// the spatial partitioning of the input.
output_layouts[0] = input_layouts.lookup(0);
}
}
} else if (llvm::isa<mlir::TF::Conv2DBackpropFilterOp,
mlir::TF::Conv2DBackpropFilterV2Op,
mlir::TF::Conv3DBackpropFilterOp,
mlir::TF::Conv3DBackpropFilterV2Op>(op)) {
if (llvm::isa<mlir::TF::Conv2DBackpropFilterOp,
mlir::TF::Conv3DBackpropFilterV2Op>(op)) {
// For the ops which take filter shape as input, just return a replicated
// output shape.
output_layouts[0] =
Layout::ReplicatedOnMesh(mesh, ValueRank(op->getOpResult(0)));
} else if (input_layouts.find(1) != input_layouts.end()) {
// For the ops taking a real filter, just copy the filter layout.
output_layouts[0] = input_layouts.lookup(1);
}
} else {
return absl::InvalidArgumentError(
llvm::formatv(
"Layout propagation for unrecognized convolution op {0} not "
"supported.",
OpName(op))
.str());
}
return output_layouts;
}
StatusOr<llvm::DenseMap<int, Layout>> ConvSPMDExpander::ComputeLayoutBackward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op));
llvm::DenseMap<int, Layout> input_layouts(op->getNumOperands());
if (llvm::isa<mlir::TF::Conv2DOp, mlir::TF::Conv3DOp, mlir::TF::MaxPoolOp,
mlir::TF::MaxPoolGradOp>(op)) {
// If suggested output layout exists, try to request input image to have the
// same layout so that all computation would be local.
if (output_layouts.find(0) != output_layouts.end()) {
const Layout output_layout = output_layouts.lookup(0);
input_layouts[0] = output_layout;
// Request replicated for filter input if Conv2D/Conv3D.
if (llvm::isa<mlir::TF::Conv2DOp, mlir::TF::Conv3DOp>(op)) {
input_layouts[1] =
Layout::ReplicatedOnMesh(mesh, ValueRank(op->getOperand(1)));
}
if (llvm::isa<mlir::TF::MaxPoolGradOp>(op)) {
input_layouts[1] = output_layout; // 'orig_output'
input_layouts[2] = output_layout; // 'grad'
}
}
} else if (llvm::isa<mlir::TF::Conv2DBackpropInputOp,
mlir::TF::Conv2DBackpropInputV2Op,
mlir::TF::Conv3DBackpropInputOp,
mlir::TF::Conv3DBackpropInputV2Op>(op)) {
// Generally mark the filter as replicated.
input_layouts[1] =
Layout::ReplicatedOnMesh(mesh, ValueRank(op->getOperand(1)));
if (llvm::isa<mlir::TF::Conv2DBackpropInputOp,
mlir::TF::Conv3DBackpropInputV2Op>(op)) {
// This input is a shape.
input_layouts[0] =
Layout::ReplicatedOnMesh(mesh, ValueRank(op->getOperand(0)));
}
if (output_layouts.find(0) != output_layouts.end()) {
Layout output_layout = output_layouts.lookup(0);
// Ask for the grad to have the same layout as the output. The reasoning
// here is that the if the output is spatially partitioned, we expect
// that the grad is spatially partitioned as well.
input_layouts[2] = output_layout;
if (llvm::isa<mlir::TF::Conv2DBackpropInputV2Op,
mlir::TF::Conv3DBackpropInputOp>(op)) {
input_layouts[0] = output_layout;
}
}
} else if (llvm::isa<mlir::TF::Conv2DBackpropFilterOp,
mlir::TF::Conv2DBackpropFilterV2Op,
mlir::TF::Conv3DBackpropFilterOp,
mlir::TF::Conv3DBackpropFilterV2Op>(op)) {
// Note for Filter op, we generally expect that the output layout would
// match the variable layout for the filter which is generally replicated.
// The gradient layout most likely needs to agree with the input layout,
// e.g. both spatially partitioned or not. This is somewhat similar to
// MatMul, for now just set both to replicated.
input_layouts[0] =
Layout::ReplicatedOnMesh(mesh, ValueRank(op->getOperand(0)));
input_layouts[2] =
Layout::ReplicatedOnMesh(mesh, ValueRank(op->getOperand(2)));
if (llvm::isa<mlir::TF::Conv2DBackpropFilterOp,
mlir::TF::Conv3DBackpropFilterV2Op>(op)) {
// For ops taking filter shape as input, just use a replicated input
// layout.
input_layouts[1] =
Layout::ReplicatedOnMesh(mesh, ValueRank(op->getOperand(1)));
} else if (output_layouts.find(0) != output_layouts.end()) {
// For ops taking filer directly as input copy the output layout to the
// filter layout.
input_layouts[1] = output_layouts.lookup(0);
}
} else {
return absl::InvalidArgumentError(
llvm::formatv(
"Layout propagation for unrecognized convolution op {0} not "
"supported.",
OpName(op))
.str());
}
return input_layouts;
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/expansions/conv_spmd_expander.h
+51
View File
@@ -0,0 +1,51 @@
/* Copyright 2022 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_DTENSOR_MLIR_EXPANSIONS_CONV_SPMD_EXPANDER_H_
#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_CONV_SPMD_EXPANDER_H_
#include "llvm/ADT/DenseMap.h"
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/spmd_expander.h"
namespace tensorflow {
namespace dtensor {
// Implement Layout propagation and SPMD expansion for Convolution ops.
//
// The extended class will be registered in spmd_expander.cc for Conv2D/3D and
// Conv2D/3D Backprop ops to enable proper DTensor behavior of them. This
// implementation is internal and specific to DTensor while upstream(python)
// users won't need to use this class directly in any fashion.
class ConvSPMDExpander : public SPMDExpanderBase {
public:
StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& input_layouts) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& output_layouts) override;
};
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_CONV_SPMD_EXPANDER_H_
tensorflow/dtensor/mlir/expansions/cumsum_spmd_expander.cc
+139
View File
@@ -0,0 +1,139 @@
/* Copyright 2022 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/dtensor/mlir/expansions/cumsum_spmd_expander.h"
#include <cassert>
#include <cstdint>
#include <string>
#include "absl/strings/str_cat.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/Support/Casting.h"
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/collectives.h"
#include "tensorflow/dtensor/mlir/layout_parsing.h"
#include "tensorflow/dtensor/mlir/op_utils.h"
#include "tensorflow/dtensor/mlir/shape_utils.h"
#include "tensorflow/dtensor/mlir/value_utils.h"
namespace tensorflow {
namespace dtensor {
namespace {
// Extract `axis` tensor from Cumsum op and return it's positive value, since
// it can be a negative index.
StatusOr<int64_t> GetAxisDimension(mlir::Operation* op) {
auto cumsum = llvm::dyn_cast<mlir::TF::CumsumOp>(op);
if (cumsum == nullptr) {
return absl::InternalError(
absl::StrCat("Expected Cumsum op but got : ", OpName(op)));
}
TF_ASSIGN_OR_RETURN(int64_t axis_dim,
ExtractConstIntFromValue(cumsum.getAxis()));
int64_t tensor_rank = ValueRank(cumsum.getX());
// Axis can be in range [-tensor_rank, tensor_rank), so we add tensor_rank
// to wrap it around.
if (axis_dim >= -tensor_rank && axis_dim < 0) {
axis_dim += tensor_rank;
} else if (axis_dim < -tensor_rank || axis_dim >= tensor_rank) {
return absl::InvalidArgumentError(
"Invalid axis; expected a value in [-tensor_rank, tensor_rank)");
}
return axis_dim;
}
} // namespace
StatusOr<mlir::Operation*> CumsumSPMDExpander::ExpandOp(mlir::Operation* op) {
StatusOr<int64_t> axis_dim = GetAxisDimension(op);
if (!axis_dim.ok()) return axis_dim.status();
TF_ASSIGN_OR_RETURN(auto output_layout, ExtractSingleLayoutFromOp(op));
assert(output_layout);
// Our intermediate computation layout is the output layout with
// the axis dimension replicated. So set both the operand and output layout
// to this intermediate layout.
TF_ASSIGN_OR_RETURN(Layout intermediate_layout,
output_layout->GetLayoutWithReducedDims(
{axis_dim.value()}, /*keep_dims=*/true));
// Relayout operand to intermediate layout.
mlir::OpBuilder builder(op);
const auto operand = op->getOperand(0);
TF_ASSIGN_OR_RETURN(auto operand_layout, ExtractLayoutFromOperand(operand));
if (!operand_layout)
return absl::InvalidArgumentError(
"input layout of Cumsum op must be known before SPMD "
"expansion.");
TF_ASSIGN_OR_RETURN(
const auto new_operand,
EmitRelayout(operand, operand_layout.value(), intermediate_layout));
op->setOperand(0, new_operand);
op = InferSPMDExpandedLocalShape(op);
// Relayout output to intermediate layout.
llvm::SmallPtrSet<mlir::Operation*, 4> newly_created_ops;
builder.setInsertionPointAfter(op);
TF_ASSIGN_OR_RETURN(auto final_output,
EmitRelayout(op->getOpResult(0), intermediate_layout,
output_layout.value(), &newly_created_ops));
op->getOpResult(0).replaceAllUsesExcept(final_output, newly_created_ops);
return final_output.getDefiningOp();
}
StatusOr<llvm::DenseMap<int, Layout>> CumsumSPMDExpander::ComputeLayoutForward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
TF_ASSIGN_OR_RETURN(const auto mesh, ExtractDeviceMeshEnclosingCluster(op));
TF_ASSIGN_OR_RETURN(int64_t axis_dim, GetAxisDimension(op));
if (input_layouts.find(0) == input_layouts.end())
return llvm::DenseMap<int, Layout>();
auto input_layout = input_layouts.lookup(0);
TF_ASSIGN_OR_RETURN(
Layout input_layout_reduced_dims,
input_layout.GetLayoutWithReducedDims({axis_dim},
/*keep_dims=*/true));
return llvm::DenseMap<int, Layout>({{0, input_layout_reduced_dims}});
}
StatusOr<llvm::DenseMap<int, Layout>> CumsumSPMDExpander::ComputeLayoutBackward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
TF_ASSIGN_OR_RETURN(const auto mesh, ExtractDeviceMeshEnclosingCluster(op));
TF_ASSIGN_OR_RETURN(int64_t axis_dim, GetAxisDimension(op));
if (output_layouts.find(0) == output_layouts.end())
return llvm::DenseMap<int, Layout>();
auto output_layout = output_layouts.lookup(0);
TF_ASSIGN_OR_RETURN(
Layout output_layout_reduced_dims,
output_layout.GetLayoutWithReducedDims({axis_dim},
/*keep_dims=*/true));
return llvm::DenseMap<int, Layout>({{0, output_layout_reduced_dims}});
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/expansions/cumsum_spmd_expander.h
+44
View File
@@ -0,0 +1,44 @@
/* Copyright 2022 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_DTENSOR_MLIR_EXPANSIONS_CUMSUM_SPMD_EXPANDER_H_
#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_CUMSUM_SPMD_EXPANDER_H_
#include "llvm/ADT/DenseMap.h"
#include "mlir/IR/Operation.h" // from @llvm-project
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/spmd_expander.h"
namespace tensorflow {
namespace dtensor {
class CumsumSPMDExpander : public SPMDExpanderBase {
private:
// SPMD expand the op for local computation.
StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& input_layouts) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& output_layouts) override;
};
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_CUMSUM_SPMD_EXPANDER_H_
tensorflow/dtensor/mlir/expansions/dataparallel_spmd_expander.cc
+365
View File
@@ -0,0 +1,365 @@
/* Copyright 2022 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/dtensor/mlir/expansions/dataparallel_spmd_expander.h"
#include <string>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/FormatVariadic.h"
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.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 "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/collectives.h"
#include "tensorflow/dtensor/mlir/layout_parsing.h"
#include "tensorflow/dtensor/mlir/shape_utils.h"
#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
#include "tensorflow/dtensor/mlir/value_utils.h"
#include "tensorflow/dtensor/proto/layout.pb.h"
namespace tensorflow {
namespace dtensor {
namespace {
// Check all tensors are batch parallel
bool AllBatchParallel(const std::vector<Layout>& layouts,
const llvm::DenseMap<int, int>& batchable_indices) {
for (int i = 0; i < layouts.size(); ++i) {
if (!layouts[i].IsBatchParallel(batchable_indices.lookup(i))) return false;
}
return true;
}
// Check all Layouts have the same batch rank
bool SameBatchRank(const std::vector<Layout>& layouts,
const llvm::DenseMap<int, int>& batchable_indices) {
absl::flat_hash_set<int> batch_ranks;
// Add all batch ranks of layouts
for (auto const& idx_and_non_batch_rank : batchable_indices) {
auto const& idx = idx_and_non_batch_rank.first;
auto const& non_batch_rank = idx_and_non_batch_rank.second;
batch_ranks.insert(layouts[idx].rank() - non_batch_rank);
}
return batch_ranks.size() <= 1;
}
// Check if any layout from a set of indices is not a nullopt
bool AnyLayoutExist(const llvm::DenseMap<int, Layout>& layouts,
const llvm::DenseMap<int, int>& indices) {
for (auto const& idx_and_unused : indices) {
auto const& idx = idx_and_unused.first;
if (layouts.find(idx) != layouts.end()) return true;
}
return false;
}
// Given layouts to merge and a map of {indices of the batchable layouts, rank
// of non batch dimensions} merges those batchable layouts to produce one single
// layout. Assumes all batch ranks are the same for all batchable layouts, which
// is enforced before this
//
// Merged together so that we have the layout which is sharded in a tensor dim
// if and only if all layouts are sharded in the same sharding_spec.
StatusOr<Layout> MergeBatchLayouts(
const llvm::DenseMap<int, Layout>& layouts,
const llvm::DenseMap<int, int>& batchable_args, const Mesh& mesh) {
// Get the batch rank
int layout_idx = -1;
for (auto const& idx_and_unused : batchable_args) {
auto const& idx = idx_and_unused.first;
if (layouts.find(idx) != layouts.end()) layout_idx = idx;
}
int batch_rank =
layouts.lookup(layout_idx).rank() - batchable_args.lookup(layout_idx);
// Initialize with replicated
std::vector<std::string> merged_specs(batch_rank, Layout::kUnshardedDim);
// Merge layouts. If any dimension don't agree on sharding dim, then replicate
for (int i = 0; i < batch_rank; ++i) {
absl::flat_hash_set<std::string> spec_set;
for (auto const& arg_idx_and_unused : batchable_args) {
auto const& arg_idx = arg_idx_and_unused.first;
if (layouts.find(arg_idx) == layouts.end()) continue;
const std::string spec = layouts.lookup(arg_idx).sharding_spec(i);
if (spec != Layout::kUnshardedDim) {
spec_set.insert(spec);
}
}
if (spec_set.size() == 1) {
merged_specs[i] = *spec_set.begin();
} else {
merged_specs[i] = Layout::kUnshardedDim;
}
}
// Deduplicate same usage of mesh dims. [x,x] -> [unsharded, unsharded]
absl::flat_hash_map<std::string, int> counter;
for (const std::string& spec : merged_specs) counter[spec] += 1;
for (std::string& spec : merged_specs) {
if (counter[spec] > 1) {
spec = Layout::kUnshardedDim;
}
}
return Layout::GetLayout(merged_specs, mesh);
}
// Choose an intermediate layout to relayout. Picks the most frequently
// sharded mesh dimension for every batch dimension, then deduplicates (n-1)
// of all repeated mesh dimensions, leaving the rightmost duplicate sharded
//
// Note that this assumes the number of batch dims for every batchable
// tensor is the same and is enforced before this
//
// Examples:
// Given layouts: [x,y],[x,z],[y,z], produces [x,z]
// Deduplication: [x,x] -> will become [*, x]
StatusOr<Layout> IntermediateBatchLayout(
const std::vector<Layout>& operand_layouts,
const llvm::DenseMap<int, int>& batchable_operands,
const std::vector<Layout>& output_layouts,
const llvm::DenseMap<int, int>& batchable_outputs, const Mesh& mesh) {
if (batchable_operands.empty()) {
return absl::UnimplementedError(
llvm::formatv("There must be at least one batchable operand").str());
}
int first_batcharg_index = batchable_outputs.begin()->first;
int batch_rank = operand_layouts[first_batcharg_index].rank() -
batchable_operands.find(first_batcharg_index)->second;
std::vector<std::string> batch_specs(batch_rank, Layout::kUnshardedDim);
// For each batch dimension, finds the most commonly used mesh dimension
// and sets that to batch_specs[i].
for (int i = 0; i < batch_rank; ++i) {
std::string mesh_dim = Layout::kUnshardedDim;
int max_count = 0;
absl::flat_hash_map<std::string, int> counter;
// add operand counts
for (auto const& idx_and_unused : batchable_operands) {
auto const& idx = idx_and_unused.first;
std::string spec = operand_layouts[idx].sharding_spec(i);
if (spec != Layout::kUnshardedDim) counter[spec]++;
if (counter[spec] > max_count) {
max_count = counter[spec];
mesh_dim = spec;
}
}
// add output counts
for (auto const& idx_and_unused : batchable_outputs) {
auto const& idx = idx_and_unused.first;
std::string spec = output_layouts[idx].sharding_spec(i);
if (spec != Layout::kUnshardedDim) counter[spec]++;
if (counter[spec] > max_count) {
max_count = counter[spec];
mesh_dim = spec;
}
}
batch_specs[i] = mesh_dim;
}
// deduplicate
absl::flat_hash_map<std::string, int> counter;
for (const std::string& spec : batch_specs) counter[spec] += 1;
for (std::string& spec : batch_specs) {
if (counter[spec] > 1) {
counter[spec]--;
spec = Layout::kUnshardedDim;
}
}
return Layout::GetLayout(batch_specs, mesh);
}
} // namespace
// Relayout all operands that have batch dimensions to batch sharded
// The outputs will get the correct inferred shape from the operands
StatusOr<mlir::Operation*> DataparallelSPMDExpander::RelayoutOperandsAndOutputs(
mlir::Operation* op, const std::vector<Layout>& operand_layouts,
const std::vector<Layout>& output_layouts) {
TF_ASSIGN_OR_RETURN(const auto mesh, ExtractDeviceMeshEnclosingCluster(op));
TF_ASSIGN_OR_RETURN(
const Layout intermediate_batch_layout,
IntermediateBatchLayout(operand_layouts, batchable_operands_,
output_layouts, batchable_outputs_, mesh));
// Relayout batchable operands
for (auto i = 0; i < operand_layouts.size(); ++i) {
// Relayout operands that have a batch dimension to intermediate layout
if (batchable_operands_.find(i) != batchable_operands_.end()) {
int replicated_rank =
ValueRank(op->getOperand(i)) - intermediate_batch_layout.rank();
TF_ASSIGN_OR_RETURN(
auto new_layout,
ConcatenateLayouts(intermediate_batch_layout,
Layout::ReplicatedOnMesh(mesh, replicated_rank)));
TF_ASSIGN_OR_RETURN(
const auto new_operand,
EmitRelayout(op->getOperand(i), operand_layouts[i], new_layout));
op->setOperand(i, new_operand);
}
}
// Expand to local shape
op = InferSPMDExpandedLocalShape(op);
llvm::SmallPtrSet<mlir::Operation*, 4> newly_created_ops;
llvm::SmallVector<mlir::Value, 4> generated_outputs;
llvm::SmallVector<mlir::Type, 4> generated_types;
// Track the op that comes last after splitting.
mlir::Operation* last_op_after_splitting = op;
// Relayout batchable outputs
for (auto i = 0; i < output_layouts.size(); ++i) {
// Relayout to batch shard if tensor has batch dim
if (batchable_outputs_.find(i) != batchable_outputs_.end()) {
int replicated_rank =
ValueRank(op->getResult(i)) - intermediate_batch_layout.rank();
TF_ASSIGN_OR_RETURN(
auto new_layout,
ConcatenateLayouts(intermediate_batch_layout,
Layout::ReplicatedOnMesh(mesh, replicated_rank)));
TF_ASSIGN_OR_RETURN(auto new_output,
EmitRelayout(op->getOpResult(i), new_layout,
output_layouts[i], &newly_created_ops));
generated_outputs.emplace_back(new_output);
generated_types.emplace_back(new_output.getType());
if (last_op_after_splitting->isBeforeInBlock(
new_output.getDefiningOp())) {
last_op_after_splitting = new_output.getDefiningOp();
}
} else {
generated_outputs.push_back(op->getResult(i));
generated_types.push_back(op->getResult(i).getType());
}
}
mlir::OpBuilder builder(op);
builder.setInsertionPointAfter(last_op_after_splitting);
// Tie all outputs together with identity_n
auto identity_op = mlir::TF::IdentityNOp::create(
builder, op->getLoc(), generated_types, generated_outputs);
newly_created_ops.insert(identity_op);
for (int i = 0; i < output_layouts.size(); ++i) {
op->getOpResult(i).replaceAllUsesExcept(identity_op.getResult(i),
newly_created_ops);
}
return identity_op.getOperation();
}
StatusOr<mlir::Operation*> DataparallelSPMDExpander::ExpandOp(
mlir::Operation* op) {
TF_ASSIGN_OR_RETURN(const auto output_layouts,
ExtractRequiredLayoutFromOp(op));
TF_ASSIGN_OR_RETURN(const auto operand_layouts,
ExtractRequiredLayoutFromOperands(op));
// Check all input and output are batch parallel
if (!AllBatchParallel(operand_layouts, batchable_operands_) ||
!AllBatchParallel(output_layouts, batchable_outputs_)) {
return absl::UnimplementedError(
llvm::formatv("All operands and outputs must be batch parallel.")
.str());
}
// Check that the rank of batch dimensions are same for all batchable tensors
if (!SameBatchRank(operand_layouts, batchable_operands_) ||
!SameBatchRank(output_layouts, batchable_outputs_)) {
return absl::UnimplementedError(
llvm::formatv("All operands and outputs with batch dimensions must "
"have same batch dimension rank")
.str());
}
if (AllReplicated(output_layouts) && AllReplicated(operand_layouts))
return InferSPMDExpandedLocalShape(op);
return RelayoutOperandsAndOutputs(op, operand_layouts, output_layouts);
}
// Take all layouts of batchable operands, and merge them to produce a single
// layout for all batchable outputs.
StatusOr<llvm::DenseMap<int, Layout>>
DataparallelSPMDExpander::ComputeLayoutForward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
TF_ASSIGN_OR_RETURN(const auto mesh, ExtractDeviceMeshEnclosingCluster(op));
llvm::DenseMap<int, Layout> output_layouts;
// Compute output layouts
if (AnyLayoutExist(input_layouts, batchable_operands_)) {
TF_ASSIGN_OR_RETURN(
const Layout& batch_output_layout,
MergeBatchLayouts(input_layouts, batchable_operands_, mesh));
for (const auto& output_and_index : llvm::enumerate(op->getOpResults())) {
const int output_index = output_and_index.index();
auto output = output_and_index.value();
int rank = ValueRank(output);
if (batchable_outputs_.find(output_index) != batchable_outputs_.end()) {
int replicated_rank = batchable_outputs_[output_index];
TF_ASSIGN_OR_RETURN(auto new_layout,
ConcatenateLayouts(batch_output_layout,
Layout::ReplicatedOnMesh(
mesh, replicated_rank)));
output_layouts[output_index] = new_layout;
} else {
output_layouts[output_index] = Layout::ReplicatedOnMesh(mesh, rank);
}
}
}
return output_layouts;
}
// Take all layouts of batchable outputs, and merge them to produce a single
// layout for all batchable operands.
StatusOr<llvm::DenseMap<int, Layout>>
DataparallelSPMDExpander::ComputeLayoutBackward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
TF_ASSIGN_OR_RETURN(const auto mesh, ExtractDeviceMeshEnclosingCluster(op));
llvm::DenseMap<int, Layout> input_layouts;
// Compute input layouts in the following way: For operand indices that
// have a batch dimension, batch shard it the same way as the output layouts.
// Otherwise, replicate.
if (AnyLayoutExist(output_layouts, batchable_outputs_)) {
TF_ASSIGN_OR_RETURN(
const Layout& batch_operand_layout,
MergeBatchLayouts(output_layouts, batchable_outputs_, mesh));
for (const auto& operand_and_index : llvm::enumerate(op->getOperands())) {
const int operand_index = operand_and_index.index();
auto operand = operand_and_index.value();
int rank = ValueRank(operand);
if (batchable_operands_.find(operand_index) !=
batchable_operands_.end()) {
int replicated_rank = batchable_operands_[operand_index];
TF_ASSIGN_OR_RETURN(auto new_layout,
ConcatenateLayouts(batch_operand_layout,
Layout::ReplicatedOnMesh(
mesh, replicated_rank)));
input_layouts[operand_index] = new_layout;
} else {
input_layouts[operand_index] = Layout::ReplicatedOnMesh(mesh, rank);
}
}
}
return input_layouts;
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/expansions/dataparallel_spmd_expander.h
+70
View File
@@ -0,0 +1,70 @@
/* Copyright 2022 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_DTENSOR_MLIR_EXPANSIONS_DATAPARALLEL_SPMD_EXPANDER_H_
#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_DATAPARALLEL_SPMD_EXPANDER_H_
#include <utility>
#include <vector>
#include "llvm/ADT/DenseMap.h"
#include "mlir/IR/Operation.h" // from @llvm-project
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/spmd_expander.h"
namespace tensorflow {
namespace dtensor {
// General SPMD Expander for data parallel ops.
// We define data parallel ops as ops that have tensors possibly with a batch
// dimension. Assumes batch dimensions start from the left. Tensors may
// may have multiple batch dimensions, including zero
class DataparallelSPMDExpander : public SPMDExpanderBase {
protected:
// These maps contain {arg_index, non_batch_rank}
// Example is for TF:FFT2D, the batchable_operands and batchable_outputs has
// {0, 2} because the first argument is batchable and the last 2 dimensions
// are the non-batch dimensions
llvm::DenseMap<int, int> batchable_operands_;
llvm::DenseMap<int, int> batchable_outputs_;
StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& input_layouts) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& output_layouts) override;
public:
explicit DataparallelSPMDExpander(llvm::DenseMap<int, int> batchable_operands,
llvm::DenseMap<int, int> batchable_outputs)
: batchable_operands_(std::move(batchable_operands)),
batchable_outputs_(std::move(batchable_outputs)) {}
private:
// Relayouts all operands and outputs with a batch dimensions to a batch
// sharded layout. This should only be called when there is at least one
// batch sharded operand or batch sharded output
StatusOr<mlir::Operation*> RelayoutOperandsAndOutputs(
mlir::Operation* op, const std::vector<Layout>& operand_layouts,
const std::vector<Layout>& output_layouts);
};
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_DATAPARALLEL_SPMD_EXPANDER_H_
tensorflow/dtensor/mlir/expansions/disable_copy_on_read_spmd_expander.cc
+48
View File
@@ -0,0 +1,48 @@
/* Copyright 2022 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/dtensor/mlir/expansions/disable_copy_on_read_spmd_expander.h"
#include "llvm/ADT/DenseMap.h"
#include "mlir/IR/Operation.h" // from @llvm-project
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/shape_utils.h"
namespace tensorflow {
namespace dtensor {
StatusOr<mlir::Operation*> DisableCopyOnReadSPMDExpander::ExpandOp(
mlir::Operation* op) {
return InferSPMDExpandedLocalShape(op);
}
StatusOr<llvm::DenseMap<int, Layout>>
DisableCopyOnReadSPMDExpander::ComputeLayoutForward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
// DisableCopyOnRead has no outputs;
return llvm::DenseMap<int, Layout>();
}
StatusOr<llvm::DenseMap<int, Layout>>
DisableCopyOnReadSPMDExpander::ComputeLayoutBackward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& operand_layouts,
const llvm::DenseMap<int, Layout>& output_layouts) {
// Prefer the layout from operand zero.
return operand_layouts;
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/expansions/disable_copy_on_read_spmd_expander.h
+44
View File
@@ -0,0 +1,44 @@
/* Copyright 2022 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_DTENSOR_MLIR_EXPANSIONS_DISABLE_COPY_ON_READ_SPMD_EXPANDER_H_
#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_DISABLE_COPY_ON_READ_SPMD_EXPANDER_H_
#include "llvm/ADT/DenseMap.h"
#include "mlir/IR/Operation.h" // from @llvm-project
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/spmd_expander.h"
namespace tensorflow {
namespace dtensor {
class DisableCopyOnReadSPMDExpander : public SPMDExpanderBase {
public:
StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& input_layouts) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& operand_layouts,
const llvm::DenseMap<int, Layout>& output_layouts) override;
};
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_DISABLE_COPY_ON_READ_SPMD_EXPANDER_H_
tensorflow/dtensor/mlir/expansions/dtensor_op_spmd_expander.cc
+342
View File
@@ -0,0 +1,342 @@
/* Copyright 2022 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/dtensor/mlir/expansions/dtensor_op_spmd_expander.h"
#include <optional>
#include <string>
#include <vector>
#include "absl/container/flat_hash_set.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/types/optional.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/FormatVariadic.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/collectives.h"
#include "tensorflow/dtensor/mlir/dtensor_send_recv.h"
#include "tensorflow/dtensor/mlir/ir/tf_dtensor.h"
#include "tensorflow/dtensor/mlir/layout_parsing.h"
#include "tensorflow/dtensor/mlir/op_utils.h"
namespace tensorflow {
namespace dtensor {
namespace {
// FIXME(feyu): This function should take layouts as arguments. It doesn't need
// the ops.
// Validates send/recv layout and mesh configurations. Among other things, this
// checks for below constraints.
// 1. Src/target layouts have non empty mesh.
// 2. Src/target layouts have the same host.
// 3. Src/target layouts are from different mesh.
// 4. One of scr/target layout is from host mesh cluster.
// 5. CPU host cluster mesh has 1 device.
absl::Status ValidateSendRecvLayoutConfiguration(
mlir::TF::DTensorSend dtensor_send, mlir::TF::DTensorRecv dtensor_recv) {
// If either one of the send/recv ops has already been lowered, then send/recv
// configuration has already been verified.
if (!dtensor_send || !dtensor_recv) return absl::OkStatus();
TF_ASSIGN_OR_RETURN(const absl::optional<Layout> send_layout_or_null,
ExtractLayoutFromOperand(dtensor_send.getInput()));
if (!send_layout_or_null.has_value())
return absl::InvalidArgumentError(
"Input to DTensorSend must have specified layout.");
TF_ASSIGN_OR_RETURN(const Layout output_layout,
ExtractRequiredSingleLayoutFromOp(dtensor_recv));
const Layout& send_layout = send_layout_or_null.value();
const Layout& recv_layout = output_layout;
const Mesh& recv_mesh = dtensor_recv.getMesh();
const Mesh& send_mesh = send_layout.mesh();
// If any one of send/recv mesh are empty, return error.
if (send_mesh.IsEmpty() || recv_mesh.IsEmpty())
return absl::InvalidArgumentError(
"Found empty mesh when sending/receiving tensor across clusters.");
// If send host not found in list of receiving hosts, return error.
std::vector<std::string> send_hosts = send_layout.ReducedMesh().hosts();
std::vector<std::string> recv_hosts = recv_layout.ReducedMesh().hosts();
if (send_hosts != recv_hosts)
return absl::InvalidArgumentError("Send and receive hosts don't match");
// Check shards in sending host match those in the receiving host.
const auto send_host_shard_map = send_layout.HostShardMap();
const auto recv_host_shard_map = recv_layout.HostShardMap();
for (const std::string& host : send_hosts) {
const ShardVector& shards_in_send_host =
send_host_shard_map.find(host)->second;
ShardVector shards_in_recv_host = recv_host_shard_map.find(host)->second;
if (shards_in_send_host != shards_in_recv_host)
return absl::InvalidArgumentError(absl::StrCat(
"Send and receive host shard vectors don't match. Send shard_vector:",
shards_in_send_host.ToString(),
" / Recv host spec : ", shards_in_recv_host.ToString()));
}
// Send/Recv mesh must be different.
if (recv_mesh == send_mesh)
return absl::InvalidArgumentError(
"Found CopyToMesh op sending tensor to same mesh. Only use "
"CopyToMesh to transfer data across different mesh cluster. For "
"changing layout within the same mesh, use tf.Relayout op.");
// Either one of send/recv pair must be to/from CPU mesh.
// For example, TPU mesh -> GPU mesh or TPU mesh -> another TPU mesh
// is disallowed.
if (!send_mesh.is_cpu_mesh() && !recv_mesh.is_cpu_mesh())
return absl::InvalidArgumentError(
"tf.CopyToMesh op must be used to send data from/to host mesh.");
return absl::OkStatus();
}
template <typename RelayoutOp>
StatusOr<mlir::Operation*> ExpandRelayoutOp(RelayoutOp relayout,
Layout target_layout,
Layout input_layout,
Layout output_layout) {
bool match_present = false;
for (const std::string& sharding_spec : target_layout.sharding_spec_strs())
if (sharding_spec == Layout::kMatch) match_present = true;
if (!match_present && output_layout != target_layout)
return absl::InternalError(
"output layout of Relayout op after layout propagation does not match "
"layout specified by Relayout op.");
auto value_or_status =
EmitRelayout(relayout.getInput(), input_layout, output_layout);
if (!value_or_status.ok())
return absl::InvalidArgumentError(
llvm::formatv("Unsupported layout received for {0} op. Trying "
"to set tensor "
"to layout : {1}. Found error {2}",
relayout->getName().getStringRef(),
target_layout.ToString(),
value_or_status.status().message())
.str());
mlir::Value output = value_or_status.value();
relayout.getOutput().replaceAllUsesWith(output);
relayout.erase();
return output.getDefiningOp();
}
// Takes relayout which may have kMatch dimensions and uses it to mask input.
// Here source_layout
StatusOr<Layout> MergeLayouts(
const absl::flat_hash_set<std::string>& used_mesh_dimensions,
const Layout& mask_layout, const Layout& target_layout) {
std::vector<std::string> sharding_specs(mask_layout.sharding_spec_strs());
for (int i = 0; i < target_layout.rank(); ++i) {
if (sharding_specs[i] == Layout::kMatch &&
!used_mesh_dimensions.contains(target_layout.sharding_spec(i)))
sharding_specs[i] = target_layout.sharding_spec(i);
}
return Layout::GetLayout(sharding_specs, target_layout.mesh());
}
// Computes the layout of Relayout's (or RelayoutLike's) input or output, based
// on the layout from the corresponding output or input (as `incoming_layout`).
// Note that this implies that we compute the same layout for the
// operand and output.
// `mask_layout` is set to the user-supplied layout attribute on the op.
StatusOr<llvm::DenseMap<int, Layout>> ComputeRelayoutLayout(
const Layout& mask_layout, std::optional<const Layout> incoming_layout) {
absl::flat_hash_set<std::string> used_dimensions;
bool match_present = false;
for (const std::string& sharding_spec : mask_layout.sharding_spec_strs()) {
if (sharding_spec == Layout::kMatch)
match_present = true;
else if (Layout::IsShardedDimension(sharding_spec))
used_dimensions.insert(sharding_spec);
}
if (!match_present) {
return llvm::DenseMap<int, Layout>({{0, mask_layout}});
}
if (incoming_layout) {
TF_ASSIGN_OR_RETURN(
Layout new_layout,
MergeLayouts(used_dimensions, mask_layout, *incoming_layout));
return llvm::DenseMap<int, Layout>({{0, new_layout}});
}
return llvm::DenseMap<int, Layout>();
}
} // namespace
StatusOr<mlir::Operation*> RelayoutSPMDExpander::ExpandOp(mlir::Operation* op) {
auto relayout = mlir::cast<mlir::TF::RelayoutOp>(op);
TF_ASSIGN_OR_RETURN(const Layout target_layout,
Layout::FromString(relayout.getLayout().str()));
TF_ASSIGN_OR_RETURN(const Layout output_layout,
ExtractRequiredSingleLayoutFromOp(op));
TF_ASSIGN_OR_RETURN(const Layout input_layout,
ExtractRequiredLayoutFromOperand(relayout.getInput()));
return ExpandRelayoutOp<mlir::TF::RelayoutOp>(relayout, target_layout,
input_layout, output_layout);
}
StatusOr<llvm::DenseMap<int, Layout>>
RelayoutSPMDExpander::ComputeLayoutForward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
auto relayout = llvm::cast<mlir::TF::RelayoutOp>(op);
TF_ASSIGN_OR_RETURN(const Layout mask_layout,
Layout::FromString(relayout.getLayout().str()));
std::optional<const Layout> incoming_layout;
if (input_layouts.find(0) != input_layouts.end())
incoming_layout.emplace(input_layouts.lookup(0));
return ComputeRelayoutLayout(mask_layout, incoming_layout);
}
StatusOr<llvm::DenseMap<int, Layout>>
RelayoutSPMDExpander::ComputeLayoutBackward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
auto relayout = llvm::cast<mlir::TF::RelayoutOp>(op);
TF_ASSIGN_OR_RETURN(const Layout mask_layout,
Layout::FromString(relayout.getLayout().str()));
std::optional<const Layout> incoming_layout;
if (output_layouts.find(0) != output_layouts.end())
incoming_layout.emplace(output_layouts.lookup(0));
return ComputeRelayoutLayout(mask_layout, incoming_layout);
}
StatusOr<mlir::Operation*> RelayoutLikeSPMDExpander::ExpandOp(
mlir::Operation* op) {
auto relayout_grad = mlir::cast<mlir::TF::RelayoutLikeOp>(op);
TF_ASSIGN_OR_RETURN(
const Layout target_layout,
ExtractRequiredLayoutFromOperand(relayout_grad.getLayoutInput()));
TF_ASSIGN_OR_RETURN(const Layout output_layout,
ExtractRequiredSingleLayoutFromOp(op));
TF_ASSIGN_OR_RETURN(
const Layout input_layout,
ExtractRequiredLayoutFromOperand(relayout_grad.getInput()));
return ExpandRelayoutOp<mlir::TF::RelayoutLikeOp>(
relayout_grad, target_layout, input_layout, output_layout);
}
StatusOr<llvm::DenseMap<int, Layout>>
RelayoutLikeSPMDExpander::ComputeLayoutForward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
// RelayoutLike's output has the same layout as the corresponding Relayout's
// input operand.
if (input_layouts.find(1) == input_layouts.end())
return llvm::DenseMap<int, Layout>();
return llvm::DenseMap<int, Layout>({{0, input_layouts.lookup(1)}});
}
StatusOr<llvm::DenseMap<int, Layout>>
RelayoutLikeSPMDExpander::ComputeLayoutBackward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
if (output_layouts.find(0) == output_layouts.end())
return llvm::DenseMap<int, Layout>();
const Layout output_layout = output_layouts.lookup(0);
return llvm::DenseMap<int, Layout>({
// Return replicated layout for the input operand since we do not want to
// enforce any particular layout on it.
{0, Layout::ReplicatedLike(output_layout)},
// Set layout for the forward pass's input operand to match the output of
// the RelayoutLike op.
{1, output_layout},
});
}
StatusOr<mlir::Operation*> DTensorSendSPMDExpander::ExpandOp(
mlir::Operation* op) {
mlir::ModuleOp module = op->getParentOfType<mlir::ModuleOp>();
auto dtensor_send = llvm::cast<mlir::TF::DTensorSend>(op);
TF_ASSIGN_OR_RETURN(mlir::Operation * recv_op,
GetCorrespondingDTensorSendRecvOp<mlir::TF::DTensorSend>(
module, dtensor_send));
auto dtensor_recv = llvm::dyn_cast<mlir::TF::DTensorRecv>(recv_op);
TF_RETURN_IF_ERROR(
ValidateSendRecvLayoutConfiguration(dtensor_send, dtensor_recv));
return LowerDTensorSend(op, recv_op);
}
// DTensorSend op respects input layout from input operations and does not
// set any preferred inputs layouts. During SPMD expansion, however, tensor
// values are changed to replicated layout before transferring data across mesh
// cluster.
StatusOr<llvm::DenseMap<int, Layout>>
DTensorSendSPMDExpander::ComputeLayoutForward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
return llvm::DenseMap<int, Layout>();
}
StatusOr<llvm::DenseMap<int, Layout>>
DTensorSendSPMDExpander::ComputeLayoutBackward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
return llvm::DenseMap<int, Layout>();
}
StatusOr<mlir::Operation*> DTensorRecvSPMDExpander::ExpandOp(
mlir::Operation* op) {
mlir::ModuleOp module = op->getParentOfType<mlir::ModuleOp>();
auto dtensor_recv = llvm::cast<mlir::TF::DTensorRecv>(op);
TF_ASSIGN_OR_RETURN(mlir::Operation * send_op,
GetCorrespondingDTensorSendRecvOp<mlir::TF::DTensorRecv>(
module, dtensor_recv));
auto dtensor_send = llvm::dyn_cast<mlir::TF::DTensorSend>(send_op);
TF_RETURN_IF_ERROR(
ValidateSendRecvLayoutConfiguration(dtensor_send, dtensor_recv));
return LowerDTensorRecv(send_op, op);
}
// DTensorRecv always returns tensors with fully replicated layout.
StatusOr<llvm::DenseMap<int, Layout>>
DTensorRecvSPMDExpander::ComputeLayoutForward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
mlir::TF::DTensorRecv dtensor_recv =
mlir::dyn_cast<mlir::TF::DTensorRecv>(op);
if (!dtensor_recv) {
return absl::InvalidArgumentError(
llvm::formatv("Expecting DTensorRecvOp but got {0}", OpName(op)).str());
}
return llvm::DenseMap<int, Layout>();
}
StatusOr<llvm::DenseMap<int, Layout>>
DTensorRecvSPMDExpander::ComputeLayoutBackward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
return llvm::DenseMap<int, Layout>();
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/expansions/dtensor_op_spmd_expander.h
+98
View File
@@ -0,0 +1,98 @@
/* Copyright 2022 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_DTENSOR_MLIR_EXPANSIONS_DTENSOR_OP_SPMD_EXPANDER_H_
#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_DTENSOR_OP_SPMD_EXPANDER_H_
#include "llvm/ADT/DenseMap.h"
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/spmd_expander.h"
namespace tensorflow {
namespace dtensor {
// Converts layout of input tensor to target layout inserting split or reduction
// ops if necessary.
class RelayoutSPMDExpander : public SPMDExpanderBase {
public:
StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& input_layouts) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& output_layouts) override;
};
// Converts layout of gradient tensor to the layout of the original Relayout's
// input tensor, using the same expansion logic as RelayoutOp.
class RelayoutLikeSPMDExpander : public SPMDExpanderBase {
public:
StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& input_layouts) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& output_layouts) override;
};
// Lowers DTensorSend op to backend specific TF send/ xla send operation.
// Following is the semantics for DTensorSend/Recv.
// a) Both replicated/sharded DTensors can be sent/received.
// b) When sharded DTensor is sent to another mesh, the DTensor is first
// all-to-all'ed to replicated tensor and sent to target mesh.
// c) Send/Recv mesh must be from or to CPU mesh. That is, TPU<->TPU or
// GPU<->GTU is not supported.
// d) Cross host send/recv is not supported. That is, sending tensor from
// TPU device of TPUWorker 0 to host of TPUWorker 1 is unsupported.
class DTensorSendSPMDExpander : public SPMDExpanderBase {
public:
StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& input_layouts) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& output_layouts) override;
};
// Lowers DTensorRecv op to backend specific TF recv/ xla recv operation.
class DTensorRecvSPMDExpander : public SPMDExpanderBase {
public:
StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& input_layouts) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& output_layouts) override;
};
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_DTENSOR_OP_SPMD_EXPANDER_H_
tensorflow/dtensor/mlir/expansions/einsum_spmd_expander.cc
+571
View File
@@ -0,0 +1,571 @@
/* Copyright 2022 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/dtensor/mlir/expansions/einsum_spmd_expander.h"
#include <cassert>
#include <cstddef>
#include <string>
#include <utility>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/status/status.h"
#include "absl/strings/str_split.h"
#include "absl/strings/string_view.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/FormatVariadic.h"
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/IRMapping.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/collectives.h"
#include "tensorflow/dtensor/mlir/layout_parsing.h"
#include "tensorflow/dtensor/mlir/shape_utils.h"
#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
#include "tensorflow/dtensor/mlir/value_utils.h"
namespace tensorflow {
namespace dtensor {
// Einsum, like reductions, is implemented as a local operation followed by
// an all-reduce over dimensions that have been reduced.
StatusOr<mlir::Operation*> EinsumSPMDExpander::ExpandOp(mlir::Operation* op) {
std::vector<Layout> input_layouts(op->getNumOperands());
for (int i = 0; i < op->getNumOperands(); ++i) {
TF_ASSIGN_OR_RETURN(auto layout,
ExtractLayoutFromOperand(op->getOperand(i)));
if (!layout)
return absl::InvalidArgumentError(
absl::StrCat("missing layout for input ", i));
input_layouts[i] = layout.value();
}
TF_ASSIGN_OR_RETURN(auto output_layout, ExtractSingleLayoutFromOp(op));
if (!output_layout)
return absl::InvalidArgumentError("is missing output layout.");
std::vector<mlir::Value> new_inputs;
Layout layout_after_einsum;
absl::flat_hash_set<std::string> reduce_dims;
TF_RETURN_IF_ERROR(MaybeRelayoutInputs(input_layouts, op,
output_layout.value(), reduce_dims,
layout_after_einsum, new_inputs));
mlir::OpBuilder builder(op);
mlir::IRMapping mapping;
for (int i = 0; i < op->getNumOperands(); ++i)
mapping.map(op->getOperand(i), new_inputs[i]);
mlir::Operation* new_op = builder.clone(*op, mapping);
// Note that the output shape of new_op is cloned from op, so we need to
// update to the local shape.
new_op = InferSPMDExpandedLocalShape(new_op);
if (!reduce_dims.empty()) {
TF_ASSIGN_OR_RETURN(
new_op, EmitAllReduce(builder, layout_after_einsum, reduce_dims, new_op,
kReduceOpAdd));
}
TF_ASSIGN_OR_RETURN(auto final_output,
EmitRelayout(new_op->getOpResult(0), layout_after_einsum,
output_layout.value()));
op->getOpResult(0).replaceAllUsesWith(final_output);
op->erase();
return final_output.getDefiningOp();
}
// TODO(power) -- we use a simplified equation parser here. consider
// refactoring einsum_spmd_expander and reusing the TF parser.
//
// Given the input equation, this has 3 outputs:
// reduced_dims: The set of mesh dimesions that we need to all reduce over.
// input_mappings: for each equation input, the map from the equation labels
// to the tensor dimension of that label.
// output_mapping: as above, but for the equation output.
absl::Status ExtractEquationRelations(
absl::string_view equation, absl::flat_hash_set<char>& reduced_dims,
std::vector<absl::flat_hash_map<char, std::vector<int>>>& input_mappings,
absl::flat_hash_map<char, std::vector<int>>& output_mapping) {
std::pair<std::string, std::string> parts = absl::StrSplit(equation, "->");
absl::flat_hash_set<char> non_reduced_dims;
// Mark kept dimensions from the output.
for (const auto& char_and_index : llvm::enumerate(parts.second)) {
// TODO(b/172691887): Support Broadcasting for einsum.
if (char_and_index.value() == '.')
return absl::UnimplementedError(absl::StrCat(
"Broadcasting is unimplemented for einsum. Received equation ",
equation));
non_reduced_dims.insert(char_and_index.value());
// Construct the output mapping, note that output is not allowed to have
// duplicate labels. This would mean that the datatype of output_mapping
// should really be absl::flat_hash_map<char, int>, but having the same
// type as the input_mapping keeps GetSpecsFromLabelsAndMap simpler.
if (output_mapping.contains(char_and_index.value()))
return errors::InvalidArgument("received label ", char_and_index.value(),
" multiple times in the "
"output of einsum equation ",
equation);
output_mapping[char_and_index.value()].emplace_back(char_and_index.index());
}
std::vector<std::string> inputs = absl::StrSplit(parts.first, ',');
// Note that the TF einsum op only supports at most 2 inputs. This is slightly
// confusing as the tf.einsum interface actually supports > 2 inputs.
if (inputs.size() > 2)
return absl::InvalidArgumentError(
absl::StrCat("einsum only supports at most 2 inputs received equation ",
equation, " which has ", inputs.size(), " inputs"));
input_mappings.resize(inputs.size());
// Compute the input mappings and keep track of labels which are reduced.
for (int i = 0; i < inputs.size(); ++i) {
for (const auto& char_and_index : llvm::enumerate(inputs[i])) {
input_mappings[i][char_and_index.value()].emplace_back(
char_and_index.index());
if (!non_reduced_dims.contains(char_and_index.value()))
reduced_dims.insert(char_and_index.value());
}
}
return absl::OkStatus();
}
// For a set of layouts and mappings from labels to offsets in the layouts,
// return a mappings of labels to ShardingSpecs.
// If the label appears multiples with different mesh dimensions in the
// sharding specs we raise an error if replicate_incompatible_dimensions is
// false. Otherwise we treat the dimension as if it were unsharded.
// Labels with unsharded dimensions are not recorded in the output.
StatusOr<absl::flat_hash_map<char, std::string>> GetLabelToShardingSpec(
bool replicate_incompatible_dimensions, const std::vector<Layout>& layouts,
const std::vector<absl::flat_hash_map<char, std::vector<int>>>& mappings) {
absl::flat_hash_map<char, std::string> label_to_sharding_spec;
absl::flat_hash_set<char> incompatible_labels;
// For each mapping, identify the mesh dimension and whether it has been
// reduced away.
for (int index = 0; index < layouts.size(); ++index) {
for (const auto& mapping : mappings[index]) {
for (int offset : mapping.second) {
if (offset >= layouts[index].rank())
return absl::InvalidArgumentError(
llvm::formatv(
"specified einsum equation for operand {0} tried to "
"read layout at offset {1}, but layout is {2} with rank "
"{3}",
index, offset, layouts[index].ToString(),
layouts[index].rank())
.str());
const std::string& sharding_spec = layouts[index].sharding_spec(offset);
if (label_to_sharding_spec.contains(mapping.first)) {
if (Layout::IsShardedDimension(sharding_spec) &&
label_to_sharding_spec[mapping.first] != sharding_spec) {
if (!replicate_incompatible_dimensions)
return absl::InvalidArgumentError(
llvm::formatv(
"incompatible mesh dimensions in equation, label '{0}' "
"is mapped to mesh dimension '{1}' and '{2}'",
mapping.first, sharding_spec,
label_to_sharding_spec[mapping.first])
.str());
else
incompatible_labels.insert(mapping.first);
}
} else if (Layout::IsShardedDimension(sharding_spec)) {
label_to_sharding_spec[mapping.first] = sharding_spec;
}
}
}
}
// For labels that had incompatible dimensions, treat them as replicated.
// We would need to insert some all to all in the SPMD expansion for these.
for (char label : incompatible_labels) label_to_sharding_spec.erase(label);
return label_to_sharding_spec;
}
// The layout we generated may be invalid as the same dimension may be used
// multiple times. E.g. ab,bc->ac (i.e. matmul) with a and c sharded over the
// same dim. In this case we mark all such dimensions as replicated.
StatusOr<Layout> VerifyOrFixLayout(
std::pair<std::vector<std::string>, absl::flat_hash_map<std::string, int>>
pair,
const Mesh& mesh) {
std::vector<std::string> sharding_specs = pair.first;
absl::flat_hash_map<std::string, int> dimension_use_count = pair.second;
for (int i = 0; i < sharding_specs.size(); ++i)
if (Layout::IsShardedDimension(sharding_specs[i]) &&
dimension_use_count[sharding_specs[i]] > 1)
sharding_specs[i] = Layout::kUnshardedDim;
return Layout::GetLayout(sharding_specs, mesh);
}
// Construct a layout on a given mesh from the label to tensor dimension map
// and the label to mesh_dimension map.
std::pair<std::vector<std::string>, absl::flat_hash_map<std::string, int>>
GetSpecsFromLabelsAndMap(
const absl::flat_hash_map<char, std::vector<int>>& label_to_index,
const absl::flat_hash_map<char, std::string>& label_to_sharding_spec) {
int layout_rank = 0;
for (const auto& label_and_indices : label_to_index)
layout_rank += label_and_indices.second.size();
std::vector<std::string> sharding_specs(layout_rank);
absl::flat_hash_map<std::string, int> dimension_use_count;
for (const auto& label_and_indices : label_to_index) {
const auto& loc = label_to_sharding_spec.find(label_and_indices.first);
if (loc != label_to_sharding_spec.end()) {
const std::string& sharding_spec = loc->second;
for (int index : label_and_indices.second)
sharding_specs[index] = sharding_spec;
dimension_use_count[sharding_spec] += label_and_indices.second.size();
} else {
for (int index : label_and_indices.second)
sharding_specs[index] = Layout::kUnshardedDim;
}
}
return std::make_pair(sharding_specs, dimension_use_count);
}
StatusOr<llvm::DenseMap<int, Layout>> EinsumSPMDExpander::ComputeLayoutForward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
if (input_layouts.empty()) return llvm::DenseMap<int, Layout>();
TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op));
// Need the mapping of input and output labels from the equation.
auto einsum_op = mlir::cast<mlir::TF::EinsumOp>(op);
size_t num_inputs = einsum_op.getNumOperands();
std::string equation = einsum_op.getEquation().str();
absl::flat_hash_set<char> reduced_dim_labels;
std::vector<absl::flat_hash_map<char, std::vector<int>>> input_mappings;
absl::flat_hash_map<char, std::vector<int>> output_mapping;
TF_RETURN_IF_ERROR(ExtractEquationRelations(equation, reduced_dim_labels,
input_mappings, output_mapping));
if (input_mappings.size() != num_inputs)
return absl::InvalidArgumentError(absl::StrCat(
"Einsum equation ", equation, " has ", input_mappings.size(),
" inputs but this op has ", num_inputs, " inputs."));
// GetLabelToShardingSpec requires two inputs if the einsum equation needs
// two inputs. We may only have one layout, so make other replicated. This
// will have the same effect as only using the defined layout and using
// replicated for all the missing dimensions.
std::vector<Layout> layouts;
for (int k = 0; k < num_inputs; ++k) {
if (input_layouts.find(k) != input_layouts.end()) {
layouts.emplace_back(input_layouts.lookup(k));
} else {
int rank = ValueRank(op->getOperand(k));
if (rank < 0)
return absl::InvalidArgumentError(
absl::StrCat("No rank for input ", k));
// This case can only happen when there are two inputs. Input 1 - k
// is the other input. In this case of the if, input k is missing, so
// this means that input 1 - k must be there.
layouts.emplace_back(Layout::ReplicatedOnMesh(mesh, rank));
}
}
// For each input, identify the mesh dimension
TF_ASSIGN_OR_RETURN(
auto input_label_to_sharding_spec,
GetLabelToShardingSpec(
/*replicate_incompatible_dimensions=*/true, layouts, input_mappings));
// Compute output layout based on retained mesh dimensions
TF_ASSIGN_OR_RETURN(
const auto& output_layout,
VerifyOrFixLayout(GetSpecsFromLabelsAndMap(output_mapping,
input_label_to_sharding_spec),
mesh));
return llvm::DenseMap<int, Layout>({{0, output_layout}});
}
StatusOr<llvm::DenseMap<int, Layout>> EinsumSPMDExpander::ComputeLayoutBackward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
if (output_layouts.find(0) == output_layouts.end())
return llvm::DenseMap<int, Layout>();
const Layout output_layout = output_layouts.lookup(0);
TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op));
// Need the mapping of input and output labels from the equation.
auto einsum_op = mlir::cast<mlir::TF::EinsumOp>(op);
size_t num_inputs = einsum_op.getNumOperands();
std::string equation = einsum_op.getEquation().str();
absl::flat_hash_set<char> reduced_dim_labels;
std::vector<absl::flat_hash_map<char, std::vector<int>>> input_mappings;
absl::flat_hash_map<char, std::vector<int>> output_mapping;
TF_RETURN_IF_ERROR(ExtractEquationRelations(equation, reduced_dim_labels,
input_mappings, output_mapping));
if (input_mappings.size() != num_inputs)
return absl::InvalidArgumentError(absl::StrCat(
"Einsum equation ", equation, " has ", input_mappings.size(),
" inputs but this op has ", num_inputs, " inputs."));
// Using the output mapping, construct an equation label to mesh dimension
// mapping.
TF_ASSIGN_OR_RETURN(auto output_label_to_sharding_spec,
GetLabelToShardingSpec(
/*replicate_incompatible_dimensions=*/false,
{output_layout}, {output_mapping}));
// Defines a set of labels that could be set to Any. The conditions for an
// operand label to be set to any are that 1) is not present in the output
// and 2) is not repeated in any operand.
absl::flat_hash_set<char> labels_for_any;
for (const auto& operand_mapping : input_mappings)
for (const auto& label_to_indices : operand_mapping)
labels_for_any.insert(label_to_indices.first);
// Filter repeated labels.
for (const auto& operand_mapping : input_mappings)
for (const auto& label_to_indices : operand_mapping)
if (label_to_indices.second.size() > 1)
labels_for_any.erase(label_to_indices.first);
// Filter labels in output.
for (const auto& label_to_indices : output_mapping)
labels_for_any.erase(label_to_indices.first);
llvm::DenseMap<int, Layout> input_layouts(num_inputs);
// Derive operand sharding specs from output's sharding specs.
for (size_t i = 0; i < num_inputs; ++i) {
absl::flat_hash_map<char, std::vector<int>> labels_to_indices =
input_mappings[i];
std::pair<std::vector<std::string>, absl::flat_hash_map<std::string, int>>
sharding_specs_and_dim_count = GetSpecsFromLabelsAndMap(
labels_to_indices, output_label_to_sharding_spec);
std::vector<std::string> sharding_specs =
sharding_specs_and_dim_count.first;
absl::flat_hash_map<std::string, int> dim_count =
sharding_specs_and_dim_count.second;
// Flip "unsharded" specs to "any" if they are present in the set.
for (const auto& label_to_indices : labels_to_indices) {
char label = label_to_indices.first;
if (labels_for_any.contains(label)) {
int index = label_to_indices.second[0];
sharding_specs[index] = Layout::kAny;
}
}
TF_ASSIGN_OR_RETURN(
const auto& layout,
VerifyOrFixLayout(std::make_pair(sharding_specs, dim_count), mesh));
input_layouts[i] = layout;
}
return input_layouts;
}
// A few things we don't support, but could:
// * "xx->" or "xx->x": Trace like operation where at least one input dimension
// for x is sharded. If both are sharded, we can compute the einsum on the
// diagonal machines in the mesh and 0s on the off diagonals and then all
// the much smaller matrix.
absl::Status EinsumSPMDExpander::MaybeRelayoutInputs(
const std::vector<Layout>& input_layouts, mlir::Operation* op,
const Layout& output_layout, absl::flat_hash_set<std::string>& reduce_dims,
Layout& einsum_layout, std::vector<mlir::Value>& new_inputs) {
if (!mlir::isa<mlir::TF::EinsumOp>(op))
return absl::InvalidArgumentError(
"called einsum spmd expander but op is not Einsum.");
mlir::TF::EinsumOp einsum = mlir::cast<mlir::TF::EinsumOp>(op);
std::vector<absl::flat_hash_map<char, std::vector<int>>> input_mappings;
absl::flat_hash_map<char, std::vector<int>> output_mapping;
absl::flat_hash_set<char> contracting_labels;
absl::flat_hash_set<char> all_labels;
TF_RETURN_IF_ERROR(ExtractEquationRelations(einsum.getEquation().str(),
contracting_labels,
input_mappings, output_mapping));
for (const auto& input_mapping : input_mappings)
for (const auto& char_and_positions : input_mapping)
all_labels.emplace(char_and_positions.first);
// We will update this array throughout this function with the following rules
// 1. The sharding of a label which is not in the map is unknown.
// 2. Once the sharding of label becomes known and is unsharded, we
// won't change that.
TF_ASSIGN_OR_RETURN(auto input_label_to_sharding_spec,
GetLabelToShardingSpec(
/*replicate_incompatible_dimensions=*/false,
input_layouts, input_mappings));
TF_ASSIGN_OR_RETURN(const auto output_label_to_sharding_spec,
GetLabelToShardingSpec(
/*replicate_incompatible_dimensions=*/false,
{output_layout}, {output_mapping}));
for (const char label : all_labels) {
if (input_label_to_sharding_spec.contains(label) &&
output_label_to_sharding_spec.contains(label) &&
input_label_to_sharding_spec[label] !=
output_label_to_sharding_spec.find(label)->second)
return errors::InvalidArgument(
"for label ", label, " input and output layouts are sharded on ",
" non-equal dimensions ", input_label_to_sharding_spec[label],
" and ", output_label_to_sharding_spec.find(label)->second,
"respectively");
}
// First priority is to ensure that labels which occur at least twice on one
// side never get sharded, as we cannot deal with that. This corresponds to
// taking a trace on that input, which will require us to be unsharded.
for (const auto& input_mapping : input_mappings)
for (const auto& char_and_positions : input_mapping)
if (char_and_positions.second.size() > 1)
input_label_to_sharding_spec[char_and_positions.first] =
Layout::kUnshardedDim;
absl::flat_hash_map<std::string, absl::flat_hash_set<char>>
sharding_dim_to_non_contracting_labels;
absl::flat_hash_map<std::string, absl::flat_hash_set<char>>
sharding_dim_to_contracting_labels;
for (const auto& label_and_spec : input_label_to_sharding_spec) {
if (Layout::IsShardedDimension(label_and_spec.second)) {
if (contracting_labels.contains(label_and_spec.first))
sharding_dim_to_contracting_labels[label_and_spec.second].insert(
label_and_spec.first);
else
sharding_dim_to_non_contracting_labels[label_and_spec.second].insert(
label_and_spec.first);
}
}
// If a non-contracting dimension is sharded in the output and non-sharded
// in the input and no other label is sharded on that dimension, then shard
// it.
// This handles the *,x . x,* -> *,y case and also if batch dimensions are
// sharded on the output but not the input.
for (const char label : all_labels) {
// Note that only sharded labels are in output_label_to_sharding_spec, so
// there is no need to check that the spec is sharded.
if (!contracting_labels.contains(label) &&
output_label_to_sharding_spec.contains(label) &&
!input_label_to_sharding_spec.contains(label)) {
const std::string& string_spec =
output_label_to_sharding_spec.find(label)->second;
if (!sharding_dim_to_non_contracting_labels.contains(string_spec) &&
!sharding_dim_to_contracting_labels.contains(string_spec)) {
input_label_to_sharding_spec[label] = string_spec;
sharding_dim_to_non_contracting_labels[string_spec].insert(label);
}
}
}
// Handle the case when two non-contracting dimensions are have the same
// sharding spec.
// Note that the case of three non-contracting dimensions having the same
// sharding spec is impossible. Since there are at most two inputs, at least
// one input would have two dimensions with the same sharing spec.
// This handles the y,x . x,y -> *,y case.
absl::flat_hash_set<std::string> dims_with_multiple_labels;
for (const auto& spec_and_labels : sharding_dim_to_non_contracting_labels) {
if (spec_and_labels.second.size() > 1) {
assert(spec_and_labels.second.size() == 2);
dims_with_multiple_labels.insert(spec_and_labels.first);
}
}
for (const auto& dim : dims_with_multiple_labels) {
// TODO(bfontain): Update this to pick default label to keep based on shape.
char label_to_keep = 0xFF;
// Note that all these conditions evaluated in the loop below are mutually
// as exclusive as no two labels in the output have the same sharding
// spec.
// If the no label is found we choose the lexicographically least label to
// keep this stable with respect to ordering.
for (const char label : sharding_dim_to_non_contracting_labels[dim]) {
if (output_label_to_sharding_spec.contains(label) &&
output_label_to_sharding_spec.find(label)->second == dim) {
label_to_keep = label;
break;
} else if (label < label_to_keep) {
label_to_keep = label;
}
}
for (const char label : sharding_dim_to_non_contracting_labels[dim])
if (label != label_to_keep)
input_label_to_sharding_spec[label] = Layout::kUnshardedDim;
sharding_dim_to_non_contracting_labels[dim].clear();
sharding_dim_to_non_contracting_labels[dim].insert(label_to_keep);
}
// Handle the case where a non-contracting and contracting dim have the same
// sharding spec. For now we always unshard the contracting axis. Note that
// this is safe.
// This handles the case x,y . *,y -> x,y
for (const auto& spec_and_labels : sharding_dim_to_contracting_labels) {
if (!spec_and_labels.second.empty() &&
!sharding_dim_to_non_contracting_labels[spec_and_labels.first]
.empty()) {
assert(spec_and_labels.second.size() == 1);
assert(sharding_dim_to_non_contracting_labels[spec_and_labels.first]
.size() == 1);
input_label_to_sharding_spec[*spec_and_labels.second.begin()] =
Layout::kUnshardedDim;
}
}
// Relayout the inputs
mlir::OpBuilder builder(op);
new_inputs.resize(input_mappings.size());
for (int i = 0; i < input_mappings.size(); ++i) {
TF_ASSIGN_OR_RETURN(
const Layout new_input_layout,
VerifyOrFixLayout(GetSpecsFromLabelsAndMap(
input_mappings[i], input_label_to_sharding_spec),
output_layout.mesh()));
TF_ASSIGN_OR_RETURN(
new_inputs[i],
EmitRelayout(op->getOperand(i), input_layouts[i], new_input_layout));
}
TF_ASSIGN_OR_RETURN(
einsum_layout,
VerifyOrFixLayout(GetSpecsFromLabelsAndMap(output_mapping,
input_label_to_sharding_spec),
output_layout.mesh()));
for (const auto& contracting : contracting_labels)
reduce_dims.emplace(input_label_to_sharding_spec[contracting]);
return absl::OkStatus();
}
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/mlir/expansions/einsum_spmd_expander.h
+67
View File
@@ -0,0 +1,67 @@
/* Copyright 2022 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_DTENSOR_MLIR_EXPANSIONS_EINSUM_SPMD_EXPANDER_H_
#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_EINSUM_SPMD_EXPANDER_H_
#include <string>
#include <vector>
#include "absl/container/flat_hash_set.h"
#include "absl/status/status.h"
#include "llvm/ADT/DenseMap.h"
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "tensorflow/core/platform/status.h"
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/spmd_expander.h"
namespace tensorflow {
namespace dtensor {
class EinsumSPMDExpander : public SPMDExpanderBase {
public:
StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& input_layouts) override;
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& output_layouts) override;
private:
// This function will take the inputs, possibly slice or add AllConcat along
// various mesh dimensions before the einsum operation takes place. This also
// returns:
// * The mesh dimensions, if any, that the output of the einsum should be
// summed along.
// * The resulting output layout of the einsum operation, so we can insert an
// AllConcat/split to make the output have the desired layout.
// * The new inputs to fed into the einsum.
absl::Status MaybeRelayoutInputs(
const std::vector<Layout>& input_layouts, mlir::Operation* op,
const Layout& output_layout,
absl::flat_hash_set<std::string>& reduce_dims, Layout& einsum_layout,
std::vector<mlir::Value>& new_inputs);
};
} // namespace dtensor
} // namespace tensorflow
#endif // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_EINSUM_SPMD_EXPANDER_H_
tensorflow/dtensor/mlir/expansions/elementwise_spmd_expander.cc
+184
View File
@@ -0,0 +1,184 @@
/* Copyright 2022 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/dtensor/mlir/expansions/elementwise_spmd_expander.h"
#include <cassert>
#include <cstdint>
#include <optional>
#include "absl/container/flat_hash_set.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/SmallVector.h"
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/UseDefLists.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/mlir/collectives.h"
#include "tensorflow/dtensor/mlir/layout_parsing.h"
#include "tensorflow/dtensor/mlir/shape_utils.h"
#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
#include "tensorflow/dtensor/mlir/value_utils.h"
namespace tensorflow {
namespace dtensor {
namespace {
StatusOr<llvm::SmallVector<int64_t, 4>> GetShape(mlir::Value value) {
TF_ASSIGN_OR_RETURN(const auto shape, GetShapeOfValue(value));
return llvm::SmallVector<int64_t, 4>{shape.begin(), shape.end()};
}
} // namespace
StatusOr<mlir::Operation*> ElementwiseSPMDExpander::ExpandOp(
mlir::Operation* op) {
TF_ASSIGN_OR_RETURN(auto output_layout, ExtractSingleLayoutFromOp(op));
assert(output_layout);
mlir::OpBuilder builder(op);
for (auto& operand : op->getOpOperands()) {
// Verify that all output dimensions (including the dimensions added by
// broadcasting) is more sharded then the correspdonding layout
// configuration of the same dimension of every operands.
TF_ASSIGN_OR_RETURN(auto operand_layout,
ExtractLayoutFromOperand(operand.get()));
if (!operand_layout)
return absl::InvalidArgumentError(
"input layout of elementwise op must be known before SPMD "
"expansion.");
// For scalar operands, splitting is not needed.
if (operand_layout->rank() == 0) continue;
llvm::SmallPtrSet<mlir::Operation*, 4> newly_created_ops;
// Note that due to broacasting, inputs and output tensors are aligned to
// the right. Therefore, we truncate the output layout.
const int rank_offset = output_layout->rank() - operand_layout->rank();
// Get the desired layout for this operand.
//
// - Truncate: It should be the same layout as the output, truncated to
// take into account broadcasting on the rank and removing sharding in
// dimensions where the operand has size 1 (where we have broadcasting on
// the dimension size).
//
// - Relayout: If the output and operand have different sharding spec, we
// adjust the operands. For example, if operand is 'z,*' and output is
// '*.y', relayout operand to conform output. This means the SPMD safer
// and easeier. In future, we might do certain optimization to save FLops.
// For example, if all operands are 'x,y' and output is '*,*', relayouting
// output could be the choice (saving communications).
auto truncated_layout = output_layout->Truncate(rank_offset, /*end=*/true);
mlir::Value output;
TF_ASSIGN_OR_RETURN(const auto& shape, ExtractGlobalInputShape(operand));
absl::flat_hash_set<int> size_one_dims;
for (int i = 0; i < shape.size(); ++i)
if (shape[i] == 1) size_one_dims.emplace(i);
TF_ASSIGN_OR_RETURN(truncated_layout,
truncated_layout.GetLayoutWithReducedDims(
size_one_dims, /*keep_dims=*/true));
TF_ASSIGN_OR_RETURN(
output, EmitRelayout(operand.get(), *operand_layout, truncated_layout));
operand.set(output);
}
// If result is a resource, the shape of the result should be adjusted to
// local value of the resource, based on the layout for output.
// This logic is similar to VarHandle op SPMD expansion.
//
// Resource output is only likely to be for identity op. However, keeping
// the checkgeneric here.
auto op_result = op->getOpResult(0);
if (IsResourceType(op_result)) {
TF_RETURN_IF_ERROR(InferSPMDExpandedLocalShapeForResourceOutput(
&op_result, output_layout.value(), builder.getContext()));
}
// For element-wise op SPMD expansion, given that operand layouts are
// compatible to op's layout, op can simply be executed without any changes.
return InferSPMDExpandedLocalShape(op);
}
// Computes output layouts of elementwise operation using broadcast logic for
// operands.
StatusOr<llvm::DenseMap<int, Layout>>
ElementwiseSPMDExpander::ComputeLayoutForward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
TF_ASSIGN_OR_RETURN(std::optional<Layout> merged_operand_layout,
GetMergedOperandLayout(input_layouts, op));
if (merged_operand_layout) {
const int output_rank = ValueRank(op->getOpResult(0));
if (output_rank == -1)
return absl::InvalidArgumentError("Output has unknown rank");
// We assume that all elementwise operations output a single tensor.
return llvm::DenseMap<int, Layout>(
{{0, merged_operand_layout->LeftPad(output_rank)}});
}
return llvm::DenseMap<int, Layout>();
}
// Computes input layouts of elementwise operation using broadcast logic for
// operands.
StatusOr<llvm::DenseMap<int, Layout>>
ElementwiseSPMDExpander::ComputeLayoutBackward(
mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
// Do not infer any operand layout if no output layout is present.
if (output_layouts.find(0) == output_layouts.end())
return llvm::DenseMap<int, Layout>();
Layout output_layout = output_layouts.lookup(0);
llvm::DenseMap<int, Layout> input_layouts;
for (const auto& operand_and_index : llvm::enumerate(op->getOperands())) {
const int operand_index = operand_and_index.index();
auto operand = operand_and_index.value();
TF_ASSIGN_OR_RETURN(auto operand_shape, GetShape(operand));
Layout output_layout_truncated = output_layout.Truncate(
output_layout.sharding_spec_strs().size() - operand_shape.size(),
/*end=*/true);
auto inferred_operand_layout_strs =
output_layout_truncated.sharding_spec_strs();
if (inferred_operand_layout_strs.size() != operand_shape.size())
return absl::FailedPreconditionError(
"Mismatch of operand shape size and layout size.");
for (const auto& dim_shape_and_index : llvm::enumerate(operand_shape)) {
const int dim_index = dim_shape_and_index.index();
const int dim_shape = dim_shape_and_index.value();
if (dim_shape <= 1) {
inferred_operand_layout_strs[dim_index] = Layout::kUnshardedDim;
}
}
TF_ASSIGN_OR_RETURN(
auto inferred_operand_layout,
Layout::GetLayout(inferred_operand_layout_strs, output_layout.mesh()));
input_layouts[operand_index] = inferred_operand_layout;
}
return input_layouts;
}
} // namespace dtensor
} // namespace tensorflow

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