优化SIMD版本的范围生成算法

Question

在

C

中的一个项目中，我想快速生成具有起点、增量值和终点的双精度浮点数的一维向量。对我来说重要的是它必须尽可能快。我打算使用

SIMD

内在函数来优化它。我已经使用

C

内在函数在

SIMD

中编写了它。它按预期工作。例如，当我像

generate_range_avx(0.0, 0.5, 3.0)

这样称呼它时，它会生成一个向量

[0.0 0.5 1.0 1.5 2. 2.5 3.0]

。然而，

SIMD

版本比naive版本慢2-3倍。

我不精通

SIMD

指令或内在函数。所以，我的矢量化可能是错误的，而且效率肯定很低。

这是代码。

#include <stdio.h>
#include <stdlib.h>
#include <memory.h>

#include <math.h>
#include <immintrin.h> 

#define CALCELTS(start, increment, end) ((size_t)round(((end) - (start)) / (increment)) + 1)

static size_t align(size_t num, size_t alignment) {
    return (num + alignment - 1) & ~(alignment - 1);
}

typedef struct
{
    double* elts;
    size_t len;
} vector_t;

vector_t generate_range_naive(double start, double increment, double end) {
    // Calculate the unaligned number of elements
    size_t num_elements = CALCELTS(start, increment, end);

    size_t aligned_num_elts = align(num_elements, 4);

    double* data = malloc(sizeof * data * aligned_num_elts);
    if (!data)
    {
        fprintf(stderr, "Failed to allocate memory.");
        exit(1);
    }
    for (size_t i = 0; i < num_elements; i++) {
        data[i] = start;
        start += increment;
    }
    data[num_elements - 1] = end;
    return (vector_t) { data, num_elements };
}
vector_t generate_range_avx(double start, double increment, double end) {
    // Calculate the unaligned number of elements
    size_t num_elements = CALCELTS(start, increment, end);

    if (num_elements < 4)
        return generate_range_naive(start, increment, end);

    size_t aligned_num_elts = align(num_elements, 4);

    // Allocate memory for the vector.
    double* data = _mm_malloc(sizeof(*data) * aligned_num_elts, 32);

    if (!data)
    {
        fprintf(stderr, "Failed to allocate memory.");
        exit(1);
    } 
    __m256d start_v = _mm256_set1_pd(start);
    __m256d indices_v = _mm256_set_pd(3.0, 2.0, 1.0, 0.0);
    __m256d increment_v = _mm256_set1_pd(increment);
    __m256d x_v = _mm256_add_pd(start_v, _mm256_mul_pd(indices_v, increment_v));
    increment_v = _mm256_mul_pd(increment_v, _mm256_set1_pd(4.0));

    for (size_t i = 0; i < aligned_num_elts; i += 4) {
        _mm256_stream_pd(data + i, x_v);
        x_v = _mm256_add_pd(x_v, increment_v);
    }
    data[num_elements - 1] = end;
    return (vector_t) { data, num_elements };
} 

static void print_range(vector_t vec)
{
    for (size_t i = 0; i < vec.len; i++)
    {
        printf("%lf ", vec.elts[i]);
    }
    putchar('\n');
}

int main(int argc, char* argv[])
{
    // I benchmarked the functions using my simple benchmark library but didn't include it here. 
    // If the need to use it arises, I can attach it too. 
    vector_t vec1 = generate_range_avx(1.0, 0.5, 100.0);
    print_range(vec1);
    vector_t vec2 = generate_range_naive(1.0, 0.5, 100.0);
    print_range(vec2);
    return 0;
}

我的开发环境是

Windows

/

MSVC

，目标架构是

x86_64

。因此，我使用了编译器内在函数。

• 我的代码有什么问题？
• 可以优化吗？
• 或者是否没有使用 SIMD 来解决此类范围生成问题的优化方法？

注意：问题中的语言是

C

，但如果您愿意，可以用

x86_64

汇编代码来回答。

Answer 1

尝试以下版本。与您的实现相同的想法，有两个变化。

1. 用常规矢量存储替换了

   _mm256_stream_pd

   。

2. 将循环展开 4。我认为您的版本可能会在

   _mm256_add_pd

   指令的延迟上遇到瓶颈。根据 uops.info，在大多数现代处理器上，延迟为 3-4 个时钟周期。这意味着无论循环中还有什么，由于连续的数据依赖链，原始循环都会限制为 3-4 个周期/迭代。出于同样的原因，我的循环也这样做，但是我的循环的 1 次迭代生成 16 个数字，而不是您的版本中的 4 个数字。

#include <immintrin.h>
#include <math.h>

struct vector_t
{
    double* elts;
    size_t len;
};

size_t computeCount( double start, double increment, double end )
{
    double len = end - start;
    len /= increment;
    long num = lround( ceil( len ) );
    return (size_t)num;
}

struct vector_t generateRange( double start, double increment, double end )
{
    size_t lenExact = computeCount( start, increment, end );
    size_t lenAligned = ( lenExact + 15 ) & ( ~(size_t)15 );

    double* rdi = (double*)_mm_malloc( lenAligned * sizeof( double ), 32 );

    struct vector_t result;
    result.len = lenExact;
    result.elts = rdi;

    __m256d vIndices = _mm256_set_pd( 3.0, 2.0, 1.0, 0.0 );
    __m256d vInc1 = _mm256_set1_pd( increment );
    vIndices = _mm256_mul_pd( vIndices, vInc1 );

    __m256d v0 = _mm256_set1_pd( start );
    v0 = _mm256_add_pd( v0, vIndices );

    __m256d vInc4 = _mm256_mul_pd( vInc1, _mm256_set1_pd( 4 ) );
    __m256d v1 = _mm256_add_pd( v0, vInc4 );
    __m256d vInc8 = _mm256_add_pd( vInc4, vInc4 );
    __m256d v2 = _mm256_add_pd( v0, vInc8 );
    __m256d v3 = _mm256_add_pd( v2, vInc4 );
    __m256d vInc = _mm256_add_pd( vInc8, vInc8 );

    double* endPointer = rdi + lenAligned;
    for( ; rdi < endPointer; rdi += 16 )
    {
        _mm256_store_pd( rdi, v0 );
        _mm256_store_pd( rdi + 4, v1 );
        _mm256_store_pd( rdi + 8, v2 );
        _mm256_store_pd( rdi + 12, v3 );

        v0 = _mm256_add_pd( v0, vInc );
        v1 = _mm256_add_pd( v1, vInc );
        v2 = _mm256_add_pd( v2, vInc );
        v3 = _mm256_add_pd( v3, vInc );
    }

    return result;
}