[LoopVectorize][Fix] Crash when invariant store address is calculated inside loop

[igor.kirillov]

Summary

Fixes #57572

Generally LICM pass is responsible for sinking out code that calculates
invariant address inside loop as it only needed to be calculated once.
But in rare case it does not happen we will not be vectorizing the
loop.

[igor.kirillov]

This is a bug fix for - https://github.com/llvm/llvm-project/issues/57572

[igor.kirillov]

Add an empty line to the test and a reference to the github issue

[mgabka]

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

938

nit: "check its" -> "check if"

llvm/test/Transforms/LoopVectorize/reduction-with-invariant-store.ll

517

nit: maybe better to call it : reduc_store_invariant_addr_not_hoisted ?

[david-arm]

LGTM with nits addressed!

llvm/test/Transforms/LoopVectorize/reduction-with-invariant-store.ll

518

nit: I think this should be calculated instead of calucalated

[fhahn]

You could add something like Fixes #57572 to the description, so the issue will be auto-closed once the fix is pushed.

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

955

I think the place for the check means we now may not vectorize cases we vectorized before, specifically if we have reductions and a store to an invariant address that's defined in the loop body that's not used by a reduction, as in the example below.

You likely want to drop the && !getReductionVars().empty()) { check above and only check the address if isInvariantStoreOfReduction is true.

define i32 @test(i32* %dst, i32* readonly %src) {
entry:
  br label %for.body

for.body:                                         ; preds = %for.body, %entry
  %sum = phi i32 [ 0, %entry ], [ %add, %for.body ]
  %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
  %gep.src = getelementptr inbounds i32, i32* %src, i64 %iv
  %0 = load i32, i32* %gep.src, align 4
  %add = add nsw i32 %sum, %0
  %gep.dst = getelementptr inbounds i32, i32* %dst, i64 42
  store i32 0, i32* %gep.dst, align 4
  %iv.next = add nuw nsw i64 %iv, 1
  %exitcond = icmp eq i64 %iv.next, 1000
  br i1 %exitcond, label %exit, label %for.body

exit:                                             ; preds = %for.body
  %add.lcssa = phi i32 [ %add, %for.body ]
  ret i32 %add.lcssa
}

[igor.kirillov]

Fixes

[igor.kirillov]

Update commit message

[fhahn]

LGTM with the inline comment addressed, thanks!

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

955

To simplify the code, please add an early continue if the store is not isInvariantStoreOfReduction(SI) and remove the 2 separate checks on line 953 and 939

Please also add the test above to ensure the regression is guarded against.

[igor.kirillov]

Refactor `isInvariantStoreOfReduction' check

[fhahn]

Did you add the extra test for the regression in the original version of the patch? It looks like it wasn't included in 2d60d7ba1a26/

[igor.kirillov]

In D133687#3826724, @fhahn wrote:

Did you add the extra test for the regression in the original version of the patch? It looks like it wasn't included in 2d60d7ba1a26/

Done - rGa94a8555/

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

bool LoopVectorizationLegality::canVectorizeMemory() {
   if (!LAI->canVectorizeMemory())
     return false;
 
   // We can vectorize stores to invariant address when final reduction value is
   // guaranteed to be stored at the end of the loop. Also, if decision to
   // vectorize loop is made, runtime checks are added so as to make sure that
   // invariant address won't alias with any other objects.
   if (!LAI->getStoresToInvariantAddresses().empty()) {
-    // For each invariant address, check its last stored value is unconditional.
+    // For each invariant address, check if last stored value is unconditional
+    // and the address is not calculated inside the loop.
     for (StoreInst *SI : LAI->getStoresToInvariantAddresses()) {
-      if (isInvariantStoreOfReduction(SI) &&
-          blockNeedsPredication(SI->getParent())) {
+      if (!isInvariantStoreOfReduction(SI))
+        continue;
+
+      if (blockNeedsPredication(SI->getParent())) {
         reportVectorizationFailure(
             "We don't allow storing to uniform addresses",
             "write of conditional recurring variant value to a loop "
             "invariant address could not be vectorized",
             "CantVectorizeStoreToLoopInvariantAddress", ORE, TheLoop);
         return false;
       }
+
+      // Invariant address should be defined outside of loop. LICM pass usually
+      // makes sure it happens, but in rare cases it does not, we do not want
+      // to overcomplicate vectorization to support this case.
+      if (Instruction *Ptr = dyn_cast<Instruction>(SI->getPointerOperand())) {
+        if (TheLoop->contains(Ptr)) {
+          reportVectorizationFailure(
+              "Invariant address is calculated inside the loop",
+              "write to a loop invariant address could not "
+              "be vectorized",
+              "CantVectorizeStoreToLoopInvariantAddress", ORE, TheLoop);
+          return false;
+        }
+      }
     }
 
     if (LAI->hasDependenceInvolvingLoopInvariantAddress()) {
       // For each invariant address, check its last stored value is the result
       // of one of our reductions.
       //
       // We do not check if dependence with loads exists because they are
       // currently rejected earlier in LoopAccessInfo::analyzeLoop. In case this

llvm/test/Transforms/LoopVectorize/reduction-with-invariant-store.ll

body:
   store i64 %red.next, i64* %dst, align 8
   %iv.next = add nuw nsw i32 %iv, 1
   %ec = icmp ugt i32 %iv, 363
   br i1 %ec, label %exit, label %body
 
 exit:
   ret void
 }
+
+define void @reduc_store_invariant_addr_not_hoisted(i32* %dst, i32* readonly %src) {
+; CHECK-LABEL: @reduc_store_invariant_addr_not_hoisted
+; CHECK-NOT: vector.body:
+entry:
+  br label %for.body
+
+for.body:
+  %sum = phi i32 [ 0, %entry ], [ %add, %for.body ]
+  %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
+  %gep.src = getelementptr inbounds i32, i32* %src, i64 %iv
+  %0 = load i32, i32* %gep.src, align 4
+  %add = add nsw i32 %sum, %0
+  %gep.dst = getelementptr inbounds i32, i32* %dst, i64 42
+  store i32 %add, i32* %gep.dst, align 4
+  %iv.next = add nuw nsw i64 %iv, 1
+  %exitcond = icmp eq i64 %iv.next, 1000
+  br i1 %exitcond, label %exit, label %for.body
+
+exit:
+  ret void
+}