dpf.cuh

Go to the documentation of this file.

// SPDX-License-Identifier: Apache-2.0
#pragma once
#include <cuda_runtime.h>
#include <type_traits>
#include <cstddef>
#include <cassert>
#include <omp.h>
#include <fss/group.cuh>
#include <fss/prg.cuh>
#include <fss/util.cuh>
 
namespace fss {
 
template <int in_bits, typename Group, typename Prg, typename In = uint, int par_depth = -1>
  requires((std::is_unsigned_v<In> || std::is_same_v<In, __uint128_t>) && in_bits <= sizeof(In) * 8 &&
      Groupable<Group> && Prgable<Prg, 2>)
class Dpf {
public:
  Prg prg;
 
  struct __align__(32) Cw {
    int4 s;
    bool tr;
  };
  // For only 1 and aligned memory access on GPU
  static_assert(sizeof(Cw) == 32);
 
  __host__ __device__ void Gen(Cw cws[], const int4 s0s[2], In a, int4 b_buf) {
    int4 s0 = s0s[0];
    s0 = util::SetLsb(s0, false);
    int4 s1 = s0s[1];
    s1 = util::SetLsb(s1, false);
    bool t0 = false;
    bool t1 = true;
    b_buf = util::SetLsb(b_buf, false);
 
    for (int i = 0; i < in_bits; ++i) {
      auto [s0l, s0r] = prg.Gen(s0);
      auto [s1l, s1r] = prg.Gen(s1);
 
      bool t0l = util::GetLsb(s0l);
      s0l = util::SetLsb(s0l, false);
      bool t0r = util::GetLsb(s0r);
      s0r = util::SetLsb(s0r, false);
      bool t1l = util::GetLsb(s1l);
      s1l = util::SetLsb(s1l, false);
      bool t1r = util::GetLsb(s1r);
      s1r = util::SetLsb(s1r, false);
 
      bool a_bit = (a >> (in_bits - 1 - i)) & 1;
 
      int4 s_cw;
      if (!a_bit) s_cw = util::Xor(s0r, s1r);
      else s_cw = util::Xor(s0l, s1l);
 
      bool tl_cw = t0l ^ t1l ^ a_bit ^ 1;
      bool tr_cw = t0r ^ t1r ^ a_bit;
 
      if (!a_bit) {
        s0 = s0l;
        if (t0) s0 = util::Xor(s0, s_cw);
        s1 = s1l;
        if (t1) s1 = util::Xor(s1, s_cw);
 
        if (t0) t0 = t0l ^ tl_cw;
        else t0 = t0l;
        if (t1) t1 = t1l ^ tl_cw;
        else t1 = t1l;
      } else {
        s0 = s0r;
        if (t0) s0 = util::Xor(s0, s_cw);
        s1 = s1r;
        if (t1) s1 = util::Xor(s1, s_cw);
 
        if (t0) t0 = t0r ^ tr_cw;
        else t0 = t0r;
        if (t1) t1 = t1r ^ tr_cw;
        else t1 = t1r;
      }
 
      s_cw = util::SetLsb(s_cw, tl_cw);
      // Write both 16B halves so the full 32B slot is covered by unmasked stores.
      // {s_cw, tr_cw} alone leaves the padding undefined: NVCC emits a masked
      // 256-bit store, and the partial-sector write forces a write-allocate
      // round trip per level on GPU (measured ~2x Gen time).
      int4 *cw_slot = reinterpret_cast<int4 *>(&cws[i]);
      cw_slot[0] = s_cw;
      cw_slot[1] = make_int4(tr_cw, 0, 0, 0);
    }
 
    auto v_cw_np1 = Group::From(b_buf) + (-Group::From(s0)) + Group::From(s1);
    if (t1) v_cw_np1 = -v_cw_np1;
    cws[in_bits] = {v_cw_np1.Into(), false};
  }
 
  __host__ __device__ int4 Eval(bool b, int4 s0, const Cw cws[], In x) {
    int4 s = s0;
    s = util::SetLsb(s, false);
    bool t = b;
 
    for (int i = 0; i < in_bits; ++i) {
      Cw cw = cws[i];
      int4 s_cw = cw.s;
      bool tl_cw = util::GetLsb(s_cw);
      s_cw = util::SetLsb(s_cw, false);
      bool tr_cw = cw.tr;
 
      auto [sl, sr] = prg.Gen(s);
 
      bool tl = util::GetLsb(sl);
      sl = util::SetLsb(sl, false);
      bool tr = util::GetLsb(sr);
      sr = util::SetLsb(sr, false);
 
      if (t) {
        sl = util::Xor(sl, s_cw);
        sr = util::Xor(sr, s_cw);
        tl = tl ^ tl_cw;
        tr = tr ^ tr_cw;
      }
 
      bool x_bit = (x >> (in_bits - 1 - i)) & 1;
 
      if (!x_bit) {
        s = sl;
        t = tl;
      } else {
        s = sr;
        t = tr;
      }
    }
 
    auto y = Group::From(s);
    int4 v_cw_np1 = cws[in_bits].s;
    assert((v_cw_np1.w & 1) == 0);
    if (t) y = y + Group::From(v_cw_np1);
    if (b) y = -y;
 
    return y.Into();
  }
 
  void EvalAll(bool b, int4 s0, const Cw cws[], int4 ys[]) {
    int4 st = s0;
    bool t = b;
    st = util::SetLsb(st, t);
 
    assert(in_bits < sizeof(size_t) * 8);
    size_t l = 0;
    size_t r = 1ULL << in_bits;
    int i = 0;
 
    int par_depth_ = util::ResolveParDepth(par_depth);
 
#pragma omp parallel
#pragma omp single
    EvalTree(b, st, cws, ys, l, r, i, par_depth_);
  }
 
private:
  void EvalTree(bool b, int4 st, const Cw cws[], int4 ys[], size_t l, size_t r, int i, int par_depth_) {
    bool t = util::GetLsb(st);
    int4 s = st;
    s = util::SetLsb(s, false);
 
    if (i == in_bits) {
      auto y = Group::From(s);
      int4 v_cw_np1 = cws[in_bits].s;
      assert((v_cw_np1.w & 1) == 0);
      if (t) y = y + Group::From(v_cw_np1);
      if (b) y = -y;
      assert(l + 1 == r);
      ys[l] = y.Into();
      return;
    }
 
    Cw cw = cws[i];
    int4 s_cw = cw.s;
    bool tl_cw = util::GetLsb(s_cw);
    s_cw = util::SetLsb(s_cw, false);
    bool tr_cw = cw.tr;
 
    auto [sl, sr] = prg.Gen(s);
 
    bool tl = util::GetLsb(sl);
    sl = util::SetLsb(sl, false);
    bool tr = util::GetLsb(sr);
    sr = util::SetLsb(sr, false);
 
    if (t) {
      sl = util::Xor(sl, s_cw);
      sr = util::Xor(sr, s_cw);
      tl = tl ^ tl_cw;
      tr = tr ^ tr_cw;
    }
 
    int4 stl = sl;
    stl = util::SetLsb(stl, tl);
    int4 str = sr;
    str = util::SetLsb(str, tr);
 
    size_t mid = (l + r) / 2;
 
    if (i < par_depth_) {
#pragma omp task
      EvalTree(b, stl, cws, ys, l, mid, i + 1, par_depth_);
#pragma omp task
      EvalTree(b, str, cws, ys, mid, r, i + 1, par_depth_);
#pragma omp taskwait
    } else {
      EvalTree(b, stl, cws, ys, l, mid, i + 1, par_depth_);
      EvalTree(b, str, cws, ys, mid, r, i + 1, par_depth_);
    }
  }
};
 
}  // namespace fss