Better GPGPU integration

Greetings,

I'm trying to figure out how dask can efficiently implement a hybrid GPGPU algorithm. Apologies if this is going to be a long post - as I think it's best if I share my analysis of the problem before I get to my proposed solution.

## Memory management

GPGPU on dask falls in 4 big use cases:
1. discrete graphics card with seamless unified memory (e.g. NVidia Pascal and later).
2. integrated GPU with shared memory - e.g. Intel / AMD CPUs with a OpenCL-capable GPU, NVidia Tegra, etc.
3. discrete graphics card _without_ seamless-unified memory, but implementing algorithms where copying all memory from host to GPU and from GPU to host every single time you run a kernel only causes a minor performance loss. Again, the dask kernels always acquire in input and return plain numpy arrays.
4. discrete graphics card _without_ seamless-unified memory, where copying memory back and forth has a major cost, so it's necessary to find a way to explicitly keep the data in GPU space.

In use cases 1, 2, and 3, the input and output of all dask kernels is a plain numpy array, whose underlying RAM _may_ be in GPU or host space.

In the discussion below, I'm going to deliberately ignore use case 4, as it would add a massive amount of complication both in terms of handling non-numpy data between kernels and in preventing out-of-memory errors in GPU space.

## Limits of the current dask design

To my knowledge, the only support that dask offers for GPGPU is provided by http://distributed.readthedocs.io/en/latest/resources.html.

This has several limitations:
- No support (as far as I can tell) for DataArray et al. structures
- Not compatible with plain dask (i.e. not distributed) -> doesn't scale down.
- No way to minimize memory transfers on multi-GPU systems.
There's no way to tell dask "this function is going to return a numpy array whose data is in GPU space; if the output is going to be processed by another GPU-based kernel, it would be best to use the same GPU as before".
- no way to pipeline CPU and GPU calculations. Let's say for example that you have a legacy algorithm which is entirely CPU based, e.g. a two-stages embarrassingly parallel problem:

```
C C C C C C C C
^ ^ ^ ^ ^ ^ ^ ^
C C C C C C C C
```
C = dask kernel that consumes 100% of 1 CPU core

Now you want to convert only the second stage to GPGPU, as converting the first stage would have a poor cost/benefit ratio - e.g., 80% of the complexity of the algorithm is in stage 1, but 80% of the time is spent in stage 2.

```
G G G G G G G G
^ ^ ^ ^ ^ ^ ^ ^
C C C C C C C C
```
G = dask GPU kernel

Let's say you run the above on a 4x CPU, 1x GPU machine. In theory, one could run 4 stage-1 kernels on CPU, and then the other 4 stage-1 kernels in parallel with the first batch of stage-2 GPU kernels. However, there's no way to tell this to dask, as (to my understanding, please correct me if I'm wrong) if you run a task with ``resources={GPU: 1}``, it will still lock 1 CPU.
You could, of course, work around it by starting your worker with ``--resources CPU=4 --threads 999`` and then explicitly setting ``resources={CPU: 1}`` on every CPU-based kernel - but when working with legacy code and/or higher-level libraries (e.g. xarray, or even just the dask DataArray framework) this can become incredibly hard.

This problem is also present in non-GPU problems where you have a slow I/O task - there's no way to tell dask that the CPU consumption of that task is going to be negligible.

## Proposed design - first iteration

1. Move non-CPU resource management to plain dask (e.g. not distributed) - see below.
2. Write a ``@dask.resources(res1=quantity, res2=quantity...)`` decorator, which can then be applied to dask kernels. The implementation would be to simply add a ``dask_resources`` attribute to the decorated function, in order to retain seamless compatibility with dask dicts.
3. Treat "CPU" as a special, hardcoded resource, which defaults to 1 if omitted and can be set to 0.
4. When using dask.multithreading.get, the available resources will default to CPU=. When using dask.multiprocessing.get, the available resources will default to CPU=1. Add extra resources with ``dask.set_options(resources={"GPU": 1})``.

```
# Don't lock a CPU core, as this kernel is going to be 99% GPU
@dask.resources(GPU=1, CPU=0)
def my_dask_kernel(my_input):
my_output = numpy.empty_like(my_input)
# Use numba.cuda, PyCUDA, PyOpenCL, or ctypes/cffi linking to a .cu library
my_gpu_kernel(my_output, my_input)
return my_output
```

## Second iteration
This second batch of changes is specifically aimed at multi-GPU systems. Without it, multi-GPU systems that run single-GPU problems should probably start 1 independent worker per GPU - although it will cost needless pickling-unpickling and IPC, and it will make it impossible to use nvlink to transfer memory across GPUs.

1. In the ``@dask.resources`` decorator, accept an optional parameter ``affinity=["GPU"]`` which indicates that the task _should_ (if possible) run on the same GPU Resource ID as the task that produced its inputs.
2. Offer a function, ``dask.resource_id(resource_name)`` which returns a thread-local _resource ID_.
e.g.
```
# In multi-GPU systems, it's preferred to run on the same GPU that produced my_input in order to avoid a costly memory transfer.
@dask.resources(GPU=1, CPU=0, affinity=["GPU"])
def my_dask_cuda_kernel(my_input):
gpuid = dask.resource_id("GPU") # e.g. 0 or 1
my_output = numpy.empty_like(my_input)
my_cuda_kernel(my_output, my_input, gpuid)
return my_output
```
The resource ID is
- an integer counting from zero, for all resources that were requested as an integer =1
e.g. ``@dask.resources(GPU=1)`` -> ``dask.resource_id("GPU")`` may return 0, 1, 2, 3 on a quad-GPU host
- a tuple, for all resources that were requested as an integer > 1.
e.g. ``@dask.resources(GPU=2)`` -> ``dask.resource_id("GPU")`` may return (0, 3) on a quad-GPU host
- won't work for float resources (e.g. RAM)

3. When defining which resources are available, either from the command-line or with dask.set_options, allow grouping together resource IDs in _affinity groups_. The idea is that, if you have a 2 GPUs system, you really do want to have 4 single-GPU tasks running at the same time, in order to obtain optimal GPU usage and minimize the overhead.

e.g. ``dask.set_options(resources={"GPU": ((0, 1), (2, 3))}`` describes a 2-GPU system with 4 GPU tasks. When a task specifies ``@dask_resources(GPU=1, affinity=["GPU"])``, and its input ran on GPU 0, then it means it's going to be optimal to run the task either on GPU 0 or GPU 1.

4. When affinity groups are defined and ``@dask.resources`` requests more than 1 GPU, it should always mean that they must be within the same affinity group if there's a group large enough. e.g. in the configuration above, ``@dask.resources(GPU=2)`` means that the task will run on one whole graphics card, whereas ``@dask.resources(GPU=4)`` will run on both graphics cards.

Opinions? Did I miss major designs already existing or in the works? Can somebody come up with simpler ideas to tackle the problem?

Thanks