Skip to content

Commit

Permalink
[msan][aarch64] Improve argument classification
Browse files Browse the repository at this point in the history
Arm64 use multiple registers (varg slots) to pass arrays.

Reviewers: kstoimenov, thurstond

Reviewed By: thurstond

Pull Request: #72728
  • Loading branch information
vitalybuka committed Nov 18, 2023
1 parent faaea79 commit 66e9429
Show file tree
Hide file tree
Showing 3 changed files with 46 additions and 30 deletions.
4 changes: 2 additions & 2 deletions compiler-rt/test/msan/vararg_shadow.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,8 @@
// Without -fno-sanitize-memory-param-retval we can't even pass poisoned values.
// RUN: %clangxx_msan -fno-sanitize-memory-param-retval -fsanitize-memory-track-origins=0 -O3 %s -o %t

// Nothing works yet.
// XFAIL: target={{(aarch64|loongarch64|mips|powerpc64).*}}
// FIXME: The rest is likely still broken.
// XFAIL: target={{(loongarch64|mips|powerpc64).*}}

#include <sanitizer/msan_interface.h>
#include <stdarg.h>
Expand Down
45 changes: 30 additions & 15 deletions llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -5071,14 +5071,27 @@ struct VarArgAArch64Helper : public VarArgHelperBase {
MemorySanitizerVisitor &MSV)
: VarArgHelperBase(F, MS, MSV, /*VAListTagSize=*/32) {}

ArgKind classifyArgument(Value *arg) {
Type *T = arg->getType();
if (T->isFPOrFPVectorTy())
return AK_FloatingPoint;
if ((T->isIntegerTy() && T->getPrimitiveSizeInBits() <= 64) ||
(T->isPointerTy()))
return AK_GeneralPurpose;
return AK_Memory;
// A very rough approximation of aarch64 argument classification rules.
std::pair<ArgKind, uint64_t> classifyArgument(Type *T) {
if (T->isIntOrPtrTy() && T->getPrimitiveSizeInBits() <= 64)
return {AK_GeneralPurpose, 1};
if (T->isFloatingPointTy() && T->getPrimitiveSizeInBits() <= 128)
return {AK_FloatingPoint, 1};

if (T->isArrayTy()) {
auto R = classifyArgument(T->getArrayElementType());
R.second *= T->getScalarType()->getArrayNumElements();
return R;
}

if (const FixedVectorType *FV = dyn_cast<FixedVectorType>(T)) {
auto R = classifyArgument(FV->getScalarType());
R.second *= FV->getNumElements();
return R;
}

LLVM_DEBUG(errs() << "Unknown vararg type: " << *T << "\n");
return {AK_Memory, 0};
}

// The instrumentation stores the argument shadow in a non ABI-specific
Expand All @@ -5098,20 +5111,22 @@ struct VarArgAArch64Helper : public VarArgHelperBase {
const DataLayout &DL = F.getParent()->getDataLayout();
for (const auto &[ArgNo, A] : llvm::enumerate(CB.args())) {
bool IsFixed = ArgNo < CB.getFunctionType()->getNumParams();
ArgKind AK = classifyArgument(A);
if (AK == AK_GeneralPurpose && GrOffset >= AArch64GrEndOffset)
auto [AK, RegNum] = classifyArgument(A->getType());
if (AK == AK_GeneralPurpose &&
(GrOffset + RegNum * 8) > AArch64GrEndOffset)
AK = AK_Memory;
if (AK == AK_FloatingPoint && VrOffset >= AArch64VrEndOffset)
if (AK == AK_FloatingPoint &&
(VrOffset + RegNum * 16) > AArch64VrEndOffset)
AK = AK_Memory;
Value *Base;
switch (AK) {
case AK_GeneralPurpose:
Base = getShadowPtrForVAArgument(A->getType(), IRB, GrOffset, 8);
GrOffset += 8;
Base = getShadowPtrForVAArgument(A->getType(), IRB, GrOffset);
GrOffset += 8 * RegNum;
break;
case AK_FloatingPoint:
Base = getShadowPtrForVAArgument(A->getType(), IRB, VrOffset, 8);
VrOffset += 16;
Base = getShadowPtrForVAArgument(A->getType(), IRB, VrOffset);
VrOffset += 16 * RegNum;
break;
case AK_Memory:
// Don't count fixed arguments in the overflow area - va_start will
Expand Down
27 changes: 14 additions & 13 deletions llvm/test/Instrumentation/MemorySanitizer/AArch64/vararg_shadow.ll
Original file line number Diff line number Diff line change
Expand Up @@ -304,8 +304,8 @@ define linkonce_odr dso_local void @_Z4testI10Int64Int64EvT_([2 x i64] %arg.coer
; CHECK-NEXT: store [2 x i64] [[TMP19]], ptr @__msan_param_tls, align 8
; CHECK-NEXT: store i32 0, ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_param_tls to i64), i64 16) to ptr), align 8
; CHECK-NEXT: store [2 x i64] [[TMP19]], ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_param_tls to i64), i64 24) to ptr), align 8
; CHECK-NEXT: store [2 x i64] [[TMP19]], ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_va_arg_tls to i64), i64 192) to ptr), align 8
; CHECK-NEXT: store i64 16, ptr @__msan_va_arg_overflow_size_tls, align 8
; CHECK-NEXT: store [2 x i64] [[TMP19]], ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_va_arg_tls to i64), i64 24) to ptr), align 8
; CHECK-NEXT: store i64 0, ptr @__msan_va_arg_overflow_size_tls, align 8
; CHECK-NEXT: call void ([2 x i64], i32, ...) @_Z5test2I10Int64Int64EvT_iz([2 x i64] [[DOTFCA_1_INSERT3]], i32 noundef 1, [2 x i64] [[DOTFCA_1_INSERT3]])
; CHECK-NEXT: ret void
;
Expand Down Expand Up @@ -370,8 +370,8 @@ define linkonce_odr dso_local void @_Z4testI12DoubleDoubleEvT_([2 x double] alig
; CHECK-NEXT: store [2 x i64] [[TMP19]], ptr @__msan_param_tls, align 8
; CHECK-NEXT: store i32 0, ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_param_tls to i64), i64 16) to ptr), align 8
; CHECK-NEXT: store [2 x i64] [[TMP19]], ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_param_tls to i64), i64 24) to ptr), align 8
; CHECK-NEXT: store [2 x i64] [[TMP19]], ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_va_arg_tls to i64), i64 192) to ptr), align 8
; CHECK-NEXT: store i64 16, ptr @__msan_va_arg_overflow_size_tls, align 8
; CHECK-NEXT: store [2 x i64] [[TMP19]], ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_va_arg_tls to i64), i64 96) to ptr), align 8
; CHECK-NEXT: store i64 0, ptr @__msan_va_arg_overflow_size_tls, align 8
; CHECK-NEXT: call void ([2 x double], i32, ...) @_Z5test2I12DoubleDoubleEvT_iz([2 x double] alignstack(8) [[DOTFCA_1_INSERT3]], i32 noundef 1, [2 x double] alignstack(8) [[DOTFCA_1_INSERT3]])
; CHECK-NEXT: ret void
;
Expand Down Expand Up @@ -466,8 +466,8 @@ define linkonce_odr dso_local void @_Z4testI7Double4EvT_([4 x double] alignstack
; CHECK-NEXT: store [4 x i64] [[TMP35]], ptr @__msan_param_tls, align 8
; CHECK-NEXT: store i32 0, ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_param_tls to i64), i64 32) to ptr), align 8
; CHECK-NEXT: store [4 x i64] [[TMP35]], ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_param_tls to i64), i64 40) to ptr), align 8
; CHECK-NEXT: store [4 x i64] [[TMP35]], ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_va_arg_tls to i64), i64 192) to ptr), align 8
; CHECK-NEXT: store i64 32, ptr @__msan_va_arg_overflow_size_tls, align 8
; CHECK-NEXT: store [4 x i64] [[TMP35]], ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_va_arg_tls to i64), i64 128) to ptr), align 8
; CHECK-NEXT: store i64 0, ptr @__msan_va_arg_overflow_size_tls, align 8
; CHECK-NEXT: call void ([4 x double], i32, ...) @_Z5test2I7Double4EvT_iz([4 x double] alignstack(8) [[DOTFCA_3_INSERT7]], i32 noundef 1, [4 x double] alignstack(8) [[DOTFCA_3_INSERT7]])
; CHECK-NEXT: ret void
;
Expand Down Expand Up @@ -542,8 +542,8 @@ define linkonce_odr dso_local void @_Z4testI11DoubleFloatEvT_([2 x i64] %arg.coe
; CHECK-NEXT: store [2 x i64] [[TMP19]], ptr @__msan_param_tls, align 8
; CHECK-NEXT: store i32 0, ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_param_tls to i64), i64 16) to ptr), align 8
; CHECK-NEXT: store [2 x i64] [[TMP19]], ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_param_tls to i64), i64 24) to ptr), align 8
; CHECK-NEXT: store [2 x i64] [[TMP19]], ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_va_arg_tls to i64), i64 192) to ptr), align 8
; CHECK-NEXT: store i64 16, ptr @__msan_va_arg_overflow_size_tls, align 8
; CHECK-NEXT: store [2 x i64] [[TMP19]], ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_va_arg_tls to i64), i64 24) to ptr), align 8
; CHECK-NEXT: store i64 0, ptr @__msan_va_arg_overflow_size_tls, align 8
; CHECK-NEXT: call void ([2 x i64], i32, ...) @_Z5test2I11DoubleFloatEvT_iz([2 x i64] [[DOTFCA_1_INSERT3]], i32 noundef 1, [2 x i64] [[DOTFCA_1_INSERT3]])
; CHECK-NEXT: ret void
;
Expand Down Expand Up @@ -608,8 +608,8 @@ define linkonce_odr dso_local void @_Z4testI11LongDouble2EvT_([2 x fp128] aligns
; CHECK-NEXT: store [2 x i128] [[TMP19]], ptr @__msan_param_tls, align 8
; CHECK-NEXT: store i32 0, ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_param_tls to i64), i64 32) to ptr), align 8
; CHECK-NEXT: store [2 x i128] [[TMP19]], ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_param_tls to i64), i64 40) to ptr), align 8
; CHECK-NEXT: store [2 x i128] [[TMP19]], ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_va_arg_tls to i64), i64 192) to ptr), align 8
; CHECK-NEXT: store i64 32, ptr @__msan_va_arg_overflow_size_tls, align 8
; CHECK-NEXT: store [2 x i128] [[TMP19]], ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_va_arg_tls to i64), i64 96) to ptr), align 8
; CHECK-NEXT: store i64 0, ptr @__msan_va_arg_overflow_size_tls, align 8
; CHECK-NEXT: call void ([2 x fp128], i32, ...) @_Z5test2I11LongDouble2EvT_iz([2 x fp128] alignstack(16) [[DOTFCA_1_INSERT5]], i32 noundef 1, [2 x fp128] alignstack(16) [[DOTFCA_1_INSERT5]])
; CHECK-NEXT: ret void
;
Expand Down Expand Up @@ -704,8 +704,8 @@ define linkonce_odr dso_local void @_Z4testI11LongDouble4EvT_([4 x fp128] aligns
; CHECK-NEXT: store [4 x i128] [[TMP35]], ptr @__msan_param_tls, align 8
; CHECK-NEXT: store i32 0, ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_param_tls to i64), i64 64) to ptr), align 8
; CHECK-NEXT: store [4 x i128] [[TMP35]], ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_param_tls to i64), i64 72) to ptr), align 8
; CHECK-NEXT: store [4 x i128] [[TMP35]], ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_va_arg_tls to i64), i64 192) to ptr), align 8
; CHECK-NEXT: store i64 64, ptr @__msan_va_arg_overflow_size_tls, align 8
; CHECK-NEXT: store [4 x i128] [[TMP35]], ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_va_arg_tls to i64), i64 128) to ptr), align 8
; CHECK-NEXT: store i64 0, ptr @__msan_va_arg_overflow_size_tls, align 8
; CHECK-NEXT: call void ([4 x fp128], i32, ...) @_Z5test2I11LongDouble4EvT_iz([4 x fp128] alignstack(16) [[DOTFCA_3_INSERT7]], i32 noundef 1, [4 x fp128] alignstack(16) [[DOTFCA_3_INSERT7]])
; CHECK-NEXT: ret void
;
Expand Down Expand Up @@ -1850,6 +1850,7 @@ define linkonce_odr dso_local void @_Z4test2I11LongDouble4EvT_([4 x fp128] align
; CHECK-NEXT: store [4 x i128] [[TMP35]], ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_param_tls to i64), i64 584) to ptr), align 8
; CHECK-NEXT: store [4 x i128] [[TMP35]], ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_param_tls to i64), i64 648) to ptr), align 8
; CHECK-NEXT: store [4 x i128] [[TMP35]], ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_param_tls to i64), i64 712) to ptr), align 8
; CHECK-NEXT: store [4 x i128] [[TMP35]], ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_va_arg_tls to i64), i64 128) to ptr), align 8
; CHECK-NEXT: store [4 x i128] [[TMP35]], ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_va_arg_tls to i64), i64 192) to ptr), align 8
; CHECK-NEXT: store [4 x i128] [[TMP35]], ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_va_arg_tls to i64), i64 256) to ptr), align 8
; CHECK-NEXT: store [4 x i128] [[TMP35]], ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_va_arg_tls to i64), i64 320) to ptr), align 8
Expand All @@ -1860,7 +1861,7 @@ define linkonce_odr dso_local void @_Z4test2I11LongDouble4EvT_([4 x fp128] align
; CHECK-NEXT: store [4 x i128] [[TMP35]], ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_va_arg_tls to i64), i64 640) to ptr), align 8
; CHECK-NEXT: store [4 x i128] [[TMP35]], ptr inttoptr (i64 add (i64 ptrtoint (ptr @__msan_va_arg_tls to i64), i64 704) to ptr), align 8
; CHECK-NEXT: call void @llvm.memset.p0.i32(ptr align 8 inttoptr (i64 add (i64 ptrtoint (ptr @__msan_va_arg_tls to i64), i64 768) to ptr), i8 0, i32 32, i1 false)
; CHECK-NEXT: store i64 1280, ptr @__msan_va_arg_overflow_size_tls, align 8
; CHECK-NEXT: store i64 1216, ptr @__msan_va_arg_overflow_size_tls, align 8
; CHECK-NEXT: call void ([4 x fp128], i32, ...) @_Z5test2I11LongDouble4EvT_iz([4 x fp128] alignstack(16) [[DOTFCA_3_INSERT121]], i32 noundef 20, [4 x fp128] alignstack(16) [[DOTFCA_3_INSERT121]], [4 x fp128] alignstack(16) [[DOTFCA_3_INSERT121]], [4 x fp128] alignstack(16) [[DOTFCA_3_INSERT121]], [4 x fp128] alignstack(16) [[DOTFCA_3_INSERT121]], [4 x fp128] alignstack(16) [[DOTFCA_3_INSERT121]], [4 x fp128] alignstack(16) [[DOTFCA_3_INSERT121]], [4 x fp128] alignstack(16) [[DOTFCA_3_INSERT121]], [4 x fp128] alignstack(16) [[DOTFCA_3_INSERT121]], [4 x fp128] alignstack(16) [[DOTFCA_3_INSERT121]], [4 x fp128] alignstack(16) [[DOTFCA_3_INSERT121]], [4 x fp128] alignstack(16) [[DOTFCA_3_INSERT121]], [4 x fp128] alignstack(16) [[DOTFCA_3_INSERT121]], [4 x fp128] alignstack(16) [[DOTFCA_3_INSERT121]], [4 x fp128] alignstack(16) [[DOTFCA_3_INSERT121]], [4 x fp128] alignstack(16) [[DOTFCA_3_INSERT121]], [4 x fp128] alignstack(16) [[DOTFCA_3_INSERT121]], [4 x fp128] alignstack(16) [[DOTFCA_3_INSERT121]], [4 x fp128] alignstack(16) [[DOTFCA_3_INSERT121]], [4 x fp128] alignstack(16) [[DOTFCA_3_INSERT121]], [4 x fp128] alignstack(16) [[DOTFCA_3_INSERT121]])
; CHECK-NEXT: ret void
;
Expand Down

0 comments on commit 66e9429

Please sign in to comment.