From 3a170b9f26b4951bce1ae9709a3f58350fb3b4e4 Mon Sep 17 00:00:00 2001 From: Alexis Engelke Date: Tue, 30 Jul 2024 13:25:23 +0000 Subject: [PATCH 1/4] [Support] Assert that DomTree nodes share parent 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 ++ llvm/lib/Analysis/ValueTracking.cpp | 4 ++++ llvm/lib/Transforms/Scalar/LoopFuse.cpp | 8 ++++++-- 4 files changed, 14 insertions(+), 2 deletions(-) diff --git a/llvm/include/llvm/Support/GenericDomTree.h b/llvm/include/llvm/Support/GenericDomTree.h index 7e2b68e6faea2..45ef38b965b75 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 67ce1540112bb..9ec0785eb5034 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/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp index 202eaad57d1e3..409006e79973d 100644 --- a/llvm/lib/Analysis/ValueTracking.cpp +++ b/llvm/lib/Analysis/ValueTracking.cpp @@ -526,6 +526,10 @@ bool llvm::isValidAssumeForContext(const Instruction *Inv, return AllowEphemerals || !isEphemeralValueOf(Inv, CxtI); } + // Inv and CxtI are in different functions. + if (Inv->getFunction() != CxtI->getFunction()) + return false; + // Inv and CxtI are in different blocks. if (DT) { if (DT->dominates(Inv, CxtI)) diff --git a/llvm/lib/Transforms/Scalar/LoopFuse.cpp b/llvm/lib/Transforms/Scalar/LoopFuse.cpp index 8512b2accbe7c..fe0e30d1965e0 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. From c6cf51214b03263696107b29606e5cfa3e0917dd Mon Sep 17 00:00:00 2001 From: Alexis Engelke Date: Wed, 31 Jul 2024 17:12:29 +0000 Subject: [PATCH 2/4] Address comments --- llvm/lib/Analysis/ValueTracking.cpp | 4 ---- llvm/lib/Transforms/Scalar/AlignmentFromAssumptions.cpp | 3 +++ 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp index 409006e79973d..202eaad57d1e3 100644 --- a/llvm/lib/Analysis/ValueTracking.cpp +++ b/llvm/lib/Analysis/ValueTracking.cpp @@ -526,10 +526,6 @@ bool llvm::isValidAssumeForContext(const Instruction *Inv, return AllowEphemerals || !isEphemeralValueOf(Inv, CxtI); } - // Inv and CxtI are in different functions. - if (Inv->getFunction() != CxtI->getFunction()) - return false; - // Inv and CxtI are in different blocks. if (DT) { if (DT->dominates(Inv, CxtI)) diff --git a/llvm/lib/Transforms/Scalar/AlignmentFromAssumptions.cpp b/llvm/lib/Transforms/Scalar/AlignmentFromAssumptions.cpp index f3422a705dca7..d66c1c597103b 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); } @@ -267,6 +268,8 @@ bool AlignmentFromAssumptionsPass::processAssumption(CallInst *ACall, for (auto &U : J->uses()) { if (U->getType()->isPointerTy()) { Instruction *K = cast(U.getUser()); + if (K->getFunction() != ACall->getFunction()) + continue; StoreInst *SI = dyn_cast(K); if (SI && SI->getPointerOperandIndex() != U.getOperandNo()) continue; From f0305cdc0d87b1976f3e73f29dce764460cf2136 Mon Sep 17 00:00:00 2001 From: Alexis Engelke Date: Wed, 31 Jul 2024 17:20:03 +0000 Subject: [PATCH 3/4] Add test case --- .../AlignmentFromAssumptions/domtree-crash.ll | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) create mode 100644 llvm/test/Transforms/AlignmentFromAssumptions/domtree-crash.ll diff --git a/llvm/test/Transforms/AlignmentFromAssumptions/domtree-crash.ll b/llvm/test/Transforms/AlignmentFromAssumptions/domtree-crash.ll new file mode 100644 index 0000000000000..df7ae6d036ef1 --- /dev/null +++ b/llvm/test/Transforms/AlignmentFromAssumptions/domtree-crash.ll @@ -0,0 +1,25 @@ +; RUN: opt -passes=alignment-from-assumptions -disable-output < %s + +; REQUIRES: asserts + +; 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. + +target triple = "x86_64-unknown-linux-gnu" + +@global = external constant [192 x i8] + +define void @fn1() { + call void @llvm.assume(i1 false) [ "align"(ptr @global, i64 1) ] + ret void +} + +define void @fn2() { + ret void + +loop: + %gep = getelementptr inbounds i8, ptr @global, i64 0 + %load = load i64, ptr %gep, align 1 + br label %loop +} From be63e6d29010d6f2284ebc5bc1e902134f0316fa Mon Sep 17 00:00:00 2001 From: Alexis Engelke Date: Thu, 1 Aug 2024 11:13:11 +0000 Subject: [PATCH 4/4] Address comments --- .../Scalar/AlignmentFromAssumptions.cpp | 2 -- .../AlignmentFromAssumptions/domtree-crash.ll | 18 +++++++++++++----- 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/llvm/lib/Transforms/Scalar/AlignmentFromAssumptions.cpp b/llvm/lib/Transforms/Scalar/AlignmentFromAssumptions.cpp index d66c1c597103b..8555ef5c22f82 100644 --- a/llvm/lib/Transforms/Scalar/AlignmentFromAssumptions.cpp +++ b/llvm/lib/Transforms/Scalar/AlignmentFromAssumptions.cpp @@ -268,8 +268,6 @@ bool AlignmentFromAssumptionsPass::processAssumption(CallInst *ACall, for (auto &U : J->uses()) { if (U->getType()->isPointerTy()) { Instruction *K = cast(U.getUser()); - if (K->getFunction() != ACall->getFunction()) - continue; StoreInst *SI = dyn_cast(K); if (SI && SI->getPointerOperandIndex() != U.getOperandNo()) continue; diff --git a/llvm/test/Transforms/AlignmentFromAssumptions/domtree-crash.ll b/llvm/test/Transforms/AlignmentFromAssumptions/domtree-crash.ll index df7ae6d036ef1..c7fc1dc699671 100644 --- a/llvm/test/Transforms/AlignmentFromAssumptions/domtree-crash.ll +++ b/llvm/test/Transforms/AlignmentFromAssumptions/domtree-crash.ll @@ -1,21 +1,29 @@ -; RUN: opt -passes=alignment-from-assumptions -disable-output < %s - -; REQUIRES: asserts +; 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. -target triple = "x86_64-unknown-linux-gnu" - @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: