| /external/tensorflow/tensorflow/compiler/xla/tests/ |
| D | multiple_devices_on_host_test.cc | 37 LocalExecutable* executable, int device_ordinal, LocalClient* client, in CompileAndExecute() argument
43 execute_options.set_device_ordinal(device_ordinal); in CompileAndExecute()
52 results->emplace_back(device_ordinal, std::move(result)); in CompileAndExecute()
76 for (int device_ordinal = 0; device_ordinal < device_count; in TestWithDeviceCount() local
77 device_ordinal++) { in TestWithDeviceCount()
79 tensorflow::ThreadOptions{}, absl::StrCat("thread-", device_ordinal), in TestWithDeviceCount()
80 [&executable, device_ordinal, client, &results_mutex, &results] { in TestWithDeviceCount()
81 CompileAndExecute(executable.get(), device_ordinal, client, in TestWithDeviceCount()
87 for (int device_ordinal = 0; device_ordinal < device_count; in TestWithDeviceCount() local
88 device_ordinal++) { in TestWithDeviceCount()
[all …]
|
| D | local_client_test_base.cc | 44 StatusOr<se::OwningDeviceMemory> TestAllocator::Allocate(int device_ordinal, in Allocate() argument
48 VLOG(2) << "Allocate(" << device_ordinal << ", " << size << ")"; in Allocate()
52 device_allocation_count_[device_ordinal]++; in Allocate()
55 device_ordinal, size, retry_on_failure, memory_space); in Allocate()
58 Status TestAllocator::Deallocate(int device_ordinal, se::DeviceMemoryBase mem) { in Deallocate() argument
59 VLOG(2) << "Deallocate(" << device_ordinal << ")"; in Deallocate()
63 device_deallocation_count_[device_ordinal]++; in Deallocate()
65 return se::StreamExecutorMemoryAllocator::Deallocate(device_ordinal, mem); in Deallocate()
73 int64 TestAllocator::allocation_count(int device_ordinal) const { in allocation_count()
75 auto it = device_allocation_count_.find(device_ordinal); in allocation_count()
[all …]
|
| /external/tensorflow/tensorflow/stream_executor/ |
| D | tf_allocator_adapter.h | 42 port::StatusOr<OwningDeviceMemory> Allocate(int device_ordinal, uint64 size,
46 port::Status Deallocate(int device_ordinal, DeviceMemoryBase mem) override;
57 port::StatusOr<Stream *> GetStream(int device_ordinal) override;
81 port::StatusOr<OwningDeviceMemory> Allocate(int device_ordinal, uint64 size, in Allocate() argument
84 CHECK_LT(device_ordinal, per_device_allocators_.size()); in Allocate()
85 return per_device_allocators_[device_ordinal].Allocate( in Allocate()
86 device_ordinal, size, retry_on_failure, memory_space); in Allocate()
89 port::Status Deallocate(int device_ordinal, DeviceMemoryBase mem) override { in Deallocate() argument
90 CHECK_LT(device_ordinal, per_device_allocators_.size()); in Deallocate()
91 return per_device_allocators_[device_ordinal].Deallocate(device_ordinal, in Deallocate()
[all …]
|
| D | device_memory_allocator.h | 60 ScopedDeviceMemory(DeviceMemoryBase mem, int device_ordinal, in ScopedDeviceMemory() argument
62 : wrapped_(mem), device_ordinal_(device_ordinal), allocator_(allocator) { in ScopedDeviceMemory()
139 int device_ordinal() const { return device_ordinal_; } in device_ordinal() function
176 virtual port::StatusOr<OwningDeviceMemory> Allocate(int device_ordinal,
187 port::StatusOr<OwningDeviceMemory> Allocate(int device_ordinal, uint64 size) { in Allocate() argument
188 return Allocate(device_ordinal, size, /*retry_on_failure=*/true, in Allocate()
193 port::StatusOr<OwningDeviceMemory> Allocate(int device_ordinal, uint64 size, in Allocate() argument
195 return Allocate(device_ordinal, size, retry_on_failure, in Allocate()
202 int device_ordinal, uint64 size, bool retry_on_failure = true,
204 return Allocate(device_ordinal, size, retry_on_failure, memory_space);
[all …]
|
| D | tf_allocator_adapter.cc | 39 int device_ordinal, uint64 size, bool retry_on_failure, in Allocate() argument
53 return OwningDeviceMemory(DeviceMemoryBase(data, size), device_ordinal, this); in Allocate()
56 port::Status TfAllocatorAdapter::Deallocate(int device_ordinal, in Deallocate() argument
62 port::StatusOr<Stream *> TfAllocatorAdapter::GetStream(int device_ordinal) { in GetStream() argument
63 CHECK_EQ(stream_->parent()->device_ordinal(), device_ordinal); in GetStream()
|
| /external/tensorflow/tensorflow/compiler/xrt/ |
| D | xrt_memory_manager.cc | 37 int64 MakeDeviceHandle(int64 device_ordinal, int64 rnd_value) { in MakeDeviceHandle() argument
39 return (device_ordinal << (64 - kDeviceBits)) | (rnd_value & kUidMask); in MakeDeviceHandle()
61 int64 handle = MakeDeviceHandle(tuple->device_ordinal(), CreateUid()); in Register()
219 DeviceContext* device_context = GetDeviceContext(tuple->device_ordinal(), in Register()
226 int device_ordinal = GetDeviceFromHandle(handle); in Lookup() local
227 DeviceContext* device_context = GetDeviceContext(device_ordinal, in Lookup()
240 int device_ordinal = GetDeviceFromHandle(handle); in Release() local
241 DeviceContext* device_context = GetDeviceContext(device_ordinal, in Release()
250 int device_ordinal) { in CompactAllocations() argument
251 DeviceContext* device_context = GetDeviceContext(device_ordinal, in CompactAllocations()
[all …]
|
| D | xrt_state.cc | 73 XRTMemoryManager* memory_manager, xla::Backend* backend, int device_ordinal, in AllocateScopedShapedBuffer() argument
76 TF_ASSIGN_OR_RETURN(auto stream, backend->BorrowStream(device_ordinal)); in AllocateScopedShapedBuffer()
96 shape, on_device_shape, backend->memory_allocator(), device_ordinal); in AllocateScopedShapedBuffer()
103 memory_manager->Allocate(backend, device_ordinal, size)); in AllocateScopedShapedBuffer()
120 int device_ordinal, in XRTBufferAllocation() argument
123 device_ordinal_(device_ordinal), in XRTBufferAllocation()
147 XRTTupleAllocation::XRTTupleAllocation(int device_ordinal, in XRTTupleAllocation() argument
151 : device_ordinal_(device_ordinal), in XRTTupleAllocation()
171 xla::Backend* backend, int device_ordinal, in CreateAndTransfer() argument
176 device_ordinal, literal.shape(), in CreateAndTransfer()
[all …]
|
| D | xrt_memory_manager.h | 98 Status CompactAllocations(xla::Backend* backend, int device_ordinal);
108 int device_ordinal,
118 xla::Backend* backend, int device_ordinal,
131 MemoryReclaimContext(xla::Backend* backend, int device_ordinal, in MemoryReclaimContext()
134 device_ordinal(device_ordinal), in MemoryReclaimContext()
138 const int device_ordinal = 0; member
145 DeviceContext* GetDeviceContext(int device_ordinal, bool create_if_missing);
160 int device_ordinal, size_t requested_free_size) { in Run() argument
161 MemoryReclaimContext mrctx(backend, device_ordinal, requested_free_size); in Run()
|
| D | xrt_tpu_device.cc | 37 Status XRTTpuDeviceAccessor::ScopedRef::Acquire(int device_ordinal) { in Acquire() argument
39 tpu::TpuNodeContext::Create(device_ordinal)); in Acquire()
40 ordinal_ = device_ordinal; in Acquire()
47 return Acquire(metadata->device_ordinal()); in Acquire()
51 OpKernelContext* /*unused ctx*/, int device_ordinal, in InitScopedRef() argument
53 return scoped_ref->Acquire(device_ordinal); in InitScopedRef()
|
| D | xrt_state.h | 54 int device_ordinal, se::DeviceMemoryAllocator* allocator);
85 xla::Backend* backend, int device_ordinal,
93 xla::Backend* backend, int device_ordinal,
98 xla::Backend* backend, int device_ordinal,
108 xla::Backend* backend, int device_ordinal,
140 xla::Backend* backend, int device_ordinal,
193 int device_ordinal() const;
237 XRTTupleAllocation(int device_ordinal, se::DeviceMemoryAllocator* allocator,
245 int device_ordinal);
268 const xla::ShapeTree<ExpandedTupleInput>& elements, int device_ordinal,
|
| D | xrt_device.cc | 71 OpKernelContext* ctx, int device_ordinal, ScopedRef* scoped_ref) { in InitScopedRef() argument
74 if (device_ordinal != metadata->device_ordinal()) { in InitScopedRef()
75 return errors::Internal("XRT device ordinal requested ", device_ordinal, in InitScopedRef()
77 metadata->device_ordinal()); in InitScopedRef()
79 scoped_ref->Acquire(metadata->client(), device_ordinal); in InitScopedRef()
87 scoped_ref->Acquire(metadata->client(), metadata->device_ordinal()); in InitScopedRef()
|
| /external/tensorflow/tensorflow/compiler/xla/client/ |
| D | local_client.cc | 35 StatusOr<StreamPool::Ptr> BorrowStreamForDevice(int device_ordinal, in BorrowStreamForDevice() argument
37 if (device_ordinal < 0) { in BorrowStreamForDevice()
38 device_ordinal = backend->default_device_ordinal(); in BorrowStreamForDevice()
40 return backend->BorrowStream(device_ordinal); in BorrowStreamForDevice()
50 CHECK_GE(build_options_.device_ordinal(), 0) in LocalExecutable()
72 if (run_options.device_ordinal() != -1) { in ValidateExecutionOptions()
81 int run_device_ordinal = run_options.device_ordinal(); in ValidateExecutionOptions()
84 ? run_options.stream()->parent()->device_ordinal() in ValidateExecutionOptions()
89 run_device_ordinal, build_options_.device_ordinal())); in ValidateExecutionOptions()
155 stream, BorrowStreamForDevice(run_options.device_ordinal(), backend_)); in RunHelper()
[all …]
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | backend.h | 111 StatusOr<se::StreamExecutor*> stream_executor(int device_ordinal) const;
124 StatusOr<StreamPool::Ptr> BorrowStream(int device_ordinal);
131 return [this](int device_ordinal) { return BorrowStream(device_ordinal); }; in StreamBorrower()
135 bool device_ordinal_supported(int device_ordinal) const { in device_ordinal_supported() argument
136 return (device_ordinal >= 0 && device_ordinal < device_count() && in device_ordinal_supported()
137 stream_executors_[device_ordinal] != nullptr); in device_ordinal_supported()
141 string device_name(int device_ordinal) const { in device_name() argument
142 return absl::StrCat(platform_->Name(), ":", device_ordinal); in device_name()
|
| D | backend.cc | 108 StatusOr<StreamPool::Ptr> Backend::BorrowStream(int device_ordinal) { in BorrowStream() argument
109 TF_ASSIGN_OR_RETURN(auto executor, stream_executor(device_ordinal)); in BorrowStream()
148 return default_stream_executor()->device_ordinal(); in default_device_ordinal()
167 int device_ordinal) const { in stream_executor()
168 if (device_ordinal < 0 || in stream_executor()
169 device_ordinal > stream_executors_.back()->device_ordinal()) { in stream_executor()
172 device_ordinal, stream_executors_.back()->device_ordinal()); in stream_executor()
175 if (executor->device_ordinal() == device_ordinal) { in stream_executor()
180 device_name(device_ordinal)); in stream_executor()
|
| D | shaped_buffer.cc | 34 ShapedBuffer::ShapedBuffer(Shape on_device_shape, int device_ordinal) in ShapedBuffer() argument
36 device_ordinal_(device_ordinal), in ShapedBuffer()
42 int device_ordinal) in ShapedBuffer() argument
43 : ShapedBuffer(on_device_shape, device_ordinal) {} in ShapedBuffer()
90 absl::StrCat("ShapedBuffer(", device_ordinal(), in ToString()
118 int device_ordinal) in ScopedShapedBuffer() argument
119 : ShapedBuffer(std::move(on_device_shape), device_ordinal), in ScopedShapedBuffer()
125 int device_ordinal) in ScopedShapedBuffer() argument
127 device_ordinal) {} in ScopedShapedBuffer()
170 TF_CHECK_OK(allocator_->Deallocate(device_ordinal(), memory_base)); in Deallocate()
[all …]
|
| D | shaped_buffer.h | 46 ShapedBuffer(Shape on_device_shape, int device_ordinal);
49 ShapedBuffer(Shape on_host_shape, Shape on_device_shape, int device_ordinal);
71 int device_ordinal() const { return device_ordinal_; } in device_ordinal() function
154 int device_ordinal);
158 int device_ordinal);
184 CHECK_EQ(buffer.device_ordinal(), device_ordinal()); in set_buffer()
|
| /external/tensorflow/tensorflow/core/tpu/ |
| D | tpu_execute.cc | 266 allocator->Allocate(stream->parent()->device_ordinal(), in UpdateDynamicInputs()
296 void TPUCancelExecution(Env* env, int device_ordinal) { in TPUCancelExecution() argument
299 << device_ordinal; in TPUCancelExecution()
302 << " on device " << device_ordinal; in TPUCancelExecution()
310 LOG(INFO) << "TPUCancelExecution CloseTPUHost on device " << device_ordinal; in TPUCancelExecution()
313 << " on device " << device_ordinal; in TPUCancelExecution()
315 LOG(INFO) << "TPUCancelExecution CloseTPUHost on device " << device_ordinal in TPUCancelExecution()
322 int device_ordinal) { in RegisterCancellation() argument
354 [device_ordinal, env]() { TPUCancelExecution(env, device_ordinal); }); in RegisterCancellation()
360 se::Stream* stream, int device_ordinal, CancellationToken token, in UnregisterCancellation() argument
[all …]
|
| /external/tensorflow/tensorflow/compiler/xla/service/cpu/ |
| D | cpu_runtime.cc | 52 XfeedManager* GetXfeedManager(int device_ordinal) { in GetXfeedManager() argument
57 auto it = managers->find(device_ordinal); in GetXfeedManager()
59 it = managers->emplace(device_ordinal, new XfeedManager()).first; in GetXfeedManager()
235 int device_ordinal = in __xla_cpu_runtime_AcquireInfeedBufferForDequeue() local
236 run_options ? run_options->stream()->parent()->device_ordinal() : 0; in __xla_cpu_runtime_AcquireInfeedBufferForDequeue()
240 << device_ordinal; in __xla_cpu_runtime_AcquireInfeedBufferForDequeue()
243 xla::cpu::runtime::GetXfeedManager(device_ordinal); in __xla_cpu_runtime_AcquireInfeedBufferForDequeue()
259 int device_ordinal = in __xla_cpu_runtime_ReleaseInfeedBufferAfterDequeue() local
260 run_options ? run_options->stream()->parent()->device_ordinal() : 0; in __xla_cpu_runtime_ReleaseInfeedBufferAfterDequeue()
264 << device_ordinal; in __xla_cpu_runtime_ReleaseInfeedBufferAfterDequeue()
[all …]
|
| /external/tensorflow/tensorflow/python/tpu/ops/ |
| D | tpu_ops.py | 286 device_ordinal, argument
310 device_ordinal=device_ordinal,
323 device_ordinal, argument
366 device_ordinal=device_ordinal,
381 device_ordinal, argument
444 device_ordinal=device_ordinal,
461 device_ordinal, argument
524 device_ordinal=device_ordinal,
|
| /external/tensorflow/tensorflow/core/tpu/kernels/ |
| D | host_compute_ops.cc | 39 int device_ordinal = 0; in RecvAtHostOp() local
41 OP_REQUIRES_OK(ctx, ctx->GetAttr("device_ordinal", &device_ordinal)); in RecvAtHostOp()
43 ctx, device_ordinal >= 0, in RecvAtHostOp()
68 parsed_name.id = device_ordinal; in RecvAtHostOp()
84 const int device_ordinal = device_ordinal_tensor.flat<int64>()(0); in ComputeAsync() local
86 ctx, device_ordinal >= 0, in ComputeAsync()
94 parsed_name.id = device_ordinal; in ComputeAsync()
190 int device_ordinal = 0; in SendFromHostOp() local
192 OP_REQUIRES_OK(ctx, ctx->GetAttr("device_ordinal", &device_ordinal)); in SendFromHostOp()
194 ctx, device_ordinal >= 0, in SendFromHostOp()
[all …]
|
| /external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| D | nccl_utils.cc | 104 ncclComm_t NcclClique::GetCommForDeviceOrdinal(int device_ordinal) const { in GetCommForDeviceOrdinal()
105 return comms_by_device_ordinal_.at(device_ordinal).get(); in GetCommForDeviceOrdinal()
144 local_participant.device_ordinal, in LocalParticipantsToString()
184 cudaSetDevice(local_participants[i].device_ordinal)); in CreateNcclClique()
197 int device_ordinal = local_participants[i].device_ordinal; in CreateNcclClique() local
199 device_ordinal, raw_comms[i]); in CreateNcclClique()
201 comms_by_device_ordinal.emplace(device_ordinal, in CreateNcclClique()
270 for (int device_ordinal = 0; device_ordinal < local_devices->size(); in GetLocalParticipants() local
271 ++device_ordinal) { in GetLocalParticipants()
272 auto it = device_ranks.find((*local_devices)[device_ordinal]); in GetLocalParticipants()
[all …]
|
| /external/tensorflow/tensorflow/stream_executor/tpu/ |
| D | tpu_node_context.cc | 30 int device_ordinal) { in Create() argument
33 tpu::OpsApiFn()->TpuNodeContext_CreateFn(device_ordinal, status.c_status); in Create()
40 return std::make_unique<TpuNodeContext>(device_ordinal, node_context); in Create()
62 Status TpuNodeContext::Initialize(int device_ordinal) { in Initialize() argument
64 tpu::OpsApiFn()->TpuNodeContext_InitializeFn(device_ordinal, status.c_status); in Initialize()
73 int TpuNodeContext::device_ordinal() const { return device_ordinal_; } in device_ordinal() function in tensorflow::tpu::TpuNodeContext
|
| D | tpu_node_context.h | 48 static StatusOr<std::unique_ptr<TpuNodeContext>> Create(int device_ordinal);
50 explicit TpuNodeContext(int device_ordinal, XLA_TpuNodeContext* node_context) in TpuNodeContext() argument
51 : device_ordinal_(device_ordinal), node_context_(node_context) { in TpuNodeContext()
60 static Status Initialize(int device_ordinal);
64 int device_ordinal() const;
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
| D | tpu_extract_outside_compilation.cc | 94 Value device_ordinal, in CreateSendFromHostOp() argument
96 if (device_ordinal) in CreateSendFromHostOp()
99 /*dynamic_key=*/compilation_key, device_ordinal, in CreateSendFromHostOp()
112 Value device_ordinal, in CreateRecvAtHostOp() argument
114 if (device_ordinal) in CreateRecvAtHostOp()
116 loc, output_types, /*dynamic_key=*/compilation_key, device_ordinal, in CreateRecvAtHostOp()
267 Value device_ordinal, int& communication_key_index) { in MoveOpsToHost() argument
303 device_ordinal, args_communication_key); in MoveOpsToHost()
309 compilation_key, device_ordinal, in MoveOpsToHost()
358 Value compilation_key, Value device_ordinal, in DecomposeControlFlow() argument
[all …]
|
| /external/tensorflow/tensorflow/compiler/xrt/kernels/ |
| D | tpu_execute_op.cc | 141 xla::ScopedShapedBuffer output_scoped_buffer, int device_ordinal) { in AllocateOutputTuple() argument
157 node_context->backend(), device_ordinal, &output_tuple)); in AllocateOutputTuple()
177 xla::ScopedShapedBuffer output_scoped_buffer, int device_ordinal) { in AllocateOutputTensors() argument
182 device_ordinal)); in AllocateOutputTensors()
193 int device_ordinal, string rendezvous_key_base) { in RunExecutable() argument
211 backend->BorrowStream(device_ordinal)); in RunExecutable()
235 int device_ordinal, string rendezvous_key_base) { in ExecuteTPUProgram() argument
241 device_ordinal, rendezvous_key_base); in ExecuteTPUProgram()
244 runfn, backend, device_ordinal, /*requested_free_size=*/0); in ExecuteTPUProgram()
273 const int device_ordinal = metadata->device_ordinal(); in DoWork() local
[all …]
|