cubecl_ir/
metadata.rs

1use core::fmt::Display;
2
3use crate::TypeHash;
4
5use crate::{OperationReflect, Variable};
6
7/// All metadata that can be accessed in a shader.
8#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
9#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationReflect)]
10#[operation(opcode_name = MetadataOpCode, pure)]
11#[allow(missing_docs)]
12pub enum Metadata {
13    /// The rank of an array.
14    Rank { var: Variable },
15    /// The stride of an array at the given dimension.
16    Stride { dim: Variable, var: Variable },
17    /// The shape of an array at the given dimension.
18    Shape { dim: Variable, var: Variable },
19    /// The length of an array.
20    Length { var: Variable },
21    /// The length of an array's underlying buffer.
22    BufferLength { var: Variable },
23}
24
25impl Display for Metadata {
26    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
27        match self {
28            Metadata::Rank { var } => write!(f, "rank({var})"),
29            Metadata::Stride { dim, var } => write!(f, "{var}.strides[{dim}]"),
30            Metadata::Shape { dim, var } => write!(f, "{var}.shape[{dim}]"),
31            Metadata::Length { var } => write!(f, "{var}.len()"),
32            Metadata::BufferLength { var } => write!(f, "buffer_len({var})"),
33        }
34    }
35}