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
54 class CompressedBufferWriter {
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>
175 class CompressedIterator {
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
xgboost::common::detail::CheckBit
T CheckBit(const T &byte, int bit_idx)
Definition: compressed_iterator.h:27
xgboost::common::CompressedIterator::value_type
T value_type
Definition: compressed_iterator.h:180
xgboost::common::detail::ptrdiff_t
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:124
xgboost::common::detail::SymbolBits
XGBOOST_DEVICE size_t SymbolBits(size_t num_symbols)
Definition: compressed_iterator.h:37
xgboost::common::detail::ClearBit
void ClearBit(CompressedByteT *byte, int bit_idx)
Definition: compressed_iterator.h:30
xgboost::common::CompressedIterator::CompressedIterator
CompressedIterator()=default
xgboost::common::Max
XGBOOST_DEVICE T Max(T a, T b)
Definition: common.h:71
xgboost::common::CompressedIterator::CompressedIterator
CompressedIterator(const CompressedByteT *buffer, size_t num_symbols)
Definition: compressed_iterator.h:191
xgboost::common::CompressedByteT
unsigned char CompressedByteT
Definition: compressed_iterator.h:20
xgboost::common::CompressedBufferWriter::Write
void Write(CompressedByteT *buffer, IterT input_begin, IterT input_end)
Definition: compressed_iterator.h:124
xgboost::common::CompressedBufferWriter
Writes bit compressed symbols to a memory buffer. Use CompressedIterator to read symbols back from bu...
Definition: compressed_iterator.h:54
base.h
defines configuration macros of xgboost.
xgboost::common::CompressedIterator
Read symbols from a bit compressed memory buffer. Usable on device and host.
Definition: compressed_iterator.h:175
xgboost::common::CompressedBufferWriter::CalculateBufferSize
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
xgboost::common::CompressedIterator::pointer
value_type * pointer
Definition: compressed_iterator.h:181
xgboost::common::CompressedIterator::self_type
CompressedIterator< T > self_type
Definition: compressed_iterator.h:178
xgboost::common::CompressedBufferWriter::CompressedBufferWriter
XGBOOST_DEVICE CompressedBufferWriter(size_t num_symbols)
Definition: compressed_iterator.h:58
xgboost::common::detail::SetBit
void SetBit(CompressedByteT *byte, int bit_idx)
Definition: compressed_iterator.h:23
common.h
Common utilities.
xgboost::common::CompressedIterator::operator[]
XGBOOST_DEVICE reference operator[](size_t idx) const
Definition: compressed_iterator.h:216
xgboost::common::CompressedIterator::operator*
XGBOOST_DEVICE reference operator*() const
Definition: compressed_iterator.h:196
xgboost::common::CompressedIterator::difference_type
ptrdiff_t difference_type
Definition: compressed_iterator.h:179
XGBOOST_DEVICE
#define XGBOOST_DEVICE
Tag function as usable by device.
Definition: base.h:84
xgboost::common::CompressedIterator::reference
value_type reference
Definition: compressed_iterator.h:182
xgboost
namespace of xgboost
Definition: base.h:110
xgboost::common::CompressedBufferWriter::WriteSymbol
void WriteSymbol(CompressedByteT *buffer, T symbol, size_t offset)
Definition: compressed_iterator.h:90