half_tree_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, 1>)
class HalfTreeDpf {
public:
  Prg prg;
  int4 hash_key;
 
  struct __align__(32) Cw {
    int4 s;
    bool extra;
  };
  static_assert(sizeof(Cw) == 32);
 
  __host__ __device__ void Gen(Cw cws[], int4 &ocw, const int4 s0s[2], In a, int4 b_buf) {
    b_buf = util::SetLsb(b_buf, false);
 
    // Initialize: node0 has t=0, node1 has t=1
    int4 node0 = util::SetLsb(s0s[0], false);
    int4 node1 = util::SetLsb(s0s[1], true);
    int4 delta = util::Xor(node0, node1);  // LSB = 0^1 = 1
 
    // Levels 1 to n-1 (index i = 0 to in_bits-2)
    for (int i = 0; i < in_bits - 1; ++i) {
      int4 h0 = prg.Gen(util::Xor(hash_key, node0))[0];
      int4 h1 = prg.Gen(util::Xor(hash_key, node1))[0];
 
      bool a_bit = (a >> (in_bits - 1 - i)) & 1;
 
      // CW = h0 ^ h1 ^ (!a_bit ? delta : 0)
      // When a_bit=0 (go left): non-alpha is right, CW = h0^h1^delta makes right0=right1
      // When a_bit=1 (go right): non-alpha is left, CW = h0^h1 makes left0=left1
      int4 cw = util::Xor(h0, h1);
      if (!a_bit) cw = util::Xor(cw, delta);
 
      cws[i] = {cw, false};
 
      bool t0 = util::GetLsb(node0);
      bool t1 = util::GetLsb(node1);
 
      // node_b = h_b ^ (a_bit ? node_b : 0) ^ (t_b ? cw : 0)
      int4 zero4 = {0, 0, 0, 0};
      int4 ab_mask0 = a_bit ? node0 : zero4;
      int4 ab_mask1 = a_bit ? node1 : zero4;
      int4 t0_mask = t0 ? cw : zero4;
      int4 t1_mask = t1 ? cw : zero4;
 
      node0 = util::Xor(util::Xor(h0, ab_mask0), t0_mask);
      node1 = util::Xor(util::Xor(h1, ab_mask1), t1_mask);
 
      // delta = node0 ^ node1 for next level
      delta = util::Xor(node0, node1);
    }
 
    // Level n (last level, index i = in_bits-1)
    {
      bool a_n = (a >> 0) & 1;  // last bit of alpha
      bool t0 = util::GetLsb(node0);
      bool t1 = util::GetLsb(node1);
 
      // Hash with sigma in {0, 1}
      int4 h0_0 = prg.Gen(util::Xor(hash_key, util::SetLsb(node0, false)))[0];
      int4 h0_1 = prg.Gen(util::Xor(hash_key, util::SetLsb(node0, true)))[0];
      int4 h1_0 = prg.Gen(util::Xor(hash_key, util::SetLsb(node1, false)))[0];
      int4 h1_1 = prg.Gen(util::Xor(hash_key, util::SetLsb(node1, true)))[0];
 
      // Extract high (s) and low (t) parts
      int4 high0_0 = util::SetLsb(h0_0, false);
      bool low0_0 = util::GetLsb(h0_0);
      int4 high0_1 = util::SetLsb(h0_1, false);
      bool low0_1 = util::GetLsb(h0_1);
      int4 high1_0 = util::SetLsb(h1_0, false);
      bool low1_0 = util::GetLsb(h1_0);
      int4 high1_1 = util::SetLsb(h1_1, false);
      bool low1_1 = util::GetLsb(h1_1);
 
      // HCW corrects the non-alpha direction so both parties converge.
      // HCW = high{!a_n}_0 ^ high{!a_n}_1
      int4 HCW;
      if (a_n) HCW = util::Xor(high0_0, high1_0);
      else HCW = util::Xor(high0_1, high1_1);
 
      // LCW ensures:
      //   Alpha direction (sigma=a_n): low0 ^ low1 = 1 (exactly one adds ocw)
      //   Non-alpha direction (sigma=!a_n): low0 ^ low1 = 0 (cancels)
      // LCW_0 = low{0}_0 ^ low{0}_1 ^ !a_n
      // LCW_1 = low{1}_0 ^ low{1}_1 ^ a_n
      bool LCW_0 = low0_0 ^ low1_0 ^ !a_n;
      bool LCW_1 = low0_1 ^ low1_1 ^ a_n;
 
      // Store CW_n. Both 16B halves written so the full 32B slot gets unmasked
      // stores: {s, LCW_1} alone leaves padding undefined, NVCC emits a masked
      // 256-bit store, and the partial-sector write forces a write-allocate
      // round trip on GPU (measured ~2x Gen time for the same pattern).
      int4 *cw_slot = reinterpret_cast<int4 *>(&cws[in_bits - 1]);
      cw_slot[0] = util::SetLsb(HCW, LCW_0);
      cw_slot[1] = make_int4(LCW_1, 0, 0, 0);
 
      // Compute leaf for each party
      // leaf_b = (a_n ? high{1}_b||low{1}_b : high{0}_b||low{0}_b)
      int4 leaf0, leaf1;
      if (a_n) {
        leaf0 = util::SetLsb(high0_1, low0_1);
        leaf1 = util::SetLsb(high1_1, low1_1);
      } else {
        leaf0 = util::SetLsb(high0_0, low0_0);
        leaf1 = util::SetLsb(high1_0, low1_0);
      }
 
      // Apply CW correction: if t_b: leaf_b ^= SetLsb(HCW, lcw_an)
      bool lcw_an = a_n ? LCW_1 : LCW_0;
      int4 leaf_cw = util::SetLsb(HCW, lcw_an);
      if (t0) leaf0 = util::Xor(leaf0, leaf_cw);
      if (t1) leaf1 = util::Xor(leaf1, leaf_cw);
 
      // Output CW: v_cw = Group::From(b_buf) + (-Group::From(SetLsb(leaf0,false))) + Group::From(SetLsb(leaf1,false))
      auto v_cw =
          Group::From(b_buf) + (-Group::From(util::SetLsb(leaf0, false))) + Group::From(util::SetLsb(leaf1, false));
      if (util::GetLsb(leaf1)) v_cw = -v_cw;
      ocw = v_cw.Into();
    }
  }
 
  __host__ __device__ int4 Eval(bool b, int4 s0, const Cw cws[], int4 ocw, In x) {
    int4 node = util::SetLsb(s0, b);
 
    // Levels 1 to n-1 (index i = 0 to in_bits-2)
    for (int i = 0; i < in_bits - 1; ++i) {
      bool x_bit = (x >> (in_bits - 1 - i)) & 1;
      bool t = util::GetLsb(node);
 
      int4 h = prg.Gen(util::Xor(hash_key, node))[0];
 
      int4 zero4 = {0, 0, 0, 0};
      int4 xb_mask = x_bit ? node : zero4;
      int4 t_mask = t ? cws[i].s : zero4;
 
      node = util::Xor(util::Xor(h, xb_mask), t_mask);
    }
 
    // Level n (last level)
    {
      bool x_n = (x >> 0) & 1;
      bool t = util::GetLsb(node);
 
      int4 h = prg.Gen(util::Xor(hash_key, util::SetLsb(node, x_n)))[0];
 
      // Unpack CW_n
      int4 hcw = util::SetLsb(cws[in_bits - 1].s, false);
      bool lcw_xn;
      if (x_n) lcw_xn = cws[in_bits - 1].extra;
      else lcw_xn = util::GetLsb(cws[in_bits - 1].s);
 
      int4 high = util::SetLsb(h, false);
      bool low = util::GetLsb(h);
 
      if (t) {
        high = util::Xor(high, hcw);
        low = low ^ lcw_xn;
      }
 
      auto y = Group::From(high);
      if (low) y = y + Group::From(ocw);
      if (b) y = -y;
 
      return y.Into();
    }
  }
 
  void EvalAll(bool b, int4 s0, const Cw cws[], int4 ocw, int4 ys[]) {
    int4 node = util::SetLsb(s0, b);
 
    assert(in_bits < sizeof(size_t) * 8);
 
    int par_depth_ = util::ResolveParDepth(par_depth);
 
    if constexpr (in_bits == 1) {
      // Only level n (last level), no tree traversal
#pragma omp parallel
#pragma omp single
      EvalLastLevel(b, node, cws, ocw, ys);
      return;
    }
 
    // Phase 1: tree traversal for levels 1..n-1, stores nodes at level n-1
    // We use ys[] as scratch space for intermediate nodes.
    // After phase 1, ys[0..2^(in_bits-1)-1] hold the level n-1 nodes (packed s||t).
    size_t num_leaves = 1ULL << (in_bits - 1);
 
    // Recursive tree traversal
#pragma omp parallel
#pragma omp single
    EvalTree(node, cws, ys, 0, num_leaves, 0, par_depth_);
 
    // Phase 2: level n + output conversion
    int4 hcw = util::SetLsb(cws[in_bits - 1].s, false);
    bool lcw_0 = util::GetLsb(cws[in_bits - 1].s);
    bool lcw_1 = cws[in_bits - 1].extra;
    auto ocw_group = Group::From(ocw);
 
    // Iterate backward to avoid overwriting unprocessed parent nodes.
    for (size_t j = num_leaves; j-- > 0;) {
      ConvertLastLevel(b, ys[j], hcw, lcw_0, lcw_1, ocw_group, ys[2 * j], ys[2 * j + 1]);
    }
  }
 
private:
  void EvalTree(int4 node, const Cw cws[], int4 ys[], size_t l, size_t r, int i, int par_depth_) {
    // i is the level index (0-based), we traverse levels 0..in_bits-2
    // At level in_bits-1, we store the node
    if (i == in_bits - 1) {
      assert(l + 1 == r);
      ys[l] = node;
      return;
    }
 
    bool t = util::GetLsb(node);
    int4 h = prg.Gen(util::Xor(hash_key, node))[0];
 
    int4 zero4 = {0, 0, 0, 0};
    int4 t_mask = t ? cws[i].s : zero4;
 
    // Left child: left = H_S(parent) ^ (t ? cw : 0)
    int4 left = util::Xor(h, t_mask);
    // Right child: right = left ^ parent
    int4 right = util::Xor(left, node);
 
    size_t mid = (l + r) / 2;
 
    if (i < par_depth_) {
#pragma omp task
      EvalTree(left, cws, ys, l, mid, i + 1, par_depth_);
#pragma omp task
      EvalTree(right, cws, ys, mid, r, i + 1, par_depth_);
#pragma omp taskwait
    } else {
      EvalTree(left, cws, ys, l, mid, i + 1, par_depth_);
      EvalTree(right, cws, ys, mid, r, i + 1, par_depth_);
    }
  }
 
  void EvalLastLevel(bool b, int4 node, const Cw cws[], int4 ocw, int4 ys[]) {
    int4 hcw = util::SetLsb(cws[0].s, false);
    bool lcw_0 = util::GetLsb(cws[0].s);
    bool lcw_1 = cws[0].extra;
    ConvertLastLevel(b, node, hcw, lcw_0, lcw_1, Group::From(ocw), ys[0], ys[1]);
  }
 
  void ConvertLastLevel(
      bool b, int4 parent, int4 hcw, bool lcw_0, bool lcw_1, Group ocw_group, int4 &y0_out, int4 &y1_out) {
    bool t_parent = util::GetLsb(parent);
 
    int4 h0 = prg.Gen(util::Xor(hash_key, util::SetLsb(parent, false)))[0];
    int4 h1 = prg.Gen(util::Xor(hash_key, util::SetLsb(parent, true)))[0];
 
    int4 high0 = util::SetLsb(h0, false);
    bool low0 = util::GetLsb(h0);
    int4 high1 = util::SetLsb(h1, false);
    bool low1 = util::GetLsb(h1);
 
    if (t_parent) {
      high0 = util::Xor(high0, hcw);
      low0 = low0 ^ lcw_0;
      high1 = util::Xor(high1, hcw);
      low1 = low1 ^ lcw_1;
    }
 
    auto y0 = Group::From(high0);
    if (low0) y0 = y0 + ocw_group;
    if (b) y0 = -y0;
 
    auto y1 = Group::From(high1);
    if (low1) y1 = y1 + ocw_group;
    if (b) y1 = -y1;
 
    y0_out = y0.Into();
    y1_out = y1.Into();
  }
};
 
}  // namespace fss