Unroll vector binary searches for different search depths

This commit is contained in:
Marshall Lochbaum 2023-07-05 09:57:30 -04:00
parent 38671e3fe9
commit e0621f99f8

View File

@ -26,7 +26,7 @@ def for_vec_overlap{vl}{vars,begin==0,n,iter} = {
setlabel{end} setlabel{end}
} }
def ceil_log2{n:u64} = 64 - clz{n+1} def ceil_log2{n:u64} = 64 - clz{n-1}
# Shift as u16, since x86 is missing 8-bit shifts # Shift as u16, since x86 is missing 8-bit shifts
def shr16{v, n} = type{v}~~(([width{type{v}}/16]u16~~v) >> n) def shr16{v, n} = type{v}~~(([width{type{v}}/16]u16~~v) >> n)
@ -44,23 +44,25 @@ def bin_search_vec{up, w:*i8, wn, x:*i8, n, res:*i8} = {
{i} => sel{H, v, i} {i} => sel{H, v, i}
} }
if (hasarch{'AVX2'} and wn < 16) { if (hasarch{'AVX2'} and wn < 16) {
log := ceil_log2{wn-1} log := ceil_log2{wn+1}
l := 1<<log l := 1<<log
gap := l - cast_i{u8, wn} gap := l - cast_i{u8, wn}
off := U**(gap - 1) off := U**(gap - 1)
wv := homBlend{load{*H~~(w-gap), 0}, H**pre, maskOf{H,gap}} wv := homBlend{load{*H~~(w-gap), 0}, H**pre, maskOf{H,gap}}
def selw = getsel{wv} def selw = getsel{wv}
h0 := U**(l/2) h0 := U**(l/2)
@for_vec_overlap{vl} (j to n) { @unroll (klog from 2 to 5) {
if (log==klog) @for_vec_overlap{vl} (j to n) {
xv:= load{*V~~(x+j), 0} xv:= load{*V~~(x+j), 0}
s := U**0 s := U**0
h := h0 h := h0
@for (promote{u64,log}) { @unroll (klog) {
s |= h &~ lt{xv, selw{s | h}} s |= h &~ lt{xv, selw{s | h}}
h = shr16{h, 1} h = shr16{h, 1}
} }
store{*U~~(res+j), 0, s - off} store{*U~~(res+j), 0, s - off}
} }
}
} else { } else {
assert{wn < 128} # Total must fit in i8 assert{wn < 128} # Total must fit in i8
t0:*i8 = copy{256,0} t0:*i8 = copy{256,0}