[SimplifyCFG] Preserve common TBAA metadata when hoisting instructions.

Update FoldTwoEntryPHINode to collect common TBAA metadata for instructions
that match in all if-blocks and have the same TBAA metadata. If that is
the case, they access the same type on all paths and the TBAA info can
be preserved after hoisting.

I think we should be able to preserve most metadata, if it is available
on matching instructions in all blocks, i.e. preserve the intersection
of metadata on all matching instructions. I couldn't find any utility
that already computes that intersection. At the moment, the order of of
matching instructions must be the same.

Author: fhahn

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

@@ -3624,6 +3624,29 @@ static bool FoldTwoEntryPHINode(PHINode *PN, const TargetTransformInfo &TTI,
                     << "  T: " << IfTrue->getName()
                     << "  F: " << IfFalse->getName() << "\n");
 
+  // Collect common TBAA metadata, for instructions that match in all if-blocks
+  // and have the same TBAA metadata. If that is the case, they access the same
+  // type on all paths and the TBAA info can be preserved after hoisting.
+  // TODO: preserve other common metadata.
+  LockstepReverseIterator LRI(IfBlocks);
+  DenseMap<Instruction *, MDNode *> CommonTBAA;
+  while (LRI.isValid()) {
+    auto Insts = *LRI;
+    Instruction *I0 = Insts.front();
+    MDNode *MD = I0->getMetadata(LLVMContext::MD_tbaa);
+    if (!MD || any_of(Insts, [I0, MD](Instruction *I) {
+          return !I->isSameOperationAs(I0) ||
+                 !equal(I->operands(), I0->operands()) ||
+                 I->getMetadata(LLVMContext::MD_tbaa) != MD;
+        })) {
+      --LRI;
+      continue;
+    }
+    for (Instruction *I : Insts)
+      CommonTBAA[I] = MD;
+    --LRI;
+  }
+
   // If we can still promote the PHI nodes after this gauntlet of tests,
   // do all of the PHI's now.
 
@@ -3632,6 +3655,10 @@ static bool FoldTwoEntryPHINode(PHINode *PN, const TargetTransformInfo &TTI,
   for (BasicBlock *IfBlock : IfBlocks)
       hoistAllInstructionsInto(DomBlock, DomBI, IfBlock);
 
+  for (Instruction &I : *DomBlock)
+    if (auto *MD = CommonTBAA.lookup(&I))
+      I.setMetadata(LLVMContext::MD_tbaa, MD);
+
   IRBuilder<NoFolder> Builder(DomBI);
   // Propagate fast-math-flags from phi nodes to replacement selects.
   IRBuilder<>::FastMathFlagGuard FMFGuard(Builder);

llvm/test/Transforms/SimplifyCFG/hoisting-metadata.ll

@@ -8,10 +8,8 @@ define i64 @hoist_load_with_matching_pointers_and_tbaa(i1 %c) {
 ; CHECK-NEXT:  [[ENTRY:.*:]]
 ; CHECK-NEXT:    [[TMP:%.*]] = alloca i64, align 8
 ; CHECK-NEXT:    call void @init(ptr [[TMP]])
-; CHECK-NEXT:    [[TMP0:%.*]] = load i64, ptr [[TMP]], align 8
-; CHECK-NOT:       !tbaa
-; CHECK-NEXT:    [[TMP1:%.*]] = load i64, ptr [[TMP]], align 8
-; CHECK-NOT:       !tbaa
+; CHECK-NEXT:    [[TMP0:%.*]] = load i64, ptr [[TMP]], align 8, !tbaa [[M:!.+]]
+; CHECK-NEXT:    [[TMP1:%.*]] = load i64, ptr [[TMP]], align 8, !tbaa [[M]]
 ; CHECK-NEXT:    [[P:%.*]] = select i1 [[C]], i64 [[TMP0]], i64 [[TMP1]]
 ; CHECK-NEXT:    ret i64 [[P]]
 ;