uint128.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 <cassert>
#include <cstdint>
 
namespace fss::group {

struct Uint128Modulus {
  uint64_t low;  
  uint64_t high;  
};

template <Uint128Modulus mod = Uint128Modulus{0, uint64_t{1} << 63}>
  requires(((static_cast<__uint128_t>(mod.high) << 64) | mod.low) > 0 &&
      ((static_cast<__uint128_t>(mod.high) << 64) | mod.low) <= static_cast<__uint128_t>(1) << 127)
class Uint128 {
public:
  // Groupable interface.

  __host__ __device__ Uint128 operator+(Uint128 rhs) const {
    if (val_ >= kMod - rhs.val_) return {val_ + rhs.val_ - kMod};
    else return {val_ + rhs.val_};
  }

  __host__ __device__ Uint128 operator-() const {
    if (val_ == 0) return {0};
    else return {kMod - val_};
  }

  __host__ __device__ Uint128() : val_(0) {}

  __host__ __device__ static Uint128 From(int4 buf) {
    assert((buf.w & 1) == 0);
 
    __uint128_t value = static_cast<__uint128_t>(static_cast<unsigned int>(buf.x)) |
        static_cast<__uint128_t>(static_cast<unsigned int>(buf.y)) << 32 |
        static_cast<__uint128_t>(static_cast<unsigned int>(buf.z)) << 64 |
        // The LSB of buf.w is the clamped bit
        static_cast<__uint128_t>(static_cast<unsigned int>(buf.w) >> 1) << 96;
 
    return {value % kMod};
  }

  __host__ __device__ int4 Into() const {
    int4 buf = {0, 0, 0, 0};
    buf.x = static_cast<int>(val_ & 0xffffffff);
    buf.y = static_cast<int>((val_ >> 32) & 0xffffffff);
    buf.z = static_cast<int>((val_ >> 64) & 0xffffffff);
    // The LSB of buf.w is the clamped bit
    buf.w = static_cast<int>((val_ >> 96) << 1);
    return buf;
  }
 
  // Additional public API.
  __host__ __device__ __uint128_t Value() const { return val_; }
 
private:
  static constexpr __uint128_t kMod = (static_cast<__uint128_t>(mod.high) << 64) | mod.low;
 
  __uint128_t val_;
 
  __host__ __device__ Uint128(__uint128_t v) : val_(v) {}
};
static_assert(Groupable<Uint128<>>);
static_assert(Groupable<Uint128<Uint128Modulus{.low = 0xffffffffffffffff, .high = 0x3fffffff}>>);
 
}  // namespace fss::group