[Attributor][NFC] Ignore benign uses in AAMemoryBehaviorFloating

In AAMemoryBehaviorFloating we used to track benign uses in a SetVector.
With this change we look through benign uses eagerly to reduce the
number of elements (=Uses) we look at during an update.

The test does actually not fail prior to this commit but I already wrote
it so I kept it.

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

@@ -5861,9 +5861,7 @@ struct AAMemoryBehaviorFloating : AAMemoryBehaviorImpl {
   /// See AbstractAttribute::initialize(...).
   void initialize(Attributor &A) override {
     AAMemoryBehaviorImpl::initialize(A);
-    // Initialize the use vector with all direct uses of the associated value.
+    addUsesOf(A, getAssociatedValue());
-    for (const Use &U : getAssociatedValue().uses())
-      Uses.insert(&U);
   }
 
   /// See AbstractAttribute::updateImpl(...).
@@ -5889,8 +5887,14 @@ private:
   void analyzeUseIn(Attributor &A, const Use *U, const Instruction *UserI);
 
 protected:
+  /// Add the uses of \p V to the `Uses` set we look at during the update step.
+  void addUsesOf(Attributor &A, const Value &V);
+
   /// Container for (transitive) uses of the associated argument.
-  SetVector<const Use *> Uses;
+  SmallVector<const Use *, 8> Uses;
+
+  /// Set to remember the uses we already traversed.
+  SmallPtrSet<const Use *, 8> Visited;
 };
 
 /// Memory behavior attribute for function argument.
@@ -5915,9 +5919,7 @@ struct AAMemoryBehaviorArgument : AAMemoryBehaviorFloating {
     if (!Arg || !A.isFunctionIPOAmendable(*(Arg->getParent()))) {
       indicatePessimisticFixpoint();
     } else {
-      // Initialize the use vector with all direct uses of the associated value.
+      addUsesOf(A, *Arg);
-      for (const Use &U : Arg->uses())
-        Uses.insert(&U);
     }
   }
 
@@ -6169,8 +6171,7 @@ ChangeStatus AAMemoryBehaviorFloating::updateImpl(Attributor &A) {
 
     // Check if the users of UserI should also be visited.
     if (followUsersOfUseIn(A, U, UserI))
-      for (const Use &UserIUse : UserI->uses())
+      addUsesOf(A, *UserI);
-        Uses.insert(&UserIUse);
 
     // If UserI might touch memory we analyze the use in detail.
     if (UserI->mayReadOrWriteMemory())
@@ -6181,6 +6182,28 @@ ChangeStatus AAMemoryBehaviorFloating::updateImpl(Attributor &A) {
                                         : ChangeStatus::UNCHANGED;
 }
 
+void AAMemoryBehaviorFloating::addUsesOf(Attributor &A, const Value &V) {
+  SmallVector<const Use *, 8> WL;
+  for (const Use &U : V.uses())
+    WL.push_back(&U);
+
+  while (!WL.empty()) {
+    const Use *U = WL.pop_back_val();
+    if (!Visited.insert(U).second)
+      continue;
+
+    const Instruction *UserI = cast<Instruction>(U->getUser());
+    if (UserI->mayReadOrWriteMemory()) {
+      Uses.push_back(U);
+      continue;
+    }
+    if (!followUsersOfUseIn(A, U, UserI))
+      continue;
+    for (const Use &UU : UserI->uses())
+      WL.push_back(&UU);
+  }
+}
+
 bool AAMemoryBehaviorFloating::followUsersOfUseIn(Attributor &A, const Use *U,
                                                   const Instruction *UserI) {
   // The loaded value is unrelated to the pointer argument, no need to

llvm/test/Transforms/Attributor/readattrs.ll

@@ -403,3 +403,49 @@ define void @ptr_uses(i8* %ptr) {
   call void @val_use(i8 %call_val)
   ret void
 }
+
+define void @ptr_use_chain(i8* %ptr) {
+; CHECK-LABEL: define {{[^@]+}}@ptr_use_chain
+; CHECK-SAME: (i8* [[PTR:%.*]])
+; CHECK-NEXT:    [[BC0:%.*]] = bitcast i8* [[PTR]] to i32*
+; CHECK-NEXT:    [[BC1:%.*]] = bitcast i32* [[BC0]] to i8*
+; CHECK-NEXT:    [[BC2:%.*]] = bitcast i8* [[BC1]] to i32*
+; CHECK-NEXT:    [[BC3:%.*]] = bitcast i32* [[BC2]] to i8*
+; CHECK-NEXT:    [[BC4:%.*]] = bitcast i8* [[BC3]] to i32*
+; CHECK-NEXT:    [[BC5:%.*]] = bitcast i32* [[BC4]] to i8*
+; CHECK-NEXT:    [[BC6:%.*]] = bitcast i8* [[BC5]] to i32*
+; CHECK-NEXT:    [[BC7:%.*]] = bitcast i32* [[BC6]] to i8*
+; CHECK-NEXT:    [[BC8:%.*]] = bitcast i8* [[BC7]] to i32*
+; CHECK-NEXT:    [[BC9:%.*]] = bitcast i32* [[BC8]] to i8*
+; CHECK-NEXT:    [[ABC2:%.*]] = bitcast i8* [[BC9]] to i32*
+; CHECK-NEXT:    [[ABC3:%.*]] = bitcast i32* [[ABC2]] to i8*
+; CHECK-NEXT:    [[ABC4:%.*]] = bitcast i8* [[ABC3]] to i32*
+; CHECK-NEXT:    [[ABC5:%.*]] = bitcast i32* [[ABC4]] to i8*
+; CHECK-NEXT:    [[ABC6:%.*]] = bitcast i8* [[ABC5]] to i32*
+; CHECK-NEXT:    [[ABC7:%.*]] = bitcast i32* [[ABC6]] to i8*
+; CHECK-NEXT:    [[ABC8:%.*]] = bitcast i8* [[ABC7]] to i32*
+; CHECK-NEXT:    [[ABC9:%.*]] = bitcast i32* [[ABC8]] to i8*
+; CHECK-NEXT:    call void @escape_i8(i8* [[ABC9]])
+; CHECK-NEXT:    ret void
+;
+  %bc0 = bitcast i8* %ptr to i32*
+  %bc1 = bitcast i32* %bc0 to i8*
+  %bc2 = bitcast i8* %bc1 to i32*
+  %bc3 = bitcast i32* %bc2 to i8*
+  %bc4 = bitcast i8* %bc3 to i32*
+  %bc5 = bitcast i32* %bc4 to i8*
+  %bc6 = bitcast i8* %bc5 to i32*
+  %bc7 = bitcast i32* %bc6 to i8*
+  %bc8 = bitcast i8* %bc7 to i32*
+  %bc9 = bitcast i32* %bc8 to i8*
+  %abc2 = bitcast i8* %bc9 to i32*
+  %abc3 = bitcast i32* %abc2 to i8*
+  %abc4 = bitcast i8* %abc3 to i32*
+  %abc5 = bitcast i32* %abc4 to i8*
+  %abc6 = bitcast i8* %abc5 to i32*
+  %abc7 = bitcast i32* %abc6 to i8*
+  %abc8 = bitcast i8* %abc7 to i32*
+  %abc9 = bitcast i32* %abc8 to i8*
+  call void @escape_i8(i8* %abc9)
+  ret void
+}