libmir-cuda 0.2.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

use std::collections::HashSet;

use models::{
    layout::VisionConfig,
    weights::{TensorCatalog, TensorInfo, VisionTensorSchema},
};

use crate::{CudaBackend, CudaTensorSet, Error, Result};

pub fn load_vision_tensors(
    backend: &CudaBackend,
    config: &VisionConfig,
    catalog: &TensorCatalog,
) -> Result<CudaTensorSet> {
    let schema = VisionTensorSchema::discover(config);
    let mut seen = HashSet::with_capacity(schema.requirements.len());
    let mut tensors = Vec::with_capacity(schema.requirements.len());
    for requirement in schema.requirements {
        let tensor = requirement
            .aliases
            .iter()
            .find_map(|alias| find(catalog, alias))
            .ok_or_else(|| Error::MissingTensor(requirement.missing_label()))?;
        if seen.insert(tensor.name.as_str()) {
            tensors.push(tensor);
        }
    }
    for tensor in &catalog.tensors {
        if is_optional_vision_tensor(config, &tensor.name) && seen.insert(tensor.name.as_str()) {
            tensors.push(tensor);
        }
    }
    let mut upload = backend.begin_tensor_upload();
    for tensor in tensors {
        upload.enqueue(tensor)?;
    }
    upload.finish()
}

#[allow(clippy::case_sensitive_file_extension_comparisons)]
fn is_optional_vision_tensor(config: &VisionConfig, name: &str) -> bool {
    match config {
        VisionConfig::PooledEncoder(_) => {
            (name.starts_with("model.vision_tower.") || name.starts_with("model.embed_vision."))
                && name.ends_with(".bias")
        },
        VisionConfig::SpatialMergeEncoder(_) => false,
    }
}

fn find<'a>(catalog: &'a TensorCatalog, name: &str) -> Option<&'a TensorInfo> {
    catalog.tensors.iter().find(|tensor| tensor.name == name)
}