From 5b17ee44d60f005de513759c668a42bcf4cc6e61 Mon Sep 17 00:00:00 2001 From: Marshall Lochbaum Date: Mon, 27 Mar 2023 22:12:12 -0400 Subject: [PATCH] SIMD kernel for Reorder Axes transposing last two axes --- src/builtins/transpose.c | 80 +++++++++++++++++++++++++++------------- 1 file changed, 55 insertions(+), 25 deletions(-) diff --git a/src/builtins/transpose.c b/src/builtins/transpose.c index 15584b90..57cb8f45 100644 --- a/src/builtins/transpose.c +++ b/src/builtins/transpose.c @@ -32,6 +32,7 @@ #endif #endif +typedef void (*TranspFn)(void*,void*,u64,u64); #if SINGELI #define transposeFns simd_transpose #define DECL_BASE(T) \ @@ -50,7 +51,7 @@ } DECL_BASE(i8) DECL_BASE(i16) DECL_BASE(i32) DECL_BASE(i64) #undef DECL_BASE - static void (*transposeFns[])(void*,void*,u64,u64) = { + static TranspFn transposeFns[] = { transpose_i8, transpose_i16, transpose_i32, transpose_i64 }; #endif @@ -227,36 +228,65 @@ B transp_c2(B t, B w, B x) { u8 xe = TI(x,elType); usz csz = shProd(xsh, ar, xr); - if (csz >= (32*8) >> elWidthLogBits(xe)) { // cell >= 32 bytes - usz ria = csz * shProd(rsh, 0, na); - MAKE_MUT_INIT(rm, ria, xe); MUTG_INIT(rm); - for (usz i=0; i=0 - ur a = na - 1; - usz str = st[a]; - usz l = rsh[a]; - for (usz k=0; k=1 */ \ + ur a = a0; \ + usz str = st[a]; \ + usz l = rsh[a]; \ + for (usz k=0; k= (32*8) >> xlw) { // cell >= 32 bytes + usz ria = csz * shProd(rsh, 0, na); + AXIS_LOOP(MAKE_MUT_INIT(rm, ria, xe); MUTG_INIT(rm); , + na, csz, mut_copyG(rm, i, x, j, csz)); Arr* ra = mut_fp(rm); shSet_alloc(ra, rr, rsh); r = withFill(taga(ra), getFillQ(x)); decG(x); goto ret_decst; } + if ((csz & (csz-1))==0 && csz<=64>>xlw && csz<=8 // CPU-sized cells + && !dup && na>=2 && p[na-1]==ar-2 && p[na-2]==ar-1 // Last two axes transposed + && rsh[na-2]*rsh[na-1] >= (256*8) >> xlw // And large-ish + && xe!=el_B) { + TranspFn tran = transposeFns[CTZ(csz<