Repository navigation
Compute pairwise_distance with matrix multiplications - #1365
Billyzhang1229 wants to merge 1 commit into
Conversation
81711c4 to
4c7f998
Compare
jeromekelleher
left a comment
jeromekelleher
left a comment
There was a problem hiding this comment.
Looks like a straightforward win to me - nice job @Billyzhang1229 !
Write both metrics as matrix products of the values, with missing entries set to zero, and the mask of observed entries. For example, the squared Euclidean distance over the entries observed in both rows is sum(x_i^2) + sum(x_j^2) - 2 * sum(x_i * x_j), where each sum is a matrix product. Build them from Dask array operations instead of Numba gufuncs, CUDA kernels and hand-written blockwise and reduction steps, with each metric as a function in sgkit.distance.metrics, looked up by name in pairwise_distance, so that new metrics can be added there. This is several times faster on the CPU, and no longer passes any functions for Dask to tokenize (see issue 1267). The GPU path now uses CuPy instead of Numba's CUDA support: arrays with CuPy chunks are computed on the GPU, and device="gpu" moves the chunks to the GPU and the result back to NumPy. This fixes wrong results for int8 input on the GPU, where the partial sums were stored in an int8 array. The GPU workflow now installs CuPy for CUDA 12, with the CUDA libraries it uses from PyPI, instead of the CUDA toolkit from conda, which needs NVIDIA driver 525 or later on the runner, and checks that CuPy can use the GPU before running the tests, which are skipped otherwise. The distances are computed in the floating point type of the input, or float64 for other types, and integer input now reports the float64 dtype of the result before it is computed. The diagonal is now exactly zero, or nan where the correlation is undefined. This needs Dask 2024.12 or later: the blockwise fusion in earlier releases computed wrong correlations from these expressions.
4c7f998 to
7615957
Compare
|
Thanks, there is something goes wrong about floating point precision (when numbers are small). I'm checking these. |
|
Looks like all tests are passing though? |
The existing tests don't cover this. The old code computed the euclidean distance from |
|
That sounds good. Perfectly fine to just say "possible loss of precision with float32" too though. |
pairwise_distanceis the last function affected by #1267. It passes locally defined functions toda.blockwiseandda.reduction, and newer Dask tokenizes them by pickling, so the Numba gufuncs they reference are rebuilt and recompiled on every call (comment). #1275 fixed the other functions, but this one needed a different approach.This builds on #1364, which removes the Dask pin. Its two commits are included here, so only the last commit is new.
This PR:
sum(x_i^2) + sum(x_j^2) - 2 * sum(x_i * x_j), where each sum is a matrix product.sgkit.distance.metrics, looked up by name inpairwise_distance, so new metrics can be added there.device="gpu"moves the chunks to the GPU and the result back to NumPy. This fixes wrong results for int8 input on the GPU, where the partial sums were stored in an int8 array.Other changes in behavior:
device="gpu"requires CuPy instead of Numba's CUDA support.With this, sgkit uses CuPy for
pairwise_distanceon the GPU, so the CuPy mention that #579 questioned now applies.Fixes #1267.
Testing:
GUFunc._rebuildwhile running the full test suite with Dask 2026.8.0: 98 before this change, all intest_distance.py, and none after.cupy-cuda12x[ctk], the package the workflow installs. It gave the same results as the current code on all three, but it would still be good to dispatch the workflow to run the GPU tests on its T4.Benchmarks
pairwise_distancebefore and after this PR, on an Intel i9-12900K (16 threads under WSL2) and an NVIDIA RTX 5090. The times before don't include the recompilation from #1267, so the speedups come from the new computation alone:The speedup grows with the input. On the CPU at 4k × 40k, euclidean goes from 75 s to 9.3 s (8.0×) and correlation from 311 s to 17 s (18×). On the GPU at 8k × 80k, float32 is 21× and 16× faster. For the smallest input, whose 16 chunks are 125 × 1250, overhead dominates: euclidean on the CPU is only 1.2× faster, and correlation on the GPU in float32 is slower (112 ms before, 174 ms after), with 1101 tasks in the new graph instead of 241.
In float64 the RTX 5090 is only 2.2× and 2.4× faster at 8k × 80k. It computes float64 at 1/64 of its float32 rate, and the new code spends its time in matrix products: an 8192 × 8192 product with CuPy runs at 1.7 TFLOPS in float64 and 76 TFLOPS in float32. To check that this comes from the GPU and not from the code, I ran the GPU benchmarks on an A100 and an H100 on Modal, whose float64 matrix products are as fast as their float32 ones (18.6 and 58.5 TFLOPS). There, float64 at 8k × 80k is 23× and 15× faster on the A100, and 31× and 28× faster on the H100:
The CPU times are with default settings, where Dask's threaded scheduler runs a task on every core and OpenBLAS also uses every core for each matrix product. With
OPENBLAS_NUM_THREADS=1, which Dask's distributed workers set by default, the new code is faster again and varies much less between runs: 4k × 40k takes 5.0 s (15×) and 9.6 s (33×). The old code doesn't use BLAS, and its times barely change (within 3% at 2k × 20k).Setup and all timings
pairwise_distance(x, metric=metric, device=device).compute(), including building the graph, with Dask's threaded scheduler. Each time is the median of 2 to 23 runs, after a warm-up call on a small array that compiles the Numba functions and the CUDA and CuPy kernels. On the GPU the input starts in NumPy chunks, so the times include copying the chunks to the GPU and the result back. The first call on a large input in a new process is slower than these medians while CuPy's memory pool grows, for example 1.7 s instead of 0.71 s for euclidean in float32 at 8k × 80k.cupy-cuda12x[ctk],numba-cuda[cu12]), since the driver is older than the local CUDA 13.4 wheels. Dask used 16 threads, as on the local machine. The A100 is the SXM4 80GB model and the H100 the SXM5 80GB model. The results of both versions agree there too.CPU, float64:
CPU, float64, with
OPENBLAS_NUM_THREADS=1for the new code (speedup over before in parentheses):Other input types on the CPU, at 2k × 20k:
Matrix product throughput, 8192 × 8192 with CuPy (no TF32):
GPU times on the three GPUs, before → after (speedup):
float32, euclidean:
float32, correlation:
float64, euclidean:
float64, correlation: