Skip to main content

kime_cuda/
compat.rs

1//! The Laya compat graph on one GPU.
2
3use kime_model::Model;
4use kime_tensor::{Buckets, Executor, HostTensor, Result};
5
6use crate::{CudaBackend, Precision};
7
8/// The compat graph of `model` on GPU `ordinal` at `precision`, with the default bucket table.
9///
10/// # Errors
11///
12/// [`kime_tensor::Error::Device`] when the GPU cannot be opened, or an error from lowering, which
13/// for a checkpoint that loaded means a bug.
14pub fn executor(
15    model: &Model,
16    ordinal: usize,
17    precision: Precision,
18) -> Result<Executor<CudaBackend>> {
19    let t = &model.tensors;
20    let host: Vec<HostTensor<'_>> = (0..t.entries().len())
21        .map(|i| {
22            let v = t.view(i);
23            HostTensor { dtype: v.dtype, shape: v.shape, bytes: v.bytes }
24        })
25        .collect();
26    let plan = model.graph.plan(&model.spec);
27    let vocab = model.spec.encoder.vocab;
28    Executor::new(
29        CudaBackend::new(ordinal, precision)?,
30        &host,
31        plan,
32        &Buckets::default(),
33        "compat",
34        vocab,
35        3,
36    )
37}