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