CUDAMathCompat.h 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152
  1. #pragma once
  2. /* This file defines math functions compatible across different gpu
  3. * platforms (currently CUDA and HIP).
  4. */
  5. #if defined(__CUDACC__) || defined(__HIPCC__)
  6. #include <c10/macros/Macros.h>
  7. #include <c10/util/Exception.h>
  8. #ifdef __HIPCC__
  9. #define __MATH_FUNCTIONS_DECL__ inline C10_DEVICE
  10. #else /* __HIPCC__ */
  11. #ifdef __CUDACC_RTC__
  12. #define __MATH_FUNCTIONS_DECL__ C10_HOST_DEVICE
  13. #else /* __CUDACC_RTC__ */
  14. #define __MATH_FUNCTIONS_DECL__ inline C10_HOST_DEVICE
  15. #endif /* __CUDACC_RTC__ */
  16. #endif /* __HIPCC__ */
  17. namespace c10::cuda::compat {
  18. __MATH_FUNCTIONS_DECL__ float abs(float x) {
  19. return ::fabsf(x);
  20. }
  21. __MATH_FUNCTIONS_DECL__ double abs(double x) {
  22. return ::fabs(x);
  23. }
  24. __MATH_FUNCTIONS_DECL__ float exp(float x) {
  25. return ::expf(x);
  26. }
  27. __MATH_FUNCTIONS_DECL__ double exp(double x) {
  28. return ::exp(x);
  29. }
  30. __MATH_FUNCTIONS_DECL__ float ceil(float x) {
  31. return ::ceilf(x);
  32. }
  33. __MATH_FUNCTIONS_DECL__ double ceil(double x) {
  34. return ::ceil(x);
  35. }
  36. __MATH_FUNCTIONS_DECL__ float copysign(float x, float y) {
  37. #if defined(__CUDA_ARCH__) || defined(__HIPCC__)
  38. return ::copysignf(x, y);
  39. #else
  40. // std::copysign gets ICE/Segfaults with gcc 7.5/8 on arm64
  41. // (e.g. Jetson), see PyTorch PR #51834
  42. // This host function needs to be here for the compiler but is never used
  43. TORCH_INTERNAL_ASSERT(
  44. false, "CUDAMathCompat copysign should not run on the CPU");
  45. #endif
  46. }
  47. __MATH_FUNCTIONS_DECL__ double copysign(double x, double y) {
  48. #if defined(__CUDA_ARCH__) || defined(__HIPCC__)
  49. return ::copysign(x, y);
  50. #else
  51. // see above
  52. TORCH_INTERNAL_ASSERT(
  53. false, "CUDAMathCompat copysign should not run on the CPU");
  54. #endif
  55. }
  56. __MATH_FUNCTIONS_DECL__ float floor(float x) {
  57. return ::floorf(x);
  58. }
  59. __MATH_FUNCTIONS_DECL__ double floor(double x) {
  60. return ::floor(x);
  61. }
  62. __MATH_FUNCTIONS_DECL__ float log(float x) {
  63. return ::logf(x);
  64. }
  65. __MATH_FUNCTIONS_DECL__ double log(double x) {
  66. return ::log(x);
  67. }
  68. __MATH_FUNCTIONS_DECL__ float log1p(float x) {
  69. return ::log1pf(x);
  70. }
  71. __MATH_FUNCTIONS_DECL__ double log1p(double x) {
  72. return ::log1p(x);
  73. }
  74. __MATH_FUNCTIONS_DECL__ float max(float x, float y) {
  75. return ::fmaxf(x, y);
  76. }
  77. __MATH_FUNCTIONS_DECL__ double max(double x, double y) {
  78. return ::fmax(x, y);
  79. }
  80. __MATH_FUNCTIONS_DECL__ float min(float x, float y) {
  81. return ::fminf(x, y);
  82. }
  83. __MATH_FUNCTIONS_DECL__ double min(double x, double y) {
  84. return ::fmin(x, y);
  85. }
  86. __MATH_FUNCTIONS_DECL__ float pow(float x, float y) {
  87. return ::powf(x, y);
  88. }
  89. __MATH_FUNCTIONS_DECL__ double pow(double x, double y) {
  90. return ::pow(x, y);
  91. }
  92. __MATH_FUNCTIONS_DECL__ void sincos(float x, float* sptr, float* cptr) {
  93. return ::sincosf(x, sptr, cptr);
  94. }
  95. __MATH_FUNCTIONS_DECL__ void sincos(double x, double* sptr, double* cptr) {
  96. return ::sincos(x, sptr, cptr);
  97. }
  98. __MATH_FUNCTIONS_DECL__ float sqrt(float x) {
  99. return ::sqrtf(x);
  100. }
  101. __MATH_FUNCTIONS_DECL__ double sqrt(double x) {
  102. return ::sqrt(x);
  103. }
  104. __MATH_FUNCTIONS_DECL__ float rsqrt(float x) {
  105. return ::rsqrtf(x);
  106. }
  107. __MATH_FUNCTIONS_DECL__ double rsqrt(double x) {
  108. return ::rsqrt(x);
  109. }
  110. __MATH_FUNCTIONS_DECL__ float tan(float x) {
  111. return ::tanf(x);
  112. }
  113. __MATH_FUNCTIONS_DECL__ double tan(double x) {
  114. return ::tan(x);
  115. }
  116. __MATH_FUNCTIONS_DECL__ float tanh(float x) {
  117. return ::tanhf(x);
  118. }
  119. __MATH_FUNCTIONS_DECL__ double tanh(double x) {
  120. return ::tanh(x);
  121. }
  122. __MATH_FUNCTIONS_DECL__ float normcdf(float x) {
  123. return ::normcdff(x);
  124. }
  125. __MATH_FUNCTIONS_DECL__ double normcdf(double x) {
  126. return ::normcdf(x);
  127. }
  128. } // namespace c10::cuda::compat
  129. #endif