cubecl-opt 0.11.0-pre.4

Compiler optimizations for CubeCL
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56

use alloc::boxed::Box;

use cubecl_ir::{
    attributes::{
        ATTR_BUFFER_BINDING, ATTR_BUFFER_IO, BufferBindingAttr, BufferIOAttr, FuncInterface,
    },
    prelude::*,
};
use pliron::builtin::ops::FuncOp;

use crate::analyses::pointer_source::GlobalVisibility;

/// Annotate global buffers with the correct visibility. Should be run on a module scope,
/// so non-inlined functions can contribute to the analysis and be annotated.
pub struct AnnotateGlobalVisibilityPass;

#[pass_name]
impl Pass for AnnotateGlobalVisibilityPass {
    fn run(
        &mut self,
        op: Ptr<Operation>,
        ctx: &mut Context,
        analyses: &mut AnalysisManager,
    ) -> Result<PassResult> {
        let global_visibility = analyses.get_analysis::<GlobalVisibility>(op, ctx)?;

        visit_all_ops_of_type::<FuncOp, _>(
            ctx,
            &mut &global_visibility,
            op,
            |ctx, global_visibility, func| {
                let num_args = func.get_entry_block(ctx).deref(ctx).get_num_arguments();
                for i in 0..num_args {
                    let buffer_binding = func
                        .get_arg_attr::<BufferBindingAttr>(ctx, i, &ATTR_BUFFER_BINDING)
                        .map(|it| *it);
                    if let Some(buffer) = buffer_binding
                        && let Some(visibility) =
                            global_visibility.visibility.get(&buffer.buffer_pos)
                    {
                        let io = match (visibility.readable, visibility.writable) {
                            (true, true) => BufferIOAttr::ReadWrite,
                            (true, false) => BufferIOAttr::ReadOnly,
                            (false, true) => BufferIOAttr::WriteOnly,
                            (false, false) => BufferIOAttr::Dead,
                        };
                        func.set_arg_attr(ctx, i, &ATTR_BUFFER_IO, Box::new(io));
                    }
                }
            },
        );
        let mut res = PassResult::default();
        res.ir_changed = IRStatus::Changed;
        Ok(res)
    }
}