[NFCI][IRBuilder] Add CreateAllocationSize helper (#178346)

Create a new `IRBuilderBase::CreateAllocationSize` method to compute the
runtime size of an alloca as a Value*. This handles both static and
dynamic allocas by computing `ArraySize * ElementSize`, and using
CreateTypeSize to properly handle scalable vectors.

This de-duplicates code across multiple instrumentation and analysis
passes and increases consistency.

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
Jameson Nash 2026-01-28 22:46:25 -05:00 committed by GitHub
parent 3404537393
commit eef6b62dcc
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
11 changed files with 42 additions and 55 deletions

View File

@ -969,6 +969,10 @@ public:
/// in poison if type \p Ty is not big enough to hold the value.
LLVM_ABI Value *CreateTypeSize(Type *Ty, TypeSize Size);
/// Get allocation size of an alloca as a runtime Value* (handles both static
/// and dynamic allocas and vscale factor).
LLVM_ABI Value *CreateAllocationSize(Type *DestTy, AllocaInst *AI);
/// Creates a vector of type \p DstType with the linear sequence <0, 1, ...>
LLVM_ABI Value *CreateStepVector(Type *DstType, const Twine &Name = "");

View File

@ -1346,15 +1346,11 @@ SizeOffsetValue ObjectSizeOffsetEvaluator::visitAllocaInst(AllocaInst &I) {
// If needed, adjust the alloca's operand size to match the pointer indexing
// size. Subsequent math operations expect the types to match.
Value *ArraySize = Builder.CreateZExtOrTrunc(
I.getArraySize(),
DL.getIndexType(I.getContext(), DL.getAllocaAddrSpace()));
assert(ArraySize->getType() == Zero->getType() &&
Type *IndexTy = DL.getIndexType(I.getContext(), DL.getAllocaAddrSpace());
assert(IndexTy == Zero->getType() &&
"Expected zero constant to have pointer index type");
Value *Size = Builder.CreateTypeSize(
ArraySize->getType(), DL.getTypeAllocSize(I.getAllocatedType()));
Size = Builder.CreateMul(Size, ArraySize);
Value *Size = Builder.CreateAllocationSize(IndexTy, &I);
return SizeOffsetValue(Size, Zero);
}

View File

@ -669,13 +669,8 @@ void SafeStack::moveDynamicAllocasToUnsafeStack(
IRBuilder<> IRB(AI);
// Compute the new SP value (after AI).
Value *ArraySize = AI->getArraySize();
if (ArraySize->getType() != IntPtrTy)
ArraySize = IRB.CreateIntCast(ArraySize, IntPtrTy, false);
Type *Ty = AI->getAllocatedType();
uint64_t TySize = DL.getTypeAllocSize(Ty);
Value *Size = IRB.CreateMul(ArraySize, ConstantInt::get(IntPtrTy, TySize));
Value *Size = IRB.CreateAllocationSize(IntPtrTy, AI);
Value *SP = IRB.CreatePtrToInt(IRB.CreateLoad(StackPtrTy, UnsafeStackPtr),
IntPtrTy);

View File

@ -135,6 +135,15 @@ Value *IRBuilderBase::CreateTypeSize(Type *Ty, TypeSize Size) {
return CreateVScaleMultiple(*this, Ty, Size.getKnownMinValue());
}
Value *IRBuilderBase::CreateAllocationSize(Type *DestTy, AllocaInst *AI) {
const DataLayout &DL = BB->getDataLayout();
TypeSize ElemSize = DL.getTypeAllocSize(AI->getAllocatedType());
Value *Size = CreateTypeSize(DestTy, ElemSize);
if (AI->isArrayAllocation())
Size = CreateMul(CreateZExtOrTrunc(AI->getArraySize(), DestTy), Size);
return Size;
}
Value *IRBuilderBase::CreateStepVector(Type *DstType, const Twine &Name) {
Type *STy = DstType->getScalarType();
if (isa<ScalableVectorType>(DstType)) {

View File

@ -3837,11 +3837,7 @@ void FunctionStackPoisoner::handleDynamicAllocaCall(AllocaInst *AI) {
// redzones, and OldSize is number of allocated blocks with
// ElementSize size, get allocated memory size in bytes by
// OldSize * ElementSize.
const unsigned ElementSize =
F.getDataLayout().getTypeAllocSize(AI->getAllocatedType());
Value *OldSize =
IRB.CreateMul(IRB.CreateIntCast(AI->getArraySize(), IntptrTy, false),
ConstantInt::get(IntptrTy, ElementSize));
Value *OldSize = IRB.CreateAllocationSize(IntptrTy, AI);
// PartialSize = OldSize % 32
Value *PartialSize = IRB.CreateAnd(OldSize, AllocaRzMask);

View File

@ -7308,12 +7308,7 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
if (!InsPoint)
InsPoint = &I;
NextNodeIRBuilder IRB(InsPoint);
const DataLayout &DL = F.getDataLayout();
TypeSize TS = DL.getTypeAllocSize(I.getAllocatedType());
Value *Len = IRB.CreateTypeSize(MS.IntptrTy, TS);
if (I.isArrayAllocation())
Len = IRB.CreateMul(Len,
IRB.CreateZExtOrTrunc(I.getArraySize(), MS.IntptrTy));
Value *Len = IRB.CreateAllocationSize(MS.IntptrTy, &I);
if (MS.CompileKernel)
poisonAllocaKmsan(I, IRB, Len);

View File

@ -854,13 +854,6 @@ bool TypeSanitizer::instrumentMemInst(Value *V, Instruction *ShadowBase,
bool NeedsMemMove = false;
IRBuilder<> IRB(BB, IP);
auto GetAllocaSize = [&](AllocaInst *AI) {
return IRB.CreateMul(
IRB.CreateZExtOrTrunc(AI->getArraySize(), IntptrTy),
ConstantInt::get(IntptrTy,
DL.getTypeAllocSize(AI->getAllocatedType())));
};
if (auto *A = dyn_cast<Argument>(V)) {
assert(A->hasByValAttr() && "Type reset for non-byval argument?");
@ -887,7 +880,7 @@ bool TypeSanitizer::instrumentMemInst(Value *V, Instruction *ShadowBase,
if (!AI)
return false;
Size = GetAllocaSize(AI);
Size = IRB.CreateAllocationSize(IntptrTy, AI);
Dest = II->getArgOperand(0);
} else if (auto *AI = dyn_cast<AllocaInst>(I)) {
// We need to clear the types for new stack allocations (or else we might
@ -896,7 +889,7 @@ bool TypeSanitizer::instrumentMemInst(Value *V, Instruction *ShadowBase,
IRB.SetInsertPoint(&*std::next(BasicBlock::iterator(I)));
IRB.SetInstDebugLocation(I);
Size = GetAllocaSize(AI);
Size = IRB.CreateAllocationSize(IntptrTy, AI);
Dest = I;
} else {
return false;

View File

@ -20,7 +20,7 @@ target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f3
define void @f1(i64 %x) nounwind {
; TR-LABEL: define void @f1(
; TR-SAME: i64 [[X:%.*]]) #[[ATTR0:[0-9]+]] {
; TR-NEXT: [[TMP1:%.*]] = mul i64 16, [[X]]
; TR-NEXT: [[TMP1:%.*]] = mul i64 [[X]], 16
; TR-NEXT: [[TMP2:%.*]] = alloca i128, i64 [[X]], align 8
; TR-NEXT: [[TMP3:%.*]] = sub i64 [[TMP1]], 0, !nosanitize [[META0:![0-9]+]]
; TR-NEXT: [[TMP4:%.*]] = icmp ult i64 [[TMP3]], 16, !nosanitize [[META0]]
@ -36,7 +36,7 @@ define void @f1(i64 %x) nounwind {
;
; RT-LABEL: define void @f1(
; RT-SAME: i64 [[X:%.*]]) #[[ATTR0:[0-9]+]] {
; RT-NEXT: [[TMP1:%.*]] = mul i64 16, [[X]]
; RT-NEXT: [[TMP1:%.*]] = mul i64 [[X]], 16
; RT-NEXT: [[TMP2:%.*]] = alloca i128, i64 [[X]], align 8
; RT-NEXT: [[TMP3:%.*]] = sub i64 [[TMP1]], 0, !nosanitize [[META0:![0-9]+]]
; RT-NEXT: [[TMP4:%.*]] = icmp ult i64 [[TMP3]], 16, !nosanitize [[META0]]
@ -52,7 +52,7 @@ define void @f1(i64 %x) nounwind {
;
; TR-NOMERGE-LABEL: define void @f1(
; TR-NOMERGE-SAME: i64 [[X:%.*]]) #[[ATTR0:[0-9]+]] {
; TR-NOMERGE-NEXT: [[TMP1:%.*]] = mul i64 16, [[X]]
; TR-NOMERGE-NEXT: [[TMP1:%.*]] = mul i64 [[X]], 16
; TR-NOMERGE-NEXT: [[TMP2:%.*]] = alloca i128, i64 [[X]], align 8
; TR-NOMERGE-NEXT: [[TMP3:%.*]] = sub i64 [[TMP1]], 0, !nosanitize [[META0:![0-9]+]]
; TR-NOMERGE-NEXT: [[TMP4:%.*]] = icmp ult i64 [[TMP3]], 16, !nosanitize [[META0]]
@ -68,7 +68,7 @@ define void @f1(i64 %x) nounwind {
;
; RT-NOMERGE-LABEL: define void @f1(
; RT-NOMERGE-SAME: i64 [[X:%.*]]) #[[ATTR0:[0-9]+]] {
; RT-NOMERGE-NEXT: [[TMP1:%.*]] = mul i64 16, [[X]]
; RT-NOMERGE-NEXT: [[TMP1:%.*]] = mul i64 [[X]], 16
; RT-NOMERGE-NEXT: [[TMP2:%.*]] = alloca i128, i64 [[X]], align 8
; RT-NOMERGE-NEXT: [[TMP3:%.*]] = sub i64 [[TMP1]], 0, !nosanitize [[META0:![0-9]+]]
; RT-NOMERGE-NEXT: [[TMP4:%.*]] = icmp ult i64 [[TMP3]], 16, !nosanitize [[META0]]
@ -84,7 +84,7 @@ define void @f1(i64 %x) nounwind {
;
; RTABORT-NOMERGE-LABEL: define void @f1(
; RTABORT-NOMERGE-SAME: i64 [[X:%.*]]) #[[ATTR0:[0-9]+]] {
; RTABORT-NOMERGE-NEXT: [[TMP1:%.*]] = mul i64 16, [[X]]
; RTABORT-NOMERGE-NEXT: [[TMP1:%.*]] = mul i64 [[X]], 16
; RTABORT-NOMERGE-NEXT: [[TMP2:%.*]] = alloca i128, i64 [[X]], align 8
; RTABORT-NOMERGE-NEXT: [[TMP3:%.*]] = sub i64 [[TMP1]], 0, !nosanitize [[META0:![0-9]+]]
; RTABORT-NOMERGE-NEXT: [[TMP4:%.*]] = icmp ult i64 [[TMP3]], 16, !nosanitize [[META0]]
@ -100,7 +100,7 @@ define void @f1(i64 %x) nounwind {
;
; MINRT-PRESERVE-NOMERGE-LABEL: define void @f1(
; MINRT-PRESERVE-NOMERGE-SAME: i64 [[X:%.*]]) #[[ATTR0:[0-9]+]] {
; MINRT-PRESERVE-NOMERGE-NEXT: [[TMP1:%.*]] = mul i64 16, [[X]]
; MINRT-PRESERVE-NOMERGE-NEXT: [[TMP1:%.*]] = mul i64 [[X]], 16
; MINRT-PRESERVE-NOMERGE-NEXT: [[TMP2:%.*]] = alloca i128, i64 [[X]], align 8
; MINRT-PRESERVE-NOMERGE-NEXT: [[TMP3:%.*]] = sub i64 [[TMP1]], 0, !nosanitize [[META0:![0-9]+]]
; MINRT-PRESERVE-NOMERGE-NEXT: [[TMP4:%.*]] = icmp ult i64 [[TMP3]], 16, !nosanitize [[META0]]
@ -116,7 +116,7 @@ define void @f1(i64 %x) nounwind {
;
; MINRT-NOMERGE-LABEL: define void @f1(
; MINRT-NOMERGE-SAME: i64 [[X:%.*]]) #[[ATTR0:[0-9]+]] {
; MINRT-NOMERGE-NEXT: [[TMP1:%.*]] = mul i64 16, [[X]]
; MINRT-NOMERGE-NEXT: [[TMP1:%.*]] = mul i64 [[X]], 16
; MINRT-NOMERGE-NEXT: [[TMP2:%.*]] = alloca i128, i64 [[X]], align 8
; MINRT-NOMERGE-NEXT: [[TMP3:%.*]] = sub i64 [[TMP1]], 0, !nosanitize [[META0:![0-9]+]]
; MINRT-NOMERGE-NEXT: [[TMP4:%.*]] = icmp ult i64 [[TMP3]], 16, !nosanitize [[META0]]
@ -132,7 +132,7 @@ define void @f1(i64 %x) nounwind {
;
; MINRTABORT-NOMERGE-LABEL: define void @f1(
; MINRTABORT-NOMERGE-SAME: i64 [[X:%.*]]) #[[ATTR0:[0-9]+]] {
; MINRTABORT-NOMERGE-NEXT: [[TMP1:%.*]] = mul i64 16, [[X]]
; MINRTABORT-NOMERGE-NEXT: [[TMP1:%.*]] = mul i64 [[X]], 16
; MINRTABORT-NOMERGE-NEXT: [[TMP2:%.*]] = alloca i128, i64 [[X]], align 8
; MINRTABORT-NOMERGE-NEXT: [[TMP3:%.*]] = sub i64 [[TMP1]], 0, !nosanitize [[META0:![0-9]+]]
; MINRTABORT-NOMERGE-NEXT: [[TMP4:%.*]] = icmp ult i64 [[TMP3]], 16, !nosanitize [[META0]]
@ -148,7 +148,7 @@ define void @f1(i64 %x) nounwind {
;
; TR-GUARD-COMMON-LABEL: define void @f1(
; TR-GUARD-COMMON-SAME: i64 [[X:%.*]]) #[[ATTR0:[0-9]+]] {
; TR-GUARD-COMMON-NEXT: [[TMP1:%.*]] = mul i64 16, [[X]]
; TR-GUARD-COMMON-NEXT: [[TMP1:%.*]] = mul i64 [[X]], 16
; TR-GUARD-COMMON-NEXT: [[TMP2:%.*]] = alloca i128, i64 [[X]], align 8
; TR-GUARD-COMMON-NEXT: [[TMP3:%.*]] = sub i64 [[TMP1]], 0, !nosanitize [[META0:![0-9]+]]
; TR-GUARD-COMMON-NEXT: [[TMP4:%.*]] = icmp ult i64 [[TMP3]], 16, !nosanitize [[META0]]
@ -172,7 +172,7 @@ define void @f1(i64 %x) nounwind {
;
; RT-GUARD-LABEL: define void @f1(
; RT-GUARD-SAME: i64 [[X:%.*]]) #[[ATTR0:[0-9]+]] {
; RT-GUARD-NEXT: [[TMP1:%.*]] = mul i64 16, [[X]]
; RT-GUARD-NEXT: [[TMP1:%.*]] = mul i64 [[X]], 16
; RT-GUARD-NEXT: [[TMP2:%.*]] = alloca i128, i64 [[X]], align 8
; RT-GUARD-NEXT: [[TMP3:%.*]] = sub i64 [[TMP1]], 0, !nosanitize [[META0:![0-9]+]]
; RT-GUARD-NEXT: [[TMP4:%.*]] = icmp ult i64 [[TMP3]], 16, !nosanitize [[META0]]

View File

@ -182,7 +182,7 @@ define void @f6(i64 %x) nounwind {
define void @f7(i64 %x) nounwind {
; CHECK-LABEL: @f7(
; CHECK-NEXT: [[TMP1:%.*]] = mul i64 16, [[X:%.*]]
; CHECK-NEXT: [[TMP1:%.*]] = mul i64 [[X:%.*]], 16
; CHECK-NEXT: [[TMP2:%.*]] = alloca i128, i64 [[X]], align 8
; CHECK-NEXT: [[TMP3:%.*]] = sub i64 [[TMP1]], 0, !nosanitize [[META0]]
; CHECK-NEXT: [[TMP4:%.*]] = icmp ult i64 [[TMP3]], 16, !nosanitize [[META0]]
@ -231,9 +231,9 @@ define void @f9(ptr %arg) nounwind {
define void @f10(i64 %x, i64 %y) nounwind {
; CHECK-LABEL: @f10(
; CHECK-NEXT: [[TMP1:%.*]] = mul i64 16, [[X:%.*]]
; CHECK-NEXT: [[TMP1:%.*]] = mul i64 [[X:%.*]], 16
; CHECK-NEXT: [[TMP2:%.*]] = alloca i128, i64 [[X]], align 8
; CHECK-NEXT: [[TMP3:%.*]] = mul i64 16, [[Y:%.*]]
; CHECK-NEXT: [[TMP3:%.*]] = mul i64 [[Y:%.*]], 16
; CHECK-NEXT: [[TMP4:%.*]] = alloca i128, i64 [[Y]], align 8
; CHECK-NEXT: [[TMP5:%.*]] = select i1 undef, i64 [[TMP1]], i64 [[TMP3]]
; CHECK-NEXT: [[TMP6:%.*]] = select i1 undef, ptr [[TMP2]], ptr [[TMP4]]
@ -493,7 +493,7 @@ define void @scalable_alloca(i64 %y) nounwind {
; CHECK-LABEL: @scalable_alloca(
; CHECK-NEXT: [[TMP1:%.*]] = call i64 @llvm.vscale.i64()
; CHECK-NEXT: [[TMP2:%.*]] = mul nuw i64 [[TMP1]], 8
; CHECK-NEXT: [[TMP3:%.*]] = mul i64 [[TMP2]], 5
; CHECK-NEXT: [[TMP3:%.*]] = mul i64 5, [[TMP2]]
; CHECK-NEXT: [[TMP4:%.*]] = alloca <vscale x 4 x i16>, i32 5, align 8
; CHECK-NEXT: [[TMP5:%.*]] = call i64 @llvm.vscale.i64()
; CHECK-NEXT: [[TMP6:%.*]] = mul nuw i64 [[TMP5]], 8
@ -526,7 +526,6 @@ define void @scalable_alloca2(i64 %y) nounwind {
; CHECK-LABEL: @scalable_alloca2(
; CHECK-NEXT: [[TMP1:%.*]] = call i64 @llvm.vscale.i64()
; CHECK-NEXT: [[TMP2:%.*]] = mul nuw i64 [[TMP1]], 32
; CHECK-NEXT: [[TMP3:%.*]] = mul i64 [[TMP2]], 1
; CHECK-NEXT: [[TMP4:%.*]] = alloca <vscale x 4 x i64>, align 32
; CHECK-NEXT: [[TMP5:%.*]] = call i64 @llvm.vscale.i64()
; CHECK-NEXT: [[TMP6:%.*]] = mul nuw i64 [[TMP5]], 32
@ -535,14 +534,14 @@ define void @scalable_alloca2(i64 %y) nounwind {
; CHECK-NEXT: [[TMP8:%.*]] = getelementptr inbounds <vscale x 4 x i64>, ptr [[TMP4]], i64 [[Y]]
; CHECK-NEXT: [[TMP9:%.*]] = call i64 @llvm.vscale.i64(), !nosanitize [[META0]]
; CHECK-NEXT: [[TMP10:%.*]] = mul nuw i64 [[TMP9]], 32, !nosanitize [[META0]]
; CHECK-NEXT: [[TMP11:%.*]] = sub i64 [[TMP3]], [[TMP7]], !nosanitize [[META0]]
; CHECK-NEXT: [[TMP12:%.*]] = icmp ult i64 [[TMP3]], [[TMP7]], !nosanitize [[META0]]
; CHECK-NEXT: [[TMP11:%.*]] = sub i64 [[TMP2]], [[TMP7]], !nosanitize [[META0]]
; CHECK-NEXT: [[TMP12:%.*]] = icmp ult i64 [[TMP2]], [[TMP7]], !nosanitize [[META0]]
; CHECK-NEXT: [[TMP13:%.*]] = icmp ult i64 [[TMP11]], [[TMP10]], !nosanitize [[META0]]
; CHECK-NEXT: [[TMP14:%.*]] = or i1 [[TMP12]], [[TMP13]], !nosanitize [[META0]]
; CHECK-NEXT: [[TMP15:%.*]] = icmp slt i64 [[TMP7]], 0, !nosanitize [[META0]]
; CHECK-NEXT: [[TMP16:%.*]] = or i1 [[TMP15]], [[TMP14]], !nosanitize [[META0]]
; CHECK-NEXT: br i1 [[TMP16]], label [[TRAP:%.*]], label [[TMP17:%.*]]
; CHECK: 17:
; CHECK: 16:
; CHECK-NEXT: [[TMP18:%.*]] = load <vscale x 4 x i64>, ptr [[TMP8]], align 4
; CHECK-NEXT: ret void
; CHECK: trap:

View File

@ -79,7 +79,7 @@ entry:
}
; CHECK-LABEL: define void @array_non_const(
; CHECK: %[[A:.*]] = mul i64 4, %cnt
; CHECK: %[[A:.*]] = mul i64 %cnt, 4
; INLINE: call void @llvm.memset.p0.i64(ptr align 4 {{.*}}, i8 -1, i64 %[[A]], i1 false)
; CALL: call void @__msan_poison_stack(ptr {{.*}}, i64 %[[A]])
; ORIGIN: call void @__msan_set_alloca_origin_with_descr(ptr {{.*}}, i64 %[[A]],
@ -95,7 +95,7 @@ entry:
; CHECK-LABEL: define void @array_non_const32(
; CHECK: %[[Z:.*]] = zext i32 %cnt to i64
; CHECK: %[[A:.*]] = mul i64 4, %[[Z]]
; CHECK: %[[A:.*]] = mul i64 %[[Z]], 4
; INLINE: call void @llvm.memset.p0.i64(ptr align 4 {{.*}}, i8 -1, i64 %[[A]], i1 false)
; CALL: call void @__msan_poison_stack(ptr {{.*}}, i64 %[[A]])
; ORIGIN: call void @__msan_set_alloca_origin_with_descr(ptr {{.*}}, i64 %[[A]],
@ -169,7 +169,7 @@ entry:
; CHECK-LABEL: entry:
; CHECK: %x = alloca i32, i64 %cnt
; CHECK: call void @llvm.lifetime.start
; CHECK: %[[A:.*]] = mul i64 4, %cnt
; CHECK: %[[A:.*]] = mul i64 %cnt, 4
; INLINE: call void @llvm.memset.p0.i64(ptr align 4 {{.*}}, i8 -1, i64 %[[A]], i1 false)
; CALL: call void @__msan_poison_stack(ptr {{.*}}, i64 %[[A]])
; ORIGIN: call void @__msan_set_alloca_origin_with_descr(ptr {{.*}}, i64 %[[A]],

View File

@ -119,7 +119,7 @@ define i64 @test_objectsize_byref_arg(ptr byref([42 x i8]) %ptr) {
define i64 @vla_pointer_size_mismatch(i42 %x) {
; CHECK-LABEL: @vla_pointer_size_mismatch(
; CHECK-NEXT: [[TMP1:%.*]] = zext i42 [[X:%.*]] to i64
; CHECK-NEXT: [[TMP2:%.*]] = mul i64 1, [[TMP1]]
; CHECK-NEXT: [[TMP2:%.*]] = mul i64 [[TMP1]], 1
; CHECK-NEXT: [[A:%.*]] = alloca i8, i42 [[X]], align 1
; CHECK-NEXT: [[G1:%.*]] = getelementptr i8, ptr [[A]], i8 17
; CHECK-NEXT: [[TMP3:%.*]] = sub i64 [[TMP2]], 17