Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
77 changes: 57 additions & 20 deletions src/lang/eval.c
Original file line number Diff line number Diff line change
Expand Up @@ -1432,6 +1432,45 @@ ray_t* call_fn2(ray_t* fn, ray_t* a, ray_t* b) {
* Sorting builtins
* ══════════════════════════════════════════ */

/* Parallel fixed-width index gather for gather_by_idx: value memcpy loops
* split across the worker pool. Null-bit propagation stays serial (bit
* writes within a shared word would race across range boundaries).
*
* gather_by_idx is a LEAF utility with 35+ call sites; unlike op-level
* dispatch points it cannot know its execution context, so the parallel
* path additionally requires no dispatch in flight (ray_parallel_flag) —
* ray_pool_dispatch is single-producer, and a nested dispatch from a
* worker would corrupt the task ring. */
extern _Atomic(uint32_t) ray_parallel_flag; /* core/platform.h */

static inline bool gbi_par_ok(ray_pool_t* p, int64_t n) {
return p && p->n_workers > 0 && n >= RAY_PARALLEL_THRESHOLD &&
atomic_load_explicit(&ray_parallel_flag, memory_order_relaxed) == 0;
}

typedef struct {
const char* src;
char* dst;
const int64_t* idx;
uint8_t esz;
} gbi_ctx_t;

static void gbi_fill_fn(void* ctxp, uint32_t wid, int64_t start, int64_t end) {
(void)wid;
gbi_ctx_t* c = (gbi_ctx_t*)ctxp;
const char* src = c->src;
char* dst = c->dst;
const int64_t* idx = c->idx;
switch (c->esz) {
case 8: for (int64_t i = start; i < end; i++) memcpy(dst + i*8, src + idx[i]*8, 8); break;
case 4: for (int64_t i = start; i < end; i++) memcpy(dst + i*4, src + idx[i]*4, 4); break;
case 2: for (int64_t i = start; i < end; i++) memcpy(dst + i*2, src + idx[i]*2, 2); break;
case 1: for (int64_t i = start; i < end; i++) dst[i] = src[idx[i]]; break;
case 16: for (int64_t i = start; i < end; i++) memcpy(dst + i*16, src + idx[i]*16, 16); break;
default: for (int64_t i = start; i < end; i++) memcpy(dst + i*c->esz, src + idx[i]*c->esz, c->esz); break;
}
}

/* Reorder vector elements by an index array */
ray_t* gather_by_idx(ray_t* vec, int64_t* idx, int64_t n) {
int8_t type = vec->type;
Expand Down Expand Up @@ -1465,15 +1504,14 @@ ray_t* gather_by_idx(ray_t* vec, int64_t* idx, int64_t n) {
if (RAY_IS_ERR(result)) return result;
result->len = n;
uint8_t esz = (uint8_t)RAY_SYM_ELEM(w);
char* src = (char*)ray_data(vec);
char* dst = (char*)ray_data(result);
switch (esz) {
case 8: for (int64_t i = 0; i < n; i++) memcpy(dst + i*8, src + idx[i]*8, 8); break;
case 4: for (int64_t i = 0; i < n; i++) memcpy(dst + i*4, src + idx[i]*4, 4); break;
case 2: for (int64_t i = 0; i < n; i++) memcpy(dst + i*2, src + idx[i]*2, 2); break;
case 1: for (int64_t i = 0; i < n; i++) dst[i] = src[idx[i]]; break;
default: for (int64_t i = 0; i < n; i++) memcpy(dst + i*esz, src + idx[i]*esz, esz); break;
}
gbi_ctx_t gc = { .src = (const char*)ray_data(vec),
.dst = (char*)ray_data(result),
.idx = idx, .esz = esz };
ray_pool_t* gpool = ray_pool_get();
if (gbi_par_ok(gpool, n))
ray_pool_dispatch(gpool, gbi_fill_fn, &gc, n);
else
gbi_fill_fn(&gc, 0, 0, n);
if (has_nulls) {
for (int64_t i = 0; i < n; i++)
if (ray_vec_is_null(vec, idx[i]))
Expand Down Expand Up @@ -1506,17 +1544,16 @@ ray_t* gather_by_idx(ray_t* vec, int64_t* idx, int64_t n) {
if (RAY_IS_ERR(result)) return result;
result->len = n;
uint8_t esz = ray_type_sizes[type];
char* src = (char*)ray_data(vec);
char* dst = (char*)ray_data(result);
/* Typed gather — compiler constant esz enables vectorization, alias-safe */
switch (esz) {
case 8: for (int64_t i = 0; i < n; i++) memcpy(dst + i*8, src + idx[i]*8, 8); break;
case 4: for (int64_t i = 0; i < n; i++) memcpy(dst + i*4, src + idx[i]*4, 4); break;
case 2: for (int64_t i = 0; i < n; i++) memcpy(dst + i*2, src + idx[i]*2, 2); break;
case 1: for (int64_t i = 0; i < n; i++) dst[i] = src[idx[i]]; break;
default: for (int64_t i = 0; i < n; i++) memcpy(dst + i*esz, src + idx[i]*esz, esz); break;
case 16: for (int64_t i = 0; i < n; i++) memcpy(dst + i*16, src + idx[i]*16, 16); break;
}
/* Typed gather — compiler constant esz enables vectorization, alias-safe;
* pool-parallel over disjoint output ranges for large gathers. */
gbi_ctx_t gc = { .src = (const char*)ray_data(vec),
.dst = (char*)ray_data(result),
.idx = idx, .esz = esz };
ray_pool_t* gpool = ray_pool_get();
if (gbi_par_ok(gpool, n))
ray_pool_dispatch(gpool, gbi_fill_fn, &gc, n);
else
gbi_fill_fn(&gc, 0, 0, n);

/* Propagate null bitmap */
if (has_nulls) {
Expand Down
88 changes: 88 additions & 0 deletions src/ops/builtins.c
Original file line number Diff line number Diff line change
Expand Up @@ -2100,11 +2100,99 @@ ray_t* ray_nil_fn(ray_t* x) {
}

/* (where bool-vec) -> indices of true values */

/* Parallel where: 3-phase chunk compaction (parallel per-chunk popcount →
* tiny serial prefix → parallel per-chunk emit at disjoint offsets). */
#define WHERE_CHUNK_ROWS 65536

typedef struct {
const uint8_t* data;
int64_t n;
int64_t chunk_rows;
int64_t* chunk_counts; /* popcount, then exclusive-prefix offsets */
int64_t* out; /* emit-phase output (NULL during count) */
} where_ctx_t;

/* Dispatched via ray_pool_dispatch_n — each task is ONE chunk ([i,i+1)). */
static void where_count_fn(void* ctxp, uint32_t wid, int64_t start, int64_t end) {
(void)wid;
where_ctx_t* c = (where_ctx_t*)ctxp;
for (int64_t ch = start; ch < end; ch++) {
int64_t lo = ch * c->chunk_rows;
int64_t hi = lo + c->chunk_rows;
if (hi > c->n) hi = c->n;
int64_t cnt = 0;
for (int64_t i = lo; i < hi; i++) cnt += (c->data[i] != 0);
c->chunk_counts[ch] = cnt;
}
}

/* Word-at-a-time zero skip: sparse masks — the common WHERE shape —
* run at memory speed, falling to per-byte only inside a live word. */
static void where_emit_fn(void* ctxp, uint32_t wid, int64_t start, int64_t end) {
(void)wid;
where_ctx_t* c = (where_ctx_t*)ctxp;
for (int64_t ch = start; ch < end; ch++) {
int64_t lo = ch * c->chunk_rows;
int64_t hi = lo + c->chunk_rows;
if (hi > c->n) hi = c->n;
int64_t j = c->chunk_counts[ch];
int64_t i = lo;
int64_t h8 = lo + ((hi - lo) & ~7LL);
for (; i < h8; i += 8) {
uint64_t w;
memcpy(&w, c->data + i, 8);
if (!w) continue;
for (int64_t k = 0; k < 8; k++)
if (c->data[i + k]) c->out[j++] = i + k;
}
for (; i < hi; i++)
if (c->data[i]) c->out[j++] = i;
}
}

ray_t* ray_where_fn(ray_t* x) {
if (!ray_is_vec(x) || x->type != RAY_BOOL)
return ray_error("type", "where: argument must be a b8 vector, got %s", ray_type_name(x->type));
const uint8_t* data = (const uint8_t*)ray_data(x);
int64_t n = x->len;

ray_pool_t* pool = ray_pool_get();
if (pool && pool->n_workers > 0 && n >= RAY_PARALLEL_THRESHOLD) {
/* One task per chunk (ray_pool_dispatch_n). Chunk count is capped
* at 1024 — the pool's INITIAL ring capacity — so dispatch_n never
* needs to grow the ring (its clamp-on-grow-failure silently DROPS
* tasks, which here would mean uninitialized prefix entries and
* out-of-bounds emit offsets). 1024 tasks is ample granularity. */
int64_t chunk_rows = WHERE_CHUNK_ROWS;
while ((n + chunk_rows - 1) / chunk_rows > 1024) chunk_rows *= 2;
int64_t n_chunks = (n + chunk_rows - 1) / chunk_rows;
/* ops/internal.h's scratch_alloc idiom is unavailable here (that
* header clashes with builtins.c-local helpers); ray_alloc +
* ray_release below is its exact expansion. */
ray_t* cc_hdr = ray_alloc((size_t)n_chunks * sizeof(int64_t));
int64_t* chunk_counts = cc_hdr ? (int64_t*)ray_data(cc_hdr) : NULL;
if (chunk_counts) {
where_ctx_t wc = { .data = data, .n = n, .chunk_rows = chunk_rows,
.chunk_counts = chunk_counts, .out = NULL };
ray_pool_dispatch_n(pool, where_count_fn, &wc, (uint32_t)n_chunks);
int64_t cum = 0;
for (int64_t ch = 0; ch < n_chunks; ch++) {
int64_t v = chunk_counts[ch];
chunk_counts[ch] = cum;
cum += v;
}
ray_t* result = ray_vec_new(RAY_I64, cum);
if (RAY_IS_ERR(result)) { ray_release(cc_hdr); return result; }
wc.out = (int64_t*)ray_data(result);
ray_pool_dispatch_n(pool, where_emit_fn, &wc, (uint32_t)n_chunks);
result->len = cum;
ray_release(cc_hdr);
return result;
}
/* scratch OOM → sequential fallback below */
}

/* Count pass: `cnt += byte != 0` vectorizes; the branchy form costs
* a mispredict-prone compare per element. */
int64_t cnt = 0;
Expand Down
Loading
Loading