/external/tensorflow/tensorflow/compiler/xla/service/ |
D | transfer_manager.cc | 158 const Shape on_device_shape = HostShapeToDeviceShape(literal.shape()); in TransferArrayToDeviceAsync() local 159 TF_RET_CHECK(on_device_shape.IsArray()) in TransferArrayToDeviceAsync() 162 << " is not an array: " << ShapeUtil::HumanString(on_device_shape); in TransferArrayToDeviceAsync() 163 if (dest.size() < GetByteSizeRequirement(on_device_shape)) { in TransferArrayToDeviceAsync() 167 dest.size(), GetByteSizeRequirement(on_device_shape)); in TransferArrayToDeviceAsync() 169 ShapedBuffer shaped_buffer(/*on_host_shape=*/literal.shape(), on_device_shape, in TransferArrayToDeviceAsync() 245 device_buffer.on_device_shape(), in WriteTupleIndexTablesAsync() 271 TF_RET_CHECK(device_buffer.on_device_shape().IsTuple()); in WriteRootTupleIndexTable() 272 if (ShapeUtil::TupleElementCount(device_buffer.on_device_shape()) == 0) { in WriteRootTupleIndexTable() 276 TF_RET_CHECK(GetByteSizeRequirement(device_buffer.on_device_shape()) == in WriteRootTupleIndexTable() [all …]
|
D | shaped_buffer.cc | 34 ShapedBuffer::ShapedBuffer(Shape on_host_shape, Shape on_device_shape, in ShapedBuffer() argument 37 on_device_shape_(std::move(on_device_shape)), in ShapedBuffer() 74 ShapeUtil::TryGetSubshape(on_device_shape(), index)); in SubShapedBuffer() 95 ShapeUtil::HumanStringWithLayout(on_device_shape()), in ToString() 98 on_device_shape(), in ToString() 120 Shape on_device_shape, in ScopedShapedBuffer() argument 123 : ShapedBuffer(std::move(on_host_shape), std::move(on_device_shape), in ScopedShapedBuffer() 177 xla::ShapeUtil::GetSubshape(on_device_shape(), {index}); in TakeSubTree()
|
D | allocation_tracker.cc | 81 shaped_buffer.on_device_shape(), in RegisterInternal() 114 shaped_buffer->on_device_shape(), in Unregister() 152 TF_RET_CHECK(shaped_buffer->on_device_shape().IsTuple()); in DeconstructTuple() 154 if (ShapeUtil::IsNestedTuple(shaped_buffer->on_device_shape())) { in DeconstructTuple() 160 i < ShapeUtil::TupleElementCount(shaped_buffer->on_device_shape()); in DeconstructTuple() 164 ShapeUtil::GetTupleElementShape(shaped_buffer->on_device_shape(), i), in DeconstructTuple()
|
D | shaped_buffer.h | 45 ShapedBuffer(Shape on_host_shape, Shape on_device_shape, 66 const Shape& on_device_shape() const { return on_device_shape_; } in on_device_shape() function 139 explicit ScopedShapedBuffer(Shape on_host_shape, Shape on_device_shape,
|
D | generic_transfer_manager.cc | 74 TF_RET_CHECK(ShapeUtil::Equal(device_buffer.on_device_shape(), in TransferLiteralFromDevice() 108 TF_RET_CHECK(ShapeUtil::Equal(device_buffer.on_device_shape(), in TransferLiteralToDeviceAsync()
|
D | shaped_buffer_test.cc | 172 EXPECT_EQ(ssb.on_device_shape(), array_shape); in TEST()
|
D | executable.cc | 49 ShapeTree<MaybeOwningDeviceMemory> result(shaped_buffer.on_device_shape()); in MakeMaybeOwningDeviceMemoryTree()
|
/external/tensorflow/tensorflow/compiler/xrt/ |
D | xrt_state.cc | 87 xla::Shape on_device_shape = transfer_manager->HostShapeToDeviceShape(shape); in AllocateScopedShapedBuffer() local 90 << xla::ShapeUtil::HumanStringWithLayout(on_device_shape); in AllocateScopedShapedBuffer() 96 shape, on_device_shape, backend->memory_allocator(), device_ordinal); in AllocateScopedShapedBuffer() 99 xla::ShapeUtil::GetSubshape(on_device_shape, index_to_buffer.first); in AllocateScopedShapedBuffer() 150 const xla::Shape& on_device_shape) in XRTTupleAllocation() argument 154 on_device_shape_(on_device_shape), in XRTTupleAllocation() 189 shaped_buffer.on_host_shape(), shaped_buffer.on_device_shape()); in CreateAndTransfer() 212 shaped_buffer.on_host_shape(), shaped_buffer.on_device_shape()); in CreateUninitialized() 222 const xla::Shape& on_device_shape, xla::Backend* backend, in CreateFromBuffer() argument 227 on_device_shape); in CreateFromBuffer() [all …]
|
D | xrt_state.h | 106 const xla::Shape& on_device_shape, 198 const xla::Shape& on_device_shape() const; 239 const xla::Shape& on_device_shape);
|
D | xrt_util.cc | 213 xla::ShapeUtil::TupleElementCount(output_tuple->on_device_shape()); in CreateExecuteOutput()
|
/external/tensorflow/tensorflow/compiler/xla/python/ |
D | shared_device_buffer.cc | 59 const Shape& on_host_shape, const Shape& on_device_shape, in BufferFromScopedShapedBufferIterator() argument 76 int num_children = ShapeUtil::TupleElementCount(on_device_shape); in BufferFromScopedShapedBufferIterator() 80 on_host_shape.tuple_shapes(i), on_device_shape.tuple_shapes(i), in BufferFromScopedShapedBufferIterator() 87 on_device_shape, in BufferFromScopedShapedBufferIterator() 102 shaped_buffer->on_host_shape(), shaped_buffer->on_device_shape(), in FromScopedShapedBuffer() 131 Shape on_device_shape, TransferManager* transfer_manager, in MakeArray() argument 136 on_device_shape, [&](const Shape& subshape, const ShapeIndex&) -> Status { in MakeArray() 167 const Shape& on_device_shape, in AsShapedBuffer() argument 169 ShapedBuffer shaped_buffer(on_host_shape, on_device_shape, platform, in AsShapedBuffer()
|
D | local_client.cc | 340 Shape on_device_shape = scoped_buffer.on_device_shape(); in FromHostBuffer() local 343 shape, compact_shape, on_device_shape, data, in FromHostBuffer() 350 compact_shape, on_device_shape, client->client()->platform()); in FromHostBuffer() 398 compact_shape, std::move(on_device_shape), std::move(device_buffer), in FromHostBuffer() 426 device_shapes.push_back(buffer->on_device_shape()); in MakeTuple() 469 PyLocalBuffer::PyLocalBuffer(Shape on_host_shape, Shape on_device_shape, in PyLocalBuffer() argument 475 on_device_shape_(std::move(on_device_shape)), in PyLocalBuffer() 622 if (dst_buffer.on_device_shape().IsTuple()) { in CopyToDevice() 642 dst_buffer.on_host_shape(), dst_buffer.on_device_shape(), in CopyToDevice() 826 result_buffer.on_host_shape(), result_buffer.on_device_shape(), in ExecuteHelper()
|
D | shared_device_buffer.h | 107 Shape on_device_shape, TransferManager* transfer_manager, 115 const Shape& on_device_shape,
|
D | local_client.h | 214 PyLocalBuffer(Shape on_host_shape, Shape on_device_shape, 225 const Shape& on_device_shape() const { return on_device_shape_; } in on_device_shape() function
|
D | xla.cc | 186 if (!buffer.on_device_shape().IsArray()) { in PyLocalBufferGetBuffer() 278 if (!buffer.on_device_shape().IsArray()) { in PyLocalBufferCudaArrayInterface() 677 buffer->on_device_shape().IsArray() && in PYBIND11_MODULE() 678 buffer->on_device_shape().element_type() != BF16) { in PYBIND11_MODULE() 703 if (shaped_buffer.on_device_shape().IsTuple()) { in PYBIND11_MODULE()
|
D | shared_device_buffer_test.cc | 109 EXPECT_EQ(shaped_buffer.on_device_shape(), abc_tuple_device_shape); in TEST()
|
D | dlpack.cc | 254 if (buffer->on_device_shape().IsTuple()) { in BufferToDLPackManagedTensor()
|
/external/tensorflow/tensorflow/compiler/jit/ |
D | xla_tensor.cc | 50 xla::Shape on_device_shape = in AllocateShapedBuffer() local 54 xla::ScopedShapedBuffer shaped_buffer(on_host_shape, on_device_shape, in AllocateShapedBuffer() 59 xla::ShapeUtil::GetSubshape(on_device_shape, index_to_buffer.first); in AllocateShapedBuffer()
|
D | xla_launch_util.cc | 228 const xla::Shape on_device_shape = in PopulateInputs() local 230 if (on_device_shape.IsTuple()) { in PopulateInputs() 236 on_device_shape)) in PopulateInputs() 238 << xla::ShapeUtil::HumanStringWithLayout(on_device_shape) in PopulateInputs() 387 << output.on_device_shape().DebugString(); in PopulateOutputs() 398 xla::ShapeUtil::MakeTupleShape({nontuple_buffer.on_device_shape()}), in PopulateOutputs()
|
D | xla_device.cc | 130 *shape = shaped_buffer.on_device_shape(); in DefaultPaddedShapeFn()
|
/external/tensorflow/tensorflow/compiler/xla/service/interpreter/ |
D | executable.cc | 90 const auto& actual_shape = argument_buffers[i].on_device_shape(); in ExecuteAsyncOnStream()
|
/external/tensorflow/tensorflow/compiler/xla/tests/ |
D | local_client_execute_test.cc | 135 x_array.on_device_shape().layout(), LayoutUtil::MakeLayout({0, 1}))); in XLA_TEST_F() 141 y_array.on_device_shape().layout(), LayoutUtil::MakeLayout({1, 0}))); in XLA_TEST_F() 176 result_colmaj.on_device_shape().layout(), in XLA_TEST_F() 189 result_rowmaj.on_device_shape().layout(), in XLA_TEST_F()
|