ruda_kernel/dsl/post_processing/
unroll.rs

1use alloc::{vec, vec::Vec};
2use ruda_core::ir::{
3    Allocator, Arithmetic, BinaryOperator, Branch, CoopMma, CopyMemoryBulkOperator,
4    IndexAssignOperator, IndexOperator, Instruction, ManagedVariable, MatrixLayout, Metadata,
5    Operation, OperationReflect, Operator, Processor, ScopeProcessing, Type, Variable,
6    VariableKind, VectorSize,
7};
8use hashbrown::HashMap;
9
10/// The action that should be performed on an instruction, returned by ``IrTransformer::maybe_transform``
11pub enum TransformAction {
12    /// The transformer doesn't apply to this instruction
13    Ignore,
14    /// Replace this instruction with one or more other instructions
15    Replace(Vec<Instruction>),
16}
17
18#[derive(new, Debug)]
19pub struct UnrollProcessor {
20    max_vector_size: VectorSize,
21}
22
23struct Mappings(HashMap<Variable, Vec<ManagedVariable>>);
24
25impl Mappings {
26    fn get(
27        &mut self,
28        alloc: &Allocator,
29        var: Variable,
30        unroll_factor: usize,
31        vector_size: VectorSize,
32    ) -> Vec<Variable> {
33        self.0
34            .entry(var)
35            .or_insert_with(|| create_unrolled(alloc, &var, vector_size, unroll_factor))
36            .iter()
37            .map(|it| **it)
38            .collect()
39    }
40}
41
42impl UnrollProcessor {
43    fn maybe_transform(
44        &self,
45        alloc: &Allocator,
46        inst: &Instruction,
47        mappings: &mut Mappings,
48    ) -> TransformAction {
49        if let Operation::TypedMemory(memory) = &inst.operation {
50            let lowered = inst.clone().lower_typed_memory().expect("valid typed memory instruction");
51            return match self.maybe_transform(alloc, &lowered, mappings) {
52                TransformAction::Ignore => TransformAction::Ignore,
53                TransformAction::Replace(instructions) => TransformAction::Replace(
54                    instructions.into_iter().map(|instruction| memory.retag(instruction)).collect()),
55            };
56        }
57        if matches!(inst.operation, Operation::Marker(_)) {
58            return TransformAction::Ignore;
59        }
60
61        if inst.operation.args().is_none() {
62            // Detect unhandled ops that can't be reflected
63            match &inst.operation {
64                Operation::CoopMma(op) => match op {
65                    // Stride is in scalar elems
66                    CoopMma::Load {
67                        value,
68                        stride,
69                        offset,
70                        layout,
71                    } if value.vector_size() > self.max_vector_size => {
72                        return TransformAction::Replace(self.transform_cmma_load(
73                            alloc,
74                            inst.out(),
75                            value,
76                            stride,
77                            offset,
78                            layout,
79                        ));
80                    }
81                    CoopMma::Store {
82                        mat,
83                        stride,
84                        offset,
85                        layout,
86                    } if inst.out().vector_size() > self.max_vector_size => {
87                        return TransformAction::Replace(self.transform_cmma_store(
88                            alloc,
89                            inst.out(),
90                            mat,
91                            stride,
92                            offset,
93                            layout,
94                        ));
95                    }
96                    _ => return TransformAction::Ignore,
97                },
98                Operation::Branch(_) | Operation::NonSemantic(_) | Operation::Marker(_) => {
99                    return TransformAction::Ignore;
100                }
101                _ => {
102                    panic!("Need special handling for unrolling non-reflectable operations")
103                }
104            }
105        }
106
107        let args = inst.operation.args().unwrap_or_default();
108        if (inst.out.is_some() && inst.ty().vector_size() > self.max_vector_size)
109            || args
110                .iter()
111                .any(|arg| arg.vector_size() > self.max_vector_size)
112        {
113            let vector_size = max_vector_size(&inst.out, &args);
114            let unroll_factor = vector_size / self.max_vector_size;
115
116            match &inst.operation {
117                Operation::Operator(Operator::CopyMemoryBulk(op)) => TransformAction::Replace(
118                    self.transform_memcpy(alloc, op, inst.out(), unroll_factor),
119                ),
120                Operation::Operator(Operator::CopyMemory(op)) => {
121                    TransformAction::Replace(self.transform_memcpy(
122                        alloc,
123                        &CopyMemoryBulkOperator {
124                            out_index: op.out_index,
125                            input: op.input,
126                            in_index: op.in_index,
127                            len: 1,
128                            offset_input: 0.into(),
129                            offset_out: 0.into(),
130                        },
131                        inst.out(),
132                        unroll_factor,
133                    ))
134                }
135                Operation::Operator(Operator::Index(op)) if op.list.is_array() => {
136                    TransformAction::Replace(self.transform_array_index(
137                        alloc,
138                        inst.out(),
139                        op,
140                        Operator::Index,
141                        unroll_factor,
142                        mappings,
143                    ))
144                }
145                Operation::Operator(Operator::UncheckedIndex(op)) if op.list.is_array() => {
146                    TransformAction::Replace(self.transform_array_index(
147                        alloc,
148                        inst.out(),
149                        op,
150                        Operator::UncheckedIndex,
151                        unroll_factor,
152                        mappings,
153                    ))
154                }
155                Operation::Operator(Operator::Index(op)) => {
156                    TransformAction::Replace(self.transform_composite_index(
157                        alloc,
158                        inst.out(),
159                        op,
160                        Operator::Index,
161                        unroll_factor,
162                        mappings,
163                    ))
164                }
165                Operation::Operator(Operator::UncheckedIndex(op)) => {
166                    TransformAction::Replace(self.transform_composite_index(
167                        alloc,
168                        inst.out(),
169                        op,
170                        Operator::UncheckedIndex,
171                        unroll_factor,
172                        mappings,
173                    ))
174                }
175                Operation::Operator(Operator::IndexAssign(op)) if inst.out().is_array() => {
176                    TransformAction::Replace(self.transform_array_index_assign(
177                        alloc,
178                        inst.out(),
179                        op,
180                        Operator::IndexAssign,
181                        unroll_factor,
182                        mappings,
183                    ))
184                }
185                Operation::Operator(Operator::UncheckedIndexAssign(op))
186                    if inst.out().is_array() =>
187                {
188                    TransformAction::Replace(self.transform_array_index_assign(
189                        alloc,
190                        inst.out(),
191                        op,
192                        Operator::UncheckedIndexAssign,
193                        unroll_factor,
194                        mappings,
195                    ))
196                }
197                Operation::Operator(Operator::IndexAssign(op)) => {
198                    TransformAction::Replace(self.transform_composite_index_assign(
199                        alloc,
200                        inst.out(),
201                        op,
202                        Operator::IndexAssign,
203                        unroll_factor,
204                        mappings,
205                    ))
206                }
207                Operation::Operator(Operator::UncheckedIndexAssign(op)) => {
208                    TransformAction::Replace(self.transform_composite_index_assign(
209                        alloc,
210                        inst.out(),
211                        op,
212                        Operator::UncheckedIndexAssign,
213                        unroll_factor,
214                        mappings,
215                    ))
216                }
217                Operation::Metadata(op) => {
218                    TransformAction::Replace(self.transform_metadata(inst.out(), op, args))
219                }
220                _ => TransformAction::Replace(self.transform_basic(
221                    alloc,
222                    inst,
223                    args,
224                    unroll_factor,
225                    mappings,
226                )),
227            }
228        } else {
229            TransformAction::Ignore
230        }
231    }
232
233    /// Transform CMMA load offset and array
234    fn transform_cmma_load(
235        &self,
236        alloc: &Allocator,
237        out: Variable,
238        value: &Variable,
239        stride: &Variable,
240        offset: &Variable,
241        layout: &Option<MatrixLayout>,
242    ) -> Vec<Instruction> {
243        let vector_size = value.vector_size();
244        let unroll_factor = vector_size / self.max_vector_size;
245
246        let value = unroll_array(*value, self.max_vector_size, unroll_factor);
247        let (mul, offset) = mul_index(alloc, *offset, unroll_factor);
248        let load = Instruction::new(
249            Operation::CoopMma(CoopMma::Load {
250                value,
251                stride: *stride,
252                offset: *offset,
253                layout: *layout,
254            }),
255            out,
256        );
257        vec![mul, load]
258    }
259
260    /// Transform CMMA store offset and array
261    fn transform_cmma_store(
262        &self,
263        alloc: &Allocator,
264        out: Variable,
265        mat: &Variable,
266        stride: &Variable,
267        offset: &Variable,
268        layout: &MatrixLayout,
269    ) -> Vec<Instruction> {
270        let vector_size = out.vector_size();
271        let unroll_factor = vector_size / self.max_vector_size;
272
273        let out = unroll_array(out, self.max_vector_size, unroll_factor);
274        let (mul, offset) = mul_index(alloc, *offset, unroll_factor);
275        let store = Instruction::new(
276            Operation::CoopMma(CoopMma::Store {
277                mat: *mat,
278                stride: *stride,
279                offset: *offset,
280                layout: *layout,
281            }),
282            out,
283        );
284        vec![mul, store]
285    }
286
287    /// Transforms memcpy into one with higher length and adjusted indices/offsets
288    fn transform_memcpy(
289        &self,
290        alloc: &Allocator,
291        op: &CopyMemoryBulkOperator,
292        out: Variable,
293        unroll_factor: usize,
294    ) -> Vec<Instruction> {
295        let (mul1, in_index) = mul_index(alloc, op.in_index, unroll_factor);
296        let (mul2, offset_input) = mul_index(alloc, op.offset_input, unroll_factor);
297        let (mul3, out_index) = mul_index(alloc, op.out_index, unroll_factor);
298        let (mul4, offset_out) = mul_index(alloc, op.offset_out, unroll_factor);
299
300        let input = unroll_array(op.input, self.max_vector_size, unroll_factor);
301        let out = unroll_array(out, self.max_vector_size, unroll_factor);
302
303        vec![
304            mul1,
305            mul2,
306            mul3,
307            mul4,
308            Instruction::new(
309                Operator::CopyMemoryBulk(CopyMemoryBulkOperator {
310                    input,
311                    in_index: *in_index,
312                    out_index: *out_index,
313                    len: op.len * unroll_factor,
314                    offset_input: *offset_input,
315                    offset_out: *offset_out,
316                }),
317                out,
318            ),
319        ]
320    }
321
322    /// Transforms indexing into multiple index operations, each offset by 1 from the base. The base
323    /// is also multiplied by the unroll factor to compensate for the lower actual vectorization.
324    fn transform_array_index(
325        &self,
326        alloc: &Allocator,
327        out: Variable,
328        op: &IndexOperator,
329        operator: impl Fn(IndexOperator) -> Operator,
330        unroll_factor: usize,
331        mappings: &mut Mappings,
332    ) -> Vec<Instruction> {
333        let (mul, start_idx) = mul_index(alloc, op.index, unroll_factor);
334        let mut indices = (0..unroll_factor).map(|i| add_index(alloc, *start_idx, i));
335
336        let list = unroll_array(op.list, self.max_vector_size, unroll_factor);
337
338        let out = mappings.get(alloc, out, unroll_factor, self.max_vector_size);
339        let mut instructions = vec![mul];
340        instructions.extend((0..unroll_factor).flat_map(|i| {
341            let (add, idx) = indices.next().unwrap();
342            let index = Instruction::new(
343                operator(IndexOperator {
344                    list,
345                    index: *idx,
346                    vector_size: 0,
347                    unroll_factor,
348                }),
349                out[i],
350            );
351            [add, index]
352        }));
353
354        instructions
355    }
356
357    /// Transforms index assign into multiple index assign operations, each offset by 1 from the base.
358    /// The base is also multiplied by the unroll factor to compensate for the lower actual vectorization.
359    fn transform_array_index_assign(
360        &self,
361        alloc: &Allocator,
362        out: Variable,
363        op: &IndexAssignOperator,
364        operator: impl Fn(IndexAssignOperator) -> Operator,
365        unroll_factor: usize,
366        mappings: &mut Mappings,
367    ) -> Vec<Instruction> {
368        let (mul, start_idx) = mul_index(alloc, op.index, unroll_factor);
369        let mut indices = (0..unroll_factor).map(|i| add_index(alloc, *start_idx, i));
370
371        let out = unroll_array(out, self.max_vector_size, unroll_factor);
372
373        let value = mappings.get(alloc, op.value, unroll_factor, self.max_vector_size);
374
375        let mut instructions = vec![mul];
376        instructions.extend((0..unroll_factor).flat_map(|i| {
377            let (add, idx) = indices.next().unwrap();
378            let index = Instruction::new(
379                operator(IndexAssignOperator {
380                    index: *idx,
381                    vector_size: 0,
382                    value: value[i],
383                    unroll_factor,
384                }),
385                out,
386            );
387
388            [add, index]
389        }));
390
391        instructions
392    }
393
394    /// Transforms a composite index (i.e. `Vector`) that always returns a scalar. Translates the index
395    /// to a local index and an unroll index, then indexes the proper variable. Note that this requires
396    /// the index to be constant - it needs to be decomposed at compile time, otherwise it wouldn't
397    /// work.
398    fn transform_composite_index(
399        &self,
400        alloc: &Allocator,
401        out: Variable,
402        op: &IndexOperator,
403        operator: impl Fn(IndexOperator) -> Operator,
404        unroll_factor: usize,
405        mappings: &mut Mappings,
406    ) -> Vec<Instruction> {
407        let index = op
408            .index
409            .as_const()
410            .expect("Can't unroll non-constant vector index")
411            .as_usize();
412
413        let unroll_idx = index / self.max_vector_size;
414        let sub_idx = index % self.max_vector_size;
415
416        let value = mappings.get(alloc, op.list, unroll_factor, self.max_vector_size);
417
418        vec![Instruction::new(
419            operator(IndexOperator {
420                list: value[unroll_idx],
421                index: sub_idx.into(),
422                vector_size: 1,
423                unroll_factor,
424            }),
425            out,
426        )]
427    }
428
429    /// Transforms a composite index assign (i.e. `Vector`) that always takes a scalar. Translates the index
430    /// to a local index and an unroll index, then indexes the proper variable. Note that this requires
431    /// the index to be constant - it needs to be decomposed at compile time, otherwise it wouldn't
432    /// work.
433    fn transform_composite_index_assign(
434        &self,
435        alloc: &Allocator,
436        out: Variable,
437        op: &IndexAssignOperator,
438        operator: impl Fn(IndexAssignOperator) -> Operator,
439        unroll_factor: usize,
440        mappings: &mut Mappings,
441    ) -> Vec<Instruction> {
442        let index = op
443            .index
444            .as_const()
445            .expect("Can't unroll non-constant vector index")
446            .as_usize();
447
448        let unroll_idx = index / self.max_vector_size;
449        let sub_idx = index % self.max_vector_size;
450
451        let out = mappings.get(alloc, out, unroll_factor, self.max_vector_size);
452
453        vec![Instruction::new(
454            operator(IndexAssignOperator {
455                index: sub_idx.into(),
456                vector_size: 1,
457                value: op.value,
458                unroll_factor,
459            }),
460            out[unroll_idx],
461        )]
462    }
463
464    /// Transforms metadata by just replacing the type of the buffer. The values are already
465    /// properly calculated on the CPU.
466    fn transform_metadata(
467        &self,
468        out: Variable,
469        op: &Metadata,
470        args: Vec<Variable>,
471    ) -> Vec<Instruction> {
472        let op_code = op.op_code();
473        let args = args
474            .into_iter()
475            .map(|mut var| {
476                if var.vector_size() > self.max_vector_size {
477                    var.ty = var.ty.with_vector_size(self.max_vector_size);
478                }
479                var
480            })
481            .collect::<Vec<_>>();
482        let operation = Metadata::from_code_and_args(op_code, &args).unwrap();
483        vec![Instruction::new(operation, out)]
484    }
485
486    /// Transforms generic instructions, i.e. comparison, arithmetic. Unrolls each vectorized variable
487    /// to `unroll_factor` replacements, and executes the operation `unroll_factor` times.
488    fn transform_basic(
489        &self,
490        alloc: &Allocator,
491        inst: &Instruction,
492        args: Vec<Variable>,
493        unroll_factor: usize,
494        mappings: &mut Mappings,
495    ) -> Vec<Instruction> {
496        let op_code = inst.operation.op_code();
497        let out = inst
498            .out
499            .map(|out| mappings.get(alloc, out, unroll_factor, self.max_vector_size));
500        let args = args
501            .into_iter()
502            .map(|arg| {
503                if arg.vector_size() > 1 {
504                    mappings.get(alloc, arg, unroll_factor, self.max_vector_size)
505                } else {
506                    // Preserve scalars
507                    vec![arg]
508                }
509            })
510            .collect::<Vec<_>>();
511
512        (0..unroll_factor)
513            .map(|i| {
514                let out = out.as_ref().map(|out| out[i]);
515                let args = args
516                    .iter()
517                    .map(|arg| if arg.len() == 1 { arg[0] } else { arg[i] })
518                    .collect::<Vec<_>>();
519                let operation = Operation::from_code_and_args(op_code, &args)
520                    .expect("Failed to reconstruct operation");
521                Instruction {
522                    out,
523                    source_loc: inst.source_loc.clone(),
524                    modes: inst.modes,
525                    operation,
526                }
527            })
528            .collect()
529    }
530
531    fn transform_instructions(
532        &self,
533        allocator: &Allocator,
534        instructions: Vec<Instruction>,
535        mappings: &mut Mappings,
536    ) -> Vec<Instruction> {
537        let mut new_instructions = Vec::with_capacity(instructions.len());
538
539        for mut instruction in instructions {
540            if let Operation::Branch(branch) = &mut instruction.operation {
541                match branch {
542                    Branch::If(op) => {
543                        op.scope.instructions = self.transform_instructions(
544                            allocator,
545                            op.scope.instructions.drain(..).collect(),
546                            mappings,
547                        );
548                    }
549                    Branch::IfElse(op) => {
550                        op.scope_if.instructions = self.transform_instructions(
551                            allocator,
552                            op.scope_if.instructions.drain(..).collect(),
553                            mappings,
554                        );
555                        op.scope_else.instructions = self.transform_instructions(
556                            allocator,
557                            op.scope_else.instructions.drain(..).collect(),
558                            mappings,
559                        );
560                    }
561                    Branch::Switch(op) => {
562                        for (_, case) in &mut op.cases {
563                            case.instructions = self.transform_instructions(
564                                allocator,
565                                case.instructions.drain(..).collect(),
566                                mappings,
567                            );
568                        }
569                        op.scope_default.instructions = self.transform_instructions(
570                            allocator,
571                            op.scope_default.instructions.drain(..).collect(),
572                            mappings,
573                        );
574                    }
575                    Branch::RangeLoop(op) => {
576                        op.scope.instructions = self.transform_instructions(
577                            allocator,
578                            op.scope.instructions.drain(..).collect(),
579                            mappings,
580                        );
581                    }
582                    Branch::Loop(op) => {
583                        op.scope.instructions = self.transform_instructions(
584                            allocator,
585                            op.scope.instructions.drain(..).collect(),
586                            mappings,
587                        );
588                    }
589                    _ => {}
590                }
591            }
592            match self.maybe_transform(allocator, &instruction, mappings) {
593                TransformAction::Ignore => {
594                    new_instructions.push(instruction);
595                }
596                TransformAction::Replace(replacement) => {
597                    new_instructions.extend(replacement);
598                }
599            }
600        }
601
602        new_instructions
603    }
604}
605
606impl Processor for UnrollProcessor {
607    fn transform(&self, processing: ScopeProcessing, allocator: Allocator) -> ScopeProcessing {
608        let mut mappings = Mappings(Default::default());
609
610        let instructions =
611            self.transform_instructions(&allocator, processing.instructions, &mut mappings);
612
613        ScopeProcessing {
614            variables: processing.variables,
615            instructions,
616            typemap: processing.typemap.clone(),
617        }
618    }
619}
620
621fn max_vector_size(out: &Option<Variable>, args: &[Variable]) -> VectorSize {
622    let vector_size = args.iter().map(|it| it.vector_size()).max().unwrap();
623    vector_size.max(out.map(|out| out.vector_size()).unwrap_or(1))
624}
625
626fn create_unrolled(
627    allocator: &Allocator,
628    var: &Variable,
629    max_vector_size: VectorSize,
630    unroll_factor: usize,
631) -> Vec<ManagedVariable> {
632    // Preserve scalars
633    if var.vector_size() == 1 {
634        return vec![ManagedVariable::Plain(*var); unroll_factor];
635    }
636
637    let item = Type::new(var.storage_type()).with_vector_size(max_vector_size);
638    (0..unroll_factor)
639        .map(|_| match var.kind {
640            VariableKind::LocalMut { .. } | VariableKind::Versioned { .. } => {
641                allocator.create_local_mut(item)
642            }
643            VariableKind::Shared { .. } => {
644                let id = allocator.new_local_index();
645                let shared = VariableKind::Shared { id };
646                ManagedVariable::Plain(Variable::new(shared, item))
647            }
648            VariableKind::LocalConst { .. } => allocator.create_local(item),
649            other => panic!("Out must be local, found {other:?}"),
650        })
651        .collect()
652}
653
654fn add_index(alloc: &Allocator, idx: Variable, i: usize) -> (Instruction, ManagedVariable) {
655    let add_idx = alloc.create_local(idx.ty);
656    let add = Instruction::new(
657        Arithmetic::Add(BinaryOperator {
658            lhs: idx,
659            rhs: i.into(),
660        }),
661        *add_idx,
662    );
663    (add, add_idx)
664}
665
666fn mul_index(
667    alloc: &Allocator,
668    idx: Variable,
669    unroll_factor: usize,
670) -> (Instruction, ManagedVariable) {
671    let mul_idx = alloc.create_local(idx.ty);
672    let mul = Instruction::new(
673        Arithmetic::Mul(BinaryOperator {
674            lhs: idx,
675            rhs: unroll_factor.into(),
676        }),
677        *mul_idx,
678    );
679    (mul, mul_idx)
680}
681
682fn unroll_array(mut var: Variable, max_vector_size: VectorSize, factor: usize) -> Variable {
683    var.ty = var.ty.with_vector_size(max_vector_size);
684
685    match &mut var.kind {
686        VariableKind::LocalArray { unroll_factor, .. }
687        | VariableKind::ConstantArray { unroll_factor, .. }
688        | VariableKind::SharedArray { unroll_factor, .. } => {
689            *unroll_factor = factor;
690        }
691        _ => {}
692    }
693
694    var
695}
ruda_kernel/dsl/post_processing/unroll.rs

ruda_kernel/dsl/post_processing/
unroll.rs
