Searched refs:PjRtStreamExecutorBuffer (Results 1 – 2 of 2) sorted by relevance
328 void RecordUsage(PjRtStreamExecutorBuffer::ScopedHold device_buffer, in RecordUsage()363 StatusOr<std::unique_ptr<PjRtStreamExecutorBuffer>> AllocateDestinationBuffer( in AllocateDestinationBuffer()456 auto py_buffer = absl::make_unique<PjRtStreamExecutorBuffer>( in AllocateDestinationBuffer()481 PjRtStreamExecutorBuffer::ScopedHold device_buffer, in AddDestinationBufferSynchronization()506 PjRtStreamExecutorBuffer::ScopedHold::~ScopedHold() { in ~ScopedHold()512 PjRtStreamExecutorBuffer::ScopedHold::ScopedHold(ScopedHold&& other) in ScopedHold()522 void PjRtStreamExecutorBuffer::ScopedHold::Acquire( in Acquire()537 PjRtStreamExecutorBuffer::ScopedHold::ForClosure538 PjRtStreamExecutorBuffer::ScopedHold::ToClosure() { in ToClosure()546 void PjRtStreamExecutorBuffer::ScopedHold::ConvertUsageHold( in ConvertUsageHold()[all …]
235 friend class PjRtStreamExecutorBuffer;283 class PjRtStreamExecutorBuffer : public PjRtBuffer {418 friend class PjRtStreamExecutorBuffer;423 using ForClosure = std::tuple<PjRtStreamExecutorBuffer*, Type, State,426 ScopedHold(PjRtStreamExecutorBuffer* parent, Type type) in ScopedHold()449 PjRtStreamExecutorBuffer* const parent_;459 PjRtStreamExecutorBuffer(Shape on_device_shape,462 ~PjRtStreamExecutorBuffer() override;464 PjRtStreamExecutorBuffer(const PjRtStreamExecutorBuffer&) = delete;465 PjRtStreamExecutorBuffer(PjRtStreamExecutorBuffer&&) = delete;[all …]