xlog_cuda/
semantic_work.rs

1//! Original first-recording work geometry and device-produced work. Units are
2//! logical operations, not elapsed time, instructions or FLOPs. Every recorded
3//! occurrence is retained; conditional producers supply actual units separately
4//! from their cold upper geometry.
5
6use std::collections::HashSet;
7
8#[repr(u64)]
9#[derive(Clone, Copy, Debug, PartialEq, Eq)]
10pub enum ModelWorkKind {
11    Pointwise = 1,
12    Reduction = 2,
13    Scan = 3,
14    Fill = 4,
15    Contraction = 5,
16    Softmax = 6,
17    LogSoftmax = 7,
18    LogSumExp = 8,
19    Scatter = 9,
20    Gather = 10,
21    CopyBytes = 11,
22    Cast = 12,
23    SortKeys = 13,
24    BucketQueries = 14,
25    SavedCopyBytes = 15,
26    /// An executed model-call marker, separate from logical device work.
27    ModelInvocation = 16,
28}
29
30impl ModelWorkKind {
31    pub fn from_code(code: u64) -> Option<Self> {
32        Some(match code {
33            1 => Self::Pointwise,
34            2 => Self::Reduction,
35            3 => Self::Scan,
36            4 => Self::Fill,
37            5 => Self::Contraction,
38            6 => Self::Softmax,
39            7 => Self::LogSoftmax,
40            8 => Self::LogSumExp,
41            9 => Self::Scatter,
42            10 => Self::Gather,
43            11 => Self::CopyBytes,
44            12 => Self::Cast,
45            13 => Self::SortKeys,
46            14 => Self::BucketQueries,
47            15 => Self::SavedCopyBytes,
48            16 => Self::ModelInvocation,
49            _ => return None,
50        })
51    }
52}
53
54/// Immutable kernel input. Unused dimensions are zero, while a scalar has rank
55/// zero and one unit. Saved-copy coordinates name the original content witness
56/// and tensor occurrence; ordinary operations have no saved witness.
57#[repr(C)]
58#[derive(Clone, Copy, Debug, PartialEq, Eq)]
59pub(crate) struct ModelWorkEvent {
60    pub kind: u64,
61    pub witness: u64,
62    pub tensor: u64,
63    pub rank: u64,
64    pub dimensions: [u64; 8],
65    /// Zero for exact geometry, otherwise the original retained device slot
66    /// containing [actual units, sticky overflow, producer writes].
67    pub actual: u64,
68}
69
70// SAFETY: the CUDA ABI is thirteen padding-free u64 words, with no host pointers.
71unsafe impl crate::DeviceRepr for ModelWorkEvent {}
72const _: () = assert!(std::mem::size_of::<ModelWorkEvent>() == 104);
73
74impl ModelWorkEvent {
75    pub(crate) fn operation(kind: ModelWorkKind, dimensions: &[u64]) -> Result<Self, &'static str> {
76        if matches!(
77            kind,
78            ModelWorkKind::SavedCopyBytes | ModelWorkKind::ModelInvocation
79        ) {
80            return Err(
81                "saved copies and model invocations require their original native producers",
82            );
83        }
84        Self::new(kind, u64::MAX, u64::MAX, dimensions)
85    }
86
87    pub(crate) fn device_operation(
88        kind: ModelWorkKind,
89        upper_dimensions: &[u64],
90        actual: u64,
91    ) -> Result<Self, &'static str> {
92        if actual == 0 || !actual.is_multiple_of(8) || actual > u64::MAX - 24 {
93            return Err("model work requires its complete aligned native device slot");
94        }
95        let mut event = Self::operation(kind, upper_dimensions)?;
96        event.actual = actual;
97        Ok(event)
98    }
99
100    pub(crate) fn saved(
101        witness: u64,
102        tensor: u64,
103        dimensions: &[u64],
104        element_bytes: u64,
105    ) -> Result<Self, &'static str> {
106        if witness == u64::MAX
107            || tensor == u64::MAX
108            || dimensions.len() >= 8
109            || !matches!(element_bytes, 1 | 2 | 4 | 8)
110        {
111            return Err("saved-copy work has invalid original tensor geometry");
112        }
113        let mut shape = [0; 8];
114        shape[..dimensions.len()].copy_from_slice(dimensions);
115        shape[dimensions.len()] = element_bytes;
116        Self::new(
117            ModelWorkKind::SavedCopyBytes,
118            witness,
119            tensor,
120            &shape[..dimensions.len() + 1],
121        )
122    }
123
124    pub(crate) fn invocation(actual: u64) -> Result<Self, &'static str> {
125        let event = Self {
126            kind: ModelWorkKind::ModelInvocation as u64,
127            witness: u64::MAX,
128            tensor: u64::MAX,
129            rank: 0,
130            dimensions: [0; 8],
131            actual,
132        };
133        event.units()?;
134        Ok(event)
135    }
136
137    fn new(
138        kind: ModelWorkKind,
139        witness: u64,
140        tensor: u64,
141        dimensions: &[u64],
142    ) -> Result<Self, &'static str> {
143        if dimensions.len() > 8 {
144            return Err("model work geometry exceeds eight dimensions");
145        }
146        let mut event = Self {
147            kind: kind as u64,
148            witness,
149            tensor,
150            rank: dimensions.len() as u64,
151            dimensions: [0; 8],
152            actual: 0,
153        };
154        event.dimensions[..dimensions.len()].copy_from_slice(dimensions);
155        event.units()?;
156        Ok(event)
157    }
158
159    fn units(&self) -> Result<u64, &'static str> {
160        let kind = ModelWorkKind::from_code(self.kind).ok_or("model work unit kind is unknown")?;
161        if self.rank > 8
162            || self.dimensions[self.rank as usize..]
163                .iter()
164                .any(|&extent| extent != 0)
165            || (kind == ModelWorkKind::SavedCopyBytes
166                && (self.witness == u64::MAX
167                    || self.tensor == u64::MAX
168                    || self.actual != 0
169                    || self.rank == 0
170                    || !matches!(self.dimensions[self.rank as usize - 1], 1 | 2 | 4 | 8)))
171            || (kind != ModelWorkKind::SavedCopyBytes
172                && (self.witness != u64::MAX || self.tensor != u64::MAX))
173        {
174            return Err("model work event has noncanonical geometry or occurrence coordinates");
175        }
176        if kind == ModelWorkKind::ModelInvocation {
177            if self.rank != 0
178                || self.actual == 0
179                || !self.actual.is_multiple_of(8)
180                || self.actual > u64::MAX - 24
181            {
182                return Err("model invocation requires its original scalar native device slot");
183            }
184            // A call is its own resource quantity, not another work tariff.
185            return Ok(0);
186        }
187        let dimensions = &self.dimensions[..self.rank as usize];
188        if dimensions.contains(&0) {
189            return Ok(0);
190        }
191        dimensions
192            .iter()
193            .try_fold(1u64, |units, &extent| units.checked_mul(extent))
194            .ok_or("model work geometry exceeds the exact u64 counter")
195    }
196}
197
198/// Capacity is reserved cold. Neither host vector growth nor device allocation
199/// is allowed when original producer callbacks append their first-record events.
200pub(crate) struct ModelWorkRecording {
201    capacity: usize,
202    events: Vec<ModelWorkEvent>,
203    saved: HashSet<(u64, u64)>,
204    bound: u64,
205    minimum: u64,
206    frozen: bool,
207}
208
209impl ModelWorkRecording {
210    pub(crate) fn new(capacity: usize) -> Result<Self, &'static str> {
211        if capacity == 0 {
212            return Err("model work requires a positive cold event capacity");
213        }
214        let mut events = Vec::new();
215        events
216            .try_reserve_exact(capacity)
217            .map_err(|_| "model work event reservation failed")?;
218        let mut saved = HashSet::new();
219        saved
220            .try_reserve(capacity)
221            .map_err(|_| "saved occurrence reservation failed")?;
222        Ok(Self {
223            capacity,
224            events,
225            saved,
226            bound: 0,
227            minimum: 0,
228            frozen: false,
229        })
230    }
231
232    pub(crate) fn push(&mut self, event: ModelWorkEvent) -> Result<(), &'static str> {
233        if self.frozen {
234            return Err("original model work is already frozen");
235        }
236        if self.events.len() == self.capacity {
237            return Err("model work exceeds its cold event capacity");
238        }
239        let bound = self
240            .bound
241            .checked_add(event.units()?)
242            .ok_or("model work sum exceeds the exact u64 counter")?;
243        let saved = event.kind == ModelWorkKind::SavedCopyBytes as u64;
244        if saved && self.saved.contains(&(event.witness, event.tensor)) {
245            return Err("saved-tensor occurrence was already charged");
246        }
247        self.events.push(event);
248        if saved {
249            self.saved.insert((event.witness, event.tensor));
250        }
251        if event.actual == 0 {
252            // This exact contribution is a subset of the checked total bound.
253            self.minimum += bound - self.bound;
254        }
255        self.bound = bound;
256        Ok(())
257    }
258
259    #[cfg(any(test, feature = "semantic-policy"))]
260    pub(crate) fn freeze(&mut self) -> Result<u64, &'static str> {
261        if self.frozen || self.events.is_empty() || self.bound == 0 {
262            return Err("model work requires one nonempty original recording and a single freeze");
263        }
264        self.frozen = true;
265        Ok(self.bound)
266    }
267
268    #[cfg(feature = "semantic-policy")]
269    pub(crate) fn require_model_invocations(&self) -> Result<(), &'static str> {
270        if !self
271            .events
272            .iter()
273            .any(|event| event.kind == ModelWorkKind::ModelInvocation as u64)
274        {
275            return Err("model work lacks its original executed-call marker; zero is not an unobserved call count");
276        }
277        Ok(())
278    }
279
280    pub(crate) fn events(&self) -> &[ModelWorkEvent] {
281        &self.events
282    }
283    pub(crate) fn frozen_bound(&self) -> Option<u64> {
284        self.frozen.then_some(self.bound)
285    }
286}
287
288/// Checked device accumulator. The model and native contributions are kept
289/// separate; neither the old edit facts nor an estimated capacity is a charge.
290#[repr(C)]
291#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
292pub(crate) struct ExecutionWork {
293    pub raw: u64,
294    pub model_once: u64,
295    pub native_attempt: u64,
296    pub model_events: u64,
297    pub model_bound: u64,
298    pub overflow: u64,
299    /// Fixed native event order follows the producer's unit-tariff registry.
300    pub native_events: [u64; 9],
301    /// Attribution of existing native units; never added a second time to raw.
302    pub candidates: [u64; 3],
303    /// Zero is shared work; one through three select an original candidate.
304    pub active_candidate: u64,
305}
306
307impl ExecutionWork {
308    #[cfg(any(test, feature = "semantic-policy"))]
309    pub(crate) fn validate_model(
310        &self,
311        recording: &ModelWorkRecording,
312        refused_overflow: bool,
313    ) -> bool {
314        recording.frozen_bound() == Some(self.model_bound)
315            && self.model_once.checked_add(self.native_attempt) == Some(self.raw)
316            && self
317                .native_events
318                .iter()
319                .try_fold(0u64, |sum, value| sum.checked_add(*value))
320                == Some(self.native_attempt)
321            && self
322                .candidates
323                .iter()
324                .try_fold(0u64, |sum, value| sum.checked_add(*value))
325                .is_some_and(|sum| sum <= self.native_attempt)
326            && self.active_candidate == 0
327            && if refused_overflow {
328                self.overflow == 1 && self.model_events <= recording.events().len() as u64
329            } else {
330                self.overflow == 0
331                    && self.model_events == recording.events().len() as u64
332                    && (recording.minimum..=self.model_bound).contains(&self.model_once)
333            }
334    }
335}
336
337const _: () = assert!(std::mem::size_of::<ExecutionWork>() == 152);
338
339#[cfg(test)]
340mod tests {
341    use super::*;
342
343    #[test]
344    fn event_geometry_counts_declared_units_without_dense_mask_discount() {
345        let event = ModelWorkEvent::operation(ModelWorkKind::Contraction, &[2, 3, 5]).unwrap();
346        assert_eq!(event.units().unwrap(), 30);
347        assert_eq!(
348            ModelWorkEvent::operation(ModelWorkKind::Pointwise, &[])
349                .unwrap()
350                .units()
351                .unwrap(),
352            1
353        );
354        assert_eq!(
355            ModelWorkEvent::operation(ModelWorkKind::Gather, &[4, 0, 7])
356                .unwrap()
357                .units()
358                .unwrap(),
359            0
360        );
361        assert!(ModelWorkEvent::operation(ModelWorkKind::Contraction, &[u64::MAX, 2]).is_err());
362        assert!(ModelWorkEvent::operation(ModelWorkKind::Pointwise, &[1; 9]).is_err());
363    }
364
365    #[test]
366    fn saved_occurrences_keep_original_witness_coordinates_and_byte_extents() {
367        let mut recording = ModelWorkRecording::new(3).unwrap();
368        recording
369            .push(ModelWorkEvent::saved(7, 0, &[2, 3], 4).unwrap())
370            .unwrap();
371        recording
372            .push(ModelWorkEvent::saved(8, 0, &[2, 3], 4).unwrap())
373            .unwrap();
374        assert!(recording
375            .push(ModelWorkEvent::saved(7, 0, &[2, 3], 4).unwrap())
376            .is_err());
377        assert_eq!(recording.freeze().unwrap(), 48);
378        assert_eq!(recording.events()[0].witness, 7);
379        assert_eq!(recording.events()[1].witness, 8);
380    }
381
382    #[test]
383    fn recording_freezes_once_and_rejects_capacity_and_sum_overflow() {
384        let mut recording = ModelWorkRecording::new(1).unwrap();
385        let event = ModelWorkEvent::operation(ModelWorkKind::CopyBytes, &[12]).unwrap();
386        recording.push(event).unwrap();
387        assert!(recording.push(event).is_err());
388        assert_eq!(recording.freeze().unwrap(), 12);
389        assert!(recording.freeze().is_err());
390        assert!(recording.push(event).is_err());
391        let mut overflow = ModelWorkRecording::new(2).unwrap();
392        overflow
393            .push(ModelWorkEvent::operation(ModelWorkKind::CopyBytes, &[u64::MAX]).unwrap())
394            .unwrap();
395        assert!(overflow
396            .push(ModelWorkEvent::operation(ModelWorkKind::Fill, &[1]).unwrap())
397            .is_err());
398        assert!(ModelWorkRecording::new(0).is_err());
399        assert!(ModelWorkRecording::new(1).unwrap().freeze().is_err());
400    }
401
402    #[test]
403    fn completion_consumes_the_original_frozen_geometry_and_checked_common_sum() {
404        let mut recording = ModelWorkRecording::new(1).unwrap();
405        recording
406            .push(ModelWorkEvent::operation(ModelWorkKind::Contraction, &[3, 7]).unwrap())
407            .unwrap();
408        let mut actual = ExecutionWork {
409            model_once: 21,
410            native_attempt: 5,
411            raw: 26,
412            model_events: 1,
413            model_bound: 21,
414            native_events: [5, 0, 0, 0, 0, 0, 0, 0, 0],
415            ..ExecutionWork::default()
416        };
417        assert!(!actual.validate_model(&recording, false));
418        recording.freeze().unwrap();
419        assert!(actual.validate_model(&recording, false));
420        actual.raw = 21;
421        assert!(!actual.validate_model(&recording, false));
422        actual.raw = 26;
423        actual.model_events = 0;
424        assert!(!actual.validate_model(&recording, false));
425        actual.overflow = 1;
426        assert!(actual.validate_model(&recording, true));
427        actual.model_bound = 22;
428        assert!(!actual.validate_model(&recording, true));
429    }
430
431    #[test]
432    fn native_event_sum_and_candidate_attribution_are_not_additional_work() {
433        let mut recording = ModelWorkRecording::new(1).unwrap();
434        recording
435            .push(ModelWorkEvent::operation(ModelWorkKind::Fill, &[7]).unwrap())
436            .unwrap();
437        recording.freeze().unwrap();
438        let mut actual = ExecutionWork {
439            raw: 17,
440            model_once: 7,
441            model_events: 1,
442            model_bound: 7,
443            native_attempt: 10,
444            native_events: [2, 3, 0, 0, 0, 0, 0, 1, 4],
445            candidates: [1, 3, 4],
446            ..ExecutionWork::default()
447        };
448        assert!(actual.validate_model(&recording, false));
449        actual.native_events[1] += 1;
450        assert!(!actual.validate_model(&recording, false));
451        actual.native_events[1] -= 1;
452        actual.candidates[0] = 4;
453        assert!(!actual.validate_model(&recording, false));
454        actual.candidates[0] = 1;
455        actual.active_candidate = 2;
456        assert!(!actual.validate_model(&recording, false));
457    }
458}
xlog_cuda/semantic_work.rs

xlog_cuda/
semantic_work.rs
