[llvm] [NVPTX][AA] Traverse use-def chain to find non-generic addrspace (PR #106477)

[Artem Belevich via llvm-commits]

================
@@ -47,6 +53,28 @@ void NVPTXAAWrapperPass::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.setPreservesAll();
 }
 
+static unsigned getAddressSpace(const Value *V, unsigned MaxLookup) {
+  // Find the first non-generic address space traversing the UD chain.
+  // It is undefined behaviour if a pointer belongs to more than one
+  // non-overlapping address spaces along a valid execution path.
+  for (unsigned Count = 0; MaxLookup == 0 || Count < MaxLookup; ++Count) {
+    const auto *PTy = dyn_cast<PointerType>(V->getType());
+    if (!PTy)
+      return AddressSpace::ADDRESS_SPACE_GENERIC;
+
+    const unsigned AS = PTy->getAddressSpace();
+    if (AS != AddressSpace::ADDRESS_SPACE_GENERIC)
+      return AS;
+
+    // Continue traversing if address space is generic
+    const Value *VNext = getUnderlyingObject(V, 1);
+    if (VNext == V)
+      return AddressSpace::ADDRESS_SPACE_GENERIC;
+    V = VNext;
+  }
+  return AddressSpace::ADDRESS_SPACE_GENERIC;
----------------
Artem-B wrote:

I think the loop logic is flawed. If `V` is set in the last loop iteration to something in non-generic AS, the assertion will fail.

I think you may want to extract "get Value's AS" into a helper lambda, use it as part of the loop condition, and apply it to the last V's value if we terminate the loop without returning early.

https://github.com/llvm/llvm-project/pull/106477