Reland "[Support] Assert that DomTree nodes share parent" (#102782)

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 three 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.

Reverts #102780
Reland #101198
Fixes #102784

Co-authored-by: Alexis Engelke <[email protected]>

Author: vitalybuka

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 different parent");
     if (auto Idx = getNodeIndex(BB); Idx && *Idx < DomTreeNodes.size())
       return DomTreeNodes[*Idx].get();
     return nullptr;

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)) {

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);
   }
 

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.

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

@@ -6108,14 +6108,19 @@ BoUpSLP::collectUserStores(const BoUpSLP::TreeEntry *TE) const {
   DenseMap<Value *, SmallVector<StoreInst *>> PtrToStoresMap;
   for (unsigned Lane : seq<unsigned>(0, TE->Scalars.size())) {
     Value *V = TE->Scalars[Lane];
+    // Don't iterate over the users of constant data.
+    if (isa<ConstantData>(V))
+      continue;
     // To save compilation time we don't visit if we have too many users.
     if (V->hasNUsesOrMore(UsesLimit))
       break;
 
     // Collect stores per pointer object.
     for (User *U : V->users()) {
       auto *SI = dyn_cast<StoreInst>(U);
-      if (SI == nullptr || !SI->isSimple() ||
+      // Test whether we can handle the store. V might be a global, which could
+      // be used in a different function.
+      if (SI == nullptr || !SI->isSimple() || SI->getFunction() != F ||
           !isValidElementType(SI->getValueOperand()->getType()))
         continue;
       // Skip entry if already

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 different
+; function. Test that in this case the dominator tree is only queried with
+; blocks from the same function.
+
+@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
+}

llvm/test/Transforms/SLPVectorizer/X86/const-in-different-functions.ll

@@ -0,0 +1,48 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -mtriple=x86_64 -passes=slp-vectorizer < %s | FileCheck %s
+
+; Test that SLP vectorize doesn't crash if a stored constant is used in multiple
+; functions.
+
+@p = external global [64 x float]
+
+define void @_Z1hPfl() {
+; CHECK-LABEL: define void @_Z1hPfl() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr i8, ptr @p, i64 28
+; CHECK-NEXT:    store <2 x float> <float 0.000000e+00, float 1.000000e+00>, ptr [[TMP0]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  %0 = getelementptr i8, ptr @p, i64 28
+  store float 0.000000e+00, ptr %0, align 4
+  %1 = getelementptr i8, ptr @p, i64 32
+  store float 1.000000e+00, ptr %1, align 16
+  ret void
+}
+
+define void @_Z1mv(i64 %arrayidx4.i.2.idx) {
+; CHECK-LABEL: define void @_Z1mv(
+; CHECK-SAME: i64 [[ARRAYIDX4_I_2_IDX:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    ret void
+; CHECK:       [[FOR_COND1_PREHEADER_LR_PH_I:.*:]]
+; CHECK-NEXT:    br label %[[FOR_COND1_PREHEADER_I:.*]]
+; CHECK:       [[FOR_COND1_PREHEADER_I]]:
+; CHECK-NEXT:    store float 1.000000e+00, ptr @p, align 4
+; CHECK-NEXT:    [[ARRAYIDX4_I_2:%.*]] = getelementptr i8, ptr @p, i64 [[ARRAYIDX4_I_2_IDX]]
+; CHECK-NEXT:    store float 0.000000e+00, ptr [[ARRAYIDX4_I_2]], align 4
+; CHECK-NEXT:    br label %[[FOR_COND1_PREHEADER_I]]
+;
+entry:
+  ret void
+
+for.cond1.preheader.lr.ph.i:                      ; No predecessors!
+  br label %for.cond1.preheader.i
+
+for.cond1.preheader.i:                            ; preds = %for.cond1.preheader.i, %for.cond1.preheader.lr.ph.i
+  store float 1.000000e+00, ptr @p, align 4
+  %arrayidx4.i.2 = getelementptr i8, ptr @p, i64 %arrayidx4.i.2.idx
+  store float 0.000000e+00, ptr %arrayidx4.i.2, align 4
+  br label %for.cond1.preheader.i
+}