Path for pluggable low-level computational routines

The goal of this issue is to discuss the design and prototype a way to register alternative implementations for core low level routines in scikit-learn, in particular to benefit from hardware optimized implementations (e.g. using GPUs efficiently).

## Motivation

scikit-learn aims to provide reasonably easy to maintain and portable implementations of standard machine learning algorithms. Those implementations are typically written in Python (with the help of NumPy and SciPy) or in Cython when the overhead of the Python interpreter prevents us to efficiently implement algorithms with (nested) tight loops. This allows us reasonably fast implementations as binary packages (installable with pip/PyPI, conda/conda-forge, conda/anaconda or Linux distros) for a variety of platforms (Linux / macOS / Windows) x (x86_64, i686, arm64, ppc64le) from a single code-base with no external runtime dependencies beyond Python, NumPy and SciPy.

Recently, GPU hardware have proven very competitive for many machine learning related workloads, either from a pure latency standpoint, or from the standpoint of a better computation/energy trade-off (irrespective of the raw speed considerations). However, hardware optimized implementation are typically not portable and mandates additional dependencies.

We therefore propose to design a way for our users to register alternative implementations of low-level computation routines in scikit-learn provided they has installed the required extension package(s) that matches their specific hardware.

## Relationship to adopting the Array API spec

This proposal is related and complementary to another effort, namely:

- #22352

The Array API spec is useful to make it possible to have some scikit-learn estimators written using a pure numpy syntax to delegate the computation to alternative Array API compatible libraries such as CuPy.

However, some algorithms in scikit-learn cannot be efficiently written using NumPy operation only, for instance the main [K-Means loop is written in Cython](https://github.com/scikit-learn/scikit-learn/blob/0de3b0d1eaacee9f7b15cabc05752cba945c7621/sklearn/cluster/_k_means_lloyd.pyx#L116-L152) to process chunks of samples in parallel (using `prange` and OpenMP), compute distance with centroids and reduce those distances to find assign each sample to its closest centroid on-the-fly while preventing unnecessary memory transfer between CPU cache and RAM.

If we want to run this algorithm efficiently on GPU hardware, one would need to dispatch the computation of this low level function to an alternative implementation that can work on GPU, either written in C/C++ with GPU-specific supporting runtime libraries and compilers (e.g. OpenCL, NVIDIA Cuda, Intel oneAPI DPC++, AMD ROCm...)  or using a Python syntax with the help of GPU support provided in numba for instance.

## List of candidate routines

- the main k-means loop (pairwise distances between samples and centroids with on-the-fly argmin reduction) 
- the core k-nearest neighbors computation loop (pairwise distances with on-the-fly arg-k-min reduction) 
- pairwise distances (without arg-k-min reduction)
- pairwise kernels computation (e.g. for use in Nystroem)
- ...

## Explicit registration API design ideas

I started to draft some API ideas in:

- https://hackmd.io/@4rHCRgfySZSdd5eMtfUJiA/S1RDJ3HCF

Feel free to comment here, or there.

This design is subject to evolve, in particular to make it possible to register both Array API extensions and non-Array API extensions with the same registration API.

## Next steps

- start a proof of concept implementation in a branch + external repo using either numba or C++ for k-means. Edit: this is happening here:
  - #25535
  - https://github.com/soda-inria/sklearn-numba-dpex
- once the API design has converged on the PoC, write a formal SLEP.

COMMENT:
Thanks for participating in this conversation, @betatim.

> I like this way of thinking about this. In particular because you want to replace more than just one method (fit()) but might also want to replace preprocessing stuff. So it feels more like "install a plugin sklearn-XYZ-foobar to get speed ups for estimator XYZ if you have a FooBar accelerator". The plugin then uses predefined hooks to replace/augment(?) the code that is already in scikit-learn. Or maybe there will just be a plugin per accelerator.

This probably is not (yet) explicit, but I think we somewhat plan to have a grosser granularity, i.e. have a plugin per accelerator/vendor (`foobar`) that would ship implementations for several methods/functions (`XYZ`).

> I think there must be prior art for this kind of stuff out there that we could learn from. And maybe it is very unlikely that someone has more than one accelerator in the same computer/would want to disable acceleration for some estimators but not others?

I think also it is appropriate to consider some UX regarding plugins (beforehand). Do you know projects that are using pluggy, so as to get learn from their experience? I will try to search for such projects.

Labels: API, Performance