[DAGCombine] Improve Lifetime node chains.

[niravd]

Summary

Improve both start and end lifetime nodes chain dependencies.

[courbet]

Nice !

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

15575

Please refactor with the code in visitLIFETIME_END

15585

I have a hard time convincing myself that we can omit this check. Can you add a comment to the code to explain why ?

15594

please clang-format the patch.

[niravd]

Extract copied code into helper and format

[niravd]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

15585

It's not strictly; we could accidentally skip a store because the TokenFactors aren't properly minimized. There needs to be an extra check that the found store is the last semantically meaningful before the LIFETIME_END. I'll factor this out separately from the chain improvement logic.

[courbet]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

1597

Maybe avoid clang-formatting this for consistency with the rest of the function ? (sigh...)

15585

we could accidentally skip a store because the TokenFactors aren't properly minimized.

What happens then ? :)

Can you add this comment to the code so that future readers are aware of the choices made here ?

[niravd]

Now that lifetime node hashing and TokenFactor cleanup has happened, factor out Dead Store improvements as they are no longer needed to prevent regressions.

[courbet]

Thanks, just one nit.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

15591

This can be simplified to:

return ImproveLifetimeNodeChain(N);

[hans]

This was reverted together with r357309 in r357563 as it caused
https://bugs.llvm.org/show_bug.cgi?id=41352

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

private:
     SDValue visitBRCOND(SDNode *N);
     SDValue visitBR_CC(SDNode *N);
     SDValue visitLOAD(SDNode *N);
 
     SDValue replaceStoreChain(StoreSDNode *ST, SDValue BetterChain);
     SDValue replaceStoreOfFPConstant(StoreSDNode *ST);
 
     SDValue visitSTORE(SDNode *N);
+    SDValue visitLIFETIME_START(SDNode *N);
     SDValue visitLIFETIME_END(SDNode *N);
     SDValue visitINSERT_VECTOR_ELT(SDNode *N);
     SDValue visitEXTRACT_VECTOR_ELT(SDNode *N);
     SDValue visitBUILD_VECTOR(SDNode *N);
     SDValue visitCONCAT_VECTORS(SDNode *N);
     SDValue visitEXTRACT_SUBVECTOR(SDNode *N);
     SDValue visitVECTOR_SHUFFLE(SDNode *N);
     SDValue visitSCALAR_TO_VECTOR(SDNode *N);
SDValue DAGCombiner::visit(SDNode *N) {
   case ISD::EXTRACT_SUBVECTOR:  return visitEXTRACT_SUBVECTOR(N);
   case ISD::VECTOR_SHUFFLE:     return visitVECTOR_SHUFFLE(N);
   case ISD::SCALAR_TO_VECTOR:   return visitSCALAR_TO_VECTOR(N);
   case ISD::INSERT_SUBVECTOR:   return visitINSERT_SUBVECTOR(N);
   case ISD::MGATHER:            return visitMGATHER(N);
   case ISD::MLOAD:              return visitMLOAD(N);
   case ISD::MSCATTER:           return visitMSCATTER(N);
   case ISD::MSTORE:             return visitMSTORE(N);
+  case ISD::LIFETIME_START:     return visitLIFETIME_START(N);
   case ISD::LIFETIME_END:       return visitLIFETIME_END(N);
   case ISD::FP_TO_FP16:         return visitFP_TO_FP16(N);
   case ISD::FP16_TO_FP:         return visitFP16_TO_FP(N);
   case ISD::VECREDUCE_FADD:
   case ISD::VECREDUCE_FMUL:
   case ISD::VECREDUCE_ADD:
   case ISD::VECREDUCE_MUL:
   case ISD::VECREDUCE_AND:
   case ISD::VECREDUCE_OR:
SDValue DAGCombiner::visitSTORE(SDNode *N) {
   }
 
   if (SDValue NewSt = splitMergedValStore(ST))
     return NewSt;
 
   return ReduceLoadOpStoreWidth(N);
 }
 
+SDValue DAGCombiner::visitLIFETIME_START(SDNode *N) {
+  auto Chain = N->getOperand(0);
+  auto NewChain = FindBetterChain(N, Chain);
+  if (NewChain != Chain) {
+    SDNode *N2 = DAG.UpdateNodeOperands(N, NewChain, N->getOperand(1));
+    //Make sure users of new n still depend on Chain
+    auto TF = DAG.getNode(ISD::TokenFactor, SDLoc(N2), MVT::Other, Chain, SDValue(N2,0));
+    DAG.ReplaceAllUsesOfValueWith(SDValue(N2,0), TF);
+    AddToWorklist(DAG.UpdateNodeOperands(TF.getNode(), Chain, SDValue(N2,0)));
+    AddToWorklist(N2);
+    return SDValue(N, 0);
+  }
+  return SDValue();
+}
+
 SDValue DAGCombiner::visitLIFETIME_END(SDNode *N) {
+  auto Chain = N->getOperand(0);
+  auto NewChain = FindBetterChain(N, Chain);
+  if (NewChain != Chain) {
+    SDNode *N2 = DAG.UpdateNodeOperands(N, NewChain, N->getOperand(1));
+    //Make sure users of new n still depend on Chain
+    auto TF = DAG.getNode(ISD::TokenFactor, SDLoc(N2), MVT::Other, Chain, SDValue(N2,0));
+    DAG.ReplaceAllUsesOfValueWith(SDValue(N2,0), TF);
+    AddToWorklist(DAG.UpdateNodeOperands(TF.getNode(), Chain, SDValue(N2,0)));
+    AddToWorklist(N2);
+    return SDValue(N, 0);
+  }
+
   const auto *LifetimeEnd = cast<LifetimeSDNode>(N);
   if (!LifetimeEnd->hasOffset())
     return SDValue();
 
   const BaseIndexOffset LifetimeEndBase(N->getOperand(1), SDValue(),
                                         LifetimeEnd->getOffset(), false);
 
   // We walk up the chains to find stores.
   SmallVector<SDValue, 8> Chains = {N->getOperand(0)};
   while (!Chains.empty()) {
-    SDValue Chain = Chains.back();
-    Chains.pop_back();
-    if (!Chain.hasOneUse())
-      continue;
+    SDValue Chain = Chains.pop_back_val();
     switch (Chain.getOpcode()) {
     case ISD::TokenFactor:
+      AddToWorklist(Chain.getNode());
       for (unsigned Nops = Chain.getNumOperands(); Nops;)
         Chains.push_back(Chain.getOperand(--Nops));
       break;
     case ISD::LIFETIME_START:
     case ISD::LIFETIME_END:
       // We can forward past any lifetime start/end that can be proven not to
       // alias the node.
       if (!isAlias(Chain.getNode(), N))
         Chains.push_back(Chain.getOperand(0));
       break;
     case ISD::STORE: {
       StoreSDNode *ST = dyn_cast<StoreSDNode>(Chain);
       if (ST->isVolatile() || ST->isIndexed())
         continue;
       const BaseIndexOffset StoreBase = BaseIndexOffset::match(ST, DAG);
       // If we store purely within object bounds just before its lifetime ends,
       // we can remove the store.
       if (LifetimeEndBase.contains(DAG, LifetimeEnd->getSize() * 8, StoreBase,
                                    ST->getMemoryVT().getStoreSizeInBits())) {
         LLVM_DEBUG(dbgs() << "\nRemoving store:"; StoreBase.dump();
                    dbgs() << "\nwithin LIFETIME_END of : ";
                    LifetimeEndBase.dump(); dbgs() << "\n");
         CombineTo(ST, ST->getChain());
+        AddToWorklist(N);
         return SDValue(N, 0);
       }
     }
     }
   }
   return SDValue();
 }
 

llvm/lib/CodeGen/SelectionDAG/SelectionDAGAddressAnalysis.cpp

bool BaseIndexOffset::contains(const SelectionDAG &DAG, int64_t BitSize,
   // Other starts strictly before *this, it cannot be fully contained.
   //    [-------*this---------]
   // [--Other--]
   return false;
 }
 
 /// Parses tree in Ptr for base, index, offset addresses.
 static BaseIndexOffset matchLSNode(const LSBaseSDNode *N,
-                                       const SelectionDAG &DAG) {
   SDValue Ptr = N->getBasePtr();
 
   // (((B + I*M) + c)) + c ...
   SDValue Base = DAG.getTargetLoweringInfo().unwrapAddress(Ptr);
   SDValue Index = SDValue();
   int64_t Offset = 0;
   bool IsIndexSignExt = false;
 

llvm/test/CodeGen/X86/swap.ll

entry:
   call void @llvm.lifetime.end.p0i8(i64 16, i8* nonnull %part2)
   ret void
 }
 
 define dso_local void @twoallocs(i8* nocapture %a, i8* nocapture %b) local_unnamed_addr {
 ; NOAA-LABEL: twoallocs:
 ; NOAA:       # %bb.0: # %entry
 ; NOAA-NEXT:    vmovups (%rdi), %xmm0
-; NOAA-NEXT:    vmovaps %xmm0, -{{[0-9]+}}(%rsp)
 ; NOAA-NEXT:    vmovups %xmm0, (%rsi)
 ; NOAA-NEXT:    retq
 ;
 ; AA-LABEL: twoallocs:
 ; AA:       # %bb.0: # %entry
 ; AA-NEXT:    vmovups (%rdi), %xmm0
-; AA-NEXT:    vmovaps %xmm0, -{{[0-9]+}}(%rsp)
 ; AA-NEXT:    vmovups %xmm0, (%rsi)
 ; AA-NEXT:    retq
 entry:
   %alloc1 = alloca [16 x i8], align 1
   %alloc2 = alloca [16 x i8], align 1
   %part1 = getelementptr inbounds [16 x i8], [16 x i8]* %alloc1, i64 0, i64 0
   %part2 = getelementptr inbounds [16 x i8], [16 x i8]* %alloc2, i64 0, i64 0
   %part2_alias = getelementptr inbounds [16 x i8], [16 x i8]* %alloc2, i64 0, i64 0
   call void @llvm.lifetime.start.p0i8(i64 16, i8* nonnull %part1)
   call void @llvm.lifetime.start.p0i8(i64 16, i8* nonnull %part2)
   call void @llvm.memcpy.p0i8.p0i8.i64(i8* nonnull align 1 %part2, i8* align 1 %a, i64 16, i1 false)
   call void @llvm.lifetime.end.p0i8(i64 16, i8* nonnull %part1)
   tail call void @llvm.memcpy.p0i8.p0i8.i64(i8* nonnull align 1 %b, i8* align 1 %part2_alias, i64 16, i1 false)
   call void @llvm.lifetime.end.p0i8(i64 16, i8* nonnull %part2)
   ret void
 }
 
 define dso_local void @onealloc_readback_1(i8* nocapture %a, i8* nocapture %b) local_unnamed_addr {
 ; NOAA-LABEL: onealloc_readback_1:
 ; NOAA:       # %bb.0: # %entry
-; NOAA-NEXT:    vmovups (%rdi), %xmm0
-; NOAA-NEXT:    vmovaps %xmm0, -{{[0-9]+}}(%rsp)
 ; NOAA-NEXT:    vmovups (%rsi), %xmm0
-; NOAA-NEXT:    vmovaps %xmm0, -{{[0-9]+}}(%rsp)
-; NOAA-NEXT:    vmovaps -{{[0-9]+}}(%rsp), %xmm0
 ; NOAA-NEXT:    vmovups %xmm0, (%rdi)
 ; NOAA-NEXT:    retq
 ;
 ; AA-LABEL: onealloc_readback_1:
 ; AA:       # %bb.0: # %entry
-; AA-NEXT:    vmovups (%rdi), %xmm0
-; AA-NEXT:    vmovaps %xmm0, -{{[0-9]+}}(%rsp)
 ; AA-NEXT:    vmovups (%rsi), %xmm0
-; AA-NEXT:    vmovaps %xmm0, -{{[0-9]+}}(%rsp)
 ; AA-NEXT:    vmovups %xmm0, (%rdi)
 ; AA-NEXT:    retq
 entry:
   %alloc = alloca [16 x i8], i8 2, align 1
   %part2 = getelementptr inbounds [16 x i8], [16 x i8]* %alloc, i64 0, i64 0
   %part1 = getelementptr inbounds [16 x i8], [16 x i8]* %alloc, i64 1, i64 0
   call void @llvm.lifetime.start.p0i8(i64 16, i8* nonnull %part1)
   call void @llvm.lifetime.start.p0i8(i64 16, i8* nonnull %part2)
   call void @llvm.memcpy.p0i8.p0i8.i64(i8* nonnull align 1 %part1, i8* align 1 %a, i64 16, i1 false)
   call void @llvm.memcpy.p0i8.p0i8.i64(i8* nonnull align 1 %part2, i8* align 1 %b, i64 16, i1 false)
   call void @llvm.lifetime.end.p0i8(i64 16, i8* nonnull %part1)
   tail call void @llvm.memcpy.p0i8.p0i8.i64(i8* nonnull align 1 %a, i8* align 1 %part2, i64 16, i1 false)
   call void @llvm.lifetime.end.p0i8(i64 16, i8* nonnull %part2)
   ret void
 }
 
 define dso_local void @onealloc_readback_2(i8* nocapture %a, i8* nocapture %b) local_unnamed_addr {
 ; NOAA-LABEL: onealloc_readback_2:
 ; NOAA:       # %bb.0: # %entry
-; NOAA-NEXT:    vmovups (%rdi), %xmm0
-; NOAA-NEXT:    vmovaps %xmm0, -{{[0-9]+}}(%rsp)
 ; NOAA-NEXT:    vmovups (%rsi), %xmm0
-; NOAA-NEXT:    vmovaps %xmm0, -{{[0-9]+}}(%rsp)
-; NOAA-NEXT:    vmovaps -{{[0-9]+}}(%rsp), %xmm0
 ; NOAA-NEXT:    vmovups %xmm0, (%rdi)
 ; NOAA-NEXT:    retq
 ;
 ; AA-LABEL: onealloc_readback_2:
 ; AA:       # %bb.0: # %entry
-; AA-NEXT:    vmovups (%rdi), %xmm0
-; AA-NEXT:    vmovaps %xmm0, -{{[0-9]+}}(%rsp)
 ; AA-NEXT:    vmovups (%rsi), %xmm0
-; AA-NEXT:    vmovaps %xmm0, -{{[0-9]+}}(%rsp)
 ; AA-NEXT:    vmovups %xmm0, (%rdi)
 ; AA-NEXT:    retq
 entry:
   %alloc = alloca [16 x i8], i8 2, align 1
   %part1 = getelementptr inbounds [16 x i8], [16 x i8]* %alloc, i64 0, i64 0
   %part2 = getelementptr inbounds [16 x i8], [16 x i8]* %alloc, i64 1, i64 0
   call void @llvm.lifetime.start.p0i8(i64 16, i8* nonnull %part1)
   call void @llvm.lifetime.start.p0i8(i64 16, i8* nonnull %part2)