wehub-resource-sync
1898 lines
78 KiB
C++
1898 lines
78 KiB
C++
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
==============================================================================*/
|
|
|
|
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
|
|
|
|
#include <algorithm>
|
|
#include <array>
|
|
#include <cstdint>
|
|
#include <map>
|
|
#include <memory>
|
|
#include <numeric>
|
|
#include <optional>
|
|
#include <set>
|
|
#include <string>
|
|
#include <unordered_map>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
#include "tensorflow/compiler/mlir/tf2xla/mlir_bridge_rollout_policy.h"
|
|
#include "absl/algorithm/container.h"
|
|
#include "absl/container/flat_hash_map.h"
|
|
#include "absl/memory/memory.h"
|
|
#include "absl/status/status.h"
|
|
#include "absl/strings/string_view.h"
|
|
#include "absl/synchronization/mutex.h"
|
|
#include "absl/types/variant.h"
|
|
#include "tensorflow/compiler/jit/defs.h"
|
|
#include "tensorflow/compiler/jit/flags.h"
|
|
#include "tensorflow/compiler/jit/shape_inference.h"
|
|
#include "tensorflow/compiler/jit/xla_compile_util.h"
|
|
#include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h"
|
|
#include "tensorflow/compiler/mlir/tf2xla/api/v1/compile_mlir_util.h"
|
|
#include "tensorflow/compiler/mlir/utils/array_container_utils.h"
|
|
#include "tensorflow/compiler/tf2xla/graph_compiler.h"
|
|
#include "tensorflow/compiler/tf2xla/layout_util.h"
|
|
#include "tensorflow/compiler/tf2xla/rearrange_function_argument.h"
|
|
#include "tensorflow/compiler/tf2xla/shape_util.h"
|
|
#include "tensorflow/compiler/tf2xla/sharding_util.h"
|
|
#include "tensorflow/compiler/tf2xla/side_effect_util.h"
|
|
#include "tensorflow/compiler/tf2xla/tf2xla_util.h"
|
|
#include "tensorflow/compiler/tf2xla/type_util.h"
|
|
#include "tensorflow/compiler/tf2xla/xla_compilation_device.h"
|
|
#include "tensorflow/compiler/tf2xla/xla_context.h"
|
|
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
|
|
#include "xla/client/client_library.h"
|
|
#include "xla/hlo/builder/xla_builder.h"
|
|
#include "xla/hlo/builder/xla_computation.h"
|
|
#include "xla/protobuf_util.h"
|
|
#include "xla/service/hlo.pb.h"
|
|
#include "xla/service/spmd/shardy/constants.h"
|
|
#include "xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_import.h"
|
|
#include "xla/service/spmd/shardy/utils.h"
|
|
#include "xla/shape.h"
|
|
#include "xla/shape_util.h"
|
|
#include "xla/tsl/platform/errors.h"
|
|
#include "xla/tsl/platform/statusor.h"
|
|
#include "xla/util.h"
|
|
#include "xla/xla_data.pb.h"
|
|
#include "tensorflow/core/common_runtime/device.h"
|
|
#include "tensorflow/core/common_runtime/executor.h"
|
|
#include "tensorflow/core/common_runtime/function.h"
|
|
#include "tensorflow/core/common_runtime/function_def_utils.h"
|
|
#include "tensorflow/core/common_runtime/graph_constructor.h"
|
|
#include "tensorflow/core/common_runtime/graph_optimizer.h"
|
|
#include "tensorflow/core/framework/attr_value_util.h"
|
|
#include "tensorflow/core/framework/function.h"
|
|
#include "tensorflow/core/framework/graph_debug_info.pb.h"
|
|
#include "tensorflow/core/framework/metrics.h"
|
|
#include "tensorflow/core/framework/node_def_util.h"
|
|
#include "tensorflow/core/framework/types.h"
|
|
#include "tensorflow/core/graph/node_builder.h"
|
|
#include "tensorflow/core/lib/core/errors.h"
|
|
#include "tensorflow/core/lib/gtl/cleanup.h"
|
|
#include "tensorflow/core/lib/hash/hash.h"
|
|
#include "tensorflow/core/platform/logging.h"
|
|
#include "tensorflow/core/platform/status.h"
|
|
#include "tensorflow/core/protobuf/error_codes.pb.h"
|
|
#include "tensorflow/core/tpu/tpu_defs.h"
|
|
#include "tensorflow/core/util/debug_data_dumper.h"
|
|
#include "tensorflow/core/util/dump_graph.h"
|
|
#include "tsl/platform/errors.h"
|
|
#include "tsl/platform/tensor_float_32_utils.h"
|
|
|
|
namespace tensorflow {
|
|
namespace {
|
|
|
|
// Name of component for error logging. This name is fixed and required to
|
|
// enable logging.
|
|
constexpr char kSingleOpComponent[] = "TF2XLA_XLA_COMPILER_COMPILE_SINGLE_OP";
|
|
constexpr char kCompileFunctionComponent[] =
|
|
"TF2XLA_XLA_COMPILER_COMPILE_FUNCTION";
|
|
|
|
// Checks that arguments `args` match types `types`.
|
|
absl::Status CheckSignature(const DataTypeVector& types,
|
|
absl::Span<const XlaCompiler::Argument> args) {
|
|
if (args.size() != types.size()) {
|
|
return absl::InternalError(
|
|
absl::StrCat("Compilation arguments have ", args.size(),
|
|
" elements while function has ", types.size()));
|
|
}
|
|
for (int i = 0, end = types.size(); i < end; ++i) {
|
|
// Don't perform type checks on resource variables and tensor
|
|
// lists (DT_VARIANT) as we have to trick the type system in order to
|
|
// plumb them through. DT_VARIANTS are wrapped in a DT_UINT8 tensor.
|
|
if (types[i] != args[i].type && types[i] != DT_RESOURCE &&
|
|
types[i] != DT_VARIANT) {
|
|
return absl::InternalError(absl::StrCat(
|
|
"Argument ", i, " has declared type ", DataTypeString(args[i].type),
|
|
" but function parameter has type ", DataTypeString(types[i])));
|
|
}
|
|
}
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
// Uses the _Arg and _Retval nodes in the graph to determine an OpSharding for
|
|
// each argument and return value.
|
|
absl::StatusOr<
|
|
std::pair<std::map<int, xla::OpSharding>, std::map<int, xla::OpSharding>>>
|
|
ComputeArgAndRetvalShardings(const Graph& graph) {
|
|
auto get_sharding_for_node =
|
|
[](const Node* n) -> absl::StatusOr<std::optional<xla::OpSharding>> {
|
|
TF_ASSIGN_OR_RETURN(
|
|
auto sharding,
|
|
ParseShardingFromDevice(*n, std::numeric_limits<int32_t>::max(),
|
|
/*add_metadata=*/false));
|
|
return sharding;
|
|
};
|
|
std::map<int, xla::OpSharding> arg_shardings;
|
|
std::map<int, xla::OpSharding> retval_shardings;
|
|
for (const Node* n : graph.nodes()) {
|
|
if (n->IsArg()) {
|
|
TF_ASSIGN_OR_RETURN(auto sharding, get_sharding_for_node(n));
|
|
if (!sharding.has_value()) continue;
|
|
int index;
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), "index", &index));
|
|
TF_RET_CHECK(index >= 0) << "Negative _Arg index";
|
|
arg_shardings[index] = std::move(*sharding);
|
|
} else if (n->IsRetval()) {
|
|
TF_ASSIGN_OR_RETURN(auto sharding, get_sharding_for_node(n));
|
|
if (!sharding.has_value()) continue;
|
|
int index;
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), "index", &index));
|
|
TF_RET_CHECK(index >= 0) << "Negative _Retval index";
|
|
retval_shardings[index] = std::move(*sharding);
|
|
}
|
|
}
|
|
return std::make_pair(std::move(arg_shardings), std::move(retval_shardings));
|
|
}
|
|
|
|
// Due to the wonkiness with Resource Cleanup, changing how resources are
|
|
// cleaned up here need to change how resources are cleaned up in
|
|
// graph_compiler_test.
|
|
// LINT.IfChange(ExecuteGraph)
|
|
absl::Status ExecuteGraph(XlaContext* xla_context, std::unique_ptr<Graph> graph,
|
|
XlaCompilationDevice* device,
|
|
FunctionLibraryRuntime* flib, int64_t step_id) {
|
|
// Resource cleanup is a bit messy. XlaContext is a ref-countd resource; the
|
|
// resource manager takes ownership via Create, and unrefs via Cleanup. We
|
|
// explicitly add a reference to ensure the refcount at entry is maintained at
|
|
// all exit points; Create and Cleanup are always called in this function.
|
|
//
|
|
// The Executor requires us to use ScopedStepContainer. We wrap it in a
|
|
// unique_ptr so we can capture the cleanup status in the end.
|
|
xla_context->Ref();
|
|
absl::Status status;
|
|
auto step_container = std::make_unique<ScopedStepContainer>(
|
|
step_id, [&status, device](const std::string& name) {
|
|
status = device->resource_manager()->Cleanup(name);
|
|
});
|
|
TF_RETURN_IF_ERROR(step_container->Create(device->resource_manager(),
|
|
XlaContext::kXlaContextResourceName,
|
|
xla_context));
|
|
|
|
GraphCompiler graph_compiler(device, graph.get(), flib, step_container.get());
|
|
TF_RETURN_IF_ERROR(graph_compiler.Compile());
|
|
// Explicitly clean up the step container, to capture the cleanup status.
|
|
step_container.reset();
|
|
return status;
|
|
}
|
|
// LINT.ThenChange(//tensorflow/compiler/tf2xla/graph_compiler_test.cc)
|
|
|
|
// Builds the XLA computation.
|
|
// - `args` is the list of input arguments
|
|
// - `retvals` is the list of retvals produced by _Retval operators, in index
|
|
// order.
|
|
// - `arg_shardings` and `retval_shardings` are mapping from arg/return indices
|
|
// to sharding.
|
|
// - If `return_updated_values_for_all_resources` is true, all resources will be
|
|
// included in `resource_updates`, regardless of whether their value changed.
|
|
// - Sets `*num_nonconst_outputs` to the number of outputs of the `computation`.
|
|
// - Sets `*resource_updates` to a description of resources whose values are
|
|
// written by the computation; the variable writes are the last
|
|
// - `resource_updates.size()` return values from the computation. Each entry in
|
|
// `resource_updates` is a ResourceUpdate, whose `index` is the index of a
|
|
// resource variable argument to the computation to be updated, and `type` is
|
|
// the type of the final output.
|
|
absl::Status BuildComputation(
|
|
const std::vector<XlaCompiler::Argument>& args,
|
|
const std::vector<XlaExpression>& retvals,
|
|
const std::map<int, xla::OpSharding>& arg_shardings,
|
|
const std::map<int, xla::OpSharding>& retval_shardings,
|
|
const std::vector<std::unique_ptr<XlaResource>>& resources,
|
|
std::unique_ptr<xla::XlaOp> token_output,
|
|
const XlaShapeLayoutHelpers::ShapeDeterminationFns& shape_determination_fns,
|
|
bool is_entry_computation, bool return_updated_values_for_all_resources,
|
|
bool always_return_tuple, bool use_tuple_arg, bool alias_resource_update,
|
|
xla::XlaBuilder* builder, xla::XlaComputation* computation,
|
|
int* num_computation_outputs, int* num_nonconst_outputs,
|
|
std::vector<XlaCompiler::OutputDescription>* outputs,
|
|
std::vector<XlaCompiler::ResourceUpdate>* resource_updates,
|
|
xla::Shape* output_shape, absl::Span<int const> input_mapping,
|
|
std::string& result_tuple_sdy_sharding, bool use_shardy_partitioner) {
|
|
// Attach a common operator name as metadata. This has no semantic effect — it
|
|
// merely makes the HLO graph more readable when visualized via TensorBoard,
|
|
// since TensorBoard forms groups out of operators with similar names.
|
|
xla::OpMetadata retval_metadata;
|
|
retval_metadata.set_op_name("XLA_Retvals");
|
|
builder->SetOpMetadata(retval_metadata);
|
|
VLOG(1) << "Building new computation";
|
|
auto cleanup = gtl::MakeCleanup([builder]() { builder->ClearOpMetadata(); });
|
|
|
|
// Builds a no-op XLA computation. We need to set the sharding of outputs, but
|
|
// cannot change the sharding of the existing output op. To do this, we build
|
|
// a new identity op to which shardings can be applied.
|
|
auto identity_op = [builder, is_entry_computation, use_shardy_partitioner](
|
|
xla::XlaOp op,
|
|
const std::optional<xla::OpSharding>& sharding)
|
|
-> absl::StatusOr<xla::XlaOp> {
|
|
xla::XlaScopedShardingAssignment assign_sharding(builder, sharding);
|
|
xla::XlaOp copy_op = xla::Copy(op);
|
|
if (is_entry_computation && use_shardy_partitioner) {
|
|
// Add shardy shardings only for entry computation, as we only need to add
|
|
// it for wrapper main added in tensorflow.
|
|
TF_RETURN_IF_ERROR(addSdyShardingFrontendAttribute(
|
|
builder, copy_op, builder->GetShape(copy_op).value()));
|
|
}
|
|
return copy_op;
|
|
};
|
|
|
|
std::vector<xla::XlaOp> elems;
|
|
elems.reserve(retvals.size());
|
|
|
|
// Keeps track of sharding of each retval. If a retval is not in this list,
|
|
// replicate sharding is used. The first element is the output index, second
|
|
// element is the sharding.
|
|
std::unordered_map<int, xla::OpSharding> retval_index_and_sharding;
|
|
for (int i = 0, end = retvals.size(); i < end; ++i) {
|
|
XlaCompiler::OutputDescription& output = (*outputs)[i];
|
|
const XlaExpression& retval = retvals[i];
|
|
output.type = retval.dtype();
|
|
switch (retval.kind()) {
|
|
case XlaExpression::Kind::kConstant:
|
|
output.is_constant = true;
|
|
output.constant_value = *retval.constant_value();
|
|
output.shape = output.constant_value.shape();
|
|
break;
|
|
|
|
case XlaExpression::Kind::kTensorList: {
|
|
output.is_tensor_list = true;
|
|
xla::XlaOp value = retval.handle();
|
|
elems.push_back(value);
|
|
break;
|
|
}
|
|
|
|
case XlaExpression::Kind::kXlaOp: {
|
|
output.is_constant = false;
|
|
TF_ASSIGN_OR_RETURN(output.shape, retval.GetShape());
|
|
xla::XlaOp value = retval.handle();
|
|
auto it = retval_shardings.find(i);
|
|
std::optional<xla::OpSharding> sharding =
|
|
it == retval_shardings.end() ? std::optional<xla::OpSharding>()
|
|
: it->second;
|
|
if (it != retval_shardings.end()) {
|
|
retval_index_and_sharding[elems.size()] = it->second;
|
|
}
|
|
if (shape_determination_fns.shape_representation_fn) {
|
|
TF_ASSIGN_OR_RETURN(auto original_shape, builder->GetShape(value));
|
|
TF_ASSIGN_OR_RETURN(value,
|
|
ReshapeWithCorrectRepresentationAndSharding(
|
|
builder, value, original_shape,
|
|
shape_determination_fns, sharding,
|
|
/*fast_mem=*/false));
|
|
}
|
|
if (it != retval_shardings.end()) {
|
|
// Apply the sharding to the output, if there is a core assignment.
|
|
TF_ASSIGN_OR_RETURN(value, identity_op(value, sharding));
|
|
}
|
|
|
|
elems.push_back(value);
|
|
break;
|
|
}
|
|
|
|
case XlaExpression::Kind::kResource:
|
|
// Resources will be pushed into elems later when processing resource
|
|
// arguments below.
|
|
output.is_constant = false;
|
|
output.input_index = retval.resource()->arg_num();
|
|
output.shape = retval.resource()->shape();
|
|
break;
|
|
|
|
case XlaExpression::Kind::kInvalid:
|
|
return absl::InvalidArgumentError(
|
|
"Invalid expression returned by computation. "
|
|
"This probably means a return value was not set.");
|
|
}
|
|
}
|
|
*num_nonconst_outputs = elems.size();
|
|
|
|
// Add return values for resources whose values have changed.
|
|
std::vector<const XlaResource*> arg_resources;
|
|
arg_resources.reserve(resources.size());
|
|
for (const auto& resource : resources) {
|
|
if (resource->arg_num() >= 0) {
|
|
arg_resources.push_back(resource.get());
|
|
}
|
|
}
|
|
std::sort(arg_resources.begin(), arg_resources.end(),
|
|
[](const XlaResource* a, const XlaResource* b) {
|
|
return a->arg_num() < b->arg_num();
|
|
});
|
|
|
|
absl::flat_hash_map<int, int> argument_to_xla_arg;
|
|
for (int xla_arg = 0; xla_arg < input_mapping.size(); xla_arg++) {
|
|
argument_to_xla_arg[input_mapping[xla_arg]] = xla_arg;
|
|
}
|
|
|
|
std::vector<xla::XlaBuilder::InputOutputAlias> aliases;
|
|
for (const XlaResource* resource : arg_resources) {
|
|
DCHECK_LT(resource->arg_num(), args.size());
|
|
const XlaCompiler::Argument& arg = args[resource->arg_num()];
|
|
auto it = arg_shardings.find(resource->arg_num());
|
|
bool modified = !resource->value().IsIdenticalTo(resource->initial_value());
|
|
// TensorArray gradients were modified if their values changed or there are
|
|
// any newly created gradients.
|
|
for (const auto& grad : resource->tensor_array_gradients()) {
|
|
modified =
|
|
modified ||
|
|
!grad.second->value().IsIdenticalTo(grad.second->initial_value()) ||
|
|
arg.tensor_array_gradients.count(grad.first) == 0;
|
|
}
|
|
|
|
if (return_updated_values_for_all_resources || modified ||
|
|
arg.requires_broadcast) {
|
|
resource_updates->emplace_back();
|
|
XlaCompiler::ResourceUpdate& update = resource_updates->back();
|
|
update.input_index = resource->arg_num();
|
|
update.type = resource->type();
|
|
update.shape = resource->shape();
|
|
update.modified = modified;
|
|
int param_num = use_tuple_arg ? 0 : update.input_index;
|
|
if (is_entry_computation &&
|
|
arg.resource_kind != XlaResource::kTensorArray &&
|
|
alias_resource_update && argument_to_xla_arg.count(param_num)) {
|
|
// Assuming tuple arg and results are used.
|
|
xla::ShapeIndex param_index =
|
|
use_tuple_arg ? xla::ShapeIndex({update.input_index})
|
|
: xla::ShapeIndex{};
|
|
int xla_param_num = argument_to_xla_arg[param_num];
|
|
int64_t output_index_num = elems.size();
|
|
xla::ShapeIndex output_index = xla::ShapeIndex({output_index_num});
|
|
VLOG(3) << "Storing alias: " << output_index.ToString() << ": ("
|
|
<< xla_param_num << ", " << param_index.ToString() << ")";
|
|
aliases.push_back({output_index, xla_param_num, param_index});
|
|
}
|
|
for (const auto& grad : resource->tensor_array_gradients()) {
|
|
update.tensor_array_gradients_accessed.insert(grad.first);
|
|
}
|
|
|
|
xla::XlaOp handle;
|
|
TF_RETURN_IF_ERROR(resource->Pack(&handle, builder));
|
|
auto sharding = it == arg_shardings.end()
|
|
? std::optional<xla::OpSharding>()
|
|
: it->second;
|
|
// Set layout of the retval to device representation layout.
|
|
if (shape_determination_fns.layout_preference_fn &&
|
|
shape_determination_fns.shape_representation_fn) {
|
|
TF_ASSIGN_OR_RETURN(auto original_shape, builder->GetShape(handle));
|
|
TF_ASSIGN_OR_RETURN(
|
|
handle, ReshapeWithCorrectRepresentationAndSharding(
|
|
builder, handle, original_shape,
|
|
shape_determination_fns, sharding, arg.fast_mem));
|
|
}
|
|
|
|
// Request that the value be returned on a specific core.
|
|
if (it != arg_shardings.end()) {
|
|
retval_index_and_sharding[elems.size()] = it->second;
|
|
}
|
|
// Ensures the correct sharding is applied to the output.
|
|
TF_ASSIGN_OR_RETURN(handle, identity_op(handle, sharding));
|
|
elems.push_back(handle);
|
|
}
|
|
}
|
|
|
|
// If we have token output, append it as the last one.
|
|
if (token_output) {
|
|
elems.push_back(*token_output);
|
|
}
|
|
|
|
*num_computation_outputs = elems.size();
|
|
|
|
// Builds the XLA computation. We *always* form a tuple here to ensure that
|
|
// the output value is the last thing added into the XLA computation, even
|
|
// if there is only one output value.
|
|
xla::XlaOp tuple;
|
|
if (retval_index_and_sharding.empty() || !is_entry_computation) {
|
|
tuple = xla::Tuple(builder, elems);
|
|
} else {
|
|
std::vector<xla::Shape> elem_shapes;
|
|
for (const auto& elem : elems) {
|
|
TF_ASSIGN_OR_RETURN(xla::Shape elem_shape,
|
|
elem.builder()->GetShape(elem));
|
|
elem_shapes.push_back(elem_shape);
|
|
}
|
|
xla::Shape shape = xla::ShapeUtil::MakeTupleShape(elem_shapes);
|
|
// Copy specified sharding from retval_index_and_sharding.
|
|
std::vector<xla::HloSharding> sharding_elems;
|
|
for (int i = 0, end = elems.size(); i < end; i++) {
|
|
const auto& iter = retval_index_and_sharding.find(i);
|
|
TF_RET_CHECK(iter != retval_index_and_sharding.end());
|
|
const xla::OpSharding& sub_op_sharding = iter->second;
|
|
TF_ASSIGN_OR_RETURN(xla::HloSharding sub_sharding,
|
|
xla::HloSharding::FromProto(sub_op_sharding));
|
|
if (elem_shapes[i].IsTuple()) {
|
|
const std::vector<xla::HloSharding> sub_sharding_elems =
|
|
sub_sharding.tuple_elements();
|
|
const int64_t sub_sharding_elems_size = sub_sharding_elems.size();
|
|
TF_RET_CHECK(sub_sharding_elems_size ==
|
|
xla::ShapeUtil::GetLeafCount(elem_shapes[i]));
|
|
for (const auto& sub_sharding_elem : sub_sharding_elems) {
|
|
sharding_elems.push_back(sub_sharding_elem);
|
|
}
|
|
} else {
|
|
sharding_elems.push_back(sub_sharding);
|
|
}
|
|
}
|
|
xla::HloSharding modified_sharding =
|
|
xla::HloSharding::Tuple(shape, sharding_elems);
|
|
xla::OpSharding op_sharding = modified_sharding.ToProto();
|
|
// Assign proper sharding to the tuple instruction.
|
|
xla::XlaScopedShardingAssignment assign_sharding(builder, op_sharding);
|
|
tuple = xla::Tuple(builder, elems);
|
|
if (use_shardy_partitioner && builder->sharding().has_value()) {
|
|
result_tuple_sdy_sharding =
|
|
xla::sdy::convertToSdySharding(op_sharding, shape,
|
|
/*openDims=*/false,
|
|
/*inlineMesh=*/true);
|
|
}
|
|
}
|
|
bool returns_tuple = always_return_tuple || elems.size() != 1;
|
|
VLOG(3) << "Computation returns a tuple=" << returns_tuple;
|
|
if (!returns_tuple) {
|
|
xla::GetTupleElement(tuple, 0);
|
|
|
|
for (xla::XlaBuilder::InputOutputAlias& alias : aliases) {
|
|
if (alias.output_index == xla::ShapeIndex({0})) {
|
|
VLOG(3) << "For aliased parameter " << alias.param_number << ": "
|
|
<< alias.param_index.ToString()
|
|
<< " normalizing output_index from {0} to {}, as a scalar is "
|
|
"returned from the cluster";
|
|
alias.output_index = xla::ShapeIndex({});
|
|
}
|
|
}
|
|
}
|
|
|
|
for (xla::XlaBuilder::InputOutputAlias& alias : aliases) {
|
|
builder->SetUpAlias(alias.output_index, alias.param_number,
|
|
alias.param_index);
|
|
}
|
|
TF_ASSIGN_OR_RETURN(*computation, builder->Build());
|
|
|
|
TF_ASSIGN_OR_RETURN(auto program_shape, computation->GetProgramShape());
|
|
*output_shape = program_shape.result();
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
} // namespace
|
|
|
|
std::string XlaCompiler::Argument::HumanString() const {
|
|
std::string common;
|
|
if (!name.empty()) {
|
|
common = absl::StrCat(" name=", name);
|
|
}
|
|
absl::StrAppend(&common, " type=", DataTypeString(type),
|
|
" shape=", ShapeHumanString());
|
|
absl::StrAppend(
|
|
&common, " is_same_data_across_replicas=", is_same_data_across_replicas);
|
|
switch (kind) {
|
|
case kInvalid:
|
|
return "invalid";
|
|
case kConstant:
|
|
return absl::StrCat("kind=constant", common,
|
|
" value=", constant_value.DebugString());
|
|
case kConstantResource:
|
|
return absl::StrCat("kind=constant-resource", common,
|
|
" value=", constant_value.DebugString());
|
|
case kResource: {
|
|
std::string output = absl::StrCat(
|
|
"kind=resource", common,
|
|
" resource_kind=", XlaResource::KindToString(resource_kind),
|
|
" initialized=", initialized, " is_fast_mem=", fast_mem);
|
|
if (max_array_size >= 0) {
|
|
absl::StrAppend(&output, " max_array_size=", max_array_size);
|
|
}
|
|
if (!tensor_array_gradients.empty()) {
|
|
absl::StrAppend(&output, " tensor_array_gradients=",
|
|
absl::StrJoin(tensor_array_gradients, ","));
|
|
}
|
|
return output;
|
|
}
|
|
case kParameter:
|
|
return absl::StrCat("kind=parameter", common);
|
|
case kTensorList:
|
|
return absl::StrCat("kind=tensorlist", common);
|
|
case kToken:
|
|
return absl::StrCat("token", common);
|
|
}
|
|
}
|
|
|
|
std::vector<int64_t> XlaCompiler::Argument::DimensionSizes() const {
|
|
if (absl::holds_alternative<TensorShape>(shape)) {
|
|
return xla::InlinedVectorToVector(std::get<TensorShape>(shape).dim_sizes());
|
|
} else {
|
|
return xla::SpanToVector(std::get<xla::Shape>(shape).dimensions());
|
|
}
|
|
}
|
|
|
|
absl::InlinedVector<int64_t, 4>
|
|
XlaCompiler::Argument::DimensionSizesAsInlinedVector() const {
|
|
if (absl::holds_alternative<TensorShape>(shape)) {
|
|
return std::get<TensorShape>(shape).dim_sizes();
|
|
} else {
|
|
auto v = std::get<xla::Shape>(shape).dimensions();
|
|
return absl::InlinedVector<int64_t, 4>(v.begin(), v.end());
|
|
}
|
|
}
|
|
|
|
std::string XlaCompiler::Argument::ShapeHumanString() const {
|
|
if (absl::holds_alternative<TensorShape>(shape)) {
|
|
return std::get<TensorShape>(shape).DebugString();
|
|
} else {
|
|
return std::get<xla::Shape>(shape).ToProto().DebugString();
|
|
}
|
|
}
|
|
|
|
XlaCompiler::XlaCompiler(XlaCompiler::Options options)
|
|
: options_(options),
|
|
initialization_status_(absl::OkStatus()),
|
|
next_step_id_(1),
|
|
device_(new XlaCompilationDevice(SessionOptions(), options_.device_type)),
|
|
device_mgr_(absl::WrapUnique(device_)) {
|
|
CHECK(!options_.device_type.type_string().empty());
|
|
if (options_.populate_resource_manager) {
|
|
initialization_status_ =
|
|
(*options_.populate_resource_manager)(device_->resource_manager());
|
|
}
|
|
|
|
local_flib_def_.reset(new FunctionLibraryDefinition(OpRegistry::Global(),
|
|
FunctionDefLibrary()));
|
|
local_pflr_.reset(new ProcessFunctionLibraryRuntime(
|
|
&device_mgr_, Env::Default(), /*config=*/nullptr,
|
|
options.graph_def_version, local_flib_def_.get(), OptimizerOptions()));
|
|
pflr_.reset(new ProcessFunctionLibraryRuntime(
|
|
&device_mgr_, Env::Default(), /*config=*/nullptr,
|
|
options.graph_def_version, options.flib_def, OptimizerOptions()));
|
|
|
|
local_flib_runtime_ = local_pflr_->GetFLR(device_->name());
|
|
flib_runtime_ = pflr_->GetFLR(device_->name());
|
|
|
|
// The default layout preference is no preference and the default shape
|
|
// representation function is the identity.
|
|
XlaShapeLayoutHelpers::ShapeDeterminationFns& shape_determination_fns =
|
|
options_.shape_determination_fns;
|
|
if (!shape_determination_fns.shape_representation_fn) {
|
|
shape_determination_fns.shape_representation_fn =
|
|
IdentityShapeRepresentationFn();
|
|
}
|
|
if (!shape_determination_fns.layout_preference_fn) {
|
|
shape_determination_fns.layout_preference_fn = UseNoPreferenceLayoutFn();
|
|
}
|
|
}
|
|
|
|
XlaCompiler::~XlaCompiler() = default;
|
|
|
|
int64_t XlaCompiler::NextStepId() { return next_step_id_++; }
|
|
|
|
uint64_t XlaCompiler::SignatureHash::operator()(
|
|
const std::pair<std::string, std::vector<Argument>>& signature) const {
|
|
return std::hash<std::string>()(signature.first);
|
|
}
|
|
|
|
static absl::Status GetFunctionBody(const NameAttrList& function,
|
|
FunctionLibraryRuntime* flib_runtime,
|
|
const FunctionBody** fbody) {
|
|
FunctionLibraryRuntime::Handle handle;
|
|
TF_RETURN_IF_ERROR(flib_runtime->Instantiate(
|
|
function.name(), AttrSlice(&function.attr()), &handle));
|
|
|
|
*fbody = flib_runtime->GetFunctionBody(handle);
|
|
TF_RET_CHECK(*fbody);
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
absl::Status XlaCompiler::FindFunctionBody(const NameAttrList& function,
|
|
const FunctionBody** fbody,
|
|
const ConfigProto** config_proto) {
|
|
// The function may be in either the local_flib_runtime_ or flib_runtime_.
|
|
// Look up the function in local first and if it is not found then look up the
|
|
// function in flib_runtime_.
|
|
auto status = GetFunctionBody(function, local_flib_runtime_, fbody);
|
|
if (!status.ok()) {
|
|
if (!absl::IsNotFound(status)) {
|
|
return status;
|
|
}
|
|
TF_RETURN_WITH_CONTEXT_IF_ERROR(
|
|
GetFunctionBody(function, flib_runtime_, fbody),
|
|
"Local lookup failed with: ", status.message());
|
|
if (config_proto) {
|
|
*config_proto = flib_runtime_->config_proto();
|
|
}
|
|
VLOG(4) << "Function " << function.name() << " in flib_runtime_";
|
|
} else {
|
|
if (config_proto) {
|
|
*config_proto = local_flib_runtime_->config_proto();
|
|
}
|
|
VLOG(4) << "Function " << function.name() << " in local_flib_runtime_";
|
|
}
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
std::unique_ptr<Graph> XlaCompiler::GetGraph(const FunctionBody* fbody) {
|
|
std::unique_ptr<Graph> graph(new Graph(options_.flib_def));
|
|
CopyGraph(*fbody->graph, graph.get());
|
|
|
|
bool is_inside_mustcompile = false;
|
|
if (fbody->record) {
|
|
PruneFunctionBody(fbody->record->fdef(), graph.get(), {});
|
|
TryGetNodeAttr(AttrSlice(&fbody->record->fdef().attr()),
|
|
kXlaMustCompileAttr, &is_inside_mustcompile);
|
|
}
|
|
|
|
// Performs a first function inlining pass before shape inference, since
|
|
// otherwise shape inference can't see inside functions and a comprehensive
|
|
// shape_map, including function ops, is needed to constant-propagate Shape
|
|
// Ops below.
|
|
auto flags = GetBuildXlaOpsPassFlags();
|
|
OptimizerOptions opts;
|
|
opts.set_opt_level(OptimizerOptions::L0);
|
|
opts.set_do_common_subexpression_elimination(false);
|
|
opts.set_do_function_inlining(true);
|
|
opts.set_do_constant_folding(!flags->tf_xla_disable_constant_folding);
|
|
GraphOptimizer optimizer(opts);
|
|
// Do not constant fold nodes that output DT_VARIANT type tensors.
|
|
// XLA does not support Const nodes of Variant type since it needs
|
|
// to know the original ops to be able to compile them to the relevant
|
|
// XLA form.
|
|
// TODO(srbs): This filter is a little conservative. E.g. a subgraph of
|
|
// the form:
|
|
// Const
|
|
// |
|
|
// EmptyTensorList -> TensorListPushBack -> TensorListPopBack -> Op
|
|
// |
|
|
// (Discard popped list)
|
|
//
|
|
// Would have been reduced to "Const -> Op" without this filter.
|
|
// However since we are only allowed to specify the filter at the "Node"
|
|
// level there is no good way to allow the above behavior. So we
|
|
// disallow any sort of constant folding on Variant nodes for now.
|
|
//
|
|
// Also do not consider constant folding Shape ops. When there is a dynamic
|
|
// dimension in a tensor, TF2XLA currently represent them as the static
|
|
// upperbound shape, which can be constant folded and then lose the info
|
|
// that this Shape is dynamic.
|
|
auto cf_consider_fn = [](const Node* n) {
|
|
for (const auto& output_arg : n->op_def().output_arg()) {
|
|
if (output_arg.type() == DT_VARIANT) {
|
|
return false;
|
|
}
|
|
}
|
|
const auto& ts = n->type_string();
|
|
// XLA has special logic to handle dynamic shapes, don't constant fold
|
|
// them.
|
|
if (ts == "Shape" || ts == "ShapeN" || ts == "Size") {
|
|
return false;
|
|
}
|
|
return true;
|
|
};
|
|
GraphOptimizer::Options graph_optimizer_options;
|
|
graph_optimizer_options.cf_consider_fn = cf_consider_fn;
|
|
graph_optimizer_options.inline_multi_device_functions = true;
|
|
graph_optimizer_options.inline_impl_selection_group_functions = true;
|
|
graph_optimizer_options.inline_with_single_device_body_placer = true;
|
|
graph_optimizer_options.ignore_noinline = is_inside_mustcompile;
|
|
|
|
{
|
|
GraphShapeInfo shape_info;
|
|
InferShapes(graph.get(), /*arg_shapes=*/{},
|
|
flib_runtime_->GetFunctionLibraryDefinition(), &shape_info)
|
|
.IgnoreError();
|
|
auto node_name_index = graph->BuildNodeNameIndex();
|
|
std::unordered_map<std::string, std::vector<PartialTensorShape>> shape_map;
|
|
for (const auto& node_shape_info : shape_info) {
|
|
const std::string& node_name = node_shape_info.first;
|
|
const std::vector<InferredShape>& output_shapes = node_shape_info.second;
|
|
const auto& node_iter = node_name_index.find(node_name);
|
|
if (node_iter != node_name_index.end()) {
|
|
auto& partial_shapes = shape_map[node_name];
|
|
for (const auto& inferred_shape : output_shapes) {
|
|
partial_shapes.push_back(inferred_shape.shape);
|
|
}
|
|
}
|
|
}
|
|
graph_optimizer_options.shape_map = &shape_map;
|
|
optimizer.Optimize(flib_runtime_, flib_runtime_->env(),
|
|
/*device=*/nullptr, &graph, graph_optimizer_options);
|
|
}
|
|
|
|
// Run shape inference on the graph and optimize the graph again.
|
|
GraphShapeInfo shape_info;
|
|
InferShapes(graph.get(), /*arg_shapes=*/{},
|
|
flib_runtime_->GetFunctionLibraryDefinition(), &shape_info)
|
|
.IgnoreError();
|
|
auto node_name_index = graph->BuildNodeNameIndex();
|
|
std::unordered_map<std::string, std::vector<PartialTensorShape>> shape_map;
|
|
for (const auto& node_shape_info : shape_info) {
|
|
const std::string& node_name = node_shape_info.first;
|
|
const std::vector<InferredShape>& output_shapes = node_shape_info.second;
|
|
const auto& node_iter = node_name_index.find(node_name);
|
|
if (node_iter != node_name_index.end()) {
|
|
auto& partial_shapes = shape_map[node_name];
|
|
for (const auto& inferred_shape : output_shapes) {
|
|
partial_shapes.push_back(inferred_shape.shape);
|
|
}
|
|
}
|
|
}
|
|
graph_optimizer_options.shape_map = &shape_map;
|
|
optimizer.Optimize(flib_runtime_, flib_runtime_->env(),
|
|
/*device=*/nullptr, &graph, graph_optimizer_options);
|
|
|
|
return graph;
|
|
}
|
|
|
|
// Collects all control rets from `orig_control_ret_nodes` that are still valid,
|
|
// keeping the same order.
|
|
std::vector<std::string> GetValidControlRets(
|
|
absl::Span<Node* const> orig_control_ret_nodes, const Graph& graph) {
|
|
// Build map from control ret node name to index.
|
|
// We use Node name instead of Node* here to index into the map as we populate
|
|
// the map with nodes in FunctionDef control_ret_nodes and later query it
|
|
// using the nodes in `graph`. The Node pointers would be different but the
|
|
// Node name is expected to remain the same between the two.
|
|
absl::flat_hash_map<std::string, int> control_ret_nodes_map;
|
|
for (int i = 0; i < orig_control_ret_nodes.size(); ++i) {
|
|
const Node* n = orig_control_ret_nodes[i];
|
|
control_ret_nodes_map[n->name()] = i;
|
|
}
|
|
// Check which control rets are still valid.
|
|
std::vector<bool> is_valid_control_ret(orig_control_ret_nodes.size(), false);
|
|
int num_valid_control_rets = 0;
|
|
for (const Node* n : graph.nodes()) {
|
|
auto iter = control_ret_nodes_map.find(n->name());
|
|
if (iter != control_ret_nodes_map.end()) {
|
|
++num_valid_control_rets;
|
|
is_valid_control_ret[iter->second] = true;
|
|
}
|
|
}
|
|
// Return valid control rets in same order as they appear in
|
|
// `orig_control_ret_nodes`.
|
|
std::vector<std::string> valid_control_rets;
|
|
valid_control_rets.reserve(num_valid_control_rets);
|
|
for (int i = 0; i < orig_control_ret_nodes.size(); ++i) {
|
|
if (is_valid_control_ret[i]) {
|
|
valid_control_rets.push_back(orig_control_ret_nodes[i]->name());
|
|
}
|
|
}
|
|
return valid_control_rets;
|
|
}
|
|
|
|
absl::Status XlaCompiler::CompileSingleOp(
|
|
const XlaCompiler::CompileOptions& compile_options,
|
|
const XlaCompiler::SingleOpCompileArgument& single_op_compile_argument,
|
|
absl::Span<const Argument> args, XlaCompiler::CompilationResult* result) {
|
|
const NodeDef& node_def = single_op_compile_argument.node_def;
|
|
TF_ASSIGN_OR_RETURN(
|
|
auto graph,
|
|
CreateSingleOpGraph(node_def, args,
|
|
single_op_compile_argument.output_dtypes));
|
|
|
|
*result = {};
|
|
absl::Status status = ADD_SOURCE_LOCATION(CompileGraph(
|
|
compile_options, node_def.name(), std::move(graph), args, result));
|
|
if (status.ok()) {
|
|
tensorflow::metrics::IncrementPhase2XlaCompilerCounter(
|
|
tensorflow::metrics::Phase2XlaCompilerMetric::
|
|
kCompileSingleOpXlaBuilderSuccess);
|
|
} else {
|
|
tensorflow::metrics::IncrementPhase2XlaCompilerCounter(
|
|
tensorflow::metrics::Phase2XlaCompilerMetric::
|
|
kCompileSingleOpXlaBuilderFailure);
|
|
tsl::error_logging::Log(mlir::TF::kBridgeComponent, kSingleOpComponent,
|
|
status.ToString())
|
|
.IgnoreError();
|
|
}
|
|
return status;
|
|
}
|
|
|
|
absl::Status XlaCompiler::CompileFunction(
|
|
const XlaCompiler::CompileOptions& options,
|
|
const NameAttrList& fn_name_attrs,
|
|
absl::Span<const XlaCompiler::Argument> args,
|
|
XlaCompiler::CompilationResult* result) {
|
|
std::string function_id =
|
|
Canonicalize(fn_name_attrs.name(), AttrSlice(&fn_name_attrs.attr()));
|
|
VLOG(1) << "XlaCompiler::CompileFunction " << function_id;
|
|
|
|
const std::vector<XlaCompiler::Argument> arg_vector(args.begin(), args.end());
|
|
auto it = cache_.find({function_id, arg_vector});
|
|
if (it != cache_.end()) {
|
|
*result = it->second;
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
const FunctionBody* fbody;
|
|
const ConfigProto* config = nullptr;
|
|
TF_RETURN_IF_ERROR(FindFunctionBody(fn_name_attrs, &fbody, &config));
|
|
|
|
std::optional<ConfigProto> config_proto;
|
|
if (config) {
|
|
config_proto = *config;
|
|
}
|
|
|
|
TF_RETURN_WITH_CONTEXT_IF_ERROR(
|
|
CheckSignature(fbody->arg_types, args),
|
|
"Signature check failure while compiling: ", fn_name_attrs.name());
|
|
|
|
// Set shapes for _Arg nodes. They are useful for constant folding (e.g. an
|
|
// Xla op requires a compile-time constant input, and that input is shape of
|
|
// an _Arg node.
|
|
for (int i = 0, end = args.size(); i < end; i++) {
|
|
// Skip resource variables and tensor lists.
|
|
DataType dtype;
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(fbody->arg_nodes[i]->def(), "T", &dtype));
|
|
if (dtype == DT_RESOURCE || dtype == DT_VARIANT) {
|
|
continue;
|
|
}
|
|
|
|
if (absl::holds_alternative<xla::Shape>(args[i].shape)) {
|
|
xla::Shape xla_shape = std::get<xla::Shape>(args[i].shape);
|
|
TensorShape tensor_shape;
|
|
// If xla_shape is dynamic, prevent constant folding by not setting
|
|
// output_shapes.
|
|
if (XLAShapeToTensorShape(xla_shape, &tensor_shape).ok() &&
|
|
xla_shape.is_static()) {
|
|
fbody->arg_nodes[i]->ClearAttr("_output_shapes");
|
|
fbody->arg_nodes[i]->AddAttr("_output_shapes",
|
|
std::vector<TensorShape>{tensor_shape});
|
|
}
|
|
} else {
|
|
TensorShape tensor_shape = std::get<TensorShape>(args[i].shape);
|
|
fbody->arg_nodes[i]->ClearAttr("_output_shapes");
|
|
fbody->arg_nodes[i]->AddAttr("_output_shapes",
|
|
std::vector<TensorShape>{tensor_shape});
|
|
}
|
|
}
|
|
|
|
std::unique_ptr<Graph> graph = GetGraph(fbody);
|
|
|
|
// _Arg and _Retval nodes don't exist in the stored subgraph for the function;
|
|
// they are added by the function body looked up. Therefore, they don't have
|
|
// core assignments here.
|
|
// Attempt to assign a core to each _Retval and _Arg. Chooses the
|
|
// lowest-numbered core that consumes the argument. We choose the
|
|
// lowest-numbered core so the assignment is deterministic.
|
|
for (Node* n : graph->nodes()) {
|
|
if (n->IsArg()) {
|
|
TF_RETURN_IF_ERROR(SetNodeShardingFromNeighbors(n, /*out_edges=*/true));
|
|
}
|
|
}
|
|
// Do _Retval as a second loop, in case the retval's input is an _Arg (which
|
|
// may have gotten a device assignment from the first loop).
|
|
for (Node* n : graph->nodes()) {
|
|
if (n->IsRetval()) {
|
|
TF_RETURN_IF_ERROR(SetNodeShardingFromNeighbors(n, /*out_edges=*/false));
|
|
}
|
|
}
|
|
|
|
if (VLOG_IS_ON(2)) {
|
|
VLOG(2) << "XlaCompiler::CompileFunction: "
|
|
<< DumpGraphToFile(
|
|
absl::StrCat("xla_compile_function_", function_id), *graph);
|
|
}
|
|
|
|
VLOG(1) << "====================================================";
|
|
|
|
VLOG(1) << "CompileFunction with XlaBuilder";
|
|
auto status =
|
|
CompileGraph(options, function_id, std::move(graph), args, result);
|
|
if (!status.ok()) {
|
|
tensorflow::metrics::IncrementPhase2XlaCompilerCounter(
|
|
tensorflow::metrics::Phase2XlaCompilerMetric::
|
|
kCompileFunctionXlaBuilderFailure);
|
|
::tsl::errors::AppendToMessage(
|
|
&status, "tf2xla conversion failed while converting ",
|
|
std::move(function_id),
|
|
". Run with TF_DUMP_GRAPH_PREFIX=/path/to/dump/dir and "
|
|
"--vmodule=xla_compiler=2 to obtain a dump of the compiled "
|
|
"functions.");
|
|
tsl::error_logging::Log(mlir::TF::kBridgeComponent,
|
|
kCompileFunctionComponent, status.ToString())
|
|
.IgnoreError();
|
|
return status;
|
|
}
|
|
|
|
tensorflow::metrics::IncrementPhase2XlaCompilerCounter(
|
|
tensorflow::metrics::Phase2XlaCompilerMetric::
|
|
kCompileFunctionXlaBuilderSuccess);
|
|
VLOG(1) << "====================================================";
|
|
cache_[{function_id, arg_vector}] = *result;
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
// Computes the XLA shape for argument 'arg'.
|
|
absl::Status XlaCompiler::XLAShapeForArgument(
|
|
const XlaCompiler::Argument& arg, bool is_entry_computation,
|
|
const std::optional<xla::HloSharding>& arg_sharding,
|
|
xla::Shape* xla_shape) const {
|
|
switch (arg.kind) {
|
|
case XlaCompiler::Argument::kConstant:
|
|
LOG(FATAL) << "Unreachable case";
|
|
case XlaCompiler::Argument::kParameter: {
|
|
if (is_entry_computation) {
|
|
TensorShape shape;
|
|
if (std::holds_alternative<TensorShape>(arg.shape)) {
|
|
shape = std::get<TensorShape>(arg.shape);
|
|
} else {
|
|
TF_RETURN_IF_ERROR(
|
|
XLAShapeToTensorShape(std::get<xla::Shape>(arg.shape), &shape));
|
|
}
|
|
auto layout_preference =
|
|
options_.shape_determination_fns.layout_preference_fn(
|
|
shape, arg.type, arg.kind);
|
|
TF_ASSIGN_OR_RETURN(
|
|
*xla_shape,
|
|
options_.shape_determination_fns.shape_representation_fn(
|
|
shape, arg.type,
|
|
/*use_fast_memory=*/false, layout_preference));
|
|
TF_RETURN_IF_ERROR(RewriteLayoutWithShardedShape(
|
|
arg_sharding, /*use_fast_memory=*/false,
|
|
options_.shape_determination_fns, xla_shape));
|
|
// If the arg is dynamic then we update the shape to reflect that. The
|
|
// layout etc above lose it by forcing a swap to TensorShape.
|
|
if (std::holds_alternative<xla::Shape>(arg.shape) &&
|
|
std::get<xla::Shape>(arg.shape).is_dynamic()) {
|
|
xla::Shape dynamic_shape = std::get<xla::Shape>(arg.shape);
|
|
for (int i = 0; i < xla_shape->dimensions().size(); ++i) {
|
|
xla_shape->set_dynamic_dimension(
|
|
i, dynamic_shape.is_dynamic_dimension(i));
|
|
}
|
|
}
|
|
} else {
|
|
if (std::holds_alternative<xla::Shape>(arg.shape)) {
|
|
*xla_shape = std::get<xla::Shape>(arg.shape);
|
|
} else {
|
|
TF_RETURN_IF_ERROR(TensorShapeToXLAShape(
|
|
arg.type, std::get<TensorShape>(arg.shape), xla_shape));
|
|
}
|
|
}
|
|
return absl::OkStatus();
|
|
}
|
|
case XlaCompiler::Argument::kTensorList: {
|
|
TF_RET_CHECK(absl::holds_alternative<xla::Shape>(arg.shape));
|
|
*xla_shape = std::get<xla::Shape>(arg.shape);
|
|
return absl::OkStatus();
|
|
}
|
|
case XlaCompiler::Argument::kConstantResource:
|
|
case XlaCompiler::Argument::kResource: {
|
|
TF_RET_CHECK(arg.initialized);
|
|
|
|
switch (arg.resource_kind) {
|
|
case XlaResource::kVariable: {
|
|
TF_RET_CHECK(absl::holds_alternative<TensorShape>(arg.shape));
|
|
auto layout_preference =
|
|
options_.shape_determination_fns.layout_preference_fn(
|
|
std::get<TensorShape>(arg.shape), arg.type, arg.kind);
|
|
TF_ASSIGN_OR_RETURN(
|
|
*xla_shape,
|
|
options_.shape_determination_fns.shape_representation_fn(
|
|
std::get<TensorShape>(arg.shape), arg.type,
|
|
/*use_fast_memory=*/arg.fast_mem, layout_preference));
|
|
TF_RETURN_IF_ERROR(RewriteLayoutWithShardedShape(
|
|
arg_sharding, arg.fast_mem, options_.shape_determination_fns,
|
|
xla_shape));
|
|
return absl::OkStatus();
|
|
}
|
|
case XlaResource::kTensorArray: {
|
|
if (arg.max_array_size < 0) {
|
|
return absl::InvalidArgumentError(
|
|
"Negative max_array_size in XLAShapeForArgument");
|
|
}
|
|
TF_RET_CHECK(absl::holds_alternative<TensorShape>(arg.shape));
|
|
TensorShape shape;
|
|
TF_RETURN_IF_ERROR(shape.AddDimWithStatus(arg.max_array_size));
|
|
shape.AppendShape(std::get<TensorShape>(arg.shape));
|
|
TF_RETURN_IF_ERROR(TensorShapeToXLAShape(arg.type, shape, xla_shape));
|
|
|
|
if (!arg.tensor_array_gradients.empty()) {
|
|
std::vector<xla::Shape> tuple_shape(
|
|
arg.tensor_array_gradients.size() + 1, *xla_shape);
|
|
*xla_shape = xla::ShapeUtil::MakeTupleShape(tuple_shape);
|
|
}
|
|
return absl::OkStatus();
|
|
}
|
|
case XlaResource::kStack: {
|
|
if (arg.max_array_size < 0) {
|
|
return absl::InvalidArgumentError(
|
|
"Negative max_array_size in XLAShapeForArgument");
|
|
}
|
|
TF_RET_CHECK(absl::holds_alternative<TensorShape>(arg.shape));
|
|
TensorShape shape;
|
|
TF_RETURN_IF_ERROR(shape.AddDimWithStatus(arg.max_array_size));
|
|
shape.AppendShape(std::get<TensorShape>(arg.shape));
|
|
xla::Shape buffer_shape;
|
|
TF_RETURN_IF_ERROR(
|
|
TensorShapeToXLAShape(arg.type, shape, &buffer_shape));
|
|
*xla_shape = xla::ShapeUtil::MakeTupleShape(
|
|
{buffer_shape, xla::ShapeUtil::MakeShape(xla::S32, {})});
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
case XlaResource::kInvalid:
|
|
return absl::InternalError(
|
|
"Invalid resource type in XLAShapeForArgument()");
|
|
}
|
|
}
|
|
case XlaCompiler::Argument::kToken: {
|
|
*xla_shape = xla::ShapeUtil::MakeTokenShape();
|
|
return absl::OkStatus();
|
|
}
|
|
case XlaCompiler::Argument::kInvalid:
|
|
return absl::InternalError(
|
|
"Invalid argument type in XLAShapeForArgument()");
|
|
}
|
|
}
|
|
|
|
/* static */
|
|
void XlaCompiler::PopulateArgumentFromResource(const XlaResource& resource,
|
|
Argument* arg) {
|
|
arg->initialized = resource.initialized();
|
|
arg->kind = XlaCompiler::Argument::kResource;
|
|
arg->resource_kind = resource.kind();
|
|
|
|
arg->type = resource.type();
|
|
arg->shape = resource.shape();
|
|
arg->max_array_size = resource.max_array_size();
|
|
for (const auto& gradient : resource.tensor_array_gradients()) {
|
|
arg->tensor_array_gradients.insert(gradient.first);
|
|
}
|
|
arg->name = resource.name();
|
|
}
|
|
|
|
XlaCompiler::SingleOpCompileArgument::SingleOpCompileArgument(
|
|
const OpKernelContext& ctx) {
|
|
std::vector<DataType> output_dtypes(ctx.num_outputs());
|
|
for (int i = 0; i < output_dtypes.size(); ++i) {
|
|
output_dtypes[i] = ctx.expected_output_dtype(i);
|
|
}
|
|
this->output_dtypes = output_dtypes;
|
|
this->node_def = ctx.op_kernel().def();
|
|
auto* config_proto = ctx.function_library()->config_proto();
|
|
if (config_proto != nullptr) {
|
|
this->config_proto = *config_proto;
|
|
}
|
|
}
|
|
|
|
// Builds XLA computations for each of the arguments to the computation.
|
|
// `args` are the arguments to the computation.
|
|
absl::Status XlaCompiler::BuildArguments(
|
|
const Graph& graph, const std::vector<XlaCompiler::Argument>& args,
|
|
bool use_tuple_arg, xla::XlaBuilder* builder, XlaContext* context,
|
|
const std::map<int, xla::OpSharding>& arg_shardings,
|
|
std::vector<XlaExpression>* arg_expressions,
|
|
std::vector<int>* input_to_args, std::vector<xla::Shape>* input_shapes,
|
|
std::string& args_tuple_sdy_sharding, bool is_entry_computation) {
|
|
arg_expressions->resize(args.size());
|
|
|
|
// Argument numbers of arguments and resources that are to be passed to the
|
|
// XLA computation as runtime parameters. `input_to_args[a] = b` means that
|
|
// the a'th XLA input corresponds to the b'th original arg indexes.
|
|
input_to_args->clear();
|
|
input_to_args->reserve(args.size());
|
|
|
|
// Fills in constant arguments, and computes non-constant argument order.
|
|
for (std::vector<XlaCompiler::Argument>::size_type i = 0; i < args.size();
|
|
++i) {
|
|
const XlaCompiler::Argument& arg = args[i];
|
|
XlaExpression& arg_expression = (*arg_expressions)[i];
|
|
switch (arg.kind) {
|
|
case XlaCompiler::Argument::kConstantResource:
|
|
case XlaCompiler::Argument::kResource: {
|
|
TF_RET_CHECK(arg.resource_kind != XlaResource::kInvalid);
|
|
TF_RET_CHECK(absl::holds_alternative<TensorShape>(arg.shape));
|
|
// TODO(phawkins): this code assumes that resource arguments do not
|
|
// alias.
|
|
XlaResource* resource =
|
|
context->AddResource(std::make_unique<XlaResource>(
|
|
arg.resource_kind, i, arg.name, arg.type,
|
|
std::get<TensorShape>(arg.shape), xla::XlaOp(),
|
|
/*max_array_size=*/arg.max_array_size,
|
|
/*tensor_array_gradients=*/arg.tensor_array_gradients,
|
|
/*tensor_array_multiple_writes_aggregate=*/true,
|
|
arg.definition_stack_trace));
|
|
arg_expression =
|
|
arg.kind == XlaCompiler::Argument::kResource
|
|
? XlaExpression::Resource(resource)
|
|
: XlaExpression::ConstantResource(arg.constant_value, resource);
|
|
if (arg.initialized) {
|
|
input_to_args->push_back(i);
|
|
}
|
|
break;
|
|
}
|
|
case XlaCompiler::Argument::kParameter:
|
|
case XlaCompiler::Argument::kTensorList:
|
|
case XlaCompiler::Argument::kToken: {
|
|
input_to_args->push_back(i);
|
|
break;
|
|
}
|
|
case XlaCompiler::Argument::kConstant:
|
|
arg_expression = XlaExpression::Constant(arg.constant_value);
|
|
break;
|
|
case XlaCompiler::Argument::kInvalid:
|
|
return absl::InternalError(
|
|
"Unreachable case in BuildArguments() while filling constant args");
|
|
}
|
|
}
|
|
|
|
if (input_to_args->empty() && !use_tuple_arg) {
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
// `arg_to_inputs[c] = d` means that the c'th original arg index corresponds
|
|
// to the d'th XLA input. Note that the value -1 corresponds to constants, or
|
|
// other args that don't correspond to an input.
|
|
std::vector<int> arg_to_inputs(args.size(), -1);
|
|
for (int i = 0, end = input_to_args->size(); i < end; i++) {
|
|
arg_to_inputs[input_to_args->at(i)] = i;
|
|
}
|
|
|
|
std::vector<xla::Shape> arg_shapes(input_to_args->size());
|
|
for (std::vector<int>::size_type i = 0; i < input_to_args->size(); ++i) {
|
|
// Computes the shapes of non-constant arguments.
|
|
auto arg_sharding = arg_shardings.find((*input_to_args)[i]);
|
|
std::optional<xla::HloSharding> sharding;
|
|
if (arg_sharding != arg_shardings.end()) {
|
|
TF_ASSIGN_OR_RETURN(auto hlo_sharding,
|
|
xla::HloSharding::FromProto(arg_sharding->second));
|
|
sharding = hlo_sharding;
|
|
}
|
|
TF_RETURN_IF_ERROR(XLAShapeForArgument(args[(*input_to_args)[i]],
|
|
is_entry_computation, sharding,
|
|
&arg_shapes[i]));
|
|
}
|
|
|
|
if (use_tuple_arg) {
|
|
input_shapes->push_back(xla::ShapeUtil::MakeTupleShape(arg_shapes));
|
|
} else {
|
|
*input_shapes = arg_shapes;
|
|
}
|
|
|
|
// Attach a common operator name as metadata. This has no semantic effect — it
|
|
// merely makes the HLO graph more readable when visualized via TensorBoard,
|
|
// since TensorBoard forms groups out of operators with similar names.
|
|
xla::OpMetadata arg_metadata;
|
|
arg_metadata.set_op_name("XLA_Args");
|
|
builder->SetOpMetadata(arg_metadata);
|
|
|
|
// Build parameter handles for non-constant arguments.
|
|
std::vector<xla::XlaOp> arg_handles(input_to_args->size());
|
|
if (use_tuple_arg) {
|
|
xla::XlaOp tuple;
|
|
if (is_entry_computation) {
|
|
xla::OpSharding tuple_sharding;
|
|
tuple_sharding.set_type(xla::OpSharding::TUPLE);
|
|
for (int64_t parameter : *input_to_args) {
|
|
auto it = arg_shardings.find(parameter);
|
|
*tuple_sharding.add_tuple_shardings() =
|
|
it == arg_shardings.end() ? xla::sharding_builder::SingleDevice(0)
|
|
: it->second;
|
|
}
|
|
std::vector<bool> is_same_across_replicas;
|
|
for (int i = 0, end = input_to_args->size(); i < end; ++i) {
|
|
// Add an entry to is_same_across_replicas for every leaf buffer.
|
|
is_same_across_replicas.insert(
|
|
is_same_across_replicas.end(),
|
|
xla::ShapeUtil::GetLeafCount(arg_shapes[i]),
|
|
args[input_to_args->at(i)].is_same_data_across_replicas);
|
|
}
|
|
xla::XlaScopedShardingAssignment assign_tuple_sharding(
|
|
builder, input_to_args->empty() ? std::optional<xla::OpSharding>()
|
|
: tuple_sharding);
|
|
tuple = xla::Parameter(builder, 0, (*input_shapes)[0], "arg_tuple",
|
|
is_same_across_replicas);
|
|
if (options_.use_shardy_partitioner && builder->sharding().has_value()) {
|
|
args_tuple_sdy_sharding =
|
|
xla::sdy::convertToSdySharding(tuple_sharding, (*input_shapes)[0],
|
|
/*openDims=*/false,
|
|
/*inlineMesh=*/true);
|
|
}
|
|
} else {
|
|
tuple = xla::Parameter(builder, 0, (*input_shapes)[0], "arg_tuple");
|
|
}
|
|
|
|
for (std::vector<int>::size_type i = 0; i < input_to_args->size(); ++i) {
|
|
auto it = arg_shardings.find(i);
|
|
xla::XlaScopedShardingAssignment assign_sharding(
|
|
builder, it == arg_shardings.end() ? std::optional<xla::OpSharding>()
|
|
: it->second);
|
|
auto& arg = args[input_to_args->at(i)];
|
|
|
|
xla::OpMetadata arg_metadata;
|
|
arg_metadata.set_op_name(arg.node_name);
|
|
builder->SetOneShotOpMetadata(arg_metadata);
|
|
arg_handles[i] = xla::GetTupleElement(tuple, i);
|
|
if (is_entry_computation && options_.use_shardy_partitioner) {
|
|
// Add shardy shardings only for entry computation, as we only need to
|
|
// add it for wrapper main added in tensorflow.
|
|
TF_RETURN_IF_ERROR(addSdyShardingFrontendAttribute(
|
|
builder, arg_handles[i], arg_shapes[i]));
|
|
}
|
|
}
|
|
} else {
|
|
for (std::vector<int>::size_type i = 0; i < input_to_args->size(); ++i) {
|
|
auto it = arg_shardings.find(i);
|
|
xla::XlaScopedShardingAssignment assign_sharding(
|
|
builder, it == arg_shardings.end() ? std::optional<xla::OpSharding>()
|
|
: it->second);
|
|
if (is_entry_computation) {
|
|
// Add an entry to is_same_across_replicas for every leaf buffer.
|
|
std::vector<bool> is_same_across_replicas(
|
|
xla::ShapeUtil::GetLeafCount((*input_shapes)[i]),
|
|
args[input_to_args->at(i)].is_same_data_across_replicas);
|
|
arg_handles[i] =
|
|
xla::Parameter(builder, i, (*input_shapes)[i],
|
|
absl::StrCat("arg", i), is_same_across_replicas);
|
|
if (options_.use_shardy_partitioner) {
|
|
TF_RETURN_IF_ERROR(addSdyShardingFrontendAttribute(
|
|
builder, arg_handles[i], (*input_shapes)[i],
|
|
/*is_single_arg=*/true));
|
|
}
|
|
} else {
|
|
arg_handles[i] = xla::Parameter(builder, i, (*input_shapes)[i],
|
|
absl::StrCat("arg", i));
|
|
}
|
|
}
|
|
}
|
|
|
|
builder->ClearOpMetadata();
|
|
|
|
// Fill in the handles in non-constant arguments, and reshape parameters
|
|
// back to their correct shapes.
|
|
VLOG(2) << "XLA computation inputs:";
|
|
for (std::vector<int>::size_type i = 0; i < input_to_args->size(); ++i) {
|
|
const XlaCompiler::Argument& arg = args[input_to_args->at(i)];
|
|
VLOG(2) << " XLA arg " << i
|
|
<< " shape: " << xla::ShapeUtil::HumanString(arg_shapes[i])
|
|
<< " name: " << arg.name << " TF arg " << input_to_args->at(i)
|
|
<< " node name: " << arg.node_name
|
|
<< (arg_shardings.find(i) == arg_shardings.end()
|
|
? ""
|
|
: absl::StrCat(" sharding: ",
|
|
arg_shardings.at(i).DebugString()));
|
|
XlaExpression& arg_expression = (*arg_expressions)[input_to_args->at(i)];
|
|
switch (arg.kind) {
|
|
case XlaCompiler::Argument::kConstantResource:
|
|
case XlaCompiler::Argument::kResource: {
|
|
TF_RET_CHECK(arg.initialized);
|
|
XlaResource* resource = arg_expression.resource();
|
|
TF_RETURN_IF_ERROR(resource->SetFromPack(arg.tensor_array_gradients,
|
|
arg_handles[i], builder));
|
|
VLOG(2) << " resource: num_gradients: "
|
|
<< arg.tensor_array_gradients.size();
|
|
break;
|
|
}
|
|
case XlaCompiler::Argument::kParameter:
|
|
// Reshape parameters back to their correct shapes.
|
|
// TODO(b/76097077): propagate device assignments onto arguments and
|
|
// return values of functions, and then reshape unconditionally.
|
|
if (is_entry_computation) {
|
|
arg_expression = XlaExpression::XlaOp(
|
|
xla::Reshape(arg_handles[i], arg.DimensionSizes()), arg.type);
|
|
} else {
|
|
arg_expression = XlaExpression::XlaOp(arg_handles[i], arg.type);
|
|
if (arg.value_bound) {
|
|
TF_RET_CHECK(arg.value_dynamism);
|
|
// Propagate upper bound and value dynamism to arg_expression.
|
|
arg_expression.set_value_bound(arg.value_bound.value());
|
|
arg_expression.set_value_dynamism(arg.value_dynamism.value());
|
|
}
|
|
}
|
|
break;
|
|
case XlaCompiler::Argument::kTensorList: {
|
|
arg_expression = XlaExpression::TensorList(arg_handles[i]);
|
|
break;
|
|
}
|
|
case XlaCompiler::Argument::kToken: {
|
|
arg_expression = XlaExpression::XlaOp(arg_handles[i], arg.type);
|
|
break;
|
|
}
|
|
case XlaCompiler::Argument::kConstant:
|
|
case XlaCompiler::Argument::kInvalid:
|
|
return absl::InternalError(
|
|
"Unreachable case in BuildArguments() while filling handles");
|
|
}
|
|
}
|
|
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
namespace {
|
|
|
|
// Check that the ops of all non-functional nodes have been registered.
|
|
absl::Status ValidateFunctionDef(const FunctionDef* fdef,
|
|
const FunctionLibraryDefinition& flib_def) {
|
|
for (const NodeDef& node : fdef->node_def()) {
|
|
const std::string& op = node.op();
|
|
if (op == FunctionLibraryDefinition::kGradientOp || flib_def.Find(op)) {
|
|
continue;
|
|
}
|
|
const OpDef* op_def;
|
|
TF_RETURN_IF_ERROR(OpRegistry::Global()->LookUpOpDef(op, &op_def));
|
|
}
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
// If node is PartitionedCall or StatefulPartitionedCall, returns the
|
|
// name from the "f" attr, else returns node.def().op().
|
|
// Returned pointer points to the internal string either in node's attributes
|
|
// or in its NodeDef. This pointer is valid as long as the node has not been
|
|
// modified.
|
|
absl::Status GetPotentialFunctionName(const Node& node,
|
|
const std::string** name) {
|
|
if (node.IsPartitionedCall()) {
|
|
const AttrValue* attr_value;
|
|
TF_RETURN_IF_ERROR(
|
|
node.attrs().Find(FunctionLibraryDefinition::kFuncAttr, &attr_value));
|
|
if (!attr_value->has_func()) {
|
|
return absl::InvalidArgumentError(
|
|
absl::StrCat("The attribute value for attribute 'f' in node ",
|
|
node.DebugString(), " does not have 'func' field set"));
|
|
}
|
|
*name = &attr_value->func().name();
|
|
return absl::OkStatus();
|
|
}
|
|
*name = &node.type_string();
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
// Check that the graph doesn't have any invalid nodes (e.g. incompatible with
|
|
// given device_type, invalid data type, missing attributes...)
|
|
absl::Status ValidateGraph(const Graph* graph,
|
|
const FunctionLibraryDefinition& flib_def,
|
|
const DeviceType& device_type,
|
|
const std::string& name) {
|
|
// Make sure the XLA compilation kernels are registered. This operation is
|
|
// idempotent so it is fine if someone called it already.
|
|
XlaOpRegistry::RegisterCompilationKernels();
|
|
|
|
auto maybe_error = [&](const Node* node,
|
|
const absl::Status& s) -> absl::Status {
|
|
if (!s.ok()) {
|
|
std::string errmsg = absl::StrCat(
|
|
"Detected unsupported operations when trying to compile graph ", name,
|
|
" on ", device_type.type_string(), ": ", node->def().op(), " (",
|
|
s.message(), ")", FormatNodeForError(*node));
|
|
if (absl::StrContains(device_type.type_string(), "TPU")) {
|
|
absl::StrAppend(&errmsg,
|
|
"\nOne approach is to outside compile the unsupported "
|
|
"ops to run on CPUs by enabling soft placement "
|
|
"`tf.config.set_soft_device_placement(True)`."
|
|
" This has a potential performance penalty.\n");
|
|
}
|
|
if (std::shared_ptr<AbstractStackTrace> stack_trace =
|
|
node->GetStackTrace()) {
|
|
absl::StrAppend(
|
|
&errmsg, "\nThe op is created at: \n",
|
|
stack_trace->ToString({/*show_line_contents =*/true,
|
|
/*filter_common_prefix =*/true,
|
|
/*drop_internal_frames =*/true}));
|
|
}
|
|
|
|
return absl::InvalidArgumentError(errmsg);
|
|
}
|
|
return absl::OkStatus();
|
|
};
|
|
|
|
for (const Node* node : graph->nodes()) {
|
|
if (node->type_string() == FunctionLibraryDefinition::kGradientOp) {
|
|
continue;
|
|
}
|
|
const std::string* function_name;
|
|
TF_RETURN_IF_ERROR(GetPotentialFunctionName(*node, &function_name));
|
|
const FunctionDef* fdef = flib_def.Find(*function_name);
|
|
absl::Status s;
|
|
if (fdef) {
|
|
s = ValidateFunctionDef(fdef, flib_def);
|
|
TF_RETURN_IF_ERROR(maybe_error(node, s));
|
|
continue;
|
|
}
|
|
const OpDef* op_def;
|
|
s = OpRegistry::Global()->LookUpOpDef(node->def().op(), &op_def);
|
|
TF_RETURN_IF_ERROR(maybe_error(node, s));
|
|
TF_RETURN_IF_ERROR(ValidateNodeDef(node->def(), *op_def));
|
|
s = FindKernelDef(device_type, node->def(), nullptr, nullptr);
|
|
TF_RETURN_IF_ERROR(maybe_error(node, s));
|
|
}
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
void ConvertConstantsToExpressions(xla::XlaBuilder* builder,
|
|
absl::Span<XlaExpression> expressions) {
|
|
for (XlaExpression& expression : expressions) {
|
|
if (expression.kind() == XlaExpression::Kind::kConstant) {
|
|
expression =
|
|
XlaExpression::XlaOp(expression.AsXlaOp(builder), expression.dtype());
|
|
}
|
|
}
|
|
}
|
|
|
|
} // namespace
|
|
|
|
// A temporary dummy stack trace, used to identify locations where stack trace
|
|
// info is being lost, and to clarify how stack trace info is otherwise being
|
|
// handled in individual passes. This class and its usage below will be removed
|
|
// once we have robust end-to-end metadata handling.
|
|
// TODO(b/265059672): Remove when end-to-end stack trace handling is in place
|
|
class DummyStackTrace : public AbstractStackTrace {
|
|
absl::Span<StackFrame const> ToFrames() const override { return frames_; }
|
|
std::vector<StackFrame> ToUncachedFrames() const override { return frames_; }
|
|
|
|
StackFrame LastUserFrame() const override { return frames_.back(); }
|
|
|
|
std::vector<StackFrame> GetUserFrames(int /*limit*/) const override {
|
|
return frames_;
|
|
}
|
|
|
|
std::string ToString(const TracePrintingOptions& opts) const override {
|
|
auto frame = LastUserFrame();
|
|
return absl::StrCat(frame.file_name, ":", frame.line_number, ":",
|
|
frame.function_name);
|
|
}
|
|
|
|
std::vector<StackFrame> frames_{
|
|
StackFrame({"dummy_file_name", 10, "dummy_function_name"})};
|
|
};
|
|
|
|
namespace {
|
|
const xla::HloInstructionProto* FindInstructionById(
|
|
const xla::HloComputationProto& computation, int64_t id) {
|
|
auto iter =
|
|
absl::c_find_if(computation.instructions(),
|
|
[id](const xla::HloInstructionProto& instruction) {
|
|
return instruction.id() == id;
|
|
});
|
|
if (iter == computation.instructions().end()) {
|
|
return nullptr;
|
|
}
|
|
return &(*iter);
|
|
}
|
|
|
|
bool ShouldAddPrecisionToInstruction(
|
|
const xla::HloInstructionProto& instruction,
|
|
const xla::HloComputationProto& computation) {
|
|
static constexpr std::array<absl::string_view, 2> kOpsPossiblyUsingTF32 = {
|
|
"dot", "convolution"};
|
|
if (!absl::c_linear_search(kOpsPossiblyUsingTF32, instruction.opcode())) {
|
|
return false;
|
|
}
|
|
if (instruction.shape().element_type() == xla::F32) {
|
|
return true;
|
|
}
|
|
return absl::c_any_of(instruction.operand_ids(), [&](int64_t operand_id) {
|
|
const xla::HloInstructionProto* operand =
|
|
FindInstructionById(computation, operand_id);
|
|
return operand && operand->shape().element_type() == xla::F32;
|
|
});
|
|
}
|
|
|
|
// Add precisions configs to the HLO module to avoid TensorFloat32 computations
|
|
// in XLA.
|
|
//
|
|
// Some operations, such as Einsum are converted through MlirXlaOpKernel, which
|
|
// doesn't set the precisions, so we set them all here.
|
|
//
|
|
void IncreasePrecisionsToAvoidTF32(xla::HloModuleProto& module) {
|
|
xla::PrecisionConfig precision_config;
|
|
precision_config.add_operand_precision(xla::PrecisionConfig::HIGHEST);
|
|
precision_config.add_operand_precision(xla::PrecisionConfig::HIGHEST);
|
|
|
|
for (xla::HloComputationProto& computation : *module.mutable_computations()) {
|
|
for (xla::HloInstructionProto& instruction :
|
|
*computation.mutable_instructions()) {
|
|
if (ShouldAddPrecisionToInstruction(instruction, computation)) {
|
|
*instruction.mutable_precision_config() = precision_config;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
} // namespace
|
|
|
|
absl::Status XlaCompiler::CompileGraph(
|
|
const XlaCompiler::CompileOptions& options, const std::string& name,
|
|
std::unique_ptr<Graph> graph, absl::Span<const XlaCompiler::Argument> args,
|
|
CompilationResult* result) {
|
|
VLOG(1) << "Executing graph symbolically to populate XlaBuilder.: " << name;
|
|
if (VLOG_IS_ON(2) || DEBUG_DATA_DUMPER()->ShouldDump(name, kDebugGroupMain)) {
|
|
VLOG(2) << "XlaCompiler::CompileGraph: "
|
|
<< DumpGraphToFile(absl::StrCat("xla_compile_graph_", name), *graph,
|
|
flib_runtime_->GetFunctionLibraryDefinition());
|
|
}
|
|
|
|
DummyStackTrace stack_trace;
|
|
for (auto node : graph->nodes()) {
|
|
if (node->GetStackTrace() == nullptr) {
|
|
node->SetStackTrace(std::make_shared<DummyStackTrace>(stack_trace));
|
|
}
|
|
}
|
|
|
|
TF_RETURN_IF_ERROR(PropagateConstIntoFunctionalNodes(
|
|
graph.get(), options_.flib_def, local_flib_def_.get()));
|
|
TF_RETURN_IF_ERROR(RearrangeFunctionArguments(
|
|
[this](const NameAttrList& function, const FunctionBody** fbody) {
|
|
return FindFunctionBody(function, fbody);
|
|
},
|
|
graph.get(), local_flib_def_.get(),
|
|
pflr_->GetFunctionLibraryDefinition()));
|
|
|
|
// Report the error here if initialization failed.
|
|
TF_RETURN_IF_ERROR(initialization_status_);
|
|
|
|
// Detect invalid nodes.
|
|
// FunctionalizeControlFlow may remove some nodes from the graph.
|
|
TF_RETURN_IF_ERROR(ValidateGraph(graph.get(), *options_.flib_def,
|
|
options_.device_type, name));
|
|
auto builder = std::make_unique<xla::XlaBuilder>(name);
|
|
XlaContext* context = new XlaContext(this, builder.get(), graph.get());
|
|
core::ScopedUnref context_unref(context);
|
|
|
|
std::vector<XlaCompiler::Argument> real_args(args.begin(), args.end());
|
|
int token_input_index = -1;
|
|
std::unique_ptr<xla::XlaOp> token_output;
|
|
if (options.add_token_input_output) {
|
|
// Add extra token input.
|
|
token_input_index = real_args.size();
|
|
|
|
XlaCompiler::Argument token_arg;
|
|
token_arg.kind = XlaCompiler::Argument::kToken;
|
|
real_args.push_back(token_arg);
|
|
}
|
|
|
|
std::map<int, xla::OpSharding> arg_shardings;
|
|
std::map<int, xla::OpSharding> retval_shardings;
|
|
TF_ASSIGN_OR_RETURN(std::tie(arg_shardings, retval_shardings),
|
|
ComputeArgAndRetvalShardings(*graph));
|
|
|
|
std::vector<XlaExpression> arg_expressions;
|
|
std::string args_tuple_sdy_sharding;
|
|
TF_RETURN_IF_ERROR(
|
|
BuildArguments(*graph, real_args, options.use_tuple_arg, builder.get(),
|
|
context, arg_shardings, &arg_expressions,
|
|
&result->input_mapping, &result->xla_input_shapes,
|
|
args_tuple_sdy_sharding, options.is_entry_computation));
|
|
context->set_args(std::move(arg_expressions));
|
|
|
|
PushNodeTokenMapping();
|
|
// Use std::set instead of std::unordered_set to ensure determinism.
|
|
std::set<std::string> output_node_token_inputs;
|
|
if (token_input_index != -1) {
|
|
// Original token comes from input.
|
|
auto arg_expression = context->args()[token_input_index];
|
|
TF_RETURN_IF_ERROR(
|
|
SetNodeToken(kXlaTokenArgNodeName, arg_expression.handle()));
|
|
|
|
// Calculate token inputs for output token.
|
|
output_node_token_inputs = CalculateTokenInputsForOutputToken(*graph);
|
|
|
|
// If there's no side-effecting op in the graph, use token input as token
|
|
// output.
|
|
if (output_node_token_inputs.empty()) {
|
|
output_node_token_inputs.insert(kXlaTokenArgNodeName);
|
|
}
|
|
} else if (options.is_entry_computation) {
|
|
// Original token is manually created.
|
|
if (HasSideEffectingNodes(*graph)) {
|
|
TF_RETURN_IF_ERROR(
|
|
SetNodeToken(kXlaTokenArgNodeName, xla::CreateToken(builder.get())));
|
|
}
|
|
}
|
|
|
|
absl::Status execute_status = ExecuteGraph(context, std::move(graph), device_,
|
|
flib_runtime_, NextStepId());
|
|
if (!execute_status.ok()) {
|
|
VLOG(1) << "Failed executing graph " << name;
|
|
return execute_status;
|
|
}
|
|
if (token_input_index != -1) {
|
|
// Add extra token output.
|
|
std::vector<xla::XlaOp> token_inputs;
|
|
for (const auto& node_name : output_node_token_inputs) {
|
|
auto token_or = GetNodeToken(node_name);
|
|
TF_RETURN_IF_ERROR(token_or.status());
|
|
token_inputs.push_back(token_or.value());
|
|
}
|
|
token_output = std::make_unique<xla::XlaOp>(
|
|
xla::AfterAll(builder.get(), token_inputs));
|
|
}
|
|
TF_RETURN_IF_ERROR(PopNodeTokenMapping());
|
|
|
|
int num_nonconst_outputs;
|
|
int num_computation_outputs;
|
|
result->computation = std::make_shared<xla::XlaComputation>();
|
|
result->outputs.resize(context->retvals().size());
|
|
std::vector<XlaExpression> retvals = context->retvals();
|
|
ConvertConstantsToExpressions(builder.get(),
|
|
absl::Span<XlaExpression>(retvals));
|
|
XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns{
|
|
UseNoPreferenceLayoutFn(), IdentityShapeRepresentationFn()};
|
|
std::string result_tuple_sdy_sharding;
|
|
TF_RETURN_IF_ERROR(BuildComputation(
|
|
real_args, retvals, arg_shardings, retval_shardings, context->resources(),
|
|
std::move(token_output),
|
|
options.is_entry_computation ? options_.shape_determination_fns
|
|
: shape_determination_fns,
|
|
options.is_entry_computation,
|
|
options.return_updated_values_for_all_resources,
|
|
options.always_return_tuple, options.use_tuple_arg,
|
|
options.alias_resource_update, builder.get(), result->computation.get(),
|
|
&num_computation_outputs, &num_nonconst_outputs, &result->outputs,
|
|
&result->resource_updates, &result->xla_output_shape,
|
|
result->input_mapping, result_tuple_sdy_sharding,
|
|
options_.use_shardy_partitioner));
|
|
|
|
// Add Shardy sharding attributes for tuple arguments and results to the
|
|
// module's frontend attributes. These are stored at the module level
|
|
// because instruction-level frontend attributes are lost when tuples are
|
|
// flattened during HLO to StableHLO conversion.
|
|
auto setModuleFrontendAttribute = [&](absl::string_view key,
|
|
absl::string_view value) {
|
|
result->computation->mutable_proto()
|
|
->mutable_frontend_attributes()
|
|
->mutable_map()
|
|
->try_emplace(key, value);
|
|
};
|
|
if (!args_tuple_sdy_sharding.empty()) {
|
|
setModuleFrontendAttribute(
|
|
xla::sdy::toStringView(xla::sdy::kInTupleShardings),
|
|
args_tuple_sdy_sharding);
|
|
setModuleFrontendAttribute(xla::sdy::toStringView(xla::sdy::kUseTupleArgs),
|
|
"True");
|
|
}
|
|
if (!result_tuple_sdy_sharding.empty()) {
|
|
setModuleFrontendAttribute(
|
|
xla::sdy::toStringView(xla::sdy::kOutTupleShardings),
|
|
result_tuple_sdy_sharding);
|
|
}
|
|
|
|
for (const auto& [key, send] : host_compute_sends_) {
|
|
auto* d2h = result->host_compute_metadata.add_device_to_host();
|
|
*d2h = send;
|
|
|
|
for (int i = 0; i < d2h->metadata_size(); ++i) {
|
|
const std::string channel_name =
|
|
GetDeviceToHostChannelName(d2h->key(), i);
|
|
xla::ChannelHandle handle;
|
|
TF_RETURN_IF_ERROR(GetDeviceToHostChannelHandle(channel_name, &handle));
|
|
d2h->mutable_metadata(i)->set_channel_id(handle.handle());
|
|
}
|
|
}
|
|
for (const auto& [key, recv] : host_compute_recvs_) {
|
|
auto* h2d = result->host_compute_metadata.add_host_to_device();
|
|
*h2d = recv;
|
|
|
|
for (int i = 0; i < h2d->metadata_size(); ++i) {
|
|
const std::string channel_name =
|
|
GetHostToDeviceChannelName(h2d->key(), i);
|
|
xla::ChannelHandle handle;
|
|
TF_RETURN_IF_ERROR(GetHostToDeviceChannelHandle(channel_name, &handle));
|
|
h2d->mutable_metadata(i)->set_channel_id(handle.handle());
|
|
}
|
|
}
|
|
|
|
if (!tsl::tensor_float_32_execution_enabled()) {
|
|
IncreasePrecisionsToAvoidTF32(*result->computation->mutable_proto());
|
|
}
|
|
|
|
VLOG(2) << "Outputs: total: " << context->retvals().size()
|
|
<< " nonconstant: " << num_nonconst_outputs;
|
|
VLOG(2) << "XLA output shape: "
|
|
<< xla::ShapeUtil::HumanStringWithLayout(result->xla_output_shape);
|
|
result->collective_info = context->GetCollectiveInfo();
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
xla::ChannelHandle XlaCompiler::NewChannel(
|
|
xla::ChannelHandle::ChannelType type) {
|
|
xla::ChannelHandle new_handle;
|
|
absl::MutexLock lock(channel_mutex_);
|
|
// Create a new channel handle with a unique value.
|
|
new_handle.set_handle(next_channel_++);
|
|
new_handle.set_type(type);
|
|
return new_handle;
|
|
}
|
|
|
|
absl::Status XlaCompiler::GetChannelHandle(const std::string& key,
|
|
xla::ChannelHandle* channel) {
|
|
auto result = channels_.emplace(key, xla::ChannelHandle());
|
|
if (result.second) {
|
|
result.first->second = NewChannel(xla::ChannelHandle::DEVICE_TO_DEVICE);
|
|
}
|
|
*channel = result.first->second;
|
|
VLOG(1) << "Channel: " << key << " " << channel->DebugString();
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
absl::Status XlaCompiler::GetHostToDeviceChannelHandle(
|
|
const std::string& key, xla::ChannelHandle* channel) {
|
|
auto result = channels_.emplace(key, xla::ChannelHandle());
|
|
if (result.second) {
|
|
result.first->second = NewChannel(xla::ChannelHandle::HOST_TO_DEVICE);
|
|
}
|
|
*channel = result.first->second;
|
|
VLOG(1) << "Host to device channel: " << key << " " << channel->DebugString();
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
absl::Status XlaCompiler::GetDeviceToHostChannelHandle(
|
|
const std::string& key, xla::ChannelHandle* channel) {
|
|
auto result = channels_.emplace(key, xla::ChannelHandle());
|
|
if (result.second) {
|
|
result.first->second = NewChannel(xla::ChannelHandle::DEVICE_TO_HOST);
|
|
}
|
|
*channel = result.first->second;
|
|
VLOG(1) << "Device to host channel: " << key << " " << channel->DebugString();
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
namespace {
|
|
|
|
void SetTransfer(const std::string& key, absl::Span<const DataType> types,
|
|
absl::Span<const TensorShape> shapes,
|
|
tf2xla::HostTransferMetadata* transfer) {
|
|
transfer->set_key(key);
|
|
CHECK(types.size() == shapes.size());
|
|
for (int i = 0, end = types.size(); i < end; ++i) {
|
|
tf2xla::TensorMetadata* metadata = transfer->add_metadata();
|
|
metadata->set_type(types[i]);
|
|
shapes[i].AsProto(metadata->mutable_shape());
|
|
}
|
|
}
|
|
|
|
} // namespace
|
|
|
|
absl::Status XlaCompiler::SetDeviceToHostMetadata(
|
|
const std::string& key, absl::Span<const DataType> types,
|
|
absl::Span<const TensorShape> shapes) {
|
|
if (host_compute_sends_.find(key) != host_compute_sends_.end()) {
|
|
tf2xla::HostTransferMetadata& existing_transfer = host_compute_sends_[key];
|
|
tf2xla::HostTransferMetadata new_transfer;
|
|
SetTransfer(key, types, shapes, &new_transfer);
|
|
if (xla::protobuf_util::HaveSameSerialization(existing_transfer,
|
|
new_transfer)) {
|
|
return absl::OkStatus();
|
|
} else {
|
|
return absl::InvalidArgumentError(absl::StrCat(
|
|
"Duplicate calls to SetDeviceToHostMetadata with key ", key));
|
|
}
|
|
}
|
|
tf2xla::HostTransferMetadata& transfer = host_compute_sends_[key];
|
|
SetTransfer(key, types, shapes, &transfer);
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
absl::Status XlaCompiler::GetDeviceToHostShapes(
|
|
const std::string& key, std::vector<TensorShape>* shapes) const {
|
|
const auto iter = host_compute_sends_.find(key);
|
|
if (iter == host_compute_sends_.end()) {
|
|
return absl::InvalidArgumentError(
|
|
absl::StrCat("No host compute send shapes registered for key ", key));
|
|
}
|
|
shapes->clear();
|
|
for (int i = 0; i < iter->second.metadata_size(); ++i) {
|
|
TensorShape shape(iter->second.metadata(i).shape());
|
|
shapes->push_back(shape);
|
|
}
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
absl::Status XlaCompiler::SetHostToDeviceMetadata(
|
|
const std::string& key, absl::Span<const DataType> types,
|
|
absl::Span<const TensorShape> shapes) {
|
|
if (host_compute_recvs_.find(key) != host_compute_recvs_.end()) {
|
|
tf2xla::HostTransferMetadata& existing_transfer = host_compute_recvs_[key];
|
|
tf2xla::HostTransferMetadata new_transfer;
|
|
SetTransfer(key, types, shapes, &new_transfer);
|
|
if (xla::protobuf_util::HaveSameSerialization(existing_transfer,
|
|
new_transfer)) {
|
|
return absl::OkStatus();
|
|
} else {
|
|
return absl::InvalidArgumentError(absl::StrCat(
|
|
"Duplicate calls to SetHostToDeviceMetadata with key ", key));
|
|
}
|
|
}
|
|
tf2xla::HostTransferMetadata& transfer = host_compute_recvs_[key];
|
|
SetTransfer(key, types, shapes, &transfer);
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
absl::Status XlaCompiler::GetHostComputeControlDependency(
|
|
const std::string& host_compute_name, xla::XlaOp* handle) {
|
|
const auto iter = host_compute_control_output_.find(host_compute_name);
|
|
if (iter == host_compute_control_output_.end()) {
|
|
return absl::InvalidArgumentError(
|
|
absl::StrCat("No registered control handle for host compute Op '",
|
|
host_compute_name, "'"));
|
|
} else {
|
|
*handle = iter->second;
|
|
}
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
absl::Status XlaCompiler::SetHostComputeControlDependency(
|
|
const std::string& host_compute_name, const xla::XlaOp handle) {
|
|
if (host_compute_control_output_.find(host_compute_name) !=
|
|
host_compute_control_output_.end()) {
|
|
return absl::InvalidArgumentError(absl::StrCat(
|
|
"Duplicate control handles registered for host compute Op ",
|
|
host_compute_name));
|
|
}
|
|
host_compute_control_output_[host_compute_name] = handle;
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
void XlaCompiler::PushNodeTokenMapping() {
|
|
node_token_mapping_stack_.emplace(std::map<std::string, xla::XlaOp>{});
|
|
}
|
|
|
|
absl::Status XlaCompiler::PopNodeTokenMapping() {
|
|
if (node_token_mapping_stack_.empty()) {
|
|
return absl::FailedPreconditionError(
|
|
"Calling PopNodeTokenMapping() when node_token_mapping_stack_ is "
|
|
"empty.");
|
|
}
|
|
node_token_mapping_stack_.pop();
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
absl::Status XlaCompiler::SetNodeToken(const std::string& node_name,
|
|
const xla::XlaOp op) {
|
|
if (node_token_mapping_stack_.empty()) {
|
|
return absl::FailedPreconditionError(
|
|
"Calling SetNodeToken() when node_token_mapping_stack_ is "
|
|
"empty.");
|
|
}
|
|
auto insert_result = node_token_mapping_stack_.top().insert({node_name, op});
|
|
if (!insert_result.second) {
|
|
return absl::FailedPreconditionError(
|
|
absl::StrCat("Token mapping already exists for node ", node_name));
|
|
}
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
absl::StatusOr<xla::XlaOp> XlaCompiler::GetNodeToken(
|
|
const std::string& node_name) {
|
|
if (node_token_mapping_stack_.empty()) {
|
|
return absl::FailedPreconditionError(
|
|
"Calling GetNodeToken() when node_token_mapping_stack_ is "
|
|
"empty.");
|
|
}
|
|
auto iter = node_token_mapping_stack_.top().find(node_name);
|
|
if (iter == node_token_mapping_stack_.top().end()) {
|
|
return absl::FailedPreconditionError(
|
|
absl::StrCat("Cannot find token mapping for node ", node_name));
|
|
}
|
|
return iter->second;
|
|
}
|
|
|
|
} // namespace tensorflow
|