vkml 0.0.2

High-level Vulkan-based machine learning library
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
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143

use onnx_extractor::DataType;

use crate::{instruction, tensor::TensorDesc, utils::error::VKMLError};

use super::{Layer, execution::LayerExecution};

pub struct ConcatLayer {
    pub dim: usize,
}

impl ConcatLayer {
    pub fn new(dim: usize) -> Self {
        Self { dim }
    }
}

impl Layer for ConcatLayer {
    fn name(&self) -> String {
        "Concat".to_string()
    }

    fn config_string(&self) -> Option<String> {
        Some(format!("dim={}", self.dim))
    }

    fn requires_parameters(&self) -> bool {
        false
    }

    fn input_requirements(&self) -> (usize, Option<usize>) {
        (2, None) // At least 2 inputs, no maximum
    }

    fn parameter_shapes(&self, _input_shapes: &[&TensorDesc]) -> Option<(TensorDesc, TensorDesc)> {
        None // No parameters
    }

    fn output_shapes(
        &self,
        _batch_size: i64,
        input_shapes: &[&TensorDesc],
    ) -> Result<Vec<TensorDesc>, VKMLError> {
        if input_shapes.len() < 2 {
            return Err(VKMLError::Layer(format!(
                "Concat layer requires at least 2 inputs, got {}",
                input_shapes.len()
            )));
        }

        // Check that all inputs have the same number of dimensions
        let first_dims = input_shapes[0].dims();
        let ndim = first_dims.len();
        for shape in input_shapes.iter().skip(1) {
            if shape.dims().len() != ndim {
                return Err(VKMLError::Layer(
                    "All inputs to Concat must have same number of dimensions".to_string(),
                ));
            }
        }

        // Check that concatenation dimension is valid
        if self.dim >= ndim {
            return Err(VKMLError::Layer(format!(
                "Concat dimension {} out of range for {}-dimensional tensors",
                self.dim, ndim
            )));
        }

        // For all dimensions except concat_dim, sizes must match
        for (d, &size) in first_dims.iter().enumerate().take(ndim) {
            if d == self.dim {
                continue;
            }

            for shape in input_shapes.iter().skip(1) {
                if shape.dims()[d] != size {
                    return Err(VKMLError::Layer(format!(
                        "Dimension {} must have same size for all inputs to Concat",
                        d
                    )));
                }
            }
        }

        // Calculate the output shape
        let mut output_dims = first_dims.to_vec();

        // Sum the sizes along the concat dimension
        output_dims[self.dim] = input_shapes
            .iter()
            .map(|shape| shape.dims()[self.dim])
            .sum();

        // Create output tensor descriptor of the appropriate type
        let output_shape = TensorDesc::new(output_dims, DataType::Float);

        Ok(vec![output_shape])
    }

    fn build_layer_exec(
        &self,
        batch_size: i64,
        input_shapes: &[&TensorDesc],
    ) -> Result<LayerExecution, VKMLError> {
        if input_shapes.len() < 2 {
            return Err(VKMLError::Layer(format!(
                "Concat layer requires at least 2 inputs, got {}",
                input_shapes.len()
            )));
        }

        let mut tensors = Vec::new();
        let mut input_tensor_indices = Vec::new();

        // Add input tensors
        for shape in input_shapes {
            let idx = tensors.len();
            tensors.push((*shape).clone());
            input_tensor_indices.push(idx);
        }

        // Calculate output shape
        let output_shapes = self.output_shapes(batch_size, input_shapes)?;
        let output_shape = output_shapes[0].clone();

        // Add output tensor
        let output_idx = tensors.len();
        tensors.push(output_shape);

        // Create Concat instruction
        let instruction = instruction::concat(input_tensor_indices, output_idx, self.dim);

        // Get input mappings using the trait method
        let input_mappings = self.map_input_tensors(input_shapes.len());

        Ok(LayerExecution {
            tensors,
            instructions: vec![instruction],
            outputs: vec![output_idx],
            input_mappings,
        })
    }
}