Allow bitcast + struct GEP transform to work with addrspacecast

[arsenm]

Summary

Make getelementptr (addrspacecast (bitcast X)) behave similarly to getelementptr (bitcast X)

[arsenm]

Ping

[rnk]

lgtm

lib/Transforms/InstCombine/InstructionCombining.cpp

1678

Indentation seems off, but it could be phab.

1715

I guess it depends on your perspective, but I consider the addrspacecast case exceptional, so I'd handle it first and return early for it.

1733

Ditto with this condition.

[arsenm]

Committed with comment formatting fixed + inverted conditions in r215467

lib/Transforms/InstCombine/InstructionCombining.cpp

if (HasZeroPointerIndex) {
         }
       }
     }
   }
 
   if (!DL)
     return nullptr;
 
+  // addrspacecast between types is canonicalized as a bitcast, then an
+  // addrspacecast. To take advantage of the below bitcast + struct GEP, look
+  // through the addrspacecast.
+  if (AddrSpaceCastInst *ASC = dyn_cast<AddrSpaceCastInst>(PtrOp)) {
+   ///   X = bitcast A addrspace(1)* to B addrspace(1)*
+   ///   Y = addrspacecast A addrspace(1)* to B addrspace(2)*
+   ///   Z = gep Y, <...constant indices...>
+   /// Into an addrspacecasted GEP of the struct.
+    if (BitCastInst *BC = dyn_cast<BitCastInst>(ASC->getOperand(0)))
+      PtrOp = BC;
+  }
+
   /// See if we can simplify:
   ///   X = bitcast A* to B*
   ///   Y = gep X, <...constant indices...>
   /// into a gep of the original struct.  This is important for SROA and alias
   /// analysis of unions.  If "A" is also a bitcast, wait for A/X to be merged.
   if (BitCastInst *BCI = dyn_cast<BitCastInst>(PtrOp)) {
     Value *Operand = BCI->getOperand(0);
     PointerType *OpType = cast<PointerType>(Operand->getType());
-    unsigned OffsetBits = DL->getPointerTypeSizeInBits(OpType);
+    unsigned OffsetBits = DL->getPointerTypeSizeInBits(GEP.getType());
     APInt Offset(OffsetBits, 0);
     if (!isa<BitCastInst>(Operand) &&
-        GEP.accumulateConstantOffset(*DL, Offset) &&
-        StrippedPtrTy->getAddressSpace() == GEP.getPointerAddressSpace()) {
+        GEP.accumulateConstantOffset(*DL, Offset)) {
 
       // If this GEP instruction doesn't move the pointer, just replace the GEP
       // with a bitcast of the real input to the dest type.
       if (!Offset) {
         // If the bitcast is of an allocation, and the allocation will be
         // converted to match the type of the cast, don't touch this.
         if (isa<AllocaInst>(Operand) || isAllocationFn(Operand, TLI)) {
           // See if the bitcast simplifies, if so, don't nuke this GEP yet.
           if (Instruction *I = visitBitCast(*BCI)) {
             if (I != BCI) {
               I->takeName(BCI);
               BCI->getParent()->getInstList().insert(BCI, I);
               ReplaceInstUsesWith(*BCI, I);
             }
             return &GEP;
           }
         }
+
+        if (Operand->getType()->getPointerAddressSpace() == GEP.getAddressSpace())
           return new BitCastInst(Operand, GEP.getType());
+        return new AddrSpaceCastInst(Operand, GEP.getType());
       }
 
       // Otherwise, if the offset is non-zero, we need to find out if there is a
       // field at Offset in 'A's type.  If so, we can pull the cast through the
       // GEP.
       SmallVector<Value*, 8> NewIndices;
       if (FindElementAtOffset(OpType, Offset.getSExtValue(), NewIndices)) {
         Value *NGEP = GEP.isInBounds() ?
           Builder->CreateInBoundsGEP(Operand, NewIndices) :
           Builder->CreateGEP(Operand, NewIndices);
 
         if (NGEP->getType() == GEP.getType())
           return ReplaceInstUsesWith(GEP, NGEP);
         NGEP->takeName(&GEP);
+
+        if (NGEP->getType()->getPointerAddressSpace() == GEP.getAddressSpace())
           return new BitCastInst(NGEP, GEP.getType());
+        return new AddrSpaceCastInst(NGEP, GEP.getType());
       }
     }
   }
 
   return nullptr;
 }
 
 static bool

test/Transforms/InstCombine/cast.ll

 define i32* @test41(i32* %tmp1) {
         %tmp64 = bitcast i32* %tmp1 to { i32 }*
         %tmp65 = getelementptr { i32 }* %tmp64, i32 0, i32 0
         ret i32* %tmp65
 ; CHECK-LABEL: @test41(
 ; CHECK: ret i32* %tmp1
 }
 
+define i32 addrspace(1)* @test41_addrspacecast_smaller(i32* %tmp1) {
+  %tmp64 = addrspacecast i32* %tmp1 to { i32 } addrspace(1)*
+  %tmp65 = getelementptr { i32 } addrspace(1)* %tmp64, i32 0, i32 0
+  ret i32 addrspace(1)* %tmp65
+; CHECK-LABEL: @test41_addrspacecast_smaller(
+; CHECK: addrspacecast i32* %tmp1 to i32 addrspace(1)*
+; CHECK-NEXT: ret i32 addrspace(1)*
+}
+
+define i32* @test41_addrspacecast_larger(i32 addrspace(1)* %tmp1) {
+  %tmp64 = addrspacecast i32 addrspace(1)* %tmp1 to { i32 }*
+  %tmp65 = getelementptr { i32 }* %tmp64, i32 0, i32 0
+  ret i32* %tmp65
+; CHECK-LABEL: @test41_addrspacecast_larger(
+; CHECK: addrspacecast i32 addrspace(1)* %tmp1 to i32*
+; CHECK-NEXT: ret i32*
+}
+
 define i32 @test42(i32 %X) {
         %Y = trunc i32 %X to i8         ; <i8> [#uses=1]
         %Z = zext i8 %Y to i32          ; <i32> [#uses=1]
         ret i32 %Z
 ; CHECK-LABEL: @test42(
 ; CHECK: %Z = and i32 %X, 255
 }
 

test/Transforms/InstCombine/getelementptr.ll

 define i32 addrspace(1)* @test33_array_struct_as1([10 x %struct.Key] addrspace(1)* %A) {
 ; CHECK-LABEL: @test33_array_struct_as1(
 ; CHECK: getelementptr [10 x %struct.Key] addrspace(1)* %A, i16 0, i16 1, i32 0, i32 0
   %B = bitcast [10 x %struct.Key] addrspace(1)* %A to [20 x i32] addrspace(1)*
   %C = getelementptr [20 x i32] addrspace(1)* %B, i32 0, i32 2
   ret i32 addrspace(1)* %C
 }
 
+define i32 addrspace(1)* @test33_addrspacecast(%struct.Key* %A) {
+; CHECK-LABEL: @test33_addrspacecast(
+; CHECK: %C = getelementptr %struct.Key* %A, i64 0, i32 0, i32 1
+; CHECK-NEXT: addrspacecast i32* %C to i32 addrspace(1)*
+; CHECK-NEXT: ret
+  %B = addrspacecast %struct.Key* %A to %struct.anon addrspace(1)*
+  %C = getelementptr %struct.anon addrspace(1)* %B, i32 0, i32 2
+  ret i32 addrspace(1)* %C
+}
+
   %T2 = type { i8*, i8 }
 define i8* @test34(i8* %Val, i64 %V) nounwind {
 entry:
   %A = alloca %T2, align 8
   %mrv_gep = bitcast %T2* %A to i64*
   %B = getelementptr %T2* %A, i64 0, i32 0
 
         store i64 %V, i64* %mrv_gep