From c9dfec9664c5ea30d1dea3ed4dbcb77968bdac00 Mon Sep 17 00:00:00 2001
From: Teresa Johnson <tejohnson@google.com>
Date: Fri, 12 Sep 2025 18:35:25 -0700
Subject: [PATCH] [MemProf] Optionally allow transformation of nobuiltin
 operator new

For cases where we can guarantee the application does not override
operator new.
---
 .../llvm/Transforms/Utils/SimplifyLibCalls.h  |  2 +-
 .../lib/Transforms/Utils/SimplifyLibCalls.cpp | 43 +++++++++----
 .../InstCombine/simplify-libcalls-new.ll      | 60 ++++++++++++-------
 3 files changed, 70 insertions(+), 35 deletions(-)

diff --git a/llvm/include/llvm/Transforms/Utils/SimplifyLibCalls.h b/llvm/include/llvm/Transforms/Utils/SimplifyLibCalls.h
index deb3d6c44ef09..4e7c97194cc59 100644
--- a/llvm/include/llvm/Transforms/Utils/SimplifyLibCalls.h
+++ b/llvm/include/llvm/Transforms/Utils/SimplifyLibCalls.h
@@ -189,7 +189,7 @@ class LibCallSimplifier {
   Value *optimizeMemSet(CallInst *CI, IRBuilderBase &B);
   Value *optimizeRealloc(CallInst *CI, IRBuilderBase &B);
   Value *optimizeNew(CallInst *CI, IRBuilderBase &B, LibFunc &Func);
-  Value *optimizeExistingHotColdNew(CallInst *CI, IRBuilderBase &B);
+  Value *maybeOptimizeNoBuiltinOperatorNew(CallInst *CI, IRBuilderBase &B);
   Value *optimizeWcslen(CallInst *CI, IRBuilderBase &B);
   Value *optimizeBCopy(CallInst *CI, IRBuilderBase &B);
 
diff --git a/llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp b/llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp
index 8acebbaa5458b..4a1565977b91c 100644
--- a/llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp
@@ -61,6 +61,9 @@ static cl::opt<bool> OptimizeExistingHotColdNew(
     "optimize-existing-hot-cold-new", cl::Hidden, cl::init(false),
     cl::desc(
         "Enable optimization of existing hot/cold operator new library calls"));
+static cl::opt<bool> OptimizeNoBuiltinHotColdNew(
+    "optimize-nobuiltin-hot-cold-new-new", cl::Hidden, cl::init(false),
+    cl::desc("Enable transformation of nobuiltin operator new library calls"));
 
 namespace {
 
@@ -1723,13 +1726,11 @@ Value *LibCallSimplifier::optimizeRealloc(CallInst *CI, IRBuilderBase &B) {
   return nullptr;
 }
 
-// Allow existing calls to operator new() that takes a __hot_cold_t parameter to
-// be updated with a compiler-determined hot cold hint value. This is used in
-// cases where the call is marked nobuiltin (because operator new called
-// explicitly) and therefore cannot be replaced with a different callee.
-Value *LibCallSimplifier::optimizeExistingHotColdNew(CallInst *CI,
-                                                     IRBuilderBase &B) {
-  if (!OptimizeHotColdNew || !OptimizeExistingHotColdNew)
+// Optionally allow optimization of nobuiltin calls to operator new and its
+// variants.
+Value *LibCallSimplifier::maybeOptimizeNoBuiltinOperatorNew(CallInst *CI,
+                                                            IRBuilderBase &B) {
+  if (!OptimizeHotColdNew)
     return nullptr;
   Function *Callee = CI->getCalledFunction();
   if (!Callee)
@@ -1738,6 +1739,22 @@ Value *LibCallSimplifier::optimizeExistingHotColdNew(CallInst *CI,
   if (!TLI->getLibFunc(*Callee, Func))
     return nullptr;
   switch (Func) {
+  case LibFunc_Znwm:
+  case LibFunc_ZnwmRKSt9nothrow_t:
+  case LibFunc_ZnwmSt11align_val_t:
+  case LibFunc_ZnwmSt11align_val_tRKSt9nothrow_t:
+  case LibFunc_Znam:
+  case LibFunc_ZnamRKSt9nothrow_t:
+  case LibFunc_ZnamSt11align_val_t:
+  case LibFunc_ZnamSt11align_val_tRKSt9nothrow_t:
+  case LibFunc_size_returning_new:
+  case LibFunc_size_returning_new_aligned:
+    // By default normal operator new calls (not already passing a hot_cold_t
+    // parameter) are not mutated if the call is not marked builtin. Optionally
+    // enable that in cases where it is known to be safe.
+    if (!OptimizeNoBuiltinHotColdNew)
+      return nullptr;
+    break;
   case LibFunc_Znwm12__hot_cold_t:
   case LibFunc_ZnwmRKSt9nothrow_t12__hot_cold_t:
   case LibFunc_ZnwmSt11align_val_t12__hot_cold_t:
@@ -1748,10 +1765,15 @@ Value *LibCallSimplifier::optimizeExistingHotColdNew(CallInst *CI,
   case LibFunc_ZnamSt11align_val_tRKSt9nothrow_t12__hot_cold_t:
   case LibFunc_size_returning_new_hot_cold:
   case LibFunc_size_returning_new_aligned_hot_cold:
-    return optimizeNew(CI, B, Func);
+    // If the nobuiltin call already passes a hot_cold_t parameter, allow update
+    // of that parameter when enabled.
+    if (!OptimizeExistingHotColdNew)
+      return nullptr;
+    break;
   default:
     return nullptr;
   }
+  return optimizeNew(CI, B, Func);
 }
 
 // When enabled, replace operator new() calls marked with a hot or cold memprof
@@ -4121,9 +4143,8 @@ Value *LibCallSimplifier::optimizeCall(CallInst *CI, IRBuilderBase &Builder) {
   //       we can all non-FP calls with the StrictFP attribute to be
   //       optimized.
   if (CI->isNoBuiltin()) {
-    // If this is an existing call to a hot cold operator new, we can update the
-    // hint parameter value, which doesn't change the callee.
-    return optimizeExistingHotColdNew(CI, Builder);
+    // Optionally update operator new calls.
+    return maybeOptimizeNoBuiltinOperatorNew(CI, Builder);
   }
 
   LibFunc Func;
diff --git a/llvm/test/Transforms/InstCombine/simplify-libcalls-new.ll b/llvm/test/Transforms/InstCombine/simplify-libcalls-new.ll
index 41db7f929dfdf..5a4fb04f5f2c0 100644
--- a/llvm/test/Transforms/InstCombine/simplify-libcalls-new.ll
+++ b/llvm/test/Transforms/InstCombine/simplify-libcalls-new.ll
@@ -6,14 +6,18 @@
 ; OFF-LABEL: @new_hot_cold()
 
 ;; First check with the default hint values (254 = -2, 128 = -128, 222 = -34).
-; RUN: opt < %s -passes=instcombine -optimize-hot-cold-new -S | FileCheck %s --check-prefix=HOTCOLD -DCOLD=1 -DHOT=-2 -DNOTCOLD=-128 -DAMBIG=-34 -DPREVHINTCOLD=7 -DPREVHINTNOTCOLD=7 -DPREVHINTHOT=7 -DPREVHINTAMBIG=7
+; RUN: opt < %s -passes=instcombine -optimize-hot-cold-new -S | FileCheck %s --check-prefixes=HOTCOLD,NOBUILTIN-OFF -DCOLD=1 -DHOT=-2 -DNOTCOLD=-128 -DAMBIG=-34 -DPREVHINTCOLD=7 -DPREVHINTNOTCOLD=7 -DPREVHINTHOT=7 -DPREVHINTAMBIG=7
 
 ;; Next check with the non-default cold and hot hint values (200 =-56).
-; RUN: opt < %s -passes=instcombine -optimize-hot-cold-new -cold-new-hint-value=5 -hot-new-hint-value=200 -notcold-new-hint-value=99 -ambiguous-new-hint-value=44 -S | FileCheck %s --check-prefix=HOTCOLD -DCOLD=5 -DHOT=-56 -DAMBIG=44 -DNOTCOLD=99 -DPREVHINTCOLD=7 -DPREVHINTNOTCOLD=7 -DPREVHINTHOT=7 -DPREVHINTAMBIG=7
+; RUN: opt < %s -passes=instcombine -optimize-hot-cold-new -cold-new-hint-value=5 -hot-new-hint-value=200 -notcold-new-hint-value=99 -ambiguous-new-hint-value=44 -S | FileCheck %s --check-prefixes=HOTCOLD,NOBUILTIN-OFF -DCOLD=5 -DHOT=-56 -DAMBIG=44 -DNOTCOLD=99 -DPREVHINTCOLD=7 -DPREVHINTNOTCOLD=7 -DPREVHINTHOT=7 -DPREVHINTAMBIG=7
+
+;; Next check with the same non-default cold and hot hint values (200 =-56),
+;; but with transformation of nobuiltin calls enabled.
+; RUN: opt < %s -passes=instcombine -optimize-hot-cold-new -optimize-nobuiltin-hot-cold-new-new -cold-new-hint-value=5 -hot-new-hint-value=200 -notcold-new-hint-value=99 -ambiguous-new-hint-value=44 -S | FileCheck %s --check-prefixes=HOTCOLD,NOBUILTIN-ON -DCOLD=5 -DHOT=-56 -DAMBIG=44 -DNOTCOLD=99 -DPREVHINTCOLD=7 -DPREVHINTNOTCOLD=7 -DPREVHINTHOT=7 -DPREVHINTAMBIG=7
 
 ;; Try again with the non-default cold and hot hint values (200 =-56), and this
 ;; time specify that existing hints should be updated.
-; RUN: opt < %s -passes=instcombine -optimize-hot-cold-new -cold-new-hint-value=5 -notcold-new-hint-value=100 -hot-new-hint-value=200 -ambiguous-new-hint-value=44 -optimize-existing-hot-cold-new -S | FileCheck %s --check-prefix=HOTCOLD -DCOLD=5 -DHOT=-56 -DNOTCOLD=100 -DAMBIG=44 -DPREVHINTCOLD=5 -DPREVHINTNOTCOLD=100 -DPREVHINTHOT=-56 -DPREVHINTAMBIG=44
+; RUN: opt < %s -passes=instcombine -optimize-hot-cold-new -cold-new-hint-value=5 -notcold-new-hint-value=100 -hot-new-hint-value=200 -ambiguous-new-hint-value=44 -optimize-existing-hot-cold-new -S | FileCheck %s --check-prefixes=HOTCOLD,NOBUILTIN-OFF -DCOLD=5 -DHOT=-56 -DNOTCOLD=100 -DAMBIG=44 -DPREVHINTCOLD=5 -DPREVHINTNOTCOLD=100 -DPREVHINTHOT=-56 -DPREVHINTAMBIG=44
 
 ;; Make sure that values not in 0..255 are flagged with an error
 ; RUN: not opt < %s -passes=instcombine -optimize-hot-cold-new -cold-new-hint-value=256 -S 2>&1 | FileCheck %s --check-prefix=ERROR
@@ -40,8 +44,9 @@ define void @new() {
   ; HOTCOLD: @_Znwm12__hot_cold_t(i64 10, i8 [[AMBIG]])
   %call4 = call ptr @_Znwm(i64 10) #7
   call void @dummy(ptr %call4)
-  ;; Attribute cold on a nobuiltin call has no effect.
-  ; HOTCOLD: @_Znwm(i64 10)
+  ;; Attribute cold on a nobuiltin call has no effect, unless optionally enabled.
+  ; NOBUILTIN-OFF: @_Znwm(i64 10)
+  ; NOBUILTIN-ON: @_Znwm12__hot_cold_t(i64 10, i8 [[COLD]])
   %call3 = call ptr @_Znwm(i64 10) #6
   call void @dummy(ptr %call3)
   ret void
@@ -68,8 +73,9 @@ define void @new_align() {
   ; HOTCOLD: @_ZnwmSt11align_val_t12__hot_cold_t(i64 10, i64 8, i8 [[AMBIG]])
   %call4 = call ptr @_ZnwmSt11align_val_t(i64 10, i64 8) #7
   call void @dummy(ptr %call4)
-  ;; Attribute cold on a nobuiltin call has no effect.
-  ; HOTCOLD: @_ZnwmSt11align_val_t(i64 10, i64 8)
+  ;; Attribute cold on a nobuiltin call has no effect, unless optionally enabled.
+  ; NOBUILTIN-OFF: @_ZnwmSt11align_val_t(i64 10, i64 8)
+  ; NOBUILTIN-ON: @_ZnwmSt11align_val_t12__hot_cold_t(i64 10, i64 8, i8 [[COLD]])
   %call3 = call ptr @_ZnwmSt11align_val_t(i64 10, i64 8) #6
   call void @dummy(ptr %call3)
   ret void
@@ -97,8 +103,9 @@ define void @new_nothrow() {
   ; HOTCOLD: @_ZnwmRKSt9nothrow_t12__hot_cold_t(i64 10, ptr nonnull %nt, i8 [[AMBIG]])
   %call4 = call ptr @_ZnwmRKSt9nothrow_t(i64 10, ptr %nt) #7
   call void @dummy(ptr %call4)
-  ;; Attribute cold on a nobuiltin call has no effect.
-  ; HOTCOLD: @_ZnwmRKSt9nothrow_t(i64 10, ptr nonnull %nt)
+  ;; Attribute cold on a nobuiltin call has no effect, unless optionally enabled.
+  ; NOBUILTIN-OFF: @_ZnwmRKSt9nothrow_t(i64 10, ptr nonnull %nt)
+  ; NOBUILTIN-ON: @_ZnwmRKSt9nothrow_t12__hot_cold_t(i64 10, ptr nonnull %nt, i8 [[COLD]])
   %call3 = call ptr @_ZnwmRKSt9nothrow_t(i64 10, ptr %nt) #6
   call void @dummy(ptr %call3)
   ret void
@@ -127,8 +134,9 @@ define void @new_align_nothrow() {
   ; HOTCOLD: @_ZnwmSt11align_val_tRKSt9nothrow_t12__hot_cold_t(i64 10, i64 8, ptr nonnull %nt, i8 [[AMBIG]])
   %call4 = call ptr @_ZnwmSt11align_val_tRKSt9nothrow_t(i64 10, i64 8, ptr %nt) #7
   call void @dummy(ptr %call4)
-  ;; Attribute cold on a nobuiltin call has no effect.
-  ; HOTCOLD: @_ZnwmSt11align_val_tRKSt9nothrow_t(i64 10, i64 8, ptr nonnull %nt)
+  ;; Attribute cold on a nobuiltin call has no effect, unless optionally enabled.
+  ; NOBUILTIN-OFF: @_ZnwmSt11align_val_tRKSt9nothrow_t(i64 10, i64 8, ptr nonnull %nt)
+  ; NOBUILTIN-ON: @_ZnwmSt11align_val_tRKSt9nothrow_t12__hot_cold_t(i64 10, i64 8, ptr nonnull %nt, i8 [[COLD]])
   %call3 = call ptr @_ZnwmSt11align_val_tRKSt9nothrow_t(i64 10, i64 8, ptr %nt) #6
   call void @dummy(ptr %call3)
   ret void
@@ -154,8 +162,9 @@ define void @array_new() {
   ; HOTCOLD: @_Znam12__hot_cold_t(i64 10, i8 [[AMBIG]])
   %call4 = call ptr @_Znam(i64 10) #7
   call void @dummy(ptr %call4)
-  ;; Attribute cold on a nobuiltin call has no effect.
-  ; HOTCOLD: @_Znam(i64 10)
+  ;; Attribute cold on a nobuiltin call has no effect, unless optionally enabled.
+  ; NOBUILTIN-OFF: @_Znam(i64 10)
+  ; NOBUILTIN-ON: @_Znam12__hot_cold_t(i64 10, i8 [[COLD]])
   %call3 = call ptr @_Znam(i64 10) #6
   call void @dummy(ptr %call3)
   ret void
@@ -182,8 +191,9 @@ define void @array_new_align() {
   ; HOTCOLD: @_ZnamSt11align_val_t12__hot_cold_t(i64 10, i64 8, i8 [[AMBIG]])
   %call4 = call ptr @_ZnamSt11align_val_t(i64 10, i64 8) #7
   call void @dummy(ptr %call4)
-  ;; Attribute cold on a nobuiltin call has no effect.
-  ; HOTCOLD: @_ZnamSt11align_val_t(i64 10, i64 8)
+  ;; Attribute cold on a nobuiltin call has no effect, unless optionally enabled.
+  ; NOBUILTIN-OFF: @_ZnamSt11align_val_t(i64 10, i64 8)
+  ; NOBUILTIN-ON: @_ZnamSt11align_val_t12__hot_cold_t(i64 10, i64 8, i8 [[COLD]])
   %call3 = call ptr @_ZnamSt11align_val_t(i64 10, i64 8) #6
   call void @dummy(ptr %call3)
   ret void
@@ -211,8 +221,9 @@ define void @array_new_nothrow() {
   ; HOTCOLD: @_ZnamRKSt9nothrow_t12__hot_cold_t(i64 10, ptr nonnull %nt, i8 [[AMBIG]])
   %call4 = call ptr @_ZnamRKSt9nothrow_t(i64 10, ptr %nt) #7
   call void @dummy(ptr %call4)
-  ;; Attribute cold on a nobuiltin call has no effect.
-  ; HOTCOLD: @_ZnamRKSt9nothrow_t(i64 10, ptr nonnull %nt)
+  ;; Attribute cold on a nobuiltin call has no effect, unless optionally enabled.
+  ; NOBUILTIN-OFF: @_ZnamRKSt9nothrow_t(i64 10, ptr nonnull %nt)
+  ; NOBUILTIN-ON: @_ZnamRKSt9nothrow_t12__hot_cold_t(i64 10, ptr nonnull %nt, i8 [[COLD]])
   %call3 = call ptr @_ZnamRKSt9nothrow_t(i64 10, ptr %nt) #6
   call void @dummy(ptr %call3)
   ret void
@@ -241,8 +252,9 @@ define void @array_new_align_nothrow() {
   ; HOTCOLD: @_ZnamSt11align_val_tRKSt9nothrow_t12__hot_cold_t(i64 10, i64 8, ptr nonnull %nt, i8 [[AMBIG]])
   %call4 = call ptr @_ZnamSt11align_val_tRKSt9nothrow_t(i64 10, i64 8, ptr %nt) #7
   call void @dummy(ptr %call4)
-  ;; Attribute cold on a nobuiltin call has no effect.
-  ; HOTCOLD: @_ZnamSt11align_val_tRKSt9nothrow_t(i64 10, i64 8, ptr nonnull %nt)
+  ;; Attribute cold on a nobuiltin call has no effect, unless optionally enabled.
+  ; NOBUILTIN-OFF: @_ZnamSt11align_val_tRKSt9nothrow_t(i64 10, i64 8, ptr nonnull %nt)
+  ; NOBUILTIN-ON: @_ZnamSt11align_val_tRKSt9nothrow_t12__hot_cold_t(i64 10, i64 8, ptr nonnull %nt, i8 [[COLD]])
   %call3 = call ptr @_ZnamSt11align_val_tRKSt9nothrow_t(i64 10, i64 8, ptr %nt) #6
   call void @dummy(ptr %call3)
   ret void
@@ -492,8 +504,9 @@ define void @size_returning_test() {
   %call4 = call {ptr, i64} @__size_returning_new(i64 10) #8
   %p4  = extractvalue {ptr, i64} %call4, 0
   call void @dummy(ptr %p4)
-  ;; Attribute cold on a nobuiltin call has no effect.
-  ; HOTCOLD: @__size_returning_new(i64 10)
+  ;; Attribute cold on a nobuiltin call has no effect, unless optionally enabled.
+  ; NOBUILTIN-OFF: @__size_returning_new(i64 10)
+  ; NOBUILTIN-ON: @__size_returning_new_hot_cold(i64 10, i8 [[COLD]])
   %call3 = call {ptr, i64} @__size_returning_new(i64 10) #6
   %p3 = extractvalue {ptr, i64} %call3, 0
   call void @dummy(ptr %p3)
@@ -524,8 +537,9 @@ define void @size_returning_aligned_test() {
   %call4 = call {ptr, i64} @__size_returning_new_aligned(i64 10, i64 8) #8
   %p4  = extractvalue {ptr, i64} %call4, 0
   call void @dummy(ptr %p4)
-  ;; Attribute cold on a nobuiltin call has no effect.
-  ; HOTCOLD: @__size_returning_new_aligned(i64 10, i64 8)
+  ;; Attribute cold on a nobuiltin call has no effect, unless optionally enabled.
+  ; NOBUILTIN-OFF: @__size_returning_new_aligned(i64 10, i64 8)
+  ; NOBUILTIN-ON: @__size_returning_new_aligned_hot_cold(i64 10, i64 8, i8 [[COLD]])
   %call3 = call {ptr, i64} @__size_returning_new_aligned(i64 10, i64 8) #6
   %p3 = extractvalue {ptr, i64} %call3, 0
   call void @dummy(ptr %p3)