[NVPTX][AA] Traverse use-def chain to find non-generic addrspace (#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.

Author: AlexMaclean

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,24 @@ 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.
+  auto GetAS = [](const Value *V) -> unsigned {
+    if (const auto *PTy = dyn_cast<PointerType>(V->getType()))
+      return PTy->getAddressSpace();
+    return ADDRESS_SPACE_GENERIC;
+  };
+  while (MaxLookup-- && GetAS(V) == ADDRESS_SPACE_GENERIC) {
+    const Value *NewV = getUnderlyingObject(V, 1);
+    if (NewV == V)
+      break;
+    V = NewV;
+  }
+  return GetAS(V);
+}
+
 static AliasResult::Kind getAliasResult(unsigned AS1, unsigned AS2) {
   if ((AS1 == ADDRESS_SPACE_GENERIC) || (AS2 == ADDRESS_SPACE_GENERIC))
     return AliasResult::MayAlias;
@@ -70,8 +94,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 +111,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;

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
+}