Spine-Triton查看kernel的汇编代码——矩阵乘法没用上矢量指令

成功地测量了Spine-Triton中编写的融合算子的性能之后,下一步就是要找到性能不如PyTorch算子的原因。这时候,就需要查看自定义的算子对应的汇编代码。

先上结论:

  1. PyTorch的算子torch.matmul()使用了矢量指令完成矩阵乘法
  2. Spine-Triton编写的融合算子中矩阵乘法部分是用标量指令实现的

再上图:

  1. PyTorch算子调用了void at::native::cpublas::(anonymous namespace)::gemm_core_<float, float, float>,汇编代码:
  2. Spine-Triton编写的融合算子汇编代码:

最后给出查看汇编指令的方法:

  1. 通过perf查看PyTorch算子对应的汇编语句
    安装和使用perf的命令可以参考Bianbu RDK的文档perf + FlameGraph 使用教程和论坛的帖子Almost same execution time when enabling auto-vectorization or not。向大佬们致敬!
perf record -e u_mode_cycle python3 python/examples/test_fused_mm_softmax.single_run.py
perf report


  1. 通过ir_dumps获得编译后的二进制文件,反汇编查看指令
export SPINE_TRITON_DUMP_PATH=./ir_dumps
mkdir ir_dumps
python3 python/examples/test_fused_mm_softmax.single_run.py
objdump -D -b binary -m riscv:rv64 ir_dumps/bare_fused_mm_softmax_.o > ir_dumps/bare_fused_mm_softmax_.o.dmp


ir_dumps.zip (96.7 KB)

汇编指令说明
vfmacc :: Vector Floating-Point Instructions — RVV-web-page documentation
fmadd.d :: RISC-V Specification for example_rv64_with_overlay

题外话:perf生成的trace是可以导出为火焰图的。但是对于RISC-V架构中Python函数名解析的支持还未生效。

PYTHON_PERF_JIT_SUPPORT=1 perf record -F 99 -a -g -e cpu-clock python3 python/examples/test_fused_mm_softmax.single_run.py
perf script > out.perf
mv out.perf ../perf/.
cd ../perf/.
./git-FlameGraph/stackcollapse-perf.pl out.perf > out.folded
./git-FlameGraph/flamegraph.pl out.folded > out.svg

用浏览器打开out.svg,搜索gemm,就可以找到PyTorch算子对应的二进制程序,如下图中紫色部分所示:


上面的命令中,虽然使用环境变量PYTHON_PERF_JIT_SUPPORT=1打开了对Python函数名解析的支持,但是上图中还是存在大量的[python3.13]这样的缺省命名。
尝试查看perf对Python函数名解析的支持,得到如下的结果:

虽然Python已经支持perf了:Python 对 Linux perf 性能分析器的支持 — Python 3.12.12 文档
但是RISC-V架构上还是空白:聊聊 Python 3.12 中 perf 的原生支持 | Manjusaka

1 个赞

您可以指定arch_id,看看会不会成功使用矢量指令,参考 spine-triton/python/accuracy/flag_gems_abs.py: from triton.backends.spine_triton.driver import CPUDriver
triton.runtime.driver.set_active(CPUDriver())
driver = CPUDriver()
driver.set_current_arch_id(“0xA03C”)
triton.runtime.driver.set_active(driver)。 下周我们会发布新的版本不需要人为指定arch_id

请问,为什么会有不同的arch_id?是进迭不同的芯片架构吗?后续会给出不同架构的差异说明吗?

确保每次都能生成IR和Binary:

rm -rf ~/.triton/cache/
SPINE_TRITON_DUMP_PATH=./ir_dumps python3 python/examples/bare_matmul.py