From 58f4f15f8faf331bd8ae2ccb7693842735bd78c6 Mon Sep 17 00:00:00 2001 From: Marshall Lochbaum Date: Fri, 29 Dec 2023 20:11:18 -0500 Subject: [PATCH] Move more functionality to avx512.singeli --- src/singeli/src/avx512.singeli | 24 +++++++++++++++++++++++- src/singeli/src/base.singeli | 2 +- src/singeli/src/scan.singeli | 11 +++++------ src/singeli/src/slash.singeli | 11 +++++------ 4 files changed, 34 insertions(+), 14 deletions(-) diff --git a/src/singeli/src/avx512.singeli b/src/singeli/src/avx512.singeli index d73ab10f..9f7936fb 100644 --- a/src/singeli/src/avx512.singeli +++ b/src/singeli/src/avx512.singeli @@ -1,11 +1,33 @@ +local { + def ismask{M} = if (isvec{M}) u1==eltype{M} else 0 + def suf{T} = { + if (isfloat{T}) (if (width{T}==32) 'ps' else 'pd') + else merge{'epi', fmtnat{width{T}}} + } + def suf{V & isvec{V}} = suf{eltype{V}} + def pref{w} = merge{'_mm', if (w==128) '' else fmtnat{w}, '_'} + def pref{V & isvec{V}} = pref{width{V}} +} + local def re_mask{M, sub} = { def l = vcount{M}; def w = max{32,l} sub{fmtnat{l}, fmtnat{w}, ty_u{w}} } -local def ismask{M} = if (isvec{M}) u1==eltype{M} else 0 def reinterpret{M, a:T & ismask{M} & width{T}==width{M}} = { re_mask{M, {l,w,W} => emit{M, merge{'_cvtu',w,'_mask',l}, promote{W, a}}} } def reinterpret{T, a:M & ismask{M} & width{T}==width{M}} = { re_mask{M, {l,w,W} => cast_i{T, emit{W, merge{'_cvtmask',l,'_u',w}, a}}} } + +def maskStore{p:*V, m:M, v:V & ismask{M} & isvec{V} & vcount{M}==vcount{V}} = { + emit{void, merge{pref{V}, 'mask_storeu_', suf{V}}, p, m, v} +} + +def topMaskReg{x:V} = emit{[vcount{V}]u1, merge{pref{V},'mov',suf{V},'_mask'}, x} +def topMask{x:T & isvec{T} & 512==width{T}} = ty_u{vcount{T}}~~topMaskReg{x} +def homMask{x:T & isvec{T} & 512==width{T}} = topMask{x} + +def maskToHom{T, x:M & ismask{M} & isvec{T} & vcount{M}==vcount{T}} = { + emit{T, merge{pref{T},'movm_',suf{T}}, x} +} diff --git a/src/singeli/src/base.singeli b/src/singeli/src/base.singeli index 46d6532e..68ded7f3 100644 --- a/src/singeli/src/base.singeli +++ b/src/singeli/src/base.singeli @@ -166,7 +166,7 @@ def lvec{T, n, w & isvec{T} & vcount{T}==n & elwidth{T}==w} = 1 def { absu,andAllZero,andnz,b_getBatch,clmul,cvt,extract,fold_addw,half, homAll,homAny,homBlend,homMask,homMaskStore,homMaskStoreF,loadBatchBit, - loadLow,make,mulw,mulh,narrow,narrowPair,packHi,packLo,packQ,pair,pdep, + 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 diff --git a/src/singeli/src/scan.singeli b/src/singeli/src/scan.singeli index ace6b2d3..9deac4cd 100644 --- a/src/singeli/src/scan.singeli +++ b/src/singeli/src/scan.singeli @@ -105,25 +105,24 @@ fn scan_neq{& hasarch{'PCLMUL'}}(init:u64, x:*u64, r:*u64, nw:u64) : void = { clmul_scan_ne_any{}(*void~~x, *void~~r, init, nw, -(u64~~1)) } fn scan_neq{& hasarch{'AVX512BW', 'VPCLMULQDQ', 'GFNI'}}(init:u64, x:*u64, r:*u64, nw:u64) : void = { - def emitM{T,s,...a} = emit{T, merge{'_mm512_',s}, ...a} def V = [8]u64 def sse{a} = make{[2]u64, a, 0} carry := sse{init} # xor-scan on bytes xmat := V**base{256, 1<<(8-iota{8}) - 1} - def xor8 = emitM{V, 'gf2p8affine_epi64_epi8', ., xmat, 0} + def xor8 = emit{V, '_mm512_gf2p8affine_epi64_epi8', ., xmat, 0} # Exclusive xor-scan on one word def exor64 = clmul{., sse{1<<64 - 2}, 0} @for (xv in *V~~x, rv in *V~~r over i to cdiv{nw,vcount{V}}) { x8 := xor8{xv} - hb := sse{u64~~emitM{[64]u1, 'movepi8_mask', x8}} - xh := exor64{hb} # Exclusive xor of high bits + hb := sse{topMask{[64]u8~~x8}} + xh := exor64{hb} # Exclusive xor of high bits (xh ^ hb for inclusive) xc := xh ^ carry - v := x8 ^ emitM{V, 'movm_epi8', [64]u1~~extract{xc,0}} + v := x8 ^ V~~maskToHom{[64]u8, [64]u1~~extract{xc,0}} carry = (xc ^ hb) ^ shuf{[4]u32, xh, 4b3232} rem:= nw - 8*i if (rem < 8) { - emitM{void, 'mask_storeu_epi64', *V~~r+i, [8]u1~~(~(u8~~0xff<=32) 'AVX512F' else 'AVX512VBMI2'}}(w:*u64, def vl = 512/wt def V = [vl]T def X = getter{c, V, x} - def wu = max{32,vl} + def I = ty_u{vl} @for (w in *(ty_u{vl})~~w over cdiv{l,vl}) { - def I = ty_u{wu} def emitT{O, name, ...a} = emit{O, merge{'_mm512_',name,'_epi',f{wt}}, ...a} - def to_mask{a} = emit{[vl]u1, merge{'_cvtu',f{wu},'_mask',f{vl}}, a} - m := to_mask{promote{I,w}} + m := [vl]u1~~w c := popc{w} x := X{} # The compress-store instruction performs very poorly on Zen4, # and is also a lot worse than the following on Tiger Lake # emitT{void, 'mask_compressstoreu', r, m, x} cs := cast_i{I,promote{i64,1}<<(c%64) - 1} - if (wu==64) cs -= cast_i{I,c}>>6 + if (vl==64) cs -= cast_i{I,c}>>6 v := emitT{V, 'mask_compress', x, m, x} - emitT{void, 'mask_storeu', r, to_mask{cs}, v} + maskStore{*V~~r, [vl]u1~~cs, v} r += c } }