| 456 | } |
| 457 | |
| 458 | Result<std::unique_ptr<Buffer>> CudaMemoryManager::CopyNonOwnedFrom( |
| 459 | const Buffer& buf, const std::shared_ptr<MemoryManager>& from) { |
| 460 | if (from->is_cpu()) { |
| 461 | // CPU-to-device copy |
| 462 | ARROW_ASSIGN_OR_RAISE(auto to_context, cuda_device()->GetContext()); |
| 463 | ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> dest, to_context->Allocate(buf.size())); |
| 464 | RETURN_NOT_OK(to_context->CopyHostToDevice(dest->address(), buf.data(), buf.size())); |
| 465 | return dest; |
| 466 | } |
| 467 | if (IsCudaMemoryManager(*from)) { |
| 468 | // Device-to-device copy |
| 469 | ARROW_ASSIGN_OR_RAISE(auto to_context, cuda_device()->GetContext()); |
| 470 | ARROW_ASSIGN_OR_RAISE( |
| 471 | auto from_context, |
| 472 | checked_cast<const CudaMemoryManager&>(*from).cuda_device()->GetContext()); |
| 473 | ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> dest, to_context->Allocate(buf.size())); |
| 474 | if (to_context->handle() == from_context->handle()) { |
| 475 | // Same context |
| 476 | RETURN_NOT_OK( |
| 477 | to_context->CopyDeviceToDevice(dest->address(), buf.address(), buf.size())); |
| 478 | } else { |
| 479 | // Other context |
| 480 | RETURN_NOT_OK(from_context->CopyDeviceToAnotherDevice(to_context, dest->address(), |
| 481 | buf.address(), buf.size())); |
| 482 | } |
| 483 | return dest; |
| 484 | } |
| 485 | return nullptr; |
| 486 | } |
| 487 | |
| 488 | Result<std::shared_ptr<Buffer>> CudaMemoryManager::ViewBufferTo( |
| 489 | const std::shared_ptr<Buffer>& buf, const std::shared_ptr<MemoryManager>& to) { |
nothing calls this directly
no test coverage detected