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