Running in Parallel

A power diagram can be computed by several devices or processes at once. The call is still PowerDiagram( ... ). What decides is the positions: split along the seeds between several ranks, they give one box of seeds per rank.

JAX ( devices )

import os
os.environ[ "XLA_FLAGS" ] = "--xla_force_host_platform_device_count=4"  # 4 CPU devices, to try it on a laptop ( not on GPUs )

import jax
import numpy as np
import sdot
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P

jax.config.update( "jax_enable_x64", True )
mesh = Mesh( np.array( jax.devices() ), ( "p", ) )

rng = np.random.default_rng( 0 )
positions = jax.device_put( rng.random( ( 10000, 2 ) ), NamedSharding( mesh, P( "p", None ) ) )  # split along the seeds
weights = rng.random( 10000 ) * 1e-3

pd = sdot.PowerDiagram( positions, weights, boundaries = sdot.box_half_spaces( np.zeros( 2 ), np.ones( 2 ) ) )
m = pd.measures         # one jax array, split between the devices like the positions
print( m.sharding.spec, m.sum() )

PyTorch ( processes )

# torchrun --nproc-per-node 4 parallel_torch.py
import loom
import numpy as np
import sdot
import torch
import torch.distributed as dist

dist.init_process_group( "gloo" )  # "nccl" on GPUs
ranks = loom.TorchRanks()
rank, size = dist.get_rank(), dist.get_world_size()

# each process holds its block of the seeds ( cut here from a cloud they all draw the same way )
rng = np.random.default_rng( 0 )
positions, weights = rng.random( ( 10000, 2 ) ), rng.random( 10000 ) * 1e-3
n = len( positions ) // size
block = lambda a: torch.as_tensor( a[ rank * n:( rank + 1 ) * n ] )

pd = sdot.PowerDiagram( loom.spread( block( positions ), ranks ), loom.spread( block( weights ), ranks ),
                        boundaries = sdot.box_half_spaces( [ 0, 0 ], [ 1, 1 ] ) )
m = pd.measures         # spread like the positions: this process holds the measures of its seeds ( `m.value` )
total = m.sum()         # over all the processes: every one gets it
if rank == 0:
    print( float( total ) )  # 1: the area of the unit square
dist.destroy_process_group()

MPI ( processes )

# mpiexec -n 4 python parallel_mpi.py
import loom
import numpy as np
import sdot
from mpi4py import MPI

comm = MPI.COMM_WORLD
ranks = loom.MpiRanks( comm )
rank, size = comm.Get_rank(), comm.Get_size()

# each process holds its block of the seeds ( cut here from a cloud they all draw the same way )
rng = np.random.default_rng( 0 )
positions, weights = rng.random( ( 10000, 2 ) ), rng.random( 10000 ) * 1e-3
n = len( positions ) // size
block = lambda a: a[ rank * n:( rank + 1 ) * n ]

pd = sdot.PowerDiagram( loom.spread( block( positions ), ranks ), loom.spread( block( weights ), ranks ),
                        boundaries = sdot.box_half_spaces( [ 0, 0 ], [ 1, 1 ] ) )
m = pd.measures         # spread like the positions: this process holds the measures of its seeds ( `m.value` )
total = m.sum()         # over all the processes: every one gets it
if rank == 0:
    print( float( total ) )  # 1: the area of the unit square

One device ( boxes in turn )

import numpy as np
import sdot

rng = np.random.default_rng( 0 )
positions, weights = rng.random( ( 10000, 2 ) ), rng.random( 10000 ) * 1e-3

# 8 boxes of seeds, handled one after the other on this device
pd = sdot.PowerDiagram( positions, weights, boundaries = sdot.box_half_spaces( [ 0, 0 ], [ 1, 1 ] ), nb_boxes = 8 )
m = pd.measures
print( m.sum() )  # 1: the area of the unit square

Context	The ranks	How the positions are split	Launched with	measures
JAX	the devices of a Mesh ( loom.JaxRanks( mesh ) )	a JAX array with a NamedSharding along the seeds	python script.py	one JAX array, sharded like the positions
PyTorch	the processes of torch.distributed ( loom.TorchRanks() )	loom.spread( block, ranks ) in each process	torchrun --nproc-per-node 4 script.py	this process's block
MPI ( NumPy, CuPy )	the processes of a communicator ( loom.MpiRanks( comm ) )	loom.spread( block, ranks ) in each process	mpiexec -n 4 python script.py	this process's block
One device	—	nb_boxes = K	python script.py	an array

What is spread

• JAX is one process driving every device. A JAX array already sharded along its first axis ( NamedSharding( mesh, P( "p", None ) ) ) counts as spread as it is. loom.spread( array, loom.JaxRanks( mesh ) ) does the same from a host array.
• PyTorch and MPI run one process per rank. Each process holds its own block of the seeds and says so with loom.spread( block, ranks ). loom.spread_host( whole_array, ranks ) is the other way round: every process passes the same whole array and keeps its block.
• The result is spread like the positions. With PyTorch and MPI it is a loom tensor whose .value is the block of the process, and whose reductions are over all the processes: m.sum(), m.max(), m.mean()... give every process the same result.
• loom.gather_to_host( m ) gives the whole of a spread result on the host.

With PyTorch and MPI, a reduction and gather_to_host are collectives: every process must call them.

The weights follow the positions: spread the same way, or given whole.

How it works

Each rank builds the cells of its own seeds, starting from the whole domain. A cell that its seeds can no longer cut, but that the seeds of another rank still could, travels there: its vertices go in a message, and the cell goes on being cut on the other side. A cell only shrinks, so it never goes back to a rank it has visited. When no cell is travelling any more, the measure of each cell goes back to the rank of its seed.

No seed is copied from one rank to another: what crosses the boundaries are cells. Two sizes are not known in advance: the vertices of a travelling cell ( msg_capacity, 16 by default ) and the cells a rank sends to another in one round ( inbox_capacity, 64 by default ). When one of them is too small, the computation runs again with twice as much, and the diagram keeps the new size for its next calls.

The same code for every framework

sdot writes the distributed computation once, as a program per rank, and loom runs it in each context: in a shard_map under jit for JAX, as it is in each process for PyTorch and MPI. A program reads the blocks of its rank and talks to the others through a few operations:

• ranks.rank(), ranks.size, ranks.all_sum( x ), ranks.exchange( buf ) ( buf[ j ] goes to rank j );
• loom.route( dest, rows, size, capacity ) packs rows by destination for an exchange, and loom.scatter_add( target, idx, vals ) adds values at indices. Both are loom kernels, the same code everywhere;
• loom.while_loop( cond, body, carry ), and loom.Capacity for the sizes found too small along the way.

import os
os.environ[ "XLA_FLAGS" ] = "--xla_force_host_platform_device_count=4"

import jax
import loom
import numpy as np
from jax.sharding import Mesh

ranks = loom.JaxRanks( Mesh( np.array( jax.devices() ), ( "p", ) ) )
slots = loom.Capacity( 2 )  # rows per destination: doubled, and the program run again, if too small


def program( ranks, values ):
    """each rank sends each of its values to the rank `value % size`, which adds up what it receives"""
    o = loom.ops( values )
    dest = values % ranks.size
    out = loom.route( dest, o.reshape( values, ( -1, 1 ) ), ranks.size, slots )  # [ size, slots, 1 ], -1 where empty
    got = o.reshape( ranks.exchange( out ), ( -1, ) )                            # what the others sent to this rank
    return o.reshape( o.sum( o.where( got >= 0, got, 0 ) ), ( 1, ) )


values = loom.spread_host( np.arange( 16, dtype = np.int32 ), ranks )  # 4 values per rank
sums = loom.spmd( program, ranks )( values )           # one sum per rank
print( loom.gather_to_host( sums ) )                    # [ 0+4+8+12, 1+5+9+13, 2+6+10+14, 3+7+11+15 ]

For now

• 2D, with a domain ( boundaries ), and measures only: no gradient yet, and the other arguments of PowerDiagram ( a distribution, a periodicity... ) are refused.
• The number of seeds must split evenly between the ranks, and the ranks should be a power of two.
• The boxes are still made on the host, from the whole cloud ( gathered from the processes when they each hold a block ).
• At most 52 ranks or boxes.

The steps that remain are on the roadmap.