oneDNN在AArch64平台上matmul内核断言失败问题分析与解决

问题背景

在AArch64架构的CPU平台上，使用oneDNN进行矩阵乘法(matmul)运算时，当调用acl_matmul内核时会出现断言失败的问题。该问题主要影响与自注意力机制(SDPA/MHA/MQA)相关的图计算测试用例，导致多个测试用例被跳过或失败。

问题现象

当运行特定形状的矩阵乘法运算时，程序会在acl_matmul内核中触发断言失败，错误信息如下：

Assertion `FixedFormat || _B_transposed' failed.
Aborted (core dumped)

环境信息

• 硬件平台：AArch64架构CPU
• 操作系统：Ubuntu 22.04.1
• 编译器：GCC 11.4.0
• oneDNN版本：主分支最新提交(af1410c2)
• Compute Library版本：v24.11.1

问题复现

通过一个特定的矩阵乘法示例可以稳定复现该问题：

1. 矩阵维度：

   • 输入矩阵A：形状为[4,16,384,384]
   • 输入矩阵B：形状为[4,16,384,64]
   • 输出矩阵C：形状为[4,16,384,64]

2. 关键特征：

   • 使用OpenMP多线程并行(4线程)
   • 矩阵B使用了abdc格式(即对最后两个维度进行了转置)

问题分析

经过深入分析，发现问题根源在于Compute Library(ACL)的矩阵乘法内核实现。具体来说：

1. 当使用多线程并行执行时，某些线程配置会触发断言失败
2. 问题与矩阵B的转置状态有关
3. 某些特定形状(如384x64和64x384)可以正常工作
4. 减少线程数量可以暂时规避该问题

解决方案

ARM团队已经在Compute Library v52.0.1版本中修复了这个问题。主要修复内容包括：

1. 改进了矩阵转置处理逻辑
2. 增强了多线程环境下的稳定性
3. 修复了断言条件检查

验证结果

在升级到Compute Library v52.0.1后：

1. 原始复现用例可以正常执行
2. 相关断言不再触发
3. 多线程环境下运行稳定

后续工作

虽然大部分问题已经解决，但仍有少量测试用例需要进一步处理：

1. 完善剩余两个测试用例的支持
2. 在oneDNN中升级默认的Compute Library版本
3. 重新启用之前跳过的图计算测试用例

技术建议

对于使用oneDNN在AArch64平台上的开发者：

1. 建议使用最新版本的Compute Library(v52.0.1或更高)
2. 对于矩阵乘法运算，注意输入矩阵的布局格式
3. 在多线程环境下进行充分测试
4. 关注oneDNN的版本更新，及时获取稳定性改进

该问题的解决为AArch64平台上的深度学习计算提供了更好的稳定性和性能，特别是在处理自注意力机制等复杂模型时。