Skip to content

Make map_overlap compatible with cupy #413

@GenevieveBuckley

Description

@GenevieveBuckley

Marvin has done a bunch of work implementing a map_overlap function for dask-image: #237

Currently it supports numpy backed dask arrays, and does not support cupy GPU arrays. I believe it would be possible to extend this work to support cupy input arrays as well.

Suggested approach:
One way to approach this could be to:

  1. Take the the _map_single_coordinates_array_chunk function, and make a second version of it using cupy/cupyx functions replacing all the lines that call numpy directly.
  2. Then we could use the same dispatch mechanism as we do with the other modules (see the dask_image/dispatch folder) to register our new numpy and cupy versions.
  3. Add a test for the new functionality
  4. Update the coverage.rst table, adding a tick mark to indicate this function now has GPU support

All of the numpy specific functions I can see in the _map_single_coordinates_array_chunk function seem to have cupy/cupyx equivalents. If we need to pin to a certain version of cupy, or check before the code runs with packaging.version.parse that's ok.

Profiling results
A flow on possible complication might be how good or poor performance is in cases with larger than memory arrays. Since GPUs generally have more limited memory, is it more likely people will have larger than memory arrays under these circumstances.
That's not a blocker for this work, I think it would be valuable to have regardless. But it would be something interesting to look into.
In the CPU implementation, Marvin did some profiling of different use cases here: #237 (comment)

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions