Searched refs:host_memory_allocator (Results 1 – 7 of 7) sorted by relevance
39 Allocator* host_memory_allocator) in GPUDeviceContext() argument48 host_memory_allocator_(host_memory_allocator) { in GPUDeviceContext()63 Allocator* host_memory_allocator() const override { in host_memory_allocator() function
341 Allocator* host_memory_allocator = device_context->host_memory_allocator(); in CopyCPUTensorToGPU() local353 if (host_memory_allocator == nullptr) { in CopyCPUTensorToGPU()363 staging_buffer = host_memory_allocator->AllocateRaw( in CopyCPUTensorToGPU()379 host_memory_allocator]() { in CopyCPUTensorToGPU()381 host_memory_allocator->DeallocateRaw(staging_buffer); in CopyCPUTensorToGPU()
453 Allocator* host_memory_allocator = GetAllocator(attr); in Init() local461 stream_->device_to_device, host_memory_allocator); in Init()
209 std::unique_ptr<tensorflow::Allocator> host_memory_allocator, in PjRtStreamExecutorClient() argument215 host_memory_allocator_(std::move(host_memory_allocator)), in PjRtStreamExecutorClient()752 void* staging_buffer = host_memory_allocator()->AllocateRaw( in BufferFromHostBuffer()759 on_delete_callback = [staging_buffer, host_memory_allocator = in BufferFromHostBuffer()760 host_memory_allocator()]() { in BufferFromHostBuffer()761 host_memory_allocator->DeallocateRaw(staging_buffer); in BufferFromHostBuffer()792 void* ptr = host_memory_allocator()->AllocateRaw( in BufferFromHostBuffer()795 ptr, [host_memory_allocator = host_memory_allocator()](void* ptr) { in BufferFromHostBuffer()796 host_memory_allocator->DeallocateRaw(ptr); in BufferFromHostBuffer()
437 auto host_memory_allocator = in GetGpuClient() local453 std::move(host_memory_allocator), in GetGpuClient()
152 std::unique_ptr<tensorflow::Allocator> host_memory_allocator,279 tensorflow::Allocator* host_memory_allocator() const { in host_memory_allocator() function
123 virtual Allocator* host_memory_allocator() const { return nullptr; } in host_memory_allocator() function