[Inliner] Add support for preserving nocapture param attr

Currently if we have:
```
define @foo(ptr nocapture %p) {
entry:
    ...
    bar(ptr %p)
    ...
}
```

When inlining `foo`, we will lose the `nocapture` on `%p` which might
not be recoverable.

The goal of this patch is to preserve the `nocapture` if some
conservative analysis indicates we can.

1) Return value of `bar` is either unused or only used as return of
   `foo` (this rules of capture via return).

2) No `alloca` (or scratch memory of any sort) in `foo` s.t there is a
   path from `entry` to `bar` that goes an `alloca`. This helps rule
   out `bar` capturing `%p` in memory in a way that wouldn't be
   capturing outside of the scope of `foo`.

3) No paths in `foo` that go through `bar` have any instructions with
   side-effects other than `bar`. This rules out `bar` capturing `%p`
   in memory, but then some later instructions clearing the memory
   capture s.t `nocapture` in `foo` still holds. It also rules out
   some function (i.e `malloc`) creating scratch memory that `bar`
   could capture `%p` in but still only visible in the scope of `foo`.

Ultimately these three checks are highly conservative, but should
allow some preservation.

Author: goldsteinn

llvm/lib/Transforms/Utils/InlineFunction.cpp

@@ -68,6 +68,7 @@
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/ModRef.h"
 #include "llvm/Transforms/Utils/AssumeBundleBuilder.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/Local.h"
@@ -1364,6 +1365,104 @@ static bool MayContainThrowingOrExitingCallAfterCB(CallBase *Begin,
       ++BeginIt, End->getIterator(), InlinerAttributeWindow + 1);
 }
 
+template <typename RangeT> static bool ContainsSideEffects(RangeT Range) {
+  // Any instruction that may clear local scratch space CB stored
+  // into.
+  return any_of(Range, [](Instruction &I) { return I.mayHaveSideEffects(); });
+}
+
+template <typename RangeT> static bool ContainsScratchSpace(RangeT Range) {
+  return any_of(Range, [](Instruction &I) {
+    // Any instruction that may create local scratch space CB can store
+    // into.
+    return I.mayHaveSideEffects() || isa<AllocaInst>(&I);
+  });
+}
+
+template <typename NextFn, typename CheckFn>
+static bool CheckPathFromBBRecurse(DenseMap<BasicBlock *, bool> &CachedRes,
+                                   bool First, BasicBlock *BB, NextFn Next,
+                                   CheckFn Check) {
+  if (!First) {
+    // Initialize to true (okay to propagate) `nocapture`. This means that loops
+    // will be okay.
+    auto [Iter, Inserted] = CachedRes.try_emplace(BB, true);
+    // If we already have a result, return it.
+    if (!Inserted)
+      return Iter->second;
+
+    if (!Check(BB->instructionsWithoutDebug())) {
+      Iter->second = false;
+      return false;
+    }
+  }
+  auto NextBBs = Next(BB);
+  // Check all Succs/Preds
+  for (BasicBlock *NextBB : NextBBs) {
+    if (!CheckPathFromBBRecurse(CachedRes, /*First=*/false, NextBB, Next,
+                                Check)) {
+      CachedRes[BB] = false;
+      return false;
+    }
+  }
+
+  return true;
+}
+
+// Assuming we have:
+// define @foo(ptr nocapture %p) {
+// entry:
+//	...
+//  bar (ptr %p)
+//	...
+// }
+//
+// Determine if we can propagate `nocapture` to the `%p` at the
+// `bar`.
+static bool
+CanPropagateNoCaptureAtCB(DenseMap<BasicBlock *, bool> &PureFromBB,
+                          DenseMap<BasicBlock *, bool> &NoLocalStateToBB,
+                          BasicBlock *BB, CallBase *CB) {
+  // If CB returns and its used by anything other than `ret`, assume it may be
+  // capturing.
+  // Potential TODO: We could allow many operations.
+  if (!CB->getType()->isVoidTy())
+    for (auto Use : CB->users())
+      if (!isa<ReturnInst>(Use))
+        return false;
+
+  // Can't capture via return, so if no side-effects we are set.
+  if (!CB->mayHaveSideEffects())
+    return true;
+
+  auto It = CB->getIterator();
+  ++It;
+
+  // Check that CB instruction with side-effects on all paths from
+  // `entry` that go through the CB and there are no `alloca`
+  // instructions. This accomplishes two things. 1) It ensures that
+  // after CB, there is no way a store/other could "clean up" any
+  // captures from CB. 2) There is no local state (i.e `alloca` or a
+  // local `malloc`) that could CB could have stored in params in.
+  if (ContainsSideEffects(make_range(It, BB->end())) ||
+      ContainsScratchSpace(make_range(BB->begin(), CB->getIterator())))
+    return false;
+
+  if (!CheckPathFromBBRecurse(
+          PureFromBB, /*First=*/true, BB,
+          [](BasicBlock *CheckedBB) { return successors(CheckedBB); },
+          [](const auto &Region) { return !ContainsSideEffects(Region); }))
+    return false;
+
+  if (!CheckPathFromBBRecurse(
+          PureFromBB, /*First=*/true, BB,
+          [](BasicBlock *CheckedBB) { return predecessors(CheckedBB); },
+          [](const auto &Region) { return !ContainsScratchSpace(Region); }))
+    return false;
+
+  return true;
+}
+
 // Add attributes from CB params and Fn attributes that can always be propagated
 // to the corresponding argument / inner callbases.
 static void AddParamAndFnBasicAttributes(const CallBase &CB,
@@ -1376,6 +1475,9 @@ static void AddParamAndFnBasicAttributes(const CallBase &CB,
   SmallVector<AttrBuilder> ValidObjParamAttrs, ValidExactParamAttrs;
   bool HasAttrToPropagate = false;
 
+  DenseMap<BasicBlock *, bool> PureFromBB{};
+  DenseMap<BasicBlock *, bool> NoLocalStateToBB{};
+
   // Attributes we can only propagate if the exact parameter is forwarded.
   // We can propagate both poison generating and UB generating attributes
   // without any extra checks. The only attribute that is tricky to propagate
@@ -1394,6 +1496,8 @@ static void AddParamAndFnBasicAttributes(const CallBase &CB,
       ValidObjParamAttrs.back().addAttribute(Attribute::ReadNone);
     if (CB.paramHasAttr(I, Attribute::ReadOnly))
       ValidObjParamAttrs.back().addAttribute(Attribute::ReadOnly);
+    if (CB.doesNotCapture(I))
+      ValidObjParamAttrs.back().addCapturesAttr(CaptureInfo::none());
 
     for (Attribute::AttrKind AK : ExactAttrsToPropagate) {
       Attribute Attr = CB.getParamAttr(I, AK);
@@ -1478,9 +1582,16 @@ static void AddParamAndFnBasicAttributes(const CallBase &CB,
           continue;
         }
 
-        // If so, propagate its access attributes.
-        AL = AL.addParamAttributes(Context, I, ValidObjParamAttrs[ArgNo]);
+        AttributeSet AS = AttributeSet::get(Context, ValidObjParamAttrs[ArgNo]);
+        // Check if we can propagate `captures(none)`.
+        if (capturesNothing(AS.getCaptureInfo()) &&
+            (NewInnerCB->doesNotCapture(I) ||
+             !CanPropagateNoCaptureAtCB(PureFromBB, NoLocalStateToBB, &BB,
+                                        cast<CallBase>(&Ins))))
+          AS = AS.removeAttribute(Context, Attribute::Captures);
 
+        // If so, propagate its access attributes.
+        AL = AL.addParamAttributes(Context, I, AttrBuilder{Context, AS});
         // We can have conflicting attributes from the inner callsite and
         // to-be-inlined callsite. In that case, choose the most
         // restrictive.

llvm/test/Transforms/Inline/prop-nocapture.ll

@@ -20,7 +20,7 @@ define void @simple_nocapture_prop(ptr captures(none) %p) {
 define void @simple_nocapture_prop_caller(ptr %p) {
 ; CHECK-LABEL: define {{[^@]+}}@simple_nocapture_prop_caller
 ; CHECK-SAME: (ptr [[P:%.*]]) {
-; CHECK-NEXT:    call void @void.call.p0(ptr [[P]])
+; CHECK-NEXT:    call void @void.call.p0(ptr captures(none) [[P]])
 ; CHECK-NEXT:    ret void
 ;
   call void @simple_nocapture_prop(ptr %p)
@@ -40,7 +40,7 @@ define i32 @nocapture_with_return_prop(ptr captures(none) %p) {
 define i32 @nocapture_with_return_prop_caller(ptr %p) {
 ; CHECK-LABEL: define {{[^@]+}}@nocapture_with_return_prop_caller
 ; CHECK-SAME: (ptr [[P:%.*]]) {
-; CHECK-NEXT:    [[R_I:%.*]] = call i32 @ret.call.p0(ptr [[P]])
+; CHECK-NEXT:    [[R_I:%.*]] = call i32 @ret.call.p0(ptr captures(none) [[P]])
 ; CHECK-NEXT:    ret i32 [[R_I]]
 ;
   %r = call i32 @nocapture_with_return_prop(ptr %p)
@@ -193,7 +193,7 @@ define void @nocapture_prop_okay_seperate_alloca_caller(ptr %p, i1 %c) {
 ; CHECK-NEXT:    call void @llvm.stackrestore.p0(ptr [[SAVEDSTACK]])
 ; CHECK-NEXT:    br label [[NOCAPTURE_PROP_OKAY_SEPERATE_ALLOCA_EXIT:%.*]]
 ; CHECK:       F.i:
-; CHECK-NEXT:    call void @void.call.p0(ptr [[P]])
+; CHECK-NEXT:    call void @void.call.p0(ptr captures(none) [[P]])
 ; CHECK-NEXT:    call void @llvm.stackrestore.p0(ptr [[SAVEDSTACK]])
 ; CHECK-NEXT:    br label [[NOCAPTURE_PROP_OKAY_SEPERATE_ALLOCA_EXIT]]
 ; CHECK:       nocapture_prop_okay_seperate_alloca.exit:
@@ -286,10 +286,10 @@ F:
 define i32 @nocapture_prop_okay_no_sideeffects_caller(ptr %p, i1 %c) {
 ; CHECK-LABEL: define {{[^@]+}}@nocapture_prop_okay_no_sideeffects_caller
 ; CHECK-SAME: (ptr [[P:%.*]], i1 [[C:%.*]]) {
-; CHECK-NEXT:    call void @void.call.p0(ptr [[P]])
+; CHECK-NEXT:    call void @void.call.p0(ptr captures(none) [[P]])
 ; CHECK-NEXT:    br i1 [[C]], label [[T_I:%.*]], label [[F_I:%.*]]
 ; CHECK:       T.i:
-; CHECK-NEXT:    [[R_I:%.*]] = call i32 @ret.call.p0(ptr [[P]]) #[[ATTR3]]
+; CHECK-NEXT:    [[R_I:%.*]] = call i32 @ret.call.p0(ptr captures(none) [[P]]) #[[ATTR3]]
 ; CHECK-NEXT:    br label [[NOCAPTURE_PROP_OKAY_NO_SIDEEFFECTS_EXIT:%.*]]
 ; CHECK:       F.i:
 ; CHECK-NEXT:    br label [[NOCAPTURE_PROP_OKAY_NO_SIDEEFFECTS_EXIT]]