[OpaquePtr] Create API to make a copy of a PointerType with some address space

Some existing places use getPointerElementType() to create a copy of a
pointer type with some new address space.

Reviewed By: dblaikie

Differential Revision: https://reviews.llvm.org/D103429

include/llvm/IR/DerivedTypes.h

@@ -660,6 +660,18 @@ public:
     return PointerType::get(C, 0);
   }
 
+  /// This constructs a pointer type with the same pointee type as input
+  /// PointerType (or opaque pointer is the input PointerType is opaque) and the
+  /// given address space. This is only useful during the opaque pointer
+  /// transition.
+  /// TODO: remove after opaque pointer transition is complete.
+  static PointerType *getWithSamePointeeType(PointerType *PT,
+                                             unsigned AddressSpace) {
+    if (PT->isOpaque())
+      return get(PT->getContext(), AddressSpace);
+    return get(PT->getElementType(), AddressSpace);
+  }
+
   Type *getElementType() const {
     assert(!isOpaque() && "Attempting to get element type of opaque pointer");
     return PointeeTy;

lib/Analysis/ConstantFolding.cpp

@@ -885,8 +885,9 @@ Constant *StripPtrCastKeepAS(Constant *Ptr, Type *&ElemTy,
 
   // Preserve the address space number of the pointer.
   if (NewPtrTy->getAddressSpace() != OldPtrTy->getAddressSpace()) {
-    NewPtrTy = ElemTy->getPointerTo(OldPtrTy->getAddressSpace());
-    Ptr = ConstantExpr::getPointerCast(Ptr, NewPtrTy);
+    Ptr = ConstantExpr::getPointerCast(
+        Ptr, PointerType::getWithSamePointeeType(NewPtrTy,
+                                                 OldPtrTy->getAddressSpace()));
   }
   return Ptr;
 }

lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

@@ -960,8 +960,8 @@ bool AMDGPUPromoteAllocaImpl::handleAlloca(AllocaInst &I, bool SufficientLDS) {
     if (!Call) {
       if (ICmpInst *CI = dyn_cast<ICmpInst>(V)) {
         Value *Src0 = CI->getOperand(0);
-        Type *EltTy = Src0->getType()->getPointerElementType();
-        PointerType *NewTy = PointerType::get(EltTy, AMDGPUAS::LOCAL_ADDRESS);
+        PointerType *NewTy = PointerType::getWithSamePointeeType(
+            cast<PointerType>(Src0->getType()), AMDGPUAS::LOCAL_ADDRESS);
 
         if (isa<ConstantPointerNull>(CI->getOperand(0)))
           CI->setOperand(0, ConstantPointerNull::get(NewTy));
@@ -977,8 +977,8 @@ bool AMDGPUPromoteAllocaImpl::handleAlloca(AllocaInst &I, bool SufficientLDS) {
       if (isa<AddrSpaceCastInst>(V))
         continue;
 
-      Type *EltTy = V->getType()->getPointerElementType();
-      PointerType *NewTy = PointerType::get(EltTy, AMDGPUAS::LOCAL_ADDRESS);
+      PointerType *NewTy = PointerType::getWithSamePointeeType(
+          cast<PointerType>(V->getType()), AMDGPUAS::LOCAL_ADDRESS);
 
       // FIXME: It doesn't really make sense to try to do this for all
       // instructions.
@@ -1035,11 +1035,11 @@ bool AMDGPUPromoteAllocaImpl::handleAlloca(AllocaInst &I, bool SufficientLDS) {
       continue;
     case Intrinsic::objectsize: {
       Value *Src = Intr->getOperand(0);
-      Type *SrcTy = Src->getType()->getPointerElementType();
-      Function *ObjectSize = Intrinsic::getDeclaration(Mod,
-        Intrinsic::objectsize,
-        { Intr->getType(), PointerType::get(SrcTy, AMDGPUAS::LOCAL_ADDRESS) }
-      );
+      Function *ObjectSize = Intrinsic::getDeclaration(
+          Mod, Intrinsic::objectsize,
+          {Intr->getType(),
+           PointerType::getWithSamePointeeType(
+               cast<PointerType>(Src->getType()), AMDGPUAS::LOCAL_ADDRESS)});
 
       CallInst *NewCall = Builder.CreateCall(
           ObjectSize,

lib/Target/NVPTX/NVPTXLowerArgs.cpp

@@ -183,8 +183,8 @@ static void convertToParamAS(Value *OldUser, Value *Param) {
       return NewGEP;
     }
     if (auto *BC = dyn_cast<BitCastInst>(I.OldInstruction)) {
-      auto *NewBCType = BC->getType()->getPointerElementType()->getPointerTo(
-          ADDRESS_SPACE_PARAM);
+      auto *NewBCType = PointerType::getWithSamePointeeType(
+          cast<PointerType>(BC->getType()), ADDRESS_SPACE_PARAM);
       return BitCastInst::Create(BC->getOpcode(), I.NewParam, NewBCType,
                                  BC->getName(), BC);
     }
@@ -315,8 +315,9 @@ void NVPTXLowerArgs::markPointerAsGlobal(Value *Ptr) {
   }
 
   Instruction *PtrInGlobal = new AddrSpaceCastInst(
-      Ptr, PointerType::get(Ptr->getType()->getPointerElementType(),
-                            ADDRESS_SPACE_GLOBAL),
+      Ptr,
+      PointerType::getWithSamePointeeType(cast<PointerType>(Ptr->getType()),
+                                          ADDRESS_SPACE_GLOBAL),
       Ptr->getName(), &*InsertPt);
   Value *PtrInGeneric = new AddrSpaceCastInst(PtrInGlobal, Ptr->getType(),
                                               Ptr->getName(), &*InsertPt);

unittests/IR/TypesTest.cpp

@@ -34,4 +34,27 @@ TEST(TypesTest, LayoutIdenticalEmptyStructs) {
   EXPECT_TRUE(Foo->isLayoutIdentical(Bar));
 }
 
+TEST(TypesTest, CopyPointerType) {
+  LLVMContext C;
+
+  PointerType *P1 = PointerType::get(C, 1);
+  EXPECT_TRUE(P1->isOpaque());
+  PointerType *P1C = PointerType::getWithSamePointeeType(P1, 1);
+  EXPECT_EQ(P1, P1C);
+  EXPECT_TRUE(P1C->isOpaque());
+  PointerType *P1C0 = PointerType::getWithSamePointeeType(P1, 0);
+  EXPECT_NE(P1, P1C0);
+  EXPECT_TRUE(P1C0->isOpaque());
+
+  Type *Int8 = Type::getInt8Ty(C);
+  PointerType *P2 = PointerType::get(Int8, 1);
+  EXPECT_FALSE(P2->isOpaque());
+  PointerType *P2C = PointerType::getWithSamePointeeType(P2, 1);
+  EXPECT_EQ(P2, P2C);
+  EXPECT_FALSE(P2C->isOpaque());
+  PointerType *P2C0 = PointerType::getWithSamePointeeType(P2, 0);
+  EXPECT_NE(P2, P2C0);
+  EXPECT_FALSE(P2C0->isOpaque());
+}
+
 }  // end anonymous namespace