PipeComm

Pipeline-aware collective communication synthesis framework. This is the open-source release of the PipeComm paper, with the original Gurobi-based ILP replaced by HiGHS (via scipy.optimize.milp).

Quick Start

python3 example.py

Python API (pipecomm.py)

Communication Primitives

from pipecomm import *

buf = Buffer(1 << 20)  # 1 MiB buffer
S = create_schedule()

with tile(buf, 3) as t:          # split buffer into 3 chunks
    with pipeline(1):            # pipeline interval II=1
        S.reduce_at(-1, t[0])           # reduce, auto-select root
        S.reduce_at(-1, t[1])
        S.broadcast_from(-1, t[2])

# Feasibility (default)
result = S.synthesis(topo_name="mesh2d_3x3")

# Minimize max depth
result = S.synthesis(topo_name="mesh2d_3x3", minimize_depth=True)

for t in result.trees.trees:
    print(t.direction, t.root, t.edges)

Construction Primitives

t1 = Tree(0, "reduce", [(1, 0), (2, 0)])
t2 = Tree(0, "broadcast", [(0, 1), (0, 2)])
trees = Trees([t1, t2], II=1)

rev   = reverse(trees)        # swap reduce <-> broadcast, flip edges
both  = interleave(trees, rev) # alternate tree insertion

# Hierarchical construction: extend existing patterns to new nodes
ext   = extend(trees, subset=[0, 1, 2, 3], new_ranks=[4, 5, 6, 7, 8],
               topo_name="mesh2d_3x3")

Hierarchical Construction (extend)

extend(trees, subset, new_ranks, topo_name) implements the hierarchical pattern expansion described in the paper. Given existing trees covering a sub-topology (subset), it extends them to new nodes (new_ranks) by treating the sub-topology as a super-rank and running the ILP on a reduced topology.

Full topology (16 nodes)          Reduced topology (S + 7 new)
┌───┬───┬───┬───┐                ┌───┬───┬───┬───┐
│ 0 │ 1 │ 2 │ 9 │                │   │   │   │ 9 │
├───┼───┼───┼───┤    collapse    ├───┼───┼───┼───┤
│ 3 │ 4 │ 5 │10 │    ──────►     │   │   │   │10 │
├───┼───┼───┼───┤   subset→S     │ S │   │   ├───┤
│ 6 │ 7 │ 8 │11 │                │   │   │   │11 │
└───┴───┴───┴───┘                └───┴───┴───┴───┘
  Existing patterns               ILP finds optimal
  cover nodes 0-8                 S ↔ new_node edges

# First synthesize on sub-topology
result = S.synthesis(topo_name="mesh2d_3x3")  # 9 nodes

# Extend patterns to 4×4 mesh (16 nodes)
extended = extend(result.trees,
                  subset=list(range(9)),        # existing nodes
                  new_ranks=list(range(9, 16)), # new column
                  topo_name="mesh2d_4x4")       # full topology

The reduced ILP has 1 + len(new_ranks) nodes, making it tractable even when the full topology is large. This decomposes the synthesis task into smaller sub-problems as described in the paper.

Classes

Class	Description
Buffer(size)	Data buffer with byte size
Tree(root, direction, edges)	Single communication tree
Trees(trees, II)	Collection of trees with pipeline interval
Schedule()	Builder: reduce_at, broadcast_from, synthesis
SynthesisResult	Result (.trees, .schedule, .II)

Schedule.synthesis() Options

Parameter	Default	Description
topo_name	(required)	Topology file name without .txt
data_size	1 << 20	Data size in bytes
minimize_depth	False	If True, minimize max depth instead of feasibility

Context Managers

Manager	Usage
tile(buffer, factor)	with tile(buf, 3) as t: ... t[0], t[1], t[2]
pipeline(II)	with pipeline(1): ... sets II on current schedule

Solver

The paper's original implementation uses Gurobi for the MILP formulation. This open-source release replaces Gurobi with HiGHS via scipy.optimize.milp. The MILP finds optimal edge-disjoint spanning trees under the initiation interval (II) constraint, then schedule.cpp performs Modulo-II reservation table scheduling.

File Layout

PipeComm/
  pipecomm.py              # Python frontend (HiGHS ILP + Modulo-II scheduling)
  schedule.cpp             # Modulo-II scheduler (compiles to ./schedule)
  example.py               # Usage examples
  topology/                # Generated topology files (*.txt)