| 406 | } |
| 407 | |
| 408 | Maybe<void> InstructionsBuilder::ReleaseTensor( |
| 409 | const std::shared_ptr<vm::EagerBlobObject>& eager_blob_object) { |
| 410 | const auto& last_used_stream = JUST(eager_blob_object->last_used_stream()); |
| 411 | const auto& producer_stream = JUST(eager_blob_object->producer_stream()); |
| 412 | if (pthread_fork::IsForkedSubProcess() |
| 413 | && producer_stream->device()->enum_type() != DeviceType::kCPU) { |
| 414 | return Maybe<void>::Ok(); |
| 415 | } |
| 416 | Optional<Symbol<Stream>> stream{}; |
| 417 | if (*one::CurrentDevVmDepObjectConsumeMode() == one::DevVmDepObjectConsumeMode::NONE) { |
| 418 | stream = Optional<Symbol<Stream>>(NullOpt); |
| 419 | } else if (IsCommNetStream::Visit(last_used_stream->stream_type())) { |
| 420 | // Disable inter-device instruction sequential for tensor used by communicative stream. |
| 421 | // It's not acceptable for us that cuda compute stream is blocked by cuda nccl stream. |
| 422 | stream = Optional<Symbol<Stream>>(NullOpt); |
| 423 | } else if (IsCommNetStream::Visit(producer_stream->stream_type())) { |
| 424 | // Disable inter-device instruction sequential for tensor produced by communicative stream. |
| 425 | stream = Optional<Symbol<Stream>>(NullOpt); |
| 426 | } else { |
| 427 | stream = producer_stream; |
| 428 | } |
| 429 | struct EnableStreamWaitOnReleaseTensor final |
| 430 | : public StreamTypeVisitor<EnableStreamWaitOnReleaseTensor> { |
| 431 | static bool VisitCompute() { return true; } |
| 432 | static bool VisitHost2Device() { return true; } |
| 433 | static bool VisitDevice2Host() { return true; } |
| 434 | static bool VisitCcl() { return false; } |
| 435 | static bool VisitBarrier() { return false; } |
| 436 | static bool VisitCriticalSection() { return false; } |
| 437 | static bool VisitLazyJobLauncher() { return false; } |
| 438 | static bool VisitPinnedCompute() { return VisitCompute(); } |
| 439 | }; |
| 440 | const auto& EnableStreamWait = [&] { |
| 441 | if (last_used_stream->device() != producer_stream->device()) { return false; } |
| 442 | if (last_used_stream->stream_type() == producer_stream->stream_type()) { return true; } |
| 443 | return EnableStreamWaitOnReleaseTensor::Visit(last_used_stream->stream_type()) |
| 444 | && EnableStreamWaitOnReleaseTensor::Visit(producer_stream->stream_type()); |
| 445 | }; |
| 446 | if (last_used_stream != producer_stream) { |
| 447 | if (stream.has_value() && EnableStreamWait()) { |
| 448 | JUST(SoftSyncStreamBetween({JUST(eager_blob_object->compute_local_dep_object())}, |
| 449 | last_used_stream, JUST(stream))); |
| 450 | } else { |
| 451 | JUST(RecordEvent({JUST(eager_blob_object->compute_local_dep_object())}, last_used_stream)); |
| 452 | } |
| 453 | eager_blob_object->set_last_used_stream(producer_stream); |
| 454 | } |
| 455 | auto vm_stream = stream.map([](Symbol<Stream> stream) -> vm::Stream* { |
| 456 | return CHECK_JUST(Singleton<VirtualMachine>::Get()->GetVmStream(stream)); |
| 457 | }); |
| 458 | StreamType stream_type = producer_stream->stream_type(); |
| 459 | auto instruction = intrusive::make_shared<vm::Instruction>( |
| 460 | JUST(Singleton<VirtualMachine>::Get()->GetVmStream(producer_stream)), |
| 461 | JUST(vm::MakeReleaseTensorInstructionPolicy::Visit(stream_type, eager_blob_object, |
| 462 | vm_stream))); |
| 463 | instruction_list_->EmplaceBack(std::move(instruction)); |
| 464 | |
| 465 | return Maybe<void>::Ok(); |
no test coverage detected