Reusing for-loop variable name in while loop gives wrong result

[shi-eric]

Bug Description

Warp gives an incorrect result when a variable name used in a for loop is reused later as the induction variable of a while loop.

In the example below, both kernels should compute the same value. The only difference is that same_name() reuses i in the while loop, while different_name() uses j.

import warp as wp

wp.init()


@wp.kernel
def same_name(out: wp.array[wp.int32]):
    acc = int(0)

    for i in range(2):
        acc = acc + i

    i = int(0)
    while i < 3:
        acc = acc + 10
        i = i + 1

    out[0] = acc


@wp.kernel
def different_name(out: wp.array[wp.int32]):
    acc = int(0)

    for i in range(2):
        acc = acc + i

    j = int(0)
    while j < 3:
        acc = acc + 10
        j = j + 1

    out[0] = acc


device = "cpu"

out = wp.zeros(1, dtype=wp.int32, device=device)
wp.launch(same_name, dim=1, inputs=[out], device=device)
print("same_name", int(out.numpy()[0]))

out = wp.zeros(1, dtype=wp.int32, device=device)
wp.launch(different_name, dim=1, inputs=[out], device=device)
print("different_name", int(out.numpy()[0]))

Expected:

same_name 31
different_name 31

Observed locally on current main:

same_name 0
different_name 31

Expected calculation:

for loop:   acc = 0 + 0 + 1 = 1
while loop: acc = 1 + 10 + 10 + 10 = 31

System Information

Confirmed locally on:

• Warp: 1.15.0.dev0
• Python: 3.12.13
• OS/platform: Linux-6.8.0-1051-gke-x86_64-with-glibc2.35
• Device: CPU, x86_64

[snico432]

I’m interested in taking this issue on.

[shi-eric]

Hey @snico432, I assigned it to you. Thanks!

[snico432]

Hey @shi-eric, I have a local fix for this, but I wanted to confirm the intended semantics before opening a PR.

Should variable scoping within a Warp kernel follow Python's loop scoping behaviour?

@wp.kernel
def test_static_for_iter_reused(result: wp.array[wp.int32]):
    acc = int(0)

    for i in range(2):
        acc = acc + i
    
    # In Python, `i` is still in scope here and has the final value `1`.
    # Should Warp allow this, or should users re-declare `i` as a dynamic
    # variable before using it in a dynamic loop, e.g. `i = int(1)`?
    while i > 0:
        acc = acc + 10
        i = i - 1

    result[0] = acc

[shi-eric]

@snico432 Sorry, I think this issue was more complicated than I realized as there are a class of expressions like this that don't work in kernels, and I'm not sure if we've considered it properly about what we want to support. I'll let you know once I have a better clue of the landscape. I'm inclined to say that we don't want to support how i is used in your example, though. @c0d1f1ed might weigh in with a more definitive answer.