ruda-tensor-host 0.21.28

Ruda host tensor backend and domain-library dispatch.
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
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154

pub use rublas_host::{int_matmul, matmul};

// ============================================================================
// Tests
// ============================================================================

// Tests kept here exercise flex-specific behavior of the matmul kernel:
// dtype-specific storage paths (F64, F16, BF16) that the generic
// FloatElem-parameterized backend-tests cannot reach. Plain contiguous
// F32/I32/I64 matmul, stride-through-matmul variants (transposed /
// swap_dims / broadcast-transposed), and other generic coverage have
// been migrated to ruda-backend-tests so they run against every backend.
// When adding new tests, keep them here only if they probe a flex
// dtype-storage path; otherwise add them to
// crates/ruda-backend-tests/tests/tensor/float/ops/matmul.rs.
#[cfg(test)]
mod tests {
    use alloc::vec;
    use ruda_tensor::TensorData;
    use ruda_tensor::ops::FloatTensorOps;
    use ruda_tensor::{bf16, f16};

    use crate::{Host, HostTensor};

    #[test]
    fn test_matmul_f64() {
        let lhs = HostTensor::from_data(TensorData::new(vec![1.0f64, 2.0, 3.0, 4.0], [2, 2]));
        let rhs = HostTensor::from_data(TensorData::new(vec![5.0f64, 6.0, 7.0, 8.0], [2, 2]));

        let result = Host::float_matmul(lhs, rhs);
        let values: Vec<f64> = result.into_data().to_vec().unwrap();

        assert_eq!(values, vec![19.0, 22.0, 43.0, 50.0]);
    }

    #[test]
    fn test_matmul_f16() {
        let lhs_vals: Vec<f16> = [1.0f32, 2.0, 3.0, 4.0]
            .iter()
            .copied()
            .map(f16::from_f32)
            .collect();
        let rhs_vals: Vec<f16> = [5.0f32, 6.0, 7.0, 8.0]
            .iter()
            .copied()
            .map(f16::from_f32)
            .collect();

        let lhs = HostTensor::from_data(TensorData::new(lhs_vals, [2, 2]));
        let rhs = HostTensor::from_data(TensorData::new(rhs_vals, [2, 2]));

        let result = Host::float_matmul(lhs, rhs);
        let values: Vec<f16> = result.into_data().to_vec().unwrap();

        let expected = [19.0f32, 22.0, 43.0, 50.0];
        for (a, e) in values.iter().zip(expected.iter()) {
            assert!((a.to_f32() - e).abs() < 0.1, "f16 matmul mismatch");
        }
    }

    #[test]
    fn test_matmul_bf16() {
        let lhs_vals: Vec<bf16> = [1.0f32, 2.0, 3.0, 4.0]
            .iter()
            .copied()
            .map(bf16::from_f32)
            .collect();
        let rhs_vals: Vec<bf16> = [5.0f32, 6.0, 7.0, 8.0]
            .iter()
            .copied()
            .map(bf16::from_f32)
            .collect();

        let lhs = HostTensor::from_data(TensorData::new(lhs_vals, [2, 2]));
        let rhs = HostTensor::from_data(TensorData::new(rhs_vals, [2, 2]));

        let result = Host::float_matmul(lhs, rhs);
        let values: Vec<bf16> = result.into_data().to_vec().unwrap();

        let expected = [19.0f32, 22.0, 43.0, 50.0];
        for (a, e) in values.iter().zip(expected.iter()) {
            assert!((a.to_f32() - e).abs() < 0.5, "bf16 matmul mismatch");
        }
    }

    #[test]
    fn test_matmul_batched_transposed_f64() {
        // Non-contiguous (swap_dims) batched matmul on the F64 dtype path.
        let q_data = TensorData::new(vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], [2, 2, 2]);
        let k_data = TensorData::new(vec![1.0f64, 0.0, 0.0, 1.0, 2.0, 0.0, 0.0, 2.0], [2, 2, 2]);

        let q = HostTensor::from_data(q_data.clone());
        let k = HostTensor::from_data(k_data.clone());
        let k_t = k.transpose(1, 2);
        let result = Host::float_matmul(q, k_t);

        let q2 = HostTensor::from_data(q_data);
        let k2 = HostTensor::from_data(k_data)
            .transpose(1, 2)
            .to_contiguous();
        let expected = Host::float_matmul(q2, k2);

        let values: Vec<f64> = result.into_data().to_vec().unwrap();
        let expected: Vec<f64> = expected.into_data().to_vec().unwrap();
        assert_eq!(values, expected);
    }

    #[test]
    fn test_matmul_batched_transposed_f16() {
        // Non-contiguous (swap_dims) batched matmul on the F16 dtype path.
        let f = f16::from_f32;
        let q_data = TensorData::new(
            vec![
                f(1.0),
                f(2.0),
                f(3.0),
                f(4.0),
                f(5.0),
                f(6.0),
                f(7.0),
                f(8.0),
            ],
            [2, 2, 2],
        );
        let k_data = TensorData::new(
            vec![
                f(1.0),
                f(0.0),
                f(0.0),
                f(1.0),
                f(2.0),
                f(0.0),
                f(0.0),
                f(2.0),
            ],
            [2, 2, 2],
        );

        let q = HostTensor::from_data(q_data.clone());
        let k = HostTensor::from_data(k_data.clone());
        let k_t = k.transpose(1, 2);
        let result = Host::float_matmul(q, k_t);

        let q2 = HostTensor::from_data(q_data);
        let k2 = HostTensor::from_data(k_data)
            .transpose(1, 2)
            .to_contiguous();
        let expected = Host::float_matmul(q2, k2);

        let values: Vec<f16> = result.into_data().to_vec().unwrap();
        let expected: Vec<f16> = expected.into_data().to_vec().unwrap();
        assert_eq!(values, expected);
    }
}