[VPlan] Handle invariant GEPs in isUniformAfterVectorization.

[fhahn]

Summary

This fixes a crash caused by legal treating a scalable GEP as invariant,
but isUniformAfterVectorization does not handle GEPs.

Fixes https://github.com/llvm/llvm-project/issues/60831.

[fhahn]

Generate check lines for test.

[reames]

I think there's a deeper logic issue here. I'm pretty sure this is a bug I introduced, so, sorry!

LoopAccessInfo::isUniform uses SCEV's notion of loop invariant. This is both things which are currently loop invariant, but also things whose computation is loop invariant (e.g. uniform).

As you noticed, VPReplicateRecipe::execute uses only whether the recipe was defined. This is essentially a proxy for whether an instructions *placement* was loop invariant. This explicitly *does not* reason about computation. (I think... this part of things I'm less confident in.)

I think the basic issue here is that a Recipe can correspond to the isUniform Value without itself being uniform-per-part or UniformAfterVectorization. So, I think your change is patching over a bug, not fixing one.

[kartcq]

This patch does not solve the issue completely.
But I don't have expertise enough to comment on what needs to be fixed.

A simplified test case which still breaks the Loop vectorizer even after applying this candidate patch is as follows

test_loop2.ll

target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
target triple = "aarch64-unknown-linux-gnu"

define dso_local fastcc void @_test_loop2(i64 %init, i64 %n, i8* %store_dest) unnamed_addr #0 {
    br label %1.
  critedge5:                                       ; preds = %1
    ret void
  1:                                               ; preds = %0, %1
    %2 = phi i64 [ %7, %1 ], [ %init, %0 ]
    %3 = sub nsw i64 %n, %2
    %4 = trunc i64 %n to i8
    %5 = add i64 %2, %3
    %6 = getelementptr i8, i8* %store_dest, i64 %5
    store i8 %4, i8* %6, align 1
    %7 = add nsw i64 %2, 1
    %8 = icmp sle i64 %7, 1000
    br i1 %8, label %1, label %.critedge5
}

attributes #0 = { "target-features"="+crc,+crypto,+fp-armv8,+neon,+sve,+v8.1a,+v8.2a,+v8.3a,+v8.4a,+v8a" }

Issue an be reproduced by command

opt -passes=loop-vectorize test_loop2.ll

[fhahn]

In D144434#4141906, @reames wrote:

I think there's a deeper logic issue here. I'm pretty sure this is a bug I introduced, so, sorry!

LoopAccessInfo::isUniform uses SCEV's notion of loop invariant. This is both things which are currently loop invariant, but also things whose computation is loop invariant (e.g. uniform).

As you noticed, VPReplicateRecipe::execute uses only whether the recipe was defined. This is essentially a proxy for whether an instructions *placement* was loop invariant. This explicitly *does not* reason about computation. (I think... this part of things I'm less confident in.)

At the moment, isUniformAfterVectorization reasons about more than just placement by using information about uniform operands, but relies on uniforms being represented as uniform VPReplicateRecipes. I think the underlying issue is that we don't classify the GEP as uniform correctly. I put up D150991 to use Legal::isUniform to bring the reasoning during classification in line with isUniformMemOp which is used in other places.

In D144434#4149935, @kartcq wrote:

This patch does not solve the issue completely.
But I don't have expertise enough to comment on what needs to be fixed.

A simplified test case which still breaks the Loop vectorizer even after applying this candidate patch is as follows

Thanks, I added the new test in 07e5f57df4bf

[Ayal]

This looks good to me. @reames, @kartcq - ok with you too?

llvm/lib/Transforms/Vectorize/VPlan.h

2708

Propagating uniformity forwards on-demand from operands is fine, for GEPs or other non-phi recipes, where not cached as a result of SCEV/Divergence Analysis.

Should this function still be "inline"?

The "AfterVectorization" suffix should be dropped if favor of "onlyFirstLaneUsed()" above (independent of this patch).

[reames]

I'm fine with this.

[fhahn]

llvm/lib/Transforms/Vectorize/VPlan.h

2708

Will move into VPlan.cpp now that this has grown and rename accordingly.

llvm/lib/Transforms/Vectorize/VPlan.h

inline bool isUniformAfterVectorization(VPValue *VPV) {
   // A value defined outside the vector region must be uniform after
   // vectorization inside a vector region.
   if (VPV->isDefinedOutsideVectorRegions())
     return true;
   VPRecipeBase *Def = VPV->getDefiningRecipe();
   assert(Def && "Must have definition for value defined inside vector region");
   if (auto Rep = dyn_cast<VPReplicateRecipe>(Def))
     return Rep->isUniform();
+  if (auto *GEP = dyn_cast<VPWidenGEPRecipe>(Def))
+    return all_of(GEP->operands(), isUniformAfterVectorization);
   return false;
 }
 } // end namespace vputils
 
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_VECTORIZE_VPLAN_H

llvm/test/Transforms/LoopVectorize/AArch64/sve-inv-store.ll

 ; CHECK-NEXT:    br label [[FOR_BODY14:%.*]]
 ; CHECK:       for.body14:
 ; CHECK-NEXT:    [[INDVARS_IV:%.*]] = phi i64 [ [[BC_RESUME_VAL]], [[SCALAR_PH]] ], [ [[INDVARS_IV_NEXT:%.*]], [[FOR_BODY14]] ]
 ; CHECK-NEXT:    [[ARRAYIDX:%.*]] = getelementptr inbounds i16, ptr [[SRC]], i64 [[INDVARS_IV]]
 ; CHECK-NEXT:    [[LD:%.*]] = load i16, ptr [[ARRAYIDX]], align 2
 ; CHECK-NEXT:    store i16 [[LD]], ptr [[DST]], align 2
 ; CHECK-NEXT:    [[INDVARS_IV_NEXT]] = add nuw nsw i64 [[INDVARS_IV]], 1
 ; CHECK-NEXT:    [[EXITCOND_NOT:%.*]] = icmp eq i64 [[INDVARS_IV_NEXT]], [[N]]
-; CHECK-NEXT:    br i1 [[EXITCOND_NOT]], label [[FOR_INC24]], label [[FOR_BODY14]], !llvm.loop [[LOOP3:![0-9]+]]
+; CHECK-NEXT:    br i1 [[EXITCOND_NOT]], label [[FOR_INC24]], label [[FOR_BODY14]], !llvm.loop [[LOOP4:![0-9]+]]
 ; CHECK:       for.inc24:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   br label %for.body14
 
 for.body14:                                       ; preds = %for.body14, %entry
   %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body14 ]
for.inc:                                          ; preds = %for.body, %if.then
   %inc = add nuw nsw i64 %i.09, 1
   %exitcond.not = icmp eq i64 %inc, %N
   br i1 %exitcond.not, label %for.end, label %for.body, !llvm.loop !0
 
 for.end:                                          ; preds = %for.inc, %entry
   ret void
 }
 
+; Test case for PR60831.
+define void @test_invar_gep(i8* %store_dest) #0 {
+; CHECK-LABEL: @test_invar_gep(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = call i64 @llvm.vscale.i64()
+; CHECK-NEXT:    [[TMP1:%.*]] = mul i64 [[TMP0]], 4
+; CHECK-NEXT:    [[MIN_ITERS_CHECK:%.*]] = icmp ult i64 100, [[TMP1]]
+; CHECK-NEXT:    br i1 [[MIN_ITERS_CHECK]], label [[SCALAR_PH:%.*]], label [[VECTOR_PH:%.*]]
+; CHECK:       vector.ph:
+; CHECK-NEXT:    [[TMP2:%.*]] = call i64 @llvm.vscale.i64()
+; CHECK-NEXT:    [[TMP3:%.*]] = mul i64 [[TMP2]], 4
+; CHECK-NEXT:    [[N_MOD_VF:%.*]] = urem i64 100, [[TMP3]]
+; CHECK-NEXT:    [[N_VEC:%.*]] = sub i64 100, [[N_MOD_VF]]
+; CHECK-NEXT:    br label [[VECTOR_BODY:%.*]]
+; CHECK:       vector.body:
+; CHECK-NEXT:    [[INDEX:%.*]] = phi i64 [ 0, [[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT:    [[TMP4:%.*]] = call <vscale x 4 x i64> @llvm.experimental.stepvector.nxv4i64()
+; CHECK-NEXT:    [[DOTSPLATINSERT:%.*]] = insertelement <vscale x 4 x i64> poison, i64 [[INDEX]], i64 0
+; CHECK-NEXT:    [[DOTSPLAT:%.*]] = shufflevector <vscale x 4 x i64> [[DOTSPLATINSERT]], <vscale x 4 x i64> poison, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT:    [[TMP5:%.*]] = add <vscale x 4 x i64> zeroinitializer, [[TMP4]]
+; CHECK-NEXT:    [[TMP6:%.*]] = mul <vscale x 4 x i64> [[TMP5]], shufflevector (<vscale x 4 x i64> insertelement (<vscale x 4 x i64> poison, i64 1, i64 0), <vscale x 4 x i64> poison, <vscale x 4 x i32> zeroinitializer)
+; CHECK-NEXT:    [[TMP7:%.*]] = add <vscale x 4 x i64> [[DOTSPLAT]], [[TMP6]]
+; CHECK-NEXT:    [[TMP8:%.*]] = add i64 [[INDEX]], 0
+; CHECK-NEXT:    [[TMP9:%.*]] = add i64 [[INDEX]], 1
+; CHECK-NEXT:    [[TMP10:%.*]] = add i64 [[INDEX]], 2
+; CHECK-NEXT:    [[TMP11:%.*]] = add i64 [[INDEX]], 3
+; CHECK-NEXT:    [[TMP12:%.*]] = getelementptr i8, ptr [[STORE_DEST:%.*]], i64 0
+; CHECK-NEXT:    [[DOTSPLATINSERT1:%.*]] = insertelement <vscale x 4 x ptr> poison, ptr [[TMP12]], i64 0
+; CHECK-NEXT:    [[DOTSPLAT2:%.*]] = shufflevector <vscale x 4 x ptr> [[DOTSPLATINSERT1]], <vscale x 4 x ptr> poison, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT:    [[TMP13:%.*]] = call i32 @llvm.vscale.i32()
+; CHECK-NEXT:    [[TMP14:%.*]] = mul i32 [[TMP13]], 4
+; CHECK-NEXT:    [[TMP15:%.*]] = sub i32 [[TMP14]], 1
+; CHECK-NEXT:    [[TMP16:%.*]] = extractelement <vscale x 4 x i64> [[TMP7]], i32 [[TMP15]]
+; CHECK-NEXT:    [[TMP17:%.*]] = extractelement <vscale x 4 x ptr> [[DOTSPLAT2]], i32 0
+; CHECK-NEXT:    store i64 [[TMP16]], ptr [[TMP17]], align 1
+; CHECK-NEXT:    [[TMP18:%.*]] = call i64 @llvm.vscale.i64()
+; CHECK-NEXT:    [[TMP19:%.*]] = mul i64 [[TMP18]], 4
+; CHECK-NEXT:    [[INDEX_NEXT]] = add nuw i64 [[INDEX]], [[TMP19]]
+; CHECK-NEXT:    [[TMP20:%.*]] = icmp eq i64 [[INDEX_NEXT]], [[N_VEC]]
+; CHECK-NEXT:    br i1 [[TMP20]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP7:![0-9]+]]
+; CHECK:       middle.block:
+; CHECK-NEXT:    [[CMP_N:%.*]] = icmp eq i64 100, [[N_VEC]]
+; CHECK-NEXT:    br i1 [[CMP_N]], label [[EXIT:%.*]], label [[SCALAR_PH]]
+; CHECK:       scalar.ph:
+; CHECK-NEXT:    [[BC_RESUME_VAL:%.*]] = phi i64 [ [[N_VEC]], [[MIDDLE_BLOCK]] ], [ 0, [[ENTRY:%.*]] ]
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[IV:%.*]] = phi i64 [ [[BC_RESUME_VAL]], [[SCALAR_PH]] ], [ [[IV_NEXT:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[GEP_INVAR:%.*]] = getelementptr i8, ptr [[STORE_DEST]], i64 0
+; CHECK-NEXT:    store i64 [[IV]], ptr [[GEP_INVAR]], align 1
+; CHECK-NEXT:    [[IV_NEXT]] = add nsw i64 [[IV]], 1
+; CHECK-NEXT:    [[EC:%.*]] = icmp eq i64 [[IV_NEXT]], 100
+; CHECK-NEXT:    br i1 [[EC]], label [[EXIT]], label [[LOOP]], !llvm.loop [[LOOP8:![0-9]+]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %loop
+
+loop:
+  %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
+  %gep.invar = getelementptr i8, i8* %store_dest, i64 0
+  store i64 %iv, i8* %gep.invar, align 1
+  %iv.next = add nsw i64 %iv, 1
+  %ec = icmp eq i64 %iv.next, 100
+  br i1 %ec, label %exit, label %loop, !llvm.loop !0
+
+exit:
+  ret void
+}
+
 attributes #0 = { "target-features"="+neon,+sve" vscale_range(1, 16) }
 
 !0 = distinct !{!0, !1, !2, !3, !4, !5}
 !1 = !{!"llvm.loop.mustprogress"}
 !2 = !{!"llvm.loop.vectorize.width", i32 4}
 !3 = !{!"llvm.loop.vectorize.scalable.enable", i1 true}
 !4 = !{!"llvm.loop.vectorize.enable", i1 true}
 !5 = !{!"llvm.loop.interleave.count", i32 1}