luma-cuda 0.3.2

luma cuda implement
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

use luma_tensor::Layout;
use crate::{Cuda, CudaError, CudaResult};
use cudarc::{
    cublas::{CudaBlas, Gemm, GemmConfig, StridedBatchedConfig, sys::cublasOperation_t},
    driver::{CudaSlice, DeviceRepr},
};

fn matmul_config<T: Copy>(
    alpha: T,
    beta: T,
    (b, m, n, k): (usize, usize, usize, usize),
    lhs_l: &Layout,
    rhs_l: &Layout,
) -> CudaResult<StridedBatchedConfig<T>> {
    let lhs_stride = lhs_l.stride();
    let rhs_stride = rhs_l.stride();
    let lhs_dims = lhs_l.dims();
    let rhs_dims = rhs_l.dims();

    if lhs_dims.len() < 2 || rhs_dims.len() < 2 {
        return Err(CudaError::MatMulNonContiguous { msg: "matmul requires at least 2D".into() });
    }

    let (lhs_m2, lhs_m1) = (lhs_stride[lhs_stride.len() - 2], lhs_stride[lhs_stride.len() - 1]);
    let (rhs_m2, rhs_m1) = (rhs_stride[rhs_stride.len() - 2], rhs_stride[rhs_stride.len() - 1]);

    let (ldb, transb) = if (lhs_m1 == 1 || k == 1) && (lhs_m2 == k || m == 1) {
        (k as i32, cublasOperation_t::CUBLAS_OP_N)
    } else if (lhs_m1 == m || k == 1) && (lhs_m2 == 1 || m == 1) {
        (m as i32, cublasOperation_t::CUBLAS_OP_T)
    } else {
        return Err(CudaError::MatMulNonContiguous {
            msg: format!("LHS stride {:?} invalid for shape {:?} (m={}, k={})", lhs_stride, lhs_dims, m, k),
        });
    };

    let (lda, transa) = if (rhs_m1 == 1 || n == 1) && (rhs_m2 == n || k == 1) {
        (n as i32, cublasOperation_t::CUBLAS_OP_N)
    } else if (rhs_m1 == k || n == 1) && (rhs_m2 == 1 || k == 1) {
        (k as i32, cublasOperation_t::CUBLAS_OP_T)
    } else {
        return Err(CudaError::MatMulNonContiguous {
            msg: format!("RHS stride {:?} invalid for shape {:?} (k={}, n={})", rhs_stride, rhs_dims, k, n),
        });
    };

    let stride_a: i64 = if rhs_stride.len() >= 3 { rhs_stride[rhs_stride.len() - 3] as i64 } else { (n * k) as i64 };
    let stride_b: i64 = if lhs_stride.len() >= 3 { lhs_stride[lhs_stride.len() - 3] as i64 } else { (m * k) as i64 };
    let stride_c: i64 = (m * n) as i64;

    Ok(StridedBatchedConfig {
        batch_size: b as i32,
        gemm: GemmConfig { alpha, beta, m: n as i32, n: m as i32, k: k as i32, lda, ldb, ldc: n as i32, transa, transb },
        stride_a,
        stride_b,
        stride_c,
    })
}

pub(crate) fn launch_matmul<T: Copy + DeviceRepr>(
    device: &Cuda,
    alpha: T,
    beta: T,
    (b, m, n, k): (usize, usize, usize, usize),
    lhs: &CudaSlice<T>,
    lhs_l: &Layout,
    rhs: &CudaSlice<T>,
    rhs_l: &Layout,
) -> CudaResult<CudaSlice<T>>
where
    CudaBlas: Gemm<T>,
{
    let cfg = matmul_config(alpha, beta, (b, m, n, k), lhs_l, rhs_l)?;
    let mut out = device.alloc::<T>(b * m * n)?;
    let blas = device.0.blas.lock().unwrap();
    unsafe {
        blas.gemm_strided_batched(cfg, rhs, lhs, &mut out).map_err(CudaError::Cublas)?;
    }
    Ok(out)
}

pub(crate) fn launch_add_matmul_<T: Copy + DeviceRepr>(
    device: &Cuda,
    alpha: T,
    beta: T,
    dst: &mut CudaSlice<T>,
    _dst_l: &Layout,
    lhs: &CudaSlice<T>,
    lhs_l: &Layout,
    rhs: &CudaSlice<T>,
    rhs_l: &Layout,
    (b, m, n, k): (usize, usize, usize, usize),
) -> CudaResult<()>
where
    CudaBlas: Gemm<T>,
{
    let cfg = matmul_config(alpha, beta, (b, m, n, k), lhs_l, rhs_l)?;
    let blas = device.0.blas.lock().unwrap();
    unsafe {
        blas.gemm_strided_batched(cfg, rhs, lhs, dst).map_err(CudaError::Cublas)?;
    }
    Ok(())
}