[llvm-branch-commits] [llvm] [SelectionDAG][X86] Split via Concat <n x T> vector types for atomic load (PR #120640)

[Matt Arsenault via llvm-branch-commits]

================
@@ -1421,6 +1424,35 @@ void DAGTypeLegalizer::SplitVectorResult(SDNode *N, unsigned ResNo) {
     SetSplitVector(SDValue(N, ResNo), Lo, Hi);
 }
 
+void DAGTypeLegalizer::SplitVecRes_ATOMIC_LOAD(AtomicSDNode *LD) {
+  SDLoc dl(LD);
+
+  EVT MemoryVT = LD->getMemoryVT();
+  unsigned NumElts = MemoryVT.getVectorMinNumElements();
+
+  EVT IntMemoryVT = EVT::getVectorVT(*DAG.getContext(), MVT::i16, NumElts);
+  EVT ElemVT =
+      EVT::getVectorVT(*DAG.getContext(), MemoryVT.getVectorElementType(), 1);
+
+  // Create a single atomic to load all the elements at once.
+  SDValue Atomic =
+      DAG.getAtomic(ISD::ATOMIC_LOAD, dl, IntMemoryVT, IntMemoryVT,
+                    LD->getChain(), LD->getBasePtr(), LD->getMemOperand());
+
+  // Instead of splitting, put all the elements back into a vector.
+  SmallVector<SDValue, 4> Ops;
+  for (unsigned i = 0; i < NumElts; ++i) {
+    SDValue Elt = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, dl, MVT::i16, Atomic,
+                              DAG.getVectorIdxConstant(i, dl));
+    Elt = DAG.getBitcast(ElemVT, Elt);
+    Ops.push_back(Elt);
+  }
+  SDValue Concat = DAG.getNode(ISD::CONCAT_VECTORS, dl, MemoryVT, Ops);
+
+  ReplaceValueWith(SDValue(LD, 0), Concat);
----------------
arsenm wrote:

You don't need to associate the Hi and Lo halves, you need to report the Hi and Lo halves to the caller. The only special thing you need to do here is replace the chain uses, as is done in the non-atomic load case 

https://github.com/llvm/llvm-project/pull/120640