Shuffle selection using >1 data register
This commit is contained in:
parent
b651ca73d3
commit
dc6e49b289
@ -18,59 +18,90 @@ def gather{def:T, b:B, idx:[4]i32, M & w256{T,64}} = {
|
||||
else T ~~ emit{[4]i64, '_mm256_i32gather_epi64', *i64~~b, idx, width{eltype{B}}/8}
|
||||
}
|
||||
|
||||
def wrapChk{cw0, VI,xlf, M} = {
|
||||
cw:= cw0 + (xlf & VI~~(cw0<broadcast{VI, 0}))
|
||||
if (any{M{ty_u{cw} >= ty_u{xlf}}}) return{0}
|
||||
cw
|
||||
}
|
||||
|
||||
def shuf_select{ri, rd, TI, w, r, wl, xl, selx} = {
|
||||
def VI = [ri]TI
|
||||
def ext = ri/rd
|
||||
xlf:= broadcast{VI, cast_i{TI, xl}}
|
||||
maskedLoop{ri, wl, {i, M} => {
|
||||
cw:= wrapChk{loadBatch{w, i, VI}, VI,xlf, M}
|
||||
is:= (if (ext>1) i<<lb{ext}; else i)
|
||||
def se{e, c, o} = {
|
||||
c2:= shuf{[4]u64, c+c, 4b3120}
|
||||
each{
|
||||
{c,o} => se{e*2, VI~~c, o},
|
||||
unpackQ{c2, c2+broadcast{VI,1}},
|
||||
2*o + iota{2}
|
||||
}
|
||||
}
|
||||
def se{e==ext, c, o} = {
|
||||
io:= is+o
|
||||
got:= selx{c}
|
||||
def s{M} = storeBatch{r, io, got, M}
|
||||
if (ext==1 or not M{0}) s{M}
|
||||
else if (io*rd+rd <= wl) s{maskNone}
|
||||
else { if (io*rd < wl) s{maskAfter{wl & (rd-1)}}; return{1} }
|
||||
}
|
||||
se{1, cw, 0}
|
||||
}}
|
||||
}
|
||||
|
||||
def makesel{VI,VD, x0,logv} = {
|
||||
x:= *VD~~x0
|
||||
def halves{v} = each{bind{shuf, [4]u64, v}, tup{4b1010, 4b3232}}
|
||||
def readx{l,o} = each{bind{readx,l-1}, o + iota{2}<<(l-2)}
|
||||
def readx{l==0,o} = shuf{[4]u64, load{x}, 4b1010}
|
||||
def readx{l==1,o} = halves{load{x, o}}
|
||||
xd:= readx{logv, 0}
|
||||
|
||||
def bblend {m}{ft} = blend{tupsel{0,ft}, tupsel{1,ft}, m}
|
||||
def bblendn{m}{tf} = bblend{m}{reverse{tf}}
|
||||
def bb{c}{f, v} = (if (f) bblendn{c<v}; else bblend{(c&v)==v})
|
||||
|
||||
def bs{b, c, x} = sel{[16]i8, x, c}
|
||||
def bs{b, c, x & tuplen{b}>0} = {
|
||||
tupsel{0,b}{each{bind{bs, slice{b,1}, c}, x}}
|
||||
}
|
||||
|
||||
def i = iota{logv}
|
||||
def vs = each{bind{broadcast,VI}, 16<<reverse{i}}
|
||||
{c} => VD~~bs{each{bb{c},i==0,vs}, c, xd}
|
||||
}
|
||||
|
||||
select{rw, TI, TD}(w0:*void, x0:*void, r0:*void, wl:u64, xl:u64) : u1 = {
|
||||
w:= *TI ~~ w0
|
||||
x:= *TD ~~ x0
|
||||
r:= *TD ~~ r0
|
||||
|
||||
def wrap{cw0, VI,xlf} = cw0 + (xlf & VI~~(cw0<broadcast{VI, 0}))
|
||||
|
||||
def wd = width{TD}
|
||||
def wi = width{TI}
|
||||
if (wi==8 and wd<=32 and xl*wd<=128) {
|
||||
def ri = rw/wi; def VI = [ri]TI
|
||||
def rd = rw/wd; def VD = [rd]TD
|
||||
def xlf = broadcast{VI, cast_i{TI, xl}}
|
||||
xd := shuf{[4]u64, load{*VD~~x}, 4b1010}
|
||||
maskedLoop{ri, wl, {i, M} => {
|
||||
cw:= wrap{loadBatch{w, i, VI}, VI,xlf}
|
||||
if (any{M{ty_u{cw} >= ty_u{xlf}}}) return{0}
|
||||
is := (if (wd>wi) i<<lb{wd/wi}; else i)
|
||||
def se{w, c, o} = {
|
||||
c2 := shuf{[4]u64, c+c, 4b3120}
|
||||
each{
|
||||
{c,o} => se{w*2, VI~~c, o},
|
||||
unpackQ{c2, c2+broadcast{VI,1}},
|
||||
2*o + iota{2}
|
||||
}
|
||||
}
|
||||
def se{w==wd, c, o} = {
|
||||
io:= is+o
|
||||
got:= VD~~sel{[16]i8, xd, c}
|
||||
def s{M} = storeBatch{r, io, got, M}
|
||||
if (wd==wi or not M{0}) s{M}
|
||||
else if (io*rd+rd <= wl) s{maskNone}
|
||||
else if (io*rd < wl) s{maskAfter{wl & (rd-1)}}
|
||||
}
|
||||
se{wi, cw, 0}
|
||||
}}
|
||||
return{1}
|
||||
def wd = width{TD}; def rd = rw/wd
|
||||
def wi = width{TI}; def ri = rw/wi
|
||||
def shuf_select{l} = {
|
||||
shuf_select{ri, rd, TI, w, r, wl, xl, makesel{[ri]TI,[rd]TD, x,l}}
|
||||
}
|
||||
if (wi==8 and wd<=32 and xl*wd<=128 ) { shuf_select{0} }
|
||||
else if (wi==8 and wd<=16 and xl*wd<=128<<1) { shuf_select{1} }
|
||||
else if (wi==8 and wd<=16 and xl*wd<=128<<2) { shuf_select{2} }
|
||||
else if (wi==8 and wd<= 8 and xl*wd<=128<<3) { shuf_select{3} }
|
||||
else {
|
||||
def TIE = i32
|
||||
def TDE = tern{wd<32, u32, TD}
|
||||
def bulk = rw / width{TDE}
|
||||
def VI = [bulk]TIE
|
||||
def VD = [bulk]TDE
|
||||
def xlf = broadcast{VI, cast_i{TIE, xl}}
|
||||
|
||||
maskedLoop{bulk, wl, {i, M} => {
|
||||
cw:= wrapChk{loadBatch{w, i, VI}, VI,xlf, M}
|
||||
got:= gather{broadcast{VD,0}, x, cw, M}
|
||||
if (TDE!=TD) got&= broadcast{VD, (1<<wd)-1}
|
||||
storeBatch{r, i, got, M}
|
||||
}}
|
||||
}
|
||||
|
||||
def TIE = i32
|
||||
def TDE = tern{width{TD}<32, u32, TD}
|
||||
def bulk = rw / width{TDE}
|
||||
def VI = [bulk]TIE
|
||||
def VD = [bulk]TDE
|
||||
def xlf = broadcast{VI, cast_i{TIE, xl}}
|
||||
|
||||
maskedLoop{bulk, wl, {i, M} => {
|
||||
cw:= wrap{loadBatch{w, i, VI}, VI,xlf}
|
||||
if (any{M{ty_u{cw} >= ty_u{xlf}}}) return{0}
|
||||
got:= gather{broadcast{VD,0}, x, cw, M}
|
||||
if (TDE!=TD) got&= broadcast{VD, (1<<width{TD})-1}
|
||||
storeBatch{r, i, got, M}
|
||||
}}
|
||||
1
|
||||
}
|
||||
def select{TI, TD} = select{256, TI, TD}
|
||||
|
||||
Loading…
Reference in New Issue
Block a user