1 | /** |
---|---|
2 | * This file is adapted from PyTorch/XLA |
3 | * https://github.com/pytorch/xla/blob/master/third_party/xla_client/metrics.h |
4 | */ |
5 | |
6 | #pragma once |
7 | |
8 | #include <functional> |
9 | #include <memory> |
10 | #include <thread> |
11 | |
12 | #include <c10/macros/Export.h> |
13 | |
14 | namespace torch { |
15 | namespace lazy { |
16 | |
17 | class TORCH_API Completion { |
18 | public: |
19 | class Data; |
20 | |
21 | explicit Completion(std::shared_ptr<Data> data); |
22 | |
23 | ~Completion(); |
24 | |
25 | void Wait(); |
26 | |
27 | private: |
28 | std::shared_ptr<Data> data_; |
29 | }; |
30 | |
31 | // Schedules a closure which might wait for IO or other events/conditions. |
32 | TORCH_API void ScheduleIoClosure(std::function<void()> closure); |
33 | TORCH_API Completion |
34 | ScheduleIoClosureWithCompletion(std::function<void()> closure); |
35 | |
36 | } // namespace lazy |
37 | } // namespace torch |
38 |