[Attributor] Look through allocated heap memory

[jdoerfert]

Summary

AAPointerInfo, and thereby other places, can look already through
internal global and stack memory. This patch enables them to look
through heap memory returned by functions with a noalias return.

In the future we can look through noalias arguments as well but that
will require AAIsDead to learn that such memory can be inspected by the
caller later on. We also need teach AAPointerInfo about dominance to
actually deal with memory that might not be null or undef
initialized. D106397 is a first step in that direction already.

[jdoerfert]

Add tests

[kuter]

Note: Some tests seem broken (heap_to_stack.ll) and clang format is needed.

Also, do you think we can optimize nested stuff ? for example if a pointer was stored to memory.
I think checkForAllUses can look through values that escape to memory, right ?

llvm/lib/Transforms/IPO/Attributor.cpp

210

I don't think this check is needed. isNoAliasFn seems to always return true if this is a allocation function.

Source code for isNoAliasFn

/// Tests if a value is a call or invoke to a function that returns a
/// NoAlias pointer (including malloc/calloc/realloc/strdup-like functions).
bool llvm::isNoAliasFn(const Value *V, const TargetLibraryInfo *TLI,
                       bool LookThroughBitCast) {
  // it's safe to consider realloc as noalias since accessing the original
  // pointer is undefined behavior
  return isAllocationFn(V, TLI, LookThroughBitCast) ||
         hasNoAliasAttr(V, LookThroughBitCast);
}

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

2344

Do you think we should use the AAReachability here ?
If the function can reach itself, then it is recursive.

[jdoerfert]

In D109170#2988599, @kuter wrote:

Also, do you think we can optimize nested stuff ? for example if a pointer was stored to memory.
I think checkForAllUses can look through values that escape to memory, right ?

It "works" but I discovered that we need one more piece that I will put in here or separately.
Basically, we need to connect the use that is stored with the copies that are loaded.

llvm/lib/Transforms/IPO/Attributor.cpp

210

I don't understand why this would not be needed.

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

2344

Yes, I'm looking into that already.

[jdoerfert]

Update test

[jdoerfert]

Add a (currently failing) nested memory test

[kuter]

LGTM

llvm/include/llvm/Transforms/IPO/Attributor.h

 ///        X + none  => X
 /// not_none + undef => not_none
 ///          V1 + V2 => nullptr
 Optional<Value *>
 combineOptionalValuesInAAValueLatice(const Optional<Value *> &A,
                                      const Optional<Value *> &B, Type *Ty);
 
 /// Return the initial value of \p Obj with type \p Ty if that is a constant.
-Constant *getInitialValueForObj(Value &Obj, Type &Ty);
+Constant *getInitialValueForObj(Value &Obj, Type &Ty,
+                                const TargetLibraryInfo *TLI);
 
 /// Collect all potential underlying objects of \p Ptr at position \p CtxI in
 /// \p Objects. Assumed information is used and dependences onto \p QueryingAA
 /// are added appropriately.
 ///
 /// \returns True if \p Objects contains all assumed underlying objects, and
 ///          false if something went wrong and the objects could not be
 ///          determined.

llvm/lib/Transforms/IPO/Attributor.cpp

 
 #include "llvm/ADT/GraphTraits.h"
 #include "llvm/ADT/PointerIntPair.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/TinyPtrVector.h"
 #include "llvm/Analysis/InlineCost.h"
 #include "llvm/Analysis/LazyValueInfo.h"
+#include "llvm/Analysis/MemoryBuiltins.h"
 #include "llvm/Analysis/MemorySSAUpdater.h"
 #include "llvm/Analysis/MustExecute.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/GlobalVariable.h"
if (auto *A = dyn_cast<Argument>(&V))
     Scope = A->getParent();
   if (!Scope)
     return false;
   auto &NoRecurseAA = A.getAAFor<AANoRecurse>(
       QueryingAA, IRPosition::function(*Scope), DepClassTy::OPTIONAL);
   return NoRecurseAA.isAssumedNoRecurse();
 }
 
-Constant *AA::getInitialValueForObj(Value &Obj, Type &Ty) {
+Constant *AA::getInitialValueForObj(Value &Obj, Type &Ty,
+                                    const TargetLibraryInfo *TLI) {
   if (isa<AllocaInst>(Obj))
     return UndefValue::get(&Ty);
+  if (isNoAliasFn(&Obj, TLI)) {
+    if (isMallocLikeFn(&Obj, TLI) || isAlignedAllocLikeFn(&Obj, TLI))
+      return UndefValue::get(&Ty);
+    if (isCallocLikeFn(&Obj, TLI))
+      return Constant::getNullValue(&Ty);
+    return nullptr;
+  }
   auto *GV = dyn_cast<GlobalVariable>(&Obj);
   if (!GV || !GV->hasLocalLinkage())
     return nullptr;
   if (!GV->hasInitializer())
     return UndefValue::get(&Ty);
   return dyn_cast_or_null<Constant>(getWithType(*GV->getInitializer(), Ty));
 }
 
if (!AA::getAssumedUnderlyingObjects(A, Ptr, Objects, QueryingAA, &SI)) {
     LLVM_DEBUG(
         dbgs() << "Underlying objects stored into could not be determined\n";);
     return false;
   }
 
   SmallVector<const AAPointerInfo *> PIs;
   SmallVector<Value *> NewCopies;
 
+  const auto *TLI =
+      A.getInfoCache().getTargetLibraryInfoForFunction(*SI.getFunction());
   for (Value *Obj : Objects) {
     LLVM_DEBUG(dbgs() << "Visit underlying object " << *Obj << "\n");
     if (isa<UndefValue>(Obj))
       continue;
     if (isa<ConstantPointerNull>(Obj)) {
       // A null pointer access can be undefined but any offset from null may
       // be OK. We do not try to optimize the latter.
       if (!NullPointerIsDefined(SI.getFunction(),
                                 Ptr.getType()->getPointerAddressSpace()) &&
           A.getAssumedSimplified(Ptr, QueryingAA, UsedAssumedInformation) ==
               Obj)
         continue;
       LLVM_DEBUG(
           dbgs() << "Underlying object is a valid nullptr, giving up.\n";);
       return false;
     }
-    if (!isa<AllocaInst>(Obj) && !isa<GlobalVariable>(Obj)) {
+    if (!isa<AllocaInst>(Obj) && !isa<GlobalVariable>(Obj) &&
+        !isNoAliasFn(Obj, TLI)) {
       LLVM_DEBUG(dbgs() << "Underlying object is not supported yet: " << *Obj
                         << "\n";);
       return false;
     }
     if (auto *GV = dyn_cast<GlobalVariable>(Obj))
       if (!GV->hasLocalLinkage()) {
         LLVM_DEBUG(dbgs() << "Underlying object is global with external "
                              "linkage, not supported yet: "

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

ChangeStatus updateImpl(Attributor &A) override {
     auto HandlePassthroughUser = [&](Value *Usr, OffsetInfo &PtrOI,
                                      bool &Follow) {
       OffsetInfo &UsrOI = OffsetInfoMap[Usr];
       UsrOI = PtrOI;
       Follow = true;
       return true;
     };
 
+    const auto *TLI = getAnchorScope()
+                          ? A.getInfoCache().getTargetLibraryInfoForFunction(
+                                *getAnchorScope())
+                          : nullptr;
     auto UsePred = [&](const Use &U, bool &Follow) -> bool {
       Value *CurPtr = U.get();
       User *Usr = U.getUser();
       LLVM_DEBUG(dbgs() << "[AAPointerInfo] Analyze " << *CurPtr << " in "
                         << *Usr << "\n");
 
       OffsetInfo &PtrOI = OffsetInfoMap[CurPtr];
 
auto UsePred = [&](const Use &U, bool &Follow) -> bool {
             *StoreI->getValueOperand(), *this, UsedAssumedInformation);
         return handleAccess(A, *StoreI, *CurPtr, Content, AccessKind::AK_WRITE,
                             PtrOI.Offset, Changed,
                             StoreI->getValueOperand()->getType());
       }
       if (auto *CB = dyn_cast<CallBase>(Usr)) {
         if (CB->isLifetimeStartOrEnd())
           return true;
+        if (TLI && isFreeCall(CB, TLI))
+          return true;
         if (CB->isArgOperand(&U)) {
           unsigned ArgNo = CB->getArgOperandNo(&U);
           const auto &CSArgPI = A.getAAFor<AAPointerInfo>(
               *this, IRPosition::callsite_argument(*CB, ArgNo),
               DepClassTy::REQUIRED);
           Changed = translateAndAddCalleeState(A, CSArgPI, PtrOI.Offset, *CB) |
                     Changed;
           return true;

 
 struct AANoRecurseFunction final : AANoRecurseImpl {
   AANoRecurseFunction(const IRPosition &IRP, Attributor &A)
       : AANoRecurseImpl(IRP, A) {}
 
   /// See AbstractAttribute::initialize(...).
   void initialize(Attributor &A) override {
     AANoRecurseImpl::initialize(A);
     if (const Function *F = getAnchorScope())
       if (A.getInfoCache().getSccSize(*F) != 1)
         indicatePessimisticFixpoint();
   }
 
   /// See AbstractAttribute::updateImpl(...).
   ChangeStatus updateImpl(Attributor &A) override {
 
     // If all live call sites are known to be no-recurse, we are as well.
auto UnionWrapper = [&](Value &V, Value &Obj) {
       return Union(V);
     };
 
     Value &Ptr = *L.getPointerOperand();
     SmallVector<Value *, 8> Objects;
     if (!AA::getAssumedUnderlyingObjects(A, Ptr, Objects, AA, &L))
       return false;
 
+    const auto *TLI =
+        A.getInfoCache().getTargetLibraryInfoForFunction(*L.getFunction());
     for (Value *Obj : Objects) {
       LLVM_DEBUG(dbgs() << "Visit underlying object " << *Obj << "\n");
       if (isa<UndefValue>(Obj))
         continue;
       if (isa<ConstantPointerNull>(Obj)) {
         // A null pointer access can be undefined but any offset from null may
         // be OK. We do not try to optimize the latter.
         bool UsedAssumedInformation = false;
         if (!NullPointerIsDefined(L.getFunction(),
                                   Ptr.getType()->getPointerAddressSpace()) &&
             A.getAssumedSimplified(Ptr, AA, UsedAssumedInformation) == Obj)
           continue;
         return false;
       }
-      if (!isa<AllocaInst>(Obj) && !isa<GlobalVariable>(Obj))
+      if (!isa<AllocaInst>(Obj) && !isa<GlobalVariable>(Obj) &&
+          !isNoAliasFn(Obj, TLI))
         return false;
-      Constant *InitialVal = AA::getInitialValueForObj(*Obj, *L.getType());
+      Constant *InitialVal = AA::getInitialValueForObj(*Obj, *L.getType(), TLI);
       if (!InitialVal || !Union(*InitialVal))
         return false;
 
       LLVM_DEBUG(dbgs() << "Underlying object amenable to load-store "
                            "propagation, checking accesses next.\n");
 
       auto CheckAccess = [&](const AAPointerInfo::Access &Acc, bool IsExact) {
         LLVM_DEBUG(dbgs() << " - visit access " << Acc << "\n");

llvm/test/Transforms/Attributor/heap_to_stack_gpu.ll

 ; CHECK-NEXT:    [[TMP5:%.*]] = load i32, i32* [[TMP2]], align 4
 ; CHECK-NEXT:    [[TMP6:%.*]] = add nsw i32 [[TMP5]], -1
 ; CHECK-NEXT:    store i32 [[TMP6]], i32* [[TMP2]], align 4
 ; CHECK-NEXT:    [[TMP7:%.*]] = icmp sgt i32 [[TMP6]], 0
 ; CHECK-NEXT:    br i1 [[TMP7]], label [[TMP8:%.*]], label [[TMP11:%.*]]
 ; CHECK:       8:
 ; CHECK-NEXT:    [[TMP9:%.*]] = call noalias i8* @malloc(i64 noundef 4)
 ; CHECK-NEXT:    [[TMP10:%.*]] = bitcast i8* [[TMP9]] to i32*
-; CHECK-NEXT:    store i32 1, i32* [[TMP10]], align 8
 ; CHECK-NEXT:    br label [[TMP4]]
 ; CHECK:       11:
 ; CHECK-NEXT:    ret i32 5
 ;
   %2 = alloca i32, align 4
   %3 = alloca i32*, align 8
   store i32 %0, i32* %2, align 4
   br label %4

llvm/test/Transforms/Attributor/memory_locations.ll

 
 return:                                           ; preds = %if.end, %if.then
   %retval.0 = phi i8* [ %call, %if.then ], [ %call1, %if.end ]
   ret i8* %retval.0
 }
 
 define dso_local i8* @internal_only_rec_static_helper_malloc_noescape(i32 %arg) {
 ; FIXME: This is actually inaccessiblememonly because the malloced memory does not escape
-; CHECK-LABEL: define {{[^@]+}}@internal_only_rec_static_helper_malloc_noescape
-; CHECK-SAME: (i32 [[ARG:%.*]]) {
-; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[CALL:%.*]] = call noalias i8* @internal_only_rec_static_malloc_noescape(i32 [[ARG]])
-; CHECK-NEXT:    ret i8* [[CALL]]
+; IS__TUNIT____-LABEL: define {{[^@]+}}@internal_only_rec_static_helper_malloc_noescape
+; IS__TUNIT____-SAME: (i32 [[ARG:%.*]]) {
+; IS__TUNIT____-NEXT:  entry:
+; IS__TUNIT____-NEXT:    [[CALL:%.*]] = call noalias i8* @internal_only_rec_static_malloc_noescape(i32 [[ARG]])
+; IS__TUNIT____-NEXT:    ret i8* [[CALL]]
+;
+; IS__CGSCC____: Function Attrs: inaccessiblememonly
+; IS__CGSCC____-LABEL: define {{[^@]+}}@internal_only_rec_static_helper_malloc_noescape
+; IS__CGSCC____-SAME: (i32 [[ARG:%.*]]) #[[ATTR0]] {
+; IS__CGSCC____-NEXT:  entry:
+; IS__CGSCC____-NEXT:    [[CALL:%.*]] = call noalias i8* @internal_only_rec_static_malloc_noescape(i32 [[ARG]])
+; IS__CGSCC____-NEXT:    ret i8* [[CALL]]
 ;
 entry:
   %call = call i8* @internal_only_rec_static_malloc_noescape(i32 %arg)
   ret i8* %call
 }
 
 define internal i8* @internal_only_rec_static_malloc_noescape(i32 %arg) {
 ; FIXME: This is actually inaccessiblememonly because the malloced memory does not escape
-; CHECK-LABEL: define {{[^@]+}}@internal_only_rec_static_malloc_noescape
-; CHECK-SAME: (i32 [[ARG:%.*]]) {
-; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[REM:%.*]] = srem i32 [[ARG]], 2
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[REM]], 1
-; CHECK-NEXT:    br i1 [[CMP]], label [[IF_THEN:%.*]], label [[IF_END:%.*]]
-; CHECK:       if.then:
-; CHECK-NEXT:    [[DIV:%.*]] = sdiv i32 [[ARG]], 2
-; CHECK-NEXT:    [[CALL:%.*]] = call noalias i8* @internal_only_rec(i32 [[DIV]])
-; CHECK-NEXT:    br label [[RETURN:%.*]]
-; CHECK:       if.end:
-; CHECK-NEXT:    [[CONV:%.*]] = sext i32 [[ARG]] to i64
-; CHECK-NEXT:    [[CALL1:%.*]] = call noalias i8* @malloc(i64 [[CONV]])
-; CHECK-NEXT:    store i8 0, i8* [[CALL1]], align 1
-; CHECK-NEXT:    br label [[RETURN]]
-; CHECK:       return:
-; CHECK-NEXT:    [[RETVAL_0:%.*]] = phi i8* [ [[CALL]], [[IF_THEN]] ], [ null, [[IF_END]] ]
-; CHECK-NEXT:    ret i8* [[RETVAL_0]]
+; IS__TUNIT____-LABEL: define {{[^@]+}}@internal_only_rec_static_malloc_noescape
+; IS__TUNIT____-SAME: (i32 [[ARG:%.*]]) {
+; IS__TUNIT____-NEXT:  entry:
+; IS__TUNIT____-NEXT:    [[REM:%.*]] = srem i32 [[ARG]], 2
+; IS__TUNIT____-NEXT:    [[CMP:%.*]] = icmp eq i32 [[REM]], 1
+; IS__TUNIT____-NEXT:    br i1 [[CMP]], label [[IF_THEN:%.*]], label [[IF_END:%.*]]
+; IS__TUNIT____:       if.then:
+; IS__TUNIT____-NEXT:    [[DIV:%.*]] = sdiv i32 [[ARG]], 2
+; IS__TUNIT____-NEXT:    [[CALL:%.*]] = call noalias i8* @internal_only_rec(i32 [[DIV]])
+; IS__TUNIT____-NEXT:    br label [[RETURN:%.*]]
+; IS__TUNIT____:       if.end:
+; IS__TUNIT____-NEXT:    [[CONV:%.*]] = sext i32 [[ARG]] to i64
+; IS__TUNIT____-NEXT:    [[CALL1:%.*]] = call noalias i8* @malloc(i64 [[CONV]])
+; IS__TUNIT____-NEXT:    br label [[RETURN]]
+; IS__TUNIT____:       return:
+; IS__TUNIT____-NEXT:    [[RETVAL_0:%.*]] = phi i8* [ [[CALL]], [[IF_THEN]] ], [ null, [[IF_END]] ]
+; IS__TUNIT____-NEXT:    ret i8* [[RETVAL_0]]
+;
+; IS__CGSCC____: Function Attrs: inaccessiblememonly
+; IS__CGSCC____-LABEL: define {{[^@]+}}@internal_only_rec_static_malloc_noescape
+; IS__CGSCC____-SAME: (i32 [[ARG:%.*]]) #[[ATTR0]] {
+; IS__CGSCC____-NEXT:  entry:
+; IS__CGSCC____-NEXT:    [[REM:%.*]] = srem i32 [[ARG]], 2
+; IS__CGSCC____-NEXT:    [[CMP:%.*]] = icmp eq i32 [[REM]], 1
+; IS__CGSCC____-NEXT:    br i1 [[CMP]], label [[IF_THEN:%.*]], label [[IF_END:%.*]]
+; IS__CGSCC____:       if.then:
+; IS__CGSCC____-NEXT:    [[DIV:%.*]] = sdiv i32 [[ARG]], 2
+; IS__CGSCC____-NEXT:    [[CALL:%.*]] = call noalias i8* @internal_only_rec(i32 [[DIV]])
+; IS__CGSCC____-NEXT:    br label [[RETURN:%.*]]
+; IS__CGSCC____:       if.end:
+; IS__CGSCC____-NEXT:    [[CONV:%.*]] = sext i32 [[ARG]] to i64
+; IS__CGSCC____-NEXT:    [[CALL1:%.*]] = call noalias i8* @malloc(i64 [[CONV]])
+; IS__CGSCC____-NEXT:    br label [[RETURN]]
+; IS__CGSCC____:       return:
+; IS__CGSCC____-NEXT:    [[RETVAL_0:%.*]] = phi i8* [ [[CALL]], [[IF_THEN]] ], [ null, [[IF_END]] ]
+; IS__CGSCC____-NEXT:    ret i8* [[RETVAL_0]]
 ;
 entry:
   %rem = srem i32 %arg, 2
   %cmp = icmp eq i32 %rem, 1
   br i1 %cmp, label %if.then, label %if.end
 
 if.then:                                          ; preds = %entry
   %div = sdiv i32 %arg, 2