NVIDIA/cuda-tile

解决的问题

CUDA Tile IR 通过提供专用的中间表示（IR）和编译器基础设施，简化了高性能 CUDA 内核的开发。它特别针对基于分块的计算模式优化和 NVIDIA 张量核心单元的利用，减少了手动管理内存层次结构和 GPU 特定优化的复杂性。

工作原理

基于 MLIR（多级中间表示）框架构建，该项目提供了一种用于表达分块计算的领域特定方言。它包含用于程序化 IR 构建的 Python 绑定和用于高效序列化的字节码格式。程序可以从 MLIR 编译为字节码，然后通过 CUDA 驱动 API 进行即时编译（JIT），或使用 tileiras 工具提前编译（AoT）为 cubin 文件。

适用人群

专为希望使用更高层次抽象进行分块和张量核心优化，而非编写原始 CUDA 代码的高性能 GPU 内核开发者设计。

主要亮点

• 基于 MLIR 的方言：为基于分块的计算提供第一类操作和类型。
• Python API：支持对 IR 进行程序化操作和转换。
• 灵活的编译方式：支持即时编译（JIT）和提前编译（AoT）两种路径。
• 张量核心优化：专门针对 NVIDIA 张量核心单元进行优化，以实现最大性能。