[llvm] [SimplifyCFG] Preserve common TBAA metadata when hoisting instructions. (PR #97158)

[Florian Hahn via llvm-commits]

https://github.com/fhahn created https://github.com/llvm/llvm-project/pull/97158

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.

>From 910a8578a37c4eb6b6a9aac34fa91a5b05958b11 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Sat, 29 Jun 2024 10:43:28 +0100
Subject: [PATCH] [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.
---
 llvm/lib/Transforms/Utils/SimplifyCFG.cpp     | 27 +++++++++++++++++++
 .../SimplifyCFG/hoisting-metadata.ll          |  6 ++---
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index 6847bb7502429..c2774b8b74b4c 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/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);
diff --git a/llvm/test/Transforms/SimplifyCFG/hoisting-metadata.ll b/llvm/test/Transforms/SimplifyCFG/hoisting-metadata.ll
index 026002a4942af..4aea8634bafcb 100644
--- a/llvm/test/Transforms/SimplifyCFG/hoisting-metadata.ll
+++ b/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]]
 ;