Fix helper function with while loop and return

Signed-off-by: Jay Gu <jagu@nvidia.com>

Author: haijieg

changelog.d/while-loop-return.md

@@ -0,0 +1,5 @@
+<!--- SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved. -->
+<!--- SPDX-License-Identifier: Apache-2.0 -->
+
+- Fix "potentially undefined variable `$retval`" error when a helper function
+  returns after a ``while`` loop that contains no early return.

src/cuda/tile/_passes/ast2hir.py

@@ -929,8 +929,8 @@ def _ast2hir(func_def: ast.FunctionDef | ast.Lambda, ctx: _Context) -> hir.Block
         elif isinstance(func_def, ast.FunctionDef):
             # To enable early returns in a helper function, wrap the body in a loop.
             # Thus, we can use "break" to implement the return statement.
-            ctx.store("$returning", False)
             with ctx.new_block() as body_block:
+                ctx.store("$returning", False)
                 _stmt_list(func_def.body, ctx)
                 if body_block.jump is None:
                     ctx.store("$retval", None)

test/test_helper_function.py

@@ -249,6 +249,20 @@ def test_early_return_inside_for_loop():
         ct.launch(torch.cuda.current_stream(), (1,), kernel, (n, out))
 
 
+def return_after_while_loop(n):
+    while n > 0:
+        n = n - 1
+    return n
+
+
+def test_return_after_while_loop():
+    n = torch.tensor([3], dtype=torch.int32, device="cuda")
+    out = torch.zeros_like(n)
+    kernel = early_return_inside_loop(return_after_while_loop)
+    ct.launch(torch.cuda.current_stream(), (1,), kernel, (n, out))
+    assert out.cpu().item() == 0
+
+
 def loops(n):
     a = 0
     for i in range(n):