如何让MSVC生成将内存缓存到寄存器的汇编代码？

背景与现有实现 #

我定义了一个代表float[4][4]的mat4类型，内部复用512位寄存器存储，代码如下：

union alignas(16 * sizeof(float)) mat4 {
private: 
    __m512 m512; 
    __m512d m512d; 
    ALWAYS_INLINE mat4(__m512 m512) : m512{m512} {} 
    ALWAYS_INLINE mat4(__m512d m512d) : m512d{m512d} {} 
    ALWAYS_INLINE operator __m512&() { return m512; } 
    ALWAYS_INLINE operator __m512d&() { return m512d; } 
    ALWAYS_INLINE operator const __m512&() const { return m512; } 
    ALWAYS_INLINE operator const __m512d&() const { return m512d; } 
    ALWAYS_INLINE mat4& operator=(__m512 _m512) { m512 = _m512; return *this; } 
    ALWAYS_INLINE mat4& operator=(__m512d _m512d) { m512d = _m512d; return *this; } 
public: 
    friend void __vectorcall transform_children(mat4 parent, std::span<mat4> children); 
};

基于这个类型，我实现了transform_children函数，用AVX512F intrinsics优化，原地将每个子矩阵与父矩阵相乘：

void __vectorcall transform_children(mat4 parent, std::span<mat4> children) { 
    mat4* const __restrict bs = children.data(); 
    const size_t n = children.size(); 
    ASSUME(n != 0); 
    const mat4 zmm1 = _mm512_permute_ps(parent, 0); 
    const mat4 zmm2 = _mm512_permute_ps(parent, 85); 
    const mat4 zmm3 = _mm512_permute_ps(parent, 170); 
    const mat4 zmm0 = _mm512_permute_ps(parent, 255); 
    for (int i = 0; i < n; ++i) { 
        mat4& __restrict zmm4 = bs[i]; 
        mat4 zmm5 = _mm512_shuffle_f64x2(zmm4, zmm4, 85); 
        zmm5 = _mm512_mul_ps(zmm5, zmm2); 
        mat4 zmm6 = _mm512_shuffle_f64x2(zmm4, zmm4, 0); 
        zmm6 = _mm512_fmadd_ps(zmm1, zmm6, zmm5); 
        zmm5 = _mm512_shuffle_f64x2(zmm4, zmm4, 170); 
        zmm4 = _mm512_shuffle_f64x2(zmm4, zmm4, 255); 
        zmm4 = _mm512_fmadd_ps(zmm0, zmm4, zmm6); 
        zmm4 = _mm512_fmadd_ps(zmm3, zmm5, zmm4); 
    } 
}

MSVC的编译问题 #

GCC和Clang能将这段代码编译为最优汇编：仅对每个子矩阵做一次内存加载和一次内存写回。但MSVC的编译结果不符合预期，它对内存进行了4次访问（加载1次，写回3次），生成的汇编如下：

void transform_children(mat4,std::span<mat4,4294967295>) PROC ; transform_children, COMDAT
    mov ecx, DWORD PTR _children$[esp]
    vpermilps zmm4, zmm0, 0
    vpermilps zmm5, zmm0, 85
    vpermilps zmm6, zmm0, 170
    vpermilps zmm7, zmm0, 255
    test ecx, ecx
    je SHORT $LN36@transform_
    mov eax, DWORD PTR _children$[esp-4]
    npad    8
$LL4@transform_:
    lea eax, DWORD PTR [eax+64]
    vmovupd zmm3, ZMMWORD PTR [eax-64] ; 第一次加载
    vshuff64x2 zmm0, zmm3, zmm3, 85
    vmulps zmm0, zmm0, zmm5
    vshuff64x2 zmm1, zmm3, zmm3, 0
    vmovups zmm2, zmm4
    vfmadd213ps zmm2, zmm1, zmm0
    vshuff64x2 zmm0, zmm3, zmm3, 255
    vmovupd ZMMWORD PTR [eax-64], zmm0 ; 第一次写回
    vfmadd231ps zmm2, zmm7, ZMMWORD PTR [eax-64] ; 再次读取内存
    vshuff64x2 zmm1, zmm3, zmm3, 170
    vmovups zmm0, zmm6
    vfmadd213ps zmm0, zmm1, zmm2
    vmovups ZMMWORD PTR [eax-64], zmm0 ; 第二次写回
    sub ecx, 1
    jne SHORT $LL4@transform_
$LN36@transform_:
    vzeroupper
    ret  8
void transform_children(mat4,std::span<mat4,4294967295>) ENDP ; transform_children

解决方案 #

问题根源是MSVC对mat4& __restrict zmm4 = bs[i];的引用优化不足，没有将其绑定到寄存器，而是反复读写内存。可以通过显式将数据加载到局部寄存器变量，完成所有计算后再一次性写回内存，强制MSVC优化：

方案1：直接使用__m512 intrinsics显式加载/存储 #

void __vectorcall transform_children(mat4 parent, std::span<mat4> children) { 
    mat4* const __restrict bs = children.data(); 
    const size_t n = children.size(); 
    ASSUME(n != 0); 
    const __m512 zmm1 = _mm512_permute_ps(parent, 0); 
    const __m512 zmm2 = _mm512_permute_ps(parent, 85); 
    const __m512 zmm3 = _mm512_permute_ps(parent, 170); 
    const __m512 zmm0 = _mm512_permute_ps(parent, 255); 
    for (int i = 0; i < n; ++i) { 
        // 显式从内存加载到寄存器
        __m512 zmm4 = _mm512_loadu_ps(reinterpret_cast<float*>(&bs[i]));
        
        __m512 zmm5 = _mm512_shuffle_f64x2(zmm4, zmm4, 85);
        zmm5 = _mm512_mul_ps(zmm5, zmm2);
        __m512 zmm6 = _mm512_shuffle_f64x2(zmm4, zmm4, 0);
        zmm6 = _mm512_fmadd_ps(zmm1, zmm6, zmm5);
        zmm5 = _mm512_shuffle_f64x2(zmm4, zmm4, 170);
        __m512 temp = _mm512_shuffle_f64x2(zmm4, zmm4, 255);
        temp = _mm512_fmadd_ps(zmm0, temp, zmm6);
        temp = _mm512_fmadd_ps(zmm3, zmm5, temp);
        
        // 最后一次性写回内存
        _mm512_storeu_ps(reinterpret_cast<float*>(&bs[i]), temp);
    } 
}

方案2：显式拷贝到局部mat4变量 #

如果希望保留mat4类型的封装，也可以先把bs[i]拷贝到局部变量（编译器会将其分配到寄存器），操作完成后再写回：

void __vectorcall transform_children(mat4 parent, std::span<mat4> children) { 
    mat4* const __restrict bs = children.data(); 
    const size_t n = children.size(); 
    ASSUME(n != 0); 
    const mat4 zmm1 = _mm512_permute_ps(parent, 0); 
    const mat4 zmm2 = _mm512_permute_ps(parent, 85); 
    const mat4 zmm3 = _mm512_permute_ps(parent, 170); 
    const mat4 zmm0 = _mm512_permute_ps(parent, 255); 
    for (int i = 0; i < n; ++i) { 
        // 显式拷贝到局部变量（寄存器存储）
        mat4 zmm4 = bs[i];
        
        mat4 zmm5 = _mm512_shuffle_f64x2(zmm4, zmm4, 85);
        zmm5 = _mm512_mul_ps(zmm5, zmm2);
        mat4 zmm6 = _mm512_shuffle_f64x2(zmm4, zmm4, 0);
        zmm6 = _mm512_fmadd_ps(zmm1, zmm6, zmm5);
        zmm5 = _mm512_shuffle_f64x2(zmm4, zmm4, 170);
        zmm4 = _mm512_shuffle_f64x2(zmm4, zmm4, 255);
        zmm4 = _mm512_fmadd_ps(zmm0, zmm4, zmm6);
        zmm4 = _mm512_fmadd_ps(zmm3, zmm5, zmm4);
        
        // 最后一次性写回内存
        bs[i] = zmm4;
    } 
}

这两种方案都会强制MSVC将数据留在寄存器中完成所有计算，仅执行一次内存加载和一次写回，与GCC/Clang的优化效果一致。

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