[LoopVectorize] When tail-folding, don't always predicate uniform loads

[david-arm]

Summary

In VPRecipeBuilder::handleReplication if we believe the instruction
is predicated we then proceed to create new VP region blocks even
when the load is uniform and only predicated due to tail-folding.

I have updated isPredicatedInst to avoid treating a uniform load as
predicated when tail-folding, which means we can do a single scalar
load and a vector splat of the value.

Tests added here:

Transforms/LoopVectorize/AArch64/tail-fold-uniform-memops.ll

[fhahn]

Does the test need to be aarch64 specific or could it be target independent?

[david-arm]

In D112552#3088838, @fhahn wrote:

Does the test need to be aarch64 specific or could it be target independent?

Hi @fhahn, I think that's a really good idea. The main problem I had when trying this is that tail-folding requires masked load/store support, so when you make this target-independent we end up scalarising the loads/stores and never expose the path I changed in this patch. If there is a way to force masked loads to be legal then I could try that?

[sdesmalen]

In D112552#3090066, @david-arm wrote:

In D112552#3088838, @fhahn wrote:

Does the test need to be aarch64 specific or could it be target independent?

Hi @fhahn, I think that's a really good idea. The main problem I had when trying this is that tail-folding requires masked load/store support, so when you make this target-independent we end up scalarising the loads/stores and never expose the path I changed in this patch. If there is a way to force masked loads to be legal then I could try that?

From what I can see, there doesn't seem to be a way to force masked loads to be legal, and personally I don't think it is a big-enough deal to create a new option for it.
Maybe you can add a comment to the test describing that the test isn't really target-specific (other than the need for TTI.isLegalMaskedLoad to return true) ?

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9073

I see you've followed the message above, but I'm confused by the terminology here. From the test it does the opposite of scalarising (e.g. load scalar value + splat), which I personally associate more with 'widening'. So perhaps "Replicating uniform load in predicated loop:" would be more appropriate? (assuming that replicating can be done in multiple ways, not necessarily by scalarizing)

llvm/test/Transforms/LoopVectorize/AArch64/tail-fold-uniform-memops.ll

48

The output of this test is unchanged by your code. Is it worth adding this test as a NFC change ("adding tests for uniform memops") and then rebasing your patch? That way it's clear what your patch has changed from the test.

[david-arm]

• Changed debug wording.
• Rebased off NFC patch.

[sdesmalen]

This looks good to me. We know the load will be executed to populate at least one of the lanes so the scalar uniform load can be performed unconditionally within the vector body.

llvm/test/Transforms/LoopVectorize/AArch64/tail-fold-uniform-memops.ll

10

nit: I'd suggest adding this test in D11261 as well, so that you can see the difference in this patch.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9070

The problem with the workaround is that IsPredicated is now inaccurate for the recipe I think.

It looks like the cost-model already accurately estimates the cost as not requiring predication. So perhaps it would be better to not claim the instruction requires predication in the first place in LoopVectorizationCostModel::isPredicated?

[david-arm]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9070

Hi @fhahn, are you suggesting that IsPredicated is being set incorrectly above? If I understand correctly, I believe you're suggesting the answer would be to add more logic to isPredicatedInst? i.e. something like:

bool isPredicatedInst(Instruction *I, bool IsUniform) {
  if (IsUniform && isa<LoadInst>(I) &&
    !Legal->blockNeedsPredication(I->getParent()))
    return false;

  if (!blockNeedsPredication(I->getParent()))
  ...
}

That feels a bit messy. Would it be better to simply adjust the IsPredicated boolean value instead in this function?

[david-arm]

• Made sure we don't set IsPredicated in the VPReplicateRecipe.

[david-arm]

Hi @fhahn, I've fixed the immediate problem with IsPredicated and now it is set correctly. If you still think it's better to change the isPredicatedInst interface I'm happy to give it a try?

[sdesmalen]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9068

nit: unrelated whitespace change?

9070

Hi @david-arm, I think that could just be:

// Returns true if \p I is an instruction that will be predicated either
// through scalar predication or masked load/store or masked gather/scatter.
// Superset of instructions that return true for isScalarWithPredication.
bool isPredicatedInst(Instruction *I, bool IsUniform = false) {
  if (!blockNeedsPredication(I->getParent()))
    return false;
  // Uniform loads don't require predication.
  if (isa<LoadInst>(I) && IsUniform)
    return false;
  // Loads and stores that need some form of masked operation are predicated
  // instructions.
  if (isa<LoadInst>(I) || isa<StoreInst>(I))
    return Legal->isMaskRequired(I);
  return isScalarWithPredication(I);
}

Which seems like a sensible change to me and isn't really that messy?

Looking at the uses of that function, should the then-block of if (isPredicatedInst(I)) in getMemInstScalarizationCost have an assert that the the value is not uniform?

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9070

That feels a bit messy. Would it be better to simply adjust the IsPredicated boolean value instead in this function?

Why do you think this would be messy? If we treat those loads as unpredicted during codegen, claiming they will be predicated earlier seems inconsistent and may lead to surprises down the road for cases they are out-of-sync.

[david-arm]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9070

It's because of the confusion between CM.blockNeedsPredication and Legal->blockNeedsPredication. Vectorised blocks created when tail-folding are something artificial we have introduced during vectorisation and we know that there is always at least one active lane. I can fix this as you suggest *if* I distinguish between the two, i.e.

bool isPredicatedInst(Instruction *I, bool IsUniform) {
  if (IsUniform && isa<LoadInst>(I) &&
    !Legal->blockNeedsPredication(I->getParent()))
    return false;

  if (!blockNeedsPredication(I->getParent()))
  ...
}

It just looks very confusing that's all due to having to call two different forms of blockNeedsPredication. I will go with the code above unless you can think of a cleaner solution?

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9070

Sounds good, especially with the renaming to Legal's blockNeedsPredication.

[david-arm]

• Changed the patch to add the uniform load check in isPredicatedInst instead.

[david-arm]

gentle ping

[fhahn]

LGTM, thanks! Please make sure to update the commit message; the current one is very much out-of-date.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

1539

I think it would be good to add the explanation from your earlier comment here to help the reader.

[sdesmalen]

I agree with @fhahn that an extra comment would be useful.

[david-arm]

• Added comment.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

public:
   /// If a non-zero VF has been calculated, we check if I will be scalarized
   /// predication for that VF.
   bool isScalarWithPredication(Instruction *I) const;
 
   // Returns true if \p I is an instruction that will be predicated either
   // through scalar predication or masked load/store or masked gather/scatter.
   // Superset of instructions that return true for isScalarWithPredication.
   bool isPredicatedInst(Instruction *I) {
     if (!blockNeedsPredication(I->getParent()))
       return false;
     // Loads and stores that need some form of masked operation are predicated
     // instructions.
     if (isa<LoadInst>(I) || isa<StoreInst>(I))
       return Legal->isMaskRequired(I);
     return isScalarWithPredication(I);
   }
 
VPBasicBlock *VPRecipeBuilder::handleReplication(
   }
 
   // Finalize the recipe for Instr, first if it is not predicated.
   if (!IsPredicated) {
     LLVM_DEBUG(dbgs() << "LV: Scalarizing:" << *I << "\n");
     VPBB->appendRecipe(Recipe);
     return VPBB;
   }
+
+  if (IsUniform && isa<LoadInst>(I) &&
+      !Legal->blockNeedsPredication(I->getParent()) && Range.Start.isVector()) {
+    assert(CM.foldTailByMasking() &&
+           "Uniform load is unexpectedly marked as predicated!");
+    LLVM_DEBUG(dbgs() << "LV: Replicating uniform load in predicated loop:"
+                      << *I << "\n");
+    VPBB->appendRecipe(Recipe);
+    return VPBB;
+  }
+
   LLVM_DEBUG(dbgs() << "LV: Scalarizing and predicating:" << *I << "\n");
   assert(VPBB->getSuccessors().empty() &&
          "VPBB has successors when handling predicated replication.");
   // Record predicated instructions for above packing optimizations.
   VPBlockBase *Region = createReplicateRegion(I, Recipe, Plan);
   VPBlockUtils::insertBlockAfter(Region, VPBB);
   auto *RegSucc = new VPBasicBlock();
   VPBlockUtils::insertBlockAfter(RegSucc, Region);

llvm/test/Transforms/LoopVectorize/AArch64/tail-fold-uniform-memops.ll

 ; RUN: opt -loop-vectorize -scalable-vectorization=off -force-vector-width=4 -prefer-predicate-over-epilogue=predicate-dont-vectorize -S < %s 2>&1 | FileCheck %s
 
 ; NOTE: These tests aren't really target-specific, but it's convenient to target AArch64
 ; so that TTI.isLegalMaskedLoad can return true.
 
 target triple = "aarch64-linux-gnu"
 
+; The original loop had an unconditional uniform load. Let's make sure
+; we don't artificially create new predicated blocks for the load.
+define void @uniform_load(i32* noalias %dst, i32* noalias readonly %src, i64 %n) #0 {
+; CHECK-LABEL: @uniform_load(
+; CHECK:       vector.body:
+; CHECK-NEXT:    [[IDX:%.*]] = phi i64 [ 0, %vector.ph ], [ [[IDX_NEXT:%.*]], %vector.body ]
+; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <4 x i64> poison, i64 [[IDX]], i32 0
+; CHECK-NEXT:    [[TMP2:%.*]] = shufflevector <4 x i64> [[TMP1]], <4 x i64> poison, <4 x i32> zeroinitializer
+; CHECK-NEXT:    [[INDUCTION:%.*]] = add <4 x i64> [[TMP2]], <i64 0, i64 1, i64 2, i64 3>
+; CHECK-NEXT:    [[TMP3:%.*]] = add i64 [[IDX]], 0
+; CHECK-NEXT:    [[LOOP_PRED:%.*]] = icmp ule <4 x i64> [[INDUCTION]]
+; CHECK-NEXT:    [[LOAD_VAL:%.*]] = load i32, i32* %src, align 4
+; CHECK-NOT:     load i32, i32* %src, align 4
+; CHECK-NEXT:    [[TMP4:%.*]] = insertelement <4 x i32> poison, i32 [[LOAD_VAL]], i32 0
+; CHECK-NEXT:    [[TMP5:%.*]] = shufflevector <4 x i32> [[TMP4]], <4 x i32> poison, <4 x i32> zeroinitializer
+; CHECK-NEXT:    [[TMP6:%.*]] = getelementptr inbounds i32, i32* %dst, i64 [[TMP3]]
+; CHECK-NEXT:    [[TMP7:%.*]] = getelementptr inbounds i32, i32* [[TMP6]], i32 0
+; CHECK-NEXT:    [[STORE_PTR:%.*]] = bitcast i32* [[TMP7]] to <4 x i32>*
+; CHECK-NEXT:    call void @llvm.masked.store.v4i32.p0v4i32(<4 x i32> [[TMP5]], <4 x i32>* [[STORE_PTR]], i32 4, <4 x i1> [[LOOP_PRED]])
+; CHECK-NEXT:    [[IDX_NEXT]] = add i64 [[IDX]], 4
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i64 [[IDX_NEXT]], %n.vec
+; CHECK-NEXT:    br i1 [[CMP]], label %middle.block, label %vector.body
+
+entry:
+  br label %for.body
+
+for.body:                                         ; preds = %entry, %for.body
+  %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ]
+  %val = load i32, i32* %src, align 4
+  %arrayidx = getelementptr inbounds i32, i32* %dst, i64 %indvars.iv
+  store i32 %val, i32* %arrayidx, align 4
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %exitcond.not = icmp eq i64 %indvars.iv.next, %n
+  br i1 %exitcond.not, label %for.end, label %for.body
+
+for.end:                                          ; preds = %for.body, %entry
+  ret void
+}
+
+
 ; The original loop had a conditional uniform load. In this case we actually
 ; do need to perform conditional loads and so we end up using a gather instead.
 ; However, we at least ensure the mask is the overlap of the loop predicate
 ; and the original condition.
 define void @cond_uniform_load(i32* nocapture %dst, i32* nocapture readonly %src, i32* nocapture readonly %cond, i64 %n) #0 {
 ; CHECK-LABEL: @cond_uniform_load(
 ; CHECK:       vector.ph:
 ; CHECK:         [[TMP1:%.*]] = insertelement <4 x i32*> poison, i32* %src, i32 0
 ; CHECK-NEXT:    [[SRC_SPLAT:%.*]] = shufflevector <4 x i32*> [[TMP1]], <4 x i32*> poison, <4 x i32> zeroinitializer