1 /* Copyright 2019 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_MLIR_TENSORFLOW_ANALYSIS_RESOURCE_ALIAS_ANALYSIS_H_
17 #define TENSORFLOW_COMPILER_MLIR_TENSORFLOW_ANALYSIS_RESOURCE_ALIAS_ANALYSIS_H_
18 
19 #include <cstddef>
20 #include <cstdint>
21 #include <memory>
22 
23 #include "llvm/ADT/ArrayRef.h"
24 #include "llvm/ADT/STLExtras.h"
25 #include "llvm/ADT/SetVector.h"
26 #include "llvm/ADT/SmallSet.h"
27 #include "llvm/ADT/SmallVector.h"
28 #include "llvm/ADT/StringMap.h"
29 #include "mlir/IR/Operation.h"  // from @llvm-project
30 #include "mlir/IR/Region.h"  // from @llvm-project
31 #include "mlir/IR/TypeUtilities.h"  // from @llvm-project
32 #include "tensorflow/compiler/mlir/tensorflow/analysis/per_function_aggregate_analysis.h"
33 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
34 
35 namespace mlir {
36 namespace TF {
37 namespace detail {
38 class BacktrackAnalysis;
39 class BacktrackAnalysisInfo;
40 
41 // Resource alias analysis information for a single function.
42 class ResourceAliasAnalysisInfo {
43  public:
44   // Constructs analysis info by analyzing the given function.
45   ResourceAliasAnalysisInfo(func::FuncOp func,
46                             const BacktrackAnalysis& backtrack_analysis,
47                             SymbolTableCollection& symbol_table_collection);
48 
49   ResourceAliasAnalysisInfo(ResourceAliasAnalysisInfo&&) = default;
50 
51   // Returns if the analysis fails to resolve a resource-type value.
52   bool IsUnknownResource(Value resource) const;
53 
54   // Returns the set of unique IDs which `resource` could alias. Requires that
55   // IsUnknownResource(resource) == false.
56   const llvm::SmallSet<int64_t, 8>& GetResourceUniqueIds(Value resource) const;
57 
58   // Returns the set of values that are potentially aliases of `value`. Requires
59   // `IsUnknownResource(resource) == false`.
60   llvm::SmallSetVector<Value, 8> GetResourceAliases(Value resource) const;
61 
62   // Returns true iff given resource is allocated by op with
63   // `UniqueResourceAllocation` trait. This can be utilized for while-loop
64   // parallelization.
IsUniqueResourceAllocationId(int64_t resource_id)65   bool IsUniqueResourceAllocationId(int64_t resource_id) const {
66     return unique_resource_allocation_ids_.contains(resource_id);
67   }
68 
69  private:
70   // Maps resource value to unique ID and vice-versa. Returns true if the
71   // mapping has changed.
AddValueUniqueIDMapping(Value value,int64_t id)72   bool AddValueUniqueIDMapping(Value value, int64_t id) {
73     resource_value_to_ids_[value].insert(id);
74     return id_to_resource_values_[id].insert(value);
75   }
76 
77   // Returns the set unique Values which map to `id`.
78   const llvm::SmallSetVector<Value, 8>& GetUniqueIdResources(int64_t id) const;
79 
80   // Propagates the resource IDs from an input operand to a result. Returns
81   // true of the mapping has changed.
82   bool PropagateInputToOutput(const Value& operand, const OpResult& result);
83 
84   // Analyzes while loops to compute resource IDs for the loop results.
85   // `body_info` is the backtrack analysis info for the loop body.
86   void AnalyzeWhileLoop(Operation* while_op,
87                         const BacktrackAnalysisInfo& body_info);
88 
89   // Analyzes tf.Case/tf.If ops to compute resource IDs.
90   template <class CaseOrIfOp>
91   void AnalyzeFunctionalCaseOrIfOp(CaseOrIfOp case_or_if_op,
92                                    llvm::ArrayRef<func::FuncOp> functions,
93                                    const BacktrackAnalysis& backtrack_analysis);
94 
95   // Analyzes tf.CaseRegion/tf.IfRegion ops to compute resource IDs.
96   void AnalyzeRegionCaseOrIfOp(Operation* case_or_if_op,
97                                const BacktrackAnalysis& backtrack_analysis);
98 
99   // Maps each resource-type value to a set of unique IDs that it could alias.
100   llvm::SmallDenseMap<Value, llvm::SmallSet<int64_t, 8>, 8>
101       resource_value_to_ids_;
102 
103   // Maps each unique ID to a set of resource-type values that could alias to
104   // it. This is inverse of `resource_value_to_ids_` map.
105   llvm::SmallDenseMap<int64_t, llvm::SmallSetVector<Value, 8>, 8>
106       id_to_resource_values_;
107 
108   // Maps MLIR type IDs for resource types to internal resource type IDs.
109   llvm::SmallDenseMap<TypeID, int64_t> type_id_to_internal_type_id_;
110 
111   // Contains IDs of all resources that are allocated by ops with
112   // `UniqueResourceAllocation` trait.
113   llvm::SmallDenseSet<int64_t, 32> unique_resource_allocation_ids_;
114 
115  public:
116   // Resource IDs have the following semantics:
117   // a) -1 represents an unknown resource (both instance and type unknown)
118   // b) IDs in range [0,kMaxResourceTypeId] represent resource type IDs; we use
119   //    such IDs when we know the resource type but not the instance
120   // c) IDs > kMaxResourceTypeId represent resource instance IDs (i.e., we know
121   //    the specific resource instance)
122   //
123   // Note: In general, there can be different ops allocating a resource of the
124   // same type, for one we might assign a resource type ID and for the other
125   // a resource instance ID. That means, they will be treated as non-aliasing.
126   // This is correct for all current cases. A problematic case could be if we
127   // had two ops A and B, A has the `ResourceHandleAllocatorInterface` and B has
128   // not, and both ops might return a handle to the same resource (depending on
129   // attributes). In this case, the return value of A would get a different ID
130   // than the return value of B although both could point to the same resource.
131   // It seems highly unlikely to encounter such a case but, to be safe, this
132   // should be revisited for new resource-allocators that might potentially
133   // break our currently guaranteed correctness.
134   // For context, we are very conservative here compared to
135   // `auto_control_deps.py` where it is assumed that allocated resource values
136   // NEVER alias. We should align our assumptions in the future.
137   static constexpr int64_t kUnknownResourceId = -1;
138   static constexpr int64_t kInvalidResourceId = -2;
139   static constexpr int64_t kMaxResourceTypeId = 9999;
140 };
141 
142 }  // namespace detail
143 
144 // An analysis that runs on a module and maps each resource-type value to a
145 // set of unique IDs representing the possible resources it could alias.
146 //
147 // Note that this is not an inter-procedural or inter-regional analysis, i.e.,
148 // each function and region are handled separately and cross-function or cross-
149 // region aliasing cannot be checked by this analysis.
150 class ResourceAliasAnalysis : public detail::PerFunctionAggregateAnalysis<
151                                   detail::ResourceAliasAnalysisInfo> {
152  public:
153   // Constructs analysis by analyzing the given module operation.
154   explicit ResourceAliasAnalysis(ModuleOp module);
155 };
156 
157 }  // namespace TF
158 }  // namespace mlir
159 
160 #endif  // TENSORFLOW_COMPILER_MLIR_TENSORFLOW_ANALYSIS_RESOURCE_ALIAS_ANALYSIS_H_
161