[IVDescriptors] Support FOR where we have multiple sink pointed

[Allen]

Summary

Handles the case where Previous doesn't come before LastPrev incorrectly.

Fix https://github.com/llvm/llvm-project/issues/53483

[fhahn]

Thanks for the patch, response inline.

llvm/lib/Analysis/IVDescriptors.cpp

932

I think this handles the case where Previous doesn't come before LastPrev incorrectly. I added an additional test that crashes with the current version of the patch (02ee3fbff816).

To start with, you should be able to use return Previous->comesBefore(LastPrev); here. The reasoning is that if Previous comes before LastPrev, SinkCandidate already gets sunk after Previous, so there's nothing to do here.

I put up a follow-up patch (D118642), that should handle the other case correctly.

[Allen]

@fhahn ping

llvm/lib/Analysis/IVDescriptors.cpp

932

Yes, Thanks a lot for extending.

what need I do next step ? do you mind squash your patch D118642, or just fix the crash of your new added case 02ee3fbff816 ?

[fhahn]

llvm/lib/Analysis/IVDescriptors.cpp

932

what need I do next step ? do you mind squash your patch D118642, or just fix the crash of your new added case 02ee3fbff816 ?

As suggested in my comment, you should be able to update the patch to use

return Previous->comesBefore(LastPrev)

instead of

if (Previous->comesBefore(LastPrev))
  return true;

to fix the issue.

[Allen]

update to use return Previous->comesBefore(LastPrev); according comment

[Allen]

This comment was removed by Allen.

llvm/lib/Analysis/IVDescriptors.cpp

932

thanks for your comment very much!

[fhahn]

LGTM, thanks with the inline suggestions.

Handles the case where Previous doesn't come before LastPrev incorrectly.

Could you update the commit messages before landing this change? It should definitely handle this correctly rather than incorrectly :)

llvm/lib/Analysis/IVDescriptors.cpp

927

might be more compact to just check LastPrev->getParent() != Previous->getParent().

930

Maybe add a comment along the lines If LastPrev comes after the current Previous, SinkCandidate already gets sunk past Previous and there is nothing to do.

[fhahn]

Oh, and could you double check this version on its own fixes https://github.com/llvm/llvm-project/issues/53483 if you include it in the commit message?

[Allen]

In D118558#3312344, @fhahn wrote:

Oh, and could you double check this version on its own fixes https://github.com/llvm/llvm-project/issues/53483 if you include it in the commit message?

Yes, it still active for this test case.

llvm/lib/Analysis/IVDescriptors.cpp

927

done

930

done, much thanks for detail comment content to add.

[fhahn]

LGTM, thanks!

[fhahn]

If you do not yet have commit access, please let me know if you would like me to commit the patch on your behalf. In that case, please provide the name + email to use for the git authorship.

[Allen]

This comment was removed by Allen.

[Allen]

In D118558#3317006, @fhahn wrote:

If you do not yet have commit access, please let me know if you would like me to commit the patch on your behalf. In that case, please provide the name + email to use for the git authorship.

sorry for late reply, I pushed now as I'm on vacation yesterday.

llvm/lib/Analysis/IVDescriptors.cpp

if (DT->dominates(Previous,
                       SinkCandidate)) // We already are good w/o sinking.
       return true;
 
     if (SinkCandidate->getParent() != PhiBB ||
         SinkCandidate->mayHaveSideEffects() ||
         SinkCandidate->mayReadFromMemory() || SinkCandidate->isTerminator())
       return false;
 
-    // Do not try to sink an instruction multiple times (if multiple operands
-    // are first order recurrences).
-    // TODO: We can support this case, by sinking the instruction after the
-    // 'deepest' previous instruction.
-    if (SinkAfter.find(SinkCandidate) != SinkAfter.end())
+    // Try to sink an instruction after the 'deepest' previous instruction,
+    // which has multiple operands for first order recurrences.
+    auto It = SinkAfter.find(SinkCandidate);
+    if (It != SinkAfter.end()) {
+      auto LastPrev = It->second;
+      if (LastPrev->getParent() != Previous->getParent())
         return false;
 
+      // If LastPrev comes after the current Previous, SinkCandidate already
+      // gets sunk past Previous and nothing left to do.
+      return Previous->comesBefore(LastPrev);
+    }
+
     // If we reach a PHI node that is not dominated by Previous, we reached a
     // header PHI. No need for sinking.
     if (isa<PHINode>(SinkCandidate))
       return true;
 
     // Sink User tentatively and check its users
     InstrsToSink.insert(SinkCandidate);
     WorkList.push_back(SinkCandidate);

llvm/test/Transforms/LoopVectorize/first-order-recurrence-complex.ll

for:
   %iv.next = add nuw nsw i64 %iv, 1
   %exitcond = icmp eq i64 %iv.next, 2000
   br i1 %exitcond, label %exit, label %for
 
 exit:
   ret void
 }
 
-; TODO: We should be able to sink %tmp38 after %tmp60.
+; Sink %tmp38 after %tmp60, then it enable the loop vectorization.
 define void @instruction_with_2_FOR_operands() {
 ; CHECK-LABEL: @instruction_with_2_FOR_operands(
 ; CHECK-NEXT:  bb:
+; CHECK-NEXT:    [[SMAX:%.*]] = call i64 @llvm.smax.i64(i64 undef, i64 0)
+; CHECK-NEXT:    [[TMP0:%.*]] = add nuw i64 [[SMAX]], 1
+; CHECK-NEXT:    [[MIN_ITERS_CHECK:%.*]] = icmp ult i64 [[TMP0]], 4
+; CHECK-NEXT:    br i1 [[MIN_ITERS_CHECK]], label [[SCALAR_PH:%.*]], label [[VECTOR_PH:%.*]]
+; CHECK:       vector.ph:
+; CHECK-NEXT:    [[N_MOD_VF:%.*]] = urem i64 [[TMP0]], 4
+; CHECK-NEXT:    [[N_VEC:%.*]] = sub i64 [[TMP0]], [[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:    [[VECTOR_RECUR:%.*]] = phi <4 x float> [ <float poison, float poison, float poison, float undef>, [[VECTOR_PH]] ], [ [[BROADCAST_SPLAT3:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT:    [[VECTOR_RECUR1:%.*]] = phi <4 x float> [ <float poison, float poison, float poison, float undef>, [[VECTOR_PH]] ], [ [[BROADCAST_SPLAT:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT:    [[TMP1:%.*]] = add i64 [[INDEX]], 0
+; CHECK-NEXT:    [[TMP2:%.*]] = load float, float* undef, align 4
+; CHECK-NEXT:    [[BROADCAST_SPLATINSERT:%.*]] = insertelement <4 x float> poison, float [[TMP2]], i32 0
+; CHECK-NEXT:    [[BROADCAST_SPLAT]] = shufflevector <4 x float> [[BROADCAST_SPLATINSERT]], <4 x float> poison, <4 x i32> zeroinitializer
+; CHECK-NEXT:    [[TMP3:%.*]] = shufflevector <4 x float> [[VECTOR_RECUR1]], <4 x float> [[BROADCAST_SPLAT]], <4 x i32> <i32 3, i32 4, i32 5, i32 6>
+; CHECK-NEXT:    [[TMP4:%.*]] = load float, float* undef, align 4
+; CHECK-NEXT:    [[BROADCAST_SPLATINSERT2:%.*]] = insertelement <4 x float> poison, float [[TMP4]], i32 0
+; CHECK-NEXT:    [[BROADCAST_SPLAT3]] = shufflevector <4 x float> [[BROADCAST_SPLATINSERT2]], <4 x float> poison, <4 x i32> zeroinitializer
+; CHECK-NEXT:    [[TMP5:%.*]] = shufflevector <4 x float> [[VECTOR_RECUR]], <4 x float> [[BROADCAST_SPLAT3]], <4 x i32> <i32 3, i32 4, i32 5, i32 6>
+; CHECK-NEXT:    [[TMP6:%.*]] = fmul fast <4 x float> [[TMP5]], [[TMP3]]
+; CHECK-NEXT:    [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 4
+; CHECK-NEXT:    [[TMP7:%.*]] = icmp eq i64 [[INDEX_NEXT]], [[N_VEC]]
+; CHECK-NEXT:    br i1 [[TMP7]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP8:![0-9]+]]
+; CHECK:       middle.block:
+; CHECK-NEXT:    [[CMP_N:%.*]] = icmp eq i64 [[TMP0]], [[N_VEC]]
+; CHECK-NEXT:    [[VECTOR_RECUR_EXTRACT:%.*]] = extractelement <4 x float> [[BROADCAST_SPLAT3]], i32 3
+; CHECK-NEXT:    [[VECTOR_RECUR_EXTRACT_FOR_PHI:%.*]] = extractelement <4 x float> [[BROADCAST_SPLAT3]], i32 2
+; CHECK-NEXT:    [[VECTOR_RECUR_EXTRACT4:%.*]] = extractelement <4 x float> [[BROADCAST_SPLAT]], i32 3
+; CHECK-NEXT:    [[VECTOR_RECUR_EXTRACT_FOR_PHI5:%.*]] = extractelement <4 x float> [[BROADCAST_SPLAT]], i32 2
+; CHECK-NEXT:    br i1 [[CMP_N]], label [[BB74:%.*]], label [[SCALAR_PH]]
+; CHECK:       scalar.ph:
+; CHECK-NEXT:    [[SCALAR_RECUR_INIT6:%.*]] = phi float [ undef, [[BB:%.*]] ], [ [[VECTOR_RECUR_EXTRACT4]], [[MIDDLE_BLOCK]] ]
+; CHECK-NEXT:    [[SCALAR_RECUR_INIT:%.*]] = phi float [ undef, [[BB]] ], [ [[VECTOR_RECUR_EXTRACT]], [[MIDDLE_BLOCK]] ]
+; CHECK-NEXT:    [[BC_RESUME_VAL:%.*]] = phi i64 [ [[N_VEC]], [[MIDDLE_BLOCK]] ], [ 0, [[BB]] ]
 ; CHECK-NEXT:    br label [[BB13:%.*]]
 ; CHECK:       bb13:
-; CHECK-NEXT:    [[TMP37:%.*]] = phi float [ [[TMP60:%.*]], [[BB13]] ], [ undef, [[BB:%.*]] ]
-; CHECK-NEXT:    [[TMP27:%.*]] = phi float [ [[TMP49:%.*]], [[BB13]] ], [ undef, [[BB]] ]
-; CHECK-NEXT:    [[INDVARS_IV:%.*]] = phi i64 [ [[INDVARS_IV_NEXT:%.*]], [[BB13]] ], [ 0, [[BB]] ]
-; CHECK-NEXT:    [[TMP38:%.*]] = fmul fast float [[TMP37]], [[TMP27]]
+; CHECK-NEXT:    [[SCALAR_RECUR:%.*]] = phi float [ [[TMP60:%.*]], [[BB13]] ], [ [[SCALAR_RECUR_INIT]], [[SCALAR_PH]] ]
+; CHECK-NEXT:    [[SCALAR_RECUR7:%.*]] = phi float [ [[TMP49:%.*]], [[BB13]] ], [ [[SCALAR_RECUR_INIT6]], [[SCALAR_PH]] ]
+; CHECK-NEXT:    [[INDVARS_IV:%.*]] = phi i64 [ [[INDVARS_IV_NEXT:%.*]], [[BB13]] ], [ [[BC_RESUME_VAL]], [[SCALAR_PH]] ]
+; CHECK-NEXT:    [[TMP38:%.*]] = fmul fast float [[SCALAR_RECUR]], [[SCALAR_RECUR7]]
 ; CHECK-NEXT:    [[INDVARS_IV_NEXT]] = add nuw nsw i64 [[INDVARS_IV]], 1
 ; CHECK-NEXT:    [[TMP49]] = load float, float* undef, align 4
 ; CHECK-NEXT:    [[TMP60]] = load float, float* undef, align 4
 ; CHECK-NEXT:    [[TMP12:%.*]] = icmp slt i64 [[INDVARS_IV]], undef
-; CHECK-NEXT:    br i1 [[TMP12]], label [[BB13]], label [[BB74:%.*]]
+; CHECK-NEXT:    br i1 [[TMP12]], label [[BB13]], label [[BB74]], !llvm.loop [[LOOP9:![0-9]+]]
 ; CHECK:       bb74:
 ; CHECK-NEXT:    ret void
 ;
 
 
 bb:
   br label %bb13