Rank-agnostic select_cells
This commit is contained in:
parent
ab15eaeeeb
commit
8766810db8
@ -178,52 +178,58 @@ NOINLINE B leading_axis_arith(FC2 fc2, B w, B x, usz* wsh, usz* xsh, ur mr) { //
|
||||
|
||||
// fast special-case implementations
|
||||
extern void (*const si_select_cells_bit_lt64)(uint64_t*,uint64_t*,uint32_t,uint32_t,uint32_t); // from fold.c (fold.singeli)
|
||||
static NOINLINE B select_cells(usz n, B x, usz cam, usz k, bool leaf) { // n {leaf? <∘⊑; ⊏}⎉¯k x; TODO probably can share some parts with takedrop_highrank and/or call ⊏?
|
||||
static NOINLINE B select_cells(usz ind, B x, usz cam, usz k, bool leaf) { // ind {leaf? <∘⊑; ⊏}⎉¯k x; TODO probably can share some parts with takedrop_highrank and/or call ⊏?
|
||||
ur xr = RNK(x);
|
||||
assert(xr>1 && k<xr);
|
||||
usz* xsh = SH(x);
|
||||
usz csz = shProd(xsh, k+1, xr);
|
||||
usz take = leaf? 1 : csz;
|
||||
usz jump = xsh[k] * csz;
|
||||
assert(cam*jump == IA(x));
|
||||
usz l = xsh[k];
|
||||
assert(0<=ind && ind<l);
|
||||
assert(cam*l*csz == IA(x));
|
||||
Arr* ra;
|
||||
if (take==jump) {
|
||||
usz take = leaf? 1 : csz;
|
||||
if (l==1 && take==csz) {
|
||||
ra = cpyWithShape(incG(x));
|
||||
arr_shErase(ra, 1);
|
||||
} else if (take==1) {
|
||||
} else {
|
||||
u8 xe = TI(x,elType);
|
||||
if (xe==el_B) {
|
||||
u8 ewl= elwBitLog(xe);
|
||||
u8 xl = leaf? ewl : multWidthLog(csz, ewl);
|
||||
usz ria = cam*take;
|
||||
if (xl>=7 || (xl<3 && xl>0)) { // generic case
|
||||
MAKE_MUT_INIT(rm, ria, TI(x,elType)); MUTG_INIT(rm);
|
||||
usz jump = l * csz;
|
||||
usz xi = take*ind;
|
||||
usz ri = 0;
|
||||
for (usz i = 0; i < cam; i++) {
|
||||
mut_copyG(rm, ri, x, xi, take);
|
||||
xi+= jump;
|
||||
ri+= take;
|
||||
}
|
||||
ra = mut_fp(rm);
|
||||
} else if (xe==el_B) {
|
||||
assert(take == 1);
|
||||
SGet(x)
|
||||
HArr_p rp = m_harrUv(cam);
|
||||
for (usz i = 0; i < cam; i++) rp.a[i] = Get(x, i*jump+n);
|
||||
HArr_p rp = m_harrUv(ria);
|
||||
for (usz i = 0; i < cam; i++) rp.a[i] = Get(x, i*l+ind);
|
||||
NOGC_E; ra = (Arr*)rp.c;
|
||||
} else {
|
||||
void* rp = m_tyarrlbp(&ra, elwBitLog(xe), cam, el2t(xe));
|
||||
void* rp = m_tyarrlbp(&ra, ewl, ria, el2t(xe));
|
||||
void* xp = tyany_ptr(x);
|
||||
switch(xe) {
|
||||
case el_bit:
|
||||
switch(xl) {
|
||||
case 0:
|
||||
#if SINGELI
|
||||
if (jump < 64) si_select_cells_bit_lt64(xp, rp, cam, jump, n);
|
||||
if (l < 64) si_select_cells_bit_lt64(xp, rp, cam, l, ind);
|
||||
else
|
||||
#endif
|
||||
for (usz i=0; i<cam; i++) bitp_set(rp, i, bitp_get(xp, i*jump+n));
|
||||
for (usz i=0; i<cam; i++) bitp_set(rp, i, bitp_get(xp, i*l+ind));
|
||||
break;
|
||||
case el_i8: case el_c8: PLAINLOOP for (usz i=0; i<cam; i++) ((u8* )rp)[i] = ((u8* )xp)[i*jump+n]; break;
|
||||
case el_i16: case el_c16: PLAINLOOP for (usz i=0; i<cam; i++) ((u16*)rp)[i] = ((u16*)xp)[i*jump+n]; break;
|
||||
case el_i32: case el_c32: PLAINLOOP for (usz i=0; i<cam; i++) ((u32*)rp)[i] = ((u32*)xp)[i*jump+n]; break;
|
||||
case el_f64: PLAINLOOP for (usz i=0; i<cam; i++) ((f64*)rp)[i] = ((f64*)xp)[i*jump+n]; break;
|
||||
case 3: PLAINLOOP for (usz i=0; i<cam; i++) ((u8* )rp)[i] = ((u8* )xp)[i*l+ind]; break;
|
||||
case 4: PLAINLOOP for (usz i=0; i<cam; i++) ((u16*)rp)[i] = ((u16*)xp)[i*l+ind]; break;
|
||||
case 5: PLAINLOOP for (usz i=0; i<cam; i++) ((u32*)rp)[i] = ((u32*)xp)[i*l+ind]; break;
|
||||
case 6: PLAINLOOP for (usz i=0; i<cam; i++) ((f64*)rp)[i] = ((f64*)xp)[i*l+ind]; break;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
MAKE_MUT_INIT(rm, cam*take, TI(x,elType)); MUTG_INIT(rm);
|
||||
usz xi = take*n;
|
||||
usz ri = 0;
|
||||
for (usz i = 0; i < cam; i++) {
|
||||
mut_copyG(rm, ri, x, xi, take);
|
||||
xi+= jump;
|
||||
ri+= take;
|
||||
}
|
||||
ra = mut_fp(rm);
|
||||
}
|
||||
usz* rsh = arr_shAlloc(ra, leaf? k : xr-1);
|
||||
if (rsh) {
|
||||
@ -461,12 +467,12 @@ B for_cells_c1(B f, u32 xr, u32 cr, u32 k, B x, u32 chr) { // F⎉cr x, with arr
|
||||
if (rtid==n_const) { f=fd->f; goto const_f; }
|
||||
usz *sh = SH(x);
|
||||
if (((rtid==n_fold && cr==1) || rtid==n_insert) && TI(x,elType)!=el_B
|
||||
&& isFun(fd->f) && 1==shProd(sh, k+1, xr) && sh[k] > 0) {
|
||||
&& isFun(fd->f) && sh[k] > 0) {
|
||||
usz m = sh[k];
|
||||
u8 frtid = v(fd->f)->flags-1;
|
||||
if (m==1 || frtid==n_ltack) return select_cells(0 , x, cam, k, false);
|
||||
if ( frtid==n_rtack) return select_cells(m-1, x, cam, k, false);
|
||||
if (isPervasiveDyExt(fd->f)) {
|
||||
if (isPervasiveDyExt(fd->f) && 1==shProd(sh, k+1, xr)) {
|
||||
if (TI(x,elType)==el_bit) {
|
||||
incG(x); // keep shape alive
|
||||
B r = fold_rows_bit(fd, x, shProd(sh, 0, k), m);
|
||||
|
||||
@ -73,11 +73,15 @@ SHOULD_INLINE void arr_check_size(u64 sz, u8 type, usz ia) {
|
||||
#endif
|
||||
}
|
||||
// Log of width in bits: max of 7, and also return 7 if not power of 2
|
||||
SHOULD_INLINE u8 cellWidthLog(B x) {
|
||||
SHOULD_INLINE u8 multWidthLog(usz n, u8 lw) { // Of n elements, 1<<lw bit
|
||||
if (n & (n-1)) return 7; // Not power of 2
|
||||
return lw + CTZ(n | 128>>lw); // Max of 7; also handle n==0
|
||||
}
|
||||
SHOULD_INLINE u8 kCellWidthLog(B x, ur k) {
|
||||
assert(isArr(x) && RNK(x)>=1);
|
||||
u8 lw = arrTypeBitsLog(TY(x));
|
||||
if (LIKELY(RNK(x)==1)) return lw;
|
||||
usz csz = arr_csz(x);
|
||||
if (csz & (csz-1)) return 7; // Not power of 2
|
||||
return lw + CTZ(csz | 128>>lw); // Max of 7; also handle csz==0
|
||||
ur xr = RNK(x);
|
||||
if (LIKELY(xr <= k)) return lw;
|
||||
return multWidthLog(shProd(SH(x), k, xr), lw);
|
||||
}
|
||||
SHOULD_INLINE u8 cellWidthLog(B x) { return kCellWidthLog(x, 1); }
|
||||
|
||||
Loading…
Reference in New Issue
Block a user