refactor: remove coremltools dependency and use torch.utils.cpp_extension to load kmeans1d#31
Draft
guru-desh wants to merge 9 commits into
Draft
Conversation
guru-desh
force-pushed
the
u/gurudesh/remove-coremltools-dependency-torch-cpp
branch
from
cca4b22 to
9910570
Compare
guru-desh
commented
Jul 9, 2026
| _DLL = None | ||
|
|
||
|
|
||
| def _dll(): |
Contributor
Author
There was a problem hiding this comment.
This file is mainly the same as coremltools._deps.kmeans1d except for this function that uses torch.utils.cpp_extension to compile _core.cpp
guru-desh
commented
Jul 9, 2026
| # be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause | ||
|
|
||
|
|
||
| # MIT License |
Contributor
Author
There was a problem hiding this comment.
TODO: Is this licensing correct?
guru-desh
force-pushed
the
u/gurudesh/remove-coremltools-dependency-torch-cpp
branch
3 times, most recently
from
8c4d108 to
a17f5c0
Compare
guru-desh
commented
Jul 10, 2026
| # NOT inherit sysconfig flags and defaults to -O0, which made this | ||
| # tight-inner-loop, template-heavy DP code 10-15x slower per call than | ||
| # coremltools' prebuilt extension. Pass the same flags explicitly to match. | ||
| _EXTRA_CFLAGS = ["-std=c++11", "-O2", "-DNDEBUG"] + ( |
Contributor
Author
There was a problem hiding this comment.
Future PRs can look into using -O3 compilation flag and benchmark performance, but since this PR is meant as a pure refactor, we stick with -O2
Vendor coremltools' kmeans1d C++ core (_core.cpp, byte-for-byte from apple/coremltools; upstream kmeans1d 0.3.1, MIT) and JIT-compile it at runtime with torch.utils.cpp_extension, invoking its extern "C" entry points via ctypes. Provides cluster()/Clustered as a drop-in replacement for coremltools._deps._kmeans1d, with fail-loud input validation.
Point _efficient_kmeans and kmeans_fake_palettize at coreai_opt._utils._kmeans1d instead of coremltools._deps._kmeans1d. The cluster()/Clustered call contract is unchanged.
Drop the optional third-party kmeans1d import and _HAS_KMEANS1D flag in palettize_utils; import coreai_opt._utils._kmeans1d unconditionally. The sklearn fallback for vector k-means is unchanged.
Compare cluster() output against the coremltools oracle (unweighted, weighted, large-n, collapse, fp16 ties), assert fixed known-good cases and input validation, and add a canary that hard-fails if the vendored _core.cpp drifts from apple/coremltools main.
coremltools is no longer needed at runtime; move it (and its numpy<2.4 cap) into a coreml optional-dependency + self-referencing dependency-group, kept in default-groups so the test oracle is available. Ship the vendored _core.cpp and LICENSE as package data.
Palettization spawns many worker processes that each call _kmeans1d.cluster(). Since the vendored extension JIT-compiles on first use per-process, every worker independently raced the same on-disk build lock: one worker compiled while the rest spin-waited for it instead of just loading the finished .so, serializing a one-time multi-minute compile across the critical path. Add _kmeans1d.warmup() and call it once in the parent process before each worker pool is created (KMeansPalettizer and the coreai-torch palettize_utils compress pool), so every worker's load() call finds the extension already built.
…orkers" This reverts commit 42f8a69.
The vendored kmeans1d C++ extension was 10-15x slower per call than coremltools' precompiled version (measured directly: ~0.9s vs ~0.06s on a realistic-size weighted cluster() call), even with the actual JIT-compile lock/cache fully warmed. This was the real cause of the 1.2x-2.6x palettization wall-clock regression seen in an end-to-end LLM eval, not JIT-compile lock contention across parallel workers (a prior "warmup()" fix targeting that, since reverted, made no measurable difference). Root cause: torch.utils.cpp_extension.load() does not inherit CPython's sysconfig OPT/CFLAGS the way distutils/setuptools does, so it defaults to -O0. coremltools' setup.py sets no explicit -O flag either, but distutils prepends sysconfig's "-O2 -DNDEBUG" (the standard python.org/conda default) ahead of its extra_compile_args, so its shipped .so is effectively an -O2 build. This tight-inner-loop, template-heavy DP algorithm is extremely sensitive to that gap. Pass the same flags explicitly so torch's JIT build matches coremltools' effective compile line. Numeric output is unaffected (verified against the existing coremltools-oracle equivalence tests).
guru-desh
force-pushed
the
u/gurudesh/remove-coremltools-dependency-torch-cpp
branch
from
e858d3a to
a64cf0c
Compare
pkmandke
reviewed
Jul 11, 2026
| _DLL = None | ||
|
|
||
|
|
||
| def _dll(): |
Member
There was a problem hiding this comment.
Were you able to identify if/why this re-compiles each time during parallel kmeans?
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This PR removes
coremltoolsas a required dependency forcoreai-opt.The only part of
coremltoolsthat was being used waskmeans1d. We have copiedkmeans1dincoreai-optusingtorch.utils.cpp_extensionto compile the C++ code upon invocation.Using
torch.utils.cpp_extensionallows us to have a more faithful refactor forkmeans1d. From here, we can then choose to explore alternatives likenumbaor optimizations for thekmeans1dC++ code. This PR is scoped to have the minimal amount of changes needed to removecoremltoolsas a dependency.Thanks @pkmandke for the discussion and suggestions!
Testing
kmeans1dis equivalent to coremltoolskmeans1d(tolerance is 1e-9)coremltoolscoremltools. This means that we have a pure refactor ofkmeans1d