Metal (Apple Silicon) solver backend via MLX#23
Conversation
New pathfinder/metal_dijkstra.py: GPU SSSP on Apple Silicon's unified memory - the whole graph lives in one address space, so there is no host/device transfer layer. CSR is converted once to a dense padded neighbor table (K = max degree, small on a Manhattan lattice) and paths come from frontier-driven scatter-min relaxation with a parent- resolution phase (MLX GPU scatter lacks int64, so no packed keys). Opt-in via ORTHO_BACKEND=metal or config.use_metal; try_attach_metal grafts onto SimpleDijkstra with the CPU fallback retained - Metal can never fail a net the CPU could route. Touches unified_pathfinder.py in exactly one 5-line block to minimize conflicts with ongoing work. Verified on M4 Pro: 30/30 path-cost parity vs CPU Dijkstra (full-graph, ROI-restricted, ownership-penalty), full-suite green, and end-to-end benchmark parity (2x16: 16/16 overuse 0 iters 7; 4x40x6L: 80/80 overuse 0 iters 8). Honest status: SLOWER than CPU at synthetic scale (25.7s vs 15.8s at 80 nets) - the per-wave frontier sync dominates; the target regime is full-graph solves on monster boards where waves are wide. Optimization headroom: on-GPU frontier compaction, batched rounds, custom Metal kernel via mx.fast.metal_kernel.
Replace scatter-min waves (host sync per round) with dense incoming- neighbor gather relaxation: dist'[v] = min over arriving edges, parent from argmin, everything GPU-resident with a fixpoint check every 8 rounds - the Metal translation of the CUDA backend's resident-state design. Small solves stay on CPU via a MIN_ROI_NODES threshold (same tradeoff as CUDA's gpu_roi_min_nodes): dense rounds cost O(N*K) regardless of wave width and only pay off at scale. Parity suite forces the threshold to 0 to keep exercising Metal on tiny graphs.
Replace the dense O(N*K)-per-round relaxation with a custom Metal kernel (mx.fast.metal_kernel): threads relax only the frontier's outgoing edges, atomically maxing INVERTED uint32-encoded distances (== min on real distances, and a single init_value=0 then serves all three atomic outputs: distances, next frontier, frontier count). Parents are resolved race-free in ONE dense in-table pass after convergence - same float op order as the kernel, so distance equality is bitwise-safe. One host sync per round (the frontier count). Measured on the real ThinkinMachine backplane (M4 Pro 48GB, 8.1M nodes / 22.8M edges, graph builds in 19s at 2.0GB RSS, Metal tables +1.3s / +1.7GB): longest net (764 hops): CPU 5.3s -> Metal 1.4s (3.9x), paths identical median net (462 hops): CPU 5.5s -> Metal 1.1s (5.2x), paths identical The dense path this replaces measured 9x SLOWER than CPU (51s) - the wavefront is ~50k nodes wide on an 8.1M node graph, so dense rounds did ~160x wasted work. Small solves still route to CPU via MIN_ROI_NODES. Parity suite: 4/4; full suite: 133 passed.
|
Status update with real numbers. The Metal solver went through two architectures: Dense gather relaxation (first push): perfect parity, but 9x SLOWER than CPU on the real backplane - O(N*K) work per round while the wavefront is only ~50k nodes wide on an 8.1M-node graph. Frontier scatter kernel (current, via mx.fast.metal_kernel with atomic_fetch_min on uint32-encoded distances): measured on the actual ThinkinMachine backplane (9,481 nets, 32 layers, 8.1M nodes / 22.8M edges) on an M4 Pro:
Graph builds in 19s at 2.0GB RSS; Metal tables add 1.3s / 1.7GB. The board that once needed a rented 40GB A100 now runs entirely in a laptop's unified memory with GPU solves in ~1 second. Small ROIs still route to CPU via a size threshold (dense GPU rounds only pay off at scale - same tradeoff as the CUDA gpu_roi_min_nodes). Next: wire Metal into the full negotiation loop and attempt a complete monster route on-device. |
- Kernel loop now stops when every target's settled distance beats the frontier's best (uint32 encodings are monotonic, costs >= 0) instead of relaxing the whole graph to fixpoint. Under CPU/GPU contention the median monster solve holds 1.1s (7.0x vs contended CPU 8.2s). - _route_all: cache the full-graph identity roi/global_to_roi arrays - previously two fresh 32MB np.arange(8.1M) allocations per long net. Parity suite + engine smoke: 53 passed.
|
Banking this PR for battlestation integration day. Final state: Done and measured (M4 Pro, real ThinkinMachine backplane, 8.1M nodes / 22.8M edges):
Known remaining bottleneck (NOT Metal's): per-net Python BFS ROI extraction (~1-2s/net, extracts 75k nodes then truncates to 50k) dominates routing wall-clock for all backends. Flagging for whoever picks it up - a geometric/bbox corridor ROI on the regular lattice is pure index arithmetic. Integration checklist: rebase onto current main, re-run the parity suite, re-measure the solve table, then decide backend selection policy (CUDA/Metal/CPU) in one seam. |
1 participant
Adds an opt-in Metal GPU solver for Apple Silicon, exploiting unified memory: the entire routing graph sits in one address space shared by CPU and GPU, so the host/device transfer layer that dominates the CUDA backend simply doesn't exist here.
What
pathfinder/metal_dijkstra.py: MLX-based SSSP. CSR converted once to a dense padded neighbor table (max degree is a small constant on the Manhattan lattice); frontier-driven scatter-min relaxation with a separate parent-resolution phase (MLX GPU scatter has no int64, so no packed dist|parent keys). Supports ROI restriction and ownership node penalties with the same semantics as SimpleDijkstra.try_attach_metal()grafts onto SimpleDijkstra keeping the CPU fallback: Metal can never fail a net the CPU could route.ORTHO_BACKEND=metalorconfig.use_metal. One 5-line touch to unified_pathfinder.py; everything else is new files.tests/test_metal_solver.py: parity suite, skips cleanly off-Mac.Verified (M4 Pro, 48GB)
Honest status
Currently SLOWER than CPU on synthetic boards (25.7s vs 15.8s at 80 nets): each relaxation wave round-trips the frontier through Python. The target regime is full-graph solves on huge boards where waves are wide and the graph would not fit conventional VRAM. Next steps: on-GPU frontier compaction, batched rounds, custom kernel via mx.fast.metal_kernel, then a full ThinkinMachine backplane (9,481 nets / 32 layers / ~9M nodes - parses in 4.9s, graph fits trivially in unified memory) attempt.
Kept as a PR rather than pushed to main so it can be validated without disturbing ongoing engine work.