Static evaluation on warp constants

[steinraf]

Are there any ways to perform 'compile time' computations or branching, like for example expressions involving wp.constant?

I am working on an example where I would like to do some computations in which, depending on the dimensionality of the problem, I need to run different code.

Reduced to a simple example that looks something like this

import warp as wp
wp.init()

dim = wp.constant(2)
# dim = wp.constant(3)


@wp.kernel
def test(x: wp.array(dtype=float, ndim=dim)):

    if dim == 2:
        wp.atomic_add(x, 0, 0, 0.0)
    if dim == 3:
        wp.atomic_add(x, 0, 0, 0, 0.0)


wp.launch(test, dim=1, inputs=[wp.zeros(shape=(1,) * dim)])

Here a runtime error is raised at the line wp.atomic_add(x, 0, 0, 0, 0.0) with Num indices > num dimensions for atomic array operation, received 3, but array only has 2 is raised since it tries to call the atomic add on a 2d array with 3 indices, even though it could potentially already be deduced during parsing that this branch would never be executed.

I know there for example exist mechanisms in taichi to do things like that, through ti.static or ti.grouped (in this particular case of atomic adds the behaviour is enabled by being able to index using vectors), so I wondered if there was a similar thing in warp.

Thanks!