[NVPTX] Fixed invalid use of immediate arguments for .f16x2 instructions.

[tra]

Summary

This is a minor change on top of D30057 which fixes a bug that got fixed after D30057 was approved.

Previously (in D30057) ReplaceUses() resulted in a constant being passed as an
immediate argument to an .f16x2 instruction which does not accept
them. It also triggers a crash in ptxas.

This patch moves handling of constant v2f16 BUILD_VECTOR to lowering phase
which allows instruction selection to work correctly.

[tra]

Landed as part of rL296032: [NVPTX] Added support for .f16x2 instructions.

lib/Target/NVPTX/NVPTXISelDAGToDAG.h

#include "NVPTXGenDAGISel.inc"
   bool tryStoreParam(SDNode *N);
   void SelectAddrSpaceCast(SDNode *N);
   bool tryTextureIntrinsic(SDNode *N);
   bool trySurfaceIntrinsic(SDNode *N);
   bool tryBFE(SDNode *N);
   bool tryConstantFP16(SDNode *N);
   bool SelectSETP_F16X2(SDNode *N);
   bool tryEXTRACT_VECTOR_ELEMENT(SDNode *N);
-  bool tryBUILD_VECTOR(SDNode *N);
 
   inline SDValue getI32Imm(unsigned Imm, const SDLoc &DL) {
     return CurDAG->getTargetConstant(Imm, DL, MVT::i32);
   }
 
   // Match direct address complex pattern.
   bool SelectDirectAddr(SDValue N, SDValue &Address);
 

lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp

void NVPTXDAGToDAGISel::Select(SDNode *N) {
   case ISD::STORE:
     if (tryStore(N))
       return;
     break;
   case ISD::EXTRACT_VECTOR_ELT:
     if (tryEXTRACT_VECTOR_ELEMENT(N))
       return;
     break;
-  case ISD::BUILD_VECTOR:
-    if (tryBUILD_VECTOR(N))
-      return;
-    break;
   case NVPTXISD::SETP_F16X2:
     SelectSETP_F16X2(N);
     return;
 
   case NVPTXISD::LoadV2:
   case NVPTXISD::LoadV4:
     if (tryLoadVector(N))
       return;
bool NVPTXDAGToDAGISel::tryEXTRACT_VECTOR_ELEMENT(SDNode *N) {
   for (auto *Node : E0)
     ReplaceUses(SDValue(Node, 0), SDValue(ScatterOp, 0));
   for (auto *Node : E1)
     ReplaceUses(SDValue(Node, 0), SDValue(ScatterOp, 1));
 
   return true;
 }
 
-// We can init constant f16x2 with a single .b32 move.  Normally it
-// would get lowered as two constant loads and vector-packing move.
-//        mov.b16         %h1, 0x4000;
-//        mov.b16         %h2, 0x3C00;
-//        mov.b32         %hh2, {%h2, %h1};
-// Instead we want just a constant move:
-//        mov.b32         %hh2, 0x40003C00
-//
-// This results in better SASS code with CUDA 7.x. Ptxas in CUDA 8.0
-// generates good SASS in both cases.
-
-bool NVPTXDAGToDAGISel::tryBUILD_VECTOR(SDNode *N) {
-  if (!(N->getValueType(0) == MVT::v2f16 &&
-        isa<ConstantFPSDNode>(N->getOperand(0)) &&
-        isa<ConstantFPSDNode>(N->getOperand(1))))
-    return false;
-
-  APInt E0 =
-      cast<ConstantFPSDNode>(N->getOperand(0))->getValueAPF().bitcastToAPInt();
-  APInt E1 =
-      cast<ConstantFPSDNode>(N->getOperand(1))->getValueAPF().bitcastToAPInt();
-  SDValue Const = CurDAG->getTargetConstant(E1.zext(32).shl(16) | E0.zext(32),
-                                            SDLoc(N), MVT::i32);
-  ReplaceUses(SDValue(N, 0), Const);
-  return true;
-}
-
 static unsigned int getCodeAddrSpace(MemSDNode *N) {
   const Value *Src = N->getMemOperand()->getValue();
 
   if (!Src)
     return NVPTX::PTXLdStInstCode::GENERIC;
 
   if (auto *PT = dyn_cast<PointerType>(Src->getType())) {
     switch (PT->getAddressSpace()) {

lib/Target/NVPTX/NVPTXISelLowering.h

public:
   // and avoids calling ctlz in that case.  We have a dedicated ctlz
   // instruction, so we say that ctlz is cheap to speculate.
   bool isCheapToSpeculateCtlz() const override { return true; }
 
 private:
   const NVPTXSubtarget &STI; // cache the subtarget here
   SDValue getParamSymbol(SelectionDAG &DAG, int idx, EVT) const;
 
+  SDValue LowerBUILD_VECTOR(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerCONCAT_VECTORS(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerEXTRACT_VECTOR_ELT(SDValue Op, SelectionDAG &DAG) const;
 
   SDValue LowerLOAD(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerLOADi1(SDValue Op, SelectionDAG &DAG) const;
 
   SDValue LowerSTORE(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerSTOREi1(SDValue Op, SelectionDAG &DAG) const;

lib/Target/NVPTX/NVPTXISelLowering.cpp

NVPTXTargetLowering::NVPTXTargetLowering(const NVPTXTargetMachine &TM,
   addRegisterClass(MVT::f32, &NVPTX::Float32RegsRegClass);
   addRegisterClass(MVT::f64, &NVPTX::Float64RegsRegClass);
   addRegisterClass(MVT::f16, &NVPTX::Float16RegsRegClass);
   addRegisterClass(MVT::v2f16, &NVPTX::Float16x2RegsRegClass);
 
   // Conversion to/from FP16/FP16x2 is always legal.
   setOperationAction(ISD::SINT_TO_FP, MVT::f16, Legal);
   setOperationAction(ISD::FP_TO_SINT, MVT::f16, Legal);
-  setOperationAction(ISD::BUILD_VECTOR, MVT::v2f16, Legal);
+  setOperationAction(ISD::BUILD_VECTOR, MVT::v2f16, Custom);
   setOperationAction(ISD::EXTRACT_VECTOR_ELT, MVT::v2f16, Custom);
 
   setFP16OperationAction(ISD::SETCC, MVT::f16, Legal, Promote);
   setFP16OperationAction(ISD::SETCC, MVT::v2f16, Legal, Expand);
 
   // Operations not directly supported by NVPTX.
   setOperationAction(ISD::SELECT_CC, MVT::f16, Expand);
   setOperationAction(ISD::SELECT_CC, MVT::v2f16, Expand);
for (unsigned i = 0; i < NumOperands; ++i) {
     for (unsigned j = 0; j < NumSubElem; ++j) {
       Ops.push_back(DAG.getNode(ISD::EXTRACT_VECTOR_ELT, dl, EltVT, SubOp,
                                 DAG.getIntPtrConstant(j, dl)));
     }
   }
   return DAG.getBuildVector(Node->getValueType(0), dl, Ops);
 }
 
+// We can init constant f16x2 with a single .b32 move.  Normally it
+// would get lowered as two constant loads and vector-packing move.
+//        mov.b16         %h1, 0x4000;
+//        mov.b16         %h2, 0x3C00;
+//        mov.b32         %hh2, {%h2, %h1};
+// Instead we want just a constant move:
+//        mov.b32         %hh2, 0x40003C00
+//
+// This results in better SASS code with CUDA 7.x. Ptxas in CUDA 8.0
+// generates good SASS in both cases.
+SDValue NVPTXTargetLowering::LowerBUILD_VECTOR(SDValue Op,
+                                               SelectionDAG &DAG) const {
+  //return Op;
+  if (!(Op->getValueType(0) == MVT::v2f16 &&
+        isa<ConstantFPSDNode>(Op->getOperand(0)) &&
+        isa<ConstantFPSDNode>(Op->getOperand(1))))
+    return Op;
+
+  APInt E0 =
+      cast<ConstantFPSDNode>(Op->getOperand(0))->getValueAPF().bitcastToAPInt();
+  APInt E1 =
+      cast<ConstantFPSDNode>(Op->getOperand(1))->getValueAPF().bitcastToAPInt();
+  SDValue Const =
+      DAG.getConstant(E1.zext(32).shl(16) | E0.zext(32), SDLoc(Op), MVT::i32);
+  return DAG.getNode(ISD::BITCAST, SDLoc(Op), MVT::v2f16, Const);
+}
 SDValue NVPTXTargetLowering::LowerEXTRACT_VECTOR_ELT(SDValue Op,
                                                      SelectionDAG &DAG) const {
   SDValue Index = Op->getOperand(1);
   // Constant index will be matched by tablegen.
   if (isa<ConstantSDNode>(Index.getNode()))
     return Op;
 
   // Extract individual elements and select one of them.
case ISD::RETURNADDR:
     return SDValue();
   case ISD::FRAMEADDR:
     return SDValue();
   case ISD::GlobalAddress:
     return LowerGlobalAddress(Op, DAG);
   case ISD::INTRINSIC_W_CHAIN:
     return Op;
   case ISD::BUILD_VECTOR:
+    return LowerBUILD_VECTOR(Op, DAG);
   case ISD::EXTRACT_SUBVECTOR:
     return Op;
   case ISD::EXTRACT_VECTOR_ELT:
     return LowerEXTRACT_VECTOR_ELT(Op, DAG);
   case ISD::CONCAT_VECTORS:
     return LowerCONCAT_VECTORS(Op, DAG);
   case ISD::STORE:
     return LowerSTORE(Op, DAG);

test/CodeGen/NVPTX/f16x2-instructions.ll

 ; ## Full FP16 support enabled by default.
 ; RUN: llc < %s -mtriple=nvptx64-nvidia-cuda -mcpu=sm_53 -asm-verbose=false \
 ; RUN:          -O0 -disable-post-ra -disable-fp-elim \
 ; RUN: | FileCheck -check-prefixes CHECK,CHECK-F16 %s
 ; ## FP16 support explicitly disabled.
 ; RUN: llc < %s -mtriple=nvptx64-nvidia-cuda -mcpu=sm_53 -asm-verbose=false \
 ; RUN:          -O0 -disable-post-ra -disable-fp-elim --nvptx-no-f16-math \
 ; RUN: | FileCheck -check-prefixes CHECK,CHECK-NOF16 %s
 ; ## FP16 is not supported by hardware.
 ; RUN: llc < %s -O0 -mtriple=nvptx64-nvidia-cuda -mcpu=sm_52 -asm-verbose=false \
 ; RUN:          -disable-post-ra -disable-fp-elim \
 ; RUN: | FileCheck -check-prefixes CHECK,CHECK-NOF16 %s
 
 target datalayout = "e-m:o-i64:64-i128:128-n32:64-S128"
 
 ; CHECK-LABEL: test_ret_const(
-; CHECK:     st.param.b32    [func_retval0+0], 1073757184;
+; CHECK:     mov.u32         [[T:%r[0-9+]]], 1073757184;
+; CHECK:     mov.b32         [[R:%hh[0-9+]]], [[T]];
+; CHECK:     st.param.b32    [func_retval0+0], [[R]];
 ; CHECK-NEXT: ret;
 define <2 x half> @test_ret_const() #0 {
   ret <2 x half> <half 1.0, half 2.0>
 }
 
 ; CHECK-LABEL: test_extract_0(
 ; CHECK:      ld.param.b32    [[A:%hh[0-9]+]], [test_extract_0_param_0];
 ; CHECK:      mov.b32         {[[R:%h[0-9]+]], %tmp_hi}, [[A]];
define <2 x half> @test_fadd(<2 x half> %a, <2 x half> %b) #0 {
   %r = fadd <2 x half> %a, %b
   ret <2 x half> %r
 }
 
 ; Check that we can lower fadd with immediate arguments.
 ; CHECK-LABEL: test_fadd_imm_0(
 ; CHECK-DAG:  ld.param.b32    [[A:%hh[0-9]+]], [test_fadd_imm_0_param_0];
 ;
-; CHECK-F16:        add.rn.f16x2   [[R:%hh[0-9]+]], [[A]], 1073757184;
+; CHECK-F16:        mov.u32        [[I:%r[0-9+]]], 1073757184;
+; CHECK-F16:        mov.b32        [[IHH:%hh[0-9+]]], [[I]];
+; CHECK-F16:        add.rn.f16x2   [[R:%hh[0-9]+]], [[A]], [[IHH]];
 ;
 ; CHECK-NOF16-DAG:  mov.b32        {[[A0:%h[0-9]+]], [[A1:%h[0-9]+]]}, [[A]]
 ; CHECK-NOF16-DAG:  cvt.f32.f16    [[FA0:%f[0-9]+]], [[A0]]
 ; CHECK-NOF16-DAG:  cvt.f32.f16    [[FA1:%f[0-9]+]], [[A1]]
 ; CHECK-NOF16-DAG:  add.rn.f32     [[FR0:%f[0-9]+]], [[FA0]], 0f3F800000;
 ; CHECK-NOF16-DAG:  add.rn.f32     [[FR1:%f[0-9]+]], [[FA1]], 0f40000000;
 ; CHECK-NOF16-DAG:  cvt.rn.f16.f32 [[R0:%h[0-9]+]], [[FR0]]
 ; CHECK-NOF16-DAG:  cvt.rn.f16.f32 [[R1:%h[0-9]+]], [[FR1]]
 ; CHECK-NOF16:      mov.b32        [[R:%hh[0-9]+]], {[[R0]], [[R1]]}
 ;
 ; CHECK-NEXT: st.param.b32    [func_retval0+0], [[R]];
 ; CHECK-NEXT: ret;
 define <2 x half> @test_fadd_imm_0(<2 x half> %a) #0 {
   %r = fadd <2 x half> <half 1.0, half 2.0>, %a
   ret <2 x half> %r
 }
 
 ; CHECK-LABEL: test_fadd_imm_1(
 ; CHECK-DAG:  ld.param.b32    [[B:%hh[0-9]+]], [test_fadd_imm_1_param_0];
 ;
-; CHECK-F16:        add.rn.f16x2   [[R:%hh[0-9]+]], [[B]], 1073757184;
+; CHECK-F16:        mov.u32        [[I:%r[0-9+]]], 1073757184;
+; CHECK-F16:        mov.b32        [[IHH:%hh[0-9+]]], [[I]];
+; CHECK-F16:        add.rn.f16x2   [[R:%hh[0-9]+]], [[B]], [[IHH]];
 ;
 ; CHECK-NOF16-DAG:  mov.b32        {[[A0:%h[0-9]+]], [[A1:%h[0-9]+]]}, [[A]]
 ; CHECK-NOF16-DAG:  cvt.f32.f16    [[FA0:%f[0-9]+]], [[A0]]
 ; CHECK-NOF16-DAG:  cvt.f32.f16    [[FA1:%f[0-9]+]], [[A1]]
 ; CHECK-NOF16-DAG:  add.rn.f32     [[FR0:%f[0-9]+]], [[FA0]], 0f3F800000;
 ; CHECK-NOF16-DAG:  add.rn.f32     [[FR1:%f[0-9]+]], [[FA1]], 0f40000000;
 ; CHECK-NOF16-DAG:  cvt.rn.f16.f32 [[R0:%h[0-9]+]], [[FR0]]
 ; CHECK-NOF16-DAG:  cvt.rn.f16.f32 [[R1:%h[0-9]+]], [[FR1]]

 define <2 x half> @test_fsub(<2 x half> %a, <2 x half> %b) #0 {
   %r = fsub <2 x half> %a, %b
   ret <2 x half> %r
 }
 
 ; CHECK-LABEL: test_fneg(
 ; CHECK-DAG:  ld.param.b32    [[A:%hh[0-9]+]], [test_fneg_param_0];
 ;
-; CHECK-F16-NEXT:   sub.rn.f16x2   [[R:%hh[0-9]+]], 0, [[A]];
+; CHECK-F16:        mov.u32        [[I0:%r[0-9+]]], 0;
+; CHECK-F16:        mov.b32        [[IHH0:%hh[0-9+]]], [[I0]];
+; CHECK-F16-NEXT:   sub.rn.f16x2   [[R:%hh[0-9]+]], [[IHH0]], [[A]];
 ;
 ; CHECK-NOF16-DAG:  mov.b32        {[[A0:%h[0-9]+]], [[A1:%h[0-9]+]]}, [[A]]
 ; CHECK-NOF16-DAG:  cvt.f32.f16    [[FA0:%f[0-9]+]], [[A0]]
 ; CHECK-NOF16-DAG:  cvt.f32.f16    [[FA1:%f[0-9]+]], [[A1]]
 ; CHECK-NOF16-DAG:  mov.f32        [[Z:%f[0-9]+]], 0f00000000;
 ; CHECK-NOF16-DAG:  sub.rn.f32     [[FR0:%f[0-9]+]], [[Z]], [[FA0]];
 ; CHECK-NOF16-DAG:  sub.rn.f32     [[FR1:%f[0-9]+]], [[Z]], [[FA1]];
 ; CHECK-NOF16-DAG:  cvt.rn.f16.f32 [[R0:%h[0-9]+]], [[FR0]]