[BUG] Typed Tuple Returns in Warp Cause Error

yzx9 commented 2 weeks ago

Bug Description

Warp supports multi-value returns, but it seems that they cannot be properly typed.

Reproduction Steps

import warp as wp

@wp.kernel
def kernel():
    i, j = wp.tid()
    mod(i, j + 1)

@wp.func
def mod(a: int, b: int) -> tuple[int, int]:
    return a // b, a % b

wp.launch(kernel, dim=(100, 10))

Actual Outcome

The code fails with the following error message:

Warp NVRTC compilation error 6: NVRTC_ERROR_COMPILATION (/builds/omniverse/warp/warp/native/warp.cu:2577)
default_program(96): error: namespace "wp" has no member "tuple"
          wp::tuple var_4;
              ^

default_program(104): error: too few arguments in function call
          var_4 = mod(var_0, var_3);
                                  ^

default_program(121): error: namespace "wp" has no member "tuple"
          wp::tuple var_4;
              ^

default_program(128): error: namespace "wp" has no member "tuple"
          wp::tuple adj_4 = {};
              ^

default_program(136): error: too few arguments in function call
          var_4 = mod(var_0, var_3);
                                  ^

default_program(139): error: too few arguments in function call
          adj_mod(var_0, var_3, adj_0, adj_3, adj_4);
                                                   ^

6 errors detected in the compilation of "default_program".
Module __main__ b019979 load on device 'cuda:0' took 70.02 ms  (error)
Traceback (most recent call last):
  File "/home/yzx9/git/morphtesser/morphtesser/test.py", line 15, in <module>
    wp.launch(kernel, dim=(100, 10))
  File "/home/yzx9/.cache/pypoetry/virtualenvs/morphtesser-EpsGKO6L-py3.12/lib/python3.12/site-packages/warp/context.py", line 4693, i
n launch
    if not module.load(device):
           ^^^^^^^^^^^^^^^^^^^
  File "/home/yzx9/.cache/pypoetry/virtualenvs/morphtesser-EpsGKO6L-py3.12/lib/python3.12/site-packages/warp/context.py", line 1828, i
n load
    raise (e)
  File "/home/yzx9/.cache/pypoetry/virtualenvs/morphtesser-EpsGKO6L-py3.12/lib/python3.12/site-packages/warp/context.py", line 1816, i
n load
    warp.build.build_cuda(
  File "/home/yzx9/.cache/pypoetry/virtualenvs/morphtesser-EpsGKO6L-py3.12/lib/python3.12/site-packages/warp/build.py", line 30, in bu
ild_cuda
    raise Exception(f"CUDA kernel build failed with error code {err}")
Exception: CUDA kernel build failed with error code 6

Change tuple to typing.Tuple receive same results. However, if the typing is removed, the code works as expected.

Excepted

The code should work correctly with typed multi-value returns.

System Information

Warp 1.3.1 initialized: CUDA Toolkit 12.5, Driver 12.5 Devices: "cpu" : "CPU" "cuda:0" : "NVIDIA GeForce RTX 4090" (24 GiB, sm_89, mempool enabled) Kernel cache: /home/yzx9/.cache/warp/1.3.1

christophercrouzet commented 2 weeks ago

Hi @yzx9, thanks for reporting this issue, we'll look into it!

christophercrouzet commented 1 week ago

Hi @yzx9, this has now been addressed with commit e539355 which has been merged to main and will be included in a future release. Please let us know if you find any other issue!

NVIDIA / warp