xlog_prob/
neural_fast_path.rs

1//! GPU neural fast-path helpers (device slot mapping + AD-chain glue).
2//!
3//! This module contains GPU-resident tables used to map neural predicate outputs
4//! (probability vectors) to CNF variable ids in the compiled circuit.
5
6use xlog_core::{Result, XlogError};
7use xlog_cuda::memory::{DeviceMemoryView, TrackedCudaSlice};
8use xlog_cuda::CudaKernelProvider;
9
10/// Configuration for GPU neural fast-path weight injection.
11///
12/// Controls numerical stability parameters for mapping neural-network
13/// output probabilities to CNF variable log-weights.
14#[derive(Debug, Clone, Copy)]
15#[non_exhaustive]
16pub struct NeuralFastPathConfig {
17    /// Probability mass reserved for the implicit "none" outcome.
18    pub eps: f64,
19    /// Minimum probability clamp used for numerical stability.
20    pub min_p: f64,
21}
22
23impl Default for NeuralFastPathConfig {
24    fn default() -> Self {
25        Self {
26            eps: 1e-7,
27            min_p: 1e-12,
28        }
29    }
30}
31
32/// Device-resident mapping from neural output slots to CNF variable ids.
33///
34/// Slots are grouped (one group per neural predicate instance). Each slot is a
35/// CNF var id (DIMACS, 1-based) whose log-weights should be updated from the
36/// group’s probability vector.
37pub struct GpuWeightSlots {
38    group_offsets_host: Vec<u32>,
39    group_offsets: TrackedCudaSlice<u32>, // len = num_groups + 1
40    slot_cnf_var: TrackedCudaSlice<u32>,  // len = total_slots
41}
42
43impl GpuWeightSlots {
44    /// Upload a slot mapping from host vectors.
45    ///
46    /// `groups[g][i]` is the CNF variable id corresponding to label/slot `i` of group `g`.
47    pub fn upload(provider: &CudaKernelProvider, groups: &[Vec<u32>]) -> Result<Self> {
48        if groups.len() > u32::MAX as usize {
49            return Err(XlogError::Compilation(
50                "Neural fast-path group count exceeds GPU u32 index space".to_string(),
51            ));
52        }
53        let offset_count = groups.len().checked_add(1).ok_or_else(|| {
54            XlogError::Compilation("Neural fast-path group offset count overflow".to_string())
55        })?;
56        let total_slots = groups.iter().try_fold(0usize, |acc, group| {
57            acc.checked_add(group.len()).ok_or_else(|| {
58                XlogError::Compilation("Neural fast-path slot count overflow".to_string())
59            })
60        })?;
61        if total_slots > u32::MAX as usize {
62            return Err(XlogError::Compilation(
63                "Neural fast-path slot count exceeds GPU u32 index space".to_string(),
64            ));
65        }
66
67        let mut offsets: Vec<u32> = Vec::with_capacity(offset_count);
68        offsets.push(0);
69
70        let mut flat: Vec<u32> = Vec::with_capacity(total_slots);
71        for g in groups {
72            flat.extend_from_slice(g);
73            let next_offset = u32::try_from(flat.len()).map_err(|_| {
74                XlogError::Compilation(
75                    "Neural fast-path slot offset exceeds GPU u32 index space".to_string(),
76                )
77            })?;
78            offsets.push(next_offset);
79        }
80
81        let memory = provider.memory().clone();
82
83        let mut d_offsets = memory.alloc::<u32>(offsets.len())?;
84        provider
85            .htod_sync_copy_into_tracked(&offsets, &mut d_offsets)
86            .map_err(|e| {
87                XlogError::Kernel(format!("Failed to upload weight slot offsets: {}", e))
88            })?;
89
90        let mut d_vars = memory.alloc::<u32>(flat.len())?;
91        provider
92            .htod_sync_copy_into_tracked(&flat, &mut d_vars)
93            .map_err(|e| XlogError::Kernel(format!("Failed to upload weight slot vars: {}", e)))?;
94
95        Ok(Self {
96            group_offsets_host: offsets,
97            group_offsets: d_offsets,
98            slot_cnf_var: d_vars,
99        })
100    }
101
102    /// Returns the number of neural-output groups as a GPU-sized count.
103    pub fn num_groups(&self) -> u32 {
104        debug_assert!(!self.group_offsets_host.is_empty());
105        (self.group_offsets_host.len() - 1) as u32
106    }
107
108    /// Returns the number of neural-output groups as a host index.
109    pub fn num_groups_usize(&self) -> usize {
110        debug_assert!(!self.group_offsets_host.is_empty());
111        self.group_offsets_host.len() - 1
112    }
113
114    /// Returns the total number of mapped neural output slots.
115    pub fn total_slots(&self) -> u32 {
116        debug_assert!(!self.group_offsets_host.is_empty());
117        self.group_offsets_host[self.group_offsets_host.len() - 1]
118    }
119
120    /// Returns the device-resident CSR offsets for all groups.
121    pub fn group_offsets(&self) -> &TrackedCudaSlice<u32> {
122        &self.group_offsets
123    }
124
125    /// Returns the device-resident CNF variable identifier for every slot.
126    pub fn slot_cnf_var(&self) -> &TrackedCudaSlice<u32> {
127        &self.slot_cnf_var
128    }
129
130    /// Device view over `slot_cnf_var` for a single group.
131    pub fn group_slot_cnf_var(&self, group_idx: usize) -> Result<DeviceMemoryView<u32>> {
132        let start = *self
133            .group_offsets_host
134            .get(group_idx)
135            .ok_or_else(|| XlogError::Compilation("Group index out of bounds".to_string()))?
136            as usize;
137        let end = *self
138            .group_offsets_host
139            .get(group_idx + 1)
140            .ok_or_else(|| XlogError::Compilation("Group index out of bounds".to_string()))?
141            as usize;
142        if end < start || end > self.slot_cnf_var.len() {
143            return Err(XlogError::Compilation(
144                "Invalid group slot range in GpuWeightSlots".to_string(),
145            ));
146        }
147        Ok(self.slot_cnf_var.slice(start..end))
148    }
149}
xlog_prob/neural_fast_path.rs

xlog_prob/
neural_fast_path.rs
