Skip to main content

xlog_cuda/provider/
fj.rs

1//! GPU Free Join provider: level-synchronous factorized join execution.
2//!
3//! Design: `docs/plans/2026-06-12-d2-free-join-design.md`. The paper's
4//! (Wang/Willsey/Suciu, SIGMOD 2023) depth-first recursion over lazy
5//! hash tries is replaced by:
6//!
7//!   * **Flat sorted-range tries** (§2.1): every input is
8//!     layout-normalized (lex-sorted + deduped, the existing WCOJ
9//!     layout); a trie node is a contiguous `[lo, hi)` row range and
10//!     `get(key)` is a binary-search refinement of that range on the
11//!     next column. No per-level structure is ever built.
12//!   * **Level-synchronous frontier execution** (§2.2): a bindings
13//!     frontier (SoA u32 columns: bound variables plus per-live-atom
14//!     `(lo, hi)` range pairs) is rebuilt per plan node by bulk
15//!     two-phase EXPAND (count → device scan → emit) over the node's
16//!     cover subatom, followed by one PROBE refinement kernel per
17//!     probe subatom and a single mask compaction (reusing the
18//!     existing mask + scan + gather kernels).
19//!
20//! Invariants (§2.3, non-negotiable):
21//!   * all inputs layout-normalized per dispatch;
22//!   * no atomics in any emit path — output positions come from
23//!     exclusive scans, so output order is deterministic (parent-row
24//!     order × lex order of the plan's variable sequence);
25//!   * set semantics by construction: distinct cover groups over
26//!     deduped inputs keep the frontier duplicate-free; a final dedup
27//!     runs only when the head projects a strict subset of the bound
28//!     variables;
29//!   * zero tracked transfers — host reads are limited to the
30//!     sanctioned `dtoh_scalar_untracked` metadata scalars (scan
31//!     totals, compaction counts); recorded launches throughout.
32//!
33//! Width classes: u32/Symbol (`free_join_execute_u32_recorded`) and
34//! u64 (`free_join_execute_u64_recorded`) share one
35//! width-parameterized pipeline — frontier VAR columns carry
36//! width-sized data values while RANGE columns are u32 row indices in
37//! every width class (the staging/compaction/projection helpers are
38//! schema-driven per column, so mixed-width frontiers need no special
39//! casing). Full expansion runs at the last node (the factorized
40//! trailing-range enumeration of §2.4 remains future work).
41
42use std::ffi::c_void;
43
44use xlog_core::{Result, ScalarType, Schema, XlogError};
45
46use super::{wcoj_kernels, CudaKernelProvider, WCOJ_MODULE};
47use crate::device_runtime::StreamId;
48use crate::launch::LaunchRecorder;
49use crate::memory::{CudaColumn, TrackedCudaSlice};
50use crate::CudaBuffer;
51use crate::{AsKernelParam, LaunchAsync, LaunchConfig};
52
53const BLOCK_SIZE: u32 = 256;
54
55/// One subatom: an atom (`input_idx`) restricted to the variables its
56/// next `var_positions.len()` physical columns bind/probe. Across the
57/// whole plan, each atom's subatoms consume its columns in order and
58/// must partition them exactly (design §3).
59#[derive(Debug, Clone)]
60pub struct FjSubAtom {
61    /// Index into the `inputs` slice of
62    /// [`CudaKernelProvider::free_join_execute_u32_recorded`].
63    pub input_idx: usize,
64    /// Global variable ids bound (cover) or matched (probe) by this
65    /// subatom's columns, in column order.
66    pub var_positions: Vec<usize>,
67}
68
69/// One plan node: iterate the cover subatom (bulk EXPAND over the
70/// whole frontier), then refine every probe subatom (PROBE +
71/// compaction). Probe variables must already be bound.
72#[derive(Debug, Clone)]
73pub struct FjNode {
74    pub cover: FjSubAtom,
75    pub probes: Vec<FjSubAtom>,
76}
77
78/// A host-side Free Join plan over `inputs` (design §3). Callers hand-build
79/// the plan today; planner construction from binary joins (`binary2fj`) is a
80/// downstream integration surface.
81#[derive(Debug, Clone)]
82pub struct FjPlan {
83    /// Number of distinct join variables (ids `0..num_vars`).
84    pub num_vars: usize,
85    /// Plan nodes, executed in order.
86    pub nodes: Vec<FjNode>,
87    /// Head projection: variable ids in output column order. When
88    /// this is a strict subset of the bound variables the result is
89    /// deduplicated (set semantics).
90    pub output_vars: Vec<usize>,
91}
92
93/// Frontier column tag: which logical value a SoA column holds. The
94/// live-range set is tracked statically per node (design §5.1) —
95/// exhausted atoms' range columns are dropped, untouched atoms have
96/// no columns at all (their range is the constant `[0, n)`).
97#[derive(Debug, Clone, Copy, PartialEq, Eq)]
98enum ColTag {
99    /// Bound value of a join variable.
100    Var(usize),
101    /// Current trie-range lower bound for an atom.
102    RangeLo(usize),
103    /// Current trie-range upper bound for an atom.
104    RangeHi(usize),
105}
106
107type FrontierCol = (ColTag, TrackedCudaSlice<u8>);
108
109/// Data width class of a Free Join execution. Row indices, trie
110/// ranges, work prefixes, and group marks are u32 in every class;
111/// only DATA columns (cover/probe/var values) take this width.
112#[derive(Debug, Clone, Copy, PartialEq, Eq)]
113enum FjWidth {
114    U32,
115    U64,
116}
117
118impl FjWidth {
119    fn var_bytes(self) -> usize {
120        match self {
121            Self::U32 => std::mem::size_of::<u32>(),
122            Self::U64 => std::mem::size_of::<u64>(),
123        }
124    }
125
126    fn var_type(self) -> ScalarType {
127        match self {
128            Self::U32 => ScalarType::U32,
129            Self::U64 => ScalarType::U64,
130        }
131    }
132
133    fn count_kernel(self) -> &'static str {
134        match self {
135            Self::U32 => wcoj_kernels::FJ_EXPAND_COUNT_U32,
136            Self::U64 => wcoj_kernels::FJ_EXPAND_COUNT_U64,
137        }
138    }
139
140    fn emit_kernel(self) -> &'static str {
141        match self {
142            Self::U32 => wcoj_kernels::FJ_EXPAND_EMIT_U32,
143            Self::U64 => wcoj_kernels::FJ_EXPAND_EMIT_U64,
144        }
145    }
146
147    fn probe_kernel(self) -> &'static str {
148        match self {
149            Self::U32 => wcoj_kernels::FJ_PROBE_REFINE_U32,
150            Self::U64 => wcoj_kernels::FJ_PROBE_REFINE_U64,
151        }
152    }
153}
154
155/// Per-tag column type within a frontier of the given width: VAR
156/// columns are width-sized data, RANGE columns are u32 row indices.
157fn tag_type(tag: ColTag, width: FjWidth) -> ScalarType {
158    match tag {
159        ColTag::Var(_) => width.var_type(),
160        ColTag::RangeLo(_) | ColTag::RangeHi(_) => ScalarType::U32,
161    }
162}
163
164/// What the pipeline produces after the last plan node.
165#[derive(Debug, Clone, Copy, PartialEq, Eq)]
166enum FjMode {
167    /// Materialize the projected row set (`output_vars` columns).
168    Materialize,
169    /// Design §2.4 factorized count: reduce to
170    /// `(output_vars[0], count)` where each frontier row contributes
171    /// the PRODUCT of its remaining live trie-range lengths
172    /// (unconsumed trailing columns never expand). Plans may
173    /// partially consume atoms, but every atom must be touched.
174    CountByRoot,
175}
176
177fn owned_col_ptr(buf: &CudaBuffer, idx: usize, ctx: &str) -> Result<u64> {
178    match buf.column(idx) {
179        Some(CudaColumn::Owned(s)) => Ok(*s.device_ptr()),
180        Some(_) => Err(XlogError::Kernel(format!(
181            "{ctx}: input column {idx} must be owned"
182        ))),
183        None => Err(XlogError::Kernel(format!(
184            "{ctx}: input column {idx} not found"
185        ))),
186    }
187}
188
189fn find_col<'a>(
190    cols: &'a [FrontierCol],
191    tag: ColTag,
192    ctx: &str,
193) -> Result<&'a TrackedCudaSlice<u8>> {
194    cols.iter()
195        .find(|(t, _)| *t == tag)
196        .map(|(_, s)| s)
197        .ok_or_else(|| XlogError::Kernel(format!("{ctx}: frontier column {tag:?} missing")))
198}
199
200/// Validate the plan against the input arities. Returns the variable
201/// binding order (covers' variables in plan order).
202fn validate_plan(plan: &FjPlan, arities: &[usize], mode: FjMode, ctx: &str) -> Result<Vec<usize>> {
203    if plan.nodes.is_empty() {
204        return Err(XlogError::Kernel(format!("{ctx}: plan has no nodes")));
205    }
206    let mut bound = vec![false; plan.num_vars];
207    let mut bind_order: Vec<usize> = Vec::new();
208    let mut consumed = vec![0usize; arities.len()];
209    let check_sub = |sub: &FjSubAtom, consumed: &[usize], what: &str| -> Result<()> {
210        if sub.input_idx >= arities.len() {
211            return Err(XlogError::Kernel(format!(
212                "{ctx}: {what} input_idx {} out of bounds ({} inputs)",
213                sub.input_idx,
214                arities.len()
215            )));
216        }
217        if sub.var_positions.is_empty() {
218            return Err(XlogError::Kernel(format!(
219                "{ctx}: {what} on input {} has no variables",
220                sub.input_idx
221            )));
222        }
223        if consumed[sub.input_idx] + sub.var_positions.len() > arities[sub.input_idx] {
224            return Err(XlogError::Kernel(format!(
225                "{ctx}: {what} over-consumes input {} (arity {}, consumed {}, +{})",
226                sub.input_idx,
227                arities[sub.input_idx],
228                consumed[sub.input_idx],
229                sub.var_positions.len()
230            )));
231        }
232        for &v in &sub.var_positions {
233            if v >= plan.num_vars {
234                return Err(XlogError::Kernel(format!(
235                    "{ctx}: {what} variable {v} out of bounds (num_vars {})",
236                    plan.num_vars
237                )));
238            }
239        }
240        Ok(())
241    };
242    for (k, node) in plan.nodes.iter().enumerate() {
243        check_sub(&node.cover, &consumed, "cover")?;
244        for &v in &node.cover.var_positions {
245            if bound[v] {
246                return Err(XlogError::Kernel(format!(
247                    "{ctx}: node {k} cover rebinds variable {v}"
248                )));
249            }
250            bound[v] = true;
251            bind_order.push(v);
252        }
253        consumed[node.cover.input_idx] += node.cover.var_positions.len();
254        let mut seen_atoms = vec![node.cover.input_idx];
255        for probe in &node.probes {
256            check_sub(probe, &consumed, "probe")?;
257            if seen_atoms.contains(&probe.input_idx) {
258                return Err(XlogError::Kernel(format!(
259                    "{ctx}: node {k} touches input {} more than once",
260                    probe.input_idx
261                )));
262            }
263            seen_atoms.push(probe.input_idx);
264            for &v in &probe.var_positions {
265                if !bound[v] {
266                    return Err(XlogError::Kernel(format!(
267                        "{ctx}: node {k} probes unbound variable {v}"
268                    )));
269                }
270            }
271            consumed[probe.input_idx] += probe.var_positions.len();
272        }
273    }
274    for (i, (&used, &arity)) in consumed.iter().zip(arities.iter()).enumerate() {
275        match mode {
276            FjMode::Materialize => {
277                if used != arity {
278                    return Err(XlogError::Kernel(format!(
279                        "{ctx}: plan consumes {used}/{arity} columns of input {i} \
280                         (materialization requires full consumption)"
281                    )));
282                }
283            }
284            // §2.4 factorized counting: unconsumed trailing columns
285            // contribute their live range lengths as multiplicities,
286            // but an untouched atom has no range to read — reject.
287            FjMode::CountByRoot => {
288                if used == 0 {
289                    return Err(XlogError::Kernel(format!(
290                        "{ctx}: count plan never touches input {i} \
291                         (untouched atoms have no live range)"
292                    )));
293                }
294            }
295        }
296    }
297    if plan.output_vars.is_empty() {
298        return Err(XlogError::Kernel(format!("{ctx}: empty output_vars")));
299    }
300    if mode == FjMode::CountByRoot && plan.output_vars.len() != 1 {
301        return Err(XlogError::Kernel(format!(
302            "{ctx}: count plans take exactly one output (group) variable, got {}",
303            plan.output_vars.len()
304        )));
305    }
306    for &v in &plan.output_vars {
307        if v >= plan.num_vars || !bound[v] {
308            return Err(XlogError::Kernel(format!(
309                "{ctx}: output variable {v} is never bound"
310            )));
311        }
312    }
313    Ok(bind_order)
314}
315
316impl CudaKernelProvider {
317    /// Execute a hand-built Free Join plan over u32/Symbol relations
318    /// via the level-synchronous frontier engine. See the module docs
319    /// for the algorithm and invariants; the plan contract is
320    /// documented on [`FjPlan`].
321    ///
322    /// Inputs are layout-normalized per dispatch (sorted + deduped via
323    /// the existing WCOJ layout entries — already-normalized inputs
324    /// take the recorded fast-path check). The output contains one
325    /// column per `output_vars` entry (all `U32`) holding the join's
326    /// projected row set under set semantics.
327    ///
328    /// # Errors
329    /// * `XlogError::Kernel` if the manager has no runtime, the launch
330    ///   stream does not resolve, an input violates the u32
331    ///   width-class layout contract, the plan is invalid (unbound
332    ///   probe variables, over/under-consumed atom columns, rebound
333    ///   variables, unknown output variables), the frontier exceeds
334    ///   the u32 work-index space, or any kernel launch fails.
335    pub fn free_join_execute_u32_recorded(
336        &self,
337        inputs: &[&CudaBuffer],
338        plan: &FjPlan,
339        launch_stream: StreamId,
340    ) -> Result<CudaBuffer> {
341        self.free_join_execute_recorded_impl(
342            inputs,
343            plan,
344            launch_stream,
345            FjWidth::U32,
346            FjMode::Materialize,
347            "free_join_execute_u32_recorded",
348        )
349    }
350
351    /// u64 width-class twin of [`Self::free_join_execute_u32_recorded`]:
352    /// identical pipeline, contract, and invariants; every
353    /// input column must be `U64` and the output columns are `U64`.
354    pub fn free_join_execute_u64_recorded(
355        &self,
356        inputs: &[&CudaBuffer],
357        plan: &FjPlan,
358        launch_stream: StreamId,
359    ) -> Result<CudaBuffer> {
360        self.free_join_execute_recorded_impl(
361            inputs,
362            plan,
363            launch_stream,
364            FjWidth::U64,
365            FjMode::Materialize,
366            "free_join_execute_u64_recorded",
367        )
368    }
369
370    /// Design §2.4 factorized count-by-root over the Free Join
371    /// frontier: runs the same pipeline but reduces to
372    /// `(group, count)` instead of materializing rows. The plan's
373    /// `output_vars` must be exactly `[group_var]`; atoms may be
374    /// PARTIALLY consumed — each surviving frontier row contributes
375    /// the product of its remaining live trie-range lengths (the
376    /// d-representation count), so trailing private variables never
377    /// expand the frontier. Output schema: `(group: U32, count: U64)`.
378    ///
379    /// u32/Symbol width-class only: the reduction reuses the recorded
380    /// groupby, whose KEY columns are bounded engine-wide to
381    /// U32/Symbol (multi-type recorded sort is deferred there) — u64
382    /// bodies stay on the materialize path.
383    pub fn free_join_count_by_root_u32_recorded(
384        &self,
385        inputs: &[&CudaBuffer],
386        plan: &FjPlan,
387        launch_stream: StreamId,
388    ) -> Result<CudaBuffer> {
389        self.free_join_execute_recorded_impl(
390            inputs,
391            plan,
392            launch_stream,
393            FjWidth::U32,
394            FjMode::CountByRoot,
395            "free_join_count_by_root_u32_recorded",
396        )
397    }
398
399    #[expect(
400        clippy::too_many_lines,
401        reason = "the recorded free-join transaction keeps validation, allocation, launch recording, and commit in one auditable lifetime scope"
402    )]
403    fn free_join_execute_recorded_impl(
404        &self,
405        inputs: &[&CudaBuffer],
406        plan: &FjPlan,
407        launch_stream: StreamId,
408        width: FjWidth,
409        mode: FjMode,
410        ctx: &str,
411    ) -> Result<CudaBuffer> {
412        if self.memory().runtime().is_none() {
413            return Err(XlogError::Kernel(format!(
414                "{ctx} requires a runtime-backed GpuMemoryManager \
415                 (constructed via with_runtime)"
416            )));
417        }
418        if inputs.is_empty() {
419            return Err(XlogError::Kernel(format!("{ctx}: no inputs")));
420        }
421        let arities: Vec<usize> = inputs.iter().map(|b| b.arity()).collect();
422        let bind_order = validate_plan(plan, &arities, mode, ctx)?;
423
424        // Layout-normalize every input per dispatch. Arity-2 inputs
425        // go through the triangle-grade
426        // entry (it has the sorted+unique recorded fast-path);
427        // wider inputs use the generic full-row WCOJ sort+dedup entry.
428        let mut norm: Vec<CudaBuffer> = Vec::with_capacity(inputs.len());
429        for input in inputs {
430            let normalized = match (width, input.arity()) {
431                (FjWidth::U32, 2) => self.wcoj_layout_u32_recorded(input, launch_stream)?,
432                (FjWidth::U32, _) => self.wcoj_layout_sort_u32_recorded(input, launch_stream)?,
433                (FjWidth::U64, 2) => self.wcoj_layout_u64_recorded(input, launch_stream)?,
434                (FjWidth::U64, _) => self.wcoj_layout_sort_u64_recorded(input, launch_stream)?,
435            };
436            norm.push(normalized);
437        }
438        let mut n_rows: Vec<u32> = Vec::with_capacity(norm.len());
439        for buf in &norm {
440            let n = match buf.cached_row_count() {
441                Some(c) => c,
442                None => self.dtoh_scalar_untracked::<u32>(buf.num_rows_device(), 0)?,
443            };
444            n_rows.push(n);
445        }
446
447        let out_schema = match mode {
448            FjMode::Materialize => Schema::new(
449                plan.output_vars
450                    .iter()
451                    .map(|v| (format!("v{v}"), width.var_type()))
452                    .collect(),
453            ),
454            FjMode::CountByRoot => Schema::new(vec![
455                (format!("v{}", plan.output_vars[0]), width.var_type()),
456                ("count".to_string(), ScalarType::U64),
457            ]),
458        };
459        // Inner-join semantics: any empty atom empties the result.
460        if n_rows.contains(&0) {
461            return self.create_empty_buffer(out_schema);
462        }
463
464        let runtime = self.memory().runtime().ok_or_else(|| {
465            XlogError::Kernel(format!("{ctx} requires a runtime-backed GpuMemoryManager"))
466        })?;
467        let cu_stream = runtime
468            .stream_pool()
469            .resolve(launch_stream)
470            .ok_or_else(|| {
471                XlogError::Kernel(format!(
472                    "{ctx}: launch_stream StreamId({}) does not resolve",
473                    launch_stream.0
474                ))
475            })?;
476
477        // Frontier: starts as the single empty binding with every
478        // atom untouched (constant range [0, n)).
479        let mut frontier: Vec<FrontierCol> = Vec::new();
480        let mut count: u32 = 1;
481        // Column CAPACITY of the current frontier (the last node's
482        // n_children): compaction shrinks the logical count without
483        // reallocating, and every CudaBuffer built over frontier
484        // columns must carry row_cap == capacity (columns are
485        // row_cap × elem bytes by contract), with the logical count
486        // riding on num_rows_device.
487        let mut frontier_cap: u32 = 1;
488        let mut consumed = vec![0usize; inputs.len()];
489
490        for node in &plan.nodes {
491            let a = node.cover.input_idx;
492            let c = node.cover.var_positions.len();
493            let depth = consumed[a];
494            let cover_live = frontier.iter().any(|(t, _)| *t == ColTag::RangeLo(a));
495
496            // ---- EXPAND phase 0: total candidate work. Live cover
497            // atoms need a per-row work prefix (range lengths →
498            // exclusive scan); untouched cover atoms share the
499            // constant range, so the mapping is uniform and the
500            // total is known on host.
501            let (total_work, work_prefix) = if cover_live {
502                let mut wp = self.memory().alloc::<u32>(count as usize + 1)?;
503                let lo_col = find_col(&frontier, ColTag::RangeLo(a), ctx)?;
504                let hi_col = find_col(&frontier, ColTag::RangeHi(a), ctx)?;
505                let mut scan_scratch = self.multiblock_scan_u32_scratch_for_len(count + 1)?;
506                let mut rec = LaunchRecorder::new_strict(launch_stream);
507                rec.read(lo_col);
508                rec.read(hi_col);
509                rec.write(&wp);
510                for level in scan_scratch.levels() {
511                    rec.read_write(level);
512                }
513                rec.preflight(runtime)
514                    .map_err(|e| XlogError::Kernel(format!("{ctx}: wp preflight failed: {e}")))?;
515                let kernel = self
516                    .device()
517                    .inner()
518                    .get_func(WCOJ_MODULE, wcoj_kernels::FJ_EXPAND_WORK_PREFIX_U32)
519                    .ok_or_else(|| {
520                        XlogError::Kernel("fj_expand_work_prefix_u32 kernel not found".to_string())
521                    })?;
522                let grid = count.div_ceil(BLOCK_SIZE);
523                // SAFETY: fj_expand_work_prefix_u32(parent_lo,
524                // parent_hi, n_frontier, work_prefix); buffers are
525                // device-resident and preflighted.
526                let enqueue = |stream: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
527                    unsafe {
528                        kernel
529                            .clone()
530                            .launch_in(
531                                stream,
532                                LaunchConfig {
533                                    grid_dim: (grid, 1, 1),
534                                    block_dim: (BLOCK_SIZE, 1, 1),
535                                    shared_mem_bytes: 0,
536                                },
537                                (lo_col, hi_col, count, &mut wp),
538                            )
539                            .map_err(|e| {
540                                XlogError::Kernel(format!(
541                                    "fj_expand_work_prefix_u32 launch failed: {e}"
542                                ))
543                            })?;
544                    }
545                    self.multiblock_scan_u32_inplace_on_stream(
546                        &mut wp,
547                        count + 1,
548                        stream,
549                        &mut scan_scratch,
550                    )?;
551                    Ok(())
552                };
553                // SAFETY: preflight retained the accessed buffers and bound this stream.
554                let rec = unsafe { rec.enqueue_prepared_with(&cu_stream, enqueue) }
555                    .map_err(|error| error.into_xlog_error())?;
556                rec.commit()
557                    .map_err(|e| XlogError::Kernel(format!("{ctx}: wp commit failed: {e}")))?;
558                cu_stream
559                    .synchronize()
560                    .map_err(|e| XlogError::Kernel(format!("{ctx}: wp sync failed: {e}")))?;
561                let total = self.dtoh_scalar_untracked::<u32>(&wp, count as usize)?;
562                (u64::from(total), Some(wp))
563            } else {
564                ((count as u64) * (n_rows[a] as u64), None)
565            };
566            if total_work == 0 {
567                return self.create_empty_buffer(out_schema);
568            }
569            if total_work > u64::from(u32::MAX - 1) {
570                return Err(XlogError::Kernel(format!(
571                    "{ctx}: expansion work {total_work} exceeds the u32 work-index \
572                     space (frontier budget)"
573                )));
574            }
575            let total_work = total_work as u32;
576
577            // ---- EXPAND phase 1: mark distinct cover-prefix group
578            // starts and scan them into output offsets.
579            //
580            // Identity-group fast path: when the cover consumes through
581            // the atom's LAST column, rows within any trie range have
582            // distinct column suffixes (inputs are full-row deduped and
583            // the range fixes all preceding columns), so every candidate
584            // position is its own group — the marks pass, its device
585            // scan, and its host sync are skipped and
586            // n_children == total_work (the emit kernel takes its
587            // out == w branch via a null group_offsets pointer).
588            let identity = depth + c >= arities[a];
589            // ---- Fused-probe analysis: a probe folds into the count
590            // pass iff (a) its key variables are all bound by THIS
591            // node's cover (the kernel reads keys from cover_cols at
592            // the candidate position; earlier bindings are already
593            // encoded in the probe's carried range), and (b) it
594            // consumes through its atom's last column (existence-only —
595            // no refined range survives for later nodes). Fused probes
596            // skip the separate probe kernel and the mask compaction
597            // entirely: the emit pass materializes exactly the
598            // surviving children.
599            let is_fusable = |pr: &FjSubAtom| {
600                consumed[pr.input_idx] + pr.var_positions.len() >= arities[pr.input_idx]
601                    && pr
602                        .var_positions
603                        .iter()
604                        .all(|v| node.cover.var_positions.contains(v))
605            };
606            let fused: Vec<&FjSubAtom> = node.probes.iter().filter(|p| is_fusable(p)).collect();
607            // The count pass runs when groups need marking (non-identity
608            // cover) OR fused probes need evaluating.
609            let count_ran = !identity || !fused.is_empty();
610            // Pack one descriptor per fused probe, sequentially:
611            // [n_cols, has_range, in_lo_ptr, in_hi_ptr, n_atom_rows,
612            //  data_col_ptr * n_cols, cover_var_idx * n_cols].
613            let mut fused_desc: Vec<u64> = Vec::new();
614            for pr in &fused {
615                let p = pr.input_idx;
616                let live = frontier.iter().any(|(t, _)| *t == ColTag::RangeLo(p));
617                fused_desc.push(pr.var_positions.len() as u64);
618                fused_desc.push(u64::from(live));
619                if live {
620                    fused_desc.push(*find_col(&frontier, ColTag::RangeLo(p), ctx)?.device_ptr());
621                    fused_desc.push(*find_col(&frontier, ColTag::RangeHi(p), ctx)?.device_ptr());
622                } else {
623                    fused_desc.push(0);
624                    fused_desc.push(0);
625                }
626                fused_desc.push(u64::from(n_rows[p]));
627                for i in consumed[p]..consumed[p] + pr.var_positions.len() {
628                    fused_desc.push(owned_col_ptr(&norm[p], i, ctx)?);
629                }
630                for v in &pr.var_positions {
631                    fused_desc.push(
632                        node.cover
633                            .var_positions
634                            .iter()
635                            .position(|cv| cv == v)
636                            .expect("fusable probe keys are cover variables")
637                            as u64,
638                    );
639                }
640            }
641            let d_fused_desc: Option<TrackedCudaSlice<u64>> = if fused_desc.is_empty() {
642                None
643            } else {
644                let mut tbl = self.memory().alloc::<u64>(fused_desc.len())?;
645                self.htod_launch_metadata_sync_copy_into(&fused_desc, &mut tbl)
646                    .map_err(|e| {
647                        XlogError::Kernel(format!("{ctx}: htod fused-probe table failed: {e}"))
648                    })?;
649                Some(tbl)
650            };
651            let cover_ptrs: Vec<u64> = (depth..depth + c)
652                .map(|i| owned_col_ptr(&norm[a], i, ctx))
653                .collect::<Result<_>>()?;
654            let mut d_cover_tbl = self.memory().alloc::<u64>(c)?;
655            self.htod_launch_metadata_sync_copy_into(&cover_ptrs, &mut d_cover_tbl)
656                .map_err(|e| XlogError::Kernel(format!("{ctx}: htod cover table failed: {e}")))?;
657            let mut marks = self.memory().alloc::<u32>(total_work as usize + 1)?;
658            let has_parent_range: u32 = u32::from(cover_live);
659            let const_lo: u32 = 0;
660            let const_hi: u32 = n_rows[a];
661            let null_ptr: u64 = 0;
662            if count_ran {
663                let mut scan_scratch = self.multiblock_scan_u32_scratch_for_len(total_work + 1)?;
664                let mut rec = LaunchRecorder::new_strict(launch_stream);
665                rec.read(norm[a].num_rows_device());
666                for i in depth..depth + c {
667                    rec.read_column(norm[a].column(i).expect("validated cover column"));
668                }
669                rec.read(&d_cover_tbl);
670                if let Some(wp) = work_prefix.as_ref() {
671                    rec.read(wp);
672                    rec.read(find_col(&frontier, ColTag::RangeLo(a), ctx)?);
673                }
674                if let Some(d) = d_fused_desc.as_ref() {
675                    rec.read(d);
676                    for pr in &fused {
677                        let p = pr.input_idx;
678                        rec.read(norm[p].num_rows_device());
679                        for i in consumed[p]..consumed[p] + pr.var_positions.len() {
680                            rec.read_column(norm[p].column(i).expect("validated probe column"));
681                        }
682                        if frontier.iter().any(|(t, _)| *t == ColTag::RangeLo(p)) {
683                            rec.read(find_col(&frontier, ColTag::RangeLo(p), ctx)?);
684                            rec.read(find_col(&frontier, ColTag::RangeHi(p), ctx)?);
685                        }
686                    }
687                }
688                rec.write(&marks);
689                for level in scan_scratch.levels() {
690                    rec.read_write(level);
691                }
692                rec.preflight(runtime).map_err(|e| {
693                    XlogError::Kernel(format!("{ctx}: count preflight failed: {e}"))
694                })?;
695                let kernel = self
696                    .device()
697                    .inner()
698                    .get_func(WCOJ_MODULE, width.count_kernel())
699                    .ok_or_else(|| {
700                        XlogError::Kernel(format!("{} kernel not found", width.count_kernel()))
701                    })?;
702                let grid = total_work.div_ceil(BLOCK_SIZE);
703                let c_u32 = c as u32;
704                let identity_u32: u32 = u32::from(identity);
705                let n_fused_u32: u32 = fused.len() as u32;
706                // SAFETY: fj_expand_count_u32(cover_cols,
707                // n_cover_cols, parent_lo, work_prefix,
708                // has_parent_range, const_lo, const_hi, n_frontier,
709                // total_work, group_marks). parent_lo/work_prefix
710                // are null when the cover atom is untouched; the
711                // kernel never dereferences them on that branch.
712                let enqueue = |stream: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
713                    unsafe {
714                        let parent_lo_param = match work_prefix.as_ref() {
715                            Some(_) => {
716                                find_col(&frontier, ColTag::RangeLo(a), ctx)?.as_kernel_param()
717                            }
718                            None => null_ptr.as_kernel_param(),
719                        };
720                        let wp_param = match work_prefix.as_ref() {
721                            Some(wp) => wp.as_kernel_param(),
722                            None => null_ptr.as_kernel_param(),
723                        };
724                        let mut params: Vec<*mut c_void> = vec![
725                            (&d_cover_tbl).as_kernel_param(),
726                            c_u32.as_kernel_param(),
727                            parent_lo_param,
728                            wp_param,
729                            has_parent_range.as_kernel_param(),
730                            const_lo.as_kernel_param(),
731                            const_hi.as_kernel_param(),
732                            count.as_kernel_param(),
733                            total_work.as_kernel_param(),
734                            identity_u32.as_kernel_param(),
735                            match d_fused_desc.as_ref() {
736                                Some(d) => d.as_kernel_param(),
737                                None => null_ptr.as_kernel_param(),
738                            },
739                            n_fused_u32.as_kernel_param(),
740                            (&marks).as_kernel_param(),
741                        ];
742                        kernel
743                            .clone()
744                            .launch_in(
745                                stream,
746                                LaunchConfig {
747                                    grid_dim: (grid, 1, 1),
748                                    block_dim: (BLOCK_SIZE, 1, 1),
749                                    shared_mem_bytes: 0,
750                                },
751                                &mut params,
752                            )
753                            .map_err(|e| {
754                                XlogError::Kernel(format!(
755                                    "{} launch failed: {e}",
756                                    width.count_kernel()
757                                ))
758                            })?;
759                    }
760                    self.multiblock_scan_u32_inplace_on_stream(
761                        &mut marks,
762                        total_work + 1,
763                        stream,
764                        &mut scan_scratch,
765                    )?;
766                    Ok(())
767                };
768                // SAFETY: preflight retained the accessed buffers and bound this stream.
769                let rec = unsafe { rec.enqueue_prepared_with(&cu_stream, enqueue) }
770                    .map_err(|error| error.into_xlog_error())?;
771                rec.commit()
772                    .map_err(|e| XlogError::Kernel(format!("{ctx}: count commit failed: {e}")))?;
773            }
774            let n_children = if count_ran {
775                cu_stream
776                    .synchronize()
777                    .map_err(|e| XlogError::Kernel(format!("{ctx}: count sync failed: {e}")))?;
778                self.dtoh_scalar_untracked::<u32>(&marks, total_work as usize)?
779            } else {
780                total_work
781            };
782            if n_children == 0 {
783                return self.create_empty_buffer(out_schema);
784            }
785
786            // ---- EXPAND phase 2 + PROBE: allocate the child
787            // frontier (copied parent columns minus the cover's
788            // ranges, new cover variables, refined cover range when
789            // the cover atom keeps unconsumed columns), the probe
790            // range outputs, and the survival mask — all before the
791            // recorder, per the established discipline.
792            // VAR columns are width-sized data; RANGE columns are u32
793            // row indices in every width class. The emit kernel takes
794            // the two copy groups separately so one launch shape
795            // serves both widths.
796            let var_bytes = (n_children as usize) * width.var_bytes();
797            let range_bytes = (n_children as usize) * std::mem::size_of::<u32>();
798            let mut parent_copy_var_ptrs: Vec<u64> = Vec::new();
799            let mut child_copy_var_ptrs: Vec<u64> = Vec::new();
800            let mut parent_copy_range_ptrs: Vec<u64> = Vec::new();
801            let mut child_copy_range_ptrs: Vec<u64> = Vec::new();
802            let mut child_cols: Vec<FrontierCol> = Vec::new();
803            for (tag, slice) in &frontier {
804                if matches!(tag, ColTag::RangeLo(x) | ColTag::RangeHi(x) if *x == a) {
805                    continue; // cover range is refined, not copied
806                }
807                let is_var = matches!(tag, ColTag::Var(_));
808                let dst =
809                    self.memory()
810                        .alloc::<u8>(if is_var { var_bytes } else { range_bytes })?;
811                if is_var {
812                    parent_copy_var_ptrs.push(*slice.device_ptr());
813                    child_copy_var_ptrs.push(*dst.device_ptr());
814                } else {
815                    parent_copy_range_ptrs.push(*slice.device_ptr());
816                    child_copy_range_ptrs.push(*dst.device_ptr());
817                }
818                child_cols.push((*tag, dst));
819            }
820            let n_copy_var = parent_copy_var_ptrs.len();
821            let n_copy_range = parent_copy_range_ptrs.len();
822            let mut child_var_ptrs: Vec<u64> = Vec::with_capacity(c);
823            for &v in &node.cover.var_positions {
824                let dst = self.memory().alloc::<u8>(var_bytes)?;
825                child_var_ptrs.push(*dst.device_ptr());
826                child_cols.push((ColTag::Var(v), dst));
827            }
828            let keep_cover = depth + c < arities[a];
829            if keep_cover {
830                let lo = self.memory().alloc::<u8>(range_bytes)?;
831                let hi = self.memory().alloc::<u8>(range_bytes)?;
832                child_cols.push((ColTag::RangeLo(a), lo));
833                child_cols.push((ColTag::RangeHi(a), hi));
834            }
835            // Pointer tables (launch metadata; bounded by plan width).
836            let upload_tbl = |ptrs: &[u64]| -> Result<TrackedCudaSlice<u64>> {
837                let mut tbl = self.memory().alloc::<u64>(ptrs.len().max(1))?;
838                if !ptrs.is_empty() {
839                    self.htod_launch_metadata_sync_copy_into(ptrs, &mut tbl)
840                        .map_err(|e| {
841                            XlogError::Kernel(format!("{ctx}: htod pointer table failed: {e}"))
842                        })?;
843                }
844                Ok(tbl)
845            };
846            let d_parent_copy_var_tbl = upload_tbl(&parent_copy_var_ptrs)?;
847            let d_child_copy_var_tbl = upload_tbl(&child_copy_var_ptrs)?;
848            let d_parent_copy_range_tbl = upload_tbl(&parent_copy_range_ptrs)?;
849            let d_child_copy_range_tbl = upload_tbl(&child_copy_range_ptrs)?;
850            let d_child_var_tbl = upload_tbl(&child_var_ptrs)?;
851
852            // Probe pre-allocations (key tables, data tables,
853            // refined ranges, mask).
854            struct ProbePlan {
855                input_idx: usize,
856                n_cols: u32,
857                data_tbl: TrackedCudaSlice<u64>,
858                key_tbl: TrackedCudaSlice<u64>,
859                live: bool,
860                keep: bool,
861                out_lo: Option<TrackedCudaSlice<u8>>,
862                out_hi: Option<TrackedCudaSlice<u8>>,
863            }
864            let mut probe_plans: Vec<ProbePlan> = Vec::with_capacity(node.probes.len());
865            for probe in node.probes.iter().filter(|pr| !is_fusable(pr)) {
866                let p = probe.input_idx;
867                let p_len = probe.var_positions.len();
868                let p_depth = consumed[p];
869                let data_ptrs: Vec<u64> = (p_depth..p_depth + p_len)
870                    .map(|i| owned_col_ptr(&norm[p], i, ctx))
871                    .collect::<Result<_>>()?;
872                let key_ptrs: Vec<u64> = probe
873                    .var_positions
874                    .iter()
875                    .map(|&v| Ok(*find_col(&child_cols, ColTag::Var(v), ctx)?.device_ptr()))
876                    .collect::<Result<_>>()?;
877                let live = child_cols.iter().any(|(t, _)| *t == ColTag::RangeLo(p));
878                let keep = p_depth + p_len < arities[p];
879                let (out_lo, out_hi) = if keep {
880                    (
881                        Some(self.memory().alloc::<u8>(range_bytes)?),
882                        Some(self.memory().alloc::<u8>(range_bytes)?),
883                    )
884                } else {
885                    (None, None)
886                };
887                probe_plans.push(ProbePlan {
888                    input_idx: p,
889                    n_cols: p_len as u32,
890                    data_tbl: upload_tbl(&data_ptrs)?,
891                    key_tbl: upload_tbl(&key_ptrs)?,
892                    live,
893                    keep,
894                    out_lo,
895                    out_hi,
896                });
897            }
898            let mask: Option<TrackedCudaSlice<u8>> = if probe_plans.is_empty() {
899                None
900            } else {
901                Some(self.memory().alloc::<u8>(n_children as usize)?)
902            };
903
904            {
905                let mut rec = LaunchRecorder::new_strict(launch_stream);
906                for i in depth..depth + c {
907                    rec.read_column(norm[a].column(i).expect("validated cover column"));
908                }
909                rec.read(&d_cover_tbl);
910                if count_ran {
911                    rec.read(&marks);
912                }
913                if let Some(wp) = work_prefix.as_ref() {
914                    rec.read(wp);
915                    rec.read(find_col(&frontier, ColTag::RangeLo(a), ctx)?);
916                    rec.read(find_col(&frontier, ColTag::RangeHi(a), ctx)?);
917                }
918                for (_, slice) in &frontier {
919                    rec.read(slice);
920                }
921                rec.read(&d_parent_copy_var_tbl);
922                rec.read(&d_child_copy_var_tbl);
923                rec.read(&d_parent_copy_range_tbl);
924                rec.read(&d_child_copy_range_tbl);
925                rec.read(&d_child_var_tbl);
926                for (_, slice) in &child_cols {
927                    rec.write(slice);
928                }
929                for pp in probe_plans.iter() {
930                    let p = pp.input_idx;
931                    for i in consumed[p]..consumed[p] + pp.n_cols as usize {
932                        rec.read_column(norm[p].column(i).expect("validated probe column"));
933                    }
934                    rec.read(norm[p].num_rows_device());
935                    rec.read(&pp.data_tbl);
936                    rec.read(&pp.key_tbl);
937                    if let Some(lo) = pp.out_lo.as_ref() {
938                        rec.write(lo);
939                    }
940                    if let Some(hi) = pp.out_hi.as_ref() {
941                        rec.write(hi);
942                    }
943                }
944                if let Some(m) = mask.as_ref() {
945                    rec.write(m);
946                }
947                rec.preflight(runtime)
948                    .map_err(|e| XlogError::Kernel(format!("{ctx}: emit preflight failed: {e}")))?;
949
950                let emit_kernel = self
951                    .device()
952                    .inner()
953                    .get_func(WCOJ_MODULE, width.emit_kernel())
954                    .ok_or_else(|| {
955                        XlogError::Kernel(format!("{} kernel not found", width.emit_kernel()))
956                    })?;
957                let grid = total_work.div_ceil(BLOCK_SIZE);
958                let c_u32 = c as u32;
959                let n_copy_var_u32 = n_copy_var as u32;
960                let n_copy_range_u32 = n_copy_range as u32;
961                let keep_cover_u32 = u32::from(keep_cover);
962                // SAFETY: fj_expand_emit_{u32,u64}(cover_cols,
963                // n_cover_cols, parent_lo, parent_hi, work_prefix,
964                // has_parent_range, const_lo, const_hi, n_frontier,
965                // total_work, group_offsets, parent_copy_var_cols,
966                // child_copy_var_cols, n_copy_var_cols,
967                // parent_copy_range_cols, child_copy_range_cols,
968                // n_copy_range_cols, child_var_cols, keep_cover_range,
969                // child_cover_lo, child_cover_hi). Nullable pointers
970                // are only dereferenced behind their flags.
971                let enqueue = |stream: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
972                    unsafe {
973                        let parent_lo_param = match work_prefix.as_ref() {
974                            Some(_) => {
975                                find_col(&frontier, ColTag::RangeLo(a), ctx)?.as_kernel_param()
976                            }
977                            None => null_ptr.as_kernel_param(),
978                        };
979                        let parent_hi_param = match work_prefix.as_ref() {
980                            Some(_) => {
981                                find_col(&frontier, ColTag::RangeHi(a), ctx)?.as_kernel_param()
982                            }
983                            None => null_ptr.as_kernel_param(),
984                        };
985                        let wp_param = match work_prefix.as_ref() {
986                            Some(wp) => wp.as_kernel_param(),
987                            None => null_ptr.as_kernel_param(),
988                        };
989                        let cover_lo_param = if keep_cover {
990                            find_col(&child_cols, ColTag::RangeLo(a), ctx)?.as_kernel_param()
991                        } else {
992                            null_ptr.as_kernel_param()
993                        };
994                        let cover_hi_param = if keep_cover {
995                            find_col(&child_cols, ColTag::RangeHi(a), ctx)?.as_kernel_param()
996                        } else {
997                            null_ptr.as_kernel_param()
998                        };
999                        let mut params: Vec<*mut c_void> = vec![
1000                            (&d_cover_tbl).as_kernel_param(),
1001                            c_u32.as_kernel_param(),
1002                            parent_lo_param,
1003                            parent_hi_param,
1004                            wp_param,
1005                            has_parent_range.as_kernel_param(),
1006                            const_lo.as_kernel_param(),
1007                            const_hi.as_kernel_param(),
1008                            count.as_kernel_param(),
1009                            total_work.as_kernel_param(),
1010                            // No-count path: null offsets select the kernel's
1011                            // out == w branch (every position is its own group).
1012                            if count_ran {
1013                                (&marks).as_kernel_param()
1014                            } else {
1015                                null_ptr.as_kernel_param()
1016                            },
1017                            (&d_parent_copy_var_tbl).as_kernel_param(),
1018                            (&d_child_copy_var_tbl).as_kernel_param(),
1019                            n_copy_var_u32.as_kernel_param(),
1020                            (&d_parent_copy_range_tbl).as_kernel_param(),
1021                            (&d_child_copy_range_tbl).as_kernel_param(),
1022                            n_copy_range_u32.as_kernel_param(),
1023                            (&d_child_var_tbl).as_kernel_param(),
1024                            keep_cover_u32.as_kernel_param(),
1025                            cover_lo_param,
1026                            cover_hi_param,
1027                        ];
1028                        emit_kernel
1029                            .clone()
1030                            .launch_in(
1031                                stream,
1032                                LaunchConfig {
1033                                    grid_dim: (grid, 1, 1),
1034                                    block_dim: (BLOCK_SIZE, 1, 1),
1035                                    shared_mem_bytes: 0,
1036                                },
1037                                &mut params,
1038                            )
1039                            .map_err(|e| {
1040                                XlogError::Kernel(format!(
1041                                    "{} launch failed: {e}",
1042                                    width.emit_kernel()
1043                                ))
1044                            })?;
1045                    }
1046
1047                    // PROBE refinements over the expanded frontier.
1048                    let probe_kernel = self
1049                        .device()
1050                        .inner()
1051                        .get_func(WCOJ_MODULE, width.probe_kernel())
1052                        .ok_or_else(|| {
1053                            XlogError::Kernel(format!("{} kernel not found", width.probe_kernel()))
1054                        })?;
1055                    let probe_grid = n_children.div_ceil(BLOCK_SIZE);
1056                    for (probe_idx, pp) in probe_plans.iter().enumerate() {
1057                        let p = pp.input_idx;
1058                        let has_range = u32::from(pp.live);
1059                        let p_const_lo: u32 = 0;
1060                        let p_const_hi: u32 = n_rows[p];
1061                        let keep_u32 = u32::from(pp.keep);
1062                        let combine: u32 = u32::from(probe_idx > 0);
1063                        // SAFETY: fj_probe_refine_u32(probe_cols,
1064                        // n_probe_cols, key_cols, in_lo, in_hi, has_range,
1065                        // const_lo, const_hi, n_frontier, keep_range,
1066                        // out_lo, out_hi, mask, combine_mask). Nullable
1067                        // pointers only dereferenced behind their flags.
1068                        unsafe {
1069                            let in_lo_param = if pp.live {
1070                                find_col(&child_cols, ColTag::RangeLo(p), ctx)?.as_kernel_param()
1071                            } else {
1072                                null_ptr.as_kernel_param()
1073                            };
1074                            let in_hi_param = if pp.live {
1075                                find_col(&child_cols, ColTag::RangeHi(p), ctx)?.as_kernel_param()
1076                            } else {
1077                                null_ptr.as_kernel_param()
1078                            };
1079                            let out_lo_param = match pp.out_lo.as_ref() {
1080                                Some(lo) => lo.as_kernel_param(),
1081                                None => null_ptr.as_kernel_param(),
1082                            };
1083                            let out_hi_param = match pp.out_hi.as_ref() {
1084                                Some(hi) => hi.as_kernel_param(),
1085                                None => null_ptr.as_kernel_param(),
1086                            };
1087                            let mask_ref = mask.as_ref().expect("mask exists when probes exist");
1088                            let mut params: Vec<*mut c_void> = vec![
1089                                (&pp.data_tbl).as_kernel_param(),
1090                                pp.n_cols.as_kernel_param(),
1091                                (&pp.key_tbl).as_kernel_param(),
1092                                in_lo_param,
1093                                in_hi_param,
1094                                has_range.as_kernel_param(),
1095                                p_const_lo.as_kernel_param(),
1096                                p_const_hi.as_kernel_param(),
1097                                n_children.as_kernel_param(),
1098                                keep_u32.as_kernel_param(),
1099                                out_lo_param,
1100                                out_hi_param,
1101                                mask_ref.as_kernel_param(),
1102                                combine.as_kernel_param(),
1103                            ];
1104                            probe_kernel
1105                                .clone()
1106                                .launch_in(
1107                                    stream,
1108                                    LaunchConfig {
1109                                        grid_dim: (probe_grid, 1, 1),
1110                                        block_dim: (BLOCK_SIZE, 1, 1),
1111                                        shared_mem_bytes: 0,
1112                                    },
1113                                    &mut params,
1114                                )
1115                                .map_err(|e| {
1116                                    XlogError::Kernel(format!(
1117                                        "{} launch failed: {e}",
1118                                        width.probe_kernel()
1119                                    ))
1120                                })?;
1121                        }
1122                    }
1123                    Ok(())
1124                };
1125                // SAFETY: preflight retained the accessed buffers and bound this stream.
1126                let rec = unsafe { rec.enqueue_prepared_with(&cu_stream, enqueue) }
1127                    .map_err(|error| error.into_xlog_error())?;
1128                rec.commit()
1129                    .map_err(|e| XlogError::Kernel(format!("{ctx}: emit commit failed: {e}")))?;
1130            }
1131
1132            // ---- Bookkeeping: consumed columns, live-range set.
1133            consumed[a] += c;
1134            for probe in &node.probes {
1135                consumed[probe.input_idx] += probe.var_positions.len();
1136            }
1137            // Fused probes exhaust their atoms in the count pass: drop
1138            // their stale copied ranges from the child frontier.
1139            for pr in &fused {
1140                let p = pr.input_idx;
1141                child_cols.retain(
1142                    |(t, _)| !matches!(t, ColTag::RangeLo(x) | ColTag::RangeHi(x) if *x == p),
1143                );
1144            }
1145            // Replace probed atoms' stale (copied) ranges with the
1146            // refined outputs; exhausted atoms drop their ranges.
1147            for pp in &mut probe_plans {
1148                let p = pp.input_idx;
1149                child_cols.retain(
1150                    |(t, _)| !matches!(t, ColTag::RangeLo(x) | ColTag::RangeHi(x) if *x == p),
1151                );
1152                if pp.keep {
1153                    child_cols.push((
1154                        ColTag::RangeLo(p),
1155                        pp.out_lo.take().expect("keep implies out_lo"),
1156                    ));
1157                    child_cols.push((
1158                        ColTag::RangeHi(p),
1159                        pp.out_hi.take().expect("keep implies out_hi"),
1160                    ));
1161                }
1162            }
1163
1164            // ---- Compaction (single mask pass per node).
1165            if let Some(mask) = mask {
1166                let tags: Vec<ColTag> = child_cols.iter().map(|(t, _)| *t).collect();
1167                // Per-tag column types: the compaction helper sizes
1168                // its per-column copies from the schema, so the mixed
1169                // VAR/RANGE width classes need no special casing.
1170                let schema = Schema::new(
1171                    tags.iter()
1172                        .enumerate()
1173                        .map(|(i, t)| (format!("f{i}"), tag_type(*t, width)))
1174                        .collect(),
1175                );
1176                let d_nr = self.memory().alloc::<u32>(1)?;
1177                let mut rec_num_rows = LaunchRecorder::new_strict(launch_stream);
1178                rec_num_rows.write(&d_nr);
1179                rec_num_rows.preflight(runtime).map_err(|error| {
1180                    XlogError::Kernel(format!("{ctx}: num_rows preflight failed: {error}"))
1181                })?;
1182                let rec_num_rows = unsafe {
1183                    rec_num_rows.enqueue_prepared_with(&cu_stream, |enqueue| {
1184                        self.initialize_launch_metadata_u32(
1185                            n_children,
1186                            &d_nr,
1187                            enqueue,
1188                            &format!("{ctx}: frontier num_rows"),
1189                        )
1190                    })
1191                }
1192                .map_err(crate::launch::LaunchEnqueueError::into_xlog_error)?;
1193                rec_num_rows.commit().map_err(|error| {
1194                    XlogError::Kernel(format!("{ctx}: num_rows commit failed: {error}"))
1195                })?;
1196                let columns: Vec<CudaColumn> =
1197                    child_cols.drain(..).map(|(_, s)| s.into()).collect();
1198                let staging = CudaBuffer::from_columns_with_host_count(
1199                    columns,
1200                    u64::from(n_children),
1201                    d_nr,
1202                    schema,
1203                    n_children,
1204                );
1205                let compacted = self.compact_buffer_by_device_mask_counted_recorded(
1206                    &staging,
1207                    &mask,
1208                    launch_stream,
1209                )?;
1210                let new_count = compacted.cached_row_count().ok_or_else(|| {
1211                    XlogError::Kernel(format!("{ctx}: compaction lost its row count"))
1212                })?;
1213                if new_count == 0 {
1214                    return self.create_empty_buffer(out_schema);
1215                }
1216                let mut new_frontier: Vec<FrontierCol> = Vec::with_capacity(tags.len());
1217                for (tag, col) in tags.into_iter().zip(compacted.columns) {
1218                    let CudaColumn::Owned(slice) = col else {
1219                        return Err(XlogError::Kernel(format!(
1220                            "{ctx}: compaction produced a non-owned column"
1221                        )));
1222                    };
1223                    new_frontier.push((tag, slice));
1224                }
1225                frontier = new_frontier;
1226                count = new_count;
1227            } else {
1228                frontier = child_cols;
1229                count = n_children;
1230            }
1231            frontier_cap = n_children;
1232        }
1233
1234        // ---- COUNT epilogue (§2.4): per-row multiplicity = product
1235        // of remaining live trie-range lengths, then the existing
1236        // recorded groupby Sum reduces (group, multiplicity) to
1237        // (group, count). Unconsumed trailing columns never expand
1238        // the frontier — this is the d-representation count.
1239        if mode == FjMode::CountByRoot {
1240            let group_var = plan.output_vars[0];
1241            let mut lo_ptrs: Vec<u64> = Vec::new();
1242            let mut hi_ptrs: Vec<u64> = Vec::new();
1243            for (t, s) in &frontier {
1244                if let ColTag::RangeLo(x) = t {
1245                    lo_ptrs.push(*s.device_ptr());
1246                    hi_ptrs.push(*find_col(&frontier, ColTag::RangeHi(*x), ctx)?.device_ptr());
1247                }
1248            }
1249            let upload_tbl = |ptrs: &[u64]| -> Result<TrackedCudaSlice<u64>> {
1250                let mut tbl = self.memory().alloc::<u64>(ptrs.len().max(1))?;
1251                if !ptrs.is_empty() {
1252                    self.htod_launch_metadata_sync_copy_into(ptrs, &mut tbl)
1253                        .map_err(|e| {
1254                            XlogError::Kernel(format!("{ctx}: htod range table failed: {e}"))
1255                        })?;
1256                }
1257                Ok(tbl)
1258            };
1259            let d_lo_tbl = upload_tbl(&lo_ptrs)?;
1260            let d_hi_tbl = upload_tbl(&hi_ptrs)?;
1261            // Sized to the frontier CAPACITY so the staging buffer's
1262            // row_cap invariant holds; only the logical `count`
1263            // prefix is written/read.
1264            let mut mult = self
1265                .memory()
1266                .alloc::<u8>(frontier_cap as usize * std::mem::size_of::<u64>())?;
1267            {
1268                let mut rec = LaunchRecorder::new_strict(launch_stream);
1269                for (t, s) in &frontier {
1270                    if matches!(t, ColTag::RangeLo(_) | ColTag::RangeHi(_)) {
1271                        rec.read(s);
1272                    }
1273                }
1274                rec.read(&d_lo_tbl);
1275                rec.read(&d_hi_tbl);
1276                rec.write(&mult);
1277                rec.preflight(runtime).map_err(|e| {
1278                    XlogError::Kernel(format!("{ctx}: multiplicity preflight failed: {e}"))
1279                })?;
1280                let kernel = self
1281                    .device()
1282                    .inner()
1283                    .get_func(WCOJ_MODULE, wcoj_kernels::FJ_COUNT_MULTIPLICITY)
1284                    .ok_or_else(|| {
1285                        XlogError::Kernel("fj_count_multiplicity kernel not found".to_string())
1286                    })?;
1287                let grid = count.div_ceil(BLOCK_SIZE);
1288                let n_ranges = lo_ptrs.len() as u32;
1289                // SAFETY: fj_count_multiplicity(range_lo_cols,
1290                // range_hi_cols, n_ranges, n_frontier, mult);
1291                // device-resident, preflighted.
1292                let enqueue = |stream: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
1293                    unsafe {
1294                        kernel
1295                            .clone()
1296                            .launch_in(
1297                                stream,
1298                                LaunchConfig {
1299                                    grid_dim: (grid, 1, 1),
1300                                    block_dim: (BLOCK_SIZE, 1, 1),
1301                                    shared_mem_bytes: 0,
1302                                },
1303                                (&d_lo_tbl, &d_hi_tbl, n_ranges, count, &mut mult),
1304                            )
1305                            .map_err(|e| {
1306                                XlogError::Kernel(format!(
1307                                    "fj_count_multiplicity launch failed: {e}"
1308                                ))
1309                            })?;
1310                    }
1311                    Ok(())
1312                };
1313                // SAFETY: preflight retained the accessed buffers and bound this stream.
1314                let rec = unsafe { rec.enqueue_prepared_with(&cu_stream, enqueue) }
1315                    .map_err(|error| error.into_xlog_error())?;
1316                rec.commit().map_err(|e| {
1317                    XlogError::Kernel(format!("{ctx}: multiplicity commit failed: {e}"))
1318                })?;
1319            }
1320            let key_idx = frontier
1321                .iter()
1322                .position(|(t, _)| *t == ColTag::Var(group_var))
1323                .ok_or_else(|| {
1324                    XlogError::Kernel(format!("{ctx}: group variable {group_var} missing"))
1325                })?;
1326            let (_, key_col) = frontier.swap_remove(key_idx);
1327            let d_nr = self.memory().alloc::<u32>(1)?;
1328            let mut rec_num_rows = LaunchRecorder::new_strict(launch_stream);
1329            rec_num_rows.write(&d_nr);
1330            rec_num_rows.preflight(runtime).map_err(|error| {
1331                XlogError::Kernel(format!("{ctx}: num_rows preflight failed: {error}"))
1332            })?;
1333            let rec_num_rows = unsafe {
1334                rec_num_rows.enqueue_prepared_with(&cu_stream, |enqueue| {
1335                    self.initialize_launch_metadata_u32(
1336                        count,
1337                        &d_nr,
1338                        enqueue,
1339                        &format!("{ctx}: staging num_rows"),
1340                    )
1341                })
1342            }
1343            .map_err(crate::launch::LaunchEnqueueError::into_xlog_error)?;
1344            rec_num_rows.commit().map_err(|error| {
1345                XlogError::Kernel(format!("{ctx}: num_rows commit failed: {error}"))
1346            })?;
1347            let staging_schema = Schema::new(vec![
1348                (format!("v{group_var}"), width.var_type()),
1349                ("count".to_string(), ScalarType::U64),
1350            ]);
1351            let staging = CudaBuffer::from_columns_with_host_count(
1352                vec![key_col.into(), mult.into()],
1353                u64::from(frontier_cap),
1354                d_nr,
1355                staging_schema,
1356                count,
1357            );
1358            return self.groupby_multi_agg_recorded(
1359                &staging,
1360                &[0],
1361                &[(1, xlog_core::AggOp::Sum)],
1362                launch_stream,
1363            );
1364        }
1365
1366        // ---- Final materialization: project the head variables out
1367        // of the frontier (recorded dtod copies). The frontier holds
1368        // exactly the bound-variable columns at this point (all atoms
1369        // exhausted ⇒ no range columns survive).
1370        let perm: Vec<usize> = plan
1371            .output_vars
1372            .iter()
1373            .map(|&v| {
1374                frontier
1375                    .iter()
1376                    .position(|(t, _)| *t == ColTag::Var(v))
1377                    .ok_or_else(|| XlogError::Kernel(format!("{ctx}: output variable {v} missing")))
1378            })
1379            .collect::<Result<_>>()?;
1380        let schema = Schema::new(
1381            frontier
1382                .iter()
1383                .enumerate()
1384                .map(|(i, (t, _))| (format!("f{i}"), tag_type(*t, width)))
1385                .collect(),
1386        );
1387        let d_nr = self.memory().alloc::<u32>(1)?;
1388        let mut rec_num_rows = LaunchRecorder::new_strict(launch_stream);
1389        rec_num_rows.write(&d_nr);
1390        rec_num_rows.preflight(runtime).map_err(|error| {
1391            XlogError::Kernel(format!("{ctx}: num_rows preflight failed: {error}"))
1392        })?;
1393        let rec_num_rows = unsafe {
1394            rec_num_rows.enqueue_prepared_with(&cu_stream, |enqueue| {
1395                self.initialize_launch_metadata_u32(
1396                    count,
1397                    &d_nr,
1398                    enqueue,
1399                    &format!("{ctx}: result num_rows"),
1400                )
1401            })
1402        }
1403        .map_err(crate::launch::LaunchEnqueueError::into_xlog_error)?;
1404        rec_num_rows.commit().map_err(|error| {
1405            XlogError::Kernel(format!("{ctx}: num_rows commit failed: {error}"))
1406        })?;
1407        let columns: Vec<CudaColumn> = frontier.into_iter().map(|(_, s)| s.into()).collect();
1408        // row_cap = frontier CAPACITY (compaction shrinks the logical
1409        // count without reallocating columns); the logical count rides
1410        // on num_rows_device + the host cache.
1411        let src = CudaBuffer::from_columns_with_host_count(
1412            columns,
1413            u64::from(frontier_cap),
1414            d_nr,
1415            schema,
1416            count,
1417        );
1418        let projected =
1419            self.wcoj_project_output_columns_recorded(&src, &perm, out_schema, launch_stream)?;
1420        // Set semantics: projecting a strict subset of the bound
1421        // variables can introduce duplicates — dedup exactly then
1422        // (the full projection is duplicate-free by construction).
1423        let distinct_outputs: std::collections::BTreeSet<usize> =
1424            plan.output_vars.iter().copied().collect();
1425        if distinct_outputs.len() < bind_order.len() {
1426            return self.dedup_full_row_recorded(&projected, launch_stream);
1427        }
1428        Ok(projected)
1429    }
1430}