[llvm] [msan] Fix vscale alloca poisoning (PR #90912)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Vitaly Buka (vitalybuka)

<details>
<summary>Changes</summary>



---
Full diff: https://github.com/llvm/llvm-project/pull/90912.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp (+2-2) 
- (modified) llvm/test/Instrumentation/MemorySanitizer/vscale.ll (+58-1) 


``````````diff
diff --git a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
index 2b504b893ddb0d..2cd005125a9cca 100644
--- a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
@@ -4458,8 +4458,8 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
       InsPoint = &I;
     NextNodeIRBuilder IRB(InsPoint);
     const DataLayout &DL = F.getParent()->getDataLayout();
-    uint64_t TypeSize = DL.getTypeAllocSize(I.getAllocatedType());
-    Value *Len = ConstantInt::get(MS.IntptrTy, TypeSize);
+    TypeSize TS = DL.getTypeAllocSize(I.getAllocatedType());
+    Value *Len = IRB.CreateTypeSize(IRB.getInt32Ty(), TS);
     if (I.isArrayAllocation())
       Len = IRB.CreateMul(Len,
                           IRB.CreateZExtOrTrunc(I.getArraySize(), MS.IntptrTy));
diff --git a/llvm/test/Instrumentation/MemorySanitizer/vscale.ll b/llvm/test/Instrumentation/MemorySanitizer/vscale.ll
index e1a4a9b7aa68e5..8e6e110557e5a3 100644
--- a/llvm/test/Instrumentation/MemorySanitizer/vscale.ll
+++ b/llvm/test/Instrumentation/MemorySanitizer/vscale.ll
@@ -173,7 +173,7 @@ define void @test_param(ptr %a, ptr %b) sanitize_memory {
 ; CHECK-NEXT:    [[_MSCMP:%.*]] = icmp ne i32 [[TMP6]], 0
 ; CHECK-NEXT:    br i1 [[_MSCMP]], label [[TMP7:%.*]], label [[TMP8:%.*]], !prof [[PROF0:![0-9]+]]
 ; CHECK:       7:
-; CHECK-NEXT:    call void @__msan_warning_noreturn() #[[ATTR4:[0-9]+]]
+; CHECK-NEXT:    call void @__msan_warning_noreturn() #[[ATTR5:[0-9]+]]
 ; CHECK-NEXT:    unreachable
 ; CHECK:       8:
 ; CHECK-NEXT:    call void @fn_param(<vscale x 2 x float> [[TMP2]], ptr [[B]])
@@ -183,6 +183,63 @@ define void @test_param(ptr %a, ptr %b) sanitize_memory {
   call void @fn_param(<vscale x 2 x float> %1, ptr %b)
   ret void
 }
+
+define void @test_alloca1() sanitize_memory {
+; CHECK-LABEL: define void @test_alloca1(
+; CHECK-SAME: ) #[[ATTR0]] {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @llvm.donothing()
+; CHECK-NEXT:    [[X:%.*]] = alloca <vscale x 64 x i1>, align 4
+; CHECK-NEXT:    [[TMP0:%.*]] = call i32 @llvm.vscale.i32()
+; CHECK-NEXT:    [[TMP1:%.*]] = mul i32 [[TMP0]], 8
+; CHECK-NEXT:    [[TMP2:%.*]] = ptrtoint ptr [[X]] to i64
+; CHECK-NEXT:    [[TMP3:%.*]] = xor i64 [[TMP2]], 87960930222080
+; CHECK-NEXT:    [[TMP4:%.*]] = inttoptr i64 [[TMP3]] to ptr
+; CHECK-NEXT:    call void @llvm.memset.p0.i32(ptr align 4 [[TMP4]], i8 -1, i32 [[TMP1]], i1 false)
+; CHECK-NEXT:    ret void
+;
+entry:
+  %x = alloca <vscale x 64 x i1>, align 4
+  ret void
+}
+
+define void @test_alloca2() sanitize_memory {
+; CHECK-LABEL: define void @test_alloca2(
+; CHECK-SAME: ) #[[ATTR0]] {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @llvm.donothing()
+; CHECK-NEXT:    [[X:%.*]] = alloca <vscale x 64 x double>, align 4
+; CHECK-NEXT:    [[TMP0:%.*]] = call i32 @llvm.vscale.i32()
+; CHECK-NEXT:    [[TMP1:%.*]] = mul i32 [[TMP0]], 512
+; CHECK-NEXT:    [[TMP2:%.*]] = ptrtoint ptr [[X]] to i64
+; CHECK-NEXT:    [[TMP3:%.*]] = xor i64 [[TMP2]], 87960930222080
+; CHECK-NEXT:    [[TMP4:%.*]] = inttoptr i64 [[TMP3]] to ptr
+; CHECK-NEXT:    call void @llvm.memset.p0.i32(ptr align 4 [[TMP4]], i8 -1, i32 [[TMP1]], i1 false)
+; CHECK-NEXT:    ret void
+;
+entry:
+  %x = alloca <vscale x 64 x double>, align 4
+  ret void
+}
+
+define void @test_alloca3() sanitize_memory {
+; CHECK-LABEL: define void @test_alloca3(
+; CHECK-SAME: ) #[[ATTR0]] {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @llvm.donothing()
+; CHECK-NEXT:    [[X:%.*]] = alloca <vscale x 1 x i1>, align 4
+; CHECK-NEXT:    [[TMP0:%.*]] = call i32 @llvm.vscale.i32()
+; CHECK-NEXT:    [[TMP1:%.*]] = ptrtoint ptr [[X]] to i64
+; CHECK-NEXT:    [[TMP2:%.*]] = xor i64 [[TMP1]], 87960930222080
+; CHECK-NEXT:    [[TMP3:%.*]] = inttoptr i64 [[TMP2]] to ptr
+; CHECK-NEXT:    call void @llvm.memset.p0.i32(ptr align 4 [[TMP3]], i8 -1, i32 [[TMP0]], i1 false)
+; CHECK-NEXT:    ret void
+;
+entry:
+  %x = alloca <vscale x 1 x i1>, align 4
+  ret void
+}
+
 ;.
 ; CHECK: [[PROF0]] = !{!"branch_weights", i32 1, i32 1048575}
 ;.

``````````

</details>


https://github.com/llvm/llvm-project/pull/90912