Fix wp.static() capturing global variables instead of loop variables

When a global Python variable had the same name as a kernel for-loop
variable, wp.static() incorrectly used the global value. The fix tracks
loop variables during AST traversal and defers wp.static() evaluation
when the expression references a loop variable.

Fixes GH-1139

Signed-off-by: Eric Shi <[email protected]>

Author: shi-eric

CHANGELOG.md

@@ -15,6 +15,8 @@
 
 ### Fixed
 
+- Fix `wp.static()` incorrectly capturing global Python variables instead of loop variables when used inside for-loops
+  in kernels ([GH-1139](https://github.com/NVIDIA/warp/issues/1139)).
 - Fix `--llvm-path` build option to use existing LLVM installation when building `warp-clang` library instead of
   downloading from packman.
 - Fix excessive memory usage in CUDA graphs with multiple allocations/deallocations

warp/_src/codegen.py

@@ -3736,9 +3736,37 @@ def evaluate_static_expression(adj, node) -> tuple[Any, str]:
     # expression can be evaluated
     def replace_static_expressions(adj):
         class StaticExpressionReplacer(ast.NodeTransformer):
+            def __init__(self):
+                # Track loop variable names from enclosing for loops. This prevents
+                # wp.static() from capturing a global variable that shadows a loop variable.
+                # Uses a counter (not a set) to handle nested loops that reuse the same variable name.
+                self.loop_vars = {}
+
+            def visit_For(self, node):
+                # Track loop variable while visiting loop body (simple names only;
+                # tuple unpacking like `for x, y in ...` is rare in Warp kernels)
+                var_name = node.target.id if isinstance(node.target, ast.Name) else None
+                if var_name:
+                    self.loop_vars[var_name] = self.loop_vars.get(var_name, 0) + 1
+                result = self.generic_visit(node)
+                if var_name:
+                    self.loop_vars[var_name] -= 1
+                    if self.loop_vars[var_name] == 0:
+                        del self.loop_vars[var_name]
+                return result
+
             def visit_Call(self, node):
                 func, _ = adj.resolve_static_expression(node.func, eval_types=False)
                 if adj.is_static_expression(func):
+                    # If the static expression references an enclosing loop variable,
+                    # defer evaluation to codegen time when the loop constant is available
+                    expr_node = node.args[0] if node.args else (node.keywords[0].value if node.keywords else None)
+                    if expr_node:
+                        referenced = {n.id for n in ast.walk(expr_node) if isinstance(n, ast.Name)}
+                        if referenced & self.loop_vars.keys():
+                            adj.has_unresolved_static_expressions = True
+                            return self.generic_visit(node)
+
                     try:
                         # the static expression will execute as long as the static expression is valid and
                         # only depends on global or captured variables

warp/tests/test_static.py

@@ -659,6 +659,90 @@ def test_unresolved_static_expression(test, device):
         test.assertEqual(output2.numpy()[0], 2)
 
 
+# Global variables used to test wp.static() loop variable handling
+_global_test_idx = 999
+_global_test_j = 888
+
+
+@wp.kernel
+def static_loop_var_kernel(results: wp.array(dtype=int)):
+    """Kernel where wp.static() should capture the loop variable, not the global."""
+    for _global_test_idx in range(3):
+        results[_global_test_idx] = wp.static(_global_test_idx)
+
+
+@wp.kernel
+def static_loop_var_in_expr_kernel(results: wp.array(dtype=int)):
+    """Kernel where loop variable is used in an arithmetic expression."""
+    for _global_test_idx in range(3):
+        # Even in complex expressions, the loop variable should be used
+        results[_global_test_idx] = wp.static(_global_test_idx * 2 + 1)
+
+
+@wp.kernel
+def static_nested_loop_kernel(results: wp.array(dtype=int)):
+    """Kernel with nested loops - both loop variables should be protected."""
+    for _global_test_idx in range(2):
+        for _global_test_j in range(2):
+            idx = _global_test_idx * 2 + _global_test_j
+            results[idx] = wp.static(_global_test_idx * 10 + _global_test_j)
+
+
+@wp.kernel
+def static_nested_loop_same_var_kernel(results: wp.array(dtype=int)):
+    """Kernel with nested loops reusing the same variable name.
+
+    Tests counter-based loop variable tracking: when inner and outer loops
+    use the same variable name, the global should still not be captured.
+    Per Python semantics, after the inner loop the variable has the inner
+    loop's final value.
+    """
+    idx = 0
+    for _global_test_idx in range(2):
+        for _global_test_idx in range(3):  # intentional shadowing for test
+            pass
+        # Per Python semantics, _global_test_idx is now 2 (inner loop's final value)
+        # Key: we should NOT capture the global value (999)
+        results[idx] = wp.static(_global_test_idx)
+        idx += 1
+
+
+def test_static_loop_variable_not_shadowed_by_global(test, device):
+    """Test that wp.static() inside a for loop correctly captures the loop variable.
+
+    When a global Python variable exists with the same name as a kernel loop variable,
+    wp.static() should use the loop variable's compile-time constant value (0, 1, 2, ...),
+    not the unrelated global variable. This prevents confusing behavior where the
+    presence of a global variable silently changes the kernel's output.
+    """
+    with wp.ScopedDevice(device):
+        # Test 1: Simple loop variable
+        results = wp.zeros(3, dtype=int)
+        wp.launch(static_loop_var_kernel, dim=1, inputs=[results])
+        np.testing.assert_array_equal(results.numpy(), np.array([0, 1, 2]), err_msg="Simple loop variable test failed")
+
+        # Test 2: Loop variable in arithmetic expression
+        results2 = wp.zeros(3, dtype=int)
+        wp.launch(static_loop_var_in_expr_kernel, dim=1, inputs=[results2])
+        np.testing.assert_array_equal(
+            results2.numpy(), np.array([1, 3, 5]), err_msg="Loop variable in expression test failed"
+        )
+
+        # Test 3: Nested loops - both variables protected
+        results3 = wp.zeros(4, dtype=int)
+        wp.launch(static_nested_loop_kernel, dim=1, inputs=[results3])
+        np.testing.assert_array_equal(results3.numpy(), np.array([0, 1, 10, 11]), err_msg="Nested loop test failed")
+
+        # Test 4: Nested loops reusing the same variable name
+        # Tests counter-based tracking: global should not be captured even with shadowing
+        results4 = wp.zeros(2, dtype=int)
+        wp.launch(static_nested_loop_same_var_kernel, dim=1, inputs=[results4])
+        # Per Python semantics: inner loop shadows outer, final value (2) persists
+        np.testing.assert_array_equal(
+            results4.numpy(), np.array([2, 2]), err_msg="Nested loop with same variable test failed"
+        )
+
+
 devices = get_test_devices()
 
 
@@ -689,6 +773,12 @@ def test_static_python_call(self):
 add_function_test(TestStatic, "test_static_constant_hash", test_static_constant_hash, devices=None)
 add_function_test(TestStatic, "test_static_function_hash", test_static_function_hash, devices=None)
 add_function_test(TestStatic, "test_static_len_query", test_static_len_query, devices=None)
+add_function_test(
+    TestStatic,
+    "test_static_loop_variable_not_shadowed_by_global",
+    test_static_loop_variable_not_shadowed_by_global,
+    devices=devices,
+)
 
 
 if __name__ == "__main__":