[VectorCombine] Insert addrspacecast when crossing address space boundaries

[frasercrmck]

Summary

We can not bitcast pointers across different address spaces. This was
previously fixed in D89577 but then in D93229 an enhancement was added
which peeks further through the ponter operand, opening up the
possibility that address-space violations could be introduced.

Instead of bailing as the previous fix did, simply insert an
addrspacecast cast instruction.

[tra]

It's unclear to me whether we need to bail in this situation as I've
done, or whether we can issue an addrspace cast, at least under certain
circumstances.

You could use TTI.isNoopAddrSpaceCast() to allow some other combinations.
E.g. if both AS can be cast to generic AS, then we may be allowed to carry on, after adding appropriate ASC to generic AS.

[lebedev.ri]

Does adding an offset of i32 %x bytes do the same thing to the pointer regardless of the address space?
Can't we just emit said missing ASC?

[frasercrmck]

In D121787#3387186, @lebedev.ri wrote:

Does adding an offset of i32 %x bytes do the same thing to the pointer regardless of the address space?
Can't we just emit said missing ASC?

That was my initial thought. I don't see why we're forced into a bitcast when we have addrspacecast. But given the original fix kept us down the bitcast lines, I assumed there must be something I'm missing.

AFAICT, address spaces have no effect on GEPs. But then again, are non-integral pointer types a concern?

There may well be a question of performance, as @tra seems to imply, and that's why we don't issue addrspacecasts. Obviously if we find ourselves in this situation there's already some kind of cast going on. But it's possible that a load from one address space is much worse than a load from the original address space. isNoopAddrSpaceCast is an easy win in any case, as @tra points out.

[lebedev.ri]

In D121787#3388573, @frasercrmck wrote:

In D121787#3387186, @lebedev.ri wrote:

Does adding an offset of i32 %x bytes do the same thing to the pointer regardless of the address space?
Can't we just emit said missing ASC?

That was my initial thought. I don't see why we're forced into a bitcast when we have addrspacecast. But given the original fix kept us down the bitcast lines, I assumed there must be something I'm missing.

AFAICT, address spaces have no effect on GEPs.

Then let's just emit said ASC :)

But then again, are non-integral pointer types a concern?

We don't do any integers<->pointers casts here, so i don't see why it would be?

There may well be a question of performance, as @tra seems to imply,

I don't follow. Performance in what sense?

and that's why we don't issue addrspacecasts. Obviously if we find ourselves in this situation there's already some kind of cast going on. But it's possible that a load from one address space is much worse than a load from the original address space. isNoopAddrSpaceCast is an easy win in any case, as @tra points out.

[frasercrmck]

use addrspace cast, move test to common file

[frasercrmck]

In D121787#3388635, @lebedev.ri wrote:

In D121787#3388573, @frasercrmck wrote:

AFAICT, address spaces have no effect on GEPs.

Then let's just emit said ASC :)

I've updated the patch to do this now. No impact on any tests other than the one I added.

We don't do any integers<->pointers casts here, so i don't see why it would be?

Yes I now see I misunderstood the semantics of non-integral pointer types. I thought they could be a different bitwidth than what they claimed to be which would have impacted GEP offsetting. But they're just about the bitwise representation.

I don't follow. Performance in what sense?

I'm not sure exactly what, just picking up on @tra's comment and previous fix (why we didn't just use addrspacecast before). But like I said, if we find ourselves with a different AS to the original, musn't we have an addrspace cast in the original IR, so we're not introducing a new one? Therefore I don't think the cost of addrspacecast is important, and we already take the cost of the scalar-vs-vector load into account.

[lebedev.ri]

LG, thanks.

[frasercrmck]

rebase

[frasercrmck]

update AMDGPU tests. @arsenm do these test changes look okay to you?

[tra]

In D121787#3388739, @frasercrmck wrote:

I don't follow. Performance in what sense?

I'm not sure exactly what, just picking up on @tra's comment and previous fix (why we didn't just use addrspacecast before). But like I said, if we find ourselves with a different AS to the original, musn't we have an addrspace cast in the original IR, so we're not introducing a new one? Therefore I don't think the cost of addrspacecast is important, and we already take the cost of the scalar-vs-vector load into account.

My comment was about relaxing the AS mismatch check and allowing different AS with no-cost ASC. With that check gone, my suggestion is now irrelevant.

[frasercrmck]

fix AMDGPU test checks, which were updated with old opt binary

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

if (!Load || !Load->isSimple() || !Load->hasOneUse() ||
       Load->getFunction()->hasFnAttribute(Attribute::SanitizeMemTag) ||
       mustSuppressSpeculation(*Load))
     return false;
 
   const DataLayout &DL = I.getModule()->getDataLayout();
   Value *SrcPtr = Load->getPointerOperand()->stripPointerCasts();
   assert(isa<PointerType>(SrcPtr->getType()) && "Expected a pointer type");
 
-  // If original AS != Load's AS, we can't bitcast the original pointer and have
-  // to use Load's operand instead. Ideally we would want to strip pointer casts
-  // without changing AS, but there's no API to do that ATM.
   unsigned AS = Load->getPointerAddressSpace();
-  if (AS != SrcPtr->getType()->getPointerAddressSpace())
-    SrcPtr = Load->getPointerOperand();
 
   // We are potentially transforming byte-sized (8-bit) memory accesses, so make
   // sure we have all of our type-based constraints in place for this target.
   Type *ScalarTy = Scalar->getType();
   uint64_t ScalarSize = ScalarTy->getPrimitiveSizeInBits();
   unsigned MinVectorSize = TTI.getMinVectorRegisterBitWidth();
   if (!ScalarSize || !MinVectorSize || MinVectorSize % ScalarSize != 0 ||
       ScalarSize % 8 != 0)
bool VectorCombine::vectorizeLoadInsert(Instruction &I) {
   // We can aggressively convert to the vector form because the backend can
   // invert this transform if it does not result in a performance win.
   if (OldCost < NewCost || !NewCost.isValid())
     return false;
 
   // It is safe and potentially profitable to load a vector directly:
   // inselt undef, load Scalar, 0 --> load VecPtr
   IRBuilder<> Builder(Load);
-  Value *CastedPtr = Builder.CreateBitCast(SrcPtr, MinVecTy->getPointerTo(AS));
+  Value *CastedPtr = Builder.CreatePointerBitCastOrAddrSpaceCast(
+      SrcPtr, MinVecTy->getPointerTo(AS));
   Value *VecLd = Builder.CreateAlignedLoad(MinVecTy, CastedPtr, Alignment);
   VecLd = Builder.CreateShuffleVector(VecLd, Mask);
 
   replaceValue(I, *VecLd);
   ++NumVecLoad;
   return true;
 }
 

llvm/test/Transforms/VectorCombine/AMDGPU/as-transition-inseltpoison.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -passes=vector-combine -S -mtriple=amdgcn-amd-amdhsa | FileCheck %s --check-prefixes=CHECK
 
 ; ModuleID = 'load-as-transition.ll'
 target datalayout = "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64-S32-A5-ni:7"
 target triple = "amdgcn-amd-amdhsa"
 
 %struct.hoge = type { float }
 
 define protected amdgpu_kernel void @load_from_other_as(<4 x float>* nocapture nonnull %resultptr) local_unnamed_addr #0 {
 ; CHECK-LABEL: @load_from_other_as(
 ; CHECK-NEXT:  bb:
 ; CHECK-NEXT:    [[A:%.*]] = alloca [[STRUCT_HOGE:%.*]], align 4, addrspace(5)
-; CHECK-NEXT:    [[B:%.*]] = addrspacecast [[STRUCT_HOGE]] addrspace(5)* [[A]] to %struct.hoge*
-; CHECK-NEXT:    [[C:%.*]] = getelementptr inbounds [[STRUCT_HOGE]], %struct.hoge* [[B]], i64 0, i32 0
-; CHECK-NEXT:    [[TMP0:%.*]] = bitcast float* [[C]] to <1 x float>*
+; CHECK-NEXT:    [[TMP0:%.*]] = addrspacecast [[STRUCT_HOGE]] addrspace(5)* [[A]] to <1 x float>*
 ; CHECK-NEXT:    [[TMP1:%.*]] = load <1 x float>, <1 x float>* [[TMP0]], align 4
 ; CHECK-NEXT:    [[E:%.*]] = shufflevector <1 x float> [[TMP1]], <1 x float> poison, <4 x i32> <i32 0, i32 undef, i32 undef, i32 undef>
 ; CHECK-NEXT:    store <4 x float> [[E]], <4 x float>* [[RESULTPTR:%.*]], align 16
 ; CHECK-NEXT:    ret void
 ;
 bb:
   %a = alloca %struct.hoge, align 4, addrspace(5)
   %b = addrspacecast %struct.hoge addrspace(5)* %a to %struct.hoge*

llvm/test/Transforms/VectorCombine/AMDGPU/as-transition.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -passes=vector-combine -S -mtriple=amdgcn-amd-amdhsa | FileCheck %s --check-prefixes=CHECK
 
 ; ModuleID = 'load-as-transition.ll'
 target datalayout = "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64-S32-A5-ni:7"
 target triple = "amdgcn-amd-amdhsa"
 
 %struct.hoge = type { float }
 
 define protected amdgpu_kernel void @load_from_other_as(<4 x float>* nocapture nonnull %resultptr) local_unnamed_addr #0 {
 ; CHECK-LABEL: @load_from_other_as(
 ; CHECK-NEXT:  bb:
 ; CHECK-NEXT:    [[A:%.*]] = alloca [[STRUCT_HOGE:%.*]], align 4, addrspace(5)
-; CHECK-NEXT:    [[B:%.*]] = addrspacecast [[STRUCT_HOGE]] addrspace(5)* [[A]] to %struct.hoge*
-; CHECK-NEXT:    [[C:%.*]] = getelementptr inbounds [[STRUCT_HOGE]], %struct.hoge* [[B]], i64 0, i32 0
-; CHECK-NEXT:    [[TMP0:%.*]] = bitcast float* [[C]] to <1 x float>*
+; CHECK-NEXT:    [[TMP0:%.*]] = addrspacecast [[STRUCT_HOGE]] addrspace(5)* [[A]] to <1 x float>*
 ; CHECK-NEXT:    [[TMP1:%.*]] = load <1 x float>, <1 x float>* [[TMP0]], align 4
 ; CHECK-NEXT:    [[E:%.*]] = shufflevector <1 x float> [[TMP1]], <1 x float> poison, <4 x i32> <i32 0, i32 undef, i32 undef, i32 undef>
 ; CHECK-NEXT:    store <4 x float> [[E]], <4 x float>* [[RESULTPTR:%.*]], align 16
 ; CHECK-NEXT:    ret void
 ;
 bb:
   %a = alloca %struct.hoge, align 4, addrspace(5)
   %b = addrspacecast %struct.hoge addrspace(5)* %a to %struct.hoge*

llvm/test/Transforms/VectorCombine/X86/load-inseltpoison.ll

 ; CHECK-NEXT:    ret <4 x float> [[R]]
 ;
   %gep = getelementptr inbounds <4 x float>, <4 x float> addrspace(44)* %p, i64 0, i64 0
   %s = load float, float addrspace(44)* %gep, align 16
   %r = insertelement <4 x float> poison, float %s, i64 0
   ret <4 x float> %r
 }
 
+; Should work with addrspace even when peeking past unsafe loads through geps
+
+define <4 x i32> @unsafe_load_i32_insert_v4i32_addrspace(i32* align 16 dereferenceable(16) %v3) {
+; CHECK-LABEL: @unsafe_load_i32_insert_v4i32_addrspace(
+; CHECK-NEXT:    [[TMP1:%.*]] = addrspacecast i32* [[V3:%.*]] to <4 x i32> addrspace(42)*
+; CHECK-NEXT:    [[TMP2:%.*]] = load <4 x i32>, <4 x i32> addrspace(42)* [[TMP1]], align 16
+; CHECK-NEXT:    [[INSELT:%.*]] = shufflevector <4 x i32> [[TMP2]], <4 x i32> poison, <4 x i32> <i32 2, i32 undef, i32 undef, i32 undef>
+; CHECK-NEXT:    ret <4 x i32> [[INSELT]]
+;
+  %t0 = getelementptr inbounds i32, i32* %v3, i32 1
+  %t1 = addrspacecast i32* %t0 to i32 addrspace(42)*
+  %t2 = getelementptr inbounds i32, i32 addrspace(42)* %t1, i64 1
+  %val = load i32, i32 addrspace(42)* %t2, align 4
+  %inselt = insertelement <4 x i32> poison, i32 %val, i32 0
+  ret <4 x i32> %inselt
+}
+
 ; If there are enough dereferenceable bytes, we can offset the vector load.
 
 define <8 x i16> @gep01_load_i16_insert_v8i16(<8 x i16>* align 16 dereferenceable(18) %p) nofree nosync {
 ; CHECK-LABEL: @gep01_load_i16_insert_v8i16(
 ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds <8 x i16>, <8 x i16>* [[P:%.*]], i64 0, i64 1
 ; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i16* [[GEP]] to <8 x i16>*
 ; CHECK-NEXT:    [[TMP2:%.*]] = load <8 x i16>, <8 x i16>* [[TMP1]], align 2
 ; CHECK-NEXT:    [[R:%.*]] = shufflevector <8 x i16> [[TMP2]], <8 x i16> poison, <8 x i32> <i32 0, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef>