Home
last modified time | relevance | path

Searched refs:shaped_buffer (Results 1 – 25 of 33) sorted by relevance

12

/external/tensorflow/tensorflow/compiler/xla/service/
Dallocation_tracker.cc34 ScopedShapedBuffer shaped_buffer, const string& tag) { in Register() argument
38 replicated_buffers.emplace_back(std::move(shaped_buffer)); in Register()
66 for (const auto& shaped_buffer : replicated_buffers) { in RegisterInternal() local
67 VLOG(2) << "shaped_buffer:" << shaped_buffer; in RegisterInternal()
68 if (shaped_buffer.platform() != backend_->platform()) { in RegisterInternal()
72 backend_->platform()->Name(), shaped_buffer.platform()->Name()); in RegisterInternal()
77 for (auto& shaped_buffer : replicated_buffers) { in RegisterInternal() local
80 shaped_buffer.on_device_shape(), in RegisterInternal()
87 AddAllocationOrIncrementRefCount(shaped_buffer.buffer(index), in RegisterInternal()
88 shaped_buffer.device_ordinal()); in RegisterInternal()
[all …]
Dtransfer_manager.cc169 ShapedBuffer shaped_buffer(/*on_host_shape=*/literal.shape(), on_device_shape, in TransferArrayToDeviceAsync() local
172 shaped_buffer.set_buffer(dest, /*index=*/{}); in TransferArrayToDeviceAsync()
173 return TransferLiteralToDevice(stream, literal, shaped_buffer, in TransferArrayToDeviceAsync()
193 ShapedBuffer shaped_buffer(/*on_host_shape=*/shape, shape, in TransferArrayFromDevice() local
196 shaped_buffer.set_buffer(source, /*index=*/{}); in TransferArrayFromDevice()
197 return TransferLiteralFromDevice(stream, shaped_buffer, literal, in TransferArrayFromDevice()
320 ScopedShapedBuffer shaped_buffer(on_host_shape, on_device_shape, allocator, in AllocateScopedShapedBuffer() local
325 for (auto& pair : shaped_buffer.buffers()) { in AllocateScopedShapedBuffer()
330 allocator->Allocate(shaped_buffer.device_ordinal(), in AllocateScopedShapedBuffer()
336 return std::move(shaped_buffer); in AllocateScopedShapedBuffer()
Dshaped_buffer.cc114 ScopedShapedBuffer::ScopedShapedBuffer(ShapedBuffer shaped_buffer, in ScopedShapedBuffer() argument
116 : ShapedBuffer(std::move(shaped_buffer)), allocator_(allocator) {} in ScopedShapedBuffer()
137 ShapedBuffer shaped_buffer(static_cast<ShapedBuffer&&>(*this)); in release() local
139 return shaped_buffer; in release()
Dservice.cc246 const ShapedBuffer* shaped_buffer = replicated_buffers[replica]; in ResolveAndValidateArguments() local
249 if (shaped_buffer->platform() != execute_backend_->platform() || in ResolveAndValidateArguments()
250 shaped_buffer->device_ordinal() != replica_device_ordinal) { in ResolveAndValidateArguments()
254 i, shaped_buffer->platform()->Name(), in ResolveAndValidateArguments()
255 shaped_buffer->device_ordinal(), in ResolveAndValidateArguments()
258 replicated_arguments[replica].push_back(shaped_buffer); in ResolveAndValidateArguments()
941 TF_ASSIGN_OR_RETURN(const ShapedBuffer* shaped_buffer, in TransferToClient()
951 return_shape = Shape(shaped_buffer->on_host_shape()); in TransferToClient()
955 shaped_buffer->device_ordinal())); in TransferToClient()
960 stream.get(), *shaped_buffer)); in TransferToClient()
[all …]
Dshaped_buffer_test.cc161 xla::ScopedShapedBuffer shaped_buffer(shape, shape, /*allocator=*/&allocator, in BM_TakeSubTree() local
167 (void)shaped_buffer.TakeSubTree(/*index=*/{fan_out / 2}).release(); in BM_TakeSubTree()
Dallocation_tracker.h49 StatusOr<GlobalDataHandle> Register(ScopedShapedBuffer shaped_buffer,
Dshaped_buffer.h143 explicit ScopedShapedBuffer(ShapedBuffer shaped_buffer,
/external/tensorflow/tensorflow/compiler/xrt/
Dxrt_state.cc200 auto shaped_buffer = scoped_buffer->release(); in CreateAndTransfer() local
202 shaped_buffer.on_host_shape(), in CreateAndTransfer()
203 shaped_buffer.on_device_shape()); in CreateAndTransfer()
205 ->InitializeFromShapedBuffer(shaped_buffer, allocator, device_ordinal); in CreateAndTransfer()
210 const xla::ShapedBuffer& shaped_buffer, xla::Backend* backend, in CreateFromBuffer() argument
215 shaped_buffer.on_host_shape(), in CreateFromBuffer()
216 shaped_buffer.on_device_shape()); in CreateFromBuffer()
218 ->InitializeFromShapedBuffer(shaped_buffer, allocator, device_ordinal); in CreateFromBuffer()
229 xla::ShapedBuffer shaped_buffer = ToShapedBuffer(); in ToLiteral() local
230 for (auto& index_buffer : shaped_buffer.buffers()) { in ToLiteral()
[all …]
Dxrt_state.h96 static Status CreateFromBuffer(const xla::ShapedBuffer& shaped_buffer,
202 void InitializeFromShapedBuffer(const xla::ShapedBuffer& shaped_buffer,
DBUILD70 "//tensorflow/compiler/xla/service:shaped_buffer",
/external/tensorflow/tensorflow/compiler/jit/
Dxla_tensor.cc39 return xla_tensor->shaped_buffer().root_buffer(); in DeviceMemoryFromTensor()
54 xla::ScopedShapedBuffer shaped_buffer(on_host_shape, on_device_shape, in AllocateShapedBuffer() local
57 for (auto& index_to_buffer : shaped_buffer.buffers()) { in AllocateShapedBuffer()
69 VLOG(4) << shaped_buffer.ToString(); in AllocateShapedBuffer()
71 set_shaped_buffer(std::move(shaped_buffer)); in AllocateShapedBuffer()
Dxla_tensor.h63 const xla::ShapedBuffer& shaped_buffer() const { in shaped_buffer() function
67 xla::ShapedBuffer& shaped_buffer() { in shaped_buffer() function
72 void set_shaped_buffer(xla::ScopedShapedBuffer shaped_buffer) { in set_shaped_buffer() argument
74 absl::make_unique<xla::ScopedShapedBuffer>(std::move(shaped_buffer)); in set_shaped_buffer()
Dxla_device_context.cc152 << xla_tensor->shaped_buffer().ToString(); in CopyCPUTensorToDevice()
155 stream_->parent(), xla_tensor->shaped_buffer())) { in CopyCPUTensorToDevice()
162 host_to_device_stream_.get(), literal, xla_tensor->shaped_buffer())); in CopyCPUTensorToDevice()
241 xla_tensor->shaped_buffer().on_host_shape()), in CopyDeviceTensorToCPU()
248 device_to_host_stream.get(), xla_tensor->shaped_buffer(), literal, in CopyDeviceTensorToCPU()
252 << xla_tensor->shaped_buffer().ToString(); in CopyDeviceTensorToCPU()
Dxla_device.cc124 const xla::ShapedBuffer& shaped_buffer = xla_tensor->shaped_buffer(); in DefaultPaddedShapeFn() local
125 *shape = shaped_buffer.on_device_shape(); in DefaultPaddedShapeFn()
/external/tensorflow/tensorflow/compiler/xla/client/
Dlocal_client.cc213 const ShapedBuffer& shaped_buffer) { in LiteralFromShapedBuffer() argument
215 backend_->BorrowStream(shaped_buffer.device_ordinal())); in LiteralFromShapedBuffer()
217 shaped_buffer); in LiteralFromShapedBuffer()
279 const ShapedBuffer& shaped_buffer) { in ShapedBufferToLiteral() argument
281 shaped_buffer.device_ordinal())); in ShapedBufferToLiteral()
283 shaped_buffer); in ShapedBufferToLiteral()
318 ::xla::ScopedShapedBuffer shaped_buffer, in TransferToLocalServer()
324 stream.get(), literal, shaped_buffer)); in TransferToLocalServer()
326 replicated_buffer.emplace_back(std::move(shaped_buffer)); in TransferToLocalServer()
Dlocal_client.h87 StatusOr<Literal> LiteralFromShapedBuffer(const ShapedBuffer& shaped_buffer);
138 StatusOr<Literal> ShapedBufferToLiteral(const ShapedBuffer& shaped_buffer);
/external/tensorflow/tensorflow/compiler/xla/python/
Dlocal_computation_builder.cc119 LocalShapedBuffer::LocalShapedBuffer(ScopedShapedBuffer shaped_buffer, in LocalShapedBuffer() argument
121 : shaped_buffer_(std::move(shaped_buffer)), client_(client) {} in LocalShapedBuffer()
123 const ScopedShapedBuffer* LocalShapedBuffer::shaped_buffer() const { in shaped_buffer() function in xla::swig::LocalShapedBuffer
130 return shaped_buffer()->on_device_shape(); in shape()
134 return client_->ShapedBufferToLiteral(*shaped_buffer()); in ToLiteral()
175 DeviceMemoryAllocator* allocator = shaped_buffer()->memory_allocator(); in DestructureTuple()
188 ShapedBuffer shaped_buffer(subshape, subshape, platform, device_ordinal); in DestructureTuple() local
196 shaped_buffer.set_buffer(*device_memory, index); in DestructureTuple()
202 ScopedShapedBuffer(std::move(shaped_buffer), allocator), client_)); in DestructureTuple()
242 argument_buffers.push_back(handle->shaped_buffer()); in Execute()
[all …]
Dlocal_computation_builder.h81 LocalShapedBuffer(ScopedShapedBuffer shaped_buffer, xla::LocalClient* client);
84 const ScopedShapedBuffer* shaped_buffer() const;
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dgpu_executable.cc290 ScopedShapedBuffer shaped_buffer(root->shape(), root->shape(), in ExecuteOnStream() local
296 TF_RETURN_IF_ERROR(shaped_buffer.buffers().ForEachMutableElementWithStatus( in ExecuteOnStream()
346 return std::move(shaped_buffer); in ExecuteOnStream()
Dbuffer_comparator.cc106 auto shaped_buffer, ([&]() -> StatusOr<ScopedShapedBuffer> { in Create()
121 std::move(shaped_buffer)); in Create()
/external/tensorflow/tensorflow/compiler/xla/tests/
Dlocal_client_test_base.cc140 const ShapedBuffer& shaped_buffer) { in ShapedBufferToLiteral() argument
141 return local_client_->ShapedBufferToLiteral(shaped_buffer) in ShapedBufferToLiteral()
Dlocal_client_test_base.h89 Literal ShapedBufferToLiteral(const ShapedBuffer& shaped_buffer);
Dlocal_client_execute_test.cc786 auto shaped_buffer, in XLA_TEST_F()
791 local_client_->ShapedBufferToLiteral(shaped_buffer)); in XLA_TEST_F()
826 auto shaped_buffer, in XLA_TEST_F()
831 local_client_->ShapedBufferToLiteral(shaped_buffer)); in XLA_TEST_F()
/external/tensorflow/tensorflow/compiler/xrt/kernels/
Dxrt_execute_op.cc226 auto shaped_buffer = scoped_buffer.release(); in DoWork() local
229 shaped_buffer, device_ref.backend(), device_ref.device_ordinal(), in DoWork()
/external/tensorflow/tensorflow/compiler/xla/service/interpreter/
DBUILD92 "//tensorflow/compiler/xla/service:shaped_buffer",

12