From 597e25af4d77bca8937f373cf750e42a1bed2975 Mon Sep 17 00:00:00 2001 From: Marshall Lochbaum Date: Wed, 29 Mar 2023 20:08:36 -0400 Subject: [PATCH] Full simplification pass for Reorder Axes --- src/builtins/transpose.c | 51 +++++++++++++++++++++++++--------------- 1 file changed, 32 insertions(+), 19 deletions(-) diff --git a/src/builtins/transpose.c b/src/builtins/transpose.c index 5c75aadb..c3f3d533 100644 --- a/src/builtins/transpose.c +++ b/src/builtins/transpose.c @@ -172,19 +172,23 @@ B transp_c2(B t, B w, B x) { decG(w); } + B r; + // Compute shape for the given axes usz* xsh = SH(x); usz *rsh = (usz*)(p + xr); // Length xr - usz dup = 0, max = 0; + usz dup = 0, max = 0, id = 0; usz no_sh = -(usz)1; for (usz j=0; jmax? j : max; if (xla, rsh, rr); } - B r; - // Empty result if (IA(x) == 0) { Arr* ra = m_fillarrpEmpty(getFillQ(x)); @@ -211,23 +213,33 @@ B transp_c2(B t, B w, B x) { r = taga(ra); goto ret; } - // Number of axes that move - ur ar = max+1+dup; - if (!dup) while (ar>1 && p[ar-1]==ar-1) ar--; // Unmoved trailing - if (ar <= 1) { if (rr>1) ptr_dec(sh); r = x; goto ret; } - ur na = ar - dup; // Add up stride for each axis - usz* st = rsh + xr; // Length ar - for (usz j=0; j=1 */ \ @@ -280,8 +292,9 @@ B transp_c2(B t, B w, B x) { usz i_skip = (w-1)*hs*csz; usz end = rf*csz - i_skip; ur a0 = na - 1; - for (usz i=0; i>= 3-xlw; + else if (xlw>3) for (usz i=0; ia[0] = c; + zsh->a[0] = IA(x)/csz; zsh->a[1] = csz; Arr* z = TI(x,slice)(x, 0, IA(x)); arr_shSetU(z, 2, zsh); @@ -312,7 +325,7 @@ B transp_c2(B t, B w, B x) { // (+⌜´st×⟜↕¨rsh)⊏⥊𝕩 B ind = bi_N; for (ur k=na; k--; ) { - B v = C2(mul, m_f64(st[k]), C1(ud, m_f64(rsh[k]))); + B v = C2(mul, m_f64(st[k]/csz), C1(ud, m_f64(rsh[k]))); if (q_N(ind)) ind = v; else ind = M1C2(tbl, add, v, ind); }