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

[loladiro]

Summary

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.

[eugenis]

LGTM

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

void instrumentAlloca(AllocaInst &I, Instruction *InsPoint = nullptr) {
     if (!InsPoint)
       InsPoint = &I;
     IRBuilder<> IRB(InsPoint->getNextNode());
     IRB.SetCurrentDebugLocation(InsPoint->getDebugLoc());
     const DataLayout &DL = F.getParent()->getDataLayout();
     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);
     else
       poisonAllocaUserspace(I, IRB, Len);
   }
 
   void visitAllocaInst(AllocaInst &I) {

llvm/test/Instrumentation/MemorySanitizer/alloca.ll

 ; CHECK-LABEL: define void @array(
 ; 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 @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
   ret void
 }
 
 ; CHECK-LABEL: define void @array_non_const(
 ; CHECK: %[[A:.*]] = mul i64 4, %cnt
 ; 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
 
+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:
   %x = alloca i32, i64 5, align 4
   ret void
 }
 
 ; CHECK-LABEL: define void @unpoison_local(