Performance portable Python tasks

[elliottslaughter]

Some of my code uses Python tasks, because I launch a large number of operations sized for an individual GPU each, and I want to take separate streams of these and place them on different GPUs. (E.g., the first K * 1,000 operations go to the first GPU, the second K * 1,000 go to the second GPU, etc. up to hundreds of GPUs.)

Now the question is how do I make this performance portable.

Python tasks are not permitted to use cuPyNumeric themselves, so they have to use something else. My initial version used CuPy for the implementation of my GPU Python task. As a first pass, it was easy for me to add a top-level switch to enable NumPy as an alternative to CuPy.

That gets me CPU and GPU, but still leaves out the multi-CPU/OpenMP case. I cannot afford to run a rank per CPU core (nor would that likely be efficient), and vanilla NumPy has very limited support for OpenMP (essentially only what BLAS will do on its own). Keep in mind I run a variety of NumPy operations, not limited to what ultimately boils down to a BLAS call.

I would like to do this without writing my code multiple times. CuPy and NumPy are close enough I can stick to one implementation and stub out the differences. Other solutions like Numba would require a code rewrite, and would not actually be more ergonomic (i.e., the CuPy/NumPy-style implementation is actually how I want to write this code, I do not want to be writing loops, especially with all the implicit broadcasting in the code).

I can think of a couple options:

1. Somehow get NumPy to run on OpenMP. As far as I'm aware NumPy does not have complete support for this and I would only get a subset of operations parallelized this way.
2. Permit Python tasks to call cuPyNumeric operations, and somehow set things up so that, inside of a Python task, these calls directly call the leaf implementations for the appropriate architecture. See Nested parallelism #993 for details.

(For LANL/SLAC.)

[elliottslaughter]

A bit more research:

NumPy does NOT support parallelizing any operations, except those that would be parallelized by an underlying BLAS layer:

NumPy itself is normally intentionally limited to a single thread during function calls, however it does support multiple Python threads running at the same time. Note that for performant linear algebra NumPy uses a BLAS backend such as OpenBLAS or MKL, which may use multiple threads that may be controlled by environment variables such as OMP_NUM_THREADS depending on what is used.
https://numpy.org/devdocs/reference/global_state.html#number-of-threads-used-for-linear-algebra

Any solution would need to support NumPy ufuncs (i.e., whole-array operations like element-wise +, -, *, etc.), so I think my needs exceed what can be done automatically via BLAS.

[manopapad]

One possibly sustainable solution would be to add support for Legate task launching to work with PhysicalStores directly (and in that case just jump to the task entry directly). Then we can have the same cuPyNumeric calls working both at the top-level, and in tasks (in the latter case only executing on the device or set of cores managed by the current processor).

Some of this path already exists, in the "fast-path" experimental mode, but more refactoring may be required.

I'd like @Jacobfaib's advice here. Do you think 2 weeks would be sufficient (assuming we prioritized this) to have a proof of concept? If so, then we can at least present this is a WIP solution to the customer.

[Jacobfaib]

If we dropped everything else, we could maybe have a very very rough proof of concept in 2-3 weeks, sure.

[elliottslaughter]

My two cents are this is not worth dropping everything else for. However it would be nice to have a roadmap and timeline I can present to the customer.

[manopapad]

@Jacobfaib promised to put together an outline for this task, so we should have a plan to share by end of the week (?)

[sbahirnv]

Python tasks are not permitted to use cuPyNumeric themselves, so they have to use something else. My initial version used CuPy for the implementation of my GPU Python task. As a first pass, it was easy for me to add a top-level switch to enable NumPy as an alternative to CuPy.

@elliottslaughter , Is it possible to share a sample code how you are doing this? We can start here and work backward to add necessary legate support.

[elliottslaughter]

What do you want a code sample for? The switch in the quote is literally just an if at top-level:

if ...:
    import numpy as np
else:
    import cupy as np

It's a bit more complicated than that because there's the condition and a few differences I have to smooth over, but otherwise it works (besides the fact that NumPy is largely single-core only).

[elliottslaughter]

This is coming up again because I have a design decision to make based on whether this is available or not.

Basically, if I can assume we'll have this, there's some existing code I can reuse for a new feature. If not, I'll have to design the code to work around it. So I'll end up with more code and/or more complexity, which will impact downstream users who later need to come and modify this (and/or reuse it for other stuff).

Code portability also allows the code to be more modular because I can reuse pieces more aggressively without duplicating stuff in the code base.