Sparse matrices are widely used in graph neural networks, but multiplying them can be computationally expensive. Researchers have developed various methods to speed up this process, including workstealing, a simple and efficient strategy that reduces communication between processes by randomly stealing work from other processes. In this article, we explore how workstealing works and its advantages over other methods.
Workstealing: A Simple Strategy
Workstealing is the simplest method for reducing communication in graph neural network training. It works by randomly stealing work from other processes without regard to locality. This means that many blocks are available to be stolen, but performing stolen work is usually more expensive than regular work since all tiles of A, B, and C must be communicated.
Pseudocode for this workstealing algorithm is shown in Algorithm 3 using the stationary A method. In this algorithm, each process maintains a 2D work grid corresponding to the tiles of the stationary matrix. Each element in the work grid contains an integer that is initialized to zero. To claim a piece of work involving a tile, processes perform a fetch-and-add operation on the corresponding element of the work grid. The integer value returned corresponds to the piece of work that has been claimed.
Advantages of Workstealing
Workstealing has several advantages over other methods. Firstly, it reduces communication between processes, which can significantly speed up training time. Secondly, it is simple and easy to implement, making it a good choice for large-scale systems. Finally, workstealing can be combined with other methods, such as parallel processing, to further improve efficiency.
Conclusion
In conclusion, workstealing is a powerful technique for reducing communication in graph neural network training. By randomly stealing work from other processes without regard to locality, it simplifies the training process and improves overall efficiency. Although it has some limitations, workstealing remains a valuable method for optimizing graph neural network training.
