xgboost
compressed_iterator.h
Go to the documentation of this file.
1 
5 #pragma once
6 #include <xgboost/base.h>
7 #include <cmath>
8 #include <cstddef>
9 #include <algorithm>
10 
11 #ifdef __CUDACC__
12 #include "device_helpers.cuh"
13 #endif // __CUDACC__
14 
15 namespace xgboost {
16 namespace common {
17 
18 using CompressedByteT = unsigned char;
19 
20 namespace detail {
21 inline void SetBit(CompressedByteT *byte, int bit_idx) {
22  *byte |= 1 << bit_idx;
23 }
24 template <typename T>
25 inline T CheckBit(const T &byte, int bit_idx) {
26  return byte & (1 << bit_idx);
27 }
28 inline void ClearBit(CompressedByteT *byte, int bit_idx) {
29  *byte &= ~(1 << bit_idx);
30 }
31 static const int kPadding = 4; // Assign padding so we can read slightly off
32  // the beginning of the array
33 
34 // The number of bits required to represent a given unsigned range
35 static size_t SymbolBits(size_t num_symbols) {
36  auto bits = std::ceil(std::log2(num_symbols));
37  return std::max(static_cast<size_t>(bits), size_t(1));
38 }
39 } // namespace detail
40 
53  private:
54  size_t symbol_bits_;
55  size_t offset_;
56 
57  public:
58  explicit CompressedBufferWriter(size_t num_symbols) : offset_(0) {
59  symbol_bits_ = detail::SymbolBits(num_symbols);
60  }
61 
78  static size_t CalculateBufferSize(size_t num_elements, size_t num_symbols) {
79  const int bits_per_byte = 8;
80  size_t compressed_size = static_cast<size_t>(std::ceil(
81  static_cast<double>(detail::SymbolBits(num_symbols) * num_elements) /
82  bits_per_byte));
83  return compressed_size + detail::kPadding;
84  }
85 
86  template <typename T>
87  void WriteSymbol(CompressedByteT *buffer, T symbol, size_t offset) {
88  const int bits_per_byte = 8;
89 
90  for (size_t i = 0; i < symbol_bits_; i++) {
91  size_t byte_idx = ((offset + 1) * symbol_bits_ - (i + 1)) / bits_per_byte;
92  byte_idx += detail::kPadding;
93  size_t bit_idx =
94  ((bits_per_byte + i) - ((offset + 1) * symbol_bits_)) % bits_per_byte;
95 
96  if (detail::CheckBit(symbol, i)) {
97  detail::SetBit(&buffer[byte_idx], bit_idx);
98  } else {
99  detail::ClearBit(&buffer[byte_idx], bit_idx);
100  }
101  }
102  }
103 
104 #ifdef __CUDACC__
105  __device__ void AtomicWriteSymbol
106  (CompressedByteT* buffer, uint64_t symbol, size_t offset) {
107  size_t ibit_start = offset * symbol_bits_;
108  size_t ibit_end = (offset + 1) * symbol_bits_ - 1;
109  size_t ibyte_start = ibit_start / 8, ibyte_end = ibit_end / 8;
110 
111  symbol <<= 7 - ibit_end % 8;
112  for (ptrdiff_t ibyte = ibyte_end; ibyte >= (ptrdiff_t)ibyte_start; --ibyte) {
113  dh::AtomicOrByte(reinterpret_cast<unsigned int*>(buffer + detail::kPadding),
114  ibyte, symbol & 0xff);
115  symbol >>= 8;
116  }
117  }
118 #endif // __CUDACC__
119 
120  template <typename IterT>
121  void Write(CompressedByteT *buffer, IterT input_begin, IterT input_end) {
122  uint64_t tmp = 0;
123  size_t stored_bits = 0;
124  const size_t max_stored_bits = 64 - symbol_bits_;
125  size_t buffer_position = detail::kPadding;
126  const size_t num_symbols = input_end - input_begin;
127  for (size_t i = 0; i < num_symbols; i++) {
128  typename std::iterator_traits<IterT>::value_type symbol = input_begin[i];
129  if (stored_bits > max_stored_bits) {
130  // Eject only full bytes
131  size_t tmp_bytes = stored_bits / 8;
132  for (size_t j = 0; j < tmp_bytes; j++) {
133  buffer[buffer_position] = static_cast<CompressedByteT>(
134  tmp >> (stored_bits - (j + 1) * 8));
135  buffer_position++;
136  }
137  stored_bits -= tmp_bytes * 8;
138  tmp &= (1 << stored_bits) - 1;
139  }
140  // Store symbol
141  tmp <<= symbol_bits_;
142  tmp |= symbol;
143  stored_bits += symbol_bits_;
144  }
145 
146  // Eject all bytes
147  int tmp_bytes =
148  static_cast<int>(std::ceil(static_cast<float>(stored_bits) / 8));
149  for (int j = 0; j < tmp_bytes; j++) {
150  int shift_bits = static_cast<int>(stored_bits) - (j + 1) * 8;
151  if (shift_bits >= 0) {
152  buffer[buffer_position] =
153  static_cast<CompressedByteT>(tmp >> shift_bits);
154  } else {
155  buffer[buffer_position] =
156  static_cast<CompressedByteT>(tmp << std::abs(shift_bits));
157  }
158  buffer_position++;
159  }
160  }
161 };
162 
163 template <typename T>
164 
176  public:
177  // Type definitions for thrust
178  typedef CompressedIterator<T> self_type; // NOLINT
179  typedef ptrdiff_t difference_type; // NOLINT
180  typedef T value_type; // NOLINT
181  typedef value_type *pointer; // NOLINT
182  typedef value_type reference; // NOLINT
183 
184  private:
185  CompressedByteT *buffer_;
186  size_t symbol_bits_;
187  size_t offset_;
188 
189  public:
190  CompressedIterator() : buffer_(nullptr), symbol_bits_(0), offset_(0) {}
191  CompressedIterator(CompressedByteT *buffer, int num_symbols)
192  : buffer_(buffer), offset_(0) {
193  symbol_bits_ = detail::SymbolBits(num_symbols);
194  }
195 
196  XGBOOST_DEVICE reference operator*() const {
197  const int bits_per_byte = 8;
198  size_t start_bit_idx = ((offset_ + 1) * symbol_bits_ - 1);
199  size_t start_byte_idx = start_bit_idx / bits_per_byte;
200  start_byte_idx += detail::kPadding;
201 
202  // Read 5 bytes - the maximum we will need
203  uint64_t tmp = static_cast<uint64_t>(buffer_[start_byte_idx - 4]) << 32 |
204  static_cast<uint64_t>(buffer_[start_byte_idx - 3]) << 24 |
205  static_cast<uint64_t>(buffer_[start_byte_idx - 2]) << 16 |
206  static_cast<uint64_t>(buffer_[start_byte_idx - 1]) << 8 |
207  buffer_[start_byte_idx];
208  int bit_shift =
209  (bits_per_byte - ((offset_ + 1) * symbol_bits_)) % bits_per_byte;
210  tmp >>= bit_shift;
211  // Mask off unneeded bits
212  uint64_t mask = (static_cast<uint64_t>(1) << symbol_bits_) - 1;
213  return static_cast<T>(tmp & mask);
214  }
215 
216  XGBOOST_DEVICE reference operator[](size_t idx) const {
217  self_type offset = (*this);
218  offset.offset_ += idx;
219  return *offset;
220  }
221 };
222 } // namespace common
223 } // namespace xgboost
byte
Definition: span.h:112
T CheckBit(const T &byte, int bit_idx)
Definition: compressed_iterator.h:25
value_type * pointer
Definition: compressed_iterator.h:181
Writes bit compressed symbols to a memory buffer. Use CompressedIterator to read symbols back from bu...
Definition: compressed_iterator.h:52
CompressedIterator()
Definition: compressed_iterator.h:190
CompressedBufferWriter(size_t num_symbols)
Definition: compressed_iterator.h:58
XGBOOST_DEVICE reference operator[](size_t idx) const
Definition: compressed_iterator.h:216
static size_t CalculateBufferSize(size_t num_elements, size_t num_symbols)
Calculates number of bytes requiredm for a given number of elements and a symbol range.
Definition: compressed_iterator.h:78
void Write(CompressedByteT *buffer, IterT input_begin, IterT input_end)
Definition: compressed_iterator.h:121
void ClearBit(CompressedByteT *byte, int bit_idx)
Definition: compressed_iterator.h:28
CompressedIterator(CompressedByteT *buffer, int num_symbols)
Definition: compressed_iterator.h:191
value_type reference
Definition: compressed_iterator.h:182
XGBOOST_DEVICE reference operator*() const
Definition: compressed_iterator.h:196
void SetBit(CompressedByteT *byte, int bit_idx)
Definition: compressed_iterator.h:21
#define XGBOOST_DEVICE
Tag function as usable by device.
Definition: base.h:84
namespace of xgboost
Definition: base.h:102
CompressedIterator< T > self_type
Definition: compressed_iterator.h:178
typename std::conditional< std::is_same< std::ptrdiff_t, std::int64_t >::value, std::ptrdiff_t, std::int64_t >::type ptrdiff_t
Definition: span.h:102
defines configuration macros of xgboost.
ptrdiff_t difference_type
Definition: compressed_iterator.h:179
unsigned char CompressedByteT
Definition: compressed_iterator.h:18
T value_type
Definition: compressed_iterator.h:180
Read symbols from a bit compressed memory buffer. Usable on device and host.
Definition: compressed_iterator.h:175
void WriteSymbol(CompressedByteT *buffer, T symbol, size_t offset)
Definition: compressed_iterator.h:87