[llvm] 30d7d74 - [MSAN] Handle array alloca with non-i64 size specification

[Keno Fischer via llvm-commits]

Author: Keno Fischer
Date: 2022-08-24T03:24:21Z
New Revision: 30d7d74d5c358d0040a94735c3a46aa00ca22841

URL: https://github.com/llvm/llvm-project/commit/30d7d74d5c358d0040a94735c3a46aa00ca22841
DIFF: https://github.com/llvm/llvm-project/commit/30d7d74d5c358d0040a94735c3a46aa00ca22841.diff

LOG: [MSAN] Handle array alloca with non-i64 size specification

The array size specification of the an alloca can be any integer,
so zext or trunc it to intptr before attempting to multiply it
with an intptr constant.

Reviewed By: eugenis

Differential Revision: https://reviews.llvm.org/D131846

Added: 
    

Modified: 
    llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
    llvm/test/Instrumentation/MemorySanitizer/alloca.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
index 6b155ec6d1ce9..7d66b423deae6 100644
--- a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
@@ -3948,7 +3948,8 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
     uint64_t TypeSize = DL.getTypeAllocSize(I.getAllocatedType());
     Value *Len = ConstantInt::get(MS.IntptrTy, TypeSize);
     if (I.isArrayAllocation())
-      Len = IRB.CreateMul(Len, I.getArraySize());
+      Len = IRB.CreateMul(Len,
+                          IRB.CreateZExtOrTrunc(I.getArraySize(), MS.IntptrTy));
 
     if (MS.CompileKernel)
       poisonAllocaKmsan(I, IRB, Len);

diff  --git a/llvm/test/Instrumentation/MemorySanitizer/alloca.ll b/llvm/test/Instrumentation/MemorySanitizer/alloca.ll
index 894df189397f0..0149a2fc02798 100644
--- a/llvm/test/Instrumentation/MemorySanitizer/alloca.ll
+++ b/llvm/test/Instrumentation/MemorySanitizer/alloca.ll
@@ -65,6 +65,20 @@ entry:
 ; KMSAN: call void @__msan_poison_alloca(i8* {{.*}}, i64 20,
 ; CHECK: ret void
 
+define void @array32() sanitize_memory {
+entry:
+  %x = alloca i32, i32 5, align 4
+  ret void
+}
+
+; CHECK-LABEL: define void @array32(
+; INLINE: call void @llvm.memset.p0i8.i64(i8* align 4 {{.*}}, i8 -1, i64 20, i1 false)
+; CALL: call void @__msan_poison_stack(i8* {{.*}}, i64 20)
+; ORIGIN: call void @__msan_set_alloca_origin_with_descr(i8* {{.*}}, i64 20,
+; ORIGIN-LEAN: call void @__msan_set_alloca_origin_no_descr(i8* {{.*}}, i64 20,
+; KMSAN: call void @__msan_poison_alloca(i8* {{.*}}, i64 20,
+; CHECK: ret void
+
 define void @array_non_const(i64 %cnt) sanitize_memory {
 entry:
   %x = alloca i32, i64 %cnt, align 4
@@ -80,6 +94,22 @@ entry:
 ; KMSAN: call void @__msan_poison_alloca(i8* {{.*}}, i64 %[[A]],
 ; CHECK: ret void
 
+define void @array_non_const32(i32 %cnt) sanitize_memory {
+entry:
+  %x = alloca i32, i32 %cnt, align 4
+  ret void
+}
+
+; CHECK-LABEL: define void @array_non_const32(
+; CHECK: %[[Z:.*]] = zext i32 %cnt to i64
+; CHECK: %[[A:.*]] = mul i64 4, %[[Z]]
+; INLINE: call void @llvm.memset.p0i8.i64(i8* align 4 {{.*}}, i8 -1, i64 %[[A]], i1 false)
+; CALL: call void @__msan_poison_stack(i8* {{.*}}, i64 %[[A]])
+; ORIGIN: call void @__msan_set_alloca_origin_with_descr(i8* {{.*}}, i64 %[[A]],
+; ORIGIN-LEAN: call void @__msan_set_alloca_origin_no_descr(i8* {{.*}}, i64 %[[A]],
+; KMSAN: call void @__msan_poison_alloca(i8* {{.*}}, i64 %[[A]],
+; CHECK: ret void
+
 ; Check that the local is unpoisoned in the absence of sanitize_memory
 define void @unpoison_local() {
 entry: