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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-compiler-rt-sanitizer

Author: None (Camsyn)

<details>
<summary>Changes</summary>

This PR is based on my last PR #<!-- -->132752 (the first commit of this PR), but addressing a different issue.

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 of `PointerMayBeCaptured` may miss captures of the underlying `alloca` when analyzing derived pointers, leading to some FNs in TSan, as follows:
```c
void *Thread(void *a) {
  ((int*)a)[1] = 43;
  return 0;
}

int main() {
  int Arr[2] = {41, 42};
  pthread_t t;
  pthread_create(&t, 0, Thread, &Arr[0]);
  // Missed instrumentation here due to the FN of PointerMayBeCaptured
  Arr[1] = 43;
  barrier_wait(&barrier);
  pthread_join(t, 0);
}
```
Refer to this [godbolt page](https://godbolt.org/z/n67GrxdcE) to get the compilation result of TSan.

Even when `PointerMayBeCaptured` working correctly, it should backtrack to the original `alloca` firstly during analysis, causing redundancy to the outer's `findAllocaForValue`.
```cpp
    const AllocaInst *AI = findAllocaForValue(Addr);
    // Instead of Addr, we should check whether its base pointer is captured.
    if (AI && !PointerMayBeCaptured(Addr, true)) ...
```

Key changes:
Directly analyze the capture status of the underlying `alloca` instead of derived pointers to ensure accurate capture detection
```cpp
    const AllocaInst *AI = findAllocaForValue(Addr);
    // Instead of Addr, we should check whether its base pointer is captured.
    if (AI && !PointerMayBeCaptured(Addr, AI)) ...
```


---
Full diff: https://github.com/llvm/llvm-project/pull/132756.diff


4 Files Affected:

- (added) compiler-rt/test/tsan/stack_race3.cpp (+22) 
- (modified) llvm/lib/Transforms/Instrumentation/SanitizerBinaryMetadata.cpp (+2-2) 
- (modified) llvm/lib/Transforms/Instrumentation/ThreadSanitizer.cpp (+3-2) 
- (modified) llvm/test/Instrumentation/ThreadSanitizer/capture.ll (+31) 


``````````diff
diff --git a/compiler-rt/test/tsan/stack_race3.cpp b/compiler-rt/test/tsan/stack_race3.cpp
new file mode 100644
index 0000000000000..5078d81f3cc58
--- /dev/null
+++ b/compiler-rt/test/tsan/stack_race3.cpp
@@ -0,0 +1,22 @@
+// 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 c5932f0d65ee1..f22a87421319b 100644
--- a/llvm/lib/Transforms/Instrumentation/SanitizerBinaryMetadata.cpp
+++ b/llvm/lib/Transforms/Instrumentation/SanitizerBinaryMetadata.cpp
@@ -393,8 +393,8 @@ bool maybeSharedMutable(const Value *Addr) {
   if (!Addr)
     return true;
 
-  if (isa<AllocaInst>(getUnderlyingObject(Addr)) &&
-      !PointerMayBeCaptured(Addr, /*ReturnCaptures=*/true))
+  const AllocaInst *AI = findAllocaForValue(Addr);
+  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 2b403b695c1d2..7c5fc40d51b78 100644
--- a/llvm/lib/Transforms/Instrumentation/ThreadSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/ThreadSanitizer.cpp
@@ -448,8 +448,9 @@ void ThreadSanitizer::chooseInstructionsToInstrument(
       }
     }
 
-    if (isa<AllocaInst>(getUnderlyingObject(Addr)) &&
-        !PointerMayBeCaptured(Addr, /*ReturnCaptures=*/true)) {
+    const AllocaInst *AI = findAllocaForValue(Addr);
+    // 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 8edf310df9823..9cd5d77f4753e 100644
--- a/llvm/test/Instrumentation/ThreadSanitizer/capture.ll
+++ b/llvm/test/Instrumentation/ThreadSanitizer/capture.ll
@@ -88,4 +88,35 @@ entry:
 ; CHECK: __tsan_write
 ; CHECK: ret void
 
+define void @notcaptured3(i1 %cond) nounwind uwtable sanitize_thread {
+entry:
+  %stkobj = alloca [2 x i32], align 8
+  %derived = getelementptr inbounds i32, ptr %stkobj, i64 1
+  %ptr = select i1 %cond, ptr %derived, ptr %stkobj
+  store i32 42, ptr %ptr, align 4
+  ret void
+}
+; CHECK-LABEL: define void @notcaptured3
+; CHECK-NOT: call void @__tsan_write4(ptr %ptr)
+; CHECK: ret void
 
+define void @notcaptured4() nounwind uwtable sanitize_thread {
+entry:
+  %stkobj = alloca [10 x i8], align 1
+  br label %loop
+
+exit:
+  ret void
+
+loop:
+  %count = phi i32 [ 0, %entry ], [ %addone, %loop ]
+  %derived = phi ptr [ %stkobj, %entry ], [ %ptraddone, %loop ]
+  store i32 %count, ptr %derived, align 4
+  %ptraddone = getelementptr inbounds i32, ptr %derived, i64 1
+  %addone = add nuw nsw i32 %count, 1
+  %eq10 = icmp eq i32 %addone, 10
+  br i1 %eq10, label %exit, label %loop
+}
+; CHECK-LABEL: define void @notcaptured4
+; CHECK: ret void
+; CHECK-NOT: call void @__tsan_write4(ptr %derived)

``````````

</details>


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