Fusing large operations to reduce memory usage

[elliottslaughter]

I have a large code I'm working on. Somewhere deep in the bowels of this code I have a line that essentially amounts to:

long_array * tall_array

Much later, 3-4 levels up the callstack, I have a function that applies a cpn.sum over one of the axes, essentially throwing away one of the (large) dimensions.

Right now, this code scales very inefficiently. I'm forced to use multiple nodes purely for memory capacity reasons, but the code actually doesn't need it: it runs far too quickly for distributed execution to make sense, and I'm essentially throwing away the compute.

The ideal solution would be to fuse the multiply and the sum to avoid the memory bloat. Doing this in user code is painful, because the code is specifically factored to be reusable. The offending multiply is multiple levels down the call stack, every one of which provides a conceptually distinct purpose, and could be called in arbitrary other code. I'm essentially breaking down the code's abstractions in order to apply this optimization manually.

If this could be done automatically with reasonable overhead that would be far more effective from a code reuse and readability perspective.

LANL/SLAC, medium priority.

[manopapad]

To confirm, the * here is elementwise multiplication, not matmul, correct? So the code looks something like this?

long_array = np.ones((1,1000))
tall_array = np.ones((2000,1))
x = long_array * tall_array
...
x.sum(axis=1)

I.e. the size bloat is happening due to broadcasting?

[elliottslaughter]

The basic skeleton is correct, but the core operation is a matmul (which expands the data) and then additional (binary and/or unary) operations are applied before the sum. I believe all of the inputs are small-ish so any additional operations would have to be broadcasted to match.

[manopapad]

but the core operation is a matmul (which expands the data)

In that case, is it rather tall_array @ long_array? Because long_array @ tall_array I would expect to return a small result.

[manopapad]

@elliottslaughter conceptually, if the intermediate result of the multiplication isn't directly used, but only summed over, could you have done the same using an einsum (ignoring issues of code refactoring):

E.g. rather than

x = tall_array @ long_array
y = x.sum(axis=1)

could you have done?

y = np.einsum("ab,bc->a", tall_array, long_array)

[elliottslaughter]

For the specific code path I saw, yes this would work. The potential obstacles are:

1. There's a couple hundred lines of code here that I haven't had time to inspect deeply, and I'm not sure exactly what all goes into it. It's entirely possible that other operations are involved that may not cleanly integrate into this formulation (e.g., I can't just multiply into the inputs).
2. As mentioned previously, the code is split across multiple functions and the factoring is not totally obvious.