From d492cd08650b64760ddc195fd051b941bce3c898 Mon Sep 17 00:00:00 2001 From: Marshall Lochbaum Date: Wed, 29 Mar 2023 11:31:52 -0400 Subject: [PATCH] And strided height, so all orders are fast if axes are long --- src/builtins/transpose.c | 40 ++++++++++++++++++++++++++++++---------- 1 file changed, 30 insertions(+), 10 deletions(-) diff --git a/src/builtins/transpose.c b/src/builtins/transpose.c index da9eaba2..5c75aadb 100644 --- a/src/builtins/transpose.c +++ b/src/builtins/transpose.c @@ -257,30 +257,50 @@ B transp_c2(B t, B w, B x) { r = withFill(taga(ra), getFillQ(x)); decG(x); goto ret; } + #undef AXIS_LOOP if ((csz & (csz-1))==0 && csz<=64>>xlw && csz<=8 // CPU-sized cells - && na>=2 && st[na-2]==1 // Last axis ends up second-to-last - && rsh[na-2]*rsh[na-1] >= (256*8) >> xlw // Large-ish axes - && xe!=el_B) { + && xe!=el_B && na>=2) { + // If some result axis has stride 1 (guaranteed if dup==0), then it + // corresponds to the last argument axis and we have a strided + // transpose swapping that with the last result axis + usz rai = na-1; + usz xai=rai; while (st[--xai]!=1) if (xai==0) goto skip_2d; + if (rsh[xai]*rsh[rai] < (256*8) >> xlw) goto skip_2d; TranspFn tran = transposeFns[CTZ(csz<