Searched refs:GPUDeviceContext (Results 1 – 6 of 6) sorted by relevance
26 void GPUDeviceContext::CopyCPUTensorToDevice(const Tensor* cpu_tensor, in CopyCPUTensorToDevice()33 void GPUDeviceContext::CopyDeviceTensorToCPU(const Tensor* device_tensor, in CopyDeviceTensorToCPU()40 void GPUDeviceContext::CopyTensorInSameDevice(const Tensor* input_tensor, in CopyTensorInSameDevice()48 Status GPUDeviceContext::ThenExecute(Device* device, se::Stream* stream, in ThenExecute()
78 auto gs = static_cast<const GPUDeviceContext*>(ctx)->stream(); in PrepareCopy()130 static_cast<const GPUDeviceContext*>(device_context) in SetProtoFromGPU()202 static_cast<const GPUDeviceContext*>(send_dev_context) in DeviceToDeviceCopy()219 static_cast<const GPUDeviceContext*>(recv_dev_context)->stream(); in DeviceToDeviceCopy()270 static_cast<const GPUDeviceContext*>(device_context) in CopyGPUTensorToCPU()315 static_cast<const GPUDeviceContext*>(device_context) in CopyCPUTensorToGPU()
390 device_contexts_.push_back(new GPUDeviceContext( in Init()516 GPUDeviceContext* gpu_device_context = device_contexts_[0]; in ComputeHelper()519 static_cast<GPUDeviceContext*>(context->op_device_context()); in ComputeHelper()537 const GPUDeviceContext* idc = in ComputeHelper()538 static_cast<GPUDeviceContext*>(context->input_device_context(i)); in ComputeHelper()603 GPUDeviceContext* gpu_device_context = device_contexts_[0]; in ConsumeListOfAccessedTensors()605 gpu_device_context = static_cast<GPUDeviceContext*>(device_context); in ConsumeListOfAccessedTensors()618 GPUDeviceContext* gpu_device_context = device_contexts_[0]; in ComputeAsync()621 static_cast<GPUDeviceContext*>(context->op_device_context()); in ComputeAsync()952 const GPUDeviceContext* gpu_dc = static_cast<GPUDeviceContext*>(dc); in ReinitializeGpuDevice()
149 std::vector<GPUDeviceContext*> device_contexts_;
29 class GPUDeviceContext : public DeviceContext {32 GPUDeviceContext(int stream_id, se::Stream* stream, in GPUDeviceContext() function42 ~GPUDeviceContext() override {} in ~GPUDeviceContext()
332 static_cast<GPUDeviceContext*>(ctx->op_device_context()); in PopulateOutputs()