AVX-512 浮点比较和掩码

Question

我对 SIMD 不太熟悉，但我之前用 AVX 写过一些非常简单的东西。 现在我也想用 AVX-512 实现一些旧的 AVX 代码。

我打算做什么：

// SIZE, LOW_THRESHOLD, HIGH_THRESHOLD and array are defined
// the code works with float data

for ( int index = 0; index < SIZE; ++index )
{
    array[ index ] = LOW_THRESHOLD < array[ index ] && array[ index ] < HIGH_THRESHOLD ? 1.0f : 0.0f;
}

我用 AVX 做了什么：

const __m256 lowThreshold  = _mm256_set1_ps( LOW_THRESHOLD );
const __m256 highThreshold = _mm256_set1_ps( HIGH_THRESHOLD );
const __m256 trueValue     = _mm256_set1_ps( 1.0f );
const __m256 falseValue    = _mm256_set1_ps( 0.0f );

for ( int index = 0; index < SIZE; index += 8 )
{
    // aligned load
    const __m256 val = _mm256_load_ps( array + index );
    // compare
    const __m256 comp1 = _mm256_cmp_ps( lowThreshold, val , _CMP_LT_OQ );
    const __m256 comp2 = _mm256_cmp_ps( val , highThreshold, _CMP_LT_OQ );
    // AND
    const __m256 mask = _mm256_and_ps( comp1, comp2 );
    // blend
    const __m256 result = _mm256_blendv_ps( falseValue, trueValue, mask );
    // aligned store
    _mm256_store_ps( array + index, result );
}

现在我被困在 AVX-512 上。

const __m512 lowThreshold  = _mm512_set1_ps( LOW_THRESHOLD );
const __m512 highThreshold = _mm512_set1_ps( HIGH_THRESHOLD );
const __m512 trueValue     = _mm512_set1_ps( 1.0f );
const __m512 falseValue    = _mm512_set1_ps( 0.0f );

for ( int index = 0; index < SIZE; index += 16 )
{
    // aligned load
    const __m512 val = _mm512_load_ps( array + index );

    // the result of the comparison goes into a mask?
    const __mmask16 comp1 = _mm512_cmplt_ps_mask( lowThreshold, val );
    const __mmask16 comp2 = _mm512_cmplt_ps_mask( val, highThreshold );

    // how to use these masks?
}

使用__m512 _mm512_and_ps (__m512 a, __m512 b)会很好，但是比较后只有 __mask16 变量，我没有找到任何 _mm512 function 例如_mm256_cmp_ps 。 对于更有经验的 AVX 用户来说，这可能是一个简单的问题。 谢谢！

Answer 1

如果您查看__mmask16的类型定义，您会看到： typedef unsigned short __mmask16; . 因此，您可以将这种类型视为 uint16_t 并使用 '&'。 然后你可以使用__m512 _mm512_mask_blend_ps (__mmask16 k, __m512 a, __m512 b) 。

这对我来说非常有效：

#include <immintrin.h>
#include <x86intrin.h>
#include <stdio.h>

#define LOW_THRESHOLD 1
#define HIGH_THRESHOLD 3

int main() {
    __m512 lowThreshold  = _mm512_set1_ps( LOW_THRESHOLD );
    __m512 highThreshold = _mm512_set1_ps( HIGH_THRESHOLD );
    __m512 trueValue     = _mm512_set1_ps( 1.0f );
    __m512 falseValue    = _mm512_set1_ps( 0.0f );

    float array[16] = {-5.0f,6.0f,4.0f,1.5f,0.7f,1.0f,1.0f,-5.0f,6.0f,4.0f,1.5f,0.7f,1.0f,1.0f,-5.0f,6.0f};

    for (int i = 0; i < 16; i += 1) {
        printf("%5.1f ", array[i]);
    }
    printf("\n");

    for ( int index = 0; index < 16; index += 16 )
    {
        __m512 val = _mm512_loadu_ps( array + index );
        __mmask16 comp1 = _mm512_cmplt_ps_mask( lowThreshold, val );
        __mmask16 comp2 = _mm512_cmplt_ps_mask( val, highThreshold );
        __mmask16 mask = comp1 & comp2;
        __m512 result = _mm512_mask_blend_ps(mask, falseValue, trueValue);
        _mm512_storeu_ps( array + index, result );
    }
    for (int i = 0; i < 16; i += 1) {
        printf("%5.1f ", array[i]);
    }
    printf("\n");
}