ruprim/block/
shuffle.rs

1use ruda_kernel::dsl as kernel_dsl;
2use ruda_kernel::dsl::prelude::*;
3use crate::collective::record::{RudaRead, RudaReadExpand, RudaWrite, RudaWriteExpand};
4
5#[ruda]
6pub fn offset_access<T: RudaType<ExpandType: Assign> + Copy, S: RudaWrite<T>>(
7    input: T, output: &mut T, scratch: &mut S, distance: i32, #[comptime] threads: u32,
8) {
9    scratch.write(UNIT_POS as usize, input);
10    sync_ruda();
11    let source = UNIT_POS as i64 + distance as i64;
12    if source >= 0 && source < threads as i64 { *output = scratch.read(source as usize); }
13    sync_ruda();
14}
15
16#[ruda]
17pub fn rotate_access<T: RudaType<ExpandType: Assign> + Copy, S: RudaWrite<T>>(
18    input: T, scratch: &mut S, distance: u32, #[comptime] threads: u32,
19) -> T {
20    scratch.write(UNIT_POS as usize, input);
21    sync_ruda();
22    let source = (UNIT_POS + distance % threads) % threads;
23    let output = scratch.read(source as usize);
24    sync_ruda();
25    output
26}
27
28/// Shift records in blocked order. The unpaired output item is unchanged;
29/// return the discarded prefix (down) or suffix (up) to every thread.
30#[ruda]
31pub fn shift_access<T: RudaType<ExpandType: Assign> + Copy, I: RudaRead<T>, W: RudaWrite<T>, S: RudaWrite<T>>(
32    input: &I, output: &mut W, scratch: &mut S,
33    #[comptime] threads: usize, #[comptime] items: usize, #[comptime] down: bool,
34) -> T {
35    let start = UNIT_POS as usize * items;
36    let total = threads * items;
37    #[unroll]
38    for item in 0..items { scratch.write(start + item, input.read(item)); }
39    sync_ruda();
40    #[unroll]
41    for item in 0..items {
42        let index = start + item;
43        if down {
44            if index + 1 < total { output.write(item, scratch.read(index + 1)); }
45        } else if index > 0 { output.write(item, scratch.read(index - 1)); }
46    }
47    let boundary = if down { scratch.read(0) } else { scratch.read(total - 1) };
48    sync_ruda();
49    boundary
50}
51
52/// Offset one scalar per thread. Out-of-block destinations preserve `output`.
53#[ruda]
54pub fn offset<T: RudaPrimitive>(
55    input: T,
56    output: &mut T,
57    scratch: &mut SharedMemory<T>,
58    distance: i32,
59    #[comptime] threads: u32,
60) {
61    scratch[UNIT_POS as usize] = input;
62    sync_ruda();
63    let source = UNIT_POS as i32 + distance;
64    if source >= 0 && source < threads as i32 {
65        *output = scratch[source as usize];
66    }
67    sync_ruda();
68}
69
70/// Circular offset one scalar per thread.
71#[ruda]
72pub fn rotate<T: RudaPrimitive>(
73    input: T,
74    scratch: &mut SharedMemory<T>,
75    distance: u32,
76    #[comptime] threads: u32,
77) -> T {
78    scratch[UNIT_POS as usize] = input;
79    sync_ruda();
80    let source = (UNIT_POS + distance % threads) % threads;
81    let output = scratch[source as usize];
82    sync_ruda();
83    output
84}
85
86/// Shift a blocked register tile one item. Boundary output is preserved.
87/// Returns the block's discarded prefix (down) or suffix (up).
88#[ruda]
89pub fn shift<T: RudaPrimitive>(
90    input: &Array<T>,
91    output: &mut Array<T>,
92    scratch: &mut SharedMemory<T>,
93    #[comptime] threads: usize,
94    #[comptime] items_per_thread: usize,
95    #[comptime] down: bool,
96) -> T {
97    let start = UNIT_POS as usize * items_per_thread;
98    let total = threads * items_per_thread;
99    #[unroll]
100    for item in 0..items_per_thread { scratch[start + item] = input[item]; }
101    sync_ruda();
102    #[unroll]
103    for item in 0..items_per_thread {
104        let index = start + item;
105        if down {
106            if index + 1 < total { output[item] = scratch[index + 1]; }
107        } else {
108            if index > 0 { output[item] = scratch[index - 1]; }
109        }
110    }
111    let boundary = if down { scratch[0] } else { scratch[total - 1] };
112    sync_ruda();
113    boundary
114}
ruprim/block/shuffle.rs

ruprim/block/
shuffle.rs
