Nvidia Matmul Tiling：复杂度、缓存局部性与分块实现

题意与假设

这道练习题保留了面经中值得训练的核心约束，并将需要澄清的条件单独列出。现场拿到变体时，先和面试官确认这些条件。

• 训练代码以 row-major 连续 int64 数组表示 A 为 M×K、B 为 K×N，结果新分配为 M×N；整数乘加不溢出是输入前提。题目未明确数值类型与 API。
• 块边长 B 默认 32，仅作教学参数；最优块大小取决于缓存容量、向量化、线程数或 GPU shared memory，需要在目标硬件基准测试。
• 零维矩阵允许空输入并返回大小为 M×N 的零矩阵；形状乘积溢出、数组长度不符或 B=0 抛参数错误。

解题路线

1. 先写定义：C[i,j]=Σₜ A[i,t]×B[t,j]，i∈[0,M)、t∈[0,K)、j∈[0,N)。朴素三层循环执行 MKN 次乘加，时间 O(MKN)，输出空间 O(MN)。
2. row-major 下 A[i,t] 偏移 iK+t，B[t,j] 偏移 tN+j，C[i,j] 偏移 iN+j。按 i、t、j 顺序，内层 j 对 B 和 C 连续访问，并把 A[i,t] 保存在寄存器变量中。
3. 分块把 i、t、j 空间切成 B 大小的 tile，在块内完成乘加，让一小段 A/B/C 在较短时间内重复使用。每个块结束下标取 min(维度,起点+B)，边缘不满块也正确。
4. 正确性来自乘法定义：每个 (i,j,t) 三元组恰好落入一个块并贡献一次乘积；循环只改变累加顺序，不改变覆盖集合。浮点实现可能因加法顺序不同出现舍入差异，此整数训练版没有该问题。

手动走一遍

A=[[1,2,3],[4,5,6]]、B=[[7,8],[9,10],[11,12]]，输出 [[58,64],[139,154]]。B=2 时 K=3 的最后一个 t 块只含 t=2；C[0,0] 先累加 1×7+2×9，再加 3×11=58。

C++ 参考实现

下面的代码按上文假设实现，适合本地运行并追加自己的测试用例。

代码区可横向滚动；使用键盘时，先聚焦代码区再按方向键查看。

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <stdexcept>
#include <vector>

// Row-major A: m x k; B: k x n; result: m x n.
std::vector<std::int64_t> blocked_matmul(
    const std::vector<std::int64_t>& a,
    const std::vector<std::int64_t>& b,
    std::size_t m,
    std::size_t k,
    std::size_t n,
    std::size_t block_size = 32) {
    const std::size_t max_size = static_cast<std::size_t>(-1);
    if (block_size == 0 || (k != 0 && m > max_size / k) ||
        (n != 0 && k > max_size / n) || (n != 0 && m > max_size / n) ||
        a.size() != m * k || b.size() != k * n) {
        throw std::invalid_argument("invalid shape or block size");
    }
    std::vector<std::int64_t> result(m * n, 0);
    for (std::size_t i0 = 0; i0 < m;) {
        const std::size_t i_end = i0 + std::min(block_size, m - i0);
        for (std::size_t k0 = 0; k0 < k;) {
            const std::size_t k_end = k0 + std::min(block_size, k - k0);
            for (std::size_t j0 = 0; j0 < n;) {
                const std::size_t j_end = j0 + std::min(block_size, n - j0);
                for (std::size_t i = i0; i < i_end; ++i) {
                    for (std::size_t t = k0; t < k_end; ++t) {
                        const std::int64_t value = a[i * k + t];
                        for (std::size_t j = j0; j < j_end; ++j) {
                            result[i * n + j] += value * b[t * n + j];
                        }
                    }
                }
                j0 = j_end;
            }
            k0 = k_end;
        }
        i0 = i_end;
    }
    return result;
}

复杂度

运算量仍为 O(MKN)，输出 O(MN)；分块增加索引控制的常数开销，目标是降低缓存/显存往返而非改变渐近乘加量。实际 CPU 优化常要考虑向量化、线程划分和 cache blocking；GPU 还需分析线程映射、合并访存、shared memory、寄存器与同步。

易错点与追问

容易失分的边界

• 只回答方阵 O(N³)，没有给 M×K 与 K×N 的一般复杂度。
• 把 B[t,j] 按 t 连续读取；row-major 时其下标相隔 N，循环顺序会影响访存。
• 只测维度刚好是块大小倍数，漏掉 K、M、N 的边缘块。
• 认为分块必然提速；小矩阵或不合适块大小可能比简单循环更慢。

继续深挖

• 若换成 float/BF16，解释累加精度、溢出/下溢、舍入误差与容差测试，不要求逐位相同。
• 若转到 GPU，先量测访存还是计算瓶颈，再选择 tile 形状和线程布局；shared memory 装载与同步需要避免越界和 bank conflict。
• 若矩阵很大或稀疏，比较现成 BLAS/cuBLAS、分布式切分和稀疏格式，而不是默认手写 kernel。