in-GPU-memory interoperability with Machine Learning frameworks?

[szabolcsdombi]

Discussed in https://github.com/moderngl/moderngl/discussions/682

^{Originally posted by hrbigelow March 30, 2024}
I am interested in using both moderngl and jax in an interactive application - the moderngl is needed to support displaying images and perhaps also interacting with a windowing system to allow the user to annotate them in various ways.

But in addition to this I would like to use Jax to do other processing on the image data.

I know that I can take the route of:

1. read GPU memory Texture data into a system memory Python bytearray
2. instantiate a Jax tensor from the system memory bytearray
3. move the Jax tensor onto GPU memory
4. process the tensor on GPU
5. extract tensor data into a system memory bytearray
6. load the system memory bytearray data into the Texture

But, I wondered if there's any way to achieve this flow without involving system memory?

Thanks in advance!

[hrbigelow]

Thank you @szabolcsdombi.

By the way, @jakevdp responded to this question on the Jax end here:

The primary method for zero-copy transfer of GPU buffers into JAX is via the dlpack data structure, using the jax.from_dlpack function.

If dlpack is not an option, JAX also has support for __cuda_array_interface__.

Perhaps moderngl could be enhanced to support a Buffer.from_dlpack and Buffer.to_dlpack functionality, what do you think?

[szabolcsdombi]

Hey, I look into that, i guess we support it through the external buffer and external texture options. But I first have to verify

[hrbigelow]

Thank you @szabolcsdombi, this would be very cool if it can work.

[hrbigelow]

Just one additional comment from @jakevdp:

I should note that the Python Array API specification landed on dlpack as the preferred option for this (see discussion here: https://data-apis.org/array-api/latest/design_topics/data_interchange.html) so if you're implementing something, I'd lean toward that rather than the cuda array interface.

[hrbigelow]

Hi @szabolcsdombi, just curious if you had a chance to look at this. I'd also be willing to look into it, but would be interested in your intuitions.