[LoopVectorize] Fix bug where predicated loads/stores were dropped

[joechrisellis]

Summary

This commit fixes a bug where the loop vectoriser fails to predicate
loads/stores when interleaving for targets that support masked
loads and stores.

Code such as:

1  void foo(int *restrict data1, int *restrict data2)
2  {
3    int counter = 1024;
4    while (counter--)
5      if (data1[counter] > data2[counter])
6        data1[counter] = data2[counter];
7  }

... could previously be transformed in such a way that the predicated
store implied by:

if (data1[counter] > data2[counter])
   data1[counter] = data2[counter];

... was lost, resulting in miscompiles.

This bug was causing some tests in llvm-test-suite to fail when built
for SVE.

[joechrisellis]

Hmm. I anticipated there might be some test failures for this one.

[joechrisellis]

Minor change to x86-predication.ll test.

The codegen from before the patch had a simpler CFG because we weren’t
predicating on loads. In this particular case, it turns out that it’s actually
fine to not explicitly predicate on the loads, but only because they exist in a
basic block next to a udiv which does require predication, so we get it ‘for
free’ in a sense. This is not guaranteed to always be the case, though, so we
have to be careful.

The result of this patch is that the IR CFG for scalarize_and_sink_gather in
x86-predication.ll is larger, and so the test needs to be updated to reflect
that. Ultimately, though, the machine code lowering is the same, so no loss in
performance.

[peterwaller-arm]

llvm/test/Transforms/LoopVectorize/AArch64/scalarize-store-with-predication.ll

4

I suspect using -debug needs a REQUIRES: asserts line, since the tests will be run presumably with a release build where -debug is unavailable.

The output of:

rg -C3 -- 'RUN.* -debug-only'

... in the llvm/test directory appears to show that they generally have such a REQUIRES line.

I find this a little surprising since presumably this doesn't work on an asserts+release build; but my reading is that enabling assertions may enable debug output:

https://github.com/llvm/llvm-project/blob/8396aeb07cddd8ab9a6a154a4ab7ac56fc24bda5/llvm/cmake/modules/HandleLLVMOptions.cmake#L61-L76

[joechrisellis]

Improve test and address review comments.

• @peterwaller-arm: add REQUIRES: asserts line to new test.

[david-arm]

llvm/test/Transforms/LoopVectorize/AArch64/scalarize-store-with-predication.ll

4

@peterwaller-arm yep you're right this needs:

REQUIRES: asserts

and it will enable debug.

llvm/test/Transforms/LoopVectorize/X86/x86-predication.ll

69

Hi @joechrisellis you've lost the variable T0 here that was used to ensure the load fed into the udiv. It would be good if you could re-add this.

[joechrisellis]

Address comments.

• @david-arm: re-add T0 variable into test.

[joechrisellis]

Gentle ping. 🙂

llvm/test/Transforms/LoopVectorize/X86/x86-predication.ll

69

Sure, will re-add this. For context, I removed this because the control flow is now slightly different now that the loads are predicated, so I thought the cleanest thing to do would be to remove the use of this variable. I've added it back now.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

1520

I'm not sure about the default. Why does it make sense to default to VF = 1? Can all users of the function instead pass the right VF?

llvm/test/Transforms/LoopVectorize/AArch64/scalarize-store-with-predication.ll

2

as you are only checking specifically for small IR snippets of the predicated blocks, do we actually need -dce -instcombine? Running those potentially make the test more fragile, because it may be impacted by changes to -dce -instcombine.

6

can we instead just do 2>&1 | FileCheck %s in one go? Splitting them up makes it a bit more inconvenient to reproduce a failure, because 2 lines need to copied/executed?

22

this seems already out of sync with the IR below.... Shouldn't the restrict be translated to noalias in the function definition below? FWIW I think it would be more helpful if you'd add a comment to the test saying this checks that the store in the if.then block gets properly predicated if VF = 1 or something like that, as the IR is already quite compact :)

46

Is this check crucial? Seems like the IR checks already check this and we could remove it and execute it even with asserts.

57

do we need 2 loads here or could one be removed and compare to a constant?

73

If it's not SVE specific, can we have 2 run lines, one with -mattr=+sve and one without?

[joechrisellis]

Address review comments.

• @fhahn: test changes.

[joechrisellis]

Hi @fhahn -- thanks for the review!

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

1520

The reason I added this optional parameter is because isPredicatedInst is used in this context:

bool IsPredicated = LoopVectorizationPlanner::getDecisionAndClampRange(
    [&](ElementCount VF) { return CM.isPredicatedInst(I, VF); }, Range);

I wanted to make sure that the VF (which is used to make decisions on how to clamp the range) is made available to isScalarWithPredication in the fallback case for this particular call, but keep the same signature for calls elsewhere. Given what I think the function is supposed to do, I can see why it might not make sense to pass a VF parameter for the general case, but here I _think_ we need it. Other references to isPredicatedInst do not supply a VF parameter, so in the fallback cases, they will default to VF = 1 when we reach the call for isScalarWithPredication. This is just pushing the optional parameter up the stack so that it can be used for isPredicatedInst too.

If there's a better thing to do here I'm happy to make changes. 😄

llvm/test/Transforms/LoopVectorize/AArch64/scalarize-store-with-predication.ll

22

Good spot. I'll remove the restrict in the C code example since it isn't strictly necessary here anyways.

46

Fair enough, done!

57

I would like to keep this IR as-is for consistency with the C code at the top of the test. 🙂

[david-arm]

LGTM! I think there is potentially a regression in the X86 test 'scalarize_and_sink_gather' in terms of number of generated lines of assembly, but the fix looks right for now. We could potentially improve the code quality in a future patch that looks for instructions with matching predicates and folds them into the same predicated block.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

1520

his is just pushing the optional parameter up the stack so that it can be used for isPredicatedInst too.

That's fair enough, but my point was that there are only 2 uses of the function I think, so the default parameter is not really that useful and I think it would be better to just explicitly pass the VF in at both call sites. Also, the default of VF(1) is surprising (but I realize it is used elsewhere) and it's probably better to force the caller to think about what VF to pass in.

llvm/test/Transforms/LoopVectorize/AArch64/scalarize-store-with-predication.ll

57

I would like to keep this IR as-is for consistency with the C code at the top of the test. 🙂

Personally I think that the C example code should not stand in the way of having a simpler IR test. If you really want to keep the C code, can't you simplify the C as well? (With the even simpler IR I think the C code adds more noise than value it adds).

[joechrisellis]

Address review comments.

• @fhahn:
  • make VF parameter to isPredicatedInst required.
  • simplify IR test.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

public:
   /// predication for that VF.
   bool
   isScalarWithPredication(Instruction *I,
                           ElementCount VF = ElementCount::getFixed(1)) 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) {
+  bool isPredicatedInst(Instruction *I, ElementCount VF) {
     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);
+    return isScalarWithPredication(I, VF);
   }
 
   /// Returns true if \p I is a memory instruction with consecutive memory
   /// access that can be widened.
   bool
   memoryInstructionCanBeWidened(Instruction *I,
                                 ElementCount VF = ElementCount::getFixed(1));
 
bool LoopVectorizationCostModel::useEmulatedMaskMemRefHack(Instruction *I){
   // TODO: Cost model for emulated masked load/store is completely
   // broken. This hack guides the cost model to use an artificially
   // high enough value to practically disable vectorization with such
   // operations, except where previously deployed legality hack allowed
   // using very low cost values. This is to avoid regressions coming simply
   // from moving "masked load/store" check from legality to cost model.
   // Masked Load/Gather emulation was previously never allowed.
   // Limited number of Masked Store/Scatter emulation was allowed.
-  assert(isPredicatedInst(I) && "Expecting a scalar emulated instruction");
+  assert(isPredicatedInst(I, ElementCount::getFixed(1)) &&
+         "Expecting a scalar emulated instruction");
   return isa<LoadInst>(I) ||
          (isa<StoreInst>(I) &&
           NumPredStores > NumberOfStoresToPredicate);
 }
 
 void LoopVectorizationCostModel::collectInstsToScalarize(ElementCount VF) {
   // If we aren't vectorizing the loop, or if we've already collected the
   // instructions to scalarize, there's nothing to do. Collection may already
LoopVectorizationCostModel::getMemInstScalarizationCost(Instruction *I,
 
   // Get the overhead of the extractelement and insertelement instructions
   // we might create due to scalarization.
   Cost += getScalarizationOverhead(I, VF);
 
   // If we have a predicated load/store, it will need extra i1 extracts and
   // conditional branches, but may not be executed for each vector lane. Scale
   // the cost by the probability of executing the predicated block.
-  if (isPredicatedInst(I)) {
+  if (isPredicatedInst(I, ElementCount::getFixed(1))) {
     Cost /= getReciprocalPredBlockProb();
 
     // Add the cost of an i1 extract and a branch
     auto *Vec_i1Ty =
         VectorType::get(IntegerType::getInt1Ty(ValTy->getContext()), VF);
     Cost += TTI.getScalarizationOverhead(
         Vec_i1Ty, APInt::getAllOnesValue(VF.getKnownMinValue()),
         /*Insert=*/false, /*Extract=*/true);

 VPBasicBlock *VPRecipeBuilder::handleReplication(
     Instruction *I, VFRange &Range, VPBasicBlock *VPBB,
     VPlanPtr &Plan) {
   bool IsUniform = LoopVectorizationPlanner::getDecisionAndClampRange(
       [&](ElementCount VF) { return CM.isUniformAfterVectorization(I, VF); },
       Range);
 
   bool IsPredicated = LoopVectorizationPlanner::getDecisionAndClampRange(
-      [&](ElementCount VF) { return CM.isScalarWithPredication(I, VF); },
-      Range);
+      [&](ElementCount VF) { return CM.isPredicatedInst(I, VF); }, Range);
 
   auto *Recipe = new VPReplicateRecipe(I, Plan->mapToVPValues(I->operands()),
                                        IsUniform, IsPredicated);
   setRecipe(I, Recipe);
   Plan->addVPValue(I, Recipe);
 
   // Find if I uses a predicated instruction. If so, it will use its scalar
   // value. Avoid hoisting the insert-element which packs the scalar value into

llvm/test/Transforms/LoopVectorize/AArch64/scalarize-store-with-predication.ll

+; RUN: opt -loop-vectorize -force-vector-width=1 -force-vector-interleave=2 -S -o - < %s | FileCheck %s
+; RUN: opt -mattr=+sve -loop-vectorize -force-vector-width=1 -force-vector-interleave=2 -S -o - < %s | FileCheck %s
+
+target triple = "aarch64-unknown-linux-gnu"
+
+; This test is defending against a bug that appeared when we have a target
+; configuration where masked loads/stores are legal -- e.g. AArch64 with SVE.
+; Predication would not be applied during interleaving, enabling the
+; possibility of superfluous loads/stores which could result in miscompiles.
+; This test checks that, when we disable vectorisation and force interleaving,
+; stores are predicated properly.
+;
+; This is _not_ an SVE-specific test. The same bug could manifest on any
+; architecture with masked loads/stores, but we use SVE for testing purposes
+; here.
+
+define void @foo(i32* %data1, i32* %data2) {
+; CHECK-LABEL: @foo(
+; CHECK:       vector.body:
+; CHECK:         br i1 {{%.*}}, label %pred.store.if, label %pred.store.continue
+; CHECK:       pred.store.if:
+; CHECK-NEXT:    store i32 {{%.*}}, i32* {{%.*}}
+; CHECK-NEXT:    br label %pred.store.continue
+; CHECK:       pred.store.continue:
+; CHECK-NEXT:    br i1 {{%.*}}, label %pred.store.if2, label %pred.store.continue3
+; CHECK:       pred.store.if2:
+; CHECK-NEXT:    store i32 {{%.*}}, i32* {{%.*}}
+; CHECK-NEXT:    br label %pred.store.continue3
+; CHECK:       pred.store.continue3:
+
+entry:
+  br label %while.body
+
+while.body:
+  %i = phi i64 [ 1023, %entry ], [ %i.next, %if.end ]
+  %arrayidx = getelementptr inbounds i32, i32* %data1, i64 %i
+  %ld = load i32, i32* %arrayidx, align 4
+  %cmp = icmp sgt i32 %ld, %ld
+  br i1 %cmp, label %if.then, label %if.end
+
+if.then:
+  store i32 %ld, i32* %arrayidx, align 4
+  br label %if.end
+
+if.end:
+  %i.next = add nsw i64 %i, -1
+  %tobool.not = icmp eq i64 %i, 0
+  br i1 %tobool.not, label %while.end, label %while.body
+
+while.end:
+  ret void
+}

llvm/test/Transforms/LoopVectorize/X86/x86-predication.ll

for.end:
   ret i32 %tmp8
 }
 
 ; This test ensures that a load, which would have been widened otherwise is
 ; instead scalarized if Cost-Model so decided as part of its
 ; sink-scalar-operands optimization for predicated instructions.
 ;
 ; SINK-GATHER: vector.body:
-; SINK-GATHER: pred.udiv.if:
+; SINK-GATHER: pred.load.if:
 ; SINK-GATHER:   %[[T0:.+]] = load i32, i32* %{{.*}}, align 4
-; SINK-GATHER:   %{{.*}} = udiv i32 %[[T0]], %{{.*}}
+; SINK-GATHER: pred.load.continue:
+; SINK-GATHER:   %[[T1:.+]] = phi i32 [ poison, %vector.body ], [ %[[T0]], %pred.load.if ]
+; SINK-GATHER: pred.udiv.if:
+; SINK-GATHER:   %{{.*}} = udiv i32 %[[T1]], %{{.*}}
 ; SINK-GATHER: pred.udiv.continue:
 define i32 @scalarize_and_sink_gather(i32* %a, i1 %c, i32 %x, i64 %n) {
 entry:
   br label %for.body
 
 for.body:
   %i = phi i64 [ 0, %entry ], [ %i.next, %for.inc ]
   %r = phi i32 [ 0, %entry ], [ %tmp6, %for.inc ]