图像沿列方向降维的AVX2实现讨论

图像沿列方向降维的AVX2实现讨论

对于二维图像，取每一列像素序列的和或均值是一种基础算子。

本文主要介绍本人的优化思路，因此本文中要求图像高度不能超过256，这样就可以保证对于常见的8bit图像，每列求和的结果使用uint16_t保存即可保证不溢出，同时要求图像宽度是32或64的倍数等，省略对每行尾部无法并行的部分的计算。

普通指针实现

void ReduceColSum(const uint8_t* src, uint16_t* dst, uint32_t height,
                   uint32_t width) {
  for (size_t y = 0; y < height; ++y) {
    for (size_t x = 0; x < width; ++x) {
      dst[x] += (uint16_t)src[y * width + x];
    }
  }
}

当前OpenCV在4.6.0中的cv::Reduce代码中使用的便是使用通用寄存器执行for循环，不过对不同的算术运算封装成了模板，并使用了×4的循环展开。

AVX指令集实现

两种基本实现

第一种AVX指令集的实现如下函数ReduceColSum1所示。循环内部avx指令相关的操作主要包括load256->permute->unpacklo/unpackhi->adds->storeu。

void ReduceColSum1(const uint8_t* src, uint16_t* dst, uint32_t height,
                   uint32_t width) {
  assert(width % 32 == 0);
  assert(height <= 256);
  __m256i mSrc, mSum, mSum2;
  int nCount = (width / 32);
  const uint8_t* psrc = src;
  const __m256i m0 = _mm256_setzero_si256();
  for (size_t y = 0; y < height; ++y) {
    uint16_t* pdst = dst;
    for (size_t x = 0; x < nCount; ++x) {
      mSrc = _mm256_loadu_si256((__m256i*)(psrc));
      mSrc = _mm256_permute4x64_epi64(mSrc, 0xD8);

      mSum = _mm256_loadu_si256((__m256i*)(pdst));
      mSum2 = _mm256_loadu_si256((__m256i*)(pdst + 16));

      mSum = _mm256_adds_epu16(mSum, _mm256_unpacklo_epi8(mSrc, m0));
      _mm256_store_si256((__m256i*)(pdst), mSum);
      mSum2 = _mm256_adds_epu16(mSum2, _mm256_unpackhi_epi8(mSrc, m0));
      _mm256_storeu_si256((__m256i*)(pdst + 16), mSum2);
      pdst += 32;
      psrc += 32;
    }
  }
}

第二种实现如下ReduceColSum2函数所示，实现是将permute->unpacklo/unpackhi替换为AVX2指令集中的convert指令。可以代码更简洁易懂。但每个内层循环只能读取并处理16字节。经测试两种实现耗时大体相当。

void ReduceColSum2(const uint8_t* src, uint16_t* dst, uint32_t height,
                   uint32_t width) {
  assert(width % 32 == 0);
  assert(height <= 256);
  __m128i mSrc;
  __m256i mSum;
  const uint8_t* psrc = src;
  const int nCount = (width / 16);
  for (size_t y = 0; y < height; ++y) {
    uint16_t* pdst = dst;
    for (size_t x = 0; x < nCount; ++x) {
      mSum = _mm256_loadu_si256((__m256i*)(pdst));
      mSrc = _mm_loadu_si128((const __m128i*)(psrc));
      mSum = _mm256_adds_epu16(mSum, _mm256_cvtepu8_epi16(mSrc));
      _mm256_storeu_si256((__m256i*)(pdst), mSum);
      pdst += 16;
      psrc += 16;
    }
  }
}

测试_mm256_cvtepu8_epi32

将ReduceColSum2中_mm256_cvtepu8_epi16替换为_mm256_cvtepu8_epi32，取消图像高度不大于256的限制，经测试耗时增加了一倍左右。

void ReduceColSum3(const uint8_t* src, uint16_t* dst, uint32_t height,
                   uint32_t width) {
  assert(width % 32 == 0);
  assert(height <= 256);
  __m128i mSrc;
  __m256i mSum;
  const uint8_t* psrc = src;
  const int nCount = (width / 8);
  for (size_t y = 0; y < height; ++y) {
    uint16_t* pdst = dst;
    for (size_t x = 0; x < nCount; ++x) {
      mSum = _mm256_loadu_si256((__m256i*)(pdst));
      mSrc = _mm_loadu_si128((const __m128i*)(psrc));
      mSum = _mm256_add_epi32(mSum, _mm256_cvtepu8_epi32(mSrc));
      _mm256_storeu_si256((__m256i*)(pdst), mSum);
      pdst += 8;
      psrc += 8;
    }
  }
}

使用循环展开优化

可以看到列求和操作计算部分包括cvt和add两步骤，两条指令的Latency是比较低的，反之读写操作loadu和storeu的Latency相对较高，因此可以使用循环展开进行优化。ReduceColSum3和ReduceColSum4分别是展开2倍和4倍的代码，经测试耗时约是不使用循环展开的70%和50%。

void ReduceColSum3(const uint8_t* src, uint16_t* dst, uint32_t height,
                   uint32_t width) {
  assert(width % 32 == 0);
  assert(height <= 256);
 __m128i mSrc, mSrc1;
  __m256i mSum, mSum1;
  const uint8_t* psrc = src;
  const int nCount = (width / 32);
  for (size_t y = 0; y < height; ++y) {
    uint16_t* pdst = dst;
    for (size_t x = 0; x < nCount; ++x) {
      mSum = _mm256_loadu_si256((__m256i*)(pdst));
      mSum1 = _mm256_loadu_si256((__m256i*)(pdst+16));
      mSrc = _mm_loadu_si128((const __m128i*)(psrc));
      mSrc1 = _mm_loadu_si128((const __m128i*)(psrc + 16));

      mSum = _mm256_adds_epu16(mSum, _mm256_cvtepu8_epi16(mSrc));
      _mm256_storeu_si256((__m256i*)(pdst), mSum);
      mSum1 = _mm256_adds_epu16(mSum1, _mm256_cvtepu8_epi16(mSrc1));
      _mm256_storeu_si256((__m256i*)(pdst + 16), mSum1);
      pdst += 32;
      psrc += 32;
    }
  }
}

void ReduceColSum4(const uint8_t* src, uint16_t* dst, uint16_t height,
                   uint16_t width) {
  assert(width % 64 == 0);
  assert(height <= 256);
  __m128i mSrc, mSrc1, mSrc2, mSrc3;
  __m256i mSum, mSum1, mSum2, mSum3;
  const uint8_t* psrc = src;
  const int nCount = (width / 64);
  for (size_t y = 0; y < height; ++y) {
    uint16_t* pdst = dst;
    for (size_t x = 0; x < nCount; ++x) {
      mSum = _mm256_loadu_si256((__m256i*)(pdst));
      mSum1 = _mm256_loadu_si256((__m256i*)(pdst+16));
      mSum2 = _mm256_loadu_si256((__m256i*)(pdst+32));
      mSum3 = _mm256_loadu_si256((__m256i*)(pdst+48));

      mSrc = _mm_loadu_si128((const __m128i*)(pdst));
      mSrc1 = _mm_loadu_si128((const __m128i*)(pdst+16));
      mSrc2 = _mm_loadu_si128((const __m128i*)(pdst+32));
      mSrc3 = _mm_loadu_si128((const __m128i*)(pdst+48));

      mSum = _mm256_adds_epu16(mSum, _mm256_cvtepu8_epi16(mSrc));
      _mm256_storeu_si256((__m256i*)(pdst), mSum);

      mSum1 = _mm256_adds_epu16(mSum1, _mm256_cvtepu8_epi16(mSrc1));
      _mm256_storeu_si256((__m256i*)(pdst + 16), mSum1);

      mSum2 = _mm256_adds_epu16(mSum2, _mm256_cvtepu8_epi16(mSrc2));
      _mm256_storeu_si256((__m256i*)(pdst + 32), mSum2);

      mSum3 = _mm256_adds_epu16(mSum3, _mm256_cvtepu8_epi16(mSrc3));
      _mm256_storeu_si256((__m256i*)(pdst + 48), mSum3);
      pdst += 64;
      psrc += 64;
    }
  }
}

总结

除了本文之外，我实际进行了更多测试，总体由于不同指令的测试效果并不稳定，在本文并未体现。总体上可以看到使用AVX指令集并行只不过是第一步，尤其是对于简单的图像计算也是最简单的一步，全部依赖编译器优化也是不合理，比如使用的MSVC并不会使用循环展开。只有深入理解计算机系统，理解汇编代码、清楚计算机工作原理，才能更好的优化代码执行。