cubecl-ir 0.11.0-pre.3

Intermediate representation 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
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98

use core::cell::Ref;
use std::string::String;

use pliron::{
    builtin::attributes::{StringAttr, UnitAttr},
    opts::dce::SideEffects,
};

use crate::{
    CanMaterialize,
    interfaces::{MemoryEffect, MemoryEffects},
    prelude::*,
};

#[pliron_op(name = "cube.asm",
    format,
    attributes = (cube_asm_asm: StringAttr, cube_asm_pure: UnitAttr, cube_asm_nomem: UnitAttr, cube_asm_readonly: UnitAttr),
    verifier = "succ"
)]
#[op_traits(CanMaterialize)]
pub struct InlineAsmOp;

impl InlineAsmOp {
    pub fn new(
        ctx: &mut Context,
        result_types: Vec<TypeHandle>,
        asm: String,
        arguments: Vec<Value>,
    ) -> Self {
        let op = Operation::new(
            ctx,
            Self::get_concrete_op_info(),
            result_types,
            arguments,
            vec![],
            0,
        );
        let this = Self { op };
        this.set_attr_cube_asm_asm(ctx, asm.into());
        this
    }

    pub fn asm<'a>(&self, ctx: &'a Context) -> Ref<'a, StringAttr> {
        self.get_attr_cube_asm_asm(ctx).unwrap()
    }

    pub fn inputs(&self, ctx: &Context) -> Vec<Value> {
        self.get_operation().operands(ctx)
    }

    pub fn results(&self, ctx: &Context) -> Vec<Value> {
        self.get_operation().results(ctx)
    }

    pub fn pure(&self, ctx: &Context) -> bool {
        self.get_attr_cube_asm_pure(ctx).is_some()
    }

    pub fn set_pure(&self, ctx: &Context) {
        self.set_attr_cube_asm_pure(ctx, UnitAttr::new());
    }

    pub fn nomem(&self, ctx: &Context) -> bool {
        self.get_attr_cube_asm_nomem(ctx).is_some()
    }

    pub fn set_nomem(&self, ctx: &Context) {
        self.set_attr_cube_asm_nomem(ctx, UnitAttr::new());
    }

    pub fn readonly(&self, ctx: &Context) -> bool {
        self.get_attr_cube_asm_readonly(ctx).is_some()
    }

    pub fn set_readonly(&self, ctx: &Context) {
        self.set_attr_cube_asm_readonly(ctx, UnitAttr::new());
    }
}

#[op_interface_impl]
impl SideEffects for InlineAsmOp {
    fn has_side_effects(&self, ctx: &Context) -> bool {
        !self.pure(ctx)
    }
}

#[op_interface_impl]
impl MemoryEffects for InlineAsmOp {
    fn memory_effects(&self, ctx: &Context) -> Vec<MemoryEffect> {
        if self.nomem(ctx) {
            vec![]
        } else if self.readonly(ctx) {
            vec![MemoryEffect::ReadAll]
        } else {
            vec![MemoryEffect::ReadAll, MemoryEffect::WriteAll]
        }
    }
}