From c64cdb2c61fed3ae2aefa2408db0e7980428d320 Mon Sep 17 00:00:00 2001 From: dzaima Date: Mon, 3 Feb 2025 04:18:23 +0200 Subject: [PATCH] make store_(blended|masked)_(hom|top) accept element pointer --- src/singeli/README.md | 4 ++-- src/singeli/src/avx2.singeli | 12 ++++++------ src/singeli/src/base.singeli | 10 +++++++++- src/singeli/src/mask.singeli | 2 +- src/singeli/src/neon.singeli | 5 ++--- src/singeli/src/replicate.singeli | 2 +- src/singeli/src/scan.singeli | 6 +++--- src/singeli/src/slash.singeli | 4 ++-- src/singeli/src/sse2.singeli | 2 +- 9 files changed, 27 insertions(+), 20 deletions(-) diff --git a/src/singeli/README.md b/src/singeli/README.md index 93b66413..a13069e2 100644 --- a/src/singeli/README.md +++ b/src/singeli/README.md @@ -186,8 +186,8 @@ For unaligned scalar loads & stores, `loadu` & `storeu` should be used. - `store{p:*E, a:[k]E} : void` - store full vector - `load{[k]E, p:*E, vl} : [k]E` - load first `vl` elements of vector (memory of upper ones won't be touched, and their values are unspecified) - `store{p:*E, a:[k]E, vl} : void` - store first `vl` elements (upper ones won't be touched) -- `store_masked_hom{p:*V, m:mt{V}, a:V}` - conditionally store elements based on mask; won't touch masked-off elements -- `store_masked_top{p:*V, m:V, a:V}` - conditionally store elements based on top bit of `m`; won't touch masked-off elements +- `store_masked_hom{p:*E, m:mt{V}, a:V=[_]E}` - conditionally store elements based on mask; won't touch masked-off elements; `p` may also be `p:*V` +- `store_masked_top{p:*E, m:V, a:V=[_]E}` - conditionally store elements based on top bit of `m`; won't touch masked-off elements; `p` may also be `p:*V` - `store_blended_hom` - `store_masked_hom` but may touch masked-off elements and thus be supported on more types - `store_blended_top` - `store_masked_top` but may touch masked-off elements and thus be supported on more types diff --git a/src/singeli/src/avx2.singeli b/src/singeli/src/avx2.singeli index 2b7502c0..f79213ac 100644 --- a/src/singeli/src/avx2.singeli +++ b/src/singeli/src/avx2.singeli @@ -24,13 +24,13 @@ def sel{L, x:T, i:I if w256{T} and lvec{L,8,32} and w256{I,32}} = vec_shuffle{L, def sel{L, x:T, i:I if w256{T} and lvec{L,16,8} and w256{I, 8}} = vec_shuffle{L, x, i} # masked store; F variants may not be a single instruction -def store_masked_top{p:*T, m:M, v:T if w256i{M, 32} and w256{T,elwidth{M}}} = emit{void, '_mm256_maskstore_epi32', *i32~~p, m, [8]i32~~v} -def store_masked_top{p:*T, m:M, v:T if w256i{M, 64} and w256{T,elwidth{M}}} = emit{void, '_mm256_maskstore_pd', *f64~~p, m, [4]f64~~v} -def store_masked_hom{p:*T, m:M, v:T if w256i{M} and w256{T,elwidth{M}}} = store_masked_top{p, m, v} +def store_masked_top{p:*E, m:M=[k]_, v:V=[k]E if w256{V, 32} and w256i{M}} = emit{void, '_mm256_maskstore_epi32', *i32~~p, m, [8]i32~~v} +def store_masked_top{p:*E, m:M=[k]_, v:V=[k]E if w256{V, 64} and w256i{M}} = emit{void, '_mm256_maskstore_pd', *f64~~p, m, [4]f64~~v} +def store_masked_hom{p:*E, m:M=[k]_, v:V=[k]E if w256{V } and w256i{M}} = store_masked_top{p,m,v} -def store_blended_top{p:*T, m:M, v:T if w256i{M} and elwidth{T}>=32} = store_masked_top{p,m,v} -def store_blended_hom{p:*T, m:M, v:T if w256i{M} and elwidth{T}>=32} = store_masked_top{p,m,v} -def store_blended_hom{p:*T, m:M, v:T if w256i{M} and elwidth{T}<=16 and w256{T,elwidth{M}}} = store{p, 0, blend_hom{load{p}, v, m}} +def store_blended_top{p:*E, m:M=[k]_, v:V=[k]E if w256i{M} and width{E}>=32} = store_masked_top{p,m,v} +def store_blended_hom{p:*E, m:M=[k]_, v:V=[k]E if w256i{M} and width{E}>=32} = store_masked_top{p,m,v} +def store_blended_hom{p:*E, m:M=[k]_, v:V=[k]E if w256i{M} and width{E}<=16 and w256{V,elwidth{M}}} = store{p, blend_hom{load{V, p}, v, m}} # mask stuff def top_to_int{x:T if w256{T, 8}} = emit{u32, '_mm256_movemask_epi8', x} diff --git a/src/singeli/src/base.singeli b/src/singeli/src/base.singeli index 24d23b4b..fa421d77 100644 --- a/src/singeli/src/base.singeli +++ b/src/singeli/src/base.singeli @@ -202,6 +202,14 @@ def load{V=[k]E, ptr:*E } = load{*V~~ptr, 0} def load{V=[k]E, ptr:*E, k} = load{*V~~ptr, 0} def store{ptr:*E, val:V=[k]E } = store{*V~~ptr, 0, val} def store{ptr:*E, val:V=[k]E, k} = store{*V~~ptr, 0, val} +local def extend accept_vec_ptr{op} = { + def op{ptr:*[k]E, m:M, val:[k]E} = op{*E~~ptr, m, val} +} +def extend_each{G}{...fs} = each{G, fs} +extend (extend_each{accept_vec_ptr}){ + store_masked_hom,store_blended_hom, + store_masked_top,store_blended_top, +} @@ -234,7 +242,7 @@ def pack {a:T, b:T, k} = unzip{el_s{T}~~a, el_s{T}~~b, k} local def extend kpair{op} = { def op{a:T, b:T} = tup{op{a,b,0}, op{a,b,1}} } -extend ({...f}=>each{kpair,f}){pack, zip, zip128, mzip, mzip128, unzip, unzip128} +extend (extend_each{kpair}){pack, zip, zip128, mzip, mzip128, unzip, unzip128} def packQ{{a, b}} = packQ{a, b} def pair{{a, b}} = pair{a, b} diff --git a/src/singeli/src/mask.singeli b/src/singeli/src/mask.singeli index c6f2a5f9..056441e8 100644 --- a/src/singeli/src/mask.singeli +++ b/src/singeli/src/mask.singeli @@ -49,7 +49,7 @@ def store_narrow{ptr:*E0, i, x:[k]E1, M} = { def TF = re_el{E0, [k]E1} xu:= narrow{E0, x} - if (M{0}) store_blended_hom{*TF~~rpos, M{TF, 'to homogeneous bits'}, undef_promote{TF, xu}} + if (M{0}) store_blended_hom{rpos, M{TF, 'to homogeneous bits'}, undef_promote{TF, xu}} else store{rpos, xu, k} } diff --git a/src/singeli/src/neon.singeli b/src/singeli/src/neon.singeli index 0ca386bf..dd425e73 100644 --- a/src/singeli/src/neon.singeli +++ b/src/singeli/src/neon.singeli @@ -77,6 +77,8 @@ def load{V=[_]E, ptr:*E, vl if nvec{V} and vl*width{E}<=64} = { # implemented vi V ~~ emit{L, ntyp{'vld1', '_dup', L}, *U~~ptr} } +def store_blended_hom{p:*E, m:M=[k]_, v:V=[k]E if nveci{M,width{E}} and nvec{V}} = store{p, blend_hom{load{V, p}, v, m}} + def undef_promote{T=[_]E, x:X=[_]E if w64{X} and w128{T}} = emit{T, ntyp{'vcombine', X}, x, x} # arm_neon.h doesn't actually provide a way to do this in a 0-instruction way. ¯\_(ツ)_/¯ @@ -181,6 +183,3 @@ def and_bit_none{x:T, y:T if nveci{T}} = ~any_bit{x&y} def hom_to_int_ext{a:V=[_]E if w128u{V}} = tup{width{E}/2, extract{[1]u64~~narrow_trunc{w_h{E}, a}, 0}} # E==u8 fully handled by the below cases def hom_to_int_ext{a:V=[_]E if w64u{V}} = tup{width{E}, extract{[1]u64~~a, 0}} def hom_to_int_ext{a:([16]u8)} = tup{4, extract{[1]u64~~shrn{re_el{u16,a}, 4}, 0}} - - -def store_blended_hom{p:*T, m:M, v:T if nveci{M} and nvec{T,elwidth{M}}} = store{p, 0, blend_hom{load{p}, v, m}} diff --git a/src/singeli/src/replicate.singeli b/src/singeli/src/replicate.singeli index 1ce850eb..b4ef1ba5 100644 --- a/src/singeli/src/replicate.singeli +++ b/src/singeli/src/replicate.singeli @@ -217,7 +217,7 @@ fn rep_const_shuffle_partial4(wv:u64, ellw:u64, x:*i8, r:*i8, n:u64) : void = { setlabel{end} q := (re+step) - r - if (q!=0) store_blended_hom{*V~~r, mask_of_first{V, q}, s} + if (q!=0) store_blended_hom{r, mask_of_first{V, q}, s} } diff --git a/src/singeli/src/scan.singeli b/src/singeli/src/scan.singeli index e3e970e2..a0d21407 100644 --- a/src/singeli/src/scan.singeli +++ b/src/singeli/src/scan.singeli @@ -16,8 +16,8 @@ def scan_loop{init, x:*T, r:*T, len:(u64), scan, scan_last} = { def step = arch_defvw/width{T} def V = [step]T p:= V**init - xv:= *V ~~ x - rv:= *V ~~ r + xv:= *V~~x + rv:= *V~~r e:= len/step @for (xv, rv over e) rv = scan{xv,p} q:= len & (step-1) @@ -135,7 +135,7 @@ fn scan_neq{if hasarch{'AVX512BW', 'VPCLMULQDQ', 'GFNI'}}(init:u64, x:*u64, r:*u carry = (xc ^ hb) ^ shuf{u64, xh, 1,1} rem:= nw - 8*i if (rem < 8) { - store_masked_hom{*V~~r+i, [8]u1~~(~(u8~~0xff<>0) if (end-buf>vc) { store{r0, load{R, buf}}; r0+=vc; buf+=vc } assert{bufw % width{R} == 0} # to make sure the below doesn't read out-of-bounds on the stack - store_blended_hom{*R~~r0, mask_of_first{R, end-buf}, load{R, buf}} + store_blended_hom{r0, mask_of_first{R, end-buf}, load{R, buf}} } else { @for (r0, buf over u64~~(end-buf)) r0 = buf } @@ -255,7 +255,7 @@ fn slash{c, T if hasarch{if (width{T}>=32) 'AVX512F' else 'AVX512VBMI2'}}(w:*u64 cs := cast_i{I,promote{i64,1}<<(c%64) - 1} if (vl==64) cs -= cast_i{I,c}>>6 v := emitT{V, 'mask_compress', x, m, x} - store_masked_hom{*V~~r, [vl]u1~~cs, v} + store_masked_hom{r, [vl]u1~~cs, v} r += c } } diff --git a/src/singeli/src/sse2.singeli b/src/singeli/src/sse2.singeli index 5adcc9bb..6d6640c3 100644 --- a/src/singeli/src/sse2.singeli +++ b/src/singeli/src/sse2.singeli @@ -65,7 +65,7 @@ def packQ{a:T,b:T if w128i{T}} = packs{a,b} def x86_shufps_range{is, hi} = in_range{is,0,2, 0,4} and in_range{is,2,4, hi,hi+4} def shuf_ind{a:T, b:T=[4]E, {...is} if width{E}==32 and length{is}==4 and x86_shufps_range{is, 4}} = vec_shuffle{[4]f32, tup{a, b}, is&3} -def store_blended_hom{p:*T, m:M, v:T if w128i{M} and w128{T,elwidth{M}}} = store{p, 0, blend_hom{load{p}, v, m}} +def store_blended_hom{p:*E, m:M=[k]_, v:V=[k]E if w128i{M,width{E}} and w128{V}} = store{p, blend_hom{load{V, p}, v, m}} def widen{T, x:X if w128i{T} and w128i{X} and w128s{T}==w128s{X} and elwidth{T}>elwidth{X}} = { def s{v} = s{mzip{v,v,0}}