[LegalizeTypes][RISCV] Control sign-extend for atomicrmw input argument

Similar to the argument sign-extend control provided by
getExtendForAtomicCmpSwapArg for `cmpxchg`, this change adds
getExtendForAtomicRMWArg for `atomicrmw <op>`.

This mechanism is used to correct the RISCV code generation when using atomic
pseudo expansions that perform sub-word comparions inside a LR/SC block, and
expect a correctly sign-extended input argument.  A shorter instruction
sequence is possible for umin/max by leveraging the fact that two sign-extended
arguments still order correctly during an unsigned comparison.

Author: slachowsky

llvm/include/llvm/CodeGen/TargetLowering.h

@@ -2459,6 +2459,12 @@ class LLVM_ABI TargetLoweringBase {
     return ISD::ANY_EXTEND;
   }
 
+  /// Returns how the platform's atomic rmw operations expect their input
+  /// argument to be extended (ZERO_EXTEND, SIGN_EXTEND, or ANY_EXTEND).
+  virtual ISD::NodeType getExtendForAtomicRMWArg(unsigned Op) const {
+    return ISD::ANY_EXTEND;
+  }
+
   /// @}
 
   /// Returns true if we should normalize

llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

@@ -429,7 +429,20 @@ SDValue DAGTypeLegalizer::PromoteIntRes_Atomic0(AtomicSDNode *N) {
 }
 
 SDValue DAGTypeLegalizer::PromoteIntRes_Atomic1(AtomicSDNode *N) {
-  SDValue Op2 = GetPromotedInteger(N->getOperand(2));
+  SDValue Op2 = N->getOperand(2);
+  switch (TLI.getExtendForAtomicRMWArg(N->getOpcode())) {
+  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");
+  }
   SDValue Res = DAG.getAtomic(N->getOpcode(), SDLoc(N),
                               N->getMemoryVT(),
                               N->getChain(), N->getBasePtr(),

llvm/lib/Target/RISCV/RISCVExpandAtomicPseudoInsts.cpp

@@ -521,13 +521,6 @@ static void insertSext(const RISCVInstrInfo *TII, DebugLoc DL,
       .addReg(ShamtReg);
 }
 
-static void insertZext(const RISCVInstrInfo *TII, DebugLoc DL,
-                       MachineBasicBlock *MBB, Register ValReg, Register SrcReg,
-                       int64_t Shamt) {
-  BuildMI(MBB, DL, TII->get(RISCV::SLLI), ValReg).addReg(SrcReg).addImm(Shamt);
-  BuildMI(MBB, DL, TII->get(RISCV::SRLI), ValReg).addReg(ValReg).addImm(Shamt);
-}
-
 static void doAtomicMinMaxOpExpansion(
     const RISCVInstrInfo *TII, MachineInstr &MI, DebugLoc DL,
     MachineBasicBlock *ThisMBB, MachineBasicBlock *LoopHeadMBB,
@@ -538,9 +531,6 @@ static void doAtomicMinMaxOpExpansion(
   Register ScratchReg = MI.getOperand(1).getReg();
   Register AddrReg = MI.getOperand(2).getReg();
   Register IncrReg = MI.getOperand(3).getReg();
-  bool IsUnsigned =
-      BinOp == AtomicRMWInst::UMin || BinOp == AtomicRMWInst::UMax;
-  bool Zext = IsUnsigned && STI->is64Bit() && Width == 32;
   AtomicOrdering Ordering =
       static_cast<AtomicOrdering>(MI.getOperand(4).getImm());
 
@@ -550,12 +540,9 @@ static void doAtomicMinMaxOpExpansion(
   //   ifnochangeneeded scratch, incr, .looptail
   BuildMI(LoopHeadMBB, DL, TII->get(getLRForRMW(Ordering, Width, STI)), DestReg)
       .addReg(AddrReg);
-  if (Zext)
-    insertZext(TII, DL, LoopHeadMBB, ScratchReg, DestReg, 32);
-  else
-    BuildMI(LoopHeadMBB, DL, TII->get(RISCV::ADDI), ScratchReg)
-        .addReg(DestReg)
-        .addImm(0);
+  BuildMI(LoopHeadMBB, DL, TII->get(RISCV::ADDI), ScratchReg)
+      .addReg(DestReg)
+      .addImm(0);
   switch (BinOp) {
   default:
     llvm_unreachable("Unexpected AtomicRMW BinOp");

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

@@ -24485,6 +24485,25 @@ ISD::NodeType RISCVTargetLowering::getExtendForAtomicCmpSwapArg() const {
   return Subtarget.hasStdExtZacas() ? ISD::ANY_EXTEND : ISD::SIGN_EXTEND;
 }
 
+ISD::NodeType RISCVTargetLowering::getExtendForAtomicRMWArg(unsigned Op) const {
+  // Zaamo will use amo<op>.w which does not require extension.
+  if (Subtarget.hasStdExtZaamo() || Subtarget.hasForcedAtomics())
+    return ISD::ANY_EXTEND;
+
+  // Zalrsc pseudo expansions with comparison require sign-extension.
+  assert(Subtarget.hasStdExtZalrsc());
+  switch (Op) {
+  case ISD::ATOMIC_LOAD_MIN:
+  case ISD::ATOMIC_LOAD_MAX:
+  case ISD::ATOMIC_LOAD_UMIN:
+  case ISD::ATOMIC_LOAD_UMAX:
+    return ISD::SIGN_EXTEND;
+  default:
+    break;
+  }
+  return ISD::ANY_EXTEND;
+}
+
 Register RISCVTargetLowering::getExceptionPointerRegister(
     const Constant *PersonalityFn) const {
   return RISCV::X10;

llvm/lib/Target/RISCV/RISCVISelLowering.h

@@ -245,6 +245,7 @@ class RISCVTargetLowering : public TargetLowering {
   }
 
   ISD::NodeType getExtendForAtomicCmpSwapArg() const override;
+  ISD::NodeType getExtendForAtomicRMWArg(unsigned Op) const override;
 
   bool shouldTransformSignedTruncationCheck(EVT XVT,
                                             unsigned KeptBits) const override;

llvm/lib/Target/RISCV/RISCVInstrInfoA.td

@@ -321,11 +321,9 @@ def PseudoAtomicLoadXor32 : PseudoAMO;
 let Size = 24 in {
 def PseudoAtomicLoadMax32 : PseudoAMO;
 def PseudoAtomicLoadMin32 : PseudoAMO;
-} // Size = 24
-let Size = 28 in {
 def PseudoAtomicLoadUMax32 : PseudoAMO;
 def PseudoAtomicLoadUMin32 : PseudoAMO;
-} // Size = 28
+} // Size = 24
 
 defm : PseudoAMOPat<"atomic_swap_i32", PseudoAtomicSwap32>;
 defm : PseudoAMOPat<"atomic_load_add_i32", PseudoAtomicLoadAdd32>;