cubecl_std/tensor/layout/
as_dyn.rs

1use cubecl::prelude::*;
2use cubecl_core::{self as cubecl, unexpanded};
3use variadics_please::all_tuples;
4
5use crate::tensor::layout::*;
6
7/// Coordinates that can be converted to a dynamic sequence of signed coordinates.
8/// Can be used to convert any set of coordinates to a comptime-sized sequence for use with TMA.
9#[cube]
10pub trait IntoDyn: Coordinates + LaunchArg {
11    fn into_dyn(self) -> Sequence<i32> {
12        unexpanded!()
13    }
14}
15
16macro_rules! as_ty {
17    ($T: ident, $dummy: ident) => {
18        $T
19    };
20}
21
22macro_rules! impl_tuple {
23    ($ty: ident, $($t: ident),*) => {
24        impl IntoDyn for ($(as_ty!($ty, $t)),*) {}
25
26        impl IntoDynExpand for ($(ExpandElementTyped<as_ty!($ty, $t)>),*) {
27            fn __expand_into_dyn_method(self, scope: &mut Scope) -> SequenceExpand<i32> {
28                let mut seq = Sequence::__expand_new(scope);
29                let ($($t),*) = self;
30                let ($($t),*) = ($(i32::__expand_cast_from(scope, $t)),*);
31                $(seq.__expand_push_method(scope, $t);)*
32                seq
33            }
34        }
35    };
36}
37
38macro_rules! impl_tuples {
39    ($($t: ident),*) => {
40        impl_tuple!(u32, $($t),*);
41        impl_tuple!(i32, $($t),*);
42    };
43}
44
45all_tuples!(impl_tuples, 2, 12, t);
46
47#[cube]
48impl IntoDyn for Sequence<i32> {
49    fn into_dyn(self) -> Sequence<i32> {
50        self
51    }
52}
53
54#[cube]
55impl IntoDyn for Sequence<u32> {
56    fn into_dyn(self) -> Sequence<i32> {
57        let mut seq = Sequence::new();
58        for x in self {
59            seq.push(i32::cast_from(x));
60        }
61        seq
62    }
63}
64
65#[derive(CubeType, CubeLaunch)]
66pub struct IntoDynLayout<L: Layout<SourceCoordinates: IntoDyn> + LaunchArg> {
67    layout: L,
68}
69
70impl<L: Layout<SourceCoordinates: IntoDyn> + LaunchArg> IntoDynLayout<L> {
71    pub fn new(layout: L) -> Self {
72        IntoDynLayout { layout }
73    }
74}
75
76#[cube]
77impl<L: Layout<SourceCoordinates: IntoDyn> + LaunchArg> Layout for IntoDynLayout<L> {
78    type Coordinates = L::Coordinates;
79    type SourceCoordinates = Sequence<i32>;
80
81    fn to_source_pos(&self, pos: Self::Coordinates) -> Self::SourceCoordinates {
82        let pos = self.layout.to_source_pos(pos);
83        pos.into_dyn()
84    }
85
86    fn is_in_bounds(&self, pos: Self::Coordinates) -> bool {
87        self.layout.is_in_bounds(pos)
88    }
89
90    fn to_source_pos_checked(&self, pos: Self::Coordinates) -> (Self::SourceCoordinates, bool) {
91        let (pos, in_bounds) = self.layout.to_source_pos_checked(pos);
92        (pos.into_dyn(), in_bounds)
93    }
94
95    fn shape(&self) -> Self::Coordinates {
96        self.layout.shape()
97    }
98}
cubecl_std/tensor/layout/as_dyn.rs

cubecl_std/tensor/layout/
as_dyn.rs
