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

cubecl_ir/attributes/
entrypoint.rs
