RISC-V矩阵扩展深度解析:从向量瓶颈到二维tile的GEMM性能跃升
2026/9/17 7:39:46 网站建设 项目流程

几个月前我在一个 RISC-V 核上折腾向量指令集(RVV)的 SGEMM(单精度矩阵乘),一开始觉得 32 个向量寄存器、最高 512 位向量宽度摆在那里,性能怎么说也不会太差。结果发现无论怎么调循环展开和寄存器分块,浮点峰值利用率就是上不去。后来我把目光转向 RISC-V 正在推进的矩阵扩展,才意识到问题不在向量指令用得不熟,而在于向量这个抽象本身就不适合矩阵运算。矩阵扩展不是给向量“加宽”,而是对整个计算模型的一次升维:从一维的向量,走向二维的矩阵。

我写这篇文章的目标很直接:把 RISC-V 矩阵扩展背后“为什么要这么设计”讲清楚。你会看到向量扩展在哪里触到了天花板、矩阵寄存器(tile)到底解决了什么结构性问题、它与 ARM SME、x86 AMX 有哪些同与异,以及上手写 GEMM 时真正需要注意的工程细节。无论你是做 AI 加速器、嵌入式高性能计算的架构师,还是正在评估下一代 RISC-V 处理器的应用开发者,这篇文章应该都能给你一些可用的判断框架。

1. 向量时代的天花板:RVV 在矩阵计算上的瓶颈

1.1 向量是一条线,矩阵是一块面

RVV 的设计有很多优点,可变向量长度(VLEN)、丰富的元素宽度(int8 到 fp64)、可配置的 LMUL(向量长度乘数),这些让它在数据流处理、多媒体编解码、信号处理等场景里表现得很出色。但它的核心抽象是“一维寄存器数组”,v0 到 v31 这 32 个寄存器无论怎么组合,本质上都是线性的。

矩阵运算则是典型的二维数据流。一个 C[m][n] += A[m][k] * B[k][n] 的计算,输出是一个平面,参与计算的 A 和 B 分别在不同维度上被反复复用。当我们用向量指令去实现矩阵乘时,实际上是在做什么?是把二维的计算强行“拍扁”成一维的数据流,用手工循环去管理行和列的对应关系。

拿一个最朴素的外积版本来说,计算 4x4 累加块,你需要先把 A 的一列广播到向量里,再把 B 的一行装进向量,然后执行乘加,循环 4 次。这几条指令看起来不复杂,但问题在于:向量寄存器根本没有描述“二维坐标”的能力,你必须在指令外面维护循环变量,每一次迭代都要重新加载或重排数据。寄存器利用率一高,编译器就很容易陷入局部溢出,性能断崖式下跌。

1.2 复用率:矩阵计算性能的真瓶颈

做矩阵乘调优的人都知道,GEMM 的浮点峰值利用率很大程度取决于数据复用。具体说有三层复用:A 的每个元素被同一行的所有输出列复用,B 的每个元素被同一列的所有输出行复用,中间结果则在寄存器里持续累积。理想的硬件应该让这些复用在靠近运算单元的地方完成,而不是每一轮都从 L1 或 L2 重新拉数据。

在 RVV 的模型下,数据复用完全靠软件控制。假设 VLEN=256 位,一次能装 8 个 fp32,那么一个 8x8 的输出块就需要 8 个累加向量寄存器、8 个 A 切片、以及不断轮换的 B 切片。这个分配本身就接近寄存器文件的上限了。更大的输出块(比如 16x16)基本不可能纯用寄存器做,必须拆到 L1 层做分块。这就是为什么纯 RVV 写 GEMM 的代码,再怎么优化,峰值利用率通常卡在 60% 到 70%,不是指令执行不够快,而是数据搬运指令占了太多发射槽位。

1.3 向量指令的另一个隐性成本:取指与解码带宽

还有一个容易被忽视的问题:向量指令需要显式地描述数据移动。加载一个 8x8 的块,你需要 8 条 load 指令,每条还带不同的地址偏移。这些指令虽然能在流水线里重叠执行,但它们占据了取指带宽、解码带宽和发射队列。相比之下,如果硬件能提供一条“加载整个二维 tile”的指令,把 64 个数据一次性装进一个二维寄存器阵列,那总指令数可以下降一个数量级。

这就是矩阵扩展出现的最根本原因:不是向量扩展不好,而是矩阵计算值得拥有一个更高层次的抽象。把二维的数据块、二维的寄存器、二维的复用关系直接写进指令集语义,让硬件有能力在微架构层面做数据重排和调度,这比软件在高维循环里艰难挣扎要高效得多。

2. 二维寄存器视野:矩阵扩展的核心抽象设计

2.1 Tile:把寄存器堆从“数组”升级为“棋盘”

RISC-V 矩阵扩展最核心的抽象,是一个二维的寄存器块,社区通常叫它 tile。你可以把它想象成一块长宽可配置的“棋盘”,每个格子是一个标量数据。和 RVV 的向量寄存器一样,tile 的尺寸不应该是固定的,否则不同微架构的硬件实现会很难取舍。

具体来说,tile 的维度通常由配置寄存器(CSR)决定,比如定义 M、N 和元素位宽。软件在进入矩阵运算前先配置好 tile,硬件在此基础上把物理寄存器文件划分成若干个二维块。这个思路和 ARM SME 的 ZA 矩阵寄存器、x86 AMX 的 TMM 寄存器是一脉相承的。

如果你用惯了向量寄存器,可以这样理解:向量寄存器是一排连续编号的货架,矩阵寄存器是一块平面仓库,地址由“行+列”两个坐标组成。前者对一维数组友好,后者对二维矩阵天然友好。矩阵加载/存储指令可以一次性把一块子矩阵从内存搬进仓库,或者在仓库里做行列变换,而不用软件逐行倒腾。

2.2 为什么 tile 的尺寸必须是“可配置”的

RISC-V 的一个设计哲学是“不要用单一固定规格约束所有实现”。矩阵扩展如果定义成固定的 16x16 或 8x8,那么某些追求低面积的嵌入式核就吃不消了。所以社区草案普遍采用“可伸缩”的设计:tile 的行数、列数、元素位宽都允许实现者在一定范围内选择,软件在运行时通过 CSR 查询当前硬件支持的形状。

这一点很像 RVV 里的 VLEN。写代码时你不能假定向量一定是 128 位还是 512 位,而是用 vsetvli 动态查询。矩阵扩展也会提供类似的机制,比如配置指令里带上 M、N,硬件返回当前实际能够支持的 tile 形状。这意味着同一个二进制,理论上可以跑在 8x8 的最小实现上,也能跑在 64x64 的高性能实现上。

从编译器和运行时库的角度看,可变 tile 大小并没有想象中那么可怕。GEMM 库里本来就有分块维度参数——直接用 tile 的硬件能力作为分块参数,代码会更干净,甚至可能消除一部分运行时分支。

2.3 Tile 与向量寄存器之间应该有“桥梁”

矩阵扩展并不是要取代 RVV。相反,它的设计需要和向量寄存器紧密配合。为什么?因为矩阵运算的边界和预处理(比如量化、layout 转换、激活函数)依然要大量使用逐元素操作,这些用向量指令更顺手。

理想情况下,tile 的每一行(或者每一列)可以被映射到向量寄存器上。这样当你需要做一次逐元素的 ReLU、归一化或类型转换时,可以把 tile 的某一行“掰”出来,交给向量指令处理,处理完再放回 tile。如果把矩阵扩展设计成一套完全独立的寄存器文件,与向量引擎老死不相往来,那软件就得在两种状态间反复搬数据,性能和易用性都会大打折扣。

我在评估 ARM SME 时对这个“桥梁”设计印象很深。SME 的 ZA 寄存器按行切片之后,可以直接和 SVE2 的向量寄存器交互,编译器在一个基本块内自由切换矩阵和向量视角。RISC-V 矩阵扩展也应该走这条路线,否则就违背了 RISC-V 模块化、可组合的初衷。

3. 指令集层面的关键取舍:从 MMA 到 tile 管理

3.1 矩阵乘加指令的核心语义

矩阵扩展要解决的主要问题,肯定不是“做一条能算矩阵乘的指令”这么简单,而是要定义一组完整、可组合的操作原语。最核心的当然是矩阵乘加(MMA)类指令,语义类似:

tile_C += tile_A * tile_B

这条指令看起来很直观,但里面藏着几个关键设计决策。

第一个是累加器宽度。如果是整数矩阵乘,输入是 int8 或 int16,累加器一般需要加宽到 int32,否则多次累加会溢出。这在 GEMM 里几乎是标配。浮点场景则常见 fp16/bf16 输入、fp32 累加的组合。所以 MMA 指令必须以“输入类型+累加类型”的方式定义,而不是单一类型。

第二个是操作数的形状。一个 MxN 的 C 块,乘以一个 MxK 的 A 块和 KxN 的 B 块,这里的 M、N、K 分别取多大,会直接影响硬件乘法阵列的规模和微架构设计。常见的选择是 8x8x8、16x8x8、8x16x16 这类组合。M 和 N 决定输出 tile 的尺寸,K 决定每次乘加累积的深度。K 越大,数据复用率越高,但硬件上的输入缓冲区也越大。

3.2 Tile 的加载、存储与重排

光有计算指令远远不够,tile 怎么进寄存器、怎么出寄存器、怎么在寄存器里做布局变换,这些都是指令集层必须回答的问题。

矩阵加载/存储指令需要支持二维寻址。具体来说,能从内存加载一个“按行跨步”的二维块,也就是说每一行的起始地址可以通过 stride 参数描述,而不要求数据在内存里连续排列。没有这条指令的话,加载一个矩阵子块就退化成逐行 load,跟向量方案没有本质区别。

此外,GEMM 里有一件特别烦人的事:数据布局的转换。A 矩阵和 B 矩阵在内存中的排列通常是行主序,但某些切分策略(特别是对 B 做列访问时)需要转置或重排。如果矩阵扩展提供了在寄存器或内存层面处理转置的指令,软件就不需要额外做一遍完整的 pack 操作来处理转置。这个概念上类似向量指令里的 vtranspose,只不过是把维度从一维数组扩展到二维块。

3.3 边界场景怎么处理:tile 无法整除时

矩阵扩展最容易被忽略却又特别实际的问题,是边界处理。真实世界的矩阵尺寸千奇百怪,很少能恰好被 tile 尺寸整除。一个 15x17 的矩阵,用 8x8 的 tile 切分,必然会留下边角料。

一种方案是让软件自己处理边界:用 tile 指令先算完整部分,剩余小矩阵回退到 RVV 向量实现。这要求 tile 指令对操作数尺寸有严格约束,或者提供“仅使用 tile 左上角子区域”的计算模式。第二种思路是像 RVV 的尾巴处理(tail)那样提供掩码机制,但二维掩码的编码复杂度和硬件成本比一维多不少,所以多数草案倾向于让软件处理边界。

从实际工程经验看,让软件处理边界是合理的。GEMM 的高性能路径本来就要做“清理”循环,硬件为了边界情况去支持动态掩码,会显著增加乘法阵列控制逻辑的复杂度,收益却不高。真正需要注意的是,回退到向量的路径要足够高效,不能因为矩阵扩展的引入,把原本 RVV 的边界路径搞得比之前还慢,那就得不偿失了。

3.4 与向量扩展如何协同

在目前可见的 RISC-V 矩阵扩展讨论中,一个比较成共识的观点是:矩阵扩展建立在一组新 CSR 和指令之上,但它和向量扩展的协同方式必须从一开始就定义好。

我前面提到 tile 行可以映射到向量寄存器,这是协同的关键。另一方面,矩阵运算中的数据搬运(比如把数据按 K 维展开、广播到乘法阵列)不应该让软件逐元素去做,而是由硬件在 tile 内部自行完成。这就把指令集的功能边界划分清楚了:向量负责灵活的逐元素与数据预处理,矩阵负责数据密集的乘加流,两者各司其职。

对于跑在嵌入式平台上的轻量实现,矩阵扩展可以只提供最小子集:一个小尺寸的 tile、一条 MMA、一条 load/store,甚至可以直接复用向量寄存器作为累加器。这不是偷工减料,而是 RISC-V 模块化哲学的体现——没有哪个扩展是必须“全有或全无”的。

4. 与 ARM SME / x86 AMX 的同与异

4.1 三套方案的共同底层逻辑

在 RISC-V 之外,ARM SME 和 x86 AMX 是两组非常典型的矩阵扩展参考实现。三套方案虽然指令编码和寄存器布局完全不同,但底层的思路高度一致:把二维数据块放到离乘法阵列最近的寄存器文件里,用专门指令驱动大块乘加运算,从而减少数据搬移和指令开销。

ARM SME 建立在 SVE2 基础之上,引入的是 ZA(Z Array)矩阵寄存器。ZA 是一个可伸缩的二维数组,尺寸与 SVE 的向量长度相关。x86 AMX 则是 Intel 在服务器处理器上推的方案,提供一组固定大小的 TMM 寄存器(典型实现是 16 行乘以 64 字节,也能表达成 16x16 的 fp32 或 16x64 的 int8),配合 TMUL 单元执行矩阵乘。

RISC-V 的设计虽然尚未完全定稿,但各厂商展示的预研实现(包括玄铁 C908 提供的 INT8 矩阵运算单元)也都落到同一套“二维 tile + 乘加指令 + 显式搬入搬出”的框架里。可以说,矩阵扩展不是某个公司的奇技淫巧,而是整个行业对“矩阵计算密集化”这个趋势的共同回答。

4.2 关键差异:固定尺寸与可伸缩尺寸的路线之争

三套方案最明显的分叉点,在于 tile 尺寸的固定性。

x86 AMX 选择了固定尺寸,TMM 寄存器形状在架构上确定,软件不需要查询动态参数。这对编译器、操作系统的上下文切换逻辑来说都更简单,代价是后续硬件如果要扩展尺寸,指令集层面必须做兼容性处理,而且固定尺寸的硬件在灵活性上一定有所牺牲。

ARM SME 选择让 ZA 尺寸跟随 SVE 向量长度,也就是实现了“可伸缩但受实现约束”的模型。软件在运行时通过指令查询当前实现支持的形状,动态调整分块策略。这对跨微架构的代码移植更友好,但也让编译器后端多了一堆形状推断和分支处理。

RISC-V 目前讨论的方向更偏向 SME 这类可伸缩路线。理由很容易理解:RISC-V 的生态环境极度碎片化,从极低功耗的微控制器到高性能服务器核都会用到,固定尺寸根本无法兼容这些差异巨大的实现。软件必须做好“拿到什么形状就用什么形状”的准备。

4.3 对 RISC-V 的启示:不要重造轮子,但要保持开放

我经常看到一些人问:“既然 SME 和 AMX 都做出来了,RISC-V 为什么不直接抄?”这句话只对了一半。矩阵扩展的宏观方向确实可以借鉴,但在细节上,RISC-V 有自己的考量和优势。

ARM SME 和 SVE2 是强绑定的,要么都支持,要么都不支持。x86 AMX 则依赖 Intel 的固定实现。RISC-V 的矩阵扩展完全可以做成一个独立的、仅依赖少量基础指令的扩展,让只想要整数矩阵加速的嵌入式核不必为了浮点 SME 那套复杂机制付出代价。这种“按需组装”的灵活性,是其他指令集架构很难复制的。

另一方面,RISC-V 的矩阵扩展有机会把“从向量到矩阵的转换路径”设计得更顺滑。在 SME 和 AMX 的早期实现里,把普通内存数据填充进 tile 的高速通道都是经过反复优化的,RISC-V 可以吸取这些经验,直接在指令语义层面把二维加载、跨步寻址、以及 tile 到向量的映射定义清楚,减少后续软件栈的补丁式适配。

5. 编程模型落地:切块循环与数据布局怎么做

5.1 从数学 GEMM 到 tile 循环的伪代码

说了这么多架构层面的思路,最终代码还是要落到循环结构上。经典的 GEMM 有三层循环(i/j/k 或 k 在外层的外积形式),到了矩阵扩展上,循环结构往往变成:外层遍历输出 tile,内层遍历 K 的切块。一个概念性的伪代码如下:

// 配置 tile 形状:假设 M_tile=16, N_tile=16, K_tile=8 config_tile(M=16, N=16, K=8, dtype=fp32); for (m0 = 0; m0 < M; m0 += 16) { for (n0 = 0; n0 < N; n0 += 16) { tile_c.clear(); // 把 16x16 累加块清零 for (k0 = 0; k0 < K; k0 += 8) { tile_a.load(&A[m0][k0], stride=K); // 16x8 tile_b.load(&B[k0][n0], stride=N); // 8x16 tile_c.mma(tile_a, tile_b); // 16x16 += 16x8 * 8x16 } tile_c.store(&C[m0][n0], stride=N); } }

这段代码的关键在于:所有二维数据的搬运和计算都发生在 tile 层面,循环体里没有逐指令的行列处理。编译器只需把 tile 寄存器映射到物理寄存器文件,就能生成一段非常紧凑的指令序列。实际性能是否理想,则取决于 tile_a、tile_b 的数据是否能被硬件高效地送往乘法阵列——这正好是矩阵扩展微架构发挥的地方。

5.2 打包与布局:共享内存和 K 维度 buffering

即使有二维 load 指令,GEMM 里 PACK(打包)步骤也没有完全消失。原因很简单:硬件 tile load 之后,往往还是希望数据在内存中呈现某种连续或跨步一致的布局,以减少地址计算和 TLB 开销。

所以实践中通常的做法是:在进入热循环之前,把 A 和 B 的分块预取到一块精心设计的 workspace(共享内存区域)里,按照硬件友好的布局重新排列,然后再用 tile load 指令把数据搬到寄存器。这一步看起来绕,但很值得。数据如果排列得足够规整,K 维度上重复循环时,tile load 几乎可以做到零额外地址计算,这对流水线连续运行非常有帮助。

另外,我建议在做矩阵扩展编程时,把“K 维作为最高频缓存维”记在脑子里。A 和 B 矩阵在 K 方向上每推进一个 tile 块,所有已加载的输出 tile 都要复用一组新数据。K 分块的大小直接决定了寄存器里能同时保持多少个输出 tile,K 太大,寄存器压力飙升;K 太小,数据搬运占比上升。没有一个万能参数,只能在具体芯片上用性能计数器逐个试。

5.3 编译器支持与 intrinsics

指令集再好,如果编译器这块跟不上,落到实际开发的体验会很难受。目前 LLVM 社区对 RISC-V 矩阵扩展的讨论还处于特性设计阶段,离成熟而且大家熟悉的 auto-vectorization(自动向量化)还有距离。作为一个实际写代码的人,我给出的建议是:现阶段不需要期待编译器帮你从标量循环自动生成矩阵指令,而是专注于intrinsics(内建函数)。

内建函数的设计一般和指令一一对应,比如__riscv_mma_tile_c += tile_a * tile_b这类语义。写起来比纯汇编舒服,又能精确控制 tile 的加载和释放。提前熟悉这种开发方式,等规范正式落地时,你的代码基础就已经在了。

有一个工程细节值得注意:tile 配置 CSR 的切换开销。假如你在同一个循环里要反复切换不同形状的 tile(比如先算 16x16 的 fp32,又算 32x8 的 int8),每次切换 CSR 都可能让流水线停顿。更好的做法是尽量把同形状的运算归类到一起,减少配置切换次数。这一点很多文档不会主动提,但在调优阶段几乎一定会遇到。

6. 软硬件协同的工程深水区与我的看法

6.1 硬件实现:矩阵单元到底贵不贵

我常被问到的一个问题:“想要矩阵扩展,芯片面积和功耗上要付出多大代价?”这个问题的答案取决于矩阵单元的规模。一个支持 4096 次乘加/周期的矩阵单元(对应 64x64x1 的 int8 阵列)显然会很贵,但如果只做一个 8x8 的 INT8 矩阵单元,面积大概只相当于一个中等规模的浮点向量 ALU。

从这里可以看出 RISC-V 矩阵扩展理想的落地方式:让每个微架构自己去权衡。低成本 AIoT 芯片只需要支持很小的 tile 形状,服务器芯片则可以把 tile 做得很大,并增加更复杂的乘法阵列。指令集层面提供统一抽象,硬件实现各自发挥。

过去我在评估类似设计时发现,最容易被低估的是寄存器文件的端口(port)和重命名压力。tile 寄存器不像向量寄存器那样一条指令访问一两个寄存器就够了,MMA 指令一次要读多个 tile 数据。如果微架构没有给矩阵寄存器文件足够的读写端口,硬件利用率会非常难看。这个问题不像指令集设计那样在文档里一眼能看到,但它是实现性能的关键。

6.2 操作系统与工具链的隐性负担

矩阵扩展引入一个新的可配置状态(CSR),这给操作系统带来了一个很现实的问题:上下文切换时必须保存和恢复这些状态。如果操作系统不感知矩阵扩展状态,线程切换后 tile 配置和内容就会错乱,程序直接崩溃。

这和当年 AVX-512 刚出来时的遭遇如出一辙:频繁的寄存器状态保存会拉高上下文切换开销,操作系统需要提供一种机制,要么惰性保存,要么让应用显式标记自己的寄存器使用范围。RISC-V 的设计者有必要把这一点考虑在内,否则高性能矩阵应用会在多进程环境下被系统调用的开销拖累。

工具链方面,调试器和性能分析器需要对 tile 状态有可观测性。以前用 gh 或 perf 看一个向量程序,寄存器状态一目了然;到了矩阵扩展,如果分析工具不知道 tile 布局,定位性能瓶颈会很吃力。这些问题虽然不直接影响指令集设计,但决定了一个新扩展能否被开发者在真实产品中广泛采用。

6.3 我对于现阶段要不要投入的判断

经常有人问:“RISC-V 矩阵扩展规范还没定稿,我现在就学是不是太早?”我的回答是:如果你是在做 AI 推理库、高性能计算底层库,或者计划在两三年后推出 RISC-V 处理器芯片,现在就是最好的研究时机。

但也要做好心理准备:各家草案可能存在较大差异,你基于某一种草案写的代码可能在最终规范里不能直接编译。有效策略是抽象出一层薄薄的矩阵计算接口,把指令相关的细节隔离起来,为未来替换留下余地。我在自己项目里就是用一个 C++ 的MatmulKernel封装,底下可以编译到 RVV 向量实现,也可以编译到某个矩阵扩展的草案实现,上层完全无感。

矩阵扩展给 RISC-V 带来的不只是多几条新指令,而是让整个生态意识到:矩阵计算的密度已经是现代计算的核心指标了。作为一个已经写了多年高性能代码的人,我非常看好这个方向,但我也清楚它真正成熟还需要时间。如果你也想在这个领域施展拳脚,我建议先从一手草案读起,然后在模拟器或 FPGA 上亲手跑一个最小的 GEMM 用例,把整套工具链走通一次。踩过一次坑之后,你对矩阵扩展的理解会远超那些只看论文的人。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询