[llvm] [MemProf] Improve metadata cleanup in LTO backend (PR #113039)

[Teresa Johnson via llvm-commits]

https://github.com/teresajohnson created https://github.com/llvm/llvm-project/pull/113039

Previously we were attempting to remove the memprof-related metadata
when iterating through instructions in the LTO backend. However, we
missed some as there are a number of cases where we skip instructions,
or even entire functions. Simplify the cleanup and ensure all is removed
by doing a full sweep over all instructions after completing cloning.

This is largely NFC except with -memprof-report-hinted-sizes enabled,
because we were propagating and simplifying the metadata after inlining
in the LTO backend, which caused some stray messages as metadata was
re-converted to attributes.


>From a5cbba5239d615e442ab7448f824476619dcab49 Mon Sep 17 00:00:00 2001
From: Teresa Johnson <tejohnson at google.com>
Date: Sat, 19 Oct 2024 06:46:55 -0700
Subject: [PATCH] [MemProf] Improve metadata cleanup in LTO backend

Previously we were attempting to remove the memprof-related metadata
when iterating through instructions in the LTO backend. However, we
missed some as there are a number of cases where we skip instructions,
or even entire functions. Simplify the cleanup and ensure all is removed
by doing a full sweep over all instructions after completing cloning.

This is largely NFC except with -memprof-report-hinted-sizes enabled,
because we were propagating and simplifying the metadata after inlining
in the LTO backend, which caused some stray messages as metadata was
re-converted to attributes.
---
 .../IPO/MemProfContextDisambiguation.cpp      | 22 ++++++++++++++-----
 llvm/test/ThinLTO/X86/memprof-icp.ll          |  5 ++++-
 2 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp b/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp
index 5ade0db343f278..17f29dc0d580d3 100644
--- a/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp
+++ b/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp
@@ -4264,9 +4264,6 @@ bool MemProfContextDisambiguation::applyImport(Module &M) {
           AllocVersionsThinBackend++;
           if (!MaxAllocVersionsThinBackend)
             MaxAllocVersionsThinBackend = 1;
-          // Remove any remaining callsite metadata and we can skip the rest of
-          // the handling for this instruction, since no cloning needed.
-          I.setMetadata(LLVMContext::MD_callsite, nullptr);
           continue;
         }
 
@@ -4419,9 +4416,6 @@ bool MemProfContextDisambiguation::applyImport(Module &M) {
             CloneCallsite(Callsite->second, CB, CalledFunction);
           }
         }
-        // Memprof and callsite metadata on memory allocations no longer needed.
-        I.setMetadata(LLVMContext::MD_memprof, nullptr);
-        I.setMetadata(LLVMContext::MD_callsite, nullptr);
       }
     }
 
@@ -4429,6 +4423,22 @@ bool MemProfContextDisambiguation::applyImport(Module &M) {
     performICP(M, FS->callsites(), VMaps, ICallAnalysisInfo, ORE);
   }
 
+  // We skip some of the functions and instructions above, so remove all the
+  // metadata in a single sweep here.
+  for (auto &F : M) {
+    // We can skip memprof clones because createFunctionClones already strips
+    // the metadata from the newly created clones.
+    if (F.isDeclaration() || isMemProfClone(F))
+      continue;
+    for (auto &BB : F) {
+      for (auto &I : BB) {
+        // Memprof and callsite metadata on memory allocations no longer needed.
+        I.setMetadata(LLVMContext::MD_memprof, nullptr);
+        I.setMetadata(LLVMContext::MD_callsite, nullptr);
+      }
+    }
+  }
+
   return Changed;
 }
 
diff --git a/llvm/test/ThinLTO/X86/memprof-icp.ll b/llvm/test/ThinLTO/X86/memprof-icp.ll
index 2e976794425bbe..f17e19e1f77ef2 100644
--- a/llvm/test/ThinLTO/X86/memprof-icp.ll
+++ b/llvm/test/ThinLTO/X86/memprof-icp.ll
@@ -176,7 +176,10 @@
 ; RUN:  -pass-remarks=. -save-temps \
 ; RUN:  -o %t.noicp.out 2>&1 | FileCheck %s --implicit-check-not "created clone"
 
-; RUN: llvm-dis %t.noicp.out.2.4.opt.bc -o - | FileCheck %s --implicit-check-not "_Z3fooR2B0j.memprof"
+;; Verify that we did not do any cloning of the function with the indirect call
+;; when memprof ICP is off. However, we should still have removed the callsite
+;; metadata.
+; RUN: llvm-dis %t.noicp.out.2.4.opt.bc -o - | FileCheck %s --implicit-check-not "_Z3fooR2B0j.memprof" --implicit-check-not "!callsite"
 
 ; REMARKS-MAIN: call in clone main assigned to call function clone _Z3fooR2B0j.memprof.1
 ; REMARKS-MAIN: call in clone main assigned to call function clone _Z3fooR2B0j.memprof.1