| 707 | } |
| 708 | |
| 709 | Result<std::shared_ptr<CudaBuffer>> CudaContext::OpenIpcBuffer( |
| 710 | const CudaIpcMemHandle& ipc_handle) { |
| 711 | if (ipc_handle.memory_size() > 0) { |
| 712 | ContextSaver set_temporary(*this); |
| 713 | uint8_t* data = nullptr; |
| 714 | RETURN_NOT_OK(impl_->OpenIpcBuffer(ipc_handle, &data)); |
| 715 | // Need to ask the device how big the buffer is |
| 716 | size_t allocation_size = 0; |
| 717 | CU_RETURN_NOT_OK("cuMemGetAddressRange", |
| 718 | cuMemGetAddressRange(nullptr, &allocation_size, |
| 719 | reinterpret_cast<CUdeviceptr>(data))); |
| 720 | return std::make_shared<CudaBuffer>(data, allocation_size, this->shared_from_this(), |
| 721 | true, true); |
| 722 | } else { |
| 723 | // zero-sized buffer does not own data (which is nullptr), hence |
| 724 | // CloseIpcBuffer will not be called (see CudaBuffer::Close). |
| 725 | return std::make_shared<CudaBuffer>(nullptr, 0, this->shared_from_this(), false, |
| 726 | true); |
| 727 | } |
| 728 | } |
| 729 | |
| 730 | Status CudaContext::CloseIpcBuffer(CudaBuffer* buf) { |
| 731 | ContextSaver set_temporary(*this); |