Structs in functions

[xelatihy]

So far it seems structs are supported as arguments to kennels but cannot be used in functions. It would be nice to be able to add support for structs in functions for both arguments and return values. This should be considered low priority. I put it here in case others find the same issue I did.

[mmacklin]

Thanks Fabio, yes that makes total sense. I will look into it.

[daedalus5]

Hi Fabio, the issue should now be fixed. The following works in 0.6.1

import numpy as np
import warp as wp

wp.init()

@wp.struct
class TestStruct:
    x: int
    a: wp.array(dtype=int)
    b: wp.array(dtype=int)

@wp.func
def foo(s: TestStruct, idx: int):
    s.a[idx] = s.x + idx

@wp.kernel
def test_kernel(s: TestStruct):
    tid = wp.tid()

    foo(s, tid)
    s.b[tid] = s.a[tid]

# create struct
ts = TestStruct()

# set members
ts.x = 1
ts.a = wp.array(
    np.zeros(3),
    dtype=int,
    requires_grad=True,
)
ts.b = wp.zeros(3, dtype=int, requires_grad=True)

wp.launch(test_kernel, dim=3, inputs=[ts])
print(ts.b.numpy())
# returns [1 2 3]