SSE: обратный, если не равен нулю

Как я могу взять обратную (обратную) поплавки с инструкциями SSE, но только для ненулевых значений?

Сводка ниже:

Я хочу нормализовать массив векторов, чтобы каждый размер имел одинаковое среднее значение. В C это может быть закодировано как:

float vectors[num * dim]; // input data

// step 1. compute the sum on each dimension
float norm[dim];
memset(norm, 0, dim * sizeof(float));
for(int i = 0; i < num; i++) for(int j = 0; j < dims; j++)
    norm[j] += vectors[i * dims + j];
// step 2. convert sums to reciprocal of average
for(int j = 0; j < dims; j++) if(norm[j]) norm[j] = float(num) / norm[j];
// step 3. normalize the data
for(int i = 0; i < num; i++) for(int j = 0; j < dims; j++)
    vectors[i * dims + j] *= norm[j];

Теперь по соображениям производительности я хочу сделать это, используя SSE intinsics. Setp 1 et step 3 легко, но я застрял на шаге 2. Кажется, я не нахожу какой-либо образец кода или очевидную инструкцию SSE, чтобы взять recirpocal значения , если он не равен нулю, Для деления _mm_rcp_ps делает трюк и, возможно, сочетает его с условным движением, но как получить маску, указывающую, какой компонент равен нулю?

Мне не нужен код для описанного выше алгоритма, просто функция "обратное, если не ноль":

__m128 rcp_nz_ps(__m128 input) {
    // ????
}

Спасибо!

Ответ 1

__m128 rcp_nz_ps(__m128 input) {
    __m128 mask = _mm_cmpeq_ps(_mm_set1_ps(0.0), input);
    __m128 recip = _mm_rcp_ps(input);
    return _mm_andnot_ps(mask, recip);
}

Каждая полоса mask устанавливается либо на b111...11, если вход равен нулю, а b000...00 в противном случае. И не эта маска заменяет элементы обратного, соответствующие нулевому входу с нулем.