vini-fda / heat-wgpu

Solving the heat equation with wgpu in Rust.

heat-wgpu

This is a Rust program which leverages computing and rendering capabilities of modern GPUs to efficiently solve the heat equation and display the result in real time. It is based on the wgpu crate, which is a Rust wrapper around the WebGPU API.

The heat equation is a partial differential equation which describes the flow of heat in a given domain. For a 2D domain, it is given by the following equation:

$$ \frac{\partial u}{\partial t} = \alpha \left( \frac{\partial^2 u}{\partial x^2} + \frac{\partial^2 u}{\partial y^2} \right) $$

where $u$ is the temperature, $t$ is time, $\alpha$ is the thermal diffusivity and $x$ and $y$ are the spatial coordinates.

The equation is solved numerically using the finite difference method, which consists in discretizing the domain into a grid and approximating the derivatives by finite differences. The discretization scheme used here is the Crank-Nicolson method, which is unconditionally stable and second-order accurate in time and space.

The resulting system of equations is then solved using the Conjugate Gradient method, which is written as a series of WGSL compute shaders and runs on the GPU through wgpu.

The mathematical formulation, including explanations of the finite difference scheme and the conjugate gradient method, is described in discretization.

Sample output

Usage

From binary

Download the latest executable from the GitHub release page. Supported platforms are Windows, Linux (mainly Ubuntu) and macOS. You can start the binary on its own by running it with no arguments. On all systems, this can be done by double-clicking the executable.

On Linux and macOS, you can also run it from the terminal:

./heat-wgpu

On Windows, you can run it from the command prompt:

> heat-wgpu.exe

From source

You need to have the Rust toolchain installed. Check out the official language website for instructions in how to install. Then, run the following command:

cargo run --release

References and useful resources

• LeVeque, R. J. (2007). Finite difference methods for ordinary and partial differential equations: steady-state and time-dependent problems. Society for Industrial and Applied Mathematics.
• Shewchuk, J. R. (1994). An introduction to the conjugate gradient method without the agonizing pain.
• Bell, N., & Garland, M. (2008). Efficient sparse matrix-vector multiplication on CUDA (Vol. 2, No. 5). Nvidia Technical Report NVR-2008-004, Nvidia Corporation.
• GPU Gems 2: Programming Techniques for High-Performance Graphics and General-Purpose Computation. 2005. Addison-Wesley Professional.
  • Specifically Chapter 44: A GPU Framework for Solving Systems of Linear Equations, which also cites the following papers:
    • Bolz, J., I. Farmer, E. Grinspun, and P. Schröder. 2003. "Sparse Matrix Solvers on the GPU: Conjugate Gradients and Multigrid." ACM Transactions on Graphics (Proceedings of SIGGRAPH 2003) 22(3), pp. 917–924.
    • Krüger, Jens, and Rüdiger Westermann. 2003. "Linear Algebra Operators for GPU Implementation of Numerical Algorithms." ACM Transactions on Graphics (Proceedings of SIGGRAPH 2003) 22(3), pp. 908–916.
• Harris, M. (2007). Optimizing parallel reduction in CUDA. Nvidia developer technology, 2(4), 70.
• Mikhailov, A. (2019). Turbo, An Improved Rainbow Colormap for Visualization.
• WebGPU specification.

heat-wgpu's Issues

Improve coloring

The coloring is inconsistent between operating systems: on Windows, the color tone is a bit muted in comparison with Linux, for instance.

Edit: Upon investigation, it seems that the difference is due to the fact that the Windows implementation has an sRBG surface format, where on my Linux it falls back to Brgba8Unorm

Proper compute step in a separate thread

Currently, the compute step is done synchronously with the render step, but it should be changed so that it does not depend on things like Window events (e.g. keyboard input, resizing etc).

Implement egui integration

Useful links:

• Custom 3d in egui with wgpu

Implement the conjugate gradient method on GPU

Work list:

• Parallel dot product ✅
  • Parallel sum-reduce
  • Parallel element-wise product
• Parallel SpMV ✅
• Create buffers and compute pipeline ✅

Add parallel dot product implementation