[DAGCombine] MergeConsecutiveStores - improve non-temporal load\store handling (PR42123)

[RKSimon]

Summary

This patch is the first step towards ensuring MergeConsecutiveStores correctly handles non-temporal stores:

1 - When merging load\stores we must ensure that they all have the same non-temporal flag. This is unlikely to occur, but can in strange cases where we're storing at the end of one page and the beginning of another.

2 - The merged load\store node must retain the non-temporal flag.

After this I need to improve alignment handling - in most cases we shouldn't merge nt-ops unless they are naturally aligned - either we just prevent it in all cases or we refactor canMergeStoresTo\allowsMemoryAccess to include access to the alignment and MMO flags as well?

[niravd]

Don't we need to do something for non-temporal loads as well?

[RKSimon]

In D62910#1531055, @niravd wrote:

Don't we need to do something for non-temporal loads as well?

Yes we will - do you want that as part of this patch?

[niravd]

Yes, I think that keeping the flag stuff to an atomic makes the most sense. We should probably make sure we're doing the right thing for the dereferenceable and invariant flags as well.

[RKSimon]

Add nt-load handling as well

[niravd]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

15600 ↗

(On Diff #203205)

Looks like we were erroneously transferring dereferencable from load to store on merge before. I think we need to do a store version of the isDereferencable calculation or we'll only preserve it on truncstores.

[RKSimon]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

15600 ↗

(On Diff #203205)

Were we? The getStore() didn't use any MMO flags before.

[niravd]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

15600 ↗

(On Diff #203205)

Oh right! We're still being a bit weird about Dereferencable handling and that seems worth more of a look, but if we're being consistent with the past, there's no sense in delaying this patch.

[RKSimon]

Thanks @niravd - any thoughts on how to tackle under-aligned nt-ops? My gut feeling is that just bailing on the merge if they aren't naturally aligned would be fine - do you know of any nt vector instructions that can handle non-aligned pointers?

llvm/trunk/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if (IsLoadSrc) {
     LBasePtr = BaseIndexOffset::match(Ld, DAG);
     LoadVT = Ld->getMemoryVT();
     // Load and store should be the same type.
     if (MemVT != LoadVT)
       return;
     // Loads must only have one use.
     if (!Ld->hasNUsesOfValue(1, 0))
       return;
-    // The memory operands must not be volatile.
+    // The memory operands must not be volatile/indexed.
     if (Ld->isVolatile() || Ld->isIndexed())
       return;
   }
   auto CandidateMatch = [&](StoreSDNode *Other, BaseIndexOffset &Ptr,
                             int64_t &Offset) -> bool {
+    // The memory operands must not be volatile/indexed.
     if (Other->isVolatile() || Other->isIndexed())
       return false;
+    // Don't mix temporal stores with non-temporal stores.
+    if (St->isNonTemporal() != Other->isNonTemporal())
+      return false;
     SDValue OtherBC = peekThroughBitcasts(Other->getValue());
     // Allow merging constants of different types as integers.
     bool NoTypeMatch = (MemVT.isInteger()) ? !MemVT.bitsEq(Other->getMemoryVT())
                                            : Other->getMemoryVT() != MemVT;
     if (IsLoadSrc) {
       if (NoTypeMatch)
         return false;
       // The Load's Base Ptr must also match
       if (LoadSDNode *OtherLd = dyn_cast<LoadSDNode>(OtherBC)) {
-        auto LPtr = BaseIndexOffset::match(OtherLd, DAG);
+        BaseIndexOffset LPtr = BaseIndexOffset::match(OtherLd, DAG);
         if (LoadVT != OtherLd->getMemoryVT())
           return false;
         // Loads must only have one use.
         if (!OtherLd->hasNUsesOfValue(1, 0))
           return false;
-        // The memory operands must not be volatile.
+        // The memory operands must not be volatile/indexed.
         if (OtherLd->isVolatile() || OtherLd->isIndexed())
           return false;
+        // Don't mix temporal loads with non-temporal loads.
+        if (cast<LoadSDNode>(Val)->isNonTemporal() != OtherLd->isNonTemporal())
+          return false;
         if (!(LBasePtr.equalBaseIndex(LPtr, DAG)))
           return false;
       } else
         return false;
     }
     if (IsConstantSrc) {
       if (NoTypeMatch)
         return false;
bool DAGCombiner::MergeConsecutiveStores(StoreSDNode *St) {
   // Perform an early exit check. Do not bother looking at stored values that
   // are not constants, loads, or extracted vector elements.
   SDValue StoredVal = peekThroughBitcasts(St->getValue());
   bool IsLoadSrc = isa<LoadSDNode>(StoredVal);
   bool IsConstantSrc = isa<ConstantSDNode>(StoredVal) ||
                        isa<ConstantFPSDNode>(StoredVal);
   bool IsExtractVecSrc = (StoredVal.getOpcode() == ISD::EXTRACT_VECTOR_ELT ||
                           StoredVal.getOpcode() == ISD::EXTRACT_SUBVECTOR);
+  bool IsNonTemporalStore = St->isNonTemporal();
+  bool IsNonTemporalLoad =
+      IsLoadSrc && cast<LoadSDNode>(StoredVal)->isNonTemporal();
 
   if (!IsConstantSrc && !IsLoadSrc && !IsExtractVecSrc)
     return false;
 
   SmallVector<MemOpLink, 8> StoreNodes;
   SDNode *RootNode;
   // Find potential store merge candidates by searching through chain sub-DAG
   getStoreMergeCandidates(St, StoreNodes, RootNode);
while (NumConsecutiveStores >= 2 && LoadNodes.size() >= 2) {
       SDLoc StoreDL(StoreNodes[0].MemNode);
 
       // The merged loads are required to have the same incoming chain, so
       // using the first's chain is acceptable.
 
       SDValue NewStoreChain = getMergeStoreChains(StoreNodes, NumElem);
       AddToWorklist(NewStoreChain.getNode());
 
-      MachineMemOperand::Flags MMOFlags =
+      MachineMemOperand::Flags LdMMOFlags =
           isDereferenceable ? MachineMemOperand::MODereferenceable
                             : MachineMemOperand::MONone;
+      if (IsNonTemporalLoad)
+        LdMMOFlags |= MachineMemOperand::MONonTemporal;
+
+      MachineMemOperand::Flags StMMOFlags =
+          IsNonTemporalStore ? MachineMemOperand::MONonTemporal
+                             : MachineMemOperand::MONone;
 
       SDValue NewLoad, NewStore;
       if (UseVectorTy || !DoIntegerTruncate) {
         NewLoad =
             DAG.getLoad(JointMemOpVT, LoadDL, FirstLoad->getChain(),
                         FirstLoad->getBasePtr(), FirstLoad->getPointerInfo(),
-                        FirstLoadAlign, MMOFlags);
+                        FirstLoadAlign, LdMMOFlags);
         NewStore = DAG.getStore(
             NewStoreChain, StoreDL, NewLoad, FirstInChain->getBasePtr(),
-            FirstInChain->getPointerInfo(), FirstStoreAlign);
+            FirstInChain->getPointerInfo(), FirstStoreAlign, StMMOFlags);
       } else { // This must be the truncstore/extload case
         EVT ExtendedTy =
             TLI.getTypeToTransformTo(*DAG.getContext(), JointMemOpVT);
         NewLoad = DAG.getExtLoad(ISD::EXTLOAD, LoadDL, ExtendedTy,
                                  FirstLoad->getChain(), FirstLoad->getBasePtr(),
                                  FirstLoad->getPointerInfo(), JointMemOpVT,
-                                 FirstLoadAlign, MMOFlags);
+                                 FirstLoadAlign, LdMMOFlags);
         NewStore = DAG.getTruncStore(NewStoreChain, StoreDL, NewLoad,
                                      FirstInChain->getBasePtr(),
                                      FirstInChain->getPointerInfo(),
                                      JointMemOpVT, FirstInChain->getAlignment(),
                                      FirstInChain->getMemOperand()->getFlags());
       }
 
       // Transfer chain users from old loads to the new load.

llvm/trunk/test/CodeGen/X86/merge-consecutive-stores-nt.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc < %s -mtriple=i686-unknown-unknown -mattr=+sse2 | FileCheck %s --check-prefixes=X86
 ; RUN: llc < %s -mtriple=x86_64-unknown-unknown | FileCheck %s --check-prefixes=X64,X64-SSE,X64-SSE2
 ; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+sse4a | FileCheck %s --check-prefixes=X64,X64-SSE,X64-SSE4A
 ; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+sse4.1 | FileCheck %s --check-prefixes=X64,X64-SSE,X64-SSE41
 ; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+avx  | FileCheck %s --check-prefixes=X64,X64-AVX,X64-AVX1
 ; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+avx2 | FileCheck %s --check-prefixes=X64,X64-AVX,X64-AVX2
 
 ;
 ; PR42123
 ;
 
-; FIXME: AVX doesn't retain NT flag on load/store.
-; AVX1 load should be 2 x VMOVNTDQA xmm.
-; AVX2 load should be VMOVNTDQA ymm.
-; AVX store should be VMOVNTPS ymm.
 define void @merge_2_v4f32_align32(<4 x float>* %a0, <4 x float>* %a1)  {
 ; X86-LABEL: merge_2_v4f32_align32:
 ; X86:       # %bb.0:
 ; X86-NEXT:    movl {{[0-9]+}}(%esp), %eax
 ; X86-NEXT:    movl {{[0-9]+}}(%esp), %ecx
 ; X86-NEXT:    movaps (%ecx), %xmm0
 ; X86-NEXT:    movaps 16(%ecx), %xmm1
 ; X86-NEXT:    movntps %xmm0, (%eax)

 ; X64-SSE41-LABEL: merge_2_v4f32_align32:
 ; X64-SSE41:       # %bb.0:
 ; X64-SSE41-NEXT:    movntdqa (%rdi), %xmm0
 ; X64-SSE41-NEXT:    movntdqa 16(%rdi), %xmm1
 ; X64-SSE41-NEXT:    movntdq %xmm0, (%rsi)
 ; X64-SSE41-NEXT:    movntdq %xmm1, 16(%rsi)
 ; X64-SSE41-NEXT:    retq
 ;
-; X64-AVX-LABEL: merge_2_v4f32_align32:
-; X64-AVX:       # %bb.0:
-; X64-AVX-NEXT:    vmovaps (%rdi), %ymm0
-; X64-AVX-NEXT:    vmovaps %ymm0, (%rsi)
-; X64-AVX-NEXT:    vzeroupper
-; X64-AVX-NEXT:    retq
+; X64-AVX1-LABEL: merge_2_v4f32_align32:
+; X64-AVX1:       # %bb.0:
+; X64-AVX1-NEXT:    vmovntdqa (%rdi), %xmm0
+; X64-AVX1-NEXT:    vmovntdqa 16(%rdi), %xmm1
+; X64-AVX1-NEXT:    vmovntdq %xmm1, 16(%rsi)
+; X64-AVX1-NEXT:    vmovntdq %xmm0, (%rsi)
+; X64-AVX1-NEXT:    retq
+;
+; X64-AVX2-LABEL: merge_2_v4f32_align32:
+; X64-AVX2:       # %bb.0:
+; X64-AVX2-NEXT:    vmovntdqa (%rdi), %ymm0
+; X64-AVX2-NEXT:    vmovntdq %ymm0, (%rsi)
+; X64-AVX2-NEXT:    vzeroupper
+; X64-AVX2-NEXT:    retq
   %1 = getelementptr inbounds <4 x float>, <4 x float>* %a0, i64 1, i64 0
   %2 = bitcast float* %1 to <4 x float>*
   %3 = load <4 x float>, <4 x float>* %a0, align 32, !nontemporal !0
   %4 = load <4 x float>, <4 x float>* %2, align 16, !nontemporal !0
   %5 = getelementptr inbounds <4 x float>, <4 x float>* %a1, i64 1, i64 0
   %6 = bitcast float* %5 to <4 x float>*
   store <4 x float> %3, <4 x float>* %a1, align 32, !nontemporal !0
   store <4 x float> %4, <4 x float>* %6, align 16, !nontemporal !0
   ret void
 }
 
-; FIXME: shouldn't attempt to merge nt and non-nt loads even if aligned.
-; Must be kept seperate as VMOVNTDQA xmm + VMOVDQA xmm.
+; Don't merge nt and non-nt loads even if aligned.
 define void @merge_2_v4f32_align32_mix_ntload(<4 x float>* %a0, <4 x float>* %a1)  {
 ; X86-LABEL: merge_2_v4f32_align32_mix_ntload:
 ; X86:       # %bb.0:
 ; X86-NEXT:    movl {{[0-9]+}}(%esp), %eax
 ; X86-NEXT:    movl {{[0-9]+}}(%esp), %ecx
 ; X86-NEXT:    movaps (%ecx), %xmm0
 ; X86-NEXT:    movaps 16(%ecx), %xmm1
 ; X86-NEXT:    movaps %xmm0, (%eax)

 ; X64-SSE41-NEXT:    movntdqa (%rdi), %xmm0
 ; X64-SSE41-NEXT:    movaps 16(%rdi), %xmm1
 ; X64-SSE41-NEXT:    movdqa %xmm0, (%rsi)
 ; X64-SSE41-NEXT:    movaps %xmm1, 16(%rsi)
 ; X64-SSE41-NEXT:    retq
 ;
 ; X64-AVX-LABEL: merge_2_v4f32_align32_mix_ntload:
 ; X64-AVX:       # %bb.0:
-; X64-AVX-NEXT:    vmovaps (%rdi), %ymm0
-; X64-AVX-NEXT:    vmovaps %ymm0, (%rsi)
-; X64-AVX-NEXT:    vzeroupper
+; X64-AVX-NEXT:    vmovntdqa (%rdi), %xmm0
+; X64-AVX-NEXT:    vmovaps 16(%rdi), %xmm1
+; X64-AVX-NEXT:    vmovdqa %xmm0, (%rsi)
+; X64-AVX-NEXT:    vmovaps %xmm1, 16(%rsi)
 ; X64-AVX-NEXT:    retq
   %1 = getelementptr inbounds <4 x float>, <4 x float>* %a0, i64 1, i64 0
   %2 = bitcast float* %1 to <4 x float>*
   %3 = load <4 x float>, <4 x float>* %a0, align 32, !nontemporal !0
   %4 = load <4 x float>, <4 x float>* %2, align 16
   %5 = getelementptr inbounds <4 x float>, <4 x float>* %a1, i64 1, i64 0
   %6 = bitcast float* %5 to <4 x float>*
   store <4 x float> %3, <4 x float>* %a1, align 32
   store <4 x float> %4, <4 x float>* %6, align 16
   ret void
 }
 
-; FIXME: shouldn't attempt to merge nt and non-nt stores even if aligned.
-; Must be kept seperate as VMOVNTPS xmm + VMOVAPS xmm.
+; Don't merge nt and non-nt stores even if aligned.
 define void @merge_2_v4f32_align32_mix_ntstore(<4 x float>* %a0, <4 x float>* %a1)  {
 ; X86-LABEL: merge_2_v4f32_align32_mix_ntstore:
 ; X86:       # %bb.0:
 ; X86-NEXT:    movl {{[0-9]+}}(%esp), %eax
 ; X86-NEXT:    movl {{[0-9]+}}(%esp), %ecx
 ; X86-NEXT:    movaps (%ecx), %xmm0
 ; X86-NEXT:    movaps 16(%ecx), %xmm1
 ; X86-NEXT:    movntps %xmm0, (%eax)
 ; X86-NEXT:    movaps %xmm1, 16(%eax)
 ; X86-NEXT:    retl
 ;
 ; X64-SSE-LABEL: merge_2_v4f32_align32_mix_ntstore:
 ; X64-SSE:       # %bb.0:
 ; X64-SSE-NEXT:    movaps (%rdi), %xmm0
 ; X64-SSE-NEXT:    movaps 16(%rdi), %xmm1
 ; X64-SSE-NEXT:    movntps %xmm0, (%rsi)
 ; X64-SSE-NEXT:    movaps %xmm1, 16(%rsi)
 ; X64-SSE-NEXT:    retq
 ;
 ; X64-AVX-LABEL: merge_2_v4f32_align32_mix_ntstore:
 ; X64-AVX:       # %bb.0:
-; X64-AVX-NEXT:    vmovaps (%rdi), %ymm0
-; X64-AVX-NEXT:    vmovaps %ymm0, (%rsi)
-; X64-AVX-NEXT:    vzeroupper
+; X64-AVX-NEXT:    vmovaps (%rdi), %xmm0
+; X64-AVX-NEXT:    vmovaps 16(%rdi), %xmm1
+; X64-AVX-NEXT:    vmovntps %xmm0, (%rsi)
+; X64-AVX-NEXT:    vmovaps %xmm1, 16(%rsi)
 ; X64-AVX-NEXT:    retq
   %1 = getelementptr inbounds <4 x float>, <4 x float>* %a0, i64 1, i64 0
   %2 = bitcast float* %1 to <4 x float>*
   %3 = load <4 x float>, <4 x float>* %a0, align 32
   %4 = load <4 x float>, <4 x float>* %2, align 16
   %5 = getelementptr inbounds <4 x float>, <4 x float>* %a1, i64 1, i64 0
   %6 = bitcast float* %5 to <4 x float>*
   store <4 x float> %3, <4 x float>* %a1, align 32, !nontemporal !0
   store <4 x float> %4, <4 x float>* %6, align 16
   ret void
 }
 
-; FIXME: AVX can't perform NT-load-ymm on 16-byte aligned memory.
+; FIXME: AVX2 can't perform NT-load-ymm on 16-byte aligned memory.
 ; Must be kept seperate as VMOVNTDQA xmm.
 define void @merge_2_v4f32_align16_ntload(<4 x float>* %a0, <4 x float>* %a1)  {
 ; X86-LABEL: merge_2_v4f32_align16_ntload:
 ; X86:       # %bb.0:
 ; X86-NEXT:    movl {{[0-9]+}}(%esp), %eax
 ; X86-NEXT:    movl {{[0-9]+}}(%esp), %ecx
 ; X86-NEXT:    movaps (%ecx), %xmm0
 ; X86-NEXT:    movaps 16(%ecx), %xmm1

 ; X64-SSE41-LABEL: merge_2_v4f32_align16_ntload:
 ; X64-SSE41:       # %bb.0:
 ; X64-SSE41-NEXT:    movntdqa (%rdi), %xmm0
 ; X64-SSE41-NEXT:    movntdqa 16(%rdi), %xmm1
 ; X64-SSE41-NEXT:    movdqa %xmm0, (%rsi)
 ; X64-SSE41-NEXT:    movdqa %xmm1, 16(%rsi)
 ; X64-SSE41-NEXT:    retq
 ;
-; X64-AVX-LABEL: merge_2_v4f32_align16_ntload:
-; X64-AVX:       # %bb.0:
-; X64-AVX-NEXT:    vmovups (%rdi), %ymm0
-; X64-AVX-NEXT:    vmovups %ymm0, (%rsi)
-; X64-AVX-NEXT:    vzeroupper
-; X64-AVX-NEXT:    retq
+; X64-AVX1-LABEL: merge_2_v4f32_align16_ntload:
+; X64-AVX1:       # %bb.0:
+; X64-AVX1-NEXT:    vmovntdqa (%rdi), %xmm0
+; X64-AVX1-NEXT:    vmovntdqa 16(%rdi), %xmm1
+; X64-AVX1-NEXT:    vmovdqa %xmm1, 16(%rsi)
+; X64-AVX1-NEXT:    vmovdqa %xmm0, (%rsi)
+; X64-AVX1-NEXT:    retq
+;
+; X64-AVX2-LABEL: merge_2_v4f32_align16_ntload:
+; X64-AVX2:       # %bb.0:
+; X64-AVX2-NEXT:    vmovups (%rdi), %ymm0
+; X64-AVX2-NEXT:    vmovups %ymm0, (%rsi)
+; X64-AVX2-NEXT:    vzeroupper
+; X64-AVX2-NEXT:    retq
   %1 = getelementptr inbounds <4 x float>, <4 x float>* %a0, i64 1, i64 0
   %2 = bitcast float* %1 to <4 x float>*
   %3 = load <4 x float>, <4 x float>* %a0, align 16, !nontemporal !0
   %4 = load <4 x float>, <4 x float>* %2, align 16, !nontemporal !0
   %5 = getelementptr inbounds <4 x float>, <4 x float>* %a1, i64 1, i64 0
   %6 = bitcast float* %5 to <4 x float>*
   store <4 x float> %3, <4 x float>* %a1, align 16
   store <4 x float> %4, <4 x float>* %6, align 16