mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-14 20:57:59 +00:00
71ca600af9
Summary: Since caffe2 and torch have been consolidated, CAFFE2_API should be merged with TORCH_API. Addresses a TODO. Manually edited some references of the removed `CAFFE2_API`: * `CONTRIBUTING.md` * `caffe2/proto/CMakeLists.txt` * `cmake/ProtoBuf.cmake` * `c10/macros/Export.h` * `torch/csrc/WindowsTorchApiMacro.h` Pull Request resolved: https://github.com/pytorch/pytorch/pull/49496 Reviewed By: malfet, samestep Differential Revision: D25600726 Pulled By: janeyx99 fbshipit-source-id: 7e068d959e397ac183c097d7e9a9afeca5ddd782
42 lines
993 B
C++
42 lines
993 B
C++
#ifndef CAFFE2_CORE_NET_ASYNC_SCHEDULING_H_
|
|
#define CAFFE2_CORE_NET_ASYNC_SCHEDULING_H_
|
|
|
|
#include "caffe2/core/net_async_base.h"
|
|
|
|
namespace caffe2 {
|
|
|
|
class TORCH_API AsyncSchedulingNet : public AsyncNetBase {
|
|
public:
|
|
AsyncSchedulingNet(
|
|
const std::shared_ptr<const NetDef>& net_def,
|
|
Workspace* ws);
|
|
~AsyncSchedulingNet() override;
|
|
|
|
void Wait() override;
|
|
|
|
void Cancel() override;
|
|
|
|
protected:
|
|
bool RunAsync() override;
|
|
|
|
void pollAndSchedule(int task_id);
|
|
void schedule(int task_id, bool run_inline = false) noexcept;
|
|
void reset() override;
|
|
virtual void finishRun();
|
|
void parentCallback(int parent_id);
|
|
bool isInlineTask(int parent_id, int child_id) const;
|
|
|
|
void CancelAndFinishAsyncTasks();
|
|
|
|
std::mutex running_mutex_;
|
|
std::condition_variable running_cv_;
|
|
std::atomic<bool> running_;
|
|
|
|
std::atomic<int> processed_tasks_num_;
|
|
|
|
C10_DISABLE_COPY_AND_ASSIGN(AsyncSchedulingNet);
|
|
};
|
|
|
|
} // namespace caffe2
|
|
|
|
#endif // CAFFE2_CORE_NET_ASYNC_SCHEDULING_H_
|