Looking for advice on package for lazily evaluated kernel matrix on the GPU

General Usage
, ,
1

I am working on creating a package for lazily evaluated kernel matrices, which will work on the GPU via KernelAbstractions. The idea is to create a central object (LazyKernelMatrix) which does not allocate and evaluates the kernel function k(x_i, y_i) only when necessary. This is quite easy to do on the CPU by overloading Base.get_index but my main goal is to implement matrix-vector multiplication (or maybe other functions) on the GPU, along with various sparse approximation methods, such as that of HMatrices.jl or AcceleratedRPCholesky.jl, all without allocating. My goal is to make the package compatible with Krylov.jl so one can do data science tasks such as ridge-regression on matrices that wouldn’t otherwise fit in the GPU memory.

My main motivation comes from my own projects, which primarily arise from kernel machine learning (ridge-regression, Gaussian process regression, etc.). I do not have a GPU with a crazy amount of memory, so this kind of thing allows me to handle a much bigger system than I would be able to otherwise. For example, with the current demo package I have going, I was able to do mat-vec multiplies with a 100,000\times 100,000 matrix; much larger than what would fit in my measly GPU memory.

I am looking for advice on how to structure the package and, if this package might benefit you, what specific applications you might have for it. Or if you think the idea is flawed, let me know! I am still relatively new to julia, so I would appreciate some advice from people who have a better sense for the linear algebra and machine learning ecosystem.

The current demo package I have put together is structured in the following way:

  • The central object LazyKernelMatrix <: AbstractMatrix stores the points on which the kernel is evaluated, a KernelOperator object, and an ExecutionPlan object
  • The KernelOperator object contains the mathematical kernel function, along with any parameters.
  • The ExecutionPlan object contains the backend and workgroup size to be supplied to KernelAbstractions. The idea is to be able to easily change the execution plan (gpu/cpu, blocking strategy, or sparse/low-rank approximation method) without changing the exposed interface.

Some specific aspects I am not sure about is whether this overall structure seems efficient and scalable, or if there are aspects that I am not noticing which might make this a poor choice. Also I am having a bit of trouble with how to handle the parameters in the KernelOperator object. It is currently implemented as a simple functor, with an immutable array of parameters, but for some applications (Gaussian process hyperparameter optimization for example), it would be useful to be able to easily change these parameters without having to create the whole matrix object again. I am also currently forced to use a lot of parametrized types in the structures, which is probably necessary to easily compile to different GPU vendors, but it feels like I could be specifying supertypes better or something along those lines.

Any advice or comments on this project would be greatly appreciated! The basic package that I have currently set up (only tested CUDA currently) is here.

2

I think this would be a useful package. I took a stab at a similar idea here (aided by LLMs):

Right now it grabs into some internals of the Inti library, but it would nice to have the lazy kernel matrix logic in a separate package :+1: