Compute pairwise_distance with matrix multiplications

.github/workflows/build-gpu.yml

@@ -45,15 +45,16 @@ jobs:
           conda info
           conda list
 
-      - name: Installing cudatoolkit and dependencies
+      - name: Install dependencies
         run: |
-          conda install -c nvidia cudatoolkit
-          # Avoid building SciPy from source in CI.
-          uv pip install --only-binary scipy --python "$CONDA_PREFIX/bin/python" -e . --group test
+          # Avoid building SciPy from source in CI, and install the CUDA
+          # libraries that CuPy uses with its ctk extra.
+          uv pip install --only-binary scipy --python "$CONDA_PREFIX/bin/python" -e . --group test "cupy-cuda12x[ctk]"
 
-      - name: Numba Information
+      - name: CuPy information
+        # Fail here if CuPy can't use the GPU, since the GPU tests would be skipped.
         run: |
-          numba -s
+          python -c "import cupy; cupy.show_config(); print(cupy.ones((2, 2)) @ cupy.ones((2, 2)))"
 
       - name: Run GPU tagged tests
         run: |

docs/changelog.rst

@@ -35,6 +35,8 @@ Breaking changes
   Use bio2zarr for conversion to Zarr and vcztools for export.
   See :ref:`plink_conversion_migration`, including companion-file metadata
   differences. (:pr:`1324`, :pr:`1362`, :issue:`1318`)
+- :func:`pairwise_distance` with ``device="gpu"`` now requires CuPy instead of
+  Numba's CUDA support.
 
 Deprecations
 ~~~~~~~~~~~~
@@ -53,8 +55,12 @@ Deprecations
   `vcztools <https://github.com/sgkit-dev/vcztools>`_ packages instead.
   (:user:`tomwhite`, :pr:`1324`, :issue:`1318`)
 
-.. Improvements
-.. ~~~~~~~~~~~~
+Improvements
+~~~~~~~~~~~~
+
+- :func:`pairwise_distance` is computed with matrix multiplications instead of
+  Numba functions, which is several times faster, and arrays with CuPy chunks
+  are computed on the GPU. (:issue:`1267`)
 
 Bug fixes
 ~~~~~~~~~
@@ -70,6 +76,9 @@ Bug fixes
   (:user:`timothymillar`, :pr:`1335`, :issue:`1333`)
 - :func:`Tajimas_D` now works with newer versions of Dask, and no longer
   computes allele counts when it is called.
+- :func:`pairwise_distance` with ``device="gpu"`` no longer gives wrong results
+  for int8 input, and the dtype it reports for integer input now matches the
+  computed result.
 
 .. Documentation
 .. ~~~~~~~~~~~~~

docs/contributing.rst

@@ -99,9 +99,10 @@ the dependencies needed to test compatibility with bio2zarr. Use ``uv run`` to r
 commands in this environment without activating it.
 
 
-If you have a Nvidia GPU you will need to make sure that it is configured properly,
-as in you have cudatoolkit installed, the instructions for the same can be found on
-`nvidia docs. <https://developer.nvidia.com/cuda-toolkit>`_
+The tests marked ``gpu`` run on an NVIDIA GPU with CuPy, and are skipped if CuPy
+or a GPU is not available. To run them, install the CuPy package that matches
+your CUDA version, as described in the
+`CuPy installation guide <https://docs.cupy.dev/en/stable/install.html>`_.
 
 
 Install the pre-commit hook, which is used to enforce coding standards::

sgkit/distance/api.py

@@ -1,9 +1,11 @@
 import typing
 from typing import Literal
 
+import dask
 import dask.array as da
 import numpy as np
 
+from sgkit.distance.metrics import METRICS
 from sgkit.typing import ArrayLike
 
 MetricTypes = Literal["euclidean", "correlation"]
@@ -49,6 +51,12 @@ def pairwise_distance(
     Negative and nan values are considered as missing values. They are ignored
     for all distance metric calculations.
 
+    The distances are computed with matrix multiplications, in the floating
+    point type of ``x``, or in float64 if ``x`` has another type. Since they
+    are computed from sums of squares and products, the distance between two
+    nearly identical rows can be inaccurate in float32. Arrays with CuPy
+    chunks are computed on the GPU.
+
     Parameters
     ----------
     x
@@ -60,13 +68,14 @@ def pairwise_distance(
         'euclidean' or 'correlation'.
     split_every
         Determines the depth of the recursive aggregation in the reduction
-        step. This argument is directly passed to the call to``dask.reduction``
-        function in the reduce step of the map reduce.
+        step over the chunks of the second dimension of ``x``.
 
         Omit to let dask heuristically decide a good default. A default can
         also be set globally with the split_every key in dask.config.
     device
-        The architecture to run the calculation on, either of "cpu" or "gpu"
+        The architecture to run the calculation on, either of "cpu" or "gpu".
+        With "gpu", the chunks of ``x`` are moved to the GPU with CuPy, and
+        the result is moved back to NumPy, so this requires CuPy.
 
     Returns
     -------
@@ -97,106 +106,42 @@ def pairwise_distance(
 
     >>> x = np.array([[6, 4, 1,], [4, 5, 2], [9, 7, 3]])
     >>> pairwise_distance(x, metric='correlation').compute()
-    array([[-4.44089210e-16,  2.62956526e-01,  2.82353505e-03],
-           [ 2.62956526e-01,  0.00000000e+00,  2.14285714e-01],
-           [ 2.82353505e-03,  2.14285714e-01,  0.00000000e+00]])
+    array([[0.        , 0.26295653, 0.00282354],
+           [0.26295653, 0.        , 0.21428571],
+           [0.00282354, 0.21428571, 0.        ]])
     """
-    from sgkit.distance import metrics  # defer numba compilation
-
     valid_devices = DeviceTypes.__args__  # type: ignore[attr-defined]
     if device not in valid_devices:
         raise ValueError(
             f"Invalid Device, expected one of {valid_devices}, got: {device}"
         )
-    try:
-        map_func_name = f"{metric}_map_{device}"
-        reduce_func_name = f"{metric}_reduce_{device}"
-        map_func = getattr(metrics, map_func_name)
-        reduce_func = getattr(metrics, reduce_func_name)
-        n_map_param = metrics.N_MAP_PARAM[metric]
-    except AttributeError:
-        raise NotImplementedError(
-            f"Given metric: '{metric}' is not implemented for '{device}'."
-        )
+    if metric not in METRICS:
+        raise NotImplementedError(f"Given metric: '{metric}' is not implemented.")
 
     x = da.asarray(x)
     if x.ndim != 2:
         raise ValueError(f"2-dimensional array expected, got '{x.ndim}'")
-
-    # setting this variable outside of _pairwise to avoid it's recreation
-    # in every iteration, which eventually leads to increase in dask
-    # graph serialisation/deserialisation time significantly
-    metric_param = np.empty(n_map_param, dtype=x.dtype)
-
-    def _pairwise_cpu(f: ArrayLike, g: ArrayLike) -> ArrayLike:
-        result: ArrayLike = map_func(f[:, None, :], g, metric_param)
-        # Adding a new axis to help combine chunks along this axis in the
-        # reduction step (see the _aggregate and _combine functions below).
-        return result[..., np.newaxis]
-
-    def _pairwise_gpu(f: ArrayLike, g: ArrayLike) -> ArrayLike:  # pragma: no cover
-        result = map_func(f, g)
-        return result[..., np.newaxis]
-
-    pairwise_func = _pairwise_cpu
-    if device == "gpu":
-        pairwise_func = _pairwise_gpu  # pragma: no cover
-
-    # concatenate in blockwise leads to high memory footprints, so instead
-    # we perform blockwise without contraction followed by reduction.
-    # More about this issue: https://github.com/dask/dask/issues/6874
-    out = da.blockwise(
-        pairwise_func,
-        "ijk",
-        x,
-        "ik",
-        x,
-        "jk",
-        dtype=x.dtype,
-        concatenate=False,
-    )
-
-    def _aggregate(x_chunk: ArrayLike, **_: typing.Any) -> ArrayLike:
-        """Last function to be executed when resolving the dask graph,
-        producing the final output. It is always invoked, even when the reduced
-        Array counts a single chunk along the reduced axes.
-
-        Parameters
-        ----------
-        x_chunk
-            [array-like, shape: (M, M, N, 1)]
-            An array like two dimensional matrix. The dimension is as follows:
-            M is the chunk size on axis=0 for `x` (input for `pairwise_distance`
-            function).
-            N: is the number of chunks along axis=1
-        """
-        x_chunk = x_chunk.reshape(x_chunk.shape[:-2] + (-1, n_map_param))
-        result: ArrayLike = reduce_func(x_chunk)
-        return result
-
-    def _chunk(x_chunk: ArrayLike, **_: typing.Any) -> ArrayLike:
-        return x_chunk
-
-    def _combine(x_chunk: ArrayLike, **_: typing.Any) -> ArrayLike:
-        """Function used for intermediate recursive aggregation (see
-        split_every argument to ``da.reduction below``).  If the
-        reduction can be performed in less than 3 steps, it will
-        not be invoked at all."""
-        # reduce chunks by summing along the -2 axis
-        x_chunk_reshaped = x_chunk.reshape(x_chunk.shape[:-2] + (-1, n_map_param))
-        return x_chunk_reshaped.sum(axis=-2)[..., np.newaxis]
-
-    r = da.reduction(
-        out,
-        chunk=_chunk,
-        combine=_combine,
-        aggregate=_aggregate,
-        axis=-1,
-        dtype=x.dtype,
-        meta=np.ndarray((0, 0), dtype=x.dtype),
-        split_every=split_every,
-        name="pairwise",
-    )
-
-    t = da.triu(r)
-    return t + t.T
+    if device == "gpu":  # pragma: no cover
+        import cupy
+
+        x = x.map_blocks(cupy.asarray)
+
+    dtype = x.dtype if x.dtype.kind == "f" else np.dtype(np.float64)
+    observed = x >= 0  # false for negative and nan values
+    # zero the missing values so they drop out of the matrix products in the
+    # metrics
+    x = da.where(observed, x, 0).astype(dtype)
+    observed = observed.astype(dtype)
+    # split_every controls the sums over the column chunks inside da.matmul
+    config = {} if split_every is None else {"split_every": split_every}
+    with dask.config.set(config):
+        d = METRICS[metric](x, observed)
+
+    # make the result exactly symmetric, with zeros on the diagonal (or nan
+    # where the distance of a vector to itself is undefined)
+    u = da.triu(d, 1)
+    diagonal = da.where(da.isnan(da.diagonal(d)), np.nan, 0).astype(dtype)
+    d = u + u.T + da.diag(diagonal)
+    if device == "gpu":  # pragma: no cover
+        d = d.map_blocks(cupy.asnumpy)
+    return d