[llvm-branch-commits] [llvm] [SelectionDAG] Keep split vector atomic store value in a vector register (PR #201566)

[Matt Arsenault via llvm-branch-commits]

================
@@ -4742,16 +4742,48 @@ SDValue DAGTypeLegalizer::SplitVecOp_STORE(StoreSDNode *N, unsigned OpNo) {
 
 SDValue DAGTypeLegalizer::SplitVecOp_ATOMIC_STORE(AtomicSDNode *N) {
   SDLoc DL(N);
+  LLVMContext &Ctx = *DAG.getContext();
   SDValue StVal = N->getVal();
   EVT VT = StVal.getValueType();
+  EVT MemIntVT = EVT::getIntegerVT(Ctx, N->getMemoryVT().getSizeInBits());
+
+  // The store needs a single value spanning the full memory width. If the
+  // value can be held in a legal vector register, keep it there and extract
+  // the low integer element of the memory width. This lets the store be issued
+  // directly from a vector register (e.g. a single MOVQ/MOVD) instead of
+  // bitcasting the split vector straight to a scalar integer, which would
+  // reassemble the value element by element in GPRs.
+  //
+  // Reinterpret the value as a same-shaped integer vector first: an FP element
+  // type may not have a legal vector form (e.g. bfloat on SSE2) while the
+  // integer-of-element-size form does.
+  unsigned NumElts = VT.getVectorNumElements();
+  EVT IntEltVT = EVT::getIntegerVT(Ctx, VT.getScalarSizeInBits());
+  EVT IntVecVT = VT.changeVectorElementTypeToInteger();
+  EVT IntEltVT = IntVecVT.getVectorElementType();
+  if (DAG.getDataLayout().isLittleEndian() && TLI.isTypeLegal(MemIntVT) &&
+      IntEltVT.getSizeInBits() <= MemIntVT.getSizeInBits()) {
+    EVT WideVT = IntVecVT;
+    while (!TLI.isTypeLegal(WideVT) && WideVT.getSizeInBits() < 512)
+      WideVT =
+          EVT::getVectorVT(Ctx, IntEltVT, WideVT.getVectorNumElements() * 2);
+    if (TLI.isTypeLegal(WideVT) &&
+        WideVT.getSizeInBits() % MemIntVT.getSizeInBits() == 0) {
----------------
arsenm wrote:

Not a fan of this custom type finding loop. Can you use one of the existing find legal vector type functions. This would also depend on if the extract is free?

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