Improve the implementation of dask.array.block

[hmaarrfk]

dask/dask/array/core.py

Line 2899 in 14b4c38

def block(arrays, allow_unknown_chunksizes=False):

I recently rewrote what will likely be numpy's new implementation of block. @eric-wieser suggested an enhancement which should make it really trivial for dask to piggy back off the results and use it in it's implementation.

numpy/numpy#11971 (comment)

The first time I tried to use dask.block, it was really slow, so I had to write my own. I think this might help it since it avoid calls to concatenate all together.

Not sure about unknown dimensions.

[mrocklin]

Improvements are always welcome. Two questions:

1. That implementation seems to use __setitem__ syntax, which probably won't be as useful for dask array. Does that affect your plans?
2. Why was dask.array.block slow? Was it slow in graph construction overhead, or in execution?

[hmaarrfk]

1. Eventually we will have a helper function that returns the dtype, shape, slices and numpy arrays. These could easily be dask arrays. An other function will take all of those and create a new array and do the assignment. I Think that is within the realm of dask since assignment only happens to newly created arrays.

2. Not too sure. It probably wasn't any slower than block. I think ultimately it failed at handling one of my cases, so I couldn't use it.

I just wanted to point somebody to a no concatenate solution to this problem.

[mrocklin]

This might also be helpful in replacing the current concatenate3 function in dask.array.core, which also does a no-concatenate blocking solution, but is probably uglier than what numpy has. If numpy.block is now convenient and efficient then it might be good to switch over, especially if there is a performance increase.

[hmaarrfk]

Here is some more concrete information about the challenge:

The current implementation adds 24 graph operations for a a 3D blocking operation with 8 sub arrays.
Ideally, everything can be precomputed and the graph would just have to "place" the blocks in the correct location. This would cause maybe 8 operations to happen on the dask graph.

# import numpy as np
from dask import array as da

n = 10
# Slow setup method: hence separated from the others above
a000 = da.full((2 * n, 2 * n, 2 * n), 1, int, chunks=-1)

a100 = da.full((3 * n, 2 * n, 2 * n), 2, int, chunks=-1)
a010 = da.full((2 * n, 3 * n, 2 * n), 3, int, chunks=-1)
a001 = da.full((2 * n, 2 * n, 3 * n), 4, int, chunks=-1)

a011 = da.full((2 * n, 3 * n, 3 * n), 5, int, chunks=-1)
a101 = da.full((3 * n, 2 * n, 3 * n), 6, int, chunks=-1)
a110 = da.full((3 * n, 3 * n, 2 * n), 7, int, chunks=-1)

a111 = da.full((3 * n, 3 * n, 3 * n), 8, int, chunks=-1)

block = [
[
[a000, a001],
[a010, a011],
],
[
[a100, a101],
[a110, a111],
]
]
blocked = da.block(block)
blocked.visualize()

I'm not too sure how to handle unknown dimensions. That is really a unique (and quite cool feature) of dask.

I'll try to followup with a real benchmark soon and look into concatenate3

[mrocklin]

Any update on this @hmaarrfk ?

[hmaarrfk]

I started with a benchmark. It shows basically that it suffers a little bit from overhead. But not from copying data.

[jakirkham]

FWIW a lot of the code for block was lifted from NumPy. This was done to avoid NumPy usages of asarray, which trigger a computation, or things like concatenate, where a Dask friendly implementation is needed. If there have been changes in NumPy, that would help Dask, it is probably reasonable to update our vendored copy. Better yet if we can consolidate code between the two, but that likely isn't possible without NumPy addressing asarray and concatenate for this use case in a nicer way.

[hmaarrfk]

I've been working on getting a numpy implementation without concatenate. I kinda messed up with the review process, but it might help if you put your thoughts here: numpy/numpy#11971 You can just read the top and bottom and skip the middle

The secondary challenge was that the existing benchmarks were hard to interpret numpy/numpy#12001 and a bug in the existing implementation made one benchmark made my proposed implementation look so much worse than the existing one numpy/numpy#11979

Ultimately, different algorithms are applicable for different datasets. For blocking small arrays (512x512 final size), it actually turned out that concatenate is much faster than scanning over all the data.

Dask also offers alot of optimizations like this notion of squashing linear operations. I still need to experiment with it, but initial benchmarks dask/dask-benchmarks#18 show that there isn't much more than 10% to be gained.

[mrocklin]

Checking in here. Any update on this @hmaarrfk ? Is there anything that you are blocked on here on the dask side?

[hmaarrfk]

No, more I'm waiting for the numpy PR to go through. A few peculiarities have come up.

[mrocklin]

OK, just checking in. Thanks for the update.

[hmaarrfk]

Yeah, so while the dask operation is built on multiple dask implementations of concatenate, those aren't actually costly as data isn't copied until the very end. Maybe graph operations can be reduced, making the graph cleaner, but it isn't something I am enthousiastic about because that numpy PR took so long to get through, probably in large part due to my novice development style. I am secretly hoping now that numpy is expanding their high level interface, that maybe we can expose the internal algorithm of block and use that for dask. Dask would depend on numpy 1.16 or 1.17 I'm not sure what would happen when you block two dask arrays with chunks that don't line up. An operation that is probably allowed now.

…

On Mon, Oct 22, 2018 at 5:12 PM Matthew Rocklin ***@***.***> wrote: OK, just checking in. Thanks for the update. — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#3987 (comment)>, or mute the thread <https://github.com/notifications/unsubscribe-auth/AAFfmMnFsj5OHyn1oZegBF6catxuy0dfks5unjTFgaJpZM4WrZoo> .