From 27b7c0746c5c399d2c9b32ff3d12a779d78519bb Mon Sep 17 00:00:00 2001 From: dzaima Date: Thu, 12 Sep 2024 06:09:17 +0300 Subject: [PATCH] full interleave_cells needed to make all rank-1 inputs work on non-Singeli builds, but beneficial in general anyway --- src/builtins/cells.c | 15 ++------ src/builtins/transpose.c | 62 +++++++++++++++++-------------- src/singeli/src/transpose.singeli | 2 +- 3 files changed, 39 insertions(+), 40 deletions(-) diff --git a/src/builtins/cells.c b/src/builtins/cells.c index f3920454..ddb081df 100644 --- a/src/builtins/cells.c +++ b/src/builtins/cells.c @@ -23,7 +23,7 @@ B scan_rows_bit(u8, B x, usz m); B takedrop_highrank(bool take, B w, B x); B rotate_highrank(bool inv, B w, B x); -B try_interleave_cells(B w, B x, ur xr, ur xk, usz* xsh); // from transpose.c +B interleave_cells(B w, B x, ur xr, ur xk, usz* xsh); // from transpose.c // from select.c: B select_rows_B(B x, ux csz, ux cam, B inds); @@ -700,10 +700,7 @@ NOINLINE B for_cells_AS(B f, B w, B x, ur wcr, ur wr, u32 chr) { // F⟜x⎉wcr case n_rtack: return const_cells(w, wk, wsh, x, chr); case n_couple: if (RNK(w)==1) { x = taga(arr_shVec(reshape_one(IA(w), x))); - B r = try_interleave_cells(w, x, 1, 1, wsh); - assert(!q_N(r)); - decG(w); decG(x); - return r; + return interleave_cells(w, x, 1, 1, wsh); } break; } if (IA(w)!=0 && isPervasiveDy(f)) { @@ -754,10 +751,7 @@ NOINLINE B for_cells_SA(B f, B w, B x, ur xcr, ur xr, u32 chr) { // w⊸F⎉xcr break; case n_couple: if (RNK(x)==1) { w = taga(arr_shVec(reshape_one(IA(x), w))); - B r = try_interleave_cells(w, x, 1, 1, xsh); - assert(!q_N(r)); - decG(w); decG(x); - return r; + return interleave_cells(w, x, 1, 1, xsh); } break; case n_pick: if (isF64(w) && xcr==1 && TI(x,arrD1)) { usz l = xsh[xk]; @@ -861,8 +855,7 @@ NOINLINE B for_cells_AA(B f, B w, B x, ur wcr, ur xcr, u32 chr) { // w F⎉wcr 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; } + return interleave_cells(w, x, xr, xk, xsh); } } if (isPervasiveDy(f)) { diff --git a/src/builtins/transpose.c b/src/builtins/transpose.c index d01e2da2..2457e8f3 100644 --- a/src/builtins/transpose.c +++ b/src/builtins/transpose.c @@ -99,7 +99,7 @@ B toBPtrAny(B x) { } NOINLINE -B toElTypeArr(u8 re, B x) { // consumes x; returns an array with the given element type (re==el_B guarantees TO_BPTR working) +B toElTypeArr(u8 re, B x) { // consumes; returns an array with the given element type (re==el_B guarantees TO_BPTR working) switch (re) { default: UD; case el_bit: return toBitAny(x); case el_i8: return toI8Any(x); @@ -114,61 +114,67 @@ B toElTypeArr(u8 re, B x) { // consumes x; returns an array with the given eleme } -// Interleave arrays, 𝕨≍⎉(-xk)𝕩. Doesn't consume. -// Assumes w and x have same shape. -// Return bi_N if there isn't fast code. Guaranteed to succeed on rank 1 w & x. -B try_interleave_cells(B w, B x, ur xr, ur xk, usz* xsh) { +// interleave arrays, 𝕨≍⎉(-xk)𝕩 +B interleave_cells(B w, B x, ur xr, ur xk, usz* xsh) { // consumes w,x assert(RNK(w)==xr && xr>=1); u8 we = TI(w,elType); u8 xe = TI(x,elType); - usz csz = shProd(xsh, xk, xr); - if (csz & (csz-1)) return bi_N; // Not power of 2 u8 re = we==xe? we : el_or(we, xe); - if (0) { to_equal_types:; + if (0) { goto to_equal_types; to_equal_types:; // delay doing this until it's known that there will be code that can utilize it - incG(w); B w2 = re==we? w : toElTypeArr(re, w); - incG(x); B x2 = re==xe? x : toElTypeArr(re, x); - B r = try_interleave_cells(w2, x2, xr, xk, SH(x)); - assert(!q_N(r)); - decG(w2); decG(x2); - return r; + if (re!=we) w = toElTypeArr(re, w); + if (re!=xe) x = toElTypeArr(re, x); + return interleave_cells(w, x, xr, xk, SH(x)); } + Arr *r; u8 xlw = elwBitLog(re); usz n = shProd(xsh, 0, xk); + usz csz = shProd(xsh, xk, xr); usz ia = 2*n*csz; - Arr *r; - if (csz==1 && xlw==0) { // we & xe are trivially el_bit + + if (csz & (csz-1)) { + goto generic; + } else if (csz==1 && xlw==0) { // we & xe are trivially el_bit u64* rp; r=m_bitarrp(&rp, ia); interleave_bits(rp, bitany_ptr(w), bitany_ptr(x), ia); - } #if SINGELI - else if (csz==1 && re==el_B) { + } else if (csz==1 && re==el_B) { if (we!=xe) goto to_equal_types; - incG(w); incG(x); - B* wp = TO_BPTR(w); B* xp = TO_BPTR(x); + B* wp = TO_BPTR(w); B* xp = TO_BPTR_RUN(x, xsh = SH(x)); + HArr_p p = m_harrUv(ia); // Debug build complains with harrUp - si_interleave[3](p.a, wp, xp, n); - decG(w); decG(x); + interleave_fns[3](p.a, wp, xp, n); for (usz i=0; i>xlw && csz<=8) { // Require CPU-sized cells if (we!=xe) goto to_equal_types; assert(re!=el_B); void* rv; if (xlw==0) { u64* rp; r = m_bitarrp(&rp, ia); rv=rp; } else rv = m_tyarrp(&r,elWidth(re),ia,el2t(re)); - si_interleave[CTZ(csz<