META: Roadmap for the future of xsf #220

I'm collecting ideas here for the future direction of xsf that are currently spread out between various issues, pull request comments, and other discussions. I know this is very long but I think it's important to have all of this somewhere in one place. I'd appreciate any feedback that can be offered.

What should xsf be?

An important matter is deciding kind of code should live in xsf. xsf was conceived as a home for mathematical special function implementations written in C++ which work on both CPU and GPU for use as a dependency of both SciPy and CuPy. This idea originated with Irwin Zaid (@izaid) who also came up with the name xsf¹ and played a major role in getting the library off the ground. Standardizing on the use of C++ for SciPy special's ufunc and gufunc kernels also allows for simplifying SciPy's build processes, soxsf has also become a home for many of the kernels behind SciPy's special functions even if they cannot be made to work on GPU in a straightforward way. The end state we seem to be working towards is that all kernels for backing SciPy's special functions are written in C++ and live either in xsf or Boost.Math. Boost.Math is also using the same techniques xsf is using to allow kernels to work on GPU when possible.

Why does xsf need to exist when there is already Boost Math?

One question I've commonly heard is: "why can't all of this C++ simply live in Boost.Math?". Some ready answers are that:

• We are free to use C++17 in xsf, but Boost.Math must still support C++14.
• We are free to support only double and float inputs in xsf without concern for extended precision and arbitrary precision types.
• The scipy.special codebase contained a lot of legacy special function implementations that were not up to the standards of Boost.Math, but were still serviceable for many input domains. It was faster to translate and add things to xsf because the expected standards of quality are inherited from the existing scipy.special codebase.

One could still imagine working towards an asymptotic state where the implementations in xsf eventually find their way into Boost.Math. At times I had this in mind too, but I think there is a reduction in scope we are free to make inxsf that opens up an important niche for this emerging library:

• xsf can focus purely on providing kernels for universal functions (ufuncs) and generalized universal functions (gufuncs) without concern for producing standalone function implementations for use in C++ programs.

This may not appear like too meaningful of a difference, but actually is for reasons that will become clear below.

Ufuncs and Gufuncs with in-loop caching to eliminate redundant computation

Along with this reduction in scope, I envision a corresponding expansion in scope. xsf need not be restricted to just special function implementations, but to kernels for any universal functions and generalized universal functions which may be useful in SciPy. It may appear like there is scarcely any difference between these two classes of things. Ufuncs and gufuncs seem to suffer from a serious limitation that has largely restricted their use in SciPy to special and some internal special functions used in stats: when iterating over the elements or core slices of the inputs, each core calculation must be completely independent of all the others. This means, for example, that if there are two input arguments p and k and an expensive to compute intermediate array or scalar that is generated based on each core slice of p, this intermediate result will need to be recomputed again and again across iterations over k.

This situation is not uncommon in scipy.special, and the corresponding (and sometimes massive) degradation in performance has been seen as acceptable for relatively niche special functions. These wasteful computations appear fatal, however, to attempts to use the ufunc/gufunc paradigm more widely. Fortunately, in scipy/scipy#24894 and scipy/scipy#25128 I've found a workaround for this limitation. It is possible to cache intermediate results within ufunc/gufunc loops and work with transposes of the input and output arrays in such a way that the iteration order proceeds in optimal fashion for cache reuse. The inputs needed to generate the cached intermediate results can be made to vary most slowly, with the other inputs varying more quickly. With this optimization, I was able to obtain up to a 1000x speedup in practical workflows involving computation of Mathieu functions. The general idea is something I had asked Irwin about early on in our work, but which I have only recently been able to accomplish. @WarrenWeckesser had also suggested something along these lines here, scipy/scipy#22080 (comment). Taking a look at scipy/scipy#25128 and #144 is currently probably the best way to see how this in ufunc caching + iteration order control method works for a nontrivial example.

By removing this limitation, it's possible to significantly expand the class of things that can be efficiently computed as a ufunc or gufunc.

Gufuncs on GPU

For problems which can be efficiently computed as ufuncs or gufuncs on CPU using this pattern of caching intermediate results and controlling iteration order to ensure optimal cache behavior, I propose on GPU to use a pattern of composition of expansion and reduction gufuncs that was originally suggested by Irwin Zaid. At the moment, there's no straightforward way to implement a gufunc in CuPy which allows for batchwise parallelization of an inherently serial core computation, but I have a PR, cupy/cupy#9959, which adds support for inputs and outputs with nonzero core dimensionality to CuPy's ElementwiseKernel. With this, I was able to achieve over a 100x speedup over the optimized CPU implementations of the Poisson Binomial distribution functions mentioned above, which had already achieved a 250x speedup over the unoptimized CPU implementations from before I worked out the in-ufunc caching strategy.

An example of the expansion-reduction pattern.

Here is an example of how this expansion-reduction pattern works.

The Poisson-Binomial distribution is the distribution of the number of successes in $n$ bernoulli trials, where the success probability $p_i$ can vary between the trials. The pmf must be computed with an inherently serial recurrence relation for all $k$ in the support $\left\{0, 1, \ldots, n\right\}$.

"""Assorted statistical functions"""
from __future__ import annotations

import cupy
from cupy._core import ElementwiseKernel

preamble = """
#include <cupy/xsf/stats.h>
"""

# The gufunc `_poisson_binom_pmf_all` uses a recurrence relation
# to compute the full pmf for each 1d core slice of the input `p`.
# `in_params="T(n) p"` tells `ElementwiseKernel` that the core slices of `p` 
# are 1-dimensional of length n (containing the `n` success probabilities) and
# `out_params="T(n+1) out"`tells `ElementwiseKernel` that the corresponding
# core slices of the output are also 1 dimensional and of length n + 1.
# This is the expansion gufunc.
_poisson_binom_pmf_all = ElementwiseKernel(
    in_params="T(n) p",
    out_params="T(n+1) out",
    operation=(
        "xsf::poisson_binom_pmf_all(p.as_mdspan(), out.as_mdspan());"
    ),
    name="cupy_poisson_binom_pmf_all",
    preamble=preamble,
)


# This is the reduction gufunc. It selects the desired entries of the
# full pmf corresponding to integers in the array `k`.
_take_from_pmf = ElementwiseKernel(
    in_params="T(n) pmf, int64 k",
    out_params="T out",
    operation="out = xsf::take_from_pmf(pmf.as_mdspan(), k);",
    name="cupy_take_from_pmf",
    preamble=preamble,
)


# The full implementation is a composition of the expansion and the reduction.
def _poisson_binom_pmf(k, p):
    """Returns pmf of Poisson Binomial distribution.

    Parameters
    ----------
    k : array
        Number of successes at which to evaluate pmf.

    p : array
        Success probabilities of independent Bernoulli trials.

    Notes
    -----
    This is equivalent to a gufunc with signature ``()(i)->()``.
    The last dimension of `p` contains success probabilities and
    the preceding dimensions are batch dimensions. The batch
    dimensions are broadcast against ``k``.
    """
    out_shape = p.shape[:-1] + (p.shape[-1] + 1,)
    # By using Fortran style row major contiguity for the intermediate
    # array that contains full pmfs for each core slice of `p`, we unlock
    # significant performance gains by allowing CUDA threads in the same
    # warp to write to adjacent memory locations.
    pmf = cupy.empty(out_shape, dtype=p.dtype, order='F')
    _poisson_binom_pmf_all(p, pmf)
    return _take_from_pmf(pmf, k)

Code sharing

I chose the Poisson Binomial distribution functions as the first cases to tackle using this pair of strategies for developing efficient CPU and GPU implementation because they were the simplest nontrivial relevant functions I could find. These dual strategies should also work well for more complicated examples as well. I see the greatest strength here as the possibility of relying on the same underlying numerical routines on both CPU and GPU, with only the method of orchestration differing in order to exploit the strengths and weaknesses of the respective hardware architectures. Below is an example of what SciPy's kernel looks like for the optimized CPU implementation of the Poisson Binomial pmf. This is taken from https://github.com/scipy/scipy/blob/ea194aeb77d8e7f1b1c4221ef8e086c3de7e056b/scipy/special/_gufuncs.cpp#L72-L108.

/* gufunc kernels which use internal caches.
 *
 * Such caches live only during the course of a single call to the ufunc
 * in order to eliminate redundant calculations. This does not concern
 * user-facing caching across ufunc calls.
 *
 * In order for such caches to serve their purpose, the iteration order over the
 * course of the ufunc loops must be chosen so that elements (or core slices) of
 * the inputs that go into the cached computation vary most slowly and elements
 * of the inputs for which the cached result can be reused vary most quickly.
 * This can be accomplished by wrapping the gufunc with `_with_cache_optimization`
 * from ./_ufunc_tools.py. See the docstring of `_with_cache_optimization` for
 * more information.
 */

template <typename T_1d>
struct _poisson_binom_pmf_kernel {
    using T = typename T_1d::value_type;

    std::vector<T> dist;
    T* last_p_ptr = nullptr;
    T operator()(long long int k, T_1d p) {
        if (!last_p_ptr) {
            dist.resize(p.extent(0) + 1);
        }
        if (p.data_handle() != last_p_ptr) {
            /* If p is stepped through in the outer loops and k in the inner loops,
             * this will yield a cache hit for each inner iteration. The cache is
             * overwritten unconditionally whenever this steps to a new slice of p. */
            xsf::poisson_binom_pmf_all(p, _as_mdspan(dist));
            last_p_ptr = p.data_handle();
        }

        return xsf::take_from_pmf(_as_mdspan(dist), k);
    }
};

Both SciPy's CPU implementation and CuPy's GPU implementation rely on the core functions xsf::poisson_binom_pmf_all and xsf::take_from_pmf. CuPy uses them directly as gufunc kernels and performs and expansion followed by a reduction. SciPy applies the pattern of in-gufunc caching and uses the two core functions together in a single fused gufunc kernel. It appears that there should be a straightforward mapping between such CPU-friendly implementations and GPU-friendly implementations by means of which one can be mechanically converted to the other. This allows us to essentially extend xsf's ability to allow the same implementation to be used on CPU as on GPU for scalar ufuncs to gufuncs as well, since the challenging numerical components will be shared, and the surrounding orchestration will follow standard patterns that I hope to document and make increasingly legible over time, both for human and LLM-based developers.

Working within GPU hardware limitations

I've encountered skepticism that the functions I seek to add to CuPy with the expansion-reduction gufunc pattern can be meaningfully accelerated with GPU hardware as well as skepticism that even branch-heavy elementwise special functions can be meaningfully accelerated. A common refrain is "if this is feasible, why has no one done it yet?". The answer appears to be that xsf has been made feasible by advances in GPU technology which have been made over the past ten years. When Irwin first proposed xsf I was also surprised that computing branch heavy special functions like hyp2f1 and gdtrib can be meaningfully accelerated because I had an outdated mental model of GPU capabilities.

The first limitation we run into relates to branching and divergence. GPUs employ a single instruction multiple thread (SIMT) model, where threads are lumped into groups called warps. Each thread in a single warp iterates through a single sequence of instructions, but unlike in the single instruction multiple data (SIMD) model, conditional branches can be handled within the instructions executing within a warp. This leads to what is known as warp divergence, and does not impact correctness, but can impact performance.

When there is warp divergence, say if half of the threads in a warp take path A and half take path B, then what can occur is that first the threads taking path A execute, while the other threads wait, then the threads taking path B execute, typically with some schedule overhead involved. Note however that the CUDA C best practices guide states that

For branches including just a few instructions, warp divergence generally results in marginal performance losses. For example, the compiler may use predication to avoid an actual branch. Instead, all instructions are scheduled, but a per-thread condition code or predicate controls which threads execute the instructions. Threads with a false predicate do not write results, and also do not evaluate addresses or read operands.

branch predication eliminates the overhead of warp divergence and allows the threads in a diverging warp to continue executing instructions in lock step, although some clock cycles will be spent on no-ops. Branch predication is the main reason why I was able to get a significant GPU speedup computing gdtrib with the branch heavy Chandrupatla's root finding algorithm here.

Branch predication has been around since the introduction of the Tesla architecture in 2006, but has improved over time. Earlier versions of the CUDA best practices guide, such as https://docs.nvidia.com/cuda/archive/8.0/cuda-c-best-practices-guide/index.html#control-flow, refer to a hard limit on the number of instructions that can be predicated over:

The compiler replaces a branch instruction with predicated instructions only if the number of instructions controlled by the branch condition is less than or equal to a certain threshold: If the compiler determines that the condition is likely to produce many divergent warps, this threshold is 7; otherwise it is 4.

But this hard limit has been removed, and more recent CUDA toolkits have an optimizer that can better identify situations where predication can be beneficial. The workings of this optimizer are a proprietary detail of NVIDIA's ptxas assembler so I'm not sure exactly what it's doing, but through applying NVIDIA's nsight compute profiler to CuPy kernels backed by xsf, I've found it to do an admirable job optimizing branchy xsf code.

The Volta architecture, released in 2017, introduced independent thread scheduling giving each thread in a warp its own independent thread counter. Although xsf does not benefit from independent thread scheduling's primary use-case of allowing sharing of data between diverging threads in a warp without deadlocking, the CUDA C programming guide also mentions that independent thread scheduling allows for better use of execution resources:

With independent thread scheduling, the GPU maintains execution state per thread, including a program counter and call stack, and can yield execution at a per-thread granularity, either to make better use of execution resources or to allow one thread to wait for data to be produced by another.

This benefit is nebulously worded, but I think it refers to allowing for improved latency hiding, which is valuable for the memory intensive expansion/reduction gufunc pattern. The expansion/reduction pattern also benefits from massive improvements in memory bandwidth from Volta, to Ampere, to Hopper, to Blackwell class GPUs driven by the demands of machine learning.

Another criticism I have received is based on the common wisdom that one should employ kernel fusion to develop optimal GPU algorithms, but the gufunc expansion/reduction pattern seems to perversely call for kernel fission. For some, I think the pattern rings alarm bells; it sounds too good to be true and must be making significant sacrifices in performance for the sake of its elegant code sharing model. I think this gut feeling that something must be off is based on a mental model that is thinking in terms of the acceleration of a single large problem such as matrix multiplication. What we're interested in for xsf though is parallelization of batches of many small instances of an inherently serial problem where there are intermediate results that are shared between some of the instances. It seems every time I show a fresh LLM my Poisson-Binomial distribution CuPy examples, it suggests I fuse the expansion and reduction steps, even though in reality its suggestion significantly harms performance by forcing thousands of redundant $O(n^2)$ computations of the full pmf. Due to the increasingly common pattern of human cognitive surrender to LLMs, there is a risk that this attitude will be amplified and passed on to humans as well. I think the benchmark results in posted in cupy/cupy#9959 (comment) compared to the CPU-only results from cupy/cupy#9959 (comment) speak for themselves and show that the expansion/reduction pattern actually does work well.

When I first began working on xsf, I didn't know any of this. I just trusted that Irwin knew what he was talking about. It seems that xsf is happening now and not earlier because the time is ripe for it. The rapid pace of improvement in GPU technology that has enabled the transformer revolution has also enabled the rise of xsf.

What functions can we hope to tackle?

Bessel functions

Evaluation of Bessel functions of high order is typically accomplished with recurrence relations which reduce the problem to computing Bessel functions of lower order ². Such recurrence relations can be used for integer and noninteger orders (because Bessel functions for non-integer order of sufficiently small magnitude can be evaluated directly), but for integer orders this naturally falls under the expansion/reduction pattern. This obviously applies to UFuncs scipy.special for Bessel functions which restrict the order to being an integer, such as yn, kn, spherical_jn, spherical_yn, spherical_in, and spherical_kn, but we could also support
integer case for other Bessel functions.

It's also planned to implement _all versions of the Bessel functions which can return sequences of Bessel function values of consecutive orders, such as in
#181. These could all naturally be added as gufuncs in CuPy.

Gufunc expansion/reduction implementations of Bessel functions of non-integer order could be implemented and would work from a standpoint of correctness, but
I'm actually not sure if the GPU would offer meaningful speedup for these. Fortunately, it would take very little additional beyond the implementation of
the _all versions to try this, so there will not be much cost to trying it out to see if it works.

Hypergeometric functions

While already possible to compute hypergeometric functions with shape parameters (e.g. a, b, c for hyp2f1) of small magnitude on GPU using a CuPy ElementwiseKernel, but the case of shape parameters of large magnitude poses similar challenges to the case of Bessel functions of large order. In both cases one must rely on argument reduction through a backwards recurrence algorithm. The same caveats listed above for Bessel functions also apply here.

Mathieu functions

Thanks to the work of @brorson, with scipy/scipy#25128 and #144, we're actually most of the way to getting GPU accelerated Mathieu functions. The primary blocker now is that we'd need to write a gufunc kernel for a symmetric tridiagonal eigenvector/eigenvalue solver; this is feasible and I will have more to say about linear algebra below. Handling modified Mathieu functions will also require the _all versions Bessel functions of the first and second kind and their derivatives which is also feasible.

Legendre functions

Irwin already did a lot of work to modernize most of these, and they could naturally be added as gufuncs to CuPy.

Other special functions giving solutions of differential equations

The ellipsoidal harmonics, spheroidal wave functions, parabolic cylinder functions, could all be handled, the first two cases being similar to the Mathieu functions, and the later being computable in terms of hypergeometric functions, or directly through a recurrence. A lot of work could be spent adding support for everything in special. What I'd like to do is clearly document the processes for writing single-source kernels on CPU and GPU explained above, making them legible to community developers and agent-based workflows, to remove barriers towards the community filling remaining gaps in support.

Statistical functions

Beyond the Poisson-Binomial distribution functions, #98, gives a number of special functions that would be helpful for scipy.stats. Of these, Kendall's Tau,
Wilcoxon Distribution functions (for which gufunc kernels have already been written by @fbourgey), the distribution of the Mann Whitney U statistic, the distribution of the Ansari-Bradley W statistic, distributions of Noncentral Hypergeometric distributions, the Cramer Von Mises two sample distributions, and the Kolmogorov-Smirnov two sample distributions can all make use of the strategy
described above. The one-sample Kolmogorov-Smirnov distribution can be computed in elementwise fashion. I believe that the multivariate normal and t distributions,
(besides the bivariate case which has already been handled in elementwise fashion) can be managed, but would need to do some research to figure out exactly how.

Linear algebra

One might expect linear algebra to be very well covered by existing libraries like cuSOLVER, and that if something is not available in an existing library, it likely isn't feasible to accelerate. For parallelization of batches of an inherently serial problem involving small matrices, there
actually is a lot we could do that would be highly impactful. I first noticed this when looking into Mathieu functions and trying to figure out how to handle their need for solving eigenvector and eigenvalue problems for symmetric tridiagonal matrices. I found that the Implicit QR algorithm with Wilkinson shifts (algorithm 8.3.2 of Golub and Van Loan's Matrix Computations seems like it should work perfectly fine in a gufunc kernel, and that branch predication should mitigate the warp divergence if the matrices aren't too big. Beyond Mathieu functions, eigh_tridiagonal is also useful for Ellipsoidal Harmonics, Spherodial Wave Functions, Discrete Prolate Spheroidal Sequences (dspss in scipy.signal.windows).

This lead me to start thinking of what other linear algebra functions we could handle for batches of small matrices. While going through all the functions in SciPy's public API to assess their potential for alternative backend support, I found many functions that are dependent on the Schur decomposition scipy.linalg.schur, this includes matrix transformations like scipy.linalg.logm, scipy.linalg.fractional_matrix_power, scipy.linalg.sqrtm, and scipy.linalg.funm; and matrix equation solvers like scipy.linalg.solve_continuous_lyapunov, scipy.linalg.solve_discrete_lyapunov, and scipy.linalg.solve_sylvester. Methods for computing the Schur decomposition are inherently serial, and therefore one does not see it in any CUDA libraries. I think we could absolutely parallelize batches of Schur decompositions of small matrices. The same should also hold for the generalized Schur decomposition scipy.linalg.qz, which is used in the matrix equation solvers scipy.linalg.solve_continuous_are and scipy.linalg.solve_discrete_are. Large batches of such matrix equations come up in applications of systems of ordinary differential equations in subjects like mathematical biology, for simulations over a sweep of possible parameter values. These Schur decompositions seem challenging but doable.

Going over other linear algebra functions that aren't available in CuPy according to this comparison table, I think scipy.linalg.cholesky_banded would be a very strong candidate for this approach, and certainly less ambitious than the Schur decomposition. scipy.linalg.cholesky could of course be done as well. scipy.linalg.hessenberg would fall out essentially for free if we can do the Schur decomposition, and should probably be seen as a stepping stone towards scipy.linalg.schur.

Another problem we could solve is the handling large batches of interpolative matrix decompositions (scipy.linalg.interpolate.interp_decomp and also scipy.linalg.interpolate.id_to_svd). These are not currently available in CuPy. Unlike the Schur decomposition, it's feasible to parallelize this for the case of a single large matrix³, it just doesn't seem to be available in any of NVIDIA's libraries. It may also be beneficial to add versions of functions that already exist in CuPy which are tailored for handling large ensembles of small problems. I could see this being useful for qr, svd, lstsq, and maybe others. I'm not sure how best to handle choosing between a method better suited for a large number of small problems vs one suited for a small number of large problems from an API perspective in CuPy; it's at least something to think about.

Optimization and root finding

Handling large batches of scalar root finding and scalar minimization problems might initially seem out of reach using ufuncs, because one must pass in user callables, but there is actually a path forward here. CuPy's ElementwiseKernel dynamically compiles kernels using the NVRTC JIT at runtime. If we can compile user defined def functions that use a restricted set of Python to strings of C++ code, these strings could be dynamically inserted into the operation string of an ElementwiseKernel to allow the creation of kernel which can do vectorized rootfinding or scalar minimization of a user defined function. We could want to stick to GPU friendly algorithms such Chandrupatla's algorithm for rootfinding. I suspect this could probably done by ripping out appropriate parts of Pythran.

I think there's a lot of potential for handling batches of least squares problems, each of the functions in https://docs.scipy.org/doc/scipy/reference/optimize.html#least-squares-and-curve-fitting seem like viable candidates. I think there's a lot of potential for adding better GPU support for ensemble learning in Scikit-Learn.

Integration and differentation

If we can compile user Python callables to C++ as mentioned above, there's a lot of potential for integration and differentiation methods as well. Parallelizing the solution of ensembles of ordinary differential equations also seems possible but very hard; it is something I'll keep an eye on though.

More

One could surely find more use cases. I hope the above is a good sample. Where else to spend time can be decided by community interest.

Support for alternative array libraries beyond CuPy

I think the above is a good strategy for adding many new useful things to CuPy. What about other array libraries like JAX and PyTorch? What I want to propose is relying on CuPy as a dependency to serve as a delivery mechanism to bring functions to JAX and PyTorch as well. We can get much of the way there using DLPack to perform zero-copy transfers of arrays between libraries. There's still some work that needs to be done in this direction, but in principal it should work. To enable JAX-JIT support, I've developed a JAX/CuPy bridge which was discussed here, data-apis/array-api-extra#602, and is awaiting feedback from JAX developers here, jax-ml/jax#34732. The JAX/CuPy bridge doesn't allow JITing through these CuPy kernels obviously, but lets them work in JIT-compiled JAX code by letting the compiler treat them as opaque black boxes. Autodifferentiation could in principle be supported through the JAX FFI as well if we write functions for the derivatives ourselves. The current limitation of the JAX/CuPy bridge is that it requires the GIL to be held to make the CuPy calls, but my benchmarking shows it still seems to work well enough.

Support for platforms beyond CUDA.

This approach also already works with AMD's ROCm (as of cupy/cupy#10005). In addition, although CuPy doesn't support SYCL, these xsf kernels themselves could in principle be used with SYCL if desired, although the surrounding orchestration with ElementwiseKernel would not be available.

The importance of documentation

I don't just want to follow through on this pattern for efficient single-source kernels for CPU and GPU. I think it will be very important to document it and make the pattern legible for the community. Right now, the knowledge of how to implement any particular function with this strategy is still obscure, but I think it could become a standard pattern that is straightforward to follow for either human or LLM. By focusing first on doing this targeted high impact examples and writing some clear and comprehensive tutorial style documentation, I think we open up the possibility for a lot of valuable community contributions. For so long, community contributions were locked out by the obscurity of SciPy's old Fortran codebase. No more.