模乘的向量化

Posted 2023-02-16

【中文标题】模乘的向量化

【英文标题】：Vectorization of modulo multiplication

【发布时间】：2017-10-17 12:33:45

【问题描述】：

我有一个函数：

void Func(const int * a, const int * b, size_t size, int p, int * c)

    for (size_t i = 0; i < size; ++i)
        c[i] = (a[i]*b[i])%p;

此函数对整数数组执行多次模乘。 所有整数都是正数。 我需要提高它的性能。

我想到了 SSE 和 AVX。但是他们没有向量化模乘法的操作。 还是我错了？

也许有人知道解决这个问题的任何可能性吗？

【问题讨论】：

它是一个素数。

您是否重复使用同一个p？ size 通常有多大？如果它足够大或使用相同的p 重复足够频繁，则甚至可能值得使用硬编码向量移位和the fixed-point multiplicative inverse 来JIT 编译循环。或者使用libdivide.com 在没有 JIT 的情况下使用乘法逆运算，但这会产生更多开销（psrlq 最好使用 imm8 计数）。不过，它可能只有 SSE2 版本，而不是 AVX2 或 AVX512。

a[i]*b[i] 是否会溢出？如果是，那可以吗，还是您想要 64 位结果 mod p 的结果？

@chtz：我认为我们可以假设这是现有的工作+测试源，其中带符号的溢出将是 UB（或者实际上会以 32 位包装）。

【参考方案1】：

首先我要说明的是，模运算可以通过使用浮点数来实现：

d % p = d - int(float(d)/float(p))*p.

虽然右侧部分的操作量比左侧大，但此部分更可取，因为它可以使用 SSE/AVX 进行矢量化。

32x32 => 32-bit integer multiplication 的 SSE4.1 实现。请注意，从 FP 转换回整数是通过四舍五入完成的；如果您想要 C 浮点->整数转换等语义，请使用向零截断 (cvttps_epi32)。

void Func(const int * a, const int * b, size_t size, int p, int * c)

    __m128 _k = _mm_set1_ps(1.0f / p);
    __m128i _p = _mm_set1_epi32(p);
    for (size_t i = 0; i < size; i += 4)
    
        __m128i _a = _mm_loadu_si128((__m128i*)(a + i));
        __m128i _b = _mm_loadu_si128((__m128i*)(b + i));
        __m128i _d = _mm_mullo_epi32(_a, _b);
        __m128i _e = _mm_cvtps_epi32(_mm_mul_ps(_mm_cvtepi32_ps(_d), _k)); // e = int(float(d)/float(p));
        __m128i _c = _mm_sub_epi32(_d, _mm_mullo_epi32(_e, _p));
        _mm_storeu_si128((__m128i*)(c + i), _c);
                

使用 AVX 的实现：

void Func(const int * a, const int * b, size_t size, int p, int * c)

    __m256 _k = _mm256_set1_ps(1.0f / p);
    __m256i _p = _mm256_set1_epi32(p);
    for (size_t i = 0; i < size; i += 8)
    
        __m256i _a = _mm256_loadu_si128((__m256i*)(a + i));
        __m256i _b = _mm256_loadu_si128((__m256i*)(b + i));
        __m256i _d = _mm256_mullo_epi32(_a, _b);
        __m256i _e = _mm256_cvtps_epi32(_mm256_mul_ps(_mm256_cvtepi32_ps(_d), _k)); // e = int(float(d)/float(p));
        __m256i _c = _mm256_sub_epi32(_d, _mm256_mullo_epi32(_e, _p));
        _mm256_storeu_si128((__m256i*)(c + i), _c);
                

【讨论】：

在 pmulld 为 2 微秒的 CPU（例如 HSW 和 SKL）上，可能值得在乘法之前将整数输入转换为浮点数。但是cvtdq2ps 在添加单元上运行，因此它本身的吞吐量有限。实际上，它可能在 HSW/SKL 上实现收支平衡。

@ErmIg 出于好奇，您是否可能对不同的实现进行了一些基准测试，如果是，它们如何比较？

不幸的是我刚刚写了这段代码。我还没有检查它的性能。

你的意思是_mm_mullo_epi32(_e, _p)而不是`_mm_mul_ps(_e, _p)? Also, I guess you will have some wrong results for big a, b, p`，因为你在转换为浮点数时会丢失8位（我没有实际测试你的代码）

【参考方案2】：

其实有一个内在函数在执行这个操作： _mm256_irem_epi32

https://software.intel.com/sites/landingpage/IntrinsicsGuide/#text=_mm256_irem_epi32

【讨论】：

这是一个英特尔 SVML 库函数，而不是一条机器指令。此外，对于常量p，您需要通过一次找到乘法逆元并重用它来进行优化。 （喜欢libdivide.com）