缓存敏感矩阵乘法的多级平铺

如何解决缓存敏感矩阵乘法的多级平铺

为了好玩，我正在编写自己的 GeMM 子程序。我已经设法使用 AVX256 内核在 L1 缓存上实现了平铺版本。我知道有一些循环不变量我可以从下面提升。

template<size_t N,size_t P,size_t M>
void intrinsic_multiply(double* A,double* B,void* C) {
    alignas(32) double _A[PACK_SIZE * PACK_SIZE];
    alignas(32) double _B[PACK_SIZE * PACK_SIZE];
    double* _C = (double*)(C);

    constexpr size_t N_rem = N%PACK_SIZE;
    constexpr size_t M_rem = M%PACK_SIZE;
    constexpr size_t P_rem = P%PACK_SIZE;

    constexpr size_t N_spill = N-N_rem+(PACK_SIZE)*(size_t)(N_rem!=0);
    constexpr size_t M_spill = M-M_rem+(PACK_SIZE)*(size_t)(M_rem!=0);
    constexpr size_t P_spill = P-P_rem+(PACK_SIZE)*(size_t)(P_rem!=0);

    for(size_t i = 0; i < N_spill; i += PACK_SIZE) {
        for(size_t j = 0; j < M_spill; j += PACK_SIZE) {
            for(size_t k = 0; k < P_spill; k += PACK_SIZE) {
                pad_cols<N>(A + i + k * N,_A,((size_t)(i+PACK_SIZE>N))*N_rem,((size_t)(k+PACK_SIZE>P))*P_rem);
                pad_cols<P>(B + k + j * P,_B,((size_t)(k+PACK_SIZE>P))*P_rem,((size_t)(j+PACK_SIZE>M))*M_rem);
                macro_kernal_intrinsic<N_spill>(_A,_C + i + (j * N_spill));
            }
        }
    }
}

我很难实现多级缓存平铺，因为缓存不是彼此的倍数。每个缓存的步幅大小的估计计算如下，CACHE_SIZE 以字节为单位给出。

static constexpr size_t L3_CACHESIZE = 6291456;
static size_t L3_STRIDE_SIZE = (size_t)floor((sqrt(L3_CACHESIZE/sizeof(double))));

static constexpr size_t L2_CACHESIZE = 262144;
static size_t L2_STRIDE_SIZE = (size_t)floor((sqrt(L2_CACHESIZE/sizeof(double))));

static constexpr size_t L1_CACHESIZE = 32768;
static size_t L1_STRIDE_SIZE = (size_t)floor((sqrt(L1_CACHESIZE/sizeof(double))));

static constexpr size_t PACK_SIZE = 64;

这给出了 L1 步幅大小为 64 和 L2 步幅大小为 181。显然这些不是彼此的倍数。我有两个选择 -

每次 L2 迭代只适合 4 个 L1 块，从 0 到 63 到 127。这似乎是我在利用我的 L2 缓存。
使用整个 L2 缓存并在最后一次迭代中用 0 填充 (64*3-180) 个元素。这会引入很多冗余操作，但只会将 L2 步幅大小减少 1。
为下一次迭代预取小数块。
有一种我不知道的规范方法。

在实践中处理不是彼此倍数的块大小的最佳方法是什么？

编辑 - 响应 MSalters：

我跑了

nm -AP -t d --demangle ./src/example-GeMM | grep "intrinsic"

给出

./src/example-GeMM: void intrinsic_multiply<2048ul,3136ul,2240ul>(double*,double*,void*) W 17081 434
./src/example-GeMM: void macro_kernal_intrinsic<2048ul>(double*,double*) W 20409 234
./src/example-GeMM: void micro_kernal_intrinsic<2048ul>(double*,double*) W 21343 454

意味着相关代码部分占用 (234+454+434)/262144

解决方法

我的假设是 L1 被拆分，然后 32Kb 只是 L1d 缓存。 L2 将是统一的，这意味着您不应该将其全部用于您的数据。

另外，我还不会太担心循环提升。有了所有的 constexpr，优化器就有机会发现提升机。一个问题是 C 似乎是输出变量。因为那是 void*，所以它可以作为 size_t 的别名。如果你写了 double* C，它就不能给 size_t 起别名。这是一个例子，类型安全不仅对安全有好处，而且对速度也有好处。

请参阅以下论文（以及scholar.google.com 上的类似论文）。我认为这个很适合你的问题。（顺便说一句，我想你已经使用了这个“_mm256_mul_pd”。）反正你会在互联网上找到pdf（不想在这里链接）。

用于在 GPU 上实现高效 GEMM 的协调平铺和批处理框架， X Li,Y Liang,S Yan,L Jia,Y Li - 2019 年第 24 届研讨会论文集 - dl.acm.org

"基本问题是平铺、批处理以及它们的协同交互。平铺意味着将每个 GEMM 平铺成许多平铺。我们允许不同的 GEMM 有不同的平铺策略，而不是共享统一的平铺策略。如何统一不同的平铺策略进入单个内核是一个挑战。”

"给定一个大小为 M×N×K 的 GEMM，C 矩阵被划分为多个大小为 BY×BX 的瓦片。C 的每个瓦片需要访问大小为 BY×K 的 A 矩阵的整行部分和一个图 1(a) 中大小为 K×BX 的 B 矩阵的整个列部分。但是，A 的整个行带和 B 的列带太大，无法容纳在共享存储器和寄存器文件中。要使用片上内存，沿K维的工作负载必须被划分为许多段，如图1（b）所示。A的行段的每个段称为一个大小为BY×BK的A瓦片，列段的每个段的 B 块称为大小为 BK×BX 的 B 瓦片。最终结果可以通过沿 K 维累加每个线段的部分结果得到。"

评论：我喜欢考虑不同解决方案的想法 - 或者让它“优化”以找到最佳解决方案。（比如：你为什么要找出哪种瓷砖尺寸效果最好，让计算机找出来。好吧，很明显，编程一个瓷砖尺寸的可变设置更费力）。嗯，这是研究（我只知道你最好按常规顺序存储和访问 RAM 中的数据 - 否则你会看到“缓存未命中”，这可能会导致巨大的性能损失。另一个想法：其他一些进程可能会喜欢也住在缓存中。我不确定你是如何考虑到这一点的 - 它没有被提及。）

缓存敏感矩阵乘法的多级平铺

如何解决缓存敏感矩阵乘法的多级平铺

解决方法

相关推荐