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

//! 軸付きリダクション演算 — 全て GPU カーネルで完結 (to_vec ゼロ)

use crate::cuda_sys::cudaStream_t;
use crate::tensor::CudaTensor;
use crate::DType;
use tl_backend::{BackendError, BackendResult};

extern "C" {
    fn launch_reduce_axis_sum_kernel(
        input: *const f32,
        output: *mut f32,
        outer: i32,
        axis_size: i32,
        inner: i32,
        stream: cudaStream_t,
    );
    fn launch_reduce_axis_max_kernel(
        input: *const f32,
        output: *mut f32,
        outer: i32,
        axis_size: i32,
        inner: i32,
        stream: cudaStream_t,
    );
    fn launch_reduce_axis_min_kernel(
        input: *const f32,
        output: *mut f32,
        outer: i32,
        axis_size: i32,
        inner: i32,
        stream: cudaStream_t,
    );
    fn launch_argmax_axis_kernel(
        input: *const f32,
        output: *mut f32,
        outer: i32,
        axis_size: i32,
        inner: i32,
        stream: cudaStream_t,
    );
    fn launch_argmin_axis_kernel(
        input: *const f32,
        output: *mut f32,
        outer: i32,
        axis_size: i32,
        inner: i32,
        stream: cudaStream_t,
    );
}

impl CudaTensor {
    /// 軸パラメータを正規化し (outer, axis_size, inner, out_shape) を返す
    fn resolve_axis(&self, axis: i32) -> BackendResult<(usize, usize, usize, Vec<usize>)> {
        let shape = self.shape().to_vec();
        let ndim = shape.len();
        let axis = if axis < 0 {
            (ndim as i32 + axis) as usize
        } else {
            axis as usize
        };
        if axis >= ndim {
            return Err(BackendError::ArgumentError(format!(
                "axis {} out of range for ndim {}",
                axis, ndim
            )));
        }
        let outer: usize = shape[..axis].iter().product::<usize>().max(1);
        let axis_size = shape[axis];
        let inner: usize = shape[axis + 1..].iter().product::<usize>().max(1);
        let mut out_shape = shape.clone();
        out_shape.remove(axis);
        if out_shape.is_empty() {
            out_shape.push(1);
        }
        Ok((outer, axis_size, inner, out_shape))
    }

    /// 軸リダクションの GPU カーネル共通パターン
    fn reduce_axis_kernel_op(
        &self,
        axis: i32,
        launch: unsafe extern "C" fn(*const f32, *mut f32, i32, i32, i32, cudaStream_t),
    ) -> BackendResult<CudaTensor> {
        let (outer, axis_size, inner, out_shape) = self.resolve_axis(axis)?;
        let output = CudaTensor::uninit(&out_shape, DType::F32);
        let stream = crate::stream::get_stream().raw();
        unsafe {
            launch(
                self.buffer.ptr() as *const f32,
                output.buffer.ptr() as *mut f32,
                outer as i32,
                axis_size as i32,
                inner as i32,
                stream,
            );
        }
        crate::stream::sync_stream();
        Ok(output)
    }

    /// 軸に沿った合計 — GPU カーネル
    pub fn sum_impl(&self, axis: i32) -> BackendResult<CudaTensor> {
        self.reduce_axis_kernel_op(axis, launch_reduce_axis_sum_kernel)
    }

    /// 軸に沿った平均 — GPU カーネル
    pub fn mean_impl(&self, axis: i32) -> BackendResult<CudaTensor> {
        let (_, axis_size, _, _) = self.resolve_axis(axis)?;
        let sum = self.sum_impl(axis)?;
        sum.div_scalar_impl(axis_size as f32)
    }

    /// 軸に沿った最大値 — GPU カーネル
    pub fn max_impl(&self, axis: i32) -> BackendResult<CudaTensor> {
        self.reduce_axis_kernel_op(axis, launch_reduce_axis_max_kernel)
    }

    /// 軸に沿った最小値 — GPU カーネル
    pub fn min_impl(&self, axis: i32) -> BackendResult<CudaTensor> {
        self.reduce_axis_kernel_op(axis, launch_reduce_axis_min_kernel)
    }

    /// 軸に沿った argmax — GPU カーネル
    pub fn argmax_impl(&self, axis: i32) -> BackendResult<CudaTensor> {
        self.reduce_axis_kernel_op(axis, launch_argmax_axis_kernel)
    }

    /// 軸に沿った argmin — GPU カーネル
    pub fn argmin_impl(&self, axis: i32) -> BackendResult<CudaTensor> {
        self.reduce_axis_kernel_op(axis, launch_argmin_axis_kernel)
    }
}