[sanitizer][msan] fix AArch64 vararg support for KMSAN

Cast StackSaveAreaPtr, GrRegSaveAreaPtr, VrRegSaveAreaPtr to pointers to
fix assertions in getShadowOriginPtrKernel().

Also add an isPointerTy() assertion to getShadowOriginPtrUserspace() to
ensure both the userspace and the kernel implementations of
getShadowOriginPtr() have the same expectations.

Fixes: #69738

Patch by Mark Johnston.

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

@@ -1717,6 +1717,12 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
   std::pair<Value *, Value *>
   getShadowOriginPtrUserspace(Value *Addr, IRBuilder<> &IRB, Type *ShadowTy,
                               MaybeAlign Alignment) {
+    VectorType *VectTy = dyn_cast<VectorType>(Addr->getType());
+    if (!VectTy) {
+      assert(Addr->getType()->isPointerTy());
+    } else {
+      assert(VectTy->getElementType()->isPointerTy());
+    }
     Type *IntptrTy = ptrToIntPtrType(Addr->getType());
     Value *ShadowOffset = getShadowPtrOffset(Addr, IRB);
     Value *ShadowLong = ShadowOffset;
@@ -5258,21 +5264,25 @@ struct VarArgAArch64Helper : public VarArgHelper {
       // we need to adjust the offset for both GR and VR fields based on
       // the __{gr,vr}_offs value (since they are stores based on incoming
       // named arguments).
+      Type *RegSaveAreaPtrTy = IRB.getInt8PtrTy();
 
       // Read the stack pointer from the va_list.
-      Value *StackSaveAreaPtr = getVAField64(IRB, VAListTag, 0);
+      Value *StackSaveAreaPtr =
+          IRB.CreateIntToPtr(getVAField64(IRB, VAListTag, 0), RegSaveAreaPtrTy);
 
       // Read both the __gr_top and __gr_off and add them up.
       Value *GrTopSaveAreaPtr = getVAField64(IRB, VAListTag, 8);
       Value *GrOffSaveArea = getVAField32(IRB, VAListTag, 24);
 
-      Value *GrRegSaveAreaPtr = IRB.CreateAdd(GrTopSaveAreaPtr, GrOffSaveArea);
+      Value *GrRegSaveAreaPtr = IRB.CreateIntToPtr(
+          IRB.CreateAdd(GrTopSaveAreaPtr, GrOffSaveArea), RegSaveAreaPtrTy);
 
       // Read both the __vr_top and __vr_off and add them up.
       Value *VrTopSaveAreaPtr = getVAField64(IRB, VAListTag, 16);
       Value *VrOffSaveArea = getVAField32(IRB, VAListTag, 28);
 
-      Value *VrRegSaveAreaPtr = IRB.CreateAdd(VrTopSaveAreaPtr, VrOffSaveArea);
+      Value *VrRegSaveAreaPtr = IRB.CreateIntToPtr(
+          IRB.CreateAdd(VrTopSaveAreaPtr, VrOffSaveArea), RegSaveAreaPtrTy);
 
       // It does not know how many named arguments is being used and, on the
       // callsite all the arguments were saved.  Since __gr_off is defined as

llvm/test/Instrumentation/MemorySanitizer/AArch64/vararg-kmsan.ll

@@ -0,0 +1,51 @@
+; RUN: opt < %s -S -passes=msan -msan-kernel=1 2>&1 | FileCheck %s
+
+target datalayout = "e-m:e-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64-unknown-linux-gnu"
+
+%struct.__va_list = type { ptr, ptr, ptr, i32, i32 }
+
+define i32 @foo(i32 %guard, ...) {
+  %vl = alloca %struct.__va_list, align 8
+  call void @llvm.lifetime.start.p0(i64 32, ptr %vl)
+  call void @llvm.va_start(ptr %vl)
+  call void @llvm.va_end(ptr %vl)
+  call void @llvm.lifetime.end.p0(i64 32, ptr %vl)
+  ret i32 0
+}
+
+; First check if the variadic shadow values are saved in stack with correct
+; size (192 is total of general purpose registers size, 64, plus total of
+; floating-point registers size, 128).
+
+; CHECK-LABEL: @foo
+; CHECK: [[A:%.*]] = load {{.*}} ptr %va_arg_overflow_size
+; CHECK: [[B:%.*]] = add i64 192, [[A]]
+; CHECK: alloca {{.*}} [[B]]
+
+; We expect three memcpy operations: one for the general purpose registers,
+; one for floating-point/SIMD ones, and one for thre remaining arguments.
+
+; Propagate the GR shadow values on for the va_list::__gp_top, adjust the 
+; offset in the __msan_va_arg_tls based on va_list:__gp_off, and finally
+; issue the memcpy.
+; CHECK: [[GRP:%.*]] = getelementptr inbounds i8, ptr {{%.*}}, i64 {{%.*}}
+; CHECK: [[GRSIZE:%.*]] = sub i64 64, {{%.*}}
+; CHECK: call void @llvm.memcpy.p0.p0.i64(ptr align 8 {{%.*}}, ptr align 8 [[GRP]], i64 [[GRSIZE]], i1 false)
+
+; Propagate the VR shadow values on for the va_list::__vr_top, adjust the 
+; offset in the __msan_va_arg_tls based on va_list:__vr_off, and finally
+; issue the memcpy.
+; CHECK: [[VRP:%.*]] = getelementptr inbounds i8, ptr {{%.*}}, i64 {{%.*}}
+; CHECK: [[VRSIZE:%.*]] = sub i64 128, {{%.*}}
+; CHECK: call void @llvm.memcpy.p0.p0.i64(ptr align 8 {{%.*}}, ptr align 8 [[VRP]], i64 [[VRSIZE]], i1 false)
+
+; Copy the remaining shadow values on the va_list::__stack position (it is
+; on the constant offset of 192 from __msan_va_arg_tls).
+; CHECK: [[STACK:%.*]] = getelementptr inbounds i8, ptr {{%.*}}, i32 192
+; CHECK: call void @llvm.memcpy.p0.p0.i64(ptr align 16 {{%.*}}, ptr align 16 [[STACK]], i64 {{%.*}}, i1 false)
+
+declare void @llvm.lifetime.start.p0(i64, ptr nocapture) #1
+declare void @llvm.va_start(ptr) #2
+declare void @llvm.va_end(ptr) #2
+declare void @llvm.lifetime.end.p0(i64, ptr nocapture) #1

llvm/test/Instrumentation/MemorySanitizer/X86/vararg.ll

@@ -1,4 +1,5 @@
 ; RUN: opt < %s -msan-check-access-address=0 -S -passes=msan 2>&1
+; RUN: opt < %s -msan-check-access-address=0 -S -passes=msan -msan-kernel=1 2>&1
 ; Test that code using va_start can be compiled on i386.
 
 target datalayout = "e-p:32:32:32-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:32:64-f32:32:32-f64:32:64-v64:64:64-v128:128:128-a0:0:64-f80:32:32-n8:16:32-S128"