为啥我的卷积实现比 Tensorflow 的慢？

Posted 2023-02-16

【中文标题】为啥我的卷积实现比 Tensorflow 的慢？

【英文标题】：Why is my convolution implementation so slow compared to the Tensorflow's one?为什么我的卷积实现比 Tensorflow 的慢？

【发布时间】：2020-03-06 17:28:32

【问题描述】：

我已经在 C++ 中使用 SIMD 指令实现了 VGG19 网络，仅用于推理。我想优化一个推理请求的延迟。

由于 VGG19 主要由卷积层组成，因此我主要专注于实现高效的卷积层。我在做的时候跟着这篇论文：Anatomy Of High-Performance Deep Learning Convolutions On SIMD Architectures.

我的实施提供了正确的结果。我使用 SIMD Intrisics 和论文中描述的算法。所有的重量都是预先加载的。每层的输入和输出缓冲区是在运行实际推理之前分配的。

---

我们以 VGG19 网络的第二个卷积层为例：

输入：（224、224、64）（填充后的226、226、64） 输出：(224, 224, 64) 内核：(3, 3, 64, 64) (KH, KW, C_IN, C_OUT)

下面是代码对应的代码：

void conv2d_block1_conv2(const float* in, const float* weights, float* out) 
    constexpr int VLEN = 8; // to use _mm256_* intrisics
    constexpr int C_OUT_B = VLEN;
    constexpr int C_IN_B = VLEN;

    constexpr int H = 226;           // Input Height
    constexpr int W = 226;           // Input Width
    constexpr int C_IN = 64;         // Input Channels

    constexpr int KH = 3;            // Kernel Height
    constexpr int KW = 3;            // Kernel Width

    constexpr int H_OUT = 224;       // Output Height
    constexpr int W_OUT = 224;       // Output Width
    constexpr int C_OUT = 64;        // Output Channels

    __m256 in_vec, weights_vec, out_vec;
    for (int c_out = 0; c_out < C_OUT / C_OUT_B; c_out++)
    for (int c_in_b = 0; c_in_b < C_IN / C_IN_B; c_in_b++)
    for (int h_out = 0; h_out < H_OUT; h_out++)
    for (int w_out = 0; w_out < W_OUT; w_out++)
        const int outIdx = LINEAR_4(c_out, h_out, w_out, 0, H_OUT, W_OUT, C_OUT_B);
        out_vec = _mm256_load_ps (&out[outIdx]);
        for (int kh = 0; kh < KH; kh++)
            for (int kw = 0; kw < KW; kw++)
                for (int c_in = 0; c_in < C_IN_B; c_in++)
                    const int inIdx = LINEAR_4(c_in_b, h_out + kh, w_out + kw, c_in, H, W, C_IN_B);
                    const int weightsIdx = LINEAR_6(c_out, c_in_b, kh, kw, c_in, 0, C_IN / C_IN_B, KH, KW, C_IN_B, C_OUT_B);
                    in_vec = _mm256_set1_ps (in[inIdx]);
                    weights_vec = _mm256_load_ps(&weights[weightsIdx]); 
                    out_vec = _mm256_fmadd_ps (in_vec, weights_vec, out_vec);
                    _mm256_store_ps(&out[outIdx], out_vec);
                
    

注意：我正在研究线性地址空间。函数LINEAR4 和LINEAR6 将多维索引映射到一维索引。

array[c_out][h_out][w_out][0]         <-> LINEAR_4(c_out, h_out, w_out, 0, H_OUT, W_OUT, C_OUT_B); 
array[c_out][c_in_b][kh][kw][c_in][0] <-> LINEAR_6(c_out, c_in_b, kh, kw, c_in, 0, C_IN / C_IN_B, KH, KW, C_IN_B, C_OUT_B);

---

我为每个卷积层创建了一个类似上述的函数，为编译器提供最佳优化可能性。

但是，执行时间相当糟糕。 对于整个 VGG19 网络（都是单线程执行）：

我的实现：2400ms 带有 Tensorflow 后端的 Keras 使用 model.predict(image)：600 毫秒

这种巨大的性能差距让我想知道我做错了什么。我正在使用带有-O3 标志的clang。

所以我的问题是：

是否有我没有考虑到的关键因素？ Keras/TensorFlow 使用的是哪个实现。他们怎么这么快？

【问题讨论】：

TensorFlow 在多个线程上工作，您似乎无能为力来防止这种情况 (see this question)...

我想我成功地限制了 Tensorflow 只使用一个核心，代码来自：question我的任务管理器至少没有显示其他核心的负载。

TensorFlow 是开源的，为什么不自己去挖掘一下使用的是什么实现呢？

【参考方案1】：

我找到了性能不佳的原因。 clang 编译器只使用了 2 个 SSE 寄存器，而不是所有可用的。这导致对 L1 缓存进行不必要的写入和读取。

我手动展开了两个内部循环，编译器现在使用所有可用的 16 个 SSE 寄存器。性能大幅提升。

如果您使用 SSE Intrisics，请务必检查生成的程序集。

【讨论】：

如果没有-ffast-math，clang 无法更改您的循环以使用多个累加器，即使它想这样做； FP 加法/FMA 不是严格关联的，所以你会得到不同的舍入，而不是做一个 FMA 链与分布在多个依赖链上以隐藏 FP 延迟。当使用-ffast-math 进行自动矢量化时，clang 通常将 使用 4 个依赖链展开。 （尽管如果您可以获得 2 个/时钟 FMA，这还不足以隐藏 FMA 延迟）顺便说一句，__m256 是一个 AVX 256 位向量，可以存在于 YMM 寄存器中，比 SSE XMM 寄存器宽。

相关：Why does mulss take only 3 cycles on Haswell, different from Agner's instruction tables? (Unrolling FP loops with multiple accumulators) 是关于展开多个累加器的性能细节。

非常感谢您提供这些详细信息！这很有意义。但是我实际上已经设置了这个标志并且编译器只使用了 2 个 YMM 寄存器。我还观察到，当使用常规多维数组而不是我自己进行索引计算时，编译器能够生成更好的代码。

IIRC，即使使用 -ffast-math（或整数 SIMD），clang 也不会对手动矢量化循环进行额外的展开。但是，当 auto-矢量化一个小循环时，它确实会展开 4。 （或 2 个小循环）。当然，关联性对于矢量化来说是必要的。 （有趣的事实：GCC 只启用带有配置文件引导优化的-funroll-loops，不像clang）