[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

15596

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

15596

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

[niravd]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

15596

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?

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 Val = 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
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();
 
   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;
 
+      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.

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=+avx | FileCheck %s --check-prefixes=X64,X64-AVX
 
 ;
 ; PR42123
 ;
 
-; FIXME: AVX doesn't retain NT flag on store.
-; Should be VMOVNTPS ymm.
+; AVX 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)
 ; X86-NEXT:    movntps %xmm1, 16(%eax)
 ; X86-NEXT:    retl
 ;
 ; X64-SSE-LABEL: merge_2_v4f32_align32:
 ; 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:    movntps %xmm1, 16(%rsi)
 ; X64-SSE-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:    vmovntps %ymm0, (%rsi)
 ; X64-AVX-NEXT:    vzeroupper
 ; 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, !nontemporal !0
   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.
+; AVX must be kept seperate as VMOVNTPS xmm + VMOVAPS xmm.
 define void @merge_2_v4f32_align32_mix(<4 x float>* %a0, <4 x float>* %a1)  {
 ; X86-LABEL: merge_2_v4f32_align32_mix:
 ; 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:
 ; 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:
 ; 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