Home
last modified time | relevance | path

Searched refs:alloc_attrs (Results 1 – 25 of 34) sorted by relevance

12

/external/tensorflow/tensorflow/core/common_runtime/
Drendezvous_util.cc24 const std::vector<AllocatorAttributes>& alloc_attrs, in SendTensorsToRendezvous() argument
31 if (!alloc_attrs.empty() && (keys.size() != alloc_attrs.size())) { in SendTensorsToRendezvous()
35 "; alloc_attrs.size() = ", alloc_attrs.size()); in SendTensorsToRendezvous()
46 if (!alloc_attrs.empty()) { in SendTensorsToRendezvous()
47 rendez_args.alloc_attrs = alloc_attrs[i]; in SendTensorsToRendezvous()
58 const std::vector<AllocatorAttributes>& alloc_attrs, in RecvOutputsFromRendezvousAsync() argument
65 if (!alloc_attrs.empty() && (keys.size() != alloc_attrs.size())) { in RecvOutputsFromRendezvousAsync()
68 keys.size(), "; alloc_attrs.size() = ", alloc_attrs.size())); in RecvOutputsFromRendezvousAsync()
84 if (!alloc_attrs.empty()) { in RecvOutputsFromRendezvousAsync()
85 alloc_attr = alloc_attrs[i]; in RecvOutputsFromRendezvousAsync()
[all …]
Drendezvous_mgr.cc69 (send_args.alloc_attrs.on_host() || parsed.src.type == "CPU"); in SameWorkerRecvDone()
71 (recv_args.alloc_attrs.on_host() || parsed.dst.type == "CPU"); in SameWorkerRecvDone()
100 AllocatorAttributes attr = recv_args.alloc_attrs; in SameWorkerRecvDone()
101 attr.set_gpu_compatible(send_args.alloc_attrs.gpu_compatible() || in SameWorkerRecvDone()
102 recv_args.alloc_attrs.gpu_compatible()); in SameWorkerRecvDone()
112 send_args.alloc_attrs, recv_args.alloc_attrs, &in, out, in SameWorkerRecvDone()
Drendezvous_util.h35 const std::vector<AllocatorAttributes>& alloc_attrs,
44 const std::vector<AllocatorAttributes>& alloc_attrs,
Drenamed_device.h92 const AllocatorAttributes alloc_attrs, in MakeTensorFromProto() argument
94 return underlying_->MakeTensorFromProto(tensor_proto, alloc_attrs, tensor); in MakeTensorFromProto()
Dthreadpool_device.h39 const AllocatorAttributes alloc_attrs,
Dprocess_function_library_runtime.h61 const std::vector<AllocatorAttributes>& alloc_attrs,
74 const std::vector<AllocatorAttributes>& alloc_attrs,
Dsingle_threaded_cpu_device.cc66 const AllocatorAttributes alloc_attrs, in MakeTensorFromProto() argument
Dthreadpool_device.cc92 const TensorProto& tensor_proto, const AllocatorAttributes alloc_attrs, in MakeTensorFromProto() argument
/external/tensorflow/tensorflow/core/grappler/optimizers/
Dgpu_swapping_kernels.cc37 AllocatorAttributes alloc_attrs; in ComputeAsync() local
38 alloc_attrs.set_gpu_compatible(true); in ComputeAsync()
39 alloc_attrs.set_on_host(true); in ComputeAsync()
42 ctx, ctx->allocate_output(0, input.shape(), &output, alloc_attrs), in ComputeAsync()
Devaluation_utils.h41 const AllocatorAttributes alloc_attrs,
/external/tensorflow/tensorflow/core/kernels/
Dstack.cc47 AllocatorAttributes alloc_attrs; member
236 AllocatorAttributes alloc_attrs = ctx->input_alloc_attr(1); in ComputeAsync() local
242 if (swap_memory_ && !alloc_attrs.on_host() && in ComputeAsync()
246 Allocator* allocator = device->GetAllocator(alloc_attrs); in ComputeAsync()
263 AllocatorAttributes alloc_attrs = ctx->input_alloc_attr(1); in ComputeAsync() local
264 ctx->SetStatus(stack->Push({*cpu_tensor, alloc_attrs, true})); in ComputeAsync()
277 OP_REQUIRES_OK_ASYNC(ctx, stack->Push({tensor, alloc_attrs, false}), done); in ComputeAsync()
304 Allocator* gpu_allocator = device->GetAllocator(value.alloc_attrs); in ComputeAsync()
Dsendrecv_ops.cc88 args.alloc_attrs = ctx->input_alloc_attr(0); in Compute()
178 args.alloc_attrs = ctx->output_alloc_attr(0); in ComputeAsync()
Dpad_op.cc117 AllocatorAttributes alloc_attrs; in Compute() local
118 alloc_attrs.set_on_host(context->input_memory_type(0) == HOST_MEMORY); in Compute()
122 &collapsed_output, alloc_attrs)); in Compute()
/external/tensorflow/tensorflow/contrib/gdr/
Dgdr_worker.cc91 const bool on_host = send_args.alloc_attrs.on_host(); in GrpcRecvTensorAsync()
118 AllocatorAttributes alloc_attrs; in GrpcRecvTensorAsync() local
119 alloc_attrs.set_gpu_compatible(true); in GrpcRecvTensorAsync()
120 alloc_attrs.set_on_host(true); in GrpcRecvTensorAsync()
121 Allocator* alloc = src_dev->GetAllocator(alloc_attrs); in GrpcRecvTensorAsync()
Dgdr_memory_manager.cc375 AllocatorAttributes alloc_attrs; in TransportOptionsFromTensor() local
376 alloc_attrs.set_gpu_compatible(true); in TransportOptionsFromTensor()
377 alloc_attrs.set_nic_compatible(true); in TransportOptionsFromTensor()
378 alloc_attrs.set_on_host(true); in TransportOptionsFromTensor()
379 Allocator* alloc = device->GetAllocator(alloc_attrs); in TransportOptionsFromTensor()
458 AllocatorAttributes alloc_attrs; in TensorFromTransportOptions() local
459 alloc_attrs.set_gpu_compatible(true); in TensorFromTransportOptions()
460 alloc_attrs.set_nic_compatible(true); in TensorFromTransportOptions()
461 alloc_attrs.set_on_host(true); in TensorFromTransportOptions()
462 Allocator* alloc = device->GetAllocator(alloc_attrs); in TensorFromTransportOptions()
Dgdr_rendezvous_mgr.cc58 resp_.InitAlloc(dst_device_, recv_args_.alloc_attrs); in Start()
63 const bool on_host = recv_args_.alloc_attrs.on_host(); in Start()
/external/tensorflow/tensorflow/core/common_runtime/sycl/
Dsycl_device.cc46 const AllocatorAttributes alloc_attrs, in MakeTensorFromProto() argument
58 if (alloc_attrs.on_host()) { in MakeTensorFromProto()
61 Tensor copy(GetAllocator(alloc_attrs), parsed.dtype(), parsed.shape()); in MakeTensorFromProto()
/external/tensorflow/tensorflow/core/distributed_runtime/
Dbase_rendezvous_mgr.cc236 (send_args.alloc_attrs.on_host() || parsed.src.type == "CPU"); in SameWorkerRecvDone()
238 (recv_args.alloc_attrs.on_host() || parsed.dst.type == "CPU"); in SameWorkerRecvDone()
268 AllocatorAttributes attr = recv_args.alloc_attrs; in SameWorkerRecvDone()
269 attr.set_gpu_compatible(send_args.alloc_attrs.gpu_compatible() || in SameWorkerRecvDone()
270 recv_args.alloc_attrs.gpu_compatible()); in SameWorkerRecvDone()
283 send_args.alloc_attrs, recv_args.alloc_attrs, &in, out, in SameWorkerRecvDone()
/external/tensorflow/tensorflow/core/distributed_runtime/rpc/
Drpc_rendezvous_mgr.cc62 AllocatorAttributes alloc_attrs, Device* dst_device, in Init() argument
65 alloc_attrs_ = alloc_attrs; in Init()
253 call->Init(rwi, step_id_, parsed.FullKey(), recv_args.alloc_attrs, dst_device, in RecvFromRemoteAsync()
Dgrpc_worker_service.cc491 const bool on_host = send_args.alloc_attrs.on_host(); in GrpcRecvTensorAsync()
496 AllocatorAttributes alloc_attrs; in GrpcRecvTensorAsync() local
497 alloc_attrs.set_gpu_compatible(true); in GrpcRecvTensorAsync()
498 alloc_attrs.set_on_host(true); in GrpcRecvTensorAsync()
499 Allocator* alloc = src_dev->GetAllocator(alloc_attrs); in GrpcRecvTensorAsync()
/external/tensorflow/tensorflow/contrib/mpi/
Dmpi_rendezvous_mgr.cc121 recv_args.alloc_attrs, &val); in RecvFromRemoteAsync()
124 tr.InitAlloc(dst_device, recv_args.alloc_attrs); in RecvFromRemoteAsync()
242 (!send_args.alloc_attrs.on_host())) { in AddRequest()
/external/tensorflow/tensorflow/compiler/tf2xla/
Dxla_compilation_device.h60 const AllocatorAttributes alloc_attrs,
/external/tensorflow/tensorflow/compiler/jit/
Dxla_device.cc441 const AllocatorAttributes alloc_attrs, in MakeTensorFromProto() argument
452 if (alloc_attrs.on_host()) { in MakeTensorFromProto()
458 Allocator* allocator = GetAllocatorLocked(alloc_attrs); in MakeTensorFromProto()
/external/tensorflow/tensorflow/core/common_runtime/gpu/
Dgpu_device.h90 const AllocatorAttributes alloc_attrs,
177 Status MaybeCopyTensorToGPU(const AllocatorAttributes& alloc_attrs,
/external/tensorflow/tensorflow/core/framework/
Drendezvous.h50 AllocatorAttributes alloc_attrs; member

12