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.
- euclimix-0.1.0/MANIFEST.in +6 -0
- euclimix-0.1.0/PKG-INFO +215 -0
- euclimix-0.1.0/README.md +199 -0
- euclimix-0.1.0/docs/implementation.md +66 -0
- euclimix-0.1.0/docs/paper_sources.json +100 -0
- euclimix-0.1.0/docs/validation.md +61 -0
- euclimix-0.1.0/examples/native.cpp +12 -0
- euclimix-0.1.0/examples/python_api.py +15 -0
- euclimix-0.1.0/pyproject.toml +36 -0
- euclimix-0.1.0/setup.cfg +4 -0
- euclimix-0.1.0/src/euclimix/__init__.py +17 -0
- euclimix-0.1.0/src/euclimix/_api.py +166 -0
- euclimix-0.1.0/src/euclimix/_extension.py +62 -0
- euclimix-0.1.0/src/euclimix/_result.py +46 -0
- euclimix-0.1.0/src/euclimix/_validation.py +56 -0
- euclimix-0.1.0/src/euclimix/csrc/CMakeLists.txt +32 -0
- euclimix-0.1.0/src/euclimix/csrc/common.h +65 -0
- euclimix-0.1.0/src/euclimix/csrc/controlled.cu +142 -0
- euclimix-0.1.0/src/euclimix/csrc/controlled_bindings.cpp +11 -0
- euclimix-0.1.0/src/euclimix/csrc/include/euclimix/distance.h +81 -0
- euclimix-0.1.0/src/euclimix/csrc/native_example.cpp +17 -0
- euclimix-0.1.0/src/euclimix/csrc/paper.cu +150 -0
- euclimix-0.1.0/src/euclimix/csrc/paper_bindings.cpp +10 -0
- euclimix-0.1.0/src/euclimix/csrc/state_example.cpp +43 -0
- euclimix-0.1.0/src/euclimix/py.typed +0 -0
- euclimix-0.1.0/src/euclimix.egg-info/PKG-INFO +215 -0
- euclimix-0.1.0/src/euclimix.egg-info/SOURCES.txt +51 -0
- euclimix-0.1.0/src/euclimix.egg-info/dependency_links.txt +1 -0
- euclimix-0.1.0/src/euclimix.egg-info/requires.txt +10 -0
- euclimix-0.1.0/src/euclimix.egg-info/top_level.txt +1 -0
- euclimix-0.1.0/tests/conftest.py +12 -0
- euclimix-0.1.0/tests/reference/README.md +5 -0
- euclimix-0.1.0/tests/reference/center_kernels.cu +270 -0
- euclimix-0.1.0/tests/reference/center_kernels.cuh +28 -0
- euclimix-0.1.0/tests/reference/euclidean_cuda.cpp +33 -0
- euclimix-0.1.0/tests/reference/euclidean_cuda_advanced.cpp +32 -0
- euclimix-0.1.0/tests/reference/euclidean_cuda_mix.cpp +96 -0
- euclimix-0.1.0/tests/reference/euclidean_cuda_tiled.cpp +66 -0
- euclimix-0.1.0/tests/reference/euclidean_cuda_uniform.cpp +42 -0
- euclimix-0.1.0/tests/reference/euclidean_kernels.cu +152 -0
- euclimix-0.1.0/tests/reference/euclidean_kernels.cuh +32 -0
- euclimix-0.1.0/tests/reference/euclidean_kernels_advanced.cu +460 -0
- euclimix-0.1.0/tests/reference/euclidean_kernels_advanced.cuh +72 -0
- euclimix-0.1.0/tests/reference/euclidean_kernels_mix.cu +1626 -0
- euclimix-0.1.0/tests/reference/euclidean_kernels_mix.cuh +578 -0
- euclimix-0.1.0/tests/reference/euclidean_kernels_tiled.cu +528 -0
- euclimix-0.1.0/tests/reference/euclidean_kernels_tiled.cuh +81 -0
- euclimix-0.1.0/tests/reference/euclidean_kernels_tiled_tc.cu +300 -0
- euclimix-0.1.0/tests/reference/euclidean_kernels_uniform.cu +532 -0
- euclimix-0.1.0/tests/reference/euclidean_kernels_uniform.cuh +57 -0
- euclimix-0.1.0/tests/test_cuda.py +240 -0
- euclimix-0.1.0/tests/test_paper_regression.py +48 -0
- euclimix-0.1.0/tests/test_python.py +308 -0
euclimix-0.1.0/PKG-INFO
ADDED
|
@@ -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
|
+
[](https://github.com/inEXASCALE/euclid-cuda/actions/workflows/cpu.yml)
|
|
20
|
+
[](https://github.com/inEXASCALE/euclid-cuda/actions/workflows/gpu.yml)
|
|
21
|
+
[](docs/validation.md#python-tests)
|
|
22
|
+
|
|
23
|
+
[](pyproject.toml)
|
|
24
|
+
[](https://pytorch.org/)
|
|
25
|
+
[](https://developer.nvidia.com/cuda-toolkit)
|
|
26
|
+
[](#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
|
+
```
|
euclimix-0.1.0/README.md
ADDED
|
@@ -0,0 +1,199 @@
|
|
|
1
|
+
# euclimix
|
|
2
|
+
|
|
3
|
+
[](https://github.com/inEXASCALE/euclid-cuda/actions/workflows/cpu.yml)
|
|
4
|
+
[](https://github.com/inEXASCALE/euclid-cuda/actions/workflows/gpu.yml)
|
|
5
|
+
[](docs/validation.md#python-tests)
|
|
6
|
+
|
|
7
|
+
[](pyproject.toml)
|
|
8
|
+
[](https://pytorch.org/)
|
|
9
|
+
[](https://developer.nvidia.com/cuda-toolkit)
|
|
10
|
+
[](#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
|
+
|