From 14bb933f4250cf7c7eda5be7cf23ed942850c22c Mon Sep 17 00:00:00 2001 From: Trent Lo Date: Thu, 15 Aug 2019 17:23:31 -0700 Subject: [PATCH] [TF:XLA] Polish naming and comments for cluster_scoping_pass. --- .../compiler/jit/cluster_scoping_pass.cc | 20 +++++++++---------- .../compiler/jit/cluster_scoping_pass.h | 12 +++++------ .../compiler/jit/cluster_scoping_pass_test.cc | 11 +++++----- tensorflow/compiler/jit/defs.cc | 4 ++-- .../compiler/jit/mark_for_compilation_pass.cc | 10 ++++------ 5 files changed, 28 insertions(+), 29 deletions(-) diff --git a/tensorflow/compiler/jit/cluster_scoping_pass.cc b/tensorflow/compiler/jit/cluster_scoping_pass.cc index d51a872f898..7e1da87e5d5 100644 --- a/tensorflow/compiler/jit/cluster_scoping_pass.cc +++ b/tensorflow/compiler/jit/cluster_scoping_pass.cc @@ -52,7 +52,7 @@ class ClusterScopingPassImpl { size_t unique_scope_id_; }; -absl::optional GetXlaScope(Node* node) { +absl::optional GetXlaAutoJitScope(Node* node) { string scope; if (GetNodeAttr(node->attrs(), kXlaAutoJitScopeAttr, &scope).ok()) { return scope; @@ -61,23 +61,23 @@ absl::optional GetXlaScope(Node* node) { return absl::nullopt; } -void SetXlaScope(Node* node, StringPiece scope) { +void SetXlaAutoJitScope(Node* node, StringPiece scope) { node->AddAttr(kXlaAutoJitScopeAttr, scope); } -// NB! We append new scope as suffix to the XlaScope attribute instead of -// overriding the old value. In this way, we respect the original scopes. +// NB! We append a new scope as suffix to the XlaAutoJitScope attribute instead +// of overriding the old value. In this way, we respect the original scopes. // In other words, appending X to Y creates the conjunction of the scopes X // and Y (i.e, X & Y in effect). -void AddOrAppendScope(Node* node, absl::string_view suffix) { +void AddOrAppendXlaAutoJitScope(Node* node, absl::string_view suffix) { string updated_scope; - absl::optional cur_scope = GetXlaScope(node); + absl::optional cur_scope = GetXlaAutoJitScope(node); if (cur_scope == absl::nullopt) { updated_scope = std::string(suffix); } else { updated_scope = absl::StrCat(cur_scope.value(), "&", suffix); } - SetXlaScope(node, updated_scope); + SetXlaAutoJitScope(node, updated_scope); } void ClusterScopingPassImpl::AddScopeToAllPredecessors(Node* start) { @@ -85,7 +85,7 @@ void ClusterScopingPassImpl::AddScopeToAllPredecessors(Node* start) { std::vector starts; starts.push_back(start); - auto enter = [&](Node* n) { AddOrAppendScope(n, unique_suffix); }; + auto enter = [&](Node* n) { AddOrAppendXlaAutoJitScope(n, unique_suffix); }; ReverseDFSFrom(*graph_, starts, enter, /*leave=*/nullptr, /*stable_comparator=*/NodeComparatorName()); } @@ -95,7 +95,7 @@ void ClusterScopingPassImpl::AddScopeToAllSuccessors(Node* start) { std::vector starts; starts.push_back(start); - auto enter = [&](Node* n) { AddOrAppendScope(n, unique_suffix); }; + auto enter = [&](Node* n) { AddOrAppendXlaAutoJitScope(n, unique_suffix); }; auto not_back_edge = [](const Edge& edge) -> bool { return !edge.src()->IsNextIteration(); }; @@ -129,7 +129,7 @@ Status ClusterScopingPassImpl::Run() { // 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. + // latencies add up if Node_X and Node_Y are in the same cluster. // // IteratorGetNext -> Node_X -> Stage // diff --git a/tensorflow/compiler/jit/cluster_scoping_pass.h b/tensorflow/compiler/jit/cluster_scoping_pass.h index 63a2812e744..340837d6cc3 100644 --- a/tensorflow/compiler/jit/cluster_scoping_pass.h +++ b/tensorflow/compiler/jit/cluster_scoping_pass.h @@ -20,14 +20,14 @@ limitations under the License. namespace tensorflow { -// This pass adds xla scopes to graphs to guide the later clustering passes. -// A major reason to do this is to prevent the clustering from losing -// the important parallelism in the Tensorflow graph, which can incur -// great performance degradation. +// This pass adds scopes to nodes in the _XlaAutoJitScope 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 in the XlaScope attributes, which MarkForCompilationPass -// will need to respect for clustering decision. +// scoping information that MarkForCompilationPass will need to respect for +// clustering decision. class ClusterScopingPass : public GraphOptimizationPass { public: Status Run(const GraphOptimizationPassOptions& options) override; diff --git a/tensorflow/compiler/jit/cluster_scoping_pass_test.cc b/tensorflow/compiler/jit/cluster_scoping_pass_test.cc index 9653d1e65bb..7d099b5b561 100644 --- a/tensorflow/compiler/jit/cluster_scoping_pass_test.cc +++ b/tensorflow/compiler/jit/cluster_scoping_pass_test.cc @@ -49,7 +49,7 @@ Status ClusterScoping(std::unique_ptr* graph) { return pass.Run(opt_options); } -absl::flat_hash_map GetXlaScopes(const Graph& graph) { +absl::flat_hash_map GetXlaAutoJitScopes(const Graph& graph) { absl::flat_hash_map scopes; for (Node* node : graph.nodes()) { string scope; @@ -70,8 +70,9 @@ absl::flat_hash_map GetXlaScopes(const Graph& graph) { Node* BuildStageNode(GraphDefBuilder& builder, string name, std::initializer_list dtypes, gtl::ArraySlice values) { - auto opts = - builder.opts().WithName(std::move(name)).WithAttr("dtypes", dtypes); + auto opts = builder.opts() + .WithName(std::move(name)) + .WithAttr("dtypes", std::move(dtypes)); if (opts.HaveError()) { return nullptr; } @@ -119,7 +120,7 @@ TEST(XlaCompilationTest, StagePipelinePreserved) { TF_ASSERT_OK(ClusterScoping(&graph)); - auto scopes = GetXlaScopes(*graph); + auto scopes = GetXlaAutoJitScopes(*graph); EXPECT_NE(scopes["add0"], scopes["add1"]); EXPECT_EQ(scopes["add0"], scopes["relu0"]); EXPECT_EQ(scopes["add1"], scopes["relu1"]); @@ -171,7 +172,7 @@ TEST(XlaCompilationTest, StagePipelinePreservedAndInitialScopesRespected) { TF_ASSERT_OK(ClusterScoping(&graph)); - auto scopes = GetXlaScopes(*graph); + auto scopes = GetXlaAutoJitScopes(*graph); EXPECT_NE(scopes["add0"], scopes["add1"]); EXPECT_NE(scopes["add0"], scopes["relu0"]); EXPECT_NE(scopes["add1"], scopes["relu1"]); diff --git a/tensorflow/compiler/jit/defs.cc b/tensorflow/compiler/jit/defs.cc index e71011d8c5d..81aab02518b 100644 --- a/tensorflow/compiler/jit/defs.cc +++ b/tensorflow/compiler/jit/defs.cc @@ -19,11 +19,11 @@ namespace tensorflow { const char* const kXlaCompileAttr = "_XlaCompile"; -// User-provided through jit_scope. Effective when auto_jit is OFF. +// User-provided through jit_scope. Effective only when auto_jit is OFF. const char* const kXlaScopeAttr = "_XlaScope"; // Automatically inserted by auto_jit to guide clustering results. Effective -// when auto_jit is ON. +// only when auto_jit is ON. const char* const kXlaAutoJitScopeAttr = "_XlaAutoJitScope"; } // namespace tensorflow diff --git a/tensorflow/compiler/jit/mark_for_compilation_pass.cc b/tensorflow/compiler/jit/mark_for_compilation_pass.cc index 8cfa69ab768..8dd1e67f48c 100644 --- a/tensorflow/compiler/jit/mark_for_compilation_pass.cc +++ b/tensorflow/compiler/jit/mark_for_compilation_pass.cc @@ -932,24 +932,22 @@ absl::optional MarkForCompilationPassImpl::GetXlaScope(Node* node) { // The difference between _XlaScope and _XlaAutoJitScope is that _XlaScope is // provided by users through jit_scope APIs, while _XlaAutoJitScope is // automatically generated by the ClusterScopingPass when auto_jit is on. As - // such, we want to respect _kXlaScope when auto_jit is off, while respecting - // _kXlaAutoJitScope when auto_jit is on. + // such, we respect _kXlaScope only when auto_jit is off, while respecting + // _kXlaAutoJitScope only when auto_jit is on. // // We may want to restrict the _XlaScope behavior to require all nodes marked // with _XlaCompile=true to also have a _XlaScope property set (and raise an // error otherwise); but for now we don't do this. if (global_jit_level_ != OptimizerOptions::OFF) { - // If global_jit_level_ is ON, respect kXlaAutoJitScope (and ignore - // kXlaScope). + // If global_jit_level_ is ON, respect only kXlaAutoJitScope. const string& scope = GetNodeAttrString(node->attrs(), kXlaAutoJitScopeAttr); if (!scope.empty()) { return scope; } } else { - // If global_jit_level_ is OFF, respect kXlaScope (and ignore - // kXlaAutoJitScope). + // If global_jit_level_ is OFF, respect only kXlaScope. const string& scope = GetNodeAttrString(node->attrs(), kXlaScopeAttr); if (!scope.empty()) { return scope;