diff --git a/llvm/lib/Transforms/Instrumentation/TypeSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/TypeSanitizer.cpp index 46b56737e906e..4edf25c054b1d 100644 --- a/llvm/lib/Transforms/Instrumentation/TypeSanitizer.cpp +++ b/llvm/lib/Transforms/Instrumentation/TypeSanitizer.cpp @@ -789,6 +789,13 @@ 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(V)) { assert(A->hasByValAttr() && "Type reset for non-byval argument?"); @@ -811,7 +818,11 @@ bool TypeSanitizer::instrumentMemInst(Value *V, Instruction *ShadowBase, } } } else if (auto *II = dyn_cast(I)) { - Size = II->getArgOperand(0); + auto *AI = dyn_cast(II->getArgOperand(1)); + if (!AI) + return false; + + Size = GetAllocaSize(AI); Dest = II->getArgOperand(1); } else if (auto *AI = dyn_cast(I)) { // We need to clear the types for new stack allocations (or else we might @@ -820,10 +831,7 @@ bool TypeSanitizer::instrumentMemInst(Value *V, Instruction *ShadowBase, IRB.SetInsertPoint(&*std::next(BasicBlock::iterator(I))); IRB.SetInstDebugLocation(I); - Size = IRB.CreateMul( - IRB.CreateZExtOrTrunc(AI->getArraySize(), IntptrTy), - ConstantInt::get(IntptrTy, - DL.getTypeAllocSize(AI->getAllocatedType()))); + Size = GetAllocaSize(AI); Dest = I; } else { return false; diff --git a/llvm/test/Instrumentation/TypeSanitizer/alloca.ll b/llvm/test/Instrumentation/TypeSanitizer/alloca.ll index c53b00650cdcd..fc7263193dad6 100644 --- a/llvm/test/Instrumentation/TypeSanitizer/alloca.ll +++ b/llvm/test/Instrumentation/TypeSanitizer/alloca.ll @@ -74,3 +74,56 @@ loop: exit: ret void } + +define void @dynamic_alloca_lifetime_test(i1 %c, i64 %n) sanitize_type { +; CHECK-LABEL: @dynamic_alloca_lifetime_test( +; CHECK-NEXT: entry: +; CHECK-NEXT: [[APP_MEM_MASK:%.*]] = load i64, ptr @__tysan_app_memory_mask, align 8 +; CHECK-NEXT: [[SHADOW_BASE:%.*]] = load i64, ptr @__tysan_shadow_memory_address, align 8 +; CHECK-NEXT: [[X:%.*]] = alloca i32, i64 [[N:%.*]], align 1 +; CHECK-NEXT: [[TMP0:%.*]] = mul i64 [[N]], 4 +; CHECK-NEXT: [[TMP1:%.*]] = ptrtoint ptr [[X]] to i64 +; CHECK-NEXT: [[TMP2:%.*]] = and i64 [[TMP1]], [[APP_MEM_MASK]] +; CHECK-NEXT: [[TMP3:%.*]] = shl i64 [[TMP2]], 3 +; CHECK-NEXT: [[TMP4:%.*]] = add i64 [[TMP3]], [[SHADOW_BASE]] +; CHECK-NEXT: [[TMP5:%.*]] = inttoptr i64 [[TMP4]] to ptr +; CHECK-NEXT: [[TMP6:%.*]] = shl i64 [[TMP0]], 3 +; CHECK-NEXT: call void @llvm.memset.p0.i64(ptr align 8 [[TMP5]], i8 0, i64 [[TMP6]], i1 false) +; CHECK-NEXT: br label [[LOOP:%.*]] +; CHECK: loop: +; CHECK-NEXT: [[TMP7:%.*]] = mul i64 [[N]], 4 +; CHECK-NEXT: [[TMP8:%.*]] = ptrtoint ptr [[X]] to i64 +; CHECK-NEXT: [[TMP9:%.*]] = and i64 [[TMP8]], [[APP_MEM_MASK]] +; CHECK-NEXT: [[TMP10:%.*]] = shl i64 [[TMP9]], 3 +; CHECK-NEXT: [[TMP11:%.*]] = add i64 [[TMP10]], [[SHADOW_BASE]] +; CHECK-NEXT: [[TMP12:%.*]] = inttoptr i64 [[TMP11]] to ptr +; CHECK-NEXT: [[TMP13:%.*]] = shl i64 [[TMP7]], 3 +; CHECK-NEXT: call void @llvm.memset.p0.i64(ptr align 8 [[TMP12]], i8 0, i64 [[TMP13]], i1 false) +; CHECK-NEXT: call void @llvm.lifetime.start.p0(i64 -1, ptr [[X]]) +; CHECK-NEXT: call void @alloca_test_use(ptr [[X]]) +; CHECK-NEXT: [[TMP14:%.*]] = mul i64 [[N]], 4 +; CHECK-NEXT: [[TMP15:%.*]] = ptrtoint ptr [[X]] to i64 +; CHECK-NEXT: [[TMP16:%.*]] = and i64 [[TMP15]], [[APP_MEM_MASK]] +; CHECK-NEXT: [[TMP17:%.*]] = shl i64 [[TMP16]], 3 +; CHECK-NEXT: [[TMP18:%.*]] = add i64 [[TMP17]], [[SHADOW_BASE]] +; CHECK-NEXT: [[TMP19:%.*]] = inttoptr i64 [[TMP18]] to ptr +; CHECK-NEXT: [[TMP20:%.*]] = shl i64 [[TMP14]], 3 +; CHECK-NEXT: call void @llvm.memset.p0.i64(ptr align 8 [[TMP19]], i8 0, i64 [[TMP20]], i1 false) +; CHECK-NEXT: call void @llvm.lifetime.end.p0(i64 -1, ptr [[X]]) +; CHECK-NEXT: br i1 [[C:%.*]], label [[LOOP]], label [[EXIT:%.*]] +; CHECK: exit: +; CHECK-NEXT: ret void +; +entry: + %x = alloca i32, i64 %n, align 1 + br label %loop + +loop: + call void @llvm.lifetime.start.p0(i64 -1, ptr %x) + call void @alloca_test_use(ptr %x) + call void @llvm.lifetime.end.p0(i64 -1, ptr %x) + br i1 %c, label %loop, label %exit + +exit: + ret void +}