NEON constrep

This commit is contained in:
dzaima 2023-07-10 16:55:29 +03:00
parent eac297d4d9
commit cd7807fba0
5 changed files with 115 additions and 66 deletions

View File

@ -637,7 +637,7 @@ cachedBin‿linkerCache ← {
"xag""src/builtins/sort.c""bins" "xag""src/builtins/sort.c""bins"
"2..""src/builtins/select.c""select", "2..""src/builtins/scan.c""scan", "2..""src/builtins/select.c""select", "2..""src/builtins/scan.c""scan",
"2..""src/builtins/slash.c""constrep", "2..""src/builtins/scan.c""neq", "2a.""src/builtins/slash.c""constrep", "2..""src/builtins/scan.c""neq",
"2..""src/builtins/slash.c""slash", "2..""src/builtins/slash.c""count" "2..""src/builtins/slash.c""slash", "2..""src/builtins/slash.c""count"
objs objs

View File

@ -81,10 +81,12 @@
#endif #endif
#endif #endif
#if SINGELI_AVX2 #if SINGELI_AVX2 || SINGELI_NEON
#define SINGELI_FILE constrep #define SINGELI_FILE constrep
#include "../utils/includeSingeli.h" #include "../utils/includeSingeli.h"
#endif
#if SINGELI_AVX2
#define SINGELI_FILE count #define SINGELI_FILE count
#include "../utils/includeSingeli.h" #include "../utils/includeSingeli.h"
@ -828,7 +830,7 @@ B slash_c2(B t, B w, B x) {
u8 xk = xl-3; u8 xk = xl-3;
void* rv = m_tyarrv(&r, 1<<xk, s, xt); void* rv = m_tyarrv(&r, 1<<xk, s, xt);
void* xv = tyany_ptr(x); void* xv = tyany_ptr(x);
#if SINGELI_AVX2 #if SINGELI_AVX2 || SINGELI_NEON
simd_constrep[xk](wv, xv, rv, xlen); simd_constrep[xk](wv, xv, rv, xlen);
#else #else
#define CASE(L,T) case L: { REP_BY_SCAN(T, wv) break; } #define CASE(L,T) case L: { REP_BY_SCAN(T, wv) break; }

View File

@ -203,6 +203,7 @@ def b_getBatch{...x} = assert{'b_getBatch not supported', show{...x}}
def extract{...x} = assert{'extract not supported', show{...x}} def extract{...x} = assert{'extract not supported', show{...x}}
def abs{...x} = assert{'abs not supported', show{...x}} def abs{...x} = assert{'abs not supported', show{...x}}
def homBlend{...x} = assert{'homBlend not supported', show{...x}} def homBlend{...x} = assert{'homBlend not supported', show{...x}}
def zip{...x} = assert{'zip not supported', show{...x}}
def andnot{a, b & anyNum{a} & anyNum{b}} = a & ~b def andnot{a, b & anyNum{a} & anyNum{b}} = a & ~b
oper &~ andnot infix none 35 oper &~ andnot infix none 35

View File

@ -1,7 +1,12 @@
include './debug'
include './base' include './base'
if (hasarch{'AVX2'}) {
include './sse' include './sse'
include './avx' include './avx'
include './avx2' include './avx2'
} else if (hasarch{'AARCH64'}) {
include './neon'
}
include './mask' include './mask'
include 'util/tup' include 'util/tup'
def incl{a,b} = slice{iota{b+1},a} def incl{a,b} = slice{iota{b+1},a}
@ -17,50 +22,52 @@ def fact_inds = slice{iota{fact_size},8}
def fact_tab = makefact{basic_rep, fact_inds} def fact_tab = makefact{basic_rep, fact_inds}
factors:*u8 = fact_tab factors:*u8 = fact_tab
def sdtype = [arch_defvw/8]i8 # shuf data type
def get_shufs{step, wv, len} = { def get_shufs{step, wv, len} = {
def i = iota{len*step} def i = iota{len*step}
split{step, (i - i%wv)/wv} split{step, (i - i%wv)/wv}
} }
def get_shuf_data{wv, len} = get_shufs{32, wv, len} def get_shuf_data{wv, len} = get_shufs{vcount{sdtype}, wv, len} # [len] byte-selector vectors for wv/sdtype (expanded to wider types by read_shuf_vecs)
def get_shuf_data{wv} = get_shuf_data{wv, wv} def get_shuf_data{wv} = get_shuf_data{wv, wv}
def rep_iter_from_sh{sh} = {
def l = tuplen{sh} # all shuffle vectors for 𝕨∊3…7
def h = l >> 1 def rcsh_vals = slice{basic_rep, 1} # Handle 2 specially
{x, gen} => { rcsh_offs:*u8 = shiftright{0, scan{+,rcsh_vals}}
def fs{v, s} = gen{sel{[16]i8, v, s}} rcsh_data:*i8 = join{join{each{get_shuf_data, rcsh_vals}}}
a := shuf{[4]u64, x, 4b1010}; each{bind{fs,a}, slice{sh,0,h}}
if (l%2) fs{x, tupsel{h, sh}} # first 4 shuffle vectors for 11≤𝕨≤61; only uses the low half of the input
b := shuf{[4]u64, x, 4b3232}; each{bind{fs,b}, slice{sh,-h}} def rcsh4_dom = replicate{bind{>=,64}, replicate{fact_tab==1, fact_inds}}
} rcsh4_dat:*i8 = join{join{each{{wv}=>get_shuf_data{wv, 4}, rcsh4_dom}}}
} rcsh4_lkup:*i8 = shiftright{0, scan{+, fold{|, table{==, rcsh4_dom, iota{64}}}}}
def get_rep_iter{V, wv==2}{x, gen} = {
def s = shuf{[4]u64, x, 4b3120} def read_shuf_vecs{l, ellw:u64, shp:P} = { # tuple of byte selectors in 1<<ellw
each{{q}=>gen{V~~q}, unpackQ{s, s}} def V = eltype{P}
} def double{x:X & hasarch{'AVX2'}} = {
def get_rep_iter{V==[4]u64, wv} = {
def step = 4
def sh = each{bind{base,4}, get_shufs{step, wv, wv}}
{x, gen} => each{{s}=>gen{shuf{V, x, s}}, sh}
}
def read_shuf_vecs{l, elbytes:u64, shp:*[32]i8} = {
def double{x} = {
s:=shuf{[4]u64, x, 4b3120}; s+=s s:=shuf{[4]u64, x, 4b3120}; s+=s
each{bind{~~,[32]i8},unpackQ{s, s + type{s}**1}} r:=each{bind{~~,[32]i8},unpackQ{s, s + X**1}}
r
}
def double{x:X & hasarch{'AARCH64'}} = {
s:= x+x
zip{s, s + X**1}
} }
def doubles{n,tup} = slice{join{each{double,tup}}, 0, n} def doubles{n,tup} = slice{join{each{double,tup}}, 0, n}
def sh = each{{v}=>{r:=v}, copy{l, [32]i8**0}}
def tlen{e} = (l+(-l)%e)/e # Length for e bytes, rounded up def sh = each{{v}=>{r:=v}, l**V**0}
def tlen{e} = cdiv{l, e} # Length for e bytes, rounded up
def set{i} = { tupsel{i,sh} = each{bind{load,shp},i} } def set{i} = { tupsel{i,sh} = each{bind{load,shp},i} }
def ext{e} = { def ext{e} = {
def m = tlen{2*e}; def n = tlen{e} # m<n def m = tlen{2*e}; def n = tlen{e} # m<n
if (elbytes<=e) set{slice{iota{n},m}} if (ellw <= lb{e}) set{slice{iota{n},m}}
else slice{sh,0,n} = doubles{n,slice{sh,0,m}} else slice{sh,0,n} = doubles{n,slice{sh,0,m}}
} }
set{iota{tlen{8}}}; ext{4}; ext{2}; ext{1} set{iota{tlen{8}}}; ext{4}; ext{2}; ext{1}
sh sh
} }
def rep_const_shuffle{V, wv, onreps, xv:*V, rv:*V, n:u64} = { def rep_const_shuffle{V, wv, onreps, xv:*V, rv:*V, n:u64} = { # onreps{inputVector, {nextOutputVector} => ...}
def step = vcount{V} def step = vcount{V}
nv := n / step nv := n / step
j:u64 = 0 j:u64 = 0
@ -81,64 +88,102 @@ def rep_const_shuffle{V, wv, onreps, xv:*V, rv:*V, n:u64} = {
if (q) homMaskStoreF{rv+e, maskOf{V, q}, s} if (q) homMaskStoreF{rv+e, maskOf{V, q}, s}
} }
} }
def rep_const_shuffle{V, wv, xv:*V, rv:*V, n:u64} = rep_const_shuffle{V, wv, get_rep_iter{V, wv}, xv, rv, n}
def rcsh_vals = slice{basic_rep, 1} # Handle 2 specially if (hasarch{'AVX2'}) {
rcsh_offs:*u8 = shiftright{0, scan{+,rcsh_vals}} def rep_iter_from_sh{sh}{x, gen} = {
rcsh_data:*i8 = join{join{each{get_shuf_data, rcsh_vals}}} def l = tuplen{sh}
fn rcsh_sub{wv}(elbytes:u64, x:*i8, r:*i8, n:u64, sh:*[32]i8) : void = { def h = l>>1
def V = [32]i8 def fs{v, s} = gen{sel{[16]i8, v, s}}
def st = read_shuf_vecs{wv, elbytes, sh} a := shuf{[4]u64, x, 4b1010}; each{bind{fs,a}, slice{sh,0,h}}
rep_const_shuffle{V, wv, rep_iter_from_sh{st}, *V~~x, *V~~r, n} if (l%2) fs{x, tupsel{h, sh}}
b := shuf{[4]u64, x, 4b3232}; each{bind{fs,b}, slice{sh,-h}}
} }
def rcsh4_dom = replicate{bind{>=,64}, replicate{fact_tab==1, fact_inds}} def get_rep_iter{V, wv==2}{x, gen} = {
rcsh4_dat:*i8 = join{join{each{{wv}=>get_shuf_data{wv, 4}, rcsh4_dom}}} def s = shuf{[4]u64, x, 4b3120}
rchs4_lkup:*i8 = shiftright{0, scan{+, fold{|, table{==, rcsh4_dom, iota{64}}}}} each{{q}=>gen{V~~q}, unpackQ{s, s}}
fn rep_const_shuffle_partial4(wv:u64, elbytes:u64, x:*i8, r:*i8, n:u64) : void = { }
def get_rep_iter{V==[4]u64, wv} = {
def step = 4
def sh = each{bind{base,4}, get_shufs{step, wv, wv}}
{x, gen} => each{{s}=>gen{shuf{V, x, s}}, sh}
}
def rep_const_shuffle{V, wv, xv:*V, rv:*V, n:u64} = rep_const_shuffle{V, wv, get_rep_iter{V, wv}, xv, rv, n}
} else if (hasarch{'AARCH64'}) {
def rep_iter_from_sh{sh}{x, gen} = {
each{{s} => gen{sel{[16]u8, x, s}}, sh}
}
def rep_const_shuffle{V, wv==2, xv0:*V, rv0:*V, n:u64} = {
def E = ty_u{eltype{V}}
rv:= *E~~rv0
@for (x in *E~~xv0 over i to n) { # autovectorized well enough, probably
store{rv, i*2, x}
store{rv, i*2+1, x}
}
}
}
fn rep_const_shuffle_partial4(wv:u64, ellw:u64, x:*i8, r:*i8, n:u64) : void = {
def h = 4 def h = 4
def sh = read_shuf_vecs{h, elbytes, *[32]i8~~rcsh4_dat + 4*load{rchs4_lkup,wv}} def V = sdtype
def V = [32]i8 def sh = read_shuf_vecs{h, ellw, *V~~rcsh4_dat + h*load{rcsh4_lkup,wv}}
def step = vcount{V} # Bytes written def step = vcount{V} # Bytes written
def wvb = wv * elbytes def wvb = wv << ellw
def hs = (h*step) / wvb # Actual step size in argument elements def hs = (h*step) / wvb # Actual step size in argument elements
re := r + n*wvb - h*step def shufbase{i & hasarch{'AVX2'}} = shuf{[4]u64, load{*V~~(x+i)}, 4b1010}
def shufbase{i & hasarch{'AARCH64'}} = load{*V~~(x+i)}
def shufrun{a, s} = sel{[16]i8, a, s} # happens to be the same across AVX2 & NEON
i:u64 = 0 i:u64 = 0
re := r + n*wvb - h*step
while (r <= re) { while (r <= re) {
a := shuf{[4]u64, load{*V~~(x+i),0}, 4b1010} a := shufbase{i}
@unroll (j to h) store{*V~~r, j, sel{[16]i8, a, tupsel{j,sh}}} @unroll (j to h) store{*V~~r, j, shufrun{a, tupsel{j,sh}}}
i += hs*elbytes i += hs << ellw
r += hs*wvb r += hs*wvb
} }
re+= (h-1)*step re+= (h-1)*step
a := shuf{[4]u64, load{*V~~(x+i),0}, 4b1010} a:= shufbase{i}
s:= V**0 s:= V**0
def end = makelabel{} def end = makelabel{}
@unroll (j to h) { @unroll (j to h) {
s = sel{[16]i8, a, tupsel{j,sh}} s = shufrun{a, tupsel{j,sh}}
if (r > re) goto{end} if (r > re) goto{end}
store{*V~~r, 0, s} store{*V~~r, 0, s}
r+= step r+= step
} }
setlabel{end} setlabel{end}
q := (re+step) - r q := (re+step) - r
if (q) homMaskStoreF{*V~~r, maskOf{V, q}, s} if (q) homMaskStoreF{*V~~r, maskOf{V, q}, s}
} }
fn rep_const_shuffle_any(wv:u64, elbytes:u64, x:*i8, r:*i8, n:u64) : void = {
if (wv > tupsel{-1,rcsh_vals}) {
return{rep_const_shuffle_partial4(wv, elbytes, x, r, n)} fn rcsh_sub{wv, V}(ellw:u64, x:*i8, r:*i8, n:u64, sh:*V) : void = {
def st = read_shuf_vecs{wv, ellw, sh}
rep_const_shuffle{V, wv, rep_iter_from_sh{st}, *V~~x, *V~~r, n}
} }
n *= elbytes
fn rep_const_shuffle_any(wv:u64, ellw:u64, x:*i8, r:*i8, n:u64) : void = {
if (wv > tupsel{-1,rcsh_vals}) {
return{rep_const_shuffle_partial4(wv, ellw, x, r, n)}
}
n <<= ellw
ri := wv - tupsel{0,rcsh_vals} ri := wv - tupsel{0,rcsh_vals}
sh := *[32]i8~~rcsh_data + load{rcsh_offs,ri} sh := *sdtype~~rcsh_data + load{rcsh_offs,ri}
def try{k} = { if (wv==k) rcsh_sub{k}(elbytes, x, r, n, sh) } def try{k} = { if (wv==k) rcsh_sub{k, sdtype}(ellw, x, r, n, sh) }
each{try, rcsh_vals} each{try, rcsh_vals}
} }
def rep_const_broadcast{T, kv, loop, wv:u64, x:*T, r:*T, n:u64} = { def rep_const_broadcast{T, kv, loop, wv:u64, x:*T, r:*T, n:u64} = {
assert{kv > 0} assert{kv > 0}
def V = [256/width{T}]T def V = [arch_defvw/width{T}]T
@for (x over n) { @for (x over n) {
v := V**x v := V**x
@loop (j to kv) store{*V~~r, j, v} @loop (j to kv) store{*V~~r, j, v}
@ -162,7 +207,7 @@ fn rep_const{T}(wv:u64, x:*void, r:*void, n:u64) : void = {
} }
} }
def wT = width{T} def wT = width{T}
def vn = 256/wT def vn = arch_defvw/wT
def V = [vn]T def V = [vn]T
def max_shuffle = 2*vn def max_shuffle = 2*vn
if (wv <= max_shuffle) { if (wv <= max_shuffle) {
@ -170,7 +215,7 @@ fn rep_const{T}(wv:u64, x:*void, r:*void, n:u64) : void = {
if (wv==k) return{rep_const_shuffle{V, k, *V~~x, *V~~r, n}} if (wv==k) return{rep_const_shuffle{V, k, *V~~x, *V~~r, n}}
} }
specialize{2} specialize{2}
rep_const_shuffle_any(wv, wT/8, *i8~~x, *i8~~r, n) rep_const_shuffle_any(wv, lb{wT/8}, *i8~~x, *i8~~r, n)
} else { } else {
kv := wv / vn kv := wv / vn
@unroll (k from (max_shuffle/vn) to 4) { @unroll (k from (max_shuffle/vn) to 4) {

View File

@ -121,6 +121,7 @@ def vshl{a:T, b:T, n & knum{n}} = emit{T, ntyp{'vext', T}, a, b, n}
def zip1{a:T, b:T & nvec{T}} = emit{T, ntyp{'vzip1', T}, a, b} def zip1{a:T, b:T & nvec{T}} = emit{T, ntyp{'vzip1', T}, a, b}
def zip2{a:T, b:T & nvec{T}} = emit{T, ntyp{'vzip2', T}, a, b} def zip2{a:T, b:T & nvec{T}} = emit{T, ntyp{'vzip2', T}, a, b}
def zip{a:T, b:T & nvec{T}} = tup{zip1{a,b}, zip2{a,b}}
def packLo{x:T, y:T & nvec{T}} = { def H=ty_half{T}; emit{H, ntyp{'vuzp1', H}, H~~x, H~~y} } def packLo{x:T, y:T & nvec{T}} = { def H=ty_half{T}; emit{H, ntyp{'vuzp1', H}, H~~x, H~~y} }
def packHi{x:T, y:T & nvec{T}} = { def H=ty_half{T}; emit{H, ntyp{'vuzp2', H}, H~~x, H~~y} } def packHi{x:T, y:T & nvec{T}} = { def H=ty_half{T}; emit{H, ntyp{'vuzp2', H}, H~~x, H~~y} }