Optimization of reduction kernel based on Triton compiler

With the rapid development of deep learning models, the development of high-performance operators has become increasingly important and complex. The Triton compiler offers a more convenient and efficient solution for GPU operator development, striking a balance between flexibility and performance. Reduction operators, as fundamental components of deep learning models, play a critical role in enhancing the stability and performance of both training and inference. However, the current thread-level handling of reduction operators in the Triton compiler does not fully exploit the multilayered memory hierarchy and memory access resources of GPU. To address this issue, we propose optimizations for the code generation process of reduction operators in the Triton compiler. Specifically, we utilize Warp Shuffle instructions to optimize intra-warp processing for non-contiguous axis reduction and adopt a data-packing strategy to enhance the downgrading process of FP16 data types. These strategies effectively reduce shared memory usage during reduction and improve the parallel efficiency of computational units. Additionally, we address the warp divergence caused by masking in loops by employing a redundant masking tail-loop optimization technique, which reduces the overhead of mask loading and improves memory access efficiency. Experimental results on reduction operators with varying data scales demonstrate that the Warp Shuffle-based optimization achieves bandwidth improvements of 6.2%–32.9% for FP32 and FP16 data types. Furthermore, the packing optimization yields an average bandwidth improvement of 5.1% for FP16 data types, while the redundant masking tail-loop optimization achieves an average bandwidth improvement of 6.6%.