From 40bf3bfd1ceca442d1e7f363ba340802a8bbd690 Mon Sep 17 00:00:00 2001 From: Marshall Lochbaum Date: Tue, 4 Jun 2024 20:35:46 -0400 Subject: [PATCH] AVX2-based boolean short-row scans --- src/singeli/src/base.singeli | 2 +- src/singeli/src/scan.singeli | 52 +++++++++++++++++++++++++++++++++--- 2 files changed, 50 insertions(+), 4 deletions(-) diff --git a/src/singeli/src/base.singeli b/src/singeli/src/base.singeli index 7c50724c..b9b6ec8c 100644 --- a/src/singeli/src/base.singeli +++ b/src/singeli/src/base.singeli @@ -158,7 +158,7 @@ def lvec = match { {[n]T, n, (width{T})} => 1; {T, n, w} => 0 } # base cases def { - absu,andAllZero,andnz,b_getBatch,clmul,cvt,extract,fold_addw,half, + absu,andAllZero,andnz,b_getBatch,blend,clmul,cvt,extract,fold_addw,half, homAll,homAny,homBlend,homMask,homMaskStore,homMaskStoreF,loadBatchBit, loadLow,make,maskStore,maskToHom,mulw,mulh,narrow,narrowPair,packHi,packLo,packQ,pair,pdep, pext,popcRand,sel,shl,shr,shuf,shuf16Hi,shuf16Lo,shufHalves,storeLow, diff --git a/src/singeli/src/scan.singeli b/src/singeli/src/scan.singeli index 02d3c90f..4b65c5ae 100644 --- a/src/singeli/src/scan.singeli +++ b/src/singeli/src/scan.singeli @@ -5,6 +5,7 @@ if_inline (hasarch{'X86_64'}) { } include './mask' include './f64' +include 'util/tup' include './scan_common' # Initialized scan, generic implementation @@ -300,10 +301,13 @@ export{'si_scan_plus_i32_f64', plus_scanG{i32, f64}} # Row-wise boolean scan def aligned_mask{l} = (~u64~~0) / ((u64~~1 << l)-1) +def unaligned_mask{l} = { + def d = 64 % l + def m = (~u64~~0 >> d) / ((u64~~1 << l)-1) + tup{m<> d) / ((u64~~1 << l)-1) - m = m<>d | m<<(l-d) } } +def avx2_loop_with_unaligned_mask{xp, rp, nw, l, scan_words, apply_carry} = { + {ms, d} := unaligned_mask{l} + def V = [4]u64 + d4:usz = width{V} % l + m:= make{V, scan{{a,_} => a>>d | a<<(l-d), tup{ms, ...iota{3}}}} + c:= V**0 + @maskedLoop{4} (x in tup{V, xp}, + r in tup{V, rp} over promote{u64,nw}) { + s := scan_words{x, m} + pc:= c; c = shuf{V, -(s>>63), 4b2103} + r = apply_carry{s, blend{V, c, pc, 2b0001}, (m-V**1)&~m} + m = m>>d4 | m<<(l-d4) + } +} + fn scan_rows_andor{id}(src:*u64, dst:*u64, n:usz, l:usz) : void = { def qand = not id assert{l > 0} @@ -337,6 +356,17 @@ fn scan_rows_andor{id}(src:*u64, dst:*u64, n: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 scan_words{x, m:V} = { + mb:= m | V**1 + p:= if (qand) (x &~ m) >> 1 else ~(x | m) >> 1 + a:= if (qand) p + (mb & x) else p + (mb &~ x) + if (qand) p ^ a else ~(p ^ a) + } + 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_words, apply_carry} } else { loop_with_unaligned_mask{src, dst, nw, l, {x, c, m} => { s:= (if (qand) (x &~ m) >> 1 else ~(x | m) >> 1 ) @@ -398,6 +428,18 @@ fn scan_rows_neq(x:*u64, r:*u64, n:usz, l:usz) : void = { b:= s<<1 & m # last bit of previous row r = s ^ (b< pre{v, 1} + } + s:= vec_prefix_byshift{^, <<}{x} + b:= s<<1 & m # last bit of previous row + s ^ (b< { s:= scan_word{x} @@ -432,6 +474,10 @@ fn scan_rows_left(x:*u64, r:*u64, n:usz, l:usz) : void = { if ((l & (l-1)) == 0) { m:u64 = aligned_mask{l} @for (r, x over nw) { b:= x & m; r = b< { f:= (m-1)&~m # bits before first full row