读取/写入步幅远大于其宽度的矩阵会导致性能大幅下降

Question

我正在对 1024x1024 矩阵进行密集矩阵乘法。我使用 64x64 瓷砖使用循环阻塞/平铺来做到这一点。我创建了一个高度优化的 64x64 矩阵乘法函数（代码见我问题的结尾）。

gemm64(float *a, float *b, float *c, int stride).  

这是在图块上运行的代码。一个 1024x1204 矩阵，包含 16x16 块。

for(int i=0; i<16; i++) {
    for(int j=0; j<16; j++) {
        for(int k=0; k<16; k++) {
            gemm64(&a[64*(i*1024 + k)], &b[64*(k*1024 + j)], &c[64*(i*1024 + j)], 1024);
        }
    }
}

但是，作为一种猜测，我尝试在新矩阵 b2 中重新排列 b 矩阵的所有图块的内存（代码参见本问题的末尾），以便每个图块的步幅为64 而不是 1024。这有效地生成了一个 64x64 矩阵的数组，stride=64。

float *b2 = (float*)_mm_malloc(1024*1024*sizeof(float), 64);
reorder(b, b2, 1024);
for(int i=0; i<16; i++) {
    for(int j=0; j<16; j++) {
        for(int k=0; k<16; k++) {
            gemm64_v2(&a[64*(i*1024 + k)], &b2[64*64*(k*16 + j)], &c[64*(i*1024 + j)], 64);
        }
    }
}
_mm_free(b2);

注意 b 的偏移量如何从 &b[64*(k*1024 + j)] 更改为 &b2[64*64*(k*16 + j)] 并且传递给 gemm64 的步幅已从 1024 更改为 64。

在我的 Sandy Bridge 系统上，我的代码性能从低于峰值 flop 的 20% 跃升至 70%！

为什么以这种方式重新排列 b 矩阵中的图块会产生如此巨大的差异？

数组 a、b、b2 和 c 是 64 字节对齐的。

extern "C" void gemm64(float *a, float*b, float*c, int stride) {
    for(int i=0; i<64; i++) {
        row_m64x64(&a[1024*i], b, &c[1024*i], stride);
    }
}

void row_m64x64(const float *a, const float *b, float *c, int stride) {
    __m256 tmp0, tmp1, tmp2, tmp3, tmp4, tmp5, tmp6, tmp7;
    tmp0 = _mm256_loadu_ps(&c[ 0]);
    tmp1 = _mm256_loadu_ps(&c[ 8]);
    tmp2 = _mm256_loadu_ps(&c[16]);
    tmp3 = _mm256_loadu_ps(&c[24]);
    tmp4 = _mm256_loadu_ps(&c[32]);
    tmp5 = _mm256_loadu_ps(&c[40]);
    tmp6 = _mm256_loadu_ps(&c[48]);
    tmp7 = _mm256_loadu_ps(&c[56]);

    for(int i=0; i<64; i++) {
        __m256 areg0 = _mm256_broadcast_ss(&a[i]);

        __m256 breg0 = _mm256_loadu_ps(&b[stride*i +  0]);
        tmp0 = _mm256_add_ps(_mm256_mul_ps(areg0,breg0), tmp0);    
        __m256 breg1 = _mm256_loadu_ps(&b[stride*i +  8]);
        tmp1 = _mm256_add_ps(_mm256_mul_ps(areg0,breg1), tmp1);
        __m256 breg2 = _mm256_loadu_ps(&b[stride*i + 16]);
        tmp2 = _mm256_add_ps(_mm256_mul_ps(areg0,breg2), tmp2);    
        __m256 breg3 = _mm256_loadu_ps(&b[stride*i + 24]);
        tmp3 = _mm256_add_ps(_mm256_mul_ps(areg0,breg3), tmp3);   
        __m256 breg4 = _mm256_loadu_ps(&b[stride*i + 32]);
        tmp4 = _mm256_add_ps(_mm256_mul_ps(areg0,breg4), tmp4);    
        __m256 breg5 = _mm256_loadu_ps(&b[stride*i + 40]);
        tmp5 = _mm256_add_ps(_mm256_mul_ps(areg0,breg5), tmp5);    
        __m256 breg6 = _mm256_loadu_ps(&b[stride*i + 48]);
        tmp6 = _mm256_add_ps(_mm256_mul_ps(areg0,breg6), tmp6);    
        __m256 breg7 = _mm256_loadu_ps(&b[stride*i + 56]);
        tmp7 = _mm256_add_ps(_mm256_mul_ps(areg0,breg7), tmp7);    
    }
    _mm256_storeu_ps(&c[ 0], tmp0);
    _mm256_storeu_ps(&c[ 8], tmp1);
    _mm256_storeu_ps(&c[16], tmp2);
    _mm256_storeu_ps(&c[24], tmp3);
    _mm256_storeu_ps(&c[32], tmp4);
    _mm256_storeu_ps(&c[40], tmp5);
    _mm256_storeu_ps(&c[48], tmp6);
    _mm256_storeu_ps(&c[56], tmp7);
}

重新排序 b 矩阵的代码。

reorder(float *b, float *b2, int stride) {
    //int k = 0;
    for(int i=0; i<16; i++) {
        for(int j=0; j<16; j++) {
            for(int i2=0; i2<64; i2++) {
                for(int j2=0; j2<64; j2++) {
                    //b2[k++] = b[1024*(64*i+i2) + 64*j + j2];
                    b2[64*64*(i*16 + j) + 64*i2+j2] = b[1024*(64*i+i2) + 64*j + j2];
                }
            }
        }
    }
}

Answer 1

我想我可能已经找到了答案。我认为这与 Translation Look-aside Buffer (TLB) 有关。

他们在 Goto 和 Van de Geijn 的论文 Anatomy of High-Performance Matrix Multiplication 中写道

缓存未命中和 TLB 未命中之间最显着的区别是 高速缓存未命中不一定会使 CPU 停顿……相比之下，TLB 未命中会导致 CPU 停止直到 TLB 被新地址更新。 换句话说， 预取可以屏蔽缓存未命中，但不能屏蔽 TLB 未命中。”

在 4.2.3 节标题为“包装”之后不久，他们写道

现在的基本问题是 A 通常是较大矩阵的子矩阵，因此在内存中不连续。这反过来意味着解决它需要的不仅仅是最小数量的 TLB 条目。 解决方案是将 A 打包到一个连续的工作数组中

Answer 2

我认为问题在于最里面的循环，并且与预取有关。您正在执行以下操作：

1. 从 c 加载 256 个连续字节（4 个缓存行）
2. 从 b 加载 256 个连续字节 64 次（总共 256 个缓存行）
3. 将 256 个连续字节保存到 c

在第 2 步中，当 stride 为 64 时，您实际上是在读取一个长的连续内存块。当 stride 为 1024 时，您正在以不连续的步骤读取内存，一次 256 个字节。

因此，当 stride 为 64 时，预取器会提前将连续块读取到缓存中，并且每行最多有一次缓存未命中。当 stride 为 1024 时，预取器会被不连续读取（交替读取连续高速缓存行和广泛分离的高速缓存行）混淆，并且每行有 64 次高速缓存未命中。