cubecl_core/frontend/
base.rs

1use cubecl_ir::Scope;
2
3use crate::frontend::{CubePrimitive, NativeExpand, init_mut_of_type};
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: CubePrimitive>(scope: &Scope) -> NativeExpand<T> {
71    init_mut_of_type(scope, T::__expand_as_type(scope)).into()
72}
73
74#[cfg(test)]
75mod tests {
76    use crate as cubecl;
77    use crate::prelude::*;
78    use cubecl_ir::{
79        AddressType, FloatKind,
80        settings::{Dim3, ExecutionMode, KernelSettings},
81    };
82
83    /// What `T` names after the call one depth down has named it again: `f16` at the bottom,
84    /// `f32` above it.
85    #[cube]
86    fn named_after_reentry(#[comptime] depth: u32) -> comptime_type!(ElemType) {
87        let define!(T) = comptime!(match depth {
88            0 => ElemType::Float(FloatKind::F16),
89            _ => ElemType::Float(FloatKind::F32),
90        });
91        if comptime!(depth > 0) {
92            named_after_reentry(comptime!(depth - 1));
93        }
94        elem_type_of::<T>()
95    }
96
97    /// A function re-entered on another element finds its own definition again when the inner
98    /// call returns: an op on one tile that runs itself on a tile of another element must not
99    /// go on at the inner element.
100    #[test]
101    fn a_definition_lasts_to_the_end_of_its_block() {
102        let scope = Scope::root(KernelSettings::new(
103            Dim3::new_single(),
104            ExecutionMode::Checked,
105            AddressType::U32,
106        ));
107        let named = named_after_reentry::expand(&scope, 1);
108        assert_eq!(named, ElemType::Float(FloatKind::F32));
109    }
110
111    /// What `N` sizes after the call one depth down has sized it again: two at the bottom, four
112    /// above it.
113    #[cube]
114    fn sized_after_reentry(#[comptime] depth: u32) -> comptime_type!(usize) {
115        let size!(N) = comptime!(match depth {
116            0 => 2usize,
117            _ => 4usize,
118        });
119        if comptime!(depth > 0) {
120            sized_after_reentry(comptime!(depth - 1));
121        }
122        intrinsic!(|scope| N::__expand_value(scope))
123    }
124
125    /// A size's definition lasts to the end of its block as a type's does: a function re-entered
126    /// on another size finds its own again when the inner call returns.
127    #[test]
128    fn a_size_lasts_to_the_end_of_its_block() {
129        let scope = Scope::root(KernelSettings::new(
130            Dim3::new_single(),
131            ExecutionMode::Checked,
132            AddressType::U32,
133        ));
134        assert_eq!(sized_after_reentry::expand(&scope, 1), 4);
135    }
136}
cubecl_core/frontend/base.rs

cubecl_core/frontend/
base.rs
