cubecl_core/post_processing/
unroll.rs

1use alloc::{vec, vec::Vec};
2use cubecl_ir::{
3    Allocator, Arithmetic, BinaryOperands, CoopMma, GlobalState, IndexOperands, Instruction,
4    MatrixLayout, Memory, Metadata, Operation, OperationReflect, Operator, Scope, Type, Value,
5    ValueKind, VectorInsertOperands, VectorSize,
6};
7use hashbrown::HashMap;
8
9use crate::post_processing::{
10    analysis_helper::GlobalAnalyses,
11    util::AtomicCounter,
12    visitor::{InstructionVisitor, visit_scope},
13};
14
15/// The action that should be performed on an instruction, returned by ``IrTransformer::maybe_transform``
16pub enum TransformAction {
17    /// The transformer doesn't apply to this instruction
18    Ignore,
19    /// Replace this instruction with one or more other instructions
20    Replace(Vec<Instruction>),
21}
22
23#[derive(Debug)]
24pub struct UnrollVisitor {
25    max_vector_size: VectorSize,
26    mappings: Mappings,
27}
28
29#[derive(Default, Debug)]
30struct Mappings(HashMap<Value, Vec<Value>>);
31
32impl Mappings {
33    fn get(
34        &mut self,
35        alloc: &Allocator,
36        val: Value,
37        unroll_factor: usize,
38        vector_size: VectorSize,
39    ) -> Vec<Value> {
40        self.0
41            .entry(val)
42            .or_insert_with(|| create_unrolled(alloc, &val, vector_size, unroll_factor))
43            .to_vec()
44    }
45}
46
47impl UnrollVisitor {
48    pub fn new(max_vector_size: VectorSize) -> Self {
49        Self {
50            max_vector_size,
51            mappings: Default::default(),
52        }
53    }
54
55    pub fn apply(&mut self, scope: &Scope) {
56        let changes = AtomicCounter::new(0);
57        // We don't care about pointer sources or used variables at this point
58        let analyses = GlobalAnalyses::default();
59        self.visit_scope(scope, &analyses, &changes);
60    }
61}
62
63impl InstructionVisitor for UnrollVisitor {
64    fn visit_instruction(
65        &mut self,
66        instruction: Instruction,
67        global_state: &GlobalState,
68        _analyses: &GlobalAnalyses,
69        _changes: &AtomicCounter,
70    ) -> Vec<Instruction> {
71        match self.maybe_transform(&global_state.borrow().allocator, &instruction) {
72            TransformAction::Ignore => {
73                vec![instruction]
74            }
75            TransformAction::Replace(replacement) => replacement,
76        }
77    }
78
79    fn visit_scope(&mut self, scope: &Scope, analyses: &GlobalAnalyses, changes: &AtomicCounter) {
80        visit_scope(self, scope, analyses, changes);
81    }
82}
83
84impl UnrollVisitor {
85    fn maybe_transform(&mut self, alloc: &Allocator, inst: &Instruction) -> TransformAction {
86        if matches!(inst.operation, Operation::Marker(_)) {
87            return TransformAction::Ignore;
88        }
89
90        if inst.operation.args().is_none() {
91            // Detect unhandled ops that can't be reflected
92            match &inst.operation {
93                Operation::CoopMma(op) => match op {
94                    // Stride is in scalar elems
95                    CoopMma::Load {
96                        ptr,
97                        stride,
98                        layout,
99                    } if ptr.vector_size() > self.max_vector_size => {
100                        return TransformAction::Replace(self.transform_cmma_load(
101                            alloc,
102                            inst.out(),
103                            ptr,
104                            stride,
105                            layout,
106                        ));
107                    }
108                    CoopMma::Store {
109                        mat,
110                        stride,
111                        destination,
112                        layout,
113                    } if destination.vector_size() > self.max_vector_size => {
114                        return TransformAction::Replace(self.transform_cmma_store(
115                            alloc,
116                            mat,
117                            stride,
118                            destination,
119                            layout,
120                        ));
121                    }
122                    _ => return TransformAction::Ignore,
123                },
124                Operation::TensorIndexing(_) => return TransformAction::Ignore,
125                Operation::Branch(_) | Operation::NonSemantic(_) | Operation::Marker(_) => {
126                    return TransformAction::Ignore;
127                }
128                Operation::Operator(Operator::ReadBuiltin(_)) => {
129                    return TransformAction::Ignore;
130                }
131                Operation::DeclareVariable {
132                    value_ty: Type::Array(inner_ty, len),
133                    addr_space,
134                    alignment,
135                } => {
136                    if inner_ty.vector_size() > self.max_vector_size {
137                        let unroll_factor = inner_ty.vector_size() / self.max_vector_size;
138                        let vector_size = self.max_vector_size;
139                        let inner_ty = inner_ty.with_vector_size(vector_size);
140                        let new_ty = Type::Array(inner_ty.intern(), *len * unroll_factor);
141                        let new_ptr_ty = Type::Pointer(new_ty.intern(), *addr_space);
142                        let mut out = inst.out.unwrap();
143                        out.ty = new_ptr_ty;
144
145                        return TransformAction::Replace(vec![Instruction::new(
146                            Operation::DeclareVariable {
147                                value_ty: new_ty,
148                                addr_space: *addr_space,
149                                alignment: *alignment,
150                            },
151                            out,
152                        )]);
153                    } else {
154                        return TransformAction::Ignore;
155                    }
156                }
157                Operation::DeclareVariable {
158                    value_ty,
159                    addr_space,
160                    alignment,
161                } => {
162                    if value_ty.vector_size() > self.max_vector_size {
163                        let unroll_factor = value_ty.vector_size() / self.max_vector_size;
164                        let vector_size = self.max_vector_size;
165                        let value_ty = value_ty.with_vector_size(vector_size);
166                        let out = self
167                            .mappings
168                            .get(alloc, inst.out(), unroll_factor, vector_size);
169                        let declare = |out| {
170                            Instruction::new(
171                                Operation::DeclareVariable {
172                                    value_ty,
173                                    addr_space: *addr_space,
174                                    alignment: *alignment,
175                                },
176                                out,
177                            )
178                        };
179                        return TransformAction::Replace(out.into_iter().map(declare).collect());
180                    } else {
181                        return TransformAction::Ignore;
182                    }
183                }
184                other => {
185                    panic!(
186                        "Need special handling for unrolling non-reflectable operations.\nFound: {other}"
187                    )
188                }
189            }
190        }
191
192        let args = inst.operation.args().unwrap_or_default();
193        if (inst.out.is_some() && inst.ty().vector_size() > self.max_vector_size)
194            || args
195                .iter()
196                .any(|arg| arg.vector_size() > self.max_vector_size)
197        {
198            let vector_size = max_vector_size(&inst.out, &args);
199            let unroll_factor = vector_size / self.max_vector_size;
200
201            match &inst.operation {
202                Operation::Memory(Memory::Index(op)) => TransformAction::Replace(
203                    self.transform_array_index(alloc, inst.out(), op, Memory::Index, unroll_factor),
204                ),
205                Operation::Operator(Operator::ExtractComponent(op)) => TransformAction::Replace(
206                    self.transform_composite_extract(alloc, inst.out(), op, unroll_factor),
207                ),
208                Operation::Operator(Operator::InsertComponent(op)) => TransformAction::Replace(
209                    self.transform_composite_insert(alloc, inst.out(), op, unroll_factor),
210                ),
211                Operation::Metadata(op) => {
212                    TransformAction::Replace(self.transform_metadata(inst.out(), op, args))
213                }
214                _ => {
215                    TransformAction::Replace(self.transform_basic(alloc, inst, args, unroll_factor))
216                }
217            }
218        } else {
219            TransformAction::Ignore
220        }
221    }
222
223    /// Transform CMMA load offset and array
224    fn transform_cmma_load(
225        &mut self,
226        alloc: &Allocator,
227        out: Value,
228        ptr: &Value,
229        stride: &Value,
230        layout: &Option<MatrixLayout>,
231    ) -> Vec<Instruction> {
232        let vector_size = ptr.vector_size();
233        let unroll_factor = vector_size / self.max_vector_size;
234
235        let ptr = self
236            .mappings
237            .get(alloc, *ptr, unroll_factor, self.max_vector_size);
238        let out = unroll_array(out, self.max_vector_size, unroll_factor);
239
240        let load = Instruction::new(
241            Operation::CoopMma(CoopMma::Load {
242                ptr: ptr[0],
243                stride: *stride,
244                layout: *layout,
245            }),
246            out,
247        );
248        vec![load]
249    }
250
251    /// Transform CMMA store offset and array
252    fn transform_cmma_store(
253        &mut self,
254        alloc: &Allocator,
255        mat: &Value,
256        stride: &Value,
257        destination: &Value,
258        layout: &MatrixLayout,
259    ) -> Vec<Instruction> {
260        let vector_size = destination.vector_size();
261        let unroll_factor = vector_size / self.max_vector_size;
262
263        let destination =
264            self.mappings
265                .get(alloc, *destination, unroll_factor, self.max_vector_size);
266
267        let store = Instruction::no_out(Operation::CoopMma(CoopMma::Store {
268            mat: *mat,
269            stride: *stride,
270            destination: destination[0],
271            layout: *layout,
272        }));
273        vec![store]
274    }
275
276    /// Transforms indexing into multiple index operations, each offset by 1 from the base. The base
277    /// is also multiplied by the unroll factor to compensate for the lower actual vectorization.
278    fn transform_array_index(
279        &mut self,
280        alloc: &Allocator,
281        out: Value,
282        op: &IndexOperands,
283        operator: impl Fn(IndexOperands) -> Memory,
284        unroll_factor: usize,
285    ) -> Vec<Instruction> {
286        let (mul, start_idx) = mul_index(alloc, op.index, unroll_factor);
287        let mut indices = (0..unroll_factor).map(|i| add_index(alloc, start_idx, i));
288
289        let list = unroll_array(op.list, self.max_vector_size, unroll_factor);
290
291        let out = self
292            .mappings
293            .get(alloc, out, unroll_factor, self.max_vector_size);
294        let mut instructions = vec![mul];
295        instructions.extend((0..unroll_factor).flat_map(|i| {
296            let (add, idx) = indices.next().unwrap();
297            let index = Instruction::new(
298                operator(IndexOperands {
299                    list,
300                    index: idx,
301                    unroll_factor,
302                    checked: op.checked,
303                }),
304                out[i],
305            );
306            [add, index]
307        }));
308
309        instructions
310    }
311
312    /// Transforms a composite index (i.e. `Vector`) that always returns a scalar. Translates the index
313    /// to a local index and an unroll index, then indexes the proper variable. Note that this requires
314    /// the index to be constant - it needs to be decomposed at compile time, otherwise it wouldn't
315    /// work.
316    fn transform_composite_extract(
317        &mut self,
318        alloc: &Allocator,
319        out: Value,
320        op: &BinaryOperands,
321        unroll_factor: usize,
322    ) -> Vec<Instruction> {
323        let index = op
324            .rhs
325            .as_const()
326            .expect("Can't unroll non-constant vector index")
327            .as_usize();
328
329        let unroll_idx = index / self.max_vector_size;
330        let sub_idx = index % self.max_vector_size;
331
332        let value = self
333            .mappings
334            .get(alloc, op.lhs, unroll_factor, self.max_vector_size);
335
336        vec![Instruction::new(
337            Operator::ExtractComponent(BinaryOperands {
338                lhs: value[unroll_idx],
339                rhs: sub_idx.into(),
340            }),
341            out,
342        )]
343    }
344
345    /// Transforms a composite index assign (i.e. `Vector`) that always takes a scalar. Translates the index
346    /// to a local index and an unroll index, then indexes the proper variable. Note that this requires
347    /// the index to be constant - it needs to be decomposed at compile time, otherwise it wouldn't
348    /// work.
349    fn transform_composite_insert(
350        &mut self,
351        alloc: &Allocator,
352        out: Value,
353        op: &VectorInsertOperands,
354        unroll_factor: usize,
355    ) -> Vec<Instruction> {
356        let index = op
357            .index
358            .as_const()
359            .expect("Can't unroll non-constant vector index")
360            .as_usize();
361
362        let unroll_idx = index / self.max_vector_size;
363        let sub_idx = index % self.max_vector_size;
364
365        let vector = self
366            .mappings
367            .get(alloc, op.vector, unroll_factor, self.max_vector_size);
368        let out = self
369            .mappings
370            .get(alloc, out, unroll_factor, self.max_vector_size);
371        (0..unroll_factor)
372            .map(|i| {
373                if i == unroll_idx {
374                    Instruction::new(
375                        Operator::InsertComponent(VectorInsertOperands {
376                            vector: vector[i],
377                            index: sub_idx.into(),
378                            value: op.value,
379                        }),
380                        out[i],
381                    )
382                } else {
383                    Instruction::new(Operation::Copy(vector[i]), out[i])
384                }
385            })
386            .collect()
387    }
388
389    /// Transforms metadata by just replacing the type of the buffer. The values are already
390    /// properly calculated on the CPU.
391    fn transform_metadata(&self, out: Value, op: &Metadata, args: Vec<Value>) -> Vec<Instruction> {
392        let op_code = op.op_code();
393        let args = args
394            .into_iter()
395            .map(|mut val| {
396                if val.vector_size() > self.max_vector_size {
397                    val.ty = val.ty.with_vector_size(self.max_vector_size);
398                }
399                val
400            })
401            .collect::<Vec<_>>();
402        let operation = Metadata::from_code_and_args(op_code, &args).unwrap();
403        vec![Instruction::new(operation, out)]
404    }
405
406    /// Transforms generic instructions, i.e. comparison, arithmetic. Unrolls each vectorized variable
407    /// to `unroll_factor` replacements, and executes the operation `unroll_factor` times.
408    fn transform_basic(
409        &mut self,
410        alloc: &Allocator,
411        inst: &Instruction,
412        args: Vec<Value>,
413        unroll_factor: usize,
414    ) -> Vec<Instruction> {
415        let op_code = inst.operation.op_code();
416        let out = inst.out.map(|out| {
417            self.mappings
418                .get(alloc, out, unroll_factor, self.max_vector_size)
419        });
420        let args = args
421            .into_iter()
422            .map(|arg| {
423                if arg.vector_size() > 1 {
424                    self.mappings
425                        .get(alloc, arg, unroll_factor, self.max_vector_size)
426                } else {
427                    // Preserve scalars
428                    vec![arg]
429                }
430            })
431            .collect::<Vec<_>>();
432
433        (0..unroll_factor)
434            .map(|i| {
435                let out = out.as_ref().map(|out| out[i]);
436                let args = args
437                    .iter()
438                    .map(|arg| if arg.len() == 1 { arg[0] } else { arg[i] })
439                    .collect::<Vec<_>>();
440                let operation = Operation::from_code_and_args(op_code, &args)
441                    .expect("Failed to reconstruct operation");
442                Instruction {
443                    out,
444                    source_loc: inst.source_loc.clone(),
445                    modes: inst.modes,
446                    operation,
447                }
448            })
449            .collect()
450    }
451}
452
453fn max_vector_size(out: &Option<Value>, args: &[Value]) -> VectorSize {
454    let vector_size = args.iter().map(|it| it.vector_size()).max().unwrap();
455    vector_size.max(out.map(|out| out.vector_size()).unwrap_or(1))
456}
457
458fn create_unrolled(
459    allocator: &Allocator,
460    val: &Value,
461    max_vector_size: VectorSize,
462    unroll_factor: usize,
463) -> Vec<Value> {
464    // Preserve scalars
465    if val.vector_size() == 1 {
466        return vec![*val; unroll_factor];
467    }
468
469    let item = val.ty.with_vector_size(max_vector_size);
470    (0..unroll_factor)
471        .map(|_| match val.kind {
472            ValueKind::Value { .. } => allocator.create_value(item),
473            other => panic!("Out must be local, found {other:?}"),
474        })
475        .collect()
476}
477
478fn add_index(alloc: &Allocator, idx: Value, i: usize) -> (Instruction, Value) {
479    let add_idx = alloc.create_value(idx.ty);
480    let add = Instruction::new(
481        Arithmetic::Add(BinaryOperands {
482            lhs: idx,
483            rhs: i.into(),
484        }),
485        add_idx,
486    );
487    (add, add_idx)
488}
489
490fn mul_index(alloc: &Allocator, idx: Value, unroll_factor: usize) -> (Instruction, Value) {
491    let mul_idx = alloc.create_value(idx.ty);
492    let mul = Instruction::new(
493        Arithmetic::Mul(BinaryOperands {
494            lhs: idx,
495            rhs: unroll_factor.into(),
496        }),
497        mul_idx,
498    );
499    (mul, mul_idx)
500}
501
502fn unroll_array(mut val: Value, max_vector_size: VectorSize, factor: usize) -> Value {
503    val.ty = val.ty.with_vector_size(max_vector_size);
504
505    if let Type::Array(_, size) = &mut val.ty {
506        *size *= factor;
507    }
508
509    val
510}
cubecl_core/post_processing/unroll.rs

cubecl_core/post_processing/
unroll.rs
