[FunctionAttrs] Infer precise FMRB

[nikic]

Summary

This updates checkFunctionMemoryAccess() to infer a precise FunctionModRefBehavior, rather than an approximation split into read/write and argmemonly.

Afterwards, we still map this back to imprecise function attributes. This still allows us to infer some cases that we previously did not handle, namely inaccessiblememonly and inaccessiblemem_or_argmemonly. In practice, this means we get better memory attributes in the presence of intrinsics like @llvm.assume.

[reames]

LGTM w/suggestion.

llvm/lib/Transforms/IPO/FunctionAttrs.cpp

216

The claim that no instruction other than a call can access inaccessiblememory seems reasonable on the surface, but is a somewhat significant change. The corner cases would be something like maybe an indirect br.

If possible, could you separate this into it's own commit for risk reduction? Having something easy to revert could be helpful.

[nikic]

llvm/lib/Transforms/IPO/FunctionAttrs.cpp

216

Interesting point. Thinking about this again, I think there are two cases here:

1. If we have a MemoryLocation, then only that location can be accessed and an IR location can't be inaccessible by definition.

2. If we don't have a MemoryLocation (but still access memory), then things are less clear cut. Checking the implementation, the (non-call) instructions for which this happens are fence, catchpad and catchret.

In the interest of being conservative, it probably makes sense to keep the inaccessiblemem access in the latter case. I'm not familiar enough with the semantics of these instructions to exclude this -- e.g. could a fence synchronize atomic operations on an inaccessible object?

[antondaubert]

This change introduced a problem, or exposed an already existing bug.
Here you can find a reproducer: https://reviews.llvm.org/P8291
Could we get this change reverted until you a working on a fix?

llvm/lib/Transforms/IPO/FunctionAttrs.cpp

 /// If ThisBody is true, this function may examine the function body and will
 /// return a result pertaining to this copy of the function. If it is false, the
 /// result will be based only on AA results for the function declaration; it
 /// will be assumed that some other (perhaps less optimized) version of the
 /// function may be selected at link time.
 static FunctionModRefBehavior
 checkFunctionMemoryAccess(Function &F, bool ThisBody, AAResults &AAR,
                           const SCCNodeSet &SCCNodes) {
-  FunctionModRefBehavior MRB = AAR.getModRefBehavior(&F);
-  if (MRB.doesNotAccessMemory())
+  FunctionModRefBehavior OrigMRB = AAR.getModRefBehavior(&F);
+  if (OrigMRB.doesNotAccessMemory())
     // Already perfect!
-    return MRB;
+    return OrigMRB;
 
   if (!ThisBody)
-    return MRB;
+    return OrigMRB;
 
-  // TODO: We should directly populate a FunctionModRefBehavior here.
   // Scan the function body for instructions that may read or write memory.
-  ModRefInfo MR = ModRefInfo::NoModRef;
-  // Track if the function accesses memory not based on pointer arguments or
-  // allocas.
-  bool AccessesNonArgsOrAlloca = false;
+  FunctionModRefBehavior MRB = FunctionModRefBehavior::none();
   // Returns true if Ptr is not based on a function argument.
   auto IsArgumentOrAlloca = [](const Value *Ptr) {
     const Value *UO = getUnderlyingObject(Ptr);
     return isa<Argument>(UO) || isa<AllocaInst>(UO);
   };
   for (Instruction &I : instructions(F)) {
     // Some instructions can be ignored even if they read or write memory.
     // Detect these now, skipping to the next instruction if one is found.
     if (auto *Call = dyn_cast<CallBase>(&I)) {
       // Ignore calls to functions in the same SCC, as long as the call sites
       // don't have operand bundles.  Calls with operand bundles are allowed to
       // have memory effects not described by the memory effects of the call
       // target.
       if (!Call->hasOperandBundles() && Call->getCalledFunction() &&
           SCCNodes.count(Call->getCalledFunction()))
         continue;
-      FunctionModRefBehavior MRB = AAR.getModRefBehavior(Call);
-      ModRefInfo MRI = MRB.getModRef();
+      FunctionModRefBehavior CallMRB = AAR.getModRefBehavior(Call);
 
       // If the call doesn't access memory, we're done.
-      if (isNoModRef(MRI))
+      if (CallMRB.doesNotAccessMemory())
         continue;
 
       // A pseudo probe call shouldn't change any function attribute since it
       // doesn't translate to a real instruction. It comes with a memory access
       // tag to prevent itself being removed by optimizations and not block
       // other instructions being optimized.
       if (isa<PseudoProbeInst>(I))
         continue;
 
-      if (!MRB.onlyAccessesArgPointees()) {
-        // The call could access any memory.
-        MR |= MRI;
-        AccessesNonArgsOrAlloca = true;
-        continue;
-      }
+      MRB |= CallMRB.getWithoutLoc(FunctionModRefBehavior::ArgMem);
 
       // Check whether all pointer arguments point to local memory, and
       // ignore calls that only access local memory.
+      ModRefInfo ArgMR = CallMRB.getModRef(FunctionModRefBehavior::ArgMem);
+      if (ArgMR != ModRefInfo::NoModRef) {
         for (const Use &U : Call->args()) {
           const Value *Arg = U;
           if (!Arg->getType()->isPtrOrPtrVectorTy())
             continue;
 
           MemoryLocation Loc =
               MemoryLocation::getBeforeOrAfter(Arg, I.getAAMetadata());
           // Skip accesses to local or constant memory as they don't impact the
           // externally visible mod/ref behavior.
           if (AAR.pointsToConstantMemory(Loc, /*OrLocal=*/true))
             continue;
 
-        AccessesNonArgsOrAlloca |= !IsArgumentOrAlloca(Loc.Ptr);
-        MR |= MRI;
+          MRB |= FunctionModRefBehavior::argMemOnly(ArgMR);
+          if (!IsArgumentOrAlloca(Loc.Ptr))
+            MRB |= FunctionModRefBehavior(FunctionModRefBehavior::Other, ArgMR);
+        }
       }
       continue;
     }
 
-    if (!I.mayReadOrWriteMemory())
+    ModRefInfo MR = ModRefInfo::NoModRef;
+    if (I.mayWriteToMemory())
+      MR |= ModRefInfo::Mod;
+    if (I.mayReadFromMemory())
+      MR |= ModRefInfo::Ref;
+    if (MR == ModRefInfo::NoModRef)
       continue;
 
     Optional<MemoryLocation> Loc = MemoryLocation::getOrNone(&I);
-    // Ignore non-volatile accesses from local memory. (Atomic is okay here.)
-    if (Loc && !I.isVolatile() &&
-        AAR.pointsToConstantMemory(*Loc, /*OrLocal=*/true))
+    if (!Loc) {
+      // If no location is known, conservatively assume anything can be
+      // accessed.
+      MRB |= FunctionModRefBehavior(MR);
       continue;
+    }
 
-    AccessesNonArgsOrAlloca |= !Loc || !IsArgumentOrAlloca(Loc->Ptr);
-
-    // Any remaining instructions need to be taken seriously!  Check if they
-    // read or write memory.
-    //
-    // Writes memory, remember that.
-    if (I.mayWriteToMemory())
-      MR |= ModRefInfo::Mod;
+    // Ignore non-volatile accesses from local memory. (Atomic is okay here.)
+    if (!I.isVolatile() && AAR.pointsToConstantMemory(*Loc, /*OrLocal=*/true))
+      continue;
 
-    // If this instruction may read memory, remember that.
-    if (I.mayReadFromMemory())
-      MR |= ModRefInfo::Ref;
+    // The accessed location can be either only argument memory, or
+    // argument & other memory, but never inaccessible memory.
+    MRB |= FunctionModRefBehavior::argMemOnly(MR);
+    if (!IsArgumentOrAlloca(Loc->Ptr))
+      MRB |= FunctionModRefBehavior(FunctionModRefBehavior::Other, MR);
   }
 
-  if (!AccessesNonArgsOrAlloca)
-    return FunctionModRefBehavior::argMemOnly(MR);
-  return FunctionModRefBehavior(MR);
+  return OrigMRB & MRB;
 }
 
 FunctionModRefBehavior llvm::computeFunctionBodyMemoryAccess(Function &F,
                                                              AAResults &AAR) {
   return checkFunctionMemoryAccess(F, /*ThisBody=*/true, AAR, {});
 }
 
 /// Deduce readonly/readnone/writeonly attributes for the SCC.
 template <typename AARGetterT>
 static void addMemoryAttrs(const SCCNodeSet &SCCNodes, AARGetterT &&AARGetter,
                            SmallSet<Function *, 8> &Changed) {
-  // Check if any of the functions in the SCC read or write memory.  If they
-  // write memory then they can't be marked readnone or readonly.
-  bool ReadsMemory = false;
-  bool WritesMemory = false;
-  // Check if all functions only access memory through their arguments.
-  bool ArgMemOnly = true;
+  FunctionModRefBehavior FMRB = FunctionModRefBehavior::none();
   for (Function *F : SCCNodes) {
     // Call the callable parameter to look up AA results for this function.
     AAResults &AAR = AARGetter(*F);
     // Non-exact function definitions may not be selected at link time, and an
     // alternative version that writes to memory may be selected.  See the
     // comment on GlobalValue::isDefinitionExact for more details.
-    FunctionModRefBehavior FMRB =
-        checkFunctionMemoryAccess(*F, F->hasExactDefinition(), AAR, SCCNodes);
-    if (FMRB.doesNotAccessMemory())
-      continue;
-    ModRefInfo MR = FMRB.getModRef();
-    ReadsMemory |= isRefSet(MR);
-    WritesMemory |= isModSet(MR);
-    ArgMemOnly &= FMRB.onlyAccessesArgPointees();
-    // Reached neither readnone, readonly, writeonly nor argmemonly can be
-    // inferred. Exit.
-    if (ReadsMemory && WritesMemory && !ArgMemOnly)
+    FMRB |= checkFunctionMemoryAccess(*F, F->hasExactDefinition(), AAR,
+                                      SCCNodes);
+    // Reached bottom of the lattice, we will not be able to improve the result.
+    if (FMRB == FunctionModRefBehavior::unknown())
       return;
   }
 
-  assert((!ReadsMemory || !WritesMemory || ArgMemOnly) &&
-         "no memory attributes can be added for this SCC, should have exited "
-         "earlier");
-  // Success!  Functions in this SCC do not access memory, only read memory,
-  // only write memory, or only access memory through its arguments. Give them
-  // the appropriate attribute.
+  ModRefInfo MR = FMRB.getModRef();
 
   for (Function *F : SCCNodes) {
-    // If possible add argmemonly attribute to F, if it accesses memory.
-    if (ArgMemOnly && !F->onlyAccessesArgMemory() &&
-        (ReadsMemory || WritesMemory)) {
+    if (F->doesNotAccessMemory())
+      // Already perfect!
+      continue;
+
+    if (FMRB.doesNotAccessMemory()) {
+      // For readnone, remove all other memory attributes.
+      AttributeMask AttrsToRemove;
+      AttrsToRemove.addAttribute(Attribute::ReadOnly);
+      AttrsToRemove.addAttribute(Attribute::WriteOnly);
+      AttrsToRemove.addAttribute(Attribute::ArgMemOnly);
+      AttrsToRemove.addAttribute(Attribute::InaccessibleMemOnly);
+      AttrsToRemove.addAttribute(Attribute::InaccessibleMemOrArgMemOnly);
+
+      ++NumReadNone;
+      F->removeFnAttrs(AttrsToRemove);
+      F->addFnAttr(Attribute::ReadNone);
+      Changed.insert(F);
+      continue;
+    }
+
+    // Add argmemonly, inaccessiblememonly, or inaccessible_or_argmemonly
+    // attributes if possible.
+    AttributeMask AttrsToRemove;
+    AttrsToRemove.addAttribute(Attribute::ArgMemOnly);
+    AttrsToRemove.addAttribute(Attribute::InaccessibleMemOnly);
+    AttrsToRemove.addAttribute(Attribute::InaccessibleMemOrArgMemOnly);
+    if (FMRB.onlyAccessesArgPointees() && !F->onlyAccessesArgMemory()) {
       NumArgMemOnly++;
+      F->removeFnAttrs(AttrsToRemove);
       F->addFnAttr(Attribute::ArgMemOnly);
       Changed.insert(F);
+    } else if (FMRB.onlyAccessesInaccessibleMem() &&
+               !F->onlyAccessesInaccessibleMemory()) {
+      F->removeFnAttrs(AttrsToRemove);
+      F->addFnAttr(Attribute::InaccessibleMemOnly);
+      Changed.insert(F);
+    } else if (FMRB.onlyAccessesInaccessibleOrArgMem() &&
+               !F->onlyAccessesInaccessibleMemOrArgMem()) {
+      F->removeFnAttrs(AttrsToRemove);
+      F->addFnAttr(Attribute::InaccessibleMemOrArgMemOnly);
+      Changed.insert(F);
     }
 
     // The SCC contains functions both writing and reading from memory. We
     // cannot add readonly or writeonline attributes.
-    if (ReadsMemory && WritesMemory)
-      continue;
-    if (F->doesNotAccessMemory())
-      // Already perfect!
+    if (MR == ModRefInfo::ModRef)
       continue;
 
-    if (F->onlyReadsMemory() && ReadsMemory)
-      // No change.
+    if (F->onlyReadsMemory() && MR == ModRefInfo::Ref)
       continue;
 
-    if (F->onlyWritesMemory() && WritesMemory)
+    if (F->onlyWritesMemory() && MR == ModRefInfo::Mod)
       continue;
 
     Changed.insert(F);
 
-    // Clear out any existing attributes.
-    AttributeMask AttrsToRemove;
-    AttrsToRemove.addAttribute(Attribute::ReadOnly);
-    AttrsToRemove.addAttribute(Attribute::ReadNone);
-    AttrsToRemove.addAttribute(Attribute::WriteOnly);
-
-    if (!WritesMemory && !ReadsMemory) {
-      // Clear out any "access range attributes" if readnone was deduced.
-      AttrsToRemove.addAttribute(Attribute::ArgMemOnly);
-      AttrsToRemove.addAttribute(Attribute::InaccessibleMemOnly);
-      AttrsToRemove.addAttribute(Attribute::InaccessibleMemOrArgMemOnly);
-    }
-    F->removeFnAttrs(AttrsToRemove);
-
     // Add in the new attribute.
-    if (WritesMemory && !ReadsMemory)
-      F->addFnAttr(Attribute::WriteOnly);
-    else
-      F->addFnAttr(ReadsMemory ? Attribute::ReadOnly : Attribute::ReadNone);
-
-    if (WritesMemory && !ReadsMemory)
+    if (MR == ModRefInfo::Mod) {
       ++NumWriteOnly;
-    else if (ReadsMemory)
+      F->removeFnAttr(Attribute::ReadOnly);
+      F->addFnAttr(Attribute::WriteOnly);
+    } else {
       ++NumReadOnly;
-    else
-      ++NumReadNone;
+      assert(MR == ModRefInfo::Ref);
+      F->removeFnAttr(Attribute::WriteOnly);
+      F->addFnAttr(Attribute::ReadOnly);
+    }
   }
 }
 
 // Compute definitive function attributes for a function taking into account
 // prevailing definitions and linkage types
 static FunctionSummary *calculatePrevailingSummary(
     ValueInfo VI,
     DenseMap<ValueInfo, FunctionSummary *> &CachedPrevailingSummary,

llvm/test/Transforms/FunctionAttrs/argmemonly.ll

entry:
   store i32 %l, i32* %a
   %l.2 = load i32, i32* %a
   ret i32 %l.2
 }
 
 declare void @fn_inaccessiblememonly() inaccessiblememonly
 
 define void @test_inaccessiblememonly() {
+; CHECK: Function Attrs: inaccessiblememonly
 ; CHECK-LABEL: @test_inaccessiblememonly(
 ; CHECK-NEXT:    call void @fn_inaccessiblememonly()
 ; CHECK-NEXT:    ret void
 ;
   call void @fn_inaccessiblememonly()
   ret void
 }
 
 define void @test_inaccessiblememonly_readonly() {
-; CHECK: Function Attrs: nofree readonly
+; CHECK: Function Attrs: inaccessiblememonly nofree readonly
 ; CHECK-LABEL: @test_inaccessiblememonly_readonly(
-; CHECK-NEXT:    call void @fn_inaccessiblememonly() #[[ATTR13:[0-9]+]]
+; CHECK-NEXT:    call void @fn_inaccessiblememonly() #[[ATTR15:[0-9]+]]
 ; CHECK-NEXT:    ret void
 ;
   call void @fn_inaccessiblememonly() readonly
   ret void
 }
 
 define void @test_inaccessibleorargmemonly_readonly(i32* %arg) {
-; CHECK: Function Attrs: nofree readonly
+; CHECK: Function Attrs: inaccessiblemem_or_argmemonly nofree readonly
 ; CHECK-LABEL: @test_inaccessibleorargmemonly_readonly(
 ; CHECK-NEXT:    [[TMP1:%.*]] = load i32, i32* [[ARG:%.*]], align 4
-; CHECK-NEXT:    call void @fn_inaccessiblememonly() #[[ATTR13]]
+; CHECK-NEXT:    call void @fn_inaccessiblememonly() #[[ATTR15]]
 ; CHECK-NEXT:    ret void
 ;
   load i32, i32* %arg
   call void @fn_inaccessiblememonly() readonly
   ret void
 }
 
 define void @test_inaccessibleorargmemonly_readwrite(i32* %arg) {
+; CHECK: Function Attrs: inaccessiblemem_or_argmemonly
 ; CHECK-LABEL: @test_inaccessibleorargmemonly_readwrite(
 ; CHECK-NEXT:    store i32 0, i32* [[ARG:%.*]], align 4
-; CHECK-NEXT:    call void @fn_inaccessiblememonly() #[[ATTR13]]
+; CHECK-NEXT:    call void @fn_inaccessiblememonly() #[[ATTR15]]
 ; CHECK-NEXT:    ret void
 ;
   store i32 0, i32* %arg
   call void @fn_inaccessiblememonly() readonly
   ret void
 }