From 5e06b2514aed37e49224c7086468bd4fa5080086 Mon Sep 17 00:00:00 2001
From: Johannes Doerfert <johannes@jdoerfert.de>
Date: Sat, 9 May 2020 12:48:29 -0500
Subject: [PATCH] [Attributor][FIX] Carefully handle/ignore/forget `argmemonly`

When we have an existing `argmemonly` or `inaccessiblememorargmemonly`
we used to "know" that information. However, interprocedural constant
propagation can invalidate these attributes. We now ignore and remove
these attributes for internal functions (which may be affected by IP
constant propagation), if we are deriving new attributes for the
function.
---
 llvm/include/llvm/Transforms/IPO/Attributor.h |  5 ++++
 .../Transforms/IPO/AttributorAttributes.cpp   | 27 +++++++++++++++----
 .../Transforms/Attributor/memory_locations.ll | 25 ++++++++++++++++-
 3 files changed, 51 insertions(+), 6 deletions(-)

diff --git a/llvm/include/llvm/Transforms/IPO/Attributor.h b/llvm/include/llvm/Transforms/IPO/Attributor.h
index 42675f89070b..505e40b2c283 100644
--- a/llvm/include/llvm/Transforms/IPO/Attributor.h
+++ b/llvm/include/llvm/Transforms/IPO/Attributor.h
@@ -903,6 +903,11 @@ struct Attributor {
            Functions.size() == Functions.front()->getParent()->size();
   }
 
+  /// Return true if we derive attributes for \p Fn
+  bool isRunOn(Function &Fn) const {
+    return Functions.empty() || Functions.count(&Fn);
+  }
+
   /// Determine opportunities to derive 'default' attributes in \p F and create
   /// abstract attribute objects for them.
   ///
diff --git a/llvm/lib/Transforms/IPO/AttributorAttributes.cpp b/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
index 1c228138de36..f1a28585acea 100644
--- a/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
+++ b/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
@@ -6017,14 +6017,25 @@ struct AAMemoryLocationImpl : public AAMemoryLocation {
   /// See AbstractAttribute::initialize(...).
   void initialize(Attributor &A) override {
     intersectAssumedBits(BEST_STATE);
-    getKnownStateFromValue(getIRPosition(), getState());
+    getKnownStateFromValue(A, getIRPosition(), getState());
     IRAttribute::initialize(A);
   }
 
   /// Return the memory behavior information encoded in the IR for \p IRP.
-  static void getKnownStateFromValue(const IRPosition &IRP,
+  static void getKnownStateFromValue(Attributor &A, const IRPosition &IRP,
                                      BitIntegerState &State,
                                      bool IgnoreSubsumingPositions = false) {
+    // For internal functions we ignore `argmemonly` and
+    // `inaccessiblememorargmemonly` as we might break it via interprocedural
+    // constant propagation. It is unclear if this is the best way but it is
+    // unlikely this will cause real performance problems. If we are deriving
+    // attributes for the anchor function we even remove the attribute in
+    // addition to ignoring it.
+    bool UseArgMemOnly = true;
+    Function *AnchorFn = IRP.getAnchorScope();
+    if (AnchorFn && A.isRunOn(*AnchorFn))
+      UseArgMemOnly = !AnchorFn->hasLocalLinkage();
+
     SmallVector<Attribute, 2> Attrs;
     IRP.getAttrs(AttrKinds, Attrs, IgnoreSubsumingPositions);
     for (const Attribute &Attr : Attrs) {
@@ -6036,11 +6047,17 @@ struct AAMemoryLocationImpl : public AAMemoryLocation {
         State.addKnownBits(inverseLocation(NO_INACCESSIBLE_MEM, true, true));
         break;
       case Attribute::ArgMemOnly:
-        State.addKnownBits(inverseLocation(NO_ARGUMENT_MEM, true, true));
+        if (UseArgMemOnly)
+          State.addKnownBits(inverseLocation(NO_ARGUMENT_MEM, true, true));
+        else
+          IRP.removeAttrs({Attribute::ArgMemOnly});
         break;
       case Attribute::InaccessibleMemOrArgMemOnly:
-        State.addKnownBits(
-            inverseLocation(NO_INACCESSIBLE_MEM | NO_ARGUMENT_MEM, true, true));
+        if (UseArgMemOnly)
+          State.addKnownBits(inverseLocation(
+              NO_INACCESSIBLE_MEM | NO_ARGUMENT_MEM, true, true));
+        else
+          IRP.removeAttrs({Attribute::InaccessibleMemOrArgMemOnly});
         break;
       default:
         llvm_unreachable("Unexpected attribute!");
diff --git a/llvm/test/Transforms/Attributor/memory_locations.ll b/llvm/test/Transforms/Attributor/memory_locations.ll
index 535b48757466..3e6facb44577 100644
--- a/llvm/test/Transforms/Attributor/memory_locations.ll
+++ b/llvm/test/Transforms/Attributor/memory_locations.ll
@@ -5,6 +5,8 @@
 ; RUN: opt -aa-pipeline=basic-aa -passes=attributor-cgscc -attributor-manifest-internal  -attributor-annotate-decl-cs -S < %s | FileCheck %s --check-prefixes=CHECK,NOT_TUNIT_NPM,NOT_TUNIT_OPM,NOT_CGSCC_OPM,IS__CGSCC____,IS________NPM,IS__CGSCC_NPM
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 
+@G = external dso_local global i32, align 4
+
 ; CHECK: Function Attrs: inaccessiblememonly
 declare noalias i8* @malloc(i64) inaccessiblememonly
 
@@ -397,7 +399,6 @@ define void @callerE(i8* %arg) {
   ret void
 }
 
-@G = external dso_local global i32, align 4
 
 ; CHECK: Function Attrs:
 ; CHECK-SAME: writeonly
@@ -601,3 +602,25 @@ define i8 @readnone_caller2(i1 %c) {
   %r = call i8 @recursive_not_readnone_internal2(i8* undef, i1 %c)
   ret i8 %r
 }
+
+; CHECK: Function Attrs:
+; CHECK-NOT: argmemonly
+define internal void @argmemonly_before_ipconstprop(i32* %p) argmemonly {
+; CHECK-LABEL: define {{[^@]+}}@argmemonly_before_ipconstprop()
+; CHECK-NEXT:    store i32 0, i32* @G, align 4
+; CHECK-NEXT:    ret void
+;
+  store i32 0, i32* %p
+  ret void
+}
+
+; CHECK: Function Attrs:
+; CHECK-NOT: argmemonly
+define void @argmemonky_caller() {
+; CHECK-LABEL: define {{[^@]+}}@argmemonky_caller()
+; CHECK-NEXT:    call void @argmemonly_before_ipconstprop()
+; CHECK-NEXT:    ret void
+;
+  call void @argmemonly_before_ipconstprop(i32* @G)
+  ret void
+}
-- 
GitLab