llvmMathExtras.h 29 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910
  1. #if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
  2. //===-- llvm/Support/MathExtras.h - Useful math functions -------*- C++ -*-===//
  3. //
  4. // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
  5. // See https://llvm.org/LICENSE.txt for license information.
  6. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
  7. //
  8. //===----------------------------------------------------------------------===//
  9. //
  10. // This file contains some functions that are useful for math stuff.
  11. //
  12. //===----------------------------------------------------------------------===//
  13. #pragma once
  14. #include <c10/util/bit_cast.h>
  15. #include <algorithm>
  16. #include <cassert>
  17. #include <climits>
  18. #include <cmath>
  19. #include <cstdint>
  20. #include <cstring>
  21. #include <limits>
  22. #include <type_traits>
  23. #ifdef __ANDROID_NDK__
  24. #include <android/api-level.h>
  25. #endif
  26. #ifndef __has_builtin
  27. #define __has_builtin(x) 0
  28. #endif
  29. #ifndef LLVM_GNUC_PREREQ
  30. #if defined(__GNUC__) && defined(__GNUC_MINOR__) && defined(__GNUC_PATCHLEVEL__)
  31. #define LLVM_GNUC_PREREQ(maj, min, patch) \
  32. ((__GNUC__ << 20) + (__GNUC_MINOR__ << 10) + __GNUC_PATCHLEVEL__ >= \
  33. ((maj) << 20) + ((min) << 10) + (patch))
  34. #elif defined(__GNUC__) && defined(__GNUC_MINOR__)
  35. #define LLVM_GNUC_PREREQ(maj, min, patch) \
  36. ((__GNUC__ << 20) + (__GNUC_MINOR__ << 10) >= ((maj) << 20) + ((min) << 10))
  37. #else
  38. #define LLVM_GNUC_PREREQ(maj, min, patch) 0
  39. #endif
  40. #endif
  41. #ifdef _MSC_VER
  42. // Declare these intrinsics manually rather including intrin.h. It's very
  43. // expensive, and MathExtras.h is popular.
  44. // #include <intrin.h>
  45. extern "C" {
  46. unsigned char _BitScanForward(unsigned long* _Index, unsigned long _Mask);
  47. unsigned char _BitScanForward64(unsigned long* _Index, unsigned __int64 _Mask);
  48. unsigned char _BitScanReverse(unsigned long* _Index, unsigned long _Mask);
  49. unsigned char _BitScanReverse64(unsigned long* _Index, unsigned __int64 _Mask);
  50. }
  51. #endif
  52. namespace c10::llvm {
  53. /// The behavior an operation has on an input of 0.
  54. enum ZeroBehavior {
  55. /// The returned value is undefined.
  56. ZB_Undefined,
  57. /// The returned value is numeric_limits<T>::max()
  58. ZB_Max,
  59. /// The returned value is numeric_limits<T>::digits
  60. ZB_Width
  61. };
  62. namespace detail {
  63. template <typename T, std::size_t SizeOfT>
  64. struct TrailingZerosCounter {
  65. static std::size_t count(T Val, ZeroBehavior /*unused*/) {
  66. if (!Val)
  67. return std::numeric_limits<T>::digits;
  68. if (Val & 0x1)
  69. return 0;
  70. // Bisection method.
  71. std::size_t ZeroBits = 0;
  72. T Shift = std::numeric_limits<T>::digits >> 1;
  73. T Mask = std::numeric_limits<T>::max() >> Shift;
  74. while (Shift) {
  75. if ((Val & Mask) == 0) {
  76. Val >>= Shift;
  77. ZeroBits |= Shift;
  78. }
  79. Shift >>= 1;
  80. Mask >>= Shift;
  81. }
  82. return ZeroBits;
  83. }
  84. };
  85. #if (defined(__GNUC__) && __GNUC__ >= 4) || defined(_MSC_VER)
  86. template <typename T>
  87. struct TrailingZerosCounter<T, 4> {
  88. static std::size_t count(T Val, ZeroBehavior ZB) {
  89. if (ZB != ZB_Undefined && Val == 0)
  90. return 32;
  91. #if __has_builtin(__builtin_ctz) || LLVM_GNUC_PREREQ(4, 0, 0)
  92. return __builtin_ctz(Val);
  93. #elif defined(_MSC_VER)
  94. unsigned long Index;
  95. _BitScanForward(&Index, Val);
  96. return Index;
  97. #endif
  98. }
  99. };
  100. #if !defined(_MSC_VER) || defined(_M_X64)
  101. template <typename T>
  102. struct TrailingZerosCounter<T, 8> {
  103. static std::size_t count(T Val, ZeroBehavior ZB) {
  104. if (ZB != ZB_Undefined && Val == 0)
  105. return 64;
  106. #if __has_builtin(__builtin_ctzll) || LLVM_GNUC_PREREQ(4, 0, 0)
  107. return __builtin_ctzll(Val);
  108. #elif defined(_MSC_VER)
  109. unsigned long Index;
  110. _BitScanForward64(&Index, Val);
  111. return Index;
  112. #endif
  113. }
  114. };
  115. #endif
  116. #endif
  117. } // namespace detail
  118. /// Count number of 0's from the least significant bit to the most
  119. /// stopping at the first 1.
  120. ///
  121. /// Only unsigned integral types are allowed.
  122. ///
  123. /// \param ZB the behavior on an input of 0. Only ZB_Width and ZB_Undefined are
  124. /// valid arguments.
  125. template <typename T>
  126. std::size_t countTrailingZeros(T Val, ZeroBehavior ZB = ZB_Width) {
  127. static_assert(
  128. std::numeric_limits<T>::is_integer && !std::numeric_limits<T>::is_signed,
  129. "Only unsigned integral types are allowed.");
  130. return llvm::detail::TrailingZerosCounter<T, sizeof(T)>::count(Val, ZB);
  131. }
  132. namespace detail {
  133. template <typename T, std::size_t SizeOfT>
  134. struct LeadingZerosCounter {
  135. static std::size_t count(T Val, ZeroBehavior /*unused*/) {
  136. if (!Val)
  137. return std::numeric_limits<T>::digits;
  138. // Bisection method.
  139. std::size_t ZeroBits = 0;
  140. for (T Shift = std::numeric_limits<T>::digits >> 1; Shift; Shift >>= 1) {
  141. T Tmp = Val >> Shift;
  142. if (Tmp)
  143. Val = Tmp;
  144. else
  145. ZeroBits |= Shift;
  146. }
  147. return ZeroBits;
  148. }
  149. };
  150. #if (defined(__GNUC__) && __GNUC__ >= 4) || defined(_MSC_VER)
  151. template <typename T>
  152. struct LeadingZerosCounter<T, 4> {
  153. static std::size_t count(T Val, ZeroBehavior ZB) {
  154. if (ZB != ZB_Undefined && Val == 0)
  155. return 32;
  156. #if __has_builtin(__builtin_clz) || LLVM_GNUC_PREREQ(4, 0, 0)
  157. return __builtin_clz(Val);
  158. #elif defined(_MSC_VER)
  159. unsigned long Index;
  160. _BitScanReverse(&Index, Val);
  161. return Index ^ 31;
  162. #endif
  163. }
  164. };
  165. #if !defined(_MSC_VER) || defined(_M_X64)
  166. template <typename T>
  167. struct LeadingZerosCounter<T, 8> {
  168. static std::size_t count(T Val, ZeroBehavior ZB) {
  169. if (ZB != ZB_Undefined && Val == 0)
  170. return 64;
  171. #if __has_builtin(__builtin_clzll) || LLVM_GNUC_PREREQ(4, 0, 0)
  172. return __builtin_clzll(Val);
  173. #elif defined(_MSC_VER)
  174. unsigned long Index;
  175. _BitScanReverse64(&Index, Val);
  176. return Index ^ 63;
  177. #endif
  178. }
  179. };
  180. #endif
  181. #endif
  182. } // namespace detail
  183. /// Count number of 0's from the most significant bit to the least
  184. /// stopping at the first 1.
  185. ///
  186. /// Only unsigned integral types are allowed.
  187. ///
  188. /// \param ZB the behavior on an input of 0. Only ZB_Width and ZB_Undefined are
  189. /// valid arguments.
  190. template <typename T>
  191. std::size_t countLeadingZeros(T Val, ZeroBehavior ZB = ZB_Width) {
  192. static_assert(
  193. std::numeric_limits<T>::is_integer && !std::numeric_limits<T>::is_signed,
  194. "Only unsigned integral types are allowed.");
  195. return llvm::detail::LeadingZerosCounter<T, sizeof(T)>::count(Val, ZB);
  196. }
  197. /// Get the index of the first set bit starting from the least
  198. /// significant bit.
  199. ///
  200. /// Only unsigned integral types are allowed.
  201. ///
  202. /// \param ZB the behavior on an input of 0. Only ZB_Max and ZB_Undefined are
  203. /// valid arguments.
  204. template <typename T>
  205. T findFirstSet(T Val, ZeroBehavior ZB = ZB_Max) {
  206. if (ZB == ZB_Max && Val == 0)
  207. return std::numeric_limits<T>::max();
  208. return countTrailingZeros(Val, ZB_Undefined);
  209. }
  210. /// Create a bitmask with the N right-most bits set to 1, and all other
  211. /// bits set to 0. Only unsigned types are allowed.
  212. template <typename T>
  213. T maskTrailingOnes(unsigned N) {
  214. static_assert(std::is_unsigned_v<T>, "Invalid type!");
  215. const unsigned Bits = CHAR_BIT * sizeof(T);
  216. assert(N <= Bits && "Invalid bit index");
  217. return N == 0 ? 0 : (T(-1) >> (Bits - N));
  218. }
  219. /// Create a bitmask with the N left-most bits set to 1, and all other
  220. /// bits set to 0. Only unsigned types are allowed.
  221. template <typename T>
  222. T maskLeadingOnes(unsigned N) {
  223. return ~maskTrailingOnes<T>(CHAR_BIT * sizeof(T) - N);
  224. }
  225. /// Create a bitmask with the N right-most bits set to 0, and all other
  226. /// bits set to 1. Only unsigned types are allowed.
  227. template <typename T>
  228. T maskTrailingZeros(unsigned N) {
  229. return maskLeadingOnes<T>(CHAR_BIT * sizeof(T) - N);
  230. }
  231. /// Create a bitmask with the N left-most bits set to 0, and all other
  232. /// bits set to 1. Only unsigned types are allowed.
  233. template <typename T>
  234. T maskLeadingZeros(unsigned N) {
  235. return maskTrailingOnes<T>(CHAR_BIT * sizeof(T) - N);
  236. }
  237. /// Get the index of the last set bit starting from the least
  238. /// significant bit.
  239. ///
  240. /// Only unsigned integral types are allowed.
  241. ///
  242. /// \param ZB the behavior on an input of 0. Only ZB_Max and ZB_Undefined are
  243. /// valid arguments.
  244. template <typename T>
  245. T findLastSet(T Val, ZeroBehavior ZB = ZB_Max) {
  246. if (ZB == ZB_Max && Val == 0)
  247. return std::numeric_limits<T>::max();
  248. // Use ^ instead of - because both gcc and llvm can remove the associated ^
  249. // in the __builtin_clz intrinsic on x86.
  250. return countLeadingZeros(Val, ZB_Undefined) ^
  251. (std::numeric_limits<T>::digits - 1);
  252. }
  253. /// Macro compressed bit reversal table for 256 bits.
  254. ///
  255. /// http://graphics.stanford.edu/~seander/bithacks.html#BitReverseTable
  256. /// NOLINTNEXTLINE(*c-arrays*)
  257. static constexpr unsigned char BitReverseTable256[256] = {
  258. #define R2(n) n, n + 2 * 64, n + 1 * 64, n + 3 * 64
  259. #define R4(n) R2(n), R2(n + 2 * 16), R2(n + 1 * 16), R2(n + 3 * 16)
  260. #define R6(n) R4(n), R4(n + 2 * 4), R4(n + 1 * 4), R4(n + 3 * 4)
  261. R6(0),
  262. R6(2),
  263. R6(1),
  264. R6(3)
  265. #undef R2
  266. #undef R4
  267. #undef R6
  268. };
  269. /// Reverse the bits in \p Val.
  270. template <typename T>
  271. T reverseBits(T Val) {
  272. // NOLINTNEXTLINE(*c-arrays*)
  273. unsigned char in[sizeof(Val)];
  274. // NOLINTNEXTLINE(*c-arrays*)
  275. unsigned char out[sizeof(Val)];
  276. std::memcpy(in, &Val, sizeof(Val));
  277. for (unsigned i = 0; i < sizeof(Val); ++i)
  278. out[(sizeof(Val) - i) - 1] = BitReverseTable256[in[i]];
  279. std::memcpy(&Val, out, sizeof(Val));
  280. return Val;
  281. }
  282. // NOTE: The following support functions use the _32/_64 extensions instead of
  283. // type overloading so that signed and unsigned integers can be used without
  284. // ambiguity.
  285. /// Return the high 32 bits of a 64 bit value.
  286. constexpr inline uint32_t Hi_32(uint64_t Value) {
  287. return static_cast<uint32_t>(Value >> 32);
  288. }
  289. /// Return the low 32 bits of a 64 bit value.
  290. constexpr inline uint32_t Lo_32(uint64_t Value) {
  291. return static_cast<uint32_t>(Value);
  292. }
  293. /// Make a 64-bit integer from a high / low pair of 32-bit integers.
  294. constexpr inline uint64_t Make_64(uint32_t High, uint32_t Low) {
  295. return ((uint64_t)High << 32) | (uint64_t)Low;
  296. }
  297. /// Checks if an integer fits into the given bit width.
  298. template <unsigned N>
  299. constexpr inline bool isInt(int64_t x) {
  300. return N >= 64 ||
  301. (-(INT64_C(1) << (N - 1)) <= x && x < (INT64_C(1) << (N - 1)));
  302. }
  303. // Template specializations to get better code for common cases.
  304. template <>
  305. constexpr inline bool isInt<8>(int64_t x) {
  306. return static_cast<int8_t>(x) == x;
  307. }
  308. template <>
  309. constexpr inline bool isInt<16>(int64_t x) {
  310. return static_cast<int16_t>(x) == x;
  311. }
  312. template <>
  313. constexpr inline bool isInt<32>(int64_t x) {
  314. return static_cast<int32_t>(x) == x;
  315. }
  316. /// Checks if a signed integer is an N bit number shifted left by S.
  317. template <unsigned N, unsigned S>
  318. constexpr inline bool isShiftedInt(int64_t x) {
  319. static_assert(
  320. N > 0, "isShiftedInt<0> doesn't make sense (refers to a 0-bit number.");
  321. static_assert(N + S <= 64, "isShiftedInt<N, S> with N + S > 64 is too wide.");
  322. return isInt<N + S>(x) && (x % (UINT64_C(1) << S) == 0);
  323. }
  324. /// Checks if an unsigned integer fits into the given bit width.
  325. ///
  326. /// This is written as two functions rather than as simply
  327. ///
  328. /// return N >= 64 || X < (UINT64_C(1) << N);
  329. ///
  330. /// to keep MSVC from (incorrectly) warning on isUInt<64> that we're shifting
  331. /// left too many places.
  332. template <unsigned N>
  333. constexpr inline std::enable_if_t<(N < 64), bool> isUInt(uint64_t X) {
  334. static_assert(N > 0, "isUInt<0> doesn't make sense");
  335. return X < (UINT64_C(1) << N);
  336. }
  337. template <unsigned N>
  338. constexpr inline std::enable_if_t<N >= 64, bool> isUInt(uint64_t /*X*/) {
  339. return true;
  340. }
  341. // Template specializations to get better code for common cases.
  342. template <>
  343. constexpr inline bool isUInt<8>(uint64_t x) {
  344. return static_cast<uint8_t>(x) == x;
  345. }
  346. template <>
  347. constexpr inline bool isUInt<16>(uint64_t x) {
  348. return static_cast<uint16_t>(x) == x;
  349. }
  350. template <>
  351. constexpr inline bool isUInt<32>(uint64_t x) {
  352. return static_cast<uint32_t>(x) == x;
  353. }
  354. /// Checks if a unsigned integer is an N bit number shifted left by S.
  355. template <unsigned N, unsigned S>
  356. constexpr inline bool isShiftedUInt(uint64_t x) {
  357. static_assert(
  358. N > 0, "isShiftedUInt<0> doesn't make sense (refers to a 0-bit number)");
  359. static_assert(
  360. N + S <= 64, "isShiftedUInt<N, S> with N + S > 64 is too wide.");
  361. // Per the two static_asserts above, S must be strictly less than 64. So
  362. // 1 << S is not undefined behavior.
  363. return isUInt<N + S>(x) && (x % (UINT64_C(1) << S) == 0);
  364. }
  365. /// Gets the maximum value for a N-bit unsigned integer.
  366. inline uint64_t maxUIntN(uint64_t N) {
  367. assert(N > 0 && N <= 64 && "integer width out of range");
  368. // uint64_t(1) << 64 is undefined behavior, so we can't do
  369. // (uint64_t(1) << N) - 1
  370. // without checking first that N != 64. But this works and doesn't have a
  371. // branch.
  372. return UINT64_MAX >> (64 - N);
  373. }
  374. // Ignore the false warning "Arithmetic overflow" for MSVC
  375. #ifdef _MSC_VER
  376. #pragma warning(push)
  377. #pragma warning(disable : 4146)
  378. #endif
  379. /// Gets the minimum value for a N-bit signed integer.
  380. inline int64_t minIntN(int64_t N) {
  381. assert(N > 0 && N <= 64 && "integer width out of range");
  382. // NOLINTNEXTLINE(*-narrowing-conversions)
  383. return -(UINT64_C(1) << (N - 1));
  384. }
  385. #ifdef _MSC_VER
  386. #pragma warning(pop)
  387. #endif
  388. /// Gets the maximum value for a N-bit signed integer.
  389. inline int64_t maxIntN(int64_t N) {
  390. assert(N > 0 && N <= 64 && "integer width out of range");
  391. // This relies on two's complement wraparound when N == 64, so we convert to
  392. // int64_t only at the very end to avoid UB.
  393. // NOLINTNEXTLINE(*-narrowing-conversions)
  394. return (UINT64_C(1) << (N - 1)) - 1;
  395. }
  396. /// Checks if an unsigned integer fits into the given (dynamic) bit width.
  397. inline bool isUIntN(unsigned N, uint64_t x) {
  398. return N >= 64 || x <= maxUIntN(N);
  399. }
  400. /// Checks if an signed integer fits into the given (dynamic) bit width.
  401. inline bool isIntN(unsigned N, int64_t x) {
  402. return N >= 64 || (minIntN(N) <= x && x <= maxIntN(N));
  403. }
  404. /// Return true if the argument is a non-empty sequence of ones starting at the
  405. /// least significant bit with the remainder zero (32 bit version).
  406. /// Ex. isMask_32(0x0000FFFFU) == true.
  407. constexpr inline bool isMask_32(uint32_t Value) {
  408. return Value && ((Value + 1) & Value) == 0;
  409. }
  410. /// Return true if the argument is a non-empty sequence of ones starting at the
  411. /// least significant bit with the remainder zero (64 bit version).
  412. constexpr inline bool isMask_64(uint64_t Value) {
  413. return Value && ((Value + 1) & Value) == 0;
  414. }
  415. /// Return true if the argument contains a non-empty sequence of ones with the
  416. /// remainder zero (32 bit version.) Ex. isShiftedMask_32(0x0000FF00U) == true.
  417. constexpr inline bool isShiftedMask_32(uint32_t Value) {
  418. return Value && isMask_32((Value - 1) | Value);
  419. }
  420. /// Return true if the argument contains a non-empty sequence of ones with the
  421. /// remainder zero (64 bit version.)
  422. constexpr inline bool isShiftedMask_64(uint64_t Value) {
  423. return Value && isMask_64((Value - 1) | Value);
  424. }
  425. /// Return true if the argument is a power of two > 0.
  426. /// Ex. isPowerOf2_32(0x00100000U) == true (32 bit edition.)
  427. constexpr inline bool isPowerOf2_32(uint32_t Value) {
  428. return Value && !(Value & (Value - 1));
  429. }
  430. /// Return true if the argument is a power of two > 0 (64 bit edition.)
  431. constexpr inline bool isPowerOf2_64(uint64_t Value) {
  432. return Value && !(Value & (Value - 1));
  433. }
  434. /// Count the number of ones from the most significant bit to the first
  435. /// zero bit.
  436. ///
  437. /// Ex. countLeadingOnes(0xFF0FFF00) == 8.
  438. /// Only unsigned integral types are allowed.
  439. ///
  440. /// \param ZB the behavior on an input of all ones. Only ZB_Width and
  441. /// ZB_Undefined are valid arguments.
  442. template <typename T>
  443. std::size_t countLeadingOnes(T Value, ZeroBehavior ZB = ZB_Width) {
  444. static_assert(
  445. std::numeric_limits<T>::is_integer && !std::numeric_limits<T>::is_signed,
  446. "Only unsigned integral types are allowed.");
  447. return countLeadingZeros<T>(~Value, ZB);
  448. }
  449. /// Count the number of ones from the least significant bit to the first
  450. /// zero bit.
  451. ///
  452. /// Ex. countTrailingOnes(0x00FF00FF) == 8.
  453. /// Only unsigned integral types are allowed.
  454. ///
  455. /// \param ZB the behavior on an input of all ones. Only ZB_Width and
  456. /// ZB_Undefined are valid arguments.
  457. template <typename T>
  458. std::size_t countTrailingOnes(T Value, ZeroBehavior ZB = ZB_Width) {
  459. static_assert(
  460. std::numeric_limits<T>::is_integer && !std::numeric_limits<T>::is_signed,
  461. "Only unsigned integral types are allowed.");
  462. return countTrailingZeros<T>(~Value, ZB);
  463. }
  464. namespace detail {
  465. template <typename T, std::size_t SizeOfT>
  466. struct PopulationCounter {
  467. static unsigned count(T Value) {
  468. // Generic version, forward to 32 bits.
  469. static_assert(SizeOfT <= 4, "Not implemented!");
  470. #if defined(__GNUC__) && __GNUC__ >= 4
  471. return __builtin_popcount(Value);
  472. #else
  473. uint32_t v = Value;
  474. v = v - ((v >> 1) & 0x55555555);
  475. v = (v & 0x33333333) + ((v >> 2) & 0x33333333);
  476. return ((v + (v >> 4) & 0xF0F0F0F) * 0x1010101) >> 24;
  477. #endif
  478. }
  479. };
  480. template <typename T>
  481. struct PopulationCounter<T, 8> {
  482. static unsigned count(T Value) {
  483. #if defined(__GNUC__) && __GNUC__ >= 4
  484. return __builtin_popcountll(Value);
  485. #else
  486. uint64_t v = Value;
  487. v = v - ((v >> 1) & 0x5555555555555555ULL);
  488. v = (v & 0x3333333333333333ULL) + ((v >> 2) & 0x3333333333333333ULL);
  489. v = (v + (v >> 4)) & 0x0F0F0F0F0F0F0F0FULL;
  490. return unsigned((uint64_t)(v * 0x0101010101010101ULL) >> 56);
  491. #endif
  492. }
  493. };
  494. } // namespace detail
  495. /// Count the number of set bits in a value.
  496. /// Ex. countPopulation(0xF000F000) = 8
  497. /// Returns 0 if the word is zero.
  498. template <typename T>
  499. inline unsigned countPopulation(T Value) {
  500. static_assert(
  501. std::numeric_limits<T>::is_integer && !std::numeric_limits<T>::is_signed,
  502. "Only unsigned integral types are allowed.");
  503. return detail::PopulationCounter<T, sizeof(T)>::count(Value);
  504. }
  505. /// Return the log base 2 of the specified value.
  506. inline double Log2(double Value) {
  507. #if defined(__ANDROID_API__) && __ANDROID_API__ < 18
  508. return __builtin_log(Value) / __builtin_log(2.0);
  509. #else
  510. return log2(Value);
  511. #endif
  512. }
  513. /// Return the floor log base 2 of the specified value, -1 if the value is zero.
  514. /// (32 bit edition.)
  515. /// Ex. Log2_32(32) == 5, Log2_32(1) == 0, Log2_32(0) == -1, Log2_32(6) == 2
  516. inline unsigned Log2_32(uint32_t Value) {
  517. return static_cast<unsigned>(31 - countLeadingZeros(Value));
  518. }
  519. /// Return the floor log base 2 of the specified value, -1 if the value is zero.
  520. /// (64 bit edition.)
  521. inline unsigned Log2_64(uint64_t Value) {
  522. return static_cast<unsigned>(63 - countLeadingZeros(Value));
  523. }
  524. /// Return the ceil log base 2 of the specified value, 32 if the value is zero.
  525. /// (32 bit edition).
  526. /// Ex. Log2_32_Ceil(32) == 5, Log2_32_Ceil(1) == 0, Log2_32_Ceil(6) == 3
  527. inline unsigned Log2_32_Ceil(uint32_t Value) {
  528. return static_cast<unsigned>(32 - countLeadingZeros(Value - 1));
  529. }
  530. /// Return the ceil log base 2 of the specified value, 64 if the value is zero.
  531. /// (64 bit edition.)
  532. inline unsigned Log2_64_Ceil(uint64_t Value) {
  533. return static_cast<unsigned>(64 - countLeadingZeros(Value - 1));
  534. }
  535. /// Return the greatest common divisor of the values using Euclid's algorithm.
  536. inline uint64_t GreatestCommonDivisor64(uint64_t A, uint64_t B) {
  537. while (B) {
  538. uint64_t T = B;
  539. B = A % B;
  540. A = T;
  541. }
  542. return A;
  543. }
  544. /// This function takes a 64-bit integer and returns the bit equivalent double.
  545. inline double BitsToDouble(uint64_t Bits) {
  546. double D = 0;
  547. static_assert(sizeof(uint64_t) == sizeof(double), "Unexpected type sizes");
  548. memcpy(&D, &Bits, sizeof(Bits));
  549. return D;
  550. }
  551. /// This function takes a 32-bit integer and returns the bit equivalent float.
  552. inline float BitsToFloat(uint32_t Bits) {
  553. // TODO: Use std::bit_cast once C++20 becomes available.
  554. return c10::bit_cast<float>(Bits);
  555. }
  556. /// This function takes a double and returns the bit equivalent 64-bit integer.
  557. /// Note that copying doubles around changes the bits of NaNs on some hosts,
  558. /// notably x86, so this routine cannot be used if these bits are needed.
  559. inline uint64_t DoubleToBits(double Double) {
  560. // NOLINTNEXTLINE(cppcoreguidelines-init-variables)
  561. uint64_t Bits;
  562. static_assert(sizeof(uint64_t) == sizeof(double), "Unexpected type sizes");
  563. memcpy(&Bits, &Double, sizeof(Double));
  564. return Bits;
  565. }
  566. /// This function takes a float and returns the bit equivalent 32-bit integer.
  567. /// Note that copying floats around changes the bits of NaNs on some hosts,
  568. /// notably x86, so this routine cannot be used if these bits are needed.
  569. inline uint32_t FloatToBits(float Float) {
  570. // NOLINTNEXTLINE(cppcoreguidelines-init-variables)
  571. uint32_t Bits;
  572. static_assert(sizeof(uint32_t) == sizeof(float), "Unexpected type sizes");
  573. memcpy(&Bits, &Float, sizeof(Float));
  574. return Bits;
  575. }
  576. /// A and B are either alignments or offsets. Return the minimum alignment that
  577. /// may be assumed after adding the two together.
  578. constexpr inline uint64_t MinAlign(uint64_t A, uint64_t B) {
  579. // The largest power of 2 that divides both A and B.
  580. //
  581. // Replace "-Value" by "1+~Value" in the following commented code to avoid
  582. // MSVC warning C4146
  583. // return (A | B) & -(A | B);
  584. return (A | B) & (1 + ~(A | B));
  585. }
  586. /// Aligns \c Addr to \c Alignment bytes, rounding up.
  587. ///
  588. /// Alignment should be a power of two. This method rounds up, so
  589. /// alignAddr(7, 4) == 8 and alignAddr(8, 4) == 8.
  590. inline uintptr_t alignAddr(const void* Addr, size_t Alignment) {
  591. assert(
  592. Alignment && isPowerOf2_64((uint64_t)Alignment) &&
  593. "Alignment is not a power of two!");
  594. assert((uintptr_t)Addr + Alignment - 1 >= (uintptr_t)Addr);
  595. return (((uintptr_t)Addr + Alignment - 1) & ~(uintptr_t)(Alignment - 1));
  596. }
  597. /// Returns the necessary adjustment for aligning \c Ptr to \c Alignment
  598. /// bytes, rounding up.
  599. inline size_t alignmentAdjustment(const void* Ptr, size_t Alignment) {
  600. return alignAddr(Ptr, Alignment) - (uintptr_t)Ptr;
  601. }
  602. /// Returns the next power of two (in 64-bits) that is strictly greater than A.
  603. /// Returns zero on overflow.
  604. inline uint64_t NextPowerOf2(uint64_t A) {
  605. A |= (A >> 1);
  606. A |= (A >> 2);
  607. A |= (A >> 4);
  608. A |= (A >> 8);
  609. A |= (A >> 16);
  610. A |= (A >> 32);
  611. return A + 1;
  612. }
  613. /// Returns the power of two which is less than or equal to the given value.
  614. /// Essentially, it is a floor operation across the domain of powers of two.
  615. inline uint64_t PowerOf2Floor(uint64_t A) {
  616. if (!A)
  617. return 0;
  618. return 1ull << (63 - countLeadingZeros(A, ZB_Undefined));
  619. }
  620. /// Returns the power of two which is greater than or equal to the given value.
  621. /// Essentially, it is a ceil operation across the domain of powers of two.
  622. inline uint64_t PowerOf2Ceil(uint64_t A) {
  623. if (!A)
  624. return 0;
  625. return NextPowerOf2(A - 1);
  626. }
  627. /// Returns the next integer (mod 2**64) that is greater than or equal to
  628. /// \p Value and is a multiple of \p Align. \p Align must be non-zero.
  629. ///
  630. /// If non-zero \p Skew is specified, the return value will be a minimal
  631. /// integer that is greater than or equal to \p Value and equal to
  632. /// \p Align * N + \p Skew for some integer N. If \p Skew is larger than
  633. /// \p Align, its value is adjusted to '\p Skew mod \p Align'.
  634. ///
  635. /// Examples:
  636. /// \code
  637. /// alignTo(5, 8) = 8
  638. /// alignTo(17, 8) = 24
  639. /// alignTo(~0LL, 8) = 0
  640. /// alignTo(321, 255) = 510
  641. ///
  642. /// alignTo(5, 8, 7) = 7
  643. /// alignTo(17, 8, 1) = 17
  644. /// alignTo(~0LL, 8, 3) = 3
  645. /// alignTo(321, 255, 42) = 552
  646. /// \endcode
  647. inline uint64_t alignTo(uint64_t Value, uint64_t Align, uint64_t Skew = 0) {
  648. assert(Align != 0u && "Align can't be 0.");
  649. Skew %= Align;
  650. return (Value + Align - 1 - Skew) / Align * Align + Skew;
  651. }
  652. /// Returns the next integer (mod 2**64) that is greater than or equal to
  653. /// \p Value and is a multiple of \c Align. \c Align must be non-zero.
  654. template <uint64_t Align>
  655. constexpr inline uint64_t alignTo(uint64_t Value) {
  656. static_assert(Align != 0u, "Align must be non-zero");
  657. return (Value + Align - 1) / Align * Align;
  658. }
  659. /// Returns the integer ceil(Numerator / Denominator).
  660. inline uint64_t divideCeil(uint64_t Numerator, uint64_t Denominator) {
  661. return alignTo(Numerator, Denominator) / Denominator;
  662. }
  663. /// \c alignTo for contexts where a constant expression is required.
  664. /// \sa alignTo
  665. ///
  666. /// \todo FIXME: remove when \c constexpr becomes really \c constexpr
  667. template <uint64_t Align>
  668. struct AlignTo {
  669. static_assert(Align != 0u, "Align must be non-zero");
  670. template <uint64_t Value>
  671. struct from_value {
  672. static const uint64_t value = (Value + Align - 1) / Align * Align;
  673. };
  674. };
  675. /// Returns the largest uint64_t less than or equal to \p Value and is
  676. /// \p Skew mod \p Align. \p Align must be non-zero
  677. inline uint64_t alignDown(uint64_t Value, uint64_t Align, uint64_t Skew = 0) {
  678. assert(Align != 0u && "Align can't be 0.");
  679. Skew %= Align;
  680. return (Value - Skew) / Align * Align + Skew;
  681. }
  682. /// Returns the offset to the next integer (mod 2**64) that is greater than
  683. /// or equal to \p Value and is a multiple of \p Align. \p Align must be
  684. /// non-zero.
  685. inline uint64_t OffsetToAlignment(uint64_t Value, uint64_t Align) {
  686. return alignTo(Value, Align) - Value;
  687. }
  688. /// Sign-extend the number in the bottom B bits of X to a 32-bit integer.
  689. /// Requires 0 < B <= 32.
  690. template <unsigned B>
  691. constexpr inline int32_t SignExtend32(uint32_t X) {
  692. static_assert(B > 0, "Bit width can't be 0.");
  693. static_assert(B <= 32, "Bit width out of range.");
  694. return int32_t(X << (32 - B)) >> (32 - B);
  695. }
  696. /// Sign-extend the number in the bottom B bits of X to a 32-bit integer.
  697. /// Requires 0 < B < 32.
  698. inline int32_t SignExtend32(uint32_t X, unsigned B) {
  699. assert(B > 0 && "Bit width can't be 0.");
  700. assert(B <= 32 && "Bit width out of range.");
  701. return int32_t(X << (32 - B)) >> (32 - B);
  702. }
  703. /// Sign-extend the number in the bottom B bits of X to a 64-bit integer.
  704. /// Requires 0 < B < 64.
  705. template <unsigned B>
  706. constexpr inline int64_t SignExtend64(uint64_t x) {
  707. static_assert(B > 0, "Bit width can't be 0.");
  708. static_assert(B <= 64, "Bit width out of range.");
  709. return int64_t(x << (64 - B)) >> (64 - B);
  710. }
  711. /// Sign-extend the number in the bottom B bits of X to a 64-bit integer.
  712. /// Requires 0 < B < 64.
  713. inline int64_t SignExtend64(uint64_t X, unsigned B) {
  714. assert(B > 0 && "Bit width can't be 0.");
  715. assert(B <= 64 && "Bit width out of range.");
  716. return int64_t(X << (64 - B)) >> (64 - B);
  717. }
  718. /// Subtract two unsigned integers, X and Y, of type T and return the absolute
  719. /// value of the result.
  720. template <typename T>
  721. std::enable_if_t<std::is_unsigned_v<T>, T> AbsoluteDifference(T X, T Y) {
  722. return std::max(X, Y) - std::min(X, Y);
  723. }
  724. /// Add two unsigned integers, X and Y, of type T. Clamp the result to the
  725. /// maximum representable value of T on overflow. ResultOverflowed indicates if
  726. /// the result is larger than the maximum representable value of type T.
  727. template <typename T>
  728. std::enable_if_t<std::is_unsigned_v<T>, T> SaturatingAdd(
  729. T X,
  730. T Y,
  731. bool* ResultOverflowed = nullptr) {
  732. // NOLINTNEXTLINE(cppcoreguidelines-init-variables)
  733. bool Dummy;
  734. bool& Overflowed = ResultOverflowed ? *ResultOverflowed : Dummy;
  735. // Hacker's Delight, p. 29
  736. T Z = X + Y;
  737. Overflowed = (Z < X || Z < Y);
  738. if (Overflowed)
  739. return std::numeric_limits<T>::max();
  740. else
  741. return Z;
  742. }
  743. /// Multiply two unsigned integers, X and Y, of type T. Clamp the result to the
  744. /// maximum representable value of T on overflow. ResultOverflowed indicates if
  745. /// the result is larger than the maximum representable value of type T.
  746. template <typename T>
  747. std::enable_if_t<std::is_unsigned_v<T>, T> SaturatingMultiply(
  748. T X,
  749. T Y,
  750. bool* ResultOverflowed = nullptr) {
  751. // NOLINTNEXTLINE(cppcoreguidelines-init-variables)
  752. bool Dummy;
  753. bool& Overflowed = ResultOverflowed ? *ResultOverflowed : Dummy;
  754. // Hacker's Delight, p. 30 has a different algorithm, but we don't use that
  755. // because it fails for uint16_t (where multiplication can have undefined
  756. // behavior due to promotion to int), and requires a division in addition
  757. // to the multiplication.
  758. Overflowed = false;
  759. // Log2(Z) would be either Log2Z or Log2Z + 1.
  760. // Special case: if X or Y is 0, Log2_64 gives -1, and Log2Z
  761. // will necessarily be less than Log2Max as desired.
  762. int Log2Z = Log2_64(X) + Log2_64(Y);
  763. const T Max = std::numeric_limits<T>::max();
  764. int Log2Max = Log2_64(Max);
  765. if (Log2Z < Log2Max) {
  766. return X * Y;
  767. }
  768. if (Log2Z > Log2Max) {
  769. Overflowed = true;
  770. return Max;
  771. }
  772. // We're going to use the top bit, and maybe overflow one
  773. // bit past it. Multiply all but the bottom bit then add
  774. // that on at the end.
  775. T Z = (X >> 1) * Y;
  776. if (Z & ~(Max >> 1)) {
  777. Overflowed = true;
  778. return Max;
  779. }
  780. Z <<= 1;
  781. if (X & 1)
  782. return SaturatingAdd(Z, Y, ResultOverflowed);
  783. return Z;
  784. }
  785. /// Multiply two unsigned integers, X and Y, and add the unsigned integer, A to
  786. /// the product. Clamp the result to the maximum representable value of T on
  787. /// overflow. ResultOverflowed indicates if the result is larger than the
  788. /// maximum representable value of type T.
  789. template <typename T>
  790. std::enable_if_t<std::is_unsigned_v<T>, T> SaturatingMultiplyAdd(
  791. T X,
  792. T Y,
  793. T A,
  794. bool* ResultOverflowed = nullptr) {
  795. // NOLINTNEXTLINE(cppcoreguidelines-init-variables)
  796. bool Dummy;
  797. bool& Overflowed = ResultOverflowed ? *ResultOverflowed : Dummy;
  798. T Product = SaturatingMultiply(X, Y, &Overflowed);
  799. if (Overflowed)
  800. return Product;
  801. return SaturatingAdd(A, Product, &Overflowed);
  802. }
  803. /// Use this rather than HUGE_VALF; the latter causes warnings on MSVC.
  804. extern const float huge_valf;
  805. } // namespace c10::llvm
  806. #else
  807. #error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
  808. #endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)