Swift: High Parallelism Program Generation of Tensor Operators for Accelerating Deep Learning Inference

Optimizing deep learning inference, particularly reducing the execution latency of tensor computations at small batch sizes, is crucial for the successful and widespread adoption of deep neural network (DNN) models. However, current deep learning compilers and hand-tuned libraries often fail to achieve high hardware efficiency when executing small-batch workloads. The primary reason is the inherently sequential nature of reductions (e.g., along the hidden dimension in the flattened GEMM for LLM decoding), which is difficult to parallelize and therefore fails to fully utilize available hardware resources. In this article, we propose Swift, a novel search-based approach for efficiently generating high-performance programs for GPUs by maximizing hardware utilization. The key insight is that reduction parallelization can be incorporated into a unified representation alongside the existing tile structure, significantly expanding the search space for high-performance programs. Concretely, by enumerating all possible parallel mappings of loops, we first generate a large search space that contains high-performance programs. Then, to efficiently explore the extended search space, we employ subspace shifting exploration to identify promising regions, effectively prune large portions of the less-promising search space. We conduct experiments on three distinct GPU architectures using a diverse set of benchmarks representative of typical application scenarios. Experimental results demonstrate that Swift achieves an average speedup of 1.19× over the state-of-the-art compiler-based approaches. Moreover, compared with vendor-provided hand-tuned libraries, Swift achieves an average speedup of 2.40×.