diff --git a/src/singeli/src/fold.singeli b/src/singeli/src/fold.singeli index 32e9597e..16a2064f 100644 --- a/src/singeli/src/fold.singeli +++ b/src/singeli/src/fold.singeli @@ -74,7 +74,7 @@ fn fold_assoc_0{T==f64, op if has_simd}(x:*T, len:u64) : T = { export{'si_sum_f64', fold_assoc_0{f64,+}} -def select_rows_pow2{T, x0, r0, nv, k} = { +def extract_column_pow2{T, x0, r0, nv, k} = { def V = [arch_defvw / width{T}]T xv := *V~~x0 @for (r in *V~~r0 over i to nv) { @@ -109,66 +109,55 @@ def select_rows_pow2{T, x0, r0, nv, k} = { } } -fn select_rows_byte(x0:*void, r0:*void, n:usz, l:usz, e:u8) : usz = { - n <<= e - def vl = arch_defvw / 8 - def vh = vl / 2 - def thr = min{vl+2, 20} - if ((not has_simd) or n < vl or l > usz~~thr>>e) return{0} - if (has_simd and (l & (l-1)) == 0) { - nv := n / vl - def try_unzip{T, k} = if (k < thr and l == k) { - select_rows_pow2{T, x0, r0, nv, k} - goto{'ret'} - } - # 10 loops: i8 2,4,8,16; i16 2,4,8; i32 2,4; i64 2 - @unroll (ek to 4) if (e == ek) { - def T = ty_s{8<>e) * nv} +def extract_column_modperm{x0, r0, nv, l, el, vl} = { + # Build modular permutations + def V = [vl]u8 + def H = [vl/2]u16 + p2 := ctz{l}; l >>= p2 # Decompose into l<x) else { - ms := (V**16 & sii) == (V**16 &~ iota{V}) - {x} => homBlend{x, shuf{[4]u64, x, 2,3,0,1}, ms} - } - # Blend masks - def mg = { # Iteration i should select where mg == V**i - ss := (si < V**(l8<x) else { + ms := (V**16 & sii) == (V**16 &~ iota{V}) + {x} => homBlend{x, shuf{[4]u64, x, 2,3,0,1}, ms} + } + + # Blend masks + def mg = { # Iteration i should select where mg == V**i + ss := (si < V**(l8<>e) * nv} } - 0 + + # Handle odd and even strides separately + if (p2 == 0) { + loop{{r, i, write} => write{swap_ms{r}}} + } else { + # Store results + def add_res = { + ra := V**0 # Accumulator + i := make{V, iota{vl}%16} + o := V**(u8~~1< { ra = homBlend{shuf{[16]u8, ra, sh}, r, bl} } + } + il := make{V, iota{vl} % 16} + # Shuffle to undo interleaving of add_res + def __shr{x:(V), sh if hasarch{'X86_64'}} = V~~(H~~x >> sh) + def __shl{x:(V), sh if hasarch{'X86_64'}} = V~~(H~~x << sh) + def __shr{x:(V), sh:T if hasarch{'AARCH64'} and not isvec{T}} = x << V**cast_i{u8,-sh} + def __shl{x:(V), sh:T if hasarch{'AARCH64'} and not isvec{T}} = x << V**cast_i{u8, sh} + def uz_lane = { + l := V**(u8~~1<>p2) + dz := (il & l) | (h &~ il)>>(4-e) # low, high->middle + dz |= (il &~ (l | h))<high + shuf{[16]u8, ., dz} + } + # Adjust modular permutation to apply after unzipping + si = uz_lane{(si & V**(vl - u8~~1<> p2} + si ^= il &~ V**((16 - u8~~1<> p2) + # Cross-lane follow-up + def cross = if (vl == 16) { {x}=>x } else { + si = shuf{[4]u64, si, 0,1,0,1} + assert{p2 <= 2} + cr := make{[8]u32, tr_iota{0,2,1}} + if (p2 > 1) cr = make{[8]u32, tr_iota{2,0,1}} + shuf{[8]u32, ., cr} + } + # Run, writing every 1< { + def ra = add_res{r} + if ((plo &~ i) == 0) write{cross{uz_lane{ra}}} + }} + } +} + +# Select one element out of every l, element width 1< usz~~thr>>el or l<= thr) return{0} + nv := n / vl + if (has_simd and (l & (l-1)) == 0) { + def try_unzip{T, k} = if (k < thr and l == k) { + extract_column_pow2{T, x0, r0, nv, k} + goto{'ret'} + } + # 10 loops: i8 2,4,8,16; i16 2,4,8; i32 2,4; i64 2 + @unroll (ek to 4) if (el == ek) { + def T = ty_s{8<>el) * nv } @@ -294,7 +356,7 @@ def fold_rows_bit_lt64{ } } -fn select_rows_bit_lt64(xp:*u64, rp:*u64, n:usz, l:usz, o:usz) : void = { +fn extract_column_bit_lt64(xp:*u64, rp:*u64, n:usz, l:usz, o:usz) : void = { assert{l < 64}; assert{o < l} # Row length, and offset within row def run_loop2{loop} = loop{{a,b} => a>>o} def run_loop4{m, t, loop} = loop{{x} => x<<(l-1-o)} @@ -439,5 +501,5 @@ fn or_rows_bit(xp:*u64, rp:*u64, n:usz, l:usz, op_and:u1) : void = { } export{'si_xor_rows_bit', xor_rows_bit} export{'si_or_rows_bit', or_rows_bit} -export{'si_select_cells_bit_lt64', select_rows_bit_lt64} -export{'si_select_cells_byte', select_rows_byte} +export{'si_select_cells_bit_lt64', extract_column_bit_lt64} +export{'si_select_cells_byte', extract_column}