Towards a high-performance AI compiler with upstream MLIR

TL;DR

使用MLIR框架实现高性能AI编译器,性能达到手写代码的90%。

cs.PL 🔴 高级 2024-04-15 9 次浏览
Renato Golin Lorenzo Chelini Adam Siemieniuk Kavitha Madhu Niranjan Hasabnis Hans Pabst Evangelos Georganas Alexander Heinecke
AI编译器 MLIR 线性代数 缓存优化 微内核

核心发现

方法论

该研究提出了一种基于MLIR的编译流程,通过使用TensorFlow和PyTorch的Linalg-on-Tensor输入IR进行缓存级优化和微内核降级,实现高效向量化。核心组件包括张量方言的打包原语、线性代数流水线和微内核降级机制。

关键结果

  • 实验结果显示,编译器生成的代码在性能上达到了手写代码的90%以上,尤其是在多核环境下表现优异。
  • 通过使用VNNI、BFDOT、BFMMLA等指令,编译器在不同CPU上实现了高效的张量分布。
  • 在单线程和多线程环境中,编译器的性能与手写代码相当,特别是在Intel和AMD架构上。

研究意义

该研究为高性能计算和机器学习领域提供了一种新的编译方法,解决了现有编译器在处理复杂线性代数运算时的性能瓶颈。通过自动化编译流程,降低了对手动优化的依赖。

技术贡献

技术贡献包括引入了新的张量打包和解包操作,改进了Linalg方言的上游编译流程,并实现了与libxsmm库的接口,支持多种CPU架构的优化。

新颖性

这是首次在MLIR框架中实现如此高效的线性代数编译流程,与现有方法相比,显著提高了编译器的性能和灵活性。

局限性

  • 在某些特定的硬件配置下,编译器的性能可能会受到内存带宽的限制。
  • 当前的编译器优化策略在某些情况下可能无法充分利用硬件特性。
  • 未来工作需要进一步优化编译器的自动化决策模型。

未来方向

未来方向包括进一步优化编译器的成本模型,探索更多硬件架构的支持,以及改进微内核的自动选择策略。

AI 总览摘要

该研究提出了一种基于MLIR的高性能AI编译器,旨在解决现有编译器在处理复杂线性代数运算时的性能瓶颈。通过使用TensorFlow和PyTorch的Linalg-on-Tensor输入IR,该编译器实现了缓存级优化和微内核降级,在多核环境下表现优异。

编译器的核心技术包括张量方言的打包原语、线性代数流水线和微内核降级机制。实验结果显示,编译器生成的代码在性能上达到了手写代码的90%以上,特别是在Intel和AMD架构上表现出色。

该研究为高性能计算和机器学习领域提供了一种新的编译方法,降低了对手动优化的依赖,并为未来的编译器优化提供了新的方向。尽管在某些硬件配置下存在性能限制,但通过进一步优化编译器的成本模型和自动化决策策略,未来有望实现更高的性能提升。

深度分析

研究背景

随着机器学习和高性能计算的发展,编译器在处理复杂线性代数运算时的性能瓶颈日益显现。现有编译器通常依赖于手动优化和特定硬件的深度知识,这限制了其通用性和灵活性。

核心问题

核心问题在于如何在不依赖手动优化的情况下,实现高效的线性代数运算编译。这需要解决缓存优化、微内核选择以及多核环境下的性能瓶颈。

核心创新

该研究的核心创新在于引入了基于MLIR的编译流程,通过张量方言的打包原语和线性代数流水线,实现了自动化的缓存优化和微内核降级。

方法详解

  • �� 使用TensorFlow和PyTorch的Linalg-on-Tensor输入IR
  • �� 进行缓存级优化以提高数据访问效率
  • �� 降级到微内核以实现高效向量化
  • �� 支持多种CPU架构的优化

实验设计

实验设计包括在不同CPU架构上测试编译器的性能,使用MLP模型进行基准测试,比较编译器生成代码与手写代码的性能。

结果分析

实验结果显示,编译器在多核环境下性能达到手写代码的90%以上,特别是在Intel和AMD架构上表现优异。

应用场景

该编译器可用于高性能计算和机器学习领域的线性代数运算,降低对手动优化的依赖,提高开发效率。

局限与展望

尽管编译器性能优异,但在某些硬件配置下可能受到内存带宽的限制。未来需要进一步优化编译器的自动化决策模型。

通俗解读 非专业人士也能看懂

想象一个厨房,厨师需要快速准备多道菜。传统方法需要厨师手动切菜、烹饪,而新方法则是使用自动化设备,快速切菜并进行烹饪。这个编译器就像厨房里的自动化设备,通过优化流程,快速处理复杂的计算任务。

简单解释 像给14岁少年讲一样

想象你在玩一个游戏,需要快速升级角色。传统方法需要你手动选择装备和技能,而新方法则是游戏自动帮你选择最佳装备和技能。这个编译器就像游戏里的自动升级系统,帮助你快速提升计算性能!

术语表

MLIR (多级中间表示)

一种用于编译器开发的中间表示框架,支持多级IR重写。

用于实现编译器的核心优化流程。

Linalg-on-Tensor (张量上的线性代数)

一种用于表示线性代数运算的MLIR方言。

作为输入IR进行缓存优化和微内核降级。

VNNI (向量神经网络指令)

一种用于加速神经网络计算的指令集。

用于提高编译器生成代码的性能。

BFMMLA (BF16矩阵乘法累加)

一种用于BF16数据类型的矩阵乘法指令。

在支持BF16的CPU上实现高效矩阵运算。

libxsmm (微内核库)

一个用于高性能计算的开源微内核库。

编译器通过XSMM方言与libxsmm库接口。

开放问题 这项研究留下的未解疑问

  • 1 如何进一步优化编译器的自动化决策模型,以提高不同硬件配置下的性能。
  • 2 在支持BF16的CPU上,如何更有效地利用硬件特性。

应用场景

近期应用

高性能计算

该编译器可用于加速高性能计算任务,特别是在处理复杂线性代数运算时。

远期愿景

机器学习优化

通过进一步优化编译器的自动化决策模型,实现更高效的机器学习模型训练和推理。

原文摘要

This work proposes a compilation flow using open-source compiler passes to build a framework to achieve ninja performance from a generic linear algebra high-level abstraction. We demonstrate this flow with a proof-of-concept MLIR project that uses input IR in Linalg-on-Tensor from TensorFlow and PyTorch, performs cache-level optimizations and lowering to micro-kernels for efficient vectorization, achieving over 90% of the performance of ninja-written equivalent programs. The contributions of this work include: (1) Packing primitives on the tensor dialect and passes for cache-aware distribution of tensors (single and multi-core) and type-aware instructions (VNNI, BFDOT, BFMMLA), including propagation of shapes across the entire function; (2) A linear algebra pipeline, including tile, fuse and bufferization strategies to get model-level IR into hardware friendly tile calls; (3) A mechanism for micro-kernel lowering to an open source library that supports various CPUs.

cs.PL cs.AI cs.AR cs.DC cs.LG