branch_selection.hpp

#ifndef BV_BRANCH_SELECTION_HPP
#define BV_BRANCH_SELECTION_HPP
 
#include <cassert>
#include <cstdint>
#include <cstring>
 
#ifndef CACHE_LINE
// Apparently the most common cache line size is 64.
#define CACHE_LINE 64
#endif
 
namespace bv {
 
template <class dtype, dtype branches>
class branchless_scan {
   protected:
    dtype elems_[branches];
 
    static_assert((branches == 8) || (branches == 16) || (branches == 32) ||
                      (branches == 64) || (branches == 128),
                  "branching factor needs to be a reasonable power of 2");
 
   public:
    branchless_scan() {
        for (dtype i = 0; i < branches; i++) {
            elems_[i] = (~dtype(0)) >> 1;
        }
    }
 
    dtype get(uint8_t index) const { return elems_[index]; }
 
    void set(uint8_t index, dtype value) { elems_[index] = value; }
 
    template <class T>
    void increment(uint8_t from, uint8_t array_size, T change) {
        for (uint8_t i = from; i < array_size; i++) {
            elems_[i] += change;
        }
    }
 
    void insert(uint8_t index, uint8_t array_size, dtype a_value,
                dtype b_value) {
        assert(index > 0);
        for (uint8_t i = array_size; i > index; i--) {
            elems_[i] = elems_[i - 1];
        }
        elems_[index - 1] = index != 1 ? elems_[index - 2] + a_value : a_value;
        elems_[index] = elems_[index - 1] + b_value;
    }
 
    void remove(uint8_t index, uint8_t array_size) {
        for (uint8_t i = index; i < array_size - 1; i++) {
            elems_[i] = elems_[i + 1];
        }
        elems_[array_size - 1] = (~dtype(0)) >> 1;
    }
 
    void clear_first(uint8_t n, uint8_t array_size) {
        dtype subtrahend = elems_[n - 1];
        for (uint8_t i = n; i < array_size; i++) {
            elems_[i - n] = elems_[i] - subtrahend;
            elems_[i] = (~dtype(0)) >> 1;
        }
    }
 
    void clear_last(uint8_t n, uint8_t array_size) {
        for (uint8_t i = array_size - n; i < array_size; i++) {
            elems_[i] = (~dtype(0)) >> 1;
        }
    }
 
    void append(uint8_t n_elems, uint8_t array_size, branchless_scan* other) {
        dtype addend = array_size != 0 ? elems_[array_size - 1] : 0;
        for (uint8_t i = 0; i < n_elems; i++) {
            elems_[i + array_size] = addend + other->get(i);
        }
    }
 
    void append(uint8_t index, dtype value) {
        if (index == 0) {
            elems_[index] = value;
        } else {
            elems_[index] = elems_[index - 1] + value;
        }
    }
 
    void prepend(uint8_t n_elems, uint8_t array_size, uint8_t o_size,
                 branchless_scan* other) {
        memmove(elems_ + n_elems, elems_, array_size * sizeof(dtype));
        dtype subtrahend =
            n_elems < o_size ? other->get(o_size - n_elems - 1) : 0;
        for (uint8_t i = 0; i < n_elems; i++) {
            elems_[i] = other->get(i + o_size - n_elems) - subtrahend;
        }
        for (uint8_t i = n_elems; i < array_size + n_elems; i++) {
            elems_[i] += elems_[n_elems - 1];
        }
    }
 
    uint8_t find(dtype q) const {
        constexpr dtype SIGN_BIT = ~((~dtype(0)) >> 1);
        constexpr dtype num_bits = sizeof(dtype) * 8;
        constexpr dtype lines = CACHE_LINE / sizeof(dtype);
        for (dtype i = 0; i < branches; i += lines) {
            __builtin_prefetch(elems_ + i);
        }
        uint8_t idx;
        if constexpr (branches == 128) {
            idx = (uint8_t(1) << 6) - 1;
            idx ^= (dtype((elems_[idx] - q) & SIGN_BIT) >> (num_bits - 7)) |
                   (uint8_t(1) << 5);
            idx ^= (dtype((elems_[idx] - q) & SIGN_BIT) >> (num_bits - 6)) |
                   (uint8_t(1) << 4);
            idx ^= (dtype((elems_[idx] - q) & SIGN_BIT) >> (num_bits - 5)) |
                   (uint8_t(1) << 3);
            idx ^= (dtype((elems_[idx] - q) & SIGN_BIT) >> (num_bits - 4)) |
                   (uint8_t(1) << 2);
        } else if constexpr (branches == 64) {
            idx = (uint8_t(1) << 5) - 1;
            idx ^= (dtype((elems_[idx] - q) & SIGN_BIT) >> (num_bits - 6)) |
                   (uint8_t(1) << 4);
            idx ^= (dtype((elems_[idx] - q) & SIGN_BIT) >> (num_bits - 5)) |
                   (uint8_t(1) << 3);
            idx ^= (dtype((elems_[idx] - q) & SIGN_BIT) >> (num_bits - 4)) |
                   (uint8_t(1) << 2);
        } else if constexpr (branches == 32) {
            idx = (uint8_t(1) << 4) - 1;
            idx ^= (dtype((elems_[idx] - q) & SIGN_BIT) >> (num_bits - 5)) |
                   (uint8_t(1) << 3);
            idx ^= (dtype((elems_[idx] - q) & SIGN_BIT) >> (num_bits - 4)) |
                   (uint8_t(1) << 2);
        } else if constexpr (branches == 16) {
            idx = (uint8_t(1) << 3) - 1;
            idx ^= (dtype((elems_[idx] - q) & SIGN_BIT) >> (num_bits - 4)) |
                   (uint8_t(1) << 2);
        } else {
            idx = (uint8_t(1) << 2) - 1;
        }
        idx ^= (dtype((elems_[idx] - q) & SIGN_BIT) >> (num_bits - 3)) |
               (uint8_t(1) << 1);
        idx ^= (dtype((elems_[idx] - q) & SIGN_BIT) >> (num_bits - 2)) |
               uint8_t(1);
        return idx ^ (dtype((elems_[idx] - q) & SIGN_BIT) >> (num_bits - 1));
    }
};
 
}  // namespace bv
 
#endif