support f32x2 instructions for Blackwell

This is a rewrite of previous work that legalized v2f32 into an i64
register. Here we keep the type non-legal, and selectively legalize it
for certain operations (FADD, FSUB, FMUL, FMA). Additional operations
are handled to improve codegen quality.

Author: Prince781

llvm/include/llvm/Target/TargetSelectionDAG.td

@@ -818,6 +818,10 @@ def step_vector : SDNode<"ISD::STEP_VECTOR", SDTypeProfile<1, 1,
 def scalar_to_vector : SDNode<"ISD::SCALAR_TO_VECTOR", SDTypeProfile<1, 1, []>,
                               []>;
 
+def build_pair : SDNode<"ISD::BUILD_PAIR", SDTypeProfile<1, 2,
+                        [SDTCisInt<0>, SDTCisInt<1>, SDTCisInt<2>]>, []>;
+
+
 // vector_extract/vector_insert are deprecated. extractelt/insertelt
 // are preferred.
 def vector_extract : SDNode<"ISD::EXTRACT_VECTOR_ELT",

llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp

@@ -121,6 +121,12 @@ void NVPTXDAGToDAGISel::Select(SDNode *N) {
   case NVPTXISD::SETP_BF16X2:
     SelectSETP_BF16X2(N);
     return;
+  case NVPTXISD::FADD_F32X2:
+  case NVPTXISD::FSUB_F32X2:
+  case NVPTXISD::FMUL_F32X2:
+  case NVPTXISD::FMA_F32X2:
+    SelectF32X2Op(N);
+    return;
   case NVPTXISD::LoadV2:
   case NVPTXISD::LoadV4:
     if (tryLoadVector(N))
@@ -295,6 +301,30 @@ bool NVPTXDAGToDAGISel::SelectSETP_BF16X2(SDNode *N) {
   return true;
 }
 
+void NVPTXDAGToDAGISel::SelectF32X2Op(SDNode *N) {
+  unsigned Opcode;
+  switch (N->getOpcode()) {
+  case NVPTXISD::FADD_F32X2:
+    Opcode = NVPTX::FADD_F32X2;
+    break;
+  case NVPTXISD::FSUB_F32X2:
+    Opcode = NVPTX::FSUB_F32X2;
+    break;
+  case NVPTXISD::FMUL_F32X2:
+    Opcode = NVPTX::FMUL_F32X2;
+    break;
+  case NVPTXISD::FMA_F32X2:
+    Opcode = NVPTX::FMA_F32X2;
+    break;
+  default:
+    llvm_unreachable("Unexpected opcode!");
+  }
+  SDLoc DL(N);
+  SmallVector<SDValue> NewOps(N->ops());
+  SDNode *NewNode = CurDAG->getMachineNode(Opcode, DL, MVT::i64, NewOps);
+  ReplaceNode(N, NewNode);
+}
+
 // Find all instances of extract_vector_elt that use this v2f16 vector
 // and coalesce them into a scattering move instruction.
 bool NVPTXDAGToDAGISel::tryEXTRACT_VECTOR_ELEMENT(SDNode *N) {

llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.h

@@ -88,6 +88,7 @@ class LLVM_LIBRARY_VISIBILITY NVPTXDAGToDAGISel : public SelectionDAGISel {
   bool tryConstantFP(SDNode *N);
   bool SelectSETP_F16X2(SDNode *N);
   bool SelectSETP_BF16X2(SDNode *N);
+  void SelectF32X2Op(SDNode *N);
   bool tryEXTRACT_VECTOR_ELEMENT(SDNode *N);
   void SelectV2I64toI128(SDNode *N);
   void SelectI128toV2I64(SDNode *N);

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

@@ -866,6 +866,14 @@ NVPTXTargetLowering::NVPTXTargetLowering(const NVPTXTargetMachine &TM,
   setBF16OperationAction(ISD::FNEG, MVT::v2bf16, Legal, Expand);
   // (would be) Library functions.
 
+  if (STI.hasF32x2Instructions()) {
+    // Handle custom lowering for: v2f32 = OP v2f32, v2f32
+    for (const auto &Op : {ISD::FADD, ISD::FSUB, ISD::FMUL, ISD::FMA})
+      setOperationAction(Op, MVT::v2f32, Custom);
+    // Handle custom lowering for: i64 = bitcast v2f32
+    setOperationAction(ISD::BITCAST, MVT::v2f32, Custom);
+  }
+
   // These map to conversion instructions for scalar FP types.
   for (const auto &Op : {ISD::FCEIL, ISD::FFLOOR, ISD::FNEARBYINT, ISD::FRINT,
                          ISD::FROUNDEVEN, ISD::FTRUNC}) {
@@ -1066,6 +1074,10 @@ const char *NVPTXTargetLowering::getTargetNodeName(unsigned Opcode) const {
     MAKE_CASE(NVPTXISD::STACKSAVE)
     MAKE_CASE(NVPTXISD::SETP_F16X2)
     MAKE_CASE(NVPTXISD::SETP_BF16X2)
+    MAKE_CASE(NVPTXISD::FADD_F32X2)
+    MAKE_CASE(NVPTXISD::FSUB_F32X2)
+    MAKE_CASE(NVPTXISD::FMUL_F32X2)
+    MAKE_CASE(NVPTXISD::FMA_F32X2)
     MAKE_CASE(NVPTXISD::Dummy)
     MAKE_CASE(NVPTXISD::MUL_WIDE_SIGNED)
     MAKE_CASE(NVPTXISD::MUL_WIDE_UNSIGNED)
@@ -2099,24 +2111,58 @@ SDValue NVPTXTargetLowering::LowerBITCAST(SDValue Op, SelectionDAG &DAG) const {
   // Handle bitcasting from v2i8 without hitting the default promotion
   // strategy which goes through stack memory.
   EVT FromVT = Op->getOperand(0)->getValueType(0);
-  if (FromVT != MVT::v2i8) {
-    return Op;
-  }
-
-  // Pack vector elements into i16 and bitcast to final type
-  SDLoc DL(Op);
-  SDValue Vec0 = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, MVT::i8,
-                             Op->getOperand(0), DAG.getIntPtrConstant(0, DL));
-  SDValue Vec1 = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, MVT::i8,
-                             Op->getOperand(0), DAG.getIntPtrConstant(1, DL));
-  SDValue Extend0 = DAG.getNode(ISD::ZERO_EXTEND, DL, MVT::i16, Vec0);
-  SDValue Extend1 = DAG.getNode(ISD::ZERO_EXTEND, DL, MVT::i16, Vec1);
-  SDValue Const8 = DAG.getConstant(8, DL, MVT::i16);
-  SDValue AsInt = DAG.getNode(
-      ISD::OR, DL, MVT::i16,
-      {Extend0, DAG.getNode(ISD::SHL, DL, MVT::i16, {Extend1, Const8})});
   EVT ToVT = Op->getValueType(0);
-  return MaybeBitcast(DAG, DL, ToVT, AsInt);
+  SDLoc DL(Op);
+
+  if (FromVT == MVT::v2i8) {
+    // Pack vector elements into i16 and bitcast to final type
+    SDValue Vec0 = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, MVT::i8,
+                               Op->getOperand(0), DAG.getIntPtrConstant(0, DL));
+    SDValue Vec1 = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, MVT::i8,
+                               Op->getOperand(0), DAG.getIntPtrConstant(1, DL));
+    SDValue Extend0 = DAG.getNode(ISD::ZERO_EXTEND, DL, MVT::i16, Vec0);
+    SDValue Extend1 = DAG.getNode(ISD::ZERO_EXTEND, DL, MVT::i16, Vec1);
+    SDValue Const8 = DAG.getConstant(8, DL, MVT::i16);
+    SDValue AsInt = DAG.getNode(
+        ISD::OR, DL, MVT::i16,
+        {Extend0, DAG.getNode(ISD::SHL, DL, MVT::i16, {Extend1, Const8})});
+    EVT ToVT = Op->getValueType(0);
+    return MaybeBitcast(DAG, DL, ToVT, AsInt);
+  }
+
+  if (FromVT == MVT::v2f32) {
+    assert(ToVT == MVT::i64);
+
+    // A bitcast to i64 from v2f32.
+    // See if we can legalize the operand.
+    const SDValue &Operand = Op->getOperand(0);
+    if (Operand.getOpcode() == ISD::BUILD_VECTOR) {
+      const SDValue &BVOp0 = Operand.getOperand(0);
+      const SDValue &BVOp1 = Operand.getOperand(1);
+
+      auto CastToAPInt = [](SDValue Op) -> APInt {
+        if (Op->isUndef())
+          return APInt(64, 0); // undef values default to 0
+        return cast<ConstantFPSDNode>(Op)->getValueAPF().bitcastToAPInt().zext(
+            64);
+      };
+
+      if ((BVOp0->isUndef() || isa<ConstantFPSDNode>(BVOp0)) &&
+          (BVOp1->isUndef() || isa<ConstantFPSDNode>(BVOp1))) {
+        // cast two constants
+        APInt Value(64, 0);
+        Value = CastToAPInt(BVOp0) | CastToAPInt(BVOp1).shl(32);
+        SDValue Const = DAG.getConstant(Value, DL, MVT::i64);
+        return DAG.getBitcast(ToVT, Const);
+      }
+
+      // otherwise build an i64
+      return DAG.getNode(ISD::BUILD_PAIR, DL, MVT::i64,
+                         DAG.getBitcast(MVT::i32, BVOp0),
+                         DAG.getBitcast(MVT::i32, BVOp1));
+    }
+  }
+  return Op;
 }
 
 // We can init constant f16x2/v2i16/v4i8 with a single .b32 move.  Normally it
@@ -3055,6 +3101,13 @@ bool NVPTXTargetLowering::splitValueIntoRegisterParts(
   return false;
 }
 
+const TargetRegisterClass *
+NVPTXTargetLowering::getRegClassFor(MVT VT, bool isDivergent) const {
+  if (VT == MVT::v2f32)
+    return &NVPTX::Int64RegsRegClass;
+  return TargetLowering::getRegClassFor(VT, isDivergent);
+}
+
 // This creates target external symbol for a function parameter.
 // Name of the symbol is composed from its index and the function name.
 // Negative index corresponds to special parameter (unsized array) used for
@@ -5055,10 +5108,10 @@ static SDValue PerformEXTRACTCombine(SDNode *N,
       IsPTXVectorType(VectorVT.getSimpleVT()))
     return SDValue(); // Native vector loads already combine nicely w/
                       // extract_vector_elt.
-  // Don't mess with singletons or v2*16, v4i8 and v8i8 types, we already
+  // Don't mess with singletons or v2*16, v4i8, v8i8, or v2f32 types, we already
   // handle them OK.
   if (VectorVT.getVectorNumElements() == 1 || Isv2x16VT(VectorVT) ||
-      VectorVT == MVT::v4i8 || VectorVT == MVT::v8i8)
+      VectorVT == MVT::v4i8 || VectorVT == MVT::v8i8 || VectorVT == MVT::v2f32)
     return SDValue();
 
   // Don't mess with undef values as sra may be simplified to 0, not undef.
@@ -5478,6 +5531,45 @@ static void ReplaceCopyFromReg_128(SDNode *N, SelectionDAG &DAG,
   Results.push_back(NewValue.getValue(3));
 }
 
+static void ReplaceF32x2Op(SDNode *N, SelectionDAG &DAG,
+                           SmallVectorImpl<SDValue> &Results,
+                           bool UseFTZ) {
+  SDLoc DL(N);
+  EVT OldResultTy = N->getValueType(0); // <2 x float>
+  assert(OldResultTy == MVT::v2f32 && "Unexpected result type for F32x2 op!");
+
+  SmallVector<SDValue> NewOps;
+
+  // whether we use FTZ (TODO)
+
+  // replace with NVPTX F32x2 op:
+  unsigned Opcode;
+  switch (N->getOpcode()) {
+  case ISD::FADD:
+    Opcode = NVPTXISD::FADD_F32X2;
+    break;
+  case ISD::FSUB:
+    Opcode = NVPTXISD::FSUB_F32X2;
+    break;
+  case ISD::FMUL:
+    Opcode = NVPTXISD::FMUL_F32X2;
+    break;
+  case ISD::FMA:
+    Opcode = NVPTXISD::FMA_F32X2;
+    break;
+  default:
+    llvm_unreachable("Unexpected opcode");
+  }
+
+  // bitcast operands: <2 x float> -> i64
+  for (const SDValue &Op : N->ops())
+    NewOps.push_back(DAG.getNode(ISD::BITCAST, DL, MVT::i64, Op));
+
+  // cast i64 result of new op back to <2 x float>
+  SDValue NewValue = DAG.getNode(Opcode, DL, MVT::i64, NewOps);
+  Results.push_back(DAG.getBitcast(OldResultTy, NewValue));
+}
+
 void NVPTXTargetLowering::ReplaceNodeResults(
     SDNode *N, SmallVectorImpl<SDValue> &Results, SelectionDAG &DAG) const {
   switch (N->getOpcode()) {
@@ -5495,6 +5587,12 @@ void NVPTXTargetLowering::ReplaceNodeResults(
   case ISD::CopyFromReg:
     ReplaceCopyFromReg_128(N, DAG, Results);
     return;
+  case ISD::FADD:
+  case ISD::FSUB:
+  case ISD::FMUL:
+  case ISD::FMA:
+    ReplaceF32x2Op(N, DAG, Results, useF32FTZ(DAG.getMachineFunction()));
+    return;
   }
 }
 

llvm/lib/Target/NVPTX/NVPTXISelLowering.h

@@ -55,6 +55,10 @@ enum NodeType : unsigned {
   FSHR_CLAMP,
   MUL_WIDE_SIGNED,
   MUL_WIDE_UNSIGNED,
+  FADD_F32X2,
+  FMUL_F32X2,
+  FSUB_F32X2,
+  FMA_F32X2,
   SETP_F16X2,
   SETP_BF16X2,
   BFE,
@@ -311,6 +315,9 @@ class NVPTXTargetLowering : public TargetLowering {
                               SDValue *Parts, unsigned NumParts, MVT PartVT,
                               std::optional<CallingConv::ID> CC) const override;
 
+  const TargetRegisterClass *getRegClassFor(MVT VT,
+                                            bool isDivergent) const override;
+
   void ReplaceNodeResults(SDNode *N, SmallVectorImpl<SDValue> &Results,
                           SelectionDAG &DAG) const override;
   SDValue PerformDAGCombine(SDNode *N, DAGCombinerInfo &DCI) const override;

llvm/lib/Target/NVPTX/NVPTXInstrInfo.td

@@ -165,6 +165,7 @@ def hasHWROT32 : Predicate<"Subtarget->hasHWROT32()">;
 def noHWROT32 : Predicate<"!Subtarget->hasHWROT32()">;
 def hasDotInstructions : Predicate<"Subtarget->hasDotInstructions()">;
 def hasTcgen05Instructions : Predicate<"Subtarget->hasTcgen05Instructions()">;
+def hasF32x2Instructions : Predicate<"Subtarget->hasF32x2Instructions()">;
 
 def True : Predicate<"true">;
 def False : Predicate<"false">;
@@ -2638,13 +2639,13 @@ class LastCallArgInstVT<NVPTXRegClass regclass, ValueType vt> :
   NVPTXInst<(outs), (ins regclass:$a), "$a",
             [(LastCallArg (i32 0), vt:$a)]>;
 
-def CallArgI64     : CallArgInst<Int64Regs>;
+def CallArgI64     : CallArgInstVT<Int64Regs, i64>;
 def CallArgI32     : CallArgInstVT<Int32Regs, i32>;
 def CallArgI16     : CallArgInstVT<Int16Regs, i16>;
 def CallArgF64     : CallArgInst<Float64Regs>;
 def CallArgF32     : CallArgInst<Float32Regs>;
 
-def LastCallArgI64 : LastCallArgInst<Int64Regs>;
+def LastCallArgI64 : LastCallArgInstVT<Int64Regs, i64>;
 def LastCallArgI32 : LastCallArgInstVT<Int32Regs, i32>;
 def LastCallArgI16 : LastCallArgInstVT<Int16Regs, i16>;
 def LastCallArgF64 : LastCallArgInst<Float64Regs>;
@@ -3371,6 +3372,9 @@ let hasSideEffects = false in {
   def V2F32toF64 : NVPTXInst<(outs Float64Regs:$d),
                              (ins Float32Regs:$s1, Float32Regs:$s2),
                              "mov.b64 \t$d, {{$s1, $s2}};", []>;
+  def V2F32toI64 : NVPTXInst<(outs Int64Regs:$d),
+                             (ins Float32Regs:$s1, Float32Regs:$s2),
+                             "mov.b64 \t$d, {{$s1, $s2}};", []>;
 
   // unpack a larger int register to a set of smaller int registers
   def I64toV4I16 : NVPTXInst<(outs Int16Regs:$d1, Int16Regs:$d2,
@@ -3435,6 +3439,10 @@ def : Pat<(v2bf16 (build_vector bf16:$a, bf16:$b)),
           (V2I16toI32 $a, $b)>;
 def : Pat<(v2i16 (build_vector i16:$a, i16:$b)),
           (V2I16toI32 $a, $b)>;
+def : Pat<(v2f32 (build_vector f32:$a, f32:$b)),
+          (V2F32toI64 $a, $b)>;
+def : Pat<(i64 (build_pair i32:$a, i32:$b)),
+          (V2I32toI64 $a, $b)>;
 
 def: Pat<(v2i16 (scalar_to_vector i16:$a)),
          (CVT_u32_u16 $a, CvtNONE)>;

llvm/lib/Target/NVPTX/NVPTXIntrinsics.td

@@ -1581,6 +1581,28 @@ def INT_NVVM_ADD_RM_D : F_MATH_2<"add.rm.f64 \t$dst, $src0, $src1;",
 def INT_NVVM_ADD_RP_D : F_MATH_2<"add.rp.f64 \t$dst, $src0, $src1;",
   Float64Regs, Float64Regs, Float64Regs, int_nvvm_add_rp_d>;
 
+// F32x2 ops (sm_100+)
+
+def FADD_F32X2 : NVPTXInst<(outs Int64Regs:$res),
+                           (ins Int64Regs:$a, Int64Regs:$b),
+                           "add.rn.f32x2 \t$res, $a, $b;", []>,
+                 Requires<[hasF32x2Instructions]>;
+
+def FSUB_F32X2 : NVPTXInst<(outs Int64Regs:$res),
+                           (ins Int64Regs:$a, Int64Regs:$b),
+                           "sub.rn.f32x2 \t$res, $a, $b;", []>,
+                 Requires<[hasF32x2Instructions]>;
+
+def FMUL_F32X2 : NVPTXInst<(outs Int64Regs:$res),
+                           (ins Int64Regs:$a, Int64Regs:$b),
+                           "mul.rn.f32x2 \t$res, $a, $b;", []>,
+                 Requires<[hasF32x2Instructions]>;
+
+def FMA_F32X2  : NVPTXInst<(outs Int64Regs:$res),
+                           (ins Int64Regs:$a, Int64Regs:$b, Int64Regs:$c),
+                           "fma.rn.f32x2 \t$res, $a, $b;", []>,
+                 Requires<[hasF32x2Instructions]>;
+
 //
 // BFIND
 //

llvm/lib/Target/NVPTX/NVPTXRegisterInfo.td

@@ -62,7 +62,9 @@ def Int16Regs : NVPTXRegClass<[i16, f16, bf16], 16, (add (sequence "RS%u", 0, 4)
 def Int32Regs : NVPTXRegClass<[i32, v2f16, v2bf16, v2i16, v4i8], 32,
                               (add (sequence "R%u", 0, 4),
                               VRFrame32, VRFrameLocal32)>;
-def Int64Regs : NVPTXRegClass<[i64], 64, (add (sequence "RL%u", 0, 4), VRFrame64, VRFrameLocal64)>;
+def Int64Regs : NVPTXRegClass<[i64, v2f32], 64,
+                              (add (sequence "RL%u", 0, 4),
+                              VRFrame64, VRFrameLocal64)>;
 // 128-bit regs are not defined as general regs in NVPTX. They are used for inlineASM only.
 def Int128Regs : NVPTXRegClass<[i128], 128, (add (sequence "RQ%u", 0, 4))>;
 def Float32Regs : NVPTXRegClass<[f32], 32, (add (sequence "F%u", 0, 4))>;

llvm/lib/Target/NVPTX/NVPTXSubtarget.h

@@ -97,6 +97,7 @@ class NVPTXSubtarget : public NVPTXGenSubtargetInfo {
   bool hasDotInstructions() const {
     return SmVersion >= 61 && PTXVersion >= 50;
   }
+
   // Tcgen05 instructions in Blackwell family
   bool hasTcgen05Instructions() const {
     bool HasTcgen05 = false;
@@ -112,6 +113,8 @@ class NVPTXSubtarget : public NVPTXGenSubtargetInfo {
     return HasTcgen05 && PTXVersion >= 86;
   }
 
+  bool hasF32x2Instructions() const { return SmVersion >= 100; }
+
   // Prior to CUDA 12.3 ptxas did not recognize that the trap instruction
   // terminates a basic block. Instead, it would assume that control flow
   // continued to the next instruction. The next instruction could be in the