tenferro-gpu 0.1.0

CubeCL-backed CUDA and WebGPU provider backends for tenferro tensors.
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
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170

// Run with: cargo test --features cuda -- --ignored
use num_complex::{Complex32, Complex64};

use crate::DotGeneralConfig;
use crate::Tensor;
use tenferro_tensor::TensorDot;

use super::{
    assert_tensor_close, cpu_backend, download, gpu_backend, tensor_c32, tensor_c64, tensor_f32,
    tensor_f64, upload,
};

fn run_dot_general_case(lhs: Tensor, rhs: Tensor, config: DotGeneralConfig, tol: f64) {
    let mut cpu = cpu_backend();
    let mut gpu = gpu_backend();

    let expected = cpu.dot_general(&lhs, &rhs, &config).unwrap();
    let gpu_lhs = upload(&gpu, &lhs);
    let gpu_rhs = upload(&gpu, &rhs);
    let actual_gpu = gpu.dot_general(&gpu_lhs, &gpu_rhs, &config).unwrap();
    let actual = download(&gpu, &actual_gpu);

    assert_eq!(actual.shape(), expected.shape());
    assert_tensor_close(&actual, &expected, tol);
}

fn matmul_config() -> DotGeneralConfig {
    DotGeneralConfig {
        lhs_contracting_dims: vec![1],
        rhs_contracting_dims: vec![0],
        lhs_batch_dims: vec![],
        rhs_batch_dims: vec![],
    }
}

#[test]
#[ignore]
fn test_dot_general_matmul_f32() {
    run_dot_general_case(
        tensor_f32(vec![2, 3], vec![1.0, 4.0, 2.0, 5.0, 3.0, 6.0]),
        tensor_f32(
            vec![3, 4],
            vec![
                1.0, 5.0, 9.0, 2.0, 6.0, 10.0, 3.0, 7.0, 11.0, 4.0, 8.0, 12.0,
            ],
        ),
        matmul_config(),
        1e-4,
    );
}

#[test]
#[ignore]
fn test_dot_general_matmul_f64() {
    run_dot_general_case(
        tensor_f64(vec![2, 3], vec![1.0, 4.0, 2.0, 5.0, 3.0, 6.0]),
        tensor_f64(
            vec![3, 4],
            vec![
                1.0, 5.0, 9.0, 2.0, 6.0, 10.0, 3.0, 7.0, 11.0, 4.0, 8.0, 12.0,
            ],
        ),
        matmul_config(),
        1e-9,
    );
}

#[test]
#[ignore]
fn test_dot_general_batched_matmul_f32() {
    run_dot_general_case(
        tensor_f32(
            vec![2, 3, 5],
            vec![
                1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 3.0, 4.0, 5.0, 6.0,
                7.0, 8.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0,
            ],
        ),
        tensor_f32(
            vec![3, 4, 5],
            vec![
                1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 2.0, 3.0, 4.0, 5.0,
                6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0,
                10.0, 11.0, 12.0, 13.0, 14.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0,
                14.0, 15.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0,
            ],
        ),
        DotGeneralConfig {
            lhs_contracting_dims: vec![1],
            rhs_contracting_dims: vec![0],
            lhs_batch_dims: vec![2],
            rhs_batch_dims: vec![2],
        },
        1e-4,
    );
}

#[test]
#[ignore]
fn test_dot_general_complex_matmul_c32() {
    run_dot_general_case(
        tensor_c32(
            vec![2, 3],
            vec![
                Complex32::new(1.0, 0.5),
                Complex32::new(2.0, -1.0),
                Complex32::new(3.0, 0.25),
                Complex32::new(4.0, 1.0),
                Complex32::new(5.0, -0.5),
                Complex32::new(6.0, 0.75),
            ],
        ),
        tensor_c32(
            vec![3, 4],
            vec![
                Complex32::new(1.0, -1.0),
                Complex32::new(2.0, 0.25),
                Complex32::new(3.0, 0.5),
                Complex32::new(4.0, -0.75),
                Complex32::new(5.0, 0.0),
                Complex32::new(6.0, 1.0),
                Complex32::new(7.0, -0.5),
                Complex32::new(8.0, 0.75),
                Complex32::new(9.0, -1.25),
                Complex32::new(10.0, 0.5),
                Complex32::new(11.0, 1.5),
                Complex32::new(12.0, -0.25),
            ],
        ),
        matmul_config(),
        1e-4,
    );
}

#[test]
#[ignore]
fn test_dot_general_complex_matmul_c64() {
    run_dot_general_case(
        tensor_c64(
            vec![2, 3],
            vec![
                Complex64::new(1.0, 0.5),
                Complex64::new(2.0, -1.0),
                Complex64::new(3.0, 0.25),
                Complex64::new(4.0, 1.0),
                Complex64::new(5.0, -0.5),
                Complex64::new(6.0, 0.75),
            ],
        ),
        tensor_c64(
            vec![3, 4],
            vec![
                Complex64::new(1.0, -1.0),
                Complex64::new(2.0, 0.25),
                Complex64::new(3.0, 0.5),
                Complex64::new(4.0, -0.75),
                Complex64::new(5.0, 0.0),
                Complex64::new(6.0, 1.0),
                Complex64::new(7.0, -0.5),
                Complex64::new(8.0, 0.75),
                Complex64::new(9.0, -1.25),
                Complex64::new(10.0, 0.5),
                Complex64::new(11.0, 1.5),
                Complex64::new(12.0, -0.25),
            ],
        ),
        matmul_config(),
        1e-9,
    );
}