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