1use std::collections::HashMap;
7use std::path::PathBuf;
8use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
9use std::sync::{Arc, Mutex, OnceLock};
10
11use xlog_core::{resolve_bool, Result, Schema, XlogError};
12
13use crate::{
14 cuda_compat::{DeviceRepr, DeviceSlice, LaunchAsync, LaunchConfig},
15 cuda_graph::{CapturedCudaGraph, CsmCudaGraphKey, CudaGraphNode},
16 memory::{
17 validate_logical_row_count, CudaColumn, DeviceMemoryView, DeviceRead, TrackedCudaSlice,
18 },
19 CudaBuffer, CudaDevice, GpuMemoryManager,
20};
21
22static NEXT_PROVIDER_IDENTITY: AtomicU64 = AtomicU64::new(1);
23
24mod arithmetic;
25mod builder;
26mod filter;
27mod fj;
28mod fj_delta;
29mod fj_delta_sparse;
30mod groupby;
31mod ilp;
32mod ilp_exact;
33mod ilp_exact_nary;
34mod io;
35mod kernel_loading;
36pub use builder::CudaProviderBuilder;
37pub mod kernel_paths;
38mod launch_safe;
39mod probabilistic;
40mod relational;
41pub mod resident_filter_project;
42pub mod resident_relational;
43pub mod resident_schedule;
44mod transfer;
45mod wcoj;
46mod wcoj_metadata;
47mod wcoj_project;
48
49pub use fj::{FjNode, FjPlan, FjSubAtom};
50pub use fj_delta::{FjDeltaCols, FJ_DELTA_MAX_DOMAIN};
51pub use ilp_exact_nary::{IlpExactNaryPatterns, IlpExactNaryRequest};
52
53#[derive(Debug, Clone, Default)]
55pub struct PtxLoadProfile {
56 pub total_sec: f64,
57 pub per_module_sec: Vec<(String, f64)>,
58 pub cubin_loaded: u32,
59 pub ptx_fallback: u32,
60}
61
62fn warmup_profiling_enabled() -> Result<bool> {
63 resolve_bool(None, "XLOG_WARMUP_PROFILE", false)
64}
65
66pub(crate) fn detect_compute_capability(device: &Arc<CudaDevice>) -> Result<u32> {
68 let major = device
69 .inner()
70 .attribute(
71 cudarc::driver::sys::CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
72 )
73 .map_err(|e| XlogError::Kernel(format!("Failed to query SM major: {}", e)))?;
74 let minor = device
75 .inner()
76 .attribute(
77 cudarc::driver::sys::CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
78 )
79 .map_err(|e| XlogError::Kernel(format!("Failed to query SM minor: {}", e)))?;
80 Ok((major as u32) * 10 + (minor as u32))
81}
82
83#[derive(Debug)]
84pub(crate) enum KernelModuleSource {
85 File { path: PathBuf, is_cubin: bool },
86 EmbeddedPortablePtx { ptx: &'static str },
87}
88
89pub(crate) fn resolve_module_sources_with_locator(
90 name: &str,
91 cc: u32,
92 locator: &kernel_paths::KernelArtifactLocator,
93) -> Vec<KernelModuleSource> {
94 let mut sources: Vec<KernelModuleSource> = locator
95 .resolve_module_paths(name, cc)
96 .into_iter()
97 .filter(|(path, _)| !staged_artifact_is_stale(path))
102 .map(|(path, is_cubin)| KernelModuleSource::File { path, is_cubin })
103 .collect();
104
105 if let Some(ptx) = crate::embedded_kernel_data::portable_ptx(name) {
112 sources.push(KernelModuleSource::EmbeddedPortablePtx { ptx });
113 }
114 sources
115}
116
117fn staged_artifact_is_stale(path: &std::path::Path) -> bool {
126 let Some(file_name) = path.file_name().and_then(|n| n.to_str()) else {
127 return false;
128 };
129 let Some(expected) = crate::embedded_kernel_data::canonical_artifact_hash(file_name) else {
130 return false;
131 };
132 match std::fs::read(path) {
133 Ok(bytes) => fnv1a_64(&bytes) != expected,
134 Err(_) => false,
135 }
136}
137
138fn fnv1a_64(bytes: &[u8]) -> u64 {
140 let mut hash: u64 = 0xcbf2_9ce4_8422_2325;
141 for &byte in bytes {
142 hash ^= byte as u64;
143 hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
144 }
145 hash
146}
147
148#[cfg(test)]
149mod kernel_source_resolution_tests {
150 use super::{
151 kernel_paths::KernelArtifactLocator, resolve_module_sources_with_locator,
152 KernelModuleSource,
153 };
154 use std::fs;
155
156 #[test]
157 fn keeps_portable_ptx_fallback_when_cubin_exists() {
158 let root = std::env::temp_dir().join(format!(
159 "xlog-kernel-fallback-{}-{}",
160 std::process::id(),
161 std::time::SystemTime::now()
162 .duration_since(std::time::UNIX_EPOCH)
163 .expect("system clock before UNIX_EPOCH")
164 .as_nanos()
165 ));
166 let kernels = root.join("kernels");
167 fs::create_dir_all(&kernels).expect("create kernels dir");
168 fs::write(kernels.join("fakekernel.sm_86.cubin"), b"cubin").expect("write cubin");
173 fs::write(kernels.join("fakekernel.portable.ptx"), b"ptx").expect("write ptx");
174 let expected_cubin = kernels.join("fakekernel.sm_86.cubin");
175 let expected_ptx = kernels.join("fakekernel.portable.ptx");
176
177 let locator = KernelArtifactLocator::new(None, Some(kernels.clone()), None);
178 let sources = resolve_module_sources_with_locator("fakekernel", 86, &locator);
179
180 assert_eq!(sources.len(), 2);
181 assert!(matches!(
182 &sources[0],
183 KernelModuleSource::File {
184 path,
185 is_cubin: true
186 } if path == &expected_cubin
187 ));
188 assert!(matches!(
189 &sources[1],
190 KernelModuleSource::File {
191 path,
192 is_cubin: false
193 } if path == &expected_ptx
194 ));
195
196 fs::remove_dir_all(root).expect("remove temp kernels");
197 }
198
199 #[test]
204 fn fnv1a_64_matches_known_vectors() {
205 assert_eq!(super::fnv1a_64(b""), 0xcbf2_9ce4_8422_2325);
206 assert_eq!(super::fnv1a_64(b"a"), 0xaf63_dc4c_8601_ec8c);
207 assert_eq!(super::fnv1a_64(b"foobar"), 0x85944171_f73967e8);
208 }
209
210 #[test]
214 fn staged_artifact_not_stale_without_canonical_hash() {
215 let root = std::env::temp_dir().join(format!(
216 "xlog-kernel-stale-{}-{}",
217 std::process::id(),
218 std::time::SystemTime::now()
219 .duration_since(std::time::UNIX_EPOCH)
220 .expect("system clock before UNIX_EPOCH")
221 .as_nanos()
222 ));
223 fs::create_dir_all(&root).expect("create dir");
224 let unknown = root.join("definitely_not_a_real_kernel.sm_86.cubin");
225 fs::write(&unknown, b"bytes").expect("write");
226 assert!(!super::staged_artifact_is_stale(&unknown));
227 assert!(!super::staged_artifact_is_stale(
228 &root.join("missing.portable.ptx")
229 ));
230 fs::remove_dir_all(root).expect("remove temp dir");
231 }
232}
233
234pub(crate) fn load_module_sources(name: &str, cc: u32) -> Result<Vec<KernelModuleSource>> {
239 debug_assert!(
240 crate::kernel_manifest_data::KERNEL_CU_NAMES.contains(&name),
241 "kernel module '{name}' is not in KERNEL_CU_NAMES manifest — update kernel_manifest_data.rs"
242 );
243 let locator = kernel_paths::KernelArtifactLocator::from_env();
244 let sources = resolve_module_sources_with_locator(name, cc, &locator);
245 if sources.is_empty() {
246 Err(XlogError::Kernel(format!(
247 "{name}: no cubin, sidecar portable PTX, or embedded portable PTX found"
248 )))
249 } else {
250 Ok(sources)
251 }
252}
253
254pub(crate) struct MultiblockScanScratchU32 {
267 levels: Vec<TrackedCudaSlice<u32>>,
268}
269
270impl MultiblockScanScratchU32 {
271 pub(crate) fn levels(&self) -> &[TrackedCudaSlice<u32>] {
272 &self.levels
273 }
274}
275
276pub(crate) struct CsmCudaGraphNodes {
277 pub(crate) count: CudaGraphNode,
278 pub(crate) total: CudaGraphNode,
279 pub(crate) materialize: CudaGraphNode,
280 pub(crate) node_count: usize,
281}
282
283pub(crate) struct CsmCudaGraphEntry {
284 pub(crate) graph: CapturedCudaGraph,
285 pub(crate) nodes: CsmCudaGraphNodes,
286 pub(crate) per_probe_count: TrackedCudaSlice<u32>,
287 pub(crate) per_probe_offsets: TrackedCudaSlice<u32>,
288 pub(crate) d_logical_count: TrackedCudaSlice<u32>,
289 pub(crate) d_overflow: TrackedCudaSlice<u8>,
290 pub(crate) d_output_left: TrackedCudaSlice<u32>,
291 pub(crate) d_output_right: TrackedCudaSlice<u32>,
292 pub(crate) scan_scratch: MultiblockScanScratchU32,
293 pub(crate) probe_capacity: u32,
294 pub(crate) output_capacity: u32,
295}
296
297pub struct RadixSortScratch {
299 keys_b: TrackedCudaSlice<u32>,
300 values_b: TrackedCudaSlice<u32>,
301 hist: TrackedCudaSlice<u32>,
302 prefix: TrackedCudaSlice<u32>,
303 ranks: TrackedCudaSlice<u32>,
304 len: u32,
305}
306
307impl RadixSortScratch {
308 pub fn new(provider: &CudaKernelProvider, n: u32) -> Result<Self> {
309 let memory = provider.memory();
310 let len = n.max(1);
311 let keys_b = memory.alloc::<u32>(len as usize)?;
312 let values_b = memory.alloc::<u32>(len as usize)?;
313 let ranks = memory.alloc::<u32>(len as usize)?;
314 let block_size = CudaKernelProvider::SORT_BLOCK_SIZE;
315 let grid_size = len.div_ceil(block_size).max(1);
316 let hist = memory.alloc::<u32>((grid_size as usize) * 16)?;
317 let prefix = memory.alloc::<u32>(16)?;
318 Ok(Self {
319 keys_b,
320 values_b,
321 hist,
322 prefix,
323 ranks,
324 len,
325 })
326 }
327
328 pub fn ensure_capacity(&mut self, provider: &CudaKernelProvider, n: u32) -> Result<()> {
329 if n <= self.len {
330 return Ok(());
331 }
332 *self = Self::new(provider, n)?;
333 Ok(())
334 }
335}
336
337pub const JOIN_MODULE: &str = "xlog_join";
339pub const DEDUP_MODULE: &str = "xlog_dedup";
340pub const GROUPBY_MODULE: &str = "xlog_groupby";
341pub const SCAN_MODULE: &str = "xlog_scan";
342pub const SORT_MODULE: &str = "xlog_sort";
343pub const FILTER_MODULE: &str = "xlog_filter";
344pub const SET_OPS_MODULE: &str = "xlog_set_ops";
345pub const PACK_MODULE: &str = "xlog_pack";
346pub const CIRCUIT_MODULE: &str = "xlog_circuit";
347pub const MC_SAMPLE_MODULE: &str = "xlog_mc_sample";
348pub const MC_EVAL_MODULE: &str = "xlog_mc_eval";
349pub const MC_RESIDENT_MODULE: &str = "xlog_mc_resident";
350pub const ARITH_MODULE: &str = "xlog_arith";
351pub const SAT_MODULE: &str = "xlog_sat";
352pub const D4_MODULE: &str = "xlog_d4";
353pub const NEURAL_MODULE: &str = "xlog_neural";
354pub const PIR_MODULE: &str = "xlog_pir";
355pub const CNF_MODULE: &str = "xlog_cnf";
356pub const CACHE_MODULE: &str = "xlog_cache";
357pub const WEIGHTS_MODULE: &str = "xlog_weights";
358pub const ILP_MODULE: &str = "xlog_ilp";
359pub const ILP_CREDIT_MODULE: &str = "xlog_ilp_credit";
360pub const ILP_EXACT_MODULE: &str = "xlog_ilp_exact";
361pub const ILP_EXACT_NARY_MODULE: &str = "xlog_ilp_exact_nary";
362pub const EPISTEMIC_MODULE: &str = "xlog_epistemic";
363pub const WCOJ_MODULE: &str = "xlog_wcoj";
364pub const JOINT_SOLVE_MODULE: &str = "xlog_joint_solve";
365
366pub mod wcoj_kernels {
368 pub const WCOJ_BUILD_METADATA_MARK_BOUNDARIES_U32: &str =
369 "wcoj_build_metadata_mark_boundaries_u32";
370 pub const WCOJ_BUILD_METADATA_MARK_BOUNDARIES_U64: &str =
371 "wcoj_build_metadata_mark_boundaries_u64";
372 pub const WCOJ_BUILD_METADATA_SCATTER_U32: &str = "wcoj_build_metadata_scatter_u32";
373 pub const WCOJ_BUILD_METADATA_SCATTER_U64: &str = "wcoj_build_metadata_scatter_u64";
374 pub const WCOJ_TRIANGLE_BUILD_HG_WORK_PLAN_U32: &str = "wcoj_triangle_build_hg_work_plan_u32";
375 pub const WCOJ_TRIANGLE_COUNT_HG_U32: &str = "wcoj_triangle_count_hg_u32";
376 pub const WCOJ_TRIANGLE_GROUPBY_ROOT_COUNT_HG_U32: &str =
377 "wcoj_triangle_groupby_root_count_hg_u32";
378 pub const WCOJ_TRIANGLE_GROUPBY_ROOT_SUM_HG_U32: &str = "wcoj_triangle_groupby_root_sum_hg_u32";
379 pub const WCOJ_TRIANGLE_GROUPBY_ROOT_MIN_HG_U32: &str = "wcoj_triangle_groupby_root_min_hg_u32";
380 pub const WCOJ_TRIANGLE_GROUPBY_ROOT_MAX_HG_U32: &str = "wcoj_triangle_groupby_root_max_hg_u32";
381 pub const WCOJ_TRIANGLE_MATERIALIZE_HG_U32: &str = "wcoj_triangle_materialize_hg_u32";
382 pub const WCOJ_TRIANGLE_BUILD_HG_WORK_PLAN_U64: &str = "wcoj_triangle_build_hg_work_plan_u64";
383 pub const WCOJ_TRIANGLE_COUNT_HG_U64: &str = "wcoj_triangle_count_hg_u64";
384 pub const WCOJ_TRIANGLE_GROUPBY_ROOT_COUNT_HG_U64: &str =
385 "wcoj_triangle_groupby_root_count_hg_u64";
386 pub const WCOJ_TRIANGLE_GROUPBY_ROOT_SUM_HG_U64: &str = "wcoj_triangle_groupby_root_sum_hg_u64";
387 pub const WCOJ_TRIANGLE_GROUPBY_ROOT_MIN_HG_U64: &str = "wcoj_triangle_groupby_root_min_hg_u64";
388 pub const WCOJ_TRIANGLE_GROUPBY_ROOT_MAX_HG_U64: &str = "wcoj_triangle_groupby_root_max_hg_u64";
389 pub const WCOJ_GROUPBY_ROOT_SEGMENT_SUM_COUNTS_U32: &str =
390 "wcoj_groupby_root_segment_sum_counts_u32";
391 pub const WCOJ_GROUPBY_ROOT_SEGMENT_SUM_VALUES_U64: &str =
392 "wcoj_groupby_root_segment_sum_values_u64";
393 pub const WCOJ_GROUPBY_ROOT_SEGMENT_MIN_VALUES_U64: &str =
394 "wcoj_groupby_root_segment_min_values_u64";
395 pub const WCOJ_GROUPBY_ROOT_SEGMENT_MAX_VALUES_U64: &str =
396 "wcoj_groupby_root_segment_max_values_u64";
397 pub const WCOJ_TRIANGLE_MATERIALIZE_HG_U64: &str = "wcoj_triangle_materialize_hg_u64";
398 pub const WCOJ_TRIANGLE_COUNT_HG_CACHED_U32: &str = "wcoj_triangle_count_hg_cached_u32";
399 pub const WCOJ_TRIANGLE_MATERIALIZE_HG_CACHED_U32: &str =
400 "wcoj_triangle_materialize_hg_cached_u32";
401 pub const WCOJ_SCAN_HG_BLOCK_COUNTS_U32: &str = "wcoj_scan_hg_block_counts_u32";
402 pub const WCOJ_COMPUTE_TOTAL: &str = "wcoj_compute_total";
403 pub const WCOJ_LAYOUT_CHECK_SORTED_UNIQUE_U32: &str = "wcoj_layout_check_sorted_unique_u32";
404 pub const WCOJ_LAYOUT_CHECK_SORTED_UNIQUE_U64: &str = "wcoj_layout_check_sorted_unique_u64";
405 pub const WCOJ_4CYCLE_BUILD_E2_WORK_PREFIX_U32: &str = "wcoj_4cycle_build_e2_work_prefix_u32";
406 pub const WCOJ_4CYCLE_BUILD_HG_WORK_PLAN_U32: &str = "wcoj_4cycle_build_hg_work_plan_u32";
407 pub const WCOJ_4CYCLE_COUNT_HG_U32: &str = "wcoj_4cycle_count_hg_u32";
408 pub const WCOJ_4CYCLE_GROUPBY_ROOT_COUNT_HG_U32: &str = "wcoj_4cycle_groupby_root_count_hg_u32";
409 pub const WCOJ_4CYCLE_GROUPBY_ROOT_SUM_HG_U32: &str = "wcoj_4cycle_groupby_root_sum_hg_u32";
410 pub const WCOJ_4CYCLE_GROUPBY_ROOT_MIN_HG_U32: &str = "wcoj_4cycle_groupby_root_min_hg_u32";
411 pub const WCOJ_4CYCLE_GROUPBY_ROOT_MAX_HG_U32: &str = "wcoj_4cycle_groupby_root_max_hg_u32";
412 pub const WCOJ_4CYCLE_MATERIALIZE_HG_U32: &str = "wcoj_4cycle_materialize_hg_u32";
413 pub const WCOJ_4CYCLE_BUILD_E2_WORK_PREFIX_U64: &str = "wcoj_4cycle_build_e2_work_prefix_u64";
414 pub const WCOJ_4CYCLE_BUILD_HG_WORK_PLAN_U64: &str = "wcoj_4cycle_build_hg_work_plan_u64";
415 pub const WCOJ_4CYCLE_COUNT_HG_U64: &str = "wcoj_4cycle_count_hg_u64";
416 pub const WCOJ_4CYCLE_GROUPBY_ROOT_COUNT_HG_U64: &str = "wcoj_4cycle_groupby_root_count_hg_u64";
417 pub const WCOJ_4CYCLE_MATERIALIZE_HG_U64: &str = "wcoj_4cycle_materialize_hg_u64";
418 pub const WCOJ_CLIQUE5_COUNT_HG_U32: &str = "wcoj_clique5_count_hg_u32";
420 pub const WCOJ_CLIQUE5_MATERIALIZE_HG_U32: &str = "wcoj_clique5_materialize_hg_u32";
421 pub const WCOJ_CLIQUE5_COUNT_HG_U64: &str = "wcoj_clique5_count_hg_u64";
422 pub const WCOJ_CLIQUE5_MATERIALIZE_HG_U64: &str = "wcoj_clique5_materialize_hg_u64";
423 pub const WCOJ_CLIQUE6_COUNT_HG_U32: &str = "wcoj_clique6_count_hg_u32";
424 pub const WCOJ_CLIQUE6_MATERIALIZE_HG_U32: &str = "wcoj_clique6_materialize_hg_u32";
425 pub const WCOJ_CLIQUE6_COUNT_HG_U64: &str = "wcoj_clique6_count_hg_u64";
426 pub const WCOJ_CLIQUE6_MATERIALIZE_HG_U64: &str = "wcoj_clique6_materialize_hg_u64";
427 pub const WCOJ_CLIQUE7_COUNT_HG_U32: &str = "wcoj_clique7_count_hg_u32";
428 pub const WCOJ_CLIQUE7_MATERIALIZE_HG_U32: &str = "wcoj_clique7_materialize_hg_u32";
429 pub const WCOJ_CLIQUE7_COUNT_HG_U64: &str = "wcoj_clique7_count_hg_u64";
430 pub const WCOJ_CLIQUE7_MATERIALIZE_HG_U64: &str = "wcoj_clique7_materialize_hg_u64";
431 pub const WCOJ_CLIQUE8_COUNT_HG_U32: &str = "wcoj_clique8_count_hg_u32";
432 pub const WCOJ_CLIQUE8_MATERIALIZE_HG_U32: &str = "wcoj_clique8_materialize_hg_u32";
433 pub const WCOJ_CLIQUE8_COUNT_HG_U64: &str = "wcoj_clique8_count_hg_u64";
434 pub const WCOJ_CLIQUE8_MATERIALIZE_HG_U64: &str = "wcoj_clique8_materialize_hg_u64";
435 pub const WCOJ_CLIQUE5_GROUPBY_ROOT_COUNT_HG_U32: &str =
436 "wcoj_clique5_groupby_root_count_hg_u32";
437 pub const WCOJ_CLIQUE6_GROUPBY_ROOT_COUNT_HG_U32: &str =
438 "wcoj_clique6_groupby_root_count_hg_u32";
439 pub const FJ_EXPAND_WORK_PREFIX_U32: &str = "fj_expand_work_prefix_u32";
443 pub const FJ_EXPAND_COUNT_U32: &str = "fj_expand_count_u32";
444 pub const FJ_EXPAND_EMIT_U32: &str = "fj_expand_emit_u32";
445 pub const FJ_PROBE_REFINE_U32: &str = "fj_probe_refine_u32";
446 pub const FJ_EXPAND_COUNT_U64: &str = "fj_expand_count_u64";
447 pub const FJ_EXPAND_EMIT_U64: &str = "fj_expand_emit_u64";
448 pub const FJ_PROBE_REFINE_U64: &str = "fj_probe_refine_u64";
449 pub const FJ_COUNT_MULTIPLICITY: &str = "fj_count_multiplicity";
450 pub const FJ_DELTA_RANGE_U32: &str = "fj_delta_range_u32";
452 pub const FJ_DELTA_MARK_U32: &str = "fj_delta_mark_u32";
453 pub const FJ_DELTA_SUBTRACT_U32: &str = "fj_delta_subtract_u32";
454 pub const FJ_DELTA_POPCOUNT: &str = "fj_delta_popcount";
455 pub const FJ_DELTA_EMIT_U32: &str = "fj_delta_emit_u32";
456 pub const FJ_DELTA_MAX_U32: &str = "fj_delta_max_u32";
457 pub const FJ_DELTA_SPARSE_ESTIMATE: &str = "fj_delta_sparse_estimate";
458 pub const FJ_DELTA_SPARSE_LOAD_R: &str = "fj_delta_sparse_load_r";
459 pub const FJ_DELTA_SPARSE_INSERT_CANDIDATES: &str = "fj_delta_sparse_insert_candidates";
460 pub const FJ_DELTA_SPARSE_MARK: &str = "fj_delta_sparse_mark";
461 pub const FJ_DELTA_SPARSE_EMIT: &str = "fj_delta_sparse_emit";
462}
463
464pub mod mc_sample_kernels {
466 pub const MC_SAMPLE_BERNOULLI: &str = "mc_sample_bernoulli";
467}
468
469pub mod mc_eval_kernels {
471 pub const MC_EVAL_MASK_VAR: &str = "mc_eval_mask_var";
472 pub const MC_EVAL_MASK_AD: &str = "mc_eval_mask_ad_choice";
473 pub const MC_EVAL_QUERY_EVIDENCE_TRUTH: &str = "mc_eval_query_evidence_truth";
474 pub const MC_EVAL_ACCUMULATE_COUNTS: &str = "mc_accumulate_counts";
475}
476
477pub mod mc_resident_kernels {
479 pub const MC_RESIDENT_ENGINE: &str = "mc_resident_engine";
482}
483
484pub mod arith_kernels {
486 pub const ARITH_BINARY_I64: &str = "arith_binary_i64";
487 pub const ARITH_BINARY_I32: &str = "arith_binary_i32";
488 pub const ARITH_BINARY_U64: &str = "arith_binary_u64";
489 pub const ARITH_BINARY_U32: &str = "arith_binary_u32";
490 pub const ARITH_BINARY_F64: &str = "arith_binary_f64";
491 pub const ARITH_BINARY_F32: &str = "arith_binary_f32";
492 pub const ARITH_ABS_I64: &str = "arith_abs_i64";
493 pub const ARITH_ABS_I32: &str = "arith_abs_i32";
494 pub const ARITH_ABS_F64: &str = "arith_abs_f64";
495 pub const ARITH_ABS_F32: &str = "arith_abs_f32";
496 pub const ARITH_POW_F64: &str = "arith_pow_f64";
497 pub const ARITH_CAST: &str = "arith_cast";
498 pub const ARITH_FILL_CONST_U32: &str = "arith_fill_const_u32";
499 pub const ARITH_FILL_CONST_U64: &str = "arith_fill_const_u64";
500 pub const ARITH_FILL_CONST_I64: &str = "arith_fill_const_i64";
501 pub const ARITH_FILL_CONST_I32: &str = "arith_fill_const_i32";
502 pub const ARITH_FILL_CONST_F64: &str = "arith_fill_const_f64";
503 pub const ARITH_FILL_CONST_F32: &str = "arith_fill_const_f32";
504 pub const ARITH_FILL_CONST_U8: &str = "arith_fill_const_u8";
505 pub const ARITH_SELECT_I64: &str = "arith_select_i64";
507 pub const ARITH_SELECT_I32: &str = "arith_select_i32";
508 pub const ARITH_SELECT_U64: &str = "arith_select_u64";
509 pub const ARITH_SELECT_U32: &str = "arith_select_u32";
510 pub const ARITH_SELECT_F64: &str = "arith_select_f64";
511 pub const ARITH_SELECT_F32: &str = "arith_select_f32";
512}
513
514pub mod epistemic_kernels {
516 pub const EPISTEMIC_GENERATE_CANDIDATE_ASSUMPTIONS_U8: &str =
518 "epistemic_generate_candidate_assumptions_u8";
519 pub const EPISTEMIC_PROPAGATE_CANDIDATES_U8: &str = "epistemic_propagate_candidates_u8";
521 pub const EPISTEMIC_VALIDATE_CANDIDATE_BITS_U8: &str = "epistemic_validate_candidate_bits_u8";
523 pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_U8: &str =
525 "epistemic_populate_model_membership_u8";
526 pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_FROM_TUPLE_SOURCE_U8: &str =
528 "epistemic_populate_model_membership_from_tuple_source_u8";
529 pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_FROM_TUPLE_SOURCE_ARITY1_U8: &str =
531 "epistemic_populate_model_membership_from_tuple_source_arity1_u8";
532 pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_FROM_TUPLE_SOURCE_ARITY2_U8: &str =
534 "epistemic_populate_model_membership_from_tuple_source_arity2_u8";
535 pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_FROM_TUPLE_SOURCE_ARITY3_U8: &str =
537 "epistemic_populate_model_membership_from_tuple_source_arity3_u8";
538 pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_FROM_TUPLE_SOURCE_ARITY_N_U8: &str =
540 "epistemic_populate_model_membership_from_tuple_source_arity_n_u8";
541 pub const EPISTEMIC_VALIDATE_WORLD_VIEWS_U8: &str = "epistemic_validate_world_views_u8";
543 pub const EPISTEMIC_VALIDATE_CONSTRAINTS_U8: &str = "epistemic_validate_constraints_u8";
545 pub const EPISTEMIC_MATERIALIZE_ACCEPTED_CANDIDATES_U8: &str =
547 "epistemic_materialize_accepted_candidates_u8";
548
549 pub const EPISTEMIC_MATERIALIZE_FINAL_RESULT_FLAGS_U8: &str =
551 "epistemic_materialize_final_result_flags_u8";
552 pub const EPISTEMIC_MATERIALIZE_FINAL_TUPLE_COLUMN_U8: &str =
554 "epistemic_materialize_final_tuple_column_u8";
555 pub const EPISTEMIC_BUILD_FINAL_TUPLE_ROW_MAP_U8: &str =
557 "epistemic_build_final_tuple_row_map_u8";
558 pub const EPISTEMIC_CLOSE_FINAL_TUPLE_REJECTIONS_U8: &str =
560 "epistemic_close_final_tuple_rejections_u8";
561}
562
563pub mod neural_kernels {
565 pub const NEURAL_FILL_AD_CHAIN_F32: &str = "neural_fill_ad_chain_f32";
566 pub const NEURAL_SCATTER_AD_CHAIN_GRADS_F32: &str = "neural_scatter_ad_chain_grads_f32";
567}
568
569pub mod ilp_kernels {
571 pub const EXTRACT_NONZERO_INDICES: &str = "extract_nonzero_indices";
572 pub const ILP_MARK_SELECTED_IDS_U32: &str = "ilp_mark_selected_ids_u32";
573 pub const ILP_MARK_SELECTED_IDS_I32: &str = "ilp_mark_selected_ids_i32";
574 pub const ILP_MARK_SELECTED_IDS_I64: &str = "ilp_mark_selected_ids_i64";
575 pub const ILP_MARK_SELECTED_IDS_U64: &str = "ilp_mark_selected_ids_u64";
576 pub const ILP_VALIDATE_SELECTED_IDS_U32: &str = "ilp_validate_selected_ids_u32";
577 pub const ILP_VALIDATE_SELECTED_IDS_I32: &str = "ilp_validate_selected_ids_i32";
578 pub const ILP_VALIDATE_SELECTED_IDS_I64: &str = "ilp_validate_selected_ids_i64";
579 pub const ILP_VALIDATE_SELECTED_IDS_U64: &str = "ilp_validate_selected_ids_u64";
580 pub const ILP_BROADCAST_CANDIDATE_FLAG: &str = "ilp_broadcast_candidate_flag";
581 pub const ILP_COO_FILL_FROM_MASK: &str = "ilp_coo_fill_from_mask";
582 pub const ILP_CSR_HISTOGRAM: &str = "ilp_csr_histogram";
583 pub const ILP_REDUCE_SUM_F32: &str = "ilp_reduce_sum_f32";
584 pub const ILP_REDUCE_SUM_F64: &str = "ilp_reduce_sum_f64";
585}
586
587pub mod ilp_credit_kernels {
589 pub const ILP_COO_FILL: &str = "ilp_coo_fill";
590 pub const ILP_CREDIT_FORWARD_F32: &str = "ilp_credit_forward_f32";
591 pub const ILP_CREDIT_FORWARD_F64: &str = "ilp_credit_forward_f64";
592 pub const ILP_CREDIT_BACKWARD_F32: &str = "ilp_credit_backward_f32";
593 pub const ILP_CREDIT_BACKWARD_F64: &str = "ilp_credit_backward_f64";
594}
595
596pub mod ilp_exact_kernels {
598 pub const ILP_EXACT_SCORE: &str = "ilp_exact_score";
599 pub const ILP_EXACT_SCORE_U32: &str = "ilp_exact_score_u32";
600 pub const ILP_EXACT_SCORE_CHAIN_SMEM: &str = "ilp_exact_score_chain_smem";
601 pub const ILP_EXACT_SCORE_CHAIN_SMEM_U32: &str = "ilp_exact_score_chain_smem_u32";
602 pub const ILP_EXACT_SELECT_TOPK: &str = "ilp_exact_select_topk";
603}
604
605pub mod ilp_exact_nary_kernels {
607 pub const ILP_EXACT_NARY_SCORE: &str = "ilp_exact_nary_score";
608}
609
610pub mod pir_kernels {
612 pub const PIR_PACK_KEYS: &str = "pir_pack_keys";
613 pub const PIR_HASH_KEYS: &str = "pir_hash_keys";
614 pub const PIR_MARK_UNIQUE: &str = "pir_mark_unique";
615 pub const PIR_FIND_EXISTING: &str = "pir_find_existing";
616 pub const PIR_MARK_NEW_GROUPS: &str = "pir_mark_new_groups";
617 pub const PIR_BUILD_GROUP_IDS: &str = "pir_build_group_ids";
618 pub const PIR_FILL_CHILD_PARENTS: &str = "pir_fill_child_parents";
619 pub const PIR_MARK_UNIQUE_PAIRS: &str = "pir_mark_unique_pairs";
620 pub const PIR_COMPACT_PAIRS: &str = "pir_compact_pairs";
621 pub const PIR_COUNT_CHILDREN: &str = "pir_count_children";
622 pub const PIR_WRITE_CHILD_OFFSETS: &str = "pir_write_child_offsets";
623 pub const PIR_GATHER_CHILDREN: &str = "pir_gather_children";
624 pub const PIR_BUILD_GRAPH_CHILD_COUNTS: &str = "pir_build_graph_child_counts";
625 pub const PIR_SUM_COUNTS: &str = "pir_sum_counts";
626 pub const PIR_EMIT_NODES_AND_IDS: &str = "pir_emit_nodes_and_ids";
627 pub const PIR_UPDATE_COUNTS: &str = "pir_update_counts";
628}
629
630pub mod cnf_kernels {
632 pub const CNF_REACHABILITY_INIT: &str = "cnf_reachability_init";
633 pub const CNF_REACHABILITY_BFS: &str = "cnf_reachability_bfs";
634 pub const CNF_MARK_LEAF_CHOICE: &str = "cnf_mark_leaf_choice";
635 pub const CNF_ASSIGN_LEAF_VAR: &str = "cnf_assign_leaf_var";
636 pub const CNF_ASSIGN_CHOICE_VAR: &str = "cnf_assign_choice_var";
637 pub const CNF_MARK_NODE_VARS: &str = "cnf_mark_node_vars";
638 pub const CNF_COUNT_CLAUSES: &str = "cnf_count_clauses";
639 pub const CNF_CAPTURE_LAST_COUNTS: &str = "cnf_capture_last_counts";
640 pub const CNF_COMPUTE_LEAF_CHOICE_TOTALS: &str = "cnf_compute_leaf_choice_totals";
641 pub const CNF_COMPUTE_TOTALS: &str = "cnf_compute_totals";
642 pub const CNF_ASSIGN_NODE_VAR: &str = "cnf_assign_node_var";
643 pub const CNF_EMIT_CLAUSES: &str = "cnf_emit_clauses";
644 pub const CNF_SET_CLAUSE_END: &str = "cnf_set_clause_end";
645}
646
647pub mod weights_kernels {
649 pub const WEIGHTS_FILL_LEAF: &str = "weights_fill_leaf";
650 pub const WEIGHTS_FILL_CHOICE: &str = "weights_fill_choice";
651 pub const WEIGHTS_COUNT_LIFT_EXACT: &str = "weights_count_lift_exact";
652 pub const WEIGHTS_SET_EVIDENCE_FROM_NODES: &str = "weights_set_evidence_from_nodes";
653 pub const WEIGHTS_APPLY_EVIDENCE: &str = "weights_apply_evidence";
654 pub const WEIGHTS_MAP_NODES_TO_VARS: &str = "weights_map_nodes_to_vars";
655 pub const WEIGHTS_FORCE_VAR_FALSE: &str = "weights_force_var_false";
656 pub const WEIGHTS_RESTORE_VAR_FALSE: &str = "weights_restore_var_false";
657 pub const WEIGHTS_FORCE_VAR_TRUE: &str = "weights_force_var_true";
658 pub const WEIGHTS_RESTORE_VAR_TRUE: &str = "weights_restore_var_true";
659 pub const WEIGHTS_COPY_SLOT_TO_BATCH: &str = "weights_copy_slot_to_batch";
660 pub const WEIGHTS_APPLY_QUERY_VARS: &str = "weights_apply_query_vars";
661 pub const WEIGHTS_RESTORE_QUERY_VARS: &str = "weights_restore_query_vars";
662 pub const WEIGHTS_APPLY_QUERY_VARS_FALSE_BATCHED: &str =
663 "weights_apply_query_vars_false_batched";
664 pub const WEIGHTS_RESTORE_QUERY_VARS_FALSE_BATCHED: &str =
665 "weights_restore_query_vars_false_batched";
666 pub const WEIGHTS_APPLY_QUERY_VARS_TRUE_BATCHED: &str = "weights_apply_query_vars_true_batched";
667 pub const WEIGHTS_RESTORE_QUERY_VARS_TRUE_BATCHED: &str =
668 "weights_restore_query_vars_true_batched";
669}
670
671pub mod d4_kernels {
674 pub const D4_VALIDATE_CNF: &str = "d4_validate_cnf";
675 pub const D4_LEVELIZE_COUNTS: &str = "d4_levelize_counts";
676 pub const D4_LEVELIZE_EMIT: &str = "d4_levelize_emit";
677 pub const D4_FRONTIER_PREPARE: &str = "d4_frontier_prepare";
679 pub const D4_FRONTIER_EXPAND: &str = "d4_frontier_expand";
680 pub const D4_COMPILE_COUNT: &str = "d4_compile_count";
682 pub const D4_COMPILE_EMIT: &str = "d4_compile_emit";
683 pub const D4_CAPTURE_EMIT_META: &str = "d4_capture_emit_meta";
684 pub const D4_SUPPORT_LEVEL: &str = "d4_support_level";
686 pub const D4_SUPPORT_SET_ROOT_BITS: &str = "d4_support_set_root_bits";
687 pub const D4_SMOOTH_COUNT: &str = "d4_smooth_count";
688 pub const D4_SMOOTH_WRAPPER_COUNTS: &str = "d4_smooth_wrapper_counts";
689 pub const D4_SMOOTH_WRAPPER_EDGE_COUNTS_OR: &str = "d4_smooth_wrapper_edge_counts_or";
690 pub const D4_SMOOTH_WRAPPER_EDGE_COUNTS_DEC: &str = "d4_smooth_wrapper_edge_counts_dec";
691 pub const D4_SMOOTH_INIT_NODES: &str = "d4_smooth_init_nodes";
692 pub const D4_SMOOTH_EMIT_LEVEL: &str = "d4_smooth_emit_level";
693 pub const D4_SMOOTH_CHECK_EDGE_CAP: &str = "d4_smooth_check_edge_cap";
694 pub const D4_MARK_VARS_IN_CLAUSES: &str = "d4_mark_vars_in_clauses";
696 pub const D4_MARK_VARS_IN_CIRCUIT: &str = "d4_mark_vars_in_circuit";
697 pub const D4_BUILD_FREE_VAR_MASK: &str = "d4_build_free_var_mask";
698 pub const D4_ASSERT_U32_EQ: &str = "d4_assert_u32_eq";
700 pub const D4_ASSERT_BITSET_VAR: &str = "d4_assert_bitset_var";
701 pub const D4_ASSERT_LEAF_ROOT_AND_DEGREE: &str = "d4_assert_leaf_root_and_degree";
702}
703
704pub mod join_kernels {
706 pub const HASH_JOIN_BUILD: &str = "hash_join_build";
707 pub const HASH_JOIN_PROBE: &str = "hash_join_probe";
708 pub const COMPUTE_COMPOSITE_HASH: &str = "compute_composite_hash";
710 pub const HASH_JOIN_BUCKET_COUNT_V2: &str = "hash_join_bucket_count_v2";
711 pub const HASH_JOIN_SCATTER_V2: &str = "hash_join_scatter_v2";
712 pub const HASH_JOIN_PROBE_V2: &str = "hash_join_probe_v2";
713 pub const HASH_JOIN_PROBE_V2_COUNT_PER_ROW: &str = "hash_join_probe_v2_count_per_row";
714 pub const HASH_JOIN_PROBE_V2_MATERIALIZE: &str = "hash_join_probe_v2_materialize";
715 pub const HASH_JOIN_TOTAL_FROM_SCAN: &str = "hash_join_total_from_scan";
716 pub const HASH_JOIN_CSM_UNMATCHED_MASK: &str = "hash_join_csm_unmatched_mask";
717 pub const HASH_JOIN_SEMI: &str = "hash_join_semi";
718 pub const HASH_JOIN_ANTI: &str = "hash_join_anti";
719 pub const INIT_HASH_TABLE: &str = "init_hash_table";
720 pub const NESTED_LOOP_JOIN_INNER_U32_1KEY_PAIRS: &str = "nested_loop_join_inner_u32_1key_pairs";
726 pub const SORT_MERGE_JOIN_INNER_U32_1KEY_PAIRS: &str = "sort_merge_join_inner_u32_1key_pairs";
734}
735
736pub mod dedup_kernels {
738 pub const MARK_DUPLICATES: &str = "mark_duplicates";
739 pub const MARK_UNIQUE_COLUMNAR: &str = "mark_unique_columnar";
740 pub const MARK_UNIQUE_AND_SCAN_COLUMNAR: &str = "mark_unique_and_scan_columnar";
741 pub const COMPACT_ROWS: &str = "compact_rows";
742 pub const MARK_UNIQUE_FULL_ROW_BYTEWISE: &str = "mark_unique_full_row_bytewise";
743 pub const MARK_DIFF_FULL_ROW_TYPED_SORTED: &str = "mark_diff_full_row_typed_sorted";
744 pub const SMALL_SORT_FULL_ROW_INDICES_TYPED: &str = "small_sort_full_row_indices_typed";
745}
746
747pub mod groupby_kernels {
749 pub const DETECT_GROUP_BOUNDARIES: &str = "detect_group_boundaries";
750 pub const DETECT_BOUNDARIES: &str = "detect_boundaries";
751 pub const EXTRACT_GROUP_KEYS: &str = "extract_group_keys";
752 pub const GROUP_IDS_FROM_BOUNDARIES: &str = "group_ids_from_boundaries";
753 pub const GROUP_START_INDICES: &str = "group_start_indices";
754 pub const CAPTURE_NUM_GROUPS: &str = "capture_num_groups";
755 pub const GROUPBY_COUNT: &str = "groupby_count";
756 pub const GROUPBY_SUM: &str = "groupby_sum";
757 pub const GROUPBY_SUM_U64: &str = "groupby_sum_u64";
758 pub const GROUPBY_MIN: &str = "groupby_min";
759 pub const GROUPBY_MIN_U64: &str = "groupby_min_u64";
760 pub const GROUPBY_MAX: &str = "groupby_max";
761 pub const GROUPBY_MAX_U64: &str = "groupby_max_u64";
762 pub const GROUPBY_LOGSUMEXP_MAX: &str = "groupby_logsumexp_max";
763 pub const GROUPBY_LOGSUMEXP_SUMEXP: &str = "groupby_logsumexp_sumexp";
764 pub const GROUPBY_LOGSUMEXP_FINAL: &str = "groupby_logsumexp_final";
765}
766
767pub mod scan_kernels {
769 pub const BLOCK_INCLUSIVE_SCAN: &str = "block_inclusive_scan";
770 pub const ADD_BLOCK_OFFSETS: &str = "add_block_offsets";
771 pub const EXCLUSIVE_SCAN_MASK: &str = "exclusive_scan_mask";
772 pub const COUNT_MASK: &str = "count_mask";
773 pub const MULTIBLOCK_SCAN_PHASE1: &str = "multiblock_scan_phase1";
775 pub const MULTIBLOCK_SCAN_U32_PHASE1: &str = "multiblock_scan_u32_phase1";
776 pub const MULTIBLOCK_SCAN_PHASE2: &str = "multiblock_scan_phase2";
777 pub const MULTIBLOCK_SCAN_PHASE3: &str = "multiblock_scan_phase3";
778}
779
780pub mod sort_kernels {
782 pub const RADIX_HISTOGRAM: &str = "radix_histogram";
783 pub const RADIX_SCATTER: &str = "radix_scatter";
784 pub const COMPUTE_RANKS: &str = "compute_ranks";
785 pub const RADIX_SCATTER_STABLE: &str = "radix_scatter_stable";
786 pub const COMPUTE_DIGIT_PREFIX_SUMS: &str = "compute_digit_prefix_sums";
787 pub const INIT_INDICES: &str = "init_indices";
788 pub const APPLY_PERMUTATION_U32: &str = "apply_permutation_u32";
789 pub const APPLY_PERMUTATION_BYTES: &str = "apply_permutation_bytes";
790
791 pub const GATHER_KEYS_I32_ORDERED_U32: &str = "gather_keys_i32_ordered_u32";
792 pub const GATHER_KEYS_F32_ORDERED_U32: &str = "gather_keys_f32_ordered_u32";
793 pub const GATHER_KEYS_BOOL_ORDERED_U32: &str = "gather_keys_bool_ordered_u32";
794
795 pub const GATHER_KEYS_U64_LO_U32: &str = "gather_keys_u64_lo_u32";
796 pub const GATHER_KEYS_U64_HI_U32: &str = "gather_keys_u64_hi_u32";
797
798 pub const GATHER_KEYS_I64_LO_U32: &str = "gather_keys_i64_lo_u32";
799 pub const GATHER_KEYS_I64_HI_U32: &str = "gather_keys_i64_hi_u32";
800
801 pub const GATHER_KEYS_F64_LO_U32: &str = "gather_keys_f64_lo_u32";
802 pub const GATHER_KEYS_F64_HI_U32: &str = "gather_keys_f64_hi_u32";
803 pub const CHECK_ASCENDING_SORTED_U32: &str = "check_ascending_sorted_u32";
811}
812
813pub mod filter_kernels {
815 pub const FILTER_COMPARE_U32: &str = "filter_compare_u32";
816 pub const FILTER_COMPARE_I64: &str = "filter_compare_i64";
817 pub const FILTER_COMPARE_F64: &str = "filter_compare_f64";
818 pub const FILTER_COMPARE_I32: &str = "filter_compare_i32";
819 pub const FILTER_COMPARE_U64: &str = "filter_compare_u64";
820 pub const FILTER_COMPARE_F32: &str = "filter_compare_f32";
821 pub const FILTER_COMPARE_U8: &str = "filter_compare_u8";
822 pub const FILTER_COMPARE_U32_SCAN_PHASE1: &str = "filter_compare_u32_scan_phase1";
823 pub const FILTER_COMPARE_F64_SCAN_PHASE1: &str = "filter_compare_f64_scan_phase1";
824 pub const FILTER_COMPARE_F32_SCAN_PHASE1: &str = "filter_compare_f32_scan_phase1";
825 pub const FILTER_COMPARE_U32_COL: &str = "filter_compare_u32_col";
826 pub const FILTER_COMPARE_I32_COL: &str = "filter_compare_i32_col";
827 pub const FILTER_COMPARE_I64_COL: &str = "filter_compare_i64_col";
828 pub const FILTER_COMPARE_U64_COL: &str = "filter_compare_u64_col";
829 pub const FILTER_COMPARE_F32_COL: &str = "filter_compare_f32_col";
830 pub const FILTER_COMPARE_F64_COL: &str = "filter_compare_f64_col";
831 pub const FILTER_COMPARE_U8_COL: &str = "filter_compare_u8_col";
832 pub const FILL_U32_IOTA: &str = "fill_u32_iota";
833 pub const FILL_U32_CONST: &str = "fill_u32_const";
834 pub const MARK_RANDOM_VARS: &str = "mark_random_vars";
835 pub const RANDOM_VAR_TO_BIT_FROM_LIST: &str = "random_var_to_bit_from_list";
836 pub const CHECK_RANDOM_VAR_COUNT: &str = "check_random_var_count";
837 pub const COMPACT_U32_BY_MASK: &str = "compact_u32_by_mask";
838 pub const COMPACT_I64_BY_MASK: &str = "compact_i64_by_mask";
839 pub const COMPACT_F64_BY_MASK: &str = "compact_f64_by_mask";
840 pub const COMPACT_BYTES_BY_MASK: &str = "compact_bytes_by_mask";
841 pub const CAPTURE_COMPACT_COUNT: &str = "capture_compact_count";
842 pub const MASK_CLAMP_ROWS: &str = "mask_clamp_rows";
843 pub const MASK_AND: &str = "mask_and";
844 pub const MASK_OR: &str = "mask_or";
845 pub const MASK_NOT: &str = "mask_not";
846}
847
848pub mod set_ops_kernels {
850 pub const CONCAT_U32: &str = "concat_u32";
851 pub const CONCAT_BYTES: &str = "concat_bytes";
852 pub const SORTED_DIFF_MARK: &str = "sorted_diff_mark";
853}
854
855pub mod pack_kernels {
857 pub const PACK_KEYS: &str = "pack_keys";
859 pub const HASH_PACKED_KEYS: &str = "hash_packed_keys";
861 pub const PACK_AND_HASH_KEYS: &str = "pack_and_hash_keys";
863 pub const PACK_AND_HASH_KEYS_GENERIC: &str = "pack_and_hash_keys_generic";
865 pub const PACK_KEYS_ALIGNED: &str = "pack_keys_aligned";
867 pub const UNPACK_COLUMN: &str = "unpack_column";
869 pub const UNPACK_COLUMN_COUNTED: &str = "unpack_column_counted";
871 pub const GATHER_PACKED_ROWS: &str = "gather_packed_rows";
873 pub const GATHER_PACKED_ROWS_COUNTED: &str = "gather_packed_rows_counted";
875 pub const SCATTER_PACKED_ROWS: &str = "scatter_packed_rows";
877 pub const COMPARE_PACKED_KEYS: &str = "compare_packed_keys";
879 pub const PACK_BOOLS_TO_BITMAP: &str = "pack_bools_to_bitmap";
881}
882
883pub mod circuit_kernels {
885 pub const XGCF_FORWARD_LEVEL: &str = "xgcf_forward_level";
886 pub const XGCF_BACKWARD_LEVEL_PROPAGATE: &str = "xgcf_backward_level_propagate";
887 pub const XGCF_BACKWARD_LEVEL_DECISION_GRAD: &str = "xgcf_backward_level_decision_grad";
888 pub const XGCF_BACKWARD_LEVEL_LIT_GRAD: &str = "xgcf_backward_level_lit_grad";
889 pub const XGCF_FREE_VAR_APPLY_GRAD: &str = "xgcf_free_var_apply_grad";
890 pub const XGCF_FREE_VAR_REDUCE_STAGE: &str = "xgcf_free_var_reduce_stage";
891 pub const XGCF_ADD_SCALAR: &str = "xgcf_add_scalar";
892 pub const XGCF_FORWARD_LEVEL_CACHED: &str = "xgcf_forward_level_cached";
893 pub const XGCF_EVAL_ALL_LEVELS_CACHED: &str = "xgcf_eval_all_levels_cached";
894 pub const XGCF_EVAL_ALL_LEVELS_CACHED_BATCHED: &str = "xgcf_eval_all_levels_cached_batched";
895 pub const XGCF_BACKWARD_LEVEL_PROPAGATE_CACHED: &str = "xgcf_backward_level_propagate_cached";
896 pub const XGCF_BACKWARD_LEVEL_DECISION_GRAD_CACHED: &str =
897 "xgcf_backward_level_decision_grad_cached";
898 pub const XGCF_BACKWARD_LEVEL_LIT_GRAD_CACHED: &str = "xgcf_backward_level_lit_grad_cached";
899 pub const XGCF_BACKWARD_ALL_LEVELS_CACHED: &str = "xgcf_backward_all_levels_cached";
900 pub const XGCF_BACKWARD_ALL_LEVELS_CACHED_BATCHED: &str =
901 "xgcf_backward_all_levels_cached_batched";
902 pub const XGCF_FREE_VAR_APPLY_GRAD_CACHED: &str = "xgcf_free_var_apply_grad_cached";
903 pub const XGCF_FREE_VAR_REDUCE_STAGE_CACHED: &str = "xgcf_free_var_reduce_stage_cached";
904 pub const XGCF_ADD_SCALAR_CACHED: &str = "xgcf_add_scalar_cached";
905 pub const XGCF_SET_ROOT_ADJ_CACHED_BATCHED: &str = "xgcf_set_root_adj_cached_batched";
906 pub const XGCF_COPY_ROOT_CACHED: &str = "xgcf_copy_root_cached";
907 pub const XGCF_COPY_ROOT_CACHED_META: &str = "xgcf_copy_root_cached_meta";
908 pub const XGCF_COPY_ROOT_CACHED_META_BATCHED: &str = "xgcf_copy_root_cached_meta_batched";
909}
910
911pub mod cache_kernels {
913 pub const CACHE_CNF_HASH: &str = "cache_cnf_hash";
914 pub const CACHE_LOOKUP_OR_INSERT: &str = "cache_lookup_or_insert";
915 pub const CACHE_EVICT_LRU: &str = "cache_evict_lru";
916 pub const CACHE_STORE_U8: &str = "cache_store_u8";
917 pub const CACHE_STORE_U32: &str = "cache_store_u32";
918 pub const CACHE_STORE_I32: &str = "cache_store_i32";
919 pub const CACHE_STORE_F64: &str = "cache_store_f64";
920 pub const CACHE_STORE_META: &str = "cache_store_meta";
921}
922
923pub mod sat_kernels {
925 pub const SAT_CDCL_SOLVE: &str = "sat_cdcl_solve";
926 pub const SAT_CHECK_MODEL: &str = "sat_check_model";
927 pub const SAT_PROOF_MARK_NEEDED: &str = "sat_proof_mark_needed";
928 pub const SAT_PROOF_CHECK: &str = "sat_proof_check";
929 pub const SAT_ASSERT_STATUS: &str = "sat_assert_status";
930 pub const SAT_ASSERT_OK: &str = "sat_assert_ok";
931 pub const SAT_XGCF_CNF_COUNTS: &str = "sat_xgcf_cnf_counts";
932 pub const SAT_XGCF_CNF_EMIT: &str = "sat_xgcf_cnf_emit";
933 pub const SAT_XGCF_CNF_CAPTURE_LAST_COUNTS: &str = "sat_xgcf_cnf_capture_last_counts";
934 pub const SAT_XGCF_CNF_COMPUTE_TOTALS: &str = "sat_xgcf_cnf_compute_totals";
935 pub const SAT_CNF_WRITE_TERMINATOR: &str = "sat_cnf_write_terminator";
936 pub const SAT_CNF_COPY_INTO: &str = "sat_cnf_copy_into";
937 pub const SAT_SHIFT_OFFSETS: &str = "sat_shift_offsets";
938 pub const SAT_XGCF_WRITE_ROOT_UNIT_CLAUSE: &str = "sat_xgcf_write_root_unit_clause";
939 pub const SAT_NOT_PHI_COUNTS: &str = "sat_not_phi_counts";
940 pub const SAT_EMIT_NOT_PHI: &str = "sat_emit_not_phi";
941}
942
943pub const DEFAULT_JOIN_MAX_OUTPUT: usize = 1_000_000;
946
947pub const NESTED_LOOP_TOTAL_THRESHOLD: u64 = 4_000_000;
968
969#[derive(Debug, Clone, Copy, PartialEq, Eq)]
971#[repr(u8)]
972pub enum CompareOp {
973 Eq = 0,
974 Ne = 1,
975 Lt = 2,
976 Le = 3,
977 Gt = 4,
978 Ge = 5,
979}
980
981#[derive(Debug, Clone, Copy, PartialEq, Eq)]
983pub enum JoinType {
984 Inner,
986 Semi,
988 Anti,
990 LeftOuter,
992}
993
994struct PackedKeyData {
996 hashes: crate::memory::TrackedCudaSlice<u64>,
998 packed_keys: crate::memory::TrackedCudaSlice<u8>,
1000 key_bytes: u32,
1002}
1003
1004struct JoinHashTableV2 {
1005 bucket_counts: crate::memory::TrackedCudaSlice<u32>,
1006 bucket_offsets: crate::memory::TrackedCudaSlice<u32>,
1007 bucket_entries: crate::memory::TrackedCudaSlice<u32>,
1008 bucket_entry_hashes: crate::memory::TrackedCudaSlice<u64>,
1009 bucket_mask: u32,
1010}
1011
1012pub struct HashTableU64 {
1014 pub bucket_counts: crate::memory::TrackedCudaSlice<u32>,
1015 pub bucket_offsets: crate::memory::TrackedCudaSlice<u32>,
1016 pub bucket_entries: crate::memory::TrackedCudaSlice<u32>,
1017 pub bucket_entry_hashes: crate::memory::TrackedCudaSlice<u64>,
1018 pub bucket_mask: u32,
1019}
1020
1021pub struct JoinIndexV2 {
1026 right_num_rows: u32,
1027 right_keys: Vec<usize>,
1028 key_bytes: u32,
1029 packed_keys: crate::memory::TrackedCudaSlice<u8>,
1030 table: JoinHashTableV2,
1031}
1032
1033impl JoinIndexV2 {
1034 pub fn right_keys(&self) -> &[usize] {
1036 &self.right_keys
1037 }
1038
1039 pub fn right_num_rows(&self) -> u32 {
1041 self.right_num_rows
1042 }
1043
1044 pub fn estimated_bytes(&self) -> u64 {
1046 let mut bytes = 0u64;
1047 bytes = bytes.saturating_add(self.packed_keys.len() as u64);
1048 bytes = bytes.saturating_add(self.table.bucket_counts.len() as u64 * 4);
1049 bytes = bytes.saturating_add(self.table.bucket_offsets.len() as u64 * 4);
1050 bytes = bytes.saturating_add(self.table.bucket_entries.len() as u64 * 4);
1051 bytes = bytes.saturating_add(self.table.bucket_entry_hashes.len() as u64 * 8);
1052 bytes
1053 }
1054}
1055
1056pub struct CudaKernelProvider {
1074 provider_identity: u64,
1076 device: Arc<CudaDevice>,
1078 memory: Arc<GpuMemoryManager>,
1080 transfer_tracker: HostTransferTracker,
1082 final_observation_transfer_tracker: HostTransferTracker,
1084 final_observation_pinned_receipts: AtomicU64,
1086 ptx_load_profile: Option<PtxLoadProfile>,
1088 d2h_transfer_count: AtomicU64,
1090 untracked_metadata_dtoh_count: AtomicU64,
1097 strict_deterministic_d2h: AtomicBool,
1103 deterministic_d2h_violations: AtomicU64,
1107 recorded_op_stream: OnceLock<crate::device_runtime::StreamId>,
1116 csm_invocations: AtomicU64,
1124 csm_cuda_graph_captures: AtomicU64,
1126 csm_cuda_graph_launches: AtomicU64,
1128 csm_cuda_graph_fallbacks: AtomicU64,
1130 csm_cuda_graph_cache_hits: AtomicU64,
1132 small_full_row_sort_invocations: AtomicU64,
1136 csm_cuda_graph_cache: Mutex<HashMap<CsmCudaGraphKey, CsmCudaGraphEntry>>,
1138 wcoj_layout_fast_path_hit_count: AtomicU64,
1145 wcoj_layout_sort_invocation_count: AtomicU64,
1150 kclique_metadata_build_count: AtomicU64,
1152 kclique_metadata_build_nanos: AtomicU64,
1155 wcoj_triangle_hg_dispatch_count: AtomicU64,
1158 #[cfg(feature = "wcoj-phase-timing")]
1165 last_triangle_phase_timing:
1166 std::sync::Mutex<Option<crate::wcoj_phase_timing::WcojTrianglePhaseTiming>>,
1167}
1168
1169#[derive(Default)]
1170struct HostTransferTracker {
1171 dtoh_bytes: AtomicU64,
1172 htod_bytes: AtomicU64,
1173 dtoh_calls: AtomicU64,
1174 htod_calls: AtomicU64,
1175 launch_metadata_htod_bytes: AtomicU64,
1176 launch_metadata_htod_calls: AtomicU64,
1177}
1178
1179#[derive(Clone, Copy, Debug, Default)]
1183pub struct ResidentHostTransferStats {
1184 pub tracked_htod_calls: u64,
1185 pub tracked_htod_bytes: u64,
1186 pub tracked_dtoh_calls: u64,
1187 pub tracked_dtoh_bytes: u64,
1188 pub provider_dtoh_calls: u64,
1189 pub untracked_metadata_dtoh_calls: u64,
1190 pub deterministic_dtoh_violations: u64,
1191 pub final_dtoh_calls: u64,
1192 pub final_dtoh_bytes: u64,
1193 pub final_pinned_receipts: u64,
1194}
1195
1196type ResidentTransferCell = std::rc::Rc<std::cell::Cell<ResidentHostTransferStats>>;
1197
1198thread_local! {
1199 static RESIDENT_TRANSFER_SCOPES: std::cell::RefCell<Vec<(u64, ResidentTransferCell)>> =
1200 const { std::cell::RefCell::new(Vec::new()) };
1201}
1202
1203pub struct ResidentHostTransferScope {
1206 provider_identity: u64,
1207 counters: ResidentTransferCell,
1208}
1209
1210impl ResidentHostTransferScope {
1211 pub fn snapshot(&self) -> ResidentHostTransferStats {
1212 self.counters.get()
1213 }
1214}
1215
1216impl Drop for ResidentHostTransferScope {
1217 fn drop(&mut self) {
1218 RESIDENT_TRANSFER_SCOPES.with_borrow_mut(|scopes| {
1219 if let Some(index) = scopes.iter().rposition(|(identity, counters)| {
1220 *identity == self.provider_identity && std::rc::Rc::ptr_eq(counters, &self.counters)
1221 }) {
1222 scopes.remove(index);
1223 }
1224 });
1225 }
1226}
1227
1228fn record_resident_transfer(
1229 provider_identity: u64,
1230 update: impl Fn(&mut ResidentHostTransferStats),
1231) {
1232 RESIDENT_TRANSFER_SCOPES.with_borrow(|scopes| {
1233 for (identity, counters) in scopes.iter() {
1234 if *identity == provider_identity {
1235 let mut stats = counters.get();
1236 update(&mut stats);
1237 counters.set(stats);
1238 }
1239 }
1240 });
1241}
1242
1243#[derive(Debug, Clone, Copy)]
1244pub struct HostTransferStats {
1245 pub dtoh_bytes: u64,
1246 pub htod_bytes: u64,
1247 pub dtoh_calls: u64,
1248 pub htod_calls: u64,
1249}
1250
1251#[derive(Debug, Clone, Copy, Default)]
1253pub struct FinalObservationTransferStats {
1254 pub dtoh_bytes: u64,
1256 pub dtoh_calls: u64,
1258 pub pinned_receipts: u64,
1260}
1261
1262#[derive(Debug, Clone, Copy, Default)]
1263pub struct HostLaunchMetadataTransferStats {
1264 pub htod_bytes: u64,
1265 pub htod_calls: u64,
1266}
1267
1268impl HostTransferTracker {
1269 fn record_dtoh(&self, bytes: u64) {
1270 self.dtoh_calls.fetch_add(1, Ordering::Relaxed);
1271 self.dtoh_bytes.fetch_add(bytes, Ordering::Relaxed);
1272 }
1273
1274 fn record_htod(&self, bytes: u64) {
1275 self.htod_calls.fetch_add(1, Ordering::Relaxed);
1276 self.htod_bytes.fetch_add(bytes, Ordering::Relaxed);
1277 }
1278
1279 fn record_htod_launch_metadata(&self, bytes: u64) {
1280 self.launch_metadata_htod_calls
1281 .fetch_add(1, Ordering::Relaxed);
1282 self.launch_metadata_htod_bytes
1283 .fetch_add(bytes, Ordering::Relaxed);
1284 }
1285
1286 fn snapshot(&self) -> HostTransferStats {
1287 HostTransferStats {
1288 dtoh_bytes: self.dtoh_bytes.load(Ordering::Relaxed),
1289 htod_bytes: self.htod_bytes.load(Ordering::Relaxed),
1290 dtoh_calls: self.dtoh_calls.load(Ordering::Relaxed),
1291 htod_calls: self.htod_calls.load(Ordering::Relaxed),
1292 }
1293 }
1294
1295 fn launch_metadata_snapshot(&self) -> HostLaunchMetadataTransferStats {
1296 HostLaunchMetadataTransferStats {
1297 htod_bytes: self.launch_metadata_htod_bytes.load(Ordering::Relaxed),
1298 htod_calls: self.launch_metadata_htod_calls.load(Ordering::Relaxed),
1299 }
1300 }
1301
1302 fn reset(&self) {
1303 self.dtoh_bytes.store(0, Ordering::Relaxed);
1304 self.htod_bytes.store(0, Ordering::Relaxed);
1305 self.dtoh_calls.store(0, Ordering::Relaxed);
1306 self.htod_calls.store(0, Ordering::Relaxed);
1307 self.launch_metadata_htod_bytes.store(0, Ordering::Relaxed);
1308 self.launch_metadata_htod_calls.store(0, Ordering::Relaxed);
1309 }
1310}
1311
1312impl CudaKernelProvider {
1313 fn from_loaded_runtime_parts(
1314 device: Arc<CudaDevice>,
1315 memory: Arc<GpuMemoryManager>,
1316 ptx_load_profile: Option<PtxLoadProfile>,
1317 ) -> Result<Self> {
1318 let runtime = memory.runtime().ok_or_else(|| {
1319 XlogError::Kernel(
1320 "canonical CUDA provider requires a memory manager with an owned runtime"
1321 .to_string(),
1322 )
1323 })?;
1324 if !Arc::ptr_eq(&device, memory.device()) || !Arc::ptr_eq(&device, runtime.device()) {
1325 return Err(XlogError::Kernel(
1326 "canonical CUDA provider device, memory manager, and runtime must share the exact CUDA device handle"
1327 .to_string(),
1328 ));
1329 }
1330 let device_ordinal = u32::try_from(device.ordinal()).map_err(|_| {
1331 XlogError::Kernel(format!(
1332 "CUDA device ordinal {} is not representable as u32",
1333 device.ordinal()
1334 ))
1335 })?;
1336 if runtime.device_ordinal() != device_ordinal {
1337 return Err(XlogError::Kernel(format!(
1338 "canonical CUDA provider device ordinal mismatch: device={device_ordinal} runtime={}",
1339 runtime.device_ordinal()
1340 )));
1341 }
1342 if !runtime.supports_block_use_tracking() {
1343 return Err(XlogError::Kernel(
1344 "canonical CUDA provider requires asynchronous cross-stream block-use tracking"
1345 .to_string(),
1346 ));
1347 }
1348
1349 Ok(Self::from_loaded_device(device, memory, ptx_load_profile))
1350 }
1351
1352 fn from_loaded_device(
1355 device: Arc<CudaDevice>,
1356 memory: Arc<GpuMemoryManager>,
1357 ptx_load_profile: Option<PtxLoadProfile>,
1358 ) -> Self {
1359 Self {
1360 provider_identity: NEXT_PROVIDER_IDENTITY.fetch_add(1, Ordering::Relaxed),
1361 device,
1362 memory,
1363 transfer_tracker: HostTransferTracker::default(),
1364 final_observation_transfer_tracker: HostTransferTracker::default(),
1365 final_observation_pinned_receipts: AtomicU64::new(0),
1366 ptx_load_profile,
1367 d2h_transfer_count: AtomicU64::new(0),
1368 untracked_metadata_dtoh_count: AtomicU64::new(0),
1369 strict_deterministic_d2h: AtomicBool::new(false),
1370 deterministic_d2h_violations: AtomicU64::new(0),
1371 recorded_op_stream: OnceLock::new(),
1372 csm_invocations: AtomicU64::new(0),
1373 csm_cuda_graph_captures: AtomicU64::new(0),
1374 csm_cuda_graph_launches: AtomicU64::new(0),
1375 csm_cuda_graph_fallbacks: AtomicU64::new(0),
1376 csm_cuda_graph_cache_hits: AtomicU64::new(0),
1377 small_full_row_sort_invocations: AtomicU64::new(0),
1378 csm_cuda_graph_cache: Mutex::new(HashMap::new()),
1379 wcoj_layout_fast_path_hit_count: AtomicU64::new(0),
1380 wcoj_layout_sort_invocation_count: AtomicU64::new(0),
1381 kclique_metadata_build_count: AtomicU64::new(0),
1382 kclique_metadata_build_nanos: AtomicU64::new(0),
1383 wcoj_triangle_hg_dispatch_count: AtomicU64::new(0),
1384 #[cfg(feature = "wcoj-phase-timing")]
1385 last_triangle_phase_timing: std::sync::Mutex::new(None),
1386 }
1387 }
1388
1389 fn env_flag(name: &str) -> Result<bool> {
1391 resolve_bool(None, name, false)
1392 }
1393
1394 fn any_env_flags(names: &[&str]) -> Result<bool> {
1395 let mut enabled = false;
1396 for name in names {
1397 enabled |= Self::env_flag(name)?;
1398 }
1399 Ok(enabled)
1400 }
1401
1402 pub(crate) fn use_recorded_filters_env() -> Result<bool> {
1414 Self::any_env_flags(&["XLOG_USE_RECORDED_FILTERS", "XLOG_USE_RECORDED_OPS"])
1415 }
1416
1417 pub(crate) fn use_recorded_sort_env() -> Result<bool> {
1424 Self::any_env_flags(&["XLOG_USE_RECORDED_SORT", "XLOG_USE_RECORDED_OPS"])
1425 }
1426
1427 pub(crate) fn use_recorded_dedup_env() -> Result<bool> {
1432 Self::any_env_flags(&["XLOG_USE_RECORDED_DEDUP", "XLOG_USE_RECORDED_OPS"])
1433 }
1434
1435 pub(crate) fn use_recorded_groupby_env() -> Result<bool> {
1441 Self::any_env_flags(&["XLOG_USE_RECORDED_GROUPBY", "XLOG_USE_RECORDED_OPS"])
1442 }
1443
1444 pub(crate) fn use_recorded_hash_join_env() -> Result<bool> {
1452 Self::any_env_flags(&["XLOG_USE_RECORDED_HASH_JOIN", "XLOG_USE_RECORDED_OPS"])
1453 }
1454
1455 pub(crate) fn use_recorded_csm_env() -> Result<bool> {
1463 Self::any_env_flags(&["XLOG_USE_RECORDED_CSM", "XLOG_USE_RECORDED_OPS"])
1464 }
1465
1466 pub(crate) fn use_csm_cuda_graph_env() -> Result<bool> {
1472 Self::any_env_flags(&["XLOG_USE_CSM_CUDA_GRAPH", "XLOG_USE_CUDA_GRAPHS"])
1473 }
1474
1475 #[doc(hidden)]
1489 pub fn csm_invocations(&self) -> u64 {
1490 self.csm_invocations.load(Ordering::Relaxed)
1491 }
1492
1493 #[doc(hidden)]
1494 pub fn csm_cuda_graph_captures(&self) -> u64 {
1495 self.csm_cuda_graph_captures.load(Ordering::Relaxed)
1496 }
1497
1498 #[doc(hidden)]
1499 pub fn csm_cuda_graph_launches(&self) -> u64 {
1500 self.csm_cuda_graph_launches.load(Ordering::Relaxed)
1501 }
1502
1503 #[doc(hidden)]
1504 pub fn csm_cuda_graph_fallbacks(&self) -> u64 {
1505 self.csm_cuda_graph_fallbacks.load(Ordering::Relaxed)
1506 }
1507
1508 #[doc(hidden)]
1509 pub fn csm_cuda_graph_cache_hits(&self) -> u64 {
1510 self.csm_cuda_graph_cache_hits.load(Ordering::Relaxed)
1511 }
1512
1513 #[doc(hidden)]
1514 pub fn small_full_row_sort_invocations(&self) -> u64 {
1515 self.small_full_row_sort_invocations.load(Ordering::Relaxed)
1516 }
1517
1518 pub(crate) fn recorded_op_stream_or_init(&self) -> Option<crate::device_runtime::StreamId> {
1538 if let Some(s) = self.recorded_op_stream.get() {
1539 return Some(*s);
1540 }
1541 let runtime = self.memory.runtime()?;
1542 let stream = runtime.stream_pool().acquire().ok()?;
1543 let _ = self.recorded_op_stream.set(stream);
1544 self.recorded_op_stream.get().copied()
1545 }
1546
1547 #[cfg(feature = "wcoj-phase-timing")]
1557 pub fn take_wcoj_triangle_phase_timing(
1558 &self,
1559 ) -> Option<crate::wcoj_phase_timing::WcojTrianglePhaseTiming> {
1560 self.last_triangle_phase_timing
1561 .lock()
1562 .ok()
1563 .and_then(|mut g| g.take())
1564 }
1565
1566 pub fn wcoj_layout_fast_path_hit_count(&self) -> u64 {
1573 self.wcoj_layout_fast_path_hit_count.load(Ordering::Relaxed)
1574 }
1575
1576 pub fn wcoj_triangle_hg_dispatch_count(&self) -> u64 {
1579 self.wcoj_triangle_hg_dispatch_count.load(Ordering::Relaxed)
1580 }
1581
1582 pub fn reset_wcoj_layout_fast_path_hit_count(&self) {
1585 self.wcoj_layout_fast_path_hit_count
1586 .store(0, Ordering::Relaxed);
1587 }
1588
1589 pub fn wcoj_layout_sort_invocation_count(&self) -> u64 {
1592 self.wcoj_layout_sort_invocation_count
1593 .load(Ordering::Relaxed)
1594 }
1595
1596 pub fn reset_wcoj_layout_sort_invocation_count(&self) {
1598 self.wcoj_layout_sort_invocation_count
1599 .store(0, Ordering::Relaxed);
1600 }
1601
1602 pub fn kclique_metadata_build_count(&self) -> u64 {
1605 self.kclique_metadata_build_count.load(Ordering::Relaxed)
1606 }
1607
1608 pub fn kclique_metadata_build_nanos(&self) -> u64 {
1611 self.kclique_metadata_build_nanos.load(Ordering::Relaxed)
1612 }
1613
1614 pub fn reset_kclique_metadata_build_metrics(&self) {
1616 self.kclique_metadata_build_count
1617 .store(0, Ordering::Relaxed);
1618 self.kclique_metadata_build_nanos
1619 .store(0, Ordering::Relaxed);
1620 }
1621
1622 pub(crate) fn record_wcoj_layout_fast_path_hit(&self) {
1626 self.wcoj_layout_fast_path_hit_count
1627 .fetch_add(1, Ordering::Relaxed);
1628 }
1629
1630 pub(crate) fn record_wcoj_layout_sort_invocation(&self) {
1632 self.wcoj_layout_sort_invocation_count
1633 .fetch_add(1, Ordering::Relaxed);
1634 }
1635
1636 pub(crate) fn record_kclique_metadata_build_nanos(&self, nanos: u128) {
1638 self.kclique_metadata_build_count
1639 .fetch_add(1, Ordering::Relaxed);
1640 let nanos = u64::try_from(nanos).unwrap_or(u64::MAX);
1641 self.kclique_metadata_build_nanos
1642 .fetch_add(nanos, Ordering::Relaxed);
1643 }
1644
1645 #[doc(hidden)]
1648 pub fn record_wcoj_triangle_hg_dispatch(&self) {
1649 self.wcoj_triangle_hg_dispatch_count
1650 .fetch_add(1, Ordering::Relaxed);
1651 }
1652
1653 pub fn device(&self) -> &Arc<CudaDevice> {
1655 &self.device
1656 }
1657
1658 pub(crate) fn provider_identity(&self) -> u64 {
1659 self.provider_identity
1660 }
1661
1662 pub fn memory(&self) -> &Arc<GpuMemoryManager> {
1664 &self.memory
1665 }
1666
1667 pub fn ptx_load_profile(&self) -> Option<&PtxLoadProfile> {
1669 self.ptx_load_profile.as_ref()
1670 }
1671
1672 pub fn reset_host_transfer_stats(&self) {
1674 self.transfer_tracker.reset();
1675 }
1676
1677 pub fn host_transfer_stats(&self) -> HostTransferStats {
1679 self.transfer_tracker.snapshot()
1680 }
1681
1682 pub fn begin_resident_host_transfer_scope(&self) -> ResidentHostTransferScope {
1685 let counters = std::rc::Rc::new(std::cell::Cell::new(ResidentHostTransferStats::default()));
1686 RESIDENT_TRANSFER_SCOPES.with_borrow_mut(|scopes| {
1687 scopes.push((self.provider_identity, counters.clone()));
1688 });
1689 ResidentHostTransferScope {
1690 provider_identity: self.provider_identity,
1691 counters,
1692 }
1693 }
1694
1695 pub fn reset_final_observation_transfer_stats(&self) {
1697 self.final_observation_transfer_tracker.reset();
1698 self.final_observation_pinned_receipts
1699 .store(0, Ordering::Relaxed);
1700 }
1701
1702 pub fn final_observation_transfer_stats(&self) -> FinalObservationTransferStats {
1704 let transfers = self.final_observation_transfer_tracker.snapshot();
1705 FinalObservationTransferStats {
1706 dtoh_bytes: transfers.dtoh_bytes,
1707 dtoh_calls: transfers.dtoh_calls,
1708 pinned_receipts: self
1709 .final_observation_pinned_receipts
1710 .load(Ordering::Relaxed),
1711 }
1712 }
1713
1714 #[doc(hidden)]
1720 pub fn record_final_observation_transfer(&self, bytes: u64) {
1721 self.final_observation_transfer_tracker.record_dtoh(bytes);
1722 self.final_observation_pinned_receipts
1723 .fetch_add(1, Ordering::Relaxed);
1724 record_resident_transfer(self.provider_identity, |stats| {
1725 stats.final_dtoh_calls += 1;
1726 stats.final_dtoh_bytes += bytes;
1727 stats.final_pinned_receipts += 1;
1728 });
1729 }
1730
1731 #[doc(hidden)]
1738 pub fn finalize_resident_logical_counts(&self, entries: &[(&CudaBuffer, u32)]) -> Result<()> {
1739 for (buffer, count) in entries {
1740 if u64::from(*count) > buffer.num_rows() {
1741 return Err(XlogError::Kernel(format!(
1742 "resident receipt logical row count {} exceeds buffer capacity {}",
1743 count,
1744 buffer.num_rows()
1745 )));
1746 }
1747 if let Some(cached) = buffer.cached_row_count() {
1748 if cached != *count {
1749 return Err(XlogError::Kernel(format!(
1750 "resident receipt logical row count {} conflicts with cached count {}",
1751 count, cached
1752 )));
1753 }
1754 }
1755 }
1756 for (buffer, count) in entries {
1757 buffer.set_cached_row_count_if_unset(*count);
1758 }
1759 Ok(())
1760 }
1761
1762 pub fn host_launch_metadata_transfer_stats(&self) -> HostLaunchMetadataTransferStats {
1765 self.transfer_tracker.launch_metadata_snapshot()
1766 }
1767
1768 pub fn d2h_transfer_count(&self) -> u64 {
1774 self.d2h_transfer_count.load(Ordering::Relaxed)
1775 }
1776
1777 pub fn reset_d2h_transfer_count(&self) {
1779 self.d2h_transfer_count.store(0, Ordering::Relaxed);
1780 }
1781
1782 pub fn untracked_metadata_dtoh_count(&self) -> u64 {
1785 self.untracked_metadata_dtoh_count.load(Ordering::Relaxed)
1786 }
1787
1788 pub fn reset_untracked_metadata_dtoh_count(&self) {
1790 self.untracked_metadata_dtoh_count
1791 .store(0, Ordering::Relaxed);
1792 }
1793
1794 pub fn enable_strict_deterministic_d2h(&self) {
1811 self.strict_deterministic_d2h.store(true, Ordering::Relaxed);
1812 }
1813
1814 pub fn disable_strict_deterministic_d2h(&self) {
1816 self.strict_deterministic_d2h
1817 .store(false, Ordering::Relaxed);
1818 }
1819
1820 pub fn strict_deterministic_d2h_enabled(&self) -> bool {
1822 self.strict_deterministic_d2h.load(Ordering::Relaxed)
1823 }
1824
1825 pub fn deterministic_d2h_violation_count(&self) -> u64 {
1827 self.deterministic_d2h_violations.load(Ordering::Relaxed)
1828 }
1829
1830 pub fn reset_deterministic_d2h_violations(&self) {
1832 self.deterministic_d2h_violations
1833 .store(0, Ordering::Relaxed);
1834 }
1835
1836 pub(crate) fn check_deterministic_d2h(&self, op: &'static str, bytes: u64) -> Result<()> {
1842 if self.strict_deterministic_d2h.load(Ordering::Relaxed) {
1843 self.deterministic_d2h_violations
1844 .fetch_add(1, Ordering::Relaxed);
1845 record_resident_transfer(self.provider_identity, |stats| {
1846 stats.deterministic_dtoh_violations += 1;
1847 });
1848 return Err(XlogError::Execution(format!(
1849 "deterministic D2H gate: {} attempted to copy {} bytes from device to host",
1850 op, bytes
1851 )));
1852 }
1853 Ok(())
1854 }
1855
1856 pub(crate) fn dtoh_sync_copy_into_tracked<T: DeviceRepr, Src: crate::memory::DeviceRead<T>>(
1857 &self,
1858 src: &Src,
1859 dst: &mut [T],
1860 ) -> Result<()> {
1861 let bytes = std::mem::size_of::<T>()
1862 .checked_mul(dst.len())
1863 .ok_or_else(|| XlogError::Kernel("dtoh size overflow".to_string()))?;
1864 self.check_deterministic_d2h("dtoh_sync_copy_into_tracked", bytes as u64)?;
1865 self.transfer_tracker.record_dtoh(bytes as u64);
1866 record_resident_transfer(self.provider_identity, |stats| {
1867 stats.tracked_dtoh_calls += 1;
1868 stats.tracked_dtoh_bytes += bytes as u64;
1869 });
1870 self.device
1871 .inner()
1872 .dtoh_sync_copy_into(src, dst)
1873 .map_err(|e| XlogError::Kernel(format!("Failed to copy from device: {}", e)))
1874 }
1875
1876 pub const DTOH_SMALL_METADATA_MAX_BYTES: usize = 4096;
1881
1882 pub fn dtoh_small_metadata_untracked<T: DeviceRepr + Default + Copy>(
1907 &self,
1908 src: &crate::memory::TrackedCudaSlice<T>,
1909 count: usize,
1910 ) -> Result<Vec<T>> {
1911 let bytes = count.checked_mul(std::mem::size_of::<T>()).ok_or_else(|| {
1912 XlogError::Kernel("dtoh_small_metadata_untracked: byte size overflow".to_string())
1913 })?;
1914 if bytes > Self::DTOH_SMALL_METADATA_MAX_BYTES {
1915 return Err(XlogError::Kernel(format!(
1916 "dtoh_small_metadata_untracked: requested {} bytes exceeds metadata cap of {} bytes \
1917 (this is metadata-only; use download_column* for data-plane transfers)",
1918 bytes,
1919 Self::DTOH_SMALL_METADATA_MAX_BYTES
1920 )));
1921 }
1922 if count > src.len() {
1923 return Err(XlogError::Kernel(format!(
1924 "dtoh_small_metadata_untracked: count={count} > src.len={}",
1925 src.len()
1926 )));
1927 }
1928 if count == 0 {
1929 return Ok(Vec::new());
1930 }
1931 let slice = src.try_slice(0..count).ok_or_else(|| {
1932 XlogError::Kernel(format!(
1933 "dtoh_small_metadata_untracked: try_slice(0..{count}) failed"
1934 ))
1935 })?;
1936 let mut buf: Vec<T> = vec![T::default(); count];
1937 self.untracked_metadata_dtoh_count
1938 .fetch_add(1, Ordering::Relaxed);
1939 record_resident_transfer(self.provider_identity, |stats| {
1940 stats.untracked_metadata_dtoh_calls += 1;
1941 });
1942 self.device
1943 .inner()
1944 .dtoh_sync_copy_into(&slice, &mut buf)
1945 .map_err(|e| {
1946 XlogError::Kernel(format!("dtoh_small_metadata_untracked: copy failed: {}", e))
1947 })?;
1948 Ok(buf)
1949 }
1950
1951 pub fn dtoh_scalar_untracked<T: DeviceRepr + Default + Copy>(
1959 &self,
1960 src: &crate::memory::TrackedCudaSlice<T>,
1961 index: usize,
1962 ) -> Result<T> {
1963 if index >= src.len() {
1964 return Err(XlogError::Kernel(format!(
1965 "dtoh_scalar_untracked: index={} >= len={}",
1966 index,
1967 src.len()
1968 )));
1969 }
1970 let slice = src.try_slice(index..index + 1).ok_or_else(|| {
1971 XlogError::Kernel(format!(
1972 "dtoh_scalar_untracked: slice failed at index={}",
1973 index
1974 ))
1975 })?;
1976 let mut buf = [T::default()];
1977 self.untracked_metadata_dtoh_count
1978 .fetch_add(1, Ordering::Relaxed);
1979 record_resident_transfer(self.provider_identity, |stats| {
1980 stats.untracked_metadata_dtoh_calls += 1;
1981 });
1982 self.device
1983 .inner()
1984 .dtoh_sync_copy_into(&slice, &mut buf)
1985 .map_err(|e| XlogError::Kernel(format!("dtoh_scalar_untracked: copy failed: {}", e)))?;
1986 Ok(buf[0])
1987 }
1988
1989 pub fn htod_sync_copy_into_tracked<T: DeviceRepr, Dst: crate::memory::DeviceWrite<T>>(
1991 &self,
1992 src: &[T],
1993 dst: &mut Dst,
1994 ) -> Result<()> {
1995 let bytes = std::mem::size_of::<T>()
1996 .checked_mul(src.len())
1997 .ok_or_else(|| XlogError::Kernel("htod size overflow".to_string()))?;
1998 self.transfer_tracker.record_htod(bytes as u64);
1999 record_resident_transfer(self.provider_identity, |stats| {
2000 stats.tracked_htod_calls += 1;
2001 stats.tracked_htod_bytes += bytes as u64;
2002 });
2003 self.device
2004 .inner()
2005 .htod_sync_copy_into(src, dst)
2006 .map_err(|e| XlogError::Kernel(format!("Failed to copy to device: {}", e)))
2007 }
2008
2009 pub fn htod_sync_copy_tracked<T: DeviceRepr + 'static>(
2012 &self,
2013 src: &[T],
2014 ) -> Result<crate::memory::DeviceMemoryView<T>> {
2015 let bytes = std::mem::size_of::<T>()
2016 .checked_mul(src.len())
2017 .ok_or_else(|| XlogError::Kernel("htod size overflow".to_string()))?;
2018 self.transfer_tracker.record_htod(bytes as u64);
2019 record_resident_transfer(self.provider_identity, |stats| {
2020 stats.tracked_htod_calls += 1;
2021 stats.tracked_htod_bytes += bytes as u64;
2022 });
2023 self.device
2024 .inner()
2025 .htod_sync_copy(src)
2026 .map_err(|e| XlogError::Kernel(format!("Failed to copy to device: {}", e)))
2027 }
2028
2029 pub fn htod_launch_metadata_sync_copy_into<
2032 T: DeviceRepr,
2033 Dst: crate::memory::DeviceWrite<T>,
2034 >(
2035 &self,
2036 src: &[T],
2037 dst: &mut Dst,
2038 ) -> Result<()> {
2039 let bytes = std::mem::size_of::<T>()
2040 .checked_mul(src.len())
2041 .ok_or_else(|| XlogError::Kernel("launch metadata htod size overflow".to_string()))?;
2042 self.transfer_tracker
2043 .record_htod_launch_metadata(bytes as u64);
2044 self.device
2045 .inner()
2046 .htod_sync_copy_into(src, dst)
2047 .map_err(|e| {
2048 XlogError::Kernel(format!("Failed to copy launch metadata to device: {}", e))
2049 })
2050 }
2051
2052 pub(crate) fn initialize_launch_metadata_u32(
2055 &self,
2056 value: u32,
2057 dst: &TrackedCudaSlice<u32>,
2058 enqueue: &crate::launch::CudaEnqueue<'_>,
2059 context: &str,
2060 ) -> Result<()> {
2061 let destination_context = dst.stream().context().cu_ctx() as usize;
2062 let range = crate::device_runtime::resource::MemoryUse::new(
2063 dst.device_ptr_value(),
2064 std::mem::size_of::<u32>(),
2065 crate::device_runtime::Access::Write,
2066 )
2067 .map_err(|error| XlogError::Kernel(format!("{context}: {error}")))?;
2068 if dst.is_empty()
2069 || destination_context != enqueue.stream().context().cu_ctx() as usize
2070 || !enqueue.manifest().covers(&[(destination_context, range)])
2071 {
2072 return Err(XlogError::Kernel(format!(
2073 "{context}: metadata destination lacks an admitted write"
2074 )));
2075 }
2076 enqueue
2077 .submit(|_| unsafe {
2078 cudarc::driver::sys::cuMemsetD32Async(
2081 dst.device_ptr_value(),
2082 value,
2083 1,
2084 enqueue.stream().cu_stream(),
2085 )
2086 .result()
2087 })
2088 .map_err(|error| {
2089 XlogError::Kernel(format!(
2090 "{context}: metadata initialization failed: {error}"
2091 ))
2092 })
2093 }
2094
2095 pub fn exclusive_scan_u32_inplace(
2124 &self,
2125 data: &mut crate::memory::TrackedCudaSlice<u32>,
2126 n: u32,
2127 ) -> Result<()> {
2128 if n as usize > data.len() {
2129 return Err(XlogError::Kernel(format!(
2130 "exclusive_scan_u32_inplace: n={} exceeds slice len={}",
2131 n,
2132 data.len()
2133 )));
2134 }
2135 self.multiblock_scan_u32_inplace(data, n)
2136 }
2137
2138 fn multiblock_scan_u32_inplace(
2139 &self,
2140 data: &mut crate::memory::TrackedCudaSlice<u32>,
2141 n: u32,
2142 ) -> Result<()> {
2143 if n == 0 {
2144 return Ok(());
2145 }
2146
2147 let device = self.device.inner();
2148 let block_size = 256u32;
2149
2150 if n <= block_size {
2151 let phase2_fn = device
2152 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE2)
2153 .ok_or_else(|| {
2154 XlogError::Kernel("Failed to get multiblock_scan_phase2 kernel".to_string())
2155 })?;
2156
2157 unsafe {
2159 phase2_fn.clone().launch(
2160 LaunchConfig {
2161 grid_dim: (1, 1, 1),
2162 block_dim: (block_size, 1, 1),
2163 shared_mem_bytes: 0,
2164 },
2165 (&mut *data, n),
2166 )
2167 }
2168 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase2 failed: {}", e)))?;
2169
2170 self.device.synchronize()?;
2171 return Ok(());
2172 }
2173
2174 let num_blocks = n.div_ceil(block_size);
2175 let mut block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
2176
2177 let phase1_u32_fn = device
2178 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_U32_PHASE1)
2179 .ok_or_else(|| {
2180 XlogError::Kernel("Failed to get multiblock_scan_u32_phase1 kernel".to_string())
2181 })?;
2182
2183 unsafe {
2185 phase1_u32_fn.clone().launch(
2186 LaunchConfig {
2187 grid_dim: (num_blocks, 1, 1),
2188 block_dim: (block_size, 1, 1),
2189 shared_mem_bytes: 0,
2190 },
2191 (&mut *data, &mut block_sums, n),
2192 )
2193 }
2194 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_u32_phase1 failed: {}", e)))?;
2195 self.device.synchronize()?;
2196
2197 if num_blocks > 1 {
2198 self.multiblock_scan_u32_inplace(&mut block_sums, num_blocks)?;
2199 }
2200
2201 let phase3_fn = device
2202 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
2203 .ok_or_else(|| {
2204 XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
2205 })?;
2206
2207 unsafe {
2209 phase3_fn.clone().launch(
2210 LaunchConfig {
2211 grid_dim: (num_blocks, 1, 1),
2212 block_dim: (block_size, 1, 1),
2213 shared_mem_bytes: 0,
2214 },
2215 (&mut *data, &block_sums, n),
2216 )
2217 }
2218 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase3 failed: {}", e)))?;
2219
2220 self.device.synchronize()?;
2221 Ok(())
2222 }
2223
2224 pub(crate) fn multiblock_scan_u32_scratch_for_len(
2227 &self,
2228 mut n: u32,
2229 ) -> Result<MultiblockScanScratchU32> {
2230 let block_size = 256u32;
2231 let mut levels = Vec::new();
2232 while n > block_size {
2233 let num_blocks = n.div_ceil(block_size);
2234 levels.push(self.memory.alloc::<u32>(num_blocks as usize)?);
2235 n = num_blocks;
2236 }
2237 Ok(MultiblockScanScratchU32 { levels })
2238 }
2239
2240 pub(crate) fn multiblock_scan_u32_inplace_on_stream<D: crate::memory::DeviceWrite<u32>>(
2243 &self,
2244 data: &mut D,
2245 n: u32,
2246 enqueue: &crate::launch::CudaEnqueue<'_>,
2247 scratch: &mut MultiblockScanScratchU32,
2248 ) -> Result<()> {
2249 self.multiblock_scan_u32_inplace_on_stream_with_scratch_levels(
2250 &mut data.device_view(),
2251 n,
2252 enqueue,
2253 &mut scratch.levels,
2254 )
2255 }
2256
2257 fn multiblock_scan_u32_inplace_on_stream_with_scratch_levels(
2258 &self,
2259 data: &mut DeviceMemoryView<u32>,
2260 n: u32,
2261 enqueue: &crate::launch::CudaEnqueue<'_>,
2262 scratch_levels: &mut [TrackedCudaSlice<u32>],
2263 ) -> Result<()> {
2264 if n == 0 {
2265 return Ok(());
2266 }
2267 let device = self.device.inner();
2268 let block_size = 256u32;
2269
2270 if n <= block_size {
2271 let phase2_fn = device
2272 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE2)
2273 .ok_or_else(|| {
2274 XlogError::Kernel("Failed to get multiblock_scan_phase2 kernel".to_string())
2275 })?;
2276 unsafe {
2278 phase2_fn.clone().launch_in(
2279 enqueue,
2280 LaunchConfig {
2281 grid_dim: (1, 1, 1),
2282 block_dim: (block_size, 1, 1),
2283 shared_mem_bytes: 0,
2284 },
2285 (&mut *data, n),
2286 )
2287 }
2288 .map_err(|e| {
2289 XlogError::Kernel(format!(
2290 "multiblock_scan_phase2 (graph scratch) failed: {}",
2291 e
2292 ))
2293 })?;
2294 return Ok(());
2295 }
2296
2297 let num_blocks = n.div_ceil(block_size);
2298 let (block_sums, rest) = scratch_levels.split_first_mut().ok_or_else(|| {
2299 XlogError::Kernel(format!(
2300 "multiblock_scan_u32_inplace_on_stream_with_scratch: missing scratch level \
2301 for n={n}, num_blocks={num_blocks}"
2302 ))
2303 })?;
2304 if block_sums.len() < num_blocks as usize {
2305 return Err(XlogError::Kernel(format!(
2306 "multiblock_scan_u32_inplace_on_stream_with_scratch: scratch level too small \
2307 (have {}, need {})",
2308 block_sums.len(),
2309 num_blocks
2310 )));
2311 }
2312
2313 let phase1_u32_fn = device
2314 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_U32_PHASE1)
2315 .ok_or_else(|| {
2316 XlogError::Kernel("Failed to get multiblock_scan_u32_phase1 kernel".to_string())
2317 })?;
2318 unsafe {
2320 phase1_u32_fn.clone().launch_in(
2321 enqueue,
2322 LaunchConfig {
2323 grid_dim: (num_blocks, 1, 1),
2324 block_dim: (block_size, 1, 1),
2325 shared_mem_bytes: 0,
2326 },
2327 (&mut *data, &mut *block_sums, n),
2328 )
2329 }
2330 .map_err(|e| {
2331 XlogError::Kernel(format!(
2332 "multiblock_scan_u32_phase1 (graph scratch) failed: {}",
2333 e
2334 ))
2335 })?;
2336
2337 if num_blocks > 1 {
2338 self.multiblock_scan_u32_inplace_on_stream_with_scratch_levels(
2339 &mut block_sums.view(),
2340 num_blocks,
2341 enqueue,
2342 rest,
2343 )?;
2344 }
2345
2346 let phase3_fn = device
2347 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
2348 .ok_or_else(|| {
2349 XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
2350 })?;
2351 unsafe {
2353 phase3_fn.clone().launch_in(
2354 enqueue,
2355 LaunchConfig {
2356 grid_dim: (num_blocks, 1, 1),
2357 block_dim: (block_size, 1, 1),
2358 shared_mem_bytes: 0,
2359 },
2360 (&mut *data, &*block_sums, n),
2361 )
2362 }
2363 .map_err(|e| {
2364 XlogError::Kernel(format!(
2365 "multiblock_scan_phase3 (graph scratch) failed: {}",
2366 e
2367 ))
2368 })?;
2369 Ok(())
2370 }
2371
2372 fn multiblock_scan_u32_view_inplace(
2373 &self,
2374 data: &mut DeviceMemoryView<u32>,
2375 n: u32,
2376 ) -> Result<()> {
2377 if n == 0 {
2378 return Ok(());
2379 }
2380
2381 let device = self.device.inner();
2382 let block_size = 256u32;
2383
2384 if n <= block_size {
2385 let phase2_fn = device
2386 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE2)
2387 .ok_or_else(|| {
2388 XlogError::Kernel("Failed to get multiblock_scan_phase2 kernel".to_string())
2389 })?;
2390
2391 unsafe {
2393 phase2_fn.clone().launch(
2394 LaunchConfig {
2395 grid_dim: (1, 1, 1),
2396 block_dim: (block_size, 1, 1),
2397 shared_mem_bytes: 0,
2398 },
2399 (data, n),
2400 )
2401 }
2402 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase2 failed: {}", e)))?;
2403
2404 self.device.synchronize()?;
2405 return Ok(());
2406 }
2407
2408 let num_blocks = n.div_ceil(block_size);
2409 let mut block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
2410
2411 let phase1_u32_fn = device
2412 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_U32_PHASE1)
2413 .ok_or_else(|| {
2414 XlogError::Kernel("Failed to get multiblock_scan_u32_phase1 kernel".to_string())
2415 })?;
2416
2417 unsafe {
2419 phase1_u32_fn.clone().launch(
2420 LaunchConfig {
2421 grid_dim: (num_blocks, 1, 1),
2422 block_dim: (block_size, 1, 1),
2423 shared_mem_bytes: 0,
2424 },
2425 (&mut *data, &mut block_sums, n),
2426 )
2427 }
2428 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_u32_phase1 failed: {}", e)))?;
2429 self.device.synchronize()?;
2430
2431 if num_blocks > 1 {
2432 self.multiblock_scan_u32_inplace(&mut block_sums, num_blocks)?;
2433 }
2434
2435 let phase3_fn = device
2436 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
2437 .ok_or_else(|| {
2438 XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
2439 })?;
2440
2441 unsafe {
2443 phase3_fn.clone().launch(
2444 LaunchConfig {
2445 grid_dim: (num_blocks, 1, 1),
2446 block_dim: (block_size, 1, 1),
2447 shared_mem_bytes: 0,
2448 },
2449 (&mut *data, &block_sums, n),
2450 )
2451 }
2452 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase3 failed: {}", e)))?;
2453
2454 self.device.synchronize()?;
2455 Ok(())
2456 }
2457
2458 pub fn device_row_count(&self, buffer: &CudaBuffer) -> Result<usize> {
2463 if let Some(n) = buffer.cached_row_count() {
2464 return Ok(n as usize);
2465 }
2466 let host_rows = self.dtoh_small_metadata_untracked(buffer.num_rows_device(), 1)?;
2467 buffer.set_cached_row_count_if_unset(host_rows[0]);
2468 Ok(host_rows[0] as usize)
2469 }
2470
2471 pub fn validated_logical_row_count(&self, buffer: &CudaBuffer) -> Result<usize> {
2476 let logical_rows = self.device_row_count(buffer)?;
2477 validate_logical_row_count(buffer.num_rows(), logical_rows)
2478 }
2479
2480 fn clone_device_row_count(&self, buffer: &CudaBuffer) -> Result<TrackedCudaSlice<u32>> {
2481 let mut d_num_rows = self.memory.alloc::<u32>(1)?;
2482 self.device
2483 .inner()
2484 .dtod_copy(buffer.num_rows_device(), &mut d_num_rows)
2485 .map_err(|e| XlogError::Kernel(format!("Failed to copy row count: {}", e)))?;
2486 Ok(d_num_rows)
2487 }
2488
2489 pub fn overwrite_resident_source(
2492 &self,
2493 target: &mut CudaBuffer,
2494 source: &CudaBuffer,
2495 ) -> Result<()> {
2496 if target.schema() != source.schema()
2497 || target.num_rows() != source.num_rows()
2498 || target.columns().len() != source.columns().len()
2499 || target
2500 .columns()
2501 .iter()
2502 .zip(source.columns())
2503 .any(|(target, source)| target.num_bytes() != source.num_bytes())
2504 || !target.canonical_full_row_set_certified()
2505 || !source.canonical_full_row_set_certified()
2506 {
2507 return Err(XlogError::Execution(
2508 "resident source replacement changes captured shape or set semantics".into(),
2509 ));
2510 }
2511 for (source_column, target_column) in source.columns().iter().zip(target.columns_mut()) {
2512 if source_column.num_bytes() == 0 {
2513 continue;
2514 }
2515 self.device
2516 .inner()
2517 .dtod_copy(source_column, target_column)
2518 .map_err(|error| {
2519 XlogError::Kernel(format!("resident source copy failed: {error}"))
2520 })?;
2521 }
2522 self.device
2523 .inner()
2524 .dtod_copy(source.num_rows_device(), target.num_rows_device_mut())
2525 .map_err(|error| {
2526 XlogError::Kernel(format!("resident source count copy failed: {error}"))
2527 })?;
2528 if let Some(source_rows) = source.cached_row_count() {
2529 target.set_cached_row_count_after_device_copy(source_rows);
2530 } else {
2531 target.invalidate_cached_row_count();
2532 }
2533 target.certify_canonical_full_row_set();
2534 Ok(())
2535 }
2536
2537 fn upload_device_row_count(&self, row_count: u32) -> Result<TrackedCudaSlice<u32>> {
2538 let mut d_num_rows = self.memory.alloc::<u32>(1)?;
2539 self.htod_launch_metadata_sync_copy_into(&[row_count], &mut d_num_rows)
2540 .map_err(|e| XlogError::Kernel(format!("Failed to upload row count: {}", e)))?;
2541 Ok(d_num_rows)
2542 }
2543
2544 fn buffer_from_columns_with_device_count(
2545 &self,
2546 columns: Vec<CudaColumn>,
2547 row_cap: u64,
2548 schema: Schema,
2549 src: &CudaBuffer,
2550 ) -> Result<CudaBuffer> {
2551 let d_num_rows = self.clone_device_row_count(src)?;
2552 Ok(match src.cached_row_count() {
2553 Some(row_count) => CudaBuffer::from_columns_with_host_count(
2554 columns, row_cap, d_num_rows, schema, row_count,
2555 ),
2556 None => CudaBuffer::from_columns(columns, row_cap, d_num_rows, schema),
2557 })
2558 }
2559
2560 fn column_bytes_view(
2561 &self,
2562 col: &CudaColumn,
2563 num_bytes: usize,
2564 ) -> Result<DeviceMemoryView<u8>> {
2565 if col.num_bytes() < num_bytes {
2566 return Err(XlogError::Kernel(format!(
2567 "Column has {} bytes but {} required",
2568 col.num_bytes(),
2569 num_bytes
2570 )));
2571 }
2572 Ok(col.device_view().slice(..num_bytes))
2573 }
2574
2575 fn bytes_as_u32_view(
2576 &self,
2577 bytes: &TrackedCudaSlice<u8>,
2578 num_elements: usize,
2579 ) -> Result<DeviceMemoryView<u32>> {
2580 let required_bytes = num_elements * std::mem::size_of::<u32>();
2581 if bytes.len() < required_bytes {
2582 return Err(XlogError::Kernel(format!(
2583 "Packed keys have {} bytes but {} required for {} u32 elements",
2584 bytes.len(),
2585 required_bytes,
2586 num_elements
2587 )));
2588 }
2589 let ptr = *bytes.device_ptr();
2590 if !(ptr as usize).is_multiple_of(std::mem::align_of::<u32>()) {
2591 return Err(XlogError::Kernel(
2592 "Packed keys device pointer is not u32-aligned".to_string(),
2593 ));
2594 }
2595 unsafe { bytes.device_view().slice(..required_bytes).cast() }
2598 .ok_or_else(|| XlogError::Kernel("device column cast is invalid".into()))
2599 }
2600
2601 fn column_as_u32_view(
2603 &self,
2604 col: &CudaColumn,
2605 num_elements: usize,
2606 ) -> Result<DeviceMemoryView<u32>> {
2607 let required_bytes = num_elements * std::mem::size_of::<u32>();
2608 if col.num_bytes() < required_bytes {
2609 return Err(XlogError::Kernel(format!(
2610 "Column has {} bytes but {} required for {} u32 elements",
2611 col.num_bytes(),
2612 required_bytes,
2613 num_elements
2614 )));
2615 }
2616 let ptr = *col.device_ptr();
2617 if !(ptr as usize).is_multiple_of(std::mem::align_of::<u32>()) {
2618 return Err(XlogError::Kernel(
2619 "Column device pointer is not u32-aligned".to_string(),
2620 ));
2621 }
2622 unsafe { col.device_view().slice(..required_bytes).cast() }
2625 .ok_or_else(|| XlogError::Kernel("device column cast is invalid".into()))
2626 }
2627
2628 fn column_as_u64_view(
2629 &self,
2630 col: &CudaColumn,
2631 num_elements: usize,
2632 ) -> Result<DeviceMemoryView<u64>> {
2633 let required_bytes = num_elements * std::mem::size_of::<u64>();
2634 if col.num_bytes() < required_bytes {
2635 return Err(XlogError::Kernel(format!(
2636 "Column has {} bytes but {} required for {} u64 elements",
2637 col.num_bytes(),
2638 required_bytes,
2639 num_elements
2640 )));
2641 }
2642 let ptr = *col.device_ptr();
2643 if !(ptr as usize).is_multiple_of(std::mem::align_of::<u64>()) {
2644 return Err(XlogError::Kernel(
2645 "Column device pointer is not u64-aligned".to_string(),
2646 ));
2647 }
2648 unsafe { col.device_view().slice(..required_bytes).cast() }
2651 .ok_or_else(|| XlogError::Kernel("device column cast is invalid".into()))
2652 }
2653
2654 fn column_as_f64_view(
2656 &self,
2657 col: &CudaColumn,
2658 num_elements: usize,
2659 ) -> Result<DeviceMemoryView<f64>> {
2660 let required_bytes = num_elements * std::mem::size_of::<f64>();
2661 if col.num_bytes() < required_bytes {
2662 return Err(XlogError::Kernel(format!(
2663 "Column has {} bytes but {} required for {} f64 elements",
2664 col.num_bytes(),
2665 required_bytes,
2666 num_elements
2667 )));
2668 }
2669 let ptr = *col.device_ptr();
2670 if !(ptr as usize).is_multiple_of(std::mem::align_of::<f64>()) {
2671 return Err(XlogError::Kernel(
2672 "Column device pointer is not f64-aligned".to_string(),
2673 ));
2674 }
2675 unsafe { col.device_view().slice(..required_bytes).cast() }
2678 .ok_or_else(|| XlogError::Kernel("device column cast is invalid".into()))
2679 }
2680
2681 pub fn create_empty_buffer(&self, schema: Schema) -> Result<CudaBuffer> {
2692 let mut buffers = self.create_empty_buffers(&[schema])?;
2693 Ok(buffers.pop().expect("one schema produces one empty buffer"))
2694 }
2695
2696 pub fn create_empty_buffers(&self, schemas: &[Schema]) -> Result<Vec<CudaBuffer>> {
2698 if schemas.is_empty() {
2699 return Ok(Vec::new());
2700 }
2701
2702 let count_bytes = schemas
2703 .len()
2704 .checked_mul(std::mem::size_of::<u32>())
2705 .ok_or_else(|| XlogError::Kernel("empty relation count size overflow".into()))?;
2706 let mut counts = self.memory.alloc::<u8>(count_bytes)?;
2707 self.htod_launch_metadata_sync_copy_into(&vec![0_u8; count_bytes], &mut counts)?;
2708 let empty_column = schemas
2711 .iter()
2712 .any(|schema| schema.arity() != 0)
2713 .then(|| self.memory.alloc::<u8>(0))
2714 .transpose()?;
2715 let mut buffers = Vec::with_capacity(schemas.len());
2716 for (index, schema) in schemas.iter().enumerate() {
2717 let mut columns = Vec::with_capacity(schema.arity());
2718 for _ in 0..schema.arity() {
2719 let column = empty_column
2720 .as_ref()
2721 .expect("non-nullary schema has an empty column owner")
2722 .try_owned_subslice::<u8>(0..0)
2723 .ok_or_else(|| XlogError::Kernel("empty column view is invalid".into()))?;
2724 columns.push(column.into());
2725 }
2726 let start = index
2727 .checked_mul(std::mem::size_of::<u32>())
2728 .ok_or_else(|| XlogError::Kernel("empty relation count offset overflow".into()))?;
2729 let end = start
2730 .checked_add(std::mem::size_of::<u32>())
2731 .ok_or_else(|| XlogError::Kernel("empty relation count offset overflow".into()))?;
2732 let count = counts
2733 .try_owned_subslice::<u32>(start..end)
2734 .ok_or_else(|| XlogError::Kernel("empty relation count view is invalid".into()))?;
2735 let mut buffer =
2736 CudaBuffer::from_columns_with_host_count(columns, 0, count, schema.clone(), 0);
2737 buffer.certify_canonical_full_row_set();
2738 buffers.push(buffer);
2739 }
2740 Ok(buffers)
2741 }
2742
2743 pub fn create_zero_arity_buffer(&self, schema: Schema, rows: u32) -> Result<CudaBuffer> {
2751 debug_assert_eq!(
2752 schema.arity(),
2753 0,
2754 "create_zero_arity_buffer requires arity 0"
2755 );
2756 self.buffer_from_columns(Vec::new(), u64::from(rows), schema)
2757 }
2758
2759 pub(crate) fn buffer_from_columns(
2760 &self,
2761 columns: Vec<CudaColumn>,
2762 row_cap: u64,
2763 schema: Schema,
2764 ) -> Result<CudaBuffer> {
2765 let row_u32 = u32::try_from(row_cap)
2766 .map_err(|_| XlogError::Kernel(format!("Row capacity {} exceeds u32::MAX", row_cap)))?;
2767 let mut d_num_rows = self.memory.alloc::<u32>(1)?;
2768 self.htod_launch_metadata_sync_copy_into(&[row_u32], &mut d_num_rows)
2769 .map_err(|e| XlogError::Kernel(format!("Failed to set row count: {}", e)))?;
2770 let mut result =
2771 CudaBuffer::from_columns_with_host_count(columns, row_cap, d_num_rows, schema, row_u32);
2772 if row_u32 <= 1 {
2773 result.certify_canonical_full_row_set();
2774 }
2775 Ok(result)
2776 }
2777
2778 fn combine_schemas(&self, left: &Schema, right: &Schema) -> Schema {
2780 let mut columns = left.columns.clone();
2781 columns.extend(right.columns.iter().cloned());
2782 let mut sort_labels = left.sort_labels().to_vec();
2783 sort_labels.extend(right.sort_labels().iter().cloned());
2784 Schema::new(columns)
2785 .with_sort_labels(sort_labels)
2786 .expect("combined schema sort labels match column arity")
2787 }
2788
2789 fn schemas_type_compatible(&self, a: &Schema, b: &Schema) -> bool {
2794 if a.arity() != b.arity() {
2795 return false;
2796 }
2797 for i in 0..a.arity() {
2798 if a.column_type(i) != b.column_type(i) {
2799 return false;
2800 }
2801 }
2802 true
2803 }
2804}
2805
2806#[cfg(test)]
2807mod tests {
2808 use super::*;
2809 use crate::device_runtime::{NullSink, XlogDeviceRuntime};
2810 use xlog_core::{AggOp, MemoryBudget, ScalarType};
2811
2812 fn has_cuda_device() -> bool {
2813 CudaDevice::new(0).is_ok()
2814 }
2815
2816 struct ProviderEnvRestore {
2817 name: &'static str,
2818 previous: Option<std::ffi::OsString>,
2819 }
2820
2821 impl ProviderEnvRestore {
2822 fn capture(name: &'static str) -> Self {
2823 Self {
2824 name,
2825 previous: std::env::var_os(name),
2826 }
2827 }
2828 }
2829
2830 impl Drop for ProviderEnvRestore {
2831 fn drop(&mut self) {
2832 unsafe {
2833 match self.previous.take() {
2834 Some(value) => std::env::set_var(self.name, value),
2835 None => std::env::remove_var(self.name),
2836 }
2837 }
2838 }
2839 }
2840
2841 #[test]
2842 fn provider_boolean_environment_uses_strict_shared_semantics() {
2843 const NAME: &str = "XLOG_CUDA_TEST_BOOLEAN";
2844 static LOCK: Mutex<()> = Mutex::new(());
2845 let _guard = LOCK.lock().unwrap();
2846 let _restore = ProviderEnvRestore::capture(NAME);
2847
2848 unsafe {
2849 std::env::set_var(NAME, "yes");
2850 }
2851 assert!(CudaKernelProvider::env_flag(NAME).unwrap());
2852
2853 unsafe {
2854 std::env::set_var(NAME, "invalid");
2855 }
2856 assert!(matches!(
2857 CudaKernelProvider::env_flag(NAME),
2858 Err(XlogError::Configuration { ref name, .. }) if name == NAME
2859 ));
2860 }
2861
2862 #[test]
2863 fn test_kernel_artifact_locator_precedence_order() {
2864 use super::kernel_paths::KernelArtifactLocator;
2865 use std::fs;
2866 use std::path::PathBuf;
2867
2868 let root = std::env::temp_dir().join(format!(
2869 "xlog-kernel-paths-{}-{}",
2870 std::process::id(),
2871 std::time::SystemTime::now()
2872 .duration_since(std::time::UNIX_EPOCH)
2873 .expect("system clock before UNIX_EPOCH")
2874 .as_nanos()
2875 ));
2876 let cubin_dir = root.join("cubin");
2877 let package_dir = root.join("bin").join("kernels");
2878 let out_dir = root.join("out");
2879 fs::create_dir_all(&cubin_dir).expect("create cubin dir");
2880 fs::create_dir_all(&package_dir).expect("create package kernels dir");
2881 fs::create_dir_all(&out_dir).expect("create out dir");
2882
2883 let name = "xlog_join";
2884 let cc = 75;
2885 let cubin_path = cubin_dir.join(format!("{name}.sm_{cc}.cubin"));
2886 let package_path = package_dir.join(format!("{name}.sm_{cc}.cubin"));
2887 let out_path = out_dir.join(format!("{name}.sm_{cc}.cubin"));
2888 fs::write(&cubin_path, b"cubin").expect("write cubin file");
2889 fs::write(&package_path, b"package").expect("write package file");
2890 fs::write(&out_path, b"out").expect("write out file");
2891
2892 let locator = KernelArtifactLocator::new(
2893 Some(cubin_dir.clone()),
2894 Some(package_dir.clone()),
2895 Some(out_dir.clone()),
2896 );
2897
2898 let (path, is_cubin) = locator
2899 .resolve_module_path(name, cc)
2900 .expect("expected a kernel artifact");
2901 assert_eq!(path, cubin_path);
2902 assert!(is_cubin);
2903
2904 fs::remove_file(&cubin_path).expect("remove cubin file");
2905 let (path, is_cubin) = locator
2906 .resolve_module_path(name, cc)
2907 .expect("expected package kernel artifact");
2908 assert_eq!(path, package_path);
2909 assert!(is_cubin);
2910
2911 fs::remove_file(&package_path).expect("remove package file");
2912 let (path, is_cubin) = locator
2913 .resolve_module_path(name, cc)
2914 .expect("expected out dir kernel artifact");
2915 assert_eq!(path, out_path);
2916 assert!(is_cubin);
2917
2918 let _ = fs::remove_dir_all(PathBuf::from(&root));
2919 }
2920
2921 #[test]
2922 fn test_module_resolution_falls_back_to_embedded_portable_ptx() {
2923 use super::kernel_paths::KernelArtifactLocator;
2924
2925 let locator = KernelArtifactLocator::new(None, None, None);
2926 for name in crate::kernel_manifest_data::KERNEL_CU_NAMES {
2927 let sources = resolve_module_sources_with_locator(name, 999, &locator);
2928 assert_eq!(
2929 sources.len(),
2930 1,
2931 "{name}: expected only embedded portable PTX fallback"
2932 );
2933
2934 match &sources[0] {
2935 KernelModuleSource::EmbeddedPortablePtx { ptx } => {
2936 assert!(
2937 ptx.contains(".entry"),
2938 "{name}: embedded PTX should contain CUDA entry points"
2939 );
2940 }
2941 KernelModuleSource::File { path, .. } => {
2942 panic!(
2943 "{name}: expected embedded portable PTX fallback, got file {}",
2944 path.display()
2945 );
2946 }
2947 }
2948 }
2949 }
2950
2951 #[test]
2952 fn test_embedded_portable_ptx_manifest_matches_kernel_manifest() {
2953 let embedded_names: std::collections::BTreeSet<_> =
2954 crate::embedded_kernel_data::EMBEDDED_PORTABLE_PTX
2955 .iter()
2956 .map(|artifact| artifact.name)
2957 .collect();
2958 let manifest_names: std::collections::BTreeSet<_> =
2959 crate::kernel_manifest_data::KERNEL_CU_NAMES
2960 .iter()
2961 .copied()
2962 .collect();
2963
2964 assert_eq!(
2965 embedded_names, manifest_names,
2966 "embedded portable PTX table should cover every runtime kernel module"
2967 );
2968 }
2969
2970 #[test]
2971 fn test_kernel_provider_creation() {
2972 if !has_cuda_device() {
2973 eprintln!("Skipping test: no CUDA device available");
2974 return;
2975 }
2976
2977 let budget = MemoryBudget::with_limit(1024 * 1024 * 1024); let provider = crate::CudaProviderBuilder::new(0, budget).build();
2979 assert!(
2980 provider.is_ok(),
2981 "Failed to create kernel provider: {:?}",
2982 provider.err()
2983 );
2984
2985 let provider = provider.unwrap();
2986 assert!(Arc::ptr_eq(provider.device(), provider.memory().device()));
2987 }
2988
2989 #[test]
2990 fn test_kernel_functions_accessible() {
2991 if !has_cuda_device() {
2992 eprintln!("Skipping test: no CUDA device available");
2993 return;
2994 }
2995
2996 let budget = MemoryBudget::with_limit(1024 * 1024 * 1024);
2997 let provider = crate::CudaProviderBuilder::new(0, budget)
2998 .build()
2999 .expect("Failed to create provider");
3000
3001 let inner = provider.device().inner();
3003
3004 let build_fn = inner.get_func(JOIN_MODULE, join_kernels::HASH_JOIN_BUILD);
3006 assert!(
3007 build_fn.is_some(),
3008 "hash_join_build function should be accessible"
3009 );
3010
3011 let probe_fn = inner.get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE);
3012 assert!(
3013 probe_fn.is_some(),
3014 "hash_join_probe function should be accessible"
3015 );
3016
3017 let mark_fn = inner.get_func(DEDUP_MODULE, dedup_kernels::MARK_DUPLICATES);
3019 assert!(
3020 mark_fn.is_some(),
3021 "mark_duplicates function should be accessible"
3022 );
3023
3024 let compact_fn = inner.get_func(DEDUP_MODULE, dedup_kernels::COMPACT_ROWS);
3025 assert!(
3026 compact_fn.is_some(),
3027 "compact_rows function should be accessible"
3028 );
3029
3030 let boundaries_fn =
3032 inner.get_func(GROUPBY_MODULE, groupby_kernels::DETECT_GROUP_BOUNDARIES);
3033 assert!(
3034 boundaries_fn.is_some(),
3035 "detect_group_boundaries function should be accessible"
3036 );
3037
3038 let count_fn = inner.get_func(GROUPBY_MODULE, groupby_kernels::GROUPBY_COUNT);
3039 assert!(
3040 count_fn.is_some(),
3041 "groupby_count function should be accessible"
3042 );
3043
3044 let sum_fn = inner.get_func(GROUPBY_MODULE, groupby_kernels::GROUPBY_SUM);
3045 assert!(
3046 sum_fn.is_some(),
3047 "groupby_sum function should be accessible"
3048 );
3049
3050 let min_fn = inner.get_func(GROUPBY_MODULE, groupby_kernels::GROUPBY_MIN);
3051 assert!(
3052 min_fn.is_some(),
3053 "groupby_min function should be accessible"
3054 );
3055
3056 let max_fn = inner.get_func(GROUPBY_MODULE, groupby_kernels::GROUPBY_MAX);
3057 assert!(
3058 max_fn.is_some(),
3059 "groupby_max function should be accessible"
3060 );
3061
3062 let xgcf_forward = inner.get_func(CIRCUIT_MODULE, "xgcf_forward_level");
3064 assert!(
3065 xgcf_forward.is_some(),
3066 "xgcf_forward_level function should be accessible"
3067 );
3068
3069 let xgcf_backward_propagate =
3070 inner.get_func(CIRCUIT_MODULE, "xgcf_backward_level_propagate");
3071 assert!(
3072 xgcf_backward_propagate.is_some(),
3073 "xgcf_backward_level_propagate function should be accessible"
3074 );
3075
3076 let xgcf_backward_decision_grad =
3077 inner.get_func(CIRCUIT_MODULE, "xgcf_backward_level_decision_grad");
3078 assert!(
3079 xgcf_backward_decision_grad.is_some(),
3080 "xgcf_backward_level_decision_grad function should be accessible"
3081 );
3082
3083 let xgcf_backward_lit_grad = inner.get_func(CIRCUIT_MODULE, "xgcf_backward_level_lit_grad");
3084 assert!(
3085 xgcf_backward_lit_grad.is_some(),
3086 "xgcf_backward_level_lit_grad function should be accessible"
3087 );
3088
3089 let neural_fill = inner.get_func("xlog_neural", "neural_fill_ad_chain_f32");
3091 assert!(
3092 neural_fill.is_some(),
3093 "neural_fill_ad_chain_f32 function should be accessible"
3094 );
3095 let neural_scatter = inner.get_func("xlog_neural", "neural_scatter_ad_chain_grads_f32");
3096 assert!(
3097 neural_scatter.is_some(),
3098 "neural_scatter_ad_chain_grads_f32 function should be accessible"
3099 );
3100 }
3101
3102 #[test]
3103 fn test_module_names_unique() {
3104 assert_ne!(JOIN_MODULE, DEDUP_MODULE);
3106 assert_ne!(JOIN_MODULE, GROUPBY_MODULE);
3107 assert_ne!(DEDUP_MODULE, GROUPBY_MODULE);
3108 }
3109
3110 fn create_test_provider() -> Option<CudaKernelProvider> {
3112 let budget = MemoryBudget::with_limit(1024 * 1024 * 1024);
3113 crate::CudaProviderBuilder::new(0, budget).build().ok()
3114 }
3115
3116 fn create_test_provider_with_runtime() -> Option<(CudaKernelProvider, Arc<XlogDeviceRuntime>)> {
3117 if !has_cuda_device() {
3118 return None;
3119 }
3120 let sink = Arc::new(NullSink::new());
3121 let provider =
3122 crate::CudaProviderBuilder::new(0, MemoryBudget::with_limit(1024 * 1024 * 1024))
3123 .with_logging_sink(sink)
3124 .build()
3125 .ok()?;
3126 let runtime = Arc::clone(provider.memory().runtime()?);
3127 Some((provider, runtime))
3128 }
3129
3130 #[test]
3131 fn test_recorded_join_index_build_runs_on_runtime_stream() {
3132 let (provider, runtime) = match create_test_provider_with_runtime() {
3133 Some(fixture) => fixture,
3134 None => {
3135 eprintln!("Skipping test: no CUDA device available");
3136 return;
3137 }
3138 };
3139 let stream = runtime.stream_pool().acquire().expect("recorded stream");
3140 let left = create_test_buffer(&provider, &[1, 2, 3, 4], "key");
3141 let right = create_test_buffer(&provider, &[1, 2, 3, 4], "key");
3142
3143 let index = provider
3144 .build_join_index_v2_recorded(&right, &[0], stream)
3145 .expect("recorded join-index build");
3146 let joined = provider
3147 .hash_join_v2_with_index_recorded(
3148 &left,
3149 &right,
3150 &[0],
3151 &[0],
3152 JoinType::Inner,
3153 &index,
3154 None,
3155 stream,
3156 )
3157 .expect("recorded indexed join consumes recorded build");
3158 runtime
3159 .stream_pool()
3160 .resolve(stream)
3161 .expect("stream resolves")
3162 .synchronize()
3163 .expect("recorded stream synchronized");
3164
3165 assert_eq!(index.right_num_rows(), 4);
3166 assert_eq!(index.right_keys(), &[0]);
3167 assert_eq!(provider.device_row_count(&joined).expect("joined rows"), 4);
3168 }
3169
3170 fn create_test_buffer(
3172 provider: &CudaKernelProvider,
3173 data: &[u32],
3174 col_name: &str,
3175 ) -> CudaBuffer {
3176 let schema = Schema::new(vec![(col_name.to_string(), ScalarType::U32)]);
3177 let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
3178
3179 let mut col = provider.memory().alloc::<u8>(bytes.len()).expect("alloc");
3180 provider
3181 .device()
3182 .inner()
3183 .htod_sync_copy_into(&bytes, &mut col)
3184 .expect("htod");
3185
3186 provider
3187 .buffer_from_columns(vec![col.into()], data.len() as u64, schema)
3188 .expect("buffer")
3189 }
3190
3191 fn create_empty_test_buffer(provider: &CudaKernelProvider, schema: Schema) -> CudaBuffer {
3193 let mut columns = Vec::with_capacity(schema.arity());
3194 for _ in 0..schema.arity() {
3195 columns.push(provider.memory().alloc::<u8>(0).expect("alloc").into());
3196 }
3197 provider
3198 .buffer_from_columns(columns, 0, schema)
3199 .expect("buffer")
3200 }
3201
3202 fn read_buffer_u32(provider: &CudaKernelProvider, buffer: &CudaBuffer, col: usize) -> Vec<u32> {
3204 if buffer.is_empty() || buffer.column(col).is_none() {
3205 return vec![];
3206 }
3207 let num_rows = buffer.num_rows() as usize;
3208 let mut bytes = vec![0u8; num_rows * 4];
3209 provider
3210 .device()
3211 .inner()
3212 .dtoh_sync_copy_into(buffer.column(col).unwrap(), &mut bytes)
3213 .expect("dtoh");
3214 bytes
3215 .as_chunks::<4>()
3216 .0
3217 .iter()
3218 .map(|c| u32::from_le_bytes([c[0], c[1], c[2], c[3]]))
3219 .collect()
3220 }
3221
3222 #[test]
3223 fn test_compact_device_mask_respects_mask_len_smaller_than_row_cap() {
3224 let provider = match create_test_provider() {
3225 Some(p) => p,
3226 None => {
3227 eprintln!("Skipping test: no CUDA device available");
3228 return;
3229 }
3230 };
3231
3232 let schema = Schema::new(vec![("id".to_string(), ScalarType::U32)]);
3233 let base = create_test_buffer(&provider, &[1, 2, 3, 4, 5, 6, 7, 8], "id");
3234
3235 let row_cap = 16u64;
3236 let data: Vec<u32> = (0..row_cap as u32).collect();
3237 let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
3238 let mut col = provider.memory().alloc::<u8>(bytes.len()).expect("alloc");
3239 provider
3240 .device()
3241 .inner()
3242 .htod_sync_copy_into(&bytes, &mut col)
3243 .expect("htod");
3244 let expanded = provider
3245 .buffer_from_columns_with_device_count(vec![col.into()], row_cap, schema, &base)
3246 .expect("buffer");
3247
3248 let mask: Vec<u8> = vec![1, 0, 1, 0, 1, 0, 1, 0];
3249 let (prefix_sum, count) = provider.prefix_sum_mask(&mask).expect("prefix sum");
3250
3251 let mut d_mask = provider.memory().alloc::<u8>(mask.len()).expect("alloc");
3252 provider
3253 .device()
3254 .inner()
3255 .htod_sync_copy_into(&mask, &mut d_mask)
3256 .expect("mask htod");
3257
3258 let mut d_prefix = provider
3259 .memory()
3260 .alloc::<u32>(prefix_sum.len())
3261 .expect("alloc");
3262 provider
3263 .device()
3264 .inner()
3265 .htod_sync_copy_into(&prefix_sum, &mut d_prefix)
3266 .expect("prefix htod");
3267
3268 let mut d_out_count = provider.memory().alloc::<u32>(1).expect("alloc");
3269 provider
3270 .device()
3271 .inner()
3272 .htod_sync_copy_into(&[count], &mut d_out_count)
3273 .expect("count htod");
3274
3275 let compacted = provider
3276 .compact_buffer_by_device_mask_device_count(&expanded, &d_mask, &d_prefix, d_out_count)
3277 .expect("compact");
3278
3279 assert_eq!(compacted.num_rows(), mask.len() as u64);
3280 let device_rows = provider.device_row_count(&compacted).expect("row count");
3281 assert_eq!(device_rows as u32, count);
3282 }
3283
3284 #[test]
3285 fn test_clone_buffer_preserves_device_count() {
3286 let provider = match create_test_provider() {
3287 Some(p) => p,
3288 None => {
3289 eprintln!("Skipping test: no CUDA device available");
3290 return;
3291 }
3292 };
3293
3294 let schema = Schema::new(vec![("id".to_string(), ScalarType::U32)]);
3295 let ids: Vec<u32> = vec![10, 20, 30];
3296 let buffer = provider
3297 .create_buffer_from_slices(&[bytemuck::cast_slice(&ids)], schema)
3298 .unwrap();
3299
3300 let cloned = provider.clone_buffer(&buffer).unwrap();
3301
3302 let mut host_count = [0u32];
3303 provider
3304 .device()
3305 .inner()
3306 .dtoh_sync_copy_into(cloned.num_rows_device(), &mut host_count)
3307 .unwrap();
3308 assert_eq!(host_count[0], 3);
3309 }
3310
3311 #[test]
3320 fn test_clone_buffer_preserves_cached_row_count() {
3321 let provider = match create_test_provider() {
3322 Some(p) => p,
3323 None => {
3324 eprintln!("Skipping test: no CUDA device available");
3325 return;
3326 }
3327 };
3328
3329 let schema = Schema::new(vec![("id".to_string(), ScalarType::U32)]);
3330 let ids: Vec<u32> = vec![7, 11, 13, 17];
3331 let source = provider
3332 .create_buffer_from_slices(&[bytemuck::cast_slice(&ids)], schema)
3333 .unwrap();
3334 assert_eq!(
3338 source.cached_row_count(),
3339 Some(4),
3340 "source buffer should have its cached row count populated by \
3341 create_buffer_from_slices"
3342 );
3343
3344 let cloned = provider.clone_buffer(&source).unwrap();
3345
3346 assert_eq!(
3347 cloned.cached_row_count(),
3348 Some(4),
3349 "clone_buffer must propagate cached_row_count from source to clone",
3350 );
3351 }
3352
3353 #[test]
3354 fn certified_singleton_union_avoids_cold_row_count_read() {
3355 let provider = match create_test_provider() {
3356 Some(p) => p,
3357 None => {
3358 eprintln!("Skipping test: no CUDA device available");
3359 return;
3360 }
3361 };
3362
3363 let source = create_test_buffer(&provider, &[1, 2, 3], "id");
3364 let CudaBuffer {
3365 columns,
3366 row_cap,
3367 d_num_rows,
3368 schema,
3369 ..
3370 } = source;
3371 let mut cold_set = CudaBuffer::from_columns(columns, row_cap, d_num_rows, schema);
3372 cold_set.certify_canonical_full_row_set();
3373 assert_eq!(cold_set.cached_row_count(), None);
3374 let empty = provider
3375 .create_empty_buffer(cold_set.schema().clone())
3376 .unwrap();
3377
3378 provider.reset_host_transfer_stats();
3379 provider.reset_untracked_metadata_dtoh_count();
3380 provider.memory().reset_alloc_count();
3381 let result = provider
3382 .union_many_gpu(&[&empty, &cold_set, &empty])
3383 .unwrap();
3384
3385 assert_eq!(
3386 cold_set.cached_row_count(),
3387 None,
3388 "the certified fast path must not populate the cold count via D2H"
3389 );
3390 assert_eq!(provider.host_transfer_stats().dtoh_calls, 0);
3391 assert_eq!(provider.untracked_metadata_dtoh_count(), 0);
3392 assert_eq!(provider.memory().alloc_count(), 2);
3393 assert!(result.canonical_full_row_set_certified());
3394 assert_eq!(read_buffer_u32(&provider, &result, 0), vec![1, 2, 3]);
3395 }
3396
3397 #[test]
3398 fn cold_row_count_read_uses_audited_metadata_counter_once() {
3399 let provider = match create_test_provider() {
3400 Some(p) => p,
3401 None => {
3402 eprintln!("Skipping test: no CUDA device available");
3403 return;
3404 }
3405 };
3406
3407 let source = create_test_buffer(&provider, &[1, 2, 3], "id");
3408 let CudaBuffer {
3409 columns,
3410 row_cap,
3411 d_num_rows,
3412 schema,
3413 ..
3414 } = source;
3415 let cold_bag = CudaBuffer::from_columns(columns, row_cap, d_num_rows, schema);
3416
3417 provider.reset_untracked_metadata_dtoh_count();
3418 assert_eq!(provider.device_row_count(&cold_bag).unwrap(), 3);
3419
3420 assert_eq!(provider.untracked_metadata_dtoh_count(), 1);
3421 assert_eq!(cold_bag.cached_row_count(), Some(3));
3422 provider.reset_untracked_metadata_dtoh_count();
3423 let result = provider.union_many_gpu(&[&cold_bag]).unwrap();
3424 assert_eq!(provider.untracked_metadata_dtoh_count(), 1);
3425 assert_eq!(read_buffer_u32(&provider, &result, 0), vec![1, 2, 3]);
3426 }
3427
3428 #[test]
3429 fn zero_arity_certificate_does_not_bypass_unit_set_semantics() {
3430 let provider = match create_test_provider() {
3431 Some(p) => p,
3432 None => {
3433 eprintln!("Skipping test: no CUDA device available");
3434 return;
3435 }
3436 };
3437
3438 let d_num_rows = provider.upload_device_row_count(3).unwrap();
3439 let mut multiplicity = CudaBuffer::from_columns(
3440 Vec::new(),
3441 3,
3442 d_num_rows,
3443 xlog_core::Schema::new(Vec::new()),
3444 );
3445 multiplicity.certify_canonical_full_row_set();
3446
3447 let unit = provider.union_many_gpu(&[&multiplicity]).unwrap();
3448 assert_eq!(provider.device_row_count(&unit).unwrap(), 1);
3449 }
3450
3451 #[test]
3454 fn test_hash_join_empty_inputs() {
3455 let provider = match create_test_provider() {
3456 Some(p) => p,
3457 None => {
3458 eprintln!("Skipping test: no CUDA device available");
3459 return;
3460 }
3461 };
3462
3463 let schema = Schema::new(vec![("key".to_string(), ScalarType::U32)]);
3464 let empty = create_empty_test_buffer(&provider, schema.clone());
3465
3466 let result = provider.hash_join(&empty, &empty, &[0], &[0]);
3468 assert!(result.is_ok());
3469 assert!(result.unwrap().is_empty());
3470 }
3471
3472 #[test]
3473 fn test_hash_join_validation() {
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 left = create_test_buffer(&provider, &[1, 2, 3], "left_key");
3483 let right = create_test_buffer(&provider, &[2, 3, 4], "right_key");
3484
3485 let result = provider.hash_join(&left, &right, &[], &[0]);
3487 assert!(result.is_err());
3488
3489 let result = provider.hash_join(&left, &right, &[0], &[0, 0]);
3491 assert!(result.is_err());
3492 }
3493
3494 #[test]
3497 fn test_dedup_empty_input() {
3498 let provider = match create_test_provider() {
3499 Some(p) => p,
3500 None => {
3501 eprintln!("Skipping test: no CUDA device available");
3502 return;
3503 }
3504 };
3505
3506 let schema = Schema::new(vec![("key".to_string(), ScalarType::U32)]);
3507 let empty = create_empty_test_buffer(&provider, schema);
3508
3509 let result = provider.dedup(&empty, &[0]);
3510 assert!(result.is_ok());
3511 assert!(result.unwrap().is_empty());
3512 }
3513
3514 #[test]
3515 fn test_dedup_validation() {
3516 let provider = match create_test_provider() {
3517 Some(p) => p,
3518 None => {
3519 eprintln!("Skipping test: no CUDA device available");
3520 return;
3521 }
3522 };
3523
3524 let buffer = create_test_buffer(&provider, &[1, 1, 2, 2, 3], "key");
3525
3526 let result = provider.dedup(&buffer, &[]);
3528 assert!(result.is_err());
3529 }
3530
3531 #[test]
3532 fn test_dedup_with_duplicates() {
3533 let provider = match create_test_provider() {
3534 Some(p) => p,
3535 None => {
3536 eprintln!("Skipping test: no CUDA device available");
3537 return;
3538 }
3539 };
3540
3541 let buffer = create_test_buffer(&provider, &[3, 1, 2, 1, 3, 2], "key");
3543 let deduped = provider.dedup(&buffer, &[0]).unwrap();
3544
3545 let dedup_count = provider
3546 .device_row_count(&deduped)
3547 .expect("read dedup row count");
3548 assert_eq!(dedup_count, 3, "Should have 3 unique values");
3549
3550 let result = provider.download_column::<u32>(&deduped, 0).unwrap();
3551 assert_eq!(result, vec![1, 2, 3]);
3553 }
3554
3555 #[test]
3556 fn test_dedup_larger_input() {
3557 let provider = match create_test_provider() {
3558 Some(p) => p,
3559 None => {
3560 eprintln!("Skipping test: no CUDA device available");
3561 return;
3562 }
3563 };
3564
3565 let a: Vec<u32> = (0..500).collect();
3567 let b: Vec<u32> = (250..750).collect();
3568 let input: Vec<u32> = a.iter().chain(b.iter()).copied().collect();
3569
3570 let buffer = create_test_buffer(&provider, &input, "key");
3571 let deduped = provider.dedup(&buffer, &[0]).unwrap();
3572
3573 let dedup_count = provider
3574 .device_row_count(&deduped)
3575 .expect("read dedup row count");
3576 assert_eq!(dedup_count, 750, "Should have 750 unique values (0..750)");
3577
3578 let result = provider.download_column::<u32>(&deduped, 0).unwrap();
3580 let is_sorted = result.windows(2).all(|w| w[0] <= w[1]);
3581 assert!(is_sorted, "Output should be sorted");
3582
3583 let expected: Vec<u32> = (0..750).collect();
3585 assert_eq!(result, expected);
3586 }
3587
3588 #[test]
3591 fn test_union_empty_inputs() {
3592 let provider = match create_test_provider() {
3593 Some(p) => p,
3594 None => {
3595 eprintln!("Skipping test: no CUDA device available");
3596 return;
3597 }
3598 };
3599
3600 let schema = Schema::new(vec![("key".to_string(), ScalarType::U32)]);
3601 let empty = create_empty_test_buffer(&provider, schema.clone());
3602
3603 let result = provider.union(&empty, &empty);
3605 assert!(result.is_ok());
3606 assert!(result.unwrap().is_empty());
3607
3608 let a = create_test_buffer(&provider, &[1, 2, 3], "key");
3610 let empty2 = create_empty_test_buffer(&provider, schema);
3611 let result = provider.union(&a, &empty2);
3612 assert!(result.is_ok());
3613 let result = result.unwrap();
3614 assert_eq!(result.num_rows(), 3);
3615 }
3616
3617 #[test]
3618 fn test_union_schema_type_mismatch() {
3619 let provider = match create_test_provider() {
3620 Some(p) => p,
3621 None => {
3622 eprintln!("Skipping test: no CUDA device available");
3623 return;
3624 }
3625 };
3626
3627 let a = create_test_buffer(&provider, &[1, 2], "col_a");
3628 let b = create_test_buffer(&provider, &[3, 4], "col_b");
3629
3630 let result = provider.union(&a, &b);
3632 assert!(result.is_ok());
3633
3634 let two_col_schema = Schema::new(vec![
3636 ("x".to_string(), ScalarType::U32),
3637 ("y".to_string(), ScalarType::U32),
3638 ]);
3639 let c = provider
3640 .create_buffer_from_u32_columns(&[&[1, 2], &[3, 4]], two_col_schema)
3641 .unwrap();
3642 let result = provider.union(&a, &c);
3643 assert!(result.is_err());
3644 }
3645
3646 #[test]
3649 fn test_diff_empty_inputs() {
3650 let provider = match create_test_provider() {
3651 Some(p) => p,
3652 None => {
3653 eprintln!("Skipping test: no CUDA device available");
3654 return;
3655 }
3656 };
3657
3658 let schema = Schema::new(vec![("key".to_string(), ScalarType::U32)]);
3659 let empty = create_empty_test_buffer(&provider, schema.clone());
3660
3661 let result = provider.diff(&empty, &empty);
3663 assert!(result.is_ok());
3664 assert!(result.unwrap().is_empty());
3665
3666 let a = create_test_buffer(&provider, &[1, 2, 3], "key");
3668 let empty2 = create_empty_test_buffer(&provider, schema);
3669 let result = provider.diff(&a, &empty2);
3670 assert!(result.is_ok());
3671 let result = result.unwrap();
3672 assert_eq!(result.num_rows(), 3);
3673 }
3674
3675 #[test]
3676 fn test_diff_basic() {
3677 let provider = match create_test_provider() {
3678 Some(p) => p,
3679 None => {
3680 eprintln!("Skipping test: no CUDA device available");
3681 return;
3682 }
3683 };
3684
3685 let a = create_test_buffer(&provider, &[1, 2, 3, 4, 5], "key");
3686 let b = create_test_buffer(&provider, &[2, 4], "key");
3687
3688 let result = provider.diff(&a, &b);
3689 assert!(result.is_ok());
3690 let result = result.unwrap();
3691 assert_eq!(result.num_rows(), 3); let values = read_buffer_u32(&provider, &result, 0);
3694 assert_eq!(values, vec![1, 3, 5]);
3695 }
3696
3697 #[test]
3698 fn test_diff_all_filtered_out() {
3699 let provider = match create_test_provider() {
3700 Some(p) => p,
3701 None => {
3702 eprintln!("Skipping test: no CUDA device available");
3703 return;
3704 }
3705 };
3706
3707 let a = create_test_buffer(&provider, &[1, 2, 3], "key");
3708 let b = create_test_buffer(&provider, &[1, 2, 3, 4, 5], "key");
3709
3710 let result = provider.diff(&a, &b);
3711 assert!(result.is_ok());
3712 assert!(result.unwrap().is_empty());
3713 }
3714
3715 #[test]
3716 fn test_diff_schema_mismatch() {
3717 let provider = match create_test_provider() {
3718 Some(p) => p,
3719 None => {
3720 eprintln!("Skipping test: no CUDA device available");
3721 return;
3722 }
3723 };
3724
3725 let a = create_test_buffer(&provider, &[1, 2], "col_a");
3727 let b = create_test_buffer(&provider, &[1, 2], "col_b");
3728 let result = provider.diff(&a, &b);
3729 assert!(
3730 result.is_ok(),
3731 "Same types with different names should succeed"
3732 );
3733
3734 let schema_2col = Schema::new(vec![
3736 ("c0".to_string(), ScalarType::U32),
3737 ("c1".to_string(), ScalarType::U32),
3738 ]);
3739
3740 let bytes_2col: Vec<u8> = [1u32, 2, 3, 4]
3741 .iter()
3742 .flat_map(|v| v.to_le_bytes())
3743 .collect();
3744 let mut col0 = provider
3745 .memory()
3746 .alloc::<u8>(bytes_2col.len() / 2)
3747 .expect("alloc");
3748 let mut col1 = provider
3749 .memory()
3750 .alloc::<u8>(bytes_2col.len() / 2)
3751 .expect("alloc");
3752 provider
3753 .device()
3754 .inner()
3755 .htod_sync_copy_into(&bytes_2col[..8], &mut col0)
3756 .expect("htod");
3757 provider
3758 .device()
3759 .inner()
3760 .htod_sync_copy_into(&bytes_2col[8..], &mut col1)
3761 .expect("htod");
3762 let buffer_2col = provider
3763 .buffer_from_columns(vec![col0.into(), col1.into()], 2, schema_2col)
3764 .expect("buffer");
3765
3766 let buffer_1col = create_test_buffer(&provider, &[1, 2], "c0");
3767
3768 let result = provider.diff(&buffer_2col, &buffer_1col);
3769 assert!(result.is_err(), "Different arities should fail");
3770 }
3771
3772 #[test]
3775 fn test_groupby_empty_input() {
3776 let provider = match create_test_provider() {
3777 Some(p) => p,
3778 None => {
3779 eprintln!("Skipping test: no CUDA device available");
3780 return;
3781 }
3782 };
3783
3784 let schema = Schema::new(vec![("key".to_string(), ScalarType::U32)]);
3785 let empty = create_empty_test_buffer(&provider, schema);
3786
3787 let result = provider.groupby_agg(&empty, &[0], AggOp::Count, 0);
3788 assert!(result.is_ok());
3789 assert!(result.unwrap().is_empty());
3790 }
3791
3792 #[test]
3793 fn test_groupby_validation() {
3794 let provider = match create_test_provider() {
3795 Some(p) => p,
3796 None => {
3797 eprintln!("Skipping test: no CUDA device available");
3798 return;
3799 }
3800 };
3801
3802 let buffer = create_test_buffer(&provider, &[1, 1, 2, 2, 3], "key");
3803
3804 let result = provider.groupby_agg(&buffer, &[], AggOp::Count, 0);
3806 assert!(result.is_err());
3807
3808 let result = provider.groupby_agg(&buffer, &[0], AggOp::Count, 5);
3810 assert!(result.is_err());
3811 }
3812
3813 #[test]
3814 fn test_groupby_logsumexp() {
3815 let provider = match create_test_provider() {
3816 Some(p) => p,
3817 None => {
3818 eprintln!("Skipping test: no CUDA device available");
3819 return;
3820 }
3821 };
3822
3823 let keys: Vec<u32> = vec![1, 1, 2, 2];
3827 let values: Vec<f64> = vec![1.0, 2.0, 3.0, 4.0];
3828
3829 let schema = Schema::new(vec![
3830 ("key".to_string(), ScalarType::U32),
3831 ("value".to_string(), ScalarType::F64),
3832 ]);
3833
3834 let key_bytes: Vec<u8> = keys.iter().flat_map(|v| v.to_le_bytes()).collect();
3836 let mut key_col = provider
3837 .memory()
3838 .alloc::<u8>(key_bytes.len())
3839 .expect("alloc key");
3840 provider
3841 .device()
3842 .inner()
3843 .htod_sync_copy_into(&key_bytes, &mut key_col)
3844 .expect("upload key");
3845
3846 let val_bytes: Vec<u8> = values.iter().flat_map(|v| v.to_le_bytes()).collect();
3848 let mut val_col = provider
3849 .memory()
3850 .alloc::<u8>(val_bytes.len())
3851 .expect("alloc val");
3852 provider
3853 .device()
3854 .inner()
3855 .htod_sync_copy_into(&val_bytes, &mut val_col)
3856 .expect("upload val");
3857
3858 let buffer = provider
3859 .buffer_from_columns(vec![key_col.into(), val_col.into()], 4, schema)
3860 .expect("buffer");
3861
3862 let result = provider.groupby_agg(&buffer, &[0], AggOp::LogSumExp, 1);
3864 assert!(
3865 result.is_ok(),
3866 "groupby_agg with LogSumExp should succeed: {:?}",
3867 result.err()
3868 );
3869
3870 let result = result.unwrap();
3871 let group_count = provider
3872 .device_row_count(&result)
3873 .expect("read group count");
3874 assert_eq!(group_count, 2, "Should have 2 groups");
3875
3876 let result_values = provider
3878 .download_column::<f64>(&result, 1)
3879 .expect("download result");
3880
3881 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;
3888 assert!(
3889 (result_values[0] - expected_0).abs() < tolerance,
3890 "Group 0 logsumexp mismatch: got {}, expected {}",
3891 result_values[0],
3892 expected_0
3893 );
3894 assert!(
3895 (result_values[1] - expected_1).abs() < tolerance,
3896 "Group 1 logsumexp mismatch: got {}, expected {}",
3897 result_values[1],
3898 expected_1
3899 );
3900 }
3901
3902 #[test]
3905 fn test_combine_schemas() {
3906 let provider = match create_test_provider() {
3907 Some(p) => p,
3908 None => {
3909 eprintln!("Skipping test: no CUDA device available");
3910 return;
3911 }
3912 };
3913
3914 let left = Schema::new(vec![("a".to_string(), ScalarType::U32)]);
3915 let right = Schema::new(vec![("b".to_string(), ScalarType::U64)]);
3916
3917 let combined = provider.combine_schemas(&left, &right);
3918 assert_eq!(combined.arity(), 2);
3919 assert_eq!(combined.column_type(0), Some(ScalarType::U32));
3920 assert_eq!(combined.column_type(1), Some(ScalarType::U64));
3921 }
3922
3923 #[test]
3924 fn test_groupby_result_schema() {
3925 let provider = match create_test_provider() {
3926 Some(p) => p,
3927 None => {
3928 eprintln!("Skipping test: no CUDA device available");
3929 return;
3930 }
3931 };
3932
3933 let input = Schema::new(vec![
3934 ("key".to_string(), ScalarType::U32),
3935 ("value".to_string(), ScalarType::U32),
3936 ]);
3937
3938 let count_schema =
3940 provider.groupby_multi_agg_result_schema(&input, &[0], &[(1, AggOp::Count)]);
3941 assert_eq!(count_schema.arity(), 2);
3942 assert_eq!(count_schema.column_type(1), Some(ScalarType::U64));
3943
3944 let sum_schema = provider.groupby_multi_agg_result_schema(&input, &[0], &[(1, AggOp::Sum)]);
3946 assert_eq!(sum_schema.arity(), 2);
3947 assert_eq!(sum_schema.column_type(1), Some(ScalarType::U64));
3948
3949 let min_schema = provider.groupby_multi_agg_result_schema(&input, &[0], &[(1, AggOp::Min)]);
3951 assert_eq!(min_schema.arity(), 2);
3952 assert_eq!(min_schema.column_type(1), Some(ScalarType::U32));
3953 }
3954
3955 #[test]
3956 fn test_groupby_multi_agg_sum_returns_u64_schema() {
3957 let provider = match create_test_provider() {
3958 Some(p) => p,
3959 None => {
3960 eprintln!("Skipping test: no CUDA device");
3961 return;
3962 }
3963 };
3964
3965 let schema = Schema::new(vec![
3966 ("key".to_string(), ScalarType::U32),
3967 ("val".to_string(), ScalarType::U32),
3968 ]);
3969
3970 let result_schema =
3971 provider.groupby_multi_agg_result_schema(&schema, &[0], &[(1, AggOp::Sum)]);
3972
3973 assert_eq!(
3975 result_schema.column_type(1),
3976 Some(ScalarType::U64),
3977 "Sum aggregation should return U64 type, not U32"
3978 );
3979 }
3980
3981 #[test]
3982 fn test_join_custom_max_output() {
3983 let provider = match create_test_provider() {
3984 Some(p) => p,
3985 None => {
3986 eprintln!("Skipping test: no CUDA device available");
3987 return;
3988 }
3989 };
3990
3991 let left = create_test_buffer(&provider, &[1, 1, 1, 1, 2, 2, 2, 2], "left_key");
3996 let right = create_test_buffer(&provider, &[1, 1, 1, 2, 2, 2], "right_key");
3997
3998 let result_limited = provider
4000 .hash_join_v2_with_limit(&left, &right, &[0], &[0], JoinType::Inner, Some(10))
4001 .expect("join with limit should succeed");
4002 assert!(
4003 result_limited.num_rows() <= 10,
4004 "With limit 10, got {} rows but expected at most 10",
4005 result_limited.num_rows()
4006 );
4007
4008 let result_unlimited = provider
4010 .hash_join_v2_with_limit(&left, &right, &[0], &[0], JoinType::Inner, None)
4011 .expect("join without limit should succeed");
4012 assert_eq!(
4013 result_unlimited.num_rows(),
4014 24,
4015 "Without limit, expected 24 rows but got {}",
4016 result_unlimited.num_rows()
4017 );
4018
4019 let result_legacy = provider
4021 .hash_join_v2(&left, &right, &[0], &[0], JoinType::Inner)
4022 .expect("legacy hash_join_v2 should succeed");
4023 assert_eq!(
4024 result_legacy.num_rows(),
4025 24,
4026 "Legacy API without limit, expected 24 rows but got {}",
4027 result_legacy.num_rows()
4028 );
4029 }
4030
4031 fn create_arith_test_provider() -> Option<CudaKernelProvider> {
4035 let budget = MemoryBudget::with_limit(1024 * 1024 * 1024);
4036 crate::CudaProviderBuilder::new(0, budget).build().ok()
4037 }
4038
4039 fn create_i64_buffer(provider: &CudaKernelProvider, data: &[i64]) -> CudaBuffer {
4041 let schema = Schema::new(vec![("col".to_string(), ScalarType::I64)]);
4042 provider
4043 .create_buffer_from_slice::<i64>(data, schema)
4044 .unwrap()
4045 }
4046
4047 fn create_f64_buffer(provider: &CudaKernelProvider, data: &[f64]) -> CudaBuffer {
4049 let schema = Schema::new(vec![("col".to_string(), ScalarType::F64)]);
4050 provider
4051 .create_buffer_from_slice::<f64>(data, schema)
4052 .unwrap()
4053 }
4054
4055 #[test]
4056 fn test_add_columns_i64() {
4057 let Some(provider) = create_arith_test_provider() else {
4058 eprintln!("Skipping test: no CUDA device available");
4059 return;
4060 };
4061
4062 let a = create_i64_buffer(&provider, &[1, 2, 3, 4, 5]);
4063 let b = create_i64_buffer(&provider, &[10, 20, 30, 40, 50]);
4064
4065 let result = provider.add_columns(&a, &b).unwrap();
4066 let values = provider.download_column::<i64>(&result, 0).unwrap();
4067
4068 assert_eq!(values, vec![11, 22, 33, 44, 55]);
4069 }
4070
4071 #[test]
4072 fn test_sub_columns_i64() {
4073 let Some(provider) = create_arith_test_provider() else {
4074 eprintln!("Skipping test: no CUDA device available");
4075 return;
4076 };
4077
4078 let a = create_i64_buffer(&provider, &[10, 20, 30, 40, 50]);
4079 let b = create_i64_buffer(&provider, &[1, 2, 3, 4, 5]);
4080
4081 let result = provider.sub_columns(&a, &b).unwrap();
4082 let values = provider.download_column::<i64>(&result, 0).unwrap();
4083
4084 assert_eq!(values, vec![9, 18, 27, 36, 45]);
4085 }
4086
4087 #[test]
4088 fn test_mul_columns_i64() {
4089 let Some(provider) = create_arith_test_provider() else {
4090 eprintln!("Skipping test: no CUDA device available");
4091 return;
4092 };
4093
4094 let a = create_i64_buffer(&provider, &[2, 3, 4, 5, 6]);
4095 let b = create_i64_buffer(&provider, &[3, 4, 5, 6, 7]);
4096
4097 let result = provider.mul_columns(&a, &b).unwrap();
4098 let values = provider.download_column::<i64>(&result, 0).unwrap();
4099
4100 assert_eq!(values, vec![6, 12, 20, 30, 42]);
4101 }
4102
4103 #[test]
4104 fn test_div_columns_i64() {
4105 let Some(provider) = create_arith_test_provider() else {
4106 eprintln!("Skipping test: no CUDA device available");
4107 return;
4108 };
4109
4110 let a = create_i64_buffer(&provider, &[100, 200, 300, 400]);
4111 let b = create_i64_buffer(&provider, &[10, 20, 30, 40]);
4112
4113 let result = provider.div_columns(&a, &b).unwrap();
4114 let values = provider.download_column::<i64>(&result, 0).unwrap();
4115
4116 assert_eq!(values, vec![10, 10, 10, 10]);
4117 }
4118
4119 #[test]
4120 fn test_div_columns_by_zero() {
4121 let Some(provider) = create_arith_test_provider() else {
4122 eprintln!("Skipping test: no CUDA device available");
4123 return;
4124 };
4125
4126 let a = create_i64_buffer(&provider, &[10, 20, 30]);
4127 let b = create_i64_buffer(&provider, &[2, 0, 3]); let result = provider.div_columns(&a, &b).unwrap();
4130 let values = provider.download_column::<i64>(&result, 0).unwrap();
4131
4132 assert_eq!(values, vec![5, i64::MAX, 10]);
4134 }
4135
4136 #[test]
4137 fn test_mod_columns_i64() {
4138 let Some(provider) = create_arith_test_provider() else {
4139 eprintln!("Skipping test: no CUDA device available");
4140 return;
4141 };
4142
4143 let a = create_i64_buffer(&provider, &[17, 23, 100, 7]);
4144 let b = create_i64_buffer(&provider, &[5, 7, 30, 3]);
4145
4146 let result = provider.mod_columns(&a, &b).unwrap();
4147 let values = provider.download_column::<i64>(&result, 0).unwrap();
4148
4149 assert_eq!(values, vec![2, 2, 10, 1]);
4150 }
4151
4152 #[test]
4153 fn test_mod_columns_by_zero() {
4154 let Some(provider) = create_arith_test_provider() else {
4155 eprintln!("Skipping test: no CUDA device available");
4156 return;
4157 };
4158
4159 let a = create_i64_buffer(&provider, &[10, 20]);
4160 let b = create_i64_buffer(&provider, &[3, 0]); let result = provider.mod_columns(&a, &b).unwrap();
4163 let values = provider.download_column::<i64>(&result, 0).unwrap();
4164
4165 assert_eq!(values, vec![1, 0]);
4167 }
4168
4169 #[test]
4170 fn test_abs_column_i64() {
4171 let Some(provider) = create_arith_test_provider() else {
4172 eprintln!("Skipping test: no CUDA device available");
4173 return;
4174 };
4175
4176 let a = create_i64_buffer(&provider, &[-5, 10, -15, 20, 0]);
4177
4178 let result = provider.abs_column(&a).unwrap();
4179 let values = provider.download_column::<i64>(&result, 0).unwrap();
4180
4181 assert_eq!(values, vec![5, 10, 15, 20, 0]);
4182 }
4183
4184 #[test]
4185 fn test_min_columns_i64() {
4186 let Some(provider) = create_arith_test_provider() else {
4187 eprintln!("Skipping test: no CUDA device available");
4188 return;
4189 };
4190
4191 let a = create_i64_buffer(&provider, &[5, 10, 15, 20]);
4192 let b = create_i64_buffer(&provider, &[3, 12, 10, 25]);
4193
4194 let result = provider.min_columns(&a, &b).unwrap();
4195 let values = provider.download_column::<i64>(&result, 0).unwrap();
4196
4197 assert_eq!(values, vec![3, 10, 10, 20]);
4198 }
4199
4200 #[test]
4201 fn test_max_columns_i64() {
4202 let Some(provider) = create_arith_test_provider() else {
4203 eprintln!("Skipping test: no CUDA device available");
4204 return;
4205 };
4206
4207 let a = create_i64_buffer(&provider, &[5, 10, 15, 20]);
4208 let b = create_i64_buffer(&provider, &[3, 12, 10, 25]);
4209
4210 let result = provider.max_columns(&a, &b).unwrap();
4211 let values = provider.download_column::<i64>(&result, 0).unwrap();
4212
4213 assert_eq!(values, vec![5, 12, 15, 25]);
4214 }
4215
4216 #[test]
4217 fn test_add_columns_f64() {
4218 let Some(provider) = create_arith_test_provider() else {
4219 eprintln!("Skipping test: no CUDA device available");
4220 return;
4221 };
4222
4223 let a = create_f64_buffer(&provider, &[1.5, 2.5, 3.5]);
4224 let b = create_f64_buffer(&provider, &[0.5, 1.5, 2.5]);
4225
4226 let result = provider.add_columns(&a, &b).unwrap();
4227 let values = provider.download_column::<f64>(&result, 0).unwrap();
4228
4229 assert_eq!(values, vec![2.0, 4.0, 6.0]);
4230 }
4231
4232 #[test]
4233 fn test_mul_columns_f64() {
4234 let Some(provider) = create_arith_test_provider() else {
4235 eprintln!("Skipping test: no CUDA device available");
4236 return;
4237 };
4238
4239 let a = create_f64_buffer(&provider, &[2.0, 3.0, 4.0]);
4240 let b = create_f64_buffer(&provider, &[1.5, 2.0, 2.5]);
4241
4242 let result = provider.mul_columns(&a, &b).unwrap();
4243 let values = provider.download_column::<f64>(&result, 0).unwrap();
4244
4245 assert_eq!(values, vec![3.0, 6.0, 10.0]);
4246 }
4247
4248 #[test]
4249 fn test_div_columns_f64_by_zero() {
4250 let Some(provider) = create_arith_test_provider() else {
4251 eprintln!("Skipping test: no CUDA device available");
4252 return;
4253 };
4254
4255 let a = create_f64_buffer(&provider, &[1.0, -1.0, 0.0]);
4256 let b = create_f64_buffer(&provider, &[0.0, 0.0, 0.0]);
4257
4258 let result = provider.div_columns(&a, &b).unwrap();
4259 let values = provider.download_column::<f64>(&result, 0).unwrap();
4260
4261 assert!(values[0].is_infinite() && values[0].is_sign_positive());
4263 assert!(values[1].is_infinite() && values[1].is_sign_negative());
4264 assert!(values[2].is_nan());
4265 }
4266
4267 #[test]
4268 fn test_pow_columns() {
4269 let Some(provider) = create_arith_test_provider() else {
4270 eprintln!("Skipping test: no CUDA device available");
4271 return;
4272 };
4273
4274 let base = create_i64_buffer(&provider, &[2, 3, 4, 5]);
4275 let exp = create_i64_buffer(&provider, &[3, 2, 2, 1]);
4276
4277 let result = provider.pow_columns(&base, &exp).unwrap();
4278 let values = provider.download_column::<f64>(&result, 0).unwrap();
4279
4280 assert_eq!(values, vec![8.0, 9.0, 16.0, 5.0]);
4282 }
4283
4284 #[test]
4285 fn test_pow_columns_fractional_exp() {
4286 let Some(provider) = create_arith_test_provider() else {
4287 eprintln!("Skipping test: no CUDA device available");
4288 return;
4289 };
4290
4291 let base = create_f64_buffer(&provider, &[4.0, 9.0, 27.0]);
4292 let exp = create_f64_buffer(&provider, &[0.5, 0.5, 1.0 / 3.0]);
4293
4294 let result = provider.pow_columns(&base, &exp).unwrap();
4295 let values = provider.download_column::<f64>(&result, 0).unwrap();
4296
4297 assert!((values[0] - 2.0).abs() < 1e-10);
4299 assert!((values[1] - 3.0).abs() < 1e-10);
4300 assert!((values[2] - 3.0).abs() < 1e-10);
4301 }
4302
4303 #[test]
4304 fn test_cast_i64_to_f64() {
4305 let Some(provider) = create_arith_test_provider() else {
4306 eprintln!("Skipping test: no CUDA device available");
4307 return;
4308 };
4309
4310 let a = create_i64_buffer(&provider, &[1, 2, 3, 4, 5]);
4311
4312 let result = provider.cast_column(&a, ScalarType::F64).unwrap();
4313 let values = provider.download_column::<f64>(&result, 0).unwrap();
4314
4315 assert_eq!(values, vec![1.0, 2.0, 3.0, 4.0, 5.0]);
4316 }
4317
4318 #[test]
4319 fn test_cast_f64_to_i64() {
4320 let Some(provider) = create_arith_test_provider() else {
4321 eprintln!("Skipping test: no CUDA device available");
4322 return;
4323 };
4324
4325 let a = create_f64_buffer(&provider, &[1.9, 2.1, 3.5, 4.0, 5.7]);
4326
4327 let result = provider.cast_column(&a, ScalarType::I64).unwrap();
4328 let values = provider.download_column::<i64>(&result, 0).unwrap();
4329
4330 assert_eq!(values, vec![1, 2, 3, 4, 5]);
4332 }
4333
4334 #[test]
4335 fn test_cast_i64_to_i32() {
4336 let Some(provider) = create_arith_test_provider() else {
4337 eprintln!("Skipping test: no CUDA device available");
4338 return;
4339 };
4340
4341 let a = create_i64_buffer(&provider, &[1, 2, 3, 100, 200]);
4342
4343 let result = provider.cast_column(&a, ScalarType::I32).unwrap();
4344 let values = provider.download_column::<i32>(&result, 0).unwrap();
4345
4346 assert_eq!(values, vec![1, 2, 3, 100, 200]);
4347 }
4348
4349 #[test]
4350 fn test_arithmetic_row_count_mismatch() {
4351 let Some(provider) = create_arith_test_provider() else {
4352 eprintln!("Skipping test: no CUDA device available");
4353 return;
4354 };
4355
4356 let a = create_i64_buffer(&provider, &[1, 2, 3]);
4357 let b = create_i64_buffer(&provider, &[1, 2]); let result = provider.add_columns(&a, &b);
4360 assert!(result.is_err());
4361 let err = result.err().unwrap();
4362 assert!(err.to_string().contains("Row count mismatch"));
4363 }
4364
4365 #[test]
4366 fn test_arithmetic_empty_buffers() {
4367 let Some(provider) = create_arith_test_provider() else {
4368 eprintln!("Skipping test: no CUDA device available");
4369 return;
4370 };
4371
4372 let a = create_i64_buffer(&provider, &[]);
4373 let b = create_i64_buffer(&provider, &[]);
4374
4375 let result = provider.add_columns(&a, &b).unwrap();
4376 let values = provider.download_column::<i64>(&result, 0).unwrap();
4377
4378 assert_eq!(values, Vec::<i64>::new());
4379 }
4380
4381 #[test]
4382 fn test_wrapping_arithmetic_overflow() {
4383 let Some(provider) = create_arith_test_provider() else {
4384 eprintln!("Skipping test: no CUDA device available");
4385 return;
4386 };
4387
4388 let a = create_i64_buffer(&provider, &[i64::MAX, i64::MIN]);
4389 let b = create_i64_buffer(&provider, &[1, -1]);
4390
4391 let add_result = provider.add_columns(&a, &b).unwrap();
4393 let add_values = provider.download_column::<i64>(&add_result, 0).unwrap();
4394 assert_eq!(add_values[0], i64::MIN); assert_eq!(add_values[1], i64::MAX); }
4397
4398 #[test]
4399 fn test_abs_column_f64() {
4400 let Some(provider) = create_arith_test_provider() else {
4401 eprintln!("Skipping test: no CUDA device available");
4402 return;
4403 };
4404
4405 let a = create_f64_buffer(&provider, &[-1.5, 2.5, -3.5, 0.0]);
4406
4407 let result = provider.abs_column(&a).unwrap();
4408 let values = provider.download_column::<f64>(&result, 0).unwrap();
4409
4410 assert_eq!(values, vec![1.5, 2.5, 3.5, 0.0]);
4411 }
4412
4413 #[test]
4414 fn test_min_max_columns_f64() {
4415 let Some(provider) = create_arith_test_provider() else {
4416 eprintln!("Skipping test: no CUDA device available");
4417 return;
4418 };
4419
4420 let a = create_f64_buffer(&provider, &[1.5, 5.0, 3.0]);
4421 let b = create_f64_buffer(&provider, &[2.0, 3.0, 4.0]);
4422
4423 let min_result = provider.min_columns(&a, &b).unwrap();
4424 let min_values = provider.download_column::<f64>(&min_result, 0).unwrap();
4425 assert_eq!(min_values, vec![1.5, 3.0, 3.0]);
4426
4427 let max_result = provider.max_columns(&a, &b).unwrap();
4428 let max_values = provider.download_column::<f64>(&max_result, 0).unwrap();
4429 assert_eq!(max_values, vec![2.0, 5.0, 4.0]);
4430 }
4431}