Memory access coalescing optimization based on Triton compiler
With the advent of the era of large artificial intelligence models, the massive number of parameters and complex data processing have led to an increasing proportion of memory-intensive operators. The abundance of high-latency memory access behaviors results in underutilization of computational resources. Triton programming language and compiler enable efficient development and execution of GPU operators to accelerate model training. Currently, the memory access optimization module of the compiler faces issues where some tensor sizes cannot be coalesced for memory access or vectorized for efficient memory access. To address these challenges, we studied the implementation process of memory access coalescing optimization passes. By optimizing the computation logic of memory access patterns, we achieved coalesced memory access for operators with odd tensor sizes. For the problem of various tensor sizes not fully leveraging vectorized memory access instructions for efficient access, we implemented support for vectorized memory access by adding axis attribute value types during the degradation of the abstract syntax tree. Experiments were conducted using kernel examples from the Triton official tutorial for testing. The experimental results demonstrate that the average speedup achieved by thread memory access pattern optimization is 2.4, and the average speedup for vectorized memory access optimization across multiple tensor sizes is 1.5.
Authors
- Wei Gao (ORCID: https://orcid.org/0000-0003-2028-2407)
- Lin Han (ORCID: https://orcid.org/0000-0002-8569-7972)
- Zhen Fang (ORCID: https://orcid.org/0000-0003-0602-6255)
- Jing Li (ORCID: https://orcid.org/0000-0002-8044-2284)
- Yingying Li (ORCID: https://orcid.org/0000-0002-3999-2647)
Institutions
- PLA Information Engineering University (CN)
- Zhengzhou University (CN)
- National Supercomputing Center in Wuxi (CN)
Publication Details
- Published
- 2025-05-09
- DOI
- https://doi.org/10.1117/12.3067061
- Citations
- 1
- Primary Topic
- Advanced Computing and Algorithms
- Type
- article
- Field-Weighted Citation Impact
- 6.18