仿射方言与分析

优化计算核需要准确了解内存相对于循环迭代的访问方式。尽管通用中间表示（IR）允许数组索引中使用任意算术运算，但这种灵活性通常使编译器无法证明循环转换的安全性。如果编译器不能确保重新排序循环能保持数据依赖关系，它必须保守地避免这种优化。MLIR的仿射方言通过施加一个限制性强但功能强大的结构约束来解决此问题：所有控制流和内存访问都必须使用循环归纳变量和常量的仿射（线性）函数来定义。

这种规范的数学结构使得多面体编译技术得以应用。通过将循环嵌套建模为几何多面体，编译器可以进行精确的依赖性分析，从而实现诸如循环分块、倾斜和向量 (vector)化等积极转换，而不会改变程序语义的风险。

仿射映射与集合

这种方言的基本组成部分是仿射映射和整数集合。这些构造将索引计算逻辑与操作本身分开，使编译器能够通过数学方法分析内存模式。

仿射映射定义了一个从维度列表和符号列表到结果列表的转换。它遵循以下形式：

$$(d_0, d_1, \dots, d_n)[s_0, s_1, \dots, s_m] \rightarrow (r_0, r_1, \dots, r_k)$$

这里，$d$ 表示维度（通常是循环归纳变量），$s$ 表示符号（在作用域内保持不变的运行时值，例如图像宽度或批次大小）。结果表达式 $r$ 必须是 $d$、$s$ 和整数常量的线性组合。

考虑一个用于访问步长为 $W$ 的扁平化二维数组的映射：

#flatten_map = affine_map<(i, j)[W] -> (i * W + j)>

在这个定义中，i 和 j 是维度，W 是一个符号。编译器可以静态分析这个映射，以判断增加 j 会使内存访问移动 1 个单位，而增加 i 会使其移动 W 个单位。

整数集合有着相似的用途，但它定义的是有效点的范围，而不是转换。它们用于条件执行，例如处理分块循环中的边界条件。一个集合由一系列必须成立的仿射约束（等式或不等式）组成。

$$S = \{ (i, j) \mid 0 \le i < N, \ 0 \le j < M, \ i + j \ge 0 \}$$

结构化控制流与内存访问

仿射方言引入了用于控制流（affine.for、affine.if）和内存操作（affine.load、affine.store）的特定操作。与标准控制流不同，affine.for 循环强制其边界是外部循环变量和符号的仿射映射。

这一限制确保了迭代空间形成一个凸多面体。下面的MLIR片段展示了一个计算矩阵相加的嵌套循环，其中约束条件明确嵌入 (embedding)在IR结构中：

func.func @matrix_add(%A: memref<128x128xf32>, 
                      %B: memref<128x128xf32>, 
                      %C: memref<128x128xf32>) {
  affine.for %i = 0 to 128 {
    affine.for %j = 0 to 128 {
      %a = affine.load %A[%i, %j] : memref<128x128xf32>
      %b = affine.load %B[%i, %j] : memref<128x128xf32>
      %sum = arith.addf %a, %b : f32
      affine.store %sum, %C[%i, %j] : memref<128x128xf32>
    }
  }
  return
}

在这个例子中，编译器知道内存访问模式严格遵循 (i, j) -> (i, j)。这种简单的映射简化了依赖性分析。如果我们应用分块转换，编译器将引入新的映射来处理块内和块间坐标。

下图展示了仿射方言如何构造循环嵌套。迭代空间由边界定义，内存访问则通过应用于归纳变量的仿射映射得到。

仿射循环嵌套的结构，显示了归纳变量和符号如何传入仿射映射以确定精确的内存地址。

依赖性分析与多面体优化

仿射方言的主要优点是它支持精确的依赖性分析。在传统优化中，检查两个指针是否别名通常是无法判定的。在仿射语境下，数据依赖性检查简化为求解一个线性不等式系统，通常使用整数线性规划（ILP）或傅里叶-莫茨金消元法。

如果编译器希望并行化上述循环，它会检查是否存在任何迭代 $(i', j')$ 依赖于迭代 $(i, j)$ 中计算出的结果。由于读写发生在完全相同的坐标 $(i, j)$，并且没有跨迭代依赖（如 C[i, j] = C[i-1, j]），编译器可以在数学上证明所有迭代都是独立的。

这种能力使得循环分块（或阻塞）成为可能。分块将循环迭代空间划分为更小的块，以便将工作集放入CPU缓存或GPU共享内存。在仿射方言中，分块不仅仅是一种启发式方法；它是一种坐标转换。

下面的图表展示了一个循环的迭代空间，其中 $i$ 的范围是 0 到 16，$j$ 的范围是 0 到 16。颜色编码代表了一种 4x4 的分块策略。编译器将原始循环转换为一系列按顺序访问这些彩色块的循环，从而提升缓存局部性。

平铺迭代空间的可视化。编译器将执行分组到块（颜色）中以最大化数据重用，这种转换通过仿射分析得到验证。

渐进式下沉

仿射方言很少作为最终目标。它作为一种高级分析能力。一旦应用了分块、展开-合并或向量 (vector)化等优化，仿射构造通常会下沉到 SCF（结构化控制流） 方言或直接到 LLVM 方言。

在下沉过程中，抽象的 affine.for 循环会被转换为具有明确索引计算算术的标准循环。affine.load 操作变为标准指针算术。这种职责分离使得仿射方言能够专注于有效的转换逻辑，而无需关注指针位转换或机器特定内在函数等低级实现细节。

当你需要对密集张量进行大量循环转换时，通常会使用仿射方言。但是，如果你的计算涉及间接访问（例如，访问稀疏矩阵中的 $A[B[i]]$），仿射方言无法表示该模式，此时你必须回退到更通用的 SCF 或标准方言，从而失去执行多面体优化的能力。