[LibCallSimplifier] fold memset(malloc(x), 0, x) --> calloc(1, x)

[spatel]

Summary

This is a step towards solving PR25892:
https://llvm.org/bugs/show_bug.cgi?id=25892

It won't handle the reported case. As noted by the 'TODO' comments in the patch, we need to relax the hasOneUse() constraint and also match patterns that include memset_chk() and the llvm.memset() intrinsic in addition to memset().

[mcrosier]

What happens if there's an intervening store to the malloc'ed memory? (dumb I know.. but..)

Slightly modified version of the test case in memset_chk-1.ll:

define float* @pr25892(i64 %size) #0 {
entry:

%call = tail call i8* @malloc(i64 %size) #1
%0 = bitcast i8* %call to i32*        ;; 
store i32 1, i32* %0, align 4           ;; fdata[0] = 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 i64 @llvm.objectsize.i64.p0i8(i8* nonnull %call, i1 false)
%call3 = tail call i8* @__memset_chk(i8* nonnull %call, i32 0, i64 %size, i64 %call2) #1
br label %cleanup

cleanup:

%retval.0 = phi float* [ %bc, %if.end ], [ null, %entry ]
ret float* %retval.0

}

This does get transformed incorrectly with this patch.

Chad

lib/Transforms/Utils/SimplifyLibCalls.cpp

978

Shouldn't you check that Calloc is non-null here?

[spatel]

In D16337#331118, @mcrosier wrote:

What happens if there's an intervening store to the malloc'ed memory? (dumb I know.. but..)

Ah, good catch. I checked 'hasOneUse()' in an earlier draft of this to avoid that scenario, but then I noticed that the example in the bug report had that 'llvm.objectsize' and other uses between the malloc and memset, so I removed that condition. I need to find a way around that.

lib/Transforms/Utils/SimplifyLibCalls.cpp

978

Yes - that won't work if calloc doesn't exist for the target.

[spatel]

Patch updated: now with less ambition and more TODOs!

1. Restrict the optimization to hasOneUse() of the malloc. This isn't going to fire for most real code AFAICT because neither the llvm.memset() intrinsic nor the memset_chk() libcall cases will qualify. I don't know how to avoid the intermediate store problem noted by Chad yet, so take a baby step. Sorry PR25892.

2. The earlier draft added a DataLayout member to FortifiedLibCallSimplifier, but we can just pull that from the module and use it as needed. This way we don't have to change the constructor and affect unrelated code (CodeGenPrepare).

3. Fixed to make sure that calloc() exists and is available for the transform; if not, bail.

[mcrosier]

This looks to be in pretty good shape, but I wouldn't mind hearing other's opinions.

lib/Transforms/Utils/SimplifyLibCalls.cpp

942

Would it make more sense to check for calloc here?

if (!TLI.getLibFunc("calloc", Func) || !TLI.has(Func))

return nullptr;

It would avoid some unnecessary work in the event we don't have calloc. However, I assume we generally have calloc, so this might not be a great suggestion... Just a suggestion that doesn't necessarily have to be acted upon.

[spatel]

lib/Transforms/Utils/SimplifyLibCalls.cpp

942

If we view this in isolation, that sounds reasonable, but the calloc existence check is in emitCalloc because that's the pattern used by the other emit* functions in BuildLibCalls.

I'd prefer to leave this as-is pending resolution of the 'TODO' at line 918 about moving emitCalloc() over to BuildLibCalls to be with its siblings (even though that entire separate file is currently unnecessary because it's only ever used here AFAICT).

It's also possible (there's always a chance!) that someone will come up with some other transform here in LibCallSimplifier that needs a calloc(). In that case, I think it'd also be better to leave the existence check where it is.

[mcrosier]

LGTM assuming you've done the necessary correctness/performance due diligence.

lib/Transforms/Utils/SimplifyLibCalls.cpp

942

All very reasonable comments, which is why I left this to your discretion.

[spatel]

In D16337#334992, @mcrosier wrote:

LGTM assuming you've done the necessary correctness/performance due diligence.

Thanks, Chad. I ran test-suite and see no fails (and don't expect this to fire anyway due to the 1-use restriction). I'm not able to measure any perf differences even with a micro-benchmark on OSX, but at least the code is smaller.

[davide]

lgtm, Sanjay, and sorry for the delay reviewing this.

lib/Transforms/Utils/SimplifyLibCalls.cpp

if (!checkStringCopyLibFuncSignature(Callee, LibFunc::memmove))
     return nullptr;
 
   // memmove(x, y, n) -> llvm.memmove(x, y, n, 1)
   B.CreateMemMove(CI->getArgOperand(0), CI->getArgOperand(1),
                   CI->getArgOperand(2), 1);
   return CI->getArgOperand(0);
 }
 
+// TODO: Does this belong in BuildLibCalls or should all of those similar
+// functions be moved here?
+static Value *emitCalloc(Value *Num, Value *Size, const AttributeSet &Attrs,
+                         IRBuilder<> &B, const TargetLibraryInfo &TLI) {
+  LibFunc::Func Func;
+  if (!TLI.getLibFunc("calloc", Func) || !TLI.has(Func))
+    return nullptr;
+
+  Module *M = B.GetInsertBlock()->getModule();
+  const DataLayout &DL = M->getDataLayout();
+  IntegerType *PtrType = DL.getIntPtrType((B.GetInsertBlock()->getContext()));
+  Value *Calloc = M->getOrInsertFunction("calloc", Attrs, B.getInt8PtrTy(),
+                                         PtrType, PtrType, nullptr);
+  CallInst *CI = B.CreateCall(Calloc, { Num, Size }, "calloc");
+
+  if (const auto *F = dyn_cast<Function>(Calloc->stripPointerCasts()))
+    CI->setCallingConv(F->getCallingConv());
+
+  return CI;
+}
+
+/// 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())
+    return nullptr;
+
+  // Is the inner call really malloc()?
+  Function *InnerCallee = Malloc->getCalledFunction();
+  LibFunc::Func Func;
+  if (!TLI.getLibFunc(InnerCallee->getName(), Func) || !TLI.has(Func) ||
+      Func != LibFunc::malloc)
+    return nullptr;
+
+  // Matching the name is not good enough. Make sure the parameter and return
+  // type match the standard library signature.
+  FunctionType *FT = InnerCallee->getFunctionType();
+  if (FT->getNumParams() != 1 || !FT->getParamType(0)->isIntegerTy())
+    return nullptr;
+
+  auto *RetType = dyn_cast<PointerType>(FT->getReturnType());
+  if (!RetType || !RetType->getPointerElementType()->isIntegerTy(8))
+    return nullptr;
+
+  // The memset must cover the same number of bytes that are malloc'd.
+  if (Memset->getArgOperand(2) != Malloc->getArgOperand(0))
+    return nullptr;
+
+  // Replace the malloc with a calloc. We need the data layout to know what the
+  // 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);
+  if (!Calloc)
+    return nullptr;
+
+  Malloc->replaceAllUsesWith(Calloc);
+  Malloc->eraseFromParent();
+
+  return Calloc;
+}
+
 Value *LibCallSimplifier::optimizeMemSet(CallInst *CI, IRBuilder<> &B) {
   Function *Callee = CI->getCalledFunction();
 
   if (!checkStringCopyLibFuncSignature(Callee, LibFunc::memset))
     return nullptr;
 
+  if (auto *Calloc = foldMallocMemset(CI, B, *TLI))
+    return Calloc;
+
   // memset(p, v, n) -> llvm.memset(p, v, n, 1)
   Value *Val = B.CreateIntCast(CI->getArgOperand(1), B.getInt8Ty(), false);
   B.CreateMemSet(CI->getArgOperand(0), Val, CI->getArgOperand(2), 1);
   return CI->getArgOperand(0);
 }
 
 //===----------------------------------------------------------------------===//
 // Math Library Optimizations
if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(CI)) {
     case Intrinsic::exp2:
       return optimizeExp2(CI, Builder);
     case Intrinsic::fabs:
       return optimizeFabs(CI, Builder);
     case Intrinsic::log:
       return optimizeLog(CI, Builder);
     case Intrinsic::sqrt:
       return optimizeSqrt(CI, Builder);
+    // TODO: Use foldMallocMemset() with memset intrinsic.
     default:
       return nullptr;
     }
   }
 
   // Also try to simplify calls to fortified library functions.
   if (Value *SimplifiedFortifiedCI = FortifiedSimplifier.optimizeCall(CI)) {
     // Try to further simplify the result.

 
 Value *FortifiedLibCallSimplifier::optimizeMemSetChk(CallInst *CI,
                                                      IRBuilder<> &B) {
   Function *Callee = CI->getCalledFunction();
 
   if (!checkStringCopyLibFuncSignature(Callee, LibFunc::memset_chk))
     return nullptr;
 
+  // TODO: Try foldMallocMemset() here.
+
   if (isFortifiedCallFoldable(CI, 3, 2, false)) {
     Value *Val = B.CreateIntCast(CI->getArgOperand(1), B.getInt8Ty(), false);
     B.CreateMemSet(CI->getArgOperand(0), Val, CI->getArgOperand(2), 1);
     return CI->getArgOperand(0);
   }
   return nullptr;
 }
 

test/Transforms/InstCombine/memset-1.ll

 define i8* @test_simplify1(i8* %mem, i32 %val, i32 %size) {
 ; CHECK-LABEL: @test_simplify1(
   %ret = call i8* @memset(i8* %mem, i32 %val, i32 %size)
 ; CHECK: call void @llvm.memset
   ret i8* %ret
 ; CHECK: ret i8* %mem
 }
 
+define i8* @pr25892_lite(i32 %size) #0 {
+  %call1 = call i8* @malloc(i32 %size) #1
+  %call2 = call i8* @memset(i8* %call1, i32 0, i32 %size) #1
+  ret i8* %call2
+
+; CHECK-LABEL: @pr25892_lite(
+; CHECK-NEXT:  %calloc = call i8* @calloc(i32 1, i32 %size)
+; CHECK-NEXT:  ret i8* %calloc
+}
+
 ; FIXME: memset(malloc(x), 0, x) -> calloc(1, x)
+; This doesn't fire currently because the malloc has more than one use.
 
 define float* @pr25892(i32 %size) #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*