1 /* Copyright 2017 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_XLA_SERVICE_HLO_EXECUTION_PROFILE_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_EXECUTION_PROFILE_H_ 18 19 #include <vector> 20 21 #include "absl/container/flat_hash_map.h" 22 #include "tensorflow/compiler/xla/map_util.h" 23 #include "tensorflow/compiler/xla/service/hlo_cost_analysis.h" 24 #include "tensorflow/compiler/xla/service/hlo_execution_profile_data.pb.h" 25 #include "tensorflow/compiler/xla/service/hlo_profile_printer.h" 26 #include "tensorflow/compiler/xla/types.h" 27 28 namespace xla { 29 30 class HloInstruction; 31 32 // Maps all HloInstructions and HloComputations in an HloModule to integers. 33 // These integers form the contiguous range [0, total_count()). 34 class HloProfileIndexMap { 35 public: 36 // Scans `module` to populate this instance of HloProfileIndexMap. HloProfileIndexMap(const HloModule & module)37 explicit HloProfileIndexMap(const HloModule& module) 38 : HloProfileIndexMap(module, {}) {} 39 explicit HloProfileIndexMap(const HloModule& module, 40 absl::Span<const std::string> extra_metrics); 41 42 HloProfileIndexMap(const HloProfileIndexMap&) = default; 43 HloProfileIndexMap(HloProfileIndexMap&&) = default; 44 45 HloProfileIndexMap& operator=(const HloProfileIndexMap&) = default; 46 HloProfileIndexMap& operator=(HloProfileIndexMap&&) = default; 47 GetProfileIndexFor(const HloInstruction & instruction)48 size_t GetProfileIndexFor(const HloInstruction& instruction) const { 49 return FindOrDie(instruction_to_profile_idx(), &instruction); 50 } 51 GetProfileIndexFor(const HloComputation & computation)52 size_t GetProfileIndexFor(const HloComputation& computation) const { 53 return FindOrDie(computation_to_profile_idx(), &computation); 54 } 55 GetProfileIndexFor(const std::string & key)56 size_t GetProfileIndexFor(const std::string& key) const { 57 return xla::FindOrDie(extra_metric_to_profile_idx(), key); 58 } 59 instruction_count()60 size_t instruction_count() const { 61 return instruction_to_profile_idx().size(); 62 } 63 computation_count()64 size_t computation_count() const { 65 return computation_to_profile_idx().size(); 66 } 67 extra_metrics_count()68 size_t extra_metrics_count() const { 69 return extra_metric_to_profile_idx().size(); 70 } 71 total_count()72 size_t total_count() const { 73 return instruction_count() + computation_count() + extra_metrics_count(); 74 } 75 76 const absl::flat_hash_map<const HloInstruction*, int64_t>& instruction_to_profile_idx()77 instruction_to_profile_idx() const { 78 return instruction_to_profile_idx_; 79 } 80 81 const absl::flat_hash_map<const HloComputation*, int64_t>& computation_to_profile_idx()82 computation_to_profile_idx() const { 83 return computation_to_profile_idx_; 84 } 85 extra_metric_to_profile_idx()86 const absl::flat_hash_map<std::string, int64_t>& extra_metric_to_profile_idx() 87 const { 88 return extra_metric_to_profile_idx_; 89 } 90 91 private: 92 absl::flat_hash_map<const HloInstruction*, int64_t> 93 instruction_to_profile_idx_; 94 absl::flat_hash_map<const HloComputation*, int64_t> 95 computation_to_profile_idx_; 96 absl::flat_hash_map<std::string, int64_t> extra_metric_to_profile_idx_; 97 }; 98 99 // Create an instance of `HloProfilePrinterData`. 100 std::unique_ptr<HloProfilePrinterData> CreateHloProfilePrinterData( 101 const HloProfileIndexMap& hlo_profile_index_map, 102 const HloCostAnalysis& cost_analysis, 103 const std::string& entry_computation_name); 104 105 // Describes how much time each HLO operation took. 106 // 107 // Each HloComputation takes a certain number of cycles. This class helps break 108 // down how much time each HLO took. 109 class HloExecutionProfile { 110 public: 111 HloExecutionProfile(const HloProfilePrinterData* hlo_profile_printer_data, 112 const HloProfileIndexMap* hlo_profile_index_map); 113 114 // Record how many cycles this HLO took to execute. 115 void SetCyclesTakenBy(const HloInstruction* hlo, uint64_t cycles_taken); 116 117 // Record how many cycles this HLO took to execute. 118 void SetCyclesTakenBy(size_t index, uint64_t cycles_taken); 119 120 // Returns how many cycles this HLO took to execute. Profiling information 121 // may not be available for some instructions in which case zero is returned. 122 uint64_t GetCyclesTakenBy(const HloInstruction& hlo) const; 123 124 // Returns how many cycles this HLO took to execute. Profiling information 125 // may not be available for some instructions in which case zero is returned. 126 uint64_t GetCyclesTakenBy(size_t index) const; 127 128 // Return the number of cycles this computation took to execute. total_cycles_executed(const HloComputation & computation)129 uint64_t total_cycles_executed(const HloComputation& computation) const { 130 return profile_counters_[hlo_profile_index_map_.GetProfileIndexFor( 131 computation)]; 132 } 133 134 // Record how many cycles a computation took to execute. set_total_cycles_executed(const HloComputation & computation,uint64_t total_cycles_executed)135 void set_total_cycles_executed(const HloComputation& computation, 136 uint64_t total_cycles_executed) { 137 profile_counters_[hlo_profile_index_map_.GetProfileIndexFor(computation)] = 138 total_cycles_executed; 139 } 140 141 // Record extra metric. set_extra_metrics(const std::string & metric,uint64_t value)142 void set_extra_metrics(const std::string& metric, uint64_t value) { 143 profile_counters_[hlo_profile_index_map_.GetProfileIndexFor(metric)] = 144 value; 145 } 146 147 // Returns a version of the execution profile suitable for performance 148 // debugging; e.g. emits cycle counts, execution time at the nominal device 149 // frequency, and the effective throughput given the provided cost_analysis 150 // for the operations in a given computation. Returns an empty string if it 151 // wasn't possible to generate a printable version. ToString(float clock_rate_ghz)152 std::string ToString(float clock_rate_ghz) const { 153 return PrintHloProfile(hlo_profile_printer_data_, profile_counters_.data(), 154 clock_rate_ghz); 155 } 156 mutable_profile_counters()157 std::vector<int64_t>* mutable_profile_counters() { 158 return &profile_counters_; 159 } profile_counters()160 const std::vector<int64_t>& profile_counters() const { 161 return profile_counters_; 162 } 163 164 HloExecutionProfileData ToProto() const; 165 166 private: 167 const HloProfilePrinterData& hlo_profile_printer_data_; 168 const HloProfileIndexMap& hlo_profile_index_map_; 169 170 // Stores per-Hlo profile counters. This is the only thing that changes when 171 // we execute an XLA computation. 172 std::vector<int64_t> profile_counters_; 173 }; 174 175 } // namespace xla 176 177 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_EXECUTION_PROFILE_H_ 178