diff --git a/src/singeli/src/scan.singeli b/src/singeli/src/scan.singeli index 9b384a59..d7fead2f 100644 --- a/src/singeli/src/scan.singeli +++ b/src/singeli/src/scan.singeli @@ -350,6 +350,9 @@ fn scan_rows_andor{id}(src:*u64, dst:*u64, nl:usz, l:usz) : void = { if (qand) x &~ ((x+c) & (x+m)) else x | ((-x-c) &~ (x-m)) } + def apply_carry{s, c, f} = { + if (qand) s & (~f | c) else s | (f & c) + } if (l < 64) { if ((l & (l-1)) == 0) { if (l == 2) { @@ -363,9 +366,6 @@ fn scan_rows_andor{id}(src:*u64, dst:*u64, nl:usz, l:usz) : void = { # could use for l>=8; not much faster and takes up space # def rowwise{T} = @for (r in *T~~dst, x in *T~~src over (64/width{T})*nw) r = x &~ (x+1) } else if (hasarch{'AVX2'}) { - def apply_carry{s, c, f} = { - if (qand) s & (~f | c) else s | (f & c) - } avx2_loop_with_unaligned_mask{src, dst, nw, l, scan_mask, apply_carry} } else { loop_with_unaligned_mask{src, dst, nw, l, {x, c, m} => { @@ -376,14 +376,40 @@ fn scan_rows_andor{id}(src:*u64, dst:*u64, nl:usz, l:usz) : void = { }} } } else if (l < 160) { - q:usz = 0 # distance to next row boundary - c:u64 = id # carry - @for (r in dst, x in src over nw) { - b:= q<64 # whether there's a boundary - p:= q%64 # its position - q-= 64 - (l &- b) - r = res_m1{x, c, promote{u64, b} << p} - c = r >> 63 + if (hasarch{'AVX2'}) { + assert{l >= 64} + def V = [4]u64 + c := V**id # carry, 0 or 1 + v1:= V**1; xk:= c - v1 + q := -make{V, 64*iota{4}} # distance to next row boundary + def q_mod{} = { q = blend_top{q,q+V**l, q} } + o:u64 = 256; while (o>l) { o-=l; q_mod{} } + @for_masked{4} (x in tup{V, src}, r in tup{V, dst} over nw) { + # Get mask; <=1 bit per word + m:= v1 << q + q-= V**o; q_mod{} + # Within-word scan and carry info + r = (if (qand) x &~ ((x+v1) & (x+m)) + else x | ((-x) &~ (x-m))) + p:= (if (qand) x&~m else x|m) == xk + k:= r>>63 | p # Carry of 0 or 1, but -1 to propagate previous + # Propagate carries and adjust result + def bl{b,a} = blend_top{b,a, b} + k = bl{make_scan_idem{f64, bl}{k}, c} # Can't be -1 now + t:= blend{shuf{V, k, 3,0,1,2}, c, 1,0,0,0} + r = apply_carry{r, -t, m-v1} + c = shuf{V, k, 4**3} + } + } else { + q:usz = 0 # distance to next row boundary + c:u64 = id # carry + @for (r in dst, x in src over nw) { + b:= q<64 # whether there's a boundary + p:= q%64 # its position + q-= 64 - (l &- b) + r = res_m1{x, c, promote{u64, b} << p} + c = r >> 63 + } } } else { i :usz = 0 # row bit index