ruda_kernel/dsl/post_processing/
checked_io.rs

1use alloc::{string::String, vec::Vec};
2use ruda_core::ir::{Allocator, Instruction, ManagedVariable, Operation, Operator, Processor, Scope};
3use ruda::runtime::server::ExecutionMode;
4
5use crate::dsl::{
6    define_scalar, define_size,
7    io::{
8        expand_checked_index_assign, expand_validate_index_assign, read_tensor_atomic_checked,
9        read_tensor_atomic_validate, read_tensor_checked, read_tensor_validate,
10    },
11    prelude::Vector,
12};
13
14define_scalar!(ElemA);
15define_size!(SizeA);
16
17#[derive(new, Debug)]
18pub struct CheckedIoProcessor {
19    mode: ExecutionMode,
20    kernel_name: String,
21}
22
23impl Processor for CheckedIoProcessor {
24    fn transform(
25        &self,
26        processing: ruda_core::ir::ScopeProcessing,
27        allocator: Allocator,
28    ) -> ruda_core::ir::ScopeProcessing {
29        match self.mode {
30            ExecutionMode::Checked => self.transform_checked(processing, allocator),
31            ExecutionMode::Unchecked => processing,
32            ExecutionMode::Validate => self.transform_validate(processing, allocator),
33        }
34    }
35}
36
37impl CheckedIoProcessor {
38    fn transform_checked(
39        &self,
40        mut processing: ruda_core::ir::ScopeProcessing,
41        allocator: Allocator,
42    ) -> ruda_core::ir::ScopeProcessing {
43        let mut instructions = Vec::new();
44        core::mem::swap(&mut processing.instructions, &mut instructions);
45
46        for instruction in instructions {
47            let lowered = matches!(instruction.operation, Operation::TypedMemory(_))
48                .then(|| instruction.clone().lower_typed_memory().expect("valid typed memory instruction"));
49            let operation = lowered.as_ref().map_or(&instruction.operation, |instruction| &instruction.operation);
50            if let Operation::Operator(operator) = operation {
51                match operator {
52                    Operator::Index(op) => {
53                        let has_length = op.list.has_length();
54
55                        if has_length {
56                            let list = ManagedVariable::Plain(op.list);
57                            let index = ManagedVariable::Plain(op.index);
58                            let mut scope = Scope::root(false)
59                                .with_allocator(allocator.clone())
60                                .with_types(processing.typemap.clone());
61                            scope.register_type::<ElemA>(op.list.storage_type());
62                            scope.register_size::<SizeA>(op.list.vector_size());
63
64                            let input = if op.list.ty.is_atomic() {
65                                // Atomic can't really be checked, since the pointer needs to be
66                                // valid, so the kernel will probably not output the correct value if
67                                // not manually checked later, but will at least avoid out-of-bounds
68                                // memory access.
69                                read_tensor_atomic_checked::expand::<ElemA>(
70                                    &mut scope,
71                                    list.into(),
72                                    index.into(),
73                                    op.unroll_factor,
74                                )
75                                .expand
76                            } else {
77                                read_tensor_checked::expand::<Vector<ElemA, SizeA>>(
78                                    &mut scope,
79                                    list.into(),
80                                    index.into(),
81                                    op.unroll_factor,
82                                )
83                                .expand
84                            };
85                            let tmp_processing = scope.process([]);
86
87                            for inst in tmp_processing.instructions {
88                                processing.instructions.push(instruction.retag_typed_memory(inst));
89                            }
90                            for var in tmp_processing.variables {
91                                processing.variables.push(var);
92                            }
93
94                            processing
95                                .instructions
96                                .push(Instruction::new(Operation::Copy(*input), instruction.out()));
97                            continue;
98                        }
99                    }
100                    Operator::IndexAssign(op) => {
101                        let out = instruction.out();
102
103                        if out.has_length() {
104                            let mut scope = Scope::root(false)
105                                .with_allocator(allocator.clone())
106                                .with_types(processing.typemap.clone());
107                            expand_checked_index_assign(
108                                &mut scope,
109                                op.index,
110                                op.value,
111                                out,
112                                op.unroll_factor,
113                            );
114
115                            let tmp_processing = scope.process([]);
116
117                            for inst in tmp_processing.instructions {
118                                processing.instructions.push(instruction.retag_typed_memory(inst));
119                            }
120                            for var in tmp_processing.variables {
121                                processing.variables.push(var);
122                            }
123
124                            continue;
125                        }
126                    }
127                    _ => {}
128                }
129            }
130
131            // When we have nothing to do.
132            processing.instructions.push(instruction);
133        }
134        processing
135    }
136
137    fn transform_validate(
138        &self,
139        mut processing: ruda_core::ir::ScopeProcessing,
140        allocator: Allocator,
141    ) -> ruda_core::ir::ScopeProcessing {
142        let mut instructions = Vec::new();
143        core::mem::swap(&mut processing.instructions, &mut instructions);
144
145        for instruction in instructions {
146            let lowered = matches!(instruction.operation, Operation::TypedMemory(_))
147                .then(|| instruction.clone().lower_typed_memory().expect("valid typed memory instruction"));
148            let operation = lowered.as_ref().map_or(&instruction.operation, |instruction| &instruction.operation);
149            if let Operation::Operator(operator) = operation {
150                match operator {
151                    Operator::Index(op) => {
152                        let has_length = op.list.has_length();
153
154                        if has_length {
155                            let list = ManagedVariable::Plain(op.list);
156                            let index = ManagedVariable::Plain(op.index);
157                            let mut scope = Scope::root(false)
158                                .with_allocator(allocator.clone())
159                                .with_types(processing.typemap.clone());
160                            scope.register_type::<ElemA>(op.list.storage_type());
161                            scope.register_size::<SizeA>(op.list.vector_size());
162
163                            let input = if op.list.ty.is_atomic() {
164                                // Atomic can't really be checked, since the pointer needs to be
165                                // valid, so the kernel will probably not output the correct value if
166                                // not manually checked later, but will at least avoid out-of-bounds
167                                // memory access.
168                                read_tensor_atomic_validate::expand::<ElemA>(
169                                    &mut scope,
170                                    list.into(),
171                                    index.into(),
172                                    op.unroll_factor,
173                                    self.kernel_name.clone(),
174                                )
175                                .expand
176                            } else {
177                                read_tensor_validate::expand::<Vector<ElemA, SizeA>>(
178                                    &mut scope,
179                                    list.into(),
180                                    index.into(),
181                                    op.unroll_factor,
182                                    self.kernel_name.clone(),
183                                )
184                                .expand
185                            };
186                            let tmp_processing = scope.process([]);
187
188                            for inst in tmp_processing.instructions {
189                                processing.instructions.push(instruction.retag_typed_memory(inst));
190                            }
191                            for var in tmp_processing.variables {
192                                processing.variables.push(var);
193                            }
194
195                            processing
196                                .instructions
197                                .push(Instruction::new(Operation::Copy(*input), instruction.out()));
198                            continue;
199                        }
200                    }
201                    Operator::IndexAssign(op) => {
202                        let out = instruction.out();
203
204                        if out.has_length() {
205                            let mut scope = Scope::root(false)
206                                .with_allocator(allocator.clone())
207                                .with_types(processing.typemap.clone());
208                            expand_validate_index_assign(
209                                &mut scope,
210                                op.index,
211                                op.value,
212                                out,
213                                op.unroll_factor,
214                                &self.kernel_name,
215                            );
216
217                            let tmp_processing = scope.process([]);
218
219                            for inst in tmp_processing.instructions {
220                                processing.instructions.push(instruction.retag_typed_memory(inst));
221                            }
222                            for var in tmp_processing.variables {
223                                processing.variables.push(var);
224                            }
225
226                            continue;
227                        }
228                    }
229                    _ => {}
230                }
231            }
232
233            // When we have nothing to do.
234            processing.instructions.push(instruction);
235        }
236        processing
237    }
238}
ruda_kernel/dsl/post_processing/checked_io.rs

ruda_kernel/dsl/post_processing/
checked_io.rs
