Julia: >>(::Int32, ::Int32) not inlined

[dextorious]

Given a simple function that applies the right shift operator to each element of a vector

A = rand(Int32, 1000)
f!(A, n) = @inbounds for i in eachindex(A) A[i] = A[i] >> n end

it appears that there is a vast performance discrepancy between supplying an Int32 shift parameter and an Int64 one:

@benchmark f!(A, Int32(5)) # slow, about 2.5us on my machine
@benchmark f!(A, Int64(5)) # fast, about 115ns on my machine

The reason for the discrepancy becomes apparent after inspecting the output of @code_native:

julia> @code_native f!(A, Int32(5))
        .text
Filename: REPL[7]
        pushq   %rbp
        movq    %rsp, %rbp
        pushq   %r15
        pushq   %r14
        pushq   %rsi
        pushq   %rdi
        pushq   %rbx
        subq    $40, %rsp
        movl    %edx, %r14d
        movq    %rcx, %rdi
Source line: 67
        movq    24(%rdi), %rbx
Source line: 1
        testq   %rbx, %rbx
        jle     L72
        movq    (%rdi), %rcx
        xorl    %esi, %esi
        movabsq $">>", %r15
        nopl    (%rax)
L48:
        movl    (%rcx,%rsi,4), %ecx
        movl    %r14d, %edx
        callq   *%r15
        movq    (%rdi), %rcx
        movl    %eax, (%rcx,%rsi,4)
        addq    $1, %rsi
        cmpq    %rsi, %rbx
        jne     L48
L72:
        addq    $40, %rsp
        popq    %rbx
        popq    %rdi
        popq    %rsi
        popq    %r14
        popq    %r15
        popq    %rbp
        retq
        nopw    %cs:(%rax,%rax)

and, for the Int64 case:

julia> @code_native f!(A, Int64(5))
        .text
Filename: REPL[7]
        pushq   %rbp
        movq    %rsp, %rbp
Source line: 67
        pushq   %rsi
        movq    24(%rcx), %rax
Source line: 1
        testq   %rax, %rax
        jle     L719
        movq    (%rcx), %r9
        cmpq    $31, %rdx
        movl    $31, %ecx
        cmovbl  %edx, %ecx
        movq    %rdx, %r8
        negq    %r8
        cmpq    $31, %r8
        jbe     L241
        testq   %rdx, %rdx
        js      L417
        movl    $1, %r11d
        testq   %rax, %rax
        je      L208
        leaq    1(%rax), %r8
        movq    %rax, %rdx
        andq    $-32, %rdx
        leaq    1(%rdx), %r10
        movl    $1, %r11d
        cmpq    $1, %r10
        je      L199
        vmovd   %ecx, %xmm0
        vpbroadcastd    %xmm0, %ymm0
        leaq    96(%r9), %rsi
        nopw    %cs:(%rax,%rax)
L128:
        vmovdqu -96(%rsi), %ymm1
        vmovdqu -64(%rsi), %ymm2
        vmovdqu -32(%rsi), %ymm3
        vmovdqu (%rsi), %ymm4
        vpsravd %ymm0, %ymm1, %ymm1
        vpsravd %ymm0, %ymm2, %ymm2
        vpsravd %ymm0, %ymm3, %ymm3
        vpsravd %ymm0, %ymm4, %ymm4
        vmovdqu %ymm1, -96(%rsi)
        vmovdqu %ymm2, -64(%rsi)
        vmovdqu %ymm3, -32(%rsi)
        vmovdqu %ymm4, (%rsi)
        subq    $-128, %rsi
        addq    $-32, %rdx
        jne     L128
        movq    %r10, %r11
L199:
        cmpq    %r11, %r8
        je      L719
L208:
        leaq    -4(%r9,%r11,4), %rdx
        addq    $1, %rax
        subq    %r11, %rax
        nopl    (%rax)
L224:
        sarl    %cl, (%rdx)
        addq    $4, %rdx
        addq    $-1, %rax
        jne     L224
        jmp     L719
L241:
        testq   %rdx, %rdx
        js      L549
        movl    $1, %r11d
        testq   %rax, %rax
        je      L384
        leaq    1(%rax), %r8
        movq    %rax, %rsi
        andq    $-32, %rsi
        leaq    1(%rsi), %r10
        movl    $1, %r11d
        cmpq    $1, %r10
        je      L375
        vmovd   %ecx, %xmm0
        vpbroadcastd    %xmm0, %ymm0
        leaq    96(%r9), %rdx
        nopl    (%rax)
L304:
        vmovdqu -96(%rdx), %ymm1
        vmovdqu -64(%rdx), %ymm2
        vmovdqu -32(%rdx), %ymm3
        vmovdqu (%rdx), %ymm4
        vpsravd %ymm0, %ymm1, %ymm1
        vpsravd %ymm0, %ymm2, %ymm2
        vpsravd %ymm0, %ymm3, %ymm3
        vpsravd %ymm0, %ymm4, %ymm4
        vmovdqu %ymm1, -96(%rdx)
        vmovdqu %ymm2, -64(%rdx)
        vmovdqu %ymm3, -32(%rdx)
        vmovdqu %ymm4, (%rdx)
        subq    $-128, %rdx
        addq    $-32, %rsi
        jne     L304
        movq    %r10, %r11
L375:
        cmpq    %r11, %r8
        je      L719
L384:
        addq    $1, %rax
        subq    %r11, %rax
        leaq    -4(%r9,%r11,4), %rdx
        nopl    (%rax)
L400:
        sarl    %cl, (%rdx)
        addq    $4, %rdx
        addq    $-1, %rax
        jne     L400
        jmp     L719
L417:
        movl    $1, %ecx
        testq   %rax, %rax
        je      L505
        leaq    1(%rax), %r8
        movq    %rax, %rdx
        andq    $-32, %rdx
        leaq    1(%rdx), %r10
        movl    $1, %ecx
        cmpq    $1, %r10
        je      L496
        leaq    96(%r9), %rcx
        vpxor   %ymm0, %ymm0, %ymm0
        nopl    (%rax)
L464:
        vmovdqu %ymm0, -96(%rcx)
        vmovdqu %ymm0, -64(%rcx)
        vmovdqu %ymm0, -32(%rcx)
        vmovdqu %ymm0, (%rcx)
        subq    $-128, %rcx
        addq    $-32, %rdx
        jne     L464
        movq    %r10, %rcx
L496:
        cmpq    %rcx, %r8
        je      L719
L505:
        leaq    -4(%r9,%rcx,4), %rdx
        addq    $1, %rax
        subq    %rcx, %rax
        nopw    %cs:(%rax,%rax)
L528:
        movl    $0, (%rdx)
        addq    $4, %rdx
        addq    $-1, %rax
        jne     L528
        jmp     L719
L549:
        movl    $1, %ecx
        testq   %rax, %rax
        je      L684
        leaq    1(%rax), %r10
        movq    %rax, %rdx
        andq    $-32, %rdx
        leaq    1(%rdx), %r11
        movl    $1, %ecx
        cmpq    $1, %r11
        je      L679
        vmovd   %r8d, %xmm0
        vpbroadcastd    %xmm0, %ymm0
        leaq    96(%r9), %rcx
        nopl    (%rax,%rax)
L608:
        vmovdqu -96(%rcx), %ymm1
        vmovdqu -64(%rcx), %ymm2
        vmovdqu -32(%rcx), %ymm3
        vmovdqu (%rcx), %ymm4
        vpsllvd %ymm0, %ymm1, %ymm1
        vpsllvd %ymm0, %ymm2, %ymm2
        vpsllvd %ymm0, %ymm3, %ymm3
        vpsllvd %ymm0, %ymm4, %ymm4
        vmovdqu %ymm1, -96(%rcx)
        vmovdqu %ymm2, -64(%rcx)
        vmovdqu %ymm3, -32(%rcx)
        vmovdqu %ymm4, (%rcx)
        subq    $-128, %rcx
        addq    $-32, %rdx
        jne     L608
        movq    %r11, %rcx
L679:
        cmpq    %rcx, %r10
        je      L719
L684:
        leaq    -4(%r9,%rcx,4), %rdx
        addq    $1, %rax
        subq    %rcx, %rax
        nopl    (%rax,%rax)
L704:
        movb    %r8b, %cl
        shll    %cl, (%rdx)
        addq    $4, %rdx
        addq    $-1, %rax
        jne     L704
L719:
        popq    %rsi
        popq    %rbp
        vzeroupper
        retq
        nopw    %cs:(%rax,%rax)

Only the case of Int64 shift parameters is properly vectorized and unrolled, while the Int32 (and the same happens for Int16/Int8 as well) case remains unoptimized.

There's no inherent architectural reason for this that I'm aware of and I would expect any performance difference to favor the Int32 case rather than Int64s due to SIMD throughput, so I'm reporting the issue here. This testing was done on a Haswell generation quad core using Julia v0.5.1 (with a locally rebuilt system image) and should be reproducible elsewhere.

Thanks to @JeffreySarnoff and @oschulz for originally reporting this case in the Gitter chat.

[yuyichao]

This is yet another case you shouldn't use code_native. This is an inlining issue which is very obvious from code_warntype/code_llvm.

[dextorious]

Good point, thanks. Looking carefully at the output of @code_warntype does indeed make the inlining issue obvious.

Is this behavior a subset of a more general known issue or is this new? It's not obvious to me why the use of a non-Int64 parameter would confuse the inlining here.

[yuyichao]

The generic version of the function that works for ::Integer is longer in julia AST (even though it optimizes to the same instruction) so is not inlined. A more specific version is defined for ::Int likely for bootstrap purpose and is inlined.

[martinholters]

Probably worth @inlineing the ::Integer version.

[dextorious]

Okay, so the immediate naive solution is to simply do

import Base.>>
@inline function >>(x::Integer, c::Integer)                   
    typemin(Int) <= c <= typemax(Int) && return x >> (c % Int)
    (x >= 0 || c < 0) && return zero(x)                       
    oftype(x, -1)                                             
end                                                           

and then run the original code. I've verified that this eliminates the performance issue and Int32 parameters now run at the same speed as Int64 ones.

This does seem like something that should be fixed in Base, however. Not to mention that the same issue affects the left shift function, and possibly the unsigned versions as well (didn't check).