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 #include "common.h"
12 
13 #ifdef __CUDACC__
14 #include "device_helpers.cuh"
15 #endif // __CUDACC__
16 
17 namespace xgboost {
18 namespace common {
19 
20 using CompressedByteT = unsigned char;
21 
22 namespace detail {
23 inline void SetBit(CompressedByteT *byte, int bit_idx) {
24  *byte |= 1 << bit_idx;
25 }
26 template <typename T>
27 inline T CheckBit(const T &byte, int bit_idx) {
28  return byte & (1 << bit_idx);
29 }
30 inline void ClearBit(CompressedByteT *byte, int bit_idx) {
31  *byte &= ~(1 << bit_idx);
32 }
33 static const int kPadding = 4; // Assign padding so we can read slightly off
34  // the beginning of the array
35 
36 // The number of bits required to represent a given unsigned range
37 inline XGBOOST_DEVICE size_t SymbolBits(size_t num_symbols) {
38  auto bits = std::ceil(log2(static_cast<double>(num_symbols)));
39  return common::Max(static_cast<size_t>(bits), size_t(1));
40 }
41 } // namespace detail
42 
55  size_t symbol_bits_;
56 
57  public:
58  XGBOOST_DEVICE explicit CompressedBufferWriter(size_t num_symbols)
59  : symbol_bits_(detail::SymbolBits(num_symbols)) {}
60 
76  static size_t CalculateBufferSize(size_t num_elements, size_t num_symbols) {
77  constexpr int kBitsPerByte = 8;
78  size_t compressed_size = static_cast<size_t>(std::ceil(
79  static_cast<double>(detail::SymbolBits(num_symbols) * num_elements) /
80  kBitsPerByte));
81  // Handle atomicOr where input must be unsigned int, hence 4 bytes aligned.
82  size_t ret =
83  std::ceil(static_cast<double>(compressed_size + detail::kPadding) /
84  static_cast<double>(sizeof(unsigned int))) *
85  sizeof(unsigned int);
86  return ret;
87  }
88 
89  template <typename T>
90  void WriteSymbol(CompressedByteT *buffer, T symbol, size_t offset) {
91  const int bits_per_byte = 8;
92 
93  for (size_t i = 0; i < symbol_bits_; i++) {
94  size_t byte_idx = ((offset + 1) * symbol_bits_ - (i + 1)) / bits_per_byte;
95  byte_idx += detail::kPadding;
96  size_t bit_idx =
97  ((bits_per_byte + i) - ((offset + 1) * symbol_bits_)) % bits_per_byte;
98 
99  if (detail::CheckBit(symbol, i)) {
100  detail::SetBit(&buffer[byte_idx], bit_idx);
101  } else {
102  detail::ClearBit(&buffer[byte_idx], bit_idx);
103  }
104  }
105  }
106 
107 #ifdef __CUDACC__
108  __device__ void AtomicWriteSymbol
109  (CompressedByteT* buffer, uint64_t symbol, size_t offset) {
110  size_t ibit_start = offset * symbol_bits_;
111  size_t ibit_end = (offset + 1) * symbol_bits_ - 1;
112  size_t ibyte_start = ibit_start / 8, ibyte_end = ibit_end / 8;
113 
114  symbol <<= 7 - ibit_end % 8;
115  for (ptrdiff_t ibyte = ibyte_end; ibyte >= static_cast<ptrdiff_t>(ibyte_start); --ibyte) {
116  dh::AtomicOrByte(reinterpret_cast<unsigned int*>(buffer + detail::kPadding),
117  ibyte, symbol & 0xff);
118  symbol >>= 8;
119  }
120  }
121 #endif // __CUDACC__
122 
123  template <typename IterT>
124  void Write(CompressedByteT *buffer, IterT input_begin, IterT input_end) {
125  uint64_t tmp = 0;
126  size_t stored_bits = 0;
127  const size_t max_stored_bits = 64 - symbol_bits_;
128  size_t buffer_position = detail::kPadding;
129  const size_t num_symbols = input_end - input_begin;
130  for (size_t i = 0; i < num_symbols; i++) {
131  typename std::iterator_traits<IterT>::value_type symbol = input_begin[i];
132  if (stored_bits > max_stored_bits) {
133  // Eject only full bytes
134  size_t tmp_bytes = stored_bits / 8;
135  for (size_t j = 0; j < tmp_bytes; j++) {
136  buffer[buffer_position] = static_cast<CompressedByteT>(
137  tmp >> (stored_bits - (j + 1) * 8));
138  buffer_position++;
139  }
140  stored_bits -= tmp_bytes * 8;
141  tmp &= (1 << stored_bits) - 1;
142  }
143  // Store symbol
144  tmp <<= symbol_bits_;
145  tmp |= symbol;
146  stored_bits += symbol_bits_;
147  }
148 
149  // Eject all bytes
150  int tmp_bytes =
151  static_cast<int>(std::ceil(static_cast<float>(stored_bits) / 8));
152  for (int j = 0; j < tmp_bytes; j++) {
153  int shift_bits = static_cast<int>(stored_bits) - (j + 1) * 8;
154  if (shift_bits >= 0) {
155  buffer[buffer_position] =
156  static_cast<CompressedByteT>(tmp >> shift_bits);
157  } else {
158  buffer[buffer_position] =
159  static_cast<CompressedByteT>(tmp << std::abs(shift_bits));
160  }
161  buffer_position++;
162  }
163  }
164 };
165 
174 template <typename T>
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  const CompressedByteT *buffer_ {nullptr};
186  size_t symbol_bits_ {0};
187  size_t offset_ {0};
188 
189  public:
190  CompressedIterator() = default;
191  CompressedIterator(const CompressedByteT *buffer, size_t num_symbols)
192  : buffer_(buffer) {
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:123
T CheckBit(const T &byte, int bit_idx)
Definition: compressed_iterator.h:27
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:54
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 required for a given number of elements and a symbol range.
Definition: compressed_iterator.h:76
void Write(CompressedByteT *buffer, IterT input_begin, IterT input_end)
Definition: compressed_iterator.h:124
void ClearBit(CompressedByteT *byte, int bit_idx)
Definition: compressed_iterator.h:30
CompressedIterator(const CompressedByteT *buffer, size_t num_symbols)
Definition: compressed_iterator.h:191
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:113
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:23
#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
XGBOOST_DEVICE size_t SymbolBits(size_t num_symbols)
Definition: compressed_iterator.h:37
defines configuration macros of xgboost.
XGBOOST_DEVICE CompressedBufferWriter(size_t num_symbols)
Definition: compressed_iterator.h:58
Common utilities.
ptrdiff_t difference_type
Definition: compressed_iterator.h:179
XGBOOST_DEVICE T Max(T a, T b)
Definition: common.h:68
unsigned char CompressedByteT
Definition: compressed_iterator.h:20
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:90