cubecl_ir/
operation.rs

1use core::fmt::Display;
2
3use super::{Branch, CoopMma, NonSemantic, Plane, Synchronization, Type, Value};
4use crate::{
5    AddressSpace, Arithmetic, AtomicOp, Bitwise, Id, InstructionModes, Memory, Metadata,
6    OperationArgs, OperationReflect, Operator, Scope, TensorIndexingOps, TmaOps,
7    comparison::Comparison, marker::Marker,
8};
9use crate::{BarrierOps, SourceLoc, TypeHash};
10use alloc::{
11    format,
12    string::{String, ToString},
13    vec::Vec,
14};
15use derive_more::derive::From;
16use itertools::Itertools;
17
18/// All operations that can be used in a GPU compute shader.
19///
20/// Notes:
21///
22/// [Operator] can be vectorized, but other operations can't.
23/// Therefore, during tracing, only operators can be registered.
24///
25#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
26#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, From, OperationReflect)]
27#[operation(opcode_name = OpCode)]
28#[allow(dead_code, missing_docs, clippy::large_enum_variant)] // Some variants might not be used with different flags
29pub enum Operation {
30    #[operation(pure)]
31    #[from(ignore)]
32    Copy(#[args(allow_ptr)] Value),
33    /// Construct an aggregate (i.e. fat pointer) that's later disaggregated into normal values
34    /// supported by codegen. Not allowed to exist after the disaggregation pass.
35    ConstructAggregate(#[args(allow_ptr)] Vec<Value>),
36    /// Extract a specific field from an aggregate (i.e. read the length from a slice ptr).
37    /// Not allowed to exist after the disaggregation pass.
38    ExtractAggregateField(AggregateExtractOperands),
39    DeclareVariable {
40        #[args(skip)]
41        value_ty: Type,
42        #[args(skip)]
43        addr_space: AddressSpace,
44        #[args(skip)]
45        alignment: usize,
46    },
47    #[operation(nested)]
48    Memory(Memory),
49    #[operation(nested)]
50    Arithmetic(Arithmetic),
51    #[operation(nested)]
52    Comparison(Comparison),
53    #[operation(nested)]
54    Bitwise(Bitwise),
55    #[operation(nested)]
56    Operator(Operator),
57    #[operation(nested)]
58    Atomic(AtomicOp),
59    #[operation(nested)]
60    Metadata(Metadata),
61    #[operation(nested)]
62    Branch(Branch),
63    #[operation(nested)]
64    Synchronization(Synchronization),
65    /// Barrier followed by a load whose result is workgroup-uniform.
66    ///
67    /// Mirrors WGSL's `workgroupUniformLoad`: the input is a reference into
68    /// workgroup-shared memory (`&list[i]`, produced by [`Memory::Index`]),
69    /// and both the non-atomic and atomic WGSL overloads map here. Other
70    /// backends lower this to a `sync_cube` followed by a regular (or atomic)
71    /// load — uniformity is implicit there.
72    #[from(ignore)]
73    WorkgroupUniformLoad(#[args(allow_ptr, ptr_read)] Value),
74    #[operation(nested)]
75    Plane(Plane),
76    #[operation(nested)]
77    CoopMma(CoopMma),
78    #[operation(nested)]
79    Barrier(BarrierOps),
80    #[operation(nested)]
81    Tma(TmaOps),
82    #[operation(nested)]
83    TensorIndexing(TensorIndexingOps),
84    /// Non-semantic instructions (i.e. comments, debug info)
85    #[operation(nested)]
86    NonSemantic(NonSemantic),
87    // Markers used by compilers to update state or modes, but don't emit instructions
88    #[operation(nested)]
89    Marker(Marker),
90}
91
92/// An instruction that contains a right hand side [`Operation`] and an optional out variable.
93#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
94#[derive(Debug, Clone, PartialEq, Eq, Hash, TypeHash)]
95pub struct Instruction {
96    pub out: Option<Value>,
97    pub source_loc: Option<SourceLoc>,
98    pub modes: InstructionModes,
99    pub operation: Operation,
100}
101
102impl Instruction {
103    pub fn new(operation: impl Into<Operation>, out: Value) -> Self {
104        Instruction {
105            out: Some(out),
106            operation: operation.into(),
107            source_loc: None,
108            modes: Default::default(),
109        }
110    }
111
112    pub fn no_out(operation: impl Into<Operation>) -> Self {
113        Instruction {
114            out: None,
115            operation: operation.into(),
116            source_loc: None,
117            modes: Default::default(),
118        }
119    }
120
121    #[track_caller]
122    pub fn out(&self) -> Value {
123        self.out.unwrap()
124    }
125
126    #[track_caller]
127    pub fn ty(&self) -> Type {
128        self.out().ty
129    }
130}
131
132impl Display for Instruction {
133    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
134        match &self.operation {
135            Operation::Operator(Operator::Cast(op)) => {
136                write!(
137                    f,
138                    "{} = cast<{}>({}) : ({}) -> ({})",
139                    self.out(),
140                    self.ty(),
141                    op.input,
142                    op.input.ty,
143                    self.out().ty,
144                )
145            }
146            Operation::Operator(Operator::Reinterpret(op)) => {
147                write!(f, "{} = bitcast<{}>({})", self.out(), self.ty(), op.input)
148            }
149            _ => {
150                if let Some(out) = self.out {
151                    let mut vars_str = String::new();
152                    for (i, val) in self.operation.args().unwrap_or_default().iter().enumerate() {
153                        if i != 0 {
154                            vars_str.push_str(", ");
155                        }
156                        vars_str.push_str(&val.ty.to_string());
157                    }
158                    write!(
159                        f,
160                        "{out} = {} : ({}) -> ({})",
161                        self.operation, vars_str, out.ty
162                    )
163                } else {
164                    write!(f, "{}", self.operation)
165                }
166            }
167        }
168    }
169}
170
171impl Display for Operation {
172    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
173        match self {
174            Operation::Copy(variable) => write!(f, "{variable}"),
175            Operation::ConstructAggregate(variables) => {
176                write!(f, "aggregate({})", variables.iter().join(", "))
177            }
178            Operation::ExtractAggregateField(AggregateExtractOperands { aggregate, field }) => {
179                write!(f, "extract({aggregate}, {field})")
180            }
181            Operation::DeclareVariable {
182                value_ty,
183                addr_space,
184                alignment,
185            } => {
186                write!(
187                    f,
188                    "declare_var<{value_ty}>({addr_space}, align: {alignment})"
189                )
190            }
191            Operation::Memory(memory) => write!(f, "{memory}"),
192            Operation::Arithmetic(arithmetic) => write!(f, "{arithmetic}"),
193            Operation::Comparison(comparison) => write!(f, "{comparison}"),
194            Operation::Bitwise(bitwise) => write!(f, "{bitwise}"),
195            Operation::Operator(operator) => write!(f, "{operator}"),
196            Operation::Atomic(atomic) => write!(f, "{atomic}"),
197            Operation::Metadata(metadata) => write!(f, "{metadata}"),
198            Operation::Branch(branch) => write!(f, "{branch}"),
199            Operation::Synchronization(synchronization) => write!(f, "{synchronization}"),
200            Operation::WorkgroupUniformLoad(val) => {
201                write!(f, "workgroup_uniform_load({val})")
202            }
203            Operation::Plane(plane) => write!(f, "{plane}"),
204            Operation::CoopMma(coop_mma) => write!(f, "{coop_mma}"),
205            Operation::NonSemantic(non_semantic) => write!(f, "{non_semantic}"),
206            Operation::Barrier(barrier_ops) => write!(f, "{barrier_ops}"),
207            Operation::Tma(tma_ops) => write!(f, "{tma_ops}"),
208            Operation::TensorIndexing(ops) => write!(f, "{ops}"),
209            Operation::Marker(marker) => write!(f, "{marker}"),
210        }
211    }
212}
213
214pub fn fmt_vararg(args: &[impl Display]) -> String {
215    if args.is_empty() {
216        "".to_string()
217    } else {
218        let str = args
219            .iter()
220            .map(|it| it.to_string())
221            .collect::<Vec<_>>()
222            .join(", ");
223        format!(", {str}")
224    }
225}
226
227#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
228#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationArgs)]
229#[allow(missing_docs)]
230pub struct IndexOperands {
231    #[args(allow_ptr)]
232    pub list: Value,
233    pub index: Value,
234    pub unroll_factor: usize, // Adjustment factor for bounds check
235    pub checked: bool,
236}
237
238#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
239#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationArgs)]
240#[allow(missing_docs)]
241pub struct StoreOperands {
242    #[args(allow_ptr, ptr_write)]
243    pub ptr: Value,
244    pub value: Value,
245}
246
247#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
248#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationArgs)]
249#[allow(missing_docs)]
250pub struct BinaryOperands {
251    pub lhs: Value,
252    pub rhs: Value,
253}
254
255#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
256#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationArgs)]
257#[allow(missing_docs)]
258pub struct AtomicBinaryOperands {
259    #[args(allow_ptr, ptr_read, ptr_write)]
260    pub ptr: Value,
261    pub value: Value,
262}
263
264#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
265#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationArgs)]
266#[allow(missing_docs)]
267pub struct AggregateExtractOperands {
268    #[args(allow_ptr)]
269    pub aggregate: Value,
270    pub field: usize,
271}
272
273/// Closure passed to an intrinsic
274#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
275#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash)]
276pub struct Function {
277    /// Explicit parameters passed to the function. Does not contain closure captures.
278    pub explicit_params: Vec<Value>,
279    /// Scope containing closure instructions. Unknown values that aren't explicit params are
280    /// assumed to be captures.
281    pub scope: Scope,
282}
283
284/// Closures are functions invoked by the runtime, not us. So we only use the Id, no params.
285pub type Closure = Id;
286
287impl Display for Function {
288    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
289        let params = self.explicit_params.iter().join(", ");
290        let instructions = self.scope.instructions.borrow().iter().join("\n");
291        write!(f, "|{params}| {{\n{instructions}\n}}")
292    }
293}
294
295#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
296#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationArgs)]
297#[allow(missing_docs)]
298pub struct UnaryOperands {
299    pub input: Value,
300}
301
302impl From<Branch> for Instruction {
303    fn from(value: Branch) -> Self {
304        Instruction::no_out(value)
305    }
306}
307
308impl From<Synchronization> for Instruction {
309    fn from(value: Synchronization) -> Self {
310        Instruction::no_out(value)
311    }
312}
313
314impl From<NonSemantic> for Instruction {
315    fn from(value: NonSemantic) -> Self {
316        Instruction::no_out(value)
317    }
318}
cubecl_ir/operation.rs

cubecl_ir/
operation.rs
