From b347a720bfc4dfd6ff16cfde54f9a56f00efdb3c Mon Sep 17 00:00:00 2001
From: NAKAMURA Takumi <geek4civic@gmail.com>
Date: Wed, 26 Jun 2024 12:07:56 +0900
Subject: [PATCH] [MC/DC][Coverage] Make tvbitmapupdate capable of atomic write
 (#96042)

This also introduces "Test and conditional Read-Modify-Write". The flow
to `atomicrmw or` is marked as `unlikely`.
---
 .../Instrumentation/InstrProfiling.cpp        | 24 +++++++++++++++++++
 .../Instrumentation/InstrProfiling/mcdc.ll    | 11 +++++++++
 2 files changed, 35 insertions(+)

diff --git a/llvm/lib/Transforms/Instrumentation/InstrProfiling.cpp b/llvm/lib/Transforms/Instrumentation/InstrProfiling.cpp
index 505ca9f7256f..f994f8a62c32 100644
--- a/llvm/lib/Transforms/Instrumentation/InstrProfiling.cpp
+++ b/llvm/lib/Transforms/Instrumentation/InstrProfiling.cpp
@@ -38,6 +38,7 @@
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Type.h"
 #include "llvm/InitializePasses.h"
@@ -967,6 +968,29 @@ Function *InstrLowerer::createRMWOrFunc() {
   //  %mcdc.bits = load i8, ptr %4, align 1
   auto *Bitmap = Builder.CreateLoad(Int8Ty, ArgAddr, "mcdc.bits");
 
+  if (Options.Atomic || AtomicCounterUpdateAll) {
+    // If ((Bitmap & Val) != Val), then execute atomic (Bitmap |= Val).
+    // Note, just-loaded Bitmap might not be up-to-date. Use it just for
+    // early testing.
+    auto *Masked = Builder.CreateAnd(Bitmap, ArgVal);
+    auto *ShouldStore = Builder.CreateICmpNE(Masked, ArgVal);
+    auto *ThenTerm = BasicBlock::Create(Ctx, "", Fn);
+    auto *ElseTerm = BasicBlock::Create(Ctx, "", Fn);
+    // Assume updating will be rare.
+    auto *Unlikely = MDBuilder(Ctx).createUnlikelyBranchWeights();
+    Builder.CreateCondBr(ShouldStore, ThenTerm, ElseTerm, Unlikely);
+
+    IRBuilder<> ThenBuilder(ThenTerm);
+    ThenBuilder.CreateAtomicRMW(AtomicRMWInst::Or, ArgAddr, ArgVal,
+                                MaybeAlign(), AtomicOrdering::Monotonic);
+    ThenBuilder.CreateRetVoid();
+
+    IRBuilder<> ElseBuilder(ElseTerm);
+    ElseBuilder.CreateRetVoid();
+
+    return Fn;
+  }
+
   // Perform logical OR of profile bitmap byte and shifted bit offset.
   //  %8 = or i8 %mcdc.bits, %7
   auto *Result = Builder.CreateOr(Bitmap, ArgVal);
diff --git a/llvm/test/Instrumentation/InstrProfiling/mcdc.ll b/llvm/test/Instrumentation/InstrProfiling/mcdc.ll
index c7732a4781b6..20c002e0565d 100644
--- a/llvm/test/Instrumentation/InstrProfiling/mcdc.ll
+++ b/llvm/test/Instrumentation/InstrProfiling/mcdc.ll
@@ -1,5 +1,6 @@
 ; Check that MC/DC intrinsics are properly lowered
 ; RUN: opt < %s -passes=instrprof -S | FileCheck %s --check-prefixes=CHECK,BASIC
+; RUN: opt < %s -passes=instrprof -S -instrprof-atomic-counter-update-all | FileCheck %s --check-prefixes=CHECK,ATOMIC
 ; RUN: opt < %s -passes=instrprof -runtime-counter-relocation -S 2>&1 | FileCheck %s --check-prefix RELOC
 
 ; RELOC: Runtime counter relocation is presently not supported for MC/DC bitmaps
@@ -9,6 +10,7 @@ target triple = "x86_64-unknown-linux-gnu"
 @__profn_test = private constant [4 x i8] c"test"
 
 ; BASIC: [[PROFBM_ADDR:@__profbm_test]] = private global [1 x i8] zeroinitializer, section "__llvm_prf_bits", comdat, align 1
+; ATOMIC: [[PROFBM_ADDR:@__profbm_test]] = private global [1 x i8] zeroinitializer, section "__llvm_prf_bits", comdat, align 1
 
 define dso_local void @test(i32 noundef %A) {
 entry:
@@ -38,8 +40,17 @@ entry:
 ; CHECK:      %[[BITS:.+]] = load i8, ptr %[[ARGPTR]], align 1
 ; BASIC-NEXT: %[[LAB11:[0-9]+]] = or i8 %[[BITS]], %[[ARGVAL]]
 ; BASIC-NEXT: store i8 %[[LAB11]], ptr %[[ARGPTR]], align 1
+; ATOMIC-NEXT: %[[MASKED:.+]] = and i8 %[[BITS]], %[[ARGVAL]]
+; ATOMIC-NEXT: %[[SHOULDWRITE:.+]] = icmp ne i8 %[[MASKED]], %[[ARGVAL]]
+; ATOMIC-NEXT: br i1 %[[SHOULDWRITE]], label %[[WRITE:.+]], label %[[SKIP:.+]], !prof ![[MDPROF:[0-9]+]]
+; ATOMIC: [[WRITE]]:
+; ATOMIC-NEXT: %{{.+}} = atomicrmw or ptr %[[ARGPTR]], i8 %[[ARGVAL]] monotonic, align 1
+; ATOMIC-NEXT: ret void
+; ATOMIC: [[SKIP]]:
 ; CHECK-NEXT: ret void
 
+; ATOMIC: ![[MDPROF]] = !{!"branch_weights", i32 1, i32 1048575}
+
 declare void @llvm.instrprof.cover(ptr, i64, i32, i32)
 
 declare void @llvm.instrprof.mcdc.parameters(ptr, i64, i32)