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 namespace detail { 35 36 // Side effect analysis info for a single function. 37 class SideEffectAnalysisInfo { 38 public: 39 SideEffectAnalysisInfo() = default; 40 41 // Constructs analysis info by analyzing the given function. SideEffectAnalysisInfo(FuncOp func_op,const TF::ResourceAliasAnalysis::Info & alias_analysis)42 SideEffectAnalysisInfo( 43 FuncOp func_op, const TF::ResourceAliasAnalysis::Info& alias_analysis) { 44 AnalyzeFunction(func_op, alias_analysis); 45 } 46 47 // Constructs analysis info by analyzing the given region. SideEffectAnalysisInfo(Region * region,const TF::ResourceAliasAnalysis::Info & alias_analysis)48 SideEffectAnalysisInfo( 49 Region* region, const TF::ResourceAliasAnalysis::Info& alias_analysis) { 50 AnalyzeRegion(region, alias_analysis); 51 } 52 53 SideEffectAnalysisInfo(SideEffectAnalysisInfo&&) = default; 54 55 // Returns a vector of ops that are direct control predecessors of `op`, 56 // sorted in program order. If `filter` is provided, only predecessors that 57 // pass the filter (returning true) will be included. 58 llvm::SmallVector<Operation*, 4> DirectControlPredecessors( 59 Operation* op, 60 llvm::function_ref<bool(Operation*)> filter = nullptr) const; 61 62 // Returns a vector of ops that are direct control successors of `op`, 63 // sorted in program order. If `filter` is provided, only successors that 64 // pass the filter (returning true) will be included. 65 llvm::SmallVector<Operation*, 4> DirectControlSuccessors( 66 Operation* op, 67 llvm::function_ref<bool(Operation*)> filter = nullptr) const; 68 69 private: 70 // Runs the analysis on `func_op` and populates sorted_control_predecessors_ 71 // and sorted_control_successors_. 72 void AnalyzeFunction(FuncOp func_op, 73 const TF::ResourceAliasAnalysis::Info& alias_analysis); 74 75 // Runs the analysis on `region` and populates control_predecessors_. 76 void AnalyzeRegion(Region* region, 77 const TF::ResourceAliasAnalysis::Info& alias_analysis); 78 79 // Updates control_predecessors_ for `op` that is being visited, on the given 80 // `resource_id`. 81 void AddPredecessorsForAccess(int64_t resource_id, Operation* op, 82 bool read_only); 83 84 // Adds op's access to per_resource_access_info_. 85 void TrackAccess(int64_t resource_id, Operation* op, bool read_only); 86 87 // Maps from an op to its control predecessors. 88 llvm::SmallDenseMap<Operation*, llvm::SmallPtrSet<Operation*, 4>, 8> 89 control_predecessors_; 90 // Maps from an op to its control predecessors sorted in program order. 91 llvm::SmallDenseMap<Operation*, llvm::SmallVector<Operation*, 4>, 8> 92 sorted_control_predecessors_; 93 // Maps from an op to its control successors sorted in program order. 94 llvm::SmallDenseMap<Operation*, llvm::SmallVector<Operation*, 4>, 8> 95 sorted_control_successors_; 96 97 // Internal per-resource data structure when we build the dependencies. 98 struct PerResourceAccessInfo { 99 // Last op that writes the resource before the current op being analyzed. 100 Operation* last_write = nullptr; 101 // Read ops since last_write before the current op being analyzed. 102 llvm::SmallVector<Operation*, 8> reads_since_last_write; 103 // Whether previous accesses of this resource already tracked last unknown 104 // read for the current access being analyzed. 105 bool tracked_last_unknown_read = false; 106 // Whether previous accesses of this resource already tracked last unknown 107 // write for a the current read being analyzed. 108 bool tracked_last_unknown_write_for_read = false; 109 // Whether previous accesses of this resource already tracked last unknown 110 // write for a the current write being analyzed. 111 bool tracked_last_unknown_write_for_write = false; 112 }; 113 114 llvm::SmallDenseMap<int64_t, PerResourceAccessInfo, 8> 115 per_resource_access_info_; 116 }; 117 118 } // namespace detail 119 120 // An analysis that runs on a function and infers the control predecessors and 121 // successors for each op, based on side-effects on known and unknown resources. 122 // Side-effecting ops on unknown resources are conservatively treated as 123 // interfering with all known resource op accesses. It distinguishes accesses 124 // based on whether they are read-only, and read-only ops do not interfere with 125 // each other. 126 // 127 // If there are nested regions, each region is handled separately, and control 128 // dependencies are only tracked for ops under the same parent op. 129 class SideEffectAnalysis : public detail::PerFunctionAggregateAnalysis< 130 detail::SideEffectAnalysisInfo> { 131 public: 132 // Constructs analysis by analyzing the given module operation. 133 explicit SideEffectAnalysis(ModuleOp module); 134 }; 135 136 } // namespace TF 137 } // namespace mlir 138 139 #endif // TENSORFLOW_COMPILER_MLIR_TENSORFLOW_ANALYSIS_SIDE_EFFECT_ANALYSIS_H_ 140