teems/Enclave/OblivAlgs/TightCompaction_v2.tcc

#ifndef __NOP_TIGHTCOMPACTION_V2_TCC__
#define __NOP_TIGHTCOMPACTION_V2_TCC__

#include "pthread.h"

/*
   TightCompaction (Order Preserving Tight Compaction):

   Order Preserving TightCompaction can take an input array of blocks of
   block_size bytes each, and an array of "marked" elements with ones at the
   indices corresponding to the blocks that need to be compacted.
   It returns the input array TightCompact-ed, i.e. all the real blocks are
   moved to the start of the array, or compacted to an input index
   (with wraparound)

   This is the Oblivious Recusrive Compaction (ORCompact) algorithm from
   our CCS 2022 paper "Fast Fully Oblivious Compaction and Shuffling".

 */

template <OSwap_Style oswap_style>
void TightCompact_2power(unsigned char *buf, size_t N, size_t block_size, size_t offset, bool *selected) {
  // Compute counts of selected items at power-of-two intervals
  FOAV_SAFE2_CNTXT(TC_2power_summing_selected_count, N, block_size)
  uint32_t *selected_count = NULL;
  FOAV_SAFE_CNTXT(TC_2power, TC_PRECOMPUTE_COUNTS)
  if (TC_PRECOMPUTE_COUNTS) {
    // Allocate array to hold counts
    selected_count = new uint32_t[N+1];
    selected_count[0] = 0;
    // Compute cumulative counts
    for (size_t i=0; i<N; i++){
      FOAV_SAFE2_CNTXT(TC_2power_summing_selected_count, i, N)
      selected_count[i+1] = selected[i] + selected_count[i];
    }
    TightCompact_2power_inner<oswap_style>(buf, N, block_size, offset, selected, selected_count);
    delete[] selected_count;
  } else {
    TightCompact_2power_inner<oswap_style>(buf, N, block_size, offset, selected, selected_count);
  }
}

template <OSwap_Style oswap_style>
void TightCompact_2power_inner(unsigned char *buf, size_t N, size_t block_size, size_t offset, bool *selected, uint32_t *selected_count) {
  FOAV_SAFE2_CNTXT(TC_inner_base_cases_of_recursion, N, block_size)
  if (N==1) {
    return;
  }
  if (N==2) {
    bool swap = (!selected[0] & selected[1]) ^ offset;
    oswap_buffer<oswap_style>(buf, buf+block_size, block_size, swap);
    return;
  }

  // Number of selected items in left half
  size_t m1;
  FOAV_SAFE_CNTXT(TC_2power, TC_PRECOMPUTE_COUNTS)
  if (TC_PRECOMPUTE_COUNTS) {
    m1 = selected_count[N/2] - selected_count[0];
  } else {
    m1=0;
    FOAV_SAFE_CNTXT(TC_2power, N)
    for(size_t i=0; i<N/2; i++){
      FOAV_SAFE_CNTXT(TC_2power, i)
      m1+=selected[i];
    }
  }

  size_t offset_mod = (offset & ((N/2)-1));
  size_t offset_m1_mod = (offset+m1) & ((N/2)-1);
  bool offset_right = (offset >= N/2);
  bool left_wrapped = ((offset_mod + m1) >= (N/2));

  TightCompact_2power_inner<oswap_style>(buf, N/2, block_size, offset_mod, selected, selected_count);
  TightCompact_2power_inner<oswap_style>(buf + ((N/2)*block_size), N/2, block_size, offset_m1_mod, (selected + (N/2)), selected_count + N/2);

  unsigned char *buf1_ptr = buf, *buf2_ptr = (buf + (N/2)*block_size);
  FOAV_SAFE_CNTXT(TC_2power_inner, TC_OPT_SWAP_FLAG)
  if (TC_OPT_SWAP_FLAG) {
    bool swap_flag = left_wrapped ^ offset_right;
    size_t num_swap = N/2;
    FOAV_SAFE2_CNTXT(TC_2power_inner, num_swap, block_size)
    for(size_t i=0; i<num_swap; i++){
      FOAV_SAFE2_CNTXT(TC_2power_inner_N/2_swaps, i, num_swap)
      swap_flag = swap_flag ^ (i == offset_m1_mod);
      oswap_buffer<oswap_style>(buf1_ptr, buf2_ptr, block_size, swap_flag);
      buf1_ptr+=block_size;
      buf2_ptr+=block_size;
      FOAV_SAFE2_CNTXT(TC_2power_inner, num_swap, block_size)
    }
  } else {
    FOAV_SAFE_CNTXT(TC_2power_inner, N)
    for(size_t i=0; i<N/2; i++){
      FOAV_SAFE_CNTXT(TC_2power_inner, i)
      bool swap_flag = (i>=offset_m1_mod) ^ left_wrapped ^ offset_right;
      oswap_buffer<oswap_style>(buf1_ptr, buf2_ptr, block_size, swap_flag);
      buf1_ptr+=block_size;
      buf2_ptr+=block_size;
      FOAV_SAFE2_CNTXT(TC_2power_inner, i, N)
    }
  }
}

struct TightCompact_2power_inner_parallel_args {
    unsigned char *buf;
    size_t N, block_size, offset;
    bool *selected;
    uint32_t *selected_count;
    size_t nthreads;
};

template <OSwap_Style oswap_style>
static void* TightCompact_2power_inner_parallel_launch(void *voidargs) {
    struct TightCompact_2power_inner_parallel_args *args =
        (TightCompact_2power_inner_parallel_args *)voidargs;
    TightCompact_2power_inner_parallel<oswap_style>(args->buf, args->N,
        args->block_size, args->offset, args->selected, args->selected_count,
        args->nthreads);
    return NULL;
}

struct oswap_range_args {
    size_t block_size;
    size_t swap_start, swap_end;
    size_t offset_m1_mod;
    unsigned char *buf1, *buf2;
    bool swap_flag;
};

template <OSwap_Style oswap_style>
static void* oswap_range(void *voidargs) {
    struct oswap_range_args *args = (oswap_range_args*)voidargs;
    size_t block_size = args->block_size;
    size_t swap_start = args->swap_start;
    size_t swap_end = args->swap_end;
    size_t offset_m1_mod = args->offset_m1_mod;
    unsigned char *buf1 = args->buf1 + swap_start*block_size;
    unsigned char *buf2 = args->buf2 + swap_start*block_size;
    bool swap_flag = args->swap_flag;
    FOAV_SAFE2_CNTXT(oswap_range, block_size, swap_start)
    FOAV_SAFE_CNTXT(oswap_range, swap_end)
    //FOAV_SAFE_CNTXT(oswap_range, &OSWAP_COUNTER)
    //printf("start oswap range %p %lu %lu\n", buf1, swap_start, swap_end);
    for(size_t i=swap_start; i<swap_end; i++){
      FOAV_SAFE2_CNTXT(oswap_range, i, swap_end)
      oswap_buffer<oswap_style>(buf1, buf2, block_size, swap_flag ^ (i >= offset_m1_mod));
      buf1+=block_size;
      buf2+=block_size;
      //FOAV_SAFE_CNTXT(oswap_range, &OSWAP_COUNTER)
      FOAV_SAFE2_CNTXT(oswap_range, swap_end, block_size)
      FOAV_SAFE_CNTXT(oswap_range, i)
    }
    //printf("end oswap range %p %lu %lu\n", buf1, swap_start, swap_end);
    return NULL;
}


template <OSwap_Style oswap_style>
void TightCompact_2power_inner_parallel(unsigned char *buf, size_t N, size_t block_size, size_t offset, bool *selected, uint32_t *selected_count, size_t nthreads) {
  FOAV_SAFE_CNTXT(TC_inner_base_cases_of_recursion, g_thread_id)
  FOAV_SAFE2_CNTXT(TC_inner_base_cases_of_recursion, N, block_size)
  FOAV_SAFE_CNTXT(TC_inner_base_cases_of_recursion, nthreads)
  if (nthreads <= 1) {
    FOAV_SAFE2_CNTXT(TC_inner_base_cases_of_recursion, N, block_size)
    FOAV_SAFE_CNTXT(TC_inner_base_cases_of_recursion, nthreads)
    unsigned long start = printf_with_rtclock("Thread %u starting TightCompact_2power_inner(buf=%p, N=%lu, offset=%lu, nthreads=%lu)\n", g_thread_id, buf, N, offset, nthreads);
    TightCompact_2power_inner<oswap_style>(buf, N, block_size, offset, selected, selected_count);
    printf_with_rtclock_diff(start, "Thread %u ending TightCompact_2power_inner(buf=%p, N=%lu, offset=%lu, nthreads=%lu)\n", g_thread_id, buf, N, offset, nthreads);
    return;
  }
  FOAV_SAFE_CNTXT(TC_inner_base_cases_of_recursion, N)
  if (N==1) {
    return;
  }
  FOAV_SAFE_CNTXT(TC_inner_base_cases_of_recursion, N)
  if (N==2) {
    bool swap = (!selected[0] & selected[1]) ^ offset;
    oswap_buffer<oswap_style>(buf, buf+block_size, block_size, swap);
    return;
  }

  unsigned long start = printf_with_rtclock("Thread %u starting TightCompact_2power_inner_parallel(buf=%p, N=%lu, offset=%lu, nthreads=%lu)\n", g_thread_id, buf, N, offset, nthreads);
  // Number of selected items in left half
  size_t m1;
  m1 = selected_count[N/2] - selected_count[0];

  size_t offset_mod = (offset & ((N/2)-1));
  size_t offset_m1_mod = (offset+m1) & ((N/2)-1);
  bool offset_right = (offset >= N/2);
  bool left_wrapped = ((offset_mod + m1) >= (N/2));
  size_t lthreads = nthreads/2;
  size_t rthreads = nthreads - lthreads;

  threadid_t rightthreadid = g_thread_id + lthreads;
  /* Dispatch the right half to thread g_thread_id + lthreads; it will inherit threads
     g_thread_id + lthreads .. g_thread_id + nthreads-1. */
  struct TightCompact_2power_inner_parallel_args rightargs = {
    buf+ ((N/2)*block_size), N/2, block_size, offset_m1_mod, selected + N/2, selected_count + N/2, rthreads
  };
  threadpool_dispatch(rightthreadid,
    TightCompact_2power_inner_parallel_launch<oswap_style>,
    &rightargs);
  /* Do the left half ourselves (threads g_thread_id .. g_thread_id + lthreads-1) */
  TightCompact_2power_inner_parallel<oswap_style>(buf, N/2, block_size, offset_mod, selected, selected_count, lthreads);
  threadpool_join(rightthreadid, NULL);

  unsigned char *buf1_ptr = buf, *buf2_ptr = (buf + (N/2)*block_size);
  bool swap_flag = left_wrapped ^ offset_right;
  size_t num_swap = N/2;
  FOAV_SAFE2_CNTXT(TC_2power_inner, num_swap, block_size)

  oswap_range_args args[nthreads];
  size_t inc = num_swap / nthreads;
  size_t extra = num_swap % nthreads;
  size_t last = 0;
  for (size_t i=0; i<nthreads; ++i) {
    size_t next = last + inc + (i < extra);
    args[i] = { block_size, last, next, offset_m1_mod, buf1_ptr, buf2_ptr, swap_flag };
    last = next;
  }
  for (size_t i=0; i<nthreads-1; ++i) {
    threadpool_dispatch(g_thread_id+1+i, oswap_range<oswap_style>, args+i);
  }
  // Do the last section ourselves
  oswap_range<oswap_style>((void*)(args+nthreads-1));
  for (size_t i=0; i<nthreads-1; ++i) {
    FOAV_SAFE2_CNTXT(TC_2power_inner_parallel, i, nthreads)
    threadpool_join(g_thread_id+1+i, NULL);
  }
  printf_with_rtclock_diff(start, "Thread %u ending TightCompact_2power_inner_parallel(buf=%p, N=%lu, offset=%lu, nthreads=%lu)\n", g_thread_id, buf, N, offset, nthreads);
}

/*
  NOTE: TightCompact can only be invoked with offset 0.
  To invoke with a non-0 offset, use TightCompact_2power, with N = 2^x.
*/

template <OSwap_Style oswap_style>
void TightCompact(unsigned char *buf, size_t N, size_t block_size,
       bool *selected) {
  FOAV_SAFE2_CNTXT(TC_inner_base_cases_of_recursion, N, block_size)
  uint32_t *selected_count = NULL;
  if (TC_PRECOMPUTE_COUNTS) {

    // Allocate array to hold counts 
    try {
      selected_count = new uint32_t[N+1];
    } catch (std::bad_alloc&){
      printf("Allocating memory failed in TC\n");
    }
    selected_count[0] = 0;

    // Compute cumulative counts
    for (size_t i=0; i<N; i++){
      selected_count[i+1] = selected[i] + selected_count[i];
    }
    TightCompact_inner<oswap_style>(buf, N, block_size, selected, selected_count);
    delete[] selected_count;
  } else {
    TightCompact_inner<oswap_style>(buf, N, block_size, selected, selected_count);
  }
}

template <OSwap_Style oswap_style>
void TightCompact_parallel(unsigned char *buf, size_t N, size_t block_size,
       bool *selected, size_t nthreads) {
  FOAV_SAFE2_CNTXT(TC_inner_base_cases_of_recursion, N, block_size)
  uint32_t *selected_count = NULL;
  // Allocate array to hold counts
  try {
    selected_count = new uint32_t[N+1];
  } catch (std::bad_alloc&){
    printf("Allocating memory failed in TC\n");
  }
  selected_count[0] = 0;

  // Compute cumulative counts
  for (size_t i=0; i<N; i++){
    selected_count[i+1] = selected[i] + selected_count[i];
  }
  //printf("TightCompact_parallel(nthreads=%lu)\n", nthreads);
  TightCompact_inner_parallel<oswap_style>(buf, N, block_size, selected, selected_count, nthreads);
  delete[] selected_count;
}

template <OSwap_Style oswap_style>
void TightCompact_inner(unsigned char *buf, size_t N, size_t block_size, bool *selected, uint32_t *selected_count){
  FOAV_SAFE2_CNTXT(TC_inner_base_cases_of_recursion, N, block_size)
  if(N==0){
    return;
  }
  else if(N==1){
    return;
  }
  else if(N==2){
    bool swap = (!selected[0] & selected[1]);
    oswap_buffer<oswap_style>(buf, buf+block_size, block_size, swap);
    return;
  }

  size_t gt_pow2;
  size_t split_index;

  // Find largest power of 2 < N 
  gt_pow2 = pow2_lt(N);

  // For Order-preserving ORCompact
  // This will be right (R) of the recursion, and the leftover non-power of 2 left (L)
  split_index = N - gt_pow2;

  // Number of selected items in the non-power of 2 side (left)
  size_t mL;
  if (TC_PRECOMPUTE_COUNTS) {
    mL = selected_count[split_index] - selected_count[0];
  } else {
    mL = 0;
    for(size_t i=0; i<split_index; i++){
      mL+=selected[i];
    }
  }

  unsigned char *L_ptr = buf;
  unsigned char *R_ptr = buf + (split_index * block_size);

  //printf("Lsize = %ld, Rsize = %ld, Rside offset = %ld\n", split_index, gt_pow2, (gt_pow2 - split_index + mL));
  TightCompact_inner<oswap_style>(L_ptr, split_index, block_size, selected, selected_count);
  TightCompact_2power_inner<oswap_style>(R_ptr, gt_pow2, block_size, (gt_pow2 - split_index + mL) % gt_pow2, selected+split_index, selected_count+split_index);

  // For OP we CnS the first n_2 elements (split_size) against the suffix n_2 elements of the n_1 (2 power elements)
  R_ptr = buf + (gt_pow2 * block_size); 

  // Perform N-split_index oblivious swaps for this level
  FOAV_SAFE_CNTXT(TC_inner_oswap_loop, split_index)
  for (size_t i=0; i<split_index; i++){
    FOAV_SAFE2_CNTXT(TC_inner_oswap_loop, i, split_index)
    // Oswap blocks at L_start, R_start conditional on marked_items
    bool swap_flag = i>=mL;
    oswap_buffer<oswap_style>(L_ptr, R_ptr, block_size, swap_flag);
    L_ptr+=block_size;
    R_ptr+=block_size;
    FOAV_SAFE2_CNTXT(TC_inner_oswap_loop, i, split_index)
  }
}


template <OSwap_Style oswap_style>
void TightCompact_inner_parallel(unsigned char *buf, size_t N, size_t block_size, bool *selected, uint32_t *selected_count, size_t nthreads){
  FOAV_SAFE2_CNTXT(TC_inner_base_cases_of_recursion, N, block_size)
  FOAV_SAFE_CNTXT(TC_inner_base_cases_of_recursion, nthreads)
  if (nthreads <= 1 || N < 16) {
    unsigned long start = printf_with_rtclock("Thread %u starting TightCompact_inner(N=%lu)\n", g_thread_id, N);
    TightCompact_inner<oswap_style>(buf, N, block_size, selected, selected_count);
    printf_with_rtclock_diff(start, "Thread %u ending TightCompact_inner(N=%lu)\n", g_thread_id, N);
    return;
  }
  if(N==0){
    return;
  }
  else if(N==1){
    return;
  }
  else if(N==2){
    bool swap = (!selected[0] & selected[1]);
    oswap_buffer<oswap_style>(buf, buf+block_size, block_size, swap);
    return;
  }

  unsigned long start = printf_with_rtclock("Thread %u starting TightCompact_inner_parallel(N=%lu, nthreads=%lu)\n", g_thread_id, N, nthreads);

  size_t split_index, n1, n2;

  // Find largest power of 2 < N
  // This will be right (R) n1 of the recursion, and the leftover left (L) n2
  n1 = pow2_lt(N);
  n2 = N - n1;

  // Number of selected items in left
  size_t m2;
  m2 = selected_count[n2] - selected_count[0];

  unsigned char *L_ptr = buf;
  unsigned char *R_ptr = buf + (n2 * block_size);

  size_t lthreads = nthreads/2;
  size_t rthreads = nthreads - lthreads;

  struct TightCompact_2power_inner_parallel_args rightargs = {
    R_ptr, n1, block_size, n1 - n2 + m2, selected + n2, selected_count + n2,
    rthreads
  };
  threadpool_dispatch(g_thread_id+lthreads,
    TightCompact_2power_inner_parallel_launch<oswap_style>,
    &rightargs);
  TightCompact_inner_parallel<oswap_style>(L_ptr, n2, block_size, selected, selected_count, lthreads);
  threadpool_join(g_thread_id+lthreads, NULL);

  size_t num_swap = N-n1;
  FOAV_SAFE2_CNTXT(TC_inner_parallel, nthreads, num_swap)
  oswap_range_args args[nthreads];
  size_t inc = num_swap / nthreads;
  size_t extra = num_swap % nthreads;
  size_t last = 0;

  // We tweak R_ptr before we compare, to set compare the n2 prefix in L with the n2 suffix of R
  R_ptr = buf + (n1 * block_size);

  for (size_t i=0; i<nthreads; ++i) {
    size_t next = last + inc + (i < extra);
    args[i] = { block_size, last, next, m2, L_ptr, R_ptr, false };
    last = next;
    FOAV_SAFE2_CNTXT(TC_inner_parallel, i, nthreads)
  }
  for (size_t i=0; i<nthreads-1; ++i) {
    threadpool_dispatch(g_thread_id+1+i, oswap_range<oswap_style>, args+i);
  }
  // Do the last section ourselves
  oswap_range<oswap_style>((void*)(args+nthreads-1));
  FOAV_SAFE_CNTXT(TC_inner_parallel, nthreads)
  for (size_t i=0; i<nthreads-1; ++i) {
    threadpool_join(g_thread_id+1+i, NULL);
    FOAV_SAFE2_CNTXT(TC_inner_parallel, i, nthreads)
  }

  printf_with_rtclock_diff(start, "Thread %u ending TightCompact_inner_parallel(N=%lu, nthreads=%lu)\n", g_thread_id, N, nthreads);
}

  #ifndef BEFTS_MODE 

  // Perform the oswaps for input level over the OP_Tight Compaction Network
  template <OSwap_Style oswap_style>
  void process_TCN(uint64_t N, uint8_t level, unsigned char *bfr_ptr, size_t block_size,
          uint64_t *LS_distance) { 
    FOAV_SAFE2_CNTXT(process_TCN, N, level)
    FOAV_SAFE_CNTXT(process_TCN, block_size)
    uint64_t comparator_dist = (1<<level);
    // bfr_fop = bfr_first_operand_pointer, bfr_sop = bfr_second_operand_pointer
    unsigned char *bfr_fop = bfr_ptr;
    unsigned char *bfr_sop = bfr_ptr + (comparator_dist * block_size);

    // Number of oblivious swaps
    uint64_t num_oswaps = N - comparator_dist;
    uint64_t sop_index = comparator_dist;
    uint64_t fop_index = 0;

    FOAV_SAFE_CNTXT(process_TCN, num_oswaps)
    for(uint64_t i=0; i<num_oswaps; i++){ 
      uint64_t move_dist = LS_distance[sop_index] & (1 << (level+1)-1);
      // Obliviously if sop!=dummy AND move_dist!=0, set move_flag
      uint8_t dist_flag = ogt_set_flag(move_dist, 0);
      // but appropriate for 8 bytes oswaps.
      oswap_buffer<oswap_style>(bfr_fop, bfr_sop, block_size, dist_flag);

      // Adjust LS_distance after an oswap based on move_dist:
      // Obliviously if dist_flag, set LS_distance[thread][fop_index] to 
      // (LS_distance[thread][sop_index]-move_dist)
      LS_distance[sop_index]-= move_dist;
      oset_value(&(LS_distance[fop_index]), LS_distance[sop_index], dist_flag);
      oset_value(&(LS_distance[sop_index]), 0, dist_flag);

      bfr_fop+=block_size;
      bfr_sop+=block_size;
      sop_index++;
      fop_index++;
      FOAV_SAFE2_CNTXT(process_TCN, i, num_oswaps)
    }
  }


  template <OSwap_Style oswap_style>
  void OP_TightCompact(unsigned char *buf, uint64_t N, size_t block_size, bool *selected_list){
    
    FOAV_SAFE2_CNTXT(OP_TightCompact, N, block_size)
    FOAV_SAFE2_CNTXT(OP_TightCompact, buf, selected_list)
    uint64_t *LS_distance = new uint64_t[N];

    int TCN_l = calculatelog2(N);
    compute_LS_distances(N, buf, block_size, selected_list, LS_distance);
      
    FOAV_SAFE_CNTXT(OP_TightCompact, TCN_l)
    for(int l=0; l<TCN_l; l++) {
      process_TCN<oswap_style>(N, l, buf, block_size, LS_distance); 
      FOAV_SAFE2_CNTXT(OP_TightCompact, l, TCN_l)
    } 

    delete[] LS_distance;
  }

  #endif

#endif