[InstCombine] Canonicalize gep of addrspacecast

[nikic]

Summary

With typed pointers, InstCombine will often convert gep of addrspacecast into addrspacecast of gep, see e.g. https://llvm.godbolt.org/z/fYa3oKvdh. This isn't done consistently, because it happens as a side-effect of other folds. The most common one is trivial index descaling (where no descaling actually happens, and only the addrspacecast/gep end up being swapped). These folds are disabled for opaque pointers, so we don't canonicalize in either direction there.

This adds canonicalization of gep(ac) to ac(gep). I've limited this to the case where index types match: While this is always legal, changing index types might introduced truncs/sexts, which might be non-profitable (test diffs were a wash).

After typed pointer support is dropped, we might want to consider inverting this canonicalization and going from ac(gep) to gep(ac) instead, which seems better to me. But for now I'm matching what typed pointers were (mostly) doing.

[arsenm]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

2473

We probably shouldn't be reassociating if either address space is non-integral. I think that was one of the fixes that never happened for the existing code

2473

But yes, moving the addrspacecast up makes more sense

[nikic]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

2473

We probably shouldn't be reassociating if either address space is non-integral.

Hm, why? I don't think we really consider GEP to operate on the integer representation?

But yes, moving the addrspacecast up makes more sense

I think in that case I'll probably just ignore this discrepancy for now (and update tests with current opaque pointer behavior) and then return to this once typed pointers are removed and we don't have to worry about opposing transforms.

[arsenm]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

2473

The non-integral GEP could have some special clamping or wrapping behavior in one address space and not the other

[nikic]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

2473

The non-integral GEP could have some special clamping or wrapping behavior in one address space and not the other

Hm, is that legal? That doesn't really sound compatible with how we deal with GEP offsets in BasicAA (and stripAndAccumulateConstantOffsets for that matter). We treat offset operations in different address spaces as interchangeable, apart from trunc/sext.

[arsenm]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

2473

Key word is nonintegral address spaces, which were only half implemented

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

if (auto *PN = dyn_cast<PHINode>(PtrOp)) {
     NewGEP->insertInto(GEP.getParent(), GEP.getParent()->getFirstInsertionPt());
     return replaceOperand(GEP, 0, NewGEP);
   }
 
   if (auto *Src = dyn_cast<GEPOperator>(PtrOp))
     if (Instruction *I = visitGEPOfGEP(GEP, Src))
       return I;
 
+  // Canonicalize gep of addrspacecast to addrspacecast of gep. Don't do this
+  // if index sizes don't match, to avoid introducing sext/trunc.
+  // TODO: It might make more sense to canonicalize in the reverse direction.
+  if (auto *AC = dyn_cast<AddrSpaceCastInst>(PtrOp)) {
+    if (AC->hasOneUse() &&
+        DL.getIndexTypeSizeInBits(GEP.getType()) ==
+            DL.getIndexTypeSizeInBits(AC->getOperand(0)->getType())) {
+      Value *NewGEP = Builder.CreateGEP(GEPEltType, AC->getOperand(0), Indices,
+                                        "", GEP.isInBounds());
+      NewGEP->takeName(&GEP);
+      Value *NewAC = Builder.CreateAddrSpaceCast(NewGEP, GEP.getType());
+      NewAC->takeName(AC);
+      return replaceInstUsesWith(GEP, NewAC);
+    }
+  }
+
   // Skip if GEP source element type is scalable. The type alloc size is unknown
   // at compile-time.
   if (GEP.getNumIndices() == 1 && !IsGEPSrcEleScalable) {
     unsigned AS = GEP.getPointerAddressSpace();
     if (GEP.getOperand(1)->getType()->getScalarSizeInBits() ==
         DL.getIndexSizeInBits(AS)) {
       uint64_t TyAllocSize = DL.getTypeAllocSize(GEPEltType).getFixedSize();
 

llvm/test/Transforms/InstCombine/gep-addrspace.ll

 declare void @escape_alloca(i16*)
 
 ; check that addrspacecast is stripped when trying to mark a GEP as inbounds
 define { i8, i8 } @inbounds_after_addrspacecast() {
 ; CHECK-LABEL: @inbounds_after_addrspacecast(
 ; CHECK-NEXT:    [[T0:%.*]] = alloca i16, align 2
 ; CHECK-NEXT:    call void @escape_alloca(i16* nonnull [[T0]])
 ; CHECK-NEXT:    [[TMPCAST:%.*]] = bitcast i16* [[T0]] to [2 x i8]*
-; CHECK-NEXT:    [[T1:%.*]] = addrspacecast [2 x i8]* [[TMPCAST]] to [2 x i8] addrspace(11)*
-; CHECK-NEXT:    [[T2:%.*]] = getelementptr inbounds [2 x i8], [2 x i8] addrspace(11)* [[T1]], i64 0, i64 1
-; CHECK-NEXT:    [[T3:%.*]] = load i8, i8 addrspace(11)* [[T2]], align 1
+; CHECK-NEXT:    [[T2:%.*]] = getelementptr inbounds [2 x i8], [2 x i8]* [[TMPCAST]], i64 0, i64 1
+; CHECK-NEXT:    [[T1:%.*]] = addrspacecast i8* [[T2]] to i8 addrspace(11)*
+; CHECK-NEXT:    [[T3:%.*]] = load i8, i8 addrspace(11)* [[T1]], align 1
 ; CHECK-NEXT:    [[INSERT:%.*]] = insertvalue { i8, i8 } zeroinitializer, i8 [[T3]], 1
 ; CHECK-NEXT:    ret { i8, i8 } [[INSERT]]
 ;
   %t0 = alloca i16, align 2
   call void @escape_alloca(i16* %t0)
   %tmpcast = bitcast i16* %t0 to [2 x i8]*
   %t1 = addrspacecast [2 x i8]* %tmpcast to [2 x i8] addrspace(11)*
   %t2 = getelementptr [2 x i8], [2 x i8] addrspace(11)* %t1, i64 0, i64 1

llvm/test/Transforms/InstCombine/gep-vector.ll

;
   ret i32* %gep
 }
 
 ; Negative test - datalayout's alloc size for the 2 types must match.
 
 define i32 addrspace(3)* @bitcast_vec_to_array_addrspace(<7 x i32>* %x, i64 %y, i64 %z) {
 ; CHECK-LABEL: @bitcast_vec_to_array_addrspace(
 ; CHECK-NEXT:    [[ARR_PTR:%.*]] = bitcast <7 x i32>* [[X:%.*]] to [7 x i32]*
-; CHECK-NEXT:    [[ASC:%.*]] = addrspacecast [7 x i32]* [[ARR_PTR]] to [7 x i32] addrspace(3)*
-; CHECK-NEXT:    [[GEP:%.*]] = getelementptr [7 x i32], [7 x i32] addrspace(3)* [[ASC]], i64 [[Y:%.*]], i64 [[Z:%.*]]
-; CHECK-NEXT:    ret i32 addrspace(3)* [[GEP]]
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr [7 x i32], [7 x i32]* [[ARR_PTR]], i64 [[Y:%.*]], i64 [[Z:%.*]]
+; CHECK-NEXT:    [[ASC:%.*]] = addrspacecast i32* [[GEP]] to i32 addrspace(3)*
+; CHECK-NEXT:    ret i32 addrspace(3)* [[ASC]]
 ;
   %arr_ptr = bitcast <7 x i32>* %x to [7 x i32]*
   %asc = addrspacecast [7 x i32]* %arr_ptr to [7 x i32] addrspace(3)*
   %gep = getelementptr [7 x i32], [7 x i32] addrspace(3)* %asc, i64 %y, i64 %z
   ret i32 addrspace(3)* %gep
 }
 
 ; Negative test - datalayout's alloc size for the 2 types must match.
 
 define i32 addrspace(3)* @inbounds_bitcast_vec_to_array_addrspace(<7 x i32>* %x, i64 %y, i64 %z) {
 ; CHECK-LABEL: @inbounds_bitcast_vec_to_array_addrspace(
 ; CHECK-NEXT:    [[ARR_PTR:%.*]] = bitcast <7 x i32>* [[X:%.*]] to [7 x i32]*
-; CHECK-NEXT:    [[ASC:%.*]] = addrspacecast [7 x i32]* [[ARR_PTR]] to [7 x i32] addrspace(3)*
-; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds [7 x i32], [7 x i32] addrspace(3)* [[ASC]], i64 [[Y:%.*]], i64 [[Z:%.*]]
-; CHECK-NEXT:    ret i32 addrspace(3)* [[GEP]]
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds [7 x i32], [7 x i32]* [[ARR_PTR]], i64 [[Y:%.*]], i64 [[Z:%.*]]
+; CHECK-NEXT:    [[ASC:%.*]] = addrspacecast i32* [[GEP]] to i32 addrspace(3)*
+; CHECK-NEXT:    ret i32 addrspace(3)* [[ASC]]
 ;
   %arr_ptr = bitcast <7 x i32>* %x to [7 x i32]*
   %asc = addrspacecast [7 x i32]* %arr_ptr to [7 x i32] addrspace(3)*
   %gep = getelementptr inbounds [7 x i32], [7 x i32] addrspace(3)* %asc, i64 %y, i64 %z
   ret i32 addrspace(3)* %gep
 }
 
 ; Sizes and types match - safe to remove bitcast.
 
 define i32 addrspace(3)* @bitcast_vec_to_array_addrspace_matching_alloc_size(<4 x i32>* %x, i64 %y, i64 %z) {
 ; CHECK-LABEL: @bitcast_vec_to_array_addrspace_matching_alloc_size(
 ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr <4 x i32>, <4 x i32>* [[X:%.*]], i64 [[Y:%.*]], i64 [[Z:%.*]]
-; CHECK-NEXT:    [[TMP1:%.*]] = addrspacecast i32* [[GEP]] to i32 addrspace(3)*
-; CHECK-NEXT:    ret i32 addrspace(3)* [[TMP1]]
+; CHECK-NEXT:    [[ASC:%.*]] = addrspacecast i32* [[GEP]] to i32 addrspace(3)*
+; CHECK-NEXT:    ret i32 addrspace(3)* [[ASC]]
 ;
   %arr_ptr = bitcast <4 x i32>* %x to [4 x i32]*
   %asc = addrspacecast [4 x i32]* %arr_ptr to [4 x i32] addrspace(3)*
   %gep = getelementptr [4 x i32], [4 x i32] addrspace(3)* %asc, i64 %y, i64 %z
   ret i32 addrspace(3)* %gep
 }
 
 ; Sizes and types match - safe to remove bitcast.
 
 define i32 addrspace(3)* @inbounds_bitcast_vec_to_array_addrspace_matching_alloc_size(<4 x i32>* %x, i64 %y, i64 %z) {
 ; CHECK-LABEL: @inbounds_bitcast_vec_to_array_addrspace_matching_alloc_size(
 ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds <4 x i32>, <4 x i32>* [[X:%.*]], i64 [[Y:%.*]], i64 [[Z:%.*]]
-; CHECK-NEXT:    [[TMP1:%.*]] = addrspacecast i32* [[GEP]] to i32 addrspace(3)*
-; CHECK-NEXT:    ret i32 addrspace(3)* [[TMP1]]
+; CHECK-NEXT:    [[ASC:%.*]] = addrspacecast i32* [[GEP]] to i32 addrspace(3)*
+; CHECK-NEXT:    ret i32 addrspace(3)* [[ASC]]
 ;
   %arr_ptr = bitcast <4 x i32>* %x to [4 x i32]*
   %asc = addrspacecast [4 x i32]* %arr_ptr to [4 x i32] addrspace(3)*
   %gep = getelementptr inbounds [4 x i32], [4 x i32] addrspace(3)* %asc, i64 %y, i64 %z
   ret i32 addrspace(3)* %gep
 }
 
 ; Negative test - avoid doing bitcast on i8*, because '16' should be scaled by 'vscale'.
 
 define i8* @test_accumulate_constant_offset_vscale_nonzero(<vscale x 16 x i1> %pg, i8* %base) {
-; CHECK-LABEL: @test_accumulate_constant_offset_vscale_nonzero
-; CHECK-NEXT:   %bc = bitcast i8* %base to <vscale x 16 x i8>*
-; CHECK-NEXT:   %gep = getelementptr <vscale x 16 x i8>, <vscale x 16 x i8>* %bc, i64 1, i64 4
-; CHECK-NEXT:   ret i8* %gep
+; CHECK-LABEL: @test_accumulate_constant_offset_vscale_nonzero(
+; CHECK-NEXT:    [[BC:%.*]] = bitcast i8* [[BASE:%.*]] to <vscale x 16 x i8>*
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr <vscale x 16 x i8>, <vscale x 16 x i8>* [[BC]], i64 1, i64 4
+; CHECK-NEXT:    ret i8* [[GEP]]
+;
   %bc = bitcast i8* %base to <vscale x 16 x i8>*
   %gep = getelementptr <vscale x 16 x i8>, <vscale x 16 x i8>* %bc, i64 1, i64 4
   ret i8* %gep
 }
 
 define i8* @test_accumulate_constant_offset_vscale_zero(<vscale x 16 x i1> %pg, i8* %base) {
-; CHECK-LABEL: @test_accumulate_constant_offset_vscale_zero
-; CHECK-NEXT:   %[[RES:.*]] = getelementptr i8, i8* %base, i64 4
-; CHECK-NEXT:   ret i8* %[[RES]]
+; CHECK-LABEL: @test_accumulate_constant_offset_vscale_zero(
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i8, i8* [[BASE:%.*]], i64 4
+; CHECK-NEXT:    ret i8* [[GEP]]
+;
   %bc = bitcast i8* %base to <vscale x 16 x i8>*
   %gep = getelementptr <vscale x 16 x i8>, <vscale x 16 x i8>* %bc, i64 0, i64 4
   ret i8* %gep
 }

llvm/test/Transforms/InstCombine/non-integral-pointers.ll

entry:
   store ptr addrspace(3) %val, ptr %ptr1
   ret void
 }
 
 define i64 @g(ptr %gp) {
 ; CHECK-LABEL: @g(
 ; CHECK-NEXT:    [[DOTPRE:%.*]] = load ptr addrspace(4), ptr [[GP:%.*]], align 8
 ; CHECK-NEXT:    [[V74:%.*]] = call ptr addrspace(4) @alloc()
-; CHECK-NEXT:    [[V75:%.*]] = addrspacecast ptr addrspace(4) [[V74]] to ptr
-; CHECK-NEXT:    [[V77:%.*]] = getelementptr ptr addrspace(4), ptr [[V75]], i64 -1
-; CHECK-NEXT:    store ptr addrspace(4) [[DOTPRE]], ptr [[V77]], align 8
-; CHECK-NEXT:    [[V81:%.*]] = load i64, ptr [[V77]], align 8
+; CHECK-NEXT:    [[V77:%.*]] = getelementptr ptr addrspace(4), ptr addrspace(4) [[V74]], i64 -1
+; CHECK-NEXT:    [[V75:%.*]] = addrspacecast ptr addrspace(4) [[V77]] to ptr
+; CHECK-NEXT:    store ptr addrspace(4) [[DOTPRE]], ptr [[V75]], align 8
+; CHECK-NEXT:    [[V81:%.*]] = load i64, ptr [[V75]], align 8
 ; CHECK-NEXT:    ret i64 [[V81]]
 ;
   %.pre = load ptr addrspace(4), ptr %gp, align 8
   %v74 = call ptr addrspace(4) @alloc()
   %v75 = addrspacecast ptr addrspace(4) %v74 to ptr
   %v77 = getelementptr ptr addrspace(4), ptr %v75, i64 -1
   store ptr addrspace(4) %.pre, ptr %v77, align 8
   %v81 = load i64, ptr %v77, align 8

llvm/test/Transforms/InstCombine/opaque-ptr.ll

bb:
   br i1 %cond, label %bb2, label %bb
 
 bb2:
   ret ptr %RHS
 }
 
 define ptr addrspace(1) @gep_of_addrspace_cast(ptr %ptr) {
 ; CHECK-LABEL: @gep_of_addrspace_cast(
-; CHECK-NEXT:    [[CAST1:%.*]] = addrspacecast ptr [[PTR:%.*]] to ptr addrspace(1)
-; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds i32, ptr addrspace(1) [[CAST1]], i64 1
-; CHECK-NEXT:    ret ptr addrspace(1) [[GEP]]
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds i32, ptr [[PTR:%.*]], i64 1
+; CHECK-NEXT:    [[CAST1:%.*]] = addrspacecast ptr [[GEP]] to ptr addrspace(1)
+; CHECK-NEXT:    ret ptr addrspace(1) [[CAST1]]
 ;
   %cast1 = addrspacecast ptr %ptr to ptr addrspace(1)
   %gep = getelementptr inbounds i32, ptr addrspace(1) %cast1, i64 1
   ret ptr addrspace(1) %gep
 }
 
 define i1 @cmp_gep_same_base_same_type(ptr %ptr, i64 %idx1, i64 %idx2) {
 ; CHECK-LABEL: @cmp_gep_same_base_same_type(