Fast ≍˘ for matching-type type-sized cells
This commit is contained in:
parent
9931c1756c
commit
eb34213f0f
@ -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;
|
||||
|
||||
@ -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);
|
||||
|
||||
@ -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}}}
|
||||
|
||||
Loading…
Reference in New Issue
Block a user