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

//! tl_cuda - CUDA GPU Tensor Library
//!
//! Metal バックエンド (tl_metal) と同じ API を CUDA で実装する。
#![allow(unused)]

#[macro_use]
pub mod cuda_sys;
pub mod buffer_pool;
pub mod device;
pub mod stream;
pub mod tensor;
pub mod ops {
    pub mod activation;
    pub mod binary;
    pub mod broadcast;
    pub mod fused;
    pub mod index;
    pub mod llm;
    pub mod matmul;
    pub mod nn;
    pub mod quantized;
    pub mod reduce;
    pub mod reduce_axis;
    pub mod scalar;
    pub mod shape;
    pub mod special;
    pub mod unary;
}
pub mod autograd;
pub mod backend_impl;
pub mod device_impl;
pub mod ffi;
pub mod ffi_ops;
pub mod graph;
pub mod optim;

pub use autograd::GradFn;
pub use buffer_pool::CudaBufferPool;
pub use device::{get_device, CudaDevice};
pub use optim::{clip_grad_norm, Adam, AdamW, SGD};
pub use tensor::CudaTensor;

// Re-export tl_backend traits
pub use ffi::OpaqueTensor;
pub use tl_backend::GpuTensor;

/// データ型
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum DType {
    F32,
    F16,
    I32,
    I64,
    U8,
}

impl DType {
    /// バイト数
    pub fn size_in_bytes(&self) -> usize {
        match self {
            DType::F32 => 4,
            DType::F16 => 2,
            DType::I32 => 4,
            DType::I64 => 8,
            DType::U8 => 1,
        }
    }
}

/// 形状からバイト数を計算
pub fn shape_to_bytes(shape: &[usize], dtype: DType) -> usize {
    let elem_count: usize = shape.iter().product();
    elem_count * dtype.size_in_bytes()
}