#pragma once
|
|
#include <torch/types.h>
|
|
#include <c10/util/Exception.h>
|
|
#include <chrono>
|
#include <condition_variable>
|
#include <cstddef>
|
#include <mutex>
|
#include <queue>
|
|
namespace torch {
|
namespace data {
|
namespace detail {
|
|
/// A basic locked, blocking MPMC queue.
|
///
|
/// Every `push` and `pop` is guarded by a mutex. A condition variable is used
|
/// to communicate insertion of new elements, such that waiting threads will be
|
/// woken up if they are currently waiting inside a call to `pop()`.
|
///
|
/// Note that this data structure is written specifically for use with the
|
/// `DataLoader`. Its behavior is tailored to this use case and may not be
|
/// applicable to more general uses.
|
template <typename T>
|
class Queue {
|
public:
|
/// Pushes a new value to the back of the `Queue` and notifies one thread on
|
/// the waiting side about this event.
|
void push(T value) {
|
{
|
std::lock_guard<std::mutex> lock(mutex_);
|
queue_.push(std::move(value));
|
}
|
cv_.notify_one();
|
}
|
|
/// Blocks until at least one element is ready to be popped from the front of
|
/// the queue. An optional `timeout` in seconds can be used to limit the time
|
/// spent waiting for an element. If the wait times out, an exception is
|
/// raised.
|
T pop(optional<std::chrono::milliseconds> timeout = nullopt) {
|
std::unique_lock<std::mutex> lock(mutex_);
|
if (timeout) {
|
if (!cv_.wait_for(
|
lock, *timeout, [this] { return !this->queue_.empty(); })) {
|
// clang-format off
|
AT_ERROR(
|
"Timeout in DataLoader queue while waiting for next batch"
|
" (timeout was ", timeout->count(), " ms)");
|
// clang-format on
|
}
|
} else {
|
cv_.wait(lock, [this] { return !this->queue_.empty(); });
|
}
|
AT_ASSERT(!queue_.empty());
|
T value = queue_.front();
|
queue_.pop();
|
lock.unlock();
|
return value;
|
}
|
|
/// Empties the queue and returns the number of elements that were present at
|
/// the start of the function. No threads are notified about this event as it
|
/// is assumed to be used to drain the queue during shutdown of a
|
/// `DataLoader`.
|
size_t clear() {
|
std::lock_guard<std::mutex> lock(this->mutex_);
|
const auto size = queue_.size();
|
while (!queue_.empty()) {
|
queue_.pop();
|
}
|
return size;
|
}
|
|
private:
|
std::queue<T> queue_;
|
std::mutex mutex_;
|
std::condition_variable cv_;
|
};
|
} // namespace detail
|
} // namespace data
|
} // namespace torch
|
