[llvm] [InstrProf] Fix bug when clearing traces with samples (PR #92310)

[Ellis Hoag via llvm-commits]

https://github.com/ellishg created https://github.com/llvm/llvm-project/pull/92310

The `--temporal-profile-max-trace-length=0` flag in the `llvm-profdata merge` command is used to remove traces from a profile. There was a bug where traces would not be cleared if the profile was already sampled. This patch fixes that.

>From dc2303877f194680051de4704570747dd64f9a7b Mon Sep 17 00:00:00 2001
From: Ellis Hoag <ellis.sparky.hoag at gmail.com>
Date: Wed, 15 May 2024 12:54:22 -0700
Subject: [PATCH] [InstrProf] Fix bug when clearing traces with samples

The `--temporal-profile-max-trace-length=0` flag in the `llvm-profdata merge` command is used to remove traces from a profile. There was a bug where traces would not be cleared if the profile was already sampled. This patch fixes that.
---
 llvm/lib/ProfileData/InstrProfWriter.cpp           | 11 ++++++-----
 llvm/test/tools/llvm-profdata/trace-limit.proftext |  6 +++++-
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/ProfileData/InstrProfWriter.cpp b/llvm/lib/ProfileData/InstrProfWriter.cpp
index b61c59aacc0f9..c941b9d89df38 100644
--- a/llvm/lib/ProfileData/InstrProfWriter.cpp
+++ b/llvm/lib/ProfileData/InstrProfWriter.cpp
@@ -320,11 +320,8 @@ void InstrProfWriter::addBinaryIds(ArrayRef<llvm::object::BuildID> BIs) {
 }
 
 void InstrProfWriter::addTemporalProfileTrace(TemporalProfTraceTy Trace) {
-  if (Trace.FunctionNameRefs.size() > MaxTemporalProfTraceLength)
-    Trace.FunctionNameRefs.resize(MaxTemporalProfTraceLength);
-  if (Trace.FunctionNameRefs.empty())
-    return;
-
+  assert(Trace.FunctionNameRefs.size() <= MaxTemporalProfTraceLength);
+  assert(!Trace.FunctionNameRefs.empty());
   if (TemporalProfTraceStreamSize < TemporalProfTraceReservoirSize) {
     // Simply append the trace if we have not yet hit our reservoir size limit.
     TemporalProfTraces.push_back(std::move(Trace));
@@ -341,6 +338,10 @@ void InstrProfWriter::addTemporalProfileTrace(TemporalProfTraceTy Trace) {
 
 void InstrProfWriter::addTemporalProfileTraces(
     SmallVectorImpl<TemporalProfTraceTy> &SrcTraces, uint64_t SrcStreamSize) {
+  for (auto &Trace : SrcTraces)
+    if (Trace.FunctionNameRefs.size() > MaxTemporalProfTraceLength)
+      Trace.FunctionNameRefs.resize(MaxTemporalProfTraceLength);
+  llvm::erase_if(SrcTraces, [](auto &T) { return T.FunctionNameRefs.empty(); });
   // Assume that the source has the same reservoir size as the destination to
   // avoid needing to record it in the indexed profile format.
   bool IsDestSampled =
diff --git a/llvm/test/tools/llvm-profdata/trace-limit.proftext b/llvm/test/tools/llvm-profdata/trace-limit.proftext
index cf6edd648b23b..e246ee890ba31 100644
--- a/llvm/test/tools/llvm-profdata/trace-limit.proftext
+++ b/llvm/test/tools/llvm-profdata/trace-limit.proftext
@@ -1,13 +1,17 @@
 # RUN: llvm-profdata merge --temporal-profile-max-trace-length=0 %s -o %t.profdata
 # RUN: llvm-profdata show --temporal-profile-traces %t.profdata | FileCheck %s --check-prefix=NONE
 
+# RUN: llvm-profdata merge --temporal-profile-trace-reservoir-size=2 %s %s %s %s -o %t.profdata
+# RUN: llvm-profdata merge --temporal-profile-trace-reservoir-size=2 --temporal-profile-max-trace-length=0 %t.profdata -o %t.profdata
+# RUN: llvm-profdata show --temporal-profile-traces %t.profdata | FileCheck %s --check-prefix=NONE
+
 # RUN: llvm-profdata merge --temporal-profile-max-trace-length=2 %s -o %t.profdata
 # RUN: llvm-profdata show --temporal-profile-traces %t.profdata | FileCheck %s --check-prefixes=CHECK,SOME
 
 # RUN: llvm-profdata merge --temporal-profile-max-trace-length=1000 %s -o %t.profdata
 # RUN: llvm-profdata show --temporal-profile-traces %t.profdata | FileCheck %s --check-prefixes=CHECK,ALL
 
-# NONE: Temporal Profile Traces (samples=0 seen=0):
+# NONE: Temporal Profile Traces (samples=0
 # CHECK: Temporal Profile Traces (samples=1 seen=1):
 # SOME:   Trace 0 (weight=1 count=2):
 # ALL:    Trace 0 (weight=1 count=3):