ruprim/block/
adjacent.rs

1use ruda_kernel::dsl as kernel_dsl;
2use ruda_kernel::dsl::prelude::*;
3use crate::collective::{RudaBinaryOp, RudaBinaryOpExpand, RudaKeyEqual, RudaKeyEqualExpand};
4use crate::collective::record::{RudaRead, RudaReadExpand, RudaWrite, RudaWriteExpand};
5
6#[ruda]
7pub fn difference_access<T: RudaType<ExpandType: Assign> + Copy, I: RudaRead<T>, W: RudaWrite<T>, S: RudaWrite<T>, O: RudaBinaryOp<T>>(
8    input: &I, output: &mut W, scratch: &mut S, op: &O, neighbour: T, valid: usize,
9    #[comptime] items: usize, #[comptime] right: bool, #[comptime] has_neighbour: bool,
10) {
11    let start = UNIT_POS as usize * items;
12    #[unroll]
13    for item in 0..items {
14        if start + item < valid { scratch.write(start + item, input.read(item)); }
15    }
16    sync_ruda();
17    #[unroll]
18    for item in 0..items {
19        let index = start + item;
20        if index < valid {
21            let current = scratch.read(index);
22            let mut value = current;
23            let boundary = if right { index + 1 == valid } else { index == 0 };
24            if boundary {
25                if has_neighbour { value = op.combine(current, neighbour); }
26            } else {
27                let next = if right { index + 1 } else { index - 1 };
28                value = op.combine(current, scratch.read(next));
29            }
30            output.write(item, value);
31        }
32    }
33    sync_ruda();
34}
35
36#[ruda]
37pub trait RudaDiscontinuity<T: RudaType>: RudaType {
38    fn flag(&self, left: T, right: T, right_index: usize) -> bool;
39}
40
41#[ruda]
42pub fn discontinuity_access<T: RudaType<ExpandType: Assign> + Copy,
43    I: RudaRead<T>, H: RudaWrite<bool>, W: RudaWrite<bool>, S: RudaWrite<T>, O: RudaDiscontinuity<T>>(
44    input: &I, heads: &mut H, tails: &mut W, scratch: &mut S,
45    op: &O, predecessor: T, successor: T, valid: usize,
46    #[comptime] items: usize, #[comptime] has_predecessor: bool, #[comptime] has_successor: bool,
47) {
48    let start = UNIT_POS as usize * items;
49    #[unroll]
50    for item in 0..items {
51        if start + item < valid { scratch.write(start + item, input.read(item)); }
52    }
53    sync_ruda();
54    #[unroll]
55    for item in 0..items {
56        let index = start + item;
57        let mut head = false;
58        let mut tail = false;
59        if index < valid {
60            let value = scratch.read(index);
61            if index == 0 {
62                head = true;
63                if has_predecessor { head = op.flag(predecessor, value, index); }
64            } else { head = op.flag(scratch.read(index - 1), value, index); }
65            if index + 1 == valid {
66                tail = true;
67                if has_successor { tail = op.flag(value, successor, index + 1); }
68            } else { tail = op.flag(value, scratch.read(index + 1), index + 1); }
69        }
70        heads.write(item, head);
71        tails.write(item, tail);
72    }
73    sync_ruda();
74}
75
76/// Index-aware head/tail flags; right_index is the right item's tile rank,
77/// including valid_items for an external successor.
78#[ruda]
79pub fn discontinuity_indexed<T: RudaPrimitive, O: RudaDiscontinuity<T>>(
80    input: &Array<T>, heads: &mut Array<bool>, tails: &mut Array<bool>, scratch: &mut SharedMemory<T>,
81    op: &O, predecessor: T, successor: T, valid_items: usize,
82    #[comptime] items_per_thread: usize, #[comptime] has_predecessor: bool, #[comptime] has_successor: bool,
83) {
84    let start = UNIT_POS as usize * items_per_thread;
85    #[unroll]
86    for item in 0..items_per_thread {
87        if start + item < valid_items { scratch[start + item] = input[item]; }
88    }
89    sync_ruda();
90    #[unroll]
91    for item in 0..items_per_thread {
92        let index = start + item;
93        let mut head = false;
94        let mut tail = false;
95        if index < valid_items {
96            if index == 0 {
97                head = true;
98                if has_predecessor { head = op.flag(predecessor, scratch[index], index); }
99            } else { head = op.flag(scratch[index - 1], scratch[index], index); }
100            if index + 1 == valid_items {
101                tail = true;
102                if has_successor { tail = op.flag(scratch[index], successor, index + 1); }
103            } else { tail = op.flag(scratch[index], scratch[index + 1], index + 1); }
104        }
105        heads[item] = head;
106        tails[item] = tail;
107    }
108    sync_ruda();
109}
110
111/// Adjacent transform in blocked order. With no external neighbour, preserve
112/// the boundary input. The operator receives (current, neighbouring).
113#[ruda]
114pub fn difference<T: RudaPrimitive, O: RudaBinaryOp<T>>(
115    input: &Array<T>,
116    output: &mut Array<T>,
117    scratch: &mut SharedMemory<T>,
118    op: &O,
119    neighbour: T,
120    valid_items: usize,
121    #[comptime] items_per_thread: usize,
122    #[comptime] right: bool,
123    #[comptime] has_neighbour: bool,
124) {
125    let start = UNIT_POS as usize * items_per_thread;
126    #[unroll]
127    for item in 0..items_per_thread {
128        if start + item < valid_items {
129            scratch[start + item] = input[item];
130        }
131    }
132    sync_ruda();
133    #[unroll]
134    for item in 0..items_per_thread {
135        let index = start + item;
136        if index < valid_items {
137            let current = scratch[index];
138            let boundary = if right { index + 1 == valid_items } else { index == 0 };
139            let mut value = current;
140            if boundary {
141                if has_neighbour { value = op.combine(current, neighbour); }
142            } else {
143                let adjacent = if right { scratch[index + 1] } else { scratch[index - 1] };
144                value = op.combine(current, adjacent);
145            }
146            output[item] = value;
147        }
148    }
149    sync_ruda();
150}
151
152/// Flag heads and tails of equal-key runs. Optional neighbours connect tiles.
153#[ruda]
154pub fn discontinuity<T: RudaPrimitive, E: RudaKeyEqual<T>>(
155    input: &Array<T>,
156    heads: &mut Array<bool>,
157    tails: &mut Array<bool>,
158    scratch: &mut SharedMemory<T>,
159    equal: &E,
160    predecessor: T,
161    successor: T,
162    valid_items: usize,
163    #[comptime] items_per_thread: usize,
164    #[comptime] has_predecessor: bool,
165    #[comptime] has_successor: bool,
166) {
167    let start = UNIT_POS as usize * items_per_thread;
168    #[unroll]
169    for item in 0..items_per_thread {
170        if start + item < valid_items { scratch[start + item] = input[item]; }
171    }
172    sync_ruda();
173    #[unroll]
174    for item in 0..items_per_thread {
175        let index = start + item;
176        let mut head = false;
177        let mut tail = false;
178        if index < valid_items {
179            let value = scratch[index];
180            if index > 0 {
181                head = !equal.equal(scratch[index - 1], value);
182            } else {
183                head = true;
184                if has_predecessor { head = !equal.equal(predecessor, value); }
185            }
186            if index + 1 < valid_items {
187                tail = !equal.equal(value, scratch[index + 1]);
188            } else {
189                tail = true;
190                if has_successor { tail = !equal.equal(value, successor); }
191            }
192        }
193        heads[item] = head;
194        tails[item] = tail;
195    }
196    sync_ruda();
197}
ruprim/block/adjacent.rs

ruprim/block/
adjacent.rs
