[PR27599] [SystemZ] [SelectionDAG] Fix extension of atomic cmpxchg result.

[• koriakin]

Summary

Currently, SelectionDAG assumes 8/16-bit cmpxchg returns either a sign
extended result, or a zero extended result. SystemZ takes a third
option by returning junk in the high bits (rotated contents of the other
bytes in the memory word). In that case, don't use Assert*ext, and
zero-extend the result ourselves if a comparison is needed.

[• koriakin]

Ping?

[dsanders]

lib/CodeGen/SelectionDAG/LegalizeDAG.cpp

2863

ISD::ZERO_EXTEND only takes one operand so the DAG.getValueType(AtomicType) argument isn't necessary.

2871

I'm not sure why the ANY_EXTEND case returns ExtRes.getValue(0) instead of LHS.getValue(0). Is there a reason for it?

I'm currently thinking that returning LHS allows SelectionDAG to take advantage of the additional known-zero bits from the zero extend but returning ExtRes allows more scheduling freedom because it can be used before the comparison completes. I'm not sure which is the better choice though.

[• koriakin]

lib/CodeGen/SelectionDAG/LegalizeDAG.cpp

2871

The comparison result may be unused and DCEd - in this case, it's better not to do the zero extension at all.

[• koriakin]

Fixed inline comment.

[dsanders]

LGTM with the tests changes.

lib/CodeGen/SelectionDAG/LegalizeDAG.cpp

2871

That makes sense to me. ATOMIC_CMP_SWAP_WITH_SUCCESS is always a strong operation so it's not strange to ignore the success result.

test/CodeGen/SystemZ/cmpxchg-05.ll

5–11

This should check that the zero extend and comparison is not emitted in this test (and likewise below).

I think there should also be a test that doesn't ignore the success result to check the comparison is as expected and also one that zero extends the success result to check that multiple zero extends don't occur.

include/llvm/Target/TargetLowering.h

   /// it succeeds, and nullptr otherwise.
   /// If shouldExpandAtomicLoadInIR returns true on that load, it will undergo
   /// another round of expansion.
   virtual LoadInst *
   lowerIdempotentRMWIntoFencedLoad(AtomicRMWInst *RMWI) const {
     return nullptr;
   }
 
-  /// Returns true if the platform's atomic operations are sign extended.
-  virtual bool hasSignExtendedAtomicOps() const {
-    return false;
+  /// 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 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.

lib/CodeGen/SelectionDAG/LegalizeDAG.cpp

case ISD::ATOMIC_CMP_SWAP_WITH_SUCCESS: {
     SDValue Res = DAG.getAtomicCmpSwap(
         ISD::ATOMIC_CMP_SWAP, dl, cast<AtomicSDNode>(Node)->getMemoryVT(), VTs,
         Node->getOperand(0), Node->getOperand(1), Node->getOperand(2),
         Node->getOperand(3), cast<MemSDNode>(Node)->getMemOperand(),
         cast<AtomicSDNode>(Node)->getSuccessOrdering(),
         cast<AtomicSDNode>(Node)->getFailureOrdering(),
         cast<AtomicSDNode>(Node)->getSynchScope());
 
+    SDValue ExtRes = Res;
     SDValue LHS = Res;
     SDValue RHS = Node->getOperand(1);
 
     EVT AtomicType = cast<AtomicSDNode>(Node)->getMemoryVT();
     EVT OuterType = Node->getValueType(0);
-    if (TLI.hasSignExtendedAtomicOps()) {
+    switch (TLI.getExtendForAtomicOps()) {
+    case ISD::SIGN_EXTEND:
       LHS = DAG.getNode(ISD::AssertSext, dl, OuterType, Res,
                         DAG.getValueType(AtomicType));
       RHS = DAG.getNode(ISD::SIGN_EXTEND_INREG, dl, OuterType,
                         Node->getOperand(2), DAG.getValueType(AtomicType));
-    } else {
-      LHS = DAG.getNode(ISD::AssertZext, dl, OuterType, Res, DAG.getValueType(AtomicType));
+      ExtRes = LHS;
+      break;
+    case ISD::ZERO_EXTEND:
+      LHS = DAG.getNode(ISD::AssertZext, dl, OuterType, Res,
+                        DAG.getValueType(AtomicType));
       RHS = DAG.getNode(ISD::ZERO_EXTEND, dl, OuterType, Node->getOperand(2));
+      ExtRes = LHS;
+      break;
+    case ISD::ANY_EXTEND:
+      LHS = DAG.getZeroExtendInReg(Res, dl, AtomicType);
+      RHS = DAG.getNode(ISD::ZERO_EXTEND, dl, OuterType, Node->getOperand(2));
+      break;
+    default:
+      llvm_unreachable("Invalid atomic op extension");
     }
 
     SDValue Success =
         DAG.getSetCC(dl, Node->getValueType(1), LHS, RHS, ISD::SETEQ);
 
-    Results.push_back(LHS.getValue(0));
+    Results.push_back(ExtRes.getValue(0));
     Results.push_back(Success);
     Results.push_back(Res.getValue(1));
     break;
   }
   case ISD::DYNAMIC_STACKALLOC:
     ExpandDYNAMIC_STACKALLOC(Node, Results);
     break;
   case ISD::MERGE_VALUES:

lib/Target/Mips/MipsISelLowering.h

public:
 
     MVT getScalarShiftAmountTy(const DataLayout &, EVT) const override {
       return MVT::i32;
     }
 
     bool isCheapToSpeculateCttz() const override;
     bool isCheapToSpeculateCtlz() const override;
 
-    bool hasSignExtendedAtomicOps() const override {
-      return true;
+    ISD::NodeType getExtendForAtomicOps() const override {
+      return ISD::SIGN_EXTEND;
     }
 
     void LowerOperationWrapper(SDNode *N,
                                SmallVectorImpl<SDValue> &Results,
                                SelectionDAG &DAG) const override;
 
     /// LowerOperation - Provide custom lowering hooks for some operations.
     SDValue LowerOperation(SDValue Op, SelectionDAG &DAG) const override;

lib/Target/SystemZ/SystemZISelLowering.h

public:
   SDValue LowerReturn(SDValue Chain, CallingConv::ID CallConv, bool IsVarArg,
                       const SmallVectorImpl<ISD::OutputArg> &Outs,
                       const SmallVectorImpl<SDValue> &OutVals,
                       SDLoc DL, SelectionDAG &DAG) const override;
   SDValue prepareVolatileOrAtomicLoad(SDValue Chain, SDLoc DL,
                                       SelectionDAG &DAG) const override;
   SDValue PerformDAGCombine(SDNode *N, DAGCombinerInfo &DCI) const override;
 
+  ISD::NodeType getExtendForAtomicOps() const override {
+    return ISD::ANY_EXTEND;
+  }
+
   bool supportSwiftError() const override {
     return true;
   }
 
 private:
   const SystemZSubtarget &Subtarget;
 
   // Implement LowerOperation for individual opcodes.

test/CodeGen/SystemZ/cmpxchg-05.ll

+; Test proper extension of 8-bit/16-bit cmpxchg.
+;
+; RUN: llc < %s -mtriple=s390x-linux-gnu | FileCheck %s
+
+; CHECK-LABEL: f1
+; CHECK: llgcr %r2, [[RES:%r[0-9]+]]
+define zeroext i8 @f1(i8* nocapture, i8 zeroext, i8 zeroext) {
+  %cx = cmpxchg i8* %0, i8 %1, i8 %2 seq_cst seq_cst
+  %res = extractvalue { i8, i1 } %cx, 0
+  ret i8 %res
+}
+
+; CHECK-LABEL: f2
+; CHECK: llghr %r2, [[RES:%r[0-9]+]]
+define zeroext i16 @f2(i16* nocapture, i16 zeroext, i16 zeroext) {
+  %cx = cmpxchg i16* %0, i16 %1, i16 %2 seq_cst seq_cst
+  %res = extractvalue { i16, i1 } %cx, 0
+  ret i16 %res
+}
+
+; CHECK-LABEL: f3
+; CHECK: lgbr %r2, [[RES:%r[0-9]+]]
+define signext i8 @f3(i8* nocapture, i8 signext, i8 signext) {
+  %cx = cmpxchg i8* %0, i8 %1, i8 %2 seq_cst seq_cst
+  %res = extractvalue { i8, i1 } %cx, 0
+  ret i8 %res
+}
+
+; CHECK-LABEL: f4
+; CHECK: lghr %r2, [[RES:%r[0-9]+]]
+define signext i16 @f4(i16* nocapture, i16 signext, i16 signext) {
+  %cx = cmpxchg i16* %0, i16 %1, i16 %2 seq_cst seq_cst
+  %res = extractvalue { i16, i1 } %cx, 0
+  ret i16 %res
+}