enable ftz support

And temporarily disable -O3 in testing as it exposes an existing bug
with how test_extract_i() is lowered when optimized.

Author: Prince781

llvm/include/llvm/Target/TargetSelectionDAG.td

@@ -115,6 +115,9 @@ def SDTPtrAddOp : SDTypeProfile<1, 2, [     // ptradd
 def SDTIntBinOp : SDTypeProfile<1, 2, [     // add, and, or, xor, udiv, etc.
   SDTCisSameAs<0, 1>, SDTCisSameAs<0, 2>, SDTCisInt<0>
 ]>;
+def SDTIntTernaryOp : SDTypeProfile<1, 3, [  // fma32x2
+  SDTCisSameAs<0, 1>, SDTCisSameAs<0, 2>, SDTCisSameAs<0, 3>, SDTCisInt<0>
+]>;
 def SDTIntShiftOp : SDTypeProfile<1, 2, [   // shl, sra, srl
   SDTCisSameAs<0, 1>, SDTCisInt<0>, SDTCisInt<2>
 ]>;

llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp

@@ -121,12 +121,6 @@ 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))
@@ -305,30 +299,6 @@ 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,7 +88,6 @@ 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/NVPTXIntrinsics.td

@@ -1581,27 +1581,50 @@ 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]>;
+// packed f32 ops (sm_100+)
+class F32x2Op2<string OpcStr, Predicate Pred>
+: NVPTXInst<(outs Int64Regs:$res),
+            (ins Int64Regs:$a, Int64Regs:$b),
+            OpcStr # ".f32x2 \t$res, $a, $b;", []>,
+  Requires<[hasF32x2Instructions, Pred]>;
+class F32x2Op3<string OpcStr, Predicate Pred>
+: NVPTXInst<(outs Int64Regs:$res),
+            (ins Int64Regs:$a, Int64Regs:$b, Int64Regs:$c),
+            OpcStr # ".f32x2 \t$res, $a, $b, $c;", []>,
+  Requires<[hasF32x2Instructions, Pred]>;
+
+def fadd32x2_nvptx : SDNode<"NVPTXISD::FADD_F32X2", SDTIntBinOp>;
+def fsub32x2_nvptx : SDNode<"NVPTXISD::FSUB_F32X2", SDTIntBinOp>;
+def fmul32x2_nvptx : SDNode<"NVPTXISD::FMUL_F32X2", SDTIntBinOp>;
+def fma32x2_nvptx  : SDNode<"NVPTXISD::FMA_F32X2", SDTIntTernaryOp>;
+
+def FADD32x2     : F32x2Op2<"add.rn", doNoF32FTZ>;
+def FSUB32x2     : F32x2Op2<"sub.rn", doNoF32FTZ>;
+def FMUL32x2     : F32x2Op2<"mul.rn", doNoF32FTZ>;
+def FMA32x2      : F32x2Op3<"fma.rn", doNoF32FTZ>;
+
+def : Pat<(fadd32x2_nvptx i64:$a, i64:$b),
+          (FADD32x2 $a, $b)>, Requires<[doNoF32FTZ]>;
+def : Pat<(fsub32x2_nvptx i64:$a, i64:$b),
+          (FSUB32x2 $a, $b)>, Requires<[doNoF32FTZ]>;
+def : Pat<(fmul32x2_nvptx i64:$a, i64:$b),
+          (FMUL32x2 $a, $b)>, Requires<[doNoF32FTZ]>;
+def : Pat<(fma32x2_nvptx i64:$a, i64:$b, i64:$c),
+          (FMA32x2 $a, $b, $c)>, Requires<[doNoF32FTZ]>;
+
+def FADD32x2_ftz : F32x2Op2<"add.rn.ftz", doF32FTZ>;
+def FSUB32x2_ftz : F32x2Op2<"sub.rn.ftz", doF32FTZ>;
+def FMUL32x2_ftz : F32x2Op2<"mul.rn.ftz", doF32FTZ>;
+def FMA32x2_ftz  : F32x2Op3<"fma.rn.ftz", doF32FTZ>;
+
+def : Pat<(fadd32x2_nvptx i64:$a, i64:$b),
+          (FADD32x2_ftz $a, $b)>, Requires<[doF32FTZ]>;
+def : Pat<(fsub32x2_nvptx i64:$a, i64:$b),
+          (FSUB32x2_ftz $a, $b)>, Requires<[doF32FTZ]>;
+def : Pat<(fmul32x2_nvptx i64:$a, i64:$b),
+          (FMUL32x2_ftz $a, $b)>, Requires<[doF32FTZ]>;
+def : Pat<(fma32x2_nvptx i64:$a, i64:$b, i64:$c),
+          (FMA32x2_ftz $a, $b, $c)>, Requires<[doF32FTZ]>;
 
 //
 // BFIND