def gats

bend-parallel@0.1.0.0 · sort.bend:456

import Base

456# The gather's forks: d levels down to blocks of 2^e slots under block i,
457# with two handles each to perm, src and dst. Block i reads only its
458# entries of perm, reads src anywhere (no lane writes src) and writes only
459# its slots of dst.
460def gats(~T: Data, +d: Nat, +i: U32, +e: Nat, fp: Array<U32> & Array<U32>, fs: Array<T> & Array<T>, fd: Array<T> & Array<T>) -> Moved<T>:
461  match d:
462    case 0n:
463      (pm, pm2) = fp
464      (src, src2) = fs
465      (dst, dst2) = fd
466      +at = U32.shln(i, e)
467      moved.rejoin(~T, pm2, src2, dst2, gat.go(~T, U32.to_nat((U32.shln(1, e) - 1 : U32)), at,
468        Sd{src, dst}, pm[at]))
469    case 1n+p:
470      (p0, p1) = fp
471      (s0, s1) = fs
472      (d0, d1) = fd
473      +j = (i * 2 : U32)
474      l r = gats(~T, p, j, e, Array.fork(U32, p0), Array.fork(T, s0), Array.fork(T, d0))
475        gats(~T, p, (j + 1 : U32), e, Array.fork(U32, p1), Array.fork(T, s1), Array.fork(T, d1))
476      moved.join(~T, l, r)