[SimplifyLibcalls] Fold malloc + memset to calloc even for llvm.memset

[dblaikie]

Test case?

[lebedev.ri]

@xbolva00 please always always upload all the patches with the full context (-U99999)

test/Transforms/InstCombine/memset-1.ll

30

There is no https://bugs.llvm.org/show_bug.cgi?id=45344
You mean D45344?
Then it may be better not to use such a specific name.

[xbolva00]

Better test

[lebedev.ri]

(The diff still seem to be lacking context)

[xbolva00]

Fixed previous test, now it passes

[xbolva00]

Added context

[xbolva00]

In D45344#1059408, @lebedev.ri wrote:

(The diff still seem to be lacking context)

Yes, fixed now :D I am sorry for too many updated revisions, it is my first contribution to llvm :D

[spatel]

IIUC, these are 2 independent changes, so they should be split into separate patches:

1. Handle the memset intrinsic identically to the way we handle the memset libcall. What about memset_chk()?
2. Remove the hasOneUse() limitation on the malloc. The reason that's there is discussed in D16337 - the transform must be made safe from intermediate stores. This patch as-is will miscompile that pattern.

I've added tests at rL329412 and regenerated the check lines. Please fix/limit this patch and rebase.

[xbolva00]

In D45344#1059906, @spatel wrote:

IIUC, these are 2 independent changes, so they should be split into separate patches:

1. Handle the memset intrinsic identically to the way we handle the memset libcall. What about memset_chk()?
2. Remove the hasOneUse() limitation on the malloc. The reason that's there is discussed in D16337 - the transform must be made safe from intermediate stores. This patch as-is will miscompile that pattern.

I've added tests at rL329412 and regenerated the check lines. Please fix/limit this patch and rebase.

1. memset_chk change was removed from this commit.
2. https://reviews.llvm.org/D45381

[spatel]

Tests?

[xbolva00]

In D45344#1060243, @spatel wrote:

Tests?

Without https://reviews.llvm.org/D45381 is llvm.memset not folded to calloc.

[xbolva00]

In D45344#1060243, @spatel wrote:

Tests?

Your patch https://reviews.llvm.org/rL329412 (@malloc_and_memset_intrinsic) needs to be updated after these two patches are applied.

[xbolva00]

Added test

[spatel]

I'm confused. This patch is supposed to show a calloc transform.

@malloc_and_memset_intrinsic has multiple uses of the malloc result, so the transform doesn't work there. Is there some other way to show the transform? If not, then sorry that I misunderstood, but we can't split the patches.

But D45381 just removes the one-use check and miscompiles @buffer_is_modified_then_memset. Why is that ok?

[xbolva00]

In D45344#1060345, @spatel wrote:

I'm confused. This patch is supposed to show a calloc transform.

@malloc_and_memset_intrinsic has multiple uses of the malloc result, so the transform doesn't work there. Is there some other way to show the transform? If not, then sorry that I misunderstood, but we can't split the patches.

But D45381 just removes the one-use check and miscompiles @buffer_is_modified_then_memset. Why is that ok?

I closed D45381 & reupdated this patch, added test case. Test @buffer_is_modified_then_memset is ok now too.

[spatel]

If I apply this patch, the test case fails.

Do you understand the problem with the intervening store? You can't solve that by simply removing the one-use check.

[xbolva00]

In D45344#1060689, @spatel wrote:

If I apply this patch, the test case fails.

Do you understand the problem with the intervening store? You can't solve that by simply removing the one-use check.

Oh, yes :( Any idea how to work around it?

[xbolva00]

I added solution.

Simply check any Instruction between malloc and memset if use value from malloc, if so, do not fold.

[xbolva00]

WIP

[xbolva00]

Removed old test and added new one. Works for me.

[xbolva00]

Anything to be done? Or is this ok? @efriedma , @spatel...

[xbolva00]

Patch reworked.

• Malloc + llvm.memset folding
• Store instructions eliminated between malloc and memset calls.

Can you look at it? @spatel @efriedma

[efriedma]

This transform really belongs in DeadStoreElimination, not here; it has the relevant logic and some similar transforms.

[xbolva00]

In D45344#1081661, @efriedma wrote:

This transform really belongs in DeadStoreElimination, not here; it has the relevant logic and some similar transforms.

But folding part for intrinsic memset can be here, no?

[efriedma]

I mean the whole foldMallocMemset transform should be in DSE. I mean, arguably the transform isn't actually eliminating a dead store, but the code you need is very similar to eliminateNoopStore.

[xbolva00]

In D45344#1081728, @efriedma wrote:

I mean the whole foldMallocMemset transform should be in DSE. I mean, arguably the transform isn't actually eliminating a dead store, but the code you need is very similar to eliminateNoopStore.

Right, so I will close this for now..

include/llvm/Transforms/Utils/SimplifyLibCalls.h

private:
   Value *optimizeStrSpn(CallInst *CI, IRBuilder<> &B);
   Value *optimizeStrCSpn(CallInst *CI, IRBuilder<> &B);
   Value *optimizeStrStr(CallInst *CI, IRBuilder<> &B);
   Value *optimizeMemChr(CallInst *CI, IRBuilder<> &B);
   Value *optimizeMemCmp(CallInst *CI, IRBuilder<> &B);
   Value *optimizeMemCpy(CallInst *CI, IRBuilder<> &B);
   Value *optimizeMemMove(CallInst *CI, IRBuilder<> &B);
   Value *optimizeMemSet(CallInst *CI, IRBuilder<> &B);
+  Value *optimizeIntrinsicMemSet(CallInst *CI, IRBuilder<> &B);
   Value *optimizeRealloc(CallInst *CI, IRBuilder<> &B);
   Value *optimizeWcslen(CallInst *CI, IRBuilder<> &B);
   // Wrapper for all String/Memory Library Call Optimizations
   Value *optimizeStringMemoryLibCall(CallInst *CI, IRBuilder<> &B);
 
   // Math Library Optimizations
   Value *optimizeCAbs(CallInst *CI, IRBuilder<> &B);
   Value *optimizeCos(CallInst *CI, IRBuilder<> &B);

lib/Transforms/Utils/BuildLibCalls.cpp

bool llvm::inferLibFuncAttributes(Function &F, const TargetLibraryInfo &TLI) {
   case LibFunc_msvc_new_array_longlong: // new[](unsigned long long)
     // Operator new always returns a nonnull noalias pointer
     Changed |= setRetNonNull(F);
     Changed |= setRetDoesNotAlias(F);
     return Changed;
   // TODO: add LibFunc entries for:
   // case LibFunc_memset_pattern4:
   // case LibFunc_memset_pattern8:
+  case LibFunc_memset:
   case LibFunc_memset_pattern16:
     Changed |= setOnlyAccessesArgMemory(F);
     Changed |= setDoesNotCapture(F, 0);
     Changed |= setDoesNotCapture(F, 1);
     Changed |= setOnlyReadsMemory(F, 1);
     return Changed;
   // int __nvvm_reflect(const char *)
   case LibFunc_nvvm_reflect:

lib/Transforms/Utils/SimplifyLibCalls.cpp

 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/Analysis/ConstantFolding.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/Utils/Local.h"
 #include "llvm/Analysis/ValueTracking.h"
+#include "llvm/Analysis/CaptureTracking.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PatternMatch.h"

 /// Fold memset[_chk](malloc(n), 0, n) --> calloc(1, n).
 static Value *foldMallocMemset(CallInst *Memset, IRBuilder<> &B,
                                const TargetLibraryInfo &TLI) {
   // This has to be a memset of zeros (bzero).
   auto *FillValue = dyn_cast<ConstantInt>(Memset->getArgOperand(1));
   if (!FillValue || FillValue->getZExtValue() != 0)
     return nullptr;
 
-  // TODO: We should handle the case where the malloc has more than one use.
-  // This is necessary to optimize common patterns such as when the result of
-  // the malloc is checked against null or when a memset intrinsic is used in
-  // place of a memset library call.
   auto *Malloc = dyn_cast<CallInst>(Memset->getArgOperand(0));
-  if (!Malloc || !Malloc->hasOneUse())
+  if (!Malloc)
     return nullptr;
 
   // Is the inner call really malloc()?
   Function *InnerCallee = Malloc->getCalledFunction();
   if (!InnerCallee)
     return nullptr;
 
   LibFunc Func;
   if (!TLI.getLibFunc(*InnerCallee, Func) || !TLI.has(Func) ||
       Func != LibFunc_malloc)
     return nullptr;
 
+  // Check if malloc value is captured
+  if (PointerMayBeCaptured(Malloc, /* ReturnCaptures */ false,
+                           /* StoreCaptures */ true))
+    return nullptr;
+
   // The memset must cover the same number of bytes that are malloc'd.
   if (Memset->getArgOperand(2) != Malloc->getArgOperand(0))
     return nullptr;
 
+  Instruction *Inst;
+  bool MemsetFound = false;
+  for (Use &U : Malloc->uses()) {
+    Inst = cast<Instruction>(U.getUser());
+
+    if (Inst == Memset)
+      MemsetFound = true;
+
+    // Remove stores between malloc and memset
+    if (MemsetFound && isa<StoreInst>(U.getUser())) {
+      if (Inst->getOperand(1) == Malloc)
+        Inst->eraseFromParent();
+    }
+  }
+
   // Replace the malloc with a calloc. We need the data layout to know what the
-  // actual size of a 'size_t' parameter is. 
+  // actual size of a 'size_t' parameter is.
   B.SetInsertPoint(Malloc->getParent(), ++Malloc->getIterator());
   const DataLayout &DL = Malloc->getModule()->getDataLayout();
   IntegerType *SizeType = DL.getIntPtrType(B.GetInsertBlock()->getContext());
-  Value *Calloc = emitCalloc(ConstantInt::get(SizeType, 1),
-                             Malloc->getArgOperand(0), Malloc->getAttributes(),
-                             B, TLI);
+  Value *Calloc =
+      emitCalloc(ConstantInt::get(SizeType, 1), Malloc->getArgOperand(0),
+                 Malloc->getAttributes(), B, TLI);
   if (!Calloc)
     return nullptr;
 
   Malloc->replaceAllUsesWith(Calloc);
   Malloc->eraseFromParent();
 
   return Calloc;
 }

 
 Value *LibCallSimplifier::optimizeRealloc(CallInst *CI, IRBuilder<> &B) {
   if (isa<ConstantPointerNull>(CI->getArgOperand(0)))
     return emitMalloc(CI->getArgOperand(1), B, DL, TLI);
 
   return nullptr;
 }
 
+Value *LibCallSimplifier::optimizeIntrinsicMemSet(CallInst *CI,
+                                                   IRBuilder<> &B) {
+  if (auto *Calloc = foldMallocMemset(CI, B, *TLI))
+    return Calloc;
+
+  return nullptr;
+}
+
 //===----------------------------------------------------------------------===//
 // Math Library Optimizations
 //===----------------------------------------------------------------------===//
 
 /// Return a variant of Val with float type.
 /// Currently this works in two cases: If Val is an FPExtension of a float
 /// value to something bigger, simply return the operand.
 /// If Val is a ConstantFP but can be converted to a float ConstantFP without
if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(CI)) {
     case Intrinsic::pow:
       return optimizePow(CI, Builder);
     case Intrinsic::exp2:
       return optimizeExp2(CI, Builder);
     case Intrinsic::log:
       return optimizeLog(CI, Builder);
     case Intrinsic::sqrt:
       return optimizeSqrt(CI, Builder);
-    // TODO: Use foldMallocMemset() with memset intrinsic.
+    case Intrinsic::memset:
+      return optimizeIntrinsicMemSet(CI, Builder);
     default:
       return nullptr;
     }
   }
 
   // Also try to simplify calls to fortified library functions.
   if (Value *SimplifiedFortifiedCI = FortifiedSimplifier.optimizeCall(CI)) {
     // Try to further simplify the result.

test/Transforms/InstCombine/memset-1.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; Test that the memset library call simplifier works correctly.
 ;
 ; RUN: opt < %s -instcombine -S | FileCheck %s
 
 target datalayout = "e-p:32:32:32-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:32:64-f32:32:32-f64:32:64-v64:64:64-v128:128:128-a0:0:64-f80:128:128"
 
-declare i8* @memset(i8*, i32, i32)
+declare i8* @memset(i8* nocapture, i32, i32)
 declare void @llvm.memset.p0i8.i32(i8* nocapture writeonly, i8, i32, i32, i1)
 declare noalias i8* @malloc(i32) #1
 
 ; Check memset(mem1, val, size) -> llvm.memset(mem1, val, size, 1).
 
 define i8* @test_simplify1(i8* %mem, i32 %val, i32 %size) {
 ; CHECK-LABEL: @test_simplify1(
 ; CHECK-NEXT:    [[TMP1:%.*]] = trunc i32 [[VAL:%.*]] to i8
 ; CHECK-NEXT:    call void @llvm.memset.p0i8.i32(i8* align 1 [[MEM:%.*]], i8 [[TMP1]], i32 [[SIZE:%.*]], i1 false)
 ; CHECK-NEXT:    ret i8* [[MEM]]
 ;
   %ret = call i8* @memset(i8* %mem, i32 %val, i32 %size)
   ret i8* %ret
 }
 
 define i8* @pr25892_lite(i32 %size) #0 {
 ; CHECK-LABEL: @pr25892_lite(
 ; CHECK-NEXT:    [[CALLOC:%.*]] = call i8* @calloc(i32 1, i32 [[SIZE:%.*]])
 ; CHECK-NEXT:    ret i8* [[CALLOC]]
 ;
   %call1 = call i8* @malloc(i32 %size) #1
   %call2 = call i8* @memset(i8* %call1, i32 0, i32 %size) #1
   ret i8* %call2
 }
 
 ; FIXME: A memset intrinsic should be handled similarly to a memset() libcall.
 
 define i8* @malloc_and_memset_intrinsic(i32 %n) #0 {
 ; CHECK-LABEL: @malloc_and_memset_intrinsic(
-; CHECK-NEXT:    [[CALL:%.*]] = call i8* @malloc(i32 [[N:%.*]])
-; CHECK-NEXT:    call void @llvm.memset.p0i8.i32(i8* align 1 [[CALL]], i8 0, i32 [[N]], i1 false)
-; CHECK-NEXT:    ret i8* [[CALL]]
+; CHECK-NEXT:    [[CALLOC:%.*]] = call i8* @calloc(i32 1, i32 [[N:%.*]])
+; CHECK-NEXT:    ret i8* [[CALLOC]]
 ;
   %call = call i8* @malloc(i32 %n)
   call void @llvm.memset.p0i8.i32(i8* %call, i8 0, i32 %n, i32 1, i1 false)
   ret i8* %call
 }
 
 ; This should not create a calloc and should not crash the compiler.
 
 define i8* @notmalloc_memset(i32 %size, i8*(i32)* %notmalloc) {
 ; CHECK-LABEL: @notmalloc_memset(
 ; CHECK-NEXT:    [[CALL1:%.*]] = call i8* [[NOTMALLOC:%.*]](i32 [[SIZE:%.*]]) #0
 ; CHECK-NEXT:    call void @llvm.memset.p0i8.i32(i8* align 1 [[CALL1]], i8 0, i32 [[SIZE]], i1 false)
 ; CHECK-NEXT:    ret i8* [[CALL1]]
 ;
   %call1 = call i8* %notmalloc(i32 %size) #1
   %call2 = call i8* @memset(i8* %call1, i32 0, i32 %size) #1
   ret i8* %call2
 }
 
-; FIXME: memset(malloc(x), 0, x) -> calloc(1, x)
-; This doesn't fire currently because the malloc has more than one use.
+; Malloc has more than one use.
 
 define float* @pr25892(i32 %size) #0 {
 ; CHECK-LABEL: @pr25892(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[CALL:%.*]] = tail call i8* @malloc(i32 [[SIZE:%.*]]) #0
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8* [[CALL]], null
+; CHECK-NEXT:    [[CALLOC:%.*]] = call i8* @calloc(i32 1, i32 [[SIZE:%.*]])
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8* [[CALLOC]], null
 ; CHECK-NEXT:    br i1 [[CMP]], label [[CLEANUP:%.*]], label [[IF_END:%.*]]
 ; CHECK:       if.end:
-; CHECK-NEXT:    [[BC:%.*]] = bitcast i8* [[CALL]] to float*
-; CHECK-NEXT:    call void @llvm.memset.p0i8.i32(i8* nonnull align 1 [[CALL]], i8 0, i32 [[SIZE]], i1 false)
+; CHECK-NEXT:    [[BC:%.*]] = bitcast i8* [[CALLOC]] to float*
 ; CHECK-NEXT:    br label [[CLEANUP]]
 ; CHECK:       cleanup:
 ; CHECK-NEXT:    [[RETVAL_0:%.*]] = phi float* [ [[BC]], [[IF_END]] ], [ null, [[ENTRY:%.*]] ]
 ; CHECK-NEXT:    ret float* [[RETVAL_0]]
 ;
 entry:
   %call = tail call i8* @malloc(i32 %size) #1
   %cmp = icmp eq i8* %call, null
   br i1 %cmp, label %cleanup, label %if.end
 if.end:
   %bc = bitcast i8* %call to float*
   %call2 = tail call i8* @memset(i8* nonnull %call, i32 0, i32 %size) #1
   br label %cleanup
 cleanup:
   %retval.0 = phi float* [ %bc, %if.end ], [ null, %entry ]
   ret float* %retval.0
 }
 
 ; If there's a calloc transform, the store must also be eliminated.
 
 define i8* @buffer_is_modified_then_memset(i32 %size) {
 ; CHECK-LABEL: @buffer_is_modified_then_memset(
-; CHECK-NEXT:    [[PTR:%.*]] = tail call i8* @malloc(i32 [[SIZE:%.*]]) #0
-; CHECK-NEXT:    store i8 1, i8* [[PTR]], align 1
-; CHECK-NEXT:    call void @llvm.memset.p0i8.i32(i8* align 1 [[PTR]], i8 0, i32 [[SIZE]], i1 false)
-; CHECK-NEXT:    ret i8* [[PTR]]
+; CHECK-NEXT:    [[CALLOC:%.*]] = call i8* @calloc(i32 1, i32 [[SIZE:%.*]])
+; CHECK-NEXT:    ret i8* [[CALLOC]]
 ;
   %ptr = tail call i8* @malloc(i32 %size) #1
   store i8 1, i8* %ptr           ;; fdata[0] = 1;
   %memset = tail call i8* @memset(i8* nonnull %ptr, i32 0, i32 %size) #1
   ret i8* %memset
 }
 
+define i8* @buffer_preserve_store_after_memset(i32 %size) {
+; CHECK-LABEL: @buffer_preserve_store_after_memset(
+; CHECK-NEXT:    [[CALLOC:%.*]] = call i8* @calloc(i32 1, i32 [[SIZE:%.*]])
+; CHECK-NEXT:    store i8 1, i8* [[CALLOC]], align 1
+; CHECK-NEXT:    ret i8* [[CALLOC]]
+;
+  %ptr = tail call i8* @malloc(i32 %size) #1
+  %memset = tail call i8* @memset(i8* nonnull %ptr, i32 0, i32 %size) #1
+  store i8 1, i8* %ptr           ;; fdata[0] = 1;
+  ret i8* %memset
+}
+
+
+
 attributes #0 = { nounwind ssp uwtable }
 attributes #1 = { nounwind }
 attributes #2 = { nounwind readnone }
-