cubecl_ir/attributes/
entrypoint.rs

1use core::cell::{Ref, RefMut};
2
3use alloc::boxed::Box;
4use derive_new::new;
5use pliron::{
6    attribute::{AttrObj, Attribute, AttributeDict},
7    builtin::{
8        attributes::{DictAttr, UnitAttr, VecAttr},
9        ops::FuncOp,
10    },
11    dict_key,
12    identifier::Identifier,
13};
14
15use crate::{prelude::*, settings::Dim3};
16
17#[pliron_attr(
18    name = "cube.entrypoint_abi",
19    format = "`<cube_dim: ` $cube_dim opt($cluster_dim) `>`",
20    verifier = "succ"
21)]
22#[derive(new, PartialEq, Clone, Debug, Hash)]
23pub struct EntrypointAbiAttr {
24    pub cube_dim: Dim3,
25    pub cluster_dim: Option<Dim3>,
26}
27
28dict_key!(
29    /// Key for entry point attribute when the operation defines an entry point.
30    ATTR_KEY_ENTRY_POINT, "entry_point"
31);
32dict_key!(ATTR_BUFFER_BINDING, "buffer_binding");
33dict_key!(ATTR_TENSOR_MAP_BINDING, "tensor_map_binding");
34
35/// What a compiled kernel does with one buffer binding, stamped on the entry
36/// function's arguments by the visibility analysis and carried to the launch
37/// site, where the write scope stages exactly what the kernel writes and
38/// checks exactly what it reads.
39///
40/// Four states because the launch path asks two independent questions.
41/// Writable decides what a failure taints; readable decides what is checked
42/// before launching. An aliased or accumulating argument answers yes to both,
43/// a pure output answers only the first — which is what lets a relaunch
44/// repair a tainted buffer — and a buffer the kernel never touches answers
45/// neither.
46#[pliron_attr(name = "cube.buffer_io", format, verifier = "succ")]
47#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
48#[derive(new, PartialEq, Eq, Clone, Copy, Debug, Hash)]
49pub enum BufferIOAttr {
50    /// Read, never written: checked before launch, never claimed by a failure.
51    ReadOnly,
52    /// Written, never read: claimed by a failure, never checked.
53    WriteOnly,
54    /// Both — the aliased and accumulator cases included.
55    ReadWrite,
56    /// Neither: no check on the way in, no claim on the way out.
57    Dead,
58}
59
60impl BufferIOAttr {
61    /// Whether the kernel reads the buffer, so its contents have to be
62    /// trustworthy before the launch runs.
63    pub fn is_readable(&self) -> bool {
64        matches!(self, BufferIOAttr::ReadOnly | BufferIOAttr::ReadWrite)
65    }
66
67    /// Whether the kernel writes the buffer, so a failed launch leaves it
68    /// carrying the failure.
69    pub fn is_writable(&self) -> bool {
70        matches!(self, BufferIOAttr::WriteOnly | BufferIOAttr::ReadWrite)
71    }
72
73    /// Whether the kernel never touches the buffer at all.
74    pub fn is_dead(&self) -> bool {
75        matches!(self, BufferIOAttr::Dead)
76    }
77}
78
79dict_key!(ATTR_BUFFER_IO, "binding_io");
80
81#[pliron_attr(
82    name = "cube.buffer_binding",
83    format = "`<(` $buffer_pos `, ` opt($ext_meta_pos) `)`",
84    verifier = "succ"
85)]
86#[derive(new, PartialEq, Clone, Copy, Debug, Hash)]
87pub struct BufferBindingAttr {
88    pub buffer_pos: usize,
89    pub ext_meta_pos: Option<usize>,
90}
91
92/// [Op] that may define an entry point.
93///
94/// ### Attribute(s):
95/// | Name | Static Name Identifier | Type |
96/// |------|------------------------| -----|
97/// | entry_point | [ATTR_KEY_ENTRY_POINT] | [EntrypointAbiAttr](crate::attributes::EntrypointAbiAttr) |
98#[op_interface]
99pub trait EntrypointInterface {
100    verify_op_succ!();
101
102    /// Get the entry point ABI defined by this operation.
103    fn get_entrypoint_abi(&self, ctx: &Context) -> Option<EntrypointAbiAttr> {
104        let self_op = self.get_operation().deref(ctx);
105        let s_attr = self_op
106            .attributes
107            .get::<EntrypointAbiAttr>(&ATTR_KEY_ENTRY_POINT);
108        s_attr.cloned()
109    }
110
111    /// Set the entry point ABI defined by this operation.
112    fn set_entrypoint_abi(&self, ctx: &mut Context, entry_point_abi: EntrypointAbiAttr) {
113        let mut self_op = self.get_operation().deref_mut(ctx);
114        self_op
115            .attributes
116            .set(ATTR_KEY_ENTRY_POINT.clone(), entry_point_abi);
117    }
118}
119
120#[op_interface_impl]
121impl EntrypointInterface for FuncOp {}
122
123dict_key!(ATTR_KEY_ARG_ATTRS, "arg_attrs");
124dict_key!(ATTR_KEY_RES_ATTRS, "res_attrs");
125
126/// [Op] that may define an entry point.
127///
128/// ### Attribute(s):
129/// | Name | Static Name Identifier | Type |
130/// |------|------------------------| -----|
131/// | arg_attrs | [ATTR_KEY_ARG_ATTRS] | [VecAttr](pliron::builtin::attributes::VecAttr) |
132/// | res_attrs | [ATTR_KEY_RES_ATTRS] | [VecAttr](pliron::builtin::attributes::VecAttr) |
133pub trait FuncInterface: Op {
134    fn get_arg_attrs<'a>(&self, ctx: &'a Context, arg_idx: usize) -> Option<Ref<'a, DictAttr>> {
135        let self_op = self.get_operation().deref(ctx);
136        Ref::filter_map(self_op, |self_op| {
137            let args_attrs = self_op.attributes.get::<VecAttr>(&ATTR_KEY_ARG_ATTRS)?;
138            args_attrs.0.get(arg_idx)?.downcast_ref()
139        })
140        .ok()
141    }
142
143    fn get_arg_attr<'a, T: Attribute>(
144        &self,
145        ctx: &'a Context,
146        arg_idx: usize,
147        key: &Identifier,
148    ) -> Option<Ref<'a, T>> {
149        let arg_attrs = self.get_arg_attrs(ctx, arg_idx)?;
150        Ref::filter_map(arg_attrs, |arg_attrs| {
151            arg_attrs.lookup(key).and_then(|attr| attr.downcast_ref())
152        })
153        .ok()
154    }
155
156    fn has_arg_attr(&self, ctx: &Context, arg_idx: usize, key: &Identifier) -> bool {
157        self.get_arg_attrs(ctx, arg_idx)
158            .is_some_and(|arg_attrs| arg_attrs.lookup(key).is_some())
159    }
160
161    fn get_arg_attrs_mut<'a>(&self, ctx: &'a Context, arg_idx: usize) -> RefMut<'a, DictAttr> {
162        let self_op = self.get_operation().deref_mut(ctx);
163        RefMut::map(self_op, |self_op| {
164            get_arg_or_init_mut(&mut self_op.attributes, arg_idx)
165        })
166    }
167
168    fn get_res_attrs<'a>(&self, ctx: &'a Context, res_idx: usize) -> Option<Ref<'a, DictAttr>> {
169        let self_op = self.get_operation().deref(ctx);
170        Ref::filter_map(self_op, |self_op| {
171            let res_attrs = self_op.attributes.get::<VecAttr>(&ATTR_KEY_RES_ATTRS)?;
172            res_attrs.0.get(res_idx)?.downcast_ref()
173        })
174        .ok()
175    }
176
177    fn get_res_attr<'a, T: Attribute>(
178        &self,
179        ctx: &'a Context,
180        res_idx: usize,
181        key: &Identifier,
182    ) -> Option<Ref<'a, T>> {
183        let res_attrs = self.get_res_attrs(ctx, res_idx)?;
184        Ref::filter_map(res_attrs, |res_attrs| {
185            res_attrs.lookup(key).and_then(|attr| attr.downcast_ref())
186        })
187        .ok()
188    }
189
190    fn has_res_attr(&self, ctx: &Context, res_idx: usize, key: &Identifier) -> bool {
191        self.get_res_attrs(ctx, res_idx)
192            .is_some_and(|res_attrs| res_attrs.lookup(key).is_some())
193    }
194
195    fn get_res_attrs_mut<'a>(&self, ctx: &'a Context, res_idx: usize) -> RefMut<'a, DictAttr> {
196        let self_op = self.get_operation().deref_mut(ctx);
197        RefMut::map(self_op, |self_op| {
198            get_res_or_init_mut(&mut self_op.attributes, res_idx)
199        })
200    }
201
202    fn set_arg_attrs(&self, ctx: &Context, arg_idx: usize, dict: DictAttr) {
203        *self.get_arg_attrs_mut(ctx, arg_idx) = dict;
204    }
205
206    fn set_arg_attr(&self, ctx: &Context, arg_idx: usize, key: &Identifier, value: AttrObj) {
207        self.get_arg_attrs_mut(ctx, arg_idx).insert(key, value);
208    }
209
210    fn set_arg_attr_unit(&self, ctx: &Context, arg_idx: usize, key: &Identifier) {
211        self.get_arg_attrs_mut(ctx, arg_idx)
212            .insert(key, Box::new(UnitAttr::new()));
213    }
214
215    fn remove_arg_attr(&self, ctx: &Context, arg_idx: usize, key: &Identifier) {
216        self.get_arg_attrs_mut(ctx, arg_idx).remove(key);
217    }
218
219    fn set_res_attrs(&self, ctx: &Context, res_idx: usize, dict: DictAttr) {
220        *self.get_res_attrs_mut(ctx, res_idx) = dict;
221    }
222
223    fn set_res_attr(&self, ctx: &Context, res_idx: usize, key: &Identifier, value: AttrObj) {
224        self.get_res_attrs_mut(ctx, res_idx).insert(key, value);
225    }
226
227    fn remove_res_attr(&self, ctx: &Context, res_idx: usize, key: &Identifier) {
228        self.get_res_attrs_mut(ctx, res_idx).remove(key);
229    }
230
231    fn set_res_attr_unit(&self, ctx: &Context, res_idx: usize, key: &Identifier) {
232        self.get_res_attrs_mut(ctx, res_idx)
233            .insert(key, Box::new(UnitAttr::new()));
234    }
235}
236
237/// The per-buffer IO the visibility pass stamped on `func`'s arguments,
238/// indexed by buffer position — the vector a compiled kernel carries to the
239/// launch site.
240///
241/// Every argument with a [`BufferBindingAttr`] contributes. One without a
242/// stamped [`BufferIOAttr`] answers `ReadWrite`: that covers a tensor map,
243/// whose global side the visibility analysis cannot see, and any argument the
244/// pass never reached — the conservative direction for the launch path, since
245/// over-claiming costs a spurious loud failure and under-claiming costs a
246/// silent clean read of garbage.
247pub fn buffer_io_by_position(ctx: &Context, func: FuncOp) -> alloc::vec::Vec<BufferIOAttr> {
248    let num_args = func.get_entry_block(ctx).deref(ctx).get_num_arguments();
249    let mut io = alloc::vec::Vec::new();
250    for arg in 0..num_args {
251        let binding = func
252            .get_arg_attr::<BufferBindingAttr>(ctx, arg, &ATTR_BUFFER_BINDING)
253            .map(|it| *it);
254        let Some(binding) = binding else {
255            continue;
256        };
257        let stamped = func
258            .get_arg_attr::<BufferIOAttr>(ctx, arg, &ATTR_BUFFER_IO)
259            .map(|it| *it)
260            .unwrap_or(BufferIOAttr::ReadWrite);
261        if io.len() <= binding.buffer_pos {
262            io.resize(binding.buffer_pos + 1, BufferIOAttr::ReadWrite);
263        }
264        io[binding.buffer_pos] = stamped;
265    }
266    io
267}
268
269fn get_arg_or_init_mut(dict: &mut AttributeDict, arg_idx: usize) -> &mut DictAttr {
270    let args_attrs = attr_get_or_insert_mut(dict, &ATTR_KEY_ARG_ATTRS, || VecAttr::new(vec![]))
271        .expect("Should be `VecAttr`");
272    vec_get_or_insert_mut(args_attrs, arg_idx, || DictAttr::new(vec![]))
273        .expect("Should be `DictAttr`")
274}
275
276fn get_res_or_init_mut(dict: &mut AttributeDict, arg_idx: usize) -> &mut DictAttr {
277    let res_attrs = attr_get_or_insert_mut(dict, &ATTR_KEY_RES_ATTRS, || VecAttr::new(vec![]))
278        .expect("Should be `VecAttr`");
279    vec_get_or_insert_mut(res_attrs, arg_idx, || DictAttr::new(vec![]))
280        .expect("Should be `DictAttr`")
281}
282
283fn attr_get_or_insert_mut<'a, T: Attribute>(
284    dict: &'a mut AttributeDict,
285    key: &Identifier,
286    init: impl FnOnce() -> T,
287) -> Option<&'a mut T> {
288    if !dict.0.contains_key(key) {
289        dict.set(key.clone(), init());
290    }
291    dict.get_mut(key)
292}
293
294fn vec_get_or_insert_mut<T: Attribute>(
295    vec: &mut VecAttr,
296    idx: usize,
297    mut init: impl FnMut() -> T,
298) -> Option<&mut T> {
299    if vec.0.len() <= idx {
300        vec.0.resize_with(idx + 1, || Box::new(init()));
301    }
302    vec.0[idx].downcast_mut()
303}
304
305impl FuncInterface for FuncOp {}
cubecl_ir/attributes/entrypoint.rs

cubecl_ir/attributes/
entrypoint.rs
