我可以使用AVX FMA单元进行精确的52位整数乘法吗？

Question

AXV2没有任何整数乘法,其源大于32位.它提供32 x 32 - > 32乘法,以及32 x 32 - > 64乘以¹,但没有64位源.

假设我需要一个输入大于32位但小于或等于52位的无符号乘法 - 我可以简单地使用浮点DP乘法或FMA指令,并且当整数输入和输出时输出将是位精确的结果可以用52或更少的比特表示(即,在[0,2 ^ 52-1]范围内)？

如果我想要产品的所有104位更一般的情况怎么样？或整数乘积超过52位的情况(即,产品在位索引中的非零值> 52) - 但我只想要低52位？在后一种情况下,它MUL会给我更高的位并舍去一些低位(也许这就是IFMA帮助的？).

编辑:事实上,根据这个答案,也许它可以做任何高达2 ^ 53的事情- 我忘记了1在尾数之前隐含的领先有效地给了你一点.

---

¹有趣的是,正如Mysticial 在评论中所解释的那样,64位产品PMULDQ操作的延迟是32位PMULLD版本的一半,吞吐量是32位版本的两倍.

Answer 1

是的,这是可能的.但是从AVX2开始,它不太可能比使用MULX/ADCX/ADOX的标量方法更好.

对于不同的输入/输出域,这种方法实际上有无限多种变体.我只会覆盖其中的3个,但一旦你知道它们是如何工作的,它们很容易概括.

免责声明:

• 这里的所有解决方案都假设舍入模式是圆形到均匀的.
• 建议不要使用快速数学优化标志,因为这些解决方案依赖于严格的IEEE.

在范围签署双:[-2 ⁵¹,2 ⁵¹ ]

//  A*B = L + H*2^52
//  Input:  A and B are in the range [-2^51, 2^51]
//  Output: L and H are in the range [-2^51, 2^51]
void mul52_signed(__m256d& L, __m256d& H, __m256d A, __m256d B){
    const __m256d ROUND = _mm256_set1_pd(30423614405477505635920876929024.);    //  3 * 2^103
    const __m256d SCALE = _mm256_set1_pd(1. / 4503599627370496);                //  1 / 2^52

    //  Multiply and add normalization constant. This forces the multiply
    //  to be rounded to the correct number of bits.
    H = _mm256_fmadd_pd(A, B, ROUND);

    //  Undo the normalization.
    H = _mm256_sub_pd(H, ROUND);

    //  Recover the bottom half of the product.
    L = _mm256_fmsub_pd(A, B, H);

    //  Correct the scaling of H.
    H = _mm256_mul_pd(H, SCALE);
}

这是最简单的一种,也是唯一一种与标量方法竞争的产品.最终缩放是可选的,具体取决于您要对输出执行的操作.所以这只能被认为是3条指令.但它也是最不实用的,因为输入和输出都是浮点值.

两个FMA保持融合绝对至关重要.这就是快速数学优化可以破坏事物的地方.如果第一个FMA被分解,则L不再保证在该范围内[-2^51, 2^51].如果第二个FMA被打破,L将是完全错误的.

在范围带符号的整数:[-2 ⁵¹,2 ⁵¹ ]

//  A*B = L + H*2^52
//  Input:  A and B are in the range [-2^51, 2^51]
//  Output: L and H are in the range [-2^51, 2^51]
void mul52_signed(__m256i& L, __m256i& H, __m256i A, __m256i B){
    const __m256d CONVERT_U = _mm256_set1_pd(6755399441055744);     //  3*2^51
    const __m256d CONVERT_D = _mm256_set1_pd(1.5);

    __m256d l, h, a, b;

    //  Convert to double
    A = _mm256_add_epi64(A, _mm256_castpd_si256(CONVERT_U));
    B = _mm256_add_epi64(B, _mm256_castpd_si256(CONVERT_D));
    a = _mm256_sub_pd(_mm256_castsi256_pd(A), CONVERT_U);
    b = _mm256_sub_pd(_mm256_castsi256_pd(B), CONVERT_D);

    //  Get top half. Convert H to int64.
    h = _mm256_fmadd_pd(a, b, CONVERT_U);
    H = _mm256_sub_epi64(_mm256_castpd_si256(h), _mm256_castpd_si256(CONVERT_U));

    //  Undo the normalization.
    h = _mm256_sub_pd(h, CONVERT_U);

    //  Recover bottom half.
    l = _mm256_fmsub_pd(a, b, h);

    //  Convert L to int64
    l = _mm256_add_pd(l, CONVERT_D);
    L = _mm256_sub_epi64(_mm256_castpd_si256(l), _mm256_castpd_si256(CONVERT_D));
}

在第一个例子的基础上,我们将它与快速double <-> int64转换技巧的通用版本结合起来.

这个更有用,因为你正在使用整数.但即使使用快速转换技巧,大部分时间都将用于转换.幸运的是,如果多次乘以相同的操作数,则可以消除一些输入转换.

范围内的无符号整数:[0,2 ⁵²)

//  A*B = L + H*2^52
//  Input:  A and B are in the range [0, 2^52)
//  Output: L and H are in the range [0, 2^52)
void mul52_unsigned(__m256i& L, __m256i& H, __m256i A, __m256i B){
    const __m256d CONVERT_U = _mm256_set1_pd(4503599627370496);     //  2^52
    const __m256d CONVERT_D = _mm256_set1_pd(1);
    const __m256d CONVERT_S = _mm256_set1_pd(1.5);

    __m256d l, h, a, b;

    //  Convert to double
    A = _mm256_or_si256(A, _mm256_castpd_si256(CONVERT_U));
    B = _mm256_or_si256(B, _mm256_castpd_si256(CONVERT_D));
    a = _mm256_sub_pd(_mm256_castsi256_pd(A), CONVERT_U);
    b = _mm256_sub_pd(_mm256_castsi256_pd(B), CONVERT_D);

    //  Get top half. Convert H to int64.
    h = _mm256_fmadd_pd(a, b, CONVERT_U);
    H = _mm256_xor_si256(_mm256_castpd_si256(h), _mm256_castpd_si256(CONVERT_U));

    //  Undo the normalization.
    h = _mm256_sub_pd(h, CONVERT_U);

    //  Recover bottom half.
    l = _mm256_fmsub_pd(a, b, h);

    //  Convert L to int64
    l = _mm256_add_pd(l, CONVERT_S);
    L = _mm256_sub_epi64(_mm256_castpd_si256(l), _mm256_castpd_si256(CONVERT_S));

    //  Make Correction
    H = _mm256_sub_epi64(H, _mm256_srli_epi64(L, 63));
    L = _mm256_and_si256(L, _mm256_set1_epi64x(0x000fffffffffffff));
}

最后,我们得到了原始问题的答案.这通过调整转换并添加校正步骤来构建有符号整数解.

但在这一点上,我们有13条指令 - 其中一半是高延迟指令,不包括众多FP <-> int旁路延迟.所以这不可能赢得任何基准.相比之下,64 x 64 -> 128-bitSIMD乘法可以在16条指令中完成(如果您预处理输入则为14条).

如果舍入模式是向下舍入或舍入为零,则可以省略校正步骤.唯一重要的指示是h = _mm256_fmadd_pd(a, b, CONVERT_U);.因此,在AVX512上,您可以覆盖该指令的舍入并单独保留舍入模式.

最后的想法:

值得注意的是,通过调整魔术常数可以减少2 ⁵²的操作范围.这可能对第一个解决方案(浮点数)有用,因为它为您提供额外的尾数用于浮点累积.这使您可以避免在最后2个解决方案中不断地在int64和double之间来回转换的需要.

虽然这里的3个示例不太可能比标量方法更好,但AVX512几乎肯定会取得平衡.Knights Landing尤其是ADCX和ADOX的吞吐量很低.

当然,当AVX512-IFMA问世时,所有这一切都没有实际意义.这会将完整52 x 52 -> 104-bit产品减少为2条指令,并免费提供积累.