From fe196380ccd4f2ed0e78151e11387a299a72b337 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov@redhat.com>
Date: Fri, 23 Sep 2022 13:55:07 +0200
Subject: [PATCH] [FunctionAttrs] Use MemoryLocation::getOrNone() when infering
 memory attrs

MemoryLocation::getOrNone() already has the necessary logic to
handle different instruction types. Use it, rather than repeating
a subset of the logic. This adds support for previously unhandled
instructions like atomicrmw.
---
 llvm/lib/Transforms/IPO/FunctionAttrs.cpp    | 35 +++++++++-------------------
 llvm/test/Transforms/FunctionAttrs/nosync.ll |  6 ++---
 2 files changed, 14 insertions(+), 27 deletions(-)

diff --git a/llvm/lib/Transforms/IPO/FunctionAttrs.cpp b/llvm/lib/Transforms/IPO/FunctionAttrs.cpp
index 843d663..5bb9412 100644
--- a/llvm/lib/Transforms/IPO/FunctionAttrs.cpp
+++ b/llvm/lib/Transforms/IPO/FunctionAttrs.cpp
@@ -194,32 +194,19 @@ checkFunctionMemoryAccess(Function &F, bool ThisBody, AAResults &AAR,
         MR |= MRI;
       }
       continue;
-    } else if (LoadInst *LI = dyn_cast<LoadInst>(&I)) {
-      MemoryLocation Loc = MemoryLocation::get(LI);
-      // Ignore non-volatile loads from local memory. (Atomic is okay here.)
-      if (!LI->isVolatile() &&
-          AAR.pointsToConstantMemory(Loc, /*OrLocal=*/true))
-        continue;
-      AccessesNonArgsOrAlloca |= !IsArgumentOrAlloca(Loc.Ptr);
-    } else if (StoreInst *SI = dyn_cast<StoreInst>(&I)) {
-      MemoryLocation Loc = MemoryLocation::get(SI);
-      // Ignore non-volatile stores to local memory. (Atomic is okay here.)
-      if (!SI->isVolatile() &&
-          AAR.pointsToConstantMemory(Loc, /*OrLocal=*/true))
-        continue;
-      AccessesNonArgsOrAlloca |= !IsArgumentOrAlloca(Loc.Ptr);
-    } else if (VAArgInst *VI = dyn_cast<VAArgInst>(&I)) {
-      // Ignore vaargs on local memory.
-      MemoryLocation Loc = MemoryLocation::get(VI);
-      if (AAR.pointsToConstantMemory(Loc, /*OrLocal=*/true))
-        continue;
-      AccessesNonArgsOrAlloca |= !IsArgumentOrAlloca(Loc.Ptr);
-    } else {
-      // If AccessesNonArgsOrAlloca has not been updated above, set it
-      // conservatively.
-      AccessesNonArgsOrAlloca |= I.mayReadOrWriteMemory();
     }
 
+    if (!I.mayReadOrWriteMemory())
+      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))
+      continue;
+
+    AccessesNonArgsOrAlloca |= !Loc || !IsArgumentOrAlloca(Loc->Ptr);
+
     // Any remaining instructions need to be taken seriously!  Check if they
     // read or write memory.
     //
diff --git a/llvm/test/Transforms/FunctionAttrs/nosync.ll b/llvm/test/Transforms/FunctionAttrs/nosync.ll
index 64225dad..262334a 100644
--- a/llvm/test/Transforms/FunctionAttrs/nosync.ll
+++ b/llvm/test/Transforms/FunctionAttrs/nosync.ll
@@ -71,7 +71,7 @@ define i8 @test6(i8* %p) {
 
 ; negative case - explicit sync
 define void @test7(i8* %p) {
-; CHECK: Function Attrs: mustprogress nofree norecurse nounwind willreturn
+; CHECK: Function Attrs: argmemonly mustprogress nofree norecurse nounwind willreturn
 ; CHECK-LABEL: @test7(
 ; CHECK-NEXT:    [[TMP1:%.*]] = atomicrmw add i8* [[P:%.*]], i8 0 seq_cst, align 1
 ; CHECK-NEXT:    ret void
@@ -211,7 +211,7 @@ declare void @nosync_function() noinline nounwind uwtable nosync
 define void @call_nosync_function() nounwind uwtable noinline {
 ; CHECK: Function Attrs: noinline nosync nounwind uwtable
 ; CHECK-LABEL: @call_nosync_function(
-; CHECK-NEXT:    tail call void @nosync_function() #[[ATTR9:[0-9]+]]
+; CHECK-NEXT:    tail call void @nosync_function() #[[ATTR10:[0-9]+]]
 ; CHECK-NEXT:    ret void
 ;
   tail call void @nosync_function() noinline nounwind uwtable
@@ -225,7 +225,7 @@ declare void @might_sync() noinline nounwind uwtable
 define void @call_might_sync() nounwind uwtable noinline {
 ; CHECK: Function Attrs: noinline nounwind uwtable
 ; CHECK-LABEL: @call_might_sync(
-; CHECK-NEXT:    tail call void @might_sync() #[[ATTR9]]
+; CHECK-NEXT:    tail call void @might_sync() #[[ATTR10]]
 ; CHECK-NEXT:    ret void
 ;
   tail call void @might_sync() noinline nounwind uwtable
-- 
2.7.4