Home
last modified time | relevance | path

Searched refs:on_device_shape (Results 1 – 22 of 22) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/
Dtransfer_manager.cc158 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 …]
Dshaped_buffer.cc34 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()
Dallocation_tracker.cc81 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()
Dshaped_buffer.h45 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,
Dgeneric_transfer_manager.cc74 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()
Dshaped_buffer_test.cc172 EXPECT_EQ(ssb.on_device_shape(), array_shape); in TEST()
Dexecutable.cc49 ShapeTree<MaybeOwningDeviceMemory> result(shaped_buffer.on_device_shape()); in MakeMaybeOwningDeviceMemoryTree()
/external/tensorflow/tensorflow/compiler/xrt/
Dxrt_state.cc87 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 …]
Dxrt_state.h106 const xla::Shape& on_device_shape,
198 const xla::Shape& on_device_shape() const;
239 const xla::Shape& on_device_shape);
Dxrt_util.cc213 xla::ShapeUtil::TupleElementCount(output_tuple->on_device_shape()); in CreateExecuteOutput()
/external/tensorflow/tensorflow/compiler/xla/python/
Dshared_device_buffer.cc59 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()
Dlocal_client.cc340 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()
Dshared_device_buffer.h107 Shape on_device_shape, TransferManager* transfer_manager,
115 const Shape& on_device_shape,
Dlocal_client.h214 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
Dxla.cc186 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()
Dshared_device_buffer_test.cc109 EXPECT_EQ(shaped_buffer.on_device_shape(), abc_tuple_device_shape); in TEST()
Ddlpack.cc254 if (buffer->on_device_shape().IsTuple()) { in BufferToDLPackManagedTensor()
/external/tensorflow/tensorflow/compiler/jit/
Dxla_tensor.cc50 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()
Dxla_launch_util.cc228 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()
Dxla_device.cc130 *shape = shaped_buffer.on_device_shape(); in DefaultPaddedShapeFn()
/external/tensorflow/tensorflow/compiler/xla/service/interpreter/
Dexecutable.cc90 const auto& actual_shape = argument_buffers[i].on_device_shape(); in ExecuteAsyncOnStream()
/external/tensorflow/tensorflow/compiler/xla/tests/
Dlocal_client_execute_test.cc135 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()