1use 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#[derive(Debug, Clone)]
60pub struct FjSubAtom {
61 pub input_idx: usize,
64 pub var_positions: Vec<usize>,
67}
68
69#[derive(Debug, Clone)]
73pub struct FjNode {
74 pub cover: FjSubAtom,
75 pub probes: Vec<FjSubAtom>,
76}
77
78#[derive(Debug, Clone)]
82pub struct FjPlan {
83 pub num_vars: usize,
85 pub nodes: Vec<FjNode>,
87 pub output_vars: Vec<usize>,
91}
92
93#[derive(Debug, Clone, Copy, PartialEq, Eq)]
98enum ColTag {
99 Var(usize),
101 RangeLo(usize),
103 RangeHi(usize),
105}
106
107type FrontierCol = (ColTag, TrackedCudaSlice<u8>);
108
109#[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
155fn 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
166enum FjMode {
167 Materialize,
169 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
200fn 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 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 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 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 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 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 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 let mut frontier: Vec<FrontierCol> = Vec::new();
480 let mut count: u32 = 1;
481 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 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 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 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 let identity = depth + c >= arities[a];
589 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 let count_ran = !identity || !fused.is_empty();
610 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 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 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 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; }
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 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 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 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 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 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 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 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 consumed[a] += c;
1134 for probe in &node.probes {
1135 consumed[probe.input_idx] += probe.var_positions.len();
1136 }
1137 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 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 if let Some(mask) = mask {
1166 let tags: Vec<ColTag> = child_cols.iter().map(|(t, _)| *t).collect();
1167 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 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 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 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 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 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 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 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}