diff --git a/.gitignore b/.gitignore index 7798f7e..6f7ba39 100644 --- a/.gitignore +++ b/.gitignore @@ -47,7 +47,6 @@ uv.lock *.nc .python-version /.tox -deps/ scratch/ /local_stash /model_weights diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index d1c94c6..39d2410 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -23,6 +23,7 @@ repos: language: system entry: scripts/pre_commit/add_license_header.py --license-file configs/BSD-3-LICENSE-HEADER-TEMPLATE --start-year 2026 files: '\.(py|sh|js|css|html)$|(^|/)(GNUmakefile|[Mm]akefile)$' + exclude: ^src/coreai_opt/deps/ # ---------------------------------------------------------------------------- # 0.2 Update LICENSE year diff --git a/NOTICE.txt b/NOTICE.txt new file mode 100644 index 0000000..0abc072 --- /dev/null +++ b/NOTICE.txt @@ -0,0 +1,25 @@ +Copyright 2026 Apple Inc. + +This project contains content adapted from kmeans1d (https://github.com/dstein64/kmeans1d), the license for which follows: + +MIT License + +Copyright (c) 2019 Daniel Steinberg + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/changelog.d/31.changed b/changelog.d/31.changed new file mode 100644 index 0000000..43305b9 --- /dev/null +++ b/changelog.d/31.changed @@ -0,0 +1 @@ +Replace the coremltools-based 1D k-means used by palettization with a vendored C++ core that is JIT-compiled at runtime via `torch.utils.cpp_extension`. `coremltools` is no longer a runtime dependency (it is now an optional dependency, installable via the `coreml` extra). This requires a C++ compiler to be available on the host at runtime. diff --git a/pyproject.toml b/pyproject.toml index ec29c6d..374a210 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,11 +31,16 @@ classifiers = [ ] dynamic = [ "version" ] dependencies = [ - "coremltools>=8.3", - "numpy>=2,<2.4", # TODO: Remove once coremltools has released a version > 9.0 + # Required at runtime by torch.utils.cpp_extension to JIT-compile the vendored + # kmeans1d C++ core (a C++ toolchain must also be present on the host). + "ninja>=1.11", + "numpy>=2", "pydantic>=2.0.0", + "pyyaml>=6.0", "rich>=13.0.0", "safetensors>=0.5.3,<=0.7.0", + # Required at runtime by torch.utils.cpp_extension + "setuptools>=42", # PyTorch >= 2.9.0 requires torchao >= 0.15.0 # Python's standard dependency specification (PEP 508) doesn't support conditional dependencies # based on other package versions. We can either 1) add a stricter check to require the newer torchao @@ -54,6 +59,10 @@ coreai = [ "coreai-torch==0.4.1", "scikit-learn>=1.7.2", ] +coreml = [ + "coremltools>=8.3", + "numpy>=2,<2.4", # TODO: Remove once coremltools has released a version > 9.0 +] [project.urls] Changelog = "https://github.com/apple/coreai-optimization/blob/main/CHANGELOG.md" Repository = "https://github.com/apple/coreai-optimization" @@ -95,7 +104,9 @@ docs = [ # installs CoreAI by default (see tool.uv.default-groups). The extra is the # single source of truth; this group self-references it to avoid duplicating the # package list. +# This applies to the `coreml` group as well. coreai = [ "coreai-opt[coreai]" ] +coreml = [ "coreai-opt[coreml]" ] # Used in CI to force latest mimimum supported torch version # These torch versions must be in bounds of torch versions listed in project dependencies highest_tested_torch = [ @@ -145,6 +156,7 @@ tutorial = [ "papermill>=2.7.0", "torchinfo>=1.8.0", { include-group = "coreai" }, + { include-group = "coreml" }, { include-group = "torchvision" }, ] @@ -153,14 +165,14 @@ package-dir = { "" = "src" } [tool.setuptools.dynamic] version = { attr = "coreai_opt._about.__version__" } [tool.setuptools.package-data] -coreai_opt = [ "py.typed" ] +coreai_opt = [ "deps/_kmeans1d/_core.cpp", "deps/_kmeans1d/LICENSE", "py.typed" ] [tool.setuptools.packages] find.exclude = [ "coreai_opt_benchmarking*", "tests*" ] find.where = [ "src" ] # make env group installed by default with uv sync [tool.uv] -default-groups = [ "dev", "coreai" ] +default-groups = [ "dev", "coreai", "coreml" ] index = [ { explicit = true, name = "pytorch-cpu", url = "https://download.pytorch.org/whl/cpu" }, { explicit = true, name = "pytorch-cu128", url = "https://download.pytorch.org/whl/cu128" }, diff --git a/src/coreai_opt/coreai_utils/_utils/palettize_utils.py b/src/coreai_opt/coreai_utils/_utils/palettize_utils.py index 49ab2a1..a535ccf 100644 --- a/src/coreai_opt/coreai_utils/_utils/palettize_utils.py +++ b/src/coreai_opt/coreai_utils/_utils/palettize_utils.py @@ -24,19 +24,12 @@ ) from coreai_opt.coreai_utils._utils.graph_utils import _select_input_output_channel_axis from coreai_opt.coreai_utils.common import CompressionGranularity as _CompressionGranularity +from coreai_opt.deps import _kmeans1d logger = logging.getLogger(__name__) _CONV2D_OP = "coreai.conv2d" -try: - import kmeans1d as _kmeans1d - - _HAS_KMEANS1D = True -except ImportError: - _kmeans1d = None # type: ignore[assignment] - _HAS_KMEANS1D = False - LutParams = namedtuple("LutParams", "indices lut vector_axis") _SUPPORTED_NBITS: tuple[int, ...] = (1, 2, 3, 4, 6, 8) @@ -99,8 +92,7 @@ def _get_kmeans_lookup_table_and_weight( weight.shape[1] == 1 and num_weights >= 10_000 and weight.dtype == np.float16 ) - if (is_better_to_use_kmeans1d and _HAS_KMEANS1D) or force_kmeans1d: - assert _HAS_KMEANS1D, "Unable to import kmeans1d, please make sure it's installed." + if is_better_to_use_kmeans1d or force_kmeans1d: values, indices, counts = np.unique(weight, return_inverse=True, return_counts=True) indices = indices.flatten() n_clusters = min(len(values), lut_len) diff --git a/src/coreai_opt/deps/_kmeans1d/LICENSE b/src/coreai_opt/deps/_kmeans1d/LICENSE new file mode 100644 index 0000000..fef4b8e --- /dev/null +++ b/src/coreai_opt/deps/_kmeans1d/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2019 Daniel Steinberg + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/src/coreai_opt/deps/_kmeans1d/__init__.py b/src/coreai_opt/deps/_kmeans1d/__init__.py new file mode 100644 index 0000000..4d81dad --- /dev/null +++ b/src/coreai_opt/deps/_kmeans1d/__init__.py @@ -0,0 +1,27 @@ +# MIT License +# +# Copyright (c) 2019 Daniel Steinberg +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +# Copyright © 2026 Apple Inc. + +from coreai_opt.deps._kmeans1d.core import Clustered, cluster + +__all__ = ["Clustered", "cluster"] diff --git a/src/coreai_opt/deps/_kmeans1d/_core.cpp b/src/coreai_opt/deps/_kmeans1d/_core.cpp new file mode 100644 index 0000000..aa0b433 --- /dev/null +++ b/src/coreai_opt/deps/_kmeans1d/_core.cpp @@ -0,0 +1,372 @@ +// MIT License +// +// Copyright (c) 2019 Daniel Steinberg +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// +// Copyright © 2026 Apple Inc. + +#include + +#include +#include +#include +#include +#include +#include + +using namespace std; + +typedef unsigned long ulong; + +/* + * Internal implementation of the SMAWK algorithm. + */ +template +void _smawk( + const vector& rows, + const vector& cols, + const function& lookup, + vector* result) { + // Recursion base case + if (rows.size() == 0) return; + + // ******************************** + // * REDUCE + // ******************************** + + vector _cols; // Stack of surviving columns + for (ulong col : cols) { + while (true) { + if (_cols.size() == 0) break; + ulong row = rows[_cols.size() - 1]; + if (lookup(row, col) >= lookup(row, _cols.back())) + break; + _cols.pop_back(); + } + if (_cols.size() < rows.size()) + _cols.push_back(col); + } + + // Call recursively on odd-indexed rows + vector odd_rows; + for (ulong i = 1; i < rows.size(); i += 2) { + odd_rows.push_back(rows[i]); + } + _smawk(odd_rows, _cols, lookup, result); + + unordered_map col_idx_lookup; + for (ulong idx = 0; idx < _cols.size(); ++idx) { + col_idx_lookup[_cols[idx]] = idx; + } + + // ******************************** + // * INTERPOLATE + // ******************************** + + // Fill-in even-indexed rows + ulong start = 0; + for (ulong r = 0; r < rows.size(); r += 2) { + ulong row = rows[r]; + ulong stop = _cols.size() - 1; + if (r < rows.size() - 1) + stop = col_idx_lookup[(*result)[rows[r + 1]]]; + ulong argmin = _cols[start]; + T min = lookup(row, argmin); + for (ulong c = start + 1; c <= stop; ++c) { + T value = lookup(row, _cols[c]); + if (c == start || value < min) { + argmin = _cols[c]; + min = value; + } + } + (*result)[row] = argmin; + start = stop; + } +} + +/* + * Interface for the SMAWK algorithm, for finding the minimum value in each row + * of an implicitly-defined totally monotone matrix. + */ +template +vector smawk( + const ulong num_rows, + const ulong num_cols, + const function& lookup) { + vector result; + result.resize(num_rows); + vector rows(num_rows); + iota(begin(rows), end(rows), 0); + vector cols(num_cols); + iota(begin(cols), end(cols), 0); + _smawk(rows, cols, lookup, &result); + return result; +} + +/* + * Calculates cluster costs in O(1) using prefix sum arrays. + */ +class CostCalculator { + vector cumsum; + vector cumsum2; + + public: + CostCalculator(const vector& vec, ulong n, const vector& /*sort_idxs*/) { + cumsum.push_back(0.0); + cumsum2.push_back(0.0); + for (ulong i = 0; i < n; ++i) { + double x = vec[i]; + cumsum.push_back(x + cumsum[i]); + cumsum2.push_back(x * x + cumsum2[i]); + } + } + + double weight(ulong i, ulong j) { + return (i <= j) ? 1 + j - i : 0; + } + + double calc(ulong i, ulong j) { + if (j < i) return 0.0; + double mu = (cumsum[j + 1] - cumsum[i]) / (j - i + 1); + double result = cumsum2[j + 1] - cumsum2[i]; + result += (j - i + 1) * (mu * mu); + result -= (2 * mu) * (cumsum[j + 1] - cumsum[i]); + return result; + } +}; + +/* + * Weighted version of the CostCalculator + */ +class WeightedCostCalculator { + vector cumw; + vector cumsum; + vector cumsum2; + + public: + WeightedCostCalculator( + const vector& vec, + ulong n, + const vector& sort_idxs, + const double* unsorted_weights) { + vector sorted_weights(n); + for (ulong i = 0; i < n; ++i) { + sorted_weights[i] = unsorted_weights[sort_idxs[i]]; + } + cumw.push_back(0.0); + cumsum.push_back(0.0); + cumsum2.push_back(0.0); + for (ulong i = 0; i < n; ++i) { + double x = vec[i]; + double w = sorted_weights[i]; + cumw.push_back(w + cumw[i]); + cumsum.push_back(w * x + cumsum[i]); + cumsum2.push_back(w * x * x + cumsum2[i]); + } + } + + double weight(ulong i, ulong j) { + return (i <= j) ? cumw[j + 1] - cumw[i] : 0.0; + } + + double calc(ulong i, ulong j) { + if (j < i) return 0.0; + double w = weight(i, j); + double mu = (cumsum[j + 1] - cumsum[i]) / w; + double result = cumsum2[j + 1] - cumsum2[i]; + result += w * (mu * mu); + result -= (2 * mu) * (cumsum[j + 1] - cumsum[i]); + return result; + } +}; + +template +class Matrix { + vector data; + ulong num_rows; + ulong num_cols; + + public: + Matrix(ulong num_rows, ulong num_cols) { + this->num_rows = num_rows; + this->num_cols = num_cols; + data.resize(num_rows * num_cols); + } + + inline T get(ulong i, ulong j) { + return data[i * num_cols + j]; + } + + inline void set(ulong i, ulong j, T value) { + data[i * num_cols + j] = value; + } +}; + +template +void cluster_impl( + const double* array, + ulong n, + ulong k, + ulong* clusters, + double* centroids, + CostArgsTypes... args) { + // *************************************************** + // * Sort input array and save info for de-sorting + // *************************************************** + + vector sort_idxs(n); + iota(sort_idxs.begin(), sort_idxs.end(), 0); + sort( + sort_idxs.begin(), + sort_idxs.end(), + [&array](ulong a, ulong b) {return array[a] < array[b];}); + vector undo_sort_lookup(n); + vector sorted_array(n); + for (ulong i = 0; i < n; ++i) { + sorted_array[i] = array[sort_idxs[i]]; + undo_sort_lookup[sort_idxs[i]] = i; + } + + // *************************************************** + // * Set D and T using dynamic programming algorithm + // *************************************************** + + // Algorithm as presented in section 2.2 of (Gronlund et al., 2017). + + CostCalculatorType cost_calculator(sorted_array, n, sort_idxs, args...); + Matrix D(k, n); + Matrix T(k, n); + + for (ulong i = 0; i < n; ++i) { + D.set(0, i, cost_calculator.calc(0, i)); + T.set(0, i, 0); + } + + for (ulong k_ = 1; k_ < k; ++k_) { + auto C = [&D, &k_, &cost_calculator](ulong i, ulong j) -> double { + ulong col = i < j - 1 ? i : j - 1; + return D.get(k_ - 1, col) + cost_calculator.calc(j, i); + }; + vector row_argmins = smawk(n, n, C); + for (ulong i = 0; i < row_argmins.size(); ++i) { + ulong argmin = row_argmins[i]; + double min = C(i, argmin); + D.set(k_, i, min); + T.set(k_, i, argmin); + } + } + + // *************************************************** + // * Extract cluster assignments by backtracking + // *************************************************** + + // TODO: This step requires O(kn) memory usage due to saving the entire + // T matrix. However, it can be modified so that the memory usage is O(n). + // D and T would not need to be retained in full (D already doesn't need + // to be fully retained, although it currently is). + // Details are in section 3 of (Grønlund et al., 2017). + + vector sorted_clusters(n); + + ulong t = n; + ulong k_ = k - 1; + ulong n_ = n - 1; + // The do/while loop was used in place of: + // for (k_ = k - 1; k_ >= 0; --k_) + // to avoid wraparound of an unsigned type. + do { + ulong t_ = t; + t = T.get(k_, n_); + double centroid = 0.0; + for (ulong i = t; i < t_; ++i) { + sorted_clusters[i] = k_; + // Mean computation: this is only for squared L2 cost calculators + centroid += ( + (sorted_array[i] - centroid) + * cost_calculator.weight(i, i) + / cost_calculator.weight(t, i) + ); + } + centroids[k_] = centroid; + k_ -= 1; + n_ = t - 1; + } while (t > 0); + + // *************************************************** + // * Order cluster assignments to match de-sorted + // * ordering + // *************************************************** + + for (ulong i = 0; i < n; ++i) { + clusters[i] = sorted_clusters[undo_sort_lookup[i]]; + } +} + +extern "C" { +// "__declspec(dllexport)" causes the function to be exported when compiling on Windows. +// Otherwise, the function is not exported and the code raises +// "AttributeError: function 'cluster' not found". +// Exporting is a Windows platform requirement, not just a Visual Studio requirement +// (https://stackoverflow.com/a/22288874/1509433). The _WIN32 macro covers the Visual +// Studio compiler (MSVC) and MinGW. The __CYGWIN__ macro covers gcc and clang under +// Cygwin. +#if defined(_WIN32) || defined(__CYGWIN__) +__declspec(dllexport) +#endif +void cluster( + double* array, + ulong n, + ulong k, + ulong* clusters, + double* centroids) { + cluster_impl(array, n, k, clusters, centroids); +} + +#if defined(_WIN32) || defined(__CYGWIN__) +__declspec(dllexport) +#endif +void cluster_with_weights( + double* array, + double* weights, + ulong n, + ulong k, + ulong* clusters, + double* centroids) { + cluster_impl( + array, n, k, clusters, centroids, weights + ); +} +} // extern "C" + +static PyMethodDef module_methods[] = { + {NULL, NULL, 0, NULL} +}; + +static struct PyModuleDef _coremodule = { + PyModuleDef_HEAD_INIT, + "kmeans1d._core", + NULL, + -1, + module_methods, +}; + +PyMODINIT_FUNC PyInit__core(void) { + return PyModule_Create(&_coremodule); +} diff --git a/src/coreai_opt/deps/_kmeans1d/core.py b/src/coreai_opt/deps/_kmeans1d/core.py new file mode 100644 index 0000000..a0275be --- /dev/null +++ b/src/coreai_opt/deps/_kmeans1d/core.py @@ -0,0 +1,88 @@ +# MIT License +# +# Copyright (c) 2019 Daniel Steinberg +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +# Copyright © 2026 Apple Inc. + +import ctypes +import os +import sys +from collections import namedtuple +from collections.abc import Sequence + +from torch.utils.cpp_extension import load + +Clustered = namedtuple("Clustered", "clusters centroids") + +_EXTRA_CFLAGS = ["-std=c++11", "-O2", "-DNDEBUG"] + ( + ["-stdlib=libc++"] if sys.platform == "darwin" else [] +) +_EXTRA_LDFLAGS = ["-stdlib=libc++"] if sys.platform == "darwin" else [] + +_DLL = None + + +def _dll(): + global _DLL + if _DLL is None: + _core = load( + name="_core", + sources=[os.path.join(os.path.dirname(__file__), "_core.cpp")], + extra_cflags=_EXTRA_CFLAGS, + extra_ldflags=_EXTRA_LDFLAGS, + ) + _DLL = ctypes.cdll.LoadLibrary(_core.__file__) + return _DLL + + +def cluster(array: Sequence[float], k: int, *, weights: Sequence[float] | None = None) -> Clustered: + """ + :param array: A sequence of floats + :param k: Number of clusters (int) + :param weights: Sequence of weights (if provided, must have same length as `array`) + :return: A tuple with (clusters, centroids) + """ + assert k > 0, f"Invalid k: {k}" + n = len(array) + assert n > 0, f"Invalid len(array): {n}" + k = min(k, n) + + if weights is not None: + assert len(weights) == n, f"len(weights)={len(weights)} != len(array)={n}" + + c_array = (ctypes.c_double * n)(*array) + c_n = ctypes.c_ulong(n) + c_k = ctypes.c_ulong(k) + c_clusters = (ctypes.c_ulong * n)() + c_centroids = (ctypes.c_double * k)() + + if weights is None: + _dll().cluster(c_array, c_n, c_k, c_clusters, c_centroids) + else: + c_weights = (ctypes.c_double * n)(*weights) + _dll().cluster_with_weights(c_array, c_weights, c_n, c_k, c_clusters, c_centroids) + + clusters = list(c_clusters) + centroids = list(c_centroids) + + output = Clustered(clusters=clusters, centroids=centroids) + + return output diff --git a/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py b/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py index f953ce6..478f7aa 100644 --- a/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py +++ b/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py @@ -8,9 +8,9 @@ import numpy as np import torch -from coremltools._deps import _kmeans1d from coreai_opt.config.spec import CompressionTargetTensor +from coreai_opt.deps import _kmeans1d from coreai_opt.palettization.spec import ( PalettizationGranularity, PerGroupedChannelGranularity, @@ -202,9 +202,7 @@ def _calculate_centroids( # float32, matching the block-weight handling in _cluster_weights_1d. if sensitivities.dtype == torch.bfloat16: sensitivities = sensitivities.float() - sensitivities = self.reshape_strategy.reshape_for_kmeans( - sensitivities, axis - ) + sensitivities = self.reshape_strategy.reshape_for_kmeans(sensitivities, axis) block_sensitivities = self.granularity.get_blocks_to_cluster(sensitivities) else: block_sensitivities = [None] * len(block_weights_to_cluster) @@ -443,9 +441,7 @@ def _cluster_weights_1d( and (np.max(block_weight_flatten)) <= np.finfo(np.float16).max and np.min(block_weight_flatten) >= np.finfo(np.float16).min ): - values, indices, counts = self._reduce_weights_to_cluster( - block_weight_flatten - ) + values, indices, counts = self._reduce_weights_to_cluster(block_weight_flatten) num_clusters = min(len(values), num_clusters) if block_sensitivity_flatten is not None: counts = np.bincount(indices, weights=block_sensitivity_flatten) @@ -570,9 +566,7 @@ def _reduce_weights_to_cluster(self, block_weight_flatten: np.ndarray): # Rounding scale = 10**self.rounding_precision - block_weight_flatten = ( - np.round(block_weight_flatten.astype(np.float32) * scale) / scale - ) + block_weight_flatten = np.round(block_weight_flatten.astype(np.float32) * scale) / scale # To speed up parallel kmeans, use numpy.unique instead of # torch.unique in multiprocessing setting. @@ -610,9 +604,7 @@ def _scale_by_per_channel_scale(self, weight: torch.Tensor) -> torch.Tensor: Also scale the parameter using the computed scales. """ flattened_weight = weight.flatten(1) - per_channel_scale = torch.max( - torch.abs(flattened_weight), dim=1, keepdim=True - ).values + per_channel_scale = torch.max(torch.abs(flattened_weight), dim=1, keepdim=True).values # Handle zero scales per_channel_scale[per_channel_scale == 0] = 1 scaled_weight = flattened_weight / per_channel_scale @@ -622,9 +614,7 @@ def _scale_by_per_channel_scale(self, weight: torch.Tensor) -> torch.Tensor: return scaled_weight - def _unscale_by_per_channel_scale( - self, scaled_weight: torch.Tensor - ) -> torch.Tensor: + def _unscale_by_per_channel_scale(self, scaled_weight: torch.Tensor) -> torch.Tensor: """ Re-scale the parameter back to its original range by multiplying per channel scales. diff --git a/tests/palettization/assets/kmeans1d/fp16_duplicates.npz b/tests/palettization/assets/kmeans1d/fp16_duplicates.npz new file mode 100644 index 0000000..91268f0 Binary files /dev/null and b/tests/palettization/assets/kmeans1d/fp16_duplicates.npz differ diff --git a/tests/palettization/assets/kmeans1d/large_n.npz b/tests/palettization/assets/kmeans1d/large_n.npz new file mode 100644 index 0000000..992dd7c Binary files /dev/null and b/tests/palettization/assets/kmeans1d/large_n.npz differ diff --git a/tests/palettization/assets/kmeans1d/unweighted.npz b/tests/palettization/assets/kmeans1d/unweighted.npz new file mode 100644 index 0000000..779c0ae Binary files /dev/null and b/tests/palettization/assets/kmeans1d/unweighted.npz differ diff --git a/tests/palettization/assets/kmeans1d/weighted.npz b/tests/palettization/assets/kmeans1d/weighted.npz new file mode 100644 index 0000000..cd2ebb2 Binary files /dev/null and b/tests/palettization/assets/kmeans1d/weighted.npz differ diff --git a/tests/palettization/assets/kmeans1d/weighted_integer_counts.npz b/tests/palettization/assets/kmeans1d/weighted_integer_counts.npz new file mode 100644 index 0000000..d48ab87 Binary files /dev/null and b/tests/palettization/assets/kmeans1d/weighted_integer_counts.npz differ diff --git a/tests/palettization/test_kmeans1d.py b/tests/palettization/test_kmeans1d.py new file mode 100644 index 0000000..6f012fa --- /dev/null +++ b/tests/palettization/test_kmeans1d.py @@ -0,0 +1,228 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Fixed-value tests for the vendored 1D k-means, backed by hardcoded +literals for small inputs and checked-in golden ``.npz`` files (see +``assets/kmeans1d/``) for larger ones. Our extension JIT-compiles +``_core.cpp`` on first use and caches the artifact for the session, so the +compile cost is paid once. +""" + +from pathlib import Path + +import numpy as np +import pytest + +from coreai_opt.deps import _kmeans1d + +_RTOL = 1e-9 +_ATOL = 1e-9 + +_ASSETS_DIR = Path(__file__).parent / "assets" / "kmeans1d" + + +def _inertia(points, weights, clusters, centroids): + """Weighted within-cluster sum of squared deviations.""" + points = np.asarray(points, dtype=np.float64) + weights = np.ones_like(points) if weights is None else np.asarray(weights, dtype=np.float64) + centroids = np.asarray(centroids, dtype=np.float64) + clusters = np.asarray(clusters) + return float(np.sum(weights * (points - centroids[clusters]) ** 2)) + + +class TestKmeans1D: + """Fixed-value and behavior tests for the vendored cluster() surface.""" + + # -- Small, distinct-valued inputs: hardcode expected values directly -- + + def test_cluster(self): + array = [4.0, 4.1, 4.2, -50, 200.2, 200.4, 200.9, 80, 100, 102] + result = _kmeans1d.cluster(array, 4) + np.testing.assert_array_equal(result.clusters, [1, 1, 1, 0, 3, 3, 3, 2, 2, 2]) + np.testing.assert_allclose( + result.centroids, [-50.0, 4.1, 94.0, 200.5], rtol=_RTOL, atol=_ATOL + ) + + def test_cluster_with_weights(self): + array = [4.0, 4.1, 4.2, -50, 200.2, 200.4, 200.9, 80, 100, 102] + weights = [1, 1, 1, 0.125, 4, 1, 1, 3, 2, 2] + result = _kmeans1d.cluster(array, 4, weights=weights) + np.testing.assert_array_equal(result.clusters, [0, 0, 0, 0, 3, 3, 3, 1, 2, 2]) + np.testing.assert_allclose( + result.centroids, [1.936, 80.0, 101.0, 200.35], rtol=_RTOL, atol=_ATOL + ) + + def test_weights_vs_repetition(self): + values = [10, 24, 16, 12, 20] + weights = [3, 1, 4, 2, 3] + weighted = _kmeans1d.cluster(values, 2, weights=weights) + np.testing.assert_array_equal(weighted.clusters, [0, 1, 1, 0, 1]) + np.testing.assert_allclose(weighted.centroids, [10.8, 18.5], rtol=_RTOL, atol=_ATOL) + np.testing.assert_allclose( + _inertia(values, weights, weighted.clusters, weighted.centroids), + 66.8, + rtol=_RTOL, + atol=_ATOL, + ) + + # Weighting by integer counts equals repeating each value that many times. + repeated = np.repeat(values, weights) + repeated_result = _kmeans1d.cluster(repeated, 2) + np.testing.assert_allclose( + np.sort(repeated_result.centroids), + np.sort(weighted.centroids), + rtol=_RTOL, + atol=_ATOL, + ) + np.testing.assert_allclose( + _inertia(repeated, None, repeated_result.clusters, repeated_result.centroids), + 66.8, + rtol=_RTOL, + atol=_ATOL, + ) + + def test_k_greater_than_n_is_clamped(self): + # 5 distinct values, k=16 clamps to 5: each value is trivially its own cluster. + array = np.array([1.0, 5.0, 9.0, 13.0, 17.0]) + result = _kmeans1d.cluster(array, 16) + np.testing.assert_array_equal(result.clusters, [0, 1, 2, 3, 4]) + np.testing.assert_allclose(result.centroids, array, rtol=_RTOL, atol=_ATOL) + + @pytest.mark.parametrize( + "array, k, expected_centroids, expected_reconstruction", + [ + ([5.0, 5.0, 5.0], 3, [0.0, 0.0, 5.0], [5.0, 5.0, 5.0]), + ([1.0, 1.0, 9.0], 3, [0.0, 1.0, 9.0], [1.0, 1.0, 9.0]), + ([-5.0, -5.0, -5.0], 3, [0.0, 0.0, -5.0], [-5.0, -5.0, -5.0]), + ([2.0, 2.0, 2.0, 2.0], 4, [0.0, 0.0, 0.0, 2.0], [2.0, 2.0, 2.0, 2.0]), + ], + ids=["all-equal", "two-distinct", "all-negative-equal", "four-equal"], + ) + def test_collapse_fewer_than_k_clusters( + self, array, k, expected_centroids, expected_reconstruction + ): + # When k exceeds the number of distinct values, the optimum uses fewer than + # k clusters; unused leading centroids are zero-filled (including the + # non-ascending padding in "all-negative-equal" below). Tied values' exact + # labels aren't compared directly, since the C++ unstable sort makes their + # assignment non-reproducible; only the reconstructed value each point maps + # to is guaranteed stable. + result = _kmeans1d.cluster(np.array(array), k) + centroids = np.asarray(result.centroids) + clusters = np.asarray(result.clusters) + np.testing.assert_allclose(centroids, expected_centroids, rtol=_RTOL, atol=_ATOL) + np.testing.assert_allclose( + centroids[clusters], expected_reconstruction, rtol=_RTOL, atol=_ATOL + ) + + def test_explicit_duplicates_reconstruction_matches(self): + array = np.array([1.0, 1.0, 1.0, 5.0, 5.0, 9.0, 9.0, 9.0, 9.0]) + result = _kmeans1d.cluster(array, 3) + centroids = np.asarray(result.centroids) + clusters = np.asarray(result.clusters) + np.testing.assert_allclose(centroids, [1.0, 5.0, 9.0], rtol=_RTOL, atol=_ATOL) + np.testing.assert_allclose(centroids[clusters], array, rtol=_RTOL, atol=_ATOL) + + # -- Larger/random inputs: compare against golden values in assets/kmeans1d/ -- + + @pytest.mark.parametrize("n", [2, 5, 50, 2_000]) + @pytest.mark.parametrize("k", [2, 4, 16, 64, 256]) + def test_matches_known_unweighted_clustering(self, n, k): + golden = np.load(_ASSETS_DIR / "unweighted.npz") + result = _kmeans1d.cluster(golden[f"array_n{n}_k{k}"], k) + + np.testing.assert_array_equal(result.clusters, golden[f"clusters_n{n}_k{k}"]) + np.testing.assert_allclose( + result.centroids, golden[f"centroids_n{n}_k{k}"], rtol=_RTOL, atol=_ATOL + ) + + @pytest.mark.parametrize("n", [5, 500, 2_000]) + @pytest.mark.parametrize("k", [2, 4, 16, 64, 256]) + def test_matches_known_weighted_clustering(self, n, k): + golden = np.load(_ASSETS_DIR / "weighted.npz") + result = _kmeans1d.cluster( + golden[f"array_n{n}_k{k}"], k, weights=golden[f"weights_n{n}_k{k}"] + ) + + np.testing.assert_array_equal(result.clusters, golden[f"clusters_n{n}_k{k}"]) + np.testing.assert_allclose( + result.centroids, golden[f"centroids_n{n}_k{k}"], rtol=_RTOL, atol=_ATOL + ) + + @pytest.mark.parametrize("n", [10_000, 20_000, 50_000, 100_000]) + @pytest.mark.parametrize("k", [4, 64]) + def test_matches_known_clustering_at_scale(self, n, k): + # O(k*n) time and memory: n=1e5 x k=64 is ~100 MB per DP table. 1e6/1e7 are + # infeasible for the exact DP. + golden = np.load(_ASSETS_DIR / "large_n.npz") + result = _kmeans1d.cluster(golden[f"array_n{n}_k{k}"], k) + + np.testing.assert_array_equal(result.clusters, golden[f"clusters_n{n}_k{k}"]) + np.testing.assert_allclose( + result.centroids, golden[f"centroids_n{n}_k{k}"], rtol=_RTOL, atol=_ATOL + ) + + def test_matches_known_fp16_clustering(self): + # fp16 has few representable values, so duplicates (ties) are expected; only + # the reconstruction is compared, per test_collapse_fewer_than_k_clusters. + golden = np.load(_ASSETS_DIR / "fp16_duplicates.npz") + result = _kmeans1d.cluster(golden["array"], 16) + centroids = np.asarray(result.centroids) + clusters = np.asarray(result.clusters) + + np.testing.assert_allclose(centroids, golden["centroids"], rtol=_RTOL, atol=_ATOL) + np.testing.assert_allclose( + centroids[clusters], golden["reconstruction"], rtol=_RTOL, atol=_ATOL + ) + + def test_matches_known_weighted_integer_counts_clustering(self): + # Mirrors the fast/dedup path: cluster unique values weighted by counts. + golden = np.load(_ASSETS_DIR / "weighted_integer_counts.npz") + result = _kmeans1d.cluster(golden["values"], 16, weights=golden["counts"]) + + np.testing.assert_array_equal(result.clusters, golden["clusters"]) + np.testing.assert_allclose(result.centroids, golden["centroids"], rtol=_RTOL, atol=_ATOL) + + # -- Property-based behavior (no oracle, no golden file needed) -- + + def test_weights_equal_repetition(self): + # Weighting a value by w is equivalent to repeating it w times. + repeated = np.array([1.0, 1.0, 5.0, 5.0, 5.0, 9.0]) + unique_values = np.array([1.0, 5.0, 9.0]) + counts = np.array([2.0, 3.0, 1.0]) + + repeated_result = _kmeans1d.cluster(repeated, 2) + weighted_result = _kmeans1d.cluster(unique_values, 2, weights=counts) + + np.testing.assert_allclose( + np.sort(np.unique(repeated_result.centroids)), + np.sort(weighted_result.centroids), + rtol=_RTOL, + atol=_ATOL, + ) + + def test_centroids_are_ascending_and_labels_in_range(self): + rng = np.random.default_rng(11) + array = rng.standard_normal(1000) + result = _kmeans1d.cluster(array, 16) + + centroids = np.asarray(result.centroids) + clusters = np.asarray(result.clusters) + assert np.all(np.diff(centroids) >= 0), "centroids must be ascending" + assert clusters.min() >= 0 and clusters.max() < len(centroids) + assert clusters.shape == array.shape + assert centroids.dtype == np.float64 + + def test_accepts_list_and_preserves_original_order(self): + # A plain Python list (the doc-script call style) must work, and labels are + # returned in the original (unsorted) order. + result = _kmeans1d.cluster([9.0, 0.1, 9.1, 0.0], 2) + clusters = np.asarray(result.clusters) + centroids = np.asarray(result.centroids) + # Points 0,2 are the large cluster; 1,3 the small one. + assert clusters[0] == clusters[2] + assert clusters[1] == clusters[3] + assert clusters[0] != clusters[1] + np.testing.assert_allclose(centroids, [0.05, 9.05], rtol=_RTOL, atol=_ATOL)