From e96948cb64340eb9713fd9915bf74d76e7187856 Mon Sep 17 00:00:00 2001 From: dzaima Date: Fri, 5 Jul 2024 21:42:21 +0300 Subject: [PATCH] Singeli zip128 --- src/singeli/README.md | 4 ++-- src/singeli/src/avx2.singeli | 10 ++++------ src/singeli/src/base.singeli | 20 +++++++++++++++----- src/singeli/src/count.singeli | 2 +- src/singeli/src/dyarith.singeli | 8 ++++---- src/singeli/src/neon.singeli | 1 - src/singeli/src/replicate.singeli | 4 ++-- src/singeli/src/scan.singeli | 2 +- src/singeli/src/select.singeli | 2 +- src/singeli/src/sse2.singeli | 14 +++++--------- src/singeli/src/transpose.singeli | 6 +++++- 11 files changed, 40 insertions(+), 33 deletions(-) diff --git a/src/singeli/README.md b/src/singeli/README.md index e6b66cee..fcb00a56 100644 --- a/src/singeli/README.md +++ b/src/singeli/README.md @@ -152,6 +152,7 @@ Some may also support one scalar argument or arguments with different widths. - `zipHi{a:V,b:V} : V` - `select{zip{a, b}, 1}` - `mzipLo{a:V,b:V} : V` - `select{mzip{a, b}, 0}` - `mzipHi{a:V,b:V} : V` - `select{mzip{a, b}, 1}` +- `zip128{,Lo,Hi}`, `mzip128{,Lo,Hi}` - within-128-bit-lane versions of the above ## Mask stuff @@ -211,7 +212,6 @@ For float conversions, the used rounding mode is unspecified. - `packQ` - pack 128-bit lanes (`packs`/`packus`) for 16-bit & 32-bit elements - `packQQ` - `packQ` but also defined for 64-bit elements, assuming the high halves are zeroes - `packs` - 128-bit `packs`/`packus` -- `unpackQ` - unpack in 128-bit lanes (like `mzip`) - `shuf16Hi`, `shuf16Lo` - 16-bit shuffles with immediate - `shufHalves` @@ -242,7 +242,7 @@ For float conversions, the used rounding mode is unspecified. - `mla` - multiply-add - `mls` - multiply-subtract - `ornot` - a|~b -- `packLo`, packHi +- `packLo`, `packHi` - `popc` - popcount of each byte - `rbit` - reverse bits in bytes - `rev` - reverse elements in lanes diff --git a/src/singeli/src/avx2.singeli b/src/singeli/src/avx2.singeli index d983fefc..9484034a 100644 --- a/src/singeli/src/avx2.singeli +++ b/src/singeli/src/avx2.singeli @@ -1,9 +1,7 @@ -# questionable pack -def unpackQ{a:T,b:T==[32]i8 } = { tup{emit{[16]i16, '_mm256_unpacklo_epi8', a, b}, emit{[16]i16, '_mm256_unpackhi_epi8', a, b}}} -def unpackQ{a:T,b:T==[16]i16} = { tup{emit{[ 8]i32, '_mm256_unpacklo_epi16', a, b}, emit{[ 8]i32, '_mm256_unpackhi_epi16', a, b}}} -def unpackQ{a:T,b:T==[ 8]i32} = { tup{emit{[ 4]i64, '_mm256_unpacklo_epi32', a, b}, emit{[ 4]i64, '_mm256_unpackhi_epi32', a, b}}} -def unpackQ{a:T,b:T==[ 4]i64} = { tup{emit{[ 4]i64, '_mm256_unpacklo_epi64', a, b}, emit{[ 4]i64, '_mm256_unpackhi_epi64', a, b}}} -# inverse of questionable pack; these saturate the argument +def zip128Lo{a:T, b:T if w256{T}} = emit{T, merge{'_mm256_unpacklo',x86_postfix{eltype{T}}}, a, b} +def zip128Hi{a:T, b:T if w256{T}} = emit{T, merge{'_mm256_unpackhi',x86_postfix{eltype{T}}}, a, b} + +# questionable pack; these saturate the argument def packQ{a:T,b:T==[16]i16} = emit{[32]i8, '_mm256_packs_epi16', a, b} def packQ{a:T,b:T==[ 8]i32} = emit{[16]i16, '_mm256_packs_epi32', a, b} def packQ{a:T,b:T==[16]u16} = emit{[32]u8, '_mm256_packus_epi16', a, b} diff --git a/src/singeli/src/base.singeli b/src/singeli/src/base.singeli index b9b6ec8c..320cc1af 100644 --- a/src/singeli/src/base.singeli +++ b/src/singeli/src/base.singeli @@ -162,8 +162,8 @@ def { 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, - topBlend,topMask,topMaskStore,topMaskStoreF,unord,unpackQ,vfold,widen, - zip,zipHi,zipLo + topBlend,topMask,topMaskStore,topMaskStoreF,unord,vfold,widen, + zipHi,zipLo } def homMaskX{a:T} = tup{1, homMask{a}} # tup{n,mask}; mask with each bit repeated n times @@ -191,9 +191,19 @@ if_inline (hasarch{'X86_64'}) { def {__adds,__subs,__sqrt,vec_broadcast,vec_make} } -def mzip {a:T, b:T} = each{{v}=>el_m{T}~~v, zip{a, b}} -def mzipLo{a:T, b:T} = el_m{T} ~~ zipLo{a, b} -def mzipHi{a:T, b:T} = el_m{T} ~~ zipHi{a, b} +def zip128Lo{a:T, b:T if width{T}==128} = zipLo{a, b} +def zip128Hi{a:T, b:T if width{T}==128} = zipHi{a, b} +def zip {a:T, b:T} = tup{zipLo {a,b}, zipHi {a,b}} +def zip128{a:T, b:T} = tup{zip128Lo{a,b}, zip128Hi{a,b}} + +local def make_mzip{l, h} = tup{ + {a:T, b:T} => el_m{T} ~~ l{a, b}, + {a:T, b:T} => el_m{T} ~~ h{a, b}, + {a:T, b:T} => each{{z} => el_m{T}~~z{a, b}, tup{l, h}} +} +def {mzipLo, mzipHi, mzip } = make_mzip{zipLo, zipHi} +def {mzip128Lo, mzip128Hi, mzip128} = make_mzip{zip128Lo, zip128Hi} + def packQ{{a, b}} = packQ{a, b} def pair{{a, b}} = pair{a, b} diff --git a/src/singeli/src/count.singeli b/src/singeli/src/count.singeli index f333ef3a..e82db0e0 100644 --- a/src/singeli/src/count.singeli +++ b/src/singeli/src/count.singeli @@ -3,7 +3,7 @@ include 'util/tup' include './vecfold' if_inline (hasarch{'SSE2'}) { - fn sum_vec{T}(v:T) = vfold{+, fold{+, unpackQ{v, T**0}}} + fn sum_vec{T}(v:T) = vfold{+, fold{+, mzip128{v, T**0}}} def fold_addw{v:T=[_](i8)} = sum_vec{T}(v) } diff --git a/src/singeli/src/dyarith.singeli b/src/singeli/src/dyarith.singeli index 6057e615..bbd56fd7 100644 --- a/src/singeli/src/dyarith.singeli +++ b/src/singeli/src/dyarith.singeli @@ -34,8 +34,8 @@ def arithChk2{F, M, w:T, x:T if is_s{T} and (same{F,__add} or same{F,__sub})} = # × def arithChk2{F=(__mul), M, w:T=[_](i8), x:T if hasarch{'X86_64'}} = { - def wp = unpackQ{w, T ~~ (T**0 > w)} - def xp = unpackQ{x, T ~~ (T**0 > x)} + def wp = mzip128{w, T ~~ (T**0 > w)} + def xp = mzip128{x, T ~~ (T**0 > x)} def rp = each{__mul, wp, xp} if (M{0}) { # masked check def bad = each{{v} => ty_s{(v<<8)>>8 != v}, rp} @@ -58,8 +58,8 @@ def arithChk2{F=(__mul), M, w:T=[_](i32), x:T if hasarch{'X86_64'}} = { tup{w*x, tup{'homAny', M{abs{f32mul} >= max}}} # TODO fallback to the below if the above fails # TODO don't do this, but instead shuffle one half, do math, unshuffle that half - # def wp = unpackQ{w, T**0} - # def xp = unpackQ{x, T**0} + # def wp = mzip128{w, T**0} + # def xp = mzip128{x, T**0} # def rp = each{mul32, wp, xp} # def T2 = re_el{i64, T} # def bad = each{{v} => { diff --git a/src/singeli/src/neon.singeli b/src/singeli/src/neon.singeli index 63f99a5b..eb286679 100644 --- a/src/singeli/src/neon.singeli +++ b/src/singeli/src/neon.singeli @@ -72,7 +72,6 @@ def vshl{a:T, b:T, n if knum{n}} = emit{T, ntyp{'vext', T}, a, b, n} def zipLo{a:T, b:T if nvec{T}} = emit{T, ntyp{'vzip1', T}, a, b} def zipHi{a:T, b:T if nvec{T}} = emit{T, ntyp{'vzip2', T}, a, b} -def zip{a:T, b:T if nvec{T}} = tup{zipLo{a,b}, zipHi{a,b}} def packLo{x:T, y:T if nvec{T}} = { def H=el_s{T}; emit{H, ntyp{'vuzp1', H}, H~~x, H~~y} } def packHi{x:T, y:T if nvec{T}} = { def H=el_s{T}; emit{H, ntyp{'vuzp2', H}, H~~x, H~~y} } diff --git a/src/singeli/src/replicate.singeli b/src/singeli/src/replicate.singeli index 30046b1d..7f279186 100644 --- a/src/singeli/src/replicate.singeli +++ b/src/singeli/src/replicate.singeli @@ -104,7 +104,7 @@ rcsh4_lkup:*i8 = shiftright{0, scan{+, fold{|, table{==, rcsh4_dom, iota{64}}}}} def read_shuf_vecs{l, ellw:(u64), shp:*V} = { # tuple of byte selectors in 1<gen{V~~q}, unpackQ{s, s}} + each{{q}=>gen{V~~q}, mzip128{s, s}} } def get_rep_iter{V==[4]u64, wv} = { def step = 4 diff --git a/src/singeli/src/scan.singeli b/src/singeli/src/scan.singeli index 3d15a81b..889ab06d 100644 --- a/src/singeli/src/scan.singeli +++ b/src/singeli/src/scan.singeli @@ -149,7 +149,7 @@ fn bcs{T if hasarch{'AVX2'}}(x:*u64, r:*T, l:u64) : void = { def ii32 = iota{32}; def bit{k}=bit{k,ii32}; def tail{k}=tail{k,ii32} def sums{n} = (if (n==0) tup{0} else { def s=sums{n-1}; merge{s,s+1} }) - def widen{v:T} = unpackQ{shuf{[4]u64, v, 4b3120}, T**0} + def widen{v:T} = mzip128{shuf{[4]u64, v, 4b3120}, T**0} def sumlanes{x:(u32)} = { b:= [8]u32**x >> make{[8]u32, 4*tail{1, iota{8}}} diff --git a/src/singeli/src/select.singeli b/src/singeli/src/select.singeli index 9d5799e3..8df6fede 100644 --- a/src/singeli/src/select.singeli +++ b/src/singeli/src/select.singeli @@ -42,7 +42,7 @@ def shuf_select{ri, rd, TI, w, r, wl, xl, selx} = { c2:= shuf{[4]u64, c+c, 4b3120} each{ {c,o} => se{e*2, VI~~c, o}, - unpackQ{c2, c2 + VI**1}, + mzip128{c2, c2 + VI**1}, 2*o + iota{2} } } diff --git a/src/singeli/src/sse2.singeli b/src/singeli/src/sse2.singeli index a4f8c371..83b74f5f 100644 --- a/src/singeli/src/sse2.singeli +++ b/src/singeli/src/sse2.singeli @@ -68,15 +68,11 @@ def packs{a:T,b:T==[4]i32} = emit{[ 8]i16, '_mm_packs_epi32', a, b} def packs{a:T,b:T==[8]u16} = emit{[16]u8, '_mm_packus_epi16', a, b} def packQ{a:T,b:T if w128i{T}} = packs{a,b} - -def zipLo{a:T, b:T if w128i{T}} = emit{T, merge{'_mm_unpacklo_epi',fmtnat{elwidth{T}}}, a, b} -def zipHi{a:T, b:T if w128i{T}} = emit{T, merge{'_mm_unpackhi_epi',fmtnat{elwidth{T}}}, a, b} -def zipLo{a:T, b:T if w128f{T}} = emit{T, merge{'_mm_unpacklo_p',if (elwidth{T}==32) 's' else 'd'}, a, b} -def zipHi{a:T, b:T if w128f{T}} = emit{T, merge{'_mm_unpackhi_p',if (elwidth{T}==32) 's' else 'd'}, a, b} -def zip{a:T, b:T if w128i{T}} = tup{zipLo{a,b}, zipHi{a,b}} - - -def unpackQ{a:T, b:T if w128{T}} = mzip{a, b} +def x86_postfix{T if isint{T}} = merge{'_epi',fmtnat{width{T}}} +def x86_postfix{(f32)} = '_ps' +def x86_postfix{(f64)} = '_pd' +def zipLo{a:T, b:T if w128{T}} = emit{T, merge{'_mm_unpacklo',x86_postfix{eltype{T}}}, a, b} +def zipHi{a:T, b:T if w128{T}} = emit{T, merge{'_mm_unpackhi',x86_postfix{eltype{T}}}, a, b} def shuf{L, x:T, n if w128{T} and lvec{L,4,32} and knum{n}} = T ~~ emit{[4]i32, '_mm_shuffle_epi32', v2i{x}, n} def shuf16Lo{x:T, n} = T~~emit{[8]i16, '_mm_shufflelo_epi16', x, n} diff --git a/src/singeli/src/transpose.singeli b/src/singeli/src/transpose.singeli index 244bd49a..6ad7e97e 100644 --- a/src/singeli/src/transpose.singeli +++ b/src/singeli/src/transpose.singeli @@ -11,7 +11,11 @@ def pairs{o, x} = { def g = 2*i - i%o select{x, tup{g, g+o}} } -def unpack_pass{o, x} = merge{...each{unpackQ, ...pairs{o, x}}} + +def unpack_typed{a, b} = mzip128{a, b} +def unpack_typed{a:T, b:T if elwidth{T}==64} = zip128{a, b} +def unpack_pass{o, x} = merge{...each{unpack_typed, ...pairs{o, x}}} + def permute_pass{o, x} = { def p = pairs{o, x} def h{s} = each{{a,b}=>emit{[8]i32, '_mm256_permute2f128_si256', a,b,s}, ...p}