[MLIR] Linalg Dialect 设计思想

本文档旨在揭示构建更高级别的中间表示（IR）和方言以促进代码生成的关键设计原则。首先介绍Linalg方言的初衷，旨在解决高级分层优化问题，并在混合的export编译器环境中进行有效交互。设计工作基于大量现有技术的启发，旨在学习关键经验。Linalg的架构和设计原则，以及在高级张量计算原语方言和变换开发工作组的提案背景下，帮助定义这些高级张量计算原语的路径。

Linalg的最初目标是启动MLIR中的代码生成，通过消除复杂的代码生成问题，并引入利用编译器优化和高效库实现的调用能力。这种方法构建了高级的以生产力为先的代码生成解决方案。例如，如果某个高性能计算库或指令集架构具有在某种内存中运行的矩阵乘法原语，其运行效率可达到可实现峰值性能的95%，则在可能的情况下应使用该原语。否则，生成代码。

随着Linalg设计与MLIR设计共同演进，其应用范围超越了仅仅在密集张量上进行机器学习的领域。Linalg采用一种代码生成友好的方法，其中IR和变换同时演变。关键思想是声明和传输传统上由编译器分析获得的信息，这些信息捕获了变换的合法性和适用性。关键的变换被设计为在必要时间内保留此信息。例如，经过平铺和融合后，linalg.matmul仍然是linalg.matmul。

Linalg将变换的有效性与盈利能力分离，并在第一次迭代中自愿放弃后者。这些想法在EuroLLVM 2019开发者会议的MLIR教程中作为示例呈现。自最初的实现以来，Linalg设计随着核心MLIR基础设施的演变而演变，部分驱动了它们的演变。Linalg采用的方法被扩展为StructuredOps抽象，成为在张量和缓冲区上的实现。它由Vector方言补充，该方言定义了向量上的结构化操作，遵循与Linalg相同的基本原理和设计原则。

Linalg方言超越了线性代数操作，变得更加具有表现力，通过提供循环嵌套的抽象来支持并行化、归约和滑动窗口等任意MLIR区域。它还具有支持更丰富的数据类型，如稀疏和不规则张量和缓冲区的潜力。Linalg设计开放于演化和与其他方言和方法的交叉污染，成功用作代码生成相关抽象的暂存区，产生了以下泛化。

在Linalg中仍有一些设计问题尚未解决，它并不声称是解决所有编译问题的通用解决方案。它的目标是推动领域特定抽象的思考和实现，其中程序员的意图可以在IR中以非常高的层次直接被捕获。随着范围的演变，比“Linalg”更好的名称可能会消除与方言和其基础方法相关的某些混淆，以及它的目标和限制。

Linalg设计汲取了数十年的先前技术的灵感，以设计出现代实用的解决方案。以下非详尽列表列举了一些影响Linalg设计的项目，包括ONNX、LIFT、XLA、Halide、TVM、Tensor Comprehensions和多面体编译的经验教训。Linalg希望解决的问题包括操作定义中考虑编译器转换和下降、最小化可用操作集、在操作定义中考虑编译器转换和下降、操作之间的不重叠、完美操作知识的重要性、操作语义的局限性和固定性、多面体模型的集成挑战、保持程序信息的重要性、避免过度扩散和冗余以及渐进降低代码以保留信息。

核心编程规则变换和简单性是Linalg IR及其操作的主要目的。处理的问题基本上是由于高性能计算和并行硬件架构上的编译驱动的HPC编译问题。相关变换的选择采用协同设计方法，并涉及避免过度扩散和冗余、评估现有系统、认识到简单性和可维护性的重要性、避免在系统Y中重写X、在高级编译中保留更多可用信息并利用它们、将代码生成视为降低潜在函数的过程、将高级多面体编译器中的提升概念作为避免的场景以及将代码快速降至仿射形式、通用循环或CFG形式以保留结构信息。

Linalg设计的核心原则包括声明性规范和避免raising，以及渐进式lowering和不要过快丢失信息。复杂而具有影响力的转换不一定难以操作、编写或维护。将XLA风格的高级op语义知识与通用属性混合在一起，在MLIR中直接描述这些语义，可以适用于搜索和机器学习。可组合的声明性转换、可扩展性和未来性是Linalg设计的关键要素。MLIR允许定义结构化控制流和结构化数据类型，我们选择利用这些特性，特别是在MemRefType表示密集的非连续内存区域，以便扩展到更高级的数据类型，如不规则、稀疏和混合密集/稀疏张量、树、哈希表、记录表甚至图。

本文档总结了Linalg设计的核心原则、其在基于MLIR的实现中的应用，以及如何帮助协调核心指导原则和真实世界的需求。算法+数据结构=程序的概念强调了控制流与数据结构之间的强关系。方言不需要在转换下保持封闭的观点，以及考虑潜在和单调性，进一步展示了Linalg设计的灵活性和适应性。最后，结合核心指导原则和先前艺术作品的观察，我们呈现了一张图，以粗略展示如何看待现有系统中结构信息的分布情况，展示了Linalg方言及其未来演进的目标。