Brook TEASAR skeletonization on NVIDIA GPUs

Brook turns a 3D label volume, such as a neuron segmentation in connectomics, into one skeleton per object. Connected components, distance transforms, path tracing and, when memory allows, skeleton assembly run on the NVIDIA GPU. It takes the arguments of kimimaro.skeletonize and returns osteoid.Skeleton objects.

Get startedHow it works
TEASAR on CPU cores and on the GPU with Brook, and a look inside Brook’s path loopTop: TEASAR on CPU cores traces one path of one object at a time. On the NVIDIA GPU, Brook traces the next path of many objects at once, in a loop that runs as one CUDA graph. Bottom: each path grows in rings from its target until it touches the skeleton, and on a long object several paths are drafted in parallel, checked, and traced again only if rejected.time (schematic, not to scale)TEASAR on CPU coreseach core traces one object, one path at a timevolumecomponentsdistance transformcore 1…core 2…core 3…⋮Brook on the NVIDIA GPUone call; the path loop runs as one CUDA graphhostskeletonsno round trip per pathGPUcomponentsdistance transformone floodsecond lane: longest objectspath loop: each stepdraws the next pathof many objectsinside the path loopEach path: a ring searchfirst contact:the search stopsrings spread fromthe target, each ring’svoxels in paralleltarget: farthest remaining voxelDrafting ahead on a long objectDraft in parallelseveral next paths at once???Checkquick tests on what the draftalready computed — target stillthe farthest? no earlier pathcloser? — no new searchKeep or trace againaccepted: same asone-at-a-time tracingrejected: traced again123
TEASAR on CPU cores and on the GPU with Brook, and a look inside Brook’s path loopTop: TEASAR on CPU cores traces one path of one object at a time. On the NVIDIA GPU, Brook traces the next path of many objects at once, in a loop that runs as one CUDA graph. Bottom: on a long object several paths are drafted in parallel, checked, and traced again only if rejected.time (schematic, not to scale)TEASAR on CPU coreseach core traces one object, one path at a timewhole volumecomponentsdistancetransformcore 1…core 2…core 3…⋮Brook on the NVIDIA GPUone call; the path loop is one CUDA graphhostskeletonsno round trip per pathGPUcomponentsdistancetransformonefloodsecond lane: longest objectspath loopinside the path loopDrafting ahead on a long objectDraft in parallelseveral next paths at onceCheckquick tests on what the draft alreadycomputed — target still the farthest?no earlier path closer? — no new search???Keep or trace againaccepted: same as one-at-a-time tracingrejected: traced again123
Schematic, not to scale.

Get started

A drop-in replacement for Kimimaro’s skeletonize, also usable from C and C++. Labels can already be on the GPU: CuPy and PyTorch arrays are read without a trip through host memory.

Linux x86-64; NVIDIA GPU of compute capability 8.0+ (Ampere or later); NVIDIA driver supporting CUDA 12.3+; wheels for CPython 3.12 to 3.14. There are no macOS wheels, and Windows is untested.

Install

python -m pip install brook-cu12

# or from source (CUDA Toolkit 12.3+)
pip install build
python -m build --wheel
pip install dist/brook-*.whl

Python

import brook

skeletons = brook.skeletonize(
    labels, anisotropy=(16, 16, 40))

# many volumes in one call
batch = brook.skeletonize_batch(volumes)

C++

#include <brook/brook.hpp>

brook_volume v{};  // host or device labels
v.struct_size = sizeof(v);
v.abi_version = BROOK_ABI_VERSION;
v.data = labels;   v.dtype = BROOK_U32;
v.memory = BROOK_DEVICE;
// v.shape, v.strides in elements / bytes

brook::api::Context context;
auto skeletons = context.skeletonize(v);

How TEASAR becomes more parallel

Rings instead of a queue

A new path searches back to the skeleton one ring of voxels at a time, in parallel, and stops at first contact.

Searching back to the skeleton: one voxel per step against one ring per stepEach new path starts at a target and searches for the nearest part of the skeleton already drawn. On the CPU a priority queue settles one voxel per step. Brook settles a whole ring of voxels per step in parallel and stops at the first ring that touches the skeleton.TEASAR on the CPUa queue settles one voxel per step12345678…Brooka whole ring of voxels per stepstep 1step 2step 3targetskeleton so farfirst contact: the search stops

One wave, many objects

Distance fields for many objects come from one flood over the volume. The wave never crosses from one object into another.

Distance fields: one object at a time against one flood for many objectsOn the CPU each object's distance field is computed by its own flood, one object after another. Brook runs a single flood over the whole volume; a wave may only move between voxels of the same object, so each object still gets exactly its own field.TEASAR on the CPUone object’s flood at a timenownextlaterBrookone flood over the whole volumeWaves stop at object borders, so each object gets its own field.

Drafting ahead

Like speculative decoding: Brook drafts the next paths of a long object in parallel and keeps only those one-at-a-time tracing would have drawn.

Drafting ahead on a long object, like speculative decodingOn the CPU a long object is traced one path after another. Brook drafts the next paths in parallel inside small private windows, then keeps a draft only if one-at-a-time tracing would have produced the same target and route, the draft stayed inside its window and it does not overlap an earlier draft. A draft that fails is dropped and traced again, so the skeleton is the one that tracing one path at a time gives.TEASAR on the CPUone path, then the nextpath 1path 2path 3path 4Brookdraft ahead, then check3 drafts in parallelcheckkeep 2, redo 1A draft is kept only if it has the same target and routethat one-at-a-time tracing would give, stays inside itswindow and does not overlap an earlier draft.

Two lanes

The longest objects move to a second loop that runs alongside the first, so they stop setting the pace.

Two concurrent loops: a main lane and a long laneThe main lane advances many objects by one path per step; its steps carry a lot of work. The objects with the longest remaining chains move to a second loop that runs at the same time with short, frequent steps and drafts ahead, so they finish sooner.Brook runs two loops side by sideMain lanemany objects, one path each per stepLong lanethe longest objects, drafting aheadtimeShort, frequent steps for the longest chains.

Batches in one call

skeletonize_batch traces several volumes together, with the same result as separate calls.

25 volumes: 4.60 s as separate calls, 1.30 s as one batchMedian of three runs each; each batch result matches its single call bit for bit (25/25).25 separate calls4.60 sOne batch call1.30 s
25 volumes of 1 × 1972 × 2024 voxels on an NVIDIA RTX 4090, median of three runs each.

Twelve datasets

On these twelve datasets Brook returns the same objects as Kimimaro. Brook’s total skeleton length is 0.5% to 7.9% shorter, mostly because equal-cost routes are broken differently.

Time per volume on twelve datasets, Brook on one NVIDIA RTX 4090 against Kimimaro on 8 CPU coresKimimaro benchmark: Brook 3.08 s, Kimimaro 413.4 s, 134.3× faster; Hemibrain: Brook 2.69 s, Kimimaro 331.1 s, 123.0× faster; MICrONS: Brook 1.56 s, Kimimaro 103.9 s, 66.8× faster; FIB-25: Brook 2.38 s, Kimimaro 131.6 s, 55.3× faster; H01: Brook 1.19 s, Kimimaro 54.0 s, 45.3× faster; Scroll fibres 16 µm: Brook 0.63 s, Kimimaro 26.7 s, 42.1× faster; Kasthuri11: Brook 1.06 s, Kimimaro 37.2 s, 35.0× faster; CREMI A: Brook 1.05 s, Kimimaro 35.9 s, 34.0× faster; FAFB FFN1: Brook 1.18 s, Kimimaro 37.0 s, 31.2× faster; CREMI B: Brook 1.03 s, Kimimaro 23.0 s, 22.5× faster; CREMI C: Brook 1.08 s, Kimimaro 22.6 s, 20.9× faster; Scroll fibres 8 µm: Brook 0.57 s, Kimimaro 5.6 s, 9.9× faster.BrookKimimaro, 8 coresKimimaro benchmark134.3×Hemibrain123.0×MICrONS66.8×FIB-2555.3×H0145.3×Scroll fibres 16 µm42.1×Kasthuri1135.0×CREMI A34.0×FAFB FFN131.2×CREMI B22.5×CREMI C20.9×Scroll fibres 8 µm9.9×1 s10 s100 s1,000 s
Time per volume on twelve datasets, Brook on one NVIDIA RTX 4090 against Kimimaro on 8 CPU coresKimimaro benchmark: Brook 3.08 s, Kimimaro 413.4 s, 134.3× faster; Hemibrain: Brook 2.69 s, Kimimaro 331.1 s, 123.0× faster; MICrONS: Brook 1.56 s, Kimimaro 103.9 s, 66.8× faster; FIB-25: Brook 2.38 s, Kimimaro 131.6 s, 55.3× faster; H01: Brook 1.19 s, Kimimaro 54.0 s, 45.3× faster; Scroll fibres 16 µm: Brook 0.63 s, Kimimaro 26.7 s, 42.1× faster; Kasthuri11: Brook 1.06 s, Kimimaro 37.2 s, 35.0× faster; CREMI A: Brook 1.05 s, Kimimaro 35.9 s, 34.0× faster; FAFB FFN1: Brook 1.18 s, Kimimaro 37.0 s, 31.2× faster; CREMI B: Brook 1.03 s, Kimimaro 23.0 s, 22.5× faster; CREMI C: Brook 1.08 s, Kimimaro 22.6 s, 20.9× faster; Scroll fibres 8 µm: Brook 0.57 s, Kimimaro 5.6 s, 9.9× faster.BrookKimimaro, 8 coresKimimaro benchmark134.3×Hemibrain123.0×MICrONS66.8×FIB-2555.3×H0145.3×Scroll fibres 16 µm42.1×Kasthuri1135.0×CREMI A34.0×FAFB FFN131.2×CREMI B22.5×CREMI C20.9×Scroll fibres 8 µm9.9×1 s10 s100 s1,000 s
Seconds per volume, log scale. Brook 0.1.0 on an NVIDIA RTX 4090, median of three runs after one warmup; Kimimaro 5.8.1 with 8 workers on 8 cores of an Intel Core i9-14900KF, one run.