[x86] Fold extract_vector_elt of a load into the load's address computation.

[Bigcheese]

Summary

An address only use of an extract element of a load can be simplified to a
load. Without this the result of the extract element is spilled to the
stack so that an address is available.

[grosbach]

This looks like a really nice improvement, and you're right, probably not just for x86.

I suggest looking into making this a target-independent DAG combine that can be predicated on queries to TargetLowering about what types and operations are legal to see when it's profitable. I suspect this is going to be the right thing to do for pretty much all CPU targets.

[qcolombet]

Hi Michael,

DAGCombiner has a load slicing mechanism what tries to produce target friendly loads. Maybe we could extend this mechanism to handle this kind of cases.
Do not know how easy that would be, but that seems to me like the natural place to start :).

You can take a look at DAGCombiner::SliceUpLoad.

Thanks,
-Quentin

[Bigcheese]

Moved to DAGCombiner and refactored the existing load slicing code.

[Bigcheese]

ping

[qcolombet]

Hi Michael,

This LGTM with one remark, I leave it to you whether or not you want to address it since Nadav already gave his LGTM.

I do not particularly like the use of the lambda here. I think it would have been clearer to explicitly add the few more variables that we use (ConstEltNode, NewLowd, EltNo) and have a new (documented) private member function.
I see two advantages of using a member function instead of the lambda:

• DAGCombiner::visitEXTRACT_VECTOR_ELT would be easier to read and understand.
• The member function can be properly documented.

What do you think?

Thanks,
-Quentin

[Bigcheese]

Committed as r209788 with fixes.

[qcolombet]

Thanks!

-Quentin

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if (InVec.getOpcode() == ISD::VECTOR_SHUFFLE
 
     if (!LegalOperations) {
       EVT IndexTy = TLI.getVectorIdxTy();
       return DAG.getNode(ISD::EXTRACT_VECTOR_ELT, SDLoc(N), NVT,
                          SVInVec, DAG.getConstant(OrigElt, IndexTy));
     }
   }
 
-  // Perform only after legalization to ensure build_vector / vector_shuffle
-  // optimizations have already been done.
-  if (!LegalOperations) return SDValue();
-
-  // (vextract (v4f32 load $addr), c) -> (f32 load $addr+c*size)
-  // (vextract (v4f32 s2v (f32 load $addr)), c) -> (f32 load $addr+c*size)
-  // (vextract (v4f32 shuffle (load $addr), <1,u,u,u>), 0) -> (f32 load $addr)
-
-  if (ConstEltNo) {
-    int Elt = cast<ConstantSDNode>(EltNo)->getZExtValue();
-    bool NewLoad = false;
-    bool BCNumEltsChanged = false;
-    EVT ExtVT = VT.getVectorElementType();
-    EVT LVT = ExtVT;
 
-    // If the result of load has to be truncated, then it's not necessarily
-    // profitable.
-    if (NVT.bitsLT(LVT) && !TLI.isTruncateFree(LVT, NVT))
-      return SDValue();
 
-    if (InVec.getOpcode() == ISD::BITCAST) {
-      // Don't duplicate a load with other uses.
-      if (!InVec.hasOneUse())
-        return SDValue();
 
-      EVT BCVT = InVec.getOperand(0).getValueType();
-      if (!BCVT.isVector() || ExtVT.bitsGT(BCVT.getVectorElementType()))
-        return SDValue();
-      if (VT.getVectorNumElements() != BCVT.getVectorNumElements())
-        BCNumEltsChanged = true;
-      InVec = InVec.getOperand(0);
-      ExtVT = BCVT.getVectorElementType();
-      NewLoad = true;
-    }
 
+  auto ReplaceLoad = [&](EVT LoadVT, LoadSDNode *OriginalLoad,
+                         SDValue EltIndex) {
+    unsigned Align = OriginalLoad->getAlignment();
+    if (NewLoad) {
+      // Check the resultant load doesn't need a higher alignment than the
+      // original load.
+      unsigned NewAlign =
+        TLI.getDataLayout()
+            ->getABITypeAlignment(LoadVT.getTypeForEVT(*DAG.getContext()));
+
+      if (NewAlign > Align || !TLI.isOperationLegalOrCustom(ISD::LOAD, LoadVT))
+        return SDValue();
+
+      Align = NewAlign;
+    }
+
+    SDValue NewPtr = OriginalLoad->getBasePtr();
+    SDValue Offset;
+    EVT PtrType = NewPtr.getValueType();
+    MachinePointerInfo MPI;
+    if (ConstEltNo) {
+      int Elt = cast<ConstantSDNode>(EltNo)->getZExtValue();
+      unsigned PtrOff = LoadVT.getSizeInBits() * Elt / 8;
+      if (TLI.isBigEndian())
+        PtrOff = VT.getSizeInBits() / 8 - PtrOff;
+      Offset = DAG.getConstant(PtrOff, PtrType);
+      MPI = OriginalLoad->getPointerInfo().getWithOffset(PtrOff);
+    } else {
+      Offset = DAG.getNode(
+          ISD::MUL, SDLoc(N), EltNo.getValueType(), EltNo,
+          DAG.getConstant(LoadVT.getStoreSize(), EltNo.getValueType()));
+      if (TLI.isBigEndian())
+        Offset = DAG.getNode(ISD::SUB, SDLoc(N), EltNo.getValueType(),
+                             DAG.getConstant(VT.getStoreSize(),
+                                             EltNo.getValueType()),
+                             Offset);
+      MPI = OriginalLoad->getPointerInfo();
+    }
+    NewPtr = DAG.getNode(ISD::ADD, SDLoc(N), PtrType, NewPtr, Offset);
+
+    // 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;
+    if (NVT.bitsGT(LoadVT)) {
+      // If the result type of vextract is wider than the load, then issue an
+      // extending load instead.
+      ISD::LoadExtType ExtType = TLI.isLoadExtLegal(ISD::ZEXTLOAD, LoadVT)
+                                     ? ISD::ZEXTLOAD
+                                     : ISD::EXTLOAD;
+      Load = DAG.getExtLoad(ExtType, SDLoc(N), NVT, OriginalLoad->getChain(),
+                            NewPtr, MPI, LoadVT, OriginalLoad->isVolatile(),
+                            OriginalLoad->isNonTemporal(), Align,
+                            OriginalLoad->getTBAAInfo());
+      Chain = Load.getValue(1);
+    } else {
+      Load = DAG.getLoad(
+          LoadVT, SDLoc(N), OriginalLoad->getChain(), NewPtr, MPI,
+          OriginalLoad->isVolatile(), OriginalLoad->isNonTemporal(),
+          OriginalLoad->isInvariant(), Align, OriginalLoad->getTBAAInfo());
+      Chain = Load.getValue(1);
+      if (NVT.bitsLT(LoadVT))
+        Load = DAG.getNode(ISD::TRUNCATE, SDLoc(N), NVT, Load);
+      else
+        Load = DAG.getNode(ISD::BITCAST, SDLoc(N), NVT, Load);
+    }
+    WorkListRemover DeadNodes(*this);
+    SDValue From[] = { SDValue(N, 0), SDValue(OriginalLoad, 1) };
+    SDValue To[] = { Load, Chain };
+    DAG.ReplaceAllUsesOfValuesWith(From, To, 2);
+    // Since we're explicitly calling ReplaceAllUses, add the new node to the
+    // worklist explicitly as well.
+    AddToWorkList(Load.getNode());
+    AddUsersToWorkList(Load.getNode()); // Add users too
+    // Make sure to revisit this node to clean it up; it will usually be dead.
+    AddToWorkList(N);
+    ++OpsNarrowed;
+    return SDValue(N, 0);
+  };
+
+  // (vextract (vN[if]M load $addr), i) -> ([if]M load $addr + i * size)
+  if (!LegalOperations && !ConstEltNo && InVec.hasOneUse() &&
+      ISD::isNormalLoad(InVec.getNode())) {
+    SDValue Index = N->getOperand(1);
+    if (LoadSDNode *OrigLoad = dyn_cast<LoadSDNode>(InVec))
+      return ReplaceLoad(LVT, OrigLoad, Index);
+  }
+
+  // Perform only after legalization to ensure build_vector / vector_shuffle
+  // optimizations have already been done.
+  if (!LegalOperations) return SDValue();
+
+  // (vextract (v4f32 load $addr), c) -> (f32 load $addr+c*size)
+  // (vextract (v4f32 s2v (f32 load $addr)), c) -> (f32 load $addr+c*size)
+  // (vextract (v4f32 shuffle (load $addr), <1,u,u,u>), 0) -> (f32 load $addr)
+
+  if (ConstEltNo) {
+    int Elt = cast<ConstantSDNode>(EltNo)->getZExtValue();
+
     LoadSDNode *LN0 = nullptr;
     const ShuffleVectorSDNode *SVN = nullptr;
     if (ISD::isNormalLoad(InVec.getNode())) {
       LN0 = cast<LoadSDNode>(InVec);
     } else if (InVec.getOpcode() == ISD::SCALAR_TO_VECTOR &&
                InVec.getOperand(0).getValueType() == ExtVT &&
                ISD::isNormalLoad(InVec.getOperand(0).getNode())) {
       // Don't duplicate a load with other uses.
if (ISD::isNormalLoad(InVec.getNode())) {
         if (!InVec.hasOneUse())
           return SDValue();
 
         InVec = InVec.getOperand(0);
       }
       if (ISD::isNormalLoad(InVec.getNode())) {
         LN0 = cast<LoadSDNode>(InVec);
         Elt = (Idx < (int)NumElems) ? Idx : Idx - (int)NumElems;
+        EltNo = DAG.getConstant(Elt, EltNo.getValueType());
       }
     }
 
     // Make sure we found a non-volatile load and the extractelement is
     // the only use.
     if (!LN0 || !LN0->hasNUsesOfValue(1,0) || LN0->isVolatile())
       return SDValue();
 
     // If Idx was -1 above, Elt is going to be -1, so just return undef.
     if (Elt == -1)
       return DAG.getUNDEF(LVT);
 
-    unsigned Align = LN0->getAlignment();
-    if (NewLoad) {
-      // Check the resultant load doesn't need a higher alignment than the
-      // original load.
-      unsigned NewAlign =
-        TLI.getDataLayout()
-            ->getABITypeAlignment(LVT.getTypeForEVT(*DAG.getContext()));
-
-      if (NewAlign > Align || !TLI.isOperationLegalOrCustom(ISD::LOAD, LVT))
-        return SDValue();
-
-      Align = NewAlign;
-    }
-
-    SDValue NewPtr = LN0->getBasePtr();
-    unsigned PtrOff = 0;
-
-    if (Elt) {
-      PtrOff = LVT.getSizeInBits() * Elt / 8;
-      EVT PtrType = NewPtr.getValueType();
-      if (TLI.isBigEndian())
-        PtrOff = VT.getSizeInBits() / 8 - PtrOff;
-      NewPtr = DAG.getNode(ISD::ADD, SDLoc(N), PtrType, NewPtr,
-                           DAG.getConstant(PtrOff, PtrType));
-    }
-
-    // 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;
-    if (NVT.bitsGT(LVT)) {
-      // If the result type of vextract is wider than the load, then issue an
-      // extending load instead.
-      ISD::LoadExtType ExtType = TLI.isLoadExtLegal(ISD::ZEXTLOAD, LVT)
-        ? ISD::ZEXTLOAD : ISD::EXTLOAD;
-      Load = DAG.getExtLoad(ExtType, SDLoc(N), NVT, LN0->getChain(),
-                            NewPtr, LN0->getPointerInfo().getWithOffset(PtrOff),
-                            LVT, LN0->isVolatile(), LN0->isNonTemporal(),
-                            Align, LN0->getTBAAInfo());
-      Chain = Load.getValue(1);
-    } else {
-      Load = DAG.getLoad(LVT, SDLoc(N), LN0->getChain(), NewPtr,
-                         LN0->getPointerInfo().getWithOffset(PtrOff),
-                         LN0->isVolatile(), LN0->isNonTemporal(),
-                         LN0->isInvariant(), Align, LN0->getTBAAInfo());
-      Chain = Load.getValue(1);
-      if (NVT.bitsLT(LVT))
-        Load = DAG.getNode(ISD::TRUNCATE, SDLoc(N), NVT, Load);
-      else
-        Load = DAG.getNode(ISD::BITCAST, SDLoc(N), NVT, Load);
-    }
-    WorkListRemover DeadNodes(*this);
-    SDValue From[] = { SDValue(N, 0), SDValue(LN0,1) };
-    SDValue To[] = { Load, Chain };
-    DAG.ReplaceAllUsesOfValuesWith(From, To, 2);
-    // Since we're explcitly calling ReplaceAllUses, add the new node to the
-    // worklist explicitly as well.
-    AddToWorkList(Load.getNode());
-    AddUsersToWorkList(Load.getNode()); // Add users too
-    // Make sure to revisit this node to clean it up; it will usually be dead.
-    AddToWorkList(N);
-    return SDValue(N, 0);
+    return ReplaceLoad(LVT, LN0, EltNo);
   }
 
   return SDValue();
 }
 
 // Simplify (build_vec (ext )) to (bitcast (build_vec ))
 SDValue DAGCombiner::reduceBuildVecExtToExtBuildVec(SDNode *N) {
   // We perform this optimization post type-legalization because

test/CodeGen/X86/vec_splat.ll

 ; RUN: llc < %s -march=x86 -mcpu=pentium4 -mattr=+sse2 | FileCheck %s -check-prefix=SSE2
 ; RUN: llc < %s -march=x86 -mcpu=pentium4 -mattr=+sse3 | FileCheck %s -check-prefix=SSE3
+; RUN: llc < %s -march=x86-64 -mattr=+avx | FileCheck %s -check-prefix=AVX
 
 define void @test_v4sf(<4 x float>* %P, <4 x float>* %Q, float %X) nounwind {
   %tmp = insertelement <4 x float> zeroinitializer, float %X, i32 0    ; <<4 x float>> [#uses=1]
   %tmp2 = insertelement <4 x float> %tmp, float %X, i32 1    ; <<4 x float>> [#uses=1]
   %tmp4 = insertelement <4 x float> %tmp2, float %X, i32 2    ; <<4 x float>> [#uses=1]
   %tmp6 = insertelement <4 x float> %tmp4, float %X, i32 3    ; <<4 x float>> [#uses=1]
   %tmp8 = load <4 x float>* %Q    ; <<4 x float>> [#uses=1]
   %tmp10 = fmul <4 x float> %tmp8, %tmp6    ; <<4 x float>> [#uses=1]
define void @test_v2sd(<2 x double>* %P, <2 x double>* %Q, double %X) nounwind {
   ret void
 
 ; SSE2-LABEL: test_v2sd:
 ; SSE2: shufpd $0
 
 ; SSE3-LABEL: test_v2sd:
 ; SSE3: movddup
 }
+
+; Fold extract of a load into the load's address computation. This avoids spilling to the stack.
+define <4 x float> @load_extract_splat(<4 x float>* nocapture readonly %ptr, i64 %i, i64 %j) nounwind {
+  %1 = getelementptr inbounds <4 x float>* %ptr, i64 %i
+  %2 = load <4 x float>* %1, align 16
+  %3 = trunc i64 %j to i32
+  %4 = extractelement <4 x float> %2, i32 %3
+  %5 = insertelement <4 x float> undef, float %4, i32 0
+  %6 = insertelement <4 x float> %5, float %4, i32 1
+  %7 = insertelement <4 x float> %6, float %4, i32 2
+  %8 = insertelement <4 x float> %7, float %4, i32 3
+  ret <4 x float> %8
+  
+; AVX-LABEL: load_extract_splat
+; AVX-NOT: rsp
+; AVX: vbroadcastss
+}