[Metadata] Add TBAA struct metadata to `AAMDNode`

Summary:
Make `AAMDNodes`' `getAAMetadata()` and `setAAMetadata()` to take `!tbaa.struct`
into account as well as `!tbaa`. This impacts llvm.org/pr42022.
This is a temprorary fix needed to keep `!tbaa.struct` tag by SROA pass.
New field `TBAAStruct` should be deleted when `!tbaa` tag replaces `!tbaa.struct`.
Merging two `!tbaa.struct`'s to one is conservatively considered to be `nullptr`
(giving `MayAlias`) -- this could be enhanced, but relying on the said future
replacement.

Reviewers: RKSimon, spatel, vporpo

Subscribers: hiraditya, kosarev, llvm-commits

Tags: #llvm

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

llvm/include/llvm/IR/Metadata.h

@@ -642,26 +642,32 @@ public:
 /// A collection of metadata nodes that might be associated with a
 /// memory access used by the alias-analysis infrastructure.
 struct AAMDNodes {
-  explicit AAMDNodes(MDNode *T = nullptr, MDNode *S = nullptr,
+  explicit AAMDNodes() = default;
-                     MDNode *N = nullptr)
+  explicit AAMDNodes(MDNode *T, MDNode *TS, MDNode *S, MDNode *N)
-      : TBAA(T), Scope(S), NoAlias(N) {}
+      : TBAA(T), TBAAStruct(TS), Scope(S), NoAlias(N) {}
 
   bool operator==(const AAMDNodes &A) const {
-    return TBAA == A.TBAA && Scope == A.Scope && NoAlias == A.NoAlias;
+    return TBAA == A.TBAA && TBAAStruct == A.TBAAStruct && Scope == A.Scope &&
+           NoAlias == A.NoAlias;
   }
 
   bool operator!=(const AAMDNodes &A) const { return !(*this == A); }
 
-  explicit operator bool() const { return TBAA || Scope || NoAlias; }
+  explicit operator bool() const {
+    return TBAA || TBAAStruct || Scope || NoAlias;
+  }
 
   /// The tag for type-based alias analysis.
-  MDNode *TBAA;
+  MDNode *TBAA = nullptr;
+
+  /// The tag for type-based alias analysis (tbaa struct).
+  MDNode *TBAAStruct = nullptr;
 
   /// The tag for alias scope specification (used with noalias).
-  MDNode *Scope;
+  MDNode *Scope = nullptr;
 
   /// The tag specifying the noalias scope.
-  MDNode *NoAlias;
+  MDNode *NoAlias = nullptr;
 
   /// Given two sets of AAMDNodes that apply to the same pointer,
   /// give the best AAMDNodes that are compatible with both (i.e. a set of
@@ -671,6 +677,7 @@ struct AAMDNodes {
   AAMDNodes intersect(const AAMDNodes &Other) {
     AAMDNodes Result;
     Result.TBAA = Other.TBAA == TBAA ? TBAA : nullptr;
+    Result.TBAAStruct = Other.TBAAStruct == TBAAStruct ? TBAAStruct : nullptr;
     Result.Scope = Other.Scope == Scope ? Scope : nullptr;
     Result.NoAlias = Other.NoAlias == NoAlias ? NoAlias : nullptr;
     return Result;
@@ -682,16 +689,17 @@ template<>
 struct DenseMapInfo<AAMDNodes> {
   static inline AAMDNodes getEmptyKey() {
     return AAMDNodes(DenseMapInfo<MDNode *>::getEmptyKey(),
-                     nullptr, nullptr);
+                     nullptr, nullptr, nullptr);
   }
 
   static inline AAMDNodes getTombstoneKey() {
     return AAMDNodes(DenseMapInfo<MDNode *>::getTombstoneKey(),
-                     nullptr, nullptr);
+                     nullptr, nullptr, nullptr);
   }
 
   static unsigned getHashValue(const AAMDNodes &Val) {
     return DenseMapInfo<MDNode *>::getHashValue(Val.TBAA) ^
+           DenseMapInfo<MDNode *>::getHashValue(Val.TBAAStruct) ^
            DenseMapInfo<MDNode *>::getHashValue(Val.Scope) ^
            DenseMapInfo<MDNode *>::getHashValue(Val.NoAlias);
   }

llvm/lib/Analysis/TypeBasedAliasAnalysis.cpp

@@ -521,24 +521,21 @@ static const MDNode *getLeastCommonType(const MDNode *A, const MDNode *B) {
 }
 
 void Instruction::getAAMetadata(AAMDNodes &N, bool Merge) const {
-  if (Merge)
+  if (Merge) {
     N.TBAA =
         MDNode::getMostGenericTBAA(N.TBAA, getMetadata(LLVMContext::MD_tbaa));
-  else
+    N.TBAAStruct = nullptr;
-    N.TBAA = getMetadata(LLVMContext::MD_tbaa);
-
-  if (Merge)
     N.Scope = MDNode::getMostGenericAliasScope(
         N.Scope, getMetadata(LLVMContext::MD_alias_scope));
-  else
-    N.Scope = getMetadata(LLVMContext::MD_alias_scope);
-
-  if (Merge)
     N.NoAlias =
         MDNode::intersect(N.NoAlias, getMetadata(LLVMContext::MD_noalias));
-  else
+  } else {
+    N.TBAA = getMetadata(LLVMContext::MD_tbaa);
+    N.TBAAStruct = getMetadata(LLVMContext::MD_tbaa_struct);
+    N.Scope = getMetadata(LLVMContext::MD_alias_scope);
     N.NoAlias = getMetadata(LLVMContext::MD_noalias);
   }
+}
 
 static const MDNode *createAccessTag(const MDNode *AccessType) {
   // If there is no access type or the access type is the root node, then

llvm/lib/IR/Metadata.cpp

@@ -1262,6 +1262,7 @@ void Instruction::setMetadata(unsigned KindID, MDNode *Node) {
 
 void Instruction::setAAMetadata(const AAMDNodes &N) {
   setMetadata(LLVMContext::MD_tbaa, N.TBAA);
+  setMetadata(LLVMContext::MD_tbaa_struct, N.TBAAStruct);
   setMetadata(LLVMContext::MD_alias_scope, N.Scope);
   setMetadata(LLVMContext::MD_noalias, N.NoAlias);
 }

llvm/test/Transforms/SROA/tbaa-struct.ll

@@ -11,7 +11,7 @@ define void @bar(%vector* %y2) {
 ; CHECK-LABEL: @bar(
 ; CHECK-NEXT:    [[X14:%.*]] = call <2 x float> @foo(%vector* [[Y2:%.*]])
 ; CHECK-NEXT:    [[X7_SROA_0_0_X18_SROA_CAST:%.*]] = bitcast %vector* [[Y2]] to <2 x float>*
-; CHECK-NEXT:    store <2 x float> [[X14]], <2 x float>* [[X7_SROA_0_0_X18_SROA_CAST]], align 4
+; CHECK-NEXT:    store <2 x float> [[X14]], <2 x float>* [[X7_SROA_0_0_X18_SROA_CAST]], align 4, !tbaa.struct !0
 ; CHECK-NEXT:    ret void
 ;
   %x7 = alloca %vector