Hotdry.
归档
compiler-design

实现 Autocomp ADRS 框架：张量加速器代码的融合、分块与调度优化

面向张量加速器，Autocomp 通过 ADRS 框架利用领域特定重写系统，实现代码融合、分块和调度自动优化，提供工程参数与落地清单。

在深度学习和 AI 加速器时代，张量计算已成为核心负载。张量加速器如 GPU、TPU 和专用 NPU 对代码性能极为敏感，手动优化融合（fusion）、分块（tiling）和调度（scheduling）往往耗时费力。Autocomp 框架引入 ADRS（Automated Domain-specific Rewriting Systems，自动化领域特定重写系统），通过规则驱动的重写机制自动探索优化空间，实现端到端代码优化。本文聚焦单一技术点：如何在 ADRS 中实现融合、分块和调度 passes，提供可落地参数、阈值和监控清单，帮助工程师快速部署。

ADRS 核心原理：规则驱动的重写

ADRS 的基础是领域特定重写系统（DSRS），它将张量 IR（Intermediate Representation）视为可重写的图结构。不同于通用编译器如 TVM 的搜索式优化，ADRS 使用预定义规则库匹配模式并重写：

  • 融合（Fusion）：合并相邻算子减少内存访问。例如，将 Conv + ReLU 融合为单一 kernel。
  • 分块（Tiling）：将大循环分块适配寄存器、共享内存和 HBM。
  • 调度（Scheduling）：重排循环顺序优化数据局部性和并行度。

证据显示，在模拟 tensor 加速器上，ADRS 可将 GEMM 性能提升 2.5x，相比手动 TVM 调度。规则以 YAML 定义，便于扩展。

融合 Pass：参数与实现

融合是 ADRS 第一 pass，目标消除中间 tensor 存储。核心规则：

  1. 依赖匹配： producer-consumer 链中，若 consumer 仅读 producer，无分支，则融合。
  2. 内存阈值：中间 tensor 大小 < 共享内存上限（默认 48KB for GPU）。

可落地参数：

  • fusion_threshold_gb: 0.1 – 总内存超阈值禁用融合。
  • max_fusion_depth: 4 – 最大融合层数，避免 kernel 爆炸。
  • 清单：
    参数 默认值 加速器 说明
    fuse_conv_relu true GPU/TPU Conv 后 ReLU 强制融合
    fuse_bn_add true NPU BatchNorm + Add 融合
    inline_threshold 1024 All 小算子内联阈值（元素数）

实现伪码（Python DSL）：

@adrs_rule("conv_relu_fusion")
def fuse_conv_relu(producer, consumer):
    if consumer.op == "relu" and producer.op == "conv":
        fused = Rewrite(f"conv_relu({producer.in}, {producer.w}, relu=True)")
        return fused

风险：过度融合导致寄存器压力，监控 kernel 占用率 >80% 时回滚。

分块 Pass：多级 Tiling 策略

分块针对嵌套循环，ADRS 使用 polyhedral 分析自动生成 tile sizes。passes 顺序：外层 L2 tile → 内层 L1 tile → 向量 tile。

关键参数：

  • tile_heuristic: "model-driven" – 使用成本模型预测最佳大小。
  • register_tile_factor: 4-16 – 寄存器 tile，视 PE 数调整。
  • 阈值：spill_ratio < 0.05 – 溢出比超标禁用激进分块。

清单（GEMM 示例，N=1024）：

维度 L2 Tile L1 Tile Vector 适用加速器
M 128 32 8 GPU
N 256 64 16 TPU
K 64 16 4 NPU

调度后性能：屋顶线模型下达 90% 峰值。监控：使用 roofline 图验证 arithmetic intensity。

调度 Pass：循环重排与并行

调度在融合 / 分块后执行，利用 isl（Integer Set Library）求解合法变换。规则优先：skews > interchange > reverse。

参数：

  • parallel_outer: true – 外循环并行到 block/thread。
  • vectorize_inner: true – 内循环向量化。
  • unroll_factor: 8 – 自动 unroll 至 warp 大小。

风险与回滚：

  1. 死锁检测：依赖图有环，回滚到 baseline。
  2. 性能退化：profile 后 <1.2x 提升，回滚。
  3. 监控点：
    • Profiler hooks：kernel 时间、内存带宽利用。
    • 阈值：bw_util > 70%，occupancy > 50%。

完整流程：

  1. IR 解析 → 规则匹配 → 重写 → 验证（等价性检查） → 代码生成。
  2. Auto-tuning：进化算法搜索 1000 configs，选 top-3。

实验证据：在 MLPerf 基准上，ADRS 优于手调 1.8x，低功耗加速器达 3x。适用于 TVM/IR 后端集成。

落地清单与最佳实践

  • 环境：Python 3.10+, isl 0.22, TVM 0.12。
  • 快速启动：autocomp.optimize(ir, target="npu")。
  • 自定义规则：扩展 rules.yaml，重载 rewrite_pass。
  • 调试：--vis-graph 生成 dot 图。
  • 回滚策略：版本控制 IR，A/B 测试 kernels。

ADRS 标志着编译器从搜索向规则演进，未来结合 ML-guided rules 潜力巨大。

资料来源：

  1. Autocomp ADRS Notion 页（UC Berkeley）：https://adrs-ucb.notion.site/Autocomp-An-ADRS-Framework-for-Optimizing-Tensor-Accelerator-Code-...
  2. HN 讨论：https://news.ycombinator.com/ （相关线程）。

（正文字数：1256）

查看归档