1use std::collections::HashMap;
7use std::marker::PhantomData;
8use std::path::PathBuf;
9use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
10use std::sync::{Arc, Mutex, OnceLock};
11
12use std::ffi::c_void;
13use xlog_core::{Result, Schema, XlogError};
14
15use crate::{
16 cuda_compat::{
17 AsKernelParam, DeviceParamStorage, DevicePtr, DeviceRepr, DeviceSlice,
18 IntoKernelParamStorage, LaunchAsync, LaunchConfig,
19 },
20 cuda_graph::{CapturedCudaGraph, CsmCudaGraphKey, CudaGraphNode},
21 memory::{validate_logical_row_count, CudaColumn, TrackedCudaSlice},
22 CudaBuffer, CudaDevice, CudaStream, CudaViewMut, GpuMemoryManager,
23};
24
25mod arithmetic;
26mod filter;
27mod fj;
28mod fj_delta;
29mod fj_delta_sparse;
30mod groupby;
31mod ilp;
32mod ilp_exact;
33mod io;
34mod kernel_loading;
35pub mod kernel_paths;
36mod launch_safe;
37mod probabilistic;
38mod relational;
39mod transfer;
40mod wcoj;
41mod wcoj_metadata;
42mod wcoj_project;
43
44pub use fj::{FjNode, FjPlan, FjSubAtom};
45pub use fj_delta::{FjDeltaCols, FJ_DELTA_MAX_DOMAIN};
46
47#[derive(Debug, Clone, Default)]
49pub struct PtxLoadProfile {
50 pub total_sec: f64,
51 pub per_module_sec: Vec<(String, f64)>,
52 pub cubin_loaded: u32,
53 pub ptx_fallback: u32,
54}
55
56fn warmup_profiling_enabled() -> bool {
57 std::env::var("XLOG_WARMUP_PROFILE")
58 .map(|v| v == "1")
59 .unwrap_or(false)
60}
61
62pub(crate) fn detect_compute_capability(device: &Arc<CudaDevice>) -> Result<u32> {
64 let major = device
65 .inner()
66 .attribute(
67 cudarc::driver::sys::CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
68 )
69 .map_err(|e| XlogError::Kernel(format!("Failed to query SM major: {}", e)))?;
70 let minor = device
71 .inner()
72 .attribute(
73 cudarc::driver::sys::CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
74 )
75 .map_err(|e| XlogError::Kernel(format!("Failed to query SM minor: {}", e)))?;
76 Ok((major as u32) * 10 + (minor as u32))
77}
78
79#[cfg(test)]
80fn resolve_module_path(name: &str, cc: u32) -> Option<(std::path::PathBuf, bool)> {
81 kernel_paths::KernelArtifactLocator::from_env().resolve_module_path(name, cc)
82}
83
84#[derive(Debug)]
85pub(crate) enum KernelModuleSource {
86 File { path: PathBuf, is_cubin: bool },
87 EmbeddedPortablePtx { ptx: &'static str },
88}
89
90pub(crate) fn resolve_module_sources_with_locator(
91 name: &str,
92 cc: u32,
93 locator: &kernel_paths::KernelArtifactLocator,
94) -> Vec<KernelModuleSource> {
95 let mut sources: Vec<KernelModuleSource> = locator
96 .resolve_module_paths(name, cc)
97 .into_iter()
98 .filter(|(path, _)| !staged_artifact_is_stale(path))
103 .map(|(path, is_cubin)| KernelModuleSource::File { path, is_cubin })
104 .collect();
105
106 if let Some(ptx) = crate::embedded_kernel_data::portable_ptx(name) {
113 sources.push(KernelModuleSource::EmbeddedPortablePtx { ptx });
114 }
115 sources
116}
117
118fn staged_artifact_is_stale(path: &std::path::Path) -> bool {
127 let Some(file_name) = path.file_name().and_then(|n| n.to_str()) else {
128 return false;
129 };
130 let Some(expected) = crate::embedded_kernel_data::canonical_artifact_hash(file_name) else {
131 return false;
132 };
133 match std::fs::read(path) {
134 Ok(bytes) => fnv1a_64(&bytes) != expected,
135 Err(_) => false,
136 }
137}
138
139fn fnv1a_64(bytes: &[u8]) -> u64 {
141 let mut hash: u64 = 0xcbf2_9ce4_8422_2325;
142 for &byte in bytes {
143 hash ^= byte as u64;
144 hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
145 }
146 hash
147}
148
149#[cfg(test)]
150mod kernel_source_resolution_tests {
151 use super::{
152 kernel_paths::KernelArtifactLocator, resolve_module_sources_with_locator,
153 KernelModuleSource,
154 };
155 use std::fs;
156
157 #[test]
158 fn keeps_portable_ptx_fallback_when_cubin_exists() {
159 let root = std::env::temp_dir().join(format!(
160 "xlog-kernel-fallback-{}-{}",
161 std::process::id(),
162 std::time::SystemTime::now()
163 .duration_since(std::time::UNIX_EPOCH)
164 .expect("system clock before UNIX_EPOCH")
165 .as_nanos()
166 ));
167 let kernels = root.join("kernels");
168 fs::create_dir_all(&kernels).expect("create kernels dir");
169 fs::write(kernels.join("fakekernel.sm_86.cubin"), b"cubin").expect("write cubin");
174 fs::write(kernels.join("fakekernel.portable.ptx"), b"ptx").expect("write ptx");
175 let expected_cubin = kernels.join("fakekernel.sm_86.cubin");
176 let expected_ptx = kernels.join("fakekernel.portable.ptx");
177
178 let locator = KernelArtifactLocator::new(None, Some(kernels.clone()), None);
179 let sources = resolve_module_sources_with_locator("fakekernel", 86, &locator);
180
181 assert_eq!(sources.len(), 2);
182 assert!(matches!(
183 &sources[0],
184 KernelModuleSource::File {
185 path,
186 is_cubin: true
187 } if path == &expected_cubin
188 ));
189 assert!(matches!(
190 &sources[1],
191 KernelModuleSource::File {
192 path,
193 is_cubin: false
194 } if path == &expected_ptx
195 ));
196
197 fs::remove_dir_all(root).expect("remove temp kernels");
198 }
199
200 #[test]
205 fn fnv1a_64_matches_known_vectors() {
206 assert_eq!(super::fnv1a_64(b""), 0xcbf2_9ce4_8422_2325);
207 assert_eq!(super::fnv1a_64(b"a"), 0xaf63_dc4c_8601_ec8c);
208 assert_eq!(super::fnv1a_64(b"foobar"), 0x85944171_f73967e8);
209 }
210
211 #[test]
215 fn staged_artifact_not_stale_without_canonical_hash() {
216 let root = std::env::temp_dir().join(format!(
217 "xlog-kernel-stale-{}-{}",
218 std::process::id(),
219 std::time::SystemTime::now()
220 .duration_since(std::time::UNIX_EPOCH)
221 .expect("system clock before UNIX_EPOCH")
222 .as_nanos()
223 ));
224 fs::create_dir_all(&root).expect("create dir");
225 let unknown = root.join("definitely_not_a_real_kernel.sm_86.cubin");
226 fs::write(&unknown, b"bytes").expect("write");
227 assert!(!super::staged_artifact_is_stale(&unknown));
228 assert!(!super::staged_artifact_is_stale(
229 &root.join("missing.portable.ptx")
230 ));
231 fs::remove_dir_all(root).expect("remove temp dir");
232 }
233}
234
235pub(crate) fn load_module_sources(name: &str, cc: u32) -> Result<Vec<KernelModuleSource>> {
240 debug_assert!(
241 crate::kernel_manifest_data::KERNEL_CU_NAMES.contains(&name),
242 "kernel module '{name}' is not in KERNEL_CU_NAMES manifest — update kernel_manifest_data.rs"
243 );
244 let locator = kernel_paths::KernelArtifactLocator::from_env();
245 let sources = resolve_module_sources_with_locator(name, cc, &locator);
246 if sources.is_empty() {
247 Err(XlogError::Kernel(format!(
248 "{name}: no cubin, sidecar portable PTX, or embedded portable PTX found"
249 )))
250 } else {
251 Ok(sources)
252 }
253}
254
255#[derive(Clone)]
256pub(crate) struct RawCudaView<'a, T> {
257 ptr: cudarc::driver::sys::CUdeviceptr,
258 len: usize,
259 stream: Arc<CudaStream>,
260 #[allow(dead_code)]
273 source_block: Option<&'a crate::device_runtime::DeviceBlock>,
274 _marker: PhantomData<&'a [T]>,
275}
276
277pub(crate) struct MultiblockScanScratchU32 {
284 levels: Vec<TrackedCudaSlice<u32>>,
285}
286
287impl MultiblockScanScratchU32 {
288 pub(crate) fn levels(&self) -> &[TrackedCudaSlice<u32>] {
289 &self.levels
290 }
291}
292
293pub(crate) struct CsmCudaGraphNodes {
294 pub(crate) count: CudaGraphNode,
295 pub(crate) total: CudaGraphNode,
296 pub(crate) materialize: CudaGraphNode,
297 pub(crate) node_count: usize,
298}
299
300pub(crate) struct CsmCudaGraphEntry {
301 pub(crate) graph: CapturedCudaGraph,
302 pub(crate) nodes: CsmCudaGraphNodes,
303 pub(crate) per_probe_count: TrackedCudaSlice<u32>,
304 pub(crate) per_probe_offsets: TrackedCudaSlice<u32>,
305 pub(crate) d_logical_count: TrackedCudaSlice<u32>,
306 pub(crate) d_overflow: TrackedCudaSlice<u8>,
307 pub(crate) d_output_left: TrackedCudaSlice<u32>,
308 pub(crate) d_output_right: TrackedCudaSlice<u32>,
309 pub(crate) scan_scratch: MultiblockScanScratchU32,
310 pub(crate) probe_capacity: u32,
311 pub(crate) output_capacity: u32,
312}
313
314impl<'a, T> DeviceSlice<T> for RawCudaView<'a, T> {
315 fn len(&self) -> usize {
316 self.len
317 }
318
319 fn stream(&self) -> &Arc<CudaStream> {
320 &self.stream
321 }
322}
323
324impl<'a, T> DevicePtr<T> for RawCudaView<'a, T> {
325 fn device_ptr<'b>(
326 &'b self,
327 _stream: &'b CudaStream,
328 ) -> (
329 cudarc::driver::sys::CUdeviceptr,
330 cudarc::driver::SyncOnDrop<'b>,
331 ) {
332 (self.ptr, cudarc::driver::SyncOnDrop::Sync(None))
333 }
334}
335
336impl<'a, T> RawCudaView<'a, T> {
337 pub fn device_ptr(&self) -> &cudarc::driver::sys::CUdeviceptr {
338 &self.ptr
339 }
340
341 #[allow(dead_code)]
350 pub fn runtime_block(&self) -> Option<&'a crate::device_runtime::DeviceBlock> {
351 self.source_block
352 }
353}
354
355impl<'a, T: DeviceRepr> AsKernelParam for &RawCudaView<'a, T> {
356 fn as_kernel_param(&self) -> *mut c_void {
357 ((*self).device_ptr() as *const cudarc::driver::sys::CUdeviceptr)
358 .cast_mut()
359 .cast()
360 }
361}
362
363impl<'a, T: DeviceRepr> IntoKernelParamStorage for &'a RawCudaView<'a, T> {
364 type Storage = DeviceParamStorage<'a>;
365
366 fn into_kernel_param_storage(self) -> Self::Storage {
367 DeviceParamStorage::unsynced(self.ptr)
368 }
369}
370
371pub struct RadixSortScratch {
373 keys_b: TrackedCudaSlice<u32>,
374 values_b: TrackedCudaSlice<u32>,
375 hist: TrackedCudaSlice<u32>,
376 prefix: TrackedCudaSlice<u32>,
377 ranks: TrackedCudaSlice<u32>,
378 len: u32,
379}
380
381impl RadixSortScratch {
382 pub fn new(provider: &CudaKernelProvider, n: u32) -> Result<Self> {
383 let memory = provider.memory();
384 let len = n.max(1);
385 let keys_b = memory.alloc::<u32>(len as usize)?;
386 let values_b = memory.alloc::<u32>(len as usize)?;
387 let ranks = memory.alloc::<u32>(len as usize)?;
388 let block_size = CudaKernelProvider::SORT_BLOCK_SIZE;
389 let grid_size = len.div_ceil(block_size).max(1);
390 let hist = memory.alloc::<u32>((grid_size as usize) * 16)?;
391 let prefix = memory.alloc::<u32>(16)?;
392 Ok(Self {
393 keys_b,
394 values_b,
395 hist,
396 prefix,
397 ranks,
398 len,
399 })
400 }
401
402 pub fn ensure_capacity(&mut self, provider: &CudaKernelProvider, n: u32) -> Result<()> {
403 if n <= self.len {
404 return Ok(());
405 }
406 *self = Self::new(provider, n)?;
407 Ok(())
408 }
409}
410
411pub const JOIN_MODULE: &str = "xlog_join";
413pub const DEDUP_MODULE: &str = "xlog_dedup";
414pub const GROUPBY_MODULE: &str = "xlog_groupby";
415pub const SCAN_MODULE: &str = "xlog_scan";
416pub const SORT_MODULE: &str = "xlog_sort";
417pub const FILTER_MODULE: &str = "xlog_filter";
418pub const SET_OPS_MODULE: &str = "xlog_set_ops";
419pub const PACK_MODULE: &str = "xlog_pack";
420pub const CIRCUIT_MODULE: &str = "xlog_circuit";
421pub const MC_SAMPLE_MODULE: &str = "xlog_mc_sample";
422pub const MC_EVAL_MODULE: &str = "xlog_mc_eval";
423pub const MC_RESIDENT_MODULE: &str = "xlog_mc_resident";
424pub const ARITH_MODULE: &str = "xlog_arith";
425pub const SAT_MODULE: &str = "xlog_sat";
426pub const D4_MODULE: &str = "xlog_d4";
427pub const NEURAL_MODULE: &str = "xlog_neural";
428pub const PIR_MODULE: &str = "xlog_pir";
429pub const CNF_MODULE: &str = "xlog_cnf";
430pub const CACHE_MODULE: &str = "xlog_cache";
431pub const WEIGHTS_MODULE: &str = "xlog_weights";
432pub const ILP_MODULE: &str = "xlog_ilp";
433pub const ILP_CREDIT_MODULE: &str = "xlog_ilp_credit";
434pub const ILP_EXACT_MODULE: &str = "xlog_ilp_exact";
435pub const EPISTEMIC_MODULE: &str = "xlog_epistemic";
436pub const WCOJ_MODULE: &str = "xlog_wcoj";
437pub const JOINT_SOLVE_MODULE: &str = "xlog_joint_solve";
438
439const _: () = assert!(crate::kernel_manifest_data::KERNEL_CU_NAMES.len() == 26);
441
442pub mod wcoj_kernels {
444 pub const WCOJ_BUILD_METADATA_MARK_BOUNDARIES_U32: &str =
445 "wcoj_build_metadata_mark_boundaries_u32";
446 pub const WCOJ_BUILD_METADATA_MARK_BOUNDARIES_U64: &str =
447 "wcoj_build_metadata_mark_boundaries_u64";
448 pub const WCOJ_BUILD_METADATA_SCATTER_U32: &str = "wcoj_build_metadata_scatter_u32";
449 pub const WCOJ_BUILD_METADATA_SCATTER_U64: &str = "wcoj_build_metadata_scatter_u64";
450 pub const WCOJ_TRIANGLE_BUILD_HG_WORK_PLAN_U32: &str = "wcoj_triangle_build_hg_work_plan_u32";
451 pub const WCOJ_TRIANGLE_COUNT_HG_U32: &str = "wcoj_triangle_count_hg_u32";
452 pub const WCOJ_TRIANGLE_GROUPBY_ROOT_COUNT_HG_U32: &str =
453 "wcoj_triangle_groupby_root_count_hg_u32";
454 pub const WCOJ_TRIANGLE_GROUPBY_ROOT_SUM_HG_U32: &str = "wcoj_triangle_groupby_root_sum_hg_u32";
455 pub const WCOJ_TRIANGLE_GROUPBY_ROOT_MIN_HG_U32: &str = "wcoj_triangle_groupby_root_min_hg_u32";
456 pub const WCOJ_TRIANGLE_GROUPBY_ROOT_MAX_HG_U32: &str = "wcoj_triangle_groupby_root_max_hg_u32";
457 pub const WCOJ_TRIANGLE_MATERIALIZE_HG_U32: &str = "wcoj_triangle_materialize_hg_u32";
458 pub const WCOJ_TRIANGLE_BUILD_HG_WORK_PLAN_U64: &str = "wcoj_triangle_build_hg_work_plan_u64";
459 pub const WCOJ_TRIANGLE_COUNT_HG_U64: &str = "wcoj_triangle_count_hg_u64";
460 pub const WCOJ_TRIANGLE_GROUPBY_ROOT_COUNT_HG_U64: &str =
461 "wcoj_triangle_groupby_root_count_hg_u64";
462 pub const WCOJ_TRIANGLE_GROUPBY_ROOT_SUM_HG_U64: &str = "wcoj_triangle_groupby_root_sum_hg_u64";
463 pub const WCOJ_TRIANGLE_GROUPBY_ROOT_MIN_HG_U64: &str = "wcoj_triangle_groupby_root_min_hg_u64";
464 pub const WCOJ_TRIANGLE_GROUPBY_ROOT_MAX_HG_U64: &str = "wcoj_triangle_groupby_root_max_hg_u64";
465 pub const WCOJ_GROUPBY_ROOT_SEGMENT_SUM_COUNTS_U32: &str =
466 "wcoj_groupby_root_segment_sum_counts_u32";
467 pub const WCOJ_GROUPBY_ROOT_SEGMENT_SUM_VALUES_U64: &str =
468 "wcoj_groupby_root_segment_sum_values_u64";
469 pub const WCOJ_GROUPBY_ROOT_SEGMENT_MIN_VALUES_U64: &str =
470 "wcoj_groupby_root_segment_min_values_u64";
471 pub const WCOJ_GROUPBY_ROOT_SEGMENT_MAX_VALUES_U64: &str =
472 "wcoj_groupby_root_segment_max_values_u64";
473 pub const WCOJ_TRIANGLE_MATERIALIZE_HG_U64: &str = "wcoj_triangle_materialize_hg_u64";
474 pub const WCOJ_TRIANGLE_COUNT_HG_CACHED_U32: &str = "wcoj_triangle_count_hg_cached_u32";
475 pub const WCOJ_TRIANGLE_MATERIALIZE_HG_CACHED_U32: &str =
476 "wcoj_triangle_materialize_hg_cached_u32";
477 pub const WCOJ_SCAN_HG_BLOCK_COUNTS_U32: &str = "wcoj_scan_hg_block_counts_u32";
478 pub const WCOJ_COMPUTE_TOTAL: &str = "wcoj_compute_total";
479 pub const WCOJ_LAYOUT_CHECK_SORTED_UNIQUE_U32: &str = "wcoj_layout_check_sorted_unique_u32";
480 pub const WCOJ_LAYOUT_CHECK_SORTED_UNIQUE_U64: &str = "wcoj_layout_check_sorted_unique_u64";
481 pub const WCOJ_4CYCLE_BUILD_E2_WORK_PREFIX_U32: &str = "wcoj_4cycle_build_e2_work_prefix_u32";
482 pub const WCOJ_4CYCLE_BUILD_HG_WORK_PLAN_U32: &str = "wcoj_4cycle_build_hg_work_plan_u32";
483 pub const WCOJ_4CYCLE_COUNT_HG_U32: &str = "wcoj_4cycle_count_hg_u32";
484 pub const WCOJ_4CYCLE_GROUPBY_ROOT_COUNT_HG_U32: &str = "wcoj_4cycle_groupby_root_count_hg_u32";
485 pub const WCOJ_4CYCLE_GROUPBY_ROOT_SUM_HG_U32: &str = "wcoj_4cycle_groupby_root_sum_hg_u32";
486 pub const WCOJ_4CYCLE_GROUPBY_ROOT_MIN_HG_U32: &str = "wcoj_4cycle_groupby_root_min_hg_u32";
487 pub const WCOJ_4CYCLE_GROUPBY_ROOT_MAX_HG_U32: &str = "wcoj_4cycle_groupby_root_max_hg_u32";
488 pub const WCOJ_4CYCLE_MATERIALIZE_HG_U32: &str = "wcoj_4cycle_materialize_hg_u32";
489 pub const WCOJ_4CYCLE_BUILD_E2_WORK_PREFIX_U64: &str = "wcoj_4cycle_build_e2_work_prefix_u64";
490 pub const WCOJ_4CYCLE_BUILD_HG_WORK_PLAN_U64: &str = "wcoj_4cycle_build_hg_work_plan_u64";
491 pub const WCOJ_4CYCLE_COUNT_HG_U64: &str = "wcoj_4cycle_count_hg_u64";
492 pub const WCOJ_4CYCLE_GROUPBY_ROOT_COUNT_HG_U64: &str = "wcoj_4cycle_groupby_root_count_hg_u64";
493 pub const WCOJ_4CYCLE_MATERIALIZE_HG_U64: &str = "wcoj_4cycle_materialize_hg_u64";
494 pub const WCOJ_CLIQUE5_COUNT_HG_U32: &str = "wcoj_clique5_count_hg_u32";
496 pub const WCOJ_CLIQUE5_MATERIALIZE_HG_U32: &str = "wcoj_clique5_materialize_hg_u32";
497 pub const WCOJ_CLIQUE5_COUNT_HG_U64: &str = "wcoj_clique5_count_hg_u64";
498 pub const WCOJ_CLIQUE5_MATERIALIZE_HG_U64: &str = "wcoj_clique5_materialize_hg_u64";
499 pub const WCOJ_CLIQUE6_COUNT_HG_U32: &str = "wcoj_clique6_count_hg_u32";
500 pub const WCOJ_CLIQUE6_MATERIALIZE_HG_U32: &str = "wcoj_clique6_materialize_hg_u32";
501 pub const WCOJ_CLIQUE6_COUNT_HG_U64: &str = "wcoj_clique6_count_hg_u64";
502 pub const WCOJ_CLIQUE6_MATERIALIZE_HG_U64: &str = "wcoj_clique6_materialize_hg_u64";
503 pub const WCOJ_CLIQUE7_COUNT_HG_U32: &str = "wcoj_clique7_count_hg_u32";
504 pub const WCOJ_CLIQUE7_MATERIALIZE_HG_U32: &str = "wcoj_clique7_materialize_hg_u32";
505 pub const WCOJ_CLIQUE7_COUNT_HG_U64: &str = "wcoj_clique7_count_hg_u64";
506 pub const WCOJ_CLIQUE7_MATERIALIZE_HG_U64: &str = "wcoj_clique7_materialize_hg_u64";
507 pub const WCOJ_CLIQUE8_COUNT_HG_U32: &str = "wcoj_clique8_count_hg_u32";
508 pub const WCOJ_CLIQUE8_MATERIALIZE_HG_U32: &str = "wcoj_clique8_materialize_hg_u32";
509 pub const WCOJ_CLIQUE8_COUNT_HG_U64: &str = "wcoj_clique8_count_hg_u64";
510 pub const WCOJ_CLIQUE8_MATERIALIZE_HG_U64: &str = "wcoj_clique8_materialize_hg_u64";
511 pub const WCOJ_CLIQUE5_GROUPBY_ROOT_COUNT_HG_U32: &str =
512 "wcoj_clique5_groupby_root_count_hg_u32";
513 pub const WCOJ_CLIQUE6_GROUPBY_ROOT_COUNT_HG_U32: &str =
514 "wcoj_clique6_groupby_root_count_hg_u32";
515 pub const FJ_EXPAND_WORK_PREFIX_U32: &str = "fj_expand_work_prefix_u32";
519 pub const FJ_EXPAND_COUNT_U32: &str = "fj_expand_count_u32";
520 pub const FJ_EXPAND_EMIT_U32: &str = "fj_expand_emit_u32";
521 pub const FJ_PROBE_REFINE_U32: &str = "fj_probe_refine_u32";
522 pub const FJ_EXPAND_COUNT_U64: &str = "fj_expand_count_u64";
523 pub const FJ_EXPAND_EMIT_U64: &str = "fj_expand_emit_u64";
524 pub const FJ_PROBE_REFINE_U64: &str = "fj_probe_refine_u64";
525 pub const FJ_COUNT_MULTIPLICITY: &str = "fj_count_multiplicity";
526 pub const FJ_DELTA_RANGE_U32: &str = "fj_delta_range_u32";
528 pub const FJ_DELTA_MARK_U32: &str = "fj_delta_mark_u32";
529 pub const FJ_DELTA_SUBTRACT_U32: &str = "fj_delta_subtract_u32";
530 pub const FJ_DELTA_POPCOUNT: &str = "fj_delta_popcount";
531 pub const FJ_DELTA_EMIT_U32: &str = "fj_delta_emit_u32";
532 pub const FJ_DELTA_MAX_U32: &str = "fj_delta_max_u32";
533 pub const FJ_DELTA_SPARSE_ESTIMATE: &str = "fj_delta_sparse_estimate";
534 pub const FJ_DELTA_SPARSE_LOAD_R: &str = "fj_delta_sparse_load_r";
535 pub const FJ_DELTA_SPARSE_INSERT_CANDIDATES: &str = "fj_delta_sparse_insert_candidates";
536 pub const FJ_DELTA_SPARSE_MARK: &str = "fj_delta_sparse_mark";
537 pub const FJ_DELTA_SPARSE_EMIT: &str = "fj_delta_sparse_emit";
538}
539
540pub mod mc_sample_kernels {
542 pub const MC_SAMPLE_BERNOULLI: &str = "mc_sample_bernoulli";
543}
544
545pub mod mc_eval_kernels {
547 pub const MC_EVAL_MASK_VAR: &str = "mc_eval_mask_var";
548 pub const MC_EVAL_MASK_AD: &str = "mc_eval_mask_ad_choice";
549 pub const MC_EVAL_QUERY_EVIDENCE_TRUTH: &str = "mc_eval_query_evidence_truth";
550 pub const MC_EVAL_ACCUMULATE_COUNTS: &str = "mc_accumulate_counts";
551}
552
553pub mod mc_resident_kernels {
555 pub const MC_RESIDENT_ENGINE: &str = "mc_resident_engine";
558}
559
560pub mod arith_kernels {
562 pub const ARITH_BINARY_I64: &str = "arith_binary_i64";
563 pub const ARITH_BINARY_I32: &str = "arith_binary_i32";
564 pub const ARITH_BINARY_U64: &str = "arith_binary_u64";
565 pub const ARITH_BINARY_U32: &str = "arith_binary_u32";
566 pub const ARITH_BINARY_F64: &str = "arith_binary_f64";
567 pub const ARITH_BINARY_F32: &str = "arith_binary_f32";
568 pub const ARITH_ABS_I64: &str = "arith_abs_i64";
569 pub const ARITH_ABS_I32: &str = "arith_abs_i32";
570 pub const ARITH_ABS_F64: &str = "arith_abs_f64";
571 pub const ARITH_ABS_F32: &str = "arith_abs_f32";
572 pub const ARITH_POW_F64: &str = "arith_pow_f64";
573 pub const ARITH_CAST: &str = "arith_cast";
574 pub const ARITH_FILL_CONST_U32: &str = "arith_fill_const_u32";
575 pub const ARITH_FILL_CONST_U64: &str = "arith_fill_const_u64";
576 pub const ARITH_FILL_CONST_I64: &str = "arith_fill_const_i64";
577 pub const ARITH_FILL_CONST_I32: &str = "arith_fill_const_i32";
578 pub const ARITH_FILL_CONST_F64: &str = "arith_fill_const_f64";
579 pub const ARITH_FILL_CONST_F32: &str = "arith_fill_const_f32";
580 pub const ARITH_FILL_CONST_U8: &str = "arith_fill_const_u8";
581 pub const ARITH_SELECT_I64: &str = "arith_select_i64";
583 pub const ARITH_SELECT_I32: &str = "arith_select_i32";
584 pub const ARITH_SELECT_U64: &str = "arith_select_u64";
585 pub const ARITH_SELECT_U32: &str = "arith_select_u32";
586 pub const ARITH_SELECT_F64: &str = "arith_select_f64";
587 pub const ARITH_SELECT_F32: &str = "arith_select_f32";
588}
589
590pub mod epistemic_kernels {
592 pub const EPISTEMIC_GENERATE_CANDIDATE_ASSUMPTIONS_U8: &str =
594 "epistemic_generate_candidate_assumptions_u8";
595 pub const EPISTEMIC_PROPAGATE_CANDIDATES_U8: &str = "epistemic_propagate_candidates_u8";
597 pub const EPISTEMIC_VALIDATE_CANDIDATE_BITS_U8: &str = "epistemic_validate_candidate_bits_u8";
599 pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_U8: &str =
601 "epistemic_populate_model_membership_u8";
602 pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_FROM_TUPLE_SOURCE_U8: &str =
604 "epistemic_populate_model_membership_from_tuple_source_u8";
605 pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_FROM_TUPLE_SOURCE_ARITY1_U8: &str =
607 "epistemic_populate_model_membership_from_tuple_source_arity1_u8";
608 pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_FROM_TUPLE_SOURCE_ARITY2_U8: &str =
610 "epistemic_populate_model_membership_from_tuple_source_arity2_u8";
611 pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_FROM_TUPLE_SOURCE_ARITY3_U8: &str =
613 "epistemic_populate_model_membership_from_tuple_source_arity3_u8";
614 pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_FROM_TUPLE_SOURCE_ARITY_N_U8: &str =
616 "epistemic_populate_model_membership_from_tuple_source_arity_n_u8";
617 pub const EPISTEMIC_VALIDATE_WORLD_VIEWS_U8: &str = "epistemic_validate_world_views_u8";
619 pub const EPISTEMIC_VALIDATE_CONSTRAINTS_U8: &str = "epistemic_validate_constraints_u8";
621 pub const EPISTEMIC_MATERIALIZE_ACCEPTED_CANDIDATES_U8: &str =
623 "epistemic_materialize_accepted_candidates_u8";
624
625 pub const EPISTEMIC_MATERIALIZE_FINAL_RESULT_FLAGS_U8: &str =
627 "epistemic_materialize_final_result_flags_u8";
628 pub const EPISTEMIC_MATERIALIZE_FINAL_TUPLE_COLUMN_U8: &str =
630 "epistemic_materialize_final_tuple_column_u8";
631 pub const EPISTEMIC_BUILD_FINAL_TUPLE_ROW_MAP_U8: &str =
633 "epistemic_build_final_tuple_row_map_u8";
634 pub const EPISTEMIC_CLOSE_FINAL_TUPLE_REJECTIONS_U8: &str =
636 "epistemic_close_final_tuple_rejections_u8";
637}
638
639pub mod neural_kernels {
641 pub const NEURAL_FILL_AD_CHAIN_F32: &str = "neural_fill_ad_chain_f32";
642 pub const NEURAL_SCATTER_AD_CHAIN_GRADS_F32: &str = "neural_scatter_ad_chain_grads_f32";
643}
644
645pub mod ilp_kernels {
647 pub const EXTRACT_NONZERO_INDICES: &str = "extract_nonzero_indices";
648 pub const ILP_MARK_SELECTED_IDS_U32: &str = "ilp_mark_selected_ids_u32";
649 pub const ILP_MARK_SELECTED_IDS_I32: &str = "ilp_mark_selected_ids_i32";
650 pub const ILP_MARK_SELECTED_IDS_I64: &str = "ilp_mark_selected_ids_i64";
651 pub const ILP_MARK_SELECTED_IDS_U64: &str = "ilp_mark_selected_ids_u64";
652 pub const ILP_VALIDATE_SELECTED_IDS_U32: &str = "ilp_validate_selected_ids_u32";
653 pub const ILP_VALIDATE_SELECTED_IDS_I32: &str = "ilp_validate_selected_ids_i32";
654 pub const ILP_VALIDATE_SELECTED_IDS_I64: &str = "ilp_validate_selected_ids_i64";
655 pub const ILP_VALIDATE_SELECTED_IDS_U64: &str = "ilp_validate_selected_ids_u64";
656 pub const ILP_BROADCAST_CANDIDATE_FLAG: &str = "ilp_broadcast_candidate_flag";
657 pub const ILP_COO_FILL_FROM_MASK: &str = "ilp_coo_fill_from_mask";
658 pub const ILP_CSR_HISTOGRAM: &str = "ilp_csr_histogram";
659 pub const ILP_REDUCE_SUM_F32: &str = "ilp_reduce_sum_f32";
660 pub const ILP_REDUCE_SUM_F64: &str = "ilp_reduce_sum_f64";
661}
662
663pub mod ilp_credit_kernels {
665 pub const ILP_COO_FILL: &str = "ilp_coo_fill";
666 pub const ILP_CREDIT_FORWARD_F32: &str = "ilp_credit_forward_f32";
667 pub const ILP_CREDIT_FORWARD_F64: &str = "ilp_credit_forward_f64";
668 pub const ILP_CREDIT_BACKWARD_F32: &str = "ilp_credit_backward_f32";
669 pub const ILP_CREDIT_BACKWARD_F64: &str = "ilp_credit_backward_f64";
670}
671
672pub mod ilp_exact_kernels {
674 pub const ILP_EXACT_SCORE: &str = "ilp_exact_score";
675 pub const ILP_EXACT_SCORE_U32: &str = "ilp_exact_score_u32";
676 pub const ILP_EXACT_SCORE_CHAIN_SMEM: &str = "ilp_exact_score_chain_smem";
677 pub const ILP_EXACT_SCORE_CHAIN_SMEM_U32: &str = "ilp_exact_score_chain_smem_u32";
678 pub const ILP_EXACT_SELECT_TOPK: &str = "ilp_exact_select_topk";
679}
680
681pub mod pir_kernels {
683 pub const PIR_PACK_KEYS: &str = "pir_pack_keys";
684 pub const PIR_HASH_KEYS: &str = "pir_hash_keys";
685 pub const PIR_MARK_UNIQUE: &str = "pir_mark_unique";
686 pub const PIR_FIND_EXISTING: &str = "pir_find_existing";
687 pub const PIR_MARK_NEW_GROUPS: &str = "pir_mark_new_groups";
688 pub const PIR_BUILD_GROUP_IDS: &str = "pir_build_group_ids";
689 pub const PIR_FILL_CHILD_PARENTS: &str = "pir_fill_child_parents";
690 pub const PIR_MARK_UNIQUE_PAIRS: &str = "pir_mark_unique_pairs";
691 pub const PIR_COMPACT_PAIRS: &str = "pir_compact_pairs";
692 pub const PIR_COUNT_CHILDREN: &str = "pir_count_children";
693 pub const PIR_WRITE_CHILD_OFFSETS: &str = "pir_write_child_offsets";
694 pub const PIR_GATHER_CHILDREN: &str = "pir_gather_children";
695 pub const PIR_BUILD_GRAPH_CHILD_COUNTS: &str = "pir_build_graph_child_counts";
696 pub const PIR_SUM_COUNTS: &str = "pir_sum_counts";
697 pub const PIR_EMIT_NODES_AND_IDS: &str = "pir_emit_nodes_and_ids";
698 pub const PIR_UPDATE_COUNTS: &str = "pir_update_counts";
699}
700
701pub mod cnf_kernels {
703 pub const CNF_REACHABILITY_INIT: &str = "cnf_reachability_init";
704 pub const CNF_REACHABILITY_BFS: &str = "cnf_reachability_bfs";
705 pub const CNF_MARK_LEAF_CHOICE: &str = "cnf_mark_leaf_choice";
706 pub const CNF_ASSIGN_LEAF_VAR: &str = "cnf_assign_leaf_var";
707 pub const CNF_ASSIGN_CHOICE_VAR: &str = "cnf_assign_choice_var";
708 pub const CNF_MARK_NODE_VARS: &str = "cnf_mark_node_vars";
709 pub const CNF_COUNT_CLAUSES: &str = "cnf_count_clauses";
710 pub const CNF_CAPTURE_LAST_COUNTS: &str = "cnf_capture_last_counts";
711 pub const CNF_COMPUTE_LEAF_CHOICE_TOTALS: &str = "cnf_compute_leaf_choice_totals";
712 pub const CNF_COMPUTE_TOTALS: &str = "cnf_compute_totals";
713 pub const CNF_ASSIGN_NODE_VAR: &str = "cnf_assign_node_var";
714 pub const CNF_EMIT_CLAUSES: &str = "cnf_emit_clauses";
715 pub const CNF_SET_CLAUSE_END: &str = "cnf_set_clause_end";
716}
717
718pub mod weights_kernels {
720 pub const WEIGHTS_FILL_LEAF: &str = "weights_fill_leaf";
721 pub const WEIGHTS_FILL_CHOICE: &str = "weights_fill_choice";
722 pub const WEIGHTS_COUNT_LIFT_EXACT: &str = "weights_count_lift_exact";
723 pub const WEIGHTS_SET_EVIDENCE_FROM_NODES: &str = "weights_set_evidence_from_nodes";
724 pub const WEIGHTS_APPLY_EVIDENCE: &str = "weights_apply_evidence";
725 pub const WEIGHTS_MAP_NODES_TO_VARS: &str = "weights_map_nodes_to_vars";
726 pub const WEIGHTS_FORCE_VAR_FALSE: &str = "weights_force_var_false";
727 pub const WEIGHTS_RESTORE_VAR_FALSE: &str = "weights_restore_var_false";
728 pub const WEIGHTS_FORCE_VAR_TRUE: &str = "weights_force_var_true";
729 pub const WEIGHTS_RESTORE_VAR_TRUE: &str = "weights_restore_var_true";
730 pub const WEIGHTS_COPY_SLOT_TO_BATCH: &str = "weights_copy_slot_to_batch";
731 pub const WEIGHTS_APPLY_QUERY_VARS: &str = "weights_apply_query_vars";
732 pub const WEIGHTS_RESTORE_QUERY_VARS: &str = "weights_restore_query_vars";
733 pub const WEIGHTS_APPLY_QUERY_VARS_FALSE_BATCHED: &str =
734 "weights_apply_query_vars_false_batched";
735 pub const WEIGHTS_RESTORE_QUERY_VARS_FALSE_BATCHED: &str =
736 "weights_restore_query_vars_false_batched";
737 pub const WEIGHTS_APPLY_QUERY_VARS_TRUE_BATCHED: &str = "weights_apply_query_vars_true_batched";
738 pub const WEIGHTS_RESTORE_QUERY_VARS_TRUE_BATCHED: &str =
739 "weights_restore_query_vars_true_batched";
740}
741
742pub mod d4_kernels {
745 pub const D4_VALIDATE_CNF: &str = "d4_validate_cnf";
746 pub const D4_LEVELIZE_COUNTS: &str = "d4_levelize_counts";
747 pub const D4_LEVELIZE_EMIT: &str = "d4_levelize_emit";
748 pub const D4_FRONTIER_PREPARE: &str = "d4_frontier_prepare";
750 pub const D4_FRONTIER_EXPAND: &str = "d4_frontier_expand";
751 pub const D4_FRONTIER_PREPARE_DENSE: &str = "d4_frontier_prepare_dense";
752 pub const D4_FRONTIER_EXPAND_DENSE: &str = "d4_frontier_expand_dense";
753 pub const D4_COMPILE_COUNT: &str = "d4_compile_count";
755 pub const D4_COMPILE_EMIT: &str = "d4_compile_emit";
756 pub const D4_CAPTURE_EMIT_META: &str = "d4_capture_emit_meta";
757 pub const D4_SUPPORT_LEVEL: &str = "d4_support_level";
759 pub const D4_SUPPORT_SET_ROOT_BITS: &str = "d4_support_set_root_bits";
760 pub const D4_SMOOTH_COUNT: &str = "d4_smooth_count";
761 pub const D4_SMOOTH_WRAPPER_COUNTS: &str = "d4_smooth_wrapper_counts";
762 pub const D4_SMOOTH_WRAPPER_EDGE_COUNTS_OR: &str = "d4_smooth_wrapper_edge_counts_or";
763 pub const D4_SMOOTH_WRAPPER_EDGE_COUNTS_DEC: &str = "d4_smooth_wrapper_edge_counts_dec";
764 pub const D4_SMOOTH_INIT_NODES: &str = "d4_smooth_init_nodes";
765 pub const D4_SMOOTH_EMIT_LEVEL: &str = "d4_smooth_emit_level";
766 pub const D4_SMOOTH_CHECK_EDGE_CAP: &str = "d4_smooth_check_edge_cap";
767 pub const D4_MARK_VARS_IN_CLAUSES: &str = "d4_mark_vars_in_clauses";
769 pub const D4_MARK_VARS_IN_CIRCUIT: &str = "d4_mark_vars_in_circuit";
770 pub const D4_BUILD_FREE_VAR_MASK: &str = "d4_build_free_var_mask";
771 pub const D4_ASSERT_U32_EQ: &str = "d4_assert_u32_eq";
773 pub const D4_ASSERT_BITSET_VAR: &str = "d4_assert_bitset_var";
774 pub const D4_ASSERT_DENSE_VAR: &str = "d4_assert_dense_var";
775 pub const D4_ASSERT_LEAF_ROOT_AND_DEGREE: &str = "d4_assert_leaf_root_and_degree";
776}
777
778pub mod join_kernels {
780 pub const HASH_JOIN_BUILD: &str = "hash_join_build";
781 pub const HASH_JOIN_PROBE: &str = "hash_join_probe";
782 pub const COMPUTE_COMPOSITE_HASH: &str = "compute_composite_hash";
784 pub const HASH_JOIN_BUCKET_COUNT_V2: &str = "hash_join_bucket_count_v2";
785 pub const HASH_JOIN_SCATTER_V2: &str = "hash_join_scatter_v2";
786 pub const HASH_JOIN_PROBE_V2: &str = "hash_join_probe_v2";
787 pub const HASH_JOIN_PROBE_V2_COUNT_PER_ROW: &str = "hash_join_probe_v2_count_per_row";
788 pub const HASH_JOIN_PROBE_V2_MATERIALIZE: &str = "hash_join_probe_v2_materialize";
789 pub const HASH_JOIN_TOTAL_FROM_SCAN: &str = "hash_join_total_from_scan";
790 pub const HASH_JOIN_CSM_UNMATCHED_MASK: &str = "hash_join_csm_unmatched_mask";
791 pub const HASH_JOIN_SEMI: &str = "hash_join_semi";
792 pub const HASH_JOIN_ANTI: &str = "hash_join_anti";
793 pub const INIT_HASH_TABLE: &str = "init_hash_table";
794 pub const NESTED_LOOP_JOIN_INNER_U32_1KEY_PAIRS: &str = "nested_loop_join_inner_u32_1key_pairs";
800 pub const SORT_MERGE_JOIN_INNER_U32_1KEY_PAIRS: &str = "sort_merge_join_inner_u32_1key_pairs";
808}
809
810pub mod dedup_kernels {
812 pub const MARK_DUPLICATES: &str = "mark_duplicates";
813 pub const MARK_UNIQUE_COLUMNAR: &str = "mark_unique_columnar";
814 pub const MARK_UNIQUE_AND_SCAN_COLUMNAR: &str = "mark_unique_and_scan_columnar";
815 pub const COMPACT_ROWS: &str = "compact_rows";
816 pub const MARK_UNIQUE_FULL_ROW_BYTEWISE: &str = "mark_unique_full_row_bytewise";
817 pub const MARK_DIFF_FULL_ROW_TYPED_SORTED: &str = "mark_diff_full_row_typed_sorted";
818 pub const SMALL_SORT_FULL_ROW_INDICES_TYPED: &str = "small_sort_full_row_indices_typed";
819}
820
821pub mod groupby_kernels {
823 pub const DETECT_GROUP_BOUNDARIES: &str = "detect_group_boundaries";
824 pub const DETECT_BOUNDARIES: &str = "detect_boundaries";
825 pub const EXTRACT_GROUP_KEYS: &str = "extract_group_keys";
826 pub const GROUP_IDS_FROM_BOUNDARIES: &str = "group_ids_from_boundaries";
827 pub const GROUP_START_INDICES: &str = "group_start_indices";
828 pub const CAPTURE_NUM_GROUPS: &str = "capture_num_groups";
829 pub const GROUPBY_COUNT: &str = "groupby_count";
830 pub const GROUPBY_SUM: &str = "groupby_sum";
831 pub const GROUPBY_SUM_U64: &str = "groupby_sum_u64";
832 pub const GROUPBY_MIN: &str = "groupby_min";
833 pub const GROUPBY_MIN_U64: &str = "groupby_min_u64";
834 pub const GROUPBY_MAX: &str = "groupby_max";
835 pub const GROUPBY_MAX_U64: &str = "groupby_max_u64";
836 pub const GROUPBY_LOGSUMEXP_MAX: &str = "groupby_logsumexp_max";
837 pub const GROUPBY_LOGSUMEXP_SUMEXP: &str = "groupby_logsumexp_sumexp";
838 pub const GROUPBY_LOGSUMEXP_FINAL: &str = "groupby_logsumexp_final";
839}
840
841pub mod scan_kernels {
843 pub const BLOCK_INCLUSIVE_SCAN: &str = "block_inclusive_scan";
844 pub const ADD_BLOCK_OFFSETS: &str = "add_block_offsets";
845 pub const EXCLUSIVE_SCAN_MASK: &str = "exclusive_scan_mask";
846 pub const COUNT_MASK: &str = "count_mask";
847 pub const MULTIBLOCK_SCAN_PHASE1: &str = "multiblock_scan_phase1";
849 pub const MULTIBLOCK_SCAN_U32_PHASE1: &str = "multiblock_scan_u32_phase1";
850 pub const MULTIBLOCK_SCAN_PHASE2: &str = "multiblock_scan_phase2";
851 pub const MULTIBLOCK_SCAN_PHASE3: &str = "multiblock_scan_phase3";
852}
853
854pub mod sort_kernels {
856 pub const RADIX_HISTOGRAM: &str = "radix_histogram";
857 pub const RADIX_SCATTER: &str = "radix_scatter";
858 pub const COMPUTE_RANKS: &str = "compute_ranks";
859 pub const RADIX_SCATTER_STABLE: &str = "radix_scatter_stable";
860 pub const COMPUTE_DIGIT_PREFIX_SUMS: &str = "compute_digit_prefix_sums";
861 pub const INIT_INDICES: &str = "init_indices";
862 pub const APPLY_PERMUTATION_U32: &str = "apply_permutation_u32";
863 pub const APPLY_PERMUTATION_BYTES: &str = "apply_permutation_bytes";
864
865 pub const GATHER_KEYS_I32_ORDERED_U32: &str = "gather_keys_i32_ordered_u32";
866 pub const GATHER_KEYS_F32_ORDERED_U32: &str = "gather_keys_f32_ordered_u32";
867 pub const GATHER_KEYS_BOOL_ORDERED_U32: &str = "gather_keys_bool_ordered_u32";
868
869 pub const GATHER_KEYS_U64_LO_U32: &str = "gather_keys_u64_lo_u32";
870 pub const GATHER_KEYS_U64_HI_U32: &str = "gather_keys_u64_hi_u32";
871
872 pub const GATHER_KEYS_I64_LO_U32: &str = "gather_keys_i64_lo_u32";
873 pub const GATHER_KEYS_I64_HI_U32: &str = "gather_keys_i64_hi_u32";
874
875 pub const GATHER_KEYS_F64_LO_U32: &str = "gather_keys_f64_lo_u32";
876 pub const GATHER_KEYS_F64_HI_U32: &str = "gather_keys_f64_hi_u32";
877 pub const CHECK_ASCENDING_SORTED_U32: &str = "check_ascending_sorted_u32";
885}
886
887pub mod filter_kernels {
889 pub const FILTER_COMPARE_U32: &str = "filter_compare_u32";
890 pub const FILTER_COMPARE_I64: &str = "filter_compare_i64";
891 pub const FILTER_COMPARE_F64: &str = "filter_compare_f64";
892 pub const FILTER_COMPARE_I32: &str = "filter_compare_i32";
893 pub const FILTER_COMPARE_U64: &str = "filter_compare_u64";
894 pub const FILTER_COMPARE_F32: &str = "filter_compare_f32";
895 pub const FILTER_COMPARE_U8: &str = "filter_compare_u8";
896 pub const FILTER_COMPARE_U32_SCAN_PHASE1: &str = "filter_compare_u32_scan_phase1";
897 pub const FILTER_COMPARE_F64_SCAN_PHASE1: &str = "filter_compare_f64_scan_phase1";
898 pub const FILTER_COMPARE_F32_SCAN_PHASE1: &str = "filter_compare_f32_scan_phase1";
899 pub const FILTER_COMPARE_U32_COL: &str = "filter_compare_u32_col";
900 pub const FILTER_COMPARE_I32_COL: &str = "filter_compare_i32_col";
901 pub const FILTER_COMPARE_I64_COL: &str = "filter_compare_i64_col";
902 pub const FILTER_COMPARE_U64_COL: &str = "filter_compare_u64_col";
903 pub const FILTER_COMPARE_F32_COL: &str = "filter_compare_f32_col";
904 pub const FILTER_COMPARE_F64_COL: &str = "filter_compare_f64_col";
905 pub const FILTER_COMPARE_U8_COL: &str = "filter_compare_u8_col";
906 pub const FILL_U32_IOTA: &str = "fill_u32_iota";
907 pub const FILL_U32_CONST: &str = "fill_u32_const";
908 pub const MARK_RANDOM_VARS: &str = "mark_random_vars";
909 pub const RANDOM_VAR_TO_BIT_FROM_LIST: &str = "random_var_to_bit_from_list";
910 pub const CHECK_RANDOM_VAR_COUNT: &str = "check_random_var_count";
911 pub const COMPACT_U32_BY_MASK: &str = "compact_u32_by_mask";
912 pub const COMPACT_I64_BY_MASK: &str = "compact_i64_by_mask";
913 pub const COMPACT_F64_BY_MASK: &str = "compact_f64_by_mask";
914 pub const COMPACT_BYTES_BY_MASK: &str = "compact_bytes_by_mask";
915 pub const CAPTURE_COMPACT_COUNT: &str = "capture_compact_count";
916 pub const MASK_CLAMP_ROWS: &str = "mask_clamp_rows";
917 pub const MASK_AND: &str = "mask_and";
918 pub const MASK_OR: &str = "mask_or";
919 pub const MASK_NOT: &str = "mask_not";
920}
921
922pub mod set_ops_kernels {
924 pub const CONCAT_U32: &str = "concat_u32";
925 pub const CONCAT_BYTES: &str = "concat_bytes";
926 pub const SORTED_DIFF_MARK: &str = "sorted_diff_mark";
927}
928
929pub mod pack_kernels {
931 pub const PACK_KEYS: &str = "pack_keys";
933 pub const HASH_PACKED_KEYS: &str = "hash_packed_keys";
935 pub const PACK_AND_HASH_KEYS: &str = "pack_and_hash_keys";
937 pub const PACK_AND_HASH_KEYS_GENERIC: &str = "pack_and_hash_keys_generic";
939 pub const PACK_KEYS_ALIGNED: &str = "pack_keys_aligned";
941 pub const UNPACK_COLUMN: &str = "unpack_column";
943 pub const UNPACK_COLUMN_COUNTED: &str = "unpack_column_counted";
945 pub const GATHER_PACKED_ROWS: &str = "gather_packed_rows";
947 pub const GATHER_PACKED_ROWS_COUNTED: &str = "gather_packed_rows_counted";
949 pub const SCATTER_PACKED_ROWS: &str = "scatter_packed_rows";
951 pub const COMPARE_PACKED_KEYS: &str = "compare_packed_keys";
953 pub const PACK_BOOLS_TO_BITMAP: &str = "pack_bools_to_bitmap";
955}
956
957pub mod circuit_kernels {
959 pub const XGCF_FORWARD_LEVEL: &str = "xgcf_forward_level";
960 pub const XGCF_BACKWARD_LEVEL_PROPAGATE: &str = "xgcf_backward_level_propagate";
961 pub const XGCF_BACKWARD_LEVEL_DECISION_GRAD: &str = "xgcf_backward_level_decision_grad";
962 pub const XGCF_BACKWARD_LEVEL_LIT_GRAD: &str = "xgcf_backward_level_lit_grad";
963 pub const XGCF_FREE_VAR_APPLY_GRAD: &str = "xgcf_free_var_apply_grad";
964 pub const XGCF_FREE_VAR_REDUCE_STAGE: &str = "xgcf_free_var_reduce_stage";
965 pub const XGCF_ADD_SCALAR: &str = "xgcf_add_scalar";
966 pub const XGCF_FORWARD_LEVEL_CACHED: &str = "xgcf_forward_level_cached";
967 pub const XGCF_EVAL_ALL_LEVELS_CACHED: &str = "xgcf_eval_all_levels_cached";
968 pub const XGCF_EVAL_ALL_LEVELS_CACHED_BATCHED: &str = "xgcf_eval_all_levels_cached_batched";
969 pub const XGCF_BACKWARD_LEVEL_PROPAGATE_CACHED: &str = "xgcf_backward_level_propagate_cached";
970 pub const XGCF_BACKWARD_LEVEL_DECISION_GRAD_CACHED: &str =
971 "xgcf_backward_level_decision_grad_cached";
972 pub const XGCF_BACKWARD_LEVEL_LIT_GRAD_CACHED: &str = "xgcf_backward_level_lit_grad_cached";
973 pub const XGCF_BACKWARD_ALL_LEVELS_CACHED: &str = "xgcf_backward_all_levels_cached";
974 pub const XGCF_BACKWARD_ALL_LEVELS_CACHED_BATCHED: &str =
975 "xgcf_backward_all_levels_cached_batched";
976 pub const XGCF_FREE_VAR_APPLY_GRAD_CACHED: &str = "xgcf_free_var_apply_grad_cached";
977 pub const XGCF_FREE_VAR_REDUCE_STAGE_CACHED: &str = "xgcf_free_var_reduce_stage_cached";
978 pub const XGCF_ADD_SCALAR_CACHED: &str = "xgcf_add_scalar_cached";
979 pub const XGCF_SET_ROOT_ADJ_CACHED_BATCHED: &str = "xgcf_set_root_adj_cached_batched";
980 pub const XGCF_COPY_ROOT_CACHED: &str = "xgcf_copy_root_cached";
981 pub const XGCF_COPY_ROOT_CACHED_META: &str = "xgcf_copy_root_cached_meta";
982 pub const XGCF_COPY_ROOT_CACHED_META_BATCHED: &str = "xgcf_copy_root_cached_meta_batched";
983}
984
985pub mod cache_kernels {
987 pub const CACHE_CNF_HASH: &str = "cache_cnf_hash";
988 pub const CACHE_LOOKUP_OR_INSERT: &str = "cache_lookup_or_insert";
989 pub const CACHE_EVICT_LRU: &str = "cache_evict_lru";
990 pub const CACHE_STORE_U8: &str = "cache_store_u8";
991 pub const CACHE_STORE_U32: &str = "cache_store_u32";
992 pub const CACHE_STORE_I32: &str = "cache_store_i32";
993 pub const CACHE_STORE_F64: &str = "cache_store_f64";
994 pub const CACHE_STORE_META: &str = "cache_store_meta";
995}
996
997pub mod sat_kernels {
999 pub const SAT_CDCL_SOLVE: &str = "sat_cdcl_solve";
1000 pub const SAT_CHECK_MODEL: &str = "sat_check_model";
1001 pub const SAT_PROOF_MARK_NEEDED: &str = "sat_proof_mark_needed";
1002 pub const SAT_PROOF_CHECK: &str = "sat_proof_check";
1003 pub const SAT_ASSERT_STATUS: &str = "sat_assert_status";
1004 pub const SAT_ASSERT_OK: &str = "sat_assert_ok";
1005 pub const SAT_XGCF_CNF_COUNTS: &str = "sat_xgcf_cnf_counts";
1006 pub const SAT_XGCF_CNF_EMIT: &str = "sat_xgcf_cnf_emit";
1007 pub const SAT_XGCF_CNF_CAPTURE_LAST_COUNTS: &str = "sat_xgcf_cnf_capture_last_counts";
1008 pub const SAT_XGCF_CNF_COMPUTE_TOTALS: &str = "sat_xgcf_cnf_compute_totals";
1009 pub const SAT_CNF_WRITE_TERMINATOR: &str = "sat_cnf_write_terminator";
1010 pub const SAT_CNF_COPY_INTO: &str = "sat_cnf_copy_into";
1011 pub const SAT_SHIFT_OFFSETS: &str = "sat_shift_offsets";
1012 pub const SAT_XGCF_WRITE_ROOT_UNIT_CLAUSE: &str = "sat_xgcf_write_root_unit_clause";
1013 pub const SAT_NOT_PHI_COUNTS: &str = "sat_not_phi_counts";
1014 pub const SAT_EMIT_NOT_PHI: &str = "sat_emit_not_phi";
1015}
1016
1017pub const DEFAULT_JOIN_MAX_OUTPUT: usize = 1_000_000;
1020
1021pub const NESTED_LOOP_TOTAL_THRESHOLD: u64 = 4_000_000;
1042
1043#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1045#[repr(u8)]
1046pub enum CompareOp {
1047 Eq = 0,
1048 Ne = 1,
1049 Lt = 2,
1050 Le = 3,
1051 Gt = 4,
1052 Ge = 5,
1053}
1054
1055#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1057pub enum JoinType {
1058 Inner,
1060 Semi,
1062 Anti,
1064 LeftOuter,
1066}
1067
1068struct PackedKeyData {
1070 hashes: crate::memory::TrackedCudaSlice<u64>,
1072 packed_keys: crate::memory::TrackedCudaSlice<u8>,
1074 key_bytes: u32,
1076}
1077
1078struct JoinHashTableV2 {
1079 bucket_counts: crate::memory::TrackedCudaSlice<u32>,
1080 bucket_offsets: crate::memory::TrackedCudaSlice<u32>,
1081 bucket_entries: crate::memory::TrackedCudaSlice<u32>,
1082 bucket_entry_hashes: crate::memory::TrackedCudaSlice<u64>,
1083 bucket_mask: u32,
1084}
1085
1086pub struct HashTableU64 {
1088 pub bucket_counts: crate::memory::TrackedCudaSlice<u32>,
1089 pub bucket_offsets: crate::memory::TrackedCudaSlice<u32>,
1090 pub bucket_entries: crate::memory::TrackedCudaSlice<u32>,
1091 pub bucket_entry_hashes: crate::memory::TrackedCudaSlice<u64>,
1092 pub bucket_mask: u32,
1093}
1094
1095pub struct JoinIndexV2 {
1100 right_num_rows: u32,
1101 right_keys: Vec<usize>,
1102 key_bytes: u32,
1103 packed_keys: crate::memory::TrackedCudaSlice<u8>,
1104 table: JoinHashTableV2,
1105}
1106
1107impl JoinIndexV2 {
1108 pub fn right_keys(&self) -> &[usize] {
1110 &self.right_keys
1111 }
1112
1113 pub fn right_num_rows(&self) -> u32 {
1115 self.right_num_rows
1116 }
1117
1118 pub fn estimated_bytes(&self) -> u64 {
1120 let mut bytes = 0u64;
1121 bytes = bytes.saturating_add(self.packed_keys.len() as u64);
1122 bytes = bytes.saturating_add(self.table.bucket_counts.len() as u64 * 4);
1123 bytes = bytes.saturating_add(self.table.bucket_offsets.len() as u64 * 4);
1124 bytes = bytes.saturating_add(self.table.bucket_entries.len() as u64 * 4);
1125 bytes = bytes.saturating_add(self.table.bucket_entry_hashes.len() as u64 * 8);
1126 bytes
1127 }
1128}
1129
1130pub struct CudaKernelProvider {
1151 device: Arc<CudaDevice>,
1153 memory: Arc<GpuMemoryManager>,
1155 transfer_tracker: HostTransferTracker,
1157 ptx_load_profile: Option<PtxLoadProfile>,
1159 d2h_transfer_count: AtomicU64,
1161 untracked_metadata_dtoh_count: AtomicU64,
1168 strict_deterministic_d2h: AtomicBool,
1174 deterministic_d2h_violations: AtomicU64,
1178 recorded_op_stream: OnceLock<crate::device_runtime::StreamId>,
1187 csm_invocations: AtomicU64,
1195 csm_cuda_graph_captures: AtomicU64,
1197 csm_cuda_graph_launches: AtomicU64,
1199 csm_cuda_graph_fallbacks: AtomicU64,
1201 csm_cuda_graph_cache_hits: AtomicU64,
1203 small_full_row_sort_invocations: AtomicU64,
1207 csm_cuda_graph_cache: Mutex<HashMap<CsmCudaGraphKey, CsmCudaGraphEntry>>,
1209 wcoj_layout_fast_path_hit_count: AtomicU64,
1216 wcoj_layout_sort_invocation_count: AtomicU64,
1221 kclique_metadata_build_count: AtomicU64,
1223 kclique_metadata_build_nanos: AtomicU64,
1226 wcoj_triangle_hg_dispatch_count: AtomicU64,
1229 #[cfg(feature = "wcoj-phase-timing")]
1236 last_triangle_phase_timing:
1237 std::sync::Mutex<Option<crate::wcoj_phase_timing::WcojTrianglePhaseTiming>>,
1238}
1239
1240#[derive(Default)]
1241struct HostTransferTracker {
1242 dtoh_bytes: AtomicU64,
1243 htod_bytes: AtomicU64,
1244 dtoh_calls: AtomicU64,
1245 htod_calls: AtomicU64,
1246 launch_metadata_htod_bytes: AtomicU64,
1247 launch_metadata_htod_calls: AtomicU64,
1248}
1249
1250#[derive(Debug, Clone, Copy)]
1251pub struct HostTransferStats {
1252 pub dtoh_bytes: u64,
1253 pub htod_bytes: u64,
1254 pub dtoh_calls: u64,
1255 pub htod_calls: u64,
1256}
1257
1258#[derive(Debug, Clone, Copy, Default)]
1259pub struct HostLaunchMetadataTransferStats {
1260 pub htod_bytes: u64,
1261 pub htod_calls: u64,
1262}
1263
1264impl HostTransferTracker {
1265 fn record_dtoh(&self, bytes: u64) {
1266 self.dtoh_calls.fetch_add(1, Ordering::Relaxed);
1267 self.dtoh_bytes.fetch_add(bytes, Ordering::Relaxed);
1268 }
1269
1270 fn record_htod(&self, bytes: u64) {
1271 self.htod_calls.fetch_add(1, Ordering::Relaxed);
1272 self.htod_bytes.fetch_add(bytes, Ordering::Relaxed);
1273 }
1274
1275 fn record_htod_launch_metadata(&self, bytes: u64) {
1276 self.launch_metadata_htod_calls
1277 .fetch_add(1, Ordering::Relaxed);
1278 self.launch_metadata_htod_bytes
1279 .fetch_add(bytes, Ordering::Relaxed);
1280 }
1281
1282 fn snapshot(&self) -> HostTransferStats {
1283 HostTransferStats {
1284 dtoh_bytes: self.dtoh_bytes.load(Ordering::Relaxed),
1285 htod_bytes: self.htod_bytes.load(Ordering::Relaxed),
1286 dtoh_calls: self.dtoh_calls.load(Ordering::Relaxed),
1287 htod_calls: self.htod_calls.load(Ordering::Relaxed),
1288 }
1289 }
1290
1291 fn launch_metadata_snapshot(&self) -> HostLaunchMetadataTransferStats {
1292 HostLaunchMetadataTransferStats {
1293 htod_bytes: self.launch_metadata_htod_bytes.load(Ordering::Relaxed),
1294 htod_calls: self.launch_metadata_htod_calls.load(Ordering::Relaxed),
1295 }
1296 }
1297
1298 fn reset(&self) {
1299 self.dtoh_bytes.store(0, Ordering::Relaxed);
1300 self.htod_bytes.store(0, Ordering::Relaxed);
1301 self.dtoh_calls.store(0, Ordering::Relaxed);
1302 self.htod_calls.store(0, Ordering::Relaxed);
1303 self.launch_metadata_htod_bytes.store(0, Ordering::Relaxed);
1304 self.launch_metadata_htod_calls.store(0, Ordering::Relaxed);
1305 }
1306}
1307
1308impl CudaKernelProvider {
1309 pub fn new(device: Arc<CudaDevice>, memory: Arc<GpuMemoryManager>) -> Result<Self> {
1328 let profiling = warmup_profiling_enabled();
1329 let ptx_load_profile = Self::load_all_kernel_modules(&device, profiling)?;
1330
1331 Ok(Self {
1332 device,
1333 memory,
1334 transfer_tracker: HostTransferTracker::default(),
1335 ptx_load_profile,
1336 d2h_transfer_count: AtomicU64::new(0),
1337 untracked_metadata_dtoh_count: AtomicU64::new(0),
1338 strict_deterministic_d2h: AtomicBool::new(false),
1339 deterministic_d2h_violations: AtomicU64::new(0),
1340 recorded_op_stream: OnceLock::new(),
1341 csm_invocations: AtomicU64::new(0),
1342 csm_cuda_graph_captures: AtomicU64::new(0),
1343 csm_cuda_graph_launches: AtomicU64::new(0),
1344 csm_cuda_graph_fallbacks: AtomicU64::new(0),
1345 csm_cuda_graph_cache_hits: AtomicU64::new(0),
1346 small_full_row_sort_invocations: AtomicU64::new(0),
1347 csm_cuda_graph_cache: Mutex::new(HashMap::new()),
1348 wcoj_layout_fast_path_hit_count: AtomicU64::new(0),
1349 wcoj_layout_sort_invocation_count: AtomicU64::new(0),
1350 kclique_metadata_build_count: AtomicU64::new(0),
1351 kclique_metadata_build_nanos: AtomicU64::new(0),
1352 wcoj_triangle_hg_dispatch_count: AtomicU64::new(0),
1353 #[cfg(feature = "wcoj-phase-timing")]
1354 last_triangle_phase_timing: std::sync::Mutex::new(None),
1355 })
1356 }
1357
1358 pub fn with_runtime(device: Arc<CudaDevice>, memory: Arc<GpuMemoryManager>) -> Result<Self> {
1406 if memory.runtime().is_none() {
1407 return Err(XlogError::Kernel(
1408 "CudaKernelProvider::with_runtime requires a GpuMemoryManager built via \
1409 GpuMemoryManager::with_runtime; got a manager with no runtime attached"
1410 .to_string(),
1411 ));
1412 }
1413 Self::new(device, memory)
1414 }
1415
1416 fn env_flag(name: &str) -> bool {
1419 std::env::var(name)
1420 .map(|v| !v.is_empty() && v != "0")
1421 .unwrap_or(false)
1422 }
1423
1424 pub(crate) fn use_recorded_filters_env() -> bool {
1436 Self::env_flag("XLOG_USE_RECORDED_FILTERS") || Self::env_flag("XLOG_USE_RECORDED_OPS")
1437 }
1438
1439 pub(crate) fn use_recorded_sort_env() -> bool {
1446 Self::env_flag("XLOG_USE_RECORDED_SORT") || Self::env_flag("XLOG_USE_RECORDED_OPS")
1447 }
1448
1449 pub(crate) fn use_recorded_dedup_env() -> bool {
1454 Self::env_flag("XLOG_USE_RECORDED_DEDUP") || Self::env_flag("XLOG_USE_RECORDED_OPS")
1455 }
1456
1457 pub(crate) fn use_recorded_groupby_env() -> bool {
1463 Self::env_flag("XLOG_USE_RECORDED_GROUPBY") || Self::env_flag("XLOG_USE_RECORDED_OPS")
1464 }
1465
1466 pub(crate) fn use_recorded_hash_join_env() -> bool {
1474 Self::env_flag("XLOG_USE_RECORDED_HASH_JOIN") || Self::env_flag("XLOG_USE_RECORDED_OPS")
1475 }
1476
1477 pub(crate) fn use_recorded_csm_env() -> bool {
1485 Self::env_flag("XLOG_USE_RECORDED_CSM") || Self::env_flag("XLOG_USE_RECORDED_OPS")
1486 }
1487
1488 pub(crate) fn use_csm_cuda_graph_env() -> bool {
1494 Self::env_flag("XLOG_USE_CSM_CUDA_GRAPH") || Self::env_flag("XLOG_USE_CUDA_GRAPHS")
1495 }
1496
1497 #[doc(hidden)]
1511 pub fn csm_invocations(&self) -> u64 {
1512 self.csm_invocations.load(Ordering::Relaxed)
1513 }
1514
1515 #[doc(hidden)]
1516 pub fn csm_cuda_graph_captures(&self) -> u64 {
1517 self.csm_cuda_graph_captures.load(Ordering::Relaxed)
1518 }
1519
1520 #[doc(hidden)]
1521 pub fn csm_cuda_graph_launches(&self) -> u64 {
1522 self.csm_cuda_graph_launches.load(Ordering::Relaxed)
1523 }
1524
1525 #[doc(hidden)]
1526 pub fn csm_cuda_graph_fallbacks(&self) -> u64 {
1527 self.csm_cuda_graph_fallbacks.load(Ordering::Relaxed)
1528 }
1529
1530 #[doc(hidden)]
1531 pub fn csm_cuda_graph_cache_hits(&self) -> u64 {
1532 self.csm_cuda_graph_cache_hits.load(Ordering::Relaxed)
1533 }
1534
1535 #[doc(hidden)]
1536 pub fn small_full_row_sort_invocations(&self) -> u64 {
1537 self.small_full_row_sort_invocations.load(Ordering::Relaxed)
1538 }
1539
1540 pub(crate) fn recorded_op_stream_or_init(&self) -> Option<crate::device_runtime::StreamId> {
1560 if let Some(s) = self.recorded_op_stream.get() {
1561 return Some(*s);
1562 }
1563 let runtime = self.memory.runtime()?;
1564 let stream = runtime.stream_pool().acquire().ok()?;
1565 let _ = self.recorded_op_stream.set(stream);
1566 self.recorded_op_stream.get().copied()
1567 }
1568
1569 #[cfg(feature = "wcoj-phase-timing")]
1579 pub fn take_wcoj_triangle_phase_timing(
1580 &self,
1581 ) -> Option<crate::wcoj_phase_timing::WcojTrianglePhaseTiming> {
1582 self.last_triangle_phase_timing
1583 .lock()
1584 .ok()
1585 .and_then(|mut g| g.take())
1586 }
1587
1588 #[cfg(feature = "wcoj-phase-timing")]
1592 #[allow(dead_code)]
1593 pub(crate) fn put_wcoj_triangle_phase_timing(
1594 &self,
1595 timing: crate::wcoj_phase_timing::WcojTrianglePhaseTiming,
1596 ) {
1597 if let Ok(mut g) = self.last_triangle_phase_timing.lock() {
1598 *g = Some(timing);
1599 }
1600 }
1601
1602 pub fn wcoj_layout_fast_path_hit_count(&self) -> u64 {
1609 self.wcoj_layout_fast_path_hit_count.load(Ordering::Relaxed)
1610 }
1611
1612 pub fn wcoj_triangle_hg_dispatch_count(&self) -> u64 {
1615 self.wcoj_triangle_hg_dispatch_count.load(Ordering::Relaxed)
1616 }
1617
1618 pub fn reset_wcoj_layout_fast_path_hit_count(&self) {
1621 self.wcoj_layout_fast_path_hit_count
1622 .store(0, Ordering::Relaxed);
1623 }
1624
1625 pub fn wcoj_layout_sort_invocation_count(&self) -> u64 {
1628 self.wcoj_layout_sort_invocation_count
1629 .load(Ordering::Relaxed)
1630 }
1631
1632 pub fn reset_wcoj_layout_sort_invocation_count(&self) {
1634 self.wcoj_layout_sort_invocation_count
1635 .store(0, Ordering::Relaxed);
1636 }
1637
1638 pub fn kclique_metadata_build_count(&self) -> u64 {
1641 self.kclique_metadata_build_count.load(Ordering::Relaxed)
1642 }
1643
1644 pub fn kclique_metadata_build_nanos(&self) -> u64 {
1647 self.kclique_metadata_build_nanos.load(Ordering::Relaxed)
1648 }
1649
1650 pub fn reset_kclique_metadata_build_metrics(&self) {
1652 self.kclique_metadata_build_count
1653 .store(0, Ordering::Relaxed);
1654 self.kclique_metadata_build_nanos
1655 .store(0, Ordering::Relaxed);
1656 }
1657
1658 pub(crate) fn record_wcoj_layout_fast_path_hit(&self) {
1662 self.wcoj_layout_fast_path_hit_count
1663 .fetch_add(1, Ordering::Relaxed);
1664 }
1665
1666 pub(crate) fn record_wcoj_layout_sort_invocation(&self) {
1668 self.wcoj_layout_sort_invocation_count
1669 .fetch_add(1, Ordering::Relaxed);
1670 }
1671
1672 pub(crate) fn record_kclique_metadata_build_nanos(&self, nanos: u128) {
1674 self.kclique_metadata_build_count
1675 .fetch_add(1, Ordering::Relaxed);
1676 let nanos = u64::try_from(nanos).unwrap_or(u64::MAX);
1677 self.kclique_metadata_build_nanos
1678 .fetch_add(nanos, Ordering::Relaxed);
1679 }
1680
1681 #[doc(hidden)]
1684 pub fn record_wcoj_triangle_hg_dispatch(&self) {
1685 self.wcoj_triangle_hg_dispatch_count
1686 .fetch_add(1, Ordering::Relaxed);
1687 }
1688
1689 pub fn device(&self) -> &Arc<CudaDevice> {
1691 &self.device
1692 }
1693
1694 pub fn memory(&self) -> &Arc<GpuMemoryManager> {
1696 &self.memory
1697 }
1698
1699 pub fn ptx_load_profile(&self) -> Option<&PtxLoadProfile> {
1701 self.ptx_load_profile.as_ref()
1702 }
1703
1704 pub fn reset_host_transfer_stats(&self) {
1706 self.transfer_tracker.reset();
1707 }
1708
1709 pub fn host_transfer_stats(&self) -> HostTransferStats {
1711 self.transfer_tracker.snapshot()
1712 }
1713
1714 pub fn host_launch_metadata_transfer_stats(&self) -> HostLaunchMetadataTransferStats {
1717 self.transfer_tracker.launch_metadata_snapshot()
1718 }
1719
1720 pub fn d2h_transfer_count(&self) -> u64 {
1726 self.d2h_transfer_count.load(Ordering::Relaxed)
1727 }
1728
1729 pub fn reset_d2h_transfer_count(&self) {
1731 self.d2h_transfer_count.store(0, Ordering::Relaxed);
1732 }
1733
1734 pub fn untracked_metadata_dtoh_count(&self) -> u64 {
1737 self.untracked_metadata_dtoh_count.load(Ordering::Relaxed)
1738 }
1739
1740 pub fn reset_untracked_metadata_dtoh_count(&self) {
1742 self.untracked_metadata_dtoh_count
1743 .store(0, Ordering::Relaxed);
1744 }
1745
1746 pub fn enable_strict_deterministic_d2h(&self) {
1763 self.strict_deterministic_d2h.store(true, Ordering::Relaxed);
1764 }
1765
1766 pub fn disable_strict_deterministic_d2h(&self) {
1768 self.strict_deterministic_d2h
1769 .store(false, Ordering::Relaxed);
1770 }
1771
1772 pub fn strict_deterministic_d2h_enabled(&self) -> bool {
1774 self.strict_deterministic_d2h.load(Ordering::Relaxed)
1775 }
1776
1777 pub fn deterministic_d2h_violation_count(&self) -> u64 {
1779 self.deterministic_d2h_violations.load(Ordering::Relaxed)
1780 }
1781
1782 pub fn reset_deterministic_d2h_violations(&self) {
1784 self.deterministic_d2h_violations
1785 .store(0, Ordering::Relaxed);
1786 }
1787
1788 pub(crate) fn check_deterministic_d2h(&self, op: &'static str, bytes: u64) -> Result<()> {
1794 if self.strict_deterministic_d2h.load(Ordering::Relaxed) {
1795 self.deterministic_d2h_violations
1796 .fetch_add(1, Ordering::Relaxed);
1797 return Err(XlogError::Execution(format!(
1798 "deterministic D2H gate: {} attempted to copy {} bytes from device to host",
1799 op, bytes
1800 )));
1801 }
1802 Ok(())
1803 }
1804
1805 fn dtoh_sync_copy_into_tracked<T: DeviceRepr, Src: DevicePtr<T>>(
1806 &self,
1807 src: &Src,
1808 dst: &mut [T],
1809 ) -> Result<()> {
1810 let bytes = std::mem::size_of::<T>()
1811 .checked_mul(dst.len())
1812 .ok_or_else(|| XlogError::Kernel("dtoh size overflow".to_string()))?;
1813 self.check_deterministic_d2h("dtoh_sync_copy_into_tracked", bytes as u64)?;
1814 self.transfer_tracker.record_dtoh(bytes as u64);
1815 self.device
1816 .inner()
1817 .dtoh_sync_copy_into(src, dst)
1818 .map_err(|e| XlogError::Kernel(format!("Failed to copy from device: {}", e)))
1819 }
1820
1821 pub const DTOH_SMALL_METADATA_MAX_BYTES: usize = 4096;
1826
1827 pub fn dtoh_small_metadata_untracked<T: DeviceRepr + Default + Copy>(
1852 &self,
1853 src: &crate::memory::TrackedCudaSlice<T>,
1854 count: usize,
1855 ) -> Result<Vec<T>> {
1856 let bytes = count.checked_mul(std::mem::size_of::<T>()).ok_or_else(|| {
1857 XlogError::Kernel("dtoh_small_metadata_untracked: byte size overflow".to_string())
1858 })?;
1859 if bytes > Self::DTOH_SMALL_METADATA_MAX_BYTES {
1860 return Err(XlogError::Kernel(format!(
1861 "dtoh_small_metadata_untracked: requested {} bytes exceeds metadata cap of {} bytes \
1862 (this is metadata-only; use download_column* for data-plane transfers)",
1863 bytes,
1864 Self::DTOH_SMALL_METADATA_MAX_BYTES
1865 )));
1866 }
1867 if count > src.len() {
1868 return Err(XlogError::Kernel(format!(
1869 "dtoh_small_metadata_untracked: count={count} > src.len={}",
1870 src.len()
1871 )));
1872 }
1873 if count == 0 {
1874 return Ok(Vec::new());
1875 }
1876 let slice = src.try_slice(0..count).ok_or_else(|| {
1877 XlogError::Kernel(format!(
1878 "dtoh_small_metadata_untracked: try_slice(0..{count}) failed"
1879 ))
1880 })?;
1881 let mut buf: Vec<T> = vec![T::default(); count];
1882 self.untracked_metadata_dtoh_count
1883 .fetch_add(1, Ordering::Relaxed);
1884 self.device
1885 .inner()
1886 .dtoh_sync_copy_into(&slice, &mut buf)
1887 .map_err(|e| {
1888 XlogError::Kernel(format!("dtoh_small_metadata_untracked: copy failed: {}", e))
1889 })?;
1890 Ok(buf)
1891 }
1892
1893 pub fn dtoh_scalar_untracked<T: DeviceRepr + Default + Copy>(
1901 &self,
1902 src: &crate::memory::TrackedCudaSlice<T>,
1903 index: usize,
1904 ) -> Result<T> {
1905 if index >= src.len() {
1906 return Err(XlogError::Kernel(format!(
1907 "dtoh_scalar_untracked: index={} >= len={}",
1908 index,
1909 src.len()
1910 )));
1911 }
1912 let slice = src.try_slice(index..index + 1).ok_or_else(|| {
1913 XlogError::Kernel(format!(
1914 "dtoh_scalar_untracked: slice failed at index={}",
1915 index
1916 ))
1917 })?;
1918 let mut buf = [T::default()];
1919 self.untracked_metadata_dtoh_count
1920 .fetch_add(1, Ordering::Relaxed);
1921 self.device
1922 .inner()
1923 .dtoh_sync_copy_into(&slice, &mut buf)
1924 .map_err(|e| XlogError::Kernel(format!("dtoh_scalar_untracked: copy failed: {}", e)))?;
1925 Ok(buf[0])
1926 }
1927
1928 pub fn htod_sync_copy_into_tracked<T: DeviceRepr, Dst: cudarc::driver::DevicePtrMut<T>>(
1930 &self,
1931 src: &[T],
1932 dst: &mut Dst,
1933 ) -> Result<()> {
1934 let bytes = std::mem::size_of::<T>()
1935 .checked_mul(src.len())
1936 .ok_or_else(|| XlogError::Kernel("htod size overflow".to_string()))?;
1937 self.transfer_tracker.record_htod(bytes as u64);
1938 self.device
1939 .inner()
1940 .htod_sync_copy_into(src, dst)
1941 .map_err(|e| XlogError::Kernel(format!("Failed to copy to device: {}", e)))
1942 }
1943
1944 pub fn htod_sync_copy_tracked<T: DeviceRepr>(
1947 &self,
1948 src: &[T],
1949 ) -> Result<cudarc::driver::CudaSlice<T>> {
1950 let bytes = std::mem::size_of::<T>()
1951 .checked_mul(src.len())
1952 .ok_or_else(|| XlogError::Kernel("htod size overflow".to_string()))?;
1953 self.transfer_tracker.record_htod(bytes as u64);
1954 self.device
1955 .inner()
1956 .htod_sync_copy(src)
1957 .map_err(|e| XlogError::Kernel(format!("Failed to copy to device: {}", e)))
1958 }
1959
1960 pub fn htod_launch_metadata_sync_copy_into<
1963 T: DeviceRepr,
1964 Dst: cudarc::driver::DevicePtrMut<T>,
1965 >(
1966 &self,
1967 src: &[T],
1968 dst: &mut Dst,
1969 ) -> Result<()> {
1970 let bytes = std::mem::size_of::<T>()
1971 .checked_mul(src.len())
1972 .ok_or_else(|| XlogError::Kernel("launch metadata htod size overflow".to_string()))?;
1973 self.transfer_tracker
1974 .record_htod_launch_metadata(bytes as u64);
1975 self.device
1976 .inner()
1977 .htod_sync_copy_into(src, dst)
1978 .map_err(|e| {
1979 XlogError::Kernel(format!("Failed to copy launch metadata to device: {}", e))
1980 })
1981 }
1982
1983 pub(crate) fn htod_launch_metadata_async_copy_one<T: DeviceRepr>(
1986 &self,
1987 src: &T,
1988 dst: &TrackedCudaSlice<T>,
1989 stream: &CudaStream,
1990 context: &str,
1991 ) -> Result<()> {
1992 let bytes = std::mem::size_of::<T>();
1993 self.transfer_tracker
1994 .record_htod_launch_metadata(bytes as u64);
1995 unsafe {
1996 let res = cudarc::driver::sys::cuMemcpyHtoDAsync_v2(
1997 *dst.device_ptr(),
1998 src as *const T as *const c_void,
1999 bytes,
2000 stream.cu_stream(),
2001 );
2002 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
2003 return Err(XlogError::Kernel(format!(
2004 "{context}: launch metadata H2D failed: {res:?}"
2005 )));
2006 }
2007 }
2008 Ok(())
2009 }
2010
2011 pub fn exclusive_scan_u32_inplace(
2040 &self,
2041 data: &mut crate::memory::TrackedCudaSlice<u32>,
2042 n: u32,
2043 ) -> Result<()> {
2044 if n as usize > data.len() {
2045 return Err(XlogError::Kernel(format!(
2046 "exclusive_scan_u32_inplace: n={} exceeds slice len={}",
2047 n,
2048 data.len()
2049 )));
2050 }
2051 self.multiblock_scan_u32_inplace(data, n)
2052 }
2053
2054 fn multiblock_scan_u32_inplace(
2055 &self,
2056 data: &mut crate::memory::TrackedCudaSlice<u32>,
2057 n: u32,
2058 ) -> Result<()> {
2059 if n == 0 {
2060 return Ok(());
2061 }
2062
2063 let device = self.device.inner();
2064 let block_size = 256u32;
2065
2066 if n <= block_size {
2067 let phase2_fn = device
2068 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE2)
2069 .ok_or_else(|| {
2070 XlogError::Kernel("Failed to get multiblock_scan_phase2 kernel".to_string())
2071 })?;
2072
2073 unsafe {
2075 phase2_fn.clone().launch(
2076 LaunchConfig {
2077 grid_dim: (1, 1, 1),
2078 block_dim: (block_size, 1, 1),
2079 shared_mem_bytes: 0,
2080 },
2081 (&mut *data, n),
2082 )
2083 }
2084 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase2 failed: {}", e)))?;
2085
2086 self.device.synchronize()?;
2087 return Ok(());
2088 }
2089
2090 let num_blocks = n.div_ceil(block_size);
2091 let mut block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
2092
2093 let phase1_u32_fn = device
2094 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_U32_PHASE1)
2095 .ok_or_else(|| {
2096 XlogError::Kernel("Failed to get multiblock_scan_u32_phase1 kernel".to_string())
2097 })?;
2098
2099 unsafe {
2101 phase1_u32_fn.clone().launch(
2102 LaunchConfig {
2103 grid_dim: (num_blocks, 1, 1),
2104 block_dim: (block_size, 1, 1),
2105 shared_mem_bytes: 0,
2106 },
2107 (&mut *data, &mut block_sums, n),
2108 )
2109 }
2110 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_u32_phase1 failed: {}", e)))?;
2111 self.device.synchronize()?;
2112
2113 if num_blocks > 1 {
2114 self.multiblock_scan_u32_inplace(&mut block_sums, num_blocks)?;
2115 }
2116
2117 let phase3_fn = device
2118 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
2119 .ok_or_else(|| {
2120 XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
2121 })?;
2122
2123 unsafe {
2125 phase3_fn.clone().launch(
2126 LaunchConfig {
2127 grid_dim: (num_blocks, 1, 1),
2128 block_dim: (block_size, 1, 1),
2129 shared_mem_bytes: 0,
2130 },
2131 (&mut *data, &block_sums, n),
2132 )
2133 }
2134 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase3 failed: {}", e)))?;
2135
2136 self.device.synchronize()?;
2137 Ok(())
2138 }
2139
2140 pub(crate) fn multiblock_scan_u32_inplace_on_stream(
2155 &self,
2156 data: &mut crate::memory::TrackedCudaSlice<u32>,
2157 n: u32,
2158 cu_stream: &cudarc::driver::CudaStream,
2159 launch_stream: crate::device_runtime::StreamId,
2160 runtime: &crate::device_runtime::XlogDeviceRuntime,
2161 ) -> Result<()> {
2162 if n == 0 {
2163 return Ok(());
2164 }
2165 let device = self.device.inner();
2166 let block_size = 256u32;
2167
2168 if n <= block_size {
2169 let phase2_fn = device
2170 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE2)
2171 .ok_or_else(|| {
2172 XlogError::Kernel("Failed to get multiblock_scan_phase2 kernel".to_string())
2173 })?;
2174 unsafe {
2176 phase2_fn.clone().launch_on_stream(
2177 cu_stream,
2178 LaunchConfig {
2179 grid_dim: (1, 1, 1),
2180 block_dim: (block_size, 1, 1),
2181 shared_mem_bytes: 0,
2182 },
2183 (&mut *data, n),
2184 )
2185 }
2186 .map_err(|e| {
2187 XlogError::Kernel(format!("multiblock_scan_phase2 (on_stream) failed: {}", e))
2188 })?;
2189 return Ok(());
2190 }
2191
2192 let num_blocks = n.div_ceil(block_size);
2193 let mut block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
2194 runtime
2201 .prepare_first_use(
2202 &block_sums,
2203 launch_stream,
2204 crate::device_runtime::Access::Write,
2205 )
2206 .map_err(|e| {
2207 XlogError::Kernel(format!(
2208 "multiblock_scan_u32_inplace_on_stream: prepare block_sums failed: {}",
2209 e
2210 ))
2211 })?;
2212
2213 let phase1_u32_fn = device
2214 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_U32_PHASE1)
2215 .ok_or_else(|| {
2216 XlogError::Kernel("Failed to get multiblock_scan_u32_phase1 kernel".to_string())
2217 })?;
2218 unsafe {
2220 phase1_u32_fn.clone().launch_on_stream(
2221 cu_stream,
2222 LaunchConfig {
2223 grid_dim: (num_blocks, 1, 1),
2224 block_dim: (block_size, 1, 1),
2225 shared_mem_bytes: 0,
2226 },
2227 (&mut *data, &mut block_sums, n),
2228 )
2229 }
2230 .map_err(|e| {
2231 XlogError::Kernel(format!(
2232 "multiblock_scan_u32_phase1 (on_stream) failed: {}",
2233 e
2234 ))
2235 })?;
2236
2237 if num_blocks > 1 {
2238 self.multiblock_scan_u32_inplace_on_stream(
2239 &mut block_sums,
2240 num_blocks,
2241 cu_stream,
2242 launch_stream,
2243 runtime,
2244 )?;
2245 }
2246
2247 let phase3_fn = device
2248 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
2249 .ok_or_else(|| {
2250 XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
2251 })?;
2252 unsafe {
2254 phase3_fn.clone().launch_on_stream(
2255 cu_stream,
2256 LaunchConfig {
2257 grid_dim: (num_blocks, 1, 1),
2258 block_dim: (block_size, 1, 1),
2259 shared_mem_bytes: 0,
2260 },
2261 (&mut *data, &block_sums, n),
2262 )
2263 }
2264 .map_err(|e| {
2265 XlogError::Kernel(format!("multiblock_scan_phase3 (on_stream) failed: {}", e))
2266 })?;
2267
2268 if let Some(b) = block_sums.runtime_block() {
2274 runtime
2275 .finish_block_use(
2276 crate::device_runtime::BlockId::from_block(b),
2277 launch_stream,
2278 crate::device_runtime::Access::Write,
2279 )
2280 .map_err(|e| {
2281 XlogError::Kernel(format!(
2282 "multiblock_scan_u32_inplace_on_stream: finish_block_use \
2283 for intermediate block_sums failed: {}",
2284 e
2285 ))
2286 })?;
2287 } else {
2288 return Err(XlogError::Kernel(
2289 "multiblock_scan_u32_inplace_on_stream: intermediate block_sums has no \
2290 runtime block — caller must use a runtime-backed manager"
2291 .to_string(),
2292 ));
2293 }
2294 Ok(())
2295 }
2296
2297 pub(crate) fn multiblock_scan_u32_scratch_for_len(
2300 &self,
2301 mut n: u32,
2302 ) -> Result<MultiblockScanScratchU32> {
2303 let block_size = 256u32;
2304 let mut levels = Vec::new();
2305 while n > block_size {
2306 let num_blocks = n.div_ceil(block_size);
2307 levels.push(self.memory.alloc::<u32>(num_blocks as usize)?);
2308 n = num_blocks;
2309 }
2310 Ok(MultiblockScanScratchU32 { levels })
2311 }
2312
2313 pub(crate) fn multiblock_scan_u32_inplace_on_stream_with_scratch(
2320 &self,
2321 data: &mut crate::memory::TrackedCudaSlice<u32>,
2322 n: u32,
2323 cu_stream: &cudarc::driver::CudaStream,
2324 scratch: &mut MultiblockScanScratchU32,
2325 ) -> Result<()> {
2326 self.multiblock_scan_u32_inplace_on_stream_with_scratch_levels(
2327 data,
2328 n,
2329 cu_stream,
2330 &mut scratch.levels,
2331 )
2332 }
2333
2334 fn multiblock_scan_u32_inplace_on_stream_with_scratch_levels(
2335 &self,
2336 data: &mut crate::memory::TrackedCudaSlice<u32>,
2337 n: u32,
2338 cu_stream: &cudarc::driver::CudaStream,
2339 scratch_levels: &mut [TrackedCudaSlice<u32>],
2340 ) -> Result<()> {
2341 if n == 0 {
2342 return Ok(());
2343 }
2344 let device = self.device.inner();
2345 let block_size = 256u32;
2346
2347 if n <= block_size {
2348 let phase2_fn = device
2349 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE2)
2350 .ok_or_else(|| {
2351 XlogError::Kernel("Failed to get multiblock_scan_phase2 kernel".to_string())
2352 })?;
2353 unsafe {
2355 phase2_fn.clone().launch_on_stream(
2356 cu_stream,
2357 LaunchConfig {
2358 grid_dim: (1, 1, 1),
2359 block_dim: (block_size, 1, 1),
2360 shared_mem_bytes: 0,
2361 },
2362 (&mut *data, n),
2363 )
2364 }
2365 .map_err(|e| {
2366 XlogError::Kernel(format!(
2367 "multiblock_scan_phase2 (graph scratch) failed: {}",
2368 e
2369 ))
2370 })?;
2371 return Ok(());
2372 }
2373
2374 let num_blocks = n.div_ceil(block_size);
2375 let (block_sums, rest) = scratch_levels.split_first_mut().ok_or_else(|| {
2376 XlogError::Kernel(format!(
2377 "multiblock_scan_u32_inplace_on_stream_with_scratch: missing scratch level \
2378 for n={n}, num_blocks={num_blocks}"
2379 ))
2380 })?;
2381 if block_sums.len() < num_blocks as usize {
2382 return Err(XlogError::Kernel(format!(
2383 "multiblock_scan_u32_inplace_on_stream_with_scratch: scratch level too small \
2384 (have {}, need {})",
2385 block_sums.len(),
2386 num_blocks
2387 )));
2388 }
2389
2390 let phase1_u32_fn = device
2391 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_U32_PHASE1)
2392 .ok_or_else(|| {
2393 XlogError::Kernel("Failed to get multiblock_scan_u32_phase1 kernel".to_string())
2394 })?;
2395 unsafe {
2397 phase1_u32_fn.clone().launch_on_stream(
2398 cu_stream,
2399 LaunchConfig {
2400 grid_dim: (num_blocks, 1, 1),
2401 block_dim: (block_size, 1, 1),
2402 shared_mem_bytes: 0,
2403 },
2404 (&mut *data, &mut *block_sums, n),
2405 )
2406 }
2407 .map_err(|e| {
2408 XlogError::Kernel(format!(
2409 "multiblock_scan_u32_phase1 (graph scratch) failed: {}",
2410 e
2411 ))
2412 })?;
2413
2414 if num_blocks > 1 {
2415 self.multiblock_scan_u32_inplace_on_stream_with_scratch_levels(
2416 block_sums, num_blocks, cu_stream, rest,
2417 )?;
2418 }
2419
2420 let phase3_fn = device
2421 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
2422 .ok_or_else(|| {
2423 XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
2424 })?;
2425 unsafe {
2427 phase3_fn.clone().launch_on_stream(
2428 cu_stream,
2429 LaunchConfig {
2430 grid_dim: (num_blocks, 1, 1),
2431 block_dim: (block_size, 1, 1),
2432 shared_mem_bytes: 0,
2433 },
2434 (&mut *data, &*block_sums, n),
2435 )
2436 }
2437 .map_err(|e| {
2438 XlogError::Kernel(format!(
2439 "multiblock_scan_phase3 (graph scratch) failed: {}",
2440 e
2441 ))
2442 })?;
2443 Ok(())
2444 }
2445
2446 pub(crate) fn multiblock_scan_u32_view_inplace_on_stream(
2454 &self,
2455 data: &mut CudaViewMut<'_, u32>,
2456 n: u32,
2457 cu_stream: &cudarc::driver::CudaStream,
2458 launch_stream: crate::device_runtime::StreamId,
2459 runtime: &crate::device_runtime::XlogDeviceRuntime,
2460 ) -> Result<()> {
2461 if n == 0 {
2462 return Ok(());
2463 }
2464 let device = self.device.inner();
2465 let block_size = 256u32;
2466
2467 if n <= block_size {
2468 let phase2_fn = device
2469 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE2)
2470 .ok_or_else(|| {
2471 XlogError::Kernel("Failed to get multiblock_scan_phase2 kernel".to_string())
2472 })?;
2473 unsafe {
2475 phase2_fn.clone().launch_on_stream(
2476 cu_stream,
2477 LaunchConfig {
2478 grid_dim: (1, 1, 1),
2479 block_dim: (block_size, 1, 1),
2480 shared_mem_bytes: 0,
2481 },
2482 (data, n),
2483 )
2484 }
2485 .map_err(|e| {
2486 XlogError::Kernel(format!(
2487 "multiblock_scan_phase2 (view on_stream) failed: {}",
2488 e
2489 ))
2490 })?;
2491 return Ok(());
2492 }
2493
2494 let num_blocks = n.div_ceil(block_size);
2495 let mut block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
2496 runtime
2500 .prepare_first_use(
2501 &block_sums,
2502 launch_stream,
2503 crate::device_runtime::Access::Write,
2504 )
2505 .map_err(|e| {
2506 XlogError::Kernel(format!(
2507 "multiblock_scan_u32_view_inplace_on_stream: prepare block_sums failed: {}",
2508 e
2509 ))
2510 })?;
2511
2512 let phase1_u32_fn = device
2513 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_U32_PHASE1)
2514 .ok_or_else(|| {
2515 XlogError::Kernel("Failed to get multiblock_scan_u32_phase1 kernel".to_string())
2516 })?;
2517 unsafe {
2519 phase1_u32_fn.clone().launch_on_stream(
2520 cu_stream,
2521 LaunchConfig {
2522 grid_dim: (num_blocks, 1, 1),
2523 block_dim: (block_size, 1, 1),
2524 shared_mem_bytes: 0,
2525 },
2526 (&mut *data, &mut block_sums, n),
2527 )
2528 }
2529 .map_err(|e| {
2530 XlogError::Kernel(format!(
2531 "multiblock_scan_u32_phase1 (view on_stream) failed: {}",
2532 e
2533 ))
2534 })?;
2535
2536 if num_blocks > 1 {
2537 self.multiblock_scan_u32_inplace_on_stream(
2538 &mut block_sums,
2539 num_blocks,
2540 cu_stream,
2541 launch_stream,
2542 runtime,
2543 )?;
2544 }
2545
2546 let phase3_fn = device
2547 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
2548 .ok_or_else(|| {
2549 XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
2550 })?;
2551 unsafe {
2553 phase3_fn.clone().launch_on_stream(
2554 cu_stream,
2555 LaunchConfig {
2556 grid_dim: (num_blocks, 1, 1),
2557 block_dim: (block_size, 1, 1),
2558 shared_mem_bytes: 0,
2559 },
2560 (&mut *data, &block_sums, n),
2561 )
2562 }
2563 .map_err(|e| {
2564 XlogError::Kernel(format!(
2565 "multiblock_scan_phase3 (view on_stream) failed: {}",
2566 e
2567 ))
2568 })?;
2569
2570 if let Some(b) = block_sums.runtime_block() {
2572 runtime
2573 .finish_block_use(
2574 crate::device_runtime::BlockId::from_block(b),
2575 launch_stream,
2576 crate::device_runtime::Access::Write,
2577 )
2578 .map_err(|e| {
2579 XlogError::Kernel(format!(
2580 "multiblock_scan_u32_view_inplace_on_stream: finish_block_use \
2581 for intermediate block_sums failed: {}",
2582 e
2583 ))
2584 })?;
2585 } else {
2586 return Err(XlogError::Kernel(
2587 "multiblock_scan_u32_view_inplace_on_stream: intermediate block_sums has no \
2588 runtime block — caller must use a runtime-backed manager"
2589 .to_string(),
2590 ));
2591 }
2592 Ok(())
2593 }
2594
2595 fn multiblock_scan_u32_view_inplace(
2596 &self,
2597 data: &mut CudaViewMut<'_, u32>,
2598 n: u32,
2599 ) -> Result<()> {
2600 if n == 0 {
2601 return Ok(());
2602 }
2603
2604 let device = self.device.inner();
2605 let block_size = 256u32;
2606
2607 if n <= block_size {
2608 let phase2_fn = device
2609 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE2)
2610 .ok_or_else(|| {
2611 XlogError::Kernel("Failed to get multiblock_scan_phase2 kernel".to_string())
2612 })?;
2613
2614 unsafe {
2616 phase2_fn.clone().launch(
2617 LaunchConfig {
2618 grid_dim: (1, 1, 1),
2619 block_dim: (block_size, 1, 1),
2620 shared_mem_bytes: 0,
2621 },
2622 (data, n),
2623 )
2624 }
2625 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase2 failed: {}", e)))?;
2626
2627 self.device.synchronize()?;
2628 return Ok(());
2629 }
2630
2631 let num_blocks = n.div_ceil(block_size);
2632 let mut block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
2633
2634 let phase1_u32_fn = device
2635 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_U32_PHASE1)
2636 .ok_or_else(|| {
2637 XlogError::Kernel("Failed to get multiblock_scan_u32_phase1 kernel".to_string())
2638 })?;
2639
2640 unsafe {
2642 phase1_u32_fn.clone().launch(
2643 LaunchConfig {
2644 grid_dim: (num_blocks, 1, 1),
2645 block_dim: (block_size, 1, 1),
2646 shared_mem_bytes: 0,
2647 },
2648 (&mut *data, &mut block_sums, n),
2649 )
2650 }
2651 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_u32_phase1 failed: {}", e)))?;
2652 self.device.synchronize()?;
2653
2654 if num_blocks > 1 {
2655 self.multiblock_scan_u32_inplace(&mut block_sums, num_blocks)?;
2656 }
2657
2658 let phase3_fn = device
2659 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
2660 .ok_or_else(|| {
2661 XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
2662 })?;
2663
2664 unsafe {
2666 phase3_fn.clone().launch(
2667 LaunchConfig {
2668 grid_dim: (num_blocks, 1, 1),
2669 block_dim: (block_size, 1, 1),
2670 shared_mem_bytes: 0,
2671 },
2672 (&mut *data, &block_sums, n),
2673 )
2674 }
2675 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase3 failed: {}", e)))?;
2676
2677 self.device.synchronize()?;
2678 Ok(())
2679 }
2680
2681 pub fn device_row_count(&self, buffer: &CudaBuffer) -> Result<usize> {
2686 if let Some(n) = buffer.cached_row_count() {
2687 return Ok(n as usize);
2688 }
2689 let mut host_rows = [0u32];
2690 self.device
2691 .inner()
2692 .dtoh_sync_copy_into(buffer.num_rows_device(), &mut host_rows)
2693 .map_err(|e| XlogError::Kernel(format!("Failed to read row count: {}", e)))?;
2694 buffer.set_cached_row_count_if_unset(host_rows[0]);
2695 Ok(host_rows[0] as usize)
2696 }
2697
2698 pub fn validated_logical_row_count(&self, buffer: &CudaBuffer) -> Result<usize> {
2703 let logical_rows = self.device_row_count(buffer)?;
2704 validate_logical_row_count(buffer.num_rows(), logical_rows)
2705 }
2706
2707 fn clone_device_row_count(&self, buffer: &CudaBuffer) -> Result<TrackedCudaSlice<u32>> {
2708 let mut d_num_rows = self.memory.alloc::<u32>(1)?;
2709 self.device
2710 .inner()
2711 .dtod_copy(buffer.num_rows_device(), &mut d_num_rows)
2712 .map_err(|e| XlogError::Kernel(format!("Failed to copy row count: {}", e)))?;
2713 Ok(d_num_rows)
2714 }
2715
2716 fn upload_device_row_count(&self, row_count: u32) -> Result<TrackedCudaSlice<u32>> {
2717 let mut d_num_rows = self.memory.alloc::<u32>(1)?;
2718 self.htod_launch_metadata_sync_copy_into(&[row_count], &mut d_num_rows)
2719 .map_err(|e| XlogError::Kernel(format!("Failed to upload row count: {}", e)))?;
2720 Ok(d_num_rows)
2721 }
2722
2723 fn buffer_from_columns_with_device_count(
2724 &self,
2725 columns: Vec<CudaColumn>,
2726 row_cap: u64,
2727 schema: Schema,
2728 src: &CudaBuffer,
2729 ) -> Result<CudaBuffer> {
2730 let d_num_rows = self.clone_device_row_count(src)?;
2731 Ok(CudaBuffer::from_columns(
2732 columns, row_cap, d_num_rows, schema,
2733 ))
2734 }
2735
2736 fn column_bytes_view<'a>(
2737 &self,
2738 col: &'a CudaColumn,
2739 num_bytes: usize,
2740 ) -> Result<RawCudaView<'a, u8>> {
2741 if col.num_bytes() < num_bytes {
2742 return Err(XlogError::Kernel(format!(
2743 "Column has {} bytes but {} required",
2744 col.num_bytes(),
2745 num_bytes
2746 )));
2747 }
2748 let ptr = *col.device_ptr();
2749 Ok(RawCudaView {
2750 ptr,
2751 len: num_bytes,
2752 stream: col.stream().clone(),
2753 source_block: col.runtime_block(),
2754 _marker: PhantomData,
2755 })
2756 }
2757
2758 fn bytes_as_u32_view<'a>(
2759 &self,
2760 bytes: &'a TrackedCudaSlice<u8>,
2761 num_elements: usize,
2762 ) -> Result<RawCudaView<'a, u32>> {
2763 let required_bytes = num_elements * std::mem::size_of::<u32>();
2764 if bytes.len() < required_bytes {
2765 return Err(XlogError::Kernel(format!(
2766 "Packed keys have {} bytes but {} required for {} u32 elements",
2767 bytes.len(),
2768 required_bytes,
2769 num_elements
2770 )));
2771 }
2772 let ptr = *bytes.device_ptr();
2773 if !(ptr as usize).is_multiple_of(std::mem::align_of::<u32>()) {
2774 return Err(XlogError::Kernel(
2775 "Packed keys device pointer is not u32-aligned".to_string(),
2776 ));
2777 }
2778 Ok(RawCudaView {
2779 ptr,
2780 len: num_elements,
2781 stream: bytes.stream().clone(),
2782 source_block: bytes.runtime_block(),
2783 _marker: PhantomData,
2784 })
2785 }
2786
2787 fn column_as_u32_view<'a>(
2789 &self,
2790 col: &'a CudaColumn,
2791 num_elements: usize,
2792 ) -> Result<RawCudaView<'a, u32>> {
2793 let required_bytes = num_elements * std::mem::size_of::<u32>();
2794 if col.num_bytes() < required_bytes {
2795 return Err(XlogError::Kernel(format!(
2796 "Column has {} bytes but {} required for {} u32 elements",
2797 col.num_bytes(),
2798 required_bytes,
2799 num_elements
2800 )));
2801 }
2802 let ptr = *col.device_ptr();
2803 if !(ptr as usize).is_multiple_of(std::mem::align_of::<u32>()) {
2804 return Err(XlogError::Kernel(
2805 "Column device pointer is not u32-aligned".to_string(),
2806 ));
2807 }
2808 Ok(RawCudaView {
2809 ptr,
2810 len: num_elements,
2811 stream: col.stream().clone(),
2812 source_block: col.runtime_block(),
2813 _marker: PhantomData,
2814 })
2815 }
2816
2817 fn column_as_u64_view<'a>(
2818 &self,
2819 col: &'a CudaColumn,
2820 num_elements: usize,
2821 ) -> Result<RawCudaView<'a, u64>> {
2822 let required_bytes = num_elements * std::mem::size_of::<u64>();
2823 if col.num_bytes() < required_bytes {
2824 return Err(XlogError::Kernel(format!(
2825 "Column has {} bytes but {} required for {} u64 elements",
2826 col.num_bytes(),
2827 required_bytes,
2828 num_elements
2829 )));
2830 }
2831 let ptr = *col.device_ptr();
2832 if !(ptr as usize).is_multiple_of(std::mem::align_of::<u64>()) {
2833 return Err(XlogError::Kernel(
2834 "Column device pointer is not u64-aligned".to_string(),
2835 ));
2836 }
2837 Ok(RawCudaView {
2838 ptr,
2839 len: num_elements,
2840 stream: col.stream().clone(),
2841 source_block: col.runtime_block(),
2842 _marker: PhantomData,
2843 })
2844 }
2845
2846 fn column_as_f64_view<'a>(
2848 &self,
2849 col: &'a CudaColumn,
2850 num_elements: usize,
2851 ) -> Result<RawCudaView<'a, f64>> {
2852 let required_bytes = num_elements * std::mem::size_of::<f64>();
2853 if col.num_bytes() < required_bytes {
2854 return Err(XlogError::Kernel(format!(
2855 "Column has {} bytes but {} required for {} f64 elements",
2856 col.num_bytes(),
2857 required_bytes,
2858 num_elements
2859 )));
2860 }
2861 let ptr = *col.device_ptr();
2862 if !(ptr as usize).is_multiple_of(std::mem::align_of::<f64>()) {
2863 return Err(XlogError::Kernel(
2864 "Column device pointer is not f64-aligned".to_string(),
2865 ));
2866 }
2867 Ok(RawCudaView {
2868 ptr,
2869 len: num_elements,
2870 stream: col.stream().clone(),
2871 source_block: col.runtime_block(),
2872 _marker: PhantomData,
2873 })
2874 }
2875
2876 pub fn create_empty_buffer(&self, schema: Schema) -> Result<CudaBuffer> {
2887 let mut columns = Vec::with_capacity(schema.arity());
2888 for _ in 0..schema.arity() {
2889 columns.push(self.memory.alloc::<u8>(0)?.into());
2891 }
2892 self.buffer_from_columns(columns, 0, schema)
2893 }
2894
2895 pub fn create_zero_arity_buffer(&self, schema: Schema, rows: u32) -> Result<CudaBuffer> {
2903 debug_assert_eq!(
2904 schema.arity(),
2905 0,
2906 "create_zero_arity_buffer requires arity 0"
2907 );
2908 self.buffer_from_columns(Vec::new(), u64::from(rows), schema)
2909 }
2910
2911 pub(crate) fn buffer_from_columns(
2912 &self,
2913 columns: Vec<CudaColumn>,
2914 row_cap: u64,
2915 schema: Schema,
2916 ) -> Result<CudaBuffer> {
2917 let row_u32 = u32::try_from(row_cap)
2918 .map_err(|_| XlogError::Kernel(format!("Row capacity {} exceeds u32::MAX", row_cap)))?;
2919 let mut d_num_rows = self.memory.alloc::<u32>(1)?;
2920 self.htod_launch_metadata_sync_copy_into(&[row_u32], &mut d_num_rows)
2921 .map_err(|e| XlogError::Kernel(format!("Failed to set row count: {}", e)))?;
2922 Ok(CudaBuffer::from_columns_with_host_count(
2923 columns, row_cap, d_num_rows, schema, row_u32,
2924 ))
2925 }
2926
2927 fn combine_schemas(&self, left: &Schema, right: &Schema) -> Schema {
2929 let mut columns = left.columns.clone();
2930 columns.extend(right.columns.iter().cloned());
2931 let mut sort_labels = left.sort_labels().to_vec();
2932 sort_labels.extend(right.sort_labels().iter().cloned());
2933 Schema::new(columns)
2934 .with_sort_labels(sort_labels)
2935 .expect("combined schema sort labels match column arity")
2936 }
2937
2938 fn schemas_type_compatible(&self, a: &Schema, b: &Schema) -> bool {
2943 if a.arity() != b.arity() {
2944 return false;
2945 }
2946 for i in 0..a.arity() {
2947 if a.column_type(i) != b.column_type(i) {
2948 return false;
2949 }
2950 }
2951 true
2952 }
2953}
2954
2955#[cfg(test)]
2956mod tests {
2957 use super::*;
2958 use crate::device_runtime::{
2959 AsyncCudaResource, DeviceMemoryResource, GlobalDeviceBudget, LoggingResource, NullSink,
2960 StreamPool, XlogDeviceRuntime,
2961 };
2962 use xlog_core::{AggOp, MemoryBudget, ScalarType};
2963
2964 fn has_cuda_device() -> bool {
2965 CudaDevice::new(0).is_ok()
2966 }
2967
2968 #[test]
2969 fn test_kernel_artifact_locator_precedence_order() {
2970 use super::kernel_paths::KernelArtifactLocator;
2971 use std::fs;
2972 use std::path::PathBuf;
2973
2974 let root = std::env::temp_dir().join(format!(
2975 "xlog-kernel-paths-{}-{}",
2976 std::process::id(),
2977 std::time::SystemTime::now()
2978 .duration_since(std::time::UNIX_EPOCH)
2979 .expect("system clock before UNIX_EPOCH")
2980 .as_nanos()
2981 ));
2982 let cubin_dir = root.join("cubin");
2983 let package_dir = root.join("bin").join("kernels");
2984 let out_dir = root.join("out");
2985 fs::create_dir_all(&cubin_dir).expect("create cubin dir");
2986 fs::create_dir_all(&package_dir).expect("create package kernels dir");
2987 fs::create_dir_all(&out_dir).expect("create out dir");
2988
2989 let name = "xlog_join";
2990 let cc = 75;
2991 let cubin_path = cubin_dir.join(format!("{name}.sm_{cc}.cubin"));
2992 let package_path = package_dir.join(format!("{name}.sm_{cc}.cubin"));
2993 let out_path = out_dir.join(format!("{name}.sm_{cc}.cubin"));
2994 fs::write(&cubin_path, b"cubin").expect("write cubin file");
2995 fs::write(&package_path, b"package").expect("write package file");
2996 fs::write(&out_path, b"out").expect("write out file");
2997
2998 let locator = KernelArtifactLocator::new(
2999 Some(cubin_dir.clone()),
3000 Some(package_dir.clone()),
3001 Some(out_dir.clone()),
3002 );
3003
3004 let (path, is_cubin) = locator
3005 .resolve_module_path(name, cc)
3006 .expect("expected a kernel artifact");
3007 assert_eq!(path, cubin_path);
3008 assert!(is_cubin);
3009
3010 fs::remove_file(&cubin_path).expect("remove cubin file");
3011 let (path, is_cubin) = locator
3012 .resolve_module_path(name, cc)
3013 .expect("expected package kernel artifact");
3014 assert_eq!(path, package_path);
3015 assert!(is_cubin);
3016
3017 fs::remove_file(&package_path).expect("remove package file");
3018 let (path, is_cubin) = locator
3019 .resolve_module_path(name, cc)
3020 .expect("expected out dir kernel artifact");
3021 assert_eq!(path, out_path);
3022 assert!(is_cubin);
3023
3024 let _ = fs::remove_dir_all(PathBuf::from(&root));
3025 }
3026
3027 #[test]
3028 fn test_module_resolution_finds_portable_ptx() {
3029 for name in crate::kernel_manifest_data::KERNEL_CU_NAMES {
3032 let result = resolve_module_path(name, 999);
3033 assert!(
3034 result.is_some(),
3035 "resolve_module_path({name}, 999) should find portable PTX"
3036 );
3037 let (path, is_cubin) = result.unwrap();
3038 assert!(
3039 !is_cubin,
3040 "{name}: expected portable PTX fallback, got cubin"
3041 );
3042 assert!(
3043 path.to_str().unwrap().ends_with(".portable.ptx"),
3044 "{name}: path should end with .portable.ptx, got {:?}",
3045 path
3046 );
3047 }
3048 }
3049
3050 #[test]
3051 fn test_module_resolution_falls_back_to_embedded_portable_ptx() {
3052 use super::kernel_paths::KernelArtifactLocator;
3053
3054 let locator = KernelArtifactLocator::new(None, None, None);
3055 for name in crate::kernel_manifest_data::KERNEL_CU_NAMES {
3056 let sources = resolve_module_sources_with_locator(name, 999, &locator);
3057 assert_eq!(
3058 sources.len(),
3059 1,
3060 "{name}: expected only embedded portable PTX fallback"
3061 );
3062
3063 match &sources[0] {
3064 KernelModuleSource::EmbeddedPortablePtx { ptx } => {
3065 assert!(
3066 ptx.contains(".entry"),
3067 "{name}: embedded PTX should contain CUDA entry points"
3068 );
3069 }
3070 KernelModuleSource::File { path, .. } => {
3071 panic!(
3072 "{name}: expected embedded portable PTX fallback, got file {}",
3073 path.display()
3074 );
3075 }
3076 }
3077 }
3078 }
3079
3080 #[test]
3081 fn test_embedded_portable_ptx_manifest_matches_kernel_manifest() {
3082 let embedded_names: std::collections::BTreeSet<_> =
3083 crate::embedded_kernel_data::EMBEDDED_PORTABLE_PTX
3084 .iter()
3085 .map(|artifact| artifact.name)
3086 .collect();
3087 let manifest_names: std::collections::BTreeSet<_> =
3088 crate::kernel_manifest_data::KERNEL_CU_NAMES
3089 .iter()
3090 .copied()
3091 .collect();
3092
3093 assert_eq!(
3094 embedded_names, manifest_names,
3095 "embedded portable PTX table should cover every runtime kernel module"
3096 );
3097 }
3098
3099 #[test]
3100 fn test_kernel_provider_creation() {
3101 if !has_cuda_device() {
3102 eprintln!("Skipping test: no CUDA device available");
3103 return;
3104 }
3105
3106 let device = Arc::new(CudaDevice::new(0).expect("Failed to create device"));
3107 let budget = MemoryBudget::with_limit(1024 * 1024 * 1024); let memory = Arc::new(GpuMemoryManager::new(device.clone(), budget));
3109
3110 let provider = CudaKernelProvider::new(device.clone(), memory.clone());
3111 assert!(
3112 provider.is_ok(),
3113 "Failed to create kernel provider: {:?}",
3114 provider.err()
3115 );
3116
3117 let provider = provider.unwrap();
3118 assert!(Arc::ptr_eq(provider.device(), &device));
3119 assert!(Arc::ptr_eq(provider.memory(), &memory));
3120 }
3121
3122 #[test]
3123 fn test_kernel_functions_accessible() {
3124 if !has_cuda_device() {
3125 eprintln!("Skipping test: no CUDA device available");
3126 return;
3127 }
3128
3129 let device = Arc::new(CudaDevice::new(0).expect("Failed to create device"));
3130 let budget = MemoryBudget::with_limit(1024 * 1024 * 1024);
3131 let memory = Arc::new(GpuMemoryManager::new(device.clone(), budget));
3132
3133 let _provider =
3134 CudaKernelProvider::new(device.clone(), memory).expect("Failed to create provider");
3135
3136 let inner = device.inner();
3138
3139 let build_fn = inner.get_func(JOIN_MODULE, join_kernels::HASH_JOIN_BUILD);
3141 assert!(
3142 build_fn.is_some(),
3143 "hash_join_build function should be accessible"
3144 );
3145
3146 let probe_fn = inner.get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE);
3147 assert!(
3148 probe_fn.is_some(),
3149 "hash_join_probe function should be accessible"
3150 );
3151
3152 let mark_fn = inner.get_func(DEDUP_MODULE, dedup_kernels::MARK_DUPLICATES);
3154 assert!(
3155 mark_fn.is_some(),
3156 "mark_duplicates function should be accessible"
3157 );
3158
3159 let compact_fn = inner.get_func(DEDUP_MODULE, dedup_kernels::COMPACT_ROWS);
3160 assert!(
3161 compact_fn.is_some(),
3162 "compact_rows function should be accessible"
3163 );
3164
3165 let boundaries_fn =
3167 inner.get_func(GROUPBY_MODULE, groupby_kernels::DETECT_GROUP_BOUNDARIES);
3168 assert!(
3169 boundaries_fn.is_some(),
3170 "detect_group_boundaries function should be accessible"
3171 );
3172
3173 let count_fn = inner.get_func(GROUPBY_MODULE, groupby_kernels::GROUPBY_COUNT);
3174 assert!(
3175 count_fn.is_some(),
3176 "groupby_count function should be accessible"
3177 );
3178
3179 let sum_fn = inner.get_func(GROUPBY_MODULE, groupby_kernels::GROUPBY_SUM);
3180 assert!(
3181 sum_fn.is_some(),
3182 "groupby_sum function should be accessible"
3183 );
3184
3185 let min_fn = inner.get_func(GROUPBY_MODULE, groupby_kernels::GROUPBY_MIN);
3186 assert!(
3187 min_fn.is_some(),
3188 "groupby_min function should be accessible"
3189 );
3190
3191 let max_fn = inner.get_func(GROUPBY_MODULE, groupby_kernels::GROUPBY_MAX);
3192 assert!(
3193 max_fn.is_some(),
3194 "groupby_max function should be accessible"
3195 );
3196
3197 let xgcf_forward = inner.get_func(CIRCUIT_MODULE, "xgcf_forward_level");
3199 assert!(
3200 xgcf_forward.is_some(),
3201 "xgcf_forward_level function should be accessible"
3202 );
3203
3204 let xgcf_backward_propagate =
3205 inner.get_func(CIRCUIT_MODULE, "xgcf_backward_level_propagate");
3206 assert!(
3207 xgcf_backward_propagate.is_some(),
3208 "xgcf_backward_level_propagate function should be accessible"
3209 );
3210
3211 let xgcf_backward_decision_grad =
3212 inner.get_func(CIRCUIT_MODULE, "xgcf_backward_level_decision_grad");
3213 assert!(
3214 xgcf_backward_decision_grad.is_some(),
3215 "xgcf_backward_level_decision_grad function should be accessible"
3216 );
3217
3218 let xgcf_backward_lit_grad = inner.get_func(CIRCUIT_MODULE, "xgcf_backward_level_lit_grad");
3219 assert!(
3220 xgcf_backward_lit_grad.is_some(),
3221 "xgcf_backward_level_lit_grad function should be accessible"
3222 );
3223
3224 let neural_fill = inner.get_func("xlog_neural", "neural_fill_ad_chain_f32");
3226 assert!(
3227 neural_fill.is_some(),
3228 "neural_fill_ad_chain_f32 function should be accessible"
3229 );
3230 let neural_scatter = inner.get_func("xlog_neural", "neural_scatter_ad_chain_grads_f32");
3231 assert!(
3232 neural_scatter.is_some(),
3233 "neural_scatter_ad_chain_grads_f32 function should be accessible"
3234 );
3235 }
3236
3237 #[test]
3238 fn test_module_names_unique() {
3239 assert_ne!(JOIN_MODULE, DEDUP_MODULE);
3241 assert_ne!(JOIN_MODULE, GROUPBY_MODULE);
3242 assert_ne!(DEDUP_MODULE, GROUPBY_MODULE);
3243 }
3244
3245 fn create_test_provider() -> Option<CudaKernelProvider> {
3247 if !has_cuda_device() {
3248 return None;
3249 }
3250 let device = Arc::new(CudaDevice::new(0).ok()?);
3251 let budget = MemoryBudget::with_limit(1024 * 1024 * 1024);
3252 let memory = Arc::new(GpuMemoryManager::new(device.clone(), budget));
3253 CudaKernelProvider::new(device, memory).ok()
3254 }
3255
3256 fn create_test_provider_with_runtime() -> Option<(CudaKernelProvider, Arc<XlogDeviceRuntime>)> {
3257 if !has_cuda_device() {
3258 return None;
3259 }
3260 let device = Arc::new(CudaDevice::new(0).ok()?);
3261 let pool = Arc::new(StreamPool::with_defaults(Arc::clone(&device)));
3262 let sink = Arc::new(NullSink::new());
3263 let async_resource: Box<dyn DeviceMemoryResource + Send + Sync> = Box::new(
3264 AsyncCudaResource::new(Arc::clone(&device), 0, Arc::clone(&pool)),
3265 );
3266 let logging: Box<dyn DeviceMemoryResource + Send + Sync> =
3267 Box::new(LoggingResource::new(async_resource, sink));
3268 let budget: Box<dyn DeviceMemoryResource + Send + Sync> =
3269 Box::new(GlobalDeviceBudget::new(logging, 1024 * 1024 * 1024));
3270 let runtime = Arc::new(XlogDeviceRuntime::with_resource(
3271 Arc::clone(&device),
3272 0,
3273 pool,
3274 budget,
3275 ));
3276 let memory = Arc::new(GpuMemoryManager::with_runtime(
3277 Arc::clone(&device),
3278 MemoryBudget::with_limit(1024 * 1024 * 1024),
3279 Arc::clone(&runtime),
3280 ));
3281 let provider = CudaKernelProvider::with_runtime(device, memory).ok()?;
3282 Some((provider, runtime))
3283 }
3284
3285 #[test]
3286 fn test_recorded_join_index_build_runs_on_runtime_stream() {
3287 let (provider, runtime) = match create_test_provider_with_runtime() {
3288 Some(fixture) => fixture,
3289 None => {
3290 eprintln!("Skipping test: no CUDA device available");
3291 return;
3292 }
3293 };
3294 let stream = runtime.stream_pool().acquire().expect("recorded stream");
3295 let left = create_test_buffer(&provider, &[1, 2, 3, 4], "key");
3296 let right = create_test_buffer(&provider, &[1, 2, 3, 4], "key");
3297
3298 let index = provider
3299 .build_join_index_v2_recorded(&right, &[0], stream)
3300 .expect("recorded join-index build");
3301 let joined = provider
3302 .hash_join_v2_with_index_recorded(
3303 &left,
3304 &right,
3305 &[0],
3306 &[0],
3307 JoinType::Inner,
3308 &index,
3309 None,
3310 stream,
3311 )
3312 .expect("recorded indexed join consumes recorded build");
3313 runtime
3314 .stream_pool()
3315 .resolve(stream)
3316 .expect("stream resolves")
3317 .synchronize()
3318 .expect("recorded stream synchronized");
3319
3320 assert_eq!(index.right_num_rows(), 4);
3321 assert_eq!(index.right_keys(), &[0]);
3322 assert_eq!(provider.device_row_count(&joined).expect("joined rows"), 4);
3323 }
3324
3325 fn create_test_buffer(
3327 provider: &CudaKernelProvider,
3328 data: &[u32],
3329 col_name: &str,
3330 ) -> CudaBuffer {
3331 let schema = Schema::new(vec![(col_name.to_string(), ScalarType::U32)]);
3332 let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
3333
3334 let mut col = provider.memory().alloc::<u8>(bytes.len()).expect("alloc");
3335 provider
3336 .device()
3337 .inner()
3338 .htod_sync_copy_into(&bytes, &mut col)
3339 .expect("htod");
3340
3341 provider
3342 .buffer_from_columns(vec![col.into()], data.len() as u64, schema)
3343 .expect("buffer")
3344 }
3345
3346 fn create_empty_test_buffer(provider: &CudaKernelProvider, schema: Schema) -> CudaBuffer {
3348 let mut columns = Vec::with_capacity(schema.arity());
3349 for _ in 0..schema.arity() {
3350 columns.push(provider.memory().alloc::<u8>(0).expect("alloc").into());
3351 }
3352 provider
3353 .buffer_from_columns(columns, 0, schema)
3354 .expect("buffer")
3355 }
3356
3357 fn read_buffer_u32(provider: &CudaKernelProvider, buffer: &CudaBuffer, col: usize) -> Vec<u32> {
3359 if buffer.is_empty() || buffer.column(col).is_none() {
3360 return vec![];
3361 }
3362 let num_rows = buffer.num_rows() as usize;
3363 let mut bytes = vec![0u8; num_rows * 4];
3364 provider
3365 .device()
3366 .inner()
3367 .dtoh_sync_copy_into(buffer.column(col).unwrap(), &mut bytes)
3368 .expect("dtoh");
3369 bytes
3370 .chunks_exact(4)
3371 .map(|c| u32::from_le_bytes([c[0], c[1], c[2], c[3]]))
3372 .collect()
3373 }
3374
3375 #[test]
3376 fn test_compact_device_mask_respects_mask_len_smaller_than_row_cap() {
3377 let provider = match create_test_provider() {
3378 Some(p) => p,
3379 None => {
3380 eprintln!("Skipping test: no CUDA device available");
3381 return;
3382 }
3383 };
3384
3385 let schema = Schema::new(vec![("id".to_string(), ScalarType::U32)]);
3386 let base = create_test_buffer(&provider, &[1, 2, 3, 4, 5, 6, 7, 8], "id");
3387
3388 let row_cap = 16u64;
3389 let data: Vec<u32> = (0..row_cap as u32).collect();
3390 let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
3391 let mut col = provider.memory().alloc::<u8>(bytes.len()).expect("alloc");
3392 provider
3393 .device()
3394 .inner()
3395 .htod_sync_copy_into(&bytes, &mut col)
3396 .expect("htod");
3397 let expanded = provider
3398 .buffer_from_columns_with_device_count(vec![col.into()], row_cap, schema, &base)
3399 .expect("buffer");
3400
3401 let mask: Vec<u8> = vec![1, 0, 1, 0, 1, 0, 1, 0];
3402 let (prefix_sum, count) = provider.prefix_sum_mask(&mask).expect("prefix sum");
3403
3404 let mut d_mask = provider.memory().alloc::<u8>(mask.len()).expect("alloc");
3405 provider
3406 .device()
3407 .inner()
3408 .htod_sync_copy_into(&mask, &mut d_mask)
3409 .expect("mask htod");
3410
3411 let mut d_prefix = provider
3412 .memory()
3413 .alloc::<u32>(prefix_sum.len())
3414 .expect("alloc");
3415 provider
3416 .device()
3417 .inner()
3418 .htod_sync_copy_into(&prefix_sum, &mut d_prefix)
3419 .expect("prefix htod");
3420
3421 let mut d_out_count = provider.memory().alloc::<u32>(1).expect("alloc");
3422 provider
3423 .device()
3424 .inner()
3425 .htod_sync_copy_into(&[count], &mut d_out_count)
3426 .expect("count htod");
3427
3428 let compacted = provider
3429 .compact_buffer_by_device_mask_device_count(&expanded, &d_mask, &d_prefix, d_out_count)
3430 .expect("compact");
3431
3432 assert_eq!(compacted.num_rows(), mask.len() as u64);
3433 let device_rows = provider.device_row_count(&compacted).expect("row count");
3434 assert_eq!(device_rows as u32, count);
3435 }
3436
3437 #[test]
3438 fn test_clone_buffer_preserves_device_count() {
3439 let provider = match create_test_provider() {
3440 Some(p) => p,
3441 None => {
3442 eprintln!("Skipping test: no CUDA device available");
3443 return;
3444 }
3445 };
3446
3447 let schema = Schema::new(vec![("id".to_string(), ScalarType::U32)]);
3448 let ids: Vec<u32> = vec![10, 20, 30];
3449 let buffer = provider
3450 .create_buffer_from_slices(&[bytemuck::cast_slice(&ids)], schema)
3451 .unwrap();
3452
3453 let cloned = provider.clone_buffer(&buffer).unwrap();
3454
3455 let mut host_count = [0u32];
3456 provider
3457 .device()
3458 .inner()
3459 .dtoh_sync_copy_into(cloned.num_rows_device(), &mut host_count)
3460 .unwrap();
3461 assert_eq!(host_count[0], 3);
3462 }
3463
3464 #[test]
3473 fn test_clone_buffer_preserves_cached_row_count() {
3474 let provider = match create_test_provider() {
3475 Some(p) => p,
3476 None => {
3477 eprintln!("Skipping test: no CUDA device available");
3478 return;
3479 }
3480 };
3481
3482 let schema = Schema::new(vec![("id".to_string(), ScalarType::U32)]);
3483 let ids: Vec<u32> = vec![7, 11, 13, 17];
3484 let source = provider
3485 .create_buffer_from_slices(&[bytemuck::cast_slice(&ids)], schema)
3486 .unwrap();
3487 assert_eq!(
3491 source.cached_row_count(),
3492 Some(4),
3493 "source buffer should have its cached row count populated by \
3494 create_buffer_from_slices"
3495 );
3496
3497 let cloned = provider.clone_buffer(&source).unwrap();
3498
3499 assert_eq!(
3500 cloned.cached_row_count(),
3501 Some(4),
3502 "clone_buffer must propagate cached_row_count from source to clone",
3503 );
3504 }
3505
3506 #[test]
3509 fn test_hash_join_empty_inputs() {
3510 let provider = match create_test_provider() {
3511 Some(p) => p,
3512 None => {
3513 eprintln!("Skipping test: no CUDA device available");
3514 return;
3515 }
3516 };
3517
3518 let schema = Schema::new(vec![("key".to_string(), ScalarType::U32)]);
3519 let empty = create_empty_test_buffer(&provider, schema.clone());
3520
3521 let result = provider.hash_join(&empty, &empty, &[0], &[0]);
3523 assert!(result.is_ok());
3524 assert!(result.unwrap().is_empty());
3525 }
3526
3527 #[test]
3528 fn test_hash_join_validation() {
3529 let provider = match create_test_provider() {
3530 Some(p) => p,
3531 None => {
3532 eprintln!("Skipping test: no CUDA device available");
3533 return;
3534 }
3535 };
3536
3537 let left = create_test_buffer(&provider, &[1, 2, 3], "left_key");
3538 let right = create_test_buffer(&provider, &[2, 3, 4], "right_key");
3539
3540 let result = provider.hash_join(&left, &right, &[], &[0]);
3542 assert!(result.is_err());
3543
3544 let result = provider.hash_join(&left, &right, &[0], &[0, 0]);
3546 assert!(result.is_err());
3547 }
3548
3549 #[test]
3552 fn test_dedup_empty_input() {
3553 let provider = match create_test_provider() {
3554 Some(p) => p,
3555 None => {
3556 eprintln!("Skipping test: no CUDA device available");
3557 return;
3558 }
3559 };
3560
3561 let schema = Schema::new(vec![("key".to_string(), ScalarType::U32)]);
3562 let empty = create_empty_test_buffer(&provider, schema);
3563
3564 let result = provider.dedup(&empty, &[0]);
3565 assert!(result.is_ok());
3566 assert!(result.unwrap().is_empty());
3567 }
3568
3569 #[test]
3570 fn test_dedup_validation() {
3571 let provider = match create_test_provider() {
3572 Some(p) => p,
3573 None => {
3574 eprintln!("Skipping test: no CUDA device available");
3575 return;
3576 }
3577 };
3578
3579 let buffer = create_test_buffer(&provider, &[1, 1, 2, 2, 3], "key");
3580
3581 let result = provider.dedup(&buffer, &[]);
3583 assert!(result.is_err());
3584 }
3585
3586 #[test]
3587 fn test_dedup_with_duplicates() {
3588 let provider = match create_test_provider() {
3589 Some(p) => p,
3590 None => {
3591 eprintln!("Skipping test: no CUDA device available");
3592 return;
3593 }
3594 };
3595
3596 let buffer = create_test_buffer(&provider, &[3, 1, 2, 1, 3, 2], "key");
3598 let deduped = provider.dedup(&buffer, &[0]).unwrap();
3599
3600 let dedup_count = provider
3601 .device_row_count(&deduped)
3602 .expect("read dedup row count");
3603 assert_eq!(dedup_count, 3, "Should have 3 unique values");
3604
3605 let result = provider.download_column::<u32>(&deduped, 0).unwrap();
3606 assert_eq!(result, vec![1, 2, 3]);
3608 }
3609
3610 #[test]
3611 fn test_dedup_larger_input() {
3612 let provider = match create_test_provider() {
3613 Some(p) => p,
3614 None => {
3615 eprintln!("Skipping test: no CUDA device available");
3616 return;
3617 }
3618 };
3619
3620 let a: Vec<u32> = (0..500).collect();
3622 let b: Vec<u32> = (250..750).collect();
3623 let input: Vec<u32> = a.iter().chain(b.iter()).copied().collect();
3624
3625 let buffer = create_test_buffer(&provider, &input, "key");
3626 let deduped = provider.dedup(&buffer, &[0]).unwrap();
3627
3628 let dedup_count = provider
3629 .device_row_count(&deduped)
3630 .expect("read dedup row count");
3631 assert_eq!(dedup_count, 750, "Should have 750 unique values (0..750)");
3632
3633 let result = provider.download_column::<u32>(&deduped, 0).unwrap();
3635 let is_sorted = result.windows(2).all(|w| w[0] <= w[1]);
3636 assert!(is_sorted, "Output should be sorted");
3637
3638 let expected: Vec<u32> = (0..750).collect();
3640 assert_eq!(result, expected);
3641 }
3642
3643 #[test]
3646 fn test_union_empty_inputs() {
3647 let provider = match create_test_provider() {
3648 Some(p) => p,
3649 None => {
3650 eprintln!("Skipping test: no CUDA device available");
3651 return;
3652 }
3653 };
3654
3655 let schema = Schema::new(vec![("key".to_string(), ScalarType::U32)]);
3656 let empty = create_empty_test_buffer(&provider, schema.clone());
3657
3658 let result = provider.union(&empty, &empty);
3660 assert!(result.is_ok());
3661 assert!(result.unwrap().is_empty());
3662
3663 let a = create_test_buffer(&provider, &[1, 2, 3], "key");
3665 let empty2 = create_empty_test_buffer(&provider, schema);
3666 let result = provider.union(&a, &empty2);
3667 assert!(result.is_ok());
3668 let result = result.unwrap();
3669 assert_eq!(result.num_rows(), 3);
3670 }
3671
3672 #[test]
3673 fn test_union_schema_type_mismatch() {
3674 let provider = match create_test_provider() {
3675 Some(p) => p,
3676 None => {
3677 eprintln!("Skipping test: no CUDA device available");
3678 return;
3679 }
3680 };
3681
3682 let a = create_test_buffer(&provider, &[1, 2], "col_a");
3683 let b = create_test_buffer(&provider, &[3, 4], "col_b");
3684
3685 let result = provider.union(&a, &b);
3687 assert!(result.is_ok());
3688
3689 let two_col_schema = Schema::new(vec![
3691 ("x".to_string(), ScalarType::U32),
3692 ("y".to_string(), ScalarType::U32),
3693 ]);
3694 let c = provider
3695 .create_buffer_from_u32_columns(&[&[1, 2], &[3, 4]], two_col_schema)
3696 .unwrap();
3697 let result = provider.union(&a, &c);
3698 assert!(result.is_err());
3699 }
3700
3701 #[test]
3704 fn test_diff_empty_inputs() {
3705 let provider = match create_test_provider() {
3706 Some(p) => p,
3707 None => {
3708 eprintln!("Skipping test: no CUDA device available");
3709 return;
3710 }
3711 };
3712
3713 let schema = Schema::new(vec![("key".to_string(), ScalarType::U32)]);
3714 let empty = create_empty_test_buffer(&provider, schema.clone());
3715
3716 let result = provider.diff(&empty, &empty);
3718 assert!(result.is_ok());
3719 assert!(result.unwrap().is_empty());
3720
3721 let a = create_test_buffer(&provider, &[1, 2, 3], "key");
3723 let empty2 = create_empty_test_buffer(&provider, schema);
3724 let result = provider.diff(&a, &empty2);
3725 assert!(result.is_ok());
3726 let result = result.unwrap();
3727 assert_eq!(result.num_rows(), 3);
3728 }
3729
3730 #[test]
3731 fn test_diff_basic() {
3732 let provider = match create_test_provider() {
3733 Some(p) => p,
3734 None => {
3735 eprintln!("Skipping test: no CUDA device available");
3736 return;
3737 }
3738 };
3739
3740 let a = create_test_buffer(&provider, &[1, 2, 3, 4, 5], "key");
3741 let b = create_test_buffer(&provider, &[2, 4], "key");
3742
3743 let result = provider.diff(&a, &b);
3744 assert!(result.is_ok());
3745 let result = result.unwrap();
3746 assert_eq!(result.num_rows(), 3); let values = read_buffer_u32(&provider, &result, 0);
3749 assert_eq!(values, vec![1, 3, 5]);
3750 }
3751
3752 #[test]
3753 fn test_diff_all_filtered_out() {
3754 let provider = match create_test_provider() {
3755 Some(p) => p,
3756 None => {
3757 eprintln!("Skipping test: no CUDA device available");
3758 return;
3759 }
3760 };
3761
3762 let a = create_test_buffer(&provider, &[1, 2, 3], "key");
3763 let b = create_test_buffer(&provider, &[1, 2, 3, 4, 5], "key");
3764
3765 let result = provider.diff(&a, &b);
3766 assert!(result.is_ok());
3767 assert!(result.unwrap().is_empty());
3768 }
3769
3770 #[test]
3771 fn test_diff_schema_mismatch() {
3772 let provider = match create_test_provider() {
3773 Some(p) => p,
3774 None => {
3775 eprintln!("Skipping test: no CUDA device available");
3776 return;
3777 }
3778 };
3779
3780 let a = create_test_buffer(&provider, &[1, 2], "col_a");
3782 let b = create_test_buffer(&provider, &[1, 2], "col_b");
3783 let result = provider.diff(&a, &b);
3784 assert!(
3785 result.is_ok(),
3786 "Same types with different names should succeed"
3787 );
3788
3789 let schema_2col = Schema::new(vec![
3791 ("c0".to_string(), ScalarType::U32),
3792 ("c1".to_string(), ScalarType::U32),
3793 ]);
3794
3795 let bytes_2col: Vec<u8> = [1u32, 2, 3, 4]
3796 .iter()
3797 .flat_map(|v| v.to_le_bytes())
3798 .collect();
3799 let mut col0 = provider
3800 .memory()
3801 .alloc::<u8>(bytes_2col.len() / 2)
3802 .expect("alloc");
3803 let mut col1 = provider
3804 .memory()
3805 .alloc::<u8>(bytes_2col.len() / 2)
3806 .expect("alloc");
3807 provider
3808 .device()
3809 .inner()
3810 .htod_sync_copy_into(&bytes_2col[..8], &mut col0)
3811 .expect("htod");
3812 provider
3813 .device()
3814 .inner()
3815 .htod_sync_copy_into(&bytes_2col[8..], &mut col1)
3816 .expect("htod");
3817 let buffer_2col = provider
3818 .buffer_from_columns(vec![col0.into(), col1.into()], 2, schema_2col)
3819 .expect("buffer");
3820
3821 let buffer_1col = create_test_buffer(&provider, &[1, 2], "c0");
3822
3823 let result = provider.diff(&buffer_2col, &buffer_1col);
3824 assert!(result.is_err(), "Different arities should fail");
3825 }
3826
3827 #[test]
3830 fn test_groupby_empty_input() {
3831 let provider = match create_test_provider() {
3832 Some(p) => p,
3833 None => {
3834 eprintln!("Skipping test: no CUDA device available");
3835 return;
3836 }
3837 };
3838
3839 let schema = Schema::new(vec![("key".to_string(), ScalarType::U32)]);
3840 let empty = create_empty_test_buffer(&provider, schema);
3841
3842 let result = provider.groupby_agg(&empty, &[0], AggOp::Count, 0);
3843 assert!(result.is_ok());
3844 assert!(result.unwrap().is_empty());
3845 }
3846
3847 #[test]
3848 fn test_groupby_validation() {
3849 let provider = match create_test_provider() {
3850 Some(p) => p,
3851 None => {
3852 eprintln!("Skipping test: no CUDA device available");
3853 return;
3854 }
3855 };
3856
3857 let buffer = create_test_buffer(&provider, &[1, 1, 2, 2, 3], "key");
3858
3859 let result = provider.groupby_agg(&buffer, &[], AggOp::Count, 0);
3861 assert!(result.is_err());
3862
3863 let result = provider.groupby_agg(&buffer, &[0], AggOp::Count, 5);
3865 assert!(result.is_err());
3866 }
3867
3868 #[test]
3869 fn test_groupby_logsumexp() {
3870 let provider = match create_test_provider() {
3871 Some(p) => p,
3872 None => {
3873 eprintln!("Skipping test: no CUDA device available");
3874 return;
3875 }
3876 };
3877
3878 let keys: Vec<u32> = vec![1, 1, 2, 2];
3882 let values: Vec<f64> = vec![1.0, 2.0, 3.0, 4.0];
3883
3884 let schema = Schema::new(vec![
3885 ("key".to_string(), ScalarType::U32),
3886 ("value".to_string(), ScalarType::F64),
3887 ]);
3888
3889 let key_bytes: Vec<u8> = keys.iter().flat_map(|v| v.to_le_bytes()).collect();
3891 let mut key_col = provider
3892 .memory()
3893 .alloc::<u8>(key_bytes.len())
3894 .expect("alloc key");
3895 provider
3896 .device()
3897 .inner()
3898 .htod_sync_copy_into(&key_bytes, &mut key_col)
3899 .expect("upload key");
3900
3901 let val_bytes: Vec<u8> = values.iter().flat_map(|v| v.to_le_bytes()).collect();
3903 let mut val_col = provider
3904 .memory()
3905 .alloc::<u8>(val_bytes.len())
3906 .expect("alloc val");
3907 provider
3908 .device()
3909 .inner()
3910 .htod_sync_copy_into(&val_bytes, &mut val_col)
3911 .expect("upload val");
3912
3913 let buffer = provider
3914 .buffer_from_columns(vec![key_col.into(), val_col.into()], 4, schema)
3915 .expect("buffer");
3916
3917 let result = provider.groupby_agg(&buffer, &[0], AggOp::LogSumExp, 1);
3919 assert!(
3920 result.is_ok(),
3921 "groupby_agg with LogSumExp should succeed: {:?}",
3922 result.err()
3923 );
3924
3925 let result = result.unwrap();
3926 let group_count = provider
3927 .device_row_count(&result)
3928 .expect("read group count");
3929 assert_eq!(group_count, 2, "Should have 2 groups");
3930
3931 let result_values = provider
3933 .download_column::<f64>(&result, 1)
3934 .expect("download result");
3935
3936 let expected_0 = 2.0_f64 + ((-1.0_f64).exp() + 1.0_f64).ln(); let expected_1 = 4.0_f64 + ((-1.0_f64).exp() + 1.0_f64).ln(); let tolerance = 1e-5;
3943 assert!(
3944 (result_values[0] - expected_0).abs() < tolerance,
3945 "Group 0 logsumexp mismatch: got {}, expected {}",
3946 result_values[0],
3947 expected_0
3948 );
3949 assert!(
3950 (result_values[1] - expected_1).abs() < tolerance,
3951 "Group 1 logsumexp mismatch: got {}, expected {}",
3952 result_values[1],
3953 expected_1
3954 );
3955 }
3956
3957 #[test]
3960 fn test_combine_schemas() {
3961 let provider = match create_test_provider() {
3962 Some(p) => p,
3963 None => {
3964 eprintln!("Skipping test: no CUDA device available");
3965 return;
3966 }
3967 };
3968
3969 let left = Schema::new(vec![("a".to_string(), ScalarType::U32)]);
3970 let right = Schema::new(vec![("b".to_string(), ScalarType::U64)]);
3971
3972 let combined = provider.combine_schemas(&left, &right);
3973 assert_eq!(combined.arity(), 2);
3974 assert_eq!(combined.column_type(0), Some(ScalarType::U32));
3975 assert_eq!(combined.column_type(1), Some(ScalarType::U64));
3976 }
3977
3978 #[test]
3979 fn test_groupby_result_schema() {
3980 let provider = match create_test_provider() {
3981 Some(p) => p,
3982 None => {
3983 eprintln!("Skipping test: no CUDA device available");
3984 return;
3985 }
3986 };
3987
3988 let input = Schema::new(vec![
3989 ("key".to_string(), ScalarType::U32),
3990 ("value".to_string(), ScalarType::U32),
3991 ]);
3992
3993 let count_schema =
3995 provider.groupby_multi_agg_result_schema(&input, &[0], &[(1, AggOp::Count)]);
3996 assert_eq!(count_schema.arity(), 2);
3997 assert_eq!(count_schema.column_type(1), Some(ScalarType::U64));
3998
3999 let sum_schema = provider.groupby_multi_agg_result_schema(&input, &[0], &[(1, AggOp::Sum)]);
4001 assert_eq!(sum_schema.arity(), 2);
4002 assert_eq!(sum_schema.column_type(1), Some(ScalarType::U64));
4003
4004 let min_schema = provider.groupby_multi_agg_result_schema(&input, &[0], &[(1, AggOp::Min)]);
4006 assert_eq!(min_schema.arity(), 2);
4007 assert_eq!(min_schema.column_type(1), Some(ScalarType::U32));
4008 }
4009
4010 #[test]
4011 fn test_groupby_multi_agg_sum_returns_u64_schema() {
4012 let provider = match create_test_provider() {
4013 Some(p) => p,
4014 None => {
4015 eprintln!("Skipping test: no CUDA device");
4016 return;
4017 }
4018 };
4019
4020 let schema = Schema::new(vec![
4021 ("key".to_string(), ScalarType::U32),
4022 ("val".to_string(), ScalarType::U32),
4023 ]);
4024
4025 let result_schema =
4026 provider.groupby_multi_agg_result_schema(&schema, &[0], &[(1, AggOp::Sum)]);
4027
4028 assert_eq!(
4030 result_schema.column_type(1),
4031 Some(ScalarType::U64),
4032 "Sum aggregation should return U64 type, not U32"
4033 );
4034 }
4035
4036 #[test]
4037 fn test_join_custom_max_output() {
4038 let provider = match create_test_provider() {
4039 Some(p) => p,
4040 None => {
4041 eprintln!("Skipping test: no CUDA device available");
4042 return;
4043 }
4044 };
4045
4046 let left = create_test_buffer(&provider, &[1, 1, 1, 1, 2, 2, 2, 2], "left_key");
4051 let right = create_test_buffer(&provider, &[1, 1, 1, 2, 2, 2], "right_key");
4052
4053 let result_limited = provider
4055 .hash_join_v2_with_limit(&left, &right, &[0], &[0], JoinType::Inner, Some(10))
4056 .expect("join with limit should succeed");
4057 assert!(
4058 result_limited.num_rows() <= 10,
4059 "With limit 10, got {} rows but expected at most 10",
4060 result_limited.num_rows()
4061 );
4062
4063 let result_unlimited = provider
4065 .hash_join_v2_with_limit(&left, &right, &[0], &[0], JoinType::Inner, None)
4066 .expect("join without limit should succeed");
4067 assert_eq!(
4068 result_unlimited.num_rows(),
4069 24,
4070 "Without limit, expected 24 rows but got {}",
4071 result_unlimited.num_rows()
4072 );
4073
4074 let result_legacy = provider
4076 .hash_join_v2(&left, &right, &[0], &[0], JoinType::Inner)
4077 .expect("legacy hash_join_v2 should succeed");
4078 assert_eq!(
4079 result_legacy.num_rows(),
4080 24,
4081 "Legacy API without limit, expected 24 rows but got {}",
4082 result_legacy.num_rows()
4083 );
4084 }
4085
4086 fn create_arith_test_provider() -> Option<CudaKernelProvider> {
4090 if !has_cuda_device() {
4091 return None;
4092 }
4093 let device = Arc::new(CudaDevice::new(0).ok()?);
4094 let budget = MemoryBudget::with_limit(1024 * 1024 * 1024);
4095 let memory = Arc::new(GpuMemoryManager::new(device.clone(), budget));
4096 CudaKernelProvider::new(device, memory).ok()
4097 }
4098
4099 fn create_i64_buffer(provider: &CudaKernelProvider, data: &[i64]) -> CudaBuffer {
4101 let schema = Schema::new(vec![("col".to_string(), ScalarType::I64)]);
4102 provider
4103 .create_buffer_from_slice::<i64>(data, schema)
4104 .unwrap()
4105 }
4106
4107 fn create_f64_buffer(provider: &CudaKernelProvider, data: &[f64]) -> CudaBuffer {
4109 let schema = Schema::new(vec![("col".to_string(), ScalarType::F64)]);
4110 provider
4111 .create_buffer_from_slice::<f64>(data, schema)
4112 .unwrap()
4113 }
4114
4115 #[test]
4116 fn test_add_columns_i64() {
4117 let Some(provider) = create_arith_test_provider() else {
4118 eprintln!("Skipping test: no CUDA device available");
4119 return;
4120 };
4121
4122 let a = create_i64_buffer(&provider, &[1, 2, 3, 4, 5]);
4123 let b = create_i64_buffer(&provider, &[10, 20, 30, 40, 50]);
4124
4125 let result = provider.add_columns(&a, &b).unwrap();
4126 let values = provider.download_column::<i64>(&result, 0).unwrap();
4127
4128 assert_eq!(values, vec![11, 22, 33, 44, 55]);
4129 }
4130
4131 #[test]
4132 fn test_sub_columns_i64() {
4133 let Some(provider) = create_arith_test_provider() else {
4134 eprintln!("Skipping test: no CUDA device available");
4135 return;
4136 };
4137
4138 let a = create_i64_buffer(&provider, &[10, 20, 30, 40, 50]);
4139 let b = create_i64_buffer(&provider, &[1, 2, 3, 4, 5]);
4140
4141 let result = provider.sub_columns(&a, &b).unwrap();
4142 let values = provider.download_column::<i64>(&result, 0).unwrap();
4143
4144 assert_eq!(values, vec![9, 18, 27, 36, 45]);
4145 }
4146
4147 #[test]
4148 fn test_mul_columns_i64() {
4149 let Some(provider) = create_arith_test_provider() else {
4150 eprintln!("Skipping test: no CUDA device available");
4151 return;
4152 };
4153
4154 let a = create_i64_buffer(&provider, &[2, 3, 4, 5, 6]);
4155 let b = create_i64_buffer(&provider, &[3, 4, 5, 6, 7]);
4156
4157 let result = provider.mul_columns(&a, &b).unwrap();
4158 let values = provider.download_column::<i64>(&result, 0).unwrap();
4159
4160 assert_eq!(values, vec![6, 12, 20, 30, 42]);
4161 }
4162
4163 #[test]
4164 fn test_div_columns_i64() {
4165 let Some(provider) = create_arith_test_provider() else {
4166 eprintln!("Skipping test: no CUDA device available");
4167 return;
4168 };
4169
4170 let a = create_i64_buffer(&provider, &[100, 200, 300, 400]);
4171 let b = create_i64_buffer(&provider, &[10, 20, 30, 40]);
4172
4173 let result = provider.div_columns(&a, &b).unwrap();
4174 let values = provider.download_column::<i64>(&result, 0).unwrap();
4175
4176 assert_eq!(values, vec![10, 10, 10, 10]);
4177 }
4178
4179 #[test]
4180 fn test_div_columns_by_zero() {
4181 let Some(provider) = create_arith_test_provider() else {
4182 eprintln!("Skipping test: no CUDA device available");
4183 return;
4184 };
4185
4186 let a = create_i64_buffer(&provider, &[10, 20, 30]);
4187 let b = create_i64_buffer(&provider, &[2, 0, 3]); let result = provider.div_columns(&a, &b).unwrap();
4190 let values = provider.download_column::<i64>(&result, 0).unwrap();
4191
4192 assert_eq!(values, vec![5, i64::MAX, 10]);
4194 }
4195
4196 #[test]
4197 fn test_mod_columns_i64() {
4198 let Some(provider) = create_arith_test_provider() else {
4199 eprintln!("Skipping test: no CUDA device available");
4200 return;
4201 };
4202
4203 let a = create_i64_buffer(&provider, &[17, 23, 100, 7]);
4204 let b = create_i64_buffer(&provider, &[5, 7, 30, 3]);
4205
4206 let result = provider.mod_columns(&a, &b).unwrap();
4207 let values = provider.download_column::<i64>(&result, 0).unwrap();
4208
4209 assert_eq!(values, vec![2, 2, 10, 1]);
4210 }
4211
4212 #[test]
4213 fn test_mod_columns_by_zero() {
4214 let Some(provider) = create_arith_test_provider() else {
4215 eprintln!("Skipping test: no CUDA device available");
4216 return;
4217 };
4218
4219 let a = create_i64_buffer(&provider, &[10, 20]);
4220 let b = create_i64_buffer(&provider, &[3, 0]); let result = provider.mod_columns(&a, &b).unwrap();
4223 let values = provider.download_column::<i64>(&result, 0).unwrap();
4224
4225 assert_eq!(values, vec![1, 0]);
4227 }
4228
4229 #[test]
4230 fn test_abs_column_i64() {
4231 let Some(provider) = create_arith_test_provider() else {
4232 eprintln!("Skipping test: no CUDA device available");
4233 return;
4234 };
4235
4236 let a = create_i64_buffer(&provider, &[-5, 10, -15, 20, 0]);
4237
4238 let result = provider.abs_column(&a).unwrap();
4239 let values = provider.download_column::<i64>(&result, 0).unwrap();
4240
4241 assert_eq!(values, vec![5, 10, 15, 20, 0]);
4242 }
4243
4244 #[test]
4245 fn test_min_columns_i64() {
4246 let Some(provider) = create_arith_test_provider() else {
4247 eprintln!("Skipping test: no CUDA device available");
4248 return;
4249 };
4250
4251 let a = create_i64_buffer(&provider, &[5, 10, 15, 20]);
4252 let b = create_i64_buffer(&provider, &[3, 12, 10, 25]);
4253
4254 let result = provider.min_columns(&a, &b).unwrap();
4255 let values = provider.download_column::<i64>(&result, 0).unwrap();
4256
4257 assert_eq!(values, vec![3, 10, 10, 20]);
4258 }
4259
4260 #[test]
4261 fn test_max_columns_i64() {
4262 let Some(provider) = create_arith_test_provider() else {
4263 eprintln!("Skipping test: no CUDA device available");
4264 return;
4265 };
4266
4267 let a = create_i64_buffer(&provider, &[5, 10, 15, 20]);
4268 let b = create_i64_buffer(&provider, &[3, 12, 10, 25]);
4269
4270 let result = provider.max_columns(&a, &b).unwrap();
4271 let values = provider.download_column::<i64>(&result, 0).unwrap();
4272
4273 assert_eq!(values, vec![5, 12, 15, 25]);
4274 }
4275
4276 #[test]
4277 fn test_add_columns_f64() {
4278 let Some(provider) = create_arith_test_provider() else {
4279 eprintln!("Skipping test: no CUDA device available");
4280 return;
4281 };
4282
4283 let a = create_f64_buffer(&provider, &[1.5, 2.5, 3.5]);
4284 let b = create_f64_buffer(&provider, &[0.5, 1.5, 2.5]);
4285
4286 let result = provider.add_columns(&a, &b).unwrap();
4287 let values = provider.download_column::<f64>(&result, 0).unwrap();
4288
4289 assert_eq!(values, vec![2.0, 4.0, 6.0]);
4290 }
4291
4292 #[test]
4293 fn test_mul_columns_f64() {
4294 let Some(provider) = create_arith_test_provider() else {
4295 eprintln!("Skipping test: no CUDA device available");
4296 return;
4297 };
4298
4299 let a = create_f64_buffer(&provider, &[2.0, 3.0, 4.0]);
4300 let b = create_f64_buffer(&provider, &[1.5, 2.0, 2.5]);
4301
4302 let result = provider.mul_columns(&a, &b).unwrap();
4303 let values = provider.download_column::<f64>(&result, 0).unwrap();
4304
4305 assert_eq!(values, vec![3.0, 6.0, 10.0]);
4306 }
4307
4308 #[test]
4309 fn test_div_columns_f64_by_zero() {
4310 let Some(provider) = create_arith_test_provider() else {
4311 eprintln!("Skipping test: no CUDA device available");
4312 return;
4313 };
4314
4315 let a = create_f64_buffer(&provider, &[1.0, -1.0, 0.0]);
4316 let b = create_f64_buffer(&provider, &[0.0, 0.0, 0.0]);
4317
4318 let result = provider.div_columns(&a, &b).unwrap();
4319 let values = provider.download_column::<f64>(&result, 0).unwrap();
4320
4321 assert!(values[0].is_infinite() && values[0].is_sign_positive());
4323 assert!(values[1].is_infinite() && values[1].is_sign_negative());
4324 assert!(values[2].is_nan());
4325 }
4326
4327 #[test]
4328 fn test_pow_columns() {
4329 let Some(provider) = create_arith_test_provider() else {
4330 eprintln!("Skipping test: no CUDA device available");
4331 return;
4332 };
4333
4334 let base = create_i64_buffer(&provider, &[2, 3, 4, 5]);
4335 let exp = create_i64_buffer(&provider, &[3, 2, 2, 1]);
4336
4337 let result = provider.pow_columns(&base, &exp).unwrap();
4338 let values = provider.download_column::<f64>(&result, 0).unwrap();
4339
4340 assert_eq!(values, vec![8.0, 9.0, 16.0, 5.0]);
4342 }
4343
4344 #[test]
4345 fn test_pow_columns_fractional_exp() {
4346 let Some(provider) = create_arith_test_provider() else {
4347 eprintln!("Skipping test: no CUDA device available");
4348 return;
4349 };
4350
4351 let base = create_f64_buffer(&provider, &[4.0, 9.0, 27.0]);
4352 let exp = create_f64_buffer(&provider, &[0.5, 0.5, 1.0 / 3.0]);
4353
4354 let result = provider.pow_columns(&base, &exp).unwrap();
4355 let values = provider.download_column::<f64>(&result, 0).unwrap();
4356
4357 assert!((values[0] - 2.0).abs() < 1e-10);
4359 assert!((values[1] - 3.0).abs() < 1e-10);
4360 assert!((values[2] - 3.0).abs() < 1e-10);
4361 }
4362
4363 #[test]
4364 fn test_cast_i64_to_f64() {
4365 let Some(provider) = create_arith_test_provider() else {
4366 eprintln!("Skipping test: no CUDA device available");
4367 return;
4368 };
4369
4370 let a = create_i64_buffer(&provider, &[1, 2, 3, 4, 5]);
4371
4372 let result = provider.cast_column(&a, ScalarType::F64).unwrap();
4373 let values = provider.download_column::<f64>(&result, 0).unwrap();
4374
4375 assert_eq!(values, vec![1.0, 2.0, 3.0, 4.0, 5.0]);
4376 }
4377
4378 #[test]
4379 fn test_cast_f64_to_i64() {
4380 let Some(provider) = create_arith_test_provider() else {
4381 eprintln!("Skipping test: no CUDA device available");
4382 return;
4383 };
4384
4385 let a = create_f64_buffer(&provider, &[1.9, 2.1, 3.5, 4.0, 5.7]);
4386
4387 let result = provider.cast_column(&a, ScalarType::I64).unwrap();
4388 let values = provider.download_column::<i64>(&result, 0).unwrap();
4389
4390 assert_eq!(values, vec![1, 2, 3, 4, 5]);
4392 }
4393
4394 #[test]
4395 fn test_cast_i64_to_i32() {
4396 let Some(provider) = create_arith_test_provider() else {
4397 eprintln!("Skipping test: no CUDA device available");
4398 return;
4399 };
4400
4401 let a = create_i64_buffer(&provider, &[1, 2, 3, 100, 200]);
4402
4403 let result = provider.cast_column(&a, ScalarType::I32).unwrap();
4404 let values = provider.download_column::<i32>(&result, 0).unwrap();
4405
4406 assert_eq!(values, vec![1, 2, 3, 100, 200]);
4407 }
4408
4409 #[test]
4410 fn test_arithmetic_row_count_mismatch() {
4411 let Some(provider) = create_arith_test_provider() else {
4412 eprintln!("Skipping test: no CUDA device available");
4413 return;
4414 };
4415
4416 let a = create_i64_buffer(&provider, &[1, 2, 3]);
4417 let b = create_i64_buffer(&provider, &[1, 2]); let result = provider.add_columns(&a, &b);
4420 assert!(result.is_err());
4421 let err = result.err().unwrap();
4422 assert!(err.to_string().contains("Row count mismatch"));
4423 }
4424
4425 #[test]
4426 fn test_arithmetic_empty_buffers() {
4427 let Some(provider) = create_arith_test_provider() else {
4428 eprintln!("Skipping test: no CUDA device available");
4429 return;
4430 };
4431
4432 let a = create_i64_buffer(&provider, &[]);
4433 let b = create_i64_buffer(&provider, &[]);
4434
4435 let result = provider.add_columns(&a, &b).unwrap();
4436 let values = provider.download_column::<i64>(&result, 0).unwrap();
4437
4438 assert_eq!(values, Vec::<i64>::new());
4439 }
4440
4441 #[test]
4442 fn test_wrapping_arithmetic_overflow() {
4443 let Some(provider) = create_arith_test_provider() else {
4444 eprintln!("Skipping test: no CUDA device available");
4445 return;
4446 };
4447
4448 let a = create_i64_buffer(&provider, &[i64::MAX, i64::MIN]);
4449 let b = create_i64_buffer(&provider, &[1, -1]);
4450
4451 let add_result = provider.add_columns(&a, &b).unwrap();
4453 let add_values = provider.download_column::<i64>(&add_result, 0).unwrap();
4454 assert_eq!(add_values[0], i64::MIN); assert_eq!(add_values[1], i64::MAX); }
4457
4458 #[test]
4459 fn test_abs_column_f64() {
4460 let Some(provider) = create_arith_test_provider() else {
4461 eprintln!("Skipping test: no CUDA device available");
4462 return;
4463 };
4464
4465 let a = create_f64_buffer(&provider, &[-1.5, 2.5, -3.5, 0.0]);
4466
4467 let result = provider.abs_column(&a).unwrap();
4468 let values = provider.download_column::<f64>(&result, 0).unwrap();
4469
4470 assert_eq!(values, vec![1.5, 2.5, 3.5, 0.0]);
4471 }
4472
4473 #[test]
4474 fn test_min_max_columns_f64() {
4475 let Some(provider) = create_arith_test_provider() else {
4476 eprintln!("Skipping test: no CUDA device available");
4477 return;
4478 };
4479
4480 let a = create_f64_buffer(&provider, &[1.5, 5.0, 3.0]);
4481 let b = create_f64_buffer(&provider, &[2.0, 3.0, 4.0]);
4482
4483 let min_result = provider.min_columns(&a, &b).unwrap();
4484 let min_values = provider.download_column::<f64>(&min_result, 0).unwrap();
4485 assert_eq!(min_values, vec![1.5, 3.0, 3.0]);
4486
4487 let max_result = provider.max_columns(&a, &b).unwrap();
4488 let max_values = provider.download_column::<f64>(&max_result, 0).unwrap();
4489 assert_eq!(max_values, vec![2.0, 5.0, 4.0]);
4490 }
4491}