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_SIDE_EFFECT_ANALYSIS_H_ 17 #define TENSORFLOW_COMPILER_MLIR_TENSORFLOW_ANALYSIS_SIDE_EFFECT_ANALYSIS_H_ 18 19 #include <cstddef> 20 #include <cstdint> 21 #include <memory> 22 23 #include "llvm/ADT/SetVector.h" 24 #include "llvm/ADT/SmallSet.h" 25 #include "llvm/ADT/SmallVector.h" 26 #include "llvm/ADT/StringMap.h" 27 #include "mlir/IR/Operation.h" // from @llvm-project 28 #include "mlir/IR/Region.h" // from @llvm-project 29 #include "mlir/Support/LogicalResult.h" // from @llvm-project 30 #include "tensorflow/compiler/mlir/tensorflow/analysis/resource_alias_analysis.h" 31 32 namespace mlir { 33 namespace TF { 34 using ResourceId = int64_t; 35 36 namespace detail { 37 38 class OpSideEffectCollector; 39 40 // Side effect analysis info for a single function. 41 // 42 // This class provides an interface for querying control predecessors and 43 // successors for ops of the given function. This information is computed from 44 // side effects, using resource alias analysis where possible. 45 // Remarks: 46 // - Control dependencies model execution order constraints for side-effecting 47 // ops. For example, two ops writing to the same resource cannot switch their 48 // order and cannot be executed in parallel. 49 // - A control dependency (A,B) means that op A has to be executed before op B. 50 // A is a control predecessor of B, and B is a control successor of A. 51 // - The control dependencies provided by side effect analysis are guaranteed to 52 // be sufficient for correct execution but they are not guaranteed to be 53 // minimal (that means, some control dependencies might not be required for 54 // correct execution). 55 class SideEffectAnalysisInfo { 56 public: 57 SideEffectAnalysisInfo() = default; 58 59 // Constructs analysis info by analyzing the given function. SideEffectAnalysisInfo(func::FuncOp func_op,const OpSideEffectCollector & op_side_effect_collector,const TF::ResourceAliasAnalysis::Info & alias_analysis)60 SideEffectAnalysisInfo(func::FuncOp func_op, 61 const OpSideEffectCollector& op_side_effect_collector, 62 const TF::ResourceAliasAnalysis::Info& alias_analysis) 63 : op_side_effect_collector_(op_side_effect_collector), 64 alias_analysis_(alias_analysis) { 65 AnalyzeFunction(func_op); 66 } 67 68 // Constructs analysis info by analyzing the given region. SideEffectAnalysisInfo(Region * region,const OpSideEffectCollector & op_side_effect_collector,const TF::ResourceAliasAnalysis::Info & alias_analysis)69 SideEffectAnalysisInfo(Region* region, 70 const OpSideEffectCollector& op_side_effect_collector, 71 const TF::ResourceAliasAnalysis::Info& alias_analysis) 72 : op_side_effect_collector_(op_side_effect_collector), 73 alias_analysis_(alias_analysis) { 74 AnalyzeRegion(region); 75 } 76 77 SideEffectAnalysisInfo(SideEffectAnalysisInfo&&) = default; 78 79 // Returns a vector of ops that are direct control predecessors of `op`, 80 // sorted in program order. If `filter` is provided, only predecessors that 81 // pass the filter (returning true) will be included. 82 llvm::SmallVector<Operation*, 4> DirectControlPredecessors( 83 Operation* op, 84 llvm::function_ref<bool(Operation*)> filter = nullptr) const; 85 86 // Returns a vector of ops that are direct control successors of `op`, 87 // sorted in program order. If `filter` is provided, only successors that 88 // pass the filter (returning true) will be included. 89 llvm::SmallVector<Operation*, 4> DirectControlSuccessors( 90 Operation* op, 91 llvm::function_ref<bool(Operation*)> filter = nullptr) const; 92 93 // Returns a vector of ops that are control sinks (i.e. side-effecting ops 94 // with no control successors). ControlSinks()95 llvm::ArrayRef<Operation*> ControlSinks() const { 96 return sorted_control_sinks_; 97 } 98 99 // Returns a vector with IDs of all resources that might be accessed by `op`. 100 // This includes both op-based and value-based resources. The bool indicates 101 // whether a resource is accessed read-only. 102 const llvm::SmallVector<std::pair<ResourceId, bool>>& GetResourceIds( 103 Operation* op) const; 104 105 // Returns true iff given resource is allocated by op with 106 // `UniqueResourceAllocation` trait. This can be utilized for while-loop 107 // parallelization. IsUniqueResourceAllocationId(ResourceId resource_id)108 bool IsUniqueResourceAllocationId(ResourceId resource_id) const { 109 return alias_analysis_.IsUniqueResourceAllocationId(resource_id); 110 } 111 112 private: 113 // Runs the analysis and populates `sorted_control_predecessors_` and 114 // `sorted_control_successors_` for `func_op`. Clears `control_predecessors_`. 115 void AnalyzeFunction(func::FuncOp func_op); 116 117 // Runs the analysis and populates `control_predecessors_` for `region`. 118 void AnalyzeRegion(Region* region); 119 120 // Runs the analysis and populates `control_predecessors_` for `op`. 121 void AnalyzeOp(Operation* op); 122 123 // Updates `control_predecessors_` for given `resource_id` and `op`. 124 void AddPredecessorsForAccess(ResourceId resource_id, Operation* op, 125 bool read_only); 126 127 // Updates resource access for given `resource_id` and `op` in 128 // `per_resource_access_info_` and `op_to_resource_ids_`. 129 void UpdateAccess(ResourceId resource_id, Operation* op, bool read_only); 130 131 // Returns true iff the last unknown resource access is already indirectly 132 // tracked by a previous `resource` access. `read_only` specifies the type of 133 // access considered. 134 bool IsUnknownAccessIndirectlyTrackedByResource(ResourceId resource, 135 bool read_only); 136 137 // Returns a set of resource IDs that are conflicting with `resource_id`, i.e. 138 // there are potentially dependencies between the corresponding resources. 139 llvm::SmallSet<ResourceId, 8> GetConflictingIds(ResourceId resource_id, 140 bool is_fetch_op) const; 141 142 // Maps from an op to its control predecessors. 143 llvm::SmallDenseMap<Operation*, llvm::SmallPtrSet<Operation*, 4>, 8> 144 control_predecessors_; 145 // Maps from an op to its control predecessors sorted in program order. 146 llvm::SmallDenseMap<Operation*, llvm::SmallVector<Operation*, 4>, 8> 147 sorted_control_predecessors_; 148 // Maps from an op to its control successors sorted in program order. 149 llvm::SmallDenseMap<Operation*, llvm::SmallVector<Operation*, 4>, 8> 150 sorted_control_successors_; 151 // Side-effecting ops with no control successors in this function. 152 llvm::SmallVector<Operation*, 4> sorted_control_sinks_; 153 154 // Maps from an op to its resource IDs along with a bool indicating if the 155 // resource is accessed `read-only`. 156 llvm::SmallDenseMap<Operation*, 157 llvm::SmallVector<std::pair<ResourceId, bool>>> 158 op_to_resource_ids_; 159 llvm::SmallVector<std::pair<ResourceId, bool>> empty_resource_ids_; 160 161 // Internal per-resource data structure for building the dependencies. 162 struct PerResourceAccessInfo { 163 // Last op that writes to resource before the current op is being analyzed. 164 Operation* last_write = nullptr; 165 // Read ops since `last_write` before the current op is being analyzed. 166 llvm::SmallVector<Operation*, 8> reads_since_last_write; 167 // Whether a previous access of this resource already tracks the last 168 // unknown read(s). 169 bool are_last_unknown_reads_tracked = false; 170 // Whether a previous write access of this resource already tracks the last 171 // unknown write. 172 bool is_last_unknown_write_tracked_by_write = false; 173 // Whether a previous read or write access of this resource already tracks 174 // the last unknown write. 175 bool is_last_unknown_write_tracked = false; 176 }; 177 178 // Resource access info per resource ID. 179 llvm::SmallDenseMap<ResourceId, PerResourceAccessInfo, 8> 180 per_resource_access_info_; 181 182 const OpSideEffectCollector& op_side_effect_collector_; 183 const TF::ResourceAliasAnalysis::Info& alias_analysis_; 184 }; 185 186 } // namespace detail 187 188 // An analysis that runs on a function and infers the control predecessors and 189 // successors for each op, based on side effects on known and unknown resources. 190 // Side-effecting ops on unknown resources are conservatively treated as 191 // interfering with all known resource op accesses. It distinguishes accesses 192 // based on whether they are read-only, and read-only ops do not interfere with 193 // each other. 194 // 195 // If there are nested regions, each region is handled separately, and control 196 // dependencies are only tracked for ops under the same parent op. 197 class SideEffectAnalysis : public detail::PerFunctionAggregateAnalysis< 198 detail::SideEffectAnalysisInfo> { 199 public: 200 // Constructs analysis by analyzing the given module operation. 201 explicit SideEffectAnalysis(ModuleOp module); 202 203 private: 204 ResourceAliasAnalysis alias_analysis_; 205 }; 206 207 } // namespace TF 208 } // namespace mlir 209 210 #endif // TENSORFLOW_COMPILER_MLIR_TENSORFLOW_ANALYSIS_SIDE_EFFECT_ANALYSIS_H_ 211