连续AVX2加载生成大量vmovaps指令的原因及优化咨询

我希望加载多个连续SIMD向量，并将其存储到编译期已知大小的容器中（这样就能为不同寄存器数量的架构编写可移植代码：AVX2架构加载3个向量，Arm Neon架构加载4个）。但查看AVX2汇编代码时，我对生成的大量vmovaps指令感到惊讶。

考虑以下C++代码：

#include <array>
#include <immintrin.h>
using InpT = const float* __restrict__;

template <size_t unroll_count>
auto load(InpT data, size_t index) {
    using vecs = std::array<__m256, unroll_count>;
    alignas(32) vecs loaded;
    for (size_t i = 0; i < unroll_count; ++i) {
        loaded[i] = _mm256_load_ps(data + index);
        index += 8;
    }
    return loaded;
}

template auto load<3ul>(InpT data, size_t index);

auto loadSingle(InpT data, size_t index) {
    auto vec = _mm256_load_ps(data + index);
    return vec;
}

通用加载操作生成的AVX2汇编代码如下：

auto load<3ul>(float const*, unsigned long):
        mov     rax, rdi
        vmovaps ymm0, ymmword ptr [rsi + 4*rdx]
        vmovaps ymm1, ymmword ptr [rsi + 4*rdx + 32]
        vmovaps ymm2, ymmword ptr [rsi + 4*rdx + 64]
        vmovaps ymmword ptr [rdi + 64], ymm2
        vmovaps ymmword ptr [rdi + 32], ymm1
        vmovaps ymmword ptr [rdi], ymm0
        vzeroupper
        ret

而仅加载单个向量的代码符合预期：

loadSingle(float const*, unsigned long):
        vmovaps ymm0, ymmword ptr [rdi + 4*rsi]
        ret

加载三个连续向量的操作似乎在两个不同内存区域间移动数据。请问load<3>的代码是否最优？若否，该如何改进？

当前load<3>的代码并非最优，问题出在返回std::array<__m256, 3>的实现方式上：编译器需要先把加载的向量存入栈上的loaded数组，再把整个数组拷贝到函数的返回值存储区（由调用方传入的rdi指针指向），这就产生了额外的vmovaps存储指令。

优化方向1：直接返回寄存器组（利用结构化绑定） #

如果调用方可以通过结构化绑定接收返回的向量，我们可以让函数直接返回多个__m256对象，而非std::array。编译器会自动将这些向量放在寄存器中返回，避免内存拷贝：

template <size_t unroll_count>
auto load(InpT data, size_t index) {
    if constexpr (unroll_count == 1) {
        return std::make_tuple(_mm256_load_ps(data + index));
    } else if constexpr (unroll_count == 2) {
        auto v0 = _mm256_load_ps(data + index);
        auto v1 = _mm256_load_ps(data + index + 8);
        return std::make_tuple(v0, v1);
    } else if constexpr (unroll_count == 3) {
        auto v0 = _mm256_load_ps(data + index);
        auto v1 = _mm256_load_ps(data + index + 8);
        auto v2 = _mm256_load_ps(data + index + 16);
        return std::make_tuple(v0, v1, v2);
    }
    // 可扩展更多unroll_count的分支
}

// 使用方式
auto [v0, v1, v2] = load<3>(data, idx);

对应的AVX2汇编会直接把三个向量放在ymm0、ymm1、ymm2中返回，不需要额外的存储操作：

auto load<3ul>(float const*, unsigned long):
        vmovaps ymm0, ymmword ptr [rsi + 4*rdx]
        vmovaps ymm1, ymmword ptr [rsi + 4*rdx + 32]
        vmovaps ymm2, ymmword ptr [rsi + 4*rdx + 64]
        vzeroupper
        ret

优化方向2：让调用方提供目标数组指针 #

如果必须使用数组形式存储，可以让调用方传入预先分配好的__m256数组指针，直接将向量加载到目标位置，避免中间拷贝：

template <size_t unroll_count>
void load(InpT data, size_t index, __m256* __restrict__ dest) {
    for (size_t i = 0; i < unroll_count; ++i) {
        dest[i] = _mm256_load_ps(data + index);
        index += 8;
    }
}

// 使用方式
alignas(32) __m256 vecs[3];
load<3>(data, idx, vecs);

这种方式生成的汇编会直接把每个向量加载到dest指向的内存，没有多余的中间存储操作。

原代码产生额外拷贝的原因 #

原代码中std::array作为返回值时，虽然C++的返回值优化（RVO）会尝试省略拷贝，但对于包含SIMD类型的数组，编译器为保证对齐要求，会先在函数栈上创建loaded数组，再将其整体拷贝到返回值的内存区域，从而产生额外的vmovaps指令。

内容的提问来源于stack exchange，提问作者fabian