Fast ≍˘ for matching-type type-sized cells

This commit is contained in:
Marshall Lochbaum 2023-12-30 20:27:05 -05:00
parent 9931c1756c
commit eb34213f0f
3 changed files with 55 additions and 4 deletions

View File

@ -9,6 +9,7 @@ B shape_c2(B, B, B);
B transp_c2(B, B, B);
B fold_rows(Md1D* d, B x); // from fold.c
B takedrop_highrank(bool take, B w, B x); // from sfns.c
B try_interleave_cells(B w, B x, ur xr, ur xk, usz* xsh); // from transpose.c
// X - variable name; XSH - its shape; K - number of leading axes that get iterated over; SLN - number of slices that will be made; DX - additional refcount count to add to x
#define S_KSLICES(X, XSH, K, SLN, DX)\
@ -609,6 +610,10 @@ NOINLINE B for_cells_AA(B f, B w, B x, ur wcr, ur xcr, u32 chr) {
if (rsh) shcpy(rsh, zsh, zk);
decG(w); decG(x); return taga(r);
}
if (rtid==n_couple && wr==xr) {
B r = try_interleave_cells(w, x, xr, xk, xsh);
if (!q_N(r)) { decG(w); decG(x); return r; }
}
}
if (isPervasiveDy(f)) {
if (TI(w,elType)==el_B || TI(x,elType)==el_B) goto generic;

View File

@ -66,6 +66,44 @@ typedef void (*TranspFn)(void*,void*,u64,u64,u64,u64);
#endif
#ifdef __BMI2__
static void interleave_bits(u64* rp, void* x0v, void* x1v, usz n) {
u32* x0 = (u32*)x0v; u32* x1 = (u32*)x1v;
for (usz i=0; i<BIT_N(n); i++) rp[i] = _pdep_u64(x0[i], 0x5555555555555555) | _pdep_u64(x1[i], 0xAAAAAAAAAAAAAAAA);
}
#endif
// Interleave arrays, 𝕨≍⎉(-xk)𝕩. Doesn't consume.
// Return bi_N if there isn't fast code.
B try_interleave_cells(B w, B x, ur xr, ur xk, usz* xsh) {
assert(RNK(w)==xr && xr>=1);
u8 xe = TI(x,elType); if (xe!=TI(w,elType) || xe==el_B) return bi_N;
usz csz = shProd(xsh, xk, xr);
if (csz & (csz-1)) return bi_N; // Not power of 2
u8 xlw = elwBitLog(xe);
usz n = shProd(xsh, 0, xk);
usz ia = 2*n*csz;
Arr *r;
#ifdef __BMI2__
if (csz==1 && xlw==0) {
u64* rp; r=m_bitarrp(&rp, ia);
interleave_bits(rp, bitarr_ptr(w), bitarr_ptr(x), ia);
} else
#endif
#if SINGELI
if (csz<=64>>xlw && csz<<xlw>=8) { // Require CPU-sized cells
void* rv;
if (xlw==0) { u64* rp; r = m_bitarrp(&rp, ia); rv=rp; }
else rv = m_tyarrp(&r,elWidth(xe),ia,el2t(xe));
si_interleave[CTZ(csz<<xlw)-3](rv, tyany_ptr(w), tyany_ptr(x), n);
} else
#endif
return bi_N;
usz* sh = arr_shAlloc(r, xr+1);
shcpy(sh, xsh, xk); sh[xk]=2; shcpy(sh+xk+1, xsh+xk, xr-xk);
return taga(r);
}
static void transpose_move(void* rv, void* xv, u8 xe, usz w, usz h) {
assert(xe!=el_bit); assert(xe!=el_B);
transposeFns[elwByteLog(xe)](rv, xv, w, h, w, h);
@ -89,11 +127,9 @@ static Arr* transpose_noshape(B* px, usz ia, usz w, usz h) {
} else if (xe==el_bit) {
#ifdef __BMI2__
if (h==2) {
u32* x0 = (u32*)bitarr_ptr(x);
u64* rp; r=m_bitarrp(&rp, ia);
Arr* x1o = TI(x,slice)(inc(x),w,w);
u32* x1 = (u32*) ((TyArr*)x1o)->a;
for (usz i=0; i<BIT_N(ia); i++) rp[i] = _pdep_u64(x0[i], 0x5555555555555555) | _pdep_u64(x1[i], 0xAAAAAAAAAAAAAAAA);
interleave_bits(rp, bitarr_ptr(x), ((TyArr*)x1o)->a, ia);
mm_free((Value*)x1o);
} else if (w==2) {
u64* xp = bitarr_ptr(x);

View File

@ -150,6 +150,14 @@ def transpose_with_kernel{T, k, kh, call_base, rp:*T, xp:*T, w, h, ws, hs} = {
}
}
# Interleave n values of type T from x0 and x1 into r
fn interleave{T}(r0:*void, x0:*void, x1:*void, n:u64) : void = {
rp := *T~~r0
@for (x0 in *T~~x0, x1 in *T~~x1 over i to n) {
store{rp, i*2, x0}; store{rp, i*2+1, x1}
}
}
fn transpose{T, k, kh}(r0:*void, x0:*void, w:u64, h:u64, ws:u64, hs:u64) : void = {
# Scalar transpose defined in C
def ts = if (T==i8) 'i8' else if (T==i16) 'i16' else if (T==i32) 'i32' else 'i64'
@ -160,7 +168,7 @@ fn transpose{T, k, kh}(r0:*void, x0:*void, w:u64, h:u64, ws:u64, hs:u64) : void
if (hasarch{'AVX2'} and w>=k and h>=k) {
transpose_with_kernel{T, k, kh, call_base, rp, xp, w, h, ws, hs}
} else {
if (h==2 and h==hs) @for (x0 in xp, x1 in xp+ws over i to w) { store{rp, i*2, x0}; store{rp, i*2+1, x1} }
if (h==2 and h==hs) interleave{T}(r0, x0, *void~~(xp+ws), w)
else if (w==2 and w==ws) @for (r0 in rp, r1 in rp+hs over i to h) { r0 = load{xp, i*2}; r1 = load{xp, i*2+1} }
else call_base{rp, xp, w, h}
}
@ -174,3 +182,5 @@ exportT{'simd_transpose', tup{
transpose{i32, 8},
transpose{i64, 4}
}}
exportT{'si_interleave', each{interleave, tup{i8, i16, i32, i64}}}