cubecl_core/post_processing/
unroll.rs

1use alloc::{boxed::Box, vec, vec::Vec};
2use cubecl_environment::collections::HashMap;
3use cubecl_ir::{
4    VectorSize,
5    attributes::{IndexAttr, ZeroAttr},
6    dialect::{
7        base::OperationPtrExt,
8        cmp::IEqualOp,
9        general::{CopyOp, ReinterpretCastOp, SelectOp},
10        math::{IAddOp, IMulOp, UDivOp, URemOp},
11        matrix,
12        memory::{DeclareVariableOp, IndexOp},
13        vector::{
14            CompositeConstructOp, CompositeExtractOp, CompositeInsertOp, VectorExtractDynamicOp,
15            VectorInsertDynamicOp,
16        },
17    },
18    interfaces::{MaybeVectorizedType, TriviallyUnrollable, TypedExt},
19    prelude::*,
20    try_cast_op,
21    types::{ArrayType, AtomicType, PointerType, RuntimeArrayType, VectorType},
22    verify_op_succ, verify_ty_succ,
23};
24use pliron::{
25    builtin::{
26        ops::{ConstantOp, FuncOp},
27        types::FunctionType,
28    },
29    graph::walkers::{WALKCONFIG_PREORDER_FORWARD, uninterruptible::mutable::walk_op},
30};
31
32type Mappings = HashMap<Value, Vec<Value>>;
33
34#[derive(Debug, new)]
35pub struct UnrollPass {
36    max_vector_size: VectorSize,
37}
38
39#[type_interface]
40pub trait UnrollableType: MaybeVectorizedType {
41    verify_ty_succ!();
42    fn with_vector_size(&self, ctx: &Context, vectorization: usize) -> TypeHandle;
43}
44
45#[type_interface_impl]
46impl UnrollableType for VectorType {
47    fn with_vector_size(&self, ctx: &Context, vectorization: usize) -> TypeHandle {
48        VectorType::get(ctx, self.inner, vectorization).into()
49    }
50}
51
52#[type_interface_impl]
53impl UnrollableType for AtomicType {
54    fn with_vector_size(&self, ctx: &Context, vectorization: usize) -> TypeHandle {
55        let inner = self.inner.deref(ctx);
56        let unrollable = type_cast::<dyn UnrollableType>(&*inner).expect("Should be implemented");
57        let new_inner = unrollable.with_vector_size(ctx, vectorization);
58        AtomicType::get(ctx, new_inner).into()
59    }
60}
61
62#[type_interface_impl]
63impl UnrollableType for PointerType {
64    fn with_vector_size(&self, ctx: &Context, vectorization: usize) -> TypeHandle {
65        let inner = self.inner.deref(ctx);
66        let unrollable = type_cast::<dyn UnrollableType>(&*inner).expect("Should be implemented");
67        let new_inner = unrollable.with_vector_size(ctx, vectorization);
68        PointerType::get(ctx, new_inner, self.address_space).into()
69    }
70}
71
72#[type_interface_impl]
73impl UnrollableType for ArrayType {
74    fn with_vector_size(&self, ctx: &Context, new_vec: usize) -> TypeHandle {
75        let current_vec = self.vector_size(ctx);
76        let inner = self.inner.deref(ctx);
77        let unrollable = type_cast::<dyn UnrollableType>(&*inner).expect("Should be implemented");
78        let new_inner = unrollable.with_vector_size(ctx, new_vec);
79        ArrayType::get(ctx, new_inner, self.length * current_vec / new_vec).into()
80    }
81}
82
83#[type_interface_impl]
84impl UnrollableType for RuntimeArrayType {
85    fn with_vector_size(&self, ctx: &Context, new_vec: usize) -> TypeHandle {
86        let inner = self.inner.deref(ctx);
87        let unrollable = type_cast::<dyn UnrollableType>(&*inner).expect("Should be implemented");
88        let new_inner = unrollable.with_vector_size(ctx, new_vec);
89        RuntimeArrayType::get(ctx, new_inner).into()
90    }
91}
92
93#[op_interface]
94pub trait CustomUnrollOp {
95    verify_op_succ!();
96    fn unroll(&self, ctx: &mut Context, state: &mut UnrollState);
97}
98
99#[op_interface_impl]
100impl CustomUnrollOp for DeclareVariableOp {
101    fn unroll(&self, ctx: &mut Context, state: &mut UnrollState) {
102        let value_ty = self.value_ty(ctx).get_type(ctx);
103        let current_vec = value_ty.try_get_vector_size(ctx).unwrap_or(1);
104        if current_vec <= state.max_vector_size {
105            return;
106        }
107
108        state.result.ir_changed |= IRStatus::Changed;
109        let result = self.get_result(ctx);
110        let addr_space = *self.addr_space(ctx);
111        // Align isn't handled properly, but targets that unroll ignore this anyways
112        let align = *self.alignment(ctx);
113        let new_value_ty = unroll_ty(ctx, value_ty, state.max_vector_size);
114        let new_ptr_ty = PointerType::get(ctx, new_value_ty, addr_space.0);
115
116        // The initializer has a type too, and it must match the declaration.
117        // Unroll it as well, or we end up with a broken
118        // `array<vec4<f32>, 8> = array<vec8<f32>, 4>()`.
119        let unrolled_init = |ctx: &Context| {
120            self.initializer(ctx)
121                .map(|init| match init.is::<ZeroAttr>() {
122                    true => ZeroAttr::new(new_value_ty).into(),
123                    false => init.clone(),
124                })
125        };
126
127        // Array doesn't change size, so no need to duplicate the declaration
128        if new_value_ty.size(ctx) == value_ty.size(ctx) {
129            let init = unrolled_init(ctx);
130            self.set_value_ty(ctx, new_value_ty);
131            self.get_result(ctx).set_type(ctx, new_ptr_ty.into());
132            if let Some(init) = init {
133                self.set_initializer(ctx, init);
134            }
135        } else {
136            let factor = current_vec / state.max_vector_size;
137            let mut results = vec![];
138            for _ in 0..factor {
139                let init = unrolled_init(ctx);
140                let new_op = DeclareVariableOp::new(ctx, new_value_ty, addr_space, align, init);
141                new_op
142                    .get_operation()
143                    .insert_before(ctx, self.get_operation());
144                results.push(new_op.get_result(ctx));
145            }
146            state.mappings.insert(result, results);
147            state.to_erase.push(self.get_operation());
148        }
149    }
150}
151
152#[op_interface_impl]
153impl CustomUnrollOp for IndexOp {
154    fn unroll(&self, ctx: &mut Context, state: &mut UnrollState) {
155        let base = self.base(ctx);
156        let checked = self.checked(ctx);
157        let current_vec = try_get_vec(ctx, self.get_result(ctx));
158        if current_vec > state.max_vector_size {
159            state.result.ir_changed |= IRStatus::Changed;
160            let unroll_factor = current_vec / state.max_vector_size;
161            let unroll_const = const_usize(ctx, self, unroll_factor);
162
163            let mul = IMulOp::new(ctx, self.index(ctx), unroll_const);
164            mul.get_operation().insert_before(ctx, self.get_operation());
165            let start_idx = mul.get_result(ctx);
166
167            let new_results = (0..unroll_factor)
168                .map(|i| {
169                    let i = const_usize(ctx, self, i);
170                    let add = IAddOp::new(ctx, start_idx, i);
171                    add.get_operation().insert_before(ctx, self.get_operation());
172                    let idx = add.get_result(ctx);
173
174                    let op = IndexOp::maybe_checked(ctx, base, idx, checked);
175                    op.get_operation().insert_before(ctx, self.get_operation());
176                    op.get_result(ctx)
177                })
178                .collect();
179
180            state.mappings.insert(self.get_result(ctx), new_results);
181            state.to_erase.push(self.get_operation());
182        }
183    }
184}
185
186fn const_usize(ctx: &mut Context, anchor: &dyn Op, value: usize) -> Value {
187    let op = ConstantOp::new(ctx, Box::new(IndexAttr::new(value)));
188    op.get_operation()
189        .insert_before(ctx, anchor.get_operation());
190    op.get_result(ctx)
191}
192
193#[op_interface_impl]
194impl CustomUnrollOp for CompositeExtractOp {
195    fn unroll(&self, ctx: &mut Context, state: &mut UnrollState) {
196        let vector = self.composite(ctx);
197        if !vector.is_vector(ctx) {
198            return;
199        }
200
201        let current_vec = vector.vector_size(ctx);
202        if current_vec > state.max_vector_size {
203            state.result.ir_changed |= IRStatus::Changed;
204            let index = self.index(ctx).0;
205
206            let unroll_idx = index / state.max_vector_size;
207            let sub_idx = index % state.max_vector_size;
208
209            let new_vector = state.mappings.get(&vector).expect("Should exist")[unroll_idx];
210            vector.replace_use_with(ctx, self.composite_as_use(ctx), &new_vector);
211            self.set_index(ctx, sub_idx);
212        }
213    }
214}
215
216#[op_interface_impl]
217impl CustomUnrollOp for CompositeInsertOp {
218    fn unroll(&self, ctx: &mut Context, state: &mut UnrollState) {
219        let vector = self.composite(ctx);
220        if !vector.is_vector(ctx) {
221            return;
222        }
223
224        let value = self.value(ctx);
225        let current_vec = vector.vector_size(ctx);
226        if current_vec > state.max_vector_size {
227            state.result.ir_changed |= IRStatus::Changed;
228            let index = self.index(ctx).0;
229
230            let unroll_idx = index / state.max_vector_size;
231            let sub_idx = index % state.max_vector_size;
232
233            let vectors = state.mappings.get(&vector).expect("Should exist");
234
235            let new_results = vectors.iter().enumerate().map(|(i, vector)| {
236                let op = if i == unroll_idx {
237                    CompositeInsertOp::new(ctx, *vector, value, sub_idx).get_operation()
238                } else {
239                    CopyOp::new(ctx, *vector).get_operation()
240                };
241                op.insert_before(ctx, self.get_operation());
242                op.deref(ctx).get_result(0)
243            });
244            let new_results = new_results.collect();
245            state.mappings.insert(self.get_result(ctx), new_results);
246            state.to_erase.push(self.get_operation());
247        }
248    }
249}
250
251/// A dynamic index reaches a lane the pass cannot name at compile time, so the
252/// unrolled parts are all searched: the index splits into the part it lands in
253/// and the lane within that part, and a select chain picks the part's answer.
254/// The frontend documents both dynamic accessors as very slow already, and the
255/// chain is over the (comptime) number of parts, not the lanes.
256#[op_interface_impl]
257impl CustomUnrollOp for VectorExtractDynamicOp {
258    fn unroll(&self, ctx: &mut Context, state: &mut UnrollState) {
259        let vector = self.vector(ctx);
260        if try_get_vec(ctx, vector) <= state.max_vector_size {
261            return;
262        }
263
264        state.result.ir_changed |= IRStatus::Changed;
265        let (part, lane) = split_index(ctx, self, state.max_vector_size, self.index(ctx));
266        let parts = state.mappings.get(&vector).expect("Should exist").clone();
267
268        let mut extracted = None;
269        for (i, part_vector) in parts.into_iter().enumerate() {
270            let op = VectorExtractDynamicOp::new(ctx, part_vector, lane);
271            op.get_operation().insert_before(ctx, self.get_operation());
272            let value = op.get_result(ctx);
273            extracted = Some(match extracted {
274                Some(previous) => select_part(ctx, self, part, i, value, previous),
275                None => value,
276            });
277        }
278
279        let extracted = extracted.expect("Unrolled vector should have parts");
280        self.get_result(ctx).replace_all_uses_with(ctx, &extracted);
281        state.to_erase.push(self.get_operation());
282    }
283}
284
285#[op_interface_impl]
286impl CustomUnrollOp for VectorInsertDynamicOp {
287    fn unroll(&self, ctx: &mut Context, state: &mut UnrollState) {
288        let vector = self.vector(ctx);
289        if try_get_vec(ctx, vector) <= state.max_vector_size {
290            return;
291        }
292
293        state.result.ir_changed |= IRStatus::Changed;
294        let value = self.value(ctx);
295        let (part, lane) = split_index(ctx, self, state.max_vector_size, self.index(ctx));
296        let parts = state.mappings.get(&vector).expect("Should exist").clone();
297
298        let mut new_results = Vec::with_capacity(parts.len());
299        for (i, part_vector) in parts.into_iter().enumerate() {
300            // Selecting the *scalar* to write, rather than between the written
301            // and untouched parts, keeps the condition off the vector operands:
302            // a select over vectors with a scalar condition needs SPIR-V 1.4.
303            let old = VectorExtractDynamicOp::new(ctx, part_vector, lane);
304            old.get_operation().insert_before(ctx, self.get_operation());
305            let written = select_part(ctx, self, part, i, value, old.get_result(ctx));
306
307            let op = VectorInsertDynamicOp::new(ctx, part_vector, written, lane);
308            op.get_operation().insert_before(ctx, self.get_operation());
309            new_results.push(op.get_result(ctx));
310        }
311
312        state.mappings.insert(self.get_result(ctx), new_results);
313        state.to_erase.push(self.get_operation());
314    }
315}
316
317/// Split a dynamic vector index into the unrolled part it lands in and the lane
318/// within that part.
319fn split_index(
320    ctx: &mut Context,
321    anchor: &dyn Op,
322    max_vector_size: VectorSize,
323    index: Value,
324) -> (Value, Value) {
325    let size = const_usize(ctx, anchor, max_vector_size);
326    let part = UDivOp::new(ctx, index, size);
327    part.get_operation()
328        .insert_before(ctx, anchor.get_operation());
329    let lane = URemOp::new(ctx, index, size);
330    lane.get_operation()
331        .insert_before(ctx, anchor.get_operation());
332    (part.get_result(ctx), lane.get_result(ctx))
333}
334
335/// `if part == i { on_match } else { otherwise }`.
336fn select_part(
337    ctx: &mut Context,
338    anchor: &dyn Op,
339    part: Value,
340    i: usize,
341    on_match: Value,
342    otherwise: Value,
343) -> Value {
344    let i = const_usize(ctx, anchor, i);
345    let is_part = IEqualOp::new(ctx, part, i);
346    is_part
347        .get_operation()
348        .insert_before(ctx, anchor.get_operation());
349    let select = SelectOp::new(ctx, is_part.get_result(ctx), on_match, otherwise);
350    select
351        .get_operation()
352        .insert_before(ctx, anchor.get_operation());
353    select.get_result(ctx)
354}
355
356/// A reinterpret changes the lane count with the lane width, so one side's pieces are not the
357/// other's: a `Vector<u32, 2>` is a `Vector<e4m3, 8>`. When only one side exceeds the maximum, the
358/// wide side is split at the unroll factor and the narrow side is cut into matching sub-vectors
359/// with extracts and constructs. When both exceed it, the input's pieces are reinterpreted in
360/// place and their lanes regrouped.
361#[op_interface_impl]
362impl CustomUnrollOp for ReinterpretCastOp {
363    fn unroll(&self, ctx: &mut Context, state: &mut UnrollState) {
364        let input = self.input(ctx);
365        let result = self.get_result(ctx);
366        let max = state.max_vector_size;
367        let (in_vec, out_vec) = (try_get_vec(ctx, input), try_get_vec(ctx, result));
368        if in_vec <= max && out_vec <= max {
369            return;
370        }
371        // When both sides unroll, neither is a whole piece of the other: each `max`-lane input
372        // piece reinterprets into the output lanes it shares bits with, and those lanes regroup
373        // into the `max`-lane pieces the rest of the pass indexes into. Equal lane counts are the
374        // common case here and stay one reinterpret per piece.
375        if in_vec > max && out_vec > max {
376            // A piece of the input covers `per_piece` lanes of the output, and both sides of the
377            // reinterpret below have to be a legal vector: one lane at least, and no wider than
378            // the maximum. Sub-dividing the input pieces would lift the upper bound, but nothing
379            // unrolls what this pass emits, so an over-wide piece would reach the backend as-is.
380            let per_piece = max * out_vec / in_vec;
381            assert!(
382                (1..=max).contains(&per_piece),
383                "Cannot unroll a reinterpret between a {in_vec}-lane and a {out_vec}-lane vector \
384                 when both exceed {max} lanes and the lane counts differ: reinterpret through a \
385                 vector of at most {max} lanes"
386            );
387            state.result.ir_changed |= IRStatus::Changed;
388            let op = self.get_operation();
389
390            let piece_ty = lanes_type(ctx, result, per_piece);
391            let pieces = state.mappings.get(&input).expect("Should exist").clone();
392            let converted = pieces
393                .into_iter()
394                .map(|piece| {
395                    let new_op = ReinterpretCastOp::new(ctx, piece_ty, piece);
396                    new_op.get_operation().insert_before(ctx, op);
397                    new_op.get_result(ctx)
398                })
399                .collect::<Vec<_>>();
400
401            let new_results = if per_piece == max {
402                converted
403            } else {
404                let lanes = converted
405                    .into_iter()
406                    .flat_map(|piece| split_lanes(ctx, piece, op))
407                    .collect::<Vec<_>>();
408                let result_ty = unroll_ty(ctx, result, max);
409                lanes
410                    .chunks(max)
411                    .map(|lanes| {
412                        let joined = CompositeConstructOp::new(ctx, result_ty, lanes.to_vec());
413                        joined.get_operation().insert_before(ctx, op);
414                        joined.get_result(ctx)
415                    })
416                    .collect()
417            };
418            state.mappings.insert(result, new_results);
419            state.to_erase.push(op);
420            return;
421        }
422
423        // The wide side unrolls into `factor` pieces, so the narrow side has to have at least
424        // that many lanes to hand one to each. A 64-bit scalar reinterpreted as eight fp8 lanes
425        // is the case that gets here: one lane cannot be cut in two.
426        let (wide, narrow) = match in_vec > max {
427            true => (in_vec, out_vec),
428            false => (out_vec, in_vec),
429        };
430        let factor = wide / max;
431        assert!(
432            narrow >= factor,
433            "Cannot unroll a reinterpret between a {in_vec}-lane and a {out_vec}-lane vector: the \
434             {wide}-lane side unrolls into {factor} pieces and the {narrow}-lane side has no lane \
435             to give each. Reinterpret through a vector of at least {factor} lanes"
436        );
437        state.result.ir_changed |= IRStatus::Changed;
438        let op = self.get_operation();
439
440        if in_vec > max {
441            let piece_ty = lanes_type(ctx, result, out_vec / factor);
442            let pieces = state.mappings.get(&input).expect("Should exist").clone();
443            let converted = pieces
444                .into_iter()
445                .map(|piece| {
446                    let new_op = ReinterpretCastOp::new(ctx, piece_ty, piece);
447                    new_op.get_operation().insert_before(ctx, op);
448                    new_op.get_result(ctx)
449                })
450                .collect::<Vec<_>>();
451            let lanes = converted
452                .iter()
453                .flat_map(|piece| split_lanes(ctx, *piece, op))
454                .collect();
455            let result_ty = result.get_type(ctx);
456            let joined = CompositeConstructOp::new(ctx, result_ty, lanes);
457            joined.get_operation().insert_before(ctx, op);
458            result.replace_all_uses_with(ctx, &joined.get_result(ctx));
459        } else {
460            let piece_ty = unroll_ty(ctx, result, max);
461            let lanes = split_lanes(ctx, input, op);
462            let new_results = lanes
463                .chunks(in_vec / factor)
464                .map(|chunk| {
465                    let piece = match chunk {
466                        [lane] => *lane,
467                        lanes => {
468                            let ty = lanes_type(ctx, input, lanes.len());
469                            let joined = CompositeConstructOp::new(ctx, ty, lanes.to_vec());
470                            joined.get_operation().insert_before(ctx, op);
471                            joined.get_result(ctx)
472                        }
473                    };
474                    let new_op = ReinterpretCastOp::new(ctx, piece_ty, piece);
475                    new_op.get_operation().insert_before(ctx, op);
476                    new_op.get_result(ctx)
477                })
478                .collect();
479            state.mappings.insert(result, new_results);
480        }
481        state.to_erase.push(op);
482    }
483}
484
485/// The type of `lanes` lanes of `value`'s element: the bare scalar for one lane.
486fn lanes_type(ctx: &mut Context, value: Value, lanes: usize) -> TypeHandle {
487    assert!(lanes > 0, "Reinterpret pieces must hold at least one lane");
488    let scalar = value.scalar_ty(ctx);
489    match lanes {
490        1 => scalar,
491        lanes => VectorType::get(ctx, scalar, lanes).into(),
492    }
493}
494
495/// Extracts every lane of `value` before `anchor`; a scalar is its own single lane.
496fn split_lanes(ctx: &mut Context, value: Value, anchor: Ptr<Operation>) -> Vec<Value> {
497    if !value.is_vector(ctx) {
498        return vec![value];
499    }
500    (0..value.vector_size(ctx))
501        .map(|lane| {
502            let extract = CompositeExtractOp::new(ctx, value, lane);
503            extract.get_operation().insert_before(ctx, anchor);
504            extract.get_result(ctx)
505        })
506        .collect()
507}
508
509#[op_interface_impl]
510impl CustomUnrollOp for matrix::LoadOp {
511    fn unroll(&self, ctx: &mut Context, state: &mut UnrollState) {
512        let source = self.source(ctx);
513        if source.vector_size(ctx) > state.max_vector_size {
514            state.result.ir_changed |= IRStatus::Changed;
515            let new_source = state.mappings.get(&source).expect("should exist")[0];
516            source.replace_use_with(ctx, self.source_as_use(ctx), &new_source);
517        }
518    }
519}
520
521#[op_interface_impl]
522impl CustomUnrollOp for matrix::StoreOp {
523    fn unroll(&self, ctx: &mut Context, state: &mut UnrollState) {
524        let dest = self.destination(ctx);
525        if dest.vector_size(ctx) > state.max_vector_size {
526            state.result.ir_changed |= IRStatus::Changed;
527            let new_dest = state.mappings.get(&dest).expect("should exist")[0];
528            dest.replace_use_with(ctx, self.destination_as_use(ctx), &new_dest);
529        }
530    }
531}
532
533pub struct UnrollState {
534    mappings: Mappings,
535    to_erase: Vec<Ptr<Operation>>,
536    max_vector_size: VectorSize,
537    result: PassResult,
538    rewriter: PassRewriter,
539}
540
541#[pass_name]
542impl Pass for UnrollPass {
543    fn run(
544        &mut self,
545        op: Ptr<Operation>,
546        ctx: &mut Context,
547        _analyses: &mut AnalysisManager,
548    ) -> Result<PassResult> {
549        self.unroll_func(ctx, op);
550
551        let mut state = UnrollState {
552            mappings: Default::default(),
553            to_erase: Default::default(),
554            max_vector_size: self.max_vector_size,
555            result: Default::default(),
556            rewriter: PassRewriter::default(),
557        };
558
559        walk_op(
560            ctx,
561            &mut state,
562            &WALKCONFIG_PREORDER_FORWARD,
563            op,
564            |ctx, state, node| {
565                if let IRNode::Operation(op) = node {
566                    let unroll_opds = op
567                        .operands(ctx)
568                        .iter()
569                        .any(|it| should_unroll(ctx, it, state.max_vector_size));
570                    let unroll_res = op
571                        .results(ctx)
572                        .iter()
573                        .any(|it| should_unroll(ctx, it, state.max_vector_size));
574                    let dyn_op = op.dyn_op(ctx);
575
576                    if let Some(custom) = op_cast::<dyn CustomUnrollOp>(&*dyn_op) {
577                        custom.unroll(ctx, state);
578                    } else if unroll_opds || unroll_res {
579                        state.result.ir_changed |= IRStatus::Changed;
580                        unroll_default(ctx, state, op);
581                    }
582                }
583            },
584        );
585
586        while !state.to_erase.is_empty() {
587            // Pop the next op that no longer has uses. This ensures we always start at the end of
588            // the def-use chain
589            let next = state
590                .to_erase
591                .iter()
592                .position(|it| !it.deref(ctx).has_use())
593                .expect("Erased ops should only have uses in other erased ops");
594            let op = state.to_erase.remove(next);
595            state.rewriter.erase_operation(ctx, op);
596        }
597
598        Ok(state.result)
599    }
600}
601
602impl UnrollPass {
603    fn unroll_func(&self, ctx: &mut Context, op: Ptr<Operation>) {
604        let func = op.as_op::<FuncOp>(ctx).expect("Should be func");
605        let entry_block = func.get_entry_block(ctx);
606        let func_ty = func.get_attr_func_type(ctx).unwrap().get_type(ctx);
607        let func_ty = func_ty.deref(ctx);
608        let func_ty = func_ty.downcast_ref::<FunctionType>().unwrap();
609
610        let mut new_func_inputs = vec![];
611
612        for (i, arg) in func_ty.arg_types().into_iter().enumerate() {
613            if should_unroll(ctx, arg, self.max_vector_size) {
614                let new_ty = unroll_ty(ctx, arg, self.max_vector_size);
615                new_func_inputs.push(new_ty);
616                let block_arg = entry_block.deref(ctx).get_argument(i);
617                block_arg.set_type(ctx, new_ty);
618            } else {
619                new_func_inputs.push(arg);
620            }
621        }
622
623        let new_func_ty = FunctionType::get(ctx, new_func_inputs, func_ty.res_types()).to_handle();
624        func.set_attr_func_type(ctx, new_func_ty.into());
625    }
626}
627
628fn unroll_default(ctx: &mut Context, state: &mut UnrollState, op: Ptr<Operation>) {
629    let values = op.operands(ctx).into_iter().chain(op.results(ctx));
630    let current_vec = values.map(|it| try_get_vec(ctx, it)).max().unwrap();
631    let factor = current_vec / state.max_vector_size;
632    let dyn_op = op.dyn_op(ctx);
633    let rematerialize = try_cast_op!(dyn_op, ctx, dyn TriviallyUnrollable);
634    let new_out_ty = op
635        .results(ctx)
636        .into_iter()
637        .map(|it| unroll_ty(ctx, it, state.max_vector_size))
638        .collect::<Vec<_>>();
639    let mut new_results = vec![];
640
641    for unroll_idx in 0..factor {
642        let opds = op.operands(ctx).into_iter().map(|opd| {
643            if should_unroll(ctx, opd, state.max_vector_size) {
644                state.mappings.get(&opd).expect("Should have mapping")[unroll_idx]
645            } else {
646                opd
647            }
648        });
649        let attrs = op.deref(ctx).attributes.clone();
650        let new_op = rematerialize.materialize(ctx, new_out_ty.clone(), opds.collect(), attrs);
651        new_results.extend(new_op.deref(ctx).results());
652        new_op.insert_before(ctx, op);
653    }
654
655    if !new_results.is_empty() {
656        state
657            .mappings
658            .insert(op.deref(ctx).get_result(0), new_results);
659    }
660    state.to_erase.push(op);
661}
662
663fn should_unroll(ctx: &Context, value: impl Typed, max_vector_size: usize) -> bool {
664    let ty = value.get_type(ctx).deref(ctx);
665    if !type_impls::<dyn UnrollableType>(&*ty) {
666        return false;
667    }
668    let Some(vector_size) = value.try_get_vector_size(ctx) else {
669        return false;
670    };
671    vector_size > max_vector_size
672}
673
674fn try_get_vec(ctx: &Context, value: impl Typed) -> usize {
675    value.try_get_vector_size(ctx).unwrap_or(1)
676}
677
678fn unroll_ty(ctx: &Context, ty: impl Typed, vectorization: usize) -> TypeHandle {
679    let ty = ty.get_type(ctx).deref(ctx);
680    type_cast::<dyn UnrollableType>(&*ty)
681        .expect("Should be unrollable")
682        .with_vector_size(ctx, vectorization)
683}
cubecl_core/post_processing/unroll.rs

cubecl_core/post_processing/
unroll.rs
