Add Mojo SIMD kernels and a mojo-kernel-overrides extension to replace ops with optimized kernels - #34
Merged
Conversation
…e ops with optimized kernels There are two ways to run Mojo SIMD kernels inside DuckDB, sharing one kernel source (duckdb/kernels/simd.mojo): - Named UDFs, part of the duckdb package: duckdb.kernels.register_simd_math(conn) registers mojo_sqrt/sin/cos/ln/exp/log10 as scalar functions. The kernels and helper ship in the precompiled duckdb-mojo package; nothing extra to build or LOAD. - Built-in overrides, a separate extension (packages/mojo-kernel-overrides): a self-contained CPP-ABI .so that rewrites the built-in sqrt/sin/cos/ln/exp/log10 and sum/avg/min/max in place via catalog mutation, with stock fallback. The Mojo SIMD kernels are emitted as an object and linked straight in. Activated with LOAD, or for embedders via the exported register_mojo_overrides(duckdb_connection). Build/benchmark with pixi run overrides-build / overrides-bench.
Drive the override extension through DuckDB's own benchmark_runner, stock vs overrides, via a single DUCKDB_BENCH_EXTENSION env-var toggle. - runner_load_extension.patch: standalone ~13-line interpreted_benchmark.cpp hook (allow unsigned + LOAD $DUCKDB_BENCH_EXTENSION at init). Needed because the runner has no --load-extension flag and allow_unsigned_extensions is startup-only. Independent of the SIMD source patch; works on stock libduckdb. - build_runner.sh: reuses the existing .duckdb-src checkout/build (no second DuckDB clone or compile), applies the hook only if absent (idempotent), copies the micro benchmarks in, builds benchmark_runner (+ tpch). - run_runner_compare.sh: runs the suite twice toggling the extension and prints a stock / mojo / speedup table. - micro/mojo_simd/*.benchmark: the math/aggregate micro group targeting the overridden ops (was untracked in .duckdb-src). - pixi tasks overrides-bench-runner-build / overrides-bench-runner; README + AGENTS docs. Examples use benchmark/tpch/sf1/.* (the broad benchmark/tpch/.* sweeps in CSV/delete benchmarks that abort the runner even on stock DuckDB).
…overrides Override the ungrouped HUGEINT and DECIMAL(19-38) sum/avg paths with a multi-accumulator Mojo int128 reduce (reduce_sum_i128), replacing DuckDB's overflow-checked per-element Hugeint::Add. falls back to stock on int128 overflow and for non-FLAT/null/grouped/empty input.
A CPP-ABI DuckDB extension that transparently offloads supported SQL to the GPU, with compute kernels in Mojo. An OptimizerExtension matches eligible plan shapes and rewrites them to custom GPU PhysicalOperators (translate-or-fallback), also exposed as gpu_* table functions. Results are decimal-exact vs stock DuckDB. Coverage: array_cosine_distance and TPC-H Q1/Q3/Q5/Q6/Q14 (filter, grouped aggregation, and 1-/3-/6-way joins). Mechanics: - Mojo kernels over a flat C-ABI; a single DeviceContext shared across kernels. - Decimal exactness via int64 arithmetic + int128 host-side cross-block reduction. - FK joins: host-built hash table, GPU probe over the fact side. - High-cardinality group-by: sort on the key + GPU segmented reduction (Apple GPUs lack 64-bit atomics). - Columns are pinned into resident DeviceBuffers (pre-sized pinned host staging), cached per table so the win is warm/repeated workloads; a single cold query still pays a one-time CPU materialization + upload. Built as a companion kernel shared-lib (mojo --emit shared-lib) linked by the C++ extension; version-locked to DuckDB v1.5.3; macOS/Apple-GPU only; unsigned.
…or-driven engine Generalize the GPU-offload extension from 6 hand-written, version-locked per-query matchers (Q1/Q3/Q5/Q6/Q14 + cosine) to a single descriptor-driven engine. A matched aggregate now flows: DuckDB plan -> C++ SerializeMatchedPlan (flattens the subtree into a flat int64 RawPlan tape) -> Mojo build_descriptor (the planner brain) -> generic LogicalGpuAgg/PhysicalGpuAgg source op -> Mojo execution shuttle (materialize_sql / feed_column / pin_finalize) -> generic GPU kernels (expr_vm + segreduce) -> decimal-exact result. C++ is reduced to DuckDB-ABI glue; all planner and execution logic lives in Mojo. Generic kernels are atomics-free (warp.sum + host int128 cross-block reduce): - expr_vm: postfix integer VM; OP_LOAD_DIM (on-GPU dense FK gather, since TPC-H FKs are dense surrogate PKs), OP_EQ (correlated dim<->dim joins), OP_SELECT. - segreduce: UNGROUPED / DENSE_GROUP / SORT_SEGREDUCE, N metrics. Joins use dense dim-array gather; transitive dim->dim folds run on the host (Q3 customer->orders; Q5 customer->orders + correlated c_nationkey=s_nationkey). All five TPC-H classes route through the engine by default (EXPLAIN shows GPU_AGG) and match stock at sf1, including under repeated (warm) execution. GPU_OP_GENERIC=off|none disables GPU aggregate offload (falls back to stock CPU); a value like "q3 q5" restricts to named kinds. The bespoke MatchQ*/LogicalQ*/PhysicalQ*/EnsureQ*Pinned/g_q*_pins, the gpu_q* table functions, and the legacy per-query mojo_q* kernels are removed (net -5246 lines across gpu_operator.cpp and gpu_kernels.mojo). Cosine (GPU_COSINE) is unchanged. Warm pin-cache: segreduce is split into segreduce_upload (-> SegResident, owns the resident DeviceBuffers) + segreduce_run (launches on resident buffers). A process-global cache keyed by a signature that includes the filter constants owns the resident buffers + result-assembly metadata, so a repeated identical query reruns the kernel on cached device memory with no re-materialize and no re-upload (pin_begin reports WARM for all classes; C++ skips feeding; a shared assemble keeps cold/warm paths identical). At TPC-H sf1 the warm path is correct for all five and beats stock on Q1/Q5/Q14 (~1.6-2.0x), ~parity on Q6, and below stock on Q3 (sort-segreduce path). build.sh is Linux-portable (.so + $ORIGIN vs .dylib + @loader_path); the atomics-free kernels are portable to NVIDIA with no atomics branch. Only the macOS/Apple path is hardware-validated so far. New: RawPlan contract (RAW_PLAN_CONTRACT.md, raw_plan.h / raw_plan_tags.mojo in lockstep), descriptor.mojo, expr_vm.mojo, segreduce.mojo, gpu_platform.mojo, and bench tests (raw_plan_roundtrip + per-class *_shuttle + generic_kernel) with gpu-op-*shuttle pixi tasks.
Extend the cosine/vector-search path from "compute all N distances" to a real GPU kNN that returns only the k nearest, and add an fp16-resident variant that roughly halves the bandwidth-bound scan. - Exact GPU top-k (mojo_gpu_pin_query_topk / _batch): per-block k-best in shared memory (insertion-sorted by (dist, rowid)) -> host merge of nblocks*k candidates, so only ~k results cross PCIe instead of N. Deterministic tie-break matches a CPU stable top-k. Scratch buffers cached in the resident pin state (allocated once, single sync per query). Uses std.memory.stack_allocation for shared memory (no `layout` package dependency, which isn't installed in every Mojo env and blocked the Linux build). - fp16 resident path (mojo_gpu_pin_f16 / _query_topk_f16 / _free_f16 + cosine_kernel_warp_f16): storage float16, dot/norm accumulate in float32, query stays fp32. RTX 4090: 1M x 384 2.08->1.27ms (1.64x), 1M x 768 3.68->2.08ms (1.77x), recall@10 0.994/0.995 with zero genuine misses (sub-1.0 is boundary FP-tie reordering, not lost neighbors); halves VRAM. - gpu_cosine_topk(table, column, query, k [, precision]) table function returns the k nearest (rowid, dist). Defaults to fp16 (the effectively-exact, faster, half-memory path); precision='fp32'/'exact' selects the exact path. fp16 pins cached separately (g_pins_f16, "table.column#f16"). Validated vs DuckDB's exact ORDER BY ... LIMIT k. - Benches: cosine_topk_test (bit-exact vs CPU), cosine_f16_test (recall + latency fp16 vs fp32), cosine_topk_latency. Validated on Apple (Metal) and NVIDIA (RTX 4090). Head-to-head at sf-scale (clustered unit-norm embeddings): GPU exact kNN beats CPU exact brute-force ~40-57x warm; vs vss/HNSW it is competitive-to-faster single-query up to ~1M rows (fp16) while being effectively exact, zero-build, and half the memory.
…ile)
Add a batched exact-cosine top-k that reads the resident N×K embedding matrix
once per query-tile (not once per query): each warp loads an embedding row's
lane slice into registers once, then loops over all queries in the tile reusing
it, so the dominant matrix-read bandwidth is amortized across the batch. Per-query
top-k is merged on the GPU (second-stage kernel), so only M*k results cross PCIe.
- mojo_gpu_pin_query_topk_batch[_f16]: batched kernel (fp16 + fp32), query tiling
capped by a 1536-slot shared candidate buffer (24KB/block, fits all targets),
high block count since it's compute-bound. Bit-for-bit identical to M
single-query mojo_gpu_pin_query_topk calls (same (dist,rowid) tie-break).
- gpu_cosine_topk_batch('emb','v','queries','qv', k [, precision]) table function
(vss_join-style): query set is a FLOAT[K] column of another table; scores all M
in one batched call. fp16 default. Returns (query_rowid, rowid, dist).
- Benches: cosine_batch_test (batched == single-query exact, both precisions),
cosine_batch_latency (per-query latency vs batch size).
This is the regime where GPU most decisively beats a per-query index: HNSW cannot
use its index for batched queries at all. RTX 4090, 1M×768 fp16 k=10: per-query
latency falls with M and the batch path (1.89ms/q, ~528 q/s) edges past the
single-query path (2.10ms) while staying exact. The naive batch kernel amortizes
the read but isn't compute-optimal at large M — a tensor-core GEMM is the next
lever (noted in DESIGN/README open frontier).
Max K 2048 (in-register row cache); max M per matrix-read pass = 1536/k (larger M
tiles, re-reading the matrix once per tile). Single-query entry points unchanged.
…, Q3 Add STRAT_HASH_GROUP: a GPU open-addressing hash-aggregate for high-cardinality integer group keys, gated to platforms with 64-bit atomics (NVIDIA/AMD). Apple keeps the sort + segmented-reduce path unchanged (no 64-bit atomics). Used by TPC-H Q3, replacing its sort+segreduce — which also drops the injected ORDER BY l_orderkey from Q3's cold path. - seg_hash_kernel (segreduce.mojo): one thread per fact row, grid-stride; applies the row-pass + dim gather + expr-VM metric, hashes the group key (Knuth multiplicative, pow2 mask), linear-probes claiming a slot via Atomic[int64].compare_exchange (claim-or-find, race-safe), then Atomic[int64].fetch_add into per-group int64 accumulator slots. The whole body is `comptime if is_nvidia_gpu() or is_amd_gpu()` (Apple compiles a no-op stub, never launched). Host widens each occupied slot to int128 for the decimal output (per-order Q3 revenue fits int64; the matcher must not route a shape where a single group could overflow int64 — documented). Table cap = next_pow2(2 * orders-dim row count) (tight bound on distinct orderkeys), load factor <= 0.5. - Strategy pick (descriptor.mojo): integer fact group key -> STRAT_HASH_GROUP when has_nvidia/amd_gpu_accelerator() (host), else STRAT_SORT_SEGREDUCE. HASH_GROUP omits the ORDER BY from the materialize SQL. - pin_finalize routes HASH_GROUP through _assemble_hash (emit revenue>0, dim group keys looked up by orderkey), reusing the resident _pin2 cache. Validated: Apple unchanged (Q3 = sort path, q3shuttle ALL PASS, exact; atomic body comptime-excluded; Q5/Q6/Q14 no regression). NVIDIA (RTX 4090): Q3 = HASH_GROUP, EXPLAIN shows GPU_AGG, result EXACT vs stock (count 11620, checksum, top-10 all match). Perf is honest: Q3 GPU warm ~14ms beats single-thread stock (~37ms, 2.6x) but at sf1 still ~0.87x vs 16-thread stock (~11.8ms) -- the GPU source op is single-threaded (materialize + orchestration) and sf1 is small enough that 16 CPU cores win. The hash-agg is the right architecture and should flip at larger scale factors (GPU linear scan vs the CPU's ~10x-per-SF growth); sf10+ not yet measured. Other queries unchanged (Q1 2.0x, Q5 6.5x, Q6 4.2x, Q14 10.3x).
…istic) The generic engine is default-on and cost-blind: it offloaded any matched aggregate, including TPC-H Q3 — which measurement showed is CPU-favorable (loses to 16-thread stock at sf1 0.87x AND sf10 0.54x; the gap widens with scale). Q3's shape — light per-row work over a large fact table producing many groups (~1.5M order groups, atomic hash contention) — is exactly what DuckDB's multithreaded join+hash-aggregate excels at, and the single-threaded GPU source op can't beat it at any tested scale. build_descriptor now DECLINES (returns None -> stock CPU) when the chosen strategy is a high-cardinality group-by (SORT_SEGREDUCE or HASH_GROUP), so the engine never makes such a query slower than the CPU. UNGROUPED / DENSE_GROUP (Q1/Q5/Q6/Q14 — few groups or join-heavy, which the GPU wins 2-11x) are unaffected. GPU_OP_FORCE_HIGHCARD=1 overrides (force-offload for A/B). The pure classifier is unchanged for testing via a build_descriptor_impl(force_highcard) param; the round-trip test passes it to validate classification independent of the policy. Verified: Q3 EXPLAIN shows no GPU_AGG (runs stock, correct), Q6/Q5 still GPU_AGG; force-override re-enables Q3; roundtrip test ALL PASS.
- DESIGN.md: update Performance with the RTX 4090 TPC-H numbers (Q14 ~11x, Q5 ~7x, Q6 ~3.8x, Q1 ~2x), a Vector search (kNN) section (~40-57x vs CPU exact; competitive-and-exact vs vss/HNSW up to ~1M; fp16; single-query bandwidth-optimal), and the cost heuristic (decline high-card group-by). Update Limitations: Q3 is CPU-favorable (measured/settled, kept on CPU); batched-kNN throughput is gated on a tensor-core GEMM (hand-tiled only ~1.4x; layout.tensor_core API is opaque) — the documented next step. - q3_shuttle_test: set GPU_OP_FORCE_HIGHCARD=1 (the heuristic now declines Q3 by default; the test validates the Q3 GPU execution path, so force-enable it).
Bump the pinned Mojo nightly (was 1.0.0b2.dev2026060206) across pixi.toml (host/build deps, [dependencies] mojo, operator-replacement feature), both conda.recipe yamls, and pixi.lock. Stays on the latest b2; b3 is a work-in-progress milestone (Int -> Scalar[DType.int], new SIMDSize). Address breaking changes / deprecations in the new nightly: - List[T] is now implicitly destructible only when T is, so generic code that builds a local List[T] and abandons it on a raise path no longer compiles. Add `& ImplicitlyDestructible` to `_DBase` and to the `get[T] -> List[T]` overloads in chunk.mojo/result.mojo, and validate before building `result` in `List._from_list_child` (typed_api.mojo). - Rename deprecated `MutExternalOrigin`/`ImmutExternalOrigin` to `MutUntrackedOrigin`/`ImmutUntrackedOrigin` across sources, the API generator, and the regenerated _libduckdb.mojo (diff is rename-only, so check-generated-api stays green). - Drop dead `struct_field_count` imports from two tests (reflection API is now Reflected[T]). Also fix a latent pixi multi-env conflict: list the stdlib path first in MODULAR_MOJO_MAX_IMPORT_PATH (`$CONDA_PREFIX/lib/mojo,$PIXI_PROJECT_ROOT`). Mojo takes the first import-path entry that yields `std` and recursively scans the project root, so leading with the root could pick up another materialized env's older std.mojoc under .pixi/ (e.g. the `full` env) and fail to load it. Stdlib-first avoids the shadowing. pixi run test passes (0 failures, 0 deprecation warnings).
…PU syncs mojo_gpu_pin_f16 converted fp32->fp16 *into* emb_dev.map_to_host() staging. map_to_host is bidirectional (DMAs device->host on enter), so for a pure upload it triples bus traffic: measured 6.75 GiB/s vs 14.3 GiB/s for a plain enqueue_copy on an RTX 4090. Convert into a plain host buffer and enqueue_copy the half buffer instead -> 2.1x faster pin upload, recall@10 unchanged (0 genuine misses). The same measurement (pageable enqueue_copy 17.5 GiB/s vs map_to_host 5.2 GiB/s) confirms the existing pageable uploads in the fp32 pin, the cosine re-upload and segreduce_upload were already optimal; comments now warn off map_to_host for these one-time/resident uploads. Also drop ctx.synchronize() calls between independent enqueue_* ops in segreduce_run / segreduce_run_hash (one sync only before the host read-back); they were full host stalls that the single runtime stream already orders. Verified bit-exact on Apple (Metal) and NVIDIA (RTX 4090): q5/q6/q14 shuttle, q3 shuttle incl. HASH_GROUP (strategy 3), cosine_topk_test (fp32 exact), cosine_f16_test (recall >= 0.99).
…OCK128 (experiment, default off) Static profiling (ptxas -v on dumped PTX; ncu unavailable on the test box) showed the 1-warp-per-block (block_dim=WARP=32) grid kernels cap occupancy at 50% on sm_89: the 24-block/SM limit binds first, while registers (30-37/thread, 0 spill) and shared mem (~0) leave full headroom. Added multi-warp variants (seg_ungrouped_kernel_mw, seg_dense_kernel_mw): SEG_NWARPS=4 warps/block, each warp reduces via warp.sum, then a shared-mem + one-barrier combine across warps writes one per-block partial. The partials layout and host int128 reduction are unchanged. Routed via the GPU_OP_BLOCK128 env flag (default off, no behavior change); bit-exact both states on Apple (Metal) and NVIDIA (RTX 4090). A/B on RTX 4090 (warm, ~6M rows, per-run incl. read-back), stable over repeats: Q6 UNGROUPED M=1: 0.87-0.95 ms (32) -> 0.832 ms (128) = ~1.15x faster Q1 DENSE G=4 M=8: 4.10 ms (32) -> 4.44 ms (128) = ~0.92x (8% slower) The occupancy headroom converts to speedup only for the latency-bound light kernel (ungrouped); the dense kernel is bound by per-thread local-memory accumulator work (640 B [G*M] stack frame), so 4x more threads add overhead with no latency to hide and it regresses. Shelved flag-gated rather than wired into strategy selection: a global flip is ruled out by the dense regression, and the ungrouped win is small next to the source-op/materialize cost in real Q6.
…CORE (NVIDIA, default off)
Batched fp16 cosine-kNN was compute-bound: the scalar warp-dot per (row, query)
dominated, so amortizing the matrix read across a query-tile never paid off
(~1.1-1.4x over single-query). Replace the inner product with a tiled tensor-core
GEMM (Q.Emb^T, m16n8k8, transpose_b, fp16->fp32) fused with a streaming per-query
top-k over N, so the M*N similarity matrix is never materialized -- only Emb is
read and only M*k results are written.
New src/tc_knn.mojo: the fused stage-1 kernel (grid-stride over N in BN chunks,
MMA into a shared S-tile, cosine distance from host-precomputed q/e norms,
persistent per-thread register top-k with the same (dist, rowid) tie-break as the
scalar path), stage-2 per-query merge, the run_tc_knn_batch host driver
(comptime-specialized over K in {384,768,1024,1536}), and tc_knn_supported(K, k).
src/gpu_kernels.mojo _run_topk_batch routes to the fused path when, at comptime,
the matrix is fp16 AND has_nvidia_gpu_accelerator(), and at runtime GPU_OP_TENSORCORE
is set and (K, k) is supported; otherwise the existing scalar kernel runs unchanged.
No exported C-ABI signature changes. Default off -> zero behavior change.
Gating uses the host-vs-kernel split (memory mojo-gpu-target-introspection): the
host branch + driver gate on has_nvidia_gpu_accelerator() (comptime-false on Apple,
so the layout/TensorCore code is never compiled for Metal), kernel bodies gate on
is_nvidia_gpu(). Apple (Metal) and NVIDIA (RTX 4090) builds both green; flag-unset
tests unchanged; flag-on on Apple falls back to scalar.
Validated on RTX 4090 through mojo_gpu_pin_query_topk_batch_f16 (N=1M, K=768, k=10):
recall@10 = 1.0, 128/128 exact id match, 0 genuine misses. Per-query latency
off->on: M=16 2239->323us (6.9x); M=100 1899->54us (35x); M=1000 1895->31.8us
(59.6x). The win grows with batch as the matrix read amortizes across the GEMM;
at M=128 the kernel is HBM-bandwidth-bound on the Emb read (~69% of peak), not
tensor-core-bound. Design notes folded into DESIGN.md (kNN section + portability).
…dgroup-matrix backend
Widens the fused matmul kNN path (GPU_OP_TENSORCORE, default off) two ways.
1) Metrics (NVIDIA). The Q.Emb^T MMA core is identical for cosine / L2 / inner-
product; only the distance epilogue + tie-break differ. tc_knn.mojo gains a
comptime `metric` switch (cosine `1-dot/(|q||e|)`, L2 `|q|²+|e|²-2dot`, IP `-dot`),
run_tc_knn_batch a runtime `metric`, and the enorm kernel a squared-norm flag.
New C-ABI `mojo_gpu_pin_query_topk_batch_f16_metric` (existing cosine symbol
unchanged, delegates metric=0). gpu_operator.cpp exposes it via
`gpu_cosine_topk_batch(..., metric := 'cosine'|'l2'|'euclidean'|'ip')`. Validated
(N=256k K=768): L2/IP recall@10 >= 0.997, 0 genuine misses; cosine unregressed
(N=1M 128/128 exact). Non-normalized L2 has an fp16 squared-norm cancellation
caveat (documented; normalized is exact) -- a robust fp32/tf32-input path costs
~2x Emb bandwidth, so left as a noted follow-up. The per-element
array_cosine_distance projection operator is a separate non-top-k path, untouched.
2) Apple backend. layout.tensor_core has no Apple support and the 16x16 Apple MMA
is M5-only, but Apple M1-M4 have a working 8x8 simdgroup_matrix MMA. New
src/tc_knn_apple.mojo implements the fused kNN with it (raw _mma8x8 intrinsic,
manual per-lane fragment loads + transposed Emb^T gather, 128-query amortization,
the streaming top-k ported from tc_knn.mojo). Routed from _run_topk_batch under
`comptime if has_apple_gpu_accelerator()` + GPU_OP_TENSORCORE + cosine; scalar
fallback otherwise. Validated on M3 Max: recall@10=1.0 (0 genuine misses) across
K in {384,768,1024}; per-query A/B vs scalar batched 2.9-7.5x (unified-memory
bandwidth-bound). Cosine only on Apple for now.
Each backend's MMA code is gated on its own host accelerator query
(has_nvidia_gpu_accelerator / has_apple_gpu_accelerator), so it compiles only for
its target and is elided on the other; the scalar kernel is the universal fallback.
Both builds green (Apple local + NVIDIA frederick); no exported C-ABI signature
changed; default off -> zero behavior change. DESIGN.md kNN + portability sections
updated; README documents the metric param.
…er to runtime JIT)
Replace the per-row postfix-VM interpreter (expr_vm.eval_program) with
compile-time-specialized, stackless straight-line kernels dispatched by
descriptor kind (Q1/Q6/Q14/Q5), keeping the interpreter as the universal
fallback for any unrecognized shape. This is the structural Mojo answer to
cuDF's runtime nvrtc AST_JIT: fused, register-resident, branch-free
evaluation with no runtime compiler and full NVIDIA/Apple/AMD portability.
Additive only: new seg_{ungrouped,dense}_kernel_q* variants + eval_program_fast
in segreduce.mojo; segreduce_run gains a kind param routing to them; GpuPinned
carries desc.kind and _assemble passes it (single WARM/COLD dispatch site).
expr_vm.mojo unchanged.
Bit-identical to the interpreter, verified end-to-end on both platforms:
- Apple M3 Max: all kernel oracles + q1/q6/q14/q5/q3 shuttle EXACT; probe 4.7x (Q1) / 1.8x (Q6)
- NVIDIA RTX 4090: all oracles + shuttle EXACT; probe 17.9x (Q1) / 5.5x (Q6)
Bench probe added: bench/expr_comptime_probe.mojo (interpreter vs comptime, asserts bit-identical).
… pushdown (measured) DESIGN.md: (1) describe the comptime-specialized expression kernels shipped in 3628cef and frame them as Mojo's compile-time answer to cuDF's runtime nvrtc JIT; (2) record the dynamic-filter / semi-join-pushdown decision as measured-and-declined. Phase-0 feasibility (bench/semijoin_pushdown_probe.mojo, TPC-H Q5, sf1+sf10): pushing the dim semi-join into the fact materialize cuts fact rows 33x but runs 1.4-5x SLOWER than a plain full scan, and EXPLAIN ANALYZE shows DuckDB already pushes its own dynamic filters into the lineitem scan (227k rows, not 6M). The optimization is something DuckDB's planner already does better and still loses to a plain scan; the only possible win is discrete-GPU cold-first-touch (Q5 only), which the warm pin amortizes and a GPU-direct scan removes. Not pursued.
…pping join-tree bug GpuPreOptimize: also disable STATISTICS_PROPAGATION (restored after, like COMPRESSED_MATERIALIZATION). On a multi-table FK-join it derives a redundant range filter from join-key stats and leaves it as a residual LOGICAL_FILTER inside the join tree, which blocked a legitimate Q3/Q5-shape DENSE_GROUP offload (verified: query is 'unsupported' with the pass on, routes with it off). IN_CLAUSE and LATE_MATERIALIZATION are deliberately left ENABLED — proven they only ever produce safe CPU fallback (mark-join/CHUNK_GET, or ORDER-BY-LIMIT above the aggregate) and never wrong-serialize or unlock an offload, so disabling them would just slow every query for no gain. Also fixes a latent wrong-result bug surfaced during the investigation: CollectJoinTree descended into a residual LOGICAL_FILTER while DISCARDING its predicate, so a join carrying a residual (e.g. OR) filter could route with the filter silently dropped. Now bails to CPU when a residual filter has expressions we don't serialize. Validated: all shuttle tests (Q1/Q3/Q5/Q6/Q14) EXACT on Apple + NVIDIA; residual-OR, large-IN/mark-join, ORDER-BY-LIMIT, ungrouped MIN/MAX all match stock CPU.
…atform-aware Q3's high-card integer fact key maps to STRAT_HASH_GROUP on NVIDIA/AMD (64-bit atomics) and STRAT_SORT_SEGREDUCE on Apple, exactly as descriptor.mojo selects. The test hardcoded SORT_SEGREDUCE, failing on NVIDIA (pre-existing, unrelated to the C++ optimizer change). Mirror the same comptime has_*_accelerator condition.
…uristic near-optimal) Refactor the binary high-cardinality decline into _should_decline(desc, force_highcard) in descriptor.mojo: keeps the existing high-card rule exactly, and adds a default-OFF transfer-vs-compute hook (GPU_OP_COSTMODEL=1, GPU_OP_COSTMODEL_MINROWS=<n>) using the fact GET's post-filter est_cardinality. Default behavior is byte-identical. Investigation (documented in DESIGN.md, mirroring the Q3 'measured and settled' note): the operator's win is the warm/repeated path; the cold path doesn't beat stock by design and is amortized by the resident pin. A small est_cardinality is exactly the case that loses cold AND wins warm-if-repeated, and there is no repeat-history signal at plan time, so cardinality alone can't separate 'tiny, loses even warm' from 'tiny, repeats and wins'. An active threshold would regress the accepted-and-winning classes; hence default-off hook + documented rationale rather than speculative complexity. Declining is always correct (CPU fallback), so the hook is correctness-safe. Verified: build clean; all shuttle (Q1/Q5/Q6/Q14) + Q3 + roundtrip EXACT on Apple + NVIDIA; hook fires when enabled. Pure host-side planner logic.
…(Mojo, portable) Attacks the #1 open frontier (the cold-path CPU materialize). New portable Mojo GPU decoders for DuckDB v1.5.3 native column segments, reached in-process by pinning the segment block live in the buffer manager (no file re-read) — simpler than the CUDA-only reference (Sirius). The decoders are plain Mojo bit arithmetic, comptime-specialized per type, so they run on NVIDIA/Apple/AMD (no Metal DuckDB-segment decoder exists elsewhere). Three phases, all validated bit-exact on Apple M3 Max + RTX 4090: - B (src/native_decode.mojo): unpack_value + uncompressed/bitpacking(CONSTANT,FOR) kernels; synthetic round-trip test bench/native_decode_test.mojo (pixi: gpu-op-native-decode-test), 8 cases incl. negative frames, width=1, short tails -> ALL PASS. - A (gpu_operator.cpp): in-process reachability via Catalog -> DuckTableEntry -> DataTable -> RowGroupCollection -> RowGroup::GetRawColumnData -> ColumnData::GetColumnSegmentInfo -> ColumnSegment::CreatePersistentSegment -> BufferManager::Pin. Debug table functions gpu_native_segment_info / gpu_native_group_dump. - C (gpu_native_decode_check + mojo_gpu_decode_segment): decode l_shipdate GPU-direct and compare to the SQL scan -> 6,001,215 values bit-identical, 0 mismatches, on BOTH platforms. Discovered beyond the plan: per-2048-row group modes vary WITHIN a segment even under a 'FOR' segment label (l_orderkey/l_linenumber are DELTA_FOR). The decoder defers DELTA_FOR (reports deferred_rows) rather than mis-decoding -- correctness-first. DELTA_FOR/RLE/ validity/strings + replacing the Connection::Query feed are the next phases (DESIGN.md). Additive only: existing kernels/operator/cold-path untouched; this is a validated parallel decode path. Existing shuttle tests still EXACT on both platforms.
… pin control Fixes the 'process-lifetime, never evicted' VRAM leak in the kNN embedding pin cache (pinning many embedding matrices would OOM the GPU) — the gap Sirius solves with cuCascade tiered memory. Replaces the two never-evicted fp32/fp16 pin maps with one bounded LRU registry (ResidentPin: handle, bytes, kind, last_use, refcount): - GPU_OP_PIN_BUDGET_MB (default 4096; 0 = unbounded/legacy); evicts LRU when a new pin would exceed the budget, freeing DeviceBuffers via the existing pin_free paths. - OOM-resilient: on device-alloc failure, evict-and-retry until it fits or nothing is evictable, then graceful CPU fallback — never a crash. - RAII PinLease (refcount) guards in-use entries: a buffer a query is actively reading is never evicted/unpinned (the GPU query runs outside the cache lock). - Control table functions: gpu_pin_status(), gpu_unpin(key), gpu_unpin_all(), gpu_pin_table(table, column [, precision]) for pre-warming. Eviction is correctness-safe: a pin is pure cache; an evicted entry just forces the next touch onto the COLD rebuild path (mojo_gpu_pin_begin returns COLD on a cache miss and re-materializes) — slower, never wrong. Deferred (documented in DESIGN.md): bounding the smaller Mojo-side aggregate _pin2 cache (needs SegResident teardown), and host-tier spill/re-promote (cleanest on Apple unified memory) — the natural next step. Verified on Apple M3 Max + RTX 4090: new gpu-op-pin-evict-test ALL PASS (bounded, eviction fires, cold-rebuild bit-correct vs stock array_cosine_distance, empty after unpin_all, no VRAM leak under churn); all shuttle (Q1/Q3/Q5/Q6/Q14) + native-decode regression still EXACT (gpu_native_* + cosine paths intact through the rewrite).
…mark_runner numbers Re-measured warm GPU-vs-stock with DuckDB's benchmark_runner (1 cold + 5 hot, Verify() on every hot run), both platforms: NVIDIA 4090: Q14 16.9x Q6 7.9x Q5 7.2x Q1 6.2x Apple M3 Max: Q14 3.1x Q1 2.6x Q5 1.9x Q6 1.2x All result-verified vs the TPC-H sf1 answer CSVs; routing confirmed via GPU_OP_SHADOW + a GPU_OP_GENERIC=off A/B control. These reconcile with (and on NVIDIA exceed) the prior numbers. Added a methodology note: the interactive DuckDB CLI cannot measure the warm regime (identical re-runs hit result handling and don't re-execute; varying a filter constant forces a cold signature since the pin key includes constants), which makes warm GPU work spuriously look 50-170x slower. Also documents that the 'empty grouped result on repeat' interactive-CLI observation is a confirmed CLI rendering artifact, NOT an operator defect: the source op allocates fresh per-execution state and the warm finalize rebuilds the full result set every call (verified 3 ways: benchmark_runner Verify() passed on all 10 hot grouped-Q1 runs per platform; piped-stdin CLI returns byte-identical full rows on repeat; source read of PhysicalGpuAgg lifecycle). No code change.
…DELTA (full fixed-width coverage) Adds the remaining fixed-width BITPACKING modes to native_decode.mojo's bitpacking_decode_kernel, completing fixed-width codec coverage (CONSTANT, CONSTANT_DELTA, FOR, DELTA_FOR; only INVALID/AUTO/unknown still defer): - CONSTANT_DELTA: [T frame][T delta] -> out[i]=frame+i*delta. - DELTA_FOR: [T frame][T width][T delta_offset] + packed; one CTA per 2048-row group, frame-decode into shared, inclusive prefix-sum seeded with delta_offset, write out. Group-wide scan is mathematically equivalent to DuckDB's per-32-block ApplyFrameOfReference+DeltaDecode chaining (verified vs v1.5.3 bitpacking.cpp). Serial shared-mem scan chosen for correctness/portability (no warp/CUB intrinsics). This was the blocker for real key columns: l_orderkey and l_linenumber are entirely DELTA_FOR, and l_extendedprice had DELTA_FOR groups inside FOR-labeled segments. Verified bit-exact on Apple M3 Max + RTX 4090: gpu_native_decode_check now reports 'PASS: 6001215 values bit-identical', mismatches=0, deferred_rows=0 for l_orderkey, l_linenumber, l_extendedprice, l_shipdate, l_quantity, l_discount (l_orderkey/ l_linenumber were fully deferred before; l_extendedprice had ~4096 deferred). Synthetic native_decode_test extended to 14 cases (CONSTANT_DELTA + DELTA_FOR: increasing, decreasing/all-negative, non-monotonic; multi-group + short tail) - ALL PASS both platforms. q6 shuttle still ALL PASS. RLE is N/A (a segment-level CompressionType, not a per-group BitpackingMode; no TPC-H fixed-width column uses it). Additive; cold feed path unchanged.
…asure-first design DESIGN.md cold-path frontier: update decoder coverage to all four fixed-width BITPACKING modes (CONSTANT/CONSTANT_DELTA/FOR/DELTA_FOR; the 'Not yet: DELTA_FOR' note was stale after Phase D), and record the measure-first Phase G decision: - Replacing the cold feed for speed is marginal (cold path already materializes all ~6M rows; no WHERE pushdown; filter is in-kernel) -> only trims transfer volume. - The real prize the decoder enables is predicate-independent residency (drop filter consts from _signature + move the filter into the kernel -> warm across constants). - Validity DEFERRED (no nullable TPC-H column), strings DEFERRED (only Q1/Q14 need them; Q6/Q5-fact decodable now), RLE N/A. - Being proven via a flag-gated (GPU_OP_NATIVE_DECODE) Q6 PoC.
…for Q6 (substrate) Behind GPU_OP_NATIVE_DECODE (default off), the COLD feed for the Q6 (UNGROUPED, all-fixed-width) fact request decodes its columns GPU-direct from DuckDB-native storage instead of the nested Connection::Query + chunk-gather. Reuses the already-validated decode helpers (ResolveNativeColumn / ForEachColumnSegment / mojo_gpu_decode_segment / CodecToMojo / ParseBitpackingGroup) and feeds via the identical mojo_gpu_feed_column ABI, so downstream (filter, kernel, signature, result) is byte-for-byte unchanged. All-or-nothing per request: any unsupported column/segment (VARCHAR/Dictionary, parquet-backed, nullable, ragged, has_updates) falls back to the SQL feed -> never partial, never wrong. Additive (+191/-0), gpu_operator.cpp only; no Mojo change. This is the low-risk substrate for Stage 2 (predicate-independent residency). On its own it only trims cold host->device traffic: Q6 sf1 moves the compressed bitpacked segments = 64,223,320 bytes vs the SQL feed's materialized 168,034,020 bytes (~2.6x less), with decode on-GPU. Verified bit-exact on Apple M3 Max + RTX 4090: flag-on Q6 = 123141078.2283 = stock (GPU_OP_GENERIC=off), GPU-direct log fires only with the flag on; flag-off q6/q1/q5/q14/q3 shuttle all EXACT (default path untouched).
…for Q6 (warm across varying constants) The real prize of the GPU-direct decoder work: a flag-gated (GPU_OP_NATIVE_DECODE) path where Q6's GPU residency is decoupled from its filter constants, so running Q6 with DIFFERENT WHERE constants reuses the resident columns (WARM) instead of cold re-materializing. Turns 'every distinct-constant query is cold' into 'first cold, rest warm' (e.g. a dashboard varying Q6's date range: 2nd+ query ~365ms->~0.4ms). How: the cold materialize already loads ALL rows (no WHERE), so the resident columns are content-predicate-independent; only the signature key + a host-prebaked pass column coupled them to a predicate. This decouples both: - segreduce.mojo: new seg_ungrouped_kernel_q6_pred (copy of the Q6 kernel) evaluates the Q6 predicate IN-KERNEL from the resident l_shipdate/l_discount/l_quantity columns + 5 bounds passed as launch params (GE/LT/GE/LE/LT over the same int64 decimal-scaled values as the host pass-bake -> bit-identical). segreduce_run routes to it when q6_pred_active. - gpu_kernels.mojo: Q6PredSpec/_q6_pred_spec resolves the canonical Q6 WHERE to slots + bounds by (column,cmp) role (eligible=False on ANY deviation -> keep per-constant path; never a wrong number). _signature excludes filter constants for the eligible Q6 path (keeps fact/cols/strategy/kind/cmp-structure). GpuPinned caches the constant-independent slots (NOT the bounds); _assemble threads the LIVE descriptor's bounds into every run (cold AND warm), so a warm hit applies THIS query's constants over the cached resident columns. COLD skips the host pass-bake. pin_begin gains a GPU_OP_PIN_LOG-gated WARM/COLD trace. All flag-gated; flag-off is byte-identical (signature still constant-keyed; non-Q6 shapes never take this path - verified Q1 still constant-keyed). Thesis verified on Apple M3 Max + RTX 4090: one session, GPU_OP_NATIVE_DECODE=1, Q6 with 3 different constant sets -> COLD, WARM, WARM (one constant-free signature); revenues 123141078.2283 / 112296798.0742 / 227447030.1782 all bit-exact vs stock (GPU_OP_GENERIC=off, same constants). flag-off q6/q1/q5/q14/q3 shuttle all EXACT.
…ralize the Q6 win) Extends the proven warm-across-constants mechanism from Q6 to Q1 (DENSE_GROUP) and Q14 (UNGROUPED + dim), still flag-gated on GPU_OP_NATIVE_DECODE, no string decode needed: the constant-dependent part is only the fixed-width fact filter; Q1's VARCHAR group keys and Q14's promo flag (p_type LIKE 'PROMO%') are constant-independent and stay resident via the current feed. Generalized the q6_pred mechanism: - segreduce.mojo: _fpred_pass_dev (in-kernel general fact-range filter, byte-mirror of the host _pred_pass over the same int64 slot values) + seg_dense_kernel_q1_pred and seg_ungrouped_kernel_q14_pred; segreduce_run gains gen_pred_active + a flattened (slot,cmp,bound) fpred_list passed by value (fresh per run) with dispatch branches. - gpu_kernels.mojo: _gen_pred_eligible (per-kind structural gate: KIND_Q1->DENSE_GROUP, KIND_Q14->UNGROUPED, all-range single-fact-column filter, and crucially NO dim filters so const-decoupling stays sound -> else fall back to the per-constant signature, never a wrong number); _signature emits a constant-free per-kind tag (gen_pred=<kind>); GpuPinned caches the constant-independent slots+cmps (not bounds); _assemble threads the LIVE descriptor's bounds every run (cold + warm); both finalize paths skip the host pass-bake when active. Thesis verified on Apple M3 Max + RTX 4090: Q1 (3 cutoffs) and Q14 (3 ranges) each COLD,WARM,WARM under one constant-free signature; results bit-identical to stock with the same constants (Q1 integer/decimal SUMs + count exact -- e.g. N,O sum_qty 74476040/73533166/75637054 = stock; Q14 promo 16.380.../16.575.../16.584... = stock). Q1 avg_disc differs only in the 17th sig-digit -- a pre-existing GPU-AVG-as-double artifact present with the flag off too (DESIGN.md: 'averages match to double rounding'), not introduced here. Flag-off byte-identical (per-constant signature, 3x COLD); q6_pred unaffected; Q3/Q5 never take the path; all shuttle tests EXACT both platforms.
…mpletes the grouped set) Extends the warm-across-constants mechanism to Q5 (the last accepted grouped query), flag-gated on GPU_OP_NATIVE_DECODE. Q5's filter is on DIMENSIONS (r_name + o_orderdate), not the fact, and the naive 're-feed dims on warm' is blocked by the cold/warm feed contract (warm feeds nothing). Path B instead makes ALL resident buffers constant-INDEPENDENT and evaluates region+date IN-KERNEL from per-run scalars, exactly like Q6/Q1/Q14: - segreduce.mojo: new seg_dense_kernel_q5_pred (same lane-stride/warp.sum/partials layout as seg_dense_kernel_q5) with in-kernel pass = (o_orderdate in [lo,hi)) AND (cust_nation==supp_nation) AND (supp_region==asia_region), gid = raw supplier nationkey; segreduce_run gains a q5_pred branch in DENSE_GROUP. - gpu_kernels.mojo: _q5_pred_eligible (canonical 5-dim Q5, only region.r_name + orders.o_orderdate constant-bearing filters; any other filter -> decline -> per-constant path) + _q5_pred_enabled (flag) + _q5_pred_params (o_lo/o_hi + asia_region resolved from an all-regions name->key map cached on cold). _signature drops region/date constants for the eligible path (+|q5_pred=1). GpuPinned caches the region map + constant-independent slots. _pin_finalize_q5 builds the constant-independent residency on COLD (gid=raw nationkey, raw dim arrays o_orderdate/ cust_nation/supp_nation/supp_region, n_name labels for all nationkeys, emit revenue!=0; G>64 -> per-constant fallback) and threads live per-query scalars on COLD+WARM. gpu_operator.cpp unchanged. Correctness: nationkey<->n_name is 1:1 (group-by-nationkey == GROUP BY n_name); in-kernel supp_region==asia gates exactly stock's r_name rows; per-gid revenue is the same int128 SUM; emit revenue!=0 yields exactly stock's row set. Verified bit-exact on Apple M3 Max + RTX 4090: one session, GPU_OP_NATIVE_DECODE=1, Q5 across ASIA/EUROPE/AMERICA(1994) + ASIA(1995) -> COLD,WARM,WARM,WARM (one constant-free signature), each full (n_name,revenue) incl. ORDER BY row order bit-identical to stock. DECISIVE probe: MIDDLE EAST(1996) -- a region NEVER materialized on the ASIA cold -- served WARM and bit-exact, proving residency is truly region/date-independent. Empty date window -> 0 rows == stock. Q5 with an extra o_totalprice filter -> eligibility declines -> per-constant fallback, still exact. Flag-off byte-identical (per-constant signature); all shuttle tests EXACT both platforms.
…Q1/Q5/Q6/Q14 Update DESIGN.md from 'being proven via a Q6 PoC' to the shipped+verified state: predicate-independent residency (warm across varying filter constants) is implemented flag-gated (GPU_OP_NATIVE_DECODE) and proven bit-exact on Apple + NVIDIA for all four GPU-accepted queries. Notes the per-kind approach (fixed-width fact filter for Q6/Q1/Q14; Path B in-kernel region/date for Q5), the strict eligibility gate + per-constant fallback, and the decisive Q5 'never-materialized region served warm' probe.
… elimination (audit batch 3) Audit Group E: the FK->dim INNER join is lowered as a dense per-key gather dims[off+key] (one dim row per fact row), which silently mishandled two cases: - NON-UNIQUE build side: a many-to-many join collapsed to 1x (picks one dim row) -> aggregate too low. - PARTIAL/non-covering dim with no filter: unmatched fact rows were kept (INNER join not enforced) -> aggregate too high. Part 1 (partial-dim): force needs_pass=True for every fact->dim edge in _pin_finalize_generic_dims so the kind-2 build's 0/1 presence array eliminates fact rows whose FK has no matching dim row. No-op for covering dims (Q14/Q3). Part 2 (non-unique fanout): decline at the C++ JOINS serializer unless every DIRECT fact->dim edge's dim-side key is provably unique -- declared PRIMARY KEY/UNIQUE constraint, else an exact count(*)==count(DISTINCT key) probe on a fresh nested Connection (dbgen TPC-H tables carry no PK; DuckDB HLL approx-distinct is too inaccurate as a uniqueness oracle). Gated to direct fact->dim edges only (dim<->dim correlated/transitive edges are not gather edges; gating them over-declines Q5). The probe result is CACHED per (db,table,col) keyed with the table's exact storage row count, so it runs once per session (not per warm iteration) and re-probes only when DML changes the row count -- preserving the warm Q5/Q14 wins. Nonzero -> decline to CPU. Verified both platforms (Apple + RTX 4090): 5 shuttle + roundtrip + or_filter off & on; live TPC-H sf1 Q5 + Q14 route (GPU_AGG) and match stock bit-for-bit; a non-unique dim join declines to correct stock (not the collapsed half); the probe fires once then cache-hits; partial-dim joins eliminate absent-FK rows.
…eturning nan/-inf (audit batch 4a) Audit Group G (NVIDIA transcendental f64 path, default-on): sum/avg of sqrt(x) with x<0, or ln/log10/log2(x) with x<=0, returned nan/-inf where stock DuckDB RAISES an Out of Range error -- one out-of-domain row silently poisoned the whole aggregate. Fix (NaN-sentinel + host raise; preserves the in-domain win class): the f64 VM maps domain violations to a NaN sentinel (sqrt(x<0)->NaN; ln/log10/log2(x<=0)->NaN, normalizing the raw-log -inf/nan) while exp overflow stays a valid Inf -- so a NaN result is an unambiguous domain-error signal. The host f64 finalize detects a NaN result cell (gated off the stats path, whose zero-variance nan is legitimate) and maps it to an rc (10 sqrt / 11 log); the C++ raises OutOfRangeException with the matching DuckDB message. Valid Inf (exp overflow) is preserved and does NOT raise; grouped (DENSE) domain errors raise the whole query like stock. (The sentinel can't tell zero from negative apart, so a log-of-zero reports "...of a negative number"; both still raise Out of Range.) Verified both platforms: Apple 5 shuttle + roundtrip + or_filter off & on (f64 declines on Apple -> no regression); NVIDIA RTX 4090 -- sqrt(neg)/ln(0)/ln(neg)/log10/log2 raise matching stock, in-domain sqrt/ln/exp/log2 route + match (rel-err <=3e-11), exp overflow -> inf preserved, grouped domain error raises.
…s (audit batch 4b) Audit Group H (NVIDIA stats f64 path, default-on, UNGROUPED): variance/stddev/covar/ corr/regr_* used the one-pass sum-of-squares form (Sx2 - Sx*Sx/n) over RAW second moments, which catastrophically cancels on large-magnitude data -- DECIMAL(18,2) ~1e8 gave NEGATIVE variance (clamped to 0 for stddev), corr=-nan, and ~7.4e-4 rel-err at 1e6 magnitude. Stock uses Welford. Fix (shifted-data centered moments, single-pass, no kernel-launch change): subtract a per-column shift c (the row-0 value, near the data magnitude) from each value before squaring, so the kernel accumulates sum(x-c), sum((x-c)^2), sum((x-cx)(y-cy)). Variance/covar/corr/regr_slope/r2 are shift-invariant and use the centered sums unchanged; the mean-returning kinds (regr_avgx/avgy/intercept) add the shift back. Built from existing OP_PUSH_CONST/OP_SUB (no new opcode); c=0 fallback (byte-identical) when no rows / skip-materialize-omitted. The negative-variance->0 clamp is now harmless. n<2-> NULL, zero-variance, and corr-of-constant->nan semantics preserved. UNGROUPED-only; grouped stats still decline. Verified both platforms: Apple 5 shuttle + roundtrip + or_filter off & on (stats decline on Apple -> no regression); NVIDIA RTX 4090 -- large-magnitude (1e8) var/stddev/covar/ corr/regr now positive + match stock (was negative/nan); medium (1e6) rel-err ~2e-13 (was 7.4e-4); small/constant-column(corr nan)/n<2(NULL) preserved.
…e exactness contract (audit Group B) Audit Group B: the segreduce int64 contract (per-row metric + per-block partial fit int64, host widens to int128) silently WRAPS where stock widens to HUGEINT / raises, for per-element products that overflow int64 (sum(x*y), x,y ~1e10) or per-block partials of huge-magnitude values. This is data-magnitude-dependent and CANNOT be declined at plan time without also declining the correct int128-result TPC-H sums (Q1/Q6 are DECIMAL(38,x) int128 results computed via this exact int64-partial + host-int128-fold path -- verified correct). The proper fix is runtime overflow detection -> CPU fallback (matching stock's overflow-checked add), deferred. Comment only -- documents the limitation in the contract where future maintainers will see it.
…ilters (GPU_OP_FILTER_OR) Extends NR3 (default-off OR-filter pass-program) from OR-of-equalities to also handle <,<=,>,>=,!= over INTEGER/DATE columns -- enabling `WHERE a<10 OR a>95`, OR-of-AND-of- ranges, date-range OR, etc. Adds 5 additive expr-VM comparison opcodes OP_LT/LE/GT/GE/NE (19..23) in lockstep raw_plan_tags.mojo + raw_plan.h (NO RP_MAGIC / section-layout change); the INT eval_program handles them like OP_EQ (pop rhs,lhs -> push 0/1). EmitOrEq maps the 6 comparison ExpressionTypes to the opcodes and inverts the comparator when the column is the RHS (`5>a` == `a<5`), keeping lhs=column. All NR3 gates unchanged (INTEGER/DATE-only, VARCHAR/DECIMAL decline, op-count<=64, const-type-must-match). BETWEEN-within-OR rides the existing AND->MUL recursion. Still default-off. (Single-column OR-of-BETWEEN folds into the scan's table filters in DuckDB -> no residual LogicalFilter -> declines, matching stock; not a gap.) Verified both platforms (Apple + RTX 4090): 5 shuttle + roundtrip + or_filter (now with an OR-of-range tape, COLD+WARM bit-exact) off & on; live A/B -- a<10 OR a>95, const-left inversion, a!=k OR, date-range OR, OR-of-AND-of-ranges all route + match stock; VARCHAR range declines; OR-of-eq regression intact; empty-match -> NULL.
…a default-on sum/avg pin-signature collision Two host-side changes (no GPU kernel change), full commit gate validated on Apple + RTX 4090: 5 q*_shuttle + raw_plan_roundtrip OFF+ON = 12/12 BOTH platforms; live-SQL nullable_sql_test routes + matches a CPU NULL-semantics reference bit-exactly; avg collision probe fixed. Flag OFF is byte-identical to before. 1. GPU_OP_NULLABLE (default-off, presence-only) — widen the operator's broadest coverage gate. It declined the ENTIRE offload whenever any projected column lacked a NOT NULL constraint (TPC-H is all-NOT-NULL, so this never bit TPC-H but excluded ~all real-world tables). Relax it for a provably-safe slice: single-table, UNGROUPED, a single SUM with an INT128 result over a nullable agg/filter column. Mechanism: a NULL agg/filter row is semantically a filtered-out row, so each row's validity is AND-ed into the EXISTING host-baked pass column in _pin_finalize_generic (reuses the NR3 + empty-sum machinery, zero GPU kernel change); ungrouped_count_m is summed over the SAME pass-gated rows so sum stays exact, and all-NULL / empty-match -> SQL NULL. Plumbing: FedColumn.validity + a new additive C-ABI mojo_gpu_feed_validity (tape HEADER untouched) + C++ FlatVector::Validity capture in the cold-feed loop; the role-aware accept/decline runs at the END of SerializeMatchedPlan where roles are resolved. ret_is_int128 is load-bearing: it keeps the slice on the int128 assembly path and EXCLUDES AVG (DuckDB rewrites a nullable avg(x) into a `sum` aggregate with a DOUBLE output -> the int assembly would write res_lo while the DOUBLE extraction reads res_f64 => 0.0) and the f64 transcendental SUM. MUST-DECLINE (all verified -> stock -> correct): grouped, joins, count(*), avg, multi-agg, MIN/MAX, stats/ transcendental, IS NULL/IS NOT NULL filters, native-decode/predicate-independent paths. New live-SQL test bench/nullable_sql_test.mojo + a gpu-op-nullable-test task. 2. Fix a pre-existing DEFAULT-ON sum/avg warm-pin signature collision (found while verifying #1). _signature folded the aggregate-program fingerprint ONLY on the f64 path, so sum(c) and avg(c) over the same column shared a signature: running sum(x) then avg(x) on the same table made avg WARM-reuse sum's resident pin (M=1, no count metric) -> avg = 0.0 (silent wrong result; TPC-H masks it -- Q1's avg is grouped + canonical, Q6 is sum). Make the agg-program fingerprint UNCONDITIONAL so the int path keys distinctly. Monotonic-safe: only MORE-distinct signatures, so a missed warm-hit just re-pins cold (still correct) and the committed warm wins cannot regress (shuttle warm-reuse intact). Same bug class as c584417. Repro: bench/avg_ungrouped_probe.mojo.
…ntal SUM + statistical aggregates) Widen the nullable safe slice (still default-off, single-table, UNGROUPED, single aggregate, no native-decode) to the NVIDIA-only f64 path: a transcendental SUM (sum(sqrt|exp|ln|log10|log2|pow|sin|cos(x))) or a statistical aggregate (stddev/var/covar/corr/regr_*). The f64 ungrouped kernel gates EVERY metric by the host-baked pass column (_row_passes, n_fpred==0 path), so the SAME validity fold from the int128 SUM slice excludes NULL rows from the shared f64 sums with no kernel change. A 2-arg stat (corr(x,y)) is NULL-excluded when EITHER arg is NULL -- exactly AND-ing both columns' validity into the one pass bit, matching SQL. AVG stays excluded (the avg-as-sum-double hazard); a transcendental SUM is told from a plain (rewritten-avg) SUM by a transcendental opcode in its program. Also fixes a validity bug in the stat shift constant that this exposed: _metric_arg_shift picked the centered-moments shift from ROW 0's raw value, but row 0 may be NULL (uninitialized garbage ~1e18) -> the centered (x-c)^2 sums catastrophically cancel (NVIDIA showed stddev=6.7e9, var=-3e19, corr=nan before the fix). Pick the FIRST row where every LOAD_COL slot the arg reads is VALID; fall back to legacy unshifted (c=0) if none. The per-row pass-gate already excludes NULL rows from the sums, so this only restores numerical stability. No-op for all-valid data (r0=0 == old behavior) -> the committed non-nullable stats path is byte-identical. Gate validated on Apple + RTX 4090 (frederick): 5 q*_shuttle + raw_plan_roundtrip OFF+ON = 12/12 each; nullable_sql_test (int) ALL PASS; new nullable_f64_sql_test on the 4090 -- sum(sqrt|ln), stddev, var, corr, regr_slope over nullable DECIMAL all match stock within rel-err <=2.2e-12 (vs garbage before). Apple declines f64 (op==stock) so it passes trivially there. New test + gpu-op-nullable-f64-test task.
Add AGG_AVG to the nullable safe slice (still default-off, single-table, UNGROUPED, single aggregate, no native-decode). A single avg(x) emits two internal metrics -- sum (m0) and count (m1) -- BOTH summed over the same pass-gated rows, so the existing validity fold excludes NULL-x rows from numerator AND denominator => avg over the non-NULL rows, exactly SQL avg(x). An all-NULL column -> count 0 -> ungrouped_count_m (== avg's m1) marks the cell SQL NULL. Covers both the int128 assembly path (avg of an int64-backed column) and the f64 path (avg(sqrt(x)) etc., NVIDIA-only). avg was previously thought unsupportable: it returned 0.0, which I attributed to a DuckDB avg-as-sum-double rewrite. EXPLAIN disproves that -- nullable avg(x) is a plain `avg(#0)` aggregate, identical to the NOT-NULL plan. The 0.0 was the sum/avg warm-pin signature COLLISION (avg(x) reused a preceding sum(x)'s resident pin, which has M=1 and no count metric), already fixed in d168cdf. With distinct signatures, avg routes and computes correctly. Validated on Apple + RTX 4090: 5 q*_shuttle + raw_plan_roundtrip OFF+ON = 12/12 each; nullable_sql_test -- avg(b), avg(b) WHERE f>50, avg(all-NULL)->NULL all route (shadow kind=1) + match the CPU reference; nullable_f64_sql_test on the 4090 -- avg(sqrt(x)) rel-err 1.4e-14 vs stock. count(*), multi-agg, grouped, joins still decline -> stock.
… (int128) The int128 DENSE_GROUP path emitted a PHANTOM row for any group whose rows ALL fail the WHERE filter: the dense gid is built over ALL materialized rows (materialize_sql has no WHERE -- the filter is host-pass-baked), so a fully-filtered group still gets a gid, and _assemble's emit loop emits a row for EVERY gid (KIND_Q1 sets emit_agg=-1, so no gate fires). SQL forms groups AFTER filtering and OMITS such a group. Verified (bench/q1_dense_filter_phantom_probe.mojo): a Q1-shaped query whose filter eliminates one (returnflag,linestatus) group returned 3 rows incl. phantom (N,O) count_order=0; stock returned 2. Default-on; TPC-H Q1's broad cutoff keeps all 4 groups so the benchmark masks it. Fix (mirrors the already-correct f64 dense path, which gates on a per-group passing count): add GpuPinned.grouped_count_m -- the metric index of a per-group FILTER-passing count, REUSED from an existing count(*) / AVG-count metric for a STRAT_DENSE_GROUP int128 query (M unchanged; Q1 has count_order). _assemble skips any dense group whose grouped_count_m metric is 0. grouped_count_m<0 (no count metric / non-dense) keeps legacy behavior; Q5's emit_gt0 revenue gate is untouched (it goes through a separate finalize where grouped_count_m stays -1). No kernel change (the dense kernels are generic over M). A guaranteed emit for a count-less dense shape is deferred to the grouped-nullable follow-up. Validated Apple + RTX 4090: phantom probe now returns 2 rows == stock; 5 q*_shuttle + raw_plan_roundtrip PASS (Q1 with all groups populated unaffected); nullable_sql_test PASS.
…+ multi-aggregate
Generalize the nullable slice from a single aggregate to count(*) + N aggregates over
DIFFERENT nullable columns, by separating FILTER-pass from per-column VALIDITY:
- The host pass column folds in only FILTER-column validity (a NULL filter row =>
predicate UNKNOWN => excluded from existence + count(*) + every metric).
- Each INT metric multiplies in its own agg-input-column validity (OP_LOAD_COL(valid)
OP_MUL appended per distinct nullable input; value*0 zeroes a NULL row, value*1 is
identity). count(*) is left UNMULTIPLIED so it counts ALL filter-passing rows.
- Per-aggregate NULL-on-empty: each SUM/AVG/stat gets a validity-PRODUCT count metric
(Σ valid_a[*valid_b..]); _assemble / _assemble_f64 mark THAT aggregate's cell SQL
NULL when its count is 0 (e.g. sum(x) all-NULL among filter-passing rows -> NULL),
independent of other aggregates. AVG's denominator IS that validity-product count.
- Routing: an UNGROUPED int multi-aggregate classifies KIND_UNKNOWN; a new accessor
mojo_gpu_desc_a1_ungrouped_ok (ungrouped, all {SUM,AVG,count(*)}, not f64) lets
TryRouteGeneric route it under GPU_OP_NULLABLE. (Single-agg int still routes via
KIND_Q6; f64 via is_transcendental/is_stats.)
f64 path keeps the validity-IN-PASS fold (ce3e7db): a metric tape evaluating a
transcendental on a NULL row's garbage (sqrt(negative) -> domain error) BEFORE a
multiply could zero it must be avoided, so NULL rows are excluded from the kernel
entirely. Pass-folding cannot give per-column exclusion across multiple f64
aggregates, so f64 nullable is restricted to a SINGLE aggregate (the shipped scope).
C++ gate: drop aggregates.size()==1; accept N aggregates each in {int128 SUM, AVG,
count(*), stat, transcendental SUM}; f64-present => still single-agg.
Validated Apple + RTX 4090: 5 q*_shuttle + raw_plan_roundtrip OFF+ON 12/12 (Apple) /
6/6 (NVIDIA); nullable_sql_test (incl. sum/avg all-NULL -> SQL NULL); new
nullable_multiagg_sql_test (count(*),sum(b); sum(b),sum(c) over different nullable
columns; count(*),sum,avg; all routed via PIN_LOG, == stock); nullable_f64_sql_test on
the 4090 (sum(sqrt|ln), avg(sqrt), stddev/var/corr/regr over nullable -- the sqrt-on-
NULL-garbage domain error is fixed by the f64 validity-in-pass path). Flag OFF byte-identical.
…lable columns
Extend the nullable slice to a DENSE GROUP BY (NOT-NULL group key) over nullable
agg/filter columns, int path (count(*) + SUM/AVG). No new kernel: the A1 per-metric
validity multiply, the per-aggregate validity-count NULL marking, and the dense
filter-count existence gate are ALL per-group (g*M-indexed), so they already do the
right thing once a grouped shape routes.
- Existence: grouped_count_m now EMITS a guaranteed per-group filter-passing count when
no count(*)/AVG-count is reusable (a count-less `sum(x) GROUP BY k`), so a fully-
filtered group is OMITTED (matches stock; only the newly-routed grouped shapes get
the extra metric -- Q1 reuses count_order, Q5 uses a separate finalize).
- Per-group NULL: a group with filter-passing rows but all-NULL x emits SQL NULL (its
validity-product count is 0); a group's sum excludes only its own NULL-x rows.
- Routing: a DENSE int GROUP BY classifies KIND_UNKNOWN; new accessor
mojo_gpu_desc_a1_grouped_ok (DENSE, all {SUM,AVG,count(*)}, not f64) lets
TryRouteGeneric route it under GPU_OP_NULLABLE_GROUPED.
- Gate (SerializeMatchedPlan): accept a grouped nullable plan iff single-table, no
native-decode, int aggs, and EVERY group key is NOT NULL (a nullable group key would
form its own SQL NULL group -- unmodeled by the dense gid build -> declines). The
group-key NOT-NULL check uses the catalog (same storage-index + NOT_NULL-constraint
basis as any_nullable_projected).
Validated Apple + RTX 4090: 5 q*_shuttle + raw_plan_roundtrip across OFF /
GPU_OP_NULLABLE / +GROUPED (18/18 Apple, 6/6 NVIDIA); new nullable_grouped_sql_test --
per-group sum(x)+count(*); a FULLY-FILTERED group OMITTED; an ALL-NULL-x group -> SQL
NULL; per-group avg; and a NULLABLE GROUP KEY DECLINES (shadow=unsupported -> stock) --
all == stock. nullable_sql_test / multiagg / f64 still pass. Flags OFF byte-identical.
…for default-on) A 25-query operator-vs-stock sweep over the full GPU_OP_NULLABLE feature: varied column TYPE (DECIMAL 15,2 / 15,0 / 18,4 signed, BIGINT, all-NULL), NULL density, sign, and zero; SUM/AVG/count(*)/multi-agg over different nullable columns; range/eq/between/ empty/self filters; DENSE GROUP BY (incl. all-NULL group -> NULL); and f64 transcendental/stats. 0 mismatch on BOTH Apple and RTX 4090 (f64 routes only on NVIDIA; declines == stock on Apple). Due-diligence + regression guard ahead of any default-on flip.
…DEFAULT-ON The nullable feature (A1 unified pass model: int SUM/AVG/count(*)/multi-agg + f64 transcendental/stats single-agg + DENSE GROUP BY over nullable columns) is now default-on with opt-out (off/0/none), mirroring GPU_OP_TRANSCENDENTAL/GPU_OP_STATS. The operator now fires on real-world (nullable) tables automatically instead of declining the whole offload. Every nullable shape outside the validated slice still fail-closes at the gate -> stock (nullable group key, f64 multi-agg, IS NULL filters, native-decode), and the per-metric validity multiply / per-group NULL-on-empty / dense filter-count existence gate make the routed shapes bit-exact vs stock. Validated Apple + RTX 4090 after the flip: nullable_adversarial_sweep (25 varied shapes) relying on the DEFAULTS = 0 mismatch on both; default-on ROUTES confirmed via GPU_OP_SHADOW (routed kinds, not unsupported); opt-out (=off) declines every nullable query -> stock (byte-identical to pre-feature); 5 q*_shuttle + raw_plan_roundtrip 12/12 under both default and opt-out; NOT-NULL phantom-group probe still MATCH; f64 sweep on the 4090.
…bmodule Toolchain bump: - Switch Mojo from the max-nightly 1.0.0b2.dev nightly to the stable 1.0.0b2 release on the conda.modular.com/max channel, across pixi.toml (default + operator-replacement feature), both conda recipes, and the CI channel. No source changes were needed; the full test suite + both extension suites pass. - Bump DuckDB to 1.5.4 (libduckdb-devel/duckdb-cli, recipe run ranges, duckdb-from-source, and the CPP-ABI extension footer-version defaults). Regenerate duckdb/_libduckdb.mojo (only a one-line docstring delta vs 1.5.3). Vendor DuckDB source as a git submodule: - Add third_party/duckdb pinned to v1.5.4 (shallow), replacing the clone_duckdb.sh shallow clone of .duckdb-src. It is the single source of truth for code generation, the duckdb-from-source build (now a path source), and DuckDB's benchmark_runner. clone-duckdb now inits the submodule and warns on version drift; the C++ extensions still build against conda's libduckdb-devel headers by default. CI inits the submodule in the check-generated-api job. operator-replacement -> reference: - Document it as a reference implementation superseded by mojo-kernel-overrides (catalog mutation for built-ins) and mojo-gpu-operator (the same OptimizerExtension interception, for GPU offload); capture the unique cast-bridge + re-bind technique. Kept wired into the full env. mojo-gpu-operator: - Drop the gpu-op-* probe/kernel-test/shuttle/nullable/bench pixi tasks (keep gpu-op-build/clean); update docs + bench comments to run the bench/ files directly with mojo run.
sbrunk
force-pushed
the
mojo-simd-kernels
branch
from
June 25, 2026 17:20
804f1de to
5fdd530
Compare
…rrides on the decline path The GPU operator (OptimizerExtension) and mojo-kernel-overrides (catalog mutation) compose orthogonally: anything the operator declines runs through DuckDB's executor, so installing the override catalog mutation makes declined queries use the Mojo SIMD kernels (min/max, transcendental, INT128 sum) before stock — turning "decline = parity" into "decline = still beat DuckDB" where the SIMD kernels win. - build.sh compiles the override SIMD kernels (mojo-kernel-overrides/src/capi_shim.mojo, a plain object with no Mojo-runtime deps) + mojo_overrides.cpp straight into the operator .so, so a single LOAD can install both. - gpu_operator.cpp: new GpuOpOverridesOn() flag (GPU_OP_OVERRIDES, DEFAULT-OFF / presence-only). LoadInternal + register_mojo_gpu_operator call RegisterMojoOverrides(db) only when set, so flag-off is byte-identical (catalog untouched; the GPU-accept paths are rewritten to GPU before catalog functions run, so they're unaffected either way). Validated on Apple + RTX 4090: with GPU_OP_OVERRIDES=1, a GPU-declined query (ungrouped min/max/sum/transcendental) returns bit-identical to stock, and the GPU accept path still routes (sum(sqrt(d)) -> [gpu-shadow] kind=1) with overrides installed; flag-off == flag-on == stock on both. The kernel/descriptor paths are untouched, so the shuttle gate (green post-bump on the 4090) is not implicated by this extension-init change. Default-on flip and the cold-INT128 hybrid (NR1 part B) are deferred follow-ups.
…D), bit-exact via eval_program reuse On the COLD finalize path for an UNGROUPED int aggregate, compute the per-metric int128 sums on the CPU instead of the GPU kernel->readback roundtrip, while the resident buffers are still uploaded so the next (WARM) query runs on GPU. This removes the catastrophic cold-first-query GPU loss for generic scalar-sum shapes. Bit-exactness is by construction: the host reuses the kernel's OWN `eval_program` (the same expression VM segreduce runs per row), accumulating the int64 per-row metric values into int128 over the SAME packed host cols + the SAME pass program. The result feeds `_assemble` via a new `cpu_sums` param which, when non-empty, is used INSTEAD of `segreduce_run` — so all the downstream res_lo/res_hi unpack and ungrouped_count_m NULL-on-empty logic is reused verbatim. Restricted scope (eligibility, else -> normal GPU path): - GPU_OP_COLD_HYBRID set (DEFAULT-OFF, presence-only -> flag-off is byte-identical: the getenv guard leaves cpu_sums empty -> _assemble runs the kernel as before). - STRAT_UNGROUPED, not float64, M>0. - HOST-baked / pass-program filter only: NOT the q6/gen/f64 in-kernel-predicate residency shapes (those zero the host pass column), so the Q6 residency path is untouched. - Plain upload path only (not _colpool_on): the host `cols` are fully packed (no skip-materialize / col-pointer omission). This generic finalize has no FK dims. Covers int64-backed sums too (cpu_sums are Int128, matching segreduce_run's output for any int agg), but the int128/HUGEINT case is where the win lives. Validated bit-exact on Apple + RTX 4090: or_filter_shuttle_test fires the hybrid (pin-log) and is ALL PASS; the full standard shuttle gate (q1/q3/q5/q6/q14 + roundtrip) is ALL PASS with the flag on (grouped/dim shapes skip the hybrid -> kernel path unchanged); a live A/B (sum(x::BIGINT) WHERE a IN (2,9,40), OR filter) routes + fires the hybrid + matches stock. Deferred: a measured cold A/B + parallelizing the host eval loop (the scalar per-row eval_program likely needs parallelize / async-overlap with the upload to beat the GPU cold roundtrip); then the default-on flip.
…LD_HYBRID)" This reverts commit 7e74377 (NR1 part B). The cold INT128 hybrid was bit-exact but a cold A/B on the RTX 4090 (6M-row ungrouped sum + OR filter) falsified its premise: stock 5ms | GPU-cold 198ms | hybrid 234ms | hybrid+parallel-pack 149ms Stage breakdown (pin-log): materialize (separate SQL scan + feed) ~100ms + pack/pass ~92ms serial / ~22ms parallel DOMINATE the cold cost. The GPU reduce is cheap (~few ms); the CPU parallel reduce (16 workers, the per-row eval_program VM) is ~24ms -- SLOWER than the GPU roundtrip it replaces. So the hybrid adds cost on the reduce and leaves the dominant materialize+pack tax untouched -> it wins no regime (slightly loses), even parallelized. The real cold lever is the materialize+pack tax (the demoted "cold-tax block": R3 parallelize-finalize / R4 native-decode), not the GPU reduce. Recorded in PERF_BACKLOG.md so it is not re-attempted. NR1 part A (GPU_OP_OVERRIDES, 397a8bb) is unaffected and stays.
…/cosine/inner_product) Override DuckDB's array_distance, array_inner_product, array_negative_inner_product, array_cosine_similarity, array_cosine_distance (FLOAT[]/DOUBLE[]). Stock folds are serial single-accumulator loops (result += x*y) that can't auto-vectorize (FP reduction) — latency bound. Three Mojo kernels in duckdb/kernels/simd.mojo (array_dot/array_l2dist/ array_cosine_sim, W-wide SIMD accumulators + 2x unroll to break the dependency chain) back all five functions; the C++ wrapper MojoArrayFold replicates ArrayGenericFold exactly (NULL-row -> NULL, NULL element -> throw, constant-vector result) and swaps only the inner reduction. OverrideArrayFold swaps f.function on the FLOAT/DOUBLE overloads, like OverrideScalar. Raw fold 7.4-12.4x (768-d f32 cosine 10.8x; stock 4.8 GB/s latency-bound -> kernel ~52 GB/s, memory-bandwidth bound). End-to-end ORDER BY ... LIMIT k scan ~1.6x (scan + top-k dominate). Extension .so links only libSystem(libm) + libc++ — no Modular runtime (two-extensions distribution basis). Adds: - pixi run overrides-test (stock-vs-override correctness: ~1 ULP across f32/f64 + tail sizes 13/17/100/128, NULL semantics, known values) - benchmark/micro/mojo_simd/array_*.benchmark (topk + sum-isolating) - CPU_SIMD_BACKLOG.md (lessons ported from the GPU operator)
Lift the all-valid restriction: FLAT columns with NULLs are now handled in place instead of falling back to stock. Masked SIMD reductions in duckdb/kernels/simd.mojo (reduce_sum_f64_masked, generic reduce_minmax_masked[dt,w,is_min], reduce_sum_i128_masked) reduce only the valid lanes — expand the validity word to a lane mask (iota + .gt()), select(mask, x, identity), accumulate, and return the valid count so AVG divides correctly and SUM/MIN/MAX leave the state unset on an all-NULL chunk. Wired into every aggregate wrapper (all-valid fast path unchanged; FLAT-with-NULLs takes the masked kernel). Scalar funcs (sqrt/sin/...) compute over all rows then copy the input validity mask to the result. Grouped aggregates still use stock. Nullable sum/avg: ~1.75x over stock (50M f64, 30% NULL, 1 thread) and now covered instead of declined. min/max correct (DuckDB answers base-column min/max from zonemaps anyway). overrides-test extended: stock-vs-override over DOUBLE/REAL/HUGEINT/DECIMAL(38,4), sum(sqrt) over a nullable column, and an all-NULL column -> NULL. Max rel diff 3.6e-14 (summation order); HUGEINT/DECIMAL exact.
…e function
Add a register-tiled CPU SIMD brute-force kNN: knn_topk kernel in
duckdb/kernels/simd.mojo (QT=4 queries per embedding load, raising arithmetic
intensity since the per-pair dot is L1/L2 read-bound), exported as
mojo_knn_{cosine,l2,ip}_f32, and a table function mojo_knn(emb_table, emb_col,
query_table, query_col, k, metric:=) registered in the catalog (works from LOAD
and the dlopen register entry).
DuckDB has no batch-kNN primitive, so the SQL alternative (cross-join + window
over M*N distance rows) is pathological: mojo_knn is ~89x faster (M=64, N=100k,
D=128, 1 thread); the kernel itself is ~2.8x (AVX-512) over naive optimized scans.
Exact ranking: recall@k=1.0 vs stock for cosine/l2/ip, self-match correct
(validated in overrides-test). l2/cosine use the |a|^2+|b|^2-2dot identity, so a
self/duplicate pair can read ~1e-3 instead of 0 (benign cancellation; ordering
unaffected). Dependency-free.
…er rewrite
Add an OptimizerExtension that transparently rewrites ungrouped sum(f(col)) /
avg(f(col)) for f in {sqrt,sin,cos,ln,exp,log10} (DOUBLE) into a custom aggregate
that applies f and reduces in one pass (no intermediate vector). Fused kernels
reduce_fsum_map[f] / reduce_fsum_map_masked[f] in duckdb/kernels/simd.mojo
(plain + A1-masked for nullable), exported as mojo_fsum_<f>_f64[_masked], wired to
templated FusedAgg<plain,masked,scalar,is_avg> aggregates built in the rewrite.
Modest win (~1.13x over the already-SIMD two-pass; ~2.25x over stock scalar) —
fusion is mostly a GPU advantage since the CPU intermediate is L1-resident and the
transcendental compute dominates. The value here is also the optimizer-rewrite
infrastructure, which item 5 reuses to route shapes to CPU/GPU backends.
Ungrouped only; nullable handled (masked); domain errors follow the scalar
fast-path caveat (NaN not throw). Opt out with MOJO_OVERRIDES_NO_FUSE=1. Validated
in overrides-test: sum/avg of all six transcendentals vs stock (max rel 6e-12),
nullable, empty->NULL, and rewrite-fires check.
Add a cardinality crossover so the optimizer declines small-N matched shapes to the co-installed CPU SIMD tier (GPU_OP_OVERRIDES) or stock, where GPU pin/transfer overhead would otherwise lose. GpuOpMinRows() (default 50000; 0 disables) + BelowGpuCrossover(node) = max child estimated_cardinality < min, gated in TryRouteGeneric (before LogicalGpuAgg) and the cosine route. Validated on Mac (Apple GPU) and frederick (RTX 4090) via EXPLAIN: cosine routes GPU above the threshold, CPU below; composition with GPU_OP_OVERRIDES returns stock-identical results.
Replace the scattered per-package benchmark scripts/groups with one home and one driver: - benchmark/sql/<group>/*.benchmark — DuckDB benchmark_runner files (mojo_simd moved from mojo-kernel-overrides; new gpu_xover, gpu_knn). - benchmark/drivers/bench_runner.sh — one compare driver (--engines toggle mode + --by-suffix per-file mode), replacing run_runner_compare / gpu_warm_sweep / knn_warm_sweep. build_runner.sh + runner_load_extension.patch moved here. - benchmark/mojo/knn_compare.mojo (+ bench_util.mojo) — Mojo client vector-search harness: single + batch cosine top-k across stock / cpu-simd / vss-HNSW / GPU, latency + recall. - benchmark/README.md — index. pixi tasks collapse to bench-build / bench-sql / bench-knn (replacing overrides-bench-runner*, knn-bench*, gpu-xover-bench). build_runner now stages all benchmark/sql groups and links the runner with -rdynamic (so it re-exports DuckDB RTTI for the dlopen'd CPP-ABI extensions). run_gpu_op_bench.sh + the package READMEs + AGENTS.md repointed to the new tasks. Validated both driver modes + bench-knn on Mac and frederick.
sbrunk
force-pushed
the
mojo-simd-kernels
branch
from
June 30, 2026 16:34
c8db2da to
8694c83
Compare
Main README: add "Accelerate DuckDB" to the overview, document the full mojo-kernel-overrides surface (array distance, nullable, fused transcendental, mojo_knn) and the general-purpose mojo-gpu-operator (offload + vector search), and point benchmarking at the consolidated bench-build/bench-sql/bench-knn harness. Add a "CPP-ABI extensions" subsection: a CPP extension is a C++ extension calling Mojo @export abi("C") kernels, with the three entry points, the two link models, and the version-lock / internal-headers / unsigned / -rdynamic caveats. append_extension_metadata.py: document the CPP --abi-type (already in choices) in the docstrings/help.
Replace bench_runner.sh with bench_runner.py: same behavior (toggle --engines and --by-suffix modes, staging benchmark/sql/* into the runner tree, median per benchmark) but far more readable than the bash + awk + inline-python it replaced. The runner already shelled out to python for the table, so this removes a language boundary. argparse + statistics.median; per-engine env overlaid on os.environ so caller-exported GPU_OP_* flags still reach the runner. pixi bench-sql, run_gpu_op_bench.sh, and benchmark/README.md repointed.
sbrunk
force-pushed
the
mojo-simd-kernels
branch
from
June 30, 2026 18:40
f03913d to
d43ccf5
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds two new CPP-ABI DuckDB extensions backed by Mojo kernels, bumps the toolchain/DuckDB pins, and consolidates the benchmark harness.
Main changes
Toolchain & DuckDB
1.0.0b2stable and DuckDB1.5.4.third_party/duckdbgit submodule (single source for code-gen, source build, andbenchmark_runner);clone-duckdbwarns on pin drift, CIcheck-generated-apiguards stale bindings.packages/mojo-gpu-operatorfor transparent GPU offloadGPU_OP_NULLABLE/_GROUPED, default-on) via unified A1 pass model.GPU_OP_MIN_ROWS) co-installs the SIMD overrides on the decline path.packages/mojo-kernel-overrides: SIMD built-in overridessqrt/sin/cos/ln/exp/log10andsum/avg/min/maxwith Mojo SIMD kernels, stock fallback for non-FLAT/null/grouped.HUGEINT/DECIMAL(19..38)sum/avg.array_distance/cosine/inner_product), fusedsum/avg(transcendental)optimizer rewrite, A1 nullable mask-multiply, andmojo_knnblocked multi-query kNN table function.Benchmarks
benchmark/{sql,drivers,mojo}withbench-build/bench-sql/bench-knntasks; compare driver rewritten in Python (bench_runner.py).benchmark_runnerfrom the submodule with a ~13-lineLOAD-extension hook (no libduckdb fork).