1 /* Copyright 2018 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 // Contains utilities for clustering compilable graph nodes via XLA. 17 18 #ifndef TENSORFLOW_COMPILER_JIT_XLA_CLUSTER_UTIL_H_ 19 #define TENSORFLOW_COMPILER_JIT_XLA_CLUSTER_UTIL_H_ 20 21 #include "absl/container/flat_hash_map.h" 22 #include "absl/container/flat_hash_set.h" 23 #include "absl/types/optional.h" 24 #include "tensorflow/compiler/jit/xla_activity.pb.h" 25 #include "tensorflow/compiler/xla/service/graphcycles/graphcycles.h" 26 #include "tensorflow/compiler/xla/statusor.h" 27 #include "tensorflow/core/common_runtime/optimization_registry.h" 28 #include "tensorflow/core/graph/algorithm.h" 29 #include "tensorflow/stream_executor/lib/statusor.h" 30 31 namespace tensorflow { 32 33 // The attribute that marks nodes to be grouped into functions by the 34 // encapsulate subgraphs pass. 35 extern const char* const kXlaClusterAttr; 36 37 // The attribute that marks nodes in a cluster to be placed outside the xla 38 // compilation by the encapsulate subgraphs pass. 39 extern const char* const kXlaOutsideCompilationAttr; 40 41 // The attribute that marks certain inputs to a Node as required to be a 42 // constant at compile time. If this attribute is present then the 43 // CompileTimeConstantInput information in the corresponding XlaOpKernel is 44 // ignored. 45 // 46 // The value for this attribute, if present, has to be a list of strings naming 47 // the inputs to the node that must be constant. 48 extern const char* const kXlaCompileTimeConstantInputsAttr; 49 50 using OrderedNodeSet = std::set<Node*, NodeComparatorID>; 51 52 // Returns true if `node` has a ref tensor input that it forwards to its output. 53 bool HasForwardedRefInput(const Node& node); 54 55 // Creates a graph representation to enable cycle detection when clustering. 56 // This representation handles loops in graph by disconnecting each loop from 57 // the enclosing graph. 58 // 59 // Returns true for success and false for valid graphs that we can't handle yet 60 // (b/127521408). 61 xla::StatusOr<bool> CreateCycleDetectionGraph(const Graph* graph, 62 GraphCycles* cycles); 63 64 // Returns the XLA cluster in which `node` is placed if it is in an XLA cluster, 65 // otherwise returns nullopt. 66 absl::optional<absl::string_view> GetXlaClusterForNode(const Node& node); 67 68 // Removes `node_def` its XLA cluster (by clearing its _XlaCluster attribute). 69 void RemoveFromXlaCluster(NodeDef* node_def); 70 71 // Removes `node` its XLA cluster (by clearing its _XlaCluster attribute). 72 void RemoveFromXlaCluster(Node* node); 73 74 // Returns true if `node` has a DT_RESOURCE typed input or output. 75 bool HasResourceInputOrOutput(const Node& node); 76 77 // Determines the global jit level based on GraphOptimizationPassOptions, 78 // --tf_xla_auto_jit and whether the graph is a single GPU graph. 79 OptimizerOptions::GlobalJitLevel GetGlobalJitLevelForGraph( 80 const GraphOptimizationPassOptions& options); 81 82 // Returns true if `g` is a single-GPU graph. A single-GPU graph uses exactly 83 // one GPU (and any number of CPUs). 84 bool IsSingleGpuGraph(const Graph& g); 85 86 // Returns true if it is possible (but not guaranteed) that `n` calls a 87 // function. 88 bool MayCallFunction(const Node& n, const FunctionLibraryDefinition* flib_def); 89 90 // Returns true if `node` an operator that consumes only the shape of its input, 91 // not the data itself. 92 bool IsShapeConsumerOp(const Node& node); 93 94 // Computes a clustering summary for `graph`. See documentation on 95 // `XlaAutoClusteringSummary` for details. 96 XlaAutoClusteringSummary GetXlaAutoClusteringSummary(const Graph& graph); 97 98 // Returns the set of nodes that have a path to or from nodes that may have ref 99 // variables as input or output. 100 // 101 // We assume each node has a trivial path to itself so the returned set includes 102 // all of the nodes that have ref variables as input or output. 103 xla::StatusOr<absl::flat_hash_set<Node*>> GetNodesRelatedToRefVariables( 104 const Graph& graph, FunctionLibraryRuntime* lib_runtime); 105 106 } // namespace tensorflow 107 108 #endif // TENSORFLOW_COMPILER_JIT_XLA_CLUSTER_UTIL_H_ 109