common.h 6.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140
  1. // !!!! PLEASE READ !!!!
  2. // Minimize (transitively) included headers from _avx*.cc because some of the
  3. // functions defined in the headers compiled with platform dependent compiler
  4. // options can be reused by other translation units generating illegal
  5. // instruction run-time error.
  6. // Common utilities for writing performance kernels and easy dispatching of
  7. // different backends.
  8. /*
  9. The general workflow shall be as follows, say we want to
  10. implement a functionality called void foo(int a, float b).
  11. In foo.h, do:
  12. void foo(int a, float b);
  13. In foo_avx512.cc, do:
  14. void foo__avx512(int a, float b) {
  15. [actual avx512 implementation]
  16. }
  17. In foo_avx2.cc, do:
  18. void foo__avx2(int a, float b) {
  19. [actual avx2 implementation]
  20. }
  21. In foo_avx.cc, do:
  22. void foo__avx(int a, float b) {
  23. [actual avx implementation]
  24. }
  25. In foo.cc, do:
  26. // The base implementation should *always* be provided.
  27. void foo__base(int a, float b) {
  28. [base, possibly slow implementation]
  29. }
  30. decltype(foo__base) foo__avx512;
  31. decltype(foo__base) foo__avx2;
  32. decltype(foo__base) foo__avx;
  33. void foo(int a, float b) {
  34. // You should always order things by their preference, faster
  35. // implementations earlier in the function.
  36. AVX512_DO(foo, a, b);
  37. AVX2_DO(foo, a, b);
  38. AVX_DO(foo, a, b);
  39. BASE_DO(foo, a, b);
  40. }
  41. */
  42. // Details: this functionality basically covers the cases for both build time
  43. // and run time architecture support.
  44. //
  45. // During build time:
  46. // The build system should provide flags CAFFE2_PERF_WITH_AVX512,
  47. // CAFFE2_PERF_WITH_AVX2, and CAFFE2_PERF_WITH_AVX that corresponds to the
  48. // __AVX512F__, __AVX512DQ__, __AVX512VL__, __AVX2__, and __AVX__ flags the
  49. // compiler provides. Note that we do not use the compiler flags but rely on
  50. // the build system flags, because the common files (like foo.cc above) will
  51. // always be built without __AVX512F__, __AVX512DQ__, __AVX512VL__, __AVX2__
  52. // and __AVX__.
  53. // During run time:
  54. // we use cpuinfo to identify cpu support and run the proper functions.
  55. #pragma once
  56. #if defined(CAFFE2_PERF_WITH_SVE) || defined(CAFFE2_PERF_WITH_AVX512) || \
  57. defined(CAFFE2_PERF_WITH_AVX2) || defined(CAFFE2_PERF_WITH_AVX)
  58. #include <cpuinfo.h>
  59. #endif
  60. // DO macros: these should be used in your entry function, similar to foo()
  61. // above, that routes implementations based on CPU capability.
  62. #define BASE_DO(funcname, ...) return funcname##__base(__VA_ARGS__);
  63. #ifdef CAFFE2_PERF_WITH_SVE
  64. #define SVE_DO(funcname, ...) \
  65. { \
  66. static const bool isDo = cpuinfo_initialize() && cpuinfo_has_arm_sve(); \
  67. if (isDo) { \
  68. return funcname##__sve(__VA_ARGS__); \
  69. } \
  70. }
  71. #else // CAFFE2_PERF_WITH_SVE
  72. #define SVE_DO(funcname, ...)
  73. #endif // CAFFE2_PERF_WITH_SVE
  74. #ifdef CAFFE2_PERF_WITH_AVX512
  75. #define AVX512_DO(funcname, ...) \
  76. { \
  77. static const bool isDo = cpuinfo_initialize() && \
  78. cpuinfo_has_x86_avx512f() && cpuinfo_has_x86_avx512dq() && \
  79. cpuinfo_has_x86_avx512vl(); \
  80. if (isDo) { \
  81. return funcname##__avx512(__VA_ARGS__); \
  82. } \
  83. }
  84. #else // CAFFE2_PERF_WITH_AVX512
  85. #define AVX512_DO(funcname, ...)
  86. #endif // CAFFE2_PERF_WITH_AVX512
  87. #ifdef CAFFE2_PERF_WITH_AVX2
  88. #define AVX2_DO(funcname, ...) \
  89. { \
  90. static const bool isDo = cpuinfo_initialize() && cpuinfo_has_x86_avx2(); \
  91. if (isDo) { \
  92. return funcname##__avx2(__VA_ARGS__); \
  93. } \
  94. }
  95. #define AVX2_FMA_DO(funcname, ...) \
  96. { \
  97. static const bool isDo = cpuinfo_initialize() && cpuinfo_has_x86_avx2() && \
  98. cpuinfo_has_x86_fma3(); \
  99. if (isDo) { \
  100. return funcname##__avx2_fma(__VA_ARGS__); \
  101. } \
  102. }
  103. #else // CAFFE2_PERF_WITH_AVX2
  104. #define AVX2_DO(funcname, ...)
  105. #define AVX2_FMA_DO(funcname, ...)
  106. #endif // CAFFE2_PERF_WITH_AVX2
  107. #ifdef CAFFE2_PERF_WITH_AVX
  108. #define AVX_DO(funcname, ...) \
  109. { \
  110. static const bool isDo = cpuinfo_initialize() && cpuinfo_has_x86_avx(); \
  111. if (isDo) { \
  112. return funcname##__avx(__VA_ARGS__); \
  113. } \
  114. }
  115. #define AVX_F16C_DO(funcname, ...) \
  116. { \
  117. static const bool isDo = cpuinfo_initialize() && cpuinfo_has_x86_avx() && \
  118. cpuinfo_has_x86_f16c(); \
  119. if (isDo) { \
  120. return funcname##__avx_f16c(__VA_ARGS__); \
  121. } \
  122. }
  123. #else // CAFFE2_PERF_WITH_AVX
  124. #define AVX_DO(funcname, ...)
  125. #define AVX_F16C_DO(funcname, ...)
  126. #endif // CAFFE2_PERF_WITH_AVX