luma-cuda 0.3.0

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

use crate::{Cuda, CudaError, CudaResult, kernel};
use crate::builder_arg;
use luma_tensor::{CmpOp, Layout};
use cudarc::driver::{CudaSlice, DeviceRepr, LaunchConfig, PushKernelArg};

pub(crate) fn cmp_kernel_name(op: CmpOp, suffix: &str) -> String {
    let op_str = match op {
        CmpOp::Eq => "eq",
        CmpOp::Ne => "ne",
        CmpOp::Lt => "lt",
        CmpOp::Le => "le",
        CmpOp::Gt => "gt",
        CmpOp::Ge => "ge",
    };
    format!("b{}_{}", op_str, suffix)
}

pub(crate) fn launch_cmp<T: DeviceRepr>(
    device: &Cuda,
    op: CmpOp,
    type_name: &str,
    module: &kernel::Module,
    lhs: &CudaSlice<T>,
    rhs: &CudaSlice<T>,
    lhs_l: &Layout,
    rhs_l: &Layout,
) -> CudaResult<CudaSlice<u8>> {
    let dims = lhs_l.dims();
    let elem_count = lhs_l.shape().element_count();
    let num_dims = dims.len();
    let kernel_name = cmp_kernel_name(op, type_name);
    let func = device.load_function(&kernel_name, module)?;

    let mut builder = func.builder();
    let dims_dev = device.memcpy_stod(dims)?;
    let lhs_strides = device.memcpy_stod(lhs_l.stride())?;
    let rhs_strides = device.memcpy_stod(rhs_l.stride())?;
    let output = device.alloc::<u8>(elem_count)?;
    let lhs_view = lhs.slice(lhs_l.start_offset()..);
    let rhs_view = rhs.slice(rhs_l.start_offset()..);
    builder_arg!(builder, elem_count);
    builder_arg!(builder, num_dims);
    builder.arg(&dims_dev);
    builder.arg(&lhs_strides);
    builder.arg(&rhs_strides);
    builder.arg(&lhs_view);
    builder.arg(&rhs_view);
    builder.arg(&output);

    let config = LaunchConfig::for_num_elems(elem_count as u32);
    unsafe { builder.launch(config) }.map_err(CudaError::CudaDriver)?;
    Ok(output)
}

pub(crate) fn launch_cmp_scalar<T: DeviceRepr>(
    device: &Cuda,
    op: CmpOp,
    type_name: &str,
    module: &kernel::Module,
    lhs: &CudaSlice<T>,
    lhs_l: &Layout,
    rhs: T,
) -> CudaResult<CudaSlice<u8>> {
    let rhs_dev = device.memcpy_stod(&[rhs])?;
    let broadcast_l = Layout::contiguous(lhs_l.shape().clone());
    launch_cmp(device, op, type_name, module, lhs, &rhs_dev, lhs_l, &broadcast_l)
}