[llvm] [InstrProf] Support conditional counter updates (PR #102542)

[via llvm-commits]

https://github.com/gulfemsavrun created https://github.com/llvm/llvm-project/pull/102542

This patch adds support for conditional counter updates in single byte counters mode to reduce the write contention by first checking whether the counter is set before overwriting it.

>From e8fd07c2af05c2e5a78f3d2e8f9d09a11f0a8914 Mon Sep 17 00:00:00 2001
From: Gulfem Savrun Yeniceri <gulfem at google.com>
Date: Thu, 8 Aug 2024 21:36:55 +0000
Subject: [PATCH] [InstrProf] Support conditional counter updates

This patch adds support of conditional counter updates
in single byte counters mode to reduce the write contention
by first checking whether the counter is set before
overwriting it.
---
 .../Instrumentation/InstrProfiling.cpp           | 16 ++++++++++++++--
 .../conditional-counter-updates.ll               | 15 +++++++++++++++
 2 files changed, 29 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/Instrumentation/InstrProfiling/conditional-counter-updates.ll

diff --git a/llvm/lib/Transforms/Instrumentation/InstrProfiling.cpp b/llvm/lib/Transforms/Instrumentation/InstrProfiling.cpp
index 1805ea89272ec7..5cf2096657547d 100644
--- a/llvm/lib/Transforms/Instrumentation/InstrProfiling.cpp
+++ b/llvm/lib/Transforms/Instrumentation/InstrProfiling.cpp
@@ -1213,8 +1213,20 @@ Value *InstrLowerer::getBitmapAddress(InstrProfMCDCTVBitmapUpdate *I) {
 void InstrLowerer::lowerCover(InstrProfCoverInst *CoverInstruction) {
   auto *Addr = getCounterAddress(CoverInstruction);
   IRBuilder<> Builder(CoverInstruction);
-  // We store zero to represent that this block is covered.
-  Builder.CreateStore(Builder.getInt8(0), Addr);
+  if (ConditionalCounterUpdate) {
+    auto &Ctx = CoverInstruction->getParent()->getContext();
+    auto *Int8Ty = llvm::Type::getInt8Ty(Ctx);
+    Value *Load = Builder.CreateLoad(Int8Ty, Addr, "pgocount");
+    Value *Cmp = Builder.CreateICmpNE(Load, ConstantInt::get(Int8Ty, 0),
+                                      "pgocount.ifnonzero");
+    Value *Sel =
+        Builder.CreateSelect(Cmp, Builder.getInt8(0), Load, "pgocount.select");
+    Builder.CreateStore(Sel, Addr);
+  } else {
+    // We store zero to represent that this block is covered.
+    Builder.CreateStore(Builder.getInt8(0), Addr);
+  }
+
   CoverInstruction->eraseFromParent();
 }
 
diff --git a/llvm/test/Instrumentation/InstrProfiling/conditional-counter-updates.ll b/llvm/test/Instrumentation/InstrProfiling/conditional-counter-updates.ll
new file mode 100644
index 00000000000000..704a0f3e64b101
--- /dev/null
+++ b/llvm/test/Instrumentation/InstrProfiling/conditional-counter-updates.ll
@@ -0,0 +1,15 @@
+; RUN: opt < %s -S -passes=instrprof -conditional-counter-update | FileCheck %s
+
+target triple = "x86_64-unknown-linux-gnu"
+
+ at __profn_foo = private constant [3 x i8] c"foo"
+
+; CHECK-LABEL: define void @foo
+; CHECK-NEXT: %pgocount = load i8, ptr @__profc_foo, align 1
+; CHECK-NEXT: %pgocount.ifnonzero = icmp ne i8 %pgocount, 0
+; CHECK-NEXT: %pgocount.select = select i1 %pgocount.ifnonzero, i8 0, i8 %pgocount
+; CHECK-NEXT: store i8 %pgocount.select, ptr @__profc_foo, align 1
+define void @foo() {
+  call void @llvm.instrprof.cover(ptr @__profn_foo, i64 0, i32 1, i32 0)
+  ret void
+}