[GlobalOpt] [MemoryBuiltins] Update getMallocType with bitcast on stored pointer

[scui]

Summary

Currently, getMallocType checks the bitcast uses of the malloc call. If there is no bitcast uses of the malloc call, malloc function’s return type is returned. Sometimes, the bitcast is on the stored-to pointer instead of on the malloc call. This patch is to also consider the cases like this:

%call = call i8* @malloc
store i8* %call, i8** bitcast (%T** @g to i8**)

[efriedma]

If you going to mess with this code, can you just kill off all the malloc use type heuristics, and just unconditionally create an [i8 x n] array? The existing type sniffing code won't work with opaque pointers anyway.

[scui]

In D107574#2929159, @efriedma wrote:

If you going to mess with this code, can you just kill off all the malloc use type heuristics, and just unconditionally create an [i8 x n] array? The existing type sniffing code won't work with opaque pointers anyway.

The existing code has been there for a while. It tries to get the type for the code in the following pattern:

%call = call i8* @malloc
 %0 = bitcast i8* to T*
 store T* %0, %T** @g

This somehow can be transformed to the code:

%call = call i8* @malloc
store i8* %call, i8** bitcast (%T** @g to i8**)

The code here is trying to recognize the code pattern.

It looks to me the IR code will be sub-optimal for some cases if we simply return Type i8* for getMallocType. I’m thinking maybe we can return just either the only bitcasted type or the malloc return type, ie, no nullptr return, but I’m not sure if this really matters for real cases.

[efriedma]

In D107574#2929796, @scui wrote:

It looks to me the IR code will be sub-optimal for some cases if we simply return Type i8* for getMallocType.

Will it? Almost all optimizations can deal with bitcasts of global variables without getting confused.

Anyway, if you want to refine the type of the global created by the transform, I guess that's okay, but confusion about the types involved shouldn't ever block the transform.

[scui]

In D107574#2929921, @efriedma wrote:

In D107574#2929796, @scui wrote:

It looks to me the IR code will be sub-optimal for some cases if we simply return Type i8* for getMallocType.

Will it? Almost all optimizations can deal with bitcasts of global variables without getting confused.

I could be wrong. One case I can imagine is to set initial value, for example, an i32 value for [4 x i8] global.

Anyway, if you want to refine the type of the global created by the transform, I guess that's okay, but confusion about the types involved shouldn't ever block the transform.

I'm also confused with the nullptr return with the existing code. I'll try another version to limit to looking into only the store from the malloc call, and see if that will make more sense.

[scui]

Rewrite getMallocType - only look at the bitcast feeding into the store (value or pointer) for malloc call, and no longer return nullptr.

[efriedma]

At this point, the code is really specific to GlobalOpt transform; maybe we want to move it into GlobalOpt.cpp?

llvm/lib/Analysis/MemoryBuiltins.cpp

354

Update here as well?

366

We still end up blocking the transform in some cases: if the malloc size isn't a multiple of the type inferred by getMallocType(), the transform fails.

[nikic]

Would it be possible to drop getMallocType, only keep getMallocAllocatedType, and determine the type based on the value type of a (possibly bitcasted) store to the address? That would be forward compatible with opaque pointers.

[nikic]

D117092 implements @efriedma's suggestion of always using an [i8 x Size] type.

[reames]

Appears to have been subsumed by other changes. Please close or rebase if there's still something missing.

llvm/include/llvm/Analysis/MemoryBuiltins.h

 extractMallocCall(const Value *I,
                   function_ref<const TargetLibraryInfo &(Function &)> GetTLI);
 inline CallInst *
 extractMallocCall(Value *I,
                   function_ref<const TargetLibraryInfo &(Function &)> GetTLI) {
   return const_cast<CallInst *>(extractMallocCall((const Value *)I, GetTLI));
 }
 
-/// getMallocType - Returns the PointerType resulting from the malloc call.
-/// The PointerType depends on the number of bitcast uses of the malloc call:
-///   0: PointerType is the malloc calls' return type.
-///   1: PointerType is the bitcast's result type.
-///  >1: Unique PointerType cannot be determined, return NULL.
+/// getMallocType - Returns the PointerType resulting from the malloc call. The
+/// PointerType is the bitcast's result type of the call feeding into store. If
+/// there is no bitcast feeding into store, the PointerType is the call's return
+/// type.
 PointerType *getMallocType(const CallInst *CI, const TargetLibraryInfo *TLI);
 
-/// getMallocAllocatedType - Returns the Type allocated by malloc call.
-/// The Type depends on the number of bitcast uses of the malloc call:
-///   0: PointerType is the malloc calls' return type.
-///   1: PointerType is the bitcast's result type.
-///  >1: Unique PointerType cannot be determined, return NULL.
+/// getMallocAllocatedType - Returns the Type allocated by malloc call. The Type
+/// is the bitcast's result type of the call feeding into store. If there is no
+/// bitcast feeding into store, the Type is the call's return type.
 Type *getMallocAllocatedType(const CallInst *CI, const TargetLibraryInfo *TLI);
 
 /// getMallocArraySize - Returns the array size of a malloc call.  If the
 /// argument passed to malloc is a multiple of the size of the malloced type,
 /// then return that multiple.  For non-array mallocs, the multiple is
 /// constant 1.  Otherwise, return NULL for mallocs whose array size cannot be
 /// determined.
 Value *getMallocArraySize(CallInst *CI, const DataLayout &DL,

llvm/lib/Analysis/MemoryBuiltins.cpp

 
 static Value *computeArraySize(const CallInst *CI, const DataLayout &DL,
                                const TargetLibraryInfo *TLI,
                                bool LookThroughSExt = false) {
   if (!CI)
     return nullptr;
 
   // The size of the malloc's result type must be known to determine array size.
   Type *T = getMallocAllocatedType(CI, TLI);
   if (!T || !T->isSized())
     return nullptr;
 
   unsigned ElementSize = DL.getTypeAllocSize(T);
   if (StructType *ST = dyn_cast<StructType>(T))
     ElementSize = DL.getStructLayout(ST)->getSizeInBytes();
 
   // If malloc call's arg can be determined to be a multiple of ElementSize,
   // return the multiple.  Otherwise, return NULL.
   Value *MallocArg = CI->getArgOperand(0);
   Value *Multiple = nullptr;
   if (ComputeMultiple(MallocArg, ElementSize, Multiple, LookThroughSExt))
     return Multiple;
 
   return nullptr;
 }
 
 /// getMallocType - Returns the PointerType resulting from the malloc call.
-/// The PointerType depends on the number of bitcast uses of the malloc call:
-///   0: PointerType is the calls' return type.
-///   1: PointerType is the bitcast's result type.
-///  >1: Unique PointerType cannot be determined, return NULL.
+/// The PointerType is the bitcast's result type of the call feeding into store.
+/// If there is no bitcast feeding into store, the PointerType is the call's
+/// return type.
 PointerType *llvm::getMallocType(const CallInst *CI,
                                  const TargetLibraryInfo *TLI) {
   assert(isMallocLikeFn(CI, TLI) && "getMallocType and not malloc call");
 
-  PointerType *MallocType = nullptr;
-  unsigned NumOfBitCastUses = 0;
-
-  // Determine if CallInst has a bitcast use.
-  for (const User *U : CI->users())
-    if (const BitCastInst *BCI = dyn_cast<BitCastInst>(U)) {
-      MallocType = cast<PointerType>(BCI->getDestTy());
-      NumOfBitCastUses++;
-    }
+  // Determine if CallInst has a bitcast feeding into store, either with value
+  // or pointer operand. The store is the only store for the malloc call.
+  for (const User *U : CI->users()) {
+    if (const BitCastInst *BCI = dyn_cast<BitCastInst>(U))
+      if (any_of(BCI->uses(), [&](const Use &U) {
+            return isa<StoreInst>(U.getUser()) && U.getOperandNo() == 0;
+          }))
+        return cast<PointerType>(BCI->getDestTy());
+
+    if (auto *ST = dyn_cast<StoreInst>(U))
+      if (ST->getValueOperand() == CI)
+        if (auto *BCO = dyn_cast<BitCastOperator>(ST->getPointerOperand()))
+          if (auto *PTy = dyn_cast<PointerType>(
+                  cast<PointerType>(BCO->getOperand(0)->getType())
+                      ->getElementType()))
+            return PTy;
 
-  // Malloc call has 1 bitcast use, so type is the bitcast's destination type.
-  if (NumOfBitCastUses == 1)
-    return MallocType;
-
   // Malloc call was not bitcast, so type is the malloc function's return type.
-  if (NumOfBitCastUses == 0)
-    return cast<PointerType>(CI->getType());
-
-  // Type could not be determined.
-  return nullptr;
 }
 
-/// getMallocAllocatedType - Returns the Type allocated by malloc call.
-/// The Type depends on the number of bitcast uses of the malloc call:
-///   0: PointerType is the malloc calls' return type.
-///   1: PointerType is the bitcast's result type.
-///  >1: Unique PointerType cannot be determined, return NULL.
+/// getMallocAllocatedType - Returns the Type allocated by malloc call. The Type
+/// is the bitcast's result type of the call feeding into store. If there is no
+/// bitcast feeding into store, the Type is the call's return type.
 Type *llvm::getMallocAllocatedType(const CallInst *CI,
                                    const TargetLibraryInfo *TLI) {
   PointerType *PT = getMallocType(CI, TLI);
-  return PT ? PT->getElementType() : nullptr;
+  assert(PT && "Expect non-null pointer type");
+  return PT->getElementType();
 }
 
 /// getMallocArraySize - Returns the array size of a malloc call.  If the
 /// argument passed to malloc is a multiple of the size of the malloced type,
 /// then return that multiple.  For non-array mallocs, the multiple is
 /// constant 1.  Otherwise, return NULL for mallocs whose array size cannot be
 /// determined.
 Value *llvm::getMallocArraySize(CallInst *CI, const DataLayout &DL,

llvm/lib/Transforms/IPO/GlobalOpt.cpp

if (Constant *SOVC = dyn_cast<Constant>(StoredOnceVal)) {
         SOVC = ConstantExpr::getBitCast(SOVC, GV->getInitializer()->getType());
 
       // Optimize away any trapping uses of the loaded value.
       if (OptimizeAwayTrappingUsesOfLoads(GV, SOVC, DL, GetTLI))
         return true;
     } else if (CallInst *CI = extractMallocCall(StoredOnceVal, GetTLI)) {
       auto *TLI = &GetTLI(*CI->getFunction());
       Type *MallocType = getMallocAllocatedType(CI, TLI);
-      if (MallocType && tryToOptimizeStoreOfMallocToGlobal(GV, CI, MallocType,
-                                                           Ordering, DL, TLI))
+      if (tryToOptimizeStoreOfMallocToGlobal(GV, CI, MallocType, Ordering, DL,
+                                             TLI))
         return true;
     }
   }
 
   return false;
 }
 
 /// At this point, we have learned that the only two values ever stored into GV

llvm/test/Transforms/GlobalOpt/malloc-promote-5.ll

+; RUN: opt -globalopt -S < %s | FileCheck %s
+; RUN: opt -passes=globalopt -S < %s | FileCheck %s
+
+@g = internal global i32* null, align 8
+; CHECK: [[G_BODY:@.*]] =  internal unnamed_addr global i32 undef
+
+define signext i32 @f() {
+; CHECK-LABEL: @f(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    store i32 1, i32* [[G_BODY]], align 4
+; CHECK-NEXT:    store i8 2, i8* bitcast (i32* [[G_BODY]] to i8*), align 4
+; CHECK-NEXT:    ret i32 1
+;
+entry:
+  %call = call i8* @malloc(i64 4)
+  store i8* %call, i8** bitcast (i32** @g to i8**), align 8
+  %0 = load i32*, i32** @g, align 8
+  store i32 1, i32* %0, align 4
+  %1 = load i8*, i8** bitcast (i32** @g to i8**), align 8
+  store i8 2, i8* %1, align 4
+  ret i32 1
+}
+
+define signext i32 @main() {
+; CHECK-LABEL: @main(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CALL:%.*]] = call signext i32 @f()
+; CHECK-NEXT:    [[TMP0:%.*]] = load i32, i32* [[G_BODY]], align 4
+; CHECK-NEXT:    ret i32 [[TMP0]]
+;
+entry:
+  %call = call signext i32 @f()
+  %0 = load i32*, i32** @g, align 8
+  %1 = load i32, i32* %0, align 4
+  ret i32 %1
+}
+
+declare noalias align 16 i8* @malloc(i64)