cubecl_core/frontend/
base.rs

1use cubecl_ir::Scope;
2
3use crate::frontend::{ConstRuntimeAssign, CubeType, NativeExpand};
4
5#[macro_export]
6macro_rules! unexpanded {
7    () => ({
8        panic!("Unexpanded Cube functions should not be called. ");
9    });
10    ($msg:expr) => ({
11        panic!($msg);
12    });
13    ($fmt:expr, $($arg:tt)*) => ({
14        panic!($fmt, $($arg)*);
15    });
16}
17
18#[macro_export]
19macro_rules! expand_error {
20    () => ({
21        panic!("An error occurred during kernel expansion");
22    });
23    ($msg:expr) => ({
24        panic!(concat!("An error occurred during kernel expansion:\n", $msg));
25    });
26    ($fmt:expr, $($arg:tt)*) => ({
27        panic!(concat!("An error occurred during kernel expansion:\n", $fmt), $($arg)*);
28    });
29}
30
31#[macro_export]
32macro_rules! expand_assert {
33    ($cond:expr) => ({
34        assert!($cond, "An error occurred during kernel expansion");
35    });
36    ($cond:expr, $msg:expr) => ({
37        assert!($cond, concat!("An error occurred during kernel expansion:\n", $msg));
38    });
39    ($cond:expr, $fmt:expr, $($arg:tt)*) => ({
40        assert!($cond, concat!("An error occurred during kernel expansion:\n", $fmt), $($arg)*);
41    });
42}
43
44/// `let size!(N) = size;` defines a size `N` to the end of its block. A function re-entered on
45/// another size gives `N` back to its outer call when the inner one returns.
46#[macro_export]
47macro_rules! size {
48    ($name: ident) => {
49        $name
50    };
51}
52
53/// `let define!(T) = elem;` defines a type `T` naming `elem` to the end of its block. A
54/// function re-entered on another element gives `T` back to its outer call when the inner one
55/// returns.
56#[macro_export]
57macro_rules! define {
58    ($name: ident) => {
59        $name
60    };
61}
62
63pub fn unexpanded_value<T>() -> T {
64    unexpanded!()
65}
66
67/// Used for `let x;` local bindings. These are technically immutable in Rust, but because they can
68/// be visible at a higher level scope than their assignment we need to allocate a slot of them at
69/// that level and treat them as mutable.
70pub fn uninit_local<T: CubeType>(scope: &Scope) -> NativeExpand<T>
71where
72    T::ExpandType: ConstRuntimeAssign<Expand = NativeExpand<T>>,
73{
74    T::ExpandType::const_init_mut(scope)
75}
76
77#[cfg(test)]
78mod tests {
79    use crate as cubecl;
80    use crate::prelude::*;
81    use cubecl_ir::{
82        AddressType, FloatKind,
83        settings::{Dim3, ExecutionMode, KernelSettings},
84    };
85
86    /// What `T` names after the call one depth down has named it again: `f16` at the bottom,
87    /// `f32` above it.
88    #[cube]
89    fn named_after_reentry(#[comptime] depth: u32) -> comptime_type!(ElemType) {
90        let define!(T) = comptime!(match depth {
91            0 => ElemType::Float(FloatKind::F16),
92            _ => ElemType::Float(FloatKind::F32),
93        });
94        if comptime!(depth > 0) {
95            named_after_reentry(comptime!(depth - 1));
96        }
97        elem_type_of::<T>()
98    }
99
100    /// A function re-entered on another element finds its own definition again when the inner
101    /// call returns: an op on one tile that runs itself on a tile of another element must not
102    /// go on at the inner element.
103    #[test]
104    fn a_definition_lasts_to_the_end_of_its_block() {
105        let scope = Scope::root(KernelSettings::new(
106            Dim3::new_single(),
107            ExecutionMode::Checked,
108            AddressType::U32,
109        ));
110        let named = named_after_reentry::expand(&scope, 1);
111        assert_eq!(named, ElemType::Float(FloatKind::F32));
112    }
113
114    /// What `N` sizes after the call one depth down has sized it again: two at the bottom, four
115    /// above it.
116    #[cube]
117    fn sized_after_reentry(#[comptime] depth: u32) -> comptime_type!(usize) {
118        let size!(N) = comptime!(match depth {
119            0 => 2usize,
120            _ => 4usize,
121        });
122        if comptime!(depth > 0) {
123            sized_after_reentry(comptime!(depth - 1));
124        }
125        intrinsic!(|scope| N::__expand_value(scope))
126    }
127
128    /// A size's definition lasts to the end of its block as a type's does: a function re-entered
129    /// on another size finds its own again when the inner call returns.
130    #[test]
131    fn a_size_lasts_to_the_end_of_its_block() {
132        let scope = Scope::root(KernelSettings::new(
133            Dim3::new_single(),
134            ExecutionMode::Checked,
135            AddressType::U32,
136        ));
137        assert_eq!(sized_after_reentry::expand(&scope, 1), 4);
138    }
139}
cubecl_core/frontend/base.rs

cubecl_core/frontend/
base.rs
