Motivation

This page captures a plan for 3-steps PR that aims to provide memory and compute improvements mostly arising from fewer temporaries, copies and collectives in DistributedArray and CG/CGLS. Each Phase below will become one commit in the PR.

Context

The rationale of this work comes from large-scale experiments that showed suspicious high peaks of memory usage and slowdowns due to use of swap memory. More specifically, when solving a large problem with pylops-mpi CGLS, about 40-50% of each iteration goes to full-vector copies and temporaries. On top of that, every small array operation runs blocking collectives.

Execution

Phase A - DistributedArray (pylops_mpi/DistributedArray.py)

A1. Build result arrays without collectives

• Add a private _like(self, local_array=None, dtype=None). It creates the instance with
  object.__new__ and copies _global_shape, _base_comm, _base_comm_nccl, _partition,
  _axis, _mask, _sub_comm (reused, never split again), _local_shape, _engine and
  dtype. It then takes local_array or allocates ncp.empty(local_shape, dtype). It runs no
  checks and no collectives, because the metadata comes from an array that is already valid.
• Cache local_shapes after the first allgather, in self._local_shapes. If __init__ was
  given local_shapes, store them there. _like copies the cache.
• Use _like in __neg__, add, multiply, conj, copy, zeros_like and empty_like.
  Also use it in StackedDistributedArray.zeros_like and .empty_like, which currently call
  distarray.local_shapes.

A2. Arithmetic with out= and real in-place operators

• Add one private helper, _ufunc(self, ufunc, other, out):
  • If other is a DistributedArray, run the existing _check_partition_shape and
    _check_mask, then use other.local_array. Otherwise use the scalar or array as is.
  • If out is None, set out = self._like(dtype=result_dtype), where result_dtype follows
    the current behaviour (self.dtype).
  • For SCATTER and UNSAFE_BROADCAST: ufunc(self.local_array, other_arr, out=out.local_array).
    This is one pass with no temporary. numpy and cupy ufuncs both accept out=.
  • For BROADCAST: keep out[:] = ufunc(...), so the rank-0 broadcast in __setitem__ still
    runs and the semantics don't change.
• Change the public methods to add(x, out=None) and multiply(x, out=None), and add
  subtract(x, out=None). All three go through _ufunc, so existing calls behave the same.
• Rewire the operators:
  • __sub__ calls subtract. Today it is add(-x), which allocates the negation first.
  • __iadd__, __isub__ and a new __imul__ use out=self.
  • iadd keeps its name and calls add(x, out=self).
  • __neg__ uses ncp.negative(..., out=...).
  • __rsub__ is missing, so scalar - x fails today. Leave it out.
• Edge case: if out= would need an unsafe cast (for example complex into a float32 out),
  numpy raises. Today [:]= casts silently with a warning. For that case, go back to the
  allocate-then-assign path so behaviour stays the same.

A3. dot and norm

• Use ravel() instead of flatten() in dot (line 836) and norm (line 997). For contiguous
  local_array, ravel() is a view.
• Add private _local_dot(other, vdot) and _local_sqnorm(), which return the un-reduced
  partial sum. dot becomes allreduce(_local_dot(...)). The BROADCAST to_dist path stays as
  it is.
• Add a fast path for 2-norms in _compute_vector_norm: compute abs(vdot(a, a)) locally,
  allreduce it as float64, then take the square root. This drops the two float64 temporaries.
  Other ord values keep the current code.
  • Precision change: the local partial sum is accumulated in float32 for float32 data, the
    same as dot already does, instead of float64 today. Cost histories will move at the level
    of float32 rounding.

Phase B - StackedDistributedArray (same file, from line 1274)

• Change add, iadd, multiply, __neg__, __sub__ and __isub__, and add subtract and
  __imul__.
  • Each method allocates its result once with self.empty_like() (now free of collectives), or
    uses out/self.
  • It then calls the child method with out=out[iarr], recursing into nested stacks.
  • Today, for example, add does self.copy(), then a child add that builds a new array plus a
    temporary, then [:] copies the result. That is three full passes where one is enough.
  • Add out=None to add, subtract and multiply, matching the DistributedArray API.
• Fused reductions in dot and norm (ord=2):
  • Fast path: walk the leaves, add up their _local_dot or _local_sqnorm, and run a single
    allreduce on the shared base_comm.
  • Only take it when every leaf is SCATTER with mask is None and they all share base_comm
    and base_comm_nccl. Otherwise use the current per-child loop.

Phase C - Solvers (pylops_mpi/optimization/cls_basic.py)

CGLS (setup at line 316, step at line 381)

• Persistent buffers: x, c, s, and self.r (the previous r, reused as scratch). q
  and r are new arrays each iteration only because Op.matvec and Op.rmatvec allocate them.
  Operators are out of scope.
• New step, using the in-place operators from Phases A and B:

  qq = self.q.dot(self.q, vdot=True)
  if self.damp != 0.0:
      qq += self.damp * self.c.dot(self.c, vdot=True)
  a = float(np.abs(self.kold / qq))
  self.c.multiply(a, out=self.r)   # previous r is dead after the last c update
  x += self.r
  self.q *= a                      # q is replaced by the next matvec
  self.s -= self.q
  r = self.Op.rmatvec(self.s)
  if self.damp != 0.0:
      x.multiply(self.damp, out=self.r)
      r -= self.r
  k = float(np.abs(r.dot(r, vdot=True)))
  b = k / self.kold
  self.c *= b
  self.c += r
  self.r = r
  self.q = self.Op.matvec(self.c)
  ...
  cost1 = self.cost[-1] if self.damp == 0.0 else sqrt(cost**2 + damp * |x.x|)

• The arithmetic is the same as today: x + a*c, s - a*q, r + b*c. The only differences
  should come from how rounding falls with out=.
• Make the same changes in setup: drop the x * damp and x.dot(x) work when damp == 0,
  and use vdot=True instead of .conj().
• Leave alone: in setup, damped_x = x * damp and cost1 use the un-squared damp, while
  step uses damp**2. Serial pylops does the same (pylops cls_basic.py lines 581 and 605).
  I'm noting it here, not fixing it in this change.

CG (step at line 114)

Opc = self.Op.matvec(self.c)
a = float(self.kold / np.abs(self.c.dot(Opc, vdot=True)))
Opc *= a
self.r -= Opc
self.c.multiply(a, out=Opc)   # Opc is free now; same space as c
x += Opc
k = float(np.abs(self.r.dot(self.r, vdot=True)))
self.c *= k / self.kold
self.c += self.r

• Watch the order of c.dot(Opc): today it is c.dot(Opc.conj()), not vdot. For complex
  data, keep self.c.dot(Opc.conj()), since its conj copy is unavoidable. Use vdot only where
  both operands are the same vector.

ISTA and FISTA (cls_sparsity.py)

No solver changes. They get the Phase A and B savings automatically.

Tests (tests/test_distributedarray.py, tests/test_stackedarray.py)

Reuse the existing par* fixtures. Add tests for:

• out= results: each method with out= matches the operator result (real and complex,
  scalar and array operand). For +=, -= and *=, check that local_array is the same object
  before and after (is).
• Masked arrays: (a + b).sub_comm is a.sub_comm, so no new communicator is created.
• BROADCAST: results are unchanged and equal on every rank.
• Stacked arrays: nested-stack arithmetic with out=. A fused dot/norm equals the sum of
  the per-child results, both on the fast path and the fallback path (mixed partitions or masked
  children).
• Solvers: the existing test_cg* and test_cgls* tests in tests/test_solver.py
  (broadcast and stacked variants) must pass unchanged.