hodu_cuda_kernels 0.2.2

hodu cuda kernels
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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

use crate::{
    compat::*,
    cuda::*,
    error::{CudaKernelError, Result},
    kernel::Kernels,
    kernels::macros::ops,
    source::Source,
};

ops!(concat, split, chunk);

/// Executes a concatenation operation on multiple input tensors along a specified dimension.
///
/// # Arguments
/// * `kernel` - Concat operation kernel (e.g., concat::F32)
/// * `kernels` - Kernel cache for managing compiled kernels
/// * `context` - CUDA context to execute on
/// * `input` - Combined input buffer containing all input tensors
/// * `output` - Output buffer for concatenated result
/// * `metadata` - Metadata describing tensor shapes, strides, and offsets
///
/// # Metadata Layout
/// The metadata slice must contain the following elements in order:
/// - `metadata[0]`: num_els (total number of output elements)
/// - `metadata[1]`: num_dims (number of dimensions)
/// - `metadata[2..2+num_dims]`: output_shape (shape of output tensor)
/// - `metadata[2+num_dims]`: concat_dim (dimension along which to concatenate)
/// - `metadata[2+num_dims+1]`: num_inputs (number of input tensors)
/// - `metadata[2+num_dims+2..2+num_dims+2+num_inputs*num_dims]`: input_shapes (flattened)
/// - `metadata[2+num_dims+2+num_inputs*num_dims..2+num_dims+2+2*num_inputs*num_dims]`: input_strides (flattened)
/// - `metadata[2+num_dims+2+2*num_inputs*num_dims..2+num_dims+2+2*num_inputs*num_dims+num_inputs]`: input_offsets
/// - `metadata[2+num_dims+2+2*num_inputs*num_dims+num_inputs..]`: input_buffer_offsets
///
/// Total metadata length: `2 + num_dims + 2 + num_inputs * (2 * num_dims + 2)`
pub fn call_ops_concat<T>(
    kernel: crate::kernels::macros::Kernel,
    kernels: &Kernels,
    context: &Arc<CudaContext>,
    input: &CudaSlice<T>,
    output: &mut CudaSlice<T>,
    metadata: &[usize],
) -> Result<()>
where
    T: cudarc::driver::DeviceRepr,
{
    let func = kernels.load_function(context, Source::OpsConcatSplit, kernel.0)?;

    let num_els = metadata[0];
    let block_size = 256u32;
    let grid_size = (num_els as u32).div_ceil(block_size).max(1);

    let cfg = LaunchConfig {
        grid_dim: (grid_size, 1, 1),
        block_dim: (block_size, 1, 1),
        shared_mem_bytes: 0,
    };

    let stream = context.default_stream();
    let metadata_dev = stream
        .memcpy_stod(metadata)
        .map_err(|e| CudaKernelError::MemoryError(format!("Failed to copy metadata: {:?}", e)))?;

    unsafe {
        func.launch(&stream, cfg, |args| {
            args.arg(input).arg(output).arg(&metadata_dev);
        })
        .map_err(|e| CudaKernelError::LaunchError(format!("Failed to launch kernel: {:?}", e)))?;
    }

    Ok(())
}

/// Executes a split operation to extract a portion of a tensor along a specified dimension.
///
/// # Arguments
/// * `kernel` - Split operation kernel (e.g., split::F32)
/// * `kernels` - Kernel cache for managing compiled kernels
/// * `context` - CUDA context to execute on
/// * `input` - Input tensor buffer
/// * `output` - Output buffer for extracted portion
/// * `metadata` - Metadata describing tensor shape, strides, and split parameters
///
/// # Metadata Layout
/// The metadata slice must contain the following elements in order:
/// - `metadata[0]`: num_els (total number of output elements)
/// - `metadata[1]`: num_dims (number of dimensions)
/// - `metadata[2..2+num_dims]`: input_shape (shape of input tensor)
/// - `metadata[2+num_dims..2+2*num_dims]`: strides (stride for each dimension)
/// - `metadata[2+2*num_dims]`: offset (starting offset in input buffer)
/// - `metadata[2+2*num_dims+1]`: split_dim (dimension along which to split)
/// - `metadata[2+2*num_dims+2]`: output_size_on_dim (size of output on split dimension)
/// - `metadata[2+2*num_dims+3]`: split_offset (offset on split dimension)
///
/// Total metadata length: `2 + num_dims * 2 + 4`
pub fn call_ops_split<T>(
    kernel: crate::kernels::macros::Kernel,
    kernels: &Kernels,
    context: &Arc<CudaContext>,
    input: &CudaSlice<T>,
    output: &mut CudaSlice<T>,
    metadata: &[usize],
) -> Result<()>
where
    T: cudarc::driver::DeviceRepr,
{
    let func = kernels.load_function(context, Source::OpsConcatSplit, kernel.0)?;

    let num_els = metadata[0];
    let block_size = 256u32;
    let grid_size = (num_els as u32).div_ceil(block_size).max(1);

    let cfg = LaunchConfig {
        grid_dim: (grid_size, 1, 1),
        block_dim: (block_size, 1, 1),
        shared_mem_bytes: 0,
    };

    let stream = context.default_stream();
    let metadata_dev = stream
        .memcpy_stod(metadata)
        .map_err(|e| CudaKernelError::MemoryError(format!("Failed to copy metadata: {:?}", e)))?;

    unsafe {
        func.launch(&stream, cfg, |args| {
            args.arg(input).arg(output).arg(&metadata_dev);
        })
        .map_err(|e| CudaKernelError::LaunchError(format!("Failed to launch kernel: {:?}", e)))?;
    }

    Ok(())
}