From 82e946227d9ebfcb1bc5bda98af357879144d776 Mon Sep 17 00:00:00 2001
From: Vedant Kumar <vsk@apple.com>
Date: Thu, 16 Aug 2018 22:24:47 +0000
Subject: [PATCH] [InstrProf] Use atomic profile counter updates for TSan

Thread sanitizer instrumentation fails to skip all loads and stores to
profile counters. This can happen if profile counter updates are merged:

  %.sink = phi i64* ...
  %pgocount5 = load i64, i64* %.sink
  %27 = add i64 %pgocount5, 1
  %28 = bitcast i64* %.sink to i8*
  call void @__tsan_write8(i8* %28)
  store i64 %27, i64* %.sink

To suppress TSan diagnostics about racy counter updates, make the
counter updates atomic when TSan is enabled. If there's general interest
in this mode it can be surfaced as a clang/swift driver option.

Testing: check-{llvm,clang,profile}

rdar://40477803

Differential Revision: https://reviews.llvm.org/D50867

llvm-svn: 339955
---
 include/llvm/Transforms/Instrumentation.h     |  3 +++
 .../Instrumentation/InstrProfiling.cpp        | 22 ++++++++++++++-----
 .../InstrProfiling/atomic-updates.ll          | 14 ++++++++++++
 3 files changed, 33 insertions(+), 6 deletions(-)
 create mode 100644 test/Instrumentation/InstrProfiling/atomic-updates.ll

diff --git a/include/llvm/Transforms/Instrumentation.h b/include/llvm/Transforms/Instrumentation.h
index 4a346c8d745..26d2bde8ed3 100644
--- a/include/llvm/Transforms/Instrumentation.h
+++ b/include/llvm/Transforms/Instrumentation.h
@@ -111,6 +111,9 @@ struct InstrProfOptions {
   // Do counter register promotion
   bool DoCounterPromotion = false;
 
+  // Use atomic profile counter increments.
+  bool Atomic = false;
+
   // Name of the profile file to use as output
   std::string InstrProfileOutput;
 
diff --git a/lib/Transforms/Instrumentation/InstrProfiling.cpp b/lib/Transforms/Instrumentation/InstrProfiling.cpp
index 4d5dfb0aa66..62da9300253 100644
--- a/lib/Transforms/Instrumentation/InstrProfiling.cpp
+++ b/lib/Transforms/Instrumentation/InstrProfiling.cpp
@@ -96,6 +96,11 @@ cl::opt<double> NumCountersPerValueSite(
     // is usually smaller than 2.
     cl::init(1.0));
 
+cl::opt<bool> AtomicCounterUpdateAll(
+    "instrprof-atomic-counter-update-all", cl::ZeroOrMore,
+    cl::desc("Make all profile counter updates atomic (for testing only)"),
+    cl::init(false));
+
 cl::opt<bool> AtomicCounterUpdatePromoted(
     "atomic-counter-update-promoted", cl::ZeroOrMore,
     cl::desc("Do counter update using atomic fetch add "
@@ -597,12 +602,17 @@ void InstrProfiling::lowerIncrement(InstrProfIncrementInst *Inc) {
   IRBuilder<> Builder(Inc);
   uint64_t Index = Inc->getIndex()->getZExtValue();
   Value *Addr = Builder.CreateConstInBoundsGEP2_64(Counters, 0, Index);
-  Value *Load = Builder.CreateLoad(Addr, "pgocount");
-  auto *Count = Builder.CreateAdd(Load, Inc->getStep());
-  auto *Store = Builder.CreateStore(Count, Addr);
-  Inc->replaceAllUsesWith(Store);
-  if (isCounterPromotionEnabled())
-    PromotionCandidates.emplace_back(cast<Instruction>(Load), Store);
+
+  if (Options.Atomic || AtomicCounterUpdateAll) {
+    Builder.CreateAtomicRMW(AtomicRMWInst::Add, Addr, Inc->getStep(),
+                            AtomicOrdering::Monotonic);
+  } else {
+    Value *Load = Builder.CreateLoad(Addr, "pgocount");
+    auto *Count = Builder.CreateAdd(Load, Inc->getStep());
+    auto *Store = Builder.CreateStore(Count, Addr);
+    if (isCounterPromotionEnabled())
+      PromotionCandidates.emplace_back(cast<Instruction>(Load), Store);
+  }
   Inc->eraseFromParent();
 }
 
diff --git a/test/Instrumentation/InstrProfiling/atomic-updates.ll b/test/Instrumentation/InstrProfiling/atomic-updates.ll
new file mode 100644
index 00000000000..7230a205dff
--- /dev/null
+++ b/test/Instrumentation/InstrProfiling/atomic-updates.ll
@@ -0,0 +1,14 @@
+; RUN: opt < %s -S -instrprof -instrprof-atomic-counter-update-all | FileCheck %s
+
+target triple = "x86_64-apple-macosx10.10.0"
+
+@__profn_foo = hidden constant [3 x i8] c"foo"
+
+; CHECK-LABEL: define void @foo
+; CHECK-NEXT: atomicrmw add i64* getelementptr inbounds ([1 x i64], [1 x i64]* @__profc_foo, i64 0, i64 0), i64 1 monotonic
+define void @foo() {
+  call void @llvm.instrprof.increment(i8* getelementptr inbounds ([3 x i8], [3 x i8]* @__profn_foo, i32 0, i32 0), i64 0, i32 1, i32 0)
+  ret void
+}
+
+declare void @llvm.instrprof.increment(i8*, i64, i32, i32)