mumu-gpu 0.1.0

GPU/Vulkan matrix and tensor operations for the mumu/lava language
Documentation 
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

#![allow(unused_imports, dead_code, unused_variables)]
use mumu::parser::types::{Value, TensorValue, ElemType};
use mumu::parser::interpreter::Interpreter;

pub fn gpu_to_tensor_bridge(
    _interp: &mut Interpreter,
    args: Vec<Value>
) -> Result<Value, String> {
    if args.len() != 1 {
        return Err(format!("gpu:to_tensor => expected 1 argument, got {}", args.len()));
    }
    let array_val = &args[0];

    let (rows, cols, float_bytes) = match array_val {
        Value::Int2DArray(mat) => {
            if mat.is_empty() {
                return Err("gpu:to_tensor => got empty 2D int array".to_string());
            }
            let r_count = mat.len();
            let c_count = mat[0].len();
            for row_i in mat {
                if row_i.len() != c_count {
                    return Err("gpu:to_tensor => inconsistent row lengths".to_string());
                }
            }
            let mut bytes = Vec::new();
            for row in mat {
                for &iv in row {
                    let fv = iv as f32;
                    bytes.extend_from_slice(&fv.to_le_bytes());
                }
            }
            (r_count, c_count, bytes)
        }
        Value::Float2DArray(mat) => {
            if mat.is_empty() {
                return Err("gpu:to_tensor => got empty 2D float array".to_string());
            }
            let r_count = mat.len();
            let c_count = mat[0].len();
            for row_i in mat {
                if row_i.len() != c_count {
                    return Err("gpu:to_tensor => inconsistent row lengths".to_string());
                }
            }
            let mut bytes = Vec::new();
            for row in mat {
                for &fv in row {
                    let f32v = fv as f32;
                    bytes.extend_from_slice(&f32v.to_le_bytes());
                }
            }
            (r_count, c_count, bytes)
        }
        other => {
            return Err(format!(
                "gpu:to_tensor => argument must be Int2DArray or Float2DArray, got {:?}",
                other
            ));
        }
    };

    let shape = vec![rows, cols];
    let tv = TensorValue {
        elem_type: ElemType::Float32,
        shape,
        data: float_bytes,
    };

    #[cfg(debug_assertions)]
    {
        crate::debug::set_last_call_info("to_tensor", false);
    }

    Ok(Value::Tensor(tv))
}

pub fn gpu_to_array_bridge(
    _interp: &mut Interpreter,
    args: Vec<Value>
) -> Result<Value, String> {
    if args.len() != 1 {
        return Err(format!("gpu:to_array => expected 1 argument, got {}", args.len()));
    }
    let v = &args[0];
    let t = match v {
        Value::Tensor(tv) => tv,
        other => {
            return Err(format!("gpu:to_array => argument must be a Tensor, got {:?}", other));
        }
    };
    if t.shape.len() != 2 {
        return Err("gpu:to_array => only rank=2 supported".to_string());
    }
    if t.elem_type != ElemType::Float32 {
        return Err("gpu:to_array => tensor must be Float32".to_string());
    }

    let rows = t.shape[0];
    let cols = t.shape[1];
    let needed = rows * cols * 4;
    if t.data.len() != needed {
        return Err("gpu:to_array => data length mismatch".to_string());
    }
    let mut offset = 0;
    let mut out = Vec::with_capacity(rows);
    for _r in 0..rows {
        let mut rowf = Vec::with_capacity(cols);
        for _c in 0..cols {
            let chunk = &t.data[offset..offset+4];
            offset += 4;
            let val = f32::from_le_bytes(chunk.try_into().unwrap());
            rowf.push(val as f64);
        }
        out.push(rowf);
    }

    #[cfg(debug_assertions)]
    {
        crate::debug::set_last_call_info("to_array", false);
    }

    Ok(Value::Float2DArray(out))
}