// Transpose and Reorder Axes (⍉) // Transpose // One length-2 axis: dedicated code // Boolean: pdep for height 2; pext for width 2 // SHOULD use a generic implementation if BMI2 not present // SHOULD optimize other short lengths with pdep/pext and shuffles // Boolean 𝕩: convert to integer // SHOULD have bit matrix transpose kernel // CPU sizes: native or SIMD code // Large SIMD kernels used when they fit, overlapping for odd sizes // i8: 16Γ—16; i16: 16Γ—8; i32: 8Γ—8; f64: 4Γ—4 // COULD use half-width or smaller kernels to improve odd sizes // Scalar transpose or loop used for overhang of 1 // Reorder Axes: generate indices and select with +⌜ and ⊏ // Transpose inverse ⍉⁼𝕩: data movement of ⍉ with different shape logic // COULD implement fast β‰βŸn // SHOULD convert ⍉ with rank to a Reorder Axes call #include "../core.h" #include "../utils/each.h" #include "../utils/talloc.h" #include "../builtins.h" #include "../utils/calls.h" #ifdef __BMI2__ #include #if USE_VALGRIND #define _pdep_u64 vg_pdep_u64 #endif #endif #define TRANSPOSE_LOOP( DST, SRC, W, H) PLAINLOOP for(usz y=0,xi=0;y< H;y++) NOVECTORIZE for(usz x=0;x< W;x++) DST[x*H+y] = SRC[xi++] #define TRANSPOSE_BLOCK(DST, SRC, BW, BH, W, H) PLAINLOOP for(usz y=0 ;ya; for (usz i=0; ia, xsh+1, xr-1); sh->a[xr-1] = h; arr_shReplace(r, xr, sh); return taga(r); } usz w = xsh[1] * shProd(xsh, 2, xr); Arr* r = transpose_noshape(&x, ia, w, h); usz* rsh = arr_shAlloc(r, xr); if (xr==2) rsh[0] = w; else shcpy(rsh, SH(x)+1, xr-1); rsh[xr-1] = h; decG(x); return taga(r); } B mul_c2(B,B,B); B ud_c1(B,B); B tbl_c2(Md1D*,B,B); B select_c2(B,B,B); B transp_c2(B t, B w, B x) { usz wia=1; if (isArr(w)) { if (RNK(w)>1) thrM("⍉: 𝕨 must have rank at most 1"); wia = IA(w); if (wia==0) { decG(w); return isArr(x)? x : m_atomUnit(x); } } ur xr; if (isAtm(x) || (xr=RNK(x))=xr) thrF("⍉: Axis %s does not exist (%i≑=𝕩)", a, xr); if (a==xr-1) { TFREE(p); return C1(transp, x); } p[0] = a; } else { SGetU(w) for (usz i=0; i=xr) thrF("⍉: Axis %s does not exist (%i≑=𝕩)", a, xr); p[i] = a; } decG(w); } // compute shape for the given axes usz* xsh = SH(x); TALLOC(usz, rsh, xr); usz dup = 0, max = 0; usz no_sh = -(usz)1; for (usz j=0; jmax? j : max; if (xl= rr) thrF("⍉: Skipped result axis"); if (wiaa, rsh, rr); arr_shSetU(ra, rr, sh); } decG(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) { r = x; goto ret; } // Add up stride for each axis TALLOC(u64, st, rr); for (usz j=0; j1) { zsh = m_shArr(zr); zsh->a[0] = c; shcpy(zsh->a+1, xsh+ar, xr-ar); } Arr* z = TI(x,slice)(x, 0, IA(x)); if (zr>1) arr_shSetU(z, zr, zsh); else arr_shVec(z); x = taga(z); } // (+⌜´stΓ—βŸœβ†•Β¨rsh)⊏β₯Šπ•© B ind = bi_N; for (ur k=ar-dup; k--; ) { B v = C2(mul, m_f64(st[k]), C1(ud, m_f64(rsh[k]))); if (q_N(ind)) ind = v; else ind = M1C2(tbl, add, v, ind); } TFREE(st); r = C2(select, ind, x); ret:; TFREE(rsh); TFREE(p); return r; } B transp_im(B t, B x) { if (isAtm(x)) thrM("⍉⁼: 𝕩 must not be an atom"); ur xr = RNK(x); if (xr<=1) return x; usz ia = IA(x); usz* xsh = SH(x); usz w = xsh[xr-1]; if (ia==0 || w==1 || w==ia /*h==1*/) { Arr* r = cpyWithShape(x); ShArr* sh = m_shArr(xr); sh->a[0] = w; shcpy(sh->a+1, xsh, xr-1); arr_shReplace(r, xr, sh); return taga(r); } usz h = xsh[0] * shProd(xsh, 1, xr-1); Arr* r = transpose_noshape(&x, ia, w, h); usz* rsh = arr_shAlloc(r, xr); rsh[0] = w; if (xr==2) rsh[1] = h; else shcpy(rsh+1, SH(x), xr-1); decG(x); return taga(r); } B transp_uc1(B t, B o, B x) { return transp_im(m_f64(0), c1(o, transp_c1(t, x))); } void transp_init(void) { c(BFn,bi_transp)->uc1 = transp_uc1; c(BFn,bi_transp)->im = transp_im; }