libmir-cuda 0.3.0

CUDA inference backend for libmir
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88

use std::{collections::HashMap, time::Duration};

use super::{CudaAutoTuner, DenseRuntimeEntry, storage::StoredDenseEntry};
use crate::{DenseExecution, DensePlanRequest, PlanSource};

impl CudaAutoTuner {
    pub(in crate::backend) fn lookup_dense(
        &self,
        request: DensePlanRequest,
    ) -> Option<(DenseExecution, PlanSource)> {
        if self.inner.config.mode == super::CudaTuningMode::Disabled {
            return None;
        }
        self.inner
            .state
            .lock()
            .ok()?
            .dense
            .get(&request)
            .map(|entry| (entry.execution, entry.source))
    }

    pub(in crate::backend) fn claim_dense(&self, request: DensePlanRequest) -> bool {
        let Ok(mut state) = self.inner.state.lock() else {
            return false;
        };
        self.inner.config.mode == super::CudaTuningMode::Startup
            && !state.sealed
            && state.budget.available()
            && !state.dense.contains_key(&request)
            && state.dense_inflight.insert(request)
    }

    pub(in crate::backend) fn record_dense(
        &self,
        request: DensePlanRequest,
        execution: DenseExecution,
        average: Duration,
        tuning_elapsed: Duration,
    ) {
        let snapshot = {
            let Ok(mut state) = self.inner.state.lock() else {
                return;
            };
            state.dense_inflight.remove(&request);
            state.budget.consume(tuning_elapsed);
            state.dense.insert(
                request,
                DenseRuntimeEntry {
                    execution,
                    source: PlanSource::MeasuredStartup,
                    average_ns: u64::try_from(average.as_nanos()).unwrap_or(u64::MAX),
                },
            );
            Self::snapshot(&state)
        };
        self.persist(snapshot);
    }

    pub(in crate::backend) fn abandon_dense(&self, request: DensePlanRequest) {
        if let Ok(mut state) = self.inner.state.lock() {
            state.dense_inflight.remove(&request);
        }
    }
}

pub(super) fn stored_entries(
    entries: &HashMap<DensePlanRequest, DenseRuntimeEntry>,
) -> Vec<StoredDenseEntry> {
    let mut stored = entries
        .iter()
        .map(|(request, entry)| StoredDenseEntry {
            request: *request,
            execution: entry.execution,
            average_ns: entry.average_ns,
        })
        .collect::<Vec<_>>();
    stored.sort_by_key(|entry| {
        (
            entry.request.phase as u8,
            entry.request.role as u8,
            entry.request.tokens,
            entry.request.input_features,
            entry.request.output_features,
        )
    });
    stored
}