MLIR中Dialect的抽象层级简介

Ref:

https://discourse.llvm.org/t/codegen-dialect-overview/2723

基本概念

我们写机器学习代码时，我们通常会使用高级别的数据结构，比如tensor。但是，这些高级别的结构不能直接在硬件上执行，需要经过一些中间步骤将它们转换为低级别的表示形式，比如buffer。这些中间步骤通常涉及到dialects，也就是不同的语言方言。在转换的过程中，我们会使用最近引入的dialects，也就是最适合当前转换的语言方言。这些dialects的目的是将高级别的结构转换为低级别的表示形式，以便最终在硬件上执行。

MLIR（Multi-Level Intermediate Representation，多级中间表示）是一个现代编译器基础设施，旨在统一不同抽象级别的程序表示、转换和优化。MLIR的核心设计理念是通过可扩展的IR（中间表示）来支持各种编程模型和硬件目标。

在MLIR中，Dialect（方言）是最核心的概念之一。Dialect可以理解为一组相关操作、类型和属性的集合，用于表示特定领域的计算或硬件抽象。不同的Dialect可以在同一个IR中共存，这使得MLIR能够在不同抽象层级之间进行灵活转换。"Dialects" 是 MLIR 中的一个概念，它可以根据它所抽象的特征的级别来组织成一个堆栈，就像一个抽象的塔一样。在这个堆栈中，越高级别的抽象表示越抽象，越低级别的抽象表示越接近实际的底层实现。将较高级别的抽象表示向下转化成较低级别的抽象表示通常比较容易，因为较低级别的抽象已经包含了较高级别抽象的所有特征，相当于是将抽象的东西转化为更具体的东西。但是，反过来将较低级别的抽象表示提升到较高级别的抽象表示通常比较困难，因为较低级别的抽象不一定包含较高级别抽象的所有特征，需要通过一些复杂的技术和算法来推导和计算出来。

Dialect的内部组成部分

一个Dialect通常由以下几个关键组成部分构成：

Dialect定义：
- 每个Dialect可以看作是一个命名空间，包含了一组相关的操作和类型
- 开发者可以通过继承Dialect类来定义新的方言，并在其中注册自定义的操作和类型
Operations（操作）：
- Operation是MLIR的基本构建块，类似于传统编程语言中的指令
- 每个Operation可以有多个输入和输出，并且可以附带属性来提供额外信息
- 在自定义Dialect中，开发者可以定义新的Operation以便更好地表达特定领域的计算
Types（类型）：
- 类型用于描述Operation的输入和输出的性质
- MLIR提供了基本类型（如整数和浮点数），同时也允许自定义类型
- 自定义类型通常用于表示领域特定的数据结构或约束
Attributes（属性）：
- 属性是不可变的键值对，附加在Operation上以提供更多上下文
- 属性可以是简单的标量值（如整数或字符串），也可以是复杂的数据结构
- 通过使用属性，可以将额外的元数据附加到Operation上，而不影响其操作数
Patterns（模式）和Rewriters（重写器）：
- 模式用于匹配和转换Operation，用于优化或转换
- 重写器是实现模式匹配和转换的具体逻辑，允许在MLIR中进行优化和代码生成

Dialect的抽象层级

MLIR的Dialect体系构成了一个从高级抽象到低级实现的逐层细化过程。以下是主要的抽象层级：

1. 高层Dialect (High-Level Dialect)

特点：

接近于源语言或领域特定语言
操作通常比较复杂，对应于高级语言中的一个函数或计算图节点
类型系统复杂，可能包含张量、动态形状等高级概念
主要用于表示计算的意图，而非具体实现细节

目的：

捕捉源语言的语义，保留原始代码的计算意图
方便进行高层次的领域特定优化

例子：

TensorFlow Dialect ：表示TensorFlow图操作，如tf.MatMul、tf.Add等
HLO Dialect：表示XLA的高级操作
TCF Dialect：用于张量计算框架

2. 中层Dialect (Mid-Level Dialect)

特点：

介于高层和底层之间的抽象
操作粒度适中，可以表示常见的计算模式
类型系统相对简单，通常包含基本类型和结构
主要用于通用优化和转换

目的：

将高层Dialect转换为更接近硬件的表示
进行一些通用的优化，如循环变换、内存访问优化等

例子：

Linalg Dialect

对结构化数据进行结构化处理的通用表示。具有以下特点：

既可以将tensor作为操作数，也可以将buffer作为操作数
包含两类主要操作类型：
- Payload操作：表示执行具体运算的操作，如矩阵乘法、卷积等
- Struct操作：表示如何进行运算的操作，如循环结构
在实际应用中，外部Dialect很多情况下会先转换到Linalg Dialect再执行后续优化
提供了高层次的线性代数抽象，便于应用领域特定优化

应用场景：

线性代数库构建
深度学习编译器后端
图像处理算法表示

Vector Dialect

对SIMD（单指令多数据）或SIMT（单指令多线程）模型的抽象。其特点包括：

作为一种向量的中间表示，可以被转换到不同Target对应的底层表示
实现Tensor在不同平台的支持
包含向量化操作，如向量加载/存储、向量计算等
支持向量掩码、向量洗牌等高级向量操作

应用场景：

多媒体处理加速
科学计算应用优化
深度学习模型向量化

Affine Dialect

对多面体编译（polyhedral compilation）的实现。其主要特点：

包含多维数据结构的控制流操作
支持多维数据的循环和条件控制，存储映射操作等
目标是实现多面体变换，如：
- 自动并行化
- 用于局部改进的循环融合和平铺
- MLIR中的循环矢量化
使用仿射表达式描述循环边界和内存访问模式

应用场景：

高性能计算优化
编译器循环优化

SCF (Structured Control Flow) Dialect

比控制流图CFG更高层的抽象：

提供结构化控制流操作，如并行的for和while循环以及条件判断
通常Affine和Linalg会降低到SCF
SCF也可以降低为Standard中的基本控制流图（CFG）操作
保留了控制流的结构信息，便于后续优化

应用场景：

编译器中间表示
程序分析和控制流优化

Async Dialect

用来表示异步操作模型：

通常为一些操作的集合，在不同的抽象层次含义有所变化
支持异步执行、token传递和异步区域
可以表示并发执行和同步机制
有助于实现并行计算模型

应用场景：

并发程序编写
分布式计算系统实现

3. 低层Dialect (Low-Level Dialect)

特点：

接近于目标硬件的抽象
操作非常细粒度，通常对应于一条机器指令
类型系统简单，主要包含硬件支持的基本类型
主要用于代码生成和硬件相关优化

目的：

将中层Dialect转换为可以直接映射到硬件的表示
进行硬件特定的优化

例子：

LLVM Dialect ：映射到LLVM IR的操作，如llvm.add、llvm.load等
NVVM Dialect：NVIDIA GPU特定操作
ROCDL Dialect：AMD GPU特定操作
AVX Dialect：Intel x86 SIMD指令集
Neon Dialect：ARM SIMD指令集
SVE Dialect：ARM可扩展向量扩展指令集
SPIR-V Dialect：映射到SPIR-V（用于图形和计算）

4. 标准层 (Standard Layer)

贯穿所有层级的是：

Standard Dialect：提供基本的、通用的操作，如算术运算、控制流等
Builtin Dialect：提供MLIR本身的核心概念，如模块、函数等

这些方言为其他所有方言提供了基础构建块。

Dialect转换流程

MLIR中的典型转换流程通常遵循从高到低的抽象层级降低过程：

高层领域特定Dialect （如TensorFlow、HLO）
↓
中层通用Dialect （如Linalg）
↓
结构化控制流 （SCF或Affine）
↓
向量化 （Vector Dialect）
↓
底层平台特定Dialect（如LLVM、NVVM、AVX等）

在这个过程中，每一步转换都会降低抽象层次，增加实现细节，同时应用特定于该层级的优化。

通俗易懂的例子：矩阵乘法的多层表示

让我们通过一个矩阵乘法的例子来说明这些抽象层级如何协同工作：

1. 高层表示（Linalg Dialect）

在最高层，我们可以使用Linalg Dialect直接表达矩阵乘法的语义：

mlir 复制代码

%C = linalg.matmul ins(%A, %B : tensor<4x8xf32>, tensor<8x4xf32>) 
               outs(%init : tensor<4x4xf32>) -> tensor<4x4xf32>

这里使用linalg.matmul直接表达了矩阵乘法的语义，无需关心具体实现细节。代码简洁明了，保留了计算的意图。这是一个典型的Linalg Dialect中的payload操作，直接表达了计算的内容。

2. 中层表示（Affine + SCF Dialects）

降低抽象层次后，矩阵乘法被展开为嵌套循环，可以使用Affine Dialect表示：

mlir 复制代码

%C = affine.for %i = 0 to 4 {
  %row = affine.for %j = 0 to 4 {
    %sum = constant 0.0 : f32
    %sum_final = affine.for %k = 0 to 8 iter_args(%sum_iter = %sum) -> f32 {
      %a_element = tensor.extract %A[%i, %k] : tensor<4x8xf32>
      %b_element = tensor.extract %B[%k, %j] : tensor<8x4xf32>
      %product = arith.mulf %a_element, %b_element : f32
      %sum_next = arith.addf %sum_iter, %product : f32
      affine.yield %sum_next : f32
    }
    tensor.insert %sum_final into %init[%i, %j] : tensor<4x4xf32>
  }
}

这一层级展示了计算的结构，包括嵌套循环和基本运算。代码变长了，但提供了更多优化机会，如循环变换、向量化等。Affine Dialect特别适合表示这种多维循环结构，因为它可以使用仿射表达式来描述循环边界和访问模式。

或者使用SCF Dialect表示：

mlir 复制代码

%C = scf.for %i = %c0 to %c4 step %c1 iter_args(%init_arg = %init) -> tensor<4x4xf32> {
  %row = scf.for %j = %c0 to %c4 step %c1 iter_args(%row_arg = %init_arg) -> tensor<4x4xf32> {
    %sum = constant 0.0 : f32
    %sum_final = scf.for %k = %c0 to %c8 step %c1 iter_args(%sum_iter = %sum) -> f32 {
      %a_element = tensor.extract %A[%i, %k] : tensor<4x8xf32>
      %b_element = tensor.extract %B[%k, %j] : tensor<8x4xf32>
      %product = arith.mulf %a_element, %b_element : f32
      %sum_next = arith.addf %sum_iter, %product : f32
      scf.yield %sum_next : f32
    }
    %updated_row = tensor.insert %sum_final into %row_arg[%i, %j] : tensor<4x4xf32>
    scf.yield %updated_row : tensor<4x4xf32>
  }
  scf.yield %row : tensor<4x4xf32>
}

SCF提供了更加通用的结构化控制流表示，适用于不需要精确仿射表达式的场景。

3. 向量化表示（Vector Dialect）

如果我们想利用SIMD指令，可以将上述代码向量化：

mlir 复制代码

%C = scf.for %i = %c0 to %c4 step %c1 iter_args(%init_arg = %init) -> tensor<4x4xf32> {
  %row = scf.for %j = %c0 to %c4 step %c1 iter_args(%row_arg = %init_arg) -> tensor<4x4xf32> {
    // 加载向量
    %a_vec = vector.transfer_read %A[%i, %c0], %cst : tensor<4x8xf32>, vector<8xf32>
    %b_vec = vector.transfer_read %B[%c0, %j], %cst : tensor<8x4xf32>, vector<8xf32>
    
    // 向量点乘
    %dot = vector.dot %a_vec, %b_vec : vector<8xf32>, vector<8xf32>
    
    // 存储结果
    %updated_row = tensor.insert %dot into %row_arg[%i, %j] : tensor<4x4xf32>
    scf.yield %updated_row : tensor<4x4xf32>
  }
  scf.yield %row : tensor<4x4xf32>
}

Vector Dialect提供了向量级操作，使得代码可以利用SIMD指令加速。这里使用了向量读取和点乘操作，直接表达了向量级计算。

4. 低层表示（LLVM Dialect）

进一步降低抽象，转换为接近机器代码的表示：

mlir 复制代码

llvm.func @matmul(%A: !llvm.ptr<float>, %B: !llvm.ptr<float>, %C: !llvm.ptr<float>) {
  // 初始化循环变量
  %c0 = llvm.constant(0 : index) : i64
  %c4 = llvm.constant(4 : index) : i64
  %c8 = llvm.constant(8 : index) : i64
  
  // 外层循环（i）
  llvm.br ^bb1(%c0 : i64)
^bb1(%i : i64):
  %cond = llvm.icmp "slt" %i, %c4 : i64
  llvm.cond_br %cond, ^bb2(%c0 : i64), ^bb7

  // 中层循环（j）
^bb2(%j : i64):
  %cond2 = llvm.icmp "slt" %j, %c4 : i64
  llvm.cond_br %cond2, ^bb3, ^bb6

  // 初始化累加器
^bb3:
  %init = llvm.constant(0.0 : f32) : f32
  llvm.br ^bb4(%c0 : i64, %init : f32)

  // 内层循环（k）
^bb4(%k : i64, %sum : f32):
  %cond3 = llvm.icmp "slt" %k, %c8 : i64
  llvm.cond_br %cond3, ^bb5, ^bb5_exit

  // 循环体：计算 A[i,k] * B[k,j] 并累加
^bb5:
  // 计算内存地址和加载值
  %a_idx = llvm.add %i, %k : i64
  %a_val = llvm.load %A[%a_idx] : !llvm.ptr<float>
  
  %b_idx = llvm.add %k, %j : i64
  %b_val = llvm.load %B[%b_idx] : !llvm.ptr<float>
  
  // 计算乘积并累加
  %prod = llvm.fmul %a_val, %b_val : f32
  %new_sum = llvm.fadd %sum, %prod : f32
  
  %k_next = llvm.add %k, %c1 : i64
  llvm.br ^bb4(%k_next, %new_sum : i64, f32)

  // 存储结果
^bb5_exit:
  %c_idx = llvm.add %i, %j : i64
  llvm.store %sum, %C[%c_idx] : f32
  
  %j_next = llvm.add %j, %c1 : i64
  llvm.br ^bb2(%j_next : i64)

^bb6:
  %i_next = llvm.add %i, %c1 : i64
  llvm.br ^bb1(%i_next : i64)

^bb7:
  llvm.return
}

在这一层级，我们看到了显式的控制流跳转、内存地址计算等低级细节。代码变得更加复杂，但更接近实际的硬件执行模式。

5. 特定硬件指令集（如AVX Dialect）

对于支持AVX指令集的x86平台，可以进一步转换为使用AVX指令：

mlir 复制代码

// 使用AVX指令进行向量化计算（示意代码）
%ymm0 = avx.load %A_ptr : vector<8xf32>
%ymm1 = avx.load %B_ptr : vector<8xf32>
%ymm2 = avx.mul %ymm0, %ymm1 : vector<8xf32>
%result = avx.hadd %ymm2 : vector<8xf32>
avx.store %C_ptr, %result : vector<8xf32>

这一层级直接映射到特定硬件的指令集，充分利用硬件特性。

Dialect转换过程的价值

这种多层级表示的价值在于：

保留语义信息：高级表示保留了原始代码的意图（"这是矩阵乘法"），而不仅仅是一系列低级指令。这使得编译器能够理解代码的目的，从而进行更智能的优化。
分层优化：每一层都可以应用特定的优化：
- 高层：领域特定优化，如算法替换、数学等价变换
- 中层：通用优化，如循环变换、内存访问优化
- 低层：硬件特定优化，如指令调度、寄存器分配
灵活性：可以在任何层级进入或退出编译管道，支持不同的编译流程和目标。
可扩展性：可以通过添加新的Dialect来支持新的领域或硬件，无需修改整个编译器框架。

通过这种方式，MLIR能够同时处理高级抽象和低级实现细节，使得从高级语言到机器代码的转换过程更加灵活和高效，特别适合处理复杂的领域特定计算，如机器学习、图形处理等。

MLIR中Dialect的抽象层级 简介