[LoopVectorize] Add support for invariant stores of ordered reductions

[david-arm]

llvm/test/Transforms/LoopVectorize/AArch64/strict-fadd.ll

1364–1375

Hi @malharJ, can you add some proper CHECK lines for this test please similar to the other positive test cases above, e.g. @fadd_strict?

[malharJ]

• Updated LIT test with more detailed CHECKs.

The main difference between this and "fadd_strict" is the presence of
a store (to the invariant destination) in the middle.block

[malharJ]

• Removed redundant blank lines from test

[fhahn]

@igor.kirillov do you remember why this was restriction was added originally?

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

1020

this comment needs updating

[malharJ]

• Updated comment in source

[igor.kirillov]

In D126772#3560561, @fhahn wrote:

@igor.kirillov do you remember why this was restriction was added originally?

I don't remember actually. I know that it didn't work properly half a year ago but it does work now.

[malharJ]

ping

[david-arm]

LGTM! From private discussions I also understand you've run the LLVM test suite using scalable vectorisation on SVE hardware and all tests passed. Can you address my test comment before merging please?

llvm/test/Transforms/LoopVectorize/AArch64/strict-fadd.ll

1371

I think this should be something like

store float %[[RDX]], float* %[[DEST_PTR]]

because otherwise the tests might break if the numbering changes in future.

[malharJ]

llvm/test/Transforms/LoopVectorize/AArch64/strict-fadd.ll

1371

good spot, somehow I seem to have missed it.
I've corrected this in the submitted patch.

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

if (!EnableStrictReductions ||
       any_of(getInductionVars(), [&](auto &Induction) -> bool {
         InductionDescriptor IndDesc = Induction.second;
         return IndDesc.getExactFPMathInst();
       }))
     return false;
 
   // We can now only vectorize if all reductions with Exact FP math also
   // have the isOrdered flag set, which indicates that we can move the
   // reduction operations in-loop, and do not have intermediate store.
   return (all_of(getReductionVars(), [&](auto &Reduction) -> bool {
     const RecurrenceDescriptor &RdxDesc = Reduction.second;
-    return !RdxDesc.hasExactFPMath() ||
-           (RdxDesc.isOrdered() && !RdxDesc.IntermediateStore);
+    return !RdxDesc.hasExactFPMath() || RdxDesc.isOrdered();
   }));
 }
 
 bool LoopVectorizationLegality::isInvariantStoreOfReduction(StoreInst *SI) {
   return any_of(getReductionVars(), [&](auto &Reduction) -> bool {
     const RecurrenceDescriptor &RdxDesc = Reduction.second;
     return RdxDesc.IntermediateStore == SI;
   });

llvm/test/Transforms/LoopVectorize/AArch64/strict-fadd.ll

for.end:
   ret float %muladd2
 }
 
 declare float @llvm.fmuladd.f32(float, float, float)
 
 ; Test case with invariant store where fadd is strict.
 define void @reduction_store_to_invariant_address(float* %dst, float* readonly %src) {
 ; CHECK-ORDERED-LABEL: @reduction_store_to_invariant_address(
-; CHECK-ORDERED-NOT: vector.body
+; CHECK-ORDERED: entry
+; CHECK-ORDERED: %[[DEST_PTR:.*]] = getelementptr inbounds float, float* %dst, i64 42
+; CHECK-ORDERED: vector.body
+; CHECK-ORDERED: %[[VEC_PHI:.*]] = phi float [ 0.000000e+00, %vector.ph ], [ %[[RDX:.*]], %vector.body ]
+; CHECK-ORDERED: %[[LOAD_VEC:.*]] = load <8 x float>, <8 x float>*
+; CHECK-ORDERED: %[[RDX]] = call float @llvm.vector.reduce.fadd.v8f32(float %[[VEC_PHI]], <8 x float> %[[LOAD_VEC]])
+; CHECK-ORDERED: middle.block
+; CHECK-ORDERED: store float %4, float* %[[DEST_PTR]]
+; CHECK-ORDERED: for.body
+; CHECK-ORDERED: %[[LOAD:.*]] = load float, float*
+; CHECK-ORDERED: %[[FADD:.*]] = fadd float %{{.*}}, %[[LOAD]]
+; CHECK-ORDERED: store float %[[FADD]], float* %[[DEST_PTR]]
 
 ; CHECK-UNORDERED-LABEL: @reduction_store_to_invariant_address(
-; CHECK-UNORDERED-NOT: vector.body
+; CHECK-UNORDERED: entry
+; CHECK-UNORDERED: %[[DEST_PTR:.*]] = getelementptr inbounds float, float* %dst, i64 42
+; CHECK-UNORDERED: vector.body
+; CHECK-UNORDERED: %[[VEC_PHI:.*]] = phi <8 x float> [ <float 0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00>, %vector.ph ], [ %[[FADD_VEC:.*]], %vector.body ]
+; CHECK-UNORDERED: %[[LOAD_VEC:.*]] = load <8 x float>, <8 x float>*
+; CHECK-UNORDERED: %[[FADD_VEC]] = fadd <8 x float> %[[VEC_PHI]], %[[LOAD_VEC]]
+; CHECK-UNORDERED-NOT: call float @llvm.vector.reduce.fadd
+; CHECK-UNORDERED: middle.block
+; CHECK-UNORDERED: %[[RDX:.*]] = call float @llvm.vector.reduce.fadd.v8f32(float -0.000000e+00, <8 x float> %[[FADD_VEC]])
+; CHECK-UNORDERED: store float %[[RDX]], float* %[[DEST_PTR]]
+; CHECK-UNORDERED: for.body
+; CHECK-UNORDERED: %[[LOAD:.*]] = load float, float*
+; CHECK-UNORDERED: %[[FADD:.*]] = fadd float {{.*}}, %[[LOAD]]
+; CHECK-UNORDERED: store float %[[FADD]], float* %[[DEST_PTR]]
 
 ; CHECK-NOT-VECTORIZED-LABEL: @reduction_store_to_invariant_address(
 ; CHECK-NOT-VECTORIZED-NOT: vector.body
 
 entry:
   %arrayidx = getelementptr inbounds float, float* %dst, i64 42
   store float 0.000000e+00, float* %arrayidx, align 4
   br label %for.body
 
 for.body:
   %0 = phi float [ 0.000000e+00, %entry ], [ %add, %for.body ]
   %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ]
   %arrayidx1 = getelementptr inbounds float, float* %src, i64 %indvars.iv
   %1 = load float, float* %arrayidx1, align 4
   %add = fadd float %0, %1
   store float %add, float* %arrayidx, align 4
   %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
   %exitcond = icmp eq i64 %indvars.iv.next, 1000
-  br i1 %exitcond, label %for.cond.cleanup, label %for.body
+  br i1 %exitcond, label %for.cond.cleanup, label %for.body, !llvm.loop !0
 
 for.cond.cleanup:
   ret void
 }
 
 !0 = distinct !{!0, !5, !9, !11}
 !1 = distinct !{!1, !5, !10, !11}
 !2 = distinct !{!2, !6, !9, !11}