随着空间架构（Spatial Architecture）在 AI 加速领域的兴起，AMD AI Engine（AIE）作为 Versal ACAP 平台的核心计算单元，为高性能线性代数运算提供了新的硬件基础。然而，要在这种非冯・诺依曼架构上实现高效的 BLAS（Basic Linear Algebra Subprograms）库，传统的矩阵分块策略需要重新设计。本文基于 AIEBLAS 开源项目，深入探讨针对 AIE 架构的矩阵分块优化策略。
AIE 架构特点与 BLAS 库需求
AMD AI Engine 采用空间数据流架构，在 Versal VCK5000 开发板上以 8×50 的网格排列 400 个 AIE 核心。每个 AIE 核心包含一个超长指令字（VLIW）向量处理器和 32KB 本地内存，相邻 AIE 之间可以直接读写对方的本地内存，非本地通信则通过 AXI4 Stream 实现。这种架构特点决定了：
- 数据局部性至关重要：32KB 本地内存容量有限，必须精心管理数据驻留
- 通信开销显著：片外内存访问（通过 PL-AIE 接口）带宽为 4GB/s，远低于片内通信
- 并行性天然存在：400 个 AIE 核心提供了大规模空间并行计算能力
AIEBLAS 项目正是为了在这样的架构上提供可重用、可定制的 BLAS 例程而设计的开源实现。项目采用 JSON 配置文件驱动代码生成，自动生成 AIE 内核、PL 内核、数据流图及构建系统。
当前设计：Windows 机制与数据流编程
AIEBLAS 的当前实现采用 "windows" 机制处理向量和矩阵数据。Windows 是存储在 AIE 本地内存中的数据块，相比通过 AXI4 Stream 逐元素传输，windows 允许使用更宽的数据通路访问，从而充分发挥 AIE 向量处理器的计算能力。
数据流编程模型
AIEBLAS 的核心设计理念是数据流编程。用户可以在 JSON 配置中指定 BLAS 例程之间的连接关系，代码生成器会相应创建数据流图定义。例如，计算axpydot（先执行向量加法axpy，再执行点积dot）时，两个例程可以直接通过片内通信连接，避免中间结果写回片外内存。
这种设计带来了显著的性能优势。实验数据显示，使用数据流实现的axpydot相比非数据流实现性能提升了一倍。然而，当前实现仍面临挑战：单例程性能相比 OpenBLAS 最多有 10 倍差距，这主要源于片外内存访问瓶颈。
矩阵分块策略的必要性与挑战
在传统 CPU/GPU 架构上，矩阵分块（tiling）是优化 BLAS 性能的经典技术。通过将大矩阵分解为适合缓存的小块，可以显著提高数据局部性，减少内存带宽需求。但在 AIE 架构上，分块策略面临独特挑战：
1. 内存层次结构差异
AIE 的内存层次相对扁平：
- 本地内存：32KB，访问延迟最低
- 相邻 AIE 内存：可通过直接读写访问
- 片外 DRAM：通过 PL 接口，带宽有限
传统 CPU 的多级缓存（L1/L2/L3）和 GPU 的全局内存 / 共享内存层次在 AIE 上不存在，需要重新设计分块策略。
2. 数据流与分块的协同
AIEBLAS 的数据流编程模型天然支持流水线执行，但如何将分块策略与数据流结合是一大挑战。理想情况下，每个数据块应该在整个计算流水线中流动，最大化硬件利用率。
3. 分块尺寸的约束
分块尺寸受多个因素限制：
- 本地内存容量：32KB 限制
- 向量宽度：最大 512 位（64 字节）
- AIE 阵列拓扑：8×50 网格影响通信模式
- PL-AIE 接口带宽：4GB/s
针对 AIE 架构的优化分块策略
基于 AIEBLAS 的设计和 AIE 架构特点，我们提出以下分块优化策略：
1. 多级分块架构
针对 AIE 的扁平内存层次，建议采用三级分块策略：
# 三级分块参数示例
TILE_DRAM = 1024×1024 # 片外内存分块，匹配PL接口突发传输
TILE_AIE = 256×256 # AIE间传输分块，匹配网络带宽
TILE_LOCAL = 64×64 # 本地计算分块，匹配32KB内存
第一级（DRAM 到 PL）：利用 PL 内核的突发传输能力，将大矩阵从 DRAM 加载到 PL 内存。分块尺寸应匹配 AXI 接口的最大突发长度，通常为 1024×1024 元素（FP32 约 4MB）。
第二级（PL 到 AIE）：通过多个 AIE-PL 接口并行传输，分块尺寸应考虑接口带宽平衡。256×256 的分块可以在多个 AIE 间均匀分布。
第三级（AIE 本地）：在单个 AIE 内部分块计算，64×64 的分块（FP32 约 16KB）可以完全容纳在 32KB 本地内存中，为输入 / 输出数据留出空间。
2. 分块形状优化
传统分块通常使用正方形分块，但在 AIE 架构上，矩形分块可能更优：
- 对于 GEMV（矩阵 - 向量乘法）：使用高瘦分块（如 256×16），减少 AIE 间的水平通信
- 对于 GEMM（矩阵乘法）：考虑 AIE 网格的 8×50 形状，使用 8×50 的倍数作为分块尺寸
3. 数据流感知的分块调度
将分块策略与数据流编程结合，需要设计智能的调度器：
class AIETileScheduler:
def schedule_tiles(self, matrix_a, matrix_b):
# 预取：提前加载下一个分块到PL内存
# 流水线：多个分块在不同AIE间并行处理
# 重叠：计算当前分块时传输下一个分块
pass
关键优化点包括：
- 预取机制：在计算当前分块时，预取下一个分块到 PL 内存
- 流水线深度：根据 AIE 数量确定最佳流水线深度，通常为 AIE 数量的 1/4 到 1/2
- 负载均衡：考虑 AIE 间通信延迟，动态调整分块分配
4. 分块参数调优指南
基于 AIE 架构约束，推荐以下分块参数：
| 参数 | 推荐值 | 约束条件 |
|---|---|---|
| 本地分块尺寸 | 64×64 | 32KB 本地内存限制 |
| AIE 间分块 | 128×128 | 相邻 AIE 通信带宽 |
| 向量化宽度 | 512 位 | AIE 向量处理器支持 |
| 流水线深度 | 8-16 级 | AIE 阵列规模 |
| 预取距离 | 2-4 个分块 | PL 内存容量 |
5. 监控与调优指标
实施分块策略后，需要监控以下关键指标：
- 本地内存利用率：目标 > 80%，但避免溢出
- AIE 计算利用率：通过性能计数器监控 VLIW 指令发射率
- 片外带宽利用率：监控 PL-AIE 接口的实际吞吐量
- 流水线气泡率：数据依赖导致的空闲周期比例
实现路径与工程挑战
将分块策略集成到 AIEBLAS 中面临多个工程挑战：
1. 代码生成器扩展
当前 AIEBLAS 的代码生成器基于模板，需要扩展以支持：
- 分块循环的自动生成
- 分块尺寸的参数化配置
- 预取和流水线控制逻辑
2. 内存管理复杂性
多级分块需要精细的内存管理：
- PL 内存分配：为多个分块预分配缓冲区
- AIE 内存映射：确保分块在 AIE 间的正确映射
- 数据一致性：分块边界的数据同步
3. 调试与验证
空间架构的调试比传统架构更复杂：
- 数据流可视化：需要工具可视化分块在 AIE 间的流动
- 性能分析：细粒度的性能计数器访问
- 正确性验证：分块算法的数学正确性证明
性能预期与优化目标
基于当前 AIEBLAS 的性能数据和 AIE 架构特点，实施优化分块策略后预期：
- 片外访问减少：通过多级分块和预取，预计减少 50-70% 的片外内存访问
- 计算利用率提升：更好的数据局部性可将 AIE 计算利用率从当前的 30-40% 提升到 60-70%
- 性能差距缩小：与 OpenBLAS 的性能差距从 10 倍缩小到 3-5 倍
关键的性能拐点出现在：
- 本地内存完全利用：当分块尺寸恰好匹配本地内存容量时
- 通信计算完全重叠：当数据传输时间完全被计算隐藏时
- 流水线完全饱和：当所有 AIE 都处于忙碌状态时
结论与展望
AMD AI Engine 为 BLAS 库的实现提供了独特的空间架构基础，但也带来了矩阵分块策略的新挑战。AIEBLAS 项目的 windows 机制和数据流编程为分块优化提供了良好起点，但需要深入的多级分块策略来充分发挥硬件潜力。
未来的优化方向包括：
- 自适应分块：根据矩阵尺寸和形状动态调整分块策略
- 混合精度支持：针对 FP16、BF16 等低精度格式优化分块
- 稀疏矩阵优化：针对稀疏 BLAS 的特殊分块策略
- 自动调优框架：基于机器学习的分块参数自动优化
随着空间架构在科学计算和 AI 领域的广泛应用，针对这类架构的 BLAS 优化将成为关键的基础软件技术。AIEBLAS 及其分块策略的研究不仅对 AMD AI Engine 有直接价值，也为其他空间架构（如 SambaNova、Cerebras）的 BLAS 实现提供了参考框架。
资料来源
- Tiziano De Matteis, "Developing a BLAS library for the AMD AI Engine Extended Abstract", arXiv:2410.00825v1, 2024
- AIEBLAS 开源项目：https://github.com/atlarge-research/AIE-BLAS
- AMD AI Engine 架构文档：UG1079 - AI Engine Kernel and Graph Programming Guide