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_profile_printer.h" 25 #include "tensorflow/compiler/xla/types.h" 26 #include "tensorflow/core/platform/stream_executor_no_cuda.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 using DeviceDescription = se::DeviceDescription; 109 110 HloExecutionProfile(const HloProfilePrinterData* hlo_profile_printer_data, 111 const HloProfileIndexMap* hlo_profile_index_map); 112 113 // Record how many cycles this HLO took to execute. 114 void SetCyclesTakenBy(const HloInstruction* hlo, uint64 cycles_taken); 115 116 // Returns how many cycles this HLO took to execute. Profiling information 117 // may not be available for some instructions in which case zero is returned. 118 uint64 GetCyclesTakenBy(const HloInstruction& hlo) const; 119 120 // Return the number of cycles this computation took to execute. total_cycles_executed(const HloComputation & computation)121 uint64 total_cycles_executed(const HloComputation& computation) const { 122 return profile_counters_[hlo_profile_index_map_.GetProfileIndexFor( 123 computation)]; 124 } 125 126 // Record how many cycles a computation took to execute. set_total_cycles_executed(const HloComputation & computation,uint64 total_cycles_executed)127 void set_total_cycles_executed(const HloComputation& computation, 128 uint64 total_cycles_executed) { 129 profile_counters_[hlo_profile_index_map_.GetProfileIndexFor(computation)] = 130 total_cycles_executed; 131 } 132 133 // Record extra metric. set_extra_metrics(const string & metric,uint64 value)134 void set_extra_metrics(const string& metric, uint64 value) { 135 profile_counters_[hlo_profile_index_map_.GetProfileIndexFor(metric)] = 136 value; 137 } 138 139 // Returns a version of the execution profile suitable for performance 140 // debugging; e.g. emits cycle counts, execution time at the nominal device 141 // frequency, and the effective throughput given the provided cost_analysis 142 // for the operations in a given computation. Returns an empty string if it 143 // wasn't possible to generate a printable version. ToString(const DeviceDescription & device_description)144 string ToString(const DeviceDescription& device_description) const { 145 return PrintHloProfile(hlo_profile_printer_data_, profile_counters_.data(), 146 device_description.clock_rate_ghz()); 147 } 148 mutable_profile_counters()149 std::vector<int64>* mutable_profile_counters() { return &profile_counters_; } profile_counters()150 const std::vector<int64>& profile_counters() const { 151 return profile_counters_; 152 } 153 154 private: 155 const HloProfilePrinterData& hlo_profile_printer_data_; 156 const HloProfileIndexMap& hlo_profile_index_map_; 157 158 // Stores per-Hlo profile counters. This is the only thing that changes when 159 // we execute an XLA computation. 160 std::vector<int64> profile_counters_; 161 }; 162 163 } // namespace xla 164 165 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_EXECUTION_PROFILE_H_ 166