cubecl_opt/passes/
alloc_shared_memory.rs

1use ::pliron::context::Context;
2use cubecl_ir::{
3    AddressSpace, Pure,
4    attributes::IndexAttr,
5    prelude::*,
6    types::{BytesType, PointerType},
7};
8use pliron::{
9    builtin::{attributes::TypeAttr, ops::FuncOp},
10    irbuild::listener::DummyListener,
11};
12
13use crate::SharedLiveness;
14
15#[cube_op(
16    name = "cube.alloc_shared",
17    format = "`size = ` attr($size, $IndexAttr) `, align = ` attr($alignment, $IndexAttr)"
18)]
19#[result_ty(fixed = PointerType::get(ctx, BytesType::get(ctx).into(), AddressSpace::Shared).to_handle())]
20#[op_traits(Pure)]
21pub struct AllocSharedOp {
22    pub size: IndexAttr,
23    pub alignment: IndexAttr,
24}
25
26#[cube_op(
27    name = "cube.slice_shared",
28    format = "$0 `[` attr($offset, $IndexAttr) `] : ` type($0)"
29)]
30#[result_ty(argument)]
31#[op_traits(Pure)]
32pub struct SliceSharedOp {
33    pub block: Value,
34    pub offset: IndexAttr,
35    pub value_ty: TypeAttr,
36}
37
38/// Allocates shared memory as a single block and attaches offsets to shared memory declarations.
39pub struct AllocateSharedMemoryBlockPass;
40
41#[pass_name]
42impl Pass for AllocateSharedMemoryBlockPass {
43    fn run(
44        &mut self,
45        op: Ptr<Operation>,
46        ctx: &mut Context,
47        analyses: &mut AnalysisManager,
48    ) -> Result<PassResult> {
49        let mut res = PassResult::default();
50        let analysis = analyses.get_analysis::<SharedLiveness>(op, ctx)?;
51        let mut rewriter = IRRewriter::<DummyListener>::default();
52        let op = op.dyn_op(ctx).downcast::<FuncOp>().unwrap();
53
54        let allocs = analysis.allocations.values().copied().collect::<Vec<_>>();
55
56        if !analysis.allocations.is_empty() {
57            let size = allocs.iter().map(|it| it.end(ctx)).max().unwrap();
58            let alignment = allocs.iter().map(|it| it.smem.alignment).max().unwrap();
59
60            let entry = op.get_entry_block(ctx);
61            let alloc = AllocSharedOp::new(ctx, size, alignment);
62            alloc.get_operation().insert_at_front(entry, ctx);
63            res.ir_changed |= IRStatus::Changed;
64            let block = alloc.get_result(ctx);
65
66            for alloc in allocs {
67                let declaration = alloc.value.defining_op().expect("Should be op");
68                let ptr_ty = declaration.result(ctx).get_type(ctx);
69                let slice = SliceSharedOp::new(ctx, ptr_ty, block, alloc.offset, alloc.value_ty);
70                slice.get_operation().insert_before(ctx, declaration);
71                rewriter.replace_operation(ctx, declaration, slice.get_operation());
72            }
73        }
74
75        Ok(res)
76    }
77}
cubecl_opt/passes/alloc_shared_memory.rs

cubecl_opt/passes/
alloc_shared_memory.rs
