探索Vortex开源GPGPU：RISC-V SIMT架构(4-2)，TCU 矩阵计算(2)

目录

前言

一、TCU模块框图

二、WMMA代码分析

2.1  WMMA矩阵分块

2.2  WMMA矩阵地址偏移计算

2.2.1  WMMA matrixA

2.2.2  WMMA matrixB

2.2.3  WMMA matrixC

2.2.4  tcu_int WMMA源代码

总结

前言

本篇分析Vortex矩阵计算的核心模块TCU WMMA。

前文：探索Vortex开源GPGPU：RISC-V SIMT架构(4-2)，TCU 矩阵计算(1)https://blog.ZEEKLOG.net/weixin_55313207/article/details/157395313

一、TCU模块框图

• dispatch_unit：把前级Issue_slice过来的请求经过仲裁器仲裁为1路
• pe_switch：分为前向和反向

前向：根据pe_select把请求分配到tcu_fp和tcu_int，pe_select是根据数据流里面的fmt_s[3]来判断(1: fp32，fp16，bf16;  0: int32， int8，uint8，int4，uint4) 反向：把tcu_fp(浮点WMMA)和tcu_int(定点WMMA)返回的结果经过仲裁器和时序缓冲器送到gather_unit

• gather_unit：根据数据流里面的isw(slice序号)把执行结果分配到对应的commit_if，再送给下游Commit模块

二、WMMA代码分析

为了简单起见，只分析tcu_int(定点) WMMA。

2.1  WMMA矩阵分块

matrixB垂直方向分为m_steps个小矩阵块

，

....

matrixA水平方向分为n_steps个小矩阵块

，

....

以n_steps=4为例，则matrixA x matrixB的最终结果可分解为4个小矩阵块的乘加。

对于每个

 x 

，如下图所示，我们有以下结论：

blkA 垂直方向有 m_steps 个小块，所以 blkA x blkB 总共需要 m_steps * n_steps周期计算出tileM行，tileN列结果，需要m_steps * n_steps 个FTR，既总共8个FTR存最终结果

  x 

 需要 n_steps周期计算出 tcM 行， (tcN*n_steps)列结果，占用n_steps个FTR寄存器

  x 

  有tcM*tcN个结果，既NT个结果，假定每个结果都是XLEN长度，因此需要分配一个线程寄存器(FTR)来存储

对于matrixC += matrixA x  matrixB，在垂直方向和水平方向各有 n_steps个

 ， 

，所以：

WMMA matrixC  总共计算周期为 m_steps * n_steps * n_steps = 8 * n_steps