Home
last modified time | relevance | path

Searched refs:GetByteSizeRequirement (Results 1 – 14 of 14) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/
Dgeneric_transfer_manager.cc53 stream, GetByteSizeRequirement(shape), element_pointers->data(), region)); in WriteSingleTupleIndexTable()
79 /*size=*/GetByteSizeRequirement(subshape)); in TransferLiteralFromDevice()
113 TF_RET_CHECK(GetByteSizeRequirement(device_subshape) == in TransferLiteralToDeviceAsync()
124 /*size=*/GetByteSizeRequirement(device_subshape), source, in TransferLiteralToDeviceAsync()
133 /*size=*/GetByteSizeRequirement(device_subshape), source, in TransferLiteralToDeviceAsync()
159 int64 GenericTransferManager::GetByteSizeRequirement(const Shape& shape) const { in GetByteSizeRequirement() function in xla::GenericTransferManager
Dtransfer_manager.cc166 if (dest.size() < GetByteSizeRequirement(on_device_shape)) { in TransferArrayToDeviceAsync()
170 dest.size(), GetByteSizeRequirement(on_device_shape)); in TransferArrayToDeviceAsync()
190 if (source.size() < GetByteSizeRequirement(shape)) { in TransferArrayFromDevice()
194 source.size(), GetByteSizeRequirement(shape))); in TransferArrayFromDevice()
303 TF_RET_CHECK(GetByteSizeRequirement(device_subshape) == in WriteTupleIndexTablesAsync()
329 TF_RET_CHECK(GetByteSizeRequirement(device_buffer.on_device_shape()) == in WriteRootTupleIndexTable()
349 TF_RET_CHECK(GetByteSizeRequirement(buffer_tree.shape()) == in WriteRootTupleIndexTable()
410 GetByteSizeRequirement(subshape), in AllocateScopedShapedBuffer()
Dgeneric_transfer_manager.h60 int64 GetByteSizeRequirement(const Shape& shape) const override;
Dtransfer_manager.h235 virtual int64 GetByteSizeRequirement(const Shape& shape) const = 0;
/external/tensorflow/tensorflow/compiler/xla/service/cpu/
Dcpu_transfer_manager.cc100 int64 size = GetByteSizeRequirement(shape); in TransferLiteralToInfeed()
123 int64 tuple_element_size = GetByteSizeRequirement(tuple_element_shape); in TransferLiteralToInfeed()
181 int64 size = GetByteSizeRequirement(literal.shape()); in TransferLiteralFromOutfeed()
195 TF_RET_CHECK(size == GetByteSizeRequirement(received_shape)); in TransferLiteralFromOutfeed()
209 int64 size = GetByteSizeRequirement(tuple_element_shape); in TransferLiteralFromOutfeed()
221 TF_RET_CHECK(GetByteSizeRequirement(literal.shape()) == in TransferLiteralFromOutfeed()
222 GetByteSizeRequirement(received_shape)); in TransferLiteralFromOutfeed()
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dgpu_transfer_manager.cc61 int64 tuple_element_size = GetByteSizeRequirement(sub_shape); in TransferLiteralToInfeed()
125 absl::make_unique<gpu::OutfeedBuffer>(GetByteSizeRequirement(shape)); in TransferLiteralFromOutfeed()
/external/tensorflow/tensorflow/stream_executor/tpu/
Dtpu_transfer_manager.h68 int64 GetByteSizeRequirement(const xla::Shape& shape) const override;
Dtpu_transfer_manager.cc211 int64 TpuTransferManager::GetByteSizeRequirement( in GetByteSizeRequirement() function in tensorflow::tpu::TpuTransferManager
/external/tensorflow/tensorflow/compiler/jit/
Dxla_tensor.cc57 client->backend().transfer_manager()->GetByteSizeRequirement(subshape); in AllocateShapedBuffer()
/external/tensorflow/tensorflow/compiler/xla/pjrt/
Dtracked_device_buffer_test.cc40 client->backend().transfer_manager()->GetByteSizeRequirement( in MakeArray()
Dpjrt_stream_executor_client.cc1022 ->GetByteSizeRequirement(on_device_shape_); in OnDeviceSizeInBytes()
1506 transfer_manager->GetByteSizeRequirement(on_device_shape))); in MakeTupleHelper()
/external/tensorflow/tensorflow/compiler/xrt/
Dxrt_state.cc100 uint64 size = transfer_manager->GetByteSizeRequirement(subshape); in AllocateScopedShapedBuffer()
487 uint64 size = transfer_manager->GetByteSizeRequirement(subshape); in MakeTuple()
/external/tensorflow/tensorflow/compiler/xla/tests/
Dtransfer_manager_test.cc45 return transfer_manager_->GetByteSizeRequirement(shape); in TransferManagerTest()
/external/tensorflow/tensorflow/core/tpu/kernels/
Dtpu_execute_op.cc286 transfer_manager->GetByteSizeRequirement(input_buffers->buffers.shape()); in BuildComputationInputs()