ruprim/block/
mod.rs

1//! Block collectives. Every thread in the block participates in every call.
2
3use ruda_kernel::dsl as kernel_dsl;
4use ruda_kernel::dsl::prelude::*;
5use crate::collective::{RudaBinaryOp, RudaBinaryOpExpand};
6
7pub mod exchange;
8pub mod io;
9pub mod adjacent;
10pub mod sort;
11pub mod shuffle;
12pub mod radix;
13pub mod histogram;
14pub mod run_length;
15pub mod rank;
16pub mod topk;
17pub mod raking;
18pub mod record;
19
20#[ruda]
21pub trait RudaBlockPrefix<T: RudaType>: RudaType {
22    fn prefix(&mut self, aggregate: T) -> T;
23}
24
25/// Ordered inclusive scan of a blocked register tile.
26/// Scratch contains `threads * items_per_thread` elements. `valid_items` is
27/// uniform and positive; invalid output elements are unspecified.
28#[ruda]
29pub fn inclusive_scan<T: RudaPrimitive, O: RudaBinaryOp<T>>(
30    input: &Array<T>,
31    output: &mut Array<T>,
32    scratch: &mut SharedMemory<T>,
33    op: &O,
34    valid_items: usize,
35    #[comptime] threads: usize,
36    #[comptime] items_per_thread: usize,
37) {
38    let start = UNIT_POS as usize * items_per_thread;
39    #[unroll]
40    for item in 0..items_per_thread {
41        if start + item < valid_items { scratch[start + item] = input[item]; }
42    }
43    sync_ruda();
44    let mut distance = 1usize;
45    while distance < threads * items_per_thread {
46        #[unroll]
47        for item in 0..items_per_thread {
48            let index = start + item;
49            if index < valid_items {
50                let mut value = scratch[index];
51                if index >= distance {
52                    value = op.combine(scratch[index - distance], value);
53                }
54                output[item] = value;
55            }
56        }
57        sync_ruda();
58        #[unroll]
59        for item in 0..items_per_thread {
60            if start + item < valid_items { scratch[start + item] = output[item]; }
61        }
62        sync_ruda();
63        distance *= 2;
64    }
65    #[unroll]
66    for item in 0..items_per_thread {
67        if start + item < valid_items { output[item] = scratch[start + item]; }
68    }
69    sync_ruda();
70}
71
72/// Ordered exclusive scan. Returns the unseeded aggregate to every thread.
73#[ruda]
74pub fn exclusive_scan<T: RudaPrimitive, O: RudaBinaryOp<T>>(
75    input: &Array<T>,
76    output: &mut Array<T>,
77    scratch: &mut SharedMemory<T>,
78    initial: T,
79    op: &O,
80    valid_items: usize,
81    #[comptime] threads: usize,
82    #[comptime] items_per_thread: usize,
83) -> T {
84    inclusive_scan::<T, O>(input, output, scratch, op, valid_items, threads, items_per_thread);
85    let aggregate = scratch[valid_items - 1];
86    let start = UNIT_POS as usize * items_per_thread;
87    #[unroll]
88    for item in 0..items_per_thread {
89        let index = start + item;
90        let mut value = initial;
91        if index > 0 && index < valid_items {
92            value = op.combine(initial, scratch[index - 1]);
93        }
94        output[item] = value;
95    }
96    sync_ruda();
97    aggregate
98}
99
100/// Ordered tree reduction; the aggregate is broadcast to all block threads.
101/// `valid_items` must be positive. Input is in blocked register order.
102#[ruda]
103pub fn reduce<T: RudaPrimitive, O: RudaBinaryOp<T>>(
104    input: &Array<T>,
105    scratch: &mut SharedMemory<T>,
106    op: &O,
107    valid_items: usize,
108    #[comptime] threads: usize,
109    #[comptime] items_per_thread: usize,
110) -> T {
111    let start = UNIT_POS as usize * items_per_thread;
112    #[unroll]
113    for item in 0..items_per_thread {
114        if start + item < valid_items { scratch[start + item] = input[item]; }
115    }
116    sync_ruda();
117    let mut distance = 1usize;
118    while distance < threads * items_per_thread {
119        #[unroll]
120        for item in 0..items_per_thread {
121            let index = start + item;
122            if index % (distance * 2) == 0 && index + distance < valid_items {
123                scratch[index] = op.combine(scratch[index], scratch[index + distance]);
124            }
125        }
126        sync_ruda();
127        distance *= 2;
128    }
129    let aggregate = scratch[0];
130    sync_ruda();
131    aggregate
132}
133
134/// Inclusive/exclusive scan with a stateful block prefix callback. The first
135/// native subgroup invokes the callback, and thread zero's result is used.
136/// Scratch requires `threads * items_per_thread + 1` entries. Valid input is
137/// positive and callback state updates occur on the participating threads.
138#[ruda]
139pub fn scan_with_prefix<T: RudaPrimitive, O: RudaBinaryOp<T>, P: RudaBlockPrefix<T>>(
140    input: &Array<T>, output: &mut Array<T>, scratch: &mut SharedMemory<T>,
141    op: &O, prefix: &mut P, valid_items: usize,
142    #[comptime] threads: usize, #[comptime] items_per_thread: usize, #[comptime] exclusive: bool,
143) -> T {
144    inclusive_scan::<T, O>(input, output, scratch, op, valid_items, threads, items_per_thread);
145    let aggregate = scratch[valid_items - 1];
146    let prefix_index = threads * items_per_thread;
147    if UNIT_POS < PLANE_DIM {
148        let value = prefix.prefix(aggregate);
149        if UNIT_POS == 0 { scratch[prefix_index] = value; }
150    }
151    sync_ruda();
152    let seed = scratch[prefix_index];
153    #[unroll]
154    for item in 0..items_per_thread {
155        let index = UNIT_POS as usize * items_per_thread + item;
156        if index < valid_items {
157            let mut value = seed;
158            if exclusive {
159                if index > 0 { value = op.combine(seed, scratch[index - 1]); }
160            } else {
161                value = op.combine(seed, output[item]);
162            }
163            output[item] = value;
164        }
165    }
166    sync_ruda();
167    aggregate
168}
ruprim/block/mod.rs

ruprim/block/
mod.rs
