Skip to content

Commit ddf5725

Browse files
chestnykhchestnykh
authored
[nsan] Emit calls to optimized functions (llvm#98900)
As previously noted in nsan.cpp we can implement optimized variants of `__nsan_copy_values` and `__nsan_set_value_unknown` if a memory operation size is known. Now the instrumentation creates calls to optimized functions if there is 4, 8 or 16-byte memory operation like `memset(X, value, 4/8/16)` or `memcpy(dst, src, 4/8/16)` nsan.cpp provides definitions of the optimized functions.
1 parent 3993a47 commit ddf5725

File tree

3 files changed

+156
-25
lines changed

3 files changed

+156
-25
lines changed

‎compiler-rt/lib/nsan/nsan.cpp

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -54,8 +54,6 @@ using namespace __nsan;
5454
constexpr int kMaxVectorWidth = 8;
5555

5656
// When copying application memory, we also copy its shadow and shadow type.
57-
// FIXME: We could provide fixed-size versions that would nicely
58-
// vectorize for known sizes.
5957
extern "C" SANITIZER_INTERFACE_ATTRIBUTE void
6058
__nsan_copy_values(const u8 *daddr, const u8 *saddr, uptr size) {
6159
internal_memmove((void *)GetShadowTypeAddrFor(daddr),
@@ -64,13 +62,33 @@ __nsan_copy_values(const u8 *daddr, const u8 *saddr, uptr size) {
6462
size * kShadowScale);
6563
}
6664

67-
// FIXME: We could provide fixed-size versions that would nicely
68-
// vectorize for known sizes.
65+
#define NSAN_COPY_VALUES_N(N) \
66+
extern "C" SANITIZER_INTERFACE_ATTRIBUTE void __nsan_copy_##N( \
67+
const u8 *daddr, const u8 *saddr) { \
68+
__builtin_memmove((void *)GetShadowTypeAddrFor(daddr), \
69+
GetShadowTypeAddrFor(saddr), N); \
70+
__builtin_memmove((void *)GetShadowAddrFor(daddr), \
71+
GetShadowAddrFor(saddr), N * kShadowScale); \
72+
}
73+
74+
NSAN_COPY_VALUES_N(4)
75+
NSAN_COPY_VALUES_N(8)
76+
NSAN_COPY_VALUES_N(16)
77+
6978
extern "C" SANITIZER_INTERFACE_ATTRIBUTE void
7079
__nsan_set_value_unknown(const u8 *addr, uptr size) {
7180
internal_memset((void *)GetShadowTypeAddrFor(addr), 0, size);
7281
}
7382

83+
#define NSAN_SET_VALUE_UNKNOWN_N(N) \
84+
extern "C" SANITIZER_INTERFACE_ATTRIBUTE void __nsan_set_value_unknown_##N( \
85+
const u8 *daddr) { \
86+
__builtin_memset((void *)GetShadowTypeAddrFor(daddr), 0, N); \
87+
}
88+
89+
NSAN_SET_VALUE_UNKNOWN_N(4)
90+
NSAN_SET_VALUE_UNKNOWN_N(8)
91+
NSAN_SET_VALUE_UNKNOWN_N(16)
7492

7593
const char *FTInfo<float>::kCppTypeName = "float";
7694
const char *FTInfo<double>::kCppTypeName = "double";

‎llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp

Lines changed: 102 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -493,6 +493,60 @@ class ValueToShadowMap {
493493
DenseMap<Value *, Value *> Map;
494494
};
495495

496+
class NsanMemOpFn {
497+
public:
498+
NsanMemOpFn(Module &M, ArrayRef<StringRef> Sized, StringRef Fallback,
499+
size_t NumArgs);
500+
FunctionCallee getFunctionFor(uint64_t MemOpSize) const;
501+
FunctionCallee getFallback() const;
502+
503+
private:
504+
SmallVector<FunctionCallee> Funcs;
505+
size_t NumSizedFuncs;
506+
};
507+
508+
NsanMemOpFn::NsanMemOpFn(Module &M, ArrayRef<StringRef> Sized,
509+
StringRef Fallback, size_t NumArgs) {
510+
LLVMContext &Ctx = M.getContext();
511+
AttributeList Attr;
512+
Attr = Attr.addFnAttribute(Ctx, Attribute::NoUnwind);
513+
Type *PtrTy = PointerType::getUnqual(Ctx);
514+
Type *VoidTy = Type::getVoidTy(Ctx);
515+
IntegerType *IntptrTy = M.getDataLayout().getIntPtrType(Ctx);
516+
FunctionType *SizedFnTy = nullptr;
517+
518+
NumSizedFuncs = Sized.size();
519+
520+
// First entry is fallback function
521+
if (NumArgs == 3) {
522+
Funcs.push_back(
523+
M.getOrInsertFunction(Fallback, Attr, VoidTy, PtrTy, PtrTy, IntptrTy));
524+
SizedFnTy = FunctionType::get(VoidTy, {PtrTy, PtrTy}, false);
525+
} else if (NumArgs == 2) {
526+
Funcs.push_back(
527+
M.getOrInsertFunction(Fallback, Attr, VoidTy, PtrTy, IntptrTy));
528+
SizedFnTy = FunctionType::get(VoidTy, {PtrTy}, false);
529+
} else {
530+
assert(!"Unexpected value of sized functions arguments");
531+
}
532+
533+
for (size_t i = 0; i < NumSizedFuncs; ++i)
534+
Funcs.push_back(M.getOrInsertFunction(Sized[i], SizedFnTy, Attr));
535+
}
536+
537+
FunctionCallee NsanMemOpFn::getFunctionFor(uint64_t MemOpSize) const {
538+
// Now `getFunctionFor` operates on `Funcs` of size 4 (at least) and the
539+
// following code assumes that the number of functions in `Func` is sufficient
540+
assert(NumSizedFuncs >= 3 && "Unexpected number of sized functions");
541+
542+
size_t Idx =
543+
MemOpSize == 4 ? 1 : (MemOpSize == 8 ? 2 : (MemOpSize == 16 ? 3 : 0));
544+
545+
return Funcs[Idx];
546+
}
547+
548+
FunctionCallee NsanMemOpFn::getFallback() const { return Funcs[0]; }
549+
496550
/// Instantiating NumericalStabilitySanitizer inserts the nsan runtime library
497551
/// API function declarations into the module if they don't exist already.
498552
/// Instantiating ensures the __nsan_init function is in the list of global
@@ -550,12 +604,16 @@ class NumericalStabilitySanitizer {
550604
LLVMContext &Context;
551605
MappingConfig Config;
552606
IntegerType *IntptrTy = nullptr;
607+
608+
// TODO: Use std::array instead?
553609
FunctionCallee NsanGetShadowPtrForStore[FTValueType::kNumValueTypes] = {};
554610
FunctionCallee NsanGetShadowPtrForLoad[FTValueType::kNumValueTypes] = {};
555611
FunctionCallee NsanCheckValue[FTValueType::kNumValueTypes] = {};
556612
FunctionCallee NsanFCmpFail[FTValueType::kNumValueTypes] = {};
557-
FunctionCallee NsanCopyValues;
558-
FunctionCallee NsanSetValueUnknown;
613+
614+
NsanMemOpFn NsanCopyFns;
615+
NsanMemOpFn NsanSetUnknownFns;
616+
559617
FunctionCallee NsanGetRawShadowTypePtr;
560618
FunctionCallee NsanGetRawShadowPtr;
561619
GlobalValue *NsanShadowRetTag = nullptr;
@@ -598,7 +656,14 @@ static GlobalValue *createThreadLocalGV(const char *Name, Module &M, Type *Ty) {
598656
}
599657

600658
NumericalStabilitySanitizer::NumericalStabilitySanitizer(Module &M)
601-
: DL(M.getDataLayout()), Context(M.getContext()), Config(Context) {
659+
: DL(M.getDataLayout()), Context(M.getContext()), Config(Context),
660+
NsanCopyFns(M, {"__nsan_copy_4", "__nsan_copy_8", "__nsan_copy_16"},
661+
"__nsan_copy_values", /*NumArgs=*/3),
662+
NsanSetUnknownFns(M,
663+
{"__nsan_set_value_unknown_4",
664+
"__nsan_set_value_unknown_8",
665+
"__nsan_set_value_unknown_16"},
666+
"__nsan_set_value_unknown", /*NumArgs=*/2) {
602667
IntptrTy = DL.getIntPtrType(Context);
603668
Type *PtrTy = PointerType::getUnqual(Context);
604669
Type *Int32Ty = Type::getInt32Ty(Context);
@@ -634,11 +699,6 @@ NumericalStabilitySanitizer::NumericalStabilitySanitizer(Module &M)
634699
Attr, VoidTy, VTTy, VTTy, ShadowTy, ShadowTy, Int32Ty, Int1Ty, Int1Ty);
635700
}
636701

637-
NsanCopyValues = M.getOrInsertFunction("__nsan_copy_values", Attr, VoidTy,
638-
PtrTy, PtrTy, IntptrTy);
639-
NsanSetValueUnknown = M.getOrInsertFunction("__nsan_set_value_unknown", Attr,
640-
VoidTy, PtrTy, IntptrTy);
641-
642702
// TODO: Add attributes nofree, nosync, readnone, readonly,
643703
NsanGetRawShadowTypePtr = M.getOrInsertFunction(
644704
"__nsan_internal_get_raw_shadow_type_ptr", Attr, PtrTy, PtrTy);
@@ -1880,7 +1940,7 @@ void NumericalStabilitySanitizer::propagateNonFTStore(
18801940
}
18811941
}
18821942
// All other stores just reset the shadow value to unknown.
1883-
Builder.CreateCall(NsanSetValueUnknown, {Dst, ValueSize});
1943+
Builder.CreateCall(NsanSetUnknownFns.getFallback(), {Dst, ValueSize});
18841944
}
18851945

18861946
void NumericalStabilitySanitizer::propagateShadowValues(
@@ -2123,21 +2183,45 @@ bool NumericalStabilitySanitizer::sanitizeFunction(
21232183
return !ValueToShadow.empty();
21242184
}
21252185

2186+
static uint64_t GetMemOpSize(Value *V) {
2187+
uint64_t OpSize = 0;
2188+
if (Constant *C = dyn_cast<Constant>(V)) {
2189+
auto *CInt = dyn_cast<ConstantInt>(C);
2190+
if (CInt && CInt->getValue().getBitWidth() <= 64)
2191+
OpSize = CInt->getValue().getZExtValue();
2192+
}
2193+
2194+
return OpSize;
2195+
}
2196+
21262197
// Instrument the memory intrinsics so that they properly modify the shadow
21272198
// memory.
21282199
bool NumericalStabilitySanitizer::instrumentMemIntrinsic(MemIntrinsic *MI) {
21292200
IRBuilder<> Builder(MI);
21302201
if (auto *M = dyn_cast<MemSetInst>(MI)) {
2131-
Builder.CreateCall(
2132-
NsanSetValueUnknown,
2133-
{/*Address=*/M->getArgOperand(0),
2134-
/*Size=*/Builder.CreateIntCast(M->getArgOperand(2), IntptrTy, false)});
2202+
FunctionCallee SetUnknownFn =
2203+
NsanSetUnknownFns.getFunctionFor(GetMemOpSize(M->getArgOperand(2)));
2204+
if (SetUnknownFn.getFunctionType()->getNumParams() == 1)
2205+
Builder.CreateCall(SetUnknownFn, {/*Address=*/M->getArgOperand(0)});
2206+
else
2207+
Builder.CreateCall(SetUnknownFn,
2208+
{/*Address=*/M->getArgOperand(0),
2209+
/*Size=*/Builder.CreateIntCast(M->getArgOperand(2),
2210+
IntptrTy, false)});
2211+
21352212
} else if (auto *M = dyn_cast<MemTransferInst>(MI)) {
2136-
Builder.CreateCall(
2137-
NsanCopyValues,
2138-
{/*Destination=*/M->getArgOperand(0),
2139-
/*Source=*/M->getArgOperand(1),
2140-
/*Size=*/Builder.CreateIntCast(M->getArgOperand(2), IntptrTy, false)});
2213+
FunctionCallee CopyFn =
2214+
NsanCopyFns.getFunctionFor(GetMemOpSize(M->getArgOperand(2)));
2215+
2216+
if (CopyFn.getFunctionType()->getNumParams() == 2)
2217+
Builder.CreateCall(CopyFn, {/*Destination=*/M->getArgOperand(0),
2218+
/*Source=*/M->getArgOperand(1)});
2219+
else
2220+
Builder.CreateCall(CopyFn, {/*Destination=*/M->getArgOperand(0),
2221+
/*Source=*/M->getArgOperand(1),
2222+
/*Size=*/
2223+
Builder.CreateIntCast(M->getArgOperand(2),
2224+
IntptrTy, false)});
21412225
}
21422226
return false;
21432227
}

‎llvm/test/Instrumentation/NumericalStabilitySanitizer/memory.ll

Lines changed: 32 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,15 +7,24 @@ target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f3
77

88
declare void @llvm.memcpy.p0i8.p0i8.i64(i8*, i8*, i64, i1)
99

10-
define void @call_memcpy_intrinsic(i8* nonnull align 8 dereferenceable(16) %a, i8* nonnull align 8 dereferenceable(16) %b) sanitize_numerical_stability {
11-
; CHECK-LABEL: @call_memcpy_intrinsic(
10+
define void @call_memcpy_intrinsics(i8* nonnull align 8 dereferenceable(16) %a, i8* nonnull align 8 dereferenceable(16) %b) sanitize_numerical_stability {
11+
; CHECK-LABEL: @call_memcpy_intrinsics(
1212
; CHECK-NEXT: entry:
13-
; CHECK-NEXT: call void @__nsan_copy_values(ptr [[A:%.*]], ptr [[B:%.*]], i64 16)
13+
; CHECK-NEXT: call void @__nsan_copy_4(ptr [[A:%.*]], ptr [[B:%.*]])
14+
; CHECK-NEXT: tail call void @llvm.memcpy.p0.p0.i64(ptr nonnull align 8 dereferenceable(16) [[A]], ptr nonnull align 8 dereferenceable(16) [[B]], i64 4, i1 false)
15+
; CHECK-NEXT: call void @__nsan_copy_8(ptr [[A:%.*]], ptr [[B:%.*]])
16+
; CHECK-NEXT: tail call void @llvm.memcpy.p0.p0.i64(ptr nonnull align 8 dereferenceable(16) [[A]], ptr nonnull align 8 dereferenceable(16) [[B]], i64 8, i1 false)
17+
; CHECK-NEXT: call void @__nsan_copy_16(ptr [[A:%.*]], ptr [[B:%.*]])
1418
; CHECK-NEXT: tail call void @llvm.memcpy.p0.p0.i64(ptr nonnull align 8 dereferenceable(16) [[A]], ptr nonnull align 8 dereferenceable(16) [[B]], i64 16, i1 false)
19+
; CHECK-NEXT: call void @__nsan_copy_values(ptr [[A:%.*]], ptr [[B:%.*]], i64 15)
20+
; CHECK-NEXT: tail call void @llvm.memcpy.p0.p0.i64(ptr nonnull align 8 dereferenceable(16) [[A]], ptr nonnull align 8 dereferenceable(16) [[B]], i64 15, i1 false)
1521
; CHECK-NEXT: ret void
1622
;
1723
entry:
24+
tail call void @llvm.memcpy.p0i8.p0i8.i64(ptr nonnull align 8 dereferenceable(16) %a, ptr nonnull align 8 dereferenceable(16) %b, i64 4, i1 false)
25+
tail call void @llvm.memcpy.p0i8.p0i8.i64(ptr nonnull align 8 dereferenceable(16) %a, ptr nonnull align 8 dereferenceable(16) %b, i64 8, i1 false)
1826
tail call void @llvm.memcpy.p0i8.p0i8.i64(ptr nonnull align 8 dereferenceable(16) %a, ptr nonnull align 8 dereferenceable(16) %b, i64 16, i1 false)
27+
tail call void @llvm.memcpy.p0i8.p0i8.i64(ptr nonnull align 8 dereferenceable(16) %a, ptr nonnull align 8 dereferenceable(16) %b, i64 15, i1 false)
1928
ret void
2029
}
2130

@@ -32,6 +41,26 @@ entry:
3241
ret void
3342
}
3443

44+
define void @call_memset_intrinsics(i8* nonnull align 8 dereferenceable(16) %a) sanitize_numerical_stability {
45+
; CHECK-LABEL: @call_memset_intrinsics(
46+
; CHECK-NEXT: entry:
47+
; CHECK-NEXT: call void @__nsan_set_value_unknown_4(ptr [[A:%.*]])
48+
; CHECK-NEXT: tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) [[A]], i8 0, i64 4, i1 false)
49+
; CHECK-NEXT: call void @__nsan_set_value_unknown_8(ptr [[A:%.*]])
50+
; CHECK-NEXT: tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) [[A]], i8 0, i64 8, i1 false)
51+
; CHECK-NEXT: call void @__nsan_set_value_unknown_16(ptr [[A:%.*]])
52+
; CHECK-NEXT: tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) [[A]], i8 0, i64 16, i1 false)
53+
; CHECK-NEXT: call void @__nsan_set_value_unknown(ptr [[A:%.*]], i64 15)
54+
; CHECK-NEXT: tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) [[A]], i8 0, i64 15, i1 false)
55+
; CHECK-NEXT: ret void
56+
;
57+
entry:
58+
tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) %a, i8 0, i64 4, i1 false)
59+
tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) %a, i8 0, i64 8, i1 false)
60+
tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) %a, i8 0, i64 16, i1 false)
61+
tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) %a, i8 0, i64 15, i1 false)
62+
ret void
63+
}
3564

3665
define void @transfer_float(float* %dst, float* %src) sanitize_numerical_stability {
3766
; CHECK-LABEL: @transfer_float(

0 commit comments

Comments
 (0)