SLPVectorizer: Fix assert with different sized address spaces

[ABataev]

lib/Transforms/Vectorize/SLPVectorizer.cpp

6402

I don't think this is correct. Instead, you should try to collect GEPs with the different address spaces into different nodes

[arsenm]

lib/Transforms/Vectorize/SLPVectorizer.cpp

6402

That's what I thought at first, but his isn't vectorizing the GEP itself as far as I can see. It only cares about the index operation, which could be the same size between different address spaces

[ABataev]

lib/Transforms/Vectorize/SLPVectorizer.cpp

6402

Yes, I see now. But I still think this not quite correct, just like the original code. As I understand, it tries to vectorize indices with the same base. So, when we gather GEPs, we just don't need to use GetUnderlyingObject(GEP->getPointerOperand(), *DL) as the key, just GEP->getPointerOperand(). I think, we should not dig deep into the base of the GEP.

[arsenm]

lib/Transforms/Vectorize/SLPVectorizer.cpp

6402

This isn't considering the bases at all, which makes sense based on this comment:

// Vectorize the index computations of getelementptr instructions. This
// is primarily intended to catch gather-like idioms ending at
// non-consecutive loads.

[ABataev]

lib/Transforms/Vectorize/SLPVectorizer.cpp

6402

No, it considers the bases when gathers the GEPs in bundles. These bundles are grouped using GetUnderlyingObject(GEP->getPointerOperand(), *DL) as the key, but instead it should use just GEP->getPointerOperand(). This may effect in best vectorization result, may reduce the compilation time + fix your problem with the compiler crash as all GEPs in bundles will have the same base + index. It means that the types of the SCEVs also will be the same.

[arsenm]

Remove GetUnderlyingObject, improve tests

[ABataev]

test/Transforms/SLPVectorizer/AMDGPU/address-space-ptr-sze-gep-index-assert.ll

2

Use utils/update_test_checks.py to generate correct checks for the test

[arsenm]

Use update_test_checks

[lebedev.ri]

test/Transforms/SLPVectorizer/AMDGPU/address-space-ptr-sze-gep-index-assert.ll

10–15

You need to manually cleanup the old manually-generated check-lines.

[arsenm]

Fix old checks

[ABataev]

LG

[arsenm]

r341215

lib/Transforms/Vectorize/SLPVectorizer.cpp

for (Instruction &I : *BB) {
     else if (auto *GEP = dyn_cast<GetElementPtrInst>(&I)) {
       auto Idx = GEP->idx_begin()->get();
       if (GEP->getNumIndices() > 1 || isa<Constant>(Idx))
         continue;
       if (!isValidElementType(Idx->getType()))
         continue;
       if (GEP->getType()->isVectorTy())
         continue;
-      GEPs[GetUnderlyingObject(GEP->getPointerOperand(), *DL)].push_back(GEP);
+      GEPs[GEP->getPointerOperand()].push_back(GEP);
     }
   }
 }
 
 bool SLPVectorizerPass::tryToVectorizePair(Value *A, Value *B, BoUpSLP &R) {
   if (!A || !B)
     return false;
   Value *VL[] = { A, B };
for (unsigned BI = 0, BE = Entry.second.size(); BI < BE; BI += 16) {
       for (int I = 0, E = GEPList.size(); I < E && Candidates.size() > 1; ++I) {
         auto *GEPI = cast<GetElementPtrInst>(GEPList[I]);
         if (!Candidates.count(GEPI))
           continue;
         auto *SCEVI = SE->getSCEV(GEPList[I]);
         for (int J = I + 1; J < E && Candidates.size() > 1; ++J) {
           auto *GEPJ = cast<GetElementPtrInst>(GEPList[J]);
           auto *SCEVJ = SE->getSCEV(GEPList[J]);
           if (isa<SCEVConstant>(SE->getMinusSCEV(SCEVI, SCEVJ))) {
             Candidates.remove(GEPList[I]);
             Candidates.remove(GEPList[J]);
           } else if (GEPI->idx_begin()->get() == GEPJ->idx_begin()->get()) {
             Candidates.remove(GEPList[J]);
           }
         }
       }
 

test/Transforms/SLPVectorizer/AMDGPU/address-space-ptr-sze-gep-index-assert.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -slp-vectorizer -slp-threshold=-18 < %s | FileCheck %s
+
+; Make sure there's no SCEV assert when the indexes are for different
+; sized address spaces
+
+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"
+
+define void @slp_scev_assert(i32 %idx, i64 %tmp3) #0 {
+; CHECK-LABEL: @slp_scev_assert(
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    [[TMP:%.*]] = addrspacecast i8 addrspace(5)* undef to i8*
+; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds i8, i8 addrspace(5)* undef, i32 [[IDX:%.*]]
+; CHECK-NEXT:    [[TMP4:%.*]] = getelementptr inbounds i8, i8* [[TMP]], i64 [[TMP3:%.*]]
+; CHECK-NEXT:    store i8 0, i8 addrspace(5)* [[TMP2]]
+; CHECK-NEXT:    store i8 0, i8* [[TMP4]]
+; CHECK-NEXT:    ret void
+;
+bb:
+  %tmp = addrspacecast i8 addrspace(5)* undef to i8*
+  %tmp2 = getelementptr inbounds i8, i8 addrspace(5)* undef, i32 %idx
+  %tmp4 = getelementptr inbounds i8, i8* %tmp, i64 %tmp3
+  store i8 0, i8 addrspace(5)* %tmp2
+  store i8 0, i8* %tmp4
+  ret void
+}
+
+define void @multi_as_reduction_different_sized(i32 addrspace(3)* %lds, i32 %idx0, i64 %idx1) #0 {
+; CHECK-LABEL: @multi_as_reduction_different_sized(
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    [[FLAT:%.*]] = addrspacecast i32 addrspace(3)* [[LDS:%.*]] to i32*
+; CHECK-NEXT:    [[ADD0:%.*]] = add i32 [[IDX0:%.*]], 2
+; CHECK-NEXT:    [[ADD1:%.*]] = add i64 [[IDX1:%.*]], 1
+; CHECK-NEXT:    [[LDS_1:%.*]] = getelementptr inbounds i32, i32 addrspace(3)* [[LDS]], i32 [[ADD0]]
+; CHECK-NEXT:    [[FLAT_1:%.*]] = getelementptr inbounds i32, i32* [[FLAT]], i64 [[ADD1]]
+; CHECK-NEXT:    [[LOAD_LDS_0:%.*]] = load i32, i32 addrspace(3)* [[LDS]], align 4
+; CHECK-NEXT:    [[LOAD_LDS_1:%.*]] = load i32, i32 addrspace(3)* [[LDS_1]], align 4
+; CHECK-NEXT:    [[LOAD_FLAT_0:%.*]] = load i32, i32* [[FLAT]], align 4
+; CHECK-NEXT:    [[LOAD_FLAT_1:%.*]] = load i32, i32* [[FLAT_1]], align 4
+; CHECK-NEXT:    [[SUB0:%.*]] = sub i32 [[LOAD_FLAT_0]], [[LOAD_LDS_0]]
+; CHECK-NEXT:    [[SUB1:%.*]] = sub i32 [[LOAD_FLAT_1]], [[LOAD_LDS_1]]
+; CHECK-NEXT:    store i32 [[SUB0]], i32* undef
+; CHECK-NEXT:    store i32 [[SUB1]], i32* undef
+; CHECK-NEXT:    ret void
+;
+bb:
+  %flat = addrspacecast i32 addrspace(3)* %lds to i32*
+  %add0 = add i32 %idx0, 2
+  %add1 = add i64 %idx1, 1
+
+  %lds.1 = getelementptr inbounds i32, i32 addrspace(3)* %lds, i32 %add0
+  %flat.1 = getelementptr inbounds i32, i32* %flat, i64 %add1
+
+  %load.lds.0 = load i32, i32 addrspace(3)* %lds, align 4
+  %load.lds.1 = load i32, i32 addrspace(3)* %lds.1, align 4
+
+  %load.flat.0 = load i32, i32* %flat, align 4
+  %load.flat.1 = load i32, i32* %flat.1, align 4
+
+  %sub0 = sub i32 %load.flat.0, %load.lds.0
+  %sub1 = sub i32 %load.flat.1, %load.lds.1
+
+  store i32 %sub0, i32* undef
+  store i32 %sub1, i32* undef
+  ret void
+}
+
+; This should vectorize if using GetUnderlyingObject
+define void @multi_as_reduction_same_size(i32 addrspace(1)* %global, i64 %idx0, i64 %idx1) #0 {
+; CHECK-LABEL: @multi_as_reduction_same_size(
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    [[FLAT:%.*]] = addrspacecast i32 addrspace(1)* [[GLOBAL:%.*]] to i32*
+; CHECK-NEXT:    [[ADD0:%.*]] = add i64 [[IDX0:%.*]], 2
+; CHECK-NEXT:    [[ADD1:%.*]] = add i64 [[IDX1:%.*]], 1
+; CHECK-NEXT:    [[GLOBAL_1:%.*]] = getelementptr inbounds i32, i32 addrspace(1)* [[GLOBAL]], i64 [[ADD0]]
+; CHECK-NEXT:    [[FLAT_1:%.*]] = getelementptr inbounds i32, i32* [[FLAT]], i64 [[ADD1]]
+; CHECK-NEXT:    [[LOAD_GLOBAL_0:%.*]] = load i32, i32 addrspace(1)* [[GLOBAL]], align 4
+; CHECK-NEXT:    [[LOAD_GLOBAL_1:%.*]] = load i32, i32 addrspace(1)* [[GLOBAL_1]], align 4
+; CHECK-NEXT:    [[LOAD_FLAT_0:%.*]] = load i32, i32* [[FLAT]], align 4
+; CHECK-NEXT:    [[LOAD_FLAT_1:%.*]] = load i32, i32* [[FLAT_1]], align 4
+; CHECK-NEXT:    [[SUB0:%.*]] = sub i32 [[LOAD_FLAT_0]], [[LOAD_GLOBAL_0]]
+; CHECK-NEXT:    [[SUB1:%.*]] = sub i32 [[LOAD_FLAT_1]], [[LOAD_GLOBAL_1]]
+; CHECK-NEXT:    store i32 [[SUB0]], i32* undef
+; CHECK-NEXT:    store i32 [[SUB1]], i32* undef
+; CHECK-NEXT:    ret void
+;
+bb:
+  %flat = addrspacecast i32 addrspace(1)* %global to i32*
+  %add0 = add i64 %idx0, 2
+  %add1 = add i64 %idx1, 1
+
+  %global.1 = getelementptr inbounds i32, i32 addrspace(1)* %global, i64 %add0
+  %flat.1 = getelementptr inbounds i32, i32* %flat, i64 %add1
+
+  %load.global.0 = load i32, i32 addrspace(1)* %global, align 4
+  %load.global.1 = load i32, i32 addrspace(1)* %global.1, align 4
+
+  %load.flat.0 = load i32, i32* %flat, align 4
+  %load.flat.1 = load i32, i32* %flat.1, align 4
+
+  %sub0 = sub i32 %load.flat.0, %load.global.0
+  %sub1 = sub i32 %load.flat.1, %load.global.1
+
+  store i32 %sub0, i32* undef
+  store i32 %sub1, i32* undef
+  ret void
+}
+
+; This should vectorize if using GetUnderlyingObject
+; The add is done in the same width, even though the address space size is smaller
+define void @multi_as_reduction_different_sized_noncanon(i32 addrspace(3)* %lds, i64 %idx0, i64 %idx1) #0 {
+; CHECK-LABEL: @multi_as_reduction_different_sized_noncanon(
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    [[FLAT:%.*]] = addrspacecast i32 addrspace(3)* [[LDS:%.*]] to i32*
+; CHECK-NEXT:    [[ADD0:%.*]] = add i64 [[IDX0:%.*]], 2
+; CHECK-NEXT:    [[ADD1:%.*]] = add i64 [[IDX1:%.*]], 1
+; CHECK-NEXT:    [[LDS_1:%.*]] = getelementptr inbounds i32, i32 addrspace(3)* [[LDS]], i64 [[ADD0]]
+; CHECK-NEXT:    [[FLAT_1:%.*]] = getelementptr inbounds i32, i32* [[FLAT]], i64 [[ADD1]]
+; CHECK-NEXT:    [[LOAD_LDS_0:%.*]] = load i32, i32 addrspace(3)* [[LDS]], align 4
+; CHECK-NEXT:    [[LOAD_LDS_1:%.*]] = load i32, i32 addrspace(3)* [[LDS_1]], align 4
+; CHECK-NEXT:    [[LOAD_FLAT_0:%.*]] = load i32, i32* [[FLAT]], align 4
+; CHECK-NEXT:    [[LOAD_FLAT_1:%.*]] = load i32, i32* [[FLAT_1]], align 4
+; CHECK-NEXT:    [[SUB0:%.*]] = sub i32 [[LOAD_FLAT_0]], [[LOAD_LDS_0]]
+; CHECK-NEXT:    [[SUB1:%.*]] = sub i32 [[LOAD_FLAT_1]], [[LOAD_LDS_1]]
+; CHECK-NEXT:    store i32 [[SUB0]], i32* undef
+; CHECK-NEXT:    store i32 [[SUB1]], i32* undef
+; CHECK-NEXT:    ret void
+;
+bb:
+  %flat = addrspacecast i32 addrspace(3)* %lds to i32*
+  %add0 = add i64 %idx0, 2
+  %add1 = add i64 %idx1, 1
+
+  %lds.1 = getelementptr inbounds i32, i32 addrspace(3)* %lds, i64 %add0
+  %flat.1 = getelementptr inbounds i32, i32* %flat, i64 %add1
+
+  %load.lds.0 = load i32, i32 addrspace(3)* %lds, align 4
+  %load.lds.1 = load i32, i32 addrspace(3)* %lds.1, align 4
+
+  %load.flat.0 = load i32, i32* %flat, align 4
+  %load.flat.1 = load i32, i32* %flat.1, align 4
+
+  %sub0 = sub i32 %load.flat.0, %load.lds.0
+  %sub1 = sub i32 %load.flat.1, %load.lds.1
+
+  store i32 %sub0, i32* undef
+  store i32 %sub1, i32* undef
+  ret void
+}