xlog_prob/compilation/
gpu_weights.rs

1//! GPU-native weight table builders for exact inference.
2
3use std::sync::Arc;
4
5use cudarc::driver::LaunchConfig;
6use xlog_core::{Result, XlogError};
7use xlog_cuda::memory::TrackedCudaSlice;
8use xlog_cuda::provider::{weights_kernels, WEIGHTS_MODULE};
9use xlog_cuda::{CudaKernelProvider, LaunchAsync};
10
11use crate::compilation::gpu_cnf::GpuCnfVarTables;
12
13/// Device-resident log weights indexed by DIMACS variable.
14pub struct GpuWeights {
15    /// Log weight for assigning each variable true.
16    pub log_true: TrackedCudaSlice<f64>,
17    /// Log weight for assigning each variable false.
18    pub log_false: TrackedCudaSlice<f64>,
19}
20
21fn kernel_count_u32(context: &str, count: usize) -> Result<u32> {
22    u32::try_from(count)
23        .map_err(|_| XlogError::Compilation(format!("{context} exceeds GPU u32 index space")))
24}
25
26fn grid_for(count: u32, block: u32) -> Result<u32> {
27    if count == 0 {
28        return Ok(0);
29    }
30    if block == 0 {
31        return Err(XlogError::Compilation(
32            "GPU weight kernel block size must be nonzero".to_string(),
33        ));
34    }
35    let grid = (count as u64).div_ceil(block as u64);
36    let step = grid
37        .checked_mul(block as u64)
38        .ok_or_else(|| XlogError::Compilation("GPU weight grid-stride overflow".to_string()))?;
39    if step > u32::MAX as u64 {
40        return Err(XlogError::Compilation(
41            "GPU weight grid-stride step exceeds u32 index space".to_string(),
42        ));
43    }
44    u32::try_from(grid).map_err(|_| {
45        XlogError::Compilation("GPU weight kernel grid exceeds u32 index space".to_string())
46    })
47}
48
49fn checked_var_table_count(var_cap: u32) -> Result<u32> {
50    var_cap.checked_add(1).ok_or_else(|| {
51        XlogError::Compilation("GPU weight var_cap exceeds u32 table index space".to_string())
52    })
53}
54
55fn weights_len_for_var_cap(var_cap: u32) -> Result<usize> {
56    (var_cap as usize)
57        .checked_add(1)
58        .ok_or_else(|| XlogError::Compilation("weight table size overflow".to_string()))
59}
60
61fn query_weights_len_for_var_cap(var_cap: u32) -> Result<usize> {
62    (var_cap as usize)
63        .checked_add(1)
64        .ok_or_else(|| XlogError::Compilation("query var_cap overflow".to_string()))
65}
66
67fn evidence_len_for_var_cap(var_cap: u32) -> Result<usize> {
68    (var_cap as usize)
69        .checked_add(1)
70        .ok_or_else(|| XlogError::Compilation("evidence var_cap overflow".to_string()))
71}
72
73/// Builds a device table that maps DIMACS variables to forced evidence values.
74pub fn build_evidence_by_var_gpu(
75    node_var: &TrackedCudaSlice<u32>,
76    evidence_nodes: &TrackedCudaSlice<u32>,
77    evidence_vals: &TrackedCudaSlice<u8>,
78    var_cap: u32,
79    provider: &Arc<CudaKernelProvider>,
80) -> Result<TrackedCudaSlice<u8>> {
81    if evidence_nodes.len() != evidence_vals.len() {
82        return Err(XlogError::Compilation(format!(
83            "GPU evidence nodes len {} != vals len {}",
84            evidence_nodes.len(),
85            evidence_vals.len()
86        )));
87    }
88    let len = evidence_len_for_var_cap(var_cap)?;
89
90    let memory = provider.memory();
91    let device = provider.device().inner();
92    let mut evidence_by_var = memory.alloc::<u8>(len)?;
93    device
94        .memset_zeros(&mut evidence_by_var)
95        .map_err(|e| XlogError::Kernel(format!("Failed to zero evidence buffer: {}", e)))?;
96
97    let count = evidence_nodes.len();
98    if count == 0 {
99        return Ok(evidence_by_var);
100    }
101    let count_u32 = kernel_count_u32("GPU evidence node count", count)?;
102
103    let func = device
104        .get_func(
105            WEIGHTS_MODULE,
106            weights_kernels::WEIGHTS_SET_EVIDENCE_FROM_NODES,
107        )
108        .ok_or_else(|| {
109            XlogError::Kernel("weights_set_evidence_from_nodes kernel not found".to_string())
110        })?;
111    let block = 256u32;
112    let grid = grid_for(count_u32, block)?;
113    // SAFETY: kernel arguments match the PTX signature; device buffers were allocated with sufficient size
114    unsafe {
115        func.clone().launch(
116            LaunchConfig {
117                grid_dim: (grid.max(1), 1, 1),
118                block_dim: (block, 1, 1),
119                shared_mem_bytes: 0,
120            },
121            (
122                node_var,
123                evidence_nodes,
124                evidence_vals,
125                count_u32,
126                var_cap,
127                &mut evidence_by_var,
128            ),
129        )
130    }
131    .map_err(|e| XlogError::Kernel(format!("weights_set_evidence_from_nodes failed: {}", e)))?;
132    // No device synchronize: returns device-resident slice; same-stream ordering suffices.
133    Ok(evidence_by_var)
134}
135
136/// Maps device-resident PIR node identifiers to their DIMACS variables.
137pub fn map_nodes_to_vars_gpu(
138    node_var: &TrackedCudaSlice<u32>,
139    node_ids: &TrackedCudaSlice<u32>,
140    var_cap: u32,
141    provider: &Arc<CudaKernelProvider>,
142) -> Result<TrackedCudaSlice<u32>> {
143    let memory = provider.memory();
144    let device = provider.device().inner();
145    let mut out = memory.alloc::<u32>(node_ids.len())?;
146    let count = node_ids.len();
147    if count == 0 {
148        return Ok(out);
149    }
150    let count_u32 = kernel_count_u32("GPU node-to-var map count", count)?;
151
152    let func = device
153        .get_func(WEIGHTS_MODULE, weights_kernels::WEIGHTS_MAP_NODES_TO_VARS)
154        .ok_or_else(|| {
155            XlogError::Kernel("weights_map_nodes_to_vars kernel not found".to_string())
156        })?;
157
158    let block = 256u32;
159    let grid = grid_for(count_u32, block)?;
160    // SAFETY: kernel arguments match the PTX signature; device buffers were allocated with sufficient size
161    unsafe {
162        func.clone().launch(
163            LaunchConfig {
164                grid_dim: (grid.max(1), 1, 1),
165                block_dim: (block, 1, 1),
166                shared_mem_bytes: 0,
167            },
168            (node_var, node_ids, count_u32, var_cap, &mut out),
169        )
170    }
171    .map_err(|e| XlogError::Kernel(format!("weights_map_nodes_to_vars failed: {}", e)))?;
172    // No device synchronize: returns device-resident slice; same-stream ordering suffices.
173    Ok(out)
174}
175
176/// Temporarily forces query variables true by replacing their false log weights.
177pub fn apply_query_vars_device(
178    provider: &Arc<CudaKernelProvider>,
179    query_vars: &TrackedCudaSlice<u32>,
180    var_cap: u32,
181    log_false: &mut TrackedCudaSlice<f64>,
182    saved: &mut TrackedCudaSlice<f64>,
183) -> Result<()> {
184    let count = query_vars.len();
185    if saved.len() < count {
186        return Err(XlogError::Compilation(format!(
187            "query restore buffer len {} < query vars len {}",
188            saved.len(),
189            count
190        )));
191    }
192    let weights_len = query_weights_len_for_var_cap(var_cap)?;
193    if log_false.len() < weights_len {
194        return Err(XlogError::Compilation(format!(
195            "log_false len {} < var_cap+1 {}",
196            log_false.len(),
197            weights_len
198        )));
199    }
200    if count == 0 {
201        return Ok(());
202    }
203    let count_u32 = kernel_count_u32("GPU query apply count", count)?;
204
205    let device = provider.device().inner();
206    let func = device
207        .get_func(WEIGHTS_MODULE, weights_kernels::WEIGHTS_APPLY_QUERY_VARS)
208        .ok_or_else(|| {
209            XlogError::Kernel("weights_apply_query_vars kernel not found".to_string())
210        })?;
211
212    let block = 256u32;
213    let grid = grid_for(count_u32, block)?;
214    // SAFETY: kernel arguments match the PTX signature; device buffers were allocated with sufficient size
215    unsafe {
216        func.clone().launch(
217            LaunchConfig {
218                grid_dim: (grid.max(1), 1, 1),
219                block_dim: (block, 1, 1),
220                shared_mem_bytes: 0,
221            },
222            (query_vars, count_u32, var_cap, log_false, saved),
223        )
224    }
225    .map_err(|e| XlogError::Kernel(format!("weights_apply_query_vars failed: {}", e)))?;
226    // No device synchronize: same-stream ordering guarantees visibility to subsequent kernels.
227    Ok(())
228}
229
230/// Restores false log weights saved by [`apply_query_vars_device`].
231pub fn restore_query_vars_device(
232    provider: &Arc<CudaKernelProvider>,
233    query_vars: &TrackedCudaSlice<u32>,
234    var_cap: u32,
235    log_false: &mut TrackedCudaSlice<f64>,
236    saved: &TrackedCudaSlice<f64>,
237) -> Result<()> {
238    let count = query_vars.len();
239    if saved.len() < count {
240        return Err(XlogError::Compilation(format!(
241            "query restore buffer len {} < query vars len {}",
242            saved.len(),
243            count
244        )));
245    }
246    let weights_len = query_weights_len_for_var_cap(var_cap)?;
247    if log_false.len() < weights_len {
248        return Err(XlogError::Compilation(format!(
249            "log_false len {} < var_cap+1 {}",
250            log_false.len(),
251            weights_len
252        )));
253    }
254    if count == 0 {
255        return Ok(());
256    }
257    let count_u32 = kernel_count_u32("GPU query restore count", count)?;
258
259    let device = provider.device().inner();
260    let func = device
261        .get_func(WEIGHTS_MODULE, weights_kernels::WEIGHTS_RESTORE_QUERY_VARS)
262        .ok_or_else(|| {
263            XlogError::Kernel("weights_restore_query_vars kernel not found".to_string())
264        })?;
265
266    let block = 256u32;
267    let grid = grid_for(count_u32, block)?;
268    // SAFETY: kernel arguments match the PTX signature; device buffers were allocated with sufficient size
269    unsafe {
270        func.clone().launch(
271            LaunchConfig {
272                grid_dim: (grid.max(1), 1, 1),
273                block_dim: (block, 1, 1),
274                shared_mem_bytes: 0,
275            },
276            (query_vars, count_u32, var_cap, log_false, saved),
277        )
278    }
279    .map_err(|e| XlogError::Kernel(format!("weights_restore_query_vars failed: {}", e)))?;
280    // No device synchronize: same-stream ordering guarantees visibility to subsequent kernels.
281    Ok(())
282}
283
284/// Builds validated device log-weight tables from probabilities and evidence.
285pub fn build_weights_gpu(
286    vars: &GpuCnfVarTables,
287    leaf_probs: &TrackedCudaSlice<f64>,
288    choice_true: &TrackedCudaSlice<f64>,
289    choice_false: &TrackedCudaSlice<f64>,
290    evidence_by_var: &TrackedCudaSlice<u8>,
291    provider: &Arc<CudaKernelProvider>,
292) -> Result<GpuWeights> {
293    let var_cap = vars.max_var;
294    let weights_len = weights_len_for_var_cap(var_cap)?;
295
296    if vars.leaf_var.len() < leaf_probs.len() {
297        return Err(XlogError::Compilation(format!(
298            "leaf_probs len {} exceeds leaf_var len {}",
299            leaf_probs.len(),
300            vars.leaf_var.len()
301        )));
302    }
303    if vars.choice_var.len() < choice_true.len() {
304        return Err(XlogError::Compilation(format!(
305            "choice_true len {} exceeds choice_var len {}",
306            choice_true.len(),
307            vars.choice_var.len()
308        )));
309    }
310    if choice_true.len() != choice_false.len() {
311        return Err(XlogError::Compilation(format!(
312            "choice_true len {} != choice_false len {}",
313            choice_true.len(),
314            choice_false.len()
315        )));
316    }
317    if evidence_by_var.len() != weights_len {
318        return Err(XlogError::Compilation(format!(
319            "evidence_by_var len {} != weights len {}",
320            evidence_by_var.len(),
321            weights_len
322        )));
323    }
324
325    let memory = provider.memory();
326    let device = provider.device().inner();
327    let mut log_true = memory.alloc::<f64>(weights_len)?;
328    let mut log_false = memory.alloc::<f64>(weights_len)?;
329
330    // Initialize to 0.0
331    device
332        .memset_zeros(&mut log_true)
333        .map_err(|e| XlogError::Kernel(format!("Failed to zero log_true weights: {}", e)))?;
334    device
335        .memset_zeros(&mut log_false)
336        .map_err(|e| XlogError::Kernel(format!("Failed to zero log_false weights: {}", e)))?;
337
338    let block = 256u32;
339
340    if !leaf_probs.is_empty() {
341        let leaf_count = kernel_count_u32("GPU leaf probability count", leaf_probs.len())?;
342        let func = device
343            .get_func(WEIGHTS_MODULE, weights_kernels::WEIGHTS_FILL_LEAF)
344            .ok_or_else(|| XlogError::Kernel("weights_fill_leaf kernel not found".to_string()))?;
345        let grid = grid_for(leaf_count, block)?;
346        // SAFETY: kernel arguments match the PTX signature; device buffers were allocated with sufficient size
347        unsafe {
348            func.clone().launch(
349                LaunchConfig {
350                    grid_dim: (grid.max(1), 1, 1),
351                    block_dim: (block, 1, 1),
352                    shared_mem_bytes: 0,
353                },
354                (
355                    &vars.leaf_var,
356                    leaf_probs,
357                    leaf_count,
358                    var_cap,
359                    &mut log_true,
360                    &mut log_false,
361                ),
362            )
363        }
364        .map_err(|e| XlogError::Kernel(format!("weights_fill_leaf failed: {}", e)))?;
365    }
366
367    if !choice_true.is_empty() {
368        let choice_count = kernel_count_u32("GPU choice probability count", choice_true.len())?;
369        let func = device
370            .get_func(WEIGHTS_MODULE, weights_kernels::WEIGHTS_FILL_CHOICE)
371            .ok_or_else(|| XlogError::Kernel("weights_fill_choice kernel not found".to_string()))?;
372        let grid = grid_for(choice_count, block)?;
373        // SAFETY: kernel arguments match the PTX signature; device buffers were allocated with sufficient size
374        unsafe {
375            func.clone().launch(
376                LaunchConfig {
377                    grid_dim: (grid.max(1), 1, 1),
378                    block_dim: (block, 1, 1),
379                    shared_mem_bytes: 0,
380                },
381                (
382                    &vars.choice_var,
383                    choice_true,
384                    choice_false,
385                    choice_count,
386                    var_cap,
387                    &mut log_true,
388                    &mut log_false,
389                ),
390            )
391        }
392        .map_err(|e| XlogError::Kernel(format!("weights_fill_choice failed: {}", e)))?;
393    }
394
395    if !evidence_by_var.is_empty() {
396        let var_table_count = checked_var_table_count(var_cap)?;
397        let func = device
398            .get_func(WEIGHTS_MODULE, weights_kernels::WEIGHTS_APPLY_EVIDENCE)
399            .ok_or_else(|| {
400                XlogError::Kernel("weights_apply_evidence kernel not found".to_string())
401            })?;
402        let grid = grid_for(var_table_count, block)?;
403        // SAFETY: kernel arguments match the PTX signature; device buffers were allocated with sufficient size
404        unsafe {
405            func.clone().launch(
406                LaunchConfig {
407                    grid_dim: (grid.max(1), 1, 1),
408                    block_dim: (block, 1, 1),
409                    shared_mem_bytes: 0,
410                },
411                (evidence_by_var, var_cap, &mut log_true, &mut log_false),
412            )
413        }
414        .map_err(|e| XlogError::Kernel(format!("weights_apply_evidence failed: {}", e)))?;
415    }
416    // No device synchronize: returns device-resident weights; same-stream ordering suffices.
417    Ok(GpuWeights {
418        log_true,
419        log_false,
420    })
421}
xlog_prob/compilation/gpu_weights.rs

xlog_prob/compilation/
gpu_weights.rs
