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

//! 単項演算（要素ごと）— CUDA カーネルで GPU 上で完結

use crate::cuda_sys::cudaStream_t;
use crate::tensor::CudaTensor;
use crate::DType;
use tl_backend::BackendResult;

extern "C" {
    fn launch_neg_kernel(x: *const f32, y: *mut f32, n: i32, stream: cudaStream_t);
    fn launch_abs_kernel(x: *const f32, y: *mut f32, n: i32, stream: cudaStream_t);
    fn launch_exp_kernel(x: *const f32, y: *mut f32, n: i32, stream: cudaStream_t);
    fn launch_log_kernel(x: *const f32, y: *mut f32, n: i32, stream: cudaStream_t);
    fn launch_sqrt_kernel(x: *const f32, y: *mut f32, n: i32, stream: cudaStream_t);
}

impl CudaTensor {
    /// GPU カーネルによる単項演算の共通パターン
    fn unary_kernel_op(
        &self,
        launch: unsafe extern "C" fn(*const f32, *mut f32, i32, cudaStream_t),
    ) -> BackendResult<CudaTensor> {
        let n = self.elem_count();
        let output = CudaTensor::uninit(self.shape(), DType::F32);
        let stream = crate::stream::get_stream().raw();
        unsafe {
            launch(
                self.buffer.ptr() as *const f32,
                output.buffer.ptr() as *mut f32,
                n as i32,
                stream,
            );
        }
        crate::stream::sync_stream();
        Ok(output)
    }

    /// 負数
    pub fn neg_impl(&self) -> BackendResult<CudaTensor> {
        self.unary_kernel_op(launch_neg_kernel)
    }

    /// 絶対値
    pub fn abs_impl(&self) -> BackendResult<CudaTensor> {
        self.unary_kernel_op(launch_abs_kernel)
    }

    /// 指数関数
    pub fn exp_impl(&self) -> BackendResult<CudaTensor> {
        self.unary_kernel_op(launch_exp_kernel)
    }

    /// 対数関数
    pub fn log_impl(&self) -> BackendResult<CudaTensor> {
        self.unary_kernel_op(launch_log_kernel)
    }

    /// 平方根
    pub fn sqrt_impl(&self) -> BackendResult<CudaTensor> {
        self.unary_kernel_op(launch_sqrt_kernel)
    }
}