[compiler-rt] [llvm] [TSan, SanitizerBinaryMetadata] Analyze the capture status for `alloca` rather than arbitrary `Addr` (PR #132756)

[via llvm-commits]

https://github.com/Camsyn updated https://github.com/llvm/llvm-project/pull/132756

>From 1a3518a2bf64a63cbb1fc5f16441902817aa5a24 Mon Sep 17 00:00:00 2001
From: Camsyn <camsyn at foxmail.com>
Date: Thu, 24 Apr 2025 15:14:14 +0800
Subject: [PATCH] [TSan, SanitizerBinaryMetadata] Analyze the capture status
 for `alloca` rather than arbitrary `Addr`

This commit addresses the limitation in `PointerMayBeCaptured` analysis
when dealing with derived pointers (e.g. arr+1) as described in
issue #132739.

The current implementation may miss captures of the underlying alloca
when analyzing derived pointers. Even when working correctly,
backtracking to the original alloca during analysis causes redundancy
to the outer's `findAllocaForValue`.

Key changes:
Directly analyze the capture status of the underlying alloca instead
   of derived pointers to ensure accurate capture detection

This patch fixes some FNs of TSan, referring to the appending
testcases for more details.
---
 compiler-rt/test/tsan/stack_race3.cpp         | 21 +++++++++++++++++++
 .../SanitizerBinaryMetadata.cpp               |  2 +-
 .../Instrumentation/ThreadSanitizer.cpp       |  3 ++-
 .../ThreadSanitizer/capture.ll                | 15 +++++++++++++
 4 files changed, 39 insertions(+), 2 deletions(-)
 create mode 100644 compiler-rt/test/tsan/stack_race3.cpp

diff --git a/compiler-rt/test/tsan/stack_race3.cpp b/compiler-rt/test/tsan/stack_race3.cpp
new file mode 100644
index 0000000000000..c50f4ad4c35cc
--- /dev/null
+++ b/compiler-rt/test/tsan/stack_race3.cpp
@@ -0,0 +1,21 @@
+// RUN: %clangxx_tsan -O1 %s -o %t && %deflake %run %t | FileCheck %s
+#include "test.h"
+
+void *Thread(void *a) {
+  barrier_wait(&barrier);
+  ((int *)a)[1] = 43;
+  return 0;
+}
+
+int main() {
+  barrier_init(&barrier, 2);
+  int Arr[2] = {41, 42};
+  pthread_t t;
+  pthread_create(&t, 0, Thread, &Arr[0]);
+  Arr[1] = 43;
+  barrier_wait(&barrier);
+  pthread_join(t, 0);
+}
+
+// CHECK: WARNING: ThreadSanitizer: data race
+// CHECK:   Location is stack of main thread.
diff --git a/llvm/lib/Transforms/Instrumentation/SanitizerBinaryMetadata.cpp b/llvm/lib/Transforms/Instrumentation/SanitizerBinaryMetadata.cpp
index b6787fd8cb21e..4801ac75f8572 100644
--- a/llvm/lib/Transforms/Instrumentation/SanitizerBinaryMetadata.cpp
+++ b/llvm/lib/Transforms/Instrumentation/SanitizerBinaryMetadata.cpp
@@ -393,7 +393,7 @@ bool maybeSharedMutable(const Value *Addr) {
     return true;
 
   const AllocaInst *AI = findAllocaForValue(Addr);
-  if (AI && !PointerMayBeCaptured(Addr, /*ReturnCaptures=*/true))
+  if (AI && !PointerMayBeCaptured(AI, /*ReturnCaptures=*/true))
     return false; // Object is on stack but does not escape.
 
   Addr = Addr->stripInBoundsOffsets();
diff --git a/llvm/lib/Transforms/Instrumentation/ThreadSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/ThreadSanitizer.cpp
index 7f846d25b8133..ec9f78edfeb1c 100644
--- a/llvm/lib/Transforms/Instrumentation/ThreadSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/ThreadSanitizer.cpp
@@ -449,7 +449,8 @@ void ThreadSanitizer::chooseInstructionsToInstrument(
     }
 
     const AllocaInst *AI = findAllocaForValue(Addr);
-    if (AI && !PointerMayBeCaptured(Addr, /*ReturnCaptures=*/true)) {
+    // Instead of Addr, we should check whether its base pointer is captured.
+    if (AI && !PointerMayBeCaptured(AI, /*ReturnCaptures=*/true)) {
       // The variable is addressable but not captured, so it cannot be
       // referenced from a different thread and participate in a data race
       // (see llvm/Analysis/CaptureTracking.h for details).
diff --git a/llvm/test/Instrumentation/ThreadSanitizer/capture.ll b/llvm/test/Instrumentation/ThreadSanitizer/capture.ll
index 9cd5d77f4753e..e1b9e03b88446 100644
--- a/llvm/test/Instrumentation/ThreadSanitizer/capture.ll
+++ b/llvm/test/Instrumentation/ThreadSanitizer/capture.ll
@@ -47,6 +47,21 @@ entry:
 ; CHECK: __tsan_write
 ; CHECK: ret void
 
+define void @captured3() nounwind uwtable sanitize_thread {
+entry:
+  %stkobj = alloca [2 x i32], align 8
+  ; escapes due to store into global
+  store ptr %stkobj, ptr @sink, align 8
+  ; derived is captured as its base object is captured
+  %derived = getelementptr inbounds i32, ptr %stkobj, i64 1
+  store i32 42, ptr %derived, align 4
+  ret void
+}
+; CHECK-LABEL: define void @captured3
+; CHECK: __tsan_write
+; CHECK: __tsan_write
+; CHECK: ret void
+
 define void @notcaptured0() nounwind uwtable sanitize_thread {
 entry:
   %ptr = alloca i32, align 4