GPU Load Balancing

TL;DR

This paper introduces a flexible, open-source GPU load-balancing framework that enhances performance and productivity for irregular parallel algorithms, demonstrated by a novel matrix multiplication method achieving significant speedups.

Contribution

It presents a new GPU load-balancing abstraction supporting static and dynamic schedules, and introduces Stream-K, a work-centric GEMM parallelization with superior performance.

Findings

Stream-K achieves up to 14x peak speedup on GEMM.

The framework improves GPU utilization for irregular workloads.

Stream-K outperforms CUTLASS and cuBLAS in speed and consistency.

Abstract

Fine-grained workload and resource balancing is the key to high performance for regular and irregular computations on the GPUs. In this dissertation, we conduct an extensive survey of existing load-balancing techniques to build an abstraction that addresses the difficulty of scheduling computations on the GPU. We propose a GPU fine-grained load-balancing abstraction that decouples load balancing from work processing and aims to support both static and dynamic schedules with a programmable interface to implement new load-balancing schedules. Prior to our work, the only way to unleash the GPU's potential on irregular problems has been to workload-balance through application-specific, tightly coupled load-balancing techniques. With our open-source framework for load-balancing, we hope to improve programmers' productivity when developing irregular-parallel algorithms on the GPU, and also improve the overall performance characteristics for such applications by allowing a quick path to experimentation with a variety of existing load-balancing techniques. Using our insights from load-balancing irregular workloads, we build Stream-K, a work-centric parallelization of matrix multiplication (GEMM) and related computations in dense linear algebra. Whereas contemporary decompositions are primarily tile-based, our method operates by partitioning an even share of the aggregate inner loop iterations among physical processing elements. This provides a near-perfect utilization of computing resources, regardless of how efficiently the output tiling for any given problem quantizes across the underlying processing elements. On GPU processors, our Stream-K parallelization of GEMM produces a peak speedup of up to 14x and 6.7x, and an average performance response that is both higher and more consistent across 32K GEMM problem geometries than state-of-the-art math libraries such as CUTLASS and cuBLAS.