GGML的CPU算子解读:矩阵乘法

🕒 2026-01-20 📁 推理框架 👤 laumy 🔥 97 热度

算子实现

调用流程

主要是ggml_compute_forward_mul_mat函数,该函数把任务拆分,最终计算执行调用ggml_compute_forward_mul_mat_one_chunk实现。

任务拆分

void ggml_compute_forward_mul_mat(
        const struct ggml_compute_params * params,
              struct ggml_tensor * dst) {

    // 1. 获取输入和输出张量的指针
    const struct ggml_tensor * src0 = dst->src[0];
    const struct ggml_tensor * src1 = dst->src[1];

    // 2. 初始化一堆本地变量 (ne0, nb0 等),方便后面写公式
    GGML_TENSOR_BINARY_OP_LOCALS

    // 3. 获取当前线程的信息
    const int ith = params->ith; // 我是第几个线程?
    const int nth = params->nth; // 总共有几个线程?

    // 4. 获取src0(权重)的数据类型,用于后续查询src1(输入/激活值)必须转换为什么类型才能进行计算。
    //在大语言模型中,为了省内存,src0(权重矩阵)通常是经过量化的比如Q4_0,Q5_1,Q8_0等格式。
    enum ggml_type const vec_dot_type = type_traits_cpu[src0->type].vec_dot_type;

    // ... (省略一堆 GGML_ASSERT 检查,确保维度合法,比如 src0 和 src1 能不能相乘) ...

#if GGML_USE_LLAMAFILE
    // 5. 【极速通道】如果系统支持 llamafile 的手写汇编优化,直接丢给它算,算完直接返回。
    // 这比下面的 C 代码快得多。
    if (src1_cont) {
        // ... 尝试调用 llamafile_sgemm ...
        if (!llamafile_sgemm(...)) goto UseGgmlGemm1; // 如果失败,就跳回下面的普通通道
        return; 
    }
UseGgmlGemm1:;
#endif

    // 6. 如果src1的类型与src0的类型不匹配,需要将src1转换成vec_dot_type类型。
    if (src1->type != vec_dot_type) {
        char * wdata = params->wdata;
        // ... (循环遍历 src1,把数据解压并复制到 wdata) ...
        // 这里用了 from_float 函数进行转换
    }

    // 7. 【同步】确保所有线程都知道数据准备好了
    //只让线程0做一次初始化,把 current_chunk 设为 nth:因为每个线程先从 current_chunk = ith 
    //开始领一个“私有起始 chunk”,所以下一个“公共可抢”的 chunk 从 nth 开始。
    if (ith == 0) {
        atomic_store_explicit(¶ms->threadpool->current_chunk, nth, memory_order_relaxed);
    }
    ggml_barrier(params->threadpool); 
    // barrier 等待:确保所有线程都完成了“src1 预处理/计数器初始化”,再进入真正的 chunk 计算阶段。
    // ... (再次尝试 llamafile 优化,省略) ...

    // 8. 【切蛋糕】决定每个任务块 (chunk) 有多大
    // nr0 是结果矩阵的行数,nr1 是其余维度的大小
    const int64_t nr0 = ne0; //nr0 代表结果矩阵在“第一维”的长度(可以理解为输出的行数/一个主维度)。
    const int64_t nr1 = ne1 * ne2 * ne3;
    //nr1 把其余维度展平成第二维(把 ne1*ne2*ne3 看成“列块/批次展平后的列方向”)。

    int chunk_size = 16; // 默认每块按 16 规模做分块(和 one_chunk 内的 16×16 tiling 呼应)。
    if (nr0 == 1 || nr1 == 1) { chunk_size = 64; } // 如果矩阵很细长,就切大一点

    // 算出总共有多少块蛋糕 (nchunk0 * nchunk1)
    int64_t nchunk0 = (nr0 + chunk_size - 1) / chunk_size;
    int64_t nchunk1 = (nr1 + chunk_size - 1) / chunk_size;

    // 9. 【NUMA 优化】如果是多路 CPU 服务器,或者切得太碎了,就改用保守策略
    if (nchunk0 * nchunk1 < nth * 4 || ggml_is_numa()) {
        nchunk0 = nr0 > nr1 ? nth : 1; 
        nchunk1 = nr0 > nr1 ? 1 : nth; 
    }

    // 计算每个块实际包含多少行/列 (dr0, dr1)
    //把 nr0/nr1 平均分给 nchunk0/nchunk1,得到每块在两维上的跨度 dr0/dr1。
    const int64_t dr0 = (nr0 + nchunk0 - 1) / nchunk0;
    const int64_t dr1 = (nr1 + nchunk1 - 1) / nchunk1;

    // 10. 【开始干活】线程开始领任务
    // 一开始,第 ith 个线程领第 ith 块蛋糕
    //每个线程先从 current_chunk = ith 开始,保证一开始每个线程都有活(减少原子争用)。
    int current_chunk = ith;

    while (current_chunk < nchunk0 * nchunk1) {
        // 算出当前蛋糕块在矩阵的哪个位置 (start 到 end)

        const int64_t ith0 = current_chunk % nchunk0;
        const int64_t ith1 = current_chunk / nchunk0;

        const int64_t ir0_start = dr0 * ith0;
        const int64_t ir0_end = MIN(ir0_start + dr0, nr0);
        // ... (同理算出 ir1 的范围) ...

        // 11. 【核心调用】把这一小块任务交给工人去算
        // num_rows_per_vec_dot 是看 CPU 是否支持一次算两行 (MMLA 指令)
        ggml_compute_forward_mul_mat_one_chunk(..., ir0_start, ir0_end, ir1_start, ir1_end);

        // 如果线程数比蛋糕块还多,干完这一块就可以下班了
        if (nth >= nchunk0 * nchunk1) { break; }

        // 12. 【抢任务】否则,去原子计数器里“取号”,领下一块没人干的蛋糕
        current_chunk = atomic_fetch_add_explicit(¶ms->threadpool->current_chunk, 1, memory_order_relaxed);
    }
}

计算执行

static void ggml_compute_forward_mul_mat_one_chunk(
    // ... 参数省略,主要是范围 ir0_start 等 ...
    ) {

    // 1. 获取数据指针
    const struct ggml_tensor * src0 = dst->src[0];
    const struct ggml_tensor * src1 = dst->src[1];

    // ... 初始化本地变量 ...

    // 2. 准备点积函数 (vec_dot)
    // type_traits_cpu 是一张表,根据数据类型查到最优的计算函数
    ggml_vec_dot_t const vec_dot = type_traits_cpu[type].vec_dot;

    // 3. 【广播因子】计算 src1 比 src0 大多少倍
    // 比如 src1 有 2 个头,src0 只有 1 个头,那 r2 就是 2。计算时 src0 要重复用。
    const int64_t r2 = ne12 / ne02;
    const int64_t r3 = ne13 / ne03;

    // 如果任务范围是空的,直接返回
    if (ir0_start >= ir0_end || ir1_start >= ir1_end) { return; }

    // 4. 【Tiling 分块参数】
    // 即使是这一小块任务,也要再切碎,为了放进 CPU 缓存 (L1 Cache)
    const int64_t blck_0 = 16;
    const int64_t blck_1 = 16;

    // 5. 定义临时缓冲区,防止多个线程写内存冲突 (False Sharing)
    float tmp[32]; 

    // 6. 【双重 Tiling 循环】开始遍历任务块
    // iir1 每次跳 16 行
    for (int64_t iir1 = ir1_start; iir1 < ir1_end; iir1 += blck_1) {
        // iir0 每次跳 16 列
        for (int64_t iir0 = ir0_start; iir0 < ir0_end; iir0 += blck_0) {

            // 7. 【处理块内的每一行】
            for (int64_t ir1 = iir1; ir1 < iir1 + blck_1 && ir1 < ir1_end; ir1 += num_rows_per_vec_dot) {

                // 8. 【计算索引】这部分数学很绕,简单说就是:
                // 根据当前的行号 ir1,反推出它在 src1 的第几层、第几页 (i13, i12, i11)
                const int64_t i13 = (ir1 / (ne12 * ne1));
                const int64_t i12 = (ir1 - i13 * ne12 * ne1) / ne1;
                const int64_t i11 = (ir1 - i13 * ne12 * ne1 - i12 * ne1);

                // 9. 【应用广播】算出对应的 src0 的位置
                // 如果 src0 比较小,就除以倍率 r2, r3 来复用数据
                const int64_t i03 = i13 / r3;
                const int64_t i02 = i12 / r2;

                // 拿到 src0 这一行的内存地址
                const char * src0_row = (const char*)src0->data + (0 + i02 * nb02 + i03 * nb03);

                // 拿到 src1 这一列的内存地址 (wdata 是之前解压好的数据)
                const char * src1_col = (const char*)wdata + ...; // (省略复杂的偏移量计算)

                // 拿到 结果 dst 的写入地址
                float * dst_col = (float*)((char*)dst->data + ...);

                // 10. 【核心计算循环】
                // 遍历 src0 的行 (分块内),调用 SIMD 指令进行计算
                for (int64_t ir0 = iir0; ir0 < iir0 + blck_0 && ir0 < ir0_end; ir0 += num_rows_per_vec_dot) {
                    // vec_dot:一次性计算 ne00 个元素的点积
                    // 结果先存在 tmp 数组里,而不是直接写回 dst
                    vec_dot(ne00, &tmp[ir0 - iir0], ..., src0_row + ir0 * nb01, ..., src1_col, ...);
                }

                // 11. 【写回结果】
                // 把算好的 tmp 里的数据,一次性拷贝 (memcpy) 到最终的目标 dst_col
                for (int cn = 0; cn < num_rows_per_vec_dot; ++cn) {
                    memcpy(&dst_col[iir0 + ...], tmp + ..., ...);
                }
            }
        }
    }
}

以一个例子来总结一下。

void ggml_compute_forward_mul_mat(
    const struct ggml_compute_params * params,  // 计算参数(包含线程信息)
    struct ggml_tensor * dst                    // 结果矩阵 C = A × B
)

功能是矩阵乘法:结果矩阵C=矩阵A x 矩阵B。其并行的策略如下:

  • 切分对象: 结果矩阵C(不是A和B)
  • 并行方式:2D切分+work_stealing(工作窃取)
  • 写入方式:每个线程直接写入结果矩阵的不同位置。
矩阵A (m×k) × 矩阵B (k×n) = 结果矩阵C (m×n)
     ↑              ↑                    ↑
   只读           只读              切分这个!

如上,矩阵A和B是只读的,所有线程都需要读取完整的A和B,结果矩阵C需要写入,每个线程写入不同的位置,切分C可以让多个线程并行计算不同的部分,互不干扰。

示例演示

输入数据

矩阵A (4行×3列):
┌─────────┐
│ 1  2  3 │
│ 4  5  6 │
│ 7  8  9 │
│ 10 11 12│
└─────────┘

矩阵B (3行×4列):
┌─────────────┐
│ 1  2  3  4  │
│ 5  6  7  8  │
│ 9  10 11 12 │
└─────────────┘

要计算:结果矩阵C = A × B (4行×4列)

分配结果矩阵

结果矩阵C(内存已分配,4×4 = 16个float):
┌─────────────────────┐
│ ?  ?  ?  ?  │ ← 行0
│ ?  ?  ?  ?  │ ← 行1
│ ?  ?  ?  ?  │ ← 行2
│ ?  ?  ?  ?  │ ← 行3
└─────────────────────┘
   列0 列1 列2 列3

计算切分参数

假设:2个线程,chunk_size = 2

nr0 = 4, nr1 = 4
nchunk0 = (4 + 2 - 1) / 2 = 3
nchunk1 = (4 + 2 - 1) / 2 = 3
总块数 = 3 × 3 = 9

检查:9 < 2 * 4 = 8? 不,9 >= 8
-> 不触发优化,保持3×3切分

dr0 = (4 + 3 - 1) / 3 = 2  // 每个块2行
dr1 = (4 + 3 - 1) / 3 = 2  // 每个块2列

切分结果矩阵

结果矩阵C (4×4) 被切分成9个块:
┌──────┬──────┬──────┐
│ 块0  │ 块1  │ 块2  │ ← 行块0(行0-1)
│[0,0] │[0,1] │[0,2] │
│[1,0] │[1,1] │[1,2] │
├──────┼──────┼──────┤
│ 块3  │ 块4  │ 块5  │ ← 行块1(行2-3)
│[2,0] │[2,1] │[2,2] │
│[3,0] │[3,1] │[3,2] │
├──────┼──────┼──────┤
│ 块6  │ 块7  │ 块8  │ ← 行块2(超出边界)
│[0,3] │[1,3] │[2,3] │
│[1,3] │[2,3] │[3,3] │
└──────┴──────┴──────┘

映射表

列0  列1  列2  列3
┌────┬────┬────┬────┐
│块0 │块0 │块3 │块3 │ ← 行0
│    │    │    │    │
├────┼────┼────┼────┤
│块0 │块0 │块3 │块3 │ ← 行1
│    │    │    │    │
├────┼────┼────┼────┤
│块1 │块1 │块4 │块4 │ ← 行2
│    │    │    │    │
├────┼────┼────┼────┤
│块1 │块1 │块4 │块4 │ ← 行3
└────┴────┴────┴────┘

块0: ith0=0, ith1=0 → C[0:2, 0:2] → 左上角 2×2
块1: ith0=1, ith1=0 → C[2:4, 0:2] → 左下角 2×2
块3: ith0=0, ith1=1 → C[0:2, 2:4] → 右上角 2×2
块4: ith0=1, ith1=1 → C[2:4, 2:4] → 右下角 2×2

转换的过程

// 步骤1:块编号 → 块坐标
current_chunk = 3
ith0 = 3 % 3 = 0  // 块在第0行(行方向的第0个块)
ith1 = 3 / 3 = 1  // 块在第1列(列方向的第1个块)

// 步骤2:块坐标 → 实际位置
ir0_start = dr0 * ith0 = 2 * 0 = 0      // 从第0行开始
ir0_end = MIN(0 + 2, 4) = 2            // 到第2行结束(不包括2,实际是0-1行)
ir1_start = dr1 * ith1 = 2 * 1 = 2     // 从第2列开始
ir1_end = MIN(2 + 2, 4) = 4            // 到第4列结束(不包括4,实际是2-3列)

// 步骤3:实际位置 → 结果矩阵元素
块3 = C[0:2, 2:4] = {
    C[0][2], C[0][3],
    C[1][2], C[1][3]
}

之所以要做一个映射,是因为work-stealing需要一维的变化,原子操作只能操作一个整数,线程并发是以块为最小运算单位。也是为了方便后面做负载均衡。

并行运算

(1)线程0的工作

current_chunk = 0  // 从块0开始

// 处理块0
ith0 = 0 % 3 = 0, ith1 = 0 / 3 = 0
ir0_start = 0, ir0_end = 2  // 行0-1
ir1_start = 0, ir1_end = 2  // 列0-1

计算:
  C[0][0] = A[0行] · B[第0列] = [1,2,3]·[1,5,9] = 38
  C[0][1] = A[0行] · B[第1列] = [1,2,3]·[2,6,10] = 44
  C[1][0] = A[1行] · B[第0列] = [4,5,6]·[1,5,9] = 83
  C[1][1] = A[1行] · B[第1列] = [4,5,6]·[2,6,10] = 98

直接写入结果矩阵:
  dst->data[0*4 + 0] = 38  // C[0][0]
  dst->data[0*4 + 1] = 44  // C[0][1]
  dst->data[1*4 + 0] = 83  // C[1][0]
  dst->data[1*4 + 1] = 98  // C[1][1]

// 完成后,抢下一个块
current_chunk = atomic_fetch_add(&current_chunk, 1) = 2

// 处理块2(无效,会被跳过)
ir0_start = 4, ir0_end = 4  // 空
→ 跳过

// 继续抢下一个块
current_chunk = atomic_fetch_add(&current_chunk, 1) = 3

// 处理块3

(2)线程1的工作

current_chunk = 1  // 从块1开始

// 处理块1
ith0 = 1 % 3 = 1, ith1 = 1 / 3 = 0
ir0_start = 2, ir0_end = 4  // 行2-3
ir1_start = 0, ir1_end = 2  // 列0-1

计算并写入 C[2:4, 0:2]...

// 完成后,抢下一个块
current_chunk = atomic_fetch_add(&current_chunk, 1) = 4

// 处理块4

(3)最终的结果

所有线程完成后,结果矩阵C:
┌─────────────────────┐
│ 38  44  50  56  │ ← 所有元素都已计算并写入
│ 83  98  113 128 │
│ 128 152 176 200 │
│ 173 206 239 272 │
└─────────────────────┘

发表你的看法

\t