如何在AVX/AVX2中实现Pair-wise与Component-wise内存布局转换？

问题背景 #

我正在为AVX/AVX2编写一个double-double算术库，遇到的问题是非SIMD版本与SIMD版本采用不同的内存布局：

两种布局定义 #

// Pair-wise 布局：内存中按{hi, lo}对连续存储，顺序为{x0,y0}, {x1,y1}, {x2,y2}, {x3,y3}
struct Float64x2 { 
    double hi;
    double lo;
};

// Component-wise 布局：所有hi分量和lo分量分别打包存储，顺序为{x0,x1,x2,x3}, {y0,y1,y2,y3}
struct __m256dx2 { 
    __m256d hi; // 存储4个hi分量：x0, x1, x2, x3
    __m256d lo; // 存储4个lo分量：y0, y1, y2, y3
};

需要实现两种布局的双向转换：从Pair-wise加载转换为Component-wise，以及从Component-wise存储转换为Pair-wise。

转换实现代码 #

1. Pair-wise → Component-wise（加载转换） #

该函数从Pair-wise布局的内存中加载数据，并转换为Component-wise的__m256dx2结构：

#include <immintrin.h>

struct __m256dx2 { 
    __m256d hi;
    __m256d lo;
};

__m256dx2 _mm256x2_loadu_pdx2(const double* mem_addr) {
    // 加载8个double到两个256位寄存器
    __m256d vec0 = _mm256_loadu_pd(mem_addr);       // [x0, y0, x1, y1]
    __m256d vec1 = _mm256_loadu_pd(mem_addr + 4);   // [x2, y2, x3, y3]

    // 提取所有hi分量的临时结果：[x0, x2, x1, x3]
    __m256d hi_temp = _mm256_unpacklo_pd(vec0, vec1);
    // 提取所有lo分量的临时结果：[y0, y2, y1, y3]
    __m256d lo_temp = _mm256_unpackhi_pd(vec0, vec1);

    // 重排hi分量顺序为[x0, x1, x2, x3]
    __m256d hi = _mm256_permute4x64_pd(hi_temp, 0x27);
    // 重排lo分量顺序为[y0, y1, y2, y3]
    __m256d lo = _mm256_permute4x64_pd(lo_temp, 0x27);

    __m256dx2 result;
    result.hi = hi;
    result.lo = lo;
    return result;
}

指令说明

• _mm256_unpacklo_pd：从两个输入寄存器的低128位段中，分别取每个double对的第0个元素，交叉组合成新的寄存器。
• _mm256_unpackhi_pd：从两个输入寄存器的低128位段中，分别取每个double对的第1个元素，交叉组合成新的寄存器。
• _mm256_permute4x64_pd：对256位寄存器中的4个double元素进行任意重排，0x27的二进制是00100111，对应将临时结果的顺序调整为目标顺序。

2. Component-wise → Pair-wise（存储转换） #

该函数将Component-wise的__m256dx2结构转换为Pair-wise布局并存储到内存：

void _mm256x2_storeu_pdx2(double* mem_addr, __m256dx2 val) {
    __m256d hi = val.hi; // [x0, x1, x2, x3]
    __m256d lo = val.lo; // [y0, y1, y2, y3]

    // 交叉组合得到Pair-wise布局的寄存器
    __m256d vec0 = _mm256_unpacklo_pd(hi, lo); // [x0, y0, x1, y1]
    __m256d vec1 = _mm256_unpackhi_pd(hi, lo); // [x2, y2, x3, y3]

    // 存储到内存
    _mm256_storeu_pd(mem_addr, vec0);
    _mm256_storeu_pd(mem_addr + 4, vec1);
}

指令说明

• _mm256_unpacklo_pd：将hi的低两个元素与lo的低两个元素交叉，组成{x0,y0}, {x1,y1}的布局。
• _mm256_unpackhi_pd：将hi的高两个元素与lo的高两个元素交叉，组成{x2,y2}, {x3,y3}的布局。

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