NVIDIA / warp

A Python framework for high performance GPU simulation and graphics
https://nvidia.github.io/warp/
Other
1.75k stars 148 forks source link

Bool vectors causing errors. #177

Closed loliverhennigh closed 3 months ago

loliverhennigh commented 4 months ago

If you make a vector of bools then indexing seems to not work. The following is a minimal example to create the issue

import warp as wp

# Initialize Warp
wp.init()

# Does not work
bool_vec = wp.vec(3, dtype=wp.bool)
c = wp.constant(bool_vec([True, False, True]))

# Works
#uint_vec = wp.vec(3, dtype=wp.uint8)
#c = wp.constant(uint_vec([1, 0, 1]))

@wp.kernel
def warp_kernel(
    f0: wp.array(dtype=float),
    f1: wp.array(dtype=float),
):

    # get index
    x = wp.tid()
    x + 1 # avoid edges

    # sample needed points
    f_add = wp.float32(0.0)
    if c[0]:
        f_add += f0[x-1]
    if c[1]:
        f_add += f0[x]
    if c[2]:
        f_add += f0[x+1]

    # set value
    f1[x] = f_add

if __name__ == "__main__":

    # Make f0, f1
    n = 256
    f0 = wp.zeros((n), dtype=wp.float32, device="cuda:0")
    f1 = wp.zeros((n), dtype=wp.float32, device="cuda:0")

    # launch kernel
    wp.launch(
        warp_kernel,
        inputs=[
            f0,
            f1,
        ],
        dim=(n-2,),
    )

This causes the following error

warp.codegen.WarpCodegenError: Error while parsing function "warp_kernel" at /home/oliver/xlb_workspace/XLB/examples/backend_comparisons/bool_vec.py:26:
    if c[0]:
;Couldn't find function overload for 'extract' that matched inputs with types: [vector(length=3, dtype=<class 'warp.types.bool'>), int32]

Works fine if using uint8 though.

c0d1f1ed commented 4 months ago

Thanks for reporting this! I'll have a closer look.

shi-eric commented 4 months ago

Yes, I neglected to test that the wp.bool type worked with vectors and matrices in the original commit.

c0d1f1ed commented 3 months ago

Fixed in v0.13.0. Thanks Eric!