[AddressSanitizer] Analyze parallel-promotable allocas before ASan's …

…instrumentation might invalidate task analysis.

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;

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 }