I spent some time benchmarking two implementations (see code in this gist, heavily lifted from the code referenced above):

array_api: uses normal dask array api (operators, dask ufuncs, etc...)
map_blocks: uses a single function and apply it to the input arrays using a single call to da.map_blocks.

Results:

On my computer the fastest implementation for a given input array size was map_blocks using chunks=1e5. For nopt=1e7 (the input size) and chunks=1e7 we can break the execution into 4 stages (build, optimize, compute, and finalize):

(dask) jcrist dask $ python bench.py --nopt=1e7 --chunksize=1e5
using_map_blocks
=======================
chunksize:       100000
numblocks:       100
-----------------------
Graph Size:      400
Optimized Size:  400
-----------------------
Build:          0.003 s
Optimize:       0.005 s
Compute:        0.314 s
Finalize:       0.088 s
Total:          0.411 s

using_array_api
=======================
chunksize:       100000
numblocks:       100
-----------------------
Graph Size:      2900
Optimized Size:  2500
-----------------------
Build:          0.017 s
Optimize:       0.059 s
Compute:        0.473 s
Finalize:       0.088 s
Total:          0.638 s

Looking at the above, the array api version is slower due to:

Higher optimize and build cost. This is mostly due to the larger graph size (2900 tasks vs 400).
Higher compute cost. This turns out to be mostly due to the increased relevance of scheduler overhead at this relatively small chunksize. The single mapped function in the map_blocks function takes much longer than a single add call, so the overhead is less relevant there.

The array api version performs much better for larger chunk sizes:

Note that for larger chunks the map_blocks call performs much worse. This is due to the single function taking longer and fewer opportunities for parallelism (fewer blocks overall). In contrast the array_api version has plenty of smaller tasks still leaving opportunities for parallelism.

Note that at larger chunk sizes the relevance of graph build and optimize time also goes down. This scales fairly linearly with number of blocks (log scale to show both lines).

For this test case using the array api leads to ~10x slowdown on build and optimize, but for chunks of sufficient size this is negligible compared to compute cost.

Opportunities for speedups

This is a weird benchmark because it's very sensitive to chunksize. The best performance occurred when all the operations were grouped together into a single task and the chunksize was relatively small. This works fine, and is fairly readable, but it'd be nice to speedup the array_api version to be more compatible.

In Fuse array elementwise operations at graph build time #2538 it was suggested that by delaying graph building from atop operations (most everything here) we can fuse tasks at build time. This would speedup build time (fewer calls to atop) and optimize time (smaller graph to optimize). It's not clear if it would speed up compute time, as the result should be similar to that of our normal task fusion. It would lower the overhead of using smaller chunk sizes though.
In dask.array.jit or dask.array.vectorize #1946 it was suggested that we might interpret fused tasks to intelligently useout keywords, or jit fused tasks with numba to speedup individual tasks. This would probably speedup compute time, but also is a bit magicical.

Uh oh!

Improve dask.array Black Scholes benchmark performance #2497

Description

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions