CHAPTER 36 / 架构、性能与并行

Tensor shape、stride 与 GEMM 分块

转置一定要复制数据吗,矩阵乘法如何通过tile复用减少搬运?

阅读与推演约 60 分钟练习时间另计

这一章要弄清楚

  • 用shape和stride定位tensor元素
  • 区分转置视图与实际数据重排
  • 验证非整除维度的分块GEMM

先备知识:CPU cache、缓存一致性与 SIMD / Reduction、scan 与数值正确性 / Kernel 索引、边界与 grid-stride loop / 合并访问、共享内存与同步范围

C++20 / macOS 与 Linux;本章的硬件模型只推演逻辑,不代表设备性能。

shape说明有多大,stride说明怎样走

一个tensor不只是连续数组加一个名字。Shape描述每个维度的长度;stride描述沿某个维度前进一步需要跨过多少存储元素,必须明确单位是元素还是字节。对行主序2×3矩阵,数据为1、2、3、4、5、6,shape是(2,3),元素stride是(3,1)。坐标(1,2)的偏移为1×3+2×1=5,读取6。

转置视图把shape换成(3,2),stride换成(1,3),同一份存储就呈现1、4;2、5;3、6。没有复制六个元素,但后续行遍历不再连续。视图避免搬运并不代表所有后续计算都更快;某些kernel要求连续布局,届时要支持stride或显式重排。持有视图还意味着必须保证底层存储的生命周期,不能让一个短暂vector返回悬空的tensor视图。

坐标合法与地址合法都要验证

先检查每个坐标落在shape范围内,再计算偏移,并检查乘加是否溢出及偏移是否落在存储范围。不能只用最后的线性偏移判断坐标是否合法。真实tensor还可能有负stride、广播的零stride和子视图offset,本章只使用非负stride和零基地址,避免把接口范围扩展得含糊。第一例验证转置视图、越界坐标及不合理stride。

GEMM把依赖写成三个维度

矩阵乘法 C=A×B 要求A的shape是(M,K),B是(K,N),C是(M,N)。输出 C[r,c] 是沿k累加 A[r,k]*B[k,c]。M决定输出行,N决定输出列,K是归约维度。示例的A为2×3,B为3×2,首个输出为1×7+2×9+3×11=58。把B当作N×K或忘记leading dimension,可能访问合法内存却算错语义。

tile通过复用改变流量

朴素实现为每个输出分别读取整行A和整列B,重复搬运相同数据。分块选择一小块C,逐段遍历K;每段先加载A小块和B小块,再让它们参与多次乘加。若tile边长为T,一轮理想加载约2T²个元素,执行约2T³个浮点运算,复用随T增加。但更大的tile占用更多寄存器或共享区,不能无限放大。

第二例用固定2×2 tile在CPU上明确加载两个暂存数组、执行局部乘加并写回,输入维度非整除时用零填充边缘。对0..5的M、N、K组合,共216种shape,与朴素整数参考逐项比较。整数值受控,避免溢出;此例验证分块算法和尾部,不是GPU性能结果。真实GPU实现仍需要前章讲的合作加载、共享区同步以及设备验证。

数值和布局进入同一个合同

实际GEMM常含 C=alpha*A*B+beta*C,还要约定输入与累加精度、布局、batch维和是否允许原地别名。本章只实现无别名的纯乘积,不能直接把例子当作完整BLAS接口。浮点版本要根据K长度、数值尺度和应用需求设定容差;NaN检查及相消案例沿用上一章。

优化实验应固定shape、dtype、布局和计时范围,分别记录数据重排、传输与核心计算。大型方阵的高吞吐不能代表细长矩阵或很小batch。面试中拿一个2×3乘3×2的具体例子,指着索引说明哪里复用、何时同步、尾部补什么,再讨论tensor core或厂商专用tile指令,会更容易证明你理解了机制。

动手改变 · 观察因果

在矩阵格子里追踪每一段 K

下一段 K 给所选 C 元素加多少?哪些输入被复用?

A 为 2×3,B 为 3×2,tile 为 2×2,使用小整数。展示合作加载的算法关系,不代表 GPU 指令、同步或性能。 对应 例题 32-b;图中的代码行是步骤提示,完整可编译源码见例题。

改输入后从第一步重新推演。Tab 选择控件,Enter/空格操作按钮;图内方向键平移,手机可横向滑动。

正在准备默认算例。下方例题包含完整源码与逐步解释。

第 1 步

先预测,再前进一步

下一段 K 给所选 C 元素加多少?哪些输入被复用?

输入、边界、状态变化

完整文字推演与当前数据
    静态推演与完整文字(便于对照、打印)
    观察 · 推演

    K 分段贡献的文字摘要(二维复用见上方实验)

    A tile1 2B tile7 9累加1×7+2×9=2501 / 03 · TILESA tile1 2B tile7 9累加1×7+2×9=2501 / 03 · TILES
    K=0..1

    先读A的1、2和B的7、9。

    1 / 3
    阅读完整推演文字
    1. K=0..1

      A tile:1 2;B tile:7 9;累加:1×7+2×9=25

      先读A的1、2和B的7、9。

    2. K=2与填充

      A tile:3 0;B tile:11 0;本轮:33

      第二段有效值为3、11,越界位置补0。

    3. 写回输出

      C[0,0]:58;C row0:58 64;C row1:139 154

      25+33=58;其他输出按相同复用计算。

    跟着例子,走完一遍

    例题 01C++20 · 本机可运行

    零复制转置视图

    2×3矩阵底层[1,2,3,4,5,6],原stride(3,1),转置stride(1,3)。

    1. 坐标先按shape检查。
    2. 用r*row_stride+c*col_stride计算偏移并检查溢出。
    3. 逐行读取转置视图,不改变底层顺序。
    32-a.cpp
    下载
    #include <cstddef>
    #include <iostream>
    #include <limits>
    #include <stdexcept>
    #include <vector>
    int at(const std::vector<int>& data,std::size_t rows,std::size_t cols,std::size_t rs,std::size_t cs,std::size_t r,std::size_t c) {
        if(r>=rows || c>=cols) throw std::out_of_range("coordinate");
        const auto maximum=std::numeric_limits<std::size_t>::max();
        if((r!=0 && rs>maximum/r)||(c!=0 && cs>maximum/c)) throw std::out_of_range("stride overflow");
        const auto a=r*rs,b=c*cs;
        if(b>maximum-a || a+b>=data.size()) throw std::out_of_range("storage");
        return data[a+b];
    }
    int main() {
        const std::vector<int> data{1,2,3,4,5,6};
        if(at(data,2,3,3,1,1,2)!=6) throw std::runtime_error("row major");
        int rejected=0;
        try {
            const int unexpected=at(data,3,2,1,3,0,2);
            std::cout<<"unexpected invalid-coordinate value="<<unexpected<<'\n';
            return 1;
        } catch(const std::out_of_range&){++rejected;}
        try {
            const int unexpected=at(data,3,2,std::numeric_limits<std::size_t>::max(),3,2,1);
            std::cout<<"unexpected invalid-stride value="<<unexpected<<'\n';
            return 1;
        } catch(const std::out_of_range&){++rejected;}
        if(rejected!=2) throw std::runtime_error("bounds");
        std::cout<<"transpose=";
        for(std::size_t r=0;r<3;++r){if(r)std::cout<<';';for(std::size_t c=0;c<2;++c)std::cout<<(c?",":"")<<at(data,3,2,1,3,r,c);}
        std::cout<<'\n';
    }

    如何编译和运行下载的 .cpp 文件 →

    结果与解释

    transpose=1,4;2,5;3,6

    视图改变寻址规则,不复制或重排底层vector。

    在本机运行这个例子

    下载后,在文件所在目录执行。需要支持 C++20 的编译器;POSIX 示例还需要章节说明中的系统条件。

    clang++ -std=c++20 -Wall -Wextra -Wpedantic -Werror -pthread 32-a.cpp -o example && ./example

    预期标准输出:

    transpose=1,4;2,5;3,6
    
    例题 02C++20 · 本机可运行

    2×2 tile处理K=3尾部

    A=[1,2,3;4,5,6],B=[7,8;9,10;11,12];并测试216组小shape。

    1. 每块C用局部累加数组。
    2. K每次前进2,越界A/B位置补0。
    3. 只写回有效C坐标,与朴素参考比较。
    32-b.cpp
    下载
    #include <array>
    #include <cstddef>
    #include <iostream>
    #include <stdexcept>
    #include <vector>
    std::vector<long long> gemm(const std::vector<int>& a,const std::vector<int>& b,std::size_t m,std::size_t n,std::size_t k) {
        if(m>100 || n>100 || k>100 || a.size()!=m*k || b.size()!=k*n) throw std::invalid_argument("bounded shape");
        constexpr std::size_t tile=2;std::vector<long long> c(m*n,0);
        for(std::size_t row=0;row<m;row+=tile) for(std::size_t col=0;col<n;col+=tile) {
            std::array<long long,4> accum{};
            for(std::size_t base=0;base<k;base+=tile) {
                std::array<int,4> aa{},bb{};
                for(std::size_t r=0;r<tile;++r) for(std::size_t q=0;q<tile;++q) {
                    if(row+r<m && base+q<k) aa[r*tile+q]=a[(row+r)*k+base+q];
                    if(base+r<k && col+q<n) bb[r*tile+q]=b[(base+r)*n+col+q];
                }
                for(std::size_t r=0;r<tile;++r) for(std::size_t q=0;q<tile;++q)
                    for(std::size_t x=0;x<tile;++x) accum[r*tile+q]+=static_cast<long long>(aa[r*tile+x])*bb[x*tile+q];
            }
            for(std::size_t r=0;r<tile;++r) for(std::size_t q=0;q<tile;++q)
                if(row+r<m && col+q<n) c[(row+r)*n+col+q]=accum[r*tile+q];
        }
        return c;
    }
    int main() {
        int shapes=0;
        for(std::size_t m=0;m<=5;++m) for(std::size_t n=0;n<=5;++n) for(std::size_t k=0;k<=5;++k) {
            std::vector<int>a(m*k),b(k*n);for(std::size_t i=0;i<a.size();++i)a[i]=static_cast<int>(i%7)-3;
            for(std::size_t i=0;i<b.size();++i)b[i]=static_cast<int>(i%5)-2;
            auto c=gemm(a,b,m,n,k);
            for(std::size_t r=0;r<m;++r) for(std::size_t col=0;col<n;++col) {
                long long expected=0;for(std::size_t x=0;x<k;++x)expected+=static_cast<long long>(a[r*k+x])*b[x*n+col];
                if(c[r*n+col]!=expected) throw std::runtime_error("oracle");
            }
            ++shapes;
        }
        const auto c=gemm({1,2,3,4,5,6},{7,8,9,10,11,12},2,2,3);
        if(c!=std::vector<long long>{58,64,139,154}) throw std::runtime_error("example");
        std::cout<<"C="<<c[0]<<','<<c[1]<<';'<<c[2]<<','<<c[3]<<" shapes="<<shapes<<'\n';
    }
    结果与解释

    C=58,64;139,154 shapes=216

    CPU分块代码明确复用已加载tile,但不代表GPU执行或加速。

    在本机运行这个例子

    下载后,在文件所在目录执行。需要支持 C++20 的编译器;POSIX 示例还需要章节说明中的系统条件。

    clang++ -std=c++20 -Wall -Wextra -Wpedantic -Werror -pthread 32-b.cpp -o example && ./example

    预期标准输出:

    C=58,64;139,154 shapes=216
    
    常见错误

    只改shape就宣称完成转置

    shape从2×3改3×2却仍按连续row-major stride读取,得到[1,2;3,4;5,6]而非转置。

    修正思路:转置视图同时置换shape与stride;若需要连续转置布局,则实际重排数据并明确成本。

    轮到你动手

    先写预测或代码,再按需打开提示。完整答案用于对照自己的推理。

    练习 1

    原矩阵shape(2,3)、stride(3,1),转置视图坐标(2,1)对应原来哪个元素?

    给我一点提示
    1. 转置stride为(1,3)。
    2. 先算元素偏移。
    查看答案与推理

    偏移2×1+1×3=5,值为6,对应原矩阵(1,2)。底层存储完全未变,转置第一维的步长变成1。

    练习 2

    2×2 tile乘法中M=3、N=1、K=3,哪些维度需要边界处理?

    给我一点提示
    1. 输出行和列都有部分tile。
    2. 最后K段也不足2。
    查看答案与推理

    最后输出行块只有1行,唯一列块只有1列,最后K块只有1个有效归约元素。加载A/B分别按实际维度补0,最终写C只允许r<3且c<1;不要只处理K尾部而遗漏输出边界。

    把理解说出来

    先用中文讲清因果,再用英文回答。问题依据技能主题编写,并非公司内部题库。

    解释

    shape和stride各表示什么?

    参考回答 / English answer

    shape是各维长度;stride是沿该维前进一步跨过的存储单位数,必须明确元素或字节单位。

    Shape gives dimension lengths. Strides tell how far storage advances along each dimension, in explicitly stated units.
    预测

    转置一定会复制数据吗?

    参考回答 / English answer

    不一定,交换shape和stride可建视图;但后续连续性要求可能需要实际重排。

    A transpose can be a metadata-only view. A later operation may still require materializing a contiguous layout.
    找错

    GEMM的K维为什么不能忽略?

    参考回答 / English answer

    它同时是A列数与B行数,是每个输出的归约长度;不匹配时乘法语义不成立。

    K is the reduction dimension shared by A’s columns and B’s rows. A mismatch makes the matrix product invalid.
    解释

    tile为何能减少内存流量?

    参考回答 / English answer

    加载的一小块A/B用于多个输出乘加,增加复用;要把加载和同步成本也算进去。

    Tiling reuses loaded A and B values across multiple outputs. The benefit must exceed loading and synchronization overhead.
    追问

    tile越大是否必然更快?

    参考回答 / English answer

    不是,更多寄存器/共享区可能降低并行驻留,边缘浪费也增加;选择依赖shape、dtype和设备。

    Larger tiles improve potential reuse but consume more resources. The best choice depends on shape, data type, and hardware.
    找错

    分块CPU程序通过216个shape是否可以宣称GPU GEMM已验证?

    参考回答 / English answer

    不能,这只验证CPU算法与尾部逻辑。设备版本仍要编译、执行、同步和数值验证以及独立性能测量。

    Those tests validate the CPU algorithm and boundary handling. A GPU implementation still needs its own execution and correctness evidence.

    继续查证

    公开资料用于查证;本章图解和例题是独立教学内容。CPU 逻辑模型不能证明设备性能。

    接着看已有的图解

    这些资料按主题补充本章内容。