Performance-Portable Tensor Transpositions in MLIR
摘要
This paper presents an optimized code generation for arbitrary out-of-place tensor transpositions using the MLIR compiler infrastructure, portable across CPU architectures. The proposed modular and reusable approach encodes optimizations such as multi-level tiling and explicit vectorization at multiple levels of abstraction as a sequence of transformation and conversion passes in MLIR. The efficient code generated is evaluated on AMD, Intel, and ARM processors and achieves performance comparable to the state-of-the-art HPTT library [27], a compelling speedup over Eigen [10], and a significant fraction of the STREAM memory bandwidth on these platforms. We further integrate this progressive lowering pipeline into COMET, an MLIR-based compiler for tensor contractions, and obtain an average speedup of \(26\%\) with the TTGT approach.