[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.

include/llvm/Transforms/Utils/SimplifyLibCalls.h

 /// \brief This class implements simplifications for calls to fortified library
 /// functions (__st*cpy_chk, __memcpy_chk, __memmove_chk, __memset_chk), to,
 /// when possible, replace them with their non-checking counterparts.
 /// Other optimizations can also be done, but it's possible to disable them and
 /// only simplify needless use of the checking versions (when the object size
 /// is unknown) by passing true for OnlyLowerUnknownSize.
 class FortifiedLibCallSimplifier {
 private:
+  const DataLayout &DL;
   const TargetLibraryInfo *TLI;
   bool OnlyLowerUnknownSize;
 
 public:
-  FortifiedLibCallSimplifier(const TargetLibraryInfo *TLI,
+  FortifiedLibCallSimplifier(const DataLayout &DL, const TargetLibraryInfo *TLI,
                              bool OnlyLowerUnknownSize = false);
 
   /// \brief Take the given call instruction and return a more
   /// optimal value to replace the instruction with or 0 if a more
   /// optimal form can't be found.
   /// The call must not be an indirect call.
   Value *optimizeCall(CallInst *CI);
 

lib/CodeGen/CodeGenPrepare.cpp

bool CodeGenPrepare::optimizeCallInst(CallInst *CI, bool& ModifiedDT) {
 
   // From here on out we're working with named functions.
   if (!CI->getCalledFunction()) return false;
 
   // Lower all default uses of _chk calls.  This is very similar
   // to what InstCombineCalls does, but here we are only lowering calls
   // to fortified library functions (e.g. __memcpy_chk) that have the default
   // "don't know" as the objectsize.  Anything else should be left alone.
-  FortifiedLibCallSimplifier Simplifier(TLInfo, true);
+  FortifiedLibCallSimplifier Simplifier(*DL, TLInfo, true);
   if (Value *V = Simplifier.optimizeCall(CI)) {
     CI->replaceAllUsesWith(V);
     CI->eraseFromParent();
     return true;
   }
   return false;
 }
 

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 DataLayout &DL,
+                         const TargetLibraryInfo &TLI) {
+  if (!TLI.has(LibFunc::calloc))
+    return nullptr;
+
+  Module *M = B.GetInsertBlock()->getModule();
+  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,
+                               const DataLayout &DL) {
+  // This has to be a memset of zeros (bzero).
+  auto *FillValue = dyn_cast<ConstantInt>(Memset->getArgOperand(1));
+  if (!FillValue || FillValue->getZExtValue() != 0)
+    return nullptr;
+
+  auto *Malloc = dyn_cast<CallInst>(Memset->getArgOperand(0));
+  if (!Malloc)
+    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.
+  B.SetInsertPoint(Malloc->getParent(), ++Malloc->getIterator());
+  IntegerType *SizeType = DL.getIntPtrType(B.GetInsertBlock()->getContext());
+  Value *Calloc = emitCalloc(ConstantInt::get(SizeType, 1),
+                             Malloc->getArgOperand(0), Malloc->getAttributes(),
+                             B, DL, TLI);
+  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;
 
+  // TODO: Use foldMallocMemset() here.
+
   // 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.
if (TLI->getLibFunc(FuncName, Func) && TLI->has(Func)) {
     }
   }
   return nullptr;
 }
 
 LibCallSimplifier::LibCallSimplifier(
     const DataLayout &DL, const TargetLibraryInfo *TLI,
     function_ref<void(Instruction *, Value *)> Replacer)
-    : FortifiedSimplifier(TLI), DL(DL), TLI(TLI), UnsafeFPShrink(false),
+    : FortifiedSimplifier(DL, TLI), DL(DL), TLI(TLI), UnsafeFPShrink(false),
       Replacer(Replacer) {}
 
 void LibCallSimplifier::replaceAllUsesWith(Instruction *I, Value *With) {
   // Indirect through the replacer used in this instance.
   Replacer(I, With);
 }
 
 // TODO:

 
 Value *FortifiedLibCallSimplifier::optimizeMemSetChk(CallInst *CI,
                                                      IRBuilder<> &B) {
   Function *Callee = CI->getCalledFunction();
 
   if (!checkStringCopyLibFuncSignature(Callee, LibFunc::memset_chk))
     return nullptr;
 
+  if (Value *Calloc = foldMallocMemset(CI, B, *TLI, DL))
+    return Calloc;
+
   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;
 }
 
case LibFunc::strncpy_chk:
     return optimizeStrpNCpyChk(CI, Builder, Func);
   default:
     break;
   }
   return nullptr;
 }
 
 FortifiedLibCallSimplifier::FortifiedLibCallSimplifier(
-    const TargetLibraryInfo *TLI, bool OnlyLowerUnknownSize)
-    : TLI(TLI), OnlyLowerUnknownSize(OnlyLowerUnknownSize) {}
+    const DataLayout &DL, const TargetLibraryInfo *TLI, bool OnlyLowerUnknownSize)
+    : DL(DL), TLI(TLI), OnlyLowerUnknownSize(OnlyLowerUnknownSize) {}

test/Transforms/InstCombine/memset_chk-1.ll

 
 declare i8* @__memmove_chk(i8*, i8*, i64, i64)
 declare i8* @strrchr(i8*, i32)
 declare i64 @strlen(i8* nocapture)
 declare i64 @llvm.objectsize.i64.p0i8(i8*, i1)
 
 declare i8* @__memset_chk(i8*, i32, i64, i64)
 
-; FIXME: memset(malloc(x), 0, x) -> calloc(1, x)
-
 define float* @pr25892(i64 %size) #0 {
 entry:
   %call = tail call i8* @malloc(i64 %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 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
 
 ; CHECK-LABEL: @pr25892(
 ; CHECK:       entry:
-; CHECK-NEXT:    %call = tail call i8* @malloc(i64 %size)
-; CHECK-NEXT:    %cmp = icmp eq i8* %call, null
+; CHECK-NEXT:    %calloc = call i8* @calloc(i64 1, i64 %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:    %call2 = tail call i64 @llvm.objectsize.i64.p0i8(i8* nonnull %call, i1 false)
-; CHECK-NEXT:    %call3 = tail call i8* @__memset_chk(i8* nonnull %call, i32 0, i64 %size, i64 %call2)
+; 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
 }
 
 declare noalias i8* @malloc(i64) #1
 
 attributes #0 = { nounwind ssp uwtable }
 attributes #1 = { nounwind }
 attributes #2 = { nounwind readnone }