mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
6302e4001a
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/11862 Reviewed By: Yangqing Differential Revision: D9942428 fbshipit-source-id: dea03f5ba0e621a047aa50bc4aa97acc834d2a39
48 lines
1.2 KiB
C++
48 lines
1.2 KiB
C++
#include "caffe2/core/event.h"
|
|
|
|
#include <atomic>
|
|
#include <condition_variable>
|
|
|
|
namespace caffe2 {
|
|
|
|
struct CPUEventWrapper {
|
|
explicit CPUEventWrapper(const DeviceOption& option)
|
|
: status_(EventStatus::EVENT_INITIALIZED) {
|
|
CAFFE_ENFORCE(
|
|
option.device_type() == PROTO_CPU ||
|
|
option.device_type() == PROTO_MKLDNN ||
|
|
option.device_type() == PROTO_IDEEP,
|
|
"Expected CPU/MKLDNN/IDEEP device type");
|
|
}
|
|
~CPUEventWrapper() {}
|
|
|
|
std::mutex mutex_;
|
|
std::condition_variable cv_completed_;
|
|
std::atomic<int> status_;
|
|
std::string err_msg_;
|
|
std::vector<EventCallbackFunction> callbacks_;
|
|
};
|
|
|
|
void EventCreateCPU(const DeviceOption& option, Event* event);
|
|
|
|
void EventRecordCPU(
|
|
Event* event,
|
|
const void* /* unused */,
|
|
const char* err_msg);
|
|
|
|
void EventFinishCPU(const Event* event);
|
|
|
|
void EventWaitCPUCPU(const Event* event, void* /* context */);
|
|
|
|
EventStatus EventQueryCPU(const Event* event);
|
|
|
|
const std::string& EventErrorMessageCPU(const Event* event);
|
|
|
|
void EventSetFinishedCPU(const Event* event, const char* err_msg);
|
|
|
|
bool EventCanScheduleCPU(const Event*, const Event*);
|
|
|
|
void EventResetCPU(Event*);
|
|
|
|
} // namespace caffe2
|