partitioner.hpp

// reference:
// - gomp: https://github.com/gcc-mirror/gcc/blob/master/libgomp/iter.c
// - komp: https://github.com/llvm-mirror/openmp/blob/master/runtime/src/kmp_dispatch.cpp
 
#pragma once

 
namespace tf {
  
enum class PartitionerType : int {
  STATIC,
  DYNAMIC
};
 
 
//template <typename C>
//class PartitionInvoker : public PartitionerBase {
//
//  protected
//
//  C _closure;
//
//  template <typename... ArgsT>
//  auto operator()(ArgsT&&... args) {
//    return std::invoke(closure, std::forward<ArgsT>(args)...);
//  }
//
//  template <typename... ArgsT>
//  auto operator()(ArgsT&&... args) const {
//    return std::invoke(closure, std::forward<ArgsT>(args)...);
//  }
//
//};

class DefaultClosureWrapper {};
 
 
 
// ----------------------------------------------------------------------------
// Partitioner Base
// ----------------------------------------------------------------------------

template <typename C = DefaultClosureWrapper>
class PartitionerBase {
 
  public:
  
  constexpr static bool is_default_wrapper_v = std::is_same_v<C, DefaultClosureWrapper>;
  
  using closure_wrapper_type = C;

  PartitionerBase() = default;

  explicit PartitionerBase(size_t chunk_size) : _chunk_size {chunk_size} {}
  
  PartitionerBase(size_t chunk_size, C&& closure_wrapper) :
    _chunk_size {chunk_size},
    _closure_wrapper {std::forward<C>(closure_wrapper)} {
  }

  size_t chunk_size() const { return _chunk_size; }
  
  void chunk_size(size_t cz) { _chunk_size = cz; }

  const C& closure_wrapper() const { return _closure_wrapper; }

  C& closure_wrapper() { return _closure_wrapper; }

  template <typename F>
  void closure_wrapper(F&& fn) { _closure_wrapper = std::forward<F>(fn); }

  template <typename F>
  TF_FORCE_INLINE decltype(auto) operator () (F&& callable) {
    if constexpr(is_default_wrapper_v) {
      return std::forward<F>(callable);
    }
    else {
      // closure wrapper is stateful - capture it by reference
      return [this, c=std::forward<F>(callable)]() mutable { _closure_wrapper(c); };
    }
  }
 
  protected:
  
  size_t _chunk_size{0};
  
  C _closure_wrapper;
};
 
// ----------------------------------------------------------------------------
// Static Partitioner
// ----------------------------------------------------------------------------

template <typename C = DefaultClosureWrapper>
class StaticPartitioner : public PartitionerBase<C> {
 
  public:
  
  static constexpr PartitionerType type() { return PartitionerType::STATIC; }

  StaticPartitioner() = default;
  
  explicit StaticPartitioner(size_t sz) : PartitionerBase<C>(sz) {}
   
  explicit StaticPartitioner(size_t sz, C&& closure) :
    PartitionerBase<C>(sz, std::forward<C>(closure)) {
  }
  
  size_t adjusted_chunk_size(size_t N, size_t W, size_t w) const {
    return this->_chunk_size ? this->_chunk_size : N/W + (w < N%W);
  }
  
  // --------------------------------------------------------------------------
  // scheduling methods
  // --------------------------------------------------------------------------
  
  template <typename F>
  void loop(
    size_t N, size_t W, size_t curr_b, size_t chunk_size, F&& func
  ) {
    size_t stride = W * chunk_size;
    while(curr_b < N) {
      size_t curr_e = (std::min)(curr_b + chunk_size, N);
      func(curr_b, curr_e);
      curr_b += stride;
    }
  }
  
  template <typename F>
  void loop_until(
    size_t N, size_t W, size_t curr_b, size_t chunk_size, F&& func
  ) {
    size_t stride = W * chunk_size;
    while(curr_b < N) {
      size_t curr_e = (std::min)(curr_b + chunk_size, N);
      if(func(curr_b, curr_e)) {
        return;
      }
      curr_b += stride;
    }
  }
  
  template <IndexRangeMDLike R, typename F>
  void loop(R& range, size_t N, size_t W, std::atomic<size_t>& next, F&& func) const {
    //size_t N = range.volume();
    size_t curr_b = next.load(std::memory_order_relaxed);
    size_t chunk_size = (this->_chunk_size == 0) ? (N+W-1) / W : this->_chunk_size;
 
    while (curr_b < N) {
      // calculates the actual valid sub-box and exact consumed amount
      auto [box, consumed] = range.consume_chunk(curr_b, chunk_size);
      if (next.compare_exchange_weak(curr_b, curr_b + consumed,
                                     std::memory_order_relaxed,
                                     std::memory_order_relaxed)) {
        func(box);
        curr_b += consumed;
      }
    }
  }
 
};
 
// ----------------------------------------------------------------------------
// Guided Partitioner
// ----------------------------------------------------------------------------

template <typename C = DefaultClosureWrapper>
class GuidedPartitioner : public PartitionerBase<C> {
 
  public:
  
  static constexpr PartitionerType type() { return PartitionerType::DYNAMIC; }
  
  GuidedPartitioner() = default;

  explicit GuidedPartitioner(size_t sz) : PartitionerBase<C> (sz) {}
  
  explicit GuidedPartitioner(size_t sz, C&& closure) :
    PartitionerBase<C>(sz, std::forward<C>(closure)) {
  }
  
  // --------------------------------------------------------------------------
  // scheduling methods
  // --------------------------------------------------------------------------
  
  template <typename F>
  void loop(
    size_t N, size_t W, std::atomic<size_t>& next, F&& func
  ) const {
 
    size_t chunk_size = (this->_chunk_size == 0) ? size_t{1} : this->_chunk_size;
 
    size_t p1 = 2 * W * (chunk_size + 1);
    float  p2 = 0.5f / static_cast<float>(W);
    size_t curr_b = next.load(std::memory_order_relaxed);
 
    while(curr_b < N) {
 
      size_t r = N - curr_b;
 
      // fine-grained
      if(r < p1) {
        while(1) {
          curr_b = next.fetch_add(chunk_size, std::memory_order_relaxed);
          if(curr_b >= N) {
            return;
          }
          func(curr_b, (std::min)(curr_b + chunk_size, N));
        }
        break;
      }
      // coarse-grained
      else {
        size_t q = static_cast<size_t>(p2 * r);
        if(q < chunk_size) {
          q = chunk_size;
        }
        //size_t curr_e = (q <= r) ? curr_b + q : N;
        size_t curr_e = (std::min)(curr_b + q, N);
        if(next.compare_exchange_strong(curr_b, curr_e, std::memory_order_relaxed,
                                                        std::memory_order_relaxed)) {
          func(curr_b, curr_e);
          //curr_b = next.load(std::memory_order_relaxed);
          curr_b = curr_e;
        }
      }
    }
  }
  
  template <typename F>
  void loop_until(
    size_t N, size_t W, std::atomic<size_t>& next, F&& func
  ) const {
 
    size_t chunk_size = (this->_chunk_size == 0) ? size_t{1} : this->_chunk_size;
 
    size_t p1 = 2 * W * (chunk_size + 1);
    float  p2 = 0.5f / static_cast<float>(W);
    size_t curr_b = next.load(std::memory_order_relaxed);
 
    while(curr_b < N) {
 
      size_t r = N - curr_b;
 
      // fine-grained
      if(r < p1) {
        while(1) {
          curr_b = next.fetch_add(chunk_size, std::memory_order_relaxed);
          if(curr_b >= N) {
            return;
          }
          if(func(curr_b, (std::min)(curr_b + chunk_size, N))) {
            return;
          }
        }
        break;
      }
      // coarse-grained
      else {
        size_t q = static_cast<size_t>(p2 * r);
        if(q < chunk_size) {
          q = chunk_size;
        }
        //size_t curr_e = (q <= r) ? curr_b + q : N;
        size_t curr_e = (std::min)(curr_b + q, N);
        if(next.compare_exchange_strong(curr_b, curr_e, std::memory_order_relaxed,
                                                        std::memory_order_relaxed)) {
          if(func(curr_b, curr_e)) {
            return;
          }
          //curr_b = next.load(std::memory_order_relaxed);
          curr_b = curr_e;
        }
      }
    }
  }
  
  template <IndexRangeMDLike R, typename F>
  void loop(R& range, size_t N, size_t W, std::atomic<size_t>& next, F&& func) const {
 
    size_t chunk_size = (this->_chunk_size == 0) ? size_t{1} : this->_chunk_size;
    size_t requested_chunk_size;
    size_t p1 = 2 * W * (chunk_size + 1);
    float  p2 = 0.5f / static_cast<float>(W);
    size_t curr_b = next.load(std::memory_order_relaxed);
 
    while(curr_b < N) {
 
      size_t r = N - curr_b;
      
      // find-grained
      if(r < p1) {
        requested_chunk_size = chunk_size;
      }
      // coarse
      else {
        requested_chunk_size = (std::max)(static_cast<size_t>(p2*r), chunk_size);
      }
      
      auto [box, consumed] = range.consume_chunk(curr_b, requested_chunk_size);
      
      if (next.compare_exchange_weak(curr_b, curr_b + consumed,
                                     std::memory_order_relaxed,
                                     std::memory_order_relaxed)) {
        func(box);
        curr_b += consumed;
      }
    }
  }
 
};
 
// ----------------------------------------------------------------------------
// Dynamic Partitioner
// ----------------------------------------------------------------------------

template <typename C = DefaultClosureWrapper>
class DynamicPartitioner : public PartitionerBase<C> {
 
  public:
  
  static constexpr PartitionerType type() { return PartitionerType::DYNAMIC; }

  DynamicPartitioner() = default;
  
  explicit DynamicPartitioner(size_t sz) : PartitionerBase<C>(sz) {}
   
  explicit DynamicPartitioner(size_t sz, C&& closure) :
    PartitionerBase<C>(sz, std::forward<C>(closure)) {
  }
  
  // --------------------------------------------------------------------------
  // scheduling methods
  // --------------------------------------------------------------------------

  template <typename F>
  void loop(
    size_t N, size_t, std::atomic<size_t>& next, F&& func
  ) const {
 
    size_t chunk_size = (this->_chunk_size == 0) ? size_t{1} : this->_chunk_size;
    size_t curr_b = next.fetch_add(chunk_size, std::memory_order_relaxed);
 
    while(curr_b < N) {
      func(curr_b, (std::min)(curr_b + chunk_size, N));
      curr_b = next.fetch_add(chunk_size, std::memory_order_relaxed);
    }
  }
  
  template <typename F>
  void loop_until(
    size_t N, size_t, std::atomic<size_t>& next, F&& func
  ) const {
 
    size_t chunk_size = (this->_chunk_size == 0) ? size_t{1} : this->_chunk_size;
    size_t curr_b = next.fetch_add(chunk_size, std::memory_order_relaxed);
 
    while(curr_b < N) {
      if(func(curr_b, (std::min)(curr_b + chunk_size, N))) {
        return;
      }
      curr_b = next.fetch_add(chunk_size, std::memory_order_relaxed);
    }
  }
  
  template <IndexRangeMDLike R, typename F>
  void loop(R& range, size_t N, size_t, std::atomic<size_t>& next, F&& func) const {
    //size_t N = range.volume();
    size_t curr_b = next.load(std::memory_order_relaxed);
    size_t chunk_size = (this->_chunk_size == 0) ? size_t{1} : this->_chunk_size;
 
    while (curr_b < N) {
      // calculates the actual valid sub-box and exact consumed amount
      auto [box, consumed] = range.consume_chunk(curr_b, chunk_size);
      if (next.compare_exchange_weak(curr_b, curr_b + consumed,
                                     std::memory_order_relaxed,
                                     std::memory_order_relaxed)) {
        func(box);
        curr_b += consumed;
      }
    }
  }
 
};
 
// ----------------------------------------------------------------------------
// RandomPartitioner
// ----------------------------------------------------------------------------

template <typename C = DefaultClosureWrapper>
class RandomPartitioner : public PartitionerBase<C> {
 
  public:
  
  static constexpr PartitionerType type() { return PartitionerType::DYNAMIC; }

  RandomPartitioner() = default;
  
  explicit RandomPartitioner(size_t sz) : PartitionerBase<C>(sz) {}
   
  explicit RandomPartitioner(size_t sz, C&& closure) :
    PartitionerBase<C>(sz, std::forward<C>(closure)) {
  }
  
  RandomPartitioner(float alpha, float beta) : _alpha{alpha}, _beta{beta} {}

  RandomPartitioner(float alpha, float beta, C&& closure) : 
    _alpha {alpha}, _beta {beta}, 
    PartitionerBase<C>(0, std::forward<C>(closure)) {
  }

  float alpha() const { return _alpha; }
  
  float beta() const { return _beta; }
  
  std::pair<size_t, size_t> chunk_size_range(size_t N, size_t W) const {
    
    size_t b1 = static_cast<size_t>(_alpha * N * W);
    size_t b2 = static_cast<size_t>(_beta  * N * W);
 
    if(b1 > b2) {
      std::swap(b1, b2);
    }
 
    b1 = (std::max)(b1, size_t{1});
    b2 = (std::max)(b2, b1 + 1);
 
    return {b1, b2};
  }
 
  // --------------------------------------------------------------------------
  // scheduling methods
  // --------------------------------------------------------------------------
  
  template <typename F>
  void loop(
    size_t N, size_t W, std::atomic<size_t>& next, F&& func
  ) const {
 
    auto [b1, b2] = chunk_size_range(N, W); 
    
    std::default_random_engine engine {std::random_device{}()};
    std::uniform_int_distribution<size_t> dist(b1, b2);
    
    size_t chunk_size = dist(engine);
    size_t curr_b = next.fetch_add(chunk_size, std::memory_order_relaxed);
 
    while(curr_b < N) {
      func(curr_b, (std::min)(curr_b + chunk_size, N));
      chunk_size = dist(engine);
      curr_b = next.fetch_add(chunk_size, std::memory_order_relaxed);
    }
  }

  template <typename F>
  void loop_until(
    size_t N, size_t W, std::atomic<size_t>& next, F&& func
  ) const {
 
    auto [b1, b2] = chunk_size_range(N, W); 
    
    std::default_random_engine engine {std::random_device{}()};
    std::uniform_int_distribution<size_t> dist(b1, b2);
    
    size_t chunk_size = dist(engine);
    size_t curr_b = next.fetch_add(chunk_size, std::memory_order_relaxed);
 
    while(curr_b < N) {
      if(func(curr_b, (std::min)(curr_b + chunk_size, N))){
        return;
      }
      chunk_size = dist(engine);
      curr_b = next.fetch_add(chunk_size, std::memory_order_relaxed);
    }
  }

  template <IndexRangeMDLike R, typename F>
  void loop(R& range, size_t N, size_t W, std::atomic<size_t>& next, F&& func) const {
 
    auto [b1, b2] = chunk_size_range(N, W);
 
    std::default_random_engine engine{std::random_device{}()};
    std::uniform_int_distribution<size_t> dist(b1, b2);
 
    size_t curr_b = next.load(std::memory_order_relaxed);
 
    while (curr_b < N) {
      auto [box, consumed] = range.consume_chunk(curr_b, dist(engine));
      if (next.compare_exchange_weak(curr_b, curr_b + consumed,
                                     std::memory_order_relaxed,
                                     std::memory_order_relaxed)) {
        func(box);
        curr_b += consumed;
      }
      // CAS failure reloads curr_b; resample chunk on next iteration naturally
    }
  }
 
  private:
 
  float _alpha {0.01f};
  float _beta  {0.50f};
};

using DefaultPartitioner = GuidedPartitioner<>;

template <typename P>
concept Partitioner = std::derived_from<P, PartitionerBase<typename P::closure_wrapper_type>>;

template <typename P>
inline constexpr bool is_partitioner_v = Partitioner<P>;
 
}  // end of namespace tf -----------------------------------------------------