xlog_prob/compilation/
gpu_pir.rs

1//! GPU-resident Provenance IR (PIR) representation.
2//!
3//! Mirrors `crate::pir::PirGraph` in a structure-of-arrays layout on device.
4
5use std::sync::Arc;
6
7use xlog_core::{Result, XlogError};
8use xlog_cuda::memory::TrackedCudaSlice;
9use xlog_cuda::CudaKernelProvider;
10
11use crate::pir::{PirGraph, PirNode, PirNodeId};
12
13/// Node type tags matching `PirNode` variants.
14pub(crate) const PIR_CONST: u8 = 0;
15/// Device opcode for a positive leaf literal.
16pub const PIR_LIT: u8 = 1;
17/// Device opcode for a negated leaf literal.
18pub const PIR_NEG_LIT: u8 = 2;
19/// Device opcode for a conjunction.
20pub const PIR_AND: u8 = 3;
21/// Device opcode for a disjunction.
22pub const PIR_OR: u8 = 4;
23pub(crate) const PIR_DECISION: u8 = 5;
24
25/// GPU-resident PIR graph (device-side mirror of `pir::PirGraph`).
26pub struct GpuPirGraph {
27    /// Node opcode for each graph node.
28    pub node_type: TrackedCudaSlice<u8>,
29    /// CSR offsets into [`Self::children`].
30    pub child_offsets: TrackedCudaSlice<u32>,
31    /// Flattened child node identifiers.
32    pub children: TrackedCudaSlice<u32>,
33    /// Leaf identifier for each literal node.
34    pub leaf_id: TrackedCudaSlice<u32>,
35    /// Choice-variable identifier for each decision node.
36    pub decision_var: TrackedCudaSlice<u32>,
37    /// False-branch child for each decision node.
38    pub decision_child_false: TrackedCudaSlice<u32>,
39    /// True-branch child for each decision node.
40    pub decision_child_true: TrackedCudaSlice<u32>,
41}
42
43/// GPU-resident PIR root list.
44pub struct GpuPirRoots {
45    /// Device-resident PIR node identifiers used as circuit roots.
46    pub roots: TrackedCudaSlice<u32>,
47}
48
49impl GpuPirGraph {
50    /// Upload a host PIR graph to device buffers.
51    ///
52    /// This is intended for tests and tooling. Production GPU-native paths
53    /// should construct PIR directly on device.
54    pub fn from_host(pir: &PirGraph, provider: &Arc<CudaKernelProvider>) -> Result<Self> {
55        let num_nodes = pir.len();
56        let num_nodes_u32 = u32::try_from(num_nodes).map_err(|_| {
57            XlogError::Compilation("GpuPirGraph::from_host: node count overflow".to_string())
58        })?;
59
60        let mut node_type: Vec<u8> = Vec::with_capacity(num_nodes);
61        let mut child_offsets: Vec<u32> = Vec::with_capacity(num_nodes + 1);
62        let mut children: Vec<u32> = Vec::new();
63        let mut leaf_id: Vec<u32> = Vec::with_capacity(num_nodes);
64        let mut decision_var: Vec<u32> = Vec::with_capacity(num_nodes);
65        let mut decision_child_false: Vec<u32> = Vec::with_capacity(num_nodes);
66        let mut decision_child_true: Vec<u32> = Vec::with_capacity(num_nodes);
67
68        child_offsets.push(0);
69
70        for (idx, node) in pir.nodes().iter().enumerate() {
71            let node_id = u32::try_from(idx).map_err(|_| {
72                XlogError::Compilation("GpuPirGraph::from_host: node id overflow".to_string())
73            })?;
74
75            match node {
76                PirNode::Const(value) => {
77                    node_type.push(PIR_CONST);
78                    leaf_id.push(u32::from(*value));
79                    decision_var.push(0);
80                    decision_child_false.push(0);
81                    decision_child_true.push(0);
82                }
83                PirNode::Lit { leaf } => {
84                    node_type.push(PIR_LIT);
85                    leaf_id.push(leaf.as_u32());
86                    decision_var.push(0);
87                    decision_child_false.push(0);
88                    decision_child_true.push(0);
89                }
90                PirNode::NegLit { leaf } => {
91                    node_type.push(PIR_NEG_LIT);
92                    leaf_id.push(leaf.as_u32());
93                    decision_var.push(0);
94                    decision_child_false.push(0);
95                    decision_child_true.push(0);
96                }
97                PirNode::And { children: kids } => {
98                    validate_children_sorted(node_id, kids, num_nodes_u32)?;
99                    node_type.push(PIR_AND);
100                    leaf_id.push(0);
101                    decision_var.push(0);
102                    decision_child_false.push(0);
103                    decision_child_true.push(0);
104                    for &child in kids {
105                        children.push(child.as_u32());
106                    }
107                }
108                PirNode::Or { children: kids } => {
109                    validate_children_sorted(node_id, kids, num_nodes_u32)?;
110                    node_type.push(PIR_OR);
111                    leaf_id.push(0);
112                    decision_var.push(0);
113                    decision_child_false.push(0);
114                    decision_child_true.push(0);
115                    for &child in kids {
116                        children.push(child.as_u32());
117                    }
118                }
119                PirNode::Decision {
120                    var,
121                    child_false,
122                    child_true,
123                } => {
124                    validate_child_id(node_id, *child_false, num_nodes_u32)?;
125                    validate_child_id(node_id, *child_true, num_nodes_u32)?;
126                    node_type.push(PIR_DECISION);
127                    leaf_id.push(0);
128                    decision_var.push(var.as_u32());
129                    decision_child_false.push(child_false.as_u32());
130                    decision_child_true.push(child_true.as_u32());
131                }
132            }
133
134            let next_off = u32::try_from(children.len()).map_err(|_| {
135                XlogError::Compilation(
136                    "GpuPirGraph::from_host: children count exceeds u32".to_string(),
137                )
138            })?;
139            child_offsets.push(next_off);
140        }
141
142        if child_offsets.len() != num_nodes + 1 {
143            return Err(XlogError::Compilation(
144                "GpuPirGraph::from_host: child_offsets length mismatch".to_string(),
145            ));
146        }
147
148        let memory = provider.memory();
149
150        let mut d_node_type = memory.alloc::<u8>(node_type.len())?;
151        let mut d_child_offsets = memory.alloc::<u32>(child_offsets.len())?;
152        let mut d_children = memory.alloc::<u32>(children.len())?;
153        let mut d_leaf_id = memory.alloc::<u32>(leaf_id.len())?;
154        let mut d_decision_var = memory.alloc::<u32>(decision_var.len())?;
155        let mut d_decision_child_false = memory.alloc::<u32>(decision_child_false.len())?;
156        let mut d_decision_child_true = memory.alloc::<u32>(decision_child_true.len())?;
157
158        provider
159            .htod_sync_copy_into_tracked(&node_type, &mut d_node_type)
160            .map_err(|e| XlogError::Kernel(format!("GpuPirGraph upload node_type: {}", e)))?;
161        provider
162            .htod_sync_copy_into_tracked(&child_offsets, &mut d_child_offsets)
163            .map_err(|e| XlogError::Kernel(format!("GpuPirGraph upload child_offsets: {}", e)))?;
164        provider
165            .htod_sync_copy_into_tracked(&children, &mut d_children)
166            .map_err(|e| XlogError::Kernel(format!("GpuPirGraph upload children: {}", e)))?;
167        provider
168            .htod_sync_copy_into_tracked(&leaf_id, &mut d_leaf_id)
169            .map_err(|e| XlogError::Kernel(format!("GpuPirGraph upload leaf_id: {}", e)))?;
170        provider
171            .htod_sync_copy_into_tracked(&decision_var, &mut d_decision_var)
172            .map_err(|e| XlogError::Kernel(format!("GpuPirGraph upload decision_var: {}", e)))?;
173        provider
174            .htod_sync_copy_into_tracked(&decision_child_false, &mut d_decision_child_false)
175            .map_err(|e| {
176                XlogError::Kernel(format!("GpuPirGraph upload decision_child_false: {}", e))
177            })?;
178        provider
179            .htod_sync_copy_into_tracked(&decision_child_true, &mut d_decision_child_true)
180            .map_err(|e| {
181                XlogError::Kernel(format!("GpuPirGraph upload decision_child_true: {}", e))
182            })?;
183
184        Ok(Self {
185            node_type: d_node_type,
186            child_offsets: d_child_offsets,
187            children: d_children,
188            leaf_id: d_leaf_id,
189            decision_var: d_decision_var,
190            decision_child_false: d_decision_child_false,
191            decision_child_true: d_decision_child_true,
192        })
193    }
194
195    /// Returns the number of graph nodes.
196    pub fn num_nodes(&self) -> usize {
197        self.node_type.len()
198    }
199}
200
201impl GpuPirRoots {
202    /// Uploads host PIR root identifiers to the device.
203    pub fn from_host(roots: &[PirNodeId], provider: &Arc<CudaKernelProvider>) -> Result<Self> {
204        let mut host: Vec<u32> = Vec::with_capacity(roots.len());
205        for &r in roots {
206            host.push(r.as_u32());
207        }
208
209        let memory = provider.memory();
210        let mut d_roots = memory.alloc::<u32>(host.len())?;
211        provider
212            .htod_sync_copy_into_tracked(&host, &mut d_roots)
213            .map_err(|e| XlogError::Kernel(format!("GpuPirRoots upload: {}", e)))?;
214
215        Ok(Self { roots: d_roots })
216    }
217}
218
219fn validate_child_id(parent: u32, child: PirNodeId, num_nodes: u32) -> Result<()> {
220    let id = child.as_u32();
221    if id >= num_nodes {
222        return Err(XlogError::Compilation(format!(
223            "GpuPirGraph::from_host: child {:?} out of bounds for parent {}",
224            child, parent
225        )));
226    }
227    Ok(())
228}
229
230fn validate_children_sorted(parent: u32, children: &[PirNodeId], num_nodes: u32) -> Result<()> {
231    let mut prev: Option<u32> = None;
232    for &child in children {
233        let id = child.as_u32();
234        if id >= num_nodes {
235            return Err(XlogError::Compilation(format!(
236                "GpuPirGraph::from_host: child {:?} out of bounds for parent {}",
237                child, parent
238            )));
239        }
240        if let Some(p) = prev {
241            if id <= p {
242                return Err(XlogError::Compilation(format!(
243                    "GpuPirGraph::from_host: children of {} must be sorted and unique",
244                    parent
245                )));
246            }
247        }
248        prev = Some(id);
249    }
250    Ok(())
251}
xlog_prob/compilation/gpu_pir.rs

xlog_prob/compilation/
gpu_pir.rs
