Support union type annotations on kernel and function parameters #2293

Description

I'd like to annotate a parameter of a @wp.array or @wp.func with a union of types, so it accepts exactly those types and nothing else.

A common use case is a kernel that should run on a whole array or on a subset of it through wp.indexedarray:

import warp as wp

@wp.kernel
def apply_gravity(
    velocities: wp.array[wp.vec3] | wp.indexedarray[wp.vec3],
    dt: float,
):
    i = wp.tid()
    velocities[i] = velocities[i] + wp.vec3(0.0, -9.81, 0.0) * dt

velocities = wp.zeros(1000, dtype=wp.vec3)

# Update every particle
wp.launch(apply_gravity, dim=velocities.shape[0], inputs=[velocities, 0.01])

# Update only the particles that aren't pinned, without copying the data
active = wp.array([3, 17, 42, 256], dtype=int)
active_velocities = wp.indexedarray(velocities, [active])
wp.launch(apply_gravity, dim=active.shape[0], inputs=[active_velocities, 0.01])

Expected Behavior

• A union parameter makes the kernel generic, just like Any. The kernel is still only specialized at runtime for the current input types.
• However if a input type matches no union member, the launch throws a clear error, independent of the code generation passing:

wp.launch(apply_gravity, dim=10, inputs=[wp.zeros(10, dtype=wp.float32), 0.01])
# TypeError: Kernel apply_gravity argument 'velocities' type mismatch:
#   expected array(ndim=1, dtype=vec3f) | indexedarray(ndim=1, dtype=vec3f),
#   got array(ndim=1, dtype=float32)

• The same should hold for precompilling with wp.overload, requesting a type not in the union throws an error.
• Unions should be compatible with the preexisting sets of types (wp.Int, wp.Float, wp.Scalar): i.e. velocities: wp.array[wp.Scalar] | wp.indexedarray[wp.Scale]

Implementation

I threw together a quick prototype, and was able to get desired behavior with minimal changes to types.py and codegen.py

Context

Currently any generic parameter must be typed with Any. This can become quite painful when working on larger projects where typing improves readability. Additionally this prohibits proper safeguards via typing, you rely on the kernel compilation to fail with unintended types, which is not always the case.

Alternatives or Workarounds

Use TYPE_CHECKING flag to get similar behavior, still leads to worse readability / more boilerplate.

from typing import TYPE_CHECKING, TypeAlias

if TYPE_CHECKING:
    Vec3Velocities: TypeAlias = wp.array[wp.vec3] | wp.indexedarray[wp.vec3]
else:
    Vec3Velocities: Any


@wp.kernel
def apply_gravity(velocities: Vec3Velocities, dt: float):
    ...