可扩展性与自定义方言

机器学习操作和目标硬件架构的复杂性和多样性，使传统、单一的中间表示的表达能力捉襟见肘。现代ML编译器，特别是MLIR，处理这种复杂性的一个核心设计理念是可扩展性。MLIR没有试图定义一套单一、通用的操作和类型，而是提供了一个用于定义和组合模块化方言的框架。

方言：ML抽象的命名空间

可以将MLIR中的方言看作一个专门的命名空间，它包含特定的一组操作、类型和属性，这些都为特定方面或抽象级别定制。这种模块化方法使MLIR能够在同一个架构中同时表示来自不同层的计算。

例如，一个ML模型最初可能使用模仿TensorFlow（tf方言）或PyTorch操作符的高级方言中的操作来表示。优化过程然后逐步将这种表示形式降低到侧重于线性代数（linalg方言）、结构化控制流和循环（affine和scf方言）、向量操作（vector方言）的方言，最终降至硬件专用方言，例如用于CPU的llvm、用于GPU的gpu，或用于跨厂商GPU编程的spirv。

方言中定义的主要组成部分包括：

1. 操作 (Ops)： 它们是IR中计算或结构的基本单元。每个操作都属于一个特定的方言（例如，tf.Conv2D、linalg.matmul、affine.for、llvm.add）。操作定义其参数、结果、属性，以及重要地，它们的语义。它们还可以有自定义汇编格式用于文本表示，以及验证器，以确保IR根据方言规则的正确性。
2. 类型： 除了标准内置类型（如整数、浮点数、向量、张量）之外，方言还可以定义自定义数据类型。例子包括量化类型（!quant.uniform<i8:f32>)或代表硬件特定状态或资源的类型。这些类型确保操作在具有预期语义和约束的数据上执行。
3. 属性： 属性提供与操作或类型相关的编译时元数据。它们是用于配置操作（例如，卷积的步幅和填充，表示为ArrayAttr）或指定类型特征（例如，MemRef类型的内存空间）的常量值。方言可以定义复杂、结构化的属性。

ML计算通过不同的MLIR方言逐步降低，从高级框架表示降至硬件特定目标。自定义硬件方言可以整合到此流程中。

实现通用性：接口

虽然方言提供了专用性，但构建一个高效的编译器需要通用分析和转换，它们可以在不同的方言之间操作，而无需了解其具体细节。MLIR通过接口实现这一点。接口定义了一个契约或一组方法，操作或方言可以实现这些方法。

例如，InferTypeOpInterface允许任何方言的操作提供逻辑，根据其操作数类型和属性推导其结果类型。一个通用的类型推断过程随后可以在任何操作上查询此接口，无论其属于哪个方言，以在IR中传播类型信息。同样，还存在用于内存效果（MemoryEffectOpInterface）、循环表示（LoopLikeOpInterface）以及许多其他常见的编译器规约的接口，这使得诸如融合、缓冲化或调度等过程能够更通用地编写。

定义自定义方言

MLIR可扩展性的真正优势在于能够定义全新的自定义方言。这对于以下几个原因非常重要：

1. 支持新型硬件： 当目标是具有独特指令或内存架构的新型AI加速器或DSP时，自定义方言可以直接表示其功能。降低过程随后可以将中级抽象（如linalg操作）转换为此自定义硬件方言的特定操作。
2. 特定方面优化： 对于某些方面（例如，稀疏计算、特定信号处理任务），自定义方言可以提供更高级的操作，更有效地捕捉方面语义，从而在降低到更通用形式之前实现更强大的专业优化。
3. 研究与实验： 通过定义自定义方言，研究人员可以快速构建新的IR结构或优化策略原型，而无需修改核心MLIR框架或干扰现有方言。

定义方言通常涉及使用MLIR的C++ API，或者更常见的是使用TableGen。TableGen是一种声明式描述语言，在LLVM和MLIR中广泛使用，用于定义代表IR组件（如操作、类型、属性和接口）的记录。通过这些TableGen描述，实现方言类、解析器、打印器和验证逻辑的C++代码会自动生成，显著减少了样板代码。

例如，定义一个自定义操作需要指定其名称（在方言命名空间内）、其参数和结果（带类型约束）、其属性，以及潜在的C++方法，用于验证、形状推断，或通过接口定义特定特性。

可扩展性的优势

MLIR的方言系统为构建复杂的ML编译器提供了重要优势：

• 模块化： 方言封装方面特定逻辑，使编译器更易于理解、维护和扩展。
• 关注点分离： 高级、与硬件无关的优化可以与低级、目标特定代码生成清晰分离。
• 可重用性： 通用架构（解析、验证、过程管理）和接口允许通用分析和转换在不同方言之间重用。
• 逐步降低： 混合方言的能力允许从抽象的框架级表示逐步转换为具体的硬件指令，并在最合适的抽象级别应用优化。
• 适应性： 该系统通过添加新方言，而不会干扰现有方言，从而很容易适应新的ML模型、算法和硬件目标。

通过使用方言和接口实现可扩展性，MLIR提供了一个依据，以应对在不断变化的硬件环境中优化多样化ML工作负载的挑战。理解这个核心原则对于理解现代ML编译堆栈的设计以及它们如何实现高性能非常重要。