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