[llvm] [Attributor]: allow more than one offset for a pass through user in AAPointerInfo (PR #94416)

[Vidush Singhal via llvm-commits]

https://github.com/vidsinghal updated https://github.com/llvm/llvm-project/pull/94416

>From a5d76876f5a43e7a6c01afed4127855905b33b6b Mon Sep 17 00:00:00 2001
From: Vidush Singhal <singhal2 at ruby964.llnl.gov>
Date: Tue, 4 Jun 2024 16:51:06 -0700
Subject: [PATCH] [Attributor]: AApointerinfo, allow more than one offset for a
 pass through user

---
 .../Transforms/IPO/AttributorAttributes.cpp   |  6 +++-
 .../Transforms/Attributor/pointer-info.ll     | 30 ++++++++++++++++---
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Transforms/IPO/AttributorAttributes.cpp b/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
index 9a5732dca5b79..3fc6e4a5ca54b 100644
--- a/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
+++ b/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
@@ -1607,11 +1607,15 @@ ChangeStatus AAPointerInfoFloating::updateImpl(Attributor &A) {
     //
     // The RHS is a reference that may be invalidated by an insertion caused by
     // the LHS. So we ensure that the side-effect of the LHS happens first.
+
+    assert(OffsetInfoMap.contains(CurPtr) &&
+           "CurPtr does not exist in the map!");
+
     auto &UsrOI = OffsetInfoMap[Usr];
     auto &PtrOI = OffsetInfoMap[CurPtr];
     assert(!PtrOI.isUnassigned() &&
            "Cannot pass through if the input Ptr was not visited!");
-    UsrOI = PtrOI;
+    UsrOI.merge(PtrOI);
     Follow = true;
     return true;
   };
diff --git a/llvm/test/Transforms/Attributor/pointer-info.ll b/llvm/test/Transforms/Attributor/pointer-info.ll
index 6afdbdaee317c..f4cebbbb05b99 100644
--- a/llvm/test/Transforms/Attributor/pointer-info.ll
+++ b/llvm/test/Transforms/Attributor/pointer-info.ll
@@ -14,7 +14,7 @@ define void @foo(ptr %ptr) {
 ; TUNIT-NEXT:    br label [[CALL_BR:%.*]]
 ; TUNIT:       call.br:
 ; TUNIT-NEXT:    [[TMP1:%.*]] = getelementptr inbounds [[STRUCT_TEST_A]], ptr [[TMP0]], i64 0, i32 2
-; TUNIT-NEXT:    tail call void @bar(ptr noalias nocapture nofree noundef nonnull readonly byval([[STRUCT_TEST_A]]) align 8 dereferenceable(24) [[TMP0]]) #[[ATTR2:[0-9]+]]
+; TUNIT-NEXT:    tail call void @bar(ptr noalias nocapture nofree noundef nonnull readonly byval([[STRUCT_TEST_A]]) align 8 dereferenceable(24) [[TMP0]]) #[[ATTR3:[0-9]+]]
 ; TUNIT-NEXT:    ret void
 ;
 ; CGSCC: Function Attrs: mustprogress nofree nosync nounwind willreturn memory(none)
@@ -26,7 +26,7 @@ define void @foo(ptr %ptr) {
 ; CGSCC:       call.br:
 ; CGSCC-NEXT:    [[TMP1:%.*]] = getelementptr inbounds [[STRUCT_TEST_A]], ptr [[TMP0]], i64 0, i32 2
 ; CGSCC-NEXT:    store ptr [[PTR]], ptr [[TMP1]], align 8
-; CGSCC-NEXT:    tail call void @bar(ptr noalias nocapture nofree noundef nonnull readnone byval([[STRUCT_TEST_A]]) align 8 dereferenceable(24) [[TMP0]]) #[[ATTR2:[0-9]+]]
+; CGSCC-NEXT:    tail call void @bar(ptr noalias nocapture nofree noundef nonnull readnone byval([[STRUCT_TEST_A]]) align 8 dereferenceable(24) [[TMP0]]) #[[ATTR3:[0-9]+]]
 ; CGSCC-NEXT:    ret void
 ;
 entry:
@@ -52,12 +52,34 @@ define void @bar(ptr noundef byval(%struct.test.a) align 8 %dev) {
   store i32 1, ptr %1
   ret void
 }
+
+
+define void @multiple_access_bins_for_select_when_pass_through_instruction(i1 %cond) {
+; CHECK: Function Attrs: mustprogress nofree norecurse nosync nounwind willreturn memory(write)
+; CHECK-LABEL: define {{[^@]+}}@multiple_access_bins_for_select_when_pass_through_instruction
+; CHECK-SAME: (i1 [[COND:%.*]]) #[[ATTR2:[0-9]+]] {
+; CHECK-NEXT:    [[TMP1:%.*]] = alloca [10 x i8], align 4
+; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds [10 x i8], ptr [[TMP1]], i64 0, i64 5
+; CHECK-NEXT:    [[TMP3:%.*]] = getelementptr inbounds [10 x i8], ptr [[TMP1]], i64 0, i64 2
+; CHECK-NEXT:    [[TMP4:%.*]] = select i1 [[COND]], ptr [[TMP2]], ptr [[TMP3]]
+; CHECK-NEXT:    store i8 100, ptr [[TMP4]], align 1
+; CHECK-NEXT:    ret void
+;
+%2 = alloca [10 x i8], align 4
+%3 = getelementptr inbounds [10 x i8], ptr %2, i64 0, i64 5
+%4 = getelementptr inbounds [10 x i8], ptr %2, i64 0, i64 2
+%5 = select i1 %cond, ptr %3, ptr %4
+store i8 100, ptr %5
+ret void
+}
 ;.
 ; TUNIT: attributes #[[ATTR0]] = { mustprogress nofree norecurse nosync nounwind willreturn memory(none) }
 ; TUNIT: attributes #[[ATTR1]] = { mustprogress nofree norecurse nosync nounwind willreturn memory(argmem: write) }
-; TUNIT: attributes #[[ATTR2]] = { nofree nosync nounwind willreturn memory(write) }
+; TUNIT: attributes #[[ATTR2]] = { mustprogress nofree norecurse nosync nounwind willreturn memory(write) }
+; TUNIT: attributes #[[ATTR3]] = { nofree nosync nounwind willreturn memory(write) }
 ;.
 ; CGSCC: attributes #[[ATTR0]] = { mustprogress nofree nosync nounwind willreturn memory(none) }
 ; CGSCC: attributes #[[ATTR1]] = { mustprogress nofree norecurse nosync nounwind willreturn memory(argmem: write) }
-; CGSCC: attributes #[[ATTR2]] = { nofree nounwind willreturn memory(write) }
+; CGSCC: attributes #[[ATTR2]] = { mustprogress nofree norecurse nosync nounwind willreturn memory(write) }
+; CGSCC: attributes #[[ATTR3]] = { nofree nounwind willreturn memory(write) }
 ;.