#ifndef CAFFE2_OPERATORS_PREFETCH_OP_H_
|
#define CAFFE2_OPERATORS_PREFETCH_OP_H_
|
|
#include <condition_variable>
|
#include <mutex>
|
#include <thread> // NOLINT
|
|
#include "caffe2/core/context.h"
|
#include "caffe2/core/operator.h"
|
|
namespace caffe2 {
|
|
// PrefetchOperator is an operator that prefetches the next batch. It should
|
// almost always be used to read things from disk, so I am setting the input to
|
// zero blobs.
|
//
|
// For any operator that is derived from PrefetchOperator, it should
|
// explicitly call the Finalize() function in its destructor, so that the
|
// prefetching thread is properly destructed.
|
|
// Note: We inherit from OperatorBase since we control the
|
// synchronization properties of this operator ourselves (we inform
|
// the waiting producer after we synchronize). This is a special-case
|
// - you should generally inherit from Operator<Context> directly.
|
template <class Context>
|
class PrefetchOperator : public OperatorBase {
|
public:
|
PrefetchOperator(const OperatorDef& operator_def, Workspace* ws)
|
: OperatorBase(operator_def, ws),
|
context_(operator_def.device_option()),
|
prefetched_(false),
|
prefetch_success_(true),
|
finalize_(false),
|
no_prefetch_(GetSingleArgument<bool>("no_prefetch", false)) {
|
context_.SwitchToDevice();
|
}
|
|
virtual ~PrefetchOperator() noexcept {
|
CHECK(finalize_ || !prefetch_thread_.get()) <<
|
"YOU MADE A PROGRAMING ERROR: derived class of PrefetchOperator "
|
"should call Finalize() in its destructor so the prefetching "
|
"thread is joined. ";
|
}
|
|
void Finalize() {
|
if (prefetch_thread_.get()) {
|
{
|
std::unique_lock<std::mutex> lock(prefetch_access_mutex_);
|
while (!prefetched_)
|
consumer_.wait(lock);
|
finalize_ = true;
|
prefetched_ = false;
|
}
|
producer_.notify_one();
|
prefetch_thread_->join();
|
prefetch_thread_.reset();
|
} else {
|
// If we never initialized the prefetch thread, just set
|
// finalize anyway.
|
finalize_ = true;
|
}
|
}
|
|
bool Run(int /* unused */ /*stream_id*/) override {
|
if (no_prefetch_) {
|
context_.SwitchToDevice();
|
bool result = Prefetch() && CopyPrefetched();
|
context_.FinishDeviceComputation();
|
return result;
|
}
|
// Note(jiayq): We only start the prefetch_thread at the Run() function
|
// instead of in the constructor, because the prefetch_thread needs to start
|
// after all derived classes' constructors finish.
|
if (!prefetch_thread_) {
|
prefetch_thread_.reset(
|
new std::thread([this] { this->PrefetchWorker(); }));
|
}
|
context_.SwitchToDevice();
|
std::unique_lock<std::mutex> lock(prefetch_access_mutex_);
|
while (!prefetched_)
|
consumer_.wait(lock);
|
if (!prefetch_success_) {
|
LOG(ERROR) << "Prefetching failed.";
|
return false;
|
}
|
if (!CopyPrefetched()) {
|
LOG(ERROR) << "Error when copying prefetched data.";
|
return false;
|
}
|
prefetched_ = false;
|
context_.FinishDeviceComputation();
|
producer_.notify_one();
|
return true;
|
}
|
|
void PrefetchWorker() {
|
context_.SwitchToDevice();
|
std::unique_lock<std::mutex> lock(prefetch_access_mutex_);
|
while (prefetched_)
|
producer_.wait(lock);
|
while (!finalize_) {
|
// We will need to run a FinishDeviceComputation() call because the
|
// prefetcher thread and the main thread are potentially using different
|
// streams (like on GPU).
|
try {
|
prefetch_success_ = Prefetch();
|
context_.FinishDeviceComputation();
|
} catch (const std::exception& e) {
|
// TODO: propagate exception_ptr to the caller side
|
LOG(ERROR) << "Prefetching error " << e.what();
|
prefetch_success_ = false;
|
}
|
prefetched_ = true;
|
consumer_.notify_one();
|
while (prefetched_)
|
producer_.wait(lock);
|
}
|
}
|
|
// You will need to implement this instead of the Run function.
|
virtual bool Prefetch() = 0;
|
virtual bool CopyPrefetched() = 0;
|
|
protected:
|
Context context_;
|
std::mutex prefetch_access_mutex_;
|
std::condition_variable producer_, consumer_;
|
// prefetched_ is used to tell the operator that it is done.
|
std::atomic<bool> prefetched_;
|
// prefetch_success_ is used to see if prefetching failed or not.
|
std::atomic<bool> prefetch_success_;
|
// finalize_ is used to tell the prefetcher to quit.
|
std::atomic<bool> finalize_;
|
unique_ptr<std::thread> prefetch_thread_;
|
|
// Whether to do prefetching or run this as a normal operator
|
const bool no_prefetch_;
|
};
|
|
} // namespace caffe2
|
|
#endif // CAFFE2_OPERATORS_PREFETCH_OP_H_
|
