Kernel Operations on the GPU, with Autodiff, without Memory Overflows
Benjamin Charlier Jean Feydy Joan Alexis Glaunès François-David Collin Ghislain Durif
cs.LG






       


Kernel Operations on the GPU, with Autodiff,
 without Memory Overflows
    Benjamin Charlier* benjamin.charlier@umontpellier.fr 

       IMAG

       Université de Montpellier, CNRS

       Montpellier, France
       Jean Feydy* jean.feydy@ens.fr 

       DMA

       École Normale Supérieure

       Paris, France
       Joan Alexis Glaunès* alexis.glaunes@parisdescartes.fr 

       MAP5 

       Université de Paris, CNRS 

       Paris, France

       François-David Collin francois-david.collin@umontpellier.fr

       Ghislain Durif ghislain.durif@umontpellier.fr

       IMAG

       Université de Montpellier, CNRS

       Montpellier, France * equal contribution
       
    December 30, 2023
================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================





The  library provides a fast and memory-efficient GPU support for tensors whose entries are given by a mathematical formula, such as kernel and distance matrices.
 alleviates the main bottleneck of tensor-centric libraries 
for kernel and geometric applications: memory consumption. 
It also supports automatic differentiation and outperforms standard GPU baselines,
including   tensors or the  and  libraries. 
 combines optimized  schemes with binders for high-level languages:  ( and ),  and . As a result, high-level “quadratic” codes can now scale up to large data sets with millions of samples processed in seconds. 

 brings graphics-like performances for kernel methods and is  freely available on standard repositories (PyPi, CRAN). To showcase its versatility, we provide tutorials in a wide range of settings online at 
https://www.kernel-operations.io.







  kernel methods, 
  Gaussian processes, GPU,
  automatic differentiation






§ INTRODUCTION



Recent advances in machine learning have been driven by the diffusion of two pieces of software: automatic differentiation and GPU backends for tensor computations. Today, thanks to e.g. the  or  libraries <cit.>, users routinely perform gradient descent on functions that involve millions of parameters. 


These high-level  frameworks unlock the
use of massively parallel hardware for machine learning research.
Under the hood, they rely on  routines that are
often supported by hardware manufacturers:
the  and  libraries edited by Nvidia
provide the binaries for linear algebra
and convolutions that power a majority of deep learning models.
In practice, the presence of a complete software stack
(from low-level binaries to well-documented  libraries)
is a prerequisite for the widespread adoption
of a research idea by the machine learning community.

The  library intends to provide such a
solid numerical foundation
for all methods that involve large distance or kernel matrices.
A motivating example is the computation of pair-wise interactions
of the form:

    a_i   ∑_j=1^exp( - x_i-y_j^2 / 2σ^2)  b_j 
         = ∑_j=1^ k(x_i,y_j)  b_j 
            for   i=1,…,

where  and  range from a hundred to a billion,
x_1, …, x_ and y_1, …, y_ all belong
to a common vector space ^,
the signals b_1, …, b_ are real numbers 
and k(x,y) is a Gaussian kernel of deviation σ > 0. 
This operation is often understood as a matrix-vector
product with a Gaussian kernel matrix (K_i,j) = (k(x_i,y_j)), or
equivalently as a convolution with the kernel function k
that is sampled on the point clouds (x_i) and (y_j).


To perform this computation, 
common practice
is to create an explicit -by- buffer (K_i,j)
and compute Eq. (<ref>)
with a linear algebra routine.
Unfortunately, this method requires the
storage of the kernel matrix 
as a contiguous array in memory
and does not scale when  and  are
in the order of 50k or more.
To work around this problem,
a common strategy is thus to decompose Eq. (<ref>)
as a collection of smaller matrix-vector products
using a  or  “” loop.
This method alleviates memory issues but remains inefficient
in spaces of dimension ⩽ 100:
the transfer of the tiles of the kernel
matrix (K_i,j) between different layers of GPU memory
remains a narrow bottleneck for computations.



 leverages efficient  schemes
from the graphics literature <cit.>
to streamline the use of registers and reach
optimal performance in this setting.
It is fast, transparent to use and targets a single yet powerful abstraction: semi-symbolic arrays whose entries M_i,j are given by a mathematical formula
“F(x_i, y_j)”, evaluated on data arrays that are indexed by line and column numbers “i“ and “j”.


It can be called from the major scripting languages used in the scientific community:  ( and ), , and . In this introduction, we focus on the  interface and
refer to our documentation for additional tutorials,
applications and benchmarks.




§  PURPOSE AND USAGE



A generic reduction framework. 
The workhorse of the  library is 
a  engine for generic reductions
on sampled data. 
Let us assume that we have at hand:





  * 



parameters: a collection p^1∈^d_p^1, …, p^P∈^d_p^P of vectors;

  * i-variables: a collection x^1∈^M× d_x^1, …, x^X∈^M× d_x^X of matrices, with rows indexed by i∈[i]1, (hence for each k and i, x^k_i is a vector in ^d_x^k);


  * j-variables: a collection y^1∈^N× d_y^1, …, y^Y∈^N× d_y^Y of matrices, with rows indexed by j∈[i]1, (hence for each k and j, y^k_j is a vector in ^d_y^k);

  * a vector-valued symbolic formula F(p^1, …, p^P, x^1_i, …, x^X_i, y^1_j, …, y^Y_j)∈^d_out;

  * a reduction operation such as a sum, max, argmin, log-sum-exp, etc.







Then, a single call to the  engine allows users to evaluate the expression:

    a_i   Reduction_j=1,…,[ F(p^1, …, p^P, x^1_i, …, x^X_i, y^1_j, …, y^Y_j)  ] 
          for   i=1,…,

efficiently, with a linear memory footprint on GPUs and CPUs.
As illustrated in our gallery of tutorials, 
this level of generality allows  to handle off-grid convolutions,
k-nearest neighbors classification, 
k-means clustering and many other tasks.















The  abstraction.
The “” wrapper for  arrays
and  tensors
lets users specify computations
along the lines of Eq. (<ref>)
with a tensor-like interface. 


For instance, we can specify the
Gaussian matrix-vector product of Eq. (<ref>) with:

[bgcolor=blue!4,linenos,mathescape,tabsize=4,firstnumber=1,stripnl=false,fontsize=]python
from pykeops.torch import LazyTensor  # Wrapper for PyTorch Tensors
x_i  = LazyTensor(x[:,None,:])      # (M,D) Tensor -> (M,1,D) Symbolic Tensor
y_j  = LazyTensor(y[None,:,:])      # (N,D) Tensor -> (1,N,D) Symbolic Tensor
D_ij = ((x_i - y_j)**2).sum(dim=2)  # (M,N,1) Symbolic matrix of squared distances
K_ij = (- D_ij / (2 * s**2)).exp()  # (M,N,1) Symbolic Gaussian kernel matrix
a    = K_ij @ b  # Genuine torch Tensor. (M,N,1) @ (N,D) = (M,D)


In the script above, no computation is performed at lines 4 and 5:
lazily, the  engine builds up a symbolic formula F
encoded as a string attribute of the  . 
The lazy chain is only terminated by
the virtual matrix product of line 6, a generic reduction
that triggers the real computation:
no intermediate -by- buffer is created in the global device memory. 

Note that variable types (i-, j-variable or parameter) 
are inferred from the shapes of the input tensors at lines 2 and 3:
in practice, symbolic tensors are
as easy to use as sparse matrices.
The  wrapper turns a dense array into a symbolic matrix whose
axes -3 and -2 are understood as “virtual” dimensions;
a reduction on these axes is the signal that triggers
a call to the  engine.













    







As showcased on our website and at the end of this paper, 
scripts for kernel and geometric applications generally outperform
their  and  counterparts
by several orders of magnitude while keeping a linear memory footprint. 

 support a wide range of mathematical operations that mimic the usual
interface for  arrays and  tensors.
They fully support broadcasting and batch dimensions, 
as well as a decent collection 
of reduction operations: , 
,  and  
but
also  or  that return
the indices of the smallest (or K-smallest) elements
of the rows of a symbolic tensor.

Inner engine.
Internally,  creates an optimized  code
for every new reduction and formula F that it encounters.
Binaries are then compiled and stored on the hard drive for later use:
compilation relies on the standard  stack (,  and/or  compilers) and is only performed once per reduction.



Backpropagation.
Crucially,  supports automatic differentiation up to arbitrary orders of differentiation:
a new binary is created automatically for every new partial derivative
that is required by the user's computations. 
This mechanism is fully integrated with the
 engine and lets users
“backprop” through  calls using
the usual 
and  methods.



§ INTERNAL ENGINE




Efficient GPU schemes.
















The   routines are based on parallel implementations of a
tiled Map-Reduce scheme: we split reductions into medium-sized blocks to best leverage the shared memory of GPU devices. This technique is well-known in GPU programming and detailed in  reference guides <cit.>. The originality of  resides in its genericity with respect to the formula F,
its math-friendly interface and its support of automatic differentiation.






















Building formulas, automatic chain-rule.
Internally,  encodes formulas as recursively templated  classes: each elementary vector-valued operation is defined in a templated  class, that specifies an evaluation code to inline in the  scheme and a symbolic expression for its gradient. As an example, the element-wise, vector-valued exponential function is encoded with: 
[bgcolor=blue!4,linenos,mathescape,tabsize=4,firstnumber=1,stripnl=false]c++
template < class F >
struct Exp : UnaryOp<Exp, F> 
  // dimension of the output: Exp(F) has the same dimension as F
  static const int DIM = F::DIM 
  // actual computation, to be inlined inside the Cuda device routine
  static DEVICE INLINE void Operation(TYPE *out, TYPE *in) 
    #pragma unroll
    for (int k=0; k<DIM; k++)  out[k] = exp(in[k]); 
  
  // templated expression for the adjoint of the differential operator
  // of Exp w.r.t. the variable V, and applied to GRADIN input vector:
  // ∇_V(Exp(F)).GRADIN = ∇_V(F).(Exp(F)×GRADIN) in math notations
  template < class V, class GRADIN >
  using DiffT = typename F::template DiffT< V, Mult<Exp<F>,GRADIN> >;
;

Using similar definitions for other mathematical operations, 
we can then express a Gaussian matrix-vector product as a sum reduction of 
the formula  where  and  are special classes that represent data loaders.
This templated engine has two main advantages: first, the code for evaluating the full formula F is built up at compile time, allowing the compiler to optimize the resulting code; second, the chain rule derivation for automatic differentiation can be performed at compile time, using the recursive template mechanics.













































§ PERFORMANCE EVALUATION


























 is geared towards computations
that fit the mould of Eq. (<ref>).
In this specific context, it combines
a fully transparent interface with state-of-the-art performance.
To illustrate this, we compare 
to similar scientific computing libraries
– ,
,  <cit.> 
and  <cit.>
– on a simple benchmark:
the Gaussian kernel matrix-vector product of Eq. (<ref>) with an increasing number of points = in dimension D=3. All experiments are performed with  precision on a Nvidia RTX 2080 Ti GPU, with the exception of the PyTorch-TPU column that was run in Google Colab;
code is available on our repository in the
https://github.com/getkeops/keops/tree/master/benchmarks
folder.
Our https://www.kernel-operations.io/keops/_auto_benchmarks/index.html also includes comparisons
with JAX <cit.> and domain-specific libraries
such as FAISS <cit.> for K-Nearest Neighbors search.







 
             
              
                 
          

    =10k 
        9 ms     10 ms
        13 ms     1.0 ms
        3.80 ms     0.7 ms
        0.4 ms 
 
    =100k 
             
        89 ms     34.1 ms
        36.8 ms     15.0 ms
        14.6 ms
 
    =1M 
             
             3.8 s
        2.79 s     1.39 s
        1.38 s
 
    Lines of code 
        5     5
        5     15
        17     5
        55  
 
    Interface

    
    

-
like     -
like    
    -
like     
    

    
    low-level
     -
like    
    

 

    



As evidenced by this table,  turns -like scripts
into highly competitive binaries.
Going further, it can be neatly interfaced with the iterative linear solvers of the  <cit.> or  <cit.> libraries
and supports the specification of cluster-wise block-sparsity patterns: this allows users to solve large kernel linear systems efficiently, with applications to geology (Kriging), imaging (splines), statistics (Gaussian processes) and data sciences (kernel methods). 





§ INTENDED USE, LIMITATIONS AND FUTURE WORKS


 fills a specific but important niche in machine learning research.
Unlike most other compilers for deep learning computations, 
such as  and , it is meant to be used directly
by theorists of the machine learning community.
Our focus on the simple yet powerful concept
of symbolic matrices allows us to keep
a transparent interface, while being more efficient
than the generalist  and  frameworks
on a wide range of computations.
In future works, we intend to add support for approximation schemes
such as the Nyström and FFM methods <cit.>,
beyond the block-sparse truncation rule
that is currently supported.
These features will be valuable to many researchers in the field,
while being out of scope for most deep learning libraries.

The core strength of the   engine is that it
optimizes the use of GPU registers for computations
that fit the template of Eq. (<ref>).
In practice, it is therefore of limited use in situations
where a kernel matrix can fit in memory and 
is re-used a large number of times, 
or when the evaluation of the formula “F(p,x_i,y_j)”
takes a significant amount of time.
Keeping in mind the motivating example of Eq. (<ref>),
we believe that  will be most useful
for computations that involve 10k points or more
in a space of dimension ⩽ 100.
Going forward, our main priority
will be to ease the deployment of pre-compiled  binaries,
reduce compilation times to at most a handful of seconds per routine
and add support for the newly released  Tensor cores.













§ ACKNOWLEDGEMENTS


The three first authors are the project leaders: they contributed equally to the library
and its documentation. 
The authors also thank Alain Trouvé, whose theoretical work in shape analysis
was the first motivation for the development of the  engine.


0.2in




