ruPRIM 0.1.7

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

use ruda_kernel::dsl as kernel_dsl;
use ruda_kernel::dsl::prelude::*;
use crate::collective::{RudaCompare, merge};

/// Stable logical-warp key/value merge sort, in blocked register order.
/// Scratch regions cover the block's threads times `items_per_lane`; each
/// logical warp uses a disjoint region. Only subgroup barriers are used.
#[ruda]
pub fn merge_sort_pairs<K: RudaPrimitive, V: RudaPrimitive, C: RudaCompare<K>>(
    keys: &mut Array<K>,
    values: &mut Array<V>,
    source_keys: &mut SharedMemory<K>,
    destination_keys: &mut SharedMemory<K>,
    source_values: &mut SharedMemory<V>,
    destination_values: &mut SharedMemory<V>,
    compare: &C,
    valid_items: usize,
    #[comptime] width: u32,
    #[comptime] items_per_lane: usize,
) {
    let lane = UNIT_POS_PLANE % width;
    let base = (UNIT_POS - lane) as usize * items_per_lane;
    let start = lane as usize * items_per_lane;
    #[unroll]
    for item in 0..items_per_lane {
        if start + item < valid_items {
            source_keys[base + start + item] = keys[item];
            source_values[base + start + item] = values[item];
        }
    }
    sync_plane();
    let mut run = 1usize;
    while run < width as usize * items_per_lane {
        #[unroll]
        for item in 0..items_per_lane {
            let index = start + item;
            if index < valid_items {
                let rank = merge::rank::<K, C>(source_keys, compare, index, valid_items, run, base);
                destination_keys[base + rank] = source_keys[base + index];
                destination_values[base + rank] = source_values[base + index];
            }
        }
        sync_plane();
        #[unroll]
        for item in 0..items_per_lane {
            let index = start + item;
            if index < valid_items {
                source_keys[base + index] = destination_keys[base + index];
                source_values[base + index] = destination_values[base + index];
            }
        }
        sync_plane();
        run *= 2;
    }
    #[unroll]
    for item in 0..items_per_lane {
        if start + item < valid_items {
            keys[item] = source_keys[base + start + item];
            values[item] = source_values[base + start + item];
        }
    }
    sync_plane();
}

/// Stable logical-warp merge sort for keys only.
#[ruda]
pub fn merge_sort_keys<K: RudaPrimitive, C: RudaCompare<K>>(
    keys: &mut Array<K>,
    source: &mut SharedMemory<K>,
    destination: &mut SharedMemory<K>,
    compare: &C,
    valid_items: usize,
    #[comptime] width: u32,
    #[comptime] items_per_lane: usize,
) {
    let lane = UNIT_POS_PLANE % width;
    let base = (UNIT_POS - lane) as usize * items_per_lane;
    let start = lane as usize * items_per_lane;
    #[unroll]
    for item in 0..items_per_lane {
        if start + item < valid_items {
            source[base + start + item] = keys[item];
        }
    }
    sync_plane();
    let mut run = 1usize;
    while run < width as usize * items_per_lane {
        #[unroll]
        for item in 0..items_per_lane {
            let index = start + item;
            if index < valid_items {
                let rank = merge::rank::<K, C>(source, compare, index, valid_items, run, base);
                destination[base + rank] = source[base + index];
            }
        }
        sync_plane();
        #[unroll]
        for item in 0..items_per_lane {
            if start + item < valid_items {
                source[base + start + item] = destination[base + start + item];
            }
        }
        sync_plane();
        run *= 2;
    }
    #[unroll]
    for item in 0..items_per_lane {
        if start + item < valid_items {
            keys[item] = source[base + start + item];
        }
    }
    sync_plane();
}