Searched refs:TrackedDeviceBuffer (Results 1 – 5 of 5) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/pjrt/ |
| D | tracked_device_buffer.cc | 96 /* static */ std::shared_ptr<TrackedDeviceBuffer>
97 TrackedDeviceBuffer::FromScopedShapedBuffer( in FromScopedShapedBuffer()
114 return std::make_shared<TrackedDeviceBuffer>( in FromScopedShapedBuffer()
120 ShapedBuffer TrackedDeviceBuffer::AsShapedBuffer( in AsShapedBuffer()
137 void TrackedDeviceBuffer::AddToInputAsImmutable( in AddToInputAsImmutable()
148 void TrackedDeviceBuffer::AddToInputAsDonated( in AddToInputAsDonated()
163 TrackedDeviceBuffer::TrackedDeviceBuffer( in TrackedDeviceBuffer() function in xla::TrackedDeviceBuffer
176 TrackedDeviceBuffer::~TrackedDeviceBuffer() { in ~TrackedDeviceBuffer()
190 void TrackedDeviceBuffer::AddUsageEvent( in AddUsageEvent()
207 TrackedDeviceBuffer::StreamAndEventContainer
[all …]
|
| D | tracked_device_buffer.h | 119 class TrackedDeviceBuffer {
135 static std::shared_ptr<TrackedDeviceBuffer> FromScopedShapedBuffer(
206 TrackedDeviceBuffer() : in_use_(true) {} in TrackedDeviceBuffer() function
207 TrackedDeviceBuffer(se::DeviceMemoryAllocator* allocator, int device_ordinal,
212 ~TrackedDeviceBuffer();
246 void GetDeviceBufferEvents(const TrackedDeviceBuffer& buffer,
251 void WaitForBufferDefinitionEventsOnStream(const TrackedDeviceBuffer& buffer,
|
| D | pjrt_stream_executor_client.h | 378 const std::shared_ptr<TrackedDeviceBuffer>& buffer() const { in buffer()
383 TrackedDeviceBuffer* operator->() const { return buffer().get(); }
384 const TrackedDeviceBuffer& operator*() const { return *buffer(); }
424 Status, std::shared_ptr<TrackedDeviceBuffer>>;
442 void Acquire(StatusOr<std::shared_ptr<TrackedDeviceBuffer>>&& buffer_or);
456 std::shared_ptr<TrackedDeviceBuffer> buffer_;
460 std::shared_ptr<TrackedDeviceBuffer> device_buffer,
541 StatusOr<std::shared_ptr<TrackedDeviceBuffer>> Release(
556 StatusOr<std::shared_ptr<TrackedDeviceBuffer>> GetBufferForHoldLocked(
567 void ConvertUsageHold(TrackedDeviceBuffer* buffer, se::Stream* usage_stream,
[all …]
|
| D | tracked_device_buffer_test.cc | 30 StatusOr<std::shared_ptr<TrackedDeviceBuffer>> MakeArray(const Shape& shape, in MakeArray()
45 return std::make_shared<TrackedDeviceBuffer>( in MakeArray()
105 std::shared_ptr<TrackedDeviceBuffer> device_buffer = in TEST()
106 TrackedDeviceBuffer::FromScopedShapedBuffer(&shaped_buffer, {}); in TEST()
|
| D | pjrt_stream_executor_client.cc | 452 std::shared_ptr<TrackedDeviceBuffer> dst_device_buffer = in AllocateDestinationBuffer()
453 TrackedDeviceBuffer::FromScopedShapedBuffer(&dst_buffer, in AllocateDestinationBuffer()
523 StatusOr<std::shared_ptr<TrackedDeviceBuffer>>&& buffer_or) { in Acquire()
616 std::shared_ptr<TrackedDeviceBuffer> tracked_device_buffer) in TrackedDeviceBufferExternalReference()
624 std::shared_ptr<TrackedDeviceBuffer> tracked_device_buffer_;
635 std::shared_ptr<TrackedDeviceBuffer> tracked_device_buffer, in ReleaseDeviceMemoryOwnership()
710 auto device_buffer = std::make_shared<TrackedDeviceBuffer>( in BufferFromHostBuffer()
963 auto device_buffer = std::make_shared<TrackedDeviceBuffer>( in CreateViewOfDeviceBuffer()
998 Shape on_device_shape, std::shared_ptr<TrackedDeviceBuffer> device_buffer, in PjRtStreamExecutorBuffer()
1039 StatusOr<std::shared_ptr<TrackedDeviceBuffer>>
[all …]
|