Launching Mojo kernels on a specific CUDA/HIP stream?

GPU Programming
1

Hi,

I am writing a set of Mojo bindings for a CUDA/HIP library that expects applications to launch their kernels on a given stream, because the library may launch dependent kernels on the same stream.

Is there a way to specify the CUDA/HIP stream on which to launch a Mojo kernel? If not, we can work around this, but it requires doing a device synchronize before/after the Mojo kernel launch, so it’s not performance-optimal.

Thanks,

Ben

2

This is not possible in Mojo today. I have created a proposal to add this feature here: Proposal: Allow importing CUDA/HIP stream handles in DeviceContext.