From 7efc4ddc07e3099bd45aff91b509795aa0bac9f2 Mon Sep 17 00:00:00 2001
From: TB Schardl <neboat@mit.edu>
Date: Wed, 18 Oct 2023 14:15:46 +0000
Subject: [PATCH] [AddressSanitizer] Analyze parallel-promotable allocas before
 ASan's instrumentation might invalidate task analysis.

---
 .../Instrumentation/AddressSanitizer.cpp      | 23 +++++++++-
 .../asan-parallel-promotable-alloca-check.ll  | 46 +++++++++++++++++++
 2 files changed, 68 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/Transforms/Tapir/asan-parallel-promotable-alloca-check.ll

diff --git a/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
index 06ccd3590819..8ae8e30b9ee0 100644
--- a/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
@@ -716,6 +716,7 @@ struct AddressSanitizer {
   bool maybeInsertAsanInitAtFunctionEntry(Function &F);
   bool maybeInsertDynamicShadowAtFunctionEntry(Function &F);
   void markEscapedLocalAllocas(Function &F);
+  void recordInterestingParallelAllocas(const Function &F);
 
 private:
   friend struct FunctionStackPoisoner;
@@ -773,6 +774,7 @@ struct AddressSanitizer {
   Value *LocalDynamicShadow = nullptr;
   const StackSafetyGlobalInfo *SSGI;
   DenseMap<const AllocaInst *, bool> ProcessedAllocas;
+  SmallPtrSet<const AllocaInst *, 16> InterestingParallelAllocas;
 
   FunctionCallee AMDGPUAddressShared;
   FunctionCallee AMDGPUAddressPrivate;
@@ -1266,7 +1268,7 @@ bool AddressSanitizer::isInterestingAlloca(const AllocaInst &AI) {
        // Promotable allocas are common under -O0.
        (!ClSkipPromotableAllocas || !isAllocaPromotable(&AI)) &&
        (!ClSkipPromotableAllocas ||
-        (TI->isSerial() || !TI->isAllocaParallelPromotable(&AI))) &&
+        (TI->isSerial() || InterestingParallelAllocas.contains(&AI))) &&
        // inalloca allocas are not treated as static, and we don't want
        // dynamic alloca instrumentation for them as well.
        !AI.isUsedWithInAlloca() &&
@@ -2673,6 +2675,21 @@ void AddressSanitizer::markEscapedLocalAllocas(Function &F) {
   }
 }
 
+void AddressSanitizer::recordInterestingParallelAllocas(const Function &F) {
+  if (!ClSkipPromotableAllocas || TI->isSerial())
+    return;
+
+  for (const BasicBlock &BB : F)
+    for (const Instruction &I : BB)
+      if (const AllocaInst *AI = dyn_cast<AllocaInst>(&I))
+        if (AI->getAllocatedType()->isSized() &&
+            ((!AI->isStaticAlloca()) || getAllocaSizeInBytes(*AI) > 0) &&
+            // We are only interested in allocas not promotable to registers.
+            // Promotable allocas are common under -O0.
+            !isAllocaPromotable(AI) && !TI->isAllocaParallelPromotable(AI))
+          InterestingParallelAllocas.insert(AI);
+}
+
 bool AddressSanitizer::suppressInstrumentationSiteForDebug(int &Instrumented) {
   bool ShouldInstrument =
       ClDebugMin < 0 || ClDebugMax < 0 ||
@@ -2715,6 +2732,10 @@ bool AddressSanitizer::instrumentFunction(Function &F,
   // can be passed to that intrinsic.
   markEscapedLocalAllocas(F);
 
+  // Record all interesting parallel allocas, using TaskInfo analysis before
+  // instrumentation may disrupt the validity of the analysis.
+  recordInterestingParallelAllocas(F);
+
   // We want to instrument every address only once per basic block (unless there
   // are calls between uses).
   SmallPtrSet<Value *, 16> TempsToInstrument;
diff --git a/llvm/test/Transforms/Tapir/asan-parallel-promotable-alloca-check.ll b/llvm/test/Transforms/Tapir/asan-parallel-promotable-alloca-check.ll
new file mode 100644
index 000000000000..d73982ca0b1d
--- /dev/null
+++ b/llvm/test/Transforms/Tapir/asan-parallel-promotable-alloca-check.ll
@@ -0,0 +1,46 @@
+; Check that ASan properly checks for parallel-promotable allocas.
+;
+; RUN: opt < %s -passes="asan<>" -S | FileCheck %s
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+; Function Attrs: nounwind willreturn memory(argmem: readwrite)
+declare token @llvm.syncregion.start() #0
+
+; Function Attrs: sanitize_address
+define void @do_timestep() #1 {
+entry:
+  %syncreg = tail call token @llvm.syncregion.start()
+  %0 = load i32, ptr null, align 8
+  %vla = alloca float, i64 0, align 16
+  detach within %syncreg, label %for.body.lr.ph.us, label %pfor.inc.us
+
+pfor.inc.us:                                      ; preds = %for.body.lr.ph.us, %entry
+  ret void
+
+for.body.lr.ph.us:                                ; preds = %entry
+  %call.us = call i32 null(ptr null, i32 0, i32 0, ptr %vla)
+  reattach within %syncreg, label %pfor.inc.us
+}
+
+; CHECK: define void @do_timestep()
+; CHECK: entry:
+; CHECK: br i1 %{{.*}}, label %[[ASAN_BLK_2:[a-zA-Z0-9_\.]+]], label %[[ASAN_BLK_5:[a-zA-Z0-9_\.]+]]
+
+; CHECK: [[ASAN_BLK_2]]:
+; CHECK: br i1 %{{.*}}, label %[[ASAN_BLK_4:[a-zA-Z0-9_\.]+]], label %[[ASAN_BLK_5]]
+
+; CHECK: [[ASAN_BLK_5]]:
+; CHECK-NEXT: load i32, ptr null
+; CHECK-NEXT: %vla = alloca float, i64 0
+; CHECK-NEXT: detach within %syncreg, label %for.body.lr.ph.us, label %pfor.inc.us
+
+; CHECK: pfor.inc.us:
+; CHECK-NEXT: ret void
+
+; CHECK: for.body.lr.ph.us:
+; CHECK-NEXT: call i32 null(ptr null, i32 0, i32 0, ptr %vla)
+; CHECK-NEXT: reattach within %syncreg, label %pfor.inc.us
+
+attributes #0 = { nounwind willreturn memory(argmem: readwrite) }
+attributes #1 = { sanitize_address }