[SelectionDAG] Re-calculate scoped AA metadata when merging stores.

Reviewed By: jeroen.dobbelaere

Differential Revision: https://reviews.llvm.org/D102821

Author: darkbuck

llvm/include/llvm/IR/Metadata.h

@@ -711,6 +711,11 @@ struct AAMDNodes {
   /// Given two sets of AAMDNodes applying to potentially different locations,
   /// determine the best AAMDNodes that apply to both.
   AAMDNodes merge(const AAMDNodes &Other) const;
+
+  /// Determine the best AAMDNodes after concatenating two different locations
+  /// together. Different from `merge`, where different locations should
+  /// overlap each other, `concat` puts non-overlapping locations together.
+  AAMDNodes concat(const AAMDNodes &Other) const;
 };
 
 // Specialize DenseMapInfo for AAMDNodes.

llvm/lib/Analysis/TypeBasedAliasAnalysis.cpp

@@ -530,6 +530,14 @@ AAMDNodes AAMDNodes::merge(const AAMDNodes &Other) const {
   return Result;
 }
 
+AAMDNodes AAMDNodes::concat(const AAMDNodes &Other) const {
+  AAMDNodes Result;
+  Result.TBAA = Result.TBAAStruct = nullptr;
+  Result.Scope = MDNode::getMostGenericAliasScope(Scope, Other.Scope);
+  Result.NoAlias = MDNode::intersect(NoAlias, Other.NoAlias);
+  return Result;
+}
+
 AAMDNodes Instruction::getAAMetadata() const {
   AAMDNodes Result;
   Result.TBAA = getMetadata(LLVMContext::MD_tbaa);

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

@@ -16970,6 +16970,22 @@ bool DAGCombiner::mergeStoresOfConstantsOrVecElts(
   unsigned SizeInBits = NumStores * ElementSizeBits;
   unsigned NumMemElts = MemVT.isVector() ? MemVT.getVectorNumElements() : 1;
 
+  Optional<MachineMemOperand::Flags> Flags;
+  AAMDNodes AAInfo;
+  for (unsigned I = 0; I != NumStores; ++I) {
+    StoreSDNode *St = cast<StoreSDNode>(StoreNodes[I].MemNode);
+    if (!Flags) {
+      Flags = St->getMemOperand()->getFlags();
+      AAInfo = St->getAAInfo();
+      continue;
+    }
+    // Skip merging if there's an inconsistent flag.
+    if (Flags != St->getMemOperand()->getFlags())
+      return false;
+    // Concatenate AA metadata.
+    AAInfo = AAInfo.concat(St->getAAInfo());
+  }
+
   EVT StoreTy;
   if (UseVector) {
     unsigned Elts = NumStores * NumMemElts;
@@ -17087,9 +17103,9 @@ bool DAGCombiner::mergeStoresOfConstantsOrVecElts(
   // make sure we use trunc store if it's necessary to be legal.
   SDValue NewStore;
   if (!UseTrunc) {
-    NewStore =
-        DAG.getStore(NewChain, DL, StoredVal, FirstInChain->getBasePtr(),
-                     FirstInChain->getPointerInfo(), FirstInChain->getAlign());
+    NewStore = DAG.getStore(NewChain, DL, StoredVal, FirstInChain->getBasePtr(),
+                            FirstInChain->getPointerInfo(),
+                            FirstInChain->getAlign(), Flags.getValue(), AAInfo);
   } else { // Must be realized as a trunc store
     EVT LegalizedStoredValTy =
         TLI.getTypeToTransformTo(*DAG.getContext(), StoredVal.getValueType());
@@ -17101,7 +17117,7 @@ bool DAGCombiner::mergeStoresOfConstantsOrVecElts(
     NewStore = DAG.getTruncStore(
         NewChain, DL, ExtendedStoreVal, FirstInChain->getBasePtr(),
         FirstInChain->getPointerInfo(), StoredVal.getValueType() /*TVT*/,
-        FirstInChain->getAlign(), FirstInChain->getMemOperand()->getFlags());
+        FirstInChain->getAlign(), Flags.getValue(), AAInfo);
   }
 
   // Replace all merged stores with the new store.

llvm/test/CodeGen/AArch64/merge-scoped-aa-store.ll

@@ -0,0 +1,58 @@
+; RUN: llc %s -mtriple=aarch64-eabi -stop-after=finalize-isel -o - | FileCheck --check-prefix=MIR %s
+
+; Ensure the scoped AA metadata is still retained after store merging.
+
+; MIR-DAG: ![[DOMAIN:[0-9]+]] = distinct !{!{{[0-9]+}}, !"domain"}
+; MIR-DAG: ![[SCOPE0:[0-9]+]] = distinct !{!{{[0-9]+}}, ![[DOMAIN]], !"scope 0"}
+; MIR-DAG: ![[SCOPE1:[0-9]+]] = distinct !{!{{[0-9]+}}, ![[DOMAIN]], !"scope 1"}
+; MIR-DAG: ![[SCOPE2:[0-9]+]] = distinct !{!{{[0-9]+}}, ![[DOMAIN]], !"scope 2"}
+; MIR-DAG: ![[SCOPE3:[0-9]+]] = distinct !{!{{[0-9]+}}, ![[DOMAIN]], !"scope 3"}
+; MIR-DAG: ![[SET0:[0-9]+]] = !{![[SCOPE0]]}
+; MIR-DAG: ![[SET1:[0-9]+]] = !{![[SCOPE1]]}
+; MIR-DAG: ![[SET2:[0-9]+]] = !{![[SCOPE2]]}
+; MIR-DAG: ![[SET3:[0-9]+]] = !{![[SCOPE3]]}
+
+define void @blam0(<3 x float>* %g0, <3 x float>* %g1) {
+; MIR-LABEL: name: blam0
+; MIR: LDRDui %0, 0 :: (load (s64) from %ir.g0, align 4, !alias.scope ![[SET0]], !noalias ![[SET1]])
+; MIR: STRDui killed %5, %1, 0 :: (store (s64) into %ir.tmp41, align 4, !alias.scope ![[SET1]], !noalias ![[SET0]])
+  %tmp4 = getelementptr inbounds <3 x float>, <3 x float>* %g1, i64 0, i64 0
+  %tmp5 = load <3 x float>, <3 x float>* %g0, align 4, !alias.scope !0, !noalias !1
+  %tmp6 = extractelement <3 x float> %tmp5, i64 0
+  store float %tmp6, float* %tmp4, align 4, !alias.scope !1, !noalias !0
+  %tmp7 = getelementptr inbounds float, float* %tmp4, i64 1
+  %tmp8 = load <3 x float>, <3 x float>* %g0, align 4, !alias.scope !0, !noalias !1
+  %tmp9 = extractelement <3 x float> %tmp8, i64 1
+  store float %tmp9, float* %tmp7, align 4, !alias.scope !1, !noalias !0
+  ret void;
+}
+
+; Ensure new scoped AA metadata are calculated after merging stores.
+define void @blam1(<3 x float>* %g0, <3 x float>* %g1) {
+; MIR-LABEL: name: blam1
+; MIR: machineMetadataNodes:
+; MIR-DAG: ![[MMSET0:[0-9]+]] = !{![[SCOPE2]], ![[SCOPE1]]}
+; MIR-DAG: ![[MMSET1:[0-9]+]] = !{}
+; MIR: body:
+; MIR: LDRDui %0, 0 :: (load (s64) from %ir.g0, align 4, !alias.scope ![[SET0]], !noalias ![[SET1]])
+; MIR: STRDui killed %5, %1, 0 :: (store (s64) into %ir.tmp41, align 4, !alias.scope ![[MMSET0]], !noalias ![[MMSET1]])
+  %tmp4 = getelementptr inbounds <3 x float>, <3 x float>* %g1, i64 0, i64 0
+  %tmp5 = load <3 x float>, <3 x float>* %g0, align 4, !alias.scope !0, !noalias !1
+  %tmp6 = extractelement <3 x float> %tmp5, i64 0
+  store float %tmp6, float* %tmp4, align 4, !alias.scope !1, !noalias !0
+  %tmp7 = getelementptr inbounds float, float* %tmp4, i64 1
+  %tmp8 = load <3 x float>, <3 x float>* %g0, align 4, !alias.scope !0, !noalias !1
+  %tmp9 = extractelement <3 x float> %tmp8, i64 1
+  store float %tmp9, float* %tmp7, align 4, !alias.scope !5, !noalias !6
+  ret void;
+}
+
+!0 = !{!2}
+!1 = !{!3}
+!2 = !{!2, !4, !"scope 0"}
+!3 = !{!3, !4, !"scope 1"}
+!4 = !{!4, !"domain"}
+!5 = !{!7}
+!6 = !{!8}
+!7 = !{!7, !4, !"scope 2"}
+!8 = !{!8, !4, !"scope 3"}