CUDA Type Signature for RNG States


I’m trying to explicitly specify a type signature for a Numba CUDA kernel that takes RNG states generated as per Random Number Generation — Numba 0.50.1 documentation, but I’m not sure what type signature to use for the RNG states. Please can anyone point me in the right direction?

More specifically, I have this:

@cuda.jit  # ("void(f4[:,:],f4[:,:],f4[:,:,:],f4[:],f4[:,:],?,...)")
def __ck_sample_correspondences(depth_map, prev_depth_map, r_flow, centre, proj_inv, rng_states, correspondences) -> None:

The states themselves were created via:

self.__rng_states = create_xoroshiro128p_states(self.__max_hypotheses, seed=12345)

I tried looking at the states in the debugger, but their dtype is “dtype([(‘s0’, ‘<u8’), (‘s1’, ‘<u8’)], align=True)”, and I’m not quite sure what to do with that.

I’m aware that I can simply avoid specifying the signature at all, but then the first invocation of the kernel is slow because of the JIT compilation. My best alternative currently is to make a dummy invocation of the kernel to warm it up, but I’d prefer to specify the signature directly if that’s possible. Any suggestions please?