Allow optimization of __size_returning_new variants. (#102258)

#101564 added support to TLI to
detect variants of operator new which provide feedback on the actual
size of memory allocated (http://wg21.link/P0901R5). This patch extends
SimplifyLibCalls to handle hot cold hinting of these variants.

llvm/include/llvm/Transforms/Utils/BuildLibCalls.h

@@ -267,6 +267,13 @@ namespace llvm {
                                       IRBuilderBase &B,
                                       const TargetLibraryInfo *TLI,
                                       LibFunc NewFunc, uint8_t HotCold);
+  Value *emitHotColdSizeReturningNew(Value *Num, IRBuilderBase &B,
+                                     const TargetLibraryInfo *TLI,
+                                     LibFunc NewFunc, uint8_t HotCold);
+  Value *emitHotColdSizeReturningNewAligned(Value *Num, Value *Align,
+                                            IRBuilderBase &B,
+                                            const TargetLibraryInfo *TLI,
+                                            LibFunc NewFunc, uint8_t HotCold);
 }
 
 #endif

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

@@ -3692,7 +3692,7 @@ Instruction *InstCombinerImpl::tryOptimizeCall(CallInst *CI) {
 
   // Skip optimizing notail and musttail calls so
   // LibCallSimplifier::optimizeCall doesn't have to preserve those invariants.
-  // LibCallSimplifier::optimizeCall should try to preseve tail calls though.
+  // LibCallSimplifier::optimizeCall should try to preserve tail calls though.
   if (CI->isMustTailCall() || CI->isNoTailCall())
     return nullptr;
 

llvm/lib/Transforms/Instrumentation/MemProfiler.cpp

@@ -20,6 +20,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Analysis/MemoryBuiltins.h"
 #include "llvm/Analysis/MemoryProfileInfo.h"
+#include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/DataLayout.h"
@@ -753,8 +754,8 @@ stackFrameIncludesInlinedCallStack(ArrayRef<Frame> ProfileCallStack,
   return InlCallStackIter == InlinedCallStack.end();
 }
 
-static bool isNewWithHotColdVariant(Function *Callee,
-                                    const TargetLibraryInfo &TLI) {
+static bool isAllocationWithHotColdVariant(Function *Callee,
+                                           const TargetLibraryInfo &TLI) {
   if (!Callee)
     return false;
   LibFunc Func;
@@ -769,6 +770,8 @@ static bool isNewWithHotColdVariant(Function *Callee,
   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:
     return true;
   case LibFunc_Znwm12__hot_cold_t:
   case LibFunc_ZnwmRKSt9nothrow_t12__hot_cold_t:
@@ -778,6 +781,8 @@ static bool isNewWithHotColdVariant(Function *Callee,
   case LibFunc_ZnamRKSt9nothrow_t12__hot_cold_t:
   case LibFunc_ZnamSt11align_val_t12__hot_cold_t:
   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 ClMemProfMatchHotColdNew;
   default:
     return false;
@@ -945,9 +950,8 @@ readMemprof(Module &M, Function &F, IndexedInstrProfReader *MemProfReader,
       // instruction's locations match the prefix Frame locations on an
       // allocation context with the same leaf.
       if (AllocInfoIter != LocHashToAllocInfo.end()) {
-        // Only consider allocations via new, to reduce unnecessary metadata,
-        // since those are the only allocations that will be targeted initially.
-        if (!isNewWithHotColdVariant(CI->getCalledFunction(), TLI))
+        // Only consider allocations which support hinting.
+        if (!isAllocationWithHotColdVariant(CI->getCalledFunction(), TLI))
           continue;
         // We may match this instruction's location list to multiple MIB
         // contexts. Add them to a Trie specialized for trimming the contexts to

llvm/lib/Transforms/Utils/BuildLibCalls.cpp

@@ -19,6 +19,7 @@
 #include "llvm/IR/CallingConv.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
+#include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Module.h"
@@ -1963,6 +1964,56 @@ Value *llvm::emitCalloc(Value *Num, Value *Size, IRBuilderBase &B,
   return CI;
 }
 
+Value *llvm::emitHotColdSizeReturningNew(Value *Num, IRBuilderBase &B,
+                                         const TargetLibraryInfo *TLI,
+                                         LibFunc SizeFeedbackNewFunc,
+                                         uint8_t HotCold) {
+  Module *M = B.GetInsertBlock()->getModule();
+  if (!isLibFuncEmittable(M, TLI, SizeFeedbackNewFunc))
+    return nullptr;
+
+  StringRef Name = TLI->getName(SizeFeedbackNewFunc);
+
+  // __sized_ptr_t struct return type { void*, size_t }
+  StructType *SizedPtrT =
+      StructType::get(M->getContext(), {B.getPtrTy(), Num->getType()});
+  FunctionCallee Func =
+      M->getOrInsertFunction(Name, SizedPtrT, Num->getType(), B.getInt8Ty());
+  inferNonMandatoryLibFuncAttrs(M, Name, *TLI);
+  CallInst *CI = B.CreateCall(Func, {Num, B.getInt8(HotCold)}, "sized_ptr");
+
+  if (const Function *F = dyn_cast<Function>(Func.getCallee()))
+    CI->setCallingConv(F->getCallingConv());
+
+  return CI;
+}
+
+Value *llvm::emitHotColdSizeReturningNewAligned(Value *Num, Value *Align,
+                                                IRBuilderBase &B,
+                                                const TargetLibraryInfo *TLI,
+                                                LibFunc SizeFeedbackNewFunc,
+                                                uint8_t HotCold) {
+  Module *M = B.GetInsertBlock()->getModule();
+  if (!isLibFuncEmittable(M, TLI, SizeFeedbackNewFunc))
+    return nullptr;
+
+  StringRef Name = TLI->getName(SizeFeedbackNewFunc);
+
+  // __sized_ptr_t struct return type { void*, size_t }
+  StructType *SizedPtrT =
+      StructType::get(M->getContext(), {B.getPtrTy(), Num->getType()});
+  FunctionCallee Func = M->getOrInsertFunction(Name, SizedPtrT, Num->getType(),
+                                               Align->getType(), B.getInt8Ty());
+  inferNonMandatoryLibFuncAttrs(M, Name, *TLI);
+  CallInst *CI =
+      B.CreateCall(Func, {Num, Align, B.getInt8(HotCold)}, "sized_ptr");
+
+  if (const Function *F = dyn_cast<Function>(Func.getCallee()))
+    CI->setCallingConv(F->getCallingConv());
+
+  return CI;
+}
+
 Value *llvm::emitHotColdNew(Value *Num, IRBuilderBase &B,
                             const TargetLibraryInfo *TLI, LibFunc NewFunc,
                             uint8_t HotCold) {

llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp

@@ -18,6 +18,7 @@
 #include "llvm/Analysis/ConstantFolding.h"
 #include "llvm/Analysis/Loads.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
+#include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/AttributeMask.h"
 #include "llvm/IR/DataLayout.h"
@@ -1745,7 +1746,7 @@ Value *LibCallSimplifier::optimizeNew(CallInst *CI, IRBuilderBase &B,
   // if cold or hot, and leave as-is for default handling if "notcold" aka warm.
   // Note that in cases where we decide it is "notcold", it might be slightly
   // better to replace the hinted call with a non hinted call, to avoid the
-  // extra paramter and the if condition check of the hint value in the
+  // extra parameter and the if condition check of the hint value in the
   // allocator. This can be considered in the future.
   switch (Func) {
   case LibFunc_Znwm12__hot_cold_t:
@@ -1844,6 +1845,30 @@ Value *LibCallSimplifier::optimizeNew(CallInst *CI, IRBuilderBase &B,
           TLI, LibFunc_ZnamSt11align_val_tRKSt9nothrow_t12__hot_cold_t,
           HotCold);
     break;
+  case LibFunc_size_returning_new:
+    if (HotCold != NotColdNewHintValue)
+      return emitHotColdSizeReturningNew(CI->getArgOperand(0), B, TLI,
+                                         LibFunc_size_returning_new_hot_cold,
+                                         HotCold);
+    break;
+  case LibFunc_size_returning_new_hot_cold:
+    if (OptimizeExistingHotColdNew)
+      return emitHotColdSizeReturningNew(CI->getArgOperand(0), B, TLI,
+                                         LibFunc_size_returning_new_hot_cold,
+                                         HotCold);
+    break;
+  case LibFunc_size_returning_new_aligned:
+    if (HotCold != NotColdNewHintValue)
+      return emitHotColdSizeReturningNewAligned(
+          CI->getArgOperand(0), CI->getArgOperand(1), B, TLI,
+          LibFunc_size_returning_new_aligned_hot_cold, HotCold);
+    break;
+  case LibFunc_size_returning_new_aligned_hot_cold:
+    if (OptimizeExistingHotColdNew)
+      return emitHotColdSizeReturningNewAligned(
+          CI->getArgOperand(0), CI->getArgOperand(1), B, TLI,
+          LibFunc_size_returning_new_aligned_hot_cold, HotCold);
+    break;
   default:
     return nullptr;
   }
@@ -3759,6 +3784,7 @@ Value *LibCallSimplifier::optimizeStringMemoryLibCall(CallInst *CI,
   Module *M = CI->getModule();
   LibFunc Func;
   Function *Callee = CI->getCalledFunction();
+
   // Check for string/memory library functions.
   if (TLI->getLibFunc(*Callee, Func) && isLibFuncEmittable(M, TLI, Func)) {
     // Make sure we never change the calling convention.
@@ -3851,6 +3877,10 @@ Value *LibCallSimplifier::optimizeStringMemoryLibCall(CallInst *CI,
     case LibFunc_ZnamRKSt9nothrow_t12__hot_cold_t:
     case LibFunc_ZnamSt11align_val_t12__hot_cold_t:
     case LibFunc_ZnamSt11align_val_tRKSt9nothrow_t12__hot_cold_t:
+    case LibFunc_size_returning_new:
+    case LibFunc_size_returning_new_hot_cold:
+    case LibFunc_size_returning_new_aligned:
+    case LibFunc_size_returning_new_aligned_hot_cold:
       return optimizeNew(CI, Builder, Func);
     default:
       break;

llvm/test/Transforms/InstCombine/simplify-libcalls-new.ll

@@ -340,6 +340,94 @@ define void @array_new_align_nothrow_hot_cold() {
   ret void
 }
 
+;; Check that operator __size_returning_new(unsigned long) converted to
+;; __size_returning_new(unsigned long, __hot_cold_t) with a hot or cold attribute.
+; HOTCOLD-LABEL: @size_returning_test()
+define void @size_returning_test() {
+  ;; Attribute cold converted to __hot_cold_t cold value.
+  ; HOTCOLD: @__size_returning_new_hot_cold(i64 10, i8 [[COLD]])
+  %call = call {ptr, i64} @__size_returning_new(i64 10) #3
+  %p  = extractvalue {ptr, i64} %call, 0
+  call void @dummy(ptr %p)
+  ;; Attribute notcold has no effect.
+  ; HOTCOLD: @__size_returning_new(i64 10)
+  %call1 = call {ptr, i64} @__size_returning_new(i64 10) #4
+  %p1  = extractvalue {ptr, i64} %call1, 0
+  call void @dummy(ptr %p1)
+  ;; Attribute hot converted to __hot_cold_t hot value.
+  ; HOTCOLD: @__size_returning_new_hot_cold(i64 10, i8 [[HOT]])
+  %call2 = call {ptr, i64} @__size_returning_new(i64 10) #5
+  %p2  = extractvalue {ptr, i64} %call2, 0
+  call void @dummy(ptr %p2)
+  ret void
+}
+
+;; Check that operator __size_returning_new_aligned(unsigned long, std::align_val_t) converted to
+;; __size_returning_new_aligned(unsigned long, std::align_val_t, __hot_cold_t) with a hot or cold attribute.
+; HOTCOLD-LABEL: @size_returning_aligned_test()
+define void @size_returning_aligned_test() {
+  ;; Attribute cold converted to __hot_cold_t cold value.
+  ; HOTCOLD: @__size_returning_new_aligned_hot_cold(i64 10, i64 8, i8 [[COLD]])
+  %call = call {ptr, i64} @__size_returning_new_aligned(i64 10, i64 8) #3
+  %p  = extractvalue {ptr, i64} %call, 0
+  call void @dummy(ptr %p)
+  ;; Attribute notcold has no effect.
+  ; HOTCOLD: @__size_returning_new_aligned(i64 10, i64 8)
+  %call1 = call {ptr, i64} @__size_returning_new_aligned(i64 10, i64 8) #4
+  %p1  = extractvalue {ptr, i64} %call1, 0
+  call void @dummy(ptr %p1)
+  ;; Attribute hot converted to __hot_cold_t hot value.
+  ; HOTCOLD: @__size_returning_new_aligned_hot_cold(i64 10, i64 8, i8 [[HOT]])
+  %call2 = call {ptr, i64} @__size_returning_new_aligned(i64 10, i64 8) #5
+  %p2  = extractvalue {ptr, i64} %call2, 0
+  call void @dummy(ptr %p2)
+  ret void
+}
+
+;; Check that __size_returning_new_hot_cold(unsigned long, __hot_cold_t)
+;; optionally has its hint updated.
+; HOTCOLD-LABEL: @size_returning_update_test()
+define void @size_returning_update_test() {
+  ;; Attribute cold converted to __hot_cold_t cold value.
+  ; HOTCOLD: @__size_returning_new_hot_cold(i64 10, i8 [[PREVHINTCOLD]])
+  %call = call {ptr, i64} @__size_returning_new_hot_cold(i64 10, i8 7) #3
+  %p  = extractvalue {ptr, i64} %call, 0
+  call void @dummy(ptr %p)
+  ;; Attribute notcold converted to __hot_cold_t notcold value.
+  ; HOTCOLD: @__size_returning_new_hot_cold(i64 10, i8 [[PREVHINTNOTCOLD]])
+  %call1 = call {ptr, i64} @__size_returning_new_hot_cold(i64 10, i8 7) #4
+  %p1 = extractvalue {ptr, i64} %call1, 0
+  call void @dummy(ptr %p1)
+  ;; Attribute hot converted to __hot_cold_t hot value.
+  ; HOTCOLD: @__size_returning_new_hot_cold(i64 10, i8 [[PREVHINTHOT]])
+  %call2 = call {ptr, i64} @__size_returning_new_hot_cold(i64 10, i8 7) #5
+  %p2 = extractvalue {ptr, i64} %call2, 0
+  call void @dummy(ptr %p2)
+  ret void
+}
+
+;; Check that __size_returning_new_aligned_hot_cold(unsigned long, __hot_cold_t)
+;; optionally has its hint updated.
+; HOTCOLD-LABEL: @size_returning_aligned_update_test()
+define void @size_returning_aligned_update_test() {
+  ;; Attribute cold converted to __hot_cold_t cold value.
+  ; HOTCOLD: @__size_returning_new_aligned_hot_cold(i64 10, i64 8, i8 [[PREVHINTCOLD]])
+  %call = call {ptr, i64} @__size_returning_new_aligned_hot_cold(i64 10, i64 8, i8 7) #3
+  %p  = extractvalue {ptr, i64} %call, 0
+  call void @dummy(ptr %p)
+  ;; Attribute notcold converted to __hot_cold_t notcold value.
+  ; HOTCOLD: @__size_returning_new_aligned_hot_cold(i64 10, i64 8, i8 [[PREVHINTNOTCOLD]])
+  %call1 = call {ptr, i64} @__size_returning_new_aligned_hot_cold(i64 10, i64 8, i8 7) #4
+  %p1 = extractvalue {ptr, i64} %call1, 0
+  call void @dummy(ptr %p1)
+  ;; Attribute hot converted to __hot_cold_t hot value.
+  ; HOTCOLD: @__size_returning_new_aligned_hot_cold(i64 10, i64 8, i8 [[PREVHINTHOT]])
+  %call2 = call {ptr, i64} @__size_returning_new_aligned_hot_cold(i64 10, i64 8, i8 7) #5
+  %p2 = extractvalue {ptr, i64} %call2, 0
+  call void @dummy(ptr %p2)
+  ret void
+}
+
 ;; So that instcombine doesn't optimize out the call.
 declare void @dummy(ptr)
 
@@ -360,6 +448,18 @@ declare ptr @_ZnamSt11align_val_t12__hot_cold_t(i64, i64, i8)
 declare ptr @_ZnamRKSt9nothrow_t12__hot_cold_t(i64, ptr, i8)
 declare ptr @_ZnamSt11align_val_tRKSt9nothrow_t12__hot_cold_t(i64, i64, ptr, i8)
 
+
+declare {ptr, i64} @__size_returning_new(i64)
+declare {ptr, i64} @__size_returning_new_hot_cold(i64, i8)
+declare {ptr, i64} @__size_returning_new_aligned(i64, i64)
+declare {ptr, i64} @__size_returning_new_aligned_hot_cold(i64, i64, i8)
+
 attributes #0 = { builtin allocsize(0) "memprof"="cold" }
 attributes #1 = { builtin allocsize(0) "memprof"="notcold" }
 attributes #2 = { builtin allocsize(0) "memprof"="hot" }
+
+;; Use separate attributes for __size_returning_new variants since they are not
+;; treated as builtins.
+attributes #3 = { "memprof" = "cold" }
+attributes #4 = { "memprof" = "notcold" }
+attributes #5 = { "memprof" = "hot" }