cubecl_tile/dialect/
level.rs

1//! `tile.level`: the tiles one instance takes at a level distributed over a scope.
2
3use cubecl::ir::{ExpandValue, OpInserter, Scope};
4use cubecl_ir::{
5    RecursiveMemoryEffects, RecursiveSideEffects, RecursivelySpeculatable,
6    dialect::branch::YieldOp, interfaces::LowerInterface, prelude::*,
7};
8use pliron::{
9    basic_block::BasicBlock, derive::op_interface_impl, irbuild::inserter::OpInsertionPoint,
10    linked_list::ContainsLinkedList,
11};
12
13use crate::CoordsExpand;
14
15/// The tiles one instance takes at a level distributed over a scope: which tile is whose.
16///
17/// Its region computes them from the instance's position in the launch, as tracing wrote it, and
18/// yields the ones known only at runtime; its results stand for them. Its lowering inlines the
19/// region; a backend that runs the level's instances another way lowers it its own way.
20///
21/// With the test-only `sequential-levels` feature, it lowers instead to a loop over the level's
22/// instances, run by the first.
23#[pliron_op(
24    name = "tile.level",
25    format = "`: ` types(CharSpace(`,`)) ` ` region($0)",
26    verifier = "succ"
27)]
28#[op_interfaces(NOpdsInterface<0>, OneRegionInterface, SingleBlockRegionInterface)]
29#[op_traits(RecursiveMemoryEffects, RecursiveSideEffects, RecursivelySpeculatable)]
30pub struct LevelOp;
31
32/// A walk's `(counts, positions, scales)`, as traced.
33type TilesExpand = (
34    CoordsExpand<usize>,
35    CoordsExpand<usize>,
36    CoordsExpand<usize>,
37);
38
39impl LevelOp {
40    /// The `(counts, positions, scales)` `body` traces, inside a `tile.level` op.
41    pub(crate) fn trace(scope: &Scope, body: impl FnOnce(&Scope) -> TilesExpand) -> TilesExpand {
42        let block = BasicBlock::new(scope.ctx_mut(), None, vec![]);
43        let child = scope.branch_child(OpInserter::new_at_block_end(block));
44        let (mut counts, mut positions, mut scales) = body(&child);
45
46        let mut runtime = counts
47            .runtime_mut()
48            .chain(positions.runtime_mut())
49            .chain(scales.runtime_mut())
50            .collect::<Vec<_>>();
51        let yielded = runtime
52            .iter()
53            .map(|v| v.expand.value(&child))
54            .collect::<Vec<_>>();
55        let ctx = scope.ctx_mut();
56        let types = yielded.iter().map(|v| v.get_type(ctx)).collect();
57        let yield_op = Operation::new(
58            ctx,
59            YieldOp::get_concrete_op_info(),
60            vec![],
61            yielded,
62            vec![],
63            0,
64        );
65        yield_op.insert_at_back(block, ctx);
66
67        let op = Operation::new(ctx, Self::get_concrete_op_info(), types, vec![], vec![], 1);
68        block.insert_at_front(op.deref(ctx).get_region(0), ctx);
69        scope.register(&LevelOp { op });
70
71        for (i, v) in runtime.iter_mut().enumerate() {
72            **v = ExpandValue::Value(op.deref(scope.ctx()).get_result(i)).into();
73        }
74        (counts, positions, scales)
75    }
76
77    /// Moves the region's code before the op, and the values it yields in place of the results.
78    fn inline(&self, ctx: &mut Context, rewriter: &mut dyn Rewriter) {
79        let op = self.get_operation();
80        let block = self.get_body(ctx, 0);
81        let ops = block.deref(ctx).iter(ctx).collect::<Vec<_>>();
82        let (terminator, body) = ops.split_last().expect("a level's region ends in a yield");
83        for &inner in body {
84            rewriter.move_operation(ctx, inner, OpInsertionPoint::BeforeOperation(op));
85        }
86        let yielded = terminator.deref(ctx).operands().collect::<Vec<_>>();
87        for (i, value) in yielded.into_iter().enumerate() {
88            let result = op.deref(ctx).get_result(i);
89            rewriter.replace_value_uses_with(ctx, result, value);
90        }
91        rewriter.erase_operation(ctx, op);
92    }
93}
94
95/// The code the region holds, in the op's place: what tracing wrote before the op existed.
96#[op_interface_impl]
97impl LowerInterface for LevelOp {
98    fn lower(&self, ctx: &mut Context, rewriter: &mut MatchRewriter) -> Result<()> {
99        #[cfg(feature = "sequential-levels")]
100        super::sequential::run_on_first_instance(ctx, rewriter, self)?;
101        self.inline(ctx, rewriter);
102        Ok(())
103    }
104}
cubecl_tile/dialect/level.rs

cubecl_tile/dialect/
level.rs
