[compiler-rt] [llvm] [nsan] Emit calls to optimized functions (PR #98900)

Dmitriy Chestnykh via llvm-commits llvm-commits at lists.llvm.org
Wed Jul 17 23:06:12 PDT 2024


https://github.com/chestnykh updated https://github.com/llvm/llvm-project/pull/98900

>From 48a51c1e7699bff6e87cfbb13193de9759c1c0e8 Mon Sep 17 00:00:00 2001
From: Dmitry Chestnykh <dm.chestnykh at gmail.com>
Date: Mon, 15 Jul 2024 16:32:57 +0300
Subject: [PATCH 01/19] [nsan] Emit calls to optimized functions

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 definition of the optimized functions.
---
 compiler-rt/lib/nsan/nsan.cpp                 | 26 +++++-
 .../lib/sanitizer_common/sanitizer_linux.cpp  | 83 ++++++++++++++++-
 .../TestCases/Posix/dump_registers.cpp        | 12 ++-
 .../NumericalStabilitySanitizer.cpp           | 87 ++++++++++++++----
 .../NumericalStabilitySanitizer/memory.ll     | 89 ++++++++++++++++++-
 5 files changed, 269 insertions(+), 28 deletions(-)

diff --git a/compiler-rt/lib/nsan/nsan.cpp b/compiler-rt/lib/nsan/nsan.cpp
index 194093c9679d0..cbf478f887429 100644
--- a/compiler-rt/lib/nsan/nsan.cpp
+++ b/compiler-rt/lib/nsan/nsan.cpp
@@ -54,8 +54,6 @@ using namespace __nsan;
 constexpr int kMaxVectorWidth = 8;
 
 // When copying application memory, we also copy its shadow and shadow type.
-// FIXME: We could provide fixed-size versions that would nicely
-// vectorize for known sizes.
 extern "C" SANITIZER_INTERFACE_ATTRIBUTE void
 __nsan_copy_values(const u8 *daddr, const u8 *saddr, uptr size) {
   internal_memmove((void *)GetShadowTypeAddrFor(daddr),
@@ -64,13 +62,33 @@ __nsan_copy_values(const u8 *daddr, const u8 *saddr, uptr size) {
                    size * kShadowScale);
 }
 
-// FIXME: We could provide fixed-size versions that would nicely
-// vectorize for known sizes.
+#define NSAN_COPY_VALUES_N(N)                                                  \
+  extern "C" SANITIZER_INTERFACE_ATTRIBUTE void __nsan_copy_##N(               \
+      const u8 *daddr, const u8 *saddr) {                                      \
+    __builtin_memmove((void *)GetShadowTypeAddrFor(daddr),                     \
+                      GetShadowTypeAddrFor(saddr), N);                         \
+    __builtin_memmove((void *)GetShadowAddrFor(daddr),                         \
+                      GetShadowAddrFor(saddr), N *kShadowScale);               \
+  }
+
+NSAN_COPY_VALUES_N(4);
+NSAN_COPY_VALUES_N(8);
+NSAN_COPY_VALUES_N(16);
+
 extern "C" SANITIZER_INTERFACE_ATTRIBUTE void
 __nsan_set_value_unknown(const u8 *addr, uptr size) {
   internal_memset((void *)GetShadowTypeAddrFor(addr), 0, size);
 }
 
+#define NSAN_SET_VALUE_UNKNOWN_N(N)                                            \
+  extern "C" SANITIZER_INTERFACE_ATTRIBUTE void __nsan_set_value_unknown_##N(  \
+      const u8 *daddr) {                                                       \
+    __builtin_memset((void *)GetShadowTypeAddrFor(daddr), 0, N);               \
+  }
+
+NSAN_SET_VALUE_UNKNOWN_N(4);
+NSAN_SET_VALUE_UNKNOWN_N(8);
+NSAN_SET_VALUE_UNKNOWN_N(16);
 
 const char *FTInfo<float>::kCppTypeName = "float";
 const char *FTInfo<double>::kCppTypeName = "double";
diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_linux.cpp b/compiler-rt/lib/sanitizer_common/sanitizer_linux.cpp
index 7935c88204a05..58233a3c0f607 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_linux.cpp
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_linux.cpp
@@ -2118,8 +2118,89 @@ bool SignalContext::IsTrueFaultingAddress() const {
   return si->si_signo == SIGSEGV && si->si_code != 128;
 }
 
+static const char *RegNumToRegName(int reg) {
+#  if defined(__x86_64__)
+  switch (reg) {
+    case REG_RAX:
+      return "rax";
+    case REG_RBX:
+      return "rbx";
+    case REG_RCX:
+      return "rcx";
+    case REG_RDX:
+      return "rdx";
+    case REG_RDI:
+      return "rdi";
+    case REG_RSI:
+      return "rsi";
+    case REG_RBP:
+      return "rbp";
+    case REG_RSP:
+      return "rsp";
+    case REG_R8:
+      return "r8";
+    case REG_R9:
+      return "r9";
+    case REG_R10:
+      return "r10";
+    case REG_R11:
+      return "r11";
+    case REG_R12:
+      return "r12";
+    case REG_R13:
+      return "r13";
+    case REG_R14:
+      return "r14";
+    case REG_R15:
+      return "r15";
+    default:
+      return NULL;
+  }
+#  endif
+  return NULL;
+}
+
 void SignalContext::DumpAllRegisters(void *context) {
-  // FIXME: Implement this.
+#  if SANITIZER_LINUX
+  ucontext_t *ucontext = (ucontext_t *)context;
+#    define DUMPREG64(r)                             \
+      Printf("%s = 0x%016llx  ", RegNumToRegName(r), \
+             ucontext->uc_mcontext.gregs[r]);
+#    define DUMPREG_(r) \
+      Printf(" ");      \
+      DUMPREG(r);
+#    define DUMPREG__(r) \
+      Printf("  ");      \
+      DUMPREG(r);
+#    define DUMPREG___(r) \
+      Printf("   ");      \
+      DUMPREG(r);
+#    if defined(__x86_64__)
+#      define DUMPREG(r) DUMPREG64(r)
+  Report("Register values:\n");
+  DUMPREG(REG_RAX);
+  DUMPREG(REG_RBX);
+  DUMPREG(REG_RCX);
+  DUMPREG(REG_RDX);
+  Printf("\n");
+  DUMPREG(REG_RDI);
+  DUMPREG(REG_RSI);
+  DUMPREG(REG_RBP);
+  DUMPREG(REG_RSP);
+  Printf("\n");
+  DUMPREG_(REG_R8);
+  DUMPREG_(REG_R9);
+  DUMPREG(REG_R10);
+  DUMPREG(REG_R11);
+  Printf("\n");
+  DUMPREG(REG_R12);
+  DUMPREG(REG_R13);
+  DUMPREG(REG_R14);
+  DUMPREG(REG_R15);
+  Printf("\n");
+#    endif
+#  endif
+  // FIXME: Implement this for other OSes and architectures.
 }
 
 static void GetPcSpBp(void *context, uptr *pc, uptr *sp, uptr *bp) {
diff --git a/compiler-rt/test/sanitizer_common/TestCases/Posix/dump_registers.cpp b/compiler-rt/test/sanitizer_common/TestCases/Posix/dump_registers.cpp
index f09b2bf4447cc..2cbd40e924263 100644
--- a/compiler-rt/test/sanitizer_common/TestCases/Posix/dump_registers.cpp
+++ b/compiler-rt/test/sanitizer_common/TestCases/Posix/dump_registers.cpp
@@ -1,16 +1,20 @@
 // Check that sanitizer prints registers dump_registers on dump_registers=1
 // RUN: %clangxx  %s -o %t
-// RUN: %env_tool_opts=dump_registers=0 %run %t 2>&1 | FileCheck %s --check-prefix=CHECK-NODUMP
-// RUN: %run %t 2>&1 | FileCheck %s --check-prefix=CHECK-DUMP
+// RUN: %env_tool_opts=dump_registers=0 not %run %t 2>&1 | FileCheck %s --check-prefix=CHECK-NODUMP
+// RUN: not %run %t 2>&1 | FileCheck %s --check-prefix=CHECK-DUMP
 //
-// FIXME: Implement.
-// XFAIL: *
+// FIXME: Implement for other OSes and architectures.
+// REQUIRES: x86_64-target-arch, linux
 
 #include <signal.h>
 
 int main() {
   raise(SIGSEGV);
   // CHECK-DUMP: Register values
+  // CHECK-DUMP-NEXT: rax = {{0x[0-9a-f]+}}  rbx = {{0x[0-9a-f]+}}  rcx = {{0x[0-9a-f]+}}  rdx = {{0x[0-9a-f]+}}
+  // CHECK-DUMP-NEXT: rdi = {{0x[0-9a-f]+}}  rsi = {{0x[0-9a-f]+}}  rbp = {{0x[0-9a-f]+}}  rsp = {{0x[0-9a-f]+}}
+  // CHECK-DUMP-NEXT:  r8 = {{0x[0-9a-f]+}}   r9 = {{0x[0-9a-f]+}}  r10 = {{0x[0-9a-f]+}}  r11 = {{0x[0-9a-f]+}}
+  // CHECK-DUMP-NEXT: r12 = {{0x[0-9a-f]+}}   r13 = {{0x[0-9a-f]+}} r14 = {{0x[0-9a-f]+}}  r15 = {{0x[0-9a-f]+}}
   // CHECK-NODUMP-NOT: Register values
   return 0;
 }
diff --git a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
index 99b1c779f3167..607b4384b5136 100644
--- a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
@@ -550,12 +550,15 @@ class NumericalStabilitySanitizer {
   LLVMContext &Context;
   MappingConfig Config;
   IntegerType *IntptrTy = nullptr;
+
+  // TODO: Use std::array instead?
   FunctionCallee NsanGetShadowPtrForStore[FTValueType::kNumValueTypes] = {};
   FunctionCallee NsanGetShadowPtrForLoad[FTValueType::kNumValueTypes] = {};
   FunctionCallee NsanCheckValue[FTValueType::kNumValueTypes] = {};
   FunctionCallee NsanFCmpFail[FTValueType::kNumValueTypes] = {};
-  FunctionCallee NsanCopyValues;
-  FunctionCallee NsanSetValueUnknown;
+
+  std::array<FunctionCallee, 4> NsanCopyFunction;
+  std::array<FunctionCallee, 4> NsanSetValueUnknownFunction;
   FunctionCallee NsanGetRawShadowTypePtr;
   FunctionCallee NsanGetRawShadowPtr;
   GlobalValue *NsanShadowRetTag = nullptr;
@@ -634,10 +637,27 @@ NumericalStabilitySanitizer::NumericalStabilitySanitizer(Module &M)
         Attr, VoidTy, VTTy, VTTy, ShadowTy, ShadowTy, Int32Ty, Int1Ty, Int1Ty);
   }
 
-  NsanCopyValues = M.getOrInsertFunction("__nsan_copy_values", Attr, VoidTy,
-                                         PtrTy, PtrTy, IntptrTy);
-  NsanSetValueUnknown = M.getOrInsertFunction("__nsan_set_value_unknown", Attr,
-                                              VoidTy, PtrTy, IntptrTy);
+  NsanCopyFunction[0] =
+      M.getOrInsertFunction("__nsan_copy_4", Attr, VoidTy, PtrTy, PtrTy);
+  NsanCopyFunction[1] =
+      M.getOrInsertFunction("__nsan_copy_8", Attr, VoidTy, PtrTy, PtrTy);
+  NsanCopyFunction[2] =
+      M.getOrInsertFunction("__nsan_copy_16", Attr, VoidTy, PtrTy, PtrTy);
+
+  NsanCopyFunction[3] = M.getOrInsertFunction("__nsan_copy_values", Attr,
+                                              VoidTy, PtrTy, PtrTy, IntptrTy);
+
+  NsanSetValueUnknownFunction[0] =
+      M.getOrInsertFunction("__nsan_set_value_unknown_4", Attr, VoidTy, PtrTy);
+
+  NsanSetValueUnknownFunction[1] =
+      M.getOrInsertFunction("__nsan_set_value_unknown_8", Attr, VoidTy, PtrTy);
+
+  NsanSetValueUnknownFunction[2] =
+      M.getOrInsertFunction("__nsan_set_value_unknown_16", Attr, VoidTy, PtrTy);
+
+  NsanSetValueUnknownFunction[3] = M.getOrInsertFunction(
+      "__nsan_set_value_unknown", Attr, VoidTy, PtrTy, IntptrTy);
 
   // TODO: Add attributes nofree, nosync, readnone, readonly,
   NsanGetRawShadowTypePtr = M.getOrInsertFunction(
@@ -1880,7 +1900,7 @@ void NumericalStabilitySanitizer::propagateNonFTStore(
     }
   }
   // All other stores just reset the shadow value to unknown.
-  Builder.CreateCall(NsanSetValueUnknown, {Dst, ValueSize});
+  Builder.CreateCall(NsanSetValueUnknownFunction.back(), {Dst, ValueSize});
 }
 
 void NumericalStabilitySanitizer::propagateShadowValues(
@@ -2123,21 +2143,56 @@ bool NumericalStabilitySanitizer::sanitizeFunction(
   return !ValueToShadow.empty();
 }
 
+static size_t GetInstrumentationCalleeIdxForMemOp(Value *V,
+                                                  std::size_t MaxIdx) {
+  uint64_t OpSize = 0;
+  if (Constant *C = dyn_cast<Constant>(V)) {
+    auto *CInt = dyn_cast<ConstantInt>(C);
+    if (CInt && CInt->getValue().getBitWidth() <= 64)
+      OpSize = CInt->getValue().getZExtValue();
+  }
+
+  size_t CandidateIdx =
+      OpSize == 4 ? 0 : (OpSize == 8 ? 1 : (OpSize == 16 ? 2 : MaxIdx));
+
+  return CandidateIdx <= MaxIdx ? CandidateIdx : MaxIdx;
+}
+
 // Instrument the memory intrinsics so that they properly modify the shadow
 // memory.
 bool NumericalStabilitySanitizer::instrumentMemIntrinsic(MemIntrinsic *MI) {
   IRBuilder<> Builder(MI);
   if (auto *M = dyn_cast<MemSetInst>(MI)) {
-    Builder.CreateCall(
-        NsanSetValueUnknown,
-        {/*Address=*/M->getArgOperand(0),
-         /*Size=*/Builder.CreateIntCast(M->getArgOperand(2), IntptrTy, false)});
+    Value *OpSizeValue = M->getArgOperand(2);
+    size_t NrSetValueUnknownFns = NsanSetValueUnknownFunction.size();
+    size_t NsanSetValueUnknownFunctionIdx = GetInstrumentationCalleeIdxForMemOp(
+        OpSizeValue, NrSetValueUnknownFns - 1);
+    if (NsanSetValueUnknownFunctionIdx != NrSetValueUnknownFns - 1)
+      Builder.CreateCall(
+          NsanSetValueUnknownFunction[NsanSetValueUnknownFunctionIdx],
+          {/*Address=*/M->getArgOperand(0)});
+    else
+      Builder.CreateCall(NsanSetValueUnknownFunction.back(),
+                         {/*Address=*/M->getArgOperand(0),
+                          /*Size=*/Builder.CreateIntCast(M->getArgOperand(2),
+                                                         IntptrTy, false)});
+
   } else if (auto *M = dyn_cast<MemTransferInst>(MI)) {
-    Builder.CreateCall(
-        NsanCopyValues,
-        {/*Destination=*/M->getArgOperand(0),
-         /*Source=*/M->getArgOperand(1),
-         /*Size=*/Builder.CreateIntCast(M->getArgOperand(2), IntptrTy, false)});
+    Value *OpSizeValue = M->getArgOperand(2);
+    size_t NrCopyFns = NsanCopyFunction.size();
+    size_t NsanCopyFunctionIdx =
+        GetInstrumentationCalleeIdxForMemOp(OpSizeValue, NrCopyFns - 1);
+
+    if (NsanCopyFunctionIdx != NrCopyFns - 1)
+      Builder.CreateCall(NsanCopyFunction[NsanCopyFunctionIdx],
+                         {/*Destination=*/M->getArgOperand(0),
+                          /*Source=*/M->getArgOperand(1)});
+    else
+      Builder.CreateCall(
+          NsanCopyFunction.back(),
+          {/*Destination=*/M->getArgOperand(0),
+           /*Source=*/M->getArgOperand(1),
+           /*Size=*/Builder.CreateIntCast(OpSizeValue, IntptrTy, false)});
   }
   return false;
 }
diff --git a/llvm/test/Instrumentation/NumericalStabilitySanitizer/memory.ll b/llvm/test/Instrumentation/NumericalStabilitySanitizer/memory.ll
index 3fe78d8b19b0a..1f238b88f49a4 100644
--- a/llvm/test/Instrumentation/NumericalStabilitySanitizer/memory.ll
+++ b/llvm/test/Instrumentation/NumericalStabilitySanitizer/memory.ll
@@ -7,10 +7,10 @@ target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f3
 
 declare void @llvm.memcpy.p0i8.p0i8.i64(i8*, i8*, i64, i1)
 
-define void @call_memcpy_intrinsic(i8* nonnull align 8 dereferenceable(16) %a, i8* nonnull align 8 dereferenceable(16) %b) sanitize_numerical_stability {
-; CHECK-LABEL: @call_memcpy_intrinsic(
+define void @call_memcpy_intrinsic_16bytes(i8* nonnull align 8 dereferenceable(16) %a, i8* nonnull align 8 dereferenceable(16) %b) sanitize_numerical_stability {
+; CHECK-LABEL: @call_memcpy_intrinsic_16bytes(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    call void @__nsan_copy_values(ptr [[A:%.*]], ptr [[B:%.*]], i64 16)
+; CHECK-NEXT:    call void @__nsan_copy_16(ptr [[A:%.*]], ptr [[B:%.*]])
 ; 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)
 ; CHECK-NEXT:    ret void
 ;
@@ -19,6 +19,42 @@ entry:
   ret void
 }
 
+define void @call_memcpy_intrinsic_8bytes(i8* nonnull align 8 dereferenceable(16) %a, i8* nonnull align 8 dereferenceable(16) %b) sanitize_numerical_stability {
+; CHECK-LABEL: @call_memcpy_intrinsic_8bytes(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @__nsan_copy_8(ptr [[A:%.*]], ptr [[B:%.*]])
+; 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)
+; CHECK-NEXT:    ret void
+;
+entry:
+  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)
+  ret void
+}
+
+define void @call_memcpy_intrinsic_4bytes(i8* nonnull align 8 dereferenceable(16) %a, i8* nonnull align 8 dereferenceable(16) %b) sanitize_numerical_stability {
+; CHECK-LABEL: @call_memcpy_intrinsic_4bytes(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @__nsan_copy_4(ptr [[A:%.*]], ptr [[B:%.*]])
+; 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)
+; CHECK-NEXT:    ret void
+;
+entry:
+  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)
+  ret void
+}
+
+define void @call_memcpy_intrinsic(i8* nonnull align 8 dereferenceable(16) %a, i8* nonnull align 8 dereferenceable(16) %b) sanitize_numerical_stability {
+; CHECK-LABEL: @call_memcpy_intrinsic(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @__nsan_copy_values(ptr [[A:%.*]], ptr [[B:%.*]], i64 15)
+; 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)
+; CHECK-NEXT:    ret void
+;
+entry:
+  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)
+  ret void
+}
+
 declare dso_local i8* @memcpy(i8*, i8*, i64) local_unnamed_addr
 
 define void @call_memcpy(i8* nonnull align 8 dereferenceable(16) %a, i8* nonnull align 8 dereferenceable(16) %b) sanitize_numerical_stability {
@@ -32,6 +68,53 @@ entry:
   ret void
 }
 
+define void @call_memset_intrinsic_16bytes(i8* nonnull align 8 dereferenceable(16) %a) sanitize_numerical_stability {
+; CHECK-LABEL: @call_memset_intrinsic_16bytes(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @__nsan_set_value_unknown_16(ptr [[A:%.*]])
+; CHECK-NEXT:    tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) [[A]], i8 0, i64 16, i1 false)
+; CHECK-NEXT:    ret void
+;
+entry:
+  tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) %a, i8 0, i64 16, i1 false)
+  ret void
+}
+
+define void @call_memset_intrinsic_8bytes(i8* nonnull align 8 dereferenceable(16) %a) sanitize_numerical_stability {
+; CHECK-LABEL: @call_memset_intrinsic_8bytes(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @__nsan_set_value_unknown_8(ptr [[A:%.*]])
+; CHECK-NEXT:    tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) [[A]], i8 0, i64 8, i1 false)
+; CHECK-NEXT:    ret void
+;
+entry:
+  tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) %a, i8 0, i64 8, i1 false)
+  ret void
+}
+
+define void @call_memset_intrinsic_4bytes(i8* nonnull align 8 dereferenceable(16) %a) sanitize_numerical_stability {
+; CHECK-LABEL: @call_memset_intrinsic_4bytes(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @__nsan_set_value_unknown_4(ptr [[A:%.*]])
+; CHECK-NEXT:    tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) [[A]], i8 0, i64 4, i1 false)
+; CHECK-NEXT:    ret void
+;
+entry:
+  tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) %a, i8 0, i64 4, i1 false)
+  ret void
+}
+
+define void @call_memset_intrinsic(i8* nonnull align 8 dereferenceable(16) %a) sanitize_numerical_stability {
+; CHECK-LABEL: @call_memset_intrinsic(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @__nsan_set_value_unknown(ptr [[A:%.*]], i64 15)
+; CHECK-NEXT:    tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) [[A]], i8 0, i64 15, i1 false)
+; CHECK-NEXT:    ret void
+;
+entry:
+  tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) %a, i8 0, i64 15, i1 false)
+  ret void
+}
 
 define void @transfer_float(float* %dst, float* %src) sanitize_numerical_stability {
 ; CHECK-LABEL: @transfer_float(

>From 57e95870b20246358d7e777bff05c51a354c3dde Mon Sep 17 00:00:00 2001
From: Dmitry Chestnykh <dm.chestnykh at gmail.com>
Date: Mon, 15 Jul 2024 17:02:16 +0300
Subject: [PATCH 02/19] [compiler-rt][nsan] Remove unneeded semicolons

---
 compiler-rt/lib/nsan/nsan.cpp | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/compiler-rt/lib/nsan/nsan.cpp b/compiler-rt/lib/nsan/nsan.cpp
index cbf478f887429..13ef3f9da4fcb 100644
--- a/compiler-rt/lib/nsan/nsan.cpp
+++ b/compiler-rt/lib/nsan/nsan.cpp
@@ -71,9 +71,9 @@ __nsan_copy_values(const u8 *daddr, const u8 *saddr, uptr size) {
                       GetShadowAddrFor(saddr), N *kShadowScale);               \
   }
 
-NSAN_COPY_VALUES_N(4);
-NSAN_COPY_VALUES_N(8);
-NSAN_COPY_VALUES_N(16);
+NSAN_COPY_VALUES_N(4)
+NSAN_COPY_VALUES_N(8)
+NSAN_COPY_VALUES_N(16)
 
 extern "C" SANITIZER_INTERFACE_ATTRIBUTE void
 __nsan_set_value_unknown(const u8 *addr, uptr size) {
@@ -86,9 +86,9 @@ __nsan_set_value_unknown(const u8 *addr, uptr size) {
     __builtin_memset((void *)GetShadowTypeAddrFor(daddr), 0, N);               \
   }
 
-NSAN_SET_VALUE_UNKNOWN_N(4);
-NSAN_SET_VALUE_UNKNOWN_N(8);
-NSAN_SET_VALUE_UNKNOWN_N(16);
+NSAN_SET_VALUE_UNKNOWN_N(4)
+NSAN_SET_VALUE_UNKNOWN_N(8)
+NSAN_SET_VALUE_UNKNOWN_N(16)
 
 const char *FTInfo<float>::kCppTypeName = "float";
 const char *FTInfo<double>::kCppTypeName = "double";

>From 7380cc50a8286a38cf2703424e6af228e3284328 Mon Sep 17 00:00:00 2001
From: Dmitry Chestnykh <dm.chestnykh at gmail.com>
Date: Mon, 15 Jul 2024 18:15:45 +0300
Subject: [PATCH 03/19] [nsan] Optimize `GetInstrumentationCalleeIdxForMemOp`

---
 .../Instrumentation/NumericalStabilitySanitizer.cpp          | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
index 607b4384b5136..41876f9c9c97f 100644
--- a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
@@ -2152,10 +2152,7 @@ static size_t GetInstrumentationCalleeIdxForMemOp(Value *V,
       OpSize = CInt->getValue().getZExtValue();
   }
 
-  size_t CandidateIdx =
-      OpSize == 4 ? 0 : (OpSize == 8 ? 1 : (OpSize == 16 ? 2 : MaxIdx));
-
-  return CandidateIdx <= MaxIdx ? CandidateIdx : MaxIdx;
+  return OpSize == 4 ? 0 : (OpSize == 8 ? 1 : (OpSize == 16 ? 2 : MaxIdx));
 }
 
 // Instrument the memory intrinsics so that they properly modify the shadow

>From 336b4b9e783b7f200646ac401c5a6fc64b223a83 Mon Sep 17 00:00:00 2001
From: Dmitry Chestnykh <dm.chestnykh at gmail.com>
Date: Tue, 16 Jul 2024 06:46:13 +0300
Subject: [PATCH 04/19] [nsan] Simplify memory.ll test

- Test all calls to `__nsan_copy_*` in one function
- Test all calls to `__nsan_set_value_unknown_*` in one function
---
 .../NumericalStabilitySanitizer/memory.ll     | 90 ++++---------------
 1 file changed, 18 insertions(+), 72 deletions(-)

diff --git a/llvm/test/Instrumentation/NumericalStabilitySanitizer/memory.ll b/llvm/test/Instrumentation/NumericalStabilitySanitizer/memory.ll
index 1f238b88f49a4..31f32d7bd8df5 100644
--- a/llvm/test/Instrumentation/NumericalStabilitySanitizer/memory.ll
+++ b/llvm/test/Instrumentation/NumericalStabilitySanitizer/memory.ll
@@ -7,50 +7,23 @@ target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f3
 
 declare void @llvm.memcpy.p0i8.p0i8.i64(i8*, i8*, i64, i1)
 
-define void @call_memcpy_intrinsic_16bytes(i8* nonnull align 8 dereferenceable(16) %a, i8* nonnull align 8 dereferenceable(16) %b) sanitize_numerical_stability {
-; CHECK-LABEL: @call_memcpy_intrinsic_16bytes(
-; CHECK-NEXT:  entry:
-; CHECK-NEXT:    call void @__nsan_copy_16(ptr [[A:%.*]], ptr [[B:%.*]])
-; 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)
-; CHECK-NEXT:    ret void
-;
-entry:
-  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)
-  ret void
-}
-
-define void @call_memcpy_intrinsic_8bytes(i8* nonnull align 8 dereferenceable(16) %a, i8* nonnull align 8 dereferenceable(16) %b) sanitize_numerical_stability {
-; CHECK-LABEL: @call_memcpy_intrinsic_8bytes(
-; CHECK-NEXT:  entry:
-; CHECK-NEXT:    call void @__nsan_copy_8(ptr [[A:%.*]], ptr [[B:%.*]])
-; 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)
-; CHECK-NEXT:    ret void
-;
-entry:
-  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)
-  ret void
-}
-
-define void @call_memcpy_intrinsic_4bytes(i8* nonnull align 8 dereferenceable(16) %a, i8* nonnull align 8 dereferenceable(16) %b) sanitize_numerical_stability {
-; CHECK-LABEL: @call_memcpy_intrinsic_4bytes(
+define void @call_memcpy_intrinsics(i8* nonnull align 8 dereferenceable(16) %a, i8* nonnull align 8 dereferenceable(16) %b) sanitize_numerical_stability {
+; CHECK-LABEL: @call_memcpy_intrinsics(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    call void @__nsan_copy_4(ptr [[A:%.*]], ptr [[B:%.*]])
 ; 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)
-; CHECK-NEXT:    ret void
-;
-entry:
-  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)
-  ret void
-}
-
-define void @call_memcpy_intrinsic(i8* nonnull align 8 dereferenceable(16) %a, i8* nonnull align 8 dereferenceable(16) %b) sanitize_numerical_stability {
-; CHECK-LABEL: @call_memcpy_intrinsic(
-; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @__nsan_copy_8(ptr [[A:%.*]], ptr [[B:%.*]])
+; 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)
+; CHECK-NEXT:    call void @__nsan_copy_16(ptr [[A:%.*]], ptr [[B:%.*]])
+; 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)
 ; CHECK-NEXT:    call void @__nsan_copy_values(ptr [[A:%.*]], ptr [[B:%.*]], i64 15)
 ; 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)
 ; CHECK-NEXT:    ret void
 ;
 entry:
+  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)
+  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)
+  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)
   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)
   ret void
 }
@@ -68,50 +41,23 @@ entry:
   ret void
 }
 
-define void @call_memset_intrinsic_16bytes(i8* nonnull align 8 dereferenceable(16) %a) sanitize_numerical_stability {
-; CHECK-LABEL: @call_memset_intrinsic_16bytes(
-; CHECK-NEXT:  entry:
-; CHECK-NEXT:    call void @__nsan_set_value_unknown_16(ptr [[A:%.*]])
-; CHECK-NEXT:    tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) [[A]], i8 0, i64 16, i1 false)
-; CHECK-NEXT:    ret void
-;
-entry:
-  tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) %a, i8 0, i64 16, i1 false)
-  ret void
-}
-
-define void @call_memset_intrinsic_8bytes(i8* nonnull align 8 dereferenceable(16) %a) sanitize_numerical_stability {
-; CHECK-LABEL: @call_memset_intrinsic_8bytes(
-; CHECK-NEXT:  entry:
-; CHECK-NEXT:    call void @__nsan_set_value_unknown_8(ptr [[A:%.*]])
-; CHECK-NEXT:    tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) [[A]], i8 0, i64 8, i1 false)
-; CHECK-NEXT:    ret void
-;
-entry:
-  tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) %a, i8 0, i64 8, i1 false)
-  ret void
-}
-
-define void @call_memset_intrinsic_4bytes(i8* nonnull align 8 dereferenceable(16) %a) sanitize_numerical_stability {
-; CHECK-LABEL: @call_memset_intrinsic_4bytes(
+define void @call_memset_intrinsics(i8* nonnull align 8 dereferenceable(16) %a) sanitize_numerical_stability {
+; CHECK-LABEL: @call_memset_intrinsics(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    call void @__nsan_set_value_unknown_4(ptr [[A:%.*]])
 ; CHECK-NEXT:    tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) [[A]], i8 0, i64 4, i1 false)
-; CHECK-NEXT:    ret void
-;
-entry:
-  tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) %a, i8 0, i64 4, i1 false)
-  ret void
-}
-
-define void @call_memset_intrinsic(i8* nonnull align 8 dereferenceable(16) %a) sanitize_numerical_stability {
-; CHECK-LABEL: @call_memset_intrinsic(
-; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @__nsan_set_value_unknown_8(ptr [[A:%.*]])
+; CHECK-NEXT:    tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) [[A]], i8 0, i64 8, i1 false)
+; CHECK-NEXT:    call void @__nsan_set_value_unknown_16(ptr [[A:%.*]])
+; CHECK-NEXT:    tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) [[A]], i8 0, i64 16, i1 false)
 ; CHECK-NEXT:    call void @__nsan_set_value_unknown(ptr [[A:%.*]], i64 15)
 ; CHECK-NEXT:    tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) [[A]], i8 0, i64 15, i1 false)
 ; CHECK-NEXT:    ret void
 ;
 entry:
+  tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) %a, i8 0, i64 4, i1 false)
+  tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) %a, i8 0, i64 8, i1 false)
+  tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) %a, i8 0, i64 16, i1 false)
   tail call void @llvm.memset.p0.i64(ptr nonnull align 8 dereferenceable(16) %a, i8 0, i64 15, i1 false)
   ret void
 }

>From 7eaef056e3fb6786b9895e6a3d7e6bd6a508c582 Mon Sep 17 00:00:00 2001
From: Dmitry Chestnykh <dm.chestnykh at gmail.com>
Date: Tue, 16 Jul 2024 21:08:59 +0300
Subject: [PATCH 05/19] [nsan] Use size_t instead of std::size_t

---
 .../Transforms/Instrumentation/NumericalStabilitySanitizer.cpp | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
index 41876f9c9c97f..28ddb7efa7d2b 100644
--- a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
@@ -2143,8 +2143,7 @@ bool NumericalStabilitySanitizer::sanitizeFunction(
   return !ValueToShadow.empty();
 }
 
-static size_t GetInstrumentationCalleeIdxForMemOp(Value *V,
-                                                  std::size_t MaxIdx) {
+static size_t GetInstrumentationCalleeIdxForMemOp(Value *V, size_t MaxIdx) {
   uint64_t OpSize = 0;
   if (Constant *C = dyn_cast<Constant>(V)) {
     auto *CInt = dyn_cast<ConstantInt>(C);

>From f10ff37b3554804240ec5aacdda50e8893a7eeee Mon Sep 17 00:00:00 2001
From: Dmitry Chestnykh <dm.chestnykh at gmail.com>
Date: Tue, 16 Jul 2024 22:21:27 +0300
Subject: [PATCH 06/19] [nsan] Add helper classes

Add base class for handling details of
some `__nsan_*` functions
Add derived `NsanCopy` and `NsanSetUnknownClasses`
with the appropriate methods
---
 .../NumericalStabilitySanitizer.cpp           | 171 +++++++++++++-----
 1 file changed, 123 insertions(+), 48 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
index 28ddb7efa7d2b..1d26b06b2c49f 100644
--- a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
@@ -45,6 +45,7 @@
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 
 #include <cstdint>
+#include <map>
 
 using namespace llvm;
 
@@ -493,6 +494,103 @@ class ValueToShadowMap {
   DenseMap<Value *, Value *> Map;
 };
 
+// Base class for handling some details of __nsan_* functions
+class NsanInstrumentationFunction {
+public:
+  virtual std::pair<FunctionCallee, int>
+  getCalleeAndNumArgsForMemoryOp(const uint64_t MemOpSize) const = 0;
+  virtual FunctionCallee getGeneralFunction() const = 0;
+  virtual ~NsanInstrumentationFunction() = default;
+};
+
+// Helper class for __nsan_copy_* family
+class NsanCopy : public NsanInstrumentationFunction {
+public:
+  NsanCopy(Module &M);
+  std::pair<FunctionCallee, int>
+  getCalleeAndNumArgsForMemoryOp(const uint64_t MemOpSize) const override;
+  FunctionCallee getGeneralFunction() const override;
+
+private:
+  std::array<FunctionCallee, 4> NsanCopyFunction;
+};
+
+NsanCopy::NsanCopy(Module &M) {
+  LLVMContext &Ctx = M.getContext();
+  AttributeList Attr;
+  Attr = Attr.addFnAttribute(Ctx, Attribute::NoUnwind);
+  Type *PtrTy = PointerType::getUnqual(Ctx);
+  Type *VoidTy = Type::getVoidTy(Ctx);
+  IntegerType *IntptrTy = M.getDataLayout().getIntPtrType(Ctx);
+
+  NsanCopyFunction[0] =
+      M.getOrInsertFunction("__nsan_copy_4", Attr, VoidTy, PtrTy, PtrTy);
+  NsanCopyFunction[1] =
+      M.getOrInsertFunction("__nsan_copy_8", Attr, VoidTy, PtrTy, PtrTy);
+  NsanCopyFunction[2] =
+      M.getOrInsertFunction("__nsan_copy_16", Attr, VoidTy, PtrTy, PtrTy);
+
+  NsanCopyFunction[3] = M.getOrInsertFunction("__nsan_copy_values", Attr,
+                                              VoidTy, PtrTy, PtrTy, IntptrTy);
+}
+
+std::pair<FunctionCallee, int>
+NsanCopy::getCalleeAndNumArgsForMemoryOp(const uint64_t MemOpSize) const {
+  size_t Idx =
+      MemOpSize == 4 ? 0 : (MemOpSize == 8 ? 1 : (MemOpSize == 16 ? 2 : 3));
+  return Idx == 3 ? std::make_pair(NsanCopyFunction[Idx], 3)
+                  : std::make_pair(NsanCopyFunction[Idx], 2);
+}
+
+FunctionCallee NsanCopy::getGeneralFunction() const {
+  return NsanCopyFunction.back();
+}
+
+// Helper class for __nsan_set_value_unknown* family
+class NsanSetUnknown : public NsanInstrumentationFunction {
+public:
+  NsanSetUnknown(Module &M);
+  std::pair<FunctionCallee, int>
+  getCalleeAndNumArgsForMemoryOp(const uint64_t MemOpSize) const override;
+  FunctionCallee getGeneralFunction() const override;
+
+private:
+  std::array<FunctionCallee, 4> NsanSetUnknownFunction;
+};
+
+NsanSetUnknown::NsanSetUnknown(Module &M) {
+  LLVMContext &Ctx = M.getContext();
+  AttributeList Attr;
+  Attr = Attr.addFnAttribute(Ctx, Attribute::NoUnwind);
+  Type *PtrTy = PointerType::getUnqual(Ctx);
+  Type *VoidTy = Type::getVoidTy(Ctx);
+  IntegerType *IntptrTy = M.getDataLayout().getIntPtrType(Ctx);
+
+  NsanSetUnknownFunction[0] =
+      M.getOrInsertFunction("__nsan_set_value_unknown_4", Attr, VoidTy, PtrTy);
+
+  NsanSetUnknownFunction[1] =
+      M.getOrInsertFunction("__nsan_set_value_unknown_8", Attr, VoidTy, PtrTy);
+
+  NsanSetUnknownFunction[2] =
+      M.getOrInsertFunction("__nsan_set_value_unknown_16", Attr, VoidTy, PtrTy);
+
+  NsanSetUnknownFunction[3] = M.getOrInsertFunction(
+      "__nsan_set_value_unknown", Attr, VoidTy, PtrTy, IntptrTy);
+}
+
+std::pair<FunctionCallee, int>
+NsanSetUnknown::getCalleeAndNumArgsForMemoryOp(const uint64_t MemOpSize) const {
+  size_t Idx =
+      MemOpSize == 4 ? 0 : (MemOpSize == 8 ? 1 : (MemOpSize == 16 ? 2 : 3));
+  return Idx == 3 ? std::make_pair(NsanSetUnknownFunction[Idx], 2)
+                  : std::make_pair(NsanSetUnknownFunction[Idx], 1);
+}
+
+FunctionCallee NsanSetUnknown::getGeneralFunction() const {
+  return NsanSetUnknownFunction.back();
+}
+
 /// Instantiating NumericalStabilitySanitizer inserts the nsan runtime library
 /// API function declarations into the module if they don't exist already.
 /// Instantiating ensures the __nsan_init function is in the list of global
@@ -557,8 +655,9 @@ class NumericalStabilitySanitizer {
   FunctionCallee NsanCheckValue[FTValueType::kNumValueTypes] = {};
   FunctionCallee NsanFCmpFail[FTValueType::kNumValueTypes] = {};
 
-  std::array<FunctionCallee, 4> NsanCopyFunction;
-  std::array<FunctionCallee, 4> NsanSetValueUnknownFunction;
+  NsanCopy NsanCopyFns;
+  NsanSetUnknown NsanSetUnknownFns;
+
   FunctionCallee NsanGetRawShadowTypePtr;
   FunctionCallee NsanGetRawShadowPtr;
   GlobalValue *NsanShadowRetTag = nullptr;
@@ -573,6 +672,7 @@ class NumericalStabilitySanitizer {
 
   std::optional<Regex> CheckFunctionsFilter;
 };
+
 } // end anonymous namespace
 
 PreservedAnalyses
@@ -601,7 +701,8 @@ static GlobalValue *createThreadLocalGV(const char *Name, Module &M, Type *Ty) {
 }
 
 NumericalStabilitySanitizer::NumericalStabilitySanitizer(Module &M)
-    : DL(M.getDataLayout()), Context(M.getContext()), Config(Context) {
+    : DL(M.getDataLayout()), Context(M.getContext()), Config(Context),
+      NsanCopyFns(M), NsanSetUnknownFns(M) {
   IntptrTy = DL.getIntPtrType(Context);
   Type *PtrTy = PointerType::getUnqual(Context);
   Type *Int32Ty = Type::getInt32Ty(Context);
@@ -637,28 +738,6 @@ NumericalStabilitySanitizer::NumericalStabilitySanitizer(Module &M)
         Attr, VoidTy, VTTy, VTTy, ShadowTy, ShadowTy, Int32Ty, Int1Ty, Int1Ty);
   }
 
-  NsanCopyFunction[0] =
-      M.getOrInsertFunction("__nsan_copy_4", Attr, VoidTy, PtrTy, PtrTy);
-  NsanCopyFunction[1] =
-      M.getOrInsertFunction("__nsan_copy_8", Attr, VoidTy, PtrTy, PtrTy);
-  NsanCopyFunction[2] =
-      M.getOrInsertFunction("__nsan_copy_16", Attr, VoidTy, PtrTy, PtrTy);
-
-  NsanCopyFunction[3] = M.getOrInsertFunction("__nsan_copy_values", Attr,
-                                              VoidTy, PtrTy, PtrTy, IntptrTy);
-
-  NsanSetValueUnknownFunction[0] =
-      M.getOrInsertFunction("__nsan_set_value_unknown_4", Attr, VoidTy, PtrTy);
-
-  NsanSetValueUnknownFunction[1] =
-      M.getOrInsertFunction("__nsan_set_value_unknown_8", Attr, VoidTy, PtrTy);
-
-  NsanSetValueUnknownFunction[2] =
-      M.getOrInsertFunction("__nsan_set_value_unknown_16", Attr, VoidTy, PtrTy);
-
-  NsanSetValueUnknownFunction[3] = M.getOrInsertFunction(
-      "__nsan_set_value_unknown", Attr, VoidTy, PtrTy, IntptrTy);
-
   // TODO: Add attributes nofree, nosync, readnone, readonly,
   NsanGetRawShadowTypePtr = M.getOrInsertFunction(
       "__nsan_internal_get_raw_shadow_type_ptr", Attr, PtrTy, PtrTy);
@@ -1900,7 +1979,7 @@ void NumericalStabilitySanitizer::propagateNonFTStore(
     }
   }
   // All other stores just reset the shadow value to unknown.
-  Builder.CreateCall(NsanSetValueUnknownFunction.back(), {Dst, ValueSize});
+  Builder.CreateCall(NsanSetUnknownFns.getGeneralFunction(), {Dst, ValueSize});
 }
 
 void NumericalStabilitySanitizer::propagateShadowValues(
@@ -2143,7 +2222,7 @@ bool NumericalStabilitySanitizer::sanitizeFunction(
   return !ValueToShadow.empty();
 }
 
-static size_t GetInstrumentationCalleeIdxForMemOp(Value *V, size_t MaxIdx) {
+static uint64_t GetMemOpSize(Value *V) {
   uint64_t OpSize = 0;
   if (Constant *C = dyn_cast<Constant>(V)) {
     auto *CInt = dyn_cast<ConstantInt>(C);
@@ -2151,7 +2230,7 @@ static size_t GetInstrumentationCalleeIdxForMemOp(Value *V, size_t MaxIdx) {
       OpSize = CInt->getValue().getZExtValue();
   }
 
-  return OpSize == 4 ? 0 : (OpSize == 8 ? 1 : (OpSize == 16 ? 2 : MaxIdx));
+  return OpSize;
 }
 
 // Instrument the memory intrinsics so that they properly modify the shadow
@@ -2159,36 +2238,32 @@ static size_t GetInstrumentationCalleeIdxForMemOp(Value *V, size_t MaxIdx) {
 bool NumericalStabilitySanitizer::instrumentMemIntrinsic(MemIntrinsic *MI) {
   IRBuilder<> Builder(MI);
   if (auto *M = dyn_cast<MemSetInst>(MI)) {
-    Value *OpSizeValue = M->getArgOperand(2);
-    size_t NrSetValueUnknownFns = NsanSetValueUnknownFunction.size();
-    size_t NsanSetValueUnknownFunctionIdx = GetInstrumentationCalleeIdxForMemOp(
-        OpSizeValue, NrSetValueUnknownFns - 1);
-    if (NsanSetValueUnknownFunctionIdx != NrSetValueUnknownFns - 1)
-      Builder.CreateCall(
-          NsanSetValueUnknownFunction[NsanSetValueUnknownFunctionIdx],
-          {/*Address=*/M->getArgOperand(0)});
+    std::pair<FunctionCallee, int> SetUnknownFn =
+        NsanSetUnknownFns.getCalleeAndNumArgsForMemoryOp(
+            GetMemOpSize(M->getArgOperand(2)));
+    if (SetUnknownFn.second == 1)
+      Builder.CreateCall(SetUnknownFn.first, {/*Address=*/M->getArgOperand(0)});
     else
-      Builder.CreateCall(NsanSetValueUnknownFunction.back(),
+      Builder.CreateCall(SetUnknownFn.first,
                          {/*Address=*/M->getArgOperand(0),
                           /*Size=*/Builder.CreateIntCast(M->getArgOperand(2),
                                                          IntptrTy, false)});
 
   } else if (auto *M = dyn_cast<MemTransferInst>(MI)) {
-    Value *OpSizeValue = M->getArgOperand(2);
-    size_t NrCopyFns = NsanCopyFunction.size();
-    size_t NsanCopyFunctionIdx =
-        GetInstrumentationCalleeIdxForMemOp(OpSizeValue, NrCopyFns - 1);
-
-    if (NsanCopyFunctionIdx != NrCopyFns - 1)
-      Builder.CreateCall(NsanCopyFunction[NsanCopyFunctionIdx],
-                         {/*Destination=*/M->getArgOperand(0),
-                          /*Source=*/M->getArgOperand(1)});
+    std::pair<FunctionCallee, int> CopyFn =
+        NsanCopyFns.getCalleeAndNumArgsForMemoryOp(
+            GetMemOpSize(M->getArgOperand(2)));
+
+    if (CopyFn.second == 2)
+      Builder.CreateCall(CopyFn.first, {/*Destination=*/M->getArgOperand(0),
+                                        /*Source=*/M->getArgOperand(1)});
     else
       Builder.CreateCall(
-          NsanCopyFunction.back(),
+          CopyFn.first,
           {/*Destination=*/M->getArgOperand(0),
            /*Source=*/M->getArgOperand(1),
-           /*Size=*/Builder.CreateIntCast(OpSizeValue, IntptrTy, false)});
+           /*Size=*/
+           Builder.CreateIntCast(M->getArgOperand(2), IntptrTy, false)});
   }
   return false;
 }

>From 01258e1175f995cf84c451e5da6dddb069a9dd98 Mon Sep 17 00:00:00 2001
From: Dmitry Chestnykh <dm.chestnykh at gmail.com>
Date: Tue, 16 Jul 2024 22:26:02 +0300
Subject: [PATCH 07/19] [nsan] Remove unneeded newline

---
 .../Transforms/Instrumentation/NumericalStabilitySanitizer.cpp   | 1 -
 1 file changed, 1 deletion(-)

diff --git a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
index 1d26b06b2c49f..9c1fec5d22ab7 100644
--- a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
@@ -672,7 +672,6 @@ class NumericalStabilitySanitizer {
 
   std::optional<Regex> CheckFunctionsFilter;
 };
-
 } // end anonymous namespace
 
 PreservedAnalyses

>From 3ac259dcd509169dc5ebd7207789f7fed0c4ba08 Mon Sep 17 00:00:00 2001
From: Dmitry Chestnykh <dm.chestnykh at gmail.com>
Date: Tue, 16 Jul 2024 22:27:25 +0300
Subject: [PATCH 08/19] [nsan] Remove const qualifier from uint64_t args

---
 .../Instrumentation/NumericalStabilitySanitizer.cpp       | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
index 9c1fec5d22ab7..f715cd6ec5f7b 100644
--- a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
@@ -498,7 +498,7 @@ class ValueToShadowMap {
 class NsanInstrumentationFunction {
 public:
   virtual std::pair<FunctionCallee, int>
-  getCalleeAndNumArgsForMemoryOp(const uint64_t MemOpSize) const = 0;
+  getCalleeAndNumArgsForMemoryOp(uint64_t MemOpSize) const = 0;
   virtual FunctionCallee getGeneralFunction() const = 0;
   virtual ~NsanInstrumentationFunction() = default;
 };
@@ -508,7 +508,7 @@ class NsanCopy : public NsanInstrumentationFunction {
 public:
   NsanCopy(Module &M);
   std::pair<FunctionCallee, int>
-  getCalleeAndNumArgsForMemoryOp(const uint64_t MemOpSize) const override;
+  getCalleeAndNumArgsForMemoryOp(uint64_t MemOpSize) const override;
   FunctionCallee getGeneralFunction() const override;
 
 private:
@@ -535,7 +535,7 @@ NsanCopy::NsanCopy(Module &M) {
 }
 
 std::pair<FunctionCallee, int>
-NsanCopy::getCalleeAndNumArgsForMemoryOp(const uint64_t MemOpSize) const {
+NsanCopy::getCalleeAndNumArgsForMemoryOp(uint64_t MemOpSize) const {
   size_t Idx =
       MemOpSize == 4 ? 0 : (MemOpSize == 8 ? 1 : (MemOpSize == 16 ? 2 : 3));
   return Idx == 3 ? std::make_pair(NsanCopyFunction[Idx], 3)
@@ -551,7 +551,7 @@ class NsanSetUnknown : public NsanInstrumentationFunction {
 public:
   NsanSetUnknown(Module &M);
   std::pair<FunctionCallee, int>
-  getCalleeAndNumArgsForMemoryOp(const uint64_t MemOpSize) const override;
+  getCalleeAndNumArgsForMemoryOp(uint64_t MemOpSize) const override;
   FunctionCallee getGeneralFunction() const override;
 
 private:

>From c58aaac9f4c935434e095a668b10c3ee6bc1b033 Mon Sep 17 00:00:00 2001
From: Dmitry Chestnykh <dm.chestnykh at gmail.com>
Date: Wed, 17 Jul 2024 07:32:11 +0300
Subject: [PATCH 09/19] [nsan] Simplify helper classes

Use one class instead of base and derived classes
---
 .../NumericalStabilitySanitizer.cpp           | 153 +++++++-----------
 1 file changed, 56 insertions(+), 97 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
index f715cd6ec5f7b..f9ce473f4d94c 100644
--- a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
@@ -494,28 +494,22 @@ class ValueToShadowMap {
   DenseMap<Value *, Value *> Map;
 };
 
-// Base class for handling some details of __nsan_* functions
-class NsanInstrumentationFunction {
+// First parameter is the number of functions
+// Second parameter is the number of fallback function arguments
+template <size_t N, int NFallbackArgs> class NsanMemOpFn {
 public:
-  virtual std::pair<FunctionCallee, int>
-  getCalleeAndNumArgsForMemoryOp(uint64_t MemOpSize) const = 0;
-  virtual FunctionCallee getGeneralFunction() const = 0;
-  virtual ~NsanInstrumentationFunction() = default;
-};
-
-// Helper class for __nsan_copy_* family
-class NsanCopy : public NsanInstrumentationFunction {
-public:
-  NsanCopy(Module &M);
-  std::pair<FunctionCallee, int>
-  getCalleeAndNumArgsForMemoryOp(uint64_t MemOpSize) const override;
-  FunctionCallee getGeneralFunction() const override;
+  NsanMemOpFn(Module &M, ArrayRef<StringRef> Sized, StringRef Fallback);
+  // Number of parameters can be extracted from FunctionCallee
+  FunctionCallee getFunctionFor(uint64_t MemOpSize) const;
+  FunctionCallee getFallback() const;
 
 private:
-  std::array<FunctionCallee, 4> NsanCopyFunction;
+  std::array<FunctionCallee, N> Funcs;
 };
 
-NsanCopy::NsanCopy(Module &M) {
+template <size_t N, int NFallbackArgs>
+NsanMemOpFn<N, NFallbackArgs>::NsanMemOpFn(Module &M, ArrayRef<StringRef> Sized,
+                                           StringRef Fallback) {
   LLVMContext &Ctx = M.getContext();
   AttributeList Attr;
   Attr = Attr.addFnAttribute(Ctx, Attribute::NoUnwind);
@@ -523,72 +517,34 @@ NsanCopy::NsanCopy(Module &M) {
   Type *VoidTy = Type::getVoidTy(Ctx);
   IntegerType *IntptrTy = M.getDataLayout().getIntPtrType(Ctx);
 
-  NsanCopyFunction[0] =
-      M.getOrInsertFunction("__nsan_copy_4", Attr, VoidTy, PtrTy, PtrTy);
-  NsanCopyFunction[1] =
-      M.getOrInsertFunction("__nsan_copy_8", Attr, VoidTy, PtrTy, PtrTy);
-  NsanCopyFunction[2] =
-      M.getOrInsertFunction("__nsan_copy_16", Attr, VoidTy, PtrTy, PtrTy);
+  for (size_t i = 0; i < N - 1; ++i) {
+    if (NFallbackArgs == 3)
+      Funcs[i] = M.getOrInsertFunction(Sized[i], Attr, VoidTy, PtrTy, PtrTy);
+    else if (NFallbackArgs == 2)
+      Funcs[i] = M.getOrInsertFunction(Sized[i], Attr, VoidTy, PtrTy);
+  }
 
-  NsanCopyFunction[3] = M.getOrInsertFunction("__nsan_copy_values", Attr,
-                                              VoidTy, PtrTy, PtrTy, IntptrTy);
+  if (NFallbackArgs == 3)
+    Funcs[N - 1] =
+        M.getOrInsertFunction(Fallback, Attr, VoidTy, PtrTy, PtrTy, IntptrTy);
+  else if (NFallbackArgs == 2)
+    Funcs[N - 1] =
+        M.getOrInsertFunction(Fallback, Attr, VoidTy, PtrTy, IntptrTy);
 }
 
-std::pair<FunctionCallee, int>
-NsanCopy::getCalleeAndNumArgsForMemoryOp(uint64_t MemOpSize) const {
+template <size_t N, int NFallbackArgs>
+FunctionCallee
+NsanMemOpFn<N, NFallbackArgs>::getFunctionFor(uint64_t MemOpSize) const {
   size_t Idx =
       MemOpSize == 4 ? 0 : (MemOpSize == 8 ? 1 : (MemOpSize == 16 ? 2 : 3));
-  return Idx == 3 ? std::make_pair(NsanCopyFunction[Idx], 3)
-                  : std::make_pair(NsanCopyFunction[Idx], 2);
-}
-
-FunctionCallee NsanCopy::getGeneralFunction() const {
-  return NsanCopyFunction.back();
-}
-
-// Helper class for __nsan_set_value_unknown* family
-class NsanSetUnknown : public NsanInstrumentationFunction {
-public:
-  NsanSetUnknown(Module &M);
-  std::pair<FunctionCallee, int>
-  getCalleeAndNumArgsForMemoryOp(uint64_t MemOpSize) const override;
-  FunctionCallee getGeneralFunction() const override;
-
-private:
-  std::array<FunctionCallee, 4> NsanSetUnknownFunction;
-};
-
-NsanSetUnknown::NsanSetUnknown(Module &M) {
-  LLVMContext &Ctx = M.getContext();
-  AttributeList Attr;
-  Attr = Attr.addFnAttribute(Ctx, Attribute::NoUnwind);
-  Type *PtrTy = PointerType::getUnqual(Ctx);
-  Type *VoidTy = Type::getVoidTy(Ctx);
-  IntegerType *IntptrTy = M.getDataLayout().getIntPtrType(Ctx);
-
-  NsanSetUnknownFunction[0] =
-      M.getOrInsertFunction("__nsan_set_value_unknown_4", Attr, VoidTy, PtrTy);
 
-  NsanSetUnknownFunction[1] =
-      M.getOrInsertFunction("__nsan_set_value_unknown_8", Attr, VoidTy, PtrTy);
-
-  NsanSetUnknownFunction[2] =
-      M.getOrInsertFunction("__nsan_set_value_unknown_16", Attr, VoidTy, PtrTy);
-
-  NsanSetUnknownFunction[3] = M.getOrInsertFunction(
-      "__nsan_set_value_unknown", Attr, VoidTy, PtrTy, IntptrTy);
-}
-
-std::pair<FunctionCallee, int>
-NsanSetUnknown::getCalleeAndNumArgsForMemoryOp(const uint64_t MemOpSize) const {
-  size_t Idx =
-      MemOpSize == 4 ? 0 : (MemOpSize == 8 ? 1 : (MemOpSize == 16 ? 2 : 3));
-  return Idx == 3 ? std::make_pair(NsanSetUnknownFunction[Idx], 2)
-                  : std::make_pair(NsanSetUnknownFunction[Idx], 1);
+  assert(Idx <= N - 1 && "Functions array is too small");
+  return Funcs[Idx];
 }
 
-FunctionCallee NsanSetUnknown::getGeneralFunction() const {
-  return NsanSetUnknownFunction.back();
+template <size_t N, int NFallbackArgs>
+FunctionCallee NsanMemOpFn<N, NFallbackArgs>::getFallback() const {
+  return Funcs.back();
 }
 
 /// Instantiating NumericalStabilitySanitizer inserts the nsan runtime library
@@ -655,8 +611,8 @@ class NumericalStabilitySanitizer {
   FunctionCallee NsanCheckValue[FTValueType::kNumValueTypes] = {};
   FunctionCallee NsanFCmpFail[FTValueType::kNumValueTypes] = {};
 
-  NsanCopy NsanCopyFns;
-  NsanSetUnknown NsanSetUnknownFns;
+  NsanMemOpFn<4, 3> NsanCopyFns;
+  NsanMemOpFn<4, 2> NsanSetUnknownFns;
 
   FunctionCallee NsanGetRawShadowTypePtr;
   FunctionCallee NsanGetRawShadowPtr;
@@ -701,7 +657,13 @@ static GlobalValue *createThreadLocalGV(const char *Name, Module &M, Type *Ty) {
 
 NumericalStabilitySanitizer::NumericalStabilitySanitizer(Module &M)
     : DL(M.getDataLayout()), Context(M.getContext()), Config(Context),
-      NsanCopyFns(M), NsanSetUnknownFns(M) {
+      NsanCopyFns(M, {"__nsan_copy_4", "__nsan_copy_8", "__nsan_copy_16"},
+                  "__nsan_copy_values"),
+      NsanSetUnknownFns(M,
+                        {"__nsan_set_value_unknown_4",
+                         "__nsan_set_value_unknown_8",
+                         "__nsan_set_value_unknown_16"},
+                        "__nsan_set_value_unknown") {
   IntptrTy = DL.getIntPtrType(Context);
   Type *PtrTy = PointerType::getUnqual(Context);
   Type *Int32Ty = Type::getInt32Ty(Context);
@@ -1978,7 +1940,7 @@ void NumericalStabilitySanitizer::propagateNonFTStore(
     }
   }
   // All other stores just reset the shadow value to unknown.
-  Builder.CreateCall(NsanSetUnknownFns.getGeneralFunction(), {Dst, ValueSize});
+  Builder.CreateCall(NsanSetUnknownFns.getFallback(), {Dst, ValueSize});
 }
 
 void NumericalStabilitySanitizer::propagateShadowValues(
@@ -2237,32 +2199,29 @@ static uint64_t GetMemOpSize(Value *V) {
 bool NumericalStabilitySanitizer::instrumentMemIntrinsic(MemIntrinsic *MI) {
   IRBuilder<> Builder(MI);
   if (auto *M = dyn_cast<MemSetInst>(MI)) {
-    std::pair<FunctionCallee, int> SetUnknownFn =
-        NsanSetUnknownFns.getCalleeAndNumArgsForMemoryOp(
-            GetMemOpSize(M->getArgOperand(2)));
-    if (SetUnknownFn.second == 1)
-      Builder.CreateCall(SetUnknownFn.first, {/*Address=*/M->getArgOperand(0)});
+    FunctionCallee SetUnknownFn =
+        NsanSetUnknownFns.getFunctionFor(GetMemOpSize(M->getArgOperand(2)));
+    if (SetUnknownFn.getFunctionType()->getNumParams() == 1)
+      Builder.CreateCall(SetUnknownFn, {/*Address=*/M->getArgOperand(0)});
     else
-      Builder.CreateCall(SetUnknownFn.first,
+      Builder.CreateCall(SetUnknownFn,
                          {/*Address=*/M->getArgOperand(0),
                           /*Size=*/Builder.CreateIntCast(M->getArgOperand(2),
                                                          IntptrTy, false)});
 
   } else if (auto *M = dyn_cast<MemTransferInst>(MI)) {
-    std::pair<FunctionCallee, int> CopyFn =
-        NsanCopyFns.getCalleeAndNumArgsForMemoryOp(
-            GetMemOpSize(M->getArgOperand(2)));
+    FunctionCallee CopyFn =
+        NsanCopyFns.getFunctionFor(GetMemOpSize(M->getArgOperand(2)));
 
-    if (CopyFn.second == 2)
-      Builder.CreateCall(CopyFn.first, {/*Destination=*/M->getArgOperand(0),
-                                        /*Source=*/M->getArgOperand(1)});
+    if (CopyFn.getFunctionType()->getNumParams() == 2)
+      Builder.CreateCall(CopyFn, {/*Destination=*/M->getArgOperand(0),
+                                  /*Source=*/M->getArgOperand(1)});
     else
-      Builder.CreateCall(
-          CopyFn.first,
-          {/*Destination=*/M->getArgOperand(0),
-           /*Source=*/M->getArgOperand(1),
-           /*Size=*/
-           Builder.CreateIntCast(M->getArgOperand(2), IntptrTy, false)});
+      Builder.CreateCall(CopyFn, {/*Destination=*/M->getArgOperand(0),
+                                  /*Source=*/M->getArgOperand(1),
+                                  /*Size=*/
+                                  Builder.CreateIntCast(M->getArgOperand(2),
+                                                        IntptrTy, false)});
   }
   return false;
 }

>From fbef05a086574c0b3e086287e4e349951de5e60b Mon Sep 17 00:00:00 2001
From: Dmitry Chestnykh <dm.chestnykh at gmail.com>
Date: Wed, 17 Jul 2024 07:54:57 +0300
Subject: [PATCH 10/19] [nsan] Pass number of fallback fn args as ctor param

---
 .../NumericalStabilitySanitizer.cpp           | 27 +++++++++----------
 1 file changed, 13 insertions(+), 14 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
index f9ce473f4d94c..601bb254d611f 100644
--- a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
@@ -496,9 +496,10 @@ class ValueToShadowMap {
 
 // First parameter is the number of functions
 // Second parameter is the number of fallback function arguments
-template <size_t N, int NFallbackArgs> class NsanMemOpFn {
+template <size_t N> class NsanMemOpFn {
 public:
-  NsanMemOpFn(Module &M, ArrayRef<StringRef> Sized, StringRef Fallback);
+  NsanMemOpFn(Module &M, ArrayRef<StringRef> Sized, StringRef Fallback,
+              int NFallbackArgs);
   // Number of parameters can be extracted from FunctionCallee
   FunctionCallee getFunctionFor(uint64_t MemOpSize) const;
   FunctionCallee getFallback() const;
@@ -507,9 +508,9 @@ template <size_t N, int NFallbackArgs> class NsanMemOpFn {
   std::array<FunctionCallee, N> Funcs;
 };
 
-template <size_t N, int NFallbackArgs>
-NsanMemOpFn<N, NFallbackArgs>::NsanMemOpFn(Module &M, ArrayRef<StringRef> Sized,
-                                           StringRef Fallback) {
+template <size_t N>
+NsanMemOpFn<N>::NsanMemOpFn(Module &M, ArrayRef<StringRef> Sized,
+                            StringRef Fallback, int NFallbackArgs) {
   LLVMContext &Ctx = M.getContext();
   AttributeList Attr;
   Attr = Attr.addFnAttribute(Ctx, Attribute::NoUnwind);
@@ -532,9 +533,8 @@ NsanMemOpFn<N, NFallbackArgs>::NsanMemOpFn(Module &M, ArrayRef<StringRef> Sized,
         M.getOrInsertFunction(Fallback, Attr, VoidTy, PtrTy, IntptrTy);
 }
 
-template <size_t N, int NFallbackArgs>
-FunctionCallee
-NsanMemOpFn<N, NFallbackArgs>::getFunctionFor(uint64_t MemOpSize) const {
+template <size_t N>
+FunctionCallee NsanMemOpFn<N>::getFunctionFor(uint64_t MemOpSize) const {
   size_t Idx =
       MemOpSize == 4 ? 0 : (MemOpSize == 8 ? 1 : (MemOpSize == 16 ? 2 : 3));
 
@@ -542,8 +542,7 @@ NsanMemOpFn<N, NFallbackArgs>::getFunctionFor(uint64_t MemOpSize) const {
   return Funcs[Idx];
 }
 
-template <size_t N, int NFallbackArgs>
-FunctionCallee NsanMemOpFn<N, NFallbackArgs>::getFallback() const {
+template <size_t N> FunctionCallee NsanMemOpFn<N>::getFallback() const {
   return Funcs.back();
 }
 
@@ -611,8 +610,8 @@ class NumericalStabilitySanitizer {
   FunctionCallee NsanCheckValue[FTValueType::kNumValueTypes] = {};
   FunctionCallee NsanFCmpFail[FTValueType::kNumValueTypes] = {};
 
-  NsanMemOpFn<4, 3> NsanCopyFns;
-  NsanMemOpFn<4, 2> NsanSetUnknownFns;
+  NsanMemOpFn<4> NsanCopyFns;
+  NsanMemOpFn<4> NsanSetUnknownFns;
 
   FunctionCallee NsanGetRawShadowTypePtr;
   FunctionCallee NsanGetRawShadowPtr;
@@ -658,12 +657,12 @@ static GlobalValue *createThreadLocalGV(const char *Name, Module &M, Type *Ty) {
 NumericalStabilitySanitizer::NumericalStabilitySanitizer(Module &M)
     : DL(M.getDataLayout()), Context(M.getContext()), Config(Context),
       NsanCopyFns(M, {"__nsan_copy_4", "__nsan_copy_8", "__nsan_copy_16"},
-                  "__nsan_copy_values"),
+                  "__nsan_copy_values", 3),
       NsanSetUnknownFns(M,
                         {"__nsan_set_value_unknown_4",
                          "__nsan_set_value_unknown_8",
                          "__nsan_set_value_unknown_16"},
-                        "__nsan_set_value_unknown") {
+                        "__nsan_set_value_unknown", 2) {
   IntptrTy = DL.getIntPtrType(Context);
   Type *PtrTy = PointerType::getUnqual(Context);
   Type *Int32Ty = Type::getInt32Ty(Context);

>From 57c755ac3718130bb678f196ca4cfe8bee50eb38 Mon Sep 17 00:00:00 2001
From: Dmitry Chestnykh <dm.chestnykh at gmail.com>
Date: Wed, 17 Jul 2024 08:04:55 +0300
Subject: [PATCH 11/19] [nsan] Remove unneeded include

---
 .../Transforms/Instrumentation/NumericalStabilitySanitizer.cpp   | 1 -
 1 file changed, 1 deletion(-)

diff --git a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
index 601bb254d611f..2eb21e3e98e6f 100644
--- a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
@@ -45,7 +45,6 @@
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 
 #include <cstdint>
-#include <map>
 
 using namespace llvm;
 

>From 381767f1156a44f1d10895960b855dd7a0de4bd7 Mon Sep 17 00:00:00 2001
From: Dmitry Chestnykh <dm.chestnykh at gmail.com>
Date: Wed, 17 Jul 2024 08:19:59 +0300
Subject: [PATCH 12/19] [compiler-rt][nsan] Fix wrong formatting

clang-format cannot recognize that this is multiplication,
not pointer declaration.
See #99271
---
 compiler-rt/lib/nsan/nsan.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/compiler-rt/lib/nsan/nsan.cpp b/compiler-rt/lib/nsan/nsan.cpp
index 13ef3f9da4fcb..bfbc22bb80b1f 100644
--- a/compiler-rt/lib/nsan/nsan.cpp
+++ b/compiler-rt/lib/nsan/nsan.cpp
@@ -68,7 +68,7 @@ __nsan_copy_values(const u8 *daddr, const u8 *saddr, uptr size) {
     __builtin_memmove((void *)GetShadowTypeAddrFor(daddr),                     \
                       GetShadowTypeAddrFor(saddr), N);                         \
     __builtin_memmove((void *)GetShadowAddrFor(daddr),                         \
-                      GetShadowAddrFor(saddr), N *kShadowScale);               \
+                      GetShadowAddrFor(saddr), N * kShadowScale);              \
   }
 
 NSAN_COPY_VALUES_N(4)

>From 16320d24a81271947139dd728cc4d192abaea9d0 Mon Sep 17 00:00:00 2001
From: Dmitry Chestnykh <dm.chestnykh at gmail.com>
Date: Wed, 17 Jul 2024 09:28:14 +0300
Subject: [PATCH 13/19] [nsan] Calculate function type once

With this approach the constructor of `NsanMemOpFn`
can be simplified
---
 .../lib/sanitizer_common/sanitizer_linux.cpp  | 83 +------------------
 .../TestCases/Posix/dump_registers.cpp        | 12 +--
 .../NumericalStabilitySanitizer.cpp           | 29 ++++---
 3 files changed, 21 insertions(+), 103 deletions(-)

diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_linux.cpp b/compiler-rt/lib/sanitizer_common/sanitizer_linux.cpp
index 58233a3c0f607..7935c88204a05 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_linux.cpp
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_linux.cpp
@@ -2118,89 +2118,8 @@ bool SignalContext::IsTrueFaultingAddress() const {
   return si->si_signo == SIGSEGV && si->si_code != 128;
 }
 
-static const char *RegNumToRegName(int reg) {
-#  if defined(__x86_64__)
-  switch (reg) {
-    case REG_RAX:
-      return "rax";
-    case REG_RBX:
-      return "rbx";
-    case REG_RCX:
-      return "rcx";
-    case REG_RDX:
-      return "rdx";
-    case REG_RDI:
-      return "rdi";
-    case REG_RSI:
-      return "rsi";
-    case REG_RBP:
-      return "rbp";
-    case REG_RSP:
-      return "rsp";
-    case REG_R8:
-      return "r8";
-    case REG_R9:
-      return "r9";
-    case REG_R10:
-      return "r10";
-    case REG_R11:
-      return "r11";
-    case REG_R12:
-      return "r12";
-    case REG_R13:
-      return "r13";
-    case REG_R14:
-      return "r14";
-    case REG_R15:
-      return "r15";
-    default:
-      return NULL;
-  }
-#  endif
-  return NULL;
-}
-
 void SignalContext::DumpAllRegisters(void *context) {
-#  if SANITIZER_LINUX
-  ucontext_t *ucontext = (ucontext_t *)context;
-#    define DUMPREG64(r)                             \
-      Printf("%s = 0x%016llx  ", RegNumToRegName(r), \
-             ucontext->uc_mcontext.gregs[r]);
-#    define DUMPREG_(r) \
-      Printf(" ");      \
-      DUMPREG(r);
-#    define DUMPREG__(r) \
-      Printf("  ");      \
-      DUMPREG(r);
-#    define DUMPREG___(r) \
-      Printf("   ");      \
-      DUMPREG(r);
-#    if defined(__x86_64__)
-#      define DUMPREG(r) DUMPREG64(r)
-  Report("Register values:\n");
-  DUMPREG(REG_RAX);
-  DUMPREG(REG_RBX);
-  DUMPREG(REG_RCX);
-  DUMPREG(REG_RDX);
-  Printf("\n");
-  DUMPREG(REG_RDI);
-  DUMPREG(REG_RSI);
-  DUMPREG(REG_RBP);
-  DUMPREG(REG_RSP);
-  Printf("\n");
-  DUMPREG_(REG_R8);
-  DUMPREG_(REG_R9);
-  DUMPREG(REG_R10);
-  DUMPREG(REG_R11);
-  Printf("\n");
-  DUMPREG(REG_R12);
-  DUMPREG(REG_R13);
-  DUMPREG(REG_R14);
-  DUMPREG(REG_R15);
-  Printf("\n");
-#    endif
-#  endif
-  // FIXME: Implement this for other OSes and architectures.
+  // FIXME: Implement this.
 }
 
 static void GetPcSpBp(void *context, uptr *pc, uptr *sp, uptr *bp) {
diff --git a/compiler-rt/test/sanitizer_common/TestCases/Posix/dump_registers.cpp b/compiler-rt/test/sanitizer_common/TestCases/Posix/dump_registers.cpp
index 2cbd40e924263..f09b2bf4447cc 100644
--- a/compiler-rt/test/sanitizer_common/TestCases/Posix/dump_registers.cpp
+++ b/compiler-rt/test/sanitizer_common/TestCases/Posix/dump_registers.cpp
@@ -1,20 +1,16 @@
 // Check that sanitizer prints registers dump_registers on dump_registers=1
 // RUN: %clangxx  %s -o %t
-// RUN: %env_tool_opts=dump_registers=0 not %run %t 2>&1 | FileCheck %s --check-prefix=CHECK-NODUMP
-// RUN: not %run %t 2>&1 | FileCheck %s --check-prefix=CHECK-DUMP
+// RUN: %env_tool_opts=dump_registers=0 %run %t 2>&1 | FileCheck %s --check-prefix=CHECK-NODUMP
+// RUN: %run %t 2>&1 | FileCheck %s --check-prefix=CHECK-DUMP
 //
-// FIXME: Implement for other OSes and architectures.
-// REQUIRES: x86_64-target-arch, linux
+// FIXME: Implement.
+// XFAIL: *
 
 #include <signal.h>
 
 int main() {
   raise(SIGSEGV);
   // CHECK-DUMP: Register values
-  // CHECK-DUMP-NEXT: rax = {{0x[0-9a-f]+}}  rbx = {{0x[0-9a-f]+}}  rcx = {{0x[0-9a-f]+}}  rdx = {{0x[0-9a-f]+}}
-  // CHECK-DUMP-NEXT: rdi = {{0x[0-9a-f]+}}  rsi = {{0x[0-9a-f]+}}  rbp = {{0x[0-9a-f]+}}  rsp = {{0x[0-9a-f]+}}
-  // CHECK-DUMP-NEXT:  r8 = {{0x[0-9a-f]+}}   r9 = {{0x[0-9a-f]+}}  r10 = {{0x[0-9a-f]+}}  r11 = {{0x[0-9a-f]+}}
-  // CHECK-DUMP-NEXT: r12 = {{0x[0-9a-f]+}}   r13 = {{0x[0-9a-f]+}} r14 = {{0x[0-9a-f]+}}  r15 = {{0x[0-9a-f]+}}
   // CHECK-NODUMP-NOT: Register values
   return 0;
 }
diff --git a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
index 2eb21e3e98e6f..78b9112ceef37 100644
--- a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
@@ -497,8 +497,7 @@ class ValueToShadowMap {
 // Second parameter is the number of fallback function arguments
 template <size_t N> class NsanMemOpFn {
 public:
-  NsanMemOpFn(Module &M, ArrayRef<StringRef> Sized, StringRef Fallback,
-              int NFallbackArgs);
+  NsanMemOpFn(Module &M, ArrayRef<StringRef> Sized, StringRef Fallback, size_t NumArgs);
   // Number of parameters can be extracted from FunctionCallee
   FunctionCallee getFunctionFor(uint64_t MemOpSize) const;
   FunctionCallee getFallback() const;
@@ -509,31 +508,34 @@ template <size_t N> class NsanMemOpFn {
 
 template <size_t N>
 NsanMemOpFn<N>::NsanMemOpFn(Module &M, ArrayRef<StringRef> Sized,
-                            StringRef Fallback, int NFallbackArgs) {
+                                           StringRef Fallback, size_t NumArgs) {
   LLVMContext &Ctx = M.getContext();
   AttributeList Attr;
   Attr = Attr.addFnAttribute(Ctx, Attribute::NoUnwind);
   Type *PtrTy = PointerType::getUnqual(Ctx);
   Type *VoidTy = Type::getVoidTy(Ctx);
   IntegerType *IntptrTy = M.getDataLayout().getIntPtrType(Ctx);
+  FunctionType *SizedFnTy;
 
-  for (size_t i = 0; i < N - 1; ++i) {
-    if (NFallbackArgs == 3)
-      Funcs[i] = M.getOrInsertFunction(Sized[i], Attr, VoidTy, PtrTy, PtrTy);
-    else if (NFallbackArgs == 2)
-      Funcs[i] = M.getOrInsertFunction(Sized[i], Attr, VoidTy, PtrTy);
-  }
+  if (NumArgs == 3)
+    SizedFnTy = FunctionType::get(VoidTy, {PtrTy, PtrTy}, false);
+  else
+    SizedFnTy = FunctionType::get(VoidTy, {PtrTy}, false);
+
+  for (size_t i = 0; i < N - 1; ++i)
+    Funcs[i] = M.getOrInsertFunction(Sized[i], SizedFnTy, Attr);
 
-  if (NFallbackArgs == 3)
+  if (NumArgs == 3)
     Funcs[N - 1] =
         M.getOrInsertFunction(Fallback, Attr, VoidTy, PtrTy, PtrTy, IntptrTy);
-  else if (NFallbackArgs == 2)
+  else
     Funcs[N - 1] =
         M.getOrInsertFunction(Fallback, Attr, VoidTy, PtrTy, IntptrTy);
 }
 
 template <size_t N>
-FunctionCallee NsanMemOpFn<N>::getFunctionFor(uint64_t MemOpSize) const {
+FunctionCallee
+NsanMemOpFn<N>::getFunctionFor(uint64_t MemOpSize) const {
   size_t Idx =
       MemOpSize == 4 ? 0 : (MemOpSize == 8 ? 1 : (MemOpSize == 16 ? 2 : 3));
 
@@ -541,7 +543,8 @@ FunctionCallee NsanMemOpFn<N>::getFunctionFor(uint64_t MemOpSize) const {
   return Funcs[Idx];
 }
 
-template <size_t N> FunctionCallee NsanMemOpFn<N>::getFallback() const {
+template <size_t N>
+FunctionCallee NsanMemOpFn<N>::getFallback() const {
   return Funcs.back();
 }
 

>From 8e6591fa082bf64871c33e029d5278a98c765e37 Mon Sep 17 00:00:00 2001
From: Dmitry Chestnykh <dm.chestnykh at gmail.com>
Date: Wed, 17 Jul 2024 09:40:10 +0300
Subject: [PATCH 14/19] [nsan] Update comment

---
 .../Transforms/Instrumentation/NumericalStabilitySanitizer.cpp | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
index 78b9112ceef37..0ca369f5d4bf9 100644
--- a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
@@ -493,8 +493,7 @@ class ValueToShadowMap {
   DenseMap<Value *, Value *> Map;
 };
 
-// First parameter is the number of functions
-// Second parameter is the number of fallback function arguments
+// Template parameter is the number of functions
 template <size_t N> class NsanMemOpFn {
 public:
   NsanMemOpFn(Module &M, ArrayRef<StringRef> Sized, StringRef Fallback, size_t NumArgs);

>From 31efb9bef097102cc2b941730ebf2baa423369e1 Mon Sep 17 00:00:00 2001
From: Dmitry Chestnykh <dm.chestnykh at gmail.com>
Date: Wed, 17 Jul 2024 11:36:12 +0300
Subject: [PATCH 15/19] [nsan] Correct code style

---
 .../Instrumentation/NumericalStabilitySanitizer.cpp   | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
index 0ca369f5d4bf9..daddc100554e9 100644
--- a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
@@ -496,7 +496,8 @@ class ValueToShadowMap {
 // Template parameter is the number of functions
 template <size_t N> class NsanMemOpFn {
 public:
-  NsanMemOpFn(Module &M, ArrayRef<StringRef> Sized, StringRef Fallback, size_t NumArgs);
+  NsanMemOpFn(Module &M, ArrayRef<StringRef> Sized, StringRef Fallback,
+              size_t NumArgs);
   // Number of parameters can be extracted from FunctionCallee
   FunctionCallee getFunctionFor(uint64_t MemOpSize) const;
   FunctionCallee getFallback() const;
@@ -507,7 +508,7 @@ template <size_t N> class NsanMemOpFn {
 
 template <size_t N>
 NsanMemOpFn<N>::NsanMemOpFn(Module &M, ArrayRef<StringRef> Sized,
-                                           StringRef Fallback, size_t NumArgs) {
+                            StringRef Fallback, size_t NumArgs) {
   LLVMContext &Ctx = M.getContext();
   AttributeList Attr;
   Attr = Attr.addFnAttribute(Ctx, Attribute::NoUnwind);
@@ -533,8 +534,7 @@ NsanMemOpFn<N>::NsanMemOpFn(Module &M, ArrayRef<StringRef> Sized,
 }
 
 template <size_t N>
-FunctionCallee
-NsanMemOpFn<N>::getFunctionFor(uint64_t MemOpSize) const {
+FunctionCallee NsanMemOpFn<N>::getFunctionFor(uint64_t MemOpSize) const {
   size_t Idx =
       MemOpSize == 4 ? 0 : (MemOpSize == 8 ? 1 : (MemOpSize == 16 ? 2 : 3));
 
@@ -542,8 +542,7 @@ NsanMemOpFn<N>::getFunctionFor(uint64_t MemOpSize) const {
   return Funcs[Idx];
 }
 
-template <size_t N>
-FunctionCallee NsanMemOpFn<N>::getFallback() const {
+template <size_t N> FunctionCallee NsanMemOpFn<N>::getFallback() const {
   return Funcs.back();
 }
 

>From 8d1d201fe0e04155f324f9174c09e63b5cf6c176 Mon Sep 17 00:00:00 2001
From: Dmitry Chestnykh <dm.chestnykh at gmail.com>
Date: Thu, 18 Jul 2024 07:38:12 +0300
Subject: [PATCH 16/19] [nsan] Remove template parameter and refactor code

---
 .../NumericalStabilitySanitizer.cpp           | 68 +++++++++++--------
 1 file changed, 39 insertions(+), 29 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
index daddc100554e9..77c3195f6ae2e 100644
--- a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
@@ -493,8 +493,9 @@ class ValueToShadowMap {
   DenseMap<Value *, Value *> Map;
 };
 
-// Template parameter is the number of functions
-template <size_t N> class NsanMemOpFn {
+// First parameter is the number of functions
+// Second parameter is the number of fallback function arguments
+class NsanMemOpFn {
 public:
   NsanMemOpFn(Module &M, ArrayRef<StringRef> Sized, StringRef Fallback,
               size_t NumArgs);
@@ -503,48 +504,57 @@ template <size_t N> class NsanMemOpFn {
   FunctionCallee getFallback() const;
 
 private:
-  std::array<FunctionCallee, N> Funcs;
+  SmallVector<FunctionCallee> Funcs;
+  size_t NumSizedFuncs;
 };
 
-template <size_t N>
-NsanMemOpFn<N>::NsanMemOpFn(Module &M, ArrayRef<StringRef> Sized,
-                            StringRef Fallback, size_t NumArgs) {
+NsanMemOpFn::NsanMemOpFn(Module &M, ArrayRef<StringRef> Sized,
+                         StringRef Fallback, size_t NumArgs) {
   LLVMContext &Ctx = M.getContext();
   AttributeList Attr;
   Attr = Attr.addFnAttribute(Ctx, Attribute::NoUnwind);
   Type *PtrTy = PointerType::getUnqual(Ctx);
   Type *VoidTy = Type::getVoidTy(Ctx);
   IntegerType *IntptrTy = M.getDataLayout().getIntPtrType(Ctx);
-  FunctionType *SizedFnTy;
+  FunctionType *SizedFnTy = nullptr;
 
-  if (NumArgs == 3)
-    SizedFnTy = FunctionType::get(VoidTy, {PtrTy, PtrTy}, false);
-  else
-    SizedFnTy = FunctionType::get(VoidTy, {PtrTy}, false);
+  NumSizedFuncs = Sized.size();
 
-  for (size_t i = 0; i < N - 1; ++i)
-    Funcs[i] = M.getOrInsertFunction(Sized[i], SizedFnTy, Attr);
+  // Reserve space for sized funcs and for fallback
+  Funcs.reserve(NumSizedFuncs + 1);
 
-  if (NumArgs == 3)
-    Funcs[N - 1] =
+  if (NumArgs == 3) {
+    Funcs[NumSizedFuncs] =
         M.getOrInsertFunction(Fallback, Attr, VoidTy, PtrTy, PtrTy, IntptrTy);
-  else
-    Funcs[N - 1] =
+    SizedFnTy = FunctionType::get(VoidTy, {PtrTy, PtrTy}, false);
+  } else if (NumArgs == 2) {
+    Funcs[NumSizedFuncs] =
         M.getOrInsertFunction(Fallback, Attr, VoidTy, PtrTy, IntptrTy);
+    SizedFnTy = FunctionType::get(VoidTy, {PtrTy}, false);
+  } else {
+    assert(!"Unexpected value of sized functions arguments");
+  }
+
+  for (size_t i = 0; i < NumSizedFuncs; ++i)
+    Funcs[i] = M.getOrInsertFunction(Sized[i], SizedFnTy, Attr);
 }
 
-template <size_t N>
-FunctionCallee NsanMemOpFn<N>::getFunctionFor(uint64_t MemOpSize) const {
-  size_t Idx =
-      MemOpSize == 4 ? 0 : (MemOpSize == 8 ? 1 : (MemOpSize == 16 ? 2 : 3));
+FunctionCallee NsanMemOpFn::getFunctionFor(uint64_t MemOpSize) const {
+  // We have NumSizedFuncs + 1 elements in `Funcs`
+  size_t MaxIdx = NumSizedFuncs;
+
+  // Now `getFunctionFor` operates on `Funcs` of size 4 (at least) and the
+  // following code assumes that the number of functions in `Func` is sufficient
+  assert(MaxIdx >= 3 && "Unexpected MaxIdx value");
+
+  size_t Idx = MemOpSize == 4
+                   ? 0
+                   : (MemOpSize == 8 ? 1 : (MemOpSize == 16 ? 2 : MaxIdx));
 
-  assert(Idx <= N - 1 && "Functions array is too small");
   return Funcs[Idx];
 }
 
-template <size_t N> FunctionCallee NsanMemOpFn<N>::getFallback() const {
-  return Funcs.back();
-}
+FunctionCallee NsanMemOpFn::getFallback() const { return Funcs[NumSizedFuncs]; }
 
 /// Instantiating NumericalStabilitySanitizer inserts the nsan runtime library
 /// API function declarations into the module if they don't exist already.
@@ -610,8 +620,8 @@ class NumericalStabilitySanitizer {
   FunctionCallee NsanCheckValue[FTValueType::kNumValueTypes] = {};
   FunctionCallee NsanFCmpFail[FTValueType::kNumValueTypes] = {};
 
-  NsanMemOpFn<4> NsanCopyFns;
-  NsanMemOpFn<4> NsanSetUnknownFns;
+  NsanMemOpFn NsanCopyFns;
+  NsanMemOpFn NsanSetUnknownFns;
 
   FunctionCallee NsanGetRawShadowTypePtr;
   FunctionCallee NsanGetRawShadowPtr;
@@ -657,12 +667,12 @@ static GlobalValue *createThreadLocalGV(const char *Name, Module &M, Type *Ty) {
 NumericalStabilitySanitizer::NumericalStabilitySanitizer(Module &M)
     : DL(M.getDataLayout()), Context(M.getContext()), Config(Context),
       NsanCopyFns(M, {"__nsan_copy_4", "__nsan_copy_8", "__nsan_copy_16"},
-                  "__nsan_copy_values", 3),
+                  "__nsan_copy_values", /*NumArgs=*/3),
       NsanSetUnknownFns(M,
                         {"__nsan_set_value_unknown_4",
                          "__nsan_set_value_unknown_8",
                          "__nsan_set_value_unknown_16"},
-                        "__nsan_set_value_unknown", 2) {
+                        "__nsan_set_value_unknown", /*NumArgs=*/2) {
   IntptrTy = DL.getIntPtrType(Context);
   Type *PtrTy = PointerType::getUnqual(Context);
   Type *Int32Ty = Type::getInt32Ty(Context);

>From 6e52b73ac1c836be7330df95b552981a2fd9eeef Mon Sep 17 00:00:00 2001
From: Dmitry Chestnykh <dm.chestnykh at gmail.com>
Date: Thu, 18 Jul 2024 07:39:37 +0300
Subject: [PATCH 17/19] [nsan] Remove wrong comments

---
 .../Transforms/Instrumentation/NumericalStabilitySanitizer.cpp  | 2 --
 1 file changed, 2 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
index 77c3195f6ae2e..49ac9a4d03b32 100644
--- a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
@@ -493,8 +493,6 @@ class ValueToShadowMap {
   DenseMap<Value *, Value *> Map;
 };
 
-// First parameter is the number of functions
-// Second parameter is the number of fallback function arguments
 class NsanMemOpFn {
 public:
   NsanMemOpFn(Module &M, ArrayRef<StringRef> Sized, StringRef Fallback,

>From d01b89932c162edf1cc26d3525e83e5cee2b2bf0 Mon Sep 17 00:00:00 2001
From: Dmitry Chestnykh <dm.chestnykh at gmail.com>
Date: Thu, 18 Jul 2024 09:07:04 +0300
Subject: [PATCH 18/19] [nsan] Reorganize work with SmallVector

---
 .../NumericalStabilitySanitizer.cpp           | 26 +++++++------------
 1 file changed, 10 insertions(+), 16 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
index 49ac9a4d03b32..ceaece0faa55c 100644
--- a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
@@ -518,41 +518,35 @@ NsanMemOpFn::NsanMemOpFn(Module &M, ArrayRef<StringRef> Sized,
 
   NumSizedFuncs = Sized.size();
 
-  // Reserve space for sized funcs and for fallback
-  Funcs.reserve(NumSizedFuncs + 1);
-
+  // First entry is fallback function
   if (NumArgs == 3) {
-    Funcs[NumSizedFuncs] =
-        M.getOrInsertFunction(Fallback, Attr, VoidTy, PtrTy, PtrTy, IntptrTy);
+    Funcs.push_back(
+        M.getOrInsertFunction(Fallback, Attr, VoidTy, PtrTy, PtrTy, IntptrTy));
     SizedFnTy = FunctionType::get(VoidTy, {PtrTy, PtrTy}, false);
   } else if (NumArgs == 2) {
-    Funcs[NumSizedFuncs] =
-        M.getOrInsertFunction(Fallback, Attr, VoidTy, PtrTy, IntptrTy);
+    Funcs.push_back(
+        M.getOrInsertFunction(Fallback, Attr, VoidTy, PtrTy, IntptrTy));
     SizedFnTy = FunctionType::get(VoidTy, {PtrTy}, false);
   } else {
     assert(!"Unexpected value of sized functions arguments");
   }
 
   for (size_t i = 0; i < NumSizedFuncs; ++i)
-    Funcs[i] = M.getOrInsertFunction(Sized[i], SizedFnTy, Attr);
+    Funcs.push_back(M.getOrInsertFunction(Sized[i], SizedFnTy, Attr));
 }
 
 FunctionCallee NsanMemOpFn::getFunctionFor(uint64_t MemOpSize) const {
-  // We have NumSizedFuncs + 1 elements in `Funcs`
-  size_t MaxIdx = NumSizedFuncs;
-
   // Now `getFunctionFor` operates on `Funcs` of size 4 (at least) and the
   // following code assumes that the number of functions in `Func` is sufficient
-  assert(MaxIdx >= 3 && "Unexpected MaxIdx value");
+  assert(NumSizedFuncs >= 3 && "Unexpected number of sized functions");
 
-  size_t Idx = MemOpSize == 4
-                   ? 0
-                   : (MemOpSize == 8 ? 1 : (MemOpSize == 16 ? 2 : MaxIdx));
+  size_t Idx =
+      MemOpSize == 4 ? 1 : (MemOpSize == 8 ? 2 : (MemOpSize == 16 ? 3 : 0));
 
   return Funcs[Idx];
 }
 
-FunctionCallee NsanMemOpFn::getFallback() const { return Funcs[NumSizedFuncs]; }
+FunctionCallee NsanMemOpFn::getFallback() const { return Funcs[0]; }
 
 /// Instantiating NumericalStabilitySanitizer inserts the nsan runtime library
 /// API function declarations into the module if they don't exist already.

>From 2f3a9a24cbbb8017c8d649f086cf12e564fb85ea Mon Sep 17 00:00:00 2001
From: Dmitry Chestnykh <dm.chestnykh at gmail.com>
Date: Thu, 18 Jul 2024 09:08:02 +0300
Subject: [PATCH 19/19] [nsan] Drop unneeded comment

---
 .../Transforms/Instrumentation/NumericalStabilitySanitizer.cpp   | 1 -
 1 file changed, 1 deletion(-)

diff --git a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
index ceaece0faa55c..0b54f571f5508 100644
--- a/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/NumericalStabilitySanitizer.cpp
@@ -497,7 +497,6 @@ class NsanMemOpFn {
 public:
   NsanMemOpFn(Module &M, ArrayRef<StringRef> Sized, StringRef Fallback,
               size_t NumArgs);
-  // Number of parameters can be extracted from FunctionCallee
   FunctionCallee getFunctionFor(uint64_t MemOpSize) const;
   FunctionCallee getFallback() const;
 



More information about the llvm-commits mailing list