| 756 | |
| 757 | template <class PropagatorStateType> |
| 758 | void ExecutorState<PropagatorStateType>::ProcessAsync( |
| 759 | const NodeItem& item, const OpKernelContext::Params& params, |
| 760 | const TaggedNode& tagged_node, Entry* first_input, |
| 761 | NodeExecStatsInterface* stats) { |
| 762 | AsyncOpKernel* async_kernel = item.kernel->AsAsync(); |
| 763 | DCHECK(async_kernel != nullptr); |
| 764 | AsyncState* state = |
| 765 | new AsyncState(params, tagged_node, &item, first_input, stats); |
| 766 | |
| 767 | ExecutorInternal::KernelStatsInfo kernel_stat_buffer; |
| 768 | kernel_stats_->StartCollectOp(&item, &kernel_stat_buffer); |
| 769 | auto done = [this, state, kernel_stat_buffer{std::move(kernel_stat_buffer)}]() { |
| 770 | Device* device = immutable_state_.params().device; |
| 771 | NodeExecStatsInterface* stats = state->stats; // Shorthand |
| 772 | Entry* first_input = state->first_input; // Shorthand |
| 773 | |
| 774 | nodestats::SetOpEnd(stats); |
| 775 | this->GetKernelStats()->StopCollectOp(state->item, |
| 776 | const_cast<ExecutorInternal::KernelStatsInfo*>(&kernel_stat_buffer)); |
| 777 | EntryVector outputs(state->item->num_outputs); |
| 778 | Status s = ProcessOutputs(*state->item, &state->ctx, outputs.data(), stats); |
| 779 | nodestats::SetMemory(stats, &state->ctx); |
| 780 | if (vlog_) { |
| 781 | VLOG(2) << "Async kernel done: " << state->item->node->id() << " step " |
| 782 | << step_id_ << " " << SummarizeNodeDef(state->item->kernel->def()) |
| 783 | << (state->tagged_node.get_is_dead() ? " is dead" : "") |
| 784 | << " device: " << device->name(); |
| 785 | } |
| 786 | |
| 787 | // Clears inputs. |
| 788 | const int num_inputs = state->item->num_inputs; |
| 789 | for (int i = 0; i < num_inputs; ++i) { |
| 790 | (first_input + i)->ClearVal(); |
| 791 | } |
| 792 | propagator_.MaybeMarkCompleted(state->tagged_node); |
| 793 | TaggedNodeSeq ready; |
| 794 | if (s.ok()) { |
| 795 | propagator_.PropagateOutputs(state->tagged_node, &outputs, &ready); |
| 796 | } |
| 797 | outputs.clear(); |
| 798 | const bool completed = NodeDone(s, &ready, stats, nullptr); |
| 799 | delete state; |
| 800 | if (completed) ScheduleFinish(); |
| 801 | }; |
| 802 | nodestats::SetOpStart(stats); |
| 803 | { |
| 804 | profiler::TraceMe activity( |
| 805 | [&] { |
| 806 | int64 id = step_id_; |
| 807 | if (step_container_ && step_container_->StepId()) { |
| 808 | id = step_container_->StepId(); |
| 809 | } |
| 810 | return strings::StrCat( |
| 811 | async_kernel->name(), ":", async_kernel->type_string(), |
| 812 | "#id=", id, ",device=", immutable_state_.params().device->name(), |
| 813 | ",async=true#"); |
| 814 | }, |
| 815 | profiler::GetTFTraceMeLevel(async_kernel->IsExpensive())); |
nothing calls this directly
no test coverage detected