Resolve NITs

Author: peterbell10

llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp

@@ -2464,18 +2464,17 @@ bool NVPTXDAGToDAGISel::tryBF16ArithToFMA(SDNode *N) {
     return false;
 
   const NVPTXSubtarget *STI = TM.getSubtargetImpl();
-  const bool IsNativelySupported =
-      STI->getSmVersion() >= 90 && STI->getPTXVersion() >= 78;
-  if (IsNativelySupported)
+  if (STI->hasNativeBF16Support(N->getOpcode()))
     return false;
 
-  assert(VT == MVT::bf16 || VT == MVT::v2bf16);
-  const bool IsVec = VT == MVT::v2bf16;
+  const bool IsVec = VT.isVector();
+  assert(!IsVec || VT.getVectorNumElements() == 2);
   SDLoc DL(N);
   SDValue N0 = N->getOperand(0);
   SDValue N1 = N->getOperand(1);
   SmallVector<SDValue, 3> Operands;
   auto GetConstant = [&](float Value) -> SDValue {
+    // BF16 immediates must be legalized to integer register values
     APFloat APF(Value);
     bool LosesInfo;
     APF.convert(APFloat::BFloat(), APFloat::rmNearestTiesToEven, &LosesInfo);

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

@@ -535,34 +535,7 @@ NVPTXTargetLowering::NVPTXTargetLowering(const NVPTXTargetMachine &TM,
 
   auto setBF16OperationAction = [&](unsigned Op, MVT VT, LegalizeAction Action,
                                     LegalizeAction NoBF16Action) {
-    bool IsOpSupported = STI.hasBF16Math();
-    switch (Op) {
-    // Several BF16 instructions are available on sm_90 only.
-    case ISD::FADD:
-    case ISD::FMUL:
-    case ISD::FSUB:
-    case ISD::SELECT:
-    case ISD::SELECT_CC:
-    case ISD::SETCC:
-    case ISD::FEXP2:
-    case ISD::FCEIL:
-    case ISD::FFLOOR:
-    case ISD::FNEARBYINT:
-    case ISD::FRINT:
-    case ISD::FROUNDEVEN:
-    case ISD::FTRUNC:
-      IsOpSupported = STI.getSmVersion() >= 90 && STI.getPTXVersion() >= 78;
-      break;
-    // Several BF16 instructions are available on sm_80 only.
-    case ISD::FMINNUM:
-    case ISD::FMAXNUM:
-    case ISD::FMAXNUM_IEEE:
-    case ISD::FMINNUM_IEEE:
-    case ISD::FMAXIMUM:
-    case ISD::FMINIMUM:
-      IsOpSupported &= STI.getSmVersion() >= 80 && STI.getPTXVersion() >= 70;
-      break;
-    }
+    bool IsOpSupported = STI.hasNativeBF16Support(Op);
     setOperationAction(
         Op, VT, IsOpSupported ? Action : NoBF16Action);
   };
@@ -862,11 +835,10 @@ NVPTXTargetLowering::NVPTXTargetLowering(const NVPTXTargetMachine &TM,
       AddPromotedToType(Op, MVT::bf16, MVT::f32);
   }
 
-  // Lower bf16 add/mul/sub as fma when it avoids promotion
+  // On SM80, we select add/mul/sub as fma to avoid promotion to float
   for (const auto &Op : {ISD::FADD, ISD::FMUL, ISD::FSUB}) {
     for (const auto &VT : {MVT::bf16, MVT::v2bf16}) {
-      if (getOperationAction(Op, VT) != Legal &&
-          getOperationAction(ISD::FMA, VT) == Legal) {
+      if (!STI.hasNativeBF16Support(Op) && STI.hasNativeBF16Support(ISD::FMA)) {
         setOperationAction(Op, VT, Custom);
       }
     }

llvm/lib/Target/NVPTX/NVPTXSubtarget.cpp

@@ -70,6 +70,38 @@ bool NVPTXSubtarget::allowFP16Math() const {
   return hasFP16Math() && NoF16Math == false;
 }
 
+bool NVPTXSubtarget::hasNativeBF16Support(int Opcode) const {
+  if (!hasBF16Math())
+    return false;
+
+  switch (Opcode) {
+  // Several BF16 instructions are available on sm_90 only.
+  case ISD::FADD:
+  case ISD::FMUL:
+  case ISD::FSUB:
+  case ISD::SELECT:
+  case ISD::SELECT_CC:
+  case ISD::SETCC:
+  case ISD::FEXP2:
+  case ISD::FCEIL:
+  case ISD::FFLOOR:
+  case ISD::FNEARBYINT:
+  case ISD::FRINT:
+  case ISD::FROUNDEVEN:
+  case ISD::FTRUNC:
+    return getSmVersion() >= 90 && getPTXVersion() >= 78;
+  // Several BF16 instructions are available on sm_80 only.
+  case ISD::FMINNUM:
+  case ISD::FMAXNUM:
+  case ISD::FMAXNUM_IEEE:
+  case ISD::FMINNUM_IEEE:
+  case ISD::FMAXIMUM:
+  case ISD::FMINIMUM:
+    return getSmVersion() >= 80 && getPTXVersion() >= 70;
+  }
+  return true;
+}
+
 void NVPTXSubtarget::failIfClustersUnsupported(
     std::string const &FailureMessage) const {
   if (hasClusters())

llvm/lib/Target/NVPTX/NVPTXSubtarget.h

@@ -118,6 +118,8 @@ class NVPTXSubtarget : public NVPTXGenSubtargetInfo {
   }
   bool hasTargetName() const { return !TargetName.empty(); }
 
+  bool hasNativeBF16Support(int Opcode) const;
+
   // Get maximum value of required alignments among the supported data types.
   // From the PTX ISA doc, section 8.2.3:
   //  The memory consistency model relates operations executed on memory