[llvm] [offload] Do not pool memory while allocation traces are requested (PR #214752)

[Robert Imschweiler via llvm-commits]

https://github.com/ro-i updated https://github.com/llvm/llvm-project/pull/214752

>From 022265c76bd74cc16714960cdb91961fd6fb42a7 Mon Sep 17 00:00:00 2001
From: Robert Imschweiler <robert.imschweiler at amd.com>
Date: Wed, 5 Aug 2026 10:24:00 -0500
Subject: [PATCH 1/2] [offload] Do not pool memory while allocation traces are
 requested

Would otherwise hide use-after-free because memory stays valid if it's
in the pool.
---
 offload/plugins-nextgen/common/src/PluginInterface.cpp | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/offload/plugins-nextgen/common/src/PluginInterface.cpp b/offload/plugins-nextgen/common/src/PluginInterface.cpp
index e539cbfc55324..d4f3d8b62a272 100644
--- a/offload/plugins-nextgen/common/src/PluginInterface.cpp
+++ b/offload/plugins-nextgen/common/src/PluginInterface.cpp
@@ -598,9 +598,13 @@ Error GenericDeviceTy::init(GenericPluginTy &Plugin) {
     GridValues.GV_Max_WG_Size =
         std::min(GridValues.GV_Max_WG_Size, uint32_t(OMP_TeamsThreadLimit));
 
-  // Enable the memory manager if required.
+  // Enable the memory manager if required. Memory that the manager pools is
+  // not returned to the driver when it is freed, so a later access to it
+  // neither faults nor is reported by the allocation tracker. Leave the pool
+  // disabled while allocation traces are requested, so that we don't mask
+  // use-after-free (since they don't fault if the memory is still in the pool).
   auto [ThresholdMM, EnableMM] = MemoryManagerTy::getSizeThresholdFromEnv();
-  if (EnableMM) {
+  if (EnableMM && !OMPX_TrackAllocationTraces) {
     if (ThresholdMM == 0)
       ThresholdMM = getMemoryManagerSizeThreshold();
     MemoryManager = new MemoryManagerTy(*this, ThresholdMM);

>From f20a75434463f07a0022165f31fb86453472e9a5 Mon Sep 17 00:00:00 2001
From: Robert Imschweiler <robert.imschweiler at amd.com>
Date: Fri, 7 Aug 2026 09:52:10 -0500
Subject: [PATCH 2/2] reduce comment

---
 offload/plugins-nextgen/common/src/PluginInterface.cpp | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/offload/plugins-nextgen/common/src/PluginInterface.cpp b/offload/plugins-nextgen/common/src/PluginInterface.cpp
index d4f3d8b62a272..70e289060a446 100644
--- a/offload/plugins-nextgen/common/src/PluginInterface.cpp
+++ b/offload/plugins-nextgen/common/src/PluginInterface.cpp
@@ -598,11 +598,9 @@ Error GenericDeviceTy::init(GenericPluginTy &Plugin) {
     GridValues.GV_Max_WG_Size =
         std::min(GridValues.GV_Max_WG_Size, uint32_t(OMP_TeamsThreadLimit));
 
-  // Enable the memory manager if required. Memory that the manager pools is
-  // not returned to the driver when it is freed, so a later access to it
-  // neither faults nor is reported by the allocation tracker. Leave the pool
-  // disabled while allocation traces are requested, so that we don't mask
-  // use-after-free (since they don't fault if the memory is still in the pool).
+  // Enable the memory manager if required. Leave the pool disabled while
+  // allocation traces are requested, so that we don't mask use-after-free
+  // (since they don't fault if the memory is still in the pool).
   auto [ThresholdMM, EnableMM] = MemoryManagerTy::getSizeThresholdFromEnv();
   if (EnableMM && !OMPX_TrackAllocationTraces) {
     if (ThresholdMM == 0)