group_data.h

Go to the documentation of this file.

 
#ifndef XGBOOST_COMMON_GROUP_DATA_H_
#define XGBOOST_COMMON_GROUP_DATA_H_
 
#include <cstddef>
#include <vector>
#include <algorithm>
#include <utility>
 
#include "xgboost/base.h"
 
namespace xgboost {
namespace common {
template<typename ValueType, typename SizeType = bst_ulong, bool is_row_major = false>
class ParallelGroupBuilder {
 public:
  ParallelGroupBuilder(std::vector<SizeType> *p_rptr,
                       std::vector<ValueType> *p_data,
                       size_t base_row_offset = 0)
      : rptr_(*p_rptr),
        data_(*p_data),
        base_row_offset_(base_row_offset) {}
 
  void InitBudget(std::size_t max_key, int nthread) {
    thread_rptr_.resize(nthread);
    const size_t full_size = is_row_major ? max_key : max_key - std::min(base_row_offset_, max_key);
    thread_displacement_ = is_row_major ? full_size / nthread : 0;
    for (std::size_t i = 0; i < thread_rptr_.size() - 1; ++i) {
      const size_t thread_size = is_row_major ? thread_displacement_ : full_size;
      thread_rptr_[i].resize(thread_size, 0);
    }
    const size_t last_thread_size = is_row_major ? (full_size - (nthread - 1)*thread_displacement_)
                                                 : full_size;
    thread_rptr_[nthread - 1].resize(last_thread_size, 0);
  }
 
  void AddBudget(std::size_t key, int threadid, SizeType nelem = 1) {
    std::vector<SizeType> &trptr = thread_rptr_[threadid];
    size_t offset_key = is_row_major ? (key - base_row_offset_ - threadid*thread_displacement_)
                                     : (key - base_row_offset_);
    if (trptr.size() < offset_key + 1) {
      trptr.resize(offset_key + 1, 0);
    }
    trptr[offset_key] += nelem;
  }
 
  inline void InitStorage() {
    if (is_row_major) {
      size_t expected_rows = 0;
      for (std::size_t tid = 0; tid < thread_rptr_.size(); ++tid) {
        expected_rows += thread_rptr_[tid].size();
      }
      // initialize rptr to be beginning of each segment
      SizeType rptr_fill_value = rptr_.empty() ? 0 : rptr_.back();
      rptr_.resize(expected_rows + base_row_offset_ + 1, rptr_fill_value);
 
      std::size_t count = 0;
      size_t offset_idx = base_row_offset_ + 1;
      for (std::size_t tid = 0; tid < thread_rptr_.size(); ++tid) {
        std::vector<SizeType> &trptr = thread_rptr_[tid];
        for (std::size_t i = 0; i < trptr.size(); ++i) {
          std::size_t thread_count = trptr[i];  // how many entries in this row
          trptr[i] = count + rptr_fill_value;
          count += thread_count;
          if (offset_idx < rptr_.size()) {
            rptr_[offset_idx++] += count;
          }
        }
      }
      data_.resize(rptr_.back());  // usage of empty allocator can help to improve performance
    } else {
      // set rptr to correct size
      SizeType rptr_fill_value = rptr_.empty() ? 0 : rptr_.back();
      for (std::size_t tid = 0; tid < thread_rptr_.size(); ++tid) {
        if (rptr_.size() <= thread_rptr_[tid].size() + base_row_offset_) {
          rptr_.resize(thread_rptr_[tid].size() + base_row_offset_ + 1,
                       rptr_fill_value);  // key + 1
        }
      }
      // initialize rptr to be beginning of each segment
      std::size_t count = 0;
      for (std::size_t i = base_row_offset_; i + 1 < rptr_.size(); ++i) {
        for (std::size_t tid = 0; tid < thread_rptr_.size(); ++tid) {
          std::vector<SizeType> &trptr = thread_rptr_[tid];
          if (i < trptr.size() +
                      base_row_offset_) {  // i^th row is assigned for this thread
            std::size_t thread_count =
                trptr[i - base_row_offset_];  // how many entries in this row
            trptr[i - base_row_offset_] = count + rptr_.back();
            count += thread_count;
          }
        }
        rptr_[i + 1] += count;  // pointer accumulated from all thread
      }
      data_.resize(rptr_.back());
    }
  }
 
  void Push(std::size_t key, ValueType&& value, int threadid) {
    size_t offset_key = is_row_major ? (key - base_row_offset_ - threadid * thread_displacement_)
                                     : (key - base_row_offset_);
    SizeType &rp = thread_rptr_[threadid][offset_key];
    data_[rp++] = std::move(value);
  }
 
 private:
  std::vector<SizeType> &rptr_;
  std::vector<ValueType> &data_;
  std::vector<std::vector<SizeType> > thread_rptr_;
  size_t base_row_offset_;
  size_t thread_displacement_;
};
}  // namespace common
}  // namespace xgboost
#endif  // XGBOOST_COMMON_GROUP_DATA_H_