From 8101d1863cc3a6ca0ca49962903f2d7651b25659 Mon Sep 17 00:00:00 2001 From: Alexis Engelke Date: Sat, 10 Aug 2024 18:19:05 +0200 Subject: [PATCH] [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. --- llvm/include/llvm/Support/GenericDomTree.h | 2 ++ llvm/lib/Analysis/TypeMetadataUtils.cpp | 2 ++ .../Scalar/AlignmentFromAssumptions.cpp | 1 + llvm/lib/Transforms/Scalar/LoopFuse.cpp | 8 +++-- .../AlignmentFromAssumptions/domtree-crash.ll | 33 +++++++++++++++++++ 5 files changed, 44 insertions(+), 2 deletions(-) create mode 100644 llvm/test/Transforms/AlignmentFromAssumptions/domtree-crash.ll diff --git a/llvm/include/llvm/Support/GenericDomTree.h b/llvm/include/llvm/Support/GenericDomTree.h index 7e2b68e6faea29c..45ef38b965b7520 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 *getNode(const NodeT *BB) const { + assert((!BB || Parent == NodeTrait::getParent(const_cast(BB))) && + "cannot get DomTreeNode of block with different 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 67ce1540112bb73..9ec0785eb5034d6 100644 --- a/llvm/lib/Analysis/TypeMetadataUtils.cpp +++ b/llvm/lib/Analysis/TypeMetadataUtils.cpp @@ -33,6 +33,8 @@ findCallsAtConstantOffset(SmallVectorImpl &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(User)) { diff --git a/llvm/lib/Transforms/Scalar/AlignmentFromAssumptions.cpp b/llvm/lib/Transforms/Scalar/AlignmentFromAssumptions.cpp index f3422a705dca7a8..8555ef5c22f8272 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(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 8512b2accbe7c2e..fe0e30d1965e05f 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 000000000000000..c7fc1dc69967189 --- /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 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 +}