ParallelGuard.h 373 B

1234567891011121314151617181920
  1. #pragma once
  2. #include <c10/macros/Macros.h>
  3. namespace c10 {
  4. // RAII thread local guard that tracks whether code is being executed in
  5. // `at::parallel_for` or `at::parallel_reduce` loop function.
  6. class C10_API ParallelGuard {
  7. public:
  8. static bool is_enabled();
  9. ParallelGuard(bool state);
  10. ~ParallelGuard();
  11. private:
  12. bool previous_state_;
  13. };
  14. } // namespace c10