cubecl_ir/dialect/
asm.rs

1use core::cell::Ref;
2use std::string::String;
3
4use derive_more::From;
5use derive_new::new;
6use pliron::{
7    builtin::attributes::{StringAttr, UnitAttr},
8    opts::dce::SideEffects,
9};
10
11use crate::{
12    AddressSpaceVecAttr, CanMaterialize,
13    interfaces::{MemoryEffect, MemoryEffects},
14    prelude::*,
15    typed_vec_attr,
16};
17
18#[format]
19#[derive(PartialEq, Eq, Hash, Clone, Debug)]
20pub enum InputKind {
21    /// Regular input
22    In,
23    /// Input that reads memory through the value
24    MemIn,
25    /// Input that writes memory through the value.
26    MemOut,
27    /// Input that reads and writes memory through the value.
28    MemInout,
29}
30
31#[pliron_attr(name = "asm.reg_class", format, verifier = "succ")]
32#[derive(PartialEq, Eq, Hash, Clone, Debug)]
33pub enum RegSpec {
34    /// Default register class for the type
35    Inferred,
36    /// Custom class mainly for CPU (i.e. `xmm_reg`)
37    Class(String),
38    /// Explicit register mainly for CPU (i.e. `"ax"`)
39    Explicit(String),
40}
41
42#[pliron_attr(name = "asm.input_spec", format, verifier = "succ")]
43#[derive(PartialEq, Eq, Hash, Clone, Debug, new)]
44pub struct InputSpec {
45    pub kind: InputKind,
46    pub class: RegSpec,
47}
48
49typed_vec_attr!(RegSpec, "asm.reg_specs", RegSpecsAttr);
50typed_vec_attr!(InputSpec, "asm.input_specs", InputSpecsAttr);
51
52#[pliron_attr(name = "asm.memory_clobbers", format, verifier = "succ")]
53#[derive(PartialEq, Eq, Hash, Clone, Debug, From)]
54pub struct MemoryClobbersAttr(pub MemoryClobbers);
55
56#[format]
57#[derive(PartialEq, Eq, Hash, Clone, Debug)]
58pub enum MemoryClobbers {
59    Nomem,
60    Readonly,
61    Explicit {
62        reads_spaces: AddressSpaceVecAttr,
63        writes_spaces: AddressSpaceVecAttr,
64    },
65    ReadWrite,
66}
67
68#[pliron_op(name = "cube.asm",
69    format,
70    attributes = (
71        cube_asm_asm: StringAttr,
72        cube_asm_pure: UnitAttr,
73        cube_asm_memory_clobbers: MemoryClobbersAttr,
74        cube_asm_out_spec: RegSpecsAttr,
75        cube_asm_in_spec: InputSpecsAttr,
76    ),
77    verifier = "succ"
78)]
79#[op_traits(CanMaterialize)]
80pub struct InlineAsmOp;
81
82impl InlineAsmOp {
83    pub fn new(
84        ctx: &mut Context,
85        result_types: Vec<TypeHandle>,
86        out_spec: Vec<RegSpec>,
87        asm: String,
88        memory_clobbers: MemoryClobbers,
89        arguments: Vec<Value>,
90        in_spec: Vec<InputSpec>,
91    ) -> Self {
92        let op = Operation::new(
93            ctx,
94            Self::get_concrete_op_info(),
95            result_types,
96            arguments,
97            vec![],
98            0,
99        );
100        let this = Self { op };
101        this.set_attr_cube_asm_asm(ctx, asm.into());
102        this.set_attr_cube_asm_memory_clobbers(ctx, memory_clobbers.into());
103        this.set_attr_cube_asm_out_spec(ctx, out_spec.into());
104        this.set_attr_cube_asm_in_spec(ctx, in_spec.into());
105        this
106    }
107
108    pub fn asm<'a>(&self, ctx: &'a Context) -> Ref<'a, StringAttr> {
109        self.get_attr_cube_asm_asm(ctx).unwrap()
110    }
111
112    pub fn inputs(&self, ctx: &Context) -> Vec<Value> {
113        self.get_operation().operands(ctx)
114    }
115
116    pub fn results(&self, ctx: &Context) -> Vec<Value> {
117        self.get_operation().results(ctx)
118    }
119
120    pub fn pure(&self, ctx: &Context) -> bool {
121        self.get_attr_cube_asm_pure(ctx).is_some()
122    }
123
124    pub fn set_pure(&self, ctx: &Context) {
125        self.set_attr_cube_asm_pure(ctx, UnitAttr::new());
126    }
127
128    pub fn memory_clobbers<'a>(&self, ctx: &'a Context) -> Ref<'a, MemoryClobbers> {
129        Ref::map(
130            self.get_attr_cube_asm_memory_clobbers(ctx).unwrap(),
131            |attr| &attr.0,
132        )
133    }
134
135    pub fn out_specs(&self, ctx: &Context) -> Vec<RegSpec> {
136        self.get_attr_cube_asm_out_spec(ctx).unwrap().0.clone()
137    }
138
139    pub fn in_specs(&self, ctx: &Context) -> Vec<InputSpec> {
140        self.get_attr_cube_asm_in_spec(ctx).unwrap().0.clone()
141    }
142}
143
144#[op_interface_impl]
145impl SideEffects for InlineAsmOp {
146    fn has_side_effects(&self, ctx: &Context) -> bool {
147        !self.pure(ctx)
148    }
149}
150
151#[op_interface_impl]
152impl MemoryEffects for InlineAsmOp {
153    fn memory_effects(&self, ctx: &Context) -> Vec<MemoryEffect> {
154        match &*self.memory_clobbers(ctx) {
155            MemoryClobbers::Nomem => vec![],
156            MemoryClobbers::Readonly => vec![MemoryEffect::ReadAll],
157            MemoryClobbers::Explicit {
158                reads_spaces,
159                writes_spaces,
160            } => {
161                let mut out = vec![];
162                for space in reads_spaces.0.iter() {
163                    out.push(MemoryEffect::ReadAllInSpace(*space));
164                }
165                for space in writes_spaces.0.iter() {
166                    out.push(MemoryEffect::WriteAllInSpace(*space));
167                }
168                for (value, spec) in self.inputs(ctx).into_iter().zip(self.in_specs(ctx)) {
169                    match spec.kind {
170                        InputKind::MemIn => {
171                            out.push(MemoryEffect::Read(value));
172                        }
173                        InputKind::MemOut => {
174                            out.push(MemoryEffect::Write(value));
175                        }
176                        InputKind::MemInout => {
177                            out.push(MemoryEffect::Read(value));
178                            out.push(MemoryEffect::Write(value));
179                        }
180                        InputKind::In => {}
181                    }
182                }
183                out
184            }
185            MemoryClobbers::ReadWrite => vec![MemoryEffect::ReadAll, MemoryEffect::WriteAll],
186        }
187    }
188}
cubecl_ir/dialect/asm.rs

cubecl_ir/dialect/
asm.rs
