[llvm] bf6986f - [TSan, SanitizerBinaryMetadata] Improve instrument for derived pointers via phis/selects (#132752)

[via llvm-commits]

Author: Camsyn
Date: 2025-04-17T10:09:07+02:00
New Revision: bf6986f9f09f79da38006a83c339226c429bb686

URL: https://github.com/llvm/llvm-project/commit/bf6986f9f09f79da38006a83c339226c429bb686
DIFF: https://github.com/llvm/llvm-project/commit/bf6986f9f09f79da38006a83c339226c429bb686.diff

LOG: [TSan, SanitizerBinaryMetadata] Improve instrument for derived pointers via phis/selects (#132752)

ThreadSanitizer.cpp and SanitizerBinaryMetadata.cpp previously used
`getUnderlyingObject` to check if pointers originate from stack objects.

However, `getUnderlyingObject()` by default only looks through linear
chains, not selects/phis. In particular, this means that we miss cases
involving pointer induction variables.

For instance,
```llvm
%stkobj = alloca [2 x i32], align 8
; getUnderlyingObject(%derived) = %derived
%derived = getelementptr inbounds i32, ptr %stkobj, i64 1
```

This will result in redundant instrumentation of TSan, resulting in
greater performance costs, especially when there are loops, referring to
this [godbolt page](https://godbolt.org/z/eaT1fPjTW) for details.
```cpp
char loop(int x) {
    char buf[10];
    char *p = buf;
    for (int i = 0; i < x && i < 10; i++) {
      // Should not instrument, as its base object is a non-captured stack
      // variable.
      // However, currectly, it is instrumented due to %p = %phi ...
      *p++ = i;
    }

    // Use buf to prevent it from being eliminated by optimization
    return buf[9];
}
```

There are TWO APIs `getUnderlyingObjectAggressive` and
`findAllocaForValue` that can backtrack the pointer via tree traversal,
supporting phis/selects.

This patch replaces `getUnderlyingObject` with `findAllocaForValue`
which:
1. Properly tracks through PHINodes and select operations
2. Directly identifies if a pointer comes from a `AllocaInst`

Performance impact:
- Compilation: Moderate cost increase due to wider value tracing, but...
- Runtime: Significant wins for code with pointer induction variables
derived from stack allocas, especially for loop-heavy code, as
instrumentation can now be safely omitted.

Added: 
    

Modified: 
    llvm/lib/Transforms/Instrumentation/SanitizerBinaryMetadata.cpp
    llvm/lib/Transforms/Instrumentation/ThreadSanitizer.cpp
    llvm/test/Instrumentation/ThreadSanitizer/capture.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Instrumentation/SanitizerBinaryMetadata.cpp b/llvm/lib/Transforms/Instrumentation/SanitizerBinaryMetadata.cpp
index ab5cdc08664d7..b6787fd8cb21e 100644
--- a/llvm/lib/Transforms/Instrumentation/SanitizerBinaryMetadata.cpp
+++ b/llvm/lib/Transforms/Instrumentation/SanitizerBinaryMetadata.cpp
@@ -392,8 +392,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(Addr, /*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 1811d145f9907..7f846d25b8133 100644
--- a/llvm/lib/Transforms/Instrumentation/ThreadSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/ThreadSanitizer.cpp
@@ -448,8 +448,8 @@ void ThreadSanitizer::chooseInstructionsToInstrument(
       }
     }
 
-    if (isa<AllocaInst>(getUnderlyingObject(Addr)) &&
-        !PointerMayBeCaptured(Addr, /*ReturnCaptures=*/true)) {
+    const AllocaInst *AI = findAllocaForValue(Addr);
+    if (AI && !PointerMayBeCaptured(Addr, /*ReturnCaptures=*/true)) {
       // The variable is addressable but not captured, so it cannot be
       // referenced from a 
diff erent 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)