wehub-resource-sync
227 lines
11 KiB
Markdown
227 lines
11 KiB
Markdown
# ARM CPU 低 bit GEMM Kernel (W2 / W3 / W4 / W8)
|
||
|
||
> 范围: ARMv8.2 **sdot** (`SRC_UNIT=4`) 与 ARMv8.6 **i8mm / smmla** (`SRC_UNIT=8`),FP32 与 FP16 输出路径,共 W4/W8 老路径 + 本次新增的 W2/W3 路径。
|
||
>
|
||
> W2/W3 来自 commit `85e3f63a252 [CPU:Feature] Support low-bit ARM quant kernels`;W4/W8 早已上线,本文一并汇总以便对比。
|
||
|
||
---
|
||
|
||
## 0. 总览
|
||
|
||
低 bit GEMM 把每个 OC 的权重压缩到 ≤ 1 byte,反量化与计算融合在汇编里完成。统一约定:
|
||
|
||
* 计算单元:`UNIT = 8` 路 OC × `SRC_UNIT` 路 IC,即 sdot 一个 cell 32 weight、smmla 一个 cell 64 weight。
|
||
* 权重在 cell 内 **同 IC、跨 OC pack** 进单 byte;OC 间隔取 2 是为了让单条 `ushr/and` 直接吐出 smmla 期望的 OC pair。
|
||
* Kernel 输出 **unsigned [0, 2^bits − 1]**;`_computeReorderQuantInfo` 写入 `originOffset = -(2^bits / 2)`,在 post‑process 用一次 `MLA_WEIGHTZERO` 把 zero point 重定心到 signed,kernel 内部不出现 sub / 符号扩展(W3 unsigned-domain 是关键)。
|
||
* Post‑process 链统一为 `Int32ToFloat → MUL_SCALE → MLA_WEIGHTZERO → MUL_INPUT_SCALE → ADD_BIAS → ReLU`,FP16 路径多一组 `fcvtn / fcvtl`。
|
||
* W4/W8 成熟阶段已把 per-block dequant 链最末的 `fadd v_acc, v_per_block` 融进前面的链(`fmla v_acc, v_per_block, scale_eff`),W2/W3 直接复用同一融合后处理。
|
||
|
||
| 位宽 | weightBytes (有效) | originOffset | smmla 单 cell 字节 | sdot 单 cell 字节 | 解包关键宏 |
|
||
|------|--------------------|--------------|----|----|----|
|
||
| W8 | 1.00 | 0 | 64 | 32 | 不需要解包 |
|
||
| W4 | 0.50 | -8 | 32 | 16 | `ushr/and #15` |
|
||
| W3 | 0.375(sdot 实占 0.5,4B pad) | -4 | 24 (16 main + 8 aux) | 12 (8 main + 4 aux,padded 16) | `W3_UNPACK_SERIAL` / `UNPACK_W3_SDOT` |
|
||
| W2 | 0.25 | -2 | 16 | 8 | `ushr/ushl/and #3` / `UNPACK_W2_SDOT` |
|
||
|
||
---
|
||
|
||
## 1. 数据排布
|
||
|
||
### 1.1 W8 — int8 直存
|
||
|
||
* 一个 byte 一个 weight,`reorderWeight` 直接打成 `(cellCount, UNIT × SRC_UNIT)` 排布,smmla cell 64 B,sdot cell 32 B。
|
||
* signed int8;无需 `originOffset`。
|
||
|
||
### 1.2 W4 — 单 byte 2 路 OC
|
||
|
||
```
|
||
W4 byte = | bits[7:4] = OC 偶号 | bits[3:0] = OC 奇号 |
|
||
```
|
||
|
||
* **smmla cell (32 B)**:`byte i (0..15)` → IC = i,OC 偶号;`byte i (16..31)` → IC = i-16,OC 奇号。
|
||
* **sdot cell (16 B)**:`byte i (0..7)` → IC = i,OC 偶号;`byte i (8..15)` → IC = i-8,OC 奇号。
|
||
* 解包仅 `ushr v,#4` + `and v,v_mask15` 两条指令,产出两个 OC pair。
|
||
|
||
### 1.3 W2 — 单 byte 4 路 OC
|
||
|
||
```
|
||
W2 byte: bits[7:6]=oc01 bits[5:4]=oc23 bits[3:2]=oc45 bits[1:0]=oc67
|
||
```
|
||
|
||
* **smmla cell (16 B)**:`byte 0..7` → IC=0..7, OC=0,2,4,6;`byte 8..15` → IC=0..7, OC=1,3,5,7。
|
||
* **sdot cell (8 B)**:`byte 0..3` → IC=0..3, OC=0,1,2,3;`byte 4..7` → IC=0..3, OC=4,5,6,7。
|
||
* OC 走 2-step 间隔是为了让 smmla 的 `ushr v0,#6 / #4 / #2 / and v0,#3` 4 步 ALU 直接拿到 4 个 OC pair。
|
||
|
||
### 1.4 W3 — bit-plane 2 + 1 拆分
|
||
|
||
W3 不能 4 路并打,改成位平面拆:**低 2 bit 与 W2 同 layout 存为 main**,**高 1 bit 单独存 aux**。
|
||
|
||
* **smmla cell (24 B)** = 16 B main + 8 B aux
|
||
* aux IC-major:`byte i (i=0..7)` 存 8 个 OC 在 IC=i 的高 bit,bit `(7-j)` 表示 OC `j`。
|
||
* 选 IC-major:i8mm 用 `ld1r {v1.2d},[x2],#8` 复制成 16B 后,通过 `ushl + and` 链 4 次抽出 4 个 OC pair 高 bit。
|
||
* **sdot cell (12 B + 4 B pad → 16 B)** = 8 B main + 4 B aux + 4 B 0-pad
|
||
* aux OC-major:`byte k` 存 OC `k` 在 IC=0..3 的高 bit;0-pad 由 reorder 的 memset 给出。
|
||
* `shape[4]` 在 `ConvInt8TiledExecutor.cpp` 向上取整到 2,否则 `shape[4]=1` 让 packer/kernel 越界写。
|
||
|
||
### 1.5 originOffset 语义
|
||
|
||
C++ 把 `wf = alpha * (wi_unsigned + originOffset) + zero` 中的 `originOffset` 注入每行 weight 的 zero point。kernel 输出 unsigned 域 → post‑process 用一次 `MLA_WEIGHTZERO` 把 zp 项打入,等价于把 unsigned 域中心化为 signed,**kernel 不再需要 sub / bic**。这是 W3 实测从 41.3 → 50.2 GB/s 的直接原因。
|
||
|
||
---
|
||
|
||
## 2. 汇编实现要点
|
||
|
||
### 2.1 smmla (i8mm) — TILE_10 共享骨架
|
||
|
||
```
|
||
LoopSz_TILE_10:
|
||
ld1 {v0/...},[x2],#W ; weight
|
||
ld1 {v3..v6},[x11],#64 ; src tile 0..7
|
||
ld1 {v7},[x11],#16 ; src tile 8..9
|
||
<UNPACK 产出 v8..v11 = 4 个 OC pair>
|
||
smmla v12,v3,v8 ... smmla v31,v7,v11 ; 5 行 src × 4 列 OC pair = 20 条 smmla
|
||
bne LoopSz_TILE_10
|
||
```
|
||
|
||
* 累加寄存器恰好用满 `v12..v31`(20 个)。
|
||
* W4/W2/W3 的差异仅体现在 `<UNPACK>` 段,后处理(`ADD_BIAS_FLOAT/MUL_SCALE/...`)逐宏复用。
|
||
|
||
各 unpack 成本(per cell):
|
||
|
||
| 位宽 | 解包指令 | 备注 |
|
||
|------|----------|------|
|
||
| W4 | `ushr v8,v0,#4 ; and v10,v0,v2` × 2 vec → 4 ALU | 2 个 16B weight 寄存器 |
|
||
| W2 | `ushr v0,#6/#4/#2 ; and v9/v10/v11,v2(=#3)` → 6 ALU | 1 个 16B weight 寄存器 |
|
||
| W3 | W2 的 6 ALU + `W3_UNPACK_SERIAL` 中 aux 4× `ushl/add/and` ≈ 14 ALU | 主 16B + aux 8B(`ld1r`) |
|
||
|
||
### 2.2 sdot — TILE_1 主战场,多 tile fall-through
|
||
|
||
```
|
||
L8LoopSz_TILE_1_lu1:
|
||
ld1 {v_w*},[x2],#cellBytes
|
||
<UNPACK> -> v3 (OC0..3 ×4 IC), v12 (OC4..7 ×4 IC)
|
||
ld1r {v0.4s},[x1],#stride ; 1 个 input 4-byte
|
||
sdot v8.4s, v3.16b, v0.4b[0]
|
||
sdot v9.4s, v12.16b, v0.4b[0]
|
||
bne L8LoopSz_TILE_1_lu1
|
||
```
|
||
|
||
* W4 path:`UNPACK = ushr/and` 极简,TILE_12/8/4 都各自有完整 schedule。
|
||
* W2/W3 sdot:目前仅 TILE_1 有专门 schedule,TILE_12/8/4 fall-through 到 TILE_1 单 batch 循环以保正确;后续优化按 W4 模板补全。
|
||
* W2 sdot UNPACK(`UNPACK_W2_SDOT`)9 条:`tbl/ext/tbl/ushl/ushl/movi #3/and/and`。
|
||
* W3 sdot UNPACK(`UNPACK_W3_SDOT`):main 8 条(同 W2)+ aux 8 条(对称的 `tbl/ext/ushl(<<2)/and(#4)`)+ 2 条 `orr` 合并 = 18 条,每 cell 32 weight。
|
||
* `lu4` 模式预取 4 cell + 1 src,做 8 条 sdot,显著减少 IC 循环开销;tail 走 `lu1`。
|
||
|
||
### 2.3 W3 `W3_UNPACK_SERIAL` 详解
|
||
|
||
```text
|
||
; aux_dup 已是 ld1r 复制后的 16B; main 是 16B
|
||
ushl oc01, aux_dup, shifts ; shifts = {-7,-6,-5,-4,-3,-2,-1,0, ...}
|
||
add t, shifts, idx (=2)
|
||
ushl oc23, aux_dup, t
|
||
add t, t, idx
|
||
ushl oc45, aux_dup, t
|
||
add t, t, idx
|
||
ushl oc67, aux_dup, t ; 4 个 OC pair 的高 bit 落在 bit2 位置
|
||
and oc**, oc**, mask1 (=#4) ; 只留 bit2
|
||
|
||
ushr t, main, #6 ; add oc01, oc01, t
|
||
ushr t, main, #4 ; and t, t, mask3 (=#3) ; add oc23, ...
|
||
ushr t, main, #2 ; and t, t, mask3 ; add oc45, ...
|
||
and t, main, mask3 ; add oc67, ...
|
||
```
|
||
|
||
* 4 次 `ushl/add` 让 aux 不依赖 tbl,与 main 的 ushr 链可乱序发射;mask3=#3、mask1=#4、idx=#2、shifts 全部是 cell 外预加载常量,inner loop 不再额外 movi。
|
||
|
||
### 2.4 与 W4 共用的循环外结构
|
||
|
||
* `REVERT_INPUT_DEQUANT_BIAS` / `REVERT_WEIGHT_KERNEL_SUM`:block 切换时把 src/weight kernel sum 指针回滚到当前 block 的起点。
|
||
* 多 block 复用 src,通过外层 `TILE10_BLOCKNUM` 循环;block 内 `LoopSz` 是 IC 循环。
|
||
* `LoopDz8 / LoopDz4`:tail OC 不足 8 时降到 4 OC 的 schedule。
|
||
|
||
---
|
||
|
||
## 3. C++ 侧 reorder (`ConvInt8TiledExecutor.cpp`)
|
||
|
||
```
|
||
loader -> signed int8 (W2 ∈ [-2,1], W3 ∈ [-4,3], W4 ∈ [-8,7])
|
||
-> tmpWeight = signed + (-originOffset) // unsigned uint8
|
||
-> reorderWeight 走标准 int8 排布得到 (cellCount, UNIT*SRC_UNIT)
|
||
-> 按 cell layout 压成最终 weightReordered:
|
||
W4: 1 byte / 2 weight
|
||
W2: 1 byte / 4 weight (单 byte 4 路 OC)
|
||
W3: main 段 (同 W2 layout) + aux 段 (位平面)
|
||
```
|
||
|
||
要点:
|
||
* `_computeReorderQuantInfo` 接受 `weightBits`,内部计算 `originOffset`,直接把 zp 加偏移项写入 alpha/zero buffer。
|
||
* `weightBytes` 决定 KV 排布步长:W2=0.25,W3 sdot=`(SRC_UNIT*3+7)/8 / SRC_UNIT`(sdot=0.5,smmla=0.375),W4=0.5,W8=1。
|
||
* W2/W3 的 reorder 加了 fused fast path,避免 W4 之前用过的 2× oc·ic staging,RSS spike 显著下降(W2: 2.71 → 1.56 GB,W3: 2.60 → 1.40 GB,见提交 `feature/support_2bit`)。
|
||
|
||
---
|
||
|
||
## 4. 函数指针注册与选择
|
||
|
||
```
|
||
struct MatmulRelatedFunctions / CoreInt8Functions:
|
||
Int8GemmKernel // W8 FP32
|
||
Int8GemmKernel_W4 / _W3 / _W2 // FP32
|
||
MNNGemmInt8AddBiasScale_Unit_FP16 // W8 FP16
|
||
MNNGemmInt8AddBiasScale_w{4,3,2}_Unit_FP16
|
||
MNNGemmInt8AddBiasScale_*_DecodeMax // SME2 单 batch 优化(目前仅 W8/W4)
|
||
```
|
||
|
||
`DenseConvInt8TiledExecutor::onResize` 按 `mWeightBits + gcore->bytes/pack + 函数指针是否非空` 选 kernel。SME2 DecodeMax 暂未提供 W2/W3 版本,fallback 到普通 W2/W3 kernel;函数指针为空的 backend(老 ARMv8 / x86)按 `canUseInt2/3/4` 不走低 bit 路径,直接 loader 展开成 int8。
|
||
|
||
---
|
||
|
||
## 5. 性能数据(LinearRoofline 4096×14336, M-Mac 4t,GB/s)
|
||
|
||
| 位宽 | i8mm (smmla) | 备注 |
|
||
|------|--------------|------|
|
||
| W8 | 109.4 | 超过 llama.cpp Q8_0 (109) |
|
||
| W4 | 100.7 | +29% vs llama.cpp Q4_0 (78) |
|
||
| W3 | 50.2 | 与 llama.cpp Q3_K 持平,unsigned-domain + 4-IDX TILE_1 后从 41.3 提升 |
|
||
| W2 | 64.5 | +58% vs llama.cpp Q2_K (41) |
|
||
|
||
W2/W3 sdot path 的 prefill 大尺寸场景因 TILE_4/8/12 fall-through 偏慢,decode 已经达预期。
|
||
|
||
---
|
||
|
||
## 6. 已知限制 & 后续
|
||
|
||
| 项 | 状态 |
|
||
|----|------|
|
||
| W2/W3 sdot TILE_4/8/12 | 暂 fall-through 到 TILE_1,prefill 偏慢,需要按 W4 模板补 schedule。 |
|
||
| W2/W3 SME2 DecodeMax | 未提供专用 kernel,直接复用普通 W2/W3。 |
|
||
| W3 sdot 解包 ALU 链 | 含 ext 串行,可参考 i8mm 的 4-IDX 多寄存器思路降一档延迟。 |
|
||
| ARMv8.0(无 sdot/i8mm) | 不进入低 bit 路径,loader 展开成 int8 走 W8 kernel。 |
|
||
|
||
---
|
||
|
||
## 7. 关键文件索引
|
||
|
||
```
|
||
source/backend/cpu/arm/arm64/low_memory/ # FP32 路径
|
||
MNNGemmInt8AddBiasScale_ARMV82_w{2,3,4}_Unit.S # sdot
|
||
MNNGemmInt8AddBiasScale_ARMV86_w{2,3,4}_Unit.S # i8mm
|
||
MNNGemmInt8AddBiasScale_16x4_w4_Unit.S # 旧 ARMv8 路径(W4)
|
||
source/backend/cpu/arm/arm64/ # FP32 W8
|
||
MNNGemmInt8AddBiasScale_ARMV82_Unit.S
|
||
MNNGemmInt8AddBiasScale_ARMV86_Unit.S
|
||
MNNGemmInt8AddBiasScale_16x4_Unit{,_FAST}.S # 老 dotprod 路径
|
||
source/backend/arm82/asm/arm64/low_memory/ # FP16 W2/W3/W4
|
||
MNNGemmInt8AddBiasScale_ARMV82_w{2,3,4}_Unit_FP16.S
|
||
MNNGemmInt8AddBiasScale_ARMV86_w{2,3,4}_Unit_FP16.S
|
||
source/backend/arm82/asm/arm64/ # FP16 W8
|
||
MNNGemmInt8AddBiasScale_ARMV82_Unit_FP16.S
|
||
MNNGemmInt8AddBiasScale_ARMV86_Unit_FP16.S
|
||
|
||
source/backend/cpu/compute/
|
||
ConvInt8TiledExecutor.cpp # reorder / originOffset / kernel 选择
|
||
CommonOptFunction.h # MatmulRelatedFunctions
|
||
Int8FunctionsOpt.{h,cpp} # CoreInt8Functions
|
||
test/speed/GemvBWTest.cpp # LinearRoofline benchmark (W2/W3/W4/W8)
|
||
skills/arm-cpu-optimize/ # 优化方法论与 step 文档
|
||
```
|