[llvm] 8101d18 - [Support] Assert that DomTree nodes share parent (#101198)

[via llvm-commits]

Author: Alexis Engelke
Date: 2024-08-10T18:19:05+02:00
New Revision: 8101d1863cc3a6ca0ca49962903f2d7651b25659

URL: https://github.com/llvm/llvm-project/commit/8101d1863cc3a6ca0ca49962903f2d7651b25659
DIFF: https://github.com/llvm/llvm-project/commit/8101d1863cc3a6ca0ca49962903f2d7651b25659.diff

LOG: [Support] Assert that DomTree nodes share parent (#101198)

A dominance query of a block that is in a different function is
ill-defined, so assert that getNode() is only called for blocks that are
in the same function.

There are two cases, where this behavior did occur. LoopFuse didn't
explicitly do this, but didn't invalidate the SCEV block dispositions,
leaving dangling pointers to free'ed basic blocks behind, causing
use-after-free. We do, however, want to be able to dereference basic
blocks inside the dominator tree, so that we can refer to them by a
number stored inside the basic block.

Added: 
    llvm/test/Transforms/AlignmentFromAssumptions/domtree-crash.ll

Modified: 
    llvm/include/llvm/Support/GenericDomTree.h
    llvm/lib/Analysis/TypeMetadataUtils.cpp
    llvm/lib/Transforms/Scalar/AlignmentFromAssumptions.cpp
    llvm/lib/Transforms/Scalar/LoopFuse.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/Support/GenericDomTree.h b/llvm/include/llvm/Support/GenericDomTree.h
index 7e2b68e6faea29..45ef38b965b752 100644
--- a/llvm/include/llvm/Support/GenericDomTree.h
+++ b/llvm/include/llvm/Support/GenericDomTree.h
@@ -397,6 +397,8 @@ class DominatorTreeBase {
   /// may (but is not required to) be null for a forward (backwards)
   /// statically unreachable block.
   DomTreeNodeBase<NodeT> *getNode(const NodeT *BB) const {
+    assert((!BB || Parent == NodeTrait::getParent(const_cast<NodeT *>(BB))) &&
+           "cannot get DomTreeNode of block with 
diff erent parent");
     if (auto Idx = getNodeIndex(BB); Idx && *Idx < DomTreeNodes.size())
       return DomTreeNodes[*Idx].get();
     return nullptr;

diff  --git a/llvm/lib/Analysis/TypeMetadataUtils.cpp b/llvm/lib/Analysis/TypeMetadataUtils.cpp
index 67ce1540112bb7..9ec0785eb5034d 100644
--- a/llvm/lib/Analysis/TypeMetadataUtils.cpp
+++ b/llvm/lib/Analysis/TypeMetadataUtils.cpp
@@ -33,6 +33,8 @@ findCallsAtConstantOffset(SmallVectorImpl<DevirtCallSite> &DevirtCalls,
     // after indirect call promotion and inlining, where we may have uses
     // of the vtable pointer guarded by a function pointer check, and a fallback
     // indirect call.
+    if (CI->getFunction() != User->getFunction())
+      continue;
     if (!DT.dominates(CI, User))
       continue;
     if (isa<BitCastInst>(User)) {

diff  --git a/llvm/lib/Transforms/Scalar/AlignmentFromAssumptions.cpp b/llvm/lib/Transforms/Scalar/AlignmentFromAssumptions.cpp
index f3422a705dca7a..8555ef5c22f827 100644
--- a/llvm/lib/Transforms/Scalar/AlignmentFromAssumptions.cpp
+++ b/llvm/lib/Transforms/Scalar/AlignmentFromAssumptions.cpp
@@ -208,6 +208,7 @@ bool AlignmentFromAssumptionsPass::processAssumption(CallInst *ACall,
       continue;
 
     if (Instruction *K = dyn_cast<Instruction>(J))
+      if (K->getFunction() == ACall->getFunction())
         WorkList.push_back(K);
   }
 

diff  --git a/llvm/lib/Transforms/Scalar/LoopFuse.cpp b/llvm/lib/Transforms/Scalar/LoopFuse.cpp
index 8512b2accbe7c2..fe0e30d1965e05 100644
--- a/llvm/lib/Transforms/Scalar/LoopFuse.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopFuse.cpp
@@ -1729,7 +1729,9 @@ struct LoopFuser {
     // mergeLatch may remove the only block in FC1.
     SE.forgetLoop(FC1.L);
     SE.forgetLoop(FC0.L);
-    SE.forgetLoopDispositions();
+    // Forget block dispositions as well, so that there are no dangling
+    // pointers to erased/free'ed blocks.
+    SE.forgetBlockAndLoopDispositions();
 
     // Move instructions from FC0.Latch to FC1.Latch.
     // Note: mergeLatch requires an updated DT.
@@ -2023,7 +2025,9 @@ struct LoopFuser {
     // mergeLatch may remove the only block in FC1.
     SE.forgetLoop(FC1.L);
     SE.forgetLoop(FC0.L);
-    SE.forgetLoopDispositions();
+    // Forget block dispositions as well, so that there are no dangling
+    // pointers to erased/free'ed blocks.
+    SE.forgetBlockAndLoopDispositions();
 
     // Move instructions from FC0.Latch to FC1.Latch.
     // Note: mergeLatch requires an updated DT.

diff  --git a/llvm/test/Transforms/AlignmentFromAssumptions/domtree-crash.ll b/llvm/test/Transforms/AlignmentFromAssumptions/domtree-crash.ll
new file mode 100644
index 00000000000000..c7fc1dc6996718
--- /dev/null
+++ b/llvm/test/Transforms/AlignmentFromAssumptions/domtree-crash.ll
@@ -0,0 +1,33 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -passes=alignment-from-assumptions -S < %s | FileCheck %s
+
+; The alignment assumption is a global, which has users in a 
diff erent
+; function. Test that in this case the dominator tree is only queried with
+; blocks from the same function.
+
+ at global = external constant [192 x i8]
+
+define void @fn1() {
+; CHECK-LABEL: define void @fn1() {
+; CHECK-NEXT:    call void @llvm.assume(i1 false) [ "align"(ptr @global, i64 1) ]
+; CHECK-NEXT:    ret void
+;
+  call void @llvm.assume(i1 false) [ "align"(ptr @global, i64 1) ]
+  ret void
+}
+
+define void @fn2() {
+; CHECK-LABEL: define void @fn2() {
+; CHECK-NEXT:    ret void
+; CHECK:       [[LOOP:.*]]:
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds i8, ptr @global, i64 0
+; CHECK-NEXT:    [[LOAD:%.*]] = load i64, ptr [[GEP]], align 1
+; CHECK-NEXT:    br label %[[LOOP]]
+;
+  ret void
+
+loop:
+  %gep = getelementptr inbounds i8, ptr @global, i64 0
+  %load = load i64, ptr %gep, align 1
+  br label %loop
+}