Skip to content

Commit 0253222

Browse files
BigVanBigVan
authored
Merge pull request #384 from liulanzheng/main
support 32bit linearized B+tree for lsmt
2 parents cd6411c + 80362a7 commit 0253222

File tree

1 file changed

+77
-44
lines changed

1 file changed

+77
-44
lines changed

‎src/overlaybd/lsmt/index.cpp‎

Lines changed: 77 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
*/
1616

1717
#include "index.h"
18+
#include "file.h"
1819
#include <vector>
1920
#include <set>
2021
#include <algorithm>
@@ -66,53 +67,87 @@ bool is_avx512f_supported() {
6667
#endif
6768
}
6869

69-
const static uint32_t ORDER = 8;
70-
const static uint32_t MAX_LEVEL = 10;
71-
static constexpr uint32_t NODES_PER_LEVEL[MAX_LEVEL] = {8, 72, 648, 5832, 52488, 472392, 4251528, 38263752, 344373768, 3099363912};
72-
static constexpr uint32_t LEVEL_START_ID[MAX_LEVEL] = {0, 8, 80, 728, 6560, 59048, 531440, 4782968, 43046720, 387420488};
70+
const static uint32_t KEYS_PER_NODE_64 = 8;
71+
const static uint32_t MAX_LEVEL_64 = 10;
72+
static constexpr uint32_t NODES_PER_LEVEL_64[MAX_LEVEL_64] = {8, 72, 648, 5832, 52488, 472392, 4251528, 38263752, 344373768, 3099363912};
73+
static constexpr uint32_t LEVEL_START_ID_64[MAX_LEVEL_64] = {0, 8, 80, 728, 6560, 59048, 531440, 4782968, 43046720, 387420488};
7374

75+
const static uint32_t KEYS_PER_NODE_32 = 16;
76+
const static uint32_t MAX_LEVEL_32 = 7;
77+
static constexpr uint32_t NODES_PER_LEVEL_32[MAX_LEVEL_32] = {16, 272, 4624, 78608, 1336336, 22717712, 386200304};
78+
static constexpr uint32_t LEVEL_START_ID_32[MAX_LEVEL_32] = { 0, 16, 288, 4912, 83520, 1419856, 24137568};
7479

80+
template<typename KeyType>
7581
struct DefaultInnerSearch {
76-
static uint32_t inner_search(const uint64_t *base, uint64_t x) {
77-
uint8_t mask = 0;
82+
static_assert(std::is_same<KeyType, uint32_t>::value || std::is_same<KeyType, uint64_t>::value,
83+
"KeyType must be uint32_t or uint64_t");
84+
85+
static constexpr uint32_t KEYS = std::is_same<KeyType, uint64_t>::value ? KEYS_PER_NODE_64 : KEYS_PER_NODE_32;
86+
87+
static uint32_t inner_search(const KeyType *base, KeyType x) {
88+
uint32_t mask = 0;
7889
#pragma GCC unroll 20
79-
for (uint32_t i = 0; i < ORDER; i++) {
90+
for (uint32_t i = 0; i < KEYS; i++) {
8091
mask |= ( (base[i] <= x) << i );
8192
}
8293
return __builtin_popcount(mask);
8394
}
8495
};
8596

8697
#ifdef __x86_64__
98+
template<typename KeyType>
8799
struct Avx512InnerSearch {
100+
static_assert(std::is_same<KeyType, uint32_t>::value || std::is_same<KeyType, uint64_t>::value,
101+
"KeyType must be uint32_t or uint64_t");
102+
88103
#ifdef __clang__
89104
#pragma clang attribute push (__attribute__((target("avx512f"))), apply_to=function)
90105
#else // __GNUC__
91106
#pragma GCC push_options
92107
#pragma GCC target ("avx512f")
93108
#endif
94109

95-
static uint32_t inner_search(const uint64_t *base, uint64_t x) {
110+
template<typename T = KeyType>
111+
static typename std::enable_if<std::is_same<T, uint64_t>::value, uint32_t>::type
112+
inner_search(const KeyType *base, KeyType x) {
96113
__m512i vx = _mm512_set1_epi64(x);
97114
__m512i data = _mm512_load_si512(base);
98115
uint8_t mask = _mm512_cmp_epu64_mask(vx, data, _MM_CMPINT_GE);
99116
return __builtin_popcount(mask);
100117
}
101118

119+
template<typename T = KeyType>
120+
static typename std::enable_if<std::is_same<T, uint32_t>::value, uint32_t>::type
121+
inner_search(const KeyType *base, KeyType x) {
122+
__m512i vx = _mm512_set1_epi32(x);
123+
__m512i data = _mm512_load_si512(base);
124+
__mmask16 mask = _mm512_cmp_epu32_mask(vx, data, _MM_CMPINT_GE);
125+
return __builtin_popcount(mask);
126+
}
127+
102128
#ifdef __clang__
103129
#pragma clang attribute pop
104130
#else // __GNUC__
105131
#pragma GCC pop_options
106132
#endif
107133
};
108134
#else // __x86_64__
109-
using Avx512InnerSearch = DefaultInnerSearch;
135+
template<typename KeyType>
136+
using Avx512InnerSearch = DefaultInnerSearch<KeyType>;
110137
#endif
111138

139+
template<typename KeyType>
112140
class LinearizedBptree {
141+
static_assert(std::is_same<KeyType, uint32_t>::value || std::is_same<KeyType, uint64_t>::value,
142+
"KeyType must be uint32_t or uint64_t");
113143
public:
144+
static constexpr const uint32_t KEYS_PER_NODE = std::is_same<KeyType, uint64_t>::value ? KEYS_PER_NODE_64 : KEYS_PER_NODE_32;
145+
static constexpr const uint32_t MAX_LEVEL = std::is_same<KeyType, uint64_t>::value ? MAX_LEVEL_64 : MAX_LEVEL_32;
146+
static constexpr const uint32_t *NODES_PER_LEVEL = std::is_same<KeyType, uint64_t>::value ? NODES_PER_LEVEL_64 : NODES_PER_LEVEL_32;
147+
static constexpr const uint32_t *LEVEL_START_ID = std::is_same<KeyType, uint64_t>::value ? LEVEL_START_ID_64 : LEVEL_START_ID_32;
148+
114149
uint64_t N;
115-
uint64_t *node = nullptr;
150+
KeyType *node = nullptr;
116151
int32_t DEPTH = -1;
117152

118153
LinearizedBptree() {}
@@ -139,9 +174,9 @@ class LinearizedBptree {
139174
LOG_ERROR_RETURN(EINVAL, -1, "linearized bptree not used: too many mappings");
140175
}
141176

142-
N = (LEVEL_START_ID[DEPTH-1] + mapping_size + ORDER - 1) / ORDER * ORDER;
143-
LOG_INFO("building Linearized B+tree ", VALUE(DEPTH), VALUE(mapping_size), VALUE(N));
144-
auto ret = posix_memalign((void**)&node, 64, N*sizeof(uint64_t));
177+
N = (LEVEL_START_ID[DEPTH-1] + mapping_size + KEYS_PER_NODE - 1) / KEYS_PER_NODE * KEYS_PER_NODE;
178+
LOG_INFO("building Linearized B+tree ", VALUE(DEPTH), VALUE(mapping_size), VALUE(N), VALUE(sizeof(KeyType)));
179+
auto ret = posix_memalign((void**)&node, 64, N*sizeof(KeyType));
145180
if (ret != 0) {
146181
LOG_ERRNO_RETURN(ENOBUFS, -1, "linearized bptree not used: failed to alloc memory");
147182
}
@@ -157,29 +192,29 @@ class LinearizedBptree {
157192
while (p < N)
158193
node[p++] = -1;
159194

160-
auto G = ORDER;
195+
auto G = KEYS_PER_NODE;
161196
for (auto level = DEPTH-1; level > 0; level--) {
162197
auto pos = LEVEL_START_ID[level - 1];
163-
for (uint32_t i = 0; i < leaf_size; i += G * (ORDER + 1)) {
164-
for (uint32_t j = 1; j <= ORDER; j++) {
198+
for (uint32_t i = 0; i < leaf_size; i += G * (KEYS_PER_NODE + 1)) {
199+
for (uint32_t j = 1; j <= KEYS_PER_NODE; j++) {
165200
uint32_t lower_id = leaf_start + i + G * j;
166201
node[pos++] = (lower_id < N) ? node[lower_id] : -1;
167202
}
168203
}
169-
G *= (ORDER + 1);
204+
G *= (KEYS_PER_NODE + 1);
170205
}
171206
LOG_INFO("building Linearized B+tree done");
172207
return 0;
173208
}
174209

175210
template<typename InnerSearchImpl>
176-
uint32_t search(const uint64_t x) const {
211+
uint32_t search(const KeyType x) const {
177212
uint32_t res = 0;
178213
#pragma GCC unroll 20
179214
for (int i = DEPTH; i > 1; --i) {
180215
auto node_base = node + res;
181216
uint32_t c = InnerSearchImpl::inner_search(node_base, x);
182-
res = (ORDER+1)*res + (c+1)*ORDER;
217+
res = (KEYS_PER_NODE+1)*res + (c+1)*KEYS_PER_NODE;
183218
}
184219
auto node_base = node + res;
185220
res += InnerSearchImpl::inner_search(node_base, x);
@@ -291,51 +326,44 @@ class Index : public IMemoryIndex {
291326
UNIMPLEMENTED_POINTER(IMemoryIndex *make_read_only_index() const override);
292327
};
293328

329+
template<typename KeyType, template<typename> class SearchPolicy = DefaultInnerSearch>
294330
class IndexLBPT : public Index {
331+
static_assert(std::is_same<KeyType, uint32_t>::value || std::is_same<KeyType, uint64_t>::value,
332+
"KeyType must be uint32_t or uint64_t");
295333
public:
296-
LinearizedBptree *lbpt = nullptr;
334+
using LBPTree = LinearizedBptree<KeyType>;
335+
LBPTree *lbpt = nullptr;
297336

298337
~IndexLBPT() {
299338
safe_delete(lbpt);
300339
}
301340

302-
IndexLBPT(vector<SegmentMapping> &&m, uint64_t vsize, LinearizedBptree *lbpt)
341+
IndexLBPT(vector<SegmentMapping> &&m, uint64_t vsize, LBPTree *lbpt)
303342
: Index(std::move(m), vsize), lbpt(lbpt) {
304343
}
305344

306345
size_t lookup(Segment s, SegmentMapping *pm, size_t n) const override {
307346
if (s.length == 0)
308347
return 0;
309-
auto lb = pbegin + lbpt->search<DefaultInnerSearch>(s.offset);;
348+
349+
auto lb = this->pbegin + this->lbpt->template search<SearchPolicy<KeyType>>(static_cast<KeyType>(s.offset));
310350
if (lb->end() <= s.offset)
311351
lb++;
312352

313-
auto m = copy_n(lb, pend, s.end(), pm, n);
353+
auto m = copy_n(lb, this->pend, s.end(), pm, n);
314354
trim_edge_mappings(pm, m, s);
315355
return m;
316356
}
317357
};
318358

319-
class IndexLBPTAcc : public IndexLBPT {
320-
public:
321-
IndexLBPTAcc(vector<SegmentMapping> &&m, uint64_t vsize, LinearizedBptree *lbpt)
322-
: IndexLBPT(std::move(m), vsize, lbpt) {
323-
}
324-
325-
size_t lookup(Segment s, SegmentMapping *pm, size_t n) const override {
326-
if (s.length == 0)
327-
return 0;
328-
auto lb = pbegin + lbpt->search<Avx512InnerSearch>(s.offset);
329-
if (lb->end() <= s.offset)
330-
lb++;
331-
auto m = copy_n(lb, pend, s.end(), pm, n);
332-
trim_edge_mappings(pm, m, s);
333-
return m;
334-
}
335-
};
359+
template<typename KeyType>
360+
using IndexLBPTAcc = IndexLBPT<KeyType, Avx512InnerSearch>;
336361

362+
template <typename KeyType>
337363
static inline Index* new_index_with_lineriazed_bptree(vector<SegmentMapping> &&m, uint64_t vsize = 0) {
338-
auto tree = new LinearizedBptree();
364+
static_assert(std::is_same<KeyType, uint32_t>::value || std::is_same<KeyType, uint64_t>::value,
365+
"KeyType must be uint32_t or uint64_t");
366+
auto tree = new LinearizedBptree<KeyType>();
339367
if (tree->build(m) < 0) {
340368
delete tree;
341369
LOG_WARN("failed to build linearized b+tree, failover to binary search");
@@ -344,9 +372,9 @@ static inline Index* new_index_with_lineriazed_bptree(vector<SegmentMapping> &&m
344372

345373
if (is_avx512f_supported()) {
346374
LOG_INFO("using accelerated search for linearized b+tree");
347-
return new IndexLBPTAcc(std::move(m), vsize, tree);
375+
return new IndexLBPTAcc<KeyType>(std::move(m), vsize, tree);
348376
}
349-
return new IndexLBPT(std::move(m), vsize, tree);
377+
return new IndexLBPT<KeyType>(std::move(m), vsize, tree);
350378
}
351379

352380
class LevelIndex : public Index {
@@ -887,6 +915,11 @@ IMemoryIndex *merge_memory_indexes(const IMemoryIndex **pindexes, size_t n) {
887915
mapping.reserve(pi[0]->size());
888916
merge_indexes(0, mapping, pi, n, 0, UINT64_MAX);
889917

890-
return new_index_with_lineriazed_bptree(std::move(mapping), pindexes[0]->vsize());
918+
if (pindexes[0]->vsize() < static_cast<uint64_t>(UINT32_MAX) * ALIGNMENT
919+
&& mapping.size() < NODES_PER_LEVEL_32[MAX_LEVEL_32-1]) {
920+
return new_index_with_lineriazed_bptree<uint32_t>(std::move(mapping), pindexes[0]->vsize());
921+
}
922+
923+
return new_index_with_lineriazed_bptree<uint64_t>(std::move(mapping), pindexes[0]->vsize());
891924
}
892925
} // namespace LSMT

0 commit comments

Comments
 (0)