cubecl_core/post_processing/
checked_io.rs

1use alloc::{vec, vec::Vec};
2
3use alloc::string::String;
4use cubecl_ir::{GlobalState, Instruction, Memory, Operation, Scope};
5use cubecl_runtime::server::ExecutionMode;
6
7use crate::{
8    io::*,
9    post_processing::{
10        analysis_helper::GlobalAnalyses, util::AtomicCounter, visitor::InstructionVisitor,
11    },
12};
13
14#[derive(new, Debug)]
15pub struct CheckedIoVisitor {
16    mode: ExecutionMode,
17    kernel_name: String,
18}
19
20impl CheckedIoVisitor {
21    pub fn apply(&mut self, scope: &Scope) {
22        let changes = AtomicCounter::new(0);
23        // We don't care about pointer sources or used variables at this point
24        let analyses = GlobalAnalyses::default();
25        self.visit_scope(scope, &analyses, &changes);
26    }
27}
28
29impl InstructionVisitor for CheckedIoVisitor {
30    fn visit_instruction(
31        &mut self,
32        instruction: Instruction,
33        global_state: &GlobalState,
34        _analyses: &GlobalAnalyses,
35        _changes: &AtomicCounter,
36    ) -> Vec<Instruction> {
37        match self.mode {
38            ExecutionMode::Checked => self.transform_checked(instruction, global_state),
39            ExecutionMode::Validate => self.transform_validate(instruction, global_state),
40            ExecutionMode::Unchecked => vec![instruction],
41        }
42    }
43}
44
45impl CheckedIoVisitor {
46    fn transform_checked(
47        &self,
48        instruction: Instruction,
49        global_state: &GlobalState,
50    ) -> Vec<Instruction> {
51        if let Operation::Memory(memory) = &instruction.operation {
52            match memory {
53                Memory::Index(op) if op.checked => {
54                    let has_length = op.list.has_buffer_length();
55
56                    if has_length {
57                        let list = op.list;
58                        let index = op.index;
59                        let scope = Scope::root(false).with_global_state(global_state.clone());
60
61                        expand_checked_index(
62                            &scope,
63                            list,
64                            index,
65                            instruction.out(),
66                            op.unroll_factor,
67                        );
68
69                        return scope.take_instructions();
70                    }
71                }
72                _ => {}
73            }
74        }
75        vec![instruction]
76    }
77
78    fn transform_validate(
79        &self,
80        instruction: Instruction,
81        global_state: &GlobalState,
82    ) -> Vec<Instruction> {
83        if let Operation::Memory(memory) = &instruction.operation {
84            match memory {
85                Memory::Index(op) if op.checked => {
86                    let has_length = op.list.has_buffer_length();
87
88                    if has_length {
89                        let list = op.list;
90                        let index = op.index;
91                        let scope = Scope::root(false).with_global_state(global_state.clone());
92
93                        expand_validate_index(
94                            &scope,
95                            list,
96                            index,
97                            instruction.out(),
98                            op.unroll_factor,
99                            &self.kernel_name,
100                        );
101
102                        return scope.take_instructions();
103                    }
104                }
105                _ => {}
106            }
107        }
108        vec![instruction]
109    }
110}
cubecl_core/post_processing/checked_io.rs

cubecl_core/post_processing/
checked_io.rs
