累积两个 Tensor Core wmma::accumulator 片段

Question

假设我有两个

wmma::fragment<wmma::accumulator, 16, 16, 16, half> a, b;

实例（即

a

和

b

）。我将如何对

a

和

b

进行按元素相加并将结果存储回

a

？

Answer 1

wmma 片段实际上存储在寄存器中。东京工业大学的科学家开发了一个库 wmma_extension（除其他功能外）可以轻松地对 wmma 片段进行算术运算。

可以在这里找到该库：https://github.com/wmmae/wmma_extension

作为简单的单行代码（加上包含）进行算术运算如下所示：https://github.com/wmmae/wmma_extension/blob/main/docs/ops.md

科学家们在2023年发布了两篇相关论文：

• https://arxiv.org/pdf/2308.15152.pdf（减少共享内存占用以利用 Tensor Core 及其灵活的 API 扩展库的高吞吐量）
• https://arxiv.org/pdf/2203.03341.pdf（从 Tensor Core 恢复单精度精度，同时超越 FP32 理论峰值性能）