[eendebakpt]

v_lshift() and v_rshift() pass a carry from one loop iteration to the next, which prevents vectorization. The carry into a digit is just the bits shifted out of its neighbour, so each output digit can be built from two input digits instead, with nothing carried.

These helpers do the normalization in x_divrem() (three passes that dominate
when the quotient is short) and the scaling in true division.

Benchmark	main	PR
a << 7 (10 digits)	35.2 ns	31.9 ns: 1.10x faster
a << 7 (100 digits)	87.3 ns	42.8 ns: 2.04x faster
a >> 7 (100 digits)	90.2 ns	42.4 ns: 2.13x faster
a << 7 (1000 digits)	639 ns	213 ns: 2.99x faster
a >> 7 (1000 digits)	631 ns	209 ns: 3.02x faster
10 // 2 digits	160 ns	149 ns: 1.08x faster
1000 // 2 digits	14.7 us	13.9 us: 1.05x faster
100 // 99 digits	535 ns	349 ns: 1.54x faster
1000 // 999 digits	4.51 us	2.95 us: 1.53x faster
1000 % 999 digits	4.53 us	2.94 us: 1.54x faster
divmod, 1000 by 999 digits	4.55 us	2.98 us: 1.52x faster
1000 / 999 digits	5.22 us	3.18 us: 1.64x faster
pyperformance pidigits	161 ms	129 ms: 1.24x faster

Benchmark script

import pyperf

runner = pyperf.Runner()
def ints(na, nb=1):
    return (f"import random; r = random.Random(1234); "
            f"a = r.getrandbits({na*30-1}) | (1 << {na*30-2}); "
            f"b = r.getrandbits({nb*30-3}) | (1 << {nb*30-4})")
for na in (2, 10, 100, 1000):
    runner.timeit(f"a << 7 ({na} digits)", "a << 7", setup=ints(na))
    runner.timeit(f"a >> 7 ({na} digits)", "a >> 7", setup=ints(na))
runner.timeit("-a >> 7 (1000 digits)", "a >> 7", setup=ints(1000) + "; a = -a")
for na, nb in [(2, 2), (10, 2), (1000, 2), (10, 9), (100, 99), (1000, 999)]:
    runner.timeit(f"{na} // {nb} digits", "a // b", setup=ints(na, nb))
runner.timeit("1000 % 999 digits", "a % b", setup=ints(1000, 999))
runner.timeit("divmod, 1000 by 999 digits", "divmod(a, b)", setup=ints(1000, 999))
runner.timeit("1000 / 999 digits", "a / b", setup=ints(1000, 999))

Note: when disabling vectorization I get a 20% slowdown. I think that is ok, since modern cpu's/compilers can do this, if not we could modify the PR to gate the change behind some macro.

Generated with Claude Code

• Issue: Improve performance of pidigits and int arithmetic #157742

[kumaraditya303]

Just curious, have you checked the assembly code to see which SIMD instructions are getting used? Maybe we can do even better with some more rearrangement?

[eendebakpt]

Just curious, have you checked the assembly code to see which SIMD instructions are getting used? Maybe we can do even better with some more rearrangement?

I did not check this. From my AI: with GCC 14 at the default -O3 on x86-64 it uses baseline SSE2, 4 digits per iteration. The left-shift loop (inlined into long_lshift1, x_divrem, long_true_divide and _PyLong_Frexp) is:

movdqu  a[i-1 ...], xmm0      ; 4 digits, offset by one
movdqu  a[i   ...], xmm1      ; 4 digits
psrld   xmm4, xmm0            ; >> (PyLong_SHIFT - d)
pslld   xmm3, xmm1            ; << d
por     xmm1, xmm0
pand    xmm2, xmm0            ; & PyLong_MASK
movups  xmm0, z[i ...]
add / cmp / jne

The loop is versioned with a runtime overlap check, the tail is done with one 2-digit vector step and then scalar code. That is 10 instructions per 4 digits, where the old scalar loop needed about 6 per digit.

[picnixz]

There is a similar issue/PR for aarch64 with similar approches (using double digits) IIRC

[eendebakpt]

There is a similar issue/PR for aarch64 with similar approches (using double digits) IIRC

I think that is #156443. It touches the same code, but is different,

[markshannon]

Impressive speedups.

How common are numbers with only a few digits, say 3 to 6? (not 1 or 2 as those are "compact")
Is it worth branching on the number of digits, or would that add too much complexity?
I assuming that vectorization only pays off for more than a few digits.

[eendebakpt]

How common are numbers with only a few digits, say 3 to 6? (not 1 or 2 as those are "compact") Is it worth branching on the number of digits, or would that add too much complexity? I assuming that vectorization only pays off for more than a few digits.

Not sure how common the 4+ digits are, but I suspect the 3 digits (needed for 64-bit integers) are used more often (but nearly as often as the small ints, 1-digit or 2-digit ints though).

Adding a branch is not worth it, as the compiler naturally adds branches as part of the vectorization (both gcc and clang, assembly and benchmarks in the details below). This indeed implies that the vectorization starts to pay off for larger number of digits m.

Details

Benchmarks for small number of digits (PGO):

Benchmark	main	PR
a << 7 (1 digits)	20.0 ns	20.0 ns: not significant
a >> 7 (1 digits)	17.0 ns	17.1 ns: 1.01x slower
a << 7 (2 digits)	22.0 ns	22.4 ns: not significant
a >> 7 (2 digits)	22.3 ns	22.2 ns: not significant
a << 7 (3 digits)	22.2 ns	22.7 ns: 1.03x slower
a >> 7 (3 digits)	24.2 ns	23.9 ns: not significant
a << 7 (4 digits)	22.8 ns	22.8 ns: not significant
a >> 7 (4 digits)	24.6 ns	24.5 ns: not significant
2 // 2 digits	42.8 ns	42.5 ns: 1.01x faster
3 // 2 digits	50.8 ns	50.0 ns: 1.02x faster
4 // 2 digits	60.2 ns	58.8 ns: 1.02x faster
2 / 2 digits	69.8 ns	71.1 ns: 1.02x slower
3 / 2 digits	68.6 ns	69.2 ns: 1.01x slower
4 / 2 digits	68.5 ns	69.5 ns: 1.02x slower
float(a) (2 digits)	29.1 ns	29.7 ns: 1.02x slower
float(a) (3 digits)	29.1 ns	29.7 ns: 1.02x slower
float(a) (4 digits)	29.1 ns	29.7 ns: 1.02x slower

This is in the noise, and also influenced by binary layout noise. The compiler already generates the branch on the shift m. GCC 14, -O3, the vectorized v_lshift():

v_lshift:
        mov     eax, [rsi]            ; z[0] = (a[0] << d) & MASK
        mov     r10d, 30
        sub     r10d, ecx             ; e = 30 - d
        sal     eax, cl
        and     eax, 0x3fffffff
        mov     [rdi], eax
        cmp     rdx, 1
        jle     .done                 ; m == 1: nothing else to do
        cmp     rdx, 2
        je      .scalar               ; m == 2: scalar loop
        mov     rax, rdi
        sub     rax, rsi
        cmp     rax, 12
        jbe     .scalar               ; z and a overlap: scalar loop
        lea     rax, [rdx-2]
        cmp     rax, 2
        jbe     .tail                 ; m <= 4: one 2-wide step + scalar
        ...                           ; set up xmm2 = MASK, xmm3 = e, xmm4 = d
.vector:                              ; m >= 5: 4 digits per iteration
        movdqu  xmm0, [rsi+rax]       ; a[i..i+3]
        movdqu  xmm1, [rsi+rax-4]     ; a[i-1..i+2]
        pslld   xmm0, xmm4            ; a[i] << d
        psrld   xmm1, xmm3            ; a[i-1] >> e
        por     xmm0, xmm1
        pand    xmm0, xmm2            ; & MASK
        movups  [rdi+rax], xmm0       ; z[i..i+3]
        add     rax, 16
        cmp     rdx, rax
        jne     .vector
        ...                           ; 2-wide tail (movq/pslld/psrld/por/pand), then scalar tail
.scalar:                              ; same body, one digit per iteration
        mov     eax, [rsi+rdx*4-4]
        mov     r11d, [rsi+rdx*4]
        mov     ecx, r10d
        shr     eax, cl
        mov     ecx, r9d
        sal     r11d, cl
        or      eax, r11d
        and     eax, 0x3fffffff
        mov     [rdi+rdx*4], eax
        add     rdx, 1
        cmp     r8, rdx
        jne     .scalar
.done:
        mov     eax, [rsi+r8*4-4]     ; return a[m-1] >> e
        mov     ecx, r10d
        shr     eax, cl
        ret

So m == 1 costs one digit, m == 2 one scalar iteration, m = 3, 4 one 2-wide step and possibly a scalar step, and only m >= 5 enters the 4-wide loop (10 instructions per 4 digits). For comparison, the carry loop on main is 11 instructions per digit with the carry on the critical path:

.loop:                                ; main, one digit per iteration
        mov     eax, [r9+rsi*4]       ; a[i]
        sal     rax, cl               ; (twodigits)a[i] << d
        or      rax, rdi              ; | carry
        mov     edi, eax
        shr     rax, 30               ; carry = acc >> 30
        and     edi, 0x3fffffff
        mov     [r8+rsi*4], edi       ; z[i] = acc & MASK
        add     rsi, 1
        mov     edi, eax              ; carry -> next iteration
        cmp     rdx, rsi
        jne     .loop

Clang does the same with different thresholds: the vector loop handles 8 digits per iteration (two SSE2 vectors, same 6 instructions per vector), it is entered for m >= 9, and the scalar path is unrolled by two. The PGO build of the PR additionally unrolls GCC's vector loop by four.

Clang 21 (-O3) vectorizes the same loop with two SSE2 vectors per iteration, entered for m >= 9; below that it runs a scalar loop unrolled by two:

.LBB2_4:                              ; clang: 8 digits per iteration
        movdqu  xmm3, [rsi+r11*4]     ; a[i-1..i+2]
        movdqu  xmm4, [rsi+r11*4+4]   ; a[i..i+3]
        movdqu  xmm5, [rsi+r11*4+16]  ; a[i+3..i+6]
        movdqu  xmm6, [rsi+r11*4+20]  ; a[i+4..i+7]
        pslld   xmm4, xmm1            ; a[i] << d
        pslld   xmm6, xmm1
        psrld   xmm3, xmm0            ; a[i-1] >> e
        psrld   xmm5, xmm0
        por     xmm3, xmm4
        por     xmm5, xmm6
        pand    xmm3, xmm2            ; & MASK
        pand    xmm5, xmm2
        movdqu  [rdi+r11*4+4], xmm3   ; z[i..i+3]
        movdqu  [rdi+r11*4+20], xmm5  ; z[i+4..i+7]
        add     r11, 8
        cmp     r10, r11
        jne     .LBB2_4

In the PGO+LTO build clang inlines both helpers into every caller (x_divrem, _PyLong_Frexp, long_true_divide, the shift methods), GCC inlines them into x_divrem and the shift methods and keeps one out-of-line copy for _PyLong_Frexp and long_true_divide.

[XiaohongGong]

Hi,

I tested the performance on my AArch64 and x86_64 machines. Here is a summary of the results:

1. The pidigits benchmark improves by approximately 25% across all of my test machines, on both AArch64 and x86_64.
2. For the Python microbenchmarks in the comment above, I observed speedups ranging from 1.10× to 3× on both architectures for larger inputs (>4 digits). I also confirmed that these cases are vectorized. For smaller inputs (<4 digits), however, I observed minor regressions of less than 5%. I’m not sure how significant these regressions are in practice, particularly given the observation above that shifting a 64-bit integer typically involves three digits.

Overall, I think this PR is a positive step toward improving the performance of long objects shifts/division. I can also inspect the generated code to identify any further optimization opportunities.

Best Regards,
Xiaohong

[XiaohongGong]

Hi,

I tested the performance on my AArch64 and x86_64 machines. Here is a summary of the results:

1. The pidigits benchmark improves by approximately 25% across all of my test machines, on both AArch64 and x86_64.
2. For the Python microbenchmarks in the comment above, I observed speedups ranging from 1.10× to 3× on both architectures for larger inputs (>4 digits). I also confirmed that these cases are vectorized. For smaller inputs (<4 digits), however, I observed minor regressions of less than 5%. I’m not sure how significant these regressions are in practice, particularly given the observation above that shifting a 64-bit integer typically involves three digits.

Overall, I think this PR is a positive step toward improving the performance of long objects shifts/division. I can also inspect the generated code to identify any further optimization opportunities.

Best Regards, Xiaohong

I recently took a closer look at this change, focusing on the possible causes of the regressions, particularly for smaller m values. A few potential influences are:

1. The additional memory load in each iteration may be relatively costly on AArch64. Since the same memory is accessed frequently, I would expect most accesses to hit in the cache, keeping the latency low. However, compared with the previous implementation, the extra load could still increase pressure on the load pipelines.

2. Auto-vectorization from compiler may introduce setup overhead, such as runtime memory-alias checks and other pre-loop checks. These costs could be more noticeable for smaller m values, where there are fewer iterations over which to amortize them.

3. The additional shift in the loop body may also contribute to the regressions on some X86 systems. On some X86 systems, there is not a direct instruction to shift a runtime variable, the shift count must be saved to a special register cl. For both shifts in (a[i] << d) | (a[i-1] >> e), switching between the counts could add instructions and register dependencies.

   Here is the code generated for a[i] << d) | (a[i-1] >> e on one my local X86 machine (AMD EPYC 9124（Zen 4 / Genoa）):

mov    -0x4(%rsi,%rax,4),%edx   // a[i-1]
mov    (%rsi,%rax,4),%r8d       // a[i]
mov    %r9d,%ecx
shr    %cl,%edx            // a[i - 1] >> e
mov    %r10d,%ecx
shl    %cl,%r8d            // a[i] << d
or     %r8d,%edx

Due to above consideration, I'm wondering whether we should special handling the cases for smaller m. For example, unrolling the loop totally for m=1,2,3,4. The code for v_lshift may look like:

static digit
v_lshift(digit *z, digit *a, Py_ssize_t m, int d)
{
    Py_ssize_t i;
    int e = PyLong_SHIFT - d;

    assert(0 <= d && d < PyLong_SHIFT);
    assert(m > 0);

    z[0] = (a[0] << d) & PyLong_MASK;
    if (m == 1) {
      return a[0] >> e;
    }

    if (m == 2) {
      z[1] = ((a[1] << d) | (a[0] >> e)) & PyLong_MASK;
      return a[1] >> e;
    }

    if (m == 3) {
      z[1] = ((a[1] << d) | (a[0] >> e)) & PyLong_MASK;
      z[2] = ((a[2] << d) | (a[1] >> e)) & PyLong_MASK;
      return a[2] >> e;
    }

    if (m == 4) {
      z[1] = ((a[1] << d) | (a[0] >> e)) & PyLong_MASK;
      z[2] = ((a[2] << d) | (a[1] >> e)) & PyLong_MASK;
      z[3] = ((a[3] << d) | (a[2] >> e)) & PyLong_MASK;
      return a[3] >> e;
    }

    // m >= 4
    // Build each digit from two digits of a, rather than passing a carry
    // from one iteration to the next, so that the loop can be vectorized.
    for (i = 1; i < m; i++) {
        z[i] = ((a[i] << d) | (a[i-1] >> e)) & PyLong_MASK;
    }
    return a[m - 1] >> e;
}

The code is more straightforward for smaller digits values:

1. It could avoid loop-control overhead and any setup overhead associated with auto-vectorization.
2. It could also enable further compiler optimizations, such as better instruction scheduling and combining memory loads, potentially improving IPC.

I tested this version with a smaller C case, and the performance can improve >10% on both aarch64 and x86 systems. Would this be a worthwhile optimization, or would the added code complexity outweigh the benefits?