NVIDIA / warp

A Python framework for high performance GPU simulation and graphics

Home Page:https://nvidia.github.io/warp/

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Breaking in a loop inside a kernel

jc211 opened this issue · comments

Adding a break inside a for loop in the kernel raises an error.

Example:

import warp as wp
@wp.kernel
def foo(a: int):
    for i in range(10):
        wp.printf("a: %d, i: %d\n", a, i)
        if i == a:
            break

wp.init()
wp.launch(kernel=foo, dim=(1,), inputs=[3])
wp.synchronize_device()

Getting around this would probably require something like: while(i != a) instead of a for loop

Thanks for reporting this and providing a complete repro! It appears to be an issue with for loops eligible for unrolling (by default up to 16 iterations). Indeed a while loop could be used as a workaround since we currently never unroll them:

i = int(0)
while i < 10 and i != a:
    wp.printf("a: %d, i: %d\n", a, i)
    i += 1

Not very elegant, I know. I'll try to fix this properly by the next release.

Hi @jc211, this has been addressed in b57868f.

Thanks,
Miles