[LV][IVDescriptors] Fix recurrence identity element for FMin and FMax reductions.

[karthiksenthil]

Summary

For a min and max reduction idioms, the identity (i.e. neutral) element
should be datatype's highest and lowest possible values respectively.
Current implementation in IVDescriptors incorrectly returns -Inf for FMin
reduction and +Inf for FMax reduction. This patch fixes this bug which
was causing incorrect reduction computation results in loops vectorized
by LV.

[dmgreen]

Do we have any tests for Max as well as Min? It looks like the llvm/test/Transforms/LoopVectorize/AArch64/scalable-reduction-inloop-cond.ll test may need updating.

[spatel]

Should these be guarded with an FMF check?

The enum definition says:

FMin,       ///< FP min implemented in terms of select(cmp()).
FMax,       ///< FP max implemented in terms of select(cmp()).

So if we have a NaN input:
fmin(NaN, Inf) --> select (fcmp olt NaN, Inf), NaN, Inf --> Inf (so the NaN input did not survive)

Is it possible to write unit-tests to check these directly? I see a llvm/unittests/Analysis/IVDescriptorsTest.cpp test file...

[karthiksenthil]

Thanks for the reviews! I've addressed the comments below -

In D137220#3901561, @dmgreen wrote:

Do we have any tests for Max as well as Min? It looks like the llvm/test/Transforms/LoopVectorize/AArch64/scalable-reduction-inloop-cond.ll test may need updating.

I've added new tests in llvm/unittests/Analysis/IVDescriptorsTest.cpp to cover both FMin and FMax reductions. Updated the LIT test as well.

In D137220#3902612, @spatel wrote:

Should these be guarded with an FMF check?

The enum definition says:

FMin,       ///< FP min implemented in terms of select(cmp()).
FMax,       ///< FP max implemented in terms of select(cmp()).

So if we have a NaN input:
fmin(NaN, Inf) --> select (fcmp olt NaN, Inf), NaN, Inf --> Inf (so the NaN input did not survive)

Do you mean that identity value computation for FMin/FMax should be guarded with a FMF.noNaNs() check? What would be the identity value when the flag is absent?

Is it possible to write unit-tests to check these directly? I see a llvm/unittests/Analysis/IVDescriptorsTest.cpp test file...

I've updated the test file with new tests, I can extend this for testing missing nnan flag scenario as well.

[spatel]

In D137220#3902612, @spatel wrote:

Should these be guarded with an FMF check?

The enum definition says:

FMin,       ///< FP min implemented in terms of select(cmp()).
FMax,       ///< FP max implemented in terms of select(cmp()).

So if we have a NaN input:
fmin(NaN, Inf) --> select (fcmp olt NaN, Inf), NaN, Inf --> Inf (so the NaN input did not survive)

Do you mean that identity value computation for FMin/FMax should be guarded with a FMF.noNaNs() check? What would be the identity value when the flag is absent?

Yes, I'm not sure what it means if we don't have FMF.noNaNs(). Is it possible to create this recurrence without that FMF? If not, can we assert that FMF.noNaNs() is set?

[karthiksenthil]

In D137220#3903584, @spatel wrote:

In D137220#3902612, @spatel wrote:

Should these be guarded with an FMF check?

The enum definition says:

FMin,       ///< FP min implemented in terms of select(cmp()).
FMax,       ///< FP max implemented in terms of select(cmp()).

So if we have a NaN input:
fmin(NaN, Inf) --> select (fcmp olt NaN, Inf), NaN, Inf --> Inf (so the NaN input did not survive)

Do you mean that identity value computation for FMin/FMax should be guarded with a FMF.noNaNs() check? What would be the identity value when the flag is absent?

Yes, I'm not sure what it means if we don't have FMF.noNaNs(). Is it possible to create this recurrence without that FMF? If not, can we assert that FMF.noNaNs() is set?

Looks like the following code in RecurrenceDescriptor::isRecurrenceInstr guarantees that nnan flag is set for FMin/FMax reductions -

case Instruction::FCmp:
case Instruction::ICmp:
case Instruction::Call:
  if (isSelectCmpRecurrenceKind(Kind))
    return isSelectCmpPattern(L, OrigPhi, I, Prev);
  if (isIntMinMaxRecurrenceKind(Kind) ||
      (((FuncFMF.noNaNs() && FuncFMF.noSignedZeros()) ||
        (isa<FPMathOperator>(I) && I->hasNoNaNs() &&
         I->hasNoSignedZeros())) &&
       isFPMinMaxRecurrenceKind(Kind)))
    return isMinMaxPattern(I, Kind, Prev);

I've added asserts in the latest diff to check that FMF.noNaNs() is set.

[spatel]

In D137220#3903744, @karthiksenthil wrote:

I've added asserts in the latest diff to check that FMF.noNaNs() is set.

Thanks - should it also assert for noSignedZeros?

[karthiksenthil]

In D137220#3904940, @spatel wrote:

In D137220#3903744, @karthiksenthil wrote:

I've added asserts in the latest diff to check that FMF.noNaNs() is set.

Thanks - should it also assert for noSignedZeros?

Yes, I think it would be better to keep this assert in-sync with the checks in RecurrenceDescriptor::isRecurrenceInstr. I have updated the asserts and unit-tests in latest diff.

[spatel]

LGTM

[karthiksenthil]

Rebase to llvm trunk. No functional updates.

[karthiksenthil]

Thanks for the reviews and approval! I don't have access to commit this change. @spatel can you commit this revision to llvm trunk?

llvm/lib/Analysis/IVDescriptors.cpp

case RecurKind::UMax:
     return ConstantInt::get(Tp, 0);
   case RecurKind::SMin:
     return ConstantInt::get(Tp,
                             APInt::getSignedMaxValue(Tp->getIntegerBitWidth()));
   case RecurKind::SMax:
     return ConstantInt::get(Tp,
                             APInt::getSignedMinValue(Tp->getIntegerBitWidth()));
   case RecurKind::FMin:
-    return ConstantFP::getInfinity(Tp, true);
+    assert((FMF.noNaNs() && FMF.noSignedZeros()) &&
+           "nnan, nsz is expected to be set for FP min reduction.");
+    return ConstantFP::getInfinity(Tp, false /*Negative*/);
   case RecurKind::FMax:
-    return ConstantFP::getInfinity(Tp, false);
+    assert((FMF.noNaNs() && FMF.noSignedZeros()) &&
+           "nnan, nsz is expected to be set for FP max reduction.");
+    return ConstantFP::getInfinity(Tp, true /*Negative*/);
   case RecurKind::SelectICmp:
   case RecurKind::SelectFCmp:
     return getRecurrenceStartValue();
     break;
   default:
     llvm_unreachable("Unknown recurrence kind");
   }
 }

llvm/test/Transforms/LoopVectorize/AArch64/scalable-reduction-inloop-cond.ll

 ; CHECK-NEXT:    [[TMP6:%.*]] = getelementptr inbounds float, float* [[TMP5]], i32 0
 ; CHECK-NEXT:    [[TMP7:%.*]] = bitcast float* [[TMP6]] to <vscale x 4 x float>*
 ; CHECK-NEXT:    [[WIDE_LOAD:%.*]] = load <vscale x 4 x float>, <vscale x 4 x float>* [[TMP7]], align 4
 ; CHECK-NEXT:    [[TMP8:%.*]] = fcmp une <vscale x 4 x float> [[WIDE_LOAD]], shufflevector (<vscale x 4 x float> insertelement (<vscale x 4 x float> poison, float 3.000000e+00, i32 0), <vscale x 4 x float> poison, <vscale x 4 x i32> zeroinitializer)
 ; CHECK-NEXT:    [[TMP9:%.*]] = getelementptr float, float* [[A:%.*]], i64 [[TMP4]]
 ; CHECK-NEXT:    [[TMP10:%.*]] = getelementptr float, float* [[TMP9]], i32 0
 ; CHECK-NEXT:    [[TMP11:%.*]] = bitcast float* [[TMP10]] to <vscale x 4 x float>*
 ; CHECK-NEXT:    [[WIDE_MASKED_LOAD:%.*]] = call <vscale x 4 x float> @llvm.masked.load.nxv4f32.p0nxv4f32(<vscale x 4 x float>* [[TMP11]], i32 4, <vscale x 4 x i1> [[TMP8]], <vscale x 4 x float> poison)
-; CHECK-NEXT:    [[TMP12:%.*]] = select fast <vscale x 4 x i1> [[TMP8]], <vscale x 4 x float> [[WIDE_MASKED_LOAD]], <vscale x 4 x float> shufflevector (<vscale x 4 x float> insertelement (<vscale x 4 x float> poison, float 0xFFF0000000000000, i32 0), <vscale x 4 x float> poison, <vscale x 4 x i32> zeroinitializer)
+; CHECK-NEXT:    [[TMP12:%.*]] = select fast <vscale x 4 x i1> [[TMP8]], <vscale x 4 x float> [[WIDE_MASKED_LOAD]], <vscale x 4 x float> shufflevector (<vscale x 4 x float> insertelement (<vscale x 4 x float> poison, float 0x7FF0000000000000, i32 0), <vscale x 4 x float> poison, <vscale x 4 x i32> zeroinitializer)
 ; CHECK-NEXT:    [[TMP13:%.*]] = call fast float @llvm.vector.reduce.fmin.nxv4f32(<vscale x 4 x float> [[TMP12]])
 ; CHECK-NEXT:    [[RDX_MINMAX_CMP:%.*]] = fcmp fast olt float [[TMP13]], [[VEC_PHI]]
 ; CHECK-NEXT:    [[RDX_MINMAX_SELECT]] = select fast i1 [[RDX_MINMAX_CMP]], float [[TMP13]], float [[VEC_PHI]]
 ; CHECK-NEXT:    [[TMP14:%.*]] = call i64 @llvm.vscale.i64()
 ; CHECK-NEXT:    [[TMP15:%.*]] = mul i64 [[TMP14]], 4
 ; CHECK-NEXT:    [[INDEX_NEXT]] = add nuw i64 [[INDEX]], [[TMP15]]
 ; CHECK-NEXT:    [[TMP16:%.*]] = icmp eq i64 [[INDEX_NEXT]], [[N_VEC]]
 ; CHECK-NEXT:    br i1 [[TMP16]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP4:![0-9]+]]

llvm/test/Transforms/LoopVectorize/reduction-inloop-cond.ll

 ; CHECK:       pred.load.if5:
 ; CHECK-NEXT:    [[TMP21:%.*]] = or i64 [[INDEX]], 3
 ; CHECK-NEXT:    [[TMP22:%.*]] = getelementptr inbounds float, float* [[A]], i64 [[TMP21]]
 ; CHECK-NEXT:    [[TMP23:%.*]] = load float, float* [[TMP22]], align 4
 ; CHECK-NEXT:    [[TMP24:%.*]] = insertelement <4 x float> [[TMP19]], float [[TMP23]], i64 3
 ; CHECK-NEXT:    br label [[PRED_LOAD_CONTINUE6]]
 ; CHECK:       pred.load.continue6:
 ; CHECK-NEXT:    [[TMP25:%.*]] = phi <4 x float> [ [[TMP19]], [[PRED_LOAD_CONTINUE4]] ], [ [[TMP24]], [[PRED_LOAD_IF5]] ]
-; CHECK-NEXT:    [[TMP26:%.*]] = select fast <4 x i1> [[TMP2]], <4 x float> [[TMP25]], <4 x float> <float 0xFFF0000000000000, float 0xFFF0000000000000, float 0xFFF0000000000000, float 0xFFF0000000000000>
+; CHECK-NEXT:    [[TMP26:%.*]] = select fast <4 x i1> [[TMP2]], <4 x float> [[TMP25]], <4 x float> <float 0x7FF0000000000000, float 0x7FF0000000000000, float 0x7FF0000000000000, float 0x7FF0000000000000>
 ; CHECK-NEXT:    [[TMP27:%.*]] = call fast float @llvm.vector.reduce.fmin.v4f32(<4 x float> [[TMP26]])
 ; CHECK-NEXT:    [[TMP28]] = call fast float @llvm.minnum.f32(float [[TMP27]], float [[VEC_PHI]])
 ; CHECK-NEXT:    [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 4
 ; CHECK-NEXT:    [[TMP29:%.*]] = icmp eq i64 [[INDEX_NEXT]], [[N_VEC]]
 ; CHECK-NEXT:    br i1 [[TMP29]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP4:![0-9]+]]
 ; CHECK:       middle.block:
 ; CHECK-NEXT:    [[CMP_N:%.*]] = icmp eq i64 [[N_VEC]], [[N]]
 ; CHECK-NEXT:    br i1 [[CMP_N]], label [[FOR_END:%.*]], label [[SCALAR_PH]]

llvm/unittests/Analysis/IVDescriptorsTest.cpp

runWithLoopInfoAndSE(
         PHINode *Inst_i = dyn_cast<PHINode>(&Header->front());
         assert(Inst_i->getName() == "tmp30");
         InductionDescriptor IndDesc;
         bool IsInductionPHI =
             InductionDescriptor::isInductionPHI(Inst_i, L, &SE, IndDesc);
         EXPECT_TRUE(IsInductionPHI);
       });
 }
+
+// This tests that correct identity value is returned for a RecurrenceDescriptor
+// that describes FMin reduction idiom.
+TEST(IVDescriptorsTest, FMinRednIdentity) {
+  // Parse the module.
+  LLVMContext Context;
+
+  std::unique_ptr<Module> M = parseIR(Context,
+                                      R"(define float @foo(float* %A, i64 %ub) {
+entry:
+  br label %for.body
+
+for.body:
+  %i = phi i64 [ 0, %entry ], [ %i.next, %for.body ]
+  %fmin = phi float [ 1.000000e+00, %entry ], [ %fmin.next, %for.body ]
+  %arrayidx = getelementptr inbounds float, float* %A, i64 %i
+  %ld = load float, float* %arrayidx
+  %fmin.cmp = fcmp nnan nsz olt float %fmin, %ld
+  %fmin.next = select nnan nsz i1 %fmin.cmp, float %fmin, float %ld
+  %i.next = add nsw i64 %i, 1
+  %cmp = icmp slt i64 %i.next, %ub
+  br i1 %cmp, label %for.body, label %for.end
+
+for.end:
+  %fmin.lcssa = phi float [ %fmin.next, %for.body ]
+  ret float %fmin.lcssa
+})");
+
+  runWithLoopInfoAndSE(
+      *M, "foo", [&](Function &F, LoopInfo &LI, ScalarEvolution &SE) {
+        Function::iterator FI = F.begin();
+        // First basic block is entry - skip it.
+        BasicBlock *Header = &*(++FI);
+        assert(Header->getName() == "for.body");
+        Loop *L = LI.getLoopFor(Header);
+        EXPECT_NE(L, nullptr);
+        BasicBlock::iterator BBI = Header->begin();
+        assert((&*BBI)->getName() == "i");
+        ++BBI;
+        PHINode *Phi = dyn_cast<PHINode>(&*BBI);
+        assert(Phi->getName() == "fmin");
+        RecurrenceDescriptor Rdx;
+        bool IsRdxPhi = RecurrenceDescriptor::isReductionPHI(Phi, L, Rdx);
+        EXPECT_TRUE(IsRdxPhi);
+        RecurKind Kind = Rdx.getRecurrenceKind();
+        EXPECT_EQ(Kind, RecurKind::FMin);
+        Type *Ty = Phi->getType();
+        Value *Id = Rdx.getRecurrenceIdentity(Kind, Ty, Rdx.getFastMathFlags());
+        // Identity value for FP min reduction is +Inf.
+        EXPECT_EQ(Id, ConstantFP::getInfinity(Ty, false /*Negative*/));
+      });
+}
+
+// This tests that correct identity value is returned for a RecurrenceDescriptor
+// that describes FMax reduction idiom.
+TEST(IVDescriptorsTest, FMaxRednIdentity) {
+  // Parse the module.
+  LLVMContext Context;
+
+  std::unique_ptr<Module> M = parseIR(Context,
+                                      R"(define float @foo(float* %A, i64 %ub) {
+entry:
+  br label %for.body
+
+for.body:
+  %i = phi i64 [ 0, %entry ], [ %i.next, %for.body ]
+  %fmax = phi float [ 1.000000e+00, %entry ], [ %fmax.next, %for.body ]
+  %arrayidx = getelementptr inbounds float, float* %A, i64 %i
+  %ld = load float, float* %arrayidx
+  %fmax.cmp = fcmp nnan nsz ogt float %fmax, %ld
+  %fmax.next = select nnan nsz i1 %fmax.cmp, float %fmax, float %ld
+  %i.next = add nsw i64 %i, 1
+  %cmp = icmp slt i64 %i.next, %ub
+  br i1 %cmp, label %for.body, label %for.end
+
+for.end:
+  %fmax.lcssa = phi float [ %fmax.next, %for.body ]
+  ret float %fmax.lcssa
+})");
+
+  runWithLoopInfoAndSE(
+      *M, "foo", [&](Function &F, LoopInfo &LI, ScalarEvolution &SE) {
+        Function::iterator FI = F.begin();
+        // First basic block is entry - skip it.
+        BasicBlock *Header = &*(++FI);
+        assert(Header->getName() == "for.body");
+        Loop *L = LI.getLoopFor(Header);
+        EXPECT_NE(L, nullptr);
+        BasicBlock::iterator BBI = Header->begin();
+        assert((&*BBI)->getName() == "i");
+        ++BBI;
+        PHINode *Phi = dyn_cast<PHINode>(&*BBI);
+        assert(Phi->getName() == "fmax");
+        RecurrenceDescriptor Rdx;
+        bool IsRdxPhi = RecurrenceDescriptor::isReductionPHI(Phi, L, Rdx);
+        EXPECT_TRUE(IsRdxPhi);
+        RecurKind Kind = Rdx.getRecurrenceKind();
+        EXPECT_EQ(Kind, RecurKind::FMax);
+        Type *Ty = Phi->getType();
+        Value *Id = Rdx.getRecurrenceIdentity(Kind, Ty, Rdx.getFastMathFlags());
+        // Identity value for FP max reduction is -Inf.
+        EXPECT_EQ(Id, ConstantFP::getInfinity(Ty, true /*Negative*/));
+      });
+}