tl_cuda 0.4.0

CUDA GPU tensor library for TL
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
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171

//! Step 3: 基本演算テスト
//! binary ops, unary ops, scalar ops, comparison ops

use serial_test::serial;
use tl_cuda::{CudaTensor, DType};

fn approx_eq(a: &[f32], b: &[f32], eps: f32) {
    assert_eq!(
        a.len(),
        b.len(),
        "length mismatch: {} vs {}",
        a.len(),
        b.len()
    );
    for (i, (&x, &y)) in a.iter().zip(b.iter()).enumerate() {
        assert!(
            (x - y).abs() < eps,
            "mismatch at index {}: {} vs {} (eps={})",
            i,
            x,
            y,
            eps
        );
    }
}

// ========== 二項演算 ==========

#[test]
#[serial]
fn test_add() {
    let a = CudaTensor::from_slice(&[1.0f32, 2.0, 3.0], &[3], DType::F32);
    let b = CudaTensor::from_slice(&[4.0f32, 5.0, 6.0], &[3], DType::F32);
    let c = a.add_impl(&b).unwrap();
    assert_eq!(c.shape(), &[3]);
    approx_eq(&c.to_vec::<f32>(), &[5.0, 7.0, 9.0], 1e-6);
}

#[test]
#[serial]
fn test_sub() {
    let a = CudaTensor::from_slice(&[5.0f32, 3.0, 1.0], &[3], DType::F32);
    let b = CudaTensor::from_slice(&[1.0f32, 2.0, 3.0], &[3], DType::F32);
    let c = a.sub_impl(&b).unwrap();
    approx_eq(&c.to_vec::<f32>(), &[4.0, 1.0, -2.0], 1e-6);
}

#[test]
#[serial]
fn test_mul() {
    let a = CudaTensor::from_slice(&[2.0f32, 3.0, 4.0], &[3], DType::F32);
    let b = CudaTensor::from_slice(&[3.0f32, 2.0, 1.0], &[3], DType::F32);
    let c = a.mul_impl(&b).unwrap();
    approx_eq(&c.to_vec::<f32>(), &[6.0, 6.0, 4.0], 1e-6);
}

#[test]
#[serial]
fn test_div() {
    let a = CudaTensor::from_slice(&[6.0f32, 8.0, 10.0], &[3], DType::F32);
    let b = CudaTensor::from_slice(&[2.0f32, 4.0, 5.0], &[3], DType::F32);
    let c = a.div_impl(&b).unwrap();
    approx_eq(&c.to_vec::<f32>(), &[3.0, 2.0, 2.0], 1e-6);
}

#[test]
#[serial]
fn test_neg() {
    let a = CudaTensor::from_slice(&[1.0f32, -2.0, 3.0], &[3], DType::F32);
    let c = a.neg_impl().unwrap();
    approx_eq(&c.to_vec::<f32>(), &[-1.0, 2.0, -3.0], 1e-6);
}

#[test]
#[serial]
fn test_abs() {
    let a = CudaTensor::from_slice(&[-1.0f32, 2.0, -3.0], &[3], DType::F32);
    let c = a.abs_impl().unwrap();
    approx_eq(&c.to_vec::<f32>(), &[1.0, 2.0, 3.0], 1e-6);
}

// ========== スカラー演算 ==========

#[test]
#[serial]
fn test_add_scalar() {
    let a = CudaTensor::from_slice(&[1.0f32, 2.0, 3.0], &[3], DType::F32);
    let c = a.add_scalar_impl(10.0).unwrap();
    approx_eq(&c.to_vec::<f32>(), &[11.0, 12.0, 13.0], 1e-6);
}

#[test]
#[serial]
fn test_mul_scalar() {
    let a = CudaTensor::from_slice(&[1.0f32, 2.0, 3.0], &[3], DType::F32);
    let c = a.mul_scalar_impl(3.0).unwrap();
    approx_eq(&c.to_vec::<f32>(), &[3.0, 6.0, 9.0], 1e-6);
}

#[test]
#[serial]
fn test_div_scalar() {
    let a = CudaTensor::from_slice(&[6.0f32, 9.0, 12.0], &[3], DType::F32);
    let c = a.div_scalar_impl(3.0).unwrap();
    approx_eq(&c.to_vec::<f32>(), &[2.0, 3.0, 4.0], 1e-6);
}

#[test]
#[serial]
fn test_pow_scalar() {
    let a = CudaTensor::from_slice(&[2.0f32, 3.0, 4.0], &[3], DType::F32);
    let c = a.pow_scalar_impl(2.0).unwrap();
    approx_eq(&c.to_vec::<f32>(), &[4.0, 9.0, 16.0], 1e-6);
}

#[test]
#[serial]
fn test_clamp() {
    let a = CudaTensor::from_slice(&[-2.0f32, 0.5, 1.5, 3.0], &[4], DType::F32);
    let c = a.clamp_impl(0.0, 2.0).unwrap();
    approx_eq(&c.to_vec::<f32>(), &[0.0, 0.5, 1.5, 2.0], 1e-6);
}

// ========== broadcast ==========

#[test]
#[serial]
fn test_add_broadcast() {
    // [2, 3] + [3] → broadcast to [2, 3]
    let a = CudaTensor::from_slice(&[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3], DType::F32);
    let b = CudaTensor::from_slice(&[10.0f32, 20.0, 30.0], &[3], DType::F32);
    let c = a.add_impl(&b).unwrap();
    assert_eq!(c.shape(), &[2, 3]);
    approx_eq(
        &c.to_vec::<f32>(),
        &[11.0, 22.0, 33.0, 14.0, 25.0, 36.0],
        1e-6,
    );
}

// ========== 2D テンソル ==========

#[test]
#[serial]
fn test_add_2d() {
    let a = CudaTensor::from_slice(&[1.0f32, 2.0, 3.0, 4.0], &[2, 2], DType::F32);
    let b = CudaTensor::from_slice(&[5.0f32, 6.0, 7.0, 8.0], &[2, 2], DType::F32);
    let c = a.add_impl(&b).unwrap();
    assert_eq!(c.shape(), &[2, 2]);
    approx_eq(&c.to_vec::<f32>(), &[6.0, 8.0, 10.0, 12.0], 1e-6);
}

// ========== 比較演算 ==========

#[test]
#[serial]
fn test_eq() {
    let a = CudaTensor::from_slice(&[1.0f32, 2.0, 3.0], &[3], DType::F32);
    let b = CudaTensor::from_slice(&[1.0f32, 0.0, 3.0], &[3], DType::F32);
    let c = a.eq_impl(&b).unwrap();
    approx_eq(&c.to_vec::<f32>(), &[1.0, 0.0, 1.0], 1e-6);
}

#[test]
#[serial]
fn test_lt() {
    let a = CudaTensor::from_slice(&[1.0f32, 2.0, 3.0], &[3], DType::F32);
    let b = CudaTensor::from_slice(&[2.0f32, 2.0, 1.0], &[3], DType::F32);
    let c = a.lt_impl(&b).unwrap();
    approx_eq(&c.to_vec::<f32>(), &[1.0, 0.0, 0.0], 1e-6);
}