循环分析与变换：Affine 与 SCF

上一篇讲了 Tensor 和 Linalg——AI 编译的高层抽象。Linalg 的结构化操作（linalg.generic）最终要降阶为实际的循环。这一篇讲循环层的两个方言：Affine（带约束的仿射循环，可做依赖分析和多面体变换）和 SCF（结构化控制流，接近机器执行模型）。

一、两个方言的分工

linalg.generic / linalg.matmul
         │
         ▼
affine.for (仿射循环——可分析、可变换)
         │
         ▼
scf.for (结构化控制流——不再做多面体分析)
         │
         ▼
cf.br / cf.cond_br (通用 CFG——最终降到 LLVM IR)

affine 方言的循环有强约束——循环边界和数组下标必须是循环变量的仿射函数。这个约束的代价是表达力有限（不能表示间接访问如 A[B[i]]），但收益是编译器可以做精确的依赖分析和多面体调度。

scf 方言的循环没有这个约束——scf.for 的边界和步长可以是任意 SSA 值。这使它更通用，但也意味着依赖分析只能靠通用方法（无仿射假设）。

二、Affine 方言的表示

2.1 affine.for

// 仿射循环：从 0 到 N-1，步长为 1
affine.for %i = 0 to N {
  // %i 是 index 类型
  %val = affine.load %A[%i] : memref<?xf32>
  %res = arith.addf %val, %val : f32
  affine.store %res, %B[%i] : memref<?xf32>
}

// 步长不为 1
affine.for %i = 0 to 100 step 2 {
  // ...
}

// 嵌套 loop 带仿射下标
affine.for %i = 0 to M {
  affine.for %j = 0 to N {
    // A[i][j] ——仿射约束：下标是循环变量的线性函数
    %val = affine.load %A[%i, %j] : memref<MxNxf32>
  }
}

关键约束： - 循环下界、上界必须是仿射表达式（0、M、M/2 是可以的；A[i] 是不允许的）。 - 所有 affine.load 和 affine.store 的下标必须是循环变量和符号常量的仿射函数。 - affine.apply 可以计算仿射表达式。

2.2 affine.apply

// %idx = 2 * %i + %j + 5
%idx = affine.apply affine_map<(d0, d1) -> (d0 * 2 + d1 + 5)>(%i, %j)

affine.apply 的参数必须是 index 类型，输出也是 index 类型。它不执行实际计算——它声明了一个仿射关系，编译器可以利用它进行分析。

2.3 affine.if

// 条件必须是仿射表达式之间的比较
affine.if affine_set<(d0) : (d0 - 10 >= 0)>(%i) {
  // then
} else {
  // else
}

2.4 AffineMap

AffineMap 是 MLIR 中表示仿射变换的核心数学对象——它把一个多维索引空间映射到另一个：

// 例子：map(d0, d1) → (d0, d1 floordiv 4, d1 mod 4)
// 表示将 j 维度按 4 分块后的映射
auto map = AffineMap::get(2, 0, {
    getAffineDimExpr(0, &ctx),                     // d0
    getAffineDimExpr(1, &ctx).floorDiv(4),         // d1 floordiv 4
    getAffineDimExpr(1, &ctx) % 4                  // d1 mod 4
}, &ctx);

AffineMap 在内存布局、tiling 映射和多面体变换中是基础工具——但它的细节和 API 深度需要单独一篇文章展开。

三、依赖分析

Affine 方言的最大优势是精确的依赖分析。给定以下循环：

affine.for %i = 1 to N-1 {
  %prev = affine.load %A[%i - 1] : memref<?xf32>  // 读取前一个元素
  %curr = affine.load %A[%i] : memref<?xf32>       // 读取当前元素
  %sum = arith.addf %prev, %curr : f32
  affine.store %sum, %A[%i] : memref<?xf32>         // 写入当前元素
}

MLIR 的 AffineDependenceAnalysis 可以确定：

1. 第 i 次迭代的 store A[i] 和第 i+1 次迭代的 load A[i] 之间有 RAW（Read-After-Write）依赖。
2. 依赖距离为 1（迭代间距离为 1）。
3. 这个依赖是”循环携带的”（loop-carried）——不能被消除，但可以被调整。

这些信息直接驱动后续的调度决策： - 依赖距离为 1 → 不能向量化（相邻迭代有依赖）。 - 依赖是循环携带的 → 不能将循环标记为全并行。 - 如果依赖距离为 0 → 循环可能可以被并行化。

依赖分析的 C++ API

以下代码为示意伪代码，展示依赖分析的使用方式；真实 API 见 mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp 与 mlir/include/mlir/Dialect/Affine/Analysis/AffineAnalysis.h。

void analyzeLoopDependence(affine::AffineForOp forOp) {
  // 获取该循环的依赖分析结果
  SmallVector<SmallVector<AffineDependence, 2>, 4> dependences;

  // 计算依赖
  for (auto &dep : getDependences(forOp)) {
    if (dep.isLoopIndependent()) {
      // 依赖发生在同一次迭代内
    } else {
      // 依赖跨迭代——获取依赖距离向量
      auto distance = dep.getDependenceDistance();
      // distance[d] 表示在第 d 维上的依赖距离
    }
  }
}

四、与多面体（Polyhedral）模型的联系

Affine 方言的约束与多面体编译（Polyhedral compilation）的输入条件一致：

MLIR 的 Affine 方言可以被现有 Polyhedral 工具（如 ISL）直接消费，也可以使用 MLIR 内建的 AffineLoopFusion、AffineLoopTiling 等 Pass 做调度。

关键的工程差异：MLIR 的内建 Affine Pass 是有损的——它们针对常见模式做优化，不覆盖多面体模型的全部理论能力。如果你需要全 PPCG（Polyhedral Parallel Code Generation）的能力，需要集成外部工具（如 ISL、PPCG）。对多数 AI 编译中的 tiling、fusion、vectorization 场景，MLIR 内建 Pass 通常够用，但复杂调度仍需外部多面体工具或手写 Transform 策略。

五、SCF：结构化控制流

SCF（Structured Control Flow）方言不依赖仿射假设，提供通用的结构化控制流操作。

5.1 核心 Op

5.2 scf.for 的 iter_args

scf.for 的 iter_args（迭代参数）是 MLIR 结构的亮点——它用 Block 参数优雅地建模跨迭代状态：

// 求和：sum += A[i]
// init 是迭代的初始值，%sum 是跨迭代携带的值
%result = scf.for %i = %c0 to %N step %c1
    iter_args(%sum = %c0_f32) -> (f32) {
  %elem = memref.load %A[%i] : memref<?xf32>
  %new_sum = arith.addf %sum, %elem : f32
  scf.yield %new_sum : f32
}
// %result 是循环结束后的最终 sum

iter_args 是 MLIR 对照 LLVM phi 节点的结构化语法——它把跨迭代的数据流显式表达为”带初始值的循环参数”。

5.3 scf.if 的 yield

%result = scf.if %cond -> i32 {
  %then_val = arith.addi %a, %b : i32
  scf.yield %then_val : i32
} else {
  scf.yield %c : i32
}

SCF 的所有控制流 Op 都通过 scf.yield 产出一个或多个值（从 Region 到外层）。

六、Affine → SCF 的降阶

-lower-affine Pass 将 Affine 方言降阶为 SCF + 标准 CF：

mlir-opt input.mlir -lower-affine

降阶过程：

// 输入（Affine）
affine.for %i = 0 to N {
  %val = affine.load %A[%i] : memref<?xf32>
  affine.store %val, %B[%i] : memref<?xf32>
}

// 输出（SCF + memref）
scf.for %i = %c0 to %N step %c1 {
  %val = memref.load %A[%i] : memref<?xf32>
  memref.store %val, %B[%i] : memref<?xf32>
}

关键变化：

1. affine.for → scf.for：仿射上下界被求值为具体的 index 值。
2. affine.load/affine.store → memref.load/memref.store：仿射下标被求值为线性的 index 值。
3. affine.if → scf.if：仿射条件被求值为 i1 布尔值。

降阶后，IR 不再包含任何仿射约束——因此编译器不能再做精确的依赖分析。这是”收益递减”的分界线：在 Affine 层做完所有需要精确依赖分析的优化（tiling、fusion、parallelization），然后降阶到 SCF 继续做通用优化（CSE、loop invariant code motion）。

七、优化策略分层

抽象层级           优化操作
──────────────────────────────────────────
tensor            算子融合（fusion）、公共子表达式消除
linalg            分块（tiling）决策、向量化方向选择
affine            依赖分析、分块调优、循环融合、并行化
scf               循环展开、循环不变代码外提
cf + arith        指令选择、寄存器分配（由 LLVM 完成）

每一层方言的优化 Pass 只利用自身能获得的语义信息，向下层传递优化后的 IR。这就是”渐进降阶”的精髓——不是一次性地把高层 IR 低化，而是分层、分批、每层做完优化再降。

八、本篇后续

上一篇（Tensor & Linalg）和这篇（Affine & SCF）覆盖了 AI 编译器核心的循环和计算抽象。下一章讲这些循环和计算最终怎么生成 GPU 代码——GPU 方言、SPIR-V 出口和内存层次抽象。

参考资料

官方文档（A 级）

• MLIR Affine Dialect — https://mlir.llvm.org/docs/Dialects/Affine/
• MLIR SCF Dialect — https://mlir.llvm.org/docs/Dialects/SCF/

论文（A 级）

• Bondhugula, U. et al. A Practical Automatic Polyhedral Parallelizer and Locality Optimizer. PLDI 2008. (PLuTo — 多面体编译器，MLIR Affine 的优化策略参考来源)

源码（A 级）

• mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
• mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
• mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
• mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp