[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:

It looks like for HIP we already emit the kernel with AS 1 pointer arguments. However, is it always the case?

```
#include <cstdint>
#include <hip/hip_runtime.h>

__device__ __constant__ int constNumber[4] = {1, 2, 3, 4};

__global__ void kernel(int *out, int *in) { *out = in[3]; }

int main(int argc, char *argv[]) {
  int out;

  int *out_dev = nullptr;
  hipError_t err = hipMalloc(&out_dev, sizeof(int));
  if (err != hipSuccess)
    return 1;

  kernel<<<1, 1>>>(out_dev, constNumber);

  err = hipMemcpyDtoH(&out, out_dev, sizeof(int));
  if (err != hipSuccess)
    return 2;

  printf("out=%d\n", out);

  return 0;
}
```

The compiler doesn't complain anything, but at runtime it crashes due to memory access fault. The IR shows the kernel signature is:

```
define protected amdgpu_kernel void @_Z6kernelPiS_(ptr addrspace(1) noundef %out.coerce, ptr addrspace(1) noundef %in.coerce)
```

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