write F32X2 result into two i32 registers

Allows better codegen as each register can be forwarded through
subsequent EXTRACT_VECTOR_ELT nodes.

Author: Prince781

llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp

@@ -196,6 +196,10 @@ void NVPTXDAGToDAGISel::Select(SDNode *N) {
       SelectI128toV2I64(N);
       return;
     }
+    if (N->getOperand(1).getValueType() == MVT::i64 && N->getNumValues() == 3) {
+      SelectI64ToV2I32(N);
+      return;
+    }
     break;
   }
   case ISD::FADD:
@@ -2795,6 +2799,17 @@ void NVPTXDAGToDAGISel::SelectI128toV2I64(SDNode *N) {
   ReplaceNode(N, Mov);
 }
 
+void NVPTXDAGToDAGISel::SelectI64ToV2I32(SDNode *N) {
+  SDValue Ch = N->getOperand(0);
+  SDValue Src = N->getOperand(1);
+  SDLoc DL(N);
+
+  SDNode *Mov = CurDAG->getMachineNode(NVPTX::I64toV2I32, DL,
+                                       {MVT::i32, MVT::i32, Ch.getValueType()},
+                                       {Src, Ch});
+  ReplaceNode(N, Mov);
+}
+
 /// GetConvertOpcode - Returns the CVT_ instruction opcode that implements a
 /// conversion from \p SrcTy to \p DestTy.
 unsigned NVPTXDAGToDAGISel::GetConvertOpcode(MVT DestTy, MVT SrcTy,

llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.h

@@ -92,6 +92,7 @@ class LLVM_LIBRARY_VISIBILITY NVPTXDAGToDAGISel : public SelectionDAGISel {
   bool tryEXTRACT_VECTOR_ELEMENT(SDNode *N);
   void SelectV2I64toI128(SDNode *N);
   void SelectI128toV2I64(SDNode *N);
+  void SelectI64ToV2I32(SDNode *N);
   void SelectCpAsyncBulkG2S(SDNode *N);
   void SelectCpAsyncBulkS2G(SDNode *N);
   void SelectCpAsyncBulkPrefetchL2(SDNode *N);

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

@@ -872,6 +872,8 @@ NVPTXTargetLowering::NVPTXTargetLowering(const NVPTXTargetMachine &TM,
       setOperationAction(Op, MVT::v2f32, Custom);
     // Handle custom lowering for: i64 = bitcast v2f32
     setOperationAction(ISD::BITCAST, MVT::v2f32, Custom);
+    // Handle custom lowering for: f32 = extract_vector_elt v2f32
+    setOperationAction(ISD::EXTRACT_VECTOR_ELT, MVT::v2f32, Custom);
   }
 
   // These map to conversion instructions for scalar FP types.
@@ -2253,6 +2255,20 @@ SDValue NVPTXTargetLowering::LowerEXTRACT_VECTOR_ELT(SDValue Op,
     return DAG.getAnyExtOrTrunc(BFE, DL, Op->getValueType(0));
   }
 
+  if (VectorVT == MVT::v2f32) {
+    if (Vector.getOpcode() == ISD::BITCAST) {
+      // peek through v2f32 = bitcast (i64 = build_pair (i32 A, i32 B))
+      // where A:i32, B:i32 = CopyFromReg (i64 = F32X2 Operation ...)
+      SDValue Pair = Vector.getOperand(0);
+      assert(Pair.getOpcode() == ISD::BUILD_PAIR);
+      return DAG.getNode(
+          ISD::BITCAST, DL, Op.getValueType(),
+          Pair.getOperand(cast<ConstantSDNode>(Index)->getZExtValue()));
+    }
+    if (Vector.getOpcode() == ISD::BUILD_VECTOR)
+      return Vector.getOperand(cast<ConstantSDNode>(Index)->getZExtValue());
+  }
+
   // Constant index will be matched by tablegen.
   if (isa<ConstantSDNode>(Index.getNode()))
     return Op;
@@ -5565,9 +5581,22 @@ static void ReplaceF32x2Op(SDNode *N, SelectionDAG &DAG,
   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 Chain = DAG.getEntryNode();
+
+  // break i64 result into two i32 registers for later instructions that may
+  // access element #0 or #1. otherwise, this code will be eliminated
   SDValue NewValue = DAG.getNode(Opcode, DL, MVT::i64, NewOps);
-  Results.push_back(DAG.getBitcast(OldResultTy, NewValue));
+  MachineRegisterInfo &RegInfo = DAG.getMachineFunction().getRegInfo();
+  Register DestReg = RegInfo.createVirtualRegister(
+      DAG.getTargetLoweringInfo().getRegClassFor(MVT::i64));
+  SDValue RegCopy = DAG.getCopyToReg(Chain, DL, DestReg, NewValue);
+  SDValue Explode = DAG.getNode(ISD::CopyFromReg, DL,
+                                {MVT::i32, MVT::i32, Chain.getValueType()},
+                                {RegCopy, DAG.getRegister(DestReg, MVT::i64)});
+  // cast i64 result of new op back to <2 x float>
+  Results.push_back(DAG.getBitcast(
+      OldResultTy, DAG.getNode(ISD::BUILD_PAIR, DL, MVT::i64,
+                               {Explode.getValue(0), Explode.getValue(1)})));
 }
 
 void NVPTXTargetLowering::ReplaceNodeResults(