Skip to content

fix: respect CPU affinity and reduce clustering parallel overhead - #132

Merged
gouyt13 merged 2 commits into
mainfrom
fix/affinity-aware-clustering-parallelism
Oct 7, 2026
Merged

gouyt13 merged 2 commits into
mainfrom
fix/affinity-aware-clustering-parallelism

Conversation

@gouyt13

@gouyt13 gouyt13 commented Oct 7, 2026

Copy link
Copy Markdown
Member

Summary

Default thread selection used the whole-machine CPU count even when the process was restricted to one NUMA node. On a dual-socket host with 96 logical CPUs, launching on node 0 could therefore still request 96 workers for its 48 available logical CPUs. Short clustering stages also paid unnecessary thread-team and synchronization overhead.

  • Resolve automatic and oversized thread requests against available CPUs, using Linux affinity or OpenMP placement information, and use the same detection in the Python clustering example. Respect OpenMP placement in IVF initialization when binding is enabled.
  • Keep the centroid-graph worker limit local to QGKMeans: min(thread_budget, max(1, k / 128)). Training-vector assignment retains its available-CPU-capped budget.
  • Remove PiPNN's redundant single-thread outer OpenMP region at the root so its tile workers can reuse the normal team. Fuse shared Lloyd centroid means, optional normalization, and shift calculation while preserving empty-centroid refill handling.
  • Add affinity/placement, exception propagation, centroid-shift, and PiPNN regression coverage; clarify threading documentation.

Public signatures, saved-index formats, and dependencies are unchanged. The library does not change the process-wide OpenMP wait policy.

Testing

  • Portable Release C++ build (RABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF): 299 passed, 1 skipped, 0 failures across 300 registered tests. The binding-only invocation skips without binding; the dedicated OpenMPBinding process passes.
  • Rebuilt and installed the wheel in the existing project environment; verified imported extension path and SHA256 against the wheel. Full Python suite: 727 passed, 2 skipped (Linux ARM64/macOS repaired-wheel checks). C++ example tests were required and passed.
  • Full clang-tidy for configured library/binding translation units, clang-format, Ruff, strict MkDocs build, documentation links, and git diff --check passed.
  • Installed CMake package built and linked successfully from tests/consumer; its test passed.
  • Local verification was on Linux x86-64; other platforms were not run locally.

QGKMeans performance

Full datasets, 4,096 centroids, L2, seed 42, 25 iterations plus Exact final assignment; three trials per version/configuration with alternating execution order. Baseline: 035f4857a2f116892d6c66f2fadd8f0bb42098de. Both versions used GCC 14.3.0, Release, native optimization disabled, runtime SIMD dispatch on an AVX-512-capable host, and the same libgomp runtime (16.2.0). Individual kernel selection was not traced.

Hardware: two Xeon Gold 6418H sockets, each with 24 physical cores / 48 logical CPUs. Launch binding used numactl --cpunodebind=0 or --cpunodebind=0,1; no memory binding/interleave. OpenMP/GOMP wait and binding overrides were cleared; BLAS thread counts were 1.

Dataset CPU binding / threads Before median After median Time reduction
msong, 994,185 x 420 node 0 / 48 10.52 s 9.11 s 13.4%
msong, 994,185 x 420 both sockets / 96 10.80 s 6.79 s 37.1%
GIST, 1,000,000 x 960 node 0 / 48 17.54 s 16.64 s 5.1%
GIST, 1,000,000 x 960 both sockets / 96 15.26 s 13.11 s 14.1%

Times cover training including Exact, excluding file loading/saving. All 12 outputs for each dataset had byte-identical centroid and label arrays. Dual-socket timings varied substantially: msong before 9.12-10.96 s / after 6.71-9.68 s; GIST before 14.12-20.62 s / after 11.88-18.95 s. These measurements establish improvement for the tested QGKMeans configurations, not every library operation or workload; the k / 128 limit is a heuristic.

@gouyt13
gouyt13 merged commit a065785 into main Oct 7, 2026
37 checks passed
@gouyt13
gouyt13 deleted the fix/affinity-aware-clustering-parallelism branch October 7, 2026 18:36
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant