Skip to main content

xlog_cuda/provider/
mod.rs

1//! CUDA kernel provider implementation
2//!
3//! This module provides the `CudaKernelProvider` which manages pre-compiled
4//! PTX kernels for GPU execution of relational operations (join, dedup, groupby).
5
6use std::collections::HashMap;
7use std::marker::PhantomData;
8use std::path::PathBuf;
9use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
10use std::sync::{Arc, Mutex, OnceLock};
11
12use std::ffi::c_void;
13use xlog_core::{Result, Schema, XlogError};
14
15use crate::{
16    cuda_compat::{
17        AsKernelParam, DeviceParamStorage, DevicePtr, DeviceRepr, DeviceSlice,
18        IntoKernelParamStorage, LaunchAsync, LaunchConfig,
19    },
20    cuda_graph::{CapturedCudaGraph, CsmCudaGraphKey, CudaGraphNode},
21    memory::{validate_logical_row_count, CudaColumn, TrackedCudaSlice},
22    CudaBuffer, CudaDevice, CudaStream, CudaViewMut, GpuMemoryManager,
23};
24
25mod arithmetic;
26mod filter;
27mod fj;
28mod fj_delta;
29mod fj_delta_sparse;
30mod groupby;
31mod ilp;
32mod ilp_exact;
33mod io;
34mod kernel_loading;
35pub mod kernel_paths;
36mod launch_safe;
37mod probabilistic;
38mod relational;
39mod transfer;
40mod wcoj;
41mod wcoj_metadata;
42mod wcoj_project;
43
44pub use fj::{FjNode, FjPlan, FjSubAtom};
45pub use fj_delta::{FjDeltaCols, FJ_DELTA_MAX_DOMAIN};
46
47/// Per-module PTX load timing (populated only when XLOG_WARMUP_PROFILE=1).
48#[derive(Debug, Clone, Default)]
49pub struct PtxLoadProfile {
50    pub total_sec: f64,
51    pub per_module_sec: Vec<(String, f64)>,
52    pub cubin_loaded: u32,
53    pub ptx_fallback: u32,
54}
55
56fn warmup_profiling_enabled() -> bool {
57    std::env::var("XLOG_WARMUP_PROFILE")
58        .map(|v| v == "1")
59        .unwrap_or(false)
60}
61
62/// Detect device compute capability as a two-digit number (e.g. 75, 80, 120).
63pub(crate) fn detect_compute_capability(device: &Arc<CudaDevice>) -> Result<u32> {
64    let major = device
65        .inner()
66        .attribute(
67            cudarc::driver::sys::CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
68        )
69        .map_err(|e| XlogError::Kernel(format!("Failed to query SM major: {}", e)))?;
70    let minor = device
71        .inner()
72        .attribute(
73            cudarc::driver::sys::CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
74        )
75        .map_err(|e| XlogError::Kernel(format!("Failed to query SM minor: {}", e)))?;
76    Ok((major as u32) * 10 + (minor as u32))
77}
78
79#[cfg(test)]
80fn resolve_module_path(name: &str, cc: u32) -> Option<(std::path::PathBuf, bool)> {
81    kernel_paths::KernelArtifactLocator::from_env().resolve_module_path(name, cc)
82}
83
84#[derive(Debug)]
85pub(crate) enum KernelModuleSource {
86    File { path: PathBuf, is_cubin: bool },
87    EmbeddedPortablePtx { ptx: &'static str },
88}
89
90pub(crate) fn resolve_module_sources_with_locator(
91    name: &str,
92    cc: u32,
93    locator: &kernel_paths::KernelArtifactLocator,
94) -> Vec<KernelModuleSource> {
95    let mut sources: Vec<KernelModuleSource> = locator
96        .resolve_module_paths(name, cc)
97        .into_iter()
98        // Skip any staged cubin/PTX whose bytes diverge from what this binary
99        // was built against. A stale staged artifact (kernel signature changed
100        // but the staged copy was never refreshed) otherwise loads "fine" and
101        // then launches a mismatched kernel into an illegal address.
102        .filter(|(path, _)| !staged_artifact_is_stale(path))
103        .map(|(path, is_cubin)| KernelModuleSource::File { path, is_cubin })
104        .collect();
105
106    // ALWAYS append the embedded portable PTX as the final fallback. It is
107    // compiled into this binary, so it can never be stale relative to the launch
108    // sites — it guarantees a signature-correct kernel even when every staged
109    // File artifact was skipped as stale or fails to load. (Previously this was
110    // suppressed whenever any portable-PTX *file* existed, which let a stale
111    // staged PTX shadow the fresh embedded one.)
112    if let Some(ptx) = crate::embedded_kernel_data::portable_ptx(name) {
113        sources.push(KernelModuleSource::EmbeddedPortablePtx { ptx });
114    }
115    sources
116}
117
118/// A staged cubin/PTX is "stale" when this binary embeds a canonical integrity
119/// hash for that artifact file name and the on-disk bytes do not match it — the
120/// staged artifact diverges from what this build produced. Loading such an
121/// artifact can launch a mismatched kernel into an illegal address, so it is
122/// skipped in favor of a fresh source. Artifacts with no embedded canonical
123/// hash (e.g. an arch this build did not produce) are NOT treated as stale — we
124/// can only validate what we built — nor are unreadable files (the loader
125/// surfaces the IO error).
126fn staged_artifact_is_stale(path: &std::path::Path) -> bool {
127    let Some(file_name) = path.file_name().and_then(|n| n.to_str()) else {
128        return false;
129    };
130    let Some(expected) = crate::embedded_kernel_data::canonical_artifact_hash(file_name) else {
131        return false;
132    };
133    match std::fs::read(path) {
134        Ok(bytes) => fnv1a_64(&bytes) != expected,
135        Err(_) => false,
136    }
137}
138
139/// FNV-1a 64-bit, matching the build-time hash in `crates/xlog-cuda/build.rs`.
140fn fnv1a_64(bytes: &[u8]) -> u64 {
141    let mut hash: u64 = 0xcbf2_9ce4_8422_2325;
142    for &byte in bytes {
143        hash ^= byte as u64;
144        hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
145    }
146    hash
147}
148
149#[cfg(test)]
150mod kernel_source_resolution_tests {
151    use super::{
152        kernel_paths::KernelArtifactLocator, resolve_module_sources_with_locator,
153        KernelModuleSource,
154    };
155    use std::fs;
156
157    #[test]
158    fn keeps_portable_ptx_fallback_when_cubin_exists() {
159        let root = std::env::temp_dir().join(format!(
160            "xlog-kernel-fallback-{}-{}",
161            std::process::id(),
162            std::time::SystemTime::now()
163                .duration_since(std::time::UNIX_EPOCH)
164                .expect("system clock before UNIX_EPOCH")
165                .as_nanos()
166        ));
167        let kernels = root.join("kernels");
168        fs::create_dir_all(&kernels).expect("create kernels dir");
169        // Use a name this build does NOT produce, so neither file carries a
170        // canonical integrity hash (the staleness skip is exercised separately).
171        // This isolates the file-resolution precedence: cubin first, then the
172        // portable-PTX file as fallback.
173        fs::write(kernels.join("fakekernel.sm_86.cubin"), b"cubin").expect("write cubin");
174        fs::write(kernels.join("fakekernel.portable.ptx"), b"ptx").expect("write ptx");
175        let expected_cubin = kernels.join("fakekernel.sm_86.cubin");
176        let expected_ptx = kernels.join("fakekernel.portable.ptx");
177
178        let locator = KernelArtifactLocator::new(None, Some(kernels.clone()), None);
179        let sources = resolve_module_sources_with_locator("fakekernel", 86, &locator);
180
181        assert_eq!(sources.len(), 2);
182        assert!(matches!(
183            &sources[0],
184            KernelModuleSource::File {
185                path,
186                is_cubin: true
187            } if path == &expected_cubin
188        ));
189        assert!(matches!(
190            &sources[1],
191            KernelModuleSource::File {
192                path,
193                is_cubin: false
194            } if path == &expected_ptx
195        ));
196
197        fs::remove_dir_all(root).expect("remove temp kernels");
198    }
199
200    // Locks the FNV-1a contract between build.rs (which embeds canonical
201    // artifact hashes) and the runtime (which re-hashes staged artifacts). If
202    // these two implementations ever diverge, every staged artifact would read
203    // as "stale" — so these canonical FNV-1a 64-bit vectors must hold.
204    #[test]
205    fn fnv1a_64_matches_known_vectors() {
206        assert_eq!(super::fnv1a_64(b""), 0xcbf2_9ce4_8422_2325);
207        assert_eq!(super::fnv1a_64(b"a"), 0xaf63_dc4c_8601_ec8c);
208        assert_eq!(super::fnv1a_64(b"foobar"), 0x85944171_f73967e8);
209    }
210
211    // A file whose name this build did not produce has no canonical hash, so it
212    // is conservatively NOT treated as stale (we only validate what we built).
213    // A nonexistent path is likewise not "stale" — the loader surfaces IO.
214    #[test]
215    fn staged_artifact_not_stale_without_canonical_hash() {
216        let root = std::env::temp_dir().join(format!(
217            "xlog-kernel-stale-{}-{}",
218            std::process::id(),
219            std::time::SystemTime::now()
220                .duration_since(std::time::UNIX_EPOCH)
221                .expect("system clock before UNIX_EPOCH")
222                .as_nanos()
223        ));
224        fs::create_dir_all(&root).expect("create dir");
225        let unknown = root.join("definitely_not_a_real_kernel.sm_86.cubin");
226        fs::write(&unknown, b"bytes").expect("write");
227        assert!(!super::staged_artifact_is_stale(&unknown));
228        assert!(!super::staged_artifact_is_stale(
229            &root.join("missing.portable.ptx")
230        ));
231        fs::remove_dir_all(root).expect("remove temp dir");
232    }
233}
234
235/// Resolve a kernel module from sidecar artifacts or embedded portable PTX.
236///
237/// Asserts (in debug builds) that `name` is present in the kernel manifest,
238/// catching name/order drift between the manifest and provider load blocks.
239pub(crate) fn load_module_sources(name: &str, cc: u32) -> Result<Vec<KernelModuleSource>> {
240    debug_assert!(
241        crate::kernel_manifest_data::KERNEL_CU_NAMES.contains(&name),
242        "kernel module '{name}' is not in KERNEL_CU_NAMES manifest — update kernel_manifest_data.rs"
243    );
244    let locator = kernel_paths::KernelArtifactLocator::from_env();
245    let sources = resolve_module_sources_with_locator(name, cc, &locator);
246    if sources.is_empty() {
247        Err(XlogError::Kernel(format!(
248            "{name}: no cubin, sidecar portable PTX, or embedded portable PTX found"
249        )))
250    } else {
251        Ok(sources)
252    }
253}
254
255#[derive(Clone)]
256pub(crate) struct RawCudaView<'a, T> {
257    ptr: cudarc::driver::sys::CUdeviceptr,
258    len: usize,
259    stream: Arc<CudaStream>,
260    /// Optional back-reference to the source [`DeviceBlock`]
261    /// when this view borrows a region of a runtime-backed
262    /// allocation. The launch recorder uses this to attach
263    /// cross-stream uses without losing identity through view
264    /// construction. `None` for views built from external
265    /// memory or legacy paths; strict-mode launch recorders
266    /// reject `None` views.
267    ///
268    /// Read by [`RawCudaView::runtime_block`]; the field
269    /// itself is intentionally not directly exposed because
270    /// the lifetime of the back-reference is bound to the
271    /// view's `'a`.
272    #[allow(dead_code)]
273    source_block: Option<&'a crate::device_runtime::DeviceBlock>,
274    _marker: PhantomData<&'a [T]>,
275}
276
277/// Preallocated scratch layout for graph-capturable u32 multi-block scans.
278///
279/// The legacy stream-aware scan helper allocates recursive `block_sums`
280/// buffers inside the helper. CUDA Graph capture records concrete allocation
281/// addresses, so bounded CSM CUDA Graph replay needs the scan topology and
282/// scratch buffers to be fixed before capture begins.
283pub(crate) struct MultiblockScanScratchU32 {
284    levels: Vec<TrackedCudaSlice<u32>>,
285}
286
287impl MultiblockScanScratchU32 {
288    pub(crate) fn levels(&self) -> &[TrackedCudaSlice<u32>] {
289        &self.levels
290    }
291}
292
293pub(crate) struct CsmCudaGraphNodes {
294    pub(crate) count: CudaGraphNode,
295    pub(crate) total: CudaGraphNode,
296    pub(crate) materialize: CudaGraphNode,
297    pub(crate) node_count: usize,
298}
299
300pub(crate) struct CsmCudaGraphEntry {
301    pub(crate) graph: CapturedCudaGraph,
302    pub(crate) nodes: CsmCudaGraphNodes,
303    pub(crate) per_probe_count: TrackedCudaSlice<u32>,
304    pub(crate) per_probe_offsets: TrackedCudaSlice<u32>,
305    pub(crate) d_logical_count: TrackedCudaSlice<u32>,
306    pub(crate) d_overflow: TrackedCudaSlice<u8>,
307    pub(crate) d_output_left: TrackedCudaSlice<u32>,
308    pub(crate) d_output_right: TrackedCudaSlice<u32>,
309    pub(crate) scan_scratch: MultiblockScanScratchU32,
310    pub(crate) probe_capacity: u32,
311    pub(crate) output_capacity: u32,
312}
313
314impl<'a, T> DeviceSlice<T> for RawCudaView<'a, T> {
315    fn len(&self) -> usize {
316        self.len
317    }
318
319    fn stream(&self) -> &Arc<CudaStream> {
320        &self.stream
321    }
322}
323
324impl<'a, T> DevicePtr<T> for RawCudaView<'a, T> {
325    fn device_ptr<'b>(
326        &'b self,
327        _stream: &'b CudaStream,
328    ) -> (
329        cudarc::driver::sys::CUdeviceptr,
330        cudarc::driver::SyncOnDrop<'b>,
331    ) {
332        (self.ptr, cudarc::driver::SyncOnDrop::Sync(None))
333    }
334}
335
336impl<'a, T> RawCudaView<'a, T> {
337    pub fn device_ptr(&self) -> &cudarc::driver::sys::CUdeviceptr {
338        &self.ptr
339    }
340
341    /// Borrow the back-reference to the source
342    /// [`crate::device_runtime::DeviceBlock`], if this view was
343    /// constructed from a runtime-backed allocation. Returns
344    /// `None` for views built from external memory or legacy
345    /// paths.
346    ///
347    /// Public API reserved for the filter-class migration; no
348    /// production caller exists yet.
349    #[allow(dead_code)]
350    pub fn runtime_block(&self) -> Option<&'a crate::device_runtime::DeviceBlock> {
351        self.source_block
352    }
353}
354
355impl<'a, T: DeviceRepr> AsKernelParam for &RawCudaView<'a, T> {
356    fn as_kernel_param(&self) -> *mut c_void {
357        ((*self).device_ptr() as *const cudarc::driver::sys::CUdeviceptr)
358            .cast_mut()
359            .cast()
360    }
361}
362
363impl<'a, T: DeviceRepr> IntoKernelParamStorage for &'a RawCudaView<'a, T> {
364    type Storage = DeviceParamStorage<'a>;
365
366    fn into_kernel_param_storage(self) -> Self::Storage {
367        DeviceParamStorage::unsynced(self.ptr)
368    }
369}
370
371/// Scratch buffers for stable radix sorting of u32 key/value pairs.
372pub struct RadixSortScratch {
373    keys_b: TrackedCudaSlice<u32>,
374    values_b: TrackedCudaSlice<u32>,
375    hist: TrackedCudaSlice<u32>,
376    prefix: TrackedCudaSlice<u32>,
377    ranks: TrackedCudaSlice<u32>,
378    len: u32,
379}
380
381impl RadixSortScratch {
382    pub fn new(provider: &CudaKernelProvider, n: u32) -> Result<Self> {
383        let memory = provider.memory();
384        let len = n.max(1);
385        let keys_b = memory.alloc::<u32>(len as usize)?;
386        let values_b = memory.alloc::<u32>(len as usize)?;
387        let ranks = memory.alloc::<u32>(len as usize)?;
388        let block_size = CudaKernelProvider::SORT_BLOCK_SIZE;
389        let grid_size = len.div_ceil(block_size).max(1);
390        let hist = memory.alloc::<u32>((grid_size as usize) * 16)?;
391        let prefix = memory.alloc::<u32>(16)?;
392        Ok(Self {
393            keys_b,
394            values_b,
395            hist,
396            prefix,
397            ranks,
398            len,
399        })
400    }
401
402    pub fn ensure_capacity(&mut self, provider: &CudaKernelProvider, n: u32) -> Result<()> {
403        if n <= self.len {
404            return Ok(());
405        }
406        *self = Self::new(provider, n)?;
407        Ok(())
408    }
409}
410
411/// Module names for loaded PTX modules
412pub const JOIN_MODULE: &str = "xlog_join";
413pub const DEDUP_MODULE: &str = "xlog_dedup";
414pub const GROUPBY_MODULE: &str = "xlog_groupby";
415pub const SCAN_MODULE: &str = "xlog_scan";
416pub const SORT_MODULE: &str = "xlog_sort";
417pub const FILTER_MODULE: &str = "xlog_filter";
418pub const SET_OPS_MODULE: &str = "xlog_set_ops";
419pub const PACK_MODULE: &str = "xlog_pack";
420pub const CIRCUIT_MODULE: &str = "xlog_circuit";
421pub const MC_SAMPLE_MODULE: &str = "xlog_mc_sample";
422pub const MC_EVAL_MODULE: &str = "xlog_mc_eval";
423pub const MC_RESIDENT_MODULE: &str = "xlog_mc_resident";
424pub const ARITH_MODULE: &str = "xlog_arith";
425pub const SAT_MODULE: &str = "xlog_sat";
426pub const D4_MODULE: &str = "xlog_d4";
427pub const NEURAL_MODULE: &str = "xlog_neural";
428pub const PIR_MODULE: &str = "xlog_pir";
429pub const CNF_MODULE: &str = "xlog_cnf";
430pub const CACHE_MODULE: &str = "xlog_cache";
431pub const WEIGHTS_MODULE: &str = "xlog_weights";
432pub const ILP_MODULE: &str = "xlog_ilp";
433pub const ILP_CREDIT_MODULE: &str = "xlog_ilp_credit";
434pub const ILP_EXACT_MODULE: &str = "xlog_ilp_exact";
435pub const EPISTEMIC_MODULE: &str = "xlog_epistemic";
436pub const WCOJ_MODULE: &str = "xlog_wcoj";
437pub const JOINT_SOLVE_MODULE: &str = "xlog_joint_solve";
438
439// Compile-time check: kernel manifest lists exactly 26 modules.
440const _: () = assert!(crate::kernel_manifest_data::KERNEL_CU_NAMES.len() == 26);
441
442/// Kernel function names in the GPU WCOJ module.
443pub mod wcoj_kernels {
444    pub const WCOJ_BUILD_METADATA_MARK_BOUNDARIES_U32: &str =
445        "wcoj_build_metadata_mark_boundaries_u32";
446    pub const WCOJ_BUILD_METADATA_MARK_BOUNDARIES_U64: &str =
447        "wcoj_build_metadata_mark_boundaries_u64";
448    pub const WCOJ_BUILD_METADATA_SCATTER_U32: &str = "wcoj_build_metadata_scatter_u32";
449    pub const WCOJ_BUILD_METADATA_SCATTER_U64: &str = "wcoj_build_metadata_scatter_u64";
450    pub const WCOJ_TRIANGLE_BUILD_HG_WORK_PLAN_U32: &str = "wcoj_triangle_build_hg_work_plan_u32";
451    pub const WCOJ_TRIANGLE_COUNT_HG_U32: &str = "wcoj_triangle_count_hg_u32";
452    pub const WCOJ_TRIANGLE_GROUPBY_ROOT_COUNT_HG_U32: &str =
453        "wcoj_triangle_groupby_root_count_hg_u32";
454    pub const WCOJ_TRIANGLE_GROUPBY_ROOT_SUM_HG_U32: &str = "wcoj_triangle_groupby_root_sum_hg_u32";
455    pub const WCOJ_TRIANGLE_GROUPBY_ROOT_MIN_HG_U32: &str = "wcoj_triangle_groupby_root_min_hg_u32";
456    pub const WCOJ_TRIANGLE_GROUPBY_ROOT_MAX_HG_U32: &str = "wcoj_triangle_groupby_root_max_hg_u32";
457    pub const WCOJ_TRIANGLE_MATERIALIZE_HG_U32: &str = "wcoj_triangle_materialize_hg_u32";
458    pub const WCOJ_TRIANGLE_BUILD_HG_WORK_PLAN_U64: &str = "wcoj_triangle_build_hg_work_plan_u64";
459    pub const WCOJ_TRIANGLE_COUNT_HG_U64: &str = "wcoj_triangle_count_hg_u64";
460    pub const WCOJ_TRIANGLE_GROUPBY_ROOT_COUNT_HG_U64: &str =
461        "wcoj_triangle_groupby_root_count_hg_u64";
462    pub const WCOJ_TRIANGLE_GROUPBY_ROOT_SUM_HG_U64: &str = "wcoj_triangle_groupby_root_sum_hg_u64";
463    pub const WCOJ_TRIANGLE_GROUPBY_ROOT_MIN_HG_U64: &str = "wcoj_triangle_groupby_root_min_hg_u64";
464    pub const WCOJ_TRIANGLE_GROUPBY_ROOT_MAX_HG_U64: &str = "wcoj_triangle_groupby_root_max_hg_u64";
465    pub const WCOJ_GROUPBY_ROOT_SEGMENT_SUM_COUNTS_U32: &str =
466        "wcoj_groupby_root_segment_sum_counts_u32";
467    pub const WCOJ_GROUPBY_ROOT_SEGMENT_SUM_VALUES_U64: &str =
468        "wcoj_groupby_root_segment_sum_values_u64";
469    pub const WCOJ_GROUPBY_ROOT_SEGMENT_MIN_VALUES_U64: &str =
470        "wcoj_groupby_root_segment_min_values_u64";
471    pub const WCOJ_GROUPBY_ROOT_SEGMENT_MAX_VALUES_U64: &str =
472        "wcoj_groupby_root_segment_max_values_u64";
473    pub const WCOJ_TRIANGLE_MATERIALIZE_HG_U64: &str = "wcoj_triangle_materialize_hg_u64";
474    pub const WCOJ_TRIANGLE_COUNT_HG_CACHED_U32: &str = "wcoj_triangle_count_hg_cached_u32";
475    pub const WCOJ_TRIANGLE_MATERIALIZE_HG_CACHED_U32: &str =
476        "wcoj_triangle_materialize_hg_cached_u32";
477    pub const WCOJ_SCAN_HG_BLOCK_COUNTS_U32: &str = "wcoj_scan_hg_block_counts_u32";
478    pub const WCOJ_COMPUTE_TOTAL: &str = "wcoj_compute_total";
479    pub const WCOJ_LAYOUT_CHECK_SORTED_UNIQUE_U32: &str = "wcoj_layout_check_sorted_unique_u32";
480    pub const WCOJ_LAYOUT_CHECK_SORTED_UNIQUE_U64: &str = "wcoj_layout_check_sorted_unique_u64";
481    pub const WCOJ_4CYCLE_BUILD_E2_WORK_PREFIX_U32: &str = "wcoj_4cycle_build_e2_work_prefix_u32";
482    pub const WCOJ_4CYCLE_BUILD_HG_WORK_PLAN_U32: &str = "wcoj_4cycle_build_hg_work_plan_u32";
483    pub const WCOJ_4CYCLE_COUNT_HG_U32: &str = "wcoj_4cycle_count_hg_u32";
484    pub const WCOJ_4CYCLE_GROUPBY_ROOT_COUNT_HG_U32: &str = "wcoj_4cycle_groupby_root_count_hg_u32";
485    pub const WCOJ_4CYCLE_GROUPBY_ROOT_SUM_HG_U32: &str = "wcoj_4cycle_groupby_root_sum_hg_u32";
486    pub const WCOJ_4CYCLE_GROUPBY_ROOT_MIN_HG_U32: &str = "wcoj_4cycle_groupby_root_min_hg_u32";
487    pub const WCOJ_4CYCLE_GROUPBY_ROOT_MAX_HG_U32: &str = "wcoj_4cycle_groupby_root_max_hg_u32";
488    pub const WCOJ_4CYCLE_MATERIALIZE_HG_U32: &str = "wcoj_4cycle_materialize_hg_u32";
489    pub const WCOJ_4CYCLE_BUILD_E2_WORK_PREFIX_U64: &str = "wcoj_4cycle_build_e2_work_prefix_u64";
490    pub const WCOJ_4CYCLE_BUILD_HG_WORK_PLAN_U64: &str = "wcoj_4cycle_build_hg_work_plan_u64";
491    pub const WCOJ_4CYCLE_COUNT_HG_U64: &str = "wcoj_4cycle_count_hg_u64";
492    pub const WCOJ_4CYCLE_GROUPBY_ROOT_COUNT_HG_U64: &str = "wcoj_4cycle_groupby_root_count_hg_u64";
493    pub const WCOJ_4CYCLE_MATERIALIZE_HG_U64: &str = "wcoj_4cycle_materialize_hg_u64";
494    // General-arity clique kernels (k=5..8 from a single template).
495    pub const WCOJ_CLIQUE5_COUNT_HG_U32: &str = "wcoj_clique5_count_hg_u32";
496    pub const WCOJ_CLIQUE5_MATERIALIZE_HG_U32: &str = "wcoj_clique5_materialize_hg_u32";
497    pub const WCOJ_CLIQUE5_COUNT_HG_U64: &str = "wcoj_clique5_count_hg_u64";
498    pub const WCOJ_CLIQUE5_MATERIALIZE_HG_U64: &str = "wcoj_clique5_materialize_hg_u64";
499    pub const WCOJ_CLIQUE6_COUNT_HG_U32: &str = "wcoj_clique6_count_hg_u32";
500    pub const WCOJ_CLIQUE6_MATERIALIZE_HG_U32: &str = "wcoj_clique6_materialize_hg_u32";
501    pub const WCOJ_CLIQUE6_COUNT_HG_U64: &str = "wcoj_clique6_count_hg_u64";
502    pub const WCOJ_CLIQUE6_MATERIALIZE_HG_U64: &str = "wcoj_clique6_materialize_hg_u64";
503    pub const WCOJ_CLIQUE7_COUNT_HG_U32: &str = "wcoj_clique7_count_hg_u32";
504    pub const WCOJ_CLIQUE7_MATERIALIZE_HG_U32: &str = "wcoj_clique7_materialize_hg_u32";
505    pub const WCOJ_CLIQUE7_COUNT_HG_U64: &str = "wcoj_clique7_count_hg_u64";
506    pub const WCOJ_CLIQUE7_MATERIALIZE_HG_U64: &str = "wcoj_clique7_materialize_hg_u64";
507    pub const WCOJ_CLIQUE8_COUNT_HG_U32: &str = "wcoj_clique8_count_hg_u32";
508    pub const WCOJ_CLIQUE8_MATERIALIZE_HG_U32: &str = "wcoj_clique8_materialize_hg_u32";
509    pub const WCOJ_CLIQUE8_COUNT_HG_U64: &str = "wcoj_clique8_count_hg_u64";
510    pub const WCOJ_CLIQUE8_MATERIALIZE_HG_U64: &str = "wcoj_clique8_materialize_hg_u64";
511    pub const WCOJ_CLIQUE5_GROUPBY_ROOT_COUNT_HG_U32: &str =
512        "wcoj_clique5_groupby_root_count_hg_u32";
513    pub const WCOJ_CLIQUE6_GROUPBY_ROOT_COUNT_HG_U32: &str =
514        "wcoj_clique6_groupby_root_count_hg_u32";
515    // Free Join frontier engine primitives. The work
516    // prefix kernel is width-agnostic (ranges are u32 row indices in
517    // every width class); count/emit/probe have u64 data twins.
518    pub const FJ_EXPAND_WORK_PREFIX_U32: &str = "fj_expand_work_prefix_u32";
519    pub const FJ_EXPAND_COUNT_U32: &str = "fj_expand_count_u32";
520    pub const FJ_EXPAND_EMIT_U32: &str = "fj_expand_emit_u32";
521    pub const FJ_PROBE_REFINE_U32: &str = "fj_probe_refine_u32";
522    pub const FJ_EXPAND_COUNT_U64: &str = "fj_expand_count_u64";
523    pub const FJ_EXPAND_EMIT_U64: &str = "fj_expand_emit_u64";
524    pub const FJ_PROBE_REFINE_U64: &str = "fj_probe_refine_u64";
525    pub const FJ_COUNT_MULTIPLICITY: &str = "fj_count_multiplicity";
526    // D3 S3 spike — factorized recursive delta novel-set pipeline.
527    pub const FJ_DELTA_RANGE_U32: &str = "fj_delta_range_u32";
528    pub const FJ_DELTA_MARK_U32: &str = "fj_delta_mark_u32";
529    pub const FJ_DELTA_SUBTRACT_U32: &str = "fj_delta_subtract_u32";
530    pub const FJ_DELTA_POPCOUNT: &str = "fj_delta_popcount";
531    pub const FJ_DELTA_EMIT_U32: &str = "fj_delta_emit_u32";
532    pub const FJ_DELTA_MAX_U32: &str = "fj_delta_max_u32";
533    pub const FJ_DELTA_SPARSE_ESTIMATE: &str = "fj_delta_sparse_estimate";
534    pub const FJ_DELTA_SPARSE_LOAD_R: &str = "fj_delta_sparse_load_r";
535    pub const FJ_DELTA_SPARSE_INSERT_CANDIDATES: &str = "fj_delta_sparse_insert_candidates";
536    pub const FJ_DELTA_SPARSE_MARK: &str = "fj_delta_sparse_mark";
537    pub const FJ_DELTA_SPARSE_EMIT: &str = "fj_delta_sparse_emit";
538}
539
540/// Kernel function names in the Monte Carlo sampling module
541pub mod mc_sample_kernels {
542    pub const MC_SAMPLE_BERNOULLI: &str = "mc_sample_bernoulli";
543}
544
545/// Kernel function names in the Monte Carlo evaluation module
546pub mod mc_eval_kernels {
547    pub const MC_EVAL_MASK_VAR: &str = "mc_eval_mask_var";
548    pub const MC_EVAL_MASK_AD: &str = "mc_eval_mask_ad_choice";
549    pub const MC_EVAL_QUERY_EVIDENCE_TRUTH: &str = "mc_eval_query_evidence_truth";
550    pub const MC_EVAL_ACCUMULATE_COUNTS: &str = "mc_accumulate_counts";
551}
552
553/// Kernel function names in the GPU-resident Datalog/MC engine module.
554pub mod mc_resident_kernels {
555    /// Single megakernel: evaluates all MC worlds to fixpoint and counts
556    /// query/evidence satisfaction with zero host interaction in-region.
557    pub const MC_RESIDENT_ENGINE: &str = "mc_resident_engine";
558}
559
560/// Kernel function names in the arithmetic module
561pub mod arith_kernels {
562    pub const ARITH_BINARY_I64: &str = "arith_binary_i64";
563    pub const ARITH_BINARY_I32: &str = "arith_binary_i32";
564    pub const ARITH_BINARY_U64: &str = "arith_binary_u64";
565    pub const ARITH_BINARY_U32: &str = "arith_binary_u32";
566    pub const ARITH_BINARY_F64: &str = "arith_binary_f64";
567    pub const ARITH_BINARY_F32: &str = "arith_binary_f32";
568    pub const ARITH_ABS_I64: &str = "arith_abs_i64";
569    pub const ARITH_ABS_I32: &str = "arith_abs_i32";
570    pub const ARITH_ABS_F64: &str = "arith_abs_f64";
571    pub const ARITH_ABS_F32: &str = "arith_abs_f32";
572    pub const ARITH_POW_F64: &str = "arith_pow_f64";
573    pub const ARITH_CAST: &str = "arith_cast";
574    pub const ARITH_FILL_CONST_U32: &str = "arith_fill_const_u32";
575    pub const ARITH_FILL_CONST_U64: &str = "arith_fill_const_u64";
576    pub const ARITH_FILL_CONST_I64: &str = "arith_fill_const_i64";
577    pub const ARITH_FILL_CONST_I32: &str = "arith_fill_const_i32";
578    pub const ARITH_FILL_CONST_F64: &str = "arith_fill_const_f64";
579    pub const ARITH_FILL_CONST_F32: &str = "arith_fill_const_f32";
580    pub const ARITH_FILL_CONST_U8: &str = "arith_fill_const_u8";
581    // Conditional select kernels
582    pub const ARITH_SELECT_I64: &str = "arith_select_i64";
583    pub const ARITH_SELECT_I32: &str = "arith_select_i32";
584    pub const ARITH_SELECT_U64: &str = "arith_select_u64";
585    pub const ARITH_SELECT_U32: &str = "arith_select_u32";
586    pub const ARITH_SELECT_F64: &str = "arith_select_f64";
587    pub const ARITH_SELECT_F32: &str = "arith_select_f32";
588}
589
590/// Kernel function names in the epistemic module.
591pub mod epistemic_kernels {
592    /// Device-side epistemic candidate-assumption generator.
593    pub const EPISTEMIC_GENERATE_CANDIDATE_ASSUMPTIONS_U8: &str =
594        "epistemic_generate_candidate_assumptions_u8";
595    /// Device-side epistemic candidate propagation staging kernel.
596    pub const EPISTEMIC_PROPAGATE_CANDIDATES_U8: &str = "epistemic_propagate_candidates_u8";
597    /// Device-side epistemic candidate bit validation kernel.
598    pub const EPISTEMIC_VALIDATE_CANDIDATE_BITS_U8: &str = "epistemic_validate_candidate_bits_u8";
599    /// Device-side model-membership staging kernel.
600    pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_U8: &str =
601        "epistemic_populate_model_membership_u8";
602    /// Device-side tuple-source-backed model-membership kernel.
603    pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_FROM_TUPLE_SOURCE_U8: &str =
604        "epistemic_populate_model_membership_from_tuple_source_u8";
605    /// Device-side arity-one tuple-key-backed model-membership kernel.
606    pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_FROM_TUPLE_SOURCE_ARITY1_U8: &str =
607        "epistemic_populate_model_membership_from_tuple_source_arity1_u8";
608    /// Device-side arity-two tuple-key-backed model-membership kernel.
609    pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_FROM_TUPLE_SOURCE_ARITY2_U8: &str =
610        "epistemic_populate_model_membership_from_tuple_source_arity2_u8";
611    /// Device-side arity-three tuple-key-backed model-membership kernel.
612    pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_FROM_TUPLE_SOURCE_ARITY3_U8: &str =
613        "epistemic_populate_model_membership_from_tuple_source_arity3_u8";
614    /// Device-side generic-arity tuple-key-backed model-membership kernel.
615    pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_FROM_TUPLE_SOURCE_ARITY_N_U8: &str =
616        "epistemic_populate_model_membership_from_tuple_source_arity_n_u8";
617    /// Device-side world-view validation kernel.
618    pub const EPISTEMIC_VALIDATE_WORLD_VIEWS_U8: &str = "epistemic_validate_world_views_u8";
619    /// Device-side world-view integrity-constraint validation kernel.
620    pub const EPISTEMIC_VALIDATE_CONSTRAINTS_U8: &str = "epistemic_validate_constraints_u8";
621    /// Device-side accepted-candidate materialization staging kernel.
622    pub const EPISTEMIC_MATERIALIZE_ACCEPTED_CANDIDATES_U8: &str =
623        "epistemic_materialize_accepted_candidates_u8";
624
625    /// Device-side final-result flag materialization staging kernel.
626    pub const EPISTEMIC_MATERIALIZE_FINAL_RESULT_FLAGS_U8: &str =
627        "epistemic_materialize_final_result_flags_u8";
628    /// Device-side final tuple materialization kernel.
629    pub const EPISTEMIC_MATERIALIZE_FINAL_TUPLE_COLUMN_U8: &str =
630        "epistemic_materialize_final_tuple_column_u8";
631    /// Device-side final tuple row-map kernel.
632    pub const EPISTEMIC_BUILD_FINAL_TUPLE_ROW_MAP_U8: &str =
633        "epistemic_build_final_tuple_row_map_u8";
634    /// Device-side final tuple rejection-close kernel.
635    pub const EPISTEMIC_CLOSE_FINAL_TUPLE_REJECTIONS_U8: &str =
636        "epistemic_close_final_tuple_rejections_u8";
637}
638
639/// Kernel function names in the neural fast-path module.
640pub mod neural_kernels {
641    pub const NEURAL_FILL_AD_CHAIN_F32: &str = "neural_fill_ad_chain_f32";
642    pub const NEURAL_SCATTER_AD_CHAIN_GRADS_F32: &str = "neural_scatter_ad_chain_grads_f32";
643}
644
645/// Kernel function names in the ILP module.
646pub mod ilp_kernels {
647    pub const EXTRACT_NONZERO_INDICES: &str = "extract_nonzero_indices";
648    pub const ILP_MARK_SELECTED_IDS_U32: &str = "ilp_mark_selected_ids_u32";
649    pub const ILP_MARK_SELECTED_IDS_I32: &str = "ilp_mark_selected_ids_i32";
650    pub const ILP_MARK_SELECTED_IDS_I64: &str = "ilp_mark_selected_ids_i64";
651    pub const ILP_MARK_SELECTED_IDS_U64: &str = "ilp_mark_selected_ids_u64";
652    pub const ILP_VALIDATE_SELECTED_IDS_U32: &str = "ilp_validate_selected_ids_u32";
653    pub const ILP_VALIDATE_SELECTED_IDS_I32: &str = "ilp_validate_selected_ids_i32";
654    pub const ILP_VALIDATE_SELECTED_IDS_I64: &str = "ilp_validate_selected_ids_i64";
655    pub const ILP_VALIDATE_SELECTED_IDS_U64: &str = "ilp_validate_selected_ids_u64";
656    pub const ILP_BROADCAST_CANDIDATE_FLAG: &str = "ilp_broadcast_candidate_flag";
657    pub const ILP_COO_FILL_FROM_MASK: &str = "ilp_coo_fill_from_mask";
658    pub const ILP_CSR_HISTOGRAM: &str = "ilp_csr_histogram";
659    pub const ILP_REDUCE_SUM_F32: &str = "ilp_reduce_sum_f32";
660    pub const ILP_REDUCE_SUM_F64: &str = "ilp_reduce_sum_f64";
661}
662
663/// Kernel function names in the ILP credit module.
664pub mod ilp_credit_kernels {
665    pub const ILP_COO_FILL: &str = "ilp_coo_fill";
666    pub const ILP_CREDIT_FORWARD_F32: &str = "ilp_credit_forward_f32";
667    pub const ILP_CREDIT_FORWARD_F64: &str = "ilp_credit_forward_f64";
668    pub const ILP_CREDIT_BACKWARD_F32: &str = "ilp_credit_backward_f32";
669    pub const ILP_CREDIT_BACKWARD_F64: &str = "ilp_credit_backward_f64";
670}
671
672/// Kernel function names in the native bounded exact-induction module.
673pub mod ilp_exact_kernels {
674    pub const ILP_EXACT_SCORE: &str = "ilp_exact_score";
675    pub const ILP_EXACT_SCORE_U32: &str = "ilp_exact_score_u32";
676    pub const ILP_EXACT_SCORE_CHAIN_SMEM: &str = "ilp_exact_score_chain_smem";
677    pub const ILP_EXACT_SCORE_CHAIN_SMEM_U32: &str = "ilp_exact_score_chain_smem_u32";
678    pub const ILP_EXACT_SELECT_TOPK: &str = "ilp_exact_select_topk";
679}
680
681/// Kernel function names in the PIR interning module.
682pub mod pir_kernels {
683    pub const PIR_PACK_KEYS: &str = "pir_pack_keys";
684    pub const PIR_HASH_KEYS: &str = "pir_hash_keys";
685    pub const PIR_MARK_UNIQUE: &str = "pir_mark_unique";
686    pub const PIR_FIND_EXISTING: &str = "pir_find_existing";
687    pub const PIR_MARK_NEW_GROUPS: &str = "pir_mark_new_groups";
688    pub const PIR_BUILD_GROUP_IDS: &str = "pir_build_group_ids";
689    pub const PIR_FILL_CHILD_PARENTS: &str = "pir_fill_child_parents";
690    pub const PIR_MARK_UNIQUE_PAIRS: &str = "pir_mark_unique_pairs";
691    pub const PIR_COMPACT_PAIRS: &str = "pir_compact_pairs";
692    pub const PIR_COUNT_CHILDREN: &str = "pir_count_children";
693    pub const PIR_WRITE_CHILD_OFFSETS: &str = "pir_write_child_offsets";
694    pub const PIR_GATHER_CHILDREN: &str = "pir_gather_children";
695    pub const PIR_BUILD_GRAPH_CHILD_COUNTS: &str = "pir_build_graph_child_counts";
696    pub const PIR_SUM_COUNTS: &str = "pir_sum_counts";
697    pub const PIR_EMIT_NODES_AND_IDS: &str = "pir_emit_nodes_and_ids";
698    pub const PIR_UPDATE_COUNTS: &str = "pir_update_counts";
699}
700
701/// Kernel function names in the GPU CNF encoder module.
702pub mod cnf_kernels {
703    pub const CNF_REACHABILITY_INIT: &str = "cnf_reachability_init";
704    pub const CNF_REACHABILITY_BFS: &str = "cnf_reachability_bfs";
705    pub const CNF_MARK_LEAF_CHOICE: &str = "cnf_mark_leaf_choice";
706    pub const CNF_ASSIGN_LEAF_VAR: &str = "cnf_assign_leaf_var";
707    pub const CNF_ASSIGN_CHOICE_VAR: &str = "cnf_assign_choice_var";
708    pub const CNF_MARK_NODE_VARS: &str = "cnf_mark_node_vars";
709    pub const CNF_COUNT_CLAUSES: &str = "cnf_count_clauses";
710    pub const CNF_CAPTURE_LAST_COUNTS: &str = "cnf_capture_last_counts";
711    pub const CNF_COMPUTE_LEAF_CHOICE_TOTALS: &str = "cnf_compute_leaf_choice_totals";
712    pub const CNF_COMPUTE_TOTALS: &str = "cnf_compute_totals";
713    pub const CNF_ASSIGN_NODE_VAR: &str = "cnf_assign_node_var";
714    pub const CNF_EMIT_CLAUSES: &str = "cnf_emit_clauses";
715    pub const CNF_SET_CLAUSE_END: &str = "cnf_set_clause_end";
716}
717
718/// Kernel function names in the weights module.
719pub mod weights_kernels {
720    pub const WEIGHTS_FILL_LEAF: &str = "weights_fill_leaf";
721    pub const WEIGHTS_FILL_CHOICE: &str = "weights_fill_choice";
722    pub const WEIGHTS_COUNT_LIFT_EXACT: &str = "weights_count_lift_exact";
723    pub const WEIGHTS_SET_EVIDENCE_FROM_NODES: &str = "weights_set_evidence_from_nodes";
724    pub const WEIGHTS_APPLY_EVIDENCE: &str = "weights_apply_evidence";
725    pub const WEIGHTS_MAP_NODES_TO_VARS: &str = "weights_map_nodes_to_vars";
726    pub const WEIGHTS_FORCE_VAR_FALSE: &str = "weights_force_var_false";
727    pub const WEIGHTS_RESTORE_VAR_FALSE: &str = "weights_restore_var_false";
728    pub const WEIGHTS_FORCE_VAR_TRUE: &str = "weights_force_var_true";
729    pub const WEIGHTS_RESTORE_VAR_TRUE: &str = "weights_restore_var_true";
730    pub const WEIGHTS_COPY_SLOT_TO_BATCH: &str = "weights_copy_slot_to_batch";
731    pub const WEIGHTS_APPLY_QUERY_VARS: &str = "weights_apply_query_vars";
732    pub const WEIGHTS_RESTORE_QUERY_VARS: &str = "weights_restore_query_vars";
733    pub const WEIGHTS_APPLY_QUERY_VARS_FALSE_BATCHED: &str =
734        "weights_apply_query_vars_false_batched";
735    pub const WEIGHTS_RESTORE_QUERY_VARS_FALSE_BATCHED: &str =
736        "weights_restore_query_vars_false_batched";
737    pub const WEIGHTS_APPLY_QUERY_VARS_TRUE_BATCHED: &str = "weights_apply_query_vars_true_batched";
738    pub const WEIGHTS_RESTORE_QUERY_VARS_TRUE_BATCHED: &str =
739        "weights_restore_query_vars_true_batched";
740}
741
742/// Kernel function names in the GPU Decision-DNNF compiler module
743/// (CNF validation + circuit levelization).
744pub mod d4_kernels {
745    pub const D4_VALIDATE_CNF: &str = "d4_validate_cnf";
746    pub const D4_LEVELIZE_COUNTS: &str = "d4_levelize_counts";
747    pub const D4_LEVELIZE_EMIT: &str = "d4_levelize_emit";
748    // BFS frontier expansion and unit propagation.
749    pub const D4_FRONTIER_PREPARE: &str = "d4_frontier_prepare";
750    pub const D4_FRONTIER_EXPAND: &str = "d4_frontier_expand";
751    pub const D4_FRONTIER_PREPARE_DENSE: &str = "d4_frontier_prepare_dense";
752    pub const D4_FRONTIER_EXPAND_DENSE: &str = "d4_frontier_expand_dense";
753    // Per-frontier Decision-DNNF DFS worker (count+emit).
754    pub const D4_COMPILE_COUNT: &str = "d4_compile_count";
755    pub const D4_COMPILE_EMIT: &str = "d4_compile_emit";
756    pub const D4_CAPTURE_EMIT_META: &str = "d4_capture_emit_meta";
757    // GPU smoothing with random-variable support and wrapper emission.
758    pub const D4_SUPPORT_LEVEL: &str = "d4_support_level";
759    pub const D4_SUPPORT_SET_ROOT_BITS: &str = "d4_support_set_root_bits";
760    pub const D4_SMOOTH_COUNT: &str = "d4_smooth_count";
761    pub const D4_SMOOTH_WRAPPER_COUNTS: &str = "d4_smooth_wrapper_counts";
762    pub const D4_SMOOTH_WRAPPER_EDGE_COUNTS_OR: &str = "d4_smooth_wrapper_edge_counts_or";
763    pub const D4_SMOOTH_WRAPPER_EDGE_COUNTS_DEC: &str = "d4_smooth_wrapper_edge_counts_dec";
764    pub const D4_SMOOTH_INIT_NODES: &str = "d4_smooth_init_nodes";
765    pub const D4_SMOOTH_EMIT_LEVEL: &str = "d4_smooth_emit_level";
766    pub const D4_SMOOTH_CHECK_EDGE_CAP: &str = "d4_smooth_check_edge_cap";
767    // GPU free-variable mask for variables in clauses versus the circuit.
768    pub const D4_MARK_VARS_IN_CLAUSES: &str = "d4_mark_vars_in_clauses";
769    pub const D4_MARK_VARS_IN_CIRCUIT: &str = "d4_mark_vars_in_circuit";
770    pub const D4_BUILD_FREE_VAR_MASK: &str = "d4_build_free_var_mask";
771    // GPU-only assertions (tests + invariant enforcement without host reads).
772    pub const D4_ASSERT_U32_EQ: &str = "d4_assert_u32_eq";
773    pub const D4_ASSERT_BITSET_VAR: &str = "d4_assert_bitset_var";
774    pub const D4_ASSERT_DENSE_VAR: &str = "d4_assert_dense_var";
775    pub const D4_ASSERT_LEAF_ROOT_AND_DEGREE: &str = "d4_assert_leaf_root_and_degree";
776}
777
778/// Kernel function names in the join module
779pub mod join_kernels {
780    pub const HASH_JOIN_BUILD: &str = "hash_join_build";
781    pub const HASH_JOIN_PROBE: &str = "hash_join_probe";
782    // V2 kernels for multi-column joins
783    pub const COMPUTE_COMPOSITE_HASH: &str = "compute_composite_hash";
784    pub const HASH_JOIN_BUCKET_COUNT_V2: &str = "hash_join_bucket_count_v2";
785    pub const HASH_JOIN_SCATTER_V2: &str = "hash_join_scatter_v2";
786    pub const HASH_JOIN_PROBE_V2: &str = "hash_join_probe_v2";
787    pub const HASH_JOIN_PROBE_V2_COUNT_PER_ROW: &str = "hash_join_probe_v2_count_per_row";
788    pub const HASH_JOIN_PROBE_V2_MATERIALIZE: &str = "hash_join_probe_v2_materialize";
789    pub const HASH_JOIN_TOTAL_FROM_SCAN: &str = "hash_join_total_from_scan";
790    pub const HASH_JOIN_CSM_UNMATCHED_MASK: &str = "hash_join_csm_unmatched_mask";
791    pub const HASH_JOIN_SEMI: &str = "hash_join_semi";
792    pub const HASH_JOIN_ANTI: &str = "hash_join_anti";
793    pub const INIT_HASH_TABLE: &str = "init_hash_table";
794    /// Nested-loop inner join (emit-pairs design). Reads
795    /// the single key column from each side; emits matched
796    /// `(left_idx, right_idx)` pairs as two parallel u32 arrays.
797    /// Payload columns are materialized after the kernel via
798    /// `gather_buffer_by_indices` in the provider fn.
799    pub const NESTED_LOOP_JOIN_INNER_U32_1KEY_PAIRS: &str = "nested_loop_join_inner_u32_1key_pairs";
800    /// Sort-merge inner join (emit-pairs design,
801    /// caller-asserted pre-sorted inputs). Reads the single
802    /// key column from each side, performs per-thread binary
803    /// search on the right side to find matched-key runs,
804    /// emits `(left_idx, right_idx)` pairs as two parallel
805    /// u32 arrays. Payload columns materialize after the
806    /// kernel via `gather_buffer_by_indices`.
807    pub const SORT_MERGE_JOIN_INNER_U32_1KEY_PAIRS: &str = "sort_merge_join_inner_u32_1key_pairs";
808}
809
810/// Kernel function names in the dedup module
811pub mod dedup_kernels {
812    pub const MARK_DUPLICATES: &str = "mark_duplicates";
813    pub const MARK_UNIQUE_COLUMNAR: &str = "mark_unique_columnar";
814    pub const MARK_UNIQUE_AND_SCAN_COLUMNAR: &str = "mark_unique_and_scan_columnar";
815    pub const COMPACT_ROWS: &str = "compact_rows";
816    pub const MARK_UNIQUE_FULL_ROW_BYTEWISE: &str = "mark_unique_full_row_bytewise";
817    pub const MARK_DIFF_FULL_ROW_TYPED_SORTED: &str = "mark_diff_full_row_typed_sorted";
818    pub const SMALL_SORT_FULL_ROW_INDICES_TYPED: &str = "small_sort_full_row_indices_typed";
819}
820
821/// Kernel function names in the groupby module
822pub mod groupby_kernels {
823    pub const DETECT_GROUP_BOUNDARIES: &str = "detect_group_boundaries";
824    pub const DETECT_BOUNDARIES: &str = "detect_boundaries";
825    pub const EXTRACT_GROUP_KEYS: &str = "extract_group_keys";
826    pub const GROUP_IDS_FROM_BOUNDARIES: &str = "group_ids_from_boundaries";
827    pub const GROUP_START_INDICES: &str = "group_start_indices";
828    pub const CAPTURE_NUM_GROUPS: &str = "capture_num_groups";
829    pub const GROUPBY_COUNT: &str = "groupby_count";
830    pub const GROUPBY_SUM: &str = "groupby_sum";
831    pub const GROUPBY_SUM_U64: &str = "groupby_sum_u64";
832    pub const GROUPBY_MIN: &str = "groupby_min";
833    pub const GROUPBY_MIN_U64: &str = "groupby_min_u64";
834    pub const GROUPBY_MAX: &str = "groupby_max";
835    pub const GROUPBY_MAX_U64: &str = "groupby_max_u64";
836    pub const GROUPBY_LOGSUMEXP_MAX: &str = "groupby_logsumexp_max";
837    pub const GROUPBY_LOGSUMEXP_SUMEXP: &str = "groupby_logsumexp_sumexp";
838    pub const GROUPBY_LOGSUMEXP_FINAL: &str = "groupby_logsumexp_final";
839}
840
841/// Kernel function names in the scan module
842pub mod scan_kernels {
843    pub const BLOCK_INCLUSIVE_SCAN: &str = "block_inclusive_scan";
844    pub const ADD_BLOCK_OFFSETS: &str = "add_block_offsets";
845    pub const EXCLUSIVE_SCAN_MASK: &str = "exclusive_scan_mask";
846    pub const COUNT_MASK: &str = "count_mask";
847    // Multi-block scan kernels for large prefix sums
848    pub const MULTIBLOCK_SCAN_PHASE1: &str = "multiblock_scan_phase1";
849    pub const MULTIBLOCK_SCAN_U32_PHASE1: &str = "multiblock_scan_u32_phase1";
850    pub const MULTIBLOCK_SCAN_PHASE2: &str = "multiblock_scan_phase2";
851    pub const MULTIBLOCK_SCAN_PHASE3: &str = "multiblock_scan_phase3";
852}
853
854/// Kernel function names in the sort module
855pub mod sort_kernels {
856    pub const RADIX_HISTOGRAM: &str = "radix_histogram";
857    pub const RADIX_SCATTER: &str = "radix_scatter";
858    pub const COMPUTE_RANKS: &str = "compute_ranks";
859    pub const RADIX_SCATTER_STABLE: &str = "radix_scatter_stable";
860    pub const COMPUTE_DIGIT_PREFIX_SUMS: &str = "compute_digit_prefix_sums";
861    pub const INIT_INDICES: &str = "init_indices";
862    pub const APPLY_PERMUTATION_U32: &str = "apply_permutation_u32";
863    pub const APPLY_PERMUTATION_BYTES: &str = "apply_permutation_bytes";
864
865    pub const GATHER_KEYS_I32_ORDERED_U32: &str = "gather_keys_i32_ordered_u32";
866    pub const GATHER_KEYS_F32_ORDERED_U32: &str = "gather_keys_f32_ordered_u32";
867    pub const GATHER_KEYS_BOOL_ORDERED_U32: &str = "gather_keys_bool_ordered_u32";
868
869    pub const GATHER_KEYS_U64_LO_U32: &str = "gather_keys_u64_lo_u32";
870    pub const GATHER_KEYS_U64_HI_U32: &str = "gather_keys_u64_hi_u32";
871
872    pub const GATHER_KEYS_I64_LO_U32: &str = "gather_keys_i64_lo_u32";
873    pub const GATHER_KEYS_I64_HI_U32: &str = "gather_keys_i64_hi_u32";
874
875    pub const GATHER_KEYS_F64_LO_U32: &str = "gather_keys_f64_lo_u32";
876    pub const GATHER_KEYS_F64_HI_U32: &str = "gather_keys_f64_hi_u32";
877    /// Sort-merge sortedness-detection kernel — single-pass adjacent-
878    /// pair check; atomically writes 0 to a u32 flag on
879    /// `keys[i] > keys[i+1]`. Caller initializes flag to 1
880    /// before launch, reads result post-launch. Used by the
881    /// dispatch-site eligibility check at `execute_join` to
882    /// validate caller-asserted sortedness before invoking
883    /// `sort_merge_join_v2_inner_u32_1key`.
884    pub const CHECK_ASCENDING_SORTED_U32: &str = "check_ascending_sorted_u32";
885}
886
887/// Kernel function names in the filter module
888pub mod filter_kernels {
889    pub const FILTER_COMPARE_U32: &str = "filter_compare_u32";
890    pub const FILTER_COMPARE_I64: &str = "filter_compare_i64";
891    pub const FILTER_COMPARE_F64: &str = "filter_compare_f64";
892    pub const FILTER_COMPARE_I32: &str = "filter_compare_i32";
893    pub const FILTER_COMPARE_U64: &str = "filter_compare_u64";
894    pub const FILTER_COMPARE_F32: &str = "filter_compare_f32";
895    pub const FILTER_COMPARE_U8: &str = "filter_compare_u8";
896    pub const FILTER_COMPARE_U32_SCAN_PHASE1: &str = "filter_compare_u32_scan_phase1";
897    pub const FILTER_COMPARE_F64_SCAN_PHASE1: &str = "filter_compare_f64_scan_phase1";
898    pub const FILTER_COMPARE_F32_SCAN_PHASE1: &str = "filter_compare_f32_scan_phase1";
899    pub const FILTER_COMPARE_U32_COL: &str = "filter_compare_u32_col";
900    pub const FILTER_COMPARE_I32_COL: &str = "filter_compare_i32_col";
901    pub const FILTER_COMPARE_I64_COL: &str = "filter_compare_i64_col";
902    pub const FILTER_COMPARE_U64_COL: &str = "filter_compare_u64_col";
903    pub const FILTER_COMPARE_F32_COL: &str = "filter_compare_f32_col";
904    pub const FILTER_COMPARE_F64_COL: &str = "filter_compare_f64_col";
905    pub const FILTER_COMPARE_U8_COL: &str = "filter_compare_u8_col";
906    pub const FILL_U32_IOTA: &str = "fill_u32_iota";
907    pub const FILL_U32_CONST: &str = "fill_u32_const";
908    pub const MARK_RANDOM_VARS: &str = "mark_random_vars";
909    pub const RANDOM_VAR_TO_BIT_FROM_LIST: &str = "random_var_to_bit_from_list";
910    pub const CHECK_RANDOM_VAR_COUNT: &str = "check_random_var_count";
911    pub const COMPACT_U32_BY_MASK: &str = "compact_u32_by_mask";
912    pub const COMPACT_I64_BY_MASK: &str = "compact_i64_by_mask";
913    pub const COMPACT_F64_BY_MASK: &str = "compact_f64_by_mask";
914    pub const COMPACT_BYTES_BY_MASK: &str = "compact_bytes_by_mask";
915    pub const CAPTURE_COMPACT_COUNT: &str = "capture_compact_count";
916    pub const MASK_CLAMP_ROWS: &str = "mask_clamp_rows";
917    pub const MASK_AND: &str = "mask_and";
918    pub const MASK_OR: &str = "mask_or";
919    pub const MASK_NOT: &str = "mask_not";
920}
921
922/// Kernel function names in the set_ops module
923pub mod set_ops_kernels {
924    pub const CONCAT_U32: &str = "concat_u32";
925    pub const CONCAT_BYTES: &str = "concat_bytes";
926    pub const SORTED_DIFF_MARK: &str = "sorted_diff_mark";
927}
928
929/// Kernel function names in the pack module (GPU-side key packing)
930pub mod pack_kernels {
931    /// Pack multiple columns into row-major byte array
932    pub const PACK_KEYS: &str = "pack_keys";
933    /// Compute FNV-1a hash from packed keys
934    pub const HASH_PACKED_KEYS: &str = "hash_packed_keys";
935    /// Fused pack + hash in single pass (optimal for join key preparation)
936    pub const PACK_AND_HASH_KEYS: &str = "pack_and_hash_keys";
937    /// Fused pack + hash for arbitrary key column counts
938    pub const PACK_AND_HASH_KEYS_GENERIC: &str = "pack_and_hash_keys_generic";
939    /// Vectorized pack for 8-byte aligned columns
940    pub const PACK_KEYS_ALIGNED: &str = "pack_keys_aligned";
941    /// Unpack single column from packed row data
942    pub const UNPACK_COLUMN: &str = "unpack_column";
943    /// Unpack single column with device-resident row count
944    pub const UNPACK_COLUMN_COUNTED: &str = "unpack_column_counted";
945    /// Gather rows from packed data based on index array
946    pub const GATHER_PACKED_ROWS: &str = "gather_packed_rows";
947    /// Gather rows with device-resident row count
948    pub const GATHER_PACKED_ROWS_COUNTED: &str = "gather_packed_rows_counted";
949    /// Scatter write: distribute packed rows to non-contiguous output positions
950    pub const SCATTER_PACKED_ROWS: &str = "scatter_packed_rows";
951    /// Compare packed keys for equality
952    pub const COMPARE_PACKED_KEYS: &str = "compare_packed_keys";
953    /// Pack u8 bools into Arrow bitmap bytes
954    pub const PACK_BOOLS_TO_BITMAP: &str = "pack_bools_to_bitmap";
955}
956
957/// Kernel function names in the circuit module
958pub mod circuit_kernels {
959    pub const XGCF_FORWARD_LEVEL: &str = "xgcf_forward_level";
960    pub const XGCF_BACKWARD_LEVEL_PROPAGATE: &str = "xgcf_backward_level_propagate";
961    pub const XGCF_BACKWARD_LEVEL_DECISION_GRAD: &str = "xgcf_backward_level_decision_grad";
962    pub const XGCF_BACKWARD_LEVEL_LIT_GRAD: &str = "xgcf_backward_level_lit_grad";
963    pub const XGCF_FREE_VAR_APPLY_GRAD: &str = "xgcf_free_var_apply_grad";
964    pub const XGCF_FREE_VAR_REDUCE_STAGE: &str = "xgcf_free_var_reduce_stage";
965    pub const XGCF_ADD_SCALAR: &str = "xgcf_add_scalar";
966    pub const XGCF_FORWARD_LEVEL_CACHED: &str = "xgcf_forward_level_cached";
967    pub const XGCF_EVAL_ALL_LEVELS_CACHED: &str = "xgcf_eval_all_levels_cached";
968    pub const XGCF_EVAL_ALL_LEVELS_CACHED_BATCHED: &str = "xgcf_eval_all_levels_cached_batched";
969    pub const XGCF_BACKWARD_LEVEL_PROPAGATE_CACHED: &str = "xgcf_backward_level_propagate_cached";
970    pub const XGCF_BACKWARD_LEVEL_DECISION_GRAD_CACHED: &str =
971        "xgcf_backward_level_decision_grad_cached";
972    pub const XGCF_BACKWARD_LEVEL_LIT_GRAD_CACHED: &str = "xgcf_backward_level_lit_grad_cached";
973    pub const XGCF_BACKWARD_ALL_LEVELS_CACHED: &str = "xgcf_backward_all_levels_cached";
974    pub const XGCF_BACKWARD_ALL_LEVELS_CACHED_BATCHED: &str =
975        "xgcf_backward_all_levels_cached_batched";
976    pub const XGCF_FREE_VAR_APPLY_GRAD_CACHED: &str = "xgcf_free_var_apply_grad_cached";
977    pub const XGCF_FREE_VAR_REDUCE_STAGE_CACHED: &str = "xgcf_free_var_reduce_stage_cached";
978    pub const XGCF_ADD_SCALAR_CACHED: &str = "xgcf_add_scalar_cached";
979    pub const XGCF_SET_ROOT_ADJ_CACHED_BATCHED: &str = "xgcf_set_root_adj_cached_batched";
980    pub const XGCF_COPY_ROOT_CACHED: &str = "xgcf_copy_root_cached";
981    pub const XGCF_COPY_ROOT_CACHED_META: &str = "xgcf_copy_root_cached_meta";
982    pub const XGCF_COPY_ROOT_CACHED_META_BATCHED: &str = "xgcf_copy_root_cached_meta_batched";
983}
984
985/// Kernel function names in the cache module
986pub mod cache_kernels {
987    pub const CACHE_CNF_HASH: &str = "cache_cnf_hash";
988    pub const CACHE_LOOKUP_OR_INSERT: &str = "cache_lookup_or_insert";
989    pub const CACHE_EVICT_LRU: &str = "cache_evict_lru";
990    pub const CACHE_STORE_U8: &str = "cache_store_u8";
991    pub const CACHE_STORE_U32: &str = "cache_store_u32";
992    pub const CACHE_STORE_I32: &str = "cache_store_i32";
993    pub const CACHE_STORE_F64: &str = "cache_store_f64";
994    pub const CACHE_STORE_META: &str = "cache_store_meta";
995}
996
997/// Kernel function names in the SAT module
998pub mod sat_kernels {
999    pub const SAT_CDCL_SOLVE: &str = "sat_cdcl_solve";
1000    pub const SAT_CHECK_MODEL: &str = "sat_check_model";
1001    pub const SAT_PROOF_MARK_NEEDED: &str = "sat_proof_mark_needed";
1002    pub const SAT_PROOF_CHECK: &str = "sat_proof_check";
1003    pub const SAT_ASSERT_STATUS: &str = "sat_assert_status";
1004    pub const SAT_ASSERT_OK: &str = "sat_assert_ok";
1005    pub const SAT_XGCF_CNF_COUNTS: &str = "sat_xgcf_cnf_counts";
1006    pub const SAT_XGCF_CNF_EMIT: &str = "sat_xgcf_cnf_emit";
1007    pub const SAT_XGCF_CNF_CAPTURE_LAST_COUNTS: &str = "sat_xgcf_cnf_capture_last_counts";
1008    pub const SAT_XGCF_CNF_COMPUTE_TOTALS: &str = "sat_xgcf_cnf_compute_totals";
1009    pub const SAT_CNF_WRITE_TERMINATOR: &str = "sat_cnf_write_terminator";
1010    pub const SAT_CNF_COPY_INTO: &str = "sat_cnf_copy_into";
1011    pub const SAT_SHIFT_OFFSETS: &str = "sat_shift_offsets";
1012    pub const SAT_XGCF_WRITE_ROOT_UNIT_CLAUSE: &str = "sat_xgcf_write_root_unit_clause";
1013    pub const SAT_NOT_PHI_COUNTS: &str = "sat_not_phi_counts";
1014    pub const SAT_EMIT_NOT_PHI: &str = "sat_emit_not_phi";
1015}
1016
1017/// Default maximum output size for join operations.
1018/// This prevents memory overflow when joining large tables with high cardinality matches.
1019pub const DEFAULT_JOIN_MAX_OUTPUT: usize = 1_000_000;
1020
1021/// Nested-loop join eligibility threshold (Cartesian product
1022/// upper bound). The dispatcher routes to nested-loop iff
1023/// `num_left * num_right <= NESTED_LOOP_TOTAL_THRESHOLD`; the
1024/// provider validates the same invariant fail-closed before any
1025/// allocation.
1026///
1027/// This is the **single source of truth** for the threshold.
1028/// `xlog-runtime`'s dispatch site imports this constant; do NOT
1029/// redeclare in xlog-runtime (would create either drift risk or
1030/// a reverse `xlog-cuda → xlog-runtime` dep cycle).
1031///
1032/// Value (`4_000_000`) is grounded in the bench-spike at
1033/// `bench-spike/w42-nested-loop` HEAD `9c0cefc6` (see
1034/// `docs/evidence/2026-05-07-w42-bench-spike/README.md`):
1035/// largest symmetric tested cell `L=R=2000` → 4M total wins by
1036/// 5.41× over hash; the algorithmic crossover is extrapolated to
1037/// ~10000×10000 = 100M; 4M leaves 6× margin to absorb
1038/// production-kernel cost asymmetry. The threshold also caps the
1039/// index-array allocation at 32 MB total (4M × 4 bytes × 2
1040/// arrays).
1041pub const NESTED_LOOP_TOTAL_THRESHOLD: u64 = 4_000_000;
1042
1043/// Comparison operators for filtering
1044#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1045#[repr(u8)]
1046pub enum CompareOp {
1047    Eq = 0,
1048    Ne = 1,
1049    Lt = 2,
1050    Le = 3,
1051    Gt = 4,
1052    Ge = 5,
1053}
1054
1055/// Join types for hash_join_v2
1056#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1057pub enum JoinType {
1058    /// Inner join: return rows where keys match on both sides
1059    Inner,
1060    /// Semi join: return left rows that have any match in right (no right columns)
1061    Semi,
1062    /// Anti join: return left rows that have NO match in right
1063    Anti,
1064    /// Left outer join: return all left rows, with nulls for non-matching right
1065    LeftOuter,
1066}
1067
1068/// Result of packing key columns and computing hashes for join operations
1069struct PackedKeyData {
1070    /// Computed hash values (one per row)
1071    hashes: crate::memory::TrackedCudaSlice<u64>,
1072    /// Packed key data in row-major format
1073    packed_keys: crate::memory::TrackedCudaSlice<u8>,
1074    /// Total bytes per row (key stride)
1075    key_bytes: u32,
1076}
1077
1078struct JoinHashTableV2 {
1079    bucket_counts: crate::memory::TrackedCudaSlice<u32>,
1080    bucket_offsets: crate::memory::TrackedCudaSlice<u32>,
1081    bucket_entries: crate::memory::TrackedCudaSlice<u32>,
1082    bucket_entry_hashes: crate::memory::TrackedCudaSlice<u64>,
1083    bucket_mask: u32,
1084}
1085
1086/// Bucketed hash table for u64 hashes.
1087pub struct HashTableU64 {
1088    pub bucket_counts: crate::memory::TrackedCudaSlice<u32>,
1089    pub bucket_offsets: crate::memory::TrackedCudaSlice<u32>,
1090    pub bucket_entries: crate::memory::TrackedCudaSlice<u32>,
1091    pub bucket_entry_hashes: crate::memory::TrackedCudaSlice<u64>,
1092    pub bucket_mask: u32,
1093}
1094
1095/// Cached build-side join index for v2 hash join.
1096///
1097/// This captures the packed key bytes and bucketed hash table layout for the build (right) side,
1098/// enabling reuse across repeated joins on the same relation + key columns.
1099pub struct JoinIndexV2 {
1100    right_num_rows: u32,
1101    right_keys: Vec<usize>,
1102    key_bytes: u32,
1103    packed_keys: crate::memory::TrackedCudaSlice<u8>,
1104    table: JoinHashTableV2,
1105}
1106
1107impl JoinIndexV2 {
1108    /// Key columns (indices) this index was built for.
1109    pub fn right_keys(&self) -> &[usize] {
1110        &self.right_keys
1111    }
1112
1113    /// Row count of the build-side buffer at index build time.
1114    pub fn right_num_rows(&self) -> u32 {
1115        self.right_num_rows
1116    }
1117
1118    /// Approximate device memory used by this cached index.
1119    pub fn estimated_bytes(&self) -> u64 {
1120        let mut bytes = 0u64;
1121        bytes = bytes.saturating_add(self.packed_keys.len() as u64);
1122        bytes = bytes.saturating_add(self.table.bucket_counts.len() as u64 * 4);
1123        bytes = bytes.saturating_add(self.table.bucket_offsets.len() as u64 * 4);
1124        bytes = bytes.saturating_add(self.table.bucket_entries.len() as u64 * 4);
1125        bytes = bytes.saturating_add(self.table.bucket_entry_hashes.len() as u64 * 8);
1126        bytes
1127    }
1128}
1129
1130/// CUDA kernel provider for xlog GPU operations
1131///
1132/// Manages pre-compiled PTX modules for relational operations:
1133/// - **Join**: Hash join with build/probe phases
1134/// - **Dedup**: Sort-based deduplication with prefix-sum compaction
1135/// - **GroupBy**: Sorted-input group aggregation (count, sum, min, max)
1136///
1137/// PTX modules are loaded at construction time and stored in the CUDA device.
1138/// Kernel functions can be retrieved using `device.get_func()`.
1139///
1140/// # Example
1141/// ```ignore
1142/// use std::sync::Arc;
1143/// use xlog_cuda::{CudaDevice, GpuMemoryManager, CudaKernelProvider};
1144/// use xlog_core::MemoryBudget;
1145///
1146/// let device = Arc::new(CudaDevice::new(0)?);
1147/// let memory = Arc::new(GpuMemoryManager::new(device.clone(), MemoryBudget::default()));
1148/// let provider = CudaKernelProvider::new(device, memory)?;
1149/// ```
1150pub struct CudaKernelProvider {
1151    /// The CUDA device with loaded PTX modules
1152    device: Arc<CudaDevice>,
1153    /// GPU memory manager for kernel allocations
1154    memory: Arc<GpuMemoryManager>,
1155    /// Tracked host transfers for diagnostics
1156    transfer_tracker: HostTransferTracker,
1157    /// PTX load profiling data (populated only when XLOG_WARMUP_PROFILE=1)
1158    ptx_load_profile: Option<PtxLoadProfile>,
1159    /// Column-level D2H transfer counter (incremented by each download_column_* call)
1160    d2h_transfer_count: AtomicU64,
1161    /// Untracked control-plane metadata D2H read counter. Incremented by every
1162    /// `dtoh_scalar_untracked` / `dtoh_small_metadata_untracked` call. These are
1163    /// bounded metadata reads (row counts, scan totals) exempt from the
1164    /// data-plane transfer contract, but the GPU-resident MC engine's no-host
1165    /// gate must prove they are *also* zero inside the measured region — hence an
1166    /// explicit, resettable counter.
1167    untracked_metadata_dtoh_count: AtomicU64,
1168    /// Strict deterministic-Datalog D2H gate. When `true`, any data-plane D2H
1169    /// transfer (column downloads or `dtoh_sync_copy_into_tracked`) increments
1170    /// the violation counter and returns `XlogError::Execution` from the
1171    /// originating call. Metadata reads via `dtoh_scalar_untracked` are NOT
1172    /// gated. See [`CudaKernelProvider::enable_strict_deterministic_d2h`].
1173    strict_deterministic_d2h: AtomicBool,
1174    /// Cumulative count of deterministic-D2H gate violations observed since
1175    /// the last reset. Increments even on the failing path (the originating
1176    /// call still returns `Err`); kept for telemetry and tests.
1177    deterministic_d2h_violations: AtomicU64,
1178    /// Lazy-initialized non-default launch stream used by
1179    /// env-gated recorded-operator dispatch (filter, sort,
1180    /// dedup, GroupBy, hash-join). Cached for the provider's
1181    /// lifetime — the [`crate::device_runtime::StreamPool`]
1182    /// never returns streams to a free-list, so per-call
1183    /// acquire would saturate it. One stream per provider is
1184    /// sufficient because the recorder serializes work on it;
1185    /// multiple operations chain through commit-order events.
1186    recorded_op_stream: OnceLock<crate::device_runtime::StreamId>,
1187    /// Test/diagnostic-only counter for CSM (count-scan-materialize)
1188    /// invocations selected by the recorded hash-join dispatch.
1189    /// **Not part of any public stability guarantee** — its existence,
1190    /// shape, exposure, and increment semantics may change in any
1191    /// release. Used by the env-dispatch test suite to prove that CSM
1192    /// was actually selected for eligible Inner / LeftOuter cases (and
1193    /// not selected for Semi / Anti or when the env gate is off).
1194    csm_invocations: AtomicU64,
1195    /// Diagnostic counter for bounded CSM CUDA Graph captures.
1196    csm_cuda_graph_captures: AtomicU64,
1197    /// Diagnostic counter for bounded CSM CUDA Graph launches.
1198    csm_cuda_graph_launches: AtomicU64,
1199    /// Diagnostic counter for bounded CSM CUDA Graph ineligibility fallbacks.
1200    csm_cuda_graph_fallbacks: AtomicU64,
1201    /// Diagnostic counter for bounded CSM CUDA Graph cache replays.
1202    csm_cuda_graph_cache_hits: AtomicU64,
1203    /// Diagnostic counter for graph-mode small full-row set-maintenance
1204    /// sorts. This is test telemetry only; production correctness must not
1205    /// depend on the value.
1206    small_full_row_sort_invocations: AtomicU64,
1207    /// Bounded CSM CUDA Graph replay cache.
1208    csm_cuda_graph_cache: Mutex<HashMap<CsmCudaGraphKey, CsmCudaGraphEntry>>,
1209    /// Per-process counter of WCOJ layout fast-path hits. The
1210    /// fast-path skips `dedup_full_row_recorded` when the input
1211    /// is already strictly lex-sorted and full-row unique.
1212    /// Tests + the phase report binary read this counter to
1213    /// confirm the fast-path actually fired vs. silently fell
1214    /// through to the existing dedup pipeline.
1215    wcoj_layout_fast_path_hit_count: AtomicU64,
1216    /// Diagnostic counter for generic WCOJ layout-sort helper
1217    /// invocations. Used by K-clique dispatch-plan certifications to
1218    /// prove K-clique runtime dispatch no longer routes every edge
1219    /// through the old all-edge `wcoj_layout_sort_*_recorded` path.
1220    wcoj_layout_sort_invocation_count: AtomicU64,
1221    /// Diagnostic counter for K-clique leader-edge metadata builds.
1222    kclique_metadata_build_count: AtomicU64,
1223    /// Diagnostic counter for cumulative nanoseconds spent building K-clique
1224    /// leader-edge metadata.
1225    kclique_metadata_build_nanos: AtomicU64,
1226    /// Histogram-guided triangle WCOJ routing counter: successful dispatches
1227    /// accepted through the block-slice provider entry.
1228    wcoj_triangle_hg_dispatch_count: AtomicU64,
1229    /// Diagnostic-only: last WCOJ triangle dispatch's per-phase
1230    /// CUDA-event timings, populated by `wcoj_triangle_*_recorded`
1231    /// when the `wcoj-phase-timing` Cargo feature is on. Read by
1232    /// the `wcoj_phase_report` binary in xlog-integration. Field
1233    /// is absent when the feature is off, so production builds
1234    /// have zero overhead.
1235    #[cfg(feature = "wcoj-phase-timing")]
1236    last_triangle_phase_timing:
1237        std::sync::Mutex<Option<crate::wcoj_phase_timing::WcojTrianglePhaseTiming>>,
1238}
1239
1240#[derive(Default)]
1241struct HostTransferTracker {
1242    dtoh_bytes: AtomicU64,
1243    htod_bytes: AtomicU64,
1244    dtoh_calls: AtomicU64,
1245    htod_calls: AtomicU64,
1246    launch_metadata_htod_bytes: AtomicU64,
1247    launch_metadata_htod_calls: AtomicU64,
1248}
1249
1250#[derive(Debug, Clone, Copy)]
1251pub struct HostTransferStats {
1252    pub dtoh_bytes: u64,
1253    pub htod_bytes: u64,
1254    pub dtoh_calls: u64,
1255    pub htod_calls: u64,
1256}
1257
1258#[derive(Debug, Clone, Copy, Default)]
1259pub struct HostLaunchMetadataTransferStats {
1260    pub htod_bytes: u64,
1261    pub htod_calls: u64,
1262}
1263
1264impl HostTransferTracker {
1265    fn record_dtoh(&self, bytes: u64) {
1266        self.dtoh_calls.fetch_add(1, Ordering::Relaxed);
1267        self.dtoh_bytes.fetch_add(bytes, Ordering::Relaxed);
1268    }
1269
1270    fn record_htod(&self, bytes: u64) {
1271        self.htod_calls.fetch_add(1, Ordering::Relaxed);
1272        self.htod_bytes.fetch_add(bytes, Ordering::Relaxed);
1273    }
1274
1275    fn record_htod_launch_metadata(&self, bytes: u64) {
1276        self.launch_metadata_htod_calls
1277            .fetch_add(1, Ordering::Relaxed);
1278        self.launch_metadata_htod_bytes
1279            .fetch_add(bytes, Ordering::Relaxed);
1280    }
1281
1282    fn snapshot(&self) -> HostTransferStats {
1283        HostTransferStats {
1284            dtoh_bytes: self.dtoh_bytes.load(Ordering::Relaxed),
1285            htod_bytes: self.htod_bytes.load(Ordering::Relaxed),
1286            dtoh_calls: self.dtoh_calls.load(Ordering::Relaxed),
1287            htod_calls: self.htod_calls.load(Ordering::Relaxed),
1288        }
1289    }
1290
1291    fn launch_metadata_snapshot(&self) -> HostLaunchMetadataTransferStats {
1292        HostLaunchMetadataTransferStats {
1293            htod_bytes: self.launch_metadata_htod_bytes.load(Ordering::Relaxed),
1294            htod_calls: self.launch_metadata_htod_calls.load(Ordering::Relaxed),
1295        }
1296    }
1297
1298    fn reset(&self) {
1299        self.dtoh_bytes.store(0, Ordering::Relaxed);
1300        self.htod_bytes.store(0, Ordering::Relaxed);
1301        self.dtoh_calls.store(0, Ordering::Relaxed);
1302        self.htod_calls.store(0, Ordering::Relaxed);
1303        self.launch_metadata_htod_bytes.store(0, Ordering::Relaxed);
1304        self.launch_metadata_htod_calls.store(0, Ordering::Relaxed);
1305    }
1306}
1307
1308impl CudaKernelProvider {
1309    /// Create a new CUDA kernel provider
1310    ///
1311    /// Loads all kernel modules into the CUDA device.
1312    /// Prefers cubin for the detected SM arch, falls back to portable PTX (sm_75+).
1313    ///
1314    /// # Arguments
1315    /// * `device` - The CUDA device to load modules into
1316    /// * `memory` - The GPU memory manager for kernel allocations
1317    ///
1318    /// # Errors
1319    /// Returns `XlogError::Kernel` if PTX loading fails
1320    ///
1321    /// # Example
1322    /// ```ignore
1323    /// let device = Arc::new(CudaDevice::new(0)?);
1324    /// let memory = Arc::new(GpuMemoryManager::new(device.clone(), MemoryBudget::default()));
1325    /// let provider = CudaKernelProvider::new(device, memory)?;
1326    /// ```
1327    pub fn new(device: Arc<CudaDevice>, memory: Arc<GpuMemoryManager>) -> Result<Self> {
1328        let profiling = warmup_profiling_enabled();
1329        let ptx_load_profile = Self::load_all_kernel_modules(&device, profiling)?;
1330
1331        Ok(Self {
1332            device,
1333            memory,
1334            transfer_tracker: HostTransferTracker::default(),
1335            ptx_load_profile,
1336            d2h_transfer_count: AtomicU64::new(0),
1337            untracked_metadata_dtoh_count: AtomicU64::new(0),
1338            strict_deterministic_d2h: AtomicBool::new(false),
1339            deterministic_d2h_violations: AtomicU64::new(0),
1340            recorded_op_stream: OnceLock::new(),
1341            csm_invocations: AtomicU64::new(0),
1342            csm_cuda_graph_captures: AtomicU64::new(0),
1343            csm_cuda_graph_launches: AtomicU64::new(0),
1344            csm_cuda_graph_fallbacks: AtomicU64::new(0),
1345            csm_cuda_graph_cache_hits: AtomicU64::new(0),
1346            small_full_row_sort_invocations: AtomicU64::new(0),
1347            csm_cuda_graph_cache: Mutex::new(HashMap::new()),
1348            wcoj_layout_fast_path_hit_count: AtomicU64::new(0),
1349            wcoj_layout_sort_invocation_count: AtomicU64::new(0),
1350            kclique_metadata_build_count: AtomicU64::new(0),
1351            kclique_metadata_build_nanos: AtomicU64::new(0),
1352            wcoj_triangle_hg_dispatch_count: AtomicU64::new(0),
1353            #[cfg(feature = "wcoj-phase-timing")]
1354            last_triangle_phase_timing: std::sync::Mutex::new(None),
1355        })
1356    }
1357
1358    /// Construct a provider whose `GpuMemoryManager` must already
1359    /// have a v0.6 [`crate::device_runtime::XlogDeviceRuntime`]
1360    /// attached via [`GpuMemoryManager::with_runtime`].
1361    ///
1362    /// Equivalent to [`Self::new`] in every respect — same kernel
1363    /// loading, same field initialization — but **rejects** managers
1364    /// that lack a runtime. This guards against the misconfiguration
1365    /// in which a caller asks for runtime-routed provider semantics
1366    /// (by calling `with_runtime`) but supplies a legacy manager
1367    /// built via [`GpuMemoryManager::new`]; without the check, the
1368    /// resulting provider would silently keep using the cudarc
1369    /// default allocator and the runtime budget/logging stack would
1370    /// never observe the allocations the caller expected to be
1371    /// routed through it.
1372    ///
1373    /// Note: a runtime-routed manager passed to [`Self::new`] still
1374    /// routes correctly — `alloc::<T>` and `alloc_raw` consult
1375    /// `memory.runtime()` regardless of which provider constructor
1376    /// was used. `with_runtime` exists for callers that want the
1377    /// requirement enforced at construction time, not for
1378    /// correctness of the routing itself.
1379    ///
1380    /// This is the **opt-in** runtime entry point for providers.
1381    /// `Self::new` continues to accept managers without a runtime
1382    /// (the legacy default) and remains the production constructor
1383    /// until the runtime stack is certified end-to-end.
1384    ///
1385    /// # Errors
1386    /// Returns `XlogError::Kernel` if `memory.runtime()` is `None`,
1387    /// or anything `Self::new` would return.
1388    ///
1389    /// # Example
1390    /// ```ignore
1391    /// let device = Arc::new(CudaDevice::new(0)?);
1392    /// let runtime = Arc::new(XlogDeviceRuntime::with_resource(
1393    ///     Arc::clone(&device),
1394    ///     0,
1395    ///     Arc::new(StreamPool::with_defaults(Arc::clone(&device))),
1396    ///     Box::new(AsyncCudaResource::new(/* ... */)),
1397    /// ));
1398    /// let memory = Arc::new(GpuMemoryManager::with_runtime(
1399    ///     Arc::clone(&device),
1400    ///     MemoryBudget::default(),
1401    ///     runtime,
1402    /// ));
1403    /// let provider = CudaKernelProvider::with_runtime(device, memory)?;
1404    /// ```
1405    pub fn with_runtime(device: Arc<CudaDevice>, memory: Arc<GpuMemoryManager>) -> Result<Self> {
1406        if memory.runtime().is_none() {
1407            return Err(XlogError::Kernel(
1408                "CudaKernelProvider::with_runtime requires a GpuMemoryManager built via \
1409                 GpuMemoryManager::with_runtime; got a manager with no runtime attached"
1410                    .to_string(),
1411            ));
1412        }
1413        Self::new(device, memory)
1414    }
1415
1416    /// Internal: parse a "boolean" env var. Empty / unset / `"0"`
1417    /// → false; any other value → true.
1418    fn env_flag(name: &str) -> bool {
1419        std::env::var(name)
1420            .map(|v| !v.is_empty() && v != "0")
1421            .unwrap_or(false)
1422    }
1423
1424    /// Whether the recorded filter dispatch is enabled via env.
1425    ///
1426    /// Returns `true` when either `XLOG_USE_RECORDED_FILTERS` or
1427    /// the umbrella `XLOG_USE_RECORDED_OPS` env var is set.
1428    /// Combined with a runtime-backed manager, this routes
1429    /// `filter::<T>` through the recorded launch path.
1430    ///
1431    /// Env-gated rather than default-on so the migration is
1432    /// opt-in for real callers; the existing legacy paths remain
1433    /// the production default until the runtime stack is
1434    /// certified end-to-end.
1435    pub(crate) fn use_recorded_filters_env() -> bool {
1436        Self::env_flag("XLOG_USE_RECORDED_FILTERS") || Self::env_flag("XLOG_USE_RECORDED_OPS")
1437    }
1438
1439    /// Whether the recorded sort dispatch is enabled via env.
1440    /// Reads `XLOG_USE_RECORDED_SORT` or the umbrella
1441    /// `XLOG_USE_RECORDED_OPS`. The recorded-sort path is narrowed
1442    /// to U32 / Symbol keys only — the public
1443    /// `sort()` dispatcher checks both this env flag AND key
1444    /// type compatibility before routing.
1445    pub(crate) fn use_recorded_sort_env() -> bool {
1446        Self::env_flag("XLOG_USE_RECORDED_SORT") || Self::env_flag("XLOG_USE_RECORDED_OPS")
1447    }
1448
1449    /// Whether the recorded full-row dedup dispatch is enabled
1450    /// via env. Reads `XLOG_USE_RECORDED_DEDUP` or the umbrella
1451    /// `XLOG_USE_RECORDED_OPS`. `dedup_full_row_recorded` is
1452    /// narrow to all-U32 / Symbol columns.
1453    pub(crate) fn use_recorded_dedup_env() -> bool {
1454        Self::env_flag("XLOG_USE_RECORDED_DEDUP") || Self::env_flag("XLOG_USE_RECORDED_OPS")
1455    }
1456
1457    /// Whether the recorded GroupBy dispatch is enabled via
1458    /// env. Reads `XLOG_USE_RECORDED_GROUPBY` or
1459    /// `XLOG_USE_RECORDED_OPS`. `groupby_multi_agg_recorded`
1460    /// supports U32 / Symbol keys + Count / Sum / Min / Max
1461    /// aggs only.
1462    pub(crate) fn use_recorded_groupby_env() -> bool {
1463        Self::env_flag("XLOG_USE_RECORDED_GROUPBY") || Self::env_flag("XLOG_USE_RECORDED_OPS")
1464    }
1465
1466    /// Whether the recorded hash-join dispatch is enabled via
1467    /// env. Reads `XLOG_USE_RECORDED_HASH_JOIN` or
1468    /// `XLOG_USE_RECORDED_OPS`. `hash_join_v2_recorded` and
1469    /// `hash_join_v2_with_index_recorded` cover all four join
1470    /// types (Inner / Semi / Anti / LeftOuter); the only
1471    /// hard constraint inherited from `pack_keys` is `≤4`
1472    /// key columns.
1473    pub(crate) fn use_recorded_hash_join_env() -> bool {
1474        Self::env_flag("XLOG_USE_RECORDED_HASH_JOIN") || Self::env_flag("XLOG_USE_RECORDED_OPS")
1475    }
1476
1477    /// Whether the recorded CSM (count-scan-materialize)
1478    /// dispatch is enabled via env. Reads `XLOG_USE_RECORDED_CSM`
1479    /// or `XLOG_USE_RECORDED_OPS`. CSM is a sub-strategy of the
1480    /// recorded hash-join: it is consulted only after the
1481    /// recorded path has already been selected, and only for
1482    /// `JoinType::Inner` / `JoinType::LeftOuter` where a CSM
1483    /// implementation exists. `Semi` / `Anti` are not affected.
1484    pub(crate) fn use_recorded_csm_env() -> bool {
1485        Self::env_flag("XLOG_USE_RECORDED_CSM") || Self::env_flag("XLOG_USE_RECORDED_OPS")
1486    }
1487
1488    /// Whether the bounded CSM CUDA Graph path is enabled.
1489    ///
1490    /// This is narrower than `XLOG_USE_RECORDED_CSM`: callers must first select
1491    /// the recorded CSM hash-join path, then opt into graph capture/replay with
1492    /// `XLOG_USE_CSM_CUDA_GRAPH=1` (or the broader `XLOG_USE_CUDA_GRAPHS=1`).
1493    pub(crate) fn use_csm_cuda_graph_env() -> bool {
1494        Self::env_flag("XLOG_USE_CSM_CUDA_GRAPH") || Self::env_flag("XLOG_USE_CUDA_GRAPHS")
1495    }
1496
1497    /// Test/diagnostic-only telemetry: number of times the recorded
1498    /// hash-join dispatch routed through a CSM (count-scan-materialize)
1499    /// method since this provider was created. Increments once per
1500    /// dispatched call across all four CSM methods (Inner / LeftOuter,
1501    /// non-indexed / indexed). Used by `test_csm_env_dispatch` to
1502    /// prove dispatch selection.
1503    ///
1504    /// **Not part of any public stability guarantee.** Hidden from
1505    /// rustdoc with `#[doc(hidden)]` so it does not appear in
1506    /// generated API docs; the symbol remains callable from
1507    /// integration tests within this crate but production callers
1508    /// must not depend on it. May be renamed, gated behind a cargo
1509    /// feature, or withdrawn in any release without notice.
1510    #[doc(hidden)]
1511    pub fn csm_invocations(&self) -> u64 {
1512        self.csm_invocations.load(Ordering::Relaxed)
1513    }
1514
1515    #[doc(hidden)]
1516    pub fn csm_cuda_graph_captures(&self) -> u64 {
1517        self.csm_cuda_graph_captures.load(Ordering::Relaxed)
1518    }
1519
1520    #[doc(hidden)]
1521    pub fn csm_cuda_graph_launches(&self) -> u64 {
1522        self.csm_cuda_graph_launches.load(Ordering::Relaxed)
1523    }
1524
1525    #[doc(hidden)]
1526    pub fn csm_cuda_graph_fallbacks(&self) -> u64 {
1527        self.csm_cuda_graph_fallbacks.load(Ordering::Relaxed)
1528    }
1529
1530    #[doc(hidden)]
1531    pub fn csm_cuda_graph_cache_hits(&self) -> u64 {
1532        self.csm_cuda_graph_cache_hits.load(Ordering::Relaxed)
1533    }
1534
1535    #[doc(hidden)]
1536    pub fn small_full_row_sort_invocations(&self) -> u64 {
1537        self.small_full_row_sort_invocations.load(Ordering::Relaxed)
1538    }
1539
1540    /// Lazily acquire one non-default launch stream from the
1541    /// runtime's [`crate::device_runtime::StreamPool`] for
1542    /// recorded-operator dispatch, and cache it for this
1543    /// provider's lifetime. Shared across all env-gated
1544    /// recorded paths (filter, sort, dedup, GroupBy,
1545    /// hash-join) — a single stream is sufficient because the
1546    /// recorder serializes work on it; multiple operations
1547    /// chain naturally through commit-order events.
1548    ///
1549    /// Returns `None` when:
1550    ///   * the manager has no runtime attached
1551    ///     (`memory.runtime() == None`), or
1552    ///   * the stream pool is at capacity and `acquire` fails.
1553    ///
1554    /// On a lost race during first init the loser leaks one
1555    /// stream (the pool keeps it alive); both winners cache
1556    /// the same `StreamId`. Acceptable cost — practical pool
1557    /// sizes are large compared to the number of providers
1558    /// per process.
1559    pub(crate) fn recorded_op_stream_or_init(&self) -> Option<crate::device_runtime::StreamId> {
1560        if let Some(s) = self.recorded_op_stream.get() {
1561            return Some(*s);
1562        }
1563        let runtime = self.memory.runtime()?;
1564        let stream = runtime.stream_pool().acquire().ok()?;
1565        let _ = self.recorded_op_stream.set(stream);
1566        self.recorded_op_stream.get().copied()
1567    }
1568
1569    /// Take the per-phase WCOJ triangle dispatch timings recorded
1570    /// by the most recent `wcoj_triangle_*_recorded` call. Reading
1571    /// clears the slot — designed for one-shot consumption by the
1572    /// `wcoj_phase_report` binary in xlog-integration. Returns
1573    /// `None` if no triangle dispatch has fired since the last
1574    /// read (or since construction).
1575    ///
1576    /// Compiled in only with the `wcoj-phase-timing` Cargo
1577    /// feature; production builds have no such method.
1578    #[cfg(feature = "wcoj-phase-timing")]
1579    pub fn take_wcoj_triangle_phase_timing(
1580        &self,
1581    ) -> Option<crate::wcoj_phase_timing::WcojTrianglePhaseTiming> {
1582        self.last_triangle_phase_timing
1583            .lock()
1584            .ok()
1585            .and_then(|mut g| g.take())
1586    }
1587
1588    /// Internal: store the phase timings produced by a triangle
1589    /// dispatch. Overwrites any prior unread slot — the report
1590    /// binary is expected to read after every `execute_plan`.
1591    #[cfg(feature = "wcoj-phase-timing")]
1592    #[allow(dead_code)]
1593    pub(crate) fn put_wcoj_triangle_phase_timing(
1594        &self,
1595        timing: crate::wcoj_phase_timing::WcojTrianglePhaseTiming,
1596    ) {
1597        if let Ok(mut g) = self.last_triangle_phase_timing.lock() {
1598            *g = Some(timing);
1599        }
1600    }
1601
1602    /// Number of times `wcoj_layout_*_recorded` short-circuited
1603    /// to the fast-path (recorded clone) instead of running
1604    /// `dedup_full_row_recorded`. Increments by 1 per
1605    /// fast-path hit (3 hits per dispatch when all inputs are
1606    /// already sorted+unique). Used by tests + the phase
1607    /// report to confirm the fast-path fired.
1608    pub fn wcoj_layout_fast_path_hit_count(&self) -> u64 {
1609        self.wcoj_layout_fast_path_hit_count.load(Ordering::Relaxed)
1610    }
1611
1612    /// Histogram-guided block-slice triangle WCOJ test/diagnostic counter:
1613    /// successful dispatches that routed through the provider entry.
1614    pub fn wcoj_triangle_hg_dispatch_count(&self) -> u64 {
1615        self.wcoj_triangle_hg_dispatch_count.load(Ordering::Relaxed)
1616    }
1617
1618    /// Reset the fast-path hit counter to 0. Tests use this to
1619    /// scope counter assertions to a single dispatch.
1620    pub fn reset_wcoj_layout_fast_path_hit_count(&self) {
1621        self.wcoj_layout_fast_path_hit_count
1622            .store(0, Ordering::Relaxed);
1623    }
1624
1625    /// Number of calls to `wcoj_layout_sort_*_recorded` since the
1626    /// last reset. Diagnostic-only; used by dispatch-plan certification.
1627    pub fn wcoj_layout_sort_invocation_count(&self) -> u64 {
1628        self.wcoj_layout_sort_invocation_count
1629            .load(Ordering::Relaxed)
1630    }
1631
1632    /// Reset the WCOJ layout-sort invocation counter to 0.
1633    pub fn reset_wcoj_layout_sort_invocation_count(&self) {
1634        self.wcoj_layout_sort_invocation_count
1635            .store(0, Ordering::Relaxed);
1636    }
1637
1638    /// Number of K-clique leader-edge metadata builds since the
1639    /// last reset.
1640    pub fn kclique_metadata_build_count(&self) -> u64 {
1641        self.kclique_metadata_build_count.load(Ordering::Relaxed)
1642    }
1643
1644    /// Cumulative nanoseconds spent building K-clique leader-edge
1645    /// metadata since the last reset.
1646    pub fn kclique_metadata_build_nanos(&self) -> u64 {
1647        self.kclique_metadata_build_nanos.load(Ordering::Relaxed)
1648    }
1649
1650    /// Reset K-clique metadata build diagnostics.
1651    pub fn reset_kclique_metadata_build_metrics(&self) {
1652        self.kclique_metadata_build_count
1653            .store(0, Ordering::Relaxed);
1654        self.kclique_metadata_build_nanos
1655            .store(0, Ordering::Relaxed);
1656    }
1657
1658    /// Internal: increment the fast-path counter. Called by
1659    /// `wcoj_layout_*_recorded` after a successful fast-path
1660    /// branch. Not part of any public stability guarantee.
1661    pub(crate) fn record_wcoj_layout_fast_path_hit(&self) {
1662        self.wcoj_layout_fast_path_hit_count
1663            .fetch_add(1, Ordering::Relaxed);
1664    }
1665
1666    /// Internal: increment the generic WCOJ layout-sort counter.
1667    pub(crate) fn record_wcoj_layout_sort_invocation(&self) {
1668        self.wcoj_layout_sort_invocation_count
1669            .fetch_add(1, Ordering::Relaxed);
1670    }
1671
1672    /// Internal: record a K-clique leader-edge metadata build.
1673    pub(crate) fn record_kclique_metadata_build_nanos(&self, nanos: u128) {
1674        self.kclique_metadata_build_count
1675            .fetch_add(1, Ordering::Relaxed);
1676        let nanos = u64::try_from(nanos).unwrap_or(u64::MAX);
1677        self.kclique_metadata_build_nanos
1678            .fetch_add(nanos, Ordering::Relaxed);
1679    }
1680
1681    /// Runtime hook: record a successful histogram-guided block-slice triangle
1682    /// dispatch.
1683    #[doc(hidden)]
1684    pub fn record_wcoj_triangle_hg_dispatch(&self) {
1685        self.wcoj_triangle_hg_dispatch_count
1686            .fetch_add(1, Ordering::Relaxed);
1687    }
1688
1689    /// Get the CUDA device
1690    pub fn device(&self) -> &Arc<CudaDevice> {
1691        &self.device
1692    }
1693
1694    /// Get the GPU memory manager
1695    pub fn memory(&self) -> &Arc<GpuMemoryManager> {
1696        &self.memory
1697    }
1698
1699    /// Get PTX load profiling data (only populated when XLOG_WARMUP_PROFILE=1).
1700    pub fn ptx_load_profile(&self) -> Option<&PtxLoadProfile> {
1701        self.ptx_load_profile.as_ref()
1702    }
1703
1704    /// Reset tracked host transfer statistics.
1705    pub fn reset_host_transfer_stats(&self) {
1706        self.transfer_tracker.reset();
1707    }
1708
1709    /// Snapshot tracked host transfer statistics.
1710    pub fn host_transfer_stats(&self) -> HostTransferStats {
1711        self.transfer_tracker.snapshot()
1712    }
1713
1714    /// Snapshot launch-parameter H2D uploads tracked separately from
1715    /// `host_transfer_stats`.
1716    pub fn host_launch_metadata_transfer_stats(&self) -> HostLaunchMetadataTransferStats {
1717        self.transfer_tracker.launch_metadata_snapshot()
1718    }
1719
1720    /// Read the column-level D2H transfer counter.
1721    ///
1722    /// This counter increments once per `download_column_*` call, enabling
1723    /// callers (e.g. the ILP trainer) to assert that no column downloads
1724    /// occurred during a performance-critical section.
1725    pub fn d2h_transfer_count(&self) -> u64 {
1726        self.d2h_transfer_count.load(Ordering::Relaxed)
1727    }
1728
1729    /// Reset the column-level D2H transfer counter to zero.
1730    pub fn reset_d2h_transfer_count(&self) {
1731        self.d2h_transfer_count.store(0, Ordering::Relaxed);
1732    }
1733
1734    /// Count of untracked control-plane metadata D2H reads
1735    /// (`dtoh_scalar_untracked` + `dtoh_small_metadata_untracked`).
1736    pub fn untracked_metadata_dtoh_count(&self) -> u64 {
1737        self.untracked_metadata_dtoh_count.load(Ordering::Relaxed)
1738    }
1739
1740    /// Reset the untracked metadata D2H read counter to zero.
1741    pub fn reset_untracked_metadata_dtoh_count(&self) {
1742        self.untracked_metadata_dtoh_count
1743            .store(0, Ordering::Relaxed);
1744    }
1745
1746    /// Enable the strict deterministic-Datalog D2H gate.
1747    ///
1748    /// While enabled, any data-plane device-to-host transfer (column downloads
1749    /// via `download_column` / `download_column_untracked`, and any internal
1750    /// transfer routed through `dtoh_sync_copy_into_tracked`) increments
1751    /// [`CudaKernelProvider::deterministic_d2h_violation_count`] and returns
1752    /// `XlogError::Execution` from the originating call.
1753    ///
1754    /// Metadata reads via [`CudaKernelProvider::dtoh_scalar_untracked`] are
1755    /// allowed and never trip the gate.
1756    ///
1757    /// Default is `false`; the runtime opts in via
1758    /// `RuntimeConfig::strict_deterministic_d2h`. v0.5.5 ships the gate
1759    /// opt-in only — known-violating relational paths (set difference,
1760    /// join count/materialize) are scheduled for replacement before the
1761    /// default flips.
1762    pub fn enable_strict_deterministic_d2h(&self) {
1763        self.strict_deterministic_d2h.store(true, Ordering::Relaxed);
1764    }
1765
1766    /// Disable the strict deterministic-Datalog D2H gate.
1767    pub fn disable_strict_deterministic_d2h(&self) {
1768        self.strict_deterministic_d2h
1769            .store(false, Ordering::Relaxed);
1770    }
1771
1772    /// Returns whether the strict deterministic-Datalog D2H gate is enabled.
1773    pub fn strict_deterministic_d2h_enabled(&self) -> bool {
1774        self.strict_deterministic_d2h.load(Ordering::Relaxed)
1775    }
1776
1777    /// Cumulative deterministic-D2H gate violations since the last reset.
1778    pub fn deterministic_d2h_violation_count(&self) -> u64 {
1779        self.deterministic_d2h_violations.load(Ordering::Relaxed)
1780    }
1781
1782    /// Reset the deterministic-D2H violation counter to zero.
1783    pub fn reset_deterministic_d2h_violations(&self) {
1784        self.deterministic_d2h_violations
1785            .store(0, Ordering::Relaxed);
1786    }
1787
1788    /// Chokepoint for the deterministic-D2H gate.
1789    ///
1790    /// If the gate is enabled, increments the violation counter and returns
1791    /// `XlogError::Execution` naming the offending operation and byte count.
1792    /// If the gate is disabled, returns `Ok(())` cheaply.
1793    pub(crate) fn check_deterministic_d2h(&self, op: &'static str, bytes: u64) -> Result<()> {
1794        if self.strict_deterministic_d2h.load(Ordering::Relaxed) {
1795            self.deterministic_d2h_violations
1796                .fetch_add(1, Ordering::Relaxed);
1797            return Err(XlogError::Execution(format!(
1798                "deterministic D2H gate: {} attempted to copy {} bytes from device to host",
1799                op, bytes
1800            )));
1801        }
1802        Ok(())
1803    }
1804
1805    fn dtoh_sync_copy_into_tracked<T: DeviceRepr, Src: DevicePtr<T>>(
1806        &self,
1807        src: &Src,
1808        dst: &mut [T],
1809    ) -> Result<()> {
1810        let bytes = std::mem::size_of::<T>()
1811            .checked_mul(dst.len())
1812            .ok_or_else(|| XlogError::Kernel("dtoh size overflow".to_string()))?;
1813        self.check_deterministic_d2h("dtoh_sync_copy_into_tracked", bytes as u64)?;
1814        self.transfer_tracker.record_dtoh(bytes as u64);
1815        self.device
1816            .inner()
1817            .dtoh_sync_copy_into(src, dst)
1818            .map_err(|e| XlogError::Kernel(format!("Failed to copy from device: {}", e)))
1819    }
1820
1821    /// Hard cap (in bytes) for [`Self::dtoh_small_metadata_untracked`].
1822    /// Set deliberately small (4 KB) so the helper cannot become a
1823    /// general-purpose vector D2H escape hatch — it's strictly for
1824    /// classifier histograms and similar small metadata round-trips.
1825    pub const DTOH_SMALL_METADATA_MAX_BYTES: usize = 4096;
1826
1827    /// Read a small metadata vector (≤ [`Self::DTOH_SMALL_METADATA_MAX_BYTES`])
1828    /// from device to host WITHOUT updating the D2H transfer tracker.
1829    ///
1830    /// Sibling of [`Self::dtoh_scalar_untracked`] for callers that need
1831    /// a few bucket counts (the WCOJ skew classifier reads a 3 × 64 ×
1832    /// `u32` = 768-byte histogram in one go) instead of `count` separate
1833    /// scalar reads. Like `dtoh_scalar_untracked`, this method is
1834    /// whitelisted by the strict deterministic-D2H gate
1835    /// ([`Self::enable_strict_deterministic_d2h`]) — it does NOT trip
1836    /// the gate, on purpose, because metadata reads are part of the
1837    /// determinism contract (just like a scalar `total` after a scan).
1838    ///
1839    /// # Hard contract — DO NOT WIDEN THE CAP
1840    /// The 4 KB cap is the contract. If a caller wants a larger D2H,
1841    /// it's a data-plane transfer and must go through the tracked
1842    /// `download_column*` path. Widening this cap turns the helper
1843    /// into a backdoor for tracked-bypass column reads, which would
1844    /// silently invalidate the strict deterministic-D2H gate.
1845    ///
1846    /// # Errors
1847    ///   * `XlogError::Kernel` if `count * size_of::<T>()` exceeds
1848    ///     `DTOH_SMALL_METADATA_MAX_BYTES`.
1849    ///   * `XlogError::Kernel` if `count` exceeds the device slice's
1850    ///     length, or if the inner sync copy fails.
1851    pub fn dtoh_small_metadata_untracked<T: DeviceRepr + Default + Copy>(
1852        &self,
1853        src: &crate::memory::TrackedCudaSlice<T>,
1854        count: usize,
1855    ) -> Result<Vec<T>> {
1856        let bytes = count.checked_mul(std::mem::size_of::<T>()).ok_or_else(|| {
1857            XlogError::Kernel("dtoh_small_metadata_untracked: byte size overflow".to_string())
1858        })?;
1859        if bytes > Self::DTOH_SMALL_METADATA_MAX_BYTES {
1860            return Err(XlogError::Kernel(format!(
1861                "dtoh_small_metadata_untracked: requested {} bytes exceeds metadata cap of {} bytes \
1862                 (this is metadata-only; use download_column* for data-plane transfers)",
1863                bytes,
1864                Self::DTOH_SMALL_METADATA_MAX_BYTES
1865            )));
1866        }
1867        if count > src.len() {
1868            return Err(XlogError::Kernel(format!(
1869                "dtoh_small_metadata_untracked: count={count} > src.len={}",
1870                src.len()
1871            )));
1872        }
1873        if count == 0 {
1874            return Ok(Vec::new());
1875        }
1876        let slice = src.try_slice(0..count).ok_or_else(|| {
1877            XlogError::Kernel(format!(
1878                "dtoh_small_metadata_untracked: try_slice(0..{count}) failed"
1879            ))
1880        })?;
1881        let mut buf: Vec<T> = vec![T::default(); count];
1882        self.untracked_metadata_dtoh_count
1883            .fetch_add(1, Ordering::Relaxed);
1884        self.device
1885            .inner()
1886            .dtoh_sync_copy_into(&slice, &mut buf)
1887            .map_err(|e| {
1888                XlogError::Kernel(format!("dtoh_small_metadata_untracked: copy failed: {}", e))
1889            })?;
1890        Ok(buf)
1891    }
1892
1893    /// Read a single scalar from device to host WITHOUT updating the
1894    /// D2H transfer tracker. Use ONLY for metadata reads (e.g. total_nnz
1895    /// after an exclusive scan), never for data-plane transfers.
1896    ///
1897    /// This makes the "metadata != data-plane" contract explicit and
1898    /// auditable: callers that bypass tracking must call this method
1899    /// (which is grep-able) rather than reaching for device().inner().
1900    pub fn dtoh_scalar_untracked<T: DeviceRepr + Default + Copy>(
1901        &self,
1902        src: &crate::memory::TrackedCudaSlice<T>,
1903        index: usize,
1904    ) -> Result<T> {
1905        if index >= src.len() {
1906            return Err(XlogError::Kernel(format!(
1907                "dtoh_scalar_untracked: index={} >= len={}",
1908                index,
1909                src.len()
1910            )));
1911        }
1912        let slice = src.try_slice(index..index + 1).ok_or_else(|| {
1913            XlogError::Kernel(format!(
1914                "dtoh_scalar_untracked: slice failed at index={}",
1915                index
1916            ))
1917        })?;
1918        let mut buf = [T::default()];
1919        self.untracked_metadata_dtoh_count
1920            .fetch_add(1, Ordering::Relaxed);
1921        self.device
1922            .inner()
1923            .dtoh_sync_copy_into(&slice, &mut buf)
1924            .map_err(|e| XlogError::Kernel(format!("dtoh_scalar_untracked: copy failed: {}", e)))?;
1925        Ok(buf[0])
1926    }
1927
1928    /// Upload host data to device while recording data-plane H2D transfer stats.
1929    pub fn htod_sync_copy_into_tracked<T: DeviceRepr, Dst: cudarc::driver::DevicePtrMut<T>>(
1930        &self,
1931        src: &[T],
1932        dst: &mut Dst,
1933    ) -> Result<()> {
1934        let bytes = std::mem::size_of::<T>()
1935            .checked_mul(src.len())
1936            .ok_or_else(|| XlogError::Kernel("htod size overflow".to_string()))?;
1937        self.transfer_tracker.record_htod(bytes as u64);
1938        self.device
1939            .inner()
1940            .htod_sync_copy_into(src, dst)
1941            .map_err(|e| XlogError::Kernel(format!("Failed to copy to device: {}", e)))
1942    }
1943
1944    /// Allocate a CUDA slice from host data while recording data-plane H2D
1945    /// transfer stats.
1946    pub fn htod_sync_copy_tracked<T: DeviceRepr>(
1947        &self,
1948        src: &[T],
1949    ) -> Result<cudarc::driver::CudaSlice<T>> {
1950        let bytes = std::mem::size_of::<T>()
1951            .checked_mul(src.len())
1952            .ok_or_else(|| XlogError::Kernel("htod size overflow".to_string()))?;
1953        self.transfer_tracker.record_htod(bytes as u64);
1954        self.device
1955            .inner()
1956            .htod_sync_copy(src)
1957            .map_err(|e| XlogError::Kernel(format!("Failed to copy to device: {}", e)))
1958    }
1959
1960    /// Upload bounded launch metadata from host to device while recording it in
1961    /// the launch-metadata subcounter.
1962    pub fn htod_launch_metadata_sync_copy_into<
1963        T: DeviceRepr,
1964        Dst: cudarc::driver::DevicePtrMut<T>,
1965    >(
1966        &self,
1967        src: &[T],
1968        dst: &mut Dst,
1969    ) -> Result<()> {
1970        let bytes = std::mem::size_of::<T>()
1971            .checked_mul(src.len())
1972            .ok_or_else(|| XlogError::Kernel("launch metadata htod size overflow".to_string()))?;
1973        self.transfer_tracker
1974            .record_htod_launch_metadata(bytes as u64);
1975        self.device
1976            .inner()
1977            .htod_sync_copy_into(src, dst)
1978            .map_err(|e| {
1979                XlogError::Kernel(format!("Failed to copy launch metadata to device: {}", e))
1980            })
1981    }
1982
1983    /// Upload one launch-metadata scalar to device on a caller-owned stream
1984    /// while recording the transfer in the launch-metadata H2D counters.
1985    pub(crate) fn htod_launch_metadata_async_copy_one<T: DeviceRepr>(
1986        &self,
1987        src: &T,
1988        dst: &TrackedCudaSlice<T>,
1989        stream: &CudaStream,
1990        context: &str,
1991    ) -> Result<()> {
1992        let bytes = std::mem::size_of::<T>();
1993        self.transfer_tracker
1994            .record_htod_launch_metadata(bytes as u64);
1995        unsafe {
1996            let res = cudarc::driver::sys::cuMemcpyHtoDAsync_v2(
1997                *dst.device_ptr(),
1998                src as *const T as *const c_void,
1999                bytes,
2000                stream.cu_stream(),
2001            );
2002            if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
2003                return Err(XlogError::Kernel(format!(
2004                    "{context}: launch metadata H2D failed: {res:?}"
2005                )));
2006            }
2007        }
2008        Ok(())
2009    }
2010
2011    /// Compute exclusive prefix sum of u8 mask, returns (prefix_sum_vec, total_count)
2012    ///
2013    /// This is useful for compaction operations where we need to know:
2014    /// 1. The output position for each input element (prefix sum)
2015    /// 2. The total number of elements that pass the mask (count)
2016    ///
2017    /// # Arguments
2018    /// * `mask` - A slice of u8 values (0 or non-zero)
2019    ///
2020    /// # Returns
2021    /// A tuple of:
2022    /// - `Vec<u32>` containing the exclusive prefix sum
2023    /// - `u32` containing the total count of non-zero mask elements
2024    ///
2025    /// # Example
2026    /// ```ignore
2027    /// let mask = vec![1u8, 0, 1, 1, 0, 1];
2028    /// let (prefix_sum, count) = provider.prefix_sum_mask(&mask)?;
2029    /// // prefix_sum = [0, 1, 1, 2, 3, 3]
2030    /// // count = 4
2031    /// ```
2032    ///
2033    /// # Note
2034    /// For small inputs (<=256 elements), a CPU scan is used for efficiency.
2035    /// For larger inputs, a three-phase multi-block GPU scan is used.
2036    ///
2037    /// # Errors
2038    /// Returns `XlogError::Kernel` if kernel execution fails
2039    pub fn exclusive_scan_u32_inplace(
2040        &self,
2041        data: &mut crate::memory::TrackedCudaSlice<u32>,
2042        n: u32,
2043    ) -> Result<()> {
2044        if n as usize > data.len() {
2045            return Err(XlogError::Kernel(format!(
2046                "exclusive_scan_u32_inplace: n={} exceeds slice len={}",
2047                n,
2048                data.len()
2049            )));
2050        }
2051        self.multiblock_scan_u32_inplace(data, n)
2052    }
2053
2054    fn multiblock_scan_u32_inplace(
2055        &self,
2056        data: &mut crate::memory::TrackedCudaSlice<u32>,
2057        n: u32,
2058    ) -> Result<()> {
2059        if n == 0 {
2060            return Ok(());
2061        }
2062
2063        let device = self.device.inner();
2064        let block_size = 256u32;
2065
2066        if n <= block_size {
2067            let phase2_fn = device
2068                .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE2)
2069                .ok_or_else(|| {
2070                    XlogError::Kernel("Failed to get multiblock_scan_phase2 kernel".to_string())
2071                })?;
2072
2073            // SAFETY: multiblock_scan_phase2(uint32_t* block_sums, uint32_t num_blocks)
2074            unsafe {
2075                phase2_fn.clone().launch(
2076                    LaunchConfig {
2077                        grid_dim: (1, 1, 1),
2078                        block_dim: (block_size, 1, 1),
2079                        shared_mem_bytes: 0,
2080                    },
2081                    (&mut *data, n),
2082                )
2083            }
2084            .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase2 failed: {}", e)))?;
2085
2086            self.device.synchronize()?;
2087            return Ok(());
2088        }
2089
2090        let num_blocks = n.div_ceil(block_size);
2091        let mut block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
2092
2093        let phase1_u32_fn = device
2094            .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_U32_PHASE1)
2095            .ok_or_else(|| {
2096                XlogError::Kernel("Failed to get multiblock_scan_u32_phase1 kernel".to_string())
2097            })?;
2098
2099        // SAFETY: multiblock_scan_u32_phase1(uint32_t* data, uint32_t* block_sums, uint32_t n)
2100        unsafe {
2101            phase1_u32_fn.clone().launch(
2102                LaunchConfig {
2103                    grid_dim: (num_blocks, 1, 1),
2104                    block_dim: (block_size, 1, 1),
2105                    shared_mem_bytes: 0,
2106                },
2107                (&mut *data, &mut block_sums, n),
2108            )
2109        }
2110        .map_err(|e| XlogError::Kernel(format!("multiblock_scan_u32_phase1 failed: {}", e)))?;
2111        self.device.synchronize()?;
2112
2113        if num_blocks > 1 {
2114            self.multiblock_scan_u32_inplace(&mut block_sums, num_blocks)?;
2115        }
2116
2117        let phase3_fn = device
2118            .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
2119            .ok_or_else(|| {
2120                XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
2121            })?;
2122
2123        // SAFETY: multiblock_scan_phase3(uint32_t* prefix_sum, const uint32_t* block_offsets, uint32_t n)
2124        unsafe {
2125            phase3_fn.clone().launch(
2126                LaunchConfig {
2127                    grid_dim: (num_blocks, 1, 1),
2128                    block_dim: (block_size, 1, 1),
2129                    shared_mem_bytes: 0,
2130                },
2131                (&mut *data, &block_sums, n),
2132            )
2133        }
2134        .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase3 failed: {}", e)))?;
2135
2136        self.device.synchronize()?;
2137        Ok(())
2138    }
2139
2140    /// Stream-aware variant of [`Self::multiblock_scan_u32_inplace`].
2141    ///
2142    /// Runs every kernel of the recursive scan on `cu_stream`
2143    /// (no `device.synchronize()`), and records each intermediate
2144    /// `block_sums` allocation against the runtime so that when
2145    /// the helper returns and the local drops, the runtime's
2146    /// deallocate can queue `cuStreamWaitEvent(alloc_stream,
2147    /// recorded_event)` BEFORE `cuMemFreeAsync` — the same
2148    /// cross-stream lifetime safety the LaunchRecorder gives
2149    /// caller-provided buffers.
2150    ///
2151    /// `data` is not recorded here: the caller already records
2152    /// its own write of `data` against the same launch_stream
2153    /// (typically via `LaunchRecorder::write` BEFORE preflight).
2154    pub(crate) fn multiblock_scan_u32_inplace_on_stream(
2155        &self,
2156        data: &mut crate::memory::TrackedCudaSlice<u32>,
2157        n: u32,
2158        cu_stream: &cudarc::driver::CudaStream,
2159        launch_stream: crate::device_runtime::StreamId,
2160        runtime: &crate::device_runtime::XlogDeviceRuntime,
2161    ) -> Result<()> {
2162        if n == 0 {
2163            return Ok(());
2164        }
2165        let device = self.device.inner();
2166        let block_size = 256u32;
2167
2168        if n <= block_size {
2169            let phase2_fn = device
2170                .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE2)
2171                .ok_or_else(|| {
2172                    XlogError::Kernel("Failed to get multiblock_scan_phase2 kernel".to_string())
2173                })?;
2174            // SAFETY: kernel signature matches; data is mutated in place.
2175            unsafe {
2176                phase2_fn.clone().launch_on_stream(
2177                    cu_stream,
2178                    LaunchConfig {
2179                        grid_dim: (1, 1, 1),
2180                        block_dim: (block_size, 1, 1),
2181                        shared_mem_bytes: 0,
2182                    },
2183                    (&mut *data, n),
2184                )
2185            }
2186            .map_err(|e| {
2187                XlogError::Kernel(format!("multiblock_scan_phase2 (on_stream) failed: {}", e))
2188            })?;
2189            return Ok(());
2190        }
2191
2192        let num_blocks = n.div_ceil(block_size);
2193        let mut block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
2194        // Fence alloc-ready → launch_stream for block_sums
2195        // before phase1 kernel writes it. The alloc was queued
2196        // on the manager's default stream; without this wait,
2197        // a launch_stream-queued kernel can begin before
2198        // cuMemAllocAsync completes and read pool-recycled
2199        // bytes when the streams differ.
2200        runtime
2201            .prepare_first_use(
2202                &block_sums,
2203                launch_stream,
2204                crate::device_runtime::Access::Write,
2205            )
2206            .map_err(|e| {
2207                XlogError::Kernel(format!(
2208                    "multiblock_scan_u32_inplace_on_stream: prepare block_sums failed: {}",
2209                    e
2210                ))
2211            })?;
2212
2213        let phase1_u32_fn = device
2214            .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_U32_PHASE1)
2215            .ok_or_else(|| {
2216                XlogError::Kernel("Failed to get multiblock_scan_u32_phase1 kernel".to_string())
2217            })?;
2218        // SAFETY: kernel signature matches.
2219        unsafe {
2220            phase1_u32_fn.clone().launch_on_stream(
2221                cu_stream,
2222                LaunchConfig {
2223                    grid_dim: (num_blocks, 1, 1),
2224                    block_dim: (block_size, 1, 1),
2225                    shared_mem_bytes: 0,
2226                },
2227                (&mut *data, &mut block_sums, n),
2228            )
2229        }
2230        .map_err(|e| {
2231            XlogError::Kernel(format!(
2232                "multiblock_scan_u32_phase1 (on_stream) failed: {}",
2233                e
2234            ))
2235        })?;
2236
2237        if num_blocks > 1 {
2238            self.multiblock_scan_u32_inplace_on_stream(
2239                &mut block_sums,
2240                num_blocks,
2241                cu_stream,
2242                launch_stream,
2243                runtime,
2244            )?;
2245        }
2246
2247        let phase3_fn = device
2248            .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
2249            .ok_or_else(|| {
2250                XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
2251            })?;
2252        // SAFETY: kernel signature matches.
2253        unsafe {
2254            phase3_fn.clone().launch_on_stream(
2255                cu_stream,
2256                LaunchConfig {
2257                    grid_dim: (num_blocks, 1, 1),
2258                    block_dim: (block_size, 1, 1),
2259                    shared_mem_bytes: 0,
2260                },
2261                (&mut *data, &block_sums, n),
2262            )
2263        }
2264        .map_err(|e| {
2265            XlogError::Kernel(format!("multiblock_scan_phase3 (on_stream) failed: {}", e))
2266        })?;
2267
2268        // Record `block_sums` use on `launch_stream` BEFORE it
2269        // drops at end-of-scope. Without this, the runtime's
2270        // deallocate would queue `cuMemFreeAsync` on alloc_stream
2271        // without waiting for the launch_stream chain that's
2272        // still reading/writing block_sums to complete.
2273        if let Some(b) = block_sums.runtime_block() {
2274            runtime
2275                .finish_block_use(
2276                    crate::device_runtime::BlockId::from_block(b),
2277                    launch_stream,
2278                    crate::device_runtime::Access::Write,
2279                )
2280                .map_err(|e| {
2281                    XlogError::Kernel(format!(
2282                        "multiblock_scan_u32_inplace_on_stream: finish_block_use \
2283                         for intermediate block_sums failed: {}",
2284                        e
2285                    ))
2286                })?;
2287        } else {
2288            return Err(XlogError::Kernel(
2289                "multiblock_scan_u32_inplace_on_stream: intermediate block_sums has no \
2290                 runtime block — caller must use a runtime-backed manager"
2291                    .to_string(),
2292            ));
2293        }
2294        Ok(())
2295    }
2296
2297    /// Allocate every recursive `block_sums` buffer needed by
2298    /// [`Self::multiblock_scan_u32_inplace_on_stream_with_scratch`].
2299    pub(crate) fn multiblock_scan_u32_scratch_for_len(
2300        &self,
2301        mut n: u32,
2302    ) -> Result<MultiblockScanScratchU32> {
2303        let block_size = 256u32;
2304        let mut levels = Vec::new();
2305        while n > block_size {
2306            let num_blocks = n.div_ceil(block_size);
2307            levels.push(self.memory.alloc::<u32>(num_blocks as usize)?);
2308            n = num_blocks;
2309        }
2310        Ok(MultiblockScanScratchU32 { levels })
2311    }
2312
2313    /// Stream-aware u32 scan with caller-owned scratch.
2314    ///
2315    /// This is the CUDA Graph compatible counterpart to
2316    /// [`Self::multiblock_scan_u32_inplace_on_stream`]: all scratch buffers are
2317    /// supplied by the caller, so graph capture sees a stable scan topology and
2318    /// stable intermediate addresses.
2319    pub(crate) fn multiblock_scan_u32_inplace_on_stream_with_scratch(
2320        &self,
2321        data: &mut crate::memory::TrackedCudaSlice<u32>,
2322        n: u32,
2323        cu_stream: &cudarc::driver::CudaStream,
2324        scratch: &mut MultiblockScanScratchU32,
2325    ) -> Result<()> {
2326        self.multiblock_scan_u32_inplace_on_stream_with_scratch_levels(
2327            data,
2328            n,
2329            cu_stream,
2330            &mut scratch.levels,
2331        )
2332    }
2333
2334    fn multiblock_scan_u32_inplace_on_stream_with_scratch_levels(
2335        &self,
2336        data: &mut crate::memory::TrackedCudaSlice<u32>,
2337        n: u32,
2338        cu_stream: &cudarc::driver::CudaStream,
2339        scratch_levels: &mut [TrackedCudaSlice<u32>],
2340    ) -> Result<()> {
2341        if n == 0 {
2342            return Ok(());
2343        }
2344        let device = self.device.inner();
2345        let block_size = 256u32;
2346
2347        if n <= block_size {
2348            let phase2_fn = device
2349                .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE2)
2350                .ok_or_else(|| {
2351                    XlogError::Kernel("Failed to get multiblock_scan_phase2 kernel".to_string())
2352                })?;
2353            // SAFETY: kernel signature matches; data is mutated in place.
2354            unsafe {
2355                phase2_fn.clone().launch_on_stream(
2356                    cu_stream,
2357                    LaunchConfig {
2358                        grid_dim: (1, 1, 1),
2359                        block_dim: (block_size, 1, 1),
2360                        shared_mem_bytes: 0,
2361                    },
2362                    (&mut *data, n),
2363                )
2364            }
2365            .map_err(|e| {
2366                XlogError::Kernel(format!(
2367                    "multiblock_scan_phase2 (graph scratch) failed: {}",
2368                    e
2369                ))
2370            })?;
2371            return Ok(());
2372        }
2373
2374        let num_blocks = n.div_ceil(block_size);
2375        let (block_sums, rest) = scratch_levels.split_first_mut().ok_or_else(|| {
2376            XlogError::Kernel(format!(
2377                "multiblock_scan_u32_inplace_on_stream_with_scratch: missing scratch level \
2378                 for n={n}, num_blocks={num_blocks}"
2379            ))
2380        })?;
2381        if block_sums.len() < num_blocks as usize {
2382            return Err(XlogError::Kernel(format!(
2383                "multiblock_scan_u32_inplace_on_stream_with_scratch: scratch level too small \
2384                 (have {}, need {})",
2385                block_sums.len(),
2386                num_blocks
2387            )));
2388        }
2389
2390        let phase1_u32_fn = device
2391            .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_U32_PHASE1)
2392            .ok_or_else(|| {
2393                XlogError::Kernel("Failed to get multiblock_scan_u32_phase1 kernel".to_string())
2394            })?;
2395        // SAFETY: kernel signature matches.
2396        unsafe {
2397            phase1_u32_fn.clone().launch_on_stream(
2398                cu_stream,
2399                LaunchConfig {
2400                    grid_dim: (num_blocks, 1, 1),
2401                    block_dim: (block_size, 1, 1),
2402                    shared_mem_bytes: 0,
2403                },
2404                (&mut *data, &mut *block_sums, n),
2405            )
2406        }
2407        .map_err(|e| {
2408            XlogError::Kernel(format!(
2409                "multiblock_scan_u32_phase1 (graph scratch) failed: {}",
2410                e
2411            ))
2412        })?;
2413
2414        if num_blocks > 1 {
2415            self.multiblock_scan_u32_inplace_on_stream_with_scratch_levels(
2416                block_sums, num_blocks, cu_stream, rest,
2417            )?;
2418        }
2419
2420        let phase3_fn = device
2421            .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
2422            .ok_or_else(|| {
2423                XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
2424            })?;
2425        // SAFETY: kernel signature matches.
2426        unsafe {
2427            phase3_fn.clone().launch_on_stream(
2428                cu_stream,
2429                LaunchConfig {
2430                    grid_dim: (num_blocks, 1, 1),
2431                    block_dim: (block_size, 1, 1),
2432                    shared_mem_bytes: 0,
2433                },
2434                (&mut *data, &*block_sums, n),
2435            )
2436        }
2437        .map_err(|e| {
2438            XlogError::Kernel(format!(
2439                "multiblock_scan_phase3 (graph scratch) failed: {}",
2440                e
2441            ))
2442        })?;
2443        Ok(())
2444    }
2445
2446    /// Stream-aware view-inplace variant of
2447    /// [`Self::multiblock_scan_u32_view_inplace`]. Same shape
2448    /// as [`Self::multiblock_scan_u32_inplace_on_stream`] but
2449    /// over a `CudaViewMut` (used by recorded radix sort
2450    /// digit loops that scan per-digit slices of the histogram
2451    /// in place). Records intermediate `block_sums` against
2452    /// the runtime before they drop at end-of-scope.
2453    pub(crate) fn multiblock_scan_u32_view_inplace_on_stream(
2454        &self,
2455        data: &mut CudaViewMut<'_, u32>,
2456        n: u32,
2457        cu_stream: &cudarc::driver::CudaStream,
2458        launch_stream: crate::device_runtime::StreamId,
2459        runtime: &crate::device_runtime::XlogDeviceRuntime,
2460    ) -> Result<()> {
2461        if n == 0 {
2462            return Ok(());
2463        }
2464        let device = self.device.inner();
2465        let block_size = 256u32;
2466
2467        if n <= block_size {
2468            let phase2_fn = device
2469                .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE2)
2470                .ok_or_else(|| {
2471                    XlogError::Kernel("Failed to get multiblock_scan_phase2 kernel".to_string())
2472                })?;
2473            // SAFETY: phase2 kernel signature.
2474            unsafe {
2475                phase2_fn.clone().launch_on_stream(
2476                    cu_stream,
2477                    LaunchConfig {
2478                        grid_dim: (1, 1, 1),
2479                        block_dim: (block_size, 1, 1),
2480                        shared_mem_bytes: 0,
2481                    },
2482                    (data, n),
2483                )
2484            }
2485            .map_err(|e| {
2486                XlogError::Kernel(format!(
2487                    "multiblock_scan_phase2 (view on_stream) failed: {}",
2488                    e
2489                ))
2490            })?;
2491            return Ok(());
2492        }
2493
2494        let num_blocks = n.div_ceil(block_size);
2495        let mut block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
2496        // Fence alloc-ready → launch_stream for block_sums
2497        // before phase1 kernel writes it. See the inplace
2498        // variant for the full rationale.
2499        runtime
2500            .prepare_first_use(
2501                &block_sums,
2502                launch_stream,
2503                crate::device_runtime::Access::Write,
2504            )
2505            .map_err(|e| {
2506                XlogError::Kernel(format!(
2507                    "multiblock_scan_u32_view_inplace_on_stream: prepare block_sums failed: {}",
2508                    e
2509                ))
2510            })?;
2511
2512        let phase1_u32_fn = device
2513            .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_U32_PHASE1)
2514            .ok_or_else(|| {
2515                XlogError::Kernel("Failed to get multiblock_scan_u32_phase1 kernel".to_string())
2516            })?;
2517        // SAFETY: phase1 kernel signature.
2518        unsafe {
2519            phase1_u32_fn.clone().launch_on_stream(
2520                cu_stream,
2521                LaunchConfig {
2522                    grid_dim: (num_blocks, 1, 1),
2523                    block_dim: (block_size, 1, 1),
2524                    shared_mem_bytes: 0,
2525                },
2526                (&mut *data, &mut block_sums, n),
2527            )
2528        }
2529        .map_err(|e| {
2530            XlogError::Kernel(format!(
2531                "multiblock_scan_u32_phase1 (view on_stream) failed: {}",
2532                e
2533            ))
2534        })?;
2535
2536        if num_blocks > 1 {
2537            self.multiblock_scan_u32_inplace_on_stream(
2538                &mut block_sums,
2539                num_blocks,
2540                cu_stream,
2541                launch_stream,
2542                runtime,
2543            )?;
2544        }
2545
2546        let phase3_fn = device
2547            .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
2548            .ok_or_else(|| {
2549                XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
2550            })?;
2551        // SAFETY: phase3 kernel signature.
2552        unsafe {
2553            phase3_fn.clone().launch_on_stream(
2554                cu_stream,
2555                LaunchConfig {
2556                    grid_dim: (num_blocks, 1, 1),
2557                    block_dim: (block_size, 1, 1),
2558                    shared_mem_bytes: 0,
2559                },
2560                (&mut *data, &block_sums, n),
2561            )
2562        }
2563        .map_err(|e| {
2564            XlogError::Kernel(format!(
2565                "multiblock_scan_phase3 (view on_stream) failed: {}",
2566                e
2567            ))
2568        })?;
2569
2570        // Record block_sums use before end-of-scope drop.
2571        if let Some(b) = block_sums.runtime_block() {
2572            runtime
2573                .finish_block_use(
2574                    crate::device_runtime::BlockId::from_block(b),
2575                    launch_stream,
2576                    crate::device_runtime::Access::Write,
2577                )
2578                .map_err(|e| {
2579                    XlogError::Kernel(format!(
2580                        "multiblock_scan_u32_view_inplace_on_stream: finish_block_use \
2581                     for intermediate block_sums failed: {}",
2582                        e
2583                    ))
2584                })?;
2585        } else {
2586            return Err(XlogError::Kernel(
2587                "multiblock_scan_u32_view_inplace_on_stream: intermediate block_sums has no \
2588                 runtime block — caller must use a runtime-backed manager"
2589                    .to_string(),
2590            ));
2591        }
2592        Ok(())
2593    }
2594
2595    fn multiblock_scan_u32_view_inplace(
2596        &self,
2597        data: &mut CudaViewMut<'_, u32>,
2598        n: u32,
2599    ) -> Result<()> {
2600        if n == 0 {
2601            return Ok(());
2602        }
2603
2604        let device = self.device.inner();
2605        let block_size = 256u32;
2606
2607        if n <= block_size {
2608            let phase2_fn = device
2609                .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE2)
2610                .ok_or_else(|| {
2611                    XlogError::Kernel("Failed to get multiblock_scan_phase2 kernel".to_string())
2612                })?;
2613
2614            // SAFETY: multiblock_scan_phase2(uint32_t* block_sums, uint32_t num_blocks)
2615            unsafe {
2616                phase2_fn.clone().launch(
2617                    LaunchConfig {
2618                        grid_dim: (1, 1, 1),
2619                        block_dim: (block_size, 1, 1),
2620                        shared_mem_bytes: 0,
2621                    },
2622                    (data, n),
2623                )
2624            }
2625            .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase2 failed: {}", e)))?;
2626
2627            self.device.synchronize()?;
2628            return Ok(());
2629        }
2630
2631        let num_blocks = n.div_ceil(block_size);
2632        let mut block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
2633
2634        let phase1_u32_fn = device
2635            .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_U32_PHASE1)
2636            .ok_or_else(|| {
2637                XlogError::Kernel("Failed to get multiblock_scan_u32_phase1 kernel".to_string())
2638            })?;
2639
2640        // SAFETY: multiblock_scan_u32_phase1(uint32_t* data, uint32_t* block_sums, uint32_t n)
2641        unsafe {
2642            phase1_u32_fn.clone().launch(
2643                LaunchConfig {
2644                    grid_dim: (num_blocks, 1, 1),
2645                    block_dim: (block_size, 1, 1),
2646                    shared_mem_bytes: 0,
2647                },
2648                (&mut *data, &mut block_sums, n),
2649            )
2650        }
2651        .map_err(|e| XlogError::Kernel(format!("multiblock_scan_u32_phase1 failed: {}", e)))?;
2652        self.device.synchronize()?;
2653
2654        if num_blocks > 1 {
2655            self.multiblock_scan_u32_inplace(&mut block_sums, num_blocks)?;
2656        }
2657
2658        let phase3_fn = device
2659            .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
2660            .ok_or_else(|| {
2661                XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
2662            })?;
2663
2664        // SAFETY: multiblock_scan_phase3(uint32_t* prefix_sum, const uint32_t* block_offsets, uint32_t n)
2665        unsafe {
2666            phase3_fn.clone().launch(
2667                LaunchConfig {
2668                    grid_dim: (num_blocks, 1, 1),
2669                    block_dim: (block_size, 1, 1),
2670                    shared_mem_bytes: 0,
2671                },
2672                (&mut *data, &block_sums, n),
2673            )
2674        }
2675        .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase3 failed: {}", e)))?;
2676
2677        self.device.synchronize()?;
2678        Ok(())
2679    }
2680
2681    // ============== Internal Helper Methods ==============
2682
2683    /// Read a buffer's logical row count, using the host cache when available
2684    /// and falling back to a metadata-only device-to-host read when needed.
2685    pub fn device_row_count(&self, buffer: &CudaBuffer) -> Result<usize> {
2686        if let Some(n) = buffer.cached_row_count() {
2687            return Ok(n as usize);
2688        }
2689        let mut host_rows = [0u32];
2690        self.device
2691            .inner()
2692            .dtoh_sync_copy_into(buffer.num_rows_device(), &mut host_rows)
2693            .map_err(|e| XlogError::Kernel(format!("Failed to read row count: {}", e)))?;
2694        buffer.set_cached_row_count_if_unset(host_rows[0]);
2695        Ok(host_rows[0] as usize)
2696    }
2697
2698    /// Read and validate a buffer's logical row count for outward-facing APIs.
2699    ///
2700    /// This keeps exported/query-visible lengths tied to the device logical row
2701    /// count while still rejecting impossible metadata (`logical_rows > row_cap`).
2702    pub fn validated_logical_row_count(&self, buffer: &CudaBuffer) -> Result<usize> {
2703        let logical_rows = self.device_row_count(buffer)?;
2704        validate_logical_row_count(buffer.num_rows(), logical_rows)
2705    }
2706
2707    fn clone_device_row_count(&self, buffer: &CudaBuffer) -> Result<TrackedCudaSlice<u32>> {
2708        let mut d_num_rows = self.memory.alloc::<u32>(1)?;
2709        self.device
2710            .inner()
2711            .dtod_copy(buffer.num_rows_device(), &mut d_num_rows)
2712            .map_err(|e| XlogError::Kernel(format!("Failed to copy row count: {}", e)))?;
2713        Ok(d_num_rows)
2714    }
2715
2716    fn upload_device_row_count(&self, row_count: u32) -> Result<TrackedCudaSlice<u32>> {
2717        let mut d_num_rows = self.memory.alloc::<u32>(1)?;
2718        self.htod_launch_metadata_sync_copy_into(&[row_count], &mut d_num_rows)
2719            .map_err(|e| XlogError::Kernel(format!("Failed to upload row count: {}", e)))?;
2720        Ok(d_num_rows)
2721    }
2722
2723    fn buffer_from_columns_with_device_count(
2724        &self,
2725        columns: Vec<CudaColumn>,
2726        row_cap: u64,
2727        schema: Schema,
2728        src: &CudaBuffer,
2729    ) -> Result<CudaBuffer> {
2730        let d_num_rows = self.clone_device_row_count(src)?;
2731        Ok(CudaBuffer::from_columns(
2732            columns, row_cap, d_num_rows, schema,
2733        ))
2734    }
2735
2736    fn column_bytes_view<'a>(
2737        &self,
2738        col: &'a CudaColumn,
2739        num_bytes: usize,
2740    ) -> Result<RawCudaView<'a, u8>> {
2741        if col.num_bytes() < num_bytes {
2742            return Err(XlogError::Kernel(format!(
2743                "Column has {} bytes but {} required",
2744                col.num_bytes(),
2745                num_bytes
2746            )));
2747        }
2748        let ptr = *col.device_ptr();
2749        Ok(RawCudaView {
2750            ptr,
2751            len: num_bytes,
2752            stream: col.stream().clone(),
2753            source_block: col.runtime_block(),
2754            _marker: PhantomData,
2755        })
2756    }
2757
2758    fn bytes_as_u32_view<'a>(
2759        &self,
2760        bytes: &'a TrackedCudaSlice<u8>,
2761        num_elements: usize,
2762    ) -> Result<RawCudaView<'a, u32>> {
2763        let required_bytes = num_elements * std::mem::size_of::<u32>();
2764        if bytes.len() < required_bytes {
2765            return Err(XlogError::Kernel(format!(
2766                "Packed keys have {} bytes but {} required for {} u32 elements",
2767                bytes.len(),
2768                required_bytes,
2769                num_elements
2770            )));
2771        }
2772        let ptr = *bytes.device_ptr();
2773        if !(ptr as usize).is_multiple_of(std::mem::align_of::<u32>()) {
2774            return Err(XlogError::Kernel(
2775                "Packed keys device pointer is not u32-aligned".to_string(),
2776            ));
2777        }
2778        Ok(RawCudaView {
2779            ptr,
2780            len: num_elements,
2781            stream: bytes.stream().clone(),
2782            source_block: bytes.runtime_block(),
2783            _marker: PhantomData,
2784        })
2785    }
2786
2787    /// Reinterpret a `CudaBuffer` column as a `u32` slice for kernel access.
2788    fn column_as_u32_view<'a>(
2789        &self,
2790        col: &'a CudaColumn,
2791        num_elements: usize,
2792    ) -> Result<RawCudaView<'a, u32>> {
2793        let required_bytes = num_elements * std::mem::size_of::<u32>();
2794        if col.num_bytes() < required_bytes {
2795            return Err(XlogError::Kernel(format!(
2796                "Column has {} bytes but {} required for {} u32 elements",
2797                col.num_bytes(),
2798                required_bytes,
2799                num_elements
2800            )));
2801        }
2802        let ptr = *col.device_ptr();
2803        if !(ptr as usize).is_multiple_of(std::mem::align_of::<u32>()) {
2804            return Err(XlogError::Kernel(
2805                "Column device pointer is not u32-aligned".to_string(),
2806            ));
2807        }
2808        Ok(RawCudaView {
2809            ptr,
2810            len: num_elements,
2811            stream: col.stream().clone(),
2812            source_block: col.runtime_block(),
2813            _marker: PhantomData,
2814        })
2815    }
2816
2817    fn column_as_u64_view<'a>(
2818        &self,
2819        col: &'a CudaColumn,
2820        num_elements: usize,
2821    ) -> Result<RawCudaView<'a, u64>> {
2822        let required_bytes = num_elements * std::mem::size_of::<u64>();
2823        if col.num_bytes() < required_bytes {
2824            return Err(XlogError::Kernel(format!(
2825                "Column has {} bytes but {} required for {} u64 elements",
2826                col.num_bytes(),
2827                required_bytes,
2828                num_elements
2829            )));
2830        }
2831        let ptr = *col.device_ptr();
2832        if !(ptr as usize).is_multiple_of(std::mem::align_of::<u64>()) {
2833            return Err(XlogError::Kernel(
2834                "Column device pointer is not u64-aligned".to_string(),
2835            ));
2836        }
2837        Ok(RawCudaView {
2838            ptr,
2839            len: num_elements,
2840            stream: col.stream().clone(),
2841            source_block: col.runtime_block(),
2842            _marker: PhantomData,
2843        })
2844    }
2845
2846    /// Reinterpret a `CudaBuffer` column as an `f64` slice for kernel access.
2847    fn column_as_f64_view<'a>(
2848        &self,
2849        col: &'a CudaColumn,
2850        num_elements: usize,
2851    ) -> Result<RawCudaView<'a, f64>> {
2852        let required_bytes = num_elements * std::mem::size_of::<f64>();
2853        if col.num_bytes() < required_bytes {
2854            return Err(XlogError::Kernel(format!(
2855                "Column has {} bytes but {} required for {} f64 elements",
2856                col.num_bytes(),
2857                required_bytes,
2858                num_elements
2859            )));
2860        }
2861        let ptr = *col.device_ptr();
2862        if !(ptr as usize).is_multiple_of(std::mem::align_of::<f64>()) {
2863            return Err(XlogError::Kernel(
2864                "Column device pointer is not f64-aligned".to_string(),
2865            ));
2866        }
2867        Ok(RawCudaView {
2868            ptr,
2869            len: num_elements,
2870            stream: col.stream().clone(),
2871            source_block: col.runtime_block(),
2872            _marker: PhantomData,
2873        })
2874    }
2875
2876    /// Create an empty buffer with the given schema (all columns are empty slices)
2877    ///
2878    /// # Arguments
2879    /// * `schema` - The schema for the empty buffer
2880    ///
2881    /// # Returns
2882    /// A new CudaBuffer with zero rows
2883    ///
2884    /// # Errors
2885    /// Returns `XlogError::Kernel` if allocation fails
2886    pub fn create_empty_buffer(&self, schema: Schema) -> Result<CudaBuffer> {
2887        let mut columns = Vec::with_capacity(schema.arity());
2888        for _ in 0..schema.arity() {
2889            // Allocate zero-length column
2890            columns.push(self.memory.alloc::<u8>(0)?.into());
2891        }
2892        self.buffer_from_columns(columns, 0, schema)
2893    }
2894
2895    /// Create a zero-arity (nullary) relation buffer carrying `rows` unit tuples.
2896    ///
2897    /// A nullary relation holds exactly when it has at least one row; its single
2898    /// possible tuple is the empty tuple `()`. `create_buffer_from_slices` with no
2899    /// column slices routes to `create_empty_buffer` (0 rows), which represents the
2900    /// relation as *absent* — wrong for an asserted nullary fact. Nullary facts must
2901    /// use this path so presence is materialized as one row.
2902    pub fn create_zero_arity_buffer(&self, schema: Schema, rows: u32) -> Result<CudaBuffer> {
2903        debug_assert_eq!(
2904            schema.arity(),
2905            0,
2906            "create_zero_arity_buffer requires arity 0"
2907        );
2908        self.buffer_from_columns(Vec::new(), u64::from(rows), schema)
2909    }
2910
2911    pub(crate) fn buffer_from_columns(
2912        &self,
2913        columns: Vec<CudaColumn>,
2914        row_cap: u64,
2915        schema: Schema,
2916    ) -> Result<CudaBuffer> {
2917        let row_u32 = u32::try_from(row_cap)
2918            .map_err(|_| XlogError::Kernel(format!("Row capacity {} exceeds u32::MAX", row_cap)))?;
2919        let mut d_num_rows = self.memory.alloc::<u32>(1)?;
2920        self.htod_launch_metadata_sync_copy_into(&[row_u32], &mut d_num_rows)
2921            .map_err(|e| XlogError::Kernel(format!("Failed to set row count: {}", e)))?;
2922        Ok(CudaBuffer::from_columns_with_host_count(
2923            columns, row_cap, d_num_rows, schema, row_u32,
2924        ))
2925    }
2926
2927    /// Combine schemas from left and right buffers for join result
2928    fn combine_schemas(&self, left: &Schema, right: &Schema) -> Schema {
2929        let mut columns = left.columns.clone();
2930        columns.extend(right.columns.iter().cloned());
2931        let mut sort_labels = left.sort_labels().to_vec();
2932        sort_labels.extend(right.sort_labels().iter().cloned());
2933        Schema::new(columns)
2934            .with_sort_labels(sort_labels)
2935            .expect("combined schema sort labels match column arity")
2936    }
2937
2938    /// Check if two schemas have compatible types (same arity and column types)
2939    ///
2940    /// This ignores column names, which is useful for Datalog operations where
2941    /// projected relations may have different column names but the same types.
2942    fn schemas_type_compatible(&self, a: &Schema, b: &Schema) -> bool {
2943        if a.arity() != b.arity() {
2944            return false;
2945        }
2946        for i in 0..a.arity() {
2947            if a.column_type(i) != b.column_type(i) {
2948                return false;
2949            }
2950        }
2951        true
2952    }
2953}
2954
2955#[cfg(test)]
2956mod tests {
2957    use super::*;
2958    use crate::device_runtime::{
2959        AsyncCudaResource, DeviceMemoryResource, GlobalDeviceBudget, LoggingResource, NullSink,
2960        StreamPool, XlogDeviceRuntime,
2961    };
2962    use xlog_core::{AggOp, MemoryBudget, ScalarType};
2963
2964    fn has_cuda_device() -> bool {
2965        CudaDevice::new(0).is_ok()
2966    }
2967
2968    #[test]
2969    fn test_kernel_artifact_locator_precedence_order() {
2970        use super::kernel_paths::KernelArtifactLocator;
2971        use std::fs;
2972        use std::path::PathBuf;
2973
2974        let root = std::env::temp_dir().join(format!(
2975            "xlog-kernel-paths-{}-{}",
2976            std::process::id(),
2977            std::time::SystemTime::now()
2978                .duration_since(std::time::UNIX_EPOCH)
2979                .expect("system clock before UNIX_EPOCH")
2980                .as_nanos()
2981        ));
2982        let cubin_dir = root.join("cubin");
2983        let package_dir = root.join("bin").join("kernels");
2984        let out_dir = root.join("out");
2985        fs::create_dir_all(&cubin_dir).expect("create cubin dir");
2986        fs::create_dir_all(&package_dir).expect("create package kernels dir");
2987        fs::create_dir_all(&out_dir).expect("create out dir");
2988
2989        let name = "xlog_join";
2990        let cc = 75;
2991        let cubin_path = cubin_dir.join(format!("{name}.sm_{cc}.cubin"));
2992        let package_path = package_dir.join(format!("{name}.sm_{cc}.cubin"));
2993        let out_path = out_dir.join(format!("{name}.sm_{cc}.cubin"));
2994        fs::write(&cubin_path, b"cubin").expect("write cubin file");
2995        fs::write(&package_path, b"package").expect("write package file");
2996        fs::write(&out_path, b"out").expect("write out file");
2997
2998        let locator = KernelArtifactLocator::new(
2999            Some(cubin_dir.clone()),
3000            Some(package_dir.clone()),
3001            Some(out_dir.clone()),
3002        );
3003
3004        let (path, is_cubin) = locator
3005            .resolve_module_path(name, cc)
3006            .expect("expected a kernel artifact");
3007        assert_eq!(path, cubin_path);
3008        assert!(is_cubin);
3009
3010        fs::remove_file(&cubin_path).expect("remove cubin file");
3011        let (path, is_cubin) = locator
3012            .resolve_module_path(name, cc)
3013            .expect("expected package kernel artifact");
3014        assert_eq!(path, package_path);
3015        assert!(is_cubin);
3016
3017        fs::remove_file(&package_path).expect("remove package file");
3018        let (path, is_cubin) = locator
3019            .resolve_module_path(name, cc)
3020            .expect("expected out dir kernel artifact");
3021        assert_eq!(path, out_path);
3022        assert!(is_cubin);
3023
3024        let _ = fs::remove_dir_all(PathBuf::from(&root));
3025    }
3026
3027    #[test]
3028    fn test_module_resolution_finds_portable_ptx() {
3029        // Verify resolve_module_path finds portable PTX for all modules.
3030        // Uses a dummy cc (999) so cubin won't match — only portable PTX.
3031        for name in crate::kernel_manifest_data::KERNEL_CU_NAMES {
3032            let result = resolve_module_path(name, 999);
3033            assert!(
3034                result.is_some(),
3035                "resolve_module_path({name}, 999) should find portable PTX"
3036            );
3037            let (path, is_cubin) = result.unwrap();
3038            assert!(
3039                !is_cubin,
3040                "{name}: expected portable PTX fallback, got cubin"
3041            );
3042            assert!(
3043                path.to_str().unwrap().ends_with(".portable.ptx"),
3044                "{name}: path should end with .portable.ptx, got {:?}",
3045                path
3046            );
3047        }
3048    }
3049
3050    #[test]
3051    fn test_module_resolution_falls_back_to_embedded_portable_ptx() {
3052        use super::kernel_paths::KernelArtifactLocator;
3053
3054        let locator = KernelArtifactLocator::new(None, None, None);
3055        for name in crate::kernel_manifest_data::KERNEL_CU_NAMES {
3056            let sources = resolve_module_sources_with_locator(name, 999, &locator);
3057            assert_eq!(
3058                sources.len(),
3059                1,
3060                "{name}: expected only embedded portable PTX fallback"
3061            );
3062
3063            match &sources[0] {
3064                KernelModuleSource::EmbeddedPortablePtx { ptx } => {
3065                    assert!(
3066                        ptx.contains(".entry"),
3067                        "{name}: embedded PTX should contain CUDA entry points"
3068                    );
3069                }
3070                KernelModuleSource::File { path, .. } => {
3071                    panic!(
3072                        "{name}: expected embedded portable PTX fallback, got file {}",
3073                        path.display()
3074                    );
3075                }
3076            }
3077        }
3078    }
3079
3080    #[test]
3081    fn test_embedded_portable_ptx_manifest_matches_kernel_manifest() {
3082        let embedded_names: std::collections::BTreeSet<_> =
3083            crate::embedded_kernel_data::EMBEDDED_PORTABLE_PTX
3084                .iter()
3085                .map(|artifact| artifact.name)
3086                .collect();
3087        let manifest_names: std::collections::BTreeSet<_> =
3088            crate::kernel_manifest_data::KERNEL_CU_NAMES
3089                .iter()
3090                .copied()
3091                .collect();
3092
3093        assert_eq!(
3094            embedded_names, manifest_names,
3095            "embedded portable PTX table should cover every runtime kernel module"
3096        );
3097    }
3098
3099    #[test]
3100    fn test_kernel_provider_creation() {
3101        if !has_cuda_device() {
3102            eprintln!("Skipping test: no CUDA device available");
3103            return;
3104        }
3105
3106        let device = Arc::new(CudaDevice::new(0).expect("Failed to create device"));
3107        let budget = MemoryBudget::with_limit(1024 * 1024 * 1024); // 1 GB
3108        let memory = Arc::new(GpuMemoryManager::new(device.clone(), budget));
3109
3110        let provider = CudaKernelProvider::new(device.clone(), memory.clone());
3111        assert!(
3112            provider.is_ok(),
3113            "Failed to create kernel provider: {:?}",
3114            provider.err()
3115        );
3116
3117        let provider = provider.unwrap();
3118        assert!(Arc::ptr_eq(provider.device(), &device));
3119        assert!(Arc::ptr_eq(provider.memory(), &memory));
3120    }
3121
3122    #[test]
3123    fn test_kernel_functions_accessible() {
3124        if !has_cuda_device() {
3125            eprintln!("Skipping test: no CUDA device available");
3126            return;
3127        }
3128
3129        let device = Arc::new(CudaDevice::new(0).expect("Failed to create device"));
3130        let budget = MemoryBudget::with_limit(1024 * 1024 * 1024);
3131        let memory = Arc::new(GpuMemoryManager::new(device.clone(), budget));
3132
3133        let _provider =
3134            CudaKernelProvider::new(device.clone(), memory).expect("Failed to create provider");
3135
3136        // Verify all kernel functions can be retrieved
3137        let inner = device.inner();
3138
3139        // Join kernels
3140        let build_fn = inner.get_func(JOIN_MODULE, join_kernels::HASH_JOIN_BUILD);
3141        assert!(
3142            build_fn.is_some(),
3143            "hash_join_build function should be accessible"
3144        );
3145
3146        let probe_fn = inner.get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE);
3147        assert!(
3148            probe_fn.is_some(),
3149            "hash_join_probe function should be accessible"
3150        );
3151
3152        // Dedup kernels
3153        let mark_fn = inner.get_func(DEDUP_MODULE, dedup_kernels::MARK_DUPLICATES);
3154        assert!(
3155            mark_fn.is_some(),
3156            "mark_duplicates function should be accessible"
3157        );
3158
3159        let compact_fn = inner.get_func(DEDUP_MODULE, dedup_kernels::COMPACT_ROWS);
3160        assert!(
3161            compact_fn.is_some(),
3162            "compact_rows function should be accessible"
3163        );
3164
3165        // GroupBy kernels
3166        let boundaries_fn =
3167            inner.get_func(GROUPBY_MODULE, groupby_kernels::DETECT_GROUP_BOUNDARIES);
3168        assert!(
3169            boundaries_fn.is_some(),
3170            "detect_group_boundaries function should be accessible"
3171        );
3172
3173        let count_fn = inner.get_func(GROUPBY_MODULE, groupby_kernels::GROUPBY_COUNT);
3174        assert!(
3175            count_fn.is_some(),
3176            "groupby_count function should be accessible"
3177        );
3178
3179        let sum_fn = inner.get_func(GROUPBY_MODULE, groupby_kernels::GROUPBY_SUM);
3180        assert!(
3181            sum_fn.is_some(),
3182            "groupby_sum function should be accessible"
3183        );
3184
3185        let min_fn = inner.get_func(GROUPBY_MODULE, groupby_kernels::GROUPBY_MIN);
3186        assert!(
3187            min_fn.is_some(),
3188            "groupby_min function should be accessible"
3189        );
3190
3191        let max_fn = inner.get_func(GROUPBY_MODULE, groupby_kernels::GROUPBY_MAX);
3192        assert!(
3193            max_fn.is_some(),
3194            "groupby_max function should be accessible"
3195        );
3196
3197        // Circuit kernels (XGCF forward/backward)
3198        let xgcf_forward = inner.get_func(CIRCUIT_MODULE, "xgcf_forward_level");
3199        assert!(
3200            xgcf_forward.is_some(),
3201            "xgcf_forward_level function should be accessible"
3202        );
3203
3204        let xgcf_backward_propagate =
3205            inner.get_func(CIRCUIT_MODULE, "xgcf_backward_level_propagate");
3206        assert!(
3207            xgcf_backward_propagate.is_some(),
3208            "xgcf_backward_level_propagate function should be accessible"
3209        );
3210
3211        let xgcf_backward_decision_grad =
3212            inner.get_func(CIRCUIT_MODULE, "xgcf_backward_level_decision_grad");
3213        assert!(
3214            xgcf_backward_decision_grad.is_some(),
3215            "xgcf_backward_level_decision_grad function should be accessible"
3216        );
3217
3218        let xgcf_backward_lit_grad = inner.get_func(CIRCUIT_MODULE, "xgcf_backward_level_lit_grad");
3219        assert!(
3220            xgcf_backward_lit_grad.is_some(),
3221            "xgcf_backward_level_lit_grad function should be accessible"
3222        );
3223
3224        // Neural fast-path kernels (AD chain weight fill + gradient scatter)
3225        let neural_fill = inner.get_func("xlog_neural", "neural_fill_ad_chain_f32");
3226        assert!(
3227            neural_fill.is_some(),
3228            "neural_fill_ad_chain_f32 function should be accessible"
3229        );
3230        let neural_scatter = inner.get_func("xlog_neural", "neural_scatter_ad_chain_grads_f32");
3231        assert!(
3232            neural_scatter.is_some(),
3233            "neural_scatter_ad_chain_grads_f32 function should be accessible"
3234        );
3235    }
3236
3237    #[test]
3238    fn test_module_names_unique() {
3239        // Ensure module names don't collide
3240        assert_ne!(JOIN_MODULE, DEDUP_MODULE);
3241        assert_ne!(JOIN_MODULE, GROUPBY_MODULE);
3242        assert_ne!(DEDUP_MODULE, GROUPBY_MODULE);
3243    }
3244
3245    // Helper function to create test provider
3246    fn create_test_provider() -> Option<CudaKernelProvider> {
3247        if !has_cuda_device() {
3248            return None;
3249        }
3250        let device = Arc::new(CudaDevice::new(0).ok()?);
3251        let budget = MemoryBudget::with_limit(1024 * 1024 * 1024);
3252        let memory = Arc::new(GpuMemoryManager::new(device.clone(), budget));
3253        CudaKernelProvider::new(device, memory).ok()
3254    }
3255
3256    fn create_test_provider_with_runtime() -> Option<(CudaKernelProvider, Arc<XlogDeviceRuntime>)> {
3257        if !has_cuda_device() {
3258            return None;
3259        }
3260        let device = Arc::new(CudaDevice::new(0).ok()?);
3261        let pool = Arc::new(StreamPool::with_defaults(Arc::clone(&device)));
3262        let sink = Arc::new(NullSink::new());
3263        let async_resource: Box<dyn DeviceMemoryResource + Send + Sync> = Box::new(
3264            AsyncCudaResource::new(Arc::clone(&device), 0, Arc::clone(&pool)),
3265        );
3266        let logging: Box<dyn DeviceMemoryResource + Send + Sync> =
3267            Box::new(LoggingResource::new(async_resource, sink));
3268        let budget: Box<dyn DeviceMemoryResource + Send + Sync> =
3269            Box::new(GlobalDeviceBudget::new(logging, 1024 * 1024 * 1024));
3270        let runtime = Arc::new(XlogDeviceRuntime::with_resource(
3271            Arc::clone(&device),
3272            0,
3273            pool,
3274            budget,
3275        ));
3276        let memory = Arc::new(GpuMemoryManager::with_runtime(
3277            Arc::clone(&device),
3278            MemoryBudget::with_limit(1024 * 1024 * 1024),
3279            Arc::clone(&runtime),
3280        ));
3281        let provider = CudaKernelProvider::with_runtime(device, memory).ok()?;
3282        Some((provider, runtime))
3283    }
3284
3285    #[test]
3286    fn test_recorded_join_index_build_runs_on_runtime_stream() {
3287        let (provider, runtime) = match create_test_provider_with_runtime() {
3288            Some(fixture) => fixture,
3289            None => {
3290                eprintln!("Skipping test: no CUDA device available");
3291                return;
3292            }
3293        };
3294        let stream = runtime.stream_pool().acquire().expect("recorded stream");
3295        let left = create_test_buffer(&provider, &[1, 2, 3, 4], "key");
3296        let right = create_test_buffer(&provider, &[1, 2, 3, 4], "key");
3297
3298        let index = provider
3299            .build_join_index_v2_recorded(&right, &[0], stream)
3300            .expect("recorded join-index build");
3301        let joined = provider
3302            .hash_join_v2_with_index_recorded(
3303                &left,
3304                &right,
3305                &[0],
3306                &[0],
3307                JoinType::Inner,
3308                &index,
3309                None,
3310                stream,
3311            )
3312            .expect("recorded indexed join consumes recorded build");
3313        runtime
3314            .stream_pool()
3315            .resolve(stream)
3316            .expect("stream resolves")
3317            .synchronize()
3318            .expect("recorded stream synchronized");
3319
3320        assert_eq!(index.right_num_rows(), 4);
3321        assert_eq!(index.right_keys(), &[0]);
3322        assert_eq!(provider.device_row_count(&joined).expect("joined rows"), 4);
3323    }
3324
3325    // Helper function to create a CudaBuffer with U32 data
3326    fn create_test_buffer(
3327        provider: &CudaKernelProvider,
3328        data: &[u32],
3329        col_name: &str,
3330    ) -> CudaBuffer {
3331        let schema = Schema::new(vec![(col_name.to_string(), ScalarType::U32)]);
3332        let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
3333
3334        let mut col = provider.memory().alloc::<u8>(bytes.len()).expect("alloc");
3335        provider
3336            .device()
3337            .inner()
3338            .htod_sync_copy_into(&bytes, &mut col)
3339            .expect("htod");
3340
3341        provider
3342            .buffer_from_columns(vec![col.into()], data.len() as u64, schema)
3343            .expect("buffer")
3344    }
3345
3346    // Helper function to create an empty buffer with correct column count
3347    fn create_empty_test_buffer(provider: &CudaKernelProvider, schema: Schema) -> CudaBuffer {
3348        let mut columns = Vec::with_capacity(schema.arity());
3349        for _ in 0..schema.arity() {
3350            columns.push(provider.memory().alloc::<u8>(0).expect("alloc").into());
3351        }
3352        provider
3353            .buffer_from_columns(columns, 0, schema)
3354            .expect("buffer")
3355    }
3356
3357    // Helper function to read U32 data from CudaBuffer
3358    fn read_buffer_u32(provider: &CudaKernelProvider, buffer: &CudaBuffer, col: usize) -> Vec<u32> {
3359        if buffer.is_empty() || buffer.column(col).is_none() {
3360            return vec![];
3361        }
3362        let num_rows = buffer.num_rows() as usize;
3363        let mut bytes = vec![0u8; num_rows * 4];
3364        provider
3365            .device()
3366            .inner()
3367            .dtoh_sync_copy_into(buffer.column(col).unwrap(), &mut bytes)
3368            .expect("dtoh");
3369        bytes
3370            .chunks_exact(4)
3371            .map(|c| u32::from_le_bytes([c[0], c[1], c[2], c[3]]))
3372            .collect()
3373    }
3374
3375    #[test]
3376    fn test_compact_device_mask_respects_mask_len_smaller_than_row_cap() {
3377        let provider = match create_test_provider() {
3378            Some(p) => p,
3379            None => {
3380                eprintln!("Skipping test: no CUDA device available");
3381                return;
3382            }
3383        };
3384
3385        let schema = Schema::new(vec![("id".to_string(), ScalarType::U32)]);
3386        let base = create_test_buffer(&provider, &[1, 2, 3, 4, 5, 6, 7, 8], "id");
3387
3388        let row_cap = 16u64;
3389        let data: Vec<u32> = (0..row_cap as u32).collect();
3390        let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
3391        let mut col = provider.memory().alloc::<u8>(bytes.len()).expect("alloc");
3392        provider
3393            .device()
3394            .inner()
3395            .htod_sync_copy_into(&bytes, &mut col)
3396            .expect("htod");
3397        let expanded = provider
3398            .buffer_from_columns_with_device_count(vec![col.into()], row_cap, schema, &base)
3399            .expect("buffer");
3400
3401        let mask: Vec<u8> = vec![1, 0, 1, 0, 1, 0, 1, 0];
3402        let (prefix_sum, count) = provider.prefix_sum_mask(&mask).expect("prefix sum");
3403
3404        let mut d_mask = provider.memory().alloc::<u8>(mask.len()).expect("alloc");
3405        provider
3406            .device()
3407            .inner()
3408            .htod_sync_copy_into(&mask, &mut d_mask)
3409            .expect("mask htod");
3410
3411        let mut d_prefix = provider
3412            .memory()
3413            .alloc::<u32>(prefix_sum.len())
3414            .expect("alloc");
3415        provider
3416            .device()
3417            .inner()
3418            .htod_sync_copy_into(&prefix_sum, &mut d_prefix)
3419            .expect("prefix htod");
3420
3421        let mut d_out_count = provider.memory().alloc::<u32>(1).expect("alloc");
3422        provider
3423            .device()
3424            .inner()
3425            .htod_sync_copy_into(&[count], &mut d_out_count)
3426            .expect("count htod");
3427
3428        let compacted = provider
3429            .compact_buffer_by_device_mask_device_count(&expanded, &d_mask, &d_prefix, d_out_count)
3430            .expect("compact");
3431
3432        assert_eq!(compacted.num_rows(), mask.len() as u64);
3433        let device_rows = provider.device_row_count(&compacted).expect("row count");
3434        assert_eq!(device_rows as u32, count);
3435    }
3436
3437    #[test]
3438    fn test_clone_buffer_preserves_device_count() {
3439        let provider = match create_test_provider() {
3440            Some(p) => p,
3441            None => {
3442                eprintln!("Skipping test: no CUDA device available");
3443                return;
3444            }
3445        };
3446
3447        let schema = Schema::new(vec![("id".to_string(), ScalarType::U32)]);
3448        let ids: Vec<u32> = vec![10, 20, 30];
3449        let buffer = provider
3450            .create_buffer_from_slices(&[bytemuck::cast_slice(&ids)], schema)
3451            .unwrap();
3452
3453        let cloned = provider.clone_buffer(&buffer).unwrap();
3454
3455        let mut host_count = [0u32];
3456        provider
3457            .device()
3458            .inner()
3459            .dtoh_sync_copy_into(cloned.num_rows_device(), &mut host_count)
3460            .unwrap();
3461        assert_eq!(host_count[0], 3);
3462    }
3463
3464    /// `clone_buffer` must propagate the host-side `cached_row_count` so
3465    /// downstream code can read the row count without a D2H round-trip.
3466    /// Without this propagation, buffers flowed through the relation store
3467    /// (`CompiledIlpProgram::put_relation` calls `clone_buffer` before
3468    /// storing) lose their host-visible count, forcing consumers to choose
3469    /// between an extra D2H (violating the native bounded exact-induction
3470    /// transfer-budget gates) and a hard error. This test pins the cache-propagation
3471    /// contract directly.
3472    #[test]
3473    fn test_clone_buffer_preserves_cached_row_count() {
3474        let provider = match create_test_provider() {
3475            Some(p) => p,
3476            None => {
3477                eprintln!("Skipping test: no CUDA device available");
3478                return;
3479            }
3480        };
3481
3482        let schema = Schema::new(vec![("id".to_string(), ScalarType::U32)]);
3483        let ids: Vec<u32> = vec![7, 11, 13, 17];
3484        let source = provider
3485            .create_buffer_from_slices(&[bytemuck::cast_slice(&ids)], schema)
3486            .unwrap();
3487        // Source's cache is populated by the `create_buffer_from_*` path;
3488        // verify the precondition so a regression in that path shows up here
3489        // rather than silently passing the real assertion below.
3490        assert_eq!(
3491            source.cached_row_count(),
3492            Some(4),
3493            "source buffer should have its cached row count populated by \
3494             create_buffer_from_slices"
3495        );
3496
3497        let cloned = provider.clone_buffer(&source).unwrap();
3498
3499        assert_eq!(
3500            cloned.cached_row_count(),
3501            Some(4),
3502            "clone_buffer must propagate cached_row_count from source to clone",
3503        );
3504    }
3505
3506    // ============== Hash Join Tests ==============
3507
3508    #[test]
3509    fn test_hash_join_empty_inputs() {
3510        let provider = match create_test_provider() {
3511            Some(p) => p,
3512            None => {
3513                eprintln!("Skipping test: no CUDA device available");
3514                return;
3515            }
3516        };
3517
3518        let schema = Schema::new(vec![("key".to_string(), ScalarType::U32)]);
3519        let empty = create_empty_test_buffer(&provider, schema.clone());
3520
3521        // Join empty with empty
3522        let result = provider.hash_join(&empty, &empty, &[0], &[0]);
3523        assert!(result.is_ok());
3524        assert!(result.unwrap().is_empty());
3525    }
3526
3527    #[test]
3528    fn test_hash_join_validation() {
3529        let provider = match create_test_provider() {
3530            Some(p) => p,
3531            None => {
3532                eprintln!("Skipping test: no CUDA device available");
3533                return;
3534            }
3535        };
3536
3537        let left = create_test_buffer(&provider, &[1, 2, 3], "left_key");
3538        let right = create_test_buffer(&provider, &[2, 3, 4], "right_key");
3539
3540        // Empty key columns
3541        let result = provider.hash_join(&left, &right, &[], &[0]);
3542        assert!(result.is_err());
3543
3544        // Mismatched key lengths
3545        let result = provider.hash_join(&left, &right, &[0], &[0, 0]);
3546        assert!(result.is_err());
3547    }
3548
3549    // ============== Dedup Tests ==============
3550
3551    #[test]
3552    fn test_dedup_empty_input() {
3553        let provider = match create_test_provider() {
3554            Some(p) => p,
3555            None => {
3556                eprintln!("Skipping test: no CUDA device available");
3557                return;
3558            }
3559        };
3560
3561        let schema = Schema::new(vec![("key".to_string(), ScalarType::U32)]);
3562        let empty = create_empty_test_buffer(&provider, schema);
3563
3564        let result = provider.dedup(&empty, &[0]);
3565        assert!(result.is_ok());
3566        assert!(result.unwrap().is_empty());
3567    }
3568
3569    #[test]
3570    fn test_dedup_validation() {
3571        let provider = match create_test_provider() {
3572            Some(p) => p,
3573            None => {
3574                eprintln!("Skipping test: no CUDA device available");
3575                return;
3576            }
3577        };
3578
3579        let buffer = create_test_buffer(&provider, &[1, 1, 2, 2, 3], "key");
3580
3581        // Empty key columns
3582        let result = provider.dedup(&buffer, &[]);
3583        assert!(result.is_err());
3584    }
3585
3586    #[test]
3587    fn test_dedup_with_duplicates() {
3588        let provider = match create_test_provider() {
3589            Some(p) => p,
3590            None => {
3591                eprintln!("Skipping test: no CUDA device available");
3592                return;
3593            }
3594        };
3595
3596        // Test dedup with duplicates: [3, 1, 2, 1, 3, 2]
3597        let buffer = create_test_buffer(&provider, &[3, 1, 2, 1, 3, 2], "key");
3598        let deduped = provider.dedup(&buffer, &[0]).unwrap();
3599
3600        let dedup_count = provider
3601            .device_row_count(&deduped)
3602            .expect("read dedup row count");
3603        assert_eq!(dedup_count, 3, "Should have 3 unique values");
3604
3605        let result = provider.download_column::<u32>(&deduped, 0).unwrap();
3606        // Result should be sorted and deduped
3607        assert_eq!(result, vec![1, 2, 3]);
3608    }
3609
3610    #[test]
3611    fn test_dedup_larger_input() {
3612        let provider = match create_test_provider() {
3613            Some(p) => p,
3614            None => {
3615                eprintln!("Skipping test: no CUDA device available");
3616                return;
3617            }
3618        };
3619
3620        // Create input with duplicates: 0..500 ++ 250..750 = 1000 elements, 750 unique
3621        let a: Vec<u32> = (0..500).collect();
3622        let b: Vec<u32> = (250..750).collect();
3623        let input: Vec<u32> = a.iter().chain(b.iter()).copied().collect();
3624
3625        let buffer = create_test_buffer(&provider, &input, "key");
3626        let deduped = provider.dedup(&buffer, &[0]).unwrap();
3627
3628        let dedup_count = provider
3629            .device_row_count(&deduped)
3630            .expect("read dedup row count");
3631        assert_eq!(dedup_count, 750, "Should have 750 unique values (0..750)");
3632
3633        // Verify output is sorted
3634        let result = provider.download_column::<u32>(&deduped, 0).unwrap();
3635        let is_sorted = result.windows(2).all(|w| w[0] <= w[1]);
3636        assert!(is_sorted, "Output should be sorted");
3637
3638        // Verify expected values
3639        let expected: Vec<u32> = (0..750).collect();
3640        assert_eq!(result, expected);
3641    }
3642
3643    // ============== Union Tests ==============
3644
3645    #[test]
3646    fn test_union_empty_inputs() {
3647        let provider = match create_test_provider() {
3648            Some(p) => p,
3649            None => {
3650                eprintln!("Skipping test: no CUDA device available");
3651                return;
3652            }
3653        };
3654
3655        let schema = Schema::new(vec![("key".to_string(), ScalarType::U32)]);
3656        let empty = create_empty_test_buffer(&provider, schema.clone());
3657
3658        // Empty union empty
3659        let result = provider.union(&empty, &empty);
3660        assert!(result.is_ok());
3661        assert!(result.unwrap().is_empty());
3662
3663        // Non-empty union empty
3664        let a = create_test_buffer(&provider, &[1, 2, 3], "key");
3665        let empty2 = create_empty_test_buffer(&provider, schema);
3666        let result = provider.union(&a, &empty2);
3667        assert!(result.is_ok());
3668        let result = result.unwrap();
3669        assert_eq!(result.num_rows(), 3);
3670    }
3671
3672    #[test]
3673    fn test_union_schema_type_mismatch() {
3674        let provider = match create_test_provider() {
3675            Some(p) => p,
3676            None => {
3677                eprintln!("Skipping test: no CUDA device available");
3678                return;
3679            }
3680        };
3681
3682        let a = create_test_buffer(&provider, &[1, 2], "col_a");
3683        let b = create_test_buffer(&provider, &[3, 4], "col_b");
3684
3685        // Different column names but same types should succeed (Datalog union semantics)
3686        let result = provider.union(&a, &b);
3687        assert!(result.is_ok());
3688
3689        // Different arity should fail - create a 2-column buffer
3690        let two_col_schema = Schema::new(vec![
3691            ("x".to_string(), ScalarType::U32),
3692            ("y".to_string(), ScalarType::U32),
3693        ]);
3694        let c = provider
3695            .create_buffer_from_u32_columns(&[&[1, 2], &[3, 4]], two_col_schema)
3696            .unwrap();
3697        let result = provider.union(&a, &c);
3698        assert!(result.is_err());
3699    }
3700
3701    // ============== Diff Tests ==============
3702
3703    #[test]
3704    fn test_diff_empty_inputs() {
3705        let provider = match create_test_provider() {
3706            Some(p) => p,
3707            None => {
3708                eprintln!("Skipping test: no CUDA device available");
3709                return;
3710            }
3711        };
3712
3713        let schema = Schema::new(vec![("key".to_string(), ScalarType::U32)]);
3714        let empty = create_empty_test_buffer(&provider, schema.clone());
3715
3716        // Empty diff empty
3717        let result = provider.diff(&empty, &empty);
3718        assert!(result.is_ok());
3719        assert!(result.unwrap().is_empty());
3720
3721        // Non-empty diff empty should return all of a
3722        let a = create_test_buffer(&provider, &[1, 2, 3], "key");
3723        let empty2 = create_empty_test_buffer(&provider, schema);
3724        let result = provider.diff(&a, &empty2);
3725        assert!(result.is_ok());
3726        let result = result.unwrap();
3727        assert_eq!(result.num_rows(), 3);
3728    }
3729
3730    #[test]
3731    fn test_diff_basic() {
3732        let provider = match create_test_provider() {
3733            Some(p) => p,
3734            None => {
3735                eprintln!("Skipping test: no CUDA device available");
3736                return;
3737            }
3738        };
3739
3740        let a = create_test_buffer(&provider, &[1, 2, 3, 4, 5], "key");
3741        let b = create_test_buffer(&provider, &[2, 4], "key");
3742
3743        let result = provider.diff(&a, &b);
3744        assert!(result.is_ok());
3745        let result = result.unwrap();
3746        assert_eq!(result.num_rows(), 3); // 1, 3, 5
3747
3748        let values = read_buffer_u32(&provider, &result, 0);
3749        assert_eq!(values, vec![1, 3, 5]);
3750    }
3751
3752    #[test]
3753    fn test_diff_all_filtered_out() {
3754        let provider = match create_test_provider() {
3755            Some(p) => p,
3756            None => {
3757                eprintln!("Skipping test: no CUDA device available");
3758                return;
3759            }
3760        };
3761
3762        let a = create_test_buffer(&provider, &[1, 2, 3], "key");
3763        let b = create_test_buffer(&provider, &[1, 2, 3, 4, 5], "key");
3764
3765        let result = provider.diff(&a, &b);
3766        assert!(result.is_ok());
3767        assert!(result.unwrap().is_empty());
3768    }
3769
3770    #[test]
3771    fn test_diff_schema_mismatch() {
3772        let provider = match create_test_provider() {
3773            Some(p) => p,
3774            None => {
3775                eprintln!("Skipping test: no CUDA device available");
3776                return;
3777            }
3778        };
3779
3780        // Different column names with same types should work (Datalog semantics)
3781        let a = create_test_buffer(&provider, &[1, 2], "col_a");
3782        let b = create_test_buffer(&provider, &[1, 2], "col_b");
3783        let result = provider.diff(&a, &b);
3784        assert!(
3785            result.is_ok(),
3786            "Same types with different names should succeed"
3787        );
3788
3789        // Create buffers with different arities (this should fail)
3790        let schema_2col = Schema::new(vec![
3791            ("c0".to_string(), ScalarType::U32),
3792            ("c1".to_string(), ScalarType::U32),
3793        ]);
3794
3795        let bytes_2col: Vec<u8> = [1u32, 2, 3, 4]
3796            .iter()
3797            .flat_map(|v| v.to_le_bytes())
3798            .collect();
3799        let mut col0 = provider
3800            .memory()
3801            .alloc::<u8>(bytes_2col.len() / 2)
3802            .expect("alloc");
3803        let mut col1 = provider
3804            .memory()
3805            .alloc::<u8>(bytes_2col.len() / 2)
3806            .expect("alloc");
3807        provider
3808            .device()
3809            .inner()
3810            .htod_sync_copy_into(&bytes_2col[..8], &mut col0)
3811            .expect("htod");
3812        provider
3813            .device()
3814            .inner()
3815            .htod_sync_copy_into(&bytes_2col[8..], &mut col1)
3816            .expect("htod");
3817        let buffer_2col = provider
3818            .buffer_from_columns(vec![col0.into(), col1.into()], 2, schema_2col)
3819            .expect("buffer");
3820
3821        let buffer_1col = create_test_buffer(&provider, &[1, 2], "c0");
3822
3823        let result = provider.diff(&buffer_2col, &buffer_1col);
3824        assert!(result.is_err(), "Different arities should fail");
3825    }
3826
3827    // ============== GroupBy Aggregation Tests ==============
3828
3829    #[test]
3830    fn test_groupby_empty_input() {
3831        let provider = match create_test_provider() {
3832            Some(p) => p,
3833            None => {
3834                eprintln!("Skipping test: no CUDA device available");
3835                return;
3836            }
3837        };
3838
3839        let schema = Schema::new(vec![("key".to_string(), ScalarType::U32)]);
3840        let empty = create_empty_test_buffer(&provider, schema);
3841
3842        let result = provider.groupby_agg(&empty, &[0], AggOp::Count, 0);
3843        assert!(result.is_ok());
3844        assert!(result.unwrap().is_empty());
3845    }
3846
3847    #[test]
3848    fn test_groupby_validation() {
3849        let provider = match create_test_provider() {
3850            Some(p) => p,
3851            None => {
3852                eprintln!("Skipping test: no CUDA device available");
3853                return;
3854            }
3855        };
3856
3857        let buffer = create_test_buffer(&provider, &[1, 1, 2, 2, 3], "key");
3858
3859        // Empty key columns
3860        let result = provider.groupby_agg(&buffer, &[], AggOp::Count, 0);
3861        assert!(result.is_err());
3862
3863        // Value column out of bounds
3864        let result = provider.groupby_agg(&buffer, &[0], AggOp::Count, 5);
3865        assert!(result.is_err());
3866    }
3867
3868    #[test]
3869    fn test_groupby_logsumexp() {
3870        let provider = match create_test_provider() {
3871            Some(p) => p,
3872            None => {
3873                eprintln!("Skipping test: no CUDA device available");
3874                return;
3875            }
3876        };
3877
3878        // Create buffer with U32 keys and F64 values
3879        // Group 0 (key=1): values 1.0, 2.0 -> logsumexp = log(e^1 + e^2) ≈ 2.31326
3880        // Group 1 (key=2): values 3.0, 4.0 -> logsumexp = log(e^3 + e^4) ≈ 4.31326
3881        let keys: Vec<u32> = vec![1, 1, 2, 2];
3882        let values: Vec<f64> = vec![1.0, 2.0, 3.0, 4.0];
3883
3884        let schema = Schema::new(vec![
3885            ("key".to_string(), ScalarType::U32),
3886            ("value".to_string(), ScalarType::F64),
3887        ]);
3888
3889        // Create key column
3890        let key_bytes: Vec<u8> = keys.iter().flat_map(|v| v.to_le_bytes()).collect();
3891        let mut key_col = provider
3892            .memory()
3893            .alloc::<u8>(key_bytes.len())
3894            .expect("alloc key");
3895        provider
3896            .device()
3897            .inner()
3898            .htod_sync_copy_into(&key_bytes, &mut key_col)
3899            .expect("upload key");
3900
3901        // Create value column
3902        let val_bytes: Vec<u8> = values.iter().flat_map(|v| v.to_le_bytes()).collect();
3903        let mut val_col = provider
3904            .memory()
3905            .alloc::<u8>(val_bytes.len())
3906            .expect("alloc val");
3907        provider
3908            .device()
3909            .inner()
3910            .htod_sync_copy_into(&val_bytes, &mut val_col)
3911            .expect("upload val");
3912
3913        let buffer = provider
3914            .buffer_from_columns(vec![key_col.into(), val_col.into()], 4, schema)
3915            .expect("buffer");
3916
3917        // Run LogSumExp aggregation grouped by key column (0), aggregating value column (1)
3918        let result = provider.groupby_agg(&buffer, &[0], AggOp::LogSumExp, 1);
3919        assert!(
3920            result.is_ok(),
3921            "groupby_agg with LogSumExp should succeed: {:?}",
3922            result.err()
3923        );
3924
3925        let result = result.unwrap();
3926        let group_count = provider
3927            .device_row_count(&result)
3928            .expect("read group count");
3929        assert_eq!(group_count, 2, "Should have 2 groups");
3930
3931        // Download results
3932        let result_values = provider
3933            .download_column::<f64>(&result, 1)
3934            .expect("download result");
3935
3936        // Expected values:
3937        // logsumexp(1.0, 2.0) = 2.0 + log(exp(1.0-2.0) + exp(2.0-2.0)) = 2.0 + log(e^-1 + 1) ≈ 2.31326
3938        // logsumexp(3.0, 4.0) = 4.0 + log(exp(3.0-4.0) + exp(4.0-4.0)) = 4.0 + log(e^-1 + 1) ≈ 4.31326
3939        let expected_0 = 2.0_f64 + ((-1.0_f64).exp() + 1.0_f64).ln(); // ≈ 2.31326
3940        let expected_1 = 4.0_f64 + ((-1.0_f64).exp() + 1.0_f64).ln(); // ≈ 4.31326
3941
3942        let tolerance = 1e-5;
3943        assert!(
3944            (result_values[0] - expected_0).abs() < tolerance,
3945            "Group 0 logsumexp mismatch: got {}, expected {}",
3946            result_values[0],
3947            expected_0
3948        );
3949        assert!(
3950            (result_values[1] - expected_1).abs() < tolerance,
3951            "Group 1 logsumexp mismatch: got {}, expected {}",
3952            result_values[1],
3953            expected_1
3954        );
3955    }
3956
3957    // ============== Schema Helper Tests ==============
3958
3959    #[test]
3960    fn test_combine_schemas() {
3961        let provider = match create_test_provider() {
3962            Some(p) => p,
3963            None => {
3964                eprintln!("Skipping test: no CUDA device available");
3965                return;
3966            }
3967        };
3968
3969        let left = Schema::new(vec![("a".to_string(), ScalarType::U32)]);
3970        let right = Schema::new(vec![("b".to_string(), ScalarType::U64)]);
3971
3972        let combined = provider.combine_schemas(&left, &right);
3973        assert_eq!(combined.arity(), 2);
3974        assert_eq!(combined.column_type(0), Some(ScalarType::U32));
3975        assert_eq!(combined.column_type(1), Some(ScalarType::U64));
3976    }
3977
3978    #[test]
3979    fn test_groupby_result_schema() {
3980        let provider = match create_test_provider() {
3981            Some(p) => p,
3982            None => {
3983                eprintln!("Skipping test: no CUDA device available");
3984                return;
3985            }
3986        };
3987
3988        let input = Schema::new(vec![
3989            ("key".to_string(), ScalarType::U32),
3990            ("value".to_string(), ScalarType::U32),
3991        ]);
3992
3993        // Count result schema (u64 to match predicate declarations)
3994        let count_schema =
3995            provider.groupby_multi_agg_result_schema(&input, &[0], &[(1, AggOp::Count)]);
3996        assert_eq!(count_schema.arity(), 2);
3997        assert_eq!(count_schema.column_type(1), Some(ScalarType::U64));
3998
3999        // Sum result schema
4000        let sum_schema = provider.groupby_multi_agg_result_schema(&input, &[0], &[(1, AggOp::Sum)]);
4001        assert_eq!(sum_schema.arity(), 2);
4002        assert_eq!(sum_schema.column_type(1), Some(ScalarType::U64));
4003
4004        // Min/Max result schema
4005        let min_schema = provider.groupby_multi_agg_result_schema(&input, &[0], &[(1, AggOp::Min)]);
4006        assert_eq!(min_schema.arity(), 2);
4007        assert_eq!(min_schema.column_type(1), Some(ScalarType::U32));
4008    }
4009
4010    #[test]
4011    fn test_groupby_multi_agg_sum_returns_u64_schema() {
4012        let provider = match create_test_provider() {
4013            Some(p) => p,
4014            None => {
4015                eprintln!("Skipping test: no CUDA device");
4016                return;
4017            }
4018        };
4019
4020        let schema = Schema::new(vec![
4021            ("key".to_string(), ScalarType::U32),
4022            ("val".to_string(), ScalarType::U32),
4023        ]);
4024
4025        let result_schema =
4026            provider.groupby_multi_agg_result_schema(&schema, &[0], &[(1, AggOp::Sum)]);
4027
4028        // Sum should return U64 to prevent overflow
4029        assert_eq!(
4030            result_schema.column_type(1),
4031            Some(ScalarType::U64),
4032            "Sum aggregation should return U64 type, not U32"
4033        );
4034    }
4035
4036    #[test]
4037    fn test_join_custom_max_output() {
4038        let provider = match create_test_provider() {
4039            Some(p) => p,
4040            None => {
4041                eprintln!("Skipping test: no CUDA device available");
4042                return;
4043            }
4044        };
4045
4046        // Create buffers that produce more than 10 results when joined
4047        // Left: [1, 1, 1, 1, 2, 2, 2, 2] - 4 copies of 1, 4 copies of 2
4048        // Right: [1, 1, 1, 2, 2, 2] - 3 copies of 1, 3 copies of 2
4049        // Join produces: 4*3 + 4*3 = 24 results
4050        let left = create_test_buffer(&provider, &[1, 1, 1, 1, 2, 2, 2, 2], "left_key");
4051        let right = create_test_buffer(&provider, &[1, 1, 1, 2, 2, 2], "right_key");
4052
4053        // Test with limit of 10 - should get at most 10
4054        let result_limited = provider
4055            .hash_join_v2_with_limit(&left, &right, &[0], &[0], JoinType::Inner, Some(10))
4056            .expect("join with limit should succeed");
4057        assert!(
4058            result_limited.num_rows() <= 10,
4059            "With limit 10, got {} rows but expected at most 10",
4060            result_limited.num_rows()
4061        );
4062
4063        // Test with None (default) - should get all 24 results
4064        let result_unlimited = provider
4065            .hash_join_v2_with_limit(&left, &right, &[0], &[0], JoinType::Inner, None)
4066            .expect("join without limit should succeed");
4067        assert_eq!(
4068            result_unlimited.num_rows(),
4069            24,
4070            "Without limit, expected 24 rows but got {}",
4071            result_unlimited.num_rows()
4072        );
4073
4074        // Test legacy API still works (backward compatibility)
4075        let result_legacy = provider
4076            .hash_join_v2(&left, &right, &[0], &[0], JoinType::Inner)
4077            .expect("legacy hash_join_v2 should succeed");
4078        assert_eq!(
4079            result_legacy.num_rows(),
4080            24,
4081            "Legacy API without limit, expected 24 rows but got {}",
4082            result_legacy.num_rows()
4083        );
4084    }
4085
4086    // ============== Arithmetic Operation Tests ==============
4087
4088    /// Helper to create a test provider for arithmetic tests
4089    fn create_arith_test_provider() -> Option<CudaKernelProvider> {
4090        if !has_cuda_device() {
4091            return None;
4092        }
4093        let device = Arc::new(CudaDevice::new(0).ok()?);
4094        let budget = MemoryBudget::with_limit(1024 * 1024 * 1024);
4095        let memory = Arc::new(GpuMemoryManager::new(device.clone(), budget));
4096        CudaKernelProvider::new(device, memory).ok()
4097    }
4098
4099    /// Helper to create an i64 buffer for arithmetic tests
4100    fn create_i64_buffer(provider: &CudaKernelProvider, data: &[i64]) -> CudaBuffer {
4101        let schema = Schema::new(vec![("col".to_string(), ScalarType::I64)]);
4102        provider
4103            .create_buffer_from_slice::<i64>(data, schema)
4104            .unwrap()
4105    }
4106
4107    /// Helper to create an f64 buffer for arithmetic tests
4108    fn create_f64_buffer(provider: &CudaKernelProvider, data: &[f64]) -> CudaBuffer {
4109        let schema = Schema::new(vec![("col".to_string(), ScalarType::F64)]);
4110        provider
4111            .create_buffer_from_slice::<f64>(data, schema)
4112            .unwrap()
4113    }
4114
4115    #[test]
4116    fn test_add_columns_i64() {
4117        let Some(provider) = create_arith_test_provider() else {
4118            eprintln!("Skipping test: no CUDA device available");
4119            return;
4120        };
4121
4122        let a = create_i64_buffer(&provider, &[1, 2, 3, 4, 5]);
4123        let b = create_i64_buffer(&provider, &[10, 20, 30, 40, 50]);
4124
4125        let result = provider.add_columns(&a, &b).unwrap();
4126        let values = provider.download_column::<i64>(&result, 0).unwrap();
4127
4128        assert_eq!(values, vec![11, 22, 33, 44, 55]);
4129    }
4130
4131    #[test]
4132    fn test_sub_columns_i64() {
4133        let Some(provider) = create_arith_test_provider() else {
4134            eprintln!("Skipping test: no CUDA device available");
4135            return;
4136        };
4137
4138        let a = create_i64_buffer(&provider, &[10, 20, 30, 40, 50]);
4139        let b = create_i64_buffer(&provider, &[1, 2, 3, 4, 5]);
4140
4141        let result = provider.sub_columns(&a, &b).unwrap();
4142        let values = provider.download_column::<i64>(&result, 0).unwrap();
4143
4144        assert_eq!(values, vec![9, 18, 27, 36, 45]);
4145    }
4146
4147    #[test]
4148    fn test_mul_columns_i64() {
4149        let Some(provider) = create_arith_test_provider() else {
4150            eprintln!("Skipping test: no CUDA device available");
4151            return;
4152        };
4153
4154        let a = create_i64_buffer(&provider, &[2, 3, 4, 5, 6]);
4155        let b = create_i64_buffer(&provider, &[3, 4, 5, 6, 7]);
4156
4157        let result = provider.mul_columns(&a, &b).unwrap();
4158        let values = provider.download_column::<i64>(&result, 0).unwrap();
4159
4160        assert_eq!(values, vec![6, 12, 20, 30, 42]);
4161    }
4162
4163    #[test]
4164    fn test_div_columns_i64() {
4165        let Some(provider) = create_arith_test_provider() else {
4166            eprintln!("Skipping test: no CUDA device available");
4167            return;
4168        };
4169
4170        let a = create_i64_buffer(&provider, &[100, 200, 300, 400]);
4171        let b = create_i64_buffer(&provider, &[10, 20, 30, 40]);
4172
4173        let result = provider.div_columns(&a, &b).unwrap();
4174        let values = provider.download_column::<i64>(&result, 0).unwrap();
4175
4176        assert_eq!(values, vec![10, 10, 10, 10]);
4177    }
4178
4179    #[test]
4180    fn test_div_columns_by_zero() {
4181        let Some(provider) = create_arith_test_provider() else {
4182            eprintln!("Skipping test: no CUDA device available");
4183            return;
4184        };
4185
4186        let a = create_i64_buffer(&provider, &[10, 20, 30]);
4187        let b = create_i64_buffer(&provider, &[2, 0, 3]); // Note: division by zero
4188
4189        let result = provider.div_columns(&a, &b).unwrap();
4190        let values = provider.download_column::<i64>(&result, 0).unwrap();
4191
4192        // Division by zero returns i64::MAX
4193        assert_eq!(values, vec![5, i64::MAX, 10]);
4194    }
4195
4196    #[test]
4197    fn test_mod_columns_i64() {
4198        let Some(provider) = create_arith_test_provider() else {
4199            eprintln!("Skipping test: no CUDA device available");
4200            return;
4201        };
4202
4203        let a = create_i64_buffer(&provider, &[17, 23, 100, 7]);
4204        let b = create_i64_buffer(&provider, &[5, 7, 30, 3]);
4205
4206        let result = provider.mod_columns(&a, &b).unwrap();
4207        let values = provider.download_column::<i64>(&result, 0).unwrap();
4208
4209        assert_eq!(values, vec![2, 2, 10, 1]);
4210    }
4211
4212    #[test]
4213    fn test_mod_columns_by_zero() {
4214        let Some(provider) = create_arith_test_provider() else {
4215            eprintln!("Skipping test: no CUDA device available");
4216            return;
4217        };
4218
4219        let a = create_i64_buffer(&provider, &[10, 20]);
4220        let b = create_i64_buffer(&provider, &[3, 0]); // Note: mod by zero
4221
4222        let result = provider.mod_columns(&a, &b).unwrap();
4223        let values = provider.download_column::<i64>(&result, 0).unwrap();
4224
4225        // Mod by zero returns 0
4226        assert_eq!(values, vec![1, 0]);
4227    }
4228
4229    #[test]
4230    fn test_abs_column_i64() {
4231        let Some(provider) = create_arith_test_provider() else {
4232            eprintln!("Skipping test: no CUDA device available");
4233            return;
4234        };
4235
4236        let a = create_i64_buffer(&provider, &[-5, 10, -15, 20, 0]);
4237
4238        let result = provider.abs_column(&a).unwrap();
4239        let values = provider.download_column::<i64>(&result, 0).unwrap();
4240
4241        assert_eq!(values, vec![5, 10, 15, 20, 0]);
4242    }
4243
4244    #[test]
4245    fn test_min_columns_i64() {
4246        let Some(provider) = create_arith_test_provider() else {
4247            eprintln!("Skipping test: no CUDA device available");
4248            return;
4249        };
4250
4251        let a = create_i64_buffer(&provider, &[5, 10, 15, 20]);
4252        let b = create_i64_buffer(&provider, &[3, 12, 10, 25]);
4253
4254        let result = provider.min_columns(&a, &b).unwrap();
4255        let values = provider.download_column::<i64>(&result, 0).unwrap();
4256
4257        assert_eq!(values, vec![3, 10, 10, 20]);
4258    }
4259
4260    #[test]
4261    fn test_max_columns_i64() {
4262        let Some(provider) = create_arith_test_provider() else {
4263            eprintln!("Skipping test: no CUDA device available");
4264            return;
4265        };
4266
4267        let a = create_i64_buffer(&provider, &[5, 10, 15, 20]);
4268        let b = create_i64_buffer(&provider, &[3, 12, 10, 25]);
4269
4270        let result = provider.max_columns(&a, &b).unwrap();
4271        let values = provider.download_column::<i64>(&result, 0).unwrap();
4272
4273        assert_eq!(values, vec![5, 12, 15, 25]);
4274    }
4275
4276    #[test]
4277    fn test_add_columns_f64() {
4278        let Some(provider) = create_arith_test_provider() else {
4279            eprintln!("Skipping test: no CUDA device available");
4280            return;
4281        };
4282
4283        let a = create_f64_buffer(&provider, &[1.5, 2.5, 3.5]);
4284        let b = create_f64_buffer(&provider, &[0.5, 1.5, 2.5]);
4285
4286        let result = provider.add_columns(&a, &b).unwrap();
4287        let values = provider.download_column::<f64>(&result, 0).unwrap();
4288
4289        assert_eq!(values, vec![2.0, 4.0, 6.0]);
4290    }
4291
4292    #[test]
4293    fn test_mul_columns_f64() {
4294        let Some(provider) = create_arith_test_provider() else {
4295            eprintln!("Skipping test: no CUDA device available");
4296            return;
4297        };
4298
4299        let a = create_f64_buffer(&provider, &[2.0, 3.0, 4.0]);
4300        let b = create_f64_buffer(&provider, &[1.5, 2.0, 2.5]);
4301
4302        let result = provider.mul_columns(&a, &b).unwrap();
4303        let values = provider.download_column::<f64>(&result, 0).unwrap();
4304
4305        assert_eq!(values, vec![3.0, 6.0, 10.0]);
4306    }
4307
4308    #[test]
4309    fn test_div_columns_f64_by_zero() {
4310        let Some(provider) = create_arith_test_provider() else {
4311            eprintln!("Skipping test: no CUDA device available");
4312            return;
4313        };
4314
4315        let a = create_f64_buffer(&provider, &[1.0, -1.0, 0.0]);
4316        let b = create_f64_buffer(&provider, &[0.0, 0.0, 0.0]);
4317
4318        let result = provider.div_columns(&a, &b).unwrap();
4319        let values = provider.download_column::<f64>(&result, 0).unwrap();
4320
4321        // IEEE 754: 1.0/0.0 = Inf, -1.0/0.0 = -Inf, 0.0/0.0 = NaN
4322        assert!(values[0].is_infinite() && values[0].is_sign_positive());
4323        assert!(values[1].is_infinite() && values[1].is_sign_negative());
4324        assert!(values[2].is_nan());
4325    }
4326
4327    #[test]
4328    fn test_pow_columns() {
4329        let Some(provider) = create_arith_test_provider() else {
4330            eprintln!("Skipping test: no CUDA device available");
4331            return;
4332        };
4333
4334        let base = create_i64_buffer(&provider, &[2, 3, 4, 5]);
4335        let exp = create_i64_buffer(&provider, &[3, 2, 2, 1]);
4336
4337        let result = provider.pow_columns(&base, &exp).unwrap();
4338        let values = provider.download_column::<f64>(&result, 0).unwrap();
4339
4340        // pow always returns f64
4341        assert_eq!(values, vec![8.0, 9.0, 16.0, 5.0]);
4342    }
4343
4344    #[test]
4345    fn test_pow_columns_fractional_exp() {
4346        let Some(provider) = create_arith_test_provider() else {
4347            eprintln!("Skipping test: no CUDA device available");
4348            return;
4349        };
4350
4351        let base = create_f64_buffer(&provider, &[4.0, 9.0, 27.0]);
4352        let exp = create_f64_buffer(&provider, &[0.5, 0.5, 1.0 / 3.0]);
4353
4354        let result = provider.pow_columns(&base, &exp).unwrap();
4355        let values = provider.download_column::<f64>(&result, 0).unwrap();
4356
4357        // sqrt(4) = 2, sqrt(9) = 3, cbrt(27) = 3
4358        assert!((values[0] - 2.0).abs() < 1e-10);
4359        assert!((values[1] - 3.0).abs() < 1e-10);
4360        assert!((values[2] - 3.0).abs() < 1e-10);
4361    }
4362
4363    #[test]
4364    fn test_cast_i64_to_f64() {
4365        let Some(provider) = create_arith_test_provider() else {
4366            eprintln!("Skipping test: no CUDA device available");
4367            return;
4368        };
4369
4370        let a = create_i64_buffer(&provider, &[1, 2, 3, 4, 5]);
4371
4372        let result = provider.cast_column(&a, ScalarType::F64).unwrap();
4373        let values = provider.download_column::<f64>(&result, 0).unwrap();
4374
4375        assert_eq!(values, vec![1.0, 2.0, 3.0, 4.0, 5.0]);
4376    }
4377
4378    #[test]
4379    fn test_cast_f64_to_i64() {
4380        let Some(provider) = create_arith_test_provider() else {
4381            eprintln!("Skipping test: no CUDA device available");
4382            return;
4383        };
4384
4385        let a = create_f64_buffer(&provider, &[1.9, 2.1, 3.5, 4.0, 5.7]);
4386
4387        let result = provider.cast_column(&a, ScalarType::I64).unwrap();
4388        let values = provider.download_column::<i64>(&result, 0).unwrap();
4389
4390        // Truncation towards zero
4391        assert_eq!(values, vec![1, 2, 3, 4, 5]);
4392    }
4393
4394    #[test]
4395    fn test_cast_i64_to_i32() {
4396        let Some(provider) = create_arith_test_provider() else {
4397            eprintln!("Skipping test: no CUDA device available");
4398            return;
4399        };
4400
4401        let a = create_i64_buffer(&provider, &[1, 2, 3, 100, 200]);
4402
4403        let result = provider.cast_column(&a, ScalarType::I32).unwrap();
4404        let values = provider.download_column::<i32>(&result, 0).unwrap();
4405
4406        assert_eq!(values, vec![1, 2, 3, 100, 200]);
4407    }
4408
4409    #[test]
4410    fn test_arithmetic_row_count_mismatch() {
4411        let Some(provider) = create_arith_test_provider() else {
4412            eprintln!("Skipping test: no CUDA device available");
4413            return;
4414        };
4415
4416        let a = create_i64_buffer(&provider, &[1, 2, 3]);
4417        let b = create_i64_buffer(&provider, &[1, 2]); // Different size
4418
4419        let result = provider.add_columns(&a, &b);
4420        assert!(result.is_err());
4421        let err = result.err().unwrap();
4422        assert!(err.to_string().contains("Row count mismatch"));
4423    }
4424
4425    #[test]
4426    fn test_arithmetic_empty_buffers() {
4427        let Some(provider) = create_arith_test_provider() else {
4428            eprintln!("Skipping test: no CUDA device available");
4429            return;
4430        };
4431
4432        let a = create_i64_buffer(&provider, &[]);
4433        let b = create_i64_buffer(&provider, &[]);
4434
4435        let result = provider.add_columns(&a, &b).unwrap();
4436        let values = provider.download_column::<i64>(&result, 0).unwrap();
4437
4438        assert_eq!(values, Vec::<i64>::new());
4439    }
4440
4441    #[test]
4442    fn test_wrapping_arithmetic_overflow() {
4443        let Some(provider) = create_arith_test_provider() else {
4444            eprintln!("Skipping test: no CUDA device available");
4445            return;
4446        };
4447
4448        let a = create_i64_buffer(&provider, &[i64::MAX, i64::MIN]);
4449        let b = create_i64_buffer(&provider, &[1, -1]);
4450
4451        // Addition should wrap
4452        let add_result = provider.add_columns(&a, &b).unwrap();
4453        let add_values = provider.download_column::<i64>(&add_result, 0).unwrap();
4454        assert_eq!(add_values[0], i64::MIN); // MAX + 1 wraps to MIN
4455        assert_eq!(add_values[1], i64::MAX); // MIN - 1 wraps to MAX
4456    }
4457
4458    #[test]
4459    fn test_abs_column_f64() {
4460        let Some(provider) = create_arith_test_provider() else {
4461            eprintln!("Skipping test: no CUDA device available");
4462            return;
4463        };
4464
4465        let a = create_f64_buffer(&provider, &[-1.5, 2.5, -3.5, 0.0]);
4466
4467        let result = provider.abs_column(&a).unwrap();
4468        let values = provider.download_column::<f64>(&result, 0).unwrap();
4469
4470        assert_eq!(values, vec![1.5, 2.5, 3.5, 0.0]);
4471    }
4472
4473    #[test]
4474    fn test_min_max_columns_f64() {
4475        let Some(provider) = create_arith_test_provider() else {
4476            eprintln!("Skipping test: no CUDA device available");
4477            return;
4478        };
4479
4480        let a = create_f64_buffer(&provider, &[1.5, 5.0, 3.0]);
4481        let b = create_f64_buffer(&provider, &[2.0, 3.0, 4.0]);
4482
4483        let min_result = provider.min_columns(&a, &b).unwrap();
4484        let min_values = provider.download_column::<f64>(&min_result, 0).unwrap();
4485        assert_eq!(min_values, vec![1.5, 3.0, 3.0]);
4486
4487        let max_result = provider.max_columns(&a, &b).unwrap();
4488        let max_values = provider.download_column::<f64>(&max_result, 0).unwrap();
4489        assert_eq!(max_values, vec![2.0, 5.0, 4.0]);
4490    }
4491}