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

[Shilei Tian via llvm-branch-commits]

================
@@ -12571,17 +12571,59 @@ 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-generic 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, we assume the source address space is
+    // the correct one.
+    Value *V = &getAssociatedValue();
+    if (auto *ASC = dyn_cast<AddrSpaceCastInst>(V)) {
+      [[maybe_unused]] bool R = takeAddressSpace(ASC->getSrcAddressSpace());
+      assert(R && "the take should happen");
+      indicateOptimisticFixpoint();
+      return;
+    }
+    if (auto *C = dyn_cast<ConstantExpr>(V)) {
+      if (C->getOpcode() == Instruction::AddrSpaceCast) {
+        [[maybe_unused]] bool R = takeAddressSpace(
+            C->getOperand(0)->getType()->getPointerAddressSpace());
+        assert(R && "the take should happen");
+        indicateOptimisticFixpoint();
+        return;
+      }
+    }
   }
 
   ChangeStatus updateImpl(Attributor &A) override {
-    int32_t OldAddressSpace = AssumedAddressSpace;
+    uint32_t OldAddressSpace = AssumedAddressSpace;
     auto *AUO = A.getOrCreateAAFor<AAUnderlyingObjects>(getIRPosition(), this,
                                                         DepClassTy::REQUIRED);
     auto Pred = [&](Value &Obj) {
       if (isa<UndefValue>(&Obj))
         return true;
+      // If an argument in generic address space has addrspace cast uses, and
+      // those casts are same, then we take the dst addrspace.
+      if (auto *Arg = dyn_cast<Argument>(&Obj)) {
----------------
shiltian wrote:

If that's the case, the `AMDGPUPromoteKernelArgumentsPass` should not exist at the first place.
Both Flang and Clang (for OpenMP target offloading) still generate generic pointer.
HIP and OpenCL front ends generate AS1 pointer, which is convenient.

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