Keeping Futhark off the GPU

Earlier this year, Elias Smedegaard worked on an efficient method to compute sparse Jacobian matrices using automatic differentiation. The BSc thesis dealt with graph colouring, which is a notoriously difficult problem in computational complexity. However, the focus was on devising a feasible solution that employed an efficient algorithm, rather than the optimal one.

The problem was that the algorithm used by Elias was inherently sequential, which posed a challenge when trying to utilize GPU parallelism for the second step of the process.

Futhark, the programming language used for this project, has a simple compilation model where arrays must reside in memory. When the GPU pipelines are activated, all arrays are transferred to GPU memory. While this design is convenient and safe, it may not be the most efficient. For instance, if an array is stored in GPU memory, the loop body would need to copy a single element from GPU to CPU memory, which would be inefficient due to the overhead of communication setup.

To address this issue, Elias explored a solution that would allow the program to decide whether to store arrays in CPU or GPU memory based on their usage. Ideally, Futhark would support separate compilation, enabling different parts of the program to use different backends. However, this feature is absent in Futhark, making it challenging to implement.

Instead, the team opted for a workaround: introducing a new attribute, #[cpu_function], that compiles the function body to CPU code and stores input, intermediate, and output arrays in CPU memory.

This approach allowed the team to successfully compute Jacobians on the GPU, albeit with a somewhat clunky solution. The team realized that a more sophisticated optimisation, which moved data instead of computation, would be a better long-term solution. However, implementing this optimisation would require considerable effort and time, leading the team to adopt a simpler workaround for the time being.