c - 如何将 32 位 float 转换为 8 位有符号字符？ (4 :1 packing of int32 to int8 __m256i)

Question

我想做的是:

1. 将输入的 float 乘以一个固定因子。
2. 将它们转换为 8 位有符号字符。

请注意，大多数输入的值的绝对范围很小，例如 [-6, 6]，因此固定因子可以将它们映射到 [-127, 127]。

我只在 avx2 指令集上工作，所以不能使用像 _mm256_cvtepi32_epi8 这样的内部函数。我想使用 _mm256_packs_epi16 但它会将两个输入混合在一起。 :(

我还编写了一些将 32 位 float 转换为 16 位整数的代码，它的工作方式正是我想要的。

void Quantize(const float* input, __m256i* output, float quant_mult, int num_rows, int width) {
  // input is a matrix actuaaly, num_rows and width represent the number of rows and columns of the matrix
  assert(width % 16 == 0);

  int num_input_chunks = width / 16;

  __m256 avx2_quant_mult = _mm256_set_ps(quant_mult, quant_mult, quant_mult, quant_mult,
                                     quant_mult, quant_mult, quant_mult, quant_mult);

  for (int i = 0; i < num_rows; ++i) {
    const float* input_row = input + i * width;
    __m256i* output_row = output + i * num_input_chunks;
    for (int j = 0; j < num_input_chunks; ++j) {
      const float* x = input_row + j * 16;
      // Process 16 floats at once, since each __m256i can contain 16 16-bit integers.

      __m256 f_0 = _mm256_loadu_ps(x);
      __m256 f_1 = _mm256_loadu_ps(x + 8);

      __m256 m_0 = _mm256_mul_ps(f_0, avx2_quant_mult);
      __m256 m_1 = _mm256_mul_ps(f_1, avx2_quant_mult);

      __m256i i_0 = _mm256_cvtps_epi32(m_0);
      __m256i i_1 = _mm256_cvtps_epi32(m_1);

      *(output_row + j) = _mm256_packs_epi32(i_0, i_1);
    }
  }
}

欢迎任何帮助，非常感谢!

Answer 1

对于多个源 vector 的良好吞吐量，_mm256_packs_epi16 有 2 个输入 vector 而不是产生较窄的输出是一件好事。 (AVX512 _mm256_cvtepi32_epi8 不一定是最有效的做事方式，因为带有内存目标的版本解码为多个微指令，或者常规版本为您提供多个需要单独存储的小输出。 )

或者您是在提示它如何在线上运作？是的，这很烦人，但是 _mm256_packs_epi32 做同样的事情。如果您的输出可以有交错的数据组，那么也可以为此做同样的事情。

最好的办法是将 4 个 vector 合并为 1 个，分 2 步进行车道内包装(因为没有车道交叉包装)。然后使用一次交叉洗牌来修复它。

#include <immintrin.h>
// loads 128 bytes = 32 floats
// converts and packs with signed saturation to 32 int8_t
__m256i pack_float_int8(const float*p) {
    __m256i a = _mm256_cvtps_epi32(_mm256_loadu_ps(p));
    __m256i b = _mm256_cvtps_epi32(_mm256_loadu_ps(p+8));
    __m256i c = _mm256_cvtps_epi32(_mm256_loadu_ps(p+16));
    __m256i d = _mm256_cvtps_epi32(_mm256_loadu_ps(p+24));
    __m256i ab = _mm256_packs_epi32(a,b);        // 16x int16_t
    __m256i cd = _mm256_packs_epi32(c,d);
    __m256i abcd = _mm256_packs_epi16(ab, cd);   // 32x int8_t
    // packed to one vector, but in [ a_lo, b_lo, c_lo, d_lo | a_hi, b_hi, c_hi, d_hi ] order
    // if you can deal with that in-memory format (e.g. for later in-lane unpack), great, you're done

    // but if you need sequential order, then vpermd:
    __m256i lanefix = _mm256_permutevar8x32_epi32(abcd, _mm256_setr_epi32(0,4, 1,5, 2,6, 3,7));
    return lanefix;
}

(Compiles nicely on the Godbolt compiler explorer)。

在循环中调用它并_mm256_store_si256生成 vector 。

---

(对于uint8_t unsigned destination，在第16->8步使用_mm256_packus_epi16，其他保持不变。我们仍然使用signed 32->16 打包，因为 16 -> u8 vpackuswb 打包仍然将其 epi16 input 签名。你需要 -1 被视为 -1，而不是 +0xFFFF，因为无符号饱和将其钳制为 0。)

---

每个 256 位存储有 4 次洗牌，每个时钟吞吐量 1 次洗牌将成为英特尔 CPU 的瓶颈。您应该获得每个时钟一个浮点 vector 的吞吐量，瓶颈在端口 5。 (https://agner.org/optimize/)。或者，如果 L2 中的数据不热，内存带宽可能会出现瓶颈。

---

如果你只有一个单个 vector 要做，你可以考虑使用_mm256_shuffle_epi8将每个epi32元素的低字节放入每个 channel 的低32位，然后 _mm256_permutevar8x32_epi32用于过马路。

另一个单 vector 替代方案(适用于 Ryzen)是 extracti128 + 128 位 packssdw + packsswb。但是，如果你只是做一个单一的 vector ，那仍然是好的。 (仍然在 Ryzen 上，您将希望在 128 位 vector 中工作以避免额外的 channel 交叉洗牌，因为 Ryzen 将每条 256 位指令拆分为(至少)2 128 位微指令。)

相关:

• SSE - AVX conversion from double to char
• How can I convert a vector of float to short int using avx instructions?