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

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

pub(crate) fn launch_copy_offset<T: DeviceRepr>(
    device: &Cuda,
    kernel_name: &str,
    module: &kernel::Module,
    src: &CudaSlice<T>,
    layout: &Layout,
    dst: &CudaSlice<T>,
    dst_offset: usize,
) -> CudaResult<()> {
    let dims = layout.dims();
    let elem_count = layout.shape().element_count();
    let num_dims = dims.len();
    let func = device.load_function(kernel_name, module)?;

    let mut builder = func.builder();
    let dims = device.memcpy_stod(dims)?;
    let strides = device.memcpy_stod(layout.stride())?;
    let src_view = src.slice(layout.start_offset()..);
    builder_arg!(builder, elem_count);
    builder_arg!(builder, num_dims);
    builder.arg(&dims);
    builder.arg(&strides);
    builder.arg(&dst_offset);
    builder.arg(&src_view);
    builder.arg(dst);

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

pub(crate) fn launch_copy2d<T: DeviceRepr>(
    device: &Cuda,
    kernel_name: &str,
    module: &kernel::Module,
    d1: usize,
    d2: usize,
    src_s: usize,
    dst_s: usize,
    src: &CudaSlice<T>,
    src_offset: usize,
    dst: &CudaSlice<T>,
    dst_offset: usize,
) -> CudaResult<()> {
    let func = device.load_function(kernel_name, module)?;
    let total = d1 * d2;

    let mut builder = func.builder();
    let src_view = src.slice(src_offset..);
    let dst_view = dst.slice(dst_offset..);
    builder_arg!(builder, d1);
    builder_arg!(builder, d2);
    builder_arg!(builder, src_s);
    builder_arg!(builder, dst_s);
    builder.arg(&src_view);
    builder.arg(&dst_view);

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