uint.cuh

Go to the documentation of this file.

// SPDX-License-Identifier: Apache-2.0
#pragma once
#include <fss/group.cuh>
#include <cuda_runtime.h>
#include <type_traits>
#include <cassert>
 
namespace fss::group {
 
template <typename T, T mod = 0>
  requires((std::is_unsigned_v<T> || std::is_same_v<T, __uint128_t>) && sizeof(T) <= 16 &&
      (sizeof(T) < 16 || (mod > 0 && mod <= static_cast<T>(1) << 127)))
struct Uint {
  T val;
 
  __host__ __device__ Uint operator+(Uint rhs) const {
    if constexpr (mod == 0) return {static_cast<T>(val + rhs.val)};
 
    if (val >= mod - rhs.val) return {static_cast<T>(val + rhs.val - mod)};
    else return {static_cast<T>(val + rhs.val)};
  }
 
  __host__ __device__ Uint operator-() const {
    if constexpr (mod == 0) return {static_cast<T>(-val)};
 
    if (val == 0) return {0};
    else return {static_cast<T>(mod - val)};
  }
 
  __host__ __device__ Uint() : val(0) {}
 
  __host__ __device__ static Uint From(int4 buf) {
    assert((buf.w & 1) == 0);
 
    T val = 0;
    if constexpr (sizeof(T) < 4) val = buf.x & ((1 << 8 * sizeof(T)) - 1);
    // Cast to unsigned int first to prevent sign extension when promoting to larger types
    else if constexpr (sizeof(T) == 4) val = static_cast<unsigned int>(buf.x);
    else if constexpr (sizeof(T) == 8)
      val = static_cast<T>(static_cast<unsigned int>(buf.x)) | static_cast<T>(static_cast<unsigned int>(buf.y)) << 32;
    else if constexpr (sizeof(T) == 16)
      val = static_cast<T>(static_cast<unsigned int>(buf.x)) | static_cast<T>(static_cast<unsigned int>(buf.y)) << 32 |
          static_cast<T>(static_cast<unsigned int>(buf.z)) << 64 |
          // For uint128, LSB of buf.w is the clamped bit
          static_cast<T>(static_cast<unsigned int>(buf.w) >> 1) << 96;
    else __builtin_unreachable();
 
    if constexpr (mod > 0) val %= mod;
 
    return {val};
  }
 
  __host__ __device__ int4 Into() const {
    int4 buf = {0, 0, 0, 0};
    if constexpr (sizeof(T) <= 4) buf.x = static_cast<int>(val);
    else if constexpr (sizeof(T) == 8) {
      buf.x = static_cast<int>(val & 0xffffffff);
      buf.y = static_cast<int>(val >> 32);
    } else if constexpr (sizeof(T) == 16) {
      buf.x = static_cast<int>(val & 0xffffffff);
      buf.y = static_cast<int>((val >> 32) & 0xffffffff);
      buf.z = static_cast<int>((val >> 64) & 0xffffffff);
      // For uint128, its LSB is the clamped bit
      buf.w = static_cast<int>((val >> 96) << 1);
    } else __builtin_unreachable();
    return buf;
  }
 
private:
  __host__ __device__ Uint(T v) : val(v) {}
};
static_assert(Groupable<Uint<uint8_t>>);
static_assert(Groupable<Uint<uint16_t>>);
static_assert(Groupable<Uint<uint32_t>>);
static_assert(Groupable<Uint<uint64_t>>);
static_assert(Groupable<Uint<__uint128_t, static_cast<__uint128_t>(1) << 127>>);
 
}  // namespace fss::group