ruPRIM 0.1.6

Parallel primitives, reductions, scans, and indexing for Ruda.
Documentation 
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

use ruda_core::ir::features::AtomicUsage;
use ruda_kernel::dsl as kernel_dsl;
use ruda_kernel::dsl::{Runtime, calculate_ruda_count_elemwise, prelude::*};
use ruda_kernel::library::{FastDivmod, tensor::layout::linear::LinearView};
use ruda_kernel::tensor::{RudaTensor, layout::{address_type, shape_divmod_range}};

pub(super) fn scatter_nd_add<R: Runtime>(
    tensor: RudaTensor<R>,
    indices: RudaTensor<R>,
    values: RudaTensor<R>,
) -> RudaTensor<R> {
    let idx_shape = indices.meta.shape();
    let m = idx_shape.len();
    let k = idx_shape[m - 1];
    let num_updates: usize = idx_shape[..m - 1].iter().product();
    let slice_size: usize = tensor.meta.shape()[k..].iter().product();
    let total = num_updates * slice_size;
    let supports_atomic_add = tensor.client.properties()
        .atomic_type_usage(Type::new(StorageType::Atomic(tensor.dtype.into())))
        .contains(AtomicUsage::Add);
    let working_units = if supports_atomic_add { total } else { slice_size };
    let ruda_dim = RudaDim::new(tensor.client.properties(), working_units);
    let ruda_count = calculate_ruda_count_elemwise(&tensor.client, working_units, ruda_dim);
    let launch = if supports_atomic_add {
        atomic_add_kernel::launch_unchecked::<R>
    } else {
        ordered_add_kernel::launch_unchecked::<R>
    };
    let dtypes = [tensor.dtype.into(), indices.dtype.into()];
    let data_slice_shape = shape_divmod_range(&tensor, k..tensor.meta.num_dims());
    let address_type = address_type!(tensor, indices, values).max(AddressType::from_len(
        total.max(indices.meta.num_elements()),
    ));

    unsafe {
        launch(
            &tensor.client,
            ruda_count,
            ruda_dim,
            address_type,
            tensor.clone().into_tensor_arg(),
            indices.into_linear_view(),
            values.into_linear_view(),
            data_slice_shape,
            slice_size,
            k,
            num_updates,
            dtypes,
        );
    }
    tensor
}

#[ruda(launch_unchecked, address_type = "dynamic")]
fn atomic_add_kernel<T: Numeric, I: Int>(
    data: &mut Tensor<Atomic<T>>,
    indices: &LinearView<I>,
    values: &LinearView<T>,
    data_slice_shape: Sequence<FastDivmod<usize>>,
    slice_size: usize,
    k: usize,
    num_updates: usize,
    #[define(T, I)] _dtypes: [StorageType; 2],
) {
    if ABSOLUTE_POS >= num_updates * slice_size {
        terminate!();
    }
    let update_idx = ABSOLUTE_POS / slice_size;
    let mut remainder = ABSOLUTE_POS % slice_size;
    let slice_rank = data_slice_shape.len().comptime();
    let mut data_idx = 0usize;
    #[unroll]
    for i in 0..slice_rank {
        let dim = slice_rank - i - 1;
        let (rem, coord) = data_slice_shape[dim].div_mod(remainder);
        remainder = rem;
        data_idx += coord * data.stride(k + dim);
    }
    for j in 0..k {
        data_idx += usize::cast_from(indices[update_idx * k + j]) * data.stride(j);
    }
    data[data_idx].fetch_add(values[ABSOLUTE_POS]);
}

#[ruda(launch_unchecked, address_type = "dynamic")]
fn ordered_add_kernel<T: Numeric, I: Int>(
    data: &mut Tensor<T>,
    indices: &LinearView<I>,
    values: &LinearView<T>,
    data_slice_shape: Sequence<FastDivmod<usize>>,
    slice_size: usize,
    k: usize,
    num_updates: usize,
    #[define(T, I)] _dtypes: [StorageType; 2],
) {
    if ABSOLUTE_POS >= slice_size {
        terminate!();
    }
    let mut remainder = ABSOLUTE_POS;
    let slice_rank = data_slice_shape.len().comptime();
    let mut slice_offset = 0usize;
    #[unroll]
    for i in 0..slice_rank {
        let dim = slice_rank - i - 1;
        let (rem, coord) = data_slice_shape[dim].div_mod(remainder);
        remainder = rem;
        slice_offset += coord * data.stride(k + dim);
    }
    for update_idx in 0..num_updates {
        let mut data_idx = slice_offset;
        for j in 0..k {
            data_idx += usize::cast_from(indices[update_idx * k + j]) * data.stride(j);
        }
        data[data_idx] = data[data_idx] + values[update_idx * slice_size + ABSOLUTE_POS];
    }
}