refactor: remove `coremltools` dependency and use `torch.utils.cpp_extension` to load kmeans1d by guru-desh · Pull Request #31 · apple/coreai-optimization

guru-desh · 2026-07-09T20:40:50Z

Description

This PR removes coremltools as a required dependency for coreai-opt.

The only part of coremltools that was being used was kmeans1d. We have copied kmeans1d in coreai-opt using torch.utils.cpp_extension to compile the C++ code upon invocation.

Using torch.utils.cpp_extension allows us to have a more faithful refactor for kmeans1d. From here, we can then choose to explore alternatives like numba or optimizations for the kmeans1d C++ code. This PR is scoped to have the minimal amount of changes needed to remove coremltools as a dependency.

Thanks @pkmandke for the discussion and suggestions!

Testing

We have equivalence tests to ensure that this copy of kmeans1d is equivalent to coremltools kmeans1d (tolerance is 1e-9)
kmeans1d tests are copied over from coremltools
We also have a hash check to check that the C++ file copied over is the same as the one that exists on coremltools. This means that we have a pure refactor of kmeans1d

guru-desh · 2026-07-09T21:09:39Z

+_DLL = None
+
+
+def _dll():


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 · 2026-07-10T09:08:08Z

+# 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"] + (


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

Drop "(copied byte-for-byte from coremltools)" -- attribution already lives in the vendored files' headers and NOTICE.txt, and the phrase will go stale once the core is further optimized (per aseemw's PR apple#31 review).

aseemw's PR apple#31 review: kmeans1d correctness tests importing coremltools as a live oracle defeats the purpose of removing it as a required dependency, and the "must stay byte-identical to upstream" canary asserts an invariant that isn't actually required (the vendored code can be modified/adapted going forward). Collapses TestEquivalence/TestKmeans1DBehavior/TestsKmeans1dFromCoremltools into one TestKmeans1D class and deletes TestVendoredSourceMatchesUpstream entirely. Small discrete-input cases (k>n clamping, fewer-than-k-cluster collapse, explicit duplicates) now assert hardcoded literal values directly, captured by running the vendored cluster() once (already validated against coremltools in CI, so a known-good golden value). Larger/random-array cases (unweighted/weighted equivalence grids, large-n up to 100k, fp16 duplicates, weighted-integer-counts) would require thousands of literal values to hardcode inline, so their expected coremltools-oracle output is instead checked into tests/palettization/assets/kmeans1d/*.npz (generated once via the new generate_assets.py, not collected by pytest) and loaded at test time -- inputs are regenerated from the same fixed RNG seeds rather than also stored on disk, keeping the assets small (~300KB total). Net effect: this test file has zero coremltools imports (confirmed via AST inspection), so coreml stops being an implicit test dependency for kmeans1d correctness.

aseemw's PR apple#31 review: kmeans1d correctness tests importing coremltools as a live oracle defeats the purpose of removing it as a required dependency, and the "must stay byte-identical to upstream" canary asserts an invariant that isn't actually required (the vendored code can be modified/adapted going forward). Collapses TestEquivalence/TestKmeans1DBehavior/TestsKmeans1dFromCoremltools into one TestKmeans1D class and deletes TestVendoredSourceMatchesUpstream entirely. Small discrete-input cases (k>n clamping, fewer-than-k-cluster collapse, explicit duplicates) now assert hardcoded literal values directly, captured by running the vendored cluster() once (already validated against coremltools in CI, so a known-good golden value). Larger/random-array cases (unweighted/weighted equivalence grids, large-n up to 100k, fp16 duplicates, weighted-integer-counts) would require thousands of literal values to hardcode inline, so their expected coremltools-oracle output is instead checked into tests/palettization/assets/kmeans1d/*.npz and loaded at test time -- inputs are regenerated from the same fixed RNG seeds rather than also stored on disk, keeping the assets small (~300KB total). Net effect: this test file has zero coremltools imports (confirmed via AST inspection), so coreml stops being an implicit test dependency for kmeans1d correctness.

aseemw's PR apple#31 review: kmeans1d correctness tests importing coremltools as a live oracle defeats the purpose of removing it as a required dependency, and the "must stay byte-identical to upstream" canary asserts an invariant that isn't actually required (the vendored code can be modified/adapted going forward). Collapses TestEquivalence/TestKmeans1DBehavior/TestsKmeans1dFromCoremltools into one TestKmeans1D class and deletes TestVendoredSourceMatchesUpstream entirely. Small discrete-input cases (k>n clamping, fewer-than-k-cluster collapse, explicit duplicates) now assert hardcoded literal values directly, captured by running the vendored cluster() once (already validated against coremltools in CI, so a known-good golden value). Larger/random-array cases (unweighted/weighted equivalence grids, large-n up to 100k, fp16 duplicates, weighted-integer-counts) would require thousands of literal values to hardcode inline, so their inputs and expected coremltools-oracle output are instead checked into tests/palettization/assets/kmeans1d/*.npz and loaded at test time (~3.3MB total, mostly the large-n case). Net effect: this test file has zero coremltools imports (confirmed via AST inspection), so coreml stops being an implicit test dependency for kmeans1d correctness.

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.

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).

torch.utils.cpp_extension.load() imports setuptools internally to JIT-compile the vendored kmeans1d core. This was previously pulled in transitively via coremltools, which is now an optional dependency, so it must be declared explicitly.

Isolates the vendored third-party kmeans1d source (and its MIT LICENSE) into a dedicated deps/ folder, separate from coreai-opt's own internal _utils modules, per Apple legal's guidance on incorporating vendored dependencies (mirrors how coremltools itself vendored the same code).

Completes the previous move: repoints pyproject.toml package-data and the three call sites at coreai_opt.deps._kmeans1d instead of the old coreai_opt._utils._kmeans1d location. Also anchors the .gitignore deps/ rule to the repo root (/deps/) so it stops shadowing the new src/coreai_opt/deps/ package.

Bring the per-file headers in line with Apple legal's guidance for incorporating vendored third-party code (mirrors coremltools' own kmeans1d headers): every file in the vendored package now carries the upstream MIT license text followed by a trailing Apple copyright line, with no coreai-opt repo-wide header layered on top since this isn't coreai-opt's own code. Bumps the trailing line from "Copyright (c) 2023 Apple Inc." (when coremltools first vendored it) to "Copyright (c) 2026 Apple Inc." (when coreai-opt vendored it from coremltools). Also excludes src/coreai_opt/deps/ from the add-license-header pre-commit hook, which otherwise unconditionally stamps coreai-opt's own BSD-3 header onto any .py file lacking one.

Attributes the vendored kmeans1d source under src/coreai_opt/deps/ to its upstream MIT license, per Apple legal's guidance for incorporating vendored dependencies (mirrors the equivalent NOTICE.txt commit in coremltools, which vendored the same code).

The vendored-source-matches-upstream canary hashes our _core.cpp against coremltools' live upstream copy. Now that coreai-opt's copy carries its own vendoring-year copyright line ("Copyright (c) 2026 Apple Inc.") rather than coremltools' ("Copyright (c) 2023 Apple Inc."), the raw byte hashes never match. Strip that one known, intentional line from both sides before hashing so the canary still catches real drift in the kmeans1d logic without permanently failing on the copyright year.

The unanchored `deps/` rule (from the initial commit, alongside other personal-scratch entries like scratch/ and /local_stash) silently shadowed the newly-added, tracked src/coreai_opt/deps/ package: search tools like ripgrep skip gitignored paths by default, and any future new file added under deps/ would need an explicit force-add. Nothing in the repo (Makefile, CI, scripts, docs) reads or writes a deps/ folder, so remove the rule outright rather than just anchoring it.

The vendored kmeans1d is JIT-compiled via torch.utils.cpp_extension, which needs a C++ toolchain on the host at runtime; previously this requirement rode in transitively via the coremltools wheel, so it's worth calling out now that coremltools is optional.

Drop "(copied byte-for-byte from coremltools)" -- attribution already lives in the vendored files' headers and NOTICE.txt, and the phrase will go stale once the core is further optimized (per aseemw's PR apple#31 review).

aseemw's review: comments explaining implementation choices by diffing against what coremltools' build does are unnecessary now that this is a fresh adaptation, not a diff against coremltools. Drop the comment entirely rather than rephrasing it to be self-contained; the flags themselves are unchanged.

Re-verified via explain-coreai-opt (aseemw asked why docs needs the coreml group): Sphinx's autodoc/autosummary machinery does genuinely import every public coreai_opt module at build time, but nothing under src/coreai_opt/ imports coremltools anywhere -- the only remaining kmeans1d/coremltools attribution now lives in comments and the vendored package's own LICENSE/header, none of which autodoc touches. No CI job builds docs either, so this has no build-time effect to verify against.

aseemw's PR apple#31 review: kmeans1d correctness tests importing coremltools as a live oracle defeats the purpose of removing it as a required dependency, and the "must stay byte-identical to upstream" canary asserts an invariant that isn't actually required (the vendored code can be modified/adapted going forward). Collapses TestEquivalence/TestKmeans1DBehavior/TestsKmeans1dFromCoremltools into one TestKmeans1D class and deletes TestVendoredSourceMatchesUpstream entirely. Small discrete-input cases (k>n clamping, fewer-than-k-cluster collapse, explicit duplicates) now assert hardcoded literal values directly, captured by running the vendored cluster() once (already validated against coremltools in CI, so a known-good golden value). Larger/random-array cases (unweighted/weighted equivalence grids, large-n up to 100k, fp16 duplicates, weighted-integer-counts) would require thousands of literal values to hardcode inline, so their inputs and expected coremltools-oracle output are instead checked into tests/palettization/assets/kmeans1d/*.npz and loaded at test time (~3.3MB total, mostly the large-n case). Net effect: this test file has zero coremltools imports (confirmed via AST inspection), so coreml stops being an implicit test dependency for kmeans1d correctness.

guru-desh force-pushed the u/gurudesh/remove-coremltools-dependency-torch-cpp branch from cca4b22 to 9910570 Compare July 9, 2026 21:03

guru-desh commented Jul 9, 2026

View reviewed changes

Comment thread src/coreai_opt/deps/_kmeans1d/core.py

guru-desh force-pushed the u/gurudesh/remove-coremltools-dependency-torch-cpp branch 3 times, most recently from 8c4d108 to a17f5c0 Compare July 10, 2026 00:05

guru-desh commented Jul 10, 2026

View reviewed changes

guru-desh force-pushed the u/gurudesh/remove-coremltools-dependency-torch-cpp branch from e858d3a to a64cf0c Compare July 10, 2026 09:08

guru-desh requested review from aseemw, dengqiaoyu and pkmandke July 10, 2026 22:56

guru-desh assigned aseemw Jul 10, 2026

guru-desh added the enhancement New feature or request label Jul 10, 2026

pkmandke reviewed Jul 11, 2026

View reviewed changes

Comment thread src/coreai_opt/deps/_kmeans1d/core.py

guru-desh force-pushed the u/gurudesh/remove-coremltools-dependency-torch-cpp branch from 1d5cdaf to a14c10c Compare July 13, 2026 00:18

guru-desh marked this pull request as ready for review July 13, 2026 20:05

pkmandke reviewed Jul 13, 2026

View reviewed changes

Comment thread changelog.d/31.changed Outdated

guru-desh force-pushed the u/gurudesh/remove-coremltools-dependency-torch-cpp branch from 74244d2 to de2bd23 Compare July 14, 2026 06:09

pkmandke approved these changes Jul 14, 2026

View reviewed changes

aseemw requested changes Jul 14, 2026

View reviewed changes

Comment thread changelog.d/31.changed Outdated

Comment thread src/coreai_opt/deps/_kmeans1d/core.py Outdated

Comment thread tests/palettization/test_kmeans1d.py Outdated

Comment thread pyproject.toml Outdated

Comment thread pyproject.toml

aseemw reviewed Jul 14, 2026

View reviewed changes

Comment thread tests/palettization/test_kmeans1d.py Outdated

aseemw reviewed Jul 14, 2026

View reviewed changes

Comment thread tests/palettization/test_kmeans1d.py Outdated

guru-desh force-pushed the u/gurudesh/remove-coremltools-dependency-torch-cpp branch from 39c0fb8 to 9541614 Compare July 14, 2026 22:49

guru-desh force-pushed the u/gurudesh/remove-coremltools-dependency-torch-cpp branch from 9541614 to 416db85 Compare July 14, 2026 22:54

aseemw approved these changes Jul 14, 2026

View reviewed changes

guru-desh added 18 commits July 15, 2026 11:24

refactor: switch palettization kmeans to vendored kmeans1d

fd3acc1

Point _efficient_kmeans and kmeans_fake_palettize at coreai_opt._utils._kmeans1d instead of coremltools._deps._kmeans1d. The cluster()/Clustered call contract is unchanged.

refactor: make coreai model palettization use vendored kmeans1d

4bb7bd5

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.

docs: add changelog fragment for coremltools optional dependency

d886a90

docs: trim changelog wording per review feedback

91787e8

Drop "(copied byte-for-byte from coremltools)" -- attribution already lives in the vendored files' headers and NOTICE.txt, and the phrase will go stale once the core is further optimized (per aseemw's PR apple#31 review).

guru-desh force-pushed the u/gurudesh/remove-coremltools-dependency-torch-cpp branch from 416db85 to cd2e650 Compare July 15, 2026 18:25

guru-desh merged commit edd4720 into apple:main Jul 15, 2026
8 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

refactor: remove `coremltools` dependency and use `torch.utils.cpp_extension` to load kmeans1d#31

refactor: remove `coremltools` dependency and use `torch.utils.cpp_extension` to load kmeans1d#31
guru-desh merged 19 commits into
apple:mainfrom
guru-desh:u/gurudesh/remove-coremltools-dependency-torch-cpp

guru-desh commented Jul 9, 2026 •

edited

Loading

Uh oh!

guru-desh Jul 9, 2026

Uh oh!

Uh oh!

guru-desh Jul 10, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		_DLL = None


		def _dll():

Uh oh!

Conversation

guru-desh commented Jul 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Testing

Uh oh!

guru-desh Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

guru-desh Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

guru-desh commented Jul 9, 2026 •

edited

Loading