xref: /aosp_15_r20/external/tensorflow/tensorflow/compiler/jit/resource_operation_safety_analysis.h (revision b6fb3261f9314811a0f4371741dbb8839866f948)
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/xla/service/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 writes to happen before resource reads
29 // are problematic.  This analysis returns the set of pairs of resource
30 // operations that cannot be put in the same cluster because XLA cannot respect
31 // the dependencies between them in the TensorFlow program.
32 //
33 // The restrictions are not transitive: it is fine to put A and C in the same
34 // cluster even if the returned set contains (A,B) and (B,C).
35 //
36 // In other words, if these pairs are seen as edges in an undirected graph of
37 // the nodes in `g` then auto-clustering is at least as constrained as the graph
38 // coloring problem on this graph.
39 //
40 //
41 // For instance if we auto-cluster all operations in this TensorFlow graph:
42 //
43 //         AssignVariablepOp0  ->  AssignVariableOp1
44 //                                      |
45 //                                      v
46 //                              ReadVariableOp0  ->  ReadVariableOp1
47 //
48 // we will lose the AssignVariablepOp1 -> ReadVariableOp0. The ReadVariableOp0
49 // -> ReadVariableOp1 and AssignVariableOp0 -> AssignVariableOp1 edges will be
50 // respected by XlaLaunchOp though because all reads happen before all writes
51 // with that limited clustering..
52 //
53 //
54 // NB!  The result computed by this analysis assumes that we don't auto-cluster
55 // back-edges (i.e. the edges from NextIteration to Merge).
56 //
57 // NB!  The result computed by this analysis assumes that we don't auto-cluster
58 // functional control flow nodes containing resource operations.
59 //
60 // If `resource_ops_to_ignore` is set then nodes for which it returns true are
61 // ignored (we pretend these nodes are not resource operations).
62 Status ComputeIncompatibleResourceOperationPairs(
63     const Graph& g, const FunctionLibraryDefinition* flib_def,
64     const std::function<Status(const Node&, bool*)>& resource_ops_to_ignore,
65     std::vector<std::pair<int, int>>* result);
66 }  // namespace tensorflow
67 
68 #endif  // TENSORFLOW_COMPILER_JIT_RESOURCE_OPERATION_SAFETY_ANALYSIS_H_
69