Add fusion for ufuncs and perhaps other operations

[agoose77]

From today's meeting:

In dask-awkward we could look to fusing ufunc operations like x * 1 + 2 using the task graph. This is effectively what NumExpr does, but without requiring the string processing.

We should also look to identify operations that are impossible or very memory intensive at the moment, such as ak.cartesian.

[jpivarski]

We were thinking that a Dask DAG could be used to identify sequences of fusable operations and fuse them.

That leaves open the scope:

1. Only ufuncs? Easy and could be valuable for users with large formulas to compute¹. However, there's nothing specifically Awkward about that: it would be equally useful in dask.array and xarray, and these have much larger userbases.
2. Filter operations of the form array[ufunc_only_expression_with_boolean_output]. The goal would be to not only fuse the ufunc-only predicate, but also the application of the predicate to array so that it never needs to make a large boolean array (and loop over the same out-of-CPU-cache data multiple times). This use-case is also equally applicable to dask.array and xarray. It's more complicated for Awkward Arrays because the predicate might have a different depth than the array. At worst, though, this operation would need to produce a new NumpyArray and ListOffsetArray offsets (without knowing in advance how large either of them are; this is a job for Numba's TypedList).
3. ak.cartesian and ak.combinations can both produce large arrays that are only used as intermediates, which can easily exceed memory constraints. As it is right now, users need to choose a different strategy (ak.cartesian/ak.combinations or Numba) depending on that size, which is not good because it's not a gradual path from simple implementations to complex implementations—it's a clean-break/rewrite. Fusing a path from combinatorial-expansion through ufunc-like operations to reduction would be ideal. It's that grow-flat-shrink plateau that might want to be in-memory or virtual depending on how much memory the plateau would take in the middle. This is what Femtocode was for, but maybe now we can actually do it. However, it's also the largest-scope option we've considered, and it's not clear to me what the boundaries are.²

Implementation:

• If only ufuncs, the fused operation can be implemented by NumExpr or Numba. NumExpr is a lighter dependency, but can never expand beyond the ufunc-only case.
• Can we just write Numba functions for each Awkward function and have Numba/LLVM automatically fuse them? I'd be surprised if compilers can be this magical, but it's worth an investigation.
• If we have to generate our own pre-fused Numba code, then we can only consider limited-scope cases.

Footnotes

1. Note that the Vector library consists of a lot of ufunc-only formulas in which ufunc-fusion could provide a big lift. It would be interesting to see how complicated formulas involving vector manipulations get in Coffea. ↩

2. Considering how hard it is to write these plateaus in Awkward anyway (c.f. Gordon's question in Slack), maybe this should be a new interface anyway, and the borders delineating that new interface would simplify the scope and implementation of this kind of project. ↩