rename gpu macro

tracel-ai · nathanielsimard · May 15, 2024 · Apr 9, 2024 · Apr 10, 2024 · Apr 10, 2024
commit 0c509759bad452cf33df465161b4c243dc35f817
diff --git a/crates/burn-cube-macros/src/prelude.rs b/crates/burn-cube-macros/src/prelude.rs
@@ -33,8 +33,6 @@ pub(crate) fn get_prelude(needed_functions: &HashSet<VariableKey>) -> proc_macro
 
 fn codegen_float_new() -> proc_macro2::TokenStream {
     quote::quote! {
-        // TODO ENCAPSULATE IMPORTS
-
         pub fn float_new<F: burn_cube::Float>(val: f32) -> F {
             F::new(val, 1)
         }

diff --git a/crates/burn-cube/tests/cast_elem.rs b/crates/burn-cube/tests/cast_elem.rs
@@ -1,6 +1,8 @@
 use burn_cube::{cube, Bool, CubeContext, Float, Int, UInt, F32, I32};
-use burn_jit::gpu;
-use burn_jit::gpu::{Elem, Item, Variable};
+use burn_jit::{
+    cube_inline,
+    gpu::{Elem, Item, Variable},
+};
 
 macro_rules! cast_test {
     ($name:ident, $module:ident, $from:expr, $to:expr) => {
@@ -15,7 +17,7 @@ macro_rules! cast_test {
 
             assert_eq!(
                 format!("{:?}", scope.operations),
-                gpu_macro_ref_cast($from, $to)
+                inline_macro_ref_cast($from, $to)
             );
         }
     };
@@ -32,7 +34,7 @@ macro_rules! cast_test {
 
             assert_eq!(
                 format!("{:?}", scope.operations),
-                gpu_macro_ref_identity($ty)
+                inline_macro_ref_identity($ty)
             );
         }
     };
@@ -246,7 +248,7 @@ cast_test!(
     Item::Scalar(Elem::Bool)
 );
 
-fn gpu_macro_ref_cast(from_item: Item, to_item: Item) -> String {
+fn inline_macro_ref_cast(from_item: Item, to_item: Item) -> String {
     let mut context = CubeContext::root();
     let x = context.create_local(from_item);
 
@@ -257,25 +259,25 @@ fn gpu_macro_ref_cast(from_item: Item, to_item: Item) -> String {
     let z = scope.create_local(to_item);
 
     match from_item.elem() {
-        Elem::Float(_) => gpu!(scope, y = x + 2f32),
-        Elem::Int(_) => gpu!(scope, y = x + 2i32),
-        Elem::UInt => gpu!(scope, y = x + 2u32),
-        Elem::Bool => gpu!(scope, y = x && false),
+        Elem::Float(_) => cube_inline!(scope, y = x + 2f32),
+        Elem::Int(_) => cube_inline!(scope, y = x + 2i32),
+        Elem::UInt => cube_inline!(scope, y = x + 2u32),
+        Elem::Bool => cube_inline!(scope, y = x && false),
     }
 
-    gpu!(scope, y_casted = cast(y));
+    cube_inline!(scope, y_casted = cast(y));
 
     match to_item.elem() {
-        Elem::Float(_) => gpu!(scope, z = y_casted + 34f32),
-        Elem::Int(_) => gpu!(scope, z = y_casted + 34i32),
-        Elem::UInt => gpu!(scope, z = y_casted + 34u32),
-        Elem::Bool => gpu!(scope, z = y_casted || true),
+        Elem::Float(_) => cube_inline!(scope, z = y_casted + 34f32),
+        Elem::Int(_) => cube_inline!(scope, z = y_casted + 34i32),
+        Elem::UInt => cube_inline!(scope, z = y_casted + 34u32),
+        Elem::Bool => cube_inline!(scope, z = y_casted || true),
     }
 
     format!("{:?}", scope.operations)
 }
 
-fn gpu_macro_ref_identity(item: Item) -> String {
+fn inline_macro_ref_identity(item: Item) -> String {
     // When staying with the same type variables are automatically reused in cube
     let mut context = CubeContext::root();
     let x = context.create_local(item);
@@ -285,19 +287,19 @@ fn gpu_macro_ref_identity(item: Item) -> String {
     let y = scope.create_local(item);
 
     match item.elem() {
-        Elem::Float(_) => gpu!(scope, y = x + 2f32),
-        Elem::Int(_) => gpu!(scope, y = x + 2i32),
-        Elem::UInt => gpu!(scope, y = x + 2u32),
-        Elem::Bool => gpu!(scope, y = x && false),
+        Elem::Float(_) => cube_inline!(scope, y = x + 2f32),
+        Elem::Int(_) => cube_inline!(scope, y = x + 2i32),
+        Elem::UInt => cube_inline!(scope, y = x + 2u32),
+        Elem::Bool => cube_inline!(scope, y = x && false),
     }
 
-    gpu!(scope, x = cast(y));
+    cube_inline!(scope, x = cast(y));
 
     match item.elem() {
-        Elem::Float(_) => gpu!(scope, y = x + 34f32),
-        Elem::Int(_) => gpu!(scope, y = x + 34i32),
-        Elem::UInt => gpu!(scope, y = x + 34u32),
-        Elem::Bool => gpu!(scope, y = x || true),
+        Elem::Float(_) => cube_inline!(scope, y = x + 34f32),
+        Elem::Int(_) => cube_inline!(scope, y = x + 34i32),
+        Elem::UInt => cube_inline!(scope, y = x + 34u32),
+        Elem::Bool => cube_inline!(scope, y = x || true),
     }
 
     format!("{:?}", scope.operations)

diff --git a/crates/burn-cube/tests/cast_kind.rs b/crates/burn-cube/tests/cast_kind.rs
@@ -1,7 +1,8 @@
 use burn_cube::{cube, CubeContext, Float, Int, F32, F64, I32, I64};
-use burn_jit::gpu;
-// use burn_jit::gpu::FloatKind;
-use burn_jit::gpu::{Elem, Item};
+use burn_jit::{
+    cube_inline,
+    gpu::{Elem, Item},
+};
 
 #[cube]
 pub fn cast_float_kind<F1: Float, F2: Float>(input: F1) {
@@ -28,7 +29,7 @@ fn cube_cast_float_kind_test() {
     cast_float_kind::expand::<F64, F32>(&mut context, input);
     let scope = context.into_scope();
 
-    assert_eq!(format!("{:?}", scope.operations), gpu_macro_ref_float());
+    assert_eq!(format!("{:?}", scope.operations), inline_macro_ref_float());
 }
 
 #[test]
@@ -41,10 +42,10 @@ fn cube_cast_int_kind_test() {
     cast_int_kind::expand::<I32, I64>(&mut context, input);
     let scope = context.into_scope();
 
-    assert_eq!(format!("{:?}", scope.operations), gpu_macro_ref_int());
+    assert_eq!(format!("{:?}", scope.operations), inline_macro_ref_int());
 }
 
-fn gpu_macro_ref_float() -> String {
+fn inline_macro_ref_float() -> String {
     let mut context = CubeContext::root();
     let float_64 = Item::Scalar(Elem::Float(F64::into_kind()));
     let float_32 = Item::Scalar(Elem::Float(F32::into_kind()));
@@ -55,14 +56,14 @@ fn gpu_macro_ref_float() -> String {
     let y = scope.create_local(float_32);
     let z = scope.create_local(float_32);
 
-    gpu!(scope, x = input + 5.9f32 as f64);
-    gpu!(scope, y = cast(x));
-    gpu!(scope, z = y + 2.3f32);
+    cube_inline!(scope, x = input + 5.9f32 as f64);
+    cube_inline!(scope, y = cast(x));
+    cube_inline!(scope, z = y + 2.3f32);
 
     format!("{:?}", scope.operations)
 }
 
-fn gpu_macro_ref_int() -> String {
+fn inline_macro_ref_int() -> String {
     let mut context = CubeContext::root();
     let int_32 = Item::Scalar(Elem::Int(I32::into_kind()));
     let int_64 = Item::Scalar(Elem::Int(I64::into_kind()));
@@ -73,9 +74,9 @@ fn gpu_macro_ref_int() -> String {
     let y = scope.create_local(int_64);
     let z = scope.create_local(int_64);
 
-    gpu!(scope, x = input + 5i32);
-    gpu!(scope, y = cast(x));
-    gpu!(scope, z = y + 2i64);
+    cube_inline!(scope, x = input + 5i32);
+    cube_inline!(scope, y = cast(x));
+    cube_inline!(scope, z = y + 2i64);
 
     format!("{:?}", scope.operations)
 }
diff --git a/crates/burn-cube/tests/for_loop.rs b/crates/burn-cube/tests/for_loop.rs
@@ -1,6 +1,8 @@
 use burn_cube::{cube, range, range_expand, Array, CubeContext, Float, UInt, F32};
-use burn_jit::gpu;
-use burn_jit::gpu::{Elem, Item, Variable};
+use burn_jit::{
+    cube_inline,
+    gpu::{Elem, Item, Variable},
+};
 
 type ElemType = F32;
 
@@ -26,7 +28,7 @@ fn test_for_loop_with_unroll() {
     for_loop::expand::<ElemType>(&mut context, lhs, rhs, end, unroll);
     let scope = context.into_scope();
 
-    assert_eq!(format!("{:?}", scope.operations), gpu_macro_ref(unroll));
+    assert_eq!(format!("{:?}", scope.operations), inline_macro_ref(unroll));
 }
 
 #[test]
@@ -41,10 +43,10 @@ fn test_for_loop_no_unroll() {
     for_loop::expand::<ElemType>(&mut context, lhs, rhs, end, unroll);
     let scope = context.into_scope();
 
-    assert_eq!(format!("{:?}", scope.operations), gpu_macro_ref(unroll));
+    assert_eq!(format!("{:?}", scope.operations), inline_macro_ref(unroll));
 }
 
-fn gpu_macro_ref(unroll: bool) -> String {
+fn inline_macro_ref(unroll: bool) -> String {
     let mut context = CubeContext::root();
     let item = Item::Scalar(Elem::Float(ElemType::into_kind()));
 
@@ -58,15 +60,15 @@ fn gpu_macro_ref(unroll: bool) -> String {
     // Kernel
     let tmp1 = scope.create_local(item);
     let tmp2 = scope.create_local(item);
-    gpu!(scope, tmp1 = rhs * rhs);
-    gpu!(scope, tmp2 = tmp1 + rhs);
+    cube_inline!(scope, tmp1 = rhs * rhs);
+    cube_inline!(scope, tmp2 = tmp1 + rhs);
 
-    gpu!(
+    cube_inline!(
         &mut scope,
         range(0u32, end, unroll).for_each(|i, scope| {
-            gpu!(scope, rhs = lhs[i]);
-            gpu!(scope, tmp1 = tmp2 + rhs);
-            gpu!(scope, lhs[i] = tmp1);
+            cube_inline!(scope, rhs = lhs[i]);
+            cube_inline!(scope, tmp1 = tmp2 + rhs);
+            cube_inline!(scope, lhs[i] = tmp1);
         })
     );
 

diff --git a/crates/burn-cube/tests/if.rs b/crates/burn-cube/tests/if.rs
@@ -1,6 +1,5 @@
 use burn_cube::{cube, if_expand, CubeContext, Float, F32};
-use burn_jit::gpu;
-use burn_jit::gpu::{Elem, Item, Variable};
+use burn_jit::{cube_inline, gpu::{Elem, Item, Variable}};
 
 type ElemType = F32;
 
@@ -20,10 +19,10 @@ fn cube_if_test() {
     if_greater::expand::<ElemType>(&mut context, lhs);
     let scope = context.into_scope();
 
-    assert_eq!(format!("{:?}", scope.operations), gpu_macro_ref());
+    assert_eq!(format!("{:?}", scope.operations), inline_macro_ref());
 }
 
-fn gpu_macro_ref() -> String {
+fn inline_macro_ref() -> String {
     let mut context = CubeContext::root();
     let item = Item::Scalar(Elem::Float(ElemType::into_kind()));
     let lhs = context.create_local(item);
@@ -33,9 +32,9 @@ fn gpu_macro_ref() -> String {
     let lhs: Variable = lhs.into();
     let y = scope.create_local(item);
 
-    gpu!(scope, cond = lhs > 0f32);
-    gpu!(&mut scope, if(cond).then(|scope| {
-        gpu!(scope, y = lhs + 4.0f32);
+    cube_inline!(scope, cond = lhs > 0f32);
+    cube_inline!(&mut scope, if(cond).then(|scope| {
+        cube_inline!(scope, y = lhs + 4.0f32);
     }));
 
     format!("{:?}", scope.operations)

diff --git a/crates/burn-cube/tests/if_else.rs b/crates/burn-cube/tests/if_else.rs
@@ -1,6 +1,8 @@
 use burn_cube::{cube, if_else_expand, CubeContext, Float, F32};
-use burn_jit::gpu;
-use burn_jit::gpu::{Elem, Item, Variable};
+use burn_jit::{
+    cube_inline,
+    gpu::{Elem, Item, Variable},
+};
 
 type ElemType = F32;
 
@@ -22,10 +24,10 @@ fn cube_if_else_test() {
     if_else::expand::<ElemType>(&mut context, lhs);
     let scope = context.into_scope();
 
-    assert_eq!(format!("{:?}", scope.operations), gpu_macro_ref());
+    assert_eq!(format!("{:?}", scope.operations), inline_macro_ref());
 }
 
-fn gpu_macro_ref() -> String {
+fn inline_macro_ref() -> String {
     let mut context = CubeContext::root();
     let item = Item::Scalar(Elem::Float(ElemType::into_kind()));
     let lhs = context.create_local(item);
@@ -35,11 +37,11 @@ fn gpu_macro_ref() -> String {
     let lhs: Variable = lhs.into();
     let y = scope.create_local(item);
 
-    gpu!(scope, cond = lhs < 0f32);
-    gpu!(&mut scope, if(cond).then(|scope| {
-        gpu!(scope, y = lhs + 4.0f32);
+    cube_inline!(scope, cond = lhs < 0f32);
+    cube_inline!(&mut scope, if(cond).then(|scope| {
+        cube_inline!(scope, y = lhs + 4.0f32);
     }).else(|scope|{
-        gpu!(scope, y = lhs - 5.0f32);
+        cube_inline!(scope, y = lhs - 5.0f32);
     }));
 
     format!("{:?}", scope.operations)

diff --git a/crates/burn-cube/tests/literal.rs b/crates/burn-cube/tests/literal.rs
@@ -1,6 +1,5 @@
 use burn_cube::{cube, CubeContext, Float, F32};
-use burn_jit::gpu;
-use burn_jit::gpu::{Elem, Item};
+use burn_jit::{cube_inline, gpu::{Elem, Item}};
 
 type ElemType = F32;
 
@@ -18,17 +17,17 @@ fn cube_literal_test() {
     literal::expand::<ElemType>(&mut context, lhs);
     let scope = context.into_scope();
 
-    assert_eq!(format!("{:?}", scope.operations), gpu_macro_ref());
+    assert_eq!(format!("{:?}", scope.operations), inline_macro_ref());
 }
 
-fn gpu_macro_ref() -> String {
+fn inline_macro_ref() -> String {
     let mut context = CubeContext::root();
     let item = Item::Scalar(Elem::Float(ElemType::into_kind()));
     let lhs = context.create_local(item);
 
     let mut scope = context.into_scope();
     let out = scope.create_local(item);
-    gpu!(scope, out = lhs + 5.9f32);
+    cube_inline!(scope, out = lhs + 5.9f32);
 
     format!("{:?}", scope.operations)
 }
diff --git a/crates/burn-cube/tests/loop.rs b/crates/burn-cube/tests/loop.rs
@@ -1,7 +1,7 @@
 use burn_cube::{
     break_expand, cube, if_expand, loop_expand, while_loop_expand, CubeContext, Int, I32,
 };
-use burn_jit::gpu;
+use burn_jit::cube_inline;
 use burn_jit::gpu::Branch;
 use burn_jit::gpu::{Elem, Item, Variable};
 
@@ -33,7 +33,7 @@ fn cube_while_test() {
     while_not::expand::<ElemType>(&mut context, lhs);
     let scope = context.into_scope();
 
-    assert_eq!(format!("{:?}", scope.operations), gpu_macro_ref());
+    assert_eq!(format!("{:?}", scope.operations), inline_macro_ref());
 }
 
 #[test]
@@ -45,10 +45,10 @@ fn cube_loop_break_test() {
     manual_loop_break::expand::<ElemType>(&mut context, lhs);
     let scope = context.into_scope();
 
-    assert_eq!(format!("{:?}", scope.operations), gpu_macro_ref());
+    assert_eq!(format!("{:?}", scope.operations), inline_macro_ref());
 }
 
-fn gpu_macro_ref() -> String {
+fn inline_macro_ref() -> String {
     let mut context = CubeContext::root();
     let item = Item::Scalar(Elem::Int(ElemType::into_kind()));
     let lhs = context.create_local(item);
@@ -58,15 +58,15 @@ fn gpu_macro_ref() -> String {
     let lhs: Variable = lhs.into();
     let rhs = scope.create_local(item);
 
-    gpu!(
+    cube_inline!(
         &mut scope,
         loop(|scope| {
-            gpu!(scope, cond = lhs != 0);
-            gpu!(scope, if(cond).then(|scope|{
+            cube_inline!(scope, cond = lhs != 0);
+            cube_inline!(scope, if(cond).then(|scope|{
                 scope.register(Branch::Break);
             }));
 
-            gpu!(scope, rhs = lhs - 1i32);
+            cube_inline!(scope, rhs = lhs - 1i32);
         })
     );