class HloExecutionProfiler {
public:
// If profiling is enabled, start an execution timer running.
- explicit HloExecutionProfiler(bool do_profile, HloExecutionProfile* profile,
- se::Stream* stream,
- const HloComputation* computation)
+ explicit HloExecutionProfiler(
+ bool do_profile, HloExecutionProfile* profile, se::Stream* stream,
+ const std::vector<Pool<se::Stream>::SmartPtr>& sub_streams,
+ const HloComputation* computation)
: do_profile_(do_profile),
profile_(profile),
stream_(stream),
+ sub_streams_(sub_streams),
computation_(computation) {
if (do_profile_) {
clock_rate_ghz_ =
CHECK(!finished_execution_) << "Call FinishExecution only once!";
finished_execution_ = true;
if (do_profile_) {
+ stream_->ThenWaitFor(&sub_streams_);
stream_->ThenStopTimer(execution_timer_.get());
stream_->BlockHostUntilDone().IgnoreError();
profile_->set_total_cycles_executed(
// that the hlo_instruction took to execute in the profile.
void FinishOperation(const HloInstruction* hlo_instruction) {
if (do_profile_) {
+ stream_->ThenWaitFor(&sub_streams_);
stream_->ThenStopTimer(per_op_timer_.get());
stream_->BlockHostUntilDone().IgnoreError();
profile_->SetCyclesTakenBy(
double clock_rate_ghz_;
HloExecutionProfile* profile_;
se::Stream* stream_;
+ const std::vector<Pool<se::Stream>::SmartPtr>& sub_streams_;
const HloComputation* computation_;
std::unique_ptr<se::Timer> execution_timer_;
std::unique_ptr<se::Timer> per_op_timer_;
LOG(WARNING) << "PROFILING: profiling is enabled";
}
- HloExecutionProfiler profiler(do_profile, hlo_execution_profile, main_stream,
- hlo_module_->entry_computation());
-
- uint64 start_micros = tensorflow::Env::Default()->NowMicros();
-
// Stream 0 indicates `main_stream` and substreams start from stream 1.
std::vector<Pool<se::Stream>::SmartPtr> sub_streams;
+ sub_streams.reserve(thunk_schedule_->StreamCount() - 1);
while (sub_streams.size() + 1 < thunk_schedule_->StreamCount()) {
sub_streams.emplace_back();
TF_ASSIGN_OR_RETURN(
run_options->BorrowStream(main_stream->parent()->device_ordinal()));
}
+ HloExecutionProfiler profiler(do_profile, hlo_execution_profile, main_stream,
+ sub_streams, hlo_module_->entry_computation());
+ uint64 start_micros = tensorflow::Env::Default()->NowMicros();
+
// The next event enqueued on stream N must not run until the thunk at
// last_blocking_thunk_for_stream[N] completes.
std::map<int32, const Thunk*> last_blocking_thunk_for_stream;