[LegalizeTypes][RISCV] Correctly sign-extend comparison for ATOMIC_CMP_XCHG

[jrtc27]

Summary

Currently, the comparison argument used for ATOMIC_CMP_XCHG is legalised
with GetPromotedInteger, which leaves the upper bits of the value
undefind. Since this is used for comparing in an LR/SC loop with a
full-width comparison, we must sign extend it. We introduce a new
getExtendForAtomicCmpSwapArg to complement getExtendForAtomicOps, since
many targets have compare-and-swap instructions (or pseudos) that
correctly handle an any-extend input, and the existing function
determines the extension of the result, whereas we are concerned with
the input.

This is related to https://reviews.llvm.org/D58829, which solved the
issue for ATOMIC_CMP_SWAP_WITH_SUCCESS, but not the simpler
ATOMIC_CMP_SWAP.

[jrtc27]

Removed atomic-rmw.ll change; nothing actually changed in there other than comments and a missing . on labels, so that churn can be postponed until it actually gets touched in a meaningful way.

[lenary]

Is there a reason this hasn't caused code changes for i16 or i8?

[jrtc27]

In D74453#1872395, @lenary wrote:

Is there a reason this hasn't caused code changes for i16 or i8?

There is no lr.b/lr.h, so they use the masked intrinsics, and hence everything is zero-extended already.

[asb]

Great fix, thanks @jrtc27. Please do go ahead and commit.

[arichardson]

As this is such a simple change and fixes a nasty code generation bug, should this be backported to the release branch?

For our CHERI version of FreeBSD, I changed the atomic.h header to use __atomic_foo instead of hand-written assembly and this causes freezes on startup without this change.
I am considering upstreaming this change to FreeBSD, so it would be nice if clang 10 had a working __atomic_compare_exchange for uint32_t.
If this is backported, I could just guard it with #if __clang_major__ < 10.

[asb]

Yes, if @jrtc27 commits I'd be supportive of a cherry-pick to the release branch.

[efriedma]

llvm/include/llvm/CodeGen/TargetLowering.h

1967

Please add an explanation for why this needs to be separate from getExtendForAtomicOps.

[luismarques]

Hey @jrtc27 . Is there anything blocking this patch being committed?

[jrtc27]

In D74453#1931264, @luismarques wrote:

Hey @jrtc27 . Is there anything blocking this patch being committed?

No, I forgot about this/was busy with other things and this got neglected. Now committed.

llvm/include/llvm/CodeGen/TargetLowering.h

1967

Added in committed version.

llvm/include/llvm/CodeGen/TargetLowering.h

public:
   }
 
   /// Returns how the platform's atomic operations are extended (ZERO_EXTEND,
   /// SIGN_EXTEND, or ANY_EXTEND).
   virtual ISD::NodeType getExtendForAtomicOps() const {
     return ISD::ZERO_EXTEND;
   }
 
+  /// Returns how the platform's atomic compare and swap expects its comparison
+  /// value to be extended (ZERO_EXTEND, SIGN_EXTEND, or ANY_EXTEND). This is
+  /// separate from getExtendForAtomicOps, which is concerned with the
+  /// sign-extension of the instruction's output, whereas here we are concerned
+  /// with the sign-extension of the input. For targets with compare-and-swap
+  /// instructions (or sub-word comparisons in their LL/SC loop expansions),
+  /// the input can be ANY_EXTEND, but the output will still have a specific
+  /// extension.
+  virtual ISD::NodeType getExtendForAtomicCmpSwapArg() const {
+    return ISD::ANY_EXTEND;
+  }
+
   /// @}
 
   /// Returns true if we should normalize
   /// select(N0&N1, X, Y) => select(N0, select(N1, X, Y), Y) and
   /// select(N0|N1, X, Y) => select(N0, select(N1, X, Y, Y)) if it is likely
   /// that it saves us from materializing N0 and N1 in an integer register.
   /// Targets that are able to perform and/or on flags should return false here.
   virtual bool shouldNormalizeToSelectSequence(LLVMContext &Context,

llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

SDValue Res = DAG.getAtomicCmpSwap(
         ISD::ATOMIC_CMP_SWAP_WITH_SUCCESS, SDLoc(N), N->getMemoryVT(), VTs,
         N->getChain(), N->getBasePtr(), N->getOperand(2), N->getOperand(3),
         N->getMemOperand());
     ReplaceValueWith(SDValue(N, 0), Res.getValue(0));
     ReplaceValueWith(SDValue(N, 2), Res.getValue(2));
     return Res.getValue(1);
   }
 
-  SDValue Op2 = GetPromotedInteger(N->getOperand(2));
+  // Op2 is used for the comparison and thus must be extended according to the
+  // target's atomic operations. Op3 is merely stored and so can be left alone.
+  SDValue Op2 = N->getOperand(2);
   SDValue Op3 = GetPromotedInteger(N->getOperand(3));
+  switch (TLI.getExtendForAtomicCmpSwapArg()) {
+  case ISD::SIGN_EXTEND:
+    Op2 = SExtPromotedInteger(Op2);
+    break;
+  case ISD::ZERO_EXTEND:
+    Op2 = ZExtPromotedInteger(Op2);
+    break;
+  case ISD::ANY_EXTEND:
+    Op2 = GetPromotedInteger(Op2);
+    break;
+  default:
+    llvm_unreachable("Invalid atomic op extension");
+  }
+
   SDVTList VTs =
       DAG.getVTList(Op2.getValueType(), N->getValueType(1), MVT::Other);
   SDValue Res = DAG.getAtomicCmpSwap(
       N->getOpcode(), SDLoc(N), N->getMemoryVT(), VTs, N->getChain(),
       N->getBasePtr(), Op2, Op3, N->getMemOperand());
   // Update the use to N with the newly created Res.
   for (unsigned i = 1, NumResults = N->getNumValues(); i < NumResults; ++i)
     ReplaceValueWith(SDValue(N, i), Res.getValue(i));

llvm/lib/Target/RISCV/RISCVISelLowering.h

Instruction *emitLeadingFence(IRBuilder<> &Builder, Instruction *Inst,
                                 AtomicOrdering Ord) const override;
   Instruction *emitTrailingFence(IRBuilder<> &Builder, Instruction *Inst,
                                  AtomicOrdering Ord) const override;
 
   ISD::NodeType getExtendForAtomicOps() const override {
     return ISD::SIGN_EXTEND;
   }
 
+  ISD::NodeType getExtendForAtomicCmpSwapArg() const override {
+    return ISD::SIGN_EXTEND;
+  }
+
   bool shouldExpandShift(SelectionDAG &DAG, SDNode *N) const override {
     if (DAG.getMachineFunction().getFunction().hasMinSize())
       return false;
     return true;
   }
   bool isDesirableToCommuteWithShift(const SDNode *N,
                                      CombineLevel Level) const override;
 

llvm/test/CodeGen/RISCV/atomic-cmpxchg.ll

 ; RV64I-NEXT:    mv a4, zero
 ; RV64I-NEXT:    call __atomic_compare_exchange_4
 ; RV64I-NEXT:    ld ra, 8(sp)
 ; RV64I-NEXT:    addi sp, sp, 16
 ; RV64I-NEXT:    ret
 ;
 ; RV64IA-LABEL: cmpxchg_i32_monotonic_monotonic:
 ; RV64IA:       # %bb.0:
+; RV64IA-NEXT:    sext.w a1, a1
 ; RV64IA-NEXT:  .LBB20_1: # =>This Inner Loop Header: Depth=1
 ; RV64IA-NEXT:    lr.w a3, (a0)
 ; RV64IA-NEXT:    bne a3, a1, .LBB20_3
 ; RV64IA-NEXT:  # %bb.2: # in Loop: Header=BB20_1 Depth=1
 ; RV64IA-NEXT:    sc.w a4, a2, (a0)
 ; RV64IA-NEXT:    bnez a4, .LBB20_1
 ; RV64IA-NEXT:  .LBB20_3:
 ; RV64IA-NEXT:    ret

 ; RV64I-NEXT:    mv a4, zero
 ; RV64I-NEXT:    call __atomic_compare_exchange_4
 ; RV64I-NEXT:    ld ra, 8(sp)
 ; RV64I-NEXT:    addi sp, sp, 16
 ; RV64I-NEXT:    ret
 ;
 ; RV64IA-LABEL: cmpxchg_i32_acquire_monotonic:
 ; RV64IA:       # %bb.0:
+; RV64IA-NEXT:    sext.w a1, a1
 ; RV64IA-NEXT:  .LBB21_1: # =>This Inner Loop Header: Depth=1
 ; RV64IA-NEXT:    lr.w.aq a3, (a0)
 ; RV64IA-NEXT:    bne a3, a1, .LBB21_3
 ; RV64IA-NEXT:  # %bb.2: # in Loop: Header=BB21_1 Depth=1
 ; RV64IA-NEXT:    sc.w a4, a2, (a0)
 ; RV64IA-NEXT:    bnez a4, .LBB21_1
 ; RV64IA-NEXT:  .LBB21_3:
 ; RV64IA-NEXT:    ret

 ; RV64I-NEXT:    addi a4, zero, 2
 ; RV64I-NEXT:    call __atomic_compare_exchange_4
 ; RV64I-NEXT:    ld ra, 8(sp)
 ; RV64I-NEXT:    addi sp, sp, 16
 ; RV64I-NEXT:    ret
 ;
 ; RV64IA-LABEL: cmpxchg_i32_acquire_acquire:
 ; RV64IA:       # %bb.0:
+; RV64IA-NEXT:    sext.w a1, a1
 ; RV64IA-NEXT:  .LBB22_1: # =>This Inner Loop Header: Depth=1
 ; RV64IA-NEXT:    lr.w.aq a3, (a0)
 ; RV64IA-NEXT:    bne a3, a1, .LBB22_3
 ; RV64IA-NEXT:  # %bb.2: # in Loop: Header=BB22_1 Depth=1
 ; RV64IA-NEXT:    sc.w a4, a2, (a0)
 ; RV64IA-NEXT:    bnez a4, .LBB22_1
 ; RV64IA-NEXT:  .LBB22_3:
 ; RV64IA-NEXT:    ret

 ; RV64I-NEXT:    mv a4, zero
 ; RV64I-NEXT:    call __atomic_compare_exchange_4
 ; RV64I-NEXT:    ld ra, 8(sp)
 ; RV64I-NEXT:    addi sp, sp, 16
 ; RV64I-NEXT:    ret
 ;
 ; RV64IA-LABEL: cmpxchg_i32_release_monotonic:
 ; RV64IA:       # %bb.0:
+; RV64IA-NEXT:    sext.w a1, a1
 ; RV64IA-NEXT:  .LBB23_1: # =>This Inner Loop Header: Depth=1
 ; RV64IA-NEXT:    lr.w a3, (a0)
 ; RV64IA-NEXT:    bne a3, a1, .LBB23_3
 ; RV64IA-NEXT:  # %bb.2: # in Loop: Header=BB23_1 Depth=1
 ; RV64IA-NEXT:    sc.w.rl a4, a2, (a0)
 ; RV64IA-NEXT:    bnez a4, .LBB23_1
 ; RV64IA-NEXT:  .LBB23_3:
 ; RV64IA-NEXT:    ret

 ; RV64I-NEXT:    addi a4, zero, 2
 ; RV64I-NEXT:    call __atomic_compare_exchange_4
 ; RV64I-NEXT:    ld ra, 8(sp)
 ; RV64I-NEXT:    addi sp, sp, 16
 ; RV64I-NEXT:    ret
 ;
 ; RV64IA-LABEL: cmpxchg_i32_release_acquire:
 ; RV64IA:       # %bb.0:
+; RV64IA-NEXT:    sext.w a1, a1
 ; RV64IA-NEXT:  .LBB24_1: # =>This Inner Loop Header: Depth=1
 ; RV64IA-NEXT:    lr.w a3, (a0)
 ; RV64IA-NEXT:    bne a3, a1, .LBB24_3
 ; RV64IA-NEXT:  # %bb.2: # in Loop: Header=BB24_1 Depth=1
 ; RV64IA-NEXT:    sc.w.rl a4, a2, (a0)
 ; RV64IA-NEXT:    bnez a4, .LBB24_1
 ; RV64IA-NEXT:  .LBB24_3:
 ; RV64IA-NEXT:    ret

 ; RV64I-NEXT:    mv a4, zero
 ; RV64I-NEXT:    call __atomic_compare_exchange_4
 ; RV64I-NEXT:    ld ra, 8(sp)
 ; RV64I-NEXT:    addi sp, sp, 16
 ; RV64I-NEXT:    ret
 ;
 ; RV64IA-LABEL: cmpxchg_i32_acq_rel_monotonic:
 ; RV64IA:       # %bb.0:
+; RV64IA-NEXT:    sext.w a1, a1
 ; RV64IA-NEXT:  .LBB25_1: # =>This Inner Loop Header: Depth=1
 ; RV64IA-NEXT:    lr.w.aq a3, (a0)
 ; RV64IA-NEXT:    bne a3, a1, .LBB25_3
 ; RV64IA-NEXT:  # %bb.2: # in Loop: Header=BB25_1 Depth=1
 ; RV64IA-NEXT:    sc.w.rl a4, a2, (a0)
 ; RV64IA-NEXT:    bnez a4, .LBB25_1
 ; RV64IA-NEXT:  .LBB25_3:
 ; RV64IA-NEXT:    ret

 ; RV64I-NEXT:    addi a4, zero, 2
 ; RV64I-NEXT:    call __atomic_compare_exchange_4
 ; RV64I-NEXT:    ld ra, 8(sp)
 ; RV64I-NEXT:    addi sp, sp, 16
 ; RV64I-NEXT:    ret
 ;
 ; RV64IA-LABEL: cmpxchg_i32_acq_rel_acquire:
 ; RV64IA:       # %bb.0:
+; RV64IA-NEXT:    sext.w a1, a1
 ; RV64IA-NEXT:  .LBB26_1: # =>This Inner Loop Header: Depth=1
 ; RV64IA-NEXT:    lr.w.aq a3, (a0)
 ; RV64IA-NEXT:    bne a3, a1, .LBB26_3
 ; RV64IA-NEXT:  # %bb.2: # in Loop: Header=BB26_1 Depth=1
 ; RV64IA-NEXT:    sc.w.rl a4, a2, (a0)
 ; RV64IA-NEXT:    bnez a4, .LBB26_1
 ; RV64IA-NEXT:  .LBB26_3:
 ; RV64IA-NEXT:    ret

 ; RV64I-NEXT:    mv a4, zero
 ; RV64I-NEXT:    call __atomic_compare_exchange_4
 ; RV64I-NEXT:    ld ra, 8(sp)
 ; RV64I-NEXT:    addi sp, sp, 16
 ; RV64I-NEXT:    ret
 ;
 ; RV64IA-LABEL: cmpxchg_i32_seq_cst_monotonic:
 ; RV64IA:       # %bb.0:
+; RV64IA-NEXT:    sext.w a1, a1
 ; RV64IA-NEXT:  .LBB27_1: # =>This Inner Loop Header: Depth=1
 ; RV64IA-NEXT:    lr.w.aqrl a3, (a0)
 ; RV64IA-NEXT:    bne a3, a1, .LBB27_3
 ; RV64IA-NEXT:  # %bb.2: # in Loop: Header=BB27_1 Depth=1
 ; RV64IA-NEXT:    sc.w.aqrl a4, a2, (a0)
 ; RV64IA-NEXT:    bnez a4, .LBB27_1
 ; RV64IA-NEXT:  .LBB27_3:
 ; RV64IA-NEXT:    ret

 ; RV64I-NEXT:    addi a4, zero, 2
 ; RV64I-NEXT:    call __atomic_compare_exchange_4
 ; RV64I-NEXT:    ld ra, 8(sp)
 ; RV64I-NEXT:    addi sp, sp, 16
 ; RV64I-NEXT:    ret
 ;
 ; RV64IA-LABEL: cmpxchg_i32_seq_cst_acquire:
 ; RV64IA:       # %bb.0:
+; RV64IA-NEXT:    sext.w a1, a1
 ; RV64IA-NEXT:  .LBB28_1: # =>This Inner Loop Header: Depth=1
 ; RV64IA-NEXT:    lr.w.aqrl a3, (a0)
 ; RV64IA-NEXT:    bne a3, a1, .LBB28_3
 ; RV64IA-NEXT:  # %bb.2: # in Loop: Header=BB28_1 Depth=1
 ; RV64IA-NEXT:    sc.w.aqrl a4, a2, (a0)
 ; RV64IA-NEXT:    bnez a4, .LBB28_1
 ; RV64IA-NEXT:  .LBB28_3:
 ; RV64IA-NEXT:    ret

 ; RV64I-NEXT:    addi a4, zero, 5
 ; RV64I-NEXT:    call __atomic_compare_exchange_4
 ; RV64I-NEXT:    ld ra, 8(sp)
 ; RV64I-NEXT:    addi sp, sp, 16
 ; RV64I-NEXT:    ret
 ;
 ; RV64IA-LABEL: cmpxchg_i32_seq_cst_seq_cst:
 ; RV64IA:       # %bb.0:
+; RV64IA-NEXT:    sext.w a1, a1
 ; RV64IA-NEXT:  .LBB29_1: # =>This Inner Loop Header: Depth=1
 ; RV64IA-NEXT:    lr.w.aqrl a3, (a0)
 ; RV64IA-NEXT:    bne a3, a1, .LBB29_3
 ; RV64IA-NEXT:  # %bb.2: # in Loop: Header=BB29_1 Depth=1
 ; RV64IA-NEXT:    sc.w.aqrl a4, a2, (a0)
 ; RV64IA-NEXT:    bnez a4, .LBB29_1
 ; RV64IA-NEXT:  .LBB29_3:
 ; RV64IA-NEXT:    ret