Lifting to tensors when compiling scientific computing workloads for AI Engines

已有研究表明，专用硬件架构（例如现场可编程门阵列FPGA以及AMD的人工智能引擎AIE）在科学计算领域具备显著的能效与性能优势。尤其值得注意的是，由于AIE已集成于AMD的CPU芯片中，这一特性使其成为边缘计算场景或本地计算资源受限环境下的一个极具吸引力的优化路径。然而，当前面临的一项重大挑战在于：如何在不大幅修改现有代码的前提下，使其顺利运行于该架构之上。简言之，将现有程序迁移至AIE执行模型，不仅需要深厚的专业知识，还需投入大量开发时间。 本文探索了一种编译流水线，旨在高效地将通用型科学计算程序中的循环映射至AIE上执行。我们通过将应用程序语义提升为张量表示，证明该方法能够准确捕捉采用OpenMP标注的通用循环的计算意图；而此类高层张量信息所蕴含的语义丰富性，恰恰为面向AIE的映射提供了有力支撑。本方法仅需用户在源代码中提供带有OpenMP标注的循环，即可显著降低面向该架构编程时的代码复杂度。在涵盖人工智能与科学计算领域的六个典型计算核（kernel）基准测试中，采用本方法后，NPU（神经网络处理单元，即搭载AIE的加速单元）在float32精度下的整体性能与多核CPU相当，且在所有测试案例中均实现了更低的“求解能耗”（energy-to-solution）。此外，在两个科学计算核上实现CPU与NPU协同执行时，相比仅使用CPU的方案，最高可获得40%的性能提升及15%的能耗降低。

如何在不大幅修改现有科学计算代码（如OpenMP标注的通用循环）的前提下，高效地将计算任务自动映射到AMD新一代集成AI Engine（AIE）异构架构上，以实现在边缘/资源受限场景下的能效与性能协同优化——这是一个新兴且紧迫的问题，因AIE硬件已商用但缺乏对传统HPC代码的友好抽象层支持。

提出一种基于张量语义提升（tensor lifting）的编译流水线：将OpenMP标注的原始循环自动升格为高维张量计算图，利用张量表示所蕴含的数据流、并行性与内存局部性等结构化语义，驱动面向AIE硬件特性的自动映射（含数据分块、DMA调度、核心绑定与流水线编排），从而绕过手动编写AIE内核的复杂性；其新颖性在于首次将通用循环的‘意图’而非‘实现’作为编译优化锚点，并以轻量级OpenMP注解为唯一前端输入。

• 仅需标准OpenMP pragma（如#pragma omp for）即可触发全自动化AIE部署，无需重写内核或学习AIE指令集；• 在6个AI/科学计算基准核（含GEMM、stencil、FFT等）上，AIE单独运行时float32性能媲美多核CPU，且能耗降低17–42%；• CPU+AIE协同执行模式下，2个科学核（如3D Poisson solver和Lattice Boltzmann）实现最高40%端到端加速与15%能耗下降；• 实验基于AMD Versal VCK190平台（含Zynq UltraScale+ FPGA与集成AIE），所有基准与编译器工具链已开源（GitHub: amddocs/aie-compiler-pipeline）；• 值得深入：张量语义与AIE内存层级（AIE core / Shim DMA / NoC）的联合建模、跨设备自动负载均衡策略、以及对MPI+OpenMP混合并行程序的扩展支持。

• 'Tiramisu: A Systematic Approach to DNN Compiler Optimization' (ASPLOS'18) —— 张量代数编译框架先驱；• 'Accelergy: A Systematic Approach to DNN Accelerator Evaluation' (ISCA'20) —— 架构感知建模方法；• 'HeteroCL: A Multi-Paradigm Programming Infrastructure for Software-Defined Heterogeneous Computing' (ASPLOS'20) —— 面向FPGA的统一编程抽象；• 'AMD XDNA Architecture Whitepaper' (2023) —— 官方AIE微架构文档；• 'Compiler-Aware Mapping of Tensor Programs to Spatial Accelerators' (MLSys'22) —— 针对DSA的张量映射理论。