[RISCV] Inefficient code generated for byte atomicrmw xchg when the input is a constant zero

[Amanieu]

Noticed while working on rust-lang/rust#114034. When the replacement value is known to be zero then we should lower to an amoand instead of a LL/SC loop.

IR

define zeroext i8 @swap_zero(ptr %0) {
  %2 = atomicrmw xchg ptr %0, i8 0 monotonic
  ret i8 %2
}

Current output

swap_zero:                              # @swap_zero
        andi    a1, a0, -4
        slli    a0, a0, 3
        li      a2, 255
        sllw    a2, a2, a0
.LBB0_1:                                # =>This Inner Loop Header: Depth=1
        lr.w    a3, (a1)
        li      a4, 0
        xor     a4, a3, a4
        and     a4, a4, a2
        xor     a4, a3, a4
        sc.w    a4, a4, (a1)
        bnez    a4, .LBB0_1
        srlw    a0, a3, a0
        andi    a0, a0, 255
        ret

Ideal output

swap_zero:                              # @swap_zero
        andi    a1, a0, -4
        slli    a0, a0, 3
        li      a2, 255
        sllw    a2, a2, a0
        not     a2, a2
        amoand  a1, a2, (a1)
        srlw    a0, a1, a0
        andi    a0, a0, 255
        ret

[llvmbot]

@llvm/issue-subscribers-backend-risc-v

[riking]

If you black_box the false, we get a single byte stack spill followed by nearly the desired codegen: https://rust.godbolt.org/z/dWPPbqEhP

[topperc]

InstCombine aggressively turns atomicrmw and with 0 into atomicrmw xchg. Maybe AtomicExpandPass needs to reverse this for partial atomics based on the target.

[appujee]

Do you have llc or opt command that will repro the above assembly?

From https://godbolt.org/z/54r8h56cq

llc generated __atomic_exchange_1 instead.

[topperc]

Do you have llc or opt command that will repro the above assembly?

From https://godbolt.org/z/54r8h56cq

llc generated __atomic_exchange_1 instead.

-mattr=+a

[appujee]

Will generating amoswap followed by andi reg, reg, 255 have the same semantics?

[topperc]

Will generating amoswap followed by andi reg, reg, 255 have the same semantics?

No. We need to write 0 to a single byte. amoswap writes 32 or 64 bits, but we don't know what value to put in the other bytes until after the read. Using amoand also writes 32 or 64 bits, but would allow to mask out the single byte that needs to be zeroed.

[appujee]

ah, i see there is only amoswap.w and amoswap.d versions

[appujee]

Just confirming that we do generate amoswap for 32 bits. https://godbolt.org/z/8MxMWe5aE