NVIDIA / warp

A Python framework for high performance GPU simulation and graphics
https://nvidia.github.io/warp/
Other
4.3k stars 246 forks source link

[REQUEST] Support for c++ style <std::map> #301

Open AntoineRichard opened 3 months ago

AntoineRichard commented 3 months ago

Description

Hi there! First I wanted to say thanks for the fantastic work you are doing with Warp, and the support you're providing for it!

As for the request, I was wondering if you could implement c++ style std::map inside warp.

This would provide the ability to do things such as:

wp.map(key_type=int, mapped_type=wp.array(dtype=float))

Which would be equivalent to:

std::map<int, <std::vector<float>>>

This could be useful to minimize CPU/GPU exchanges.

Context

The reason I'm asking this, is working that I work on very large procedurally generated images. These images are built by chunk on the CPU. Then, some processing is done on the GPU side with WARP.

Currently, the pipeline flow is the following:

We have a base image.

  1. Shift that image by a given amount of chunks.
  2. Build missing image chunks.
  3. Rebuild image using the previously generated chunks.
  4. image is shared with warp (send to GPU)
  5. Process on the gpu
  6. Repeat (1-6)

If the feature was available, the pipeline would like that:

We have an initial map containing the chunks.

  1. Build missing chunks and send them to warp.
  2. Remove chunks (from warp side).
  3. Process on the GPU with proper chunk indexing.
  4. Repeat(1-4)

Having a std::map like system would allow avoiding rebuilding the image, and also reduce the amount of data sent to the GPU. I can provide more details if needed.

Thanks again for all the work and support,

Best,

Antoine