Hope to help on GPU kernels development

[gouchangjiang]

Hi, The Virtual Brain developers and researchers, I would like to know if you need some volunteers for GPU kernel developers. I am pleased to spend my spare time on it. The technology stack can be on JAX, Numba, or others. This link may related. I can refactor some already implemented algorithms to make them work faster on GPUs, or new problems, but no guarantee of the deadline.

Hope to hear back from you.

[gouchangjiang]

a JAX related tech stack Custom_op_for_GPUs

[maedoc]

hi!

Thanks for your interest & kind offer. The main bottlenecks we've seen so far w/ Jax are related to sparse matrices for neural fields (similar to finite element methods), and we've starting working on OpenCL versions, e.g.

• standard CSR SpMV for GPU
• standard CSR SpMV for CPU
• time delayed sparse coupling, batched for GPU

In that repository we were trying to use OpenCL & ISPC to cover GPU & CPU widely (ISPC has similar programming model to OpenCL). It would be of interest to pick this up along different dimensions

• port from OpenCL to CUDA, Metal for the different platforms
• extend delay kernels to ISPC for good CPU support
• add gradients, especially for delay case, to support machine learning cases
• link to Jax as custom ops

The gradient support for delays in particular had me a bit blocked, since it might require either a different data structure or some atomics (which are likely too slow to be useful).

Anyway, there's no required timeline, so feel free to think about it and see what you might be interested in? Thanks again!

[gouchangjiang]

You are welcome. So, let me start with porting the SpMV to CUDA, then extend delay kernels. Start from the popcorn repository?

[maedoc]

Sure take a look and let me know what you think 👍

[gouchangjiang]

I just added the SpMV cuda kernel, pr, and planning to port the delayed sparse coupling to CUDA, then to improve the performance for both.