From eb34213f0fc343bafdaaa4872f6e3579eded2a16 Mon Sep 17 00:00:00 2001 From: Marshall Lochbaum Date: Sat, 30 Dec 2023 20:27:05 -0500 Subject: [PATCH] =?UTF-8?q?Fast=20=E2=89=8D=CB=98=20for=20matching-type=20?= =?UTF-8?q?type-sized=20cells?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/builtins/cells.c | 5 ++++ src/builtins/transpose.c | 42 ++++++++++++++++++++++++++++--- src/singeli/src/transpose.singeli | 12 ++++++++- 3 files changed, 55 insertions(+), 4 deletions(-) diff --git a/src/builtins/cells.c b/src/builtins/cells.c index 1ff71961..354fb8f5 100644 --- a/src/builtins/cells.c +++ b/src/builtins/cells.c @@ -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; diff --git a/src/builtins/transpose.c b/src/builtins/transpose.c index 331b0127..a071d407 100644 --- a/src/builtins/transpose.c +++ b/src/builtins/transpose.c @@ -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=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<=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<a; - for (usz i=0; ia, ia); mm_free((Value*)x1o); } else if (w==2) { u64* xp = bitarr_ptr(x); diff --git a/src/singeli/src/transpose.singeli b/src/singeli/src/transpose.singeli index 08dbc977..4fe11838 100644 --- a/src/singeli/src/transpose.singeli +++ b/src/singeli/src/transpose.singeli @@ -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}}}