Line data Source code
1 : #include "fd_wsample.h"
2 : #include <math.h> /* For sqrt */
3 : #if FD_HAS_AVX512
4 : #include "../../util/simd/fd_avx512.h"
5 : #endif
6 :
7 64680136 : #define R 9
8 : /* This sampling problem is an interesting one from a performance
9 : perspective. There are lots of interesting approaches. The
10 : header/implementation split is designed to give lots of flexibility
11 : for future optimization. The current implementation uses radix 9
12 : tree with all the leaves on the bottom. */
13 :
14 : /* I'm not sure exactly how to classify the tree that this
15 : implementation uses, but it's something like a B-tree with some
16 : tricks from binary heaps. In particular, like a B-tree, each node
17 : stores several keys, where the keys are the cumulative sums of
18 : subtrees, called left sums below. Like a binary heap, it is
19 : pointer-free, and it is stored implicitly in a flat array. the
20 : typical way that a binary heap is stored. Specifically, the root is
21 : at index 0, and node i's children are at Ri+1, Ri+2, ... Ri+R, where
22 : R is the radix. The leaves are all at the same level, at the bottom,
23 : stored implicitly, which means that a search can be done with almost
24 : no branch mispredictions.
25 :
26 : As an example, suppose R=5 and our tree contains weights 100, 80, 50,
27 : 31, 27, 14, and 6. Because that is 7 nodes, the height is set to
28 : ceil(log_5(7))=2.
29 :
30 : Root (0)
31 : / \
32 : Child (1) Child (2) -- -- --
33 : / / | \ \ | \
34 : 100, 80, 50, 31, 27 14, 6
35 :
36 : The first child of the root, node 1, has left_sum values
37 : |100|180|230|261|. Note that we only store R-1 values, and so the
38 : last child, 27, does not feature in the sums; the search process
39 : handles it implicitly, as we'll see. The second child of the root,
40 : node 2, has left sum values |14|20|20|20|. Then the root node, node
41 : 0, has left sum values |288|308|308|308|.
42 :
43 : The total sum is 308, so we start by drawing a random values in the
44 : range [0, 308). We see which child that falls under, adjust our
45 : random value, and repeat, until we reach a leaf.
46 :
47 : In general, if a node has children (left to right) with subtree
48 : weights a, b, c, d, e. Then the left sums are |a|a+b|a+b+c|a+b+c+d|
49 : and the full sum S=a+b+c+d+e. For a query value of x in [0, S), we
50 : want to pick:
51 : child 0 if x < a
52 : child 1 if a <= x < a+b
53 : child 2 if a+b <= x < a+b+c
54 : child 3 if a+b+c <= x < a+b+c+d
55 : child 4 if a+b+c+d<= x
56 :
57 : Which is equivalent to choosing child
58 : (a<=x) + (a+b<=x) + (a+b+c<=x) + (a+b+c+d<=x)
59 : which can be computed branchlessly. The value of e only comes into
60 : play in the restriction that x is in [0, S), and e is included in the
61 : value of S.
62 :
63 : There are two details left to discuss in order to search recursively.
64 : First, in order to remove an element for sampling without
65 : replacement, we need to know the weight of the element we're
66 : removing. As with e above, the weights may not be stored explicitly,
67 : but as long as we keep track of the weight of the current subtree as
68 : we search recursively, we can obtain the weight without much work.
69 : Thus, we need to know the sum S' of the chosen subtree, i in [0, R).
70 : For j in [-1, R), define
71 : / 0 if j==-1
72 : l[j] = | left_sum[ j ] if j in [0, R-1)
73 : \ S if j==R-1
74 : Essentially, we're computing the natural extension values for
75 : left_sum on the left and right side. Then observe that S' = l[i] -
76 : l[i-1].
77 :
78 : Secondly, in order to search recursively, we need to adjust the query
79 : value to put it into [0, S'). Specifically, we need to subtract off
80 : the sum of all the children to the left of the child we've chosen.
81 : If we've chosen child i, then that's just l[i-1].
82 :
83 : All the above extends easily to any radix, but 3, 5, and 9 = 1+2^n
84 : are the natural ones. I only tried 5 and 9, and 9 was substantially
85 : faster.
86 :
87 : It's worth noting that storing left sums means that you have to do
88 : the most work when you update the left-most child. Because we
89 : typically store the weights largest to smallest, the left-most child
90 : is the one we delete the most frequently, so now our most common case
91 : (probability-wise) and our worst case (performance-wise) are the
92 : same, which seems bad. I initially implemented this using right sums
93 : instead of left to address this problem. They're less intuitive, but
94 : also work. However, I found that having an unpredictable loop in the
95 : deletion method was far worse than just updating each element, which
96 : means that the "less work" advantage of right sums went away. */
97 :
98 : struct __attribute__((aligned(8UL*(R-1UL)))) tree_ele {
99 : /* left_sum stores the cumulative weight of the subtrees at this
100 : node. See the long note above for more information. */
101 : ulong left_sum[ R-1 ];
102 : };
103 : typedef struct tree_ele tree_ele_t;
104 :
105 : struct __attribute__((aligned(64UL))) fd_wsample_private {
106 : ulong total_cnt;
107 : ulong total_weight;
108 : ulong unremoved_cnt;
109 : ulong unremoved_weight; /* Initial value for S explained above */
110 :
111 : /* internal_node_cnt and height are both determined by the number of
112 : leaves in the original tree, via the following formulas:
113 : height = ceil(log_r(leaf_cnt))
114 : internal_node_cnt = sum_{i=0}^{height-1} R^i
115 : height and internal_node_cnt both exclude the leaves, which are
116 : only implicit.
117 :
118 : All the math seems to disallow leaf_cnt==0, but for convenience, we
119 : do allow it. height==internal_node_cnt==0 in that case.
120 :
121 : height actually fits in a uchar. Storing as ulong is more natural,
122 : but we don't want to spill over into another cache line.
123 :
124 : If poisoned_mode==1, then all sample calls will return poisoned.*/
125 : ulong internal_node_cnt;
126 : ulong poisoned_weight;
127 : uint height;
128 : char restore_enabled;
129 : char poisoned_mode;
130 : /* Two bytes of padding here */
131 :
132 : fd_chacha_rng_t * rng;
133 :
134 : /* tree: Here's where the actual tree is stored, at indices [0,
135 : internal_node_cnt). The indexing scheme is explained in the long
136 : comment above.
137 :
138 : If restore_enabled==1, then indices [internal_node_cnt+1,
139 : 2*internal_node_cnt+1) store a copy of the tree after construction
140 : but before any deletion so that restoring deleted elements can be
141 : implemented as a memcpy.
142 :
143 : The tree itself is surrounded by two dummy elements, dummy, and
144 : tree[internal_node_cnt], that aren't actually used. This is
145 : because searching the tree branchlessly involves some out of bounds
146 : reads, and although the value is immediately discarded, it's better
147 : to know where exactly those reads might go. */
148 : tree_ele_t dummy;
149 : tree_ele_t tree[];
150 : };
151 :
152 : typedef struct fd_wsample_private fd_wsample_t;
153 :
154 :
155 : FD_FN_CONST ulong
156 1941 : fd_wsample_align( void ) {
157 1941 : return 64UL;
158 1941 : }
159 :
160 : /* Returns -1 on failure */
161 : static inline int
162 : compute_height( ulong leaf_cnt,
163 : ulong * out_height,
164 1294 : ulong * out_internal_cnt ) {
165 : /* If we made the sentinel values larger, we could take larger
166 : leaf_cnt, but at INT_MAX this would still use an astronomical
167 : amount of memory, so we retain this max for the moment. */
168 1294 : if( FD_UNLIKELY( leaf_cnt >= (ulong)INT_MAX ) ) return -1;
169 :
170 1294 : ulong height = 0;
171 1294 : ulong internal = 0UL;
172 1294 : ulong powRh = 1UL; /* = R^height */
173 1770 : while( leaf_cnt>powRh ) {
174 476 : internal += powRh;
175 476 : powRh *= R;
176 476 : height++;
177 476 : }
178 1294 : *out_height = height;
179 1294 : *out_internal_cnt = internal;
180 1294 : return 0;
181 1294 : }
182 :
183 : FD_FN_CONST ulong
184 647 : fd_wsample_footprint( ulong ele_cnt, int restore_enabled ) {
185 647 : ulong height;
186 647 : ulong internal_cnt;
187 : /* Computing the closed form of the sum in compute_height, we get
188 : internal_cnt = 1/8 * (9^ceil(log_9( ele_cnt ) ) - 1)
189 : x <= ceil( x ) < x+1
190 : 1/8 * ele_cnt - 1/8 <= internal_cnt < 9/8 * ele_cnt - 1/8
191 : */
192 647 : if( FD_UNLIKELY( compute_height( ele_cnt, &height, &internal_cnt ) ) ) return 0UL;
193 647 : return sizeof(fd_wsample_t) + ((restore_enabled?2UL:1UL)*internal_cnt + 1UL)*sizeof(tree_ele_t);
194 647 : }
195 :
196 : fd_wsample_t *
197 647 : fd_wsample_join( void * shmem ) {
198 647 : if( FD_UNLIKELY( !shmem ) ) {
199 0 : FD_LOG_WARNING(( "NULL shmem" ));
200 0 : return NULL;
201 0 : }
202 :
203 647 : if( FD_UNLIKELY( !fd_ulong_is_aligned( (ulong)shmem, fd_wsample_align() ) ) ) {
204 0 : FD_LOG_WARNING(( "misaligned shmem" ));
205 0 : return NULL;
206 0 : }
207 647 : return (fd_wsample_t *)shmem;
208 647 : }
209 :
210 : /* Note: The following optimization insights are not used in this
211 : high radix implementation. Performance in the deletion case is much
212 : more important than in the non-deletion case, and it's not clear how
213 : to translate this. I'm leaving the code and comment because it is a
214 : useful and non-trivial insight. */
215 : #if 0
216 : /* If we assume the probability of querying node i is proportional to
217 : 1/i, then observe that the midpoint of the probability mass in the
218 : continuous approximation is the solution to (in Mathematica syntax):
219 :
220 : Integrate[ 1/i, {i, lo, hi}] = 2*Integrate[ 1/i, {i, lo, mid} ]
221 :
222 : which gives mid = sqrt(lo*hi). This is in contrast to when the
223 : integrand is a constant, which gives the normal binary search rule:
224 : mid=(lo+hi)/2.
225 :
226 : Thus, we want the search to follow this modified binary search rule,
227 : since that'll approximately split the stake weight/probability mass
228 : in half at each step.
229 :
230 : This is almost as nice to work with mathematically as the normal
231 : binary search rule. The jth entry from the left at level k is the
232 : region [ N^((1/2^k)*j), N^((1/2^k)*(j+1)) ). We're basically doing
233 : binary search in the log domain.
234 :
235 : Rather than trying to compute these transcendental functions, this
236 : simple recursive implementation gives the treap the right shape by
237 : setting prio very carefully, since higher prio means closer to the
238 : root. This may be a slight abuse of a treap, but it's easier than
239 : implementing a whole custom tree for just this purpose.
240 :
241 : We want to be careful to ensure that the recursion is tightly
242 : bounded. From a continuous perspective, that's not a problem: the
243 : widest interval at level k is the last one, and we break when the
244 : interval's width is less than 1. Solving 1=N-N^(1-(1/2^k)) for k
245 : yields k=-lg(1-log(N-1)/log(N)). k is monotonically increasing as a
246 : function of N, which means that for all N,
247 : k <= -lg(1-log(N_max-1)/log(N_max)) < 37
248 : since N_max=2^32-1.
249 :
250 : The math is more complicated with rounding and finite precision, but
251 : sqrt(lo*hi) is very different from lo and hi unless lo and hi are
252 : approximately the same. In that case, lo<mid and mid<hi ensures that
253 : both intervals are strictly smaller than the interval they came from,
254 : which prevents an infinite loop. */
255 : static inline void
256 : seed_recursive( treap_ele_t * pool,
257 : uint lo,
258 : uint hi,
259 : uint prio ) {
260 : uint mid = (uint)(sqrtf( (float)lo*(float)hi ) + 0.5f);
261 : if( (lo<mid) & (mid<hi) ) {
262 : /* since we start with lo=1, shift by 1 */
263 : pool[mid-1U].prio = prio;
264 : seed_recursive( pool, lo, mid, prio-1U );
265 : seed_recursive( pool, mid, hi, prio-1U );
266 : }
267 : }
268 : #endif
269 :
270 :
271 : void *
272 : fd_wsample_new_init( void * shmem,
273 : fd_chacha_rng_t * rng,
274 : ulong ele_cnt,
275 : int restore_enabled,
276 647 : int opt_hint ) {
277 647 : if( FD_UNLIKELY( !shmem ) ) {
278 0 : FD_LOG_WARNING(( "NULL shmem" ));
279 0 : return NULL;
280 0 : }
281 :
282 647 : if( FD_UNLIKELY( !fd_ulong_is_aligned( (ulong)shmem, fd_wsample_align() ) ) ) {
283 0 : FD_LOG_WARNING(( "misaligned shmem" ));
284 0 : return NULL;
285 0 : }
286 :
287 647 : ulong height;
288 647 : ulong internal_cnt;
289 647 : if( FD_UNLIKELY( compute_height( ele_cnt, &height, &internal_cnt ) ) ) {
290 0 : FD_LOG_WARNING(( "bad ele_cnt" ));
291 0 : return NULL;
292 0 : }
293 :
294 647 : fd_wsample_t * sampler = (fd_wsample_t *)shmem;
295 :
296 647 : sampler->total_cnt = 0UL;
297 647 : sampler->total_weight = 0UL;
298 647 : sampler->unremoved_cnt = 0UL;
299 647 : sampler->unremoved_weight = 0UL;
300 647 : sampler->internal_node_cnt = internal_cnt;
301 647 : sampler->poisoned_weight = 0UL;
302 647 : sampler->height = (uint)height;
303 647 : sampler->restore_enabled = (char)!!restore_enabled;
304 647 : sampler->poisoned_mode = 0;
305 647 : sampler->rng = rng;
306 :
307 647 : fd_memset( sampler->tree, (char)0, internal_cnt*sizeof(tree_ele_t) );
308 :
309 647 : (void)opt_hint; /* Not used at the moment */
310 :
311 647 : return shmem;
312 647 : }
313 :
314 : void *
315 : fd_wsample_new_add( void * shmem,
316 1278 : ulong weight ) {
317 1278 : fd_wsample_t * sampler = (fd_wsample_t *)shmem;
318 1278 : if( FD_UNLIKELY( !sampler ) ) return NULL;
319 :
320 1278 : if( FD_UNLIKELY( weight==0UL ) ) {
321 0 : FD_LOG_WARNING(( "zero weight entry found" ));
322 0 : return NULL;
323 0 : }
324 1278 : if( FD_UNLIKELY( sampler->total_weight+weight<weight ) ) {
325 0 : FD_LOG_WARNING(( "total weight too large" ));
326 0 : return NULL;
327 0 : }
328 :
329 1278 : tree_ele_t * tree = sampler->tree;
330 1278 : ulong i = sampler->internal_node_cnt + sampler->unremoved_cnt;
331 :
332 2427 : for( ulong h=0UL; h<sampler->height; h++ ) {
333 1149 : ulong parent = (i-1UL)/R;
334 1149 : ulong child_idx = i-1UL - R*parent; /* in [0, R) */
335 8574 : for( ulong k=child_idx; k<R-1UL; k++ ) tree[ parent ].left_sum[ k ] += weight;
336 1149 : i = parent;
337 1149 : }
338 :
339 1278 : sampler->unremoved_cnt++;
340 1278 : sampler->total_cnt++;
341 1278 : sampler->unremoved_weight += weight;
342 1278 : sampler->total_weight += weight;
343 :
344 1278 : return shmem;
345 1278 : }
346 :
347 : void *
348 : fd_wsample_new_fini( void * shmem,
349 647 : ulong poisoned_weight ) {
350 647 : fd_wsample_t * sampler = (fd_wsample_t *)shmem;
351 647 : if( FD_UNLIKELY( !sampler ) ) return NULL;
352 :
353 647 : if( FD_UNLIKELY( sampler->total_weight+poisoned_weight<sampler->total_weight ) ) {
354 0 : FD_LOG_WARNING(( "poisoned_weight caused overflow" ));
355 0 : return NULL;
356 0 : }
357 :
358 647 : sampler->poisoned_weight = poisoned_weight;
359 :
360 647 : if( sampler->restore_enabled ) {
361 : /* Copy the sampler to make restore fast. */
362 0 : fd_memcpy( sampler->tree+sampler->internal_node_cnt+1UL, sampler->tree, sampler->internal_node_cnt*sizeof(tree_ele_t) );
363 0 : }
364 :
365 647 : return (void *)sampler;
366 647 : }
367 :
368 : void *
369 647 : fd_wsample_leave( fd_wsample_t * sampler ) {
370 647 : if( FD_UNLIKELY( !sampler ) ) {
371 0 : FD_LOG_WARNING(( "NULL sampler" ));
372 0 : return NULL;
373 0 : }
374 :
375 647 : return (void *)sampler;
376 647 : }
377 :
378 : void *
379 647 : fd_wsample_delete( void * shmem ) {
380 647 : if( FD_UNLIKELY( !shmem ) ) {
381 0 : FD_LOG_WARNING(( "NULL shmem" ));
382 0 : return NULL;
383 0 : }
384 :
385 647 : if( FD_UNLIKELY( !fd_ulong_is_aligned( (ulong)shmem, fd_wsample_align() ) ) ) {
386 0 : FD_LOG_WARNING(( "misaligned shmem" ));
387 0 : return NULL;
388 0 : }
389 647 : return shmem;
390 647 : }
391 :
392 : void
393 : fd_wsample_seed_rng( fd_wsample_t * sampler,
394 0 : uchar seed[ 32 ] ) {
395 0 : fd_chacha_rng_init( sampler->rng, seed, FD_CHACHA_RNG_ALGO_CHACHA8 );
396 0 : }
397 :
398 : fd_wsample_t *
399 0 : fd_wsample_restore_all( fd_wsample_t * sampler ) {
400 0 : if( FD_UNLIKELY( !sampler->restore_enabled ) ) return NULL;
401 :
402 0 : sampler->unremoved_weight = sampler->total_weight;
403 0 : sampler->unremoved_cnt = sampler->total_cnt;
404 0 : sampler->poisoned_mode = 0;
405 :
406 0 : fd_memcpy( sampler->tree, sampler->tree+sampler->internal_node_cnt+1UL, sampler->internal_node_cnt*sizeof(tree_ele_t) );
407 0 : return sampler;
408 0 : }
409 :
410 : #define fd_ulong_if_force( c, t, f ) (__extension__({ \
411 : ulong result; \
412 : __asm__( "testl %1, %1; \n\t" \
413 : "movq %3, %0; \n\t" \
414 : "cmovne %2, %0; \n\t" \
415 : : "=&r"(result) \
416 : : "r"(c), "rm"(t), "rmi"(f) \
417 : : "cc" ); \
418 : result; \
419 : }))
420 :
421 : /* Helper methods for sampling functions */
422 : typedef struct { ulong idx; ulong weight; } idxw_pair_t; /* idx in [0, total_cnt) */
423 :
424 : /* Assumes query in [0, unremoved_weight), which implies
425 : unremoved_weight>0, so the tree can't be empty. */
426 : static inline idxw_pair_t
427 : fd_wsample_map_sample_i( fd_wsample_t const * sampler,
428 19276806 : ulong query ) {
429 19276806 : tree_ele_t const * tree = sampler->tree;
430 :
431 19276806 : ulong cursor = 0UL;
432 19276806 : ulong S = sampler->unremoved_weight;
433 35444003 : for( ulong h=0UL; h<sampler->height; h++ ) {
434 16167197 : tree_ele_t const * e = tree+cursor;
435 16167197 : ulong x = query;
436 16167197 : ulong child_idx = 0UL;
437 :
438 16167197 : #if FD_HAS_AVX512 && R==9
439 16167197 : __mmask8 mask = _mm512_cmple_epu64_mask( wwv_ld( e->left_sum ), wwv_bcast( x ) );
440 16167197 : child_idx = (ulong)fd_uchar_popcnt( mask );
441 : #else
442 : for( ulong i=0UL; i<R-1UL; i++ ) child_idx += (ulong)(e->left_sum[ i ]<=x);
443 : #endif
444 :
445 : /* See the note at the top of this file for the explanation of l[i]
446 : and l[i-1]. Because this is fd_ulong_if and not a ternary, these
447 : can read/write out of what you would think the appropriate bounds
448 : are. The dummy elements, as described along with tree makes this
449 : safe. */
450 : #if 0
451 : ulong li = fd_ulong_if( child_idx<R-1UL, e->left_sum[ child_idx ], S );
452 : ulong lm1 = fd_ulong_if( child_idx>0UL, e->left_sum[ child_idx-1UL ], 0UL );
453 : #elif 0
454 : ulong li = fd_ulong_if_force( child_idx<R-1UL, e->left_sum[ child_idx ], S );
455 : ulong lm1 = fd_ulong_if_force( child_idx>0UL, e->left_sum[ child_idx-1UL ], 0UL );
456 : #else
457 16167197 : ulong * temp = (ulong *)e->left_sum;
458 16167197 : ulong orig_m1 = temp[ -1 ]; ulong orig_Rm1 = temp[ R-1UL ];
459 16167197 : temp[ -1 ] = 0UL; temp[ R-1UL ] = S;
460 16167197 : ulong li = temp[ child_idx ];
461 16167197 : ulong lm1 = temp[ (long)child_idx-1L ];
462 16167197 : temp[ -1 ] = orig_m1; temp[ R-1UL ] = orig_Rm1;
463 16167197 : #endif
464 :
465 16167197 : query -= lm1;
466 16167197 : S = li - lm1;
467 16167197 : cursor = R*cursor + child_idx + 1UL;
468 16167197 : }
469 19276806 : idxw_pair_t to_return = { .idx = cursor - sampler->internal_node_cnt, .weight = S };
470 19276806 : return to_return;
471 19276806 : }
472 :
473 : ulong
474 : fd_wsample_map_sample( fd_wsample_t * sampler,
475 19808350 : ulong query ) {
476 19808350 : return fd_wsample_map_sample_i( sampler, query ).idx;
477 19808350 : }
478 :
479 :
480 :
481 : /* Also requires the tree to be non-empty */
482 : static inline void
483 : fd_wsample_remove( fd_wsample_t * sampler,
484 0 : idxw_pair_t to_remove ) {
485 0 : ulong cursor = to_remove.idx + sampler->internal_node_cnt;
486 0 : tree_ele_t * tree = sampler->tree;
487 :
488 0 : ulong height = sampler->height;
489 0 : for( ulong h=0UL; h<height; h++ ) {
490 0 : ulong parent = (cursor-1UL)/R;
491 0 : ulong child_idx = cursor-1UL - R*parent; /* in [0, R) */
492 0 : #if FD_HAS_AVX512 && R==9
493 0 : wwv_t weight = wwv_bcast( to_remove.weight );
494 0 : wwv_t left_sum = wwv_ld( tree[ parent ].left_sum );
495 0 : __m128i _child_idx = _mm_set1_epi16( (short) child_idx );
496 0 : __mmask8 mask = _mm_cmplt_epi16_mask( _child_idx, _mm_setr_epi16( 1, 2, 3, 4, 5, 6, 7, 8 ) );
497 0 : left_sum = _mm512_mask_sub_epi64( left_sum, mask, left_sum, weight );
498 0 : wwv_st( tree[ parent ].left_sum, left_sum );
499 : #elif 0
500 : for( ulong k=0UL; k<R-1UL; k++ ) tree[ parent ].left_sum[ k ] -= fd_ulong_if( child_idx<=k, to_remove.weight, 0UL );
501 : #elif 0
502 : for( ulong k=0UL; k<R-1UL; k++ ) tree[ parent ].left_sum[ k ] -= fd_ulong_if_force( child_idx<=k, to_remove.weight, 0UL );
503 : #elif 1
504 : /* The compiler loves inserting a difficult to predict branch for
505 : fd_ulong_if, but this forces it not to do that. */
506 : for( ulong k=0UL; k<R-1UL; k++ ) tree[ parent ].left_sum[ k ] -= (ulong)(((long)(child_idx - k - 1UL))>>63) & to_remove.weight;
507 : #else
508 : /* This version does the least work, but has a hard-to-predict
509 : branch. The branchless versions are normally substantially
510 : faster. */
511 : for( ulong k=child_idx; k<R-1UL; k++ ) tree[ parent ].left_sum[ k ] -= to_remove.weight;
512 : #endif
513 0 : cursor = parent;
514 0 : }
515 0 : sampler->unremoved_cnt--;
516 0 : sampler->unremoved_weight -= to_remove.weight;
517 0 : }
518 :
519 : static inline ulong
520 : fd_wsample_find_weight( fd_wsample_t const * sampler,
521 0 : ulong idx /* in [0, total_cnt) */) {
522 : /* The fact we don't store the weights explicitly makes this function
523 : more complicated, but this is not used very frequently. */
524 0 : tree_ele_t const * tree = sampler->tree;
525 0 : ulong cursor = idx + sampler->internal_node_cnt;
526 :
527 : /* Initialize to the 0 height case */
528 0 : ulong lm1 = 0UL;
529 0 : ulong li = sampler->unremoved_weight;
530 :
531 0 : for( ulong h=0UL; h<sampler->height; h++ ) {
532 0 : ulong parent = (cursor-1UL)/R;
533 0 : ulong child_idx = cursor-1UL - R*parent; /* in [0, R) */
534 :
535 : /* If child_idx < R-1, we can compute the weight easily. If
536 : child_idx==R-1, the computation is S - left_sum[ R-2 ], but we
537 : don't know S, so we need to continue up the tree. */
538 0 : lm1 += fd_ulong_if( child_idx>0UL, tree[ parent ].left_sum[ (long)child_idx-1L ], 0UL );
539 0 : if( FD_LIKELY( child_idx<R-1UL ) ) {
540 0 : li = tree[ parent ].left_sum[ child_idx ];
541 0 : break;
542 0 : }
543 :
544 0 : cursor = parent;
545 0 : }
546 :
547 0 : return li - lm1;
548 0 : }
549 :
550 : void
551 : fd_wsample_remove_idx( fd_wsample_t * sampler,
552 0 : ulong idx ) {
553 :
554 0 : ulong weight = fd_wsample_find_weight( sampler, idx );
555 0 : idxw_pair_t r = { .idx = idx, .weight = weight };
556 0 : fd_wsample_remove( sampler, r );
557 0 : }
558 :
559 :
560 : #if FD_HAS_AVX512
561 :
562 : /* TRAVERSE_LEVEL Takes in and updates s (a vector with all elements set
563 : to the unremoved weight in the current subtree), x or xprime (a
564 : vector with all elements set to the current query value, or the
565 : value+1, respectively), and cursor (a ulong of the current position
566 : in the traversal). See fd_wsample_map_sample_i for more about these
567 : values. Calling this macro height times is the same as calling
568 : fd_wsample_map_sample_i, except for that it declares mask{i} and
569 : cursor{i}, to save work in PROPAGATE_LEVEL, which is like
570 : fd_wsample_remove.
571 :
572 : This only works in the R=9 case, but we'll explain it as if R=5 like
573 : before. Then, v, a vector of tree elements has values
574 : ----------------------------------------
575 : | a | a+b | a+b+c | a+b+c+d |
576 : ----------------------------------------
577 : Then, like before, the desired state is
578 : If... pick, and l_{i-1} and l_i
579 : x < a child 0 0 a
580 : a <= x < a+b child 1 a a+b
581 : a+b <= x < a+b+c child 2 a+b a+b+c
582 : a+b+c <= x < a+b+c+d child 3 a+b+c a+b+c+d
583 : a+b+c+d<= x child 4 a+b+c+d a+b+c+d+e
584 :
585 : Below, there are two pretty different implementations, and this
586 : explains both:
587 :
588 : Implementation 1:
589 : We want to get information about v0<=x in as much as the vector as
590 : fast as possible. We could do some kind of compress and then bcast,
591 : but operations with mask registers are frustratingly slow. Instead,
592 : we first define x'=x+1, so then v0<=x is equivalent to v0-x'<0, or
593 : whether v0-x' has the high bit set. Because of
594 : fd_chacha_rng_ulong_roll's contract, we know x<ULONG_MAX, so forming
595 : x' is safe. We then use a _mm512_permutexvar_epi8 to essentially
596 : broadcast the high byte from each of the ulongs (really we just need
597 : the high bit) to the whole vector. A popcnt gives us our base
598 : values.
599 :
600 : Case compressed_signs popcnt popcnt-1
601 : x < a 0 0 -1
602 : a <= x < a+b 0x80 1 0
603 : a+b <= x < a+b+c 0x8080 2 1
604 : a+b+c <= x < a+b+c+d 0x808080 3 2
605 : a+b+c+d<= x 0x80808080 4 3
606 :
607 : Then, keeping in mind that _mm512_permutex2var_epi64 will select an
608 : element from the second vector if the next highest bit is set, using
609 : popcnt-1 as the selector and (vec, 0) as (a,b) gives l_{i-1}; and
610 : using popcnt as the selector with (vec, s) gives l_i.
611 :
612 : Finally, we just need to produce the mask mask{i}, which we can do
613 : with wwv_lt, off the critical path.
614 :
615 :
616 : Implementation 2
617 : We first use wwv_le to compare v0<=x to form lmask. Then We extract
618 : the most significant bit by lmask^(lmask>>1). Note that below the
619 : bits are written least to most significant, which is backwards of the
620 : normal way.
621 :
622 : Case lmask single_bit ~lmask
623 : x < a 0 0 0 0 0 0 0 0 1 1 1 1
624 : a <= x < a+b 1 0 0 0 1 0 0 0 0 1 1 1
625 : a+b <= x < a+b+c 1 1 0 0 0 1 0 0 0 0 1 1
626 : a+b+c <= x < a+b+c+d 1 1 1 0 0 0 1 0 0 0 0 1
627 : a+b+c+d<= x 1 1 1 1 0 0 0 1 0 0 0 0
628 :
629 : Then we use _mm512_mask_compress_epi64 to move the element
630 : corresponding to the first set bit to the 0th position, or the src
631 : vector if there aren't any set bits. From there, we can use
632 : _mm512_broadcastq_epi64 to fill a vector with that value. Applying
633 : this trick to single_bit gives us l_{i-1} and to ~lmask give l_i. */
634 : #define FD_WSAMPLE_IMPLEMENTATION 2
635 :
636 : #if FD_WSAMPLE_IMPLEMENTATION==1
637 : #define PREPARE() \
638 : ulong cursor = 0UL; \
639 : wwv_t xprime = wwv_bcast( unif+1UL ); \
640 : wwv_t s = wwv_bcast( sampler->unremoved_weight ); \
641 : wwv_t high_bit_mask = wwv_bcast( 0x8000000000000000UL ); \
642 : wwv_t gather_signs_idx = wwv_bcast( 0x070F171F272F373FUL )
643 :
644 : #define TRAVERSE_LEVEL(i) \
645 : __mmask8 mask##i; \
646 : ulong cursor##i = cursor; \
647 : wwv_t vec##i; \
648 : do { \
649 : wwv_t vec = wwv_ld( tree[cursor].left_sum ); \
650 : wwv_t sign = wwv_and( high_bit_mask, wwv_sub( vec, xprime ) ); \
651 : wwv_t compressed_signs = _mm512_permutexvar_epi8( gather_signs_idx, sign ); \
652 : wwv_t popcnt = _mm512_popcnt_epi64( compressed_signs ); \
653 : wwv_t li = _mm512_permutex2var_epi64( vec, popcnt, s ); \
654 : wwv_t lim1 = _mm512_permutex2var_epi64( vec, wwv_sub( popcnt, wwv_one() ), wwv_zero() ); \
655 : mask##i = _mm512_cmpge_epu64_mask( vec, xprime ); \
656 : xprime = wwv_sub( xprime, lim1 ); \
657 : s = wwv_sub( li, lim1 ); \
658 : ulong child_idx = (ulong)_mm_extract_epi64( _mm512_castsi512_si128( popcnt ), 0 ); \
659 : cursor = R*cursor + child_idx + 1UL; \
660 : vec##i = vec; \
661 : } while( 0 )
662 :
663 : #define PROPAGATE_LEVEL(i) \
664 : wwv_st( tree[cursor##i].left_sum, wwv_sub_if( mask##i, vec##i, s, vec##i ) );
665 :
666 : #define FINALIZE() \
667 : do { \
668 : sampler->unremoved_weight -= (ulong)_mm256_extract_epi64( _mm512_castsi512_si256( s ), 0 ); \
669 : sampler->unremoved_cnt--; \
670 : } while( 0 )
671 :
672 : #elif FD_WSAMPLE_IMPLEMENTATION==2
673 :
674 : #define PREPARE() \
675 0 : ulong cursor = 0UL; \
676 0 : wwv_t x = wwv_bcast( unif ); \
677 0 : wwv_t s = wwv_bcast( sampler->unremoved_weight ); \
678 :
679 : #define TRAVERSE_LEVEL(i) \
680 0 : __mmask8 mask##i; \
681 0 : ulong cursor##i = cursor; \
682 0 : wwv_t vec##i; \
683 0 : do { \
684 0 : wwv_t vec = wwv_ld( tree[cursor].left_sum ); \
685 0 : __mmask8 lmask = _mm512_cmple_epu64_mask( vec, x ); \
686 0 : cursor = R*cursor + (ulong)fd_uchar_popcnt( lmask ) + 1UL; \
687 0 : __mmask8 single_bit = _kxor_mask8( lmask, _kshiftri_mask8( lmask, 1U ) ); \
688 0 : wwv_t lim1 = _mm512_broadcastq_epi64( _mm512_castsi512_si128( _mm512_maskz_compress_epi64( single_bit, vec ) ) ); \
689 0 : x = wwv_sub( x, lim1 ); \
690 0 : mask##i = _knot_mask8( lmask ); \
691 0 : wwv_t li = _mm512_broadcastq_epi64( _mm512_castsi512_si128( _mm512_mask_compress_epi64( s, mask##i, vec ) ) ); \
692 0 : s = wwv_sub( li, lim1 ); \
693 0 : vec##i = vec; \
694 0 : } while( 0 )
695 :
696 : #define PROPAGATE_LEVEL(i) \
697 0 : wwv_st( tree[cursor##i].left_sum, wwv_sub_if( mask##i, vec##i, s, vec##i ) );
698 :
699 : #define FINALIZE() \
700 0 : do { \
701 0 : sampler->unremoved_weight -= (ulong)_mm256_extract_epi64( _mm512_castsi512_si256( s ), 0 ); \
702 0 : sampler->unremoved_cnt--; \
703 0 : } while( 0 )
704 :
705 : #endif
706 : #else
707 : #define FD_WSAMPLE_IMPLEMENTATION 0
708 : #endif
709 :
710 : /* For now, implement the _many functions as loops over the single
711 : sample functions. It is possible to do better though. */
712 :
713 : void
714 : fd_wsample_sample_many( fd_wsample_t * sampler,
715 : ulong * idxs,
716 0 : ulong cnt ) {
717 0 : for( ulong i=0UL; i<cnt; i++ ) idxs[i] = fd_wsample_sample( sampler );
718 0 : }
719 :
720 : void
721 : fd_wsample_sample_and_remove_many( fd_wsample_t * sampler,
722 : ulong * idxs,
723 0 : ulong cnt ) {
724 : /* The compiler doesn't seem to like inlining the call to
725 : fd_wsample_sample_and_remove, which hurts performance by a few
726 : percent because it triggers worse behavior in the CPUs front end.
727 : To address this, we manually inline it here. */
728 0 : for( ulong i=0UL; i<cnt; i++ ) {
729 0 : if( FD_UNLIKELY( !sampler->unremoved_weight ) ) { idxs[ i ] = FD_WSAMPLE_EMPTY; continue; }
730 0 : if( FD_UNLIKELY( sampler->poisoned_mode ) ) { idxs[ i ] = FD_WSAMPLE_INDETERMINATE; continue; }
731 0 : ulong unif = fd_chacha_rng_ulong_roll( sampler->rng, sampler->unremoved_weight+sampler->poisoned_weight );
732 0 : if( FD_UNLIKELY( unif>=sampler->unremoved_weight ) ) {
733 0 : idxs[ i ] = FD_WSAMPLE_INDETERMINATE;
734 0 : sampler->poisoned_mode = 1;
735 0 : continue;
736 0 : }
737 0 : #if FD_WSAMPLE_IMPLEMENTATION > 0
738 0 : if( FD_LIKELY( sampler->height==3UL ) ) {
739 0 : tree_ele_t * tree = sampler->tree;
740 0 : PREPARE();
741 :
742 0 : TRAVERSE_LEVEL(0);
743 0 : TRAVERSE_LEVEL(1);
744 0 : TRAVERSE_LEVEL(2);
745 0 : PROPAGATE_LEVEL(0);
746 0 : PROPAGATE_LEVEL(1);
747 0 : PROPAGATE_LEVEL(2);
748 :
749 0 : FINALIZE();
750 0 : idxs[ i ] = cursor - sampler->internal_node_cnt;
751 0 : continue;
752 0 : }
753 0 : else if( FD_LIKELY( sampler->height==4UL ) ) {
754 0 : tree_ele_t * tree = sampler->tree;
755 0 : PREPARE();
756 :
757 0 : TRAVERSE_LEVEL(0);
758 0 : TRAVERSE_LEVEL(1);
759 0 : TRAVERSE_LEVEL(2);
760 0 : TRAVERSE_LEVEL(3);
761 0 : PROPAGATE_LEVEL(0);
762 0 : PROPAGATE_LEVEL(1);
763 0 : PROPAGATE_LEVEL(2);
764 0 : PROPAGATE_LEVEL(3);
765 :
766 0 : FINALIZE();
767 0 : idxs[ i ] = cursor - sampler->internal_node_cnt;
768 0 : continue;
769 0 : }
770 0 : #endif
771 0 : idxw_pair_t p = fd_wsample_map_sample_i( sampler, unif );
772 0 : fd_wsample_remove( sampler, p );
773 0 : idxs[ i ] = p.idx;
774 0 : }
775 0 : }
776 :
777 :
778 :
779 : ulong
780 20091338 : fd_wsample_sample( fd_wsample_t * sampler ) {
781 20091338 : if( FD_UNLIKELY( !sampler->unremoved_weight ) ) return FD_WSAMPLE_EMPTY;
782 20091338 : if( FD_UNLIKELY( sampler->poisoned_mode ) ) return FD_WSAMPLE_INDETERMINATE;
783 20091338 : ulong unif = fd_chacha_rng_ulong_roll( sampler->rng, sampler->unremoved_weight+sampler->poisoned_weight );
784 20091338 : if( FD_UNLIKELY( unif>=sampler->unremoved_weight ) ) return FD_WSAMPLE_INDETERMINATE;
785 20091338 : return (ulong)fd_wsample_map_sample( sampler, unif );
786 20091338 : }
787 :
788 : ulong
789 0 : fd_wsample_sample_and_remove( fd_wsample_t * sampler ) {
790 0 : if( FD_UNLIKELY( !sampler->unremoved_weight ) ) return FD_WSAMPLE_EMPTY;
791 0 : if( FD_UNLIKELY( sampler->poisoned_mode ) ) return FD_WSAMPLE_INDETERMINATE;
792 0 : ulong unif = fd_chacha_rng_ulong_roll( sampler->rng, sampler->unremoved_weight+sampler->poisoned_weight );
793 0 : if( FD_UNLIKELY( unif>=sampler->unremoved_weight ) ) {
794 0 : sampler->poisoned_mode = 1;
795 0 : return FD_WSAMPLE_INDETERMINATE;
796 0 : }
797 :
798 0 : #if FD_WSAMPLE_IMPLEMENTATION > 0
799 0 : if( FD_LIKELY( sampler->height==3UL ) ) {
800 0 : tree_ele_t * tree = sampler->tree;
801 0 : PREPARE();
802 :
803 0 : TRAVERSE_LEVEL(0);
804 0 : TRAVERSE_LEVEL(1);
805 0 : TRAVERSE_LEVEL(2);
806 0 : PROPAGATE_LEVEL(0);
807 0 : PROPAGATE_LEVEL(1);
808 0 : PROPAGATE_LEVEL(2);
809 :
810 0 : FINALIZE();
811 0 : return cursor - sampler->internal_node_cnt;
812 0 : }
813 0 : if( FD_LIKELY( sampler->height==4UL ) ) {
814 0 : tree_ele_t * tree = sampler->tree;
815 0 : PREPARE();
816 :
817 0 : TRAVERSE_LEVEL(0);
818 0 : TRAVERSE_LEVEL(1);
819 0 : TRAVERSE_LEVEL(2);
820 0 : TRAVERSE_LEVEL(3);
821 0 : PROPAGATE_LEVEL(0);
822 0 : PROPAGATE_LEVEL(1);
823 0 : PROPAGATE_LEVEL(2);
824 0 : PROPAGATE_LEVEL(3);
825 0 : FINALIZE();
826 0 : return cursor - sampler->internal_node_cnt;
827 0 : }
828 0 : #endif
829 0 : idxw_pair_t p = fd_wsample_map_sample_i( sampler, unif );
830 0 : fd_wsample_remove( sampler, p );
831 0 : return p.idx;
832 0 : }
|
