Array API compatibility in scikit-image

Contributor & Development Discussion scikit-image
1

Hi,

The Array API standard aims to standardize a common subset of functionality of the majority of array libraries, such as NumPy, PyTorch, CuPy and JAX, with a view to remedy the fragmentation of the array computing ecosystem caused by the accumulated divergences among these—almost, but not completely dissimilar—array/tensor libraries.

Adopting the Array API standard in scikit-image gives users flexibility in choosing their software stack for array computing, and [enables performance improvements] (https://labs.quansight.org/blog/array-api-meta-blogposts) from using low-level implementations of hardware accelerated algorithms.

Much of scikit-image relies on computational kernels written in Cython. An immediate concern is whether adopting the Array API brings meaningful benefits—these handwritten kernels are and will remain being NumPy-only, while the majority of gains reported from Array API adoption are from using hardware accelerators.

First of all, to an extent of scikit-image using scipy.ndimage, the latter does benefit from GPU execution for CuPy arrays, today. This way, scikit-image functions which call scipy.ndimage functions and have their internals Array API compatible, use CUDA automatically for CuPy array inputs.

More generally, Array API compatibility gives a generic framework for, and implements the foundational infrastructure of, dispatching to specialized accelerator-enabled implementations (such as CUCIM and similar GPU libraries). Specific details of the dispatching can take multiple forms, there are several implementations in different libraries, and there are multiple further ideas floating around. What the Array API compatibility provides however, is a general and ecosystem-aligned framework for working through these (both fascinating and difficult) details.

Here is

  • a draft SKIP for adopting the Array API in scikit-image, as a PR, and in a rendered form.

  • a POC/WIP/demo level implementation which brings Array API compatibility to one function, skimage2.metrics.structured_similarity

Thoughts? Comments?
I know it’s a large topic with quite a bit of history, and it comes in the midst of the skimage2 transition. Either way, I’d be happy to go into details, either large-scale or smaller items, to see if this is something y’all would be interested in.

Cheers,
Evgeni

2

In the past I did a bit of work to add a dispatching system to scikit-image. The idea of it is that anyone can write a backend that implements some/all of the functions that they want to accelerate. Crucially you can write a backend with what ever technology that you want. This means it is possible to write a backend that is implemented in terms of the array API, even for functions that are implemented in cython in scikit-image itself.

A good place to get started if you want to look at the code is probably GitHub - betatim/skimage-cucim-backend: Cucim backend for scikit-image · GitHub. This is a backend that is implemented in terms of CuCIM. You need to install scikit-image from a fork (mentioned in the README).

I haven’t been able to work on this for a few months, but maybe what is there is already enough to be able to implement a backend in terms of the array API (or ask an AI agent to do so) and get a feeling for performance, ergonomics, etc

I’d also be interested in picking this work up again and “finishing” it. At least on the scikit-image side. But with not so much time and the push for scikit-image v2 I’ve lost track/contact with development.

I think a “plugin system” like this would be cool in general because it allows anyone to write a backend using what ever technology they want and accept input independent of its type (e.g. the CuCIM backend could choose to accept large Numpy arrays, convert to cupy, compute and return the result as Numpy again). This means having the plugin system would create room for experimentation.

3

Hi Tim,

I think a plugin system and Array API compatibility are complimentary, not mutually exclusive. I tend to think of Array API as a UI / framework, and plugins for dispatching/delegation as an implementation layer. Indeed, a plugin system can be very elaborate: from a trivial “torch CPU plugin” which just converts inputs to numpy, runs computations, and converts the result back, to your CUCIM plugin, to a mixed numpy/CUCIM plugin which decides on CPU or GPU depending on some heuristics on sizes, anything.
What an Array API brings to the table is a common framework. For example, I heard about somebody developing a CuPy / JAX bridge, so that CUDA kernels from CuPy would be understood by jax JIT. Than—if that works in the end, of course—a user would be able to use jit on a larger parts of their application, which uses scikit-image functions. Without Array API they would run into a usual problem that jax.numpy is almost but not completely a drop-in numpy replacement. With Arrray API, this point of friction is gone.

When writing the SKIP I tried to keep this framing and keep the SKIP itself mostly agnostic to specific backends / plugins. If that did not work, I’d appreciate pointers to how it can be improved, or rewrites/additions (there is no hard reason for the SKIP to not have multiple co-authors :-)).

Evgeni