| 74 | } |
| 75 | |
| 76 | void XlaTensor::WaitForDefinitionEventOnStream(se::Stream* stream) { |
| 77 | mutex_lock lock(mu_); |
| 78 | if (!definition_event_) { |
| 79 | return; |
| 80 | } |
| 81 | |
| 82 | // The set of defined streams is expected to be very small indeed (usually |
| 83 | // 1-2), so a simple linear scan should be fast enough. |
| 84 | if (std::find(streams_defined_on_.begin(), streams_defined_on_.end(), |
| 85 | stream) != streams_defined_on_.end()) { |
| 86 | // stream is in streams_defined_on_; it doesn't need to be waited on. |
| 87 | return; |
| 88 | } |
| 89 | |
| 90 | stream->ThenWaitFor(definition_event_.get()); |
| 91 | streams_defined_on_.push_back(stream); |
| 92 | } |
| 93 | |
| 94 | void XlaTensor::ResetDefinitionEvent(std::shared_ptr<se::Event> event, |
| 95 | se::Stream* stream) { |
no test coverage detected