euclimix 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (53) hide show
  1. euclimix-0.1.0/MANIFEST.in +6 -0
  2. euclimix-0.1.0/PKG-INFO +215 -0
  3. euclimix-0.1.0/README.md +199 -0
  4. euclimix-0.1.0/docs/implementation.md +66 -0
  5. euclimix-0.1.0/docs/paper_sources.json +100 -0
  6. euclimix-0.1.0/docs/validation.md +61 -0
  7. euclimix-0.1.0/examples/native.cpp +12 -0
  8. euclimix-0.1.0/examples/python_api.py +15 -0
  9. euclimix-0.1.0/pyproject.toml +36 -0
  10. euclimix-0.1.0/setup.cfg +4 -0
  11. euclimix-0.1.0/src/euclimix/__init__.py +17 -0
  12. euclimix-0.1.0/src/euclimix/_api.py +166 -0
  13. euclimix-0.1.0/src/euclimix/_extension.py +62 -0
  14. euclimix-0.1.0/src/euclimix/_result.py +46 -0
  15. euclimix-0.1.0/src/euclimix/_validation.py +56 -0
  16. euclimix-0.1.0/src/euclimix/csrc/CMakeLists.txt +32 -0
  17. euclimix-0.1.0/src/euclimix/csrc/common.h +65 -0
  18. euclimix-0.1.0/src/euclimix/csrc/controlled.cu +142 -0
  19. euclimix-0.1.0/src/euclimix/csrc/controlled_bindings.cpp +11 -0
  20. euclimix-0.1.0/src/euclimix/csrc/include/euclimix/distance.h +81 -0
  21. euclimix-0.1.0/src/euclimix/csrc/native_example.cpp +17 -0
  22. euclimix-0.1.0/src/euclimix/csrc/paper.cu +150 -0
  23. euclimix-0.1.0/src/euclimix/csrc/paper_bindings.cpp +10 -0
  24. euclimix-0.1.0/src/euclimix/csrc/state_example.cpp +43 -0
  25. euclimix-0.1.0/src/euclimix/py.typed +0 -0
  26. euclimix-0.1.0/src/euclimix.egg-info/PKG-INFO +215 -0
  27. euclimix-0.1.0/src/euclimix.egg-info/SOURCES.txt +51 -0
  28. euclimix-0.1.0/src/euclimix.egg-info/dependency_links.txt +1 -0
  29. euclimix-0.1.0/src/euclimix.egg-info/requires.txt +10 -0
  30. euclimix-0.1.0/src/euclimix.egg-info/top_level.txt +1 -0
  31. euclimix-0.1.0/tests/conftest.py +12 -0
  32. euclimix-0.1.0/tests/reference/README.md +5 -0
  33. euclimix-0.1.0/tests/reference/center_kernels.cu +270 -0
  34. euclimix-0.1.0/tests/reference/center_kernels.cuh +28 -0
  35. euclimix-0.1.0/tests/reference/euclidean_cuda.cpp +33 -0
  36. euclimix-0.1.0/tests/reference/euclidean_cuda_advanced.cpp +32 -0
  37. euclimix-0.1.0/tests/reference/euclidean_cuda_mix.cpp +96 -0
  38. euclimix-0.1.0/tests/reference/euclidean_cuda_tiled.cpp +66 -0
  39. euclimix-0.1.0/tests/reference/euclidean_cuda_uniform.cpp +42 -0
  40. euclimix-0.1.0/tests/reference/euclidean_kernels.cu +152 -0
  41. euclimix-0.1.0/tests/reference/euclidean_kernels.cuh +32 -0
  42. euclimix-0.1.0/tests/reference/euclidean_kernels_advanced.cu +460 -0
  43. euclimix-0.1.0/tests/reference/euclidean_kernels_advanced.cuh +72 -0
  44. euclimix-0.1.0/tests/reference/euclidean_kernels_mix.cu +1626 -0
  45. euclimix-0.1.0/tests/reference/euclidean_kernels_mix.cuh +578 -0
  46. euclimix-0.1.0/tests/reference/euclidean_kernels_tiled.cu +528 -0
  47. euclimix-0.1.0/tests/reference/euclidean_kernels_tiled.cuh +81 -0
  48. euclimix-0.1.0/tests/reference/euclidean_kernels_tiled_tc.cu +300 -0
  49. euclimix-0.1.0/tests/reference/euclidean_kernels_uniform.cu +532 -0
  50. euclimix-0.1.0/tests/reference/euclidean_kernels_uniform.cuh +57 -0
  51. euclimix-0.1.0/tests/test_cuda.py +240 -0
  52. euclimix-0.1.0/tests/test_paper_regression.py +48 -0
  53. euclimix-0.1.0/tests/test_python.py +308 -0
@@ -0,0 +1,6 @@
1
+ include README.md pyproject.toml
2
+ recursive-include src/euclimix *.py *.typed *.cpp *.cu *.h CMakeLists.txt
3
+ recursive-include docs *.md *.json
4
+ recursive-include examples *.py *.cpp
5
+ recursive-include tests *.py *.cu *.cuh *.cpp *.md
6
+
@@ -0,0 +1,215 @@
1
+ Metadata-Version: 2.4
2
+ Name: euclimix
3
+ Version: 0.1.0
4
+ Summary: Mixed-precision squared Euclidean distances with selective CUDA recomputation
5
+ Author: Xinye Chen
6
+ Requires-Python: >=3.11
7
+ Description-Content-Type: text/markdown
8
+ Requires-Dist: torch<3,>=2.11
9
+ Requires-Dist: ninja>=1.11
10
+ Provides-Extra: test
11
+ Requires-Dist: pytest>=7.4; extra == "test"
12
+ Requires-Dist: pytest-cov>=4.1; extra == "test"
13
+ Requires-Dist: numpy>=1.26; extra == "test"
14
+ Provides-Extra: build
15
+ Requires-Dist: build>=1.2; extra == "build"
16
+
17
+ # euclimix
18
+
19
+ [![CI](https://img.shields.io/github/actions/workflow/status/inEXASCALE/euclid-cuda/cpu.yml?branch=main&label=CI&logo=githubactions&logoColor=white&style=flat-square)](https://github.com/inEXASCALE/euclid-cuda/actions/workflows/cpu.yml)
20
+ [![CUDA tests](https://img.shields.io/github/actions/workflow/status/inEXASCALE/euclid-cuda/gpu.yml?branch=main&label=CUDA%20tests&logo=nvidia&logoColor=white&style=flat-square)](https://github.com/inEXASCALE/euclid-cuda/actions/workflows/gpu.yml)
21
+ [![Python coverage gate: 100%](https://img.shields.io/badge/Python%20coverage%20gate-100%25-2C974B?logo=pytest&logoColor=white&style=flat-square)](docs/validation.md#python-tests)
22
+
23
+ [![Python 3.11+](https://img.shields.io/badge/Python-3.11%2B-3776AB?logo=python&logoColor=white&style=flat-square)](pyproject.toml)
24
+ [![PyTorch 2.11+](https://img.shields.io/badge/PyTorch-2.11%2B-EE4C2C?logo=pytorch&logoColor=white&style=flat-square)](https://pytorch.org/)
25
+ [![CUDA 12.8+](https://img.shields.io/badge/CUDA-12.8%2B-76B900?logo=nvidia&logoColor=white&style=flat-square)](https://developer.nvidia.com/cuda-toolkit)
26
+ [![arXiv:2407.12208](https://img.shields.io/badge/arXiv-2407.12208-B31B1B?logo=arxiv&logoColor=white&style=flat-square)](#citation)
27
+
28
+ Mixed-precision Euclidean distances on NVIDIA GPUs, with Python/pybind11 and
29
+ C++/CUDA interfaces. Low-precision expanded distances are combined with
30
+ selective higher-precision recomputation.
31
+
32
+ ## Installation
33
+
34
+ Requirements: Linux, Python 3.11+, CUDA-enabled PyTorch 2.11+, CUDA toolkit
35
+ 12.8+ with `nvcc`, and an Ampere or newer GPU (SM80+). A compatible C++17
36
+ compiler is required; newer LibTorch versions may require C++20.
37
+ The PyTorch CUDA wheel provides runtime libraries, not the CUDA compiler.
38
+ cuDNN is not required.
39
+
40
+ Example with CUDA 12.8:
41
+
42
+ ```bash
43
+ git clone https://github.com/inEXASCALE/euclid-cuda.git
44
+ cd euclid-cuda
45
+ conda create -n euclimix python=3.12 -y
46
+ conda activate euclimix
47
+ python -m pip install torch==2.11.0 --index-url https://download.pytorch.org/whl/cu128
48
+ python -m pip install .
49
+ python -c 'import euclimix; euclimix.build_cuda()'
50
+ ```
51
+
52
+ Set `CUDA_HOME` if the toolkit is not detected. Extensions compile on first
53
+ use and are cached by PyTorch. Set `TORCH_CUDA_ARCH_LIST` for the target GPU
54
+ (for example, `8.0` for A100), `MAX_JOBS` to limit build parallelism, and
55
+ `EUCLIMIX_BUILD_VERBOSE=1` for compiler output. Build and warm up the functions
56
+ before timing them.
57
+
58
+ ## Python interface
59
+
60
+ ### Squared and Euclidean distances
61
+
62
+ ```python
63
+ import torch
64
+ from euclimix import pairwise_squared_distances, pairwise_distances
65
+
66
+ X = torch.tensor([[0., 0.], [3., 4.]], device='cuda', dtype=torch.float32)
67
+
68
+ squared = pairwise_squared_distances(X)
69
+ distances = pairwise_distances(X)
70
+
71
+ print(squared.cpu()) # tensor([[ 0., 25.], [25., 0.]])
72
+ print(distances.cpu()) # tensor([[0., 5.], [5., 0.]])
73
+ ```
74
+
75
+ Omitting `Y` computes all distances within `X`. Pass a second matrix to
76
+ compute an `(n_samples_X, n_samples_Y)` result.
77
+
78
+ ### Precision and recomputation statistics
79
+
80
+ ```python
81
+ import torch
82
+ from euclimix import pairwise_squared_distances
83
+
84
+ X = torch.randn(1024, 128, device='cuda', dtype=torch.float32)
85
+ Y = torch.randn(512, 128, device='cuda', dtype=torch.float32)
86
+
87
+ result = pairwise_squared_distances(
88
+ X, Y, precision='fp16', correction='fp64', rho=5.0, return_info=True,
89
+ )
90
+
91
+ print(result.distances.shape) # torch.Size([1024, 512])
92
+ print(result.distances.dtype) # torch.float32
93
+ print(result.fallback_count)
94
+ print(f'Recomputed: {result.fallback_fraction:.2%}')
95
+ ```
96
+
97
+ ### Storage and threshold selection
98
+
99
+ ```python
100
+ import torch
101
+ from euclimix import controlled_squared_distances, direct_squared_distances
102
+
103
+ X = torch.randn(256, 128, device='cuda', dtype=torch.float32)
104
+ Y = torch.randn(128, 128, device='cuda', dtype=torch.float32)
105
+
106
+ result = controlled_squared_distances(
107
+ X, Y, precision='bf16', storage='fp32', test='extended', return_info=True,
108
+ )
109
+ reference = direct_squared_distances(X, Y, precision='fp64')
110
+ max_absolute_error = (result.distances.double() - reference).abs().max()
111
+
112
+ print(result.model)
113
+ print(result.fallback_fraction)
114
+ print(max_absolute_error.item())
115
+ ```
116
+
117
+ Use `storage='low'` to store intermediates in the input format and
118
+ `test='original'` to select the empirical threshold. The FP64 direct
119
+ reference retains FP64 output.
120
+
121
+ | Function | Result | Implementation |
122
+ |---|---|---|
123
+ | `pairwise_squared_distances` | FP32 squared distances | Expanded formula, Original test, index-list recomputation |
124
+ | `controlled_squared_distances` | FP32 squared distances | Configurable storage and Original/Extended test, mask-based recomputation |
125
+ | `direct_squared_distances` | FP32 or FP64 squared distances | Direct subtraction and FMA accumulation |
126
+ | `pairwise_distances` | FP32 Euclidean distances | Square root of `pairwise_squared_distances` |
127
+
128
+ Inputs must be finite, dense FP32 CUDA matrices on the same device. Omit `Y`
129
+ to compare `X` with itself. Noncontiguous inputs and negative views are
130
+ materialized. Empty row sets are supported; zero features and more than
131
+ `2**31-1` pairs are rejected. Autograd and CUDA graph capture are unsupported.
132
+
133
+ FP16/BF16 cross products use FP32 accumulation. FP32/FP64 recomputation reads
134
+ the original FP32 coordinates; mixed outputs are stored in FP32. The
135
+ `correction='same'` option uses quantized coordinates and low-format arithmetic,
136
+ which can lose accuracy or overflow. Finite inputs can also produce squared
137
+ distances outside the FP32 output range.
138
+
139
+ The Extended error bound assumes a bound on the inner-product arithmetic;
140
+ its cuBLAS application is conditional on that assumption. A relative-error
141
+ bound for distances does not ensure identical nearest-center assignments.
142
+ See [Numerical implementation](docs/implementation.md) for the precision paths
143
+ and threshold definitions.
144
+
145
+ Input validation and reading fallback statistics synchronize with the GPU.
146
+ Native kernels use the current CUDA stream. Callers using multiple streams
147
+ must establish producer/consumer dependencies with events or `wait_stream`.
148
+ The FP32 index-list path follows PyTorch's TF32/matrix-multiplication policy;
149
+ autocast does not override explicit native precision choices.
150
+
151
+ ## C++/CUDA interface
152
+
153
+ The native API uses LibTorch CUDA tensors and the types in
154
+ `euclimix/distance.h`. See [examples/native.cpp](examples/native.cpp).
155
+
156
+ ```cpp
157
+ #include <euclimix/distance.h>
158
+
159
+ euclimix::DistanceResult distances(torch::Tensor X, torch::Tensor Y) {
160
+ euclimix::PaperOptions options;
161
+ options.precision = euclimix::InputPrecision::FP16;
162
+ options.correction = euclimix::Correction::FP64;
163
+ options.rho = 5.0;
164
+ return euclimix::paper_squared_distances(X, Y, options);
165
+ }
166
+ ```
167
+
168
+ Build the native library and examples:
169
+
170
+ ```bash
171
+ cmake -S src/euclimix/csrc -B .native-build \
172
+ -DCMAKE_PREFIX_PATH="$(python -c 'import torch; print(torch.utils.cmake_prefix_path)')" \
173
+ -DCMAKE_CUDA_ARCHITECTURES=80 -DEUCLIMIX_BUILD_EXAMPLE=ON
174
+ cmake --build .native-build -j2
175
+ .native-build/euclimix_native_example
176
+ .native-build/euclimix_state_example
177
+ ```
178
+
179
+ Native inputs must be contiguous with negative views resolved. The CMake
180
+ library can be used without Python or pybind11.
181
+
182
+ ## Testing
183
+
184
+ ```bash
185
+ python -m pip install -e '.[test,build]'
186
+ python -m pytest -m 'not cuda' --cov=euclimix --cov-branch --cov-report=term-missing
187
+ EUCLIMIX_REQUIRE_CUDA=1 python -m pytest -m cuda
188
+ EUCLIMIX_REQUIRE_CUDA=1 EUCLIMIX_COMPARE_PAPER=1 python -m pytest -m cuda
189
+ ```
190
+
191
+ CPU CI requires 100% Python statement and branch coverage. CUDA integration
192
+ tests require a GPU; the GPU workflow uses a self-hosted Linux runner labeled
193
+ `self-hosted, linux, x64, gpu`. Python coverage does not measure CUDA code.
194
+ See [Testing](docs/validation.md) for device-memory checks and runner setup.
195
+
196
+ Reference CUDA sources in `tests/reference` are used for regression tests.
197
+ Their source hashes are recorded in [paper_sources.json](docs/paper_sources.json).
198
+
199
+ ## Citation
200
+
201
+ If you use euclimix in your research, please cite the accompanying
202
+ [paper](https://arxiv.org/abs/2407.12208):
203
+
204
+ ```bibtex
205
+ @misc{carson2024computing,
206
+ title = {Computing $k$-means in mixed precision},
207
+ author = {Erin Carson and Xinye Chen and Xiaobo Liu},
208
+ year = {2024},
209
+ eprint = {2407.12208},
210
+ archivePrefix = {arXiv},
211
+ primaryClass = {math.NA},
212
+ doi = {10.48550/arXiv.2407.12208},
213
+ url = {https://arxiv.org/abs/2407.12208}
214
+ }
215
+ ```
@@ -0,0 +1,199 @@
1
+ # euclimix
2
+
3
+ [![CI](https://img.shields.io/github/actions/workflow/status/inEXASCALE/euclid-cuda/cpu.yml?branch=main&label=CI&logo=githubactions&logoColor=white&style=flat-square)](https://github.com/inEXASCALE/euclid-cuda/actions/workflows/cpu.yml)
4
+ [![CUDA tests](https://img.shields.io/github/actions/workflow/status/inEXASCALE/euclid-cuda/gpu.yml?branch=main&label=CUDA%20tests&logo=nvidia&logoColor=white&style=flat-square)](https://github.com/inEXASCALE/euclid-cuda/actions/workflows/gpu.yml)
5
+ [![Python coverage gate: 100%](https://img.shields.io/badge/Python%20coverage%20gate-100%25-2C974B?logo=pytest&logoColor=white&style=flat-square)](docs/validation.md#python-tests)
6
+
7
+ [![Python 3.11+](https://img.shields.io/badge/Python-3.11%2B-3776AB?logo=python&logoColor=white&style=flat-square)](pyproject.toml)
8
+ [![PyTorch 2.11+](https://img.shields.io/badge/PyTorch-2.11%2B-EE4C2C?logo=pytorch&logoColor=white&style=flat-square)](https://pytorch.org/)
9
+ [![CUDA 12.8+](https://img.shields.io/badge/CUDA-12.8%2B-76B900?logo=nvidia&logoColor=white&style=flat-square)](https://developer.nvidia.com/cuda-toolkit)
10
+ [![arXiv:2407.12208](https://img.shields.io/badge/arXiv-2407.12208-B31B1B?logo=arxiv&logoColor=white&style=flat-square)](#citation)
11
+
12
+ Mixed-precision Euclidean distances on NVIDIA GPUs, with Python/pybind11 and
13
+ C++/CUDA interfaces. Low-precision expanded distances are combined with
14
+ selective higher-precision recomputation.
15
+
16
+ ## Installation
17
+
18
+ Requirements: Linux, Python 3.11+, CUDA-enabled PyTorch 2.11+, CUDA toolkit
19
+ 12.8+ with `nvcc`, and an Ampere or newer GPU (SM80+). A compatible C++17
20
+ compiler is required; newer LibTorch versions may require C++20.
21
+ The PyTorch CUDA wheel provides runtime libraries, not the CUDA compiler.
22
+ cuDNN is not required.
23
+
24
+ Example with CUDA 12.8:
25
+
26
+ ```bash
27
+ git clone https://github.com/inEXASCALE/euclid-cuda.git
28
+ cd euclid-cuda
29
+ conda create -n euclimix python=3.12 -y
30
+ conda activate euclimix
31
+ python -m pip install torch==2.11.0 --index-url https://download.pytorch.org/whl/cu128
32
+ python -m pip install .
33
+ python -c 'import euclimix; euclimix.build_cuda()'
34
+ ```
35
+
36
+ Set `CUDA_HOME` if the toolkit is not detected. Extensions compile on first
37
+ use and are cached by PyTorch. Set `TORCH_CUDA_ARCH_LIST` for the target GPU
38
+ (for example, `8.0` for A100), `MAX_JOBS` to limit build parallelism, and
39
+ `EUCLIMIX_BUILD_VERBOSE=1` for compiler output. Build and warm up the functions
40
+ before timing them.
41
+
42
+ ## Python interface
43
+
44
+ ### Squared and Euclidean distances
45
+
46
+ ```python
47
+ import torch
48
+ from euclimix import pairwise_squared_distances, pairwise_distances
49
+
50
+ X = torch.tensor([[0., 0.], [3., 4.]], device='cuda', dtype=torch.float32)
51
+
52
+ squared = pairwise_squared_distances(X)
53
+ distances = pairwise_distances(X)
54
+
55
+ print(squared.cpu()) # tensor([[ 0., 25.], [25., 0.]])
56
+ print(distances.cpu()) # tensor([[0., 5.], [5., 0.]])
57
+ ```
58
+
59
+ Omitting `Y` computes all distances within `X`. Pass a second matrix to
60
+ compute an `(n_samples_X, n_samples_Y)` result.
61
+
62
+ ### Precision and recomputation statistics
63
+
64
+ ```python
65
+ import torch
66
+ from euclimix import pairwise_squared_distances
67
+
68
+ X = torch.randn(1024, 128, device='cuda', dtype=torch.float32)
69
+ Y = torch.randn(512, 128, device='cuda', dtype=torch.float32)
70
+
71
+ result = pairwise_squared_distances(
72
+ X, Y, precision='fp16', correction='fp64', rho=5.0, return_info=True,
73
+ )
74
+
75
+ print(result.distances.shape) # torch.Size([1024, 512])
76
+ print(result.distances.dtype) # torch.float32
77
+ print(result.fallback_count)
78
+ print(f'Recomputed: {result.fallback_fraction:.2%}')
79
+ ```
80
+
81
+ ### Storage and threshold selection
82
+
83
+ ```python
84
+ import torch
85
+ from euclimix import controlled_squared_distances, direct_squared_distances
86
+
87
+ X = torch.randn(256, 128, device='cuda', dtype=torch.float32)
88
+ Y = torch.randn(128, 128, device='cuda', dtype=torch.float32)
89
+
90
+ result = controlled_squared_distances(
91
+ X, Y, precision='bf16', storage='fp32', test='extended', return_info=True,
92
+ )
93
+ reference = direct_squared_distances(X, Y, precision='fp64')
94
+ max_absolute_error = (result.distances.double() - reference).abs().max()
95
+
96
+ print(result.model)
97
+ print(result.fallback_fraction)
98
+ print(max_absolute_error.item())
99
+ ```
100
+
101
+ Use `storage='low'` to store intermediates in the input format and
102
+ `test='original'` to select the empirical threshold. The FP64 direct
103
+ reference retains FP64 output.
104
+
105
+ | Function | Result | Implementation |
106
+ |---|---|---|
107
+ | `pairwise_squared_distances` | FP32 squared distances | Expanded formula, Original test, index-list recomputation |
108
+ | `controlled_squared_distances` | FP32 squared distances | Configurable storage and Original/Extended test, mask-based recomputation |
109
+ | `direct_squared_distances` | FP32 or FP64 squared distances | Direct subtraction and FMA accumulation |
110
+ | `pairwise_distances` | FP32 Euclidean distances | Square root of `pairwise_squared_distances` |
111
+
112
+ Inputs must be finite, dense FP32 CUDA matrices on the same device. Omit `Y`
113
+ to compare `X` with itself. Noncontiguous inputs and negative views are
114
+ materialized. Empty row sets are supported; zero features and more than
115
+ `2**31-1` pairs are rejected. Autograd and CUDA graph capture are unsupported.
116
+
117
+ FP16/BF16 cross products use FP32 accumulation. FP32/FP64 recomputation reads
118
+ the original FP32 coordinates; mixed outputs are stored in FP32. The
119
+ `correction='same'` option uses quantized coordinates and low-format arithmetic,
120
+ which can lose accuracy or overflow. Finite inputs can also produce squared
121
+ distances outside the FP32 output range.
122
+
123
+ The Extended error bound assumes a bound on the inner-product arithmetic;
124
+ its cuBLAS application is conditional on that assumption. A relative-error
125
+ bound for distances does not ensure identical nearest-center assignments.
126
+ See [Numerical implementation](docs/implementation.md) for the precision paths
127
+ and threshold definitions.
128
+
129
+ Input validation and reading fallback statistics synchronize with the GPU.
130
+ Native kernels use the current CUDA stream. Callers using multiple streams
131
+ must establish producer/consumer dependencies with events or `wait_stream`.
132
+ The FP32 index-list path follows PyTorch's TF32/matrix-multiplication policy;
133
+ autocast does not override explicit native precision choices.
134
+
135
+ ## C++/CUDA interface
136
+
137
+ The native API uses LibTorch CUDA tensors and the types in
138
+ `euclimix/distance.h`. See [examples/native.cpp](examples/native.cpp).
139
+
140
+ ```cpp
141
+ #include <euclimix/distance.h>
142
+
143
+ euclimix::DistanceResult distances(torch::Tensor X, torch::Tensor Y) {
144
+ euclimix::PaperOptions options;
145
+ options.precision = euclimix::InputPrecision::FP16;
146
+ options.correction = euclimix::Correction::FP64;
147
+ options.rho = 5.0;
148
+ return euclimix::paper_squared_distances(X, Y, options);
149
+ }
150
+ ```
151
+
152
+ Build the native library and examples:
153
+
154
+ ```bash
155
+ cmake -S src/euclimix/csrc -B .native-build \
156
+ -DCMAKE_PREFIX_PATH="$(python -c 'import torch; print(torch.utils.cmake_prefix_path)')" \
157
+ -DCMAKE_CUDA_ARCHITECTURES=80 -DEUCLIMIX_BUILD_EXAMPLE=ON
158
+ cmake --build .native-build -j2
159
+ .native-build/euclimix_native_example
160
+ .native-build/euclimix_state_example
161
+ ```
162
+
163
+ Native inputs must be contiguous with negative views resolved. The CMake
164
+ library can be used without Python or pybind11.
165
+
166
+ ## Testing
167
+
168
+ ```bash
169
+ python -m pip install -e '.[test,build]'
170
+ python -m pytest -m 'not cuda' --cov=euclimix --cov-branch --cov-report=term-missing
171
+ EUCLIMIX_REQUIRE_CUDA=1 python -m pytest -m cuda
172
+ EUCLIMIX_REQUIRE_CUDA=1 EUCLIMIX_COMPARE_PAPER=1 python -m pytest -m cuda
173
+ ```
174
+
175
+ CPU CI requires 100% Python statement and branch coverage. CUDA integration
176
+ tests require a GPU; the GPU workflow uses a self-hosted Linux runner labeled
177
+ `self-hosted, linux, x64, gpu`. Python coverage does not measure CUDA code.
178
+ See [Testing](docs/validation.md) for device-memory checks and runner setup.
179
+
180
+ Reference CUDA sources in `tests/reference` are used for regression tests.
181
+ Their source hashes are recorded in [paper_sources.json](docs/paper_sources.json).
182
+
183
+ ## Citation
184
+
185
+ If you use euclimix in your research, please cite the accompanying
186
+ [paper](https://arxiv.org/abs/2407.12208):
187
+
188
+ ```bibtex
189
+ @misc{carson2024computing,
190
+ title = {Computing $k$-means in mixed precision},
191
+ author = {Erin Carson and Xinye Chen and Xiaobo Liu},
192
+ year = {2024},
193
+ eprint = {2407.12208},
194
+ archivePrefix = {arXiv},
195
+ primaryClass = {math.NA},
196
+ doi = {10.48550/arXiv.2407.12208},
197
+ url = {https://arxiv.org/abs/2407.12208}
198
+ }
199
+ ```
@@ -0,0 +1,66 @@
1
+ # Numerical implementation
2
+
3
+ The mixed-precision functions evaluate the expanded squared-distance formula,
4
+ test each result, and recompute rejected entries using the direct formula.
5
+
6
+ ## Index-list implementation
7
+
8
+ `pairwise_squared_distances` supports FP16, BF16 and FP32 default paths.
9
+ FP16/BF16 inputs are quantized from the original FP32 coordinates. Squared
10
+ norms accumulate in FP32 and are stored in the selected format. Cross
11
+ products use cuBLAS GemmEx with FP32 compute and low-format output. The FP32
12
+ path uses LibTorch `mm` and its matrix-multiplication policy.
13
+
14
+ Distance formation, clamping to zero and threshold comparison use FP32.
15
+ The Original test recomputes distances at or below `rho * c(r) * (p + q)`,
16
+ where `p` and `q` are stored squared norms and `r` is the feature count:
17
+
18
+ | Precision / correction | `c(r)` |
19
+ |---|---|
20
+ | FP16 | `(2*r + 4) * 2**-11` |
21
+ | BF16 | `(2*r + 4) * 2**-7` |
22
+ | FP32 / FP32 or same | `(2*r + 4) * 2**-23` |
23
+ | FP32 / FP64 | `2*r * 2**-23` |
24
+
25
+ These are empirical coefficients. The BF16 coefficient uses `2**-7`,
26
+ whereas BF16 unit roundoff is `2**-8`.
27
+
28
+ Rejected entries are appended to an index list and recomputed by a separate
29
+ kernel. FP32/FP64 correction uses original FP32 coordinates; `same` uses
30
+ quantized coordinates and low-format arithmetic. Mixed outputs are FP32,
31
+ including after FP64 correction. Nonfinite expanded results or thresholds
32
+ also trigger recomputation.
33
+
34
+ ## Mask-based implementation
35
+
36
+ `controlled_squared_distances` supports FP16/BF16 inputs, low/FP32
37
+ intermediate storage and Original/Extended tests. Norms use sequential FP32
38
+ FMA. Cross products use either cuBLAS or coordinate-order FP32 FMA.
39
+
40
+ The Extended bound includes input-conversion residuals, inner-product
41
+ arithmetic, intermediate storage and gradual underflow. Bounds and thresholds
42
+ are evaluated conservatively using directed FP64 rounding. The test accepts
43
+ a finite squared distance `d` only when `d > upward_round(rho * B)`.
44
+ Under the inner-product error assumption, accepted distances have relative
45
+ error at most `1/(rho-1)`. For cuBLAS, the coefficient is an assumption about
46
+ the library reduction; FP32 compute alone does not establish it.
47
+
48
+ Both tests compute diagnostic bounds. Rejected entries are recomputed by
49
+ scanning a mask. FP32/FP64 correction reads original FP32 inputs, with FP32
50
+ output in the Python interface. The native interface also supports FP64 output.
51
+ The accepted-path bound does not cover rounding after direct recomputation.
52
+
53
+ ## Execution
54
+
55
+ All native calls use a CUDA device guard and the current PyTorch stream.
56
+ Input storage is recorded on that stream. PyTorch configures the cuBLAS
57
+ handle's stream and workspace; euclimix preserves them and restores modified
58
+ math/pointer modes. Rebinding the stream would reset the workspace, as
59
+ documented by [cuBLAS](https://docs.nvidia.com/cuda/archive/12.8.0/cublas/index.html#cublassetstream).
60
+
61
+ The index-list extension uses `--use_fast_math`. The mask-based extension
62
+ uses `--ftz=false --fmad=false --prec-div=true --prec-sqrt=true`. Separate
63
+ modules keep these compiler settings independent. Temporary buffers and
64
+ outputs are managed by LibTorch tensors.
65
+
66
+ Algorithm and analysis: [arXiv:2407.12208](https://arxiv.org/abs/2407.12208).
@@ -0,0 +1,100 @@
1
+ {
2
+ "paper_commit": "eb195775a59930c146bbbcf283f85608585333b0",
3
+ "original_files": [
4
+ {
5
+ "file": "tests/reference/center_kernels.cu",
6
+ "paper_file": "exps/csrc/center_kernels.cu",
7
+ "sha256": "16ac5f909f07d6c154a47b035fb2ee36d51fd925b017e3d0412381b21cbb3d56"
8
+ },
9
+ {
10
+ "file": "tests/reference/center_kernels.cuh",
11
+ "paper_file": "exps/csrc/center_kernels.cuh",
12
+ "sha256": "13e15ea723e358e3bcab28a12394e95889137fd17e6cd1d3f946f503641c9be4"
13
+ },
14
+ {
15
+ "file": "tests/reference/euclidean_cuda.cpp",
16
+ "paper_file": "exps/csrc/euclidean_cuda.cpp",
17
+ "sha256": "bb5e2c338bdd1e4e350427eecc054dcd5b16bd8bc34f7319ad1f7d1b4cf69314"
18
+ },
19
+ {
20
+ "file": "tests/reference/euclidean_cuda_advanced.cpp",
21
+ "paper_file": "exps/csrc/euclidean_cuda_advanced.cpp",
22
+ "sha256": "533f534a1eb19646ce61508a7fb6ae2b1e3ad16caa17e744d88ea0119a51cf1c"
23
+ },
24
+ {
25
+ "file": "tests/reference/euclidean_cuda_mix.cpp",
26
+ "paper_file": "exps/csrc/euclidean_cuda_mix.cpp",
27
+ "sha256": "1d1f8d62a2e3a68c4ea7e05992f793bf12b6d6b7577947e047d0cb3ca5933061"
28
+ },
29
+ {
30
+ "file": "tests/reference/euclidean_cuda_tiled.cpp",
31
+ "paper_file": "exps/csrc/euclidean_cuda_tiled.cpp",
32
+ "sha256": "f00f3a9ca93a6467bd25e532b8ba8f48c43f024a73ba5cc93fda3eb07d025c1f"
33
+ },
34
+ {
35
+ "file": "tests/reference/euclidean_cuda_uniform.cpp",
36
+ "paper_file": "exps/csrc/euclidean_cuda_uniform.cpp",
37
+ "sha256": "93e2bcf5fe4db75a4c57ad16ac799db61f6d1082627d21857a10d138057ae8a1"
38
+ },
39
+ {
40
+ "file": "tests/reference/euclidean_kernels.cu",
41
+ "paper_file": "exps/csrc/euclidean_kernels.cu",
42
+ "sha256": "49978f3e11a8a5099bea80779d6724680db20d4385b8048f0559b2137dfd443e"
43
+ },
44
+ {
45
+ "file": "tests/reference/euclidean_kernels.cuh",
46
+ "paper_file": "exps/csrc/euclidean_kernels.cuh",
47
+ "sha256": "fac538925d7f057170548a656eee6b7f022e304a574847e7ec2f48d826d27375"
48
+ },
49
+ {
50
+ "file": "tests/reference/euclidean_kernels_advanced.cu",
51
+ "paper_file": "exps/csrc/euclidean_kernels_advanced.cu",
52
+ "sha256": "0ade485978bfa3ef2ab57f617bae3bc3ba19b949cf0d0fc40037b4f165bdd145"
53
+ },
54
+ {
55
+ "file": "tests/reference/euclidean_kernels_advanced.cuh",
56
+ "paper_file": "exps/csrc/euclidean_kernels_advanced.cuh",
57
+ "sha256": "d8eb4b6f4392599fd0b03c4ebd37f4f2b57deafabaa066ebddd3a4bc2da5d77d"
58
+ },
59
+ {
60
+ "file": "tests/reference/euclidean_kernels_mix.cu",
61
+ "paper_file": "exps/csrc/euclidean_kernels_mix.cu",
62
+ "sha256": "b8632295983c8d3a053f3e644005a12dc45a1d4187c094fd33c8eb60bb954b92"
63
+ },
64
+ {
65
+ "file": "tests/reference/euclidean_kernels_mix.cuh",
66
+ "paper_file": "exps/csrc/euclidean_kernels_mix.cuh",
67
+ "sha256": "8f6b40e436a533deb80329f9aace4c1c8a4e8a5fc93fb4ff8ab5fa963df36194"
68
+ },
69
+ {
70
+ "file": "tests/reference/euclidean_kernels_tiled.cu",
71
+ "paper_file": "exps/csrc/euclidean_kernels_tiled.cu",
72
+ "sha256": "aa181dd1f340ba1b5d4bf589b27afba1be46b9d2530a0449cfad7623dc16c74c"
73
+ },
74
+ {
75
+ "file": "tests/reference/euclidean_kernels_tiled.cuh",
76
+ "paper_file": "exps/csrc/euclidean_kernels_tiled.cuh",
77
+ "sha256": "c3effae1d0181406f8ef9f72640f74c1f32b8af4df5a8499a90f70a6b35ee2a3"
78
+ },
79
+ {
80
+ "file": "tests/reference/euclidean_kernels_tiled_tc.cu",
81
+ "paper_file": "exps/csrc/euclidean_kernels_tiled_tc.cu",
82
+ "sha256": "fba17ae0715e6316a3f512b78ad87e917fe17d365ca8f268ffa48df43237dae0"
83
+ },
84
+ {
85
+ "file": "tests/reference/euclidean_kernels_uniform.cu",
86
+ "paper_file": "exps/csrc/euclidean_kernels_uniform.cu",
87
+ "sha256": "d95233345b911d95ad3e4d659114252fb4d3a8e041eb713408c3c66ecf547778"
88
+ },
89
+ {
90
+ "file": "tests/reference/euclidean_kernels_uniform.cuh",
91
+ "paper_file": "exps/csrc/euclidean_kernels_uniform.cuh",
92
+ "sha256": "7d4bc4c2a6cad6c08783c0dd8b8d5be964f34182cde05b96acb107438387c04a"
93
+ }
94
+ ],
95
+ "controlled_source": {
96
+ "paper_file": "exps/reviewer_experiments/csrc/reviewer_distance.cu",
97
+ "sha256": "aa289df82785abbfdff4065a3acef64bf056c92c71feebc69c8f9462c0d0a678"
98
+ },
99
+ "repository": "https://github.com/chenxinye/mp_kmeans_paper"
100
+ }
@@ -0,0 +1,61 @@
1
+ # Testing
2
+
3
+ Install test and build dependencies:
4
+
5
+ ```bash
6
+ python -m pip install -e '.[test,build]'
7
+ ```
8
+
9
+ ## Python tests
10
+
11
+ ```bash
12
+ python -m pytest -m 'not cuda' --cov=euclimix --cov-branch \
13
+ --cov-report=term-missing --cov-report=xml
14
+ ```
15
+
16
+ CI requires 100% Python statement and branch coverage. Unit tests cover
17
+ input validation, dispatch, extension loading and result statistics. Native
18
+ entry points are stubbed in these tests; CUDA arithmetic is tested separately.
19
+
20
+ ## CUDA tests
21
+
22
+ Use a supported NVIDIA GPU, a compatible host compiler and the CUDA toolkit:
23
+
24
+ ```bash
25
+ export TORCH_CUDA_ARCH_LIST=8.0 # A100
26
+ export MAX_JOBS=2
27
+ python -c 'import euclimix; euclimix.build_cuda()'
28
+ EUCLIMIX_REQUIRE_CUDA=1 python -m pytest -m cuda
29
+ ```
30
+
31
+ Tests cover precision/correction modes, storage, threshold selection, near
32
+ pairs, overflow, underflow, tensor views, autocast and nondefault streams.
33
+ FP32 matrix-multiplication tests disable TF32 temporarily and restore the
34
+ previous policy. The device-guard test requires two GPUs.
35
+
36
+ Compare with the reference CUDA implementation:
37
+
38
+ ```bash
39
+ EUCLIMIX_REQUIRE_CUDA=1 EUCLIMIX_COMPARE_PAPER=1 python -m pytest -m cuda
40
+ ```
41
+
42
+ Check device memory accesses:
43
+
44
+ ```bash
45
+ EUCLIMIX_REQUIRE_CUDA=1 compute-sanitizer --tool memcheck --error-exitcode=1 \
46
+ python -m pytest -m cuda tests/test_cuda.py
47
+ ```
48
+
49
+ The native CMake examples in the README check the C++ interface and cuBLAS
50
+ mode/stream preservation. Source checksums for the regression fixtures are
51
+ stored in `docs/paper_sources.json`.
52
+
53
+ ## CI
54
+
55
+ The CPU workflow runs on pushes and pull requests, builds and installs a
56
+ wheel, and uploads test and coverage reports. The GPU workflow is manually
57
+ dispatched on a self-hosted Linux runner labeled `self-hosted, linux, x64, gpu`.
58
+ It requires CUDA 12.8, CMake 3.24+, compute-sanitizer and a compatible compiler.
59
+
60
+ `EUCLIMIX_REQUIRE_CUDA=1` makes a missing GPU an error. Python coverage does
61
+ not measure CUDA code.
@@ -0,0 +1,12 @@
1
+ // Include this header from an application linked with the CMake euclimix target.
2
+ #include <euclimix/distance.h>
3
+
4
+ euclimix::DistanceResult example(torch::Tensor x, torch::Tensor y) {
5
+ // Original FP32 CUDA coordinates, contiguous on one device.
6
+ // FP16 expanded path, Original criterion, direct FP64 correction.
7
+ euclimix::PaperOptions options;
8
+ options.precision = euclimix::InputPrecision::FP16;
9
+ options.correction = euclimix::Correction::FP64;
10
+ return euclimix::paper_squared_distances(x, y, options);
11
+ }
12
+
@@ -0,0 +1,15 @@
1
+ """Small CUDA example, including a direct FP64 reference."""
2
+
3
+ import torch
4
+
5
+ from euclimix import controlled_squared_distances, direct_squared_distances, pairwise_squared_distances
6
+
7
+ X = torch.tensor([[0., 0.], [3., 4.]], device="cuda")
8
+ index_list = pairwise_squared_distances(X, return_info=True)
9
+ extended = controlled_squared_distances(X, return_info=True)
10
+ reference = direct_squared_distances(X)
11
+ print("Squared distances:", index_list.distances)
12
+ print("Extended squared distances:", extended.distances)
13
+ print("Direct FP64 reference:", reference)
14
+ print("Recomputation fraction:", index_list.fallback_fraction)
15
+