Revert variable reuse optimization in loops (#121)

crates/cubecl-core/src/frontend/branch.rs

@@ -415,7 +415,8 @@ pub fn return_expand(context: &mut CubeContext) {
     context.register(Branch::Return);
 }
 
-pub fn loop_expand(context: &mut CubeContext, block: impl FnOnce(&mut CubeContext)) {
+// Don't make this `FnOnce`, it must be executable multiple times
+pub fn loop_expand(context: &mut CubeContext, mut block: impl FnMut(&mut CubeContext)) {
     let mut inside_loop = context.child();
 
     block(&mut inside_loop);

crates/cubecl-core/tests/frontend/loop.rs

@@ -87,6 +87,7 @@ mod tests {
 
         let mut scope = context.into_scope();
         let cond = scope.create_local(Item::new(Elem::Bool));
+        let y = scope.create_local(item);
         let lhs: Variable = lhs.into();
 
         cpa!(
@@ -97,8 +98,8 @@ mod tests {
                 cpa!(scope, if(cond).then(|scope|{
                         scope.register(Branch::Break)
                 }));
-
-                cpa!(scope, lhs = lhs % 1i32);
+                // Must not mutate `lhs` because it is used in every iteration
+                cpa!(scope, y = lhs % 1i32);
             })
         );
 
@@ -112,6 +113,7 @@ mod tests {
 
         let mut scope = context.into_scope();
         let cond = scope.create_local(Item::new(Elem::Bool));
+        let y = scope.create_local(item);
         let lhs: Variable = lhs.into();
 
         cpa!(
@@ -124,8 +126,8 @@ mod tests {
                         false => scope.register(Branch::Break)
                     }
                 }));
-
-                cpa!(scope, lhs = lhs % 1i32);
+                // Must not mutate `lhs` because it is used in every iteration
+                cpa!(scope, y = lhs % 1i32);
             })
         );
 

crates/cubecl-macros/src/expression.rs

@@ -82,7 +82,10 @@ pub enum Expression {
         block: Block,
         scope: Scope,
     },
-    Loop(Block),
+    Loop {
+        block: Block,
+        scope: Scope,
+    },
     If {
         condition: Box<Expression>,
         then_block: Block,

crates/cubecl-macros/src/generate/expression.rs

@@ -261,9 +261,9 @@ impl Expression {
                     }
                 }
             }
-            Expression::Loop(block) => {
+            Expression::Loop { block, scope } => {
                 let loop_ty = frontend_type("branch");
-                let block = block.to_tokens(context);
+                let block = context.in_fn_mut(scope, |ctx| block.to_tokens(ctx));
 
                 quote![#loop_ty::loop_expand(context, |context| #block);]
             }

crates/cubecl-macros/src/parse/branch.rs

@@ -67,8 +67,8 @@ fn expand_for_in_loop(
 }
 
 pub fn expand_loop(loop_expr: ExprLoop, context: &mut Context) -> syn::Result<Expression> {
-    let (block, _) = context.in_scope(|ctx| Block::from_block(loop_expr.body, ctx))?;
-    Ok(Expression::Loop(block))
+    let (block, scope) = context.in_scope(|ctx| Block::from_block(loop_expr.body, ctx))?;
+    Ok(Expression::Loop { block, scope })
 }
 
 pub fn expand_if(if_expr: ExprIf, context: &mut Context) -> syn::Result<Expression> {