Unroll.h 1.1 KB

1234567891011121314151617181920212223242526272829303132333435
  1. #if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
  2. #pragma once
  3. #include <c10/macros/Macros.h>
  4. #include <type_traits>
  5. // Utility to guarantee complete unrolling of a loop where the bounds are known
  6. // at compile time. Various pragmas achieve similar effects, but are not as
  7. // portable across compilers.
  8. // Example: c10::ForcedUnroll<4>{}(f); is equivalent to f(0); f(1); f(2); f(3);
  9. namespace c10 {
  10. template <int n>
  11. struct ForcedUnroll {
  12. template <typename Func, typename... Args>
  13. C10_ALWAYS_INLINE void operator()(const Func& f, Args... args) const {
  14. ForcedUnroll<n - 1>{}(f, args...);
  15. f(std::integral_constant<int, n - 1>{}, args...);
  16. }
  17. };
  18. template <>
  19. struct ForcedUnroll<1> {
  20. template <typename Func, typename... Args>
  21. C10_ALWAYS_INLINE void operator()(const Func& f, Args... args) const {
  22. f(std::integral_constant<int, 0>{}, args...);
  23. }
  24. };
  25. } // namespace c10
  26. #else
  27. #error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
  28. #endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)