1 /* Copyright 2018 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_GPU_HLO_EXECUTION_PROFILER_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_HLO_EXECUTION_PROFILER_H_ 18 19 #include <memory> 20 #include <stack> 21 #include <unordered_set> 22 #include <vector> 23 24 #include "tensorflow/compiler/xla/service/hlo_computation.h" 25 #include "tensorflow/compiler/xla/service/hlo_execution_profile.h" 26 #include "tensorflow/compiler/xla/service/hlo_instruction.h" 27 #include "tensorflow/compiler/xla/service/stream_pool.h" 28 #include "tensorflow/core/platform/stream_executor_no_cuda.h" 29 30 namespace xla { 31 namespace gpu { 32 33 class ScopedInstructionProfiler; 34 35 // A helper class for profiling HLO in the course of GPU program execution. 36 // All of the profiling is guarded internally, to avoid the caller needing to 37 // have lots of conditionals sprinkled around. 38 class HloExecutionProfiler { 39 public: 40 // If profiling is enabled, start an execution timer running. 41 explicit HloExecutionProfiler(bool do_profile, HloExecutionProfile* profile, 42 se::Stream* stream, 43 const std::vector<StreamPool::Ptr>& sub_streams, 44 const HloComputation* computation); 45 46 // If profiling is enabled, sets the total cycle count on the profile from the 47 // execution timer. 48 void FinishExecution(); 49 50 // If profiling is enabled, starts a timer for a (sub)computation. 51 void StartHloComputation(); 52 53 // If profiling is enabled stops the timer for a (sub)computation and records 54 // the time that the computation took to execute in the profile. 55 void FinishHloComputation(const HloComputation* computation); 56 57 // If profiling is enabled, starts a per-operation timer. 58 void StartHloInstruction(); 59 60 // If profiling is enabled, stops the per-operation timer and records the time 61 // that the hlo_instruction took to execute in the profile. 62 void FinishHloInstruction(const HloInstruction* hlo_instruction); 63 64 // Returns a ScopedInstructionProfiler and triggers a call to 65 // StartHloInstruction(). Once the returned ScopedInstructionProfiler goes 66 // out of scope, it triggers a call to FinishHloInstruction(). 67 std::unique_ptr<ScopedInstructionProfiler> MakeScopedInstructionProfiler( 68 const HloInstruction* hlo_instruction); 69 70 private: 71 const bool do_profile_; 72 double clock_rate_ghz_; 73 HloExecutionProfile* profile_; 74 se::Stream* stream_; 75 const std::vector<StreamPool::Ptr>& sub_streams_; 76 const HloComputation* computation_; 77 std::stack<std::unique_ptr<se::Timer>> timers_; 78 // Contains the HLO instructions for which we are currently measuring the 79 // time. 80 std::unordered_set<const HloInstruction*> hlo_instructions_; 81 bool finished_execution_ = false; 82 }; 83 84 // This class can be used within the ExecuteOnStream() implementations of 85 // Thunks. It ensures that we always have a pair of matching 86 // StartHloInstruction() and FinishHloInstruction() calls to the profiler. 87 class ScopedInstructionProfiler { 88 public: ScopedInstructionProfiler(HloExecutionProfiler * profiler,const HloInstruction * hlo_instruction)89 ScopedInstructionProfiler(HloExecutionProfiler* profiler, 90 const HloInstruction* hlo_instruction) 91 : profiler_(profiler), hlo_instruction_(hlo_instruction) { 92 if (hlo_instruction != nullptr) { 93 profiler->StartHloInstruction(); 94 } 95 } ~ScopedInstructionProfiler()96 ~ScopedInstructionProfiler() { 97 if (hlo_instruction_ != nullptr) { 98 profiler_->FinishHloInstruction(hlo_instruction_); 99 } 100 } 101 102 private: 103 HloExecutionProfiler* profiler_; 104 const HloInstruction* hlo_instruction_; 105 }; 106 107 } // namespace gpu 108 } // namespace xla 109 110 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_HLO_EXECUTION_PROFILER_H_ 111