[llvm] [Attributor] Take the address space from addrspacecast directly (PR #108258)

[Matt Arsenault via llvm-commits]

================
@@ -12571,8 +12571,35 @@ struct AAAddressSpaceImpl : public AAAddressSpace {
   void initialize(Attributor &A) override {
     assert(getAssociatedType()->isPtrOrPtrVectorTy() &&
            "Associated value is not a pointer");
-    if (getAssociatedType()->getPointerAddressSpace())
+    // If the pointer already has non-flat address space, we assume it is the
+    // correct one.
+    if (getAssociatedType()->getPointerAddressSpace()) {
+      [[maybe_unused]] bool R =
+          takeAddressSpace(getAssociatedType()->getPointerAddressSpace());
+      assert(R && "the take should happen");
       indicateOptimisticFixpoint();
+      return;
+    }
+    // If the pointer is an addrspacecast, it has to be from a non-flat to flat.
+    // We assume the source address space is the correct one.
+    Value *V = &getAssociatedValue();
+    if (auto *ASCI = dyn_cast<AddrSpaceCastInst>(V)) {
+      assert(ASCI->getDestAddressSpace() == 0 &&
----------------
arsenm wrote:

Don't hardcoded the generic = 0, query TTI 

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