Memory requirement of ak.sum vs np.sum

[alexander-held]

Following some investigations of @ekauffma into memory use, I observed that there seems to be a significant difference between using awkward and numpy for doing partial sums of higher-dimensional (large) arrays that I do not understand. I am curious if this is expected.

import numpy as np
import awkward as ak

x = np.random.random(size=(1000, 1000, 100))

ak_res = ak.sum(x, axis=1)
np_res = np.sum(x, axis=1)

assert np.allclose(ak_res, np_res)

I am running the script above through memray via the following mini script:

rm -r output.bin
rm -r memray-flamegraph-output.html
memray run -o output.bin test.py
memray flamegraph output.bin
open memray-flamegraph-output.html

I see a 763 MiB allocation for the array x, which makes sense given float64 and 100M floats. With ak.sum removed from the script, this dominates in terms of memory.

Surprisingly, the addition of ak.sum results in an additional 2.3 GiB of allocation via various nextcontent / nextparents / nextcarry objects, see the screenshot below (can open in new tab to have something more readable, or reproduce with memray).

I had naively assumed thus far that awkward just hands this off to numpy and would essentially act the same otherwise.

I am using Python 3.9.16, awkward 2.2.1 and numpy 1.24.2 on macOS.

[jpivarski]

I had naively assumed thus far that awkward just hands this off to numpy and would essentially act the same otherwise.

It's the other way around: if np.sum is called on an ak.Array, it checks to see if ak.Array has a __array_function__, which it does, and then NumPy calls that instead of its usual np.sum operation.

What you're doing here is calling ak.sum on a NumPy array, which converts the NumPy array into an Awkward Array (which is not expensive) and then does general variable-length reduction on that Awkward Array (which is expensive). There are more intermediate arrays involved in ak.sum, and that will account for more memory.

We could install a check in all of the NumPy functions we NEP18-overload to see if there are no Awkward Arrays in the arguments and then pass it over to the equivalent NumPy function (here). But, there's an exception to that: one of the arguments might not be an Awkward Array, but it might be an iterable that should be interpreted as an Awkward Array instead of a NumPy array (e.g. it's a list of variable-length lists), in which case you would not want the NumPy version to take over. So the hardest part of this optimization would be determining if the arguments of ak.whatever ought to be Awkward-like. Also, ak.sum should return an Awkward Array, regardless of whether it is implemented by np.sum or not.

[alexander-held]

Thanks for clarifying this Jim!

x = np.random.random(size=(1000, 1000, 100))
y = ak.Array(x)
np_res = np.sum(y, axis=1)

also reproduces the large memory requirement, so that is all consistent then.

At the risk of falling down a rabbit hole: how crucial are these intermediate arrays for performance? I imagine that there are still ways of doing this with negligible memory overhead via a jit-compiled iteration over the array? This example requires 3x additional memory compared to the array itself, are there other API calls that might require even more (and if so, could other approaches be useful there)?

Further down the rabbit hole: I assume other calculations also may use similar intermediate arrays (?) — if so, can dask-awkward help avoid overhead from creating the same (or similar) objects by graph optimization when users call multiple awkward functions on the same objects?

[jpivarski]

Some intermediate arrays can be avoided in the axis=-1 case, although that would mean adding another implementation path. (axis=-1 is the most important case, I know.) JIT-compilation is unnecessary for that: there is enough information available for everything to be precompiled.

dask-awkward can't fuse Awkward operations, if that's what you're asking.