Memory spikes

[jvdoorss]

A number of changes intended to avoid short-lived memory allocations for intermediate variables.

More complicated one-liners can imply that several like-sized Tensors are stored until the computation is completed, even though a more sequential approach could reduce that number. For multi-GB tensors in the trunk-module, containing $\mathcal{O}(N_{residue}^2)$ elements, this can prove a crucial optimization.

Predicting the structure of 9b9j on a 40GB A100 gpu failed without these changes, but succeeds with them.

Typical changes:

• alter z = z + a to z += a (or z[:] = z + a) to write directly to already allocated memory
• reuse allocated memory in the outer scope by not initializing a new variable if the shape is constant:

  def inner_scope(z):
    z[:] = shape_preserving_op(z)
    ...
  

  instead of

  def inner_scope(z):
    z = shape_preserving_op(z) #now there's 2 z's
    ...
  

• delete redundant tensors
• split up one-liners

Clearly this approach has some drawbacks, reducing speed (probably) and versatility of methods (that now modify the outer scope), but might be an acceptable trade-off.

NOTE While I think changes like this will in general optimize memory movements, as e.g. argued here and as seen from the actual memory consumption improvement of Boltz, it could be that not all proposed changes have an actual effect.