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

[Alex MacLean via llvm-commits]

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

Address space information may be encoded anywhere along the use-def chain. take advantage of this by traversing the chain until we find a non-generic addrspace.

>From d8dcccf4ef363786355424cb638120a155aad74b Mon Sep 17 00:00:00 2001
From: Alex MacLean <amaclean at nvidia.com>
Date: Fri, 23 Aug 2024 16:18:46 +0000
Subject: [PATCH] [NVPTXAA] Traverse use-def chain to find non-generic
 addrspace

---
 llvm/lib/Target/NVPTX/NVPTXAliasAnalysis.cpp | 38 ++++++++++++++++----
 llvm/test/CodeGen/NVPTX/nvptx-aa.ll          | 37 +++++++++++++++++++
 2 files changed, 68 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Target/NVPTX/NVPTXAliasAnalysis.cpp b/llvm/lib/Target/NVPTX/NVPTXAliasAnalysis.cpp
index 4f106584eb0a94..e86dfee875fd89 100644
--- a/llvm/lib/Target/NVPTX/NVPTXAliasAnalysis.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXAliasAnalysis.cpp
@@ -15,11 +15,17 @@
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/CallingConv.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/Support/CommandLine.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "NVPTX-aa"
 
+static cl::opt<unsigned> TraverseAddressSpacesLimit(
+    "nvptx-traverse-address-aliasing-limit", cl::Hidden,
+    cl::desc("Depth limit for finding address space through traversal"),
+    cl::init(6));
+
 AnalysisKey NVPTXAA::Key;
 
 char NVPTXAAWrapperPass::ID = 0;
@@ -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;
+}
+
 static AliasResult::Kind getAliasResult(unsigned AS1, unsigned AS2) {
   if ((AS1 == ADDRESS_SPACE_GENERIC) || (AS2 == ADDRESS_SPACE_GENERIC))
     return AliasResult::MayAlias;
@@ -70,8 +98,8 @@ static AliasResult::Kind getAliasResult(unsigned AS1, unsigned AS2) {
 AliasResult NVPTXAAResult::alias(const MemoryLocation &Loc1,
                                  const MemoryLocation &Loc2, AAQueryInfo &AAQI,
                                  const Instruction *) {
-  unsigned AS1 = Loc1.Ptr->getType()->getPointerAddressSpace();
-  unsigned AS2 = Loc2.Ptr->getType()->getPointerAddressSpace();
+  unsigned AS1 = getAddressSpace(Loc1.Ptr, TraverseAddressSpacesLimit);
+  unsigned AS2 = getAddressSpace(Loc2.Ptr, TraverseAddressSpacesLimit);
 
   return getAliasResult(AS1, AS2);
 }
@@ -87,11 +115,7 @@ static bool isConstOrParam(unsigned AS) {
 ModRefInfo NVPTXAAResult::getModRefInfoMask(const MemoryLocation &Loc,
                                             AAQueryInfo &AAQI,
                                             bool IgnoreLocals) {
-  if (isConstOrParam(Loc.Ptr->getType()->getPointerAddressSpace()))
-    return ModRefInfo::NoModRef;
-
-  const Value *Base = getUnderlyingObject(Loc.Ptr);
-  if (isConstOrParam(Base->getType()->getPointerAddressSpace()))
+  if (isConstOrParam(getAddressSpace(Loc.Ptr, TraverseAddressSpacesLimit)))
     return ModRefInfo::NoModRef;
 
   return ModRefInfo::ModRef;
diff --git a/llvm/test/CodeGen/NVPTX/nvptx-aa.ll b/llvm/test/CodeGen/NVPTX/nvptx-aa.ll
index 79bd06ee677ac0..074e741dc3e949 100644
--- a/llvm/test/CodeGen/NVPTX/nvptx-aa.ll
+++ b/llvm/test/CodeGen/NVPTX/nvptx-aa.ll
@@ -111,3 +111,40 @@ loop:
 done:
   ret i8 %v2
 }
+
+;; Address space information may be encoded anywhere along the UD chain.
+;; We define a set of tests that:
+;;  1. Perform some number of address space casts on pointer A and B
+;;  2. Store a value to address A
+;;  3. Store a value to address B (that we know does not alias with A)
+
+;; generic->space
+; CHECK-ALIAS-LABEL: Function: test_traversal_gen_space
+; CHECK-ALIAS: NoAlias: i32 addrspace(1)* %global, i32 addrspace(5)* %local
+define void @test_traversal_gen_space(ptr %gen, ptr addrspace(1) %global) {
+  %local = addrspacecast ptr %gen to ptr addrspace(5)
+  store i32 1, ptr addrspace(5) %local, align 8
+  store i32 5, ptr addrspace(1) %global, align 8
+  ret void
+}
+
+;; space->generic
+; CHECK-ALIAS-LABEL: Function: test_traversal_space_gen
+; CHECK-ALIAS: NoAlias: i32* %gen, i32 addrspace(1)* %global
+define void @test_traversal_space_gen(ptr addrspace(5) %local, ptr addrspace(1) %global) {
+  %gen = addrspacecast ptr addrspace(5) %local to ptr
+  store i32 2, ptr %gen, align 8
+  store i32 5, ptr addrspace(1) %global, align 8
+  ret void
+}
+
+;; generic->space->generic
+; CHECK-ALIAS-LABEL: Function: test_traversal_gen_space_gen
+; CHECK-ALIAS: NoAlias: i32* %gen2, i32 addrspace(1)* %global
+define void @test_traversal_gen_space_gen(ptr %gen1, ptr addrspace(1) %global) {
+  %local = addrspacecast ptr %gen1 to ptr addrspace(5)
+  %gen2 = addrspacecast ptr addrspace(5) %local to ptr
+  store i32 3, ptr %gen2, align 8
+  store i32 5, ptr addrspace(1) %global, align 8
+  ret void
+}