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 #ifndef TENSORFLOW_COMPILER_JIT_RESOURCE_OPERATION_SAFETY_ANALYSIS_H_ 17 #define TENSORFLOW_COMPILER_JIT_RESOURCE_OPERATION_SAFETY_ANALYSIS_H_ 18 19 #include "tensorflow/compiler/jit/graphcycles/graphcycles.h" 20 #include "tensorflow/core/framework/function.h" 21 #include "tensorflow/core/graph/graph.h" 22 23 namespace tensorflow { 24 // An XLA cluster hoists all resource reads to be beginning of the cluster 25 // execution and all the resource writes to the end. This means it cannot 26 // enforce arbitrary ordering dependencies (via control or data edges) between 27 // resource operations. Since all resource reads happen before all resource 28 // writes, edges constraining resource reads to happen before resource writes 29 // are fine, but all other kinds of edges are problematic. This analysis 30 // returns the set of pairs of resource operations that cannot be put in the 31 // same cluster because XLA cannot respect the dependencies between them in the 32 // TensorFlow program. 33 // 34 // The restrictions are not transitive: it is fine to put A and C in the same 35 // cluster even if the returned set contains (A,B) and (B,C). 36 // 37 // In other words, if these pairs are seen as edges in an undirected graph of 38 // the nodes in `g` then auto-clustering is at least as constrained as the graph 39 // coloring problem on this graph. 40 // 41 // 42 // For instance if we auto-cluster all operations in this TensorFlow graph: 43 // 44 // ReadVariablepOp0 -> ReadVariableOp1 45 // | 46 // v 47 // AssignVariableOp0 -> AssignVariableOp1 48 // 49 // we will lose the ReadVariablepOp0 -> ReadVariableOp1 and the 50 // AssignVariableOp0 -> AssignVariableOp1 dependencies. I.e. it is possible for 51 // XlaLaunchOp to issue ReadVariableOp1 before ReadVariablepOp0 since it reads 52 // all the resource variables when the cluster starts executing without any 53 // particular ordering between them; same holds for the AssignVariableOp0 -> 54 // AssignVariableOp1 edge. The ReadVariableOp1 -> AssignVariableOp0 edge will 55 // be respected by XlaLaunchOp though because all reads happen before all 56 // writes. 57 // 58 // 59 // NB! The result computed by this analysis assumes that we don't auto-cluster 60 // back-edges (i.e. the edges from NextIteration to Merge). 61 // 62 // NB! The result computed by this analysis assumes that we don't auto-cluster 63 // functional control flow nodes containing resource operations. 64 // 65 // If `resource_ops_to_ignore` is set then nodes for which it returns true are 66 // ignored (we pretend these nodes are not resource operations). 67 Status ComputeIncompatibleResourceOperationPairs( 68 const Graph& g, const FunctionLibraryDefinition* flib_def, 69 const std::function<Status(const Node&, bool*)>& resource_ops_to_ignore, 70 std::vector<std::pair<int, int>>* result); 71 } // namespace tensorflow 72 73 #endif // TENSORFLOW_COMPILER_JIT_RESOURCE_OPERATION_SAFETY_ANALYSIS_H_ 74