ruprim/collective/
tile.rs

1use ruda_kernel::dsl as kernel_dsl;
2use ruda_kernel::dsl::prelude::*;
3use super::{RudaUnaryOp, RudaUnaryOpExpand};
4use super::record::{RudaRead, RudaReadExpand, RudaWrite, RudaWriteExpand};
5
6/// Map a register to blocked, striped, or logical-warp-striped tile order.
7#[ruda]
8pub fn index(lane: usize, item: usize, #[comptime] lanes: usize, #[comptime] items: usize,
9    #[comptime] striped: bool, #[comptime] warp_striped: bool, #[comptime] width: usize,
10) -> usize {
11    if warp_striped { lane / width * width * items + item * width + lane % width }
12    else if striped { item * lanes + lane }
13    else { lane * items + item }
14}
15
16/// Iterator load with explicit conversion and tail semantics. lane is a block
17/// lane or logical-warp lane, respectively; input/output must not overlap.
18#[ruda]
19pub fn load<T: RudaType, U: RudaType + Copy, I: RudaRead<T>, W: RudaWrite<U>, O: RudaUnaryOp<T, U>>(
20    input: &I, output: &mut W, convert: &O, offset: usize, valid: usize, lane: usize, padding: U,
21    #[comptime] lanes: usize, #[comptime] items: usize, #[comptime] striped: bool,
22    #[comptime] warp_striped: bool, #[comptime] width: usize, #[comptime] pad: bool,
23) {
24    #[unroll]
25    for item in 0..items {
26        let position = index(lane, item, lanes, items, striped, warp_striped, width);
27        if position < valid { output.write(item, convert.apply(input.read(offset + position))); }
28        else if pad { output.write(item, padding); }
29    }
30}
31
32#[ruda]
33pub fn store<T: RudaType, U: RudaType, I: RudaRead<T>, W: RudaWrite<U>, O: RudaUnaryOp<T, U>>(
34    input: &I, output: &mut W, convert: &O, offset: usize, valid: usize, lane: usize,
35    #[comptime] lanes: usize, #[comptime] items: usize, #[comptime] striped: bool,
36    #[comptime] warp_striped: bool, #[comptime] width: usize,
37) {
38    #[unroll]
39    for item in 0..items {
40        let position = index(lane, item, lanes, items, striped, warp_striped, width);
41        if position < valid { output.write(offset + position, convert.apply(input.read(item))); }
42    }
43}
44
45/// Collective layout conversion for scalar/record storage and arbitrary
46/// readers/writers. Scratch has lanes * items entries. For warp scope its
47/// base is disjoint for each logical group; all native subgroup lanes call it.
48#[ruda]
49pub fn exchange<T: RudaType, I: RudaRead<T>, W: RudaWrite<T>, S: RudaWrite<T>>(
50    input: &I, output: &mut W, scratch: &mut S, lane: usize, base: usize,
51    #[comptime] lanes: usize, #[comptime] items: usize, #[comptime] width: usize,
52    #[comptime] input_striped: bool, #[comptime] input_warp_striped: bool,
53    #[comptime] output_striped: bool, #[comptime] output_warp_striped: bool, #[comptime] warp_scope: bool,
54) {
55    #[unroll]
56    for item in 0..items {
57        let position = index(lane, item, lanes, items, input_striped, input_warp_striped, width);
58        scratch.write(base + position, input.read(item));
59    }
60    if warp_scope { sync_plane(); } else { sync_ruda(); }
61    #[unroll]
62    for item in 0..items {
63        let position = index(lane, item, lanes, items, output_striped, output_warp_striped, width);
64        output.write(item, scratch.read(base + position));
65    }
66    if warp_scope { sync_plane(); } else { sync_ruda(); }
67}
68
69/// Scatter-to-blocked/striped/warp-striped with optional valid flags and
70/// negative/out-of-range rank guarding. Written ranks must be unique. A caller
71/// using guarded holes must initialize scratch at every subsequently read hole.
72#[ruda]
73pub fn scatter<T: RudaType, I: RudaRead<T>, W: RudaWrite<T>, S: RudaWrite<T>>(
74    input: &I, ranks: &Array<i64>, flags: &Array<bool>, output: &mut W, scratch: &mut S,
75    lane: usize, base: usize, #[comptime] lanes: usize, #[comptime] items: usize, #[comptime] width: usize,
76    #[comptime] striped: bool, #[comptime] warp_striped: bool, #[comptime] guarded: bool,
77    #[comptime] flagged: bool, #[comptime] warp_scope: bool,
78) {
79    #[unroll]
80    for item in 0..items {
81        let rank = ranks[item];
82        let mut valid = true;
83        if guarded { valid = rank >= 0 && (rank as u64) < (lanes * items) as u64; }
84        if flagged { valid = valid && flags[item]; }
85        if valid { scratch.write(base + rank as usize, input.read(item)); }
86    }
87    if warp_scope { sync_plane(); } else { sync_ruda(); }
88    #[unroll]
89    for item in 0..items {
90        let position = index(lane, item, lanes, items, striped, warp_striped, width);
91        output.write(item, scratch.read(base + position));
92    }
93    if warp_scope { sync_plane(); } else { sync_ruda(); }
94}
95
96#[derive(Clone, Copy, RudaType, RudaLaunch)]
97pub struct RudaIdentity;
98
99impl<R: Runtime> Clone for RudaIdentityLaunch<R> {
100    fn clone(&self) -> Self { Self::new() }
101}
102
103#[ruda]
104impl<T: RudaType> RudaUnaryOp<T, T> for RudaIdentity {
105    fn apply(&self, value: T) -> T { value }
106}
ruprim/collective/tile.rs

ruprim/collective/
tile.rs
