diff --git a/llvm/lib/Transforms/IPO/Attributor.cpp b/llvm/lib/Transforms/IPO/Attributor.cpp --- a/llvm/lib/Transforms/IPO/Attributor.cpp +++ b/llvm/lib/Transforms/IPO/Attributor.cpp @@ -4335,7 +4335,7 @@ if (IsMalloc) { if (auto *Size = dyn_cast(I.getOperand(0))) - if (Size->getValue().sle(MaxHeapToStackSize)) + if (Size->getValue().ule(MaxHeapToStackSize)) if (UsesCheck(I) || FreeCheck(I)) { MallocCalls.insert(&I); return true; @@ -4345,7 +4345,7 @@ if (auto *Num = dyn_cast(I.getOperand(0))) if (auto *Size = dyn_cast(I.getOperand(1))) if ((Size->getValue().umul_ov(Num->getValue(), Overflow)) - .sle(MaxHeapToStackSize)) + .ule(MaxHeapToStackSize)) if (!Overflow && (UsesCheck(I) || FreeCheck(I))) { MallocCalls.insert(&I); return true; diff --git a/llvm/test/Transforms/Attributor/heap_to_stack.ll b/llvm/test/Transforms/Attributor/heap_to_stack.ll --- a/llvm/test/Transforms/Attributor/heap_to_stack.ll +++ b/llvm/test/Transforms/Attributor/heap_to_stack.ll @@ -313,15 +313,27 @@ define i32 @test_sle() { %1 = tail call noalias i8* @malloc(i64 -1) - ; FIXME: This should not be transformed - ; CHECK: %1 = alloca i8, i64 -1 - ; CHECK-NEXT: @no_sync_func(i8* noalias nocapture nofree %1) + ; CHECK: %1 = tail call noalias i8* @malloc(i64 -1) + ; CHECK-NEXT: @no_sync_func(i8* noalias nofree %1) tail call void @no_sync_func(i8* %1) %2 = bitcast i8* %1 to i32* store i32 10, i32* %2 %3 = load i32, i32* %2 tail call void @free(i8* %1) - ; CHECK-NOT: tail call void @free(i8* noalias %1) + ; CHECK: tail call void @free(i8* noalias %1) + ret i32 %3 +} + +define i32 @test_overflow() { + %1 = tail call noalias i8* @calloc(i64 65537, i64 65537) + ; CHECK: %1 = tail call noalias i8* @calloc(i64 65537, i64 65537) + ; CHECK-NEXT: @no_sync_func(i8* noalias nofree %1) + tail call void @no_sync_func(i8* %1) + %2 = bitcast i8* %1 to i32* + store i32 10, i32* %2 + %3 = load i32, i32* %2 + tail call void @free(i8* %1) + ; CHECK: tail call void @free(i8* noalias %1) ret i32 %3 }