[AMDGPU] Fix spurious NoAlias results

[frasercrmck]

After a30e50f, AMDGPUAAResult is being called in more situations where BasicAA isn't sure. This exposed some regressions where NoAlias is being incorrectly returned for two identical pointers.

The fix is to check the underlying objects for equality before returning NoAlias.

[llvmbot]

@llvm/pr-subscribers-backend-amdgpu

Author: Fraser Cormack (frasercrmck)

Changes

After a30e50f, AMDGPUAAResult is being called in more situations where BasicAA isn't sure. This exposed some regressions where NoAlias is being incorrectly returned for two identical pointers.

The fix is to check the underlying objects for equality before returning NoAlias.

---

Full diff: https://github.com/llvm/llvm-project/pull/122309.diff

2 Files Affected:

• (modified) llvm/lib/Target/AMDGPU/AMDGPUAliasAnalysis.cpp (+5-2)
• (modified) llvm/test/CodeGen/AMDGPU/amdgpu-alias-analysis.ll (+17)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUAliasAnalysis.cpp b/llvm/lib/Target/AMDGPU/AMDGPUAliasAnalysis.cpp
index 8d3eac6868318e..c2177c18f3cc3d 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUAliasAnalysis.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUAliasAnalysis.cpp
@@ -80,10 +80,13 @@ AliasResult AMDGPUAAResult::alias(const MemoryLocation &LocA,
     } else if (const Argument *Arg = dyn_cast<Argument>(ObjA)) {
       const Function *F = Arg->getParent();
       switch (F->getCallingConv()) {
-      case CallingConv::AMDGPU_KERNEL:
+      case CallingConv::AMDGPU_KERNEL: {
         // In the kernel function, kernel arguments won't alias to (local)
         // variables in shared or private address space.
-        return AliasResult::NoAlias;
+        const auto *ObjB =
+            getUnderlyingObject(B.Ptr->stripPointerCastsForAliasAnalysis());
+        return (ObjA == ObjB) ? AliasResult::MustAlias : AliasResult::NoAlias;
+      }
       default:
         // TODO: In the regular function, if that local variable in the
         // location B is not captured, that argument pointer won't alias to it
diff --git a/llvm/test/CodeGen/AMDGPU/amdgpu-alias-analysis.ll b/llvm/test/CodeGen/AMDGPU/amdgpu-alias-analysis.ll
index a13eb5c6d085f8..851c18dfccf5af 100644
--- a/llvm/test/CodeGen/AMDGPU/amdgpu-alias-analysis.ll
+++ b/llvm/test/CodeGen/AMDGPU/amdgpu-alias-analysis.ll
@@ -318,3 +318,20 @@ define void @test_9_9(ptr addrspace(9) %p, ptr addrspace(9) %p1) {
   load i8, ptr addrspace(9) %p1
   ret void
 }
+
+; CHECK-LABEL: Function: test_kernel_arg_local_ptr
+; CHECK: MayAlias:    i32 addrspace(3)* %arg, i32 addrspace(3)* %arg1
+; CHECK: MustAlias:   i32 addrspace(3)* %arg, i32* %arg2
+; CHECK: MustAlias:   i32 addrspace(3)* %arg1, i32* %arg2
+define amdgpu_kernel void @test_kernel_arg_local_ptr(ptr addrspace(3) noundef align 4 %arg) {
+entry:
+  %load1 = load i32, ptr addrspace(3) %arg, align 4
+  %arg.plus.1 = getelementptr inbounds nuw i8, ptr addrspace(3) %arg, i64 1
+  %arg1 = getelementptr inbounds nuw i8, ptr addrspace(3) %arg.plus.1, i64 -1
+  %load2 = load i32, ptr addrspace(3) %arg1, align 4
+  %arg.plus.4 = getelementptr inbounds nuw i8, ptr addrspace(3) %arg, i64 4
+  %acast = addrspacecast ptr addrspace(3) %arg.plus.4 to ptr
+  %arg2 = getelementptr inbounds i8, ptr %acast, i64 -4
+  %load3 = load i32, ptr %arg2, align 4
+  ret void
+}

[frasercrmck]

I wasn't sure if the same fix needs done on the other case here:

    if (const LoadInst *LI = dyn_cast<LoadInst>(ObjA)) {
      // If a generic pointer is loaded from the constant address space, it
      // could only be a GLOBAL or CONSTANT one as that address space is solely
      // prepared on the host side, where only GLOBAL or CONSTANT variables are
      // visible. Note that this even holds for regular functions.
      if (LI->getPointerAddressSpace() == AMDGPUAS::CONSTANT_ADDRESS)
        return AliasResult::NoAlias;

or whether the comment can be relied upon. I thought I'd defer it to the PR discussion.

[nikic]

This looks good to me from an AA perspective. Can't comment on whether this matches AMDGPU semantics :)

[frasercrmck]

ping, thanks