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(FuncOp func,
46 const BacktrackAnalysis& backtrack_analysis);
47
48 ResourceAliasAnalysisInfo(ResourceAliasAnalysisInfo&&) = default;
49
50 // Returns if the analysis fails to resolve a resource-type value.
51 bool IsUnknownResource(Value resource) const;
52
53 // Returns the set unique IDs which `resource` could alias. Requires that
54 // IsUnknownResource(resource) == false.
55 const llvm::SmallSet<int64_t, 8>& GetResourceUniqueIds(Value resource) const;
56
57 // Returns the set of values that are potentially aliases of `value`. Requires
58 // that IsUnknownResource(resource) == false.
59 llvm::SmallSetVector<Value, 8> GetResourceAliases(Value resource) const;
60
61 private:
62 // Maps resource value to unique ID and vice-versa. Returns true of the
63 // mapping has changed.
AddValueUniqueIDMapping(Value value,int64_t id)64 bool AddValueUniqueIDMapping(Value value, int64_t id) {
65 resource_value_to_ids_[value].insert(id);
66 return id_to_resource_values_[id].insert(value);
67 }
68
69 // Returns the set unique Values which map to `id`.
70 const llvm::SmallSetVector<Value, 8>& GetUniqueIdResources(int64_t id) const;
71
72 // Propagates the resource ID's from an input operand to a result. Returns
73 // true of the mapping has changed.
74 bool PropagateInputToOutput(const Value& operand, const OpResult& result);
75
76 // Analyzes while loops to compute resourceID's for the loop results.
77 // `body_info` is the backtrack analysis info for the loop body.
78 void AnalyzeWhileLoop(Operation* while_op,
79 const BacktrackAnalysisInfo& body_info);
80
81 // Analyzes tf.Case/tf.If ops to compute resourceID's.
82 template <class CaseOrIfOp>
83 void AnalyzeFunctionalCaseOrIfOp(CaseOrIfOp case_or_if_op,
84 llvm::ArrayRef<FuncOp> functions,
85 const BacktrackAnalysis& backtrack_analysis);
86
87 // Analyzes tf.CaseRegion/tf.IfRegion ops to compute resourceID's.
88 void AnalyzeRegionCaseOrIfOp(Operation* case_or_if_op,
89 const BacktrackAnalysis& backtrack_analysis);
90
91 // Maps each resource-type value to a set of unique IDs that it could alias.
92 llvm::SmallDenseMap<Value, llvm::SmallSet<int64_t, 8>, 8>
93 resource_value_to_ids_;
94
95 // Maps each unique ID to a set of resource-type values that could alias to
96 // it. This is inverse of `resource_value_to_ids_` map.
97 llvm::SmallDenseMap<int64_t, llvm::SmallSetVector<Value, 8>, 8>
98 id_to_resource_values_;
99
100 public:
101 static constexpr int64_t kUnknownResourceId = -1;
102 };
103
104 } // namespace detail
105
106 // An analysis that runs on a module and maps each resource-type value to a
107 // set of unique IDs representing the possible resources it could alias.
108 //
109 // Note that this is not an inter-procedural or inter-regional analysis, i.e.,
110 // each function and region are handled separately and cross-function or cross-
111 // region aliasing cannot be checked by this analysis.
112 class ResourceAliasAnalysis : public detail::PerFunctionAggregateAnalysis<
113 detail::ResourceAliasAnalysisInfo> {
114 public:
115 // Constructs analysis by analyzing the given module operation.
116 explicit ResourceAliasAnalysis(ModuleOp module);
117 };
118
119 // Returns a range with just resource type values from the input range
120 // preserved.
121 template <typename RangeT>
filter_resources(RangeT && range)122 auto filter_resources(RangeT&& range) {
123 return llvm::make_filter_range(std::forward<RangeT>(range), [](Value val) {
124 return getElementTypeOrSelf(val.getType()).isa<TF::ResourceType>();
125 });
126 }
127
128 } // namespace TF
129 } // namespace mlir
130
131 #endif // TENSORFLOW_COMPILER_MLIR_TENSORFLOW_ANALYSIS_RESOURCE_ALIAS_ANALYSIS_H_
132