wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 5 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/compiler/jit/BUILD
+2115
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/jit/build_xla_ops_pass.cc
+624
View File
@@ -0,0 +1,624 @@
/* 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/jit/build_xla_ops_pass.h"
#include "absl/algorithm/container.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.h"
#include "tensorflow/cc/framework/ops.h"
#include "tensorflow/cc/framework/scope_internal.h"
#include "tensorflow/cc/ops/array_ops.h"
#include "tensorflow/cc/ops/const_op.h"
#include "tensorflow/cc/ops/control_flow_ops.h"
#include "tensorflow/cc/ops/functional_ops.h"
#include "tensorflow/cc/ops/logging_ops.h"
#include "tensorflow/compiler/jit/defs.h"
#include "tensorflow/compiler/jit/device_util.h"
#include "tensorflow/compiler/jit/encapsulate_subgraphs_pass.h"
#include "tensorflow/compiler/jit/flags.h"
#include "tensorflow/compiler/jit/xla_cluster_util.h"
#include "tensorflow/compiler/tf2xla/cc/ops/xla_jit_ops.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "xla/status_macros.h"
#include "tensorflow/core/common_runtime/function.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/common_runtime/optimization_registry.h"
#include "tensorflow/core/framework/graph_def_util.h"
#include "tensorflow/core/framework/memory_types.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/graph/algorithm.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/hash/hash.h"
#include "tensorflow/core/public/version.h"
#include "tensorflow/core/util/dump_graph.h"
namespace tensorflow {
namespace {
struct DebuggingOpts {
// If true, insert Print nodes to print every output from an XLA cluster.
bool print_outputs;
// If true, insert CheckNumerics nodes for every floating point typed input to
// an XLA cluster.
bool check_input_numerics;
// If true, insert CheckNumerics nodes for every floating point typed output
// from an XLA cluster.
bool check_output_numerics;
};
void MoveOutgoingEdges(Graph* g, Node* old_node, Node* new_node) {
std::vector<const Edge*> out_edges(old_node->out_edges().begin(),
old_node->out_edges().end());
for (const Edge* edge : out_edges) {
// TODO(sanjoy): This does not update NodeDef inputs. To be able to update
// NodeDef inputs we first need to fix encapsulate_subgraphs_pass to fix up
// the NodeDef inputs to the function call nodes.
g->AddEdge(new_node, edge->src_output(), edge->dst(), edge->dst_input());
g->RemoveEdge(edge);
}
}
// Returns a data value that is dead iff `control` is dead.
Output ControlToData(const Scope& scope, Node* control) {
// The choice of data type here is important.
//
// We implement a "control merge", which is a control edge that is alive if
// either of two nodes (denoted as A and B below) are alive, in the following
// manner:
//
// A --ctrl--> Const0 --data--> Merge --data--> Identity
// ^ |
// | ctrl
// B --ctrl--> Const1 --data-----+ |
// v
// ***
//
// where *** denotes the merged control output.
//
// We want everything starting from Const{0/1} to Identity to either wholly
// live on the host or wholly live on device so we need to pick a data type
// that is either consistently assigned to the device (e.g. float) or
// consistently assigned to the host (e.g. int32). We should *not* pick a
// data type that partly placed on the host and partly on the device
// (e.g. bool constants are placed on the device but bool Identity is placed
// on the host).
Output data = ops::Const(scope.WithOpName("ctrl_as_data"),
Tensor(DT_INT32, TensorShape({0})));
scope.graph()->AddControlEdge(control, data.node());
return Output(data.node());
}
// Returns an operation that can be control-depended on that is dead iff `data`
// is dead.
Operation DataToControl(const Scope& scope, Output data) {
return Operation(
ops::Identity(scope.WithOpName("data_as_ctrl"), data).node());
}
// Replaces each outgoing edge from `old_node` with a merge node that merges in
// the corresponding output from `new_node`.
void MergeOutgoingDataEdges(const Scope& s, Node* old_node, Node* new_node,
absl::string_view cluster_name,
const DebuggingOpts& debugging_opts) {
if (!s.status().ok()) {
return;
}
std::vector<Output> merged_outputs(old_node->num_outputs(), Output(nullptr));
std::vector<const Edge*> data_edges;
absl::c_copy_if(old_node->out_edges(), std::back_inserter(data_edges),
[](const Edge* e) { return !e->IsControlEdge(); });
for (const Edge* e : data_edges) {
int oidx = e->src_output();
Output merged_output = merged_outputs[oidx];
if (merged_output.node() == nullptr) {
Output new_output(new_node, oidx);
if (debugging_opts.print_outputs) {
std::string cpu_device = "/job:localhost/replica:0/task:0/device:CPU:0";
ops::Print print_op(s.WithOpName("print_", oidx)
.WithDevice(cpu_device)
.WithAssignedDevice(cpu_device),
new_output, {new_output},
ops::Print::Attrs{}
.Message(absl::StrCat("output ", oidx, " from ",
old_node->name(), " is "))
.FirstN(1000)
.Summarize(-1));
new_output = print_op;
}
if (debugging_opts.check_output_numerics &&
DataTypeIsFloating(new_output.type())) {
ops::CheckNumerics check_numerics_op(
s.WithOpName("check_output_", oidx)
.WithDevice(new_node->requested_device())
.WithAssignedDevice(new_node->assigned_device_name()),
new_output,
absl::StrCat("CheckNumerics failed for output ", oidx, "(",
new_output.name(), ") from cluster ", cluster_name));
new_output = check_numerics_op;
}
ops::_XlaMerge xla_merge_op(s.WithOpName("merge_oidx_", oidx),
Output(old_node, oidx), new_output);
merged_output = merged_outputs[oidx] = xla_merge_op.output;
}
Node* dst = e->dst();
int dst_idx = e->dst_input();
s.graph()->RemoveEdge(e);
s.graph()->AddEdge(merged_output.node(), merged_output.index(), dst,
dst_idx);
}
}
// Replaces each control successor of `old_node` to execute whenever either
// `old_node` or `new_node` is executed.
void MergeOutgoingControlEdges(const Scope& s, Node* old_node, Node* new_node) {
if (!s.status().ok()) {
return;
}
std::vector<const Edge*> ctrl_edges;
absl::c_copy_if(old_node->out_edges(), std::back_inserter(ctrl_edges),
[](const Edge* e) { return e->IsControlEdge(); });
if (ctrl_edges.empty()) {
return;
}
if (ctrl_edges.size() == 1 && ctrl_edges.front()->dst()->IsSink()) {
// Avoid creating a Merge node if we can just add an edge to _SINK
// instead.
s.graph()->AddControlEdge(new_node, s.graph()->sink_node());
return;
}
// We can't merge control edges directly so we instead first "convert" them to
// normal values that can be merged, merge the values and then "convert" the
// merged value back into control.
//
// NB! We need to copy out the outgoing control edges before constructing
// old_ctrl_as_data otherwise the control edge from old_node to the constant
// in ControlToData will be present in ctrl_edges.
Output old_ctrl_as_data = ControlToData(s, old_node);
Output new_ctrl_as_data = ControlToData(s, new_node);
ops::Merge ctrl_merge_as_data(s.WithOpName("ctrl_merge"),
{old_ctrl_as_data, new_ctrl_as_data});
Operation ctrl_merge = DataToControl(s, ctrl_merge_as_data.output);
for (const Edge* e : ctrl_edges) {
s.graph()->AddControlEdge(ctrl_merge.node(), e->dst());
s.graph()->RemoveControlEdge(e);
}
}
struct XlaClusterInfo {
std::vector<Output> constant_inputs;
std::vector<Output> non_constant_inputs;
std::vector<Output> resource_inputs;
NameAttrList function;
};
Output IncomingEdgeAsOutput(const Edge* e) {
return Output(e->src(), e->src_output());
}
absl::Status GetXlaClusterInfo(Node* n, XlaClusterInfo* result) {
int num_constant_inputs, num_resource_inputs;
TF_RETURN_IF_ERROR(
GetNodeAttr(n->attrs(), kXlaNumConstantArgsAttr, &num_constant_inputs));
TF_RETURN_IF_ERROR(
GetNodeAttr(n->attrs(), kXlaNumResourceArgsAttr, &num_resource_inputs));
if (num_constant_inputs < 0 || num_resource_inputs < 0 ||
num_constant_inputs + num_resource_inputs > n->num_inputs()) {
return absl::InvalidArgumentError(
"Invalid number of constant/resource arguments to XLA kernel.");
}
int num_non_constant_inputs =
n->num_inputs() - num_constant_inputs - num_resource_inputs;
std::vector<const Edge*> input_edges_vector;
TF_RETURN_IF_ERROR(n->input_edges(&input_edges_vector));
absl::Span<const Edge*> input_edges(input_edges_vector);
absl::c_transform(input_edges.subspan(0, num_constant_inputs),
std::back_inserter(result->constant_inputs),
IncomingEdgeAsOutput);
absl::c_transform(
input_edges.subspan(num_constant_inputs, num_non_constant_inputs),
std::back_inserter(result->non_constant_inputs), IncomingEdgeAsOutput);
absl::c_transform(
input_edges.subspan(num_constant_inputs + num_non_constant_inputs,
num_resource_inputs),
std::back_inserter(result->resource_inputs), IncomingEdgeAsOutput);
result->function.set_name(n->type_string());
*result->function.mutable_attr() = n->def().attr();
return absl::OkStatus();
}
absl::Status CopyIncomingControlEdges(Graph* g, Node* from, Node* to) {
for (const Edge* e : from->in_edges()) {
if (e->IsControlEdge()) {
g->AddControlEdge(e->src(), to);
}
}
return absl::OkStatus();
}
void RemoveAllIncomingControlEdges(Graph* g, Node* n) {
std::vector<const Edge*> incoming_ctrl_edges;
absl::c_copy_if(n->in_edges(), std::back_inserter(incoming_ctrl_edges),
[](const Edge* e) { return e->IsControlEdge(); });
for (const Edge* e : incoming_ctrl_edges) {
g->RemoveControlEdge(e);
}
}
// Returns true (into `result`) if a node placed on `device` must be compiled.
absl::Status DeviceRequiresCompilation(
const jit::DeviceInfoCache& device_info_cache, jit::DeviceId device,
bool* result) {
const XlaOpRegistry::DeviceRegistration* registration =
device_info_cache.GetCompilationDevice(device);
*result = registration->autoclustering_policy ==
XlaOpRegistry::AutoclusteringPolicy::kAlways;
return absl::OkStatus();
}
// Replaces `n` with a `PartitionedCall` op that calls the same function.
absl::StatusOr<Node*> ReplaceFunctionCallWithPartitionedCall(
const GraphOptimizationPassOptions& options,
const FunctionLibraryDefinition& flib_def, Node* n, Graph* g,
const NameAttrList& func, const Scope& root) {
std::string config_string =
options.session_options->config.SerializeAsString();
int input_count = absl::c_count_if(
n->in_edges(), [](const Edge* e) { return !e->IsControlEdge(); });
std::vector<Output> args(input_count);
for (const Edge* e : n->in_edges()) {
if (!e->IsControlEdge()) {
args[e->dst_input()] = Output(e->src(), e->src_output());
}
}
// In theory we can use PartitionedCall if the XLA cluster does not have any
// stateful operations. However, for now we choose to be conservative since
// we don't have any evidence that choosing a stateless partitioned call helps
// for performance.
ops::StatefulPartitionedCall call(
root.WithOpName("stateful_partitioned_call"), args, n->output_types(),
func, ops::StatefulPartitionedCall::Attrs{}.ConfigProto(config_string));
for (const Edge* e : n->in_edges()) {
if (e->IsControlEdge()) {
g->AddControlEdge(e->src(), call.operation.node());
}
}
std::vector<const Edge*> edges_to_delete;
for (const Edge* e : n->out_edges()) {
edges_to_delete.push_back(e);
if (e->IsControlEdge()) {
g->AddControlEdge(call.operation.node(), e->dst());
} else {
g->AddEdge(call.operation.node(), e->src_output(), e->dst(),
e->dst_input());
}
}
for (const Edge* e : edges_to_delete) {
g->RemoveEdge(e);
}
g->RemoveNode(n);
return call.operation.node();
}
absl::StatusOr<jit::DeviceId> InferDeviceForCluster(
jit::DeviceInfoCache* device_info_cache, Node* n,
const std::string& function_name,
const FunctionLibraryDefinition& flib_def) {
const FunctionDef* func_def = flib_def.Find(function_name);
TF_RET_CHECK(func_def) << "Could not find " << function_name;
jit::DeviceSet device_set;
for (const NodeDef& ndef : func_def->node_def()) {
VLOG(3) << ndef.DebugString();
if (!ndef.device().empty()) {
TF_ASSIGN_OR_RETURN(jit::DeviceId device_id,
device_info_cache->GetIdFor(ndef.device()));
device_set.Insert(device_id);
}
}
if (!n->assigned_device_name().empty()) {
// TODO(sanjoy): We need this because EncapsulateSubgraphsPass drops device
// assignment when constant folding. We should fix EncapsulateSubgraphsPass
// instead.
TF_ASSIGN_OR_RETURN(jit::DeviceId device_id,
device_info_cache->GetIdFor(n->assigned_device_name()));
device_set.Insert(device_id);
}
TF_ASSIGN_OR_RETURN(jit::DeviceId result,
PickDeviceForXla(*device_info_cache, device_set,
/*allow_mixing_unknown_and_cpu=*/true));
VLOG(2) << "For " << function_name << " PickDeviceForXla("
<< device_info_cache->DebugString(device_set) << ") -> "
<< device_info_cache->GetNameFor(result);
return result;
}
std::vector<Output> GetXlaRunArgs(const Scope& s,
const XlaClusterInfo& cluster_info,
const DebuggingOpts& debugging_opts) {
std::vector<Output> xla_run_args;
xla_run_args.reserve(cluster_info.non_constant_inputs.size() +
cluster_info.resource_inputs.size());
int input_idx = 0;
for (const Output& o : cluster_info.non_constant_inputs) {
if (debugging_opts.check_input_numerics && DataTypeIsFloating(o.type())) {
ops::CheckNumerics check_numerics_op(
s.WithOpName("check_input_", input_idx), o,
absl::StrCat("CheckNumerics failed for input ", input_idx, "(",
o.name(), ") into ", cluster_info.function.name()));
xla_run_args.push_back(check_numerics_op);
} else {
xla_run_args.push_back(o);
}
input_idx++;
}
absl::c_copy(cluster_info.resource_inputs, std::back_inserter(xla_run_args));
return xla_run_args;
}
absl::StatusOr<MemoryTypeVector> GetOutputMemoryTypes(const Scope& root,
Node* n) {
MemoryTypeVector input_mtypes, output_mtypes;
DeviceType device_type("");
TF_RETURN_IF_ERROR(
DeviceNameToDeviceType(n->assigned_device_name(), &device_type));
TF_RETURN_IF_ERROR(MemoryTypesForNode(root.graph()->op_registry(),
device_type, n->def(), &input_mtypes,
&output_mtypes));
return output_mtypes;
}
// Predicate INT32 typed inputs to `n` on the deadness of
// `predicate_as_control`.
//
// This is a performance optimization. Since INT32 arguments to a
// PartitionedCall are placed on the host, a producer that produces them on the
// device will incur a D2H copy, even if the PartitionedCall is not executed
// (i.e. even if we choose to execute the XLA compiled computation via _XlaRun).
// To prevent this, we add control dependencies to make the int32 input edges
// into the PartitionedCall dead. With this change the D2H copy only happens if
// the PartitionedCall is actually executed.
absl::Status PredicateInt32Inputs(const Scope& root, Node* n,
Operation predicate_as_control) {
std::vector<Output> int32_inputs;
std::vector<int> int32_inputs_input_idxs;
for (const Edge* e : n->in_edges()) {
if (e->IsControlEdge()) {
continue;
}
if (e->src()->output_type(e->src_output()) == DT_INT32) {
TF_ASSIGN_OR_RETURN(MemoryTypeVector source_output_mem_types,
GetOutputMemoryTypes(root, e->src()));
if (source_output_mem_types[e->src_output()] == DEVICE_MEMORY) {
int32_inputs.push_back(Output(e->src(), e->src_output()));
int32_inputs_input_idxs.push_back(e->dst_input());
}
}
}
if (int32_inputs.empty()) {
return absl::OkStatus();
}
// Create a single IdentityN that is dead if and only if
// `predicate_as_control` is dead.
//
// IdentityN is also special in that, unlike `Identity`, it does not place
// int32 inputs in host memory. Placing int32 inputs in host memory would
// defeat the purpose of adding this indirection.
ops::IdentityN identity_n(root.WithOpName("int32_id_n"), int32_inputs);
root.graph()->AddControlEdge(predicate_as_control.node(),
identity_n.operation.node());
for (int i = 0, end = int32_inputs.size(); i < end; i++) {
TF_RETURN_IF_ERROR(root.graph()->UpdateEdge(identity_n[i].node(), i, n,
int32_inputs_input_idxs[i]));
}
return absl::OkStatus();
}
absl::Status ReplaceNodeWithXlaCompileAndXlaRun(
jit::DeviceInfoCache* device_info_cache,
const GraphOptimizationPassOptions& options,
const FunctionLibraryDefinition& flib_def, bool lazy_compilation_enabled,
const DebuggingOpts& debugging_opts, Graph* g, Node* n) {
XlaClusterInfo cluster_info;
TF_RETURN_IF_ERROR(GetXlaClusterInfo(n, &cluster_info));
TF_ASSIGN_OR_RETURN(
jit::DeviceId device,
InferDeviceForCluster(device_info_cache, n, cluster_info.function.name(),
flib_def));
bool requires_compilation;
TF_RETURN_IF_ERROR(DeviceRequiresCompilation(*device_info_cache, device,
&requires_compilation));
if (!lazy_compilation_enabled) {
requires_compilation = true;
}
std::string device_name_str =
std::string(device_info_cache->GetNameFor(device));
absl::Status status;
Scope root = NewInternalScope(g, &status, /*refiner=*/nullptr)
.NewSubScope(n->name())
.WithDevice(n->requested_device())
.WithAssignedDevice(device_name_str);
ops::_XlaCompile xla_compile(root.WithOpName("xla_compile"),
/*constants=*/cluster_info.constant_inputs,
/*args=*/cluster_info.non_constant_inputs,
/*resources=*/cluster_info.resource_inputs,
/*must_compile=*/requires_compilation,
cluster_info.function);
bool has_ref_attr;
TF_RETURN_IF_ERROR(
GetNodeAttr(n->attrs(), kXlaHasReferenceVarsAttr, &has_ref_attr));
xla_compile.operation.node()->AddAttr(kXlaHasReferenceVarsAttr, has_ref_attr);
TF_RETURN_IF_ERROR(
CopyIncomingControlEdges(g, /*from=*/n, /*to=*/xla_compile.key.node()));
std::vector<Output> xla_run_args =
GetXlaRunArgs(root, cluster_info, debugging_opts);
if (requires_compilation) {
// "Strict" compilation: every _XlaCompile invocation must compile the
// cluster.
ops::_XlaRun xla_run(root.WithOpName("xla_run"), xla_run_args,
xla_compile.key, n->output_types());
MoveOutgoingEdges(g, /*old_node=*/n,
/*new_node=*/xla_run.operation.node());
g->RemoveNode(n);
} else {
// "Lazy" compilation: an _XlaCompile invocation may decide not to compile
// the cluster based on profitability heuristics.
// We generate the following graph:
//
// (use_tf_call, use_xla_run) =
// Switch(pred=xla_compile.compilation_successful,
// value=xla_compile.key)
//
// tf_call_outputs = cluster_N(..., ^use_tf_call)
// xla_run_outputs = _XlaRun(..., key=use_xla_run)
// outputs = Merge(tf_call_outputs, xla_run_outputs).
ops::Switch s(root.WithOpName("predicated_compilation_key"),
xla_compile.key, xla_compile.compilation_successful);
Output predicated_compilation_key = s.output_true;
Output inverse_predicated_compilation_key = s.output_false;
ops::_XlaRun xla_run(root.WithOpName("xla_run"), xla_run_args,
predicated_compilation_key, n->output_types());
MergeOutgoingControlEdges(root, /*old_node=*/n,
/*new_node=*/xla_run.operation.node());
MergeOutgoingDataEdges(root, /*old_node=*/n,
/*new_node=*/xla_run.operation.node(),
cluster_info.function.name(), debugging_opts);
TF_RETURN_IF_ERROR(root.status());
// We already have a TensorFlow function call into the cluster -- the
// original node we set out to rewrite. We just wire in the correct control
// deps and we're done.
RemoveAllIncomingControlEdges(g, n);
Operation inverse_predicate_as_control =
DataToControl(root, inverse_predicated_compilation_key);
g->AddControlEdge(inverse_predicate_as_control.node(), n);
n->ClearAttr(kXlaCompiledKernelAttr);
TF_ASSIGN_OR_RETURN(Node* const pco, ReplaceFunctionCallWithPartitionedCall(
options, flib_def, n, g,
cluster_info.function, root));
TF_RETURN_IF_ERROR(
PredicateInt32Inputs(root, pco, inverse_predicate_as_control));
}
return absl::OkStatus();
}
} // namespace
absl::Status BuildXlaOpsPass::Run(const GraphOptimizationPassOptions& options) {
Graph* graph = options.graph->get();
// Copy out the nodes we want to rewrite to avoid modifying the graph while we
// iterate on graph->op_nodes().
std::vector<Node*> xla_compiled_kernels;
absl::c_copy_if(graph->op_nodes(), std::back_inserter(xla_compiled_kernels),
[](const Node* n) {
if (n->IsSend() || n->IsRecv() || n->IsControlFlow()) {
return false;
}
// Only compile nodes that are marked for compilation by the
// compilation-marking pass (via 'attr_name').
return IsXlaCompiledKernel(*n);
});
bool lazy_compilation_enabled =
enable_lazy_compilation_
? *enable_lazy_compilation_
: GetBuildXlaOpsPassFlags()->tf_xla_enable_lazy_compilation;
jit::DeviceInfoCache device_info_cache;
const BuildXlaOpsPassFlags& flags = *GetBuildXlaOpsPassFlags();
DebuggingOpts debugging_opts;
debugging_opts.print_outputs = flags.tf_xla_print_cluster_outputs;
debugging_opts.check_input_numerics =
flags.tf_xla_check_cluster_input_numerics;
debugging_opts.check_output_numerics =
flags.tf_xla_check_cluster_output_numerics;
VLOG(1) << "print_outputs = " << debugging_opts.print_outputs;
VLOG(1) << "check_input_numerics = " << debugging_opts.check_input_numerics;
VLOG(1) << "check_output_numerics = " << debugging_opts.check_output_numerics;
for (Node* n : xla_compiled_kernels) {
TF_RETURN_IF_ERROR(ReplaceNodeWithXlaCompileAndXlaRun(
&device_info_cache, options, *options.flib_def,
lazy_compilation_enabled, debugging_opts, graph, n));
}
if (VLOG_IS_ON(1)) {
DumpGraphToFile("build_xla_ops", *graph, options.flib_def);
}
return absl::OkStatus();
}
} // namespace tensorflow
tensorflow/compiler/jit/build_xla_ops_pass.h
+45
View File
@@ -0,0 +1,45 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_JIT_BUILD_XLA_OPS_PASS_H_
#define TENSORFLOW_COMPILER_JIT_BUILD_XLA_OPS_PASS_H_
#include "absl/types/optional.h"
#include "tensorflow/core/common_runtime/optimization_registry.h"
#include "tensorflow/core/lib/core/status.h"
namespace tensorflow {
// Replaces TF function calls marked with `_XlaCompiledKernel` with _XlaCompile
// and _XlaRun nodes (which compile and launch, respectively, the corresponding
// HLO module).
class BuildXlaOpsPass : public GraphOptimizationPass {
public:
// If enable_lazy_compilation is not nullopt then *enable_lazy_compilation
// overrides --tf_xla_enable_lazy_compilation flag in deciding whether lazy
// compilation is enabled.
explicit BuildXlaOpsPass(
std::optional<bool> enable_lazy_compilation = std::nullopt)
: enable_lazy_compilation_(enable_lazy_compilation) {}
absl::Status Run(const GraphOptimizationPassOptions& options) override;
private:
std::optional<bool> enable_lazy_compilation_;
};
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_BUILD_XLA_OPS_PASS_H_
tensorflow/compiler/jit/build_xla_ops_pass_test.cc
+328
View File
@@ -0,0 +1,328 @@
/* Copyright 2018 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/jit/build_xla_ops_pass.h"
#include "tensorflow/cc/framework/ops.h"
#include "tensorflow/cc/ops/array_ops.h"
#include "tensorflow/cc/ops/resource_variable_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/compiler/jit/defs.h"
#include "tensorflow/compiler/jit/encapsulate_subgraphs_pass.h"
#include "tensorflow/compiler/jit/node_matchers.h"
#include "tensorflow/compiler/jit/test_util.h"
#include "tensorflow/core/common_runtime/device_factory.h"
#include "tensorflow/core/graph/algorithm.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/public/session_options.h"
namespace tensorflow {
namespace {
class BuildXlaOpsTest : public ::testing::Test {
protected:
void SetUp() override {
// This is needed to register the XLA_* devices.
CHECK(DeviceFactory::AddDevices(
SessionOptions(), "/job:localhost/replica:0/task:0", &devices_)
.ok());
}
private:
std::vector<std::unique_ptr<Device>> devices_;
};
using ::tensorflow::testing::FindNodeByName;
using ::tensorflow::testing::matchers::Attr;
using ::tensorflow::testing::matchers::CtrlDeps;
using ::tensorflow::testing::matchers::Inputs;
using ::tensorflow::testing::matchers::NodeWith;
using ::tensorflow::testing::matchers::Op;
using ::tensorflow::testing::matchers::Out;
using ::testing::_;
absl::Status BuildXlaOps(const Scope& s, const FunctionDefLibrary& fdef_lib,
std::unique_ptr<Graph>* result) {
auto graph = std::make_unique<Graph>(OpRegistry::Global());
TF_RETURN_IF_ERROR(s.ToGraph(graph.get()));
FunctionLibraryDefinition flib_def(graph->op_registry(), fdef_lib);
// Assign all nodes to the CPU device.
static const char* kCpuDevice = "/job:localhost/replica:0/task:0/cpu:0";
for (Node* n : graph->nodes()) {
if (n->requested_device().empty()) {
n->set_assigned_device_name(kCpuDevice);
} else {
n->set_assigned_device_name(n->requested_device());
}
}
FixupSourceAndSinkEdges(graph.get());
GraphOptimizationPassWrapper wrapper;
GraphOptimizationPassOptions opt_options =
wrapper.CreateGraphOptimizationPassOptions(&graph);
opt_options.flib_def = &flib_def;
BuildXlaOpsPass pass(/*enable_lazy_compilation=*/true);
TF_RETURN_IF_ERROR(pass.Run(opt_options));
VLOG(3) << graph->ToGraphDefDebug().DebugString();
*result = std::move(graph);
return absl::OkStatus();
}
absl::Status MakeXlaCompiledKernel(Graph* graph, const std::string& callee_name,
const std::string& node_name,
int num_constant_args, int num_resource_args,
Node** result) {
NodeDef call_node;
call_node.set_name(node_name);
call_node.set_op(callee_name);
AddNodeAttr(kXlaCompiledKernelAttr, true, &call_node);
AddNodeAttr(kXlaNumConstantArgsAttr, num_constant_args, &call_node);
AddNodeAttr(kXlaNumResourceArgsAttr, num_resource_args, &call_node);
TF_ASSIGN_OR_RETURN(*result, graph->AddNode(call_node));
return absl::OkStatus();
}
absl::Status MakeXlaCompiledKernel(Graph* graph, const std::string& callee_name,
const std::string& node_name,
Node** result) {
return MakeXlaCompiledKernel(graph, callee_name, node_name,
/*num_constant_args=*/0, /*num_resource_args=*/0,
result);
}
Node* MakeWrite(const Scope& scope, Output value_to_write,
const std::string& id) {
Output var_handle = ops::VarHandleOp(scope.WithOpName("Var_" + id), DT_FLOAT,
TensorShape({}));
ops::AssignVariableOp assign_op(scope.WithOpName("Assignee_" + id),
var_handle, value_to_write);
return assign_op.operation.node();
}
Node* MakeWrite(const Scope& scope, const std::string& id) {
return MakeWrite(
scope, ops::Const(scope.WithOpName("ValueToAssign" + id), 1.0f), id);
}
FunctionDefLibrary CreateFunctionDefLibWithConstFunction(
const std::string& name) {
FunctionDefLibrary fdef_lib;
FunctionDef func = FunctionDefHelper::Create(
/*function_name=*/name, /*in_def=*/{}, /*out_def=*/{"out: float"},
/*attr_def*/
{}, /*node_def=*/{FunctionDefHelper::Const("one", 1.0f)},
/*ret_def=*/{{"out", "out:output:0"}});
*fdef_lib.add_function() = std::move(func);
return fdef_lib;
}
TEST_F(BuildXlaOpsTest, ControlDepsPreserved) {
const char* kXlaDeviceName = "/job:worker/replica:0/task:0/device:XLA_CPU:0";
Scope root = Scope::NewRootScope().WithDevice(kXlaDeviceName).ExitOnError();
FunctionDefLibrary fdef_lib =
CreateFunctionDefLibWithConstFunction("cluster_0");
TF_ASSERT_OK(root.graph()->AddFunctionLibrary(fdef_lib));
Node* call;
TF_ASSERT_OK(MakeXlaCompiledKernel(root.graph(), "cluster_0", "C", &call));
call->AddAttr(kXlaHasReferenceVarsAttr, false);
call->set_requested_device(kXlaDeviceName);
Node* write_op = MakeWrite(root, "write");
write_op->AddAttr(kXlaHasReferenceVarsAttr, false);
root.graph()->AddControlEdge(call, write_op);
std::unique_ptr<Graph> graph;
TF_ASSERT_OK(BuildXlaOps(root, fdef_lib, &graph));
Node* write_op_new = FindNodeByName(graph.get(), write_op->name());
ASSERT_NE(write_op_new, nullptr);
EXPECT_THAT(write_op_new, NodeWith(CtrlDeps(NodeWith(Op("_XlaRun")))));
}
TEST_F(BuildXlaOpsTest, CleanFailureOnBogusAttr) {
Scope root = Scope::NewRootScope().ExitOnError();
FunctionDefLibrary fdef_lib =
CreateFunctionDefLibWithConstFunction("cluster_0");
TF_ASSERT_OK(root.graph()->AddFunctionLibrary(fdef_lib));
Node* call;
TF_ASSERT_OK(
MakeXlaCompiledKernel(root.graph(), "cluster_0", "C", 100, 100, &call));
Node* write_op = MakeWrite(root, "write");
root.graph()->AddControlEdge(call, write_op);
std::unique_ptr<Graph> graph;
absl::Status failure_status = BuildXlaOps(root, fdef_lib, &graph);
ASSERT_FALSE(failure_status.ok());
EXPECT_EQ(failure_status.code(), error::INVALID_ARGUMENT);
}
TEST_F(BuildXlaOpsTest, OnNonXlaDevice) {
Scope root = Scope::NewRootScope().ExitOnError();
FunctionDefLibrary fdef_lib =
CreateFunctionDefLibWithConstFunction("cluster_0");
TF_ASSERT_OK(root.graph()->AddFunctionLibrary(fdef_lib));
Node* call;
TF_ASSERT_OK(MakeXlaCompiledKernel(root.graph(), "cluster_0", "C", &call));
TF_ASSERT_OK(root.DoShapeInference(call));
call->AddAttr(kXlaHasReferenceVarsAttr, false);
Node* write_op = MakeWrite(root, Output(call), "write_result");
write_op->AddAttr(kXlaHasReferenceVarsAttr, false);
auto xla_compile = NodeWith(Op("_XlaCompile"), Attr("must_compile", false));
auto predicated_compilation_key =
NodeWith(Op("Switch"), Inputs(Out(0, xla_compile), Out(1, xla_compile)));
auto xla_run =
NodeWith(Op("_XlaRun"), Inputs(Out(1, predicated_compilation_key)));
auto tf_call =
NodeWith(Op("StatefulPartitionedCall"),
CtrlDeps(NodeWith(Op("Identity"),
Inputs(Out(0, predicated_compilation_key)))));
auto merge = NodeWith(Op("_XlaMerge"), Inputs(Out(tf_call), Out(xla_run)));
auto assign_var = NodeWith(Op("AssignVariableOp"), Inputs(_, Out(merge)));
std::unique_ptr<Graph> graph;
TF_ASSERT_OK(BuildXlaOps(root, fdef_lib, &graph));
Node* write_op_new = FindNodeByName(graph.get(), write_op->name());
ASSERT_NE(write_op_new, nullptr);
EXPECT_THAT(write_op_new, assign_var);
}
TEST_F(BuildXlaOpsTest, OnXlaDevice) {
const char* kXlaDeviceName = "/job:worker/replica:0/task:0/device:XLA_CPU:0";
Scope root = Scope::NewRootScope().WithDevice(kXlaDeviceName).ExitOnError();
FunctionDefLibrary fdef_lib =
CreateFunctionDefLibWithConstFunction("cluster_0");
TF_ASSERT_OK(root.graph()->AddFunctionLibrary(fdef_lib));
Node* call;
TF_ASSERT_OK(MakeXlaCompiledKernel(root.graph(), "cluster_0", "C", &call));
call->set_requested_device(kXlaDeviceName);
TF_ASSERT_OK(root.DoShapeInference(call));
call->AddAttr(kXlaHasReferenceVarsAttr, false);
Node* write_op = MakeWrite(root, Output(call), "write_result");
write_op->AddAttr(kXlaHasReferenceVarsAttr, false);
std::unique_ptr<Graph> graph;
TF_ASSERT_OK(BuildXlaOps(root, fdef_lib, &graph));
auto xla_op =
NodeWith(Op("_XlaRun"), Inputs(Out(NodeWith(Op("_XlaCompile")))));
auto assign_var =
NodeWith(Op("AssignVariableOp"), Inputs(Out(NodeWith()), Out(xla_op)));
Node* write_op_new = FindNodeByName(graph.get(), write_op->name());
ASSERT_NE(write_op_new, nullptr);
EXPECT_THAT(write_op_new, assign_var);
}
TEST_F(BuildXlaOpsTest, NoExtraMergeForEdgeToSink) {
Scope root = Scope::NewRootScope().ExitOnError();
FunctionDefLibrary fdef_lib =
CreateFunctionDefLibWithConstFunction("cluster_0");
TF_ASSERT_OK(root.graph()->AddFunctionLibrary(fdef_lib));
Node* call;
TF_ASSERT_OK(MakeXlaCompiledKernel(root.graph(), "cluster_0", "C", &call));
call->AddAttr(kXlaHasReferenceVarsAttr, false);
std::unique_ptr<Graph> graph;
TF_ASSERT_OK(BuildXlaOps(root, fdef_lib, &graph));
Node* sink_node = graph->sink_node();
EXPECT_THAT(sink_node,
NodeWith(CtrlDeps(NodeWith(Op("_XlaRun")),
NodeWith(Op("StatefulPartitionedCall")),
NodeWith(Op("NoOp")))));
}
#ifdef GOOGLE_CUDA
FunctionDefLibrary CreateFunctionDefLibWithInt32Input(const std::string& name) {
FunctionDefLibrary fdef_lib;
FunctionDef func = FunctionDefHelper::Create(
/*function_name=*/name, /*in_def=*/{"in: int32"},
/*out_def=*/{"out: int32"},
/*attr_def=*/{}, /*node_def=*/{{{"out"}, "Identity", {"in"}}},
/*ret_def=*/{{"out", "out:output:0"}});
*fdef_lib.add_function() = std::move(func);
return fdef_lib;
}
// This tests a rewrite that only makes sense and is active in a CUDA-enabled
// build. Specifically we check that we insert an IdentityN op to avoid extra
// device-to-host copies.
TEST_F(BuildXlaOpsTest, NoDeviceToHostCopiesForClustersWithInt32Inputs) {
const char* kXlaDeviceName = "/job:worker/replica:0/task:0/device:GPU:0";
Scope root = Scope::NewRootScope()
.WithDevice(kXlaDeviceName)
.WithAssignedDevice(kXlaDeviceName)
.ExitOnError();
FunctionDefLibrary fdef_lib =
CreateFunctionDefLibWithInt32Input("cluster_int32");
TF_ASSERT_OK(root.graph()->AddFunctionLibrary(fdef_lib));
Node* call;
TF_ASSERT_OK(
MakeXlaCompiledKernel(root.graph(), "cluster_int32", "C", &call));
call->set_requested_device(kXlaDeviceName);
call->AddAttr(kXlaHasReferenceVarsAttr, false);
auto var =
ops::VarHandleOp(root.WithOpName("var"), DT_INT32, TensorShape({}));
auto int32_on_device =
ops::ReadVariableOp(root.WithOpName("int32_on_device"), var, DT_INT32);
root.graph()->AddEdge(int32_on_device.node(), 0, call, 0);
std::unique_ptr<Graph> graph;
TF_ASSERT_OK(BuildXlaOps(root, fdef_lib, &graph));
Node* stateful_partitioned_call_op = nullptr;
for (Node* n : graph->op_nodes()) {
if (n->type_string() == "StatefulPartitionedCall") {
ASSERT_EQ(stateful_partitioned_call_op, nullptr);
stateful_partitioned_call_op = n;
}
}
ASSERT_NE(stateful_partitioned_call_op, nullptr);
auto xla_compile = NodeWith(Op("_XlaCompile"));
auto switch_on_compilation_pred =
NodeWith(Op("Switch"), Inputs(Out(0, xla_compile), Out(1, xla_compile)));
auto ctrl_dep =
NodeWith(Op("Identity"), Inputs(Out(0, switch_on_compilation_pred)));
// Check that we pipe int32 inputs through an IdentityN to avoid extra D2H
// copies.
EXPECT_THAT(
stateful_partitioned_call_op,
NodeWith(Inputs(Out(NodeWith(Op("IdentityN"), CtrlDeps(ctrl_dep))))));
}
#endif
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/clone_constants_for_better_clustering.cc
+225
View File
@@ -0,0 +1,225 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/jit/clone_constants_for_better_clustering.h"
#include <string>
#include "absl/algorithm/container.h"
#include "tensorflow/compiler/jit/xla_cluster_util.h"
#include "xla/status_macros.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/util/dump_graph.h"
namespace tensorflow {
using tsl::StatusOr;
class CloneConstantsForBetterClusteringPassImpl {
public:
explicit CloneConstantsForBetterClusteringPassImpl(Graph* graph)
: graph_(graph), unique_name_counter_(0) {}
absl::Status Run();
private:
absl::Status CloneSmallConstantInputs(
const absl::flat_hash_set<std::string>& name_set, Node* n);
std::string GenerateUniqueName(
const absl::flat_hash_set<std::string>& name_set,
absl::string_view prefix);
absl::StatusOr<Node*> CloneNode(
const absl::flat_hash_set<std::string>& name_set, Node* n);
Graph* graph_;
int unique_name_counter_;
};
std::string CloneConstantsForBetterClusteringPassImpl::GenerateUniqueName(
const absl::flat_hash_set<std::string>& name_set,
absl::string_view prefix) {
std::string candidate;
do {
candidate = absl::StrCat(prefix, "/clone_", unique_name_counter_++);
} while (name_set.contains(candidate));
return candidate;
}
absl::StatusOr<Node*> CloneConstantsForBetterClusteringPassImpl::CloneNode(
const absl::flat_hash_set<std::string>& name_set, Node* n) {
NodeDef new_in_def = n->def();
new_in_def.clear_input();
new_in_def.set_name(GenerateUniqueName(name_set, new_in_def.name()));
TF_ASSIGN_OR_RETURN(Node * new_in, graph_->AddNode(new_in_def));
for (const Edge* e : n->in_edges()) {
if (e->IsControlEdge()) {
graph_->AddControlEdge(e->src(), new_in);
} else {
graph_->AddEdge(e->src(), e->src_output(), new_in, e->dst_input());
}
}
new_in->set_assigned_device_name(n->assigned_device_name());
return new_in;
}
namespace {
absl::StatusOr<bool> IsConstantSmall(Node* n) {
const TensorProto* proto = nullptr;
TF_RETURN_IF_ERROR(GetNodeAttr(n->def(), "value", &proto));
int64_t total_elements = 1;
for (const auto& dim : proto->tensor_shape().dim()) {
if (dim.size() < 0) {
return absl::InternalError(absl::StrCat(
"Unknown dimension size in constant tensor ", n->name()));
}
total_elements *= dim.size();
}
// TODO(sanjoy): It may make sense to combine this threshold with XLA's "large
// constant" threshold, if there is one.
const int kSmallTensorThreshold = 16;
return total_elements < kSmallTensorThreshold;
}
// We only clone small constants since we want to avoid increasing memory
// pressure on GPUs.
absl::StatusOr<bool> IsSmallConstant(Node* n) {
if (!n->IsConstant()) {
return false;
}
return IsConstantSmall(n);
}
bool IsInPlaceOp(absl::string_view op_name) {
return op_name == "InplaceUpdate" || op_name == "InplaceAdd" ||
op_name == "InplaceSub";
}
} // namespace
absl::Status
CloneConstantsForBetterClusteringPassImpl::CloneSmallConstantInputs(
const absl::flat_hash_set<std::string>& name_set, Node* n) {
std::vector<const Edge*> in_edges;
// Get the edges and sort them so we clone in a deterministic order.
absl::c_copy(n->in_edges(), std::back_inserter(in_edges));
absl::c_stable_sort(in_edges, [](const Edge* e1, const Edge* e2) {
return e1->id() < e2->id();
});
for (const Edge* e : in_edges) {
Node* input = e->src();
TF_ASSIGN_OR_RETURN(bool is_small_constant, IsSmallConstant(input));
if (is_small_constant && input->out_edges().size() != 1) {
VLOG(2) << "Cloning small constant " << input->name();
TF_ASSIGN_OR_RETURN(Node* const input_cloned, CloneNode(name_set, input));
if (e->IsControlEdge()) {
graph_->AddControlEdge(input_cloned, e->dst());
} else {
int dst_input = e->dst_input();
TF_RET_CHECK(e->src_output() == 0)
<< "expected constant to have exactly one non-control output, but "
"found output index = "
<< e->src_output();
graph_->RemoveEdge(e);
graph_->AddEdge(input_cloned, 0, n, dst_input);
}
}
}
return absl::OkStatus();
}
absl::Status CloneConstantsForBetterClusteringPassImpl::Run() {
absl::flat_hash_set<std::string> name_set;
absl::c_transform(graph_->nodes(), std::inserter(name_set, name_set.begin()),
[](Node* n) { return n->name(); });
std::vector<Node*> nodes;
for (Node* n : graph_->nodes()) {
// We rely on the immutability of Tensors to safely clone Const operations.
// However, "in place" ops do not respect the immutability of Tensors so we
// avoid this transformation when such ops are present in the graph.
//
// In-place operations are problematic because they break the semantic
// illusion that tensorflow::Tensor instances are immutable. For instance
// if we have the following graph:
//
// digraph {
// SRC -> Const
// SRC -> I
// SRC -> V
// Const -> Identity
// Const -> InplaceAdd [label="x"]
// I -> InplaceAdd [label="i"]
// V -> InplaceAdd [label="v"]
// InplaceAdd -> Identity [style=dotted]
// }
//
// then the value produced by `Identity` is Const+I*V since InplaceAdd
// modifies the tensor in place. However, if we clone `Const` and turn the
// graph into:
//
// digraph {
// SRC -> "Const/clone_1"
// SRC -> "Const/clone_2"
// SRC -> I
// SRC -> V
// "Const/clone_1" -> Identity
// "Const/clone_2" -> InplaceAdd [label="x"]
// I -> InplaceAdd [label="i"]
// V -> InplaceAdd [label="v"]
// InplaceAdd -> Identity [style=dotted]
// }
//
// then `Identity` no longer produces Const+I*V because the InplaceAdd
// operation only modifies Const/clone_2 in place.
if (IsInPlaceOp(n->type_string())) {
return absl::OkStatus();
}
nodes.push_back(n);
}
// Iterate over a copy of the nodes to avoid iterating over g->nodes() while
// creating more nodes.
for (Node* n : nodes) {
TF_RETURN_IF_ERROR(CloneSmallConstantInputs(name_set, n));
}
return absl::OkStatus();
}
absl::Status CloneConstantsForBetterClusteringPass::Run(
const GraphOptimizationPassOptions& options) {
if (GetGlobalJitLevelForGraph(options) == OptimizerOptions::OFF) {
return absl::OkStatus();
}
Graph* g = options.graph->get();
if (VLOG_IS_ON(1)) {
DumpGraphToFile("before_clone_constants_for_better_clustering", *g);
}
TF_RETURN_IF_ERROR(CloneConstantsForBetterClusteringPassImpl{g}.Run());
if (VLOG_IS_ON(1)) {
DumpGraphToFile("after_clone_constants_for_better_clustering", *g);
}
return absl::OkStatus();
}
} // namespace tensorflow
tensorflow/compiler/jit/clone_constants_for_better_clustering.h
+62
View File
@@ -0,0 +1,62 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_JIT_CLONE_CONSTANTS_FOR_BETTER_CLUSTERING_H_
#define TENSORFLOW_COMPILER_JIT_CLONE_CONSTANTS_FOR_BETTER_CLUSTERING_H_
#include "absl/container/flat_hash_set.h"
#include "tensorflow/core/common_runtime/optimization_registry.h"
namespace tensorflow {
// Clones small host constants in the graph to make it easier to form larger
// clusters.
//
// This helps us in two ways:
//
// - It reduces dependencies between clusters. Let's say a constant C is used
// by nodes X and Y. If X and Y are put in different clusters (for whatever
// reason) Y's cluster now has to wait for all the operations in X's cluster
// to finish before it starts running.
//
// - It lets us create bigger clusters in multi-GPU benchmarks. Consider the
// following graph:
//
// digraph {
// Const -> GPU_1
// Const -> GPU_0_Y
// GPU_0_X -> GPU_0_Y
// }
//
// We'd cluster Const and GPU_1 together (and place it on GPU_1), and this
// will block us from clustering GPU_0_X and GPU_0_Y together since that
// would increase the amount of work on GPU 0 waiting on work on GPU 1.
// However, cloning Const into two copies, one for GPU_0_Y and one for GPU_1
// will let us create one cluster containing {Const/copy_0, GPU_1} and
// another containing {Const/copy_1, GPU_0_X, GPU_0_Y}.
//
// We only clone small host constants now to avoid increasing memory consumption
// too much. Moreover, in practice the constants we have to duplicate are
// things like the `perm` input to `Transpose` and the `size` input to `Slice`
// which tend to be small anyway.
class CloneConstantsForBetterClusteringPass : public GraphOptimizationPass {
public:
CloneConstantsForBetterClusteringPass() = default;
absl::Status Run(const GraphOptimizationPassOptions& options) override;
};
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_CLONE_CONSTANTS_FOR_BETTER_CLUSTERING_H_
tensorflow/compiler/jit/clone_constants_for_better_clustering_test.cc
+232
View File
@@ -0,0 +1,232 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/jit/clone_constants_for_better_clustering.h"
#include "tensorflow/cc/framework/ops.h"
#include "tensorflow/cc/ops/array_ops.h"
#include "tensorflow/cc/ops/const_op.h"
#include "tensorflow/cc/ops/math_ops.h"
#include "tensorflow/compiler/jit/node_matchers.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/public/session_options.h"
namespace tensorflow {
namespace {
using ::tensorflow::testing::FindNodeByName;
absl::Status CloneConstantsForBetterClustering(const Scope& s,
std::unique_ptr<Graph>* result) {
auto graph = std::make_unique<Graph>(OpRegistry::Global());
SessionOptions session_options;
session_options.config.mutable_graph_options()
->mutable_optimizer_options()
->set_global_jit_level(OptimizerOptions::ON_2);
GraphOptimizationPassOptions options;
options.graph = &graph;
options.session_options = &session_options;
// Scope::ToGraph seems to drop assigned devices, probably because it goes
// through a GraphDef. So explicitly maintain the device assignment.
// std::unordered_map<string, string> assigned_device_names;
// for (Node* n : s.graph()->nodes()) {
// assigned_device_names[n->name()] = n->assigned_device_name();
// }
GraphConstructorOptions opts;
opts.expect_device_spec = true;
TF_RETURN_IF_ERROR(s.ToGraph(graph.get(), opts));
// for (Node* n : graph->nodes()) {
// n->set_assigned_device_name(assigned_device_names[n->name()]);
// }
CloneConstantsForBetterClusteringPass rewriter;
TF_RETURN_IF_ERROR(rewriter.Run(options));
*result = std::move(graph);
return absl::OkStatus();
}
const char* kCPU = "/job:localhost/replica:0/task:0/device:CPU:0";
const char* kGPU = "/job:localhost/replica:0/task:0/device:GPU:0";
TEST(CloneConstantsForBetterClusteringTest, ScalarConstantPlacedOnGpu) {
Scope root = Scope::NewRootScope().ExitOnError();
Scope on_gpu = root.WithAssignedDevice(kGPU).WithDevice(kGPU);
Output in = ops::Placeholder(on_gpu.WithOpName("in"), DT_FLOAT);
Output c = ops::Const(on_gpu.WithOpName("const"), 1.0f, {});
Output add1 = ops::AddV2(on_gpu.WithOpName("add1"), in, c);
Output add2 = ops::AddV2(on_gpu.WithOpName("add2"), add1, c);
std::unique_ptr<Graph> result;
TF_ASSERT_OK(CloneConstantsForBetterClustering(root, &result));
OutputTensor add1_operand;
TF_ASSERT_OK(
FindNodeByName(result.get(), "add1")->input_tensor(1, &add1_operand));
OutputTensor add2_operand;
TF_ASSERT_OK(
FindNodeByName(result.get(), "add2")->input_tensor(1, &add2_operand));
EXPECT_NE(add1_operand.node, add2_operand.node);
}
TEST(CloneConstantsForBetterClusteringTest, HostConstantPlacedOnCpu) {
Scope root = Scope::NewRootScope().ExitOnError();
Scope on_gpu = root.WithAssignedDevice(kGPU).WithDevice(kGPU);
Scope on_cpu = root.WithAssignedDevice(kCPU).WithDevice(kCPU);
Output in0 = ops::Placeholder(on_gpu.WithOpName("in0"), DT_FLOAT);
Output in1 = ops::Placeholder(on_gpu.WithOpName("in1"), DT_FLOAT);
Output perm = ops::Const(on_cpu.WithOpName("perm"), {3, 1, 2, 0});
{
Output tr0 = ops::Transpose(on_gpu.WithOpName("tr0"), in0, perm);
Output tr1 = ops::Transpose(on_gpu.WithOpName("tr1"), in1, perm);
}
std::unique_ptr<Graph> result;
TF_ASSERT_OK(CloneConstantsForBetterClustering(root, &result));
OutputTensor tr0_perm;
TF_ASSERT_OK(FindNodeByName(result.get(), "tr0")->input_tensor(1, &tr0_perm));
OutputTensor tr1_perm;
TF_ASSERT_OK(FindNodeByName(result.get(), "tr1")->input_tensor(1, &tr1_perm));
EXPECT_NE(tr0_perm.node, tr1_perm.node);
}
TEST(CloneConstantsForBetterClusteringTest, HostConstantPlacedOnGpu) {
Scope root = Scope::NewRootScope().ExitOnError();
Scope on_gpu = root.WithAssignedDevice(kGPU).WithDevice(kGPU);
Output in0 = ops::Placeholder(on_gpu.WithOpName("in0"), DT_FLOAT);
Output in1 = ops::Placeholder(on_gpu.WithOpName("in1"), DT_FLOAT);
Output perm = ops::Const(on_gpu.WithOpName("perm"), {3, 1, 2, 0});
{
Output tr0 = ops::Transpose(on_gpu.WithOpName("tr0"), in0, perm);
Output tr1 = ops::Transpose(on_gpu.WithOpName("tr1"), in1, perm);
}
std::unique_ptr<Graph> result;
TF_ASSERT_OK(CloneConstantsForBetterClustering(root, &result));
OutputTensor tr0_perm;
TF_ASSERT_OK(FindNodeByName(result.get(), "tr0")->input_tensor(1, &tr0_perm));
OutputTensor tr1_perm;
TF_ASSERT_OK(FindNodeByName(result.get(), "tr1")->input_tensor(1, &tr1_perm));
EXPECT_NE(tr0_perm.node, tr1_perm.node);
}
TEST(CloneConstantsForBetterClusteringTest, CloneSmallDeviceConstants) {
Scope root = Scope::NewRootScope().ExitOnError();
Scope on_gpu = root.WithAssignedDevice(kGPU).WithDevice(kGPU);
Output in0 = ops::Placeholder(on_gpu.WithOpName("in0"), DT_FLOAT);
Output in1 = ops::Placeholder(on_gpu.WithOpName("in1"), DT_FLOAT);
Output perm_f32 = ops::Const(on_gpu.WithOpName("perm"), {3.0, 1.0, 2.0, 0.0});
Output perm_int0 =
ops::Cast(on_gpu.WithOpName("perm_cast_0"), perm_f32, DT_INT32);
Output perm_int1 =
ops::Cast(on_gpu.WithOpName("perm_cast_1"), perm_f32, DT_INT32);
{
Output tr0 = ops::Transpose(on_gpu.WithOpName("tr0"), in0, perm_int0);
Output tr1 = ops::Transpose(on_gpu.WithOpName("tr1"), in1, perm_int1);
}
std::unique_ptr<Graph> result;
TF_ASSERT_OK(CloneConstantsForBetterClustering(root, &result));
OutputTensor tr0_perm;
TF_ASSERT_OK(
FindNodeByName(result.get(), "perm_cast_0")->input_tensor(0, &tr0_perm));
OutputTensor tr1_perm;
TF_ASSERT_OK(
FindNodeByName(result.get(), "perm_cast_1")->input_tensor(0, &tr1_perm));
EXPECT_NE(tr0_perm.node, tr1_perm.node);
}
TEST(CloneConstantsForBetterClusteringTest, DontCloneLargeConstants) {
Scope root = Scope::NewRootScope().ExitOnError();
Scope on_gpu = root.WithAssignedDevice(kGPU).WithDevice(kGPU);
Scope on_cpu = root.WithAssignedDevice(kCPU).WithDevice(kCPU);
Output in0 = ops::Placeholder(on_gpu.WithOpName("in0"), DT_FLOAT);
Output in1 = ops::Placeholder(on_gpu.WithOpName("in1"), DT_FLOAT);
Output perm = ops::Const(
on_cpu.WithOpName("perm"),
{17, 16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0});
{
Output tr0 = ops::Transpose(on_gpu.WithOpName("tr0"), in0, perm);
Output tr1 = ops::Transpose(on_gpu.WithOpName("tr1"), in1, perm);
}
std::unique_ptr<Graph> result;
TF_ASSERT_OK(CloneConstantsForBetterClustering(root, &result));
OutputTensor tr0_perm;
TF_ASSERT_OK(FindNodeByName(result.get(), "tr0")->input_tensor(1, &tr0_perm));
OutputTensor tr1_perm;
TF_ASSERT_OK(FindNodeByName(result.get(), "tr1")->input_tensor(1, &tr1_perm));
EXPECT_EQ(tr0_perm.node, tr1_perm.node);
}
TEST(CloneConstantsForBetterClusteringTest, InplaceOps) {
Scope root = Scope::NewRootScope().ExitOnError();
Scope on_gpu = root.WithAssignedDevice(kGPU).WithDevice(kGPU);
Scope on_cpu = root.WithAssignedDevice(kCPU).WithDevice(kCPU);
Output in0 = ops::Placeholder(on_gpu.WithOpName("in0"), DT_FLOAT);
Output in1 = ops::Placeholder(on_gpu.WithOpName("in1"), DT_FLOAT);
Output perm = ops::Const(on_cpu.WithOpName("perm"), {3, 1, 2, 0});
{
Output tr0 = ops::Transpose(on_gpu.WithOpName("tr0"), in0, perm);
Output tr1 = ops::Transpose(on_gpu.WithOpName("tr1"), in1, perm);
}
Output in_place_add =
ops::InplaceAdd(on_cpu.WithOpName("tr0"), perm,
ops::Placeholder(on_cpu.WithOpName("i"), DT_INT32), perm);
std::unique_ptr<Graph> result;
TF_ASSERT_OK(CloneConstantsForBetterClustering(root, &result));
OutputTensor tr0_perm;
TF_ASSERT_OK(FindNodeByName(result.get(), "tr0")->input_tensor(1, &tr0_perm));
OutputTensor tr1_perm;
TF_ASSERT_OK(FindNodeByName(result.get(), "tr1")->input_tensor(1, &tr1_perm));
EXPECT_EQ(tr0_perm.node, tr1_perm.node);
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/cluster_scoping_pass.cc
+164
View File
@@ -0,0 +1,164 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/jit/cluster_scoping_pass.h"
#include "absl/algorithm/container.h"
#include "absl/container/flat_hash_set.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/compiler/jit/defs.h"
#include "tensorflow/compiler/jit/xla_cluster_util.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/graph/algorithm.h"
namespace tensorflow {
namespace {
class ClusterScopingPassImpl {
public:
ClusterScopingPassImpl(Graph* graph,
OptimizerOptions::GlobalJitLevel global_jit_level)
: graph_(graph),
global_jit_level_(global_jit_level),
unique_scope_id_(0) {}
absl::Status Run();
private:
absl::Status ScopingForPipelineStages();
size_t GetUniqueScopeId() { return unique_scope_id_++; }
void AddScopeToAllTransitivePredecessors(Node* start);
void AddScopeToAllTransitiveSuccessors(Node* start);
private:
Graph* graph_;
OptimizerOptions::GlobalJitLevel global_jit_level_;
size_t unique_scope_id_;
};
std::optional<std::string> GetXlaInternalScope(Node* node) {
std::string scope;
if (GetNodeAttr(node->attrs(), kXlaInternalScopeAttr, &scope).ok()) {
return scope;
}
return std::nullopt;
}
void SetXlaInternalScope(Node* node, absl::string_view scope) {
node->AddAttr(kXlaInternalScopeAttr, scope);
}
// NB! We append a new scope as suffix to the _XlaInternalScope attribute
// instead of overriding the old value. In other words, appending scope B to
// scope A creates the conjunction of the scopes A and B (i.e, A & B) and,
// in effect, the node gets both the old and new scopes. As a unique scope
// disallows a node being merged with nodes in other scopes, the scope
// conjunction preserves the semantic of the old scope (i.e., the node still
// cannot be merged with the previously incompatible nodes.)
//
// For example, the below case should be rare in practice but can serve for the
// purpose of discussion. After adding scopes for both Stage and Unstage,
// Node_Y will receive both scopes "unstage" and "stage", while Node_X receives
// only scope "stage". The semantic of scope "unstage" is preserved although
// scope "stage" is later appended. As a result, Node_X and Node_Y will be put
// into different clusters.
//
// Unstage -> Node_Y (scope "unstage & stage")
// |
// V
// Node_X (scope "stage") -> Stage
//
void AddOrAppendXlaInternalScope(Node* node, absl::string_view suffix) {
std::string updated_scope;
std::optional<std::string> cur_scope = GetXlaInternalScope(node);
if (cur_scope == std::nullopt) {
updated_scope = std::string(suffix);
} else {
updated_scope = absl::StrCat(cur_scope.value(), "&", suffix);
}
SetXlaInternalScope(node, updated_scope);
}
void ClusterScopingPassImpl::AddScopeToAllTransitivePredecessors(Node* start) {
const std::string unique_suffix = absl::StrCat("_", GetUniqueScopeId());
std::vector<Node*> starts;
starts.push_back(start);
auto enter = [&](Node* n) { AddOrAppendXlaInternalScope(n, unique_suffix); };
ReverseDFSFrom(*graph_, starts, enter, /*leave=*/nullptr,
/*stable_comparator=*/NodeComparatorName());
}
void ClusterScopingPassImpl::AddScopeToAllTransitiveSuccessors(Node* start) {
const std::string unique_suffix = absl::StrCat("_", GetUniqueScopeId());
std::vector<Node*> starts;
starts.push_back(start);
auto enter = [&](Node* n) { AddOrAppendXlaInternalScope(n, unique_suffix); };
DFSFrom(*graph_, starts, enter, /*leave=*/nullptr,
/*stable_comparator=*/NodeComparatorName(),
// Do not filter any edges to better capture the semantics of
// transitive closure of successors. We may revisit this when
// we see more cases needing cluster scoping in the future.
/*edge_filter=*/nullptr);
}
// This preserves the parallelism between pipeline stages. For example, below
// is a typical pattern of input pipelining in Tensorflow and this heuristic
// ensures Node_X and Node_Y are put into different clusters. Without the
// heuristic, they may be put into the same cluster and it can introduce
// artificial dependencies and incur great performance loss. In this example,
// Node_Y becomes dependent on IteratorGetNext and the latencies add up if
// Node_X and Node_Y are in the same cluster.
//
// IteratorGetNext -> Node_X -> Stage
//
// Unstage -> Node_Y
//
absl::Status ClusterScopingPassImpl::ScopingForPipelineStages() {
for (Node* n : graph_->nodes()) {
DCHECK(n);
if (n->type_string() == "Unstage") {
AddScopeToAllTransitiveSuccessors(n);
}
if (n->type_string() == "Stage") {
AddScopeToAllTransitivePredecessors(n);
}
}
return absl::OkStatus();
}
absl::Status ClusterScopingPassImpl::Run() {
if (global_jit_level_ == OptimizerOptions::OFF) {
return absl::OkStatus();
}
return ScopingForPipelineStages();
}
} // namespace
absl::Status ClusterScopingPass::Run(
const GraphOptimizationPassOptions& options) {
Graph* graph = options.graph->get();
return ClusterScopingPassImpl{graph, GetGlobalJitLevelForGraph(options)}
.Run();
}
} // namespace tensorflow
tensorflow/compiler/jit/cluster_scoping_pass.h
+38
View File
@@ -0,0 +1,38 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_JIT_CLUSTER_SCOPING_PASS_H_
#define TENSORFLOW_COMPILER_JIT_CLUSTER_SCOPING_PASS_H_
#include "tensorflow/core/common_runtime/optimization_registry.h"
namespace tensorflow {
// This pass adds scopes to nodes in the _XlaInternalScope attribute to guide
// the later clustering passes. A major reason to do this is to prevent the
// clustering from losing critical parallelism in the Tensorflow graph, which
// can incur great performance degradation.
//
// This pass must be run before MarkForCompilationPass, as it stores the
// scoping information that MarkForCompilationPass will need to respect for
// clustering decision.
class ClusterScopingPass : public GraphOptimizationPass {
public:
absl::Status Run(const GraphOptimizationPassOptions& options) override;
};
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_CLUSTER_SCOPING_PASS_H_
tensorflow/compiler/jit/cluster_scoping_pass_test.cc
+180
View File
@@ -0,0 +1,180 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/jit/cluster_scoping_pass.h"
#include "absl/container/flat_hash_map.h"
#include "tensorflow/compiler/jit/defs.h"
#include "tensorflow/compiler/jit/test_util.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/common_runtime/graph_def_builder_util.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/graph/algorithm.h"
#include "tensorflow/core/graph/graph_def_builder.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/public/session_options.h"
namespace tensorflow {
namespace {
absl::Status ClusterScoping(std::unique_ptr<Graph>* graph) {
FixupSourceAndSinkEdges(graph->get());
GraphOptimizationPassWrapper wrapper;
wrapper.session_options.config.mutable_graph_options()
->mutable_optimizer_options()
->set_global_jit_level(OptimizerOptions::ON_2);
GraphOptimizationPassOptions opt_options =
wrapper.CreateGraphOptimizationPassOptions(graph);
ClusterScopingPass pass;
return pass.Run(opt_options);
}
absl::flat_hash_map<std::string, std::string> GetXlaInternalScopes(
const Graph& graph) {
absl::flat_hash_map<std::string, std::string> scopes;
for (Node* node : graph.nodes()) {
std::string scope;
if (GetNodeAttr(node->attrs(), kXlaInternalScopeAttr, &scope).ok()) {
scopes[node->name()] = scope;
}
}
if (VLOG_IS_ON(2)) {
VLOG(2) << "_XlaInternalScopes:";
for (const auto& p : scopes) {
VLOG(2) << " " << p.first << " -> " << p.second;
}
}
return scopes;
}
Node* BuildStageNode(GraphDefBuilder& builder, std::string name,
std::initializer_list<DataType> dtypes,
absl::Span<const ops::NodeOut> values) {
auto opts = builder.opts()
.WithName(std::move(name))
.WithAttr("dtypes", std::move(dtypes));
if (opts.HaveError()) {
return nullptr;
}
NodeBuilder node_builder(name, "Stage", opts.op_registry());
node_builder.Input(values);
return opts.FinalizeBuilder(&node_builder);
}
TEST(XlaCompilationTest, StagePipelinePreserved) {
std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
{
// Graph:
// b
// |
// v
// a -> add0 (ClusterX) -> relu0 (ClusterX) -> stage
//
// b
// |
// v
// unstage -> add1 (ClusterY) -> relu1 (ClusterY)
GraphDefBuilder builder(GraphDefBuilder::kFailImmediately);
Node* a = ops::SourceOp("Const", builder.opts()
.WithName("a")
.WithAttr("dtype", DT_FLOAT)
.WithAttr("value", Tensor()));
Node* b = ops::SourceOp("Const", builder.opts()
.WithName("b")
.WithAttr("dtype", DT_FLOAT)
.WithAttr("value", Tensor()));
Node* unstage = ops::SourceOp(
"Unstage",
builder.opts().WithName("unstage").WithAttr("dtypes", {DT_FLOAT}));
Node* add0 = ops::BinaryOp("Add", a, b, builder.opts().WithName("add0"));
Node* add1 =
ops::BinaryOp("Add", unstage, b, builder.opts().WithName("add1"));
Node* relu0 = ops::UnaryOp("Relu", add0, builder.opts().WithName("relu0"));
ops::UnaryOp("Relu", add1, builder.opts().WithName("relu1"));
BuildStageNode(builder, "stage", {DT_FLOAT}, {relu0});
TF_EXPECT_OK(GraphDefBuilderToGraph(builder, graph.get()));
}
TF_ASSERT_OK(ClusterScoping(&graph));
auto scopes = GetXlaInternalScopes(*graph);
EXPECT_NE(scopes["add0"], scopes["add1"]);
EXPECT_EQ(scopes["add0"], scopes["relu0"]);
EXPECT_EQ(scopes["add1"], scopes["relu1"]);
}
TEST(XlaCompilationTest, StagePipelinePreservedAndInitialScopesRespected) {
std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
{
// Graph:
// b
// |
// v
// a -> add0 (ClusterA) -> relu0 (ClusterB) -> stage
//
// b
// |
// v
// unstage -> add1 (ClusterC) -> relu1 (ClusterD)
GraphDefBuilder builder(GraphDefBuilder::kFailImmediately);
Node* a = ops::SourceOp("Const", builder.opts()
.WithName("a")
.WithAttr("dtype", DT_FLOAT)
.WithAttr("value", Tensor()));
Node* b = ops::SourceOp("Const", builder.opts()
.WithName("b")
.WithAttr("dtype", DT_FLOAT)
.WithAttr("value", Tensor()));
Node* unstage = ops::SourceOp(
"Unstage",
builder.opts().WithName("unstage").WithAttr("dtypes", {DT_FLOAT}));
// Intentionally give add0 and add1 the same initial scope but they should
// be separated by the ClusterScopingPass.
Node* add0 = ops::BinaryOp("Add", a, b,
builder.opts().WithName("add0").WithAttr(
kXlaInternalScopeAttr, "ClusterA"));
Node* add1 = ops::BinaryOp("Add", unstage, b,
builder.opts().WithName("add1").WithAttr(
kXlaInternalScopeAttr, "ClusterA"));
Node* relu0 = ops::UnaryOp("Relu", add0,
builder.opts().WithName("relu0").WithAttr(
kXlaInternalScopeAttr, "ClusterB"));
ops::UnaryOp("Relu", add1,
builder.opts().WithName("relu1").WithAttr(
kXlaInternalScopeAttr, "ClusterD"));
BuildStageNode(builder, "stage", {DT_FLOAT}, {relu0});
TF_EXPECT_OK(GraphDefBuilderToGraph(builder, graph.get()));
}
TF_ASSERT_OK(ClusterScoping(&graph));
auto scopes = GetXlaInternalScopes(*graph);
EXPECT_NE(scopes["add0"], scopes["add1"]);
EXPECT_NE(scopes["add0"], scopes["relu0"]);
EXPECT_NE(scopes["add1"], scopes["relu1"]);
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/compilability_check_util.cc
+758
View File
@@ -0,0 +1,758 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/jit/compilability_check_util.h"
#include <cstddef>
#include <iterator>
#include <string>
#include <utility>
#include <vector>
#include "absl/algorithm/container.h"
#include "absl/container/flat_hash_set.h"
#include "absl/log/check.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/jit/defs.h"
#include "tensorflow/compiler/jit/xla_activity.pb.h"
#include "tensorflow/compiler/jit/xla_activity_listener.h"
#include "tensorflow/compiler/jit/xla_cluster_util.h"
#include "tensorflow/compiler/tf2xla/const_analysis.h"
#include "tensorflow/compiler/tf2xla/tf2xla_defs.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "xla/tsl/platform/errors.h"
#include "tensorflow/core/common_runtime/function_body.h"
#include "tensorflow/core/common_runtime/function_utils.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/graph_def_util.h"
#include "tensorflow/core/framework/memory_types.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/lib/gtl/cleanup.h"
namespace tensorflow {
namespace {
bool HasResourceInput(const Node& node) {
return absl::c_count(node.input_types(), DT_RESOURCE) != 0;
}
void LogNotCompilable(const Node& node, absl::string_view reason = "") {
VLOG(3) << "Found uncompilable node " << node.name() << " (op "
<< node.type_string() << ")" << (reason.empty() ? "" : ": ")
<< reason;
}
bool IsInOutsideCompilationCluster(const Node& n) {
return n.attrs().Find(kXlaOutsideCompilationAttr) != nullptr;
}
absl::Status MakeCallNodeFromAttribute(const Node& node,
const std::string& attr_name,
NodeDef* node_def) {
const NameAttrList* name_attr;
TF_RETURN_IF_ERROR(GetNodeAttr(node.attrs(), attr_name, &name_attr));
node_def->set_op(name_attr->name());
*(node_def->mutable_attr()) = name_attr->attr();
return absl::OkStatus();
}
absl::StatusOr<std::vector<NodeDef>> MakeCallNodesFromAttribute(
const Node& node, absl::string_view attr_name,
absl::string_view call_name) {
std::vector<NameAttrList> attr_lists;
TF_RETURN_IF_ERROR(GetNodeAttr(node.attrs(), attr_name, &attr_lists));
std::vector<NodeDef> out;
out.reserve(attr_lists.size());
for (int i = 0; i < attr_lists.size(); i++) {
out.emplace_back();
NodeDef& inserted = out.back();
inserted.set_name(absl::StrCat(call_name, "_", i));
inserted.set_op(attr_lists[i].name());
*inserted.mutable_attr() = attr_lists[i].attr();
}
return out;
}
// Utility which searches for values in a sorted list by scanning over it once.
// No matter how many times ScanForValue is called, the list is scanned at most
// once. However, if a call to ScanForValue skips over a value, that value is
// not revisited in future calls to ScanForValue, so callers must take
// care to order their calls.
//
// Useful for merging multiple sorted lists in O(n) time.
class SinglePassSearch {
public:
// Creates a SinglePassSearch object that can be used to search in `values`.
// Does not take ownership of `values`. `values` must outlive this.
// `values` must be sorted.
explicit SinglePassSearch(absl::Span<int const> values)
: current_index_(0), values_(values) {}
// Scans forward in the vector looking for "value", updating the internal
// position in to the vector.
// Returns true iff the vector contains the given value at or after current
// position.
// Not thread-safe.
bool ScanForValue(int value) {
while (current_index_ < values_.size() &&
values_[current_index_] <= value) {
if (values_[current_index_] == value) {
current_index_++;
return true;
}
current_index_++;
}
return false;
}
private:
int current_index_;
const absl::Span<int const> values_;
};
} // anonymous namespace
RecursiveCompilabilityChecker::UncompilableNodesMap
RecursiveCompilabilityChecker::FindUncompilableNodes(
const Node& node, FunctionLibraryRuntime* lib_runtime,
const std::vector<RecursiveCompilabilityChecker::StackFrame>*
node_stack_trace) const {
std::vector<StackFrameView> stack_trace;
// If `node_stack_trace` is provided, that means `node` is inside
// a function body, and therefore, arg nodes and retval nodes are
// not considered uncompilable.
if (node_stack_trace != nullptr) {
for (const auto& frame : *node_stack_trace) {
stack_trace.emplace_back(
StackFrameView{frame.name, frame.function_name, frame.stack_trace});
}
}
stack_trace.emplace_back(
StackFrameView{node.name(), "", node.GetStackTrace()});
RecursiveCompilabilityChecker::UncompilableNodesMap uncompilable_nodes;
IsCompilableNode(node, lib_runtime, &stack_trace,
/*encapsulating_function=*/nullptr, &uncompilable_nodes);
return uncompilable_nodes;
}
bool RecursiveCompilabilityChecker::HasXLAKernel(
const Node& node, std::string* uncompilable_reason) const {
// There is a SymbolicGradient kernel on the XLA_JIT device, but the gradient
// is really a kind of function call and will be handled by
// IsCompilableCall().
if (node.type_string() == "SymbolicGradient") {
*uncompilable_reason =
"SymbolicGradient should be handled by IsCompilableCall().";
return false;
}
if (node.type_string() == "Const") {
const AttrValue* attr = node.attrs().Find("dtype");
if (!op_filter_.allow_string_consts && attr != nullptr &&
attr->type() == DT_STRING) {
*uncompilable_reason =
"Const op with type DT_STRING is not supported by XLA.";
return false;
}
}
// XLA does not offer guaranteed aliasing between the input and output of the
// XLA cluster so it can't implement the forward-tensor-ref semantic. Leave
// such nodes out of XLA clusters.
if (HasForwardedRefInput(node)) {
VLOG(2) << "Rejecting " << node.name() << ": Identity with unsafe cast.";
*uncompilable_reason = "Identity with unsafe cast.";
return false;
}
absl::Status s =
FindKernelDef(jit_device_type_, node.def(), nullptr, nullptr);
if (!s.ok()) {
*uncompilable_reason = s.message();
return false;
}
return true;
}
// Tests whether 'if_node' is compilable. Every operator in the then_branch and
// else_branch functions must be compilable for 'if_node' to be compilable.
bool RecursiveCompilabilityChecker::IsCompilableIf(
const Node& if_node, FunctionLibraryRuntime* lib_runtime,
std::vector<StackFrameView>* stack_trace,
NameAttrList* encapsulating_function,
RecursiveCompilabilityChecker::UncompilableNodesMap* uncompilable_nodes)
const {
bool is_compilable = true;
is_compilable &= ExtractNodeDefAndCheckCompilability(
if_node, "then_branch", "if_then", encapsulating_function, lib_runtime,
stack_trace, uncompilable_nodes);
if (!uncompilable_nodes && !is_compilable) return is_compilable;
is_compilable &= ExtractNodeDefAndCheckCompilability(
if_node, "else_branch", "if_else", encapsulating_function, lib_runtime,
stack_trace, uncompilable_nodes);
return is_compilable;
}
bool RecursiveCompilabilityChecker::IsCompilableCase(
const Node& case_node, FunctionLibraryRuntime* lib_runtime,
std::vector<StackFrameView>* stack_trace,
NameAttrList* encapsulating_function,
RecursiveCompilabilityChecker::UncompilableNodesMap* uncompilable_nodes)
const {
absl::StatusOr<std::vector<NodeDef>> calls =
MakeCallNodesFromAttribute(case_node, "branches", "branch");
if (!calls.ok()) {
VLOG(2) << "Rejecting node " << case_node.name() << ": "
<< "missing attribute 'branches'";
return false;
}
bool is_compilable = true;
for (const NodeDef& call : *calls) {
is_compilable &=
IsCompilableCall(call, lib_runtime, stack_trace, encapsulating_function,
uncompilable_nodes);
}
return is_compilable;
}
// Tests whether 'while_node' is a completely compilable loop.
// Every operator in the condition and body functions must be compilable for a
// while loop to be compilable.
bool RecursiveCompilabilityChecker::IsCompilableWhile(
const Node& while_node, FunctionLibraryRuntime* lib_runtime,
std::vector<StackFrameView>* stack_trace,
NameAttrList* encapsulating_function,
RecursiveCompilabilityChecker::UncompilableNodesMap* uncompilable_nodes)
const {
bool is_compilable = true;
is_compilable &= ExtractNodeDefAndCheckCompilability(
while_node, "cond", "while_cond", encapsulating_function, lib_runtime,
stack_trace, uncompilable_nodes);
if (!uncompilable_nodes && !is_compilable) return is_compilable;
is_compilable &= ExtractNodeDefAndCheckCompilability(
while_node, "body", "while_body", encapsulating_function, lib_runtime,
stack_trace, uncompilable_nodes);
return is_compilable;
}
bool RecursiveCompilabilityChecker::ExtractNodeDefAndCheckCompilability(
const Node& node, const std::string& attr_name,
const std::string& call_name, NameAttrList* encapsulating_function,
FunctionLibraryRuntime* lib_runtime,
std::vector<StackFrameView>* stack_trace,
RecursiveCompilabilityChecker::UncompilableNodesMap* uncompilable_nodes)
const {
NodeDef call;
call.set_name(call_name);
if (!MakeCallNodeFromAttribute(node, attr_name, &call).ok()) {
const auto uncompilable_reason = absl::StrCat(
"missing '", attr_name, "' attribute from node", node.name());
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
VLOG(2) << "Rejecting node " << node.name() << ": " << uncompilable_reason
<< ".";
return false;
}
if (!IsCompilableCall(call, lib_runtime, stack_trace, encapsulating_function,
uncompilable_nodes)) {
VLOG(2) << "Rejecting node " << node.name()
<< ": can't compile : " << call.op();
return false;
}
return true;
}
// Tests whether 'call_def' is a call to a completely compilable function.
// Every operator in the function must be compilable for a function to be
// compilable.
bool RecursiveCompilabilityChecker::IsCompilableCall(
const NodeDef& call_def, FunctionLibraryRuntime* lib_runtime,
std::vector<StackFrameView>* stack_trace,
NameAttrList* encapsulating_function,
RecursiveCompilabilityChecker::UncompilableNodesMap* uncompilable_nodes)
const {
if (stack_trace->size() > kMaxRecursionDepth) {
std::string uncompilable_reason = "function depth limit exceeded";
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
VLOG(2) << "Rejecting " << call_def.op() << ": " << uncompilable_reason
<< ".";
return false;
}
FunctionLibraryRuntime::Handle handle;
absl::Status s;
NameAttrList function;
s = NameAndAttrsFromFunctionCall(call_def, &function);
if (s.ok()) {
s = lib_runtime->Instantiate(function.name(), AttrSlice(&function.attr()),
&handle);
}
if (!s.ok()) {
std::string uncompilable_reason =
absl::StrCat("could not instantiate call: '", function.name(), "'");
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
VLOG(2) << "Rejecting " << call_def.DebugString() << ": "
<< uncompilable_reason << " : " << s;
return false;
}
auto release_handle_on_return =
gtl::MakeCleanup([&] { CHECK_OK(lib_runtime->ReleaseHandle(handle)); });
const FunctionBody* fbody = lib_runtime->GetFunctionBody(handle);
bool is_compilable = true;
for (const Node* node : fbody->graph->op_nodes()) {
stack_trace->emplace_back(
StackFrameView{node->name(), function.name(), node->GetStackTrace()});
is_compilable &= IsCompilableNode(*node, lib_runtime, stack_trace,
&function, uncompilable_nodes);
stack_trace->pop_back();
if (!uncompilable_nodes && !is_compilable) return is_compilable;
}
return is_compilable;
}
bool RecursiveCompilabilityChecker::OpIsInaccurate(const Node& node) const {
// b/127344411: SelfAdjointEigV2 and Svd precision issues.
return node.type_string() == "SelfAdjointEigV2" ||
node.type_string() == "Svd";
}
bool RecursiveCompilabilityChecker::OpIsSlow(const Node& node) const {
// b/128001705: SelfAdjointEigV2 and Svd performance issues.
// b/135640736: MatrixInverse performance issues.
// b/111271662: MatrixSolve performance issues.
// https://github.com/tensorflow/tensorflow/pull/31012:
// ResizeNearestNeighbor, ResizeBilinear, and ResizeBilinearGrad sometimes
// create convolutions too large for CuDNN to handle.
// NonMaxSuppressionV3/V4 in XLA runs significantly slower than TF kernel in
// object detection models, specially when there are a lot of proposed
// bounding boxes.
return node.type_string() == "SelfAdjointEigV2" ||
node.type_string() == "Svd" || node.type_string() == "Qr" ||
node.type_string() == "MatrixInverse" ||
node.type_string() == "MatrixSolve" ||
node.type_string() == "ResizeBilinearGrad" ||
node.type_string() == "NonMaxSuppressionV3" ||
node.type_string() == "NonMaxSuppressionV4";
}
bool RecursiveCompilabilityChecker::IsCompilableNode(
const Node& node, FunctionLibraryRuntime* lib_runtime,
std::vector<StackFrameView>* stack_trace,
NameAttrList* encapsulating_function,
RecursiveCompilabilityChecker::UncompilableNodesMap* uncompilable_nodes)
const {
auto stack_depth = stack_trace->size();
if (op_filter_.allow_outside_compiled && IsInOutsideCompilationCluster(node))
return true;
if (node.IsSource() || node.IsSink()) {
absl::string_view uncompilable_reason = "source or sink node";
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
// _Arg nodes in a top-level function represent feeds and _Retval nodes in a
// top-level function represent fetches.
if (stack_depth == 1 &&
(node.type_string() == "_Arg" || node.type_string() == "_Retval")) {
absl::string_view uncompilable_reason = "top level _Arg or _Retval";
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
if (node.attrs().Find("_scoped_allocator") ||
node.attrs().Find("_forward_from")) {
// TODO(b/128858118): XLA does not support _scoped_allocator and
// _forward_from.
absl::string_view uncompilable_reason =
"_scoped_allocator or _forward_from attribute";
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
std::string uncompilable_reason;
if (IsFunctionCall(*lib_runtime->GetFunctionLibraryDefinition(), node)) {
if (!IsCompilableCall(node.def(), lib_runtime, stack_trace,
encapsulating_function, uncompilable_nodes)) {
LogNotCompilable(node, "unsupported function");
return false;
}
} else if (!HasXLAKernel(node, &uncompilable_reason)) {
MaybeMarkUncompilableNode(
absl::StrCat("unsupported op: ", uncompilable_reason), *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
if (node.IsWhileNode() &&
!IsCompilableWhile(node, lib_runtime, stack_trace, encapsulating_function,
uncompilable_nodes)) {
LogNotCompilable(node, "unsupported while");
return false;
}
if (node.IsIfNode() &&
!IsCompilableIf(node, lib_runtime, stack_trace, encapsulating_function,
uncompilable_nodes)) {
LogNotCompilable(node, "unsupported if");
return false;
}
if (op_filter_.require_always_compilable && node.IsCaseNode() &&
!IsCompilableCase(node, lib_runtime, stack_trace, encapsulating_function,
uncompilable_nodes)) {
LogNotCompilable(node, "unsupported case");
return false;
}
if (!op_filter_.allow_stateful_rng_ops &&
IsStatefulRandomOp(node.type_string())) {
absl::string_view uncompilable_reason = "stateful random op";
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
if (!op_filter_.allow_control_trigger && node.IsControlTrigger()) {
absl::string_view uncompilable_reason = "not allowed control trigger";
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
if (!op_filter_.allow_eliding_assert_and_checknumerics_ops &&
IsAssertOrCheckNumerics(node.type_string())) {
absl::string_view uncompilable_reason = "Assert or CheckNumerics";
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
if (!op_filter_.allow_collective_reduce_v2 &&
node.type_string() == "CollectiveReduceV2") {
absl::string_view uncompilable_reason = "Collective op";
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
if (!op_filter_.allow_where_op && node.type_string() == "Where") {
absl::string_view uncompilable_reason = "Where op";
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
if (!op_filter_.allow_unique_op && node.type_string() == "Unique") {
absl::string_view uncompilable_reason = "Unique op";
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
if (!op_filter_.allow_ops_producing_or_consuming_variant &&
OpProducesOrConsumesVariant(node)) {
absl::string_view uncompilable_reason = "DT_VARIANT producer/consumer";
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
if (!op_filter_.allow_stack_ops && IsStackOp(node)) {
absl::string_view uncompilable_reason = "Stack op";
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
if (!op_filter_.allow_tensor_array_ops && IsTensorArrayOp(node)) {
absl::string_view uncompilable_reason = "TensorArray op";
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
if (!op_filter_.allow_resource_ops_in_called_functions && stack_depth > 1 &&
HasResourceInput(node)) {
absl::string_view uncompilable_reason =
"resource variable op in called function";
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
if (!op_filter_.allow_inaccurate_ops && OpIsInaccurate(node)) {
absl::string_view uncompilable_reason =
"operation with numerical accuracy issues";
BroadcastOptimizationRemark(XlaOptimizationRemark::INACCURATE_OPERATION,
node.DebugString())
.IgnoreError();
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
if (!op_filter_.allow_slow_ops && OpIsSlow(node)) {
absl::string_view uncompilable_reason = "slow operation";
BroadcastOptimizationRemark(XlaOptimizationRemark::SLOW_OPERATION,
node.DebugString())
.IgnoreError();
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
return true;
}
RecursiveCompilabilityChecker::OperationFilter CreateOperationFilter(
const XlaOpRegistry::DeviceRegistration& registration) {
RecursiveCompilabilityChecker::OperationFilter op_filter;
op_filter.allow_resource_ops_in_called_functions =
registration.cluster_resource_variable_ops_unsafely;
op_filter.allow_stack_ops = registration.cluster_stack_ops;
op_filter.allow_tensor_array_ops = registration.cluster_tensor_array_ops;
op_filter.allow_stateful_rng_ops = registration.cluster_stateful_rng_ops;
op_filter.allow_control_trigger = registration.cluster_control_trigger;
op_filter.allow_eliding_assert_and_checknumerics_ops =
registration.elide_assert_and_checknumerics;
op_filter.allow_ops_producing_or_consuming_variant =
registration.cluster_variant_ops;
op_filter.allow_slow_ops = registration.cluster_slow_ops;
op_filter.allow_inaccurate_ops = registration.cluster_inaccurate_ops;
return op_filter;
}
/*static*/ void RecursiveCompilabilityChecker::MaybeMarkUncompilableNode(
const absl::string_view reason,
const std::vector<StackFrameView>& stack_trace,
NameAttrList* encapsulating_function,
RecursiveCompilabilityChecker::UncompilableNodesMap* uncompilable_nodes) {
if (!uncompilable_nodes) return;
UncompilableNodeInfo node_info;
node_info.uncompilable_reason = std::string(reason);
absl::c_transform(stack_trace, std::back_inserter(node_info.stack_trace),
[](const StackFrameView& stack_element) {
return StackFrame{
std::string(stack_element.name),
std::string(stack_element.function_name),
stack_element.stack_trace};
});
node_info.name = std::string(stack_trace.back().name);
auto function =
encapsulating_function ? *encapsulating_function : NameAttrList();
auto function_identifier = function.ShortDebugString();
auto it = uncompilable_nodes->find(function_identifier);
if (it == uncompilable_nodes->end()) {
std::vector<RecursiveCompilabilityChecker::UncompilableNodeInfo>
uncompilable_node_info{std::move(node_info)};
uncompilable_nodes->emplace(
std::move(function_identifier),
std::make_pair(function, std::move(uncompilable_node_info)));
} else {
it->second.second.emplace_back(std::move(node_info));
}
}
// Returns `true` iff node has a given `attr` set to `true`. Returns `false`
// both for the missing attr, and the attr set to `false`.
static bool HasBoolAttr(const NodeDef& node, const char* attr) {
const auto& it = node.attr().find(attr);
return it != node.attr().end() && it->second.b();
}
bool CanCreateXlaKernel(const NodeDef& node_def) {
return HasBoolAttr(node_def, kXlaMustCompileAttr);
}
absl::Status GetBodyAndConstantsAndResources(
FunctionLibraryRuntime* flr, const NameAttrList& function,
const FunctionBody** fbody, std::vector<int>* constant_arg_indices,
std::vector<int>* resource_arg_indices) {
FunctionLibraryRuntime::Handle handle;
TF_RETURN_IF_ERROR(
flr->Instantiate(function.name(), AttrSlice(&function.attr()), &handle));
*fbody = flr->GetFunctionBody(handle);
CHECK(*fbody); // Can't be nullptr since we just instantiated it.
const DataTypeVector& arg_types = (*fbody)->arg_types;
std::vector<bool> const_args(arg_types.size());
// If we can't analyze the const args. Bail out.
TF_RETURN_IF_ERROR(
BackwardsConstAnalysis(*((*fbody)->graph), &const_args,
/*compile_time_const_nodes=*/nullptr, flr));
for (size_t i = 0; i < const_args.size(); ++i) {
if (const_args[i]) {
constant_arg_indices->push_back(i);
}
}
// There can be hundreds of resource variables. Reserve the space for them.
// We don't reserve for constants above as they are usually few.
resource_arg_indices->reserve(arg_types.size());
for (size_t i = 0; i < arg_types.size(); ++i) {
if (arg_types[i] == DT_RESOURCE) {
resource_arg_indices->push_back(i);
}
}
return absl::OkStatus();
}
tensorflow::MemoryTypeVector GetInputMemoryTypes(
const tensorflow::FunctionBody* fbody,
absl::Span<int const> constant_arg_indices,
absl::Span<int const> resource_arg_indices) {
// Set input and output memory types.
tensorflow::MemoryTypeVector input_memory_types(fbody->arg_types.size(),
tensorflow::DEVICE_MEMORY);
// These indices are used only for optimization purposes. They allow us
// to loop over constant_arg_indices and resource_arg_indices only once
// while iterating over all the function arguments checking if it is a
// resource or a constant.
// The reason we optimized this code is because functions can have a lot of
// captured arguments. For example, the backward pass of ResNet50 takes in all
// 214 variables and a similar number of activations.
SinglePassSearch constants_search(constant_arg_indices);
SinglePassSearch resources_search(resource_arg_indices);
for (size_t i = 0; i < fbody->arg_types.size(); ++i) {
if (resources_search.ScanForValue(i) || constants_search.ScanForValue(i)) {
// Compile-time constants and resource handles are expected to be in
// host memory.
input_memory_types[i] = tensorflow::HOST_MEMORY;
}
}
return input_memory_types;
}
tensorflow::MemoryTypeVector GetOutputMemoryTypes(
const tensorflow::FunctionBody* fbody) {
tensorflow::MemoryTypeVector output_memory_types(fbody->ret_types.size(),
tensorflow::DEVICE_MEMORY);
for (size_t i = 0; i < fbody->ret_types.size(); ++i) {
if (fbody->ret_types[i] == tensorflow::DT_RESOURCE) {
output_memory_types[i] = tensorflow::HOST_MEMORY;
}
}
return output_memory_types;
}
static auto const ops_triggering_xla_compilation =
new absl::flat_hash_set<std::string>{"XlaBroadcastHelper",
"XlaCallModule",
"XlaConv",
"XlaConvV2",
"XlaDequantize",
"XlaDot",
"XlaDotV2",
"XlaDynamicSlice",
"XlaDynamicUpdateSlice",
"XlaEinsum",
"XlaGather",
"XlaIf",
"XlaKeyValueSort",
"XlaPad",
"XlaRecv",
"XlaReduce",
"XlaReduceWindow",
"XlaReplicaId",
"XlaRngBitGenerator",
"XlaScatter",
"XlaSelectAndScatter",
"XlaSelfAdjointEig",
"XlaSend",
"XlaSharding",
"XlaSort",
"XlaSpmdFullToShardShape",
"XlaSpmdShardToFullShape",
"XlaSvd",
"XlaVariadicReduceV2",
"XlaVariadicSort",
"XlaWhile"};
bool NodeCanTriggerXlaCompilation(const NodeDef& node) {
return node.attr().find(kXlaClusterIdAttr) != node.attr().end() ||
HasBoolAttr(node, kXlaMustCompileAttr) ||
HasBoolAttr(node, kXlaCompileAttr) ||
HasBoolAttr(node, kXlaScopeAttr) ||
HasBoolAttr(node, kXlaInternalScopeAttr) ||
ops_triggering_xla_compilation->count(node.op());
}
bool CanTriggerXlaCompilation(const GraphDef& graph) {
for (const FunctionDef& function : graph.library().function()) {
for (const NodeDef& node : function.node_def()) {
if (NodeCanTriggerXlaCompilation(node)) {
return true;
}
}
}
for (const NodeDef& node : graph.node()) {
if (NodeCanTriggerXlaCompilation(node)) {
return true;
}
}
return false;
}
} // namespace tensorflow
tensorflow/compiler/jit/compilability_check_util.h
+340
View File
@@ -0,0 +1,340 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_JIT_COMPILABILITY_CHECK_UTIL_H_
#define TENSORFLOW_COMPILER_JIT_COMPILABILITY_CHECK_UTIL_H_
#include <string>
#include "absl/algorithm/container.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "absl/types/optional.h"
#include "tensorflow/compiler/jit/defs.h"
#include "tensorflow/compiler/jit/device_util.h"
#include "tensorflow/compiler/jit/flags.h"
#include "tensorflow/compiler/jit/resource_operation_safety_analysis.h"
#include "tensorflow/compiler/tf2xla/const_analysis.h"
#include "tensorflow/compiler/tf2xla/resource_operation_table.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "xla/service/graphcycles/graphcycles.h"
#include "xla/union_find.h"
#include "xla/util.h"
#include "tensorflow/core/common_runtime/function.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/bounds_check.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/graph_def_util.h"
#include "tensorflow/core/framework/memory_types.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/graph/algorithm.h"
#include "tensorflow/core/graph/control_flow.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/lib/gtl/cleanup.h"
#include "tensorflow/core/lib/strings/stringprintf.h"
#include "tensorflow/core/public/version.h"
#include "tensorflow/core/util/dump_graph.h"
namespace tensorflow {
// Checks whether a TF node can be compiled or not. "Recursive" as in for call
// and functional while nodes it recursively checks whether the callee functions
// can be compiled.
class RecursiveCompilabilityChecker {
public:
// Contains node name and function name. If the node is not inside a function
// body, function name is an empty string.
struct StackFrame {
std::string name;
std::string function_name;
std::shared_ptr<AbstractStackTrace> stack_trace;
};
// Contains information about uncompilable node inside a function body.
struct UncompilableNodeInfo {
std::string name;
// A list representing a stacktrace from the highest level node in
// increasing call depth to immediate node that fails the
// compilability checker.
std::vector<StackFrame> stack_trace;
std::string uncompilable_reason;
};
// Aggregates information about what kinds of ops are allowed.
struct OperationFilter { // TODO(lzr): Add AllowEverything() helper.
// Whether resource variable ops are allowed are allowed in callees. We do
// not allow resource variable ops in called functions (either as direct TF
// calls or as higher order control flow ops) because we do not yet model
// their memory effects in jit/resource_operation_safety_analysis.
bool allow_resource_ops_in_called_functions = false;
// Whether Stack operations are allowed. We avoid auto-clustering Stack
// operations in general because we do not support snapshotting them.
//
// TODO(b/112837194): This restriction can be lifted with some work.
bool allow_stack_ops = false;
// Whether TensorArray operations are allowed. We avoid auto-clustering
// TensorArray operations in general because we do not support snapshotting
// them.
//
// TODO(b/112837194): This restriction can be lifted with some work.
bool allow_tensor_array_ops = false;
// Whether stateful RNG ops are allowed. XLA's RNG does not have the same
// seeding behavior as TensorFlow's RNG (b/34749654). So we avoid
// auto-clustering stateful RNG ops.
bool allow_stateful_rng_ops = false;
// TODO(b/118970344): Whether ControlTrigger ops are allowed. It is unsound
// to cluster ControlTrigger because of how we use deadness analysis.
bool allow_control_trigger = false;
// Whether it is okay to "cluster" Assert and CheckNumerics by simply
// removing them (they're not removed during clustering, but their
// XlaOpKernel is a no-op kernel). We avoid auto-clustering these ops so
// that the user is not surprised when XLA is implicitly enabled. If the
// user explicitly specifies to use XLA, it is fine to resort to a dummy
// implementation. Currently Assert and CheckNumerics ops have dummy XLA
// implementations.
bool allow_eliding_assert_and_checknumerics_ops = false;
// Whether ops that produce or consume DT_VARIANT values are allowed. We
// don't auto-cluster these ops because we don't yet support live-in or
// live-out DT_VARIANT values.
bool allow_ops_producing_or_consuming_variant = false;
// Whether ops known to be slow on XLA-GPU should be considered compilable.
bool allow_slow_ops = false;
// Whether ops known to have numerical accuracy issues should be considered
// compilable..
bool allow_inaccurate_ops = false;
// Require the function to be always compilable, regardless whether some
// control flow branches might be dead for a given input.
bool require_always_compilable = false;
// Whether string constants are compilable.
bool allow_string_consts = true;
// Whether to allow the compilation of CollectiveReduceV2Op.
bool allow_collective_reduce_v2 = true;
// Whether to allow the compilation of WhereOp.
bool allow_where_op = true;
// Whether to allow the compilation of UniqueOp. Compilation of the UniqueOp
// generates output with bounded dynamic shape that may cause failures with
// auto clustering.
// TODO(b/209813421): Enable tf.unique during
// autoclustering once all failures are rfixed.
bool allow_unique_op = true;
// Whether ops that are marked as outside compiled are always considered
// compilable.
// TODO(b/191502757): Make this behavior true by default and remove this
// option once inference converter supports outside compilation.
bool allow_outside_compiled = false;
};
RecursiveCompilabilityChecker(OperationFilter op_filter,
DeviceType jit_device_type)
: op_filter_(std::move(op_filter)),
jit_device_type_(std::move(jit_device_type)) {}
using UncompilableNodesMap =
std::map<std::string,
std::pair<NameAttrList, std::vector<UncompilableNodeInfo>>>;
// Returns a map where the key is the function identifier(short debug
// string) of the function encapsulating the uncompilable nodes, and the
// value is a pair of NameAttrList of the function and a vector of
// uncompilable node info. When uncompilable node is not inside any
// function call nodes, then key is a ShortDebugString() of an empty
// NameAttrList.
//
// Also, when `node` is inside a function body, users can set
// `node_stack_trace` to provide an additional context for `node`'s
// placement within the outer most graph.
UncompilableNodesMap FindUncompilableNodes(
const Node& node, FunctionLibraryRuntime* lib_runtime,
const std::vector<StackFrame>* node_stack_trace = nullptr) const;
// Returns true if `node` can be compiled by XLA.
bool IsCompilableNode(const Node& node,
FunctionLibraryRuntime* lib_runtime) const {
std::vector<StackFrameView> stack_trace;
stack_trace.emplace_back(StackFrameView{node.name(), ""});
return IsCompilableNode(node, lib_runtime, &stack_trace);
}
// Returns true if XLA supports this Op, but we don't want to cluster it (ie:
// due to performance or correctness concerns).
bool OpIsInaccurate(const Node& node) const;
bool OpIsSlow(const Node& node) const;
private:
struct StackFrameView {
absl::string_view name;
absl::string_view function_name;
std::shared_ptr<AbstractStackTrace> stack_trace;
};
bool IsCompilableNode(
const Node& node, FunctionLibraryRuntime* lib_runtime,
std::vector<StackFrameView>* stack_trace,
NameAttrList* encapsulating_function = nullptr,
UncompilableNodesMap* uncompilable_nodes = nullptr) const;
bool IsCompilableCall(
const NodeDef& call_def, FunctionLibraryRuntime* lib_runtime,
std::vector<StackFrameView>* stack_trace,
NameAttrList* encapsulating_function = nullptr,
UncompilableNodesMap* uncompilable_nodes = nullptr) const;
bool IsCompilableIf(const Node& if_node, FunctionLibraryRuntime* lib_runtime,
std::vector<StackFrameView>* stack_trace,
NameAttrList* encapsulating_function,
UncompilableNodesMap* uncompilable_nodes) const;
bool IsCompilableWhile(const Node& while_node,
FunctionLibraryRuntime* lib_runtime,
std::vector<StackFrameView>* stack_trace,
NameAttrList* encapsulating_function,
UncompilableNodesMap* uncompilable_nodes) const;
// Tests whether 'case_node' is compilable. Every operator in all branches
// must be compilable.
bool IsCompilableCase(const Node& case_node,
FunctionLibraryRuntime* lib_runtime,
std::vector<StackFrameView>* stack_trace,
NameAttrList* encapsulating_function,
UncompilableNodesMap* uncompilable_nodes) const;
// Returns compilability of node def retrieved from `node`'s attribute with
// name `attr_name`.
bool ExtractNodeDefAndCheckCompilability(
const Node& node, const std::string& attr_name,
const std::string& call_name, NameAttrList* encapsulating_function,
FunctionLibraryRuntime* lib_runtime,
std::vector<StackFrameView>* stack_trace,
UncompilableNodesMap* uncompilable_nodes) const;
bool IsStackOp(const Node& node) const {
const XlaResourceOpInfo* op_info =
GetResourceOpInfoForOp(node.type_string());
return op_info && op_info->resource_kind() == XlaResourceKind::kStack;
}
bool IsTensorArrayOp(const Node& node) const {
const XlaResourceOpInfo* op_info =
GetResourceOpInfoForOp(node.type_string());
return op_info && op_info->resource_kind() == XlaResourceKind::kTensorArray;
}
bool IsAssertOrCheckNumerics(absl::string_view op_name) const {
return op_name == "Assert" || op_name == "CheckNumerics";
}
bool IsStatefulRandomOp(absl::string_view op_name) const {
return op_name == "RandomUniform" || op_name == "RandomShuffle" ||
op_name == "RandomUniformInt" || op_name == "RandomStandardNormal" ||
op_name == "TruncatedNormal" || op_name == "Multinomial";
}
bool OpProducesOrConsumesVariant(const Node& node) const {
auto is_variant = [](DataType dtype) { return dtype == DT_VARIANT; };
return absl::c_any_of(node.input_types(), is_variant) ||
absl::c_any_of(node.output_types(), is_variant);
}
bool HasXLAKernel(const Node& node,
std::string* uncompilable_reason = nullptr) const;
static void MaybeMarkUncompilableNode(
const absl::string_view reason,
const std::vector<StackFrameView>& stack_trace,
NameAttrList* encapsulating_function,
UncompilableNodesMap* uncompilable_nodes_map);
// Make sure we don't recurse infinitely on recursive functions.
const size_t kMaxRecursionDepth = 50;
const OperationFilter op_filter_;
const DeviceType jit_device_type_;
};
RecursiveCompilabilityChecker::OperationFilter CreateOperationFilter(
const XlaOpRegistry::DeviceRegistration& registration);
// Given a FunctionLibraryRuntime and a `function`, returns this function's body
// in `fbody` as well as the indices of its constant and resource arguments.
// `fbody` is owned by `flr`.
// `constant_arg_indices` and `resource_arg_indices` should be empty vector.
// They are sorted in ascending order on this function's return.
absl::Status GetBodyAndConstantsAndResources(
FunctionLibraryRuntime* flr, const NameAttrList& function,
const FunctionBody** fbody, std::vector<int>* constant_arg_indices,
std::vector<int>* resource_arg_indices);
// Given a NodeDef `node_def` returns true iff `node_def` has kXlaCompileAttr
// set.
bool CanCreateXlaKernel(const NodeDef& node_def);
// Returns memory types for the input.
// `constant_arg_indices` and `resource_arg_indices` are sorted arrays of
// indices corresponding to constant and resource arguments respectively.
//
// One might wonder, about the case where a compile-time constant argument
// (which must be in host memory) is also used as an input into an op,
// e.g. `Add`, that expects its inputs in device memory. Here is how it
// works now.
// First, what do we mean by "op expects an input in XYZ memory"?
// There are two types of "ops" here: the tf2xla kernel and the HLO
// computation it builds. The tf2xla kernel needs to retrieve the actual
// numeric value of the compile-time constant tensors, so it really expects
// them to be on in host memory. However, for other inputs, it refers to them
// using xla::ComputationDataHandle, which is just a symbolic handle that
// xla::ComputationBuilder assigns. How does this handle gets assigned for
// constant arguments? Even constant arguments get an _Arg node in the graph
// instantiated for Function compilation. The tf2xla kernel for constant _Arg
// nodes takes the constant value, converts it to XlaLiteral, and feeds it
// to xla::ComputationBuilder.ConstantLiteral, which returns the handle. This
// constant XlaLiteral is included in the HLO graph, and subsequently, in
// the actual executable, which is copied to the device before being
// executed. Thus, when this executable runs, the constant is available in
// device memory.
tensorflow::MemoryTypeVector GetInputMemoryTypes(
const tensorflow::FunctionBody* fbody,
absl::Span<int const> constant_arg_indices,
absl::Span<int const> resource_arg_indices);
// Returns output memory types.
//
// XlaLaunch kernel keeps all outputs (including constants, which it copies),
// in device memory except for resources.
tensorflow::MemoryTypeVector GetOutputMemoryTypes(
const tensorflow::FunctionBody* fbody);
// Check whether graph can trigger XLA compilation.
bool CanTriggerXlaCompilation(const GraphDef& graph);
// Returns true iff the node can trigger XLA compilation.
bool NodeCanTriggerXlaCompilation(const NodeDef& node);
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_COMPILABILITY_CHECK_UTIL_H_
tensorflow/compiler/jit/compilability_check_util_test.cc
+656
View File
@@ -0,0 +1,656 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/jit/compilability_check_util.h"
#include <algorithm>
#include <memory>
#include <vector>
#include "absl/log/check.h"
#include "absl/strings/match.h"
#include "absl/types/span.h"
#include "tensorflow/cc/framework/ops.h"
#include "tensorflow/cc/framework/scope.h"
#include "tensorflow/cc/ops/array_ops.h"
#include "tensorflow/cc/ops/functional_ops.h"
#include "tensorflow/compiler/jit/defs.h"
#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "tensorflow/core/common_runtime/graph_def_builder_util.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/graph/graph_def_builder.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/public/version.h"
namespace tensorflow {
namespace {
AttrValue FuncListAttr(const absl::Span<const char* const> names) {
AttrValue attr;
for (const char* name : names) {
attr.mutable_list()->add_func()->set_name(name);
}
return attr;
}
constexpr char kFunctionalIfNodeName[] = "If";
constexpr char kFunctionalCaseNodeName[] = "Case";
constexpr char kFunctionalWhileNodeName[] = "While";
constexpr char kCompilableFunctionName[] = "CompilableFn";
constexpr char kCompilableFunctionNodeName[] = "n_c";
constexpr char kUncompilableFunctionName[] = "UncompilableFn";
constexpr char kUncompilableFunctionNodeName[] = "n_c_uncompilable";
constexpr char kUncompilableFunctionTwoName[] = "UncompilableFnTwo";
constexpr char kUncompilableFunctionNodeTwoName[] = "n_d_uncompilable";
constexpr char kNonMaxSuppressionNodeName[] = "NonMaxSuppression";
// A dummy OpKernel for testing.
class DummyCompilableOp : public XlaOpKernel {
public:
explicit DummyCompilableOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
void Compile(XlaOpKernelContext* ctx) override {
ctx->SetOutput(0, ctx->Input(0));
}
};
// Register the DummyCompilableOp kernel for CPU.
REGISTER_OP("InputFloatOp").Output("o: float");
REGISTER_OP("InputInt32Op").Output("o: int32");
REGISTER_OP("CompilableOp").Input("i: float").Output("o: float");
REGISTER_XLA_OP(Name("CompilableOp").Device(DEVICE_CPU_XLA_JIT),
DummyCompilableOp);
// Dummy op that is uncompilable in CPU.
REGISTER_OP("MissingKernel").Input("i: float").Output("o: float");
class CompilabilityCheckUtilTest : public ::testing::Test {
protected:
void SetUp() override {
XlaOpRegistry::RegisterCompilationKernels();
op_filter_.allow_resource_ops_in_called_functions = false;
op_filter_.allow_stack_ops = false;
op_filter_.allow_tensor_array_ops = false;
op_filter_.allow_stateful_rng_ops = false;
op_filter_.allow_control_trigger = false;
op_filter_.allow_eliding_assert_and_checknumerics_ops = false;
op_filter_.allow_ops_producing_or_consuming_variant = false;
op_filter_.allow_inaccurate_ops = false;
op_filter_.allow_slow_ops = false;
op_filter_.allow_outside_compiled = false;
checker_ = CreateCompilabilityChecker();
}
std::unique_ptr<RecursiveCompilabilityChecker> CreateCompilabilityChecker() {
return std::make_unique<RecursiveCompilabilityChecker>(op_filter_,
device_type_);
}
FunctionLibraryRuntime* GetFunctionLibraryRuntime() {
OptimizerOptions opts;
pflr_ = std::make_unique<ProcessFunctionLibraryRuntime>(
nullptr, Env::Default(), /*config=*/nullptr, TF_GRAPH_DEF_VERSION,
flib_def_.get(), opts);
return pflr_->GetFLR(ProcessFunctionLibraryRuntime::kDefaultFLRDevice);
}
RecursiveCompilabilityChecker::OperationFilter op_filter_;
DeviceType device_type_ = DeviceType(DEVICE_CPU_XLA_JIT);
std::unique_ptr<FunctionDefLibrary> func_library_ =
std::make_unique<FunctionDefLibrary>();
std::unique_ptr<FunctionLibraryDefinition> flib_def_ =
std::make_unique<FunctionLibraryDefinition>(OpRegistry::Global(),
*func_library_);
std::unique_ptr<RecursiveCompilabilityChecker> checker_;
std::unique_ptr<ProcessFunctionLibraryRuntime> pflr_;
};
TEST_F(CompilabilityCheckUtilTest, CheckNonFunctionalNodes) {
GraphDefBuilder builder(GraphDefBuilder::kFailImmediately);
auto opts = builder.opts();
Node* const0 = ops::SourceOp("InputFloatOp", opts);
Node* compilable_op = ops::UnaryOp("CompilableOp", const0, opts);
Node* uncompilable_op = ops::UnaryOp("MissingKernel", compilable_op, opts);
GraphDef graph_def;
TF_EXPECT_OK(builder.ToGraphDef(&graph_def));
auto* flib_runtime = GetFunctionLibraryRuntime();
// Source node is not compilable.
EXPECT_FALSE(checker_->IsCompilableNode(*const0, flib_runtime));
EXPECT_TRUE(checker_->IsCompilableNode(*compilable_op, flib_runtime));
// Uncompilable as we are only checking compilability in CPU device type.
EXPECT_FALSE(checker_->IsCompilableNode(*uncompilable_op, flib_runtime));
const auto uncompilable_nodes =
checker_->FindUncompilableNodes(*uncompilable_op, flib_runtime);
ASSERT_EQ(1, uncompilable_nodes.size());
auto node_info_it =
uncompilable_nodes.find(NameAttrList().ShortDebugString());
ASSERT_NE(uncompilable_nodes.end(), node_info_it);
const auto& uncompilable_nodes_inside_function = node_info_it->second.second;
ASSERT_EQ(1, uncompilable_nodes_inside_function.size());
const auto& uncompilable_node_info = uncompilable_nodes_inside_function.at(0);
EXPECT_TRUE(absl::StrContains(uncompilable_node_info.uncompilable_reason,
"unsupported op"));
ASSERT_EQ(1, uncompilable_node_info.stack_trace.size());
ASSERT_EQ("", uncompilable_node_info.stack_trace.at(0).function_name);
}
TEST_F(CompilabilityCheckUtilTest, CheckOutsideCompiledNode) {
GraphDefBuilder builder(GraphDefBuilder::kFailImmediately);
auto opts = builder.opts();
Node* const0 = ops::SourceOp("InputFloatOp", opts);
Node* uncompilable_op = ops::UnaryOp("MissingKernel", const0, opts);
uncompilable_op->AddAttr("_xla_outside_compilation", "0");
GraphDef graph_def;
TF_EXPECT_OK(builder.ToGraphDef(&graph_def));
auto* flib_runtime = GetFunctionLibraryRuntime();
// Outside compiled ops are considered by default..
EXPECT_FALSE(checker_->IsCompilableNode(*uncompilable_op, flib_runtime));
const auto uncompilable_nodes =
checker_->FindUncompilableNodes(*uncompilable_op, flib_runtime);
ASSERT_EQ(1, uncompilable_nodes.size());
op_filter_.allow_outside_compiled = true;
checker_ = CreateCompilabilityChecker();
// With filter option outside compiled ops are ignored and considered
// compilable.
EXPECT_TRUE(checker_->IsCompilableNode(*uncompilable_op, flib_runtime));
const auto uncompilable_nodes2 =
checker_->FindUncompilableNodes(*uncompilable_op, flib_runtime);
ASSERT_EQ(0, uncompilable_nodes2.size());
}
TEST_F(CompilabilityCheckUtilTest, CheckSimpleFunctionNode) {
FunctionDefLibrary flib;
*flib.add_function() = FunctionDefHelper::Define(
/*Function*/ kUncompilableFunctionName,
/*Inputs*/ {"n_a:float"},
/*Outputs*/ {"n_c_uncompilable:float"},
/*Attributes*/ {},
// Node info
{{{kUncompilableFunctionNodeName}, "MissingKernel", {"n_a"}}});
flib_def_.reset(new FunctionLibraryDefinition(OpRegistry::Global(), flib));
GraphDefBuilder builder(GraphDefBuilder::kFailImmediately, flib_def_.get());
std::unique_ptr<Graph> graph(new Graph(flib_def_.get()));
Node* const0 = ops::SourceOp("InputFloatOp", builder.opts());
Node* functional_node = ops::UnaryOp(kUncompilableFunctionName, const0,
builder.opts().WithName("D"));
TF_EXPECT_OK(GraphDefBuilderToGraph(builder, graph.get()));
auto* flib_runtime = GetFunctionLibraryRuntime();
EXPECT_FALSE(checker_->IsCompilableNode(*functional_node, flib_runtime));
const auto uncompilable_nodes =
checker_->FindUncompilableNodes(*functional_node, flib_runtime);
EXPECT_EQ(1, uncompilable_nodes.size());
NameAttrList function;
function.set_name(kUncompilableFunctionName);
const auto node_info_it =
uncompilable_nodes.find(function.ShortDebugString());
ASSERT_NE(uncompilable_nodes.end(), node_info_it);
const auto& uncompilable_node_list = node_info_it->second.second;
ASSERT_EQ(1, uncompilable_node_list.size());
const auto& node_info = uncompilable_node_list.at(0);
const auto& node_stack = node_info.stack_trace;
ASSERT_EQ(2, node_stack.size());
EXPECT_EQ("D", node_stack.at(0).name);
EXPECT_EQ(kUncompilableFunctionNodeName, node_stack.at(1).name);
EXPECT_EQ(kUncompilableFunctionNodeName, node_info.name);
EXPECT_TRUE(
absl::StrContains(node_info.uncompilable_reason, "unsupported op"));
}
TEST_F(CompilabilityCheckUtilTest, CheckFunctionalWhileNode) {
FunctionDefLibrary flib;
*flib.add_function() = FunctionDefHelper::Define(
/*Function*/ kCompilableFunctionName,
/*Inputs*/ {"n_a:float", "n_b:float"},
/*Outputs*/ {"n_c:float"},
/*Attribute*/ {},
// Node info
{{{kCompilableFunctionNodeName},
"Add",
{"n_a", "n_b"},
{{"T", DT_FLOAT}}}});
*flib.add_function() = FunctionDefHelper::Define(
/*Function*/ kUncompilableFunctionName,
/*Inputs*/ {"n_a:float"},
/*Outputs*/ {"n_c_uncompilable:float"},
/*Attributes*/ {},
// Node info
{{{kUncompilableFunctionNodeName}, "MissingKernel", {"n_a"}}});
flib_def_.reset(new FunctionLibraryDefinition(OpRegistry::Global(), flib));
GraphDefBuilder builder(GraphDefBuilder::kFailImmediately, flib_def_.get());
Node* const0 = ops::SourceOp("InputFloatOp", builder.opts());
Node* input_node = ops::UnaryOp("CompilableOp", const0, builder.opts());
NameAttrList compilable;
compilable.set_name(kCompilableFunctionName);
NameAttrList uncompilable;
uncompilable.set_name(kUncompilableFunctionName);
NodeBuilder while_builder(kFunctionalWhileNodeName, "While",
builder.opts().op_registry());
while_builder.Input({input_node, input_node})
.Attr("cond", compilable)
.Attr("body", uncompilable);
builder.opts().FinalizeBuilder(&while_builder);
GraphDef graph_def;
TF_EXPECT_OK(builder.ToGraphDef(&graph_def));
std::unique_ptr<Graph> graph(new Graph(flib_def_.get()));
CHECK_OK(GraphDefBuilderToGraph(builder, graph.get()));
auto while_node_it = std::find_if(
graph->nodes().begin(), graph->nodes().end(),
[&](const Node* n) { return n->name() == kFunctionalWhileNodeName; });
EXPECT_NE(while_node_it, graph->nodes().end());
auto* flib_runtime = GetFunctionLibraryRuntime();
EXPECT_FALSE(checker_->IsCompilableNode(**while_node_it, flib_runtime));
const auto uncompilable_nodes =
checker_->FindUncompilableNodes(**while_node_it, flib_runtime);
ASSERT_EQ(1, uncompilable_nodes.size());
NameAttrList function;
function.set_name(kUncompilableFunctionName);
const auto node_info_it =
uncompilable_nodes.find(function.ShortDebugString());
ASSERT_NE(uncompilable_nodes.end(), node_info_it);
const auto& uncompilable_node_list = node_info_it->second.second;
ASSERT_EQ(1, uncompilable_node_list.size());
const auto& node_info = uncompilable_node_list.at(0);
const auto& node_stack = node_info.stack_trace;
ASSERT_EQ(2, node_stack.size());
const auto& stacktrace_first_node_info = node_stack.at(0);
EXPECT_EQ(kFunctionalWhileNodeName, stacktrace_first_node_info.name);
EXPECT_EQ("", stacktrace_first_node_info.function_name);
const auto& stacktrace_second_node_info = node_stack.at(1);
EXPECT_EQ(kUncompilableFunctionNodeName, stacktrace_second_node_info.name);
EXPECT_EQ(kUncompilableFunctionName,
stacktrace_second_node_info.function_name);
EXPECT_EQ(kUncompilableFunctionNodeName, node_info.name);
EXPECT_TRUE(
absl::StrContains(node_info.uncompilable_reason, "unsupported op"));
}
TEST_F(CompilabilityCheckUtilTest, CheckFunctionalIfNode) {
FunctionDefLibrary flib;
*flib.add_function() = FunctionDefHelper::Define(
/*Function*/ kUncompilableFunctionName,
/*Inputs*/ {"n_a:float"},
/*Outputs*/ {"n_c_uncompilable:float"},
/*Attributes*/ {},
// Node info
{{{kUncompilableFunctionNodeName}, "MissingKernel", {"n_a"}}});
*flib.add_function() = FunctionDefHelper::Define(
/*Function*/ kUncompilableFunctionTwoName,
/*Inputs*/ {"n_a:float"},
/*Outputs*/ {"n_d_uncompilable:float"},
/*Attribute*/ {},
// Node info
{{{kUncompilableFunctionNodeTwoName}, "MissingKernel", {"n_a"}}});
NameAttrList uncompilable_fn1_attr;
uncompilable_fn1_attr.set_name(kUncompilableFunctionName);
NameAttrList uncompilable_fn2_attr;
uncompilable_fn2_attr.set_name(kUncompilableFunctionTwoName);
Scope root = Scope::NewRootScope().ExitOnError();
TF_ASSERT_OK(root.graph()->AddFunctionLibrary(flib));
auto predicate = ops::Placeholder(root.WithOpName("pred"), DT_BOOL);
auto placeholder = ops::Placeholder(root.WithOpName("A"), DT_INT32);
std::vector<NodeBuilder::NodeOut> if_inputs(
{NodeBuilder::NodeOut(placeholder.node())});
Node* if_node;
TF_ASSERT_OK(
NodeBuilder(kFunctionalIfNodeName, "If", &root.graph()->flib_def())
.Input(predicate.node())
.Input(if_inputs)
.Attr("then_branch", uncompilable_fn1_attr)
.Attr("else_branch", uncompilable_fn2_attr)
.Attr("Tout", {DT_INT32})
.Finalize(root.graph(), &if_node));
std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
TF_ASSERT_OK(root.ToGraph(graph.get()));
flib_def_.reset(new FunctionLibraryDefinition(OpRegistry::Global(), flib));
auto if_node_it = std::find_if(
graph->nodes().begin(), graph->nodes().end(),
[&](const Node* n) { return n->name() == kFunctionalIfNodeName; });
EXPECT_NE(if_node_it, graph->nodes().end());
auto* flib_runtime = GetFunctionLibraryRuntime();
EXPECT_FALSE(checker_->IsCompilableNode(**if_node_it, flib_runtime));
const auto uncompilable_nodes =
checker_->FindUncompilableNodes(**if_node_it, flib_runtime);
ASSERT_EQ(2, uncompilable_nodes.size());
NameAttrList function_one;
function_one.set_name(kUncompilableFunctionName);
auto it = uncompilable_nodes.find(function_one.ShortDebugString());
ASSERT_NE(uncompilable_nodes.end(), it);
const auto& uncompilable_node_list = it->second.second;
ASSERT_EQ(1, uncompilable_node_list.size());
const auto& uncompilable_node_one = uncompilable_node_list.at(0);
const auto& node_one_stack = uncompilable_node_one.stack_trace;
ASSERT_EQ(2, node_one_stack.size());
const auto& node_one_stacktrace_first_node = node_one_stack.at(0);
EXPECT_EQ(kFunctionalIfNodeName, node_one_stacktrace_first_node.name);
EXPECT_EQ("", node_one_stacktrace_first_node.function_name);
const auto& stacktrace_second_node_info = node_one_stack.at(1);
EXPECT_EQ(kUncompilableFunctionNodeName, stacktrace_second_node_info.name);
EXPECT_EQ(kUncompilableFunctionName,
stacktrace_second_node_info.function_name);
EXPECT_EQ(kUncompilableFunctionNodeName, uncompilable_node_one.name);
EXPECT_TRUE(absl::StrContains(uncompilable_node_one.uncompilable_reason,
"unsupported op"));
NameAttrList function_two;
function_two.set_name(kUncompilableFunctionTwoName);
it = uncompilable_nodes.find(function_two.ShortDebugString());
ASSERT_NE(uncompilable_nodes.end(), it);
const auto& uncompilable_node_two_list = it->second.second;
ASSERT_EQ(1, uncompilable_node_two_list.size());
const auto& uncompilable_node_two = uncompilable_node_two_list.at(0);
const auto& node_two_stack = uncompilable_node_two.stack_trace;
ASSERT_EQ(2, node_two_stack.size());
const auto& node_two_stacktrace_first_node = node_two_stack.at(0);
EXPECT_EQ(kFunctionalIfNodeName, node_two_stacktrace_first_node.name);
EXPECT_EQ("", node_two_stacktrace_first_node.function_name);
const auto& node_two_stacktrace_second_node = node_two_stack.at(1);
EXPECT_EQ(kUncompilableFunctionNodeTwoName,
node_two_stacktrace_second_node.name);
EXPECT_EQ(kUncompilableFunctionTwoName,
node_two_stacktrace_second_node.function_name);
EXPECT_EQ(kUncompilableFunctionNodeTwoName, uncompilable_node_two.name);
EXPECT_TRUE(absl::StrContains(uncompilable_node_one.uncompilable_reason,
"unsupported op"));
}
TEST_F(CompilabilityCheckUtilTest, CheckFunctionalCaseNode) {
FunctionDefLibrary flib;
*flib.add_function() = FunctionDefHelper::Define(
/*Function*/ kUncompilableFunctionName,
/*Inputs*/ {"n_a:float"},
/*Outputs*/ {"n_c_uncompilable:float"},
/*Attributes*/ {},
// Node info
{{{kUncompilableFunctionNodeName}, "MissingKernel", {"n_a"}}});
*flib.add_function() = FunctionDefHelper::Define(
/*Function*/ kUncompilableFunctionTwoName,
/*Inputs*/ {"n_a:float"},
/*Outputs*/ {"n_d_uncompilable:float"},
/*Attribute*/ {},
// Node info
{{{kUncompilableFunctionNodeTwoName}, "MissingKernel", {"n_a"}}});
Scope root = Scope::NewRootScope().ExitOnError();
TF_ASSERT_OK(root.graph()->AddFunctionLibrary(flib));
auto branch_index = ops::Placeholder(root.WithOpName("pred"), DT_INT32);
auto placeholder = ops::Placeholder(root.WithOpName("A"), DT_INT32);
std::vector<NodeBuilder::NodeOut> inputes(
{NodeBuilder::NodeOut(placeholder.node())});
Node* case_node;
TF_ASSERT_OK(
NodeBuilder(kFunctionalCaseNodeName, "Case", &root.graph()->flib_def())
.Input(branch_index.node())
.Input(inputes)
.Attr("branches", FuncListAttr({kUncompilableFunctionName,
kUncompilableFunctionTwoName}))
.Attr("Tout", {DT_INT32})
.Finalize(root.graph(), &case_node));
std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
TF_ASSERT_OK(root.ToGraph(graph.get()));
flib_def_.reset(new FunctionLibraryDefinition(OpRegistry::Global(), flib));
auto case_node_it = std::find_if(
graph->nodes().begin(), graph->nodes().end(),
[&](const Node* n) { return n->name() == kFunctionalCaseNodeName; });
EXPECT_NE(case_node_it, graph->nodes().end());
auto* flib_runtime = GetFunctionLibraryRuntime();
op_filter_.require_always_compilable = false;
checker_ = CreateCompilabilityChecker();
EXPECT_TRUE(checker_->IsCompilableNode(**case_node_it, flib_runtime));
op_filter_.require_always_compilable = true;
checker_ = CreateCompilabilityChecker();
EXPECT_FALSE(checker_->IsCompilableNode(**case_node_it, flib_runtime));
}
TEST_F(CompilabilityCheckUtilTest, TestCanNotTriggerXlaCompilation) {
GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
Scope root = Scope::NewRootScope().ExitOnError();
FunctionDefLibrary library;
FunctionDef identity_func = FunctionDefHelper::Create(
"IdentityFunc",
/*in_def=*/{"x:float"},
/*out_def=*/{"res:float"},
/*attr_def=*/{},
/*node_def=*/{{{"t0"}, "Identity", {"x"}, {{"T", DT_FLOAT}}}},
/*ret_def*/ {{"res", "t0:output"}});
*library.add_function() = identity_func;
Output in = ops::Placeholder(root, DT_FLOAT);
NameAttrList b_name_attr;
b_name_attr.set_name("IdentityFunc");
ops::PartitionedCall call(root.WithOpName("call"), {in}, {DT_FLOAT},
b_name_attr);
GraphDef graph_def;
TF_ASSERT_OK(root.graph()->AddFunctionLibrary(library));
TF_ASSERT_OK(root.ToGraphDef(&graph_def));
EXPECT_FALSE(CanTriggerXlaCompilation(graph_def));
}
TEST_F(CompilabilityCheckUtilTest, TestXlaOpsCanTriggerXlaCompilation) {
GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
Scope root = Scope::NewRootScope().ExitOnError();
FunctionDefLibrary library;
FunctionDef sort_func = FunctionDefHelper::Create(
"SortFunc",
/*in_def=*/{"x:float"},
/*out_def=*/{"res:float"},
/*attr_def=*/{},
/*node_def=*/{{{"t0"}, "XlaSort", {"x"}, {{"T", DT_FLOAT}}}},
/*ret_def*/ {{"res", "t0:output"}});
*library.add_function() = sort_func;
Output in = ops::Placeholder(root, DT_FLOAT);
NameAttrList b_name_attr;
b_name_attr.set_name("SortFunc");
ops::PartitionedCall call(root.WithOpName("call"), {in}, {DT_FLOAT},
b_name_attr);
GraphDef graph_def;
TF_ASSERT_OK(root.graph()->AddFunctionLibrary(library));
TF_ASSERT_OK(root.ToGraphDef(&graph_def));
EXPECT_TRUE(CanTriggerXlaCompilation(graph_def));
}
TEST_F(CompilabilityCheckUtilTest, TestCanTriggerXlaCompilation) {
GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
Scope root = Scope::NewRootScope().ExitOnError();
FunctionDefLibrary library;
AttrValue true_attribute;
true_attribute.set_b(true);
FunctionDef identity_func = FunctionDefHelper::Create(
"IdentityFunc",
/*in_def=*/{"x:float"},
/*out_def=*/{"res:float"},
/*attr_def=*/{},
/*node_def=*/{{{"t0"}, "Identity", {"x"}, {{"T", DT_FLOAT}}}},
/*ret_def*/ {{"res", "t0:output"}});
(*identity_func.mutable_attr())[kXlaMustCompileAttr] = true_attribute;
FunctionDef call_identity = FunctionDefHelper::Create(
"CallIdentity",
/*in_def=*/{"x:float"},
/*out_def=*/{"z:float"}, /*attr_def=*/{},
/*node_def=*/
{{{"func_call"},
"PartitionedCall",
{"x"},
{{"Tin", DataTypeSlice({DT_FLOAT})},
{"Tout", DataTypeSlice({DT_FLOAT})},
{"f",
FunctionDefHelper::FunctionRef("IdentityRef", {{"T", DT_FLOAT}})},
{kXlaMustCompileAttr, true}}}},
/*ret_def=*/{{"z", "func_call:output:0"}});
*library.add_function() = identity_func;
*library.add_function() = call_identity;
Output in = ops::Placeholder(root, DT_FLOAT);
NameAttrList b_name_attr;
b_name_attr.set_name("CallIdentity");
ops::PartitionedCall call(root.WithOpName("call"), {in}, {DT_FLOAT},
b_name_attr);
GraphDef graph_def;
TF_ASSERT_OK(root.graph()->AddFunctionLibrary(library));
TF_ASSERT_OK(root.ToGraphDef(&graph_def));
EXPECT_TRUE(CanTriggerXlaCompilation(graph_def));
}
TEST_F(CompilabilityCheckUtilTest, CheckNonMaxSuppressionV3UncompilableSlowOp) {
GraphDefBuilder builder(GraphDefBuilder::kFailImmediately);
auto opts = builder.opts();
Node* boxes = ops::SourceOp("InputFloatOp", opts);
Node* scores = ops::SourceOp("InputFloatOp", opts);
Node* max_output_size = ops::SourceOp("InputInt32Op", opts);
Node* iou_threshold = ops::SourceOp("InputFloatOp", opts);
Node* score_threshold = ops::SourceOp("InputFloatOp", opts);
NodeBuilder non_max_suppression_builder(
kNonMaxSuppressionNodeName, "NonMaxSuppressionV3", opts.op_registry());
non_max_suppression_builder.Input(boxes)
.Input(scores)
.Input(max_output_size)
.Input(iou_threshold)
.Input(score_threshold)
.Attr("T", DT_FLOAT);
Node* non_max_suppression;
non_max_suppression =
builder.opts().FinalizeBuilder(&non_max_suppression_builder);
GraphDef graph_def;
TF_EXPECT_OK(builder.ToGraphDef(&graph_def));
auto* flib_runtime = GetFunctionLibraryRuntime();
EXPECT_FALSE(checker_->IsCompilableNode(*non_max_suppression, flib_runtime));
const auto uncompilable_nodes =
checker_->FindUncompilableNodes(*non_max_suppression, flib_runtime);
ASSERT_EQ(1, uncompilable_nodes.size());
auto node_info_it =
uncompilable_nodes.find(NameAttrList().ShortDebugString());
ASSERT_NE(uncompilable_nodes.end(), node_info_it);
const auto& uncompilable_nodes_inside_function = node_info_it->second.second;
ASSERT_EQ(1, uncompilable_nodes_inside_function.size());
const auto& uncompilable_node_info = uncompilable_nodes_inside_function.at(0);
EXPECT_TRUE(absl::StrContains(uncompilable_node_info.uncompilable_reason,
"slow operation"));
}
TEST_F(CompilabilityCheckUtilTest, CheckNonMaxSuppressionV4UncompilableSlowOp) {
GraphDefBuilder builder(GraphDefBuilder::kFailImmediately);
auto opts = builder.opts();
Node* boxes = ops::SourceOp("InputFloatOp", opts);
Node* scores = ops::SourceOp("InputFloatOp", opts);
Node* max_output_size = ops::SourceOp("InputInt32Op", opts);
Node* iou_threshold = ops::SourceOp("InputFloatOp", opts);
Node* score_threshold = ops::SourceOp("InputFloatOp", opts);
NodeBuilder non_max_suppression_v4_builder(
kNonMaxSuppressionNodeName, "NonMaxSuppressionV4", opts.op_registry());
non_max_suppression_v4_builder.Input(boxes)
.Input(scores)
.Input(max_output_size)
.Input(iou_threshold)
.Input(score_threshold)
.Attr("T", DT_FLOAT);
Node* non_max_suppression_v4;
non_max_suppression_v4 =
builder.opts().FinalizeBuilder(&non_max_suppression_v4_builder);
GraphDef graph_def;
TF_EXPECT_OK(builder.ToGraphDef(&graph_def));
auto* flib_runtime = GetFunctionLibraryRuntime();
EXPECT_FALSE(
checker_->IsCompilableNode(*non_max_suppression_v4, flib_runtime));
const auto uncompilable_nodes =
checker_->FindUncompilableNodes(*non_max_suppression_v4, flib_runtime);
ASSERT_EQ(1, uncompilable_nodes.size());
auto node_info_it =
uncompilable_nodes.find(NameAttrList().ShortDebugString());
ASSERT_NE(uncompilable_nodes.end(), node_info_it);
const auto& uncompilable_nodes_inside_function = node_info_it->second.second;
ASSERT_EQ(1, uncompilable_nodes_inside_function.size());
const auto& uncompilable_node_info = uncompilable_nodes_inside_function.at(0);
EXPECT_TRUE(absl::StrContains(uncompilable_node_info.uncompilable_reason,
"slow operation"));
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/compilation_passes_test_main.cc
+58
View File
@@ -0,0 +1,58 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <memory>
#include "absl/strings/match.h"
#include "absl/strings/string_view.h"
#include "tensorflow/compiler/jit/flags.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
GTEST_API_ int main(int real_argc, char** real_argv) {
std::vector<tensorflow::Flag> flag_list;
tensorflow::AppendMarkForCompilationPassFlags(&flag_list);
auto usage = tensorflow::Flags::Usage(real_argv[0], flag_list);
std::vector<char*> args;
args.reserve(real_argc + 1);
for (int i = 0; i < real_argc; i++) {
args.push_back(real_argv[i]);
}
struct FreeDeleter {
void operator()(char* ptr) { free(ptr); }
};
std::unique_ptr<char, FreeDeleter> enable_global_jit_arg(
strdup("--tf_xla_cpu_global_jit=true"));
args.push_back(enable_global_jit_arg.get());
std::unique_ptr<char, FreeDeleter> reduce_min_cluster_size_arg(
strdup("--tf_xla_min_cluster_size=2"));
args.push_back(reduce_min_cluster_size_arg.get());
int argc = args.size();
if (!tensorflow::Flags::Parse(&argc, &args.front(), flag_list)) {
LOG(ERROR) << "\n" << usage;
return 2;
}
testing::InitGoogleTest(&argc, &args.front());
return RUN_ALL_TESTS();
}
tensorflow/compiler/jit/deadness_analysis.cc
+1610
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/jit/deadness_analysis.h
+99
View File
@@ -0,0 +1,99 @@
/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_JIT_DEADNESS_ANALYSIS_H_
#define TENSORFLOW_COMPILER_JIT_DEADNESS_ANALYSIS_H_
#include "tensorflow/core/graph/graph.h"
namespace tensorflow {
// This analyzes a TensorFlow graph to identify nodes which may have partially
// dead inputs (i.e. these nodes may have some dead inputs and some alive
// inputs).
//
// For example, the ADD node in the following graph
//
// V0 PRED0 V1 PRED1
// | | | |
// v v v v
// SWITCH SWITCH
// | |
// +---+ + ---+
// | |
// v v
// ADD
//
// can have its inputs independently dead or alive based on the runtime values
// of PRED0 and PRED1.
//
// It is tempting to call this a liveness analysis but I avoided that because
// "liveness" already has other connotations.
class DeadnessAnalysis {
public:
// An opaque representation of a predicate. DeadnessPredicate
// instances that compare equal via operator== represent predicates
// that always evaluate to the same value.
struct DeadnessPredicate {
public:
DeadnessPredicate(const DeadnessPredicate&) = default;
DeadnessPredicate(DeadnessPredicate&&) = default;
DeadnessPredicate& operator=(const DeadnessPredicate&) = default;
DeadnessPredicate& operator=(DeadnessPredicate&&) = default;
bool operator==(const DeadnessPredicate& other) const {
return other.pred_ == pred_;
}
bool operator!=(const DeadnessPredicate& other) const {
return other.pred_ != pred_;
}
private:
explicit DeadnessPredicate(void* pred) : pred_(pred) {}
// This is really a Predicate*, but we don't want to expose that
// implementation detail to our clients. `pred_` has pointer equality so we
// can just compare the pointer in operator== and operator!=.
void* pred_;
friend class DeadnessAnalysis;
};
virtual absl::StatusOr<DeadnessPredicate> GetPredicateFor(Node* n,
int oidx) const = 0;
// Prints out the internal state of this instance. For debugging purposes
// only.
virtual void Print() const = 0;
virtual ~DeadnessAnalysis();
std::string DebugString(DeadnessPredicate predicate) const;
// Run the deadness analysis over `graph` and returns an error or a populated
// instance of DeadnessAnalysis in `result`.
static absl::Status Run(const Graph& graph,
std::unique_ptr<DeadnessAnalysis>* result);
protected:
static DeadnessPredicate MakeDeadnessPredicate(void* pred) {
return DeadnessPredicate(pred);
}
};
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_DEADNESS_ANALYSIS_H_
tensorflow/compiler/jit/deadness_analysis_internal.h
+36
View File
@@ -0,0 +1,36 @@
/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_JIT_DEADNESS_ANALYSIS_INTERNAL_H_
#define TENSORFLOW_COMPILER_JIT_DEADNESS_ANALYSIS_INTERNAL_H_
#include "absl/container/flat_hash_map.h"
#include "tensorflow/core/graph/tensor_id.h"
namespace tensorflow {
namespace deadness_analysis_internal {
// Returns a map describing the predicate each Tensor was mapped to. For
// testing purposes only.
using PredicateMapTy =
absl::flat_hash_map<TensorId, std::string, TensorId::Hasher>;
absl::Status ComputePredicates(const Graph& graph,
PredicateMapTy* out_predicate_map,
bool enable_optimistic = true);
} // namespace deadness_analysis_internal
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_DEADNESS_ANALYSIS_INTERNAL_H_
tensorflow/compiler/jit/deadness_analysis_test.cc
+1313
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/jit/defs.cc
+45
View File
@@ -0,0 +1,45 @@
/* 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/jit/defs.h"
#include <atomic>
namespace tensorflow {
const char* const kXlaMustCompileAttr = "_XlaMustCompile";
const char* const kXlaCompileAttr = "_XlaCompile";
// User-provided through jit_scope APIs. Effective only when auto_jit is OFF.
const char* const kXlaScopeAttr = "_XlaScope";
// Automatically inserted by auto_jit to guide clustering results. Effective
// only when auto_jit is ON.
const char* const kXlaInternalScopeAttr = "_XlaInternalScope";
const char* const kXlaClusterIdAttr = "_xla_compile_id";
static std::atomic<bool> xla_devices_creation_required(false);
// Request XLA:GPU and XLA:CPU device creation. Deprecated, only used by XRT
// backend.
void RequestXlaDevicesCreation() { xla_devices_creation_required = true; }
// Check whether XLA:GPU and XLA:CPU device creation was requested. Deprecated,
// only used by XRT backend.
bool XlaDevicesCreationRequired() { return xla_devices_creation_required; }
} // namespace tensorflow
tensorflow/compiler/jit/defs.h
+49
View File
@@ -0,0 +1,49 @@
/* 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.
==============================================================================*/
// Provides definitions needed for use of the TensorFlow XLA
// device.
#ifndef TENSORFLOW_COMPILER_JIT_DEFS_H_
#define TENSORFLOW_COMPILER_JIT_DEFS_H_
namespace tensorflow {
// Name of attribute used to tag operators for compilation with XLA
// Implies must-compile semantics: either it will be compiled
// with XLA, or an error will be thrown.
extern const char* const kXlaMustCompileAttr; // "_XlaMustCompile"
// Implies auto-clustering: tagged nodes will be clustered and compiled with XLA
// on a best-effort basis.
extern const char* const kXlaCompileAttr; // "_XlaCompile"
// Implies auto-clustering within the given scope.
extern const char* const kXlaScopeAttr; // "_XlaScope"
extern const char* const kXlaInternalScopeAttr; // "_XlaInternalScope"
// The id of the compiled cluster.
extern const char* const kXlaClusterIdAttr; // "_xla_compile_id"
[[deprecated("XLA:CPU/GPU devices are deprecated")]] void
RequestXlaDevicesCreation();
[[deprecated("XLA:CPU/GPU devices are deprecated")]] bool
XlaDevicesCreationRequired();
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_DEFS_H_
tensorflow/compiler/jit/device_compilation_cache.h
+280
View File
@@ -0,0 +1,280 @@
/* 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_COMPILER_JIT_DEVICE_COMPILATION_CACHE_H_
#define TENSORFLOW_COMPILER_JIT_DEVICE_COMPILATION_CACHE_H_
#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include "absl/base/nullability.h"
#include "absl/container/flat_hash_map.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/compiler/jit/device_compilation_cluster_signature.h"
#include "tensorflow/compiler/jit/xla_compile_util.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/client/local_client.h"
#include "xla/pjrt/pjrt_client.h"
#include "tensorflow/core/platform/mutex.h"
namespace tensorflow {
namespace device_compilation_cache_internal {
template <typename ExecutableType>
int64_t ExecutableSize(const ExecutableType* executable) {
return 0;
}
template <>
inline int64_t ExecutableSize<xla::LocalExecutable>(
const xla::LocalExecutable* executable) {
if (executable != nullptr && executable->executable() != nullptr) {
return executable->executable()->SizeOfGeneratedCodeInBytes();
}
return 0;
}
template <>
inline int64_t ExecutableSize<xla::PjRtLoadedExecutable>(
const xla::PjRtLoadedExecutable* executable) {
if (executable != nullptr) {
return executable->SizeOfGeneratedCodeInBytes();
}
return 0;
}
} // namespace device_compilation_cache_internal
// Cache to store compiled HLO, executables and related metadata keyed by
// `DeviceCompilationClusterSignature`. The cache owns the stored
// CompilationResults and Executables.
// Currently no cache eviction policy is implemented and the cache grows without
// bound.
template <typename ExecutableType>
class DeviceCompilationCache {
public:
DeviceCompilationCache() = default;
~DeviceCompilationCache() = default;
using Key = DeviceCompilationClusterSignature;
struct Value {
DeviceCompileState compile_state = DeviceCompileState::kUncompiled;
absl::Status compilation_status;
int64_t request_count = 0;
const XlaCompiler::CompilationResult* compilation_result = nullptr;
ExecutableType* executable = nullptr;
};
// Returns std::nullopt if value for the supplied key is not found. If a value
// is found, `request_count` is incremented before returning the value.
std::optional<Value> Lookup(const Key& key) const;
// Inserts an empty value if value is not found and returns it. If a value is
// found, `request_count` is incremented before returning the value.
Value LookupOrCreate(const Key& key);
// Caches `compile_state`, `compilation_status`, `compilation_result` and
// `executable` and associates them with the provided `key`. Takes ownership
// of `compilation_result` and `executable`. Does not increment the
// corresponding `request_count`. Only arguments that are not std::nullopt are
// updated in the cache.
void Store(const Key& key, std::optional<DeviceCompileState> compile_state,
std::optional<absl::Status> compilation_status,
std::optional<std::unique_ptr<XlaCompiler::CompilationResult>>
compilation_result,
std::optional<std::unique_ptr<ExecutableType>> executable);
std::string DebugString() const;
// Erase any cache entries that have a null entry and releases all references
// to `xla::XlaComputation` in the non-null cache entries.
void Finalize() {
const mutex_lock lock(compile_cache_mu_);
absl::erase_if(
cache_,
[&](std::pair<const Key, absl_nullable std::unique_ptr<Entry>>& kv) {
Entry* absl_nullable const entry = kv.second.get();
if (entry == nullptr) {
return true;
}
const mutex_lock entry_lock(entry->mu);
if (entry->compilation_result != nullptr) {
entry->compilation_result->computation.reset();
}
return false;
});
};
private:
// The value associated with a cache entry.
struct Entry {
mutable mutex mu;
// The current compilation state for this entry.
DeviceCompileState compile_state TF_GUARDED_BY(mu) =
DeviceCompileState::kUncompiled;
// The number of times a compilation with this signature has been requested.
int64_t request_count TF_GUARDED_BY(mu) = 0;
// Did compilation succeed?
absl::Status compilation_status TF_GUARDED_BY(mu);
// Output of the XlaCompiler.
std::unique_ptr<XlaCompiler::CompilationResult> compilation_result
TF_GUARDED_BY(mu);
// The XLA executable compiled from <computation>. May be null if no
// executable has been built.
std::unique_ptr<ExecutableType> executable TF_GUARDED_BY(mu);
std::string DebugString() const {
mutex_lock lock(mu);
int64_t executable_size =
device_compilation_cache_internal::ExecutableSize<ExecutableType>(
executable.get());
int64_t hlo_module_size = 0;
if (compilation_result != nullptr &&
compilation_result->computation != nullptr) {
hlo_module_size =
compilation_result->computation->proto().ByteSizeLong();
}
return absl::StrCat(
"{compile_state: ", compile_state, ", request_count: ", request_count,
", compilation_status: ", compilation_status.ToString(),
", compilation_result?: ", compilation_result != nullptr,
", hlo_module_size: ", hlo_module_size, " bytes",
", executable?: ", executable != nullptr,
", executable_size: ", executable_size, " bytes}");
}
};
mutable mutex compile_cache_mu_;
absl::flat_hash_map<Key, std::unique_ptr<Entry>, Key::Hash> cache_
TF_GUARDED_BY(compile_cache_mu_);
DeviceCompilationCache(const DeviceCompilationCache&) = delete;
void operator=(const DeviceCompilationCache&) = delete;
};
template <typename ExecutableType>
std::optional<typename DeviceCompilationCache<ExecutableType>::Value>
DeviceCompilationCache<ExecutableType>::Lookup(const Key& key) const {
// The outer lock protects the existence of the cache entry. It does not
// protect the contents of the cache entry.
Entry* entry;
{
mutex_lock lock(compile_cache_mu_);
// Find cache entry.
auto it = cache_.find(key);
if (it == cache_.cend()) {
return std::nullopt;
}
entry = it->second.get();
}
mutex_lock lock(entry->mu);
Value value = {/*compile_state=*/entry->compile_state,
/*compilation_status=*/entry->compilation_status,
/*request_count=*/++entry->request_count,
/*compilation_result=*/entry->compilation_result.get(),
/*executable=*/entry->executable.get()};
return value;
}
template <typename ExecutableType>
typename DeviceCompilationCache<ExecutableType>::Value
DeviceCompilationCache<ExecutableType>::LookupOrCreate(const Key& key) {
// The outer lock protects the existence of the cache entry. It does not
// protect the contents of the cache entry.
Entry* entry;
{
mutex_lock lock(compile_cache_mu_);
// Emplace empty cache entry if not found.
auto it = cache_.emplace(key, std::make_unique<Entry>()).first;
entry = it->second.get();
}
mutex_lock lock(entry->mu);
Value value = {/*compile_state=*/entry->compile_state,
/*compilation_status=*/entry->compilation_status,
/*request_count=*/++entry->request_count,
/*compilation_result=*/entry->compilation_result.get(),
/*executable=*/entry->executable.get()};
return value;
}
template <typename ExecutableType>
void DeviceCompilationCache<ExecutableType>::Store(
const Key& key, std::optional<DeviceCompileState> compile_state,
std::optional<absl::Status> compilation_status,
std::optional<std::unique_ptr<XlaCompiler::CompilationResult>>
compilation_result,
std::optional<std::unique_ptr<ExecutableType>> executable) {
Entry* entry;
{
mutex_lock lock(compile_cache_mu_);
// Emplace empty cache entry if not found.
auto it = cache_.emplace(key, std::make_unique<Entry>()).first;
entry = it->second.get();
}
{
mutex_lock lock(entry->mu);
if (compile_state.has_value()) {
entry->compile_state = *compile_state;
}
if (compilation_status.has_value()) {
entry->compilation_status = *compilation_status;
}
if (compilation_result.has_value()) {
entry->compilation_result = std::move(*compilation_result);
}
if (executable.has_value()) {
entry->executable = std::move(*executable);
}
}
VLOG(4) << "Added/updated cache entry: key=" << key.HumanString()
<< ", entry=" << entry->DebugString();
}
template <typename ExecutableType>
std::string DeviceCompilationCache<ExecutableType>::DebugString() const {
std::string s = "DeviceCompilationCache<ExecutableType> {\n";
{
mutex_lock lock(compile_cache_mu_);
for (const auto& [key, entry] : cache_) {
absl::StrAppend(&s, key.HumanString(), " : ", entry->DebugString(),
",\n");
}
}
absl::StrAppend(&s, "}");
return s;
}
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_DEVICE_COMPILATION_CACHE_H_
tensorflow/compiler/jit/device_compilation_cache_test.cc
+277
View File
@@ -0,0 +1,277 @@
/* 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/compiler/jit/device_compilation_cache.h"
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "tensorflow/compiler/jit/xla_compile_util.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/hlo/builder/xla_computation.h"
#include "xla/tsl/protobuf/error_codes.pb.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace {
struct FakeExecutable {
std::string data;
explicit FakeExecutable(const std::string& s) : data(s) {}
};
using Cache = DeviceCompilationCache<FakeExecutable>;
using Signature = DeviceCompilationClusterSignature;
absl::StatusOr<Signature> BuildSampleSignature(const std::string& fn_name) {
NameAttrList fn;
fn.set_name(fn_name);
std::vector<XlaCompiler::Argument> args(1);
args[0].kind = XlaCompiler::Argument::kConstant;
args[0].type = DT_INT32;
args[0].shape = TensorShape({4, 0});
args[0].constant_value = Tensor(DT_INT32, {4, 0});
return Signature::Build(fn, args);
}
TEST(DeviceCompilationCacheTest, LookupEntryDoesntExist) {
auto cache = std::make_unique<Cache>();
TF_ASSERT_OK_AND_ASSIGN(auto key, BuildSampleSignature("foo"));
auto cache_value = cache->Lookup(key);
EXPECT_FALSE(cache_value.has_value());
}
TEST(DeviceCompilationCacheTest, LookupOrCreateEntryDoesntExist) {
auto cache = std::make_unique<Cache>();
TF_ASSERT_OK_AND_ASSIGN(auto key, BuildSampleSignature("foo"));
Cache::Value cache_value = cache->LookupOrCreate(key);
EXPECT_EQ(cache_value.compile_state, DeviceCompileState::kUncompiled);
EXPECT_EQ(cache_value.request_count, 1);
EXPECT_EQ(cache_value.compilation_result, nullptr);
EXPECT_EQ(cache_value.executable, nullptr);
}
TEST(DeviceCompilationCacheTest, IncrementRequestCountOnLookup) {
auto cache = std::make_unique<Cache>();
TF_ASSERT_OK_AND_ASSIGN(auto key, BuildSampleSignature("foo"));
Cache::Value cache_value = cache->LookupOrCreate(key);
EXPECT_EQ(cache_value.request_count, 1);
cache_value = cache->LookupOrCreate(key);
EXPECT_EQ(cache_value.request_count, 2);
cache_value = cache->LookupOrCreate(key);
EXPECT_EQ(cache_value.request_count, 3);
}
TEST(DeviceCompilationCacheTest, RequestCountUnchangedOnStore) {
auto cache = std::make_unique<Cache>();
TF_ASSERT_OK_AND_ASSIGN(auto key, BuildSampleSignature("foo"));
Cache::Value cache_value = cache->LookupOrCreate(key);
EXPECT_EQ(cache_value.request_count, 1);
cache_value = cache->LookupOrCreate(key);
EXPECT_EQ(cache_value.request_count, 2);
cache_value = cache->LookupOrCreate(key);
EXPECT_EQ(cache_value.request_count, 3);
auto compilation_result = std::make_unique<XlaCompiler::CompilationResult>();
cache->Store(key, DeviceCompileState::kCompiled, absl::OkStatus(),
std::move(compilation_result), std::nullopt);
cache_value = cache->LookupOrCreate(key);
EXPECT_EQ(cache_value.request_count, 4);
}
TEST(DeviceCompilationCacheTest, StoreLookup) {
auto cache = std::make_unique<Cache>();
TF_ASSERT_OK_AND_ASSIGN(auto key, BuildSampleSignature("foo"));
auto compilation_result = std::make_unique<XlaCompiler::CompilationResult>();
auto executable = std::make_unique<FakeExecutable>("foo_exe");
cache->Store(key, DeviceCompileState::kCompiled, absl::OkStatus(),
std::move(compilation_result), std::move(executable));
auto cache_value = cache->Lookup(key);
EXPECT_EQ(cache_value->compile_state, DeviceCompileState::kCompiled);
EXPECT_EQ(cache_value->request_count, 1);
EXPECT_TRUE(cache_value->compilation_status.ok());
EXPECT_TRUE(cache_value->compilation_result != nullptr);
EXPECT_TRUE(cache_value->executable != nullptr);
EXPECT_EQ(cache_value->executable->data, "foo_exe");
}
TEST(DeviceCompilationCacheTest, StoreLookupOrCreate) {
auto cache = std::make_unique<Cache>();
TF_ASSERT_OK_AND_ASSIGN(auto key, BuildSampleSignature("foo"));
auto compilation_result = std::make_unique<XlaCompiler::CompilationResult>();
auto executable = std::make_unique<FakeExecutable>("foo_exe");
cache->Store(key, DeviceCompileState::kCompiled, absl::OkStatus(),
std::move(compilation_result), std::move(executable));
auto cache_value = cache->LookupOrCreate(key);
EXPECT_EQ(cache_value.compile_state, DeviceCompileState::kCompiled);
EXPECT_EQ(cache_value.request_count, 1);
EXPECT_TRUE(cache_value.compilation_status.ok());
EXPECT_TRUE(cache_value.compilation_result != nullptr);
EXPECT_TRUE(cache_value.executable != nullptr);
EXPECT_EQ(cache_value.executable->data, "foo_exe");
}
TEST(DeviceCompilationCacheTest, StoreOptionalArgs) {
auto cache = std::make_unique<Cache>();
TF_ASSERT_OK_AND_ASSIGN(auto key, BuildSampleSignature("foo"));
auto compilation_result = std::make_unique<XlaCompiler::CompilationResult>();
auto executable = std::make_unique<FakeExecutable>("foo_exe");
cache->Store(key, DeviceCompileState::kCompiled, std::nullopt, std::nullopt,
std::nullopt);
auto cache_value = cache->Lookup(key);
EXPECT_EQ(cache_value->compile_state, DeviceCompileState::kCompiled);
EXPECT_TRUE(cache_value->compilation_status.ok());
EXPECT_TRUE(cache_value->compilation_result == nullptr);
EXPECT_TRUE(cache_value->executable == nullptr);
cache->Store(key, std::nullopt,
absl::InvalidArgumentError("Couldn't compile."), std::nullopt,
std::nullopt);
cache_value = cache->Lookup(key);
EXPECT_EQ(cache_value->compile_state, DeviceCompileState::kCompiled);
EXPECT_EQ(cache_value->compilation_status.code(), error::INVALID_ARGUMENT);
EXPECT_TRUE(cache_value->compilation_result == nullptr);
EXPECT_TRUE(cache_value->executable == nullptr);
cache->Store(key, std::nullopt, std::nullopt, std::move(compilation_result),
std::nullopt);
cache_value = cache->Lookup(key);
EXPECT_EQ(cache_value->compile_state, DeviceCompileState::kCompiled);
EXPECT_EQ(cache_value->compilation_status.code(), error::INVALID_ARGUMENT);
EXPECT_TRUE(cache_value->compilation_result != nullptr);
EXPECT_TRUE(cache_value->executable == nullptr);
cache->Store(key, std::nullopt, std::nullopt, std::nullopt,
std::move(executable));
cache_value = cache->Lookup(key);
EXPECT_EQ(cache_value->compile_state, DeviceCompileState::kCompiled);
EXPECT_EQ(cache_value->compilation_status.code(), error::INVALID_ARGUMENT);
EXPECT_TRUE(cache_value->compilation_result != nullptr);
EXPECT_TRUE(cache_value->executable != nullptr);
EXPECT_EQ(cache_value->executable->data, "foo_exe");
}
TEST(DeviceCompilationCacheTest, StoreMultipleEntries) {
auto cache = std::make_unique<Cache>();
TF_ASSERT_OK_AND_ASSIGN(auto key1, BuildSampleSignature("foo"));
TF_ASSERT_OK_AND_ASSIGN(auto key2, BuildSampleSignature("bar"));
auto compilation_result1 = std::make_unique<XlaCompiler::CompilationResult>();
auto compilation_result2 = std::make_unique<XlaCompiler::CompilationResult>();
auto executable1 = std::make_unique<FakeExecutable>("foo_exe");
auto executable2 = std::make_unique<FakeExecutable>("bar_exe");
cache->Store(key1, DeviceCompileState::kCompiled,
absl::InvalidArgumentError("Invalid argument."),
std::move(compilation_result1), std::move(executable1));
cache->Store(key2, DeviceCompileState::kCompiling, absl::OkStatus(),
std::move(compilation_result2), std::move(executable2));
auto cache_value_1 = cache->Lookup(key1);
auto cache_value_2 = cache->Lookup(key2);
EXPECT_EQ(cache_value_1->compile_state, DeviceCompileState::kCompiled);
EXPECT_EQ(cache_value_1->compilation_status.code(), error::INVALID_ARGUMENT);
EXPECT_TRUE(cache_value_1->compilation_result != nullptr);
EXPECT_TRUE(cache_value_1->executable != nullptr);
EXPECT_EQ(cache_value_1->executable->data, "foo_exe");
EXPECT_EQ(cache_value_2->compile_state, DeviceCompileState::kCompiling);
EXPECT_TRUE(cache_value_2->compilation_status.ok());
EXPECT_TRUE(cache_value_2->compilation_result != nullptr);
EXPECT_TRUE(cache_value_2->executable != nullptr);
EXPECT_EQ(cache_value_2->executable->data, "bar_exe");
}
TEST(DeviceCompilationCacheTest, Finalize) {
auto cache = std::make_unique<Cache>();
TF_ASSERT_OK_AND_ASSIGN(auto key, BuildSampleSignature("foo"));
auto computation = std::make_shared<xla::XlaComputation>();
auto compilation_result = std::make_unique<XlaCompiler::CompilationResult>();
compilation_result->computation = computation;
cache->Store(key, /*compile_state=*/DeviceCompileState::kCompiled,
/*compilation_status=*/absl::OkStatus(),
std::move(compilation_result),
/*executable=*/std::make_unique<FakeExecutable>("foo_exe"));
std::optional<Cache::Value> cache_value = cache->Lookup(key);
ASSERT_TRUE(cache_value.has_value());
ASSERT_TRUE(cache_value->compilation_result != nullptr);
EXPECT_EQ(cache_value->compilation_result->computation, computation);
cache->Finalize();
cache_value = cache->Lookup(key);
ASSERT_TRUE(cache_value.has_value());
ASSERT_TRUE(cache_value->compilation_result != nullptr);
EXPECT_TRUE(cache_value->compilation_result->computation == nullptr);
}
TEST(DeviceCompilationCache, FinalizeWithNullComputation) {
auto cache = std::make_unique<Cache>();
TF_ASSERT_OK_AND_ASSIGN(auto key, BuildSampleSignature("foo"));
cache->Store(
key, /*compile_state=*/DeviceCompileState::kCompiled,
/*compilation_status=*/absl::OkStatus(),
/*compilation_result=*/std::make_unique<XlaCompiler::CompilationResult>(),
/*executable=*/std::make_unique<FakeExecutable>("foo_exe"));
std::optional<Cache::Value> cache_value = cache->Lookup(key);
ASSERT_TRUE(cache_value.has_value());
ASSERT_TRUE(cache_value->compilation_result != nullptr);
EXPECT_TRUE(cache_value->compilation_result->computation == nullptr);
cache->Finalize();
cache_value = cache->Lookup(key);
ASSERT_TRUE(cache_value.has_value());
ASSERT_TRUE(cache_value->compilation_result != nullptr);
EXPECT_TRUE(cache_value->compilation_result->computation == nullptr);
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/device_compilation_cluster_signature.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/compiler/jit/device_compilation_cluster_signature.h"
#include <string>
#include <utility>
#include <variant>
#include "absl/status/statusor.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/node_def_util.h"
namespace tensorflow {
namespace {
using Signature = DeviceCompilationClusterSignature;
using TensorTypeAndShape = Signature::TensorTypeAndShape;
// Functor that converts a Signature's arg to a human readable string.
struct SignatureHumanStringAppender {
explicit SignatureHumanStringAppender(std::string* dest) : dest(dest) {}
std::string* dest;
void operator()(const Tensor& arg) {
absl::StrAppend(dest, "; ", arg.DebugString());
}
void operator()(const TensorTypeAndShape& arg) {
absl::StrAppend(dest, ",", DataTypeString(arg.first));
absl::StrAppend(dest, " [", absl::StrJoin(arg.second, ","), "]");
}
};
// Functor that compares the arg values of two different signatures. Returns
// true when the args are not equal.
struct SignatureNotEqual {
bool operator()(const Tensor& arg, const Tensor& other) {
return arg.dtype() != other.dtype() || arg.shape() != other.shape() ||
arg.tensor_data() != other.tensor_data();
}
bool operator()(const TensorTypeAndShape& arg,
const TensorTypeAndShape& other) {
return arg.first != other.first || arg.second != other.second;
}
bool operator()(const Tensor& arg, const TensorTypeAndShape& other) {
return true;
}
bool operator()(const TensorTypeAndShape& arg, const Tensor& other) {
return true;
}
};
// Functor that incrementally computes a Signature's hash given its current hash
// and one of its args.
struct SignatureHashCombiner {
explicit SignatureHashCombiner(const uint64_t h) : h(h) {}
uint64_t h;
uint64_t operator()(const Tensor& arg) {
h = Hash64Combine(h, std::hash<int>()(static_cast<int>(arg.dtype())));
h = Hash64Combine(
h, Hash64(arg.tensor_data().data(), arg.tensor_data().size()));
for (int dim = 0; dim < arg.dims(); ++dim) {
h = Hash64Combine(h, std::hash<int>()(arg.dim_size(dim)));
}
return h;
}
uint64_t operator()(const TensorTypeAndShape& arg) {
h = Hash64Combine(h, std::hash<int>()(static_cast<int>(arg.first)));
h = Hash64Combine(h, std::hash<int>()(arg.second.size()));
for (int dim : arg.second) {
h = Hash64Combine(h, std::hash<int>()(dim));
}
return h;
}
};
} // namespace
// Compute a string signature which encodes the shapes of the
// arguments in the supplied list.
std::string Signature::HumanString() const {
std::string result = name;
for (const auto& arg : args) {
std::visit(SignatureHumanStringAppender(&result), arg);
}
return result;
}
bool Signature::operator==(const Signature& other) const {
if (name != other.name) return false;
if (args.size() != other.args.size()) return false;
for (int i = 0, end = args.size(); i < end; ++i) {
if (std::visit(SignatureNotEqual(), args[i], other.args[i])) {
return false;
}
}
return true;
}
uint64_t Signature::Hash::operator()(const Signature& signature) const {
uint64_t h = std::hash<std::string>()(signature.name);
for (const auto& arg : signature.args) {
h = std::visit(SignatureHashCombiner(h), arg);
}
return h;
}
static absl::StatusOr<Signature> AppendArguments(
Signature signature, absl::Span<const XlaCompiler::Argument> args) {
signature.args.reserve(args.size());
for (const XlaCompiler::Argument& arg : args) {
switch (arg.kind) {
case XlaCompiler::Argument::kConstant:
case XlaCompiler::Argument::kConstantResource:
signature.args.push_back(arg.constant_value);
break;
case XlaCompiler::Argument::kParameter:
case XlaCompiler::Argument::kResource:
signature.args.push_back(
TensorTypeAndShape(arg.type, arg.DimensionSizesAsInlinedVector()));
break;
default:
return absl::InvalidArgumentError(
absl::StrCat("Unhandled argument kind in XlaCompilationCache: ",
arg.HumanString()));
}
}
return signature;
}
absl::StatusOr<Signature> Signature::Build(
const NameAttrList& function,
absl::Span<const XlaCompiler::Argument> args) {
Signature signature;
signature.name = Canonicalize(function.name(), AttrSlice(&function.attr()));
return AppendArguments(std::move(signature), args);
}
absl::StatusOr<DeviceCompilationClusterSignature> Signature::Build(
const DeviceCompilationCanonicalFunction& canonical_function,
absl::Span<const XlaCompiler::Argument> args) {
Signature signature;
signature.name = canonical_function.canonical;
return AppendArguments(std::move(signature), args);
}
} // namespace tensorflow
tensorflow/compiler/jit/device_compilation_cluster_signature.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_COMPILER_JIT_DEVICE_COMPILATION_CLUSTER_SIGNATURE_H_
#define TENSORFLOW_COMPILER_JIT_DEVICE_COMPILATION_CLUSTER_SIGNATURE_H_
#include <cstdint>
#include <string>
#include <utility>
#include <variant>
#include "absl/container/inlined_vector.h"
#include "absl/status/statusor.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
// Pre-computed device compilation canonical function.
struct DeviceCompilationCanonicalFunction {
std::string canonical;
};
inline DeviceCompilationCanonicalFunction Canonicalize(
const NameAttrList& function) {
return {Canonicalize(function.name(), AttrSlice(&function.attr()))};
}
// Describes the types, shapes and any compile-time constant arguments
// to a kernel. Key that uniquely identifies a compilation output.
struct DeviceCompilationClusterSignature {
// Name of the cluster, built from the function name and it's attributes.
std::string name;
// List of args (either as a TensorTypeAndShape or as a Tensor value)
// for compile-time constant arguments to the compilation, ordered by
// argument number. Tensors must be in host memory.
using TensorTypeAndShape =
std::pair<DataType, absl::InlinedVector<int64_t, 4>>;
absl::InlinedVector<std::variant<Tensor, TensorTypeAndShape>, 8> args;
bool operator==(const DeviceCompilationClusterSignature& other) const;
struct Hash {
uint64_t operator()(
const DeviceCompilationClusterSignature& signature) const;
};
// Returns a human-readable description of the signature.
std::string HumanString() const;
// Builds the signature for a compilation.
static absl::StatusOr<DeviceCompilationClusterSignature> Build(
const NameAttrList& function,
absl::Span<const XlaCompiler::Argument> args);
static absl::StatusOr<DeviceCompilationClusterSignature> Build(
const DeviceCompilationCanonicalFunction& canonical_function,
absl::Span<const XlaCompiler::Argument> args);
};
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_DEVICE_COMPILATION_CLUSTER_SIGNATURE_H_
tensorflow/compiler/jit/device_compilation_cluster_signature_test.cc
+123
View File
@@ -0,0 +1,123 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/jit/device_compilation_cluster_signature.h"
#include <utility>
#include <vector>
#include "tensorflow/compiler/jit/flags.h"
#include "tensorflow/compiler/tf2xla/shape_util.h"
#include "xla/client/client_library.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
namespace tensorflow {
namespace {
using SignatureHash = DeviceCompilationClusterSignature::Hash;
TEST(DeviceCompilationClusterSignatureTest, SignatureEquality) {
NameAttrList fn;
fn.set_name("afunction");
std::vector<XlaCompiler::Argument> args(1);
args[0].kind = XlaCompiler::Argument::kConstant;
args[0].type = DT_INT32;
args[0].shape = TensorShape({4, 0});
args[0].constant_value = Tensor(DT_INT32, {4, 0});
TF_ASSERT_OK_AND_ASSIGN(DeviceCompilationClusterSignature s1,
DeviceCompilationClusterSignature::Build(fn, args));
args[0].type = DT_FLOAT;
args[0].constant_value = Tensor(DT_FLOAT, {4, 0});
TF_ASSERT_OK_AND_ASSIGN(DeviceCompilationClusterSignature s2,
DeviceCompilationClusterSignature::Build(fn, args));
args[0].shape = TensorShape({0, 4});
args[0].constant_value = Tensor(DT_FLOAT, {0, 4});
TF_ASSERT_OK_AND_ASSIGN(DeviceCompilationClusterSignature s3,
DeviceCompilationClusterSignature::Build(fn, args));
std::vector<DeviceCompilationClusterSignature> signatures = {s1, s2, s3};
for (int i = 0; i < signatures.size(); ++i) {
for (int j = 0; j < signatures.size(); ++j) {
EXPECT_EQ(i == j, signatures[i] == signatures[j])
<< "s1: " << signatures[i].HumanString() << "\n"
<< "s2: " << signatures[j].HumanString();
EXPECT_EQ(i == j,
signatures[i].HumanString() == signatures[j].HumanString())
<< "s1: " << signatures[i].HumanString() << "\n"
<< "s2: " << signatures[j].HumanString();
EXPECT_EQ(i == j, SignatureHash()(signatures[i]) ==
SignatureHash()(signatures[j]))
<< "s1: " << signatures[i].HumanString() << "\n"
<< "s1_hash: " << SignatureHash()(signatures[i]) << "\n"
<< "s2: " << signatures[j].HumanString() << "\n"
<< "s2_hash: " << SignatureHash()(signatures[j]);
}
}
}
TEST(DeviceCompilationClusterSignatureTest, SignatureUniqueness) {
NameAttrList fn;
fn.set_name("afunction");
std::vector<XlaCompiler::Argument> args(2);
args[0].kind = XlaCompiler::Argument::kConstant;
args[0].type = DT_INT32;
args[0].constant_value = Tensor(DT_INT32, {4, 0});
args[1].kind = XlaCompiler::Argument::kParameter;
args[1].type = DT_INT32;
args[1].shape = TensorShape({4, 0});
TF_ASSERT_OK_AND_ASSIGN(DeviceCompilationClusterSignature s1,
DeviceCompilationClusterSignature::Build(fn, args));
using std::swap; // go/using-std-swap
swap(args[0], args[1]);
TF_ASSERT_OK_AND_ASSIGN(DeviceCompilationClusterSignature s2,
DeviceCompilationClusterSignature::Build(fn, args));
EXPECT_NE(s1.HumanString(), s2.HumanString());
EXPECT_NE(SignatureHash()(s1), SignatureHash()(s2));
EXPECT_FALSE(s1 == s2);
}
void BM_BuildSignature(::testing::benchmark::State& state) {
const int n_args = state.range(0);
NameAttrList fn;
fn.set_name("afunction");
for (int i = 0; i < n_args; i++) {
(*fn.mutable_attr())[absl::StrCat("T", i)].set_type(DT_FLOAT);
}
std::vector<XlaCompiler::Argument> args(n_args);
for (int i = 0; i < n_args; i++) {
args[i].kind = (((i % 3) == 0) ? XlaCompiler::Argument::kConstant
: XlaCompiler::Argument::kParameter);
args[i].type = DT_INT32;
args[i].shape = TensorShape({4, 0});
args[i].constant_value = Tensor(DT_INT32, {4, 0});
}
for (auto i : state) {
auto s = DeviceCompilationClusterSignature::Build(fn, args);
CHECK(s.ok());
DeviceCompilationClusterSignature sig = std::move(s.value());
}
}
BENCHMARK(BM_BuildSignature)->Arg(0)->Arg(1)->Arg(2)->Arg(5)->Arg(10);
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/device_compilation_profiler.cc
+233
View File
@@ -0,0 +1,233 @@
/* 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/compiler/jit/device_compilation_profiler.h"
#include <cstdint>
#include <optional>
#include <string>
#include <utility>
#include "absl/strings/str_cat.h"
#include "tensorflow/compiler/jit/xla_activity.pb.h"
#include "tensorflow/compiler/jit/xla_activity_listener.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/metrics.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tsl/platform/mutex.h"
namespace tensorflow {
namespace {
bool ShouldBeMegamorphic(int64_t compile_count, int64_t execution_count) {
const int64_t kCompileThreshold = 10;
const int64_t kMinExecutionsPerCompile = 50;
// This heuristic is trying to capture the following property: have we sunk a
// certain minimum amount of compile time into the cluster that didn't quite
// "pay off"?
return compile_count > kCompileThreshold &&
execution_count < kMinExecutionsPerCompile * compile_count;
}
void RegisterExecutionForCluster(
const NameAttrList& function,
DeviceCompilationProfiler::ClusterCompileStats* stats) {
++stats->execution_count;
// The is_megamorphic bit is "sticky". We assume clusters that have been
// observed to be megamorphic once stay megamorphic forever.
if (!stats->is_megamorphic &&
ShouldBeMegamorphic(stats->compile_count, stats->execution_count)) {
VLOG(1) << "Marking " << function.name()
<< " as megamorphic, compile_count=" << stats->compile_count
<< " execution_count=" << stats->execution_count;
stats->is_megamorphic = true;
}
}
// The number of times a lazy compilation must be requested for a specific
// signature before we attempt to compile it.
constexpr int64_t kDefaultCompilationThreshold = 2;
// Maximum number of ongoing compilations.
constexpr int64_t kMaxNumOngoingCompilations = kNumAsyncDeviceCompilerThreads;
} // namespace
DeviceCompilationProfiler::~DeviceCompilationProfiler() {
mutex_lock lock(mu_);
cluster_compile_stats_.clear();
}
absl::StatusOr<DeviceCompilationProfiler::ClusterCompileStats>
DeviceCompilationProfiler::GetCompileStats(const NameAttrList& function) const {
mutex_lock lock(mu_);
if (auto it = cluster_compile_stats_.find(function.name());
it != cluster_compile_stats_.end()) {
return it->second;
}
return absl::NotFoundError(absl::StrCat(
"Couldn't find compilation stats for cluster: ", function.name()));
}
void DeviceCompilationProfiler::RegisterExecution(
const NameAttrList& function) {
mutex_lock lock(mu_);
auto it =
cluster_compile_stats_.emplace(function.name(), ClusterCompileStats{})
.first;
RegisterExecutionForCluster(function, &it->second);
}
absl::Status DeviceCompilationProfiler::RegisterCompilation(
const NameAttrList& function, int64_t compile_time_us,
bool used_persistent_cache) {
metrics::UpdateXlaCompilationTime(compile_time_us);
const std::string& function_name = function.name();
mutex_lock lock(mu_);
// Create a stats entry if it doesn't already exist.
auto it =
cluster_compile_stats_.emplace(function.name(), ClusterCompileStats{})
.first;
const uint64_t compile_time_s = compile_time_us / 1.0e6;
it->second.compile_count++;
it->second.cumulative_compile_time_us += compile_time_us;
VLOG(1) << "Compiled " << function_name << " " << it->second.compile_count
<< " times, compile time: " << compile_time_us
<< " us, cumulative: " << it->second.cumulative_compile_time_us
<< " us ("
<< tensorflow::strings::HumanReadableElapsedTime(compile_time_s)
<< " / "
<< tensorflow::strings::HumanReadableElapsedTime(
it->second.cumulative_compile_time_us / 1.0e6)
<< ")";
XlaJitCompilationActivity jit_compilation_activity;
jit_compilation_activity.set_cluster_name(function_name);
jit_compilation_activity.set_compile_count(it->second.compile_count);
jit_compilation_activity.set_compile_time_us(compile_time_us);
jit_compilation_activity.set_cumulative_compile_time_us(
it->second.cumulative_compile_time_us);
jit_compilation_activity.set_used_persistent_cache(used_persistent_cache);
return BroadcastXlaActivity(std::move(jit_compilation_activity));
}
bool DeviceCompilationProfiler::ShouldCompileCluster(
const NameAttrList& function, DeviceCompileMode compile_mode,
int64_t current_request_count) {
std::optional<int64_t> compile_threshold;
if (compile_mode == DeviceCompileMode::kLazy) {
compile_threshold = kDefaultCompilationThreshold;
} else if (compile_mode == DeviceCompileMode::kAsync) {
compile_threshold = 0; // for now, always compile right away.
}
if (compile_mode == DeviceCompileMode::kStrict) {
// Lazy compilation is disabled.
return true;
}
mutex_lock lock(mu_);
// Create a stats entry if one isn't found and register an execution.
// Determine eligibility assuming this is the first execution of the cluster
// and this cluster has never been compiled before.
auto [it, cluster_not_found] =
cluster_compile_stats_.emplace(function.name(), ClusterCompileStats{});
if (cluster_not_found) {
RegisterExecutionForCluster(function, &it->second);
}
// We avoid compiling clusters that have "gone megamorphic" i.e. have an
// excessive amount of shape dynamism.
if (it->second.is_megamorphic) {
BroadcastOptimizationRemark(XlaOptimizationRemark::MEGAMORPHIC_FUNCTION,
function.name())
.IgnoreError();
VLOG(2) << "Not compiling cluster " << function.name()
<< " because it is megamorphic.";
return false;
}
// TODO(b/255826209): Figure out if Lazy compilation is still needed given
// that we always compile a cluster the first time it is executed (explained
// below) regardless of compilation mode. If it is not, clean up the related
// logic.
// We always compile a cluster the very first time it is executed. This is an
// optimistic guess that pays off for statically shaped TensorFlow graphs
// (since they get the benefit of XLA right away without waiting for warmup)
// and doesn't hurt much for dynamically shaped TensorFlow graphs (we "pay" at
// most one cluster-compilation's worth of compile time).
if (it->second.execution_count == 1) {
return true;
}
if (compile_mode == DeviceCompileMode::kAsync) {
// Asynchronous compilation is enabled.
if (num_ongoing_compilations_ >= kMaxNumOngoingCompilations) {
VLOG(2) << "Not asynchronously compiling cluster " << function.name()
<< " because of too many ongoing compilations.";
return false;
}
}
bool reached_compile_threshold = current_request_count >= *compile_threshold;
if (!reached_compile_threshold) {
VLOG(2) << "Not compiling cluster " << function.name()
<< " because it has not reached compile threshold; threshold is "
<< *compile_threshold << " execution count "
<< current_request_count << ".";
}
return reached_compile_threshold;
}
void DeviceCompilationProfiler::IncrementOngoingAsyncCompilations() {
mutex_lock lock(mu_);
num_ongoing_compilations_++;
}
void DeviceCompilationProfiler::DecrementOngoingAsyncCompilations() {
mutex_lock lock(mu_);
num_ongoing_compilations_--;
}
int64_t DeviceCompilationProfiler::GetNumOngoingAsyncCompilations() const {
mutex_lock lock(mu_);
return num_ongoing_compilations_;
}
std::string DeviceCompilationProfiler::DebugString() const {
std::string debug_string =
"DeviceCompilationProfiler {\ncluster_compile_stats: {\n";
{
mutex_lock lock(mu_);
for (const auto& [key, stats] : cluster_compile_stats_) {
absl::StrAppend(&debug_string, key, ": ", stats.DebugString(), "\n");
}
}
absl::StrAppend(&debug_string, "}\nnum_ongoing_compilations=",
GetNumOngoingAsyncCompilations(), "\n}\n");
return debug_string;
}
} // namespace tensorflow
tensorflow/compiler/jit/device_compilation_profiler.h
+101
View File
@@ -0,0 +1,101 @@
/* 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_COMPILER_JIT_DEVICE_COMPILATION_PROFILER_H_
#define TENSORFLOW_COMPILER_JIT_DEVICE_COMPILATION_PROFILER_H_
#include <cstdint>
#include <string>
#include "tensorflow/compiler/jit/xla_compile_util.h"
#include "tensorflow/core/framework/attr_value.pb.h"
namespace tensorflow {
// Tracks statistics for device compilation and uses these to determine whether
// the given cluster should be compiled or not.
class DeviceCompilationProfiler : public ResourceBase {
public:
DeviceCompilationProfiler() = default;
~DeviceCompilationProfiler() override;
struct ClusterCompileStats {
// Number of times the cluster has been (re-)compiled.
int64_t compile_count = 0;
// The number of times this cluster has been executed.
int64_t execution_count = 0;
// Cumulative time spent compiling the cluster.
int64_t cumulative_compile_time_us = 0;
// True if we have decided that this cluster is too dynamic (i.e. its shapes
// change too frequently) to profitably JIT compile. Once a cluster is
// tagged megamorphic, it stays megamorphic forever.
bool is_megamorphic = false;
std::string DebugString() const {
return absl::StrCat(
"DeviceCompilationProfiler::ClusterCompileStats {compile_count=",
compile_count, ", execution_count=", execution_count,
", cumulative_compile_time_us=", cumulative_compile_time_us,
", is_megamorphic=", is_megamorphic, "}");
}
};
// Returns the compilation statistics for the given cluster.
absl::StatusOr<ClusterCompileStats> GetCompileStats(
const NameAttrList& function) const;
// Determines whether the cluster should be compiled. Creates and inserts an
// entry into stats (also calls `RegisterExecution`) for `function` if it
// doesn't already exist.
virtual bool ShouldCompileCluster(const NameAttrList& function,
DeviceCompileMode compile_mode,
int64_t current_request_count);
// Registers a cluster execution. Increments the execution count for the given
// cluster and also determines whether the cluster has gone megamorphic (and
// sets the megamorphic bit accordingly).
void RegisterExecution(const NameAttrList& function);
// Registers a cluster compilation. Increments the compilation count and
// accumulates the compile time for the given cluster. Also broadcasts an
// XlaJitCompilationActivity.
virtual absl::Status RegisterCompilation(const NameAttrList& function,
int64_t compile_time_us,
bool used_persistent_cache);
void IncrementOngoingAsyncCompilations();
void DecrementOngoingAsyncCompilations();
int64_t GetNumOngoingAsyncCompilations() const;
std::string DebugString() const override;
private:
mutable mutex mu_;
// Maps cluster names to compilation statistics for said cluster.
absl::flat_hash_map<std::string, ClusterCompileStats> cluster_compile_stats_
TF_GUARDED_BY(mu_);
int64_t num_ongoing_compilations_ TF_GUARDED_BY(mu_) = 0;
DeviceCompilationProfiler(const DeviceCompilationProfiler&) = delete;
void operator=(const DeviceCompilationProfiler&) = delete;
};
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_DEVICE_COMPILATION_PROFILER_H_
tensorflow/compiler/jit/device_compilation_profiler_test.cc
+243
View File
@@ -0,0 +1,243 @@
/* 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/compiler/jit/device_compilation_profiler.h"
#include <memory>
#include <utility>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "tensorflow/compiler/jit/tests/device_compiler_test_helper.h"
#include "tensorflow/compiler/jit/xla_activity.pb.h"
#include "tensorflow/core/framework/attr_value.pb.h"
namespace tensorflow {
namespace {
TEST(DeviceCompilationProfilerTest, RegisterExecution) {
DeviceCompilationProfiler* profiler = new DeviceCompilationProfiler();
core::ScopedUnref profiler_ref(profiler);
NameAttrList function;
function.set_name("TestFunc");
for (int i = 0; i < 5; ++i) {
profiler->RegisterExecution(function);
}
TF_ASSERT_OK_AND_ASSIGN(auto stats, profiler->GetCompileStats(function));
EXPECT_EQ(stats.execution_count, 5);
}
TEST(DeviceCompilationProfilerTest, RegisterCompilation) {
DeviceCompilationProfiler* profiler = new DeviceCompilationProfiler();
core::ScopedUnref profiler_ref(profiler);
auto listener = std::make_unique<JitCompilationListener>();
auto listener_ptr = listener.get();
RegisterXlaActivityListener(std::move(listener));
NameAttrList function;
function.set_name("TestFunc");
std::vector<XlaJitCompilationActivity> expected_activities;
for (int i = 0; i < 5; ++i) {
EXPECT_TRUE(profiler->RegisterCompilation(function, 4, false).ok());
TF_ASSERT_OK_AND_ASSIGN(auto stats, profiler->GetCompileStats(function));
XlaJitCompilationActivity expected_activity;
expected_activity.set_cluster_name(function.name());
expected_activity.set_compile_count(stats.compile_count);
expected_activity.set_compile_time_us(4);
expected_activity.set_cumulative_compile_time_us(
stats.cumulative_compile_time_us);
expected_activity.set_used_persistent_cache(false);
expected_activities.push_back(expected_activity);
}
TF_ASSERT_OK_AND_ASSIGN(auto stats, profiler->GetCompileStats(function));
EXPECT_EQ(stats.compile_count, 5);
EXPECT_EQ(stats.cumulative_compile_time_us, 5 * 4);
// TODO(b/255826209): Use ::testing::EqualsProto once b/135192747 is fixed.
const auto& actual_activities = listener_ptr->GetListenerHistory();
EXPECT_EQ(actual_activities.size(), expected_activities.size());
for (size_t i = 0; i < actual_activities.size(); ++i) {
EXPECT_EQ(actual_activities[i].SerializeAsString(),
expected_activities[i].SerializeAsString());
}
}
TEST(DeviceCompilationProfilerTest, OngoingAsyncCompilations) {
DeviceCompilationProfiler* profiler = new DeviceCompilationProfiler();
core::ScopedUnref profiler_ref(profiler);
for (int i = 0; i < 5; ++i) {
profiler->IncrementOngoingAsyncCompilations();
}
EXPECT_EQ(profiler->GetNumOngoingAsyncCompilations(), 5);
for (int i = 0; i < 5; ++i) {
profiler->DecrementOngoingAsyncCompilations();
}
EXPECT_EQ(profiler->GetNumOngoingAsyncCompilations(), 0);
for (int i = 0; i < 5; ++i) {
profiler->IncrementOngoingAsyncCompilations();
profiler->DecrementOngoingAsyncCompilations();
}
EXPECT_EQ(profiler->GetNumOngoingAsyncCompilations(), 0);
}
TEST(DeviceCompilationProfilerTest, ShouldCompileClusterNotFound) {
DeviceCompilationProfiler* profiler = new DeviceCompilationProfiler();
core::ScopedUnref profiler_ref(profiler);
NameAttrList function;
function.set_name("TestFunc");
EXPECT_TRUE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kAsync, 0));
EXPECT_TRUE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kLazy, 0));
EXPECT_TRUE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kStrict, 0));
}
TEST(DeviceCompilationProfilerTest, ShouldCompileClusterFirstExecution) {
DeviceCompilationProfiler* profiler = new DeviceCompilationProfiler();
core::ScopedUnref profiler_ref(profiler);
NameAttrList function;
function.set_name("TestFunc");
profiler->RegisterExecution(function);
EXPECT_TRUE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kAsync, 0));
EXPECT_TRUE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kLazy, 0));
}
TEST(DeviceCompilationProfilerTest, ShouldCompileClusterMegamorphic) {
DeviceCompilationProfiler* profiler = new DeviceCompilationProfiler();
core::ScopedUnref profiler_ref(profiler);
NameAttrList function;
function.set_name("TestFunc");
const int64_t kCompileThreshold = 10;
const int64_t kMinExecutionsPerCompile = 50;
// Register compilation enough times (without registering executions enough
// times) so that the function is marked megamorphic.
for (int i = 0; i < kCompileThreshold + 1; ++i) {
EXPECT_TRUE(profiler->RegisterCompilation(function, 1, false).ok());
}
profiler->RegisterExecution(function);
// Shouldn't compile cluster since it has gone megamorphic.
EXPECT_FALSE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kAsync, 0));
EXPECT_FALSE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kLazy, 0));
TF_ASSERT_OK_AND_ASSIGN(auto stats, profiler->GetCompileStats(function));
EXPECT_TRUE(stats.is_megamorphic);
// Always compile for strict compile mode.
EXPECT_TRUE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kStrict, 0));
// Once a cluster has gone megamorphic, it remains megamorphic (even though
// it's being executed more frequently now) and shouldn't be compiled again.
for (int i = 0; i < kCompileThreshold * kMinExecutionsPerCompile + 1; ++i) {
profiler->RegisterExecution(function);
}
EXPECT_FALSE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kAsync, 0));
EXPECT_FALSE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kLazy, 0));
TF_ASSERT_OK_AND_ASSIGN(stats, profiler->GetCompileStats(function));
EXPECT_TRUE(stats.is_megamorphic);
// Always compile for strict compile mode.
EXPECT_TRUE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kStrict, 0));
}
TEST(DeviceCompilationProfilerTest, ShouldCompileClusterAsync) {
DeviceCompilationProfiler* profiler = new DeviceCompilationProfiler();
core::ScopedUnref profiler_ref(profiler);
NameAttrList function;
function.set_name("TestFunc");
const int64_t kMaxNumOngoingCompilations = 10;
for (int i = 0; i < kMaxNumOngoingCompilations; ++i) {
profiler->IncrementOngoingAsyncCompilations();
}
// Should allow compilation since this is the first execution.
profiler->RegisterExecution(function);
EXPECT_TRUE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kAsync, 0));
// Should not allow compilation since this is not the first execution and
// we've already reached the maximum number of ongoing compilations allowed.
profiler->RegisterExecution(function);
EXPECT_FALSE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kAsync, 0));
profiler->DecrementOngoingAsyncCompilations();
// Should allow compilation since we've decremented the number of ongoing
// compilations.
EXPECT_TRUE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kAsync, 0));
}
TEST(DeviceCompilationProfilerTest, ShouldCompileClusterLazy) {
DeviceCompilationProfiler* profiler = new DeviceCompilationProfiler();
core::ScopedUnref profiler_ref(profiler);
NameAttrList function;
function.set_name("TestFunc");
constexpr int64_t kDefaultCompilationThreshold = 2;
// Should allow compilation since this is the first execution.
profiler->RegisterExecution(function);
EXPECT_TRUE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kLazy, 0));
// Shouldn't allow compilation until compilation has been requested at least
// kDefaultCompilationThreshold times.
profiler->RegisterExecution(function);
for (int current_request_count = 0;
current_request_count < kDefaultCompilationThreshold;
++current_request_count) {
EXPECT_FALSE(profiler->ShouldCompileCluster(
function, DeviceCompileMode::kLazy, current_request_count));
}
EXPECT_TRUE(profiler->ShouldCompileCluster(function, DeviceCompileMode::kLazy,
kDefaultCompilationThreshold));
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/device_compiler.h
+573
View File
@@ -0,0 +1,573 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_JIT_DEVICE_COMPILER_H_
#define TENSORFLOW_COMPILER_JIT_DEVICE_COMPILER_H_
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>
#include "absl/algorithm/container.h"
#include "absl/base/call_once.h"
#include "absl/base/nullability.h"
#include "absl/container/flat_hash_map.h"
#include "absl/status/status.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/jit/device_compilation_cache.h"
#include "tensorflow/compiler/jit/device_compilation_cluster_signature.h"
#include "tensorflow/compiler/jit/device_compilation_profiler.h"
#include "tensorflow/compiler/jit/device_compiler_client.h"
#include "tensorflow/compiler/jit/device_executable_persistor.h"
#include "tensorflow/compiler/jit/flags.h"
#include "tensorflow/compiler/jit/tf_graph_to_hlo_compiler.h"
#include "tensorflow/compiler/jit/xla_compile_util.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/framework/metrics.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/resource_base.h"
#include "tensorflow/core/lib/core/threadpool.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/thread_annotations.h"
namespace tensorflow {
// Compiles/lowers a given Tensorflow graph/function/cluster into a compiled XLA
// compilation (HLO) using the XlaCompiler and compiles the resulting
// XlaCompilationResult into an `ExecutableType` (eg. xla::LocalExecutable) by
// calling `ClientType` (eg. xla::LocalClient).
//
// Caches the compiled XlaCompilationResult and Executable using a
// DeviceCompilationCache. Compilation is done only when there's a cache miss.
//
// Uses the DeviceExecutablePersistor class for persistence and tries to load a
// serialized executable from disk upon a request for compilation. If the
// appropriate executable isn't found on disk, compiles the given Tensorflow
// function/graph/cluster into an XlaCompilationResult (HLO) and
// `ExecutableType` and tries saving/persisting the compiled HLO and executable
// to disk.
//
// Since XLA computations must have static shapes, DeviceCompiler generates a
// new XLA computation for each new set of input shapes.
// TODO(b/255826209): De-templatize once we've moved to Device API completely.
template <typename ExecutableType, typename ClientType>
class DeviceCompiler : public ResourceBase {
public:
DeviceCompiler(
std::unique_ptr<DeviceExecutablePersistor<ExecutableType, ClientType>>
persistor,
std::unique_ptr<DeviceCompilerClient<ExecutableType, ClientType>>
compiler_client);
~DeviceCompiler() override;
enum class CompileScope {
kOp,
kFunction,
};
// Compiles a function into a XlaCompiler::CompilationResult that can be used
// to execute an XLA Computation. Compilation results are cached. Compilation
// is skipped if there is a cache hit. `function` is the name of a Tensorflow
// function to compile. `args` is a description of the arguments to the
// computation.
//
// `compile_mode` controls the behavior of the compilation cache on a cache
// miss. If `compile_mode` is `kLazy` then, based on some profitability
// heuristics, the compilation cache may decide not to compile the cluster at
// this time. In this case it returns null into both `out_compilation_result`
// and `out_executable`. If `compile_mode` is `kStrict` then the compilation
// cache always attempts the compilation on a cache miss. If compilation mode
// is 'kAsync' compilation of the cluster happens in the background while the
// fallback path executes.
//
// The result of compilation is written to `*out_compilation_result`, which
// must be non-null. If `out_executable` is non-null, also builds an
// `ExecutableType` and sets `out_executable` to point to it. The
// resulting executable pointer may be null if the computation has no
// non-constant outputs.
absl::Status CompileIfNeeded(
const XlaCompiler::Options& options, const NameAttrList& function,
const std::vector<XlaCompiler::Argument>& args,
const XlaCompiler::CompileOptions& compile_options,
DeviceCompileMode compile_mode, DeviceCompilationProfiler* profiler,
const XlaCompiler::CompilationResult** out_compilation_result,
ExecutableType** out_executable);
// As above, but for a single op.
absl::Status CompileSingleOpIfNeeded(
const XlaCompiler::Options& options,
const std::vector<XlaCompiler::Argument>& args,
const XlaCompiler::CompileOptions& compile_options, OpKernelContext* ctx,
DeviceCompilationProfiler* profiler,
const XlaCompiler::CompilationResult** out_compilation_result,
ExecutableType** out_executable);
// An override that allows the caller to specify the function explicitly.
absl::Status CompileSingleOpIfNeeded(
const XlaCompiler::Options& options, const NameAttrList& function,
const DeviceCompilationCanonicalFunction& canonical_function,
const std::vector<XlaCompiler::Argument>& args,
const XlaCompiler::CompileOptions& compile_options, OpKernelContext* ctx,
DeviceCompilationProfiler* profiler,
const XlaCompiler::CompilationResult** out_compilation_result,
ExecutableType** out_executable);
ClientType* client() const { return compiler_client_->client(); }
const DeviceType& device_type() const { return persistor_->device_type(); }
DeviceCompilationCache<ExecutableType>* cache() { return cache_.get(); }
DeviceExecutablePersistor<ExecutableType, ClientType>* persistor() {
return persistor_.get();
}
DeviceCompilerClient<ExecutableType, ClientType>* compiler_client() {
return compiler_client_.get();
}
std::string DebugString() const override;
private:
// Common implementation of Compile and CompileSingleOp. The `OpKernelContext`
// parameter is always null for the former.
absl::Status CompileImpl(
const XlaCompiler::CompileOptions& compile_options,
const XlaCompiler::Options& options, const NameAttrList& function,
const DeviceCompilationCanonicalFunction& canonical_function,
const std::vector<XlaCompiler::Argument>& args, CompileScope scope,
DeviceCompileMode compile_mode, OpKernelContext* ctx,
DeviceCompilationProfiler* profiler,
const XlaCompiler::CompilationResult** out_compilation_result,
ExecutableType** out_executable);
StatusOr<typename DeviceCompilationCache<ExecutableType>::Value>
CompileStrict(
const DeviceCompilationClusterSignature& sig,
const XlaCompiler::CompileOptions& compile_options,
const XlaCompiler::Options& options,
const std::vector<XlaCompiler::Argument>& args,
const NameAttrList& function,
typename DeviceCompilationCache<ExecutableType>::Value cache_value,
CompileScope scope, OpKernelContext* ctx,
DeviceCompilationProfiler* profiler, mutex* mu)
TF_EXCLUSIVE_LOCKS_REQUIRED(*mu);
absl::Status CompileAsynchronous(
const DeviceCompilationClusterSignature& sig,
const XlaCompiler::CompileOptions& compile_options,
const XlaCompiler::Options& options,
const std::vector<XlaCompiler::Argument>& args,
const NameAttrList& function, CompileScope scope, OpKernelContext* ctx,
DeviceCompilationProfiler* profiler);
// Releases all references held to `std::shared_ptr<xla::XlaComputation>`
// held by the cache.
//
// This is to be called during session finalization, after all compilation
// has completed and computations no longer need to be accessed through the
// cache.
void Finalize() override;
std::unique_ptr<DeviceExecutablePersistor<ExecutableType, ClientType>>
persistor_;
std::unique_ptr<DeviceCompilerClient<ExecutableType, ClientType>>
compiler_client_;
std::unique_ptr<DeviceCompilationCache<ExecutableType>> cache_;
// Pool of threads for asynchronous compilations.
std::unique_ptr<thread::ThreadPool> async_compiler_threads_;
mutex cluster_mutexes_mu_;
absl::flat_hash_map<DeviceCompilationClusterSignature, std::unique_ptr<mutex>,
DeviceCompilationClusterSignature::Hash>
cluster_mutexes_ TF_GUARDED_BY(cluster_mutexes_mu_);
DeviceCompiler(const DeviceCompiler&) = delete;
void operator=(const DeviceCompiler&) = delete;
};
namespace device_compiler_internal {
// Print something that users can search for to definitively ascertain that XLA
// was used for their TF model.
// Prints only once to avoid spamming LOG(INFO).
inline void LogOnceXlaCompiledFirstCluster() {
static absl::once_flag log_once;
absl::call_once(log_once, [] {
LOG(INFO) << "Compiled cluster using XLA! This line is logged at most "
"once for the lifetime of the process.";
});
}
template <typename ExecutableType>
inline absl::Status EligibleToPersist(DeviceCompileState compile_state,
const ExecutableType* executable) {
if (compile_state != DeviceCompileState::kCompiled) {
return absl::FailedPreconditionError(
"Cache entry to serialize is not compiled.");
}
if (executable == nullptr) {
return absl::FailedPreconditionError(
"LocalExecutable not found for cache entry to serialize.");
}
return absl::OkStatus();
}
} // namespace device_compiler_internal
template <typename ExecutableType, typename ClientType>
DeviceCompiler<ExecutableType, ClientType>::DeviceCompiler(
std::unique_ptr<DeviceExecutablePersistor<ExecutableType, ClientType>>
persistor,
std::unique_ptr<DeviceCompilerClient<ExecutableType, ClientType>>
compiler_client)
: persistor_(std::move(persistor)),
compiler_client_(std::move(compiler_client)) {
cache_ = std::make_unique<DeviceCompilationCache<ExecutableType>>();
async_compiler_threads_ = std::make_unique<tensorflow::thread::ThreadPool>(
tensorflow::Env::Default(), "async_compiler_threads",
kNumAsyncDeviceCompilerThreads);
}
template <typename ExecutableType, typename ClientType>
DeviceCompiler<ExecutableType, ClientType>::~DeviceCompiler() {
// Since programs are owned by the cache, ensure any use of our programs have
// completed by waiting for all stream executors to complete.
compiler_client_->WaitForProgramsToFinish();
// Wait for all outstanding compilations to finish.
// Resetting the pointer explicitly in the top level destructor.
// Without this, the pointer would be reset when the AsyncCompilationState
// is destructed, which is dependent on the order of the members in the
// DeviceCompiler class, which is error prone if the order changes.
async_compiler_threads_.reset();
// TODO(b/110813685): Think about the program ownership model. Programs are
// currently owned by the compilation cache which means we must wait for
// program completion in the destructor. There are multiple compilation caches
// around, which complicates things a little. Perhaps having programs be
// shared_ptrs (an invasive change) would make the model easier to reason
// about?
}
template <typename ExecutableType, typename ClientType>
std::string DeviceCompiler<ExecutableType, ClientType>::DebugString() const {
return "DeviceCompiler";
}
template <typename ExecutableType, typename ClientType>
absl::Status DeviceCompiler<ExecutableType, ClientType>::CompileIfNeeded(
const XlaCompiler::Options& options, const NameAttrList& function,
const std::vector<XlaCompiler::Argument>& args,
const XlaCompiler::CompileOptions& compile_options,
DeviceCompileMode compile_mode, DeviceCompilationProfiler* profiler,
const XlaCompiler::CompilationResult** out_compilation_result,
ExecutableType** out_executable) {
return CompileImpl(compile_options, options, function, Canonicalize(function),
args, CompileScope::kFunction, compile_mode,
/*ctx=*/nullptr, profiler, out_compilation_result,
out_executable);
}
inline NameAttrList GetDeviceCompilerFunction(const NodeDef& def) {
NameAttrList function;
function.set_name(def.op());
*function.mutable_attr() = def.attr();
// Remove the "_class" attribute from the attribute set used to create the
// compilation cache key. This attribute is information for the colocator
// and causes false uniqueness between nodes.
function.mutable_attr()->erase("_class");
return function;
}
template <typename ExecutableType, typename ClientType>
absl::Status
DeviceCompiler<ExecutableType, ClientType>::CompileSingleOpIfNeeded(
const XlaCompiler::Options& options,
const std::vector<XlaCompiler::Argument>& args,
const XlaCompiler::CompileOptions& compile_options, OpKernelContext* ctx,
DeviceCompilationProfiler* profiler,
const XlaCompiler::CompilationResult** out_compilation_result,
ExecutableType** out_executable) {
const NodeDef& def = ctx->op_kernel().def();
const NameAttrList function = GetDeviceCompilerFunction(def);
return CompileSingleOpIfNeeded(options, function, Canonicalize(function),
args, compile_options, ctx, profiler,
out_compilation_result, out_executable);
}
template <typename ExecutableType, typename ClientType>
absl::Status
DeviceCompiler<ExecutableType, ClientType>::CompileSingleOpIfNeeded(
const XlaCompiler::Options& options, const NameAttrList& function,
const DeviceCompilationCanonicalFunction& canonical_function,
const std::vector<XlaCompiler::Argument>& args,
const XlaCompiler::CompileOptions& compile_options, OpKernelContext* ctx,
DeviceCompilationProfiler* profiler,
const XlaCompiler::CompilationResult** out_compilation_result,
ExecutableType** out_executable) {
return CompileImpl(compile_options, options, function, canonical_function,
args, CompileScope::kOp, DeviceCompileMode::kStrict, ctx,
profiler, out_compilation_result, out_executable);
}
template <typename ExecutableType, typename ClientType>
StatusOr<typename DeviceCompilationCache<ExecutableType>::Value>
DeviceCompiler<ExecutableType, ClientType>::CompileStrict(
const DeviceCompilationClusterSignature& sig,
const XlaCompiler::CompileOptions& compile_options,
const XlaCompiler::Options& options,
const std::vector<XlaCompiler::Argument>& args,
const NameAttrList& function,
typename DeviceCompilationCache<ExecutableType>::Value cache_value,
CompileScope scope, OpKernelContext* ctx,
DeviceCompilationProfiler* profiler, mutex* mu) {
tensorflow::Env* env = tensorflow::Env::Default();
const uint64_t compile_start_us = env->NowMicros();
TfGraphToHloCompiler compiler(options);
cache_value.compile_state = DeviceCompileState::kCompiled;
std::unique_ptr<ExecutableType> out_executable;
auto out_compilation_result =
std::make_unique<XlaCompiler::CompilationResult>();
if (scope == CompileScope::kOp) {
cache_value.compilation_status = compiler.CompileSingleOp(
compile_options, ctx, args, out_compilation_result.get());
} else {
CHECK(scope == CompileScope::kFunction); // Crash OK
cache_value.compilation_status = compiler.Compile(
compile_options, function, args, out_compilation_result.get());
}
TF_RETURN_IF_ERROR(cache_value.compilation_status);
TF_RET_CHECK(cache_value.executable == nullptr);
TF_RET_CHECK(out_compilation_result->computation != nullptr);
auto loaded_executable = persistor_->TryToLoadExecutable(
DeviceCompilationClusterSignature::Hash()(sig), sig.HumanString(),
options, *out_compilation_result, compiler_client_.get());
if (loaded_executable.has_value()) {
cache_value.compilation_status = loaded_executable->status();
if (loaded_executable->ok()) {
out_executable = *std::move(*loaded_executable);
metrics::UpdatePersistentCacheLoadCount();
}
} else {
auto built_executable =
compiler_client_->BuildExecutable(options, *out_compilation_result);
TF_RETURN_IF_ERROR(built_executable.status());
out_executable = *std::move(built_executable);
TF_RETURN_IF_ERROR(
device_compiler_internal::EligibleToPersist<ExecutableType>(
cache_value.compile_state, out_executable.get()));
TF_RETURN_IF_ERROR(persistor_->TryToPersistExecutable(
DeviceCompilationClusterSignature::Hash()(sig), sig.HumanString(),
options, *out_compilation_result, *out_executable,
compiler_client_.get()));
}
cache_value.compilation_result = out_compilation_result.get();
cache_value.executable = out_executable.get();
cache_->Store(sig, cache_value.compile_state, cache_value.compilation_status,
std::move(out_compilation_result), std::move(out_executable));
// Finalize the cache to release the XlaComputation after it was compiled.
cache_->Finalize();
const uint64_t compile_end_us = env->NowMicros();
const uint64_t compile_time_us = compile_end_us - compile_start_us;
device_compiler_internal::LogOnceXlaCompiledFirstCluster();
TF_RETURN_IF_ERROR(profiler->RegisterCompilation(
function, compile_time_us, loaded_executable.has_value()));
return cache_value;
}
template <typename ExecutableType, typename ClientType>
absl::Status DeviceCompiler<ExecutableType, ClientType>::CompileAsynchronous(
const DeviceCompilationClusterSignature& signature,
const XlaCompiler::CompileOptions& compile_options,
const XlaCompiler::Options& options,
const std::vector<XlaCompiler::Argument>& args,
const NameAttrList& function, CompileScope scope, OpKernelContext* ctx,
DeviceCompilationProfiler* profiler) {
// Explicitly capture all required data by value for async compilation.
// Update compilation state in cache.
cache_->Store(signature, DeviceCompileState::kCompiling, std::nullopt,
std::nullopt, std::nullopt);
profiler->IncrementOngoingAsyncCompilations();
// Don't move the above code into the thread function as it synchronously
// updates the async compilation state!
// When the ThreadPool for the compilation cache is destroyed, it waits for
// compilations to have finished. This means that both 'entry' and 'this' will
// be alive for the duration of the compilation.
// !!Pay attention when additional variables must be captured by this lambda!!
// All values are captured by value. Make sure that all pointer values (like
// entry) do not get freed until the lambda has finished.
const std::string& function_name = function.name();
async_compiler_threads_->Schedule([=] {
VLOG(2) << "Starting asynchronous compilation of cluster " << function_name
<< '.';
// We don't need to lock mu, but do it anyway to satisfy thread safety
// analysis.
mutex mu;
mutex_lock lock(mu);
auto cache_value = typename DeviceCompilationCache<ExecutableType>::Value();
auto s = CompileStrict(signature, compile_options, options, args, function,
cache_value, scope, ctx, profiler, &mu);
VLOG(2) << "Finished asynchronous compililation of cluster "
<< function_name << '.';
profiler->DecrementOngoingAsyncCompilations();
// Update compilation status in cache.
if (!s.ok()) {
cache_->Store(signature, std::nullopt, s.status(), std::nullopt,
std::nullopt);
}
});
return absl::OkStatus();
}
template <typename ExecutableType, typename ClientType>
void DeviceCompiler<ExecutableType, ClientType>::Finalize() {
const mutex_lock lock(cluster_mutexes_mu_);
std::vector<mutex* absl_nonnull> cluster_mutexes;
cluster_mutexes.reserve(cluster_mutexes_.size());
for (auto& [_, mutex] : cluster_mutexes_) {
if (mutex != nullptr) {
cluster_mutexes.push_back(mutex.get());
}
}
// Sort the mutexes before locking to ensure that this happens in a
// deterministic order, consistent between resizes of the `cluster_mutexes_`
// map.
absl::c_sort(cluster_mutexes);
std::vector<mutex_lock> cluster_mutex_locks;
cluster_mutex_locks.reserve(cluster_mutexes.size());
for (mutex* absl_nonnull const mutex : cluster_mutexes) {
cluster_mutex_locks.emplace_back(*mutex);
}
cache_->Finalize();
}
template <typename ExecutableType, typename ClientType>
absl::Status DeviceCompiler<ExecutableType, ClientType>::CompileImpl(
const XlaCompiler::CompileOptions& compile_options,
const XlaCompiler::Options& options, const NameAttrList& function,
const DeviceCompilationCanonicalFunction& canonical_function,
const std::vector<XlaCompiler::Argument>& args, CompileScope scope,
DeviceCompileMode compile_mode, OpKernelContext* ctx,
DeviceCompilationProfiler* profiler,
const XlaCompiler::CompilationResult** out_compilation_result,
ExecutableType** out_executable) {
DCHECK_NE(out_executable, nullptr);
VLOG(2) << "DeviceCompiler::Compile " << DebugString();
if (VLOG_IS_ON(2)) {
VLOG(2) << "num_inputs=" << args.size();
for (int i = 0, end = args.size(); i < end; i++) {
VLOG(3) << i << ": " << args[i].HumanString();
}
}
TF_ASSIGN_OR_RETURN(auto signature, DeviceCompilationClusterSignature::Build(
canonical_function, args));
// The outer lock protects the existence of the mutex in the map.
mutex* cluster_mutex;
{
mutex_lock lock(cluster_mutexes_mu_);
auto it =
cluster_mutexes_.emplace(signature, std::make_unique<mutex>()).first;
cluster_mutex = it->second.get();
}
profiler->RegisterExecution(function);
std::string human_signature;
if (VLOG_IS_ON(2)) {
human_signature = VLOG_IS_ON(3) ? signature.HumanString() : function.name();
VLOG(2) << "DeviceCompilationClusterSignature: " << human_signature;
}
// Acquire the cache entry lock and compile, if necessary.
// TODO(phawkins): this locking will need to be restructured when we implement
// cache eviction.
mutex_lock cluster_compile_lock(*cluster_mutex);
auto cache_value = cache_->LookupOrCreate(signature);
int64_t current_request_count = cache_value.request_count;
VLOG(2) << "Compilation cache entry hit: "
<< static_cast<int>(cache_value.compile_state)
<< " signature: " << human_signature << " with request count "
<< current_request_count;
DeviceCompileState state = cache_value.compile_state;
*out_compilation_result = nullptr;
*out_executable = nullptr;
// Check if the requested entry is uncompiled and return an error if
// compilation is disabled. This will raise an error for kLazy even if we have
// not yet hit the compilation threshold and no compilation happens this
// round. This is to avoid non-determanism of when compilation is disallowed,
// for example by changing the threshold.
if (state == DeviceCompileState::kUncompiled && FailOnXlaCompilation()) {
VLOG(1) << "XLA compilation disabled: " << function.name() << "\n"
<< absl::StrJoin(
args, "\n",
[](std::string* out, const XlaCompiler::Argument& arg) {
absl::StrAppend(out, " arg: ", arg.HumanString());
});
return absl::InternalError("XLA compilation disabled");
}
if (state == DeviceCompileState::kUncompiled) {
XLA_SCOPED_LOGGING_TIMER("Compilation of XLA executable");
if (!profiler->ShouldCompileCluster(function, compile_mode,
current_request_count)) {
VLOG(2) << "Not compiling for signature: " << human_signature;
return absl::OkStatus();
} else if (compile_mode == DeviceCompileMode::kAsync) {
VLOG(2) << "Queueing asynchronous compilation for signature: "
<< human_signature;
TF_RETURN_IF_ERROR(CompileAsynchronous(signature, compile_options,
options, args, function, scope,
ctx, profiler));
return absl::OkStatus();
} else {
VLOG(2) << "Instantly compiling for signature: " << human_signature;
TF_ASSIGN_OR_RETURN(
cache_value,
CompileStrict(signature, compile_options, options, args, function,
cache_value, scope, ctx, profiler, cluster_mutex));
}
} else if (state == DeviceCompileState::kCompiling) {
VLOG(2) << "Ongoing asynchronous compilation for signature: "
<< human_signature;
return absl::OkStatus();
} else if (state == DeviceCompileState::kCompiled) {
VLOG(2) << "Already Compiled for signature: " << human_signature;
}
TF_RETURN_IF_ERROR(cache_value.compilation_status);
*out_compilation_result = cache_value.compilation_result;
*out_executable = cache_value.executable;
return absl::OkStatus();
}
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_DEVICE_COMPILER_H_
tensorflow/compiler/jit/device_compiler_client.cc
+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.
==============================================================================*/
#include "tensorflow/compiler/jit/device_compiler_client.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "tensorflow/core/util/determinism.h"
namespace tensorflow {
xla::ExecutableBuildOptions GetExecutableBuildOptions(
const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& result, int default_device_ordinal) {
xla::ExecutableBuildOptions build_options;
if (result.collective_info) {
build_options.set_num_replicas(result.collective_info->group_size);
}
if (options.device_ordinal != -1) {
build_options.set_device_ordinal(options.device_ordinal);
} else if (default_device_ordinal != -1) {
build_options.set_device_ordinal(default_device_ordinal);
}
build_options.set_result_layout(result.xla_output_shape);
build_options.set_device_allocator(options.device_allocator.get());
build_options.set_alias_passthrough_params(options.alias_passthrough_params);
build_options.mutable_debug_options()->set_xla_detailed_logging(
options.detailed_logging);
if (tensorflow::OpDeterminismRequired()) {
build_options.mutable_debug_options()->set_xla_gpu_deterministic_ops(true);
}
return build_options;
}
} // namespace tensorflow
tensorflow/compiler/jit/device_compiler_client.h
+76
View File
@@ -0,0 +1,76 @@
/* 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_COMPILER_JIT_DEVICE_COMPILER_CLIENT_H_
#define TENSORFLOW_COMPILER_JIT_DEVICE_COMPILER_CLIENT_H_
#include <optional>
#include <string>
#include <variant>
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/client/executable_build_options.h"
namespace tensorflow {
template <typename ExecutableType, typename ClientType>
class DeviceCompilerClient {
public:
DeviceCompilerClient() = default;
virtual ~DeviceCompilerClient() = default;
// Compiles `result` (HLO) to an `ExecutableType` using `ClientType` and
// returns it.
virtual StatusOr<std::unique_ptr<ExecutableType>> BuildExecutable(
const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& result) = 0;
// Serializes an available `executable` to string using `ClientType` and
// returns it.
virtual absl::StatusOr<std::string> SerializeExecutable(
const ExecutableType& executable) = 0;
// Compiles `result` (HLO) to a serializable executable (eg.
// xla::AotCompilationResult) using `ClientType`, serializes it to string and
// returns it.
virtual absl::StatusOr<std::string> BuildSerializedExecutable(
const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& result) = 0;
// Loads `serialized_executable` into an `ExecutableType` using `ClientType`.
virtual StatusOr<std::unique_ptr<ExecutableType>> LoadExecutable(
const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& result,
const std::string& serialized_executable) = 0;
// Waits for the underlying `ClientType` backend's programs to finish
// executing before returning.
virtual void WaitForProgramsToFinish() = 0;
virtual ClientType* client() const = 0;
private:
DeviceCompilerClient(const DeviceCompilerClient&) = delete;
void operator=(const DeviceCompilerClient&) = delete;
};
// Generates the ExecutableBuildOptions for compilation from HLO to
// executable.
xla::ExecutableBuildOptions GetExecutableBuildOptions(
const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& result, int default_device_ordinal);
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_DEVICE_COMPILER_CLIENT_H_
tensorflow/compiler/jit/device_compiler_client_test.cc
+76
View File
@@ -0,0 +1,76 @@
/* 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/compiler/jit/device_compiler_client.h"
#include <gtest/gtest.h>
namespace tensorflow {
namespace {
TEST(GetExecutableOptionTest, Basic) {
XlaCompiler::Options options;
options.device_ordinal = 0;
options.alias_passthrough_params = true;
options.detailed_logging = true;
XlaCompiler::CompilationResult result;
xla::Shape xla_output_shape;
result.xla_output_shape = xla_output_shape;
auto build_option =
GetExecutableBuildOptions(options, result, /*default_device_ordinal=*/-1);
EXPECT_EQ(build_option.device_ordinal(), 0);
EXPECT_EQ(build_option.result_layout()->ToString(),
xla_output_shape.ToString());
EXPECT_EQ(build_option.alias_passthrough_params(), true);
EXPECT_EQ(build_option.debug_options().xla_detailed_logging(), true);
EXPECT_EQ(build_option.debug_options().xla_enable_dumping(), true);
}
TEST(GetExecutableOptionTest, DefaultDeviceOrdinal) {
XlaCompiler::Options options;
XlaCompiler::CompilationResult result;
auto build_option =
GetExecutableBuildOptions(options, result, /*default_device_ordinal=*/0);
EXPECT_EQ(build_option.device_ordinal(), 0);
}
TEST(GetExecutableOptionTest, DeviceOrdinalNotSet) {
XlaCompiler::Options options;
XlaCompiler::CompilationResult result;
auto build_option =
GetExecutableBuildOptions(options, result, /*default_device_ordinal=*/-1);
EXPECT_EQ(build_option.device_ordinal(), -1);
}
TEST(GetExecutableOptionTest, DumpingWithoutDetailedLogging) {
XlaCompiler::Options options;
options.detailed_logging = false;
XlaCompiler::CompilationResult result;
auto build_option =
GetExecutableBuildOptions(options, result, /*default_device_ordinal=*/-1);
EXPECT_FALSE(build_option.debug_options().xla_detailed_logging());
EXPECT_TRUE(build_option.debug_options().xla_enable_dumping());
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/device_compiler_disable_test.cc
+89
View File
@@ -0,0 +1,89 @@
/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <memory>
#include <utility>
#include <vector>
#include "tensorflow/compiler/jit/device_compilation_profiler.h"
#include "tensorflow/compiler/jit/flags.h"
#include "tensorflow/compiler/jit/device_compiler.h"
#include "tensorflow/compiler/jit/xla_device_compiler_client.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/client/client_library.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace {
// This test is kept separate because it disables XLA compilation globally.
TEST(DeviceCompilerTest, TestDisabledXlaCompilation) {
NameAttrList fn;
fn.set_name("afunction");
// Create mock arguments so we see them in the VLOG when compilation fails.
std::vector<XlaCompiler::Argument> args(2);
for (int i = 0; i < 2; ++i) {
args[i].kind = XlaCompiler::Argument::kParameter;
args[i].type = DT_INT32;
args[i].shape = TensorShape({2, i + 1});
args[i].name = absl::StrCat("arg", i);
}
DisableXlaCompilation();
xla::LocalClient* client = xla::ClientLibrary::LocalClientOrDie();
DeviceType device_type = DeviceType(DEVICE_CPU_XLA_JIT);
const XlaCompiler::CompilationResult* compilation_result;
xla::LocalExecutable* executable;
using XlaDeviceExecutablePersistor =
DeviceExecutablePersistor<xla::LocalExecutable, xla::LocalClient>;
auto persistor = std::make_unique<XlaDeviceExecutablePersistor>(
XlaDeviceExecutablePersistor::Config(), device_type);
auto compiler_client = std::make_unique<XlaDeviceCompilerClient>(client);
auto xla_device_compiler =
new DeviceCompiler<xla::LocalExecutable, xla::LocalClient>(
std::move(persistor), std::move(compiler_client));
core::ScopedUnref xla_device_compiler_ref(xla_device_compiler);
auto profiler = new DeviceCompilationProfiler();
core::ScopedUnref profiler_ref(profiler);
// Check that strict compilation is disallowed.
absl::Status status = xla_device_compiler->CompileIfNeeded(
XlaCompiler::Options{}, fn, args, XlaCompiler::CompileOptions{},
DeviceCompileMode::kStrict, profiler, &compilation_result, &executable);
EXPECT_FALSE(status.ok());
EXPECT_TRUE(absl::StrContains(status.message(), "XLA compilation disabled"));
// Check that async compilation is disallowed.
status = xla_device_compiler->CompileIfNeeded(
XlaCompiler::Options{}, fn, args, XlaCompiler::CompileOptions{},
DeviceCompileMode::kAsync, profiler, &compilation_result, &executable);
EXPECT_FALSE(status.ok());
EXPECT_TRUE(absl::StrContains(status.message(), "XLA compilation disabled"));
// Check that lazy compilation is disallowed.
status = xla_device_compiler->CompileIfNeeded(
XlaCompiler::Options{}, fn, args, XlaCompiler::CompileOptions{},
DeviceCompileMode::kLazy, profiler, &compilation_result, &executable);
EXPECT_FALSE(status.ok());
EXPECT_TRUE(absl::StrContains(status.message(), "XLA compilation disabled"));
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/device_compiler_test.cc
+673
View File
@@ -0,0 +1,673 @@
/* 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/compiler/jit/device_compiler.h"
#include <iostream>
#include <iterator>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/algorithm/container.h"
#include "absl/log/check.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/synchronization/notification.h"
#include "absl/time/clock.h"
#include "absl/time/time.h"
#include "tensorflow/cc/framework/scope.h"
#include "tensorflow/cc/ops/function_ops.h"
#include "tensorflow/cc/ops/math_ops.h"
#include "tensorflow/compiler/jit/device_compilation_cluster_signature.h"
#include "tensorflow/compiler/jit/device_compiler_client.h"
#include "tensorflow/compiler/jit/tests/device_compiler_test_helper.h"
#include "tensorflow/compiler/jit/tf_graph_to_hlo_compiler.pb.h"
#include "tensorflow/compiler/jit/xla_compile_util.h"
#include "tensorflow/compiler/jit/xla_device_compiler_client.h"
#include "tensorflow/compiler/tf2xla/xla_argument.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/client/client_library.h"
#include "xla/client/local_client.h"
#include "xla/hlo/builder/xla_computation.h"
#include "xla/stream_executor/platform_manager.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/lib/strings/proto_serialization.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/framework/fake_input.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/graph_to_functiondef.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/resource_base.h"
#include "tensorflow/core/kernels/ops_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/status_matchers.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/test.h"
#include "tsl/platform/path.h"
namespace tensorflow {
namespace {
using ::testing::_;
using ::testing::ElementsAreArray;
using ::testing::EndsWith;
using ::testing::Eq;
using ::testing::EqualsProto;
using ::testing::IsTrue;
using ::testing::NotNull;
using ::testing::Return;
using ::testing::SizeIs;
using ::testing::StrEq;
using ::testing::UnorderedElementsAre;
using ::testing::status::IsOk;
using XlaDeviceCompiler =
DeviceCompiler<xla::LocalExecutable, xla::LocalClient>;
using XlaDeviceExecutablePersistor =
DeviceExecutablePersistor<xla::LocalExecutable, xla::LocalClient>;
using Signature = DeviceCompilationClusterSignature;
xla::LocalClient* GetLocalClient() {
// TODO(b/255826209): Figure out how to run this test with the CPU client as
// well.
auto platform = se::PlatformManager::PlatformWithName("cuda").value();
return xla::ClientLibrary::GetOrCreateLocalClient(platform).value();
}
XlaDeviceCompiler* CreateXlaDeviceCompiler(bool enable_persistence = false) {
auto xla_compiler_client =
std::make_unique<XlaDeviceCompilerClient>(GetLocalClient());
auto xla_persistor = std::make_unique<XlaDeviceExecutablePersistor>(
XlaDeviceExecutablePersistor::Config{
enable_persistence ? testing::TmpDir() : "", false, "xla"},
DeviceType(DEVICE_GPU_XLA_JIT));
return new XlaDeviceCompiler(std::move(xla_persistor),
std::move(xla_compiler_client));
}
absl::StatusOr<std::unique_ptr<Graph>> SampleGraphAddXY() {
std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
Scope scope = Scope::NewRootScope().ExitOnError();
auto a = ops::_Arg(scope.WithOpName("A"), DT_INT32, 0);
auto b = ops::_Arg(scope.WithOpName("B"), DT_INT32, 1);
auto c = ops::Add(scope.WithOpName("C"), a, b);
auto d = ops::_Retval(scope.WithOpName("D"), c, 0);
TF_RETURN_IF_ERROR(scope.ToGraph(graph.get()));
return graph;
}
absl::StatusOr<FunctionDef> SampleFuntionAddXY(const std::string& name) {
TF_ASSIGN_OR_RETURN(auto graph, SampleGraphAddXY());
FunctionDef fdef;
TF_RETURN_IF_ERROR(GraphToFunctionDef(*graph, name, &fdef));
return fdef;
}
std::vector<XlaCompiler::Argument> SampleArgsForAddXY() {
std::vector<XlaCompiler::Argument> args(2);
args[0].kind = XlaCompiler::Argument::kParameter;
args[0].type = DT_INT32;
args[0].shape = TensorShape({2});
args[1].kind = XlaCompiler::Argument::kParameter;
args[1].type = DT_INT32;
args[1].shape = TensorShape({2});
return args;
}
class MockXlaDeviceExecutablePersistor
: public DeviceExecutablePersistor<xla::LocalExecutable, xla::LocalClient> {
public:
MockXlaDeviceExecutablePersistor()
: DeviceExecutablePersistor<xla::LocalExecutable, xla::LocalClient>(
Config{testing::TmpDir(), false, "xla"},
DeviceType(DEVICE_CPU_XLA_JIT)) {}
MOCK_METHOD(absl::Status, TryToPersistExecutable,
(uint64_t, const std::string&, const XlaCompiler::Options&,
const XlaCompiler::CompilationResult&,
const xla::LocalExecutable&,
(DeviceCompilerClient<xla::LocalExecutable, xla::LocalClient>*)),
(const, override));
};
class MockDeviceCompilationProfiler : public DeviceCompilationProfiler {
public:
MOCK_METHOD(bool, ShouldCompileCluster,
(const NameAttrList& function, DeviceCompileMode compile_mode,
int64_t current_request_count),
(override));
MOCK_METHOD(absl::Status, RegisterCompilation,
(const NameAttrList& function, int64_t compile_time_us,
bool used_persistent_cache),
(override));
};
class DeviceCompilerTest : public ::testing::Test {
protected:
void SetUp() override {
flib_def_ = std::make_unique<FunctionLibraryDefinition>(
OpRegistry::Global(), FunctionDefLibrary());
TF_ASSERT_OK_AND_ASSIGN(auto fdef, SampleFuntionAddXY("foo"));
TF_ASSERT_OK(flib_def_->AddFunctionDef(fdef));
profiler_ = new DeviceCompilationProfiler();
profiler_ref_ = std::make_unique<core::ScopedUnref>(profiler_);
mock_profiler_ = new MockDeviceCompilationProfiler();
mock_profiler_ref_ = std::make_unique<core::ScopedUnref>(mock_profiler_);
xla_device_compiler_ = CreateXlaDeviceCompiler();
xla_device_compiler_ref_ =
std::make_unique<core::ScopedUnref>(xla_device_compiler_);
auto listener = std::make_unique<JitCompilationListener>();
listener_ = listener.get();
RegisterXlaActivityListener(std::move(listener));
}
XlaCompiler::Options GetDefaultXlaOptions() {
XlaCompiler::Options options;
options.device_type = DeviceType(DEVICE_GPU_XLA_JIT);
options.client = xla_device_compiler_->client();
options.flib_def = flib_def_.get();
return options;
}
absl::StatusOr<std::unique_ptr<xla::LocalExecutable>>
BuildSampleXlaExecutable() {
TF_ASSIGN_OR_RETURN(auto graph, SampleGraphAddXY());
auto args = SampleArgsForAddXY();
// Compiles the graph.
XlaCompiler compiler(GetDefaultXlaOptions());
XlaCompiler::CompilationResult compilation_result;
TF_RETURN_IF_ERROR(compiler.CompileGraph(XlaCompiler::CompileOptions(),
"graph", std::move(graph), args,
&compilation_result));
return xla_device_compiler_->compiler_client()->BuildExecutable(
GetDefaultXlaOptions(), compilation_result);
}
std::unique_ptr<FunctionLibraryDefinition> flib_def_;
JitCompilationListener* listener_;
DeviceCompilationProfiler* profiler_;
std::unique_ptr<core::ScopedUnref> profiler_ref_;
MockDeviceCompilationProfiler* mock_profiler_;
std::unique_ptr<core::ScopedUnref> mock_profiler_ref_;
XlaDeviceCompiler* xla_device_compiler_;
std::unique_ptr<core::ScopedUnref> xla_device_compiler_ref_;
};
TEST_F(DeviceCompilerTest, CompileStrictSuccess) {
const XlaCompiler::CompilationResult* compilation_result = nullptr;
xla::LocalExecutable* xla_executable = nullptr;
XlaCompiler::Options options = GetDefaultXlaOptions();
NameAttrList fn;
fn.set_name("foo");
TF_EXPECT_OK(xla_device_compiler_->CompileIfNeeded(
options, fn, SampleArgsForAddXY(), XlaCompiler::CompileOptions{},
DeviceCompileMode::kStrict, profiler_, &compilation_result,
&xla_executable));
EXPECT_TRUE(compilation_result != nullptr);
EXPECT_TRUE(xla_executable != nullptr);
}
TEST_F(DeviceCompilerTest, CompileShouldCompileClusterFalse) {
const XlaCompiler::CompilationResult* compilation_result = nullptr;
xla::LocalExecutable* xla_executable = nullptr;
XlaCompiler::Options options = GetDefaultXlaOptions();
NameAttrList fn;
fn.set_name("foo");
// Using a mock here since it's difficult to have ShouldCompileCluster()
// return false.
EXPECT_CALL(*mock_profiler_,
ShouldCompileCluster(_, DeviceCompileMode::kLazy, 1))
.WillOnce(Return(false));
TF_EXPECT_OK(xla_device_compiler_->CompileIfNeeded(
options, fn, SampleArgsForAddXY(), XlaCompiler::CompileOptions{},
DeviceCompileMode::kLazy, mock_profiler_, &compilation_result,
&xla_executable));
EXPECT_TRUE(compilation_result == nullptr);
EXPECT_TRUE(xla_executable == nullptr);
}
TEST_F(DeviceCompilerTest, CompileCacheHit) {
const XlaCompiler::CompilationResult* compilation_result = nullptr;
xla::LocalExecutable* xla_executable = nullptr;
XlaCompiler::Options options = GetDefaultXlaOptions();
NameAttrList fn;
fn.set_name("foo");
TF_EXPECT_OK(xla_device_compiler_->CompileIfNeeded(
options, fn, SampleArgsForAddXY(), XlaCompiler::CompileOptions{},
DeviceCompileMode::kStrict, profiler_, &compilation_result,
&xla_executable));
EXPECT_TRUE(compilation_result != nullptr);
EXPECT_TRUE(xla_executable != nullptr);
const XlaCompiler::CompilationResult* new_compilation_result = nullptr;
xla::LocalExecutable* new_xla_executable = nullptr;
// Request compiling the same function again.
TF_EXPECT_OK(xla_device_compiler_->CompileIfNeeded(
options, fn, SampleArgsForAddXY(), XlaCompiler::CompileOptions{},
DeviceCompileMode::kStrict, profiler_, &new_compilation_result,
&new_xla_executable));
// new_compilation_result and new_xla_executable should point to the
// compilation_result and executable returned after the first compilation
// request.
EXPECT_EQ(compilation_result, new_compilation_result);
EXPECT_EQ(xla_executable, new_xla_executable);
}
TEST_F(DeviceCompilerTest, CompileAsyncSuccess) {
const XlaCompiler::CompilationResult* compilation_result = nullptr;
xla::LocalExecutable* xla_executable = nullptr;
XlaCompiler::Options options = GetDefaultXlaOptions();
NameAttrList fn;
fn.set_name("foo");
// Using a mock here to determine when the async compilation finishes. This is
// to avoid using absl::SleepFor().
// `RegisterCompilation` is the last call that happens just before the async
// compilation completes. We use the completion of this call to determine when
// the compilation finshes to verify expected behavior.
absl::Notification done;
EXPECT_CALL(*mock_profiler_,
ShouldCompileCluster(_, DeviceCompileMode::kAsync, 1))
.WillOnce(Return(true));
EXPECT_CALL(*mock_profiler_, RegisterCompilation(_, _, false))
.WillOnce([&done] {
done.Notify();
return absl::OkStatus();
});
auto args = SampleArgsForAddXY();
TF_EXPECT_OK(xla_device_compiler_->CompileIfNeeded(
options, fn, args, XlaCompiler::CompileOptions{},
DeviceCompileMode::kAsync, mock_profiler_, &compilation_result,
&xla_executable));
// compilation_result and xla_executable aren't available immediately after
// requesting compilation in asynchronous mode.
EXPECT_TRUE(compilation_result == nullptr);
EXPECT_TRUE(xla_executable == nullptr);
// Check if an appropriate entry is made in xla_cache.
auto xla_cache = xla_device_compiler_->cache();
TF_ASSERT_OK_AND_ASSIGN(auto signature, Signature::Build(fn, args));
auto cache_value = xla_cache->Lookup(signature);
EXPECT_TRUE(cache_value);
EXPECT_TRUE(cache_value->compile_state != DeviceCompileState::kUncompiled);
// Wait for async compilation to complete.
done.WaitForNotification();
cache_value = xla_cache->Lookup(signature);
EXPECT_TRUE(cache_value);
EXPECT_TRUE(cache_value->compile_state == DeviceCompileState::kCompiled);
EXPECT_TRUE(cache_value->compilation_result != nullptr);
EXPECT_TRUE(cache_value->executable != nullptr);
EXPECT_TRUE(cache_value->compilation_status.ok());
}
TEST_F(DeviceCompilerTest, CompilePersistentCacheEnabled) {
auto xla_device_compiler =
CreateXlaDeviceCompiler(/*enable_persistence=*/true);
core::ScopedUnref xla_device_compiler_ref(xla_device_compiler);
NameAttrList fn;
fn.set_name("foo");
auto args = SampleArgsForAddXY();
XlaCompiler::Options options = GetDefaultXlaOptions();
const XlaCompiler::CompilationResult* compilation_result = nullptr;
xla::LocalExecutable* xla_executable = nullptr;
TF_EXPECT_OK(xla_device_compiler->CompileIfNeeded(
options, fn, args, XlaCompiler::CompileOptions{},
DeviceCompileMode::kStrict, profiler_, &compilation_result,
&xla_executable));
EXPECT_TRUE(compilation_result != nullptr);
EXPECT_TRUE(xla_executable != nullptr);
// Check if device_compiler was able to load the executable from the
// persistent cache.
std::vector<XlaJitCompilationActivity> activity_history =
listener_->GetListenerHistory();
EXPECT_EQ(activity_history.size(), 1);
EXPECT_EQ(activity_history[0].cluster_name(), fn.name());
EXPECT_EQ(activity_history[0].compile_count(), 1);
EXPECT_FALSE(activity_history[0].used_persistent_cache());
listener_->ClearListenerHistory();
// Create another DeviceCompiler object pointing to the same persistent cache
// directory. It should load the executable instead of building it.
auto xla_device_compiler_2 =
CreateXlaDeviceCompiler(/*enable_persistence=*/true);
core::ScopedUnref xla_device_compiler_ref_2(xla_device_compiler_2);
auto profiler = new DeviceCompilationProfiler();
core::ScopedUnref profiler_ref(profiler);
const XlaCompiler::CompilationResult* compilation_result_2 = nullptr;
xla::LocalExecutable* xla_executable_2 = nullptr;
TF_EXPECT_OK(xla_device_compiler_2->CompileIfNeeded(
options, fn, args, XlaCompiler::CompileOptions{},
DeviceCompileMode::kStrict, profiler, &compilation_result_2,
&xla_executable_2));
EXPECT_TRUE(compilation_result_2 != nullptr);
EXPECT_TRUE(xla_executable_2 != nullptr);
activity_history = listener_->GetListenerHistory();
EXPECT_EQ(activity_history.size(), 1);
EXPECT_EQ(activity_history[0].cluster_name(), fn.name());
EXPECT_EQ(activity_history[0].compile_count(), 1);
// Verify that the executable was loaded instead of built.
EXPECT_TRUE(activity_history[0].used_persistent_cache());
}
TEST_F(DeviceCompilerTest, CompileFailedToLoadFromPersistentCache) {
auto xla_device_compiler =
CreateXlaDeviceCompiler(/*enable_persistence=*/true);
core::ScopedUnref xla_device_compiler_ref(xla_device_compiler);
NameAttrList fn;
fn.set_name("foo");
auto args = SampleArgsForAddXY();
XlaCompiler::Options options = GetDefaultXlaOptions();
const XlaCompiler::CompilationResult* compilation_result = nullptr;
xla::LocalExecutable* xla_executable = nullptr;
// Persist an executable.
TF_EXPECT_OK(xla_device_compiler->CompileIfNeeded(
options, fn, args, XlaCompiler::CompileOptions{},
DeviceCompileMode::kStrict, profiler_, &compilation_result,
&xla_executable));
// Corrupt the file which contains the serialized executable.
std::vector<std::string> files;
TF_ASSERT_OK(Env::Default()->GetChildren(testing::TmpDir(), &files));
std::string const* serialized_executable_filename = nullptr;
for (const auto& file : files) {
if (absl::StartsWith(file, "xla__")) {
serialized_executable_filename = &file;
break;
}
}
EXPECT_TRUE(serialized_executable_filename != nullptr);
std::string serialized_executable_filepath =
io::JoinPath(testing::TmpDir(), *serialized_executable_filename);
std::unique_ptr<WritableFile> serialized_executable_file;
TF_ASSERT_OK(Env::Default()->NewWritableFile(serialized_executable_filepath,
&serialized_executable_file));
TF_ASSERT_OK(serialized_executable_file->Append("Garbage."));
TF_ASSERT_OK(serialized_executable_file->Close());
// Create another DeviceCompiler object pointing to the same persistent cache
// directory. It should error out while loading the executable from the
// corrupt file.
auto xla_device_compiler_2 =
CreateXlaDeviceCompiler(/*enable_persistence=*/true);
core::ScopedUnref xla_device_compiler_ref_2(xla_device_compiler_2);
const XlaCompiler::CompilationResult* compilation_result_2 = nullptr;
xla::LocalExecutable* xla_executable_2 = nullptr;
EXPECT_FALSE(xla_device_compiler_2
->CompileIfNeeded(options, fn, args,
XlaCompiler::CompileOptions{},
DeviceCompileMode::kStrict, profiler_,
&compilation_result_2, &xla_executable_2)
.ok());
EXPECT_TRUE(compilation_result_2 == nullptr);
EXPECT_TRUE(xla_executable_2 == nullptr);
}
TEST_F(DeviceCompilerTest, CompileStrictPersistentCacheFailedToPersist) {
auto xla_compiler_client =
std::make_unique<XlaDeviceCompilerClient>(GetLocalClient());
auto xla_persistor = std::make_unique<MockXlaDeviceExecutablePersistor>();
auto xla_device_compiler = new XlaDeviceCompiler(
std::move(xla_persistor), std::move(xla_compiler_client));
core::ScopedUnref xla_device_compiler_ref(xla_device_compiler);
NameAttrList fn;
fn.set_name("foo");
auto args = SampleArgsForAddXY();
XlaCompiler::Options options = GetDefaultXlaOptions();
const XlaCompiler::CompilationResult* compilation_result = nullptr;
xla::LocalExecutable* xla_executable = nullptr;
auto persistor = absl::down_cast<MockXlaDeviceExecutablePersistor*>(
xla_device_compiler->persistor());
TF_ASSERT_OK_AND_ASSIGN(auto signature, Signature::Build(fn, args));
EXPECT_CALL(*persistor,
TryToPersistExecutable(Signature::Hash()(signature),
signature.HumanString(), _, _, _, _))
.WillOnce(Return(absl::FailedPreconditionError("Random error.")));
EXPECT_THAT(xla_device_compiler->CompileIfNeeded(
options, fn, args, XlaCompiler::CompileOptions{},
DeviceCompileMode::kStrict, profiler_, &compilation_result,
&xla_executable),
absl_testing::StatusIs(error::FAILED_PRECONDITION,
::testing::HasSubstr("Random error.")));
EXPECT_TRUE(compilation_result == nullptr);
EXPECT_TRUE(xla_executable == nullptr);
}
class DeviceCompilerTestWithDump : public DeviceCompilerTest {
protected:
explicit DeviceCompilerTestWithDump() {
dump_dir_ = tsl::io::JoinPath(
testing::TmpDir(),
absl::StrCat("dump_test_", absl::ToUnixNanos(absl::Now())));
CHECK_OK(Env::Default()->RecursivelyCreateDir(dump_dir_));
setenv("TF_GRAPH_TO_HLO_COMPILER_DUMP_DIR", dump_dir_.c_str(), 1);
}
~DeviceCompilerTestWithDump() override {
unsetenv("TF_GRAPH_TO_HLO_COMPILER_DUMP_DIR");
}
std::string dump_dir_;
};
TEST_F(DeviceCompilerTestWithDump, CompileStrictDebugInformationDumpWorks) {
// We create a new XlaDeviceCompiler here so that we ensure that the cached
// result is not used and `CompileStrict` is called.
auto xla_device_compiler =
CreateXlaDeviceCompiler(/*enable_persistence=*/false);
ASSERT_THAT(xla_device_compiler, NotNull());
core::ScopedUnref xla_device_compiler_ref(xla_device_compiler);
// Now we run the compilation. We only care that `CompileStrict` has been
// actually called, and we don't care about whether it succeeded.
const XlaCompiler::CompilationResult* compilation_result = nullptr;
xla::LocalExecutable* xla_executable = nullptr;
XlaCompiler::Options options = GetDefaultXlaOptions();
XlaCompiler::CompileOptions compile_options;
NameAttrList fn;
fn.set_name("foo");
xla_device_compiler
->CompileIfNeeded(options, fn, SampleArgsForAddXY(), compile_options,
DeviceCompileMode::kStrict, profiler_,
&compilation_result, &xla_executable)
.IgnoreError();
// Check the directory structure.
std::vector<std::string> dump_dir_files;
EXPECT_THAT(Env::Default()->GetChildren(dump_dir_, &dump_dir_files), IsOk());
EXPECT_THAT(dump_dir_files, SizeIs(1));
absl::string_view dump_subdir = dump_dir_files[0];
std::vector<std::string> dump_subdir_files;
EXPECT_THAT(
Env::Default()->GetChildren(tsl::io::JoinPath(dump_dir_, dump_subdir),
&dump_subdir_files),
IsOk());
// `options.pb` and one pb for the compile call arguments.
EXPECT_THAT(dump_subdir_files,
UnorderedElementsAre("options.pb", EndsWith(".pb")));
// First check the options file.
std::string options_file =
tsl::io::JoinPath(dump_dir_, dump_subdir, "options.pb");
std::string options_contents;
EXPECT_THAT(ReadFileToString(Env::Default(), options_file, &options_contents),
IsOk());
TfGraphToHloCompilerOptions options_proto;
EXPECT_THAT(options_proto.ParseFromString(options_contents), IsTrue());
EXPECT_THAT(absl::StrCat(tsl::DeterministicProtoHash64(options_proto)),
StrEq(dump_subdir));
EXPECT_THAT(options_proto.device_type(),
StrEq(options.device_type.type_string()));
EXPECT_THAT(options_proto.flib_def(),
EqualsProto(options.flib_def->ToProto()));
EXPECT_THAT(options_proto.graph_def_version(), Eq(options.graph_def_version));
// Now check the compile call arguments file.
absl::string_view compile_call_args_file_name =
(dump_subdir_files[0] == "options.pb" ? dump_subdir_files[1]
: dump_subdir_files[0]);
std::string compile_call_args_file =
tsl::io::JoinPath(dump_dir_, dump_subdir, compile_call_args_file_name);
std::string compile_call_args_contents;
EXPECT_THAT(ReadFileToString(Env::Default(), compile_call_args_file,
&compile_call_args_contents),
IsOk());
TfGraphToHloCompilerCompileCallArgs compile_call_args_proto;
EXPECT_THAT(
compile_call_args_proto.ParseFromString(compile_call_args_contents),
IsTrue());
EXPECT_THAT(compile_call_args_proto.compile_options(),
StrEq(compile_options.DebugString()));
EXPECT_THAT(compile_call_args_proto.function(), EqualsProto(fn));
std::vector<XlaArgument> xla_arguments_from_proto;
absl::c_transform(compile_call_args_proto.xla_args(),
std::back_inserter(xla_arguments_from_proto),
[](const tf2xla::XlaArgumentProto& proto) -> XlaArgument {
auto arg = XlaArgument::FromProto(proto);
CHECK_OK(arg); // Crash OK
return arg.value();
});
EXPECT_THAT(xla_arguments_from_proto, ElementsAreArray(SampleArgsForAddXY()));
}
TEST_F(OpsTestBase, CompileSingleOpSuccess) {
TF_EXPECT_OK(NodeDefBuilder("identity_op", "Identity")
.Input(FakeInput(DT_FLOAT))
.Attr("T", DT_FLOAT)
.Finalize(node_def()));
TF_EXPECT_OK(InitOp());
AddInputFromArray<float>(TensorShape({1, 2}), {6.9, 4.2});
TF_EXPECT_OK(RunOpKernel());
auto xla_device_compiler = CreateXlaDeviceCompiler();
core::ScopedUnref xla_device_compiler_ref(xla_device_compiler);
auto profiler = new DeviceCompilationProfiler();
core::ScopedUnref profiler_ref(profiler);
const XlaCompiler::CompilationResult* compilation_result = nullptr;
xla::LocalExecutable* xla_executable = nullptr;
XlaOpRegistry::RegisterCompilationKernels();
auto flib_def = std::make_unique<FunctionLibraryDefinition>(
OpRegistry::Global(), FunctionDefLibrary());
XlaCompiler::Options options;
options.device_type = DeviceType(DEVICE_GPU_XLA_JIT);
options.client = GetLocalClient();
options.flib_def = flib_def.get();
std::vector<XlaCompiler::Argument> args(1);
args[0].kind = XlaCompiler::Argument::kConstant;
args[0].type = DT_FLOAT;
args[0].shape = TensorShape({1, 2});
args[0].constant_value = GetInput(0);
args[0].initialized = true;
NameAttrList fn;
fn.set_name("foo");
TF_EXPECT_OK(xla_device_compiler->CompileSingleOpIfNeeded(
options, args, XlaCompiler::CompileOptions{}, context_.get(), profiler,
&compilation_result, &xla_executable));
EXPECT_TRUE(compilation_result != nullptr);
EXPECT_TRUE(xla_executable != nullptr);
}
TEST_F(DeviceCompilerTest, Finalize) {
XlaCompiler::Options options = GetDefaultXlaOptions();
NameAttrList fn;
fn.set_name("foo");
const XlaCompiler::CompilationResult* compilation_result = nullptr;
xla::LocalExecutable* xla_executable = nullptr;
TF_EXPECT_OK(xla_device_compiler_->CompileIfNeeded(
options, fn, SampleArgsForAddXY(), XlaCompiler::CompileOptions{},
DeviceCompileMode::kStrict, profiler_, &compilation_result,
&xla_executable));
ASSERT_TRUE(compilation_result != nullptr);
// Cast to `ResourceBase` to verify that the `Finalize` implementation
// overrides the base class's.
static_cast<ResourceBase*>(xla_device_compiler_)->Finalize();
TF_EXPECT_OK(xla_device_compiler_->CompileIfNeeded(
options, fn, SampleArgsForAddXY(), XlaCompiler::CompileOptions{},
DeviceCompileMode::kStrict, profiler_, &compilation_result,
&xla_executable));
ASSERT_TRUE(compilation_result != nullptr);
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/device_context_test.cc
+96
View File
@@ -0,0 +1,96 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <memory>
#include <string_view>
#include <utility>
#include <vector>
#include <gtest/gtest.h>
#include "tensorflow/compiler/jit/flags.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "tensorflow/core/framework/tensor_testutil.h"
namespace tensorflow {
namespace {
static bool Initialized = [] {
auto& rollout_config = GetXlaOpsCommonFlags()->tf_xla_use_device_api;
rollout_config.enabled_for_xla_launch_ = true;
rollout_config.enabled_for_compile_on_demand_ = true;
tensorflow::GetXlaDeviceFlags()->tf_xla_enable_xla_devices = true;
return true;
}();
class DeviceContextTest : public ::testing::Test {
public:
void SetDevice(const std::string& device_type) {
auto& rollout_config = GetXlaOpsCommonFlags()->tf_xla_use_device_api;
rollout_config.AllowForDeviceInXlaLaunch(DeviceType(device_type));
rollout_config.AllowForDeviceInXlaCompileOnDemand(DeviceType(device_type));
auto device_factory = DeviceFactory::GetFactory(device_type);
SessionOptions options;
std::vector<std::unique_ptr<Device>> devices;
absl::Status s = device_factory->CreateDevices(
options, "/job:worker/replica:0/task:0", &devices);
device_ = std::move(devices[0]);
tensorflow::AllocatorAttributes host_alloc_attr;
host_alloc_attr.set_on_host(true);
host_allocator_ = device_->GetAllocator(host_alloc_attr);
tensorflow::AllocatorAttributes device_alloc_attr;
device_alloc_attr.set_on_host(false);
device_allocator_ = device_->GetAllocator(device_alloc_attr);
tensorflow::DeviceContext* device_context;
auto status = device_->TryGetDeviceContext(&device_context);
TF_EXPECT_OK(status);
device_context_.reset(device_context);
}
std::unique_ptr<Device> device_;
tensorflow::core::RefCountPtr<DeviceContext> device_context_;
tensorflow::Allocator* host_allocator_;
tensorflow::Allocator* device_allocator_;
};
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
TEST_F(DeviceContextTest, TestXlaGpuRoundTripTransferWithDeviceApi) {
SetDevice(DEVICE_XLA_GPU);
tensorflow::Tensor origin_cpu_tensor(host_allocator_, tensorflow::DT_FLOAT,
tensorflow::TensorShape({2, 2}));
tensorflow::test::FillValues<float>(&origin_cpu_tensor, {1.2, 2.3, 3.4, 4.5});
tensorflow::Tensor device_tensor(device_allocator_, tensorflow::DT_FLOAT,
tensorflow::TensorShape({2, 2}));
tensorflow::Tensor dest_cpu_tensor(host_allocator_, tensorflow::DT_FLOAT,
tensorflow::TensorShape({2, 2}));
TF_ASSERT_OK(device_context_->CopyCPUTensorToDeviceSync(
&origin_cpu_tensor, device_.get(), &device_tensor));
TF_ASSERT_OK(device_context_->CopyDeviceTensorToCPUSync(
&device_tensor, "", device_.get(), &dest_cpu_tensor));
LOG(INFO) << "H2D - D2H roundtrip completes. tensor: "
<< dest_cpu_tensor.DebugString(4);
tensorflow::test::ExpectClose(origin_cpu_tensor, dest_cpu_tensor);
}
#endif
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/device_executable_persistor.cc
+37
View File
@@ -0,0 +1,37 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/jit/device_executable_persistor.h"
#include <string>
#include "absl/strings/str_cat.h"
namespace tensorflow {
std::string XlaSerializedCacheKeyToFileName(const XlaSerializedCacheKey& key) {
static constexpr char kXlaSerializedCacheKeySeparator[] = "__";
return absl::StrCat(
key.prefix(), key.prefix().empty() ? "" : kXlaSerializedCacheKeySeparator,
key.signature_fingerprint(), kXlaSerializedCacheKeySeparator,
key.cluster_fingerprint(), kXlaSerializedCacheKeySeparator,
key.device_type(),
key.compiled_using_pjrt()
? absl::StrCat(kXlaSerializedCacheKeySeparator, "pjrt")
: "",
".pb");
}
} // namespace tensorflow
tensorflow/compiler/jit/device_executable_persistor.h
+404
View File
@@ -0,0 +1,404 @@
/* 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_COMPILER_JIT_DEVICE_EXECUTABLE_PERSISTOR_H_
#define TENSORFLOW_COMPILER_JIT_DEVICE_EXECUTABLE_PERSISTOR_H_
#include <optional>
#include <string>
#include <utility>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/compiler/jit/device_compiler_client.h"
#include "tensorflow/compiler/jit/xla_compilation_cache.pb.h"
#include "tensorflow/compiler/jit/xla_device_compiler_client.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/pjrt/pjrt_client.h"
#include "xla/service/hlo.pb.h"
#include "xla/util.h"
#include "tensorflow/core/framework/device.h"
#include "tensorflow/core/lib/strings/proto_serialization.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/types.h"
#include "tsl/platform/errors.h"
#include "tsl/platform/statusor.h"
namespace tensorflow {
// Returns the persisted compilation cache file name for the given key.
std::string XlaSerializedCacheKeyToFileName(const XlaSerializedCacheKey& key);
// Offers a way to persist and/or load compiled `ExecutableType`s along with the
// corresponding HLO (`CompilationResult`) to/from `persistent_cache_directory`
// (if one was provided during construction) on disk using `ClientType`.
template <typename ExecutableType, typename ClientType>
class DeviceExecutablePersistor {
public:
// Configuration for setting up persistence (directory, filename prefix, etc).
struct Config {
Config() = default;
explicit Config(absl::string_view persistent_cache_directory,
bool disable_strict_signature_checks,
absl::string_view persistence_prefix,
bool persistent_cache_directory_read_only)
: persistent_cache_directory(persistent_cache_directory),
disable_strict_signature_checks(disable_strict_signature_checks),
persistence_prefix(persistence_prefix),
persistent_cache_directory_read_only(
persistent_cache_directory_read_only) {}
explicit Config(absl::string_view persistent_cache_directory,
bool disable_strict_signature_checks,
absl::string_view persistence_prefix)
: persistent_cache_directory(persistent_cache_directory),
disable_strict_signature_checks(disable_strict_signature_checks),
persistence_prefix(persistence_prefix) {}
// If non-empty, JIT-compiled executables are saved to and loaded from the
// specified file system directory path.
std::string persistent_cache_directory;
// Disable strict signature checks for entries loaded into the cache from
// external sources.
bool disable_strict_signature_checks = false;
// The cache persistence prefix to use if serializing/deserialzing entries.
std::string persistence_prefix;
// Cache is read-only if set to true.
bool persistent_cache_directory_read_only = false;
};
DeviceExecutablePersistor(const Config& config,
const DeviceType& device_type);
virtual ~DeviceExecutablePersistor() = default;
// Returns std::nullopt if persistence is not enabled (i.e.
// `persistent_cache_directory_` is empty) or if the serialized entry is not
// found on disk. Otherwise, loads and returns the serialized executable
// (or returns a status).
// TODO(b/255826209): Take in Signature instead of hash and string once cache
// is refactored.
std::optional<StatusOr<std::unique_ptr<ExecutableType>>> TryToLoadExecutable(
uint64_t signature_hash, const std::string& signature_str,
const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& compilation_result,
DeviceCompilerClient<ExecutableType, ClientType>* client) const;
// Tries to serialize an already built `executable` and persist it on disk. If
// unable to do so, tries to build a serialized executable using the AOT
// pipeline and persists that to disk.
// TODO(b/255826209): Take in Signature instead hash and string once cache
// is refactored.
virtual absl::Status TryToPersistExecutable(
uint64_t signature_hash, const std::string& signature_str,
const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& compilation_result,
const ExecutableType& executable,
DeviceCompilerClient<ExecutableType, ClientType>* client) const;
const DeviceType& device_type() const { return device_type_; }
const std::string& persistence_prefix() const { return persistence_prefix_; }
const std::string& persistent_cache_directory() const {
return persistent_cache_directory_;
}
private:
// Returns a cache key proto that identifies an entry in the compilation
// cache.
XlaSerializedCacheKey BuildSerializedCacheKey(
uint64_t signature_hash, const xla::HloModuleProto& hlo_module) const;
XlaSerializedCacheKey BuildSerializedCacheKey(
uint64_t signature_hash, const xla::HloModuleProto& hlo_module,
bool compiled_using_pjrt) const;
// Serializes the signature and its corresponding entry to a proto message.
absl::StatusOr<XlaSerializedCacheEntry> SerializeEntry(
uint64_t signature_hash, const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& compilation_result,
const ExecutableType& executable,
DeviceCompilerClient<ExecutableType, ClientType>* compiler_client) const;
// Saves the cache entry in the file directory supplied during the
// construction of this class. Overwrites existing entries.
absl::Status SaveSerializedEntry(const XlaSerializedCacheEntry& entry) const;
// Tries to read a cache entry given a `key` by searching the file directory
// supplied during the construction of this class. Returns std::nullopt if no
// cache entry is found.
absl::StatusOr<std::optional<XlaSerializedCacheEntry>>
TryToReadSerializedEntry(const XlaSerializedCacheKey& key) const;
// Checks if the loaded `entry` matches the expected `key` and `hlo_module`.
absl::Status VerifyLoadedCacheEntry(
const XlaSerializedCacheKey& key, const xla::HloModuleProto& hlo_module,
const XlaSerializedCacheEntry& entry) const;
std::string GetFilePath(const XlaSerializedCacheKey& key) const;
const DeviceType device_type_;
const bool disable_strict_signature_checks_;
const std::string persistence_prefix_;
// If non-empty, JIT-compiled executables are saved to and loaded from the
// specified file system directory path.
const std::string persistent_cache_directory_;
// Cache is read-only if set to true.
const bool persistent_cache_directory_read_only_;
DeviceExecutablePersistor(const DeviceExecutablePersistor&) = delete;
void operator=(const DeviceExecutablePersistor&) = delete;
};
template <typename ExecutableType, typename ClientType>
DeviceExecutablePersistor<ExecutableType, ClientType>::
DeviceExecutablePersistor(const Config& config,
const DeviceType& device_type)
: device_type_(device_type),
disable_strict_signature_checks_(config.disable_strict_signature_checks),
persistence_prefix_(config.persistence_prefix),
persistent_cache_directory_(config.persistent_cache_directory),
persistent_cache_directory_read_only_(
config.persistent_cache_directory_read_only) {}
template <typename ExecutableType, typename ClientType>
std::string DeviceExecutablePersistor<ExecutableType, ClientType>::GetFilePath(
const XlaSerializedCacheKey& key) const {
const std::string file_name = XlaSerializedCacheKeyToFileName(key);
return io::JoinPath(persistent_cache_directory_, file_name);
}
template <typename ExecutableType, typename ClientType>
XlaSerializedCacheKey
DeviceExecutablePersistor<ExecutableType, ClientType>::BuildSerializedCacheKey(
uint64_t signature_hash, const xla::HloModuleProto& hlo_module,
bool compiled_using_pjrt) const {
XlaSerializedCacheKey key;
key.set_signature_fingerprint(signature_hash);
key.set_cluster_fingerprint(DeterministicProtoHash64(hlo_module));
key.set_device_type(device_type().type_string());
key.set_prefix(persistence_prefix());
key.set_compiled_using_pjrt(compiled_using_pjrt);
return key;
}
template <typename ExecutableType, typename ClientType>
XlaSerializedCacheKey
DeviceExecutablePersistor<ExecutableType, ClientType>::BuildSerializedCacheKey(
uint64_t signature_hash, const xla::HloModuleProto& hlo_module) const {
return BuildSerializedCacheKey(signature_hash, hlo_module, false);
}
// This template specialization sets compiled_using_prjt to true in the cache
// key when the template arguments are PjRtLoadedExecutable and PjRtClient.
template <>
inline XlaSerializedCacheKey
DeviceExecutablePersistor<xla::PjRtLoadedExecutable, xla::PjRtClient>::
BuildSerializedCacheKey(uint64_t signature_hash,
const xla::HloModuleProto& hlo_module) const {
return BuildSerializedCacheKey(signature_hash, hlo_module, true);
}
template <typename ExecutableType, typename ClientType>
absl::StatusOr<std::optional<XlaSerializedCacheEntry>>
DeviceExecutablePersistor<ExecutableType, ClientType>::TryToReadSerializedEntry(
const XlaSerializedCacheKey& key) const {
Env* env = Env::Default();
const std::string file_path = GetFilePath(key);
if (!env->FileExists(file_path).ok()) {
return absl::StatusOr<std::optional<XlaSerializedCacheEntry>>(std::nullopt);
}
XlaSerializedCacheEntry entry;
TF_RETURN_IF_ERROR(ReadTextOrBinaryProto(env, file_path, &entry));
return std::optional<XlaSerializedCacheEntry>(entry);
}
template <typename ExecutableType, typename ClientType>
absl::Status
DeviceExecutablePersistor<ExecutableType, ClientType>::VerifyLoadedCacheEntry(
const XlaSerializedCacheKey& key, const xla::HloModuleProto& hlo_module,
const XlaSerializedCacheEntry& entry) const {
XLA_SCOPED_LOGGING_TIMER(absl::StrCat("Verifying loaded cache entry: ",
hlo_module.entry_computation_name()));
if (!AreSerializedProtosEqual(key, entry.key())) {
VLOG(2) << "Serialized cache key does not match:\n"
<< "got:\n"
<< entry.key().DebugString() << "\nexpected:\n"
<< key.DebugString() << "\n";
return absl::InvalidArgumentError("Serialized cache key does not match.");
}
// Perform a stricter (slower) check of the snapshot to verify that they
// match exactly.
if (!disable_strict_signature_checks_) {
if (!AreSerializedProtosEqual(hlo_module, entry.hlo_module())) {
VLOG(2) << "HLOs do not match:\n"
<< "got:\n"
<< hlo_module.DebugString() << "\nexpected:\n"
<< entry.hlo_module().DebugString() << "\n";
return absl::InvalidArgumentError("Serialized HLO does not match.");
}
}
if (entry.executable().empty()) {
return absl::InvalidArgumentError("No binary found in serialized entry.");
}
return absl::OkStatus();
}
template <typename ExecutableType, typename ClientType>
absl::Status
DeviceExecutablePersistor<ExecutableType, ClientType>::SaveSerializedEntry(
const XlaSerializedCacheEntry& entry) const {
Env* env = Env::Default();
TF_RETURN_IF_ERROR(env->RecursivelyCreateDir(persistent_cache_directory_));
// The cache on the filesystem can be read while we're writing out the proto.
// To prevent reads of partially-written files, we write the proto to a temp
// file, then move it into place once we're done writing. And we warn the
// user if these moves are not known to be atomic.
bool has_atomic_move = false;
env->HasAtomicMove(persistent_cache_directory_, &has_atomic_move)
.IgnoreError();
if (!has_atomic_move) {
LOG_EVERY_POW_2(WARNING)
<< "Filesystem for XLA persistent cache at "
<< persistent_cache_directory_
<< " does not support atomic moves. Therefore the persistent cache is "
"racy if you have multiple XLA compilations occurring "
"simultaneously! You have been warned. :)";
}
// Write to temp location, then when that completes, atomically move into the
// final location.
std::string temp_path =
io::JoinPath(persistent_cache_directory_,
XlaSerializedCacheKeyToFileName(entry.key()));
if (!env->CreateUniqueFileName(&temp_path, ".tmp")) {
return absl::UnavailableError(absl::StrCat(
"Could not create a unique file inside ", persistent_cache_directory_));
}
TF_RETURN_IF_ERROR(WriteBinaryProto(env, temp_path, entry));
return env->RenameFile(temp_path, GetFilePath(entry.key()));
}
template <typename ExecutableType, typename ClientType>
absl::StatusOr<XlaSerializedCacheEntry>
DeviceExecutablePersistor<ExecutableType, ClientType>::SerializeEntry(
uint64_t signature_hash, const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& compilation_result,
const ExecutableType& executable,
DeviceCompilerClient<ExecutableType, ClientType>* compiler_client) const {
XlaSerializedCacheEntry serialized_entry;
const xla::HloModuleProto& hlo_module =
compilation_result.computation->proto();
*serialized_entry.mutable_key() =
BuildSerializedCacheKey(signature_hash, hlo_module);
*serialized_entry.mutable_hlo_module() = hlo_module;
// XLA compiler supports exporting executables as an AOT compilation result
// to avoid running potentially expensive compilation pipeline twice.
// Check if XLA compiler can export available executable.
if (auto serialized_executable =
compiler_client->SerializeExecutable(executable);
serialized_executable.ok()) {
serialized_entry.set_executable(std::move(*serialized_executable));
return serialized_entry;
} else if (serialized_executable.status().code() == error::UNIMPLEMENTED) {
VLOG(1) << "Executable export is not implemented";
} else {
return serialized_executable.status();
}
TF_ASSIGN_OR_RETURN(
auto serialized_executable,
compiler_client->BuildSerializedExecutable(options, compilation_result));
serialized_entry.set_executable(std::move(serialized_executable));
return serialized_entry;
}
template <typename ExecutableType, typename ClientType>
std::optional<StatusOr<std::unique_ptr<ExecutableType>>>
DeviceExecutablePersistor<ExecutableType, ClientType>::TryToLoadExecutable(
uint64_t signature_hash, const std::string& signature_str,
const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& compilation_result,
DeviceCompilerClient<ExecutableType, ClientType>* compiler_client) const {
if (persistent_cache_directory_.empty()) {
return std::nullopt;
}
const xla::HloModuleProto& hlo_module =
compilation_result.computation->proto();
XlaSerializedCacheKey cache_key =
BuildSerializedCacheKey(signature_hash, hlo_module);
std::optional<XlaSerializedCacheEntry> serialized_entry;
{
XLA_SCOPED_LOGGING_TIMER(
absl::StrCat("Try loading serialized cache entry:", signature_str));
TF_ASSIGN_OR_RETURN(serialized_entry, TryToReadSerializedEntry(cache_key));
}
if (!serialized_entry.has_value()) {
return std::nullopt;
}
TF_RETURN_IF_ERROR(
VerifyLoadedCacheEntry(cache_key, hlo_module, *serialized_entry));
VLOG(1) << "Loading cached entry for: " << signature_str;
return compiler_client->LoadExecutable(options, compilation_result,
serialized_entry->executable());
}
template <typename ExecutableType, typename ClientType>
absl::Status
DeviceExecutablePersistor<ExecutableType, ClientType>::TryToPersistExecutable(
uint64_t signature_hash, const std::string& signature_str,
const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& compilation_result,
const ExecutableType& executable,
DeviceCompilerClient<ExecutableType, ClientType>* client) const {
if (persistent_cache_directory_.empty() ||
persistent_cache_directory_read_only_) {
VLOG(1) << "Not persisting executable. No `persistent_cache_directory` "
"provided or cache is read-only.";
return absl::OkStatus();
}
XLA_SCOPED_LOGGING_TIMER(
absl::StrCat("Serializing and saving cache entry: ", signature_str));
TF_ASSIGN_OR_RETURN(XlaSerializedCacheEntry serialized_entry,
SerializeEntry(signature_hash, options,
compilation_result, executable, client));
TF_RETURN_IF_ERROR(SaveSerializedEntry(std::move(serialized_entry)));
VLOG(2) << "XlaSerializedCacheEntry saved for signature: [" << signature_str
<< "] with signature hash: " << signature_hash;
return absl::OkStatus();
}
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_DEVICE_EXECUTABLE_PERSISTOR_H_
tensorflow/compiler/jit/device_executable_persistor_test.cc
+637
View File
@@ -0,0 +1,637 @@
/* 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/compiler/jit/device_executable_persistor.h"
#include <stdlib.h>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "tensorflow/cc/framework/scope.h"
#include "tensorflow/cc/ops/function_ops.h"
#include "tensorflow/cc/ops/math_ops.h"
#include "tensorflow/compiler/jit/device_compiler_client.h"
#include "tensorflow/compiler/jit/pjrt_device_compiler_client.h"
#include "tensorflow/compiler/jit/xla_compilation_cache.pb.h"
#include "tensorflow/compiler/jit/xla_device_compiler_client.h"
#include "xla/client/client_library.h"
#include "xla/client/executable_build_options.h"
#include "xla/client/local_client.h"
#include "xla/pjrt/pjrt_client.h"
#include "xla/pjrt/plugin/xla_cpu/cpu_client_options.h"
#include "xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status_matchers.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/tfrt/common/create_pjrt_client_util.h"
#include "tensorflow/core/tfrt/common/pjrt_util.h"
namespace tensorflow {
namespace {
using ::testing::_;
using ::testing::ByMove;
using ::testing::Return;
using XlaDeviceExecutablePersistor =
DeviceExecutablePersistor<xla::LocalExecutable, xla::LocalClient>;
using PjRtDeviceExecutablePersistor =
DeviceExecutablePersistor<xla::PjRtLoadedExecutable, xla::PjRtClient>;
class DeviceExecutionPersistorTest : public ::testing::Test {
protected:
void SetUp() override {
xla_compiler_client_ = std::make_unique<XlaDeviceCompilerClient>(
xla::ClientLibrary::LocalClientOrDie());
TF_ASSERT_OK(CreatePjRtCompilerClient());
XlaOpRegistry::RegisterCompilationKernels();
flib_def_ = std::make_unique<FunctionLibraryDefinition>(
OpRegistry::Global(), FunctionDefLibrary());
cache_dir_ = testing::TmpDir();
TF_ASSERT_OK_AND_ASSIGN(compilation_result_add_,
BuildSampleCompilationResult());
}
absl::StatusOr<std::unique_ptr<xla::LocalExecutable>>
BuildSampleExecutable() {
return xla_compiler_client_->BuildExecutable(DefaultXlaOptions(),
compilation_result_add_);
}
absl::StatusOr<std::unique_ptr<xla::PjRtLoadedExecutable>>
BuildSamplePjRtExecutable() {
return pjrt_compiler_client_->BuildExecutable(DefaultPjRtOptions(),
compilation_result_add_);
}
absl::StatusOr<XlaCompiler::CompilationResult> BuildSampleCompilationResult(
bool mul = false) {
std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
Scope scope = Scope::NewRootScope().ExitOnError();
auto a = ops::_Arg(scope.WithOpName("A"), DT_INT32, 0);
auto b = ops::_Arg(scope.WithOpName("B"), DT_INT32, 1);
if (mul) {
auto c = ops::Mul(scope.WithOpName("C"), a, b);
auto d = ops::_Retval(scope.WithOpName("D"), c, 0);
TF_RETURN_IF_ERROR(scope.ToGraph(graph.get()));
} else {
auto c = ops::Add(scope.WithOpName("C"), a, b);
auto d = ops::_Retval(scope.WithOpName("D"), c, 0);
TF_RETURN_IF_ERROR(scope.ToGraph(graph.get()));
}
// Builds a description of the arguments.
std::vector<XlaCompiler::Argument> args(2);
args[0].kind = XlaCompiler::Argument::kParameter;
args[0].type = DT_INT32;
args[0].shape = TensorShape({2});
args[1].kind = XlaCompiler::Argument::kParameter;
args[1].type = DT_INT32;
args[1].shape = TensorShape({2});
// Compiles the graph.
XlaCompiler compiler(DefaultXlaOptions());
XlaCompiler::CompilationResult compilation_result;
TF_RETURN_IF_ERROR(compiler.CompileGraph(XlaCompiler::CompileOptions(),
"graph", std::move(graph), args,
&compilation_result));
return compilation_result;
}
XlaCompiler::Options DefaultXlaOptions() {
XlaCompiler::Options options;
options.device_type = DeviceType(DEVICE_CPU_XLA_JIT);
options.client = xla_compiler_client_->client();
options.flib_def = flib_def_.get();
return options;
}
XlaCompiler::Options DefaultPjRtOptions() {
XlaCompiler::Options options;
options.device_type = DeviceType(DEVICE_CPU_XLA_JIT);
options.client = nullptr;
options.flib_def = flib_def_.get();
return options;
}
absl::Status CreatePjRtCompilerClient() {
// Create PjRtClient manually while GetOrCreatePjRtClient() is WIP.
xla::CpuClientOptions options;
options.asynchronous = true;
options.cpu_device_count = 1;
TF_RETURN_IF_ERROR(SetPjRtClientInTFGlobalResourceManager(
DEVICE_CPU_XLA_JIT, xla::GetXlaPjrtCpuClient(options).value()));
TF_ASSIGN_OR_RETURN(auto pjrt_client,
GetOrCreatePjRtClient(DeviceType(DEVICE_CPU_XLA_JIT)));
pjrt_compiler_client_ =
std::make_unique<PjRtDeviceCompilerClient>(pjrt_client);
return absl::OkStatus();
}
std::unique_ptr<FunctionLibraryDefinition> flib_def_;
std::unique_ptr<XlaDeviceCompilerClient> xla_compiler_client_;
std::unique_ptr<PjRtDeviceCompilerClient> pjrt_compiler_client_;
XlaCompiler::CompilationResult compilation_result_add_;
std::string serialized_xla_executable_ = "serialized_xla_executable";
std::string serialized_pjrt_executable_ = "serialized_pjrt_executable";
std::string cache_dir_;
};
// Using a mock to make testing different branches and triggering errors easier.
// Currently the `XlaDeviceCompilerClient`'s load/serialize functions don't work
// with the current test setup.
// TODO(b/255826209): Look into using a real object for most tests.
class MockXlaCompilerClient : public XlaDeviceCompilerClient {
public:
MockXlaCompilerClient() : XlaDeviceCompilerClient(nullptr) {}
MOCK_METHOD(absl::StatusOr<std::string>, SerializeExecutable,
(const xla::LocalExecutable& executable), (override));
MOCK_METHOD(absl::StatusOr<std::string>, BuildSerializedExecutable,
(const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& result),
(override));
MOCK_METHOD(absl::StatusOr<std::unique_ptr<xla::LocalExecutable>>,
LoadExecutable,
(const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& result,
const std::string& serialized_executable),
(override));
};
class MockPjRtCompilerClient : public PjRtDeviceCompilerClient {
public:
MockPjRtCompilerClient() : PjRtDeviceCompilerClient(nullptr) {}
MOCK_METHOD(absl::StatusOr<std::string>, SerializeExecutable,
(const xla::PjRtLoadedExecutable& executable), (override));
MOCK_METHOD(absl::StatusOr<std::string>, BuildSerializedExecutable,
(const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& result),
(override));
MOCK_METHOD(absl::StatusOr<std::unique_ptr<xla::PjRtLoadedExecutable>>,
LoadExecutable,
(const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& result,
const std::string& serialized_executable),
(override));
};
std::string GetFilePath(XlaSerializedCacheKey key,
const std::string& persistent_cache_dir) {
static constexpr char kXlaSerializedCacheKeySeparator[] = "__";
std::string file_name = absl::StrCat(
key.prefix(), key.prefix().empty() ? "" : kXlaSerializedCacheKeySeparator,
key.signature_fingerprint(), kXlaSerializedCacheKeySeparator,
key.cluster_fingerprint(), kXlaSerializedCacheKeySeparator,
key.device_type(),
key.compiled_using_pjrt()
? absl::StrCat(kXlaSerializedCacheKeySeparator, "pjrt")
: "",
".pb");
return io::JoinPath(persistent_cache_dir, file_name);
}
absl::StatusOr<XlaSerializedCacheEntry> ReadCacheEntryFromFile(
XlaSerializedCacheKey key, const std::string& persistent_cache_dir) {
std::string file_path = GetFilePath(key, persistent_cache_dir);
XlaSerializedCacheEntry entry;
TF_RETURN_IF_ERROR(ReadTextOrBinaryProto(Env::Default(), file_path, &entry));
return entry;
}
XlaSerializedCacheKey CreateCacheKey(
uint64_t signature_hash,
const XlaCompiler::CompilationResult& compilation_result,
const DeviceType& device_type, const std::string& persistence_prefix,
bool compiled_using_pjrt = false) {
XlaSerializedCacheKey key;
key.set_signature_fingerprint(signature_hash);
key.set_cluster_fingerprint(
DeterministicProtoHash64(compilation_result.computation->proto()));
key.set_device_type(device_type.type_string());
key.set_prefix(persistence_prefix);
key.set_compiled_using_pjrt(compiled_using_pjrt);
return key;
}
TEST_F(DeviceExecutionPersistorTest, PersistCacheDirNotSet) {
XlaDeviceExecutablePersistor::Config config(
/*persistent_cache_directory=*/"",
/*disable_strict_signature_checks=*/false,
/*persistence_prefix=*/"xla");
XlaDeviceExecutablePersistor persistor(config,
DefaultXlaOptions().device_type);
MockXlaCompilerClient mock_client;
TF_ASSERT_OK_AND_ASSIGN(auto executable, BuildSampleExecutable());
TF_EXPECT_OK(persistor.TryToPersistExecutable(
/*signature_hash=*/123, "signature_string", DefaultXlaOptions(),
compilation_result_add_, *executable, &mock_client));
auto key =
CreateCacheKey(/*signature_hash=*/123, compilation_result_add_,
persistor.device_type(), persistor.persistence_prefix());
auto entry = ReadCacheEntryFromFile(key, "");
EXPECT_FALSE(entry.ok());
}
TEST_F(DeviceExecutionPersistorTest, PersistCacheDirReadOnly) {
XlaDeviceExecutablePersistor::Config config(
/*persistent_cache_directory=*/"cache_dir_",
/*disable_strict_signature_checks=*/false,
/*persistence_prefix=*/"xla",
/*persistent_cache_directory_read_only=*/true);
XlaDeviceExecutablePersistor persistor(config,
DefaultXlaOptions().device_type);
MockXlaCompilerClient mock_client;
TF_ASSERT_OK_AND_ASSIGN(auto executable, BuildSampleExecutable());
TF_EXPECT_OK(persistor.TryToPersistExecutable(
/*signature_hash=*/123, "signature_string", DefaultXlaOptions(),
compilation_result_add_, *executable, &mock_client));
auto key =
CreateCacheKey(/*signature_hash=*/123, compilation_result_add_,
persistor.device_type(), persistor.persistence_prefix());
auto entry = ReadCacheEntryFromFile(key, "");
EXPECT_FALSE(entry.ok());
}
TEST_F(DeviceExecutionPersistorTest, PersistSerializeAlreadyBuiltExecutable) {
XlaDeviceExecutablePersistor::Config config(
/*persistent_cache_directory=*/cache_dir_,
/*disable_strict_signature_checks=*/false,
/*persistence_prefix=*/"xla");
XlaDeviceExecutablePersistor persistor(config,
DefaultXlaOptions().device_type);
MockXlaCompilerClient mock_client;
EXPECT_CALL(mock_client, SerializeExecutable(_))
.WillOnce(
Return(absl::StatusOr<std::string>(serialized_xla_executable_)));
TF_ASSERT_OK_AND_ASSIGN(auto executable, BuildSampleExecutable());
TF_EXPECT_OK(persistor.TryToPersistExecutable(
/*signature_hash=*/123, "signature_string", DefaultXlaOptions(),
compilation_result_add_, *executable, &mock_client));
auto key =
CreateCacheKey(/*signature_hash=*/123, compilation_result_add_,
persistor.device_type(), persistor.persistence_prefix());
TF_ASSERT_OK_AND_ASSIGN(auto entry, ReadCacheEntryFromFile(key, cache_dir_));
EXPECT_EQ(entry.executable(), serialized_xla_executable_);
}
TEST_F(DeviceExecutionPersistorTest, PersistBuildSerializedExecutable) {
XlaDeviceExecutablePersistor::Config config(
/*persistent_cache_directory=*/cache_dir_,
/*disable_strict_signature_checks=*/false,
/*persistence_prefix=*/"xla");
XlaDeviceExecutablePersistor persistor(config,
DefaultXlaOptions().device_type);
MockXlaCompilerClient mock_client;
EXPECT_CALL(mock_client, SerializeExecutable(_))
.WillOnce(Return(absl::UnimplementedError("Unimplemented.")));
EXPECT_CALL(mock_client, BuildSerializedExecutable(_, _))
.WillOnce(Return(serialized_xla_executable_));
TF_ASSERT_OK_AND_ASSIGN(auto executable, BuildSampleExecutable());
TF_EXPECT_OK(persistor.TryToPersistExecutable(
/*signature_hash=*/123, "signature_string", DefaultXlaOptions(),
compilation_result_add_, *executable, &mock_client));
auto key =
CreateCacheKey(/*signature_hash=*/123, compilation_result_add_,
persistor.device_type(), persistor.persistence_prefix());
TF_ASSERT_OK_AND_ASSIGN(auto entry, ReadCacheEntryFromFile(key, cache_dir_));
EXPECT_EQ(entry.executable(), serialized_xla_executable_);
}
TEST_F(DeviceExecutionPersistorTest, PersistSerializeExecutableError) {
XlaDeviceExecutablePersistor::Config config(
/*persistent_cache_directory=*/cache_dir_,
/*disable_strict_signature_checks=*/false,
/*persistence_prefix=*/"xla");
XlaDeviceExecutablePersistor persistor(config,
DefaultXlaOptions().device_type);
MockXlaCompilerClient mock_client;
EXPECT_CALL(mock_client, SerializeExecutable(_))
.WillOnce(Return(absl::InvalidArgumentError("InvalidArgument.")));
TF_ASSERT_OK_AND_ASSIGN(auto executable, BuildSampleExecutable());
EXPECT_THAT(
persistor.TryToPersistExecutable(
/*signature_hash=*/123, "signature_string", DefaultXlaOptions(),
compilation_result_add_, *executable, &mock_client),
absl_testing::StatusIs(error::INVALID_ARGUMENT));
}
TEST_F(DeviceExecutionPersistorTest, PersistExecutableEmpty) {
XlaDeviceExecutablePersistor::Config config(
/*persistent_cache_directory=*/cache_dir_,
/*disable_strict_signature_checks=*/false,
/*persistence_prefix=*/"xla");
XlaDeviceExecutablePersistor persistor(config,
DefaultXlaOptions().device_type);
MockXlaCompilerClient mock_client;
xla::LocalExecutable empty_executable(
nullptr, nullptr,
GetExecutableBuildOptions(DefaultXlaOptions(), compilation_result_add_,
0));
EXPECT_CALL(mock_client, SerializeExecutable(_))
.WillOnce(Return(absl::FailedPreconditionError("Failed precondition.")));
EXPECT_THAT(
persistor.TryToPersistExecutable(
/*signature_hash=*/123, "signature_string", DefaultXlaOptions(),
compilation_result_add_, empty_executable, &mock_client),
absl_testing::StatusIs(error::FAILED_PRECONDITION));
}
TEST_F(DeviceExecutionPersistorTest, LoadCacheDirNotSet) {
XlaDeviceExecutablePersistor::Config config(
/*persistent_cache_directory=*/"",
/*disable_strict_signature_checks=*/false,
/*persistence_prefix=*/"xla");
XlaDeviceExecutablePersistor persistor(config,
DefaultXlaOptions().device_type);
MockXlaCompilerClient mock_client;
auto executable = persistor.TryToLoadExecutable(
123, "signature_string", DefaultXlaOptions(), compilation_result_add_,
&mock_client);
EXPECT_FALSE(executable.has_value());
}
TEST_F(DeviceExecutionPersistorTest, LoadSuccess) {
XlaDeviceExecutablePersistor::Config config(
/*persistent_cache_directory=*/cache_dir_,
/*disable_strict_signature_checks=*/false,
/*persistence_prefix=*/"xla");
XlaDeviceExecutablePersistor persistor(config,
DefaultXlaOptions().device_type);
MockXlaCompilerClient mock_client;
TF_ASSERT_OK_AND_ASSIGN(auto executable, BuildSampleExecutable());
EXPECT_CALL(mock_client, LoadExecutable(_, _, serialized_xla_executable_))
.WillOnce(Return(ByMove(std::move(executable))));
auto loaded_executable = persistor.TryToLoadExecutable(
/*signature_hash=*/123, "signature_string", DefaultXlaOptions(),
compilation_result_add_, &mock_client);
EXPECT_TRUE(loaded_executable.has_value());
EXPECT_TRUE(loaded_executable.value().ok());
EXPECT_TRUE((*loaded_executable.value())->executable() != nullptr);
}
TEST_F(DeviceExecutionPersistorTest, LoadFileDoesntExist) {
XlaDeviceExecutablePersistor::Config config(
/*persistent_cache_directory=*/cache_dir_,
/*disable_strict_signature_checks=*/false,
/*persistence_prefix=*/"xla");
XlaDeviceExecutablePersistor persistor(config,
DefaultXlaOptions().device_type);
MockXlaCompilerClient mock_client;
// Try to load an executable for a different signature hash (which hasn't been
// persisted).
auto loaded_executable = persistor.TryToLoadExecutable(
/*signature_hash=*/12345, "different_signature", DefaultXlaOptions(),
compilation_result_add_, &mock_client);
EXPECT_FALSE(loaded_executable.has_value());
}
TEST_F(DeviceExecutionPersistorTest, LoadSerializedKeyMismatch) {
XlaDeviceExecutablePersistor::Config config(
/*persistent_cache_directory=*/cache_dir_,
/*disable_strict_signature_checks=*/false,
/*persistence_prefix=*/"xla");
XlaDeviceExecutablePersistor persistor(config,
DefaultXlaOptions().device_type);
auto key1 =
CreateCacheKey(/*signature_hash=*/123, compilation_result_add_,
persistor.device_type(), persistor.persistence_prefix());
auto key2 =
CreateCacheKey(/*signature_hash=*/456, compilation_result_add_,
persistor.device_type(), persistor.persistence_prefix());
// File for key2 contains the same content as key1.
TF_ASSERT_OK(Env::Default()->CopyFile(
GetFilePath(key1, persistor.persistent_cache_directory()),
GetFilePath(key2, persistor.persistent_cache_directory())));
MockXlaCompilerClient mock_client;
// Try to load an executable from file corresponding to key2 (whose file
// content corresponds to key1).
auto loaded_executable = persistor.TryToLoadExecutable(
/*signature_hash=*/456, "different_signature", DefaultXlaOptions(),
compilation_result_add_, &mock_client);
EXPECT_TRUE(loaded_executable.has_value());
EXPECT_FALSE(loaded_executable->ok());
EXPECT_THAT(loaded_executable.value(),
absl_testing::StatusIs(error::INVALID_ARGUMENT));
}
TEST_F(DeviceExecutionPersistorTest, LoadSerializedHloMismatch) {
XlaDeviceExecutablePersistor::Config config(
/*persistent_cache_directory=*/cache_dir_,
/*disable_strict_signature_checks=*/false,
/*persistence_prefix=*/"xla");
XlaDeviceExecutablePersistor persistor(config,
DefaultXlaOptions().device_type);
TF_ASSERT_OK_AND_ASSIGN(auto compilation_result_mul,
BuildSampleCompilationResult(true));
auto key1 =
CreateCacheKey(/*signature_hash=*/123, compilation_result_add_,
persistor.device_type(), persistor.persistence_prefix());
auto key2 =
CreateCacheKey(/*signature_hash=*/123, compilation_result_mul,
persistor.device_type(), persistor.persistence_prefix());
// Read serialized entry corresponding to key1.
XlaSerializedCacheEntry entry;
TF_ASSERT_OK(ReadTextOrBinaryProto(
Env::Default(), GetFilePath(key1, persistor.persistent_cache_directory()),
&entry));
// Change the entry's key to key2.
*entry.mutable_key() = key2;
// Write the modified entry to file corresponding to key2.
TF_ASSERT_OK(WriteBinaryProto(
Env::Default(), GetFilePath(key2, persistor.persistent_cache_directory()),
entry));
MockXlaCompilerClient mock_client;
// Try to load executable corresponding to key2 (whose file contains HLO
// corresponding to key1).
auto loaded_executable = persistor.TryToLoadExecutable(
/*signature_hash=*/123, "signature", DefaultXlaOptions(),
compilation_result_mul, &mock_client);
EXPECT_TRUE(loaded_executable.has_value());
EXPECT_FALSE(loaded_executable->ok());
EXPECT_THAT(loaded_executable.value(),
absl_testing::StatusIs(error::INVALID_ARGUMENT));
}
TEST_F(DeviceExecutionPersistorTest, LoadStrictChecksDisabled) {
XlaDeviceExecutablePersistor::Config config(
/*persistent_cache_directory=*/cache_dir_,
/*disable_strict_signature_checks=*/true,
/*persistence_prefix=*/"xla");
XlaDeviceExecutablePersistor persistor(config,
DefaultXlaOptions().device_type);
TF_ASSERT_OK_AND_ASSIGN(auto compilation_result_mul,
BuildSampleCompilationResult(true));
auto key1 =
CreateCacheKey(/*signature_hash=*/123, compilation_result_add_,
persistor.device_type(), persistor.persistence_prefix());
auto key2 =
CreateCacheKey(/*signature_hash=*/123, compilation_result_mul,
persistor.device_type(), persistor.persistence_prefix());
// Read serialized entry corresponding to key1.
XlaSerializedCacheEntry entry;
TF_ASSERT_OK(ReadTextOrBinaryProto(
Env::Default(), GetFilePath(key1, persistor.persistent_cache_directory()),
&entry));
// Change the entry's key to key2.
*entry.mutable_key() = key2;
// Write the modified entry to file corresponding to key2.
TF_ASSERT_OK(WriteBinaryProto(
Env::Default(), GetFilePath(key2, persistor.persistent_cache_directory()),
entry));
MockXlaCompilerClient mock_client;
TF_ASSERT_OK_AND_ASSIGN(auto executable, BuildSampleExecutable());
EXPECT_CALL(mock_client, LoadExecutable(_, _, serialized_xla_executable_))
.WillOnce(Return(ByMove(std::move(executable))));
auto loaded_executable =
persistor.TryToLoadExecutable(123, "signature", DefaultXlaOptions(),
compilation_result_mul, &mock_client);
EXPECT_TRUE(loaded_executable.has_value());
EXPECT_TRUE(loaded_executable->ok());
}
TEST_F(DeviceExecutionPersistorTest, LoadSerializedExecutableEmpty) {
XlaDeviceExecutablePersistor::Config config(
/*persistent_cache_directory=*/cache_dir_,
/*disable_strict_signature_checks=*/false,
/*persistence_prefix=*/"xla");
XlaDeviceExecutablePersistor persistor(config,
DefaultXlaOptions().device_type);
auto key =
CreateCacheKey(/*signature_hash=*/123, compilation_result_add_,
persistor.device_type(), persistor.persistence_prefix());
// Read serialized entry.
XlaSerializedCacheEntry entry;
TF_ASSERT_OK(ReadTextOrBinaryProto(
Env::Default(), GetFilePath(key, persistor.persistent_cache_directory()),
&entry));
entry.clear_executable();
// Write entry to another file.
TF_ASSERT_OK(WriteBinaryProto(
Env::Default(), GetFilePath(key, persistor.persistent_cache_directory()),
entry));
MockXlaCompilerClient mock_client;
auto loaded_executable = persistor.TryToLoadExecutable(
/*signature_hash=*/123, "signature", DefaultXlaOptions(),
compilation_result_add_, &mock_client);
EXPECT_TRUE(loaded_executable.has_value());
EXPECT_FALSE(loaded_executable->ok());
EXPECT_THAT(loaded_executable.value(),
absl_testing::StatusIs(error::INVALID_ARGUMENT));
}
TEST_F(DeviceExecutionPersistorTest, PersistPjRtAndXlaExecutables) {
// Persist PJRT executable.
PjRtDeviceExecutablePersistor::Config pjrt_config(
/*persistent_cache_directory=*/cache_dir_,
/*disable_strict_signature_checks=*/false,
/*persistence_prefix=*/"xla");
PjRtDeviceExecutablePersistor pjrt_persistor(
pjrt_config, DefaultPjRtOptions().device_type);
MockPjRtCompilerClient mock_pjrt_client;
EXPECT_CALL(mock_pjrt_client, SerializeExecutable(_))
.WillOnce(Return(serialized_pjrt_executable_));
TF_ASSERT_OK_AND_ASSIGN(auto pjrt_executable, BuildSamplePjRtExecutable());
TF_EXPECT_OK(pjrt_persistor.TryToPersistExecutable(
/*signature_hash=*/123, "signature_string", DefaultPjRtOptions(),
compilation_result_add_, *pjrt_executable, &mock_pjrt_client));
// Persist XLA executable.
XlaDeviceExecutablePersistor::Config xla_config(
/*persistent_cache_directory=*/cache_dir_,
/*disable_strict_signature_checks=*/false,
/*persistence_prefix=*/"xla");
XlaDeviceExecutablePersistor xla_persistor(xla_config,
DefaultXlaOptions().device_type);
MockXlaCompilerClient mock_xla_client;
EXPECT_CALL(mock_xla_client, SerializeExecutable(_))
.WillOnce(Return(serialized_xla_executable_));
TF_ASSERT_OK_AND_ASSIGN(auto xla_executable, BuildSampleExecutable());
TF_EXPECT_OK(xla_persistor.TryToPersistExecutable(
/*signature_hash=*/123, "signature_string", DefaultXlaOptions(),
compilation_result_add_, *xla_executable, &mock_xla_client));
// Read and verify serialized PJRT executable.
auto pjrt_key = CreateCacheKey(
/*signature_hash=*/123, compilation_result_add_,
pjrt_persistor.device_type(), pjrt_persistor.persistence_prefix(), true);
TF_ASSERT_OK_AND_ASSIGN(auto entry,
ReadCacheEntryFromFile(pjrt_key, cache_dir_));
EXPECT_EQ(entry.executable(), serialized_pjrt_executable_);
// Read and verify serialized XLA executable.
auto xla_key = CreateCacheKey(/*signature_hash=*/123, compilation_result_add_,
pjrt_persistor.device_type(),
pjrt_persistor.persistence_prefix());
TF_ASSERT_OK_AND_ASSIGN(entry, ReadCacheEntryFromFile(xla_key, cache_dir_));
EXPECT_EQ(entry.executable(), serialized_xla_executable_);
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/device_util.cc
+236
View File
@@ -0,0 +1,236 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/jit/device_util.h"
#include "absl/algorithm/container.h"
#include "absl/container/flat_hash_set.h"
#include "xla/status_macros.h"
namespace tensorflow {
namespace jit {
void DeviceSet::Insert(DeviceId device_id) {
int word_index = device_id.id() / kWordSize;
int bit_index = device_id.id() % kWordSize;
const int storage_size = storage_.size();
if (word_index >= storage_size) {
storage_.resize(word_index + 1, 0);
}
storage_[word_index] |= (1ull << bit_index);
}
void DeviceSet::UnionWith(const DeviceSet& other) {
if (other.storage_.size() > storage_.size()) {
storage_.resize(other.storage_.size(), 0);
}
for (int i = 0, end = other.storage_.size(); i < end; i++) {
storage_[i] |= other.storage_[i];
}
}
bool DeviceSet::IsEmpty() const {
return absl::c_all_of(storage_, [&](uint64_t val) { return val == 0; });
}
absl::StatusOr<DeviceId> DeviceInfoCache::GetIdFor(absl::string_view name) {
TF_RET_CHECK(!name.empty());
auto it = name_to_id_.find(name);
if (it != name_to_id_.end()) {
return it->second;
}
int new_id = names_.size();
names_.push_back(std::string(name));
id_to_device_type_.push_back(std::make_unique<DeviceType>(""));
DeviceType* device_type = id_to_device_type_.back().get();
TF_RETURN_IF_ERROR(DeviceNameToDeviceType(names_.back(), device_type));
is_cpu_.push_back(device_type->type_string() == DEVICE_CPU);
is_gpu_.push_back(device_type->type_string() == DEVICE_GPU);
name_to_id_.emplace(std::string(name), DeviceId(new_id));
const XlaOpRegistry::DeviceRegistration* compilation_device;
if (!XlaOpRegistry::GetCompilationDevice(device_type->type(),
&compilation_device)) {
compilation_device = nullptr;
}
id_to_compilation_device_.push_back(compilation_device);
return DeviceId(new_id);
}
std::string DeviceInfoCache::DebugString(const DeviceSet& device_set) const {
std::vector<std::string> names;
device_set.ForEach([&](DeviceId device_id) {
names.push_back(std::string(GetNameFor(device_id)));
return true;
});
return absl::StrCat("[", absl::StrJoin(names, ","), "]");
}
} // namespace jit
absl::Status DeviceNameToDeviceType(const std::string& device,
DeviceType* device_type) {
DeviceNameUtils::ParsedName parsed;
if (!DeviceNameUtils::ParseFullName(device, &parsed)) {
return absl::InternalError(
absl::StrCat("Malformed assigned device '", device, "'"));
}
*device_type = DeviceType(parsed.type);
return absl::OkStatus();
}
absl::StatusOr<std::optional<jit::DeviceId>> PickDeviceForXlaImpl(
const jit::DeviceInfoCache& device_info_cache,
const jit::DeviceSet& devices, bool allow_mixing_unknown_and_cpu,
bool failure_to_pick_is_error) {
#define FAILED_TO_PICK_DEVICE(failing_status) \
do { \
if (failure_to_pick_is_error) { \
return failing_status; \
} else { \
return {std::nullopt}; \
} \
} while (false)
std::optional<jit::DeviceId> maybe_gpu_device;
std::optional<jit::DeviceId> maybe_cpu_device;
std::optional<jit::DeviceId> maybe_unknown_device;
bool multiple_cpu_devices = false;
bool multiple_gpu_devices = false;
bool multiple_unknown_devices = false;
// Returns 'true' if d0 and d1 are conflicting devices. If they are
// compatible, update d1 with a more specific one.
// TODO(sanjoy): Cache DeviceNameUtils::ParsedName inside device_info_cache.
const auto is_multiple_devices =
[&](const jit::DeviceId& d0, std::optional<jit::DeviceId>* d1) -> bool {
const absl::string_view name0 = device_info_cache.GetNameFor(d0);
const absl::string_view name1 = device_info_cache.GetNameFor(d1->value());
DeviceNameUtils::ParsedName parsed0, parsed1;
if (!DeviceNameUtils::ParseFullName(name0, &parsed0) ||
!DeviceNameUtils::ParseFullName(name1, &parsed1) ||
!DeviceNameUtils::AreCompatibleDevNames(parsed0, parsed1)) {
return true;
}
if (DeviceNameUtils::IsSpecification(parsed0, parsed1)) {
return false;
}
if (DeviceNameUtils::IsSpecification(parsed1, parsed0)) {
*d1 = d0;
return false;
}
return true;
};
devices.ForEach([&](jit::DeviceId device) {
if (device_info_cache.IsGpu(device)) {
if (maybe_gpu_device) {
multiple_gpu_devices = is_multiple_devices(device, &maybe_gpu_device);
if (multiple_gpu_devices) return false;
} else {
maybe_gpu_device = device;
}
} else if (device_info_cache.IsCpu(device)) {
if (maybe_cpu_device) {
multiple_cpu_devices = is_multiple_devices(device, &maybe_cpu_device);
if (multiple_cpu_devices) return false;
} else {
maybe_cpu_device = device;
}
} else {
if (maybe_unknown_device) {
multiple_unknown_devices = true;
return false;
}
maybe_unknown_device = device;
}
return true;
});
if (multiple_cpu_devices) {
FAILED_TO_PICK_DEVICE(absl::InternalError(absl::StrCat(
"Multiple CPU devices ", device_info_cache.DebugString(devices))));
}
if (multiple_gpu_devices) {
FAILED_TO_PICK_DEVICE(absl::InternalError(absl::StrCat(
"Multiple GPU devices ", device_info_cache.DebugString(devices))));
}
if (multiple_unknown_devices) {
FAILED_TO_PICK_DEVICE(absl::InternalError(absl::StrCat(
"Multiple unknown devices ", device_info_cache.DebugString(devices))));
}
if (maybe_unknown_device && maybe_gpu_device) {
FAILED_TO_PICK_DEVICE(absl::InternalError(
absl::StrCat("Found both unknown and GPU devices: ",
device_info_cache.GetNameFor(*maybe_unknown_device), ", ",
device_info_cache.GetNameFor(*maybe_gpu_device))));
}
if (!allow_mixing_unknown_and_cpu) {
if (maybe_unknown_device && maybe_cpu_device) {
FAILED_TO_PICK_DEVICE(absl::InternalError(
absl::StrCat("Found both unknown and CPU devices: ",
device_info_cache.GetNameFor(*maybe_unknown_device),
", ", device_info_cache.GetNameFor(*maybe_cpu_device))));
}
}
if (maybe_gpu_device) {
return {*maybe_gpu_device};
} else if (maybe_unknown_device) {
return {*maybe_unknown_device};
} else if (maybe_cpu_device) {
return {*maybe_cpu_device};
}
FAILED_TO_PICK_DEVICE(absl::InternalError("Empty device set!"));
#undef FAILED_TO_PICK_DEVICE
}
absl::StatusOr<jit::DeviceId> PickDeviceForXla(
const jit::DeviceInfoCache& device_info_cache,
const jit::DeviceSet& devices, bool allow_mixing_unknown_and_cpu) {
TF_ASSIGN_OR_RETURN(std::optional<jit::DeviceId> device_id,
PickDeviceForXlaImpl(device_info_cache, devices,
allow_mixing_unknown_and_cpu,
/*failure_to_pick_is_error=*/true));
return *device_id;
}
absl::StatusOr<std::optional<jit::DeviceId>> MaybePickDeviceForXla(
const jit::DeviceInfoCache& device_info_cache,
const jit::DeviceSet& devices, bool allow_mixing_unknown_and_cpu) {
return PickDeviceForXlaImpl(device_info_cache, devices,
allow_mixing_unknown_and_cpu,
/*failure_to_pick_is_error=*/false);
}
} // namespace tensorflow
tensorflow/compiler/jit/device_util.h
+203
View File
@@ -0,0 +1,203 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_JIT_DEVICE_UTIL_H_
#define TENSORFLOW_COMPILER_JIT_DEVICE_UTIL_H_
#include <functional>
#include <memory>
#include "absl/container/flat_hash_map.h"
#include "absl/numeric/bits.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "xla/status_macros.h"
#include "tensorflow/core/framework/types.h"
namespace tensorflow {
namespace jit {
class DeviceInfoCache;
class DeviceSet;
// Instances of DeviceId represent TensorFlow devices as integers.
//
// This helps avoid having to manipulate device names as strings when
// auto-clustering.
class DeviceId {
public:
DeviceId(DeviceId&&) = default;
DeviceId(const DeviceId&) = default;
DeviceId& operator=(const DeviceId&) = default;
bool operator==(const DeviceId& other) const { return id() == other.id(); }
bool operator!=(const DeviceId& other) const { return !(*this == other); }
private:
int id_;
explicit DeviceId(int id) : id_(id) {}
int id() const { return id_; }
friend class DeviceInfoCache;
friend class DeviceSet;
};
// A set of DeviceIds, represented as a bitmap.
class DeviceSet {
public:
void Insert(DeviceId device_id);
void UnionWith(const DeviceSet& other);
bool IsEmpty() const;
// Calls `func` on each DeviceId in the set. Stops iterating early if `func`
// return false.
//
// TODO(sanjoy): Change this to take a typed std::function if that's
// performance neutral.
template <typename FnTy>
void ForEach(FnTy func) const {
// This is really a poor man's iterator, we should consider writing a proper
// iterator if this ends up being used widely.
for (int word_index = 0, end = storage_.size(); word_index < end;
word_index++) {
uint64_t word = storage_[word_index];
while (word != 0) {
uint64_t only_lowest_bit_set = word & -word;
// The number of trailing zeros in a non-zero word is the index of the
// least significant 1.
int bit_index = absl::countr_zero(word);
if (!func(DeviceId(word_index * kWordSize + bit_index))) {
return;
}
word ^= only_lowest_bit_set;
}
}
}
private:
absl::InlinedVector<uint64_t, 1> storage_;
const int kWordSize = 64;
};
// Caches some miscellaneous information about TF devices. Thread compatible.
class DeviceInfoCache {
public:
bool IsGpu(DeviceId device) const { return is_gpu_[device.id()]; }
bool IsCpu(DeviceId device) const { return is_cpu_[device.id()]; }
absl::string_view GetNameFor(DeviceId device) const {
return names_[device.id()];
}
absl::StatusOr<DeviceId> GetIdFor(absl::string_view name);
using DeviceRegistration = const XlaOpRegistry::DeviceRegistration;
DeviceRegistration* GetCompilationDevice(DeviceId device) const {
return id_to_compilation_device_[device.id()];
}
absl::StatusOr<DeviceRegistration*> GetCompilationDevice(
absl::string_view name) {
TF_ASSIGN_OR_RETURN(DeviceId device_id, GetIdFor(name));
return GetCompilationDevice(device_id);
}
const DeviceType& GetDeviceTypeFor(DeviceId device) const {
return *id_to_device_type_[device.id()];
}
using DeviceTypeConstRef = std::reference_wrapper<const DeviceType>;
absl::StatusOr<DeviceTypeConstRef> GetDeviceTypeFor(
absl::string_view device_name) {
TF_ASSIGN_OR_RETURN(DeviceId device_id, GetIdFor(device_name));
return std::cref(*id_to_device_type_[device_id.id()]);
}
std::string DebugString(const DeviceSet& device_set) const;
private:
absl::flat_hash_map<std::string, DeviceId> name_to_id_;
// These fields are populated for a device in GetIdFor, *before* we give out a
// DeviceId.
std::vector<const XlaOpRegistry::DeviceRegistration*>
id_to_compilation_device_;
std::vector<std::unique_ptr<DeviceType>> id_to_device_type_;
std::vector<std::string> names_;
std::vector<bool> is_cpu_;
std::vector<bool> is_gpu_;
};
} // namespace jit
// Returns the DeviceType corresponding to 'device'.
absl::Status DeviceNameToDeviceType(const std::string& device,
DeviceType* device_type);
// Picks the device for which XLA should compile a cluster that contains
// operations placed in devices in `devices`. For instance a cluster that
// contains operations solely placed on the CPU will be compiled into a CPU
// executable by XLA, whereas a cluster that contains operations placed on the
// CPU and also operations placed on the GPU will be compiled into a GPU
// executable.
//
// Returns a non-OK Status if no unambiguous choice of device exists.
//
// We choose the device using the following rules:
//
// - It is an error for `device_names` to contain more than one device of the
// same type.
// - GPU is preferred over CPU.
// - If `allow_mixing_unknown_and_cpu` is true then unknown devices are
// preferred over CPU.
// - XLA devices count as "unrecognized devices".
//
// This set of rules above implicitly assume that XLA:GPU can compile all
// operations in the cluster that XLA:CPU can compile, and if
// `allow_mixing_unknown_and_cpu` then the unrecognized device can also compile
// all operations in the cluster that XLA:CPU can compile.
//
// We provide the `allow_mixing_unknown_and_cpu` knob so that we can do both of
// the following things:
//
// - Let MarkForCompilationPass not inject CPU-placed operations into clusters
// that will run on unknown devices (because the unknown XLA backend may not
// support every operation supported by CPU).
// - Let BuildXlaOpsPass successfully infer a compilation device for a cluster
// that contains nodes placed on both the CPU and on unknown devices. In this
// case it is the responsibility of the optimization pass that injected the
// CPU nodes into the cluster to ensure that these nodes can be compiled by
// the unknown XLA backend.
absl::StatusOr<jit::DeviceId> PickDeviceForXla(
const jit::DeviceInfoCache& device_info_cache,
const jit::DeviceSet& devices, bool allow_mixing_unknown_and_cpu);
// This is like `PickDeviceForXla` except that it returns nullopt (instead of a
// non-OK Status) if no unambiguous choice of device exists.
//
// We return a failing Status for errors unrelated to the device choice
// algorithm itself.
absl::StatusOr<std::optional<jit::DeviceId>> MaybePickDeviceForXla(
const jit::DeviceInfoCache& device_info_cache,
const jit::DeviceSet& devices, bool allow_mixing_unknown_and_cpu);
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_DEVICE_UTIL_H_
tensorflow/compiler/jit/device_util_test.cc
+144
View File
@@ -0,0 +1,144 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/jit/device_util.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace {
absl::Status PickDeviceHelper(bool allow_mixing_unknown_and_cpu,
absl::Span<const absl::string_view> device_names,
std::string* result) {
jit::DeviceInfoCache cache;
jit::DeviceSet device_set;
for (absl::string_view name : device_names) {
TF_ASSIGN_OR_RETURN(jit::DeviceId device_id, cache.GetIdFor(name));
device_set.Insert(device_id);
}
TF_ASSIGN_OR_RETURN(
jit::DeviceId result_id,
PickDeviceForXla(cache, device_set, allow_mixing_unknown_and_cpu));
*result = std::string(cache.GetNameFor(result_id));
return absl::OkStatus();
}
void CheckPickDeviceResult(absl::string_view expected_result,
bool allow_mixing_unknown_and_cpu,
absl::Span<const absl::string_view> inputs) {
std::string result;
TF_ASSERT_OK(PickDeviceHelper(allow_mixing_unknown_and_cpu, inputs, &result))
<< "inputs = [" << absl::StrJoin(inputs, ", ")
<< "], allow_mixing_unknown_and_cpu=" << allow_mixing_unknown_and_cpu
<< ", expected_result=" << expected_result;
EXPECT_EQ(result, expected_result);
}
void CheckPickDeviceHasError(bool allow_mixing_unknown_and_cpu,
absl::Span<const absl::string_view> inputs) {
std::string result;
EXPECT_FALSE(
PickDeviceHelper(allow_mixing_unknown_and_cpu, inputs, &result).ok());
}
const char* kCPU0 = "/job:localhost/replica:0/task:0/device:CPU:0";
const char* kGPU0 = "/job:localhost/replica:0/task:0/device:GPU:0";
const char* kXPU0 = "/job:localhost/replica:0/task:0/device:XPU:0";
const char* kYPU0 = "/job:localhost/replica:0/task:0/device:YPU:0";
const char* kCPU1 = "/job:localhost/replica:0/task:0/device:CPU:1";
const char* kGPU1 = "/job:localhost/replica:0/task:0/device:GPU:1";
const char* kXPU1 = "/job:localhost/replica:0/task:0/device:XPU:1";
const char* kCPU0Partial = "/device:CPU:0";
const char* kGPU0Partial = "/device:GPU:0";
const char* kXPU0Partial = "/device:XPU:0";
TEST(PickDeviceForXla, UniqueDevice) {
CheckPickDeviceResult(kGPU0, false, {kGPU0, kGPU0});
}
TEST(PickDeviceForXla, MoreSpecificDevice) {
CheckPickDeviceResult(kCPU0, false, {kCPU0, kCPU0Partial});
CheckPickDeviceResult(kGPU0, false, {kGPU0, kGPU0Partial});
// Unknown devices do not support merging of full and partial specifications.
CheckPickDeviceHasError(false, {kXPU1, kXPU0Partial});
}
TEST(PickDeviceForXla, DeviceOrder) {
CheckPickDeviceResult(kGPU0, false, {kGPU0, kCPU0});
CheckPickDeviceResult(kGPU0, false, {kCPU0, kGPU0});
CheckPickDeviceResult(kXPU0, true, {kXPU0, kCPU0});
}
TEST(PickDeviceForXla, MultipleUnknownDevices) {
CheckPickDeviceHasError(false, {kXPU0, kYPU0});
}
TEST(PickDeviceForXla, GpuAndUnknown) {
CheckPickDeviceHasError(false, {kGPU0, kXPU1});
}
TEST(PickDeviceForXla, UnknownAndCpu) {
CheckPickDeviceHasError(false, {kXPU0, kCPU1});
}
TEST(PickDeviceForXla, MultipleDevicesOfSameType) {
CheckPickDeviceHasError(true, {kCPU0, kCPU1});
CheckPickDeviceHasError(false, {kCPU0, kCPU1});
CheckPickDeviceHasError(false, {kGPU0, kGPU1});
CheckPickDeviceHasError(false, {kXPU0, kXPU1});
CheckPickDeviceHasError(false, {kCPU0, kCPU1, kGPU0});
}
void SimpleRoundTripTestForDeviceSet(int num_devices) {
jit::DeviceSet device_set;
jit::DeviceInfoCache device_info_cache;
std::vector<std::string> expected_devices, actual_devices;
for (int i = 0; i < num_devices; i++) {
std::string device_name =
absl::StrCat("/job:localhost/replica:0/task:0/device:XPU:", i);
TF_ASSERT_OK_AND_ASSIGN(jit::DeviceId device_id,
device_info_cache.GetIdFor(device_name));
device_set.Insert(device_id);
expected_devices.push_back(device_name);
}
device_set.ForEach([&](jit::DeviceId device_id) {
actual_devices.push_back(
std::string(device_info_cache.GetNameFor(device_id)));
return true;
});
EXPECT_EQ(expected_devices, actual_devices);
}
TEST(DeviceSetTest, SimpleRoundTrip_One) { SimpleRoundTripTestForDeviceSet(1); }
TEST(DeviceSetTest, SimpleRoundTrip_Small) {
SimpleRoundTripTestForDeviceSet(8);
}
TEST(DeviceSetTest, SimpleRoundTrip_Large) {
SimpleRoundTripTestForDeviceSet(800);
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc
+1334
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/jit/encapsulate_subgraphs_pass.h
+108
View File
@@ -0,0 +1,108 @@
/* 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.
==============================================================================*/
// An optimization pass that groups nodes marked with a common
// kXlaClusterAttr into functions, and replaces the original nodes by
// calls. The calls are annotated with kXlaCompiledKernelAttr.
#ifndef TENSORFLOW_COMPILER_JIT_ENCAPSULATE_SUBGRAPHS_PASS_H_
#define TENSORFLOW_COMPILER_JIT_ENCAPSULATE_SUBGRAPHS_PASS_H_
#include "tensorflow/core/common_runtime/optimization_registry.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/lib/core/status.h"
namespace tensorflow {
// EncapsulateSubgraphs pass takes all the nodes with the same cluster ID
// (derived from kXlaClusterAttr=ID (kXlaClusterAttr) attribute), puts them into
// a TF function, and replaces the subgraph in the main graph with a call to
// that TF function annotated with kXlaCompiledKernelAttr (_XlaCompiledKernel).
class EncapsulateSubgraphsPass : public GraphOptimizationPass {
public:
absl::Status Run(const GraphOptimizationPassOptions& options) override;
};
// A rewriting function to apply to each subgraph during encapsulation.
// 'arg_source_tensors' are the tensors corresponding to the arguments in the
// original source graph (*not* 'graph').
//
// 'graph' is the subgraph. The rewriting may renumber the inputs and outputs;
// 'input_permutation' is a mapping from old argument numbers to new argument
// numbers, whereas 'output_permutation' is the same for outputs. Both
// 'input_permutation' and 'output_permutation' are initialized to the identity
// permutation. 'nodedef' is the NodeDef for the call to the function under
// construction, provided to allow additional attributes to be set.
// The rewrite may also change the NodeDef's operator name, and that
// name will be used as the name of the generated function.
typedef std::function<absl::Status(
const std::vector<OutputTensor>& arg_source_tensors,
std::unique_ptr<Graph>* graph, std::vector<int>* input_permutation,
std::vector<int>* output_permutation, NodeDef* node_def)>
RewriteSubgraphFn;
// Transformation that finds subgraphs whose nodes are marked with
// 'group_attribute', splits those subgraphs into functions, and replaces
// the originals with function calls.
//
// 'group_attribute' must be a string valued-attribute that names the new
// functions to introduce.
//
// If 'rewrite_subgraph_fn' is set, it is applied to each subgraph before
// function conversion.
//
// If 'reuse_existing_functions' is set, use an existing function with the
// same name, if any.
//
// TODO(phawkins): currently, some information in control edges
// is not preserved. Suppose you have A and B in the main
// graph, C and D in a subgraph. B and C have control deps from A, D has control
// dep from B. Originally D must run after C, post-transformation this
// dependency is lost.
absl::Status EncapsulateSubgraphsInFunctions(
std::string group_attribute, const Graph& graph_in,
const RewriteSubgraphFn& rewrite_subgraph_fn, bool reuse_existing_functions,
std::unique_ptr<Graph>* graph_out, FunctionLibraryDefinition* library);
// The attribute that marks function calls produced by the encapsulate
// subgraphs pass and that should in turn be compiled via XlaLaunch operators.
extern const char* const kXlaCompiledKernelAttr;
// Does `node` have the kXlaCompiledKernelAttr attribute?
bool IsXlaCompiledKernel(const Node& node);
// Functions produced by the EncapsulateSubgraphs pass have their arguments in
// the order:
// 1) compile-time constant arguments, in host memory,
// 2) other arguments, in device memory.
// 3) resource variable arguments, in host memory. Note that only the resource
// Tensor itself is in host memory; the underlying value may be in device
// memory.
// The functions are annotated with the following attributes that describe how
// many constant and resource arguments there are:
// Name of the attribute containing the number of constant arguments.
extern const char* const kXlaNumConstantArgsAttr;
// Name of the attribute containing the number of resource variable arguments.
extern const char* const kXlaNumResourceArgsAttr;
// Name of the attribute defining whether the cluster has reference variables.
extern const char* const kXlaHasReferenceVarsAttr;
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_ENCAPSULATE_SUBGRAPHS_PASS_H_
tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc
+2788
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/jit/encapsulate_util.cc
+418
View File
@@ -0,0 +1,418 @@
/* Copyright 2018 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/jit/encapsulate_util.h"
#include <algorithm>
#include <iterator>
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/strings/str_cat.h"
#include "absl/types/optional.h"
#include "tensorflow/compiler/jit/shape_inference.h"
#include "tensorflow/compiler/tf2xla/tf2xla_util.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"
using tsl::StatusOr;
namespace tensorflow {
namespace {
// Returns string attribute value for the node if the attribute is present,
// otherwise returns empty optional value.
std::optional<std::string> GetStringAttr(const Node& n,
const std::string& attr_name) {
auto attr = n.attrs().Find(attr_name);
if (!attr) {
return std::nullopt;
} else {
return attr->s();
}
}
// Adds a value to the node's list attribute.
template <typename T>
absl::Status AppendToListAttr(Node* n, const std::string& attr_name,
const std::string& value) {
std::vector<T> attr_value;
absl::Status s = GetNodeAttr(n->attrs(), attr_name, &attr_value);
if (!s.ok() && s.code() != error::NOT_FOUND) {
return s;
}
n->ClearAttr(attr_name);
attr_value.push_back(value);
n->AddAttr(attr_name, attr_value);
return absl::OkStatus();
}
// Replaces attribute value.
template <typename T>
void ReplaceAttr(Node* n, const std::string& attr_name, const T& value) {
n->ClearAttr(attr_name);
n->AddAttr(attr_name, value);
}
// Step 1 for `PreprocessEdgesBetweenOutsideCompilations`. See comments of
// `PreprocessEdgesBetweenOutsideCompilations` for details.
absl::Status PreprocessControlEdgesBetweenOutsideCompilations(
Graph* g, const std::string& outside_compilation_attr_name) {
// Gather edges to remove. We should not remove the edge while iterating.
std::vector<const Edge*> edges_to_remove;
for (const Edge* e : g->edges()) {
if (!e->IsControlEdge()) {
continue;
}
auto src_outside_compilation =
GetStringAttr(*e->src(), outside_compilation_attr_name);
auto dst_outside_compilation =
GetStringAttr(*e->dst(), outside_compilation_attr_name);
if (src_outside_compilation && dst_outside_compilation) {
if (*src_outside_compilation != *dst_outside_compilation) {
// Case 1a: outside compilation to outside compilation control edge.
edges_to_remove.push_back(e);
TF_RETURN_IF_ERROR(AppendToListAttr<std::string>(
e->dst(), kXlaControlDependenciesWithinXlaClusterAttrName,
e->src()->name()));
}
} else if (src_outside_compilation && !dst_outside_compilation) {
// Case 1b: outside compilation to its XLA computation control edge.
ReplaceAttr(e->src(), kXlaConnectedToXlaComputationAttrName, true);
} else if (!src_outside_compilation && dst_outside_compilation) {
// Case 1b: XLA computation to outside compilation in it control edge.
ReplaceAttr(e->dst(), kXlaConnectedFromXlaComputationAttrName, true);
}
}
for (auto e : edges_to_remove) {
g->RemoveEdge(e);
}
return absl::OkStatus();
}
// Step 2 for `PreprocessEdgesBetweenOutsideCompilations`. See comments of
// `PreprocessEdgesBetweenOutsideCompilations` for details.
absl::Status PreprocessDataEdgesBetweenOutsideCompilations(
Graph* g, const std::string& outside_compilation_attr_name) {
// Gather edges between outside compilation and host computation. Notice that
// we do not store `Edge*` directly because we remove some nodes while adding
// Identity nodes, and those Edge pointers might be invalidated.
struct EdgeInfo {
int dst_input, dst_node_id;
};
std::vector<EdgeInfo> edges;
for (const Edge* e : g->edges()) {
if (e->IsControlEdge()) {
continue;
}
auto src_outside_compilation =
GetStringAttr(*e->src(), outside_compilation_attr_name);
auto dst_outside_compilation =
GetStringAttr(*e->dst(), outside_compilation_attr_name);
if (src_outside_compilation && dst_outside_compilation &&
*src_outside_compilation != *dst_outside_compilation) {
edges.push_back(EdgeInfo{e->dst_input(), e->dst()->id()});
VLOG(4) << "Oc -> oc edge: " << e->DebugString();
}
}
// Remove the edge from host to outside compilation. Add a placeholder as
// outside compilation node input.
std::map<std::pair<std::string, int>, Node*> placeholders;
for (int i = 0, end = edges.size(); i < end; i++) {
Node* dst = g->FindNodeId(edges[i].dst_node_id);
const Edge* e;
TF_RETURN_IF_ERROR(dst->input_edge(edges[i].dst_input, &e));
Node* src = e->src();
int src_output = e->src_output(), dst_input = e->dst_input();
g->RemoveEdge(e);
// Find or create placeholder node.
std::string new_name =
absl::StrCat(src->name(), "_oc_to_oc_placeholder_", src_output);
auto placeholder_index = std::make_pair(src->name(), src_output);
auto iter = placeholders.find(placeholder_index);
Node* placeholder_node;
if (iter == placeholders.end()) {
NodeDefBuilder placeholder_builder(new_name, "Placeholder");
placeholder_builder.Attr("dtype", src->output_type(src_output));
std::string outside_compilation_attr;
TF_RETURN_IF_ERROR(GetNodeAttr(dst->attrs(),
outside_compilation_attr_name,
&outside_compilation_attr));
placeholder_builder.Attr(outside_compilation_attr_name,
outside_compilation_attr);
placeholder_builder.Attr(kOutsideCompilationOriginalNodeAttrName,
src->name());
placeholder_builder.Attr(kOutsideCompilationSrcOutputAttrName,
src_output);
NodeDef placeholder_def;
TF_RETURN_IF_ERROR(placeholder_builder.Finalize(&placeholder_def));
TF_ASSIGN_OR_RETURN(placeholder_node, g->AddNode(placeholder_def));
placeholders[placeholder_index] = placeholder_node;
} else {
placeholder_node = iter->second;
}
g->AddEdge(placeholder_node, 0, dst, dst_input);
// Replace `e->dst()` because its input node changed.
NodeDef new_def = dst->def();
*new_def.mutable_input(dst_input) = placeholder_node->name();
TF_ASSIGN_OR_RETURN(Node * dst_replace_node, ReplaceNode(g, dst, new_def));
// Other edge in `edges` might have `e->dst()` as src or dst
// node. Before removing `e->dst()`, replace those edges with
// corresponding edges for `dst_replace_node`.
for (int j = i + 1, end = edges.size(); j < end; j++) {
if (edges[j].dst_node_id == edges[i].dst_node_id) {
edges[j].dst_node_id = dst_replace_node->id();
}
}
}
return absl::OkStatus();
}
// Step 1 for `PostprocessEdgesBetweenOutsideCompilations`. See comments of
// `PostprocessEdgesBetweenOutsideCompilations` for details.
absl::Status PostprocessDataEdgesBetweenOutsideCompilations(
Graph* g, const std::string& outside_compilation_attr_name) {
// Gather all outside compilation to outside compilation nodes.
std::vector<Node*> placeholder_nodes;
for (Node* n : g->nodes()) {
if (n->type_string() == "Placeholder" &&
HasNodeAttr(n->def(), kOutsideCompilationOriginalNodeAttrName)) {
placeholder_nodes.push_back(n);
}
}
// Remove the placeholder nodes, and reconnect original edge.
auto node_name_index = g->BuildNodeNameIndex();
for (auto n : placeholder_nodes) {
std::string node_name;
int node_src_output;
TF_RETURN_IF_ERROR(GetNodeAttr(
n->attrs(), kOutsideCompilationOriginalNodeAttrName, &node_name));
TF_RETURN_IF_ERROR(GetNodeAttr(
n->attrs(), kOutsideCompilationSrcOutputAttrName, &node_src_output));
auto iter = node_name_index.find(node_name);
if (iter == node_name_index.end()) {
return absl::InternalError(absl::StrCat(
"Cannot find original node for oc -> host placeholder node ",
node_name));
}
// Change all usage node to use the original node instead.
Node* original_node = iter->second;
std::vector<const Edge*> control_edges;
std::vector<OutEdgeInfo> data_edges;
for (auto e : n->out_edges()) {
if (e->IsControlEdge()) {
control_edges.push_back(e);
} else {
data_edges.push_back({e->dst(), e->src_output(), e->dst_input()});
}
}
for (const Edge* e : control_edges) {
g->AddControlEdge(original_node, e->dst());
g->RemoveEdge(e);
}
for (int i = 0, end = data_edges.size(); i < end; i++) {
Node* dst = data_edges[i].dst;
NodeDef new_def = dst->def();
int dst_input = data_edges[i].dst_input;
*new_def.mutable_input(dst_input) =
absl::StrCat(original_node->name(), ":", node_src_output);
TF_ASSIGN_OR_RETURN(Node * replace_node, ReplaceNode(g, dst, new_def));
const Edge* edge_to_replace = nullptr;
TF_RETURN_IF_ERROR(replace_node->input_edge(dst_input, &edge_to_replace));
g->RemoveEdge(edge_to_replace);
g->AddEdge(original_node, node_src_output, replace_node, dst_input);
// Other edges might have `dst` as dst node. Update those edges with
// `replace_node`.
for (int j = i + 1, end = data_edges.size(); j < end; j++) {
if (data_edges[j].dst == dst) {
data_edges[j].dst = replace_node;
}
}
// Other placeholder node might have `dst` as original node. Update
// `node_name_index` with `replace_node`.
node_name_index[replace_node->name()] = replace_node;
}
// Remove placeholder node.
g->RemoveNode(n);
}
return absl::OkStatus();
}
// Step 2 for `PostprocessEdgesBetweenOutsideCompilations`. See comments of
// `PostprocessEdgesBetweenOutsideCompilations` for details.
absl::Status PostprocessControlEdgesBetweenOutsideCompilations(
Graph* g, const std::string& outside_compilation_attr_name) {
auto node_name_index = g->BuildNodeNameIndex();
// Reconnect outside compilation to outside compilation control edge.
for (Node* n : g->nodes()) {
std::vector<std::string> control_deps;
absl::Status s =
GetNodeAttr(n->attrs(), kXlaControlDependenciesWithinXlaClusterAttrName,
&control_deps);
if (!s.ok()) {
if (s.code() != error::NOT_FOUND) {
return s;
} else {
continue;
}
} else {
n->ClearAttr(kXlaControlDependenciesWithinXlaClusterAttrName);
for (const std::string& control_input : control_deps) {
auto iter = node_name_index.find(control_input);
if (iter == node_name_index.end()) {
return absl::InternalError(
absl::StrCat("Cannot find original node for ", control_input));
}
g->AddControlEdge(iter->second, n);
}
}
}
return absl::OkStatus();
}
} // namespace
const char kXlaInferredShapesAttrName[] = "_xla_inferred_shapes";
const char kXlaConnectedToXlaComputationAttrName[] =
"_xla_connected_to_xla_computation";
const char kXlaConnectedFromXlaComputationAttrName[] =
"_xla_connected_from_xla_computation";
const char kOutsideCompilationOriginalNodeAttrName[] =
"_xla_oc_to_oc_node_name";
const char kOutsideCompilationSrcOutputAttrName[] = "_xla_oc_to_oc_src_output";
const char kXlaControlDependenciesWithinXlaClusterAttrName[] =
"_xla_control_dependencies_within_xla_cluster";
const char kXlaIsLiftedArgAttrName[] = "_xla_is_lifted_arg";
const char kXlaLiftedArgOutsideCompilationAttrName[] = "_xla_lifted_arg_oc";
const char kXlaOutsideCompilationInputsAttrName[] = "_xla_oc_inputs";
const char kXlaIsPlaceholderForArg[] = "_xla_is_placeholder_for_arg";
absl::Status PerformStaticShapeInferenceBeforeEncapsulation(Graph* g) {
// Perform shape inference.
std::map<int, InferredShape> arg_shapes;
GraphShapeInfo shape_info;
TF_RETURN_IF_ERROR(
InferShapes(g, arg_shapes, /*fnlib_def=*/nullptr, &shape_info));
// Add attribute for output shapes.
auto node_name_index = g->BuildNodeNameIndex();
for (auto iter : shape_info) {
std::vector<PartialTensorShape> output_shapes;
std::transform(iter.second.begin(), iter.second.end(),
std::back_inserter(output_shapes),
[](const InferredShape& inferred_shape) {
return inferred_shape.shape;
});
Node* n = node_name_index[iter.first];
n->AddAttr(kXlaInferredShapesAttrName, output_shapes);
}
return absl::OkStatus();
}
absl::StatusOr<
std::unique_ptr<absl::flat_hash_map<std::string, std::vector<std::string>>>>
OutsideCompilationClusterDependencies(
const Graph* g, const std::string& outside_compilation_attr_name) {
auto cluster_deps = std::make_unique<
absl::flat_hash_map<std::string, absl::flat_hash_set<std::string>>>();
for (const Edge* e : g->edges()) {
auto src_outside_compilation =
GetStringAttr(*e->src(), outside_compilation_attr_name);
auto dst_outside_compilation =
GetStringAttr(*e->dst(), outside_compilation_attr_name);
if (src_outside_compilation && dst_outside_compilation &&
*src_outside_compilation != *dst_outside_compilation) {
auto dst_deps_it = cluster_deps->find(*dst_outside_compilation);
if (dst_deps_it == cluster_deps->end()) {
cluster_deps->insert(std::make_pair(
*dst_outside_compilation,
absl::flat_hash_set<std::string>({*src_outside_compilation})));
} else {
dst_deps_it->second.insert(*src_outside_compilation);
}
}
}
auto cluster_deps_ordered = std::make_unique<
absl::flat_hash_map<std::string, std::vector<std::string>>>();
for (auto it = cluster_deps->begin(); it != cluster_deps->end(); it++) {
std::vector<std::string> ordered_deps(it->second.begin(), it->second.end());
std::sort(ordered_deps.begin(), ordered_deps.end());
cluster_deps_ordered->insert(std::make_pair(it->first, ordered_deps));
}
return std::move(cluster_deps_ordered);
}
absl::Status PreprocessEdgesBetweenOutsideCompilations(
Graph* g, const std::string& outside_compilation_attr_name) {
// Remove edges from source node to outside compilation nodes, and edges
// from outside compilation nodes to sink node.
std::vector<const Edge*> edges_to_remove;
for (const Edge* e : g->source_node()->out_edges()) {
if (HasNodeAttr(e->dst()->def(), outside_compilation_attr_name)) {
edges_to_remove.push_back(e);
}
}
for (const Edge* e : g->sink_node()->in_edges()) {
if (HasNodeAttr(e->src()->def(), outside_compilation_attr_name)) {
edges_to_remove.push_back(e);
}
}
for (auto e : edges_to_remove) {
g->RemoveEdge(e);
}
TF_RETURN_IF_ERROR(PreprocessControlEdgesBetweenOutsideCompilations(
g, outside_compilation_attr_name));
TF_RETURN_IF_ERROR(PreprocessDataEdgesBetweenOutsideCompilations(
g, outside_compilation_attr_name));
return absl::OkStatus();
}
absl::Status PostprocessEdgesBetweenOutsideCompilations(
Graph* g, const std::string& outside_compilation_attr_name) {
TF_RETURN_IF_ERROR(PostprocessDataEdgesBetweenOutsideCompilations(
g, outside_compilation_attr_name));
TF_RETURN_IF_ERROR(PostprocessControlEdgesBetweenOutsideCompilations(
g, outside_compilation_attr_name));
return absl::OkStatus();
}
} // namespace tensorflow
tensorflow/compiler/jit/encapsulate_util.h
+155
View File
@@ -0,0 +1,155 @@
/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
// This file contains some utility functions for encapsulating XLA computation
// in host graph and encapsulating outside compilation in XLA computation.
#ifndef TENSORFLOW_COMPILER_JIT_ENCAPSULATE_UTIL_H_
#define TENSORFLOW_COMPILER_JIT_ENCAPSULATE_UTIL_H_
#include "absl/container/flat_hash_map.h"
#include "tensorflow/core/graph/graph.h"
namespace tensorflow {
// Attribute marking output tensor shapes inferred by XLA. Attribute value is
// a list of PartialTensorShape objects.
extern const char kXlaInferredShapesAttrName[];
// Infers output shapes for all nodes in graph `g`. The output shapes will be
// stored in node attribute `kXlaInferredShapesAttrName`.
//
// We have to perform shape inference before encapsulation because after
// encapsulation, some nodes will be encapsulated into function call, and shape
// inference does not handle function call at the moment.
absl::Status PerformStaticShapeInferenceBeforeEncapsulation(Graph* g);
// Attribute indicating that some ops in this node's XLA computation has control
// dependency on this node. Attribute value will always be "true".
extern const char kXlaConnectedToXlaComputationAttrName[];
// Attribute indicating that this node has control dependency on some ops in
// this node's XLA computation. Attribute value will always be "true".
extern const char kXlaConnectedFromXlaComputationAttrName[];
// Attribute indicating that this is an Placeholder node added to act as a
// temporary input node for an outside compilation node. Attribute value will be
// string (original input node name).
extern const char kOutsideCompilationOriginalNodeAttrName[];
// Attribute indicating that this is an Placeholder node added to act as a
// temporary input node for an outside compilation node. Attribute value will be
// int (src_output for original edge).
extern const char kOutsideCompilationSrcOutputAttrName[];
// Attribute indicating that this node has control dependencies on some other
// nodes within the same XLA cluster. Attribute value will be a list of string
// (node names).
extern const char kXlaControlDependenciesWithinXlaClusterAttrName[];
// Attribute indicating that this node is an outside compilation node which is
// lifted out of If/While/function node. Attribute value will always be boolean
// value "true".
extern const char kXlaIsLiftedArgAttrName[];
// Attribute indicating that this node is a Placeholder node for an _Arg node
// lifted out of If/While/function node. Attribute value will be a string, which
// is the outside compilation cluster name sending the lifted arg node to host.
extern const char kXlaLiftedArgOutsideCompilationAttrName[];
// Attribute indicating that this is an IdentityN node receiving inputs for a
// outside compilation Placeholder node (the original outside compilation node
// is moved out of TPU computation, and we left a Placeholder node there).
// Attribute value will be a string, which is the outside compilation cluster
// name for the outside compilation Placeholder node.
extern const char kXlaOutsideCompilationInputsAttrName[];
// Attribute indicating that this is a Placeholder node for an _Arg node used in
// outside compilation. We should not move this node out of XLA computation.
// Attribute value will always be boolean value "true".
extern const char kXlaIsPlaceholderForArg[];
// Information for XLA computation.
struct XlaClusterInfo {
// Add an explicitly-defined default constructor for this class.
//
// The compiler may delete the default constructor here because
// host_compute_core is a const member whose type (std::map) doesn't
// necessarily have a user provided constructor -- while libc++ and
// libstdc++ 4.8 provide a user defined default constructor, libstdc++ at
// least >= 7.3 does not. See also c++11 [class.ctor] p5.
//
// TODO(klimek): In c++17 we'll be able to initialize host_compute_core
// without losing aggregate initialization, which allows us to get rid of
// the constructor definitions again.
XlaClusterInfo() {}
XlaClusterInfo(const std::string& cluster_name,
const NameAttrList& func_name_attrs, Node* node,
const std::map<std::string, int>& host_compute_core)
: cluster_name(cluster_name),
func_name_attrs(func_name_attrs),
node(node),
host_compute_core(host_compute_core) {}
// XLA cluster name. It might be different from `func_name`.
const std::string cluster_name;
// Name and attributes of XLA computation function.
const NameAttrList func_name_attrs;
// The XLA computation node in the graph.
Node* node;
// A mapping from outside compilation cluster name to its device assignment.
const std::map<std::string, int> host_compute_core;
};
// Finds dependencies between outside compilation clusters, including both data
// dependencies and control dependencies. cluster_deps maps the name name of an
// outside compilation cluster to a set of names of outside compilation clusters
// that it depends on.
absl::StatusOr<
std::unique_ptr<absl::flat_hash_map<std::string, std::vector<std::string>>>>
OutsideCompilationClusterDependencies(
const Graph* g, const std::string& outside_compilation_attr_name);
// Preprocesses edges within the same XLA cluster. It will perform the following
// operations in order:
//
// 0. Remove edges from source node to outside compilation nodes, and edges
// from outside compilation nodes to sink node.
// 1a. For edges between different outside compilation clusters, remove the edge
// and add attr "kXlaControlDependenciesWithinXlaClusterAttrName = src node
// name" to dst node.
// 1b. For control edges between outside compilation and its XLA computation,
// add attr "kXlaConnected{From, To}XlaComputationAttrName = true" to the
// outside compilation node.
// 2. For data edges between different outside compilations, remove the edge
// and create a Placeholder node as dst node's input.
absl::Status PreprocessEdgesBetweenOutsideCompilations(
Graph* g, const std::string& outside_compilation_attr_name);
// Postprocesses edges within the same XLA cluster. This function reverts what
// `PreprocessEdgesBetweenOutsideCompilations` did. It will perform the
// following operations in order:
//
// 1. Remove Placeholder nodes between different outside compilations (created
// in `PreprocessEdgesBetweenOutsideCompilations` step 2).
// 2a. Reconnect control edges between different outside compilations (marked by
// `PreprocessEdgesBetweenOutsideCompilations` step 1a).
// Notice that control edges marked by
// `PreprocessEdgesBetweenOutsideCompilations` step 1b are not handled here.
// They are handled in `RewriteOutsideCompilationSubgraphFn`.
absl::Status PostprocessEdgesBetweenOutsideCompilations(
Graph* g, const std::string& outside_compilation_attr_name);
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_ENCAPSULATE_UTIL_H_
tensorflow/compiler/jit/encapsulate_util_test.cc
+62
View File
@@ -0,0 +1,62 @@
/* Copyright 2018 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/jit/encapsulate_util.h"
#include <vector>
#include "absl/log/check.h"
#include "tensorflow/cc/framework/ops.h"
#include "tensorflow/cc/framework/scope.h"
#include "tensorflow/cc/ops/array_ops.h"
#include "tensorflow/cc/ops/const_op.h"
#include "tensorflow/cc/ops/math_ops.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
TEST(PerformStaticShapeInferenceBeforeEncapsulationTest, Basic) {
// Build the graph:
// "add" = "const_0" + "const_1"
// "identity" = "add"
tensorflow::Scope s = tensorflow::Scope::NewRootScope();
Output const_0 = ops::Const(s.WithOpName("const_0"), 1, {2});
Output const_1 = ops::Const(s.WithOpName("const_1"), 2, {2});
Output add = ops::Add(s.WithOpName("add"), const_0, const_1);
Output identity = ops::Identity(s.WithOpName("identity"), add);
Graph g(OpRegistry::Global());
CHECK_OK(s.ToGraph(&g));
CHECK_OK(PerformStaticShapeInferenceBeforeEncapsulation(&g));
// Check that "add" node now has _xla_inferred_shapes attr.
auto node_index = g.BuildNodeNameIndex();
Node *add_node = node_index["add"];
std::vector<PartialTensorShape> output_shapes;
CHECK_OK(GetNodeAttr(add_node->attrs(), kXlaInferredShapesAttrName,
&output_shapes));
EXPECT_EQ(output_shapes.size(), 1);
TensorShapeProto shape_proto;
output_shapes[0].AsProto(&shape_proto);
EXPECT_EQ(shape_proto.dim_size(), 1);
EXPECT_EQ(shape_proto.dim(0).size(), 2);
}
} // namespace tensorflow
tensorflow/compiler/jit/encapsulate_xla_computations_pass.cc
+405
View File
@@ -0,0 +1,405 @@
/* Copyright 2018 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/jit/encapsulate_xla_computations_pass.h"
#include <functional>
#include <string>
#include "absl/algorithm/container.h"
#include "absl/container/flat_hash_set.h"
#include "absl/memory/memory.h"
#include "absl/strings/ascii.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/compiler/jit/defs.h"
#include "tensorflow/compiler/jit/encapsulate_subgraphs_pass.h"
#include "tensorflow/compiler/jit/xla_cluster_util.h"
#include "xla/status_macros.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/graph/graph_node_util.h"
#include "tensorflow/core/lib/core/stringpiece.h"
#include "tensorflow/core/lib/hash/hash.h"
#include "tensorflow/core/lib/strings/proto_serialization.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/platform/fingerprint.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/util/dump_graph.h"
namespace tensorflow {
namespace {
const char* const kXlaClusterOutput = "XlaClusterOutput";
bool IsCpuGpuCompile(const Graph* graph) {
for (Node* n : graph->nodes()) {
std::string name;
// Only consider nodes being compiled.
if (!TryGetNodeAttr(n->attrs(), kXlaClusterIdAttr, &name)) continue;
// Early return for any node with a device that is not a CPU or GPU.
DeviceNameUtils::ParsedName parsed;
if (DeviceNameUtils::ParseFullName(n->requested_device(), &parsed)) {
if (parsed.type != DEVICE_CPU && parsed.type != DEVICE_GPU) {
return false;
}
}
}
return true;
}
// Checks if a graph node is marked to be a guaranteed constant.
bool is_guaranteed_constant(const Node& n) {
bool guaranteed_constant = false;
if (!TryGetNodeAttr(n.attrs(), "_is_guaranteed_constant",
&guaranteed_constant)) {
return false;
}
return guaranteed_constant;
}
// Finds the `index` of an _Arg or _Retval node.
absl::Status GetIndexAttr(const Node& n, int num_args, int* index) {
TF_RETURN_IF_ERROR(GetNodeAttr(n.attrs(), "index", index));
if (*index < 0 || *index >= num_args) {
return absl::InvalidArgumentError(
absl::StrCat("Invalid ", n.type_string(), " number ", *index));
}
return absl::OkStatus();
}
// Returns the data type of the destination of an edge.
DataType EdgeType(const Edge* edge) {
return edge->dst()->input_type(edge->dst_input());
}
// Adds the control inputs of `node` to `*deps`.
void AddControlInputs(const Node& node, absl::flat_hash_set<Node*>* deps) {
for (const Edge* edge : node.in_edges()) {
if (edge->IsControlEdge()) {
deps->insert(edge->src());
}
}
}
// Adds the control outputs of `node` to `*deps`.
void AddControlOutputs(const Node& node, absl::flat_hash_set<Node*>* deps) {
for (const Edge* edge : node.out_edges()) {
if (edge->IsControlEdge()) {
deps->insert(edge->dst());
}
}
}
// Rewrite function to be passed to EncapsulateSubgraphsInFunctions that sorts
// the arguments into the order expected by XlaLaunch computations:
// 1) arguments
// 2) resource variable arguments
// See the documentation of EncapsulateSubgraphsInFunctions for the meaning
// of the arguments.
//
// TODO(b/113166435): Ordering constraints on XlaLaunch op can be relaxed.
absl::Status RewriteSubgraph(
const std::vector<OutputTensor>& arg_source_tensors,
std::unique_ptr<Graph>* graph_ptr, std::vector<int>* input_permutation,
std::vector<int>* output_permutation, NodeDef* call_def) {
Graph* graph = graph_ptr->get();
const int num_args = input_permutation->size();
const int num_retvals = output_permutation->size();
std::vector<Node*> args;
std::vector<Node*> retvals;
args.reserve(num_args);
retvals.reserve(num_retvals);
for (Node* n : graph->nodes()) {
if (n->type_string() == "_Arg") {
// Check if this is a guaranteed constant.
if (is_guaranteed_constant(*n)) {
return absl::InvalidArgumentError(absl::StrCat(
"Guaranteed constants are not supported (", n->name(), ")"));
}
args.push_back(n);
} else if (n->type_string() == "_Retval") {
retvals.push_back(n);
}
}
if (std::find(args.begin(), args.end(), nullptr) != args.end()) {
return absl::InvalidArgumentError("Missing or non-consecutive arguments");
}
// Reorders the arguments.
std::sort(args.begin(), args.end(), [&](Node* a, Node* b) {
// Non-resources appear before resources
bool a_is_resource = (a->output_type(0) == DT_RESOURCE);
bool b_is_resource = (b->output_type(0) == DT_RESOURCE);
// Uses the name as a tiebreaker so the output is deterministic.
absl::string_view a_name(a->name());
absl::string_view b_name(b->name());
return std::tie(a_is_resource, a_name) < std::tie(b_is_resource, b_name);
});
// Sorts the retvals by name so the order is deterministic.
std::sort(retvals.begin(), retvals.end(),
[](Node* a, Node* b) { return a->name() < b->name(); });
// Computes the permutation to produce the correct argument order, and update
// the argument indices.
int variable_start_index = num_args;
for (int i = 0; i < num_args; ++i) {
int index;
TF_RETURN_IF_ERROR(GetIndexAttr(*args[i], num_args, &index));
if (args[i]->output_type(0) == DT_RESOURCE &&
variable_start_index == num_args) {
variable_start_index = i;
}
(*input_permutation)[index] = i;
args[i]->AddAttr("index", i);
}
VLOG(4) << "variable_start_index: " << variable_start_index;
// Computes the permutation to produce the correct retval order, and update
// the argument indices.
for (int i = 0; i < num_retvals; ++i) {
int index;
TF_RETURN_IF_ERROR(GetIndexAttr(*retvals[i], num_retvals, &index));
(*output_permutation)[index] = i;
retvals[i]->AddAttr("index", i);
}
AddNodeAttr(kXlaClusterIdAttr, call_def->name(), call_def);
AddNodeAttr("_variable_start_index", variable_start_index, call_def);
// Uniquify the function name by computing a fingerprint of the function.
// Nondeterminism in serialization would not lead to incorrect results, but
// may cause spurious cache misses.
TF_ASSIGN_OR_RETURN(uint64_t fingerprint, FingerprintGraph(*graph));
VLOG(1) << "Subgraph fingerprint:" << fingerprint;
call_def->set_op(absl::StrCat(call_def->op(), "_", fingerprint));
return absl::OkStatus();
}
} // namespace
/*static*/ absl::Status EncapsulateXlaComputationsPass::Encapsulate(
std::unique_ptr<Graph>* graph, FunctionLibraryDefinition* flib_def) {
// Check for undeclared outputs before Encapsulation, so we can give a better
// error message.
// TODO(phawkins): merge this with the encapsulation code to avoid the extra
// O(n) pass over the edges.
for (const Edge* e : (*graph)->edges()) {
if (!e->IsControlEdge() &&
e->src()->attrs().Find(kXlaClusterIdAttr) != nullptr &&
e->dst()->attrs().Find(kXlaClusterIdAttr) == nullptr &&
e->dst()->type_string() != kXlaClusterOutput) {
return absl::InvalidArgumentError(absl::StrCat(
"Undeclared output of XLA computation. Some common causes of this "
"error are: 1) variable initializers that depend on the XLA "
"computation; 2) gradient computations that depend on the XLA "
"computation, which can be mitigated by moving gradient computations "
"inside XLA computation. Offending edge: ",
e->src()->name(), ":", e->src_output(), " -> ", e->dst()->name(), ":",
e->dst_input()));
}
}
auto output = std::make_unique<Graph>((*graph)->op_registry());
TF_RETURN_WITH_CONTEXT_IF_ERROR(
EncapsulateSubgraphsInFunctions(
kXlaClusterIdAttr, **graph, RewriteSubgraph,
/*reuse_existing_functions=*/true, &output, flib_def),
"EncapsulateXlaComputationsPass failed");
graph->swap(output);
return absl::OkStatus();
}
/*static*/ absl::Status EncapsulateXlaComputationsPass::BuildXlaLaunchOps(
Graph* graph,
const std::function<absl::StatusOr<bool>(const Node&)>& is_xla_launch_node,
const std::function<absl::StatusOr<XlaFunctionInfo>(const Node&)>&
get_xla_function_info,
const bool add_edges_to_output_of_downstream_nodes) {
// Finds all of the XlaLaunch function calls, to avoid mutating the graph
// while iterating.
std::vector<Node*> launch_nodes;
for (Node* n : graph->nodes()) {
TF_ASSIGN_OR_RETURN(const bool is_xla_launch_node, is_xla_launch_node(*n));
if (is_xla_launch_node) launch_nodes.push_back(n);
}
// Replaces each launch function call together with its neighboring
// XlaClusterOutput nodes with a XlaLaunch node.
for (Node* launch : launch_nodes) {
TF_ASSIGN_OR_RETURN(const XlaFunctionInfo xla_function_info,
get_xla_function_info(*launch));
std::vector<const Edge*> in_edges;
TF_RETURN_IF_ERROR(launch->input_edges(&in_edges));
const int num_inputs = in_edges.size();
const int variable_start_index = xla_function_info.variable_start_index;
const int num_variables = num_inputs - variable_start_index;
const int num_args = variable_start_index;
VLOG(4) << "Launch node '" << launch->name() << "'"
<< " input edges: " << in_edges.size() << " num_args: " << num_args
<< " num_variables: " << num_variables;
std::vector<Node*> nodes_to_remove = {launch};
// Data and control inputs to the new XlaLaunch node.
std::vector<std::pair<Node*, int>> data_inputs(num_inputs);
absl::flat_hash_set<Node*> control_inputs;
DataTypeVector arg_types(num_args);
AddControlInputs(*launch, &control_inputs);
for (int i = 0; i < num_args; ++i) {
const Edge* edge = in_edges[i];
data_inputs[i] = {edge->src(), edge->src_output()};
arg_types[i] = EdgeType(edge);
}
// Appends the variable inputs.
for (int i = 0; i < num_variables; ++i) {
int pos = variable_start_index + i;
const Edge* edge = in_edges[pos];
data_inputs[pos] = {edge->src(), edge->src_output()};
}
// Outputs.
const int num_outputs = launch->output_types().size();
absl::flat_hash_set<Node*> control_outputs;
std::vector<std::vector<std::pair<Node*, int>>> data_outputs(num_outputs);
const DataTypeVector& output_types(launch->output_types());
for (const Edge* le : launch->out_edges()) {
if (le->IsControlEdge()) {
control_outputs.insert(le->dst());
} else {
TF_RET_CHECK(le->src_output() < num_outputs);
Node* output_node = le->dst();
if (add_edges_to_output_of_downstream_nodes) {
TF_RET_CHECK(output_node->type_string() == kXlaClusterOutput)
<< le->DebugString();
nodes_to_remove.push_back(output_node);
for (const Edge* oe : output_node->out_edges()) {
TF_RET_CHECK(!oe->IsControlEdge());
data_outputs[le->src_output()].push_back(
{oe->dst(), oe->dst_input()});
}
AddControlOutputs(*output_node, &control_outputs);
} else {
data_outputs[le->src_output()].push_back(
{le->dst(), le->dst_input()});
}
}
}
NodeDef def;
def.set_name(launch->name());
MergeDebugInfo(NodeDebugInfo(launch->def()), &def);
// Target the XLA CPU/GPU backends.
VLOG(2) << "Replacing with XlaLaunch";
VLOG(2) << "Device is " << launch->requested_device();
def.set_op("XlaLaunch");
def.set_device(launch->requested_device());
AddNodeAttr("Tconstants", DataTypeVector{}, &def);
AddNodeAttr("Targs", arg_types, &def);
AddNodeAttr("Nresources", num_variables, &def);
AddNodeAttr("Tresults", output_types, &def);
NameAttrList function;
function.set_name(xla_function_info.function_name);
AddNodeAttr("function", function, &def);
for (Node* node : nodes_to_remove) {
VLOG(2) << "Deleting node " << node->DebugString();
// Ensure that we do not attempt to add control edges to nodes that are
// deleted.
control_inputs.erase(node);
control_outputs.erase(node);
graph->RemoveNode(node);
}
TF_ASSIGN_OR_RETURN(Node * xla_launch, graph->AddNode(def));
for (int i = 0, end = data_inputs.size(); i < end; ++i) {
graph->AddEdge(data_inputs[i].first, data_inputs[i].second, xla_launch,
i);
}
for (Node* n : control_inputs) {
graph->AddControlEdge(n, xla_launch);
}
for (int i = 0, end = data_outputs.size(); i < end; ++i) {
for (const auto& successor : data_outputs[i]) {
graph->AddEdge(xla_launch, i, successor.first, successor.second);
}
}
for (Node* n : control_outputs) {
graph->AddControlEdge(xla_launch, n);
}
}
return absl::OkStatus();
}
/*static*/ absl::Status EncapsulateXlaComputationsPass::BuildXlaLaunchOps(
Graph* graph) {
const auto is_xla_launch_node = [](const Node& node) -> absl::StatusOr<bool> {
const std::string& name =
GetNodeAttrString(node.attrs(), kXlaClusterIdAttr);
return !name.empty();
};
const auto get_xla_function_info =
[](const Node& node) -> absl::StatusOr<XlaFunctionInfo> {
XlaFunctionInfo result;
TF_RETURN_IF_ERROR(GetNodeAttr(node.attrs(), "_variable_start_index",
&result.variable_start_index));
result.function_name = node.type_string();
return result;
};
return BuildXlaLaunchOps(graph, is_xla_launch_node, get_xla_function_info,
/*add_edges_to_output_of_downstream_nodes=*/true);
}
absl::Status EncapsulateXlaComputationsPass::Run(
const GraphOptimizationPassOptions& options) {
VLOG(1) << "EncapsulateXlaComputations(): "
<< DumpGraphToFile("encapsulate_xla_computations_before",
**options.graph, options.flib_def);
const char* additional_help =
IsCpuGpuCompile(options.graph->get())
? xla::status_macros::kPossibleAutoJitAlternative
: "";
TF_RETURN_WITH_CONTEXT_IF_ERROR(Encapsulate(options.graph, options.flib_def),
additional_help);
VLOG(1) << "EncapsulateXlaComputations() half-way: "
<< DumpGraphToFile("encapsulate_xla_computations_halfway",
**options.graph, options.flib_def);
TF_RETURN_WITH_CONTEXT_IF_ERROR(BuildXlaLaunchOps(options.graph->get()),
additional_help);
VLOG(1) << "EncapsulateXlaComputations() finished: "
<< DumpGraphToFile("encapsulate_xla_computations_after",
**options.graph, options.flib_def);
return absl::OkStatus();
}
} // namespace tensorflow
tensorflow/compiler/jit/encapsulate_xla_computations_pass.h
+85
View File
@@ -0,0 +1,85 @@
/* Copyright 2018 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.
==============================================================================*/
// Rewrites computations generated by the xla.compile() Python code into
// XlaLaunch nodes.
//
// xla.compile() does two main things:
// a) marks operators that make up an XLA computation with the attribute
// _xla_compile_id=XYZ, where XYZ is a unique key.
// b) adds XlaClusterOutput nodes to represent outputs of the computation.
// These nodes are not marked with the _xla_compile_id attribute.
#ifndef TENSORFLOW_COMPILER_JIT_ENCAPSULATE_XLA_COMPUTATIONS_PASS_H_
#define TENSORFLOW_COMPILER_JIT_ENCAPSULATE_XLA_COMPUTATIONS_PASS_H_
#include <functional>
#include <string>
#include "tensorflow/core/common_runtime/optimization_registry.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/statusor.h"
namespace tensorflow {
// Encapsulates nodes marked with the _xla_compile_id attribute into
// XlaLaunch operators.
class EncapsulateXlaComputationsPass : public GraphOptimizationPass {
public:
absl::Status Run(const GraphOptimizationPassOptions& options) override;
// The following methods are public only for unit tests.
// This pass has two stages:
// a) first, we call EncapsulateSubgraphsPass to encapsulate all nodes
// marked with the same _xla_compile_id attribute into functions. These
// functions contain the computations to be passed to XlaLaunch. During
// encapsulation, we sort the arguments into the order expected by
// XlaLaunch.
static absl::Status Encapsulate(std::unique_ptr<Graph>* graph,
FunctionLibraryDefinition* flib_def);
// b) we rewrite the function calls generated in phase (a) into XlaLaunch
// operators. We also convert the XlaClusterOutput output nodes of the
// function call into the outputs of the XlaLaunch operator.
static absl::Status BuildXlaLaunchOps(Graph* graph);
struct XlaFunctionInfo {
int variable_start_index = -1;
std::string function_name;
};
// We need to introduce this version to adapt to the output of gpu inference
// converter. The single argument overload version calls this function.
//
// When add_edges_to_output_of_downstream_nodes is true, the output edges of
// the xla_launch_node's immediate downstream nodes would be attached to the
// generated xla node. For example, if the original graph is
// StatefulPartitionedCall{_xla_compile_id=1} -> XlaClusterOutput -> NodeA
// The output graph of this function would look like the following when
// add_edges_to_output_of_downstream_nodes is true:
// XlaLaunch -> NodeA
static absl::Status BuildXlaLaunchOps(
Graph* graph,
const std::function<absl::StatusOr<bool>(const Node&)>&
is_xla_launch_node,
const std::function<absl::StatusOr<XlaFunctionInfo>(const Node&)>&
get_xla_function_info,
bool add_edges_to_output_of_downstream_nodes);
};
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_ENCAPSULATE_XLA_COMPUTATIONS_PASS_H_
tensorflow/compiler/jit/encapsulate_xla_computations_pass_test.cc
+356
View File
@@ -0,0 +1,356 @@
/* Copyright 2018 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/jit/encapsulate_xla_computations_pass.h"
#include <initializer_list>
#include <memory>
#include <string>
#include <unordered_map>
#include "absl/log/check.h"
#include "absl/status/status.h"
#include "tensorflow/cc/framework/ops.h"
#include "tensorflow/cc/framework/scope.h"
#include "tensorflow/cc/ops/array_ops.h"
#include "tensorflow/cc/ops/function_ops.h"
#include "tensorflow/cc/ops/math_ops.h"
#include "tensorflow/cc/ops/resource_variable_ops.h"
#include "tensorflow/compiler/jit/defs.h"
#include "tensorflow/compiler/jit/xla_cluster_util.h"
#include "tensorflow/compiler/tf2xla/cc/ops/xla_jit_ops.h"
#include "tensorflow/compiler/tf2xla/test_util.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/graph_to_functiondef.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/util/equal_graph_def.h"
namespace tensorflow {
static std::unique_ptr<Graph> MakeOuterGraph(
const FunctionLibraryDefinition& flib_def, const std::string& function) {
Scope scope = Scope::NewRootScope().ExitOnError();
TF_EXPECT_OK(scope.graph()->AddFunctionLibrary(flib_def.ToProto()));
auto a = ops::Placeholder(scope.WithOpName("A"), DT_INT32);
auto b = ops::Placeholder(scope.WithOpName("B"), DT_FLOAT);
auto c = ops::Placeholder(scope.WithOpName("C"), DT_INT32);
auto d = ops::Placeholder(scope.WithOpName("D"), DT_FLOAT);
auto u = ops::Placeholder(scope.WithOpName("U"), DT_RESOURCE);
auto v = ops::Placeholder(scope.WithOpName("V"), DT_RESOURCE);
auto w = ops::Placeholder(scope.WithOpName("W"), DT_RESOURCE);
NodeDef def;
CHECK_OK(NodeDefBuilder("launch0", function, &flib_def)
.Input(a.node()->name(), 0, DT_INT32)
.Input(b.node()->name(), 0, DT_FLOAT)
.Input(c.node()->name(), 0, DT_INT32)
.Input(d.node()->name(), 0, DT_FLOAT)
.Input(u.node()->name(), 0, DT_RESOURCE)
.Input(v.node()->name(), 0, DT_RESOURCE)
.Input(w.node()->name(), 0, DT_RESOURCE)
.Device("/gpu:0")
.Attr(kXlaClusterIdAttr, "launch0")
.Attr("_variable_start_index", 4)
.Finalize(&def));
absl::Status status;
Node* launch = scope.graph()->AddNode(def, &status);
CHECK_OK(status);
CHECK_OK(scope.DoShapeInference(launch));
scope.graph()->AddEdge(a.node(), 0, launch, 0);
scope.graph()->AddEdge(b.node(), 0, launch, 1);
scope.graph()->AddEdge(c.node(), 0, launch, 2);
scope.graph()->AddEdge(d.node(), 0, launch, 3);
scope.graph()->AddEdge(u.node(), 0, launch, 4);
scope.graph()->AddEdge(v.node(), 0, launch, 5);
scope.graph()->AddEdge(w.node(), 0, launch, 6);
auto out0 =
ops::XlaClusterOutput(scope.WithOpName("Out0"), Output(launch, 0));
auto out1 =
ops::XlaClusterOutput(scope.WithOpName("Out1"), Output(launch, 1));
auto out2 =
ops::XlaClusterOutput(scope.WithOpName("Out2"), Output(launch, 2));
auto out3 =
ops::XlaClusterOutput(scope.WithOpName("Out3"), Output(launch, 3));
auto consumer0_a = ops::Identity(scope.WithOpName("consumer0_a"), out0);
auto consumer0_b = ops::Identity(scope.WithOpName("consumer0_b"), out0);
auto consumer0_c = ops::Identity(scope.WithOpName("consumer0_c"), out0);
auto consumer1 = ops::Identity(scope.WithOpName("consumer1"), out1);
auto consumer2 = ops::Identity(scope.WithOpName("consumer2"), out2);
auto consumer3 = ops::Identity(scope.WithOpName("consumer3"), out3);
std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
CHECK_OK(scope.ToGraph(graph.get()));
return graph;
}
// Makes an encapsulate body graph for use in tests.
static std::unique_ptr<Graph> MakeBodyGraph() {
Scope scope = Scope::NewRootScope().ExitOnError();
auto arg0 = ops::_Arg(scope.WithOpName("a_0_arg"), DT_INT32, 0);
auto arg1 = ops::_Arg(scope.WithOpName("b_0_arg"), DT_FLOAT, 1);
auto arg2 = ops::_Arg(scope.WithOpName("c_0_arg"), DT_INT32, 2);
auto arg3 = ops::_Arg(scope.WithOpName("d_0_arg"), DT_FLOAT, 3);
auto arg4 = ops::_Arg(scope.WithOpName("u_0_arg"), DT_RESOURCE, 4);
auto arg5 = ops::_Arg(scope.WithOpName("v_0_arg"), DT_RESOURCE, 5);
auto arg6 = ops::_Arg(scope.WithOpName("w_0_arg"), DT_RESOURCE, 6);
auto add_attrs = [](Node* node) {
node->AddAttr(kXlaClusterIdAttr, "launch0");
node->set_requested_device("/gpu:0");
};
auto b_identity = ops::Identity(scope.WithOpName("B_identity"), arg1);
add_attrs(b_identity.node());
auto read_u = ops::ReadVariableOp(scope.WithOpName("ReadU"), arg4, DT_FLOAT);
add_attrs(read_u.node());
auto read_v = ops::ReadVariableOp(scope.WithOpName("ReadV"), arg5, DT_FLOAT);
add_attrs(read_v.node());
auto read_w = ops::ReadVariableOp(scope.WithOpName("ReadW"), arg6, DT_FLOAT);
add_attrs(read_w.node());
auto e = ops::Add(scope.WithOpName("E"), arg0, arg2);
add_attrs(e.node());
auto f = ops::Add(scope.WithOpName("F"), read_v, read_w);
add_attrs(f.node());
auto g = ops::Add(scope.WithOpName("G"), f, arg3);
add_attrs(g.node());
auto out0 = ops::_Retval(scope.WithOpName("b_identity_0_retval_RetVal"),
b_identity, 0);
auto out1 = ops::_Retval(scope.WithOpName("e_0_retval_RetVal"), e, 1);
auto out2 = ops::_Retval(scope.WithOpName("g_0_retval_RetVal"), g, 2);
auto out3 =
ops::_Retval(scope.WithOpName("readu_0_retval_RetVal"), read_u, 3);
std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
CHECK_OK(scope.ToGraph(graph.get()));
return graph;
}
TEST(EncapsulateXlaComputations, DeterministicEncapsulate) {
// Test that control edge insertion order doesn't affect the cache key
// (cluster name) generated by TPU encapsulate pass.
auto get_serialized_graph = [](bool control_input_reversed,
bool operand_reversed) -> std::string {
FunctionLibraryDefinition flib_def(OpRegistry::Global(),
FunctionDefLibrary());
std::unique_ptr<Graph> graph(new Graph(&flib_def));
{
Scope scope = Scope::NewRootScope().ExitOnError();
auto a0 = ops::Placeholder(scope.WithOpName("A0"), DT_INT32);
auto a1 = ops::Placeholder(scope.WithOpName("A1"), DT_INT32);
ops::Add e = operand_reversed ? ops::Add(scope.WithOpName("E"), a0, a1)
: ops::Add(scope.WithOpName("E"), a1, a0);
auto add_attrs = [](Node* node) {
node->AddAttr(kXlaClusterIdAttr, "launch0");
};
add_attrs(e.node());
CHECK_OK(scope.ToGraph(graph.get()));
auto get_node_in_graph = [&graph](Node* node) {
return graph->FindNodeId(node->id());
};
// Insert control edge in different order. The order should not affect
// the encapsulated or serialized graph.
if (!control_input_reversed) {
graph->AddControlEdge(get_node_in_graph(a0.node()),
get_node_in_graph(e.node()), true);
graph->AddControlEdge(get_node_in_graph(a1.node()),
get_node_in_graph(e.node()), true);
} else {
graph->AddControlEdge(get_node_in_graph(a1.node()),
get_node_in_graph(e.node()), true);
graph->AddControlEdge(get_node_in_graph(a0.node()),
get_node_in_graph(e.node()), true);
}
}
CHECK_OK(EncapsulateXlaComputationsPass::Encapsulate(&graph, &flib_def));
return SerializeGraphDeterministic(*graph).value();
};
// Changing the order of control input shouldn't affect the graph generated.
EXPECT_EQ(get_serialized_graph(/*control_input_reversed=*/true,
/*operand_reversed=*/false),
get_serialized_graph(/*control_input_reversed=*/false,
/*operand_reversed=*/false));
// Changing the order of data input should affect the graph generated.
EXPECT_NE(get_serialized_graph(/*control_input_reversed=*/false,
/*operand_reversed=*/true),
get_serialized_graph(/*control_input_reversed=*/false,
/*operand_reversed=*/false));
}
TEST(EncapsulateXlaComputations, Encapsulate) {
FunctionLibraryDefinition flib_def(OpRegistry::Global(),
FunctionDefLibrary());
std::unique_ptr<Graph> graph(new Graph(&flib_def));
{
Scope scope = Scope::NewRootScope().ExitOnError();
auto a = ops::Placeholder(scope.WithOpName("A"), DT_INT32);
auto b = ops::Placeholder(scope.WithOpName("B"), DT_FLOAT);
auto c = ops::Placeholder(scope.WithOpName("C"), DT_INT32);
auto d = ops::Placeholder(scope.WithOpName("D"), DT_FLOAT);
auto u = ops::Placeholder(scope.WithOpName("U"), DT_RESOURCE);
auto v = ops::Placeholder(scope.WithOpName("V"), DT_RESOURCE);
auto w = ops::Placeholder(scope.WithOpName("W"), DT_RESOURCE);
auto add_attrs = [](Node* node) {
node->AddAttr(kXlaClusterIdAttr, "launch0");
node->set_requested_device("/gpu:0");
};
auto b_identity = ops::Identity(scope.WithOpName("B_identity"), b);
add_attrs(b_identity.node());
auto read_u = ops::ReadVariableOp(scope.WithOpName("ReadU"), u, DT_FLOAT);
add_attrs(read_u.node());
auto read_v = ops::ReadVariableOp(scope.WithOpName("ReadV"), v, DT_FLOAT);
add_attrs(read_v.node());
auto read_w = ops::ReadVariableOp(scope.WithOpName("ReadW"), w, DT_FLOAT);
add_attrs(read_w.node());
auto e = ops::Add(scope.WithOpName("E"), a, c);
add_attrs(e.node());
auto f = ops::Add(scope.WithOpName("F"), read_v, read_w);
add_attrs(f.node());
auto g = ops::Add(scope.WithOpName("G"), f, d);
add_attrs(g.node());
auto out0 = ops::XlaClusterOutput(scope.WithOpName("Out0"), b_identity);
auto out1 = ops::XlaClusterOutput(scope.WithOpName("Out1"), e);
auto out2 = ops::XlaClusterOutput(scope.WithOpName("Out2"), g);
auto out3 = ops::XlaClusterOutput(scope.WithOpName("Out3"), read_u);
auto consumer0_a = ops::Identity(scope.WithOpName("consumer0_a"), out0);
auto consumer0_b = ops::Identity(scope.WithOpName("consumer0_b"), out0);
auto consumer0_c = ops::Identity(scope.WithOpName("consumer0_c"), out0);
auto consumer1 = ops::Identity(scope.WithOpName("consumer1"), out1);
auto consumer2 = ops::Identity(scope.WithOpName("consumer2"), out2);
auto consumer3 = ops::Identity(scope.WithOpName("consumer3"), out3);
TF_ASSERT_OK(scope.ToGraph(graph.get()));
}
std::unique_ptr<Graph> graph_copy(new Graph(&flib_def));
CopyGraph(*graph, graph_copy.get());
TF_ASSERT_OK(EncapsulateXlaComputationsPass::Encapsulate(&graph, &flib_def));
std::unordered_map<std::string, Node*> index = graph->BuildNodeNameIndex();
std::string function = index.at("launch0")->type_string();
// Tests the outer graph is as expected.
{
std::unique_ptr<Graph> outer = MakeOuterGraph(flib_def, function);
GraphDef expected_def;
outer->ToGraphDef(&expected_def);
GraphDef actual_def;
graph->ToGraphDef(&actual_def);
TF_EXPECT_GRAPH_EQ_INTERNAL(expected_def, actual_def);
}
// Tests the encapsulated body graph is as expected.
{
std::unique_ptr<Graph> body = MakeBodyGraph();
GraphDef expected_body_def;
body->ToGraphDef(&expected_body_def);
InstantiationResultForTest result;
TF_EXPECT_OK(InstantiateFunctionForTest(function, flib_def, &result));
EXPECT_EQ((DataTypeVector{DT_INT32, DT_FLOAT, DT_INT32, DT_FLOAT,
DT_RESOURCE, DT_RESOURCE, DT_RESOURCE}),
result.arg_types);
EXPECT_EQ((DataTypeVector{DT_FLOAT, DT_INT32, DT_FLOAT, DT_FLOAT}),
result.ret_types);
TF_EXPECT_GRAPH_EQ(expected_body_def, result.gdef);
}
// Encapsulates the same computation again, verifies we reuse the same
// function. Encapsulation should be deterministic to avoid recompilation.
TF_ASSERT_OK(
EncapsulateXlaComputationsPass::Encapsulate(&graph_copy, &flib_def));
std::unordered_map<std::string, Node*> index_copy =
graph_copy->BuildNodeNameIndex();
std::string function_copy = index_copy.at("launch0")->type_string();
EXPECT_EQ(function, function_copy);
}
TEST(EncapsulateXlaComputations, BuildXlaLaunchOp) {
std::unique_ptr<Graph> body_graph = MakeBodyGraph();
FunctionDefLibrary flib;
TF_ASSERT_OK(GraphToFunctionDef(*body_graph, "launch0", flib.add_function()));
FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
std::unique_ptr<Graph> graph = MakeOuterGraph(flib_def, "launch0");
TF_ASSERT_OK(EncapsulateXlaComputationsPass::BuildXlaLaunchOps(graph.get()));
Scope scope = Scope::DisabledShapeInferenceScope().ExitOnError();
TF_EXPECT_OK(scope.graph()->AddFunctionLibrary(flib));
auto a = ops::Placeholder(scope.WithOpName("A"), DT_INT32);
auto b = ops::Placeholder(scope.WithOpName("B"), DT_FLOAT);
auto c = ops::Placeholder(scope.WithOpName("C"), DT_INT32);
auto d = ops::Placeholder(scope.WithOpName("D"), DT_FLOAT);
auto u = ops::Placeholder(scope.WithOpName("U"), DT_RESOURCE);
auto v = ops::Placeholder(scope.WithOpName("V"), DT_RESOURCE);
auto w = ops::Placeholder(scope.WithOpName("W"), DT_RESOURCE);
NameAttrList function;
function.set_name("launch0");
auto launch = ops::XlaLaunch(
scope.WithOpName("launch0").WithDevice("/gpu:0"),
std::initializer_list<Input>{}, std::initializer_list<Input>{a, b, c, d},
std::initializer_list<Input>{u, v, w},
DataTypeVector{DT_FLOAT, DT_INT32, DT_FLOAT, DT_FLOAT}, function);
auto consumer0_a =
ops::Identity(scope.WithOpName("consumer0_a"), launch.results[0]);
auto consumer0_b =
ops::Identity(scope.WithOpName("consumer0_b"), launch.results[0]);
auto consumer0_c =
ops::Identity(scope.WithOpName("consumer0_c"), launch.results[0]);
auto consumer1 =
ops::Identity(scope.WithOpName("consumer1"), launch.results[1]);
auto consumer2 =
ops::Identity(scope.WithOpName("consumer2"), launch.results[2]);
auto consumer3 =
ops::Identity(scope.WithOpName("consumer3"), launch.results[3]);
GraphDef expected_def;
TF_ASSERT_OK(scope.ToGraphDef(&expected_def));
GraphDef actual_def;
graph->ToGraphDef(&actual_def);
TF_EXPECT_GRAPH_EQ(expected_def, actual_def);
}
} // namespace tensorflow
tensorflow/compiler/jit/extract_outside_compilation_pass.cc
+2588
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/jit/extract_outside_compilation_pass.h
+113
View File
@@ -0,0 +1,113 @@
/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_JIT_EXTRACT_OUTSIDE_COMPILATION_PASS_H_
#define TENSORFLOW_COMPILER_JIT_EXTRACT_OUTSIDE_COMPILATION_PASS_H_
#include "absl/types/optional.h"
#include "tensorflow/compiler/jit/encapsulate_util.h"
#include "xla/status_macros.h"
#include "tensorflow/core/graph/graph.h"
namespace tensorflow {
// Rewrite function for outside compilation subgraphs. It will perform the
// following steps:
//
// 1. Add a XLA computation key placeholder node (it will be used as input for
// XlaRecvAtHost and XlaSendFromHost);
// 2. Replace all _Arg nodes with one single XlaRecvAtHost node;
// 3. Replace all _Retval nodes with one single XlaSendFromHost node;
// 4. Mark all nodes except key placeholder with attr `xla_cluster_attr_name`
// and `outside_compilation_attr_name`;
// 5. For nodes marked with attr kXlaConnectedToXlaComputationAttrName, add a
// control edge from the node to XlaSendFromHost; for nodes marked with attr
// kXlaConnectedFromXlaComputationAttrName, add a control edge from
// XlaRecvAtHost node to the node;
// 6. Try pruning XlaRecvAtHost/XlaSendFromHost/key placeholder node.
// 7. Add necessary attributes to `node_def`, so we can replace it with a
// XlaHostCompute node later. If all input shapes for XlaSendFromHost are
// known, "shapes" attr will be set to the list of input shapes; otherwise
// "shape_inference_graph" attr will be set to shape inference function name.
class RewriteOutsideCompilationSubgraphFn {
public:
RewriteOutsideCompilationSubgraphFn(
const std::string& xla_cluster_attr_name,
const std::string& outside_compilation_attr_name,
const std::string& xla_cluster_name, const std::string& new_function_name)
: xla_cluster_attr_name_(xla_cluster_attr_name),
outside_compilation_attr_name_(outside_compilation_attr_name),
xla_cluster_name_(xla_cluster_name),
new_function_name_(new_function_name) {}
absl::Status operator()(const std::vector<OutputTensor>&,
std::unique_ptr<Graph>* graph,
std::vector<int>* input_permutation,
std::vector<int>* output_permutation,
NodeDef* node_def);
private:
std::string xla_cluster_attr_name_;
std::string outside_compilation_attr_name_;
std::string xla_cluster_name_;
std::string new_function_name_;
};
// For an XLA computation function, replace all outside compilations with
// XlaHostCompute nodes. Each outside compilation subgraph will be rewritten by
// `RewriteOutsideCompilationSubgraphFn`, and they will be merged into one
// single host side graph (`host_graph`).
//
// xla_cluster_attr_name and outside_compilation_attr_name: attr name for XLA
// computation and outside compilation. Required for
// `RewriteOutsideCompilationSubgraphFn`.
// xla_cluster_name: XLA cluster name for this XLA computation. We need it
// because XLA cluster name might be different from `func_name`.
// func_name_attrs: they will be used to instantiate the XLA computation func.
// new_func_name: new function name for rewritten XLA computation func.
// host_compute_core: mapping from outside compilation cluster name to XLA
// device assignment.
// fld: FunctionLibraryDefinition object.
// host_graph: Graph object to store host side graph for all outside
// compilations within this XLA computation func. If there is no outside
// compilation, it will be empty.
// shape_inference_graphs: a list of outside compilation shape inference
// function names. These functions need to be rewritten later.
// has_outside_compilation: a bool indicating whether this function has any
// outside compilation nodes.
absl::Status ExtractOutsideCompilationForFunction(
const std::string& xla_cluster_attr_name,
const std::string& outside_compilation_attr_name,
const std::string& xla_cluster_name, const NameAttrList& func_name_attrs,
const std::string& new_func_name, const std::string& host_graph_func_name,
const std::map<std::string, int>& host_compute_core,
FunctionLibraryRuntime* flr, FunctionLibraryDefinition* fld,
std::vector<std::string>* shape_inference_graphs,
bool* has_outside_compilation);
// Rewrites XLA computation in `clusters` to replace outside compilation nodes
// with XlaHostCompute, and moves those outside compilations into `g`. If shapes
// of outside compilation outputs cannot be determined now, we will store shape
// inference graph into `fld`.
absl::Status ExtractOutsideCompilation(
const std::string& xla_cluster_attr_name,
const std::string& outside_compilation_attr_name,
const std::unordered_map<std::string, XlaClusterInfo>& clusters, Graph* g,
FunctionLibraryRuntime* flr, FunctionLibraryDefinition* fld,
bool* modified);
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_EXTRACT_OUTSIDE_COMPILATION_PASS_H_
tensorflow/compiler/jit/extract_outside_compilation_pass_test.cc
+1025
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/jit/flags.cc
+557
View File
@@ -0,0 +1,557 @@
/* Copyright 2018 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/jit/flags.h"
#include <limits>
#include <mutex> // NOLINT
#include <optional>
#include <vector>
#include "absl/base/call_once.h"
#include "absl/strings/numbers.h"
#include "absl/strings/str_split.h"
#include "absl/strings/strip.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/dump_graph.h"
#include "xla/parse_flags_from_env.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/tpu/kernels/sparse_core_xla_flags_defaults.h"
#include "tensorflow/core/util/command_line_flags.h"
namespace tensorflow {
namespace {
BuildXlaOpsPassFlags* build_ops_flags;
MarkForCompilationPassFlags* mark_for_compilation_flags;
XlaDeviceFlags* device_flags;
XlaSparseCoreFlags* sparse_core_flags;
XlaOpsCommonFlags* ops_flags;
XlaCallModuleFlags* call_module_flags;
MlirCommonFlags* mlir_flags;
JitRtFlags* jitrt_flags;
std::vector<Flag>* jitrt_flag_list;
std::vector<Flag>* flag_list;
absl::once_flag flags_init;
bool SetterForXlaAutoJitFlag(const std::string& value) {
int32_t opt_level;
// We need to use the mark_for_compilation_flags directly here instead of
// going via GetMarkForCompilationPassFlags() to avoid infinite recursion. The
// latter will try to setup and parse flags, which would bring us back to this
// setter.
if (absl::SimpleAtoi(value, &opt_level)) {
mark_for_compilation_flags->xla_auto_jit_flag
.optimization_level_single_gpu = opt_level;
mark_for_compilation_flags->xla_auto_jit_flag.optimization_level_general =
opt_level;
return true;
}
if (value == "fusible") {
mark_for_compilation_flags->xla_auto_jit_flag
.optimization_level_single_gpu = 1;
mark_for_compilation_flags->xla_auto_jit_flag.optimization_level_general =
1;
mark_for_compilation_flags->tf_xla_ops_to_cluster = "FUSIBLE";
return true;
}
absl::string_view value_sv(value);
if (!absl::ConsumePrefix(&value_sv, "single-gpu(") ||
!absl::ConsumeSuffix(&value_sv, ")") ||
!absl::SimpleAtoi(value_sv, &opt_level)) {
return false;
}
mark_for_compilation_flags->xla_auto_jit_flag.optimization_level_single_gpu =
opt_level;
return true;
}
bool SetterForXlaCallModuleDisabledChecks(const std::string& value) {
auto directives = absl::StrSplit(value, ',', absl::SkipEmpty());
call_module_flags->disabled_checks.insert(directives.begin(),
directives.end());
return true;
}
void AppendMarkForCompilationPassFlagsInternal(std::vector<Flag>* flag_list) {
std::vector<Flag> new_flags = {
Flag("tf_xla_auto_jit", SetterForXlaAutoJitFlag, "0",
"Control compilation of operators into XLA computations on CPU and "
"GPU devices. 0 = use ConfigProto setting; -1 = off; 1 = on for "
"things very likely to be improved; 2 = on for everything; "
"(experimental) fusible = only for Tensorflow operations that XLA "
"knows how to fuse. "
"If set to single-gpu(<N>) then this resolves to <N> for single-GPU "
"graphs (graphs that have at least one node placed on a GPU and no "
"more than one GPU is in use through the entire graph) and 0 "
"otherwise. Experimental."),
Flag("tf_xla_min_cluster_size",
&mark_for_compilation_flags->tf_xla_min_cluster_size,
"Minimum number of operators in an XLA compilation. Ignored for "
"operators placed on an XLA device or operators explicitly marked "
"for compilation."),
Flag("tf_xla_max_cluster_size",
&mark_for_compilation_flags->tf_xla_max_cluster_size,
"Maximum number of operators in an XLA compilation."),
Flag(
"tf_xla_ops_to_cluster",
&mark_for_compilation_flags->tf_xla_ops_to_cluster,
"(experimental) "
"Limit the operations clustered by XLA to these operations. "
"If multiple, separate them with commas. Shortcuts: "
" PW: All point-wise operations."
" RED: All reduction operations."
" MISC: Mixed operations."
" PWRED: TF operations that get converted to PW+RED operation in XLA."
" REDUCEWINDOW: TF operations like MaxPool/AvgPool that get "
"converted to ReduceWindow in XLA."
" REDUCEWINDOWPW: Operation that get converted to ReduceWindow + PW "
"(LRN, LRNGrad)."
" BN: TF FusedBatchNorm* operations."
" FUSIBLE: All TF operations that XLA can fuse (All the above). "
"You can also put any TF operation name, e.g. 'FUSIBLE,MatMul'."),
Flag("tf_xla_cluster_exclude_ops",
&mark_for_compilation_flags->tf_xla_cluster_exclude_ops,
"(experimental) "
"Exclude the operations from auto-clustering. "
"If multiple, separate them with commas."
" Where, Some_other_ops"),
Flag("tf_xla_clustering_debug",
&mark_for_compilation_flags->tf_xla_clustering_debug,
"Dump graphs during XLA compilation."),
Flag("tf_xla_cpu_global_jit",
&mark_for_compilation_flags->tf_xla_cpu_global_jit,
"Enables global JIT compilation for CPU via SessionOptions."),
Flag("tf_xla_clustering_fuel",
&mark_for_compilation_flags->tf_xla_clustering_fuel,
"Places an artificial limit on the number of ops marked as "
"eligible for clustering."),
Flag("tf_xla_disable_deadness_safety_checks_for_debugging",
&mark_for_compilation_flags
->tf_xla_disable_deadness_safety_checks_for_debugging,
"Disable deadness related safety checks when clustering (this is "
"unsound)."),
Flag("tf_xla_disable_resource_variable_safety_checks_for_debugging",
&mark_for_compilation_flags
->tf_xla_disable_resource_variable_safety_checks_for_debugging,
"Disable resource variables related safety checks when clustering "
"(this is unsound)."),
Flag("tf_xla_deterministic_cluster_names",
&mark_for_compilation_flags->tf_xla_deterministic_cluster_names,
"Causes the function names assigned by auto clustering to be "
"deterministic from run to run."),
Flag("tf_xla_persistent_cache_directory",
&mark_for_compilation_flags->tf_xla_persistent_cache_directory,
"If non-empty, JIT-compiled executables are saved to and loaded "
"from the specified file system directory path. Empty by default."),
Flag("tf_xla_persistent_cache_device_types",
&mark_for_compilation_flags->tf_xla_persistent_cache_device_types,
"If non-empty, the persistent cache will only be used for the "
"specified devices (comma separated). Each device type should be "
"able to be converted to `DeviceType`."),
Flag("tf_xla_persistent_cache_read_only",
&mark_for_compilation_flags->tf_xla_persistent_cache_read_only,
"If true, the persistent cache will be read-only."),
Flag("tf_xla_disable_strict_signature_checks",
&mark_for_compilation_flags->tf_xla_disable_strict_signature_checks,
"If true, entires loaded into the XLA compile cache will not have "
"their signatures checked strictly. Defaults to false."),
Flag("tf_xla_persistent_cache_prefix",
&mark_for_compilation_flags->tf_xla_persistent_cache_prefix,
"Specifies the persistance cache prefix. Default is "
"\"xla_compile_cache\""),
Flag("tf_xla_sparse_core_disable_table_stacking",
&sparse_core_flags->tf_xla_sparse_core_disable_table_stacking,
"Disable table stacking for all the tables passed to the SparseCore"
"mid level API."),
Flag("tf_xla_sparse_core_minibatch_max_division_level",
&sparse_core_flags->tf_xla_sparse_core_minibatch_max_division_level,
"Max level of division to split input data into minibatches."),
Flag("tf_xla_sparse_core_stacking_mem_limit_bytes",
&sparse_core_flags->tf_xla_sparse_core_stacking_mem_limit_bytes,
"If non-zero, limits the size of the activations for a given table"
"to be below these many bytes."),
Flag("tf_xla_sparse_core_stacking_table_shard_limit_bytes",
&sparse_core_flags
->tf_xla_sparse_core_stacking_table_shard_limit_bytes,
"If non-zero, limits the size of any table shard to be below these"
"many bytes.")};
flag_list->insert(flag_list->end(), new_flags.begin(), new_flags.end());
}
void AllocateAndParseJitRtFlags() {
jitrt_flags = new JitRtFlags;
jitrt_flags->always_specialize = false;
jitrt_flags->cost_driven_async_parallel_for = false;
jitrt_flags->enable_crash_reproducer = false;
jitrt_flags->log_query_of_death = false;
jitrt_flags->vectorize = false;
jitrt_flag_list = new std::vector<Flag>({
Flag("always_specialize", &jitrt_flags->always_specialize, ""),
Flag("cost_driven_async_parallel_for",
&jitrt_flags->cost_driven_async_parallel_for, ""),
Flag("enable_crash_reproducer", &jitrt_flags->enable_crash_reproducer,
""),
Flag("log_query_of_death", &jitrt_flags->log_query_of_death, ""),
Flag("vectorize", &jitrt_flags->vectorize, ""),
});
xla::ParseFlagsFromEnvAndDieIfUnknown("TF_JITRT_FLAGS", *jitrt_flag_list);
}
void AllocateAndParseFlags() {
build_ops_flags = new BuildXlaOpsPassFlags;
build_ops_flags->tf_xla_enable_lazy_compilation = true;
build_ops_flags->tf_xla_print_cluster_outputs = false;
build_ops_flags->tf_xla_check_cluster_input_numerics = false;
build_ops_flags->tf_xla_check_cluster_output_numerics = false;
build_ops_flags->tf_xla_disable_constant_folding = false;
build_ops_flags->tf_xla_disable_full_embedding_pipelining = false;
build_ops_flags->tf_xla_disable_full_embedding_pipelining_with_summaries =
true;
build_ops_flags->tf_xla_embedding_parallel_iterations = 0;
mark_for_compilation_flags = new MarkForCompilationPassFlags;
mark_for_compilation_flags->xla_auto_jit_flag.optimization_level_single_gpu =
0;
mark_for_compilation_flags->xla_auto_jit_flag.optimization_level_general = 0;
mark_for_compilation_flags->tf_xla_min_cluster_size = 4;
mark_for_compilation_flags->tf_xla_max_cluster_size =
std::numeric_limits<int32_t>::max();
mark_for_compilation_flags->tf_xla_clustering_debug = false;
mark_for_compilation_flags->tf_xla_cpu_global_jit = false;
mark_for_compilation_flags->tf_xla_clustering_fuel =
std::numeric_limits<int64_t>::max();
mark_for_compilation_flags
->tf_xla_disable_deadness_safety_checks_for_debugging = false;
mark_for_compilation_flags
->tf_xla_disable_resource_variable_safety_checks_for_debugging = false;
mark_for_compilation_flags->tf_xla_deterministic_cluster_names = false;
mark_for_compilation_flags->tf_xla_persistent_cache_directory = "";
mark_for_compilation_flags->tf_xla_persistent_cache_device_types = "";
mark_for_compilation_flags->tf_xla_persistent_cache_read_only = false;
mark_for_compilation_flags->tf_xla_disable_strict_signature_checks = false;
mark_for_compilation_flags->tf_xla_persistent_cache_prefix =
"xla_compile_cache";
device_flags = new XlaDeviceFlags;
device_flags->tf_xla_compile_on_demand = false;
device_flags->tf_xla_enable_xla_devices = false;
sparse_core_flags = new XlaSparseCoreFlags;
sparse_core_flags->tf_xla_sparse_core_minibatch_max_division_level =
kDefaultSparseCoreMinibatchMaxDivisionLevel;
sparse_core_flags->tf_xla_sparse_core_disable_table_stacking =
kDefaultDisableTableStacking;
sparse_core_flags->tf_xla_sparse_core_stacking_mem_limit_bytes =
kDefaultXlaSparseCoreStackingMemLimit;
sparse_core_flags->tf_xla_sparse_core_stacking_table_shard_limit_bytes =
kDefaultXlaSparseCoreStackingTableShardLimit;
ops_flags = new XlaOpsCommonFlags;
ops_flags->tf_xla_always_defer_compilation = false;
ops_flags->tf_xla_async_compilation = false;
ops_flags->tf_xla_use_device_api.enabled_for_xla_launch_ = true;
ops_flags->tf_xla_use_device_api.enabled_for_compile_on_demand_ = true;
ops_flags->tf_xla_use_device_api.enabled_for_compile_and_run_ = true;
ops_flags->tf_xla_use_device_api.enabled_for_all_ = false;
ops_flags->tf_xla_use_device_api.enabled_for_gpu_ = true;
call_module_flags = new XlaCallModuleFlags;
// The `enable_mlir_bridge` flag allows the user to explicitly request that
// their program is (or isn't) compiled using the MLIR-based TF-to-XLA bridge.
//
// The `enable_mlir_bridge_is_explicit` variable tracks whether or not the
// user has made an explicit request. That is, if this variable is set to
// true, the program honors the user's request as per `enable_mlir_bridge`; if
// it's set to false, the default behavior is used (which may run either
// bridge, on a per-graph basis).
bool enable_mlir_bridge = false;
bool enable_mlir_bridge_is_explicit = false;
bool enable_mlir_merge_control_flow_pass = true;
bool enable_mlir_convert_control_to_data_outputs_pass = false;
bool enable_mlir_composite_tpuexecute_side_effects = false;
bool enable_mlir_strict_clusters = false;
bool enable_mlir_multiple_local_cpu_devices = false;
bool enable_mlir_debug_info_serialization = true;
// Dump graphs in TFG dialect.
bool use_tfg_graph_dumper = false;
bool enable_tpu_variable_runtime_reformatting_pass = true;
bool enable_serialize_mlir_to_compressed_bytecode = false;
flag_list = new std::vector<Flag>(
{Flag("tf_xla_enable_lazy_compilation",
&build_ops_flags->tf_xla_enable_lazy_compilation, ""),
Flag("tf_xla_print_cluster_outputs",
&build_ops_flags->tf_xla_print_cluster_outputs,
"If true then insert Print nodes to print out values produced by "
"XLA clusters."),
Flag("tf_xla_check_cluster_input_numerics",
&build_ops_flags->tf_xla_check_cluster_input_numerics,
"If true then insert CheckNumerics nodes to check all cluster "
"inputs."),
Flag("tf_xla_check_cluster_output_numerics",
&build_ops_flags->tf_xla_check_cluster_output_numerics,
"If true then insert CheckNumerics nodes to check all cluster "
"outputs."),
Flag("tf_xla_disable_constant_folding",
&build_ops_flags->tf_xla_disable_constant_folding,
"If true then disables constant folding on TF graph before XLA "
"compilation."),
Flag("tf_xla_disable_full_embedding_pipelining",
&build_ops_flags->tf_xla_disable_full_embedding_pipelining,
"If true then disables full embedding pipelining and instead use "
"strict SparseCore / TensorCore sequencing."),
Flag("tf_xla_disable_full_embedding_pipelining_with_summaries",
&build_ops_flags
->tf_xla_disable_full_embedding_pipelining_with_summaries,
"If true then disables full embedding pipelining when summary ops "
"are detected."),
Flag("tf_xla_embedding_parallel_iterations",
&build_ops_flags->tf_xla_embedding_parallel_iterations,
"If >0 then use this many parallel iterations in "
"embedding_pipelining and embedding_sequency. By default, use the "
"parallel_iterations on the original model WhileOp."),
Flag("tf_xla_compile_on_demand", &device_flags->tf_xla_compile_on_demand,
"Switch a device into 'on-demand' mode, where instead of "
"autoclustering ops are compiled one by one just-in-time."),
Flag("tf_xla_enable_xla_devices",
&device_flags->tf_xla_enable_xla_devices,
"Generate XLA_* devices, where placing a computation on such a "
"device"
"forces compilation by XLA. Deprecated."),
Flag("tf_xla_always_defer_compilation",
&ops_flags->tf_xla_always_defer_compilation, ""),
Flag("tf_xla_async_compilation", &ops_flags->tf_xla_async_compilation,
"When lazy compilation is enabled, asynchronous compilation starts "
"the cluster compilation in the background, and the fallback path "
"is executed until the compilation has finished."),
Flag("tf_xla_use_device_api_for_xla_launch",
&ops_flags->tf_xla_use_device_api.enabled_for_xla_launch_,
"If true, uses Device API (PjRt) for single device compilation and "
"execution of functions marked for JIT compilation i.e. "
"jit_compile=True. Defaults to false."),
Flag("tf_xla_use_device_api_for_compile_on_demand",
&ops_flags->tf_xla_use_device_api.enabled_for_compile_on_demand_,
"If true, uses Device API (PjRt) for compiling and executing ops "
"one by one in 'on-demand' mode. Defaults to false."),
Flag("tf_xla_use_device_api_for_auto_jit",
&ops_flags->tf_xla_use_device_api.enabled_for_compile_and_run_,
"If true, uses Device API (PjRt) for compilation and execution "
"when auto-clustering is enabled. Defaults to false."),
Flag("tf_xla_use_device_api",
&ops_flags->tf_xla_use_device_api.enabled_for_all_,
"If true, uses Device API (PjRt) for compilation and execution "
"of ops one-by-one in 'on-demand' mode, for functions marked for "
"JIT compilation, or when auto-clustering is enabled. Defaults to "
"false."),
Flag("tf_xla_enable_device_api_for_gpu",
&ops_flags->tf_xla_use_device_api.enabled_for_gpu_,
"If true, uses Device API (PjRt) for TF GPU device. This is a "
"helper flag so that individual tests can turn on PjRt for GPU "
"specifically."),
Flag("tf_xla_call_module_disabled_checks",
SetterForXlaCallModuleDisabledChecks, "",
"A comma-sepated list of directives specifying the safety checks "
"to be skipped when compiling XlaCallModuleOp. See the op "
"documentation for the recognized values."),
Flag("tf_mlir_enable_mlir_bridge", &enable_mlir_bridge,
"Enables experimental MLIR-Based TensorFlow Compiler Bridge.",
&enable_mlir_bridge_is_explicit),
Flag("tf_mlir_enable_merge_control_flow_pass",
&enable_mlir_merge_control_flow_pass,
"Enables MergeControlFlow pass for MLIR-Based TensorFlow Compiler "
"Bridge."),
Flag("tf_mlir_enable_convert_control_to_data_outputs_pass",
&enable_mlir_convert_control_to_data_outputs_pass,
"Enables `tf-executor-convert-control-to-data-outputs` pass for "
"MLIR-Based TensorFlow Compiler Bridge."),
Flag("tf_mlir_composite_tpuexecute_side_effects",
&enable_mlir_composite_tpuexecute_side_effects,
"Enables certain TPUExecute ops to run in parallel if they only "
"operate on resources that live on composite devices."),
Flag("tf_mlir_enable_strict_clusters", &enable_mlir_strict_clusters,
"Do not allow clusters that have cyclic control dependencies."),
Flag("tf_mlir_enable_multiple_local_cpu_devices",
&enable_mlir_multiple_local_cpu_devices,
"Enable multiple local CPU devices. CPU ops which are outside "
"compiled inside the tpu cluster will also be replicated across "
"multiple cpu devices."),
Flag("tf_mlir_enable_debug_info_serialization",
&enable_mlir_debug_info_serialization,
"Enable debug info serialization in MLIR."),
Flag("tf_dump_graphs_in_tfg", &use_tfg_graph_dumper,
"When tf_dump_graphs_in_tfg is true, graphs after transformations "
"are dumped in MLIR TFG dialect and not in GraphDef"),
Flag("tf_mlir_enable_tpu_variable_runtime_reformatting_pass",
&enable_tpu_variable_runtime_reformatting_pass,
"Enables TPUVariableRuntimeReformatting pass for MLIR-Based "
"TensorFlow Compiler Bridge. This enables weight update sharding "
"and creates TPUReshardVariables ops."),
Flag("tf_serialize_mlir_to_compressed_bytecode",
&enable_serialize_mlir_to_compressed_bytecode,
"If true, serialize MLIR to compressed bytecode.")});
AppendMarkForCompilationPassFlagsInternal(flag_list);
xla::ParseFlagsFromEnvAndDieIfUnknown("TF_XLA_FLAGS", *flag_list);
mlir_flags = new MlirCommonFlags;
if (!enable_mlir_bridge_is_explicit) {
mlir_flags->tf_mlir_enable_mlir_bridge =
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED;
} else if (enable_mlir_bridge) {
mlir_flags->tf_mlir_enable_mlir_bridge =
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_ENABLED;
} else {
mlir_flags->tf_mlir_enable_mlir_bridge =
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_DISABLED;
}
mlir_flags->tf_mlir_enable_merge_control_flow_pass =
enable_mlir_merge_control_flow_pass;
mlir_flags->tf_mlir_enable_convert_control_to_data_outputs_pass =
enable_mlir_convert_control_to_data_outputs_pass;
mlir_flags->tf_mlir_enable_composite_tpuexecute_side_effects =
enable_mlir_composite_tpuexecute_side_effects;
mlir_flags->tf_mlir_enable_strict_clusters = enable_mlir_strict_clusters;
mlir_flags->tf_mlir_enable_tpu_variable_runtime_reformatting_pass =
enable_tpu_variable_runtime_reformatting_pass;
mlir_flags->tf_mlir_enable_multiple_local_cpu_devices =
enable_mlir_multiple_local_cpu_devices;
mlir_flags->tf_mlir_enable_debug_info_serialization =
enable_mlir_debug_info_serialization;
mlir_flags->tf_serialize_mlir_to_compressed_bytecode =
enable_serialize_mlir_to_compressed_bytecode;
if (use_tfg_graph_dumper) {
UseMlirForGraphDump(MlirDumpConfig{}.elide_large_attributes().emit_dialect(
MlirDumpConfig::Dialect::kTFG));
}
AllocateAndParseJitRtFlags();
}
void ResetFlags() {
delete build_ops_flags;
delete mark_for_compilation_flags;
delete device_flags;
delete ops_flags;
delete mlir_flags;
delete flag_list;
delete jitrt_flags;
delete jitrt_flag_list;
AllocateAndParseFlags();
}
} // namespace
bool SetXlaAutoJitFlagFromFlagString(const std::string& value) {
absl::call_once(flags_init, &AllocateAndParseFlags);
return SetterForXlaAutoJitFlag(value);
}
BuildXlaOpsPassFlags* GetBuildXlaOpsPassFlags() {
absl::call_once(flags_init, &AllocateAndParseFlags);
return build_ops_flags;
}
MarkForCompilationPassFlags* GetMarkForCompilationPassFlags() {
absl::call_once(flags_init, &AllocateAndParseFlags);
return mark_for_compilation_flags;
}
XlaSparseCoreFlags* GetXlaSparseCoreFlags() {
absl::call_once(flags_init, &AllocateAndParseFlags);
return sparse_core_flags;
}
XlaDeviceFlags* GetXlaDeviceFlags() {
absl::call_once(flags_init, &AllocateAndParseFlags);
return device_flags;
}
XlaOpsCommonFlags* GetXlaOpsCommonFlags() {
absl::call_once(flags_init, &AllocateAndParseFlags);
return ops_flags;
}
XlaCallModuleFlags* GetXlaCallModuleFlags() {
absl::call_once(flags_init, &AllocateAndParseFlags);
return call_module_flags;
}
MlirCommonFlags* GetMlirCommonFlags() {
absl::call_once(flags_init, &AllocateAndParseFlags);
return mlir_flags;
}
void ResetJitCompilerFlags() { ResetFlags(); }
const JitRtFlags& GetJitRtFlags() {
absl::call_once(flags_init, &AllocateAndParseFlags);
return *jitrt_flags;
}
ConfigProto::Experimental::MlirBridgeRollout GetMlirBridgeRolloutState(
std::optional<const ConfigProto> config_proto) {
// TF1 graphs that do not override Sessions's ConfigProto and TF2 graphs
// can enable/disable the graph via tf_mlir_enable_mlir_bridge.
auto tf_mlir_enable_mlir_bridge =
GetMlirCommonFlags()->tf_mlir_enable_mlir_bridge;
if (tf_mlir_enable_mlir_bridge !=
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED) {
return tf_mlir_enable_mlir_bridge;
}
// If a ConfigProto was not passed in, we can assume the caller is
// checking if TF2 graph should have the bridge enabled / disabled. In that
// case, we have already checked tf_mlir_enable_mlir_bridge so it is safe to
// return UNSPECIFIED here.
if (!config_proto.has_value()) {
return ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED;
}
// TF1 graphs that do override Session's ConfigProto and set
// ConfigProto's enable_mlir_bridge or mlir_bridge_rollout fields will not
// update tf_mlir_enable_mlir_bridge so check their values.
// ConfigProto's enable_mlir_bridge defaults to false so only respect it
// when it is true.
if (config_proto.value().experimental().enable_mlir_bridge()) {
return ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_ENABLED;
}
return config_proto.value().experimental().mlir_bridge_rollout();
}
void AppendMarkForCompilationPassFlags(std::vector<Flag>* flag_list) {
absl::call_once(flags_init, &AllocateAndParseFlags);
AppendMarkForCompilationPassFlagsInternal(flag_list);
}
static std::atomic<bool> xla_compilation_disabled(false);
void DisableXlaCompilation() { xla_compilation_disabled = true; }
void EnableXlaCompilation() { xla_compilation_disabled = false; }
bool FailOnXlaCompilation() { return xla_compilation_disabled; }
} // namespace tensorflow
tensorflow/compiler/jit/flags.h
+366
View File
@@ -0,0 +1,366 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_JIT_FLAGS_H_
#define TENSORFLOW_COMPILER_JIT_FLAGS_H_
#include <cstdint>
#include <optional>
#include <string>
#include <vector>
#include "absl/container/flat_hash_set.h"
#include "absl/types/optional.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/util/command_line_flags.h"
namespace tensorflow {
struct XlaAutoJitFlag {
// Control compilation of operators into XLA computations on CPU and GPU
// devices. 0 = use ConfigProto setting; -1 = off; 1 = on for things very
// likely to be improved; 2 = on for everything.
//
// If all non-CPU ops in the graph being optimized are placed on a single GPU
// and there is at least one node placed on that GPU then
// `optimization_level_single_gpu` applies. Otherwise
// `optimization_level_general` applies.
//
// Experimental.
int32_t optimization_level_single_gpu;
int32_t optimization_level_general;
};
// Sets the xla_auto_jit_flag based on the given flag string. Supported syntax
// is:
// <number>: sets general and single_gpu setting to the provided number.
// single-gpu(<number>): sets the single_gpu setting to the provided number.
bool SetXlaAutoJitFlagFromFlagString(const std::string& value);
// Flags associated with the XLA bridge's mark_for_compilation_pass module.
struct MarkForCompilationPassFlags {
XlaAutoJitFlag xla_auto_jit_flag;
// Minimum number of operators in an XLA compilation. Ignored for operators
// placed on an XLA device or operators explicitly marked for compilation.
int32_t tf_xla_min_cluster_size;
// Maximum number of operators in an XLA compilation.
int32_t tf_xla_max_cluster_size;
// If non-empty, limit XLA clustering to the following TF operations.
std::string tf_xla_ops_to_cluster;
// If non-empty, remove following operations from XLA clustering excludelist.
std::string tf_xla_cluster_exclude_ops;
// Dump graphs during XLA compilation.
bool tf_xla_clustering_debug;
// Enables global JIT compilation for CPU via SessionOptions.
bool tf_xla_cpu_global_jit;
// "Compiler fuel" for clustering. Only this many ops will be marked as
// eligible for clustering.
int64_t tf_xla_clustering_fuel;
// If tf_xla_disable_deadness_safety_checks_for_debugging is set to true then
// we do not do deadness related safety checks. This is unsound in general,
// but can be used as a debugging aid.
bool tf_xla_disable_deadness_safety_checks_for_debugging;
// If tf_xla_disable_resource_variable_safety_checks_for_debugging is set to
// true then we do not do safety checks to preserve TensorFlow's resource
// variable concurrency semantics. This is unsound in general, but can be
// used as a debugging aid.
bool tf_xla_disable_resource_variable_safety_checks_for_debugging;
// If true names of clustered operations will be computed deterministically
// so that they remain stable from run to run of auto clusteing.
bool tf_xla_deterministic_cluster_names;
// If non-empty, JIT-compiled executables are saved to and loaded from the
// specified file system directory path.
std::string tf_xla_persistent_cache_directory;
// If non-empty, the persistent cache will only be used for the specified
// devices (comma separated). Each device type should be able to be converted
// to `DeviceType`.
std::string tf_xla_persistent_cache_device_types;
bool tf_xla_persistent_cache_read_only;
// If true, entries loaded into the XLA compile cache will not have their
// signatures checked strictly. This should generally not be disabled except
// for debugging. Defaults to false.
bool tf_xla_disable_strict_signature_checks;
// Specifies the persistance cache prefix. Default is "xla_compile_cache"
std::string tf_xla_persistent_cache_prefix;
};
// Flags associated with XLA Sparse Core.
struct XlaSparseCoreFlags {
// Max level of division to split input data into minibatches.
int tf_xla_sparse_core_minibatch_max_division_level;
// Disable table stacking for all the tables passed to the SparseCore
// mid level API.
bool tf_xla_sparse_core_disable_table_stacking;
// If non-zero, limits the size of the activations for a given table to
// be below these many bytes.
int64_t tf_xla_sparse_core_stacking_mem_limit_bytes;
// If non-zero, limits the size of any table shard to be below these
// many bytes.
int64_t tf_xla_sparse_core_stacking_table_shard_limit_bytes;
};
// Flags associated with the XLA bridge's xla_device module.
struct XlaDeviceFlags {
// Switch the CPU device into "on-demand" mode, where instead of
// auto-clustering ops are compiled one by one just-in-time.
// Enabling this mode by a legacy flag is a temporary mechanism. When this
// feature is battle-tested, we will switch this to be a session option.
bool tf_xla_compile_on_demand;
// Enables "XLA" devices if this flag is set.
bool tf_xla_enable_xla_devices;
};
// Flags common to the _Xla* ops and their kernels.
struct XlaOpsCommonFlags {
// If true, _XlaCompile always refuses to compile the cluster, which means the
// XLA clusters always run in the TF executor. Defaults to false.
bool tf_xla_always_defer_compilation;
// If true, _XlaCompile compiles the cluster asynchronously with respect to
// the main execution. The fallback path is taken while compilation happens.
bool tf_xla_async_compilation;
class PjRtForSingleDeviceCompilationRollout {
public:
// Allow using Device API (PjRt) for `device_type` in the XlaLaunch op.
// Please note that `enabled_for_xla_launch_` needs to be true in addition
// to the `device_type` being allowed in order to use the Device API for
// single device compilation and execution in the XlaLaunch op.
void AllowForDeviceInXlaLaunch(const DeviceType& device_type) {
xla_launch_allowed_devices_.insert(device_type.type_string());
}
bool IsEnabledInXlaLaunchForDevice(const DeviceType& device_type) const {
if (!enabled_for_gpu_ && device_type.type_string() == "GPU") return false;
return enabled_for_all_ ||
(enabled_for_xla_launch_ &&
xla_launch_allowed_devices_.contains(device_type.type_string()));
}
// Allow using Device API (PjRt) for `device_type` in the XlaCompileOnDemand
// op. Please note that `enabled_for_compile_on_demand_` needs to be true in
// addition to the `device_type` being allowed in order to use the Device
// API for single device compilation and execution in the XlaCompileOnDemand
// op.
void AllowForDeviceInXlaCompileOnDemand(const DeviceType& device_type) {
xla_compile_on_demand_allowed_devices_.insert(device_type.type_string());
}
bool IsEnabledInXlaCompileOnDemandForDevice(
const DeviceType& device_type) const {
if (!enabled_for_gpu_ && device_type.type_string() == "GPU") return false;
return enabled_for_all_ ||
(enabled_for_compile_on_demand_ &&
xla_compile_on_demand_allowed_devices_.contains(
device_type.type_string()));
}
// Allow using Device API (PjRt) for `device_type` in the XlaCompile and
// XlaRun ops. Please note that `enabled_for_compile_and_run_` needs to be
// true in addition to the `device_type` being allowed in order to use the
// Device API for single device compilation and execution in the XlaCompile
// and XlaRun ops.
void AllowForDeviceInXlaCompileAndRun(const DeviceType& device_type) {
xla_compile_and_run_allowed_devices_.insert(device_type.type_string());
}
bool IsEnabledInXlaCompileAndRunForDevice(
const DeviceType& device_type) const {
if (!enabled_for_gpu_ && device_type.type_string() == "GPU") return false;
return enabled_for_all_ || (enabled_for_compile_and_run_ &&
xla_compile_and_run_allowed_devices_.contains(
device_type.type_string()));
}
bool IsEnabledForGpu() const { return enabled_for_gpu_; }
// If true, uses Device API (PjRt) for single device compilation and
// execution of functions marked for JIT compilation i.e. jit_compile=True.
// Defaults to false.
bool enabled_for_xla_launch_;
// If true, uses Device API (PjRt) for compiling and executing ops one by
// one in "on-demand" mode. Defaults to false.
bool enabled_for_compile_on_demand_;
// If true, uses Device API (PjRt) for compilation and execution when
// auto-clustering is enabled. Defaults to false.
bool enabled_for_compile_and_run_;
// If true, uses Device API (PjRt) for compilation and execution everywhere
// i.e. for functions marked for JIT compilation, for ops in "on-demand"
// mode and auto-clustering. Defaults to false.
//
// Note that this flag can be overridden by device flag like
// `enabled_for_gpu_` below.
bool enabled_for_all_;
// If true, enable Device API (PjRt) for TF GPU device. This is a helper
// flag so that individual tests can turn on PjRt for GPU specifically.
// Once the rollout to GPU is complete, this flag can be deprecated.
bool enabled_for_gpu_;
private:
// Devices for which using Device API (PjRt) is allowed in the XlaLaunch op.
// This can only be modified programmatically.
absl::flat_hash_set<std::string> xla_launch_allowed_devices_;
// Devices for which using Device API (PjRt) is allowed in the
// XlaCompileOnDemand op. This can only be modified programmatically.
absl::flat_hash_set<std::string> xla_compile_on_demand_allowed_devices_;
// Devices for which using Device API (PjRt) is allowed in the
// XlaCompile and XlaRun ops. This can only be modified programmatically.
absl::flat_hash_set<std::string> xla_compile_and_run_allowed_devices_;
} tf_xla_use_device_api;
};
// Flags for the XlaCallModule kernel.
struct XlaCallModuleFlags {
// Used by XlaCallModuleOp to specify safety checks to disable.
absl::flat_hash_set<std::string> disabled_checks;
};
// Flags for the build_xla_ops pass.
struct BuildXlaOpsPassFlags {
// Enables lazy compilation for TF/XLA (only when auto-clustering) if true.
// Defaults to true.
bool tf_xla_enable_lazy_compilation;
// If true then insert Print nodes to print out values produced by XLA
// clusters. Useful for debugging.
bool tf_xla_print_cluster_outputs;
// If true, insert CheckNumerics nodes for every floating point typed input to
// an XLA cluster.
bool tf_xla_check_cluster_input_numerics;
// If true, insert CheckNumerics nodes for every floating point typed output
// from an XLA cluster.
bool tf_xla_check_cluster_output_numerics;
// Disables all constant folding. The primary use for this is for testing to
// guarantee that tests are run on XLA and not on TF's CPU implementation.
bool tf_xla_disable_constant_folding;
// Disables full embedding pipelining when true. Instead, strict SparseCore
// TensorCore sequencing will be used.
bool tf_xla_disable_full_embedding_pipelining;
// Whether to enable automatical embedding pipelining when summary ops are
// detected in the graph.
bool tf_xla_disable_full_embedding_pipelining_with_summaries;
// Force the WhileOps in embedding_pipelining and embedding_sequencing to use
// this many parallel_iterations
int tf_xla_embedding_parallel_iterations;
};
// Flags for common MLIR configurations.
struct MlirCommonFlags {
ConfigProto::Experimental::MlirBridgeRollout tf_mlir_enable_mlir_bridge;
bool tf_mlir_enable_merge_control_flow_pass;
bool tf_mlir_enable_convert_control_to_data_outputs_pass;
bool tf_mlir_enable_composite_tpuexecute_side_effects;
bool tf_mlir_enable_strict_clusters;
bool tf_mlir_enable_tpu_variable_runtime_reformatting_pass;
// TODO(pineapplejuice233): Revisit this flag once the performance impact is verified
// with different local CPU devices settings.
bool tf_mlir_enable_multiple_local_cpu_devices;
bool tf_mlir_enable_debug_info_serialization;
bool tf_serialize_mlir_to_compressed_bytecode;
};
// Flags for the JitRt pipeline -- see tf_jitrt_pipeline.h for details.
struct JitRtFlags {
bool always_specialize;
bool cost_driven_async_parallel_for;
// Enables tracking of the "live" JitRt queries to, on a crash, identify the
// "query of death". See TfJitRtQueryOfDeathLogger.
bool log_query_of_death;
// Enable vectorization, which requires tiling and peeling on different ops.
bool vectorize;
// Enables crash reproducer for JitRt MLIR pass manager.
bool enable_crash_reproducer;
};
// Return a pointer to the DumpGraphFlags struct;
// repeated calls return the same pointer.
// This should be called only after Flags::Parse() has returned.
// Getters for flags structs defined above. The first call to any of these
// parses TF_XLA_FLAGS for all of them. Those functions which return a pointer
// always return the same pointer.
MarkForCompilationPassFlags* GetMarkForCompilationPassFlags();
BuildXlaOpsPassFlags* GetBuildXlaOpsPassFlags();
XlaSparseCoreFlags* GetXlaSparseCoreFlags();
XlaDeviceFlags* GetXlaDeviceFlags();
XlaOpsCommonFlags* GetXlaOpsCommonFlags();
XlaCallModuleFlags* GetXlaCallModuleFlags();
MlirCommonFlags* GetMlirCommonFlags();
void ResetJitCompilerFlags();
const JitRtFlags& GetJitRtFlags();
// Returns the effective MLIR bridge rollout state based on the flags and the
// optional configuration.
ConfigProto::Experimental::MlirBridgeRollout GetMlirBridgeRolloutState(
std::optional<const ConfigProto> config_proto);
// Appends the flag definitions associated with
// MarkForCompilationPassFlags/DumpGraphFlags to `flag_list`.
//
// Has the side-effect of parsing TF_XLA_FLAGS if that hasn't happened yet.
void AppendMarkForCompilationPassFlags(
std::vector<tensorflow::Flag>* flag_list);
// Disables XLA compilation, forces it to return an error message instead. Can
// be used by a server to ensure that JIT compilation is opt-in.
void DisableXlaCompilation();
// Enables XLA compilation. Can be used with `DisableXlaCompilation` to
// enable/disable JIT compilation at different stages.
void EnableXlaCompilation();
// Returns `false` unless `DisableXlaCompilation` was called.
bool FailOnXlaCompilation();
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_FLAGS_H_
tensorflow/compiler/jit/force_xla_constants_on_host_pass.cc
+56
View File
@@ -0,0 +1,56 @@
/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/jit/force_xla_constants_on_host_pass.h"
#include "tensorflow/compiler/jit/compilability_check_util.h"
#include "tensorflow/compiler/jit/defs.h"
#include "tensorflow/core/common_runtime/optimization_registry.h"
namespace tensorflow {
absl::Status ForceXlaConstantsOnHostPass::Run(
const GraphOptimizationPassOptions& options) {
Graph* graph = options.graph->get();
OptimizerOptions opts;
auto pflr = std::make_unique<ProcessFunctionLibraryRuntime>(
nullptr, options.session_options->env, /*config=*/nullptr,
TF_GRAPH_DEF_VERSION, options.flib_def, opts);
FunctionLibraryRuntime* flr =
pflr->GetFLR(ProcessFunctionLibraryRuntime::kDefaultFLRDevice);
for (Node* node : graph->nodes()) {
if (CanCreateXlaKernel(node->def())) {
const FunctionBody* fbody = nullptr;
std::vector<int> constant_arg_indices;
std::vector<int> resource_arg_indices;
NameAttrList function;
TF_RETURN_IF_ERROR(NameAndAttrsFromFunctionCall(node->def(), &function));
// Force all constants to be on the host memory.
TF_RETURN_IF_ERROR(GetBodyAndConstantsAndResources(
flr, function, &fbody, &constant_arg_indices, &resource_arg_indices));
VLOG(3) << "Found constant arg indices: "
<< absl::StrJoin(constant_arg_indices, ", ");
node->AddAttr("_input_hostmem", constant_arg_indices);
}
}
return absl::OkStatus();
}
} // namespace tensorflow
tensorflow/compiler/jit/force_xla_constants_on_host_pass.h
+36
View File
@@ -0,0 +1,36 @@
/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_JIT_FORCE_XLA_CONSTANTS_ON_HOST_PASS_H_
#define TENSORFLOW_COMPILER_JIT_FORCE_XLA_CONSTANTS_ON_HOST_PASS_H_
#include "absl/container/flat_hash_set.h"
#include "tensorflow/compiler/jit/compilability_check_util.h"
#include "tensorflow/core/common_runtime/optimization_registry.h"
namespace tensorflow {
// An optimization pass which marks the constants which have to be resolved for
// XLA compilation with `_input_hostmem`.
class ForceXlaConstantsOnHostPass : public GraphOptimizationPass {
public:
ForceXlaConstantsOnHostPass() = default;
absl::Status Run(const GraphOptimizationPassOptions& options) override;
};
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_FORCE_XLA_CONSTANTS_ON_HOST_PASS_H_
tensorflow/compiler/jit/force_xla_constants_on_host_pass_test.cc
+108
View File
@@ -0,0 +1,108 @@
/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/jit/force_xla_constants_on_host_pass.h"
#include "absl/strings/match.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/cc/framework/ops.h"
#include "tensorflow/cc/ops/functional_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/compiler/jit/compilability_check_util.h"
#include "tensorflow/compiler/jit/defs.h"
#include "tensorflow/compiler/jit/test_util.h"
#include "tensorflow/core/common_runtime/function.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/framework/function_testlib.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/graph/graph_def_builder.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/public/session_options.h"
#include "tensorflow/core/public/version.h"
namespace tensorflow {
namespace {
absl::Status ForceXlaConstantsOnHost(const Scope& s,
FunctionLibraryDefinition* flib_def,
std::unique_ptr<Graph>* result) {
auto graph = std::make_unique<Graph>(OpRegistry::Global());
GraphOptimizationPassOptions options;
SessionOptions session_options;
session_options.env = Env::Default();
options.graph = &graph;
options.session_options = &session_options;
options.flib_def = flib_def;
TF_RETURN_IF_ERROR(s.ToGraph(graph.get()));
ForceXlaConstantsOnHostPass rewriter;
TF_RETURN_IF_ERROR(rewriter.Run(options));
*result = std::move(graph);
return absl::OkStatus();
}
TEST(ForceXlaConstantsOnHostPassTest, Simple) {
GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
Scope root = Scope::NewRootScope().ExitOnError();
FunctionDefLibrary library;
FunctionDef called_func =
FunctionDefHelper::Create("TransposeCall",
/*in_def=*/{"a:float", "b:int32"},
/*out_def=*/{"c:float"}, {},
{{{"t0"},
"Transpose",
{"a", "b"},
{
{"T", DT_FLOAT},
{"Tperm", DT_INT32},
}}},
{{"c", "t0:y:0"}});
AttrValue true_attribute;
true_attribute.set_b(true);
(*called_func.mutable_attr())[kXlaMustCompileAttr] = true_attribute;
*library.add_function() = called_func;
TF_ASSERT_OK(root.graph()->AddFunctionLibrary(library));
FunctionLibraryDefinition flib_def(OpRegistry::Global(), library);
Output in = ops::Placeholder(root, DT_FLOAT);
Output perm = ops::Const(root, {3, 1, 2, 0});
NameAttrList b_name_attr;
b_name_attr.set_name("TransposeCall");
ops::PartitionedCall call(root.WithOpName("call"), {in, perm}, {DT_FLOAT},
b_name_attr);
call.output.front().node()->AddAttr(kXlaMustCompileAttr, true);
std::unique_ptr<Graph> graph;
TF_ASSERT_OK(ForceXlaConstantsOnHost(root, &flib_def, &graph));
bool found = false;
for (Node* node : graph->nodes()) {
if (CanCreateXlaKernel(node->def())) {
EXPECT_FALSE(found);
found = true;
std::vector<int32_t> hostmem_attr;
EXPECT_TRUE(TryGetNodeAttr(node->def(), "_input_hostmem", &hostmem_attr));
EXPECT_EQ(hostmem_attr.size(), 1);
EXPECT_EQ(hostmem_attr[0], 1);
}
}
EXPECT_TRUE(found);
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/get_compiler_ir.cc
+482
View File
@@ -0,0 +1,482 @@
/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/jit/get_compiler_ir.h"
#include <cstdint>
#include <deque>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/algorithm/container.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/ascii.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/jit/compilability_check_util.h"
#include "tensorflow/compiler/jit/device_compiler.h"
#include "tensorflow/compiler/jit/variable_info.h"
#include "tensorflow/compiler/jit/variable_info_util.h"
#include "tensorflow/compiler/jit/xla_compiler_options_util.h"
#include "tensorflow/compiler/jit/xla_launch_util.h"
#include "tensorflow/compiler/jit/xla_platform_info.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/client/executable_build_options.h"
#include "xla/client/local_client.h"
#include "xla/hlo/translate/portable_api.h"
#include "xla/service/hlo_graph_dumper.h"
#include "xla/status_macros.h"
#include "xla/stream_executor/host/host_platform_id.h"
#include "xla/stream_executor/platform.h"
#include "tensorflow/core/common_runtime/eager/tensor_handle.h"
#include "tensorflow/core/framework/device_base.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/framework/resource_handle.pb.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/platform/refcount.h"
#include "tensorflow/core/platform/statusor.h"
#include "tsl/platform/errors.h"
#include "tsl/platform/statusor.h"
namespace tensorflow {
static absl::StatusOr<std::unique_ptr<xla::LocalExecutable>> BuildExecutable(
xla::LocalClient* local_client,
const XlaCompiler::CompilationResult& result,
const XlaCompiler::Options& options,
const bool xla_embed_ir_in_executable = false) {
std::vector<const xla::Shape*> argument_layouts(
result.xla_input_shapes.size());
for (int i = 0, end = result.xla_input_shapes.size(); i < end; ++i) {
argument_layouts[i] = &result.xla_input_shapes[i];
}
xla::ExecutableBuildOptions build_options;
if (result.collective_info) {
build_options.set_num_replicas(result.collective_info->group_size);
}
build_options.set_device_ordinal(
options.device_ordinal != -1 ? options.device_ordinal
: local_client->default_device_ordinal());
build_options.set_result_layout(result.xla_output_shape);
build_options.set_device_allocator(options.device_allocator.get());
build_options.set_alias_passthrough_params(options.alias_passthrough_params);
build_options.mutable_debug_options()->set_xla_detailed_logging(
options.detailed_logging);
// If the embed_ir_in_executable is set, hlo_proto will be dumped in
// executable. The hlo_proto contains HLO modules and buffer assignment.
build_options.mutable_debug_options()->set_xla_embed_ir_in_executable(
xla_embed_ir_in_executable);
TF_ASSIGN_OR_RETURN(
std::vector<std::unique_ptr<xla::LocalExecutable>> executables,
local_client->Compile(*result.computation, argument_layouts,
build_options));
TF_RET_CHECK(executables.size() == 1);
return std::move(executables[0]);
}
static absl::StatusOr<std::string> BuildHLOString(
IrExportStage stage, const XlaCompiler::CompilationResult& result,
xla::LocalClient* local_client, const XlaCompiler::Options& options) {
switch (stage) {
case IrExportStage::STABLEHLO:
case IrExportStage::STABLEHLO_SERIALIZED:
case IrExportStage::HLO:
case IrExportStage::HLO_NO_METADATA:
case IrExportStage::HLO_SERIALIZED: {
TF_ASSIGN_OR_RETURN(xla::ProgramShape program_shape,
result.computation->GetProgramShape());
xla::HloModuleConfig config(program_shape);
TF_ASSIGN_OR_RETURN(
std::unique_ptr<xla::HloModule> new_module,
xla::HloModule::CreateFromProto(result.computation->proto(), config));
if (stage == IrExportStage::STABLEHLO_SERIALIZED) {
TF_ASSIGN_OR_RETURN(
std::string stablehlo,
xla::ConvertHloToStablehlo(*new_module, /*emit_bytecode=*/true));
return stablehlo;
}
if (stage == IrExportStage::STABLEHLO) {
TF_ASSIGN_OR_RETURN(
std::string stablehlo,
xla::ConvertHloToStablehlo(*new_module, /*emit_bytecode=*/false));
return stablehlo;
}
xla::HloPrintOptions opts;
opts.set_print_large_constants(false);
opts.set_print_operand_shape(true);
if (stage == IrExportStage::HLO_NO_METADATA) {
opts.set_print_metadata(false);
}
if (stage == IrExportStage::HLO_SERIALIZED) {
return new_module->ToProto().SerializeAsString();
}
return new_module->ToString(opts);
}
case IrExportStage::OPTIMIZED_HLO:
case IrExportStage::OPTIMIZED_HLO_SERIALIZED: {
TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::LocalExecutable> executable,
BuildExecutable(local_client, result, options));
xla::Executable* new_executable = executable->executable();
if (stage == IrExportStage::OPTIMIZED_HLO_SERIALIZED) {
return new_executable->module().ToProto().SerializeAsString();
} else {
return new_executable->module().ToString();
}
}
case IrExportStage::OPTIMIZED_HLO_PROTO_SERIALIZED: {
TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::LocalExecutable> executable,
BuildExecutable(local_client, result, options,
/*xla_embed_ir_in_executable=*/true));
return executable->executable()->hlo_proto()->SerializeAsString();
}
case IrExportStage::OPTIMIZED_HLO_DOT: {
TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::LocalExecutable> executable,
BuildExecutable(local_client, result, options));
absl::StatusOr<std::string> graph = xla::RenderGraph(
*executable->executable()->module().entry_computation(),
"Visualization",
/*debug_options=*/{}, xla::RenderedGraphFormat::kDot,
/*hlo_render_options=*/{});
TF_RETURN_IF_ERROR(graph.status());
return *graph;
}
}
}
static absl::StatusOr<std::vector<XlaCompiler::Argument>>
BuildXlaCompilerArgumentFromTensorSpec(
const FunctionBody* fbody, absl::Span<int const> must_be_constant_idxs,
absl::Span<const Tensor* const> inputs,
absl::Span<VariableInfo const> variable_args,
absl::Span<const ArgShapeAndDType> flat_arg_shape_and_dtype) {
TF_RET_CHECK(fbody != nullptr);
auto& input_args = fbody->record->fdef().signature().input_arg();
int input_arg_size = input_args.size();
std::vector<XlaCompiler::Argument> args;
args.reserve(input_arg_size);
for (auto& arg_info : flat_arg_shape_and_dtype) {
XlaCompiler::Argument arg;
arg.kind = XlaCompiler::Argument::kParameter;
arg.type = arg_info.dtype;
arg.shape = arg_info.shape;
args.push_back(arg);
}
// Build Xla Compiler Arguments from concrete_fn.captured_inputs
absl::flat_hash_map<int, const VariableInfo*> variable_info_lookup;
TF_RETURN_IF_ERROR(
CreateVariableInfoLookup(variable_args, variable_info_lookup));
for (const VariableInfo& info : variable_args) {
TF_RET_CHECK(!info.var() || info.lock_held() || info.shared_lock_held())
<< "Need to hold the lock on resource variables "
"before calling BuildXlaCompilerArguments";
variable_info_lookup.emplace(info.index(), &info);
}
int offset = flat_arg_shape_and_dtype.size();
// Here it takes in the concrete_fn.captured_inputs and builds the appropriate
// XLA compiler arguments.
for (int64_t input_num = offset; input_num < input_arg_size; ++input_num) {
const Tensor* input = inputs[input_num];
XlaCompiler::Argument arg;
if (variable_info_lookup.count(input_num)) {
// Handles tf.resource variables.
TF_RET_CHECK(input->dtype() == DT_RESOURCE);
const VariableInfo& variable = *variable_info_lookup[input_num];
arg.kind = XlaCompiler::Argument::kResource;
arg.resource_kind = XlaResource::kVariable;
arg.definition_stack_trace = variable.definition_stack_trace();
TF_RET_CHECK(variable.var() && variable.var()->is_initialized);
const Tensor* value = variable.var()->tensor();
arg.type = value->dtype();
arg.shape = value->shape();
arg.initialized = true;
} else {
// Instead of embedding constant into HLO,
// we handle tf.constant as parameter to reduce size.
arg.kind = XlaCompiler::Argument::kParameter;
arg.type = input->dtype();
arg.shape = input->shape();
}
args.push_back(arg);
}
for (int64_t i = 0; i < input_arg_size; ++i) {
args[i].name = input_args[i].name();
}
return args;
}
absl::Status ValidateGetCompilerIrTfrtTpu(absl::string_view device_type,
stream_executor::Stream* stream,
IrExportStage stage) {
auto is_tfrt_tpu_supported_stage = [](IrExportStage stage) {
return stage == IrExportStage::HLO ||
stage == IrExportStage::HLO_NO_METADATA ||
stage == IrExportStage::HLO_SERIALIZED ||
stage == IrExportStage::STABLEHLO ||
stage == IrExportStage::STABLEHLO_SERIALIZED;
};
// TODO(b/238830423): support GetCompilerIr on TFRT TPU device for stages
// that requires compilation from HLO to executable.
if (device_type != DEVICE_CPU && stream == nullptr &&
!is_tfrt_tpu_supported_stage(stage)) {
return absl::InternalError(
"GetCompilerIr with requested stage is not supported on this device.");
}
return absl::OkStatus();
}
absl::StatusOr<std::vector<XlaCompiler::Argument>> PrepareXlaCompilerArgs(
FunctionLibraryRuntime* flr, const NameAttrList& function,
EagerContext* context, Device* dev,
absl::Span<const ArgShapeAndDType> input_arg_shape_and_dtype,
absl::Span<const TensorHandle* const> input_handles,
CompilerArgSource compiler_arg_source) {
const FunctionBody* fbody = nullptr;
std::vector<int> constant_arg_indices;
std::vector<int> resource_arg_indices;
TF_RETURN_IF_ERROR(GetBodyAndConstantsAndResources(
flr, function, &fbody, &constant_arg_indices, &resource_arg_indices));
if (dev == nullptr && !resource_arg_indices.empty()) {
return absl::InternalError(
"GetCompilerIr does not support resource arguments when device is "
"null.");
}
// `input_args` includes both concrete_fn input args and captured_input here.
auto& input_args = fbody->record->fdef().signature().input_arg();
// Here input_arg_size = len(flat_args) + len(captured_input)
int input_arg_size = input_args.size();
std::vector<const Tensor*> inputs(input_arg_size);
std::deque<Tensor> inputs_storage;
std::vector<VariableInfo> variable_infos;
int offset = input_arg_shape_and_dtype.size();
for (int i = 0; i < input_handles.size(); i++) {
const TensorHandle* th = input_handles[i];
const Tensor* t;
// Handle owns the tensor.
TF_RETURN_IF_ERROR(th->Tensor(&t));
if (absl::c_binary_search(constant_arg_indices, i)) {
// Need to make sure it's on the host.
inputs_storage.emplace_back(t->dtype(), t->shape());
TF_RETURN_IF_ERROR(
th->CopyToDevice(*context, /*d=*/nullptr, &inputs_storage.back()));
inputs[i + offset] = &inputs_storage.back();
} else {
inputs[i + offset] = t;
}
}
if (dev != nullptr) {
TF_RETURN_IF_ERROR(GetVariableInfosFromInputs(dev->resource_manager(), dev,
inputs, resource_arg_indices,
&variable_infos));
TF_RETURN_IF_ERROR(LockVariables(absl::MakeSpan(variable_infos)));
}
absl::StatusOr<std::vector<XlaCompiler::Argument>> args;
if (compiler_arg_source == CompilerArgSource::TENSOR_SPEC) {
args = BuildXlaCompilerArgumentFromTensorSpec(fbody, constant_arg_indices,
inputs, variable_infos,
input_arg_shape_and_dtype);
} else if (compiler_arg_source == CompilerArgSource::CONCRETE_INPUT) {
args = XlaComputationLaunchContext::BuildXlaCompilerArguments(
constant_arg_indices, inputs, variable_infos, dev);
}
return args;
}
absl::StatusOr<std::string> CompileAndBuildHLOString(
IrExportStage stage, const XlaCompiler::Options& options,
xla::LocalClient* local_client, const NameAttrList& function,
const std::vector<XlaCompiler::Argument>& args) {
XlaCompiler::CompileOptions compile_options;
compile_options.always_return_tuple = false;
compile_options.alias_resource_update = true;
XlaCompiler compiler(options);
XlaCompiler::CompilationResult result;
TF_RETURN_IF_ERROR(
compiler.CompileFunction(compile_options, function, args, &result));
return BuildHLOString(stage, result, local_client, options);
}
/**
* Clarifies the different meanings of 'input_arg_shape_and_dtype' and
* 'input_handles' in different cases.
*
* For TENSOR_SPEC case:
* - `input_arg_shape_and_dtype`: Contains the shape and dtype of
* concrete_fn input args.
* - `input_handles`: Contains the concrete_fn.captured_input tensors.
*
* For CONCRETE_INPUT case:
* - `input_arg_shape_and_dtype`: it is empty.
* - `input_handles`: Contains all concrete_fn inputs tensors, including
* captured inputs.
*/
absl::StatusOr<std::string> GetCompilerIr(
IrExportStage stage, ProcessFunctionLibraryRuntime* pflr,
absl::string_view func_name, Device* dev, EagerContext* context,
absl::Span<const ArgShapeAndDType> input_arg_shape_and_dtype,
absl::Span<const TensorHandle* const> input_handles,
CompilerArgSource compiler_arg_source) {
using XlaDeviceCompiler =
DeviceCompiler<xla::LocalExecutable, xla::LocalClient>;
se::Stream* stream = nullptr;
if (const DeviceBase::AcceleratorDeviceInfo* accelerator_device_info =
dev->tensorflow_accelerator_device_info()) {
stream = accelerator_device_info->stream;
}
TF_RETURN_IF_ERROR(
ValidateGetCompilerIrTfrtTpu(dev->device_type(), stream, stage));
NameAttrList function;
function.set_name(std::string{func_name});
FunctionLibraryRuntime* flr = pflr->GetFLR(dev->name());
TF_ASSIGN_OR_RETURN(
std::vector<XlaCompiler::Argument> args,
PrepareXlaCompilerArgs(flr, function, context, dev,
input_arg_shape_and_dtype, input_handles,
compiler_arg_source));
XlaPlatformInfo platform_info = XlaPlatformInfoFromDevice(dev);
auto compilation_device_type = platform_info.device_type();
if (platform_info.device_type() != DEVICE_TPU) {
TF_ASSIGN_OR_RETURN(compilation_device_type,
GetCompilationDeviceType(platform_info.device_type()));
}
XlaDeviceCompiler* xla_device_compiler;
TF_RETURN_IF_ERROR(dev->resource_manager()->LookupOrCreate<XlaDeviceCompiler>(
dev->resource_manager()->default_container(), "xla_device_compiler",
&xla_device_compiler, [&](XlaDeviceCompiler** xla_device_compiler) {
return BuildXlaDeviceCompiler(dev, flr, platform_info,
compilation_device_type,
xla_device_compiler);
}));
core::ScopedUnref xla_device_compiler_ref(xla_device_compiler);
XlaCompiler::Options options;
if (platform_info.device_type() == DEVICE_TPU && stream == nullptr) {
options = GenerateCompilerOptionsForTfrtTpu(*xla_device_compiler, *flr);
} else {
options = GenerateCompilerOptions(*xla_device_compiler, *flr, dev, stream,
platform_info,
/*has_ref_vars=*/false);
}
return CompileAndBuildHLOString(stage, options, xla_device_compiler->client(),
function, args);
}
absl::StatusOr<std::string> GetCompilerIr(
IrExportStage stage, ProcessFunctionLibraryRuntime* pflr,
absl::string_view func_name, absl::string_view platform_name,
EagerContext* context,
absl::Span<const ArgShapeAndDType> input_arg_shape_and_dtype,
absl::Span<const TensorHandle* const> input_handles,
CompilerArgSource compiler_arg_source) {
using XlaDeviceCompiler =
DeviceCompiler<xla::LocalExecutable, xla::LocalClient>;
TF_RETURN_IF_ERROR(
ValidateGetCompilerIrTfrtTpu(platform_name, /*stream=*/nullptr, stage));
NameAttrList function;
function.set_name(std::string{func_name});
std::string device_name;
if (!platform_name.empty()) {
device_name = absl::StrCat("/device:", platform_name, ":0");
}
FunctionLibraryRuntime* flr = pflr->GetFLR(device_name);
if (flr == nullptr) {
// Use CPU as the fallback to get the `FunctionLibraryRuntime`.
flr = pflr->GetFLR("/device:CPU:0");
}
TF_ASSIGN_OR_RETURN(
std::vector<XlaCompiler::Argument> args,
PrepareXlaCompilerArgs(flr, function, context, /*dev=*/nullptr,
input_arg_shape_and_dtype, input_handles,
compiler_arg_source));
se::Platform::Id platform_id = nullptr;
if (platform_name == DEVICE_CPU) {
platform_id = se::host::kHostPlatformId;
}
XlaPlatformInfo platform_info(DeviceType(platform_name), platform_id,
/*xla_device_metadata=*/nullptr,
/*pjrt_device_metadata=*/nullptr,
/*device_allocator=*/nullptr);
DeviceType compilation_device_type = platform_info.device_type();
if (platform_info.device_type() != DEVICE_TPU) {
TF_ASSIGN_OR_RETURN(compilation_device_type,
GetCompilationDeviceType(platform_info.device_type()));
}
XlaDeviceCompiler* xla_device_compiler;
const std::string xla_device_compiler_name = absl::StrCat(
absl::AsciiStrToLower(platform_name), "_xla_device_compiler");
TF_RETURN_IF_ERROR(
context->HostCPU()->resource_manager()->LookupOrCreate<XlaDeviceCompiler>(
context->HostCPU()->resource_manager()->default_container(),
xla_device_compiler_name, &xla_device_compiler,
[&](XlaDeviceCompiler** xla_device_compiler) {
return BuildXlaDeviceCompiler(/*dev=*/nullptr, flr, platform_info,
compilation_device_type,
xla_device_compiler);
}));
core::ScopedUnref xla_device_compiler_ref(xla_device_compiler);
XlaCompiler::Options options;
if (platform_info.device_type() == DEVICE_TPU) {
options = GenerateCompilerOptionsForTfrtTpu(*xla_device_compiler, *flr);
} else {
options.device_type = compilation_device_type;
options.flib_def = flr->GetFunctionLibraryDefinition();
options.graph_def_version = flr->graph_def_version();
options.allow_cpu_custom_calls =
(platform_info.platform_id() == se::host::kHostPlatformId);
options.alias_passthrough_params = !platform_info.is_on_xla_device();
}
return CompileAndBuildHLOString(stage, options, /*local_client=*/nullptr,
function, args);
}
} // namespace tensorflow
tensorflow/compiler/jit/get_compiler_ir.h
+82
View File
@@ -0,0 +1,82 @@
/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_JIT_GET_COMPILER_IR_H_
#define TENSORFLOW_COMPILER_JIT_GET_COMPILER_IR_H_
#include <string>
#include <vector>
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/platform/statusor.h"
namespace tensorflow {
class ProcessFunctionLibraryRuntime;
class Device;
class Tensor;
class TensorHandle;
class EagerContext;
enum class IrExportStage {
STABLEHLO,
STABLEHLO_SERIALIZED,
HLO,
HLO_NO_METADATA,
HLO_SERIALIZED,
OPTIMIZED_HLO,
OPTIMIZED_HLO_SERIALIZED,
OPTIMIZED_HLO_PROTO_SERIALIZED,
OPTIMIZED_HLO_DOT
};
struct ArgShapeAndDType {
TensorShape shape;
DataType dtype;
};
enum class CompilerArgSource {
TENSOR_SPEC,
CONCRETE_INPUT,
};
// Returns the IR format of the selected stage for a given function `func_name`
// using library runtime `runtime` on a device `dev` with given
// `inputs_arg_shape_and_dtype` and `input_handles`.
absl::StatusOr<std::string> GetCompilerIr(
IrExportStage stage, ProcessFunctionLibraryRuntime* pflr,
absl::string_view func_name, Device* dev, EagerContext* context,
absl::Span<const ArgShapeAndDType> input_arg_shape_and_dtype,
absl::Span<const TensorHandle* const> input_handles,
CompilerArgSource compiler_arg_source);
// Returns the IR format of the selected stage for a given function `func_name`
// using library runtime `runtime` on a platform `platform_name` with given
// `inputs_arg_shape_and_dtype` and `input_handles`.
absl::StatusOr<std::string> GetCompilerIr(
IrExportStage stage, ProcessFunctionLibraryRuntime* pflr,
absl::string_view func_name, absl::string_view platform_name,
EagerContext* context,
absl::Span<const ArgShapeAndDType> input_arg_shape_and_dtype,
absl::Span<const TensorHandle* const> input_handles,
CompilerArgSource compiler_arg_source);
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_GET_COMPILER_IR_H_
tensorflow/compiler/jit/increase_dynamism_for_auto_jit_pass.cc
+408
View File
@@ -0,0 +1,408 @@
/* Copyright 2018 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/jit/increase_dynamism_for_auto_jit_pass.h"
#include <iterator>
#include "absl/algorithm/container.h"
#include "absl/log/check.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/types/optional.h"
#include "tensorflow/cc/framework/ops.h"
#include "tensorflow/cc/framework/scope.h"
#include "tensorflow/cc/framework/scope_internal.h"
#include "tensorflow/cc/ops/array_ops.h"
#include "tensorflow/cc/ops/const_op.h"
#include "tensorflow/cc/ops/math_ops.h"
#include "tensorflow/compiler/jit/flags.h"
#include "tensorflow/compiler/jit/xla_cluster_util.h"
#include "xla/status_macros.h"
#include "tensorflow/core/common_runtime/optimization_registry.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/graph/algorithm.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/util/device_name_utils.h"
#include "tensorflow/core/util/dump_graph.h"
#include "tsl/platform/errors.h"
#include "tsl/platform/statusor.h"
namespace tensorflow {
namespace {
// StatusOrOptional<T> instances hold
//
// - A non-OK Status to indicate an error that needs to be propagated out of
// this pass (e.g. the Graph is malformed).
//
// - A nullopt to indicate the function that created the instance failed to do
// what it set out to do but this is not actually an error
// (e.g. TryToGetTensorFromConstOp was passed a non-Const node).
//
// - A T to indicate a successful operation.
template <class T>
using StatusOrOptional = StatusOr<std::optional<T>>;
StatusOrOptional<Tensor> TryToGetTensorFromConstOp(Node* n) {
if (n->type_string() != "Const") {
return {std::nullopt};
}
const TensorProto* proto = nullptr;
TF_RETURN_IF_ERROR(GetNodeAttr(n->def(), "value", &proto));
Tensor tensor(proto->dtype());
TF_RET_CHECK(tensor.FromProto(*proto));
return {tensor};
}
struct SliceInputs {
Output slice_op;
Output input;
Output begin;
Output size;
// The size of the TF slice operation as a std::vector. We can always compute
// this because we only manipulate slices with a Const size.
std::vector<int64_t> size_as_vector;
};
std::vector<int64_t> IntTensorAsVector(const Tensor& t) {
DCHECK(t.dtype() == DT_INT32 || t.dtype() == DT_INT64);
std::vector<int64_t> result;
result.reserve(t.NumElements());
for (int i = 0; i < t.NumElements(); i++) {
int64_t element = t.dtype() == DT_INT32
? static_cast<int64_t>(t.flat<int32_t>()(i))
: t.flat<int64_t>()(i);
result.push_back(element);
}
return result;
}
// Packages up the inputs to a Slice operation into an instance of
// `SliceInputs`.
StatusOrOptional<SliceInputs> GetSliceInputs(Node* slice) {
const int kSliceInputIndex = 0;
const int kSliceBeginIndex = 1;
const int kSliceSizeIndex = 2;
const Edge* slice_input_edge;
TF_RETURN_IF_ERROR(slice->input_edge(kSliceInputIndex, &slice_input_edge));
const Edge* slice_size_edge;
TF_RETURN_IF_ERROR(slice->input_edge(kSliceSizeIndex, &slice_size_edge));
const Edge* slice_begin_edge;
TF_RETURN_IF_ERROR(slice->input_edge(kSliceBeginIndex, &slice_begin_edge));
SliceInputs slice_inputs;
slice_inputs.input =
Output(slice_input_edge->src(), slice_input_edge->src_output());
slice_inputs.begin =
Output(slice_begin_edge->src(), slice_begin_edge->src_output());
slice_inputs.size =
Output(slice_size_edge->src(), slice_size_edge->src_output());
TF_ASSIGN_OR_RETURN(std::optional<Tensor> tf_slice_size,
TryToGetTensorFromConstOp(slice_inputs.size.node()));
if (!tf_slice_size.has_value()) {
return {std::nullopt};
}
if (tf_slice_size->dims() != 1) {
return {std::nullopt};
}
slice_inputs.size_as_vector = IntTensorAsVector(*tf_slice_size);
return {slice_inputs};
}
// Casts `x` to a DT_INT64 if it isn't one already.
Output MakeInt64(const Scope& host_scope, absl::string_view name,
const Output& x) {
return x.type() == DT_INT64
? x
: ops::Cast(host_scope.WithOpName(name, "_s64"), x, DT_INT64);
}
// Returns `slice_inputs` with the index and size inputs cast to DT_INT64.
SliceInputs MakeSliceIndexAndSizeInt64(const Scope& host_scope,
const SliceInputs& slice_inputs) {
SliceInputs result;
result.input = slice_inputs.input;
result.begin = MakeInt64(host_scope, "begin", slice_inputs.begin);
result.size = MakeInt64(host_scope, "size", slice_inputs.size);
result.size_as_vector = slice_inputs.size_as_vector;
return result;
}
// This class caches emitted constants to avoid creating multiple nodes for the
// same constant value. This helps make the generated GraphDef more readable.
class ConstantCache {
public:
explicit ConstantCache(const Scope& s,
const std::vector<const Edge*>& control_deps)
: scope_(s), control_deps_(control_deps) {}
Output Get1DHostConstant(int64_t constant) {
auto it = cache_.find(constant);
if (it == cache_.end()) {
Output new_const =
ops::Const(scope_.WithOpName("const_", constant), {constant});
it = cache_.insert({constant, new_const}).first;
for (const Edge* e : control_deps_) {
scope_.graph()->AddControlEdge(e->src(), new_const.node());
}
}
return it->second;
}
private:
Scope scope_;
std::unordered_map<int, Output> cache_;
std::vector<const Edge*> control_deps_;
};
// Returns a node computing the size of the Slice op with inputs `slice_inputs`.
absl::Status ComputeSliceSize(const Scope& host_scope,
const SliceInputs& slice_inputs,
std::vector<const Edge*> control_deps,
Output* size) {
// If slice_size[i] >= 0 then slice_size[i] = slice_size[i].
//
// If slice_size[i] == -1 then slice_size[i] = input_size[i] -
// begin[i].
//
// If slice_size[i] < -1 then executing the slice will throw an error, and we
// don't do anything here. We've already filtered these cases out in
// IsRewritableSlice.
if (absl::c_all_of(slice_inputs.size_as_vector,
[](int64_t i) { return i >= 0; })) {
*size = slice_inputs.size;
return absl::OkStatus();
}
Output input_shape =
ops::Shape(host_scope.WithOpName("input_shape"), slice_inputs.input,
ops::Shape::OutType(DT_INT64));
ConstantCache constant_pool(host_scope, control_deps);
std::vector<Output> slice_size;
for (int i = 0, end = slice_inputs.size_as_vector.size(); i < end; i++) {
if (slice_inputs.size_as_vector[i] >= 0) {
slice_size.push_back(
constant_pool.Get1DHostConstant(slice_inputs.size_as_vector[i]));
continue;
}
DCHECK_EQ(slice_inputs.size_as_vector[i], -1);
Output begin_i = ops::Slice(
host_scope.WithOpName("begin_", i), slice_inputs.begin,
constant_pool.Get1DHostConstant(i), constant_pool.Get1DHostConstant(1));
Output input_shape_i = ops::Slice(
host_scope.WithOpName("input_shape_", i), input_shape,
constant_pool.Get1DHostConstant(i), constant_pool.Get1DHostConstant(1));
slice_size.push_back(ops::Sub(host_scope.WithOpName("slice_size_", i),
input_shape_i, begin_i));
DCHECK_EQ(slice_size.back().type(), DT_INT64);
}
// Trivial ConcatV2 nodes (with exactly one input) are disallowed.
if (slice_size.size() == 1) {
*size = slice_size[0];
} else {
auto concat_axis = ops::Const(host_scope.WithOpName("concat_axis"), 0);
for (const Edge* e : control_deps) {
host_scope.graph()->AddControlEdge(e->src(), concat_axis.node());
}
*size = ops::Concat(host_scope.WithOpName("slice_size"), slice_size,
concat_axis);
}
return absl::OkStatus();
}
// Terminology: "static sized" slice is a slice with the
// _XlaCompileTimeConstantInputs attribute set to {2}. The output shape of
// these slices can be solely determined by their "size" input.
absl::Status ConvertTensorFlowSliceToStaticShapedSlice(
Graph* g, Node* slice, const SliceInputs& slice_inputs,
absl::string_view cluster_name, Node** result) {
std::string host_name;
TF_RETURN_IF_ERROR(DeviceNameUtils::DeviceNameToCpuDeviceName(
slice->assigned_device_name(), &host_name));
absl::Status status;
Scope main_scope =
NewInternalScope(g, &status, /*refiner=*/nullptr)
.WithXlaCluster(std::string(cluster_name))
.NewSubScope(absl::StrCat(slice->name(), "/static_shaped_slice"));
Scope host_scope = main_scope.WithAssignedDevice(host_name);
// In the future we may want to be clever here and avoid the extra Cast ops.
SliceInputs slice_inputs_int64 =
MakeSliceIndexAndSizeInt64(host_scope, slice_inputs);
// Create a list of all control dependencies to be copied when possibly
// replacing nodes related to slice_size.
Node* old_size;
std::vector<const Edge*> old_size_ctrl_deps;
TF_RETURN_IF_ERROR(slice->input_node(2, &old_size));
absl::c_copy_if(old_size->in_edges(), std::back_inserter(old_size_ctrl_deps),
[](const Edge* e) { return e->IsControlEdge(); });
Output slice_size;
TF_RETURN_IF_ERROR(ComputeSliceSize(host_scope, slice_inputs_int64,
old_size_ctrl_deps, &slice_size));
*result =
ops::Slice(main_scope.WithAssignedDevice(slice->assigned_device_name())
.WithOpName("static_shaped_slice"),
slice_inputs_int64.input, slice_inputs_int64.begin, slice_size)
.node();
TF_RETURN_IF_ERROR(main_scope.status());
std::vector<std::string> compile_time_const_inputs;
compile_time_const_inputs.push_back("size");
(*result)->AddAttr(kXlaCompileTimeConstantInputsAttr,
compile_time_const_inputs);
return status;
}
void ReplaceTensorFlowSliceWithStaticShapedSlice(Graph* g, Node* slice,
Node* static_shaped_slice) {
std::vector<const Edge*> slice_out_edges;
absl::c_copy(slice->out_edges(), std::back_inserter(slice_out_edges));
for (const Edge* e : slice_out_edges) {
DCHECK(e->src_output() == 0 || e->src_output() == Graph::kControlSlot);
int src_output = e->src_output();
int dst_input = e->dst_input();
Node* dst = e->dst();
g->RemoveEdge(e);
g->AddEdge(static_shaped_slice, src_output, dst, dst_input);
}
for (const Edge* e : slice->in_edges()) {
if (e->IsControlEdge()) {
g->AddControlEdge(e->src(), static_shaped_slice);
}
}
g->RemoveNode(slice);
}
absl::Status RewriteSlice(Graph* g, Node* slice,
const SliceInputs& slice_inputs,
absl::string_view cluster_name) {
VLOG(3) << "Rewriting slice " << slice->name()
<< " to a \"static shaped\" Slice";
Node* static_shaped_slice;
TF_RETURN_IF_ERROR(ConvertTensorFlowSliceToStaticShapedSlice(
g, slice, slice_inputs, cluster_name, &static_shaped_slice));
ReplaceTensorFlowSliceWithStaticShapedSlice(g, slice, static_shaped_slice);
return absl::OkStatus();
}
// Return true if `n` is a slice we should rewrite to have a static shape
// (i.e. have the output shape only depend on the "size" input).
absl::StatusOr<bool> ShouldRewriteSlice(Node* n) {
if (n->type_string() != "Slice") {
return false;
}
if (!GetXlaClusterForNode(*n).has_value()) {
// There is no need to change slice ops outside XLA clusters.
return false;
}
TF_ASSIGN_OR_RETURN(std::optional<SliceInputs> slice_inputs,
GetSliceInputs(n));
if (!slice_inputs.has_value()) {
return false;
}
// If slice_size[i] < -1 for any i then executing the slice will throw an
// error, and we don't do anything here.
bool slice_size_has_error =
absl::c_all_of(slice_inputs->size_as_vector,
[](int64_t size_i) { return size_i >= -1; });
if (!slice_size_has_error) {
return false;
}
// No point in rewriting slices that have both size and begin as constants.
return !slice_inputs->begin.node()->IsConstant();
}
absl::Status FindAndRewriteSlices(Graph* g, bool* changed) {
std::vector<Node*> slices_to_rewrite;
for (Node* n : g->nodes()) {
TF_ASSIGN_OR_RETURN(bool is_rewritable, ShouldRewriteSlice(n));
if (is_rewritable) {
slices_to_rewrite.push_back(n);
}
}
for (Node* n : slices_to_rewrite) {
TF_ASSIGN_OR_RETURN(std::optional<SliceInputs> slice_inputs,
GetSliceInputs(n));
TF_RET_CHECK(slice_inputs.has_value());
TF_RETURN_IF_ERROR(
RewriteSlice(g, n, *slice_inputs, *GetXlaClusterForNode(*n)));
}
if (!slices_to_rewrite.empty()) {
// We've added constants to the graph; hook them up to _SOURCE.
FixupSourceAndSinkEdges(g);
}
*changed = !slices_to_rewrite.empty();
return absl::OkStatus();
}
} // namespace
absl::Status IncreaseDynamismForAutoJitPass::Run(
const GraphOptimizationPassOptions& options) {
MarkForCompilationPassFlags* flags = GetMarkForCompilationPassFlags();
if (flags->tf_xla_clustering_debug) {
DumpGraphToFile("before_increase_dynamism_for_auto_jit_pass",
**options.graph, options.flib_def);
}
bool changed;
TF_RETURN_IF_ERROR(FindAndRewriteSlices(options.graph->get(), &changed));
if (changed && flags->tf_xla_clustering_debug) {
DumpGraphToFile("increase_dynamism_for_auto_jit_pass", **options.graph,
options.flib_def);
}
return absl::OkStatus();
}
} // namespace tensorflow
tensorflow/compiler/jit/increase_dynamism_for_auto_jit_pass.h
+59
View File
@@ -0,0 +1,59 @@
/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_JIT_INCREASE_DYNAMISM_FOR_AUTO_JIT_PASS_H_
#define TENSORFLOW_COMPILER_JIT_INCREASE_DYNAMISM_FOR_AUTO_JIT_PASS_H_
#include "absl/status/status.h"
#include "tensorflow/core/common_runtime/optimization_registry.h"
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
// Increases the amount of "dynamism" representable by XLA clusters by rewriting
// the TensorFlow graph. This pass does the following rewrites:
//
// Slice
// -----
//
// Slice(op, begin, size <must be constant>) =>
// Slice(op, begin, actual_size(op.shape(), size, begin));
// _XlaCompileTimeConstantInputs={2}
//
// where
//
// actual_size(op_shape, size, begin)[i] =
// size[i] == -1 ? (op_shape[i] - size[i])
// : size[i]
//
// This pass, combined with jit/partially_decluster_pass, reduces the number of
// unnecessary cluster recompilations in some common cases. After the rewrite
// shown above jit/partially_decluster_pass extracts the actual_size(...)
// computation to outside the XLA cluster, causing the cluster to be versioned
// only on the actual size of the XlaDynamicSlice. This avoids recompilation
// due to superficial changes that don't affect tensor shapes.
//
// Future Work TODO(b/111210515)
// -----------------------------
//
// In the future we will also translate StridedSlice and Pad a similar way.
class IncreaseDynamismForAutoJitPass : public GraphOptimizationPass {
public:
absl::Status Run(const GraphOptimizationPassOptions& options) override;
};
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_INCREASE_DYNAMISM_FOR_AUTO_JIT_PASS_H_
tensorflow/compiler/jit/increase_dynamism_for_auto_jit_pass_test.cc
+477
View File
@@ -0,0 +1,477 @@
/* Copyright 2018 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/jit/increase_dynamism_for_auto_jit_pass.h"
#include <gmock/gmock.h>
#include "absl/status/status.h"
#include "tensorflow/cc/framework/ops.h"
#include "tensorflow/cc/framework/scope.h"
#include "tensorflow/cc/ops/array_ops.h"
#include "tensorflow/cc/ops/const_op.h"
#include "tensorflow/compiler/jit/node_matchers.h"
#include "tensorflow/compiler/jit/xla_cluster_util.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "tensorflow/core/common_runtime/device_set.h"
#include "tensorflow/core/common_runtime/optimization_registry.h"
#include "tensorflow/core/framework/allocator.h"
#include "tensorflow/core/framework/device.h"
#include "tensorflow/core/framework/device_attributes.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/public/session_options.h"
#include "tsl/platform/errors.h"
namespace tensorflow {
namespace {
using ::testing::_;
using testing::matchers::AssignedDevice;
using testing::matchers::Attr;
using testing::matchers::Const;
using testing::matchers::CtrlDeps;
using testing::matchers::Inputs;
using testing::matchers::Name;
using testing::matchers::NodeWith;
using testing::matchers::Op;
using testing::matchers::Out;
// A fake device used to populate a DeviceSet.
class FakeDevice : public Device {
public:
explicit FakeDevice(const DeviceAttributes& device_attributes)
: Device(nullptr, device_attributes) {}
absl::Status Sync() override {
return absl::UnimplementedError("FakeDevice::Sync()");
}
Allocator* GetAllocator(AllocatorAttributes attr) override { return nullptr; }
static std::unique_ptr<Device> Make(const std::string& name,
const std::string& type) {
DeviceAttributes device_attributes;
device_attributes.set_name(name);
device_attributes.set_device_type(DeviceType(type).type());
return std::make_unique<FakeDevice>(device_attributes);
}
};
const char* kHostName = "/job:worker/replica:0/task:0/device:CPU:0";
const char* kDeviceName = "/job:worker/replica:0/task:0/device:GPU:0";
absl::Status IncreaseDynamismForAutoJit(const Scope& s,
std::unique_ptr<Graph>* result) {
std::vector<std::unique_ptr<Device>> devices;
devices.push_back(FakeDevice::Make(kDeviceName, DEVICE_GPU));
devices.push_back(FakeDevice::Make(kHostName, DEVICE_CPU));
std::unique_ptr<DeviceSet> device_set(new DeviceSet());
for (auto& device : devices) {
device_set->AddDevice(device.get());
}
auto graph = std::make_unique<Graph>(OpRegistry::Global());
SessionOptions session_options;
session_options.config.mutable_graph_options()
->mutable_optimizer_options()
->set_global_jit_level(OptimizerOptions::ON_2);
GraphOptimizationPassOptions options;
options.graph = &graph;
options.device_set = device_set.get();
options.session_options = &session_options;
// Scope::ToGraph seems to drop assigned devices, probably because it goes
// through a GraphDef. So explicitly maintain the device assignment.
std::unordered_map<std::string, std::string> assigned_device_names;
for (Node* n : s.graph()->nodes()) {
assigned_device_names[n->name()] = n->assigned_device_name();
}
TF_RETURN_IF_ERROR(s.ToGraph(graph.get()));
for (Node* n : graph->nodes()) {
n->set_assigned_device_name(assigned_device_names[n->name()]);
}
IncreaseDynamismForAutoJitPass rewriter;
TF_RETURN_IF_ERROR(rewriter.Run(options));
*result = std::move(graph);
return absl::OkStatus();
}
TEST(SliceToDynamicSliceRewriteTest, Basic) {
Scope root = Scope::NewRootScope()
.ExitOnError()
.WithAssignedDevice(kDeviceName)
.WithXlaCluster("cluster_0");
Output input = ops::Placeholder(root.WithOpName("input"), DT_FLOAT);
Output begin = ops::Placeholder(root.WithOpName("begin"), DT_INT32);
Output size = ops::Const(root.WithOpName("size"), {-1, 500});
Output slice = ops::Slice(root.WithOpName("slice"), input, begin, size);
std::unique_ptr<Graph> result;
TF_ASSERT_OK(IncreaseDynamismForAutoJit(root, &result));
const int64_t zero_64 = 0;
const int32_t zero_32 = 0;
const int64_t one_64 = 1;
auto m_input = Out(NodeWith(Op("Placeholder"), Name("input")));
auto m_begin_s64 = Out(NodeWith(
Op("Cast"), Inputs(Out(NodeWith(Op("Placeholder"), Name("begin"))))));
auto m_input_shape = Out(NodeWith(Op("Shape"), Inputs(m_input)));
auto m_slice_size_0 = Out(NodeWith(
Op("Sub"), AssignedDevice(kHostName),
Inputs(
Out(NodeWith(Op("Slice"), AssignedDevice(kHostName),
Inputs(m_input_shape, Const(zero_64), Const(one_64)))),
Out(NodeWith(Op("Slice"), AssignedDevice(kHostName),
Inputs(m_begin_s64, Const(zero_64), Const(one_64)))))));
auto m_dynamic_slice_size =
Out(NodeWith(Op("ConcatV2"), AssignedDevice(kHostName),
Inputs(m_slice_size_0, Const(static_cast<int64_t>(500)),
Const(zero_32))));
std::vector<std::string> compile_time_constant_inputs;
compile_time_constant_inputs.push_back("size");
auto m_dynamic_slice = NodeWith(
Op("Slice"), AssignedDevice(kDeviceName),
Attr(kXlaCompileTimeConstantInputsAttr, compile_time_constant_inputs),
Inputs(m_input, m_begin_s64, m_dynamic_slice_size));
Node* static_shaped_slice = testing::FindNodeByName(
result.get(), "slice/static_shaped_slice/static_shaped_slice");
ASSERT_NE(static_shaped_slice, nullptr);
EXPECT_THAT(static_shaped_slice, m_dynamic_slice);
}
TEST(SliceToDynamicSliceRewriteTest, SliceFromVector) {
Scope root = Scope::NewRootScope()
.ExitOnError()
.WithAssignedDevice(kDeviceName)
.WithXlaCluster("cluster_0");
Output input = ops::Placeholder(root.WithOpName("input"), DT_FLOAT);
Output begin = ops::Placeholder(root.WithOpName("begin"), DT_INT32);
Output size = ops::Const(root.WithOpName("size"), {-1});
Output slice = ops::Slice(root.WithOpName("slice"), input, begin, size);
std::unique_ptr<Graph> result;
TF_ASSERT_OK(IncreaseDynamismForAutoJit(root, &result));
Node* static_shaped_slice = testing::FindNodeByName(
result.get(), "slice/static_shaped_slice/static_shaped_slice");
EXPECT_NE(static_shaped_slice, nullptr);
EXPECT_THAT(result->nodes(), Not(Contains(NodeWith(Op("ConcatV2")))));
}
TEST(SliceToDynamicSliceRewriteTest, ControlDependencePreserved) {
Scope root = Scope::NewRootScope()
.ExitOnError()
.WithAssignedDevice(kDeviceName)
.WithXlaCluster("cluster_0");
Output input = ops::Placeholder(root.WithOpName("input"), DT_FLOAT);
Output begin = ops::Placeholder(root.WithOpName("begin"), DT_INT32);
Output size = ops::Const(root.WithOpName("size"), {-1, 500});
Output control_pred = ops::Placeholder(root.WithOpName("control"), DT_BOOL);
Output slice = ops::Slice(root.WithOpName("slice"), input, begin, size);
root.graph()->AddControlEdge(control_pred.node(), slice.node());
std::unique_ptr<Graph> result;
TF_ASSERT_OK(IncreaseDynamismForAutoJit(root, &result));
Node* static_shaped_slice = testing::FindNodeByName(
result.get(), "slice/static_shaped_slice/static_shaped_slice");
ASSERT_NE(static_shaped_slice, nullptr);
EXPECT_THAT(static_shaped_slice,
NodeWith(Op("Slice"),
CtrlDeps(NodeWith(Op("Placeholder"), Name("control")))));
}
int64_t ToInt64(int v) { return static_cast<int64_t>(v); }
TEST(SliceToDynamicSliceRewriteTest, Int64Indices) {
Scope root = Scope::NewRootScope()
.ExitOnError()
.WithAssignedDevice(kDeviceName)
.WithXlaCluster("cluster_0");
Output input = ops::Placeholder(root.WithOpName("input"), DT_FLOAT);
Output begin = ops::Placeholder(root.WithOpName("begin"), DT_INT64);
Output size =
ops::Const(root.WithOpName("size"), {ToInt64(-1), ToInt64(500)});
Output slice = ops::Slice(root.WithOpName("slice"), input, begin, size);
std::unique_ptr<Graph> result;
TF_ASSERT_OK(IncreaseDynamismForAutoJit(root, &result));
EXPECT_THAT(result->nodes(), Not(Contains(NodeWith(Op("Cast")))));
}
TEST(SliceToDynamicSliceRewriteTest, DontRewriteInvalidSlice) {
Scope root = Scope::NewRootScope()
.ExitOnError()
.WithAssignedDevice(kDeviceName)
.WithXlaCluster("cluster_0");
Output input = ops::Placeholder(root.WithOpName("input"), DT_FLOAT);
Output begin = ops::Placeholder(root.WithOpName("begin"), DT_INT32);
// The shape refiner throws an error if we use a bogus constant value for
// size. So we first use a Placeholder to placate the shape refiner, and
// later replace it with a bogus constant.
Output size_placeholder =
ops::Placeholder(root.WithOpName("size_placeholder"), DT_INT32);
Output slice =
ops::Slice(root.WithOpName("slice"), input, begin, size_placeholder);
Output size = ops::Const(root.WithOpName("size"), {-8, 500});
TF_ASSERT_OK(root.graph()->UpdateEdge(/*new_src=*/size.node(),
/*new_src_index=*/0,
/*dst=*/slice.node(), /*dst_index=*/2));
std::unique_ptr<Graph> result;
TF_ASSERT_OK(IncreaseDynamismForAutoJit(root, &result));
EXPECT_THAT(result->nodes(),
Not(Contains(NodeWith(Op("Slice"),
Attr(kXlaCompileTimeConstantInputsAttr)))));
}
TEST(SliceToDynamicSliceRewriteTest, DontRewriteUnclusteredSlice) {
Scope root =
Scope::NewRootScope().ExitOnError().WithAssignedDevice(kDeviceName);
Output input = ops::Placeholder(root.WithOpName("input"), DT_FLOAT);
Output begin = ops::Placeholder(root.WithOpName("begin"), DT_INT32);
Output size = ops::Const(root.WithOpName("size"), {-1, 500});
Output slice = ops::Slice(root.WithOpName("slice"), input, begin, size);
std::unique_ptr<Graph> result;
TF_ASSERT_OK(IncreaseDynamismForAutoJit(root, &result));
EXPECT_THAT(result->nodes(),
Not(Contains(NodeWith(Op("Slice"),
Attr(kXlaCompileTimeConstantInputsAttr)))));
}
TEST(SliceToDynamicSliceRewriteTest, DontRewriteSliceWithNonConstSize) {
Scope root = Scope::NewRootScope()
.ExitOnError()
.WithAssignedDevice(kDeviceName)
.WithXlaCluster("cluster_0");
Output input = ops::Placeholder(root.WithOpName("input"), DT_FLOAT);
Output begin = ops::Placeholder(root.WithOpName("begin"), DT_INT64);
Output size = ops::Placeholder(root.WithOpName("size"), DT_INT64);
Output slice = ops::Slice(root.WithOpName("slice"), input, begin, size);
std::unique_ptr<Graph> result;
TF_ASSERT_OK(IncreaseDynamismForAutoJit(root, &result));
EXPECT_THAT(result->nodes(),
Not(Contains(NodeWith(Op("Slice"),
Attr(kXlaCompileTimeConstantInputsAttr)))));
}
TEST(SliceToDynamicSliceRewriteTest, ScalarSlice) {
Scope root = Scope::NewRootScope()
.ExitOnError()
.WithAssignedDevice(kDeviceName)
.WithXlaCluster("cluster_0");
Output input = ops::Placeholder(root.WithOpName("input"), DT_FLOAT);
Output begin = ops::Placeholder(root.WithOpName("begin"), DT_INT64);
Output size = ops::Const<int64_t>(root.WithOpName("size"), {});
Output slice = ops::Slice(root.WithOpName("slice"), input, begin, size);
std::unique_ptr<Graph> result;
TF_ASSERT_OK(IncreaseDynamismForAutoJit(root, &result));
Node* static_shaped_slice = testing::FindNodeByName(
result.get(), "slice/static_shaped_slice/static_shaped_slice");
ASSERT_NE(static_shaped_slice, nullptr);
EXPECT_THAT(static_shaped_slice,
NodeWith(Op("Slice"), Attr(kXlaCompileTimeConstantInputsAttr),
Inputs(_, _, Out(NodeWith(Name(size.node()->name()))))));
}
TEST(SliceToDynamicSliceRewriteTest, IndicesNotVector) {
Scope root = Scope::NewRootScope()
.ExitOnError()
.WithAssignedDevice(kDeviceName)
.WithXlaCluster("cluster_0");
auto ToInt64 = [](int v) { return static_cast<int64_t>(v); };
Output input = ops::Placeholder(root.WithOpName("input"), DT_FLOAT);
Output begin = ops::Placeholder(root.WithOpName("begin"), DT_INT64);
// The C++ node bindings immediately error out when we try construct a bogus
// slice so we first use a placeholder to construct the Slice and then replace
// the input.
Output size_placeholder = ops::Placeholder(root.WithOpName("size"), DT_INT64);
Output slice =
ops::Slice(root.WithOpName("slice"), input, begin, size_placeholder);
Output size =
ops::Const(root.WithOpName("size"), {{ToInt64(-1)}, {ToInt64(500)}});
TF_ASSERT_OK(root.graph()->UpdateEdge(size.node(), 0, slice.node(), 2));
std::unique_ptr<Graph> result;
TF_ASSERT_OK(IncreaseDynamismForAutoJit(root, &result));
EXPECT_THAT(result->nodes(),
Not(Contains(NodeWith(Op("Slice"),
Attr(kXlaCompileTimeConstantInputsAttr)))));
}
TEST(SliceToDynamicSliceRewriteTest, SliceWithSliceInput) {
Scope root = Scope::NewRootScope()
.ExitOnError()
.WithAssignedDevice(kDeviceName)
.WithXlaCluster("cluster_0");
Output input = ops::Placeholder(root.WithOpName("input"), DT_FLOAT);
Output begin = ops::Placeholder(root.WithOpName("begin"), DT_INT32);
Output size_a = ops::Const(root.WithOpName("size_a"), {-1, 500});
Output slice = ops::Slice(root.WithOpName("slice"), input, begin, size_a);
Output size_b = ops::Const(root.WithOpName("size_a"), {-1, 200});
Output slice_with_slice_input = ops::Slice(
root.WithOpName("slice_with_slice_input"), slice, begin, size_b);
std::unique_ptr<Graph> result;
TF_ASSERT_OK(IncreaseDynamismForAutoJit(root, &result));
Node* static_shaped_slice = testing::FindNodeByName(
result.get(),
"slice_with_slice_input/static_shaped_slice/static_shaped_slice");
ASSERT_NE(static_shaped_slice, nullptr);
EXPECT_EQ(static_shaped_slice->output_type(0), DT_FLOAT)
<< "Expected DT_FLOAT, was "
<< DataType_Name(static_shaped_slice->output_type(0));
EXPECT_THAT(
static_shaped_slice,
NodeWith(
Op("Slice"),
Inputs(Out(NodeWith(
Op("Slice"),
Name("slice/static_shaped_slice/static_shaped_slice"))),
_, _)));
}
TEST(SliceToDynamicSliceRewriteTest, SliceWithSliceBegin) {
Scope root = Scope::NewRootScope()
.ExitOnError()
.WithAssignedDevice(kDeviceName)
.WithXlaCluster("cluster_0");
Output input_float =
ops::Placeholder(root.WithOpName("input_float"), DT_FLOAT);
Output input_i64 = ops::Placeholder(root.WithOpName("input_i64"), DT_INT64);
Output begin_begin =
ops::Placeholder(root.WithOpName("begin_begin"), DT_INT32);
Output begin_size = ops::Const(root.WithOpName("begin_size"), {-1});
Output begin =
ops::Slice(root.WithOpName("begin"), input_i64, begin_begin, begin_size);
Output size =
ops::Const(root.WithOpName("size"), {ToInt64(-1), ToInt64(200)});
Output slice_with_slice_begin = ops::Slice(
root.WithOpName("slice_with_slice_begin"), input_float, begin, size);
std::unique_ptr<Graph> result;
TF_ASSERT_OK(IncreaseDynamismForAutoJit(root, &result));
Node* static_shaped_slice = testing::FindNodeByName(
result.get(),
"slice_with_slice_begin/static_shaped_slice/static_shaped_slice");
ASSERT_NE(static_shaped_slice, nullptr);
EXPECT_EQ(static_shaped_slice->output_type(0), DT_FLOAT)
<< "Expected DT_FLOAT, was "
<< DataType_Name(static_shaped_slice->output_type(0));
EXPECT_THAT(
static_shaped_slice,
NodeWith(
Op("Slice"),
Inputs(_,
Out(NodeWith(
Op("Slice"),
Name("begin/static_shaped_slice/static_shaped_slice"))),
_)));
}
// New constants being created need to have control dependencies copied to
// ensure correct control flow analysis in TF V2.
TEST(SliceToDynamicSliceRewriteTest, WithControlDepsToConstant) {
Scope root = Scope::NewRootScope()
.ExitOnError()
.WithAssignedDevice(kDeviceName)
.WithXlaCluster("cluster_0");
Output input = ops::Placeholder(root.WithOpName("input"), DT_FLOAT);
Output begin = ops::Placeholder(root.WithOpName("begin"), DT_INT32);
Output size = ops::Const(root.WithOpName("size"), {-1});
Output slice = ops::Slice(root.WithOpName("slice"), input, begin, size);
// Add an additional dependency that should still exist in with the new size
// variables.
Output dependency = ops::Placeholder(root.WithOpName("dependency"), DT_BOOL);
root.graph()->AddControlEdge(dependency.node(), size.node());
std::unique_ptr<Graph> result;
TF_ASSERT_OK(IncreaseDynamismForAutoJit(root, &result));
// Check that the new constants have control dependencies.
Node* const_0 = testing::FindNodeByName(result.get(),
"slice/static_shaped_slice/const_0");
EXPECT_NE(const_0, nullptr);
EXPECT_THAT(const_0,
NodeWith(Op("Const"), CtrlDeps(NodeWith(Op("Placeholder"),
Name("dependency")))));
}
TEST(SliceToDynamicSliceRewriteTest, DontRewriteSliceWithConstBegin) {
Scope root = Scope::NewRootScope()
.ExitOnError()
.WithAssignedDevice(kDeviceName)
.WithXlaCluster("cluster_0");
Output input = ops::Placeholder(root.WithOpName("input"), DT_FLOAT);
Output begin = ops::Const(root.WithOpName("begin"), {10, 10});
Output size = ops::Const(root.WithOpName("size"), {-1, 500});
Output slice = ops::Slice(root.WithOpName("slice"), input, begin, size);
std::unique_ptr<Graph> result;
TF_ASSERT_OK(IncreaseDynamismForAutoJit(root, &result));
Node* slice_node = testing::FindNodeByName(result.get(), "slice");
EXPECT_THAT(slice_node,
NodeWith(Op("Slice"), Inputs(Out(NodeWith(Op("Placeholder"))),
Out(NodeWith(Op("Const"))),
Out(NodeWith(Op("Const"))))));
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/jit_compilation_pass_registration.cc
+84
View File
@@ -0,0 +1,84 @@
/* 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/jit/build_xla_ops_pass.h"
#include "tensorflow/compiler/jit/clone_constants_for_better_clustering.h"
#include "tensorflow/compiler/jit/cluster_scoping_pass.h"
#include "tensorflow/compiler/jit/encapsulate_subgraphs_pass.h"
#include "tensorflow/compiler/jit/encapsulate_xla_computations_pass.h"
#include "tensorflow/compiler/jit/force_xla_constants_on_host_pass.h"
#include "tensorflow/compiler/jit/increase_dynamism_for_auto_jit_pass.h"
#include "tensorflow/compiler/jit/mark_for_compilation_pass.h"
#include "tensorflow/compiler/jit/partially_decluster_pass.h"
#include "tensorflow/compiler/jit/report_clustering_info_pass.h"
#include "tensorflow/core/common_runtime/optimization_registry.h"
namespace tensorflow {
// PRE_PLACEMENT passes:
// EncapsulateXlaComputationsPass rewrites computations generated by the
// xla.compile() Python code into XlaLaunch nodes.
REGISTER_OPTIMIZATION(OptimizationPassRegistry::PRE_PLACEMENT, 36,
EncapsulateXlaComputationsPass);
// from
// tensorflow/compiler/tf2xla/functionalize_control_flow_pass_registration.cc
// FunctionalizeControlFlowPass: 27
//
// This pass looks at the graph and all associated FunctionDefs, and turns
// traditional control flow structure (Switch/Merge/etc.) into functional
// control flow structure (XlaIf/XlaWhile). Following passes must
// handle those FunctionDef correctly.
// POST_REWRITE_FOR_EXEC passes that support auto-clustering to enable XLA:
REGISTER_OPTIMIZATION(OptimizationPassRegistry::POST_REWRITE_FOR_EXEC, 5,
CloneConstantsForBetterClusteringPass);
REGISTER_OPTIMIZATION(OptimizationPassRegistry::POST_REWRITE_FOR_EXEC, 9,
ClusterScopingPass);
REGISTER_OPTIMIZATION(OptimizationPassRegistry::POST_REWRITE_FOR_EXEC, 10,
MarkForCompilationPass);
REGISTER_OPTIMIZATION(OptimizationPassRegistry::POST_REWRITE_FOR_EXEC, 12,
ForceXlaConstantsOnHostPass);
REGISTER_OPTIMIZATION(OptimizationPassRegistry::POST_REWRITE_FOR_EXEC, 20,
IncreaseDynamismForAutoJitPass);
REGISTER_OPTIMIZATION(OptimizationPassRegistry::POST_REWRITE_FOR_EXEC, 30,
PartiallyDeclusterPass);
// ReportClusteringInfoPass pass needs to run after all of the auto-clustering
// passes have run but before encapsulation has run. This way it can easily
// compute a summary of the clustering decisions we made and broadcast it via
// xla_activity_listener.
REGISTER_OPTIMIZATION(OptimizationPassRegistry::POST_REWRITE_FOR_EXEC, 40,
ReportClusteringInfoPass);
// The EncapsulateSubgraphs pass must run after the MarkForCompilationPass. We
// also need to run it after the graph been rewritten to have _Send nodes added
// for fetches. Before the _Send nodes are added, fetch nodes are identified by
// name, and encapsulation might remove that node from the graph.
REGISTER_OPTIMIZATION(OptimizationPassRegistry::POST_REWRITE_FOR_EXEC, 50,
EncapsulateSubgraphsPass);
// Must run after EncapsulateSubgraphsPass.
REGISTER_OPTIMIZATION(OptimizationPassRegistry::POST_REWRITE_FOR_EXEC, 60,
BuildXlaOpsPass);
} // namespace tensorflow
tensorflow/compiler/jit/kernels/BUILD
+91
View File
@@ -0,0 +1,91 @@
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//tensorflow/compiler/tf2xla:internal"],
licenses = ["notice"],
)
XLA_OPS_DEPS = [
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/synchronization",
"//tensorflow/compiler/jit:common",
"//tensorflow/compiler/jit:compilation_passes",
"//tensorflow/compiler/jit:flags",
"//tensorflow/compiler/jit:xla_activity_listener",
"//tensorflow/compiler/jit:xla_activity_proto_cc",
"//tensorflow/compiler/jit:device_compiler",
"//tensorflow/compiler/jit:variable_info",
"//tensorflow/compiler/jit:variable_info_util",
"//tensorflow/compiler/jit:xla_device_no_jit_rewrite_registration",
"//tensorflow/compiler/jit:xla_cluster_util",
"//tensorflow/compiler/jit:xla_host_recv_device_context",
"//tensorflow/compiler/jit:xla_host_send_device_context",
"//tensorflow/compiler/jit:xla_launch_util",
"//tensorflow/compiler/tf2xla:common",
"//tensorflow/compiler/tf2xla:tf2xla_util",
"//tensorflow/compiler/tf2xla:xla_compiler",
"//tensorflow/compiler/tf2xla:xla_helpers",
"//tensorflow/compiler/tf2xla:xla_op_registry",
"@xla//xla:executable_run_options",
"@xla//xla:status_macros",
"@xla//xla/client:client_library",
"@xla//xla/client:local_client",
"@xla//xla/service:compiler",
"@xla//xla/service/gpu:gpu_executable_run_options",
"//tensorflow/core:core_cpu_internal",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:state_ops_op_lib",
"//tensorflow/core/platform:stream_executor_no_cuda",
"//tensorflow/core/profiler/lib:traceme",
"@xla//xla/stream_executor/integrations:tf_allocator_adapter",
"@com_google_absl//absl/types:optional",
]
# Linked by tensorflow core, without registration of jit compilation passes.
cc_library(
name = "xla_ops_no_jit_rewrite_registration",
srcs = ["xla_ops.cc"],
hdrs = ["xla_ops.h"],
deps = XLA_OPS_DEPS + [
"//tensorflow/compiler/jit:device_compilation_cache",
"//tensorflow/compiler/jit:device_compilation_profiler",
"//tensorflow/compiler/jit:pjrt_compile_util",
"//tensorflow/compiler/jit:tf_graph_to_hlo_compiler",
"//tensorflow/compiler/jit:tf_to_hlo_compiler",
"//tensorflow/compiler/jit:xla_compile_util",
"//tensorflow/core:framework_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/platform:refcount",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/container:node_hash_map",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:string_view",
"@com_google_absl//absl/types:span",
"@tsl//tsl/profiler/lib:traceme",
"@xla//xla/hlo/ir:hlo",
"@xla//xla/pjrt:pjrt_client",
"@xla//xla/service:executable",
"@xla//xla/stream_executor/host:host_platform_id",
"@xla//xla/tsl/concurrency:async_value",
"@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
],
alwayslink = 1,
)
cc_library(
name = "xla_ops",
hdrs = ["xla_ops.h"],
deps = XLA_OPS_DEPS + [
":xla_ops_no_jit_rewrite_registration",
"//tensorflow/compiler/jit:jit_compilation_passes",
"//tensorflow/core:protos_all_cc",
],
alwayslink = 1,
)
tensorflow/compiler/jit/kernels/xla_ops.cc
+1046
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/jit/kernels/xla_ops.h
+142
View File
@@ -0,0 +1,142 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_JIT_KERNELS_XLA_OPS_H_
#define TENSORFLOW_COMPILER_JIT_KERNELS_XLA_OPS_H_
#include <atomic>
#include <vector>
#include "tensorflow/compiler/jit/device_compiler.h"
#include "tensorflow/compiler/jit/xla_device.h"
#include "tensorflow/compiler/jit/xla_launch_util.h"
#include "tensorflow/compiler/jit/xla_platform_info.h"
#include "xla/stream_executor/integrations/tf_allocator_adapter.h"
#include "tensorflow/core/framework/allocator.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/util/stream_executor_util.h"
namespace tensorflow {
// XlaLocalLaunchBase is almost the same as XlaLocalLaunchOp.
// The only difference is that it does not require arguments to follow
// the "constants, then regular args, then resources" order.
// It takes vectors of constant and resource arguments explicitly.
// It does not have corresponding OpDef because it is never present
// in the GraphDef.
// Currently, it is used by eager runtime. FunctionLibraryRuntime creates
// this kernel when asked to create a kernel for an XLA-compiled function.
//
// `has_ref_vars`: whether the input computation can have reference variables.
// TODO(cheshire): instead derive this information from the input graph.
class XlaLocalLaunchBase : public AsyncOpKernel {
public:
XlaLocalLaunchBase(OpKernelConstruction* ctx,
const std::vector<int>& constants,
const std::vector<int>& resources,
const NameAttrList& function, bool has_ref_vars);
XlaLocalLaunchBase(const XlaLocalLaunchBase&) = delete;
XlaLocalLaunchBase& operator=(const XlaLocalLaunchBase&) = delete;
~XlaLocalLaunchBase() override = default;
void ComputeAsync(OpKernelContext* ctx, DoneCallback done) override;
protected:
// Indexes of compile-time constant inputs
const std::vector<int> constants_;
// Indexes of resource inputs
const std::vector<int> resources_;
const NameAttrList function_;
const XlaPlatformInfo platform_info_;
bool has_ref_vars_;
};
// XlaLocalLaunchOp is used to replace a region of the TensorFlow graph
// which will be compiled and executed using XLA. The XlaLocalLaunchOp is
// responsible for handling interactions with the TensorFlow executor.
// Once all inputs are present, and their shapes are known, the op can
// use a 'DeviceCompiler' to compile and execute code which is specific
// to the shapes of input Tensors.
// XlaLocalLaunchOp uses xla::LocalClient::Compile() and
// xla::LocalExecutable::Run(), and passes arguments into/out of XLA in device
// memory.
class XlaLocalLaunchOp : public XlaLocalLaunchBase {
public:
explicit XlaLocalLaunchOp(OpKernelConstruction* ctx);
~XlaLocalLaunchOp() override;
private:
XlaLocalLaunchOp(const XlaLocalLaunchOp&) = delete;
void operator=(const XlaLocalLaunchOp&) = delete;
};
class XlaCompileOp : public OpKernel {
public:
explicit XlaCompileOp(OpKernelConstruction* ctx);
void Compute(OpKernelContext* ctx) override;
private:
// Indexes of compile-time constant inputs
const std::vector<int> constants_;
// Indexes of resource inputs
const std::vector<int> resources_;
const NameAttrList function_;
XlaPlatformInfo platform_info_;
const bool must_compile_;
// Whether the graph has TF reference variables.
const bool has_ref_vars_;
// cannot_compile_cluster_ is set to true if XLA returns an Unimplemented
// error when compiling the cluster this _XlaCompile is supposed to compile.
// If `cannot_compile_cluster_` is true then we avoid compiling this cluster
// on any future calls to _XlaCompile.
bool cannot_compile_cluster_ TF_GUARDED_BY(cannot_compile_cluster_mu_) =
false;
mutex cannot_compile_cluster_mu_;
};
class XlaRunOp : public OpKernel {
public:
explicit XlaRunOp(OpKernelConstruction* ctx);
void Compute(OpKernelContext* ctx) override;
private:
const XlaPlatformInfo platform_info_;
};
class XlaMergeOp : public OpKernel {
public:
explicit XlaMergeOp(OpKernelConstruction* ctx);
void Compute(OpKernelContext* ctx) override;
};
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_KERNELS_XLA_OPS_H_
tensorflow/compiler/jit/mark_for_compilation_pass.cc
+2377
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/jit/mark_for_compilation_pass.h
+63
View File
@@ -0,0 +1,63 @@
/* 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.
==============================================================================*/
// An optimization passes that marks nodes that are to be compiled with
// attribute kXlaClusterAttr. Nodes with the same cluster ID will be compiled
// together.
#ifndef TENSORFLOW_COMPILER_JIT_MARK_FOR_COMPILATION_PASS_H_
#define TENSORFLOW_COMPILER_JIT_MARK_FOR_COMPILATION_PASS_H_
#include "absl/container/flat_hash_set.h"
#include "tensorflow/compiler/jit/compilability_check_util.h"
#include "tensorflow/core/common_runtime/optimization_registry.h"
namespace tensorflow {
// The attribute that marks nodes to be grouped into functions by the
// encapsulate subgraphs pass.
extern const char* const kXlaClusterAttr;
// Marks a subset of nodes in the graph which are to be clustered
// with an attribute _XlaCluster=<cluster id> so they are picked up by the
// EncapsulateSubgraphsPass.
class MarkForCompilationPass : public GraphOptimizationPass {
public:
MarkForCompilationPass() = default;
absl::Status Run(const GraphOptimizationPassOptions& options) override;
private:
absl::Status RunForTest(const GraphOptimizationPassOptions& options,
bool disable_deadness_analysis,
bool deterministic_cluster_names);
friend class MarkForCompilationPassTestHelper;
};
absl::flat_hash_map<std::string, std::vector<std::string>>* GetAllowlistTable();
namespace testing {
// DO NOT USE IN PRODUCTION.
//
// Resets some internal state to let us write reliable unit tests.
void ResetClusterSequenceNumber();
// Return a list of operation that we choose not to put into the allowlist.
absl::flat_hash_set<std::string> GetKnownXLAAllowlistOp();
} // namespace testing
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_MARK_FOR_COMPILATION_PASS_H_
tensorflow/compiler/jit/mark_for_compilation_pass_test.cc
+2024
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/jit/mark_for_compilation_pass_test_helper.cc
+81
View File
@@ -0,0 +1,81 @@
/* Copyright 2018 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/jit/mark_for_compilation_pass_test_helper.h"
#include <memory>
#include <vector>
#include "tensorflow/compiler/jit/cluster_scoping_pass.h"
#include "tensorflow/core/common_runtime/device_factory.h"
#include "tensorflow/core/lib/gtl/cleanup.h"
#include "tensorflow/core/public/session_options.h"
namespace tensorflow {
/*static*/ absl::Status MarkForCompilationPassTestHelper::MarkForCompilation(
std::unique_ptr<Graph>* graph, FunctionLibraryDefinition* flib_def,
MarkForCompilationPassTestHelper::Options options) {
// Assign all unassigned nodes to the CPU device.
static const char* kCpuDevice = "/job:localhost/replica:0/task:0/cpu:0";
for (Node* n : (*graph)->nodes()) {
if (n->assigned_device_name().empty()) {
n->set_assigned_device_name(kCpuDevice);
}
}
SessionOptions session_options;
if (options.enable_global_jit) {
session_options.config.mutable_graph_options()
->mutable_optimizer_options()
->set_global_jit_level(OptimizerOptions::ON_2);
}
if (!options.session_name.empty()) {
session_options.config.mutable_experimental()
->mutable_session_metadata()
->set_name(options.session_name);
}
// Call AddDevices to register the XLA devices.
//
// It may be worth refactoring out XlaOpRegistry::RegisterCompilationDevice to
// make this more direct, but probably not worth it solely for this test.
std::vector<std::unique_ptr<Device>> devices;
TF_RETURN_IF_ERROR(DeviceFactory::AddDevices(session_options, "", &devices));
GraphOptimizationPassOptions opt_options;
opt_options.graph = graph;
opt_options.session_options = &session_options;
opt_options.flib_def = flib_def;
if (options.enable_cluster_scoping) {
ClusterScopingPass cluster_scoping_pass;
TF_RETURN_IF_ERROR(cluster_scoping_pass.Run(opt_options));
}
MarkForCompilationPass mark_for_compilation_pass;
return mark_for_compilation_pass.RunForTest(
opt_options,
/*disable_deadness_analysis=*/options.disable_deadness_analysis,
/*deterministic_cluster_names=*/options.deterministic_cluster_names);
}
/*static*/ absl::Status MarkForCompilationPassTestHelper::MarkForCompilation(
std::unique_ptr<Graph>* graph,
MarkForCompilationPassTestHelper::Options options) {
FunctionDefLibrary flib;
FunctionLibraryDefinition flib_def((*graph)->op_registry(), flib);
return MarkForCompilation(graph, &flib_def, options);
}
} // namespace tensorflow
tensorflow/compiler/jit/mark_for_compilation_pass_test_helper.h
+85
View File
@@ -0,0 +1,85 @@
/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_JIT_MARK_FOR_COMPILATION_PASS_TEST_HELPER_H_
#define TENSORFLOW_COMPILER_JIT_MARK_FOR_COMPILATION_PASS_TEST_HELPER_H_
#include <memory>
#include <string>
#include <utility>
#include "tensorflow/compiler/jit/mark_for_compilation_pass.h"
namespace tensorflow {
class MarkForCompilationPassTestHelper {
public:
struct Options {
bool enable_global_jit;
bool disable_deadness_analysis;
bool enable_cluster_scoping;
bool deterministic_cluster_names;
std::string session_name; // ConfigProto.Experimental.SessionMetadata.name
Options()
: enable_global_jit(true),
disable_deadness_analysis(true),
enable_cluster_scoping(true),
deterministic_cluster_names(false) {}
Options WithNoGlobalJit() {
Options copy = *this;
copy.enable_global_jit = false;
return copy;
}
Options WithDeadnessAnalysis() {
Options copy = *this;
copy.disable_deadness_analysis = false;
return copy;
}
Options WithNoClusterScoping() {
Options copy = *this;
copy.enable_cluster_scoping = false;
return copy;
}
Options WithDeterministicClusterNames() {
Options copy = *this;
copy.deterministic_cluster_names = true;
return copy;
}
Options WithSessionName(std::string name) {
Options copy = *this;
copy.session_name = std::move(name);
return copy;
}
};
// Runs the MarkForCompilation pass on `graph` after assigning all nodes in
// `graph` to the CPU device. To make testing easier, ignores device
// registration and _XlaCompile attributes.
static absl::Status MarkForCompilation(std::unique_ptr<Graph>* graph,
FunctionLibraryDefinition* flib_def,
Options options = Options());
// Like `MarkForCompilation` but creates `flib_def` from the op registry.
static absl::Status MarkForCompilation(std::unique_ptr<Graph>* graph,
Options options = Options());
};
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_MARK_FOR_COMPILATION_PASS_TEST_HELPER_H_
tensorflow/compiler/jit/node_matchers.cc
+562
View File
@@ -0,0 +1,562 @@
/* Copyright 2018 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/jit/node_matchers.h"
#include <cstdint>
#include <iterator>
#include <map>
#include <optional>
#include <ostream>
#include <sstream>
#include <string>
#include <utility>
#include <vector>
#include "absl/algorithm/container.h"
#include "absl/log/check.h"
#include "absl/log/log.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.h"
#include "absl/strings/str_split.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "tensorflow/cc/framework/ops.h"
#include "tensorflow/core/framework/attr_value_util.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/graph/graph_node_util.h"
namespace tensorflow {
namespace testing {
namespace matchers {
namespace {
using impl::NodeMatcherProperties;
using impl::OutEdge;
std::string IndentAllButFirstLine(absl::string_view text) {
std::vector<std::string> lines = absl::StrSplit(text, '\n');
for (int i = 1; i < lines.size(); i++) {
lines[i].insert(0, " ");
}
return absl::StrJoin(lines, "\n");
}
template <typename T>
bool CompareTensor(const Tensor& actual, const Tensor& expected,
::testing::MatchResultListener* listener) {
if (actual.NumElements() != expected.NumElements()) {
if (listener->IsInterested()) {
*listener << "\nwas looking for tensor with " << expected.NumElements()
<< " elements, found tensor with " << actual.NumElements()
<< " elements";
return false;
}
}
for (int64_t i = 0, e = actual.NumElements(); i < e; i++) {
if (actual.flat<T>()(i) != expected.flat<T>()(i)) {
*listener << "\nmismatch in constant tensor at index " << i
<< " expected = " << expected.flat<T>()(i)
<< " actual = " << actual.flat<T>()(i);
return false;
}
}
return true;
}
bool MatchAndExplainTensor(const Tensor& tensor, const Tensor& expected_tensor,
::testing::MatchResultListener* listener) {
if (tensor.dtype() != expected_tensor.dtype()) {
if (listener->IsInterested()) {
*listener << "\nexpected tensor of type "
<< DataType_Name(expected_tensor.dtype())
<< " but found one of type " << DataType_Name(tensor.dtype());
return false;
}
}
switch (tensor.dtype()) {
case DT_HALF:
return CompareTensor<Eigen::half>(tensor, expected_tensor, listener);
case DT_FLOAT:
return CompareTensor<float>(tensor, expected_tensor, listener);
case DT_DOUBLE:
return CompareTensor<double>(tensor, expected_tensor, listener);
case DT_INT8:
return CompareTensor<int8_t>(tensor, expected_tensor, listener);
case DT_INT16:
return CompareTensor<int16_t>(tensor, expected_tensor, listener);
case DT_INT32:
return CompareTensor<int32_t>(tensor, expected_tensor, listener);
case DT_INT64:
return CompareTensor<int64_t>(tensor, expected_tensor, listener);
case DT_UINT8:
return CompareTensor<uint8_t>(tensor, expected_tensor, listener);
case DT_UINT16:
return CompareTensor<uint16_t>(tensor, expected_tensor, listener);
case DT_UINT32:
return CompareTensor<uint32_t>(tensor, expected_tensor, listener);
case DT_UINT64:
return CompareTensor<uint64_t>(tensor, expected_tensor, listener);
default:
LOG(FATAL) << "Unsupported dtype " // Crash ok: testonly.
<< DataType_Name(tensor.dtype());
}
}
struct NodeMatcher : public ::testing::MatcherInterface<const Node*> {
bool MatchAndExplain(
const Node* node,
::testing::MatchResultListener* listener) const override {
if (op && node->type_string() != *op) {
if (listener->IsInterested()) {
*listener << "\nexpected op " << *op << " but found "
<< node->type_string();
}
return false;
}
if (assigned_device && node->assigned_device_name() != *assigned_device) {
if (listener->IsInterested()) {
*listener << "\nexpected assigned_device " << *assigned_device
<< " but found \"" << node->assigned_device_name() << "\"";
}
return false;
}
if (name && node->name() != *name) {
if (listener->IsInterested()) {
*listener << "\nexpected name " << *name << " but found "
<< node->name();
}
return false;
}
if (constant_value) {
const TensorProto* proto = nullptr;
if (!TryGetNodeAttr(node->def(), "value", &proto)) {
if (listener->IsInterested()) {
*listener << "\ncould not find \"value\" attribute in node";
}
return false;
}
Tensor tensor(proto->dtype());
if (!tensor.FromProto(*proto)) {
if (listener->IsInterested()) {
*listener << "\ncould not convert TensorProto in \"value\" attribute "
"to Tensor";
}
return false;
}
if (!MatchAndExplainTensor(/*tensor=*/tensor,
/*expected_tensor=*/*constant_value,
listener)) {
return false;
}
}
if (input_matchers) {
if (input_matchers->size() != node->num_inputs()) {
if (listener->IsInterested()) {
*listener << "\nexpected " << input_matchers->size()
<< " inputs but node has " << node->num_inputs();
}
return false;
}
for (int input_idx = 0, e = input_matchers->size(); input_idx < e;
input_idx++) {
if (!MatchAndExplainInput(node, input_idx, listener)) {
return false;
}
}
}
std::vector<const Node*> control_deps;
for (const Edge* e : node->in_edges()) {
if (e->IsControlEdge()) {
control_deps.push_back(e->src());
}
}
::testing::StringMatchResultListener inner_listener;
if (control_dep_set &&
!control_dep_set->MatchAndExplain(control_deps, &inner_listener)) {
if (listener->IsInterested()) {
std::string explanation = inner_listener.str();
if (!explanation.empty()) {
explanation = absl::StrCat(", ", explanation, ",");
}
*listener << "ctrl_deps" << explanation << " does not match expected: ";
control_dep_set->DescribeTo(listener->stream());
}
return false;
}
const AttrValueMap attr_value_map = node->def().attr();
for (const auto& attr_kv_pair : attrs) {
auto it = attr_value_map.find(attr_kv_pair.first);
if (it == attr_value_map.end()) {
if (listener->IsInterested()) {
*listener << "did not find attribute named \"" << attr_kv_pair.first
<< "\" in node";
}
return false;
}
if (attr_kv_pair.second &&
!AreAttrValuesEqual(it->second, *attr_kv_pair.second)) {
if (listener->IsInterested()) {
*listener << "attribute named " << attr_kv_pair.first
<< " does not match value; expected: \""
<< SummarizeAttrValue(*attr_kv_pair.second)
<< "\", found: \"" << SummarizeAttrValue(it->second)
<< "\"";
}
return false;
}
}
return true;
}
void DescribeTo(::std::ostream* os) const override {
std::vector<std::string> predicates;
if (name) {
predicates.push_back(absl::StrCat("name: ", *name));
}
if (op) {
predicates.push_back(absl::StrCat("op: ", *op));
}
if (assigned_device) {
predicates.push_back(absl::StrCat("assigned device: ", *assigned_device));
}
bool printed_something = !predicates.empty();
*os << absl::StrJoin(predicates, ", ");
if (constant_value) {
printed_something = true;
*os << "constant value: " << constant_value->DebugString();
}
if (input_matchers) {
if (!input_matchers->empty()) {
printed_something = true;
*os << " with " << (input_matchers->size() == 1 ? "only " : "")
<< "input" << (input_matchers->size() == 1 ? "" : "s") << " ";
}
if (input_matchers->size() == 1) {
::std::stringstream ss;
input_matchers->front().DescribeTo(&ss);
printed_something = true;
*os << "matching " << ss.str();
} else {
int edge_idx = 0;
for (const ::testing::Matcher<OutEdge>& matcher : (*input_matchers)) {
*os << "\n [" << edge_idx << "] matching (";
::std::stringstream ss;
matcher.DescribeTo(&ss);
printed_something = true;
*os << IndentAllButFirstLine(ss.str());
*os << ")";
edge_idx++;
}
}
}
if (control_dep_set) {
printed_something = true;
*os << " and control deps ";
control_dep_set->DescribeTo(os);
}
if (!attrs.empty()) {
printed_something = true;
std::vector<std::string> attrs_str;
absl::c_transform(
attrs, std::back_inserter(attrs_str),
[](const std::pair<std::string, std::optional<AttrValue>>&
attr_kv_pair) {
return absl::StrCat(attr_kv_pair.first, "->",
attr_kv_pair.second
? SummarizeAttrValue(*attr_kv_pair.second)
: "*");
});
*os << " and attr values matching [" << absl::StrJoin(attrs_str, ", ")
<< "]";
}
if (!printed_something) {
*os << "is any node";
}
}
bool MatchAndExplainInput(const Node* node, int input_idx,
::testing::MatchResultListener* listener) const {
const Edge* edge;
if (!node->input_edge(input_idx, &edge).ok()) {
if (listener->IsInterested()) {
*listener << "\ncould not find incoming edge for input " << input_idx;
}
return false;
}
::testing::StringMatchResultListener inner_listener;
OutEdge input = {edge->src(), edge->src_output()};
if ((*input_matchers)[input_idx].MatchAndExplain(input, &inner_listener)) {
return true;
}
if (listener->IsInterested()) {
*listener << "\ninput " << input_idx << " does not match expected:\n";
(*input_matchers)[input_idx].DescribeTo(listener->stream());
std::string explanation = inner_listener.str();
if (!explanation.empty()) {
*listener << ", " << explanation;
}
}
return false;
}
std::optional<std::string> op;
std::optional<std::string> name;
std::optional<std::string> assigned_device;
std::optional<Tensor> constant_value;
std::optional<std::vector<::testing::Matcher<OutEdge>>> input_matchers;
std::optional<::testing::Matcher<absl::Span<const Node* const>>>
control_dep_set;
std::map<std::string, std::optional<AttrValue>> attrs;
};
// Matches a dst and dst_output on an input edge. Today we only use this with
// dst_output=0 but we will eventually need to support multi-output operations.
class OutEdgeMatcher : public ::testing::MatcherInterface<OutEdge> {
public:
OutEdgeMatcher(::testing::Matcher<const Node*> src_matcher, int src_oidx)
: src_matcher_(std::move(src_matcher)), src_oidx_(src_oidx) {}
bool MatchAndExplain(
OutEdge out_edge,
::testing::MatchResultListener* listener) const override {
::testing::StringMatchResultListener inner_listener;
if (!src_matcher_.MatchAndExplain(out_edge.first, &inner_listener)) {
if (listener->IsInterested()) {
*listener << "\nsource does not match expected ";
src_matcher_.DescribeTo(listener->stream());
std::string explanation = inner_listener.str();
if (!explanation.empty()) {
*listener << "\n\t" << explanation;
}
}
return false;
}
if (out_edge.second != src_oidx_) {
if (listener->IsInterested()) {
*listener << "\nexpected output slot to be " << src_oidx_
<< " but found " << out_edge.second;
}
return false;
}
return true;
}
void DescribeTo(::std::ostream* os) const override {
if (src_oidx_) {
*os << "output slot: " << src_oidx_ << ", source: (";
}
src_matcher_.DescribeTo(os);
if (src_oidx_) {
*os << ")";
}
}
private:
::testing::Matcher<const Node*> src_matcher_;
int src_oidx_;
};
} // namespace
::testing::Matcher<const Node*> impl::NodeWith(
absl::Span<const NodeMatcherProperties> props) {
NodeMatcher* matcher = new NodeMatcher();
for (const NodeMatcherProperties& prop : props) {
if (prop.name()) {
DCHECK(!matcher->name);
matcher->name = prop.name();
}
if (prop.op()) {
DCHECK(!matcher->op);
matcher->op = prop.op();
}
if (prop.constant_value()) {
DCHECK(!matcher->constant_value);
matcher->constant_value = prop.constant_value();
}
if (prop.assigned_device()) {
DCHECK(!matcher->assigned_device);
matcher->assigned_device = prop.assigned_device();
}
if (prop.inputs()) {
DCHECK(!matcher->input_matchers);
matcher->input_matchers = *prop.inputs();
}
if (prop.control_deps()) {
DCHECK(!matcher->control_dep_set);
matcher->control_dep_set =
::testing::UnorderedElementsAreArray(*prop.control_deps());
}
if (prop.attr()) {
auto insert_result = matcher->attrs.insert(*prop.attr());
DCHECK(insert_result.second);
}
}
return ::testing::MakeMatcher(matcher);
}
impl::NodeMatcherProperties Name(std::string name) {
impl::NodeMatcherProperties props;
props.set_name(std::move(name));
return props;
}
// Matches a node with op `op`.
impl::NodeMatcherProperties Op(std::string op) {
impl::NodeMatcherProperties props;
props.set_op(std::move(op));
return props;
}
// Matches a node with assigned device `assigned_device`.
impl::NodeMatcherProperties AssignedDevice(std::string assigned_device) {
impl::NodeMatcherProperties props;
props.set_assigned_device(std::move(assigned_device));
return props;
}
impl::NodeMatcherProperties impl::Inputs(
absl::Span<const ::testing::Matcher<OutEdge>> inputs) {
std::vector<::testing::Matcher<OutEdge>> inputs_vector;
absl::c_copy(inputs, std::back_inserter(inputs_vector));
impl::NodeMatcherProperties props;
props.set_inputs(std::move(inputs_vector));
return props;
}
impl::NodeMatcherProperties impl::CtrlDeps(
absl::Span<const ::testing::Matcher<const Node*>> control_deps) {
std::vector<::testing::Matcher<const Node*>> control_deps_vector;
absl::c_copy(control_deps, std::back_inserter(control_deps_vector));
impl::NodeMatcherProperties props;
props.set_control_deps(std::move(control_deps_vector));
return props;
}
std::pair<std::string, AttrValue> impl::AttrLiteralHelper(
const std::pair<std::string, bool>& bool_attr) {
AttrValue attr_value;
attr_value.set_b(bool_attr.second);
return {bool_attr.first, attr_value};
}
std::pair<std::string, AttrValue> impl::AttrLiteralHelper(
const std::pair<std::string, absl::Span<const int>>& int_list_attr) {
AttrValue attr_value;
AttrValue::ListValue* list = attr_value.mutable_list();
for (int i : int_list_attr.second) {
list->add_i(i);
}
return {int_list_attr.first, attr_value};
}
std::pair<std::string, AttrValue> impl::AttrLiteralHelper(
const std::pair<std::string, absl::Span<const std::string>>&
string_list_attr) {
AttrValue attr_value;
AttrValue::ListValue* list = attr_value.mutable_list();
for (const std::string& s : string_list_attr.second) {
list->add_s(s);
}
return {string_list_attr.first, attr_value};
}
impl::NodeMatcherProperties impl::Attr(std::pair<std::string, AttrValue> attr) {
impl::NodeMatcherProperties props;
props.set_attr(std::move(attr));
return props;
}
impl::NodeMatcherProperties impl::Attr(std::string name) {
impl::NodeMatcherProperties props;
props.set_attr({std::move(name), std::nullopt});
return props;
}
NodeMatcherProperties ConstantValue(
const ::tensorflow::Input::Initializer& val) {
CHECK_OK(val.status);
NodeMatcherProperties props;
props.set_constant_value(val.tensor);
return props;
}
::testing::Matcher<impl::OutEdge> Const(
const ::tensorflow::Input::Initializer& val) {
return Out(NodeWith(ConstantValue(val)));
}
::testing::Matcher<impl::OutEdge> Out(
int oidx, ::testing::Matcher<const Node*> node_matcher) {
return ::testing::MakeMatcher(new OutEdgeMatcher(node_matcher, oidx));
}
} // namespace matchers
Node* FindNodeByName(Graph* g, absl::string_view name) {
for (Node* n : g->nodes()) {
if (n->name() == name) {
return n;
}
}
return nullptr;
}
} // namespace testing
void PrintTo(const Node* n, ::std::ostream* os) { *os << SummarizeNode(*n); }
void PrintTo(Node* n, ::std::ostream* os) { *os << SummarizeNode(*n); }
} // namespace tensorflow
tensorflow/compiler/jit/node_matchers.h
+252
View File
@@ -0,0 +1,252 @@
/* Copyright 2018 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.
==============================================================================*/
// Provides a set of matchers for tensorflow nodes.
//
// Example usage:
//
// tensorflow::Node* node = ...;
// EXPECT_THAT(node, NodeWith(Name("name"), Op("op"),
// Inputs(Out(3, NodeWith(Name("input"))))))
//
// Matchable node properties (the expressions that go inside NodeWith(...))
// are:
//
// - Name(string): matches the node name exactly. We will probably need to
// have this take a string matcher soon in the future.
//
// - Op(string): matches the op exactly.
//
// - AssignedDevice(string): matches the assigned device exactly.
//
// - Inputs(<ordered list>): matches the list of non-control inputs to the node
// exactly (i.e. does not match a suffix or a prefix) where each element
// matches an output of a node (see Out(idx, node) below).
//
// - CtrlDeps(<unordered list>): matches the list of control dependences on the
// node exactly but in any order.
//
// - ConstantValue(tensorflow::Input::Initializer init): matches a Const node
// with the constant value `init`. Implies Op("Const").
//
// - Attr(name, value): Matches a single attribute with name `name` and value
// `value`. Right now only boolean values are supported.
//
// Overlapping node properties may not be repeated in a single NodeWith(...)
// matcher. E.g. NodeWith(Op("Foo"), Op("Bar")) will CHECK-fail. Since
// ConstantValue implies Op("Const"), a single NodeWith matcher can't have both
// ConstantValue(...) and Op(...). Multiple Attr() values can be combined as
// long as the attribute names are different.
//
// Out(idx, node) matches the `idx`'th output of a node that matches `node`.
#ifndef TENSORFLOW_COMPILER_JIT_NODE_MATCHERS_H_
#define TENSORFLOW_COMPILER_JIT_NODE_MATCHERS_H_
#include <array>
#include <string>
#include <vector>
#include "absl/algorithm/container.h"
#include "absl/strings/string_view.h"
#include "absl/types/optional.h"
#include "absl/types/span.h"
#include "tensorflow/cc/framework/ops.h"
#include "xla/hlo/testlib/test.h"
#include "tensorflow/core/graph/graph.h"
namespace tensorflow {
namespace testing {
namespace matchers {
namespace impl {
using OutEdge = std::pair<const Node*, int>;
// -----------------------------------------------------------------------------
// Implementation details.
// Properties that we match on for a particular Node. If a particular property
// is nullopt then any value for it is allowed.
class NodeMatcherProperties {
public:
using NodeSeqMatcher = std::vector<::testing::Matcher<const Node*>>;
using InputSeqMatcher = std::vector<::testing::Matcher<OutEdge>>;
using AttrKeyValuePair = std::pair<std::string, std::optional<AttrValue>>;
const std::optional<std::string>& name() const { return name_; }
const std::optional<std::string>& op() const { return op_; }
const std::optional<std::string>& assigned_device() const {
return assigned_device_;
}
const std::optional<Tensor>& constant_value() const {
return constant_value_;
}
const std::optional<InputSeqMatcher>& inputs() const {
return input_matchers_;
}
const std::optional<NodeSeqMatcher>& control_deps() const {
return control_deps_;
}
const std::optional<AttrKeyValuePair>& attr() const { return attr_; }
void set_name(std::string name) {
DCHECK(IsEmpty());
name_ = std::move(name);
}
void set_op(std::string op) {
DCHECK(IsEmpty());
op_ = std::move(op);
}
void set_assigned_device(std::string assigned_device) {
DCHECK(IsEmpty());
assigned_device_ = std::move(assigned_device);
}
void set_constant_value(Tensor constant_value) {
DCHECK(IsEmpty());
constant_value_ = std::move(constant_value);
op_ = "Const";
}
void set_inputs(InputSeqMatcher inputs) {
DCHECK(IsEmpty());
input_matchers_ = std::move(inputs);
}
void set_control_deps(NodeSeqMatcher control_deps) {
DCHECK(IsEmpty());
control_deps_ = std::move(control_deps);
}
void set_attr(AttrKeyValuePair attr) {
DCHECK(IsEmpty());
attr_ = std::move(attr);
}
bool IsEmpty() const {
return !name().has_value() && !op().has_value() && !inputs().has_value() &&
!control_deps().has_value() && !attr().has_value();
}
private:
std::optional<std::string> name_;
std::optional<std::string> op_;
std::optional<std::string> assigned_device_;
std::optional<Tensor> constant_value_;
std::optional<InputSeqMatcher> input_matchers_;
std::optional<NodeSeqMatcher> control_deps_;
std::optional<AttrKeyValuePair> attr_;
};
::testing::Matcher<const Node*> NodeWith(
absl::Span<const NodeMatcherProperties> props);
impl::NodeMatcherProperties Inputs(
absl::Span<const ::testing::Matcher<OutEdge>> inputs);
impl::NodeMatcherProperties CtrlDeps(
absl::Span<const ::testing::Matcher<const Node*>> control_deps);
impl::NodeMatcherProperties Attr(std::pair<std::string, AttrValue> attrs);
impl::NodeMatcherProperties Attr(std::string name);
std::pair<std::string, AttrValue> AttrLiteralHelper(
const std::pair<std::string, bool>& bool_attr);
std::pair<std::string, AttrValue> AttrLiteralHelper(
const std::pair<std::string, absl::Span<const int>>& int_list_attr);
std::pair<std::string, AttrValue> AttrLiteralHelper(
const std::pair<std::string, absl::Span<const std::string>>&
string_list_attr);
} // namespace impl
// -----------------------------------------------------------------------------
// Public interface.
// Matches a node with name `name`.
impl::NodeMatcherProperties Name(std::string name);
// Matches a node with op `op`.
impl::NodeMatcherProperties Op(std::string op);
// Matches a node with assigned device `assigned_device`.
impl::NodeMatcherProperties AssignedDevice(std::string assigned_device);
// Matches a node with a boolean typed attribute named `name` and with value
// `value`.
template <typename ValueTy>
impl::NodeMatcherProperties Attr(const std::string& name, ValueTy value) {
return impl::Attr({impl::AttrLiteralHelper({name, value})});
}
inline impl::NodeMatcherProperties Attr(const std::string& name) {
return impl::Attr(name);
}
// Matches a node with inputs `inputs`.
//
// `inputs` are ordered; `inputs`[i] must match input i.
template <typename... Ts>
impl::NodeMatcherProperties Inputs(Ts... inputs) {
return impl::Inputs({inputs...});
}
// Matches the `idx`'th output of a node that matches `node`.
::testing::Matcher<impl::OutEdge> Out(int oidx,
::testing::Matcher<const Node*> node);
// Matches the first output of a node that matches `node`.
inline ::testing::Matcher<impl::OutEdge> Out(
::testing::Matcher<const Node*> node) {
return Out(0, node);
}
// Matches a node with control dependences `control_deps`.
//
// `control_deps` are unordered and will match the control deps of a node in any
// order.
template <typename... Ts>
impl::NodeMatcherProperties CtrlDeps(Ts... control_deps) {
return impl::CtrlDeps({control_deps...});
}
// Matches a constant node with value `val`.
impl::NodeMatcherProperties ConstantValue(
const ::tensorflow::Input::Initializer& val);
// The main gmock matcher. See file comment for example usage.
template <typename... Ts>
::testing::Matcher<const Node*> NodeWith(Ts... args) {
std::array<impl::NodeMatcherProperties, sizeof...(Ts)> array = {args...};
return impl::NodeWith(array);
}
::testing::Matcher<impl::OutEdge> Const(
const ::tensorflow::Input::Initializer& val);
} // namespace matchers
// If `g` has a node named `name` returns it, otherwise returns null.
Node* FindNodeByName(Graph* g, absl::string_view name);
} // namespace testing
void PrintTo(const Node* n, ::std::ostream* os);
void PrintTo(Node* n, ::std::ostream* os);
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_NODE_MATCHERS_H_
tensorflow/compiler/jit/node_matchers_test.cc
+230
View File
@@ -0,0 +1,230 @@
/* Copyright 2018 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/jit/node_matchers.h"
#include <string>
#include "tensorflow/cc/framework/ops.h"
#include "tensorflow/cc/ops/array_ops.h"
#include "tensorflow/cc/ops/const_op.h"
#include "tensorflow/cc/ops/control_flow_ops.h"
#include "tensorflow/cc/ops/control_flow_ops_internal.h"
#include "tensorflow/cc/ops/math_ops.h"
namespace tensorflow {
namespace testing {
namespace {
using ::testing::_;
using testing::matchers::AssignedDevice;
using testing::matchers::Attr;
using testing::matchers::ConstantValue;
using testing::matchers::CtrlDeps;
using testing::matchers::Inputs;
using testing::matchers::Name;
using testing::matchers::NodeWith;
using testing::matchers::Op;
using testing::matchers::Out;
template <typename M, typename T>
std::string Explain(const T& t, const M& m) {
::testing::StringMatchResultListener listener;
EXPECT_THAT(t, ::testing::Not(m)); // For the error message.
EXPECT_FALSE(m.MatchAndExplain(t, &listener));
return listener.str();
}
TEST(NodeMatchers, CheckAgainstConstant) {
Scope root = Scope::NewRootScope().ExitOnError();
Output placeholder =
ops::Placeholder(root.WithOpName("placeholder"), DT_FLOAT);
EXPECT_THAT(placeholder.node(), NodeWith(Op("Placeholder")));
EXPECT_THAT(placeholder.node(), NodeWith(Name("placeholder")));
EXPECT_THAT(placeholder.node(),
NodeWith(Op("Placeholder"), Name("placeholder")));
EXPECT_THAT(placeholder.node(),
NodeWith(Name("placeholder"), Op("Placeholder")));
EXPECT_THAT(placeholder.node(), NodeWith(Inputs()));
EXPECT_THAT(placeholder.node(),
NodeWith(Op("Placeholder"), Name("placeholder"), Inputs()));
EXPECT_EQ(Explain(placeholder.node(), NodeWith(Op("Add"))),
"\nexpected op Add but found Placeholder");
EXPECT_EQ(Explain(placeholder.node(), NodeWith(Name("add"))),
"\nexpected name add but found placeholder");
EXPECT_EQ(Explain(placeholder.node(), NodeWith(Inputs(Out(NodeWith())))),
"\nexpected 1 inputs but node has 0");
}
TEST(NodeMatchers, CheckAgainstBinary) {
Scope root = Scope::NewRootScope().ExitOnError();
Output placeholder_a =
ops::Placeholder(root.WithOpName("placeholder_a"), DT_FLOAT);
Output placeholder_b =
ops::Placeholder(root.WithOpName("placeholder_b"), DT_FLOAT);
Output add = ops::Add(root.WithOpName("add"), placeholder_a, placeholder_b);
EXPECT_THAT(add.node(),
NodeWith(Op("Add"), Name("add"),
Inputs(Out(NodeWith(Name("placeholder_a"))),
Out(NodeWith(Name("placeholder_b"))))));
EXPECT_EQ(Explain(add.node(), NodeWith(Inputs())),
"\nexpected 0 inputs but node has 2");
EXPECT_EQ(
Explain(add.node(), NodeWith(Inputs(Out(NodeWith(Name("blah"))), _))),
"\ninput 0 does not match expected:\nname: blah, \nsource does not match "
"expected name: blah\n\t\nexpected name blah but found placeholder_a");
EXPECT_EQ(
Explain(add.node(), NodeWith(Inputs(_, Out(NodeWith(Name("blah")))))),
"\ninput 1 does not match expected:\nname: blah, \nsource does not match "
"expected name: blah\n\t\nexpected name blah but found placeholder_b");
}
TEST(NodeMatchers, CheckControlDependence) {
Scope root = Scope::NewRootScope().ExitOnError();
Output placeholder_a =
ops::Placeholder(root.WithOpName("placeholder_a"), DT_FLOAT);
Output placeholder_b =
ops::Placeholder(root.WithOpName("placeholder_b"), DT_FLOAT);
Output placeholder_c =
ops::Placeholder(root.WithOpName("placeholder_c"), DT_FLOAT);
Output placeholder_d =
ops::Placeholder(root.WithOpName("placeholder_d"), DT_FLOAT);
root.graph()->AddControlEdge(placeholder_a.node(), placeholder_c.node());
root.graph()->AddControlEdge(placeholder_b.node(), placeholder_c.node());
EXPECT_THAT(placeholder_c.node(),
NodeWith(Name("placeholder_c"),
CtrlDeps(NodeWith(Name("placeholder_a")),
NodeWith(Name("placeholder_b")))));
EXPECT_THAT(placeholder_d.node(),
NodeWith(Name("placeholder_d"), CtrlDeps()));
// TODO(griffithjames): Exactly match these explanations.
//
// When the OSS build has been updated to include the new error messages, the
// Explain() expectations can be exact strings again.
{
const std::string explanation =
Explain(placeholder_c.node(), NodeWith(CtrlDeps()));
EXPECT_NE(explanation.find("ctrl_deps, which has 2 elements"),
std::string::npos);
EXPECT_NE(explanation.find("does not match expected: is empty"),
std::string::npos);
}
{
const std::string explanation =
Explain(placeholder_d.node(), NodeWith(CtrlDeps(NodeWith())));
EXPECT_NE(explanation.find("ctrl_deps"), std::string::npos);
EXPECT_NE(explanation.find("does not match expected: has 1 element and "
"that element is any node"),
std::string::npos);
}
}
TEST(NodeMatchers, ConstValue) {
Scope root = Scope::NewRootScope().ExitOnError();
Output placeholder =
ops::Placeholder(root.WithOpName("placeholder"), DT_FLOAT);
Output const_0d = ops::Const(root.WithOpName("const_0d"), 42);
Output const_2d = ops::Const(root.WithOpName("const_2d"), {{1, 2}, {4, 3}});
EXPECT_THAT(const_0d.node(), NodeWith(ConstantValue(42)));
EXPECT_THAT(const_0d.node(), NodeWith(ConstantValue(42), Name("const_0d")));
EXPECT_THAT(const_2d.node(), NodeWith(ConstantValue({{1, 2}, {4, 3}})));
EXPECT_EQ(Explain(placeholder.node(), NodeWith(ConstantValue(42))),
"\nexpected op Const but found Placeholder");
EXPECT_EQ(
Explain(const_0d.node(), NodeWith(ConstantValue(43))),
"\nmismatch in constant tensor at index 0 expected = 43 actual = 42");
EXPECT_EQ(
Explain(const_0d.node(), NodeWith(ConstantValue({{1, 2}, {4, 3}}))),
"\nwas looking for tensor with 4 elements, found tensor with 1 elements");
EXPECT_EQ(
Explain(const_2d.node(), NodeWith(ConstantValue(42))),
"\nwas looking for tensor with 1 elements, found tensor with 4 elements");
}
TEST(NodeMatchers, AssignedDevice) {
Scope root = Scope::NewRootScope().ExitOnError();
Output placeholder_a =
ops::Placeholder(root.WithOpName("placeholder_a"), DT_FLOAT);
Output placeholder_b =
ops::Placeholder(root.WithOpName("placeholder_b"), DT_FLOAT);
Output assigned_add =
ops::Add(root.WithOpName("assigned_add"), placeholder_a, placeholder_b);
assigned_add.node()->set_assigned_device_name(
"/job:localhost/replica:0/task:0/device:CPU:0");
Output unassigned_add =
ops::Add(root.WithOpName("unassigned_add"), placeholder_a, placeholder_b);
EXPECT_THAT(
assigned_add.node(),
NodeWith(AssignedDevice("/job:localhost/replica:0/task:0/device:CPU:0")));
EXPECT_THAT(unassigned_add.node(), NodeWith(AssignedDevice("")));
EXPECT_EQ(Explain(unassigned_add.node(),
NodeWith(AssignedDevice(
"/job:localhost/replica:0/task:0/device:CPU:0"))),
"\nexpected assigned_device "
"/job:localhost/replica:0/task:0/device:CPU:0 but found \"\"");
}
TEST(NodeMatchers, OutputIndices) {
Scope root = Scope::NewRootScope().ExitOnError();
Output pred = ops::Placeholder(root.WithOpName("pred"), DT_BOOL);
Output data = ops::Placeholder(root.WithOpName("data"), DT_FLOAT);
ops::Switch sw(root.WithOpName("switch"), data, pred);
Output add = ops::Add(root.WithOpName("add"), sw.output_true,
ops::Placeholder(root.WithOpName("addend"), DT_FLOAT));
EXPECT_THAT(add.node(), NodeWith(Inputs(Out(1, NodeWith(Op("Switch"))), _)));
EXPECT_EQ(
Explain(add.node(), NodeWith(Inputs(Out(0, NodeWith(Op("Switch"))), _))),
"\ninput 0 does not match expected:\nop: Switch, \nexpected output slot "
"to be 0 but found 1");
}
TEST(NodeMatchers, Attrs) {
Scope root = Scope::NewRootScope().ExitOnError();
Output enter = ops::internal::Enter(
root.WithOpName("enter"),
ops::Placeholder(root.WithOpName("data"), DT_FLOAT), "frame_name",
ops::internal::Enter::Attrs{}.IsConstant(true));
EXPECT_THAT(enter.node(), NodeWith(Attr("is_constant", true)));
EXPECT_EQ(Explain(enter.node(), NodeWith(Attr("is_constant", false))),
"attribute named is_constant does not match value; expected: "
"\"false\", found: \"true\"");
EXPECT_EQ(Explain(enter.node(), NodeWith(Attr("missing_attr", false))),
"did not find attribute named \"missing_attr\" in node");
}
} // namespace
} // namespace testing
} // namespace tensorflow
tensorflow/compiler/jit/ops/BUILD
+43
View File
@@ -0,0 +1,43 @@
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:tensorflow.bzl", "tf_gen_op_wrapper_py")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//tensorflow/compiler/tf2xla:internal"],
licenses = ["notice"],
)
cc_library(
name = "xla_ops",
srcs = ["xla_ops.cc"],
deps = [
"//tensorflow/core:framework",
"@com_google_absl//absl/status",
],
alwayslink = 1,
)
tf_gen_op_wrapper_py(
name = "xla_ops_wrapper_py",
out = "xla_ops.py",
extra_py_deps = [
"//tensorflow/python:pywrap_tfe",
"//tensorflow/python/util:dispatch",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
py_lib_rule = py_library,
deps = ["//tensorflow/compiler/jit/ops:xla_ops"],
)
py_library(
name = "xla_ops_grad",
srcs = ["xla_ops_grad.py"],
strict_deps = True,
visibility = [
"//tensorflow/compiler/tf2xla:internal",
"//tensorflow/python/ops:__pkg__",
],
deps = ["//tensorflow/python/framework:ops"],
)
tensorflow/compiler/jit/ops/xla_ops.cc
+127
View File
@@ -0,0 +1,127 @@
/* 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 "absl/status/status.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"
namespace tensorflow {
using shape_inference::InferenceContext;
REGISTER_OP("XlaLaunch")
.Input("constants: Tconstants")
.Attr("Tconstants: list(type) >= 0")
.Input("args: Targs")
.Attr("Targs: list(type) >= 0")
.Input("resources: Nresources * resource")
.Attr("Nresources: int >= 0")
.Output("results: Tresults")
.Attr("Tresults: list(type) >= 0")
.Attr("function: func")
// XLA random-number generation ops are stateful.
// TODO(phawkins): create stateful and non-stateful variants of XlaLaunch.
.SetIsStateful()
.Doc("XLA Launch Op. For use by the XLA JIT only.");
REGISTER_OP("XlaLaunchV2")
.Input("args: Targs")
.Output("results: Tresults")
.Attr("Targs: list(type) >= 0")
.Attr("Tresults: list(type) >= 0")
.Attr("constants: list(int) >= 0")
.Attr("resources: list(int) >= 0")
.Attr("function: func")
.Doc("XLA Launch Op. For use by the XLA JIT only.");
REGISTER_OP("XlaClusterOutput")
.Input("input: T")
// Note: when replication is supported, this op will have N outputs.
.Output("outputs: T")
.Attr("T: type")
.SetShapeFn([](InferenceContext* c) {
for (int i = 0; i < c->num_outputs(); ++i) {
c->set_output(i, c->input(0));
}
return absl::OkStatus();
})
.Doc(
"Operator that connects the output of an XLA computation to other "
"consumer graph nodes.");
REGISTER_OP("_XlaCompile")
.Input("constants: Tconstants")
.Attr("Tconstants: list(type) >= 0")
.Attr("must_compile: bool")
.Input("args: Targs")
.Attr("Targs: list(type) >= 0")
.Input("resources: Nresources * resource")
.Attr("Nresources: int >= 0")
.Output("key: string")
.Output("compilation_successful: bool")
.Attr("function: func")
// The compilation cache is stateful.
.SetIsStateful()
.Doc(R"(XLA Compile Op. For use by the XLA JIT only.
Compiles a TensorFlow function into an XLA LocalExecutable and returns a key
that _XlaRun can use to look up the LocalExecutable and execute it.
key: A key that can be used to look up the local executable compiled by the
node and associated metadata.
compilation_successful: If the `must_compile` attr is false the _XlaCompile op
can decide not to compile the clusters based on some profitability
heuristics. In that case `compilation_successful` is false if _XlaCompile
chose not to compile the cluster. If the `must_compile` attr is true then
_XlaCompile always attempts to compile the cluster and
`compilation_successful` is always true.
)");
REGISTER_OP("_XlaRun")
.Input("args: Targs")
.Attr("Targs: list(type) >= 0")
.Output("results: Tresults")
.Attr("Tresults: list(type) >= 0")
.Input("key: string")
// XLA random-number generation ops are stateful.
// TODO(phawkins): create stateful and non-stateful variants of _XlaRun.
.SetIsStateful()
.Doc(R"(XLA Run Op. For use by the XLA JIT only.
Executes a TensorFlow function previously compiled into a LocalExecutable by an
_XlaCompile op.
)");
REGISTER_OP("_XlaMerge")
.Input("partitioned_call: T")
.Input("xla_run: T")
.Output("output: T")
.Attr("T: type")
.SetShapeFn([](InferenceContext* c) {
c->set_output(0, c->input(0));
return absl::OkStatus();
})
.Doc(R"(XLA Merge Op. For use by the XLA JIT only.
Merges the outputs from the PartitionedCall node and the _XlaRun node.
Unlike the TensorFlow Merge op, which requires inputs of some types to be
placed on the host, the _XlaMerge op can merge inputs of all types when
placed on the device. This prevents the need for copy operations, in
particular when an XLA cluster has int32 outputs. The _XlaMerge up does not
have a value_index output that identifies the chosen input.
)");
} // namespace tensorflow
tensorflow/compiler/jit/ops/xla_ops_grad.py
+25
View File
@@ -0,0 +1,25 @@
"""Gradients for XLA ops."""
# Copyright 2018 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.
# ==============================================================================
from tensorflow.python.framework import ops
@ops.RegisterGradient("XlaClusterOutput")
def _XlaClusterOutputGrad(_, grad):
del grad # unused
raise RuntimeError("Gradient computation of graph in xla.compile() is "
"prohibited because it can cause performance degradation."
"Please move gradient computation inside xla.compile().")
tensorflow/compiler/jit/package_groups.bzl
+13
View File
@@ -0,0 +1,13 @@
"""Tensorflow JIT package_group definitions"""
def legacy_jit_users_package_group(name):
"""Defines visibility group for //third_party/tensorflow/compiler/jit.
Args:
name: package group name
"""
native.package_group(
name = name,
packages = ["//..."],
)
tensorflow/compiler/jit/partially_decluster_pass.cc
+435
View File
@@ -0,0 +1,435 @@
/* Copyright 2018 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/jit/partially_decluster_pass.h"
#include "absl/algorithm/container.h"
#include "absl/container/flat_hash_set.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/compiler/jit/device_util.h"
#include "tensorflow/compiler/jit/xla_cluster_util.h"
#include "tensorflow/compiler/tf2xla/const_analysis.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "tensorflow/core/common_runtime/function.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/memory_types.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/graph/graph_node_util.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/public/version.h"
namespace tensorflow {
namespace {
bool NotBackedge(const Edge& edge) { return !edge.src()->IsNextIteration(); }
namespace reduce_device_to_host_copies {
absl::Status FindNodesToDecluster(const Graph& graph,
absl::flat_hash_set<Node*>* result,
absl::Span<Node* const> post_order) {
// Find nodes that have at least one user outside their cluster that expects
// hostmem output. These nodes should be cloned to outside the cluster to
// avoid the device-host copy we'd otherwise need.
MemoryTypeVector input_mtypes, output_mtypes;
for (Node* n : post_order) {
std::optional<absl::string_view> from_cluster = GetXlaClusterForNode(*n);
if (!from_cluster) {
continue;
}
// Assume the benefit of not outputting a larger tensor outweighs the
// benefit of this check.
// TODO(tpopp): Only apply this if the value being consumed is not output
// from the cluster to another consumer.
// TODO(tpopp): See if XlaRun can be modified to avoid this issue
// completely.
if (IsShapeConsumerOp(*n)) {
continue;
}
// We assume the only XLA-auto-clusterable operations with side effects are
// resource variable updates. We can't execute these twice.
if (HasResourceInputOrOutput(*n)) {
continue;
}
DeviceType device_type("");
TF_RETURN_IF_ERROR(
DeviceNameToDeviceType(n->assigned_device_name(), &device_type));
TF_RETURN_IF_ERROR(MemoryTypesForNode(graph.op_registry(), device_type,
n->def(), &input_mtypes,
&output_mtypes));
for (const Edge* e : n->out_edges()) {
Node* dst = e->dst();
if (e->IsControlEdge()) {
continue;
}
bool edge_incurs_extra_device_to_host_copy;
if (output_mtypes[e->src_output()] == DEVICE_MEMORY) {
// If the output of the *TensorFlow* operation is in DEVICE_MEMORY then
// keep the node clustered -- XLA will also produce the output in device
// memory and we will get some benefit from clustering.
edge_incurs_extra_device_to_host_copy = false;
} else {
MemoryTypeVector dst_input_mtypes, dst_output_mtypes;
DeviceType dst_device_type("");
TF_RETURN_IF_ERROR(DeviceNameToDeviceType(dst->assigned_device_name(),
&dst_device_type));
TF_RETURN_IF_ERROR(MemoryTypesForNode(graph.op_registry(), device_type,
dst->def(), &dst_input_mtypes,
&dst_output_mtypes));
edge_incurs_extra_device_to_host_copy =
dst_input_mtypes[e->dst_input()] == HOST_MEMORY;
}
if (!edge_incurs_extra_device_to_host_copy) {
continue;
}
// Check if `dst` is in a different cluster, unclustered, or about to be
// partially declustered (here we rely on the post-order traversal order).
// If yes, decluster `n` to avoid the device-to-host memcpy.
std::optional<absl::string_view> dst_cluster =
result->count(dst) ? std::nullopt : GetXlaClusterForNode(*dst);
if (from_cluster != dst_cluster) {
CHECK(result->insert(n).second);
break;
}
}
}
return absl::OkStatus();
}
absl::Status PartiallyDeclusterNode(Graph* graph, Node* n) {
absl::string_view cluster_name = *GetXlaClusterForNode(*n);
absl::InlinedVector<const Edge*, 6> out_edges_to_clone;
for (const Edge* out_edge : n->out_edges()) {
if (out_edge->IsControlEdge()) {
continue;
}
Node* dst = out_edge->dst();
std::optional<absl::string_view> dst_cluster_name =
GetXlaClusterForNode(*dst);
if (dst_cluster_name != cluster_name) {
out_edges_to_clone.push_back(out_edge);
}
}
CHECK(!out_edges_to_clone.empty()) << n->DebugString();
NodeDef ndef = n->def();
ndef.set_name(absl::StrCat(n->name(), "/declustered"));
MergeDebugInfo(NodeDebugInfo(n->def()), &ndef);
RemoveFromXlaCluster(&ndef);
TF_ASSIGN_OR_RETURN(Node * cloned_node, graph->AddNode(ndef));
cloned_node->set_assigned_device_name(n->assigned_device_name());
for (const Edge* in_edge : n->in_edges()) {
graph->AddEdge(in_edge->src(), in_edge->src_output(), cloned_node,
in_edge->dst_input());
}
for (const Edge* out_edge_to_clone : out_edges_to_clone) {
graph->AddEdge(cloned_node, out_edge_to_clone->src_output(),
out_edge_to_clone->dst(), out_edge_to_clone->dst_input());
graph->RemoveEdge(out_edge_to_clone);
}
if (n->out_edges().empty()) {
graph->RemoveNode(n);
}
return absl::OkStatus();
}
// Clones nodes to outside their cluster to avoid device-to-host copies. For
// instance, converts this:
//
// .....
// |
// v
// A_Clustered ====> C_Unclustered
// |
// v
// B_Clustered
//
// to:
//
// .....
// | |
// | +-------------+
// | |
// v v
// A_Clustered A_Unclustered ====> C_Unclustered
// |
// v
// B_Clustered
//
// where the ===> arrow has a hostmem source and destination and would entail a
// device to host copy if the source and destination were not in the same XLA
// cluster.
absl::Status PartiallyDeclusterGraph(Graph* graph) {
// When deciding whether to decluster a particular node, we base our decision
// on if we've decided that some of its consumers have to be declustered too.
// Iterating the graph in post-order guarantees that consumers have been
// visited before producers.
std::vector<Node*> post_order;
GetPostOrder(*graph, &post_order, /*stable_comparator=*/NodeComparatorName(),
/*edge_filter=*/NotBackedge);
absl::flat_hash_set<Node*> nodes_to_partially_decluster;
TF_RETURN_IF_ERROR(
FindNodesToDecluster(*graph, &nodes_to_partially_decluster, post_order));
if (VLOG_IS_ON(3)) {
for (Node* n : post_order) {
if (nodes_to_partially_decluster.count(n)) {
VLOG(3) << n->DebugString();
}
}
}
for (Node* n : post_order) {
if (nodes_to_partially_decluster.count(n)) {
TF_RETURN_IF_ERROR(PartiallyDeclusterNode(graph, n));
}
}
// Recompute post order since PartiallyDeclusterNode may have deleted nodes.
post_order.clear();
GetPostOrder(*graph, &post_order, /*stable_comparator=*/NodeComparatorName(),
/*edge_filter=*/NotBackedge);
nodes_to_partially_decluster.clear();
TF_RETURN_IF_ERROR(
FindNodesToDecluster(*graph, &nodes_to_partially_decluster, post_order));
CHECK(nodes_to_partially_decluster.empty());
return absl::OkStatus();
}
} // namespace reduce_device_to_host_copies
namespace reduce_recompilation {
bool IsIntraClusterEdge(const Edge& edge) {
std::optional<absl::string_view> src_cluster_name =
GetXlaClusterForNode(*edge.src());
std::optional<absl::string_view> dst_cluster_name =
GetXlaClusterForNode(*edge.dst());
return src_cluster_name.has_value() && src_cluster_name == dst_cluster_name;
}
bool IsMustCompileDevice(const DeviceType& device_type) {
const XlaOpRegistry::DeviceRegistration* registration;
if (XlaOpRegistry::GetCompilationDevice(device_type.type(), &registration)) {
return registration->autoclustering_policy ==
XlaOpRegistry::AutoclusteringPolicy::kAlways;
}
return false;
}
absl::Status MustCompileNode(const Node* n, bool* must_compile) {
DeviceType device_type("");
TF_RETURN_IF_ERROR(
DeviceNameToDeviceType(n->assigned_device_name(), &device_type));
if (IsMustCompileDevice(device_type)) {
*must_compile = true;
return absl::OkStatus();
}
// We must compile `n` if it does not have a TensorFlow kernel.
*must_compile = !FindKernelDef(device_type, n->def(), nullptr, nullptr).ok();
return absl::OkStatus();
}
// Declusters nodes to reduce the number of times we think we need to recompile
// a TensorFlow graph.
//
// Abstractly, if we have a cluster of this form:
//
// x0 = arg0
// x1 = arg1
// ...
// shape = f(x0, x1, ...)
// result = Reshape(input=<something>, new_shape=shape)
//
// then pulling `f` out of the cluster may reduce the number of compilations and
// will never increase the number of compilations.
//
// We may reduce the number of compilations if f is many to one. For instance
// if f(x,y) = x-y then x=3,y=1 and x=4,y=2 will generate two different
// compilations if f is in the cluster but only one compilation if f is outside
// the cluster.
//
// Declustering f will increase the number of compilations only if f is a
// one-to-many "function" i.e. isn't a function at all. RNG is one possible
// example, depending on how we look at it. But we never create clusters where
// such f's would be marked as must-be-constant.
//
// We assume here that the extra repeated (repeated compared to a clustered f
// where it will always be constant folded) host-side computation of f does not
// regress performance in any significant manner. We will have to revisit this
// algorithm with a more complex cost model if this assumption turns out to be
// incorrect.
absl::Status PartiallyDeclusterGraph(Graph* graph,
const FunctionLibraryDefinition* flib_def,
Env* env) {
std::vector<bool> compile_time_const_nodes(graph->num_node_ids());
OptimizerOptions opts;
auto pflr = std::make_unique<ProcessFunctionLibraryRuntime>(
nullptr, env, /*config=*/nullptr, TF_GRAPH_DEF_VERSION, flib_def, opts);
FunctionLibraryRuntime* lib_runtime =
pflr->GetFLR(ProcessFunctionLibraryRuntime::kDefaultFLRDevice);
TF_RETURN_IF_ERROR(BackwardsConstAnalysis(*graph, nullptr,
&compile_time_const_nodes,
lib_runtime, IsIntraClusterEdge));
std::vector<Node*> rpo;
GetReversePostOrder(*graph, &rpo, /*stable_comparator=*/NodeComparatorName(),
/*edge_filter=*/NotBackedge);
for (Node* n : rpo) {
if (!compile_time_const_nodes[n->id()]) {
continue;
}
absl::string_view cluster_name = *GetXlaClusterForNode(*n);
bool node_on_cluster_edge =
absl::c_all_of(n->in_edges(), [&](const Edge* e) {
std::optional<absl::string_view> incoming_cluster =
GetXlaClusterForNode(*e->src());
return !incoming_cluster || *incoming_cluster != cluster_name;
});
// We don't want to decluster F in a graph like
//
// Input -> OP -> Shape -> F -> Reshape
//
// Doing so will break up the cluster. Even if we were okay with breaking
// up the cluster we will at least have to relabel the two clusters to have
// different cluster names.
//
// We may want to revisit this in the future: we may have cases where OP is
// a small computation that does not benefit from XLA while XLA can optimize
// everything that follows the Reshape. In these cases it may be wise to
// remove Input, OP, Shape and F from the cluster, if F is a many-to-one
// function.
//
// Note that we do the right thing for graphs like:
//
// Input -> F0 -> F1 -> Reshape
//
// Since we iterate in RPO, we'll first encounter F0, decluster it, then
// encounter F1, decluster it and so on.
if (node_on_cluster_edge) {
bool must_compile_node;
TF_RETURN_IF_ERROR(MustCompileNode(n, &must_compile_node));
if (!must_compile_node) {
if (n->IsConstant()) {
// We must decluster Const nodes that have an input control edge from
// a different device, because this node may be part of the
// co-ordination of while loops between devices.
for (auto it : n->in_edges()) {
if (!it->src()->assigned_device_name().empty() &&
it->src()->assigned_device_name() !=
n->assigned_device_name()) {
VLOG(3) << "Declustering Const with cross-device control input "
<< n->name();
RemoveFromXlaCluster(n);
break;
}
}
} else {
VLOG(3) << "Declustering must-be-constant node " << n->name();
RemoveFromXlaCluster(n);
}
}
}
}
return absl::OkStatus();
}
} // namespace reduce_recompilation
namespace decluster_root_shape_consumers {
absl::Status PartiallyDeclusterGraph(Graph* graph) {
std::vector<Node*> reverse_post_order;
GetReversePostOrder(*graph, &reverse_post_order,
/*stable_comparator=*/NodeComparatorName(),
/*edge_filter=*/NotBackedge);
for (Node* n : reverse_post_order) {
if (!IsShapeConsumerOp(*n)) {
continue;
}
std::optional<absl::string_view> cluster = GetXlaClusterForNode(*n);
if (!cluster.has_value()) {
continue;
}
auto input_belongs_to_same_cluster = [&](const Edge* e) {
return cluster == GetXlaClusterForNode(*e->src());
};
if (absl::c_any_of(n->in_edges(), input_belongs_to_same_cluster)) {
continue;
}
VLOG(2) << "Declustering " << n->name()
<< " because it is a root shape consumer";
RemoveFromXlaCluster(n);
}
return absl::OkStatus();
}
} // namespace decluster_root_shape_consumers
} // namespace
absl::Status PartiallyDeclusterPass::Run(
const GraphOptimizationPassOptions& options) {
// NB! In this pass we assume the only XLA-auto-clusterable operations that
// may have side effects are resource variable operations so we don't cluster
// those. The pass will have to be updated if this assumption becomes
// invalid.
Graph* graph = options.graph->get();
TF_RETURN_IF_ERROR(
reduce_device_to_host_copies::PartiallyDeclusterGraph(graph));
if (options.flib_def == nullptr) {
return absl::InvalidArgumentError(
"GraphOptimizationPassOptions::flib_def must be set for "
"PartiallyDeclusterPass.");
}
if (options.session_options == nullptr ||
options.session_options->env == nullptr) {
return absl::InvalidArgumentError(
"GraphOptimizationPassOptions::session_options::env must be set for "
"PartiallyDeclusterPass.");
}
TF_RETURN_IF_ERROR(reduce_recompilation::PartiallyDeclusterGraph(
graph, options.flib_def, options.session_options->env));
TF_RETURN_IF_ERROR(
decluster_root_shape_consumers::PartiallyDeclusterGraph(graph));
return absl::OkStatus();
}
} // namespace tensorflow
tensorflow/compiler/jit/partially_decluster_pass.h
+35
View File
@@ -0,0 +1,35 @@
/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_JIT_PARTIALLY_DECLUSTER_PASS_H_
#define TENSORFLOW_COMPILER_JIT_PARTIALLY_DECLUSTER_PASS_H_
#include "tensorflow/core/common_runtime/optimization_registry.h"
namespace tensorflow {
// Clones or moves nodes from within a cluster to outside the cluster if
// profitable. There are two reasons why we do this:
//
// - Reducing device-to-host copies.
// - Reducing the number of XLA recompilations.
class PartiallyDeclusterPass : public GraphOptimizationPass {
public:
absl::Status Run(const GraphOptimizationPassOptions& options) override;
};
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_PARTIALLY_DECLUSTER_PASS_H_
tensorflow/compiler/jit/partially_decluster_pass_test.cc
+561
View File
@@ -0,0 +1,561 @@
/* Copyright 2018 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/jit/partially_decluster_pass.h"
#include "absl/memory/memory.h"
#include "tensorflow/cc/framework/ops.h"
#include "tensorflow/cc/ops/array_ops.h"
#include "tensorflow/cc/ops/const_op.h"
#include "tensorflow/cc/ops/control_flow_ops_internal.h"
#include "tensorflow/cc/ops/function_ops.h"
#include "tensorflow/cc/ops/sendrecv_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/compiler/jit/defs.h"
#include "tensorflow/compiler/jit/test_util.h"
#include "tensorflow/compiler/jit/xla_cluster_util.h"
#include "tensorflow/compiler/tf2xla/cc/ops/xla_ops.h"
#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/common_runtime/graph_def_builder_util.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/graph/algorithm.h"
#include "tensorflow/core/graph/graph_def_builder.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace {
REGISTER_OP("FakeNullary").Output("out: int32");
REGISTER_OP("FakeBinary")
.Input("host_in: int32")
.Input("device_in: int32")
.Output("host_out: int32")
.Output("device_out: int32");
REGISTER_OP("FakeResourceVar").Output("out: resource");
REGISTER_OP("FakeResourceUpdate")
.Input("in: resource")
.Output("out: resource")
.Output("something_else: int32");
class FakeBinaryOp : public OpKernel {
public:
explicit FakeBinaryOp(OpKernelConstruction* context) : OpKernel(context) {}
void Compute(OpKernelContext* ctx) override { CHECK(false); }
};
class FakeResourceUpdateOp : public OpKernel {
public:
explicit FakeResourceUpdateOp(OpKernelConstruction* context)
: OpKernel(context) {}
void Compute(OpKernelContext* ctx) override { CHECK(false); }
};
REGISTER_KERNEL_BUILDER(Name("FakeBinary")
.Device(DEVICE_CPU)
.HostMemory("host_in")
.HostMemory("host_out"),
FakeBinaryOp);
REGISTER_KERNEL_BUILDER(
Name("FakeResourceUpdate").Device(DEVICE_CPU).HostMemory("something_else"),
FakeResourceUpdateOp);
absl::Status PartiallyDecluster(std::unique_ptr<Graph>* graph) {
FixupSourceAndSinkEdges(graph->get());
// Assign all nodes to the CPU device.
static const char* kCpuDevice = "/job:localhost/replica:0/task:0/cpu:0";
for (Node* n : (*graph)->nodes()) {
if (n->assigned_device_name().empty()) {
n->set_assigned_device_name(kCpuDevice);
}
}
GraphOptimizationPassWrapper wrapper;
GraphOptimizationPassOptions opt_options =
wrapper.CreateGraphOptimizationPassOptions(graph);
PartiallyDeclusterPass pass;
return pass.Run(opt_options);
}
Node* FindNodeByName(const Graph& graph, const std::string& name) {
for (Node* node : graph.nodes()) {
if (node->name() == name) {
return node;
}
}
return nullptr;
}
bool GetInputsForNode(const Graph& graph, const std::string& node_name,
std::vector<Node*>* inputs) {
const Node* node = FindNodeByName(graph, node_name);
if (node == nullptr) {
return false;
}
for (const Edge* e : node->in_edges()) {
inputs->push_back(e->src());
}
std::sort(inputs->begin(), inputs->end(), NodeComparatorName());
return true;
}
TEST(PartiallyDeclusterPassTest, ClusteredAndUnclustered) {
std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
{
GraphDefBuilder builder(GraphDefBuilder::kFailImmediately);
Node* input =
ops::SourceOp("FakeNullary", builder.opts().WithName("Input"));
Node* clustered_producer =
ops::BinaryOp("FakeBinary", input, input,
builder.opts().WithName("ClusteredProducer"));
ops::BinaryOp("FakeBinary", clustered_producer, input,
builder.opts().WithName("UnclusteredConsumer"));
Node* clustered_consumer =
ops::BinaryOp("FakeBinary", {clustered_producer, 1}, input,
builder.opts().WithName("ClusteredConsumer"));
clustered_producer->AddAttr(kXlaClusterAttr, "cluster_0");
clustered_consumer->AddAttr(kXlaClusterAttr, "cluster_0");
TF_EXPECT_OK(GraphDefBuilderToGraph(builder, graph.get()));
}
TF_ASSERT_OK(PartiallyDecluster(&graph));
std::vector<Node*> unclustered_consumer_inputs;
ASSERT_TRUE(GetInputsForNode(*graph, "UnclusteredConsumer",
&unclustered_consumer_inputs));
ASSERT_EQ(unclustered_consumer_inputs.size(), 2);
EXPECT_EQ(unclustered_consumer_inputs[0]->name(),
"ClusteredProducer/declustered");
EXPECT_EQ(unclustered_consumer_inputs[1]->name(), "Input");
std::vector<Node*> clustered_consumer_inputs;
ASSERT_TRUE(GetInputsForNode(*graph, "ClusteredConsumer",
&clustered_consumer_inputs));
ASSERT_EQ(clustered_consumer_inputs.size(), 2);
EXPECT_EQ(clustered_consumer_inputs[0]->name(), "ClusteredProducer");
EXPECT_EQ(clustered_consumer_inputs[1]->name(), "Input");
}
TEST(PartiallyDeclusterPassTest, DifferentClusters) {
std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
{
GraphDefBuilder builder(GraphDefBuilder::kFailImmediately);
Node* input =
ops::SourceOp("FakeNullary", builder.opts().WithName("Input"));
Node* clustered_producer =
ops::BinaryOp("FakeBinary", input, input,
builder.opts().WithName("ClusteredProducer"));
Node* consumer_in_different_cluster =
ops::BinaryOp("FakeBinary", clustered_producer, input,
builder.opts().WithName("ConsumerInDifferentCluster"));
Node* clustered_consumer =
ops::BinaryOp("FakeBinary", input, {clustered_producer, 1},
builder.opts().WithName("ClusteredConsumer"));
clustered_producer->AddAttr(kXlaClusterAttr, "cluster_0");
clustered_consumer->AddAttr(kXlaClusterAttr, "cluster_0");
consumer_in_different_cluster->AddAttr(kXlaClusterAttr, "cluster_1");
TF_EXPECT_OK(GraphDefBuilderToGraph(builder, graph.get()));
}
TF_ASSERT_OK(PartiallyDecluster(&graph));
std::vector<Node*> inputs;
ASSERT_TRUE(GetInputsForNode(*graph, "ConsumerInDifferentCluster", &inputs));
ASSERT_EQ(inputs.size(), 2);
EXPECT_EQ(inputs[0]->name(), "ClusteredProducer/declustered");
EXPECT_EQ(inputs[1]->name(), "Input");
}
TEST(PartiallyDeclusterPassTest, DontDeclusterIfUserIsDeviceMem) {
std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
{
GraphDefBuilder builder(GraphDefBuilder::kFailImmediately);
Node* input =
ops::SourceOp("FakeNullary", builder.opts().WithName("Input"));
Node* clustered_producer =
ops::BinaryOp("FakeBinary", input, input,
builder.opts().WithName("ClusteredProducer"));
// The first input is hostmem and the second input is devicemem.
Node* consumer_in_different_cluster =
ops::BinaryOp("FakeBinary", input, clustered_producer,
builder.opts().WithName("ConsumerInDifferentCluster"));
Node* clustered_consumer =
ops::BinaryOp("FakeBinary", input, {clustered_producer, 1},
builder.opts().WithName("ClusteredConsumer"));
clustered_producer->AddAttr(kXlaClusterAttr, "cluster_0");
clustered_consumer->AddAttr(kXlaClusterAttr, "cluster_0");
consumer_in_different_cluster->AddAttr(kXlaClusterAttr, "cluster_1");
TF_EXPECT_OK(GraphDefBuilderToGraph(builder, graph.get()));
}
TF_ASSERT_OK(PartiallyDecluster(&graph));
std::vector<Node*> inputs;
ASSERT_TRUE(GetInputsForNode(*graph, "ConsumerInDifferentCluster", &inputs));
ASSERT_EQ(inputs.size(), 2);
EXPECT_EQ(inputs[0]->name(), "ClusteredProducer");
EXPECT_EQ(inputs[1]->name(), "Input");
}
TEST(PartiallyDeclusterPassTest, DontDuplicateResourceVarOps) {
std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
{
GraphDefBuilder builder(GraphDefBuilder::kFailImmediately);
Node* input =
ops::SourceOp("FakeNullary", builder.opts().WithName("Input"));
Node* resource_var = ops::SourceOp("FakeResourceVar",
builder.opts().WithName("ResourceVar"));
Node* clustered_producer =
ops::UnaryOp("FakeResourceUpdate", resource_var,
builder.opts().WithName("ClusteredProducer"));
Node* consumer_in_different_cluster =
ops::BinaryOp("FakeBinary", {clustered_producer, 1}, input,
builder.opts().WithName("ConsumerInDifferentCluster"));
Node* clustered_consumer =
ops::BinaryOp("FakeBinary", input, {clustered_producer, 1},
builder.opts().WithName("ClusteredConsumer"));
clustered_producer->AddAttr(kXlaClusterAttr, "cluster_0");
clustered_consumer->AddAttr(kXlaClusterAttr, "cluster_0");
consumer_in_different_cluster->AddAttr(kXlaClusterAttr, "cluster_1");
TF_EXPECT_OK(GraphDefBuilderToGraph(builder, graph.get()));
}
TF_ASSERT_OK(PartiallyDecluster(&graph));
std::vector<Node*> inputs;
ASSERT_TRUE(GetInputsForNode(*graph, "ConsumerInDifferentCluster", &inputs));
ASSERT_EQ(inputs.size(), 2);
EXPECT_EQ(inputs[0]->name(), "ClusteredProducer");
EXPECT_EQ(inputs[1]->name(), "Input");
}
TEST(PartiallyDeclusterPassTest, DeclusterDependentNodes) {
std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
{
GraphDefBuilder builder(GraphDefBuilder::kFailImmediately);
Node* input =
ops::SourceOp("FakeNullary", builder.opts().WithName("Input"));
Node* clustered_producer_0 =
ops::BinaryOp("FakeBinary", input, input,
builder.opts().WithName("ClusteredProducer0"));
Node* clustered_producer_1 =
ops::BinaryOp("FakeBinary", clustered_producer_0, input,
builder.opts().WithName("ClusteredProducer1"));
ops::BinaryOp("FakeBinary", clustered_producer_1, input,
builder.opts().WithName("UnclusteredConsumer"));
Node* clustered_consumer =
ops::BinaryOp("FakeBinary", {clustered_producer_1, 1}, input,
builder.opts().WithName("ClusteredConsumer"));
clustered_producer_0->AddAttr(kXlaClusterAttr, "cluster_0");
clustered_producer_1->AddAttr(kXlaClusterAttr, "cluster_0");
clustered_consumer->AddAttr(kXlaClusterAttr, "cluster_0");
TF_EXPECT_OK(GraphDefBuilderToGraph(builder, graph.get()));
}
TF_ASSERT_OK(PartiallyDecluster(&graph));
std::vector<Node*> unclustered_consumer_inputs, declustered_producer_1_inputs;
ASSERT_TRUE(GetInputsForNode(*graph, "UnclusteredConsumer",
&unclustered_consumer_inputs));
ASSERT_EQ(unclustered_consumer_inputs.size(), 2);
EXPECT_EQ(unclustered_consumer_inputs[0]->name(),
"ClusteredProducer1/declustered");
EXPECT_EQ(unclustered_consumer_inputs[1]->name(), "Input");
ASSERT_TRUE(GetInputsForNode(*graph, "ClusteredProducer1/declustered",
&declustered_producer_1_inputs));
ASSERT_EQ(declustered_producer_1_inputs.size(), 2);
EXPECT_EQ(declustered_producer_1_inputs[0]->name(),
"ClusteredProducer0/declustered");
EXPECT_EQ(declustered_producer_1_inputs[1]->name(), "Input");
}
void AddToCluster(absl::Span<Node* const> nodes,
absl::string_view cluster_name) {
for (Node* n : nodes) {
n->AddAttr(kXlaClusterAttr, std::string(cluster_name));
}
}
TEST(PartiallyDeclusterPassTest, DeclusterMustBeConstantNodes) {
tensorflow::Scope s = tensorflow::Scope::NewRootScope();
Output shape_a = ops::Placeholder(s.WithOpName("shape_a"), DT_INT32,
ops::Placeholder::Attrs{});
Output shape_b = ops::Placeholder(s.WithOpName("shape_b"), DT_INT32,
ops::Placeholder::Attrs{});
Output shape = ops::Add(s.WithOpName("shape"), shape_a, shape_b);
Output reshape_input = ops::Placeholder(s.WithOpName("reshape_input"),
DT_FLOAT, ops::Placeholder::Attrs{});
Output reshape = ops::Reshape(s.WithOpName("reshape"), reshape_input, shape);
AddToCluster({shape.node(), reshape.node()}, "cluster_0");
auto graph = std::make_unique<Graph>(OpRegistry::Global());
TF_ASSERT_OK(s.ToGraph(graph.get()));
TF_ASSERT_OK(PartiallyDecluster(&graph));
const Node* n = FindNodeByName(*graph, "shape");
ASSERT_NE(n, nullptr);
EXPECT_EQ(GetXlaClusterForNode(*n), std::nullopt);
}
TEST(PartiallyDeclusterPassTest, DeclusteringStopsAtMetadataOps) {
tensorflow::Scope s = tensorflow::Scope::NewRootScope();
Output input_a = ops::Placeholder(s.WithOpName("input_a"), DT_INT32,
ops::Placeholder::Attrs{});
Output input_b = ops::Placeholder(s.WithOpName("shape_b"), DT_FLOAT,
ops::Placeholder::Attrs{});
Output mul = ops::Mul(s.WithOpName("mul"), input_b, input_b);
Output shape_of_mul = ops::Shape(s.WithOpName("shape_of_mul"), mul);
Output shape = ops::Add(s.WithOpName("shape"), shape_of_mul, input_a);
Output reshape_input = ops::Placeholder(s.WithOpName("reshape_input"),
DT_FLOAT, ops::Placeholder::Attrs{});
Output reshape = ops::Reshape(s.WithOpName("reshape"), reshape_input, shape);
AddToCluster({mul.node(), shape_of_mul.node(), shape.node(), reshape.node()},
"cluster_0");
auto graph = std::make_unique<Graph>(OpRegistry::Global());
TF_ASSERT_OK(s.ToGraph(graph.get()));
TF_ASSERT_OK(PartiallyDecluster(&graph));
const Node* n = FindNodeByName(*graph, "shape");
ASSERT_NE(n, nullptr);
EXPECT_EQ(GetXlaClusterForNode(*n), "cluster_0");
}
TEST(PartiallyDeclusterPassTest, EdgeAcrossDifferentClusters) {
tensorflow::Scope s = tensorflow::Scope::NewRootScope();
Output shape_a = ops::Placeholder(s.WithOpName("shape_a"), DT_INT32,
ops::Placeholder::Attrs{});
Output shape_b = ops::Placeholder(s.WithOpName("shape_b"), DT_INT32,
ops::Placeholder::Attrs{});
Output shape = ops::Add(s.WithOpName("shape"), shape_a, shape_b);
Output reshape_input = ops::Placeholder(s.WithOpName("reshape_input"),
DT_FLOAT, ops::Placeholder::Attrs{});
Output reshape = ops::Reshape(s.WithOpName("reshape"), reshape_input, shape);
AddToCluster({reshape.node()}, "cluster_0");
AddToCluster({shape.node()}, "cluster_1");
auto graph = std::make_unique<Graph>(OpRegistry::Global());
TF_ASSERT_OK(s.ToGraph(graph.get()));
TF_ASSERT_OK(PartiallyDecluster(&graph));
const Node* n = FindNodeByName(*graph, "shape");
ASSERT_NE(n, nullptr);
EXPECT_EQ(GetXlaClusterForNode(*n), "cluster_1");
}
TEST(PartiallyDeclusterPassTest, DontDeclusterXlaDeviceOps) {
tensorflow::Scope s = tensorflow::Scope::NewRootScope();
Output shape_a = ops::Placeholder(s.WithOpName("shape_a"), DT_INT32,
ops::Placeholder::Attrs{});
Output shape_b = ops::Placeholder(s.WithOpName("shape_b"), DT_INT32,
ops::Placeholder::Attrs{});
Output shape = ops::Add(s.WithOpName("shape"), shape_a, shape_b);
Output reshape_input = ops::Placeholder(s.WithOpName("reshape_input"),
DT_FLOAT, ops::Placeholder::Attrs{});
Output reshape = ops::Reshape(s.WithOpName("reshape"), reshape_input, shape);
AddToCluster({shape.node(), reshape.node()}, "cluster_0");
auto graph = std::make_unique<Graph>(OpRegistry::Global());
TF_ASSERT_OK(s.ToGraph(graph.get()));
// This is needed to register the XLA_GPU device.
std::vector<std::unique_ptr<Device>> devices;
TF_ASSERT_OK(DeviceFactory::AddDevices(
SessionOptions(), "/job:localhost/replica:0/task:0", &devices));
// Scope::ToGraph loses the assigned device name since it goes through
// GraphDef/NodeDef which does not have a field for the assigned device name.
Node* n = FindNodeByName(*graph, "shape");
ASSERT_NE(n, nullptr);
n->set_assigned_device_name(
"/job:localhost/replica:0/task:0/device:XLA_GPU:0");
TF_ASSERT_OK(PartiallyDecluster(&graph));
EXPECT_EQ(GetXlaClusterForNode(*n), "cluster_0");
}
TEST(PartiallyDeclusterPassTest, EliminatedUnusedNodes) {
const char* const kClusteredProducer0Name = "ClusteredProducer0";
const char* const kClusteredProducer1Name = "ClusteredProducer1";
std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
{
GraphDefBuilder builder(GraphDefBuilder::kFailImmediately);
Node* input =
ops::SourceOp("FakeNullary", builder.opts().WithName("Input"));
Node* clustered_producer_0 =
ops::BinaryOp("FakeBinary", input, input,
builder.opts().WithName(kClusteredProducer0Name));
Node* clustered_producer_1 =
ops::BinaryOp("FakeBinary", clustered_producer_0, input,
builder.opts().WithName(kClusteredProducer1Name));
ops::BinaryOp("FakeBinary", clustered_producer_1, input,
builder.opts().WithName("UnclusteredConsumer"));
clustered_producer_0->AddAttr(kXlaClusterAttr, "cluster_0");
clustered_producer_1->AddAttr(kXlaClusterAttr, "cluster_0");
TF_EXPECT_OK(GraphDefBuilderToGraph(builder, graph.get()));
}
TF_ASSERT_OK(PartiallyDecluster(&graph));
EXPECT_EQ(FindNodeByName(*graph, kClusteredProducer0Name), nullptr);
EXPECT_EQ(FindNodeByName(*graph, kClusteredProducer1Name), nullptr);
}
TEST(PartiallyDeclusterPassTest, MetadataOpsDontStartClusters) {
tensorflow::Scope root = tensorflow::Scope::NewRootScope();
tensorflow::Scope in_cluster_and = root.WithXlaCluster("cluster_0");
Output a = ops::Placeholder(root.WithOpName("a"), DT_FLOAT);
Output b = ops::Shape(in_cluster_and.WithOpName("b"), a);
Output c = ops::Rank(in_cluster_and.WithOpName("c"), b);
Output d = ops::Size(in_cluster_and.WithOpName("d"), c);
(void)ops::Shape(in_cluster_and.WithOpName("e"), d);
auto graph = std::make_unique<Graph>(OpRegistry::Global());
TF_ASSERT_OK(root.ToGraph(graph.get()));
TF_ASSERT_OK(PartiallyDecluster(&graph));
Node* n_b = FindNodeByName(*graph, "b");
ASSERT_NE(n_b, nullptr);
EXPECT_EQ(GetXlaClusterForNode(*n_b), std::nullopt);
Node* n_c = FindNodeByName(*graph, "c");
ASSERT_NE(n_c, nullptr);
EXPECT_EQ(GetXlaClusterForNode(*n_c), std::nullopt);
Node* n_d = FindNodeByName(*graph, "d");
ASSERT_NE(n_d, nullptr);
EXPECT_EQ(GetXlaClusterForNode(*n_d), std::nullopt);
Node* n_e = FindNodeByName(*graph, "e");
ASSERT_NE(n_e, nullptr);
EXPECT_EQ(GetXlaClusterForNode(*n_e), std::nullopt);
}
TEST(PartiallyDeclusterPassTest, MetaConsumersArentDeclustered) {
tensorflow::Scope root = tensorflow::Scope::NewRootScope();
tensorflow::Scope in_cluster_and = root.WithXlaCluster("cluster_0");
auto graph = std::make_unique<Graph>(OpRegistry::Global());
Output a = ops::Placeholder(root.WithOpName("a"), DT_FLOAT);
Output b = ops::Add(in_cluster_and.WithOpName("b"), a, a);
Output c = ops::Rank(in_cluster_and.WithOpName("c"), b);
Output e;
TF_ASSERT_OK(
CreateOutputWithScope("FakeBinary", {c, c}, root.WithOpName("e"), &e));
TF_ASSERT_OK(root.ToGraph(graph.get()));
TF_ASSERT_OK(PartiallyDecluster(&graph));
Node* n_b = FindNodeByName(*graph, "b");
ASSERT_NE(n_b, nullptr);
EXPECT_EQ(GetXlaClusterForNode(*n_b), "cluster_0");
Node* n_c = FindNodeByName(*graph, "c");
ASSERT_NE(n_c, nullptr);
EXPECT_EQ(GetXlaClusterForNode(*n_c), "cluster_0");
}
TEST(PartiallyDeclusterPassTest, ConstInputsToSliceArentDeclustered) {
tensorflow::Scope root = tensorflow::Scope::NewRootScope();
tensorflow::Scope in_cluster_and = root.WithXlaCluster("cluster_0");
auto graph = std::make_unique<Graph>(OpRegistry::Global());
Output a = ops::Placeholder(root.WithOpName("a"), DT_FLOAT,
ops::Placeholder::Attrs{{4}});
Output b = ops::Const(in_cluster_and.WithOpName("b"), {1});
Output c = ops::Const(in_cluster_and.WithOpName("c"), {2});
Output d = ops::Slice(in_cluster_and.WithOpName("d"), a, b, c);
TF_ASSERT_OK(root.ToGraph(graph.get()));
TF_ASSERT_OK(PartiallyDecluster(&graph));
Node* n_b = FindNodeByName(*graph, "b");
ASSERT_NE(n_b, nullptr);
EXPECT_EQ(GetXlaClusterForNode(*n_b), "cluster_0");
Node* n_c = FindNodeByName(*graph, "c");
ASSERT_NE(n_c, nullptr);
EXPECT_EQ(GetXlaClusterForNode(*n_c), "cluster_0");
}
TEST(PartiallyDeclusterPassTest,
ConstInLoopWithCrossDeviceControlInputsAreDeclustered) {
// Based on DontClusterTheSpecialIdentityDrivingConstsInLoop in
// mark_for_compilation_pass_test.cc
tensorflow::Scope root = tensorflow::Scope::NewRootScope();
tensorflow::Scope in_cluster_and = root.WithXlaCluster("cluster_0");
auto graph = std::make_unique<Graph>(OpRegistry::Global());
Output a = ops::Placeholder(root.WithOpName("a"), DT_FLOAT,
ops::Placeholder::Attrs{{4}});
Output b = ops::Const(in_cluster_and.WithOpName("b"), {1});
Output c = ops::Const(in_cluster_and.WithOpName("c"), {2});
Output slice = ops::Slice(in_cluster_and.WithOpName("slice"), a, b, c);
Output cond = ops::Placeholder(root.WithOpName("cond"), DT_BOOL);
Output value = ops::Placeholder(root.WithOpName("value"), DT_FLOAT);
Output loop_cond = ops::LoopCond(root.WithOpName("loop_cond"), cond);
ops::Switch switch_node(root.WithOpName("switch"), value, loop_cond);
Output identity =
ops::Identity(root.WithOpName("identity"), switch_node.output_true);
root.graph()->AddControlEdge(identity.node(), b.node());
TF_ASSERT_OK(root.ToGraph(graph.get()));
// This is needed to register the XLA_GPU device.
std::vector<std::unique_ptr<Device>> devices;
TF_ASSERT_OK(DeviceFactory::AddDevices(
SessionOptions(), "/job:localhost/replica:0/task:0", &devices));
// Scope::ToGraph loses the assigned device name since it goes through
// GraphDef/NodeDef which does not have a field for the assigned device name.
Node* identity_node = FindNodeByName(*graph, "identity");
ASSERT_NE(identity_node, nullptr);
identity_node->set_assigned_device_name(
"/job:localhost/replica:0/task:0/device:XLA_GPU:0");
TF_ASSERT_OK(PartiallyDecluster(&graph));
Node* n_b = FindNodeByName(*graph, "b");
ASSERT_NE(n_b, nullptr);
EXPECT_EQ(GetXlaClusterForNode(*n_b), std::nullopt);
Node* n_c = FindNodeByName(*graph, "c");
ASSERT_NE(n_c, nullptr);
EXPECT_EQ(GetXlaClusterForNode(*n_c), "cluster_0");
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/pjrt_base_device.cc
+60
View File
@@ -0,0 +1,60 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/jit/pjrt_base_device.h"
namespace tensorflow {
namespace {
DeviceAttributes BuildPjRtBaseDeviceAttributes(const std::string& name_prefix,
const std::string& device_name,
int device_ordinal) {
return Device::BuildDeviceAttributes(
absl::StrCat(name_prefix, "/device:", device_name, ":", device_ordinal),
DeviceType(device_name), Bytes(16ULL << 30), DeviceLocality(),
absl::StrCat("device: ", device_name, " device"));
}
} // namespace
PjRtBaseDevice::PjRtBaseDevice(const SessionOptions& session_options,
const Options& options)
: LocalDevice(session_options,
BuildPjRtBaseDeviceAttributes(options.device_name_prefix,
options.device_name,
options.device_ordinal)),
metadata_(DeviceType(options.compilation_device_name),
options.shape_determination_fns) {
if (options.shape_determination_fns.empty()) {
LOG(ERROR) << "shape_representation_fns must be non-empty.";
}
VLOG(1) << "Created PJRT base device " << options.compilation_device_name
<< " device_name: " << name();
}
/*static*/ absl::StatusOr<const PjRtBaseDevice::Metadata*>
PjRtBaseDevice::GetMetadataFromDevice(DeviceBase* device) {
PjRtBaseDevice* pjrt_device =
dynamic_cast<PjRtBaseDevice*>(device->UnderlyingDevice());
if (pjrt_device == nullptr) {
return absl::InternalError(absl::StrCat(
"Cannot get device metadata from non-PJRT device \"", device->name(),
"\". GetMetadata must only be called on a device derived from "
"PjRtBaseDevice. Either an internal bug has been triggered, or an "
"XLA-specific op has been placed on the wrong device."));
}
return &pjrt_device->metadata_;
}
} // namespace tensorflow
tensorflow/compiler/jit/pjrt_base_device.h
+112
View File
@@ -0,0 +1,112 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_JIT_PJRT_BASE_DEVICE_H_
#define TENSORFLOW_COMPILER_JIT_PJRT_BASE_DEVICE_H_
#include <string>
#include <utility>
#include <vector>
#include "tensorflow/compiler/tf2xla/layout_util.h"
#include "tensorflow/core/common_runtime/local_device.h"
#include "tensorflow/core/framework/device_base.h"
namespace tensorflow {
// tensorflow::PjRtBaseDevice replaces the deprecated tensorflow::XlaDevice.
// This accelerator agnostic device is mainly used to store metadata.
class PjRtBaseDevice : public LocalDevice {
public:
// Stores metadata about the PjRtBaseDevice.
class Metadata {
public:
Metadata(const DeviceType& jit_device_type,
std::vector<XlaShapeLayoutHelpers::ShapeDeterminationFns>
shape_determination_fns)
: jit_device_type_(jit_device_type),
shape_determination_fns_(std::move(shape_determination_fns)) {}
// The index of the device on this host.
int device_ordinal() const;
const DeviceType& jit_device_type() const { return jit_device_type_; }
const XlaShapeLayoutHelpers::ShapeDeterminationFns&
default_shape_determination_fns() const {
return shape_determination_fns_.at(0);
}
const XlaShapeLayoutHelpers::ShapeDeterminationFns&
shape_determination_fns_at(int i) const {
return shape_determination_fns_[i];
}
private:
const DeviceType jit_device_type_;
std::vector<XlaShapeLayoutHelpers::ShapeDeterminationFns>
shape_determination_fns_;
Metadata(const Metadata&) = delete;
void operator=(const Metadata&) = delete;
};
struct Options {
// The device name's prefix (e.g., "/task:7")
std::string device_name_prefix;
// The name of the device (e.g., "TPU")
std::string device_name;
// The index of the device.
int device_ordinal = -1;
// The name of the compilation device, also referred to as jit_device_type.
// (e.g., "XLA_CPU_JIT");
std::string compilation_device_name;
// A vector of ShapeDeterminationFn (i.e., a bundle of LayoutSelectionFn,
// ShapeRepresentationFn). Each bundle describes how the on-host shapes of
// a) argument and return value, for entry computations b) variables, for
// all computations, should be represented in XLA. Parameters/return values
// will be shaped according to the function pair, and reshaped back to/from
// their declared shapes for computations. Must be non-empty.
std::vector<XlaShapeLayoutHelpers::ShapeDeterminationFns>
shape_determination_fns;
Options(std::string device_name_prefix, std::string device_name,
int device_ordinal, std::string compilation_device_name,
std::vector<XlaShapeLayoutHelpers::ShapeDeterminationFns>
shape_determination_fns)
: device_name_prefix(device_name_prefix),
device_name(device_name),
device_ordinal(device_ordinal),
compilation_device_name(compilation_device_name),
shape_determination_fns(shape_determination_fns) {}
};
// Creates a new PJRT base device.
PjRtBaseDevice(const SessionOptions& session_options, const Options& options);
static absl::StatusOr<const PjRtBaseDevice::Metadata*> GetMetadataFromDevice(
DeviceBase* device);
private:
// The metadata of this PjRtBaseDevice.
const Metadata metadata_;
};
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_PJRT_BASE_DEVICE_H_
tensorflow/compiler/jit/pjrt_compile_util.cc
+89
View File
@@ -0,0 +1,89 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/jit/pjrt_compile_util.h"
#include <vector>
#include "tensorflow/compiler/jit/device_compilation_profiler.h"
#include "tensorflow/compiler/jit/device_compiler.h"
#include "tensorflow/compiler/jit/xla_compile_util.h"
#include "tensorflow/compiler/jit/xla_compiler_options_util.h"
#include "tensorflow/compiler/jit/xla_platform_info.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/pjrt/pjrt_client.h"
#include "tensorflow/core/framework/device_base.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/lib/core/refcount.h"
#include "tensorflow/core/platform/status.h"
#include "tsl/platform/errors.h"
#include "tsl/platform/statusor.h"
namespace tensorflow {
using PjRtDeviceCompiler =
DeviceCompiler<xla::PjRtLoadedExecutable, xla::PjRtClient>;
absl::Status CompileToPjRtLoadedExecutable(
const DeviceBase* device, const XlaPlatformInfo& platform_info,
const NameAttrList& function,
const std::vector<XlaCompiler::Argument>& args,
DeviceCompileMode compile_mode, bool has_ref_vars,
bool may_alias_resource_update, FunctionLibraryRuntime* flr,
ResourceMgr* rm, const XlaCompiler::CompilationResult** compilation_result,
xla::PjRtClient** client, xla::PjRtLoadedExecutable** executable) {
PjRtDeviceCompiler* pjrt_device_compiler;
DeviceCompilationProfiler* profiler;
TF_RETURN_IF_ERROR(GetOrCreatePjRtDeviceCompilerAndProfiler(
platform_info, rm, flr, &pjrt_device_compiler, &profiler));
// Hold the reference to the PJRT device compiler and profiler during
// evaluation. (We could probably free them sooner because the ResourceMgr
// will retain references, but this is more obviously correct.)
core::ScopedUnref pjrt_device_compiler_ref(pjrt_device_compiler);
core::ScopedUnref profiler_ref(profiler);
*client = pjrt_device_compiler->client();
XlaCompiler::Options options = GenerateCompilerOptionsForPjRt(
*flr, device, platform_info, pjrt_device_compiler);
XlaCompiler::CompileOptions compile_options =
GenerateCompileOptions(has_ref_vars, may_alias_resource_update);
return pjrt_device_compiler->CompileIfNeeded(
options, function, args, compile_options, compile_mode, profiler,
compilation_result, executable);
}
absl::Status CompileToPjRtLoadedExecutable(
const OpKernelContext& ctx, const XlaPlatformInfo& platform_info,
const NameAttrList& function,
const std::vector<XlaCompiler::Argument>& args,
DeviceCompileMode compile_mode, bool has_ref_vars,
bool may_alias_resource_update,
const XlaCompiler::CompilationResult** compilation_result,
xla::PjRtClient** client, xla::PjRtLoadedExecutable** executable) {
TF_ASSIGN_OR_RETURN(ResourceMgr * rm, GetResourceMgrForDeviceCompiler(
ctx, platform_info.device_type()));
return CompileToPjRtLoadedExecutable(
ctx.device(), platform_info, function, args, compile_mode, has_ref_vars,
may_alias_resource_update, ctx.function_library(), rm, compilation_result,
client, executable);
}
} // namespace tensorflow
tensorflow/compiler/jit/pjrt_compile_util.h
+60
View File
@@ -0,0 +1,60 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_JIT_PJRT_COMPILE_UTIL_H_
#define TENSORFLOW_COMPILER_JIT_PJRT_COMPILE_UTIL_H_
#include "tensorflow/compiler/jit/xla_compile_util.h"
#include "tensorflow/compiler/jit/xla_platform_info.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/pjrt/pjrt_client.h"
#include "tensorflow/core/framework/device_base.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
// Compiles a `function` to PjRtLoadedExecutable `executable` with `ctx`.
// The compilation result is output in `compilation_result`. The PJRT client
// used for compilation is output in `client`. The PJRT executable is output in
// `executable`.
absl::Status CompileToPjRtLoadedExecutable(
const OpKernelContext& ctx, const XlaPlatformInfo& platform_info,
const NameAttrList& function,
const std::vector<XlaCompiler::Argument>& args,
DeviceCompileMode compile_mode, bool has_ref_vars,
bool may_alias_resource_update,
const XlaCompiler::CompilationResult** compilation_result,
xla::PjRtClient** client, xla::PjRtLoadedExecutable** executable);
// Similar to the above function but it does not take a OpKernelContext.
// Instead, it takes the following arguments that are obtained from
// OpKernelContext in the above function.
// - `device`: the device used to compile the function.
// - `rm`: the resource manager for DeviceCompiler to store JIT-compiled XLA
// computation.
// - `flr`: the FunctionLibraryRuntime for the `function`.
absl::Status CompileToPjRtLoadedExecutable(
const DeviceBase* device, const XlaPlatformInfo& platform_info,
const NameAttrList& function,
const std::vector<XlaCompiler::Argument>& args,
DeviceCompileMode compile_mode, bool has_ref_vars,
bool may_alias_resource_update, FunctionLibraryRuntime* flr,
ResourceMgr* rm, const XlaCompiler::CompilationResult** compilation_result,
xla::PjRtClient** client, xla::PjRtLoadedExecutable** executable);
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_PJRT_COMPILE_UTIL_H_
tensorflow/compiler/jit/pjrt_compile_util_test.cc
+129
View File
@@ -0,0 +1,129 @@
/* 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.
==============================================================================*/
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#include "tensorflow/compiler/jit/pjrt_compile_util.h"
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "tensorflow/cc/framework/scope.h"
#include "tensorflow/cc/ops/function_ops.h"
#include "tensorflow/cc/ops/math_ops.h"
#include "tensorflow/compiler/jit/test_util.h"
#include "tensorflow/core/framework/device_base.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/graph_to_functiondef.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/statusor.h"
namespace tensorflow {
namespace {
absl::StatusOr<std::unique_ptr<Graph>> SampleGraphAddXY() {
std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
Scope scope = Scope::NewRootScope().ExitOnError();
auto a = ops::_Arg(scope.WithOpName("A"), DT_INT32, 0);
auto b = ops::_Arg(scope.WithOpName("B"), DT_INT32, 1);
auto c = ops::Add(scope.WithOpName("C"), a, b);
auto d = ops::_Retval(scope.WithOpName("D"), c, 0);
TF_RETURN_IF_ERROR(scope.ToGraph(graph.get()));
return graph;
}
absl::StatusOr<FunctionDef> SampleFuntionAddXY(const std::string& name) {
TF_ASSIGN_OR_RETURN(auto graph, SampleGraphAddXY());
FunctionDef fdef;
TF_RETURN_IF_ERROR(GraphToFunctionDef(*graph, name, &fdef));
return fdef;
}
std::vector<XlaCompiler::Argument> SampleArgsForAddXY() {
std::vector<XlaCompiler::Argument> args(2);
args[0].kind = XlaCompiler::Argument::kParameter;
args[0].type = DT_INT32;
args[0].shape = TensorShape({2});
args[1].kind = XlaCompiler::Argument::kParameter;
args[1].type = DT_INT32;
args[1].shape = TensorShape({2});
return args;
}
TEST(PjrtCompileUtilTest, CompileToPjRtLoadedExecutable) {
DeviceSetup device_setup;
TF_ASSERT_OK_AND_ASSIGN(auto fdef, SampleFuntionAddXY("foo"));
device_setup.AddDevicesAndSetUp({DEVICE_GPU}, fdef);
Device* device = device_setup.GetDevice(DEVICE_GPU);
const XlaPlatformInfo platform_info = XlaPlatformInfoFromDevice(device);
NameAttrList function;
function.set_name("foo");
ResourceMgr resource_mgr("");
const XlaCompiler::CompilationResult* compilation_result = nullptr;
xla::PjRtLoadedExecutable* pjrt_executable = nullptr;
xla::PjRtClient* pjrt_client = nullptr;
TF_EXPECT_OK(CompileToPjRtLoadedExecutable(
device, platform_info, function, SampleArgsForAddXY(),
DeviceCompileMode::kStrict, /*has_ref_vars=*/true,
/*may_alias_resource_update=*/true, device_setup.flr(), &resource_mgr,
&compilation_result, &pjrt_client, &pjrt_executable));
EXPECT_TRUE(compilation_result != nullptr);
EXPECT_TRUE(pjrt_executable != nullptr);
EXPECT_TRUE(pjrt_client != nullptr);
}
TEST(PjrtCompileUtilTest, CompileToPjRtLoadedExecutableWithOpKernelContext) {
DeviceSetup device_setup;
TF_ASSERT_OK_AND_ASSIGN(auto fdef, SampleFuntionAddXY("foo"));
device_setup.AddDevicesAndSetUp({DEVICE_GPU}, fdef);
Device* device = device_setup.GetDevice(DEVICE_GPU);
const XlaPlatformInfo platform_info = XlaPlatformInfoFromDevice(device);
NameAttrList function;
function.set_name("foo");
ResourceMgr resource_mgr("");
OpKernelContext::Params params;
params.resource_manager = &resource_mgr;
params.device = device;
params.function_library = device_setup.flr();
OpKernelContext ctx(&params, 1);
const XlaCompiler::CompilationResult* compilation_result = nullptr;
xla::PjRtLoadedExecutable* pjrt_executable = nullptr;
xla::PjRtClient* pjrt_client = nullptr;
TF_EXPECT_OK(CompileToPjRtLoadedExecutable(
ctx, platform_info, function, SampleArgsForAddXY(),
DeviceCompileMode::kStrict, /*has_ref_vars=*/true,
/*may_alias_resource_update=*/true, &compilation_result, &pjrt_client,
&pjrt_executable));
EXPECT_TRUE(compilation_result != nullptr);
EXPECT_TRUE(pjrt_executable != nullptr);
EXPECT_TRUE(pjrt_client != nullptr);
}
} // namespace
} // namespace tensorflow
#endif
tensorflow/compiler/jit/pjrt_device_compiler_client.cc
+93
View File
@@ -0,0 +1,93 @@
/* 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/compiler/jit/pjrt_device_compiler_client.h"
#include <memory>
#include <string>
#include <utility>
namespace tensorflow {
xla::CompileOptions GetPjRtCompileOptions(
const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& result) {
xla::CompileOptions pjrt_compile_options;
pjrt_compile_options.argument_layouts = result.xla_input_shapes;
pjrt_compile_options.executable_build_options =
GetExecutableBuildOptions(options, result, /*default_device_ordinal=*/-1);
if (pjrt_compile_options.executable_build_options.num_replicas() > 1 ||
pjrt_compile_options.executable_build_options.num_partitions() > 1) {
// Compile executable for sharded program
pjrt_compile_options.compile_portable_executable = false;
} else {
// Compile portable executable for single device compilation.
pjrt_compile_options.compile_portable_executable = true;
}
return pjrt_compile_options;
}
absl::StatusOr<std::unique_ptr<xla::PjRtLoadedExecutable>>
PjRtDeviceCompilerClient::BuildExecutable(
const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& result) {
VLOG(2) << "Compiling to xla::PjRtLoadedExecutable.";
TF_ASSIGN_OR_RETURN(
auto executable,
client_->CompileAndLoad(*result.computation,
GetPjRtCompileOptions(options, result)));
VLOG(2) << "Compiled PJRT executable " << executable->name()
<< " num_replicas " << executable->num_replicas()
<< " num_partitions " << executable->num_partitions();
return std::move(executable);
}
absl::StatusOr<std::string> PjRtDeviceCompilerClient::SerializeExecutable(
const xla::PjRtLoadedExecutable& executable) {
VLOG(1) << "Serializing xla::PjRtLoadedExecutable to string.";
return executable.SerializeExecutable();
}
absl::StatusOr<std::string> PjRtDeviceCompilerClient::BuildSerializedExecutable(
const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& result) {
VLOG(1) << "PJRT currently doesn't support AOT compilation. Compiling to "
"xla::PjRtLoadedExecutable and serializing it";
TF_ASSIGN_OR_RETURN(auto executable, BuildExecutable(options, result));
return executable->SerializeExecutable();
}
absl::StatusOr<std::unique_ptr<xla::PjRtLoadedExecutable>>
PjRtDeviceCompilerClient::LoadExecutable(
const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& result,
const std::string& serialized_executable) {
VLOG(1) << "Deserializing from string to xla::PjRtLoadedExecutable.";
return client_->LoadSerializedExecutable(
serialized_executable, GetPjRtCompileOptions(options, result),
xla::LoadOptions());
}
void PjRtDeviceCompilerClient::WaitForProgramsToFinish() {
// TODO(b/255826209): Modify this if PjRtClient exposes a function to wait for
// programs to finish.
LOG(INFO) << "Unimplemented: PJRT uses futures and waiting for programs to "
"finish isn't necessary.";
}
} // namespace tensorflow
tensorflow/compiler/jit/pjrt_device_compiler_client.h
+85
View File
@@ -0,0 +1,85 @@
/* 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_COMPILER_JIT_PJRT_DEVICE_COMPILER_CLIENT_H_
#define TENSORFLOW_COMPILER_JIT_PJRT_DEVICE_COMPILER_CLIENT_H_
#include <memory>
#include <optional>
#include <string>
#include "tensorflow/compiler/jit/device_compiler_client.h"
#include "xla/pjrt/pjrt_client.h"
namespace tensorflow {
// Calls into PjRtClient to provide functionality for building, serializing and
// loading PjRtLoadedExecutables.
class PjRtDeviceCompilerClient
: public DeviceCompilerClient<xla::PjRtLoadedExecutable, xla::PjRtClient> {
public:
explicit PjRtDeviceCompilerClient(xla::PjRtClient* client)
: client_(client) {}
absl::StatusOr<std::unique_ptr<xla::PjRtLoadedExecutable>> BuildExecutable(
const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& result) override;
// Returns a platform-specific serialization of `executable`. The
// serialization is not guaranteed to be stable over time. `executable` must
// have been produced by this client.
absl::StatusOr<std::string> SerializeExecutable(
const xla::PjRtLoadedExecutable& executable) override;
// PjRt doesn't support AOT compilation yet. Builds a PjRtLoadedExecutable and
// serializes it to string.
absl::StatusOr<std::string> BuildSerializedExecutable(
const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& result) override;
// Deserializes a serialized executable as produced by
// PjRtExecutable::SerializeExecutable(). `serialized_executable` must have
// been produced by a compiler of the same platform and version as this one.
//
// PjRt doesn't support AOT compilation yet. Loading a serialized executable
// is currently only implemented for TfrtTpuPjrtClient and hence, this
// function doesn't use PjRtClient::LoadSerializedExecutable() and uses
// PjRtClient::DeserializeExecutable() instead.
absl::StatusOr<std::unique_ptr<xla::PjRtLoadedExecutable>> LoadExecutable(
const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& result,
const std::string& serialized_executable) override;
// No-op. PJRT uses futures and waiting for programs to finish isn't
// necessary.
void WaitForProgramsToFinish() override;
xla::PjRtClient* client() const override { return client_; }
private:
xla::PjRtClient* const client_;
PjRtDeviceCompilerClient(const PjRtDeviceCompilerClient&) = delete;
void operator=(const PjRtDeviceCompilerClient&) = delete;
};
// Generates CompileOptions for PJRT compilation.
xla::CompileOptions GetPjRtCompileOptions(
const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& result);
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_PJRT_DEVICE_COMPILER_CLIENT_H_
tensorflow/compiler/jit/pjrt_device_context.cc
+305
View File
@@ -0,0 +1,305 @@
/* 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/compiler/jit/pjrt_device_context.h"
#include <memory>
#include <optional>
#include <utility>
#include "absl/status/status.h"
#include "tensorflow/c/experimental/next_pluggable_device/tensor_pjrt_buffer_util.h"
#include "tensorflow/compiler/jit/pjrt_tensor_buffer.h"
#include "tensorflow/compiler/jit/pjrt_tensor_buffer_util.h"
#include "tensorflow/compiler/tf2xla/literal_util.h"
#include "xla/pjrt/pjrt_client.h"
#include "xla/pjrt/pjrt_common.h"
#include "xla/tsl/c/tsl_status_internal.h"
#include "xla/tsl/framework/device_id_utils.h"
#include "tensorflow/core/common_runtime/dma_helper.h"
#include "tensorflow/core/common_runtime/next_pluggable_device/next_pluggable_device_api.h"
#include "tensorflow/core/framework/device.h"
#include "tensorflow/core/framework/device_factory.h"
#include "tensorflow/core/profiler/lib/traceme.h"
#include "tensorflow/core/tfrt/common/async_value_tensor.h"
#include "tensorflow/core/tfrt/common/create_pjrt_client_util.h"
namespace tensorflow {
namespace {
absl::StatusOr<std::unique_ptr<xla::PjRtBuffer>> HostTensorToPjRtBuffer(
const tensorflow::Tensor* cpu_tensor, tensorflow::Device* device,
xla::PjRtClient* pjrt_client,
const XlaShapeLayoutHelpers::ShapeDeterminationFns
shape_determination_fns) {
XlaLayoutPreference layout_preference =
shape_determination_fns.layout_preference_fn(
cpu_tensor->shape(), cpu_tensor->dtype(), std::nullopt);
TF_ASSIGN_OR_RETURN(xla::Shape shape,
shape_determination_fns.shape_representation_fn(
cpu_tensor->shape(), cpu_tensor->dtype(),
/*fast_mem=*/false, layout_preference));
const xla::Layout* device_layout = &(shape.layout());
// The device id should match the local_device_id in
// tensorflow/compiler/xla/pjrt/pjrt_client.h.
const int pjrt_device_id =
tsl::GetDeviceIdFromDeviceParsedName(device->parsed_name());
TF_ASSIGN_OR_RETURN(
xla::PjRtDevice * pjrt_device,
pjrt_client->LookupAddressableDevice(xla::LocalDeviceId(pjrt_device_id)));
TF_ASSIGN_OR_RETURN(xla::PjRtMemorySpace * pjrt_memory,
pjrt_device->default_memory_space());
auto first_try_buffer = pjrt_client->BufferFromHostBuffer(
cpu_tensor->data(), shape.element_type(), shape.dimensions(),
/*byte_strides=*/std::nullopt,
xla::PjRtClient::HostBufferSemantics::kImmutableZeroCopy,
/*on_done_with_host_buffer=*/
[cpu_tensor = *cpu_tensor]() { /* frees tensor */ }, pjrt_memory,
device_layout);
if (first_try_buffer.ok()) {
return std::move(*first_try_buffer);
}
if (first_try_buffer.status().code() == absl::StatusCode::kUnimplemented) {
LOG_FIRST_N(WARNING, 1)
<< first_try_buffer.status()
<< "; fallback to BufferFromHostBuffer without device layout.";
TF_ASSIGN_OR_RETURN(
std::unique_ptr<xla::PjRtBuffer> second_try_buffer,
pjrt_client->BufferFromHostBuffer(
cpu_tensor->data(), shape.element_type(), shape.dimensions(),
/*byte_strides=*/std::nullopt,
xla::PjRtClient::HostBufferSemantics::kImmutableZeroCopy,
/*on_done_with_host_buffer=*/
[cpu_tensor = *cpu_tensor]() { /* frees tensor */ }, pjrt_memory,
/*device_layout=*/nullptr));
return second_try_buffer;
} else {
return first_try_buffer.status();
}
}
} // namespace
void PjRtDeviceContext::CopyDeviceTensorToCPU(const Tensor* device_tensor,
absl::string_view tensor_name,
Device* device,
Tensor* cpu_tensor,
StatusCallback done) {
tsl::profiler::TraceMe traceme("PjRtDeviceContext::CopyDeviceTensorToCPU");
if (device_tensor->NumElements() == 0) {
VLOG(2) << "CopyDeviceTensorToCPU empty tensor";
done(absl::OkStatus());
return;
}
auto literal = std::make_unique<xla::MutableBorrowingLiteral>();
auto status = tensorflow::HostTensorToMutableBorrowingLiteral(cpu_tensor,
literal.get());
if (!status.ok()) {
done(status);
return;
}
xla::PjRtBuffer* device_buffer;
AsyncValueTensor* device_tensor_av =
tensorflow::AsyncValueTensor::FromTensor(device_tensor);
if (use_pjrt_tensor_buffer_) {
if (device_tensor_av) {
done(absl::InvalidArgumentError(
"If use_pjrt_tensor_buffer is set, the device tensor should not "
"contain an AsyncValueTensor."));
return;
}
const PjRtTensorBuffer* pjrt_tensor_buffer =
dynamic_cast<const PjRtTensorBuffer*>(DMAHelper::buffer(device_tensor));
if (pjrt_tensor_buffer == nullptr) {
done(absl::UnimplementedError(
"use_pjrt_tensor_buffer is set to true. Transferring a tensor "
"without pjrt_tensor_buffer in this case is not supported."));
return;
}
device_buffer = pjrt_tensor_buffer->pjrt_buffer();
} else {
device_buffer = device_tensor_av->GetBuffer().get();
}
if (device_buffer == nullptr) {
done(absl::InvalidArgumentError(
"The device tensor has no associated device buffer."));
return;
}
tsl::Future<void> future = device_buffer->ToLiteral(literal.get());
future.OnReady([literal = std::move(literal), done = std::move(done)](
const absl::Status& status) { done(status); });
}
void PjRtDeviceContext::CopyCPUTensorToDevice(const Tensor* cpu_tensor,
Device* device,
Tensor* device_tensor,
StatusCallback done,
bool sync_dst_compute) const {
tsl::profiler::TraceMe traceme("PjRtDeviceContext::CopyCPUTensorToDevice");
if (cpu_tensor->NumElements() == 0) {
VLOG(2) << "CopyCPUTensorToDevice empty tensor";
done(absl::OkStatus());
return;
}
absl::StatusOr<xla::PjRtClient*> pjrt_client =
GetOrCreatePjRtClient(DeviceType(device->device_type()));
if (!pjrt_client.ok()) {
done(pjrt_client.status());
return;
}
absl::StatusOr<std::unique_ptr<xla::PjRtBuffer>> buffer_or =
HostTensorToPjRtBuffer(cpu_tensor, device, *pjrt_client,
shape_determination_fns_);
if (!buffer_or.ok()) {
done(buffer_or.status());
return;
}
xla::PjRtBuffer* pjrt_buffer = (*buffer_or).get();
if (use_pjrt_tensor_buffer_) {
// Copy the newly created tensor with PjRtTensorBuffer to output device
// tensor.
absl::StatusOr<Tensor> t = MakeTensorFromPjRtBuffer(
device_tensor->dtype(), device_tensor->shape(), std::move(*buffer_or));
if (!t.ok()) {
done(t.status());
return;
}
*device_tensor = *t;
} else {
AsyncValueTensor* result_tensor =
tensorflow::AsyncValueTensor::FromTensor(device_tensor);
// The result tensor should be newly allocated, which does not point to a
// valid buffer yet.
CHECK(!result_tensor->GetBuffer()); // Crash OK
result_tensor->SetBuffer(std::move(*buffer_or));
}
pjrt_buffer->GetReadyFuture().OnReady(std::move(done));
}
void PjRtDeviceContext::CopyTensorInSameDevice(const Tensor* input_tensor,
Device* device,
Tensor* output_tensor,
StatusCallback done) const {
if (!DeviceFactory::IsPluggableDevice(device->device_type())) {
done(absl::UnimplementedError(
"Same-device copies in PjRtDeviceContext is only implemented when "
"is_pluggable_device is true."));
return;
}
// TODO(b/288585098): consider whether to support same device copy in PJRT
// API.
absl::StatusOr<PJRT_Buffer*> c_src_buffer =
GetPjRtCBufferFromTensor(input_tensor);
if (!c_src_buffer.ok()) {
done(c_src_buffer.status());
return;
}
absl::StatusOr<xla::PjRtCApiClient*> c_api_client =
tensorflow::GetPjRtCApiClient(
tensorflow::DeviceType(device->device_type()));
if (!c_api_client.ok()) {
done(c_api_client.status());
return;
}
TSL_Status c_status;
PJRT_Buffer* dst_buffer = TfnpdApi()->TFNPD_SameDevicePjRtBufferCopy(
*c_src_buffer, (*c_api_client)->pjrt_c_client(), &c_status);
if (!c_status.status.ok()) {
done(c_status.status);
return;
}
auto set_c_buffer_status =
SetPjRtCBufferToTensor(dst_buffer, *c_api_client, output_tensor);
if (!set_c_buffer_status.ok()) {
done(set_c_buffer_status);
return;
}
AsyncValueTensor* result_tensor =
tensorflow::AsyncValueTensor::FromTensor(output_tensor);
result_tensor->GetBuffer()->GetReadyFuture().OnReady(std::move(done));
}
void PjRtDeviceToDeviceCopy(DeviceContext* send_dev_context,
DeviceContext* recv_dev_context, Device* src,
Device* dst, AllocatorAttributes src_alloc_attr,
AllocatorAttributes dst_alloc_attr,
const Tensor* input, Tensor* output,
int dev_to_dev_stream_index, StatusCallback done) {
tsl::profiler::TraceMe traceme("PjRtDevice_DeviceToDeviceCopy");
if (input->NumElements() == 0) {
VLOG(2) << "PjRtDevice_DeviceToDeviceCopy empty tensor";
done(absl::OkStatus());
return;
}
absl::StatusOr<xla::PjRtClient*> pjrt_dst_client =
GetOrCreatePjRtClient(DeviceType(dst->device_type()));
if (!pjrt_dst_client.ok()) {
done(pjrt_dst_client.status());
return;
}
xla::PjRtBuffer* src_device_buffer =
tensorflow::AsyncValueTensor::FromTensor(input)->GetBuffer().get();
// The device id should match the local_device_id in
// tensorflow/compiler/xla/pjrt/pjrt_client.h.
const int pjrt_dst_device_id =
tsl::GetDeviceIdFromDeviceParsedName(dst->parsed_name());
xla::PjRtDevice* pjrt_dst_device =
(*pjrt_dst_client)
->LookupAddressableDevice(xla::LocalDeviceId(pjrt_dst_device_id))
.value();
absl::StatusOr<std::unique_ptr<xla::PjRtBuffer>> buffer_or =
src_device_buffer->CopyToMemorySpace(
*pjrt_dst_device->default_memory_space());
if (!buffer_or.ok()) {
done(buffer_or.status());
return;
}
xla::PjRtBuffer* pjrt_buffer = (*buffer_or).get();
if (static_cast<PjRtDeviceContext*>(recv_dev_context)
->use_pjrt_tensor_buffer()) {
// Copy the newly created tensor with PjRtTensorBuffer to output device
// tensor.
absl::StatusOr<Tensor> t = MakeTensorFromPjRtBuffer(
output->dtype(), output->shape(), std::move(*buffer_or));
if (!t.ok()) {
done(t.status());
return;
}
*output = *t;
} else {
AsyncValueTensor* output_tensor =
tensorflow::AsyncValueTensor::FromTensor(output);
// The result tensor should be newly allocated, which does not point to a
// valid buffer yet.
CHECK(!output_tensor->GetBuffer()); // Crash OK
output_tensor->SetBuffer(std::move(*buffer_or));
}
pjrt_buffer->GetReadyFuture().OnReady(std::move(done));
}
} // namespace tensorflow
tensorflow/compiler/jit/pjrt_device_context.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_COMPILER_JIT_PJRT_DEVICE_CONTEXT_H_
#define TENSORFLOW_COMPILER_JIT_PJRT_DEVICE_CONTEXT_H_
#include <utility>
#include "tensorflow/compiler/tf2xla/layout_util.h"
#include "tensorflow/core/framework/device_base.h"
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
// Helper class for managing data transfers between host and accelerator
// devices using PjRt.
class PjRtDeviceContext : public DeviceContext {
public:
explicit PjRtDeviceContext(
XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,
bool use_pjrt_tensor_buffer = false)
: shape_determination_fns_(std::move(shape_determination_fns)),
use_pjrt_tensor_buffer_(use_pjrt_tensor_buffer) {}
void CopyCPUTensorToDevice(const Tensor* cpu_tensor, Device* device,
Tensor* device_tensor, StatusCallback done,
bool sync_dst_compute) const override;
void CopyDeviceTensorToCPU(const Tensor* device_tensor,
absl::string_view tensor_name, Device* device,
Tensor* cpu_tensor, StatusCallback done) override;
void CopyTensorInSameDevice(const Tensor* input_tensor, Device* device,
Tensor* output_tensor,
StatusCallback done) const override;
bool use_pjrt_tensor_buffer() const { return use_pjrt_tensor_buffer_; }
private:
XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns_;
// Note: we currently assume the PjRtBuffer is a PjRtStreamExecutorBuffer.
bool use_pjrt_tensor_buffer_;
};
void PjRtDeviceToDeviceCopy(DeviceContext* send_dev_context,
DeviceContext* recv_dev_context, Device* src,
Device* dst, AllocatorAttributes src_alloc_attr,
AllocatorAttributes dst_alloc_attr,
const Tensor* input, Tensor* output,
int dev_to_dev_stream_index, StatusCallback done);
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_PJRT_DEVICE_CONTEXT_H_
tensorflow/compiler/jit/pjrt_tensor_buffer.h
+57
View File
@@ -0,0 +1,57 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_JIT_PJRT_TENSOR_BUFFER_H_
#define TENSORFLOW_COMPILER_JIT_PJRT_TENSOR_BUFFER_H_
#include <memory>
#include <utility>
#include "xla/pjrt/pjrt_client.h"
#include "tensorflow/core/framework/allocation_description.pb.h"
#include "tensorflow/core/framework/tensor.h"
namespace tensorflow {
// PjRtTensorBuffer is derived from TensorBuffer, which holds a device memory
// pointer so that legacy TF kernel can access it directly. PjRtTensorBuffer
// also owns a PjRtBuffer for XLA kernel's usage.
class PjRtTensorBuffer : public TensorBuffer {
public:
PjRtTensorBuffer(const void* ptr, size_t expected_size,
std::unique_ptr<xla::PjRtBuffer> pjrt_buffer)
: TensorBuffer(const_cast<void*>(ptr)),
expected_size_(expected_size),
pjrt_buffer_(std::move(pjrt_buffer)) {}
size_t size() const override { return expected_size_; }
TensorBuffer* root_buffer() override { return this; }
xla::PjRtBuffer* pjrt_buffer() const { return pjrt_buffer_.get(); }
// TODO(b/288965065): Implement this.
void FillAllocationDescription(AllocationDescription* proto) const override {
proto->set_requested_bytes(static_cast<int64_t>(expected_size_));
}
private:
size_t expected_size_;
std::unique_ptr<xla::PjRtBuffer> pjrt_buffer_;
};
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_PJRT_TENSOR_BUFFER_H_
tensorflow/compiler/jit/pjrt_tensor_buffer_util.cc
+89
View File
@@ -0,0 +1,89 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/jit/pjrt_tensor_buffer_util.h"
#include <cstddef>
#include <memory>
#include <utility>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "tensorflow/compiler/jit/pjrt_tensor_buffer.h"
#include "xla/pjrt/pjrt_client.h"
#include "tensorflow/core/common_runtime/dma_helper.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.h"
#include "tsl/platform/statusor.h"
namespace tensorflow {
static size_t GetTensorSize(const TensorShape& shape, const DataType dtype) {
return shape.num_elements() * DataTypeSize(dtype);
}
absl::StatusOr<Tensor> MakeTensorFromPjRtBuffer(
const DataType dtype, const TensorShape& shape,
std::unique_ptr<xla::PjRtBuffer> pjrt_buffer) {
TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::PjRtBuffer::ExternalReference> ref,
pjrt_buffer->AcquireExternalReference());
auto* tensor_buffer =
new PjRtTensorBuffer(ref->OpaqueDeviceMemoryDataPointer(),
GetTensorSize(shape, dtype), std::move(pjrt_buffer));
Tensor result(dtype, shape, tensor_buffer);
tensor_buffer->Unref();
return result;
}
// If existing_tensor does not use PjRtTensorBuffer and the opaque device memory
// is the same, the tensor should be reused so that the same device memory will
// not be double-freed.
static bool ShouldReuseTensor(void* opaque_device_memory,
const size_t expected_size,
const Tensor* existing_tensor) {
const PjRtTensorBuffer* input_pjrt_tensor_buffer =
dynamic_cast<const PjRtTensorBuffer*>(DMAHelper::buffer(existing_tensor));
if (input_pjrt_tensor_buffer != nullptr) {
return false;
}
const size_t current_size =
GetTensorSize(existing_tensor->shape(), existing_tensor->dtype());
return existing_tensor->tensor_data().data() == opaque_device_memory &&
current_size == expected_size;
}
absl::Status PjRtTensorBufferUtil::UpdateOrMakeTensorWithPjRtBuffer(
const DataType dtype, const TensorShape& shape,
std::unique_ptr<xla::PjRtBuffer> pjrt_buffer, Tensor* output_tensor) {
TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::PjRtBuffer::ExternalReference> ref,
pjrt_buffer->AcquireExternalReference());
const size_t expected_size = GetTensorSize(shape, dtype);
void* opaque_device_memory = ref->OpaqueDeviceMemoryDataPointer();
auto* tensor_buffer = new PjRtTensorBuffer(
opaque_device_memory, expected_size, std::move(pjrt_buffer));
if (ShouldReuseTensor(opaque_device_memory, expected_size, output_tensor)) {
output_tensor->buf_ = tensor_buffer;
return absl::OkStatus();
}
Tensor result(dtype, shape, tensor_buffer);
tensor_buffer->Unref();
*output_tensor = result;
return absl::OkStatus();
}
} // namespace tensorflow
tensorflow/compiler/jit/pjrt_tensor_buffer_util.h
+56
View File
@@ -0,0 +1,56 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_JIT_PJRT_TENSOR_BUFFER_UTIL_H_
#define TENSORFLOW_COMPILER_JIT_PJRT_TENSOR_BUFFER_UTIL_H_
#include <memory>
#include "absl/status/statusor.h"
#include "tensorflow/compiler/jit/pjrt_tensor_buffer.h"
#include "xla/pjrt/pjrt_client.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
namespace tensorflow {
// Takes the device memory pointer from the PjRtBuffer and create a Tensor that
// contains a PjRtTensorBuffer. The PjRtTensorBuffer holds the pointer to the
// device memory. It also owns the PjRtBuffer.
//
// TODO(b/289001822): Create a unit test to cover this function.
absl::StatusOr<Tensor> MakeTensorFromPjRtBuffer(
DataType dtype, const TensorShape& shape,
std::unique_ptr<xla::PjRtBuffer> pjrt_buffer);
// For TensorFlow internal use only.
class PjRtTensorBufferUtil {
public:
// Takes the device memory pointer from the PjRtBuffer and create a
// PjRtTensorBuffer. The PjRtTensorBuffer holds the pointer to the device
// memory. It also owns the PjRtBuffer. If output_tensor does not use
// PjRtTensorBuffer and the opaque device memory is the same, update the
// output_tensor->buf_ so that the same device memory will not be double-free.
// Otherwise a new Tensor will be created with the PjRtTensorBuffer.
//
// TODO(b/289001822): Create a unit test to cover this function.
static absl::Status UpdateOrMakeTensorWithPjRtBuffer(
DataType dtype, const TensorShape& shape,
std::unique_ptr<xla::PjRtBuffer> pjrt_buffer, Tensor* output_tensor);
};
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_PJRT_TENSOR_BUFFER_UTIL_H_
tensorflow/compiler/jit/pjrt_tensor_buffer_util_test.cc
+72
View File
@@ -0,0 +1,72 @@
/* 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/compiler/jit/pjrt_tensor_buffer_util.h"
#include <cstdint>
#include <optional>
#include <utility>
#include <vector>
#include <gtest/gtest.h>
#include "tensorflow/compiler/jit/test_util.h"
#include "xla/pjrt/pjrt_client.h"
#include "xla/shape.h"
#include "xla/shape_util.h"
#include "tensorflow/core/framework/allocator.h"
#include "tensorflow/core/framework/device.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/tfrt/common/pjrt_util.h"
#include "tsl/platform/statusor.h"
namespace tensorflow {
namespace {
TEST(PjRtTensorBufferUtilTest, MakeTensorFromPjRtBuffer) {
DeviceSetup device_setup;
device_setup.AddDevicesAndSetUp({DEVICE_GPU});
Device* device = device_setup.GetDevice(DEVICE_GPU);
std::vector<int64_t> dimensions = {2, 3};
Tensor dest_cpu_tensor(cpu_allocator(), tensorflow::DT_INT32,
tensorflow::TensorShape(dimensions));
TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client, GetPjRtClient(DEVICE_GPU));
std::vector<int32_t> data{1, 2, 3, 4, 5, 6};
xla::Shape xla_shape = xla::ShapeUtil::MakeShape(xla::S32, dimensions);
xla::PjRtDevice* pjrt_device = pjrt_client->addressable_devices()[0];
TF_ASSERT_OK_AND_ASSIGN(xla::PjRtMemorySpace * pjrt_memory,
pjrt_device->default_memory_space());
TF_ASSERT_OK_AND_ASSIGN(
auto pjrt_buffer,
pjrt_client->BufferFromHostBuffer(
data.data(), xla_shape.element_type(), xla_shape.dimensions(),
/*byte_strides=*/std::nullopt,
xla::PjRtClient::HostBufferSemantics::kImmutableOnlyDuringCall,
nullptr, pjrt_memory, /*device_layout=*/nullptr));
TF_ASSERT_OK_AND_ASSIGN(
Tensor tensor, MakeTensorFromPjRtBuffer(DT_INT32, TensorShape(dimensions),
std::move(pjrt_buffer)));
auto s = device->tensorflow_accelerator_device_info()
->pjrt_context->CopyDeviceTensorToCPUSync(&tensor, "", device,
&dest_cpu_tensor);
for (int i = 0; i < tensor.NumElements(); ++i) {
EXPECT_EQ(dest_cpu_tensor.flat<int32_t>().data()[i], data[i]);
}
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/rearrange_function_argument_pass_test.cc
+240
View File
@@ -0,0 +1,240 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <initializer_list>
#include <memory>
#include <utility>
#include <vector>
#include "absl/log/check.h"
#include "absl/status/status.h"
#include "tensorflow/cc/framework/ops.h"
#include "tensorflow/cc/framework/scope.h"
#include "tensorflow/cc/ops/const_op.h"
#include "tensorflow/cc/ops/function_ops.h"
#include "tensorflow/cc/ops/functional_ops.h"
#include "tensorflow/compiler/tf2xla/rearrange_function_argument.h"
#include "xla/tsl/platform/errors.h"
#include "tensorflow/core/common_runtime/function.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/graph_to_functiondef.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"
namespace tensorflow {
TEST(RearrangeFunctionArgumentForFunctionTest, Basic) {
FunctionDefLibrary fdl;
{
// Function for StatefulPartitionedCall's "f", If's
// "then_branch"/"else_branch".
// "arg0" (T=DT_RESOURCE), "arg1" (T=DT_BOOL)
// "ret0" = "arg1"
// "ret1" = "arg0"
tensorflow::Scope s = tensorflow::Scope::NewRootScope();
Output arg0 = ops::_Arg(s.WithOpName("arg0"), DT_RESOURCE, 0);
Output arg1 = ops::_Arg(s.WithOpName("arg1"), DT_BOOL, 1);
auto ret0 = ops::_Retval(s.WithOpName("ret0"), arg1, 0);
auto ret1 = ops::_Retval(s.WithOpName("ret1"), arg0, 1);
std::unique_ptr<Graph> g(new Graph(OpRegistry::Global()));
CHECK_OK(s.ToGraph(g.get()));
FunctionDef *xla_fdef = fdl.add_function();
CHECK_OK(GraphToFunctionDef(*g, "f1", xla_fdef));
}
{
// Function for While's "body".
// "arg0" (T=DT_RESOURCE), "arg1" (T=DT_BOOL)
// "ret0" = "arg0"
// "ret1" = "arg1"
tensorflow::Scope s = tensorflow::Scope::NewRootScope();
Output arg0 = ops::_Arg(s.WithOpName("arg0"), DT_RESOURCE, 0);
Output arg1 = ops::_Arg(s.WithOpName("arg1"), DT_BOOL, 1);
auto ret0 = ops::_Retval(s.WithOpName("ret0"), arg0, 0);
auto ret1 = ops::_Retval(s.WithOpName("ret1"), arg1, 1);
std::unique_ptr<Graph> g(new Graph(OpRegistry::Global()));
CHECK_OK(s.ToGraph(g.get()));
FunctionDef *xla_fdef = fdl.add_function();
CHECK_OK(GraphToFunctionDef(*g, "f2", xla_fdef));
}
{
// Function for While's "cond".
// "arg0" (T=DT_RESOURCE), "arg1" (T=DT_BOOL)
// "ret0" = "arg1"
tensorflow::Scope s = tensorflow::Scope::NewRootScope();
Output arg0 = ops::_Arg(s.WithOpName("arg0"), DT_RESOURCE, 0);
Output arg1 = ops::_Arg(s.WithOpName("arg1"), DT_BOOL, 1);
auto ret0 = ops::_Retval(s.WithOpName("ret0"), arg1, 0);
std::unique_ptr<Graph> g(new Graph(OpRegistry::Global()));
CHECK_OK(s.ToGraph(g.get()));
FunctionDef *xla_fdef = fdl.add_function();
CHECK_OK(GraphToFunctionDef(*g, "f3", xla_fdef));
}
FunctionLibraryDefinition fld(OpRegistry::Global(), fdl);
// Build the XLA computation graph.
// "arg0" (T=DT_RESOURCE), "arg1" (T=DT_INT32)
// "arg0", "arg1" -> "if" (If) -> "ret0", "ret1"
// "arg0", "arg1" -> "while" (While) -> "ret2", "ret3"
tensorflow::Scope s = tensorflow::Scope::NewRootScope();
Output arg0 = ops::_Arg(s.WithOpName("arg0"), DT_RESOURCE, 0);
Output arg1 = ops::_Arg(s.WithOpName("arg1"), DT_BOOL, 1);
NameAttrList f;
f.set_name("f1");
auto if_op = ops::If(s.WithOpName("if"), arg1,
std::initializer_list<Input>{arg0, arg1},
{DT_BOOL, DT_RESOURCE}, f, f);
auto ret0 = ops::_Retval(s.WithOpName("ret0"), if_op.output[0], 0);
auto ret1 = ops::_Retval(s.WithOpName("ret1"), if_op.output[1], 1);
NameAttrList cond_fn, body_fn;
cond_fn.set_name("f3");
body_fn.set_name("f2");
auto while_op =
ops::While(s.WithOpName("while"),
std::initializer_list<Input>{arg0, arg1}, cond_fn, body_fn);
auto ret2 = ops::_Retval(s.WithOpName("ret2"), while_op.output[0], 2);
auto ret3 = ops::_Retval(s.WithOpName("ret3"), while_op.output[1], 3);
std::unique_ptr<Graph> g(new Graph(OpRegistry::Global()));
CHECK_OK(s.ToGraph(g.get()));
std::vector<std::unique_ptr<FunctionBody>> fbodies;
CHECK_OK(RearrangeFunctionArguments(
[&](const NameAttrList& function, const FunctionBody** fbody) {
std::unique_ptr<FunctionBody> new_fbody;
TF_RETURN_IF_ERROR(FunctionDefToBodyHelper(*fld.Find(function.name()),
AttrSlice(&function.attr()),
&fld, &new_fbody));
*fbody = new_fbody.get();
fbodies.push_back(std::move(new_fbody));
return absl::OkStatus();
},
g.get(), &fld));
// Check function f1_rearrange_0, input types should be {DT_BOOL, DT_RESOURCE}
// and output types should be {DT_BOOL}.
const FunctionDef *f1_rewritten = fld.Find("f1_rearrange_0");
CHECK_NE(f1_rewritten, nullptr);
ASSERT_EQ(f1_rewritten->signature().input_arg_size(), 2);
EXPECT_EQ(f1_rewritten->signature().input_arg(0).type(), DT_BOOL);
EXPECT_EQ(f1_rewritten->signature().input_arg(1).type(), DT_RESOURCE);
ASSERT_EQ(f1_rewritten->signature().output_arg_size(), 1);
EXPECT_EQ(f1_rewritten->signature().output_arg(0).type(), DT_BOOL);
// Check node "if" input and output edges.
auto node_name_index = g->BuildNodeNameIndex();
const Node *if_node = node_name_index.at("if");
ASSERT_NE(if_node, nullptr);
const Node *input_node;
CHECK_OK(if_node->input_node(1, &input_node));
EXPECT_EQ(input_node->name(), "arg1");
CHECK_OK(if_node->input_node(2, &input_node));
EXPECT_EQ(input_node->name(), "arg0");
const Node *ret0_node = node_name_index.at("ret0");
ASSERT_NE(ret0_node, nullptr);
CHECK_OK(ret0_node->input_node(0, &input_node));
EXPECT_EQ(input_node->name(), "if");
const Node *ret1_node = node_name_index.at("ret1");
ASSERT_NE(ret1_node, nullptr);
CHECK_OK(ret1_node->input_node(0, &input_node));
EXPECT_EQ(input_node->name(), "arg0");
// Check node "while" input and output edges.
const Node *while_node = node_name_index.at("while");
ASSERT_NE(while_node, nullptr);
CHECK_OK(while_node->input_node(0, &input_node));
EXPECT_EQ(input_node->name(), "arg1");
CHECK_OK(while_node->input_node(1, &input_node));
EXPECT_EQ(input_node->name(), "arg0");
const Node *ret2_node = node_name_index.at("ret2");
ASSERT_NE(ret2_node, nullptr);
CHECK_OK(ret2_node->input_node(0, &input_node));
EXPECT_EQ(input_node->name(), "arg0");
const Node *ret3_node = node_name_index.at("ret3");
ASSERT_NE(ret3_node, nullptr);
CHECK_OK(ret3_node->input_node(0, &input_node));
EXPECT_EQ(input_node->name(), "while");
}
TEST(RearrangeFunctionArgumentForFunctionTest,
WhileResourceRetvalFromDifferentArgUnimplemented) {
FunctionDefLibrary fdl;
{
// Function for While's "body".
// "arg0" (T=DT_RESOURCE), "arg1" (T=DT_RESOURCE), "arg2" (T=DT_INT32)
// "ret0" = "arg1"
// "ret1" = "arg0"
tensorflow::Scope s = tensorflow::Scope::NewRootScope();
Output arg0 = ops::_Arg(s.WithOpName("arg0"), DT_RESOURCE, 0);
Output arg1 = ops::_Arg(s.WithOpName("arg1"), DT_RESOURCE, 1);
Output arg2 = ops::_Arg(s.WithOpName("arg2"), DT_INT32, 2);
auto ret0 = ops::_Retval(s.WithOpName("ret0"), arg1, 0);
auto ret1 = ops::_Retval(s.WithOpName("ret1"), arg0, 1);
auto ret2 = ops::_Retval(s.WithOpName("ret2"), arg2, 2);
std::unique_ptr<Graph> g(new Graph(OpRegistry::Global()));
CHECK_OK(s.ToGraph(g.get()));
FunctionDef *xla_fdef = fdl.add_function();
CHECK_OK(GraphToFunctionDef(*g, "f2", xla_fdef));
}
{
// Function for While's "cond".
// "arg0" (T=DT_RESOURCE), "arg1" (T=DT_RESOURCE), "arg2" (T=DT_INT32)
// "ret0" = true
tensorflow::Scope s = tensorflow::Scope::NewRootScope();
Output arg0 = ops::_Arg(s.WithOpName("arg0"), DT_RESOURCE, 0);
Output arg1 = ops::_Arg(s.WithOpName("arg1"), DT_RESOURCE, 1);
Output arg2 = ops::_Arg(s.WithOpName("arg2"), DT_INT32, 2);
Output cond = ops::Const(s.WithOpName("const"), true, TensorShape({}));
auto ret0 = ops::_Retval(s.WithOpName("ret0"), cond, 0);
std::unique_ptr<Graph> g(new Graph(OpRegistry::Global()));
CHECK_OK(s.ToGraph(g.get()));
FunctionDef *xla_fdef = fdl.add_function();
CHECK_OK(GraphToFunctionDef(*g, "f1", xla_fdef));
}
FunctionLibraryDefinition fld(OpRegistry::Global(), fdl);
// Build the XLA computation graph.
// "arg0" (T=DT_RESOURCE), "arg1" (T=DT_RESOURCE), "arg2" (T=DT_INT32)
// "arg0", "arg1" -> "while" (While)
tensorflow::Scope s = tensorflow::Scope::NewRootScope();
Output arg0 = ops::_Arg(s.WithOpName("arg0"), DT_RESOURCE, 0);
Output arg1 = ops::_Arg(s.WithOpName("arg1"), DT_RESOURCE, 1);
Output arg2 = ops::_Arg(s.WithOpName("arg2"), DT_INT32, 2);
NameAttrList cond_fn, body_fn;
cond_fn.set_name("f1");
body_fn.set_name("f2");
auto while_op = ops::While(s.WithOpName("while"),
std::initializer_list<Input>{arg0, arg1, arg2},
cond_fn, body_fn);
std::unique_ptr<Graph> g(new Graph(OpRegistry::Global()));
CHECK_OK(s.ToGraph(g.get()));
std::vector<std::unique_ptr<FunctionBody>> fbodies;
absl::Status status = RearrangeFunctionArguments(
[&](const NameAttrList &function, const FunctionBody **fbody) {
std::unique_ptr<FunctionBody> new_fbody;
TF_RETURN_IF_ERROR(FunctionDefToBodyHelper(*fld.Find(function.name()),
AttrSlice(&function.attr()),
&fld, &new_fbody));
*fbody = new_fbody.get();
fbodies.push_back(std::move(new_fbody));
return absl::OkStatus();
},
g.get(), &fld);
EXPECT_EQ(status.code(), error::UNIMPLEMENTED);
}
} // namespace tensorflow
tensorflow/compiler/jit/report_clustering_info_pass.cc
+32
View File
@@ -0,0 +1,32 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/jit/report_clustering_info_pass.h"
#include "tensorflow/compiler/jit/flags.h"
#include "tensorflow/compiler/jit/xla_activity_listener.h"
#include "tensorflow/compiler/jit/xla_cluster_util.h"
namespace tensorflow {
absl::Status ReportClusteringInfoPass::Run(
const GraphOptimizationPassOptions& options) {
XlaAutoClusteringActivity activity;
*activity.mutable_summary() = GetXlaAutoClusteringSummary(**options.graph);
activity.set_global_jit_level(GetGlobalJitLevelForGraph(options));
activity.set_cpu_global_jit_enabled(
GetMarkForCompilationPassFlags()->tf_xla_cpu_global_jit);
return BroadcastXlaActivity(std::move(activity));
}
} // namespace tensorflow
tensorflow/compiler/jit/report_clustering_info_pass.h
+32
View File
@@ -0,0 +1,32 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_JIT_REPORT_CLUSTERING_INFO_PASS_H_
#define TENSORFLOW_COMPILER_JIT_REPORT_CLUSTERING_INFO_PASS_H_
#include "tensorflow/core/common_runtime/optimization_registry.h"
namespace tensorflow {
// This is not really an optimization pass. It does not change the graph in any
// way; instead it computes a summary of the XLA clusters in the graph and
// broadcasts it via xla_activity_listener.
class ReportClusteringInfoPass : public GraphOptimizationPass {
public:
absl::Status Run(const GraphOptimizationPassOptions& options) override;
};
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_REPORT_CLUSTERING_INFO_PASS_H_
tensorflow/compiler/jit/resource_operation_safety_analysis.cc
+321
View File
@@ -0,0 +1,321 @@
/* Copyright 2018 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.
==============================================================================*/
// ALGORITHM OVERVIEW
// ==================
//
// An XLA cluster hoists all resource reads to be beginning of the cluster
// execution and all the resource writes to the end. This means it cannot
// enforce arbitrary ordering dependencies (via control or data edges) between
// resource operations. Since all resource reads happen before all resource
// writes, edges constraining resource reads to happen before resource writes
// are fine, but all other kinds of edges are problematic. This analysis
// computes the set of pairs of resource operations that cannot be put in the
// same cluster because XLA cannot respect the dependencies between them in the
// TensorFlow program.
//
// TODO(b/112856632): We can, in theory, support Read->Read and Write->Write
// dependencies.
//
// Specifically the result computed by this analysis contains the edge {W, R}
// iff all of these hold true:
//
// - In the graph (g - {edges from NextIteration to Merge}) there is a path
// from W to R.
// - IsEdgeSafe(W, R) == False [defined below]
// - W != R (note: some resource operations both read from and write to
// resource variables).
//
// The result is incorrect around loops because we ignore edges from
// NextIteration to Merge. For instance, in:
//
// Init -----> Merge <-------+
// | |
// v |
// Read |
// | |
// v |
// Write |
// | |
// v |
// NextIteration --+
//
// we won't put (Read, Write) in the returned set. This is fine if
// auto-clustering can only cluster the Read->Write edge, but it is a problem if
// it clusters the Write->NextIteration->Merge->Read edges instead. So we rely
// on auto-clustering to not cluster NextIteration->Merge edges. The same
// problem is present for the functional version of the loop above and we also
// rely on auto-clustering not clustering functional while loops containing
// resource operations.
//
// One way to think about this is that we only care about cases where two nodes,
// A and B, would normally have been put in the same cluster but cannot legally
// be in the same cluster because of resourcevar-dependencies. If A and B would
// normally have been put in the same cluster then all paths between A and B
// would have to be clusterable (otherwise we'd have introduced a cycle). Ergo
// there could not have been a NextIteration->Merge edge between A and B since
// we don't cluster these edges.
//
// IMPLEMENTATION
// --------------
//
// We traverse the graph minus backedges in reverse post order, mapping each
// node to the set of resource operation reaching that node. Since we visit
// producers before consumers, we can construct the set of reaching operations
// by taking the union of the operations reaching the input nodes. These
// "reaching resource operations" can then be used to create the pairs of
// incompatible nodes using `IsEdgeSafe`.
#include "tensorflow/compiler/jit/resource_operation_safety_analysis.h"
#include "absl/container/flat_hash_set.h"
#include "absl/memory/memory.h"
#include "absl/strings/str_join.h"
#include "absl/types/optional.h"
#include "tensorflow/compiler/jit/xla_cluster_util.h"
#include "tensorflow/compiler/tf2xla/resource_operation_table.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/graph/algorithm.h"
#include "tensorflow/core/graph/tensor_id.h"
#include "tensorflow/core/lib/hash/hash.h"
namespace tensorflow {
namespace {
// Maps `n` to the XlaResourceOpKind corresponding to its operation. If `n` is
// not a resource operation recognized by XLA then sets `out_resource_op_kind`
// to nullopt.
absl::Status XlaResourceOpKindForNode(
const Node& n, const FunctionLibraryDefinition* flib_def,
const std::function<absl::Status(const Node&, bool*)>&
resource_ops_to_ignore,
std::optional<XlaResourceOpKind>* out_resource_op_kind) {
bool should_ignore = false;
if (resource_ops_to_ignore) {
TF_RETURN_IF_ERROR(resource_ops_to_ignore(n, &should_ignore));
}
if (should_ignore) {
*out_resource_op_kind = std::nullopt;
return absl::OkStatus();
}
const XlaResourceOpInfo* op_info = GetResourceOpInfoForOp(n.type_string());
if (op_info) {
*out_resource_op_kind = op_info->kind();
return absl::OkStatus();
}
// We conservatively assume that functions will both read and write resource
// variables. In the future we may consider doing some form of
// inter-procedural analysis.
if (MayCallFunction(n, flib_def)) {
*out_resource_op_kind = XlaResourceOpKind::kReadWrite;
} else {
*out_resource_op_kind = std::nullopt;
}
return absl::OkStatus();
}
// Returns true if a control or data dependence from a TensorFlow operation of
// resource op kind `from` to a TensorFlow operation of resource op kind `to`
// can be represented by an XLA cluster and needs no special handling around
// auto-jit.
bool IsEdgeSafe(XlaResourceOpKind from, XlaResourceOpKind to) {
// XLA clusters force all reads to happen before all writes. Moreover the set
// of reads are executed as one atomic operation, and the set of writes are as
// another atomic operation. This means we can faithfully represent the
// following edges: Read->*, *->Write.
return from == XlaResourceOpKind::kRead || to == XlaResourceOpKind::kWrite;
}
using ResourceOp = std::pair<int, XlaResourceOpKind>;
std::string ResourceOpToString(const ResourceOp& resource_op) {
return absl::StrCat(
resource_op.first, ": ",
XlaResourceOpInfo::XlaResourceOpKindToString(resource_op.second));
}
// A copy-on-write set used to store the set of ResourceOps reaching a node in a
// TensorFlow graph.
//
// TODO(sanjoy): It may be useful to pull this out into its own header at some
// point.
class ResourceOpSet {
private:
using Impl = absl::flat_hash_set<ResourceOp>;
public:
ResourceOpSet() = default;
// Adds all ResourceOp s in `other` to this set.
void Add(const ResourceOpSet& other) {
CHECK(!frozen_);
if (other.impl_ == impl_) {
other.frozen_ = true;
return;
}
if (!impl_) {
other.frozen_ = true;
impl_ = other.impl_;
return;
}
for (ResourceOp resource_op : other) {
Add(resource_op);
}
}
void Add(const ResourceOp& resource_op) {
CHECK(!frozen_);
if (!IsCopy() && Contains(resource_op)) {
// We can avoid the copy if the item we want to insert already exists.
return;
}
EnsureIsCopied();
impl_->insert(resource_op);
}
Impl::const_iterator begin() const {
return impl_ ? impl_->begin() : GetEmptyImpl()->begin();
}
Impl::const_iterator end() const {
return impl_ ? impl_->end() : GetEmptyImpl()->end();
}
bool Contains(const ResourceOp& resource_op) const {
return impl_ != nullptr && impl_->count(resource_op);
}
private:
bool IsCopy() const { return storage_ != nullptr; }
void EnsureIsCopied() {
if (storage_ == nullptr) {
storage_ = std::make_unique<Impl>();
for (ResourceOp op : *this) {
storage_->insert(op);
}
impl_ = storage_.get();
}
}
static Impl* GetEmptyImpl() {
static Impl* empty_impl = new Impl;
return empty_impl;
}
Impl* impl_ = nullptr;
std::unique_ptr<Impl> storage_;
// frozen_ is true if there is another set pointing to this set's impl_. We
// can no longer add elements to this set in that case since the sets pointing
// to this set expect the contents of this set to be stable.
mutable bool frozen_ = false;
ResourceOpSet(const ResourceOpSet&) = delete;
void operator=(const ResourceOpSet&) = delete;
};
std::string ResourceOpSetToString(const ResourceOpSet& resource_op_set) {
std::vector<std::string> elements_debug_string;
std::transform(resource_op_set.begin(), resource_op_set.end(),
std::back_inserter(elements_debug_string), ResourceOpToString);
return absl::StrCat("{", absl::StrJoin(elements_debug_string, ","), "}");
}
std::string NodeToString(const Node& n, XlaResourceOpKind resource_op_kind) {
return absl::StrCat(
"[", n.name(), ": ", n.type_string(), "(",
XlaResourceOpInfo::XlaResourceOpKindToString(resource_op_kind), ")", "]");
}
} // namespace
absl::Status ComputeIncompatibleResourceOperationPairs(
const Graph& g, const FunctionLibraryDefinition* flib_def,
const std::function<absl::Status(const Node&, bool*)>&
resource_ops_to_ignore,
std::vector<std::pair<int, int>>* result) {
CHECK(result->empty());
std::vector<Node*> rpo;
GetReversePostOrder(g, &rpo, /*stable_comparator=*/NodeComparatorName(),
/*edge_filter=*/[](const Edge& edge) {
return !edge.src()->IsNextIteration();
});
auto resource_op_set_for_node =
std::make_unique<ResourceOpSet[]>(g.num_node_ids());
const bool vlog = VLOG_IS_ON(2);
for (Node* n : rpo) {
std::optional<XlaResourceOpKind> op_kind;
TF_RETURN_IF_ERROR(XlaResourceOpKindForNode(
*n, flib_def, resource_ops_to_ignore, &op_kind));
ResourceOpSet* resource_op_set = &resource_op_set_for_node[n->id()];
// Merge the reaching resource operations for all the incoming edges to
// create the set of all possible resource ops reaching `n`.
for (const Edge* e : n->in_edges()) {
if (n->IsMerge() && e->src()->IsNextIteration()) {
// Ignore back-edges (see file comment).
continue;
}
const ResourceOpSet& incoming_op_set =
resource_op_set_for_node[e->src()->id()];
resource_op_set->Add(incoming_op_set);
}
// Add to the "incompatible resource ops" set if necessary.
if (op_kind) {
for (ResourceOp incoming_op : *resource_op_set) {
if (IsEdgeSafe(incoming_op.second, *op_kind)) {
continue;
}
if (vlog) {
VLOG(2) << "Unsafe edge: "
<< NodeToString(*g.FindNodeId(incoming_op.first),
incoming_op.second)
<< " -> " << NodeToString(*n, *op_kind);
}
result->push_back({incoming_op.first, n->id()});
}
// Some graphs might have a lot of 'kRead' kinds, but they are always safe
// for incoming ops, so not storing them might save a lot of memory.
if (op_kind != XlaResourceOpKind::kRead) {
resource_op_set->Add({n->id(), *op_kind});
}
}
if (vlog) {
VLOG(3) << n->name() << " -> " << ResourceOpSetToString(*resource_op_set);
}
}
std::sort(result->begin(), result->end());
CHECK(std::unique(result->begin(), result->end()) == result->end());
return absl::OkStatus();
}
} // namespace tensorflow
tensorflow/compiler/jit/resource_operation_safety_analysis.h
+69
View File
@@ -0,0 +1,69 @@
/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_JIT_RESOURCE_OPERATION_SAFETY_ANALYSIS_H_
#define TENSORFLOW_COMPILER_JIT_RESOURCE_OPERATION_SAFETY_ANALYSIS_H_
#include "xla/service/graphcycles/graphcycles.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/graph/graph.h"
namespace tensorflow {
// An XLA cluster hoists all resource reads to be beginning of the cluster
// execution and all the resource writes to the end. This means it cannot
// enforce arbitrary ordering dependencies (via control or data edges) between
// resource operations. Since all resource reads happen before all resource
// writes, edges constraining resource writes to happen before resource reads
// are problematic. This analysis returns the set of pairs of resource
// operations that cannot be put in the same cluster because XLA cannot respect
// the dependencies between them in the TensorFlow program.
//
// The restrictions are not transitive: it is fine to put A and C in the same
// cluster even if the returned set contains (A,B) and (B,C).
//
// In other words, if these pairs are seen as edges in an undirected graph of
// the nodes in `g` then auto-clustering is at least as constrained as the graph
// coloring problem on this graph.
//
//
// For instance if we auto-cluster all operations in this TensorFlow graph:
//
// AssignVariablepOp0 -> AssignVariableOp1
// |
// v
// ReadVariableOp0 -> ReadVariableOp1
//
// we will lose the AssignVariablepOp1 -> ReadVariableOp0. The ReadVariableOp0
// -> ReadVariableOp1 and AssignVariableOp0 -> AssignVariableOp1 edges will be
// respected by XlaLaunchOp though because all reads happen before all writes
// with that limited clustering..
//
//
// NB! The result computed by this analysis assumes that we don't auto-cluster
// back-edges (i.e. the edges from NextIteration to Merge).
//
// NB! The result computed by this analysis assumes that we don't auto-cluster
// functional control flow nodes containing resource operations.
//
// If `resource_ops_to_ignore` is set then nodes for which it returns true are
// ignored (we pretend these nodes are not resource operations).
absl::Status ComputeIncompatibleResourceOperationPairs(
const Graph& g, const FunctionLibraryDefinition* flib_def,
const std::function<absl::Status(const Node&, bool*)>&
resource_ops_to_ignore,
std::vector<std::pair<int, int>>* result);
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_RESOURCE_OPERATION_SAFETY_ANALYSIS_H_

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