Hexcute: A Tile-based Programming Language with Automatic Layout and Task-Mapping Synthesis
This work addresses the problem of reducing programming effort and improving performance for GPU optimization of mixed-type operators in deep learning, representing an incremental improvement over prior compilers by automating key optimizations.
The paper tackles the challenge of optimizing mixed-data-type matrix multiplication operators for GPUs in deep learning workloads by introducing Hexcute, a tile-based programming language that automates layout and task mapping synthesis, achieving speedups of 1.7-11.28x over existing compilers for these operators and up to 2.91x in end-to-end evaluations.
Deep learning (DL) workloads mainly run on accelerators like GPUs. Recent DL quantization techniques demand a new matrix multiplication operator with mixed input data types, further complicating GPU optimization. Prior high-level compilers like Triton lack the expressiveness to implement key optimizations like fine-grained data pipelines and hardware-friendly memory layouts for these operators, while low-level programming models, such as Hidet, Graphene, and CUTLASS, require significant programming efforts. To balance expressiveness with engineering effort, we propose Hexcute, a tile-based programming language that exposes shared memory and register abstractions to enable fine-grained optimization for these operators. Additionally, Hexcute leverages task mapping to schedule the GPU program, and to reduce programming efforts, it automates layout and task mapping synthesis with a novel type-inference-based algorithm. Our evaluation shows that Hexcute generalizes to a wide range of DL operators, achieves 1.7-11.28$\times$ speedup over existing DL compilers for mixed-type operators, and brings up to 2.91$\times$ speedup in the end-to-end evaluation.