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:
I picked the kmeans example/function linked in the top comment. The only reason for that is that it was easy, I didn't think about whether it is the best example or not. This means if you have a preference/opinion, lets switch to that.

My goal was to write some code based on pluggy to be able to see it in action. I wanted to have a "base plugin" with the default sklearn implementation and then one "other plugin" (the fake cupy plugin in this case). I wanted two plugins in order to see what would happen if one of them wanted to pass on the computation or if it only implements a subset of the hooks. I haven't tried to find an optimal place for all the code to live, I think it feels a bit clunky right now. I haven't thought about performance either.

Diff of [the first commit](https://github.com/scikit-learn/scikit-learn/compare/main...81e267957325935a29ebd3d835bd874505c5a048) ([diff of the branch](https://github.com/scikit-learn/scikit-learn/compare/main...betatim:scikit-learn:pluggy-prototype-computational-routines), might change as I work on it).

I used [the bisecting kmeans example](https://scikit-learn.org/stable/modules/generated/sklearn.cluster.BisectingKMeans.html) to run the code.

I haven't really thought about how kmeans works, but it feels like the hook is in the wrong place now. Maybe it should be "higher level", maybe somewhere in `BisectingKMeans.fit()`? I'll leave it as it is for now though.

Related to "where should the hook be?": I think we should have the hooks "high enough" so that the overhead of picking a hook and calling it doesn't matter. This means we shouldn't hook a function that only takes 50ns to execute but is called 1000 times. Instead the hook should be "above" the loop that is making the 1000 calls. Ideally you call a hook with enough work so that it would take a few seconds or more to execute.

I like that this example allows the hook implementations to decide if they want to "take the call or not". For example a GPU plugin could decide that it will take action if the input data is already in the GPU memory or if the input data is not on the GPU memory yet but it is "big enough" so that transferring it to GPU memory is worth it. Basically, I think the decision of which hook implementation to use is more complicated than just the datatype of the input data.

Unlike in my hack above, all the "go faster" plugins should be their own Python package that people install.

I like that a plugin is "activated" by `pip install`ing it. It seems like as a user you'd install the plugin to get "faster computations with the magic hardware I have" and want to use it all the time. Even nicer is that I think we can use `set_/get_config()` together with `plugins.manager.subset_hook_caller("kmeans_single_lloyd", list_of_active_plugins)` to temporarily disable a plugin. Instead of calling `plugins.manager.hook.kmeans_single_lloyd()` directly, you use `subset_hook_caller()` together with a list of active plugins based on `get_config()`.

This is a super long comment already, so I'll stop now and wait to hear what you think 😃

Labels: API, Performance