[LibCallSimplifier] fold memset(malloc(x), 0, x) to calloc(1, x) if memset is intrinsic

[JudyZhu]

Summary

this patch want to fix the bug: https://llvm.org/bugs/show_bug.cgi?id=25892

When there is memset in code, llvm might transform memset to llvm.memset.*.
But llvm might not handle this well, so I use this patch to do it.

test:
check-llvm

[JudyZhu]

ping

[spatel]

Please review the comments in D16337 - especially:
https://reviews.llvm.org/D16337#331117

How does this patch guarantee that memory is not touched between the malloc and the memset? Is the transform safe on this example?

define float* @store_between_malloc_and_memset(i64 %size) {
entry:
  %call = tail call i8* @malloc(i64 %size) #1
  %bc1 = bitcast i8* %call to i32*
  store i32 1, i32* %bc1, align 4
  %cmp = icmp eq i8* %call, null
  br i1 %cmp, label %cleanup, label %if.end

if.end:
  %bc = bitcast i8* %call to float*
  %call3 = tail call i8* @memset(i8* %call, i32 0, i64 %size)
  br label %cleanup

cleanup:
  %r = phi float* [ %bc, %if.end ], [ null, %entry ]
  ret float* %r
}

[JudyZhu]

yes, that might get an error.
let me change it.

In D93360#2466015, @spatel wrote:

Please review the comments in D16337 - especially:
https://reviews.llvm.org/D16337#331117

How does this patch guarantee that memory is not touched between the malloc and the memset? Is the transform safe on this example?

define float* @store_between_malloc_and_memset(i64 %size) {
entry:
  %call = tail call i8* @malloc(i64 %size) #1
  %bc1 = bitcast i8* %call to i32*
  store i32 1, i32* %bc1, align 4
  %cmp = icmp eq i8* %call, null
  br i1 %cmp, label %cleanup, label %if.end

if.end:
  %bc = bitcast i8* %call to float*
  %call3 = tail call i8* @memset(i8* %call, i32 0, i64 %size)
  br label %cleanup

cleanup:
  %r = phi float* [ %bc, %if.end ], [ null, %entry ]
  ret float* %r
}

[xbolva00]

Fixed by @yurai007

llvm/test/Transforms/InstCombine/memset-1.ll

60

Remove fixme

[lebedev.ri]

This review seems to be stuck/dead, consider abandoning if no longer relevant.

llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp

 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Utils/SimplifyLibCalls.h"
 #include "llvm/ADT/APSInt.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/Analysis/BlockFrequencyInfo.h"
+#include "llvm/Analysis/CaptureTracking.h"
 #include "llvm/Analysis/ConstantFolding.h"
+#include "llvm/Analysis/Loads.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/Analysis/ProfileSummaryInfo.h"
-#include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Analysis/ValueTracking.h"
-#include "llvm/Analysis/CaptureTracking.h"
-#include "llvm/Analysis/Loads.h"
+#include "llvm/IR/Constants.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"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/KnownBits.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Transforms/Utils/BuildLibCalls.h"
+#include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/SizeOpts.h"
 
 using namespace llvm;
 using namespace PatternMatch;
 
 static cl::opt<bool>
     EnableUnsafeFPShrink("enable-double-float-shrink", cl::Hidden,
                          cl::init(false),
Value *LibCallSimplifier::optimizeMemMove(CallInst *CI, IRBuilderBase &B) {
   CallInst *NewCI = B.CreateMemMove(CI->getArgOperand(0), Align(1),
                                     CI->getArgOperand(1), Align(1), Size);
   NewCI->setAttributes(CI->getAttributes());
   return CI->getArgOperand(0);
 }
 
 /// Fold memset[_chk](malloc(n), 0, n) --> calloc(1, n).
 Value *LibCallSimplifier::foldMallocMemset(CallInst *Memset, IRBuilderBase &B) {
+  if (isa<IntrinsicInst>(Memset)) {
+    bool canBeFold = false;
+    do {
+      if (auto *MemSetInstr = dyn_cast_or_null<MemSetInst>(Memset)) {
+        if (MemSetInstr->getVolatileCst() ==
+            ConstantInt::getTrue(Memset->getContext()))
+          break;
+
+        if (MemSetInstr->getRawDest()->getType()->getPointerAddressSpace() != 0)
+          break;
+
+        canBeFold = true;
+      }
+    } while (0);
+
+    if (!canBeFold)
+      return nullptr;
+  }
+
   // 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 CanFoldMallocFn = [](CallInst *Malloc, CallInst *Memset) {
+    if (Malloc->hasOneUse())
+      return true;
+
+    for (Use &U : Malloc->uses()) {
+      User *I = U.getUser();
+
+      // skip memset function itself
+      if (I == Memset)
+          continue;
+
+      // considering the following ir:
+      //
+      // %call = call noalias i8* @malloc(i64 %conv) #4
+      // %cmp = icmp eq i8* %call, null
+      //
+      // when the call malloc to allocate virtual memory,
+      // it might have null-check
+      if (CmpInst *Cmp = dyn_cast_or_null<CmpInst>(I)) {
+        uint64_t index = U.getOperandNo() == 0 ? 1 : 0;
+        auto *Ops = Cmp->getOperand(index);
+        auto CurPredicate = Cmp->getPredicate();
+        if ((CurPredicate != ICmpInst::ICMP_EQ &&
+             CurPredicate != ICmpInst::ICMP_NE) ||
+            Ops != Constant::getNullValue(Ops->getType()))
+          return false;
+        continue;
+      }
+
+      // if the code satisfied the following ir:
+      //
+      // %call = call noalias i8* @malloc(i64 %conv) #4
+      // %1 = bitcast i8* %call to float*
+      //
+      // malloc's return value is void *, so it might have a bitcast
+      if (!isa<BitCastInst>(I))
+        return false;
+    }
+    return true;
+  };
+
   // 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 || !CanFoldMallocFn(Malloc, Memset))
     return nullptr;
 
   // Is the inner call really malloc()?
   Function *InnerCallee = Malloc->getCalledFunction();
   if (!InnerCallee)
     return nullptr;
 
   LibFunc Func;
Value *LibCallSimplifier::foldMallocMemset(CallInst *Memset, IRBuilderBase &B) {
   }
 
   return nullptr;
 }
 
 Value *LibCallSimplifier::optimizeMemSet(CallInst *CI, IRBuilderBase &B) {
   Value *Size = CI->getArgOperand(2);
   annotateNonNullAndDereferenceable(CI, 0, Size, DL);
-  if (isa<IntrinsicInst>(CI))
-    return nullptr;
 
   if (auto *Calloc = foldMallocMemset(CI, B))
     return Calloc;
 
+  if (isa<IntrinsicInst>(CI))
+    return nullptr;
+
   // memset(p, v, n) -> llvm.memset(align 1 p, v, n)
   Value *Val = B.CreateIntCast(CI->getArgOperand(1), B.getInt8Ty(), false);
   CallInst *NewCI = B.CreateMemSet(CI->getArgOperand(0), Val, Size, Align(1));
   NewCI->setAttributes(CI->getAttributes());
   return CI->getArgOperand(0);
 }
 
 Value *LibCallSimplifier::optimizeRealloc(CallInst *CI, IRBuilderBase &B) {

llvm/test/Transforms/InstCombine/memset-1.ll

 ; CHECK-NEXT:    call void @llvm.memset.p0i8.i32(i8* align 1 [[CALL1]], i8 0, i32 [[SIZE]], i1 false) #0
 ; 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.
 
 define float* @pr25892(i32 %size) #0 {
 ; CHECK-LABEL: @pr25892(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[CALL:%.*]] = tail call i8* @malloc(i32 [[SIZE:%.*]]) #0
+; CHECK-NEXT:    [[CALL:%.*]] = call i8* @calloc(i32 1, i32 [[SIZE:%.*]])
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8* [[CALL]], 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) #0
 ; 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
;
   %memset = tail call i8* @memset(i8* nonnull %ptr, i32 0, i32 %size) #1
   ret i8* %memset
 }
 
 
 attributes #0 = { nounwind ssp uwtable }
 attributes #1 = { nounwind }
 attributes #2 = { nounwind readnone }
-

llvm/test/Transforms/InstCombine/promot_malloc_and_memset_intrinsic.ll

+; RUN: opt < %s -O2 -S | FileCheck %s
+; RUN: opt < %s -O3 -S | FileCheck %s
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+%struct._IO_FILE = type { i32, i8*, i8*, i8*, i8*, i8*, i8*, i8*, i8*, i8*, i8*, i8*, %struct._IO_marker*, %struct._IO_FILE*, i32, i32, i64, i16, i8, [1 x i8], i8*, i64, i8*, i8*, i8*, i8*, i64, i32, [20 x i8] }
+%struct._IO_marker = type { %struct._IO_marker*, %struct._IO_FILE*, i32 }
+
+@stderr = external dso_local global %struct._IO_FILE*, align 8
+@.str = private unnamed_addr constant [28 x i8] c"Can't allocate data array.\0A\00", align 1
+
+; check the case1:
+;  1. malloc & mmeset
+;  2. null-check
+;  4. bitcast
+; Function Attrs: noinline nounwind uwtable
+define dso_local float* @foo(i32 %size) #0 {
+entry:
+; CHECK: %calloc = tail call i8* @calloc(i64 1, i64 %conv)
+; CHECK-NEXT: %cmp = icmp eq i8* %calloc, null
+
+  %size.addr = alloca i32, align 4
+  %fdata = alloca float*, align 8
+  store i32 %size, i32* %size.addr, align 4
+  %0 = load i32, i32* %size.addr, align 4
+  %conv = sext i32 %0 to i64
+  %call = call noalias i8* @malloc(i64 %conv) #5
+  %1 = bitcast i8* %call to float*
+  store float* %1, float** %fdata, align 8
+  %2 = load float*, float** %fdata, align 8
+  %cmp = icmp eq float* %2, null
+  br i1 %cmp, label %if.then, label %if.end
+
+if.then:                                          ; preds = %entry
+  %3 = load %struct._IO_FILE*, %struct._IO_FILE** @stderr, align 8
+  %call2 = call i32 (%struct._IO_FILE*, i8*, ...) @fprintf(%struct._IO_FILE* %3, i8* getelementptr inbounds ([28 x i8], [28 x i8]* @.str, i64 0, i64 0))
+  call void @exit(i32 1) #6
+  unreachable
+
+if.end:                                           ; preds = %entry
+; CHECK: %2 = bitcast i8* %calloc to float*
+; CHECK-NEXT: ret float* %2
+  %4 = load float*, float** %fdata, align 8
+  %5 = bitcast float* %4 to i8*
+  %6 = load i32, i32* %size.addr, align 4
+  %conv3 = sext i32 %6 to i64
+  call void @llvm.memset.p0i8.i64(i8* align 4 %5, i8 0, i64 %conv3, i1 false)
+  %7 = load float*, float** %fdata, align 8
+  ret float* %7
+}
+
+; Function Attrs: nounwind
+declare dso_local noalias i8* @malloc(i64) #1
+
+declare dso_local i32 @fprintf(%struct._IO_FILE*, i8*, ...) #2
+
+; Function Attrs: noreturn nounwind
+declare dso_local void @exit(i32) #3
+
+; Function Attrs: argmemonly nofree nosync nounwind willreturn writeonly
+declare void @llvm.memset.p0i8.i64(i8* nocapture writeonly, i8, i64, i1 immarg) #4
+
+attributes #0 = { noinline nounwind  uwtable "disable-tail-calls"="false" "frame-pointer"="all" "less-precise-fpmad"="false" "min-legal-vector-width"="0" "no-infs-fp-math"="false" "no-jump-tables"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" "unsafe-fp-math"="false" "use-soft-float"="false" }
+attributes #1 = { nounwind "disable-tail-calls"="false" "frame-pointer"="all" "less-precise-fpmad"="false" "no-infs-fp-math"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" "unsafe-fp-math"="false" "use-soft-float"="false" }
+attributes #2 = { "disable-tail-calls"="false" "frame-pointer"="all" "less-precise-fpmad"="false" "no-infs-fp-math"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" "unsafe-fp-math"="false" "use-soft-float"="false" }
+attributes #3 = { noreturn nounwind "disable-tail-calls"="false" "frame-pointer"="all" "less-precise-fpmad"="false" "no-infs-fp-math"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" "unsafe-fp-math"="false" "use-soft-float"="false" }
+attributes #4 = { argmemonly nofree nosync nounwind willreturn writeonly }
+attributes #5 = { nounwind }
+attributes #6 = { noreturn nounwind }
 No newline at end of file