flash-attention/csrc/cutlass/include/cute/numeric/math.hpp

/***************************************************************************************************

 * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.

 * SPDX-License-Identifier: BSD-3-Clause

 *

 * Redistribution and use in source and binary forms, with or without

 * modification, are permitted provided that the following conditions are met:

 *

 * 1. Redistributions of source code must retain the above copyright notice, this

 * list of conditions and the following disclaimer.

 *

 * 2. Redistributions in binary form must reproduce the above copyright notice,

 * this list of conditions and the following disclaimer in the documentation

 * and/or other materials provided with the distribution.

 *

 * 3. Neither the name of the copyright holder nor the names of its

 * contributors may be used to endorse or promote products derived from

 * this software without specific prior written permission.

 *

 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"

 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE

 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE

 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE

 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL

 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR

 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER

 * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,

 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE

 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

 *

 **************************************************************************************************/
#pragma once

#include <cute/config.hpp>

#include <cute/util/type_traits.hpp>

namespace cute
{

//
// Common Operations
//

template <class T, class U,
          __CUTE_REQUIRES(is_arithmetic<T>::value &&
                          is_arithmetic<U>::value)>
CUTE_HOST_DEVICE constexpr

auto

max(T const& t, U const& u) {
  return t < u ? u : t;
}

template <class T, class U,
          __CUTE_REQUIRES(is_arithmetic<T>::value &&
                          is_arithmetic<U>::value)>
CUTE_HOST_DEVICE constexpr

auto

min(T const& t, U const& u) {
  return t < u ? t : u;
}

template <class T,
          __CUTE_REQUIRES(is_arithmetic<T>::value)>
CUTE_HOST_DEVICE constexpr

auto

abs(T const& t) {
  if constexpr (is_signed<T>::value) {
    return t < T(0) ? -t : t;
  } else {
    return t;
  }

  CUTE_GCC_UNREACHABLE;
}

// Returns 1 if x > 0, -1 if x < 0, and 0 if x is zero.
template <class T,
          __CUTE_REQUIRES(is_arithmetic<T>::value)>
CUTE_HOST_DEVICE constexpr

int

signum(T const& x) {
  if constexpr (is_signed<T>::value) {
    return (T(0) < x) - (x < T(0));
  } else {
    return T(0) < x;
  }

  CUTE_GCC_UNREACHABLE;
}

//
// C++17 <numeric> operations
//

// Greatest common divisor of two positive integers
template <class T, class U,
          __CUTE_REQUIRES(is_std_integral<T>::value &&
                          is_std_integral<U>::value)>
CUTE_HOST_DEVICE constexpr

cute::common_type_t<T, U>

gcd(T t, U u) {
  while (true) {
    if (t == 0) { return u; }
    u %= t;
    if (u == 0) { return t; }
    t %= u;
  }
}

// Least common multiple of two positive integers
template <class T, class U,
          __CUTE_REQUIRES(is_std_integral<T>::value &&
                          is_std_integral<U>::value)>
CUTE_HOST_DEVICE constexpr

cute::common_type_t<T, U>

lcm(T const& t, U const& u) {
  return (t / gcd(t,u)) * u;
}

//
// C++20 <bit> operations
//

// Checks if a number is an integral power of two
template <class T>
CUTE_HOST_DEVICE constexpr

bool

has_single_bit(T x) {
  return x != 0 && (x & (x - 1)) == 0;
}

// Smallest number of bits needed to represent the given value
//   For x == 0, this is 0
//   For x != 0, this is 1 + floor(log2(x))
// bit_width( 0b0000 ) = 0
// bit_width( 0b0001 ) = 1
// bit_width( 0b0010 ) = 2
// bit_width( 0b0011 ) = 2
// bit_width( 0b0100 ) = 3
// bit_width( 0b0101 ) = 3
// bit_width( 0b0110 ) = 3
// bit_width( 0b0111 ) = 3
template <class T>
CUTE_HOST_DEVICE constexpr

T

bit_width(T x) {
  static_assert(is_unsigned<T>::value, "Only to be used for unsigned types.");
  constexpr int N = (numeric_limits<T>::digits == 64 ? 6 :
                    (numeric_limits<T>::digits == 32 ? 5 :
                    (numeric_limits<T>::digits == 16 ? 4 :
                    (numeric_limits<T>::digits ==  8 ? 3 : (assert(false),0)))));
  T r = 0;
  for (int i = N - 1; i >= 0; --i) {
    T shift = (x > ((T(1) << (T(1) << i))-1)) << i;
    x >>= shift;
    r  |= shift;
  }
  return r + (x != 0);
}

// Smallest integral power of two not less than the given value
// bit_ceil( 0b00000000 ) = 0b00000001
// bit_ceil( 0b00000001 ) = 0b00000001
// bit_ceil( 0b00000010 ) = 0b00000010
// bit_ceil( 0b00000011 ) = 0b00000100
// bit_ceil( 0b00000100 ) = 0b00000100
// bit_ceil( 0b00000101 ) = 0b00001000
// bit_ceil( 0b00000110 ) = 0b00001000
// bit_ceil( 0b00000111 ) = 0b00001000
// bit_ceil( 0b00001000 ) = 0b00001000
// bit_ceil( 0b00001001 ) = 0b00010000
template <class T>
CUTE_HOST_DEVICE constexpr

T

bit_ceil(T x) {
  return x == 0 ? T(1) : (T(1) << bit_width(x - 1));
}

// Largest integral power of two not greater than the given value
// bit_floor( 0b00000000 ) = 0b00000000
// bit_floor( 0b00000001 ) = 0b00000001
// bit_floor( 0b00000010 ) = 0b00000010
// bit_floor( 0b00000011 ) = 0b00000010
// bit_floor( 0b00000100 ) = 0b00000100
// bit_floor( 0b00000101 ) = 0b00000100
// bit_floor( 0b00000110 ) = 0b00000100
// bit_floor( 0b00000111 ) = 0b00000100
// bit_floor( 0b00001000 ) = 0b00001000
// bit_floor( 0b00001001 ) = 0b00001000
template <class T>
CUTE_HOST_DEVICE constexpr

T

bit_floor(T x) {
  return x == 0 ? 0 : (T(1) << (bit_width(x) - 1));
}

template <class T>
CUTE_HOST_DEVICE constexpr T rotl(T x, int s);
template <class T>
CUTE_HOST_DEVICE constexpr T rotr(T x, int s);

// Computes the result of circular bitwise left-rotation
template <class T>
CUTE_HOST_DEVICE constexpr

T

rotl(T x, int s) {
  constexpr int N = numeric_limits<T>::digits;
  return static_cast<T>(s == 0 ? x : s > 0 ? (x << s) | (x >> (N - s)) : rotr(x, -s));
}

// Computes the result of circular bitwise right-rotation
template <class T>
CUTE_HOST_DEVICE constexpr

T

rotr(T x, int s) {
  constexpr int N = numeric_limits<T>::digits;
  return static_cast<T>(s == 0 ? x : s > 0 ? (x >> s) | (x << (N - s)) : rotl(x, -s));
}

// Counts the number of consecutive 0 bits, starting from the most significant bit
// countl_zero( 0b00000000 ) = 8
// countl_zero( 0b11111111 ) = 0
// countl_zero( 0b00011100 ) = 3
template <class T>
CUTE_HOST_DEVICE constexpr

T

countl_zero(T x) {
  return numeric_limits<T>::digits - bit_width(x);
}

// Counts the number of consecutive 1 bits, starting from the most significant bit
// countl_one( 0b00000000 ) = 0
// countl_one( 0b11111111 ) = 8
// countl_one( 0b11100011 ) = 3
template <class T>
CUTE_HOST_DEVICE constexpr

T

countl_one(T x) {
  return countl_zero(~x);
}

// Counts the number of consecutive 0 bits, starting from the least significant bit
// countr_zero( 0b00000000 ) = 8
// countr_zero( 0b11111111 ) = 0
// countr_zero( 0b00011100 ) = 2
template <class T>
CUTE_HOST_DEVICE constexpr

T

countr_zero(T x) {
  return x == 0 ? numeric_limits<T>::digits : bit_width(T(x & T(-x))) - 1;  // bit_width of the LSB
}

// Counts the number of consecutive 1 bits, starting from the least significant bit
// countr_one( 0b00000000 ) = 0
// countr_one( 0b11111111 ) = 8
// countr_one( 0b11100011 ) = 2
template <class T>
CUTE_HOST_DEVICE constexpr

T

countr_one(T x) {
  return countr_zero(~x);
}

// Counts the number of 1 bits in an unsigned integer
// popcount( 0b00000000 ) = 0
// popcount( 0b11111111 ) = 8
// popcount( 0b00011101 ) = 4
template <class T>
CUTE_HOST_DEVICE constexpr

int

popcount(T x) {
  int c = 0;
  while (x) {
    ++c;
    x &= x - 1; // clear the least significant bit set
  }
  return c;
}

//
// Custom operations
//

// Computes the result of bitwise left-shift
template <class T>
CUTE_HOST_DEVICE constexpr

T

shiftl(T x, int s) {
  return s >= 0 ? (x << s) : (x >> -s);
}

// Computes the result of bitwise right-shift
template <class T>
CUTE_HOST_DEVICE constexpr

T

shiftr(T x, int s) {
  return s >= 0 ? (x >> s) : (x << -s);
}

// Safe divide
// @pre t % u == 0
// @result t / u
template <class T, class U,
          __CUTE_REQUIRES(is_std_integral<T>::value &&
                          is_std_integral<U>::value)>
CUTE_HOST_DEVICE constexpr

auto

safe_div(T const& t, U const& u) {
  //assert(t % u == 0);
  return t / u;
}

/**

 * log2 computation

 */

template <class T>
CUTE_HOST_DEVICE constexpr

int32_t

log_2(T x) {
  assert(x > 0);
  static_assert(is_unsigned<T>::value, "Only to be used for unsigned integral types.");
  return static_cast<int32_t>(bit_width(x)) - 1;
}

} // namespace cute