Mesh.slice_points: avoid full-mesh index allocations for sparse cell blocks - #1966
Mesh.slice_points: avoid full-mesh index allocations for sparse cell blocks#1966peterdsharpe wants to merge 7 commits into
Conversation
Mesh.slice_points builds two n_points-sized index tensors and scatter-gathers the kept rows; on a 142M-vertex memmap-backed surface that is ~2.3 GB of temporaries per 10k-cell sample. The reader already knows every surviving cell is fully referenced, so compact directly: torch.unique(..., return_inverse=True) remaps the cells and _gather_rows reads the referenced rows of points and each point_data leaf (one sequential range read when the rows span <= 256 MiB, else a plain gather). Output is identical to the old path; tests assert equality.
…n ranges slice_points remapped cells through a full-mesh lookup table (an arange and an old-to-new map, each n_points long) and fancy-indexed the kept rows out of points and point_data. On a memory-mapped surface with ~142M vertices that is ~2.3 GB of temporaries and one page fault per row for a slice of a few thousand points, in every dataloader worker. Sort the kept ids once and locate each cell vertex by binary search (a vertex is kept iff the search lands on an equal id; right-biased so a duplicated id maps to its last position, as the lookup table did). Gather kept rows with a helper that reads one contiguous range when the rows are close together and falls back to a plain gather otherwise. Boolean masks and slices normalize without an n_points-sized arange. Results are unchanged; tests pin every accepted index form (unsorted, duplicates, negatives, empty, slice, list, int, bool mask) and a memmap-loaded mesh to the lookup-table algorithm.
peterdsharpe
changed the title
Synthetic sweeps showed the binary search alone is 2-10x slower than the lookup table when slicing a full mesh (every cell vertex is searched with a log k factor), while it wins 5-19x for the reader path (a small cell block out of a 1e8-vertex mesh). Keep the lookup table when n_points is at most 64x the number of cell-vertex entries and search otherwise; the equivalence tests now run on both algorithms.
Measurements:
|
| block cells | metric | main: lookup table | this PR: table or search, by shape | speed-up |
|---|---|---|---|---|
| 10,000 | cold sample (load + block + compaction) | 2.22 s (1.94 s–8.36 s) | 2.60 s (1.94 s–62.97 s) | 0.9x |
| 10,000 | second block, same case (partly cached) | 935 ms (680 ms–5.67 s) | 976 ms (621 ms–25.14 s) | 1.0x |
| 10,000 | slice_points alone, same block (point pages just read) | 48 ms (38 ms–6.03 s) | 2 ms (2 ms–2.01 s) | 27.9x |
| 10,000 | recipe reader __getitem__ |
469 ms (314 ms–1.34 s) | 341 ms (266 ms–996 ms) | 1.4x |
| 10,000 | peak RSS increase during the sample | 3,081 (2,826–3,186) MiB | 963 (702–1,044) MiB | 3.2x |
| 100,000 | cold sample (load + block + compaction) | 3.25 s (2.20 s–8.78 s) | 2.92 s (2.24 s–6.43 s) | 1.1x |
| 100,000 | second block, same case (partly cached) | 95 ms (75 ms–6.82 s) | 46 ms (36 ms–10.60 s) | 2.1x |
| 100,000 | slice_points alone, same block (point pages just read) | 52 ms (44 ms–2.31 s) | 10 ms (9 ms–9.62 s) | 5.1x |
| 100,000 | recipe reader __getitem__ |
147 ms (133 ms–2.41 s) | 119 ms (94 ms–5.49 s) | 1.2x |
| 100,000 | peak RSS increase during the sample | 3,622 (3,510–3,942) MiB | 1,476 (1,314–1,584) MiB | 2.5x |
What this says:
- Memory is the headline. The per-sample peak RSS increase drops from ~3.1 GB to ~0.96 GB at 10k cells and from ~3.6 GB to ~1.5 GB at 100k cells. The ~2.1 GB saved is the two
n_points-long int64 tensors (2 × 142M × 8 B) that every dataloader worker used to allocate per sample. slice_pointsitself goes from 48 ms to 2 ms on a 10k-cell block (28x at the median) and from 52 ms to 10 ms at 100k cells (5x).- Cold wall time is unchanged within noise (medians 2.2–2.6 s / 2.9–3.3 s, fastest cases identical at 1.94 s). On Lustre the cold sample is dominated by reading 30k–300k scattered vertex rows out of a 1.7 GB point file, which no in-process algorithm changes. The outliers in both implementations (63 s, 25 s, 9 s, 8 s) are storage stalls on a loaded file system that hit whichever one was running; they also account for the min–max spread in the
slice_pointsrows, whose point gather touches the same memory-mapped pages. - The recipe reader's end-to-end
__getitem__on a partly cached case improves modestly (469 → 341 ms, 147 → 119 ms).
2. Synthetic: the reader's path as a function of mesh size
slice_cells(block) then slice_points(unique(cells)) on an N-vertex triangle mesh with 2N random cells, one scalar and one vector point field, saved once to a node-local memmap ("memory" rows clone everything into RAM first). The block is small; what matters is how compaction scales with the size of the mesh it came from. Median of 3 runs.
| N vertices | block cells | kept vertices | backing | main: lookup table | this PR: table or search, by shape | speed-up |
|---|---|---|---|---|---|---|
| 1,000,000 | 10,000 | 29,539 | memmap | 3 ms | 3 ms | 1.1x |
| 1,000,000 | 10,000 | 29,539 | memory | 2 ms | 2 ms | 1.0x |
| 1,000,000 | 100,000 | 259,232 | memmap | 3 ms | 3 ms | 1.0x |
| 1,000,000 | 100,000 | 259,232 | memory | 2 ms | 2 ms | 0.9x |
| 10,000,000 | 10,000 | 29,955 | memmap | 8 ms | 5 ms | 1.6x |
| 10,000,000 | 10,000 | 29,955 | memory | 4 ms | 2 ms | 1.8x |
| 10,000,000 | 100,000 | 295,678 | memmap | 8 ms | 6 ms | 1.3x |
| 10,000,000 | 100,000 | 295,678 | memory | 7 ms | 5 ms | 1.3x |
| 100,000,000 | 10,000 | 29,997 | memmap | 62 ms | 13 ms | 4.8x |
| 100,000,000 | 10,000 | 29,997 | memory | 49 ms | 2 ms | 20.8x |
| 100,000,000 | 100,000 | 299,571 | memmap | 65 ms | 17 ms | 3.8x |
| 100,000,000 | 100,000 | 299,571 | memory | 53 ms | 8 ms | 6.5x |
- With the lookup table, the cost of compacting a fixed-size block grows linearly with the size of the source mesh (2 ms → 49 ms from 1M to 100M vertices for a 10k-cell block in memory). With this PR it is flat at ~2 ms: the work depends on the block, not on the mesh.
- The 1M-vertex rows are on the lookup-table side of the threshold (33 points per cell-vertex entry) and match
main. The memmap rows at 100M vertices keep ~13 ms because the kept rows are scattered through the file and have to be read individually.
3. Synthetic: slice_points on the full mesh (keeping k random vertices)
Same meshes, now slicing the whole mesh to k random vertices, which is the usual use of slice_points (mesh transforms, cropping). Cell connectivity is 6N entries, so every row here takes the lookup-table branch of this PR. Median of 3 runs.
| N vertices | k kept | backing | main: lookup table | this PR: table or search, by shape | speed-up |
|---|---|---|---|---|---|
| 1,000,000 | 1,000 | memmap | 6 ms | 6 ms | 1.1x |
| 1,000,000 | 1,000 | memory | 15 ms | 5 ms | 3.0x |
| 1,000,000 | 30,000 | memmap | 6 ms | 6 ms | 1.0x |
| 1,000,000 | 30,000 | memory | 15 ms | 6 ms | 2.6x |
| 10,000,000 | 1,000 | memmap | 78 ms | 69 ms | 1.1x |
| 10,000,000 | 1,000 | memory | 76 ms | 66 ms | 1.1x |
| 10,000,000 | 30,000 | memmap | 80 ms | 73 ms | 1.1x |
| 10,000,000 | 30,000 | memory | 73 ms | 77 ms | 0.9x |
| 10,000,000 | 1,000,000 | memmap | 83 ms | 82 ms | 1.0x |
| 10,000,000 | 1,000,000 | memory | 78 ms | 81 ms | 1.0x |
| 100,000,000 | 1,000 | memmap | 1.01 s | 929 ms | 1.1x |
| 100,000,000 | 1,000 | memory | 961 ms | 923 ms | 1.0x |
| 100,000,000 | 30,000 | memmap | 995 ms | 947 ms | 1.1x |
| 100,000,000 | 30,000 | memory | 992 ms | 944 ms | 1.1x |
| 100,000,000 | 1,000,000 | memmap | 1.02 s | 992 ms | 1.0x |
| 100,000,000 | 1,000,000 | memory | 988 ms | 995 ms | 1.0x |
This PR matches main on full-mesh slices (0.9x–1.1x at 10M and 100M vertices; the 1M-vertex rows are all under 20 ms and within run-to-run noise): the common use of slice_points is unaffected, and the gains above come only where the mesh is far larger than the cells being kept.
Method and raw data
- HiLift:
bench_hilift_slice.pyon one GB300 node of an internal Slurm cluster (CPU work only), withPYTHONPATHswitched between a snapshot ofmainand this branch; both use the unchanged upstream reader (_subsample_mesh_cells→slice_cells→slice_points). Two dozen cases taken at a fixed stride through the dataset (12 per round, the second round on a disjoint dozen with the package order swapped), block seed 0 for the cold sample, seed 1 for the second block. Peak RSS fromru_maxrss. Raw:hilift_aga.jsonl. - Synthetic:
bench_slice_points.pyandbench_slice_points_block.py, one subprocess per (implementation, case), on the same node. The lookup-table algorithm frommainis reproduced verbatim inside the script; the other column runs the branch'sMesh.slice_points. Raw:synthetic_full_mesh.json,synthetic_reader_path.json. - Tables and figures:
analyze_bench1966.py.
  The PR appears safe to merge from a correctness perspective, with a non-blocking performance concern because the new range-read helper does not deliver its claimed sequential memory-mapped read. Findings
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
…her; upstream PR NVIDIA#1966 (Mesh.slice_points) supersedes it Reverts the effect of f955fbb, 11ab332, 4ce9a7d, a2e9970, 1b2ac66. The compaction fix now lives in Mesh.slice_points (PR NVIDIA#1966, merged from origin/pr/mesh-reader-cell-compaction at 1e605e0), so _subsample_mesh_cells is back to slice_cells + slice_points as on main. The page-wise positional read path (PHYSICSNEMO_MESH_PREAD_GATHER) is dropped, not re-hooked: on AGA Lustre it was 4-5x slower inside the training loop (320-361 s vs 70-130 s per epoch on the same node; book/18-notebook.qmd#sec-nb-hilift-io-verdict).
… workaround retired
|
/ok to test 1aaf5fd |
peterdsharpe
changed the title
|
/ok to test 1d5ad82 |
Signed-off-by: Peter Sharpe <peterdsharpe@gmail.com>
1d5ad82 to
8c0941f
Compare
|
/ok to test 8c0941f |
negin513
left a comment
negin513
left a comment
There was a problem hiding this comment.
Reviewed the two remap paths and the index normalization carefully — the sort+searchsorted branch is correct (stable sort + right=True reproduces last-occurrence semantics; the clamp + equality check filters dropped points safely), the dense path is unchanged in complexity, and the tests genuinely force both algorithms. Nice.
One real regression in the empty-selection branch, one gap in the algorithm gate, and one comment that overstates the cost model — inline below.
Two broader notes. (1) The 28x / RSS numbers in the description were measured at 1e605e0, before the row-gather helper was removed; worth a rerun on head before they're cited. (2) Since #1894, the zarr reader path (_zarr_mesh_subsampled in datapipes/readers/mesh.py) compacts via torch.unique(return_inverse=True) and never calls slice_points, so this optimization only serves the .pmsh/.pdmsh memmap reader and external callers. Still worthwhile, but it argues for keeping the gate simple rather than tuning it further.
… for point clouds, table when most points are kept Review follow-ups on NVIDIA#1966: the empty-selection branch returned the input connectivity dtype (int32 stayed int32) while the other branches always gave int64; a mesh without cells now skips the remap entirely; and the algorithm gate also considers how many points are kept, since the search's sort scales with the kept set, not with the connectivity. Comment on index normalization no longer claims boolean masks are free. Signed-off-by: Peter Sharpe <peterdsharpe@gmail.com>
|
^ fixes pushed up, benchmarks re-running |
|
/ok to test f75eeea |
2 participants
Description
Mesh.slice_pointsno longer allocates point-count-sized index tensors when remapping a small cell block from a much larger mesh. It retains the lookup-table algorithm for ordinary meshes and switches to sorting the kept point IDs plus binary search whenn_points > 64 * cells.numel()and fewer thann_points / 64points are kept. Meshes without cells skip the remap, and remapped connectivity is always int64 whatever the input dtype. This targets the reader's compaction of large meshes after cell sampling.Index normalization allocates only the selected range/IDs. Empty forward slices remain empty; int32/int64 indices, negative indices, duplicates, unsorted selections, boolean masks and legacy uint8 masks retain their selection semantics. Invalid tensor dtypes, malformed masks, multidimensional selections and out-of-range indices are rejected. Negative-step slices remain unsupported, matching PyTorch indexing.
Points and nested point data use ordinary indexed gathers, including when loaded from a memory-mapped mesh. The change does not claim sequential range reads. Duplicate point IDs retain the previous last-position cell remapping behavior.
Validation
test_slice_points_lookup_equivalence.pyandtest_slicing.py, with both remapping algorithms forced in the equivalence/regression suite.