[SDAG] clean up scalarizing load transform

[spatel]

Summary

I have not found a way to expose a difference for this patch in a test because it only triggers for a one-use load, but this is the code that was adapted into D118376 and caused miscompiles. The new code pattern is the same as what we do in narrowExtractedVectorLoad() (reduces load width for a subvector extract).

This removes seemingly unnecessary manual worklist management and fixes the chain updating via:

SDValue SelectionDAG::makeEquivalentMemoryOrdering(SDValue OldChain,
                                                   SDValue NewMemOpChain) {
  assert(isa<MemSDNode>(NewMemOpChain) && "Expected a memop node");
  assert(NewMemOpChain.getValueType() == MVT::Other && "Expected a token VT");
  // The new memory operation must have the same position as the old load in
  // terms of memory dependency. Create a TokenFactor for the old load and new
  // memory operation and update uses of the old load's output chain to use that
  // TokenFactor.
  if (OldChain == NewMemOpChain || OldChain.use_empty())
    return NewMemOpChain;

  SDValue TokenFactor = getNode(ISD::TokenFactor, SDLoc(OldChain), MVT::Other,
                                OldChain, NewMemOpChain);
  ReplaceAllUsesOfValueWith(OldChain, TokenFactor);
  UpdateNodeOperands(TokenFactor.getNode(), OldChain, NewMemOpChain);
  return TokenFactor;
}

[efriedma]

Sure, LGTM

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if (!TLI.allowsMemoryAccess(*DAG.getContext(), DAG.getDataLayout(), VecEltVT,
                               OriginalLoad->getMemOperand()->getFlags(),
                               &IsFast) ||
       !IsFast)
     return SDValue();
 
   SDValue NewPtr = TLI.getVectorElementPointer(DAG, OriginalLoad->getBasePtr(),
                                                InVecVT, EltNo);
 
-  // The replacement we need to do here is a little tricky: we need to
-  // replace an extractelement of a load with a load.
-  // Use ReplaceAllUsesOfValuesWith to do the replacement.
-  // Note that this replacement assumes that the extractvalue is the only
-  // use of the load; that's okay because we don't want to perform this
-  // transformation in other cases anyway.
-  SDValue Load;
-  SDValue Chain;
+  // We are replacing a vector load with a scalar load. The new load must have
+  // identical memory op ordering to the original.
   if (ResultVT.bitsGT(VecEltVT)) {
     // If the result type of vextract is wider than the load, then issue an
     // extending load instead.
-    ISD::LoadExtType ExtType = TLI.isLoadExtLegal(ISD::ZEXTLOAD, ResultVT,
-                                                  VecEltVT)
-                                   ? ISD::ZEXTLOAD
+    ISD::LoadExtType ExtType =
+        TLI.isLoadExtLegal(ISD::ZEXTLOAD, ResultVT, VecEltVT) ? ISD::ZEXTLOAD
                                                               : ISD::EXTLOAD;
-    Load = DAG.getExtLoad(ExtType, SDLoc(EVE), ResultVT,
-                          OriginalLoad->getChain(), NewPtr, MPI, VecEltVT,
-                          Alignment, OriginalLoad->getMemOperand()->getFlags(),
-                          OriginalLoad->getAAInfo());
-    Chain = Load.getValue(1);
-  } else {
-    Load = DAG.getLoad(
-        VecEltVT, SDLoc(EVE), OriginalLoad->getChain(), NewPtr, MPI, Alignment,
-        OriginalLoad->getMemOperand()->getFlags(), OriginalLoad->getAAInfo());
-    Chain = Load.getValue(1);
-    if (ResultVT.bitsLT(VecEltVT))
-      Load = DAG.getNode(ISD::TRUNCATE, SDLoc(EVE), ResultVT, Load);
-    else
-      Load = DAG.getBitcast(ResultVT, Load);
-  }
-  WorklistRemover DeadNodes(*this);
-  SDValue From[] = { SDValue(EVE, 0), SDValue(OriginalLoad, 1) };
-  SDValue To[] = { Load, Chain };
-  DAG.ReplaceAllUsesOfValuesWith(From, To, 2);
-  // Make sure to revisit this node to clean it up; it will usually be dead.
-  AddToWorklist(EVE);
-  // Since we're explicitly calling ReplaceAllUses, add the new node to the
-  // worklist explicitly as well.
-  AddToWorklistWithUsers(Load.getNode());
+    SDValue Load = DAG.getExtLoad(
+        ExtType, SDLoc(EVE), ResultVT, OriginalLoad->getChain(), NewPtr, MPI,
+        VecEltVT, Alignment, OriginalLoad->getMemOperand()->getFlags(),
+    DAG.makeEquivalentMemoryOrdering(OriginalLoad, Load);
+    ++OpsNarrowed;
+    return Load;
+  }
+
+  // The result type is narrower or the same width as the vector element.
+  SDValue Load = DAG.getLoad(
+  DAG.makeEquivalentMemoryOrdering(OriginalLoad, Load);
   ++OpsNarrowed;
-  return SDValue(EVE, 0);
+  return Load;
 }
 
 /// Transform a vector binary operation into a scalar binary operation by moving
 /// the math/logic after an extract element of a vector.
 static SDValue scalarizeExtractedBinop(SDNode *ExtElt, SelectionDAG &DAG,
                                        bool LegalOperations) {
   const TargetLowering &TLI = DAG.getTargetLoweringInfo();
   SDValue Vec = ExtElt->getOperand(0);