用编译器自动生成高性能矩阵乘法代码,无需手动调优。
Library Liberation: Competitive Performance Matmul Through Compiler-composed Nanokernels
- 通过MLIR方言自动组合微内核,实现高效寄存器利用。
- 生成的代码性能媲美顶尖库(如CUTLASS),可直接用于生产。
- 适合希望免去手工优化却追求极致性能的开发者。
AI与机器学习工作负载的快速发展,使得高层领域操作与硬件高效利用之间的差距日益扩大。目前达到接近峰值性能仍需深厚的硬件知识——要么手动编写特定目标的内核(如DeepSeek),要么依赖专用库(如CUTLASS),这两种方式均增加复杂性并限制可扩展性。本文提出一种编译方案,利用MLIR方言自动构建可扩展、高性能的微内核,弥合领域操作与处理器能力之间的鸿沟。核心机制是从低层IR构造中组合纳米内核,实现近乎最优的寄存器使用,生成针对每个目标架构定制的高效微内核。我们在基于MLIR的编译器中实现了该技术,支持向量和分块的CPU指令。实验表明,生成的纳米内核具备生产级质量,性能与当前最先进的微内核库相当。
原文摘要 · Abstract (English)
The rapidly evolving landscape of AI and machine learning workloads has widened the gap between high-level domain operations and efficient hardware utilization. Achieving near-peak performance still demands deep hardware expertise-experts either handcraft target-specific kernels (e.g., DeepSeek) or rely on specialized libraries (e.g., CUTLASS)-both of which add complexity and limit scalability for most ML practitioners. This paper introduces a compilation scheme that automatically generates scalable, high-performance microkernels by leveraging the MLIR dialects to bridge domain-level operations and processor capabilities. Our approach removes dependence on low-level libraries by enabling the compiler to auto-generate near-optimal code directly. At its core is a mechanism for composing nanokernels from low-level IR constructs with near-optimal register utilization, forming efficient microkernels tailored to each target. We implement this technique in an MLIR-based compiler supporting both vector and tile based CPU instructions. Experiments show that the generated nanokernels are of production-quality, and competitive with state-of-the-art microkernel libraries.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。