AArch64: folly::ConcurrentHashMap codegen could be better / LLVM is conservative around atomics

[MatzeB]

We noticed codegen in the context of folly::ConcurrentHashMap (aka code with a bunch of std::atomic) could be better. I am told this a reduction of an important function, that a colleague tried to vectorize:

#include <arm_neon.h>
#include <atomic>

struct __attribute__((packed)) mystruct {
    std::atomic<uint64_t> low_;
    std::atomic<uint64_t> hi_;
};

uint64_t occupiedMask(mystruct& tags_, uint64_t kFullMask) {
    uint64x2_t vec;
    vec[0] = tags_.low_.load(std::memory_order_relaxed);
    vec[1] = tags_.hi_.load(std::memory_order_relaxed);
    // signed shift extends top bit to all bits
    auto occupiedV =
        vreinterpretq_u8_s8(vshrq_n_s8(vreinterpretq_s8_u64(vec), 7));
    uint8x8_t maskV = vshrn_n_u16(vreinterpretq_u16_u8(occupiedV), 4);
    return vget_lane_u64(vreinterpret_u64_u8(maskV), 0) & kFullMask;
}

Currently produces:

occupiedMask(mystruct&, unsigned long):
        ldr     x8, [x0]
        ldr     x9, [x0, #8]
        fmov    d0, x8
        mov     v0.d[1], x9
        cmlt    v0.16b, v0.16b, #0
        shrn    v0.8b, v0.8h, #4
        fmov    x8, d0
        and     x0, x8, x1
        ret

(godbolt equivalent: https://godbolt.org/z/xobWMhe7W )

but we think this could ideally be (equavalent to the code you currently get when removing the atomics):

        ldr     q0, [x0]
        cmlt    v0.16b, v0.16b, #0
        shrn    v0.8b, v0.8h, #4
        fmov    x8, d0
        and     x0, x8, x1
        ret

[MatzeB]

Some aspects of the whole "optimizing atomics" story:

• The example at hand likely needs the SLP vectorizer to combine the two atomic loads.
• AArch64 codegen has a bunch of vector patterns specified for load<> which would arguably work equally well for atomic_load<>. I experimented with a prototype like this: https://gist.github.com/MatzeB/c39b1d7c734d3d0c68e630ba57cb995d
• AArch64LoadStoreOptimizer is not touching atomic loads either but could arguably form LDP.

[llvmbot]

@llvm/issue-subscribers-backend-aarch64

Author: Matthias Braun (MatzeB)

We noticed codegen in the context of folly::ConcurrentHashMap (aka code with a bunch of std::atomic) could be better. I am told this a reduction of an important function, that a colleague tried to vectorize:

#include <arm_neon.h>
#include <atomic>

struct __attribute__((packed)) mystruct {
    std::atomic<uint64_t> low_;
    std::atomic<uint64_t> hi_;
};

uint64_t occupiedMask(mystruct& tags_, uint64_t kFullMask) {
    uint64x2_t vec;
    vec[0] = tags_.low_.load(std::memory_order_relaxed);
    vec[1] = tags_.hi_.load(std::memory_order_relaxed);
    // signed shift extends top bit to all bits
    auto occupiedV =
        vreinterpretq_u8_s8(vshrq_n_s8(vreinterpretq_s8_u64(vec), 7));
    uint8x8_t maskV = vshrn_n_u16(vreinterpretq_u16_u8(occupiedV), 4);
    return vget_lane_u64(vreinterpret_u64_u8(maskV), 0) & kFullMask;
}

Currently produces:

occupiedMask(mystruct&, unsigned long):
        ldr     x8, [x0]
        ldr     x9, [x0, #<!-- -->8]
        fmov    d0, x8
        mov     v0.d[1], x9
        cmlt    v0.16b, v0.16b, #<!-- -->0
        shrn    v0.8b, v0.8h, #<!-- -->4
        fmov    x8, d0
        and     x0, x8, x1
        ret

(godbolt equivalent: https://godbolt.org/z/xobWMhe7W )

but we think this could ideally be:

        ldr     q0, [x0]
        cmlt    v0.16b, v0.16b, #<!-- -->0
        shrn    v0.8b, v0.8h, #<!-- -->4
        fmov    x8, d0
        and     x0, x8, x1
        ret

currently you only get this code when dropping the atomics.

[MatzeB]

This is admittedly a pretty broad issue. But thought it may be okay to start a discussion. CC @tmatheson-arm

[palmer-dabbelt]

Thanks for writing this up, it's where #160580 came from.

I also started poking around with SLP over here https://github.com/palmer-dabbelt/llvm-project/tree/slp-monotonic, which doesn't work yet...

[efriedma-quic]

Arm ARM (B2.2.1) says:

Element or Structure Reads to SIMD and floating-point registers of 64-bit or smaller elements, where each
element is aligned to the size of the element being loaded, have each element treated as a single-copy atomic
read.
Reads to SIMD and floating-point registers of a 128-bit value that is 64-bit aligned in memory are treated as
a pair of single-copy atomic 64-bit reads.

That said, this is a pretty narrow guarantee; it's target-specific, and it only applies to the specific instruction encodings mentioned, with a pointer with appropriate alignment. So if we want to optimize this, we probably need a pseudo-instruction/intrinsic to represent this exact operation, so we can properly preserve atomicity.

[MatzeB]

yeah would need a notion of "elementwise atomic" for a vector load which generic LLVM-IR cannot express...

[Nicoshev]

@MatzeB Should the discussion regarding gp ldp be treated separately? It may be more straghitforward