Make SROA handle addrspacecast

[wengxt]

Summary

This is basically D4501 from Matt with fixed testcase. In GPU related
world addrspacecast instructions are commonly used to avoid overhead
when access memory with generic load and store. SROA pass currently does
not take addrspaceCast into account. This patch uses the same way that SROA
handles bitcast to process addrspacecast.

[sanjoy]

lib/Transforms/Scalar/SROA.cpp

1825

Will this transform a "gep(addrspacecast X), Y)" to "addrspacecast(gep X, Y')"? If so, I think we need a target specific hook that does a legality check on the addrspacecast we're about to insert -- I don't think "addrspacecast X" is legal implies an addrspacecast on a GEP derived from X is also legal.

[chandlerc]

I'm actually quite alarmed at the idea of needing to handle addrspacecast of *alloca*s. I feel like their existence is a bigger problem.

I'd like to understand the use case that GPU folks have for this better. For an un-escaped alloca, why are there addrspacecasts at all? Why would they be useful? Why couldn't you lower them with arbitrarily fast loads and stores due to being unescaped allocas?

lib/Transforms/Scalar/SROA.cpp

1825

The whole point of addrspacecast was to avoid the need for target specific checks... I think we just shouldn't transform GEPs around them if that isn't legal everywhere.

[jingyue]

Regarding your question on why these addrspacecasts even exist, the short answer is that emitting addrspacecasts after alloca eases NVPTX's address space inference which lowers access to special address spaces into fast loads/stores.

The long answer requires some background on CUDA and how the NVPTX backend works. I tried to make it brief. Let me know if it makes sense.

First of all, address spaces are not included in CUDA's type system. All pointers in a CUDA source code are generic. Only when a variable is declared, the programmer can annotate the address space that this variable lives in. For example, __shared__ int *a = xxx; means a, the pointer itself, lives in the shared address space rather than a points to an int that lives in the shared space.

So, LLVM's middle-end and/or the NVPTX backend bear the burden of optimizing the memory accesses to variables in special address spaces. Given an address-taken variable annotated with a special address space, the clang front-end simply casts the resultant address to generic (i.e. addrspace(0)) right after the declaration. The optimizer needs to eliminate unnecessary addrspacecasts and emits fast loads/stores.

Currently, the pass that performs this address space optimization in NVPTX is NVPTXFavorNonGenericAddrSpaces. It works by pushing addrspacecast (from special to generic) towards loads and stores as far as possible. After that, the codegen can simply lower load/store <ty>, <ty> addrspace(x)* to fast loads/stores in machine code. For example, if a function reads from a shared variable, the IR clang emits without optimization looks like

@a = addrspace(3) global float
define ... {
  %1 = addrspacecast float addrspace(3)* @a to float* ; cast from special to generic
  %2 = load float, float* %1
  ...
}

then NVPTXFavorNonGenericAddrSpaces optimizes it to

@a = addrspace(3) global float
define ... {
  %2 = load float, float addrspace(3)* @a
  ...
}

and finally NVPTX's codegen lowers this "load from shared" to a ld.shared.f32 instruction.

This pass works well for GlobalVariables whose address spaces are explicitly annotated. However, alloca always returns addrspace(0) pointers, so the clang front-end doesn't (need to) emit any explicit addrspacecast from local to generic in the initial IR. This creates troubles for the address space inference, because NVPTXFavorNonGenericAddrSpaces has no addrspacecasts to push for accesses to the local address space.

To address this issue, we plan to run another pass called NVPTXLowerAlloca (D10483) that emits two addrspacecasts after an alloca, one casting generic to local and the other casting local back to generic. For example,

define ...() {
  %1 = alloca float
  store float 0, float* %1
}

becomes

define ...() {
  %1 = alloca float
  %2 = addrspacecast float* %1 to float addrspace(4)*
  %3 = addrspacecast float addrspace(4)* %2 to float*
  store float 0, float* %3
}

Then, NVPTXFavorNonGenericAddrSpaces pushes the second addrspacecast as usual towards the store.

define ...() {
  %1 = alloca float
  %2 = addrspacecast float* %1 to float addrspace(4)*
  store float 0, float addrspace(4)* %2
}

This is why NVPTX emits addrspacecasts after alloca at some stage.

If having SROA to handle addrspacecast of allocas is too alarming, I think the easiest way to fix the issue on our side is to run SROA between NVPTXLowerKernelArgs (another address space inference pass, which must run before SROA) and NVPTXLowerAlloca. It will work in the short term. However, Justin Holewinski and I have some concerns on how much we depend on correctly ordering these address-space inference passes, and think we should merge all of them into one big pass. If we want that merging to happen, we won't be able to run SROA in the middle. Still, workarounds exist such as merging most of the address space inference and only running a small part before SROA.

lib/Transforms/Scalar/SROA.cpp

1825

Thanks for pointing this out, Sanjoy. I understand your concerns. Curiously, on which existing targets such transformation from gep(addrspacecast X), Y to addrspacecast(gep X, Y) is illegal? Just asking so that I can keep a concrete example in mind when I later write target-independent optimizations on addrspacecast.

[arsenm]

The last time this came up, I was able to get instcombine to eliminate trivial addrspacecasts of alloca which worked just as well. I just never got around to actually testing / posting the patch. The main piece was just removing an outdated check for the address space mismatch from before addrspacecast existed

[sanjoy]

lib/Transforms/Scalar/SROA.cpp

1825

In practice, I think gep(addrspacecast X), Y and addrspacecast(gep X, Y) should compute the same result in all existing targets. However, this caps how complex an addrspacecast implementation can be. For instance, targets won't be able to implement addrspacecast as

addrspacecast(P) { return abs(P) }

since with this change, gep(addrspacecast(-1), 2) != addrspacecast(gep -1, 2).

I'm okay with moving ahead with this gep transform as long as we change the langref to reflect that addrspacecast instructions cannot be arbitrarily complex value modifications (so that an abs implementation for addrspacecast can be ruled out).

include/llvm/Analysis/PtrUseVisitor.h

void visitStoreInst(StoreInst &SI) {
     if (SI.getValueOperand() == U->get())
       PI.setEscaped(&SI);
   }
 
   void visitBitCastInst(BitCastInst &BC) {
     enqueueUsers(BC);
   }
 
+  void visitAddrSpaceCastInst(AddrSpaceCastInst &ASC) {
+    enqueueUsers(ASC);
+  }
+
   void visitPtrToIntInst(PtrToIntInst &I) {
     PI.setEscaped(&I);
   }
 
   void visitGetElementPtrInst(GetElementPtrInst &GEPI) {
     if (GEPI.use_empty())
       return;
 

lib/Transforms/Scalar/SROA.cpp

private:
 
   void visitBitCastInst(BitCastInst &BC) {
     if (BC.use_empty())
       return markAsDead(BC);
 
     return Base::visitBitCastInst(BC);
   }
 
+  void visitAddrSpaceCastInst(AddrSpaceCastInst &ASC) {
+    if (ASC.use_empty())
+      return markAsDead(ASC);
+
+    return Base::visitAddrSpaceCastInst(ASC);
+  }
+
   void visitGetElementPtrInst(GetElementPtrInst &GEPI) {
     if (GEPI.use_empty())
       return markAsDead(GEPI);
 
     if (SROAStrictInbounds && GEPI.isInBounds()) {
       // FIXME: This is a manually un-factored variant of the basic code inside
       // of GEPs with checking of the inbounds invariant specified in the
       // langref in a very strict sense. If we ever want to enable
do {
         Size = std::max(Size, DL.getTypeStoreSize(Op->getType()));
         continue;
       }
 
       if (GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(I)) {
         if (!GEP->hasAllZeroIndices())
           return GEP;
       } else if (!isa<BitCastInst>(I) && !isa<PHINode>(I) &&
-                 !isa<SelectInst>(I)) {
+                 !isa<SelectInst>(I) && !isa<AddrSpaceCastInst>(I)) {
         return I;
       }
 
       for (User *U : I->users())
         if (Visited.insert(cast<Instruction>(U)).second)
           Uses.push_back(std::make_pair(I, cast<Instruction>(U)));
     } while (!Uses.empty());
 
isInstInList(Instruction *I,
     do {
       if (Ptr == &AI)
         return true;
 
       if (BitCastInst *BCI = dyn_cast<BitCastInst>(Ptr))
         Ptr = BCI->getOperand(0);
       else if (GetElementPtrInst *GEPI = dyn_cast<GetElementPtrInst>(Ptr))
         Ptr = GEPI->getPointerOperand();
+      else if (AddrSpaceCastInst *ASC = dyn_cast<AddrSpaceCastInst>(Ptr))
+        Ptr = ASC->getOperand(0);
       else
         return false;
 
     } while (Visited.insert(Ptr).second);
 
     return false;
   }
 
do {
 
     // Stash this pointer if we've found an i8*.
     if (Ptr->getType()->isIntegerTy(8)) {
       Int8Ptr = Ptr;
       Int8PtrOffset = Offset;
     }
 
     // Peel off a layer of the pointer and update the offset appropriately.
-    if (Operator::getOpcode(Ptr) == Instruction::BitCast) {
+    unsigned Opc = Operator::getOpcode(Ptr);
+    if (Opc == Instruction::BitCast || Opc == Instruction::AddrSpaceCast) {
       Ptr = cast<Operator>(Ptr)->getOperand(0);
     } else if (GlobalAlias *GA = dyn_cast<GlobalAlias>(Ptr)) {
       if (GA->mayBeOverridden())
         break;
       Ptr = GA->getAliasee();
     } else {
       break;
     }
OffsetPtr = Int8PtrOffset == 0
                     ? Int8Ptr
                     : IRB.CreateInBoundsGEP(IRB.getInt8Ty(), Int8Ptr,
                                             IRB.getInt(Int8PtrOffset),
                                             NamePrefix + "sroa_raw_idx");
   }
   Ptr = OffsetPtr;
 
   // On the off chance we were targeting i8*, guard the bitcast here.
-  if (Ptr->getType() != PointerTy)
-    Ptr = IRB.CreateBitCast(Ptr, PointerTy, NamePrefix + "sroa_cast");
+  if (Ptr->getType() != PointerTy) {
+    Ptr = IRB.CreatePointerBitCastOrAddrSpaceCast(Ptr, PointerTy,
+                                                  NamePrefix + "sroa_cast");
+  }
 
   return Ptr;
 }
 
 /// \brief Compute the adjusted alignment for a load or store from an offset.
 static unsigned getAdjustedAlignment(Instruction *I, uint64_t Offset,
                                      const DataLayout &DL) {
   unsigned Alignment;
bool visitStoreInst(StoreInst &SI) {
     return true;
   }
 
   bool visitBitCastInst(BitCastInst &BC) {
     enqueueUsers(BC);
     return false;
   }
 
+  bool visitAddrSpaceCastInst(AddrSpaceCastInst &ASC) {
+    enqueueUsers(ASC);
+    return false;
+  }
+
   bool visitGetElementPtrInst(GetElementPtrInst &GEPI) {
     enqueueUsers(GEPI);
     return false;
   }
 
   bool visitPHINode(PHINode &PN) {
     enqueueUsers(PN);
     return false;

test/Transforms/SROA/basictest.ll

entry:
   store i64 %X, i64* %B
   br label %L2
 
 L2:
   %Z = load i64, i64* %B
   ret i64 %Z
 }
 
+define i64 @test2_addrspacecast(i64 %X) {
+; CHECK-LABEL: @test2_addrspacecast(
+; CHECK-NOT: alloca
+; CHECK: ret i64 %X
+
+entry:
+  %A = alloca [8 x i8]
+  %B = addrspacecast [8 x i8]* %A to i64 addrspace(1)*
+  store i64 %X, i64 addrspace(1)* %B
+  br label %L2
+
+L2:
+  %Z = load i64, i64 addrspace(1)* %B
+  ret i64 %Z
+}
+
 define void @test3(i8* %dst, i8* %src) {
 ; CHECK-LABEL: @test3(
 
 entry:
   %a = alloca [300 x i8]
 ; CHECK-NOT:  alloca
 ; CHECK:      %[[test3_a1:.*]] = alloca [42 x i8]
 ; CHECK-NEXT: %[[test3_a2:.*]] = alloca [99 x i8]
entry:
   %cast1 = bitcast %opaque* %x to i8*
   %cast2 = bitcast { i64, i8* }* %a to i8*
   call void @llvm.memcpy.p0i8.p0i8.i32(i8* %cast2, i8* %cast1, i32 16, i32 1, i1 false)
   %gep = getelementptr inbounds { i64, i8* }, { i64, i8* }* %a, i32 0, i32 0
   %val = load i64, i64* %gep
   ret i32 undef
 }
 
+declare void @llvm.memcpy.p0i8.p1i8.i32(i8* nocapture, i8 addrspace(1)* nocapture, i32, i32, i1) nounwind
+
+define i32 @test19_addrspacecast(%opaque* %x) {
+; This input will cause us to try to compute a natural GEP when rewriting
+; pointers in such a way that we try to GEP through the opaque type. Previously,
+; a check for an unsized type was missing and this crashed. Ensure it behaves
+; reasonably now.
+; CHECK-LABEL: @test19_addrspacecast(
+; CHECK-NOT: alloca
+; CHECK: ret i32 undef
+
+entry:
+  %a = alloca { i64, i8* }
+  %cast1 = addrspacecast %opaque* %x to i8 addrspace(1)*
+  %cast2 = bitcast { i64, i8* }* %a to i8*
+  call void @llvm.memcpy.p0i8.p1i8.i32(i8* %cast2, i8 addrspace(1)* %cast1, i32 16, i32 1, i1 false)
+  %gep = getelementptr inbounds { i64, i8* }, { i64, i8* }* %a, i32 0, i32 0
+  %val = load i64, i64* %gep
+  ret i32 undef
+}
+
 define i32 @test20() {
 ; Ensure we can track negative offsets (before the beginning of the alloca) and
 ; negative relative offsets from offsets starting past the end of the alloca.
 ; CHECK-LABEL: @test20(
 ; CHECK-NOT: alloca
 ; CHECK: %[[sum1:.*]] = add i32 1, 2
 ; CHECK: %[[sum2:.*]] = add i32 %[[sum1]], 3
 ; CHECK: ret i32 %[[sum2]]