Nvptx port LowerBITCAST to SelectionDAG

llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

@@ -474,7 +474,12 @@ SDValue DAGTypeLegalizer::PromoteIntRes_BITCAST(SDNode *N) {
 
   switch (getTypeAction(InVT)) {
   case TargetLowering::TypeLegal:
+    // Try and use in-register bitcast
+    if (SDValue Res = LowerBitcastInRegister(N))
+      return DAG.getNode(ISD::ANY_EXTEND, dl, NOutVT, Res);
+    // Fallback to stack load store
     break;
+
   case TargetLowering::TypePromoteInteger:
     if (NOutVT.bitsEq(NInVT) && !NOutVT.isVector() && !NInVT.isVector())
       // The input promotes to the same size.  Convert the promoted value.
@@ -2174,8 +2179,11 @@ SDValue DAGTypeLegalizer::PromoteIntOp_ATOMIC_STORE(AtomicSDNode *N) {
 }
 
 SDValue DAGTypeLegalizer::PromoteIntOp_BITCAST(SDNode *N) {
-  // This should only occur in unusual situations like bitcasting to an
-  // x86_fp80, so just turn it into a store+load
+  // Try and use in register bitcast
+  if (SDValue Res = LowerBitcastInRegister(N))
+    return Res;
+
+  // Fallback
   return CreateStackStoreLoad(N->getOperand(0), N->getValueType(0));
 }
 

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp

@@ -910,6 +910,96 @@ SDValue DAGTypeLegalizer::CreateStackStoreLoad(SDValue Op,
   return DAG.getLoad(DestVT, dl, Store, StackPtr, MachinePointerInfo(), Align);
 }
 
+SDValue DAGTypeLegalizer::PackBitcastInRegister(SDNode *N) const {
+  assert(N->getOpcode() == ISD::BITCAST && "Unexpected opcode!");
+
+  EVT FromVT = N->getOperand(0)->getValueType(0);
+  EVT ToVT = N->getValueType(0);
+
+  if (!FromVT.isVector() || !ToVT.isInteger())
+    return SDValue();
+
+  SDLoc DL(N);
+
+  // Get the number of elements we need to pack into the integer
+  unsigned NumElems = FromVT.getVectorNumElements();
+  EVT ElemVT = FromVT.getVectorElementType();
+  unsigned ElemBits = ElemVT.getSizeInBits();
+
+  EVT PackVT = EVT::getIntegerVT(*DAG.getContext(), ElemBits * NumElems);
+  SDValue Packed = DAG.getConstant(0, DL, PackVT);
+
+  // Determine endianness
+  bool IsBigEndian = DAG.getDataLayout().isBigEndian();
+
+  for (unsigned I = 0; I < NumElems; ++I) {
+    unsigned ElementIndex = IsBigEndian ? (NumElems - 1 - I) : I;
+    SDValue Elem =
+        DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, ElemVT, N->getOperand(0),
+                    DAG.getIntPtrConstant(ElementIndex, DL));
+    SDValue ExtElem = DAG.getNode(ISD::ZERO_EXTEND, DL, PackVT, Elem);
+    SDValue ShiftAmount = DAG.getShiftAmountConstant(ElemBits * I, PackVT, DL);
+    SDValue ShiftedElem =
+        DAG.getNode(ISD::SHL, DL, PackVT, ExtElem, ShiftAmount);
+
+    Packed = DAG.getNode(ISD::OR, DL, PackVT, Packed, ShiftedElem);
+  }
+
+  return DAG.getBitcast(ToVT, Packed);
+}
+
+SDValue DAGTypeLegalizer::UnpackBitcastInRegister(SDNode *N) const {
+  assert(N->getOpcode() == ISD::BITCAST && "Unexpected opcode!");
+  EVT FromVT = N->getOperand(0)->getValueType(0);
+  EVT ToVT = N->getValueType(0);
+
+  if (!FromVT.isInteger() || !ToVT.isVector())
+    return SDValue();
+
+  SDLoc DL(N);
+
+  unsigned NumElems = ToVT.getVectorNumElements();
+  EVT ElemVT = ToVT.getVectorElementType();
+  unsigned ElemBits = ElemVT.getSizeInBits();
+
+  // Ensure the integer has enough bits
+  unsigned PackedBits = FromVT.getSizeInBits();
+  assert(PackedBits >= ElemBits * NumElems &&
+         "Packed type does not have enough bits to represent the vector!");
+
+  // Determine endianness
+  bool IsBigEndian = DAG.getDataLayout().isBigEndian();
+
+  // Hold all the vector elements
+  SmallVector<SDValue, 8> Elements;
+  Elements.reserve(NumElems);
+
+  for (unsigned I = 0; I < NumElems; ++I) {
+    unsigned ElementIndex = IsBigEndian ? (NumElems - 1 - I) : I;
+    unsigned ShiftAmountVal = ElemBits * ElementIndex;
+
+    SDValue ShiftAmount =
+        DAG.getShiftAmountConstant(ShiftAmountVal, FromVT, DL);
+    SDValue Shifted =
+        DAG.getNode(ISD::SRL, DL, FromVT, N->getOperand(0), ShiftAmount);
+    SDValue Element = DAG.getNode(ISD::TRUNCATE, DL, ElemVT, Shifted);
+    Elements.push_back(Element);
+  }
+
+  return DAG.getBuildVector(ToVT, DL, Elements);
+}
+
+SDValue DAGTypeLegalizer::LowerBitcastInRegister(SDNode *N) const {
+  // Try the pack, if we aren't going from vector -> scalar it will backout
+  // immediately.
+  if (SDValue Res = PackBitcastInRegister(N)) {
+    return Res;
+  }
+
+  // If we get here then try and unpack the bitcast
+  return UnpackBitcastInRegister(N);
+}
+
 /// Replace the node's results with custom code provided by the target and
 /// return "true", or do nothing and return "false".
 /// The last parameter is FALSE if we are dealing with a node with legal

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h

@@ -216,6 +216,9 @@ class LLVM_LIBRARY_VISIBILITY DAGTypeLegalizer {
   SDValue BitConvertToInteger(SDValue Op);
   SDValue BitConvertVectorToIntegerVector(SDValue Op);
   SDValue CreateStackStoreLoad(SDValue Op, EVT DestVT);
+  SDValue PackBitcastInRegister(SDNode *N) const;
+  SDValue UnpackBitcastInRegister(SDNode *N) const;
+  SDValue LowerBitcastInRegister(SDNode *N) const;
   bool CustomLowerNode(SDNode *N, EVT VT, bool LegalizeResult);
   bool CustomWidenLowerNode(SDNode *N, EVT VT);
 

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

@@ -472,13 +472,6 @@ VectorizePTXValueVTs(const SmallVectorImpl<EVT> &ValueVTs,
   return VectorInfo;
 }
 
-static SDValue MaybeBitcast(SelectionDAG &DAG, SDLoc DL, EVT VT,
-                            SDValue Value) {
-  if (Value->getValueType(0) == VT)
-    return Value;
-  return DAG.getNode(ISD::BITCAST, DL, VT, Value);
-}
-
 // NVPTXTargetLowering Constructor.
 NVPTXTargetLowering::NVPTXTargetLowering(const NVPTXTargetMachine &TM,
                                          const NVPTXSubtarget &STI)
@@ -622,9 +615,6 @@ NVPTXTargetLowering::NVPTXTargetLowering(const NVPTXTargetMachine &TM,
   setOperationAction(ISD::INSERT_VECTOR_ELT, MVT::v4i8, Custom);
   setOperationAction(ISD::VECTOR_SHUFFLE, MVT::v4i8, Custom);
 
-  // Custom conversions to/from v2i8.
-  setOperationAction(ISD::BITCAST, MVT::v2i8, Custom);
-
   // Only logical ops can be done on v4i8 directly, others must be done
   // elementwise.
   setOperationAction(
@@ -2086,30 +2076,6 @@ NVPTXTargetLowering::LowerCONCAT_VECTORS(SDValue Op, SelectionDAG &DAG) const {
   return DAG.getBuildVector(Node->getValueType(0), dl, Ops);
 }
 
-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);
-}
-
 // We can init constant f16x2/v2i16/v4i8 with a single .b32 move.  Normally it
 // would get lowered as two constant loads and vector-packing move.
 // Instead we want just a constant move:
@@ -2618,8 +2584,6 @@ NVPTXTargetLowering::LowerOperation(SDValue Op, SelectionDAG &DAG) const {
     return Op;
   case ISD::BUILD_VECTOR:
     return LowerBUILD_VECTOR(Op, DAG);
-  case ISD::BITCAST:
-    return LowerBITCAST(Op, DAG);
   case ISD::EXTRACT_SUBVECTOR:
     return Op;
   case ISD::EXTRACT_VECTOR_ELT:
@@ -5202,28 +5166,6 @@ SDValue NVPTXTargetLowering::PerformDAGCombine(SDNode *N,
   return SDValue();
 }
 
-static void ReplaceBITCAST(SDNode *Node, SelectionDAG &DAG,
-                           SmallVectorImpl<SDValue> &Results) {
-  // Handle bitcasting to v2i8 without hitting the default promotion
-  // strategy which goes through stack memory.
-  SDValue Op(Node, 0);
-  EVT ToVT = Op->getValueType(0);
-  if (ToVT != MVT::v2i8) {
-    return;
-  }
-
-  // Bitcast to i16 and unpack elements into a vector
-  SDLoc DL(Node);
-  SDValue AsInt = MaybeBitcast(DAG, DL, MVT::i16, Op->getOperand(0));
-  SDValue Vec0 = DAG.getNode(ISD::TRUNCATE, DL, MVT::i8, AsInt);
-  SDValue Const8 = DAG.getConstant(8, DL, MVT::i16);
-  SDValue Vec1 =
-      DAG.getNode(ISD::TRUNCATE, DL, MVT::i8,
-                  DAG.getNode(ISD::SRL, DL, MVT::i16, {AsInt, Const8}));
-  Results.push_back(
-      DAG.getNode(ISD::BUILD_VECTOR, DL, MVT::v2i8, {Vec0, Vec1}));
-}
-
 /// ReplaceVectorLoad - Convert vector loads into multi-output scalar loads.
 static void ReplaceLoadVector(SDNode *N, SelectionDAG &DAG,
                               SmallVectorImpl<SDValue> &Results) {
@@ -5459,9 +5401,6 @@ void NVPTXTargetLowering::ReplaceNodeResults(
   switch (N->getOpcode()) {
   default:
     report_fatal_error("Unhandled custom legalization");
-  case ISD::BITCAST:
-    ReplaceBITCAST(N, DAG, Results);
-    return;
   case ISD::LOAD:
     ReplaceLoadVector(N, DAG, Results);
     return;

llvm/lib/Target/NVPTX/NVPTXISelLowering.h

@@ -265,8 +265,6 @@ class NVPTXTargetLowering : public TargetLowering {
   const NVPTXSubtarget &STI; // cache the subtarget here
   SDValue getParamSymbol(SelectionDAG &DAG, int idx, EVT) const;
 
-  SDValue LowerBITCAST(SDValue Op, SelectionDAG &DAG) 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;