ruPRIM 0.1.9

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

use ruda_kernel::dsl as kernel_dsl;
use ruda_kernel::dsl::prelude::*;

/// Convert blocked and striped layouts. Scratch holds the entire block tile.
#[ruda]
pub fn transpose<T: RudaPrimitive>(
    input: &Array<T>,
    output: &mut Array<T>,
    scratch: &mut SharedMemory<T>,
    #[comptime] threads: usize,
    #[comptime] items_per_thread: usize,
    #[comptime] blocked_to_striped: bool,
) {
    let lane = UNIT_POS as usize;
    #[unroll]
    for item in 0..items_per_thread {
        let index = if blocked_to_striped { lane * items_per_thread + item } else { item * threads + lane };
        scratch[index] = input[item];
    }
    sync_ruda();
    #[unroll]
    for item in 0..items_per_thread {
        let index = if blocked_to_striped { item * threads + lane } else { lane * items_per_thread + item };
        output[item] = scratch[index];
    }
    sync_ruda();
}

/// Scatter a permutation of tile ranks into blocked or striped register order.
#[ruda]
pub fn scatter<T: RudaPrimitive>(
    input: &Array<T>,
    ranks: &Array<u32>,
    output: &mut Array<T>,
    scratch: &mut SharedMemory<T>,
    #[comptime] threads: usize,
    #[comptime] items_per_thread: usize,
    #[comptime] striped: bool,
) {
    #[unroll]
    for item in 0..items_per_thread {
        let rank = ranks[item] as usize;
        scratch[rank] = input[item];
    }
    sync_ruda();
    #[unroll]
    for item in 0..items_per_thread {
        let index = if striped { item * threads + UNIT_POS as usize } else { UNIT_POS as usize * items_per_thread + item };
        output[item] = scratch[index];
    }
    sync_ruda();
}

/// Blocked <-> warp-striped exchange, including a partial final logical warp.
/// `width` is an explicit layout width rather than a fixed hardware warp size.
#[ruda]
pub fn warp_transpose<T: RudaPrimitive>(
    input: &Array<T>, output: &mut Array<T>, scratch: &mut SharedMemory<T>,
    #[comptime] threads: usize, #[comptime] items_per_thread: usize,
    #[comptime] width: usize, #[comptime] blocked_to_warp_striped: bool,
) {
    let thread = UNIT_POS as usize;
    let first = thread / width * width;
    let lane = thread - first;
    let lanes = min(width, threads - first);
    #[unroll]
    for item in 0..items_per_thread {
        let blocked = thread * items_per_thread + item;
        let striped = first * items_per_thread + item * lanes + lane;
        let index = if blocked_to_warp_striped { blocked } else { striped };
        scratch[index] = input[item];
    }
    sync_ruda();
    #[unroll]
    for item in 0..items_per_thread {
        let blocked = thread * items_per_thread + item;
        let striped = first * items_per_thread + item * lanes + lane;
        let index = if blocked_to_warp_striped { striped } else { blocked };
        output[item] = scratch[index];
    }
    sync_ruda();
}

/// Scatter valid items to striped output. Valid ranks must be unique and in
/// tile range. Scratch positions not written by a valid item are unspecified.
#[ruda]
pub fn scatter_flagged<T: RudaPrimitive, I: Int>(
    input: &Array<T>, ranks: &Array<I>, valid: &Array<bool>, output: &mut Array<T>,
    scratch: &mut SharedMemory<T>, #[comptime] threads: usize, #[comptime] items_per_thread: usize,
) {
    #[unroll]
    for item in 0..items_per_thread {
        if valid[item] {
            let rank = usize::cast_from(ranks[item]);
            scratch[rank] = input[item];
        }
    }
    sync_ruda();
    #[unroll]
    for item in 0..items_per_thread { output[item] = scratch[item * threads + UNIT_POS as usize]; }
    sync_ruda();
}

/// Negative ranks are excluded from the exchange.
#[ruda]
pub fn scatter_guarded<T: RudaPrimitive, I: Int>(
    input: &Array<T>, ranks: &Array<I>, output: &mut Array<T>, scratch: &mut SharedMemory<T>,
    #[comptime] threads: usize, #[comptime] items_per_thread: usize,
) {
    let mut valid = Array::<bool>::new(items_per_thread);
    #[unroll]
    for item in 0..items_per_thread { valid[item] = ranks[item] >= I::from_int(0); }
    scatter_flagged(input, ranks, &valid, output, scratch, threads, items_per_thread);
}