cubecl_ir/dialect/
asm.rs

1use core::cell::Ref;
2use std::string::String;
3
4use pliron::{
5    builtin::attributes::{StringAttr, UnitAttr},
6    opts::dce::SideEffects,
7};
8
9use crate::{
10    CanMaterialize,
11    interfaces::{MemoryEffect, MemoryEffects},
12    prelude::*,
13};
14
15#[pliron_op(name = "cube.asm",
16    format,
17    attributes = (cube_asm_asm: StringAttr, cube_asm_pure: UnitAttr, cube_asm_nomem: UnitAttr, cube_asm_readonly: UnitAttr),
18    verifier = "succ"
19)]
20#[op_traits(CanMaterialize)]
21pub struct InlineAsmOp;
22
23impl InlineAsmOp {
24    pub fn new(
25        ctx: &mut Context,
26        result_types: Vec<TypeHandle>,
27        asm: String,
28        arguments: Vec<Value>,
29    ) -> Self {
30        let op = Operation::new(
31            ctx,
32            Self::get_concrete_op_info(),
33            result_types,
34            arguments,
35            vec![],
36            0,
37        );
38        let this = Self { op };
39        this.set_attr_cube_asm_asm(ctx, asm.into());
40        this
41    }
42
43    pub fn asm<'a>(&self, ctx: &'a Context) -> Ref<'a, StringAttr> {
44        self.get_attr_cube_asm_asm(ctx).unwrap()
45    }
46
47    pub fn inputs(&self, ctx: &Context) -> Vec<Value> {
48        self.get_operation().operands(ctx)
49    }
50
51    pub fn results(&self, ctx: &Context) -> Vec<Value> {
52        self.get_operation().results(ctx)
53    }
54
55    pub fn pure(&self, ctx: &Context) -> bool {
56        self.get_attr_cube_asm_pure(ctx).is_some()
57    }
58
59    pub fn set_pure(&self, ctx: &Context) {
60        self.set_attr_cube_asm_pure(ctx, UnitAttr::new());
61    }
62
63    pub fn nomem(&self, ctx: &Context) -> bool {
64        self.get_attr_cube_asm_nomem(ctx).is_some()
65    }
66
67    pub fn set_nomem(&self, ctx: &Context) {
68        self.set_attr_cube_asm_nomem(ctx, UnitAttr::new());
69    }
70
71    pub fn readonly(&self, ctx: &Context) -> bool {
72        self.get_attr_cube_asm_readonly(ctx).is_some()
73    }
74
75    pub fn set_readonly(&self, ctx: &Context) {
76        self.set_attr_cube_asm_readonly(ctx, UnitAttr::new());
77    }
78}
79
80#[op_interface_impl]
81impl SideEffects for InlineAsmOp {
82    fn has_side_effects(&self, ctx: &Context) -> bool {
83        !self.pure(ctx)
84    }
85}
86
87#[op_interface_impl]
88impl MemoryEffects for InlineAsmOp {
89    fn memory_effects(&self, ctx: &Context) -> Vec<MemoryEffect> {
90        if self.nomem(ctx) {
91            vec![]
92        } else if self.readonly(ctx) {
93            vec![MemoryEffect::ReadAll]
94        } else {
95            vec![MemoryEffect::ReadAll, MemoryEffect::WriteAll]
96        }
97    }
98}
cubecl_ir/dialect/asm.rs

cubecl_ir/dialect/
asm.rs
