grotto_dcf.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/dpf.cuh>
#include <fss/group/bytes.cuh>
#include <fss/prg.cuh>
#include <fss/util.cuh>
 
namespace fss {
 
template <int in_bits, 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 && Prgable<Prg, 2>)
class GrottoDcf {
  using DpfType = Dpf<in_bits, group::Bytes, Prg, In, par_depth>;
 
public:
  using Cw = typename DpfType::Cw;
  Prg prg;
 
  __host__ __device__ void Gen(Cw cws[], const int4 s0s[2], In a) {
    DpfType dpf{prg};
    int4 beta = {0, 0, 0, 0};
    dpf.Gen(cws, s0s, a, beta);
  }
 
  struct ParityTree {
    bool *p;
    bool b;
  };
 
  void Preprocess(ParityTree &pt, int4 s0, const Cw cws[]) {
    constexpr size_t N = 1ULL << in_bits;
 
    // Phase 1: expand tree, write leaf control bits to pt.p[N-1 .. 2N-2]
    ExpandTree(pt.b, s0, cws, pt.p + (N - 1));
 
    // Phase 2a: build parity segment tree bottom-up
    for (size_t j = N - 2; j < N - 1; --j) {
      pt.p[j] = pt.p[2 * j + 1] ^ pt.p[2 * j + 2];
    }
  }
 
  __host__ __device__ static bool Eval(const ParityTree &pt, In x) {
    constexpr size_t N = 1ULL << in_bits;
    In e = static_cast<In>(x) + 1;
 
    // e == 0 means x + 1 overflowed, i.e., e = N (entire domain)
    if (e == 0 || e == N) return pt.p[0];
 
    bool pi = false;
    size_t cur = 0;
    for (int i = 0; i < in_bits; ++i) {
      bool e_bit = (e >> (in_bits - 1 - i)) & 1;
      if (e_bit) {
        pi ^= pt.p[2 * cur + 1];
        cur = 2 * cur + 2;
      } else {
        cur = 2 * cur + 1;
      }
    }
    return pi;
  }
 
  void EvalAll(bool b, int4 s0, const Cw cws[], bool ys[]) {
    constexpr size_t N = 1ULL << in_bits;
 
    // Phase 1: expand tree to get leaf control bits into ys[]
    ExpandTree(b, s0, cws, ys);
 
    // Phase 2b: prefix-sum scan (running XOR)
    // ys[x] currently holds leaf x's control bit.
    // Transform to: ys[x] = XOR of control bits [0..x] = share of 1[alpha <= x].
    for (size_t x = 1; x < N; ++x) {
      ys[x] = ys[x] ^ ys[x - 1];
    }
  }
 
private:
  void ExpandTree(bool b, int4 s0, const Cw cws[], bool t[]) {
    int4 st = s0;
    st = util::SetLsb(st, b);
 
    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
    ExpandTreeRec(st, cws, t, l, r, i, par_depth_);
  }
 
  void ExpandTreeRec(int4 st, const Cw cws[], bool t[], size_t l, size_t r, int i, int par_depth_) {
    bool tc = util::GetLsb(st);
    int4 s = st;
    s = util::SetLsb(s, false);
 
    if (i == in_bits) {
      assert(l + 1 == r);
      t[l] = tc;
      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 (tc) {
      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
      ExpandTreeRec(stl, cws, t, l, mid, i + 1, par_depth_);
#pragma omp task
      ExpandTreeRec(str, cws, t, mid, r, i + 1, par_depth_);
#pragma omp taskwait
    } else {
      ExpandTreeRec(stl, cws, t, l, mid, i + 1, par_depth_);
      ExpandTreeRec(str, cws, t, mid, r, i + 1, par_depth_);
    }
  }
};
 
}  // namespace fss