[X86][AVX] Ensure broadcast loads respect dependencies

[RKSimon]

Summary

To allow broadcast loads of a non-zero'th vector element, lowerVectorShuffleAsBroadcast can replace a load with a new load with an adjusted address, but unfortunately we weren't ensuring that the new load respected the same dependencies.

This patch adds a TokenFactor and updates all dependencies of the old load to reference the new load instead.

Bug found during internal testing.

[mkuper]

Yikes.

Thanks, Simon!
This LGTM, although maybe it can be simpler. If you think this is better as is, feel free to commit.

lib/Target/X86/X86ISelLowering.cpp

8684 ↗

(On Diff #72877)

Why do we need the TokenFactor? The load is only used in the shuffle, right? So why can't we just replace Ld with V completely?
Or do you prefer not to trust the load being DCE'd?

[RKSimon]

Thanks Michael.

lib/Target/X86/X86ISelLowering.cpp

8684 ↗

(On Diff #72877)

Yes, I don't trust the load being dealt with correctly - the problems I had with subvector broadcast all revolved around those cases. I've been wondering whether it'd be cleaner to split the broadcast op into register and memory intrinsic ops to avoid this but for now we just have (ugly) fallback cases as tablegen patterns.

llvm/trunk/lib/Target/X86/X86ISelLowering.cpp

if (V.getOpcode() == ISD::BUILD_VECTOR ||
     LoadSDNode *Ld = cast<LoadSDNode>(BC);
     SDValue BaseAddr = Ld->getOperand(1);
     EVT SVT = BroadcastVT.getScalarType();
     unsigned Offset = BroadcastIdx * SVT.getStoreSize();
     SDValue NewAddr = DAG.getMemBasePlusOffset(BaseAddr, Offset, DL);
     V = DAG.getLoad(SVT, DL, Ld->getChain(), NewAddr,
                     DAG.getMachineFunction().getMachineMemOperand(
                         Ld->getMemOperand(), Offset, SVT.getStoreSize()));
+
+    // Make sure the newly-created LOAD is in the same position as Ld in
+    // terms of dependency. We create a TokenFactor for Ld and V,
+    // and update uses of Ld's output chain to use the TokenFactor.
+    if (Ld->hasAnyUseOfValue(1)) {
+      SDValue NewChain = DAG.getNode(ISD::TokenFactor, DL, MVT::Other,
+                                     SDValue(Ld, 1), SDValue(V.getNode(), 1));
+      DAG.ReplaceAllUsesOfValueWith(SDValue(Ld, 1), NewChain);
+      DAG.UpdateNodeOperands(NewChain.getNode(), SDValue(Ld, 1),
+                             SDValue(V.getNode(), 1));
+    }
   } else if (!BroadcastFromReg) {
     // We can't broadcast from a vector register.
     return SDValue();
   } else if (BroadcastIdx != 0) {
     // We can only broadcast from the zero-element of a vector register,
     // but it can be advantageous to broadcast from the zero-element of a
     // subvector.
     if (!VT.is256BitVector() && !VT.is512BitVector())

llvm/trunk/test/CodeGen/X86/avx-vbroadcast.ll

; X64-NEXT:    retq
   %3 = insertelement <2 x double> %2, double %1, i32 1
   %4 = insertelement <2 x double> undef, double %1, i32 0
   %5 = insertelement <2 x double> %2, double %1, i32 1
   %6 = shufflevector <2 x double> %3, <2 x double> %5, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
   ret <4 x double> %6
 }
 
 ;
-; FIXME: When VBROADCAST replaces an existing load, ensure it still respects lifetime dependencies.
+; When VBROADCAST replaces an existing load, ensure it still respects lifetime dependencies.
 ;
 define float @broadcast_lifetime() nounwind {
 ; X32-LABEL: broadcast_lifetime:
 ; X32:       ## BB#0:
 ; X32-NEXT:    pushl %esi
-; X32-NEXT:    subl $40, %esp
+; X32-NEXT:    subl $56, %esp
 ; X32-NEXT:    leal {{[0-9]+}}(%esp), %esi
 ; X32-NEXT:    movl %esi, (%esp)
 ; X32-NEXT:    calll _gfunc
+; X32-NEXT:    vbroadcastss {{[0-9]+}}(%esp), %xmm0
+; X32-NEXT:    vmovaps %xmm0, {{[0-9]+}}(%esp) ## 16-byte Spill
 ; X32-NEXT:    movl %esi, (%esp)
 ; X32-NEXT:    calll _gfunc
 ; X32-NEXT:    vbroadcastss {{[0-9]+}}(%esp), %xmm0
-; X32-NEXT:    vbroadcastss {{[0-9]+}}(%esp), %xmm1
-; X32-NEXT:    vsubss %xmm0, %xmm1, %xmm0
+; X32-NEXT:    vsubss {{[0-9]+}}(%esp), %xmm0, %xmm0 ## 16-byte Folded Reload
 ; X32-NEXT:    vmovss %xmm0, {{[0-9]+}}(%esp)
 ; X32-NEXT:    flds {{[0-9]+}}(%esp)
-; X32-NEXT:    addl $40, %esp
+; X32-NEXT:    addl $56, %esp
 ; X32-NEXT:    popl %esi
 ; X32-NEXT:    retl
 ;
 ; X64-LABEL: broadcast_lifetime:
 ; X64:       ## BB#0:
-; X64-NEXT:    subq $24, %rsp
+; X64-NEXT:    subq $40, %rsp
 ; X64-NEXT:    movq %rsp, %rdi
 ; X64-NEXT:    callq _gfunc
+; X64-NEXT:    vbroadcastss {{[0-9]+}}(%rsp), %xmm0
+; X64-NEXT:    vmovaps %xmm0, {{[0-9]+}}(%rsp) ## 16-byte Spill
 ; X64-NEXT:    movq %rsp, %rdi
 ; X64-NEXT:    callq _gfunc
 ; X64-NEXT:    vbroadcastss {{[0-9]+}}(%rsp), %xmm0
-; X64-NEXT:    vbroadcastss {{[0-9]+}}(%rsp), %xmm1
-; X64-NEXT:    vsubss %xmm0, %xmm1, %xmm0
-; X64-NEXT:    addq $24, %rsp
+; X64-NEXT:    vsubss {{[0-9]+}}(%rsp), %xmm0, %xmm0 ## 16-byte Folded Reload
+; X64-NEXT:    addq $40, %rsp
 ; X64-NEXT:    retq
   %1 = alloca <4 x float>, align 16
   %2 = alloca <4 x float>, align 16
   %3 = bitcast <4 x float>* %1 to i8*
   %4 = bitcast <4 x float>* %2 to i8*
 
   call void @llvm.lifetime.start(i64 16, i8* %3)
   call void @gfunc(<4 x float>* %1)