ValueTracking: Handle nofpclass in computeKnownFPClass

[arsenm]

ping

[arsenm]

Cleanup & fcAllFlags

[arsenm]

b39deda3e1b3617650964de8a134ca828c1a2451

llvm/lib/Analysis/ValueTracking.cpp

void computeKnownFPClass(const Value *V, const APInt &DemandedElts,
   if (match(V, m_APFloat(C))) {
     // We know all of the classes for a scalar constant or a splat vector
     // constant!
     Known.KnownFPClasses = C->classify();
     Known.SignBit = C->isNegative();
     return;
   }
 
-  const Operator *Op = dyn_cast<Operator>(V);
-  if (!Op)
-    return;
-
   FPClassTest KnownNotFromFlags = fcNone;
-  if (const FPMathOperator *FPOp = dyn_cast<FPMathOperator>(Op)) {
+  if (const auto *CB = dyn_cast<CallBase>(V))
+    KnownNotFromFlags |= CB->getRetNoFPClass();
+  else if (const auto *Arg = dyn_cast<Argument>(V))
+    KnownNotFromFlags |= Arg->getNoFPClass();
+
+  const Operator *Op = dyn_cast<Operator>(V);
+  if (const FPMathOperator *FPOp = dyn_cast_or_null<FPMathOperator>(Op)) {
     if (FPOp->hasNoNaNs())
       KnownNotFromFlags |= fcNan;
     if (FPOp->hasNoInfs())
       KnownNotFromFlags |= fcInf;
+  }
 
-    // We no longer need to find out about these bits from inputs if we can
-    // assume this from flags.
-    InterestedClasses &= ~KnownNotFromFlags;
-  }
+  // assume this from flags/attributes.
 
   auto ClearClassesFromFlags = make_scope_exit([=, &Known] {
     Known.knownNot(KnownNotFromFlags);
   });
 
+  if (!Op)
+    return;
+
   // All recursive calls that increase depth must come after this.
   if (Depth == MaxAnalysisRecursionDepth)
     return;
 
   switch (Op->getOpcode()) {
   case Instruction::FNeg: {
     computeKnownFPClass(Op->getOperand(0), DemandedElts, InterestedClasses,
                         Known, Depth + 1, Q, TLI);

llvm/unittests/Analysis/ValueTrackingTest.cpp

TEST_F(ComputeKnownFPClassTest, SelectNNaNNInf) {
   parseAssembly(
       "define float @test(i1 %cond, float %arg0, float %arg1) {\n"
       "  %A = select nnan ninf i1 %cond, float %arg0, float %arg1"
       "  ret float %A\n"
       "}\n");
   expectKnownFPClass(~(fcNan | fcInf), std::nullopt);
 }
 
+TEST_F(ComputeKnownFPClassTest, SelectNoFPClassArgUnionAll) {
+  parseAssembly(
+      "define float @test(i1 %cond, float nofpclass(snan ninf nsub pzero pnorm) %arg0, float nofpclass(qnan nnorm nzero psub pinf) %arg1) {\n"
+      "  %A = select i1 %cond, float %arg0, float %arg1"
+      "  ret float %A\n"
+      "}\n");
+  expectKnownFPClass(fcAllFlags, std::nullopt);
+}
+
+TEST_F(ComputeKnownFPClassTest, SelectNoFPClassArgNoNan) {
+  parseAssembly(
+      "define float @test(i1 %cond, float nofpclass(nan) %arg0, float nofpclass(nan) %arg1) {\n"
+      "  %A = select i1 %cond, float %arg0, float %arg1"
+      "  ret float %A\n"
+      "}\n");
+  expectKnownFPClass(~fcNan, std::nullopt);
+}
+
+TEST_F(ComputeKnownFPClassTest, SelectNoFPClassArgNoPInf) {
+  parseAssembly(
+      "define float @test(i1 %cond, float nofpclass(inf) %arg0, float nofpclass(pinf) %arg1) {\n"
+      "  %A = select i1 %cond, float %arg0, float %arg1"
+      "  ret float %A\n"
+      "}\n");
+  expectKnownFPClass(~fcPosInf, std::nullopt);
+}
+
+TEST_F(ComputeKnownFPClassTest, SelectNoFPClassArgNoNInf) {
+  parseAssembly(
+      "define float @test(i1 %cond, float nofpclass(ninf) %arg0, float nofpclass(inf) %arg1) {\n"
+      "  %A = select i1 %cond, float %arg0, float %arg1"
+      "  ret float %A\n"
+      "}\n");
+  expectKnownFPClass(~fcNegInf, std::nullopt);
+}
+
+TEST_F(ComputeKnownFPClassTest, SelectNoFPClassCallSiteNoNan) {
+  parseAssembly(
+      "declare float @func()\n"
+      "define float @test() {\n"
+      "  %A = call nofpclass(nan) float @func()\n"
+      "  ret float %A\n"
+      "}\n");
+  expectKnownFPClass(~fcNan, std::nullopt);
+}
+
+TEST_F(ComputeKnownFPClassTest, SelectNoFPClassCallSiteNoZeros) {
+  parseAssembly(
+      "declare float @func()\n"
+      "define float @test() {\n"
+      "  %A = call nofpclass(zero) float @func()\n"
+      "  ret float %A\n"
+      "}\n");
+  expectKnownFPClass(~fcZero, std::nullopt);
+}
+
+TEST_F(ComputeKnownFPClassTest, SelectNoFPClassDeclarationNoNan) {
+  parseAssembly(
+      "declare nofpclass(nan) float @no_nans()\n"
+      "define float @test() {\n"
+      "  %A = call float @no_nans()\n"
+      "  ret float %A\n"
+      "}\n");
+  expectKnownFPClass(~fcNan, std::nullopt);
+}
+
+// Check nofpclass + ninf works on a callsite
+TEST_F(ComputeKnownFPClassTest, SelectNoFPClassCallSiteNoZerosNInfFlags) {
+  parseAssembly(
+      "declare float @func()\n"
+      "define float @test() {\n"
+      "  %A = call ninf nofpclass(zero) float @func()\n"
+      "  ret float %A\n"
+      "}\n");
+  expectKnownFPClass(~(fcZero | fcInf), std::nullopt);
+}
+
 TEST_F(ComputeKnownFPClassTest, FNegNInf) {
   parseAssembly(
       "define float @test(float %arg) {\n"
       "  %A = fneg ninf float %arg"
       "  ret float %A\n"
       "}\n");
   expectKnownFPClass(~fcInf, std::nullopt);
 }