[SystemZ] Don't lower ATOMIC_[LOAD|STORE] to [LOAD|STORE]

(Use PatFrags for loads.)
Try to convert to LOAD in select() instead.
Was e20dad7

Author: JonPsson1

llvm/include/llvm/CodeGen/SelectionDAGNodes.h

@@ -558,6 +558,7 @@ BEGIN_TWO_BYTE_PACK()
 
   class LoadSDNodeBitfields {
     friend class LoadSDNode;
+    friend class AtomicSDNode;
     friend class VPLoadSDNode;
     friend class VPStridedLoadSDNode;
     friend class MaskedLoadSDNode;
@@ -1473,6 +1474,16 @@ class AtomicSDNode : public MemSDNode {
             MMO->isAtomic()) && "then why are we using an AtomicSDNode?");
   }
 
+  void setExtensionType(ISD::LoadExtType ETy) {
+    assert(getOpcode() == ISD::ATOMIC_LOAD && "Only used for atomic loads.");
+    LoadSDNodeBits.ExtTy = ETy;
+  }
+
+  ISD::LoadExtType getExtensionType() const {
+    assert(getOpcode() == ISD::ATOMIC_LOAD && "Only used for atomic loads.");
+    return static_cast<ISD::LoadExtType>(LoadSDNodeBits.ExtTy);
+  }
+
   const SDValue &getBasePtr() const {
     return getOpcode() == ISD::ATOMIC_STORE ? getOperand(2) : getOperand(1);
   }

llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

@@ -340,6 +340,19 @@ SDValue DAGTypeLegalizer::PromoteIntRes_Atomic0(AtomicSDNode *N) {
                               N->getMemoryVT(), ResVT,
                               N->getChain(), N->getBasePtr(),
                               N->getMemOperand());
+  if (N->getOpcode() == ISD::ATOMIC_LOAD) {
+    ISD::LoadExtType ETy = cast<AtomicSDNode>(N)->getExtensionType();
+    if (ETy == ISD::NON_EXTLOAD) {
+      if (TLI.getExtendForAtomicOps() == ISD::SIGN_EXTEND)
+        ETy = ISD::SEXTLOAD;
+      else if (TLI.getExtendForAtomicOps() == ISD::ZERO_EXTEND)
+        ETy = ISD::ZEXTLOAD;
+      else
+        ETy = ISD::EXTLOAD;
+    }
+    cast<AtomicSDNode>(Res)->setExtensionType(ETy);
+  }
+
   // Legalize the chain result - switch anything that used the old chain to
   // use the new one.
   ReplaceValueWith(SDValue(N, 1), Res.getValue(1));

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

@@ -4037,6 +4037,9 @@ KnownBits SelectionDAG::computeKnownBits(SDValue Op, const APInt &DemandedElts,
     if (Op.getResNo() == 0) {
       if (TLI->getExtendForAtomicOps() == ISD::ZERO_EXTEND)
         Known.Zero.setBitsFrom(MemBits);
+      else if (Op->getOpcode() == ISD::ATOMIC_LOAD &&
+               cast<AtomicSDNode>(Op)->getExtensionType() == ISD::ZEXTLOAD)
+        Known.Zero.setBitsFrom(MemBits);
     }
     break;
   }
@@ -4848,6 +4851,13 @@ unsigned SelectionDAG::ComputeNumSignBits(SDValue Op, const APInt &DemandedElts,
         return VTBits - Tmp + 1;
       if (TLI->getExtendForAtomicOps() == ISD::ZERO_EXTEND)
         return VTBits - Tmp;
+      if (Op->getOpcode() == ISD::ATOMIC_LOAD) {
+        ISD::LoadExtType ETy = cast<AtomicSDNode>(Op)->getExtensionType();
+        if (ETy == ISD::SEXTLOAD)
+          return VTBits - Tmp + 1;
+        if (ETy == ISD::ZEXTLOAD)
+          return VTBits - Tmp;
+      }
     }
     break;
   }

llvm/lib/CodeGen/SelectionDAG/SelectionDAGDumper.cpp

@@ -828,6 +828,18 @@ void SDNode::print_details(raw_ostream &OS, const SelectionDAG *G) const {
   } else if (const MemSDNode *M = dyn_cast<MemSDNode>(this)) {
     OS << "<";
     printMemOperand(OS, *M->getMemOperand(), G);
+    if (auto *A = dyn_cast<AtomicSDNode>(M))
+      if (A->getOpcode() == ISD::ATOMIC_LOAD) {
+        bool doExt = true;
+        switch (A->getExtensionType()) {
+        default: doExt = false; break;
+        case ISD::EXTLOAD:  OS << ", anyext"; break;
+        case ISD::SEXTLOAD: OS << ", sext"; break;
+        case ISD::ZEXTLOAD: OS << ", zext"; break;
+        }
+        if (doExt)
+          OS << " from " << A->getMemoryVT();
+      }
     OS << ">";
   } else if (const BlockAddressSDNode *BA =
                dyn_cast<BlockAddressSDNode>(this)) {

llvm/lib/Target/SystemZ/SystemZISelDAGToDAG.cpp

@@ -347,6 +347,9 @@ class SystemZDAGToDAGISel : public SelectionDAGISel {
   // Try to expand a boolean SELECT_CCMASK using an IPM sequence.
   SDValue expandSelectBoolean(SDNode *Node);
 
+  // Convert ATOMIC_LOADs to LOADs to facilitate instruction selection.
+  void convertATOMIC_LOADs(SDNode *Node, unsigned Depth = 0);
+
 public:
   static char ID;
 
@@ -1513,6 +1516,10 @@ bool SystemZDAGToDAGISel::storeLoadIsAligned(SDNode *N) const {
   MachineMemOperand *MMO = MemAccess->getMemOperand();
   assert(MMO && "Expected a memory operand.");
 
+  // These instructions are not atomic.
+  if (MMO->isAtomic())
+    return false;
+
   // The memory access must have a proper alignment and no index register.
   if (MemAccess->getAlign().value() < StoreSize ||
       !MemAccess->getOffset().isUndef())
@@ -1545,6 +1552,37 @@ bool SystemZDAGToDAGISel::storeLoadIsAligned(SDNode *N) const {
   return true;
 }
 
+// This is a hack to convert ATOMIC_LOADs to LOADs in the last minute just
+// before instruction selection begins. It would have been easier if
+// ATOMIC_LOAD nodes would instead always be built by SelectionDAGBuilder as
+// LOADs with an atomic MMO and properly handled as such in DAGCombiner, but
+// until that changes they need to remain as ATOMIC_LOADs until all
+// DAGCombining is done.  Convert Node or any of its operands from
+// ATOMIC_LOAD to LOAD.
+void SystemZDAGToDAGISel::convertATOMIC_LOADs(SDNode *Node, unsigned Depth) {
+  if (Depth > 1) // Chain operands are also followed so this seems enough.
+    return;
+  if (Node->getOpcode() == ISD::ATOMIC_LOAD) {
+    auto *ALoad = cast<AtomicSDNode>(Node);
+    // It seems necessary to morph the node as it is not yet being selected.
+    LoadSDNode *Ld = cast<LoadSDNode>(CurDAG->MorphNodeTo(
+        ALoad, ISD::LOAD, CurDAG->getVTList(ALoad->getValueType(0), MVT::Other),
+        {ALoad->getChain(), ALoad->getBasePtr()}));
+    // Sanity check the morph.  The extension type for an extending load
+    // should have been set prior to instruction selection and remain in the
+    // morphed node.
+    assert(((SDNode *)Ld) == ((SDNode *)ALoad) && "Bad CSE on atomic load.");
+    assert(Ld->getMemOperand()->isAtomic() && "Broken MMO.");
+    ISD::LoadExtType ETy = Ld->getExtensionType();
+    bool IsNonExt = Ld->getMemoryVT().getSizeInBits() ==
+                    Ld->getValueType(0).getSizeInBits();
+    assert(IsNonExt == (ETy == ISD::NON_EXTLOAD) && "Bad extension type.");
+    return;
+  }
+  for (SDValue Op : Node->ops())
+    convertATOMIC_LOADs(Op.getNode(), ++Depth);
+}
+
 void SystemZDAGToDAGISel::Select(SDNode *Node) {
   // If we have a custom node, we already have selected!
   if (Node->isMachineOpcode()) {
@@ -1553,6 +1591,9 @@ void SystemZDAGToDAGISel::Select(SDNode *Node) {
     return;
   }
 
+  // Prepare any ATOMIC_LOAD to be selected as a LOAD with an atomic MMO.
+  convertATOMIC_LOADs(Node);
+
   unsigned Opcode = Node->getOpcode();
   switch (Opcode) {
   case ISD::OR:
@@ -1742,6 +1783,31 @@ void SystemZDAGToDAGISel::Select(SDNode *Node) {
     }
     break;
   }
+
+  case ISD::ATOMIC_STORE: {
+    auto *AtomOp = cast<AtomicSDNode>(Node);
+    // Store FP values directly without first moving to a GPR.
+    EVT SVT = AtomOp->getMemoryVT();
+    SDValue StoredVal = AtomOp->getVal();
+    if (SVT.isInteger() && StoredVal->getOpcode() == ISD::BITCAST &&
+        StoredVal->getOperand(0).getValueType().isFloatingPoint()) {
+      StoredVal = StoredVal->getOperand(0);
+      SVT = StoredVal.getValueType();
+    }
+    StoreSDNode *St = cast<StoreSDNode>(CurDAG->getTruncStore(
+        AtomOp->getChain(), SDLoc(AtomOp), StoredVal, AtomOp->getBasePtr(), SVT,
+        AtomOp->getMemOperand()));
+    assert(St->getMemOperand()->isAtomic() && "Broken MMO.");
+    SDNode *Chain = St;
+    // We have to enforce sequential consistency by performing a
+    // serialization operation after the store.
+    if (AtomOp->getSuccessOrdering() == AtomicOrdering::SequentiallyConsistent)
+      Chain = CurDAG->getMachineNode(SystemZ::Serialize, SDLoc(AtomOp),
+                                     MVT::Other, SDValue(Chain, 0));
+    ReplaceNode(Node, Chain);
+    SelectCode(St);
+    return;
+  }
   }
 
   SelectCode(Node);