From ee564abb29f44ea071be26eef1861725452bedf7 Mon Sep 17 00:00:00 2001
From: Chris Apple <cja-private@pm.me>
Date: Thu, 8 Aug 2024 18:06:43 -0700
Subject: [PATCH] [LLVM][rtsan] Add nosanitize_realtime instrumentation

---
 .../Instrumentation/RealtimeSanitizer.cpp     | 13 ++++++++++
 .../RealtimeSanitizer/rtsan_nosanitize.ll     | 25 +++++++++++++++++++
 2 files changed, 38 insertions(+)
 create mode 100644 llvm/test/Instrumentation/RealtimeSanitizer/rtsan_nosanitize.ll

diff --git a/llvm/lib/Transforms/Instrumentation/RealtimeSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/RealtimeSanitizer.cpp
index 7854cf4f2c625..1707ec298d6d2 100644
--- a/llvm/lib/Transforms/Instrumentation/RealtimeSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/RealtimeSanitizer.cpp
@@ -19,6 +19,8 @@
 
 #include "llvm/Transforms/Instrumentation/RealtimeSanitizer.h"
 
+#include <cassert>
+
 using namespace llvm;
 
 static void insertCallBeforeInstruction(Function &Fn, Instruction &Instruction,
@@ -51,6 +53,7 @@ RealtimeSanitizerPass::RealtimeSanitizerPass(
 PreservedAnalyses RealtimeSanitizerPass::run(Function &F,
                                              AnalysisManager<Function> &AM) {
   if (F.hasFnAttribute(Attribute::SanitizeRealtime)) {
+    assert(!F.hasFnAttribute(Attribute::NoSanitizeRealtime));
     insertCallAtFunctionEntryPoint(F, "__rtsan_realtime_enter");
     insertCallAtAllFunctionExitPoints(F, "__rtsan_realtime_exit");
 
@@ -59,5 +62,15 @@ PreservedAnalyses RealtimeSanitizerPass::run(Function &F,
     return PA;
   }
 
+  if (F.hasFnAttribute(Attribute::NoSanitizeRealtime)) {
+    assert(!F.hasFnAttribute(Attribute::SanitizeRealtime));
+    insertCallAtFunctionEntryPoint(F, "__rtsan_off");
+    insertCallAtAllFunctionExitPoints(F, "__rtsan_on");
+
+    PreservedAnalyses PA;
+    PA.preserveSet<CFGAnalyses>();
+    return PA;
+  }
+
   return PreservedAnalyses::all();
 }
diff --git a/llvm/test/Instrumentation/RealtimeSanitizer/rtsan_nosanitize.ll b/llvm/test/Instrumentation/RealtimeSanitizer/rtsan_nosanitize.ll
new file mode 100644
index 0000000000000..5a465df749da6
--- /dev/null
+++ b/llvm/test/Instrumentation/RealtimeSanitizer/rtsan_nosanitize.ll
@@ -0,0 +1,25 @@
+; RUN: opt < %s -passes=rtsan -S | FileCheck %s
+
+define void @nosanitized_function() #0 {
+  %1 = alloca ptr, align 8
+  %2 = call ptr @malloc(i64 noundef 2) #3
+  store ptr %2, ptr %1, align 8
+  ret void
+}
+
+declare ptr @malloc(i64 noundef) #1
+
+define noundef i32 @main() #2 {
+  %1 = alloca i32, align 4
+  store i32 0, ptr %1, align 4
+  call void @nosanitized_function() #4
+  ret i32 0
+}
+
+attributes #0 = { nosanitize_realtime }
+
+; CHECK-LABEL: @nosanitized_function()
+; CHECK-NEXT: call{{.*}}@__rtsan_off
+
+; CHECK: call{{.*}}@__rtsan_on
+; CHECK-NEXT: ret{{.*}}void