From a2a2add4c32ecb4d82d5e367c50ec345646c95ae Mon Sep 17 00:00:00 2001
From: Genna Wingert <wingertge@gmail.com>
Date: Sun, 15 Sep 2024 20:45:18 +0200
Subject: [PATCH] Revert variable reuse optimization in loops (#121)

---
 crates/cubecl-core/src/frontend/branch.rs       |  3 ++-
 crates/cubecl-core/tests/frontend/loop.rs       | 10 ++++++----
 crates/cubecl-macros/src/expression.rs          |  5 ++++-
 crates/cubecl-macros/src/generate/expression.rs |  4 ++--
 crates/cubecl-macros/src/parse/branch.rs        |  4 ++--
 5 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/crates/cubecl-core/src/frontend/branch.rs b/crates/cubecl-core/src/frontend/branch.rs
index 626aee84..4e4b1661 100644
--- a/crates/cubecl-core/src/frontend/branch.rs
+++ b/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);
diff --git a/crates/cubecl-core/tests/frontend/loop.rs b/crates/cubecl-core/tests/frontend/loop.rs
index 3eeb85c9..ac76ae25 100644
--- a/crates/cubecl-core/tests/frontend/loop.rs
+++ b/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);
             })
         );
 
diff --git a/crates/cubecl-macros/src/expression.rs b/crates/cubecl-macros/src/expression.rs
index 023cbd38..bc267de1 100644
--- a/crates/cubecl-macros/src/expression.rs
+++ b/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,
diff --git a/crates/cubecl-macros/src/generate/expression.rs b/crates/cubecl-macros/src/generate/expression.rs
index a849c504..19c3c030 100644
--- a/crates/cubecl-macros/src/generate/expression.rs
+++ b/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);]
             }
diff --git a/crates/cubecl-macros/src/parse/branch.rs b/crates/cubecl-macros/src/parse/branch.rs
index a8524cb3..4c6ecae8 100644
--- a/crates/cubecl-macros/src/parse/branch.rs
+++ b/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> {