cubecl_std/tensor/layout/
base.rs

1use cubecl::prelude::*;
2use cubecl_core::{self as cubecl};
3
4use super::Coordinates;
5
6/// A layout that represents the mapping from a conceptual multi-dimensional tensor to a linear
7/// storage. Some layouts may be transformers meant to be composed with others (i.e. swizzling),
8/// others will represent the actual underlying structure of the data.
9#[cube(expand_base_traits = "Clone")]
10pub trait Layout {
11    /// The coordinate type used by the conceptual tensor represented by this layout, i.e.
12    /// `(u32, u32, u32)` for a fixed-rank 3D tensor.
13    /// This does not have to match the rank of the underlying storage (if applicable).
14    /// It's only how the tensor is interpreted (viewed) by the code.
15    type Coordinates: Coordinates;
16    /// The coordinate type used by the inner storage wrapped in this layout, i.e. `u32` for
17    /// `Array`, or `(u32, u32)` for a 2D view.
18    type SourceCoordinates: Coordinates;
19
20    /// Transform a set of n-dimensional coordinates to a source coordinate space.
21    /// It is recommended to use absolute positions here, and handle the translation into lines
22    /// at the lowest level (global memory layout).
23    fn to_source_pos(&self, pos: Self::Coordinates) -> Self::SourceCoordinates;
24    /// Transform a set of n-dimensional coordinates to an offset into the underlying storage,
25    /// and return whether the position is in bounds of this layout.
26    /// See also [Layout::to_source_pos]
27    fn to_source_pos_checked(&self, pos: Self::Coordinates) -> (Self::SourceCoordinates, bool);
28    /// The shape of the conceptual tensor represented by this layout. Not necessarily the extent
29    /// of the underlying storage, but only this view of it.
30    fn shape(&self) -> Self::Coordinates;
31    fn is_in_bounds(&self, pos: Self::Coordinates) -> bool;
32}
cubecl_std/tensor/layout/base.rs

cubecl_std/tensor/layout/
base.rs
