[NVPTXAA] Traverse use-def chain to find 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,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;

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